fix: Update VertexAiRagMemoryService methods to use latest RAG module

PiperOrigin-RevId: 938567500

Author: sararob

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:

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")]
       )
   )