mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
fix(proxy): resolve x-litellm-call-id from response metadata when routes omit call_id
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
63386d6cc6
commit
ece2bbfed7
3 changed files with 88 additions and 1 deletions
|
|
@ -1571,6 +1571,9 @@ class ProxyBaseLLMRequestProcessing:
|
|||
) -> dict:
|
||||
exclude_values: Final = {"", None, "None"}
|
||||
hidden_params = hidden_params or {}
|
||||
resolved_call_id: Final = (
|
||||
call_id or hidden_params.get("litellm_call_id") or (request_data or {}).get("litellm_call_id")
|
||||
)
|
||||
timing_values: Final = _timing_values(
|
||||
hidden_params=hidden_params,
|
||||
logging_obj=litellm_logging_obj,
|
||||
|
|
@ -1598,7 +1601,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
classifier_cost: Final = _classifier_cost_from_request_data(request_data)
|
||||
|
||||
headers: Final = {
|
||||
"x-litellm-call-id": call_id,
|
||||
"x-litellm-call-id": resolved_call_id,
|
||||
"x-litellm-model-id": model_id,
|
||||
"x-litellm-model-name": model_name,
|
||||
"x-litellm-cache-key": cache_key,
|
||||
|
|
|
|||
|
|
@ -1402,6 +1402,51 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
assert "x-litellm-key-spend" in headers_7
|
||||
assert float(headers_7["x-litellm-key-spend"]) == 0.001 # Should use original spend on error
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("hidden_params", "request_data", "expected_call_id"),
|
||||
[
|
||||
(
|
||||
{"litellm_call_id": "call-from-hidden-params"},
|
||||
{"litellm_call_id": "call-from-request"},
|
||||
"call-from-hidden-params",
|
||||
),
|
||||
({}, {"litellm_call_id": "call-from-request"}, "call-from-request"),
|
||||
({"model_id": "m-1"}, {"litellm_call_id": "call-from-request"}, "call-from-request"),
|
||||
],
|
||||
)
|
||||
def test_get_custom_headers_call_id_falls_back_to_hidden_params_then_request_data(
|
||||
self, hidden_params, request_data, expected_call_id
|
||||
):
|
||||
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
|
||||
mock_user_api_key_dict.tpm_limit = None
|
||||
mock_user_api_key_dict.rpm_limit = None
|
||||
mock_user_api_key_dict.max_budget = None
|
||||
mock_user_api_key_dict.spend = 0.0
|
||||
|
||||
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
hidden_params=hidden_params,
|
||||
request_data=request_data,
|
||||
)
|
||||
|
||||
assert headers["x-litellm-call-id"] == expected_call_id
|
||||
|
||||
def test_get_custom_headers_explicit_call_id_wins_over_fallbacks(self):
|
||||
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
|
||||
mock_user_api_key_dict.tpm_limit = None
|
||||
mock_user_api_key_dict.rpm_limit = None
|
||||
mock_user_api_key_dict.max_budget = None
|
||||
mock_user_api_key_dict.spend = 0.0
|
||||
|
||||
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_id="explicit-call-id",
|
||||
hidden_params={"litellm_call_id": "call-from-hidden-params"},
|
||||
request_data={"litellm_call_id": "call-from-request"},
|
||||
)
|
||||
|
||||
assert headers["x-litellm-call-id"] == "explicit-call-id"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_queue_time_seconds_is_set_in_metadata(self, monkeypatch):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -12795,6 +12795,45 @@ async def test_moderations_reraises_proxy_exception_unwrapped():
|
|||
mock_logging.post_call_failure_hook.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_moderations_response_carries_litellm_call_id_header():
|
||||
from fastapi import Response
|
||||
|
||||
from litellm.types.utils import ModerationCreateResponse
|
||||
|
||||
call_id = "moderation-call-id-123"
|
||||
moderation_response = ModerationCreateResponse(id="modr-1", model="omni-moderation-latest", results=[])
|
||||
moderation_response._hidden_params = {"litellm_call_id": call_id, "model_id": "mod-deployment-1"}
|
||||
|
||||
async def fake_llm_call():
|
||||
return moderation_response
|
||||
|
||||
async def passthrough_add_litellm_data(data, **kwargs):
|
||||
return {**data, "litellm_call_id": call_id}
|
||||
|
||||
request = MagicMock()
|
||||
request.body = AsyncMock(return_value=b'{"input": "hi"}')
|
||||
fastapi_response = Response()
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", spend=0.0)
|
||||
|
||||
with (
|
||||
patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point
|
||||
patch.object(proxy_server_module, "route_request", new=AsyncMock(return_value=fake_llm_call())), # test-quality-ok: fakes the provider call so the response headers assembled by the real route are observable
|
||||
patch.object(proxy_server_module, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global, no injection point
|
||||
):
|
||||
mock_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
|
||||
mock_logging.update_request_status = AsyncMock()
|
||||
result = await proxy_server_module.moderations(
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
assert result is moderation_response
|
||||
assert fastapi_response.headers["x-litellm-call-id"] == call_id
|
||||
assert fastapi_response.headers["x-litellm-model-id"] == "mod-deployment-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_agents_in_db_rebuilds_registry_under_agent_reconcile_lock(monkeypatch):
|
||||
from litellm.proxy.agent_endpoints.agent_registry import (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue