Merge 86fc475f9c into 8588a2ea42

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)))

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]