fix(rerank): emit latency and cost headers on /rerank (#35419)

* fix(rerank): emit latency and cost headers on /rerank

Thread the logging object into the rerank httpx calls and pass hidden_params through to get_custom_headers, so x-litellm-overhead-duration-ms, x-litellm-response-duration-ms, x-litellm-response-cost, x-litellm-call-id and the LITELLM_DETAILED_TIMING x-litellm-timing-* headers show up on rerank like they do on chat completions

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(rerank): keep zero response cost in the /rerank cost header

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* ci: assign the new rerank endpoint tests to the proxy-endpoints shard

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test: suppress TQ008 on the rerank header tests with reasons

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: milan <milan@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: yassin <yassin@berri.ai>

.github/workflows/test-unit.yml

@@ -164,6 +164,7 @@ jobs:
               tests/test_litellm/proxy/public_endpoints
               tests/test_litellm/proxy/prompts
               tests/test_litellm/proxy/rag_endpoints
+              tests/test_litellm/proxy/rerank_endpoints
               tests/test_litellm/proxy/realtime_endpoints
               tests/test_litellm/proxy/ui_crud_endpoints
               tests/test_litellm/proxy/config_resolvers

litellm/llms/bedrock/rerank/handler.py

@@ -29,6 +29,7 @@ class BedrockRerankHandler(BaseAWSLLM):
     async def arerank(
         self,
         prepared_request: BedrockPreparedRequest,
+        logging_obj: LitellmLogging,
         timeout: float | httpx.Timeout | None = None,
         client: AsyncHTTPHandler | None = None,
     ):
@@ -40,6 +41,7 @@ class BedrockRerankHandler(BaseAWSLLM):
                 headers=dict(prepared_request["prepped"].headers),
                 data=prepared_request["body"],
                 timeout=timeout,
+                logging_obj=logging_obj,
             )
             response.raise_for_status()
         except httpx.HTTPStatusError as err:
@@ -98,6 +100,7 @@ class BedrockRerankHandler(BaseAWSLLM):
         if _is_async:
             return self.arerank(
                 prepared_request,
+                logging_obj=logging_obj,
                 timeout=timeout,
                 client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None,
             )

litellm/llms/custom_httpx/llm_http_handler.py

@@ -1203,6 +1203,7 @@ class BaseLLMHTTPHandler:
                 headers=headers,
                 data=json.dumps(request_data),
                 timeout=timeout,
+                logging_obj=logging_obj,
             )
         except Exception as e:
             raise self._handle_error(e=e, provider_config=provider_config)

litellm/proxy/rerank_endpoints/endpoints.py

