litellm/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py
mateo-berri 8b89c909a9 fix(proxy): keep a ProxyException's status and label 408s in the OpenAI error payload
error_status_code only read status_code, so a ProxyException raised
before routing (which stores its status as the string code) answered
500 with its 4xx type through the rerank, images, realtime, files, and
pass-through tails. It now falls back to a decimal code. A 408 maps to
timeout_error instead of invalid_request_error.

Tail regressions for rerank, images, realtime calls, and the chat
pass-through fail at the merge base with ('None', 'None'); the new
files-test helpers are fully typed.
2026-09-08 12:16:08 -07:00

172 lines
7.6 KiB
Python

"""
Tests for rerank_endpoints/endpoints.py response headers.
"""
import json
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException, 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 ProxyException, 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
async def _rerank_failure(
failure: Exception, *, raised_before_routing: bool, monkeypatch: pytest.MonkeyPatch
) -> ProxyException:
proxy_logging_obj = MagicMock()
proxy_logging_obj.pre_call_hook = AsyncMock(
side_effect=failure if raised_before_routing else lambda **kwargs: kwargs["data"]
)
proxy_logging_obj.post_call_failure_hook = AsyncMock()
async def fake_add_litellm_data_to_request(**kwargs: object) -> object:
return kwargs["data"]
async def failing_route_request(**kwargs: object) -> None:
raise failure
monkeypatch.setattr(proxy_server_mod, "add_litellm_data_to_request", fake_add_litellm_data_to_request)
monkeypatch.setattr(proxy_server_mod, "route_request", failing_route_request)
monkeypatch.setattr(proxy_server_mod, "proxy_logging_obj", proxy_logging_obj)
monkeypatch.setattr(proxy_server_mod, "llm_router", MagicMock())
monkeypatch.setattr(proxy_server_mod, "version", "1.2.3")
with pytest.raises(ProxyException) as raised:
await rerank(
request=_build_request(),
fastapi_response=Response(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
)
return raised.value
@pytest.mark.asyncio
async def test_a_model_the_router_cannot_serve_answers_an_openai_typed_error(monkeypatch: pytest.MonkeyPatch):
"""A bare HTTPException carries no type or param, so the tail used to ship the
literal string "None" in both fields."""
failure = HTTPException(status_code=404, detail={"error": "rerank: Invalid model name passed in model=rerank-model"})
error = await _rerank_failure(failure, raised_before_routing=False, monkeypatch=monkeypatch)
assert (error.type, error.param, error.code) == ("invalid_request_error", None, "404")
@pytest.mark.asyncio
async def test_a_rejection_raised_before_routing_keeps_its_own_status(monkeypatch: pytest.MonkeyPatch):
"""A ProxyException stores its status as the string ``code``, which the tail used to
miss and rewrap as a 500 while keeping the 4xx type and param."""
rejection = ProxyException(message="session_id is required", type="bad_request_error", param="session_id", code=400)
error = await _rerank_failure(rejection, raised_before_routing=True, monkeypatch=monkeypatch)
assert (error.type, error.param, error.code) == ("bad_request_error", "session_id", "400")