From 96d2914339ade4409b2366b80a0668ca8001aedb Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Wed, 8 Jul 2026 10:53:17 -0700 Subject: [PATCH] fix: Update VertexAiRagMemoryService methods to use latest RAG module PiperOrigin-RevId: 944587802 --- .../memory/vertex_ai_rag_memory_service.py | 57 ++++++++++++----- .../test_vertex_ai_rag_memory_service.py | 61 +++++++++---------- 2 files changed, 70 insertions(+), 48 deletions(-) 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 09745bcff7a..357603830e6 100644 --- a/src/google/adk/memory/vertex_ai_rag_memory_service.py +++ b/src/google/adk/memory/vertex_ai_rag_memory_service.py @@ -90,36 +90,56 @@ def _parse_source_display_name( class VertexAiRagMemoryService(BaseMemoryService): - """A memory service that uses Vertex AI RAG for storage and retrieval.""" + """A memory service that uses Agent Platform RAG for storage and retrieval.""" def __init__( self, rag_corpus: Optional[str] = None, similarity_top_k: Optional[int] = None, vector_distance_threshold: float = 10, + project: Optional[str] = None, + location: Optional[str] = None, ): """Initializes a VertexAiRagMemoryService. Args: - rag_corpus: The name of the Vertex AI RAG corpus to use. Format: + rag_corpus: The name of the Agent Platform RAG corpus to use. Format: ``projects/{project}/locations/{location}/ragCorpora/{rag_corpus_id}`` or ``{rag_corpus_id}`` similarity_top_k: The number of contexts to retrieve. vector_distance_threshold: Only returns contexts with vector distance smaller than the threshold. + project: The project to use for the Agent Platform RAG corpus. If not + set, the value of the GOOGLE_CLOUD_PROJECT environment variable is + used. + location: The location to use for the Agent Platform RAG corpus. If not + set, the value of the GOOGLE_CLOUD_LOCATION environment variable is + used. """ try: - import vertexai # noqa: F401 + import agentplatform # noqa: F401 except ImportError as e: from ..utils._dependency import missing_extra raise missing_extra("google-cloud-aiplatform", "gcp") from e + self._project = project or os.environ.get("GOOGLE_CLOUD_PROJECT") + self._location = location or os.environ.get("GOOGLE_CLOUD_LOCATION") + + # Fallback: if the fully-qualified corpus name is provided, use it to + # determine the project and location, if they are not already set. + if (not self._project or not self._location) and ( + rag_corpus and rag_corpus.startswith("projects/") + ): + parts = rag_corpus.split("/") + if len(parts) >= 4 and parts[0] == "projects" and parts[2] == "locations": + self._project = self._project or parts[1] + self._location = self._location or parts[3] + self._vertex_rag_store = types.VertexRagStore( rag_resources=[ types.VertexRagStoreRagResource(rag_corpus=rag_corpus), ], - similarity_top_k=similarity_top_k, vector_distance_threshold=vector_distance_threshold, ) @@ -153,14 +173,16 @@ async def add_session_to_memory(self, session: Session) -> None: if not self._vertex_rag_store.rag_resources: raise ValueError("Rag resources must be set.") - from ..dependencies.vertexai import rag + import agentplatform + + client = agentplatform.Client( + project=self._project, location=self._location + ) for rag_resource in self._vertex_rag_store.rag_resources: - rag.upload_file( + client.rag.upload_file( corpus_name=rag_resource.rag_corpus, path=temp_file_path, - # this is the temp workaround as upload file does not support - # adding metadata, thus use display_name to store the session info. display_name=_build_source_display_name( session.app_name, session.user_id, session.id ), @@ -173,17 +195,22 @@ async def search_memory( self, *, app_name: str, user_id: str, query: str ) -> SearchMemoryResponse: """Searches for sessions that match the query using rag.retrieval_query.""" - from ..dependencies.vertexai import rag + import agentplatform + from agentplatform import types as agentplatform_types + from ..events.event import Event - response = rag.retrieval_query( - text=query, - rag_resources=self._vertex_rag_store.rag_resources, - rag_corpora=self._vertex_rag_store.rag_corpora, - similarity_top_k=self._vertex_rag_store.similarity_top_k, - vector_distance_threshold=self._vertex_rag_store.vector_distance_threshold, + client = agentplatform.Client( + project=self._project, location=self._location ) + response = client.rag.retrieve_contexts( + vertex_rag_store=self._vertex_rag_store, + query=agentplatform_types.RagQuery( + text=query, + similarity_top_k=self._vertex_rag_store.similarity_top_k, + ), + ) memory_results = [] session_events_map: OrderedDict[str, list[list[Event]]] = OrderedDict() for context in response.contexts.contexts: 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 7c20de87f5a..a08bfaddda1 100644 --- a/tests/unittests/memory/test_vertex_ai_rag_memory_service.py +++ b/tests/unittests/memory/test_vertex_ai_rag_memory_service.py @@ -35,32 +35,29 @@ def _rag_context(source_display_name: str, text: str) -> SimpleNamespace: 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_rag = SimpleNamespace( - retrieval_query=mocker.Mock( - 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 = mocker.Mock() + 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"), + ] ) ) - mocker.patch("google.adk.dependencies.vertexai.rag", fake_rag) + + mocker.patch("agentplatform.Client", return_value=fake_client) response = await memory_service.search_memory( app_name="demo", user_id="alice", query="secret" @@ -73,9 +70,9 @@ async def test_search_memory_rejects_ambiguous_legacy_display_names(mocker): @pytest.mark.asyncio async def test_add_and_search_memory_uses_unambiguous_display_names(mocker): memory_service = VertexAiRagMemoryService(rag_corpus="unused") - upload_file = mocker.Mock() - fake_rag = SimpleNamespace(upload_file=upload_file) - mocker.patch("google.adk.dependencies.vertexai.rag", fake_rag) + + fake_client = mocker.Mock() + mocker.patch("agentplatform.Client", return_value=fake_client) await memory_service.add_session_to_memory( Session( @@ -96,15 +93,13 @@ async def test_add_and_search_memory_uses_unambiguous_display_names(mocker): ) ) - display_name = upload_file.call_args.kwargs["display_name"] + display_name = fake_client.rag.upload_file.call_args.kwargs["display_name"] assert display_name.startswith(_SOURCE_DISPLAY_NAME_PREFIX) assert display_name != "demo.app.alice.smith.session.secret" - fake_rag.retrieval_query = mocker.Mock( - return_value=SimpleNamespace( - contexts=SimpleNamespace( - contexts=[_rag_context(display_name, "sensitive memory")] - ) + fake_client.rag.retrieve_contexts.return_value = SimpleNamespace( + contexts=SimpleNamespace( + contexts=[_rag_context(display_name, "sensitive memory")] ) )