@@ -90,12 +90,15 @@ async def rerank(
         fastapi_response.headers.update(
             ProxyBaseLLMRequestProcessing.get_custom_headers(
                 user_api_key_dict=user_api_key_dict,
+                call_id=hidden_params.get("litellm_call_id", None) or data.get("litellm_call_id", None),
                 model_id=model_id,
                 cache_key=cache_key,
                 api_base=api_base,
                 version=version,
+                response_cost=hidden_params.get("response_cost", None),
                 model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
                 request_data=data,
+                hidden_params=hidden_params,
                 **additional_headers,
             )
         )

tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py

@@ -77,7 +77,7 @@ def test_bedrock_rerank_header_forwarding_sync(model):
 
     with (
         patch.object(client, "post") as mock_post,
-        patch(
+        patch(  # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
             "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
             return_value=mock_credentials_info,
         ),
@@ -170,7 +170,7 @@ async def test_bedrock_rerank_header_forwarding_async(model):
 
     with (
         patch.object(client, "post", new_callable=AsyncMock) as mock_post,
-        patch(
+        patch(  # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
             "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
             return_value=mock_credentials_info,
         ),
@@ -241,7 +241,7 @@ def test_bedrock_rerank_timeout_sync():
 
     with (
         patch.object(client, "post") as mock_post,
-        patch(
+        patch(  # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
             "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
             return_value=mock_credentials_info,
         ),
@@ -285,7 +285,7 @@ async def test_bedrock_rerank_timeout_async():
 
     with (
         patch.object(client, "post", new_callable=AsyncMock) as mock_post,
-        patch(
+        patch(  # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
             "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
             return_value=mock_credentials_info,
         ),
@@ -340,7 +340,7 @@ def test_bedrock_rerank_extra_headers_and_headers_merge():
 
     with (
         patch.object(client, "post") as mock_post,
-        patch(
+        patch(  # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
             "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
             return_value=mock_credentials_info,
         ),
@@ -400,3 +400,32 @@ def test_bedrock_rerank_extra_headers_and_headers_merge():
 
         except Exception as e:
             pytest.fail(f"Failed to merge and forward headers: {str(e)}")
+
+
+@pytest.mark.asyncio
+async def test_bedrock_rerank_records_llm_api_duration():
+    """The bedrock rerank handler must feed httpx timing into the logging obj, so the
+    proxy can emit x-litellm-overhead-duration-ms / x-litellm-timing-* on /rerank."""
+    import httpx
+
+    def handle(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(200, json=bedrock_rerank_response)
+
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
+
+    with patch(  # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
+        "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
+        return_value=create_mock_credentials(),
+    ):
+        response = await litellm.arerank(
+            model="bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0",
+            query=test_query,
+            documents=test_documents,
+            top_n=3,
+            client=client,
+            aws_region_name="us-east-1",
+        )
+
+    assert response._hidden_params["litellm_overhead_time_ms"] is not None
+    assert response._hidden_params["_response_ms"] >= response._hidden_params["litellm_overhead_time_ms"]

tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py

@@ -2528,6 +2528,37 @@ def test_only_callbacks_that_can_charge_a_frame_are_collected_for_ws_quota(monke
     assert _collect_ws_project_quota_callbacks() == (quota,)
 
 
+@pytest.mark.asyncio
+async def test_async_rerank_records_llm_api_duration():
+    """arerank must feed the httpx timing into the logging obj, so the proxy can emit
+    x-litellm-overhead-duration-ms / x-litellm-timing-* on /rerank."""
+
+    def handle(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            200,
+            json={
+                "id": "rerank-1",
+                "results": [{"index": 0, "relevance_score": 0.9}],
+                "meta": {"api_version": {"version": "2"}, "billed_units": {"search_units": 1}},
+            },
+        )
+
+    client = AsyncHTTPHandler()
+    client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
+
+    response = await litellm.arerank(
+        model="cohere/rerank-v3.5",
+        query="what is the capital of france",
+        documents=["paris", "berlin"],
+        top_n=1,
+        api_key="fake-key",
+        client=client,
+    )
+
+    assert response._hidden_params["litellm_overhead_time_ms"] is not None
+    assert response._hidden_params["_response_ms"] >= response._hidden_params["litellm_overhead_time_ms"]
+
+
 class _JSONBodyVideoConfig(OpenAIVideoConfig):
     def use_multipart_form_data(self) -> bool:
         return False

tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py

@@ -0,0 +1,120 @@
+"""
+Tests for rerank_endpoints/endpoints.py response headers.
+"""
+
+import json
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from fastapi import Request, Response
+
+import litellm.proxy.common_request_processing as common_request_processing_mod
+import litellm.proxy.proxy_server as proxy_server_mod
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.rerank_endpoints.endpoints import rerank
+from litellm.types.utils import RerankResponse
+
+HIDDEN_PARAMS = {
+    "model_id": "deployment-1",
+    "api_base": "https://bedrock-agent-runtime.us-east-1.amazonaws.com",
+    "response_cost": 0.002,
+    "_response_ms": 1500.5,
+    "litellm_overhead_time_ms": 12.5,
+    "callback_duration_ms": 1.25,
+    "timing_llm_api_ms": 1488.0,
+    "timing_pre_processing_ms": 10.0,
+    "timing_post_processing_ms": 2.5,
+    "timing_message_copy_ms": 0.01,
+}
+
+
+def _build_request() -> Request:
+    body = json.dumps({"model": "rerank-model", "query": "q", "documents": ["a", "b"]}).encode()
+
+    async def receive():
+        return {"type": "http.request", "body": body, "more_body": False}
+
+    return Request(
+        scope={
+            "type": "http",
+            "method": "POST",
+            "path": "/rerank",
+            "headers": [(b"content-type", b"application/json")],
+            "query_string": b"",
+        },
+        receive=receive,
+    )
+
+
+async def _call_rerank(hidden_params: dict = HIDDEN_PARAMS) -> Response:
+    response = RerankResponse(id="rerank-1", results=[{"index": 0, "relevance_score": 0.9}])
+    response._hidden_params = dict(hidden_params)
+
+    fastapi_response = Response()
+    proxy_logging_obj = MagicMock()
+    proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"])
+    proxy_logging_obj.update_request_status = AsyncMock()
+
+    async def fake_add_litellm_data_to_request(**kwargs):
+        return {**kwargs["data"], "litellm_call_id": "call-123"}
+
+    async def fake_route_request(**kwargs):
+        async def _call():
+            return response
+
+        return _call()
+
+    with (
+        patch.object(proxy_server_mod, "add_litellm_data_to_request", fake_add_litellm_data_to_request),  # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
+        patch.object(proxy_server_mod, "route_request", fake_route_request),  # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
+        patch.object(proxy_server_mod, "proxy_logging_obj", proxy_logging_obj),  # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
+        patch.object(proxy_server_mod, "llm_router", MagicMock()),  # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
+        patch.object(proxy_server_mod, "version", "1.2.3"),  # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
+    ):
+        await rerank(
+            request=_build_request(),
+            fastapi_response=fastapi_response,
+            user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
+        )
+
+    return fastapi_response
+
+
+@pytest.mark.asyncio
+async def test_rerank_emits_latency_and_cost_headers():
+    """/rerank must surface the same hidden_params-derived headers as /chat/completions."""
+    fastapi_response = await _call_rerank()
+
+    assert fastapi_response.headers["x-litellm-call-id"] == "call-123"
+    assert fastapi_response.headers["x-litellm-response-duration-ms"] == "1500.5"
+    assert fastapi_response.headers["x-litellm-overhead-duration-ms"] == "12.5"
+    assert fastapi_response.headers["x-litellm-callback-duration-ms"] == "1.25"
+    assert fastapi_response.headers["x-litellm-response-cost"] == "0.002"
+
+
+@pytest.mark.asyncio
+async def test_rerank_emits_detailed_timing_headers_when_enabled():
+    """LITELLM_DETAILED_TIMING must also work on /rerank, not just /chat/completions."""
+    with patch.object(common_request_processing_mod, "LITELLM_DETAILED_TIMING", True):  # test-quality-ok: LITELLM_DETAILED_TIMING is a module constant; toggling it is the behavior under test
+        fastapi_response = await _call_rerank()
+
+    assert fastapi_response.headers["x-litellm-timing-llm-api-ms"] == "1488.0"
+    assert fastapi_response.headers["x-litellm-timing-pre-processing-ms"] == "10.0"
+    assert fastapi_response.headers["x-litellm-timing-post-processing-ms"] == "2.5"
+    assert fastapi_response.headers["x-litellm-timing-message-copy-ms"] == "0.01"
+
+
+@pytest.mark.asyncio
+async def test_rerank_emits_zero_response_cost_header():
+    """A free deployment costs 0.0, which is a real cost and must not be dropped."""
+    fastapi_response = await _call_rerank({**HIDDEN_PARAMS, "response_cost": 0.0})
+
+    assert fastapi_response.headers["x-litellm-response-cost"] == "0.0"
+
+
+@pytest.mark.asyncio
+async def test_rerank_omits_detailed_timing_headers_when_disabled():
+    with patch.object(common_request_processing_mod, "LITELLM_DETAILED_TIMING", False):  # test-quality-ok: LITELLM_DETAILED_TIMING is a module constant; toggling it is the behavior under test
+        fastapi_response = await _call_rerank()
+
+    assert "x-litellm-timing-llm-api-ms" not in fastapi_response.headers