feat: GenAI - Grounding - Released VertexAiSearch and Retrieval to GA

PiperOrigin-RevId: 702481252

Author: Ark-kun

tests/unit/vertexai/test_generative_models.py

@@ -1308,20 +1308,13 @@ def test_conversion_methods(self, generative_models: generative_models):
         assert response.to_dict()["candidates"][0]["finish_reason"] == "STOP"
 
     @patch_genai_services
-    def test_generate_content_grounding_google_search_retriever_preview(self):
-        model = preview_generative_models.GenerativeModel("gemini-pro")
-        google_search_retriever_tool = (
-            preview_generative_models.Tool.from_google_search_retrieval(
-                preview_generative_models.grounding.GoogleSearchRetrieval()
-            )
-        )
-        response = model.generate_content(
-            "Why is sky blue?", tools=[google_search_retriever_tool]
-        )
-        assert response.text
-
-    @patch_genai_services
-    def test_generate_content_grounding_google_search_retriever(self):
+    @pytest.mark.parametrize(
+        "generative_models",
+        [generative_models, preview_generative_models],
+    )
+    def test_generate_content_grounding_google_search_retriever(
+        self, generative_models: generative_models
+    ):
         model = generative_models.GenerativeModel("gemini-pro")
         google_search_retriever_tool = (
             generative_models.Tool.from_google_search_retrieval(
@@ -1334,11 +1327,17 @@ def test_generate_content_grounding_google_search_retriever(self):
         assert response.text
 
     @patch_genai_services
-    def test_generate_content_grounding_vertex_ai_search_retriever(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(
+    @pytest.mark.parametrize(
+        "generative_models",
+        [generative_models, preview_generative_models],
+    )
+    def test_generate_content_grounding_vertex_ai_search_retriever(
+        self, generative_models: generative_models
+    ):
+        model = generative_models.GenerativeModel("gemini-pro")
+        vertex_ai_search_retriever_tool = generative_models.Tool.from_retrieval(
+            retrieval=generative_models.grounding.Retrieval(
+                source=generative_models.grounding.VertexAISearch(
                     datastore=f"projects/{_TEST_PROJECT}/locations/global/collections/default_collection/dataStores/test-datastore",
                 )
             )
@@ -1349,13 +1348,17 @@ def test_generate_content_grounding_vertex_ai_search_retriever(self):
         assert response.text
 
     @patch_genai_services
+    @pytest.mark.parametrize(
+        "generative_models",
+        [generative_models, preview_generative_models],
+    )
     def test_generate_content_grounding_vertex_ai_search_retriever_with_project_and_location(
-        self,
+        self, generative_models: generative_models
     ):
-        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(
+        model = generative_models.GenerativeModel("gemini-pro")
+        vertex_ai_search_retriever_tool = generative_models.Tool.from_retrieval(
+            retrieval=generative_models.grounding.Retrieval(
+                source=generative_models.grounding.VertexAISearch(
                     datastore="test-datastore",
                     project=_TEST_PROJECT,
                     location="global",

vertexai/generative_models/_generative_models.py

@@ -1931,7 +1931,7 @@ def from_function_declarations(
     @classmethod
     def from_retrieval(
         cls,
-        retrieval: Union["preview_grounding.Retrieval"],
+        retrieval: Union["grounding.Retrieval"],
     ) -> "Tool":
         raw_tool = gapic_tool_types.Tool(retrieval=retrieval._raw_retrieval)
         return cls._from_gapic(raw_tool=raw_tool)
@@ -2767,16 +2767,6 @@ def __init__(
                 else None
             )
 
-
-class preview_grounding:  # pylint: disable=invalid-name
-    """Grounding namespace (preview)."""
-
-    __name__ = "grounding"
-    __module__ = "vertexai.preview.generative_models"
-
-    def __init__(self):
-        raise RuntimeError("This class must not be instantiated.")
-
     class Retrieval:
         """Defines a retrieval tool that model can call to access external knowledge."""
 
@@ -2838,7 +2828,15 @@ def __init__(
                 datastore=datastore,
             )
 
-    GoogleSearchRetrieval = grounding.GoogleSearchRetrieval
+
+class preview_grounding(grounding):  # pylint: disable=invalid-name
+    """Grounding namespace (preview)."""
+
+    __name__ = "grounding"
+    __module__ = "vertexai.preview.generative_models"
+
+    def __init__(self):
+        raise RuntimeError("This class must not be instantiated.")
 
 
 def _to_content(