diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 71a93b757ce..bef4d52ce49 100644 --- a/litellm/cost_calculator.py +++ b/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 diff --git a/litellm/llms/azure_ai/azure_model_router/__init__.py b/litellm/llms/azure_ai/azure_model_router/__init__.py new file mode 100644 index 00000000000..0165d60b643 --- /dev/null +++ b/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"] diff --git a/litellm/llms/azure_ai/azure_model_router/transformation.py b/litellm/llms/azure_ai/azure_model_router/transformation.py new file mode 100644 index 00000000000..3d6dc53c515 --- /dev/null +++ b/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 diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 41a1797cebe..ac209904e6e 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/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