Merge e441d38d17 into c8635ecc67

litellm/llms/vertex_ai/rerank/transformation.py

@@ -20,6 +20,7 @@ from litellm.secret_managers.main import get_secret_str
 from litellm.types.rerank import (
     RerankBilledUnits,
     RerankResponse,
+    RerankResponseDocument,
     RerankResponseMeta,
     RerankResponseResult,
 )
@@ -172,6 +173,17 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
         except Exception as e:
             raise ValueError(f"Failed to parse response: {e}")
 
+        # Determine whether to return documents (defaults to True)
+        return_documents = True
+        if "return_documents" in optional_params and optional_params["return_documents"] is not None:
+            return_documents = bool(optional_params["return_documents"])
+        elif "return_documents" in request_data and request_data["return_documents"] is not None:
+            return_documents = bool(request_data["return_documents"])
+        elif "ignoreRecordDetailsInResponse" in request_data:
+            return_documents = not bool(request_data["ignoreRecordDetailsInResponse"])
+        elif "return_documents" in litellm_params and litellm_params["return_documents"] is not None:
+            return_documents = bool(litellm_params["return_documents"])
+
         # Extract records from response
         records: Final = raw_response_json.get("records", [])
 
@@ -179,23 +191,16 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
         results: Final = []
         for record in records:
             # Handle both cases: with full details and with only IDs
-            if "score" in record:
-                # Full response with score and details
-                results.append(
-                    {
-                        "index": int(record["id"]),
-                        "relevance_score": record.get("score", 0.0),
-                    }
-                )
-            else:
-                # Response with only IDs (when ignoreRecordDetailsInResponse=true)
-                # We can't provide a relevance score, so we'll use a default
-                results.append(
-                    {
-                        "index": int(record["id"]),
-                        "relevance_score": 1.0,  # Default score when details are ignored
-                    }
-                )
+            score_val = record.get("score", 0.0) if "score" in record else 1.0
+            doc_text = record.get("content")
+            result_item = {
+                "index": int(record["id"]),
+                "relevance_score": score_val,
+            }
+            if return_documents and doc_text is not None:
+                result_item["document"] = RerankResponseDocument(text=doc_text)
+
+            results.append(result_item)
 
         # Sort by relevance score (descending)
         results.sort(key=lambda x: x["relevance_score"], reverse=True)
@@ -204,9 +209,10 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
         # Convert results to proper RerankResponseResult objects
         rerank_results: Final = []
         for result in results:
-            rerank_results.append(
-                RerankResponseResult(index=result["index"], relevance_score=result["relevance_score"])
-            )
+            rerank_result = RerankResponseResult(index=result["index"], relevance_score=result["relevance_score"])
+            if "document" in result:
+                rerank_result["document"] = result["document"]
+            rerank_results.append(rerank_result)
 
         # Create meta object
         meta: Final = RerankResponseMeta(billed_units=RerankBilledUnits(search_units=len(records)))

tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py

