[Cosmos]: Adding new options for semantic reranking, and adding more test cases (#43275)

* Updating options for semantic reranking, and adding more test cases

* Resolving comments

Author: aayush3011

sdk/cosmos/azure-cosmos/azure/cosmos/_inference_service.py

@@ -22,7 +22,7 @@
 import json
 import os
 import urllib
-from typing import Any, cast, Dict, List, Optional
+from typing import Any, cast, Optional
 from urllib3.util.retry import Retry
 
 from azure.core import PipelineClient
@@ -142,24 +142,28 @@ def _create_inference_pipeline_client(self) -> PipelineClient:
     def rerank(
         self,
         reranking_context: str,
-        documents: List[str],
-        semantic_reranking_options: Optional[Dict[str, Any]] = None,
+        documents: list[str],
+        semantic_reranking_options: Optional[dict[str, Any]] = None,
     ) -> CosmosDict:
         """Rerank documents using the semantic reranking service.
 
-        :param reranking_context: Query / context string used to score documents.
-        :type reranking_context: str
-        :param documents: List of document strings to rerank.
-        :type documents: List[str]
-        :param semantic_reranking_options: Optional dictionary of tuning parameters. Supported keys:
-            * return_documents (bool): Include original document text in results. Default True.
-            * top_k (int): Limit number of scored documents returned.
-            * batch_size (int): Batch size for internal scoring operations.
-            * sort (bool): If True (default) results are ordered by descending score.
-        :type semantic_reranking_options: Optional[Dict[str, Any]]
-        :returns: Reranking result payload.
+        :param str reranking_context: The context or query string to use for reranking the documents.
+        :param list[str] documents: A list of documents (as strings) to be reranked.
+        :param dict[str, Any] semantic_reranking_options: Optional dictionary of additional options to customize the semantic reranking process.
+
+         Supported options:
+
+         * **return_documents** (bool): Whether to return the document text in the response. If False, only scores and indices are returned. Default is True.
+         * **top_k** (int): Maximum number of documents to return in the reranked results. If not specified, all documents are returned.
+         * **batch_size** (int): Number of documents to process in each batch. Used for optimizing performance with large document sets.
+         * **sort** (bool): Whether to sort the results by relevance score in descending order. Default is True.
+         * **document_type** (str): Type of documents being reranked. Supported values are "string" and "json".
+         * **target_paths** (str): If document_type is "json", the list of JSON paths to extract text from for reranking. Comma-separated string.
+
+        :type semantic_reranking_options: Optional[dict[str, Any]]
+        :returns: A CosmosDict containing the reranking results. The structure typically includes results list with reranked documents and their relevance scores. Each result contains index, relevance_score, and optionally document.
         :rtype: ~azure.cosmos.CosmosDict[str, Any]
-        :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: On HTTP or service error.
+        :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: If the semantic reranking operation fails.
         """
         try:
             body = {
@@ -168,14 +172,7 @@ def rerank(
             }
 
             if semantic_reranking_options:
-                if "return_documents" in semantic_reranking_options:
-                    body["return_documents"] = semantic_reranking_options["return_documents"]
-                if "top_k" in semantic_reranking_options:
-                    body["top_k"] = semantic_reranking_options["top_k"]
-                if "batch_size" in semantic_reranking_options:
-                    body["batch_size"] = semantic_reranking_options["batch_size"]
-                if "sort" in semantic_reranking_options:
-                    body["sort"] = semantic_reranking_options["sort"]
+                body.update(semantic_reranking_options)
 
             headers = {
                 HttpHeaders.ContentType: "application/json"

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py

@@ -1183,8 +1183,8 @@ async def upsert_item(
     async def semantic_rerank(
         self,
         reranking_context: str,
-        documents: List[str],
-        semantic_reranking_options: Optional[Dict[str, Any]] = None
+        documents: list[str],
+        semantic_reranking_options: Optional[dict[str, Any]] = None
     ) -> CosmosDict:
         """Rerank a list of documents using semantic reranking.
 
@@ -1193,7 +1193,7 @@ async def semantic_rerank(
 
         :param str reranking_context: The context or query string to use for reranking the documents.
         :param list[str] documents: A list of documents (as strings) to be reranked.
-        :param semantic_reranking_options: Optional dictionary of additional options to customize the semantic reranking process.
+        :param dict[str, Any] semantic_reranking_options: Optional dictionary of additional options to customize the semantic reranking process.
 
          Supported options:
 
@@ -1202,9 +1202,9 @@ async def semantic_rerank(
          * **batch_size** (int): Number of documents to process in each batch. Used for optimizing performance with large document sets.
          * **sort** (bool): Whether to sort the results by relevance score in descending order. Default is True.
          * **document_type** (str): Type of documents being reranked. Supported values are "string" and "json".
-         * **target_paths** (list[str]): If document_type is "json", the list of JSON paths to extract text from for reranking.
+         * **target_paths** (str): If document_type is "json", the list of JSON paths to extract text from for reranking. Comma-separated string.
 
-        :type semantic_reranking_options: Optional[Dict[str, Any]]
+        :type semantic_reranking_options: Optional[dict[str, Any]]
         :returns: A CosmosDict containing the reranking results. The structure typically includes results list with reranked documents and their relevance scores. Each result contains index, relevance_score, and optionally document.
         :rtype: ~azure.cosmos.CosmosDict[str, Any]
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: If the semantic reranking operation fails.

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_inference_service_async.py

@@ -22,7 +22,7 @@
 import json
 import os
 import urllib
-from typing import Any, cast, Dict, List, Optional
+from typing import Any, cast, Optional
 from urllib.parse import urlparse
 from urllib3.util.retry import Retry
 
@@ -171,41 +171,37 @@ def _get_ssl_verification_setting(self) -> bool:
     async def rerank(
         self,
         reranking_context: str,
-        documents: List[str],
-        semantic_reranking_options: Optional[Dict[str, Any]] = None,
+        documents: list[str],
+        semantic_reranking_options: Optional[dict[str, Any]] = None,
     ) -> CosmosDict:
         """Rerank documents using the semantic reranking service (async).
 
-        :param reranking_context: Query / context string used to score documents.
-        :type reranking_context: str
-        :param documents: List of document strings to rerank.
-        :type documents: List[str]
-        :param semantic_reranking_options: Optional dictionary of tuning parameters. Supported keys:
-            * return_documents (bool): Include original document text in results. Default True.
-            * top_k (int): Limit number of scored documents returned.
-            * batch_size (int): Batch size for internal scoring operations.
-            * sort (bool): If True (default) results are ordered by descending score.
-        :type semantic_reranking_options: Optional[Dict[str, Any]]
-        :returns: Reranking result payload.
+        :param str reranking_context: The context or query string to use for reranking the documents.
+        :param list[str] documents: A list of documents (as strings) to be reranked.
+        :param dict[str, Any] semantic_reranking_options: Optional dictionary of additional options to customize the semantic reranking process.
+
+         Supported options:
+
+         * **return_documents** (bool): Whether to return the document text in the response. If False, only scores and indices are returned. Default is True.
+         * **top_k** (int): Maximum number of documents to return in the reranked results. If not specified, all documents are returned.
+         * **batch_size** (int): Number of documents to process in each batch. Used for optimizing performance with large document sets.
+         * **sort** (bool): Whether to sort the results by relevance score in descending order. Default is True.
+         * **document_type** (str): Type of documents being reranked. Supported values are "string" and "json".
+         * **target_paths** (str): If document_type is "json", the list of JSON paths to extract text from for reranking. Comma-separated string.
+
+        :type semantic_reranking_options: Optional[dict[str, Any]]
+        :returns: A CosmosDict containing the reranking results. The structure typically includes results list with reranked documents and their relevance scores. Each result contains index, relevance_score, and optionally document.
         :rtype: ~azure.cosmos.CosmosDict[str, Any]
-        :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: On HTTP or service error.
+        :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: If the semantic reranking operation fails.
         """
         try:
             body = {
                 "query": reranking_context,
                 "documents": documents,
             }
 
-            # Add optional parameters if provided
             if semantic_reranking_options:
-                if "return_documents" in semantic_reranking_options:
-                    body["return_documents"] = semantic_reranking_options["return_documents"]
-                if "top_k" in semantic_reranking_options:
-                    body["top_k"] = semantic_reranking_options["top_k"]
-                if "batch_size" in semantic_reranking_options:
-                    body["batch_size"] = semantic_reranking_options["batch_size"]
-                if "sort" in semantic_reranking_options:
-                    body["sort"] = semantic_reranking_options["sort"]
+                body.update(semantic_reranking_options)
 
             headers = {
                 HttpHeaders.ContentType: "application/json"

sdk/cosmos/azure-cosmos/azure/cosmos/container.py

@@ -1058,8 +1058,8 @@ def query_items(  # pylint:disable=docstring-missing-param
     def semantic_rerank(
         self,
         reranking_context: str,
-        documents: List[str],
-        semantic_reranking_options: Optional[Dict[str, Any]] = None
+        documents: list[str],
+        semantic_reranking_options: Optional[dict[str, Any]] = None
     ) -> CosmosDict:
         """Rerank a list of documents using semantic reranking.
 
@@ -1068,7 +1068,7 @@ def semantic_rerank(
 
         :param str reranking_context: The context or query string to use for reranking the documents.
         :param list[str] documents: A list of documents (as strings) to be reranked.
-        :param semantic_reranking_options: Optional dictionary of additional options to customize the semantic reranking process.
+        :param dict[str, Any] semantic_reranking_options: Optional dictionary of additional options to customize the semantic reranking process.
 
          Supported options:
 
@@ -1077,9 +1077,9 @@ def semantic_rerank(
          * **batch_size** (int): Number of documents to process in each batch. Used for optimizing performance with large document sets.
          * **sort** (bool): Whether to sort the results by relevance score in descending order. Default is True.
          * **document_type** (str): Type of documents being reranked. Supported values are "string" and "json".
-         * **target_paths** (list[str]): If document_type is "json", the list of JSON paths to extract text from for reranking.
+         * **target_paths** (str): If document_type is "json", the list of JSON paths to extract text from for reranking. Comma-separated string.
 
-        :type semantic_reranking_options: Optional[Dict[str, Any]]
+        :type semantic_reranking_options: Optional[dict[str, Any]]
         :returns: A CosmosDict containing the reranking results. The structure typically includes results list with reranked documents and their relevance scores. Each result contains index, relevance_score, and optionally document.
         :rtype: ~azure.cosmos.CosmosDict[str, Any]
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: If the semantic reranking operation fails.

sdk/cosmos/azure-cosmos/tests/test_semantic_reranker.py

@@ -1,6 +1,7 @@
 # The MIT License (MIT)
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # cspell:ignore rerank reranker reranking
+import json
 import unittest
 
 import azure.cosmos.cosmos_client as cosmos_client
@@ -44,7 +45,7 @@ def tearDownClass(cls):
             pass
 
     def test_semantic_reranker(self):
-        documents = self._get_documents()
+        documents = self._get_documents(document_type="string")
         results = self.test_container.semantic_rerank(
             reranking_context="What is the capital of France?",
             documents=documents,
@@ -59,9 +60,62 @@ def test_semantic_reranker(self):
         assert len(results["Scores"]) == len(documents)
         assert results["Scores"][0]["document"] == "Paris is the capital of France."
 
-    def _get_documents(self):
-        return [
-            "Berlin is the capital of Germany.",
-            "Paris is the capital of France.",
-            "Madrid is the capital of Spain."
-        ]
+    def test_semantic_reranker_json_documents(self):
+        documents = self._get_documents(document_type="json")
+        results = self.test_container.semantic_rerank(
+            reranking_context="What is the capital of France?",
+            documents=[json.dumps(item) for item in documents],
+            semantic_reranking_options={
+                "return_documents": True,
+                "top_k": 10,
+                "batch_size": 32,
+                "sort": True,
+                "document_type": "json",
+                "target_paths": "text",
+            }
+        )
+
+        assert len(results["Scores"]) == len(documents)
+        returned_document = json.loads(results["Scores"][0]["document"])
+        assert returned_document["text"] == "Paris is the capital of France."
+
+    def test_semantic_reranker_nested_json_documents(self):
+        documents = self._get_documents(document_type="nested_json")
+        results = self.test_container.semantic_rerank(
+            reranking_context="What is the capital of France?",
+            documents=[json.dumps(item) for item in documents],
+            semantic_reranking_options={
+                "return_documents": True,
+                "top_k": 10,
+                "batch_size": 32,
+                "sort": True,
+                "document_type": "json",
+                "target_paths": "info.text",
+            }
+        )
+
+        assert len(results["Scores"]) == len(documents)
+        returned_document = json.loads(results["Scores"][0]["document"])
+        assert returned_document["info"]["text"] == "Paris is the capital of France."
+
+    def _get_documents(self, document_type: str):
+        if document_type == "string":
+            return [
+                "Berlin is the capital of Germany.",
+                "Paris is the capital of France.",
+                "Madrid is the capital of Spain."
+            ]
+        elif document_type == "json":
+            return [
+                {"id": "1", "text": "Berlin is the capital of Germany."},
+                {"id": "2", "text": "Paris is the capital of France."},
+                {"id": "3", "text": "Madrid is the capital of Spain."}
+            ]
+        elif document_type == "nested_json":
+            return [
+                {"id": "1", "info": {"text": "Berlin is the capital of Germany."}},
+                {"id": "2", "info": {"text": "Paris is the capital of France."}},
+                {"id": "3", "info": {"text": "Madrid is the capital of Spain."}}
+            ]
+        else:
+            raise ValueError("Unsupported document type")