diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index a3ea03e6da0..bab9e9b1990 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -808,6 +808,7 @@ def _store_cost_breakdown_in_logging_obj( completion_tokens_cost_usd_dollar: float, cost_for_built_in_tools_cost_usd_dollar: float, total_cost_usd_dollar: float, + azure_model_router_flat_cost: Optional[float] = None, original_cost: Optional[float] = None, discount_percent: Optional[float] = None, discount_amount: Optional[float] = None, @@ -824,6 +825,7 @@ def _store_cost_breakdown_in_logging_obj( completion_tokens_cost_usd_dollar: Cost of completion tokens (includes reasoning if applicable) cost_for_built_in_tools_cost_usd_dollar: Cost of built-in tools total_cost_usd_dollar: Total cost of request + azure_model_router_flat_cost: Azure Model Router flat infrastructure cost original_cost: Cost before discount discount_percent: Discount percentage applied (0.05 = 5%) discount_amount: Discount amount in USD @@ -841,6 +843,7 @@ def _store_cost_breakdown_in_logging_obj( output_cost=completion_tokens_cost_usd_dollar, total_cost=total_cost_usd_dollar, cost_for_built_in_tools_cost_usd_dollar=cost_for_built_in_tools_cost_usd_dollar, + azure_model_router_flat_cost=azure_model_router_flat_cost, original_cost=original_cost, discount_percent=discount_percent, discount_amount=discount_amount, @@ -1338,6 +1341,15 @@ def completion_cost( # noqa: PLR0915 service_tier=service_tier, response=completion_response, ) + + # Get Azure Model Router flat cost if available (for azure_ai provider) + azure_model_router_flat_cost: Optional[float] = None + if custom_llm_provider == "azure_ai": + from litellm.llms.azure_ai.cost_calculator import ( + get_azure_model_router_flat_cost, + ) + azure_model_router_flat_cost = get_azure_model_router_flat_cost() + _final_cost = ( prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar ) @@ -1377,6 +1389,7 @@ def completion_cost( # noqa: PLR0915 completion_tokens_cost_usd_dollar=completion_tokens_cost_usd_dollar, cost_for_built_in_tools_cost_usd_dollar=cost_for_built_in_tools, total_cost_usd_dollar=_final_cost, + azure_model_router_flat_cost=azure_model_router_flat_cost, original_cost=original_cost, discount_percent=discount_percent, discount_amount=discount_amount, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index b2fa6065156..bdc3d77b530 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1297,6 +1297,7 @@ class Logging(LiteLLMLoggingBaseClass): output_cost: float, total_cost: float, cost_for_built_in_tools_cost_usd_dollar: float, + azure_model_router_flat_cost: Optional[float] = None, original_cost: Optional[float] = None, discount_percent: Optional[float] = None, discount_amount: Optional[float] = None, @@ -1312,6 +1313,7 @@ class Logging(LiteLLMLoggingBaseClass): output_cost: Cost of output/completion tokens cost_for_built_in_tools_cost_usd_dollar: Cost of built-in tools total_cost: Total cost of request + azure_model_router_flat_cost: Azure Model Router flat infrastructure cost original_cost: Cost before discount discount_percent: Discount percentage (0.05 = 5%) discount_amount: Discount amount in USD @@ -1327,6 +1329,10 @@ class Logging(LiteLLMLoggingBaseClass): tool_usage_cost=cost_for_built_in_tools_cost_usd_dollar, ) + # Store Azure Model Router flat cost if provided + if azure_model_router_flat_cost is not None and azure_model_router_flat_cost > 0: + self.cost_breakdown["azure_model_router_flat_cost"] = azure_model_router_flat_cost + # Store discount information if provided if original_cost is not None: self.cost_breakdown["original_cost"] = original_cost diff --git a/litellm/llms/azure_ai/cost_calculator.py b/litellm/llms/azure_ai/cost_calculator.py index 7d502b3606f..ca0f2a8cc05 100644 --- a/litellm/llms/azure_ai/cost_calculator.py +++ b/litellm/llms/azure_ai/cost_calculator.py @@ -4,6 +4,7 @@ Handles Azure AI Foundry Model Router flat cost and other Azure AI specific pric """ from typing import Optional, Tuple +import threading from litellm._logging import verbose_logger from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token @@ -11,6 +12,10 @@ from litellm.types.utils import Usage from litellm.utils import get_model_info +# Thread-local storage for tracking Azure Model Router flat cost +_thread_local = threading.local() + + def _is_azure_model_router(model: str) -> bool: """ Check if the model is Azure AI Foundry Model Router. @@ -25,6 +30,22 @@ def _is_azure_model_router(model: str) -> bool: return "model-router" in model_lower or model_lower == "azure-model-router" +def get_azure_model_router_flat_cost() -> Optional[float]: + """ + Get the most recently calculated Azure Model Router flat cost from thread-local storage. + + Returns: + Optional[float]: The flat cost, or None if not available + """ + return getattr(_thread_local, 'azure_model_router_flat_cost', None) + + +def _clear_azure_model_router_flat_cost() -> None: + """Clear the thread-local Azure Model Router flat cost.""" + if hasattr(_thread_local, 'azure_model_router_flat_cost'): + delattr(_thread_local, 'azure_model_router_flat_cost') + + def cost_per_token( model: str, usage: Usage, response_time_ms: Optional[float] = 0.0 ) -> Tuple[float, float]: @@ -34,6 +55,7 @@ def cost_per_token( For Azure AI Foundry Model Router: - Adds a flat cost of $0.14 per million input tokens (from model_prices_and_context_window.json) - Plus the cost of the actual model used (handled by generic_cost_per_token) + - Stores the flat cost in thread-local storage for cost breakdown tracking Args: model: str, the model name without provider prefix @@ -43,6 +65,9 @@ def cost_per_token( Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd """ + # Clear any previous flat cost + _clear_azure_model_router_flat_cost() + # Calculate base cost using generic cost calculator prompt_cost, completion_cost = generic_cost_per_token( model=model, @@ -60,6 +85,9 @@ def cost_per_token( if router_flat_cost_per_token > 0: router_flat_cost = usage.prompt_tokens * router_flat_cost_per_token + # Store the flat cost in thread-local storage for cost breakdown + _thread_local.azure_model_router_flat_cost = router_flat_cost + verbose_logger.debug( f"Azure AI Model Router flat cost: ${router_flat_cost:.6f} " f"({usage.prompt_tokens} tokens × ${router_flat_cost_per_token:.9f}/token)" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 49c903502d1..889c21247e6 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2635,6 +2635,7 @@ class CostBreakdown(TypedDict, total=False): ) total_cost: float # Total cost (input + output + tool usage) tool_usage_cost: float # Cost of usage of built-in tools + azure_model_router_flat_cost: float # Azure AI Foundry Model Router flat cost ($0.14 per M input tokens) original_cost: float # Cost before discount (optional) discount_percent: float # Discount percentage applied (e.g., 0.05 = 5%) (optional) discount_amount: float # Discount amount in USD (optional) diff --git a/tests/test_litellm/llms/azure_ai/test_cost_calculator.py b/tests/test_litellm/llms/azure_ai/test_cost_calculator.py index 44dafae8848..9bc0e90c7eb 100644 --- a/tests/test_litellm/llms/azure_ai/test_cost_calculator.py +++ b/tests/test_litellm/llms/azure_ai/test_cost_calculator.py @@ -152,3 +152,84 @@ class TestAzureModelRouterFlatCost: f"Model Router flat cost with caching for {usage.prompt_tokens} tokens: ${expected_flat_cost:.6f}" ) print(f"Total prompt cost: ${prompt_cost:.6f}") + + +class TestAzureModelRouterCostBreakdown: + """Test that Azure Model Router flat cost is tracked in cost breakdown.""" + + def test_flat_cost_tracked_in_thread_local(self): + """Test that flat cost is stored in thread-local storage.""" + from litellm.llms.azure_ai.cost_calculator import ( + get_azure_model_router_flat_cost, + ) + + model = "azure-model-router" + usage = Usage( + prompt_tokens=10000, + completion_tokens=5000, + total_tokens=15000, + ) + + # Calculate cost (this should store flat cost in thread-local) + prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) + + # Retrieve the flat cost from thread-local storage + flat_cost = get_azure_model_router_flat_cost() + + # Expected flat cost + expected_flat_cost = ( + usage.prompt_tokens * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 + ) + + assert flat_cost is not None + assert flat_cost == pytest.approx(expected_flat_cost, rel=1e-9) + print(f"Flat cost tracked in thread-local: ${flat_cost:.6f}") + + def test_flat_cost_integration_with_completion_cost(self): + """Test that flat cost is properly integrated into completion_cost calculation.""" + import litellm + from litellm.cost_calculator import completion_cost + from litellm.types.utils import Usage, ModelResponse, Choices, Message + + # Create a mock response for azure_ai model router + response = ModelResponse( + id="test-123", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + role="assistant", + content="Test response", + ), + ) + ], + created=1234567890, + model="azure-model-router", + object="chat.completion", + usage=Usage( + prompt_tokens=5000, + completion_tokens=2000, + total_tokens=7000, + ), + ) + + # Set hidden params for provider + response._hidden_params = {"custom_llm_provider": "azure_ai"} + + # Calculate cost + cost = completion_cost( + completion_response=response, + model="azure-model-router", + custom_llm_provider="azure_ai", + ) + + # Expected flat cost + expected_flat_cost = ( + 5000 * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 + ) + + # Cost should include the flat cost + assert cost > expected_flat_cost + print(f"Total cost with flat fee: ${cost:.6f}") + print(f"Expected minimum flat cost: ${expected_flat_cost:.6f}")