feat: Add vector_similarity_threshold support within RagRetrievalConfig in rag_store and rag_retrieval GA and preview versions

PiperOrigin-RevId: 700812116

Author: vertex-sdk-bot

tests/unit/vertex_rag/test_rag_constants.py

@@ -508,3 +508,11 @@
     top_k=2,
     filter=Filter(vector_distance_threshold=0.5),
 )
+TEST_RAG_RETRIEVAL_SIMILARITY_CONFIG = RagRetrievalConfig(
+    top_k=2,
+    filter=Filter(vector_similarity_threshold=0.5),
+)
+TEST_RAG_RETRIEVAL_ERROR_CONFIG = RagRetrievalConfig(
+    top_k=2,
+    filter=Filter(vector_distance_threshold=0.5, vector_similarity_threshold=0.5),
+)

tests/unit/vertex_rag/test_rag_constants_preview.py

@@ -581,3 +581,12 @@
     filter=Filter(vector_distance_threshold=0.5),
     hybrid_search=HybridSearch(alpha=0.5),
 )
+TEST_RAG_RETRIEVAL_SIMILARITY_CONFIG = RagRetrievalConfig(
+    top_k=2,
+    filter=Filter(vector_distance_threshold=0.5),
+    hybrid_search=HybridSearch(alpha=0.5),
+)
+TEST_RAG_RETRIEVAL_ERROR_CONFIG = RagRetrievalConfig(
+    top_k=2,
+    filter=Filter(vector_distance_threshold=0.5, vector_similarity_threshold=0.5),
+)

tests/unit/vertex_rag/test_rag_retrieval.py

@@ -78,6 +78,15 @@ def test_retrieval_query_rag_resources_success(self):
         )
         retrieve_contexts_eq(response, tc.TEST_RETRIEVAL_RESPONSE)
 
+    @pytest.mark.usefixtures("retrieve_contexts_mock")
+    def test_retrieval_query_rag_resources_similarity_success(self):
+        response = rag.retrieval_query(
+            rag_resources=[tc.TEST_RAG_RESOURCE],
+            text=tc.TEST_QUERY_TEXT,
+            rag_retrieval_config=tc.TEST_RAG_RETRIEVAL_SIMILARITY_CONFIG,
+        )
+        retrieve_contexts_eq(response, tc.TEST_RETRIEVAL_RESPONSE)
+
     @pytest.mark.usefixtures("rag_client_mock_exception")
     def test_retrieval_query_failure(self):
         with pytest.raises(RuntimeError) as e:
@@ -105,3 +114,25 @@ def test_retrieval_query_multiple_rag_resources(self):
                 rag_retrieval_config=tc.TEST_RAG_RETRIEVAL_CONFIG,
             )
             e.match("Currently only support 1 RagResource")
+
+    def test_retrieval_query_similarity_multiple_rag_resources(self):
+        with pytest.raises(ValueError) as e:
+            rag.retrieval_query(
+                rag_resources=[tc.TEST_RAG_RESOURCE, tc.TEST_RAG_RESOURCE],
+                text=tc.TEST_QUERY_TEXT,
+                rag_retrieval_config=tc.TEST_RAG_RETRIEVAL_SIMILARITY_CONFIG,
+            )
+            e.match("Currently only support 1 RagResource")
+
+    def test_retrieval_query_invalid_config_filter(self):
+        with pytest.raises(ValueError) as e:
+            rag.retrieval_query(
+                rag_resources=[tc.TEST_RAG_RESOURCE],
+                text=tc.TEST_QUERY_TEXT,
+                rag_retrieval_config=tc.TEST_RAG_RETRIEVAL_ERROR_CONFIG,
+            )
+            e.match(
+                "Only one of vector_distance_threshold or"
+                " vector_similarity_threshold can be specified at a time"
+                " in rag_retrieval_config."
+            )

tests/unit/vertex_rag/test_rag_retrieval_preview.py

@@ -96,6 +96,17 @@ def test_retrieval_query_rag_resources_config_success(self):
             response, test_rag_constants_preview.TEST_RETRIEVAL_RESPONSE
         )
 
