tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-13 11:07:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix azure model router

Browse Source
This commit is contained in:
Sameer Kankute
2026-03-12 12:40:37 +05:30
parent 5b83aae715
commit 36ec80d90c
5 changed files with 122 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/litellm_core_utils/litellm_logging.py
+14
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),
litellm/proxy/common_request_processing.py
-8
View File
@@ -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",
litellm/proxy/proxy_server.py
+3 -1
View File
@@ -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 = (
tests/logging_callback_tests/test_standard_logging_payload.py
+96
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.
tests/test_litellm/proxy/test_common_request_processing.py
+9 -24
View File
@@ -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