feat: Support vector_distance_threshold filtering and file-based retr…

…ieval for RAG

PiperOrigin-RevId: 631564047

tests/unit/vertex_rag/test_rag_constants.py

@@ -18,6 +18,7 @@
 from vertexai.preview.rag.utils.resources import (
     RagCorpus,
     RagFile,
+    RagResource,
 )
 from google.cloud import aiplatform
 from google.cloud.aiplatform_v1beta1 import (
@@ -146,3 +147,11 @@
     ]
 )
 TEST_RETRIEVAL_RESPONSE = RetrieveContextsResponse(contexts=TEST_CONTEXTS)
+TEST_RAG_RESOURCE = RagResource(
+    rag_corpus=TEST_RAG_CORPUS_RESOURCE_NAME,
+    rag_file_ids=[TEST_RAG_FILE_ID],
+)
+TEST_RAG_RESOURCE_INVALID_NAME = RagResource(
+    rag_corpus="213lkj-1/23jkl/",
+    rag_file_ids=[TEST_RAG_FILE_ID],
+)

tests/unit/vertex_rag/test_rag_retrieval.py

@@ -70,30 +70,62 @@ def teardown_method(self):
         aiplatform.initializer.global_pool.shutdown(wait=True)
 
     @pytest.mark.usefixtures("retrieve_contexts_mock")
-    def test_retrieval_query_success(self):
+    def test_retrieval_query_rag_resources_success(self):
         response = rag.retrieval_query(
-            rag_corpora=[tc.TEST_RAG_CORPUS_RESOURCE_NAME],
+            rag_resources=[tc.TEST_RAG_RESOURCE],
             text=tc.TEST_QUERY_TEXT,
             similarity_top_k=2,
+            vector_distance_threshold=0.5,
+        )
+        retrieve_contexts_eq(response, tc.TEST_RETRIEVAL_RESPONSE)
+
+    @pytest.mark.usefixtures("retrieve_contexts_mock")
+    def test_retrieval_query_rag_corpora_success(self):
+        response = rag.retrieval_query(
+            rag_corpora=[tc.TEST_RAG_CORPUS_ID],
+            text=tc.TEST_QUERY_TEXT,
+            similarity_top_k=2,
+            vector_distance_threshold=0.5,
         )
         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:
             rag.retrieval_query(
-                rag_corpora=[tc.TEST_RAG_CORPUS_RESOURCE_NAME],
+                rag_resources=[tc.TEST_RAG_RESOURCE],
                 text=tc.TEST_QUERY_TEXT,
                 similarity_top_k=2,
+                vector_distance_threshold=0.5,
             )
             e.match("Failed in retrieving contexts due to")
 
     def test_retrieval_query_invalid_name(self):
         with pytest.raises(ValueError) as e:
             rag.retrieval_query(
-                # Should be RAG_CORPUS, not RAG_FILE
-                rag_corpora=[tc.TEST_RAG_FILE_RESOURCE_NAME],
+                rag_resources=[tc.TEST_RAG_RESOURCE_INVALID_NAME],
                 text=tc.TEST_QUERY_TEXT,
                 similarity_top_k=2,
+                vector_distance_threshold=0.5,
             )
             e.match("Invalid RagCorpus name")
+
+    def test_retrieval_query_multiple_rag_corpora(self):
+        with pytest.raises(ValueError) as e:
+            rag.retrieval_query(
+                rag_corpora=[tc.TEST_RAG_CORPUS_ID, tc.TEST_RAG_CORPUS_ID],
+                text=tc.TEST_QUERY_TEXT,
+                similarity_top_k=2,
+                vector_distance_threshold=0.5,
+            )
+            e.match("Currently only support 1 RagCorpus")
+
+    def test_retrieval_query_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,
+                similarity_top_k=2,
+                vector_distance_threshold=0.5,
+            )
+            e.match("Currently only support 1 RagResource")

tests/unit/vertex_rag/test_rag_store.py

