diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index 2ec4f2da79b..b745d44c566 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/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))) diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py index 7fea5ac0965..9e97bb25db4 100644 --- a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py +++ b/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." diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py index c2ea6f6fab9..bdae83f5690 100644 --- a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py +++ b/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 = {