From 36ec80d90c4a9e156dfd60c8357956744df48a31 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 12 Mar 2026 12:40:37 +0530 Subject: [PATCH] Fix azure model router --- litellm/litellm_core_utils/litellm_logging.py | 14 +++ litellm/proxy/common_request_processing.py | 8 -- litellm/proxy/proxy_server.py | 4 +- .../test_standard_logging_payload.py | 96 +++++++++++++++++++ .../proxy/test_common_request_processing.py | 33 ++----- 5 files changed, 122 insertions(+), 33 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 6f587abcdf..f9c5d74ee3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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), diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 07ea6a6043..da8f285504 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -556,10 +556,6 @@ 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", @@ -793,10 +789,6 @@ 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 de0228cdec..39a0887597 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5428,7 +5428,9 @@ def _restamp_streaming_chunk_model( 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): + if _is_azure_model_router_request( + requested_model_from_client + ): return chunk, model_mismatch_logged downstream_model = ( diff --git a/tests/logging_callback_tests/test_standard_logging_payload.py b/tests/logging_callback_tests/test_standard_logging_payload.py index b725d077e6..163e2d9435 100644 --- a/tests/logging_callback_tests/test_standard_logging_payload.py +++ b/tests/logging_callback_tests/test_standard_logging_payload.py @@ -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. diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 65489e93dd..3869a24d35 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1371,29 +1371,22 @@ class TestOverrideOpenAIResponseModel: 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. + 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" - - # 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 @@ -1404,20 +1397,16 @@ class TestOverrideOpenAIResponseModel: """ 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 @@ -1428,20 +1417,16 @@ class TestOverrideOpenAIResponseModel: """ 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