tests azure model router

tests/test_litellm/test_cost_calculator.py

@@ -14,6 +14,7 @@ from pydantic import BaseModel
 
 import litellm
 from litellm.cost_calculator import (
+    completion_cost,
     handle_realtime_stream_cost_calculation,
     response_cost_calculator,
 )
@@ -22,6 +23,33 @@ from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage
 from litellm.utils import TranscriptionResponse
 
 
+def test_completion_cost_uses_response_model_for_dynamic_routing():
+    """
+    Test that completion_cost uses the model from the response object
+    when the input model (e.g., azure-model-router) is not in model_cost.
+    This supports Azure Model Router and similar dynamic routing scenarios.
+    """
+    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
+    litellm.model_cost = litellm.get_model_cost_map(url="")
+
+    # Simulate Azure Model Router: input is generic router, response has actual model
+    response = ModelResponse(
+        id="test-id",
+        model="azure_ai/gpt-4o-2024-08-06",  # Response contains actual model used
+        choices=[],
+        usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
+    )
+
+    # Should calculate cost using the response model, not the input model
+    cost = completion_cost(
+        completion_response=response,
+        model="azure_ai/azure-model-router",  # Input model doesn't exist in model_cost
+        custom_llm_provider="azure_ai",
+    )
+
+    assert cost > 0, "Cost should be calculated using response model"
+
+
 def test_cost_calculator_with_response_cost_in_additional_headers():
     class MockResponse(BaseModel):
         _hidden_params = {

tests/test_litellm/test_utils.py

@@ -21,6 +21,7 @@ from litellm.types.utils import (
 from litellm.utils import (
     ProviderConfigManager,
     TextCompletionStreamWrapper,
+    _check_provider_match,
     get_llm_provider,
     get_optional_params_image_gen,
     is_cached_message,
@@ -29,6 +30,30 @@ from litellm.utils import (
 # Adds the parent directory to the system path
 
 
+def test_check_provider_match_azure_ai_allows_openai_and_azure():
+    """
+    Test that azure_ai provider can match openai and azure models.
+    This is needed for Azure Model Router which can route to OpenAI models.
+    """
+    # azure_ai should match openai models
+    assert _check_provider_match(
+        model_info={"litellm_provider": "openai"},
+        custom_llm_provider="azure_ai"
+    ) is True
+
+    # azure_ai should match azure models
+    assert _check_provider_match(
+        model_info={"litellm_provider": "azure"},
+        custom_llm_provider="azure_ai"
+    ) is True
+
+    # azure_ai should NOT match other providers
+    assert _check_provider_match(
+        model_info={"litellm_provider": "anthropic"},
+        custom_llm_provider="azure_ai"
+    ) is False
+
+
 def test_get_optional_params_image_gen():
     from litellm.llms.azure.image_generation import AzureGPTImageGenerationConfig