feat: Add Tool compatibility with RagRetrievalConfig in both Vertex AI SDK and GenAI SDK for use with generate_content.

PiperOrigin-RevId: 890232571

Author: vertex-sdk-bot

tests/unit/vertexai/test_generative_models.py

@@ -33,8 +33,8 @@
 from vertexai import generative_models
 from vertexai.preview import (
     generative_models as preview_generative_models,
-    rag,
 )
+from vertexai.preview.rag.utils import resources
 from vertexai.generative_models._generative_models import (
     prediction_service,
     gapic_prediction_service_types,
@@ -298,7 +298,7 @@ def mock_generate_content(
         tool.retrieval or tool.google_search_retrieval for tool in request.tools
     )
     has_rag_retrieval = any(
-        isinstance(tool.retrieval, rag.Retrieval) for tool in request.tools
+        isinstance(tool.retrieval, generative_models.grounding.Retrieval) and tool.retrieval._raw_retrieval._pb.WhichOneof("source") == "vertex_rag_store" for tool in request.tools
     )
     has_function_declarations = any(
         tool.function_declarations for tool in request.tools
@@ -367,6 +367,23 @@ def mock_generate_content(
             web_search_queries=[request.contents[0].parts[0].text],
         )
     elif has_rag_retrieval and request.contents[0].parts[0].text:
+        vertex_rag_store = request.tools[0].retrieval.vertex_rag_store
+        assert vertex_rag_store is not None
+        # Validate rag_resources or rag_corpora
+        if vertex_rag_store.rag_resources:
+            assert len(vertex_rag_store.rag_resources) == 1
+            assert vertex_rag_store.rag_resources[0].rag_corpus == f"projects/{_TEST_PROJECT}/locations/us-central1/ragCorpora/1234556"
+            assert vertex_rag_store.rag_resources[0].rag_file_ids == ["123", "456"]
+        elif vertex_rag_store.rag_corpora:
+            assert vertex_rag_store.rag_corpora == [f"projects/{_TEST_PROJECT}/locations/us-central1/ragCorpora/654321"]
+
+        # Validate rag_retrieval_config and metadata_filter
+        assert vertex_rag_store.rag_retrieval_config is not None
+        assert vertex_rag_store.rag_retrieval_config.top_k == 1
+        assert (
+            vertex_rag_store.rag_retrieval_config.filter.metadata_filter
+            == "test_metadata_filter"
+        )
         grounding_metadata = gapic_content_types.GroundingMetadata(
             retrieval_queries=[request.contents[0].parts[0].text],
         )
@@ -1400,20 +1417,29 @@ def test_generate_content_grounding_vertex_ai_search_retriever_with_project_and_
         assert response.text
 
     @patch_genai_services
-    def test_generate_content_vertex_rag_retriever(self):
-        model = preview_generative_models.GenerativeModel("gemini-pro")
+    @pytest.mark.parametrize(
+        "generative_models",
+        [generative_models, preview_generative_models],
+    )
+    def test_generate_content_with_vertex_rag_store_tool(self, generative_models: generative_models):
+        model = generative_models.GenerativeModel("gemini-pro")
         rag_resources = [
-            rag.RagResource(
+            resources.RagResource(
                 rag_corpus=f"projects/{_TEST_PROJECT}/locations/us-central1/ragCorpora/1234556",
                 rag_file_ids=["123", "456"],
             ),
         ]
-        rag_retriever_tool = preview_generative_models.Tool.from_retrieval(
-            retrieval=rag.Retrieval(
-                source=rag.VertexRagStore(
+        rag_retrieval_config = resources.RagRetrievalConfig(
+            top_k=1,
+            filter=resources.Filter(
+                metadata_filter="test_metadata_filter"
+            ),
+        )
+        rag_retriever_tool = generative_models.Tool.from_retrieval(
+            retrieval=generative_models.grounding.Retrieval(
+                source=generative_models.grounding.VertexRagStore(
                     rag_resources=rag_resources,
-                    similarity_top_k=1,
-                    vector_distance_threshold=0.4,
+                    rag_retrieval_config=rag_retrieval_config,
                 ),
             ),
         )
@@ -1422,6 +1448,103 @@ def test_generate_content_vertex_rag_retriever(self):
         )
         assert response.text
 
+    @patch_genai_services
+    @pytest.mark.parametrize(
+        "generative_models",
+        [generative_models, preview_generative_models],
+    )
+    def test_generate_content_with_vertex_rag_store_tool_with_rag_corpora(self, generative_models: generative_models):
+        model = generative_models.GenerativeModel("gemini-pro")
+        rag_corpora = [
+            f"projects/{_TEST_PROJECT}/locations/us-central1/ragCorpora/654321",
+        ]
+        rag_retrieval_config = resources.RagRetrievalConfig(
+            top_k=1,
+            filter=resources.Filter(
+                metadata_filter="test_metadata_filter"
+            ),
+        )
+        rag_retriever_tool = generative_models.Tool.from_retrieval(
+            retrieval=generative_models.grounding.Retrieval(
+                source=generative_models.grounding.VertexRagStore(
+                    rag_corpora=rag_corpora,
+                    rag_retrieval_config=rag_retrieval_config,
+                ),
+            ),
+        )
+        response = model.generate_content(
+            "Why is sky blue?", tools=[rag_retriever_tool]
+        )
+        assert response.text
+
+    @patch_genai_services
+    @pytest.mark.parametrize(
+        "generative_models",
+        [generative_models, preview_generative_models],
+    )
+    def test_generate_content_with_vertex_rag_store_tool_with_deprecated_fields(self, generative_models: generative_models):
+        model = generative_models.GenerativeModel("gemini-pro")
+        rag_resources = [
+            resources.RagResource(
+                rag_corpus=f"projects/{_TEST_PROJECT}/locations/us-central1/ragCorpora/1234556",
+                rag_file_ids=["123", "456"],
+            ),
+        ]
+        # Deprecated fields are now inside RagRetrievalConfig
+        rag_retrieval_config = resources.RagRetrievalConfig(
+            top_k=1,
+            filter=resources.Filter(
+                vector_distance_threshold=0.4,
+            ),
+        )
+        # Expect warnings for deprecated fields
+        with pytest.warns(DeprecationWarning):
+            rag_retriever_tool = generative_models.Tool.from_retrieval(
+                retrieval=generative_models.grounding.Retrieval(
+                    source=generative_models.grounding.VertexRagStore(
+                        rag_resources=rag_resources,
+                        rag_retrieval_config=rag_retrieval_config,
+                    ),
+                ),
+            )
+        response = model.generate_content("Why is sky blue?", tools=[rag_retriever_tool])
+        assert response.text
+
+    @patch_genai_services
+    @pytest.mark.parametrize(
+        "generative_models",
+        [generative_models, preview_generative_models],
+    )
+    def test_generate_content_with_vertex_rag_store_tool_invalid_config(self, generative_models: generative_models):
+        model = generative_models.GenerativeModel("gemini-pro")
+        rag_resources = [
+            resources.RagResource(
+                rag_corpus=f"projects/{_TEST_PROJECT}/locations/us-central1/ragCorpora/1234556",
+                rag_file_ids=["123", "456"],
+            ),
+        ]
+        rag_retrieval_config = resources.RagRetrievalConfig(
+            top_k=1,
+            filter=resources.Filter(
+                vector_distance_threshold=0.4,
+                vector_similarity_threshold=0.6,
+            ),
+        )
+        with pytest.raises(ValueError) as e:
+            rag_retriever_tool = generative_models.Tool.from_retrieval(
+                retrieval=generative_models.grounding.Retrieval(
+                    source=generative_models.grounding.VertexRagStore(
+                        rag_resources=rag_resources,
+                        rag_retrieval_config=rag_retrieval_config,
+                    ),
+                ),
+            )
+        e.match(
+            "Only one of vector_distance_threshold or"
+            " vector_similarity_threshold can be specified at a time"
+            " in rag_retrieval_config."
+        )
+
     @patch_genai_services
     def test_chat_automatic_function_calling_with_function_returning_dict(self):
         generative_models = preview_generative_models

vertexai/generative_models/_generative_models.py

@@ -66,6 +66,7 @@
 
 if TYPE_CHECKING:
     from vertexai.caching import CachedContent
+    from vertexai.preview import generative_models as preview_generative_models
 
 try:
     from PIL import Image as PIL_Image  # pylint: disable=g-import-not-at-top
@@ -2078,8 +2079,13 @@ def from_retrieval(
         cls,
         retrieval: Union["grounding.Retrieval"],
     ) -> "Tool":
-        raw_tool = gapic_tool_types.Tool(retrieval=retrieval._raw_retrieval)
-        return cls._from_gapic(raw_tool=raw_tool)
+        # Late import to avoid circular dependency
+        from vertexai.preview import generative_models as preview_generative_models
+        raw_retrieval = retrieval._raw_retrieval
+        raw_tool = gapic_tool_types.Tool(retrieval=raw_retrieval)
+        response = cls._from_gapic(raw_tool=raw_tool)
+        response.retrieval = retrieval
+        return response
 
     @classmethod
     def from_google_search_retrieval(
@@ -2978,24 +2984,35 @@ class Retrieval:
 
         def __init__(
             self,
-            source: Union["grounding.VertexAISearch"],
+            source: Union["grounding.VertexAISearch", "grounding.VertexRagStore"],
             disable_attribution: Optional[bool] = None,
         ):
             """Initializes a Retrieval tool.
 
             Args:
-                source (VertexAISearch):
-                    Set to use data source powered by Vertex AI Search.
+                source (Union[VertexAISearch, VertexRagStore]):
+                    Set to use data source powered by Vertex AI Search or Vertex Rag Store.
                 disable_attribution (bool):
                     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.
             """
-            self._raw_retrieval = gapic_tool_types.Retrieval(
-                vertex_ai_search=source._raw_vertex_ai_search,
-                disable_attribution=disable_attribution,
-            )
+            if isinstance(source, grounding.VertexAISearch):
+                self._raw_retrieval = gapic_tool_types.Retrieval(
+                    vertex_ai_search=source._raw_vertex_ai_search,
+                    disable_attribution=disable_attribution,
+                )
+            elif isinstance(source, grounding.VertexRagStore):
+                # Late import to avoid circular dependency
+                from vertexai.preview import generative_models as preview_generative_models
+                gapic_vertex_rag_store = preview_generative_models._preview_parse_vertex_rag_store_to_api(source)
+                self._raw_retrieval = gapic_tool_types.Retrieval(
+                    vertex_rag_store=gapic_vertex_rag_store,
+                    disable_attribution=disable_attribution,
+                )
+            else:
+                raise TypeError(f"Unexpected source type: {type(source)}")
 
     class VertexAISearch:
         r"""Retrieve from Vertex AI Search data store for grounding.
@@ -3034,6 +3051,29 @@ def __init__(
                 datastore=datastore,
             )
 
+    class VertexRagStore:
+        """Retrieve from Vertex Rag Store for grounding.
+
+        Attributes:
+            rag_resources: A list of RagResource. It can be used to specify corpus
+              only or ragfiles. Currently only support one corpus or multiple files
+              from one corpus. In the future we may open up multiple corpora support.
+            rag_corpora: If rag_resources is not specified, use rag_corpora as a list
+              of rag corpora names. Deprecated. Use rag_resources instead.
+            rag_retrieval_config: Optional. The config containing the retrieval
+              parameters, including top_k, vector_distance_threshold, and alpha.
+        """
+
+        def __init__(
+            self,
+            rag_resources: Optional[List["preview_generative_models.rag.RagResource"]] = None,
+            rag_corpora: Optional[List[str]] = None,
+            rag_retrieval_config: Optional["preview_generative_models.rag.RagRetrievalConfig"] = None,
+        ):
+            self.rag_resources = rag_resources
+            self.rag_corpora = rag_corpora
+            self.rag_retrieval_config = rag_retrieval_config
+
 
 class preview_grounding(grounding):  # pylint: disable=invalid-name
     """Grounding namespace (preview)."""

vertexai/preview/fixing_genai_generative.md

@@ -0,0 +1,46 @@
+# RAG with Generative Models in Google GenAI SDK and Vertex AI SDK
+
+This document outlines the file paths and class definitions relevant to integrating `RagRetrievalConfig` and its `metadata_filter` within the `Tool` input for `generate_content` in both the Google GenAI SDK (`google.genai`) and the Vertex AI SDK (`google.cloud.aiplatform`).
+
+## Relationship between SDKs
+
+While both SDKs target the Vertex AI platform, they have independent implementations for `generate_content`:
+
+*   **Google GenAI SDK (`google.genai`)**: A standalone, Pydantic-based SDK that maps directly to the Gemini REST API. It does **not** depend on the Vertex AI Python SDK's runtime code for generation. It uses its own types and transformation logic (in `models.py`) to call the API.
+*   **Vertex AI SDK (`google.cloud.aiplatform.vertexai`)**: A high-level SDK that uses GAPIC (proto-based) clients internally. It has its own `GenerativeModel` implementation in `_generative_models.py`.
+
+## Relevant File Paths
+
+### Google GenAI SDK (`google.genai`)
+
+*   **`google3/third_party/py/google/genai/types.py`**: Defines Pydantic models for `Tool`, `Retrieval`, `VertexRagStore`, `RagRetrievalConfig`, and `RagRetrievalConfigFilter`. **Verified**: These types already include the `metadata_filter` field.
+*   **`google3/third_party/py/google/genai/models.py`**: Contains the logic to transform GenAI types to Vertex AI REST API format (see `_Tool_to_vertex`). **Note**: Explicit mapping for all `Retrieval` children to camelCase might be needed here if the API is case-sensitive.
+
+### Vertex AI SDK (`google.cloud.aiplatform.vertexai`)
+
+*   **`google3/third_party/py/google/cloud/aiplatform/vertexai/generative_models/_generative_models.py`**: Contains the core implementation of `GenerativeModel`. Includes the `grounding` namespace with `Retrieval` and `VertexRagStore` classes.
+*   **`google3/third_party/py/google/cloud/aiplatform/vertexai/preview/generative_models.py`**: Contains the `_preview_parse_vertex_rag_store_to_api` function which converts SDK types to GAPIC proto messages.
+*   **`google3/third_party/py/google/cloud/aiplatform/vertexai/preview/rag/utils/resources.py`**: Defines the Python dataclasses for `RagRetrievalConfig`, `Filter`, etc., used by the Vertex AI SDK.
+
+## Unit Test File Locations
+
+*   **`google3/third_party/py/google/genai/tests/models/test_generate_content_tools.py`**: Tests for RAG tools in the GenAI SDK.
+*   **`google3/third_party/py/google/cloud/aiplatform/tests/unit/vertexai/test_generative_models.py`**: Tests for `GenerativeModel` in the Vertex AI SDK.
+
+## Current Implementation Status
+
+### Vertex AI SDK (google-cloud-aiplatform)
+
+1.  **`grounding.VertexRagStore`**: **Implemented**. Exists in `_generative_models.py` and correctly holds `RagResource` and `RagRetrievalConfig`.
+2.  **`grounding.Retrieval`**: **Implemented**. Accepts `VertexRagStore` as a source and correctly routes to the preview parser.
+3.  **Mapping to GAPIC**: **Implemented**. The function `_preview_parse_vertex_rag_store_to_api` in `preview/generative_models.py` correctly maps `metadata_filter` to the underlying proto message.
+
+### Google GenAI SDK (google.genai)
+
+1.  **Type Definitions**: **Implemented**. `types.py` has all required classes for RAG with `metadata_filter`.
+2.  **Mapping to REST**: **Partial/Pending**. `models.py` passes the `Retrieval` object. If the backend requires camelCase keys for nested RAG fields, explicit mapping in `_Tool_to_vertex` or a dedicated `_Retrieval_to_vertex` function is required.
+
+## Next Steps
+
+1.  **Verify GenAI SDK REST Mapping**: Ensure that `metadata_filter` and other RAG fields are correctly serialized to camelCase (e.g., `metadataFilter`) when calling the Vertex AI REST API from `google.genai`.
+2.  **Add End-to-End Tests**: Add a test case in `google.genai/tests/models/test_generate_content_tools.py` that specifically uses a `metadata_filter` and verifies the resulting request payload.