diff --git a/src/google/adk/memory/vertex_ai_rag_memory_service.py b/src/google/adk/memory/vertex_ai_rag_memory_service.py index 256abb56..7f7e2c2b 100644 --- a/src/google/adk/memory/vertex_ai_rag_memory_service.py +++ b/src/google/adk/memory/vertex_ai_rag_memory_service.py @@ -20,6 +20,7 @@ import base64 import binascii from collections import OrderedDict import json +import logging import os import tempfile from typing import Optional @@ -34,12 +35,22 @@ from .base_memory_service import SearchMemoryResponse from .memory_entry import MemoryEntry if TYPE_CHECKING: + import agentplatform + from ..events.event import Event from ..sessions.session import Session +logger = logging.getLogger("google_adk." + __name__) + _SOURCE_DISPLAY_NAME_PREFIX = "adk-memory-v1." +_RAG_FILE_PAGE_SIZE = 100 + +# Scoping walks the corpus file list, so it is capped to keep the cost of a +# search independent of how large a shared corpus grows. +_MAX_RAG_FILE_PAGES = 10 + def _encode_source_display_name_part(value: str) -> str: return ( @@ -90,6 +101,69 @@ def _parse_source_display_name( return parts[0], parts[1], parts[2] +async def _scoped_rag_resources( + client: agentplatform.AsyncClient, + rag_resources: list[types.VertexRagStoreRagResource], + app_name: str, + user_id: str, +) -> list[types.VertexRagStoreRagResource] | None: + """Returns resources naming only the files owned by one app and user. + + Returns None when a corpus cannot be listed within the page budget, in which + case the caller retrieves without narrowing the resources. + """ + from agentplatform import types as agentplatform_types + + scoped_resources: list[types.VertexRagStoreRagResource] = [] + for rag_resource in rag_resources: + rag_corpus = rag_resource.rag_corpus + if not rag_corpus: + return None + + rag_file_ids: list[str] = [] + page_token: str | None = None + for _ in range(_MAX_RAG_FILE_PAGES): + response = await client.rag.list_files( + name=rag_corpus, + config=agentplatform_types.ListRagFilesConfig( + page_size=_RAG_FILE_PAGE_SIZE, page_token=page_token + ), + ) + for rag_file in response.rag_files or []: + session_info = _parse_source_display_name(rag_file.display_name or "") + if ( + not session_info + or session_info[0] != app_name + or session_info[1] != user_id + or not rag_file.name + ): + continue + # rag_file_ids takes the bare file id, not the full resource name. + rag_file_ids.append(rag_file.name.rsplit("/", 1)[-1]) + page_token = response.next_page_token + if not page_token: + break + + if page_token: + # Scoping to the files seen so far would hide the caller's own memories, + # so an incomplete listing is abandoned instead. + logger.warning( + "Listing %s did not finish within %d pages, so retrieval is not" + " scoped to the requesting app and user.", + rag_corpus, + _MAX_RAG_FILE_PAGES, + ) + return None + if rag_file_ids: + scoped_resources.append( + types.VertexRagStoreRagResource( + rag_corpus=rag_corpus, rag_file_ids=rag_file_ids + ) + ) + + return scoped_resources + + class VertexAiRagMemoryService(BaseMemoryService): """A memory service that uses Agent Platform RAG for storage and retrieval.""" @@ -223,12 +297,38 @@ class VertexAiRagMemoryService(BaseMemoryService): from ..events.event import Event + rag_resources = self._vertex_rag_store.rag_resources + if not rag_resources: + raise ValueError("Rag resources must be set.") + client = agentplatform.Client( project=self._project, location=self._location ).aio try: + try: + scoped_resources = await _scoped_rag_resources( + client, rag_resources, app_name, user_id + ) + except Exception: # pylint: disable=broad-except + # Narrowing the resources only improves ranking and transfer; the + # response filter below keeps an unnarrowed retrieval correct. + logger.warning( + "Listing the corpus failed, so retrieval is not scoped to the" + " requesting app and user.", + exc_info=True, + ) + scoped_resources = None + + vertex_rag_store = self._vertex_rag_store + if scoped_resources is not None: + if not scoped_resources: + return SearchMemoryResponse() + vertex_rag_store = self._vertex_rag_store.model_copy( + update={"rag_resources": scoped_resources} + ) + response = await client.rag.retrieve_contexts( - vertex_rag_store=self._vertex_rag_store, + vertex_rag_store=vertex_rag_store, query=agentplatform_types.RagQuery( text=query, similarity_top_k=self._similarity_top_k, @@ -241,8 +341,7 @@ class VertexAiRagMemoryService(BaseMemoryService): memory_results = [] session_events_map: OrderedDict[str, list[list[Event]]] = OrderedDict() for context in response.contexts.contexts: - # filter out context that is not related - # TODO: Add server side filtering by app_name and user_id. + # Still required: retrieval is unscoped whenever listing did not finish. source_display_name = getattr(context, "source_display_name", "") if not isinstance(source_display_name, str): continue diff --git a/tests/unittests/memory/test_vertex_ai_rag_memory_service.py b/tests/unittests/memory/test_vertex_ai_rag_memory_service.py index 70315858..bbe19c92 100644 --- a/tests/unittests/memory/test_vertex_ai_rag_memory_service.py +++ b/tests/unittests/memory/test_vertex_ai_rag_memory_service.py @@ -14,12 +14,14 @@ import asyncio import json +import logging import os import tempfile from types import SimpleNamespace from google.adk.events.event import Event from google.adk.memory.vertex_ai_rag_memory_service import _build_source_display_name +from google.adk.memory.vertex_ai_rag_memory_service import _MAX_RAG_FILE_PAGES from google.adk.memory.vertex_ai_rag_memory_service import _SOURCE_DISPLAY_NAME_PREFIX from google.adk.memory.vertex_ai_rag_memory_service import VertexAiRagMemoryService from google.adk.sessions.session import Session @@ -35,6 +37,42 @@ def _rag_context(source_display_name: str, text: str) -> SimpleNamespace: ) +def _rag_file(rag_file_id: str, source_display_name: str) -> SimpleNamespace: + # A listing entry reports the full resource name, not the bare file id. + return SimpleNamespace( + name=( + "projects/test-project/locations/us-central1/ragCorpora/1/ragFiles/" + + rag_file_id + ), + display_name=source_display_name, + ) + + +def _memory_texts(response) -> list[str]: + return [memory.content.parts[0].text for memory in response.memories] + + +def _retrieved_store(fake_client) -> types.VertexRagStore: + return fake_client.rag.retrieve_contexts.call_args.kwargs["vertex_rag_store"] + + +def _async_client(mocker): + """A client on the SDK async surface, where every RAG call is awaited.""" + fake_client = mocker.Mock() + fake_client.aclose = mocker.AsyncMock() + fake_client.rag.list_files = mocker.AsyncMock() + fake_client.rag.retrieve_contexts = mocker.AsyncMock() + fake_client.rag.upload_file = mocker.AsyncMock() + return fake_client + + +def _unlistable_client(mocker): + """A client for the tests that exercise the unscoped retrieval path.""" + fake_client = _async_client(mocker) + fake_client.rag.list_files.side_effect = PermissionError("cannot list files") + return fake_client + + def _session() -> Session: return Session( app_name="demo.app", @@ -96,10 +134,9 @@ async def test_search_memory_forwards_similarity_top_k( rag_corpus="unused", similarity_top_k=configured_top_k, ) - fake_client = mocker.Mock() - fake_client.aclose = mocker.AsyncMock() - fake_client.rag.retrieve_contexts = mocker.AsyncMock( - return_value=SimpleNamespace(contexts=SimpleNamespace(contexts=[])) + fake_client = _unlistable_client(mocker) + fake_client.rag.retrieve_contexts.return_value = SimpleNamespace( + contexts=SimpleNamespace(contexts=[]) ) mocker.patch( "agentplatform.Client", return_value=mocker.Mock(aio=fake_client) @@ -117,32 +154,191 @@ async def test_search_memory_forwards_similarity_top_k( assert kwargs["vertex_rag_store"].similarity_top_k is None +@pytest.mark.asyncio +async def test_search_memory_scopes_retrieval_to_tenant_files(mocker): + """Ranking happens over only the requesting app and user's files.""" + memory_service = VertexAiRagMemoryService( + rag_corpus="corpus", similarity_top_k=5 + ) + + fake_client = _async_client(mocker) + fake_client.rag.list_files.side_effect = [ + SimpleNamespace( + rag_files=[ + _rag_file( + "alice-1", + _build_source_display_name("demo", "alice", "session-1"), + ), + _rag_file( + "bob-1", + _build_source_display_name("demo", "bob", "session-2"), + ), + ], + next_page_token="page-2", + ), + SimpleNamespace( + rag_files=[ + _rag_file("alice-2", "demo.alice.legacy-session"), + _rag_file( + "other-app-1", + _build_source_display_name("other", "alice", "session-3"), + ), + ], + next_page_token=None, + ), + ] + fake_client.rag.retrieve_contexts.return_value = SimpleNamespace( + contexts=SimpleNamespace( + contexts=[ + _rag_context( + _build_source_display_name("demo", "alice", "session-1"), + "ALICE_MEMORY", + ) + ] + ) + ) + mocker.patch( + "agentplatform.Client", return_value=mocker.Mock(aio=fake_client) + ) + + response = await memory_service.search_memory( + app_name="demo", user_id="alice", query="memory" + ) + + assert _memory_texts(response) == ["ALICE_MEMORY"] + retrieve_kwargs = fake_client.rag.retrieve_contexts.call_args.kwargs + scoped_store = retrieve_kwargs["vertex_rag_store"] + scoped_resources = scoped_store.rag_resources + assert [resource.rag_corpus for resource in scoped_resources] == ["corpus"] + assert scoped_resources[0].rag_file_ids == ["alice-1", "alice-2"] + # Rebuilding the store keeps top-k on the query and off the store. + assert retrieve_kwargs["query"].similarity_top_k == 5 + assert scoped_store.similarity_top_k is None + assert fake_client.rag.list_files.await_count == 2 + assert ( + fake_client.rag.list_files.call_args_list[1].kwargs["config"].page_token + == "page-2" + ) + fake_client.aclose.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_search_memory_skips_retrieval_without_tenant_files(mocker): + memory_service = VertexAiRagMemoryService(rag_corpus="corpus") + + fake_client = _async_client(mocker) + fake_client.rag.list_files.return_value = SimpleNamespace( + rag_files=[ + _rag_file( + "bob-1", + _build_source_display_name("demo", "bob", "session-2"), + ) + ], + next_page_token=None, + ) + mocker.patch( + "agentplatform.Client", return_value=mocker.Mock(aio=fake_client) + ) + + response = await memory_service.search_memory( + app_name="demo", user_id="alice", query="memory" + ) + + assert response.memories == [] + fake_client.rag.retrieve_contexts.assert_not_awaited() + # The early return still leaves through the finally that closes the client. + fake_client.aclose.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_search_memory_retrieves_unscoped_when_listing_fails(mocker): + """A deployment that cannot list files still retrieves its own memories.""" + memory_service = VertexAiRagMemoryService(rag_corpus="corpus") + + fake_client = _unlistable_client(mocker) + fake_client.rag.retrieve_contexts.return_value = SimpleNamespace( + contexts=SimpleNamespace( + contexts=[ + _rag_context( + _build_source_display_name("demo", "alice", "session-1"), + "ALICE_MEMORY", + ), + _rag_context( + _build_source_display_name("demo", "bob", "session-2"), + "BOB_MEMORY", + ), + ] + ) + ) + mocker.patch( + "agentplatform.Client", return_value=mocker.Mock(aio=fake_client) + ) + + response = await memory_service.search_memory( + app_name="demo", user_id="alice", query="memory" + ) + + assert _memory_texts(response) == ["ALICE_MEMORY"] + assert _retrieved_store(fake_client).rag_resources[0].rag_file_ids is None + + +@pytest.mark.asyncio +async def test_search_memory_retrieves_unscoped_when_corpus_is_too_large( + mocker, caplog +): + """Listing is capped so search cost stays independent of corpus size.""" + memory_service = VertexAiRagMemoryService(rag_corpus="corpus") + + fake_client = _async_client(mocker) + fake_client.rag.list_files.return_value = SimpleNamespace( + rag_files=[ + _rag_file( + "alice-1", + _build_source_display_name("demo", "alice", "session-1"), + ) + ], + next_page_token="another-page", + ) + fake_client.rag.retrieve_contexts.return_value = SimpleNamespace( + contexts=SimpleNamespace(contexts=[]) + ) + mocker.patch( + "agentplatform.Client", return_value=mocker.Mock(aio=fake_client) + ) + + with caplog.at_level(logging.WARNING): + await memory_service.search_memory( + app_name="demo", user_id="alice", query="memory" + ) + + assert fake_client.rag.list_files.await_count == _MAX_RAG_FILE_PAGES + assert _retrieved_store(fake_client).rag_resources[0].rag_file_ids is None + assert "not scoped to the requesting app and user" in caplog.text + + @pytest.mark.asyncio async def test_search_memory_rejects_ambiguous_legacy_display_names(mocker): """Ensures dotted user IDs cannot match another user's legacy memory.""" memory_service = VertexAiRagMemoryService(rag_corpus="unused") - fake_client = mocker.Mock() - fake_client.aclose = mocker.AsyncMock() - fake_client.rag.retrieve_contexts = mocker.AsyncMock( - return_value=SimpleNamespace( - contexts=SimpleNamespace( - contexts=[ - _rag_context( - "demo.alice.smith.session_secret", - "SECRET_FROM_ALICE_SMITH", - ), - _rag_context( - _build_source_display_name("demo", "alice", "session_ok"), - "NORMAL_ALICE_MEMORY", - ), - _rag_context( - "demo.alice.legacy_session", - "LEGACY_ALICE_MEMORY", - ), - _rag_context("demo.bob.session_other", "BOB_MEMORY"), - ] - ) + fake_client = _unlistable_client(mocker) + fake_client.rag.retrieve_contexts.return_value = SimpleNamespace( + contexts=SimpleNamespace( + contexts=[ + _rag_context( + "demo.alice.smith.session_secret", + "SECRET_FROM_ALICE_SMITH", + ), + _rag_context( + _build_source_display_name("demo", "alice", "session_ok"), + "NORMAL_ALICE_MEMORY", + ), + _rag_context( + "demo.alice.legacy_session", + "LEGACY_ALICE_MEMORY", + ), + _rag_context("demo.bob.session_other", "BOB_MEMORY"), + ] ) ) @@ -165,10 +361,7 @@ async def test_add_and_search_memory_uses_unambiguous_display_names( ): memory_service = VertexAiRagMemoryService(rag_corpus="unused") - fake_client = mocker.Mock() - fake_client.aclose = mocker.AsyncMock() - fake_client.rag.upload_file = mocker.AsyncMock() - fake_client.rag.retrieve_contexts = mocker.AsyncMock() + fake_client = _unlistable_client(mocker) mocker.patch( "agentplatform.Client", return_value=mocker.Mock(aio=fake_client) )