Skip to content

Commit a04727e

Browse files
committed
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
1 parent fdf1b7f commit a04727e

7 files changed

Lines changed: 1099 additions & 0 deletions

File tree

‎pyproject.toml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,8 +122,11 @@ optional-dependencies.all = [
122122
"opentelemetry-instrumentation-httpx>=0.54b0,<1",
123123
"opentelemetry-resourcedetector-gcp>=1.9.0a0,<2",
124124
"pandas>=2.2.3",
125+
"pgvector>=0.3",
125126
"pillow>=10.3",
126127
"protobuf>=6",
128+
"psycopg-pool>=3.2",
129+
"psycopg[binary]>=3.1",
127130
"pyarrow>=14",
128131
"pymongo>=4.9,<5",
129132
"pypika>=0.50",
@@ -275,6 +278,11 @@ optional-dependencies.otel-gcp = [
275278
"opentelemetry-instrumentation-grpc>=0.43b0,<1",
276279
"opentelemetry-instrumentation-httpx>=0.54b0,<1",
277280
]
281+
optional-dependencies.pgvector = [
282+
"pgvector>=0.3", # registers the pgvector type on the connection.
283+
"psycopg[binary]>=3.1", # async driver used by PgVectorMemoryService.
284+
"psycopg-pool>=3.2", # AsyncConnectionPool.
285+
]
278286
optional-dependencies.redis = [
279287
"redis>=4.2", # 4.2 is where redis.asyncio landed, the only API used.
280288
]
Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
# PostgreSQL / pgvector Memory Integration for ADK
2+
3+
This integration provides a self-hosted, semantic memory service for the Google
4+
Agent Development Kit (ADK), backed by PostgreSQL with the
5+
[pgvector](https://github.com/pgvector/pgvector) extension.
6+
7+
It complements the built-in memory services:
8+
9+
- `InMemoryMemoryService` is keyword-only and loses data on restart (prototyping
10+
only).
11+
- `VertexAiMemoryBankService` and `VertexAiRagMemoryService` provide semantic
12+
memory but require Google Cloud.
13+
14+
`PgVectorMemoryService` gives you semantic (vector) memory that persists across
15+
restarts and runs entirely on a database you control, which suits on-premise,
16+
air-gapped, and data-residency-constrained deployments.
17+
18+
## Features
19+
20+
- **Semantic search:** Ranks memories by cosine similarity over embeddings using
21+
a pgvector HNSW index, instead of keyword overlap.
22+
- **Self-hosted:** Vectors live in your own PostgreSQL database. No managed Cloud
23+
service is required for storage.
24+
- **Idempotent ingestion:** Re-ingesting a session updates existing rows in
25+
place rather than creating duplicates.
26+
- **Scoped memory:** Memories are isolated per `(app_name, user_id)`.
27+
- **Pluggable embeddings:** Embeds through the `google-genai` client by default,
28+
or any provider via an injected `embedder` callable.
29+
30+
## Installation / Dependencies
31+
32+
Install the optional `pgvector` extra alongside ADK:
33+
34+
```bash
35+
pip install "google-adk[pgvector]"
36+
```
37+
38+
You also need a PostgreSQL database with the `vector` extension available. The
39+
service creates the extension, table, and indexes on first use.
40+
41+
## Quick Start
42+
43+
```python
44+
from google.adk.integrations.pgvector import PgVectorMemoryService
45+
from google.adk.integrations.pgvector import PgVectorMemoryServiceConfig
46+
from google.adk.runners import Runner
47+
48+
# 1. Configure the memory service.
49+
memory_service = PgVectorMemoryService(
50+
PgVectorMemoryServiceConfig(
51+
dsn="postgresql://user:password@localhost:5432/adk",
52+
embedding_model="gemini-embedding-001",
53+
embedding_dimension=768,
54+
)
55+
)
56+
57+
# 2. Wire it into your Runner.
58+
runner = Runner(
59+
app_name="my_app",
60+
agent=agent,
61+
memory_service=memory_service,
62+
)
63+
64+
# Ingesting and searching:
65+
await memory_service.add_session_to_memory(session)
66+
response = await memory_service.search_memory(
67+
app_name="my_app", user_id="user1", query="what did we decide about billing?"
68+
)
69+
```
70+
71+
## Configuration
72+
73+
`PgVectorMemoryServiceConfig` supports the following options:
74+
75+
| Option | Default | Description |
76+
| --- | --- | --- |
77+
| `dsn` | `None` | PostgreSQL connection string. Required unless a `connection_pool` is injected. |
78+
| `table_name` | `adk_memory_entries` | Table that stores memory entries. |
79+
| `embedding_model` | `gemini-embedding-001` | `google-genai` embedding model. |
80+
| `embedding_dimension` | `768` | Stored vector dimension; must match the model output and stays fixed for the table. |
81+
| `top_k` | `10` | Maximum memories returned per search. |
82+
| `hnsw_m` | `16` | pgvector HNSW `m` parameter. |
83+
| `hnsw_ef_construction` | `64` | pgvector HNSW `ef_construction` parameter. |
84+
| `distance_threshold` | `None` | Optional cosine-distance ceiling; farther memories are dropped. |
85+
86+
## Advanced usage
87+
88+
Inject a pre-configured pool or a non-Google embedding provider:
89+
90+
```python
91+
from psycopg_pool import AsyncConnectionPool
92+
93+
pool = AsyncConnectionPool("postgresql://user:password@localhost:5432/adk")
94+
95+
async def my_embedder(texts):
96+
# Return one vector per text from any provider.
97+
...
98+
99+
memory_service = PgVectorMemoryService(
100+
PgVectorMemoryServiceConfig(embedding_dimension=1024),
101+
connection_pool=pool,
102+
embedder=my_embedder,
103+
)
104+
```
Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""PostgreSQL/pgvector integrations for ADK."""
16+
17+
from __future__ import annotations
18+
19+
from ._config import PgVectorMemoryServiceConfig
20+
from ._pgvector_memory_service import PgVectorMemoryService
21+
22+
__all__ = [
23+
"PgVectorMemoryService",
24+
"PgVectorMemoryServiceConfig",
25+
]
Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Configuration for the PostgreSQL/pgvector integrations."""
16+
17+
from __future__ import annotations
18+
19+
from typing import Optional
20+
21+
from pydantic import BaseModel
22+
from pydantic import Field
23+
24+
25+
class PgVectorMemoryServiceConfig(BaseModel):
26+
"""Configuration for PgVectorMemoryService.
27+
28+
The embedding model is called through the ``google-genai`` client, but the
29+
vectors themselves live in a PostgreSQL database you control, so the memory
30+
store does not depend on any managed Google Cloud service.
31+
"""
32+
33+
dsn: Optional[str] = Field(
34+
default=None,
35+
description=(
36+
"PostgreSQL connection string, e.g."
37+
" postgresql://user:password@host:5432/dbname. Required unless a"
38+
" pre-configured connection pool is passed to the service."
39+
),
40+
)
41+
table_name: str = Field(
42+
default="adk_memory_entries",
43+
description="Name of the table that stores memory entries.",
44+
)
45+
embedding_model: str = Field(
46+
default="gemini-embedding-001",
47+
description=(
48+
"google-genai embedding model used to embed events and queries."
49+
),
50+
)
51+
embedding_dimension: int = Field(
52+
default=768,
53+
description=(
54+
"Dimension of the stored embedding vectors. Must match the output"
55+
" dimension of embedding_model and stays fixed for the life of the"
56+
" table."
57+
),
58+
)
59+
top_k: int = Field(
60+
default=10,
61+
description="Maximum number of memories returned by a search.",
62+
)
63+
hnsw_m: int = Field(
64+
default=16,
65+
description="pgvector HNSW index parameter m (graph degree).",
66+
)
67+
hnsw_ef_construction: int = Field(
68+
default=64,
69+
description=(
70+
"pgvector HNSW index parameter ef_construction (build effort)."
71+
),
72+
)
73+
distance_threshold: Optional[float] = Field(
74+
default=None,
75+
description=(
76+
"Optional cosine-distance ceiling in [0, 2]. When set, memories whose"
77+
" distance to the query is greater than this are dropped from the"
78+
" results. When None, the closest top_k memories are always returned."
79+
),
80+
)

0 commit comments

Comments
 (0)