diff --git a/.circleci/config.yml b/.circleci/config.yml index 38e6d1fc332..12e3cb1f6b6 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3218,6 +3218,7 @@ jobs: -e DD_API_KEY=$DD_API_KEY \ -e DD_SITE=$DD_SITE \ -e LITELLM_LICENSE=$LITELLM_LICENSE \ + -e LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES=true \ --add-host host.docker.internal:host-gateway \ --name my-app \ -v $(pwd)/litellm/proxy/example_config_yaml/pass_through_config.yaml:/app/config.yaml \ diff --git a/CLAUDE.md b/CLAUDE.md index 5395d6d938e..d9061b5e2be 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -156,4 +156,4 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components: **Fix options:** 1. **Create a Prisma migration** (permanent) — run `prisma migrate dev --name ` in the worktree. The generated file will be picked up by `prisma migrate deploy` on next startup. 2. **Apply manually for local dev** — `psql -d litellm -c "ALTER TABLE ... ADD COLUMN IF NOT EXISTS ..."` after each proxy start. Fine for dev, not for production. -3. **Update litellm-proxy-extras** — if the package is installed from PyPI, its migration directory must include the new file. Either update the package or run the migration manually until the next release ships it. \ No newline at end of file +3. **Update litellm-proxy-extras** — if the package is installed from PyPI, its migration directory must include the new file. Either update the package or run the migration manually until the next release ships it. diff --git a/ci_cd/security_scans.sh b/ci_cd/security_scans.sh index 62440d13ebb..e0f370e0035 100755 --- a/ci_cd/security_scans.sh +++ b/ci_cd/security_scans.sh @@ -11,7 +11,7 @@ echo "Starting security scans for LiteLLM..." install_trivy() { echo "Installing Trivy and required tools..." sudo apt-get update - sudo apt-get install -y wget apt-transport-https gnupg lsb-release jq curl + sudo apt-get install -y wget apt-transport-https gnupg lsb-release jq curl bsdmainutils wget -qO - https://aquasecurity.github.io/trivy-repo/deb/public.key | sudo apt-key add - echo "deb https://aquasecurity.github.io/trivy-repo/deb $(lsb_release -sc) main" | sudo tee -a /etc/apt/sources.list.d/trivy.list sudo apt-get update diff --git a/docs/my-website/docs/providers/baseten.md b/docs/my-website/docs/providers/baseten.md index 4e42cdf0447..d8283c6df3d 100644 --- a/docs/my-website/docs/providers/baseten.md +++ b/docs/my-website/docs/providers/baseten.md @@ -10,12 +10,12 @@ LiteLLM supports both Baseten Model APIs and dedicated deployments with automati ### Model API (Default) - **URL**: `https://inference.baseten.co/v1` - **Format**: `baseten/` (e.g., `baseten/openai/gpt-oss-120b`) -- **Best for**: Quick access to popular models +- **Best for**: Quick access to popular models available on Baseten Model APIs: https://docs.baseten.co/development/model-apis/overview#supported-models ### Dedicated Deployments - **URL**: `https://model-{id}.api.baseten.co/environments/production/sync/v1` -- **Format**: `baseten/{8-digit-alphanumeric-code}` (e.g., `baseten/abcd1234`) -- **Best for**: Custom models, latency SLAs +- **Format**: `baseten/{8-digit-baseten-model-id}` (e.g., `baseten/abcd1234`) +- **Best for**: Custom models, enterprise SLAs :::tip **Automatic Routing**: LiteLLM detects the type based on model format: @@ -82,6 +82,8 @@ for chunk in response: ## Usage with LiteLLM Proxy +### Model API + 1. **Config**: ```yaml model_list: @@ -94,13 +96,42 @@ model_list: 2. **Request**: ```python import openai -client = openai.OpenAI( - api_key="sk-1234", - base_url="http://0.0.0.0:4000" +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +response = client.chat.completions.create( + model="baseten-model", + messages=[{"role": "user", "content": "Hello!"}] ) +``` + +### Dedicated Deployment + +If your dedicated deployment uses a `served_model_name` in your Baseten `config.yaml`, you must supply `served_model_name` to specify the model name sent in the request body, and supply the Baseten model id under the `model` field. + +1. **Config**: +```yaml +model_list: + - model_name: baseten-model # external user facing + litellm_params: + model: baseten/1234abcd # model id from Baseten dashboard + served_model_name: baseten-hosted/zai-org/GLM-5 # model name specified in Baseten config.yaml + api_key: os.environ/BASETEN_API_KEY +``` + +2. **Request**: +```python +import openai +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") response = client.chat.completions.create( model="baseten-model", messages=[{"role": "user", "content": "Hello!"}] ) ``` + +- `model: baseten/1234abcd` — the 8-digit deployment ID, used to route to `https://model-1234abcd.api.baseten.co/environments/production/sync/v1` +- `served_model_name` — sent as the `model` field in the request body, matching your deployment's configured model name. + +:::note +`served_model_name` is optional. If your deployment's model name is empty, you can omit it and just use `model: baseten/{deployment_id}`. +::: diff --git a/litellm/__init__.py b/litellm/__init__.py index 299bb18245c..51c66838613 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -358,7 +358,7 @@ ) blog_posts_url: str = os.getenv( "LITELLM_BLOG_POSTS_URL", - "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/blog_posts.json", + "https://docs.litellm.ai/blog/rss.xml", ) anthropic_beta_headers_url: str = os.getenv( "LITELLM_ANTHROPIC_BETA_HEADERS_URL", diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index a3ef7b264ec..e0e1e35b94e 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -660,7 +660,14 @@ def _select_model_name_for_cost_calc( if custom_pricing is True: if router_model_id is not None and router_model_id in litellm.model_cost: - return_model = router_model_id + entry = litellm.model_cost[router_model_id] + if ( + entry.get("input_cost_per_token") is not None + or entry.get("input_cost_per_second") is not None + ): + return_model = router_model_id + else: + return_model = model else: return_model = model diff --git a/litellm/litellm_core_utils/get_blog_posts.py b/litellm/litellm_core_utils/get_blog_posts.py index f54deb59290..2f9a14f1279 100644 --- a/litellm/litellm_core_utils/get_blog_posts.py +++ b/litellm/litellm_core_utils/get_blog_posts.py @@ -1,8 +1,8 @@ """ -Pulls the latest LiteLLM blog posts from GitHub. +Pulls the latest LiteLLM blog posts from the docs RSS feed. Falls back to the bundled local backup on any failure. -GitHub JSON URL is configured via litellm.blog_posts_url (or LITELLM_BLOG_POSTS_URL env var). +RSS URL is configured via litellm.blog_posts_url (or LITELLM_BLOG_POSTS_URL env var). Disable remote fetching entirely: export LITELLM_LOCAL_BLOG_POSTS=True @@ -11,8 +11,10 @@ import json import os import time +import xml.etree.ElementTree as ET +from email.utils import parsedate_to_datetime from importlib.resources import files -from typing import Any, Dict, List, Optional +from typing import Dict, List, Optional import httpx from pydantic import BaseModel @@ -37,9 +39,8 @@ class GetBlogPosts: """ Fetches, validates, and caches LiteLLM blog posts. - Mirrors the structure of GetModelCostMap: - - Fetches from GitHub with a 5-second timeout - - Validates the response has a non-empty ``posts`` list + - Fetches RSS feed from docs site with a 5-second timeout + - Parses the XML and extracts the latest blog post - Caches the result in-process for BLOG_POSTS_TTL_SECONDS (1 hour) - Falls back to the bundled local backup on any failure """ @@ -56,30 +57,67 @@ def load_local_blog_posts() -> List[Dict[str, str]]: return content.get("posts", []) @staticmethod - def fetch_remote_blog_posts(url: str, timeout: int = 5) -> dict: + def fetch_rss_feed(url: str, timeout: int = 5) -> str: """ - Fetch blog posts JSON from a remote URL. + Fetch RSS XML from a remote URL. - Returns the parsed response. Raises on network/parse errors. + Returns the raw XML text. Raises on network errors. """ response = httpx.get(url, timeout=timeout) response.raise_for_status() - return response.json() + return response.text @staticmethod - def validate_blog_posts(data: Any) -> bool: - """Return True if data is a dict with a non-empty ``posts`` list.""" - if not isinstance(data, dict): - verbose_logger.warning( - "LiteLLM: Blog posts response is not a dict (type=%s). " - "Falling back to local backup.", - type(data).__name__, + def parse_rss_to_posts(xml_text: str, max_posts: int = 1) -> List[Dict[str, str]]: + """ + Parse RSS XML and return a list of blog post dicts. + + Extracts title, description, date (YYYY-MM-DD), and url from each . + """ + root = ET.fromstring(xml_text) + channel = root.find("channel") + if channel is None: + raise ValueError("RSS feed missing element") + + posts: List[Dict[str, str]] = [] + for item in channel.findall("item"): + if len(posts) >= max_posts: + break + + title_el = item.find("title") + link_el = item.find("link") + desc_el = item.find("description") + pub_date_el = item.find("pubDate") + + if title_el is None or link_el is None: + continue + + # Parse RFC 2822 date to YYYY-MM-DD + date_str = "" + if pub_date_el is not None and pub_date_el.text: + try: + dt = parsedate_to_datetime(pub_date_el.text) + date_str = dt.strftime("%Y-%m-%d") + except Exception: + date_str = pub_date_el.text + + posts.append( + { + "title": title_el.text or "", + "description": desc_el.text or "" if desc_el is not None else "", + "date": date_str, + "url": link_el.text or "", + } ) - return False - posts = data.get("posts") + + return posts + + @staticmethod + def validate_blog_posts(posts: List[Dict[str, str]]) -> bool: + """Return True if posts is a non-empty list.""" if not isinstance(posts, list) or len(posts) == 0: verbose_logger.warning( - "LiteLLM: Blog posts response has no valid 'posts' list. " + "LiteLLM: Parsed RSS feed has no valid posts. " "Falling back to local backup.", ) return False @@ -102,7 +140,8 @@ def get_blog_posts(cls, url: str) -> List[Dict[str, str]]: return cached try: - data = cls.fetch_remote_blog_posts(url) + xml_text = cls.fetch_rss_feed(url) + posts = cls.parse_rss_to_posts(xml_text) except Exception as e: verbose_logger.warning( "LiteLLM: Failed to fetch blog posts from %s: %s. " @@ -112,10 +151,9 @@ def get_blog_posts(cls, url: str) -> List[Dict[str, str]]: ) return cls.load_local_blog_posts() - if not cls.validate_blog_posts(data): + if not cls.validate_blog_posts(posts): return cls.load_local_blog_posts() - posts = data["posts"] cls._cached_posts = posts cls._last_fetch_time = now return posts diff --git a/litellm/llms/baseten/chat.py b/litellm/llms/baseten/chat.py index 1e49b346088..f4c2b51997d 100644 --- a/litellm/llms/baseten/chat.py +++ b/litellm/llms/baseten/chat.py @@ -1,5 +1,7 @@ -from typing import Optional +from typing import List, Optional + from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig +from litellm.types.llms.openai import AllMessageValues class BasetenConfig(OpenAIGPTConfig): @@ -82,6 +84,28 @@ def map_openai_params( optional_params[param] = value return optional_params + def transform_request( + self, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + # For dedicated deployments, the model is the deployment ID (e.g. "wd1lndkw") + # but the server may expect a different model name in the request body + served_model_name = litellm_params.get("served_model_name") + if served_model_name: + model = served_model_name + + return super().transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + def _get_openai_compatible_provider_info( self, api_base: str, api_key: str ) -> tuple: diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 1bec0d23c91..ef01f027d6f 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -903,13 +903,17 @@ async def _execute_with_mcp_client( try: client_id, client_secret, scopes = _extract_credentials(request) - _oauth2_flow: Optional[ - Literal["client_credentials", "authorization_code"] - ] = ( - "client_credentials" - if client_id and client_secret and request.token_url - else None + _oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = ( + request.oauth2_flow or ( + "client_credentials" + if client_id and client_secret and request.token_url + else None + ) ) + # client_credentials requires token_url to fetch a token; without it the + # incoming auth header would be dropped with nothing to replace it. + if _oauth2_flow == "client_credentials" and not request.token_url: + _oauth2_flow = None server_model = MCPServer( server_id=request.server_id or "", diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404/index.html similarity index 100% rename from litellm/proxy/_experimental/out/404.html rename to litellm/proxy/_experimental/out/404/index.html diff --git a/litellm/proxy/_experimental/out/_not-found.html b/litellm/proxy/_experimental/out/_not-found/index.html similarity index 100% rename from litellm/proxy/_experimental/out/_not-found.html rename to litellm/proxy/_experimental/out/_not-found/index.html diff --git a/litellm/proxy/_experimental/out/api-reference.html b/litellm/proxy/_experimental/out/api-reference/index.html similarity index 100% rename from litellm/proxy/_experimental/out/api-reference.html rename to litellm/proxy/_experimental/out/api-reference/index.html diff --git a/litellm/proxy/_experimental/out/experimental/api-playground.html b/litellm/proxy/_experimental/out/experimental/api-playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/api-playground.html rename to litellm/proxy/_experimental/out/experimental/api-playground/index.html diff --git a/litellm/proxy/_experimental/out/experimental/budgets.html b/litellm/proxy/_experimental/out/experimental/budgets/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/budgets.html rename to litellm/proxy/_experimental/out/experimental/budgets/index.html diff --git a/litellm/proxy/_experimental/out/experimental/caching.html b/litellm/proxy/_experimental/out/experimental/caching/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/caching.html rename to litellm/proxy/_experimental/out/experimental/caching/index.html diff --git a/litellm/proxy/_experimental/out/experimental/claude-code-plugins.html b/litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/claude-code-plugins.html rename to litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html diff --git a/litellm/proxy/_experimental/out/experimental/old-usage.html b/litellm/proxy/_experimental/out/experimental/old-usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/old-usage.html rename to litellm/proxy/_experimental/out/experimental/old-usage/index.html diff --git a/litellm/proxy/_experimental/out/experimental/prompts.html b/litellm/proxy/_experimental/out/experimental/prompts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/prompts.html rename to litellm/proxy/_experimental/out/experimental/prompts/index.html diff --git a/litellm/proxy/_experimental/out/experimental/tag-management.html b/litellm/proxy/_experimental/out/experimental/tag-management/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/tag-management.html rename to litellm/proxy/_experimental/out/experimental/tag-management/index.html diff --git a/litellm/proxy/_experimental/out/guardrails.html b/litellm/proxy/_experimental/out/guardrails/index.html similarity index 100% rename from litellm/proxy/_experimental/out/guardrails.html rename to litellm/proxy/_experimental/out/guardrails/index.html diff --git a/litellm/proxy/_experimental/out/login.html b/litellm/proxy/_experimental/out/login/index.html similarity index 100% rename from litellm/proxy/_experimental/out/login.html rename to litellm/proxy/_experimental/out/login/index.html diff --git a/litellm/proxy/_experimental/out/logs.html b/litellm/proxy/_experimental/out/logs/index.html similarity index 100% rename from litellm/proxy/_experimental/out/logs.html rename to litellm/proxy/_experimental/out/logs/index.html diff --git a/litellm/proxy/_experimental/out/mcp/oauth/callback.html b/litellm/proxy/_experimental/out/mcp/oauth/callback/index.html similarity index 100% rename from litellm/proxy/_experimental/out/mcp/oauth/callback.html rename to litellm/proxy/_experimental/out/mcp/oauth/callback/index.html diff --git a/litellm/proxy/_experimental/out/model-hub.html b/litellm/proxy/_experimental/out/model-hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model-hub.html rename to litellm/proxy/_experimental/out/model-hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub.html rename to litellm/proxy/_experimental/out/model_hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub_table.html rename to litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/proxy/_experimental/out/models-and-endpoints.html b/litellm/proxy/_experimental/out/models-and-endpoints/index.html similarity index 100% rename from litellm/proxy/_experimental/out/models-and-endpoints.html rename to litellm/proxy/_experimental/out/models-and-endpoints/index.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding/index.html similarity index 100% rename from litellm/proxy/_experimental/out/onboarding.html rename to litellm/proxy/_experimental/out/onboarding/index.html diff --git a/litellm/proxy/_experimental/out/organizations.html b/litellm/proxy/_experimental/out/organizations/index.html similarity index 100% rename from litellm/proxy/_experimental/out/organizations.html rename to litellm/proxy/_experimental/out/organizations/index.html diff --git a/litellm/proxy/_experimental/out/playground.html b/litellm/proxy/_experimental/out/playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/playground.html rename to litellm/proxy/_experimental/out/playground/index.html diff --git a/litellm/proxy/_experimental/out/policies.html b/litellm/proxy/_experimental/out/policies/index.html similarity index 100% rename from litellm/proxy/_experimental/out/policies.html rename to litellm/proxy/_experimental/out/policies/index.html diff --git a/litellm/proxy/_experimental/out/settings/admin-settings.html b/litellm/proxy/_experimental/out/settings/admin-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/admin-settings.html rename to litellm/proxy/_experimental/out/settings/admin-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/logging-and-alerts.html b/litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/logging-and-alerts.html rename to litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html diff --git a/litellm/proxy/_experimental/out/settings/router-settings.html b/litellm/proxy/_experimental/out/settings/router-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/router-settings.html rename to litellm/proxy/_experimental/out/settings/router-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/ui-theme.html b/litellm/proxy/_experimental/out/settings/ui-theme/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/ui-theme.html rename to litellm/proxy/_experimental/out/settings/ui-theme/index.html diff --git a/litellm/proxy/_experimental/out/teams.html b/litellm/proxy/_experimental/out/teams/index.html similarity index 100% rename from litellm/proxy/_experimental/out/teams.html rename to litellm/proxy/_experimental/out/teams/index.html diff --git a/litellm/proxy/_experimental/out/test-key.html b/litellm/proxy/_experimental/out/test-key/index.html similarity index 100% rename from litellm/proxy/_experimental/out/test-key.html rename to litellm/proxy/_experimental/out/test-key/index.html diff --git a/litellm/proxy/_experimental/out/tools/mcp-servers.html b/litellm/proxy/_experimental/out/tools/mcp-servers/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/mcp-servers.html rename to litellm/proxy/_experimental/out/tools/mcp-servers/index.html diff --git a/litellm/proxy/_experimental/out/tools/vector-stores.html b/litellm/proxy/_experimental/out/tools/vector-stores/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/vector-stores.html rename to litellm/proxy/_experimental/out/tools/vector-stores/index.html diff --git a/litellm/proxy/_experimental/out/usage.html b/litellm/proxy/_experimental/out/usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/usage.html rename to litellm/proxy/_experimental/out/usage/index.html diff --git a/litellm/proxy/_experimental/out/users.html b/litellm/proxy/_experimental/out/users/index.html similarity index 100% rename from litellm/proxy/_experimental/out/users.html rename to litellm/proxy/_experimental/out/users/index.html diff --git a/litellm/proxy/_experimental/out/virtual-keys.html b/litellm/proxy/_experimental/out/virtual-keys/index.html similarity index 100% rename from litellm/proxy/_experimental/out/virtual-keys.html rename to litellm/proxy/_experimental/out/virtual-keys/index.html diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b7ac4212cbd..ecbd7314cd7 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1123,6 +1123,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None + oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None allow_all_keys: bool = False available_on_public_internet: bool = True is_byok: bool = False @@ -4262,7 +4263,7 @@ class DefaultInternalUserParams(LiteLLMPydanticObjectBase): LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ] ] = Field( - default=LitellmUserRoles.INTERNAL_USER, + default=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, description="Default role assigned to new users created", ) max_budget: Optional[float] = Field( diff --git a/litellm/proxy/example_config_yaml/custom_auth_basic.py b/litellm/proxy/example_config_yaml/custom_auth_basic.py index 4d633a54fe2..0da6105a305 100644 --- a/litellm/proxy/example_config_yaml/custom_auth_basic.py +++ b/litellm/proxy/example_config_yaml/custom_auth_basic.py @@ -1,6 +1,6 @@ from fastapi import Request -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth async def user_api_key_auth(request: Request, api_key: str) -> UserAPIKeyAuth: @@ -9,6 +9,7 @@ async def user_api_key_auth(request: Request, api_key: str) -> UserAPIKeyAuth: api_key="best-api-key-ever", user_id="best-user-id-ever", team_id="best-team-id-ever", + user_role=LitellmUserRoles.PROXY_ADMIN, ) except Exception: raise Exception diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py index 79f1992da44..f9ebf46a270 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py @@ -1,3 +1,33 @@ +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + from .dynamoai import DynamoAIGuardrails -__all__ = ["DynamoAIGuardrails"] +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + import litellm + + _dynamoai_callback = DynamoAIGuardrails( + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + guardrail_name=guardrail.get("guardrail_name", ""), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(_dynamoai_callback) + + return _dynamoai_callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.DYNAMOAI.value: initialize_guardrail, +} + + +guardrail_class_registry = { + SupportedGuardrailIntegrations.DYNAMOAI.value: DynamoAIGuardrails, +} diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 1c0c212b60b..e09d7607ebe 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -54,6 +54,7 @@ from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.management_endpoints.common_utils import ( + _is_user_org_admin_for_team, _is_user_team_admin, _set_object_metadata_field, ) @@ -71,6 +72,9 @@ ) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.spend_tracking.spend_tracking_utils import _is_master_key +from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( + get_ui_settings_cached, +) from litellm.proxy.utils import ( PrismaClient, ProxyLogging, @@ -95,6 +99,24 @@ ) +async def _check_custom_key_allowed(custom_key_value: Optional[str]) -> None: + """Raise 403 if custom API keys are disabled and a custom key was provided.""" + if custom_key_value is None: + return + + ui_settings = await get_ui_settings_cached() + if ui_settings.get("disable_custom_api_keys", False) is True: + verbose_proxy_logger.warning( + "Custom API key rejected: disable_custom_api_keys is enabled" + ) + raise HTTPException( + status_code=403, + detail={ + "error": "Custom API key values are disabled by your administrator. Keys must be auto-generated." + }, + ) + + def _is_team_key(data: Union[GenerateKeyRequest, LiteLLM_VerificationToken]): return data.team_id is not None @@ -353,6 +375,10 @@ def key_generation_check( ## check if key is for team or individual is_team_key = _is_team_key(data=data) + _is_admin = ( + user_api_key_dict.user_role is not None + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ) if is_team_key: if team_table is None and litellm.key_generation_settings is not None: raise HTTPException( @@ -360,7 +386,13 @@ def key_generation_check( detail=f"Unable to find team object in database. Team ID: {data.team_id}", ) elif team_table is None: - return True # assume user is assigning team_id without using the team table + if _is_admin: + return True # admins can assign team_id without team table + # Non-admin callers must have a valid team (LIT-1884) + raise HTTPException( + status_code=400, + detail=f"Unable to find team object in database. Team ID: {data.team_id}", + ) return _team_key_generation_check( team_table=team_table, user_api_key_dict=user_api_key_dict, @@ -660,6 +692,9 @@ async def _common_key_generation_helper( # noqa: PLR0915 prisma_client=prisma_client, ) + # Reject custom key values if disabled by admin + await _check_custom_key_allowed(data.key) + # Validate user-provided key format if data.key is not None and not data.key.startswith("sk-"): _masked = ( @@ -1213,6 +1248,19 @@ async def generate_key_fn( raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=message ) + # For non-admin internal users: auto-assign caller's user_id if not provided + # This prevents creating unbound keys with no user association (LIT-1884) + _is_proxy_admin = ( + user_api_key_dict.user_role is not None + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ) + if not _is_proxy_admin and data.user_id is None: + data.user_id = user_api_key_dict.user_id + verbose_proxy_logger.warning( + "key/generate: auto-assigning user_id=%s for non-admin caller", + user_api_key_dict.user_id, + ) + team_table: Optional[LiteLLM_TeamTableCachedObj] = None if data.team_id is not None: try: @@ -1227,6 +1275,12 @@ async def generate_key_fn( verbose_proxy_logger.debug( f"Error getting team object in `/key/generate`: {e}" ) + # For non-admin callers, team must exist (LIT-1884) + if not _is_proxy_admin: + raise HTTPException( + status_code=400, + detail=f"Team not found for team_id={data.team_id}. Non-admin users cannot create keys for non-existent teams.", + ) key_generation_check( team_table=team_table, @@ -1809,11 +1863,26 @@ async def _validate_update_key_data( user_api_key_cache: Any, ) -> None: """Validate permissions and constraints for key update.""" + _is_proxy_admin = ( + user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ) + + # Prevent non-admin from removing user_id (setting to empty string) (LIT-1884) + if ( + data.user_id is not None + and data.user_id == "" + and not _is_proxy_admin + ): + raise HTTPException( + status_code=403, + detail="Non-admin users cannot remove the user_id from a key.", + ) + # sanity check - prevent non-proxy admin user from updating key to belong to a different user if ( data.user_id is not None and data.user_id != existing_key_row.user_id - and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value + and not _is_proxy_admin ): raise HTTPException( status_code=403, @@ -1836,6 +1905,18 @@ async def _validate_update_key_data( user_api_key_cache=user_api_key_cache, ) + # Admin-only: only proxy admins, team admins, or org admins can modify max_budget + if data.max_budget is not None and data.max_budget != existing_key_row.max_budget: + if prisma_client is not None: + hashed_key = existing_key_row.token + await _check_key_admin_access( + user_api_key_dict=user_api_key_dict, + hashed_token=hashed_key, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + route="/key/update (max_budget)", + ) + # Check team limits if key has a team_id (from request or existing key) team_obj: Optional[LiteLLM_TeamTableCachedObj] = None _team_id_to_check = data.team_id or getattr(existing_key_row, "team_id", None) @@ -1847,6 +1928,13 @@ async def _validate_update_key_data( check_db_only=True, ) + # Validate team exists when non-admin sets a new team_id (LIT-1884) + if team_obj is None and data.team_id is not None and not _is_proxy_admin: + raise HTTPException( + status_code=400, + detail=f"Team not found for team_id={data.team_id}. Non-admin users cannot set keys to non-existent teams.", + ) + if team_obj is not None: await _check_team_key_limits( team_table=team_obj, @@ -2056,7 +2144,10 @@ async def update_key_fn( data=data, existing_key_row=existing_key_row ) - _validate_key_alias_format(key_alias=non_default_values.get("key_alias", None)) + # Only validate key_alias format if it's actually being changed + new_key_alias = non_default_values.get("key_alias", None) + if new_key_alias != existing_key_row.key_alias: + _validate_key_alias_format(key_alias=new_key_alias) await _enforce_unique_key_alias( key_alias=non_default_values.get("key_alias", None), @@ -3412,8 +3503,10 @@ async def _rotate_master_key( # noqa: PLR0915 ) -def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: +async def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: if data and data.new_key is not None: + # Reject custom key values if disabled by admin + await _check_custom_key_allowed(data.new_key) new_token = data.new_key if not data.new_key.startswith("sk-"): raise HTTPException( @@ -3505,7 +3598,7 @@ async def _execute_virtual_key_regeneration( """Generate new token, update DB, invalidate cache, and return response.""" from litellm.proxy.proxy_server import hash_token - new_token = get_new_token(data=data) + new_token = await get_new_token(data=data) new_token_hash = hash_token(new_token) new_token_key_name = f"sk-...{new_token[-4:]}" update_data = {"token": new_token_hash, "key_name": new_token_key_name} @@ -3515,7 +3608,10 @@ async def _execute_virtual_key_regeneration( non_default_values = await prepare_key_update_data( data=data, existing_key_row=key_in_db ) - _validate_key_alias_format(key_alias=non_default_values.get("key_alias")) + # Only validate key_alias format if it's actually being changed + new_key_alias = non_default_values.get("key_alias") + if new_key_alias != key_in_db.key_alias: + _validate_key_alias_format(key_alias=new_key_alias) verbose_proxy_logger.debug("non_default_values: %s", non_default_values) update_data.update(non_default_values) update_data = prisma_client.jsonify_object(data=update_data) @@ -4733,6 +4829,64 @@ def _get_condition_to_filter_out_ui_session_tokens() -> Dict[str, Any]: } +async def _check_key_admin_access( + user_api_key_dict: UserAPIKeyAuth, + hashed_token: str, + prisma_client: Any, + user_api_key_cache: DualCache, + route: str, +) -> None: + """ + Check that the caller has admin privileges for the target key. + + Allowed callers: + - Proxy admin + - Team admin for the key's team + - Org admin for the key's team's organization + + Raises HTTPException(403) if the caller is not authorized. + """ + + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + + # Look up the target key to find its team + target_key_row = await prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": hashed_token} + ) + if target_key_row is None: + raise HTTPException( + status_code=404, + detail={"error": f"Key not found: {hashed_token}"}, + ) + + # If the key belongs to a team, check team admin / org admin + if target_key_row.team_id: + team_obj = await get_team_object( + team_id=target_key_row.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + if team_obj is not None: + if _is_user_team_admin( + user_api_key_dict=user_api_key_dict, team_obj=team_obj + ): + return + if await _is_user_org_admin_for_team( + user_api_key_dict=user_api_key_dict, team_obj=team_obj + ): + return + + raise HTTPException( + status_code=403, + detail={ + "error": f"Only proxy admins, team admins, or org admins can call {route}. " + f"user_role={user_api_key_dict.user_role}, user_id={user_api_key_dict.user_id}" + }, + ) + + @router.post( "/key/block", tags=["key management"], dependencies=[Depends(user_api_key_auth)] ) @@ -4762,7 +4916,7 @@ async def block_key( }' ``` - Note: This is an admin-only endpoint. Only proxy admins can block keys. + Note: This is an admin-only endpoint. Only proxy admins, team admins, or org admins can block keys. """ from litellm.proxy.proxy_server import ( create_audit_log_for_update, @@ -4788,6 +4942,15 @@ async def block_key( else: hashed_token = data.key + # Admin-only: only proxy admins, team admins, or org admins can block keys + await _check_key_admin_access( + user_api_key_dict=user_api_key_dict, + hashed_token=hashed_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + route="/key/block", + ) + if litellm.store_audit_logs is True: # make an audit log for key update record = await prisma_client.db.litellm_verificationtoken.find_unique( @@ -4876,7 +5039,7 @@ async def unblock_key( }' ``` - Note: This is an admin-only endpoint. Only proxy admins can unblock keys. + Note: This is an admin-only endpoint. Only proxy admins, team admins, or org admins can unblock keys. """ from litellm.proxy.proxy_server import ( create_audit_log_for_update, @@ -4902,6 +5065,15 @@ async def unblock_key( else: hashed_token = data.key + # Admin-only: only proxy admins, team admins, or org admins can unblock keys + await _check_key_admin_access( + user_api_key_dict=user_api_key_dict, + hashed_token=hashed_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + route="/key/unblock", + ) + if litellm.store_audit_logs is True: # make an audit log for key update record = await prisma_client.db.litellm_verificationtoken.find_unique( diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 701762c834c..e5a34ae8bdd 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -340,9 +340,16 @@ def _maybe_setup_prometheus_multiproc_dir( return # Check if prometheus is in any callback list + # Each setting can be a list or a single string; normalize to list callbacks = litellm_settings.get("callbacks") or [] success_callbacks = litellm_settings.get("success_callback") or [] failure_callbacks = litellm_settings.get("failure_callback") or [] + if isinstance(callbacks, str): + callbacks = [callbacks] + if isinstance(success_callbacks, str): + success_callbacks = [success_callbacks] + if isinstance(failure_callbacks, str): + failure_callbacks = [failure_callbacks] all_callbacks = callbacks + success_callbacks + failure_callbacks if "prometheus" not in all_callbacks: return diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 0fa27905bab..60bf41709ef 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -129,6 +129,11 @@ class UISettings(BaseModel): description="If enabled, the user search endpoint (/user/filter/ui) restricts results by organization. When off, any authenticated user can search all users.", ) + disable_custom_api_keys: bool = Field( + default=False, + description="If true, users cannot specify custom key values. All keys must be auto-generated.", + ) + class UISettingsResponse(SettingsResponse): """Response model for UI settings""" @@ -149,6 +154,7 @@ class UISettingsResponse(SettingsResponse): "disable_vector_stores_for_internal_users", "allow_vector_stores_for_team_admins", "scope_user_search_to_org", + "disable_custom_api_keys", } # Flags that must be synced from the persisted UISettings into diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index d5abf5c8fbf..27fa27e6da3 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -44,6 +44,7 @@ class SupportedGuardrailIntegrations(Enum): APORIA = "aporia" BEDROCK = "bedrock" + DYNAMOAI = "dynamoai" GUARDRAILS_AI = "guardrails_ai" LAKERA = "lakera" LAKERA_V2 = "lakera_v2" diff --git a/litellm/types/router.py b/litellm/types/router.py index e8ff2115ff5..26bd381431a 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -193,6 +193,8 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): budget_duration: Optional[str] = None use_in_pass_through: Optional[bool] = False use_litellm_proxy: Optional[bool] = False + ## BASETEN ## + served_model_name: Optional[str] = None # override model name in request body (e.g. Baseten dedicated deployments) model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True) merge_reasoning_content_in_choices: Optional[bool] = False model_info: Optional[Dict] = None @@ -345,6 +347,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): # deployment budgets max_budget: Optional[float] budget_duration: Optional[str] + ## BASETEN ## + served_model_name: Optional[str] # override model name in request body (e.g. Baseten dedicated deployments) class DeploymentTypedDict(TypedDict, total=False): diff --git a/pyproject.toml b/pyproject.toml index 8efc6366128..07004f3ae16 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "1.82.2" +version = "1.82.3" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT" @@ -183,7 +183,7 @@ requires = ["poetry-core", "wheel"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "1.82.2" +version = "1.82.3" version_files = [ "pyproject.toml:^version" ] diff --git a/tests/local_testing/test_router_utils.py b/tests/local_testing/test_router_utils.py index 9d51685751a..4f7f53cef02 100644 --- a/tests/local_testing/test_router_utils.py +++ b/tests/local_testing/test_router_utils.py @@ -199,6 +199,7 @@ def test_router_get_model_info_wildcard_routes(): @pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=1) async def test_router_get_model_group_usage_wildcard_routes(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") @@ -219,7 +220,7 @@ async def test_router_get_model_group_usage_wildcard_routes(): ) print(resp) - await asyncio.sleep(1) + await asyncio.sleep(2) tpm, rpm = await router.get_model_group_usage(model_group="gemini/gemini-1.5-flash") diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 2544e06598e..a4a28215e16 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -395,6 +395,7 @@ async def test_mcp_http_transport_tool_not_found(): @pytest.mark.asyncio async def test_streamable_http_mcp_handler_mock(): """Test the streamable HTTP MCP handler functionality""" + from litellm.proxy._types import UserAPIKeyAuth # Mock the session manager and its methods mock_session_manager = AsyncMock() @@ -425,6 +426,8 @@ async def test_streamable_http_mcp_handler_mock(): ), patch( "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", AsyncMock(return_value=mock_auth_context), + ), patch( + "litellm.proxy._experimental.mcp_server.server.set_auth_context", ): from litellm.proxy._experimental.mcp_server.server import ( handle_streamable_http_mcp, diff --git a/tests/test_litellm/llms/baseten/chat/test_baseten_completions.py b/tests/test_litellm/llms/baseten/chat/test_baseten_completions.py index 9420149a8e4..1dc7a7b91fe 100644 --- a/tests/test_litellm/llms/baseten/chat/test_baseten_completions.py +++ b/tests/test_litellm/llms/baseten/chat/test_baseten_completions.py @@ -10,10 +10,10 @@ class TestBasetenRouting: def test_routing_logic(self): """Test routing between Model API and dedicated deployments""" config = BasetenConfig() - + # Dedicated deployment (8-character alphanumeric) assert config.get_api_base_for_model("abcd1234") == "https://model-abcd1234.api.baseten.co/environments/production/sync/v1" - + # Model API (non-8-character) assert config.get_api_base_for_model("openai/gpt-oss-120b") == "https://inference.baseten.co/v1" @@ -25,30 +25,116 @@ class TestBasetenModelAPI: def test_model_api_inference(self): """Test Model API inference with basic parameters""" config = BasetenConfig() - + # Test parameter mapping non_default_params = { "max_tokens": 100, "temperature": 0.7, "top_p": 0.9 } - + result = config.map_openai_params( non_default_params=non_default_params, optional_params={}, model="openai/gpt-oss-120b", drop_params=False ) - + assert result["max_tokens"] == 100 assert result["temperature"] == 0.7 assert result["top_p"] == 0.9 - + # Test provider info api_base, api_key = config._get_openai_compatible_provider_info(None, "test-key") assert api_base == "https://inference.baseten.co/v1" assert api_key == "test-key" + def test_model_api_transform_request(self): + """ + Model API happy path — no served_model_name, model passes through as-is. + + Proxy config: + model_list: + - model_name: baseten-model + litellm_params: + model: baseten/openai/gpt-oss-120b + api_key: your-baseten-api-key + """ + config = BasetenConfig() + + result = config.transform_request( + model="openai/gpt-oss-120b", + messages=[{"role": "user", "content": "Hello!"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert result["model"] == "openai/gpt-oss-120b" + + +class TestBasetenTransformRequest: + """Test Baseten transform_request for dedicated deployments""" + + def test_dedicated_deployment_with_served_model_name(self): + """ + Customer fix: dedicated deployment ID used for URL routing, + served_model_name sent in request body. + + Proxy config: + model_list: + - model_name: baseten-model + litellm_params: + model: baseten/wd1lndkw + served_model_name: baseten-hosted/zai-org/GLM-5 + api_key: os.environ/BASETEN_API_KEY + """ + config = BasetenConfig() + + result = config.transform_request( + model="wd1lndkw", + messages=[{"role": "user", "content": "Hello!"}], + optional_params={}, + litellm_params={"served_model_name": "baseten-hosted/zai-org/GLM-5"}, + headers={}, + ) + + assert result["model"] == "baseten-hosted/zai-org/GLM-5" + assert result["messages"] == [{"role": "user", "content": "Hello!"}] + + def test_dedicated_deployment_without_served_model_name(self): + """ + Dedicated deployment without served_model_name — deployment ID passes + through as the model name in the request body. Only works if the + deployment's served_model_name matches the deployment ID. + + Proxy config: + model_list: + - model_name: baseten-model + litellm_params: + model: baseten/wd1lndkw + api_key: os.environ/BASETEN_API_KEY + """ + config = BasetenConfig() + + result = config.transform_request( + model="wd1lndkw", + messages=[{"role": "user", "content": "Hello!"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert result["model"] == "wd1lndkw" + + def test_dedicated_deployment_api_base_routing(self): + """ + Dedicated deployment ID correctly builds the dedicated endpoint URL. + """ + config = BasetenConfig() + + assert config.get_api_base_for_model("wd1lndkw") == "https://model-wd1lndkw.api.baseten.co/environments/production/sync/v1" + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 3a01fe19edb..3acbe5465f2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -158,7 +158,6 @@ async def ok_operation(client): @pytest.mark.asyncio - @pytest.mark.skip(reason="PR #23187 changed has_client_credentials to require explicit oauth2_flow opt-in, but NewMCPServerRequest and _execute_with_mcp_client were not updated - needs fix") async def test_m2m_credentials_forwarded_to_server_model(self, monkeypatch): """M2M OAuth credentials (client_id, client_secret) from the nested ``credentials`` dict must be forwarded to the MCPServer model so that @@ -213,7 +212,6 @@ async def ok_operation(client): assert server.has_client_credentials is True @pytest.mark.asyncio - @pytest.mark.skip(reason="PR #23187 changed has_client_credentials to require explicit oauth2_flow opt-in, but NewMCPServerRequest and _execute_with_mcp_client were not updated - needs fix") async def test_m2m_drops_incoming_oauth2_headers(self, monkeypatch): """For M2M OAuth servers the incoming Authorization header (which carries the litellm API key) must NOT be forwarded as extra_headers — otherwise diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_dynamoai.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_dynamoai.py new file mode 100644 index 00000000000..7bc4e951a5f --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_dynamoai.py @@ -0,0 +1,81 @@ +""" +Tests for DynamoAI guardrail registration and initialization. +""" + +import os +from unittest.mock import patch + +import pytest + + +class TestDynamoAIGuardrailRegistration: + """Tests for DynamoAI guardrail registration in the guardrail system.""" + + def test_supported_guardrail_enum_entry(self): + """Test that DYNAMOAI is in SupportedGuardrailIntegrations enum.""" + from litellm.types.guardrails import SupportedGuardrailIntegrations + + assert hasattr(SupportedGuardrailIntegrations, "DYNAMOAI") + assert SupportedGuardrailIntegrations.DYNAMOAI.value == "dynamoai" + + def test_initialize_guardrail_function_exists(self): + """Test that initialize_guardrail function is properly exported.""" + from litellm.proxy.guardrails.guardrail_hooks.dynamoai import ( + guardrail_initializer_registry, + initialize_guardrail, + ) + + assert initialize_guardrail is not None + assert "dynamoai" in guardrail_initializer_registry + + def test_guardrail_class_registry_exists(self): + """Test that guardrail_class_registry is properly exported.""" + from litellm.proxy.guardrails.guardrail_hooks.dynamoai import ( + guardrail_class_registry, + ) + from litellm.proxy.guardrails.guardrail_hooks.dynamoai.dynamoai import ( + DynamoAIGuardrails, + ) + + assert "dynamoai" in guardrail_class_registry + assert guardrail_class_registry["dynamoai"] == DynamoAIGuardrails + + def test_initialize_guardrail_creates_instance(self): + """Test that initialize_guardrail creates a DynamoAIGuardrails instance.""" + from litellm.proxy.guardrails.guardrail_hooks.dynamoai import ( + initialize_guardrail, + ) + from litellm.proxy.guardrails.guardrail_hooks.dynamoai.dynamoai import ( + DynamoAIGuardrails, + ) + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams( + guardrail="dynamoai", + mode="pre_call", + api_key="test-key", + api_base="https://test.dynamo.ai", + ) + + guardrail = { + "guardrail_name": "test-dynamoai-guard", + } + + with patch( + "litellm.logging_callback_manager.add_litellm_callback" + ) as mock_add: + result = initialize_guardrail(litellm_params, guardrail) + + assert isinstance(result, DynamoAIGuardrails) + assert result.api_key == "test-key" + assert result.api_base == "https://test.dynamo.ai" + assert result.guardrail_name == "test-dynamoai-guard" + mock_add.assert_called_once_with(result) + + def test_dynamoai_in_global_registry(self): + """Test that dynamoai is discoverable in the global guardrail registry.""" + from litellm.proxy.guardrails.guardrail_registry import ( + guardrail_initializer_registry, + ) + + assert "dynamoai" in guardrail_initializer_registry diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index cfc16808afb..a3bca77ae3a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -41,12 +41,15 @@ _transform_verification_tokens_to_deleted_records, _validate_max_budget, _validate_reset_spend_value, + _validate_update_key_data, can_modify_verification_token, check_org_key_model_specific_limits, check_team_key_model_specific_limits, delete_verification_tokens, + generate_key_fn, generate_key_helper_fn, key_aliases, + key_generation_check, list_keys, prepare_key_update_data, reset_key_spend_fn, @@ -957,22 +960,34 @@ async def test_key_info_returns_object_permission(monkeypatch): ) -def test_get_new_token_with_valid_key(): +@pytest.mark.asyncio +async def test_get_new_token_with_valid_key(monkeypatch): """Test get_new_token function when provided with a valid key that starts with 'sk-'""" + from unittest.mock import AsyncMock + from litellm.proxy._types import RegenerateKeyRequest from litellm.proxy.management_endpoints.key_management_endpoints import ( get_new_token, ) + # Mock get_ui_settings_cached to return setting disabled (custom keys allowed) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={}), + ) + # Test with valid new_key data = RegenerateKeyRequest(new_key="sk-test123456789") - result = get_new_token(data) + result = await get_new_token(data) assert result == "sk-test123456789" -def test_get_new_token_with_invalid_key(): +@pytest.mark.asyncio +async def test_get_new_token_with_invalid_key(monkeypatch): """Test get_new_token function when provided with an invalid key that doesn't start with 'sk-'""" + from unittest.mock import AsyncMock + from fastapi import HTTPException from litellm.proxy._types import RegenerateKeyRequest @@ -980,16 +995,145 @@ def test_get_new_token_with_invalid_key(): get_new_token, ) + # Mock get_ui_settings_cached to return setting disabled (custom keys allowed) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={}), + ) + # Test with invalid new_key (doesn't start with 'sk-') data = RegenerateKeyRequest(new_key="invalid-key-123") with pytest.raises(HTTPException) as exc_info: - get_new_token(data) + await get_new_token(data) assert exc_info.value.status_code == 400 assert "New key must start with 'sk-'" in str(exc_info.value.detail) +@pytest.mark.asyncio +async def test_check_custom_key_allowed_when_disabled(monkeypatch): + """_check_custom_key_allowed raises 403 when disable_custom_api_keys is true.""" + from unittest.mock import AsyncMock + + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _check_custom_key_allowed, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={"disable_custom_api_keys": True}), + ) + + with pytest.raises(HTTPException) as exc_info: + await _check_custom_key_allowed("sk-custom-key-123") + + assert exc_info.value.status_code == 403 + assert "disabled" in str(exc_info.value.detail).lower() + + +@pytest.mark.asyncio +async def test_check_custom_key_allowed_when_enabled(monkeypatch): + """_check_custom_key_allowed does nothing when disable_custom_api_keys is false.""" + from unittest.mock import AsyncMock + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _check_custom_key_allowed, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={"disable_custom_api_keys": False}), + ) + + # Should not raise + await _check_custom_key_allowed("sk-custom-key-123") + + +@pytest.mark.asyncio +async def test_check_custom_key_allowed_when_unset(monkeypatch): + """_check_custom_key_allowed does nothing when setting is not present.""" + from unittest.mock import AsyncMock + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _check_custom_key_allowed, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={}), + ) + + # Should not raise + await _check_custom_key_allowed("sk-custom-key-123") + + +@pytest.mark.asyncio +async def test_check_custom_key_allowed_none_key_always_passes(monkeypatch): + """_check_custom_key_allowed does nothing when key is None, even if setting is on.""" + from unittest.mock import AsyncMock + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _check_custom_key_allowed, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={"disable_custom_api_keys": True}), + ) + + # Should not raise — None means auto-generate + await _check_custom_key_allowed(None) + + +@pytest.mark.asyncio +async def test_get_new_token_rejected_when_custom_keys_disabled(monkeypatch): + """get_new_token raises 403 when new_key is set and disable_custom_api_keys is true.""" + from unittest.mock import AsyncMock + + from fastapi import HTTPException + + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + get_new_token, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={"disable_custom_api_keys": True}), + ) + + data = RegenerateKeyRequest(new_key="sk-custom-regen-key") + + with pytest.raises(HTTPException) as exc_info: + await get_new_token(data) + + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_get_new_token_auto_generates_when_custom_keys_disabled(monkeypatch): + """get_new_token auto-generates a key when new_key is None, even if setting is on.""" + from unittest.mock import AsyncMock + + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + get_new_token, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={"disable_custom_api_keys": True}), + ) + + data = RegenerateKeyRequest() # no new_key + result = await get_new_token(data) + + assert result.startswith("sk-") + + @pytest.mark.asyncio async def test_generate_service_account_requires_team_id(): with pytest.raises(HTTPException): @@ -7185,3 +7329,773 @@ def test_update_key_request_has_organization_id(): # Also verify it defaults to None data_no_org = UpdateKeyRequest(key="sk-test-key") assert data_no_org.organization_id is None + + +# ============================================================================ +# Tests for admin-only access on /key/block, /key/unblock, /key/update max_budget +# ============================================================================ + + +def _setup_block_unblock_mocks(monkeypatch, mock_key_team_id=None): + """Helper to set up common mocks for block/unblock tests.""" + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + mock_proxy_logging_obj = MagicMock() + + test_hashed_token = ( + "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + ) + + mock_key_record = MagicMock() + mock_key_record.token = test_hashed_token + mock_key_record.blocked = False + mock_key_record.team_id = mock_key_team_id + mock_key_record.model_dump_json.return_value = ( + f'{{"token": "{test_hashed_token}", "blocked": false}}' + ) + + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=mock_key_record + ) + mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock( + return_value=mock_key_record + ) + + mock_key_object = MagicMock() + mock_key_object.blocked = True + + def mock_hash_token(token): + if token.startswith("sk-"): + return test_hashed_token + return token + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj + ) + monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token) + monkeypatch.setattr("litellm.store_audit_logs", False) + + async def mock_get_key_object(**kwargs): + return mock_key_object + + async def mock_cache_key_object(**kwargs): + pass + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_key_object", + mock_get_key_object, + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._cache_key_object", + mock_cache_key_object, + ) + + return mock_prisma_client, test_hashed_token + + +@pytest.mark.asyncio +async def test_block_key_rejected_for_internal_user(monkeypatch): + """Internal users should not be able to block keys.""" + from litellm.proxy._types import BlockKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import block_key + + _setup_block_unblock_mocks(monkeypatch) + + mock_request = MagicMock() + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ) + + with pytest.raises(HTTPException) as exc: + await block_key( + data=BlockKeyRequest(key="sk-test123456789"), + http_request=mock_request, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + assert exc.value.status_code == 403 + assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_unblock_key_rejected_for_internal_user(monkeypatch): + """Internal users should not be able to unblock keys.""" + from litellm.proxy._types import BlockKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import unblock_key + + _setup_block_unblock_mocks(monkeypatch) + + mock_request = MagicMock() + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ) + + with pytest.raises(HTTPException) as exc: + await unblock_key( + data=BlockKeyRequest(key="sk-test123456789"), + http_request=mock_request, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + assert exc.value.status_code == 403 + assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_block_key_allowed_for_proxy_admin(monkeypatch): + """Proxy admins should be able to block keys.""" + from litellm.proxy._types import BlockKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import block_key + + _setup_block_unblock_mocks(monkeypatch) + + mock_request = MagicMock() + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin_user", + ) + + result = await block_key( + data=BlockKeyRequest(key="sk-test123456789"), + http_request=mock_request, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_block_key_allowed_for_team_admin(monkeypatch): + """Team admins should be able to block keys belonging to their team.""" + from litellm.proxy._types import BlockKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import block_key + + team_id = "team-123" + _setup_block_unblock_mocks(monkeypatch, mock_key_team_id=team_id) + + # Mock get_team_object to return a team where the user is admin + team_obj = LiteLLM_TeamTableCachedObj( + team_id=team_id, + members_with_roles=[ + Member(user_id="team_admin_user", role="admin"), + ], + ) + + async def mock_get_team_object(**kwargs): + return team_obj + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + mock_get_team_object, + ) + + mock_request = MagicMock() + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-teamadmin", + user_id="team_admin_user", + ) + + result = await block_key( + data=BlockKeyRequest(key="sk-test123456789"), + http_request=mock_request, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_update_key_max_budget_rejected_for_internal_user(monkeypatch): + """Internal users should not be able to modify max_budget on keys.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = AsyncMock() + mock_proxy_logging_obj = MagicMock() + + test_hashed_token = ( + "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + ) + + # Mock existing key row + mock_existing_key = MagicMock() + mock_existing_key.token = test_hashed_token + mock_existing_key.user_id = "internal_user" + mock_existing_key.team_id = None + mock_existing_key.project_id = None + mock_existing_key.max_budget = 10.0 + mock_existing_key.models = [] + mock_existing_key.model_dump.return_value = { + "token": test_hashed_token, + "user_id": "internal_user", + "team_id": None, + "max_budget": 10.0, + } + + mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=mock_existing_key + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + + mock_request = MagicMock() + mock_request.query_params = {} + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ) + + with pytest.raises(ProxyException) as exc: + await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key=test_hashed_token, max_budget=999999), + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + assert str(exc.value.code) == "403" + assert "Only proxy admins, team admins, or org admins" in str(exc.value.message) + + +@pytest.mark.asyncio +async def test_update_key_non_budget_fields_allowed_for_internal_user(monkeypatch): + """Internal users should still be able to update non-budget fields on their own keys.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = AsyncMock() + mock_proxy_logging_obj = MagicMock() + + test_hashed_token = ( + "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + ) + + # Mock existing key row + mock_existing_key = MagicMock() + mock_existing_key.token = test_hashed_token + mock_existing_key.user_id = "internal_user" + mock_existing_key.team_id = None + mock_existing_key.project_id = None + mock_existing_key.max_budget = 10.0 + mock_existing_key.key_alias = None + mock_existing_key.models = [] + mock_existing_key.model_dump.return_value = { + "token": test_hashed_token, + "user_id": "internal_user", + "team_id": None, + "max_budget": 10.0, + } + + mock_updated_key = MagicMock() + mock_updated_key.token = test_hashed_token + mock_updated_key.key_alias = "my-alias" + + mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key) + mock_prisma_client.update_data = AsyncMock(return_value=mock_updated_key) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=mock_existing_key + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.store_audit_logs", False) + + def mock_hash_token(token): + return test_hashed_token + + monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token) + + async def mock_cache_key_object(**kwargs): + pass + + async def mock_delete_cache_key_object(**kwargs): + pass + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._cache_key_object", + mock_cache_key_object, + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + mock_delete_cache_key_object, + ) + + # Mock _enforce_unique_key_alias to avoid DB call + async def mock_enforce_unique_key_alias(**kwargs): + pass + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._enforce_unique_key_alias", + mock_enforce_unique_key_alias, + ) + + mock_request = MagicMock() + mock_request.query_params = {} + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ) + + # Updating key_alias (non-budget field) should succeed + result = await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key=test_hashed_token, key_alias="my-alias"), + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + assert result is not None + + +# ============================================================================ +# LIT-1884: Internal users cannot create invalid keys +# ============================================================================ + + +class TestLIT1884KeyGenerateValidation: + """Tests for LIT-1884: internal users should not be able to generate invalid keys.""" + + @pytest.mark.asyncio + async def test_internal_user_generate_key_no_user_id_auto_assigns(self): + """ + When an internal_user calls /key/generate without user_id, + the caller's user_id should be auto-assigned before reaching + _common_key_generation_helper. + """ + mock_prisma_client = AsyncMock() + + data = GenerateKeyRequest(key_alias="test-alias") + assert data.user_id is None + + user_api_key_dict = UserAPIKeyAuth( + user_id="internal-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + # Patch _common_key_generation_helper to avoid needing full DB mocks. + # We just want to verify user_id is set before we reach this point. + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), \ + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), \ + patch("litellm.proxy.proxy_server.user_custom_key_generate", None), \ + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper", + new_callable=AsyncMock, + return_value=MagicMock(), + ): + await generate_key_fn( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + # The data object should have been mutated to include the caller's user_id + assert data.user_id == "internal-user-123" + + @pytest.mark.asyncio + async def test_internal_user_generate_key_invalid_team_id_rejected(self): + """ + When an internal_user provides a non-existent team_id, + key/generate should raise ProxyException with status 400. + """ + mock_prisma_client = AsyncMock() + + data = GenerateKeyRequest( + key_alias="test-alias", + team_id="nonexistent-team-id", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_id="internal-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), \ + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), \ + patch("litellm.proxy.proxy_server.user_custom_key_generate", None), \ + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + AsyncMock(side_effect=Exception("Team not found")), + ): + with pytest.raises(ProxyException) as exc_info: + await generate_key_fn( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + assert str(exc_info.value.code) == "400" + assert "Team not found" in str(exc_info.value.message) + + @pytest.mark.asyncio + async def test_admin_generate_key_invalid_team_id_allowed(self): + """ + Admin callers should be allowed to create keys with any team_id, + even if the team doesn't exist (team_table=None is OK for admins). + """ + data = GenerateKeyRequest( + key_alias="admin-key", + team_id="nonexistent-team-id", + user_id="admin-user", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_id="admin-user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + mock_prisma_client = AsyncMock() + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), \ + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), \ + patch("litellm.proxy.proxy_server.user_custom_key_generate", None), \ + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + AsyncMock(side_effect=Exception("Team not found")), + ), \ + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper", + new_callable=AsyncMock, + return_value=MagicMock(), + ): + # Should NOT raise — admin bypasses team validation + result = await generate_key_fn( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + assert result is not None + + @pytest.mark.asyncio + async def test_admin_generate_key_no_user_id_not_auto_assigned(self): + """ + Admin callers should NOT have user_id auto-assigned — they may + intentionally create keys without a user_id. + """ + data = GenerateKeyRequest(key_alias="admin-unbound-key") + assert data.user_id is None + + user_api_key_dict = UserAPIKeyAuth( + user_id="admin-user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + mock_prisma_client = AsyncMock() + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), \ + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), \ + patch("litellm.proxy.proxy_server.user_custom_key_generate", None), \ + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper", + new_callable=AsyncMock, + return_value=MagicMock(), + ): + await generate_key_fn( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + # user_id should remain None for admin + assert data.user_id is None + + def test_key_generation_check_non_admin_no_team_table_raises(self): + """ + key_generation_check should raise 400 for non-admin when team_table is None + and key_generation_settings is not set. + """ + data = GenerateKeyRequest(team_id="some-team-id") + user_api_key_dict = UserAPIKeyAuth( + user_id="internal-user", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + with patch.object(litellm, "key_generation_settings", None): + with pytest.raises(HTTPException) as exc_info: + key_generation_check( + team_table=None, + user_api_key_dict=user_api_key_dict, + data=data, + route="key_generate", + ) + assert exc_info.value.status_code == 400 + assert "Unable to find team object" in str(exc_info.value.detail) + + def test_key_generation_check_admin_no_team_table_allowed(self): + """ + key_generation_check should allow admin to proceed even when team_table is None. + """ + data = GenerateKeyRequest(team_id="some-team-id") + user_api_key_dict = UserAPIKeyAuth( + user_id="admin-user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + with patch.object(litellm, "key_generation_settings", None): + result = key_generation_check( + team_table=None, + user_api_key_dict=user_api_key_dict, + data=data, + route="key_generate", + ) + assert result is True + + +class TestLIT1884KeyUpdateValidation: + """Tests for LIT-1884: internal users should not be able to update keys to remove user_id or set invalid team.""" + + @pytest.mark.asyncio + async def test_internal_user_cannot_remove_user_id(self): + """ + Non-admin users should not be able to set user_id to empty string (remove it). + """ + data = UpdateKeyRequest(key="sk-test-key", user_id="") + existing_key_row = MagicMock() + existing_key_row.user_id = "internal-user-123" + existing_key_row.token = "hashed_token" + existing_key_row.team_id = None + + user_api_key_dict = UserAPIKeyAuth( + user_id="internal-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + with pytest.raises(HTTPException) as exc_info: + await _validate_update_key_data( + data=data, + existing_key_row=existing_key_row, + user_api_key_dict=user_api_key_dict, + llm_router=None, + premium_user=False, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + assert exc_info.value.status_code == 403 + assert "cannot remove the user_id" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_internal_user_cannot_set_invalid_team_id(self): + """ + Non-admin users should not be able to update a key to a non-existent team. + get_team_object raises HTTPException(404) when team doesn't exist in DB. + """ + data = UpdateKeyRequest(key="sk-test-key", team_id="nonexistent-team") + existing_key_row = MagicMock() + existing_key_row.user_id = "internal-user-123" + existing_key_row.token = "hashed_token" + existing_key_row.team_id = None + existing_key_row.organization_id = None + existing_key_row.project_id = None + + user_api_key_dict = UserAPIKeyAuth( + user_id="internal-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + AsyncMock(side_effect=HTTPException( + status_code=404, + detail="Team doesn't exist in db. Team=nonexistent-team.", + )), + ): + with pytest.raises(HTTPException) as exc_info: + await _validate_update_key_data( + data=data, + existing_key_row=existing_key_row, + user_api_key_dict=user_api_key_dict, + llm_router=None, + premium_user=False, + prisma_client=AsyncMock(), + user_api_key_cache=MagicMock(), + ) + assert exc_info.value.status_code == 404 + assert "Team doesn't exist" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_admin_can_remove_user_id(self): + """ + Admin users should be allowed to set user_id to empty string. + """ + data = UpdateKeyRequest(key="sk-test-key", user_id="") + existing_key_row = MagicMock() + existing_key_row.user_id = "some-user" + existing_key_row.token = "hashed_token" + existing_key_row.team_id = None + existing_key_row.organization_id = None + existing_key_row.project_id = None + + user_api_key_dict = UserAPIKeyAuth( + user_id="admin-user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + mock_prisma_client = AsyncMock() + + # Should NOT raise + await _validate_update_key_data( + data=data, + existing_key_row=existing_key_row, + user_api_key_dict=user_api_key_dict, + llm_router=None, + premium_user=False, + prisma_client=mock_prisma_client, + user_api_key_cache=MagicMock(), + ) + + +class TestKeyAliasSkipValidationOnUnchanged: + """ + Test that updating/regenerating a key without changing its key_alias + does NOT re-validate the alias. This prevents legacy aliases (created + before stricter validation rules) from blocking edits to other fields. + """ + + @pytest.fixture(autouse=True) + def enable_validation(self): + litellm.enable_key_alias_format_validation = True + yield + litellm.enable_key_alias_format_validation = False + + @pytest.fixture + def mock_prisma(self): + prisma = MagicMock() + prisma.db = MagicMock() + prisma.db.litellm_verificationtoken = MagicMock() + prisma.get_data = AsyncMock(return_value=None) # no duplicate alias + prisma.update_data = AsyncMock(return_value=None) + prisma.jsonify_object = MagicMock(side_effect=lambda data: data) + return prisma + + @pytest.fixture + def existing_key_with_legacy_alias(self): + """A key whose alias contains '@' — valid now, but simulates a legacy alias.""" + return LiteLLM_VerificationToken( + token="hashed_token_123", + key_alias="user@domain.com", + team_id="team-1", + models=[], + max_budget=100.0, + ) + + @pytest.mark.asyncio + async def test_update_key_unchanged_legacy_alias_passes( + self, mock_prisma, existing_key_with_legacy_alias + ): + """ + Updating a key without changing its key_alias should skip format + validation — even if the alias wouldn't pass current rules. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_key_alias_format, + ) + + # Temporarily make the regex reject '@' to simulate stricter rules + import re + from litellm.proxy.management_endpoints import key_management_endpoints as mod + + original_pattern = mod._KEY_ALIAS_PATTERN + mod._KEY_ALIAS_PATTERN = re.compile( + r"^[a-zA-Z0-9][a-zA-Z0-9_\-/\.]{0,253}[a-zA-Z0-9]$" + ) + try: + # Confirm the alias WOULD fail validation directly + with pytest.raises(ProxyException): + _validate_key_alias_format("user@domain.com") + + # But prepare_key_update_data + the skip logic should allow it + # Simulate what update_key_fn does: alias is in non_default_values + # but matches existing_key_row.key_alias => skip validation + existing_alias = existing_key_with_legacy_alias.key_alias + new_alias = "user@domain.com" # same as existing + assert new_alias == existing_alias # unchanged + + # This is the core logic from update_key_fn: + if new_alias != existing_alias: + _validate_key_alias_format(new_alias) + # No exception raised — test passes + finally: + mod._KEY_ALIAS_PATTERN = original_pattern + + @pytest.mark.asyncio + async def test_update_key_changed_alias_still_validated( + self, mock_prisma, existing_key_with_legacy_alias + ): + """ + When the alias IS being changed, validation should still run. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_key_alias_format, + ) + + existing_alias = existing_key_with_legacy_alias.key_alias + new_alias = "!invalid!" + + assert new_alias != existing_alias + with pytest.raises(ProxyException): + if new_alias != existing_alias: + _validate_key_alias_format(new_alias) + + @pytest.mark.asyncio + async def test_update_key_changed_to_valid_alias_passes( + self, mock_prisma, existing_key_with_legacy_alias + ): + """ + Changing the alias to a new valid value should pass validation. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_key_alias_format, + ) + + existing_alias = existing_key_with_legacy_alias.key_alias + new_alias = "new-valid-alias" + + assert new_alias != existing_alias + # Should not raise + if new_alias != existing_alias: + _validate_key_alias_format(new_alias) + + @pytest.mark.asyncio + async def test_update_key_alias_none_skips_validation(self): + """ + When key_alias is not in the update payload (None), validation + should be skipped regardless. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_key_alias_format, + ) + + # None alias should always pass + _validate_key_alias_format(None) diff --git a/tests/test_litellm/proxy/test_prometheus_cleanup.py b/tests/test_litellm/proxy/test_prometheus_cleanup.py index b3d785f1133..0a67d5e64e0 100644 --- a/tests/test_litellm/proxy/test_prometheus_cleanup.py +++ b/tests/test_litellm/proxy/test_prometheus_cleanup.py @@ -67,6 +67,30 @@ def test_respects_existing_env_var(self, tmp_path): assert os.environ["PROMETHEUS_MULTIPROC_DIR"] == custom_dir assert os.path.isdir(custom_dir) + @pytest.mark.parametrize( + "litellm_settings", + [ + {"callbacks": "prometheus"}, + {"success_callback": "prometheus"}, + {"failure_callback": "prometheus"}, + {"callbacks": "custom_callback"}, # string but not prometheus + ], + ) + def test_handles_string_callbacks(self, litellm_settings): + """When callbacks are specified as a string instead of a list, should not crash.""" + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + os.environ.pop("prometheus_multiproc_dir", None) + + # Should not raise TypeError + ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir( + num_workers=4, + litellm_settings=litellm_settings, + ) + + # Cleanup + os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + @pytest.mark.parametrize( "num_workers, litellm_settings", [ diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 642d21a42f7..c5d6c45f9a5 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -677,7 +677,7 @@ def test_startup_fails_when_db_setup_fails( mock_atexit_register, mock_subprocess_run, ): - """Test that proxy exits with code 1 when PrismaManager.setup_database returns False""" + """Test that proxy exits with code 1 when PrismaManager.setup_database returns False and --enforce_prisma_migration_check is set""" from litellm.proxy.proxy_cli import run_server mock_subprocess_run.return_value = MagicMock(returncode=0) @@ -717,7 +717,7 @@ def test_startup_fails_when_db_setup_fails( with pytest.raises(SystemExit) as exc_info: run_server.main( - ["--local", "--skip_server_startup"], standalone_mode=False + ["--local", "--skip_server_startup", "--enforce_prisma_migration_check"], standalone_mode=False ) assert exc_info.value.code == 1 mock_setup_database.assert_called_once_with(use_migrate=True) diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 7378b14cdf5..bd9968ae936 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -111,6 +111,37 @@ def test_get_internal_user_settings(self, mock_proxy_config, mock_auth): assert "user_role" in data["field_schema"]["properties"] assert "description" in data["field_schema"]["properties"]["user_role"] + def test_get_internal_user_settings_fresh_db_defaults_to_viewer( + self, mock_auth, monkeypatch + ): + """ + On a fresh DB with no saved settings, the GET endpoint should return + INTERNAL_USER_VIEW_ONLY as the default role — matching the runtime + fallback in SSO/SCIM/JWT provisioning paths. + """ + # Simulate fresh DB: no default_internal_user_params in config + empty_config = { + "litellm_settings": {}, + "general_settings": {}, + "environment_variables": {}, + } + + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return empty_config + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + + response = client.get("/get/internal_user_settings") + assert response.status_code == 200 + + values = response.json()["values"] + assert values["user_role"] == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, ( + f"Fresh DB should default to INTERNAL_USER_VIEW_ONLY, got {values['user_role']}. " + "The Pydantic default must match the runtime fallback." + ) + def test_update_internal_user_settings( self, mock_proxy_config, mock_auth, monkeypatch ): diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 463f6952d23..8f5c3ece0ca 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -388,6 +388,65 @@ def test_custom_pricing_cost_calc_uses_router_model_id_from_litellm_metadata(): assert custom_model_id not in (selected_model_no_custom or "") +def test_per_request_custom_pricing_with_router(): + """When custom pricing is passed as per-request kwargs (not in model_list), + _select_model_name_for_cost_calc should fall back to the model name + (where register_model stored the pricing) instead of the router_model_id + (which has no pricing data). + + Regression test for the bug where response._hidden_params["response_cost"] + returned 0.0 for per-request custom pricing via Router. + """ + from litellm import Router + from litellm.cost_calculator import _select_model_name_for_cost_calc + + router = Router( + model_list=[ + { + "model_name": "openai/gpt-3.5-turbo", + "litellm_params": { + "model": "openai/gpt-3.5-turbo", + "api_key": "test_api_key", + }, + }, + ] + ) + + # Get the deployment's model_id (hash) that the router registered + deployment = router.model_list[0] + router_model_id = deployment["model_info"]["id"] + + # The router registered this hash in model_cost but without custom pricing + assert router_model_id in litellm.model_cost + entry = litellm.model_cost[router_model_id] + # No custom pricing was set in model_list, so these should be None + assert entry.get("input_cost_per_token") is None + + # Now simulate what completion() does: register custom pricing under the model name + litellm.register_model( + { + "openai/gpt-3.5-turbo": { + "input_cost_per_token": 2.0, + "output_cost_per_token": 2.0, + "litellm_provider": "openai", + } + } + ) + + # _select_model_name_for_cost_calc should pick the model name (which has pricing), + # NOT the router_model_id (which has no pricing) + selected = _select_model_name_for_cost_calc( + model="openai/gpt-3.5-turbo", + completion_response=None, + custom_pricing=True, + custom_llm_provider="openai", + router_model_id=router_model_id, + ) + assert selected is not None + assert router_model_id not in selected + assert "gpt-3.5-turbo" in selected + + def test_azure_realtime_cost_calculator(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") diff --git a/tests/test_litellm/test_get_blog_posts.py b/tests/test_litellm/test_get_blog_posts.py index a17d78e0bb6..b04fb4ec703 100644 --- a/tests/test_litellm/test_get_blog_posts.py +++ b/tests/test_litellm/test_get_blog_posts.py @@ -1,5 +1,4 @@ """Tests for GetBlogPosts utility class.""" -import json import time from unittest.mock import MagicMock, patch @@ -13,16 +12,26 @@ get_blog_posts, ) -SAMPLE_RESPONSE = { - "posts": [ - { - "title": "Test Post", - "description": "A test post.", - "date": "2026-01-01", - "url": "https://www.litellm.ai/blog/test", - } - ] -} +SAMPLE_RSS = """\ + + + + LiteLLM Blog + + Test Post + https://docs.litellm.ai/blog/test + A test post. + Wed, 01 Jan 2026 10:00:00 GMT + + + Second Post + https://docs.litellm.ai/blog/second + Another post. + Tue, 31 Dec 2025 10:00:00 GMT + + + +""" @pytest.fixture(autouse=True) @@ -45,26 +54,48 @@ def test_load_local_blog_posts_returns_list(): assert "url" in first -def test_validate_blog_posts_valid(): - assert GetBlogPosts.validate_blog_posts(SAMPLE_RESPONSE) is True +def test_parse_rss_to_posts(): + posts = GetBlogPosts.parse_rss_to_posts(SAMPLE_RSS, max_posts=1) + assert len(posts) == 1 + assert posts[0]["title"] == "Test Post" + assert posts[0]["url"] == "https://docs.litellm.ai/blog/test" + assert posts[0]["description"] == "A test post." + assert posts[0]["date"] == "2026-01-01" + + +def test_parse_rss_to_posts_multiple(): + posts = GetBlogPosts.parse_rss_to_posts(SAMPLE_RSS, max_posts=5) + assert len(posts) == 2 + assert posts[1]["title"] == "Second Post" -def test_validate_blog_posts_missing_posts_key(): - assert GetBlogPosts.validate_blog_posts({"other": []}) is False +def test_parse_rss_to_posts_invalid_xml(): + with pytest.raises(Exception): + GetBlogPosts.parse_rss_to_posts("not xml") + + +def test_parse_rss_to_posts_missing_channel(): + with pytest.raises(ValueError, match="missing "): + GetBlogPosts.parse_rss_to_posts("") + + +def test_validate_blog_posts_valid(): + posts = [{"title": "T", "description": "D", "date": "2026-01-01", "url": "https://x.com"}] + assert GetBlogPosts.validate_blog_posts(posts) is True def test_validate_blog_posts_empty_list(): - assert GetBlogPosts.validate_blog_posts({"posts": []}) is False + assert GetBlogPosts.validate_blog_posts([]) is False -def test_validate_blog_posts_not_dict(): - assert GetBlogPosts.validate_blog_posts("not a dict") is False +def test_validate_blog_posts_not_list(): + assert GetBlogPosts.validate_blog_posts("not a list") is False def test_get_blog_posts_success(): - """Fetches from remote on first call.""" + """Fetches from RSS on first call.""" mock_response = MagicMock() - mock_response.json.return_value = SAMPLE_RESPONSE + mock_response.text = SAMPLE_RSS mock_response.raise_for_status = MagicMock() with patch("litellm.litellm_core_utils.get_blog_posts.httpx.get", return_value=mock_response): @@ -86,10 +117,10 @@ def test_get_blog_posts_network_error_falls_back_to_local(): assert len(posts) > 0 -def test_get_blog_posts_invalid_json_falls_back_to_local(): - """Falls back when remote returns non-dict.""" +def test_get_blog_posts_invalid_xml_falls_back_to_local(): + """Falls back when remote returns invalid XML.""" mock_response = MagicMock() - mock_response.json.return_value = "not a dict" + mock_response.text = "not valid xml" mock_response.raise_for_status = MagicMock() with patch("litellm.litellm_core_utils.get_blog_posts.httpx.get", return_value=mock_response): @@ -101,7 +132,8 @@ def test_get_blog_posts_invalid_json_falls_back_to_local(): def test_get_blog_posts_ttl_cache_not_refetched(): """Within TTL window, does not re-fetch.""" - GetBlogPosts._cached_posts = SAMPLE_RESPONSE["posts"] + cached = [{"title": "Cached", "description": "D", "date": "2026-01-01", "url": "https://x.com"}] + GetBlogPosts._cached_posts = cached GetBlogPosts._last_fetch_time = time.time() # just now call_count = 0 @@ -110,7 +142,7 @@ def mock_get(*args, **kwargs): nonlocal call_count call_count += 1 m = MagicMock() - m.json.return_value = SAMPLE_RESPONSE + m.text = SAMPLE_RSS m.raise_for_status = MagicMock() return m @@ -123,11 +155,12 @@ def mock_get(*args, **kwargs): def test_get_blog_posts_ttl_expired_refetches(): """After TTL window, re-fetches from remote.""" - GetBlogPosts._cached_posts = SAMPLE_RESPONSE["posts"] + cached = [{"title": "Cached", "description": "D", "date": "2026-01-01", "url": "https://x.com"}] + GetBlogPosts._cached_posts = cached GetBlogPosts._last_fetch_time = time.time() - 7200 # 2 hours ago mock_response = MagicMock() - mock_response.json.return_value = SAMPLE_RESPONSE + mock_response.text = SAMPLE_RSS mock_response.raise_for_status = MagicMock() with patch( diff --git a/ui/litellm-dashboard/src/components/DefaultUserSettings.tsx b/ui/litellm-dashboard/src/components/DefaultUserSettings.tsx index 988a3bcec92..314946c520b 100644 --- a/ui/litellm-dashboard/src/components/DefaultUserSettings.tsx +++ b/ui/litellm-dashboard/src/components/DefaultUserSettings.tsx @@ -1,6 +1,6 @@ import React, { useState, useEffect } from "react"; -import { Card, Title, Text, Divider, Button, TextInput } from "@tremor/react"; -import { Typography, Spin, Switch, Select, InputNumber } from "antd"; +import { Card, Title, Text, Divider, TextInput } from "@tremor/react"; +import { Button, Typography, Spin, Switch, Select, InputNumber } from "antd"; import { PlusOutlined, DeleteOutlined } from "@ant-design/icons"; import { getInternalUserSettings, updateInternalUserSettings, modelAvailableCall } from "./networking"; import BudgetDurationDropdown, { getBudgetDurationLabel } from "./common_components/budget_duration_dropdown"; @@ -160,11 +160,10 @@ const DefaultUserSettings: React.FC = ({
Team {index + 1} @@ -208,7 +207,7 @@ const DefaultUserSettings: React.FC = ({
))} - @@ -462,7 +461,6 @@ const DefaultUserSettings: React.FC = ({ (isEditing ? (
-
) : ( - + ))} diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts index 67de4b69174..2ca7ad7ef31 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts @@ -11,6 +11,7 @@ import { getEntityBreakdown, handleExportCSV, handleExportJSON, + resolveEntities, } from "./utils"; vi.mock("@/utils/dataUtils", () => ({ @@ -1561,4 +1562,137 @@ describe("EntityUsageExport utils", () => { window.Blob = originalBlob; }); }); + + describe("resolveEntities and aggregated endpoint fallback", () => { + // Simulates the response from /user/daily/activity/aggregated which has + // empty entities but populated api_keys at the breakdown level. + // Derived from mockSpendData: flatten all entities' api_key_breakdowns + // into top-level api_keys, clear entities, and add a second key for team-1 + // to test multi-key grouping. + const aggregatedSpendData: EntitySpendData = { + ...mockSpendData, + results: mockSpendData.results.slice(0, 1).map((day) => ({ + ...day, + breakdown: { + entities: {}, + api_keys: { + ...Object.fromEntries( + Object.values(day.breakdown.entities as Record).flatMap((e: any) => + Object.entries(e.api_key_breakdown || {}), + ), + ), + // Extra key on team-1 to test multi-key-per-team aggregation + key1b: { + metrics: { spend: 5, api_requests: 50, successful_requests: 48, failed_requests: 2, total_tokens: 500 }, + metadata: { team_id: "team-1", key_alias: "staging-key" }, + }, + }, + models: { "gpt-4": { metrics: { spend: 35, api_requests: 350, total_tokens: 3500 } } }, + }, + })), + }; + + describe("resolveEntities", () => { + it("should return entities when populated", () => { + const breakdown = { + entities: { e1: { metrics: { spend: 1 } } }, + api_keys: { k1: { metrics: { spend: 2 }, metadata: { team_id: "t1" } } }, + }; + const result = resolveEntities(breakdown); + expect(result).toBe(breakdown.entities); + }); + + it("should aggregate api_keys into entities when entities is empty", () => { + const breakdown = aggregatedSpendData.results[0].breakdown; + const result = resolveEntities(breakdown); + + // Two teams: team-1 (key1+key2) and team-2 (key3) + expect(Object.keys(result)).toHaveLength(2); + expect(result["team-1"]).toBeDefined(); + expect(result["team-2"]).toBeDefined(); + + // team-1 spend = 10.5 (key1) + 5 (key1b) + expect(result["team-1"].metrics.spend).toBe(15.5); + expect(result["team-1"].metrics.api_requests).toBe(150); + expect(result["team-1"].metrics.total_tokens).toBe(1500); + + // team-2 spend = 20.3 (key2) + expect(result["team-2"].metrics.spend).toBe(20.3); + expect(result["team-2"].metrics.api_requests).toBe(200); + }); + + it("should use 'Unassigned' for keys without team_id", () => { + const breakdown = { + entities: {}, + api_keys: { + k1: { + metrics: { spend: 7, api_requests: 10, successful_requests: 10, failed_requests: 0, total_tokens: 100 }, + metadata: {}, + }, + }, + }; + const result = resolveEntities(breakdown); + expect(result["Unassigned"]).toBeDefined(); + expect(result["Unassigned"].metrics.spend).toBe(7); + }); + + it("should handle missing or empty api_keys gracefully", () => { + expect(Object.keys(resolveEntities({ entities: {}, api_keys: {} }))).toHaveLength(0); + expect(Object.keys(resolveEntities({ entities: {} }))).toHaveLength(0); + }); + + it("should preserve api_key_breakdown on aggregated entities", () => { + const breakdown = aggregatedSpendData.results[0].breakdown; + const result = resolveEntities(breakdown); + + // team-1 should have key1 and key1b in api_key_breakdown + expect(Object.keys(result["team-1"].api_key_breakdown)).toEqual(["key1", "key1b"]); + // team-2 should have key2 + expect(Object.keys(result["team-2"].api_key_breakdown)).toEqual(["key2"]); + }); + }); + + describe("getEntityBreakdown with aggregated data", () => { + it("should produce breakdown from api_keys when entities is empty", () => { + const result = getEntityBreakdown(aggregatedSpendData); + expect(result.length).toBeGreaterThan(0); + + // Sorted by spend desc: team-2 (20.3) then team-1 (15.5) + expect(result[0].metrics.spend).toBe(20.3); + expect(result[1].metrics.spend).toBe(15.5); + }); + + }); + + describe("generateDailyData with aggregated data", () => { + it("should produce rows from api_keys when entities is empty", () => { + const result = generateDailyData(aggregatedSpendData, "Team"); + expect(result.length).toBeGreaterThan(0); + expect(result[0]).toHaveProperty("Date"); + expect(result[0]).toHaveProperty("Team"); + }); + }); + + describe("generateDailyWithKeysData with aggregated data", () => { + it("should produce rows from api_keys when entities is empty", () => { + const result = generateDailyWithKeysData(aggregatedSpendData, "Team"); + expect(result.length).toBeGreaterThan(0); + + // Should have 3 key rows (key1, key1b, key2) + expect(result).toHaveLength(3); + const keyIds = result.map((r) => r["Key ID"]); + expect(keyIds).toContain("key1"); + expect(keyIds).toContain("key1b"); + expect(keyIds).toContain("key2"); + }); + }); + + describe("generateDailyWithModelsData with aggregated data", () => { + it("should produce rows from api_keys when entities is empty", () => { + const result = generateDailyWithModelsData(aggregatedSpendData, "Team"); + expect(result.length).toBeGreaterThan(0); + expect(result[0]).toHaveProperty("Model"); + }); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts index ebef3da8a77..45bf21a6e7d 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts @@ -17,6 +17,49 @@ const extractTeamIdFromApiKeyBreakdown = (apiKeyBreakdown: Record | return null; }; +// Mirrors backend SpendMetrics fields (litellm/types/activity_tracking.py). +// If the backend adds a field, add it here too. +const METRIC_KEYS = [ + "spend", "api_requests", "successful_requests", "failed_requests", + "total_tokens", "prompt_tokens", "completion_tokens", + "cache_read_input_tokens", "cache_creation_input_tokens", +] as const; + +// When breakdown.entities is empty (aggregated endpoint), reconstruct entities +// from breakdown.api_keys by grouping on metadata.team_id. +const aggregateApiKeysIntoEntities = (breakdown: Record): Record => { + const apiKeys = breakdown.api_keys; + if (!apiKeys || Object.keys(apiKeys).length === 0) return {}; + + const grouped: Record = {}; + + for (const [keyId, keyData] of Object.entries(apiKeys)) { + const teamId = keyData?.metadata?.team_id || "Unassigned"; + if (!grouped[teamId]) { + grouped[teamId] = { + metrics: Object.fromEntries(METRIC_KEYS.map((k) => [k, 0])), + api_key_breakdown: {}, + }; + } + const m = grouped[teamId].metrics; + const km = keyData?.metrics || {}; + for (const k of METRIC_KEYS) { + m[k] += km[k] || 0; + } + grouped[teamId].api_key_breakdown[keyId] = keyData; + } + + return grouped; +}; + +// Returns breakdown.entities if populated, otherwise falls back to +// reconstructing entities from breakdown.api_keys. +export const resolveEntities = (breakdown: Record): Record => { + const entities = breakdown.entities; + if (entities && Object.keys(entities).length > 0) return entities; + return aggregateApiKeysIntoEntities(breakdown); +}; + export const getEntityBreakdown = ( spendData: EntitySpendData, teamAliasMap: Record = {}, @@ -24,7 +67,7 @@ export const getEntityBreakdown = ( const entitySpend: { [key: string]: EntityBreakdown } = {}; spendData.results.forEach((day) => { - Object.entries(day.breakdown.entities || {}).forEach(([entity, data]: [string, any]) => { + Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { // Extract team_id from api_key_breakdown metadata (not data.metadata which is empty) const teamId = extractTeamIdFromApiKeyBreakdown(data.api_key_breakdown) || entity; // Extract key_alias from the first API key that has one @@ -80,7 +123,7 @@ export const generateDailyData = ( const dailyBreakdown: any[] = []; spendData.results.forEach((day) => { - Object.entries(day.breakdown.entities || {}).forEach(([entity, data]: [string, any]) => { + Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { // Extract team_id from api_key_breakdown metadata (not data.metadata which is empty) const teamId = extractTeamIdFromApiKeyBreakdown(data.api_key_breakdown); const teamAlias = teamId ? teamAliasMap[teamId] || null : null; @@ -129,7 +172,7 @@ export const generateDailyWithKeysData = ( } = {}; spendData.results.forEach((day) => { - Object.entries(day.breakdown.entities || {}).forEach(([entity, data]: [string, any]) => { + Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { const apiKeyBreakdown = data.api_key_breakdown || {}; // Iterate through each API key in the breakdown @@ -202,7 +245,7 @@ export const generateDailyWithModelsData = ( spendData.results.forEach((day) => { const dailyEntityModels: { [key: string]: { [key: string]: any } } = {}; - Object.entries(day.breakdown.entities || {}).forEach(([entity, entityData]: [string, any]) => { + Object.entries(resolveEntities(day.breakdown)).forEach(([entity, entityData]: [string, any]) => { if (!dailyEntityModels[entity]) { dailyEntityModels[entity] = {}; } @@ -230,7 +273,7 @@ export const generateDailyWithModelsData = ( }); Object.entries(dailyEntityModels).forEach(([entity, models]) => { - const entityData = day.breakdown.entities?.[entity]; + const entityData = resolveEntities(day.breakdown)[entity]; // Extract team_id from api_key_breakdown metadata (not entityData.metadata which is empty) const teamId = extractTeamIdFromApiKeyBreakdown(entityData?.api_key_breakdown); const teamAlias = teamId ? teamAliasMap[teamId] || null : null; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/HashicorpVault/HashicorpVaultEmptyPlaceholder.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/HashicorpVault/HashicorpVaultEmptyPlaceholder.test.tsx new file mode 100644 index 00000000000..4214d5fda7c --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/HashicorpVault/HashicorpVaultEmptyPlaceholder.test.tsx @@ -0,0 +1,29 @@ +import { describe, it, expect, vi } from "vitest"; +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import HashicorpVaultEmptyPlaceholder from "./HashicorpVaultEmptyPlaceholder"; + +describe("HashicorpVaultEmptyPlaceholder", () => { + it("should render the empty state message and configure button", () => { + render(); + expect(screen.getByText("No Vault Configuration Found")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /configure vault/i })).toBeInTheDocument(); + }); + + it("should call onAdd when the configure button is clicked", async () => { + const onAdd = vi.fn(); + const user = userEvent.setup(); + render(); + + await user.click(screen.getByRole("button", { name: /configure vault/i })); + + expect(onAdd).toHaveBeenCalledOnce(); + }); + + it("should display the description text about Vault purpose", () => { + render(); + expect( + screen.getByText(/Configure Hashicorp Vault to securely manage provider API keys/), + ).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/PageVisibilitySettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/PageVisibilitySettings.test.tsx new file mode 100644 index 00000000000..798572b2037 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/PageVisibilitySettings.test.tsx @@ -0,0 +1,77 @@ +import { describe, it, expect, vi } from "vitest"; +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import PageVisibilitySettings from "./PageVisibilitySettings"; + +vi.mock("@/components/page_utils", () => ({ + getAvailablePages: () => [ + { page: "usage", label: "Usage", description: "View usage stats", group: "Analytics" }, + { page: "models", label: "Models", description: "Manage models", group: "Analytics" }, + { page: "keys", label: "API Keys", description: "Manage API keys", group: "Access" }, + ], +})); + +describe("PageVisibilitySettings", () => { + it("should render the not-set tag when enabledPagesInternalUsers is null", () => { + render( + , + ); + expect(screen.getByText("Not set (all pages visible)")).toBeInTheDocument(); + }); + + it("should show the selected page count tag when pages are configured", () => { + render( + , + ); + expect(screen.getByText("2 pages selected")).toBeInTheDocument(); + }); + + it("should show singular 'page' when exactly one page is selected", () => { + render( + , + ); + expect(screen.getByText("1 page selected")).toBeInTheDocument(); + }); + + it("should call onUpdate with null when reset button is clicked", async () => { + const onUpdate = vi.fn(); + const user = userEvent.setup(); + render( + , + ); + + // Expand the collapse panel first to reveal the reset button + await user.click(screen.getByRole("button", { name: /configure page visibility/i })); + await user.click(await screen.findByRole("button", { name: /reset to default/i })); + + expect(onUpdate).toHaveBeenCalledWith({ enabled_ui_pages_internal_users: null }); + }); + + it("should display the property description when provided", () => { + render( + , + ); + expect(screen.getByText("Controls which pages are visible")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index a22c78c9430..fb7c38449b0 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -24,6 +24,7 @@ export default function UISettings() { const disableVectorStoresProperty = schema?.properties?.disable_vector_stores_for_internal_users; const allowVectorStoresTeamAdminsProperty = schema?.properties?.allow_vector_stores_for_team_admins; const scopeUserSearchProperty = schema?.properties?.scope_user_search_to_org; + const disableCustomApiKeysProperty = schema?.properties?.disable_custom_api_keys; const values = data?.values ?? {}; const isDisabledForInternalUsers = Boolean(values.disable_model_add_for_internal_users); const isDisabledTeamAdminDeleteTeamUser = Boolean(values.disable_team_admin_delete_team_user); @@ -182,6 +183,20 @@ export default function UISettings() { ); }; + const handleToggleDisableCustomApiKeys = (checked: boolean) => { + updateSettings( + { disable_custom_api_keys: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + return ( {isLoading ? ( @@ -382,6 +397,26 @@ export default function UISettings() { + {/* Disable custom Virtual key values */} + + + + Disable custom Virtual key values + + {disableCustomApiKeysProperty?.description ?? + "If true, users cannot specify custom key values. All keys must be auto-generated."} + + + + + + {/* Page Visibility for Internal Users */} = ({ team, teams, data, addKey, autoOp const { data: projects, isLoading: isProjectsLoading } = useProjects(); const { data: uiSettingsData } = useUISettings(); const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui); + const disableCustomApiKeys = Boolean(uiSettingsData?.values?.disable_custom_api_keys); const queryClient = useQueryClient(); const [form] = Form.useForm(); const [isModalVisible, setIsModalVisible] = useState(false); @@ -1581,6 +1582,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp "budget_duration", "tpm_limit", "rpm_limit", + ...(disableCustomApiKeys ? ["key"] : []), ]} /> diff --git a/ui/litellm-dashboard/src/components/ui/ui-loading-spinner.test.tsx b/ui/litellm-dashboard/src/components/ui/ui-loading-spinner.test.tsx new file mode 100644 index 00000000000..73b2de54c4c --- /dev/null +++ b/ui/litellm-dashboard/src/components/ui/ui-loading-spinner.test.tsx @@ -0,0 +1,22 @@ +import { describe, it, expect } from "vitest"; +import { render, screen } from "@testing-library/react"; +import { UiLoadingSpinner } from "./ui-loading-spinner"; + +describe("UiLoadingSpinner", () => { + it("should render an SVG element", () => { + render(); + expect(screen.getByTestId("spinner")).toBeInTheDocument(); + }); + + it("should apply custom className alongside default classes", () => { + render(); + const svg = screen.getByTestId("spinner"); + expect(svg).toHaveClass("text-red-500"); + expect(svg).toHaveClass("animate-spin"); + }); + + it("should spread additional SVG props onto the element", () => { + render(); + expect(screen.getByLabelText("Loading")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx index 0bb0f2a44d4..8ccf0a9e1b5 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx @@ -451,7 +451,7 @@ describe("useLogFilterLogic", () => { ); }); - it("should fall back to logs when backend filters are active but API returns empty", async () => { + it("should return empty results when backend filters are active but API returns empty", async () => { vi.mocked(uiSpendLogsCall).mockResolvedValue({ data: [], total: 0, @@ -474,8 +474,7 @@ describe("useLogFilterLogic", () => { { timeout: 500 }, ); - expect(result.current.filteredLogs.data).toHaveLength(1); - expect(result.current.filteredLogs.data[0].request_id).toBe("client-req"); + expect(result.current.filteredLogs.data).toHaveLength(0); }); it("should refetch when sortBy changes and backend filters are active", async () => { diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx index 400e86d19ee..4e7153b64cd 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx @@ -228,7 +228,7 @@ export function useLogFilterLogic({ const filteredLogs: PaginatedResponse = useMemo(() => { if (hasBackendFilters) { // Prefer backend result if present; otherwise fall back to latest logs - if (backendFilteredLogs && backendFilteredLogs.data && backendFilteredLogs.data.length > 0) { + if (backendFilteredLogs && backendFilteredLogs.data) { return backendFilteredLogs; } return ( diff --git a/ui/litellm-dashboard/src/utils/errorUtils.test.ts b/ui/litellm-dashboard/src/utils/errorUtils.test.ts new file mode 100644 index 00000000000..ccd65717f40 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/errorUtils.test.ts @@ -0,0 +1,40 @@ +import { describe, it, expect } from "vitest"; +import { extractErrorMessage } from "./errorUtils"; + +describe("extractErrorMessage", () => { + it("should return the message from an Error instance", () => { + expect(extractErrorMessage(new Error("Something broke"))).toBe("Something broke"); + }); + + it("should return detail when it is a string", () => { + expect(extractErrorMessage({ detail: "Not found" })).toBe("Not found"); + }); + + it("should join msg fields from a FastAPI 422 detail array", () => { + const err = { + detail: [ + { msg: "field required", loc: ["body", "name"], type: "value_error" }, + { msg: "invalid type", loc: ["body", "age"], type: "type_error" }, + ], + }; + expect(extractErrorMessage(err)).toBe("field required; invalid type"); + }); + + it("should extract error from nested detail object", () => { + expect(extractErrorMessage({ detail: { error: "bad request" } })).toBe("bad request"); + }); + + it("should fall back to message property on plain objects", () => { + expect(extractErrorMessage({ message: "fallback msg" })).toBe("fallback msg"); + }); + + it("should JSON.stringify unknown object shapes", () => { + expect(extractErrorMessage({ foo: "bar" })).toBe('{"foo":"bar"}'); + }); + + it("should stringify primitive non-object values", () => { + expect(extractErrorMessage(42)).toBe("42"); + expect(extractErrorMessage(null)).toBe("null"); + expect(extractErrorMessage(undefined)).toBe("undefined"); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/mcpToolCrudClassification.test.ts b/ui/litellm-dashboard/src/utils/mcpToolCrudClassification.test.ts new file mode 100644 index 00000000000..0ae8ae70bec --- /dev/null +++ b/ui/litellm-dashboard/src/utils/mcpToolCrudClassification.test.ts @@ -0,0 +1,63 @@ +import { describe, it, expect } from "vitest"; +import { classifyToolOp, groupToolsByCrud } from "./mcpToolCrudClassification"; + +describe("classifyToolOp", () => { + it("should classify read operations by name", () => { + expect(classifyToolOp("get-users")).toBe("read"); + expect(classifyToolOp("list-items")).toBe("read"); + expect(classifyToolOp("search documents")).toBe("read"); + }); + + it("should classify delete operations by name", () => { + expect(classifyToolOp("delete-user")).toBe("delete"); + expect(classifyToolOp("remove-item")).toBe("delete"); + expect(classifyToolOp("purge-cache")).toBe("delete"); + }); + + it("should classify create operations by name", () => { + expect(classifyToolOp("create-user")).toBe("create"); + expect(classifyToolOp("add-item")).toBe("create"); + expect(classifyToolOp("upload-file")).toBe("create"); + }); + + it("should classify update operations by name", () => { + expect(classifyToolOp("update-settings")).toBe("update"); + expect(classifyToolOp("edit-profile")).toBe("update"); + expect(classifyToolOp("rename-file")).toBe("update"); + }); + + it("should prioritize read over delete for names like get-removed-entries", () => { + expect(classifyToolOp("get-removed-entries")).toBe("read"); + expect(classifyToolOp("list-deleted-items")).toBe("read"); + }); + + it("should fall back to description when name is unrecognised", () => { + expect(classifyToolOp("mytool", "This will delete the record")).toBe("delete"); + expect(classifyToolOp("mytool", "fetch data from the API")).toBe("read"); + }); + + it("should return unknown when neither name nor description match", () => { + expect(classifyToolOp("my_tool")).toBe("unknown"); + expect(classifyToolOp("my_tool", "does something")).toBe("unknown"); + }); +}); + +describe("groupToolsByCrud", () => { + it("should group tools into their CRUD categories", () => { + const tools = [ + { name: "get-user", description: "" }, + { name: "create-item", description: "" }, + { name: "delete-record", description: "" }, + { name: "update-settings", description: "" }, + { name: "mysteryop", description: "" }, + ]; + + const groups = groupToolsByCrud(tools); + + expect(groups.read).toHaveLength(1); + expect(groups.create).toHaveLength(1); + expect(groups.delete).toHaveLength(1); + expect(groups.update).toHaveLength(1); + expect(groups.unknown).toHaveLength(1); + }); +});