mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix: support served_model_name for Baseten dedicated deployments (#23382)
* fix: langfuse trace leak key on model params * fix: pop sensitive keys from langfuse * fixes * fix: set oauth2_flow when building MCPServer in _execute_with_mcp_client * fix: add oauth2_flow to NewMCPServerRequest and guard auto-detect with token_url * fix: narrow oauth2_flow type to Literal in NewMCPServerRequest * fix: align DefaultInternalUserParams Pydantic default with runtime fallback The Pydantic default for user_role was INTERNAL_USER, but all runtime provisioning paths (SSO, SCIM, JWT) fall back to INTERNAL_USER_VIEW_ONLY when no settings are saved. This caused the UI to show "Internal User" on fresh instances while new users actually got "Internal Viewer". * test: add regression test for fresh-instance default role sync Asserts that GET /get/internal_user_settings returns INTERNAL_USER_VIEW_ONLY on a fresh DB with no saved settings, matching the runtime fallback in SSO/SCIM/JWT provisioning. * Update tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Add unit tests for 5 previously untested UI dashboard files Tests added for: UiLoadingSpinner, HashicorpVaultEmptyPlaceholder, PageVisibilitySettings, errorUtils, and mcpToolCrudClassification. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: remove skip decorators from m2m tests now that oauth2_flow is set * [Fix] Privilege escalation: restrict /key/block, /key/unblock, and max_budget updates to admins Non-admin users (INTERNAL_USER) could call /key/block and /key/unblock on arbitrary keys, and modify max_budget on their own keys via /key/update. These endpoints are now restricted to proxy admins, team admins, or org admins. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * chore(ui): migrate DefaultUserSettings buttons from Tremor to antd * [Infra] Merging RC Branch with Main (#23786) * fix(test): add missing mocks for test_streamable_http_mcp_handler_mock The test was missing mocks for extract_mcp_auth_context and set_auth_context, causing the handler to fail silently in the except block instead of reaching session_manager.handle_request. This mirrors the fix already applied to the sibling test_sse_mcp_handler_mock. Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> * fix(ci): route OpenAI models through chat completions in pass-through tests The test_anthropic_messages_openai_model_streaming_cost_injection test fails because the OpenAI Responses API returns 400 for requests routed through the Anthropic Messages endpoint. Setting LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES=true routes OpenAI models through the stable chat completions path instead. Cost injection still works since it happens at the proxy level. Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> * fix(ci): fix assemblyai custom auth and router wildcard test flakiness 1. custom_auth_basic.py: Add user_role='proxy_admin' so the custom auth user can access management endpoints like /key/generate. The test test_assemblyai_transcribe_with_non_admin_key was hidden behind an earlier -x failure and was never reached before. 2. test_router_utils.py: Add flaky(retries=3) and increase sleep from 1s to 2s for test_router_get_model_group_usage_wildcard_routes. The async callback needs time to write usage to cache, and 1s is insufficient on slower CI hardware. Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> * ci: retrigger CI pipeline Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> * fix(mypy): use LitellmUserRoles enum instead of raw string in custom_auth_basic Fixes mypy error: Argument 'user_role' has incompatible type 'str'; expected 'LitellmUserRoles | None' Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> * fix: don't close HTTP/SDK clients on LLMClientCache eviction (#22926) * fix: don't close HTTP/SDK clients on LLMClientCache eviction Removing the _remove_key override that eagerly called aclose()/close() on evicted clients. Evicted clients may still be held by in-flight streaming requests; closing them causes: RuntimeError: Cannot send a request, as the client has been closed. This is a regression from commitfb72979432. Clients that are no longer referenced will be garbage-collected naturally. Explicit shutdown cleanup happens via close_litellm_async_clients(). Fixes production crashes after the 1-hour cache TTL expires. * test: update LLMClientCache unit tests for no-close-on-eviction behavior Flip the assertions: evicted clients must NOT be closed. Replace test_remove_key_closes_async_client → test_remove_key_does_not_close_async_client and equivalents for sync/eviction paths. Add test_remove_key_removes_plain_values for non-client cache entries. Remove test_background_tasks_cleaned_up_after_completion (no more _background_tasks). Remove test_remove_key_no_event_loop variant that depended on old behavior. * test: add e2e tests for OpenAI SDK client surviving cache eviction Add two new e2e tests using real AsyncOpenAI clients: - test_evicted_openai_sdk_client_stays_usable: verifies size-based eviction doesn't close the client - test_ttl_expired_openai_sdk_client_stays_usable: verifies TTL expiry eviction doesn't close the client Both tests sleep after eviction so any create_task()-based close would have time to run, making the regression detectable. Also expand the module docstring to explain why the sleep is required. * docs(AGENTS.md): add rule — never close HTTP/SDK clients on cache eviction * docs(CLAUDE.md): add HTTP client cache safety guideline * [Fix] Install bsdmainutils for column command in security scans The security_scans.sh script uses `column` to format vulnerability output, but the package wasn't installed in the CI environment. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: handle string callback values in prometheus multiproc setup When callbacks are configured as a plain string (e.g., `callbacks: "my_callback"`) instead of a list, the proxy crashes on startup with: TypeError: can only concatenate str (not "list") to str Normalize each callback setting to a list before concatenating. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * bump: version 1.82.2 → 1.82.3 * fix(test): update test_startup_fails_when_db_setup_fails for opt-in enforcement The --enforce_prisma_migration_check flag is now required to trigger sys.exit(1) on DB migration failure, after #23675 flipped the default behavior to warn-and-continue. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix(cost_calculator): use model name for per-request custom pricing when router_model_id has no pricing When custom pricing is passed as per-request kwargs (input_cost_per_token/output_cost_per_token), completion() registers pricing under the model name, but _select_model_name_for_cost_calc was selecting the router deployment hash (which has no pricing data), causing response_cost to be 0.0. Now checks whether the router_model_id entry actually has pricing before preferring it. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Cursor Agent <cursoragent@cursor.com> Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> * Update litellm/proxy/management_endpoints/key_management_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix: clear oauth2_flow when client_credentials set without token_url * chore(ui): use antd danger prop instead of tailwind for Remove button * feat: fetch blog posts from docs RSS feed instead of static JSON on GitHub * fix: remove unused Any import from get_blog_posts * [Fix] UI - Logs: Fix empty filter results showing stale data Remove `.length > 0` check so that when a backend filter returns an empty result set the table correctly shows no data instead of falling back to the previous unfiltered logs. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * [Fix] Reapply empty filter fix after merge with main Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * [Fix] Prevent internal users from creating invalid keys via key/generate and key/update Internal users could exploit key/generate and key/update to create unbound keys (no user_id, no budget) or attach keys to non-existent teams. This adds validation for non-admin callers: auto-assign user_id on generate, reject invalid team_ids, and prevent removing user_id on update. Closes LIT-1884 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * [Fix] Remove duplicate get_team_object call in _validate_update_key_data Move the non-admin team validation into the existing get_team_object call site to avoid an extra DB round-trip. The existing call already fetches the team for limits checking — we now add the LIT-1884 guard there when team_obj is None for non-admin callers. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * [Fix] Skip key_alias re-validation on update/regenerate when alias unchanged When updating or regenerating a key without changing its key_alias, the existing alias was being re-validated against current format rules. This caused keys with legacy aliases (created before stricter validation) to become uneditable. Now validation only runs when the alias actually changes. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * [Fix] Update log filter test to match empty-result behavior The test expected fallback to all logs when backend filters return empty, but the source was intentionally changed to show empty results instead of stale data. Updated test to match. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * [Feature] Disable custom API key values via UI setting Add disable_custom_api_keys UI setting that prevents users from specifying custom key values during key generation and regeneration. When enabled, all keys must be auto-generated, eliminating the risk of key hash collisions in multi-tenant environments. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * [Fix] Add disable_custom_api_keys to UISettings Pydantic model Without this field on the model, GET /get/ui_settings omits the setting from the response and field_schema, preventing the UI from reading it. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: Register DynamoAI guardrail initializer and enum entry (#23752) * fix: Register DynamoAI guardrail initializer and enum entry Fix the "Unsupported guardrail: dynamoai" error by: 1. Adding DYNAMOAI to SupportedGuardrailIntegrations enum 2. Implementing initialize_guardrail() and registries in dynamoai/__init__.py The DynamoAI guardrail was added in PR #15920 but never properly registered in the initialization system. The __init__.py was missing the guardrail_initializer_registry and guardrail_class_registry dictionaries that the dynamic discovery mechanism looks for at module load time. Fixes #22773 Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com> * Update litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * test: Add tests for DynamoAI guardrail registration Verifies enum entry, initializer registry, class registry, instance creation, and global registry discovery. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Haiku 4.5 <noreply@anthropic.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * docs: add v1.82.3 release notes and update provider_endpoints_support.json (#23816) * [Feature] Add disable_custom_api_keys toggle to UI Settings page Adds a toggle switch to the admin UI Settings page so administrators can enable/disable custom API key values without making direct API calls. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Revert "docs: add v1.82.3 release notes and update provider_endpoints_support…" (#23817) This reverts commit966124966f. * [Fix] Rename toggle label to "Disable custom Virtual key values" Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * [Fix] Remove "API" from custom key description text Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix(ui): CSV export empty on Global Usage page Aggregated endpoint returns empty breakdown.entities; fall back to grouping breakdown.api_keys by team_id. * Revert "fix: langfuse trace leak key on model params" * fix: support served_model_name for Baseten dedicated deployments Baseten dedicated deployments use an 8-char deployment ID for URL routing, but the vLLM server may expect a different model name in the request body (e.g. baseten-hosted/zai-org/GLM-5 vs wd1lndkw). Add served_model_name litellm_param to override the model field in the request body, and declare it in LiteLLMParamsTypedDict and GenericLiteLLMParams for IDE support. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Co-authored-by: Harshit Jain <harshitjain0562@gmail.com> Co-authored-by: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Co-authored-by: AlexKer <AlexKer@users.noreply.github.com> Co-authored-by: joereyna <joseph.reyna@gmail.com> Co-authored-by: Ryan Crabbe <rcrabbe@berkeley.edu> Co-authored-by: ryan-crabbe <128659760+ryan-crabbe@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: yuneng-jiang <yuneng.jiang@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Cursor Agent <cursoragent@cursor.com> Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com> Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com>
This commit is contained in:
Alex Ker
GitHub
parent
47d238b228
commit
742b6be36a
74 changed files with 2115 additions and 116 deletions
|
|
@ -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 \
|
||||
|
|
|
|||
|
|
@ -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 <description>` 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.
|
||||
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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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/<model-name>` (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}`.
|
||||
:::
|
||||
|
|
|
|||
|
|
@ -358,7 +358,7 @@ model_cost_map_url: str = os.getenv(
|
|||
)
|
||||
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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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 @@ Disable remote fetching entirely:
|
|||
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 @@ class GetBlogPosts:
|
|||
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 <item>.
|
||||
"""
|
||||
root = ET.fromstring(xml_text)
|
||||
channel = root.find("channel")
|
||||
if channel is None:
|
||||
raise ValueError("RSS feed missing <channel> 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 @@ class GetBlogPosts:
|
|||
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 @@ class GetBlogPosts:
|
|||
)
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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 @@ class BasetenConfig(OpenAIGPTConfig):
|
|||
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:
|
||||
|
|
|
|||
|
|
@ -903,13 +903,17 @@ if MCP_AVAILABLE:
|
|||
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 "",
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
|||
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.team_member_permission_checks import (
|
|||
)
|
||||
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 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
|
||||
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(
|
||||
|
|
|
|||
|
|
@ -340,9 +340,16 @@ class ProxyInitializationHelpers:
|
|||
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
|
||||
|
|
|
|||
|
|
@ -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 @@ ALLOWED_UI_SETTINGS_FIELDS = {
|
|||
"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
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ guardrails:
|
|||
class SupportedGuardrailIntegrations(Enum):
|
||||
APORIA = "aporia"
|
||||
BEDROCK = "bedrock"
|
||||
DYNAMOAI = "dynamoai"
|
||||
GUARDRAILS_AI = "guardrails_ai"
|
||||
LAKERA = "lakera"
|
||||
LAKERA_V2 = "lakera_v2"
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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__])
|
||||
|
|
|
|||
|
|
@ -158,7 +158,6 @@ class TestExecuteWithMcpClient:
|
|||
|
||||
|
||||
@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 @@ class TestExecuteWithMcpClient:
|
|||
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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -41,12 +41,15 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
|
|||
_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)
|
||||
|
|
|
|||
|
|
@ -67,6 +67,30 @@ class TestMaybeSetupPrometheusMultiprocDir:
|
|||
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",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -677,7 +677,7 @@ class TestHealthAppFactory:
|
|||
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 @@ class TestHealthAppFactory:
|
|||
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -111,6 +111,37 @@ class TestProxySettingEndpoints:
|
|||
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
|
||||
):
|
||||
|
|
|
|||
|
|
@ -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="")
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
"""Tests for GetBlogPosts utility class."""
|
||||
import json
|
||||
import time
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
|
@ -13,16 +12,26 @@ from litellm.litellm_core_utils.get_blog_posts import (
|
|||
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 = """\
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<rss version="2.0">
|
||||
<channel>
|
||||
<title>LiteLLM Blog</title>
|
||||
<item>
|
||||
<title>Test Post</title>
|
||||
<link>https://docs.litellm.ai/blog/test</link>
|
||||
<description>A test post.</description>
|
||||
<pubDate>Wed, 01 Jan 2026 10:00:00 GMT</pubDate>
|
||||
</item>
|
||||
<item>
|
||||
<title>Second Post</title>
|
||||
<link>https://docs.litellm.ai/blog/second</link>
|
||||
<description>Another post.</description>
|
||||
<pubDate>Tue, 31 Dec 2025 10:00:00 GMT</pubDate>
|
||||
</item>
|
||||
</channel>
|
||||
</rss>
|
||||
"""
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
|
|
@ -45,26 +54,48 @@ def test_load_local_blog_posts_returns_list():
|
|||
assert "url" in first
|
||||
|
||||
|
||||
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_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 <channel>"):
|
||||
GetBlogPosts.parse_rss_to_posts("<rss></rss>")
|
||||
|
||||
|
||||
def test_validate_blog_posts_valid():
|
||||
assert GetBlogPosts.validate_blog_posts(SAMPLE_RESPONSE) is True
|
||||
|
||||
|
||||
def test_validate_blog_posts_missing_posts_key():
|
||||
assert GetBlogPosts.validate_blog_posts({"other": []}) is False
|
||||
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 test_get_blog_posts_ttl_cache_not_refetched():
|
|||
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 test_get_blog_posts_ttl_cache_not_refetched():
|
|||
|
||||
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(
|
||||
|
|
|
|||
|
|
@ -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<DefaultUserSettingsProps> = ({
|
|||
<div className="flex items-center justify-between mb-3">
|
||||
<Text className="font-medium">Team {index + 1}</Text>
|
||||
<Button
|
||||
size="sm"
|
||||
variant="secondary"
|
||||
icon={DeleteOutlined}
|
||||
size="small"
|
||||
danger
|
||||
icon={<DeleteOutlined />}
|
||||
onClick={() => removeTeam(index)}
|
||||
className="text-red-500 hover:text-red-700"
|
||||
>
|
||||
Remove
|
||||
</Button>
|
||||
|
|
@ -208,7 +207,7 @@ const DefaultUserSettings: React.FC<DefaultUserSettingsProps> = ({
|
|||
</div>
|
||||
))}
|
||||
|
||||
<Button variant="secondary" icon={PlusOutlined} onClick={addTeam} className="w-full">
|
||||
<Button icon={<PlusOutlined />} onClick={addTeam} className="w-full">
|
||||
Add Team
|
||||
</Button>
|
||||
</div>
|
||||
|
|
@ -462,7 +461,6 @@ const DefaultUserSettings: React.FC<DefaultUserSettingsProps> = ({
|
|||
(isEditing ? (
|
||||
<div className="flex gap-2">
|
||||
<Button
|
||||
variant="secondary"
|
||||
onClick={() => {
|
||||
setIsEditing(false);
|
||||
setEditedValues(settings.values || {});
|
||||
|
|
@ -471,12 +469,12 @@ const DefaultUserSettings: React.FC<DefaultUserSettingsProps> = ({
|
|||
>
|
||||
Cancel
|
||||
</Button>
|
||||
<Button onClick={handleSaveSettings} loading={saving}>
|
||||
<Button type="primary" onClick={handleSaveSettings} loading={saving}>
|
||||
Save Changes
|
||||
</Button>
|
||||
</div>
|
||||
) : (
|
||||
<Button onClick={() => setIsEditing(true)}>Edit Settings</Button>
|
||||
<Button type="primary" onClick={() => setIsEditing(true)}>Edit Settings</Button>
|
||||
))}
|
||||
</div>
|
||||
|
||||
|
|
|
|||
|
|
@ -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<string, any>).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");
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -17,6 +17,49 @@ const extractTeamIdFromApiKeyBreakdown = (apiKeyBreakdown: Record<string, any> |
|
|||
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<string, any>): Record<string, any> => {
|
||||
const apiKeys = breakdown.api_keys;
|
||||
if (!apiKeys || Object.keys(apiKeys).length === 0) return {};
|
||||
|
||||
const grouped: Record<string, any> = {};
|
||||
|
||||
for (const [keyId, keyData] of Object.entries<any>(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<string, any>): Record<string, any> => {
|
||||
const entities = breakdown.entities;
|
||||
if (entities && Object.keys(entities).length > 0) return entities;
|
||||
return aggregateApiKeysIntoEntities(breakdown);
|
||||
};
|
||||
|
||||
export const getEntityBreakdown = (
|
||||
spendData: EntitySpendData,
|
||||
teamAliasMap: Record<string, string> = {},
|
||||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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(<HashicorpVaultEmptyPlaceholder onAdd={vi.fn()} />);
|
||||
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(<HashicorpVaultEmptyPlaceholder onAdd={onAdd} />);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: /configure vault/i }));
|
||||
|
||||
expect(onAdd).toHaveBeenCalledOnce();
|
||||
});
|
||||
|
||||
it("should display the description text about Vault purpose", () => {
|
||||
render(<HashicorpVaultEmptyPlaceholder onAdd={vi.fn()} />);
|
||||
expect(
|
||||
screen.getByText(/Configure Hashicorp Vault to securely manage provider API keys/),
|
||||
).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
@ -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(
|
||||
<PageVisibilitySettings
|
||||
enabledPagesInternalUsers={null}
|
||||
isUpdating={false}
|
||||
onUpdate={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
expect(screen.getByText("Not set (all pages visible)")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should show the selected page count tag when pages are configured", () => {
|
||||
render(
|
||||
<PageVisibilitySettings
|
||||
enabledPagesInternalUsers={["usage", "keys"]}
|
||||
isUpdating={false}
|
||||
onUpdate={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
expect(screen.getByText("2 pages selected")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should show singular 'page' when exactly one page is selected", () => {
|
||||
render(
|
||||
<PageVisibilitySettings
|
||||
enabledPagesInternalUsers={["usage"]}
|
||||
isUpdating={false}
|
||||
onUpdate={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
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(
|
||||
<PageVisibilitySettings
|
||||
enabledPagesInternalUsers={["usage"]}
|
||||
isUpdating={false}
|
||||
onUpdate={onUpdate}
|
||||
/>,
|
||||
);
|
||||
|
||||
// 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(
|
||||
<PageVisibilitySettings
|
||||
enabledPagesInternalUsers={null}
|
||||
enabledPagesPropertyDescription="Controls which pages are visible"
|
||||
isUpdating={false}
|
||||
onUpdate={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
expect(screen.getByText("Controls which pages are visible")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
@ -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 (
|
||||
<Card title="UI Settings">
|
||||
{isLoading ? (
|
||||
|
|
@ -382,6 +397,26 @@ export default function UISettings() {
|
|||
|
||||
<Divider />
|
||||
|
||||
{/* Disable custom Virtual key values */}
|
||||
<Space align="start" size="middle">
|
||||
<Switch
|
||||
checked={Boolean(values.disable_custom_api_keys)}
|
||||
disabled={isUpdating}
|
||||
loading={isUpdating}
|
||||
onChange={handleToggleDisableCustomApiKeys}
|
||||
aria-label={disableCustomApiKeysProperty?.description ?? "Disable custom Virtual key values"}
|
||||
/>
|
||||
<Space direction="vertical" size={4}>
|
||||
<Typography.Text strong>Disable custom Virtual key values</Typography.Text>
|
||||
<Typography.Text type="secondary">
|
||||
{disableCustomApiKeysProperty?.description ??
|
||||
"If true, users cannot specify custom key values. All keys must be auto-generated."}
|
||||
</Typography.Text>
|
||||
</Space>
|
||||
</Space>
|
||||
|
||||
<Divider />
|
||||
|
||||
{/* Page Visibility for Internal Users */}
|
||||
<PageVisibilitySettings
|
||||
enabledPagesInternalUsers={values.enabled_ui_pages_internal_users}
|
||||
|
|
|
|||
|
|
@ -166,6 +166,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ 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<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
|
|||
"budget_duration",
|
||||
"tpm_limit",
|
||||
"rpm_limit",
|
||||
...(disableCustomApiKeys ? ["key"] : []),
|
||||
]}
|
||||
/>
|
||||
</AccordionBody>
|
||||
|
|
|
|||
|
|
@ -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(<UiLoadingSpinner data-testid="spinner" />);
|
||||
expect(screen.getByTestId("spinner")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should apply custom className alongside default classes", () => {
|
||||
render(<UiLoadingSpinner data-testid="spinner" className="text-red-500" />);
|
||||
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(<UiLoadingSpinner data-testid="spinner" aria-label="Loading" />);
|
||||
expect(screen.getByLabelText("Loading")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
@ -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 () => {
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
40
ui/litellm-dashboard/src/utils/errorUtils.test.ts
Normal file
40
ui/litellm-dashboard/src/utils/errorUtils.test.ts
Normal file
|
|
@ -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");
|
||||
});
|
||||
});
|
||||
|
|
@ -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);
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Reference in a new issue