@@ -23,9 +23,7 @@ class TestVertexAIRerankIntegration:
         importlib.reload(litellm) in conftest.py.
         """
         # Mock authentication at instance level
-        mock_ensure_access_token = MagicMock(
-            return_value=("test-access-token", "test-project-123")
-        )
+        mock_ensure_access_token = MagicMock(return_value=("test-access-token", "test-project-123"))
         self.config._ensure_access_token = mock_ensure_access_token
 
         # Test documents
@@ -39,9 +37,7 @@ class TestVertexAIRerankIntegration:
 
         # Step 1: Test request transformation
         # Validate environment
-        headers = self.config.validate_environment(
-            headers={}, model=self.model, api_key=None
-        )
+        headers = self.config.validate_environment(headers={}, model=self.model, api_key=None)
 
         # Transform request
         request_data = self.config.transform_rerank_request(
@@ -113,8 +109,15 @@ class TestVertexAIRerankIntegration:
         # Results should be sorted by relevance score (descending)
         assert result.results[0]["index"] == 3  # Highest score
         assert result.results[0]["relevance_score"] == 0.95
+        assert (
+            result.results[0]["document"]["text"]
+            == "Google's Gemini AI model represents a significant advancement in artificial intelligence technology."
+        )
         assert result.results[1]["index"] == 0  # Second highest score
         assert result.results[1]["relevance_score"] == 0.92
+        assert (
+            result.results[1]["document"]["text"] == "Gemini is a cutting edge large language model created by Google."
+        )
 
         # Verify metadata
         assert result.meta["billed_units"]["search_units"] == 2
@@ -157,15 +160,15 @@ class TestVertexAIRerankIntegration:
             raw_response=mock_response,
             model_response=model_response,
             logging_obj=mock_logging,
+            request_data=request_data,
         )
 
         # Verify response structure with default scores
         assert len(result.results) == 3
         for result_item in result.results:
-            assert (
-                result_item["relevance_score"] == 1.0
-            )  # Default score when details are ignored
+            assert result_item["relevance_score"] == 1.0  # Default score when details are ignored
             assert "index" in result_item
+            assert "document" not in result_item
 
     def test_document_title_generation(self):
         """Test that document titles are generated correctly from content."""
@@ -184,9 +187,7 @@ class TestVertexAIRerankIntegration:
         # Verify title generation
         assert request_data["records"][0]["title"] == "This is a"  # First 3 words
         assert request_data["records"][1]["title"] == "Short doc"  # Less than 3 words
-        assert (
-            request_data["records"][2]["title"] == "Another document with"
-        )  # First 3 words
+        assert request_data["records"][2]["title"] == "Another document with"  # First 3 words
 
     def test_dictionary_document_handling(self):
         """Test handling of dictionary-format documents."""
@@ -195,9 +196,7 @@ class TestVertexAIRerankIntegration:
                 "text": "Gemini is a cutting edge large language model created by Google.",
                 "title": "Custom Title 1",
             },
-            {
-                "text": "The Gemini zodiac symbol often depicts two figures standing side-by-side."
-            },
+            {"text": "The Gemini zodiac symbol often depicts two figures standing side-by-side."},
             {
                 "text": "Gemini is a constellation that can be seen in the night sky.",
                 "title": "Custom Title 3",
@@ -212,21 +211,15 @@ class TestVertexAIRerankIntegration:
 
         # Verify custom titles are used when provided
         assert request_data["records"][0]["title"] == "Custom Title 1"
-        assert (
-            request_data["records"][1]["title"] == "The Gemini zodiac"
-        )  # Generated from first 3 words
+        assert request_data["records"][1]["title"] == "The Gemini zodiac"  # Generated from first 3 words
         assert request_data["records"][2]["title"] == "Custom Title 3"
 
         # Verify content is extracted correctly
         assert (
-            request_data["records"][0]["content"]
-            == "Gemini is a cutting edge large language model created by Google."
+            request_data["records"][0]["content"] == "Gemini is a cutting edge large language model created by Google."
         )
         assert (
             request_data["records"][1]["content"]
             == "The Gemini zodiac symbol often depicts two figures standing side-by-side."
         )
-        assert (
-            request_data["records"][2]["content"]
-            == "Gemini is a constellation that can be seen in the night sky."
-        )
+        assert request_data["records"][2]["content"] == "Gemini is a constellation that can be seen in the night sky."

tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py

@@ -92,13 +92,9 @@ class TestVertexAIRerankTransform:
             litellm.vertex_project = None
             # Reset mock and set it to raise an error
             mock_ensure_access_token.reset_mock()
-            mock_ensure_access_token.side_effect = ValueError(
-                "Vertex AI project ID is required"
-            )
+            mock_ensure_access_token.side_effect = ValueError("Vertex AI project ID is required")
             try:
-                with pytest.raises(
-                    ValueError, match="Vertex AI project ID is required"
-                ):
+                with pytest.raises(ValueError, match="Vertex AI project ID is required"):
                     self.config.get_complete_url(api_base=None, model=self.model)
             finally:
                 litellm.vertex_project = original_project
@@ -111,14 +107,10 @@ class TestVertexAIRerankTransform:
         importlib.reload(litellm) in conftest.py.
         """
         # Mock the authentication at instance level
-        mock_ensure_access_token = MagicMock(
-            return_value=("test-access-token", "test-project-123")
-        )
+        mock_ensure_access_token = MagicMock(return_value=("test-access-token", "test-project-123"))
         self.config._ensure_access_token = mock_ensure_access_token
 
-        headers = self.config.validate_environment(
-            headers={}, model=self.model, api_key=None
-        )
+        headers = self.config.validate_environment(headers={}, model=self.model, api_key=None)
 
         expected_headers = {
             "Authorization": "Bearer test-access-token",
@@ -166,9 +158,7 @@ class TestVertexAIRerankTransform:
                     "text": "Gemini is a cutting edge large language model created by Google.",
                     "title": "Custom Title 1",
                 },
