scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Merge pull request #23432 from BerriAI/litellm_azure-model-router-show-actual-model

Browse source 
feat(azure_ai): show actual model used in Azure Model Router response
This commit is contained in:
Sameer Kankute 2026-03-12 22:18:50 +05:30 • committed by GitHub
commit d507f840d3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 379 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/litellm_core_utils/litellm_logging.py
View file

@ -5346,6 +5346,20 @@ def get_standard_logging_object_payload(
model_name = reconstruct_model_name(
kwargs.get("model", "") or "", custom_llm_provider, metadata
)
response_model_name: Optional[str] = None
if isinstance(final_response_obj, dict):
response_model_name = final_response_obj.get("model")
# For Azure Model Router, preserve the actual model in the top-level standard
# logging payload only when the user has opted in.
requested_model = kwargs.get("model")
if (
isinstance(requested_model, str)
and ("model_router" in requested_model.lower() or "model-router" in requested_model.lower())
and isinstance(response_model_name, str)
and response_model_name
):
model_name = response_model_name
payload: StandardLoggingPayload = StandardLoggingPayload(
id=str(id),

18
litellm/llms/azure_ai/azure_model_router/transformation.py
View file

@ -64,24 +64,17 @@ class AzureModelRouterConfig(AzureAIStudioConfig):
"""
Transform response for Model Router.
Preserves the original model path (including model_router/ prefix) in the response
for proper cost tracking and logging.
Extracts the actual model used from the Azure response (e.g., gpt-5-nano-2025-08-07)
and returns it with the azure_ai/ prefix for proper display and cost tracking.
"""
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
# Preserve the original model from litellm_params (includes routing prefixes like model_router/)
# This ensures cost tracking and logging use the full model path
original_model: str = litellm_params.get("model") or model
if not original_model.startswith("azure_ai/"):
# Add provider prefix if not already present
model_response.model = f"azure_ai/{original_model}"
else:
model_response.model = original_model
# Get base model for the parent call (strips routing prefixes for API compatibility)
base_model: str = AzureFoundryModelInfo.get_base_model(model)
return super().transform_response(
# Call parent transform_response first - this will extract the actual model
# from the raw response (e.g., "gpt-5-nano-2025-08-07")
model_response = super().transform_response(
model=base_model,
raw_response=raw_response,
model_response=model_response,
@ -94,6 +87,7 @@ class AzureModelRouterConfig(AzureAIStudioConfig):
api_key=api_key,
json_mode=json_mode,
)
return model_response
def calculate_additional_costs(
self, model: str, prompt_tokens: int, completion_tokens: int

39
litellm/proxy/common_request_processing.py
View file

@ -246,6 +246,29 @@ async def create_response(
)
def _is_azure_model_router_request(model: str) -> bool:
"""
Check if the requested model is an Azure Model Router.
Azure Model Router models follow the pattern:
- azure_ai/model_router/<deployment-name>
- azure_ai/model-router
- model_router/<deployment-name>
- model-router
Args:
model: The requested model name
Returns:
bool: True if this is an Azure Model Router request
"""
model_lower = model.lower()
return (
"model-router" in model_lower
or "model_router" in model_lower
)
def _override_openai_response_model(
*,
response_obj: Any,
@ -265,9 +288,11 @@ def _override_openai_response_model(
Errors are reserved for cases where the proxy cannot read/override the response model field.
Exception: If a fallback occurred (indicated by x-litellm-attempted-fallbacks header),
we should preserve the actual model that was used (the fallback model) rather than
overriding it with the originally requested model.
Exceptions:
1. If a fallback occurred (indicated by x-litellm-attempted-fallbacks header),
we preserve the actual model that was used (the fallback model).
2. If the request was to an Azure Model Router, we preserve the actual model
that was used (e.g., gpt-5-nano-2025-08-07) instead of the router model.
"""
if not requested_model:
return
@ -288,6 +313,14 @@ def _override_openai_response_model(
)
return
# Check if this is an Azure Model Router request - if so, preserve the actual model used
if _is_azure_model_router_request(requested_model):
verbose_proxy_logger.debug(
"%s: Azure Model Router detected - preserving actual model used from response instead of overriding to router model.",
log_context,
)
return
if isinstance(response_obj, dict):
downstream_model = response_obj.get("model")
if downstream_model != requested_model:

7
litellm/proxy/proxy_server.py
View file

@ -290,6 +290,7 @@ from litellm.proxy.batches_endpoints.endpoints import router as batches_router
from litellm.proxy.caching_routes import router as caching_router
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
_is_azure_model_router_request,
create_response,
)
from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy
@ -5427,6 +5428,12 @@ def _restamp_streaming_chunk_model(
if not requested_model_from_client or not isinstance(chunk, (BaseModel, dict)):
return chunk, model_mismatch_logged
# For Azure Model Router, preserve the actual model used in each chunk
if _is_azure_model_router_request(
requested_model_from_client
):
return chunk, model_mismatch_logged
downstream_model = (
chunk.get("model") if isinstance(chunk, dict) else getattr(chunk, "model", None)
)

96
tests/logging_callback_tests/test_standard_logging_payload.py
View file

@ -809,6 +809,102 @@ def test_usage_dict_roundtrip_in_payload(use_combined_usage_object):
assert usage_obj["total_tokens"] == 100
def test_standard_logging_payload_uses_actual_model_for_azure_router():
from litellm.litellm_core_utils.litellm_logging import (
Logging,
get_standard_logging_object_payload,
)
logging_obj = Logging(
model="azure_ai/model-router",
messages=[{"role": "user", "content": "Hello"}],
stream=False,
call_type="completion",
start_time=datetime.now(),
litellm_call_id="test-azure-router-opt-in",
function_id="test-fn",
)
kwargs = {
"model": "azure_ai/model-router",
"messages": [{"role": "user", "content": "Hello"}],
"response_cost": 0.00001,
"custom_llm_provider": "azure_ai",
}
mock_response = {
"id": "chatcmpl-azure-router-opt-in",
"object": "chat.completion",
"model": "azure_ai/gpt-5-nano-2025-08-07",
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hello"},
"finish_reason": "stop",
}
],
}
payload = get_standard_logging_object_payload(
kwargs=kwargs,
init_response_obj=mock_response,
start_time=datetime.now(),
end_time=datetime.now(),
logging_obj=logging_obj,
status="success",
)
assert payload is not None
assert payload["model"] == "azure_ai/gpt-5-nano-2025-08-07"
def test_standard_logging_payload_uses_actual_model_for_azure_router_with_underscore():
from litellm.litellm_core_utils.litellm_logging import (
Logging,
get_standard_logging_object_payload,
)
logging_obj = Logging(
model="azure_ai/model_router",
messages=[{"role": "user", "content": "Hello"}],
stream=False,
call_type="completion",
start_time=datetime.now(),
litellm_call_id="test-azure-router-underscore",
function_id="test-fn",
)
kwargs = {
"model": "azure_ai/model_router",
"messages": [{"role": "user", "content": "Hello"}],
"response_cost": 0.00001,
"custom_llm_provider": "azure_ai",
}
mock_response = {
"id": "chatcmpl-azure-router-underscore",
"object": "chat.completion",
"model": "azure_ai/gpt-5-nano-2025-08-07",
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hello"},
"finish_reason": "stop",
}
],
}
payload = get_standard_logging_object_payload(
kwargs=kwargs,
init_response_obj=mock_response,
start_time=datetime.now(),
end_time=datetime.now(),
logging_obj=logging_obj,
status="success",
)
assert payload is not None
assert payload["model"] == "azure_ai/gpt-5-nano-2025-08-07"
def test_merge_litellm_metadata_basic():
"""
Test that merge_litellm_metadata correctly merges metadata and litellm_metadata.

80
tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py
View file

@ -8,6 +8,9 @@ import pytest
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
from litellm.llms.azure_ai.azure_model_router.transformation import (
AzureModelRouterConfig,
)
from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig
@ -117,3 +120,80 @@ def test_azure_ai_grok_stop_parameter_handling():
# Test supported parameters for non-Grok models
gpt_params = config.get_supported_openai_params("gpt-4")
assert "stop" in gpt_params, "GPT models should support stop parameter"
def test_azure_model_router_response_shows_actual_model():
"""
Test that Azure Model Router returns the actual model used in the response,
not the router model.
According to the documentation, when using Azure Model Router, the response
should show the actual model that handled the request (e.g., gpt-5-nano-2025-08-07)
rather than the router model (e.g., model-router).
Regression test for: Azure Model Router should show actual model in response
"""
from httpx import Response
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
from litellm.types.utils import ModelResponse
config = AzureModelRouterConfig()
# Mock raw response from Azure that includes the actual model used
raw_response_json = {
"id": "chatcmpl-test123",
"object": "chat.completion",
"created": 1234567890,
"model": "gpt-5-nano-2025-08-07", # Actual model used by the router
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello!",
},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 5,
"total_tokens": 15,
},
}
# Create mock Response object
mock_response = MagicMock(spec=Response)
mock_response.json.return_value = raw_response_json
mock_response.text = json.dumps(raw_response_json)
mock_response.headers = {}
# Create ModelResponse object
model_response = ModelResponse()
# Create mock logging object with required methods
logging_obj = MagicMock(spec=LiteLLMLoggingObj)
logging_obj.post_call = MagicMock()
logging_obj.model_call_details = {}
# Call transform_response with router model
result = config.transform_response(
model="model-router", # This is the router model (without prefix)
raw_response=mock_response,
model_response=model_response,
logging_obj=logging_obj,
request_data={},
messages=[{"role": "user", "content": "Hello"}],
optional_params={},
litellm_params={"model": "azure_ai/model-router"}, # Original request model
encoding=None,
api_key="test-key",
json_mode=False,
)
# Verify that the response contains the actual model used, not the router model
assert result.model == "azure_ai/gpt-5-nano-2025-08-07", (
f"Expected model to be 'azure_ai/gpt-5-nano-2025-08-07' (actual model used), "
f"but got '{result.model}'"
)

79
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -15,6 +15,7 @@ from litellm.proxy.common_request_processing import (
ProxyConfig,
_extract_error_from_sse_chunk,
_get_cost_breakdown_from_logging_obj,
_is_azure_model_router_request,
_override_openai_response_model,
_parse_event_data_for_error,
create_response,
@ -1368,6 +1369,84 @@ class TestOverrideOpenAIResponseModel:
# Verify the model was not changed
assert response_obj.model == fallback_model
def test_override_model_preserves_azure_model_router_actual_model(self):
"""
Test that when the requested model is an Azure Model Router, the actual
model used (returned in the response) is preserved instead of being
overridden.
"""
requested_model = "azure_ai/model_router"
actual_model_used = "azure_ai/gpt-5-nano-2025-08-07"
response_obj = MagicMock()
response_obj.model = actual_model_used
response_obj._hidden_params = {"additional_headers": {}}
_override_openai_response_model(
response_obj=response_obj,
requested_model=requested_model,
log_context="test_context",
)
assert response_obj.model == actual_model_used
assert response_obj.model != requested_model
def test_override_model_preserves_azure_model_router_with_deployment_name(self):
"""
Test that Azure Model Router with deployment name pattern also preserves
the actual model used.
"""
requested_model = "azure_ai/model_router/my-deployment"
actual_model_used = "azure_ai/gpt-4.1-nano-2025-04-14"
response_obj = MagicMock()
response_obj.model = actual_model_used
response_obj._hidden_params = {"additional_headers": {}}
_override_openai_response_model(
response_obj=response_obj,
requested_model=requested_model,
log_context="test_context",
)
assert response_obj.model == actual_model_used
assert response_obj.model != requested_model
def test_override_model_preserves_azure_model_router_with_hyphen(self):
"""
Test that Azure Model Router with hyphen pattern (model-router) also preserves
the actual model used.
"""
requested_model = "azure_ai/model-router"
actual_model_used = "azure_ai/gpt-5-nano-2025-08-07"
response_obj = MagicMock()
response_obj.model = actual_model_used
response_obj._hidden_params = {"additional_headers": {}}
_override_openai_response_model(
response_obj=response_obj,
requested_model=requested_model,
log_context="test_context",
)
assert response_obj.model == actual_model_used
assert response_obj.model != requested_model
class TestIsAzureModelRouterRequest:
"""Tests for _is_azure_model_router_request helper"""
def test_detects_model_router_with_underscore(self):
assert _is_azure_model_router_request("azure_ai/model_router") is True
assert _is_azure_model_router_request("azure_ai/model_router/my-deployment") is True
def test_detects_model_router_with_hyphen(self):
assert _is_azure_model_router_request("azure_ai/model-router") is True
assert _is_azure_model_router_request("model-router") is True
def test_rejects_regular_models(self):
assert _is_azure_model_router_request("azure_ai/gpt-4") is False
assert _is_azure_model_router_request("gpt-4") is False
assert _is_azure_model_router_request("openai/gpt-3.5-turbo") is False
class TestStreamingOverheadHeader:
"""

64
tests/test_litellm/proxy/test_response_model_sanitization.py
View file

@ -23,7 +23,11 @@ def _initialize_proxy_with_config(config: dict, tmp_path) -> TestClient:
IMPORTANT: proxy_server.initialize() mutates module-level globals. We must call
cleanup_router_config_variables() before initializing to prevent cross-test bleed.
"""
from litellm.proxy.proxy_server import app, cleanup_router_config_variables, initialize
from litellm.proxy.proxy_server import (
app,
cleanup_router_config_variables,
initialize,
)
cleanup_router_config_variables()
@ -123,8 +127,8 @@ async def test_proxy_streaming_chunks_do_not_return_provider_prefixed_model(monk
client_model = "vllm-model"
internal_model = f"hosted_vllm/{client_model}"
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy import proxy_server
from litellm.proxy._types import UserAPIKeyAuth
# Patch proxy_logging_obj hooks so async_data_generator yields exactly our chunk.
async def _iterator_hook(
@ -176,8 +180,8 @@ async def test_proxy_streaming_chunks_use_client_requested_model_before_alias_ma
canonical_model = "vllm-model"
internal_model = f"hosted_vllm/{canonical_model}"
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy import proxy_server
from litellm.proxy._types import UserAPIKeyAuth
async def _iterator_hook(
user_api_key_dict: UserAPIKeyAuth,
@ -215,3 +219,57 @@ async def test_proxy_streaming_chunks_use_client_requested_model_before_alias_ma
payload = json.loads(first[len("data: ") :].strip())
assert payload["model"] == client_model_alias
assert not payload["model"].startswith("hosted_vllm/")
@pytest.mark.asyncio
async def test_proxy_streaming_azure_model_router_preserves_actual_model(monkeypatch):
"""
Regression test for Azure Model Router streaming:
When the client requests azure_ai/model_router, the streaming chunks should
preserve the actual model used (e.g., azure_ai/gpt-5-nano-2025-08-07) from
the downstream response, NOT override to the router model.
"""
router_model = "azure_ai/model_router"
actual_model_used = "azure_ai/gpt-5-nano-2025-08-07"
from litellm.proxy import proxy_server
from litellm.proxy._types import UserAPIKeyAuth
async def _iterator_hook(
user_api_key_dict: UserAPIKeyAuth,
response: AsyncGenerator,
request_data: dict,
):
yield _make_model_response_stream_chunk(model=actual_model_used)
monkeypatch.setattr(proxy_server.proxy_logging_obj, "async_post_call_streaming_iterator_hook", _iterator_hook)
monkeypatch.setattr(
proxy_server.proxy_logging_obj,
"async_post_call_streaming_hook",
AsyncMock(side_effect=lambda **kwargs: kwargs["response"]),
)
user_api_key_dict = UserAPIKeyAuth(api_key="sk-1234")
gen = proxy_server.async_data_generator(
response=MagicMock(),
user_api_key_dict=user_api_key_dict,
request_data={
"model": router_model,
"_litellm_client_requested_model": router_model,
},
)
chunks = []
async for item in gen:
chunks.append(item)
assert len(chunks) >= 2
first = chunks[0]
assert first.startswith("data: ")
payload = json.loads(first[len("data: ") :].strip())
# Azure Model Router: preserve actual model used, not the router model
assert payload["model"] == actual_model_used
assert payload["model"] != router_model