Merge branch 'BerriAI:main' into main

docs/my-website/docs/providers/bedrock.md

@@ -1488,6 +1488,91 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
 
 
 
+### OpenAI GPT OSS
+
+| Property | Details |
+|----------|---------|
+| Provider Route | `bedrock/converse/openai.gpt-oss-20b-1:0`, `bedrock/converse/openai.gpt-oss-120b-1:0` |
+| Provider Documentation | [Amazon Bedrock ↗](https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html) |
+
+<Tabs>
+<TabItem value="sdk" label="SDK">
+
+```python title="GPT OSS SDK Usage" showLineNumbers
+from litellm import completion
+import os
+
+# Set AWS credentials
+os.environ["AWS_ACCESS_KEY_ID"] = "your-aws-access-key"
+os.environ["AWS_SECRET_ACCESS_KEY"] = "your-aws-secret-key"
+os.environ["AWS_REGION_NAME"] = "us-east-1"
+
+# GPT OSS 20B model
+response = completion(
+    model="bedrock/converse/openai.gpt-oss-20b-1:0",
+    messages=[{"role": "user", "content": "Hello, how are you?"}],
+)
+print(response.choices[0].message.content)
+
+# GPT OSS 120B model  
+response = completion(
+    model="bedrock/converse/openai.gpt-oss-120b-1:0",
+    messages=[{"role": "user", "content": "Explain machine learning in simple terms"}],
+)
+print(response.choices[0].message.content)
+```
+
+</TabItem>
+
+<TabItem value="proxy" label="Proxy">
+
+**1. Add to config**
+
+```yaml title="config.yaml" showLineNumbers
+model_list:
+  - model_name: gpt-oss-20b
+    litellm_params:
+      model: bedrock/converse/openai.gpt-oss-20b-1:0
+      aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
+      aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
+      aws_region_name: os.environ/AWS_REGION_NAME
+      
+  - model_name: gpt-oss-120b
+    litellm_params:
+      model: bedrock/converse/openai.gpt-oss-120b-1:0
+      aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
+      aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
+      aws_region_name: os.environ/AWS_REGION_NAME
+```
+
+**2. Start proxy**
+
+```bash title="Start LiteLLM Proxy" showLineNumbers
+litellm --config /path/to/config.yaml
+
+# RUNNING at http://0.0.0.0:4000
+```
+
+**3. Test it!**
+
+```bash title="Test GPT OSS via Proxy" showLineNumbers
+curl --location 'http://0.0.0.0:4000/chat/completions' \
+  --header 'Authorization: Bearer sk-1234' \
+  --header 'Content-Type: application/json' \
+  --data '{
+    "model": "gpt-oss-20b",
+    "messages": [
+      {
+        "role": "user", 
+        "content": "What are the key benefits of open source AI?"
+      }
+    ]
+  }'
+```
+
+</TabItem>
+</Tabs>
+
 ## Provisioned throughput models
 To use provisioned throughput Bedrock models pass 
 - `model=bedrock/<base-model>`, example `model=bedrock/anthropic.claude-v2`. Set `model` to any of the [Supported AWS models](#supported-aws-bedrock-models)
@@ -1522,6 +1607,8 @@ Here's an example of using a bedrock model with LiteLLM. For a complete list, re
 
 | Model Name                 | Command                                                          |
 |----------------------------|------------------------------------------------------------------|
+| GPT-OSS 20B | `completion(model='bedrock/converse/openai.gpt-oss-20b-1:0', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
+| GPT-OSS 120B | `completion(model='bedrock/converse/openai.gpt-oss-120b-1:0', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
 | Deepseek R1    | `completion(model='bedrock/us.deepseek.r1-v1:0', messages=messages)`   | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`           |
 | Anthropic Claude-V3.5 Sonnet    | `completion(model='bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0', messages=messages)`   | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`           |
 | Anthropic Claude-V3  sonnet    | `completion(model='bedrock/anthropic.claude-3-sonnet-20240229-v1:0', messages=messages)`   | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`           |

docs/my-website/docs/providers/sambanova.md

@@ -307,3 +307,16 @@ response = litellm.completion(
 
 print(response.choices[0].message.content))
 ```
+
+## SambaNova - Embeddings
+
+```python
+import litellm
+
+response = litellm.embedding(
+    model="sambanova/E5-Mistral-7B-Instruct",
+    input=["sample text to embed", "another sample text to embed"]
+)
+
+print(response.data)
+```

litellm/__init__.py

@@ -519,6 +519,7 @@ anyscale_models: List = []
 cerebras_models: List = []
 galadriel_models: List = []
 sambanova_models: List = []
+sambanova_embedding_models: List = []
 novita_models: List = []
 assemblyai_models: List = []
 snowflake_models: List = []
@@ -695,6 +696,8 @@ def add_known_models():
             galadriel_models.append(key)
         elif value.get("litellm_provider") == "sambanova":
             sambanova_models.append(key)
+        elif value.get("litellm_provider") == "sambanova-embedding-models":
+            sambanova_embedding_models.append(key)
         elif value.get("litellm_provider") == "novita":
             novita_models.append(key)
         elif value.get("litellm_provider") == "nebius-chat-models":
@@ -879,7 +882,7 @@ models_by_provider: dict = {
     "anyscale": anyscale_models,
     "cerebras": cerebras_models,
     "galadriel": galadriel_models,
-    "sambanova": sambanova_models,
+    "sambanova": sambanova_models + sambanova_embedding_models,
     "novita": novita_models,
     "nebius": nebius_models + nebius_embedding_models,
     "assemblyai": assemblyai_models,
@@ -933,6 +936,7 @@ all_embedding_models = (
     + vertex_embedding_models
     + fireworks_ai_embedding_models
     + nebius_embedding_models
+    + sambanova_embedding_models
 )
 
 ####### IMAGE GENERATION MODELS ###################
@@ -1185,6 +1189,7 @@ nvidiaNimEmbeddingConfig = NvidiaNimEmbeddingConfig()
 from .llms.featherless_ai.chat.transformation import FeatherlessAIConfig
 from .llms.cerebras.chat import CerebrasConfig
 from .llms.sambanova.chat import SambanovaConfig
+from .llms.sambanova.embedding.transformation import SambaNovaEmbeddingConfig
 from .llms.ai21.chat.transformation import AI21ChatConfig
 from .llms.fireworks_ai.chat.transformation import FireworksAIConfig
 from .llms.fireworks_ai.completion.transformation import FireworksAITextCompletionConfig

litellm/litellm_core_utils/get_supported_openai_params.py

@@ -140,7 +140,10 @@ def get_supported_openai_params(  # noqa: PLR0915
             model=model
         )
     elif custom_llm_provider == "sambanova":
-        return litellm.SambanovaConfig().get_supported_openai_params(model=model)
+        if request_type == "embeddings":
+            litellm.SambaNovaEmbeddingConfig().get_supported_openai_params(model=model)
+        else:
+            return litellm.SambanovaConfig().get_supported_openai_params(model=model)
     elif custom_llm_provider == "nebius":
         if request_type == "chat_completion":
             return litellm.NebiusConfig().get_supported_openai_params(model=model)

litellm/llms/anthropic/common_utils.py

@@ -10,10 +10,11 @@ import litellm
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     get_file_ids_from_messages,
 )
-from litellm.llms.base_llm.base_utils import BaseLLMModelInfo
+from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.types.llms.anthropic import AllAnthropicToolsValues, AnthropicMcpServerTool
 from litellm.types.llms.openai import AllMessageValues
+from litellm.types.utils import TokenCountResponse
 
 
 class AnthropicError(BaseLLMException):
@@ -229,7 +230,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
             litellm_model_names.append(litellm_model_name)
         return litellm_model_names
 
-    def get_token_counter(self) -> Optional["AnthropicTokenCounter"]:
+    def get_token_counter(self) -> Optional[BaseTokenCounter]:
         """
         Factory method to create an Anthropic token counter.
         
@@ -239,32 +240,24 @@ class AnthropicModelInfo(BaseLLMModelInfo):
         return AnthropicTokenCounter()
 
 
-class AnthropicTokenCounter:
+class AnthropicTokenCounter(BaseTokenCounter):
     """Token counter implementation for Anthropic provider."""
-    
-    def supports_provider(
+
+    def should_use_token_counting_api(
         self, 
-        deployment: Optional[Dict[str, Any]] = None,
-        from_endpoint: bool = False
+        custom_llm_provider: Optional[str] = None,
     ) -> bool:
-        if not from_endpoint:
-            return False
-            
-        if deployment is None:
-            return False
-            
-        full_model = deployment.get("litellm_params", {}).get("model", "")
-        is_anthropic_provider = full_model.startswith("anthropic/") or "anthropic" in full_model.lower()
-        
-        return is_anthropic_provider
+        from litellm.types.utils import LlmProviders
+        return custom_llm_provider == LlmProviders.ANTHROPIC.value
     
     async def count_tokens(
         self,
         model_to_use: str,
         messages: Optional[List[Dict[str, Any]]],
+        contents: Optional[List[Dict[str, Any]]],
         deployment: Optional[Dict[str, Any]] = None,
         request_model: str = "",
-    ) -> Optional[Dict[str, Any]]:
+    ) -> Optional[TokenCountResponse]:
         from litellm.proxy.utils import count_tokens_with_anthropic_api
         
         result = await count_tokens_with_anthropic_api(
@@ -274,12 +267,13 @@ class AnthropicTokenCounter:
         )
         
         if result is not None:
-            return {
-                "total_tokens": result["total_tokens"],
-                "request_model": request_model,
-                "model_used": model_to_use,
-                "tokenizer_type": result["tokenizer_used"],
-            }
+            return TokenCountResponse(
+                total_tokens=result.get("total_tokens", 0),
+                request_model=request_model,
+                model_used=model_to_use,
+                tokenizer_type=result.get("tokenizer_used", ""),
+                original_response=result,
+            )
         
         return None
 

litellm/llms/base_llm/base_utils.py

@@ -5,14 +5,37 @@ Utility functions for base LLM classes.
 import copy
 import json
 from abc import ABC, abstractmethod
-from typing import List, Optional, Type, Union
+from typing import Any, Dict, List, Optional, Type, Union
 
 from openai.lib import _parsing, _pydantic
 from pydantic import BaseModel
 
 from litellm._logging import verbose_logger
 from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
-from litellm.types.utils import Message, ProviderSpecificModelInfo
+from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountResponse
+
+
+class BaseTokenCounter(ABC):
+    @abstractmethod
+    async def count_tokens(
+        self,
+        model_to_use: str,
+        messages: Optional[List[Dict[str, Any]]],
+        contents: Optional[List[Dict[str, Any]]],
+        deployment: Optional[Dict[str, Any]] = None,
+        request_model: str = "",
+    ) -> Optional[TokenCountResponse]:
+        pass
+
+    @abstractmethod
+    def should_use_token_counting_api(
+        self,
+        custom_llm_provider: Optional[str] = None,
+    ) -> bool:
+        """
+        Returns True if we should the this API for token counting for the selected `custom_llm_provider`
+        """
+        return False
 
 
 class BaseLLMModelInfo(ABC):
@@ -70,7 +93,7 @@ class BaseLLMModelInfo(ABC):
         """
         pass
 
-    def get_token_counter(self):
+    def get_token_counter(self) -> Optional[BaseTokenCounter]:
         """
         Factory method to create a token counter for this provider.
         

litellm/llms/bedrock/chat/converse_handler.py

@@ -276,8 +276,13 @@ class BedrockConverseLLM(BaseAWSLLM):
         else:
             modelId = self.encode_model_id(model_id=model)
 
-        if stream is True and "ai21" in modelId:
-            fake_stream = True
+        fake_stream = litellm.AmazonConverseConfig().should_fake_stream(
+            fake_stream=fake_stream,
+            model=model,
+            stream=stream,
+            custom_llm_provider="bedrock",
+        )
+
 
         ### SET REGION NAME ###
         aws_region_name = self._get_aws_region_name(

litellm/llms/bedrock/chat/converse_transformation.py

@@ -1154,3 +1154,37 @@ class AmazonConverseConfig(BaseConfig):
         if api_key:
             headers["Authorization"] = f"Bearer {api_key}"
         return headers
+    
+
+    def should_fake_stream(
+        self,
+        model: Optional[str],
+        stream: Optional[bool],
+        custom_llm_provider: Optional[str] = None,
+        fake_stream: Optional[bool] = None,
+    ) -> bool:
+        """
+        Returns True if the model/provider should fake stream
+        """
+        ###################################################################
+        # If an upstream method already set fake_stream to True, return True
+        ###################################################################
+        if fake_stream is True:
+            return True
+
+        ###################################################################
+        # Bedrock Converse Specific Logic
+        ###################################################################
+        if stream is True:
+            if model is not None:
+                ###################################################################
+                # GPT-OSS models do not support streaming
+                ###################################################################
+                if "gpt-oss" in model:
+                    return True
+                ###################################################################
+                # AI21 models do not support streaming
+                ###################################################################
+                if "ai21" in model:
+                    return True
+        return False

litellm/llms/gemini/common_utils.py

@@ -1,15 +1,16 @@
 import base64
 import datetime
-from typing import Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import httpx
 
 import litellm
 from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
-from litellm.llms.base_llm.base_utils import BaseLLMModelInfo
+from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import AllMessageValues
+from litellm.types.utils import TokenCountResponse
 
 
 class GeminiError(BaseLLMException):
@@ -89,6 +90,16 @@ class GeminiModelInfo(BaseLLMModelInfo):
         return GeminiError(
             status_code=status_code, message=error_message, headers=headers
         )
+    
+    def get_token_counter(self) -> Optional[BaseTokenCounter]:
+        """
+        Factory method to create a token counter for this provider.
+        
+        Returns:
+            Optional TokenCounterInterface implementation for this provider,
+            or None if token counting is not supported.
+        """
+        return GoogleAIStudioTokenCounter()
 
 
 def encode_unserializable_types(
@@ -137,3 +148,46 @@ def encode_unserializable_types(
 
 def get_api_key_from_env() -> Optional[str]:
     return get_secret_str("GOOGLE_API_KEY") or get_secret_str("GEMINI_API_KEY")
+
+
+class GoogleAIStudioTokenCounter(BaseTokenCounter):
+    """Token counter implementation for Google AI Studio provider."""
+    def should_use_token_counting_api(
+        self, 
+        custom_llm_provider: Optional[str] = None,
+    ) -> bool:
+        from litellm.types.utils import LlmProviders
+        return custom_llm_provider == LlmProviders.GEMINI.value
+    
+    async def count_tokens(
+        self,
+        model_to_use: str,
+        messages: Optional[List[Dict[str, Any]]],
+        contents: Optional[List[Dict[str, Any]]],
+        deployment: Optional[Dict[str, Any]] = None,
+        request_model: str = "",
+    ) -> Optional[TokenCountResponse]:
+        import copy
+
+        from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter
+        deployment = deployment or {}
+        count_tokens_params_request = copy.deepcopy(deployment.get("litellm_params", {}))
+        count_tokens_params = {
+            "model": model_to_use,
+            "contents": contents,
+        }
+        count_tokens_params_request.update(count_tokens_params)
+        result = await GoogleAIStudioTokenCounter().acount_tokens(
+            **count_tokens_params_request,
+        )
+        
+        if result is not None:
+            return TokenCountResponse(
+                total_tokens=result.get("totalTokens", 0),
+                request_model=request_model,
+                model_used=model_to_use,
+                tokenizer_type=result.get("tokenizer_used", ""),
+                original_response=result,
+            )
+        
+        return None

litellm/llms/gemini/count_tokens/handler.py

@@ -0,0 +1,125 @@
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+
+import httpx
+
+import litellm
+from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+from litellm.types.utils import LlmProviders
+
+if TYPE_CHECKING:
+    from litellm.types.google_genai.main import GenerateContentContentListUnionDict
+else:
+    GenerateContentContentListUnionDict = Any
+
+class GoogleAIStudioTokenCounter:
+    def validate_environment(
+        self,
+        api_key: Optional[str] = None,
+        headers: Optional[Dict[str, Any]] = None,
+        model: str = "",
+        litellm_params: Optional[Dict[str, Any]] = None,
+    ):
+        from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig
+        return GoogleGenAIConfig().validate_environment(
+            api_key=api_key,
+            headers=headers,
+            model=model,
+            litellm_params=litellm_params,
+        )
+
+
+    async def acount_tokens(
+        self,
+        contents: Any,
+        model: str,
+        api_key: Optional[str] = None,
+        api_base: Optional[str] = None,
+        timeout: Optional[Union[float, httpx.Timeout]] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """
+        Count tokens using Google Gen AI Studio countTokens endpoint.
+        
+        Args:
+            contents: The content to count tokens for (Google Gen AI format)
+                    Example: [{"parts": [{"text": "Hello world"}]}]
+            model: The model name (e.g. "gemini-1.5-flash")
+            api_key: Optional Google API key (will fall back to environment)
+            api_base: Optional API base URL (defaults to Google Gen AI Studio)
+            timeout: Optional timeout for the request
+            **kwargs: Additional parameters
+            
+        Returns:
+            Dict containing token count information from Google Gen AI Studio API.
+            Example response:
+            {
+                "totalTokens": 31,
+                "totalBillableCharacters": 96,
+                "promptTokensDetails": [
+                    {
+                        "modality": "TEXT",
+                        "tokenCount": 31
+                    }
+                ]
+            }
+            
+        Raises:
+            ValueError: If API key is missing
+            litellm.APIError: If the API call fails
+            litellm.APIConnectionError: If the connection fails
+            Exception: For any other unexpected errors
+        """
+        # Set up API base URL
+        base_url = api_base or "https://generativelanguage.googleapis.com"
+        url = f"{base_url}/v1beta/models/{model}:countTokens"
+        
+        # Prepare headers
+        headers = self.validate_environment(
+            api_key=api_key,
+            headers={},
+            model=model,
+            litellm_params=kwargs,
+        )
+        
+        # Prepare request body
+        request_body = {
+            "contents": contents
+        }
+        
+        async_httpx_client = get_async_httpx_client(
+            llm_provider=LlmProviders.GEMINI,
+        )
+        
+        try:
+            response = await async_httpx_client.post(
+                url=url,
+                headers=headers,
+                json=request_body
+            )
+            
+            # Check for HTTP errors
+            response.raise_for_status()
+            
+            # Parse response
+            result = response.json()
+            return result
+                
+        except httpx.HTTPStatusError as e:
+            error_msg = f"Google Gen AI Studio API error: {e.response.status_code} - {e.response.text}"
+            raise litellm.APIError(
+                message=error_msg,
+                llm_provider="gemini",
+                model=model,
+                status_code=e.response.status_code
+            ) from e
+        except httpx.RequestError as e:
+            error_msg = f"Request to Google Gen AI Studio failed: {str(e)}"
+            raise litellm.APIConnectionError(
+                message=error_msg,
+                llm_provider="gemini",
+                model=model
+            ) from e
+        except Exception as e:
+            error_msg = f"Unexpected error during token counting: {str(e)}"
+            raise Exception(error_msg) from e
+

litellm/llms/sambanova/common_utils.py

@@ -0,0 +1,6 @@
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+
+
+class SambaNovaError(BaseLLMException):
+    def __init__(self, status_code, message, headers):
+        super().__init__(status_code=status_code, message=message, headers=headers)

litellm/llms/sambanova/embedding/handler.py

@@ -0,0 +1,5 @@
+"""
+SambaNova Embedding - uses `llm_http_handler.py` to make httpx requests
+
+Request/Response transformation is handled in `transformation.py`
+"""

litellm/llms/sambanova/embedding/transformation.py

@@ -0,0 +1,139 @@
+"""
+This is OpenAI compatible - no transformation is applied
+
+"""
+from typing import List, Optional, Union
+
+import httpx
+
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
+from litellm.types.utils import EmbeddingResponse, Usage
+
+from ..common_utils import SambaNovaError
+
+
+class SambaNovaEmbeddingConfig(BaseEmbeddingConfig):
+    def __init__(self) -> None:
+        pass
+
+    def get_complete_url(
+        self,
+        api_base: Optional[str],
+        api_key: Optional[str],
+        model: str,
+        optional_params: dict,
+        litellm_params: dict,
+        stream: Optional[bool] = None,
+    ) -> str:
+        if api_base is None:
+            raise ValueError("api_base is required for SambaNova embeddings")
+        # Remove trailing slashes and ensure clean base URL
+        api_base = api_base.rstrip("/")
+        if not api_base.endswith("/embeddings"):
+            api_base = f"{api_base}/embeddings"
+        return api_base
+
+    def validate_environment(
+        self,
+        headers: dict,
+        model: str,
+        messages: List[AllMessageValues],
+        optional_params: dict,
+        litellm_params: dict,
+        api_key: Optional[str] = None,
+        api_base: Optional[str] = None,
+    ) -> dict:
+        if api_key is None:
+            api_key = get_secret_str("SAMBANOVA_API_KEY")
+
+        default_headers = {
+            "Authorization": f"Bearer {api_key}",
+            "accept": "application/json",
+            "Content-Type": "application/json",
+        }
+
+        # If 'Authorization' is provided in headers, it overrides the default.
+        if "Authorization" in headers:
+            default_headers["Authorization"] = headers["Authorization"]
+
+        # Merge other headers, overriding any default ones except Authorization
+        return {**default_headers, **headers}
+
+    def get_supported_openai_params(self, model: str):
+        """
+        Non additional params supported, placeholder method for future supported params
+        https://docs.sambanova.ai/cloud/api-reference/endpoints/embeddings-api
+        """
+        return []
+
+    def map_openai_params(
+        self,
+        non_default_params: dict,
+        optional_params: dict,
+        model: str,
+        drop_params: bool,
+    ):
+        """
+        No transformation is applied - SambaNova is openai compatible
+        """
+        supported_openai_params = self.get_supported_openai_params(model)
+        for param, value in non_default_params.items():
+            if param in supported_openai_params:
+                optional_params[param] = value
+        return optional_params
+
+    def transform_embedding_request(
+        self,
+        model: str,
+        input: AllEmbeddingInputValues,
+        optional_params: dict,
+        headers: dict,
+    ) -> dict:
+        return {
+            "input": input,
+            "model": model,
+            **optional_params,
+        }
+
+    def transform_embedding_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        model_response: EmbeddingResponse,
+        logging_obj: LiteLLMLoggingObj,
+        api_key: Optional[str],
+        request_data: dict,
+        optional_params: dict,
+        litellm_params: dict,
+    ) -> EmbeddingResponse:
+        try:
+            raw_response_json = raw_response.json()
+        except Exception:
+            raise SambaNovaError(
+                message=raw_response.text,
+                status_code=raw_response.status_code,
+                headers=raw_response.headers,
+            )
+
+        model_response.model = raw_response_json.get("model")
+        model_response.data = raw_response_json.get("data")
+        model_response.object = raw_response_json.get("object")
+
+        usage = Usage(
+            prompt_tokens=raw_response_json.get("usage", {}).get("prompt_tokens", 0),
+            total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
+        )
+
+        model_response.usage = usage
+        return model_response
+
+    def get_error_class(
+        self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
+    ) -> BaseLLMException:
+        return SambaNovaError(
+            message=error_message, status_code=status_code, headers=headers
+        )

litellm/main.py

@@ -1890,7 +1890,7 @@ def completion(  # type: ignore # noqa: PLR0915
                 or get_secret_str("COMETAPI_KEY")
                 or litellm.api_key
             )
-            
+
             api_base = (
                 api_base
                 or litellm.api_base
@@ -1917,7 +1917,7 @@ def completion(  # type: ignore # noqa: PLR0915
                 stream=stream,
                 provider_config=provider_config,
             )
-            
+
             ## LOGGING
             logging.post_call(
                 input=messages, api_key=api_key, original_response=response
@@ -4279,6 +4279,28 @@ def embedding(  # noqa: PLR0915
                 client=client,
                 aembedding=aembedding,
             )
+        elif custom_llm_provider == "sambanova":
+            api_key = api_key or litellm.api_key or get_secret_str("SAMBANOVA_API_KEY")
+            api_base = (
+                api_base
+                or litellm.api_base
+                or get_secret_str("SAMBANOVA_API_BASE")
+                or "https://api.sambanova.ai/v1"
+            )
+            response = base_llm_http_handler.embedding(
+                model=model,
+                input=input,
+                custom_llm_provider=custom_llm_provider,
+                api_base=api_base,
+                api_key=api_key,
+                logging_obj=logging,
+                timeout=timeout,
+                model_response=EmbeddingResponse(),
+                optional_params=optional_params,
+                client=client,
+                aembedding=aembedding,
+                litellm_params={},
+            )
         elif custom_llm_provider == "voyage":
             response = base_llm_http_handler.embedding(
                 model=model,

litellm/model_prices_and_context_window_backup.json

@@ -12588,8 +12588,6 @@
         "output_cost_per_token": 3e-07,
         "litellm_provider": "bedrock_converse",
         "mode": "chat",
-        "supports_function_calling": true,
-        "supports_vision": true,
         "supports_response_schema": true,
         "supports_tool_choice": true,
         "supports_reasoning": true
@@ -12602,8 +12600,6 @@
         "output_cost_per_token": 6e-07,
         "litellm_provider": "bedrock_converse",
         "mode": "chat",
-        "supports_function_calling": true,
-        "supports_vision": true,
         "supports_response_schema": true,
         "supports_tool_choice": true,
         "supports_reasoning": true

litellm/proxy/_types.py

@@ -2094,13 +2094,15 @@ class TokenCountRequest(LiteLLMPydanticObjectBase):
     model: str
     prompt: Optional[str] = None
     messages: Optional[List[dict]] = None
+    """
+    Anthropic token counting endpoint uses /messages
+    """
 
-
-class TokenCountResponse(LiteLLMPydanticObjectBase):
-    total_tokens: int
-    request_model: str
-    model_used: str
-    tokenizer_type: str
+    
+    contents: Optional[List[dict]] = None
+    """
+    Google /countTokens endpoint expects contents to be a list of dicts with the following structure:
+    """
 
 
 class CallInfo(LiteLLMPydanticObjectBase):

litellm/proxy/anthropic_endpoints/endpoints.py

@@ -16,6 +16,7 @@ from litellm.proxy.common_request_processing import (
 )
 from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
 from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
+from litellm.types.utils import TokenCountResponse
 
 router = APIRouter()
 
@@ -262,10 +263,18 @@ async def count_tokens(
         )
         
         # Call the internal token counter function with direct request flag set to False
-        token_response = await internal_token_counter(token_request, is_direct_request=False)
-        
+        token_response = await internal_token_counter(
+            request=token_request,
+            call_endpoint=True,
+        )
+        _token_response_dict: dict = {}
+        if isinstance(token_response, TokenCountResponse):
+            _token_response_dict = token_response.model_dump()
+        elif isinstance(token_response, dict):
+            _token_response_dict = token_response
+    
         # Convert the internal response to Anthropic API format
-        return {"input_tokens": token_response.total_tokens}
+        return {"input_tokens": _token_response_dict.get("total_tokens", 0)}
         
     except HTTPException:
         raise

litellm/proxy/google_endpoints/endpoints.py

@@ -3,6 +3,7 @@ from fastapi import APIRouter, Depends, Request, Response
 from litellm.proxy._types import *
 from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
 from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
 
 router = APIRouter(
     tags=["google genai endpoints"],
@@ -145,10 +146,61 @@ async def google_stream_generate_content(
 
 
 
-@router.post("/v1beta/models/{model_name}:countTokens", dependencies=[Depends(user_api_key_auth)])
-@router.post("/models/{model_name}:countTokens", dependencies=[Depends(user_api_key_auth)])
+@router.post(
+    "/v1beta/models/{model_name}:countTokens",
+    dependencies=[Depends(user_api_key_auth)],
+    response_model=TokenCountDetailsResponse,
+)
+@router.post(
+    "/models/{model_name}:countTokens",
+    dependencies=[Depends(user_api_key_auth)],
+    response_model=TokenCountDetailsResponse,
+)
 async def google_count_tokens(request: Request, model_name: str):
     """
-    Not Implemented, this is a placeholder for the google genai countTokens endpoint.
+    ```json
+    return {
+        "totalTokens": 31,
+        "totalBillableCharacters": 96,
+        "promptTokensDetails": [
+            {
+            "modality": "TEXT",
+            "tokenCount": 31
+            }
+        ]
+    }
+    ```
     """
-    return {}
+    from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
+    from litellm.proxy.proxy_server import token_counter as internal_token_counter
+
+    data = await _read_request_body(request=request)
+    contents = data.get("contents", [])
+    #Create TokenCountRequest for the internal endpoint
+    from litellm.proxy._types import TokenCountRequest
+
+    token_request = TokenCountRequest(
+        model=model_name,
+        contents=contents
+    )
+
+    # Call the internal token counter function with direct request flag set to False
+    token_response = await internal_token_counter(
+        request=token_request,
+        call_endpoint=True,
+    )
+    if token_response is not None:
+        # cast the response to the well known format
+        original_response: dict = token_response.original_response or {}
+        return TokenCountDetailsResponse(
+            totalTokens=original_response.get("totalTokens", 0),
+            promptTokensDetails=original_response.get("promptTokensDetails", []),
+        )
+    
+    #########################################################
+    # Return the response in the well known format
+    #########################################################
+    return TokenCountDetailsResponse(
+        totalTokens=0,
+        promptTokensDetails=[],
+    )

litellm/proxy/proxy_config.yaml

@@ -1,11 +1,4 @@
 model_list:
-  - model_name: openai/*
+  - model_name: vertex_ai/*
     litellm_params:
-      model: openai/*
-  - model_name: anthropic/*
-    litellm_params:
-      model: anthropic/*
-
-litellm_settings:
-  callbacks:
-    - langfuse_otel
+      model: gemini/*

litellm/proxy/proxy_server.py

@@ -11,7 +11,6 @@ import time
 import traceback
 import uuid
 import warnings
-from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from datetime import datetime, timedelta
 from typing import (
     TYPE_CHECKING,
@@ -35,10 +34,12 @@ from litellm.constants import (
     LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS,
     LITELLM_SETTINGS_SAFE_DB_OVERRIDES,
 )
+from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.types.utils import (
     ModelResponse,
     ModelResponseStream,
     TextCompletionResponse,
+    TokenCountResponse,
 )
 
 if TYPE_CHECKING:
@@ -2999,7 +3000,9 @@ class ProxyConfig:
             
             if should_reload:
                 # Perform the reload
-                from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
+                from litellm.litellm_core_utils.get_model_cost_map import (
+                    get_model_cost_map,
+                )
                 model_cost_map_url = litellm.model_cost_map_url
                 new_model_cost_map = get_model_cost_map(url=model_cost_map_url)
                 litellm.model_cost = new_model_cost_map
@@ -5742,9 +5745,10 @@ async def run_thread(
 #     dependencies=[Depends(user_api_key_auth)],
 # )
 # async def get_available_routes(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)):
+from litellm.llms.base_llm.base_utils import BaseTokenCounter
 
 
-def _get_provider_token_counter(deployment: dict, model_to_use: str):
+def _get_provider_token_counter(deployment: dict, model_to_use: str) -> Tuple[Optional[BaseTokenCounter], Optional[str], Optional[str]]:
     """
     Auto-route to the correct provider's token counter based on model/deployment.
     Uses the existing get_provider_model_info infrastructure with switch-case pattern.
@@ -5755,10 +5759,12 @@ def _get_provider_token_counter(deployment: dict, model_to_use: str):
     from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
 
     full_model = deployment.get("litellm_params", {}).get("model", "")
+    model: Optional[str] = None
+    custom_llm_provider: Optional[str] = None
 
     try:
         # Use existing LiteLLM logic to determine provider
-        model, provider, dynamic_api_key, api_base = get_llm_provider(
+        model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider(
             model=full_model,
             custom_llm_provider=deployment.get("litellm_params", {}).get(
                 "custom_llm_provider"
@@ -5772,7 +5778,7 @@ def _get_provider_token_counter(deployment: dict, model_to_use: str):
         from litellm.utils import ProviderConfigManager
 
         # Convert string provider to LlmProviders enum
-        llm_provider_enum = LlmProviders(provider)
+        llm_provider_enum = LlmProviders(custom_llm_provider)
         # Add more provider mappings as needed
 
         if llm_provider_enum:
@@ -5780,7 +5786,7 @@ def _get_provider_token_counter(deployment: dict, model_to_use: str):
                 model=full_model, provider=llm_provider_enum
             )
             if provider_model_info is not None:
-                return provider_model_info.get_token_counter()
+                return provider_model_info.get_token_counter(), model, custom_llm_provider
 
     except Exception:
         # If provider detection fails, fall back to manual checks
@@ -5788,9 +5794,9 @@ def _get_provider_token_counter(deployment: dict, model_to_use: str):
             from litellm.llms.anthropic.common_utils import AnthropicModelInfo
 
             anthropic_model_info = AnthropicModelInfo()
-            return anthropic_model_info.get_token_counter()
+            return anthropic_model_info.get_token_counter(), model, custom_llm_provider
 
-    return None
+    return None, None, None
 
 
 @router.post(
@@ -5799,62 +5805,82 @@ def _get_provider_token_counter(deployment: dict, model_to_use: str):
     dependencies=[Depends(user_api_key_auth)],
     response_model=TokenCountResponse,
 )
-async def token_counter(request: TokenCountRequest, is_direct_request: bool = True):
-    """ """
+async def token_counter(
+    request: TokenCountRequest, 
+    call_endpoint: bool = False
+):
+    """
+    Args:
+        request: TokenCountRequest
+        call_endpoint: bool - When set to "True" it will call the token counting endpoint - e.g Anthropic or Google AI Studio Token Counting APIs.
+
+    Returns:
+        TokenCountResponse
+    """
     from litellm import token_counter
 
     global llm_router
 
     prompt = request.prompt
     messages = request.messages
-    if prompt is None and messages is None:
+    contents = request.contents
+
+    #########################################################
+    # Validate request
+    #########################################################
+    if prompt is None and messages is None and contents is None:
         raise HTTPException(
-            status_code=400, detail="prompt or messages must be provided"
+            status_code=400, detail="prompt or messages or contents must be provided"
         )
 
-    deployment = None
+    deployment: Optional[Dict[str, Any]] = None
     litellm_model_name = None
     model_info: Optional[ModelMapInfo] = None
     if llm_router is not None:
         # get 1 deployment corresponding to the model
-        for _model in llm_router.model_list:
-            if _model["model_name"] == request.model:
-                deployment = _model
-                model_info = deployment.get("model_info", {})
-                break
+        try:
+            deployment = await llm_router.async_get_available_deployment(
+                model=request.model,
+                request_kwargs={},
+            )
+        except Exception:
+            verbose_proxy_logger.exception(
+                "litellm.proxy.proxy_server.token_counter(): Exception occured while getting deployment"
+            )
+            pass
     if deployment is not None:
         litellm_model_name = deployment.get("litellm_params", {}).get("model")
         # remove the custom_llm_provider_prefix in the litellm_model_name
         if "/" in litellm_model_name:
             litellm_model_name = litellm_model_name.split("/", 1)[1]
 
-    model_to_use = (
+    model_to_use: str = (
         litellm_model_name or request.model
     )  # use litellm model name, if it's not avalable then fallback to request.model
 
     # Try provider-specific token counting first - only for non-direct requests (from provider endpoints)
-    provider_counter = None
-    if deployment is not None and not is_direct_request:
+    provider_counter: Optional[BaseTokenCounter] = None
+    custom_llm_provider: Optional[str] = None
+    if call_endpoint is True and deployment is not None:
         # Auto-route to the correct provider based on model
-        provider_counter = _get_provider_token_counter(deployment, model_to_use)
+        provider_counter, _model, custom_llm_provider = _get_provider_token_counter(deployment, model_to_use)
+        if _model is not None:
+            model_to_use = _model
 
-    if provider_counter is not None and provider_counter.supports_provider(
-        deployment=deployment, from_endpoint=not is_direct_request
-    ):
-        result = await provider_counter.count_tokens(
-            model_to_use=model_to_use,
-            messages=messages,  # type: ignore
-            deployment=deployment,
-            request_model=request.model,
-        )
-
-        if result is not None:
-            return TokenCountResponse(
-                total_tokens=result["total_tokens"],
-                request_model=result["request_model"],
-                model_used=result["model_used"],
-                tokenizer_type=result["tokenizer_type"],
+    if provider_counter is not None:
+        if provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider) is True:
+            result = await provider_counter.count_tokens(
+                model_to_use=model_to_use or "",
+                messages=messages,  # type: ignore
+                contents=contents,
+                deployment=deployment,
+                request_model=request.model,
             )
+            #########################################################
+            # Transfrom the Response to the well known format
+            #########################################################
+            if result is not None:
+                return result
 
     # Default LiteLLM token counting
     custom_tokenizer: Optional[CustomHuggingfaceTokenizer] = None

litellm/types/google_genai/main.py

@@ -1,9 +1,10 @@
 # Import types from the Google GenAI SDK
-from typing import TYPE_CHECKING, Any, Optional, TypeAlias, TypedDict
+from typing import TYPE_CHECKING, Any, List, Optional, TypeAlias
 
 # During static type-checking we can rely on the real google-genai types.
 from google.genai import types as _genai_types  # type: ignore
 from pydantic import BaseModel
+from typing_extensions import TypedDict
 
 from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject
 
@@ -19,7 +20,7 @@ ToolConfigDict = _genai_types.ToolConfigDict
 
 class GenerateContentRequestDict(GenerateContentRequestParametersDict):  # type: ignore[misc]
     generationConfig: Optional[Any]
-    tools: Optional[ToolConfigDict]
+    tools: Optional[ToolConfigDict] # type: ignore[assignment]
 
 
 class GenerateContentResponse(GoogleGenAIGenerateContentResponse, BaseLiteLLMOpenAIResponseObject): # type: ignore[misc]

litellm/types/llms/vertex_ai.py

@@ -1,11 +1,12 @@
 import json
 from enum import Enum
-from typing import Any, Dict, List, Literal, Optional, Tuple, TypedDict, Union
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 from typing_extensions import (
     Protocol,
     Required,
     Self,
+    TypedDict,
     TypeGuard,
     get_origin,
     override,
@@ -241,6 +242,17 @@ class UsageMetadata(TypedDict, total=False):
     responseTokensDetails: List[PromptTokensDetails]
 
 
+class TokenCountDetailsResponse(TypedDict):
+    """
+    Response structure for token count details with modality breakdown.
+    
+    Example:
+        {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]}
+    """
+    totalTokens: int
+    promptTokensDetails: List[PromptTokensDetails]
+
+
 class CachedContent(TypedDict, total=False):
     ttl: TTL
     expire_time: str

litellm/types/utils.py

@@ -2356,6 +2356,17 @@ class LiteLLMLoggingBaseClass:
         pass
 
 
+class TokenCountResponse(LiteLLMPydanticObjectBase):
+    total_tokens: int
+    request_model: str
+    model_used: str
+    tokenizer_type: str
+    original_response: Optional[dict] = None
+    """
+    Original Response from upstream API call - if an API call was made for token counting
+    """
+
+
 class CustomHuggingfaceTokenizer(TypedDict):
     identifier: str
     revision: str  # usually 'main'

litellm/utils.py

@@ -2831,6 +2831,19 @@ def get_optional_params_embeddings(  # noqa: PLR0915
         optional_params = litellm.FireworksAIEmbeddingConfig().map_openai_params(
             non_default_params=non_default_params, optional_params={}, model=model
         )
+    elif custom_llm_provider == "sambanova":
+        supported_params = get_supported_openai_params(
+            model=model,
+            custom_llm_provider="sambanova",
+            request_type="embeddings",
+        )
+        _check_valid_arg(supported_params=supported_params)
+        optional_params = litellm.SambaNovaEmbeddingConfig().map_openai_params(
+            non_default_params=non_default_params,
+            optional_params={},
+            model=model,
+            drop_params=drop_params if drop_params is not None else False,
+        )
 
     elif (
         custom_llm_provider != "openai"
@@ -7006,6 +7019,8 @@ class ProviderConfigManager:
             return litellm.IBMWatsonXEmbeddingConfig()
         elif litellm.LlmProviders.INFINITY == provider:
             return litellm.InfinityEmbeddingConfig()
+        elif litellm.LlmProviders.SAMBANOVA == provider:
+            return litellm.SambaNovaEmbeddingConfig()
         elif (
             litellm.LlmProviders.COHERE == provider
             or litellm.LlmProviders.COHERE_CHAT == provider

model_prices_and_context_window.json

@@ -12588,8 +12588,6 @@
         "output_cost_per_token": 3e-07,
         "litellm_provider": "bedrock_converse",
         "mode": "chat",
-        "supports_function_calling": true,
-        "supports_vision": true,
         "supports_response_schema": true,
         "supports_tool_choice": true,
         "supports_reasoning": true
@@ -12602,8 +12600,6 @@
         "output_cost_per_token": 6e-07,
         "litellm_provider": "bedrock_converse",
         "mode": "chat",
-        "supports_function_calling": true,
-        "supports_vision": true,
         "supports_response_schema": true,
         "supports_tool_choice": true,
         "supports_reasoning": true

tests/llm_translation/test_bedrock_gpt_oss.py

@@ -0,0 +1,27 @@
+from base_llm_unit_tests import BaseLLMChatTest
+import pytest
+import sys
+import os
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+import litellm
+
+
+class TestBedrockGPTOSS(BaseLLMChatTest):
+    def get_base_completion_call_args(self) -> dict:
+        litellm._turn_on_debug()
+        return {
+            "model": "bedrock/converse/openai.gpt-oss-20b-1:0",
+        }
+    
+    def test_tool_call_no_arguments(self, tool_call_no_arguments):
+        """Test that tool calls with no arguments is translated correctly. Relevant issue: https://github.com/BerriAI/litellm/issues/6833"""
+        pass
+
+    def test_prompt_caching(self):
+        """
+        Remove override once we have access to Bedrock prompt caching
+        """
+        pass

tests/proxy_unit_tests/test_proxy_token_counter.py

@@ -24,7 +24,8 @@ from litellm._logging import verbose_proxy_logger
 
 verbose_proxy_logger.setLevel(level=logging.DEBUG)
 
-from litellm.proxy._types import TokenCountRequest, TokenCountResponse
+from litellm.proxy._types import TokenCountRequest
+from litellm.types.utils import TokenCountResponse
 
 
 from litellm import Router
@@ -169,12 +170,12 @@ async def test_anthropic_messages_count_tokens_endpoint():
     anthropic_endpoints._read_request_body = mock_read_request_body
     
     # Mock the internal token_counter function to return a controlled response
-    async def mock_token_counter(request, is_direct_request=True):
-        assert is_direct_request == False, "Should be called with is_direct_request=False for Anthropic endpoint"
+    async def mock_token_counter(request, call_endpoint=False):
+        assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint"
         assert request.model == "claude-3-sonnet-20240229"
         assert request.messages == [{"role": "user", "content": "Hello Claude!"}]
         
-        from litellm.proxy._types import TokenCountResponse
+        from litellm.types.utils import TokenCountResponse
         return TokenCountResponse(
             total_tokens=15,
             request_model="claude-3-sonnet-20240229",
@@ -236,12 +237,12 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model():
     anthropic_endpoints._read_request_body = mock_read_request_body
     
     # Mock the internal token_counter function to return a controlled response
-    async def mock_token_counter(request, is_direct_request=True):
-        assert is_direct_request == False, "Should be called with is_direct_request=False for Anthropic endpoint"
+    async def mock_token_counter(request, call_endpoint=True):
+        assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint"
         assert request.model == "gpt-4"
         assert request.messages == [{"role": "user", "content": "Hello GPT!"}]
         
-        from litellm.proxy._types import TokenCountResponse
+        from litellm.types.utils import TokenCountResponse
         return TokenCountResponse(
             total_tokens=12,
             request_model="gpt-4", 
@@ -300,7 +301,7 @@ async def test_internal_token_counter_anthropic_provider_detection():
             model="claude-test",
             messages=[{"role": "user", "content": "hello"}],
         ),
-        is_direct_request=False
+        call_endpoint=True
     )
     
     print("Anthropic provider test response:", response)
@@ -330,7 +331,7 @@ async def test_internal_token_counter_anthropic_provider_detection():
             model="gpt-test",
             messages=[{"role": "user", "content": "hello"}],
         ),
-        is_direct_request=False
+        call_endpoint=True
     )
     
     print("Non-Anthropic provider test response:", response)
@@ -385,7 +386,7 @@ async def test_anthropic_endpoint_error_handling():
 @pytest.mark.asyncio
 async def test_factory_anthropic_endpoint_calls_anthropic_counter():
     """Test that /v1/messages/count_tokens with Anthropic model uses Anthropic counter."""
-    from unittest.mock import patch
+    from unittest.mock import patch, AsyncMock
     from fastapi.testclient import TestClient
     from litellm.proxy.proxy_server import app
     
@@ -404,6 +405,13 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter():
                 "model_info": {}
             }]
             
+            # Mock the async method properly
+            mock_router.async_get_available_deployment = AsyncMock(return_value={
+                "model_name": "claude-3-5-sonnet",
+                "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"},
+                "model_info": {}
+            })
+            
             client = TestClient(app)
             
             response = client.post(
@@ -426,7 +434,7 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter():
 @pytest.mark.asyncio
 async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter():
     """Test that /v1/messages/count_tokens with GPT-4 does NOT use Anthropic counter."""
-    from unittest.mock import patch
+    from unittest.mock import patch, AsyncMock
     from fastapi.testclient import TestClient
     from litellm.proxy.proxy_server import app
     
@@ -444,6 +452,13 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter():
                     "model_info": {}
                 }]
                 
+                # Mock the async method properly
+                mock_router.async_get_available_deployment = AsyncMock(return_value={
+                    "model_name": "gpt-4",
+                    "litellm_params": {"model": "openai/gpt-4"},
+                    "model_info": {}
+                })
+                
                 client = TestClient(app)
                 
                 response = client.post(
@@ -466,7 +481,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter():
 @pytest.mark.asyncio
 async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic():
     """Test that /utils/token_counter does NOT use Anthropic counter even with Anthropic model."""
-    from unittest.mock import patch
+    from unittest.mock import patch, AsyncMock
     from fastapi.testclient import TestClient
     from litellm.proxy.proxy_server import app
     
@@ -484,6 +499,13 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic():
                     "model_info": {}
                 }]
                 
+                # Mock the async method properly
+                mock_router.async_get_available_deployment = AsyncMock(return_value={
+                    "model_name": "claude-3-5-sonnet",
+                    "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"},
+                    "model_info": {}
+                })
+                
                 client = TestClient(app)
                 
                 response = client.post(
@@ -499,7 +521,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic():
                 data = response.json()
                 assert data["total_tokens"] == 35
                 
-                # Verify that Anthropic API was NOT called (since is_direct_request=True)
+                # Verify that Anthropic API was NOT called (since call_endpoint=False)
                 mock_anthropic_count.assert_not_called()
 
 
@@ -523,40 +545,59 @@ async def test_factory_registration():
     }
     
     # Test Anthropic counter supports provider
-    assert counter.supports_provider(anthropic_deployment, from_endpoint=True)
-    assert not counter.supports_provider(anthropic_deployment, from_endpoint=False)
+    assert counter.should_use_token_counting_api(custom_llm_provider="anthropic")
+    assert not counter.should_use_token_counting_api(custom_llm_provider="openai")
     
     # Test non-Anthropic provider
-    assert not counter.supports_provider(non_anthropic_deployment, from_endpoint=True)
-    assert not counter.supports_provider(non_anthropic_deployment, from_endpoint=False)
+    assert not counter.should_use_token_counting_api(custom_llm_provider="openai")
     
     # Test None deployment
-    assert not counter.supports_provider(None, from_endpoint=True)
-    assert not counter.supports_provider(None, from_endpoint=False)
+    assert not counter.should_use_token_counting_api(custom_llm_provider=None)
 
 
-@pytest.mark.asyncio 
-async def test_factory_anthropic_counter_supports_provider():
-    """Test AnthropicTokenCounter provider detection logic."""
-    from litellm.llms.anthropic.common_utils import AnthropicModelInfo
+
+
+@pytest.mark.asyncio
+async def test_vertex_ai_gemini_token_counting_with_contents():
+    """
+    Test token counting for Vertex AI Gemini model using contents format with call_endpoint=True
+    """
+    llm_router = Router(
+        model_list=[
+            {
+                "model_name": "gemini-2.5-pro",
+                "litellm_params": {
+                    "model": "gemini/gemini-2.5-pro",
+                },
+            }
+        ]
+    )
     
-    anthropic_model_info = AnthropicModelInfo()
-    counter = anthropic_model_info.get_token_counter()
+    setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
     
-    # Test Anthropic provider detection
-    anthropic_deployment = {
-        "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}
-    }
-    assert counter.supports_provider(anthropic_deployment, from_endpoint=True)
-    assert not counter.supports_provider(anthropic_deployment, from_endpoint=False)
+    # Test with contents format and call_endpoint=True
+    response = await token_counter(
+        request=TokenCountRequest(
+            model="gemini-2.5-pro",
+            contents=[
+                {
+                    "parts": [
+                        {
+                            "text": "Hello world, how are you doing today? i am ij"
+                        }
+                    ]
+                }
+            ],
+        ),
+        call_endpoint=True
+    )
     
-    # Test non-Anthropic provider
-    openai_deployment = {
-        "litellm_params": {"model": "openai/gpt-4"}
-    }
-    assert not counter.supports_provider(openai_deployment, from_endpoint=True)
-    assert not counter.supports_provider(openai_deployment, from_endpoint=False)
-    
-    # Test None deployment
-    assert not counter.supports_provider(None, from_endpoint=True)
-    assert not counter.supports_provider(None, from_endpoint=False)
+    print("Vertex AI Gemini token counting response:", response)
+
+    # validate we have orignal response
+    assert response.original_response is not None
+    assert response.original_response.get("totalTokens") is not None
+    assert response.original_response.get("promptTokensDetails") is not None
+
+    prompt_tokens_details = response.original_response.get("promptTokensDetails")
+    assert prompt_tokens_details is not None

tests/test_litellm/llms/gemini/test_gemini_common_utils.py

@@ -1,6 +1,8 @@
+from unittest.mock import AsyncMock, patch
+
 import pytest
 
-from litellm.llms.gemini.common_utils import GeminiModelInfo
+from litellm.llms.gemini.common_utils import GeminiModelInfo, GoogleAIStudioTokenCounter
 
 
 class TestGeminiModelInfo:
@@ -84,3 +86,75 @@ class TestGeminiModelInfo:
         ]
 
         assert result == expected