-                {
-                    "text": "The Gemini zodiac symbol often depicts two figures standing side-by-side."
-                },
+                {"text": "The Gemini zodiac symbol often depicts two figures standing side-by-side."},
             ],
         }
 
@@ -178,9 +168,7 @@ class TestVertexAIRerankTransform:
 
         # Verify record structure with custom titles
         assert request_data["records"][0]["title"] == "Custom Title 1"
-        assert (
-            request_data["records"][1]["title"] == "The Gemini zodiac"
-        )  # First 3 words
+        assert request_data["records"][1]["title"] == "The Gemini zodiac"  # First 3 words
 
     def test_transform_rerank_request_return_documents_mapping(self):
         """Test return_documents to ignoreRecordDetailsInResponse mapping."""
@@ -226,9 +214,7 @@ class TestVertexAIRerankTransform:
             model=self.model,
             optional_rerank_params=optional_params,
             headers={},
-            litellm_params={
-                "metadata": {"requester_metadata": {"app": "litellm", "tier": "1"}}
-            },
+            litellm_params={"metadata": {"requester_metadata": {"app": "litellm", "tier": "1"}}},
         )
         assert request_data["userLabels"] == {"app": "litellm", "tier": "1"}
 
@@ -243,9 +229,7 @@ class TestVertexAIRerankTransform:
             )
 
         # Test missing documents
