ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

(fix) propagate x-litellm-model-id in responses (#16986)

Browse Source 
* propagate model id on errors too

* make it work for messages and streaming

* fix

* cleanup

* cleanup

* final

* cleanup

* clean up method name and fix responses api streaming

* remove comment
This commit is contained in:
Raghav Jhavar 2025-11-24 23:40:43 -05:00 committed by GitHub
parent 282ac87617
commit bd8196f982
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 359 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
litellm/proxy/anthropic_endpoints/endpoints.py
View File

@ -154,8 +154,13 @@ async def anthropic_response( # noqa: PLR0915
response = responses[1]
# Extract model_id from request metadata (set by router during routing)
litellm_metadata = data.get("litellm_metadata", {}) or {}
model_info = litellm_metadata.get("model_info", {}) or {}
model_id = model_info.get("id", "") or ""
# Get other metadata from hidden_params
hidden_params = getattr(response, "_hidden_params", {}) or {}
model_id = hidden_params.get("model_id", None) or ""
cache_key = hidden_params.get("cache_key", None) or ""
api_base = hidden_params.get("api_base", None) or ""
response_cost = hidden_params.get("response_cost", None) or ""
@ -216,12 +221,32 @@ async def anthropic_response( # noqa: PLR0915
str(e)
)
)
# Extract model_id from request metadata (same as success path)
litellm_metadata = data.get("litellm_metadata", {}) or {}
model_info = litellm_metadata.get("model_info", {}) or {}
model_id = model_info.get("id", "") or ""
# Get headers
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=data.get("litellm_call_id", ""),
model_id=model_id,
version=version,
response_cost=0,
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
request_data=data,
timeout=getattr(e, "timeout", None),
litellm_logging_obj=None,
)
error_msg = f"{str(e)}"
raise ProxyException(
message=getattr(e, "message", error_msg),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", 500),
headers=headers,
)

64
litellm/proxy/common_request_processing.py
View File

@ -344,6 +344,7 @@ class ProxyBaseLLMRequestProcessing:
user_max_tokens: Optional[int] = None,
user_api_base: Optional[str] = None,
model: Optional[str] = None,
llm_router: Optional[Router] = None,
) -> Tuple[dict, LiteLLMLoggingObj]:
start_time = datetime.now() # start before calling guardrail hooks
@ -498,6 +499,7 @@ class ProxyBaseLLMRequestProcessing:
user_api_base=user_api_base,
model=model,
route_type=route_type,
llm_router=llm_router,
)
tasks = []
@ -536,6 +538,13 @@ class ProxyBaseLLMRequestProcessing:
hidden_params = getattr(response, "_hidden_params", {}) or {}
model_id = hidden_params.get("model_id", None) or ""
# Fallback: extract model_id from litellm_metadata if not in hidden_params
if not model_id:
litellm_metadata = self.data.get("litellm_metadata", {}) or {}
model_info = litellm_metadata.get("model_info", {}) or {}
model_id = model_info.get("id", "") or ""
cache_key = hidden_params.get("cache_key", None) or ""
api_base = hidden_params.get("api_base", None) or ""
response_cost = hidden_params.get("response_cost", None) or ""
@ -756,11 +765,19 @@ class ProxyBaseLLMRequestProcessing:
_litellm_logging_obj: Optional[LiteLLMLoggingObj] = self.data.get(
"litellm_logging_obj", None
)
# Attempt to get model_id from logging object
#
# Note: We check the direct model_info path first (not nested in metadata) because that's where the router sets it.
# The nested metadata path is only a fallback for cases where model_info wasn't set at the top level.
model_id = self.maybe_get_model_id(_litellm_logging_obj)
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=(
_litellm_logging_obj.litellm_call_id if _litellm_logging_obj else None
),
model_id=model_id,
version=version,
response_cost=0,
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
@ -1073,3 +1090,50 @@ class ProxyBaseLLMRequestProcessing:
obj.setdefault("usage", {})["cost"] = cost_val
return obj
return None
def maybe_get_model_id(self, _logging_obj: Optional[LiteLLMLoggingObj]) -> Optional[str]:
"""
Get model_id from logging object or request metadata.
The router sets model_info.id when selecting a deployment. This tries multiple locations
where the ID might be stored depending on the request lifecycle stage.
"""
model_id = None
if _logging_obj:
# 1. Try getting from litellm_params (updated during call)
if (
hasattr(_logging_obj, "litellm_params")
and _logging_obj.litellm_params
):
# First check direct model_info path (set by router.py with selected deployment)
model_info = _logging_obj.litellm_params.get("model_info") or {}
model_id = model_info.get("id", None)
# Fallback to nested metadata path
if not model_id:
metadata = _logging_obj.litellm_params.get("metadata") or {}
model_info = metadata.get("model_info") or {}
model_id = model_info.get("id", None)
# 2. Fallback to kwargs (initial)
if not model_id:
_kwargs = getattr(_logging_obj, "kwargs", None)
if _kwargs:
litellm_params = _kwargs.get("litellm_params", {})
# First check direct model_info path
model_info = litellm_params.get("model_info") or {}
model_id = model_info.get("id", None)
# Fallback to nested metadata path
if not model_id:
metadata = litellm_params.get("metadata") or {}
model_info = metadata.get("model_info") or {}
model_id = model_info.get("id", None)
# 3. Final fallback to self.data["litellm_metadata"] (for routes like /v1/responses that populate data before error)
if not model_id:
litellm_metadata = self.data.get("litellm_metadata", {}) or {}
model_info = litellm_metadata.get("model_info", {}) or {}
model_id = model_info.get("id", None)
return model_id