@@ -0,0 +1,62 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from vertexai.preview import rag
+from vertexai.preview.generative_models import Tool
+import pytest
+import test_rag_constants as tc
+
+
+@pytest.mark.usefixtures("google_auth_mock")
+class TestRagStoreValidations:
+    def test_retrieval_tool_invalid_name(self):
+        with pytest.raises(ValueError) as e:
+            Tool.from_retrieval(
+                retrieval=rag.Retrieval(
+                    source=rag.VertexRagStore(
+                        rag_resources=[tc.TEST_RAG_RESOURCE_INVALID_NAME],
+                        similarity_top_k=3,
+                        vector_distance_threshold=0.4,
+                    ),
+                )
+            )
+            e.match("Invalid RagCorpus name")
+
+    def test_retrieval_tool_multiple_rag_corpora(self):
+        with pytest.raises(ValueError) as e:
+            Tool.from_retrieval(
+                retrieval=rag.Retrieval(
+                    source=rag.VertexRagStore(
+                        rag_corpora=[tc.TEST_RAG_CORPUS_ID, tc.TEST_RAG_CORPUS_ID],
+                        similarity_top_k=3,
+                        vector_distance_threshold=0.4,
+                    ),
+                )
+            )
+            e.match("Currently only support 1 RagCorpus")
+
+    def test_retrieval_tool_multiple_rag_resources(self):
+        with pytest.raises(ValueError) as e:
+            Tool.from_retrieval(
+                retrieval=rag.Retrieval(
+                    source=rag.VertexRagStore(
+                        rag_resources=[tc.TEST_RAG_RESOURCE, tc.TEST_RAG_RESOURCE],
+                        similarity_top_k=3,
+                        vector_distance_threshold=0.4,
+                    ),
+                )
+            )
+            e.match("Currently only support 1 RagResource")

tests/unit/vertex_ray/test_bigquery.py

@@ -311,7 +311,7 @@ def test_do_write_dataset_exists(self, ray_remote_function_mock):
         assert len(write_tasks_list) == 4
 
     # Ray 2.9.3 only
-    def test_write(self, ray_get_mock):
+    def test_write(self, ray_get_mock, ray_remote_function_mock):
         if _BigQueryDatasink is None:
             return
         bq_datasink = _BigQueryDatasink(
@@ -327,7 +327,7 @@ def test_write(self, ray_get_mock):
         assert status == "ok"
 
     # Ray 2.9.3 only
-    def test_write_dataset_exists(self, ray_get_mock):
+    def test_write_dataset_exists(self, ray_get_mock, ray_remote_function_mock):
         if _BigQueryDatasink is None:
             return
         bq_datasink = _BigQueryDatasink(

tests/unit/vertexai/test_generative_models.py

@@ -908,13 +908,18 @@ def test_generate_content_grounding_vertex_ai_search_retriever(self):
     )
     def test_generate_content_vertex_rag_retriever(self):
         model = preview_generative_models.GenerativeModel("gemini-pro")
+        rag_resources = [
+            rag.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_corpora=[
-                        f"projects/{_TEST_PROJECT}/locations/us-central1/ragCorpora/1234556"
-                    ],
+                    rag_resources=rag_resources,
                     similarity_top_k=1,
+                    vector_distance_threshold=0.4,
                 ),
             ),
         )

vertexai/preview/rag/__init__.py

@@ -36,6 +36,10 @@
     Retrieval,
     VertexRagStore,
 )
+from vertexai.preview.rag.utils.resources import (
+    RagResource,
+)
+
 
 __all__ = (
     "create_corpus",
@@ -51,4 +55,5 @@
     "retrieval_query",
     "Retrieval",
     "VertexRagStore",
+    "RagResource",
 )

vertexai/preview/rag/rag_data.py

@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+"""RAG data management SDK."""
 
 from typing import Optional, Union, Sequence
 from google import auth

vertexai/preview/rag/rag_retrieval.py

@@ -14,6 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+"""Retrieval query to get relevant contexts."""
+
+import re
 from typing import List, Optional
 from google.cloud.aiplatform import initializer
 