+    @pytest.mark.usefixtures("retrieve_contexts_mock")
+    def test_retrieval_query_rag_resources_similarity_config_success(self):
+        response = rag.retrieval_query(
+            rag_resources=[test_rag_constants_preview.TEST_RAG_RESOURCE],
+            text=test_rag_constants_preview.TEST_QUERY_TEXT,
+            rag_retrieval_config=test_rag_constants_preview.TEST_RAG_RETRIEVAL_SIMILARITY_CONFIG,
+        )
+        retrieve_contexts_eq(
+            response, test_rag_constants_preview.TEST_RETRIEVAL_RESPONSE
+        )
+
     @pytest.mark.usefixtures("retrieve_contexts_mock")
     def test_retrieval_query_rag_resources_default_config_success(self):
         response = rag.retrieval_query(
@@ -223,3 +234,28 @@ def test_retrieval_query_multiple_rag_resources_config(self):
                 rag_retrieval_config=test_rag_constants_preview.TEST_RAG_RETRIEVAL_CONFIG,
             )
             e.match("Currently only support 1 RagResource")
+
+    def test_retrieval_query_multiple_rag_resources_similarity_config(self):
+        with pytest.raises(ValueError) as e:
+            rag.retrieval_query(
+                rag_resources=[
+                    test_rag_constants_preview.TEST_RAG_RESOURCE,
+                    test_rag_constants_preview.TEST_RAG_RESOURCE,
+                ],
+                text=test_rag_constants_preview.TEST_QUERY_TEXT,
+                rag_retrieval_config=test_rag_constants_preview.TEST_RAG_RETRIEVAL_SIMILARITY_CONFIG,
+            )
+            e.match("Currently only support 1 RagResource")
+
+    def test_retrieval_query_invalid_config_filter(self):
+        with pytest.raises(ValueError) as e:
+            rag.retrieval_query(
+                rag_resources=[test_rag_constants_preview.TEST_RAG_RESOURCE],
+                text=test_rag_constants_preview.TEST_QUERY_TEXT,
+                rag_retrieval_config=test_rag_constants_preview.TEST_RAG_RETRIEVAL_ERROR_CONFIG,
+            )
+            e.match(
+                "Only one of vector_distance_threshold or"
+                " vector_similarity_threshold can be specified at a time"
+                " in rag_retrieval_config."
+            )

tests/unit/vertex_rag/test_rag_store.py

@@ -28,7 +28,7 @@ def test_retrieval_tool_invalid_name(self):
                 retrieval=rag.Retrieval(
                     source=rag.VertexRagStore(
                         rag_resources=[tc.TEST_RAG_RESOURCE_INVALID_NAME],
-                        rag_retrieval_config=tc.TEST_RAG_RETRIEVAL_CONFIG,
+                        rag_retrieval_config=tc.TEST_RAG_RETRIEVAL_SIMILARITY_CONFIG,
                     ),
                 )
             )
@@ -45,3 +45,19 @@ def test_retrieval_tool_multiple_rag_resources(self):
                 )
             )
             e.match("Currently only support 1 RagResource")
+
+    def test_retrieval_tool_invalid_config_filter(self):
+        with pytest.raises(ValueError) as e:
+            Tool.from_retrieval(
+                retrieval=rag.Retrieval(
+                    source=rag.VertexRagStore(
+                        rag_resources=[tc.TEST_RAG_RESOURCE],
+                        rag_retrieval_config=tc.TEST_RAG_RETRIEVAL_ERROR_CONFIG,
+                    )
+                )
+            )
+            e.match(
+                "Only one of vector_distance_threshold or"
+                " vector_similarity_threshold can be specified at a time"
+                " in rag_retrieval_config."
+            )

tests/unit/vertex_rag/test_rag_store_preview.py

@@ -47,6 +47,19 @@ def test_retrieval_tool_config_success(self):
                 )
             )
 
+    def test_retrieval_tool_similarity_config_success(self):
+        with pytest.warns(DeprecationWarning):
+            Tool.from_retrieval(
+                retrieval=rag.Retrieval(
+                    source=rag.VertexRagStore(
+                        rag_corpora=[
+                            test_rag_constants_preview.TEST_RAG_CORPUS_ID,
+                        ],
+                        rag_retrieval_config=test_rag_constants_preview.TEST_RAG_RETRIEVAL_SIMILARITY_CONFIG,
+                    ),
+                )
+            )
+
     def test_retrieval_tool_invalid_name(self):
         with pytest.raises(ValueError) as e:
             Tool.from_retrieval(
@@ -137,3 +150,19 @@ def test_retrieval_tool_multiple_rag_resources_config(self):
                 )
             )
             e.match("Currently only support 1 RagResource")
