diff --git a/docs/my-website/docs/providers/azure_ai/azure_model_router.md b/docs/my-website/docs/providers/azure_ai/azure_model_router.md index 5e14c7283f6..16bc1afb70e 100644 --- a/docs/my-website/docs/providers/azure_ai/azure_model_router.md +++ b/docs/my-website/docs/providers/azure_ai/azure_model_router.md @@ -5,19 +5,38 @@ Azure Model Router is a feature in Azure AI Foundry that automatically routes yo ## Key Features - **Automatic Model Selection**: Azure Model Router dynamically selects the best model for your request -- **Cost Tracking**: LiteLLM automatically tracks costs based on the actual model used (e.g., `gpt-4.1-nano`), not the router endpoint +- **Cost Tracking**: LiteLLM automatically tracks costs based on the actual model used (e.g., `gpt-4.1-nano`), plus the Model Router infrastructure fee - **Streaming Support**: Full support for streaming responses with accurate cost calculation +- **Simple Configuration**: Easy to set up via UI or config file + +## Model Naming Pattern + +Use the pattern: `azure_ai/model_router/` + +**Components:** +- `azure_ai` - The provider identifier +- `model_router` - Indicates this is a Model Router deployment +- `` - Your actual deployment name from Azure AI Foundry (e.g., `azure-model-router`) + +**Example:** `azure_ai/model_router/azure-model-router` + +**How it works:** +- LiteLLM automatically strips the `model_router/` prefix when sending requests to Azure +- Only your deployment name (e.g., `azure-model-router`) is sent to the Azure API +- The full path is preserved in responses and logs for proper cost tracking ## LiteLLM Python SDK ### Basic Usage +Use the pattern `azure_ai/model_router/` where `` is your Azure deployment name: + ```python import litellm import os response = litellm.completion( - model="azure_ai/azure-model-router", + model="azure_ai/model_router/azure-model-router", # Use your deployment name messages=[{"role": "user", "content": "Hello!"}], api_base="https://your-endpoint.cognitiveservices.azure.com/openai/v1/", api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"), @@ -26,6 +45,13 @@ response = litellm.completion( print(response) ``` +**Pattern Explanation:** +- `azure_ai` - The provider +- `model_router` - Indicates this is a model router deployment +- `azure-model-router` - Your actual deployment name from Azure AI Foundry + +LiteLLM will automatically strip the `model_router/` prefix when sending the request to Azure, so only `azure-model-router` is sent to the API. + ### Streaming with Usage Tracking ```python @@ -33,7 +59,7 @@ import litellm import os response = await litellm.acompletion( - model="azure_ai/azure-model-router", + model="azure_ai/model_router/azure-model-router", # Use your deployment name messages=[{"role": "user", "content": "hi"}], api_base="https://your-endpoint.cognitiveservices.azure.com/openai/v1/", api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"), @@ -51,13 +77,15 @@ async for chunk in response: ```yaml model_list: - - model_name: azure-model-router + - model_name: azure-model-router # Public name for your users litellm_params: - model: azure_ai/azure-model-router + model: azure_ai/model_router/azure-model-router # Use your deployment name api_base: https://your-endpoint.cognitiveservices.azure.com/openai/v1/ api_key: os.environ/AZURE_MODEL_ROUTER_API_KEY ``` +**Note:** Replace `azure-model-router` in the model path with your actual deployment name from Azure AI Foundry. + ### Start Proxy ```bash @@ -80,49 +108,42 @@ curl -X POST http://localhost:4000/chat/completions \ This walkthrough shows how to add an Azure Model Router endpoint to LiteLLM using the Admin Dashboard. -### Select Provider +### Quick Start + +1. Navigate to the **Models** page in the LiteLLM UI +2. Select **"Azure AI Foundry (Studio)"** as the provider +3. Enter your deployment name (e.g., `azure-model-router`) +4. LiteLLM will automatically format it as `azure_ai/model_router/azure-model-router` +5. Add your API base URL and API key +6. Test and save + +### Detailed Walkthrough + +#### Step 1: Select Provider Navigate to the Models page and select "Azure AI Foundry (Studio)" as the provider. -#### Navigate to Models Page +##### Navigate to Models Page ![Navigate to Models](./img/azure_model_router_01.jpeg) -#### Click Provider Dropdown +##### Click Provider Dropdown ![Click Provider](./img/azure_model_router_02.jpeg) -#### Choose Azure AI Foundry +##### Choose Azure AI Foundry ![Select Azure AI Foundry](./img/azure_model_router_03.jpeg) -### Configure Model Name +#### Step 2: Enter Deployment Name -Set up the model name by entering `azure_ai/` followed by your model router deployment name from Azure. +**New Simplified Method:** Just enter your deployment name directly in the text field. If your deployment name contains "model-router" or "model_router", LiteLLM will automatically format it as `azure_ai/model_router/`. -#### Click Model Name Field +**Example:** +- Enter: `azure-model-router` +- LiteLLM creates: `azure_ai/model_router/azure-model-router` -![Click Model Field](./img/azure_model_router_04.jpeg) - -#### Select Custom Model Name - -![Select Custom Model](./img/azure_model_router_05.jpeg) - -#### Enter LiteLLM Model Name - -![LiteLLM Model Name](./img/azure_model_router_06.jpeg) - -#### Click Custom Model Name Field - -![Enter Custom Name Field](./img/azure_model_router_07.jpeg) - -#### Type Model Prefix - -Type `azure_ai/` as the prefix. - -![Type azure_ai prefix](./img/azure_model_router_08.jpeg) - -#### Copy Model Name from Azure Portal +##### Copy Deployment Name from Azure Portal Switch to Azure AI Foundry and copy your model router deployment name. @@ -130,73 +151,79 @@ Switch to Azure AI Foundry and copy your model router deployment name. ![Copy Model Name](./img/azure_model_router_10.jpeg) -#### Paste Model Name +##### Enter Deployment Name in LiteLLM -Paste to get `azure_ai/azure-model-router`. +Paste your deployment name (e.g., `azure-model-router`) directly into the text field. -![Paste Model Name](./img/azure_model_router_11.jpeg) +![Enter Deployment Name](./img/azure_model_router_04.jpeg) -### Configure API Base and Key +**What happens behind the scenes:** +- You enter: `azure-model-router` +- LiteLLM automatically detects this is a model router deployment +- The full model path becomes: `azure_ai/model_router/azure-model-router` +- When making API calls, only `azure-model-router` is sent to Azure + +#### Step 3: Configure API Base and Key Copy the endpoint URL and API key from Azure portal. -#### Copy API Base URL from Azure +##### Copy API Base URL from Azure ![Copy API Base](./img/azure_model_router_12.jpeg) -#### Enter API Base in LiteLLM +##### Enter API Base in LiteLLM ![Click API Base Field](./img/azure_model_router_13.jpeg) ![Paste API Base](./img/azure_model_router_14.jpeg) -#### Copy API Key from Azure +##### Copy API Key from Azure ![Copy API Key](./img/azure_model_router_15.jpeg) -#### Enter API Key in LiteLLM +##### Enter API Key in LiteLLM ![Enter API Key](./img/azure_model_router_16.jpeg) -### Test and Add Model +#### Step 4: Test and Add Model Verify your configuration works and save the model. -#### Test Connection +##### Test Connection ![Test Connection](./img/azure_model_router_17.jpeg) -#### Close Test Dialog +##### Close Test Dialog ![Close Dialog](./img/azure_model_router_18.jpeg) -#### Add Model +##### Add Model ![Add Model](./img/azure_model_router_19.jpeg) -### Verify in Playground +#### Step 5: Verify in Playground Test your model and verify cost tracking is working. -#### Open Playground +##### Open Playground ![Go to Playground](./img/azure_model_router_20.jpeg) -#### Select Model +##### Select Model ![Select Model](./img/azure_model_router_21.jpeg) -#### Send Test Message +##### Send Test Message ![Send Message](./img/azure_model_router_22.jpeg) -#### View Logs +##### View Logs ![View Logs](./img/azure_model_router_23.jpeg) -#### Verify Cost Tracking +##### Verify Cost Tracking -Cost is tracked based on the actual model used (e.g., `gpt-4.1-nano`). +Cost is tracked based on the actual model used (e.g., `gpt-4.1-nano`), plus a flat infrastructure cost of $0.14 per million input tokens for using the Model Router. ![Verify Cost](./img/azure_model_router_24.jpeg) @@ -205,28 +232,50 @@ Cost is tracked based on the actual model used (e.g., `gpt-4.1-nano`). LiteLLM automatically handles cost tracking for Azure Model Router by: 1. **Detecting the actual model**: When Azure Model Router routes your request to a specific model (e.g., `gpt-4.1-nano-2025-04-14`), LiteLLM extracts this from the response -2. **Calculating accurate costs**: Costs are calculated based on the actual model used, not the router endpoint name +2. **Calculating accurate costs**: Costs are calculated based on: + - The actual model used (e.g., `gpt-4.1-nano` token costs) + - Plus a flat infrastructure cost of **$0.14 per million input tokens** for using the Model Router 3. **Streaming support**: Cost tracking works correctly for both streaming and non-streaming requests +### Cost Breakdown + +When you use Azure Model Router, the total cost includes: + +- **Model Cost**: Based on the actual model that handled your request (e.g., `gpt-4.1-nano`) +- **Router Flat Cost**: $0.14 per million input tokens (Azure AI Foundry infrastructure fee) + ### Example Response with Cost ```python import litellm response = litellm.completion( - model="azure_ai/azure-model-router", + model="azure_ai/model_router/azure-model-router", messages=[{"role": "user", "content": "Hello!"}], api_base="https://your-endpoint.cognitiveservices.azure.com/openai/v1/", api_key="your-api-key", ) # The response will show the actual model used -print(f"Model used: {response.model}") # e.g., "gpt-4.1-nano-2025-04-14" +print(f"Model used: {response.model}") # e.g., "azure_ai/gpt-4.1-nano-2025-04-14" -# Get cost +# Get cost (includes both model cost and router flat cost) from litellm import completion_cost cost = completion_cost(completion_response=response) -print(f"Cost: ${cost}") +print(f"Total cost: ${cost}") + +# Access detailed cost breakdown +if hasattr(response, '_hidden_params') and 'response_cost' in response._hidden_params: + print(f"Response cost: ${response._hidden_params['response_cost']}") ``` +### Viewing Cost Breakdown in UI + +When viewing logs in the LiteLLM UI, you'll see: +- **Model Cost**: The cost for the actual model used +- **Azure Model Router Flat Cost**: The $0.14/M input tokens infrastructure fee +- **Total Cost**: Sum of both costs + +This breakdown helps you understand exactly what you're paying for when using the Model Router. + diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 490f0288b00..bef4d52ce49 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -36,6 +36,9 @@ from litellm.llms.anthropic.cost_calculation import ( from litellm.llms.azure.cost_calculation import ( cost_per_token as azure_openai_cost_per_token, ) +from litellm.llms.azure_ai.cost_calculator import ( + cost_per_token as azure_ai_cost_per_token, +) from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.llms.bedrock.cost_calculation import ( cost_per_token as bedrock_cost_per_token, @@ -138,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: @@ -427,8 +475,8 @@ def cost_per_token( # noqa: PLR0915 return dashscope_cost_per_token(model=model, usage=usage_block) elif custom_llm_provider == "azure_ai": - return generic_cost_per_token( - model=model, usage=usage_block, custom_llm_provider=custom_llm_provider + return azure_ai_cost_per_token( + model=model, usage=usage_block, response_time_ms=response_time_ms ) else: model_info = _cached_get_model_info_helper( @@ -805,6 +853,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, + additional_costs: Optional[dict] = None, original_cost: Optional[float] = None, discount_percent: Optional[float] = None, discount_amount: Optional[float] = None, @@ -821,6 +870,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 + additional_costs: Free-form additional costs dict (e.g., {"azure_model_router_flat_cost": 0.00014}) original_cost: Cost before discount discount_percent: Discount percentage applied (0.05 = 5%) discount_amount: Discount amount in USD @@ -838,6 +888,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, + additional_costs=additional_costs, original_cost=original_cost, discount_percent=discount_percent, discount_amount=discount_amount, @@ -1335,6 +1386,15 @@ def completion_cost( # noqa: PLR0915 service_tier=service_tier, response=completion_response, ) + + # 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 ) @@ -1374,6 +1434,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, + additional_costs=additional_costs, 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 9d1360bf057..4ad2d1002bc 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, + additional_costs: Optional[dict] = 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 + additional_costs: Free-form additional costs dict (e.g., {"azure_model_router_flat_cost": 0.00014}) 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 additional costs if provided (free-form dict for extensibility) + if additional_costs and isinstance(additional_costs, dict) and len(additional_costs) > 0: + self.cost_breakdown["additional_costs"] = additional_costs + # 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/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/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 01a3f5766c6..748680f7e13 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -13,14 +13,21 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): self._model = model @staticmethod - def get_azure_ai_route(model: str) -> Literal["agents", "default"]: + def get_azure_ai_route(model: str) -> Literal["agents", "model_router", "default"]: """ Get the Azure AI route for the given model. Similar to BedrockModelInfo.get_bedrock_route(). + + Supported routes: + - agents: azure_ai/agents/ + - model_router: azure_ai/model_router/ + - default: standard models """ if "agents/" in model: return "agents" + if "model_router/" in model or "model-router/" in model: + return "model_router" return "default" @staticmethod @@ -75,8 +82,73 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): ######################################################### @staticmethod - def get_base_model(model: str) -> Optional[str]: - raise NotImplementedError("Azure Foundry does not support base model") + def strip_model_router_prefix(model: str) -> str: + """ + Strip the model_router prefix from model name. + + Examples: + - "model_router/gpt-4o" -> "gpt-4o" + - "model-router/gpt-4o" -> "gpt-4o" + - "gpt-4o" -> "gpt-4o" + + Args: + model: Model name potentially with model_router prefix + + Returns: + Model name without the prefix + """ + if "model_router/" in model: + return model.split("model_router/", 1)[1] + if "model-router/" in model: + return model.split("model-router/", 1)[1] + return model + + @staticmethod + def get_base_model(model: str) -> str: + """ + Get the base model name, stripping any Azure AI routing prefixes. + + Args: + model: Model name potentially with routing prefixes + + Returns: + Base model name + """ + # Strip model_router prefix if present + model = AzureFoundryModelInfo.strip_model_router_prefix(model) + return model + + @staticmethod + def get_azure_ai_config_for_model(model: str): + """ + Get the appropriate Azure AI config class for the given model. + + Routes to specialized configs based on model type: + - Model Router: AzureModelRouterConfig + - Claude models: AzureAnthropicConfig + - Default: AzureAIStudioConfig + + Args: + model: The model name + + Returns: + The appropriate config instance + """ + azure_ai_route = AzureFoundryModelInfo.get_azure_ai_route(model) + + if azure_ai_route == "model_router": + from litellm.llms.azure_ai.azure_model_router.transformation import ( + AzureModelRouterConfig, + ) + return AzureModelRouterConfig() + elif "claude" in model.lower(): + from litellm.llms.azure_ai.anthropic.transformation import ( + AzureAnthropicConfig, + ) + return AzureAnthropicConfig() + else: + from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig + return AzureAIStudioConfig() def validate_environment( self, diff --git a/litellm/llms/azure_ai/cost_calculator.py b/litellm/llms/azure_ai/cost_calculator.py new file mode 100644 index 00000000000..b6258425a1f --- /dev/null +++ b/litellm/llms/azure_ai/cost_calculator.py @@ -0,0 +1,102 @@ +""" +Azure AI cost calculation helper. +Handles Azure AI Foundry Model Router flat cost and other Azure AI specific pricing. +""" + +from typing import Optional, Tuple + +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.types.utils import Usage +from litellm.utils import get_model_info + + +def _is_azure_model_router(model: str) -> bool: + """ + Check if the model is Azure AI Foundry Model Router. + + Detects patterns like: + - "azure-model-router" + - "model-router" + - "model_router/" + - "model-router/" + + Args: + model: The model name + + Returns: + bool: True if this is a model router model + """ + model_lower = model.lower() + return ( + "model-router" in model_lower + or "model_router" in model_lower + or model_lower == "azure-model-router" + ) + + +def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> float: + """ + Calculate the flat cost for Azure AI Foundry Model Router. + + Args: + model: The model name (should be a model router model) + prompt_tokens: Number of prompt tokens + + Returns: + float: The flat cost in USD, or 0.0 if not applicable + """ + if not _is_azure_model_router(model): + return 0.0 + + # Get the model router pricing from model_prices_and_context_window.json + # Use "model_router" as the key (without actual model name suffix) + model_info = get_model_info(model="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: + return prompt_tokens * router_flat_cost_per_token + + return 0.0 + + +def cost_per_token( + model: str, usage: Usage, response_time_ms: Optional[float] = 0.0 +) -> Tuple[float, float]: + """ + Calculate the cost per token for Azure AI models. + + 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) + + Args: + model: str, the model name without provider prefix + usage: LiteLLM Usage block + response_time_ms: Optional response time in milliseconds + + Returns: + Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd + """ + # Calculate base cost using generic cost calculator + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="azure_ai", + ) + + # 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): + router_flat_cost = calculate_azure_model_router_flat_cost(model, usage.prompt_tokens) + + if router_flat_cost > 0: + verbose_logger.debug( + f"Azure AI Model Router flat cost: ${router_flat_cost:.6f} " + f"({usage.prompt_tokens} tokens × ${router_flat_cost / usage.prompt_tokens:.9f}/token)" + ) + + # Add flat cost to prompt cost + prompt_cost += router_flat_cost + + return prompt_cost, completion_cost 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 diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d874e6ba578..0f84bba941d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1517,6 +1517,14 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "azure_ai/model_router": { + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 0, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-services/", + "comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/ where deployment-name is your Azure deployment (e.g., azure-model-router)" + }, "azure/eu/gpt-4o-2024-08-06": { "deprecation_date": "2026-02-27", "cache_read_input_token_cost": 1.375e-06, diff --git a/litellm/proxy/cached_logo.jpg b/litellm/proxy/cached_logo.jpg index da4faf5b5ca..a10a1d24969 100644 Binary files a/litellm/proxy/cached_logo.jpg and b/litellm/proxy/cached_logo.jpg differ diff --git a/litellm/types/utils.py b/litellm/types/utils.py index f063418f92e..6c330d0f83c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2637,6 +2637,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 + additional_costs: Dict[str, float] # Free-form additional costs (e.g., {"azure_model_router_flat_cost": 0.00014}) 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/litellm/utils.py b/litellm/utils.py index 15fec357bf2..7c4eec7ba32 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7914,9 +7914,8 @@ class ProviderConfigManager: @staticmethod def _get_azure_ai_config(model: str) -> BaseConfig: """Get Azure AI config based on model type.""" - if "claude" in model.lower(): - return litellm.AzureAnthropicConfig() - return litellm.AzureAIStudioConfig() + from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + return AzureFoundryModelInfo.get_azure_ai_config_for_model(model) @staticmethod def _get_vertex_ai_config(model: str) -> BaseConfig: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d874e6ba578..0f84bba941d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1517,6 +1517,14 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "azure_ai/model_router": { + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 0, + "litellm_provider": "azure_ai", + "mode": "chat", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-services/", + "comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/ where deployment-name is your Azure deployment (e.g., azure-model-router)" + }, "azure/eu/gpt-4o-2024-08-06": { "deprecation_date": "2026-02-27", "cache_read_input_token_cost": 1.375e-06, diff --git a/tests/llm_translation/test_azure_ai.py b/tests/llm_translation/test_azure_ai.py index 1633fb4cc82..972ba34a179 100644 --- a/tests/llm_translation/test_azure_ai.py +++ b/tests/llm_translation/test_azure_ai.py @@ -373,10 +373,17 @@ def test_completion_azure_ai_gpt_4o_with_flexible_api_base(api_base): async def test_azure_ai_model_router(): """ Test Azure AI model router non-streaming response cost tracking. + Verifies that the flat cost of $0.14 per M input tokens is applied. + + Tests the pattern: azure_ai/model_router/ + Where deployment-name is the Azure deployment (e.g., "azure-model-router"). + The model_router prefix is stripped before sending to Azure API. """ + from litellm.llms.azure_ai.cost_calculator import calculate_azure_model_router_flat_cost + litellm._turn_on_debug() response = await litellm.acompletion( - model="azure_ai/azure-model-router", + model="azure_ai/model_router/azure-model-router", messages=[{"role": "user", "content": "hi who is this"}], api_base="https://ishaa-mh6uutut-swedencentral.cognitiveservices.azure.com/openai/v1/", api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"), @@ -387,6 +394,25 @@ async def test_azure_ai_model_router(): tracked_cost = response._hidden_params["response_cost"] assert tracked_cost > 0 print("Tracked cost: ", tracked_cost) + + # Verify flat cost is included using the helper function + usage = response.usage + if usage and usage.prompt_tokens: + expected_flat_cost = calculate_azure_model_router_flat_cost( + model="model_router/azure-model-router", + prompt_tokens=usage.prompt_tokens + ) + print(f"Prompt tokens: {usage.prompt_tokens}") + print(f"Expected flat cost: ${expected_flat_cost:.9f}") + print(f"Total tracked cost: ${tracked_cost:.9f}") + + # Total cost should be at least the flat cost + assert tracked_cost >= expected_flat_cost, ( + f"Cost ${tracked_cost:.9f} should be >= flat cost ${expected_flat_cost:.9f}" + ) + + # Verify the flat cost is non-zero + assert expected_flat_cost > 0, "Flat cost should be greater than 0" @pytest.mark.asyncio diff --git a/tests/test_litellm/llms/azure_ai/test_cost_calculator.py b/tests/test_litellm/llms/azure_ai/test_cost_calculator.py new file mode 100644 index 00000000000..aab8f8bf926 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/test_cost_calculator.py @@ -0,0 +1,335 @@ +""" +Test Azure AI cost calculator, especially Model Router flat cost. +""" + +import pytest + +from litellm.llms.azure_ai.cost_calculator import ( + _is_azure_model_router, + cost_per_token, +) +from litellm.types.utils import Usage +from litellm.utils import get_model_info + +# Get the flat cost from model_prices_and_context_window.json +_model_info = get_model_info(model="model_router", custom_llm_provider="azure_ai") +AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS = _model_info.get("input_cost_per_token", 0) * 1_000_000 + + +class TestAzureModelRouterDetection: + """Test that we correctly identify Azure Model Router models. + + Model Router deployments follow the pattern: model_router/ + where deployment-name is the Azure deployment (e.g., 'azure-model-router', 'prod-router') + """ + + @pytest.mark.parametrize( + "model,expected", + [ + # Deployment names containing 'model-router' or 'model_router' + ("azure-model-router", True), + ("AZURE-MODEL-ROUTER", True), + ("model-router", True), + ("MODEL-ROUTER", True), + ("my-model-router-deployment", True), + ("prod-model_router", True), + # New pattern: model_router/ + ("model_router/azure-model-router", True), + ("model-router/prod-router", True), + ("model_router/my-deployment", True), + ("MODEL_ROUTER/AZURE-MODEL-ROUTER", True), + # Non-router models + ("gpt-4o", False), + ("gpt-4o-mini", False), + ("claude-sonnet-4-5", False), + ("my-regular-deployment", False), + ], + ) + def test_is_azure_model_router(self, model: str, expected: bool): + """Test Azure Model Router detection.""" + assert _is_azure_model_router(model) == expected + + +class TestAzureModelRouterPrefix: + """Test Azure Model Router prefix stripping.""" + + @pytest.mark.parametrize( + "model,expected", + [ + # Model router deployments - the deployment name comes after model_router/ + ("model_router/azure-model-router", "azure-model-router"), + ("model-router/my-router-deployment", "my-router-deployment"), + ("model_router/prod-router", "prod-router"), + # Non-router models - should pass through unchanged + ("gpt-4o", "gpt-4o"), + ("azure-model-router", "azure-model-router"), + ("claude-sonnet-4", "claude-sonnet-4"), + ], + ) + def test_strip_model_router_prefix(self, model: str, expected: str): + """Test that model_router prefix is stripped correctly. + + The pattern is: model_router/ + where deployment-name is the Azure deployment (e.g., 'azure-model-router', 'prod-router') + """ + from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + + result = AzureFoundryModelInfo.strip_model_router_prefix(model) + assert result == expected + + +class TestAzureModelRouterFlatCost: + """Test Azure AI Foundry Model Router flat cost calculation.""" + + def test_model_router_flat_cost_basic(self): + """Test that flat cost is added for Model Router requests.""" + model = "azure-model-router" + usage = Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + ) + + prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) + + # Calculate expected flat cost + expected_flat_cost = ( + usage.prompt_tokens * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 + ) + + # Flat cost should be $0.00014 (1000 tokens × $0.14 / 1M tokens) + assert expected_flat_cost == pytest.approx(0.00014, rel=1e-9) + + # Prompt cost should include the flat cost + # (plus any base cost from the actual model used, which might be 0 if not in model_cost) + assert prompt_cost >= expected_flat_cost + print( + f"Model Router flat cost for {usage.prompt_tokens} tokens: ${expected_flat_cost:.6f}" + ) + print(f"Total prompt cost: ${prompt_cost:.6f}") + + def test_model_router_flat_cost_large_request(self): + """Test flat cost calculation for larger requests.""" + model = "model-router" + usage = Usage( + prompt_tokens=100_000, + completion_tokens=50_000, + total_tokens=150_000, + ) + + prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) + + # Calculate expected flat cost + expected_flat_cost = ( + usage.prompt_tokens * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 + ) + + # Flat cost should be $0.014 (100k tokens × $0.14 / 1M tokens) + assert expected_flat_cost == pytest.approx(0.014, rel=1e-9) + assert prompt_cost >= expected_flat_cost + print( + f"Model Router flat cost for {usage.prompt_tokens} tokens: ${expected_flat_cost:.6f}" + ) + print(f"Total prompt cost: ${prompt_cost:.6f}") + + def test_model_router_flat_cost_1m_tokens(self): + """Test flat cost for exactly 1 million input tokens.""" + model = "azure-model-router" + usage = Usage( + prompt_tokens=1_000_000, + completion_tokens=100_000, + total_tokens=1_100_000, + ) + + prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) + + # Calculate expected flat cost + expected_flat_cost = AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS + + # Flat cost should be exactly $0.14 for 1M tokens + assert expected_flat_cost == pytest.approx(0.14, rel=1e-9) + assert prompt_cost >= expected_flat_cost + print(f"Model Router flat cost for 1M tokens: ${expected_flat_cost:.6f}") + print(f"Total prompt cost: ${prompt_cost:.6f}") + + def test_non_model_router_no_flat_cost(self): + """Test that non-Model Router models don't get the flat cost.""" + model = "gpt-4o" + usage = Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + ) + + prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) + + # No flat cost should be added for non-Model Router models + # The cost might be 0 or based on the model's pricing + print(f"Non-Model Router prompt cost: ${prompt_cost:.6f}") + # We just ensure it doesn't crash and returns valid values + assert prompt_cost >= 0 + assert completion_cost >= 0 + + def test_model_router_with_cached_tokens(self): + """Test Model Router flat cost with cached tokens.""" + model = "azure-model-router" + usage = Usage( + prompt_tokens=2000, + completion_tokens=800, + total_tokens=2800, + cache_read_input_tokens=500, + cache_creation_input_tokens=200, + ) + + prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) + + # Flat cost is based on ALL prompt tokens (including cached) + expected_flat_cost = ( + usage.prompt_tokens * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 + ) + + assert expected_flat_cost == pytest.approx(0.00028, rel=1e-9) + assert prompt_cost >= expected_flat_cost + print( + 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_calculation_helper(self): + """Test that flat cost can be calculated using the helper function.""" + from litellm.llms.azure_ai.cost_calculator import ( + calculate_azure_model_router_flat_cost, + ) + + model = "azure-model-router" + prompt_tokens = 10000 + + # Calculate flat cost using helper function + flat_cost = calculate_azure_model_router_flat_cost( + model=model, prompt_tokens=prompt_tokens + ) + + # Expected flat cost + expected_flat_cost = ( + prompt_tokens * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 + ) + + assert flat_cost > 0 + assert flat_cost == pytest.approx(expected_flat_cost, rel=1e-9) + 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.""" + import litellm + from litellm.cost_calculator import completion_cost + from litellm.types.utils import Choices, Message, ModelResponse, Usage + + # 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}") + + def test_additional_costs_in_cost_breakdown(self): + """Test that Azure Model Router flat cost appears in additional_costs dict.""" + from litellm.cost_calculator import completion_cost + from litellm.litellm_core_utils.litellm_logging import LitellmLoggingObject + from litellm.types.utils import Choices, Message, ModelResponse, Usage + + # Create logging object + logging_obj = LitellmLoggingObject() + + # 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 with logging object + cost = completion_cost( + completion_response=response, + model="azure-model-router", + custom_llm_provider="azure_ai", + litellm_logging_obj=logging_obj, + ) + + # Check that cost breakdown contains additional_costs + assert hasattr(logging_obj, "cost_breakdown") + assert logging_obj.cost_breakdown is not None + assert "additional_costs" in logging_obj.cost_breakdown + assert isinstance(logging_obj.cost_breakdown["additional_costs"], dict) + + # Check that the Azure Model Router flat cost is in additional_costs + additional_costs = logging_obj.cost_breakdown["additional_costs"] + assert "Azure Model Router Flat Cost" in additional_costs + + # Verify the flat cost value + expected_flat_cost = ( + 5000 * AZURE_MODEL_ROUTER_FLAT_COST_PER_M_INPUT_TOKENS / 1_000_000 + ) + actual_flat_cost = additional_costs["Azure Model Router Flat Cost"] + assert actual_flat_cost == pytest.approx(expected_flat_cost, rel=1e-9) + + print(f"Additional costs in breakdown: {additional_costs}") + print(f"Azure Model Router Flat Cost: ${actual_flat_cost:.6f}") diff --git a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx new file mode 100644 index 00000000000..de853303c15 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx @@ -0,0 +1,289 @@ +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders, screen, waitFor } from "../../../../tests/test-utils"; +import UserDropdown from "./UserDropdown"; + +let mockUseAuthorizedImpl = () => ({ + userId: "test-user-id", + userEmail: "test@example.com", + userRole: "Admin", + premiumUser: false, +}); + +let mockUseDisableShowPromptsImpl = () => false; + +let mockGetLocalStorageItemImpl = (key: string): string | null => { + if (key === "disableShowNewBadge") return null; + if (key === "disableShowPrompts") return null; + return null; +}; + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorizedImpl(), +})); + +vi.mock("@/app/(dashboard)/hooks/useDisableShowPrompts", () => ({ + useDisableShowPrompts: () => mockUseDisableShowPromptsImpl(), +})); + +vi.mock("@/utils/localStorageUtils", () => ({ + LOCAL_STORAGE_EVENT: "local-storage-change", + getLocalStorageItem: (key: string) => mockGetLocalStorageItemImpl(key), + setLocalStorageItem: vi.fn(), + removeLocalStorageItem: vi.fn(), + emitLocalStorageChange: vi.fn(), +})); + +describe("UserDropdown", () => { + const mockOnLogout = vi.fn(); + + beforeEach(() => { + vi.clearAllMocks(); + mockUseAuthorizedImpl = () => ({ + userId: "test-user-id", + userEmail: "test@example.com", + userRole: "Admin", + premiumUser: false, + }); + mockUseDisableShowPromptsImpl = () => false; + mockGetLocalStorageItemImpl = (key: string): string | null => { + if (key === "disableShowNewBadge") return null; + if (key === "disableShowPrompts") return null; + return null; + }; + }); + + it("should render", () => { + renderWithProviders(); + expect(screen.getByRole("button")).toBeInTheDocument(); + }); + + it("should display user button with User text", () => { + renderWithProviders(); + expect(screen.getByText("User")).toBeInTheDocument(); + }); + + it("should show user email when dropdown is opened", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("test@example.com")).toBeInTheDocument(); + }); + }); + + it("should show user ID when dropdown is opened", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("test-user-id")).toBeInTheDocument(); + }); + }); + + it("should show user role when dropdown is opened", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("Admin")).toBeInTheDocument(); + }); + }); + + it("should display Standard badge for non-premium users", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("Standard")).toBeInTheDocument(); + }); + }); + + it("should display Premium badge for premium users", async () => { + const user = userEvent.setup(); + mockUseAuthorizedImpl = () => ({ + userId: "test-user-id", + userEmail: "test@example.com", + userRole: "Admin", + premiumUser: true, + }); + + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("Premium")).toBeInTheDocument(); + }); + }); + + it("should call onLogout when logout is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("test@example.com")).toBeInTheDocument(); + }); + + await user.click(screen.getByText("Logout")); + + expect(mockOnLogout).toHaveBeenCalledTimes(1); + }); + + it("should toggle hide new feature indicators switch", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("test@example.com")).toBeInTheDocument(); + }); + + const toggle = screen.getByLabelText("Toggle hide new feature indicators"); + expect(toggle).not.toBeChecked(); + + await user.click(toggle); + + const localStorageUtils = vi.mocked(await import("@/utils/localStorageUtils")); + expect(localStorageUtils.setLocalStorageItem).toHaveBeenCalledWith("disableShowNewBadge", "true"); + expect(localStorageUtils.emitLocalStorageChange).toHaveBeenCalledWith("disableShowNewBadge"); + }); + + it("should toggle hide new feature indicators switch off", async () => { + const user = userEvent.setup(); + mockGetLocalStorageItemImpl = (key: string): string | null => { + if (key === "disableShowNewBadge") return "true"; + return null; + }; + + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("test@example.com")).toBeInTheDocument(); + }); + + const toggle = screen.getByLabelText("Toggle hide new feature indicators"); + expect(toggle).toBeChecked(); + + await user.click(toggle); + + const localStorageUtils = vi.mocked(await import("@/utils/localStorageUtils")); + expect(localStorageUtils.removeLocalStorageItem).toHaveBeenCalledWith("disableShowNewBadge"); + expect(localStorageUtils.emitLocalStorageChange).toHaveBeenCalledWith("disableShowNewBadge"); + }); + + it("should toggle hide all prompts switch", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("test@example.com")).toBeInTheDocument(); + }); + + const toggle = screen.getByLabelText("Toggle hide all prompts"); + expect(toggle).not.toBeChecked(); + + await user.click(toggle); + + const localStorageUtils = vi.mocked(await import("@/utils/localStorageUtils")); + expect(localStorageUtils.setLocalStorageItem).toHaveBeenCalledWith("disableShowPrompts", "true"); + expect(localStorageUtils.emitLocalStorageChange).toHaveBeenCalledWith("disableShowPrompts"); + }); + + it("should toggle hide all prompts switch off", async () => { + const user = userEvent.setup(); + mockUseDisableShowPromptsImpl = () => true; + mockGetLocalStorageItemImpl = (key: string): string | null => { + if (key === "disableShowPrompts") return "true"; + return null; + }; + + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("test@example.com")).toBeInTheDocument(); + }); + + const toggle = screen.getByLabelText("Toggle hide all prompts"); + expect(toggle).toBeChecked(); + + await user.click(toggle); + + const localStorageUtils = vi.mocked(await import("@/utils/localStorageUtils")); + expect(localStorageUtils.removeLocalStorageItem).toHaveBeenCalledWith("disableShowPrompts"); + expect(localStorageUtils.emitLocalStorageChange).toHaveBeenCalledWith("disableShowPrompts"); + }); + + it("should display dash when user email is not available", async () => { + const user = userEvent.setup(); + mockUseAuthorizedImpl = () => ({ + userId: "test-user-id", + userEmail: null as any, + userRole: "Admin", + premiumUser: false, + }); + + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("-")).toBeInTheDocument(); + }); + }); + + it("should display dash when user ID is not available", async () => { + const user = userEvent.setup(); + mockUseAuthorizedImpl = () => ({ + userId: null as any, + userEmail: "test@example.com", + userRole: "Admin", + premiumUser: false, + }); + + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + const dashElements = screen.getAllByText("-"); + expect(dashElements.length).toBeGreaterThan(0); + }); + }); + + it("should initialize hide new feature indicators from localStorage", async () => { + const user = userEvent.setup(); + mockGetLocalStorageItemImpl = (key: string): string | null => { + if (key === "disableShowNewBadge") return "true"; + return null; + }; + + renderWithProviders(); + + await user.click(screen.getByText("User")); + + await waitFor(() => { + expect(screen.getByText("test@example.com")).toBeInTheDocument(); + }); + + const toggle = screen.getByLabelText("Toggle hide new feature indicators"); + expect(toggle).toBeChecked(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx new file mode 100644 index 00000000000..f80af33f9e2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx @@ -0,0 +1,161 @@ +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { useDisableShowPrompts } from "@/app/(dashboard)/hooks/useDisableShowPrompts"; +import { + emitLocalStorageChange, + getLocalStorageItem, + removeLocalStorageItem, + setLocalStorageItem, +} from "@/utils/localStorageUtils"; +import { + CrownOutlined, + DownOutlined, + LogoutOutlined, + MailOutlined, + SafetyOutlined, + UserOutlined, +} from "@ant-design/icons"; +import type { MenuProps } from "antd"; +import { Button, Divider, Dropdown, Space, Switch, Tag, Tooltip, Typography } from "antd"; +import React, { useEffect, useState } from "react"; + +const { Text } = Typography; + +interface UserDropdownProps { + onLogout: () => void; +} + +const UserDropdown: React.FC = ({ onLogout }) => { + const { userId, userEmail, userRole, premiumUser } = useAuthorized(); + const disableShowPrompts = useDisableShowPrompts(); + const [disableShowNewBadge, setDisableShowNewBadge] = useState(false); + + useEffect(() => { + const storedValue = getLocalStorageItem("disableShowNewBadge"); + setDisableShowNewBadge(storedValue === "true"); + }, []); + + const userItems: MenuProps["items"] = [ + { + key: "logout", + label: ( + + + Logout + + ), + onClick: onLogout, + }, + ]; + + const renderUserInfoSection = () => ( + + + + + {userEmail || "-"} + + {premiumUser ? ( + } + color="gold" + > + Premium + + ) : ( + + } + > + Standard + + + )} + + + + + + User ID + + + {userId || "-"} + + + + + + Role + + {userRole} + + + + Hide New Feature Indicators + { + setDisableShowNewBadge(checked); + if (checked) { + setLocalStorageItem("disableShowNewBadge", "true"); + emitLocalStorageChange("disableShowNewBadge"); + } else { + removeLocalStorageItem("disableShowNewBadge"); + emitLocalStorageChange("disableShowNewBadge"); + } + }} + aria-label="Toggle hide new feature indicators" + /> + + + Hide All Prompts + { + if (checked) { + setLocalStorageItem("disableShowPrompts", "true"); + emitLocalStorageChange("disableShowPrompts"); + } else { + removeLocalStorageItem("disableShowPrompts"); + emitLocalStorageChange("disableShowPrompts"); + } + }} + aria-label="Toggle hide all prompts" + /> + + + ); + + return ( + ( +
+ {renderUserInfoSection()} + + {React.cloneElement(menu as React.ReactElement, { + style: { boxShadow: "none" }, + })} +
+ )} + > + + + ); +}; + +export default UserDropdown; diff --git a/ui/litellm-dashboard/src/components/navbar.test.tsx b/ui/litellm-dashboard/src/components/navbar.test.tsx index 9fa32cf9cb0..a2996f70587 100644 --- a/ui/litellm-dashboard/src/components/navbar.test.tsx +++ b/ui/litellm-dashboard/src/components/navbar.test.tsx @@ -15,8 +15,14 @@ vi.mock("@/utils/proxyUtils", () => ({ // Create mock functions that can be controlled in tests let mockUseThemeImpl = () => ({ logoUrl: null as string | null }); let mockUseHealthReadinessImpl = () => ({ data: null as any }); -let mockGetLocalStorageItemImpl = () => null as string | null; +let mockGetLocalStorageItemImpl = (key: string) => null as string | null; let mockUseDisableShowPromptsImpl = () => false; +let mockUseAuthorizedImpl = () => ({ + userId: "test-user", + userEmail: "test@example.com", + userRole: "Admin", + premiumUser: false, +}); vi.mock("@/contexts/ThemeContext", () => ({ useTheme: () => mockUseThemeImpl(), @@ -30,9 +36,13 @@ vi.mock("@/app/(dashboard)/hooks/useDisableShowPrompts", () => ({ useDisableShowPrompts: () => mockUseDisableShowPromptsImpl(), })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorizedImpl(), +})); + vi.mock("@/utils/localStorageUtils", () => ({ LOCAL_STORAGE_EVENT: "local-storage-change", - getLocalStorageItem: () => mockGetLocalStorageItemImpl(), + getLocalStorageItem: (key: string) => mockGetLocalStorageItemImpl(key), setLocalStorageItem: vi.fn(), removeLocalStorageItem: vi.fn(), emitLocalStorageChange: vi.fn(), @@ -123,14 +133,27 @@ describe("Navbar", () => { it("should show premium user badge when premiumUser is true", async () => { const user = userEvent.setup(); - const premiumProps = { ...defaultProps, premiumUser: true }; - renderWithProviders(); + mockUseAuthorizedImpl = () => ({ + userId: "test-user", + userEmail: "test@example.com", + userRole: "Admin", + premiumUser: true, + }); + renderWithProviders(); await user.click(screen.getByText("User")); await waitFor(() => { expect(screen.getByText("Premium")).toBeInTheDocument(); }); + + // Reset mock + mockUseAuthorizedImpl = () => ({ + userId: "test-user", + userEmail: "test@example.com", + userRole: "Admin", + premiumUser: false, + }); }); it("should show version badge when health data contains version", () => { @@ -167,7 +190,10 @@ describe("Navbar", () => { const user = userEvent.setup(); // Initially disabled - mockGetLocalStorageItemImpl = () => "false"; + mockGetLocalStorageItemImpl = (key: string) => { + if (key === "disableShowNewBadge") return "false"; + return null; + }; renderWithProviders(); @@ -186,6 +212,9 @@ describe("Navbar", () => { const localStorageUtils = vi.mocked(await import("@/utils/localStorageUtils")); expect(localStorageUtils.setLocalStorageItem).toHaveBeenCalledWith("disableShowNewBadge", "true"); expect(localStorageUtils.emitLocalStorageChange).toHaveBeenCalledWith("disableShowNewBadge"); + + // Reset mock + mockGetLocalStorageItemImpl = (key: string) => null; }); it("should handle logout functionality", async () => { diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index 8e4fda0ba52..3649ca76238 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -1,32 +1,20 @@ import { useHealthReadiness } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadiness"; -import { useDisableShowPrompts } from "@/app/(dashboard)/hooks/useDisableShowPrompts"; import { getProxyBaseUrl } from "@/components/networking"; import { useTheme } from "@/contexts/ThemeContext"; import { clearTokenCookies } from "@/utils/cookieUtils"; -import { - emitLocalStorageChange, - getLocalStorageItem, - removeLocalStorageItem, - setLocalStorageItem, -} from "@/utils/localStorageUtils"; import { fetchProxySettings } from "@/utils/proxyUtils"; import { - CrownOutlined, GithubOutlined, - LogoutOutlined, - MailOutlined, MenuFoldOutlined, MenuUnfoldOutlined, MoonOutlined, - SafetyOutlined, SlackOutlined, SunOutlined, - UserOutlined, } from "@ant-design/icons"; -import type { MenuProps } from "antd"; -import { Button, Dropdown, Switch, Tooltip } from "antd"; +import { Button, Switch, Tag } from "antd"; import Link from "next/link"; import React, { useEffect, useState } from "react"; +import UserDropdown from "./Navbar/UserDropdown/UserDropdown"; interface NavbarProps { userID: string | null; @@ -59,8 +47,6 @@ const Navbar: React.FC = ({ }) => { const baseUrl = getProxyBaseUrl(); const [logoutUrl, setLogoutUrl] = useState(""); - const [disableShowNewBadge, setDisableShowNewBadge] = useState(false); - const disableShowPrompts = useDisableShowPrompts(); const { logoUrl } = useTheme(); const { data: healthData } = useHealthReadiness(); const version = healthData?.litellm_version; @@ -82,11 +68,6 @@ const Navbar: React.FC = ({ initializeProxySettings(); }, [accessToken]); - useEffect(() => { - const storedValue = getLocalStorageItem("disableShowNewBadge"); - setDisableShowNewBadge(storedValue === "true"); - }, []); - useEffect(() => { setLogoutUrl(proxySettings?.PROXY_LOGOUT_URL || ""); }, [proxySettings]); @@ -96,105 +77,6 @@ const Navbar: React.FC = ({ window.location.href = logoutUrl; }; - const userItems: MenuProps["items"] = [ - { - key: "user-info", - // Prevent dropdown from closing when interacting with the toggle - onClick: (info) => info.domEvent?.stopPropagation(), - label: ( -
-
-
- - {userID} -
- {premiumUser ? ( - -
- - Premium -
- - ) : ( - -
- - Standard -
- - )} -
-
-
- - Role - {userRole} -
-
- - Email - - {userEmail || "Unknown"} - -
-
e.stopPropagation()} - > - Hide New Feature Indicators - { - setDisableShowNewBadge(checked); - if (checked) { - setLocalStorageItem("disableShowNewBadge", "true"); - emitLocalStorageChange("disableShowNewBadge"); - } else { - removeLocalStorageItem("disableShowNewBadge"); - emitLocalStorageChange("disableShowNewBadge"); - } - }} - aria-label="Toggle hide new feature indicators" - /> -
-
e.stopPropagation()} - > - Hide All Prompts - { - if (checked) { - setLocalStorageItem("disableShowPrompts", "true"); - emitLocalStorageChange("disableShowPrompts"); - } else { - removeLocalStorageItem("disableShowPrompts"); - emitLocalStorageChange("disableShowPrompts"); - } - }} - aria-label="Toggle hide all prompts" - /> -
-
-
- ), - }, - { - key: "logout", - label: ( -
- - Logout -
- ), - }, - ]; - return (