mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(router): preserve _hidden_params in FallbackStreamWrapper so x-litellm-overhead-duration-ms is emitted for streaming requests
This commit is contained in:
Ishaan Jaffer
parent
9873ad62fa
commit
de414859f2
2 changed files with 226 additions and 11 deletions
|
|
@ -58,12 +58,12 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
_get_parent_otel_span_from_kwargs,
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
)
|
||||
from litellm.litellm_core_utils.thread_pool_executor import executor
|
||||
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
from litellm.litellm_core_utils.thread_pool_executor import executor
|
||||
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
|
||||
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
|
||||
from litellm.router_strategy.least_busy import LeastBusyLoggingHandler
|
||||
|
|
@ -1474,6 +1474,9 @@ class Router:
|
|||
logging_obj=model_response.logging_obj,
|
||||
)
|
||||
self._async_generator = async_generator
|
||||
# Preserve hidden params (including litellm_overhead_time_ms) from original response
|
||||
if hasattr(model_response, "_hidden_params"):
|
||||
self._hidden_params = model_response._hidden_params.copy()
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
|
@ -6582,16 +6585,28 @@ class Router:
|
|||
raise ValueError("Deployment not found")
|
||||
|
||||
## GET BASE MODEL
|
||||
base_model = deployment.get("model_info", {}).get("base_model", None)
|
||||
base_model = (deployment.get("model_info") or {}).get("base_model", None)
|
||||
if base_model is None:
|
||||
base_model = deployment.get("litellm_params", {}).get("base_model", None)
|
||||
base_model = (deployment.get("litellm_params") or {}).get("base_model", None)
|
||||
|
||||
model = base_model
|
||||
|
||||
## GET PROVIDER
|
||||
## GET PROVIDER - reuse LiteLLM_Params if already constructed
|
||||
litellm_params_data = deployment.get("litellm_params")
|
||||
litellm_params: LiteLLM_Params
|
||||
if isinstance(litellm_params_data, LiteLLM_Params):
|
||||
litellm_params = litellm_params_data
|
||||
elif isinstance(litellm_params_data, dict) and "model" in litellm_params_data:
|
||||
litellm_params = LiteLLM_Params(**litellm_params_data)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Deployment missing valid litellm_params. "
|
||||
f"Got: {type(litellm_params_data).__name__}, "
|
||||
f"deployment_id: {(deployment.get('model_info') or {}).get('id', 'unknown')}"
|
||||
)
|
||||
_model, custom_llm_provider, _, _ = litellm.get_llm_provider(
|
||||
model=deployment.get("litellm_params", {}).get("model", ""),
|
||||
litellm_params=LiteLLM_Params(**deployment.get("litellm_params", {})),
|
||||
model=litellm_params.model,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
## SET MODEL TO 'model=' - if base_model is None + not azure
|
||||
|
|
@ -6607,10 +6622,10 @@ class Router:
|
|||
if potential_models is not None:
|
||||
for potential_model in potential_models:
|
||||
try:
|
||||
if potential_model.get("model_info", {}).get(
|
||||
if (potential_model.get("model_info") or {}).get(
|
||||
"id"
|
||||
) == deployment.get("model_info", {}).get("id"):
|
||||
model = potential_model.get("litellm_params", {}).get(
|
||||
) == (deployment.get("model_info") or {}).get("id"):
|
||||
model = (potential_model.get("litellm_params") or {}).get(
|
||||
"model"
|
||||
)
|
||||
break
|
||||
|
|
@ -6631,9 +6646,10 @@ class Router:
|
|||
model_info = litellm.get_model_info(model=model_info_name)
|
||||
|
||||
## CHECK USER SET MODEL INFO
|
||||
user_model_info = deployment.get("model_info", {})
|
||||
user_model_info = deployment.get("model_info") or {}
|
||||
|
||||
model_info.update(user_model_info)
|
||||
if model_info is not None:
|
||||
model_info.update(user_model_info)
|
||||
|
||||
return model_info
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
import copy
|
||||
import datetime
|
||||
from typing import AsyncGenerator
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -1173,3 +1175,200 @@ class TestOverrideOpenAIResponseModel:
|
|||
assert response_obj.model == fallback_model
|
||||
|
||||
|
||||
class TestStreamingOverheadHeader:
|
||||
"""
|
||||
Tests that x-litellm-overhead-duration-ms is emitted in streaming responses.
|
||||
|
||||
Regression tests for: streaming requests not including overhead header.
|
||||
"""
|
||||
|
||||
def test_get_custom_headers_includes_overhead_when_set(self):
|
||||
"""
|
||||
get_custom_headers() returns x-litellm-overhead-duration-ms
|
||||
when litellm_overhead_time_ms is in hidden_params.
|
||||
"""
|
||||
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
|
||||
mock_user_api_key_dict.allowed_model_region = None
|
||||
|
||||
hidden_params = {
|
||||
"litellm_overhead_time_ms": 42.5,
|
||||
"_response_ms": 500.0,
|
||||
"model_id": "test-model-id",
|
||||
"api_base": "https://api.openai.com",
|
||||
}
|
||||
|
||||
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_id="test-call-id",
|
||||
model_id="test-model-id",
|
||||
cache_key="",
|
||||
api_base="https://api.openai.com",
|
||||
version="1.0.0",
|
||||
response_cost=0.001,
|
||||
model_region="",
|
||||
hidden_params=hidden_params,
|
||||
)
|
||||
|
||||
assert "x-litellm-overhead-duration-ms" in headers
|
||||
assert headers["x-litellm-overhead-duration-ms"] == "42.5"
|
||||
|
||||
def test_get_custom_headers_omits_overhead_when_none(self):
|
||||
"""
|
||||
get_custom_headers() omits x-litellm-overhead-duration-ms
|
||||
when litellm_overhead_time_ms is not in hidden_params.
|
||||
"""
|
||||
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
|
||||
mock_user_api_key_dict.allowed_model_region = None
|
||||
|
||||
hidden_params = {
|
||||
"_response_ms": 500.0,
|
||||
"model_id": "test-model-id",
|
||||
}
|
||||
|
||||
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_id="test-call-id",
|
||||
model_id="test-model-id",
|
||||
cache_key="",
|
||||
api_base="https://api.openai.com",
|
||||
version="1.0.0",
|
||||
response_cost=0.001,
|
||||
model_region="",
|
||||
hidden_params=hidden_params,
|
||||
)
|
||||
|
||||
# Should be absent (None gets filtered by exclude_values)
|
||||
assert "x-litellm-overhead-duration-ms" not in headers
|
||||
|
||||
def test_update_response_metadata_sets_overhead_on_stream_wrapper(self):
|
||||
"""
|
||||
update_response_metadata() sets litellm_overhead_time_ms on
|
||||
a streaming response's _hidden_params when llm_api_duration_ms is available.
|
||||
"""
|
||||
from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
|
||||
update_response_metadata,
|
||||
)
|
||||
|
||||
# Mock the logging object with llm_api_duration_ms set
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.model_call_details = {
|
||||
"llm_api_duration_ms": 200.0,
|
||||
"litellm_params": {},
|
||||
}
|
||||
mock_logging_obj.caching_details = None
|
||||
mock_logging_obj.callback_duration_ms = None
|
||||
mock_logging_obj.litellm_call_id = "test-call-id"
|
||||
mock_logging_obj._response_cost_calculator = MagicMock(return_value=0.001)
|
||||
|
||||
# Simulate a streaming result object with _hidden_params (like CustomStreamWrapper)
|
||||
stream_result = MagicMock()
|
||||
stream_result._hidden_params = {
|
||||
"model_id": "test-model-id",
|
||||
"api_base": "https://api.openai.com",
|
||||
"additional_headers": {},
|
||||
}
|
||||
|
||||
start_time = datetime.datetime.now() - datetime.timedelta(milliseconds=300)
|
||||
end_time = datetime.datetime.now()
|
||||
|
||||
update_response_metadata(
|
||||
result=stream_result,
|
||||
logging_obj=mock_logging_obj,
|
||||
model="gpt-4o",
|
||||
kwargs={},
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
assert "litellm_overhead_time_ms" in stream_result._hidden_params
|
||||
overhead = stream_result._hidden_params["litellm_overhead_time_ms"]
|
||||
assert overhead is not None
|
||||
assert isinstance(overhead, float)
|
||||
# overhead = total_response_ms (~300ms) - llm_api_duration_ms (200ms) = ~100ms
|
||||
assert overhead > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_response_includes_overhead_header(self):
|
||||
"""
|
||||
StreamingResponse returned by create_response() includes
|
||||
x-litellm-overhead-duration-ms in its headers.
|
||||
"""
|
||||
|
||||
async def mock_generator() -> AsyncGenerator[str, None]:
|
||||
yield 'data: {"id":"chatcmpl-test","choices":[{"delta":{"content":"hi"}}]}\n\n'
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
headers = {
|
||||
"x-litellm-overhead-duration-ms": "42.5",
|
||||
"x-litellm-call-id": "test-call-id",
|
||||
"x-litellm-model-id": "test-model-id",
|
||||
}
|
||||
|
||||
response = await create_response(
|
||||
generator=mock_generator(),
|
||||
media_type="text/event-stream",
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
assert isinstance(response, StreamingResponse)
|
||||
assert response.headers.get("x-litellm-overhead-duration-ms") == "42.5"
|
||||
|
||||
def test_streaming_overhead_header_in_custom_headers_from_stream_hidden_params(
|
||||
self,
|
||||
):
|
||||
"""
|
||||
Verifies that when get_custom_headers() is called with a streaming
|
||||
response's hidden_params (containing litellm_overhead_time_ms),
|
||||
the x-litellm-overhead-duration-ms header is correctly populated.
|
||||
|
||||
This tests the critical path: update_response_metadata sets the value
|
||||
→ get_custom_headers reads it → StreamingResponse header is set.
|
||||
"""
|
||||
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
|
||||
mock_user_api_key_dict.allowed_model_region = None
|
||||
|
||||
# This is what CustomStreamWrapper._hidden_params looks like after
|
||||
# update_response_metadata() has been called on it
|
||||
hidden_params = {
|
||||
"model_id": "openai-gpt4o-deployment",
|
||||
"api_base": "https://api.openai.com",
|
||||
"additional_headers": {},
|
||||
"litellm_overhead_time_ms": 55.3, # set by update_response_metadata
|
||||
"_response_ms": 280.0,
|
||||
"litellm_call_id": "test-call-id",
|
||||
"response_cost": 0.002,
|
||||
"cache_key": None,
|
||||
"fastest_response_batch_completion": None,
|
||||
"callback_duration_ms": None,
|
||||
}
|
||||
|
||||
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_id="test-call-id",
|
||||
model_id=hidden_params.get("model_id"),
|
||||
cache_key=hidden_params.get("cache_key") or "",
|
||||
api_base=hidden_params.get("api_base") or "",
|
||||
version="1.0.0",
|
||||
response_cost=hidden_params.get("response_cost"),
|
||||
model_region="",
|
||||
hidden_params=hidden_params,
|
||||
)
|
||||
|
||||
# The overhead header must be present and correct
|
||||
assert "x-litellm-overhead-duration-ms" in custom_headers, (
|
||||
"x-litellm-overhead-duration-ms header must be emitted during streaming. "
|
||||
"It was missing — this is the streaming overhead header regression."
|
||||
)
|
||||
assert custom_headers["x-litellm-overhead-duration-ms"] == "55.3"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue