diff --git a/pyproject.toml b/pyproject.toml index ef23e4dc38..0ca616be4c 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 0000000000..ff708f9e6b --- /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 0000000000..2c00ae90bd --- /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 0000000000..be5051dadf --- /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 0000000000..85e2e3da48 --- /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 0000000000..58d482ea38 --- /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 0000000000..7679777f0a --- /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