diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 8a68d74be5b..4d7b44bfbde 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -513,7 +513,8 @@ def cost_per_token( # noqa: PLR0915 return fireworks_ai_cost_per_token(model=model, usage=usage_block) elif custom_llm_provider == "azure": return azure_openai_cost_per_token( - model=model, usage=usage_block, response_time_ms=response_time_ms + model=model, usage=usage_block, response_time_ms=response_time_ms, + service_tier=service_tier, ) elif custom_llm_provider == "gemini": return gemini_cost_per_token( @@ -539,6 +540,7 @@ def cost_per_token( # noqa: PLR0915 usage=usage_block, response_time_ms=response_time_ms, request_model=request_model, + service_tier=service_tier, ) else: model_info = _cached_get_model_info_helper( diff --git a/litellm/llms/azure/cost_calculation.py b/litellm/llms/azure/cost_calculation.py index 5b411095ea1..2e2e0c4965f 100644 --- a/litellm/llms/azure/cost_calculation.py +++ b/litellm/llms/azure/cost_calculation.py @@ -12,7 +12,8 @@ from litellm.utils import get_model_info def cost_per_token( - model: str, usage: Usage, response_time_ms: Optional[float] = 0.0 + model: str, usage: Usage, response_time_ms: Optional[float] = 0.0, + service_tier: Optional[str] = None, ) -> Tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -47,4 +48,5 @@ def cost_per_token( model=model, usage=usage, custom_llm_provider="azure", + service_tier=service_tier, ) diff --git a/litellm/llms/azure_ai/cost_calculator.py b/litellm/llms/azure_ai/cost_calculator.py index 067181b946a..755d44fdef7 100644 --- a/litellm/llms/azure_ai/cost_calculator.py +++ b/litellm/llms/azure_ai/cost_calculator.py @@ -65,6 +65,7 @@ def cost_per_token( usage: Usage, response_time_ms: Optional[float] = 0.0, request_model: Optional[str] = None, + service_tier: Optional[str] = None, ) -> Tuple[float, float]: """ Calculate the cost per token for Azure AI models. @@ -102,6 +103,7 @@ def cost_per_token( model=model, usage=usage, custom_llm_provider="azure_ai", + service_tier=service_tier, ) except Exception as e: # For Model Router, the model name (e.g., "azure-model-router") may not be in the cost map diff --git a/tests/test_litellm/llms/azure/test_azure_cost_calculation.py b/tests/test_litellm/llms/azure/test_azure_cost_calculation.py new file mode 100644 index 00000000000..53c91032b34 --- /dev/null +++ b/tests/test_litellm/llms/azure/test_azure_cost_calculation.py @@ -0,0 +1,75 @@ +""" +Test Azure OpenAI cost calculator — service_tier pricing. +""" + +import pytest + +import litellm +from litellm.llms.azure.cost_calculation import cost_per_token +from litellm.types.utils import Usage + + +# Register a test model with tier-specific pricing +TEST_MODEL = "test-azure-gpt-4.1" +TEST_MODEL_COST = { + TEST_MODEL: { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "input_cost_per_token_priority": 0.01, + "output_cost_per_token_priority": 0.02, + "input_cost_per_token_flex": 0.0005, + "output_cost_per_token_flex": 0.001, + "litellm_provider": "azure", + "max_tokens": 8192, + } +} + + +class TestAzureServiceTierCostCalculation: + """Test that service_tier is passed through Azure cost calculation.""" + + @pytest.fixture(autouse=True) + def register_test_model(self): + litellm.register_model(model_cost=TEST_MODEL_COST) + + def test_service_tier_priority_higher_cost(self): + """Priority tier should cost more than standard.""" + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + standard_prompt, standard_completion = cost_per_token( + model=TEST_MODEL, usage=usage + ) + priority_prompt, priority_completion = cost_per_token( + model=TEST_MODEL, usage=usage, service_tier="priority" + ) + + assert priority_prompt > standard_prompt + assert priority_completion > standard_completion + + def test_service_tier_flex_lower_cost(self): + """Flex tier should cost less than standard.""" + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + standard_prompt, standard_completion = cost_per_token( + model=TEST_MODEL, usage=usage + ) + flex_prompt, flex_completion = cost_per_token( + model=TEST_MODEL, usage=usage, service_tier="flex" + ) + + assert flex_prompt < standard_prompt + assert flex_completion < standard_completion + + def test_service_tier_none_returns_standard(self): + """service_tier=None should return standard pricing.""" + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + none_prompt, none_completion = cost_per_token( + model=TEST_MODEL, usage=usage, service_tier=None + ) + standard_prompt, standard_completion = cost_per_token( + model=TEST_MODEL, usage=usage, service_tier="standard" + ) + + assert abs(none_prompt - standard_prompt) < 1e-10 + assert abs(none_completion - standard_completion) < 1e-10 diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py index 37add41b83f..20260c744f8 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py @@ -451,3 +451,51 @@ class TestAzureModelRouterCostBreakdown: assert logging_obj.cost_breakdown["additional_costs"][ "Azure Model Router Flat Cost" ] == pytest.approx(expected_flat_cost, rel=1e-9) + + +class TestAzureAIServiceTierCostCalculation: + """Test that service_tier is passed through Azure AI cost calculation.""" + + @pytest.fixture(autouse=True) + def register_test_model(self): + import litellm + litellm.register_model(model_cost={ + "test-azure-ai-model": { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "input_cost_per_token_priority": 0.01, + "output_cost_per_token_priority": 0.02, + "input_cost_per_token_flex": 0.0005, + "output_cost_per_token_flex": 0.001, + "litellm_provider": "azure_ai", + "max_tokens": 8192, + } + }) + + def test_service_tier_priority_higher_cost(self): + """Priority tier should cost more than standard for azure_ai.""" + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + standard_prompt, standard_completion = cost_per_token( + model="test-azure-ai-model", usage=usage + ) + priority_prompt, priority_completion = cost_per_token( + model="test-azure-ai-model", usage=usage, service_tier="priority" + ) + + assert priority_prompt > standard_prompt + assert priority_completion > standard_completion + + def test_service_tier_flex_lower_cost(self): + """Flex tier should cost less than standard for azure_ai.""" + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + standard_prompt, standard_completion = cost_per_token( + model="test-azure-ai-model", usage=usage + ) + flex_prompt, flex_completion = cost_per_token( + model="test-azure-ai-model", usage=usage, service_tier="flex" + ) + + assert flex_prompt < standard_prompt + assert flex_completion < standard_completion