+
+    def test_retrieval_tool_invalid_config_filter(self):
+        with pytest.raises(ValueError) as e:
+            Tool.from_retrieval(
+                retrieval=rag.Retrieval(
+                    source=rag.VertexRagStore(
+                        rag_resources=[test_rag_constants_preview.TEST_RAG_RESOURCE],
+                        rag_retrieval_config=test_rag_constants_preview.TEST_RAG_RETRIEVAL_ERROR_CONFIG,
+                    )
+                )
+            )
+            e.match(
+                "Only one of vector_distance_threshold or"
+                " vector_similarity_threshold can be specified at a time"
+                " in rag_retrieval_config."
+            )

vertexai/preview/rag/rag_retrieval.py

@@ -190,11 +190,12 @@ def retrieval_query(
     else:
         # If rag_retrieval_config is specified, check for missing parameters.
         api_retrival_config = aiplatform_v1beta1.RagRetrievalConfig()
-        api_retrival_config.top_k = (
-            rag_retrieval_config.top_k
-            if rag_retrieval_config.top_k
-            else similarity_top_k
-        )
+        # Set top_k to config value if specified
+        if rag_retrieval_config.top_k:
+            api_retrival_config.top_k = rag_retrieval_config.top_k
+        else:
+            api_retrival_config.top_k = similarity_top_k
+        # Set alpha to config value if specified
         if (
             rag_retrieval_config.hybrid_search
             and rag_retrieval_config.hybrid_search.alpha
@@ -204,6 +205,19 @@ def retrieval_query(
             )
         else:
             api_retrival_config.hybrid_search.alpha = vector_search_alpha
+        # Check if both vector_distance_threshold and vector_similarity_threshold
+        # are specified.
+        if (
+            rag_retrieval_config.filter
+            and rag_retrieval_config.filter.vector_distance_threshold
+            and rag_retrieval_config.filter.vector_similarity_threshold
+        ):
+            raise ValueError(
+                "Only one of vector_distance_threshold or"
+                " vector_similarity_threshold can be specified at a time"
+                " in rag_retrieval_config."
+            )
+        # Set vector_distance_threshold to config value if specified
         if (
             rag_retrieval_config.filter
             and rag_retrieval_config.filter.vector_distance_threshold
@@ -215,6 +229,15 @@ def retrieval_query(
             api_retrival_config.filter.vector_distance_threshold = (
                 vector_distance_threshold
             )
+        # Set vector_similarity_threshold to config value if specified
+        if (
+            rag_retrieval_config.filter
+            and rag_retrieval_config.filter.vector_similarity_threshold
+        ):
+            api_retrival_config.filter.vector_similarity_threshold = (
+                rag_retrieval_config.filter.vector_similarity_threshold
+            )
+
     query = aiplatform_v1beta1.RagQuery(
         text=text,
         rag_retrieval_config=api_retrival_config,

vertexai/preview/rag/rag_store.py

@@ -167,16 +167,42 @@ def __init__(
         else:
             # If rag_retrieval_config is specified, check for missing parameters.
             api_retrival_config = aiplatform_v1beta1.RagRetrievalConfig()
-            if not rag_retrieval_config.top_k:
+            # Set top_k to config value if specified
+            if rag_retrieval_config.top_k:
+                api_retrival_config.top_k = rag_retrieval_config.top_k
+            else:
                 api_retrival_config.top_k = similarity_top_k
+            # Check if both vector_distance_threshold and vector_similarity_threshold
+            # are specified.
             if (
-                not rag_retrieval_config.filter
-                or not rag_retrieval_config.filter.vector_distance_threshold
+                rag_retrieval_config.filter
+                and rag_retrieval_config.filter.vector_distance_threshold
+                and rag_retrieval_config.filter.vector_similarity_threshold
             ):
-                api_retrival_config.filter = (
-                    aiplatform_v1beta1.RagRetrievalConfig.Filter(
-                        vector_distance_threshold=vector_distance_threshold
-                    ),
+                raise ValueError(
+                    "Only one of vector_distance_threshold or"
+                    " vector_similarity_threshold can be specified at a time"
+                    " in rag_retrieval_config."
+                )
+            # Set vector_distance_threshold to config value if specified
+            if (
+                rag_retrieval_config.filter
+                and rag_retrieval_config.filter.vector_distance_threshold
+            ):
+                api_retrival_config.filter.vector_distance_threshold = (
+                    rag_retrieval_config.filter.vector_distance_threshold
+                )
+            else:
+                api_retrival_config.filter.vector_distance_threshold = (
+                    vector_distance_threshold
+                )
+            # Set vector_similarity_threshold to config value if specified
+            if (
+                rag_retrieval_config.filter
+                and rag_retrieval_config.filter.vector_similarity_threshold
+            ):
+                api_retrival_config.filter.vector_similarity_threshold = (
+                    rag_retrieval_config.filter.vector_similarity_threshold
                 )
 
         if rag_resources:

vertexai/rag/rag_retrieval.py

@@ -108,16 +108,27 @@ def retrieval_query(
         api_retrival_config = aiplatform_v1.RagRetrievalConfig()
     else:
         # If rag_retrieval_config is specified, check for missing parameters.
-        api_retrival_config = aiplatform_v1.RagRetrievalConfig(
-            top_k=rag_retrieval_config.top_k,
-        )
+        api_retrival_config = aiplatform_v1.RagRetrievalConfig()
+        api_retrival_config.top_k = rag_retrieval_config.top_k
+        # Set vector_distance_threshold to config value if specified
         if rag_retrieval_config.filter:
-            api_retrival_config.filter = aiplatform_v1.RagRetrievalConfig.Filter(
-                vector_distance_threshold=rag_retrieval_config.filter.vector_distance_threshold
+            # Check if both vector_distance_threshold and vector_similarity_threshold
+            # are specified.
+            if (
+                rag_retrieval_config.filter
+                and rag_retrieval_config.filter.vector_distance_threshold
+                and rag_retrieval_config.filter.vector_similarity_threshold
+            ):
+                raise ValueError(
+                    "Only one of vector_distance_threshold or"
+                    " vector_similarity_threshold can be specified at a time"
+                    " in rag_retrieval_config."
+                )
+            api_retrival_config.filter.vector_distance_threshold = (
+                rag_retrieval_config.filter.vector_distance_threshold
             )
-        else:
-            api_retrival_config.filter = aiplatform_v1.RagRetrievalConfig.Filter(
-                vector_distance_threshold=None
+            api_retrival_config.filter.vector_similarity_threshold = (
+                rag_retrieval_config.filter.vector_similarity_threshold
             )
 
     query = aiplatform_v1.RagQuery(

vertexai/rag/rag_store.py

@@ -103,20 +103,29 @@ def __init__(
             )
 
         # If rag_retrieval_config is not specified, set it to default values.
-        if not rag_retrieval_config:
-            api_retrival_config = aiplatform_v1.RagRetrievalConfig()
-        else:
-            # If rag_retrieval_config is specified, check for missing parameters.
-            api_retrival_config = aiplatform_v1.RagRetrievalConfig(
-                top_k=rag_retrieval_config.top_k,
-            )
+        api_retrival_config = aiplatform_v1.RagRetrievalConfig()
+        # If rag_retrieval_config is specified, populate the default config.
+        if rag_retrieval_config:
+            api_retrival_config.top_k = rag_retrieval_config.top_k
+            # Set vector_distance_threshold to config value if specified
             if rag_retrieval_config.filter:
-                api_retrival_config.filter = aiplatform_v1.RagRetrievalConfig.Filter(
-                    vector_distance_threshold=rag_retrieval_config.filter.vector_distance_threshold
+                # Check if both vector_distance_threshold and
+                # vector_similarity_threshold are specified.
+                if (
+                    rag_retrieval_config.filter
+                    and rag_retrieval_config.filter.vector_distance_threshold
+                    and rag_retrieval_config.filter.vector_similarity_threshold
+                ):
+                    raise ValueError(
+                        "Only one of vector_distance_threshold or"
+                        " vector_similarity_threshold can be specified at a time"
+                        " in rag_retrieval_config."
+                    )
+                api_retrival_config.filter.vector_distance_threshold = (
+                    rag_retrieval_config.filter.vector_distance_threshold
                 )
-            else:
-                api_retrival_config.filter = aiplatform_v1.RagRetrievalConfig.Filter(
-                    vector_distance_threshold=None
+                api_retrival_config.filter.vector_similarity_threshold = (
+                    rag_retrieval_config.filter.vector_similarity_threshold
                 )
 
         if rag_resources: