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

PiperOrigin-RevId: 890232571

Author: google-genai-bot

google/genai/models.py

@@ -4137,7 +4137,13 @@ def _Tool_to_vertex(
 ) -> dict[str, Any]:
   to_object: dict[str, Any] = {}
   if getv(from_object, ['retrieval']) is not None:
-    setv(to_object, ['retrieval'], getv(from_object, ['retrieval']))
+    setv(
+        to_object,
+        ['retrieval'],
+        _Retrieval_to_vertex(
+            getv(from_object, ['retrieval']), to_object, root_object
+        ),
+    )
 
   if getv(from_object, ['computer_use']) is not None:
     setv(to_object, ['computerUse'], getv(from_object, ['computer_use']))
@@ -4194,6 +4200,159 @@ def _Tool_to_vertex(
   return to_object
 
 
+def _Retrieval_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+    root_object: Optional[Union[dict[str, Any], object]] = None,
+) -> dict[str, Any]:
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ['disable_attribution']) is not None:
+    setv(
+        to_object, ['disableAttribution'], getv(from_object, ['disable_attribution'])
+    )
+
+  if getv(from_object, ['vertex_ai_search']) is not None:
+    setv(
+        to_object,
+        ['vertexAiSearch'],
+        _VertexAISearch_to_vertex(
+            getv(from_object, ['vertex_ai_search']), to_object, root_object
+        ),
+    )
+
+  if getv(from_object, ['vertex_rag_store']) is not None:
+    setv(
+        to_object,
+        ['vertexRagStore'],
+        _VertexRagStore_to_vertex(
+            getv(from_object, ['vertex_rag_store']), to_object, root_object
+        ),
+    )
+
+  return to_object
+
+
+def _VertexAISearch_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+    root_object: Optional[Union[dict[str, Any], object]] = None,
+) -> dict[str, Any]:
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ['data_store_specs']) is not None:
+    setv(to_object, ['dataStoreSpecs'], getv(from_object, ['data_store_specs']))
+  if getv(from_object, ['datastore']) is not None:
+    setv(to_object, ['datastore'], getv(from_object, ['datastore']))
+  if getv(from_object, ['engine']) is not None:
+    setv(to_object, ['engine'], getv(from_object, ['engine']))
+  if getv(from_object, ['filter']) is not None:
+    setv(to_object, ['filter'], getv(from_object, ['filter']))
+  if getv(from_object, ['max_results']) is not None:
+    setv(to_object, ['maxResults'], getv(from_object, ['max_results']))
+  return to_object
+
+
+def _VertexRagStore_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+    root_object: Optional[Union[dict[str, Any], object]] = None,
+) -> dict[str, Any]:
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ['rag_resources']) is not None:
+    setv(
+        to_object,
+        ['ragResources'],
+        [
+            _VertexRagStoreRagResource_to_vertex(item, to_object, root_object)
+            for item in getv(from_object, ['rag_resources'])
+        ],
+    )
+
+  if getv(from_object, ['rag_retrieval_config']) is not None:
+    setv(
+        to_object,
+        ['ragRetrievalConfig'],
+        _RagRetrievalConfig_to_vertex(
+            getv(from_object, ['rag_retrieval_config']), to_object, root_object
+        ),
+    )
+
+  if getv(from_object, ['similarity_top_k']) is not None:
+    setv(to_object, ['similarityTopK'], getv(from_object, ['similarity_top_k']))
+
+  if getv(from_object, ['vector_distance_threshold']) is not None:
+    setv(
+        to_object,
+        ['vectorDistanceThreshold'],
+        getv(from_object, ['vector_distance_threshold']),
+    )
+
+  return to_object
+
+
+def _VertexRagStoreRagResource_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+    root_object: Optional[Union[dict[str, Any], object]] = None,
+) -> dict[str, Any]:
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ['rag_corpus_name']) is not None:
+    setv(to_object, ['ragCorpus'], getv(from_object, ['rag_corpus_name']))
+  elif getv(from_object, ['rag_corpus']) is not None:
+    setv(to_object, ['ragCorpus'], getv(from_object, ['rag_corpus']))
+
+  if getv(from_object, ['rag_file_ids']) is not None:
+    setv(to_object, ['ragFileIds'], getv(from_object, ['rag_file_ids']))
+
+  return to_object
+
+
+def _RagRetrievalConfig_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+    root_object: Optional[Union[dict[str, Any], object]] = None,
+) -> dict[str, Any]:
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ['top_k']) is not None:
+    setv(to_object, ['topK'], getv(from_object, ['top_k']))
+
+  if getv(from_object, ['filter']) is not None:
+    setv(
+        to_object,
+        ['filter'],
+        _RagRetrievalConfigFilter_to_vertex(
+            getv(from_object, ['filter']), to_object, root_object
+        ),
+    )
+
+  return to_object
+
+
+def _RagRetrievalConfigFilter_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+    root_object: Optional[Union[dict[str, Any], object]] = None,
+) -> dict[str, Any]:
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ['vector_distance_threshold']) is not None:
+    setv(
+        to_object,
+        ['vectorDistanceThreshold'],
+        getv(from_object, ['vector_distance_threshold']),
+    )
+
+  if getv(from_object, ['vector_similarity_threshold']) is not None:
+    setv(
+        to_object,
+        ['vectorSimilarityThreshold'],
+        getv(from_object, ['vector_similarity_threshold']),
+    )
+
+  if getv(from_object, ['metadata_filter']) is not None:
+    setv(to_object, ['metadataFilter'], getv(from_object, ['metadata_filter']))
+
+  return to_object
+
+
 def _TunedModelInfo_from_vertex(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,

google/genai/tests/models/test_generate_content_tools.py

@@ -290,6 +290,34 @@ def divide_floats(a: float, b: float) -> float:
         ),
         exception_if_mldev='retrieval',
     ),
