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 #37770 from BerriAI/litellm_fix_model_router_spend_log_model

Browse source 
fix(proxy): store the actual selected model in spend logs for Azure Model Router
This commit is contained in:
ryan-crabbe-berri 2026-08-25 12:35:35 -07:00 • committed by GitHub
commit 3db6c5ab18
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 344 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/litellm_core_utils/litellm_logging.py
View file

@ -5930,11 +5930,15 @@ def get_standard_logging_object_payload(
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.
# logging payload.
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
requested_model: Final = kwargs.get("model")
if (
isinstance(requested_model, str)
and ("model_router" in requested_model.lower() or "model-router" in requested_model.lower())
stamped_selected_model: Final = AzureFoundryModelInfo.get_model_router_selected_model(hidden_params)
if stamped_selected_model is not None:
model_name = stamped_selected_model
elif (
AzureFoundryModelInfo.is_model_router_call(model=requested_model, hidden_params=hidden_params)
and isinstance(response_model_name, str)
and response_model_name
):

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

@ -65,15 +65,24 @@ class AzureModelRouterConfig(AzureAIStudioConfig):
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.
Also stamps that model onto ``_hidden_params`` so downstream consumers (spend logs,
response restamping) can read it instead of guessing the route from the model string.
"""
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
from litellm.llms.azure_ai.common_utils import (
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
AzureFoundryModelInfo,
)
from litellm.router_utils.add_retry_fallback_headers import (
get_hidden_params_dict,
)
# Get base model for the parent call (strips routing prefixes for API compatibility)
base_model: Final[str] = AzureFoundryModelInfo.get_base_model(model)
# 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(
transformed_response: Final = super().transform_response(
model=base_model,
raw_response=raw_response,
model_response=model_response,
@ -86,7 +95,15 @@ class AzureModelRouterConfig(AzureAIStudioConfig):
api_key=api_key,
json_mode=json_mode,
)
return model_response
selected_model: Final = transformed_response.model
if selected_model:
# Rebuilt rather than mutated in place: ModelResponseBase declares _hidden_params as a
# class-level dict, so an in-place write can bleed into unrelated responses.
transformed_response._hidden_params = { # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params setter # mutable-ok: ModelResponse requires _hidden_params to be a plain dict
**get_hidden_params_dict(transformed_response),
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model,
}
return transformed_response
def calculate_additional_costs(self, model: str, prompt_tokens: int, completion_tokens: int) -> dict | None:
"""

38
litellm/llms/azure_ai/common_utils.py
View file

@ -51,6 +51,9 @@ def get_azure_ai_auth_headers(
)
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model"
class AzureFoundryModelInfo(BaseLLMModelInfo):
"""Model info for Azure AI / Azure Foundry models."""
@ -82,6 +85,41 @@ class AzureFoundryModelInfo(BaseLLMModelInfo):
return "model_router"
return "default"
@staticmethod
def get_model_router_selected_model(hidden_params: Mapping[str, object] | None) -> str | None:
"""The model Azure Model Router actually served, stamped by ``AzureModelRouterConfig``.
Reading this beats re-deriving the route from a model string: the stamp is set on the
code path that was actually taken, so it holds no matter what the caller named the model.
"""
if not hidden_params:
return None
selected: Final = hidden_params.get(AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY)
if isinstance(selected, str) and selected:
return selected
return None
@staticmethod
def is_model_router_call(
model: str | None = None,
hidden_params: Mapping[str, object] | None = None,
) -> bool:
"""Whether a request went down the Azure Model Router route.
Prefers the response stamp, then the deployment's litellm model path, and only then the
caller-supplied name. The last two go through ``get_azure_ai_route`` so the model-router
name heuristic lives in exactly one place.
"""
if AzureFoundryModelInfo.get_model_router_selected_model(hidden_params) is not None:
return True
deployment_model: Final = (
hidden_params.get("litellm_model_name") or hidden_params.get("model") if hidden_params is not None else None
)
return any(
isinstance(candidate, str) and AzureFoundryModelInfo.get_azure_ai_route(candidate) == "model_router"
for candidate in (deployment_model, model)
)
@staticmethod
def get_api_base(api_base: str | None = None) -> str | None:
return api_base or litellm.api_base or get_secret_str("AZURE_AI_API_BASE")

21
litellm/proxy/common_request_processing.py
View file

@ -1138,24 +1138,25 @@ async def open_sse_before_first_byte(
)
def _is_azure_model_router_request(model: str) -> bool:
def _is_azure_model_router_request(model: str, hidden_params: Mapping[str, object] | None = None) -> bool:
"""
Check if the requested model is an Azure Model Router.
Check if a request went down the Azure Model Router route.
Azure Model Router models follow the pattern:
- azure_ai/model_router/<deployment-name>
- azure_ai/model-router
- model_router/<deployment-name>
- model-router
``model`` here is what the *client* sent, a model group alias with no ``model_router/``
prefix, so matching on it alone only works when the operator happened to put "model-router"
in the alias. Where the response is in hand its stamp answers this outright, so callers
should pass ``hidden_params``.
Args:
model: The requested model name
hidden_params: ``_hidden_params`` from the response, when the caller has it
Returns:
bool: True if this is an Azure Model Router request
"""
model_lower: Final = model.lower()
return "model-router" in model_lower or "model_router" in model_lower
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
return AzureFoundryModelInfo.is_model_router_call(model=model, hidden_params=hidden_params)
def _override_openai_response_model(
@ -1223,7 +1224,7 @@ 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):
if _is_azure_model_router_request(requested_model, hidden_params):
verbose_proxy_logger.debug(
"%s: Azure Model Router detected - preserving actual model used from response instead of overriding to router model.",
log_context,

4
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -444,7 +444,9 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
or None
)
raw_model: Final = cast(str, kwargs.get("model") or "")
model_name: Final = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {})
model_name: Final = (
standard_logging_payload.get("model") if standard_logging_payload is not None else None
) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {})
try:
payload: Final[SpendLogsPayload] = SpendLogsPayload(

2
ruff-strict-budget.json
View file

@ -168,7 +168,7 @@
"limit": 3
},
"RET504": {
"limit": 176
"limit": 175
},
"RUF012": {
"limit": 240

85
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -6,7 +6,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import time
import httpx
@ -3867,6 +3866,90 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob
assert payload["litellm_call_id"] == call_id
# ── Azure Model Router selected-model attribution ────────────────────────────
def _model_router_response(selected_model: str, stamp: bool):
"""A ModelResponse as AzureModelRouterConfig hands it back, with or without the stamp."""
from litellm.llms.azure_ai.common_utils import (
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
)
from litellm.types.utils import ModelResponse
response = ModelResponse(model=selected_model)
response._hidden_params = (
{AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {}
)
return response
def test_standard_logging_payload_uses_stamped_model_router_model(logging_obj):
"""
The selected model must win off the stamp, not off "model-router" appearing in the
requested model. An operator whose model group is named anything else was invisible
to the name check, so their logs and spend rows named the router instead.
"""
import datetime
from litellm.litellm_core_utils.litellm_logging import (
get_standard_logging_object_payload,
)
now = datetime.datetime.now()
payload = get_standard_logging_object_payload(
kwargs={
"model": "azure_ai/smart-pick",
"custom_llm_provider": "azure_ai",
"messages": [],
"litellm_params": {"metadata": {}},
},
init_response_obj=_model_router_response(
"azure_ai/grok-4-1-fast-reasoning", stamp=True
),
start_time=now,
end_time=now,
logging_obj=logging_obj,
status="success",
)
assert payload is not None
assert payload["model"] == "azure_ai/grok-4-1-fast-reasoning"
def test_standard_logging_payload_keeps_requested_model_without_router_stamp(
logging_obj,
):
"""
Control for the test above: an ordinary azure_ai deployment is unaffected, so the stamp
is what redirects attribution rather than the response model winning unconditionally.
"""
import datetime
from litellm.litellm_core_utils.litellm_logging import (
get_standard_logging_object_payload,
)
now = datetime.datetime.now()
payload = get_standard_logging_object_payload(
kwargs={
"model": "azure_ai/smart-pick",
"custom_llm_provider": "azure_ai",
"messages": [],
"litellm_params": {"metadata": {}},
},
init_response_obj=_model_router_response(
"azure_ai/grok-4-1-fast-reasoning", stamp=False
),
start_time=now,
end_time=now,
logging_obj=logging_obj,
status="success",
)
assert payload is not None
assert payload["model"] == "azure_ai/smart-pick"
def _make_dict_logging_obj():
"""Build a Logging instance configured for a non-streaming dict result."""
obj = LitellmLogging(

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

@ -201,6 +201,90 @@ def test_azure_model_router_response_shows_actual_model():
)
def test_azure_model_router_stamps_selected_model_on_hidden_params():
"""
The selected model must be stamped on _hidden_params, not left for downstream code to
re-derive by looking for "model-router" in the model string. Deployments whose alias
does not contain that text are invisible to the string check.
"""
from httpx import Response
from litellm.llms.azure_ai.common_utils import (
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
AzureFoundryModelInfo,
)
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
from litellm.types.utils import ModelResponse
raw_response_json = {
"id": "chatcmpl-test456",
"object": "chat.completion",
"created": 1234567890,
"model": "grok-4-1-fast-reasoning",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "pong"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
}
mock_response = MagicMock(spec=Response)
mock_response.json.return_value = raw_response_json
mock_response.text = json.dumps(raw_response_json)
mock_response.headers = {}
logging_obj = MagicMock(spec=LiteLLMLoggingObj)
logging_obj.post_call = MagicMock()
logging_obj.model_call_details = {}
result = AzureModelRouterConfig().transform_response(
model="smart-pick",
raw_response=mock_response,
model_response=ModelResponse(),
logging_obj=logging_obj,
request_data={},
messages=[{"role": "user", "content": "Reply with just pong"}],
optional_params={},
litellm_params={"model": "azure_ai/model_router/smart-pick"},
encoding=None,
api_key="test-key",
json_mode=False,
)
assert result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] == result.model
assert (
result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY]
== "azure_ai/grok-4-1-fast-reasoning"
)
assert AzureFoundryModelInfo.get_model_router_selected_model(
result._hidden_params
) == ("azure_ai/grok-4-1-fast-reasoning")
assert (
AzureFoundryModelInfo.is_model_router_call(
model="smart-pick", hidden_params=result._hidden_params
)
is True
)
def test_azure_model_router_stamp_does_not_leak_across_responses():
"""
ModelResponse declares _hidden_params as a class-level dict, so the stamp has to be written
as a fresh dict. Mutating in place would bleed the selected model into unrelated responses.
"""
from litellm.llms.azure_ai.common_utils import (
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
)
from litellm.types.utils import ModelResponse
untouched = ModelResponse()
assert AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY not in (untouched._hidden_params or {})
def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name():
"""
Regression test: Azure AI returns 400 when tools contain copilot_mcp_server_name.

49
tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py
View file

@ -3,11 +3,10 @@ import datetime
import json
from datetime import timezone
from typing import Any, Final, cast
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm.constants import (
@ -3500,6 +3499,50 @@ def test_get_logging_payload_failed_request_without_standard_logging_payload_lea
assert payload["custom_llm_provider"] == ""
class _ModelRouterSpendLogKwargs(TypedDict):
model: ReadOnly[str]
litellm_params: ReadOnly[dict[str, dict[str, str]]]
standard_logging_object: ReadOnly[StandardLoggingPayload]
def _model_router_spend_log_kwargs(slp_model: str | None) -> _ModelRouterSpendLogKwargs:
standard_logging_payload: Final = cast(
StandardLoggingPayload,
{
"model": slp_model,
"metadata": {},
"model_map_information": StandardLoggingModelInformation(
model_map_key="azure_ai/model_router", model_map_value=None
),
},
)
return {
"model": "azure_ai/model_router/model-router",
"litellm_params": {"metadata": {"user_api_key": "sk-test-key"}},
"standard_logging_object": standard_logging_payload,
}
def test_get_logging_payload_uses_standard_logging_payload_model():
payload = get_logging_payload(
kwargs=_model_router_spend_log_kwargs(slp_model="azure_ai/gpt-5-mini"),
response_obj={},
start_time=datetime.datetime.now(timezone.utc),
end_time=datetime.datetime.now(timezone.utc),
)
assert payload["model"] == "azure_ai/gpt-5-mini"
def test_get_logging_payload_falls_back_to_kwargs_model_when_slp_model_missing():
payload = get_logging_payload(
kwargs=_model_router_spend_log_kwargs(slp_model=None),
response_obj={},
start_time=datetime.datetime.now(timezone.utc),
end_time=datetime.datetime.now(timezone.utc),
)
assert payload["model"] == "azure_ai/model_router/model-router"
@patch("litellm.proxy.proxy_server.master_key", None)
@patch("litellm.proxy.proxy_server.general_settings", {})
def test_get_logging_payload_empty_key_slp_none_is_empty_string_not_none_literal():

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

@ -2294,6 +2294,54 @@ class TestOverrideOpenAIResponseModel:
assert response_obj.model == actual_model_used
assert response_obj.model != requested_model
def test_override_model_preserves_model_router_model_for_alias_without_router_in_name(
self,
):
"""
The client sends a model group alias, which carries no model_router/ prefix, so the
name check alone only fires when the operator happened to put "model-router" in the
alias. With the stamp on the response the actual model survives whatever it is named.
"""
from litellm.llms.azure_ai.common_utils import (
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
)
requested_model = "smart-pick"
actual_model_used = "azure_ai/grok-4-1-fast-reasoning"
response_obj = MagicMock()
response_obj.model = actual_model_used
response_obj._hidden_params = {
"additional_headers": {},
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: actual_model_used,
}
_override_openai_response_model(
response_obj=response_obj,
requested_model=requested_model,
log_context="test_context",
)
assert response_obj.model == actual_model_used
def test_override_model_still_restamps_non_router_alias_without_stamp(self):
"""
Control for the test above: absent the stamp, an ordinary deployment keeps being
restamped to the requested model, so the stamp is doing the work rather than the
preserve branch having gone unconditional.
"""
requested_model = "smart-pick"
response_obj = MagicMock()
response_obj.model = "azure_ai/grok-4-1-fast-reasoning"
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 == requested_model
def test_override_model_uses_winning_model_for_fastest_response(self):
"""
Test that when fastest_response batch completion is used with a

2
type-discipline-budget.json
View file

@ -30,7 +30,7 @@
"limit": 16673
},
"LIT011": {
"limit": 5588
"limit": 5587
},
"LIT012": {
"limit": 4510