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 16bc1afb70e..5e14c7283f6 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,38 +5,19 @@ 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`), plus the Model Router infrastructure fee +- **Cost Tracking**: LiteLLM automatically tracks costs based on the actual model used (e.g., `gpt-4.1-nano`), not the router endpoint - **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/model_router/azure-model-router", # Use your deployment name + model="azure_ai/azure-model-router", messages=[{"role": "user", "content": "Hello!"}], api_base="https://your-endpoint.cognitiveservices.azure.com/openai/v1/", api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"), @@ -45,13 +26,6 @@ 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 @@ -59,7 +33,7 @@ import litellm import os response = await litellm.acompletion( - model="azure_ai/model_router/azure-model-router", # Use your deployment name + model="azure_ai/azure-model-router", messages=[{"role": "user", "content": "hi"}], api_base="https://your-endpoint.cognitiveservices.azure.com/openai/v1/", api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"), @@ -77,15 +51,13 @@ async for chunk in response: ```yaml model_list: - - model_name: azure-model-router # Public name for your users + - model_name: azure-model-router litellm_params: - model: azure_ai/model_router/azure-model-router # Use your deployment name + model: azure_ai/azure-model-router 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 @@ -108,42 +80,49 @@ 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. -### 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 +### 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) -#### Step 2: Enter Deployment Name +### Configure Model Name -**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/`. +Set up the model name by entering `azure_ai/` followed by your model router deployment name from Azure. -**Example:** -- Enter: `azure-model-router` -- LiteLLM creates: `azure_ai/model_router/azure-model-router` +#### Click Model Name Field -##### Copy Deployment Name from Azure Portal +![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 Switch to Azure AI Foundry and copy your model router deployment name. @@ -151,79 +130,73 @@ Switch to Azure AI Foundry and copy your model router deployment name. ![Copy Model Name](./img/azure_model_router_10.jpeg) -##### Enter Deployment Name in LiteLLM +#### Paste Model Name -Paste your deployment name (e.g., `azure-model-router`) directly into the text field. +Paste to get `azure_ai/azure-model-router`. -![Enter Deployment Name](./img/azure_model_router_04.jpeg) +![Paste Model Name](./img/azure_model_router_11.jpeg) -**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 +### 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) -#### Step 4: Test and Add Model +### 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) -#### Step 5: Verify in Playground +### 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`), plus a flat infrastructure cost of $0.14 per million input tokens for using the Model Router. +Cost is tracked based on the actual model used (e.g., `gpt-4.1-nano`). ![Verify Cost](./img/azure_model_router_24.jpeg) @@ -232,50 +205,28 @@ Cost is tracked based on the actual model used (e.g., `gpt-4.1-nano`), plus a fl 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 (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 +2. **Calculating accurate costs**: Costs are calculated based on the actual model used, not the router endpoint name 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/model_router/azure-model-router", + model="azure_ai/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., "azure_ai/gpt-4.1-nano-2025-04-14" +print(f"Model used: {response.model}") # e.g., "gpt-4.1-nano-2025-04-14" -# Get cost (includes both model cost and router flat cost) +# Get cost from litellm import completion_cost cost = completion_cost(completion_response=response) -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']}") +print(f"Cost: ${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/docs/my-website/docs/tutorials/claude_code_plugin_marketplace.md b/docs/my-website/docs/tutorials/claude_code_plugin_marketplace.md index 9d93c717c4f..946fb47d92a 100644 --- a/docs/my-website/docs/tutorials/claude_code_plugin_marketplace.md +++ b/docs/my-website/docs/tutorials/claude_code_plugin_marketplace.md @@ -2,7 +2,7 @@ import Image from '@theme/IdealImage'; import Tabs from '@theme/Tabs'; import TabItem from '@theme/TabItem'; -# Claude Code Plugin Marketplace (Managed Skills) +# Claude Code Plugin Marketplace LiteLLM AI Gateway acts as a central registry for Claude Code plugins. Admins can govern which plugins are available across the organization, and engineers can discover and install approved plugins from a single source. @@ -252,7 +252,7 @@ curl -X POST http://localhost:4000/claude-code/plugins \ }' ``` -### 3. Use in Claude Code +### 3. Share with Your Team Send engineers the marketplace URL: diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index bef4d52ce49..490f0288b00 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -36,9 +36,6 @@ 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, @@ -141,51 +138,6 @@ 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: @@ -475,8 +427,8 @@ def cost_per_token( # noqa: PLR0915 return dashscope_cost_per_token(model=model, usage=usage_block) elif custom_llm_provider == "azure_ai": - return azure_ai_cost_per_token( - model=model, usage=usage_block, response_time_ms=response_time_ms + return generic_cost_per_token( + model=model, usage=usage_block, custom_llm_provider=custom_llm_provider ) else: model_info = _cached_get_model_info_helper( @@ -853,7 +805,6 @@ 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, @@ -870,7 +821,6 @@ 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 @@ -888,7 +838,6 @@ 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, @@ -1386,15 +1335,6 @@ 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 ) @@ -1434,7 +1374,6 @@ 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 14015225f38..e4cfefdfb2d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1297,7 +1297,6 @@ 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, @@ -1313,7 +1312,6 @@ 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 @@ -1329,10 +1327,6 @@ 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 deleted file mode 100644 index 0165d60b643..00000000000 --- a/litellm/llms/azure_ai/azure_model_router/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -"""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 deleted file mode 100644 index 3d6dc53c515..00000000000 --- a/litellm/llms/azure_ai/azure_model_router/transformation.py +++ /dev/null @@ -1,125 +0,0 @@ -""" -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 748680f7e13..01a3f5766c6 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -13,21 +13,14 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): self._model = model @staticmethod - def get_azure_ai_route(model: str) -> Literal["agents", "model_router", "default"]: + def get_azure_ai_route(model: str) -> Literal["agents", "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 @@ -82,73 +75,8 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): ######################################################### @staticmethod - 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 get_base_model(model: str) -> Optional[str]: + raise NotImplementedError("Azure Foundry does not support base model") def validate_environment( self, diff --git a/litellm/llms/azure_ai/cost_calculator.py b/litellm/llms/azure_ai/cost_calculator.py deleted file mode 100644 index b6258425a1f..00000000000 --- a/litellm/llms/azure_ai/cost_calculator.py +++ /dev/null @@ -1,102 +0,0 @@ -""" -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 ac209904e6e..41a1797cebe 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -437,23 +437,3 @@ 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 5b110074989..dd85737a092 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1517,14 +1517,6 @@ "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 a10a1d24969..da4faf5b5ca 100644 Binary files a/litellm/proxy/cached_logo.jpg and b/litellm/proxy/cached_logo.jpg differ diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6fd77fab7a7..136ce696511 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -623,16 +623,6 @@ class ProxyBaseLLMRequestProcessing: return self.data, logging_obj - @staticmethod - def _get_model_id_from_response(hidden_params: dict, data: dict) -> str: - """Extract model_id from hidden_params with fallback to litellm_metadata.""" - model_id = hidden_params.get("model_id", None) or "" - if not model_id: - litellm_metadata = data.get("litellm_metadata", {}) or {} - model_info = litellm_metadata.get("model_info", {}) or {} - model_id = model_info.get("id", "") or "" - return model_id - async def base_process_llm_request( self, request: Request, @@ -767,7 +757,13 @@ class ProxyBaseLLMRequestProcessing: response = responses[1] hidden_params = getattr(response, "_hidden_params", {}) or {} - model_id = self._get_model_id_from_response(hidden_params, self.data) + model_id = hidden_params.get("model_id", None) or "" + + # Fallback: extract model_id from litellm_metadata if not in hidden_params + if not model_id: + litellm_metadata = self.data.get("litellm_metadata", {}) or {} + model_info = litellm_metadata.get("model_info", {}) or {} + model_id = model_info.get("id", "") or "" cache_key, api_base, response_cost = ( hidden_params.get("cache_key", None) or "", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a00f1f605a9..078ce0edf27 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11,7 +11,7 @@ import sys import time import traceback import warnings -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta import enum from typing import ( TYPE_CHECKING, diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index f73a29dab80..867c18b6dd4 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -869,6 +869,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): try: chunk = self.litellm_custom_stream_wrapper.__next__() self._ensure_output_item_for_chunk(chunk) + # Emit any just-queued output_item event + if self._pending_response_events: + return self._pending_response_events.pop(0) self.collected_chat_completion_chunks.append(chunk) response_api_chunk = ( self._transform_chat_completion_chunk_to_response_api_chunk( @@ -876,9 +879,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) ) if response_api_chunk: - self._pending_response_events.append(response_api_chunk) - if self._pending_response_events: - return self._pending_response_events.pop(0) + return response_api_chunk # Otherwise, loop to next chunk except StopIteration: return self.common_done_event_logic(sync_mode=True) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6c330d0f83c..f063418f92e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2637,7 +2637,6 @@ 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 4fcff58496c..97aeba1b086 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7921,8 +7921,9 @@ class ProviderConfigManager: @staticmethod def _get_azure_ai_config(model: str) -> BaseConfig: """Get Azure AI config based on model type.""" - from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo - return AzureFoundryModelInfo.get_azure_ai_config_for_model(model) + if "claude" in model.lower(): + return litellm.AzureAnthropicConfig() + return litellm.AzureAIStudioConfig() @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 5b110074989..dd85737a092 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1517,14 +1517,6 @@ "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 972ba34a179..1633fb4cc82 100644 --- a/tests/llm_translation/test_azure_ai.py +++ b/tests/llm_translation/test_azure_ai.py @@ -373,17 +373,10 @@ 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/model_router/azure-model-router", + model="azure_ai/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"), @@ -394,25 +387,6 @@ 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/proxy_e2e_anthropic_messages_tests/test_claude_agent_sdk.py b/tests/proxy_e2e_anthropic_messages_tests/test_claude_agent_sdk.py index 317313c0553..0838faf4212 100644 --- a/tests/proxy_e2e_anthropic_messages_tests/test_claude_agent_sdk.py +++ b/tests/proxy_e2e_anthropic_messages_tests/test_claude_agent_sdk.py @@ -14,11 +14,9 @@ from claude_agent_sdk import ClaudeSDKClient, ClaudeAgentOptions # Test models from proxy_config.yaml -# Note: bedrock-converse-claude-sonnet-4.5 removed temporarily as the Bedrock Converse API -# for Claude Sonnet 4.5 may not be available in all regions/accounts TEST_MODELS = [ ("bedrock-claude-sonnet-4.5", "Bedrock Invoke API"), - # ("bedrock-converse-claude-sonnet-4.5", "Bedrock Converse API"), # Disabled: not yet available in CI + ("bedrock-converse-claude-sonnet-4.5", "Bedrock Converse API"), ("bedrock-nova-premier", "AWS Nova Premier"), ] diff --git a/tests/test_litellm/llms/azure_ai/test_cost_calculator.py b/tests/test_litellm/llms/azure_ai/test_cost_calculator.py deleted file mode 100644 index 03bf0a66a48..00000000000 --- a/tests/test_litellm/llms/azure_ai/test_cost_calculator.py +++ /dev/null @@ -1,345 +0,0 @@ -""" -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 datetime import datetime - - from litellm.cost_calculator import completion_cost - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import Choices, Message, ModelResponse, Usage - - # Create logging object with required parameters - logging_obj = Logging( - model="azure-model-router", - messages=[{"role": "user", "content": "Hello"}], - stream=False, - call_type="completion", - start_time=datetime.now(), - litellm_call_id="test-123", - function_id="test-function", - ) - - # 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/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index edfdeb08d82..5c1b4cbd38e 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -335,11 +335,11 @@ def test_advanced_tool_use_header_translation_for_opus_4_5(): def test_advanced_tool_use_header_filtered_for_non_opus_4_5(): """ - Test that advanced-tool-use-2025-11-20 header is filtered out for models - that don't support tool search on Bedrock. + Test that advanced-tool-use-2025-11-20 header is filtered out for non-Opus 4.5 models + without adding Bedrock-specific headers. - Tool search is supported on: Claude Opus 4.5, Claude Sonnet 4.5 - Tool search is NOT supported on: Claude 3.5 Sonnet and earlier + The translation to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 should + only happen for Claude Opus 4.5. """ from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeMessagesConfig, @@ -360,9 +360,9 @@ def test_advanced_tool_use_header_filtered_for_non_opus_4_5(): "anthropic-beta": "advanced-tool-use-2025-11-20" } - # Test with Claude 3.5 Sonnet (does NOT support tool search on Bedrock) + # Test with Claude Sonnet 4.5 (not Opus 4.5) result = config.transform_anthropic_messages_request( - model="anthropic.claude-3-5-sonnet-20241022-v2:0", + model="anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, litellm_params={}, @@ -374,11 +374,11 @@ def test_advanced_tool_use_header_filtered_for_non_opus_4_5(): assert "advanced-tool-use-2025-11-20" not in beta_headers, \ "advanced-tool-use header should be removed for Bedrock" - # Verify Bedrock-specific headers were NOT added (only for Opus 4.5 and Sonnet 4.5) + # Verify Bedrock-specific headers were NOT added (only for Opus 4.5) assert "tool-search-tool-2025-10-19" not in beta_headers, \ - "tool-search-tool should not be added for models without tool search support" + "tool-search-tool should not be added for non-Opus 4.5 models" assert "tool-examples-2025-10-29" not in beta_headers, \ - "tool-examples should not be added for models without tool search support" + "tool-examples should not be added for non-Opus 4.5 models" def test_advanced_tool_use_header_translation_with_multiple_beta_headers(): diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 23de7c9b906..5074bbf4397 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1637,82 +1637,4 @@ class TestStreamingIDConsistency: # Verify it matches the cached ID assert iterator._cached_item_id is not None - assert iterator._cached_item_id == text_done_id - - def test_first_chunk_text_delta_emitted_anthropic_style(self): - """ - When the first chunk has content (e.g. Anthropic content_block_start with initial text), - that content must be emitted as output_text.delta. Previously the first chunk was - dropped because we returned output_item_added and never transformed the chunk. - """ - from unittest.mock import Mock - - import litellm - from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, - ) - from litellm.types.llms.openai import ResponsesAPIStreamEvents - from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices - - chunk_with_initial_text = ModelResponseStream( - id="msg_abc", - choices=[ - StreamingChoices( - index=0, - delta=Delta(content="I", role="assistant"), - finish_reason=None, - ) - ], - created=1234567890, - model="anthropic/claude-sonnet-4-5", - object="chat.completion.chunk", - ) - chunk_continuation = ModelResponseStream( - id="msg_abc", - choices=[ - StreamingChoices( - index=0, - delta=Delta(content="'m happy", role=None), - finish_reason=None, - ) - ], - created=1234567890, - model="anthropic/claude-sonnet-4-5", - object="chat.completion.chunk", - ) - - mock_stream_wrapper = Mock(spec=litellm.CustomStreamWrapper) - mock_stream_wrapper.logging_obj = Mock() - mock_stream_wrapper.logging_obj._response_cost_calculator = Mock(return_value=0.001) - chunk_list = [chunk_with_initial_text, chunk_continuation] - - def next_chunk(): - if not chunk_list: - raise StopIteration - return chunk_list.pop(0) - - mock_stream_wrapper.__next__ = next_chunk - - iterator = LiteLLMCompletionStreamingIterator( - model="anthropic/claude-sonnet-4-5", - litellm_custom_stream_wrapper=mock_stream_wrapper, - request_input=[{"type": "message", "role": "user", "content": [{"type": "text", "text": "repeat after me: I'm happy to be here"}]}], - responses_api_request={}, - custom_llm_provider="anthropic", - ) - - collected_deltas = [] - max_events = 20 - for i, event in enumerate(iterator): - if i >= max_events: - break - if getattr(event, "type", None) == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA and hasattr(event, "delta"): - collected_deltas.append(event.delta) - if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED: - break - - full_text = "".join(collected_deltas) - assert full_text == "I'm happy", ( - f"First character from Anthropic content_block_start was truncated. " - f"Expected \"I'm happy\", got {repr(full_text)}" - ) \ No newline at end of file + assert iterator._cached_item_id == text_done_id \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx b/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx deleted file mode 100644 index de853303c15..00000000000 --- a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx +++ /dev/null @@ -1,289 +0,0 @@ -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 deleted file mode 100644 index f80af33f9e2..00000000000 --- a/ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx +++ /dev/null @@ -1,161 +0,0 @@ -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/mcp_tools/mcp_connect.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx index c48b9a755b7..5a012c1fc5c 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx @@ -32,7 +32,7 @@ const FeatureCard: React.FC = ({ description, children, serverName, - accessGroups = ["dev-group"], + accessGroups = ["dev"], }) => { const [useServerHeader, setUseServerHeader] = useState(false); @@ -42,9 +42,9 @@ const FeatureCard: React.FC = ({ }; if (useServerHeader && serverName) { const formattedServerName = serverName.replace(/\s+/g, "_"); - // Include both server name and access groups in the same header (comma-separated string) + // Include both server name and access groups in the same header const serverAndGroups = [formattedServerName, ...accessGroups].join(","); - headers["x-mcp-servers"] = serverAndGroups; + headers["x-mcp-servers"] = [serverAndGroups]; } return headers; }; @@ -77,13 +77,13 @@ const FeatureCard: React.FC = ({ description={