+    pytest_helper.TestTableItem(
+        name='test_rag_with_metadata_filter',
+        parameters=types._GenerateContentParameters(
+            model='gemini-2.5-flash',
+            contents=t.t_contents('What is the capital of France?'),
+            config={
+                'tools': [
+                    types.Tool(
+                        retrieval=types.Retrieval(
+                            vertex_rag_store=types.VertexRagStore(
+                                rag_resources=[
+                                    types.VertexRagStoreRagResource(
+                                        rag_corpus='projects/test-project/locations/us-central1/ragCorpora/test-corpus'
+                                    )
+                                ],
+                                rag_retrieval_config=types.RagRetrievalConfig(
+                                    filter=types.RagRetrievalConfigFilter(
+                                        metadata_filter='color = "red"',
+                                    ),
+                                ),
+                            )
+                        ),
+                    ),
+                ]
+            },
+        ),
+        exception_if_mldev='retrieval',
+    ),
     pytest_helper.TestTableItem(
         name='test_file_search',
         parameters=types._GenerateContentParameters(

google/genai/tests/verify_transform.py

@@ -0,0 +1,43 @@
+import sys
+import os
+
+# Set up path so we can import and run scripts as if they were in the workspace.
+sys.path.append(os.getcwd())
+
+from third_party.py.google.genai import types
+from third_party.py.google.genai import models
+from third_party.py.google.genai import _common
+import json
+
+class MockApiClient:
+    def __init__(self):
+        self.vertexai = True
+        self.location = 'us-central1'
+        self.project = 'test-project'
+
+def test_transformation():
+    tool = types.Tool(
+        retrieval=types.Retrieval(
+            vertex_rag_store=types.VertexRagStore(
+                rag_resources=[
+                    types.VertexRagStoreRagResource(
+                        rag_corpus='projects/test-project/locations/us-central1/ragCorpora/test-corpus'
+                    )
+                ],
+                rag_retrieval_config=types.RagRetrievalConfig(
+                    filter=types.RagRetrievalConfigFilter(
+                        metadata_filter='color = "red"',
+                    ),
+                ),
+            )
+        ),
+    )
+
+    api_client = MockApiClient()
+    # Mocking internal transform
+    transformed = models._Tool_to_vertex(api_client, tool)
+    final_dict = _common.convert_to_dict(transformed)
+    print(json.dumps(final_dict, indent=2))
+
+if __name__ == '__main__':
+    test_transformation()