diff --git a/docs/my-website/docs/providers/bedrock.md b/docs/my-website/docs/providers/bedrock.md index 21eb3ee6862..9797d678ebb 100644 --- a/docs/my-website/docs/providers/bedrock.md +++ b/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) | + + + + +```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) +``` + + + + + +**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?" + } + ] + }' +``` + + + + ## Provisioned throughput models To use provisioned throughput Bedrock models pass - `model=bedrock/`, 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']` | diff --git a/docs/my-website/docs/providers/sambanova.md b/docs/my-website/docs/providers/sambanova.md index 290b64a1f09..f7be5d3ce77 100644 --- a/docs/my-website/docs/providers/sambanova.md +++ b/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) +``` diff --git a/litellm/__init__.py b/litellm/__init__.py index 7c4a9211455..06d71ff4328 100644 --- a/litellm/__init__.py +++ b/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 diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index 6f2299c777f..5fcd2ddb70a 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/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) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 3845dc5a46e..68b5341e954 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/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 diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index 3961ee2b9e9..9172a05e385 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/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. diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 900fad3d043..cd351ca16a7 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/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( diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index e961433b52e..84762e0b99a 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/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 diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 53de6711cad..e53829d3329 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/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 \ No newline at end of file diff --git a/litellm/llms/gemini/count_tokens/handler.py b/litellm/llms/gemini/count_tokens/handler.py new file mode 100644 index 00000000000..1f13b0c3144 --- /dev/null +++ b/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 + diff --git a/litellm/llms/sambanova/common_utils.py b/litellm/llms/sambanova/common_utils.py new file mode 100644 index 00000000000..b622f705845 --- /dev/null +++ b/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) diff --git a/litellm/llms/sambanova/embedding/handler.py b/litellm/llms/sambanova/embedding/handler.py new file mode 100644 index 00000000000..c3629e4d75f --- /dev/null +++ b/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` +""" diff --git a/litellm/llms/sambanova/embedding/transformation.py b/litellm/llms/sambanova/embedding/transformation.py new file mode 100644 index 00000000000..eca44c7c039 --- /dev/null +++ b/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 + ) diff --git a/litellm/main.py b/litellm/main.py index ffe58503758..339d9e14406 100644 --- a/litellm/main.py +++ b/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, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 28dec7cce90..1001aca9c09 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/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 diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6a7006ea8d6..cf8b3d147f0 100644 --- a/litellm/proxy/_types.py +++ b/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): diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 8965ae7715d..a10a39a6a57 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/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 diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 226fcf05aac..373232e22d2 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/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=[], + ) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 6e9c9252572..938c0fc49eb 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/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 \ No newline at end of file + model: gemini/* diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 939084c38d2..78e9ae24832 100644 --- a/litellm/proxy/proxy_server.py +++ b/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 diff --git a/litellm/types/google_genai/main.py b/litellm/types/google_genai/main.py index 96abc5d2ae8..b875495bab0 100644 --- a/litellm/types/google_genai/main.py +++ b/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] diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 1c3eb49fb0f..2931770cd6e 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/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 diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ec56edf92a3..7ae92185e72 100644 --- a/litellm/types/utils.py +++ b/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' diff --git a/litellm/utils.py b/litellm/utils.py index 54628c05faa..908844d06fe 100644 --- a/litellm/utils.py +++ b/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 diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 28dec7cce90..1001aca9c09 100644 --- a/model_prices_and_context_window.json +++ b/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 diff --git a/tests/llm_translation/test_bedrock_gpt_oss.py b/tests/llm_translation/test_bedrock_gpt_oss.py new file mode 100644 index 00000000000..61bce04e2d0 --- /dev/null +++ b/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 diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index eefe5059b0b..e771c9330f4 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/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 diff --git a/tests/test_litellm/llms/gemini/test_gemini_common_utils.py b/tests/test_litellm/llms/gemini/test_gemini_common_utils.py index 6aca755506f..34472b3856d 100644 --- a/tests/test_litellm/llms/gemini/test_gemini_common_utils.py +++ b/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 + ) diff --git a/tests/test_litellm/llms/sambanova/tests_sambanova_embedding_transformation.py b/tests/test_litellm/llms/sambanova/tests_sambanova_embedding_transformation.py new file mode 100644 index 00000000000..8d445807fdc --- /dev/null +++ b/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" diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 7106a6f884f..22f0e29d11f 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/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 (
-
- 🚅 LiteLLM -
+
🚅 LiteLLM
Loading...
- ); + ) } export default function CreateKeyPage() { - const [userRole, setUserRole] = useState(""); - const [premiumUser, setPremiumUser] = useState(false); - const [disabledPersonalKeyCreation, setDisabledPersonalKeyCreation] = - useState(false); - const [userEmail, setUserEmail] = useState(null); - const [teams, setTeams] = useState(null); - const [keys, setKeys] = useState([]); - const [organizations, setOrganizations] = useState([]); - const [userModels, setUserModels] = useState([]); + const [userRole, setUserRole] = useState("") + const [premiumUser, setPremiumUser] = useState(false) + const [disabledPersonalKeyCreation, setDisabledPersonalKeyCreation] = useState(false) + const [userEmail, setUserEmail] = useState(null) + const [teams, setTeams] = useState(null) + const [keys, setKeys] = useState([]) + const [organizations, setOrganizations] = useState([]) + const [userModels, setUserModels] = useState([]) const [proxySettings, setProxySettings] = useState({ PROXY_BASE_URL: "", PROXY_LOGOUT_URL: "", - }); + }) - const [showSSOBanner, setShowSSOBanner] = useState(true); - const searchParams = useSearchParams()!; - const [modelData, setModelData] = useState({ data: [] }); - const [token, setToken] = useState(null); - const [createClicked, setCreateClicked] = useState(false); - const [authLoading, setAuthLoading] = useState(true); - const [userID, setUserID] = useState(null); + const [showSSOBanner, setShowSSOBanner] = useState(true) + const searchParams = useSearchParams()! + const [modelData, setModelData] = useState({ data: [] }) + const [token, setToken] = useState(null) + const [createClicked, setCreateClicked] = useState(false) + const [authLoading, setAuthLoading] = useState(true) + const [userID, setUserID] = useState(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(null); + const [accessToken, setAccessToken] = useState(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 ; + return } return ( @@ -237,226 +223,199 @@ export default function CreateKeyPage() { {invitation_id ? ( - ) : ( -
- -
-
- -
+ ) : ( +
+ +
+
+ +
- {page == "api-keys" ? ( - - ) : page == "models" ? ( - - ) : page == "llm-playground" ? ( - - ) : page == "users" ? ( - - ) : page == "teams" ? ( - - ) : page == "organizations" ? ( - - ) : page == "admin-panel" ? ( - - ) : page == "api_ref" ? ( - - ) : page == "settings" ? ( - - ) : page == "budgets" ? ( - - ) : page == "guardrails" ? ( - - ) : page == "prompts" ? ( - - ) : page == "transform-request" ? ( - - ) : page == "general-settings" ? ( - - ) : page == "ui-theme" ? ( - - ) : page == "model-hub-table" ? ( - - ) : page == "caching" ? ( - - ) : page == "pass-through-settings" ? ( - - ) : page == "logs" ? ( - - ) : page == "mcp-servers" ? ( - - ) : page == "tag-management" ? ( - - ) : page == "vector-stores" ? ( - - ) : page == "new_usage" ? ( - - ) : ( - - )} + {page == "api-keys" ? ( + + ) : page == "models" ? ( + + ) : page == "llm-playground" ? ( + + ) : page == "users" ? ( + + ) : page == "teams" ? ( + + ) : page == "organizations" ? ( + + ) : page == "admin-panel" ? ( + + ) : page == "api_ref" ? ( + + ) : page == "settings" ? ( + + ) : page == "budgets" ? ( + + ) : page == "guardrails" ? ( + + ) : page == "prompts" ? ( + + ) : page == "transform-request" ? ( + + ) : page == "general-settings" ? ( + + ) : page == "ui-theme" ? ( + + ) : page == "model-hub-table" ? ( + + ) : page == "caching" ? ( + + ) : page == "pass-through-settings" ? ( + + ) : page == "logs" ? ( + + ) : page == "mcp-servers" ? ( + + ) : page == "tag-management" ? ( + + ) : page == "vector-stores" ? ( + + ) : page == "new_usage" ? ( + + ) : ( + + )} +
-
- )} + )} - ); + ) } diff --git a/ui/litellm-dashboard/src/components/activity_metrics.tsx b/ui/litellm-dashboard/src/components/activity_metrics.tsx index 680c6004195..773822a7c8f 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.tsx +++ b/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; } -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 ( -
-

{label}

- {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 ( -
-
- -

- {formatCategoryName(dataKey)} -

-
-

- {formattedValue} -

-
- ); - })} -
- ); - } - 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 ( -
- {categories.map((category, idx) => { - const colorName = colors[idx] as keyof typeof colorNameToHex; - const hexColor = colorNameToHex[colorName] || colors[idx]; - return ( -
- -

- {formatCategoryName(category)} -

-
- ); - })} -
- ); -}; - const ModelSection = ({ modelName, metrics, @@ -215,14 +88,7 @@ const ModelSection = ({
Total Tokens - +
Requests per day - +
`$${formatNumberWithCommas(value, 2)}`} /> @@ -487,26 +348,14 @@ export const ActivityMetrics: React.FC = ({ /> -
- Total Requests Over Time - -
+ Total Requests Over Time 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; diff --git a/ui/litellm-dashboard/src/components/all_keys_table.tsx b/ui/litellm-dashboard/src/components/all_keys_table.tsx index aa2b92a585f..34116d3f4b1 100644 --- a/ui/litellm-dashboard/src/components/all_keys_table.tsx +++ b/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; - accessToken: string | null; - userID: string | null; - userRole: string | null; - organizations: Organization[] | null; - setCurrentOrg: React.Dispatch>; - 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 + accessToken: string | null + userID: string | null + userRole: string | null + organizations: Organization[] | null + setCurrentOrg: React.Dispatch> + 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 ( -
-
- Where Team is - -
+ const handleTeamChange = (value: string) => { + const team = teams?.find((t) => t.team_id === value) + setSelectedTeam(team || null) + } + + return ( +
+
+ Where Team is +
- ); - }; - +
+ ) +} /** * 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(null); - const [userList, setUserList] = useState([]); + const [selectedKeyId, setSelectedKeyId] = useState(null) + const [userList, setUserList] = useState([]) const [sorting, setSorting] = React.useState(() => { 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>({}); + return [ + { + id: "created_at", + desc: true, + }, + ] + }) + const [expandedAccordions, setExpandedAccordions] = useState>({}) // 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[] = [ { @@ -214,10 +193,7 @@ export function AllKeysTable({ header: () => null, cell: ({ row }) => row.getCanExpand() ? ( - ) : null, @@ -229,7 +205,7 @@ export function AllKeysTable({ cell: (info) => (
-
) : null}
- ); + ) }, }, { id: "rate_limits", header: "Rate Limits", cell: ({ row }) => { - const key = row.original; + const key = row.original return (
TPM: {key.tpm_limit !== null ? key.tpm_limit : "Unlimited"}
RPM: {key.rpm_limit !== null ? key.rpm_limit : "Unlimited"}
- ); + ) }, }, - ]; + ] 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 (
{selectedKeyId ? ( - 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({ ) : (
- +
- 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 - +
Page {isLoading ? "..." : pagination.currentPage} of {isLoading ? "..." : pagination.totalPages} - + - + , - ] : [ - , - , - ]} + footer={ + regeneratedKey + ? [ + , + ] + : [ + , + , + ] + } > {regeneratedKey ? ( Regenerated Key

- Please replace your old key with the new key generated. For - security reasons, you will not be able to view it again{" "} - 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,{" "} + you will not be able to view it again through your LiteLLM account. If you lose this secret key, + you will need to generate a new one.

Key Alias:
- 
-                {selectedToken?.key_alias || "No alias set"}
-
+ 
{selectedToken?.key_alias || "No alias set"}
New API Key: 
{regeneratedKey}
- message.success("API Key copied to clipboard")} - > + NotificationManager.success("API Key copied to clipboard")}> @@ -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({
Current expiry: {selectedToken?.expires ? new Date(selectedToken.expires).toLocaleString() : "Never"}
- {newExpiryTime && ( -
- New expiry: {newExpiryTime} -
- )} + {newExpiryTime &&
New expiry: {newExpiryTime}
} )} - ); -} \ No newline at end of file + ) +} diff --git a/ui/litellm-dashboard/src/components/tag_management/tag_info.tsx b/ui/litellm-dashboard/src/components/tag_management/tag_info.tsx index ac8f3672581..a727ff72695 100644 --- a/ui/litellm-dashboard/src/components/tag_management/tag_info.tsx +++ b/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 = ({ - tagId, - onClose, - accessToken, - is_admin, - editTag, -}) => { - const [form] = Form.useForm(); - const [tagDetails, setTagDetails] = useState(null); - const [isEditing, setIsEditing] = useState(editTag); - const [userModels, setUserModels] = useState([]); +const TagInfoView: React.FC = ({ tagId, onClose, accessToken, is_admin, editTag }) => { + const [form] = Form.useForm() + const [tagDetails, setTagDetails] = useState(null) + const [isEditing, setIsEditing] = useState(editTag) + const [userModels, setUserModels] = useState([]) 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
Loading...
; + return
Loading...
} return (
- + Tag Name: {tagDetails.name} {tagDetails.description || "No description"}
- {is_admin && !isEditing && ( - - )} + {is_admin && !isEditing && }
{isEditing ? ( -
- + + - + - Allowed LLMs{' '} + Allowed LLMs{" "} - + } name="models" > - + {userModels.map((modelId) => ( {getModelDisplayName(modelId)} @@ -179,9 +146,7 @@ const TagInfoView: React.FC = ({ ) : ( tagDetails.models.map((modelId) => ( - - {tagDetails.model_info?.[modelId] || modelId} - + {tagDetails.model_info?.[modelId] || modelId} )) )} @@ -200,7 +165,7 @@ const TagInfoView: React.FC = ({
)}
- ); -}; + ) +} -export default TagInfoView; \ No newline at end of file +export default TagInfoView diff --git a/ui/litellm-dashboard/src/components/team/member_permissions.tsx b/ui/litellm-dashboard/src/components/team/member_permissions.tsx index f3ee9f19c38..30dc3176753 100644 --- a/ui/litellm-dashboard/src/components/team/member_permissions.tsx +++ b/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 = ({ - teamId, - accessToken, - canEditTeam, -}) => { - const [permissions, setPermissions] = useState([]); - const [selectedPermissions, setSelectedPermissions] = useState([]); - const [loading, setLoading] = useState(true); - const [saving, setSaving] = useState(false); - const [hasChanges, setHasChanges] = useState(false); +const MemberPermissions: React.FC = ({ teamId, accessToken, canEditTeam }) => { + const [permissions, setPermissions] = useState([]) + const [selectedPermissions, setSelectedPermissions] = useState([]) + 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
Loading permissions...
; } - const hasPermissions = permissions.length > 0; + const handleReset = () => { + fetchPermissions() + } + + if (loading) { + return
Loading permissions...
+ } + + const hasPermissions = permissions.length > 0 return ( @@ -97,79 +93,66 @@ const MemberPermissions: React.FC = ({ - + Save Changes
)}
- - Control what team members can do when they are not team admins. - + Control what team members can do when they are not team admins. {hasPermissions ? ( - - - - Method - Endpoint - Description - Access - - - - {permissions.map((permission) => { - const permInfo = getPermissionInfo(permission); - return ( - - - - {permInfo.method} - - - - - {permInfo.endpoint} - - - - {permInfo.description} - - - - handlePermissionChange(permission, e.target.checked) - } - disabled={!canEditTeam} - /> - - - ); - })} - -
+
+ + + + Method + Endpoint + Description + + Allow Access + + + + + {permissions.map((permission) => { + const permInfo = getPermissionInfo(permission) + return ( + + + + {permInfo.method} + + + + {permInfo.endpoint} + + {permInfo.description} + + handlePermissionChange(permission, e.target.checked)} + disabled={!canEditTeam} + /> + + + ) + })} + +
+
) : (
)} - ); -}; + ) +} -export default MemberPermissions; +export default MemberPermissions diff --git a/ui/litellm-dashboard/src/components/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx similarity index 58% rename from ui/litellm-dashboard/src/components/key_edit_view.tsx rename to ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 69102a1591d..9462a9e1057 100644 --- a/ui/litellm-dashboard/src/components/key_edit_view.tsx +++ b/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; - teams?: any[] | null; - accessToken: string | null; - userID: string | null; - userRole: string | null; - premiumUser?: boolean; + keyData: KeyResponse + onCancel: () => void + onSubmit: (values: any) => Promise + 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([]); - const [promptsList, setPromptsList] = useState([]); - const team = teams?.find(team => team.team_id === keyData.team_id); - const [availableModels, setAvailableModels] = useState([]); - const [mcpAccessGroups, setMcpAccessGroups] = useState([]); - const [mcpAccessGroupsLoaded, setMcpAccessGroupsLoaded] = useState(false); + const [form] = Form.useForm() + const [userModels, setUserModels] = useState([]) + const [promptsList, setPromptsList] = useState([]) + const team = teams?.find((team) => team.team_id === keyData.team_id) + const [availableModels, setAvailableModels] = useState([]) + const [mcpAccessGroups, setMcpAccessGroups] = useState([]) + const [mcpAccessGroupsLoaded, setMcpAccessGroupsLoaded] = useState(false) const [disabledCallbacks, setDisabledCallbacks] = useState( - 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 = { "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 ( - + - {/* Only show All Team Models if team has models */} - {availableModels.length > 0 && ( - All Team Models - )} + {availableModels.length > 0 && All Team Models} {/* Show available team models */} - {availableModels.map(model => ( + {availableModels.map((model) => ( {model} @@ -184,7 +167,7 @@ export function KeyEditView({ - + @@ -196,31 +179,27 @@ export function KeyEditView({ - + - + - - + + - + - + - - + 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 }))} /> 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.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)" /> - {/* Only show All Team Models if team has models */} - {teams?.map(team => ( + {teams?.map((team) => ( {`${team.team_alias} (${team.team_id})`} @@ -290,25 +263,23 @@ export function KeyEditView({ 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) }} /> - - {/* Hidden form field for token */}
- +
Key Alias {currentKeyData.key_alias || "Not Set"} @@ -548,16 +526,16 @@ export default function KeyInfoView({ : "Unlimited"}
- +
Prompts {Array.isArray(currentKeyData.metadata?.prompts) && currentKeyData.metadata.prompts.length > 0 ? currentKeyData.metadata.prompts.map((prompt, index) => ( - - {prompt} - - )) + + {prompt} + + )) : "No prompts specified"}
@@ -632,5 +610,5 @@ export default function KeyInfoView({
- ); -} \ No newline at end of file + ) +} diff --git a/ui/litellm-dashboard/src/components/view_key_table.tsx b/ui/litellm-dashboard/src/components/templates/view_key_table.tsx similarity index 62% rename from ui/litellm-dashboard/src/components/view_key_table.tsx rename to ui/litellm-dashboard/src/components/templates/view_key_table.tsx index 094ae80ae65..ca4260d4f25 100644 --- a/ui/litellm-dashboard/src/components/view_key_table.tsx +++ b/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>; + setSelectedTeam: React.Dispatch> data: KeyResponse[] | null - setData: React.Dispatch>; - teams: Team[] | null; - premiumUser: boolean; - currentOrg: Organization | null; - organizations: Organization[] | null; - setCurrentOrg: React.Dispatch>; + setData: React.Dispatch> + teams: Team[] | null + premiumUser: boolean + currentOrg: Organization | null + organizations: Organization[] | null + setCurrentOrg: React.Dispatch> selectedKeyAlias: string | null setSelectedKeyAlias: Setter 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 = ({ @@ -159,22 +153,19 @@ const ViewKeyTable: React.FC = ({ createClicked, setAccessToken, }) => { - const [isButtonClicked, setIsButtonClicked] = useState(false); - const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); - const [keyToDelete, setKeyToDelete] = useState(null); - const [selectedItem, setSelectedItem] = useState(null); - const [spendData, setSpendData] = useState< - { day: string; spend: number }[] | null - >(null); - - // NEW: Declare filter states for team and key alias. - const [teamFilter, setTeamFilter] = useState(selectedTeam?.team_id || ""); + const [isButtonClicked, setIsButtonClicked] = useState(false) + const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false) + const [keyToDelete, setKeyToDelete] = useState(null) + const [selectedItem, setSelectedItem] = useState(null) + const [spendData, setSpendData] = useState<{ day: string; spend: number }[] | null>(null) + // NEW: Declare filter states for team and key alias. + const [teamFilter, setTeamFilter] = useState(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 = ({ 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(null); - const [userModels, setUserModels] = useState([]); - const initialKnownTeamIDs: Set = new Set(); - const [modelLimitModalVisible, setModelLimitModalVisible] = useState(false); - const [regenerateDialogVisible, setRegenerateDialogVisible] = useState(false); - const [regeneratedKey, setRegeneratedKey] = useState(null); - const [regenerateFormData, setRegenerateFormData] = useState(null); - const [regenerateForm] = Form.useForm(); - const [newExpiryTime, setNewExpiryTime] = useState(null); + const [editModalVisible, setEditModalVisible] = useState(false) + const [infoDialogVisible, setInfoDialogVisible] = useState(false) + const [selectedToken, setSelectedToken] = useState(null) + const [userModels, setUserModels] = useState([]) + const initialKnownTeamIDs: Set = new Set() + const [modelLimitModalVisible, setModelLimitModalVisible] = useState(false) + const [regenerateDialogVisible, setRegenerateDialogVisible] = useState(false) + const [regeneratedKey, setRegeneratedKey] = useState(null) + const [regenerateFormData, setRegenerateFormData] = useState(null) + const [regenerateForm] = Form.useForm() + const [newExpiryTime, setNewExpiryTime] = useState(null) - const [knownTeamIDs, setKnownTeamIDs] = useState(initialKnownTeamIDs); - const [guardrailsList, setGuardrailsList] = useState([]); + const [knownTeamIDs, setKnownTeamIDs] = useState(initialKnownTeamIDs) + const [guardrailsList, setGuardrailsList] = useState([]) 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 = ({ 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 = new Set(); + const teamIDSet: Set = 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 (
@@ -400,18 +381,12 @@ const ViewKeyTable: React.FC = ({ {isDeleteModalOpen && (
- - {viewMode === 'chart' ? ( + {viewMode === "chart" ? (
= ({ 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 (
@@ -161,7 +154,7 @@ const TopKeyView: React.FC = ({
- ); + ) }} />
@@ -177,42 +170,45 @@ const TopKeyView: React.FC = ({
)} - {isModalOpen && selectedKey && keyData && ( - console.log('Rendering modal with:', { isModalOpen, selectedKey, keyData }), -
-
- {/* Close button */} - + {isModalOpen && + selectedKey && + keyData && + (console.log("Rendering modal with:", { isModalOpen, selectedKey, keyData }), + ( +
+
+ {/* Close button */} + - {/* Content */} -
- + {/* Content */} +
+ +
-
- )} + ))} - ); -}; + ) +} -export default TopKeyView; +export default TopKeyView diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index e1b9b2cd811..d78978c495e 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/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" diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx index d2041b389fd..7bf49feb7e9 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx @@ -1,60 +1,52 @@ -import moment from "moment"; -import { useQuery } from "@tanstack/react-query"; -import { useState, useRef, useEffect, useCallback } from "react"; -import { useQueryClient } from "@tanstack/react-query"; +import moment from "moment" +import { useQuery } from "@tanstack/react-query" +import { useState, useRef, useEffect, useCallback } from "react" +import { useQueryClient } from "@tanstack/react-query" -import { uiSpendLogsCall, keyInfoV1Call, sessionSpendLogsCall, keyListCall, allEndUsersCall } from "../networking"; -import { DataTable } from "./table"; -import { columns, LogEntry } from "./columns"; -import { Row } from "@tanstack/react-table"; -import { prefetchLogDetails } from "./prefetch"; -import { RequestResponsePanel } from './RequestResponsePanel'; -import { ErrorViewer } from './ErrorViewer'; -import { internalUserRoles } from "../../utils/roles"; -import { ConfigInfoMessage } from './ConfigInfoMessage'; -import { Tooltip } from "antd"; -import { KeyResponse, Team } from "../key_team_helpers/key_list"; -import KeyInfoView from "../key_info_view"; -import { SessionView } from './SessionView'; -import { VectorStoreViewer } from './VectorStoreViewer'; -import { GuardrailViewer } from './GuardrailViewer'; -import FilterComponent from "../common_components/filter"; -import { FilterOption } from "../common_components/filter"; -import { useLogFilterLogic } from "./log_filter_logic"; -import { fetchAllKeyAliases } from "../key_team_helpers/filter_helpers"; -import { - Tab, - TabGroup, - TabList, - TabPanels, - TabPanel, - Text, - Switch, -} from "@tremor/react"; -import AuditLogs from "./audit_logs"; -import { getTimeRangeDisplay } from "./logs_utils"; -import { formatNumberWithCommas } from "@/utils/dataUtils"; +import { uiSpendLogsCall, keyInfoV1Call, sessionSpendLogsCall, keyListCall, allEndUsersCall } from "../networking" +import { DataTable } from "./table" +import { columns, LogEntry } from "./columns" +import { Row } from "@tanstack/react-table" +import { prefetchLogDetails } from "./prefetch" +import { RequestResponsePanel } from "./RequestResponsePanel" +import { ErrorViewer } from "./ErrorViewer" +import { internalUserRoles } from "../../utils/roles" +import { ConfigInfoMessage } from "./ConfigInfoMessage" +import { Tooltip } from "antd" +import { KeyResponse, Team } from "../key_team_helpers/key_list" +import KeyInfoView from "../templates/key_info_view" +import { SessionView } from "./SessionView" +import { VectorStoreViewer } from "./VectorStoreViewer" +import { GuardrailViewer } from "./GuardrailViewer" +import FilterComponent from "../common_components/filter" +import { FilterOption } from "../common_components/filter" +import { useLogFilterLogic } from "./log_filter_logic" +import { fetchAllKeyAliases } from "../key_team_helpers/filter_helpers" +import { Tab, TabGroup, TabList, TabPanels, TabPanel, Text, Switch } from "@tremor/react" +import AuditLogs from "./audit_logs" +import { getTimeRangeDisplay } from "./logs_utils" +import { formatNumberWithCommas } from "@/utils/dataUtils" interface SpendLogsTableProps { - accessToken: string | null; - token: string | null; - userRole: string | null; - userID: string | null; - allTeams: Team[]; - premiumUser: boolean; + accessToken: string | null + token: string | null + userRole: string | null + userID: string | null + allTeams: Team[] + premiumUser: boolean } export interface PaginatedResponse { - data: LogEntry[]; - total: number; - page: number; - page_size: number; - total_pages: number; + data: LogEntry[] + total: number + page: number + page_size: number + total_pages: number } interface PrefetchedLog { - messages: any[]; - response: any; + messages: any[] + response: any } export default function SpendLogsTable({ @@ -65,122 +57,100 @@ export default function SpendLogsTable({ allTeams, premiumUser, }: SpendLogsTableProps) { - const [searchTerm, setSearchTerm] = useState(""); - const [showFilters, setShowFilters] = useState(false); - const [showColumnDropdown, setShowColumnDropdown] = useState(false); - const [currentPage, setCurrentPage] = useState(1); - const [pageSize] = useState(50); - const dropdownRef = useRef(null); - const filtersRef = useRef(null); - const quickSelectRef = useRef(null); + const [searchTerm, setSearchTerm] = useState("") + const [showFilters, setShowFilters] = useState(false) + const [showColumnDropdown, setShowColumnDropdown] = useState(false) + const [currentPage, setCurrentPage] = useState(1) + const [pageSize] = useState(50) + const dropdownRef = useRef(null) + const filtersRef = useRef(null) + const quickSelectRef = useRef(null) // New state variables for Start and End Time - const [startTime, setStartTime] = useState( - moment().subtract(24, "hours").format("YYYY-MM-DDTHH:mm") - ); - const [endTime, setEndTime] = useState( - moment().format("YYYY-MM-DDTHH:mm") - ); + const [startTime, setStartTime] = useState(moment().subtract(24, "hours").format("YYYY-MM-DDTHH:mm")) + const [endTime, setEndTime] = useState(moment().format("YYYY-MM-DDTHH:mm")) - const [isCustomDate, setIsCustomDate] = useState(false); - const [quickSelectOpen, setQuickSelectOpen] = useState(false); - const [tempTeamId, setTempTeamId] = useState(""); - const [tempKeyHash, setTempKeyHash] = useState(""); - const [selectedTeamId, setSelectedTeamId] = useState(""); - const [selectedKeyHash, setSelectedKeyHash] = useState(""); - const [selectedModel, setSelectedModel] = useState(""); - const [selectedKeyInfo, setSelectedKeyInfo] = useState(null); - const [selectedKeyIdInfoView, setSelectedKeyIdInfoView] = useState(null); - const [selectedStatus, setSelectedStatus] = useState(""); - const [filterByCurrentUser, setFilterByCurrentUser] = useState( - userRole && internalUserRoles.includes(userRole) - ); - const [activeTab, setActiveTab] = useState("request logs"); + const [isCustomDate, setIsCustomDate] = useState(false) + const [quickSelectOpen, setQuickSelectOpen] = useState(false) + const [tempTeamId, setTempTeamId] = useState("") + const [tempKeyHash, setTempKeyHash] = useState("") + const [selectedTeamId, setSelectedTeamId] = useState("") + const [selectedKeyHash, setSelectedKeyHash] = useState("") + const [selectedModel, setSelectedModel] = useState("") + const [selectedKeyInfo, setSelectedKeyInfo] = useState(null) + const [selectedKeyIdInfoView, setSelectedKeyIdInfoView] = useState(null) + const [selectedStatus, setSelectedStatus] = useState("") + const [filterByCurrentUser, setFilterByCurrentUser] = useState(userRole && internalUserRoles.includes(userRole)) + const [activeTab, setActiveTab] = useState("request logs") - const [expandedRequestId, setExpandedRequestId] = useState(null); - const [selectedSessionId, setSelectedSessionId] = useState(null); + const [expandedRequestId, setExpandedRequestId] = useState(null) + const [selectedSessionId, setSelectedSessionId] = useState(null) - const queryClient = useQueryClient(); + const queryClient = useQueryClient() const [isLiveTail, setIsLiveTail] = useState(() => { - const storedValue = sessionStorage.getItem("isLiveTail"); + const storedValue = sessionStorage.getItem("isLiveTail") // default to true if nothing is stored - return storedValue !== null ? JSON.parse(storedValue) : true; - }); + return storedValue !== null ? JSON.parse(storedValue) : true + }) useEffect(() => { - sessionStorage.setItem("isLiveTail", JSON.stringify(isLiveTail)); - }, [isLiveTail]); + sessionStorage.setItem("isLiveTail", JSON.stringify(isLiveTail)) + }, [isLiveTail]) - const [selectedTimeInterval, setSelectedTimeInterval] = useState<{ value: number; unit: string }>({ - value: 24, - unit: "hours" - }); + const [selectedTimeInterval, setSelectedTimeInterval] = useState<{ value: number; unit: string }>({ + value: 24, + unit: "hours", + }) useEffect(() => { const fetchKeyInfo = async () => { if (selectedKeyIdInfoView && accessToken) { - const keyData = await keyInfoV1Call(accessToken, selectedKeyIdInfoView); + const keyData = await keyInfoV1Call(accessToken, selectedKeyIdInfoView) const keyResponse: KeyResponse = { ...keyData["info"], - "token": selectedKeyIdInfoView, - "api_key": selectedKeyIdInfoView, - }; - setSelectedKeyInfo(keyResponse); + token: selectedKeyIdInfoView, + api_key: selectedKeyIdInfoView, + } + setSelectedKeyInfo(keyResponse) } - }; - fetchKeyInfo(); - }, [selectedKeyIdInfoView, accessToken]); + } + fetchKeyInfo() + }, [selectedKeyIdInfoView, accessToken]) // Close dropdown when clicking outside useEffect(() => { function handleClickOutside(event: MouseEvent) { - if ( - dropdownRef.current && - !dropdownRef.current.contains(event.target as Node) - ) { - setShowColumnDropdown(false); + if (dropdownRef.current && !dropdownRef.current.contains(event.target as Node)) { + setShowColumnDropdown(false) } - if ( - filtersRef.current && - !filtersRef.current.contains(event.target as Node) - ) { - setShowFilters(false); + if (filtersRef.current && !filtersRef.current.contains(event.target as Node)) { + setShowFilters(false) } - if ( - quickSelectRef.current && - !quickSelectRef.current.contains(event.target as Node) - ) { - setQuickSelectOpen(false); + if (quickSelectRef.current && !quickSelectRef.current.contains(event.target as Node)) { + setQuickSelectOpen(false) } } - document.addEventListener("mousedown", handleClickOutside); - return () => - document.removeEventListener("mousedown", handleClickOutside); - }, []); - + document.addEventListener("mousedown", handleClickOutside) + return () => document.removeEventListener("mousedown", handleClickOutside) + }, []) useEffect(() => { if (userRole && internalUserRoles.includes(userRole)) { - setFilterByCurrentUser(true); + setFilterByCurrentUser(true) } - }, [userRole]); + }, [userRole]) const LiveTailControls = () => { return (
Live Tail - +
- ); - }; + ) + } const logs = useQuery({ queryKey: [ @@ -194,7 +164,7 @@ export default function SpendLogsTable({ selectedKeyHash, filterByCurrentUser ? userID : null, selectedStatus, - selectedModel + selectedModel, ], queryFn: async () => { if (!accessToken || !token || !userRole || !userID) { @@ -204,13 +174,13 @@ export default function SpendLogsTable({ page: 1, page_size: pageSize, total_pages: 0, - }; + } } - const formattedStartTime = moment(startTime).utc().format("YYYY-MM-DD HH:mm:ss"); - const formattedEndTime = isCustomDate + const formattedStartTime = moment(startTime).utc().format("YYYY-MM-DD HH:mm:ss") + const formattedEndTime = isCustomDate ? moment(endTime).utc().format("YYYY-MM-DD HH:mm:ss") - : moment().utc().format("YYYY-MM-DD HH:mm:ss"); + : moment().utc().format("YYYY-MM-DD HH:mm:ss") // Get base response from API const response = await uiSpendLogsCall( @@ -224,45 +194,42 @@ export default function SpendLogsTable({ pageSize, filterByCurrentUser ? userID : undefined, selectedStatus, - selectedModel - ); + selectedModel, + ) // Trigger prefetch for all logs - await prefetchLogDetails( - response.data, - formattedStartTime, - accessToken, - queryClient - ); + await prefetchLogDetails(response.data, formattedStartTime, accessToken, queryClient) // Update logs with prefetched data if available response.data = response.data.map((log: LogEntry) => { - const prefetchedData = queryClient.getQueryData( - ["logDetails", log.request_id, formattedStartTime] - ); + const prefetchedData = queryClient.getQueryData([ + "logDetails", + log.request_id, + formattedStartTime, + ]) if (prefetchedData?.messages && prefetchedData?.response) { - log.messages = prefetchedData.messages; - log.response = prefetchedData.response; - return log; + log.messages = prefetchedData.messages + log.response = prefetchedData.response + return log } - return log; - }); + return log + }) - return response; + return response }, enabled: !!accessToken && !!token && !!userRole && !!userID && activeTab === "request logs", refetchInterval: isLiveTail && currentPage === 1 ? 15000 : false, refetchIntervalInBackground: true, - }); + }) const logsData = logs.data || { data: [], total: 0, page: 1, page_size: pageSize || 10, - total_pages: 1 - }; + total_pages: 1, + } const { filters, @@ -270,7 +237,7 @@ export default function SpendLogsTable({ allTeams: hookAllTeams, allKeyAliases, handleFilterChange, - handleFilterReset + handleFilterReset, } = useLogFilterLogic({ logs: logsData, accessToken, @@ -280,63 +247,55 @@ export default function SpendLogsTable({ isCustomDate, setCurrentPage, userID, - userRole + userRole, }) - const fetchKeyHashForAlias = useCallback(async (keyAlias: string) => { - if (!accessToken) return; - - try { - const response = await keyListCall( - accessToken, - null, - null, - keyAlias, - null, - null, - currentPage, - pageSize - ); + const fetchKeyHashForAlias = useCallback( + async (keyAlias: string) => { + if (!accessToken) return - const selectedKey = response.keys.find( - (key: any) => key.key_alias === keyAlias - ); + try { + const response = await keyListCall(accessToken, null, null, keyAlias, null, null, currentPage, pageSize) - if (selectedKey) { - setSelectedKeyHash(selectedKey.token); + const selectedKey = response.keys.find((key: any) => key.key_alias === keyAlias) + + if (selectedKey) { + setSelectedKeyHash(selectedKey.token) + } + } catch (error) { + console.error("Error fetching key hash for alias:", error) } - } catch (error) { - console.error("Error fetching key hash for alias:", error); - } - }, [accessToken, currentPage, pageSize]); + }, + [accessToken, currentPage, pageSize], + ) // Add this effect to update selected filters when filter changes useEffect(() => { - if(!accessToken) return; + if (!accessToken) return - if (filters['Team ID']) { - setSelectedTeamId(filters['Team ID']); + if (filters["Team ID"]) { + setSelectedTeamId(filters["Team ID"]) } else { - setSelectedTeamId(""); + setSelectedTeamId("") } - setSelectedStatus(filters['Status'] || ""); - setSelectedModel(filters['Model'] || ""); - - if (filters['Key Hash']) { - setSelectedKeyHash(filters['Key Hash']); - } else if (filters['Key Alias']) { - fetchKeyHashForAlias(filters['Key Alias']); + setSelectedStatus(filters["Status"] || "") + setSelectedModel(filters["Model"] || "") + + if (filters["Key Hash"]) { + setSelectedKeyHash(filters["Key Hash"]) + } else if (filters["Key Alias"]) { + fetchKeyHashForAlias(filters["Key Alias"]) } else { - setSelectedKeyHash(""); + setSelectedKeyHash("") } - }, [filters, accessToken, fetchKeyHashForAlias]); + }, [filters, accessToken, fetchKeyHashForAlias]) // Fetch logs for a session if selected const sessionLogs = useQuery({ queryKey: ["sessionLogs", selectedSessionId], queryFn: async () => { - if (!accessToken || !selectedSessionId) return { data: [], total: 0, page: 1, page_size: 50, total_pages: 1 }; - const response = await sessionSpendLogsCall(accessToken, selectedSessionId); + if (!accessToken || !selectedSessionId) return { data: [], total: 0, page: 1, page_size: 50, total_pages: 1 } + const response = await sessionSpendLogsCall(accessToken, selectedSessionId) // If the API returns an array, wrap it in the same shape as PaginatedResponse return { data: response.data || response || [], @@ -344,130 +303,127 @@ export default function SpendLogsTable({ page: 1, page_size: 1000, total_pages: 1, - }; + } }, enabled: !!accessToken && !!selectedSessionId, - }); + }) // Add this effect to preserve expanded state when data refreshes useEffect(() => { if (logs.data?.data && expandedRequestId) { // Check if the expanded request ID still exists in the new data - const stillExists = logs.data.data.some(log => log.request_id === expandedRequestId); + const stillExists = logs.data.data.some((log) => log.request_id === expandedRequestId) if (!stillExists) { // If the request ID no longer exists in the data, clear the expanded state - setExpandedRequestId(null); + setExpandedRequestId(null) } } - }, [logs.data?.data, expandedRequestId]); + }, [logs.data?.data, expandedRequestId]) if (!accessToken || !token || !userRole || !userID) { - return null; + return null } const filteredData = - filteredLogs.data.filter((log) => { - const matchesSearch = - !searchTerm || - log.request_id.includes(searchTerm) || - log.model.includes(searchTerm) || - (log.user && log.user.includes(searchTerm)); - - // No need for additional filtering since we're now handling this in the API call - return matchesSearch; - - }).map(log => ({ - ...log, - duration: (Date.parse(log.endTime) - Date.parse(log.startTime)) / 1000, - onKeyHashClick: (keyHash: string) => setSelectedKeyIdInfoView(keyHash), - onSessionClick: (sessionId: string) => { - if (sessionId) setSelectedSessionId(sessionId); - }, - })) || []; + filteredLogs.data + .filter((log) => { + const matchesSearch = + !searchTerm || + log.request_id.includes(searchTerm) || + log.model.includes(searchTerm) || + (log.user && log.user.includes(searchTerm)) + + // No need for additional filtering since we're now handling this in the API call + return matchesSearch + }) + .map((log) => ({ + ...log, + duration: (Date.parse(log.endTime) - Date.parse(log.startTime)) / 1000, + onKeyHashClick: (keyHash: string) => setSelectedKeyIdInfoView(keyHash), + onSessionClick: (sessionId: string) => { + if (sessionId) setSelectedSessionId(sessionId) + }, + })) || [] // For session logs, add onKeyHashClick/onSessionClick as well const sessionData = - sessionLogs.data?.data?.map(log => ({ + sessionLogs.data?.data?.map((log) => ({ ...log, onKeyHashClick: (keyHash: string) => setSelectedKeyIdInfoView(keyHash), onSessionClick: (sessionId: string) => {}, - })) || []; + })) || [] // Add this function to handle manual refresh const handleRefresh = () => { - logs.refetch(); - }; - - + logs.refetch() + } const handleRowExpand = (requestId: string | null) => { - setExpandedRequestId(requestId); - }; + setExpandedRequestId(requestId) + } const logFilterOptions: FilterOption[] = [ { - name: 'Team ID', - label: 'Team ID', + name: "Team ID", + label: "Team ID", isSearchable: true, searchFn: async (searchText: string) => { - if (!allTeams || allTeams.length === 0) return []; - const filtered = allTeams.filter((team: Team) =>{ - return team.team_id.toLowerCase().includes(searchText.toLowerCase()) || - (team.team_alias && team.team_alias.toLowerCase().includes(searchText.toLowerCase())) - }); + if (!allTeams || allTeams.length === 0) return [] + const filtered = allTeams.filter((team: Team) => { + return ( + team.team_id.toLowerCase().includes(searchText.toLowerCase()) || + (team.team_alias && team.team_alias.toLowerCase().includes(searchText.toLowerCase())) + ) + }) return filtered.map((team: Team) => ({ label: `${team.team_alias || team.team_id} (${team.team_id})`, - value: team.team_id - })); - } + value: team.team_id, + })) + }, }, { - name:'Status', - label:'Status', + name: "Status", + label: "Status", isSearchable: false, options: [ - { label: 'Success', value: 'success' }, - { label: 'Failure', value: 'failure' } - ] + { label: "Success", value: "success" }, + { label: "Failure", value: "failure" }, + ], }, { - name: 'Model', - label: 'Model', + name: "Model", + label: "Model", isSearchable: false, }, { - name: 'Key Alias', - label: 'Key Alias', + name: "Key Alias", + label: "Key Alias", isSearchable: true, searchFn: async (searchText: string) => { - if (!accessToken) return []; - const keyAliases = await fetchAllKeyAliases(accessToken); - const filtered = keyAliases.filter(alias => - alias.toLowerCase().includes(searchText.toLowerCase()) - ); - return filtered.map(alias => ({ + if (!accessToken) return [] + const keyAliases = await fetchAllKeyAliases(accessToken) + const filtered = keyAliases.filter((alias) => alias.toLowerCase().includes(searchText.toLowerCase())) + return filtered.map((alias) => ({ label: alias, - value: alias - })); - } + value: alias, + })) + }, }, { - name: 'End User', - label: 'End User', + name: "End User", + label: "End User", isSearchable: true, searchFn: async (searchText: string) => { - if (!accessToken) return []; - const data = await allEndUsersCall(accessToken); - const users = data?.end_users || []; - const filtered = users.filter((u: string) => - u.toLowerCase().includes(searchText.toLowerCase()) - ); - return filtered.map((u: string) => ({ label: u, value: u })); - } + if (!accessToken) return [] + const data = await allEndUsersCall(accessToken) + const users = data?.end_users || [] + const filtered = users.filter((u: string) => u.toLowerCase().includes(searchText.toLowerCase())) + return filtered.map((u: string) => ({ label: u, value: u })) + }, }, { - name: 'Key Hash', - label: 'Key Hash', + name: "Key Hash", + label: "Key Hash", isSearchable: false, }, ] @@ -482,17 +438,17 @@ export default function SpendLogsTable({ onBack={() => setSelectedSessionId(null)} />
- ); + ) } const formatTimeUnit = (value: number, unit: string) => { if (value === 1) { - if (unit === 'minutes') return 'minute'; - if (unit === 'hours') return 'hour'; - if (unit === 'days') return 'day'; + if (unit === "minutes") return "minute" + if (unit === "hours") return "hour" + if (unit === "days") return "day" } - return unit; - }; + return unit + } const quickSelectOptions = [ { label: "Last 15 Minutes", value: 15, unit: "minutes" }, @@ -500,27 +456,20 @@ export default function SpendLogsTable({ { label: "Last 4 Hours", value: 4, unit: "hours" }, { label: "Last 24 Hours", value: 24, unit: "hours" }, { label: "Last 7 Days", value: 7, unit: "days" }, - ]; + ] const selectedOption = quickSelectOptions.find( - (option) => - option.value === selectedTimeInterval.value && option.unit === selectedTimeInterval.unit - ); + (option) => option.value === selectedTimeInterval.value && option.unit === selectedTimeInterval.unit, + ) - const displayLabel = isCustomDate - ? getTimeRangeDisplay(isCustomDate, startTime, endTime) - : selectedOption?.label; + const displayLabel = isCustomDate ? getTimeRangeDisplay(isCustomDate, startTime, endTime) : selectedOption?.label return (
setActiveTab(index === 0 ? "request logs" : "audit logs")}> - - Request Logs - - - Audit Logs - + Request Logs + Audit Logs @@ -542,7 +491,16 @@ export default function SpendLogsTable({
{selectedKeyInfo && selectedKeyIdInfoView && selectedKeyInfo.api_key === selectedKeyIdInfoView ? ( - setSelectedKeyIdInfoView(null)} premiumUser={premiumUser} /> + setSelectedKeyIdInfoView(null)} + premiumUser={premiumUser} + /> ) : selectedSessionId ? (
) : ( <> - -
-
-
-
-
- setSearchTerm(e.target.value)} - /> - - +
+
+
+
+
+ setSearchTerm(e.target.value)} /> - -
- -
-
- - - {quickSelectOpen && ( -
-
- {quickSelectOptions.map((option) => ( - - ))} -
- -
-
- )} -
- - - - +
+ +
+
+ + + {quickSelectOpen && ( +
+
+ {quickSelectOptions.map((option) => ( + + ))} +
+ +
+
+ )} +
+ + + + +
+ + {isCustomDate && ( +
+
+ { + setStartTime(e.target.value) + setCurrentPage(1) + }} + className="px-3 py-2 border rounded-md text-sm focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-blue-500" + /> +
+ to +
+ { + setEndTime(e.target.value) + setCurrentPage(1) + }} + className="px-3 py-2 border rounded-md text-sm focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-blue-500" + /> +
+
+ )}
- {isCustomDate && ( -
-
- { - setStartTime(e.target.value); - setCurrentPage(1); - }} - className="px-3 py-2 border rounded-md text-sm focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-blue-500" - /> -
- to -
- { - setEndTime(e.target.value); - setCurrentPage(1); - }} - className="px-3 py-2 border rounded-md text-sm focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-blue-500" - /> -
-
- )} -
- -
- - Showing{" "} - {logs.isLoading - ? "..." - : filteredLogs - ? (currentPage - 1) * pageSize + 1 - : 0}{" "} - -{" "} - {logs.isLoading - ? "..." - : filteredLogs - ? Math.min(currentPage * pageSize, filteredLogs.total) - : 0}{" "} - of{" "} - {logs.isLoading - ? "..." - : filteredLogs - ? filteredLogs.total - : 0}{" "} - results - -
+
- Page {logs.isLoading ? "..." : currentPage} of{" "} + Showing {logs.isLoading ? "..." : filteredLogs ? (currentPage - 1) * pageSize + 1 : 0} -{" "} {logs.isLoading ? "..." : filteredLogs - ? filteredLogs.total_pages - : 1} + ? Math.min(currentPage * pageSize, filteredLogs.total) + : 0}{" "} + of {logs.isLoading ? "..." : filteredLogs ? filteredLogs.total : 0} results - - +
+ + Page {logs.isLoading ? "..." : currentPage} of{" "} + {logs.isLoading ? "..." : filteredLogs ? filteredLogs.total_pages : 1} + + + +
-
- {isLiveTail && currentPage === 1 && ( -
-
- - Auto-refreshing every 15 seconds - + {isLiveTail && currentPage === 1 && ( +
+
+ Auto-refreshing every 15 seconds +
+
- -
- )} - true} - /> -
+ )} + true} + /> +
- )} + )} -
- ); + ) } export function RequestViewer({ row }: { row: Row }) { @@ -804,35 +731,38 @@ export function RequestViewer({ row }: { row: Row }) { const formatData = (input: any) => { if (typeof input === "string") { try { - return JSON.parse(input); + return JSON.parse(input) } catch { - return input; + return input } } - return input; - }; + return input + } // New helper function to get raw request const getRawRequest = () => { // First check if proxy_server_request exists in metadata if (row.original?.proxy_server_request) { - return formatData(row.original.proxy_server_request); + return formatData(row.original.proxy_server_request) } // Fall back to messages if proxy_server_request is empty - return formatData(row.original.messages); - }; + return formatData(row.original.messages) + } // Extract error information from metadata if available - const metadata = row.original.metadata || {}; - const hasError = metadata.status === "failure"; - const errorInfo = hasError ? metadata.error_information : null; - + const metadata = row.original.metadata || {} + const hasError = metadata.status === "failure" + const errorInfo = hasError ? metadata.error_information : null + // Check if request/response data is missing - const hasMessages = row.original.messages && - (Array.isArray(row.original.messages) ? row.original.messages.length > 0 : Object.keys(row.original.messages).length > 0); - const hasResponse = row.original.response && Object.keys(formatData(row.original.response)).length > 0; - const missingData = !hasMessages && !hasResponse; - + const hasMessages = + row.original.messages && + (Array.isArray(row.original.messages) + ? row.original.messages.length > 0 + : Object.keys(row.original.messages).length > 0) + const hasResponse = row.original.response && Object.keys(formatData(row.original.response)).length > 0 + const missingData = !hasMessages && !hasResponse + // Format the response with error details if present const formattedResponse = () => { if (hasError && errorInfo) { @@ -840,32 +770,35 @@ export function RequestViewer({ row }: { row: Row }) { error: { message: errorInfo.error_message || "An error occurred", type: errorInfo.error_class || "error", - code: errorInfo.error_code || "unknown", - param: null - } - }; + code: errorInfo.error_code || "unknown", + param: null, + }, + } } - return formatData(row.original.response); - }; - + return formatData(row.original.response) + } + // Extract vector store request metadata if available - const hasVectorStoreData = metadata.vector_store_request_metadata && - Array.isArray(metadata.vector_store_request_metadata) && - metadata.vector_store_request_metadata.length > 0; + const hasVectorStoreData = + metadata.vector_store_request_metadata && + Array.isArray(metadata.vector_store_request_metadata) && + metadata.vector_store_request_metadata.length > 0 // Extract guardrail information from metadata if available - const hasGuardrailData = row.original.metadata && row.original.metadata.guardrail_information; - + const hasGuardrailData = row.original.metadata && row.original.metadata.guardrail_information + // Calculate total masked entities if guardrail data exists const getTotalMaskedEntities = (): number => { if (!hasGuardrailData || !row.original.metadata?.guardrail_information.masked_entity_count) { - return 0; + return 0 } - return Object.values(row.original.metadata.guardrail_information.masked_entity_count) - .reduce((sum: number, count: any) => sum + (typeof count === 'number' ? count : 0), 0); - }; - - const totalMaskedEntities = getTotalMaskedEntities(); + return Object.values(row.original.metadata.guardrail_information.masked_entity_count).reduce( + (sum: number, count: any) => sum + (typeof count === "number" ? count : 0), + 0, + ) + } + + const totalMaskedEntities = getTotalMaskedEntities() return (
@@ -925,15 +858,22 @@ export function RequestViewer({ row }: { row: Row }) {
Tokens: - {row.original.total_tokens} ({row.original.prompt_tokens} prompt tokens + {row.original.completion_tokens} completion tokens) + + {row.original.total_tokens} ({row.original.prompt_tokens} prompt tokens +{" "} + {row.original.completion_tokens} completion tokens) +
Cache Read Tokens: - {formatNumberWithCommas(row.original.metadata?.additional_usage_values?.cache_read_input_tokens || 0)} + + {formatNumberWithCommas(row.original.metadata?.additional_usage_values?.cache_read_input_tokens || 0)} +
Cache Creation Tokens: - {formatNumberWithCommas(row.original.metadata?.additional_usage_values.cache_creation_input_tokens)} + + {formatNumberWithCommas(row.original.metadata?.additional_usage_values.cache_creation_input_tokens)} +
Cost: @@ -943,17 +883,18 @@ export function RequestViewer({ row }: { row: Row }) { Cache Hit: {row.original.cache_hit}
- +
Status: - + {(row.original.metadata?.status || "Success").toLowerCase() !== "failure" ? "Success" : "Failure"} -
Start Time: @@ -986,14 +927,10 @@ export function RequestViewer({ row }: { row: Row }) { /> {/* Guardrail Data - Show only if present */} - {hasGuardrailData && ( - - )} + {hasGuardrailData && } {/* Vector Store Request Data - Show only if present */} - {hasVectorStoreData && ( - - )} + {hasVectorStoreData && } {/* Error Card - Only show for failures */} {hasError && errorInfo && } @@ -1021,14 +958,24 @@ export function RequestViewer({ row }: { row: Row }) {

Metadata

-
)}
- ); -} \ No newline at end of file + ) +} diff --git a/ui/litellm-dashboard/src/components/view_users.tsx b/ui/litellm-dashboard/src/components/view_users.tsx index 39a5206fea3..0275454abd0 100644 --- a/ui/litellm-dashboard/src/components/view_users.tsx +++ b/ui/litellm-dashboard/src/components/view_users.tsx @@ -296,7 +296,7 @@ const ViewUserDashboard: React.FC = ({ accessToken, toke ) return ( -
+
@@ -506,7 +506,7 @@ const ViewUserDashboard: React.FC = ({ accessToken, toke
- +
= ({ accessToken, toke selectedUsers={selectedUsers} onSelectionChange={handleSelectionChange} /> +
+
diff --git a/ui/litellm-dashboard/src/components/view_users/columns.tsx b/ui/litellm-dashboard/src/components/view_users/columns.tsx index 63dee940dc3..7e114ced6fb 100644 --- a/ui/litellm-dashboard/src/components/view_users/columns.tsx +++ b/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 }) => ( {row.original.user_email || "-"} @@ -49,7 +49,7 @@ export const columns = ( ), }, { - header: "User Spend ($ USD)", + header: "Spend (USD)", accessorKey: "spend", cell: ({ row }) => ( @@ -58,7 +58,7 @@ export const columns = ( ), }, { - header: "User Max Budget ($ USD)", + header: "Budget (USD)", accessorKey: "max_budget", cell: ({ row }) => (