diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5437ce52706..960ed266e75 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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 ): diff --git a/litellm/llms/azure_ai/azure_model_router/transformation.py b/litellm/llms/azure_ai/azure_model_router/transformation.py index 1a924088390..61cbc213b11 100644 --- a/litellm/llms/azure_ai/azure_model_router/transformation.py +++ b/litellm/llms/azure_ai/azure_model_router/transformation.py @@ -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: """ diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index d09055d7671..26a90157455 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -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") diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index dbbf9cb673e..413fead9bd4 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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/ - - azure_ai/model-router - - model_router/ - - 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, diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index b6f695db512..d931d712a92 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -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( diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 1ca152985f9..bbfbbbc1cfb 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -168,7 +168,7 @@ "limit": 3 }, "RET504": { - "limit": 176 + "limit": 175 }, "RUF012": { "limit": 240 diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 1cac0cce4ef..cc738d68016 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -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( diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index d4fcbc823a6..11a727c9635 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -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. diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 6c8e641642b..843e1d1296f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -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(): diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 58714a5e319..c9b3e5faf8b 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -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 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 8799e3e466a..7165bb53812 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -30,7 +30,7 @@ "limit": 16673 }, "LIT011": { - "limit": 5588 + "limit": 5587 }, "LIT012": { "limit": 4510