feat: Modify v1 sdk to support rerankers

PiperOrigin-RevId: 738100302

Author: vertex-sdk-bot

tests/unit/vertex_rag/test_rag_constants.py

@@ -21,12 +21,15 @@
 from vertexai.rag import (
     Filter,
     LayoutParserConfig,
+    LlmRanker,
     Pinecone,
     RagCorpus,
     RagFile,
     RagResource,
     RagRetrievalConfig,
     RagVectorDbConfig,
+    Ranking,
+    RankService,
     SharePointSource,
     SharePointSources,
     SlackChannelsSource,
@@ -560,3 +563,26 @@
     top_k=2,
     filter=Filter(vector_distance_threshold=0.5, vector_similarity_threshold=0.5),
 )
+TEST_RAG_RETRIEVAL_CONFIG_RANK_SERVICE = RagRetrievalConfig(
+    top_k=2,
+    filter=Filter(vector_distance_threshold=0.5),
+    ranking=Ranking(rank_service=RankService(model_name="test-model-name")),
+)
+TEST_RAG_RETRIEVAL_CONFIG_LLM_RANKER = RagRetrievalConfig(
+    top_k=2,
+    filter=Filter(vector_distance_threshold=0.5),
+    ranking=Ranking(llm_ranker=LlmRanker(model_name="test-model-name")),
+)
+TEST_RAG_RETRIEVAL_RANKING_CONFIG = RagRetrievalConfig(
+    top_k=2,
+    filter=Filter(vector_distance_threshold=0.5),
+    ranking=Ranking(rank_service=RankService(model_name="test-rank-service")),
+)
+TEST_RAG_RETRIEVAL_ERROR_RANKING_CONFIG = RagRetrievalConfig(
+    top_k=2,
+    filter=Filter(vector_distance_threshold=0.5),
+    ranking=Ranking(
+        rank_service=RankService(model_name="test-rank-service"),
+        llm_ranker=LlmRanker(model_name="test-llm-ranker"),
+    ),
+)

tests/unit/vertex_rag/test_rag_retrieval.py

@@ -87,6 +87,24 @@ def test_retrieval_query_rag_resources_similarity_success(self):
         )
         retrieve_contexts_eq(response, tc.TEST_RETRIEVAL_RESPONSE)
 
+    @pytest.mark.usefixtures("retrieve_contexts_mock")
+    def test_retrieval_query_rag_corpora_config_rank_service_success(self):
+        response = rag.retrieval_query(
+            rag_resources=[tc.TEST_RAG_RESOURCE],
+            text=tc.TEST_QUERY_TEXT,
+            rag_retrieval_config=tc.TEST_RAG_RETRIEVAL_CONFIG_RANK_SERVICE,
+        )
+        retrieve_contexts_eq(response, tc.TEST_RETRIEVAL_RESPONSE)
+
+    @pytest.mark.usefixtures("retrieve_contexts_mock")
+    def test_retrieval_query_rag_corpora_config_llm_ranker_success(self):
+        response = rag.retrieval_query(
+            rag_resources=[tc.TEST_RAG_RESOURCE],
+            text=tc.TEST_QUERY_TEXT,
+            rag_retrieval_config=tc.TEST_RAG_RETRIEVAL_CONFIG_LLM_RANKER,
+        )
+        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:

tests/unit/vertex_rag/test_rag_store.py

@@ -55,6 +55,17 @@ def test_retrieval_tool_no_rag_resources(self):
             )
             e.match("rag_resources must be specified.")
 
+    def test_retrieval_tool_ranking_config_success(self):
+        tool = Tool.from_retrieval(
+            retrieval=rag.Retrieval(
+                source=rag.VertexRagStore(
+                    rag_resources=[tc.TEST_RAG_RESOURCE],
+                    rag_retrieval_config=tc.TEST_RAG_RETRIEVAL_RANKING_CONFIG,
+                ),
+            )
+        )
+        assert tool is not None
+
     def test_retrieval_tool_invalid_name(self):
         with pytest.raises(ValueError) as e:
             Tool.from_retrieval(
@@ -94,3 +105,18 @@ def test_retrieval_tool_invalid_config_filter(self):
                 " vector_similarity_threshold can be specified at a time"
                 " in rag_retrieval_config."
             )
