From 5b83aae71597e7cdc926742fe86aec2da3f91f16 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 12 Mar 2026 11:41:19 +0530 Subject: [PATCH] feat(azure_ai): show actual model used in Azure Model Router response - Azure Model Router transform_response: let parent extract actual model from raw response - common_request_processing: skip model override for Azure Model Router requests - proxy_server: skip streaming chunk model restamp for Azure Model Router - Add _is_azure_model_router_request helper - Add tests for non-streaming and streaming Made-with: Cursor --- .../azure_model_router/transformation.py | 18 ++-- litellm/proxy/common_request_processing.py | 47 +++++++++- litellm/proxy/proxy_server.py | 5 + .../chat/test_azure_ai_transformation.py | 80 ++++++++++++++++ .../proxy/test_common_request_processing.py | 94 +++++++++++++++++++ .../proxy/test_response_model_sanitization.py | 64 ++++++++++++- 6 files changed, 290 insertions(+), 18 deletions(-) diff --git a/litellm/llms/azure_ai/azure_model_router/transformation.py b/litellm/llms/azure_ai/azure_model_router/transformation.py index 3d6dc53c51..efda85f37b 100644 --- a/litellm/llms/azure_ai/azure_model_router/transformation.py +++ b/litellm/llms/azure_ai/azure_model_router/transformation.py @@ -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 diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index ce39ecf52d..07ea6a6043 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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/ + - azure_ai/model-router + - model_router/ + - 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: @@ -523,6 +556,10 @@ class ProxyBaseLLMRequestProcessing: "allm_passthrough_route", "avector_store_search", "avector_store_create", + "avector_store_retrieve", + "avector_store_list", + "avector_store_update", + "avector_store_delete", "avector_store_file_create", "avector_store_file_list", "avector_store_file_retrieve", @@ -756,6 +793,10 @@ class ProxyBaseLLMRequestProcessing: "allm_passthrough_route", "avector_store_search", "avector_store_create", + "avector_store_retrieve", + "avector_store_list", + "avector_store_update", + "avector_store_delete", "avector_store_file_create", "avector_store_file_list", "avector_store_file_retrieve", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e6bb3ee412..de0228cdec 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -289,6 +289,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 @@ -5426,6 +5427,10 @@ 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) ) 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 d903d7c85f..a26f7e7021 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 @@ -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}'" + ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index ba1084eafe..65489e93dd 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -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,99 @@ 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 with the router model. + + This ensures users can see which model actually handled their request + when using Azure Model Router. + """ + requested_model = "azure_ai/model_router" + actual_model_used = "azure_ai/gpt-5-nano-2025-08-07" + + # Create a mock object response with the actual model used + response_obj = MagicMock() + response_obj.model = actual_model_used + response_obj._hidden_params = {"additional_headers": {}} + + # Call the function - should preserve the actual model + _override_openai_response_model( + response_obj=response_obj, + requested_model=requested_model, + log_context="test_context", + ) + + # Verify the model was NOT overridden - should still be the actual model + 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" + + # Create a mock object response + response_obj = MagicMock() + response_obj.model = actual_model_used + response_obj._hidden_params = {"additional_headers": {}} + + # Call the function - should preserve the actual model + _override_openai_response_model( + response_obj=response_obj, + requested_model=requested_model, + log_context="test_context", + ) + + # Verify the model was NOT overridden + 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" + + # Create a mock object response + response_obj = MagicMock() + response_obj.model = actual_model_used + response_obj._hidden_params = {"additional_headers": {}} + + # Call the function - should preserve the actual model + _override_openai_response_model( + response_obj=response_obj, + requested_model=requested_model, + log_context="test_context", + ) + + # Verify the model was NOT overridden + 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: """ diff --git a/tests/test_litellm/proxy/test_response_model_sanitization.py b/tests/test_litellm/proxy/test_response_model_sanitization.py index b1bb8d0ed3..22785bbcb9 100644 --- a/tests/test_litellm/proxy/test_response_model_sanitization.py +++ b/tests/test_litellm/proxy/test_response_model_sanitization.py @@ -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