+
+
+class TestGoogleAIStudioTokenCounter:
+    """Test suite for GoogleAIStudioTokenCounter class"""
+
+    def test_should_use_token_counting_api(self):
+        """Test should_use_token_counting_api method with different provider values"""
+        from litellm.types.utils import LlmProviders
+        
+        token_counter = GoogleAIStudioTokenCounter()
+        
+        # Test with gemini provider - should return True
+        assert token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) is True
+        
+        # Test with other providers - should return False
+        assert token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) is False
+        assert token_counter.should_use_token_counting_api("anthropic") is False
+        assert token_counter.should_use_token_counting_api("vertex_ai") is False
+        
+        # Test with None - should return False
+        assert token_counter.should_use_token_counting_api(None) is False
+
+    @pytest.mark.asyncio
+    async def test_count_tokens(self):
+        """Test count_tokens method with mocked API response"""
+        from litellm.types.utils import TokenCountResponse
+        
+        token_counter = GoogleAIStudioTokenCounter()
+        
+        # Mock the GoogleAIStudioTokenCounter from handler module
+        mock_response = {
+            "totalTokens": 31,
+            "totalBillableCharacters": 96,
+            "promptTokensDetails": [
+                {
+                    "modality": "TEXT",
+                    "tokenCount": 31
+                }
+            ]
+        }
+        
+        with patch('litellm.llms.gemini.count_tokens.handler.GoogleAIStudioTokenCounter.acount_tokens', 
+                   new_callable=AsyncMock) as mock_acount_tokens:
+            mock_acount_tokens.return_value = mock_response
+            
+            # Test data
+            model_to_use = "gemini-1.5-flash"
+            contents = [{"parts": [{"text": "Hello world"}]}]
+            request_model = "gemini/gemini-1.5-flash"
+            
+            # Call the method
+            result = await token_counter.count_tokens(
+                model_to_use=model_to_use,
+                messages=None,
+                contents=contents,
+                deployment=None,
+                request_model=request_model
+            )
+            
+            # Verify the result
+            assert result is not None
+            assert isinstance(result, TokenCountResponse)
+            assert result.total_tokens == 31
+            assert result.request_model == request_model
+            assert result.model_used == model_to_use
+            assert result.original_response == mock_response
+            
+            # Verify the mock was called correctly
+            mock_acount_tokens.assert_called_once_with(
+                model=model_to_use,
+                contents=contents
+            )

tests/test_litellm/llms/sambanova/tests_sambanova_embedding_transformation.py

@@ -0,0 +1,41 @@
+from unittest.mock import patch
+
+import litellm
+
+
+def mock_embedding_response(*args, **kwargs):
+    """Mock response mimicking litellm.embedding output."""
+
+    class MockResponse:
+        def __init__(self):
+            self.data = [{"embedding": [0.1, 0.2, 0.3]}]  # Example embedding vector
+            self.usage = litellm.Usage()  # Mock Usage object
+            self.model = kwargs.get("model", "sambanova/E5-Mistral-7B-Instruct")
+            self.object = "embedding"
+
+        def __getitem__(self, key):
+            return getattr(self, key)
+
+    return MockResponse()
+
+
+def test_sambanova_embeddings():
+    """Mocked test for SambaNova embeddings using MagicMock."""
+    with patch("litellm.embedding", side_effect=mock_embedding_response) as mock_embed:
+        response = litellm.embedding(
+            model="sambanova/E5-Mistral-7B-Instruct",
+            input=["good morning from litellm"],
+        )
+
+        # Assertions to verify that the mock was called correctly
+        mock_embed.assert_called_once_with(
+            model="sambanova/E5-Mistral-7B-Instruct",
+            input=["good morning from litellm"],
+        )
+
+        # Assertions to check the structure of the mocked response
+        assert isinstance(response.data, list)
+        assert "embedding" in response.data[0]
+        assert isinstance(response.data[0]["embedding"], list)
+        assert response.model == "sambanova/E5-Mistral-7B-Instruct"
+        assert response.object == "embedding"