19
litellm/responses/streaming_iterator.py
View File

@ -8,7 +8,9 @@ import httpx
import litellm
from litellm.constants import STREAM_SSE_DONE_STRING
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.utils import ResponsesAPIRequestUtils
@ -51,6 +53,23 @@ class BaseResponsesAPIStreamingIterator:
self.litellm_metadata = litellm_metadata
self.custom_llm_provider = custom_llm_provider
# set hidden params for response headers (e.g., x-litellm-model-id)
# This matches ths stream wrapper in litellm/litellm_core_utils/streaming_handler.py
_api_base = get_api_base(
model=model or "",
optional_params=self.logging_obj.model_call_details.get(
"litellm_params", {}
),
)
_model_info: Dict = litellm_metadata.get("model_info", {}) if litellm_metadata else {}
self._hidden_params = {
"model_id": _model_info.get("id", None),
"api_base": _api_base,
}
self._hidden_params["additional_headers"] = process_response_headers(
self.response.headers or {}
) # GUARANTEE OPENAI HEADERS IN RESPONSE
def _process_chunk(self, chunk) -> Optional[ResponsesAPIStreamingResponse]:
"""Process a single chunk of data from the stream"""
if not chunk:

250
tests/test_litellm/proxy/test_model_id_header_propagation.py Normal file
View File

