Remove thread-local storage, use direct function calls for Azure Model Router flat cost

- Removed threading module and thread-local storage completely
- Changed get_azure_model_router_flat_cost to take model and usage as parameters
- Made _is_azure_model_router public (is_azure_model_router)
- Calculate flat cost directly where needed instead of storing in thread-local
- Updated tests to use new direct function call pattern
- Cleaner, simpler, and works correctly in async contexts

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

litellm/cost_calculator.py

@@ -1344,11 +1344,13 @@ def completion_cost(  # noqa: PLR0915
                 
                 # Get additional costs (e.g., Azure Model Router flat cost for azure_ai provider)
                 additional_costs: Optional[dict] = None
-                if custom_llm_provider == "azure_ai":
+                if custom_llm_provider == "azure_ai" and cost_per_token_usage_object:
                     from litellm.llms.azure_ai.cost_calculator import (
                         get_azure_model_router_flat_cost,
                     )
-                    azure_router_flat_cost = get_azure_model_router_flat_cost()
+                    azure_router_flat_cost = get_azure_model_router_flat_cost(
+                        model=model, usage=cost_per_token_usage_object
+                    )
                     if azure_router_flat_cost is not None and azure_router_flat_cost > 0:
                         additional_costs = {
                             "Azure Model Router Flat Cost": azure_router_flat_cost

litellm/llms/azure_ai/cost_calculator.py

@@ -4,7 +4,6 @@ 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
@@ -12,11 +11,7 @@ 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:
+def is_azure_model_router(model: str) -> bool:
     """
     Check if the model is Azure AI Foundry Model Router.
     
@@ -30,20 +25,35 @@ 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]:
+def get_azure_model_router_flat_cost(model: str, usage: Usage) -> Optional[float]:
     """
-    Get the most recently calculated Azure Model Router flat cost from thread-local storage.
+    Calculate the Azure Model Router flat cost if applicable.
+    
+    Args:
+        model: The model name
+        usage: Usage object with token counts
     
     Returns:
-        Optional[float]: The flat cost, or None if not available
+        Optional[float]: The flat cost, or None if not a model router
     """
-    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')
+    if not is_azure_model_router(model):
+        return None
+    
+    # Get the model router pricing from model_prices_and_context_window.json
+    model_info = get_model_info(model="azure-model-router", custom_llm_provider="azure_ai")
+    router_flat_cost_per_token = model_info.get("input_cost_per_token", 0)
+    
+    if router_flat_cost_per_token > 0:
+        router_flat_cost = usage.prompt_tokens * router_flat_cost_per_token
+        
+        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)"
+        )
+        
+        return router_flat_cost
+    
+    return None
 
 
 def cost_per_token(
@@ -55,7 +65,6 @@ 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
@@ -65,9 +74,6 @@ 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,
@@ -76,24 +82,8 @@ def cost_per_token(
     )
     
     # Add flat cost for Azure Model Router
-    # The flat cost is defined in model_prices_and_context_window.json for azure_ai/azure-model-router
-    if _is_azure_model_router(model):
-        # Get the model router pricing from model_prices_and_context_window.json
-        model_info = get_model_info(model="azure-model-router", custom_llm_provider="azure_ai")
-        router_flat_cost_per_token = model_info.get("input_cost_per_token", 0)
-        
-        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)"
-            )
-            
-            # Add flat cost to prompt cost
-            prompt_cost += router_flat_cost
+    router_flat_cost = get_azure_model_router_flat_cost(model, usage)
+    if router_flat_cost:
+        prompt_cost += router_flat_cost
     
     return prompt_cost, completion_cost

tests/test_litellm/llms/azure_ai/test_cost_calculator.py

@@ -4,8 +4,9 @@ Test Azure AI cost calculator, especially Model Router flat cost.
 
 import pytest
 from litellm.llms.azure_ai.cost_calculator import (
-    _is_azure_model_router,
+    is_azure_model_router,
     cost_per_token,
+    get_azure_model_router_flat_cost,
 )
 from litellm.types.utils import Usage
 from litellm.utils import get_model_info
@@ -33,7 +34,7 @@ class TestAzureModelRouterDetection:
     )
     def test_is_azure_model_router(self, model: str, expected: bool):
         """Test Azure Model Router detection."""
-        assert _is_azure_model_router(model) == expected
+        assert is_azure_model_router(model) == expected
 
 
 class TestAzureModelRouterFlatCost:
@@ -157,12 +158,8 @@ class TestAzureModelRouterFlatCost:
 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,
-        )
-
+    def test_flat_cost_calculation(self):
+        """Test that flat cost is calculated correctly."""
         model = "azure-model-router"
         usage = Usage(
             prompt_tokens=10000,
@@ -170,11 +167,8 @@ class TestAzureModelRouterCostBreakdown:
             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()
+        # Calculate the flat cost directly
+        flat_cost = get_azure_model_router_flat_cost(model=model, usage=usage)
 
         # Expected flat cost
         expected_flat_cost = (
@@ -183,7 +177,7 @@ class TestAzureModelRouterCostBreakdown:
 
         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}")
+        print(f"Flat cost calculated: ${flat_cost:.6f}")
 
     def test_flat_cost_integration_with_completion_cost(self):
         """Test that flat cost is properly integrated into completion_cost calculation."""