@@ -25,41 +28,89 @@
 from vertexai.preview.rag.utils import (
     _gapic_utils,
 )
+from vertexai.preview.rag.utils.resources import RagResource
 
 
 def retrieval_query(
-    rag_corpora: List[str],
     text: str,
+    rag_resources: Optional[List[RagResource]] = None,
+    rag_corpora: Optional[List[str]] = None,
     similarity_top_k: Optional[int] = 10,
+    vector_distance_threshold: Optional[float] = 0.3,
 ) -> RetrieveContextsResponse:
     """Retrieve top k relevant docs/chunks.
 
+    Example usage:
+    ```
+    import vertexai
+
+    vertexai.init(project="my-project")
+
+    results = vertexai.preview.rag.retrieval_query(
+        text="Why is the sky blue?",
+        rag_resources=[vertexai.preview.rag.RagResource(
+            rag_corpus="projects/my-project/locations/us-central1/ragCorpora/rag-corpus-1",
+            rag_file_ids=["rag-file-1", "rag-file-2", ...],
+        )],
+        similarity_top_k=2,
+        vector_distance_threshold=0.5,
+    )
+    ```
+
     Args:
-        rag_corpora: A list of full resource name or corpus_id of the RagCorpus. Format:
-        ``projects/{project}/locations/{location}/ragCorpora/{rag_corpus_id}``
         text: The query in text format to get relevant contexts.
+        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.
         similarity_top_k: The number of contexts to retrieve.
+        vector_distance_threshold: Optional. Only return contexts with vector
+            distance smaller than the threshold.
+
     Returns:
         RetrieveContextsResonse.
     """
     parent = initializer.global_config.common_location_path()
 
     client = _gapic_utils.create_rag_service_client()
-    vertex_rag_store = RetrieveContextsRequest.VertexRagStore()
-    # Currently only support 1 RagCorpus.
-    if len(rag_corpora) > 1:
-        raise ValueError("Currently only support 1 RagCorpus.")
-    if len(rag_corpora[0].split("/")) == 6:
-        rag_corpus_name = rag_corpora[0]
-    elif len(rag_corpora[0].split("/")) == 1:
-        rag_corpus_name = parent + "/ragCorpora/" + rag_corpora[0]
+
+    if rag_resources:
+        if len(rag_resources) > 1:
+            raise ValueError("Currently only support 1 RagResource.")
+        name = rag_resources[0].rag_corpus
+    elif rag_corpora:
+        if len(rag_corpora) > 1:
+            raise ValueError("Currently only support 1 RagCorpus.")
+        name = rag_corpora[0]
+    else:
+        raise ValueError("rag_resources or rag_corpora must be specified.")
+
+    data_client = _gapic_utils.create_rag_data_service_client()
+    if data_client.parse_rag_corpus_path(name):
+        rag_corpus_name = name
+    elif re.match("^{}$".format(_gapic_utils._VALID_RESOURCE_NAME_REGEX), name):
+        rag_corpus_name = parent + "/ragCorpora/" + name
     else:
         raise ValueError(
             "Invalid RagCorpus name: %s. Proper format should be: projects/{project}/locations/{location}/ragCorpora/{rag_corpus_id}",
             rag_corpora,
         )
 
-    vertex_rag_store.rag_corpora = [rag_corpus_name]
+    if rag_resources:
+        gapic_rag_resource = RetrieveContextsRequest.VertexRagStore.RagResource(
+            rag_corpus=rag_corpus_name,
+            rag_file_ids=rag_resources[0].rag_file_ids,
+        )
+        vertex_rag_store = RetrieveContextsRequest.VertexRagStore(
+            rag_resources=[gapic_rag_resource],
+        )
+    else:
+        vertex_rag_store = RetrieveContextsRequest.VertexRagStore(
+            rag_corpora=[rag_corpus_name],
+        )
+
+    vertex_rag_store.vector_distance_threshold = vector_distance_threshold
     query = RagQuery(text=text, similarity_top_k=similarity_top_k)
     request = RetrieveContextsRequest(
         vertex_rag_store=vertex_rag_store,