-        with pytest.raises(
-            ValueError, match="documents is required for Vertex AI rerank"
-        ):
+        with pytest.raises(ValueError, match="documents is required for Vertex AI rerank"):
             self.config.transform_rerank_request(
                 model=self.model,
                 optional_rerank_params={"query": "test query"},
@@ -294,12 +278,153 @@ class TestVertexAIRerankTransform:
         assert len(result.results) == 2
         assert result.results[0]["index"] == 1  # Converted back to 0-based index
         assert result.results[0]["relevance_score"] == 0.98
+        assert (
+            result.results[0]["document"]["text"]
+            == "The sky appears blue due to a phenomenon called Rayleigh scattering."
+        )
         assert result.results[1]["index"] == 0
         assert result.results[1]["relevance_score"] == 0.64
+        assert (
+            result.results[1]["document"]["text"]
+            == "A canvas stretched across the day, Where sunlight learns to dance and play."
+        )
 
         # Verify metadata
         assert result.meta["billed_units"]["search_units"] == 2
 
+    def test_transform_rerank_response_return_documents_true_populates_document_text(self):
+        """Test that return_documents=True populates document with {'text': record['content']}."""
+        response_data = {
+            "records": [
+                {
+                    "id": "1",
+                    "score": 0.95,
+                    "title": "Doc 1",
+                    "content": "Content of document 1",
+                },
+                {
+                    "id": "0",
+                    "score": 0.80,
+                    "title": "Doc 0",
+                    "content": "Content of document 0",
+                },
+            ]
+        }
+
+        mock_response = MagicMock(spec=httpx.Response)
+        mock_response.json.return_value = response_data
+        mock_response.text = json.dumps(response_data)
+        mock_logging = MagicMock()
+        model_response = RerankResponse()
+
+        # Test with optional_params={"return_documents": True}
+        result = self.config.transform_rerank_response(
+            model=self.model,
+            raw_response=mock_response,
+            model_response=model_response,
+            logging_obj=mock_logging,
+            optional_params={"return_documents": True},
+        )
+
+        assert len(result.results) == 2
+        assert result.results[0]["index"] == 1
+        assert result.results[0]["relevance_score"] == 0.95
+        assert result.results[0]["document"] == {"text": "Content of document 1"}
+        assert result.results[0]["document"]["text"] == "Content of document 1"
+
+        assert result.results[1]["index"] == 0
+        assert result.results[1]["relevance_score"] == 0.80
+        assert result.results[1]["document"] == {"text": "Content of document 0"}
+        assert result.results[1]["document"]["text"] == "Content of document 0"
+
+    def test_transform_rerank_response_return_documents_false_omits_document_text(self):
+        """Test that return_documents=False does not populate document field."""
+        response_data = {
+            "records": [
+                {
+                    "id": "1",
+                    "score": 0.95,
+                    "title": "Doc 1",
+                    "content": "Content of document 1",
+                },
+                {
+                    "id": "0",
+                    "score": 0.80,
+                    "title": "Doc 0",
+                    "content": "Content of document 0",
+                },
+            ]
+        }
+
+        mock_response = MagicMock(spec=httpx.Response)
+        mock_response.json.return_value = response_data
+        mock_response.text = json.dumps(response_data)
+        mock_logging = MagicMock()
+        model_response = RerankResponse()
+
+        # Test with optional_params={"return_documents": False}
+        result = self.config.transform_rerank_response(
+            model=self.model,
+            raw_response=mock_response,
+            model_response=model_response,
+            logging_obj=mock_logging,
+            optional_params={"return_documents": False},
+        )
+
+        assert len(result.results) == 2
+        assert result.results[0]["index"] == 1
+        assert result.results[0]["relevance_score"] == 0.95
+        assert "document" not in result.results[0]
+
+        assert result.results[1]["index"] == 0
+        assert result.results[1]["relevance_score"] == 0.80
+        assert "document" not in result.results[1]
+
+        # Test with request_data={"ignoreRecordDetailsInResponse": True}
+        result_request_data = self.config.transform_rerank_response(
+            model=self.model,
+            raw_response=mock_response,
+            model_response=model_response,
+            logging_obj=mock_logging,
+            request_data={"ignoreRecordDetailsInResponse": True},
+        )
+        assert "document" not in result_request_data.results[0]
+        assert "document" not in result_request_data.results[1]
+
+        # Test with litellm_params={"return_documents": True}
+        result_litellm_params = self.config.transform_rerank_response(
+            model=self.model,
+            raw_response=mock_response,
+            model_response=model_response,
+            logging_obj=mock_logging,
+            litellm_params={"return_documents": True},
+        )
+        assert result_litellm_params.results[0]["document"]["text"] == "Content of document 1"
+
+        # Test with request_data={"return_documents": True}
+        result_req_data_true = self.config.transform_rerank_response(
+            model=self.model,
+            raw_response=mock_response,
+            model_response=model_response,
+            logging_obj=mock_logging,
+            request_data={"return_documents": True},
+        )
+        assert result_req_data_true.results[0]["document"]["text"] == "Content of document 1"
+
+        # Test with records missing content
+        no_content_response_data = {"records": [{"id": "0", "score": 0.9}]}
+        mock_no_content = MagicMock(spec=httpx.Response)
+        mock_no_content.json.return_value = no_content_response_data
+        mock_no_content.text = json.dumps(no_content_response_data)
+        result_no_content = self.config.transform_rerank_response(
+            model=self.model,
+            raw_response=mock_no_content,
+            model_response=model_response,
+            logging_obj=mock_logging,
+            optional_params={"return_documents": True},
+        )
+        assert "document" not in result_no_content.results[0]
+
     def test_transform_rerank_response_with_ignore_record_details(self):
         """Test response transformation when ignoreRecordDetailsInResponse=true."""
         # Mock response with only IDs (when ignoreRecordDetailsInResponse=true)
@@ -387,9 +512,7 @@ class TestVertexAIRerankTransform:
         # Verify title generation
         assert request_data["records"][0]["title"] == "This is a"  # First 3 words
         assert request_data["records"][1]["title"] == "Short doc"  # Less than 3 words
-        assert (
-            request_data["records"][2]["title"] == "Another document with"
-        )  # First 3 words
+        assert request_data["records"][2]["title"] == "Another document with"  # First 3 words
 
     def test_record_id_generation(self):
         """Test that record IDs are generated correctly with 0-based indexing."""
@@ -470,9 +593,7 @@ class TestVertexAIRerankTransform:
         importlib.reload(litellm) in conftest.py.
         """
         # Mock the authentication at instance level
-        mock_ensure_access_token = MagicMock(
-            return_value=("test-access-token", "test-project-123")
-        )
+        mock_ensure_access_token = MagicMock(return_value=("test-access-token", "test-project-123"))
         self.config._ensure_access_token = mock_ensure_access_token
 
         optional_params = {
@@ -510,9 +631,7 @@ class TestVertexAIRerankTransform:
         Uses instance-level mocking to avoid class-reference issues caused by
         importlib.reload(litellm) in conftest.py.
         """
-        mock_ensure_access_token = MagicMock(
-            return_value=("test-access-token", "project-from-token")
-        )
+        mock_ensure_access_token = MagicMock(return_value=("test-access-token", "project-from-token"))
         self.config._ensure_access_token = mock_ensure_access_token
 
         optional_params = {