feat: Add metadata_filter usage to all retrieval and generator methods in rag_retrieval.py

vertex-sdk-bot · copybara-github · commit 841c597c35d9 · 2026-03-25T16:41:32.000-07:00
PiperOrigin-RevId: 889497482
diff --git a/tests/unit/vertex_rag/test_rag_retrieval_preview.py b/tests/unit/vertex_rag/test_rag_retrieval_preview.py
@@ -223,6 +223,27 @@ def test_retrieval_query_rag_corpora_config_rank_service_success(self):
         )
         retrieve_contexts_eq(response, tc.TEST_RETRIEVAL_RESPONSE)
 
+    @pytest.mark.usefixtures("retrieve_contexts_mock")
+    def test_retrieval_query_with_metadata_filter(self, retrieve_contexts_mock):
+        metadata_filter = 'doc.metadata.genre == "fiction"'
+        rag_retrieval_config = rag.RagRetrievalConfig(
+            top_k=10,
+            filter=rag.Filter(
+                vector_distance_threshold=0.5, metadata_filter=metadata_filter
+            ),
+        )
+        rag.retrieval_query(
+            rag_resources=[tc.TEST_RAG_RESOURCE],
+            text=tc.TEST_QUERY_TEXT,
+            rag_retrieval_config=rag_retrieval_config,
+        )
+        retrieve_contexts_mock.assert_called_once()
+        args, kwargs = retrieve_contexts_mock.call_args
+        request = kwargs["request"]
+        assert (
+            request.query.rag_retrieval_config.filter.metadata_filter == metadata_filter
+        )
+
     @pytest.mark.usefixtures("retrieve_contexts_mock")
     def test_retrieval_query_rag_corpora_config_llm_ranker_success(self):
         response = rag.retrieval_query(
diff --git a/vertexai/preview/rag/rag_retrieval.py b/vertexai/preview/rag/rag_retrieval.py
@@ -246,6 +246,10 @@ def retrieval_query(
             api_retrival_config.filter.vector_similarity_threshold = (
                 rag_retrieval_config.filter.vector_similarity_threshold
             )
+        if rag_retrieval_config.filter and rag_retrieval_config.filter.metadata_filter:
+            api_retrival_config.filter.metadata_filter = (
+                rag_retrieval_config.filter.metadata_filter
+            )
 
         if (
             rag_retrieval_config.ranking
@@ -495,6 +499,10 @@ async def async_retrieve_contexts(
             api_retrival_config.ranking.llm_ranker.model_name = (
                 rag_retrieval_config.ranking.llm_ranker.model_name
             )
+        if rag_retrieval_config.filter and rag_retrieval_config.filter.metadata_filter:
+            api_retrival_config.filter.metadata_filter = (
+                rag_retrieval_config.filter.metadata_filter
+            )
 
     query = aiplatform_v1beta1.RagQuery(
         text=text,
@@ -742,6 +750,10 @@ def ask_contexts(
             api_retrival_config.ranking.llm_ranker.model_name = (
                 rag_retrieval_config.ranking.llm_ranker.model_name
             )
+        if rag_retrieval_config.filter and rag_retrieval_config.filter.metadata_filter:
+            api_retrival_config.filter.metadata_filter = (
+                rag_retrieval_config.filter.metadata_filter
+            )
 
     query = aiplatform_v1beta1.RagQuery(
         text=text,
