scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

test_azure_ai_model_router

Browse source
This commit is contained in:
Ishaan Jaffer 2026-01-30 13:41:02 -08:00
parent 7f219ba18a
commit bb26f1e1e1
1 changed files with 27 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
tests/llm_translation/test_azure_ai.py
View file

@ -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/<deployment-name>
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