- Option 1: Get a specific server: "{serverName.replace(/\s+/g, "_")}" + Option 1: Get a specific server: ["{serverName.replace(/\s+/g, "_")}"]

- Option 2: Get a group of MCPs: "dev-group" + Option 2: Get a group of MCPs: ["dev-group"]

- You can also mix both: "Server1,dev-group" + You can also mix both: ["Server1,dev-group"]

} @@ -144,8 +144,8 @@ const MCPConnect: React.FC = ({ currentServerAccessGroups = [] // Format server names (replace spaces with underscores) const formattedServers = serverHeaders[type].map((s) => s.replace(/\s+/g, "_")); - // Use comma-separated string (can include both servers and access groups) - headers["x-mcp-servers"] = formattedServers.join(","); + // Use comma-separated format (can include both servers and access groups) + headers["x-mcp-servers"] = [formattedServers.join(",")]; } return headers; @@ -244,7 +244,7 @@ const MCPConnect: React.FC = ({ currentServerAccessGroups = [] title="Implementation Example" description="Complete cURL example for using the LiteLLM Proxy Responses API" serverName={currentServer} - accessGroups={["dev-group"]} + accessGroups={["dev"]} > = ({ currentServerAccessGroups = [] "require_approval": "never", "headers": { "x-litellm-api-key": "Bearer YOUR_LITELLM_VIRTUAL_KEY", - "x-mcp-servers": "Zapier_MCP,dev-group" + "x-mcp-servers": ["Zapier_MCP,dev"] } } ], @@ -327,7 +327,7 @@ const MCPConnect: React.FC = ({ currentServerAccessGroups = [] title="Implementation Example" description="Complete cURL example for using the Responses API" serverName="Zapier Gmail" - accessGroups={["dev-group"]} + accessGroups={["dev"]} > = ({ currentServerAccessGroups = [] "require_approval": "never", "headers": { "x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY", - "x-mcp-servers": "Zapier_MCP,dev-group" + "x-mcp-servers": ["Zapier_MCP,dev"] } } ], @@ -400,7 +400,7 @@ const MCPConnect: React.FC = ({ currentServerAccessGroups = [] title="Configuration" description="Cursor MCP configuration" serverName="Zapier Gmail" - accessGroups={["dev-group"]} + accessGroups={["dev"]} > = ({ currentServerAccessGroups = [] "url": "${proxyBaseUrl}/mcp", "headers": { "x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY", - "x-mcp-servers": "Zapier_MCP,dev-group" + "x-mcp-servers": ["Zapier_MCP,dev"] } } } diff --git a/ui/litellm-dashboard/src/components/navbar.test.tsx b/ui/litellm-dashboard/src/components/navbar.test.tsx index a2996f70587..9fa32cf9cb0 100644 --- a/ui/litellm-dashboard/src/components/navbar.test.tsx +++ b/ui/litellm-dashboard/src/components/navbar.test.tsx @@ -15,14 +15,8 @@ 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 = (key: string) => null as string | null; +let mockGetLocalStorageItemImpl = () => 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(), @@ -36,13 +30,9 @@ 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: (key: string) => mockGetLocalStorageItemImpl(key), + getLocalStorageItem: () => mockGetLocalStorageItemImpl(), setLocalStorageItem: vi.fn(), removeLocalStorageItem: vi.fn(), emitLocalStorageChange: vi.fn(), @@ -133,27 +123,14 @@ describe("Navbar", () => { it("should show premium user badge when premiumUser is true", async () => { const user = userEvent.setup(); - mockUseAuthorizedImpl = () => ({ - userId: "test-user", - userEmail: "test@example.com", - userRole: "Admin", - premiumUser: true, - }); - renderWithProviders(); + const premiumProps = { ...defaultProps, 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", () => { @@ -190,10 +167,7 @@ describe("Navbar", () => { const user = userEvent.setup(); // Initially disabled - mockGetLocalStorageItemImpl = (key: string) => { - if (key === "disableShowNewBadge") return "false"; - return null; - }; + mockGetLocalStorageItemImpl = () => "false"; renderWithProviders(); @@ -212,9 +186,6 @@ 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 3649ca76238..8e4fda0ba52 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -1,20 +1,32 @@ 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 { Button, Switch, Tag } from "antd"; +import type { MenuProps } from "antd"; +import { Button, Dropdown, Switch, Tooltip } from "antd"; import Link from "next/link"; import React, { useEffect, useState } from "react"; -import UserDropdown from "./Navbar/UserDropdown/UserDropdown"; interface NavbarProps { userID: string | null; @@ -47,6 +59,8 @@ 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; @@ -68,6 +82,11 @@ const Navbar: React.FC = ({ initializeProxySettings(); }, [accessToken]); + useEffect(() => { + const storedValue = getLocalStorageItem("disableShowNewBadge"); + setDisableShowNewBadge(storedValue === "true"); + }, []); + useEffect(() => { setLogoutUrl(proxySettings?.PROXY_LOGOUT_URL || ""); }, [proxySettings]); @@ -77,6 +96,105 @@ 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 (