From bb22742025065c10baba94b7d587ab000c6174ff Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 25 Aug 2026 15:54:25 -0700 Subject: [PATCH] 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 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yassin --- .github/workflows/test-unit.yml | 1 + litellm/llms/bedrock/rerank/handler.py | 3 + litellm/llms/custom_httpx/llm_http_handler.py | 1 + litellm/proxy/rerank_endpoints/endpoints.py | 3 + .../test_bedrock_rerank_header_forwarding.py | 39 +++++- .../custom_httpx/test_llm_http_handler.py | 31 +++++ .../proxy/rerank_endpoints/__init__.py | 0 .../proxy/rerank_endpoints/test_endpoints.py | 120 ++++++++++++++++++ 8 files changed, 193 insertions(+), 5 deletions(-) create mode 100644 tests/test_litellm/proxy/rerank_endpoints/__init__.py create mode 100644 tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 2dfca3d308f..c23678c51ae 100644 --- a/.github/workflows/test-unit.yml +++ b/.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 diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 1cc72f265eb..a2a1acec80c 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/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, ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index b7ddb55ae89..bb582f677be 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/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) diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py index 45b190c1f9d..dd5803796b7 100644 --- a/litellm/proxy/rerank_endpoints/endpoints.py +++ b/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, ) ) diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index b2a2046b131..253edad57e9 100644 --- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/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"] diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index a93b14d45f3..694a01cda5f 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/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 diff --git a/tests/test_litellm/proxy/rerank_endpoints/__init__.py b/tests/test_litellm/proxy/rerank_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py b/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py new file mode 100644 index 00000000000..9f11ff6f20d --- /dev/null +++ b/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