ui/litellm-dashboard/src/app/page.tsx

@@ -1,234 +1,220 @@
-"use client";
+"use client"
 
-import React, { Suspense, useEffect, useState } from "react";
-import { useSearchParams } from "next/navigation";
-import { jwtDecode } from "jwt-decode";
-import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
-import { Team } from "@/components/key_team_helpers/key_list";
-import Navbar from "@/components/navbar";
-import { ThemeProvider } from "@/contexts/ThemeContext";
-import UserDashboard from "@/components/user_dashboard";
-import ModelDashboard from "@/components/model_dashboard";
-import ViewUserDashboard from "@/components/view_users";
-import Teams from "@/components/teams";
-import Organizations from "@/components/organizations";
-import { fetchOrganizations } from "@/components/organizations";
-import AdminPanel from "@/components/admins";
-import Settings from "@/components/settings";
-import GeneralSettings from "@/components/general_settings";
-import PassThroughSettings from "@/components/pass_through_settings";
-import BudgetPanel from "@/components/budgets/budget_panel";
-import SpendLogsTable from "@/components/view_logs";
-import ModelHubTable from "@/components/model_hub_table";
-import NewUsagePage from "@/components/new_usage";
-import APIRef from "@/components/api_ref";
-import ChatUI from "@/components/chat_ui";
-import Sidebar from "@/components/leftnav";
-import Usage from "@/components/usage";
-import CacheDashboard from "@/components/cache_dashboard";
-import {
-  getUiConfig,
-  proxyBaseUrl,
-  setGlobalLitellmHeaderName,
-} from "@/components/networking";
-import { Organization } from "@/components/networking";
-import GuardrailsPanel from "@/components/guardrails";
-import PromptsPanel from "@/components/prompts";
-import TransformRequestPanel from "@/components/transform_request";
-import { fetchUserModels } from "@/components/create_key_button";
-import { fetchTeams } from "@/components/common_components/fetch_teams";
-import { MCPServers } from "@/components/mcp_tools";
-import TagManagement from "@/components/tag_management";
-import VectorStoreManagement from "@/components/vector_store_management";
-import UIThemeSettings from "@/components/ui_theme_settings";
-import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
-import { cx } from "@/lib/cva.config";
+import React, { Suspense, useEffect, useState } from "react"
+import { useSearchParams } from "next/navigation"
+import { jwtDecode } from "jwt-decode"
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query"
+import { Team } from "@/components/key_team_helpers/key_list"
+import Navbar from "@/components/navbar"
+import { ThemeProvider } from "@/contexts/ThemeContext"
+import UserDashboard from "@/components/user_dashboard"
+import ModelDashboard from "@/components/model_dashboard"
+import ViewUserDashboard from "@/components/view_users"
+import Teams from "@/components/teams"
+import Organizations from "@/components/organizations"
+import { fetchOrganizations } from "@/components/organizations"
+import AdminPanel from "@/components/admins"
+import Settings from "@/components/settings"
+import GeneralSettings from "@/components/general_settings"
+import PassThroughSettings from "@/components/pass_through_settings"
+import BudgetPanel from "@/components/budgets/budget_panel"
+import SpendLogsTable from "@/components/view_logs"
+import ModelHubTable from "@/components/model_hub_table"
+import NewUsagePage from "@/components/new_usage"
+import APIRef from "@/components/api_ref"
+import ChatUI from "@/components/chat_ui"
+import Sidebar from "@/components/leftnav"
+import Usage from "@/components/usage"
+import CacheDashboard from "@/components/cache_dashboard"
+import { getUiConfig, proxyBaseUrl, setGlobalLitellmHeaderName } from "@/components/networking"
+import { Organization } from "@/components/networking"
+import GuardrailsPanel from "@/components/guardrails"
+import PromptsPanel from "@/components/prompts"
+import TransformRequestPanel from "@/components/transform_request"
+import { fetchUserModels } from "@/components/organisms/create_key_button"
+import { fetchTeams } from "@/components/common_components/fetch_teams"
+import { MCPServers } from "@/components/mcp_tools"
+import TagManagement from "@/components/tag_management"
+import VectorStoreManagement from "@/components/vector_store_management"
+import UIThemeSettings from "@/components/ui_theme_settings"
+import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"
+import { cx } from "@/lib/cva.config"
 
 function getCookie(name: string) {
-  const cookieValue = document.cookie
-    .split("; ")
-    .find((row) => row.startsWith(name + "="));
-  return cookieValue ? cookieValue.split("=")[1] : null;
+  const cookieValue = document.cookie.split("; ").find((row) => row.startsWith(name + "="))
+  return cookieValue ? cookieValue.split("=")[1] : null
 }
 
 function formatUserRole(userRole: string) {
   if (!userRole) {
-    return "Undefined Role";
+    return "Undefined Role"
   }
   switch (userRole.toLowerCase()) {
     case "app_owner":
-      return "App Owner";
+      return "App Owner"
     case "demo_app_owner":
-      return "App Owner";
+      return "App Owner"
     case "app_admin":
-      return "Admin";
+      return "Admin"
     case "proxy_admin":
-      return "Admin";
+      return "Admin"
     case "proxy_admin_viewer":
-      return "Admin Viewer";
+      return "Admin Viewer"
     case "org_admin":
-      return "Org Admin";
+      return "Org Admin"
     case "internal_user":
-      return "Internal User";
+      return "Internal User"
     case "internal_user_viewer":
     case "internal_viewer": // TODO:remove if deprecated
-      return "Internal Viewer";
+      return "Internal Viewer"
     case "app_user":
-      return "App User";
+      return "App User"
     default:
-      return "Unknown Role";
+      return "Unknown Role"
   }
 }
 
 interface ProxySettings {
-  PROXY_BASE_URL: string;
-  PROXY_LOGOUT_URL: string;
+  PROXY_BASE_URL: string
+  PROXY_LOGOUT_URL: string
 }
 
-const queryClient = new QueryClient();
+const queryClient = new QueryClient()
 
 function LoadingScreen() {
   return (
     <div className={cx("h-screen", "flex items-center justify-center gap-4")}>
-      <div className="text-lg font-medium py-2 pr-4 border-r border-r-gray-200">
-        🚅 LiteLLM
-      </div>
+      <div className="text-lg font-medium py-2 pr-4 border-r border-r-gray-200">🚅 LiteLLM</div>
 
       <div className="flex items-center justify-center gap-2">
         <UiLoadingSpinner className="size-4" />
         <span className="text-gray-600 text-sm">Loading...</span>
       </div>
     </div>
-  );
+  )
 }
 
 export default function CreateKeyPage() {
-  const [userRole, setUserRole] = useState("");
-  const [premiumUser, setPremiumUser] = useState(false);
-  const [disabledPersonalKeyCreation, setDisabledPersonalKeyCreation] =
-    useState(false);
-  const [userEmail, setUserEmail] = useState<null | string>(null);
-  const [teams, setTeams] = useState<Team[] | null>(null);
-  const [keys, setKeys] = useState<null | any[]>([]);
-  const [organizations, setOrganizations] = useState<Organization[]>([]);
-  const [userModels, setUserModels] = useState<string[]>([]);
+  const [userRole, setUserRole] = useState("")
+  const [premiumUser, setPremiumUser] = useState(false)
+  const [disabledPersonalKeyCreation, setDisabledPersonalKeyCreation] = useState(false)
+  const [userEmail, setUserEmail] = useState<null | string>(null)
+  const [teams, setTeams] = useState<Team[] | null>(null)
+  const [keys, setKeys] = useState<null | any[]>([])
+  const [organizations, setOrganizations] = useState<Organization[]>([])
+  const [userModels, setUserModels] = useState<string[]>([])
   const [proxySettings, setProxySettings] = useState<ProxySettings>({
     PROXY_BASE_URL: "",
     PROXY_LOGOUT_URL: "",
-  });
+  })
 
-  const [showSSOBanner, setShowSSOBanner] = useState<boolean>(true);
-  const searchParams = useSearchParams()!;
-  const [modelData, setModelData] = useState<any>({ data: [] });
-  const [token, setToken] = useState<string | null>(null);
-  const [createClicked, setCreateClicked] = useState<boolean>(false);
-  const [authLoading, setAuthLoading] = useState(true);
-  const [userID, setUserID] = useState<string | null>(null);
+  const [showSSOBanner, setShowSSOBanner] = useState<boolean>(true)
+  const searchParams = useSearchParams()!
+  const [modelData, setModelData] = useState<any>({ data: [] })
+  const [token, setToken] = useState<string | null>(null)
+  const [createClicked, setCreateClicked] = useState<boolean>(false)
+  const [authLoading, setAuthLoading] = useState(true)
+  const [userID, setUserID] = useState<string | null>(null)
 
-  const invitation_id = searchParams.get("invitation_id");
+  const invitation_id = searchParams.get("invitation_id")
 
   // Get page from URL, default to 'api-keys' if not present
   const [page, setPage] = useState(() => {
-    return searchParams.get("page") || "api-keys";
-  });
+    return searchParams.get("page") || "api-keys"
+  })
 
   // Custom setPage function that updates URL
   const updatePage = (newPage: string) => {
     // Update URL without full page reload
-    const newSearchParams = new URLSearchParams(searchParams);
-    newSearchParams.set("page", newPage);
+    const newSearchParams = new URLSearchParams(searchParams)
+    newSearchParams.set("page", newPage)
 
     // Use Next.js router to update URL
-    window.history.pushState(null, "", `?${newSearchParams.toString()}`);
+    window.history.pushState(null, "", `?${newSearchParams.toString()}`)
 
-    setPage(newPage);
-  };
+    setPage(newPage)
+  }
 
-  const [accessToken, setAccessToken] = useState<string | null>(null);
+  const [accessToken, setAccessToken] = useState<string | null>(null)
 
   const addKey = (data: any) => {
-    setKeys((prevData) => (prevData ? [...prevData, data] : [data]));
-    setCreateClicked(() => !createClicked);
-  };
-  const redirectToLogin =
-    authLoading === false && token === null && invitation_id === null;
+    setKeys((prevData) => (prevData ? [...prevData, data] : [data]))
+    setCreateClicked(() => !createClicked)
+  }
+  const redirectToLogin = authLoading === false && token === null && invitation_id === null
 
   useEffect(() => {
-    const token = getCookie("token");
+    const token = getCookie("token")
     getUiConfig().then((data) => {
       // get the information for constructing the proxy base url, and then set the token and auth loading
-      setToken(token);
-      setAuthLoading(false);
-    });
-  }, []);
+      setToken(token)
+      setAuthLoading(false)
+    })
+  }, [])
 
   useEffect(() => {
     if (redirectToLogin) {
-      window.location.href = (proxyBaseUrl || "") + "/sso/key/generate";
+      window.location.href = (proxyBaseUrl || "") + "/sso/key/generate"
     }
-  }, [redirectToLogin]);
+  }, [redirectToLogin])
 
   useEffect(() => {
     if (!token) {
-      return;
+      return
     }
 
-    const decoded = jwtDecode(token) as { [key: string]: any };
+    const decoded = jwtDecode(token) as { [key: string]: any }
     if (decoded) {
       // set accessToken
-      setAccessToken(decoded.key);
+      setAccessToken(decoded.key)
 
-      setDisabledPersonalKeyCreation(
-        decoded.disabled_non_admin_personal_key_creation
-      );
+      setDisabledPersonalKeyCreation(decoded.disabled_non_admin_personal_key_creation)
 
       // check if userRole is defined
       if (decoded.user_role) {
-        const formattedUserRole = formatUserRole(decoded.user_role);
-        setUserRole(formattedUserRole);
+        const formattedUserRole = formatUserRole(decoded.user_role)
+        setUserRole(formattedUserRole)
         if (formattedUserRole == "Admin Viewer") {
-          setPage("usage");
+          setPage("usage")
         }
       }
 
       if (decoded.user_email) {
-        setUserEmail(decoded.user_email);
+        setUserEmail(decoded.user_email)
       }
 
       if (decoded.login_method) {
-        setShowSSOBanner(
-          decoded.login_method == "username_password" ? true : false
-        );
+        setShowSSOBanner(decoded.login_method == "username_password" ? true : false)
       }
 
       if (decoded.premium_user) {
-        setPremiumUser(decoded.premium_user);
+        setPremiumUser(decoded.premium_user)
       }
 
       if (decoded.auth_header_name) {
-        setGlobalLitellmHeaderName(decoded.auth_header_name);
+        setGlobalLitellmHeaderName(decoded.auth_header_name)
       }
 
       if (decoded.user_id) {
-        setUserID(decoded.user_id);
+        setUserID(decoded.user_id)
       }
     }
-  }, [token]);
+  }, [token])
 
   useEffect(() => {
     if (accessToken && userID && userRole) {
-      fetchUserModels(userID, userRole, accessToken, setUserModels);
+      fetchUserModels(userID, userRole, accessToken, setUserModels)
     }
     if (accessToken && userID && userRole) {
-      fetchTeams(accessToken, userID, userRole, null, setTeams);
+      fetchTeams(accessToken, userID, userRole, null, setTeams)
     }
     if (accessToken) {
-      fetchOrganizations(accessToken, setOrganizations);
+      fetchOrganizations(accessToken, setOrganizations)
     }
-  }, [accessToken, userID, userRole]);
+  }, [accessToken, userID, userRole])
 
   if (authLoading || redirectToLogin) {
-    return <LoadingScreen />;
+    return <LoadingScreen />
   }
 
   return (
@@ -237,226 +223,199 @@ export default function CreateKeyPage() {
         <ThemeProvider accessToken={accessToken}>
           {invitation_id ? (
             <UserDashboard
-            userID={userID}
-            userRole={userRole}
-            premiumUser={premiumUser}
-            teams={teams}
-            keys={keys}
-            setUserRole={setUserRole}
-            userEmail={userEmail}
-            setUserEmail={setUserEmail}
-            setTeams={setTeams}
-            setKeys={setKeys}
-            organizations={organizations}
-            addKey={addKey}
-            createClicked={createClicked}
-          />
-        ) : (
-          <div className="flex flex-col min-h-screen">
-            <Navbar
               userID={userID}
               userRole={userRole}
               premiumUser={premiumUser}
+              teams={teams}
+              keys={keys}
+              setUserRole={setUserRole}
               userEmail={userEmail}
-              setProxySettings={setProxySettings}
-              proxySettings={proxySettings}
-              accessToken={accessToken}
-              isPublicPage={false}
+              setUserEmail={setUserEmail}
+              setTeams={setTeams}
+              setKeys={setKeys}
+              organizations={organizations}
+              addKey={addKey}
+              createClicked={createClicked}
             />
-            <div className="flex flex-1 overflow-auto">
-              <div className="mt-8">
-                <Sidebar
-                  accessToken={accessToken}
-                  setPage={updatePage}
-                  userRole={userRole}
-                  defaultSelectedKey={page}
-                />
-              </div>
+          ) : (
+            <div className="flex flex-col min-h-screen">
+              <Navbar
+                userID={userID}
+                userRole={userRole}
+                premiumUser={premiumUser}
+                userEmail={userEmail}
+                setProxySettings={setProxySettings}
+                proxySettings={proxySettings}
+                accessToken={accessToken}
+                isPublicPage={false}
+              />
+              <div className="flex flex-1 overflow-auto">
+                <div className="mt-8">
+                  <Sidebar
+                    accessToken={accessToken}
+                    setPage={updatePage}
+                    userRole={userRole}
+                    defaultSelectedKey={page}
+                  />
+                </div>
 
-              {page == "api-keys" ? (
-                <UserDashboard
-                  userID={userID}
-                  userRole={userRole}
-                  premiumUser={premiumUser}
-                  teams={teams}
-                  keys={keys}
-                  setUserRole={setUserRole}
-                  userEmail={userEmail}
-                  setUserEmail={setUserEmail}
-                  setTeams={setTeams}
-                  setKeys={setKeys}
-                  organizations={organizations}
-                  addKey={addKey}
-                  createClicked={createClicked}
-                />
-              ) : page == "models" ? (
-                <ModelDashboard
-                  userID={userID}
-                  userRole={userRole}
-                  token={token}
-                  keys={keys}
-                  accessToken={accessToken}
-                  modelData={modelData}
-                  setModelData={setModelData}
-                  premiumUser={premiumUser}
-                  teams={teams}
-                />
-              ) : page == "llm-playground" ? (
-                <ChatUI
-                  userID={userID}
-                  userRole={userRole}
-                  token={token}
-                  accessToken={accessToken}
-                  disabledPersonalKeyCreation={disabledPersonalKeyCreation}
-                />
-              ) : page == "users" ? (
-                <ViewUserDashboard
-                  userID={userID}
-                  userRole={userRole}
-                  token={token}
-                  keys={keys}
-                  teams={teams}
-                  accessToken={accessToken}
-                  setKeys={setKeys}
-                />
-              ) : page == "teams" ? (
-                <Teams
-                  teams={teams}
-                  setTeams={setTeams}
-                  searchParams={searchParams}
-                  accessToken={accessToken}
-                  userID={userID}
-                  userRole={userRole}
-                  organizations={organizations}
-                  premiumUser={premiumUser}
-                />
-              ) : page == "organizations" ? (
-                <Organizations
-                  organizations={organizations}
-                  setOrganizations={setOrganizations}
-                  userModels={userModels}
-                  accessToken={accessToken}
-                  userRole={userRole}
-                  premiumUser={premiumUser}
-                />
-              ) : page == "admin-panel" ? (
-                <AdminPanel
-                  setTeams={setTeams}
-                  searchParams={searchParams}
-                  accessToken={accessToken}
-                  userID={userID}
-                  showSSOBanner={showSSOBanner}
-                  premiumUser={premiumUser}
-                  proxySettings={proxySettings}
-                />
-              ) : page == "api_ref" ? (
-                <APIRef proxySettings={proxySettings} />
-              ) : page == "settings" ? (
-                <Settings
-                  userID={userID}
-                  userRole={userRole}
-                  accessToken={accessToken}
-                  premiumUser={premiumUser}
-                />
-              ) : page == "budgets" ? (
-                <BudgetPanel accessToken={accessToken} />
-              ) : page == "guardrails" ? (
-                <GuardrailsPanel
-                  accessToken={accessToken}
-                  userRole={userRole}
-                />
-              ) : page == "prompts" ? (
-                <PromptsPanel
-                  accessToken={accessToken}
-                  userRole={userRole}
-                />
-              ) : page == "transform-request" ? (
-                <TransformRequestPanel accessToken={accessToken} />
-              ) : page == "general-settings" ? (
-                <GeneralSettings
-                  userID={userID}
-                  userRole={userRole}
-                  accessToken={accessToken}
-                  modelData={modelData}
-                />
-              ) : page == "ui-theme" ? (
-                <UIThemeSettings
-                  userID={userID}
-                  userRole={userRole}
-                  accessToken={accessToken}
-                />
-              ) : page == "model-hub-table" ? (
-                <ModelHubTable
-                  accessToken={accessToken}
-                  publicPage={false}
-                  premiumUser={premiumUser}
-                  userRole={userRole}
-                />
-              ) : page == "caching" ? (
-                <CacheDashboard
-                  userID={userID}
-                  userRole={userRole}
-                  token={token}
-                  accessToken={accessToken}
-                  premiumUser={premiumUser}
-                />
-              ) : page == "pass-through-settings" ? (
-                <PassThroughSettings
-                  userID={userID}
-                  userRole={userRole}
-                  accessToken={accessToken}
-                  modelData={modelData}
-                />
-              ) : page == "logs" ? (
-                <SpendLogsTable
-                  userID={userID}
-                  userRole={userRole}
-                  token={token}
-                  accessToken={accessToken}
-                  allTeams={(teams as Team[]) ?? []}
-                  premiumUser={premiumUser}
-                />
-              ) : page == "mcp-servers" ? (
-                <MCPServers
-                  accessToken={accessToken}
-                  userRole={userRole}
-                  userID={userID}
-                />
-              ) : page == "tag-management" ? (
-                <TagManagement
-                  accessToken={accessToken}
-                  userRole={userRole}
-                  userID={userID}
-                />
-              ) : page == "vector-stores" ? (
-                <VectorStoreManagement
-                  accessToken={accessToken}
-                  userRole={userRole}
-                  userID={userID}
-                />
-              ) : page == "new_usage" ? (
-                <NewUsagePage
-                  userID={userID}
-                  userRole={userRole}
-                  accessToken={accessToken}
-                  teams={(teams as Team[]) ?? []}
-                  premiumUser={premiumUser}
-                />
-              ) : (
-                <Usage
-                  userID={userID}
-                  userRole={userRole}
-                  token={token}
-                  accessToken={accessToken}
-                  keys={keys}
-                  premiumUser={premiumUser}
-                />
-              )}
+                {page == "api-keys" ? (
+                  <UserDashboard
+                    userID={userID}
+                    userRole={userRole}
+                    premiumUser={premiumUser}
+                    teams={teams}
+                    keys={keys}
+                    setUserRole={setUserRole}
+                    userEmail={userEmail}
+                    setUserEmail={setUserEmail}
+                    setTeams={setTeams}
+                    setKeys={setKeys}
+                    organizations={organizations}
+                    addKey={addKey}
+                    createClicked={createClicked}
+                  />
+                ) : page == "models" ? (
+                  <ModelDashboard
+                    userID={userID}
+                    userRole={userRole}
+                    token={token}
+                    keys={keys}
+                    accessToken={accessToken}
+                    modelData={modelData}
+                    setModelData={setModelData}
+                    premiumUser={premiumUser}
+                    teams={teams}
+                  />
+                ) : page == "llm-playground" ? (
+                  <ChatUI
+                    userID={userID}
+                    userRole={userRole}
+                    token={token}
+                    accessToken={accessToken}
+                    disabledPersonalKeyCreation={disabledPersonalKeyCreation}
+                  />
+                ) : page == "users" ? (
+                  <ViewUserDashboard
+                    userID={userID}
+                    userRole={userRole}
+                    token={token}
+                    keys={keys}
+                    teams={teams}
+                    accessToken={accessToken}
+                    setKeys={setKeys}
+                  />
+                ) : page == "teams" ? (
+                  <Teams
+                    teams={teams}
+                    setTeams={setTeams}
+                    searchParams={searchParams}
+                    accessToken={accessToken}
+                    userID={userID}
+                    userRole={userRole}
+                    organizations={organizations}
+                    premiumUser={premiumUser}
+                  />
+                ) : page == "organizations" ? (
+                  <Organizations
+                    organizations={organizations}
+                    setOrganizations={setOrganizations}
+                    userModels={userModels}
+                    accessToken={accessToken}
+                    userRole={userRole}
+                    premiumUser={premiumUser}
+                  />
+                ) : page == "admin-panel" ? (
+                  <AdminPanel
+                    setTeams={setTeams}
+                    searchParams={searchParams}
+                    accessToken={accessToken}
+                    userID={userID}
+                    showSSOBanner={showSSOBanner}
+                    premiumUser={premiumUser}
+                    proxySettings={proxySettings}
+                  />
+                ) : page == "api_ref" ? (
+                  <APIRef proxySettings={proxySettings} />
+                ) : page == "settings" ? (
+                  <Settings userID={userID} userRole={userRole} accessToken={accessToken} premiumUser={premiumUser} />
+                ) : page == "budgets" ? (
+                  <BudgetPanel accessToken={accessToken} />
+                ) : page == "guardrails" ? (
+                  <GuardrailsPanel accessToken={accessToken} userRole={userRole} />
+                ) : page == "prompts" ? (
+                  <PromptsPanel accessToken={accessToken} userRole={userRole} />
+                ) : page == "transform-request" ? (
+                  <TransformRequestPanel accessToken={accessToken} />
+                ) : page == "general-settings" ? (
+                  <GeneralSettings
+                    userID={userID}
+                    userRole={userRole}
+                    accessToken={accessToken}
+                    modelData={modelData}
+                  />
+                ) : page == "ui-theme" ? (
+                  <UIThemeSettings userID={userID} userRole={userRole} accessToken={accessToken} />
+                ) : page == "model-hub-table" ? (
+                  <ModelHubTable
+                    accessToken={accessToken}
+                    publicPage={false}
+                    premiumUser={premiumUser}
+                    userRole={userRole}
+                  />
+                ) : page == "caching" ? (
+                  <CacheDashboard
+                    userID={userID}
+                    userRole={userRole}
+                    token={token}
+                    accessToken={accessToken}
+                    premiumUser={premiumUser}
+                  />
+                ) : page == "pass-through-settings" ? (
+                  <PassThroughSettings
+                    userID={userID}
+                    userRole={userRole}
+                    accessToken={accessToken}
+                    modelData={modelData}
+                  />
+                ) : page == "logs" ? (
+                  <SpendLogsTable
+                    userID={userID}
+                    userRole={userRole}
+                    token={token}
+                    accessToken={accessToken}
+                    allTeams={(teams as Team[]) ?? []}
+                    premiumUser={premiumUser}
+                  />
+                ) : page == "mcp-servers" ? (
+                  <MCPServers accessToken={accessToken} userRole={userRole} userID={userID} />
+                ) : page == "tag-management" ? (
+                  <TagManagement accessToken={accessToken} userRole={userRole} userID={userID} />
+                ) : page == "vector-stores" ? (
+                  <VectorStoreManagement accessToken={accessToken} userRole={userRole} userID={userID} />
+                ) : page == "new_usage" ? (
+                  <NewUsagePage
+                    userID={userID}
+                    userRole={userRole}
+                    accessToken={accessToken}
+                    teams={(teams as Team[]) ?? []}
+                    premiumUser={premiumUser}
+                  />
+                ) : (
+                  <Usage
+                    userID={userID}
+                    userRole={userRole}
+                    token={token}
+                    accessToken={accessToken}
+                    keys={keys}
+                    premiumUser={premiumUser}
+                  />
+                )}
+              </div>
             </div>
-          </div>
-        )}
+          )}
         </ThemeProvider>
       </QueryClientProvider>
     </Suspense>
-  );
+  )
 }

ui/litellm-dashboard/src/components/activity_metrics.tsx

@@ -1,143 +1,16 @@
 import React from 'react';
-import { Card, Grid, Text, Title, Accordion, AccordionHeader, AccordionBody } from '@tremor/react';
+import { Card, Grid, Text, Title } from '@tremor/react';
 import { AreaChart, BarChart } from '@tremor/react';
-import { SpendMetrics, DailyData, ModelActivityData, MetricWithMetadata, KeyMetricWithMetadata, TopApiKeyData } from './usage/types';
+import { DailyData, ModelActivityData, KeyMetricWithMetadata, TopApiKeyData } from './usage/types';
 import { Collapse } from 'antd';
 import { formatNumberWithCommas } from '@/utils/dataUtils';
-import type { CustomTooltipProps } from "@tremor/react";
-import {
-  valueFormatter,
-  valueFormatterSpend,
-} from "../components/usage/utils/value_formatters";
+import { valueFormatter } from '../components/usage/utils/value_formatters';
+import { CustomTooltip, CustomLegend } from './common_components/chartUtils';
 
 interface ActivityMetricsProps {
   modelMetrics: Record<string, ModelActivityData>;
 }
 
