From 86fc475f9c591cde1c71441b21a6df6683b1bc92 Mon Sep 17 00:00:00 2001 From: ksk2023 <154514711+ksk2023@users.noreply.github.com> Date: Wed, 26 Aug 2026 23:31:39 +0800 Subject: [PATCH] fix(vertex_ai/rerank): populate document text when return_documents is true The response transformer read only id and score from each Vertex record and silently discarded content, so results[i].document was always absent regardless of the return_documents value - the request side already sets ignoreRecordDetailsInResponse = not return_documents correctly, but the response side never surfaced the record content Vertex returns. Read record['content'] and surface it as document.text (Cohere rerank format) whenever Vertex includes it. When return_documents=false Vertex returns IDs only, so the document field stays omitted, as expected. Same class of bug as #38006 (Bedrock rerank), fixed in #38007. Fixes #38239 --- .../llms/vertex_ai/rerank/transformation.py | 35 ++++---- .../test_vertex_ai_rerank_return_documents.py | 89 +++++++++++++++++++ 2 files changed, 109 insertions(+), 15 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_return_documents.py diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index 2ec4f2da79b..dac5c71d413 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, ) @@ -181,21 +182,24 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): # 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), - } - ) + result = { + "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 - } - ) + result = { + "index": int(record["id"]), + "relevance_score": 1.0, # Default score when details are ignored + } + # Vertex returns the record content when ignoreRecordDetailsInResponse=false + # (i.e. the caller passed return_documents=true). Surface it as document.text + # so the response matches the Cohere rerank format. + content = record.get("content") + if content is not None: + result["document"] = RerankResponseDocument(text=content) + results.append(result) # Sort by relevance score (descending) results.sort(key=lambda x: x["relevance_score"], reverse=True) @@ -204,9 +208,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_return_documents.py b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_return_documents.py new file mode 100644 index 00000000000..967fb4239e9 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_return_documents.py @@ -0,0 +1,89 @@ +""" +Regression tests for Vertex AI rerank return_documents behavior. + +Issue: the response transformer read only `id` and `score` from each Vertex +record and silently discarded `content`, so `results[i].document.text` was +always absent even when the caller passed return_documents=True (the default). + +The request side already sets `ignoreRecordDetailsInResponse = not +return_documents`, so when return_documents=True Vertex returns `content` on +every record. These tests verify the response transformer now surfaces that +content as `document.text`, and omits the field when Vertex returns IDs only. +""" + +import json +from unittest.mock import MagicMock + +import httpx + +from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig +from litellm.types.rerank import RerankResponse + + +def _mock_response(response_data: dict) -> MagicMock: + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = response_data + mock_response.text = json.dumps(response_data) + return mock_response + + +def _transform(response_data: dict) -> RerankResponse: + config = VertexAIRerankConfig() + return config.transform_rerank_response( + model="semantic-ranker-default@latest", + raw_response=_mock_response(response_data), + model_response=RerankResponse(), + logging_obj=MagicMock(), + ) + + +def test_vertex_rerank_return_documents_true_populates_document_text(): + """return_documents=True: Vertex returns content, it must reach document.text.""" + response_data = { + "records": [ + { + "id": "1", + "score": 0.95, + "title": "doc 1 title", + "content": "doc 1", + }, + { + "id": "0", + "score": 0.42, + "title": "doc 0 title", + "content": "doc 0", + }, + ] + } + + result = _transform(response_data) + + 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"] == "doc 1" + assert result.results[1]["document"]["text"] == "doc 0" + + +def test_vertex_rerank_return_documents_false_omits_document(): + """return_documents=False: Vertex returns IDs only, no document field.""" + response_data = {"records": [{"id": "1"}, {"id": "0"}]} + + result = _transform(response_data) + + assert len(result.results) == 2 + assert result.results[0]["index"] == 1 + assert "document" not in result.results[0] + assert "document" not in result.results[1] + + +def test_vertex_rerank_record_without_content_omits_document(): + """A record that lacks content (edge case) must not gain an empty document.""" + response_data = {"records": [{"id": "2", "score": 0.7}]} + + result = _transform(response_data) + + assert len(result.results) == 1 + assert result.results[0]["index"] == 2 + assert result.results[0]["relevance_score"] == 0.7 + assert "document" not in result.results[0]