+
+    def test_retrieval_tool_invalid_ranking_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_RANKING_CONFIG,
+                    )
+                )
+            )
+            e.match(
+                "Only one of rank_service or llm_ranker can be specified"
+                " at a time in rag_retrieval_config."
+            )

vertexai/rag/__init__.py

@@ -43,6 +43,7 @@
     JiraQuery,
     JiraSource,
     LayoutParserConfig,
+    LlmRanker,
     Pinecone,
     RagCorpus,
     RagEmbeddingModelConfig,
@@ -51,6 +52,8 @@
     RagResource,
     RagRetrievalConfig,
     RagVectorDbConfig,
+    Ranking,
+    RankService,
     SharePointSource,
     SharePointSources,
     SlackChannel,
@@ -67,6 +70,7 @@
     "JiraQuery",
     "JiraSource",
     "LayoutParserConfig",
+    "LlmRanker",
     "Pinecone",
     "RagCorpus",
     "RagEmbeddingModelConfig",
@@ -75,6 +79,8 @@
     "RagResource",
     "RagRetrievalConfig",
     "RagVectorDbConfig",
+    "Ranking",
+    "RankService",
     "Retrieval",
     "SharePointSource",
     "SharePointSources",

vertexai/rag/rag_retrieval.py

@@ -43,6 +43,11 @@ def retrieval_query(
         filter=vertexai.rag.rag_retrieval_config.filter(
             vector_distance_threshold=0.5
         ),
+        ranking=vertex.rag.Ranking(
+            llm_ranker=vertexai.rag.LlmRanker(
+                model_name="gemini-1.5-flash-002"
+            )
+        )
     )
 
     results = vertexai.rag.retrieval_query(
@@ -105,11 +110,11 @@ def retrieval_query(
 
     # If rag_retrieval_config is not specified, set it to default values.
     if not rag_retrieval_config:
-        api_retrival_config = aiplatform_v1.RagRetrievalConfig()
+        api_retrieval_config = aiplatform_v1.RagRetrievalConfig()
     else:
         # If rag_retrieval_config is specified, check for missing parameters.
-        api_retrival_config = aiplatform_v1.RagRetrievalConfig()
-        api_retrival_config.top_k = rag_retrieval_config.top_k
+        api_retrieval_config = aiplatform_v1.RagRetrievalConfig()
+        api_retrieval_config.top_k = rag_retrieval_config.top_k
         # Set vector_distance_threshold to config value if specified
         if rag_retrieval_config.filter:
             # Check if both vector_distance_threshold and vector_similarity_threshold
@@ -124,16 +129,30 @@ def retrieval_query(
                     " vector_similarity_threshold can be specified at a time"
                     " in rag_retrieval_config."
                 )
-            api_retrival_config.filter.vector_distance_threshold = (
+            api_retrieval_config.filter.vector_distance_threshold = (
                 rag_retrieval_config.filter.vector_distance_threshold
             )
-            api_retrival_config.filter.vector_similarity_threshold = (
+            api_retrieval_config.filter.vector_similarity_threshold = (
                 rag_retrieval_config.filter.vector_similarity_threshold
             )
+        if (
+            rag_retrieval_config.ranking
+            and rag_retrieval_config.ranking.rank_service
+            and rag_retrieval_config.ranking.llm_ranker
+        ):
+            raise ValueError("Only one of rank_service and llm_ranker can be set.")
+        if rag_retrieval_config.ranking and rag_retrieval_config.ranking.rank_service:
+            api_retrieval_config.ranking.rank_service.model_name = (
+                rag_retrieval_config.ranking.rank_service.model_name
+            )
+        elif rag_retrieval_config.ranking and rag_retrieval_config.ranking.llm_ranker:
+            api_retrieval_config.ranking.llm_ranker.model_name = (
+                rag_retrieval_config.ranking.llm_ranker.model_name
+            )
 
     query = aiplatform_v1.RagQuery(
         text=text,
-        rag_retrieval_config=api_retrival_config,
+        rag_retrieval_config=api_retrieval_config,
     )
     request = aiplatform_v1.RetrieveContextsRequest(
         vertex_rag_store=vertex_rag_store,

vertexai/rag/rag_store.py

@@ -62,6 +62,11 @@ def __init__(
             filter=vertexai.rag.RagRetrievalConfig.Filter(
                 vector_distance_threshold=0.5
             ),
+            ranking=vertex.rag.Ranking(
+                llm_ranker=vertexai.rag.LlmRanker(
+                    model_name="gemini-1.5-flash-002"
+                )
+            )
         )
 
         tool = Tool.from_retrieval(
@@ -127,6 +132,31 @@ def __init__(
                 api_retrieval_config.filter.vector_similarity_threshold = (
                     rag_retrieval_config.filter.vector_similarity_threshold
                 )
+            # Check if both rank_service and llm_ranker are specified.
+            if (
+                rag_retrieval_config.ranking
+                and rag_retrieval_config.ranking.rank_service
+                and rag_retrieval_config.ranking.rank_service.model_name
+                and rag_retrieval_config.ranking.llm_ranker
+                and rag_retrieval_config.ranking.llm_ranker.model_name
+            ):
+                raise ValueError(
+                    "Only one of rank_service or llm_ranker can be specified"
+                    " at a time in rag_retrieval_config."
+                )
+            # Set rank_service to config value if specified
+            if (
+                rag_retrieval_config.ranking
+                and rag_retrieval_config.ranking.rank_service
+            ):
+                api_retrieval_config.ranking.rank_service.model_name = (
+                    rag_retrieval_config.ranking.rank_service.model_name
+                )
+            # Set llm_ranker to config value if specified
+            if rag_retrieval_config.ranking and rag_retrieval_config.ranking.llm_ranker:
+                api_retrieval_config.ranking.llm_ranker.model_name = (
+                    rag_retrieval_config.ranking.llm_ranker.model_name
+                )
 
         gapic_rag_resource = gapic_tool_types.VertexRagStore.RagResource(
             rag_corpus=rag_corpus_name,

vertexai/rag/utils/resources.py

@@ -332,17 +332,56 @@ class Filter:
     metadata_filter: Optional[str] = None
 
 
+@dataclasses.dataclass
+class LlmRanker:
+    """LlmRanker.
+
+    Attributes:
+        model_name: The model name used for ranking. Only Gemini models are
+            supported for now.
+    """
+
+    model_name: Optional[str] = None
+
+
+@dataclasses.dataclass
+class RankService:
+    """RankService.
+
+    Attributes:
+        model_name: The model name of the rank service. Format:
+            ``semantic-ranker-512@latest``
+    """
+
+    model_name: Optional[str] = None
+
+
+@dataclasses.dataclass
+class Ranking:
+    """Ranking.
+
+    Attributes:
+        rank_service: Config for Rank Service.
+        llm_ranker: Config for LlmRanker.
+    """
+
+    rank_service: Optional[RankService] = None
+    llm_ranker: Optional[LlmRanker] = None
+
+
 @dataclasses.dataclass
 class RagRetrievalConfig:
     """RagRetrievalConfig.
 
     Attributes:
         top_k: The number of contexts to retrieve.
         filter: Config for filters.
+        ranking: Config for ranking.
     """
 
     top_k: Optional[int] = None
     filter: Optional[Filter] = None
+    ranking: Optional[Ranking] = None
 
 
 @dataclasses.dataclass