Add transform file

litellm/cost_calculator.py

@@ -141,6 +141,51 @@ def _cost_per_token_custom_pricing_helper(
     return None
 
 
+def _get_additional_costs(
+    model: str,
+    custom_llm_provider: Optional[str],
+    prompt_tokens: int,
+    completion_tokens: int,
+) -> Optional[dict]:
+    """
+    Calculate additional costs beyond standard token costs.
+    
+    This function delegates to provider-specific config classes to calculate
+    any additional costs like routing fees, infrastructure costs, etc.
+    
+    Args:
+        model: The model name
+        custom_llm_provider: The provider name (optional)
+        prompt_tokens: Number of prompt tokens
+        completion_tokens: Number of completion tokens
+        
+    Returns:
+        Optional dictionary with cost names and amounts, or None if no additional costs
+    """
+    if not custom_llm_provider:
+        return None
+        
+    try:
+        config_class = None
+        if custom_llm_provider == "azure_ai":
+            from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
+            config_class = AzureFoundryModelInfo.get_azure_ai_config_for_model(model)
+        # Add more providers here as needed
+        # elif custom_llm_provider == "other_provider":
+        #     config_class = get_other_provider_config(model)
+        
+        if config_class and hasattr(config_class, 'calculate_additional_costs'):
+            return config_class.calculate_additional_costs(
+                model=model,
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+            )
+    except Exception as e:
+        verbose_logger.debug(f"Error calculating additional costs: {e}")
+    
+    return None
+
+
 def _transcription_usage_has_token_details(
     usage_block: Optional[Usage],
 ) -> bool:
@@ -1342,19 +1387,13 @@ def completion_cost(  # noqa: PLR0915
                     response=completion_response,
                 )
                 
-                # 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":
-                    from litellm.llms.azure_ai.cost_calculator import (
-                        calculate_azure_model_router_flat_cost,
-                    )
-                    azure_router_flat_cost = calculate_azure_model_router_flat_cost(
-                        model=model, prompt_tokens=prompt_tokens
-                    )
-                    if azure_router_flat_cost > 0:
-                        additional_costs = {
-                            "Azure Model Router Flat Cost": azure_router_flat_cost
-                        }
+                # Get additional costs from provider (e.g., routing fees, infrastructure costs)
+                additional_costs = _get_additional_costs(
+                    model=model,
+                    custom_llm_provider=custom_llm_provider,
+                    prompt_tokens=prompt_tokens,
+                    completion_tokens=completion_tokens,
+                )
                 
                 _final_cost = (
                     prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar

litellm/llms/azure_ai/azure_model_router/__init__.py

@@ -0,0 +1,4 @@
+"""Azure AI Foundry Model Router support."""
+from .transformation import AzureModelRouterConfig
+
+__all__ = ["AzureModelRouterConfig"]

litellm/llms/azure_ai/azure_model_router/transformation.py

@@ -0,0 +1,125 @@
+"""
+Transformation for Azure AI Foundry Model Router.
+
+The Model Router is a special Azure AI deployment that automatically routes requests
+to the best available model. It has specific cost tracking requirements.
+"""
+from typing import Any, List, Optional
+
+from httpx import Response
+
+from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig
+from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
+from litellm.types.llms.openai import AllMessageValues
+from litellm.types.utils import ModelResponse
+
+
+class AzureModelRouterConfig(AzureAIStudioConfig):
+    """
+    Configuration for Azure AI Foundry Model Router.
+    
+    Handles:
+    - Stripping model_router prefix before sending to Azure API
+    - Preserving full model path in responses for cost tracking
+    - Calculating flat infrastructure costs for Model Router
+    """
+
+    def transform_request(
+        self,
+        model: str,
+        messages: List[AllMessageValues],
+        optional_params: dict,
+        litellm_params: dict,
+        headers: dict,
+    ) -> dict:
+        """
+        Transform request for Model Router.
+        
+        Strips the model_router/ prefix so only the deployment name is sent to Azure.
+        Example: model_router/azure-model-router -> azure-model-router
+        """
+        from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
+
+        # Get base model name (strips routing prefixes like model_router/)
+        base_model: str = AzureFoundryModelInfo.get_base_model(model)
+        
+        return super().transform_request(
+            base_model, messages, optional_params, litellm_params, headers
+        )
+
+    def transform_response(
+        self,
+        model: str,
+        raw_response: Response,
+        model_response: ModelResponse,
+        logging_obj: LiteLLMLoggingObj,
+        request_data: dict,
+        messages: List[AllMessageValues],
+        optional_params: dict,
+        litellm_params: dict,
+        encoding: Any,
+        api_key: Optional[str] = None,
+        json_mode: Optional[bool] = None,
+    ) -> ModelResponse:
+        """
+        Transform response for Model Router.
+        
+        Preserves the original model path (including model_router/ prefix) in the response
+        for proper cost tracking and logging.
+        """
+        from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
+
+        # Preserve the original model from litellm_params (includes routing prefixes like model_router/)
+        # This ensures cost tracking and logging use the full model path
+        original_model: str = litellm_params.get("model") or model
+        if not original_model.startswith("azure_ai/"):
+            # Add provider prefix if not already present
+            model_response.model = f"azure_ai/{original_model}"
+        else:
+            model_response.model = original_model
+        
+        # Get base model for the parent call (strips routing prefixes for API compatibility)
+        base_model: str = AzureFoundryModelInfo.get_base_model(model)
+        
+        return super().transform_response(
+            model=base_model,
+            raw_response=raw_response,
+            model_response=model_response,
+            logging_obj=logging_obj,
+            request_data=request_data,
+            messages=messages,
+            optional_params=optional_params,
+            litellm_params=litellm_params,
+            encoding=encoding,
+            api_key=api_key,
+            json_mode=json_mode,
+        )
+
+    def calculate_additional_costs(
+        self, model: str, prompt_tokens: int, completion_tokens: int
+    ) -> Optional[dict]:
+        """
+        Calculate additional costs for Azure Model Router.
+        
+        Adds a flat infrastructure cost of $0.14 per M input tokens for using the Model Router.
+        
+        Args:
+            model: The model name (should be a model router model)
+            prompt_tokens: Number of prompt tokens
+            completion_tokens: Number of completion tokens
+            
+        Returns:
+            Dictionary with additional costs, or None if not applicable.
+        """
+        from litellm.llms.azure_ai.cost_calculator import (
+            calculate_azure_model_router_flat_cost,
+        )
+        
+        flat_cost = calculate_azure_model_router_flat_cost(
+            model=model, prompt_tokens=prompt_tokens
+        )
+        
+        if flat_cost > 0:
+            return {"Azure Model Router Flat Cost": flat_cost}
+        
+        return None

litellm/llms/base_llm/chat/transformation.py

@@ -437,3 +437,23 @@ class BaseConfig(ABC):
         By default, this is true for almost all providers.
         """
         return True
+
+    def calculate_additional_costs(
+        self, model: str, prompt_tokens: int, completion_tokens: int
+    ) -> Optional[dict]:
+        """
+        Calculate any additional costs beyond standard token costs.
+        
+        This is used for provider-specific infrastructure costs, routing fees, etc.
+        
+        Args:
+            model: The model name
+            prompt_tokens: Number of prompt tokens
+            completion_tokens: Number of completion tokens
+            
+        Returns:
+            Optional dictionary with cost names and amounts, e.g.:
+            {"Infrastructure Fee": 0.001, "Routing Cost": 0.0005}
+            Returns None if no additional costs apply.
+        """
+        return None