-interface ChartDataPoint {
-  date: string;
-  metrics: SpendMetrics;
-}
-
-const colorNameToHex: { [key: string]: string } = {
-  blue: "#3b82f6",
-  cyan: "#06b6d4",
-  indigo: "#6366f1",
-  green: "#22c55e",
-  red: "#ef4444",
-  purple: "#8b5cf6",
-};
-
-export const CustomTooltip = ({
-  active,
-  payload,
-  label,
-}: CustomTooltipProps) => {
-  if (active && payload && payload.length) {
-    const formatCategoryName = (name: string): string => {
-      return name
-        .replace("metrics.", "")
-        .replace(/_/g, " ")
-        .split(" ")
-        .map((word) => word.charAt(0).toUpperCase() + word.slice(1))
-        .join(" ");
-    };
-
-    const getRawValue = (
-      dataPoint: ChartDataPoint,
-      key: string
-    ): number | undefined => {
-      // key is like "metrics.total_tokens"
-      const metricKey = key.substring(
-        key.indexOf(".") + 1
-      ) as keyof SpendMetrics;
-      if (dataPoint.metrics && metricKey in dataPoint.metrics) {
-        return dataPoint.metrics[metricKey];
-      }
-      return undefined;
-    };
-
-    return (
-      <div className="w-56 rounded-tremor-default border border-tremor-border bg-tremor-background p-2 text-tremor-default shadow-tremor-dropdown">
-        <p className="text-tremor-content-strong">{label}</p>
-        {payload.map((item) => {
-          const dataKey = item.dataKey?.toString();
-          if (!dataKey || !item.payload) return null;
-
-          const rawValue = getRawValue(item.payload, dataKey);
-          const isSpend = dataKey.includes("spend");
-          const formattedValue =
-            rawValue !== undefined
-              ? isSpend
-                ? `$${rawValue.toLocaleString(undefined, { minimumFractionDigits: 2, maximumFractionDigits: 2 })}`
-                : rawValue.toLocaleString()
-              : "N/A";
-
-          const colorName = item.color as keyof typeof colorNameToHex;
-          const hexColor = colorNameToHex[colorName] || item.color;
-          return (
-            <div
-              key={dataKey}
-              className="flex items-center justify-between space-x-4"
-            >
-              <div className="flex items-center space-x-2">
-                <span
-                  className={`h-2 w-2 shrink-0 rounded-full ring-2 ring-white drop-shadow-md`}
-                  style={{ backgroundColor: hexColor }}
-                />
-                <p className="font-medium text-tremor-content dark:text-dark-tremor-content">
-                  {formatCategoryName(dataKey)}
-                </p>
-              </div>
-              <p className="font-medium text-tremor-content-emphasis dark:text-dark-tremor-content-emphasis">
-                {formattedValue}
-              </p>
-            </div>
-          );
-        })}
-      </div>
-    );
-  }
-  return null;
-};
-
-const CustomLegend = ({
-  categories,
-  colors,
-}: {
-  categories: string[];
-  colors: string[];
-}) => {
-  const formatCategoryName = (name: string): string => {
-    return name
-      .replace("metrics.", "")
-      .replace(/_/g, " ")
-      .split(" ")
-      .map((word) => word.charAt(0).toUpperCase() + word.slice(1))
-      .join(" ");
-  };
-
-  return (
-    <div className="flex items-center justify-end space-x-4">
-      {categories.map((category, idx) => {
-        const colorName = colors[idx] as keyof typeof colorNameToHex;
-        const hexColor = colorNameToHex[colorName] || colors[idx];
-        return (
-          <div key={category} className="flex items-center space-x-2">
-            <span
-              className={`h-2 w-2 shrink-0 rounded-full ring-4 ring-white`}
-              style={{ backgroundColor: hexColor }}
-            />
-            <p className="text-sm text-tremor-content dark:text-dark-tremor-content">
-              {formatCategoryName(category)}
-            </p>
-          </div>
-        );
-      })}
-    </div>
-  );
-};
-
 const ModelSection = ({
   modelName,
   metrics,
@@ -215,14 +88,7 @@ const ModelSection = ({
         <Card>
           <div className="flex justify-between items-center">
             <Title>Total Tokens</Title>
-            <CustomLegend
-              categories={[
-                "metrics.prompt_tokens",
-                "metrics.completion_tokens",
-                "metrics.total_tokens",
-              ]}
-              colors={["blue", "cyan", "indigo"]}
-            />
+            <CustomLegend categories={["metrics.prompt_tokens", "metrics.completion_tokens", "metrics.total_tokens"]} colors={["blue", "cyan", "indigo"]} />
           </div>
           <AreaChart
             className="mt-4"
@@ -243,10 +109,7 @@ const ModelSection = ({
         <Card>
           <div className="flex justify-between items-center">
             <Title>Requests per day</Title>
-            <CustomLegend
-              categories={["metrics.api_requests"]}
-              colors={["blue"]}
-            />
+            <CustomLegend categories={["metrics.api_requests"]} colors={["blue"]} />
           </div>
           <BarChart
             className="mt-4"
@@ -271,9 +134,7 @@ const ModelSection = ({
             index="date"
             categories={["metrics.spend"]}
             colors={["green"]}
-            valueFormatter={valueFormatterSpend}
-            customTooltip={CustomTooltip}
-            showLegend={false}
+            valueFormatter={(value: number) => `$${formatNumberWithCommas(value, 2)}`}
           />
         </Card>
 
@@ -487,26 +348,14 @@ export const ActivityMetrics: React.FC<ActivityMetricsProps> = ({
             />
           </Card>
           <Card>
-            <div className="flex justify-between items-center">
-              <Title>Total Requests Over Time</Title>
-              <CustomLegend
-                categories={[
-                  "metrics.successful_requests",
-                  "metrics.failed_requests",
-                ]}
-                colors={["green", "red"]}
-              />
-            </div>
+            <Title>Total Requests Over Time</Title>
             <AreaChart
               className="mt-4"
               data={sortedDailyData}
               index="date"
-              categories={[
-                "metrics.successful_requests",
-                "metrics.failed_requests",
-              ]}
-              colors={["green", "red"]}
-              valueFormatter={valueFormatter}
+              categories={["metrics.successful_requests", "metrics.failed_requests"]}
+              colors={["emerald", "red"]}
+              valueFormatter={(number: number) => number.toLocaleString()}
               stack
               customTooltip={CustomTooltip}
               showLegend={false}
@@ -585,7 +434,6 @@ export const processActivityData = (dailyActivity: { results: DailyData[] }, key
           daily_data: []
         };
       }
-
       // Update totals
       modelMetrics[model].total_requests += modelData.metrics.api_requests;
       modelMetrics[model].prompt_tokens += modelData.metrics.prompt_tokens;

ui/litellm-dashboard/src/components/all_keys_table.tsx

@@ -1,68 +1,54 @@
-"use client";
-import React, { useEffect, useState, useCallback, useRef } from "react";
-import { ColumnDef, Row } from "@tanstack/react-table";
-import { DataTable } from "./view_logs/table";
+"use client"
+import React, { useEffect, useState, useCallback, useRef } from "react"
+import { ColumnDef, Row } from "@tanstack/react-table"
+import { DataTable } from "./view_logs/table"
 import { Select, SelectItem } from "@tremor/react"
 import { Button } from "@tremor/react"
-import KeyInfoView from "./key_info_view";
-import { Tooltip } from "antd";
-import { Team, KeyResponse } from "./key_team_helpers/key_list";
-import FilterComponent from "./common_components/filter";
-import { FilterOption } from "./common_components/filter";
-import { keyListCall, Organization, userListCall } from "./networking";
-import { createTeamSearchFunction } from "./key_team_helpers/team_search_fn";
-import { createOrgSearchFunction } from "./key_team_helpers/organization_search_fn";
-import { useFilterLogic } from "./key_team_helpers/filter_logic";
-import { Setter } from "@/types";
-import { updateExistingKeys } from "@/utils/dataUtils";
-import { debounce } from "lodash";
-import { defaultPageSize } from "./constants";
-import { fetchAllTeams } from "./key_team_helpers/filter_helpers";
-import { fetchAllOrganizations } from "./key_team_helpers/filter_helpers";
-import {
-  flexRender,
-  getCoreRowModel,
-  getSortedRowModel,
-  SortingState,
-  useReactTable,
-} from "@tanstack/react-table";
-import {
-  Table,
-  TableHead,
-  TableHeaderCell,
-  TableBody,
-  TableRow,
-  TableCell,
-  Icon,
-} from "@tremor/react";
-import { SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline";
-import { Badge, Text } from "@tremor/react";
-import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key";
-import { formatNumberWithCommas } from "@/utils/dataUtils";
+import KeyInfoView from "./templates/key_info_view"
+import { Tooltip } from "antd"
+import { Team, KeyResponse } from "./key_team_helpers/key_list"
+import FilterComponent from "./common_components/filter"
+import { FilterOption } from "./common_components/filter"
+import { keyListCall, Organization, userListCall } from "./networking"
+import { createTeamSearchFunction } from "./key_team_helpers/team_search_fn"
+import { createOrgSearchFunction } from "./key_team_helpers/organization_search_fn"
+import { useFilterLogic } from "./key_team_helpers/filter_logic"
+import { Setter } from "@/types"
+import { updateExistingKeys } from "@/utils/dataUtils"
+import { debounce } from "lodash"
+import { defaultPageSize } from "./constants"
+import { fetchAllTeams } from "./key_team_helpers/filter_helpers"
+import { fetchAllOrganizations } from "./key_team_helpers/filter_helpers"
+import { flexRender, getCoreRowModel, getSortedRowModel, SortingState, useReactTable } from "@tanstack/react-table"
+import { Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell, Icon } from "@tremor/react"
+import { SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline"
+import { Badge, Text } from "@tremor/react"
+import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key"
+import { formatNumberWithCommas } from "@/utils/dataUtils"
 
 interface AllKeysTableProps {
   keys: KeyResponse[]
   setKeys: (keys: KeyResponse[] | ((prev: KeyResponse[]) => KeyResponse[])) => void
   isLoading?: boolean
   pagination: {
-    currentPage: number;
-    totalPages: number;
-    totalCount: number;
-  };
-  onPageChange: (page: number) => void;
-  pageSize?: number;
-  teams: Team[] | null;
-  selectedTeam: Team | null;
-  setSelectedTeam: (team: Team | null) => void;
-  selectedKeyAlias: string | null;
-  setSelectedKeyAlias: Setter<string | null>;
-  accessToken: string | null;
-  userID: string | null;
-  userRole: string | null;
-  organizations: Organization[] | null;
-  setCurrentOrg: React.Dispatch<React.SetStateAction<Organization | null>>;
-  refresh?: () => void;
-  onSortChange?: (sortBy: string, sortOrder: 'asc' | 'desc') => void;
+    currentPage: number
+    totalPages: number
+    totalCount: number
+  }
+  onPageChange: (page: number) => void
+  pageSize?: number
+  teams: Team[] | null
+  selectedTeam: Team | null
+  setSelectedTeam: (team: Team | null) => void
+  selectedKeyAlias: string | null
+  setSelectedKeyAlias: Setter<string | null>
+  accessToken: string | null
+  userID: string | null
+  userRole: string | null
+  organizations: Organization[] | null
+  setCurrentOrg: React.Dispatch<React.SetStateAction<Organization | null>>
+  refresh?: () => void
+  onSortChange?: (sortBy: string, sortOrder: "asc" | "desc") => void
   currentSort?: {
     sortBy: string
     sortOrder: "asc" | "desc"
@@ -74,57 +60,55 @@ interface AllKeysTableProps {
 // Define columns similar to our logs table
 
 interface UserResponse {
-  user_id: string;
-  user_email: string;
-  user_role: string;
+  user_id: string
+  user_email: string
+  user_role: string
 }
 
-const TeamFilter = ({ 
-  teams, 
-  selectedTeam, 
-  setSelectedTeam 
-}: { 
-  teams: Team[] | null;
-  selectedTeam: Team | null;
-  setSelectedTeam: (team: Team | null) => void;
+const TeamFilter = ({
+  teams,
+  selectedTeam,
+  setSelectedTeam,
+}: {
+  teams: Team[] | null
+  selectedTeam: Team | null
+  setSelectedTeam: (team: Team | null) => void
 }) => {
-    const handleTeamChange = (value: string) => {
-      const team = teams?.find(t => t.team_id === value);
-      setSelectedTeam(team || null);
-    };
-  
-    return (
-      <div className="mb-4">
-        <div className="flex items-center gap-2">
-          <span className="text-sm text-gray-600">Where Team is</span>
-          <Select
-            value={selectedTeam?.team_id || ""}
-            onValueChange={handleTeamChange}
-            placeholder="Team ID"
-            className="w-[400px]"
-          >
-            <SelectItem value="team_id">Team ID</SelectItem>
-            {teams?.map((team) => (
-              <SelectItem key={team.team_id} value={team.team_id}>
-                <span className="font-medium">{team.team_alias}</span>{" "}
-                <span className="text-gray-500">({team.team_id})</span>
-              </SelectItem>
-            ))}
-          </Select>
-        </div>
+  const handleTeamChange = (value: string) => {
+    const team = teams?.find((t) => t.team_id === value)
+    setSelectedTeam(team || null)
+  }
+
+  return (
+    <div className="mb-4">
+      <div className="flex items-center gap-2">
+        <span className="text-sm text-gray-600">Where Team is</span>
+        <Select
+          value={selectedTeam?.team_id || ""}
+          onValueChange={handleTeamChange}
+          placeholder="Team ID"
+          className="w-[400px]"
+        >
+          <SelectItem value="team_id">Team ID</SelectItem>
+          {teams?.map((team) => (
+            <SelectItem key={team.team_id} value={team.team_id}>
+              <span className="font-medium">{team.team_alias}</span>{" "}
+              <span className="text-gray-500">({team.team_id})</span>
+            </SelectItem>
+          ))}
+        </Select>
       </div>
-    );
-  };
-  
+    </div>
+  )
+}
 
 /**
  * AllKeysTable – a new table for keys that mimics the table styling used in view_logs.
  * The team selector and filtering have been removed so that all keys are shown.
  */
 
-
-export function AllKeysTable({ 
-  keys, 
+export function AllKeysTable({
+  keys,
   setKeys,
   isLoading = false,
   pagination,
@@ -146,67 +130,62 @@ export function AllKeysTable({
   premiumUser,
   setAccessToken,
 }: AllKeysTableProps) {
-  const [selectedKeyId, setSelectedKeyId] = useState<string | null>(null);
-  const [userList, setUserList] = useState<UserResponse[]>([]);
+  const [selectedKeyId, setSelectedKeyId] = useState<string | null>(null)
+  const [userList, setUserList] = useState<UserResponse[]>([])
   const [sorting, setSorting] = React.useState<SortingState>(() => {
     if (currentSort) {
-      return [{
-        id: currentSort.sortBy,
-        desc: currentSort.sortOrder === 'desc'
-      }];
+      return [
+        {
+          id: currentSort.sortBy,
+          desc: currentSort.sortOrder === "desc",
+        },
+      ]
     }
-    return [{
-      id: "created_at",
-      desc: true
-    }];
-  });
-  const [expandedAccordions, setExpandedAccordions] = useState<Record<string, boolean>>({});
+    return [
+      {
+        id: "created_at",
+        desc: true,
+      },
+    ]
+  })
+  const [expandedAccordions, setExpandedAccordions] = useState<Record<string, boolean>>({})
 
   // Use the filter logic hook
 
-  const {
-    filters,
-    filteredKeys,
-    allKeyAliases,
-    allTeams,
-    allOrganizations,
-    handleFilterChange,
-    handleFilterReset
-  } = useFilterLogic({
-    keys,
-    teams,
-    organizations,
-    accessToken,
-  });
-
-  
+  const { filters, filteredKeys, allKeyAliases, allTeams, allOrganizations, handleFilterChange, handleFilterReset } =
+    useFilterLogic({
+      keys,
+      teams,
+      organizations,
+      accessToken,
+    })
 
   useEffect(() => {
     if (accessToken) {
-      const user_IDs = keys.map(key => key.user_id).filter(id => id !== null);
+      const user_IDs = keys.map((key) => key.user_id).filter((id) => id !== null)
       const fetchUserList = async () => {
-        const userListData = await userListCall(accessToken, user_IDs, 1, 100);
-        setUserList(userListData.users);
-      };
-      fetchUserList();
+        const userListData = await userListCall(accessToken, user_IDs, 1, 100)
+        setUserList(userListData.users)
+      }
+      fetchUserList()
     }
-  }, [accessToken, keys]);
+  }, [accessToken, keys])
 
   // Add a useEffect to call refresh when a key is created
   useEffect(() => {
     if (refresh) {
       const handleStorageChange = () => {
-        refresh();
-      };
-      
+        refresh()
+      }
+
       // Listen for storage events that might indicate a key was created
-      window.addEventListener('storage', handleStorageChange);
-      
+      window.addEventListener("storage", handleStorageChange)
+
       return () => {
-        window.removeEventListener('storage', handleStorageChange);
-      };
+        window.removeEventListener("storage", handleStorageChange)
+      }
     }
-  }, [refresh]);
+  }, [refresh])
 
   const columns: ColumnDef<KeyResponse>[] = [
     {
@@ -214,10 +193,7 @@ export function AllKeysTable({
       header: () => null,
       cell: ({ row }) =>
         row.getCanExpand() ? (
-          <button
-            onClick={row.getToggleExpandedHandler()}
-            style={{ cursor: "pointer" }}
-          >
+          <button onClick={row.getToggleExpandedHandler()} style={{ cursor: "pointer" }}>
             {row.getIsExpanded() ? "▼" : "▶"}
           </button>
         ) : null,
@@ -229,7 +205,7 @@ export function AllKeysTable({
       cell: (info) => (
         <div className="overflow-hidden">
           <Tooltip title={info.getValue() as string}>
-            <Button 
+            <Button
               size="xs"
               variant="light"
               className="font-mono text-blue-500 bg-blue-50 hover:bg-blue-100 text-xs font-normal px-2 py-0.5 text-left overflow-hidden truncate max-w-[200px]"
@@ -246,9 +222,9 @@ export function AllKeysTable({
       accessorKey: "key_alias",
       header: "Key Alias",
       cell: (info) => {
-        const value = info.getValue() as string;
+        const value = info.getValue() as string
         return <Tooltip title={value}>{value ? (value.length > 20 ? `${value.slice(0, 20)}...` : value) : "-"}</Tooltip>
-      }
+      },
     },
     {
       id: "key_name",
@@ -261,35 +237,41 @@ export function AllKeysTable({
       accessorKey: "team_id",
       header: "Team Alias",
       cell: ({ row, getValue }) => {
-        const teamId = getValue() as string;
-        const team = teams?.find(t => t.team_id === teamId);
-        return team?.team_alias || "Unknown";
+        const teamId = getValue() as string
+        const team = teams?.find((t) => t.team_id === teamId)
+        return team?.team_alias || "Unknown"
       },
     },
     {
       id: "team_id",
       accessorKey: "team_id",
       header: "Team ID",
-      cell: (info) => <Tooltip title={info.getValue() as string}>{info.getValue() ? `${(info.getValue() as string).slice(0, 7)}...` : "-"}</Tooltip>
+      cell: (info) => (
+        <Tooltip title={info.getValue() as string}>
+          {info.getValue() ? `${(info.getValue() as string).slice(0, 7)}...` : "-"}
+        </Tooltip>
+      ),
     },
     {
       id: "organization_id",
       accessorKey: "organization_id",
       header: "Organization ID",
-      cell: (info) => info.getValue() ? info.renderValue() : "-",
+      cell: (info) => (info.getValue() ? info.renderValue() : "-"),
     },
     {
       id: "user_email",
       accessorKey: "user_id",
       header: "User Email",
       cell: (info) => {
-        const userId = info.getValue() as string;
-        const user = userList.find(u => u.user_id === userId);
+        const userId = info.getValue() as string
+        const user = userList.find((u) => u.user_id === userId)
         return user?.user_email ? (
           <Tooltip title={user?.user_email}>
             <span>{user?.user_email.slice(0, 20)}...</span>
           </Tooltip>
-        ) : "-";
+        ) : (
+          "-"
+        )
       },
     },
     {
@@ -297,15 +279,15 @@ export function AllKeysTable({
       accessorKey: "user_id",
       header: "User ID",
       cell: (info) => {
-        const userId = info.getValue() as string | null;
+        const userId = info.getValue() as string | null
         if (userId && userId.length > 15) {
           return (
             <Tooltip title={userId}>
               <span>{userId.slice(0, 7)}...</span>
             </Tooltip>
-          );
+          )
         }
-        return userId ? userId : "-";
+        return userId ? userId : "-"
       },
     },
     {
@@ -313,8 +295,8 @@ export function AllKeysTable({
       accessorKey: "created_at",
       header: "Created At",
       cell: (info) => {
-        const value = info.getValue();
-        return value ? new Date(value as string).toLocaleDateString() : "-";
+        const value = info.getValue()
+        return value ? new Date(value as string).toLocaleDateString() : "-"
       },
     },
     {
@@ -322,15 +304,15 @@ export function AllKeysTable({
       accessorKey: "created_by",
       header: "Created By",
       cell: (info) => {
-        const value = info.getValue() as string | null;
+        const value = info.getValue() as string | null
         if (value && value.length > 15) {
           return (
             <Tooltip title={value}>
               <span>{value.slice(0, 7)}...</span>
             </Tooltip>
-          );
+          )
         }
-        return value;
+        return value
       },
     },
     {
@@ -338,8 +320,8 @@ export function AllKeysTable({
       accessorKey: "updated_at",
       header: "Updated At",
       cell: (info) => {
-        const value = info.getValue();
-        return value ? new Date(value as string).toLocaleDateString() : "Never";
+        const value = info.getValue()
+        return value ? new Date(value as string).toLocaleDateString() : "Never"
       },
     },
     {
@@ -347,8 +329,8 @@ export function AllKeysTable({
       accessorKey: "expires",
       header: "Expires",
       cell: (info) => {
-        const value = info.getValue();
-        return value ? new Date(value as string).toLocaleDateString() : "Never";
+        const value = info.getValue()
+        return value ? new Date(value as string).toLocaleDateString() : "Never"
       },
     },
     {
@@ -362,11 +344,11 @@ export function AllKeysTable({
       accessorKey: "max_budget",
       header: "Budget (USD)",
       cell: (info) => {
-        const maxBudget = info.getValue() as number | null;
+        const maxBudget = info.getValue() as number | null
         if (maxBudget === null) {
-          return "Unlimited";
+          return "Unlimited"
         }
-        return `$${formatNumberWithCommas(maxBudget)}`;
+        return `$${formatNumberWithCommas(maxBudget)}`
       },
     },
     {
@@ -374,8 +356,8 @@ export function AllKeysTable({
       accessorKey: "budget_reset_at",
       header: "Budget Reset",
       cell: (info) => {
-        const value = info.getValue();
-        return value ? new Date(value as string).toLocaleString() : "Never";
+        const value = info.getValue()
+        return value ? new Date(value as string).toLocaleString() : "Never"
       },
     },
     {
@@ -383,7 +365,7 @@ export function AllKeysTable({
       accessorKey: "models",
       header: "Models",
       cell: (info) => {
-        const models = info.getValue() as string[];
+        const models = info.getValue() as string[]
         return (
           <div className="flex flex-col py-2">
             {Array.isArray(models) ? (
@@ -402,68 +384,54 @@ export function AllKeysTable({
                             className="cursor-pointer"
                             size="xs"
                             onClick={() => {
-                              setExpandedAccordions(prev => ({
+                              setExpandedAccordions((prev) => ({
                                 ...prev,
-                                [info.row.id]: !prev[info.row.id]
-                              }));
+                                [info.row.id]: !prev[info.row.id],
+                              }))
                             }}
                           />
                         </div>
                       )}
                       <div className="flex flex-wrap gap-1">
-                        {models.slice(0, 3).map((model, index) => (
+                        {models.slice(0, 3).map((model, index) =>
                           model === "all-proxy-models" ? (
-                            <Badge
-                              key={index}
-                              size={"xs"}
-                              color="red"
-                            >
+                            <Badge key={index} size={"xs"} color="red">
                               <Text>All Proxy Models</Text>
                             </Badge>
                           ) : (
-                            <Badge
-                              key={index}
-                              size={"xs"}
-                              color="blue"
-                            >
+                            <Badge key={index} size={"xs"} color="blue">
                               <Text>
                                 {model.length > 30
                                   ? `${getModelDisplayName(model).slice(0, 30)}...`
                                   : getModelDisplayName(model)}
                               </Text>
                             </Badge>
-                          )
-                        ))}
+                          ),
+                        )}
                         {models.length > 3 && !expandedAccordions[info.row.id] && (
                           <Badge size={"xs"} color="gray" className="cursor-pointer">
-                            <Text>+{models.length - 3} {models.length - 3 === 1 ? 'more model' : 'more models'}</Text>
+                            <Text>
+                              +{models.length - 3} {models.length - 3 === 1 ? "more model" : "more models"}
+                            </Text>
                           </Badge>
                         )}
                         {expandedAccordions[info.row.id] && (
                           <div className="flex flex-wrap gap-1">
-                            {models.slice(3).map((model, index) => (
+                            {models.slice(3).map((model, index) =>
                               model === "all-proxy-models" ? (
-                                <Badge
-                                  key={index + 3}
-                                  size={"xs"}
-                                  color="red"
-                                >
+                                <Badge key={index + 3} size={"xs"} color="red">
                                   <Text>All Proxy Models</Text>
                                 </Badge>
                               ) : (
-                                <Badge
-                                  key={index + 3}
-                                  size={"xs"}
-                                  color="blue"
-                                >
+                                <Badge key={index + 3} size={"xs"} color="blue">
                                   <Text>
                                     {model.length > 30
                                       ? `${getModelDisplayName(model).slice(0, 30)}...`
                                       : getModelDisplayName(model)}
                                   </Text>
                                 </Badge>
-                              )
-                            ))}
+                              ),
+                            )}
                           </div>
                         )}
                       </div>
@@ -473,154 +441,156 @@ export function AllKeysTable({
               </div>
             ) : null}
           </div>
-        );
+        )
       },
     },
     {
       id: "rate_limits",
       header: "Rate Limits",
       cell: ({ row }) => {
-        const key = row.original;
+        const key = row.original
         return (
           <div>
             <div>TPM: {key.tpm_limit !== null ? key.tpm_limit : "Unlimited"}</div>
             <div>RPM: {key.rpm_limit !== null ? key.rpm_limit : "Unlimited"}</div>
           </div>
-        );
+        )
       },
     },
-  ];
+  ]
 
   const filterOptions: FilterOption[] = [
-    { 
-      name: 'Team ID', 
-      label: 'Team ID', 
-      isSearchable: true, 
+    {
+      name: "Team ID",
+      label: "Team ID",
+      isSearchable: true,
       searchFn: async (searchText: string) => {
-        if (!allTeams || allTeams.length === 0) return [];
-        
-        const filteredTeams = allTeams.filter(team => 
-          team.team_id.toLowerCase().includes(searchText.toLowerCase()) || 
-          (team.team_alias && team.team_alias.toLowerCase().includes(searchText.toLowerCase()))
-        );
-        
-        return filteredTeams.map(team => ({
+        if (!allTeams || allTeams.length === 0) return []
+
+        const filteredTeams = allTeams.filter(
+          (team) =>
+            team.team_id.toLowerCase().includes(searchText.toLowerCase()) ||
+            (team.team_alias && team.team_alias.toLowerCase().includes(searchText.toLowerCase())),
+        )
+
+        return filteredTeams.map((team) => ({
           label: `${team.team_alias || team.team_id} (${team.team_id})`,
-          value: team.team_id
-        }));
-      }
+          value: team.team_id,
+        }))
+      },
     },
-    { 
-      name: 'Organization ID', 
-      label: 'Organization ID', 
-      isSearchable: true, 
+    {
+      name: "Organization ID",
+      label: "Organization ID",
+      isSearchable: true,
       searchFn: async (searchText: string) => {
-        if (!allOrganizations || allOrganizations.length === 0) return [];
-        
-        const filteredOrgs = allOrganizations.filter(org => 
-          org.organization_id?.toLowerCase().includes(searchText.toLowerCase()) ?? false
-        );
-        
+        if (!allOrganizations || allOrganizations.length === 0) return []
+
+        const filteredOrgs = allOrganizations.filter(
+          (org) => org.organization_id?.toLowerCase().includes(searchText.toLowerCase()) ?? false,
+        )
+
         return filteredOrgs
-          .filter(org => org.organization_id !== null && org.organization_id !== undefined)
-          .map(org => ({
-            label: `${org.organization_id || 'Unknown'} (${org.organization_id})`,
-            value: org.organization_id as string
-          }));
-      }
+          .filter((org) => org.organization_id !== null && org.organization_id !== undefined)
+          .map((org) => ({
+            label: `${org.organization_id || "Unknown"} (${org.organization_id})`,
+            value: org.organization_id as string,
+          }))
+      },
     },
     {
       name: "Key Alias",
       label: "Key Alias",
       isSearchable: true,
       searchFn: async (searchText) => {
-        const filteredKeyAliases = allKeyAliases.filter(key => {
+        const filteredKeyAliases = allKeyAliases.filter((key) => {
           return key.toLowerCase().includes(searchText.toLowerCase())
-        });
+        })
 
         return filteredKeyAliases.map((key) => {
           return {
             label: key,
-            value: key
+            value: key,
           }
-        });
-      }
+        })
+      },
     },
     {
-      name: 'User ID',
-      label: 'User ID',
+      name: "User ID",
+      label: "User ID",
       isSearchable: false,
     },
     {
-      name: 'Key Hash',
-      label: 'Key Hash',
+      name: "Key Hash",
+      label: "Key Hash",
       isSearchable: false,
-    }
+    },
+  ]
 
-  ];
-  
   console.log(`keys: ${JSON.stringify(keys)}`)
 
   const table = useReactTable({
     data: filteredKeys,
-    columns: columns.filter(col => col.id !== 'expander'),
+    columns: columns.filter((col) => col.id !== "expander"),
     state: {
       sorting,
     },
     onSortingChange: (updaterOrValue) => {
-      const newSorting = typeof updaterOrValue === 'function' 
-        ? updaterOrValue(sorting)
-        : updaterOrValue;
+      const newSorting = typeof updaterOrValue === "function" ? updaterOrValue(sorting) : updaterOrValue
       console.log(`newSorting: ${JSON.stringify(newSorting)}`)
-      setSorting(newSorting);
+      setSorting(newSorting)
       if (newSorting && newSorting.length > 0) {
-        const sortState = newSorting[0];
-        const sortBy = sortState.id;
-        const sortOrder = sortState.desc ? 'desc' : 'asc';
+        const sortState = newSorting[0]
+        const sortBy = sortState.id
+        const sortOrder = sortState.desc ? "desc" : "asc"
         console.log(`sortBy: ${sortBy}, sortOrder: ${sortOrder}`)
         handleFilterChange({
           ...filters,
-          'Sort By': sortBy,
-          'Sort Order': sortOrder
-        });
-        onSortChange?.(sortBy, sortOrder);
+          "Sort By": sortBy,
+          "Sort Order": sortOrder,
+        })
+        onSortChange?.(sortBy, sortOrder)
       }
     },
     getCoreRowModel: getCoreRowModel(),
     getSortedRowModel: getSortedRowModel(),
     enableSorting: true,
     manualSorting: false,
-  });
+  })
 
   // Update local sorting state when currentSort prop changes
   React.useEffect(() => {
     if (currentSort) {
-      setSorting([{
-        id: currentSort.sortBy,
-        desc: currentSort.sortOrder === 'desc'
-      }]);
+      setSorting([
+        {
+          id: currentSort.sortBy,
+          desc: currentSort.sortOrder === "desc",
+        },
+      ])
     }
-  }, [currentSort]);
+  }, [currentSort])
 
   return (
     <div className="w-full h-full overflow-hidden">
       {selectedKeyId ? (
-        <KeyInfoView 
-          keyId={selectedKeyId} 
+        <KeyInfoView
+          keyId={selectedKeyId}
           onClose={() => setSelectedKeyId(null)}
-          keyData={filteredKeys.find(k => k.token === selectedKeyId)}
+          keyData={filteredKeys.find((k) => k.token === selectedKeyId)}
           onKeyDataUpdate={(updatedKeyData) => {
-            setKeys(keys => keys.map(key => {
-              if (key.token === updatedKeyData.token) {
-                return updateExistingKeys(key, updatedKeyData)
-              }
-              return key
-            }))
-            if (refresh) refresh(); // Minimal fix: refresh the full key list after an update
+            setKeys((keys) =>
+              keys.map((key) => {
+                if (key.token === updatedKeyData.token) {
+                  return updateExistingKeys(key, updatedKeyData)
+                }
+                return key
+              }),
+            )
+            if (refresh) refresh() // Minimal fix: refresh the full key list after an update
           }}
           onDelete={() => {
-            setKeys(keys => keys.filter(key => key.token !== selectedKeyId))
-            if (refresh) refresh(); // Minimal fix: refresh the full key list after a delete
+            setKeys((keys) => keys.filter((key) => key.token !== selectedKeyId))
+            if (refresh) refresh() // Minimal fix: refresh the full key list after a delete
           }}
           accessToken={accessToken}
           userID={userID}
@@ -632,19 +602,28 @@ export function AllKeysTable({
       ) : (
         <div className="border-b py-4 flex-1 overflow-hidden">
           <div className="w-full mb-6">
-            <FilterComponent options={filterOptions} onApplyFilters={handleFilterChange} initialValues={filters} onResetFilters={handleFilterReset}/>
+            <FilterComponent
+              options={filterOptions}
+              onApplyFilters={handleFilterChange}
+              initialValues={filters}
+              onResetFilters={handleFilterReset}
+            />
           </div>
 
           <div className="flex items-center justify-between w-full mb-4">
             <span className="inline-flex text-sm text-gray-700">
-              Showing {isLoading ? "..." : `${(pagination.currentPage - 1) * pageSize + 1} - ${Math.min(pagination.currentPage * pageSize, pagination.totalCount)}`} of {isLoading ? "..." : pagination.totalCount} results
+              Showing{" "}
+              {isLoading
+                ? "..."
+                : `${(pagination.currentPage - 1) * pageSize + 1} - ${Math.min(pagination.currentPage * pageSize, pagination.totalCount)}`}{" "}
+              of {isLoading ? "..." : pagination.totalCount} results
             </span>
-            
+
             <div className="inline-flex items-center gap-2">
               <span className="text-sm text-gray-700">
                 Page {isLoading ? "..." : pagination.currentPage} of {isLoading ? "..." : pagination.totalPages}
               </span>
-              
+
               <button
                 onClick={() => onPageChange(pagination.currentPage - 1)}
                 disabled={isLoading || pagination.currentPage === 1}
@@ -652,9 +631,9 @@ export function AllKeysTable({
               >
                 Previous
               </button>
-              
+
               <button
-                onClick={() => onPageChange(pagination.currentPage + 1)} 
+                onClick={() => onPageChange(pagination.currentPage + 1)}
                 disabled={isLoading || pagination.currentPage === pagination.totalPages}
                 className="px-3 py-1 text-sm border rounded-md hover:bg-gray-50 disabled:opacity-50 disabled:cursor-not-allowed"
               >
@@ -670,30 +649,27 @@ export function AllKeysTable({
                     {table.getHeaderGroups().map((headerGroup) => (
                       <TableRow key={headerGroup.id}>
                         {headerGroup.headers.map((header) => (
-                          <TableHeaderCell 
-                            key={header.id} 
+                          <TableHeaderCell
+                            key={header.id}
                             className={`py-1 h-8 ${
-                              header.id === 'actions' 
-                                ? 'sticky right-0 bg-white shadow-[-4px_0_8px_-6px_rgba(0,0,0,0.1)]' 
-                                : ''
+                              header.id === "actions"
+                                ? "sticky right-0 bg-white shadow-[-4px_0_8px_-6px_rgba(0,0,0,0.1)]"
+                                : ""
                             }`}
                             onClick={header.column.getToggleSortingHandler()}
                           >
                             <div className="flex items-center justify-between gap-2">
                               <div className="flex items-center">
-                                {header.isPlaceholder ? null : (
-                                  flexRender(
-                                    header.column.columnDef.header,
-                                    header.getContext()
-                                  )
-                                )}
+                                {header.isPlaceholder
+                                  ? null
+                                  : flexRender(header.column.columnDef.header, header.getContext())}
                               </div>
-                              {header.id !== 'actions' && (
+                              {header.id !== "actions" && (
                                 <div className="w-4">
                                   {header.column.getIsSorted() ? (
                                     {
                                       asc: <ChevronUpIcon className="h-4 w-4 text-blue-500" />,
-                                      desc: <ChevronDownIcon className="h-4 w-4 text-blue-500" />
+                                      desc: <ChevronDownIcon className="h-4 w-4 text-blue-500" />,
                                     }[header.column.getIsSorted() as string]
                                   ) : (
                                     <SwitchVerticalIcon className="h-4 w-4 text-gray-400" />
@@ -726,7 +702,7 @@ export function AllKeysTable({
                                 whiteSpace: "pre-wrap",
                                 overflow: "hidden",
                               }}
-                              className={`py-0.5 max-h-8 overflow-hidden text-ellipsis whitespace-nowrap ${cell.column.id === 'models' && (cell.getValue() as string[]).length > 3 ? "px-0" : ""}`}
+                              className={`py-0.5 max-h-8 overflow-hidden text-ellipsis whitespace-nowrap ${cell.column.id === "models" && (cell.getValue() as string[]).length > 3 ? "px-0" : ""}`}
                             >
                               {flexRender(cell.column.columnDef.cell, cell.getContext())}
                             </TableCell>
@@ -750,5 +726,5 @@ export function AllKeysTable({
         </div>
       )}
     </div>
-  );
-}
+  )
+}

ui/litellm-dashboard/src/components/common_components/chartUtils.tsx

@@ -0,0 +1,127 @@
+import React from "react";
+import type { CustomTooltipProps } from "@tremor/react";
+import { SpendMetrics } from "../usage/types";
+
+interface ChartDataPoint {
+  date: string;
+  metrics: SpendMetrics;
+}
+
+const colorNameToHex: { [key: string]: string } = {
+  blue: "#3b82f6",
+  cyan: "#06b6d4",
+  indigo: "#6366f1",
+  green: "#22c55e",
+  red: "#ef4444",
+  purple: "#8b5cf6",
+};
+
+export const CustomTooltip = ({
+  active,
+  payload,
+  label,
+}: CustomTooltipProps) => {
+  if (active && payload && payload.length) {
+    const formatCategoryName = (name: string): string => {
+      return name
+        .replace("metrics.", "")
+        .replace(/_/g, " ")
+        .split(" ")
+        .map((word) => word.charAt(0).toUpperCase() + word.slice(1))
+        .join(" ");
+    };
+
+    const getRawValue = (
+      dataPoint: ChartDataPoint,
+      key: string
+    ): number | undefined => {
+      // key is like "metrics.total_tokens"
+      const metricKey = key.substring(
+        key.indexOf(".") + 1
+      ) as keyof SpendMetrics;
+      if (dataPoint.metrics && metricKey in dataPoint.metrics) {
+        return dataPoint.metrics[metricKey];
+      }
+      return undefined;
+    };
+
+    return (
+      <div className="w-56 rounded-tremor-default border border-tremor-border bg-tremor-background p-2 text-tremor-default shadow-tremor-dropdown">
+        <p className="text-tremor-content-strong">{label}</p>
+        {payload.map((item) => {
+          const dataKey = item.dataKey?.toString();
+          if (!dataKey || !item.payload) return null;
+
+          const rawValue = getRawValue(item.payload, dataKey);
+          const isSpend = dataKey.includes("spend");
+          const formattedValue =
+            rawValue !== undefined
+              ? isSpend
+                ? `$${rawValue.toLocaleString(undefined, { minimumFractionDigits: 2, maximumFractionDigits: 2 })}`
+                : rawValue.toLocaleString()
+              : "N/A";
+
+          const colorName = item.color as keyof typeof colorNameToHex;
+          const hexColor = colorNameToHex[colorName] || item.color;
+          return (
+            <div
+              key={dataKey}
+              className="flex items-center justify-between space-x-4"
+            >
+              <div className="flex items-center space-x-2">
+                <span
+                  className={`h-2 w-2 shrink-0 rounded-full ring-2 ring-white drop-shadow-md`}
+                  style={{ backgroundColor: hexColor }}
+                />
+                <p className="font-medium text-tremor-content dark:text-dark-tremor-content">
+                  {formatCategoryName(dataKey)}
+                </p>
+              </div>
+              <p className="font-medium text-tremor-content-emphasis dark:text-dark-tremor-content-emphasis">
+                {formattedValue}
+              </p>
+            </div>
+          );
+        })}
+      </div>
+    );
+  }
+  return null;
+};
+
+export const CustomLegend = ({
+  categories,
+  colors,
+}: {
+  categories: string[];
+  colors: string[];
+}) => {
+  const formatCategoryName = (name: string): string => {
+    return name
+      .replace("metrics.", "")
+      .replace(/_/g, " ")
+      .split(" ")
+      .map((word) => word.charAt(0).toUpperCase() + word.slice(1))
+      .join(" ");
+  };
+
+  return (
+    <div className="flex items-center justify-end space-x-4">
+      {categories.map((category, idx) => {
+        const colorName = colors[idx] as keyof typeof colorNameToHex;
+        const hexColor = colorNameToHex[colorName] || colors[idx];
+        return (
+          <div key={category} className="flex items-center space-x-2">
+            <span
+              className={`h-2 w-2 shrink-0 rounded-full ring-4 ring-white`}
+              style={{ backgroundColor: hexColor }}
+            />
+            <p className="text-sm text-tremor-content dark:text-dark-tremor-content">
+              {formatCategoryName(category)}
+            </p>
+          </div>
+        );
+      })}
+    </div>
+  );
+};

ui/litellm-dashboard/src/components/networking.tsx

@@ -3672,7 +3672,7 @@ export const keyUpdateCall = async (
       const errorData = await response.text();
       handleError(errorData);
       console.error("Error response from the server:", errorData);
-      throw new Error("Network response was not ok");
+      throw new Error(errorData);
     }
     const data = await response.json();
     console.log("Update key Response:", data);

ui/litellm-dashboard/src/components/organisms/create_key_button.tsx

@@ -1,29 +1,12 @@
-"use client";
-import React, { useState, useEffect, useRef, useCallback } from "react";
-import { Button, TextInput, Grid, Col, Select as TremorSelect, SelectItem } from "@tremor/react";
-import {
-  Card,
-  Metric,
-  Text,
-  Title,
-  Subtitle,
-  Accordion,
-  AccordionHeader,
-  AccordionBody,
-} from "@tremor/react";
-import { CopyToClipboard } from "react-copy-to-clipboard";
-import {
-  Button as Button2,
-  Modal,
-  Form,
-  Input,
-  Select,
-  message,
-  Radio,
-} from "antd";
-import NumericalInput from "./shared/numerical_input";
-import { unfurlWildcardModelsInList, getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key";
-import SchemaFormFields from './common_components/check_openapi_schema';
+"use client"
+import React, { useState, useEffect, useRef, useCallback } from "react"
+import { Button, TextInput, Grid, Col, Select as TremorSelect, SelectItem } from "@tremor/react"
+import { Card, Metric, Text, Title, Subtitle, Accordion, AccordionHeader, AccordionBody } from "@tremor/react"
+import { CopyToClipboard } from "react-copy-to-clipboard"
+import { Button as Button2, Modal, Form, Input, Select, message, Radio } from "antd"
+import NumericalInput from "../shared/numerical_input"
+import { unfurlWildcardModelsInList, getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"
+import SchemaFormFields from "../common_components/check_openapi_schema"
 import {
   keyCreateCall,
   slackBudgetAlertsHealthCheck,
@@ -35,22 +18,21 @@ import {
   keyCreateServiceAccountCall,
   fetchMCPAccessGroups,
   getPromptsList,
-} from "./networking";
-import VectorStoreSelector from "./vector_store_management/VectorStoreSelector";
-import { Team } from "./key_team_helpers/key_list";
-import TeamDropdown from "./common_components/team_dropdown";
-import { InfoCircleOutlined } from '@ant-design/icons';
-import { Tooltip } from 'antd';
-import PremiumLoggingSettings from "./common_components/PremiumLoggingSettings";
-import Createuser from "./create_user_button";
-import debounce from 'lodash/debounce';
-import { rolesWithWriteAccess } from '../utils/roles';
-import BudgetDurationDropdown from "./common_components/budget_duration_dropdown";
-import { formatNumberWithCommas } from "@/utils/dataUtils";
-import { callback_map, mapDisplayToInternalNames } from "./callback_info_helpers";
-import MCPServerSelector from "./mcp_server_management/MCPServerSelector";
-import ModelAliasManager from "./common_components/ModelAliasManager";
-
+} from "../networking"
+import VectorStoreSelector from "../vector_store_management/VectorStoreSelector"
+import { Team } from "../key_team_helpers/key_list"
+import TeamDropdown from "../common_components/team_dropdown"
+import { InfoCircleOutlined } from "@ant-design/icons"
+import { Tooltip } from "antd"
+import PremiumLoggingSettings from "../common_components/PremiumLoggingSettings"
+import Createuser from "../create_user_button"
+import debounce from "lodash/debounce"
+import { rolesWithWriteAccess } from "../../utils/roles"
+import BudgetDurationDropdown from "../common_components/budget_duration_dropdown"
+import { formatNumberWithCommas } from "@/utils/dataUtils"
+import { callback_map, mapDisplayToInternalNames } from "../callback_info_helpers"
+import MCPServerSelector from "../mcp_server_management/MCPServerSelector"
+import ModelAliasManager from "../common_components/ModelAliasManager"
 
 const { Option } = Select;
 
@@ -1201,4 +1183,4 @@ const CreateKey: React.FC<CreateKeyProps> = ({
   );
 };
 
-export default CreateKey;
+export default CreateKey;

ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx

@@ -1,10 +1,11 @@
-import React, { useEffect, useState } from "react";
-import { Button, Text, TextInput, Title, Grid, Col } from "@tremor/react";
-import { Modal, Form, InputNumber, message } from "antd";
-import { add } from "date-fns";
-import { regenerateKeyCall } from "./networking";
-import { KeyResponse } from "./key_team_helpers/key_list";
-import { CopyToClipboard } from "react-copy-to-clipboard";
+import React, { useEffect, useState } from "react"
+import { Button, Text, TextInput, Title, Grid, Col } from "@tremor/react"
+import { Modal, Form, InputNumber, message } from "antd"
+import { add } from "date-fns"
+import { regenerateKeyCall } from "../networking"
+import { KeyResponse } from "../key_team_helpers/key_list"
+import { CopyToClipboard } from "react-copy-to-clipboard"
+import NotificationManager from "../molecules/notifications_manager"
 
 interface RegenerateKeyModalProps {
   selectedToken: KeyResponse | null
@@ -70,45 +71,49 @@ export function RegenerateKeyModal({
   const calculateNewExpiryTime = (duration: string | undefined): string | null => {
     if (!duration) return null
 
-      try {
-        const now = new Date();
-        let newExpiry: Date;
+    try {
+      const now = new Date()
+      let newExpiry: Date
 
-        if (duration.endsWith("s")) {
-          newExpiry = add(now, { seconds: parseInt(duration) });
-        } else if (duration.endsWith("h")) {
-          newExpiry = add(now, { hours: parseInt(duration) });
-        } else if (duration.endsWith("d")) {
-          newExpiry = add(now, { days: parseInt(duration) });
-        } else {
-          throw new Error("Invalid duration format");
-        }
-
-        return newExpiry.toLocaleString();
-      } catch (error) {
-        return null;
+      if (duration.endsWith("s")) {
+        newExpiry = add(now, { seconds: parseInt(duration) })
+      } else if (duration.endsWith("h")) {
+        newExpiry = add(now, { hours: parseInt(duration) })
+      } else if (duration.endsWith("d")) {
+        newExpiry = add(now, { days: parseInt(duration) })
+      } else {
+        throw new Error("Invalid duration format")
       }
-    };
+
+      return newExpiry.toLocaleString()
+    } catch (error) {
+      return null
+    }
+  }
 
   useEffect(() => {
     if (regenerateFormData?.duration) {
-      setNewExpiryTime(calculateNewExpiryTime(regenerateFormData.duration));
+      setNewExpiryTime(calculateNewExpiryTime(regenerateFormData.duration))
     } else {
-      setNewExpiryTime(null);
+      setNewExpiryTime(null)
     }
-  }, [regenerateFormData?.duration]);
+  }, [regenerateFormData?.duration])
 
   const handleRegenerateKey = async () => {
     if (!selectedToken || !currentAccessToken) return
 
-    setIsRegenerating(true);
+    setIsRegenerating(true)
     try {
       const formValues = await form.validateFields()
 
       // Use the current access token for the API call
-      const response = await regenerateKeyCall(currentAccessToken, selectedToken.token || selectedToken.token_id, formValues)
+      const response = await regenerateKeyCall(
+        currentAccessToken,
+        selectedToken.token || selectedToken.token_id,
+        formValues,
+      )
       setRegeneratedKey(response.key)
-      message.success("API Key regenerated successfully")
+      NotificationManager.success("API Key regenerated successfully")
 
       console.log("Full regenerate response:", response) // Debug log to see what's returned
 
@@ -142,11 +147,11 @@ export function RegenerateKeyModal({
 
       setIsRegenerating(false)
     } catch (error) {
-      console.error("Error regenerating key:", error);
-      message.error("Failed to regenerate API Key");
-      setIsRegenerating(false); // Reset regenerating state on error
+      console.error("Error regenerating key:", error)
+      NotificationManager.fromBackend(error)
+      setIsRegenerating(false) // Reset regenerating state on error
     }
-  };
+  }
 
   const handleClose = () => {
     setRegeneratedKey(null)
@@ -162,49 +167,43 @@ export function RegenerateKeyModal({
       title="Regenerate API Key"
       open={visible}
       onCancel={handleClose}
-      footer={regeneratedKey ? [
-        <Button key="close" onClick={handleClose}>
-          Close
-        </Button>,
-      ] : [
-        <Button key="cancel" onClick={handleClose} className="mr-2">
-          Cancel
-        </Button>,
-        <Button 
-          key="regenerate" 
-          onClick={handleRegenerateKey}
-          disabled={isRegenerating}
-        >
-          {isRegenerating ? "Regenerating..." : "Regenerate"}
-        </Button>,
-      ]}
+      footer={
+        regeneratedKey
+          ? [
+              <Button key="close" onClick={handleClose}>
+                Close
+              </Button>,
+            ]
+          : [
+              <Button key="cancel" onClick={handleClose} className="mr-2">
+                Cancel
+              </Button>,
+              <Button key="regenerate" onClick={handleRegenerateKey} disabled={isRegenerating}>
+                {isRegenerating ? "Regenerating..." : "Regenerate"}
+              </Button>,
+            ]
+      }
     >
       {regeneratedKey ? (
         <Grid numItems={1} className="gap-2 w-full">
           <Title>Regenerated Key</Title>
           <Col numColSpan={1}>
             <p>
-              Please replace your old key with the new key generated. For
-              security reasons, <b>you will not be able to view it again</b>{" "}
-              through your LiteLLM account. If you lose this secret key, you
-              will need to generate a new one.
+              Please replace your old key with the new key generated. For security reasons,{" "}
+              <b>you will not be able to view it again</b> through your LiteLLM account. If you lose this secret key,
+              you will need to generate a new one.
             </p>
           </Col>
           <Col numColSpan={1}>
             <Text className="mt-3">Key Alias:</Text>
             <div className="bg-gray-100 p-2 rounded mb-2">
-              <pre className="break-words whitespace-normal">
-                {selectedToken?.key_alias || "No alias set"}
-              </pre>
+              <pre className="break-words whitespace-normal">{selectedToken?.key_alias || "No alias set"}</pre>
             </div>
             <Text className="mt-3">New API Key:</Text>
             <div className="bg-gray-100 p-2 rounded mb-2">
               <pre className="break-words whitespace-normal">{regeneratedKey}</pre>
             </div>
-            <CopyToClipboard
-              text={regeneratedKey}
-              onCopy={() => message.success("API Key copied to clipboard")}
-            >
+            <CopyToClipboard text={regeneratedKey} onCopy={() => NotificationManager.success("API Key copied to clipboard")}>
               <Button className="mt-3">Copy API Key</Button>
             </CopyToClipboard>
           </Col>
@@ -215,7 +214,7 @@ export function RegenerateKeyModal({
           layout="vertical"
           onValuesChange={(changedValues) => {
             if ("duration" in changedValues) {
-              setRegenerateFormData((prev: { duration?: string }) => ({ ...prev, duration: changedValues.duration }));
+              setRegenerateFormData((prev: { duration?: string }) => ({ ...prev, duration: changedValues.duration }))
             }
           }}
         >
@@ -237,13 +236,9 @@ export function RegenerateKeyModal({
           <div className="mt-2 text-sm text-gray-500">
             Current expiry: {selectedToken?.expires ? new Date(selectedToken.expires).toLocaleString() : "Never"}
           </div>
-          {newExpiryTime && (
-            <div className="mt-2 text-sm text-green-600">
-              New expiry: {newExpiryTime}
-            </div>
-          )}
+          {newExpiryTime && <div className="mt-2 text-sm text-green-600">New expiry: {newExpiryTime}</div>}
         </Form>
       )}
     </Modal>
-  );
-}
+  )
+}

ui/litellm-dashboard/src/components/tag_management/tag_info.tsx

@@ -1,149 +1,116 @@
-import React, { useState, useEffect } from "react";
-import {
-  Card,
-  Text,
-  Title,
-  Button,
-  Badge,
-} from "@tremor/react";
-import {
-  Form,
-  Input,
-  Select as Select2,
-  message,
-  Tooltip,
-} from "antd";
-import { InfoCircleOutlined } from '@ant-design/icons';
-import { fetchUserModels } from "../create_key_button";
-import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
-import { tagInfoCall, tagUpdateCall } from "../networking";
-import { Tag, TagInfoResponse } from "./types";
+import React, { useState, useEffect } from "react"
+import { Card, Text, Title, Button, Badge } from "@tremor/react"
+import { Form, Input, Select as Select2, message, Tooltip } from "antd"
+import { InfoCircleOutlined } from "@ant-design/icons"
+import { fetchUserModels } from "../organisms/create_key_button"
+import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"
+import { tagInfoCall, tagUpdateCall } from "../networking"
+import { Tag, TagInfoResponse } from "./types"
 
 interface TagInfoViewProps {
-  tagId: string;
-  onClose: () => void;
-  accessToken: string | null;
-  is_admin: boolean;
-  editTag: boolean;
+  tagId: string
+  onClose: () => void
+  accessToken: string | null
+  is_admin: boolean
+  editTag: boolean
 }
 
-const TagInfoView: React.FC<TagInfoViewProps> = ({
-  tagId,
-  onClose,
-  accessToken,
-  is_admin,
-  editTag,
-}) => {
-  const [form] = Form.useForm();
-  const [tagDetails, setTagDetails] = useState<Tag | null>(null);
-  const [isEditing, setIsEditing] = useState<boolean>(editTag);
-  const [userModels, setUserModels] = useState<string[]>([]);
+const TagInfoView: React.FC<TagInfoViewProps> = ({ tagId, onClose, accessToken, is_admin, editTag }) => {
+  const [form] = Form.useForm()
+  const [tagDetails, setTagDetails] = useState<Tag | null>(null)
+  const [isEditing, setIsEditing] = useState<boolean>(editTag)
+  const [userModels, setUserModels] = useState<string[]>([])
 
   const fetchTagDetails = async () => {
-    if (!accessToken) return;
+    if (!accessToken) return
     try {
-      const response = await tagInfoCall(accessToken, [tagId]);
-      const tagData = response[tagId];
+      const response = await tagInfoCall(accessToken, [tagId])
+      const tagData = response[tagId]
       if (tagData) {
-        setTagDetails(tagData);
+        setTagDetails(tagData)
         if (editTag) {
           form.setFieldsValue({
             name: tagData.name,
             description: tagData.description,
             models: tagData.models,
-          });
+          })
         }
       }
     } catch (error) {
-      console.error("Error fetching tag details:", error);
-      message.error("Error fetching tag details: " + error);
+      console.error("Error fetching tag details:", error)
+      message.error("Error fetching tag details: " + error)
     }
-  };
+  }
 
   useEffect(() => {
-    fetchTagDetails();
-  }, [tagId, accessToken]);
+    fetchTagDetails()
+  }, [tagId, accessToken])
 
   useEffect(() => {
     if (accessToken) {
       // Using dummy values for userID and userRole since they're required by the function
       // TODO: Pass these as props if needed for the actual API implementation
-      fetchUserModels("dummy-user", "Admin", accessToken, setUserModels);
+      fetchUserModels("dummy-user", "Admin", accessToken, setUserModels)
     }
-  }, [accessToken]);
+  }, [accessToken])
 
   const handleSave = async (values: any) => {
-    if (!accessToken) return;
+    if (!accessToken) return
     try {
       await tagUpdateCall(accessToken, {
         name: values.name,
         description: values.description,
         models: values.models,
-      });
-      message.success("Tag updated successfully");
-      setIsEditing(false);
-      fetchTagDetails();
+      })
+      message.success("Tag updated successfully")
+      setIsEditing(false)
+      fetchTagDetails()
     } catch (error) {
-      console.error("Error updating tag:", error);
-      message.error("Error updating tag: " + error);
+      console.error("Error updating tag:", error)
+      message.error("Error updating tag: " + error)
     }
-  };
+  }
 
   if (!tagDetails) {
-    return <div>Loading...</div>;
+    return <div>Loading...</div>
   }
 
   return (
     <div className="p-4">
       <div className="flex justify-between items-center mb-6">
         <div>
-          <Button onClick={onClose} className="mb-4">← Back to Tags</Button>
+          <Button onClick={onClose} className="mb-4">
+            ← Back to Tags
+          </Button>
           <Title>Tag Name: {tagDetails.name}</Title>
           <Text className="text-gray-500">{tagDetails.description || "No description"}</Text>
         </div>
-        {is_admin && !isEditing && (
-          <Button onClick={() => setIsEditing(true)}>Edit Tag</Button>
-        )}
+        {is_admin && !isEditing && <Button onClick={() => setIsEditing(true)}>Edit Tag</Button>}
       </div>
 
       {isEditing ? (
         <Card>
-          <Form
-            form={form}
-            onFinish={handleSave}
-            layout="vertical"
-            initialValues={tagDetails}
-          >
-            <Form.Item
-              label="Tag Name"
-              name="name"
-              rules={[{ required: true, message: "Please input a tag name" }]}
-            >
+          <Form form={form} onFinish={handleSave} layout="vertical" initialValues={tagDetails}>
+            <Form.Item label="Tag Name" name="name" rules={[{ required: true, message: "Please input a tag name" }]}>
               <Input />
             </Form.Item>
 
-            <Form.Item
-              label="Description"
-              name="description"
-            >
+            <Form.Item label="Description" name="description">
               <Input.TextArea rows={4} />
             </Form.Item>
 
             <Form.Item
               label={
                 <span>
-                  Allowed LLMs{' '}
+                  Allowed LLMs{" "}
                   <Tooltip title="Select which LLMs are allowed to process this type of data">
-                    <InfoCircleOutlined style={{ marginLeft: '4px' }} />
+                    <InfoCircleOutlined style={{ marginLeft: "4px" }} />
                   </Tooltip>
                 </span>
               }
               name="models"
             >
-              <Select2
-                mode="multiple"
-                placeholder="Select LLMs"
-              >
+              <Select2 mode="multiple" placeholder="Select LLMs">
                 {userModels.map((modelId) => (
                   <Select2.Option key={modelId} value={modelId}>
                     {getModelDisplayName(modelId)}
@@ -179,9 +146,7 @@ const TagInfoView: React.FC<TagInfoViewProps> = ({
                   ) : (
                     tagDetails.models.map((modelId) => (
                       <Badge key={modelId} color="blue">
-                        <Tooltip title={`ID: ${modelId}`}>
-                          {tagDetails.model_info?.[modelId] || modelId}
-                        </Tooltip>
+                        <Tooltip title={`ID: ${modelId}`}>{tagDetails.model_info?.[modelId] || modelId}</Tooltip>
                       </Badge>
                     ))
                   )}
@@ -200,7 +165,7 @@ const TagInfoView: React.FC<TagInfoViewProps> = ({
         </div>
       )}
     </div>
-  );
-};
+  )
+}
 
-export default TagInfoView; 
+export default TagInfoView

ui/litellm-dashboard/src/components/team/member_permissions.tsx

@@ -1,4 +1,4 @@
-import React, { useState, useEffect } from "react";
+import React, { useState, useEffect } from "react"
 import {
   Card,
   Title,
@@ -10,83 +10,79 @@ import {
   TableBody,
   TableRow,
   TableCell,
-} from "@tremor/react";
-import { Button, message, Checkbox, Empty } from "antd";
-import { ReloadOutlined, SaveOutlined } from "@ant-design/icons";
-import { getTeamPermissionsCall, teamPermissionsUpdateCall } from "@/components/networking";
-import { getPermissionInfo } from "./permission_definitions";
+} from "@tremor/react"
+import { Button, message, Checkbox, Empty } from "antd"
+import { ReloadOutlined, SaveOutlined } from "@ant-design/icons"
+import { getTeamPermissionsCall, teamPermissionsUpdateCall } from "@/components/networking"
+import { getPermissionInfo } from "./permission_definitions"
 
 interface MemberPermissionsProps {
-  teamId: string;
-  accessToken: string | null;
-  canEditTeam: boolean;
+  teamId: string
+  accessToken: string | null
+  canEditTeam: boolean
 }
 
-const MemberPermissions: React.FC<MemberPermissionsProps> = ({
-  teamId,
-  accessToken,
-  canEditTeam,
-}) => {
-  const [permissions, setPermissions] = useState<string[]>([]);
-  const [selectedPermissions, setSelectedPermissions] = useState<string[]>([]);
-  const [loading, setLoading] = useState(true);
-  const [saving, setSaving] = useState(false);
-  const [hasChanges, setHasChanges] = useState(false);
+const MemberPermissions: React.FC<MemberPermissionsProps> = ({ teamId, accessToken, canEditTeam }) => {
+  const [permissions, setPermissions] = useState<string[]>([])
+  const [selectedPermissions, setSelectedPermissions] = useState<string[]>([])
+  const [loading, setLoading] = useState(true)
+  const [saving, setSaving] = useState(false)
+  const [hasChanges, setHasChanges] = useState(false)
 
   const fetchPermissions = async () => {
     try {
-      setLoading(true);
-      if (!accessToken) return;
-      const response = await getTeamPermissionsCall(accessToken, teamId);
-      const allPermissions = response.all_available_permissions || [];
-      setPermissions(allPermissions);
-      const teamPermissions = response.team_member_permissions || [];
-      setSelectedPermissions(teamPermissions);
-      setHasChanges(false);
+      setLoading(true)
+      if (!accessToken) return
+      const response = await getTeamPermissionsCall(accessToken, teamId)
+      const allPermissions = response.all_available_permissions || []
+      setPermissions(allPermissions)
+      const teamPermissions = response.team_member_permissions || []
+      setSelectedPermissions(teamPermissions)
+      setHasChanges(false)
     } catch (error) {
-      message.error("Failed to load permissions");
-      console.error("Error fetching permissions:", error);
+      message.error("Failed to load permissions")
+      console.error("Error fetching permissions:", error)
     } finally {
-      setLoading(false);
+      setLoading(false)
     }
-  };
+  }
 
   useEffect(() => {
-    fetchPermissions();
-  }, [teamId, accessToken]);
+    fetchPermissions()
+  }, [teamId, accessToken])
 
   const handlePermissionChange = (permission: string, checked: boolean) => {
     const newSelectedPermissions = checked
       ? [...selectedPermissions, permission]
-      : selectedPermissions.filter((p) => p !== permission);
-    setSelectedPermissions(newSelectedPermissions);
-    setHasChanges(true);
-  };
+      : selectedPermissions.filter((p) => p !== permission)
+    setSelectedPermissions(newSelectedPermissions)
+    setHasChanges(true)
+  }
 
   const handleSave = async () => {
     try {
-      if (!accessToken) return;
-      setSaving(true);
-      await teamPermissionsUpdateCall(accessToken, teamId, selectedPermissions);
-      message.success("Permissions updated successfully");
-      setHasChanges(false);
+      if (!accessToken) return
+      setSaving(true)
+      await teamPermissionsUpdateCall(accessToken, teamId, selectedPermissions)
+      message.success("Permissions updated successfully")
+      setHasChanges(false)
     } catch (error) {
-      message.error("Failed to update permissions");
-      console.error("Error updating permissions:", error);
+      message.error("Failed to update permissions")
+      console.error("Error updating permissions:", error)
     } finally {
-      setSaving(false);
+      setSaving(false)
     }
-  };
-
-  const handleReset = () => {
-    fetchPermissions();
-  };
-
-  if (loading) {
-    return <div className="p-6 text-center">Loading permissions...</div>;
   }
 
-  const hasPermissions = permissions.length > 0;
+  const handleReset = () => {
+    fetchPermissions()
+  }
+
+  if (loading) {
+    return <div className="p-6 text-center">Loading permissions...</div>
+  }
+
+  const hasPermissions = permissions.length > 0
 
   return (
     <Card className="bg-white shadow-md rounded-md p-6">
@@ -97,79 +93,66 @@ const MemberPermissions: React.FC<MemberPermissionsProps> = ({
             <Button icon={<ReloadOutlined />} onClick={handleReset}>
               Reset
             </Button>
-            <TremorButton
-              onClick={handleSave}
-              loading={saving}
-              className="flex items-center gap-2"
-            >
+            <TremorButton onClick={handleSave} loading={saving} className="flex items-center gap-2">
               <SaveOutlined /> Save Changes
             </TremorButton>
           </div>
         )}
       </div>
 
-      <Text className="mb-6 text-gray-600">
-        Control what team members can do when they are not team admins.
-      </Text>
+      <Text className="mb-6 text-gray-600">Control what team members can do when they are not team admins.</Text>
 
       {hasPermissions ? (
-        <Table className="mt-4">
-          <TableHead>
-            <TableRow>
-              <TableHeaderCell>Method</TableHeaderCell>
-              <TableHeaderCell>Endpoint</TableHeaderCell>
-              <TableHeaderCell>Description</TableHeaderCell>
-              <TableHeaderCell className="text-right">Access</TableHeaderCell>
-            </TableRow>
-          </TableHead>
-          <TableBody>
-            {permissions.map((permission) => {
-              const permInfo = getPermissionInfo(permission);
-              return (
-                <TableRow
-                  key={permission}
-                  className="hover:bg-gray-50 transition-colors"
-                >
-                  <TableCell>
-                    <span
-                      className={`px-2 py-1 rounded text-xs font-medium ${
-                        permInfo.method === "GET"
-                          ? "bg-blue-100 text-blue-800"
-                          : "bg-green-100 text-green-800"
-                      }`}
-                    >
-                      {permInfo.method}
-                    </span>
-                  </TableCell>
-                  <TableCell>
-                    <span className="font-mono text-sm text-gray-800">
-                      {permInfo.endpoint}
-                    </span>
-                  </TableCell>
-                  <TableCell className="text-gray-700">
-                    {permInfo.description}
-                  </TableCell>
-                  <TableCell className="text-right">
-                    <Checkbox
-                      checked={selectedPermissions.includes(permission)}
-                      onChange={(e) =>
-                        handlePermissionChange(permission, e.target.checked)
-                      }
-                      disabled={!canEditTeam}
-                    />
-                  </TableCell>
-                </TableRow>
-              );
-            })}
-          </TableBody>
-        </Table>
+        <div className="overflow-x-auto">
+          <Table className=" min-w-full">
+            <TableHead>
+              <TableRow>
+                <TableHeaderCell>Method</TableHeaderCell>
+                <TableHeaderCell>Endpoint</TableHeaderCell>
+                <TableHeaderCell>Description</TableHeaderCell>
+                <TableHeaderCell className="sticky right-0 bg-white shadow-[-4px_0_4px_-4px_rgba(0,0,0,0.1)] text-center">
+                  Allow Access
+                </TableHeaderCell>
+              </TableRow>
+            </TableHead>
+            <TableBody>
+              {permissions.map((permission) => {
+                const permInfo = getPermissionInfo(permission)
+                return (
+                  <TableRow key={permission} className="hover:bg-gray-50 transition-colors">
+                    <TableCell>
+                      <span
+                        className={`px-2 py-1 rounded text-xs font-medium ${
+                          permInfo.method === "GET" ? "bg-blue-100 text-blue-800" : "bg-green-100 text-green-800"
+                        }`}
+                      >
+                        {permInfo.method}
+                      </span>
+                    </TableCell>
+                    <TableCell>
+                      <span className="font-mono text-sm text-gray-800">{permInfo.endpoint}</span>
+                    </TableCell>
+                    <TableCell className="text-gray-700">{permInfo.description}</TableCell>
+                    <TableCell className="sticky right-0 bg-white shadow-[-4px_0_4px_-4px_rgba(0,0,0,0.1)] text-center">
+                      <Checkbox
+                        checked={selectedPermissions.includes(permission)}
+                        onChange={(e) => handlePermissionChange(permission, e.target.checked)}
+                        disabled={!canEditTeam}
+                      />
+                    </TableCell>
+                  </TableRow>
+                )
+              })}
+            </TableBody>
+          </Table>
+        </div>
       ) : (
         <div className="py-12">
           <Empty description="No permissions available" />
         </div>
       )}
     </Card>
-  );
-};
+  )
+}
 
-export default MemberPermissions;
+export default MemberPermissions

ui/litellm-dashboard/src/components/templates/key_edit_view.tsx

@@ -1,137 +1,131 @@
-import React, { useState, useEffect } from "react";
-import { Form, Input, Select, Button as AntdButton, Tooltip } from "antd";
-import { Button as TremorButton, TextInput } from "@tremor/react";
-import { KeyResponse } from "./key_team_helpers/key_list";
-import { fetchTeamModels } from "../components/create_key_button";
-import { modelAvailableCall, getPromptsList } from "./networking";
-import NumericalInput from "./shared/numerical_input";
-import VectorStoreSelector from "./vector_store_management/VectorStoreSelector";
-import MCPServerSelector from "./mcp_server_management/MCPServerSelector";
-import EditLoggingSettings from "./team/EditLoggingSettings";
-import { extractLoggingSettings, formatMetadataForDisplay } from "./key_info_utils";
-import { fetchMCPAccessGroups } from "./networking";
-import { mapInternalToDisplayNames, mapDisplayToInternalNames } from "./callback_info_helpers";
+import React, { useState, useEffect } from "react"
+import { Form, Input, Select, Button as AntdButton, Tooltip } from "antd"
+import { Button as TremorButton, TextInput } from "@tremor/react"
+import { KeyResponse } from "../key_team_helpers/key_list"
+import { fetchTeamModels } from "../organisms/create_key_button"
+import { modelAvailableCall, getPromptsList } from "../networking"
+import NumericalInput from "../shared/numerical_input"
+import VectorStoreSelector from "../vector_store_management/VectorStoreSelector"
+import MCPServerSelector from "../mcp_server_management/MCPServerSelector"
+import EditLoggingSettings from "../team/EditLoggingSettings"
+import { extractLoggingSettings, formatMetadataForDisplay } from "../key_info_utils"
+import { fetchMCPAccessGroups } from "../networking"
+import { mapInternalToDisplayNames, mapDisplayToInternalNames } from "../callback_info_helpers"
 
 interface KeyEditViewProps {
-  keyData: KeyResponse;
-  onCancel: () => void;
-  onSubmit: (values: any) => Promise<void>;
-  teams?: any[] | null;
-  accessToken: string | null;
-  userID: string | null;
-  userRole: string | null;
-  premiumUser?: boolean;
+  keyData: KeyResponse
+  onCancel: () => void
+  onSubmit: (values: any) => Promise<void>
+  teams?: any[] | null
+  accessToken: string | null
+  userID: string | null
+  userRole: string | null
+  premiumUser?: boolean
 }
 
 // Add this helper function
 const getAvailableModelsForKey = (keyData: KeyResponse, teams: any[] | null): string[] => {
   // If no teams data is available, return empty array
-  console.log("getAvailableModelsForKey:", teams);
+  console.log("getAvailableModelsForKey:", teams)
   if (!teams || !keyData.team_id) {
-    return [];
+    return []
   }
 
   // Find the team that matches the key's team_id
-  const keyTeam = teams.find(team => team.team_id === keyData.team_id);
-  
+  const keyTeam = teams.find((team) => team.team_id === keyData.team_id)
+
   // If team found and has models, return those models
   if (keyTeam?.models) {
-    return keyTeam.models;
+    return keyTeam.models
   }
 
-  return [];
-};
+  return []
+}
 
-export function KeyEditView({ 
-    keyData, 
-    onCancel, 
-    onSubmit, 
-    teams,
-    accessToken,
-    userID,
-    userRole,
-    premiumUser = false
+export function KeyEditView({
+  keyData,
+  onCancel,
+  onSubmit,
+  teams,
+  accessToken,
+  userID,
+  userRole,
+  premiumUser = false,
 }: KeyEditViewProps) {
-  const [form] = Form.useForm();
-  const [userModels, setUserModels] = useState<string[]>([]);
-  const [promptsList, setPromptsList] = useState<string[]>([]);
-  const team = teams?.find(team => team.team_id === keyData.team_id);
-  const [availableModels, setAvailableModels] = useState<string[]>([]);
-  const [mcpAccessGroups, setMcpAccessGroups] = useState<string[]>([]);
-  const [mcpAccessGroupsLoaded, setMcpAccessGroupsLoaded] = useState(false);
+  const [form] = Form.useForm()
+  const [userModels, setUserModels] = useState<string[]>([])
+  const [promptsList, setPromptsList] = useState<string[]>([])
+  const team = teams?.find((team) => team.team_id === keyData.team_id)
+  const [availableModels, setAvailableModels] = useState<string[]>([])
+  const [mcpAccessGroups, setMcpAccessGroups] = useState<string[]>([])
+  const [mcpAccessGroupsLoaded, setMcpAccessGroupsLoaded] = useState(false)
   const [disabledCallbacks, setDisabledCallbacks] = useState<string[]>(
-    Array.isArray(keyData.metadata?.litellm_disabled_callbacks) 
+    Array.isArray(keyData.metadata?.litellm_disabled_callbacks)
       ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks)
-      : []
-  );
+      : [],
+  )
 
   const fetchMcpAccessGroups = async () => {
-    if (!accessToken) return;
-    if (mcpAccessGroupsLoaded) return;
+    if (!accessToken) return
+    if (mcpAccessGroupsLoaded) return
     try {
-      const groups = await fetchMCPAccessGroups(accessToken);
-      setMcpAccessGroups(groups);
-      setMcpAccessGroupsLoaded(true);
+      const groups = await fetchMCPAccessGroups(accessToken)
+      setMcpAccessGroups(groups)
+      setMcpAccessGroupsLoaded(true)
     } catch (error) {
-      console.error("Failed to fetch MCP access groups:", error);
+      console.error("Failed to fetch MCP access groups:", error)
     }
-  };
+  }
 
   useEffect(() => {
     const fetchModels = async () => {
-      if (!userID || !userRole || !accessToken) return;
+      if (!userID || !userRole || !accessToken) return
 
       try {
         if (keyData.team_id === null) {
           // Fetch user models if no team
-          const model_available = await modelAvailableCall(
-            accessToken,
-            userID, 
-            userRole
-          );
-          const available_model_names = model_available["data"].map(
-            (element: { id: string }) => element.id
-          );
-          setAvailableModels(available_model_names);
+          const model_available = await modelAvailableCall(accessToken, userID, userRole)
+          const available_model_names = model_available["data"].map((element: { id: string }) => element.id)
+          setAvailableModels(available_model_names)
         } else if (team?.team_id) {
           // Fetch team models if team exists
-          const models = await fetchTeamModels(userID, userRole, accessToken, team.team_id);
-          setAvailableModels(Array.from(new Set([...team.models, ...models])));
+          const models = await fetchTeamModels(userID, userRole, accessToken, team.team_id)
+          setAvailableModels(Array.from(new Set([...team.models, ...models])))
         }
       } catch (error) {
-        console.error("Error fetching models:", error);
+        console.error("Error fetching models:", error)
       }
-    };
+    }
 
     const fetchPrompts = async () => {
-      if (!accessToken) return;
+      if (!accessToken) return
       try {
-        const response = await getPromptsList(accessToken);
-        setPromptsList(response.prompts.map(prompt => prompt.prompt_id));
+        const response = await getPromptsList(accessToken)
+        setPromptsList(response.prompts.map((prompt) => prompt.prompt_id))
       } catch (error) {
-        console.error("Failed to fetch prompts:", error);
+        console.error("Failed to fetch prompts:", error)
       }
-    };
+    }
 
-    fetchPrompts();
-    fetchModels();
-  }, [userID, userRole, accessToken, team, keyData.team_id]);
+    fetchPrompts()
+    fetchModels()
+  }, [userID, userRole, accessToken, team, keyData.team_id])
 
   // Sync disabled callbacks with form when component mounts
   useEffect(() => {
-    form.setFieldValue('disabled_callbacks', disabledCallbacks);
-  }, [form, disabledCallbacks]);
+    form.setFieldValue("disabled_callbacks", disabledCallbacks)
+  }, [form, disabledCallbacks])
 
   // Convert API budget duration to form format
   const getBudgetDuration = (duration: string | null) => {
-    if (!duration) return null;
+    if (!duration) return null
     const durationMap: Record<string, string> = {
       "24h": "daily",
       "7d": "weekly",
-      "30d": "monthly"
-    };
-    return durationMap[duration] || null;
-  };
+      "30d": "monthly",
+    }
+    return durationMap[duration] || null
+  }
 
   // Set initial form values
   const initialValues = {
@@ -143,39 +137,28 @@ export function KeyEditView({
     vector_stores: keyData.object_permission?.vector_stores || [],
     mcp_servers_and_groups: {
       servers: keyData.object_permission?.mcp_servers || [],
-      accessGroups: keyData.object_permission?.mcp_access_groups || []
+      accessGroups: keyData.object_permission?.mcp_access_groups || [],
     },
     logging_settings: extractLoggingSettings(keyData.metadata),
-    disabled_callbacks: Array.isArray(keyData.metadata?.litellm_disabled_callbacks) 
+    disabled_callbacks: Array.isArray(keyData.metadata?.litellm_disabled_callbacks)
       ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks)
-      : []
-  };
+      : [],
+  }
 
-  console.log("premiumUser:", premiumUser);
+  console.log("premiumUser:", premiumUser)
 
   return (
-    <Form
-      form={form}
-      onFinish={onSubmit}
-      initialValues={initialValues}
-      layout="vertical"
-    >
+    <Form form={form} onFinish={onSubmit} initialValues={initialValues} layout="vertical">
       <Form.Item label="Key Alias" name="key_alias">
         <TextInput />
       </Form.Item>
 
       <Form.Item label="Models" name="models">
-        <Select
-          mode="multiple"
-          placeholder="Select models"
-          style={{ width: "100%" }}
-        >
+        <Select mode="multiple" placeholder="Select models" style={{ width: "100%" }}>
           {/* Only show All Team Models if team has models */}
-          {availableModels.length > 0 && (
-            <Select.Option value="all-team-models">All Team Models</Select.Option>
-          )}
+          {availableModels.length > 0 && <Select.Option value="all-team-models">All Team Models</Select.Option>}
           {/* Show available team models */}
-          {availableModels.map(model => (
+          {availableModels.map((model) => (
             <Select.Option key={model} value={model}>
               {model}
             </Select.Option>
@@ -184,7 +167,7 @@ export function KeyEditView({
       </Form.Item>
 
       <Form.Item label="Max Budget (USD)" name="max_budget">
-        <NumericalInput step={0.01} style={{ width: "100%" }} placeholder="Enter a numerical value"/>
+        <NumericalInput step={0.01} style={{ width: "100%" }} placeholder="Enter a numerical value" />
       </Form.Item>
 
       <Form.Item label="Reset Budget" name="budget_duration">
@@ -196,31 +179,27 @@ export function KeyEditView({
       </Form.Item>
 
       <Form.Item label="TPM Limit" name="tpm_limit">
-        <NumericalInput min={0}/>
+        <NumericalInput min={0} />
       </Form.Item>
 
       <Form.Item label="RPM Limit" name="rpm_limit">
-        <NumericalInput min={0}/>
+        <NumericalInput min={0} />
       </Form.Item>
 
-      <Form.Item label="Max Parallel Requests" name="max_parallel_requests">  
-        <NumericalInput min={0}/>
+      <Form.Item label="Max Parallel Requests" name="max_parallel_requests">
+        <NumericalInput min={0} />
       </Form.Item>
 
       <Form.Item label="Model TPM Limit" name="model_tpm_limit">
-        <Input.TextArea rows={4}  placeholder='{"gpt-4": 100, "claude-v1": 200}'/>
+        <Input.TextArea rows={4} placeholder='{"gpt-4": 100, "claude-v1": 200}' />
       </Form.Item>
 
       <Form.Item label="Model RPM Limit" name="model_rpm_limit">
-        <Input.TextArea rows={4}  placeholder='{"gpt-4": 100, "claude-v1": 200}'/>
+        <Input.TextArea rows={4} placeholder='{"gpt-4": 100, "claude-v1": 200}' />
       </Form.Item>
 
-      
       <Form.Item label="Guardrails" name="guardrails">
-        <Tooltip 
-          title={!premiumUser ? "Setting guardrails by key is a premium feature" : ""}
-          placement="top"
-        >
+        <Tooltip title={!premiumUser ? "Setting guardrails by key is a premium feature" : ""} placement="top">
           <Select
             mode="tags"
             style={{ width: "100%" }}
@@ -229,7 +208,7 @@ export function KeyEditView({
               !premiumUser
                 ? "Premium feature - Upgrade to set guardrails by key"
                 : Array.isArray(keyData.metadata?.guardrails) && keyData.metadata.guardrails.length > 0
-                  ? `Current: ${keyData.metadata.guardrails.join(', ')}`
+                  ? `Current: ${keyData.metadata.guardrails.join(", ")}`
                   : "Select or enter guardrails"
             }
           />
@@ -237,10 +216,7 @@ export function KeyEditView({
       </Form.Item>
 
       <Form.Item label="Prompts" name="prompts">
-        <Tooltip 
-          title={!premiumUser ? "Setting prompts by key is a premium feature" : ""}
-          placement="top"
-        >
+        <Tooltip title={!premiumUser ? "Setting prompts by key is a premium feature" : ""} placement="top">
           <Select
             mode="tags"
             style={{ width: "100%" }}
@@ -249,18 +225,18 @@ export function KeyEditView({
               !premiumUser
                 ? "Premium feature - Upgrade to set prompts by key"
                 : Array.isArray(keyData.metadata?.prompts) && keyData.metadata.prompts.length > 0
-                  ? `Current: ${keyData.metadata.prompts.join(', ')}`
+                  ? `Current: ${keyData.metadata.prompts.join(", ")}`
                   : "Select or enter prompts"
             }
-            options={promptsList.map(name => ({ value: name, label: name }))}
+            options={promptsList.map((name) => ({ value: name, label: name }))}
           />
         </Tooltip>
       </Form.Item>
 
       <Form.Item label="Vector Stores" name="vector_stores">
         <VectorStoreSelector
-          onChange={(values: string[]) => form.setFieldValue('vector_stores', values)}
-          value={form.getFieldValue('vector_stores')}
+          onChange={(values: string[]) => form.setFieldValue("vector_stores", values)}
+          value={form.getFieldValue("vector_stores")}
           accessToken={accessToken || ""}
           placeholder="Select vector stores"
         />
@@ -268,20 +244,17 @@ export function KeyEditView({
 
       <Form.Item label="MCP Servers / Access Groups" name="mcp_servers_and_groups">
         <MCPServerSelector
-          onChange={val => form.setFieldValue('mcp_servers_and_groups', val)}
-          value={form.getFieldValue('mcp_servers_and_groups')}
-          accessToken={accessToken || ''}
+          onChange={(val) => form.setFieldValue("mcp_servers_and_groups", val)}
+          value={form.getFieldValue("mcp_servers_and_groups")}
+          accessToken={accessToken || ""}
           placeholder="Select MCP servers or access groups (optional)"
         />
       </Form.Item>
 
       <Form.Item label="Team ID" name="team_id">
-        <Select
-          placeholder="Select team"
-          style={{ width: "100%" }}
-        >
+        <Select placeholder="Select team" style={{ width: "100%" }}>
           {/* Only show All Team Models if team has models */}
-          {teams?.map(team => (
+          {teams?.map((team) => (
             <Select.Option key={team.team_id} value={team.team_id}>
               {`${team.team_alias} (${team.team_id})`}
             </Select.Option>
@@ -290,25 +263,23 @@ export function KeyEditView({
       </Form.Item>
       <Form.Item label="Logging Settings" name="logging_settings">
         <EditLoggingSettings
-          value={form.getFieldValue('logging_settings')}
-          onChange={(values) => form.setFieldValue('logging_settings', values)}
+          value={form.getFieldValue("logging_settings")}
+          onChange={(values) => form.setFieldValue("logging_settings", values)}
           disabledCallbacks={disabledCallbacks}
           onDisabledCallbacksChange={(internalValues) => {
             // Convert internal values back to display names for UI state
-            const displayNames = mapInternalToDisplayNames(internalValues);
-            setDisabledCallbacks(displayNames);
+            const displayNames = mapInternalToDisplayNames(internalValues)
+            setDisabledCallbacks(displayNames)
             // Store internal values in form for submission
-            form.setFieldValue('disabled_callbacks', internalValues);
+            form.setFieldValue("disabled_callbacks", internalValues)
           }}
         />
       </Form.Item>
 
-
       <Form.Item label="Metadata" name="metadata">
         <Input.TextArea rows={10} />
       </Form.Item>
 
-
       {/* Hidden form field for token */}
       <Form.Item name="token" hidden>
         <Input />
@@ -321,14 +292,10 @@ export function KeyEditView({
 
       <div className="sticky z-10 bg-white p-4 border-t border-gray-200 bottom-[-1.5rem] inset-x-[-1.5rem]">
         <div className="flex justify-end items-center gap-2">
-          <AntdButton onClick={onCancel}>
-            Cancel
-          </AntdButton>
-          <TremorButton type="submit">
-            Save Changes
-          </TremorButton>
+          <AntdButton onClick={onCancel}>Cancel</AntdButton>
+          <TremorButton type="submit">Save Changes</TremorButton>
         </div>
       </div>
     </Form>
-  );
-} 
+  )
+}

ui/litellm-dashboard/src/components/templates/key_info_view.tsx

@@ -1,4 +1,4 @@
-import React, { useEffect, useState } from "react";
+import React, { useEffect, useState } from "react"
 import {
   Card,
   Text,
@@ -13,21 +13,23 @@ import {
   Title,
   Badge,
   TextInput,
-  Select as TremorSelect
-} from "@tremor/react";
-import { ArrowLeftIcon, TrashIcon, RefreshIcon } from "@heroicons/react/outline";
-import { keyDeleteCall, keyUpdateCall } from "./networking";
-import { KeyResponse } from "./key_team_helpers/key_list";
-import { Form, Input, InputNumber, message, Select, Tooltip, Button as AntdButton } from "antd";
-import { KeyEditView } from "./key_edit_view";
-import { RegenerateKeyModal } from "./regenerate_key_modal";
-import { rolesWithWriteAccess } from '../utils/roles';
-import ObjectPermissionsView from "./object_permissions_view";
-import LoggingSettingsView from "./logging_settings_view";
-import { copyToClipboard as utilCopyToClipboard, formatNumberWithCommas } from "@/utils/dataUtils";
-import { extractLoggingSettings, formatMetadataForDisplay } from "./key_info_utils";
-import { CopyIcon, CheckIcon } from "lucide-react";
-import { callback_map, mapInternalToDisplayNames, mapDisplayToInternalNames } from "./callback_info_helpers";
+  Select as TremorSelect,
+} from "@tremor/react"
+import { ArrowLeftIcon, TrashIcon, RefreshIcon } from "@heroicons/react/outline"
+import { keyDeleteCall, keyUpdateCall } from "../networking"
+import { KeyResponse } from "../key_team_helpers/key_list"
+import { Form, Input, InputNumber, Select, Tooltip, Button as AntdButton } from "antd"
+import NotificationManager from "../molecules/notifications_manager"
+import { KeyEditView } from "./key_edit_view"
+import { RegenerateKeyModal } from "../organisms/regenerate_key_modal"
+import { rolesWithWriteAccess } from "../../utils/roles"
+import ObjectPermissionsView from "../object_permissions_view"
+import LoggingSettingsView from "../logging_settings_view"
+import { copyToClipboard as utilCopyToClipboard, formatNumberWithCommas } from "@/utils/dataUtils"
+import { extractLoggingSettings, formatMetadataForDisplay } from "../key_info_utils"
+import { CopyIcon, CheckIcon } from "lucide-react"
+import { callback_map, mapInternalToDisplayNames, mapDisplayToInternalNames } from "../callback_info_helpers"
+import { parseErrorMessage } from "../shared/errorUtils"
 
 interface KeyInfoViewProps {
   keyId: string
@@ -88,97 +90,84 @@ export default function KeyInfoView({
   if (!currentKeyData) {
     return (
       <div className="p-4">
-        <Button 
-          icon={ArrowLeftIcon} 
-          variant="light"
-          onClick={onClose}
-          className="mb-4"
-        >
+        <Button icon={ArrowLeftIcon} variant="light" onClick={onClose} className="mb-4">
           Back to Keys
         </Button>
         <Text>Key not found</Text>
       </div>
-    );
+    )
   }
 
   const handleKeyUpdate = async (formValues: Record<string, any>) => {
     try {
-      if (!accessToken) return;
+      if (!accessToken) return
 
-      const currentKey = formValues.token;
-      formValues.key = currentKey;
+      const currentKey = formValues.token
+      formValues.key = currentKey
 
       // Handle object_permission updates
       if (formValues.vector_stores !== undefined) {
         formValues.object_permission = {
           ...currentKeyData.object_permission,
-          vector_stores: formValues.vector_stores || []
-        };
+          vector_stores: formValues.vector_stores || [],
+        }
         // Remove vector_stores from the top level as it should be in object_permission
-        delete formValues.vector_stores;
+        delete formValues.vector_stores
       }
 
       if (formValues.mcp_servers_and_groups !== undefined) {
-        const { servers, accessGroups } = formValues.mcp_servers_and_groups || { servers: [], accessGroups: [] };
+        const { servers, accessGroups } = formValues.mcp_servers_and_groups || { servers: [], accessGroups: [] }
         formValues.object_permission = {
           ...currentKeyData.object_permission,
           mcp_servers: servers || [],
-          mcp_access_groups: accessGroups || []
-        };
+          mcp_access_groups: accessGroups || [],
+        }
         // Remove mcp_servers_and_groups from the top level as it should be in object_permission
-        delete formValues.mcp_servers_and_groups;
+        delete formValues.mcp_servers_and_groups
       }
 
       // Convert metadata back to an object if it exists and is a string
       if (formValues.metadata && typeof formValues.metadata === "string") {
         try {
-          const parsedMetadata = JSON.parse(formValues.metadata);
+          const parsedMetadata = JSON.parse(formValues.metadata)
           formValues.metadata = {
             ...parsedMetadata,
-            ...(formValues.guardrails?.length > 0
-              ? { guardrails: formValues.guardrails }
-              : {}),
-            ...(formValues.logging_settings
-              ? { logging: formValues.logging_settings }
-              : {}),
+            ...(formValues.guardrails?.length > 0 ? { guardrails: formValues.guardrails } : {}),
+            ...(formValues.logging_settings ? { logging: formValues.logging_settings } : {}),
             ...(formValues.disabled_callbacks?.length > 0
-              ? { 
-                  litellm_disabled_callbacks: mapDisplayToInternalNames(formValues.disabled_callbacks)
+              ? {
+                  litellm_disabled_callbacks: mapDisplayToInternalNames(formValues.disabled_callbacks),
                 }
               : {}),
-          };
+          }
         } catch (error) {
-          console.error("Error parsing metadata JSON:", error);
-          message.error("Invalid metadata JSON");
-          return;
+          console.error("Error parsing metadata JSON:", error)
+          NotificationManager.error("Invalid metadata JSON")
+          return
         }
       } else {
         formValues.metadata = {
           ...(formValues.metadata || {}),
-          ...(formValues.guardrails?.length > 0
-            ? { guardrails: formValues.guardrails }
-            : {}),
-          ...(formValues.logging_settings
-            ? { logging: formValues.logging_settings }
-            : {}),
+          ...(formValues.guardrails?.length > 0 ? { guardrails: formValues.guardrails } : {}),
+          ...(formValues.logging_settings ? { logging: formValues.logging_settings } : {}),
           ...(formValues.disabled_callbacks?.length > 0
-            ? { 
-                litellm_disabled_callbacks: mapDisplayToInternalNames(formValues.disabled_callbacks)
+            ? {
+                litellm_disabled_callbacks: mapDisplayToInternalNames(formValues.disabled_callbacks),
               }
             : {}),
-        };
+        }
       }
 
-      delete formValues.logging_settings;
+      delete formValues.logging_settings
 
       // Convert budget_duration to API format
       if (formValues.budget_duration) {
         const durationMap: Record<string, string> = {
           daily: "24h",
           weekly: "7d",
-          monthly: "30d"
-        };
-        formValues.budget_duration = durationMap[formValues.budget_duration];
+          monthly: "30d",
+        }
+        formValues.budget_duration = durationMap[formValues.budget_duration]
       }
 
       const newKeyValues = await keyUpdateCall(accessToken, formValues)
@@ -189,34 +178,34 @@ export default function KeyInfoView({
       if (onKeyDataUpdate) {
         onKeyDataUpdate(newKeyValues)
       }
-      message.success("Key updated successfully");
-      setIsEditing(false);
+      NotificationManager.success("Key updated successfully")
+      setIsEditing(false)
       // Refresh key data here if needed
     } catch (error) {
-      message.error("Failed to update key");
-      console.error("Error updating key:", error);
+      NotificationManager.fromBackend(parseErrorMessage(error))
+      console.error("Error updating key:", error)
     }
-  };
+  }
 
   const handleDelete = async () => {
     try {
       if (!accessToken) return
       await keyDeleteCall(accessToken as string, currentKeyData.token || currentKeyData.token_id)
-      message.success("Key deleted successfully")
+      NotificationManager.success("Key deleted successfully")
       if (onDelete) {
         onDelete()
       }
-      onClose();
+      onClose()
     } catch (error) {
-      console.error("Error deleting the key:", error);
-      message.error("Failed to delete key");
+      console.error("Error deleting the key:", error)
+      NotificationManager.fromBackend(error)
     }
-  };
+  }
 
   const copyToClipboard = async (text: string, key: string) => {
-    const success = await utilCopyToClipboard(text);
+    const success = await utilCopyToClipboard(text)
     if (success) {
-      setCopiedStates((prev) => ({ ...prev, [key]: true }));
+      setCopiedStates((prev) => ({ ...prev, [key]: true }))
       setTimeout(() => {
         setCopiedStates((prev) => ({ ...prev, [key]: false }))
       }, 2000)
@@ -268,12 +257,7 @@ export default function KeyInfoView({
     <div className="w-full h-screen p-4">
       <div className="flex justify-between items-center mb-6">
         <div>
-          <Button 
-            icon={ArrowLeftIcon} 
-            variant="light"
-            onClick={onClose}
-            className="mb-4"
-          >
+          <Button icon={ArrowLeftIcon} variant="light" onClick={onClose} className="mb-4">
             Back to Keys
           </Button>
           <Title>{currentKeyData.key_alias || "API Key"}</Title>
@@ -319,7 +303,9 @@ export default function KeyInfoView({
         </div>
         {userRole && rolesWithWriteAccess.includes(userRole) && (
           <div className="flex gap-2">
-            <Tooltip title={!premiumUser ? "This is a LiteLLM Enterprise feature, and requires a valid key to use." : ""}>
+            <Tooltip
+              title={!premiumUser ? "This is a LiteLLM Enterprise feature, and requires a valid key to use." : ""}
+            >
               <span className="inline-block">
                 <Button
                   icon={RefreshIcon}
@@ -363,34 +349,26 @@ export default function KeyInfoView({
               <div className="absolute inset-0 bg-gray-500 opacity-75"></div>
             </div>
 
-            <span className="hidden sm:inline-block sm:align-middle sm:h-screen" aria-hidden="true">&#8203;</span>
+            <span className="hidden sm:inline-block sm:align-middle sm:h-screen" aria-hidden="true">
+              &#8203;
+            </span>
 
             <div className="inline-block align-bottom bg-white rounded-lg text-left overflow-hidden shadow-xl transform transition-all sm:my-8 sm:align-middle sm:max-w-lg sm:w-full">
               <div className="bg-white px-4 pt-5 pb-4 sm:p-6 sm:pb-4">
                 <div className="sm:flex sm:items-start">
                   <div className="mt-3 text-center sm:mt-0 sm:ml-4 sm:text-left">
-                    <h3 className="text-lg leading-6 font-medium text-gray-900">
-                      Delete Key
-                    </h3>
+                    <h3 className="text-lg leading-6 font-medium text-gray-900">Delete Key</h3>
                     <div className="mt-2">
-                      <p className="text-sm text-gray-500">
-                        Are you sure you want to delete this key?
-                      </p>
+                      <p className="text-sm text-gray-500">Are you sure you want to delete this key?</p>
                     </div>
                   </div>
                 </div>
               </div>
               <div className="bg-gray-50 px-4 py-3 sm:px-6 sm:flex sm:flex-row-reverse">
-                <Button
-                  onClick={handleDelete}
-                  color="red"
-                  className="ml-2"
-                >
+                <Button onClick={handleDelete} color="red" className="ml-2">
                   Delete
                 </Button>
-                <Button onClick={() => setIsDeleteModalOpen(false)}>
-                  Cancel
-                </Button>
+                <Button onClick={() => setIsDeleteModalOpen(false)}>Cancel</Button>
               </div>
             </div>
           </div>
@@ -492,7 +470,7 @@ export default function KeyInfoView({
                     <Text className="font-medium">Key ID</Text>
                     <Text className="font-mono">{currentKeyData.token_id || currentKeyData.token}</Text>
                   </div>
-                  
+
                   <div>
                     <Text className="font-medium">Key Alias</Text>
                     <Text>{currentKeyData.key_alias || "Not Set"}</Text>
@@ -548,16 +526,16 @@ export default function KeyInfoView({
                         : "Unlimited"}
                     </Text>
                   </div>
-                  
+
                   <div>
                     <Text className="font-medium">Prompts</Text>
                     <Text>
                       {Array.isArray(currentKeyData.metadata?.prompts) && currentKeyData.metadata.prompts.length > 0
                         ? currentKeyData.metadata.prompts.map((prompt, index) => (
-                          <span key={index} className="px-2 mr-2 py-1 bg-blue-100 rounded text-xs">
-                            {prompt}
-                          </span>
-                        ))
+                            <span key={index} className="px-2 mr-2 py-1 bg-blue-100 rounded text-xs">
+                              {prompt}
+                            </span>
+                          ))
                         : "No prompts specified"}
                     </Text>
                   </div>
@@ -632,5 +610,5 @@ export default function KeyInfoView({
         </TabPanels>
       </TabGroup>
     </div>
-  );
-} 
+  )
+}

ui/litellm-dashboard/src/components/templates/view_key_table.tsx

@@ -1,26 +1,21 @@
-"use client";
-import React, { useEffect, useState, useMemo } from "react";
-import {
-  keyDeleteCall,
-  modelAvailableCall,
-  getGuardrailsList,
-  Organization,
-} from "./networking";
-import { add } from "date-fns";
+"use client"
+import React, { useEffect, useState, useMemo } from "react"
+import { keyDeleteCall, modelAvailableCall, getGuardrailsList, Organization } from "../networking"
+import { add } from "date-fns"
 import {
   InformationCircleIcon,
   StatusOnlineIcon,
   TrashIcon,
   PencilAltIcon,
   RefreshIcon,
-} from "@heroicons/react/outline";
+} from "@heroicons/react/outline"
 import {
   keySpendLogsCall,
   PredictedSpendLogsCall,
   keyUpdateCall,
   modelInfoCall,
   regenerateKeyCall,
-} from "./networking";
+} from "../networking"
 import {
   Badge,
   Card,
@@ -44,16 +39,13 @@ import {
   Textarea,
   Select,
   SelectItem,
-} from "@tremor/react";
-import { InfoCircleOutlined } from "@ant-design/icons";
+} from "@tremor/react"
+import { InfoCircleOutlined } from "@ant-design/icons"
 import {
   fetchAvailableModelsForTeamOrKey,
   getModelDisplayName,
-} from "./key_team_helpers/fetch_available_models_team_key";
-import {
-  MultiSelect,
-  MultiSelectItem,
-} from "@tremor/react";
+} from "../key_team_helpers/fetch_available_models_team_key"
+import { MultiSelect, MultiSelectItem } from "@tremor/react"
 import {
   Button as Button2,
   Modal,
@@ -64,28 +56,30 @@ import {
   message,
   Tooltip,
   DatePicker,
-} from "antd";
-import { CopyToClipboard } from "react-copy-to-clipboard";
-import TextArea from "antd/es/input/TextArea";
-import useKeyList from "./key_team_helpers/key_list";
-import { KeyResponse } from "./key_team_helpers/key_list";
-import { AllKeysTable } from "./all_keys_table";
-import { Team } from "./key_team_helpers/key_list";
-import { Setter } from "@/types";
+} from "antd"
+import { CopyToClipboard } from "react-copy-to-clipboard"
+import TextArea from "antd/es/input/TextArea"
+import useKeyList from "../key_team_helpers/key_list"
+import { KeyResponse } from "../key_team_helpers/key_list"
+import { AllKeysTable } from "../all_keys_table"
+import { Team } from "../key_team_helpers/key_list"
+import { Setter } from "@/types"
+
+import NotificationManager from "../molecules/notifications_manager"
 
 interface EditKeyModalProps {
-  visible: boolean;
-  onCancel: () => void;
-  token: any; // Assuming TeamType is a type representing your team object
-  onSubmit: (data: FormData) => void; // Assuming FormData is the type of data to be submitted
+  visible: boolean
+  onCancel: () => void
+  token: any // Assuming TeamType is a type representing your team object
+  onSubmit: (data: FormData) => void // Assuming FormData is the type of data to be submitted
 }
 
 interface ModelLimitModalProps {
-  visible: boolean;
-  onCancel: () => void;
-  token: KeyResponse;
-  onSubmit: (updatedMetadata: any) => void;
-  accessToken: string;
+  visible: boolean
+  onCancel: () => void
+  token: KeyResponse
+  onSubmit: (updatedMetadata: any) => void
+  accessToken: string
 }
 
 // Define the props type
@@ -94,14 +88,14 @@ interface ViewKeyTableProps {
   userRole: string | null
   accessToken: string | null
   selectedTeam: Team | null
-  setSelectedTeam: React.Dispatch<React.SetStateAction<any | null>>;
+  setSelectedTeam: React.Dispatch<React.SetStateAction<any | null>>
   data: KeyResponse[] | null
-  setData: React.Dispatch<React.SetStateAction<any[] | null>>;
-  teams: Team[] | null;
-  premiumUser: boolean;
-  currentOrg: Organization | null;
-  organizations: Organization[] | null;
-  setCurrentOrg: React.Dispatch<React.SetStateAction<Organization | null>>;
+  setData: React.Dispatch<React.SetStateAction<any[] | null>>
+  teams: Team[] | null
+  premiumUser: boolean
+  currentOrg: Organization | null
+  organizations: Organization[] | null
+  setCurrentOrg: React.Dispatch<React.SetStateAction<Organization | null>>
   selectedKeyAlias: string | null
   setSelectedKeyAlias: Setter<string | null>
   createClicked: boolean
@@ -109,36 +103,36 @@ interface ViewKeyTableProps {
 }
 
 interface ItemData {
-  key_alias: string | null;
-  key_name: string;
-  spend: string;
-  max_budget: string | null;
-  models: string[];
-  tpm_limit: string | null;
-  rpm_limit: string | null;
-  token: string;
-  token_id: string | null;
-  id: number;
-  team_id: string;
-  metadata: any;
-  user_id: string | null;
-  expires: any;
-  budget_duration: string | null;
-  budget_reset_at: string | null;
+  key_alias: string | null
+  key_name: string
+  spend: string
+  max_budget: string | null
+  models: string[]
+  tpm_limit: string | null
+  rpm_limit: string | null
+  token: string
+  token_id: string | null
+  id: number
+  team_id: string
+  metadata: any
+  user_id: string | null
+  expires: any
+  budget_duration: string | null
+  budget_reset_at: string | null
   // Add any other properties that exist in the item data
 }
 
 interface ModelLimits {
-  [key: string]: number;  // Index signature allowing string keys
+  [key: string]: number // Index signature allowing string keys
 }
 
 interface CombinedLimit {
-  tpm: number;
-  rpm: number;
+  tpm: number
+  rpm: number
 }
 
 interface CombinedLimits {
-  [key: string]: CombinedLimit;  // Index signature allowing string keys
+  [key: string]: CombinedLimit // Index signature allowing string keys
 }
 
 const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
@@ -159,22 +153,19 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
   createClicked,
   setAccessToken,
 }) => {
-  const [isButtonClicked, setIsButtonClicked] = useState(false);
-  const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
-  const [keyToDelete, setKeyToDelete] = useState<string | null>(null);
-  const [selectedItem, setSelectedItem] = useState<KeyResponse | null>(null);
-  const [spendData, setSpendData] = useState<
-    { day: string; spend: number }[] | null
-  >(null);
-  
-  // NEW: Declare filter states for team and key alias.
-  const [teamFilter, setTeamFilter] = useState<string>(selectedTeam?.team_id || "");
+  const [isButtonClicked, setIsButtonClicked] = useState(false)
+  const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false)
+  const [keyToDelete, setKeyToDelete] = useState<string | null>(null)
+  const [selectedItem, setSelectedItem] = useState<KeyResponse | null>(null)
+  const [spendData, setSpendData] = useState<{ day: string; spend: number }[] | null>(null)
 
+  // NEW: Declare filter states for team and key alias.
+  const [teamFilter, setTeamFilter] = useState<string>(selectedTeam?.team_id || "")
 
   // Keep the team filter in sync with the incoming prop.
   useEffect(() => {
-    setTeamFilter(selectedTeam?.team_id || "");
-  }, [selectedTeam]);
+    setTeamFilter(selectedTeam?.team_id || "")
+  }, [selectedTeam])
 
   // Build a memoized filters object for the backend call.
 
@@ -185,46 +176,45 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
     selectedKeyAlias,
     accessToken: accessToken || "",
     createClicked,
-  });
-
+  })
 
   const handlePageChange = (newPage: number) => {
-    refresh({ page: newPage });
-  };
+    refresh({ page: newPage })
+  }
 
-  const [editModalVisible, setEditModalVisible] = useState(false);
-  const [infoDialogVisible, setInfoDialogVisible] = useState(false);
-  const [selectedToken, setSelectedToken] = useState<KeyResponse | null>(null);
-  const [userModels, setUserModels] = useState<string[]>([]);
-  const initialKnownTeamIDs: Set<string> = new Set();
-  const [modelLimitModalVisible, setModelLimitModalVisible] = useState(false);
-  const [regenerateDialogVisible, setRegenerateDialogVisible] = useState(false);
-  const [regeneratedKey, setRegeneratedKey] = useState<string | null>(null);
-  const [regenerateFormData, setRegenerateFormData] = useState<any>(null);
-  const [regenerateForm] = Form.useForm();
-  const [newExpiryTime, setNewExpiryTime] = useState<string | null>(null);
+  const [editModalVisible, setEditModalVisible] = useState(false)
+  const [infoDialogVisible, setInfoDialogVisible] = useState(false)
+  const [selectedToken, setSelectedToken] = useState<KeyResponse | null>(null)
+  const [userModels, setUserModels] = useState<string[]>([])
+  const initialKnownTeamIDs: Set<string> = new Set()
+  const [modelLimitModalVisible, setModelLimitModalVisible] = useState(false)
+  const [regenerateDialogVisible, setRegenerateDialogVisible] = useState(false)
+  const [regeneratedKey, setRegeneratedKey] = useState<string | null>(null)
+  const [regenerateFormData, setRegenerateFormData] = useState<any>(null)
+  const [regenerateForm] = Form.useForm()
+  const [newExpiryTime, setNewExpiryTime] = useState<string | null>(null)
 
-  const [knownTeamIDs, setKnownTeamIDs] = useState(initialKnownTeamIDs);
-  const [guardrailsList, setGuardrailsList] = useState<string[]>([]);
+  const [knownTeamIDs, setKnownTeamIDs] = useState(initialKnownTeamIDs)
+  const [guardrailsList, setGuardrailsList] = useState<string[]>([])
 
   useEffect(() => {
     const calculateNewExpiryTime = (duration: string | undefined) => {
       if (!duration) {
-        return null;
+        return null
       }
 
       try {
-        const now = new Date();
-        let newExpiry: Date;
+        const now = new Date()
+        let newExpiry: Date
 
         if (duration.endsWith("s")) {
-          newExpiry = add(now, { seconds: parseInt(duration) });
+          newExpiry = add(now, { seconds: parseInt(duration) })
         } else if (duration.endsWith("h")) {
-          newExpiry = add(now, { hours: parseInt(duration) });
+          newExpiry = add(now, { hours: parseInt(duration) })
         } else if (duration.endsWith("d")) {
-          newExpiry = add(now, { days: parseInt(duration) });
+          newExpiry = add(now, { days: parseInt(duration) })
         } else {
-          throw new Error("Invalid duration format");
+          throw new Error("Invalid duration format")
         }
 
         return newExpiry.toLocaleString("en-US", {
@@ -235,143 +225,134 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
           minute: "numeric",
           second: "numeric",
           hour12: true,
-        });
+        })
       } catch (error) {
-        return null;
+        return null
       }
-    };
+    }
 
-    console.log("in calculateNewExpiryTime for selectedToken", selectedToken);
+    console.log("in calculateNewExpiryTime for selectedToken", selectedToken)
 
     // When a new duration is entered
     if (regenerateFormData?.duration) {
-      setNewExpiryTime(calculateNewExpiryTime(regenerateFormData.duration));
+      setNewExpiryTime(calculateNewExpiryTime(regenerateFormData.duration))
     } else {
-      setNewExpiryTime(null);
+      setNewExpiryTime(null)
     }
 
-    console.log("calculateNewExpiryTime:", newExpiryTime);
-  }, [selectedToken, regenerateFormData?.duration]);
+    console.log("calculateNewExpiryTime:", newExpiryTime)
+  }, [selectedToken, regenerateFormData?.duration])
 
   useEffect(() => {
     const fetchUserModels = async () => {
       try {
         if (userID === null || userRole === null || accessToken === null) {
-          return;
+          return
         }
 
-        const models = await fetchAvailableModelsForTeamOrKey(
-          userID,
-          userRole,
-          accessToken
-        );
+        const models = await fetchAvailableModelsForTeamOrKey(userID, userRole, accessToken)
         if (models) {
-          setUserModels(models);
+          setUserModels(models)
         }
       } catch (error) {
-        console.error("Error fetching user models:", error);
+        NotificationManager.error({ description: "Error fetching user models" })
       }
-    };
-
-    fetchUserModels();
-  }, [accessToken, userID, userRole]);
+    }
 
+    fetchUserModels()
+  }, [accessToken, userID, userRole])
 
   useEffect(() => {
     if (teams) {
-      const teamIDSet: Set<string> = new Set();
+      const teamIDSet: Set<string> = new Set()
       teams.forEach((team: any, index: number) => {
-        const team_obj: string = team.team_id;
-        teamIDSet.add(team_obj);
-      });
-      setKnownTeamIDs(teamIDSet);
+        const team_obj: string = team.team_id
+        teamIDSet.add(team_obj)
+      })
+      setKnownTeamIDs(teamIDSet)
     }
-  }, [teams]);
+  }, [teams])
 
   const confirmDelete = async () => {
     if (keyToDelete == null || data == null) {
-      return;
+      return
     }
 
     try {
       if (!accessToken) return
       await keyDeleteCall(accessToken, keyToDelete)
       // Successfully completed the deletion. Update the state to trigger a rerender.
-      const filteredData = data.filter((item) => item.token !== keyToDelete);
-      setData(filteredData);
+      const filteredData = data.filter((item) => item.token !== keyToDelete)
+      setData(filteredData)
     } catch (error) {
-      console.error("Error deleting the key:", error);
-      // Handle any error situations, such as displaying an error message to the user.
+      NotificationManager.error({ description: "Error deleting the key" })
     }
 
     // Close the confirmation modal and reset the keyToDelete
-    setIsDeleteModalOpen(false);
-    setKeyToDelete(null);
-  };
+    setIsDeleteModalOpen(false)
+    setKeyToDelete(null)
+  }
 
   const cancelDelete = () => {
     // Close the confirmation modal and reset the keyToDelete
-    setIsDeleteModalOpen(false);
-    setKeyToDelete(null);
-  };
+    setIsDeleteModalOpen(false)
+    setKeyToDelete(null)
+  }
 
   const handleRegenerateClick = (token: any) => {
-    setSelectedToken(token);
-    setNewExpiryTime(null);
+    setSelectedToken(token)
+    setNewExpiryTime(null)
     regenerateForm.setFieldsValue({
       key_alias: token.key_alias,
       max_budget: token.max_budget,
       tpm_limit: token.tpm_limit,
       rpm_limit: token.rpm_limit,
       duration: token.duration || "",
-    });
-    setRegenerateDialogVisible(true);
-  };
+    })
+    setRegenerateDialogVisible(true)
+  }
 
   const handleRegenerateFormChange = (field: string, value: any) => {
     setRegenerateFormData((prev: any) => ({
       ...prev,
       [field]: value,
-    }));
-  };
+    }))
+  }
 
   const handleRegenerateKey = async () => {
     if (!premiumUser) {
-      message.error(
-        "Regenerate API Key is an Enterprise feature. Please upgrade to use this feature."
-      );
-      return;
+      NotificationManager.warning({
+        description: "Regenerate API Key is an Enterprise feature. Please upgrade to use this feature.",
+      })
+      return
     }
 
     if (selectedToken == null) {
-      return;
+      return
     }
 
     try {
       const formValues = await regenerateForm.validateFields()
-      if (!accessToken) return;
+      if (!accessToken) return
       const response = await regenerateKeyCall(accessToken, selectedToken.token || selectedToken.token_id, formValues)
       setRegeneratedKey(response.key)
 
       // Update the data state with the new key_name
       if (data) {
         const updatedData = data.map((item) =>
-          item.token === selectedToken?.token ?
-            { ...item, key_name: response.key_name, ...formValues }
-          : item
-        );
-        setData(updatedData);
+          item.token === selectedToken?.token ? { ...item, key_name: response.key_name, ...formValues } : item,
+        )
+        setData(updatedData)
       }
 
-      setRegenerateDialogVisible(false);
-      regenerateForm.resetFields();
-      message.success("API Key regenerated successfully");
+      setRegenerateDialogVisible(false)
+      regenerateForm.resetFields()
+      NotificationManager.success({ description: "API Key regenerated successfully" })
     } catch (error) {
-      console.error("Error regenerating key:", error);
-      message.error("Failed to regenerate API Key");
+      console.error("Error regenerating key:", error)
+      NotificationManager.error({ description: "Failed to regenerate API Key" })
     }
-  };
-
+  }
 
   return (
     <div>
@@ -400,18 +381,12 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
       {isDeleteModalOpen && (
         <div className="fixed z-10 inset-0 overflow-y-auto">
           <div className="flex items-end justify-center min-h-screen pt-4 px-4 pb-20 text-center sm:block sm:p-0">
-            <div
-              className="fixed inset-0 transition-opacity"
-              aria-hidden="true"
-            >
+            <div className="fixed inset-0 transition-opacity" aria-hidden="true">
               <div className="absolute inset-0 bg-gray-500 opacity-75"></div>
             </div>
 
             {/* Modal Panel */}
-            <span
-              className="hidden sm:inline-block sm:align-middle sm:h-screen"
-              aria-hidden="true"
-            >
+            <span className="hidden sm:inline-block sm:align-middle sm:h-screen" aria-hidden="true">
               &#8203;
             </span>
 
@@ -420,13 +395,9 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
               <div className="bg-white px-4 pt-5 pb-4 sm:p-6 sm:pb-4">
                 <div className="sm:flex sm:items-start">
                   <div className="mt-3 text-center sm:mt-0 sm:ml-4 sm:text-left">
-                    <h3 className="text-lg leading-6 font-medium text-gray-900">
-                      Delete Key
-                    </h3>
+                    <h3 className="text-lg leading-6 font-medium text-gray-900">Delete Key</h3>
                     <div className="mt-2">
-                      <p className="text-sm text-gray-500">
-                        Are you sure you want to delete this key ?
-                      </p>
+                      <p className="text-sm text-gray-500">Are you sure you want to delete this key ?</p>
                     </div>
                   </div>
                 </div>
@@ -447,36 +418,32 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
         title="Regenerate API Key"
         visible={regenerateDialogVisible}
         onCancel={() => {
-          setRegenerateDialogVisible(false);
-          regenerateForm.resetFields();
+          setRegenerateDialogVisible(false)
+          regenerateForm.resetFields()
         }}
         footer={[
           <Button
             key="cancel"
             onClick={() => {
-              setRegenerateDialogVisible(false);
-              regenerateForm.resetFields();
+              setRegenerateDialogVisible(false)
+              regenerateForm.resetFields()
             }}
             className="mr-2"
           >
             Cancel
           </Button>,
-          <Button
-            key="regenerate"
-            onClick={handleRegenerateKey}
-            disabled={!premiumUser}
-          >
+          <Button key="regenerate" onClick={handleRegenerateKey} disabled={!premiumUser}>
             {premiumUser ? "Regenerate" : "Upgrade to Regenerate"}
           </Button>,
         ]}
       >
-        {premiumUser ?
+        {premiumUser ? (
           <Form
             form={regenerateForm}
             layout="vertical"
             onValuesChange={(changedValues, allValues) => {
               if ("duration" in changedValues) {
-                handleRegenerateFormChange("duration", changedValues.duration);
+                handleRegenerateFormChange("duration", changedValues.duration)
               }
             }}
           >
@@ -484,11 +451,7 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
               <TextInput disabled={true} />
             </Form.Item>
             <Form.Item name="max_budget" label="Max Budget (USD)">
-              <InputNumber
-                step={0.01}
-                precision={2}
-                style={{ width: "100%" }}
-              />
+              <InputNumber step={0.01} precision={2} style={{ width: "100%" }} />
             </Form.Item>
             <Form.Item name="tpm_limit" label="TPM Limit">
               <InputNumber style={{ width: "100%" }} />
@@ -496,39 +459,25 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
             <Form.Item name="rpm_limit" label="RPM Limit">
               <InputNumber style={{ width: "100%" }} />
             </Form.Item>
-            <Form.Item
-              name="duration"
-              label="Expire Key (eg: 30s, 30h, 30d)"
-              className="mt-8"
-            >
+            <Form.Item name="duration" label="Expire Key (eg: 30s, 30h, 30d)" className="mt-8">
               <TextInput placeholder="" />
             </Form.Item>
             <div className="mt-2 text-sm text-gray-500">
               Current expiry:{" "}
-              {selectedToken?.expires != null ?
-                new Date(selectedToken.expires).toLocaleString()
-              : "Never"}
+              {selectedToken?.expires != null ? new Date(selectedToken.expires).toLocaleString() : "Never"}
             </div>
-            {newExpiryTime && (
-              <div className="mt-2 text-sm text-green-600">
-                New expiry: {newExpiryTime}
-              </div>
-            )}
+            {newExpiryTime && <div className="mt-2 text-sm text-green-600">New expiry: {newExpiryTime}</div>}
           </Form>
-        : <div>
-            <p className="mb-2 text-gray-500 italic text-[12px]">
-              Upgrade to use this feature
-            </p>
+        ) : (
+          <div>
+            <p className="mb-2 text-gray-500 italic text-[12px]">Upgrade to use this feature</p>
             <Button variant="primary" className="mb-2">
-              <a
-                href="https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat"
-                target="_blank"
-              >
+              <a href="https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat" target="_blank">
                 Get Free Trial
               </a>
             </Button>
           </div>
-        }
+        )}
       </Modal>
 
       {/* Regenerated Key Display Modal */}
@@ -546,10 +495,9 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
             <Title>Regenerated Key</Title>
             <Col numColSpan={1}>
               <p>
-                Please replace your old key with the new key generated. For
-                security reasons, <b>you will not be able to view it again</b>{" "}
-                through your LiteLLM account. If you lose this secret key, you
-                will need to generate a new one.
+                Please replace your old key with the new key generated. For security reasons,{" "}
+                <b>you will not be able to view it again</b> through your LiteLLM account. If you lose this secret key,
+                you will need to generate a new one.
               </p>
             </Col>
             <Col numColSpan={1}>
@@ -575,13 +523,11 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
                   marginBottom: "10px",
                 }}
               >
-                <pre style={{ wordWrap: "break-word", whiteSpace: "normal" }}>
-                  {regeneratedKey}
-                </pre>
+                <pre style={{ wordWrap: "break-word", whiteSpace: "normal" }}>{regeneratedKey}</pre>
               </div>
               <CopyToClipboard
                 text={regeneratedKey}
-                onCopy={() => message.success("API Key copied to clipboard")}
+                onCopy={() => NotificationManager.success({ description: "API Key copied to clipboard" })}
               >
                 <Button className="mt-3">Copy API Key</Button>
               </CopyToClipboard>
@@ -590,15 +536,15 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
         </Modal>
       )}
     </div>
-  );
-};
+  )
+}
 
 // Update the type declaration to include the new function
 declare global {
   interface Window {
-    refreshKeysList?: () => void;
-    addNewKeyToList?: (newKey: any) => void;
+    refreshKeysList?: () => void
+    addNewKeyToList?: (newKey: any) => void
   }
 }
 
-export default ViewKeyTable;
+export default ViewKeyTable

ui/litellm-dashboard/src/components/top_key_view.tsx

@@ -1,74 +1,67 @@
-import React, { useState } from "react";
-import { BarChart } from "@tremor/react";
-import KeyInfoView from "./key_info_view";
-import { keyInfoV1Call } from "./networking";
-import { transformKeyInfo } from "../components/key_team_helpers/transform_key_info";
-import { DataTable } from "./view_logs/table";
-import { Tooltip } from "antd";
-import { Button } from "@tremor/react";
-import { formatNumberWithCommas } from "../utils/dataUtils";
+import React, { useState } from "react"
+import { BarChart } from "@tremor/react"
+import KeyInfoView from "./templates/key_info_view"
+import { keyInfoV1Call } from "./networking"
+import { transformKeyInfo } from "../components/key_team_helpers/transform_key_info"
+import { DataTable } from "./view_logs/table"
+import { Tooltip } from "antd"
+import { Button } from "@tremor/react"
+import { formatNumberWithCommas } from "../utils/dataUtils"
 
 interface TopKeyViewProps {
-  topKeys: any[];
-  accessToken: string | null;
-  userID: string | null;
-  userRole: string | null;
-  teams: any[] | null;
-  premiumUser: boolean;
+  topKeys: any[]
+  accessToken: string | null
+  userID: string | null
+  userRole: string | null
+  teams: any[] | null
+  premiumUser: boolean
 }
 
-const TopKeyView: React.FC<TopKeyViewProps> = ({ 
-  topKeys, 
-  accessToken, 
-  userID, 
-  userRole,
-  teams,
-  premiumUser
-}) => {
-  const [isModalOpen, setIsModalOpen] = useState<boolean>(false);
-  const [selectedKey, setSelectedKey] = useState<string | null>(null);
-  const [keyData, setKeyData] = useState<any | undefined>(undefined);
-  const [viewMode, setViewMode] = useState<'chart' | 'table'>('table');
+const TopKeyView: React.FC<TopKeyViewProps> = ({ topKeys, accessToken, userID, userRole, teams, premiumUser }) => {
+  const [isModalOpen, setIsModalOpen] = useState<boolean>(false)
+  const [selectedKey, setSelectedKey] = useState<string | null>(null)
+  const [keyData, setKeyData] = useState<any | undefined>(undefined)
+  const [viewMode, setViewMode] = useState<"chart" | "table">("table")
 
   const handleKeyClick = async (item: any) => {
-    if (!accessToken) return;
-    
-    try {
-      const keyInfo = await keyInfoV1Call(accessToken, item.api_key);
-      const transformedKeyData = transformKeyInfo(keyInfo);
+    if (!accessToken) return
 
-      setKeyData(transformedKeyData);
-      setSelectedKey(item.api_key);
-      setIsModalOpen(true);  // Open modal when key is clicked
+    try {
+      const keyInfo = await keyInfoV1Call(accessToken, item.api_key)
+      const transformedKeyData = transformKeyInfo(keyInfo)
+
+      setKeyData(transformedKeyData)
+      setSelectedKey(item.api_key)
+      setIsModalOpen(true) // Open modal when key is clicked
     } catch (error) {
-      console.error("Error fetching key info:", error);
+      console.error("Error fetching key info:", error)
     }
-  };
+  }
 
   const handleClose = () => {
-    setIsModalOpen(false);
-    setSelectedKey(null);
-    setKeyData(undefined);
-  };
+    setIsModalOpen(false)
+    setSelectedKey(null)
+    setKeyData(undefined)
+  }
 
   // Handle clicking outside the modal
   const handleOutsideClick = (e: React.MouseEvent<HTMLDivElement>) => {
     if (e.target === e.currentTarget) {
-      handleClose();
+      handleClose()
     }
-  };
+  }
 
   // Handle escape key
   React.useEffect(() => {
     const handleEscapeKey = (e: KeyboardEvent) => {
-      if (e.key === 'Escape' && isModalOpen) {
-        handleClose();
+      if (e.key === "Escape" && isModalOpen) {
+        handleClose()
       }
-    };
+    }
 
-    document.addEventListener('keydown', handleEscapeKey);
-    return () => document.removeEventListener('keydown', handleEscapeKey);
-  }, [isModalOpen]);
+    document.addEventListener("keydown", handleEscapeKey)
+    return () => document.removeEventListener("keydown", handleEscapeKey)
+  }, [isModalOpen])
 
   // Define columns for the table view
   const columns = [
@@ -78,7 +71,7 @@ const TopKeyView: React.FC<TopKeyViewProps> = ({
       cell: (info: any) => (
         <div className="overflow-hidden">
           <Tooltip title={info.getValue() as string}>
-            <Button 
+            <Button
               size="xs"
               variant="light"
               className="font-mono text-blue-500 bg-blue-50 hover:bg-blue-100 text-xs font-normal px-2 py-0.5 text-left overflow-hidden truncate max-w-[200px]"
@@ -91,42 +84,42 @@ const TopKeyView: React.FC<TopKeyViewProps> = ({
       ),
     },
     {
-        header: "Key Alias",
-        accessorKey: "key_alias",
-        cell: (info: any) => info.getValue() || "-",
-      },
+      header: "Key Alias",
+      accessorKey: "key_alias",
+      cell: (info: any) => info.getValue() || "-",
+    },
     {
       header: "Spend (USD)",
       accessorKey: "spend",
       cell: (info: any) => `$${formatNumberWithCommas(info.getValue(), 2)}`,
     },
-  ];
+  ]
 
-  const processedTopKeys = topKeys.map(k => ({
+  const processedTopKeys = topKeys.map((k) => ({
     ...k,
-    display_key_alias: k.key_alias && k.key_alias.length > 10 ? `${k.key_alias.slice(0, 10)}...` : (k.key_alias || '-'),
-  }));
+    display_key_alias: k.key_alias && k.key_alias.length > 10 ? `${k.key_alias.slice(0, 10)}...` : k.key_alias || "-",
+  }))
 
   return (
     <>
       <div className="mb-4 flex justify-end items-center">
         <div className="flex space-x-2">
-        <button
-            onClick={() => setViewMode('table')}
-            className={`px-3 py-1 text-sm rounded-md ${viewMode === 'table' ? 'bg-blue-100 text-blue-700' : 'bg-gray-100 text-gray-700'}`}
+          <button
+            onClick={() => setViewMode("table")}
+            className={`px-3 py-1 text-sm rounded-md ${viewMode === "table" ? "bg-blue-100 text-blue-700" : "bg-gray-100 text-gray-700"}`}
           >
             Table View
           </button>
           <button
-            onClick={() => setViewMode('chart')}
-            className={`px-3 py-1 text-sm rounded-md ${viewMode === 'chart' ? 'bg-blue-100 text-blue-700' : 'bg-gray-100 text-gray-700'}`}
+            onClick={() => setViewMode("chart")}
+            className={`px-3 py-1 text-sm rounded-md ${viewMode === "chart" ? "bg-blue-100 text-blue-700" : "bg-gray-100 text-gray-700"}`}
           >
             Chart View
           </button>
         </div>
       </div>
 
-      {viewMode === 'chart' ? (
+      {viewMode === "chart" ? (
         <div className="relative">
           <BarChart
             className="mt-4 h-40 cursor-pointer hover:opacity-90"
@@ -139,11 +132,11 @@ const TopKeyView: React.FC<TopKeyViewProps> = ({
             layout="vertical"
             showXAxis={false}
             showLegend={false}
-            valueFormatter={(value) => value ? `$${formatNumberWithCommas(value, 2)}` : "No Key Alias"}
+            valueFormatter={(value) => (value ? `$${formatNumberWithCommas(value, 2)}` : "No Key Alias")}
             onValueChange={(item) => handleKeyClick(item)}
             showTooltip={true}
             customTooltip={(props) => {
-              const item = props.payload?.[0]?.payload;
+              const item = props.payload?.[0]?.payload
               return (
                 <div className="relative z-50 p-3 bg-black/90 shadow-lg rounded-lg text-white max-w-xs">
                   <div className="space-y-1.5">
@@ -161,7 +154,7 @@ const TopKeyView: React.FC<TopKeyViewProps> = ({
                     </div>
                   </div>
                 </div>
-              );
+              )
             }}
           />
         </div>
@@ -177,42 +170,45 @@ const TopKeyView: React.FC<TopKeyViewProps> = ({
         </div>
       )}
 
-      {isModalOpen && selectedKey && keyData && (
-        console.log('Rendering modal with:', { isModalOpen, selectedKey, keyData }),
-        <div 
-          className="fixed inset-0 bg-black bg-opacity-50 flex items-center justify-center z-50"
-          onClick={handleOutsideClick}
-        >
-          <div className="bg-white rounded-lg shadow-xl relative w-11/12 max-w-6xl max-h-[90vh] overflow-y-auto min-h-[750px]">
-            {/* Close button */}
-            <button
-              onClick={handleClose}
-              className="absolute top-4 right-4 text-gray-500 hover:text-gray-700 focus:outline-none"
-              aria-label="Close"
-            >
-              <svg className="w-6 h-6" fill="none" stroke="currentColor" viewBox="0 0 24 24">
-                <path strokeLinecap="round" strokeLinejoin="round" strokeWidth="2" d="M6 18L18 6M6 6l12 12" />
-              </svg>
-            </button>
+      {isModalOpen &&
+        selectedKey &&
+        keyData &&
+        (console.log("Rendering modal with:", { isModalOpen, selectedKey, keyData }),
+        (
+          <div
+            className="fixed inset-0 bg-black bg-opacity-50 flex items-center justify-center z-50"
+            onClick={handleOutsideClick}
+          >
+            <div className="bg-white rounded-lg shadow-xl relative w-11/12 max-w-6xl max-h-[90vh] overflow-y-auto min-h-[750px]">
+              {/* Close button */}
+              <button
+                onClick={handleClose}
+                className="absolute top-4 right-4 text-gray-500 hover:text-gray-700 focus:outline-none"
+                aria-label="Close"
+              >
+                <svg className="w-6 h-6" fill="none" stroke="currentColor" viewBox="0 0 24 24">
+                  <path strokeLinecap="round" strokeLinejoin="round" strokeWidth="2" d="M6 18L18 6M6 6l12 12" />
+                </svg>
+              </button>
 
-            {/* Content */}
-            <div className="p-6 h-full">
-              <KeyInfoView
-                keyId={selectedKey}
-                onClose={handleClose}
-                keyData={keyData}
-                accessToken={accessToken}
-                userID={userID}
-                userRole={userRole}
-                teams={teams}
-                premiumUser={premiumUser}
-              />
+              {/* Content */}
+              <div className="p-6 h-full">
+                <KeyInfoView
+                  keyId={selectedKey}
+                  onClose={handleClose}
+                  keyData={keyData}
+                  accessToken={accessToken}
+                  userID={userID}
+                  userRole={userRole}
+                  teams={teams}
+                  premiumUser={premiumUser}
+                />
+              </div>
             </div>
           </div>
-        </div>
-      )}
+        ))}
     </>
-  );
-};
+  )
+}
 
-export default TopKeyView;
+export default TopKeyView

ui/litellm-dashboard/src/components/user_dashboard.tsx

@@ -13,8 +13,8 @@ import {
 } from "./networking"
 import { fetchTeams } from "./common_components/fetch_teams"
 import { Grid, Col, Card, Text, Title } from "@tremor/react"
-import CreateKey from "./create_key_button"
-import ViewKeyTable from "./view_key_table"
+import CreateKey from "./organisms/create_key_button"
+import ViewKeyTable from "./templates/view_key_table"
 import ViewUserSpend from "./view_user_spend"
 import ViewUserTeam from "./view_user_team"
 import DashboardTeam from "./dashboard_default_team"

ui/litellm-dashboard/src/components/view_users.tsx

@@ -296,7 +296,7 @@ const ViewUserDashboard: React.FC<ViewUserDashboardProps> = ({ accessToken, toke
   )
 
   return (
-    <div className="w-full p-6">
+    <div className="w-full p-6 overflow-hidden">
       <div className="flex items-center justify-between mb-4">
         <div className="flex space-x-3">
           <CreateUser userID={userID} accessToken={accessToken} teams={teams} possibleUIRoles={possibleUIRoles} />
@@ -506,7 +506,7 @@ const ViewUserDashboard: React.FC<ViewUserDashboardProps> = ({ accessToken, toke
                   </div>
                 </div>
               </div>
-
+              <div className="overflow-auto">
               <UserDataTable
                 data={userListQuery.data?.users || []}
                 columns={tableColumns}
@@ -529,6 +529,8 @@ const ViewUserDashboard: React.FC<ViewUserDashboardProps> = ({ accessToken, toke
                 selectedUsers={selectedUsers}
                 onSelectionChange={handleSelectionChange}
               />
+              </div>
+              
             </div>
           </TabPanel>
 

ui/litellm-dashboard/src/components/view_users/columns.tsx

@@ -33,7 +33,7 @@ export const columns = (
     ),
   },
   {
-    header: "User Email",
+    header: "Email",
     accessorKey: "user_email",
     cell: ({ row }) => (
       <span className="text-xs">{row.original.user_email || "-"}</span>
@@ -49,7 +49,7 @@ export const columns = (
     ),
   },
   {
-    header: "User Spend ($ USD)",
+    header: "Spend (USD)",
     accessorKey: "spend",
     cell: ({ row }) => (
       <span className="text-xs">
@@ -58,7 +58,7 @@ export const columns = (
     ),
   },
   {
-    header: "User Max Budget ($ USD)",
+    header: "Budget (USD)",
     accessorKey: "max_budget",
     cell: ({ row }) => (
       <span className="text-xs">