Add Azure Model Router flat cost to cost breakdown in logging payload

- Added azure_model_router_flat_cost field to CostBreakdown TypedDict
- Updated _store_cost_breakdown_in_logging_obj to accept and store the flat cost
- Updated set_cost_breakdown in LitellmLoggingObject to handle the flat cost
- Modified Azure AI cost calculator to track flat cost in thread-local storage
- Integrated flat cost extraction in completion_cost function for azure_ai provider
- Added comprehensive tests for cost breakdown tracking
- Flat cost is now visible in standard_logging_payload.cost_breakdown for observability

Co-authored-by: ishaan <ishaan@berri.ai>

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,

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

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)"

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)

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}")