From f2bb738ba9a0940cd855f46a1918d716d669df8f Mon Sep 17 00:00:00 2001 From: goodnight Date: Tue, 3 Nov 2026 14:18:36 +0100 Subject: [PATCH] feat(memory): add QdrantMemoryService --- pyproject.toml | 3 + src/google/adk_community/memory/__init__.py | 5 +- .../memory/qdrant_memory_service.py | 362 ++++++++++++++++++ .../memory/test_qdrant_memory_service.py | 343 +++++++++++++++++ 4 files changed, 712 insertions(+), 1 deletion(-) create mode 100644 src/google/adk_community/memory/qdrant_memory_service.py create mode 100644 tests/unittests/memory/test_qdrant_memory_service.py diff --git a/pyproject.toml b/pyproject.toml index a03bdcab..72410401 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,9 @@ sdc-agents = [ "sdc-agents>=4.3.3; python_version >= '3.11'", ] spraay = ["web3>=6.0.0"] +qdrant = [ + "qdrant-client>=1.12.0", # For QdrantMemoryService +] hitl = [ "aiosqlite>=0.20.0", "fastapi>=0.110.0", diff --git a/src/google/adk_community/memory/__init__.py b/src/google/adk_community/memory/__init__.py index 1f3442c0..f47224cd 100644 --- a/src/google/adk_community/memory/__init__.py +++ b/src/google/adk_community/memory/__init__.py @@ -16,9 +16,12 @@ from .open_memory_service import OpenMemoryService from .open_memory_service import OpenMemoryServiceConfig +from .qdrant_memory_service import QdrantMemoryService +from .qdrant_memory_service import QdrantMemoryServiceConfig __all__ = [ "OpenMemoryService", "OpenMemoryServiceConfig", + "QdrantMemoryService", + "QdrantMemoryServiceConfig", ] - diff --git a/src/google/adk_community/memory/qdrant_memory_service.py b/src/google/adk_community/memory/qdrant_memory_service.py new file mode 100644 index 00000000..078041cd --- /dev/null +++ b/src/google/adk_community/memory/qdrant_memory_service.py @@ -0,0 +1,362 @@ +# 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. + +"""Memory service implementation backed by Qdrant.""" + +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 logging +from typing import Any +from typing import Literal +from typing import Optional +from typing import TYPE_CHECKING +import uuid + +from google.adk.memory import _utils +from google.adk.memory.base_memory_service import BaseMemoryService +from google.adk.memory.base_memory_service import SearchMemoryResponse +from google.adk.memory.memory_entry import MemoryEntry +from google.genai import types +from pydantic import BaseModel +from pydantic import Field +from typing_extensions import override + +if TYPE_CHECKING: + from google.adk.events.event import Event + from google.adk.sessions.session import Session + from qdrant_client import AsyncQdrantClient + +logger = logging.getLogger('google_adk_community.' + __name__) + +EmbeddingFunction = Callable[[list[str]], Awaitable[list[list[float]]]] +"""Async callable mapping a batch of texts to one embedding vector each.""" + +_ID_NAMESPACE = uuid.UUID('6f1d5c2e-8f4b-4a57-9d3e-2b7c1a0e9f41') + + +class QdrantMemoryServiceConfig(BaseModel): + """Configuration for QdrantMemoryService. + + Attributes: + collection_name: Name of the Qdrant collection that stores memories. + vector_size: Embedding dimension. If None, it is inferred from the first + batch of embeddings when the collection is auto-created. + distance: Distance metric used when the collection is auto-created. + create_collection: Create the collection and payload indexes on first + write if they do not exist. Set to False to manage them yourself. + search_top_k: Maximum number of memories returned per search. + score_threshold: Optional minimum similarity score for search results. + """ + + collection_name: str = 'adk_memory' + vector_size: Optional[int] = Field(default=None, ge=1) + distance: Literal['Cosine', 'Dot', 'Euclid', 'Manhattan'] = 'Cosine' + create_collection: bool = True + search_top_k: int = Field(default=10, ge=1) + score_threshold: Optional[float] = None + + +class QdrantMemoryService(BaseMemoryService): + """A memory service that stores and searches memories in Qdrant. + + Each text-bearing event becomes one Qdrant point. Points carry `app_name` + and `user_id` in their payload, and every read, write and delete is scoped + by a payload filter on both fields. + + Embeddings are computed by a caller-provided async function, so the service + does not depend on any embedding provider. + """ + + def __init__( + self, + client: AsyncQdrantClient, + embedding_function: EmbeddingFunction, + config: Optional[QdrantMemoryServiceConfig] = None, + ): + """Initializes the Qdrant memory service. + + Args: + client: An `AsyncQdrantClient`. The caller owns its lifecycle. + embedding_function: Async callable that takes a list of texts and + returns a list of embedding vectors, one per text. + config: Optional QdrantMemoryServiceConfig. Defaults are used if None. + """ + try: + from qdrant_client import models + except ImportError as exc: + raise ImportError( + 'qdrant-client is required to use QdrantMemoryService. ' + 'Install it with: pip install google-adk-community[qdrant]' + ) from exc + if client is None: + raise ValueError('client is required.') + if embedding_function is None: + raise ValueError('embedding_function is required.') + self._models = models + self._client = client + self._embedding_function = embedding_function + self._config = config or QdrantMemoryServiceConfig() + self._collection_ready = False + self._collection_lock = asyncio.Lock() + + @override + async def add_session_to_memory(self, session: Session) -> None: + await self._upsert_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._upsert_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: + records = [] + for memory in memories: + text = _extract_text(memory.content) + if not text: + continue + key = memory.id or str(uuid.uuid4()) + metadata = {**memory.custom_metadata, **(custom_metadata or {})} + records.append(( + _point_id(app_name, user_id, key), + text, + { + 'app_name': app_name, + 'user_id': user_id, + 'text': text, + 'content': ( + _visible_content(memory.content).model_dump( + mode='json', exclude_none=True + ) + ), + 'author': memory.author, + 'timestamp': memory.timestamp, + 'custom_metadata': metadata, + }, + )) + await self._upsert(records) + + @override + async def search_memory( + self, *, app_name: str, user_id: str, query: str + ) -> SearchMemoryResponse: + if not await self._collection_exists(): + return SearchMemoryResponse() + [query_vector] = await self._embedding_function([query]) + response = await self._client.query_points( + collection_name=self._config.collection_name, + query=query_vector, + query_filter=self._scope_filter(app_name, user_id), + limit=self._config.search_top_k, + score_threshold=self._config.score_threshold, + with_payload=True, + ) + memories = [] + for point in response.points: + payload = point.payload or {} + memories.append( + MemoryEntry( + id=str(point.id), + content=types.Content.model_validate(payload['content']), + author=payload.get('author'), + timestamp=payload.get('timestamp'), + custom_metadata=payload.get('custom_metadata') or {}, + ) + ) + return SearchMemoryResponse(memories=memories) + + @override + async def delete_memory( + self, + *, + app_name: str, + user_id: str, + memory_id: str | None = None, + ) -> None: + if not memory_id or not await self._collection_exists(): + return + try: + uuid.UUID(memory_id) + except ValueError: + # Ids returned by search_memory are always UUIDs; anything else cannot + # match a point in this collection. + return + models = self._models + scoped = self._scope_filter(app_name, user_id) + await self._client.delete( + collection_name=self._config.collection_name, + points_selector=models.FilterSelector( + filter=models.Filter( + must=[models.HasIdCondition(has_id=[memory_id]), *scoped.must] + ) + ), + ) + + async def _upsert_events( + self, + *, + app_name: str, + user_id: str, + events: Sequence[Event], + session_id: str | None, + custom_metadata: Mapping[str, object] | None = None, + ) -> None: + records = [] + for event in events: + text = _extract_text(event.content) + if not text: + continue + records.append(( + _point_id(app_name, user_id, event.id), + text, + { + 'app_name': app_name, + 'user_id': user_id, + 'session_id': session_id, + 'event_id': event.id, + 'text': text, + 'content': ( + _visible_content(event.content).model_dump( + mode='json', exclude_none=True + ) + ), + 'author': event.author, + 'timestamp': _utils.format_timestamp(event.timestamp), + 'custom_metadata': dict(custom_metadata or {}), + }, + )) + await self._upsert(records) + + async def _upsert(self, records: list[tuple[str, str, dict[str, Any]]]): + """Embeds and upserts (point_id, text, payload) records.""" + if not records: + return + vectors = await self._embedding_function([text for _, text, _ in records]) + if len(vectors) != len(records): + raise ValueError( + f'embedding_function returned {len(vectors)} vectors for' + f' {len(records)} texts.' + ) + await self._ensure_collection(len(vectors[0])) + await self._client.upsert( + collection_name=self._config.collection_name, + points=[ + self._models.PointStruct( + id=point_id, vector=vector, payload=payload + ) + for (point_id, _, payload), vector in zip(records, vectors) + ], + ) + logger.debug( + 'Upserted %d memories into %s', + len(records), + self._config.collection_name, + ) + + async def _collection_exists(self) -> bool: + if self._collection_ready: + return True + return await self._client.collection_exists(self._config.collection_name) + + async def _ensure_collection(self, vector_size: int) -> None: + if self._collection_ready or not self._config.create_collection: + return + async with self._collection_lock: + if self._collection_ready: + return + name = self._config.collection_name + if not await self._client.collection_exists(name): + models = self._models + await self._client.create_collection( + collection_name=name, + vectors_config=models.VectorParams( + size=self._config.vector_size or vector_size, + distance=models.Distance(self._config.distance), + ), + ) + for field in ('app_name', 'user_id'): + await self._client.create_payload_index( + collection_name=name, + field_name=field, + field_schema=models.PayloadSchemaType.KEYWORD, + ) + self._collection_ready = True + + def _scope_filter(self, app_name: str, user_id: str): + models = self._models + return models.Filter( + must=[ + models.FieldCondition( + key='app_name', match=models.MatchValue(value=app_name) + ), + models.FieldCondition( + key='user_id', match=models.MatchValue(value=user_id) + ), + ] + ) + + +def _point_id(app_name: str, user_id: str, key: str) -> str: + """Deterministic point id so re-ingesting a session overwrites, not dupes.""" + return str(uuid.uuid5(_ID_NAMESPACE, f'{app_name}\x00{user_id}\x00{key}')) + + +def _visible_content(content: types.Content) -> types.Content: + """Returns a copy of content with only non-thought text parts.""" + return types.Content( + role=content.role, + parts=[ + types.Part(text=part.text) + for part in content.parts or [] + if part.text and not part.thought + ], + ) + + +def _extract_text(content: Optional[types.Content]) -> str: + if not content or not content.parts: + return '' + return ' '.join( + part.text for part in content.parts if part.text and not part.thought + ) diff --git a/tests/unittests/memory/test_qdrant_memory_service.py b/tests/unittests/memory/test_qdrant_memory_service.py new file mode 100644 index 00000000..36e98d90 --- /dev/null +++ b/tests/unittests/memory/test_qdrant_memory_service.py @@ -0,0 +1,343 @@ +# 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. + +import hashlib +from types import SimpleNamespace +from unittest.mock import AsyncMock +from unittest.mock import MagicMock +from unittest.mock import patch + +from google.adk.events.event import Event +from google.adk.memory.memory_entry import MemoryEntry +from google.adk.sessions.session import Session +from google.genai import types +import pytest + +pytest.importorskip('qdrant_client') + +from qdrant_client import AsyncQdrantClient +from qdrant_client import models + +from google.adk_community.memory.qdrant_memory_service import QdrantMemoryService +from google.adk_community.memory.qdrant_memory_service import QdrantMemoryServiceConfig + +APP = 'test-app' +USER = 'test-user' +DIM = 8 + + +def _event(event_id, author, text=None, thought=None, timestamp=12345.0): + parts = [] + if thought: + parts.append(types.Part(text=thought, thought=True)) + if text: + parts.append(types.Part(text=text)) + return Event( + id=event_id, + invocation_id=f'inv-{event_id}', + author=author, + timestamp=timestamp, + content=types.Content(role='user', parts=parts) if parts else None, + ) + + +SESSION = Session( + app_name=APP, + user_id=USER, + id='session-1', + last_update_time=1000, + events=[ + _event('e1', 'user', 'I like Python.', thought='secret reasoning'), + _event('e2', 'model', 'Python is great.'), + _event('e3', 'user'), # no content, skipped + Event( + id='e4', + invocation_id='inv-e4', + author='model', + content=types.Content( + parts=[ + types.Part(function_call=types.FunctionCall(name='tool')) + ] + ), + ), # no text, skipped + ], +) + + +async def fake_embed(texts): + """Deterministic bag-of-words embedding for tests.""" + vectors = [] + for text in texts: + vec = [0.0] * DIM + for word in text.lower().replace('.', ' ').split(): + vec[int(hashlib.md5(word.encode()).hexdigest(), 16) % DIM] += 1.0 + vectors.append(vec) + return vectors + + +@pytest.fixture +def mock_client(): + client = MagicMock(spec=AsyncQdrantClient) + client.collection_exists = AsyncMock(return_value=False) + client.create_collection = AsyncMock() + client.create_payload_index = AsyncMock() + client.upsert = AsyncMock() + client.delete = AsyncMock() + client.query_points = AsyncMock(return_value=SimpleNamespace(points=[])) + return client + + +@pytest.fixture +def service(mock_client): + return QdrantMemoryService( + client=mock_client, embedding_function=AsyncMock(side_effect=fake_embed) + ) + + +def test_requires_client_and_embedding_function(mock_client): + with pytest.raises(ValueError): + QdrantMemoryService(client=None, embedding_function=fake_embed) + with pytest.raises(ValueError): + QdrantMemoryService(client=mock_client, embedding_function=None) + + +def test_missing_qdrant_client_raises_import_error(mock_client): + with patch.dict('sys.modules', {'qdrant_client': None}): + with pytest.raises(ImportError, match=r'google-adk-community\[qdrant\]'): + QdrantMemoryService(client=mock_client, embedding_function=fake_embed) + + +async def test_add_session_creates_collection_and_upserts(service, mock_client): + await service.add_session_to_memory(SESSION) + + mock_client.create_collection.assert_awaited_once() + kwargs = mock_client.create_collection.call_args.kwargs + assert kwargs['collection_name'] == 'adk_memory' + assert kwargs['vectors_config'].size == DIM + assert kwargs['vectors_config'].distance == models.Distance.COSINE + indexed = { + c.kwargs['field_name'] + for c in mock_client.create_payload_index.call_args_list + } + assert indexed == {'app_name', 'user_id'} + + points = mock_client.upsert.call_args.kwargs['points'] + assert len(points) == 2 + payload = points[0].payload + assert payload['app_name'] == APP + assert payload['user_id'] == USER + assert payload['session_id'] == 'session-1' + assert payload['event_id'] == 'e1' + assert payload['text'] == 'I like Python.' + assert payload['content']['parts'] == [{'text': 'I like Python.'}] + assert 'secret reasoning' not in str(payload) + + +async def test_collection_created_once(service, mock_client): + await service.add_session_to_memory(SESSION) + await service.add_session_to_memory(SESSION) + assert mock_client.create_collection.await_count == 1 + assert mock_client.collection_exists.await_count == 1 + + +async def test_point_ids_are_deterministic_and_scoped(service, mock_client): + await service.add_session_to_memory(SESSION) + first = [p.id for p in mock_client.upsert.call_args.kwargs['points']] + await service.add_session_to_memory(SESSION) + second = [p.id for p in mock_client.upsert.call_args.kwargs['points']] + assert first == second + + other_user = SESSION.model_copy(update={'user_id': 'other'}) + await service.add_session_to_memory(other_user) + third = [p.id for p in mock_client.upsert.call_args.kwargs['points']] + assert set(first).isdisjoint(third) + + +async def test_no_create_when_disabled(mock_client): + service = QdrantMemoryService( + client=mock_client, + embedding_function=fake_embed, + config=QdrantMemoryServiceConfig(create_collection=False), + ) + await service.add_session_to_memory(SESSION) + mock_client.create_collection.assert_not_awaited() + mock_client.upsert.assert_awaited_once() + + +async def test_configured_vector_size_and_distance(mock_client): + service = QdrantMemoryService( + client=mock_client, + embedding_function=fake_embed, + config=QdrantMemoryServiceConfig(vector_size=2 * DIM, distance='Dot'), + ) + await service.add_session_to_memory(SESSION) + config = mock_client.create_collection.call_args.kwargs['vectors_config'] + assert config.size == 2 * DIM # configured size wins over inferred + assert config.distance == models.Distance.DOT + + +async def test_empty_session_does_not_call_qdrant(service, mock_client): + session = SESSION.model_copy(update={'events': [_event('x', 'user')]}) + await service.add_session_to_memory(session) + mock_client.upsert.assert_not_awaited() + service._embedding_function.assert_not_awaited() + + +async def test_embedding_count_mismatch_raises(mock_client): + async def bad_embed(texts): + return [[0.1] * DIM] + + service = QdrantMemoryService( + client=mock_client, embedding_function=bad_embed + ) + with pytest.raises(ValueError, match='returned 1 vectors for 2 texts'): + await service.add_session_to_memory(SESSION) + + +async def test_add_events_to_memory(service, mock_client): + await service.add_events_to_memory( + app_name=APP, + user_id=USER, + events=[_event('e9', 'user', 'New fact.')], + custom_metadata={'source': 'delta'}, + ) + [point] = mock_client.upsert.call_args.kwargs['points'] + assert point.payload['session_id'] is None + assert point.payload['custom_metadata'] == {'source': 'delta'} + + +async def test_add_memory(service, mock_client): + await service.add_memory( + app_name=APP, + user_id=USER, + memories=[ + MemoryEntry( + id='fact-1', + content=types.Content(parts=[types.Part(text='Likes tea.')]), + author='user', + custom_metadata={'kind': 'preference'}, + ), + MemoryEntry(content=types.Content(parts=[])), # skipped + ], + custom_metadata={'source': 'explicit'}, + ) + [point] = mock_client.upsert.call_args.kwargs['points'] + assert point.payload['text'] == 'Likes tea.' + assert point.payload['custom_metadata'] == { + 'kind': 'preference', + 'source': 'explicit', + } + + +async def test_search_scopes_by_app_and_user(service, mock_client): + mock_client.collection_exists.return_value = True + mock_client.query_points.return_value = SimpleNamespace( + points=[ + models.ScoredPoint( + id='5c56c793-69f3-4fbf-87e6-c4bf54c28c26', + version=1, + score=0.9, + payload={ + 'content': {'parts': [{'text': 'I like Python.'}]}, + 'author': 'user', + 'timestamp': '2025-01-01T00:00:00', + 'custom_metadata': {'k': 'v'}, + }, + ) + ] + ) + + result = await service.search_memory( + app_name=APP, user_id=USER, query='python' + ) + + kwargs = mock_client.query_points.call_args.kwargs + assert kwargs['limit'] == 10 + conditions = {c.key: c.match.value for c in kwargs['query_filter'].must} + assert conditions == {'app_name': APP, 'user_id': USER} + [memory] = result.memories + assert memory.id == '5c56c793-69f3-4fbf-87e6-c4bf54c28c26' + assert memory.content.parts[0].text == 'I like Python.' + assert memory.author == 'user' + assert memory.timestamp == '2025-01-01T00:00:00' + assert memory.custom_metadata == {'k': 'v'} + + +async def test_search_without_collection_returns_empty(service, mock_client): + result = await service.search_memory(app_name=APP, user_id=USER, query='x') + assert result.memories == [] + mock_client.query_points.assert_not_awaited() + + +async def test_delete_memory_is_scoped(service, mock_client): + mock_client.collection_exists.return_value = True + memory_id = '5c56c793-69f3-4fbf-87e6-c4bf54c28c26' + await service.delete_memory(app_name=APP, user_id=USER, memory_id=memory_id) + selector = mock_client.delete.call_args.kwargs['points_selector'] + has_id, *fields = selector.filter.must + assert has_id.has_id == [memory_id] + assert {c.key: c.match.value for c in fields} == { + 'app_name': APP, + 'user_id': USER, + } + + +@pytest.mark.parametrize('memory_id', [None, '', 'not-a-uuid']) +async def test_delete_memory_ignores_invalid_ids( + service, mock_client, memory_id +): + mock_client.collection_exists.return_value = True + await service.delete_memory(app_name=APP, user_id=USER, memory_id=memory_id) + mock_client.delete.assert_not_awaited() + + +@pytest.mark.filterwarnings('ignore:Payload indexes have no effect') +async def test_end_to_end_with_local_in_memory_qdrant(): + """Runs against qdrant-client's local in-memory mode (no server).""" + client = AsyncQdrantClient(':memory:') + service = QdrantMemoryService(client=client, embedding_function=fake_embed) + try: + await service.add_session_to_memory(SESSION) + await service.add_session_to_memory(SESSION) # idempotent + await service.add_session_to_memory( + SESSION.model_copy(update={'user_id': 'other'}) + ) + assert (await client.count('adk_memory')).count == 4 + + result = await service.search_memory( + app_name=APP, user_id=USER, query='I like Python.' + ) + assert len(result.memories) == 2 + assert result.memories[0].content.parts[0].text == 'I like Python.' + assert result.memories[0].author == 'user' + + await service.delete_memory( + app_name=APP, user_id='other', memory_id=result.memories[0].id + ) + assert (await client.count('adk_memory')).count == 4 # wrong scope + + await service.delete_memory( + app_name=APP, user_id=USER, memory_id=result.memories[0].id + ) + result = await service.search_memory( + app_name=APP, user_id=USER, query='I like Python.' + ) + assert [m.content.parts[0].text for m in result.memories] == [ + 'Python is great.' + ] + assert ( + await service.search_memory(app_name='nope', user_id=USER, query='x') + ).memories == [] + finally: + await client.close()