@ -0,0 +1,250 @@
"""
Test that x-litellm-model-id header is propagated correctly on error responses.
This test suite verifies the `maybe_get_model_id` method
which is responsible for extracting model_id from different locations
depending on the request lifecycle stage.
"""
import pytest
from unittest.mock import MagicMock
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy._types import UserAPIKeyAuth
def test_maybe_get_model_id_from_litellm_params():
"""
Test extraction of model_id from logging_obj.litellm_params (used by /v1/chat/completions).
"""
# Create a ProxyBaseLLMRequestProcessing instance
processor = ProxyBaseLLMRequestProcessing(data={})
# Create a mock logging object with model_info in litellm_params
mock_logging_obj = MagicMock()
mock_logging_obj.litellm_params = {
"model_info": {
"id": "test-model-id-from-litellm-params"
}
}
# Test extraction
model_id = processor.maybe_get_model_id(mock_logging_obj)
assert model_id == "test-model-id-from-litellm-params"
def test_maybe_get_model_id_from_litellm_params_nested():
"""
Test extraction of model_id from nested metadata in logging_obj.litellm_params.
"""
processor = ProxyBaseLLMRequestProcessing(data={})
# Create a mock logging object with model_info nested in metadata
mock_logging_obj = MagicMock()
mock_logging_obj.litellm_params = {
"metadata": {
"model_info": {
"id": "test-model-id-nested"
}
}
}
# Test extraction
model_id = processor.maybe_get_model_id(mock_logging_obj)
assert model_id == "test-model-id-nested"
def test_maybe_get_model_id_from_kwargs():
"""
Test extraction of model_id from logging_obj.kwargs (fallback path).
"""
processor = ProxyBaseLLMRequestProcessing(data={})
# Create a mock logging object with model_info in kwargs
mock_logging_obj = MagicMock()
mock_logging_obj.litellm_params = None
mock_logging_obj.kwargs = {
"litellm_params": {
"model_info": {
"id": "test-model-id-from-kwargs"
}
}
}
# Test extraction
model_id = processor.maybe_get_model_id(mock_logging_obj)
assert model_id == "test-model-id-from-kwargs"
def test_maybe_get_model_id_from_data():
"""
Test extraction of model_id from self.data (used by /v1/messages and /v1/responses).
"""
# Create a processor with model_info in data
processor = ProxyBaseLLMRequestProcessing(data={
"litellm_metadata": {
"model_info": {
"id": "test-model-id-from-data"
}
}
})
# Create a mock logging object without model_info
mock_logging_obj = MagicMock()
mock_logging_obj.litellm_params = {}
mock_logging_obj.kwargs = {}
# Test extraction - should fall back to self.data
model_id = processor.maybe_get_model_id(mock_logging_obj)
assert model_id == "test-model-id-from-data"
def test_maybe_get_model_id_no_logging_obj():
"""
Test extraction of model_id when logging_obj is None (should use self.data).
"""
# Create a processor with model_info in data
processor = ProxyBaseLLMRequestProcessing(data={
"litellm_metadata": {
"model_info": {
"id": "test-model-id-no-logging-obj"
}
}
})
# Test extraction with None logging_obj
model_id = processor.maybe_get_model_id(None)
assert model_id == "test-model-id-no-logging-obj"
def test_maybe_get_model_id_not_found():
"""
Test extraction of model_id when it's not available anywhere (should return None).
"""
processor = ProxyBaseLLMRequestProcessing(data={})
# Create a mock logging object without model_info anywhere
mock_logging_obj = MagicMock()
mock_logging_obj.litellm_params = {}
mock_logging_obj.kwargs = {}
# Test extraction - should return None
model_id = processor.maybe_get_model_id(mock_logging_obj)
assert model_id is None
def test_maybe_get_model_id_priority_litellm_params_over_data():
"""
Test that model_id from logging_obj.litellm_params takes priority over self.data.
"""
# Create a processor with model_info in both places
processor = ProxyBaseLLMRequestProcessing(data={
"litellm_metadata": {
"model_info": {
"id": "model-id-from-data"
}
}
})
# Create a mock logging object with model_info
mock_logging_obj = MagicMock()
mock_logging_obj.litellm_params = {
"model_info": {
"id": "model-id-from-litellm-params"
}
}
# Test extraction - should prefer litellm_params
model_id = processor.maybe_get_model_id(mock_logging_obj)
assert model_id == "model-id-from-litellm-params"
def test_get_custom_headers_includes_model_id():
"""
Test that get_custom_headers includes x-litellm-model-id when model_id is provided.
"""
# Create mock user_api_key_dict with all required attributes
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.user_id = "test-user"
mock_user_api_key_dict.team_id = "test-team"
mock_user_api_key_dict.tpm_limit = 1000
mock_user_api_key_dict.rpm_limit = 100
# Call get_custom_headers with a model_id
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=mock_user_api_key_dict,
model_id="test-model-123",
cache_key="test-cache-key",
api_base="https://api.example.com",
version="1.0.0",
response_cost=0.001,
request_data={},
hidden_params={}
)
# Verify model_id is in headers
assert "x-litellm-model-id" in headers
assert headers["x-litellm-model-id"] == "test-model-123"
def test_get_custom_headers_without_model_id():
"""
Test that get_custom_headers works correctly when model_id is None or empty.
"""
# Create mock user_api_key_dict with all required attributes
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.user_id = "test-user"
mock_user_api_key_dict.team_id = "test-team"
mock_user_api_key_dict.tpm_limit = 1000
mock_user_api_key_dict.rpm_limit = 100
# Call get_custom_headers without a model_id
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=mock_user_api_key_dict,
model_id=None,
cache_key="test-cache-key",
api_base="https://api.example.com",
version="1.0.0",
response_cost=0.001,
request_data={},
hidden_params={}
)
# x-litellm-model-id should not be in headers (or should be empty/None)
if "x-litellm-model-id" in headers:
assert headers["x-litellm-model-id"] in [None, ""]
def test_get_custom_headers_with_empty_string_model_id():
"""
Test that get_custom_headers handles empty string model_id correctly.
"""
# Create mock user_api_key_dict with all required attributes
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.user_id = "test-user"
mock_user_api_key_dict.team_id = "test-team"
mock_user_api_key_dict.tpm_limit = 1000
mock_user_api_key_dict.rpm_limit = 100
# Call get_custom_headers with empty string model_id
headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=mock_user_api_key_dict,
model_id="",
cache_key="test-cache-key",
api_base="https://api.example.com",
version="1.0.0",
response_cost=0.001,
request_data={},
hidden_params={}
)
# x-litellm-model-id should not be in headers (or should be empty)
if "x-litellm-model-id" in headers:
assert headers["x-litellm-model-id"] == ""