From a04727ea7a1b4fced84fa3682f07f8ea0092bb68 Mon Sep 17 00:00:00 2001 From: Nanduu24 Date: Thu, 24 Sep 2026 22:11:25 -0500 Subject: [PATCH] feat(memory): add self-hosted PostgreSQL/pgvector memory service Adds PgVectorMemoryService under google.adk.integrations.pgvector, a BaseMemoryService backed by PostgreSQL with the pgvector extension. Events are embedded (via the google-genai client by default, or an injectable embedder) and stored with their vectors in a table the user owns; search_memory returns the nearest memories by cosine distance using a pgvector HNSW index. Unlike VertexAiMemoryBankService and VertexAiRagMemoryService it needs no Google Cloud, and unlike InMemoryMemoryService it persists across restarts and ranks by semantic similarity rather than keyword overlap. This fills the self-hosted, non-GCP semantic-memory gap raised in #6254 and complements the exact-match SQLite memory service in #4116. The service follows the existing integrations/redis pattern: a pydantic config plus an injectable connection pool and embedder, which keeps it unit-testable without a live database. Ships as an optional extra, google-adk[pgvector]. Includes unit tests that exercise ingestion, idempotent upserts, semantic ranking, per-(app_name, user_id) scoping, the distance threshold, and serialization round-trips against an in-process fake pool and a deterministic embedder, so they run without psycopg, pgvector, or any network access. Resolves #7273 --- pyproject.toml | 8 + .../adk/integrations/pgvector/README.md | 104 +++++ .../adk/integrations/pgvector/__init__.py | 25 + .../adk/integrations/pgvector/_config.py | 80 ++++ .../pgvector/_pgvector_memory_service.py | 440 ++++++++++++++++++ .../integrations/pgvector/__init__.py | 13 + .../pgvector/test_pgvector_memory_service.py | 429 +++++++++++++++++ 7 files changed, 1099 insertions(+) create mode 100644 src/google/adk/integrations/pgvector/README.md create mode 100644 src/google/adk/integrations/pgvector/__init__.py create mode 100644 src/google/adk/integrations/pgvector/_config.py create mode 100644 src/google/adk/integrations/pgvector/_pgvector_memory_service.py create mode 100644 tests/unittests/integrations/pgvector/__init__.py create mode 100644 tests/unittests/integrations/pgvector/test_pgvector_memory_service.py diff --git a/pyproject.toml b/pyproject.toml index ef23e4dc380..0ca616be4c0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -122,8 +122,11 @@ optional-dependencies.all = [ "opentelemetry-instrumentation-httpx>=0.54b0,<1", "opentelemetry-resourcedetector-gcp>=1.9.0a0,<2", "pandas>=2.2.3", + "pgvector>=0.3", "pillow>=10.3", "protobuf>=6", + "psycopg-pool>=3.2", + "psycopg[binary]>=3.1", "pyarrow>=14", "pymongo>=4.9,<5", "pypika>=0.50", @@ -275,6 +278,11 @@ optional-dependencies.otel-gcp = [ "opentelemetry-instrumentation-grpc>=0.43b0,<1", "opentelemetry-instrumentation-httpx>=0.54b0,<1", ] +optional-dependencies.pgvector = [ + "pgvector>=0.3", # registers the pgvector type on the connection. + "psycopg[binary]>=3.1", # async driver used by PgVectorMemoryService. + "psycopg-pool>=3.2", # AsyncConnectionPool. +] optional-dependencies.redis = [ "redis>=4.2", # 4.2 is where redis.asyncio landed, the only API used. ] diff --git a/src/google/adk/integrations/pgvector/README.md b/src/google/adk/integrations/pgvector/README.md new file mode 100644 index 00000000000..ff708f9e6b5 --- /dev/null +++ b/src/google/adk/integrations/pgvector/README.md @@ -0,0 +1,104 @@ +# PostgreSQL / pgvector Memory Integration for ADK + +This integration provides a self-hosted, semantic memory service for the Google +Agent Development Kit (ADK), backed by PostgreSQL with the +[pgvector](https://github.com/pgvector/pgvector) extension. + +It complements the built-in memory services: + +- `InMemoryMemoryService` is keyword-only and loses data on restart (prototyping + only). +- `VertexAiMemoryBankService` and `VertexAiRagMemoryService` provide semantic + memory but require Google Cloud. + +`PgVectorMemoryService` gives you semantic (vector) memory that persists across +restarts and runs entirely on a database you control, which suits on-premise, +air-gapped, and data-residency-constrained deployments. + +## Features + +- **Semantic search:** Ranks memories by cosine similarity over embeddings using + a pgvector HNSW index, instead of keyword overlap. +- **Self-hosted:** Vectors live in your own PostgreSQL database. No managed Cloud + service is required for storage. +- **Idempotent ingestion:** Re-ingesting a session updates existing rows in + place rather than creating duplicates. +- **Scoped memory:** Memories are isolated per `(app_name, user_id)`. +- **Pluggable embeddings:** Embeds through the `google-genai` client by default, + or any provider via an injected `embedder` callable. + +## Installation / Dependencies + +Install the optional `pgvector` extra alongside ADK: + +```bash +pip install "google-adk[pgvector]" +``` + +You also need a PostgreSQL database with the `vector` extension available. The +service creates the extension, table, and indexes on first use. + +## Quick Start + +```python +from google.adk.integrations.pgvector import PgVectorMemoryService +from google.adk.integrations.pgvector import PgVectorMemoryServiceConfig +from google.adk.runners import Runner + +# 1. Configure the memory service. +memory_service = PgVectorMemoryService( + PgVectorMemoryServiceConfig( + dsn="postgresql://user:password@localhost:5432/adk", + embedding_model="gemini-embedding-001", + embedding_dimension=768, + ) +) + +# 2. Wire it into your Runner. +runner = Runner( + app_name="my_app", + agent=agent, + memory_service=memory_service, +) + +# Ingesting and searching: +await memory_service.add_session_to_memory(session) +response = await memory_service.search_memory( + app_name="my_app", user_id="user1", query="what did we decide about billing?" +) +``` + +## Configuration + +`PgVectorMemoryServiceConfig` supports the following options: + +| Option | Default | Description | +| --- | --- | --- | +| `dsn` | `None` | PostgreSQL connection string. Required unless a `connection_pool` is injected. | +| `table_name` | `adk_memory_entries` | Table that stores memory entries. | +| `embedding_model` | `gemini-embedding-001` | `google-genai` embedding model. | +| `embedding_dimension` | `768` | Stored vector dimension; must match the model output and stays fixed for the table. | +| `top_k` | `10` | Maximum memories returned per search. | +| `hnsw_m` | `16` | pgvector HNSW `m` parameter. | +| `hnsw_ef_construction` | `64` | pgvector HNSW `ef_construction` parameter. | +| `distance_threshold` | `None` | Optional cosine-distance ceiling; farther memories are dropped. | + +## Advanced usage + +Inject a pre-configured pool or a non-Google embedding provider: + +```python +from psycopg_pool import AsyncConnectionPool + +pool = AsyncConnectionPool("postgresql://user:password@localhost:5432/adk") + +async def my_embedder(texts): + # Return one vector per text from any provider. + ... + +memory_service = PgVectorMemoryService( + PgVectorMemoryServiceConfig(embedding_dimension=1024), + connection_pool=pool, + embedder=my_embedder, +) +``` diff --git a/src/google/adk/integrations/pgvector/__init__.py b/src/google/adk/integrations/pgvector/__init__.py new file mode 100644 index 00000000000..2c00ae90bd0 --- /dev/null +++ b/src/google/adk/integrations/pgvector/__init__.py @@ -0,0 +1,25 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""PostgreSQL/pgvector integrations for ADK.""" + +from __future__ import annotations + +from ._config import PgVectorMemoryServiceConfig +from ._pgvector_memory_service import PgVectorMemoryService + +__all__ = [ + "PgVectorMemoryService", + "PgVectorMemoryServiceConfig", +] diff --git a/src/google/adk/integrations/pgvector/_config.py b/src/google/adk/integrations/pgvector/_config.py new file mode 100644 index 00000000000..be5051dadfe --- /dev/null +++ b/src/google/adk/integrations/pgvector/_config.py @@ -0,0 +1,80 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Configuration for the PostgreSQL/pgvector integrations.""" + +from __future__ import annotations + +from typing import Optional + +from pydantic import BaseModel +from pydantic import Field + + +class PgVectorMemoryServiceConfig(BaseModel): + """Configuration for PgVectorMemoryService. + + The embedding model is called through the ``google-genai`` client, but the + vectors themselves live in a PostgreSQL database you control, so the memory + store does not depend on any managed Google Cloud service. + """ + + dsn: Optional[str] = Field( + default=None, + description=( + "PostgreSQL connection string, e.g." + " postgresql://user:password@host:5432/dbname. Required unless a" + " pre-configured connection pool is passed to the service." + ), + ) + table_name: str = Field( + default="adk_memory_entries", + description="Name of the table that stores memory entries.", + ) + embedding_model: str = Field( + default="gemini-embedding-001", + description=( + "google-genai embedding model used to embed events and queries." + ), + ) + embedding_dimension: int = Field( + default=768, + description=( + "Dimension of the stored embedding vectors. Must match the output" + " dimension of embedding_model and stays fixed for the life of the" + " table." + ), + ) + top_k: int = Field( + default=10, + description="Maximum number of memories returned by a search.", + ) + hnsw_m: int = Field( + default=16, + description="pgvector HNSW index parameter m (graph degree).", + ) + hnsw_ef_construction: int = Field( + default=64, + description=( + "pgvector HNSW index parameter ef_construction (build effort)." + ), + ) + distance_threshold: Optional[float] = Field( + default=None, + description=( + "Optional cosine-distance ceiling in [0, 2]. When set, memories whose" + " distance to the query is greater than this are dropped from the" + " results. When None, the closest top_k memories are always returned." + ), + ) diff --git a/src/google/adk/integrations/pgvector/_pgvector_memory_service.py b/src/google/adk/integrations/pgvector/_pgvector_memory_service.py new file mode 100644 index 00000000000..85e2e3da48e --- /dev/null +++ b/src/google/adk/integrations/pgvector/_pgvector_memory_service.py @@ -0,0 +1,440 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""PostgreSQL/pgvector-backed memory service implementation for ADK.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable +from collections.abc import Callable +from collections.abc import Mapping +from collections.abc import Sequence +import hashlib +import json +import logging +from typing import Any +from typing import Optional +from typing import TYPE_CHECKING + +from google.genai import types +from typing_extensions import override + +from ...memory import _utils +from ...memory.base_memory_service import BaseMemoryService +from ...memory.base_memory_service import SearchMemoryResponse +from ...memory.memory_entry import MemoryEntry +from ._config import PgVectorMemoryServiceConfig + +if TYPE_CHECKING: + from ...events.event import Event + from ...sessions.session import Session + +try: + from psycopg_pool import AsyncConnectionPool +except ImportError: + AsyncConnectionPool = None + +try: + from pgvector.psycopg import register_vector_async +except ImportError: + register_vector_async = None + +logger = logging.getLogger("google_adk." + __name__) + +# An async callable that turns a batch of texts into their embedding vectors. +Embedder = Callable[[Sequence[str]], Awaitable[Sequence[Sequence[float]]]] + + +def _content_text(content: Optional[types.Content]) -> str: + """Joins the text parts of a content into a single string.""" + if not content or not content.parts: + return "" + return " ".join(part.text for part in content.parts if part.text) + + +class PgVectorMemoryService(BaseMemoryService): + """A memory service backed by PostgreSQL with the pgvector extension. + + Events are embedded with a ``google-genai`` embedding model and stored, with + their vectors, in a PostgreSQL table you own. ``search_memory`` returns the + memories whose embeddings are closest to the query embedding by cosine + distance, using a pgvector HNSW index. Unlike ``VertexAiMemoryBankService`` + and ``VertexAiRagMemoryService`` the store runs entirely on your own database, + so it suits on-premise, air-gapped, and data-residency-constrained + deployments; unlike ``InMemoryMemoryService`` it persists across restarts and + ranks by semantic similarity rather than keyword overlap. + + Memories are scoped by ``(app_name, user_id)``. Ingesting the same event more + than once updates the existing row in place rather than creating duplicates. + + Example:: + + from google.adk.integrations.pgvector import PgVectorMemoryService + from google.adk.integrations.pgvector import PgVectorMemoryServiceConfig + + memory_service = PgVectorMemoryService( + PgVectorMemoryServiceConfig( + dsn="postgresql://user:password@localhost:5432/adk", + ) + ) + + The connection pool and the embedding function can both be injected, which + makes the service usable with a pre-tuned pool or a non-Google embedding + provider, and testable without a live database. + """ + + def __init__( + self, + config: Optional[PgVectorMemoryServiceConfig] = None, + *, + connection_pool: Optional[Any] = None, + genai_client: Optional[Any] = None, + embedder: Optional[Embedder] = None, + ): + """Initializes the PgVectorMemoryService. + + Args: + config: Configuration for the service. Defaults to + ``PgVectorMemoryServiceConfig()``; a ``dsn`` is required unless + ``connection_pool`` is supplied. + connection_pool: Optional pre-configured ``psycopg_pool`` + ``AsyncConnectionPool``. When omitted, one is created lazily from + ``config.dsn``. + genai_client: Optional ``google.genai.Client`` used for the default + embedder. Ignored when ``embedder`` is supplied. + embedder: Optional async callable that embeds a batch of texts. When + omitted, texts are embedded with ``config.embedding_model`` through the + ``google-genai`` client. + """ + self.config = config or PgVectorMemoryServiceConfig() + self._pool = connection_pool + self._owns_pool = connection_pool is None + self._pool_opened = False + self._genai_client = genai_client + self._embedder = embedder + self._schema_ready = False + self._schema_lock = asyncio.Lock() + + # --- Public BaseMemoryService API ---------------------------------------- + + @override + async def add_session_to_memory(self, session: Session) -> None: + await self._add_events( + app_name=session.app_name, + user_id=session.user_id, + events=session.events, + session_id=session.id, + ) + + @override + async def add_events_to_memory( + self, + *, + app_name: str, + user_id: str, + events: Sequence[Event], + session_id: str | None = None, + custom_metadata: Mapping[str, object] | None = None, + ) -> None: + await self._add_events( + app_name=app_name, + user_id=user_id, + events=events, + session_id=session_id, + custom_metadata=custom_metadata, + ) + + @override + async def add_memory( + self, + *, + app_name: str, + user_id: str, + memories: Sequence[MemoryEntry], + custom_metadata: Mapping[str, object] | None = None, + ) -> None: + shared_metadata = dict(custom_metadata) if custom_metadata else {} + rows: list[dict[str, Any]] = [] + texts: list[str] = [] + for memory in memories: + text = _content_text(memory.content) + if not text: + continue + metadata = {**shared_metadata, **(memory.custom_metadata or {})} + key = memory.id or _content_text(memory.content) + rows.append({ + "id": self._entry_id(app_name, user_id, None, key), + "app_name": app_name, + "user_id": user_id, + "session_id": None, + "author": memory.author, + "timestamp": memory.timestamp, + "text": text, + "content": memory.content, + "custom_metadata": metadata, + }) + texts.append(text) + await self._embed_and_upsert(rows, texts) + + @override + async def search_memory( + self, *, app_name: str, user_id: str, query: str + ) -> SearchMemoryResponse: + if not query or not query.strip(): + return SearchMemoryResponse() + + embedding = (await self._embed([query]))[0] + pool = await self._ensure_pool() + async with pool.connection() as conn: + await self._prepare_connection(conn) + cursor = await conn.execute( + "SELECT author, timestamp, content, custom_metadata, id," + f" embedding <=> %s::vector AS distance FROM {self._table}" + " WHERE app_name = %s AND user_id = %s" + " ORDER BY distance LIMIT %s", + (self._to_vector(embedding), app_name, user_id, self.config.top_k), + ) + rows = await cursor.fetchall() + + threshold = self.config.distance_threshold + memories: list[MemoryEntry] = [] + for author, timestamp, content, custom_metadata, entry_id, distance in rows: + if ( + threshold is not None + and distance is not None + and distance > threshold + ): + continue + memories.append( + MemoryEntry( + content=types.Content.model_validate(_as_dict(content)), + author=author, + timestamp=timestamp, + custom_metadata=_as_dict(custom_metadata), + id=entry_id, + ) + ) + return SearchMemoryResponse(memories=memories) + + # --- Internals ------------------------------------------------------------ + + @property + def _table(self) -> str: + return self.config.table_name + + async def _add_events( + self, + *, + app_name: str, + user_id: str, + events: Sequence[Event], + session_id: str | None, + custom_metadata: Mapping[str, object] | None = None, + ) -> None: + shared_metadata = dict(custom_metadata) if custom_metadata else {} + rows: list[dict[str, Any]] = [] + texts: list[str] = [] + for event in events: + text = _content_text(event.content) + if not text: + continue + key = event.id or f"{session_id}:{event.timestamp}:{text}" + rows.append({ + "id": self._entry_id(app_name, user_id, session_id, key), + "app_name": app_name, + "user_id": user_id, + "session_id": session_id, + "author": event.author, + "timestamp": _utils.format_timestamp(event.timestamp), + "text": text, + "content": event.content, + "custom_metadata": shared_metadata, + }) + texts.append(text) + await self._embed_and_upsert(rows, texts) + + async def _embed_and_upsert( + self, rows: list[dict[str, Any]], texts: list[str] + ) -> None: + if not rows: + return + embeddings = await self._embed(texts) + pool = await self._ensure_pool() + async with pool.connection() as conn: + await self._prepare_connection(conn) + for row, embedding in zip(rows, embeddings): + await conn.execute( + f"INSERT INTO {self._table} (id, app_name, user_id, session_id," + " author, timestamp, text, content, custom_metadata, embedding)" + " VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb, %s::jsonb," + " %s::vector)" + " ON CONFLICT (id) DO UPDATE SET text = EXCLUDED.text," + " content = EXCLUDED.content," + " custom_metadata = EXCLUDED.custom_metadata," + " embedding = EXCLUDED.embedding", + ( + row["id"], + row["app_name"], + row["user_id"], + row["session_id"], + row["author"], + row["timestamp"], + row["text"], + _content_json(row["content"]), + json.dumps(row["custom_metadata"]), + self._to_vector(embedding), + ), + ) + + def _entry_id( + self, + app_name: str, + user_id: str, + session_id: str | None, + key: str, + ) -> str: + """Builds a stable, collision-resistant id for a memory row. + + The id is derived from the memory's scope and a per-event key so that + re-ingesting the same event updates its row in place instead of inserting a + duplicate. + """ + digest = hashlib.sha256( + "\x00".join([app_name, user_id, session_id or "", key]).encode("utf-8") + ).hexdigest() + return digest + + async def _embed(self, texts: Sequence[str]) -> list[list[float]]: + embedder = self._embedder or self._default_embed + vectors = await embedder(list(texts)) + return [[float(v) for v in vector] for vector in vectors] + + async def _default_embed(self, texts: Sequence[str]) -> list[list[float]]: + """Embeds texts with the configured google-genai embedding model.""" + from google.genai import Client # pylint: disable=import-outside-toplevel + + client = self._genai_client or Client() + config = types.EmbedContentConfig() + if self.config.embedding_dimension: + config.output_dimensionality = self.config.embedding_dimension + try: + response = await client.aio.models.embed_content( + model=self.config.embedding_model, + contents=list(texts), + config=config, + ) + except Exception as ex: + raise RuntimeError(f"Failed to embed content: {ex!r}") from ex + return [list(embedding.values) for embedding in response.embeddings] + + def _to_vector(self, embedding: Sequence[float]) -> Any: + """Returns the embedding in the form the driver stores as a pgvector value. + + pgvector adapts a plain list of floats, so the list is passed through as-is. + Keeping this in one place lets a different adapter (for example a numpy + array) be swapped in without touching the queries. + """ + return list(embedding) + + def _get_pool(self) -> Any: + """Lazily creates and returns the connection pool.""" + if self._pool is not None: + return self._pool + if AsyncConnectionPool is None: + raise ImportError( + "PgVectorMemoryService requires the psycopg connection pool. Install" + " the optional dependencies with `pip install" + ' "google-adk[pgvector]"`.' + ) + if not self.config.dsn: + raise ValueError( + "PgVectorMemoryServiceConfig.dsn is required when a connection_pool" + " is not provided." + ) + self._pool = AsyncConnectionPool(self.config.dsn, open=False) + return self._pool + + async def _ensure_pool(self) -> Any: + """Returns the pool, opening it once if this service created it. + + An injected pool is assumed to be managed (opened and closed) by the + caller, mirroring how the Redis integration accepts a pre-configured + client. + """ + pool = self._get_pool() + if self._owns_pool and not self._pool_opened: + await pool.open() + self._pool_opened = True + return pool + + async def close(self) -> None: + """Closes the connection pool if this service created it.""" + if self._owns_pool and self._pool is not None and self._pool_opened: + await self._pool.close() + self._pool_opened = False + + async def _prepare_connection(self, conn: Any) -> None: + """Registers the vector type and makes sure the schema exists.""" + if register_vector_async is not None: + await register_vector_async(conn) + if self._schema_ready: + return + async with self._schema_lock: + if self._schema_ready: + return + await self._ensure_schema(conn) + self._schema_ready = True + + async def _ensure_schema(self, conn: Any) -> None: + """Creates the pgvector extension, table, and indexes if they are absent.""" + dimension = self.config.embedding_dimension + await conn.execute("CREATE EXTENSION IF NOT EXISTS vector") + await conn.execute( + f"CREATE TABLE IF NOT EXISTS {self._table} (" + " id text PRIMARY KEY," + " app_name text NOT NULL," + " user_id text NOT NULL," + " session_id text," + " author text," + " timestamp text," + " text text NOT NULL," + " content jsonb NOT NULL," + " custom_metadata jsonb NOT NULL DEFAULT '{}'::jsonb," + f" embedding vector({dimension}) NOT NULL," + " created_at timestamptz NOT NULL DEFAULT now()" + ")" + ) + await conn.execute( + f"CREATE INDEX IF NOT EXISTS {self._table}_app_user_idx" + f" ON {self._table} (app_name, user_id)" + ) + await conn.execute( + f"CREATE INDEX IF NOT EXISTS {self._table}_embedding_idx" + f" ON {self._table} USING hnsw (embedding vector_cosine_ops)" + f" WITH (m = {self.config.hnsw_m}," + f" ef_construction = {self.config.hnsw_ef_construction})" + ) + + +def _content_json(content: types.Content) -> str: + """Serializes a content to a JSON string for a jsonb column.""" + return json.dumps(content.model_dump(mode="json", exclude_none=True)) + + +def _as_dict(value: Any) -> dict[str, Any]: + """Normalizes a jsonb column value, which a driver may hand back as text.""" + if isinstance(value, str): + return json.loads(value) + return value or {} diff --git a/tests/unittests/integrations/pgvector/__init__.py b/tests/unittests/integrations/pgvector/__init__.py new file mode 100644 index 00000000000..58d482ea386 --- /dev/null +++ b/tests/unittests/integrations/pgvector/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unittests/integrations/pgvector/test_pgvector_memory_service.py b/tests/unittests/integrations/pgvector/test_pgvector_memory_service.py new file mode 100644 index 00000000000..7679777f0ab --- /dev/null +++ b/tests/unittests/integrations/pgvector/test_pgvector_memory_service.py @@ -0,0 +1,429 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for PgVectorMemoryService. + +The tests run against an in-process fake connection pool and a deterministic +bag-of-words embedder, so they exercise the service's ingestion, ranking, and +serialization logic without a live PostgreSQL database, the psycopg driver, or +any network calls. +""" + +from __future__ import annotations + +import hashlib +import math +import re +from typing import Any + +from google.adk.events.event import Event +from google.adk.integrations.pgvector import _pgvector_memory_service +from google.adk.integrations.pgvector import PgVectorMemoryService +from google.adk.integrations.pgvector import PgVectorMemoryServiceConfig +from google.adk.memory.memory_entry import MemoryEntry +from google.adk.sessions.session import Session +from google.genai import types +import pytest + +_DIM = 64 + + +@pytest.fixture(autouse=True) +def _no_vector_registration(monkeypatch): + """Skips pgvector's connection type registration in unit tests. + + The fake connection is not a real psycopg connection, so registering the + pgvector type (which queries a live database) is neither possible nor + meaningful here; that path is covered by the end-to-end tests. Forcing it off + also keeps the suite hermetic whether or not pgvector is installed, matching + CI where the optional driver is absent. + """ + monkeypatch.setattr( + _pgvector_memory_service, "register_vector_async", None, raising=False + ) + + +def _bag_of_words_embedder(dim: int = _DIM): + """Returns a deterministic embedder that hashes tokens into buckets. + + Texts that share words get overlapping vectors, so cosine distance ranks a + query nearest the events that share its words - enough to test ranking + without a real embedding model. + """ + + async def embed(texts): + vectors = [] + for text in texts: + vector = [0.0] * dim + for token in re.findall(r"\w+", text.lower()): + bucket = int(hashlib.md5(token.encode()).hexdigest(), 16) % dim + vector[bucket] += 1.0 + vectors.append(vector) + return vectors + + return embed + + +def _cosine_distance(a, b) -> float: + dot = sum(x * y for x, y in zip(a, b)) + norm_a = math.sqrt(sum(x * x for x in a)) + norm_b = math.sqrt(sum(y * y for y in b)) + if norm_a == 0 or norm_b == 0: + return 1.0 + return 1.0 - dot / (norm_a * norm_b) + + +class _FakeCursor: + + def __init__(self, rows): + self._rows = rows + + async def fetchall(self): + return self._rows + + +class _FakeConnection: + """Interprets the small set of statements PgVectorMemoryService issues.""" + + def __init__(self, store: dict[str, tuple[Any, ...]]): + self._store = store + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + async def execute(self, sql: str, params: tuple[Any, ...] | None = None): + keyword = sql.strip().split(None, 1)[0].upper() + if keyword in ("CREATE",): + return _FakeCursor([]) + if keyword == "INSERT": + # params: id, app, user, session, author, timestamp, text, content, + # custom_metadata, embedding + self._store[params[0]] = params + return _FakeCursor([]) + if keyword == "SELECT": + # params: embedding, app_name, user_id, top_k + query_vec, app_name, user_id, top_k = params + scored = [] + for row in self._store.values(): + if row[1] != app_name or row[2] != user_id: + continue + distance = _cosine_distance(query_vec, row[9]) + # author, timestamp, content, custom_metadata, id, distance + scored.append((row[4], row[5], row[7], row[8], row[0], distance)) + scored.sort(key=lambda r: r[5]) + return _FakeCursor(scored[:top_k]) + return _FakeCursor([]) + + +class _FakePgVectorPool: + + def __init__(self): + self.store: dict[str, tuple[Any, ...]] = {} + self.open_count = 0 + self.closed = False + + async def open(self): + self.open_count += 1 + + async def close(self): + self.closed = True + + def connection(self): + return _FakeConnection(self.store) + + +def _event(author: str, text: str, timestamp: float = 12345.0) -> Event: + return Event( + author=author, + timestamp=timestamp, + content=types.Content(parts=[types.Part(text=text)]), + ) + + +def _session(app_name: str, user_id: str, session_id: str, events) -> Session: + return Session( + app_name=app_name, + user_id=user_id, + id=session_id, + events=events, + ) + + +def _make_service(pool, **config_kwargs): + config = PgVectorMemoryServiceConfig( + dsn="postgresql://ignored", + embedding_dimension=_DIM, + **config_kwargs, + ) + return PgVectorMemoryService( + config, + connection_pool=pool, + embedder=_bag_of_words_embedder(), + ) + + +@pytest.mark.asyncio +async def test_add_session_and_search_returns_semantically_closest(): + pool = _FakePgVectorPool() + service = _make_service(pool) + session = _session( + "app1", + "user1", + "s1", + [ + _event("user", "How do I dispute a billing charge on my invoice?"), + _event("model", "The weather forecast for tomorrow is sunny."), + _event("user", "Who won the basketball game last night?"), + ], + ) + + await service.add_session_to_memory(session) + response = await service.search_memory( + app_name="app1", user_id="user1", query="billing invoice dispute" + ) + + assert response.memories + top = response.memories[0] + assert "billing" in top.content.parts[0].text + assert top.author == "user" + + +@pytest.mark.asyncio +async def test_search_ignores_other_users_and_apps(): + pool = _FakePgVectorPool() + service = _make_service(pool) + await service.add_session_to_memory( + _session("app1", "user1", "s1", [_event("user", "billing invoice")]) + ) + await service.add_session_to_memory( + _session("app1", "user2", "s2", [_event("user", "billing invoice")]) + ) + + response = await service.search_memory( + app_name="app1", user_id="user2", query="billing" + ) + + assert len(response.memories) == 1 + + +@pytest.mark.asyncio +async def test_empty_query_returns_no_memories(): + pool = _FakePgVectorPool() + service = _make_service(pool) + await service.add_session_to_memory( + _session("app1", "user1", "s1", [_event("user", "billing invoice")]) + ) + + response = await service.search_memory( + app_name="app1", user_id="user1", query=" " + ) + + assert response.memories == [] + + +@pytest.mark.asyncio +async def test_events_without_text_are_skipped(): + pool = _FakePgVectorPool() + service = _make_service(pool) + session = _session( + "app1", + "user1", + "s1", + [ + _event("user", "billing invoice"), + Event(author="user", timestamp=1.0), # no content + ], + ) + + await service.add_session_to_memory(session) + + assert len(pool.store) == 1 + + +@pytest.mark.asyncio +async def test_reingesting_a_session_is_idempotent(): + pool = _FakePgVectorPool() + service = _make_service(pool) + session = _session( + "app1", + "user1", + "s1", + [_event("user", "billing invoice"), _event("model", "sunny weather")], + ) + + await service.add_session_to_memory(session) + await service.add_session_to_memory(session) + + # Two text events, ingested twice, must not create duplicate rows. + assert len(pool.store) == 2 + + +@pytest.mark.asyncio +async def test_add_events_to_memory_persists_delta(): + pool = _FakePgVectorPool() + service = _make_service(pool) + + await service.add_events_to_memory( + app_name="app1", + user_id="user1", + events=[_event("user", "billing invoice dispute")], + session_id="s1", + ) + + response = await service.search_memory( + app_name="app1", user_id="user1", query="billing" + ) + assert len(response.memories) == 1 + + +@pytest.mark.asyncio +async def test_add_memory_writes_explicit_entries(): + pool = _FakePgVectorPool() + service = _make_service(pool) + + await service.add_memory( + app_name="app1", + user_id="user1", + memories=[ + MemoryEntry( + content=types.Content( + parts=[types.Part(text="the user prefers dark mode")] + ), + author="user", + custom_metadata={"source": "profile"}, + ) + ], + ) + + response = await service.search_memory( + app_name="app1", user_id="user1", query="dark mode preference" + ) + assert len(response.memories) == 1 + assert response.memories[0].custom_metadata["source"] == "profile" + + +@pytest.mark.asyncio +async def test_distance_threshold_drops_unrelated_memories(): + pool = _FakePgVectorPool() + service = _make_service(pool, distance_threshold=0.5) + await service.add_session_to_memory( + _session( + "app1", + "user1", + "s1", + [ + _event("user", "billing invoice dispute charge"), + _event("model", "basketball score tonight"), + ], + ) + ) + + response = await service.search_memory( + app_name="app1", user_id="user1", query="billing invoice dispute charge" + ) + + # Only the near-identical billing memory is within the distance ceiling. + assert len(response.memories) == 1 + assert "billing" in response.memories[0].content.parts[0].text + + +@pytest.mark.asyncio +async def test_content_survives_serialization_roundtrip(): + pool = _FakePgVectorPool() + service = _make_service(pool) + await service.add_session_to_memory( + _session( + "app1", + "user1", + "s1", + [_event("model", "billing details: multi part reply")], + ) + ) + + response = await service.search_memory( + app_name="app1", user_id="user1", query="billing details" + ) + + memory = response.memories[0] + assert isinstance(memory.content, types.Content) + assert memory.content.parts[0].text == "billing details: multi part reply" + assert memory.timestamp is not None + + +def test_missing_driver_raises_helpful_error(monkeypatch): + monkeypatch.setattr(_pgvector_memory_service, "AsyncConnectionPool", None) + service = PgVectorMemoryService( + PgVectorMemoryServiceConfig(dsn="postgresql://x"), + embedder=_bag_of_words_embedder(), + ) + + with pytest.raises(ImportError, match="google-adk\\[pgvector\\]"): + service._get_pool() + + +def test_dsn_required_when_no_pool(monkeypatch): + # Pretend the driver is installed so the missing dsn is what fails. + monkeypatch.setattr(_pgvector_memory_service, "AsyncConnectionPool", object) + service = PgVectorMemoryService( + PgVectorMemoryServiceConfig(dsn=None), + embedder=_bag_of_words_embedder(), + ) + + with pytest.raises(ValueError, match="dsn is required"): + service._get_pool() + + +@pytest.mark.asyncio +async def test_owned_pool_is_opened_once_and_closed(monkeypatch): + fake = _FakePgVectorPool() + monkeypatch.setattr( + _pgvector_memory_service, + "AsyncConnectionPool", + lambda dsn, open: fake, + ) + # No pool injected: the service creates and owns it, so it must open it. + service = PgVectorMemoryService( + PgVectorMemoryServiceConfig( + dsn="postgresql://x", embedding_dimension=_DIM + ), + embedder=_bag_of_words_embedder(), + ) + + await service.add_session_to_memory( + _session("app1", "user1", "s1", [_event("user", "billing invoice")]) + ) + await service.search_memory(app_name="app1", user_id="user1", query="billing") + + assert fake.open_count == 1 # opened once, not per operation + await service.close() + assert fake.closed is True + + +@pytest.mark.asyncio +async def test_injected_pool_is_not_opened_or_closed(): + # A caller-supplied pool is managed by the caller, so the service must not + # open or close it. + pool = _FakePgVectorPool() + service = _make_service(pool) + + await service.add_session_to_memory( + _session("app1", "user1", "s1", [_event("user", "billing invoice")]) + ) + await service.close() + + assert pool.open_count == 0 + assert pool.closed is False