feat: Grounding - Allow initialization of grounding.VertexAISearch with full resource name or data store ID, project ID, and location.

PiperOrigin-RevId: 667990245

Author: holtskinner

tests/unit/vertexai/test_generative_models.py

@@ -1023,15 +1023,38 @@ def test_generate_content_grounding_google_search_retriever(self):
     )
     def test_generate_content_grounding_vertex_ai_search_retriever(self):
         model = preview_generative_models.GenerativeModel("gemini-pro")
-        google_search_retriever_tool = preview_generative_models.Tool.from_retrieval(
+        vertex_ai_search_retriever_tool = preview_generative_models.Tool.from_retrieval(
             retrieval=preview_generative_models.grounding.Retrieval(
                 source=preview_generative_models.grounding.VertexAISearch(
                     datastore=f"projects/{_TEST_PROJECT}/locations/global/collections/default_collection/dataStores/test-datastore",
                 )
             )
         )
         response = model.generate_content(
-            "Why is sky blue?", tools=[google_search_retriever_tool]
+            "Why is sky blue?", tools=[vertex_ai_search_retriever_tool]
+        )
+        assert response.text
+
+    @mock.patch.object(
+        target=prediction_service.PredictionServiceClient,
+        attribute="generate_content",
+        new=mock_generate_content,
+    )
+    def test_generate_content_grounding_vertex_ai_search_retriever_with_project_and_location(
+        self,
+    ):
+        model = preview_generative_models.GenerativeModel("gemini-pro")
+        vertex_ai_search_retriever_tool = preview_generative_models.Tool.from_retrieval(
+            retrieval=preview_generative_models.grounding.Retrieval(
+                source=preview_generative_models.grounding.VertexAISearch(
+                    datastore="test-datastore",
+                    project=_TEST_PROJECT,
+                    location="global",
+                )
+            )
+        )
+        response = model.generate_content(
+            "Why is sky blue?", tools=[vertex_ai_search_retriever_tool]
         )
         assert response.text
 

vertexai/generative_models/_generative_models.py

@@ -20,6 +20,7 @@
 import io
 import json
 import pathlib
+import re
 from typing import (
     Any,
     AsyncIterable,
@@ -2273,7 +2274,7 @@ def __init__(
                 source (VertexAISearch):
                     Set to use data source powered by Vertex AI Search.
                 disable_attribution (bool):
-                    Optional. Disable using the result from this
+                    Deprecated. Disable using the result from this
                     tool in detecting grounding attribution. This
                     does not affect how the result is given to the
                     model for generation.
@@ -2284,22 +2285,38 @@ def __init__(
             )
 
     class VertexAISearch:
-        r"""Retrieve from Vertex AI Search datastore for grounding.
-        See https://cloud.google.com/vertex-ai-search-and-conversation
+        r"""Retrieve from Vertex AI Search data store for grounding.
+        See https://cloud.google.com/products/agent-builder
         """
 
         def __init__(
             self,
             datastore: str,
+            *,
+            project: Optional[str] = None,
+            location: Optional[str] = None,
         ):
             """Initializes a Vertex AI Search tool.
 
             Args:
                 datastore (str):
-                    Required. Fully-qualified Vertex AI Search's
-                    datastore resource ID.
-                    projects/<>/locations/<>/collections/<>/dataStores/<>
+                    Required. Vertex AI Search data store resource name. Format:
+                    ``projects/{project}/locations/{location}/collections/default_collection/dataStores/{data_store}``
+                    or ``{data_store}``.
+                project (str):
+                    Optional. Project ID of the data store. Must provide either the full data store resource name or data store id, project ID, and location.
+                location (str):
+                    Optional. Location of the data store. Must provide either the full data store resource name or data store id, project ID, and location.
             """
+            if not re.fullmatch(
+                r"^projects/[a-z0-9-]*/locations/[a-z0-9][a-z0-9-]*/collections/[a-z0-9][a-z0-9-_]*/dataStores/[a-z0-9][a-z0-9-_]*$",
+                datastore,
+            ):
+                if not project or not location:
+                    raise ValueError(
+                        "Must provide either the full data store resource name or data store id, project ID, and location."
+                    )
+                datastore = f"projects/{project}/locations/{location}/collections/default_collection/dataStores/{datastore}"
             self._raw_vertex_ai_search = gapic_tool_types.VertexAISearch(
                 datastore=datastore,
             )