From bb26f1e1e177cdd7a30c0d4ed3d2f24d3aba6b78 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Fri, 30 Jan 2026 13:41:02 -0800 Subject: [PATCH] test_azure_ai_model_router --- tests/llm_translation/test_azure_ai.py | 28 +++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/tests/llm_translation/test_azure_ai.py b/tests/llm_translation/test_azure_ai.py index 1633fb4cc82..972ba34a179 100644 --- a/tests/llm_translation/test_azure_ai.py +++ b/tests/llm_translation/test_azure_ai.py @@ -373,10 +373,17 @@ def test_completion_azure_ai_gpt_4o_with_flexible_api_base(api_base): async def test_azure_ai_model_router(): """ Test Azure AI model router non-streaming response cost tracking. + Verifies that the flat cost of $0.14 per M input tokens is applied. + + Tests the pattern: azure_ai/model_router/ + Where deployment-name is the Azure deployment (e.g., "azure-model-router"). + The model_router prefix is stripped before sending to Azure API. """ + from litellm.llms.azure_ai.cost_calculator import calculate_azure_model_router_flat_cost + litellm._turn_on_debug() response = await litellm.acompletion( - model="azure_ai/azure-model-router", + model="azure_ai/model_router/azure-model-router", messages=[{"role": "user", "content": "hi who is this"}], api_base="https://ishaa-mh6uutut-swedencentral.cognitiveservices.azure.com/openai/v1/", api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"), @@ -387,6 +394,25 @@ async def test_azure_ai_model_router(): tracked_cost = response._hidden_params["response_cost"] assert tracked_cost > 0 print("Tracked cost: ", tracked_cost) + + # Verify flat cost is included using the helper function + usage = response.usage + if usage and usage.prompt_tokens: + expected_flat_cost = calculate_azure_model_router_flat_cost( + model="model_router/azure-model-router", + prompt_tokens=usage.prompt_tokens + ) + print(f"Prompt tokens: {usage.prompt_tokens}") + print(f"Expected flat cost: ${expected_flat_cost:.9f}") + print(f"Total tracked cost: ${tracked_cost:.9f}") + + # Total cost should be at least the flat cost + assert tracked_cost >= expected_flat_cost, ( + f"Cost ${tracked_cost:.9f} should be >= flat cost ${expected_flat_cost:.9f}" + ) + + # Verify the flat cost is non-zero + assert expected_flat_cost > 0, "Flat cost should be greater than 0" @pytest.mark.asyncio