scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-12 03:48:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 118

Merge branch 'BerriAI:main' into main

Browse source
This commit is contained in:
Esteban Zeller 2026-02-25 15:01:07 -03:00 • committed by GitHub
commit bf874fd7ce
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
106 changed files with 8772 additions and 821 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
README.md
View file

@ -203,7 +203,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
{
"mcpServers": {
"LiteLLM": {
"url": "http://localhost:4000/mcp",
"url": "http://localhost:4000/mcp/",
"headers": {
"x-litellm-api-key": "Bearer sk-1234"
}

82
docs/my-website/docs/completion/prompt_caching.md
View file

@ -63,7 +63,6 @@ for _ in range(2):
}
],
},
# marked for caching with the cache_control parameter, so that this checkpoint can read from the previous cache.
{
"role": "user",
"content": [
@ -77,7 +76,6 @@ for _ in range(2):
"role": "assistant",
"content": "Certainly! the key terms and conditions are the following: the contract is 1 year long for $10/mo",
},
# The final turn is marked with cache-control, for continuing in followups.
{
"role": "user",
"content": [
@ -112,16 +110,16 @@ model_list:
api_key: os.environ/OPENAI_API_KEY
```
2. Start proxy
2. Start proxy
```bash
litellm --config /path/to/config.yaml
```
3. Test it!
3. Test it!
```python
from openai import OpenAI
from openai import OpenAI
import os
client = OpenAI(
@ -144,7 +142,6 @@ for _ in range(2):
}
],
},
# marked for caching with the cache_control parameter, so that this checkpoint can read from the previous cache.
{
"role": "user",
"content": [
@ -158,7 +155,6 @@ for _ in range(2):
"role": "assistant",
"content": "Certainly! the key terms and conditions are the following: the contract is 1 year long for $10/mo",
},
# The final turn is marked with cache-control, for continuing in followups.
{
"role": "user",
"content": [
@ -183,6 +179,78 @@ assert response.usage.prompt_tokens_details.cached_tokens > 0
</TabItem>
</Tabs>
### OpenAI `prompt_cache_key` and `prompt_cache_retention`
OpenAI prompt caching is [**automatic**](https://platform.openai.com/docs/guides/prompt-caching) — no `cache_control` message annotations are needed. Any request with 1024+ prompt tokens is eligible for caching.
OpenAI also supports two optional parameters for more control over caching behavior:
- **`prompt_cache_key`** (string) — A routing hint that improves cache hit rates for requests sharing long common prefixes. Requests with the same cache key are routed to the same backend, increasing the likelihood of a cache hit.
- **`prompt_cache_retention`** (`"in_memory"` or `"24h"`) — Controls cache TTL. Default is `"in_memory"` (5–10 min). Set to `"24h"` for extended caching that offloads KV tensors to GPU-local storage.
<Tabs>
<TabItem value="sdk" label="SDK">
```python
from litellm import completion
import os
os.environ["OPENAI_API_KEY"] = ""
response = completion(
model="gpt-4o",
messages=[
{
"role": "system",
"content": "You are an AI assistant tasked with analyzing legal documents. "
+ "Here is the full text of a complex legal agreement " * 400,
},
{
"role": "user",
"content": "What are the key terms and conditions?",
},
],
prompt_cache_key="legal-doc-analysis",
prompt_cache_retention="24h",
)
print(response.usage)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
```python
from openai import OpenAI
client = OpenAI(
api_key="LITELLM_PROXY_KEY",
base_url="LITELLM_PROXY_BASE",
)
response = client.chat.completions.create(
model="gpt-4o",
messages=[
{
"role": "system",
"content": "You are an AI assistant tasked with analyzing legal documents. "
+ "Here is the full text of a complex legal agreement " * 400,
},
{
"role": "user",
"content": "What are the key terms and conditions?",
},
],
extra_body={
"prompt_cache_key": "legal-doc-analysis",
"prompt_cache_retention": "24h",
},
)
print(response.usage)
```
</TabItem>
</Tabs>
### Anthropic Example
Anthropic charges for cache writes.

1
docs/my-website/docs/generateContent.md
View file

@ -15,6 +15,7 @@ Use LiteLLM to call Google AI's generateContent endpoints for text generation, m
| Streaming | ✅ | |
| Fallbacks | ✅ | between supported models |
| Loadbalancing | ✅ | between supported models |
| Metadata Tracking | ✅ | passes trace ID, metadata to observability callbacks (e.g. S3, Langfuse) |
## Usage
---

2
docs/my-website/docs/mcp.md
View file

@ -641,7 +641,7 @@ import asyncio
config = {
"mcpServers": {
"mcp_group": {
"url": "http://localhost:4000/mcp",
"url": "http://localhost:4000/mcp/",
"headers": {
"x-mcp-servers": "dev_group", # assume this gives access to github, zapier and deepwiki
"x-litellm-api-key": "Bearer sk-1234",

129
docs/my-website/docs/proxy/forward_client_headers.md
View file

@ -37,11 +37,11 @@ The following rules determine which headers are forwarded (see [`_get_forwardabl
| Rule | Example | Forwarded? |
|---|---|---|
| Headers starting with `x-` | `x-trace-id`, `x-custom-header`, `x-request-source` | ✅ Yes |
| `anthropic-beta` header | `anthropic-beta: prompt-caching-2024-07-31` | ✅ Yes |
| Headers starting with `x-stainless-*` | `x-stainless-lang`, `x-stainless-arch` | ❌ No (causes OpenAI SDK issues) |
| Standard HTTP headers | `Authorization`, `Content-Type`, `Host` | ❌ No |
| Other provider headers | `Accept`, `User-Agent` | ❌ No |
| Headers starting with `x-` | `x-trace-id`, `x-custom-header`, `x-request-source` | Yes |
| `anthropic-beta` header | `anthropic-beta: prompt-caching-2024-07-31` | Yes |
| Headers starting with `x-stainless-*` | `x-stainless-lang`, `x-stainless-arch` | No (causes OpenAI SDK issues) |
| Standard HTTP headers | `Authorization`, `Content-Type`, `Host` | No |
| Other provider headers | `Accept`, `User-Agent` | No |
### Additional Header Mechanisms
@ -61,6 +61,125 @@ general_settings:
forward_client_headers_to_llm_api: true
```
## Forward LLM Provider Authentication Headers
**New in v1.82+**: By default, LiteLLM strips authentication headers like `x-api-key`, `x-goog-api-key`, and `api-key` from client requests for security (these are typically used to authenticate with the proxy itself). However, you can enable forwarding of these LLM provider authentication headers to allow **Bring Your Own Key (BYOK)** scenarios where clients send their own API keys to the LLM provider.
### Configuration
Add `forward_llm_provider_auth_headers: true` to your `general_settings`:
```yaml
general_settings:
forward_client_headers_to_llm_api: true
forward_llm_provider_auth_headers: true # 👈 Enable BYOK
```
### Which Headers Are Forwarded
When `forward_llm_provider_auth_headers: true`, the following LLM provider authentication headers are preserved and forwarded:
| Header | Provider | Example |
|--------|----------|---------|
| `x-api-key` | Anthropic, Azure AI, Databricks | `x-api-key: sk-ant-api03-...` |
| `x-goog-api-key` | Google AI Studio | `x-goog-api-key: AIza...` |
| `api-key` | Azure OpenAI | `api-key: your-azure-key` |
| `ocp-apim-subscription-key` | Azure APIM | `ocp-apim-subscription-key: your-key` |
:::warning Important Security Note
The proxy's `Authorization` header (used for proxy authentication) is **never** forwarded to LLM providers, even with this setting enabled. This ensures your proxy authentication remains secure.
:::
### Use Case: Client-Side API Keys (BYOK)
This feature enables scenarios where:
1. **Clients bring their own LLM provider API keys** instead of using keys configured in the proxy
2. **Multi-tenant applications** where each tenant has their own Anthropic/OpenAI account
3. **Development environments** where developers use their personal API keys through a shared proxy
#### Example: Anthropic BYOK
```yaml
# proxy_config.yaml
model_list:
- model_name: claude-sonnet-4
litellm_params:
model: anthropic/claude-sonnet-4-20250514
# No api_key configured! Will use client's key
general_settings:
forward_client_headers_to_llm_api: true
forward_llm_provider_auth_headers: true # Enable BYOK
```
Client request:
```bash
curl -X POST "http://localhost:4000/v1/messages" \
-H "Authorization: Bearer sk-proxy-auth-123" \ # Proxy authentication (stripped)
-H "x-api-key: sk-ant-api03-YOUR-KEY..." \ # Client's Anthropic key (forwarded!)
-H "Content-Type: application/json" \
-d '{
"model": "claude-sonnet-4",
"messages": [{"role": "user", "content": "Hello"}],
"max_tokens": 100
}'
```
#### Example: Google AI Studio BYOK
```yaml
model_list:
- model_name: gemini-pro
litellm_params:
model: gemini/gemini-1.5-pro
# No api_key configured
general_settings:
forward_client_headers_to_llm_api: true
forward_llm_provider_auth_headers: true
```
Client request:
```bash
curl -X POST "http://localhost:4000/v1/chat/completions" \
-H "Authorization: Bearer sk-proxy-auth-123" \
-H "x-goog-api-key: AIza..." \
-d '{
"model": "gemini-pro",
"messages": [{"role": "user", "content": "Hello"}]
}'
```
### Security Considerations
**When to Use This Feature:**
- Internal tools where you trust all clients
- Development/testing environments
- Multi-tenant apps with proper client authentication
- Scenarios where you want clients to use their own API keys
**When NOT to Use:**
- Public APIs where you don't trust all clients
- When you want centralized billing/cost control
- When you need to enforce rate limits at the proxy level
### Backward Compatibility
For backward compatibility, if you have `forward_client_headers_to_llm_api: true` but don't explicitly set `forward_llm_provider_auth_headers`, the behavior is:
- **Default**: LLM provider auth headers are **NOT** forwarded (safe default)
- **Explicit `true`**: LLM provider auth headers **ARE** forwarded (BYOK enabled)
```yaml
# Safe default - auth headers NOT forwarded
general_settings:
forward_client_headers_to_llm_api: true
# BYOK enabled - auth headers ARE forwarded
general_settings:
forward_client_headers_to_llm_api: true
forward_llm_provider_auth_headers: true # 👈 Opt-in required
```
## Enable for a Model Group
Add the `forward_client_headers_to_llm_api` setting under `model_group_settings` in your configuration:

15
docs/my-website/docs/tutorials/presidio_pii_masking.md
View file

@ -592,6 +592,21 @@ def test_pii_masking_allows_normal_text():
## Part 7: Troubleshooting
### Issue: Guardrail failure: non-JSON response from Presidio
**Symptom:** You receive an error indicating `expected application/json Content-Type but received text/html` or similar.
**Root cause:** Your ingress controller or reverse proxy might be routing the `/analyze` or `/anonymize` POST request to a health endpoint (like `/health` or `/presidio-analyzer/health`) which returns plain text instead of JSON.
**Fix:** Ensure your `PRESIDIO_ANALYZER_API_BASE` and `PRESIDIO_ANONYMIZER_API_BASE` are correctly pointing directly to the Presidio API endpoints, or that your ingress routes the path correctly without stripping it and inadvertently forwarding to a plain-text health check endpoint.
**Verification:** You can verify your endpoints using `curl`. It should return a JSON array, not `text/html`:
```bash
curl -sv -X POST http://your-analyzer-endpoint/analyze \
-H "Content-Type: application/json" \
-d '{"text":"test","language":"en"}'
```
### Issue: Presidio Not Detecting PII
**Check 1: Language Configuration**

11
litellm/batches/main.py
View file

@ -37,7 +37,9 @@ from litellm.types.llms.openai import (
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import (
LIST_BATCHES_SUPPORTED_PROVIDERS,
OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS,
ListBatchesSupportedProvider,
LiteLLMBatch,
LlmProviders,
)
@ -674,7 +676,7 @@ def retrieve_batch(
async def alist_batches(
after: Optional[str] = None,
limit: Optional[int] = None,
custom_llm_provider: Literal["openai", "azure", "hosted_vllm", "vertex_ai"] = "openai",
custom_llm_provider: ListBatchesSupportedProvider = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@ -717,7 +719,7 @@ async def alist_batches(
def list_batches(
after: Optional[str] = None,
limit: Optional[int] = None,
custom_llm_provider: Literal["openai", "azure", "hosted_vllm", "vertex_ai"] = "openai",
custom_llm_provider: ListBatchesSupportedProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@ -843,8 +845,9 @@ def list_batches(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'list_batch'. Supported providers: openai, azure, vertex_ai.".format(
custom_llm_provider
message="LiteLLM doesn't support {} for 'list_batch'. Supported providers: {}.".format(
custom_llm_provider,
", ".join(sorted(LIST_BATCHES_SUPPORTED_PROVIDERS)),
),
model="n/a",
llm_provider=custom_llm_provider,

1
litellm/constants.py
View file

@ -244,6 +244,7 @@ REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer
MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100))
# Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth
LITELLM_ASYNCIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000))
TOOL_POLICY_CACHE_TTL_SECONDS = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60))
# Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger.
# Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire.
MAX_SIZE_IN_MEMORY_QUEUE = int(

29
litellm/litellm_core_utils/health_check_helpers.py
View file

@ -4,6 +4,8 @@ Helper functions for health check calls.
from typing import TYPE_CHECKING, Callable, Dict, Literal, Optional
from litellm.types.utils import LIST_BATCHES_SUPPORTED_PROVIDERS
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging
@ -82,6 +84,27 @@ class HealthCheckHelpers:
"tags": [LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME],
}
@staticmethod
async def _batch_health_check(
custom_llm_provider: str,
model_params: dict,
filtered_model_params: dict,
) -> dict:
"""
Health check for batch mode.
Calls list_batches for providers that support it (openai, hosted_vllm, azure,
vertex_ai). For all other providers (e.g. bedrock) the batch API surface doesn't
include list_batches, so we fall back to acompletion to verify connectivity and
credential validity instead.
"""
import litellm
if custom_llm_provider in LIST_BATCHES_SUPPORTED_PROVIDERS:
return await litellm.alist_batches(**filtered_model_params)
else:
return await litellm.acompletion(**model_params)
@staticmethod
def get_mode_handlers(
model: str,
@ -176,8 +199,10 @@ class HealthCheckHelpers:
api_key=model_params.get("api_key", None),
api_version=model_params.get("api_version", None),
),
"batch": lambda: litellm.alist_batches(
**_filter_model_params(model_params=model_params),
"batch": lambda: HealthCheckHelpers._batch_health_check(
custom_llm_provider=custom_llm_provider,
model_params=model_params,
filtered_model_params=_filter_model_params(model_params=model_params),
),
"responses": lambda: litellm.aresponses(
**_filter_model_params(model_params=model_params),

10
litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py
View file

@ -11,6 +11,7 @@ from typing import Any, List, Optional
import httpx
from litellm.types.utils import Usage
from litellm.llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import (
AmazonQwen3Config,
)
@ -79,10 +80,11 @@ class AmazonQwen2Config(AmazonQwen3Config):
# Set usage information if available in response
if "usage" in response_data:
usage_data = response_data["usage"]
if hasattr(model_response, 'usage'):
model_response.usage.prompt_tokens = usage_data.get("prompt_tokens", 0)
model_response.usage.completion_tokens = usage_data.get("completion_tokens", 0)
model_response.usage.total_tokens = usage_data.get("total_tokens", 0)
model_response.usage = Usage(
prompt_tokens=usage_data.get("prompt_tokens", 0),
completion_tokens=usage_data.get("completion_tokens", 0),
total_tokens=usage_data.get("total_tokens", 0),
)
return model_response

10
litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py
View file

@ -10,6 +10,7 @@ from typing import Any, List, Optional
import httpx
from litellm.types.utils import Usage
from litellm.llms.base_llm.chat.transformation import BaseConfig
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
AmazonInvokeConfig,
@ -201,10 +202,11 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig):
# Set usage information if available in response
if "usage" in response_data:
usage_data = response_data["usage"]
if hasattr(model_response, 'usage'):
model_response.usage.prompt_tokens = usage_data.get("prompt_tokens", 0)
model_response.usage.completion_tokens = usage_data.get("completion_tokens", 0)
model_response.usage.total_tokens = usage_data.get("total_tokens", 0)
model_response.usage = Usage(
prompt_tokens=usage_data.get("prompt_tokens", 0),
completion_tokens=usage_data.get("completion_tokens", 0),
total_tokens=usage_data.get("total_tokens", 0),
)
return model_response

1
litellm/llms/openai/chat/gpt_transformation.py
View file

@ -162,6 +162,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
"service_tier",
"safety_identifier",
"prompt_cache_key",
"prompt_cache_retention",
"store",
] # works across all models

6
litellm/llms/openai/chat/o_series_transformation.py
View file

@ -131,9 +131,7 @@ class OpenAIOSeriesConfig(OpenAIGPTConfig):
def is_model_o_series_model(self, model: str) -> bool:
model = model.split("/")[-1] # could be "openai/o3" or "o3"
return model in litellm.open_ai_chat_completion_models and any(
model.startswith(pfx) for pfx in ("o1", "o3", "o4")
)
return model.startswith(("o1", "o3", "o4")) and model in litellm.open_ai_chat_completion_models
@overload
def _transform_messages(
@ -173,4 +171,4 @@ class OpenAIOSeriesConfig(OpenAIGPTConfig):
else:
return super()._transform_messages(
messages, model, is_async=cast(Literal[False], False)
)
)

5
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -10,6 +10,7 @@ from litellm.proxy._experimental.mcp_server.ui_session_utils import (
)
from litellm.proxy._experimental.mcp_server.utils import merge_mcp_headers
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.types.mcp import MCPAuth
@ -685,7 +686,7 @@ if MCP_AVAILABLE:
return await _execute_with_mcp_client(
new_mcp_server_request,
_test_connection_operation,
raw_headers=dict(request.headers),
raw_headers=_safe_get_request_headers(request),
)
@router.post("/test/tools/list")
@ -744,5 +745,5 @@ if MCP_AVAILABLE:
_list_tools_operation,
mcp_auth_header=mcp_auth_header,
oauth2_headers=oauth2_headers,
raw_headers=dict(request.headers),
raw_headers=_safe_get_request_headers(request),
)

11
litellm/proxy/_types.py
View file

@ -183,6 +183,7 @@ class LitellmTableNames(str, enum.Enum):
KEY_TABLE_NAME = "LiteLLM_VerificationToken"
PROXY_MODEL_TABLE_NAME = "LiteLLM_ProxyModelTable"
MANAGED_FILE_TABLE_NAME = "LiteLLM_ManagedFileTable"
TOOL_TABLE_NAME = "LiteLLM_ToolTable"
class Litellm_EntityType(enum.Enum):
@ -850,6 +851,7 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase):
max_budget: Optional[float] = None
user_id: Optional[str] = None
team_id: Optional[str] = None
agent_id: Optional[str] = None
max_parallel_requests: Optional[int] = None
metadata: Optional[dict] = {}
tpm_limit: Optional[int] = None
@ -4123,6 +4125,15 @@ class SpendUpdateQueueItem(TypedDict, total=False):
response_cost: Optional[float]
class ToolDiscoveryQueueItem(TypedDict, total=False):
tool_name: str
origin: Optional[str] # MCP server name or "user_defined"
created_by: Optional[str]
key_hash: Optional[str] # hash of virtual key that triggered discovery
team_id: Optional[str] # team that triggered discovery
key_alias: Optional[str] # human-readable key alias
class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase):
unified_file_id: str
file_object: Optional[OpenAIFileObject] = None

28
litellm/proxy/auth/user_api_key_auth.py
View file

@ -483,7 +483,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
parent_otel_span = (
open_telemetry_logger.create_litellm_proxy_request_started_span(
start_time=start_time,
headers=dict(request.headers),
headers=_safe_get_request_headers(request),
)
)
@ -562,7 +562,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
parent_otel_span=parent_otel_span,
request_headers=dict(request.headers),
request_headers=_safe_get_request_headers(request),
)
is_proxy_admin = result["is_proxy_admin"]
@ -593,9 +593,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
user_id=user_id,
team_id=team_id,
team_alias=(
team_object.team_alias
if team_object is not None
else None
team_object.team_alias if team_object is not None else None
),
team_metadata=team_object.metadata
if team_object is not None
@ -709,12 +707,12 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
if isinstance(api_key, str):
return UserAPIKeyAuth(
api_key=api_key,
user_role=LitellmUserRoles.PROXY_ADMIN,
user_role=LitellmUserRoles.INTERNAL_USER,
parent_otel_span=parent_otel_span,
)
else:
return UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_role=LitellmUserRoles.INTERNAL_USER,
parent_otel_span=parent_otel_span,
)
elif api_key is None: # only require api key if master key is set
@ -846,7 +844,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
)
valid_token.parent_otel_span = parent_otel_span
if _end_user_object is not None:
valid_token.end_user_object_permission = _end_user_object.object_permission
valid_token.end_user_object_permission = (
_end_user_object.object_permission
)
return valid_token
@ -954,7 +954,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
if isinstance(
api_key, str
): # if generated token, make sure it starts with sk-.
_masked_key = "{}****{}".format(api_key[:4], api_key[-4:]) if len(api_key) > 8 else "****"
_masked_key = (
"{}****{}".format(api_key[:4], api_key[-4:])
if len(api_key) > 8
else "****"
)
assert api_key.startswith(
"sk-"
), "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format(
@ -1304,9 +1308,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
if _end_user_object is not None:
valid_token_dict.update(end_user_params)
valid_token_dict["end_user_object_permission"] = (
_end_user_object.object_permission
)
valid_token_dict[
"end_user_object_permission"
] = _end_user_object.object_permission
# check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions
# sso/login, ui/login, /key functions and /user functions

22
litellm/proxy/common_utils/http_parsing_utils.py
View file

@ -135,17 +135,29 @@ def _safe_set_request_parsed_body(
def _safe_get_request_headers(request: Optional[Request]) -> dict:
"""
[Non-Blocking] Safely get the request headers
[Non-Blocking] Safely get the request headers.
Caches the result on request.state to avoid re-creating dict(request.headers) per call.
Warning: Callers must NOT mutate the returned dict — it is shared across
all callers within the same request via the cache.
"""
if request is None:
return {}
cached = getattr(request.state, "_cached_headers", None)
if cached is not None:
return cached
try:
if request is None:
return {}
return dict(request.headers)
headers = dict(request.headers)
except Exception as e:
verbose_proxy_logger.debug(
"Unexpected error reading request headers - {}".format(e)
)
return {}
headers = {}
try:
request.state._cached_headers = headers
except Exception:
pass # request.state may not be available in all contexts
return headers
def check_file_size_under_limit(

3
litellm/proxy/custom_hooks/custom_ui_sso_hook.py
View file

@ -3,6 +3,7 @@ from fastapi_sso.sso.base import OpenID
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
class CustomSSOLoginHandler(CustomLogger):
@ -18,7 +19,7 @@ class CustomSSOLoginHandler(CustomLogger):
self,
request: Request,
) -> OpenID:
request_headers_dict = dict(request.headers)
request_headers_dict = _safe_get_request_headers(request)
verbose_logger.debug("inside custom ui sso sign in hook...")
return OpenID(
id=request_headers_dict.get("x-litellm-user-id") or "123",

151
litellm/proxy/db/db_spend_update_writer.py
View file

@ -13,7 +13,17 @@ import random
import time
import traceback
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast, overload
from typing import (
TYPE_CHECKING,
Any,
Dict,
List,
Literal,
Optional,
Union,
cast,
overload,
)
import litellm
from litellm._logging import verbose_proxy_logger
@ -23,18 +33,19 @@ from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
BaseDailySpendTransaction,
DailyTagSpendTransaction,
DailyOrganizationSpendTransaction,
DailyTeamSpendTransaction,
DailyEndUserSpendTransaction,
DailyUserSpendTransaction,
DailyAgentSpendTransaction,
DailyEndUserSpendTransaction,
DailyOrganizationSpendTransaction,
DailyTagSpendTransaction,
DailyTeamSpendTransaction,
DailyUserSpendTransaction,
DBSpendUpdateTransactions,
Litellm_EntityType,
LiteLLM_UserTable,
SpendLogsMetadata,
SpendLogsPayload,
SpendUpdateQueueItem,
ToolDiscoveryQueueItem,
)
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
DailySpendUpdateQueue,
@ -42,6 +53,9 @@ from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import (
ToolDiscoveryQueue,
)
from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
if TYPE_CHECKING:
@ -67,6 +81,7 @@ class DBSpendUpdateWriter:
self.redis_update_buffer = RedisUpdateBuffer(redis_cache=self.redis_cache)
self.pod_lock_manager = PodLockManager()
self.spend_update_queue = SpendUpdateQueue()
self.tool_discovery_queue = ToolDiscoveryQueue()
self.daily_spend_update_queue = DailySpendUpdateQueue()
self.daily_team_spend_update_queue = DailySpendUpdateQueue()
self.daily_end_user_spend_update_queue = DailySpendUpdateQueue()
@ -222,6 +237,13 @@ class DBSpendUpdateWriter:
)
)
self._enqueue_tool_registry_upsert(
kwargs=kwargs,
completion_response=completion_response,
hashed_token=hashed_token,
team_id=team_id,
)
verbose_proxy_logger.debug("Runs spend update on all tables")
except Exception:
verbose_proxy_logger.error(
@ -237,6 +259,104 @@ class DBSpendUpdateWriter:
traceback.format_exc(),
)
def _enqueue_tool_registry_upsert(
self,
kwargs: Optional[dict],
completion_response: Optional[Any],
hashed_token: Optional[str] = None,
team_id: Optional[str] = None,
) -> None:
"""
Extract tool names from the LLM request and response and enqueue them
for upsert into LiteLLM_ToolTable via ToolDiscoveryQueue.
Handles four sources:
- MCP tools: standard_logging_object.mcp_tool_call_metadata.namespaced_tool_name
- Response tool_calls (OpenAI / Anthropic pass-through converted to OpenAI format):
completion_response.choices[].message.tool_calls[].function.name
- Request tools array (OpenAI format): kwargs["tools"][].function.name
- Request tools array (Anthropic /messages format): kwargs["passthrough_logging_payload"]
["request_body"]["tools"][].name
"""
try:
if kwargs is None:
return
# Extract key_alias from kwargs metadata if available
key_alias: Optional[str] = None
_litellm_params = kwargs.get("litellm_params") or {}
_metadata = _litellm_params.get("metadata") or {}
key_alias = _metadata.get("user_api_key_alias") or None
def _enqueue(tool_name: str, origin: str = "user_defined") -> None:
self.tool_discovery_queue.add_update(
ToolDiscoveryQueueItem(
tool_name=tool_name,
origin=origin,
key_hash=hashed_token,
team_id=team_id,
key_alias=key_alias,
)
)
# --- MCP tool calls ---
sl_object = kwargs.get("standard_logging_object")
if sl_object is not None:
mcp_metadata = (
sl_object.get("metadata", {}) or {}
).get("mcp_tool_call_metadata")
if mcp_metadata and isinstance(mcp_metadata, dict):
tool_name = mcp_metadata.get("namespaced_tool_name") or mcp_metadata.get("name")
mcp_server_name = mcp_metadata.get("mcp_server_name")
if tool_name:
_enqueue(tool_name, origin=mcp_server_name or "user_defined")
# --- Tools from request body (OpenAI format: tools[].function.name) ---
request_tools = kwargs.get("tools") or []
for tool_def in request_tools:
if not isinstance(tool_def, dict):
continue
fn = tool_def.get("function") or {}
name = fn.get("name") if isinstance(fn, dict) else None
if name:
_enqueue(name)
# --- Tools from Anthropic /messages pass-through request body
# (Anthropic format: tools[].name, no "function" wrapper) ---
passthrough_payload = kwargs.get("passthrough_logging_payload") or {}
request_body = (
passthrough_payload.get("request_body")
if isinstance(passthrough_payload, dict)
else None
) or {}
for tool_def in request_body.get("tools") or []:
if not isinstance(tool_def, dict):
continue
name = tool_def.get("name")
if name:
_enqueue(name)
# --- Response tool_calls (OpenAI format; Anthropic pass-through converts tool_use here) ---
if completion_response is not None and hasattr(completion_response, "choices"):
for choice in completion_response.choices or []:
message = getattr(choice, "message", None)
if message is None:
continue
tool_calls = getattr(message, "tool_calls", None)
if not tool_calls:
continue
for tc in tool_calls:
fn = getattr(tc, "function", None)
if fn is None:
continue
tool_name = getattr(fn, "name", None)
if tool_name:
_enqueue(tool_name)
except Exception as e:
verbose_proxy_logger.debug(
"_enqueue_tool_registry_upsert error (non-blocking): %s", e
)
async def _update_key_db(
self,
response_cost: Optional[float],
@ -752,6 +872,25 @@ class DBSpendUpdateWriter:
daily_spend_transactions=daily_agent_spend_update_transactions,
)
################## Tool Registry Upserts ##################
await self._flush_tool_discovery_queue(prisma_client=prisma_client)
async def _flush_tool_discovery_queue(
self,
prisma_client: PrismaClient,
) -> None:
"""Flush ToolDiscoveryQueue and batch-upsert new tools into LiteLLM_ToolTable."""
from litellm.proxy.db.tool_registry_writer import batch_upsert_tools
try:
items = self.tool_discovery_queue.flush()
if items:
await batch_upsert_tools(prisma_client=prisma_client, items=items)
except Exception as e:
verbose_proxy_logger.debug(
"_flush_tool_discovery_queue error (non-blocking): %s", e
)
async def _commit_spend_updates_to_db( # noqa: PLR0915
self,
prisma_client: PrismaClient,

54
litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py Normal file
View file

@ -0,0 +1,54 @@
"""
In-memory buffer for tool registry upserts.
Unlike SpendUpdateQueue (which aggregates increments), ToolDiscoveryQueue
uses set-deduplication: each unique tool_name is only queued once per flush
cycle (~30s). The seen-set is cleared on every flush so that call_count
increments in subsequent cycles rather than stopping after the first flush.
"""
from typing import List, Set
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import ToolDiscoveryQueueItem
class ToolDiscoveryQueue:
"""
In-memory buffer for tool registry upserts.
Deduplicates by tool_name within each flush cycle: a tool is only queued
once per ~30s batch, so call_count increments once per flush cycle the
tool appears in (not once per invocation, but not once per pod lifetime
either). The seen-set is cleared on flush so subsequent batches can
re-count the same tool.
"""
def __init__(self) -> None:
self._seen_tool_names: Set[str] = set()
self._pending: List[ToolDiscoveryQueueItem] = []
def add_update(self, item: ToolDiscoveryQueueItem) -> None:
"""Enqueue a tool discovery item if tool_name has not been seen before."""
tool_name = item.get("tool_name", "")
if not tool_name:
return
if tool_name in self._seen_tool_names:
verbose_proxy_logger.debug(
"ToolDiscoveryQueue: skipping already-seen tool %s", tool_name
)
return
self._seen_tool_names.add(tool_name)
self._pending.append(item)
verbose_proxy_logger.debug(
"ToolDiscoveryQueue: queued new tool %s (origin=%s)",
tool_name,
item.get("origin"),
)
def flush(self) -> List[ToolDiscoveryQueueItem]:
"""Return and clear all pending items. Resets seen-set so the next
flush cycle can re-count the same tools."""
items, self._pending = self._pending, []
self._seen_tool_names.clear()
return items

179
litellm/proxy/db/tool_registry_writer.py Normal file
View file

@ -0,0 +1,179 @@
"""
DB helpers for LiteLLM_ToolTable — the global tool registry.
Tools are auto-discovered from LLM responses and upserted here.
Admins use the management endpoints to read and update call_policy.
NOTE: Uses raw SQL (query_raw / execute_raw) instead of Prisma model methods
because the generated Prisma Python client may not have LiteLLM_ToolTable
when running against an older generated schema.
"""
import uuid
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Dict, List, Optional
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import ToolDiscoveryQueueItem
from litellm.types.tool_management import LiteLLM_ToolTableRow, ToolCallPolicy
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
def _row_to_model(row: dict) -> LiteLLM_ToolTableRow:
return LiteLLM_ToolTableRow(
tool_id=row.get("tool_id", ""),
tool_name=row.get("tool_name", ""),
origin=row.get("origin"),
call_policy=row.get("call_policy", "untrusted"),
call_count=int(row.get("call_count") or 0),
assignments=row.get("assignments"),
key_hash=row.get("key_hash"),
team_id=row.get("team_id"),
key_alias=row.get("key_alias"),
created_at=row.get("created_at"),
updated_at=row.get("updated_at"),
created_by=row.get("created_by"),
updated_by=row.get("updated_by"),
)
async def batch_upsert_tools(
prisma_client: "PrismaClient",
items: List[ToolDiscoveryQueueItem],
) -> None:
"""
Batch-upsert tool registry rows via raw SQL.
On first insert: sets call_policy = "untrusted" (schema default), call_count = 1.
On conflict: increments call_count; preserves existing call_policy.
"""
if not items:
return
try:
data = [item for item in items if item.get("tool_name")]
if not data:
return
for item in data:
tool_name = item.get("tool_name", "")
origin = item.get("origin") or "user_defined"
created_by = item.get("created_by") or "system"
key_hash = item.get("key_hash")
team_id = item.get("team_id")
key_alias = item.get("key_alias")
now = datetime.now(timezone.utc).isoformat()
await prisma_client.db.execute_raw(
'INSERT INTO "LiteLLM_ToolTable" '
"(tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias, created_at, updated_at) "
"VALUES ($7, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6, $8, $8) "
"ON CONFLICT (tool_name) DO UPDATE SET "
"call_count = \"LiteLLM_ToolTable\".call_count + 1, "
"updated_at = $8",
tool_name,
origin,
created_by,
key_hash,
team_id,
key_alias,
str(uuid.uuid4()),
now,
)
verbose_proxy_logger.debug(
"tool_registry_writer: upserted %d tool(s)", len(data)
)
except Exception as e:
verbose_proxy_logger.error("tool_registry_writer batch_upsert_tools error: %s", e)
async def list_tools(
prisma_client: "PrismaClient",
call_policy: Optional[ToolCallPolicy] = None,
) -> List[LiteLLM_ToolTableRow]:
"""Return all tools, optionally filtered by call_policy."""
try:
if call_policy is not None:
rows = await prisma_client.db.query_raw(
'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, '
'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by '
'FROM "LiteLLM_ToolTable" WHERE call_policy = $1 ORDER BY created_at DESC',
call_policy,
)
else:
rows = await prisma_client.db.query_raw(
'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, '
'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by '
'FROM "LiteLLM_ToolTable" ORDER BY created_at DESC',
)
return [_row_to_model(row) for row in rows]
except Exception as e:
verbose_proxy_logger.error("tool_registry_writer list_tools error: %s", e)
return []
async def get_tool(
prisma_client: "PrismaClient",
tool_name: str,
) -> Optional[LiteLLM_ToolTableRow]:
"""Return a single tool row by tool_name."""
try:
rows = await prisma_client.db.query_raw(
'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, '
'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by '
'FROM "LiteLLM_ToolTable" WHERE tool_name = $1',
tool_name,
)
if not rows:
return None
return _row_to_model(rows[0])
except Exception as e:
verbose_proxy_logger.error("tool_registry_writer get_tool error: %s", e)
return None
async def update_tool_policy(
prisma_client: "PrismaClient",
tool_name: str,
call_policy: ToolCallPolicy,
updated_by: Optional[str],
) -> Optional[LiteLLM_ToolTableRow]:
"""Update the call_policy for a tool. Upserts the row if it does not exist yet."""
try:
_updated_by = updated_by or "system"
now = datetime.now(timezone.utc).isoformat()
await prisma_client.db.execute_raw(
'INSERT INTO "LiteLLM_ToolTable" (tool_id, tool_name, call_policy, created_by, updated_by, created_at, updated_at) '
"VALUES ($4, $1, $2, $3, $3, $5, $5) "
"ON CONFLICT (tool_name) DO UPDATE SET call_policy = $2, updated_by = $3, updated_at = $5",
tool_name,
call_policy,
_updated_by,
str(uuid.uuid4()),
now,
)
return await get_tool(prisma_client, tool_name)
except Exception as e:
verbose_proxy_logger.error("tool_registry_writer update_tool_policy error: %s", e)
return None
async def get_tools_by_names(
prisma_client: "PrismaClient",
tool_names: List[str],
) -> Dict[str, str]:
"""
Return a {tool_name: call_policy} map for the given tool names.
Used by the policy enforcement guardrail — single batch query, never N+1.
"""
if not tool_names:
return {}
try:
placeholders = ", ".join(f"${i+1}" for i in range(len(tool_names)))
rows = await prisma_client.db.query_raw(
f'SELECT tool_name, call_policy FROM "LiteLLM_ToolTable" WHERE tool_name IN ({placeholders})',
*tool_names,
)
return {row["tool_name"]: row["call_policy"] for row in rows}
except Exception as e:
verbose_proxy_logger.error("tool_registry_writer get_tools_by_names error: %s", e)
return {}

53
litellm/proxy/google_endpoints/endpoints.py
View file

@ -1,6 +1,10 @@
from datetime import datetime
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import ORJSONResponse, StreamingResponse
import litellm
from litellm._uuid import uuid
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
@ -17,7 +21,8 @@ router = APIRouter(
dependencies=[Depends(user_api_key_auth)],
)
@router.post(
"/models/{model_name:path}:generateContent", dependencies=[Depends(user_api_key_auth)]
"/models/{model_name:path}:generateContent",
dependencies=[Depends(user_api_key_auth)],
)
async def google_generate_content(
request: Request,
@ -36,12 +41,12 @@ async def google_generate_content(
data = await _read_request_body(request=request)
if "model" not in data:
data["model"] = model_name
# Extract generationConfig and pass it as config parameter
generation_config = data.pop("generationConfig", None)
if generation_config:
data["config"] = generation_config
# Add user authentication metadata for cost tracking
data = await add_litellm_data_to_request(
data=data,
@ -51,7 +56,19 @@ async def google_generate_content(
general_settings=general_settings,
version=version,
)
# Create logging object with full request metadata so callbacks (e.g. S3) get user/trace_id
data["litellm_call_id"] = request.headers.get(
"x-litellm-call-id", str(uuid.uuid4())
)
logging_obj, data = litellm.utils.function_setup(
original_function="agenerate_content",
rules_obj=litellm.utils.Rules(),
start_time=datetime.now(),
**data,
)
data["litellm_logging_obj"] = logging_obj
# call router
if llm_router is None:
raise HTTPException(status_code=500, detail="Router not initialized")
@ -103,6 +120,18 @@ async def google_stream_generate_content(
version=version,
)
# Create logging object with full request metadata so streaming END callbacks (e.g. S3) get user/trace_id
data["litellm_call_id"] = request.headers.get(
"x-litellm-call-id", str(uuid.uuid4())
)
logging_obj, data = litellm.utils.function_setup(
original_function="agenerate_content_stream",
rules_obj=litellm.utils.Rules(),
start_time=datetime.now(),
**data,
)
data["litellm_logging_obj"] = logging_obj
# call router
if llm_router is None:
raise HTTPException(status_code=500, detail="Router not initialized")
@ -247,11 +276,11 @@ async def create_interaction(
)
data = await _read_request_body(request=request)
# Default to gemini provider for interactions
if "custom_llm_provider" not in data:
data["custom_llm_provider"] = "gemini"
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
@ -301,7 +330,7 @@ async def get_interaction(
):
"""
Get an interaction by ID.
Per OpenAPI spec: GET /{api_version}/interactions/{interaction_id}
"""
from litellm.proxy.proxy_server import (
@ -319,7 +348,7 @@ async def get_interaction(
)
data = {"interaction_id": interaction_id, "custom_llm_provider": "gemini"}
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
@ -369,7 +398,7 @@ async def delete_interaction(
):
"""
Delete an interaction by ID.
Per OpenAPI spec: DELETE /{api_version}/interactions/{interaction_id}
"""
from litellm.proxy.proxy_server import (
@ -387,7 +416,7 @@ async def delete_interaction(
)
data = {"interaction_id": interaction_id, "custom_llm_provider": "gemini"}
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
@ -437,7 +466,7 @@ async def cancel_interaction(
):
"""
Cancel an interaction by ID.
Per OpenAPI spec: POST /{api_version}/interactions/{interaction_id}:cancel
"""
from litellm.proxy.proxy_server import (
@ -455,7 +484,7 @@ async def cancel_interaction(
)
data = {"interaction_id": interaction_id, "custom_llm_provider": "gemini"}
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(

145
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -2,6 +2,7 @@
CRUD ENDPOINTS FOR GUARDRAILS
"""
import concurrent.futures
import inspect
from typing import Any, Dict, List, Optional, Type, TypeVar, Union, cast
@ -11,9 +12,16 @@ from pydantic import BaseModel
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry
from litellm.proxy.guardrails.guardrail_hooks.custom_code.code_validator import (
CustomCodeValidationError,
validate_custom_code,
)
from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import (
get_custom_code_primitives,
)
from litellm.proxy.guardrails.usage_endpoints import router as guardrails_usage_router
from litellm.types.guardrails import (
PII_ENTITY_CATEGORIES_MAP,
@ -243,9 +251,11 @@ class CreateGuardrailRequest(BaseModel):
@router.post(
"/guardrails",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def create_guardrail(request: CreateGuardrailRequest):
async def create_guardrail(
request: CreateGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Create a new guardrail
@ -296,6 +306,12 @@ async def create_guardrail(request: CreateGuardrailRequest):
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@ -332,9 +348,12 @@ class UpdateGuardrailRequest(BaseModel):
@router.put(
"/guardrails/{guardrail_id}",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest):
async def update_guardrail(
guardrail_id: str,
request: UpdateGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Update an existing guardrail
@ -385,6 +404,12 @@ async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest):
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@ -429,9 +454,11 @@ async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest):
@router.delete(
"/guardrails/{guardrail_id}",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def delete_guardrail(guardrail_id: str):
async def delete_guardrail(
guardrail_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Delete a guardrail
@ -453,6 +480,12 @@ async def delete_guardrail(guardrail_id: str):
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@ -495,9 +528,12 @@ async def delete_guardrail(guardrail_id: str):
@router.patch(
"/guardrails/{guardrail_id}",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest):
async def patch_guardrail(
guardrail_id: str,
request: PatchGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Partially update an existing guardrail
@ -546,6 +582,12 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest):
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@ -1302,9 +1344,9 @@ async def get_provider_specific_params():
lakera_v2_fields = _get_fields_from_model(LakeraV2GuardrailConfigModel)
tool_permission_fields = _get_fields_from_model(ToolPermissionGuardrailConfigModel)
tool_permission_fields["ui_friendly_name"] = (
ToolPermissionGuardrailConfigModel.ui_friendly_name()
)
tool_permission_fields[
"ui_friendly_name"
] = ToolPermissionGuardrailConfigModel.ui_friendly_name()
# Return the provider-specific parameters
provider_params = {
@ -1364,10 +1406,12 @@ class TestCustomCodeGuardrailResponse(BaseModel):
@router.post(
"/guardrails/test_custom_code",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
response_model=TestCustomCodeGuardrailResponse,
)
async def test_custom_code_guardrail(request: TestCustomCodeGuardrailRequest):
async def test_custom_code_guardrail(
request: TestCustomCodeGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Test custom code guardrail logic without creating a guardrail.
@ -1440,63 +1484,26 @@ async def test_custom_code_guardrail(request: TestCustomCodeGuardrailRequest):
}
```
"""
import concurrent.futures
import re
from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import (
get_custom_code_primitives,
)
# Security validation patterns
FORBIDDEN_PATTERNS = [
# Import statements
(r"\bimport\s+", "import statements are not allowed"),
(r"\bfrom\s+\w+\s+import\b", "from...import statements are not allowed"),
(r"__import__\s*\(", "__import__() is not allowed"),
# Dangerous builtins
(r"\bexec\s*\(", "exec() is not allowed"),
(r"\beval\s*\(", "eval() is not allowed"),
(r"\bcompile\s*\(", "compile() is not allowed"),
(r"\bopen\s*\(", "open() is not allowed"),
(r"\bgetattr\s*\(", "getattr() is not allowed"),
(r"\bsetattr\s*\(", "setattr() is not allowed"),
(r"\bdelattr\s*\(", "delattr() is not allowed"),
(r"\bglobals\s*\(", "globals() is not allowed"),
(r"\blocals\s*\(", "locals() is not allowed"),
(r"\bvars\s*\(", "vars() is not allowed"),
(r"\bdir\s*\(", "dir() is not allowed"),
(r"\bbreakpoint\s*\(", "breakpoint() is not allowed"),
(r"\binput\s*\(", "input() is not allowed"),
# Dangerous dunder access
(r"__builtins__", "__builtins__ access is not allowed"),
(r"__globals__", "__globals__ access is not allowed"),
(r"__code__", "__code__ access is not allowed"),
(r"__subclasses__", "__subclasses__ access is not allowed"),
(r"__bases__", "__bases__ access is not allowed"),
(r"__mro__", "__mro__ access is not allowed"),
(r"__class__", "__class__ access is not allowed"),
(r"__dict__", "__dict__ access is not allowed"),
(r"__getattribute__", "__getattribute__ access is not allowed"),
(r"__reduce__", "__reduce__ access is not allowed"),
(r"__reduce_ex__", "__reduce_ex__ access is not allowed"),
# OS/system access
(r"\bos\.", "os module access is not allowed"),
(r"\bsys\.", "sys module access is not allowed"),
(r"\bsubprocess\.", "subprocess module access is not allowed"),
]
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to test custom code guardrails",
)
EXECUTION_TIMEOUT_SECONDS = 5
try:
# Step 0: Security validation - check for forbidden patterns
code = request.custom_code
for pattern, error_msg in FORBIDDEN_PATTERNS:
if re.search(pattern, code):
return TestCustomCodeGuardrailResponse(
success=False,
error=f"Security violation: {error_msg}",
error_type="compilation",
)
try:
validate_custom_code(request.custom_code)
except CustomCodeValidationError as e:
return TestCustomCodeGuardrailResponse(
success=False,
error=str(e),
error_type="compilation",
)
# Step 1: Compile the custom code with restricted environment
exec_globals = get_custom_code_primitives().copy()
@ -1612,10 +1619,10 @@ async def apply_guardrail(
from litellm.proxy.utils import handle_exception_on_proxy
try:
active_guardrail: Optional[CustomGuardrail] = (
GUARDRAIL_REGISTRY.get_initialized_guardrail_callback(
guardrail_name=request.guardrail_name
)
active_guardrail: Optional[
CustomGuardrail
] = GUARDRAIL_REGISTRY.get_initialized_guardrail_callback(
guardrail_name=request.guardrail_name
)
if active_guardrail is None:
raise HTTPException(

63
litellm/proxy/guardrails/guardrail_hooks/custom_code/code_validator.py Normal file
View file

@ -0,0 +1,63 @@
import re
from typing import List, Tuple
# Security validation patterns
FORBIDDEN_PATTERNS: List[Tuple[str, str]] = [
# Import statements
(r"\bimport\s+", "import statements are not allowed"),
(r"\bfrom\s+\w+\s+import\b", "from...import statements are not allowed"),
(r"__import__\s*\(", "__import__() is not allowed"),
# Dangerous builtins
(r"\bexec\s*\(", "exec() is not allowed"),
(r"\beval\s*\(", "eval() is not allowed"),
(r"\bcompile\s*\(", "compile() is not allowed"),
(r"\bopen\s*\(", "open() is not allowed"),
(r"\bgetattr\s*\(", "getattr() is not allowed"),
(r"\bsetattr\s*\(", "setattr() is not allowed"),
(r"\bdelattr\s*\(", "delattr() is not allowed"),
(r"\bglobals\s*\(", "globals() is not allowed"),
(r"\blocals\s*\(", "locals() is not allowed"),
(r"\bvars\s*\(", "vars() is not allowed"),
(r"\bdir\s*\(", "dir() is not allowed"),
(r"\bbreakpoint\s*\(", "breakpoint() is not allowed"),
(r"\binput\s*\(", "input() is not allowed"),
# Dangerous dunder access
(r"__builtins__", "__builtins__ access is not allowed"),
(r"__globals__", "__globals__ access is not allowed"),
(r"__code__", "__code__ access is not allowed"),
(r"__subclasses__", "__subclasses__ access is not allowed"),
(r"__bases__", "__bases__ access is not allowed"),
(r"__mro__", "__mro__ access is not allowed"),
(r"__class__", "__class__ access is not allowed"),
(r"__dict__", "__dict__ access is not allowed"),
(r"__getattribute__", "__getattribute__ access is not allowed"),
(r"__reduce__", "__reduce__ access is not allowed"),
(r"__reduce_ex__", "__reduce_ex__ access is not allowed"),
# OS/system access
(r"\bos\.", "os module access is not allowed"),
(r"\bsys\.", "sys module access is not allowed"),
(r"\bsubprocess\.", "subprocess module access is not allowed"),
(r"\bshutil\.", "shutil module access is not allowed"),
(r"\bctypes\.", "ctypes module access is not allowed"),
(r"\bsocket\.", "socket module access is not allowed"),
(r"\bpickle\.", "pickle module access is not allowed"),
]
class CustomCodeValidationError(Exception):
"""Raised when custom code fails security validation."""
pass
def validate_custom_code(code: str) -> None:
"""
Validate custom code against forbidden patterns.
Raises CustomCodeValidationError if any forbidden pattern is found.
"""
if not code:
return
for pattern, error_msg in FORBIDDEN_PATTERNS:
if re.search(pattern, code):
raise CustomCodeValidationError(f"Security violation: {error_msg}")

87
litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py
View file

@ -41,18 +41,19 @@ from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, cast
from fastapi import HTTPException
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (CustomGuardrail,
log_guardrail_information)
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.guardrails.guardrail_hooks.base import \
GuardrailConfigModel
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
from litellm.types.utils import GenericGuardrailAPIInputs
from .code_validator import CustomCodeValidationError, validate_custom_code
from .primitives import get_custom_code_primitives
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import \
Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
class CustomCodeGuardrailError(Exception):
@ -143,6 +144,33 @@ class CustomCodeGuardrail(CustomGuardrail):
"""Returns the config model for the UI."""
return CustomCodeGuardrailConfigModel
def _do_compile(self) -> None:
"""Internal compilation method without lock. Expected to run inside _compile_lock."""
# Create a restricted execution environment
# Only include our safe primitives
exec_globals = get_custom_code_primitives().copy()
# CRITICAL: Restrict __builtins__ to prevent sandbox escape
exec_globals["__builtins__"] = {}
# Execute the user code in the restricted environment
exec(compile(self.custom_code, "<guardrail>", "exec"), exec_globals)
# Extract the apply_guardrail function
if "apply_guardrail" not in exec_globals:
raise CustomCodeCompilationError(
"Custom code must define an 'apply_guardrail' function. "
"Expected signature: apply_guardrail(inputs, request_data, input_type)"
)
apply_fn = exec_globals["apply_guardrail"]
if not callable(apply_fn):
raise CustomCodeCompilationError(
"'apply_guardrail' must be a callable function"
)
self._compiled_function = apply_fn
def _compile_custom_code(self) -> None:
"""
Compile the custom code and extract the apply_guardrail function.
@ -154,27 +182,14 @@ class CustomCodeGuardrail(CustomGuardrail):
return
try:
# Create a restricted execution environment
# Only include our safe primitives
exec_globals = get_custom_code_primitives().copy()
# Step 1: Security validation — forbidden pattern check
try:
validate_custom_code(self.custom_code)
except CustomCodeValidationError as e:
raise CustomCodeCompilationError(str(e)) from e
# Execute the user code in the restricted environment
exec(compile(self.custom_code, "<guardrail>", "exec"), exec_globals)
# Extract the apply_guardrail function
if "apply_guardrail" not in exec_globals:
raise CustomCodeCompilationError(
"Custom code must define an 'apply_guardrail' function. "
"Expected signature: apply_guardrail(inputs, request_data, input_type)"
)
apply_fn = exec_globals["apply_guardrail"]
if not callable(apply_fn):
raise CustomCodeCompilationError(
"'apply_guardrail' must be a callable function"
)
self._compiled_function = apply_fn
# Step 2: Compile logic
self._do_compile()
verbose_proxy_logger.debug(
f"Custom code guardrail '{self.guardrail_name}' compiled successfully"
)
@ -390,6 +405,12 @@ class CustomCodeGuardrail(CustomGuardrail):
Raises:
CustomCodeCompilationError: If the new code fails to compile
"""
# Validate BEFORE acquiring lock / resetting state
try:
validate_custom_code(new_code)
except CustomCodeValidationError as e:
raise CustomCodeCompilationError(str(e)) from e
with self._compile_lock:
# Reset state
old_function = self._compiled_function
@ -399,12 +420,24 @@ class CustomCodeGuardrail(CustomGuardrail):
try:
self.custom_code = new_code
self._compile_custom_code()
self._do_compile()
verbose_proxy_logger.info(
f"Custom code guardrail '{self.guardrail_name}': Code updated successfully"
)
except SyntaxError as e:
# Rollback on failure
self.custom_code = old_code
self._compiled_function = old_function
self._compile_error = f"Syntax error in custom code: {e}"
raise CustomCodeCompilationError(self._compile_error) from e
except CustomCodeCompilationError:
# Rollback on failure
self.custom_code = old_code
self._compiled_function = old_function
raise
except Exception as e:
# Rollback on failure
self.custom_code = old_code
self._compiled_function = old_function
self._compile_error = f"Failed to compile custom code: {e}"
raise CustomCodeCompilationError(self._compile_error) from e

65
litellm/proxy/guardrails/guardrail_hooks/presidio.py
View file

@ -322,12 +322,6 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
analyze_payload,
)
async with session.post(analyze_url, json=analyze_payload) as response:
analyze_results = await response.json()
verbose_proxy_logger.debug("analyze_results: %s", analyze_results)
# Handle error responses from Presidio (e.g., {'error': 'No text provided'})
# Presidio may return a dict instead of a list when errors occur
def _fail_on_invalid_response(
reason: str,
) -> List[PresidioAnalyzeResponseItem]:
@ -347,6 +341,36 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
)
return []
async with session.post(
analyze_url,
json=analyze_payload,
headers={"Accept": "application/json"},
) as response:
# Validate HTTP status
if response.status >= 400:
error_body = await response.text()
return _fail_on_invalid_response(
f"HTTP {response.status} from Presidio analyzer: {error_body[:200]}"
)
# Validate Content-Type is JSON
content_type = getattr(
response,
"content_type",
response.headers.get("Content-Type", ""),
)
if "application/json" not in content_type:
error_body = await response.text()
return _fail_on_invalid_response(
f"expected application/json Content-Type but received '{content_type}'; body: '{error_body[:200]}'"
)
analyze_results = await response.json()
verbose_proxy_logger.debug("analyze_results: %s", analyze_results)
# Handle error responses from Presidio (e.g., {'error': 'No text provided'})
# Presidio may return a dict instead of a list when errors occur
if isinstance(analyze_results, dict):
if "error" in analyze_results:
return _fail_on_invalid_response(
@ -423,8 +447,29 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
}
async with session.post(
anonymize_url, json=anonymize_payload
anonymize_url,
json=anonymize_payload,
headers={"Accept": "application/json"},
) as response:
# Validate HTTP status
if response.status >= 400:
error_body = await response.text()
raise Exception(
f"Presidio anonymizer returned HTTP {response.status}: {error_body[:200]}"
)
# Validate Content-Type is JSON
content_type = getattr(
response,
"content_type",
response.headers.get("Content-Type", ""),
)
if "application/json" not in content_type:
error_body = await response.text()
raise Exception(
f"Presidio anonymizer returned non-JSON Content-Type '{content_type}'; body: '{error_body[:200]}'"
)
redacted_text = await response.json()
new_text = text
@ -456,7 +501,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
except Exception as e:
# Sanitize exception to avoid leaking the original text (which may
# contain API keys or other secrets) in error responses.
if "Invalid anonymizer response" in str(e):
error_str = str(e)
if (
"Invalid anonymizer response" in error_str
or "Presidio anonymizer returned" in error_str
):
raise
raise Exception(
f"Presidio PII anonymization failed: {type(e).__name__}"

16
litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py Normal file
View file

@ -0,0 +1,16 @@
import litellm
from litellm.types.guardrails import Guardrail, LitellmParams
def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail):
from litellm.proxy.guardrails.guardrail_hooks.tool_policy.tool_policy_guardrail import (
ToolPolicyGuardrail,
)
_callback = ToolPolicyGuardrail(
guardrail_name=guardrail.get("guardrail_name", "tool_policy"),
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
)
litellm.logging_callback_manager.add_litellm_callback(_callback)
return _callback

163
litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py Normal file
View file

@ -0,0 +1,163 @@
"""
Tool Policy Guardrail
Reads call_policy from LiteLLM_ToolTable and enforces it on LLM requests/responses.
Policy values:
"trusted" - allow through (no action)
"untrusted" - allow through (no action; default for newly discovered tools)
"blocked" - raise HTTPException, preventing the tool call
"dual_llm" - (Phase 3) send to second LLM for verification; currently treated as allowed
Configuration in proxy config YAML:
guardrails:
- guardrail_name: "tool_policy"
litellm_params:
guardrail: tool_policy
mode: post_call
or both pre and post call:
- guardrail_name: "tool_policy"
litellm_params:
guardrail: tool_policy
mode: during_call # runs before LLM and on response
"""
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional
from fastapi import HTTPException
from litellm._logging import verbose_proxy_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import TOOL_POLICY_CACHE_TTL_SECONDS
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
GUARDRAIL_NAME = "tool_policy"
class ToolPolicyGuardrail(CustomGuardrail):
"""
Guardrail that enforces per-tool call policies stored in LiteLLM_ToolTable.
Tools with call_policy="blocked" are rejected before/after the LLM call.
Tools with call_policy="trusted" or "untrusted" pass through unchanged.
"""
def __init__(self, **kwargs: Any) -> None:
if "supported_event_hooks" not in kwargs:
kwargs["supported_event_hooks"] = [
GuardrailEventHooks.pre_call,
GuardrailEventHooks.post_call,
GuardrailEventHooks.during_call,
]
super().__init__(**kwargs)
self._policy_cache: DualCache = DualCache()
@log_guardrail_information
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
"""
Enforce tool policies on both request tools and response tool_calls.
- input_type="request": check inputs["tools"] (tool definitions in the LLM request)
- input_type="response": check inputs["tool_calls"] (tool_calls in the LLM response)
Raises HTTPException (400) if any tool is "blocked".
"""
if input_type == "request":
tools = inputs.get("tools") or []
tool_names = [
t["function"]["name"]
for t in tools
if isinstance(t, dict)
and isinstance(t.get("function"), dict)
and t["function"].get("name")
]
else: # response
tool_calls = inputs.get("tool_calls") or []
tool_names = []
for tc in tool_calls:
fn = None
if isinstance(tc, dict):
fn = (tc.get("function") or {}).get("name")
elif hasattr(tc, "function"):
fn = getattr(tc.function, "name", None)
if fn:
tool_names.append(fn)
if not tool_names:
return inputs
policy_map = await self._get_policies_cached(tool_names)
blocked = [name for name in tool_names if policy_map.get(name) == "blocked"]
if blocked:
verbose_proxy_logger.warning(
"ToolPolicyGuardrail: blocking tool(s) %s (policy=blocked)", blocked
)
raise HTTPException(
status_code=400,
detail={
"error": "Violated tool policy",
"blocked_tools": blocked,
"message": f"Tool(s) {blocked} are blocked by policy.",
},
)
return inputs
async def _get_policies_cached(self, tool_names: List[str]) -> Dict[str, str]:
"""
Batch-fetch call_policy for the given tool names.
Caches per individual tool name (not per combination) so that adding
a new tool to a request doesn't invalidate the cached policies for all
the other tools already in the cache.
"""
from litellm.proxy.db.tool_registry_writer import get_tools_by_names
from litellm.proxy.proxy_server import prisma_client
if not tool_names or prisma_client is None:
return {}
result: Dict[str, str] = {}
cache_misses: List[str] = []
for name in tool_names:
cached = await self._policy_cache.async_get_cache(f"tool_policy:{name}")
if cached is not None and isinstance(cached, str):
result[name] = cached
else:
cache_misses.append(name)
if cache_misses:
fetched = await get_tools_by_names(
prisma_client=prisma_client, tool_names=cache_misses
)
for name, policy in fetched.items():
result[name] = policy
await self._policy_cache.async_set_cache(
key=f"tool_policy:{name}",
value=policy,
ttl=TOOL_POLICY_CACHE_TTL_SECONDS,
)
verbose_proxy_logger.debug(
"ToolPolicyGuardrail: fetched %d policies from DB (cache hits: %d)",
len(cache_misses),
len(tool_names) - len(cache_misses),
)
return result

12
litellm/proxy/health_check.py
View file

@ -283,12 +283,17 @@ async def perform_health_check(
model: Optional[str] = None,
cli_model: Optional[str] = None,
details: Optional[bool] = True,
model_id: Optional[str] = None,
max_concurrency: Optional[int] = None,
instrumentation_context: Optional[dict] = None,
):
"""
Perform a health check on the system.
When model_id is provided, only the deployment with that id is checked
(so models that share the same name but have different ids are checked separately).
When model (name) is provided, all deployments matching that name are checked.
Returns:
(bool): True if the health check passes, False otherwise.
"""
@ -314,7 +319,12 @@ async def perform_health_check(
cycle_start_time = time.monotonic()
requested_model_count = len(model_list)
if model is not None:
# Filter by model_id first so a single deployment is checked when id is specified
if model_id is not None:
_by_id = [x for x in model_list if (x.get("model_info") or {}).get("id") == model_id]
if _by_id:
model_list = _by_id
elif model is not None:
_new_model_list = [
x for x in model_list if x["litellm_params"]["model"] == model
]

3
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -243,7 +243,7 @@ async def health_services_endpoint( # noqa: PLR0915
service_in_success_callbacks = True
else:
for cb in litellm.success_callback:
if hasattr(cb, "callback_name") and cb.callback_name == service:
if getattr(cb, "callback_name", None) == service:
service_in_success_callbacks = True
break
cb_id = get_callback_identifier(cb)
@ -747,6 +747,7 @@ async def _perform_health_check_and_save(
model=target_model,
details=details,
max_concurrency=max_concurrency,
model_id=model_id,
)
# Optionally save health check result to database (non-blocking)

72
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -12,7 +12,7 @@ from litellm.litellm_core_utils.core_helpers import (
)
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import log_db_metrics
from litellm.proxy.auth.auth_checks import get_key_object, get_team_object, log_db_metrics
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.utils import ProxyUpdateSpend
from litellm.types.utils import (
@ -76,6 +76,10 @@ class _ProxyDBLogger(CustomLogger):
traceback_str=traceback_str,
)
_metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(
metadata=_metadata,
)
existing_metadata: dict = request_data.get("metadata", None) or {}
existing_metadata.update(_metadata)
@ -255,6 +259,72 @@ class _ProxyDBLogger(CustomLogger):
"Error in tracking cost callback - %s", str(e)
)
@staticmethod
async def _enrich_failure_metadata_with_key_info(metadata: dict) -> dict:
"""
Enriches failure spend log metadata by looking up the key object (and team object)
from cache/DB when key fields are missing.
This handles two scenarios:
1. Auth errors (401): UserAPIKeyAuth is created with only api_key set, all other
fields are null. We look up the full key object to fill in alias, user_id,
team_id, etc.
2. Post-auth failures (provider errors, rate limits): key fields are populated
but team_alias is missing because LiteLLM_VerificationTokenView SQL view
doesn't include it. We look up the team object to fill in team_alias.
"""
api_key_hash = metadata.get("user_api_key")
if not api_key_hash:
return metadata
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
# Step 1: If key fields are missing, look up the full key object
if metadata.get("user_api_key_alias") is None:
try:
key_obj = await get_key_object(
hashed_token=api_key_hash,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if metadata.get("user_api_key_alias") is None:
metadata["user_api_key_alias"] = key_obj.key_alias
if metadata.get("user_api_key_user_id") is None:
metadata["user_api_key_user_id"] = key_obj.user_id
if metadata.get("user_api_key_team_id") is None:
metadata["user_api_key_team_id"] = key_obj.team_id
if metadata.get("user_api_key_org_id") is None:
metadata["user_api_key_org_id"] = key_obj.org_id
except Exception:
verbose_proxy_logger.debug(
"Failed to enrich failure metadata with key info for api_key=%s",
api_key_hash,
)
# Step 2: If team_id is known but team_alias is missing, look up the team object
team_id = metadata.get("user_api_key_team_id")
if team_id and metadata.get("user_api_key_team_alias") is None:
try:
team_obj = await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if team_obj.team_alias is not None:
metadata["user_api_key_team_alias"] = team_obj.team_alias
except Exception:
verbose_proxy_logger.debug(
"Failed to enrich failure metadata with team_alias for team_id=%s",
team_id,
)
return metadata
@staticmethod
def _should_track_errors_in_db():
"""

72
litellm/proxy/litellm_pre_call_utils.py
View file

@ -10,10 +10,16 @@ import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.proxy._types import (AddTeamCallback, CommonProxyErrors,
LitellmDataForBackendLLMCall,
LitellmUserRoles, SpecialHeaders,
TeamCallbackMetadata, UserAPIKeyAuth)
from litellm.proxy._types import (
AddTeamCallback,
CommonProxyErrors,
LitellmDataForBackendLLMCall,
LitellmUserRoles,
SpecialHeaders,
TeamCallbackMetadata,
UserAPIKeyAuth,
)
from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
# Cache special headers as a frozenset for O(1) lookup performance
_SPECIAL_HEADERS_CACHE = frozenset(
@ -22,9 +28,12 @@ _SPECIAL_HEADERS_CACHE = frozenset(
from litellm.router import Router
from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS
from litellm.types.services import ServiceTypes
from litellm.types.utils import (LlmProviders, ProviderSpecificHeader,
StandardLoggingUserAPIKeyMetadata,
SupportedCacheControls)
from litellm.types.utils import (
LlmProviders,
ProviderSpecificHeader,
StandardLoggingUserAPIKeyMetadata,
SupportedCacheControls,
)
service_logger_obj = ServiceLogging() # used for tracking latency on OTEL
@ -227,7 +236,9 @@ def _get_dynamic_logging_metadata(
def clean_headers(
headers: Headers, litellm_key_header_name: Optional[str] = None
headers: Headers,
litellm_key_header_name: Optional[str] = None,
forward_llm_provider_auth_headers: bool = False,
) -> dict:
"""
Removes litellm api key from headers
@ -237,15 +248,18 @@ def clean_headers(
clean_headers = {}
litellm_key_lower = (
litellm_key_header_name.lower() if litellm_key_header_name is not None else None
)
)
for header, value in headers.items():
header_lower = header.lower()
# Preserve Authorization header if it contains Anthropic OAuth token (sk-ant-oat*)
# This allows OAuth tokens to be forwarded to Anthropic-compatible providers
# via add_provider_specific_headers_to_request()
if header_lower == "authorization" and is_anthropic_oauth_key(value):
clean_headers[header] = value
elif forward_llm_provider_auth_headers and header_lower in _SPECIAL_HEADERS_CACHE:
if litellm_key_lower and header_lower == litellm_key_lower:
continue
if header_lower == "authorization":
continue
clean_headers[header] = value
# Check if header should be excluded: either in special headers cache or matches custom litellm key
elif header_lower not in _SPECIAL_HEADERS_CACHE and (
litellm_key_lower is None or header_lower != litellm_key_lower
@ -653,7 +667,8 @@ class LiteLLMProxyRequestSetup:
return data
from litellm.proxy._types import (
LiteLLM_ManagementEndpoint_MetadataFields,
LiteLLM_ManagementEndpoint_MetadataFields_Premium)
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
)
# ignore any special fields
added_metadata = {}
@ -824,7 +839,12 @@ async def add_litellm_data_to_request( # noqa: PLR0915
from litellm.proxy.proxy_server import llm_router, premium_user
from litellm.types.proxy.litellm_pre_call_utils import SecretFields
_raw_headers: Dict[str, str] = dict(request.headers)
_raw_headers: Dict[str, str] = _safe_get_request_headers(request)
forward_llm_auth = False
if general_settings:
forward_llm_auth = general_settings.get("forward_llm_provider_auth_headers", False)
_headers: Dict[str, str] = clean_headers(
request.headers,
litellm_key_header_name=(
@ -832,7 +852,10 @@ async def add_litellm_data_to_request( # noqa: PLR0915
if general_settings is not None
else None
),
forward_llm_provider_auth_headers=forward_llm_auth,
)
verbose_proxy_logger.debug(f"Request Headers: {_headers}")
verbose_proxy_logger.debug(f"Raw Headers: {_raw_headers}")
##########################################################
# Init - Proxy Server Request
@ -1478,8 +1501,7 @@ async def move_guardrails_to_metadata(
# Only check policy engine if no local config (avoid import + registry lookup)
if not (has_key_config or has_team_config or has_request_config):
from litellm.proxy.policy_engine.policy_registry import \
get_policy_registry
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
if not get_policy_registry().is_initialized():
# Nothing configured anywhere - clean up request body fields and return
@ -1543,16 +1565,14 @@ async def move_guardrails_to_metadata(
def _is_policy_version_id(s: str) -> bool:
"""Return True if string is a policy version ID (starts with policy_<uuid> prefix)."""
from litellm.proxy.policy_engine.policy_registry import \
POLICY_VERSION_ID_PREFIX
from litellm.proxy.policy_engine.policy_registry import POLICY_VERSION_ID_PREFIX
return isinstance(s, str) and s.startswith(POLICY_VERSION_ID_PREFIX)
def _extract_policy_id(s: str) -> Optional[str]:
"""Extract raw UUID from policy_<uuid> string, or None if not a valid version ID."""
from litellm.proxy.policy_engine.policy_registry import \
POLICY_VERSION_ID_PREFIX
from litellm.proxy.policy_engine.policy_registry import POLICY_VERSION_ID_PREFIX
if not _is_policy_version_id(s):
return None
@ -1573,9 +1593,10 @@ def _match_and_track_policies(
"""
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_utils.callback_utils import (
add_policy_sources_to_metadata, add_policy_to_applied_policies_header)
from litellm.proxy.policy_engine.attachment_registry import \
get_attachment_registry
add_policy_sources_to_metadata,
add_policy_to_applied_policies_header,
)
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
# Get matching policies via attachments (with match reasons for attribution)
@ -1720,8 +1741,7 @@ async def add_guardrails_from_policy_engine(
user_api_key_dict: The user's API key authentication info
"""
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_utils.http_parsing_utils import \
get_tags_from_request_body
from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.types.proxy.policy_engine import PolicyMatchContext

2
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -2535,6 +2535,7 @@ async def generate_key_helper_fn( # noqa: PLR0915
user_id: Optional[str] = None,
user_alias: Optional[str] = None,
team_id: Optional[str] = None,
agent_id: Optional[str] = None,
user_email: Optional[str] = None,
user_role: Optional[str] = None,
max_parallel_requests: Optional[int] = None,
@ -2668,6 +2669,7 @@ async def generate_key_helper_fn( # noqa: PLR0915
"max_budget": key_max_budget,
"user_id": user_id,
"team_id": team_id,
"agent_id": agent_id,
"project_id": project_id,
"max_parallel_requests": max_parallel_requests,
"metadata": metadata_json,

3
litellm/proxy/management_endpoints/scim/scim_v2.py
View file

@ -21,6 +21,7 @@ from typing_extensions import TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
from litellm._uuid import uuid
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import (
@ -724,7 +725,7 @@ async def get_service_provider_config(request: Request):
"SCIM ServiceProviderConfig request: method=%s url=%s headers=%s",
request.method,
request.url,
dict(request.headers),
_safe_get_request_headers(request),
)
meta = {
"resourceType": "ServiceProviderConfig",

149
litellm/proxy/management_endpoints/tool_management_endpoints.py Normal file
View file

@ -0,0 +1,149 @@
"""
TOOL POLICY MANAGEMENT
All /tool management endpoints
GET /v1/tool/list - List all discovered tools and their policies
GET /v1/tool/{tool_name} - Get a single tool's details
POST /v1/tool/policy - Update the call_policy for a tool
"""
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.types.tool_management import (
LiteLLM_ToolTableRow,
ToolCallPolicy,
ToolListResponse,
ToolPolicyUpdateRequest,
ToolPolicyUpdateResponse,
)
router = APIRouter()
@router.get(
"/v1/tool/list",
tags=["tool management"],
dependencies=[Depends(user_api_key_auth)],
response_model=ToolListResponse,
)
async def list_tools(
call_policy: Optional[ToolCallPolicy] = None,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
List all auto-discovered tools and their call policies.
Parameters:
- call_policy: Optional filter — one of "trusted", "untrusted", "dual_llm", "blocked"
"""
from litellm.proxy.db.tool_registry_writer import list_tools as db_list_tools
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
try:
tools = await db_list_tools(prisma_client=prisma_client, call_policy=call_policy)
return ToolListResponse(tools=tools, total=len(tools))
except Exception as e:
verbose_proxy_logger.exception("Error listing tools: %s", e)
raise HTTPException(status_code=500, detail=str(e))
@router.get(
"/v1/tool/{tool_name:path}",
tags=["tool management"],
dependencies=[Depends(user_api_key_auth)],
response_model=LiteLLM_ToolTableRow,
)
async def get_tool(
tool_name: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get details for a single tool.
Parameters:
- tool_name: The tool name (supports namespaced names with slashes)
"""
from litellm.proxy.db.tool_registry_writer import get_tool as db_get_tool
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
try:
tool = await db_get_tool(prisma_client=prisma_client, tool_name=tool_name)
if tool is None:
raise HTTPException(
status_code=404, detail=f"Tool '{tool_name}' not found"
)
return tool
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception("Error getting tool: %s", e)
raise HTTPException(status_code=500, detail=str(e))
@router.post(
"/v1/tool/policy",
tags=["tool management"],
dependencies=[Depends(user_api_key_auth)],
response_model=ToolPolicyUpdateResponse,
)
async def update_tool_policy(
data: ToolPolicyUpdateRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Set the call policy for a tool.
Parameters:
- tool_name: str - The tool to update
- call_policy: "trusted" | "untrusted" | "dual_llm" | "blocked"
Setting a tool to "blocked" will cause the ToolPolicyGuardrail to remove
that tool_call from LLM responses before returning them to the client.
"""
from litellm.proxy.db.tool_registry_writer import (
update_tool_policy as db_update_tool_policy,
)
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
)
try:
updated = await db_update_tool_policy(
prisma_client=prisma_client,
tool_name=data.tool_name,
call_policy=data.call_policy,
updated_by=user_api_key_dict.user_id,
)
if updated is None:
raise HTTPException(
status_code=500, detail=f"Failed to update policy for tool '{data.tool_name}'"
)
return ToolPolicyUpdateResponse(
tool_name=updated.tool_name,
call_policy=updated.call_policy,
updated=True,
)
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception("Error updating tool policy: %s", e)
raise HTTPException(status_code=500, detail=str(e))

9
litellm/proxy/management_endpoints/usage_endpoints/__init__.py Normal file
View file

@ -0,0 +1,9 @@
"""
Usage endpoints package.
Re-exports the router from endpoints module.
"""
from litellm.proxy.management_endpoints.usage_endpoints.endpoints import ( # noqa: F401
router,
)

578
litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py Normal file
View file

@ -0,0 +1,578 @@
"""
AI Usage Chat - uses LLM tool calling to answer questions about
usage/spend data by querying the aggregated daily activity endpoints.
"""
import json
from datetime import date
from typing import Any, AsyncIterator, Callable, Dict, List, Literal, Optional
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL
from litellm.types.proxy.management_endpoints.common_daily_activity import (
SpendAnalyticsPaginatedResponse,
)
from typing_extensions import TypedDict
# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
USAGE_AI_TEMPERATURE = 0.2
TABLE_DAILY_USER_SPEND = "litellm_dailyuserspend"
TABLE_DAILY_TEAM_SPEND = "litellm_dailyteamspend"
TABLE_DAILY_TAG_SPEND = "litellm_dailytagspend"
ENTITY_FIELD_USER = "user_id"
ENTITY_FIELD_TEAM = "team_id"
ENTITY_FIELD_TAG = "tag"
PAGINATED_PAGE_SIZE = 200
MAX_CHAT_MESSAGES = 20
TOP_N_MODELS = 15
TOP_N_PROVIDERS = 10
TOP_N_KEYS = 10
# ---------------------------------------------------------------------------
# Types
# ---------------------------------------------------------------------------
class SSEStatusEvent(TypedDict):
type: Literal["status"]
message: str
class SSEToolCallEvent(TypedDict, total=False):
type: Literal["tool_call"]
tool_name: str
tool_label: str
arguments: Dict[str, str]
status: Literal["running", "complete", "error"]
error: str
class SSEChunkEvent(TypedDict):
type: Literal["chunk"]
content: str
class SSEDoneEvent(TypedDict):
type: Literal["done"]
class SSEErrorEvent(TypedDict):
type: Literal["error"]
message: str
SSEEvent = (
SSEStatusEvent | SSEToolCallEvent | SSEChunkEvent | SSEDoneEvent | SSEErrorEvent
)
class ToolHandler(TypedDict):
fetch: Callable[..., Any]
summarise: Callable[[Dict[str, Any]], str]
label: str
# ---------------------------------------------------------------------------
# Tool definitions (OpenAI function-calling schema)
# ---------------------------------------------------------------------------
_DATE_PARAMS = {
"start_date": {"type": "string", "description": "Start date in YYYY-MM-DD format"},
"end_date": {"type": "string", "description": "End date in YYYY-MM-DD format"},
}
_TOOL_USAGE = {
"type": "function",
"function": {
"name": "get_usage_data",
"description": (
"Fetch aggregated global usage/spend data. Returns daily spend, "
"token counts, request counts, and breakdowns by model, provider, "
"and API key. Use for overall spend, top models, top providers."
),
"parameters": {
"type": "object",
"properties": {
**_DATE_PARAMS,
"user_id": {
"type": "string",
"description": "Optional user ID filter. Omit for global view.",
},
},
"required": ["start_date", "end_date"],
},
},
}
_TOOL_TEAM = {
"type": "function",
"function": {
"name": "get_team_usage_data",
"description": (
"Fetch usage/spend data broken down by team. Use for questions "
"like 'which team spends the most' or 'show me team X usage'."
),
"parameters": {
"type": "object",
"properties": {
**_DATE_PARAMS,
"team_ids": {
"type": "string",
"description": "Optional comma-separated team IDs. Omit for all teams.",
},
},
"required": ["start_date", "end_date"],
},
},
}
_TOOL_TAG = {
"type": "function",
"function": {
"name": "get_tag_usage_data",
"description": (
"Fetch usage/spend data broken down by tag. Tags are labels "
"attached to requests (features, environments, credentials)."
),
"parameters": {
"type": "object",
"properties": {
**_DATE_PARAMS,
"tags": {
"type": "string",
"description": "Optional comma-separated tag names. Omit for all tags.",
},
},
"required": ["start_date", "end_date"],
},
},
}
TOOLS_BASE = [_TOOL_USAGE]
TOOLS_ADMIN = [_TOOL_USAGE, _TOOL_TEAM, _TOOL_TAG]
def get_tools_for_role(is_admin: bool) -> List[Dict[str, Any]]:
"""Return the tool list appropriate for the user's role."""
return TOOLS_ADMIN if is_admin else TOOLS_BASE
_SYSTEM_PROMPT_BASE = (
"You are an AI assistant embedded in the LiteLLM Usage dashboard. "
"You help users understand their LLM API spend and usage data.\n\n"
"ALWAYS call the appropriate tool(s) first to fetch data before answering. "
"You may call multiple tools if the question spans different dimensions.\n\n"
"Guidelines:\n"
"- Be concise and specific. Use exact numbers from the data.\n"
"- Format costs as dollar amounts (e.g. $12.34).\n"
"- When comparing entities, show a ranked list.\n"
"- If data is empty or no results found, say so clearly.\n"
"- Do not hallucinate data — only use what the tools return.\n"
"- Today's date will be provided below. Use it to interpret relative dates "
"like 'this week', 'this month', 'last 7 days', etc."
)
_TOOL_DESCRIPTIONS_ADMIN = (
"You have access to these tools:\n"
"- `get_usage_data`: Global/user-level usage (spend, models, providers, API keys)\n"
"- `get_team_usage_data`: Team-level usage breakdown\n"
"- `get_tag_usage_data`: Tag-level usage breakdown\n\n"
)
_TOOL_DESCRIPTIONS_BASE = (
"You have access to this tool:\n"
"- `get_usage_data`: Your usage data (spend, models, providers, API keys)\n\n"
)
def _build_system_prompt(is_admin: bool) -> str:
"""Build role-appropriate system prompt with today's date."""
tool_desc = _TOOL_DESCRIPTIONS_ADMIN if is_admin else _TOOL_DESCRIPTIONS_BASE
return (
f"{_SYSTEM_PROMPT_BASE}\n\n{tool_desc}"
f"Today's date: {date.today().isoformat()}"
)
# keep a public reference for test assertions
SYSTEM_PROMPT = _SYSTEM_PROMPT_BASE
# ---------------------------------------------------------------------------
# Data fetchers
# ---------------------------------------------------------------------------
def _parse_csv_ids(raw: Optional[str]) -> Optional[List[str]]:
if not raw:
return None
return [t.strip() for t in raw.split(",") if t.strip()]
async def _query_activity(
table_name: str,
entity_id_field: str,
entity_id: Optional[Any],
start_date: str,
end_date: str,
*,
use_aggregated: bool = False,
) -> SpendAnalyticsPaginatedResponse:
"""Shared helper that calls the daily activity query layer."""
from litellm.proxy.management_endpoints.common_daily_activity import (
get_daily_activity,
get_daily_activity_aggregated,
)
from litellm.proxy.proxy_server import prisma_client
if use_aggregated:
return await get_daily_activity_aggregated(
prisma_client=prisma_client,
table_name=table_name,
entity_id_field=entity_id_field,
entity_id=entity_id,
entity_metadata_field=None,
start_date=start_date,
end_date=end_date,
model=None,
api_key=None,
)
return await get_daily_activity(
prisma_client=prisma_client,
table_name=table_name,
entity_id_field=entity_id_field,
entity_id=entity_id,
entity_metadata_field=None,
start_date=start_date,
end_date=end_date,
model=None,
api_key=None,
page=1,
page_size=PAGINATED_PAGE_SIZE,
)
async def _fetch_usage_data(
start_date: str, end_date: str, user_id: Optional[str] = None
) -> Dict[str, Any]:
resp = await _query_activity(
TABLE_DAILY_USER_SPEND,
ENTITY_FIELD_USER,
user_id,
start_date,
end_date,
use_aggregated=True,
)
return resp.model_dump(mode="json")
async def _fetch_team_usage_data(
start_date: str, end_date: str, team_ids: Optional[str] = None
) -> Dict[str, Any]:
resp = await _query_activity(
TABLE_DAILY_TEAM_SPEND,
ENTITY_FIELD_TEAM,
_parse_csv_ids(team_ids),
start_date,
end_date,
)
return resp.model_dump(mode="json")
async def _fetch_tag_usage_data(
start_date: str, end_date: str, tags: Optional[str] = None
) -> Dict[str, Any]:
resp = await _query_activity(
TABLE_DAILY_TAG_SPEND,
ENTITY_FIELD_TAG,
_parse_csv_ids(tags),
start_date,
end_date,
)
return resp.model_dump(mode="json")
# ---------------------------------------------------------------------------
# Summarisers — convert raw JSON to concise text the LLM can reason over
# ---------------------------------------------------------------------------
def _accumulate_breakdown(
results: List[Dict[str, Any]], dimension: str, fields: List[str]
) -> Dict[str, Dict[str, float]]:
"""Aggregate a single breakdown dimension across days."""
totals: Dict[str, Dict[str, float]] = {}
for day in results:
for key, entry in day.get("breakdown", {}).get(dimension, {}).items():
if key not in totals:
totals[key] = {f: 0.0 for f in fields}
m = entry.get("metrics", {})
for f in fields:
totals[key][f] += m.get(f, 0)
return totals
def _ranked_lines(
totals: Dict[str, Dict[str, float]],
fmt: Callable[[str, Dict[str, float]], str],
limit: int,
) -> List[str]:
"""Sort by spend descending, format each entry, and truncate."""
return [
fmt(name, vals)
for name, vals in sorted(totals.items(), key=lambda x: -x[1].get("spend", 0))[
:limit
]
]
def _summarise_usage_data(data: Dict[str, Any]) -> str:
meta = data.get("metadata", {})
results = data.get("results", [])
header = (
f"Total Spend: ${meta.get('total_spend', 0):.4f}\n"
f"Total Requests: {meta.get('total_api_requests', 0)}\n"
f"Successful: {meta.get('total_successful_requests', 0)} | "
f"Failed: {meta.get('total_failed_requests', 0)}\n"
f"Total Tokens: {meta.get('total_tokens', 0)}"
)
models = _accumulate_breakdown(
results, "models", ["spend", "api_requests", "total_tokens"]
)
providers = _accumulate_breakdown(results, "providers", ["spend", "api_requests"])
model_lines = _ranked_lines(
models,
lambda n, d: f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs, {int(d['total_tokens'])} tokens)",
TOP_N_MODELS,
)
provider_lines = _ranked_lines(
providers,
lambda n, d: f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs)",
TOP_N_PROVIDERS,
)
sections = [header, ""]
sections += ["Top Models by Spend:"] + (model_lines or [" (no data)"]) + [""]
sections += ["Top Providers by Spend:"] + (provider_lines or [" (no data)"])
return "\n".join(sections)
def _summarise_entity_data(data: Dict[str, Any], entity_label: str) -> str:
"""Summarise team/tag entity usage data."""
results = data.get("results", [])
if not results:
return f"No {entity_label} usage data found for the given date range."
totals: Dict[str, Dict[str, Any]] = {}
for day in results:
for eid, entry in day.get("breakdown", {}).get("entities", {}).items():
if eid not in totals:
alias = entry.get("metadata", {}).get("alias", eid)
totals[eid] = {"alias": alias, "spend": 0.0, "requests": 0, "tokens": 0}
m = entry.get("metrics", {})
totals[eid]["spend"] += m.get("spend", 0)
totals[eid]["requests"] += m.get("api_requests", 0)
totals[eid]["tokens"] += m.get("total_tokens", 0)
lines = [f"{entity_label} Usage ({len(totals)} {entity_label.lower()}s):", ""]
for eid, d in sorted(totals.items(), key=lambda x: -x[1]["spend"]):
label = d["alias"] if d["alias"] != eid else eid
lines.append(
f"- {label} (ID: {eid}): ${d['spend']:.4f} | "
f"{int(d['requests'])} reqs | {int(d['tokens'])} tokens"
)
return "\n".join(lines)
# ---------------------------------------------------------------------------
# Tool dispatch registry
# ---------------------------------------------------------------------------
TOOL_HANDLERS: Dict[str, ToolHandler] = {
"get_usage_data": ToolHandler(
fetch=_fetch_usage_data,
summarise=_summarise_usage_data,
label="global usage data",
),
"get_team_usage_data": ToolHandler(
fetch=_fetch_team_usage_data,
summarise=lambda data: _summarise_entity_data(data, "Team"),
label="team usage data",
),
"get_tag_usage_data": ToolHandler(
fetch=_fetch_tag_usage_data,
summarise=lambda data: _summarise_entity_data(data, "Tag"),
label="tag usage data",
),
}
# ---------------------------------------------------------------------------
# SSE streaming
# ---------------------------------------------------------------------------
def _sse(event: SSEEvent) -> str:
return f"data: {json.dumps(event)}\n\n"
def _resolve_fetch_kwargs(
fn_name: str,
fn_args: Dict[str, str],
user_id: Optional[str],
is_admin: bool,
) -> Dict[str, Any]:
"""Build keyword arguments for a tool's fetch function."""
start_date = fn_args.get("start_date", "")
end_date = fn_args.get("end_date", "")
if not start_date or not end_date:
raise ValueError("Missing required start_date or end_date from tool arguments")
kwargs: Dict[str, Any] = {"start_date": start_date, "end_date": end_date}
if fn_name == "get_usage_data":
if not is_admin:
kwargs["user_id"] = user_id
elif fn_args.get("user_id"):
kwargs["user_id"] = fn_args["user_id"]
elif fn_name == "get_team_usage_data" and fn_args.get("team_ids"):
kwargs["team_ids"] = fn_args["team_ids"]
elif fn_name == "get_tag_usage_data" and fn_args.get("tags"):
kwargs["tags"] = fn_args["tags"]
return kwargs
async def _execute_tool_call(
handler: ToolHandler,
fn_name: str,
fn_args: Dict[str, str],
user_id: Optional[str],
is_admin: bool,
) -> str:
"""Run a single tool and return the summarised result text."""
kwargs = _resolve_fetch_kwargs(fn_name, fn_args, user_id, is_admin)
raw_data = await handler["fetch"](**kwargs)
return handler["summarise"](raw_data)
async def _process_tool_call(
tc: Any,
chat_messages: List[Dict[str, Any]],
user_id: Optional[str],
is_admin: bool,
) -> AsyncIterator[str]:
"""Execute a single tool call, yielding SSE events for status."""
fn_name = tc.function.name
fn_args = json.loads(tc.function.arguments)
allowed_names = {t["function"]["name"] for t in get_tools_for_role(is_admin)}
handler = TOOL_HANDLERS.get(fn_name)
if fn_name not in allowed_names or not handler:
chat_messages.append(
{
"role": "tool",
"tool_call_id": tc.id,
"content": f"Tool not available: {fn_name}",
}
)
return
tool_event_base = {
"type": "tool_call",
"tool_name": fn_name,
"tool_label": handler["label"],
"arguments": fn_args,
}
yield _sse({**tool_event_base, "status": "running"})
try:
tool_result = await _execute_tool_call(
handler, fn_name, fn_args, user_id, is_admin
)
yield _sse({**tool_event_base, "status": "complete"})
except Exception as e:
verbose_proxy_logger.error("Tool %s failed: %s", fn_name, e)
tool_result = f"Error fetching {handler['label']}. Please try again."
yield _sse({**tool_event_base, "status": "error"})
chat_messages.append(
{"role": "tool", "tool_call_id": tc.id, "content": tool_result}
)
async def _stream_final_response(
model: str, chat_messages: List[Dict[str, Any]]
) -> AsyncIterator[str]:
"""Stream the final LLM response after tool results are appended."""
yield _sse({"type": "status", "message": "Analyzing results..."})
response = await litellm.acompletion(
model=model,
messages=chat_messages,
stream=True,
temperature=USAGE_AI_TEMPERATURE,
)
async for chunk in response:
delta = chunk.choices[0].delta.content
if delta:
yield _sse({"type": "chunk", "content": delta})
async def stream_usage_ai_chat(
messages: List[Dict[str, str]],
model: Optional[str] = None,
user_id: Optional[str] = None,
is_admin: bool = False,
) -> AsyncIterator[str]:
"""Stream SSE events: status → tool_call → chunk → done."""
resolved_model = (model or "").strip() or DEFAULT_COMPETITOR_DISCOVERY_MODEL
truncated = (
messages[-MAX_CHAT_MESSAGES:] if len(messages) > MAX_CHAT_MESSAGES else messages
)
chat_messages: List[Dict[str, Any]] = [
{"role": "system", "content": _build_system_prompt(is_admin)},
*truncated,
]
try:
yield _sse({"type": "status", "message": "Thinking..."})
tools = get_tools_for_role(is_admin)
response = await litellm.acompletion(
model=resolved_model,
messages=chat_messages,
tools=tools,
temperature=USAGE_AI_TEMPERATURE,
)
choice = response.choices[0] # type: ignore
if not choice.message.tool_calls:
if choice.message.content:
yield _sse({"type": "chunk", "content": choice.message.content})
yield _sse({"type": "done"})
return
chat_messages.append(choice.message.model_dump())
for tc in choice.message.tool_calls:
async for event in _process_tool_call(tc, chat_messages, user_id, is_admin):
yield event
async for event in _stream_final_response(resolved_model, chat_messages):
yield event
yield _sse({"type": "done"})
except Exception as e:
verbose_proxy_logger.error("AI usage chat failed: %s", e)
yield _sse(
{
"type": "error",
"message": "An internal error occurred. Please try again.",
}
)

65
litellm/proxy/management_endpoints/usage_endpoints/endpoints.py Normal file
View file

@ -0,0 +1,65 @@
"""
USAGE AI CHAT ENDPOINTS
/usage/ai/chat - Stream AI chat responses about usage data
"""
from typing import List, Literal, Optional
from fastapi import APIRouter, Depends, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
router = APIRouter()
class ChatMessage(BaseModel):
role: Literal["user", "assistant"]
content: str
class UsageAIChatRequest(BaseModel):
messages: List[ChatMessage] = Field(
..., description="Chat messages (user/assistant history)"
)
model: Optional[str] = Field(default=None, description="Model to use for AI chat")
@router.post(
"/usage/ai/chat",
tags=["Budget & Spend Tracking"],
dependencies=[Depends(user_api_key_auth)],
)
async def usage_ai_chat(
data: UsageAIChatRequest,
request: Request,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
AI chat about usage data. Streams SSE events with the AI response.
The AI agent has access to tools that query aggregated daily activity data.
"""
from litellm.proxy.management_endpoints.common_utils import (
_user_has_admin_view,
)
from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import (
stream_usage_ai_chat,
)
is_admin = _user_has_admin_view(user_api_key_dict)
user_id = user_api_key_dict.user_id
messages = [{"role": m.role, "content": m.content} for m in data.messages]
return StreamingResponse(
stream_usage_ai_chat(
messages=messages,
model=data.model,
user_id=user_id,
is_admin=is_admin,
),
media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
)

11
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -28,6 +28,7 @@ from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
_safe_get_request_headers,
_safe_set_request_parsed_body,
get_form_data,
get_request_body,
@ -60,7 +61,7 @@ def create_request_copy(request: Request):
return {
"method": request.method,
"url": str(request.url),
"headers": dict(request.headers),
"headers": _safe_get_request_headers(request).copy(),
"cookies": request.cookies,
"query_params": dict(request.query_params),
}
@ -329,7 +330,7 @@ async def vllm_proxy_route(
method=request.method,
endpoint=endpoint,
request_query_params=request.query_params,
request_headers=dict(request.headers),
request_headers=_safe_get_request_headers(request),
stream=request_body.get("stream", False),
content=None,
data=None,
@ -1307,7 +1308,7 @@ async def azure_proxy_route(
method=request.method,
endpoint=endpoint,
request_query_params=request.query_params,
request_headers=dict(request.headers),
request_headers=_safe_get_request_headers(request),
stream=request_body.get("stream", False),
content=None,
data=None,
@ -1505,7 +1506,7 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict:
Returns:
dict: Headers dictionary with only allowed headers
"""
incoming_headers = dict(request.headers) or {}
incoming_headers = _safe_get_request_headers(request)
headers = {}
for header_name in ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS:
if header_name in incoming_headers:
@ -1621,7 +1622,7 @@ async def _prepare_vertex_auth_headers(
if (
vertex_credentials is None or vertex_credentials.vertex_project is None
) and router_credentials is None:
headers = dict(request.headers) or {}
headers = _safe_get_request_headers(request).copy()
headers_passed_through = True
verbose_proxy_logger.debug(
"default_vertex_config not set, incoming request headers %s", headers

7
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -50,7 +50,10 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
_safe_get_request_headers,
)
from litellm.proxy.utils import get_server_root_path
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.custom_http import httpxSpecialProvider
@ -649,7 +652,7 @@ async def pass_through_request( # noqa: PLR0915
url = httpx.URL(target)
headers = custom_headers
headers = HttpPassThroughEndpointHelpers.forward_headers_from_request(
request_headers=dict(request.headers),
request_headers=_safe_get_request_headers(request).copy(),
headers=headers,
forward_headers=forward_headers,
)

93
litellm/proxy/proxy_server.py
View file

@ -294,6 +294,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
_safe_get_request_headers,
check_file_size_under_limit,
get_form_data,
)
@ -391,6 +392,7 @@ from litellm.proxy.management_endpoints.organization_endpoints import (
router as organization_router,
)
from litellm.proxy.management_endpoints.policy_endpoints import router as policy_router
from litellm.proxy.management_endpoints.usage_endpoints import router as usage_ai_router
from litellm.proxy.management_endpoints.project_endpoints import (
router as project_router,
)
@ -409,6 +411,9 @@ from litellm.proxy.management_endpoints.team_endpoints import (
update_team,
validate_membership,
)
from litellm.proxy.management_endpoints.tool_management_endpoints import (
router as tool_management_router,
)
from litellm.proxy.management_endpoints.ui_sso import (
get_disabled_non_admin_personal_key_creation,
)
@ -1211,9 +1216,7 @@ try:
# Case 2: Runtime UI exists and is ready
if has_content and is_pre_restructured:
verbose_proxy_logger.info(
f"Using pre-restructured UI at {runtime_ui_path}"
)
verbose_proxy_logger.info(f"Using pre-restructured UI at {runtime_ui_path}")
ui_path = runtime_ui_path
# Case 3: Runtime UI exists but needs restructuring
@ -2989,6 +2992,10 @@ class ProxyConfig:
if master_key is not None and isinstance(master_key, str):
litellm_master_key_hash = hash_token(master_key)
else:
verbose_proxy_logger.critical(
"LITELLM_MASTER_KEY is not set! All requests will be treated as INTERNAL_USER with no admin access. Set LITELLM_MASTER_KEY for production use."
)
### USER API KEY CACHE IN-MEMORY TTL ###
user_api_key_cache_ttl = general_settings.get(
"user_api_key_cache_ttl", None
@ -3791,6 +3798,7 @@ class ProxyConfig:
parsed = value
elif isinstance(value, str):
import json
try:
parsed = yaml.safe_load(value)
except (yaml.YAMLError, json.JSONDecodeError):
@ -4376,10 +4384,12 @@ class ProxyConfig:
if self._should_load_db_object(object_type="model_cost_map"):
await self._check_and_reload_model_cost_map(prisma_client=prisma_client)
if self._should_load_db_object(object_type="anthropic_beta_headers"):
await self._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client)
await self._check_and_reload_anthropic_beta_headers(
prisma_client=prisma_client
)
if self._should_load_db_object(object_type="sso_settings"):
await self._init_sso_settings_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="cache_settings"):
@ -4609,7 +4619,9 @@ class ProxyConfig:
f"Error in _check_and_reload_model_cost_map: {str(e)}"
)
async def _check_and_reload_anthropic_beta_headers(self, prisma_client: PrismaClient):
async def _check_and_reload_anthropic_beta_headers(
self, prisma_client: PrismaClient
):
"""
Check if anthropic beta headers config needs to be reloaded based on database configuration.
This function runs every 10 seconds as part of _init_non_llm_objects_in_db.
@ -4700,7 +4712,11 @@ class ProxyConfig:
)
# Count providers in config
provider_count = sum(1 for k in new_config.keys() if k != "provider_aliases" and k != "description")
provider_count = sum(
1
for k in new_config.keys()
if k != "provider_aliases" and k != "description"
)
verbose_proxy_logger.info(
f"Anthropic beta headers config reloaded successfully. Providers: {provider_count}"
)
@ -5682,8 +5698,7 @@ class ProxyStartupEvent:
):
_db_val = _db_gs_record.param_value.get("store_model_in_db")
if _db_val is True or (
isinstance(_db_val, str)
and _db_val.lower() == "true"
isinstance(_db_val, str) and _db_val.lower() == "true"
):
store_model_in_db = True
verbose_proxy_logger.info(
@ -6150,6 +6165,7 @@ class ProxyStartupEvent:
"Pyroscope profiling will not run. Install with: pip install pyroscope-io"
)
#### API ENDPOINTS ####
@router.get(
"/v1/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"]
@ -10448,7 +10464,7 @@ async def async_queue_request(
data["proxy_server_request"] = {
"url": str(request.url),
"method": request.method,
"headers": dict(request.headers),
"headers": _safe_get_request_headers(request).copy(),
"body": copy.copy(data), # use copy instead of deepcopy
}
@ -10469,7 +10485,7 @@ async def async_queue_request(
data["metadata"] = {}
data["metadata"]["user_api_key"] = user_api_key_dict.api_key
data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata
_headers = dict(request.headers)
_headers = _safe_get_request_headers(request).copy()
_headers.pop(
"authorization", None
) # do not store the original `sk-..` api key in the db
@ -10988,18 +11004,14 @@ async def get_favicon():
from fastapi.responses import Response
current_dir = os.path.dirname(os.path.abspath(__file__))
default_favicon = os.path.join(
current_dir, "_experimental", "out", "favicon.ico"
)
default_favicon = os.path.join(current_dir, "_experimental", "out", "favicon.ico")
favicon_url = os.getenv("LITELLM_FAVICON_URL", "")
if not favicon_url:
if os.path.exists(default_favicon):
return FileResponse(default_favicon, media_type="image/x-icon")
raise HTTPException(
status_code=404, detail="Default favicon not found"
)
raise HTTPException(status_code=404, detail="Default favicon not found")
if favicon_url.startswith(("http://", "https://")):
try:
@ -11014,9 +11026,7 @@ async def get_favicon():
)
response = await async_client.get(favicon_url)
if response.status_code == 200:
content_type = response.headers.get(
"content-type", "image/x-icon"
)
content_type = response.headers.get("content-type", "image/x-icon")
return Response(
content=response.content,
media_type=content_type,
@ -11028,12 +11038,8 @@ async def get_favicon():
response.status_code,
)
if os.path.exists(default_favicon):
return FileResponse(
default_favicon, media_type="image/x-icon"
)
raise HTTPException(
status_code=404, detail="Favicon not found"
)
return FileResponse(default_favicon, media_type="image/x-icon")
raise HTTPException(status_code=404, detail="Favicon not found")
except HTTPException:
raise
except Exception as e:
@ -11041,20 +11047,14 @@ async def get_favicon():
"Error downloading favicon from %s: %s", favicon_url, e
)
if os.path.exists(default_favicon):
return FileResponse(
default_favicon, media_type="image/x-icon"
)
raise HTTPException(
status_code=404, detail="Favicon not found"
)
return FileResponse(default_favicon, media_type="image/x-icon")
raise HTTPException(status_code=404, detail="Favicon not found")
else:
if os.path.exists(favicon_url):
return FileResponse(favicon_url, media_type="image/x-icon")
if os.path.exists(default_favicon):
return FileResponse(default_favicon, media_type="image/x-icon")
raise HTTPException(
status_code=404, detail="Favicon not found"
)
raise HTTPException(status_code=404, detail="Favicon not found")
#### INVITATION MANAGEMENT ####
@ -12540,7 +12540,9 @@ async def reload_anthropic_beta_headers(
},
)
provider_count = sum(1 for k in new_config.keys() if k not in ["provider_aliases", "description"])
provider_count = sum(
1 for k in new_config.keys() if k not in ["provider_aliases", "description"]
)
verbose_proxy_logger.info(
f"Anthropic beta headers config reloaded successfully in current pod. Providers: {provider_count}"
)
@ -12552,7 +12554,9 @@ async def reload_anthropic_beta_headers(
"timestamp": current_time.isoformat(),
}
except Exception as e:
verbose_proxy_logger.exception(f"Failed to reload anthropic beta headers: {str(e)}")
verbose_proxy_logger.exception(
f"Failed to reload anthropic beta headers: {str(e)}"
)
raise HTTPException(
status_code=500, detail=f"Failed to reload anthropic beta headers: {str(e)}"
)
@ -12674,7 +12678,8 @@ async def cancel_anthropic_beta_headers_reload(
f"Failed to cancel anthropic beta headers reload: {str(e)}"
)
raise HTTPException(
status_code=500, detail=f"Failed to cancel anthropic beta headers reload: {str(e)}"
status_code=500,
detail=f"Failed to cancel anthropic beta headers reload: {str(e)}",
)
@ -12721,7 +12726,9 @@ async def get_anthropic_beta_headers_reload_status(
)
if config_record is None or config_record.param_value is None:
verbose_proxy_logger.info("No anthropic beta headers reload configuration found")
verbose_proxy_logger.info(
"No anthropic beta headers reload configuration found"
)
return {
"scheduled": False,
"interval_hours": None,
@ -12747,7 +12754,9 @@ async def get_anthropic_beta_headers_reload_status(
# Use pod's in-memory last reload time
if last_anthropic_beta_headers_reload is not None:
try:
last_reload_time = datetime.fromisoformat(last_anthropic_beta_headers_reload)
last_reload_time = datetime.fromisoformat(
last_anthropic_beta_headers_reload
)
time_since_last_reload = current_time - last_reload_time
hours_since_last_reload = time_since_last_reload.total_seconds() / 3600
@ -12868,6 +12877,7 @@ app.include_router(caching_router)
app.include_router(analytics_router)
app.include_router(guardrails_router)
app.include_router(policy_router)
app.include_router(usage_ai_router)
app.include_router(policy_crud_router)
app.include_router(policy_resolve_router)
app.include_router(search_tool_management_router)
@ -12881,6 +12891,7 @@ app.include_router(budget_management_router)
app.include_router(model_management_router)
app.include_router(model_access_group_management_router)
app.include_router(tag_management_router)
app.include_router(tool_management_router)
app.include_router(cost_tracking_settings_router)
app.include_router(router_settings_router)
app.include_router(fallback_management_router)

24
litellm/proxy/schema.prisma
View file

@ -64,6 +64,8 @@ model LiteLLM_AgentsTable {
litellm_params Json?
agent_card_params Json
agent_access_groups String[] @default([])
object_permission_id String?
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
created_at DateTime @default(now()) @map("created_at")
created_by String
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
@ -264,6 +266,7 @@ model LiteLLM_ObjectPermissionTable {
organizations LiteLLM_OrganizationTable[]
users LiteLLM_UserTable[]
end_users LiteLLM_EndUserTable[]
agents_table LiteLLM_AgentsTable[]
}
// Holds the MCP server configuration
@ -314,6 +317,7 @@ model LiteLLM_VerificationToken {
router_settings Json? @default("{}")
user_id String?
team_id String?
agent_id String?
project_id String?
permissions Json @default("{}")
max_parallel_requests Int?
@ -1051,6 +1055,26 @@ model LiteLLM_PolicyAttachmentTable {
updated_by String?
}
// Global tool registry - auto-discovered from LLM responses; admins set call_policy here
model LiteLLM_ToolTable {
tool_id String @id @default(uuid())
tool_name String @unique // e.g. "huggingface_remote-mcp__dynamic_space"
origin String? // MCP server name or "user_defined"
call_policy String @default("untrusted") // "trusted" | "untrusted" | "dual_llm" | "blocked"
call_count Int @default(0) // cumulative number of times this tool was seen
assignments Json? @default("{}")
key_hash String? // hash of the virtual key that first called this tool
team_id String? // team that first called this tool
key_alias String? // human-readable alias of the virtual key
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt
updated_by String?
@@index([call_policy])
@@index([team_id])
}
//Unified Access Groups table for storing unified access groups
model LiteLLM_AccessGroupTable {
access_group_id String @id @default(uuid())

3
litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py
View file

@ -19,6 +19,7 @@ from fastapi import APIRouter, Request, Response
import litellm
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
from litellm.proxy.litellm_pre_call_utils import _get_dynamic_logging_metadata
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
create_pass_through_route,
@ -32,7 +33,7 @@ def create_request_copy(request: Request):
return {
"method": request.method,
"url": str(request.url),
"headers": dict(request.headers),
"headers": _safe_get_request_headers(request).copy(),
"cookies": request.cookies,
"query_params": dict(request.query_params),
}

12
litellm/router.py
View file

@ -6423,6 +6423,7 @@ class Router:
self.model_list = []
self.model_id_to_deployment_index_map = {} # Reset the index
self.model_name_to_deployment_indices = {} # Reset the model_name index
self._invalidate_model_group_info_cache()
self._invalidate_access_groups_cache()
# we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works
@ -6733,6 +6734,7 @@ class Router:
"""
idx = len(self.model_list)
self.model_list.append(model)
self._invalidate_model_group_info_cache()
self._invalidate_access_groups_cache()
# Update model_id index for O(1) lookup
@ -6781,6 +6783,7 @@ class Router:
if removal_idx is not None:
self.model_list.pop(removal_idx)
self._invalidate_model_group_info_cache()
self._invalidate_access_groups_cache()
self._update_deployment_indices_after_removal(
model_id=deployment_id, removal_idx=removal_idx
@ -6815,6 +6818,7 @@ class Router:
if deployment_idx is not None:
# Pop the item from the list first
item = self.model_list.pop(deployment_idx)
self._invalidate_model_group_info_cache()
self._invalidate_access_groups_cache()
self._update_deployment_indices_after_removal(
model_id=id, removal_idx=deployment_idx
@ -7576,6 +7580,7 @@ class Router:
"""
# First populate the model_list
self.model_list = []
self._invalidate_model_group_info_cache()
self._invalidate_access_groups_cache()
for _, model in enumerate(model_list):
# Extract model_info from the model dict
@ -7920,6 +7925,13 @@ class Router:
return returned_models
def _invalidate_model_group_info_cache(self) -> None:
"""Invalidate the cached model group info.
Call this whenever self.model_list is modified to ensure the cache is rebuilt.
"""
self._cached_get_model_group_info.cache_clear()
def _invalidate_access_groups_cache(self) -> None:
"""Invalidate the cached access groups.

1
litellm/types/llms/openai.py
View file

@ -1125,6 +1125,7 @@ class ResponsesAPIOptionalRequestParams(TypedDict, total=False):
prompt: Optional[PromptObject]
max_tool_calls: Optional[int]
prompt_cache_key: Optional[str]
prompt_cache_retention: Optional[str]
stream_options: Optional[dict]
top_logprobs: Optional[int]
partial_images: Optional[

2
litellm/types/responses/main.py
View file

@ -6,7 +6,7 @@ from typing_extensions import Any, List, Optional, TypedDict
from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject
Phase = Optional[Literal["commentary", "final_answer"]] # TODO: Once openai sdk has updated, we can remove this and use the openai sdk type
Phase = Optional[Literal["commentary", "final_answer"]]
class GenericResponseOutputItemContentAnnotation(BaseLiteLLMOpenAIResponseObject):
"""Annotation for content in a message"""

42
litellm/types/tool_management.py Normal file
View file

@ -0,0 +1,42 @@
"""
Pydantic models for Tool Policy management endpoints.
"""
from datetime import datetime
from typing import Dict, List, Literal, Optional
from pydantic import BaseModel
ToolCallPolicy = Literal["trusted", "untrusted", "dual_llm", "blocked"]
class LiteLLM_ToolTableRow(BaseModel):
tool_id: str
tool_name: str
origin: Optional[str] = None
call_policy: ToolCallPolicy = "untrusted"
call_count: int = 0
assignments: Optional[Dict] = None
key_hash: Optional[str] = None
team_id: Optional[str] = None
key_alias: Optional[str] = None
created_at: Optional[datetime] = None
updated_at: Optional[datetime] = None
created_by: Optional[str] = None
updated_by: Optional[str] = None
class ToolListResponse(BaseModel):
tools: List[LiteLLM_ToolTableRow]
total: int
class ToolPolicyUpdateRequest(BaseModel):
tool_name: str
call_policy: ToolCallPolicy
class ToolPolicyUpdateResponse(BaseModel):
tool_name: str
call_policy: ToolCallPolicy
updated: bool

20
litellm/types/utils.py
View file

@ -1,7 +1,17 @@
import json
import time
from enum import Enum
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Union
from typing import (
TYPE_CHECKING,
Any,
Dict,
List,
Literal,
Mapping,
Optional,
Union,
get_args,
)
from openai._models import BaseModel as OpenAIObject
from openai.types.audio.transcription_create_params import (
@ -1826,7 +1836,7 @@ class ModelResponse(ModelResponseBase):
else:
usage = usage
elif stream is None or stream is False:
usage = Usage()
usage = None # avoid constructing throwaway Usage; set by convert_to_model_response_object
if hidden_params:
self._hidden_params = hidden_params
@ -3186,6 +3196,12 @@ OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: set[str] = {
LlmProviders.HOSTED_VLLM.value,
}
ListBatchesSupportedProvider = Literal["openai", "azure", "hosted_vllm", "vertex_ai"]
LIST_BATCHES_SUPPORTED_PROVIDERS: frozenset[str] = frozenset(
get_args(ListBatchesSupportedProvider)
)
class SearchProviders(str, Enum):
"""

5
litellm/utils.py
View file

@ -4644,11 +4644,12 @@ def add_provider_specific_params_to_optional_params(
)
is False
):
extra_body = passed_params.pop("extra_body", {})
extra_body = passed_params.pop("extra_body", None) or {}
for k in passed_params.keys():
if k not in openai_params and passed_params[k] is not None:
extra_body[k] = passed_params[k]
optional_params.setdefault("extra_body", {})
if not isinstance(optional_params.get("extra_body"), dict):
optional_params["extra_body"] = {}
initial_extra_body = {
**optional_params["extra_body"],
**extra_body,

2
requirements.txt
View file

@ -10,7 +10,7 @@ cryptography==46.0.5 #GHSA-r6ph-v2qm-q3c2
anyio==4.8.0 # openai + http req.
httpx==0.28.1
openai==2.9.0 # openai req.
openai==2.24.0 # openai req.
fastapi==0.120.1 # server dep
starlette==0.49.1 # starlette fastapi dep
backoff==2.2.1 # server dep

94
tests/litellm/proxy/guardrails/test_custom_code_security.py Normal file
View file

@ -0,0 +1,94 @@
import pytest
from litellm.proxy.guardrails.guardrail_hooks.custom_code.code_validator import (
validate_custom_code,
CustomCodeValidationError,
)
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
CustomCodeGuardrail,
)
# Phase 4.1: Test forbidden pattern validation
def test_validate_custom_code_import_os():
code = "import os\ndef apply_guardrail(inputs, req, ty):\n return allow()"
with pytest.raises(CustomCodeValidationError, match="import statements are not"):
validate_custom_code(code)
def test_validate_custom_code_from_subprocess():
code = (
"from subprocess import call\ndef apply_guardrail(i, r, t):\n return allow()"
)
with pytest.raises(
CustomCodeValidationError, match="import statements are not allowed"
):
validate_custom_code(code)
def test_validate_custom_code_exec():
code = "def apply_guardrail(i, r, t):\n exec('print(1)')\n return allow()"
with pytest.raises(CustomCodeValidationError, match=r"exec\(\) is not allowed"):
validate_custom_code(code)
def test_validate_custom_code_builtins():
code = "def apply_guardrail(i, r, t):\n print(__builtins__)\n return allow()"
with pytest.raises(
CustomCodeValidationError, match="__builtins__ access is not allowed"
):
validate_custom_code(code)
def test_validate_custom_code_subclasses():
code = "def apply_guardrail(i, r, t):\n print(''.__class__.__mro__[1].__subclasses__())\n return allow()"
with pytest.raises(
CustomCodeValidationError, match="__subclasses__ access is not allowed"
):
validate_custom_code(code)
def test_validate_custom_code_clean():
code = (
"def apply_guardrail(inputs, request_data, input_type):\n return allow()\n"
)
# Should not raise any exception
validate_custom_code(code)
# Phase 4.2: Test __builtins__ restriction in execution
def test_custom_code_compile_valid():
code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()"
guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="test")
# if it doesn't fail, we successfully compiled
assert guardrail._compiled_function is not None
def test_custom_code_override_builtins():
# Verify that even if pattern validation is bypassed, __builtins__ = {} blocks dangerous builtins.
# We test this by compiling safe code and verifying builtins are not accessible in the sandbox.
code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()"
guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="test")
# The compiled function's globals should have empty __builtins__
fn_globals = guardrail._compiled_function.__globals__
assert fn_globals.get("__builtins__") == {}
@pytest.mark.asyncio
async def test_custom_code_guardrail_apply():
code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()"
guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="test")
from litellm.types.utils import GenericGuardrailAPIInputs
result = await guardrail.apply_guardrail(
inputs=GenericGuardrailAPIInputs(texts=["test"]),
request_data={},
input_type="request",
)
assert result["texts"][0] == "test"
# The RBAC endpoint tests are harder to write right here, but the core security
# validations are fully covered by the simple tests above.

44
tests/litellm_utils_tests/test_health_check.py
View file

@ -473,6 +473,50 @@ def test_update_litellm_params_for_health_check():
)
@pytest.mark.asyncio
async def test_perform_health_check_filters_by_model_id():
"""
When model_id is passed, only that deployment is checked (not all deployments
that share the same model name).
"""
from litellm.proxy.health_check import perform_health_check
# Two deployments with same model_name but different ids
model_list = [
{
"model_name": "gpt-4",
"model_info": {"id": "deployment-id-1"},
"litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"},
},
{
"model_name": "gpt-4",
"model_info": {"id": "deployment-id-2"},
"litellm_params": {"model": "gpt-4", "api_key": "fake-key-2"},
},
]
captured_list = []
async def mock_perform_health_check(m_list, details=True):
captured_list.append(m_list)
return [{"model": "gpt-4", "api_key": m_list[0]["litellm_params"]["api_key"]}], []
with patch(
"litellm.proxy.health_check._perform_health_check",
side_effect=mock_perform_health_check,
):
healthy_endpoints, unhealthy_endpoints = await perform_health_check(
model_list=model_list, model_id="deployment-id-2", details=True
)
# Only one deployment (deployment-id-2) should have been passed to _perform_health_check
assert len(captured_list) == 1
assert len(captured_list[0]) == 1
assert (captured_list[0][0].get("model_info") or {}).get("id") == "deployment-id-2"
assert len(healthy_endpoints) == 1
assert healthy_endpoints[0]["api_key"] == "fake-key-2"
@pytest.mark.asyncio
async def test_perform_health_check_with_health_check_model():
"""

44
tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py
View file

@ -1453,3 +1453,47 @@ def test_convert_to_model_response_object_falsy_id_preserves_auto_generated(fals
)
assert result.id == original_id
assert result.id.startswith("chatcmpl-")
def test_convert_to_model_response_object_default_usage_overwritten():
"""
Regression test: convert_to_model_response_object must properly set Usage
on a ModelResponse that only has the default Usage from ModelResponse.__init__()
(i.e. no extra litellm.Usage() set via setattr beforehand).
This validates the optimization of removing the redundant
`setattr(model_response, "usage", litellm.Usage())` in completion().
"""
mr = ModelResponse()
# usage is not set by default (optimization: avoid constructing throwaway Usage)
assert not hasattr(mr, "usage")
response_object = {
"id": "chatcmpl-usage-test",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "Hello"},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 15,
"completion_tokens": 7,
"total_tokens": 22,
},
"model": "gpt-4o",
}
result = convert_to_model_response_object(
model_response_object=mr,
response_object=response_object,
stream=False,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert isinstance(result, ModelResponse)
assert result.usage.prompt_tokens == 15
assert result.usage.completion_tokens == 7
assert result.usage.total_tokens == 22

75
tests/proxy_unit_tests/test_proxy_server.py
View file

@ -334,6 +334,81 @@ def test_chat_completion_forward_headers(
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
@pytest.mark.parametrize("forward_llm_auth_headers", [True, False])
@mock_patch_acompletion()
def test_chat_completion_forward_llm_provider_auth_headers(
mock_acompletion, client_no_auth, forward_llm_auth_headers
):
"""
Test that LLM provider auth headers (x-api-key, x-goog-api-key) are forwarded
when forward_llm_provider_auth_headers=True.
This allows clients to send their own LLM provider API keys through the proxy.
"""
try:
# Configure general settings
gs = getattr(litellm.proxy.proxy_server, "general_settings")
gs["forward_client_headers_to_llm_api"] = True
gs["forward_llm_provider_auth_headers"] = forward_llm_auth_headers
setattr(litellm.proxy.proxy_server, "general_settings", gs)
# Test data
test_data = {
"model": "gpt-3.5-turbo",
"messages": [
{"role": "user", "content": "hello"},
],
"max_tokens": 10,
}
# Headers including LLM provider auth
request_headers = {
"Authorization": "Bearer sk-proxy-auth-123", # Proxy auth (should be stripped)
"x-api-key": "sk-ant-api03-test-anthropic-key", # Anthropic API key
"x-goog-api-key": "google-api-key-123", # Google API key
"X-Custom-Header": "custom-value", # Custom header (should be forwarded)
}
# Make request
response = client_no_auth.post(
"/v1/chat/completions", json=test_data, headers=request_headers
)
assert response.status_code == 200
# Check forwarded headers
forwarded_headers = mock_acompletion.call_args.kwargs.get("headers", {})
if forward_llm_auth_headers:
# LLM provider auth headers should be forwarded
assert "x-api-key" in forwarded_headers
assert forwarded_headers["x-api-key"] == "sk-ant-api03-test-anthropic-key"
assert "x-goog-api-key" in forwarded_headers
assert forwarded_headers["x-goog-api-key"] == "google-api-key-123"
else:
# LLM provider auth headers should be stripped
assert "x-api-key" not in forwarded_headers
assert "x-goog-api-key" not in forwarded_headers
# Custom headers should always be forwarded (when forward_client_headers_to_llm_api=True)
assert "x-custom-header" in forwarded_headers
assert forwarded_headers["x-custom-header"] == "custom-value"
# Proxy Authorization should never be forwarded
assert "authorization" not in forwarded_headers
print(f"✓ Test passed with forward_llm_provider_auth_headers={forward_llm_auth_headers}")
print(f" Forwarded headers: {list(forwarded_headers.keys())}")
except Exception as e:
pytest.fail(f"Test failed with forward_llm_auth_headers={forward_llm_auth_headers}: {str(e)}")
finally:
# Clean up
gs = getattr(litellm.proxy.proxy_server, "general_settings")
gs.pop("forward_llm_provider_auth_headers", None)
setattr(litellm.proxy.proxy_server, "general_settings", gs)
@mock_patch_acompletion()
@pytest.mark.asyncio
async def test_team_disable_guardrails(mock_acompletion, client_no_auth):

3
tests/test_litellm/interactions/test_openapi_compliance.py
View file

@ -147,7 +147,8 @@ class TestResponseCompliance:
"""Verify status enum values match spec."""
schema = spec_dict["components"]["schemas"]["CreateModelInteractionParams"]
status_prop = schema["properties"]["status"]
expected_statuses = ["UNSPECIFIED", "IN_PROGRESS", "REQUIRES_ACTION", "COMPLETED", "FAILED", "CANCELLED", "INCOMPLETE"]
# Google Interactions API uses lowercase status values (updated Feb 2026)
expected_statuses = ["in_progress", "requires_action", "completed", "failed", "cancelled", "incomplete"]
assert status_prop["enum"] == expected_statuses
print(f"✓ Status enum values: {expected_statuses}")

81
tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py
View file

@ -363,6 +363,87 @@ class TestProxyOAuthHeaderForwarding:
assert "authorization" not in cleaned
assert cleaned["content-type"] == "application/json"
def test_clean_headers_forwards_anthropic_api_key_when_enabled(self):
"""clean_headers should preserve x-api-key when forward_llm_provider_auth_headers=True."""
from starlette.datastructures import Headers
from litellm.proxy.litellm_pre_call_utils import clean_headers
raw_headers = Headers(
raw=[
(b"authorization", b"Bearer sk-proxy-auth"),
(b"x-api-key", b"sk-ant-api03-test-key"),
(b"content-type", b"application/json"),
]
)
cleaned = clean_headers(raw_headers, forward_llm_provider_auth_headers=True)
# x-api-key should be preserved when flag is True
assert "x-api-key" in cleaned
assert cleaned["x-api-key"] == "sk-ant-api03-test-key"
# Authorization (proxy auth) should still be stripped
assert "authorization" not in cleaned
assert cleaned["content-type"] == "application/json"
def test_clean_headers_strips_anthropic_api_key_when_disabled(self):
"""clean_headers should strip x-api-key when forward_llm_provider_auth_headers=False (default)."""
from starlette.datastructures import Headers
from litellm.proxy.litellm_pre_call_utils import clean_headers
raw_headers = Headers(
raw=[
(b"x-api-key", b"sk-ant-api03-test-key"),
(b"content-type", b"application/json"),
]
)
cleaned = clean_headers(raw_headers, forward_llm_provider_auth_headers=False)
# x-api-key should be stripped by default
assert "x-api-key" not in cleaned
assert cleaned["content-type"] == "application/json"
def test_clean_headers_forwards_google_api_key_when_enabled(self):
"""clean_headers should preserve x-goog-api-key when forward_llm_provider_auth_headers=True."""
from starlette.datastructures import Headers
from litellm.proxy.litellm_pre_call_utils import clean_headers
raw_headers = Headers(
raw=[
(b"x-goog-api-key", b"google-api-key-123"),
(b"content-type", b"application/json"),
]
)
cleaned = clean_headers(raw_headers, forward_llm_provider_auth_headers=True)
assert "x-goog-api-key" in cleaned
assert cleaned["x-goog-api-key"] == "google-api-key-123"
assert cleaned["content-type"] == "application/json"
def test_clean_headers_preserves_oauth_regardless_of_forward_flag(self):
"""clean_headers should always preserve OAuth tokens regardless of forward_llm_provider_auth_headers."""
from starlette.datastructures import Headers
from litellm.proxy.litellm_pre_call_utils import clean_headers
raw_headers = Headers(
raw=[
(b"authorization", f"Bearer {FAKE_OAUTH_TOKEN}".encode()),
(b"content-type", b"application/json"),
]
)
# Should preserve OAuth even with flag=False
cleaned_without_flag = clean_headers(raw_headers, forward_llm_provider_auth_headers=False)
assert "authorization" in cleaned_without_flag
assert cleaned_without_flag["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}"
# Should also preserve OAuth with flag=True
cleaned_with_flag = clean_headers(raw_headers, forward_llm_provider_auth_headers=True)
assert "authorization" in cleaned_with_flag
assert cleaned_with_flag["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}"
def test_add_provider_specific_headers_forwards_oauth(self):
"""add_provider_specific_headers_to_request should forward OAuth Authorization
as a ProviderSpecificHeader scoped to Anthropic-compatible providers."""

103
tests/test_litellm/llms/hosted_vllm/responses/test_hosted_vllm_responses.py Normal file
View file

@ -0,0 +1,103 @@
"""
Tests for hosted_vllm responses API support.
Regression test for: https://github.com/BerriAI/litellm/issues
Bug: client.responses.create() raised TypeError: 'NoneType' object is not a mapping
when extra_body=None was passed through the responses→completion pipeline for
hosted_vllm (and any OpenAI-compatible provider using add_provider_specific_params_to_optional_params).
"""
import json
import os
import sys
from unittest.mock import MagicMock, patch
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
import litellm
def _make_mock_chat_completion_response(content: str = "Hello! I'm doing well.") -> dict:
return {
"id": "chatcmpl-test123",
"object": "chat.completion",
"created": 1234567890,
"model": "Qwen/Qwen3-8B",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": content},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
}
def _make_mock_http_client(response_body: dict) -> MagicMock:
mock_client = MagicMock()
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.json.return_value = response_body
mock_response.text = json.dumps(response_body)
mock_client.post.return_value = mock_response
return mock_client
def test_hosted_vllm_responses_create_with_string_input():
"""
Regression test: responses.create() with string input must not raise
TypeError: 'NoneType' object is not a mapping.
Root cause: extra_body=None was passed explicitly through the
responses→completion pipeline. In add_provider_specific_params_to_optional_params(),
passed_params.pop("extra_body", {}) returned None (key existed with value None),
and **None raised TypeError at dict unpacking.
Fix: normalize None to {} for both extra_body and optional_params["extra_body"].
"""
mock_client = _make_mock_http_client(
_make_mock_chat_completion_response("I'm doing well, thanks!")
)
with patch(
"litellm.llms.custom_httpx.llm_http_handler._get_httpx_client",
return_value=mock_client,
):
response = litellm.responses(
model="hosted_vllm/Qwen/Qwen3-8B",
input="Hello, how are you?",
api_base="https://test-vllm.example.com/v1",
api_key="test-key",
)
from litellm.types.llms.openai import ResponsesAPIResponse
assert response is not None
assert isinstance(response, ResponsesAPIResponse)
assert len(response.output) > 0
output_message = response.output[0]
assert output_message.role == "assistant" # type: ignore[union-attr]
assert len(output_message.content) > 0 # type: ignore[union-attr]
assert "well" in output_message.content[0].text # type: ignore[union-attr]
def test_hosted_vllm_responses_create_with_explicit_none_extra_body():
"""
Directly verify the fix in add_provider_specific_params_to_optional_params:
extra_body=None must not crash when building optional_params.
"""
from litellm.utils import get_optional_params
# This should not raise TypeError: 'NoneType' object is not a mapping
optional_params = get_optional_params(
model="Qwen/Qwen3-8B",
custom_llm_provider="hosted_vllm",
extra_body=None,
)
# extra_body=None should be normalized to an empty dict (or absent)
assert optional_params.get("extra_body") is not None or "extra_body" not in optional_params

35
tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py
View file

@ -207,10 +207,10 @@ class TestOpenAIChatCompletionStreamingHandler:
def test_chunk_parser_maps_reasoning_to_reasoning_content(self):
"""
Test that chunk_parser maps 'reasoning' field to 'reasoning_content'.
Some OpenAI-compatible providers (e.g., GLM-5, hosted_vllm) return
delta.reasoning, but LiteLLM expects delta.reasoning_content.
Regression test for: Streaming responses with delta.reasoning field
coming back empty when using openai/ or hosted_vllm/ providers.
"""
@ -293,3 +293,34 @@ class TestPromptCacheKeyIntegration:
prompt_cache_key="test-cache-key-123",
)
assert optional_params.get("prompt_cache_key") == "test-cache-key-123"
class TestPromptCacheParams:
"""Tests for prompt_cache_key and prompt_cache_retention support."""
def setup_method(self):
self.config = OpenAIGPTConfig()
def test_prompt_cache_key_in_supported_params(self):
"""Test that prompt_cache_key is in supported params for OpenAI models."""
supported_params = self.config.get_supported_openai_params("gpt-4o")
assert "prompt_cache_key" in supported_params
def test_prompt_cache_retention_in_supported_params(self):
"""Test that prompt_cache_retention is in supported params for OpenAI models."""
supported_params = self.config.get_supported_openai_params("gpt-4o")
assert "prompt_cache_retention" in supported_params
def test_prompt_cache_params_passed_through(self):
"""Test that prompt_cache_key and prompt_cache_retention are passed through by map_openai_params."""
optional_params = self.config.map_openai_params(
non_default_params={
"prompt_cache_key": "my-cache-key",
"prompt_cache_retention": "24h",
},
optional_params={},
model="gpt-4o",
drop_params=False,
)
assert optional_params.get("prompt_cache_key") == "my-cache-key"
assert optional_params.get("prompt_cache_retention") == "24h"

67
tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py
View file

@ -761,3 +761,70 @@ async def test_request_body_with_html_script_tags():
f"Message content with HTML was modified during parsing: "
f"expected={msg['content']!r}, got={result['messages'][2]['content']!r}"
)
def test_safe_get_request_headers_caches_on_request_state():
"""
Test that _safe_get_request_headers caches the result on request.state
and returns the same object on subsequent calls.
"""
mock_request = MagicMock()
mock_request.headers = {"content-type": "application/json", "authorization": "Bearer sk-123"}
mock_request.state = MagicMock(spec=[]) # empty spec so getattr returns default
# First call — should create and cache
result1 = _safe_get_request_headers(mock_request)
assert result1 == {"content-type": "application/json", "authorization": "Bearer sk-123"}
assert mock_request.state._cached_headers is result1
# Second call — should return the cached object (same identity)
result2 = _safe_get_request_headers(mock_request)
assert result2 is result1
def test_safe_get_request_headers_none_request():
"""
Test that _safe_get_request_headers returns empty dict for None request.
"""
result = _safe_get_request_headers(None)
assert result == {}
def test_safe_get_request_headers_copy_protects_cache():
"""
Test that callers using .copy() before mutation do not corrupt the cache.
"""
mock_request = MagicMock()
mock_request.headers = {"authorization": "Bearer sk-123", "host": "localhost"}
mock_request.state = MagicMock(spec=[])
original = _safe_get_request_headers(mock_request)
# Simulate what mutation call sites do: copy then pop
mutable = _safe_get_request_headers(mock_request).copy()
mutable.pop("authorization", None)
# Cache must be unaffected
assert "authorization" in _safe_get_request_headers(mock_request)
assert _safe_get_request_headers(mock_request) is original
def test_safe_get_request_headers_state_unavailable():
"""
Test that _safe_get_request_headers still returns headers when
request.state rejects attribute writes (the except path on the cache-write).
"""
class ReadOnlyState:
"""State object that allows reads but raises on writes."""
def __setattr__(self, name, value):
raise AttributeError("read-only state")
def __getattr__(self, name):
return None # _cached_headers not found → triggers fresh read
mock_request = MagicMock()
mock_request.headers = {"content-type": "application/json"}
mock_request.state = ReadOnlyState()
result = _safe_get_request_headers(mock_request)
assert result == {"content-type": "application/json"}

75
tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py Normal file
View file

@ -0,0 +1,75 @@
"""
Unit tests for ToolDiscoveryQueue.
"""
import os
import sys
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import (
ToolDiscoveryQueue,
)
@pytest.fixture
def queue():
return ToolDiscoveryQueue()
def test_add_single_tool(queue):
queue.add_update({"tool_name": "my_tool", "origin": "user_defined"})
items = queue.flush()
assert len(items) == 1
assert items[0]["tool_name"] == "my_tool"
assert items[0]["origin"] == "user_defined"
def test_deduplication_same_name(queue):
"""Adding the same tool_name twice should only keep the first."""
queue.add_update({"tool_name": "tool_a", "origin": "mcp_server"})
queue.add_update({"tool_name": "tool_a", "origin": "user_defined"})
items = queue.flush()
assert len(items) == 1
assert items[0]["origin"] == "mcp_server" # first wins
def test_deduplication_different_names(queue):
queue.add_update({"tool_name": "tool_a"})
queue.add_update({"tool_name": "tool_b"})
items = queue.flush()
assert len(items) == 2
names = {i["tool_name"] for i in items}
assert names == {"tool_a", "tool_b"}
def test_flush_clears_pending(queue):
queue.add_update({"tool_name": "tool_x"})
items1 = queue.flush()
assert len(items1) == 1
items2 = queue.flush()
assert len(items2) == 0
def test_seen_names_reset_after_flush(queue):
"""Seen-set is cleared on flush so the same tool can re-enter the next cycle."""
queue.add_update({"tool_name": "tool_a"})
queue.flush()
queue.add_update({"tool_name": "tool_a"}) # same tool, new cycle
items = queue.flush()
assert len(items) == 1
assert items[0]["tool_name"] == "tool_a"
def test_empty_tool_name_ignored(queue):
queue.add_update({"tool_name": ""})
queue.add_update({"tool_name": None}) # type: ignore[arg-type]
items = queue.flush()
assert len(items) == 0
def test_flush_returns_list(queue):
result = queue.flush()
assert isinstance(result, list)

197
tests/test_litellm/proxy/db/test_tool_registry_writer.py Normal file
View file

@ -0,0 +1,197 @@
"""
Unit tests for tool_registry_writer.py — uses a mock prisma client
that exposes execute_raw / query_raw (matching the actual raw-SQL implementation).
"""
import os
import sys
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
from litellm.proxy.db.tool_registry_writer import (
batch_upsert_tools,
get_tool,
get_tools_by_names,
list_tools,
update_tool_policy,
)
def _make_prisma(query_rows=None):
"""Return a minimal mock prisma_client with execute_raw / query_raw."""
default_row = {
"tool_id": "uuid-1",
"tool_name": "my_tool",
"origin": "user_defined",
"call_policy": "untrusted",
"call_count": 1,
"assignments": {},
"key_hash": None,
"team_id": None,
"key_alias": None,
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
"created_by": None,
"updated_by": None,
}
rows = query_rows if query_rows is not None else [default_row]
prisma = MagicMock()
prisma.db.execute_raw = AsyncMock(return_value=None)
prisma.db.query_raw = AsyncMock(return_value=rows)
return prisma
@pytest.mark.asyncio
async def test_batch_upsert_tools_calls_execute_raw():
prisma = _make_prisma()
items = [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}]
await batch_upsert_tools(prisma, items)
prisma.db.execute_raw.assert_awaited_once()
call_args = prisma.db.execute_raw.call_args
sql = call_args.args[0]
assert "LiteLLM_ToolTable" in sql
assert "ON CONFLICT" in sql
@pytest.mark.asyncio
async def test_batch_upsert_tools_empty_list():
prisma = _make_prisma()
await batch_upsert_tools(prisma, [])
prisma.db.execute_raw.assert_not_awaited()
@pytest.mark.asyncio
async def test_batch_upsert_tools_skips_empty_names():
prisma = _make_prisma()
items = [{"tool_name": "", "origin": None}, {"tool_name": None}] # type: ignore[list-item]
await batch_upsert_tools(prisma, items)
prisma.db.execute_raw.assert_not_awaited()
@pytest.mark.asyncio
async def test_batch_upsert_multiple_tools_calls_execute_raw_per_tool():
prisma = _make_prisma()
items = [
{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None},
{"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice"},
]
await batch_upsert_tools(prisma, items)
assert prisma.db.execute_raw.await_count == 2
@pytest.mark.asyncio
async def test_list_tools_no_filter():
row = {
"tool_id": "id1",
"tool_name": "tool_a",
"origin": "mcp",
"call_policy": "untrusted",
"call_count": 5,
"assignments": {},
"key_hash": None,
"team_id": None,
"key_alias": None,
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
"created_by": None,
"updated_by": None,
}
prisma = _make_prisma(query_rows=[row])
result = await list_tools(prisma)
assert len(result) == 1
assert result[0].tool_name == "tool_a"
assert result[0].call_count == 5
prisma.db.query_raw.assert_awaited_once()
@pytest.mark.asyncio
async def test_list_tools_with_policy_filter():
row = {
"tool_id": "id1",
"tool_name": "blocked_tool",
"origin": None,
"call_policy": "blocked",
"call_count": 2,
"assignments": None,
"key_hash": None,
"team_id": None,
"key_alias": None,
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
"created_by": None,
"updated_by": None,
}
prisma = _make_prisma(query_rows=[row])
result = await list_tools(prisma, call_policy="blocked")
assert result[0].call_policy == "blocked"
call_args = prisma.db.query_raw.call_args
sql = call_args.args[0]
assert "WHERE call_policy" in sql
@pytest.mark.asyncio
async def test_get_tool_found():
prisma = _make_prisma()
result = await get_tool(prisma, "my_tool")
assert result is not None
assert result.tool_name == "my_tool"
prisma.db.query_raw.assert_awaited_once()
@pytest.mark.asyncio
async def test_get_tool_not_found():
prisma = _make_prisma(query_rows=[])
result = await get_tool(prisma, "nonexistent")
assert result is None
@pytest.mark.asyncio
async def test_update_tool_policy_calls_execute_raw():
row = {
"tool_id": "uuid-1",
"tool_name": "my_tool",
"origin": "user_defined",
"call_policy": "blocked",
"call_count": 1,
"assignments": {},
"key_hash": None,
"team_id": None,
"key_alias": None,
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
"created_by": None,
"updated_by": "admin",
}
prisma = _make_prisma(query_rows=[row])
result = await update_tool_policy(prisma, "my_tool", "blocked", "admin")
assert result is not None
assert result.call_policy == "blocked"
prisma.db.execute_raw.assert_awaited_once()
call_args = prisma.db.execute_raw.call_args
sql = call_args.args[0]
assert "ON CONFLICT" in sql
assert "call_policy" in sql
@pytest.mark.asyncio
async def test_get_tools_by_names_returns_policy_map():
rows = [
{"tool_name": "tool_a", "call_policy": "trusted"},
{"tool_name": "tool_b", "call_policy": "blocked"},
]
prisma = _make_prisma(query_rows=rows)
result = await get_tools_by_names(prisma, ["tool_a", "tool_b"])
assert result == {"tool_a": "trusted", "tool_b": "blocked"}
@pytest.mark.asyncio
async def test_get_tools_by_names_empty_list():
prisma = _make_prisma()
result = await get_tools_by_names(prisma, [])
assert result == {}
prisma.db.query_raw.assert_not_awaited()

124
tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py
View file

@ -2,10 +2,9 @@
"""
Test to verify the Google GenAI proxy API endpoints
"""
import asyncio
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, patch
import pytest
@ -13,7 +12,6 @@ sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
import litellm
def test_google_generate_content_endpoint():
@ -401,3 +399,123 @@ def test_google_generate_content_with_image_config():
assert "contents" in called_data
assert len(called_data["contents"]) == 1
assert called_data["contents"][0]["role"] == "user"
def test_google_generate_content_metadata_and_trace_id_callbacks():
"""Test that google_generate_content sets litellm_call_id and logging_obj for callbacks (e.g. S3, Langfuse)"""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock all required proxy server dependencies
with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch(
"litellm.proxy.proxy_server.general_settings", {}
), patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, patch(
"litellm.proxy.proxy_server.version", "1.0.0"
), patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
) as mock_add_data:
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
# Mock add_litellm_data_to_request to return data with metadata
async def mock_add_litellm_data(
data, request, user_api_key_dict, proxy_config, general_settings, version
):
# Simulate adding user metadata
data["litellm_metadata"] = {
"user_api_key_user_id": "test-user-id",
}
return data
mock_add_data.side_effect = mock_add_litellm_data
# Send a request to the endpoint with x-litellm-call-id header
test_call_id = "test-custom-call-id"
response = client.post(
"/v1beta/models/test-model:generateContent",
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
headers={
"Authorization": "Bearer sk-test-key",
"x-litellm-call-id": test_call_id,
},
)
assert response.status_code == 200
mock_router.agenerate_content.assert_called_once()
call_args = mock_router.agenerate_content.call_args
called_data = call_args[1]
# Verify that the litellm_logging_obj got assigned in the final called_data to router
assert "litellm_logging_obj" in called_data
assert "litellm_call_id" in called_data
assert called_data["litellm_call_id"] == test_call_id
def test_google_stream_generate_content_metadata_and_trace_id_callbacks():
"""Test that google_stream_generate_content sets litellm_call_id and logging_obj for callbacks"""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
app = FastAPI()
app.include_router(google_router)
client = TestClient(app)
mock_stream = AsyncMock()
mock_stream.__aiter__ = lambda self: mock_stream
mock_stream.__anext__.side_effect = StopAsyncIteration
with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch(
"litellm.proxy.proxy_server.general_settings", {}
), patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, patch(
"litellm.proxy.proxy_server.version", "1.0.0"
), patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
) as mock_add_data:
mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream)
async def mock_add_litellm_data(
data, request, user_api_key_dict, proxy_config, general_settings, version
):
data["litellm_metadata"] = {
"user_api_key_user_id": "test-user-id",
}
return data
mock_add_data.side_effect = mock_add_litellm_data
test_call_id = "test-custom-stream-call-id"
response = client.post(
"/v1beta/models/test-model:streamGenerateContent",
json={"contents": [{"role": "user", "parts": [{"text": "Hello stream"}]}]},
headers={
"Authorization": "Bearer sk-test-key",
"x-litellm-call-id": test_call_id,
},
)
assert response.status_code == 200
mock_router.agenerate_content_stream.assert_called_once()
call_args = mock_router.agenerate_content_stream.call_args
called_data = call_args[1]
assert "litellm_logging_obj" in called_data
assert "litellm_call_id" in called_data
assert called_data["litellm_call_id"] == test_call_id

166
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py
View file

@ -24,12 +24,29 @@ from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType
from litellm.types.utils import Choices, Message, ModelResponse
def _make_mock_session_iterator(json_response):
def _make_mock_session_iterator(
json_response, status=200, content_type="application/json", text_response=""
):
"""Create a mock _get_session_iterator that yields a session returning json_response."""
@asynccontextmanager
async def mock_iterator():
class MockResponse:
def __init__(self):
self.status = status
self.content_type = content_type
self.headers = {"Content-Type": content_type}
async def text(self):
if text_response:
return text_response
import json
try:
return json.dumps(json_response)
except Exception:
return str(json_response)
async def json(self):
return json_response
@ -41,6 +58,7 @@ def _make_mock_session_iterator(json_response):
class MockSession:
def post(self, *args, **kwargs):
self.last_kwargs = kwargs
return MockResponse()
async def __aenter__(self):
@ -1444,3 +1462,149 @@ def test_deny_list_and_score_threshold_combined():
# EMAIL_ADDRESS passes both filters
assert len(filtered) == 1
assert filtered[0]["entity_type"] == "EMAIL_ADDRESS"
@pytest.mark.asyncio
async def test_analyze_text_non_json_content_type_fail_closed():
"""
Test that analyze_text raises GuardrailRaisedException when Presidio health
endpoint returns text/html and fail-closed is enabled.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
presidio_analyzer_api_base="http://test-analyzer/",
presidio_anonymizer_api_base="http://test-anonymizer/",
pii_entities_config={"PERSON": PiiAction.BLOCK},
mock_testing=False,
)
mock_iterator = _make_mock_session_iterator(
json_response=None,
status=200,
content_type="text/html; charset=utf-8",
text_response="Presidio Analyzer service is up.",
)
with patch.object(guardrail, "_get_session_iterator", mock_iterator):
with pytest.raises(GuardrailRaisedException) as exc_info:
await guardrail.analyze_text(
text="Hello world",
presidio_config=None,
request_data={},
)
assert "expected application/json Content-Type" in str(exc_info.value)
assert "text/html" in str(exc_info.value)
@pytest.mark.asyncio
async def test_analyze_text_non_json_content_type_fail_open():
"""
Test that analyze_text returns empty list when Presidio returns text/html
and fail-closed is NOT enabled.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
presidio_analyzer_api_base="http://test-analyzer/",
presidio_anonymizer_api_base="http://test-anonymizer/",
mock_testing=False,
)
mock_iterator = _make_mock_session_iterator(
json_response=None,
status=200,
content_type="text/html; charset=utf-8",
text_response="Presidio Analyzer service is up.",
)
with patch.object(guardrail, "_get_session_iterator", mock_iterator):
results = await guardrail.analyze_text(
text="Hello world",
presidio_config=None,
request_data={},
)
assert results == []
@pytest.mark.asyncio
async def test_analyze_text_http_error_status():
"""
Test that analyze_text handles 5xx HTTP errors properly.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
presidio_analyzer_api_base="http://test-analyzer/",
presidio_anonymizer_api_base="http://test-anonymizer/",
pii_entities_config={"PERSON": PiiAction.BLOCK},
mock_testing=False,
)
mock_iterator = _make_mock_session_iterator(
json_response=None,
status=500,
content_type="text/plain",
text_response="Internal Server Error",
)
with patch.object(guardrail, "_get_session_iterator", mock_iterator):
with pytest.raises(GuardrailRaisedException) as exc_info:
await guardrail.analyze_text(
text="Hello world",
presidio_config=None,
request_data={},
)
assert "HTTP 500" in str(exc_info.value)
@pytest.mark.asyncio
async def test_anonymize_text_non_json_content_type():
"""
Test that anonymize_text raises Exception for non-JSON responses.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
presidio_analyzer_api_base="http://test-analyzer/",
presidio_anonymizer_api_base="http://test-anonymizer/",
mock_testing=False,
)
mock_iterator = _make_mock_session_iterator(
json_response=None,
status=200,
content_type="text/html",
text_response="Presidio Anonymizer service is up.",
)
with patch.object(guardrail, "_get_session_iterator", mock_iterator):
with pytest.raises(
Exception, match="Presidio anonymizer returned non-JSON Content-Type"
):
await guardrail.anonymize_text(
text="Hello world",
analyze_results=[{"start": 0, "end": 5, "entity_type": "PERSON"}],
output_parse_pii=False,
masked_entity_count={},
)
@pytest.mark.asyncio
async def test_anonymize_text_http_error_status():
"""
Test that anonymize_text raises Exception on HTTP error.
"""
guardrail = _OPTIONAL_PresidioPIIMasking(
presidio_analyzer_api_base="http://test-analyzer/",
presidio_anonymizer_api_base="http://test-anonymizer/",
mock_testing=False,
)
mock_iterator = _make_mock_session_iterator(
json_response=None,
status=502,
content_type="text/plain",
text_response="Bad Gateway",
)
with patch.object(guardrail, "_get_session_iterator", mock_iterator):
with pytest.raises(Exception, match="Presidio anonymizer returned HTTP 502"):
await guardrail.anonymize_text(
text="Hello world",
analyze_results=[{"start": 0, "end": 5, "entity_type": "PERSON"}],
output_parse_pii=False,
masked_entity_count={},
)

181
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py Normal file
View file

@ -0,0 +1,181 @@
"""
Unit tests for ToolPolicyGuardrail.
"""
import os
import sys
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
sys.path.insert(0, os.path.abspath("../../../../../.."))
from litellm.proxy.guardrails.guardrail_hooks.tool_policy.tool_policy_guardrail import (
ToolPolicyGuardrail,
)
from litellm.types.guardrails import GuardrailEventHooks
@pytest.fixture
def guardrail():
return ToolPolicyGuardrail()
# --- helpers ---
def _tool_request_inputs(tool_names: list) -> dict:
return {
"tools": [
{"type": "function", "function": {"name": name, "description": ""}}
for name in tool_names
]
}
def _tool_response_inputs(tool_names: list) -> dict:
return {
"tool_calls": [
{"type": "function", "function": {"name": name}}
for name in tool_names
]
}
# --- tests ---
def test_guardrail_supports_pre_and_post_call(guardrail):
hooks = guardrail.supported_event_hooks
assert GuardrailEventHooks.pre_call in hooks
assert GuardrailEventHooks.post_call in hooks
@pytest.mark.asyncio
async def test_no_tools_in_request_passes_through(guardrail):
inputs: Any = {"tools": []}
result = await guardrail.apply_guardrail(
inputs=inputs, request_data={}, input_type="request"
)
assert result is inputs
@pytest.mark.asyncio
async def test_no_tool_calls_in_response_passes_through(guardrail):
inputs: Any = {"tool_calls": []}
result = await guardrail.apply_guardrail(
inputs=inputs, request_data={}, input_type="response"
)
assert result is inputs
@pytest.mark.asyncio
async def test_untrusted_tools_pass_through(guardrail):
policy_map = {"search": "untrusted", "read_file": "trusted"}
with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)):
inputs: Any = _tool_request_inputs(["search", "read_file"])
result = await guardrail.apply_guardrail(
inputs=inputs, request_data={}, input_type="request"
)
assert result is inputs
@pytest.mark.asyncio
async def test_blocked_tool_in_request_raises_http_exception(guardrail):
policy_map = {"dangerous_tool": "blocked"}
with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)):
inputs: Any = _tool_request_inputs(["dangerous_tool"])
with pytest.raises(HTTPException) as exc_info:
await guardrail.apply_guardrail(
inputs=inputs, request_data={}, input_type="request"
)
assert exc_info.value.status_code == 400
assert "dangerous_tool" in exc_info.value.detail["blocked_tools"]
@pytest.mark.asyncio
async def test_blocked_tool_in_response_raises_http_exception(guardrail):
policy_map = {"exfil_tool": "blocked"}
with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)):
inputs: Any = _tool_response_inputs(["exfil_tool"])
with pytest.raises(HTTPException) as exc_info:
await guardrail.apply_guardrail(
inputs=inputs, request_data={}, input_type="response"
)
assert exc_info.value.status_code == 400
assert "exfil_tool" in exc_info.value.detail["blocked_tools"]
@pytest.mark.asyncio
async def test_mixed_blocked_and_allowed_raises_for_blocked(guardrail):
policy_map = {"safe_tool": "trusted", "bad_tool": "blocked"}
with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)):
inputs: Any = _tool_request_inputs(["safe_tool", "bad_tool"])
with pytest.raises(HTTPException) as exc_info:
await guardrail.apply_guardrail(
inputs=inputs, request_data={}, input_type="request"
)
blocked = exc_info.value.detail["blocked_tools"]
assert "bad_tool" in blocked
assert "safe_tool" not in blocked
@pytest.mark.asyncio
async def test_tool_not_in_db_passes_through(guardrail):
"""Tools not found in the DB (no entry) should not be blocked."""
with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value={})):
inputs: Any = _tool_request_inputs(["unknown_tool"])
result = await guardrail.apply_guardrail(
inputs=inputs, request_data={}, input_type="request"
)
assert result is inputs
@pytest.mark.asyncio
async def test_get_policies_cached_uses_cache(guardrail):
"""Second call with same tool names should return the cached result."""
policy_map = {"tool_a": "trusted"}
with patch(
"litellm.proxy.db.tool_registry_writer.get_tools_by_names",
new=AsyncMock(return_value=policy_map),
) as mock_db, patch(
"litellm.proxy.proxy_server.prisma_client",
new=MagicMock(),
):
# first call — should hit DB
result1 = await guardrail._get_policies_cached(["tool_a"])
assert result1 == policy_map
# second call — should hit cache, not DB again
result2 = await guardrail._get_policies_cached(["tool_a"])
assert result2 == policy_map
assert mock_db.call_count == 1
@pytest.mark.asyncio
async def test_get_policies_cached_no_prisma(guardrail):
"""Without a prisma client, returns empty dict."""
with patch(
"litellm.proxy.proxy_server.prisma_client",
None,
):
result = await guardrail._get_policies_cached(["tool_a"])
assert result == {}
@pytest.mark.asyncio
async def test_response_tool_calls_as_objects(guardrail):
"""tool_calls that are objects (not dicts) with .function.name should work."""
policy_map = {"obj_tool": "blocked"}
with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)):
fn = MagicMock()
fn.name = "obj_tool"
tc = MagicMock()
tc.function = fn
inputs: Any = {"tool_calls": [tc]}
with pytest.raises(HTTPException):
await guardrail.apply_guardrail(
inputs=inputs, request_data={}, input_type="response"
)

217
tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -169,6 +169,223 @@ async def test_track_cost_callback_skips_when_no_standard_logging_object():
mock_proxy_logging.failed_tracking_alert.assert_not_called()
@pytest.mark.asyncio
async def test_enrich_failure_metadata_with_team_alias():
"""
When team_id is set but team_alias is missing (and key_alias is present),
_enrich_failure_metadata_with_key_info should look up the team from cache
and populate user_api_key_team_alias.
"""
mock_team_obj = MagicMock()
mock_team_obj.team_alias = "my-team-alias"
with patch(
"litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
new_callable=AsyncMock,
return_value=mock_team_obj,
):
metadata = {
"user_api_key": "hashed_key",
"user_api_key_alias": "my-key-alias", # already set
"user_api_key_team_id": "test_team_id",
"user_api_key_team_alias": None,
}
result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
assert result["user_api_key_team_alias"] == "my-team-alias"
@pytest.mark.asyncio
async def test_enrich_failure_metadata_with_full_key_lookup():
"""
When all key fields are null (auth error 401 scenario), _enrich_failure_metadata_with_key_info
should look up the key object from cache/DB and populate alias, user_id, team_id,
then look up the team to get team_alias.
"""
mock_key_obj = MagicMock()
mock_key_obj.key_alias = "fetched-key-alias"
mock_key_obj.user_id = "fetched-user-id"
mock_key_obj.team_id = "fetched-team-id"
mock_key_obj.org_id = "fetched-org-id"
mock_team_obj = MagicMock()
mock_team_obj.team_alias = "fetched-team-alias"
with patch(
"litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
new_callable=AsyncMock,
return_value=mock_key_obj,
), patch(
"litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
new_callable=AsyncMock,
return_value=mock_team_obj,
):
metadata = {
"user_api_key": "hashed_key",
"user_api_key_alias": None, # all null - simulates auth error path
"user_api_key_user_id": None,
"user_api_key_team_id": None,
"user_api_key_team_alias": None,
"user_api_key_org_id": None,
}
result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
assert result["user_api_key_alias"] == "fetched-key-alias"
assert result["user_api_key_user_id"] == "fetched-user-id"
assert result["user_api_key_team_id"] == "fetched-team-id"
assert result["user_api_key_org_id"] == "fetched-org-id"
assert result["user_api_key_team_alias"] == "fetched-team-alias"
@pytest.mark.asyncio
async def test_enrich_failure_metadata_skips_when_team_alias_present():
"""
When team_alias is already populated, _enrich_failure_metadata_with_key_info
should not perform a team cache lookup.
"""
with patch(
"litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
new_callable=AsyncMock,
) as mock_get_key, patch(
"litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
new_callable=AsyncMock,
) as mock_get_team:
metadata = {
"user_api_key": "hashed_key",
"user_api_key_alias": "existing-alias",
"user_api_key_team_id": "test_team_id",
"user_api_key_team_alias": "already-set",
}
result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
assert result["user_api_key_team_alias"] == "already-set"
mock_get_key.assert_not_called()
mock_get_team.assert_not_called()
@pytest.mark.asyncio
async def test_enrich_failure_metadata_skips_when_no_api_key():
"""
When api_key hash is absent, _enrich_failure_metadata_with_key_info should
not perform any lookups.
"""
with patch(
"litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
new_callable=AsyncMock,
) as mock_get_key:
metadata = {
"user_api_key": None,
"user_api_key_alias": None,
"user_api_key_team_id": None,
"user_api_key_team_alias": None,
}
result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
mock_get_key.assert_not_called()
@pytest.mark.asyncio
async def test_async_post_call_failure_hook_enriches_auth_error_metadata():
"""
Simulates a 401 ProxyException (e.g. can_key_call_model). In this case
UserAPIKeyAuth is created with only api_key set. The failure hook should
look up the key and team from cache/DB to populate all missing fields.
"""
logger = _ProxyDBLogger()
# This is what auth_exception_handler creates for 401 errors
user_api_key_dict = UserAPIKeyAuth(
api_key="hashed_key",
# key_alias, user_id, team_id, team_alias are all None
)
request_data = {
"model": "claude-haiku-4-5",
"messages": [{"role": "user", "content": "Hello"}],
"metadata": {},
"litellm_params": {},
}
mock_key_obj = MagicMock()
mock_key_obj.key_alias = "my-key-alias"
mock_key_obj.user_id = "my-user-id"
mock_key_obj.team_id = "my-team-id"
mock_key_obj.org_id = None
mock_team_obj = MagicMock()
mock_team_obj.team_alias = "my-team-alias"
with patch(
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
) as mock_update_database, patch(
"litellm.proxy.hooks.proxy_track_cost_callback.get_key_object",
new_callable=AsyncMock,
return_value=mock_key_obj,
), patch(
"litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
new_callable=AsyncMock,
return_value=mock_team_obj,
):
await logger.async_post_call_failure_hook(
request_data=request_data,
original_exception=Exception("401 - model not allowed"),
user_api_key_dict=user_api_key_dict,
)
mock_update_database.assert_called_once()
call_args = mock_update_database.call_args[1]
metadata = call_args["kwargs"]["litellm_params"]["metadata"]
assert metadata["user_api_key_alias"] == "my-key-alias"
assert metadata["user_api_key_user_id"] == "my-user-id"
assert metadata["user_api_key_team_id"] == "my-team-id"
assert metadata["user_api_key_team_alias"] == "my-team-alias"
@pytest.mark.asyncio
async def test_async_post_call_failure_hook_enriches_missing_team_alias():
"""
When user_api_key_dict has a team_id but no team_alias, async_post_call_failure_hook
should look up the team from cache and populate user_api_key_team_alias in the
spend log metadata written to the DB.
"""
logger = _ProxyDBLogger()
user_api_key_dict = UserAPIKeyAuth(
api_key="test_api_key",
key_alias="test_alias",
user_id="test_user_id",
team_id="test_team_id",
team_alias=None, # Missing - simulates regular key auth where SQL view omits team_alias
)
request_data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hello"}],
"metadata": {},
"litellm_params": {},
}
mock_team_obj = MagicMock()
mock_team_obj.team_alias = "enriched-team-alias"
with patch(
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
) as mock_update_database, patch(
"litellm.proxy.hooks.proxy_track_cost_callback.get_team_object",
new_callable=AsyncMock,
return_value=mock_team_obj,
):
await logger.async_post_call_failure_hook(
request_data=request_data,
original_exception=Exception("Provider rate limit"),
user_api_key_dict=user_api_key_dict,
)
mock_update_database.assert_called_once()
call_args = mock_update_database.call_args[1]
metadata = call_args["kwargs"]["litellm_params"]["metadata"]
assert metadata["user_api_key_team_alias"] == "enriched-team-alias"
assert metadata["user_api_key_team_id"] == "test_team_id"
@pytest.mark.asyncio
@pytest.mark.parametrize("model_value", [None, ""])
async def test_track_cost_callback_skips_for_falsy_model_and_no_slo(model_value):

54
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -5623,6 +5623,7 @@ async def test_rotate_master_key_model_data_valid_for_prisma(
litellm_params/model_info are JSON strings (create_many expects dicts).
"""
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.key_management_endpoints import (
_rotate_master_key,
@ -6157,3 +6158,56 @@ async def test_get_member_team_ids():
# Should return team-A and team-B (user is a member of both)
# Should NOT return team-C (user is not in members list)
assert sorted(result) == ["team-A", "team-B"]
@pytest.mark.asyncio
async def test_generate_key_with_agent_id():
"""Test that agent_id is accepted in GenerateKeyRequest and passed to generate_key_helper_fn."""
from litellm.proxy._types import GenerateKeyRequest
# Verify GenerateKeyRequest accepts agent_id
request = GenerateKeyRequest(
key_alias="agent-test-key",
agent_id="test-agent-123",
models=[],
)
assert request.agent_id == "test-agent-123"
data_json = request.model_dump(exclude_unset=True, exclude_none=True)
assert data_json["agent_id"] == "test-agent-123"
@pytest.mark.asyncio
async def test_generate_key_helper_fn_agent_id():
"""Test that generate_key_helper_fn passes agent_id into the insert_data call."""
from unittest.mock import AsyncMock, MagicMock, call, patch
import litellm.proxy.management_endpoints.key_management_endpoints as km
mock_prisma_client = AsyncMock()
mock_insert = AsyncMock(
return_value=MagicMock(
token="sk-test",
created_at=None,
updated_at=None,
litellm_budget_table=None,
)
)
mock_prisma_client.insert_data = mock_insert
with patch.object(km, "prisma_client", mock_prisma_client):
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
await generate_key_helper_fn(
request_type="key",
agent_id="test-agent-456",
key_alias="test-agent-key",
models=[],
table_name="key",
)
assert mock_insert.called, "insert_data was never called"
# insert_data is called as insert_data(data=key_data, ...)
call_kwargs = mock_insert.call_args.kwargs
key_data = call_kwargs.get("data", {})
assert key_data.get("agent_id") == "test-agent-456", (
f"Expected agent_id='test-agent-456' in key_data, got: {key_data.get('agent_id')}"
)

149
tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py Normal file
View file

@ -0,0 +1,149 @@
"""
Unit tests for tool management endpoints (/v1/tool/*).
Uses FastAPI TestClient with mocked DB functions.
Patches target the source modules (litellm.proxy.db.tool_registry_writer.*
and litellm.proxy.proxy_server.prisma_client) because the endpoint code
imports these inside function bodies to avoid circular imports.
"""
import os
import sys
from datetime import datetime, timezone
from typing import Optional
from unittest.mock import AsyncMock, MagicMock, patch
from fastapi import FastAPI
from fastapi.testclient import TestClient
sys.path.insert(0, os.path.abspath("../../.."))
from litellm.proxy.management_endpoints.tool_management_endpoints import router
from litellm.types.tool_management import LiteLLM_ToolTableRow
# --- helpers ---
def _make_tool_row(
tool_name: str = "my_tool",
call_policy: str = "untrusted",
origin: Optional[str] = None,
) -> LiteLLM_ToolTableRow:
now = datetime.now(timezone.utc)
return LiteLLM_ToolTableRow(
tool_id="uuid-1",
tool_name=tool_name,
origin=origin,
call_policy=call_policy, # type: ignore[arg-type]
assignments={},
created_at=now,
updated_at=now,
)
def _make_app() -> FastAPI:
"""Build a minimal FastAPI app with the tool management router."""
app = FastAPI()
app.include_router(router)
return app
# Stub the auth dependency so we don't need a real proxy running.
def _override_auth():
from litellm.proxy._types import UserAPIKeyAuth
return UserAPIKeyAuth(api_key="sk-test", user_id="admin")
# A real (non-None) prisma stub for truthiness checks.
_MOCK_PRISMA = MagicMock()
# --- test class ---
class TestToolManagementEndpoints:
def setup_method(self):
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
app = _make_app()
app.dependency_overrides[user_api_key_auth] = _override_auth
self.client = TestClient(app, raise_server_exceptions=True)
@patch(
"litellm.proxy.db.tool_registry_writer.list_tools",
new_callable=AsyncMock,
)
@patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA)
def test_list_tools_returns_200(self, mock_db_list):
mock_db_list.return_value = [_make_tool_row()]
resp = self.client.get("/v1/tool/list")
assert resp.status_code == 200
body = resp.json()
assert body["total"] == 1
assert body["tools"][0]["tool_name"] == "my_tool"
@patch(
"litellm.proxy.db.tool_registry_writer.list_tools",
new_callable=AsyncMock,
)
@patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA)
def test_list_tools_with_policy_filter(self, mock_db_list):
mock_db_list.return_value = [_make_tool_row(call_policy="blocked")]
resp = self.client.get("/v1/tool/list?call_policy=blocked")
assert resp.status_code == 200
assert resp.json()["tools"][0]["call_policy"] == "blocked"
@patch(
"litellm.proxy.db.tool_registry_writer.get_tool",
new_callable=AsyncMock,
)
@patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA)
def test_get_tool_found(self, mock_db_get):
mock_db_get.return_value = _make_tool_row(tool_name="tool_a")
resp = self.client.get("/v1/tool/tool_a")
assert resp.status_code == 200
assert resp.json()["tool_name"] == "tool_a"
@patch(
"litellm.proxy.db.tool_registry_writer.get_tool",
new_callable=AsyncMock,
)
@patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA)
def test_get_tool_not_found_returns_404(self, mock_db_get):
mock_db_get.return_value = None
resp = self.client.get("/v1/tool/nonexistent", follow_redirects=True)
assert resp.status_code == 404
@patch(
"litellm.proxy.db.tool_registry_writer.update_tool_policy",
new_callable=AsyncMock,
)
@patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA)
def test_update_tool_policy_blocked(self, mock_db_update):
mock_db_update.return_value = _make_tool_row(call_policy="blocked")
resp = self.client.post(
"/v1/tool/policy",
json={"tool_name": "my_tool", "call_policy": "blocked"},
)
assert resp.status_code == 200
body = resp.json()
assert body["call_policy"] == "blocked"
assert body["updated"] is True
@patch("litellm.proxy.proxy_server.prisma_client", None)
def test_list_tools_no_db_returns_500(self):
resp = self.client.get("/v1/tool/list")
assert resp.status_code == 500
def test_update_tool_policy_invalid_policy_returns_422(self):
resp = self.client.post(
"/v1/tool/policy",
json={"tool_name": "my_tool", "call_policy": "invalid_value"},
)
assert resp.status_code == 422

0
tests/test_litellm/proxy/management_endpoints/usage_endpoints/__init__.py Normal file
View file

402
tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py Normal file
View file

@ -0,0 +1,402 @@
"""
Tests for AI Usage Chat module.
"""
import json
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import (
TOOL_HANDLERS,
TOOLS_ADMIN,
TOOLS_BASE,
_build_system_prompt,
_summarise_entity_data,
_summarise_usage_data,
stream_usage_ai_chat,
)
SAMPLE_AGGREGATED_RESPONSE = {
"results": [
{
"date": "2025-01-15",
"metrics": {
"spend": 50.25,
"prompt_tokens": 20000,
"completion_tokens": 10000,
"total_tokens": 30000,
"api_requests": 500,
"successful_requests": 480,
"failed_requests": 20,
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 0,
},
"breakdown": {
"models": {
"gpt-4": {
"metrics": {
"spend": 40.0,
"api_requests": 300,
"total_tokens": 25000,
},
"metadata": {},
"api_key_breakdown": {},
},
},
"providers": {
"openai": {
"metrics": {"spend": 50.25, "api_requests": 500},
"metadata": {},
"api_key_breakdown": {},
},
},
"api_keys": {
"sk-test123": {
"metrics": {"spend": 50.25},
"metadata": {"key_alias": "Production Key"},
},
},
"model_groups": {},
"mcp_servers": {},
"entities": {},
},
},
],
"metadata": {
"total_spend": 50.25,
"total_api_requests": 500,
"total_successful_requests": 480,
"total_failed_requests": 20,
"total_tokens": 30000,
},
}
SAMPLE_TEAM_RESPONSE = {
"results": [
{
"date": "2025-01-15",
"metrics": {"spend": 100.0, "api_requests": 1000, "total_tokens": 50000},
"breakdown": {
"entities": {
"team-1": {
"metrics": {
"spend": 60.0,
"api_requests": 600,
"total_tokens": 30000,
},
"metadata": {"alias": "Engineering"},
"api_key_breakdown": {},
},
"team-2": {
"metrics": {
"spend": 40.0,
"api_requests": 400,
"total_tokens": 20000,
},
"metadata": {"alias": "Marketing"},
"api_key_breakdown": {},
},
},
"models": {},
"providers": {},
"api_keys": {},
"model_groups": {},
"mcp_servers": {},
},
},
],
"metadata": {"total_spend": 100.0, "total_api_requests": 1000},
}
class TestToolSchemas:
def test_admin_tools_include_all(self):
assert len(TOOLS_ADMIN) == 3
names = {t["function"]["name"] for t in TOOLS_ADMIN}
assert "get_usage_data" in names
assert "get_team_usage_data" in names
assert "get_tag_usage_data" in names
def test_base_tools_restricted_to_usage_only(self):
assert len(TOOLS_BASE) == 1
assert TOOLS_BASE[0]["function"]["name"] == "get_usage_data"
def test_admin_prompt_mentions_all_tools(self):
prompt = _build_system_prompt(is_admin=True)
assert "get_usage_data" in prompt
assert "get_team_usage_data" in prompt
assert "get_tag_usage_data" in prompt
def test_non_admin_prompt_only_mentions_usage_tool(self):
prompt = _build_system_prompt(is_admin=False)
assert "get_usage_data" in prompt
assert "get_team_usage_data" not in prompt
assert "get_tag_usage_data" not in prompt
def test_system_prompt_includes_todays_date(self):
from datetime import date
prompt = _build_system_prompt(is_admin=True)
assert date.today().isoformat() in prompt
class TestSummariseUsageData:
def test_summarise_includes_totals(self):
summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE)
assert "$50.25" in summary
assert "500" in summary
def test_summarise_includes_models(self):
summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE)
assert "gpt-4" in summary
def test_summarise_includes_providers(self):
summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE)
assert "openai" in summary
def test_summarise_handles_empty_data(self):
empty = {"results": [], "metadata": {}}
summary = _summarise_usage_data(empty)
assert "no data" in summary.lower()
class TestSummariseEntityData:
def test_team_summary_includes_teams(self):
summary = _summarise_entity_data(SAMPLE_TEAM_RESPONSE, "Team")
assert "Engineering" in summary
assert "Marketing" in summary
assert "$60.0" in summary
assert "$40.0" in summary
def test_team_summary_empty(self):
empty = {"results": [], "metadata": {}}
summary = _summarise_entity_data(empty, "Team")
assert "No Team usage data" in summary
class TestStreamUsageAiChat:
@pytest.mark.asyncio
async def test_stream_emits_status_events(self):
mock_tool_call = MagicMock()
mock_tool_call.id = "call_123"
mock_tool_call.function.name = "get_usage_data"
mock_tool_call.function.arguments = json.dumps(
{
"start_date": "2025-01-01",
"end_date": "2025-01-31",
}
)
mock_first_response = MagicMock()
mock_first_response.choices = [MagicMock()]
mock_first_response.choices[0].message.tool_calls = [mock_tool_call]
mock_first_response.choices[0].message.model_dump.return_value = {
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_123",
"type": "function",
"function": {
"name": "get_usage_data",
"arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31"}',
},
}
],
}
async def mock_stream():
chunk = MagicMock()
chunk.choices = [MagicMock()]
chunk.choices[0].delta.content = "Total spend is $50.25"
yield chunk
with patch(
"litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm"
) as mock_litellm, patch(
"litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat._fetch_usage_data",
new_callable=AsyncMock,
) as mock_fetch:
mock_litellm.acompletion = AsyncMock(
side_effect=[
mock_first_response,
mock_stream(),
]
)
mock_fetch.return_value = SAMPLE_AGGREGATED_RESPONSE
events = []
async for event in stream_usage_ai_chat(
messages=[{"role": "user", "content": "What is my total spend?"}],
model="gpt-4o-mini",
user_id="user-123",
is_admin=True,
):
events.append(json.loads(event.replace("data: ", "").strip()))
status_events = [e for e in events if e["type"] == "status"]
tool_call_events = [e for e in events if e["type"] == "tool_call"]
chunk_events = [e for e in events if e["type"] == "chunk"]
done_events = [e for e in events if e["type"] == "done"]
assert len(status_events) >= 1
assert "Thinking" in status_events[0]["message"]
assert len(tool_call_events) >= 1
assert tool_call_events[0]["tool_name"] == "get_usage_data"
assert tool_call_events[0]["status"] in ("running", "complete")
assert len(chunk_events) >= 1
assert len(done_events) == 1
@pytest.mark.asyncio
async def test_stream_handles_team_tool(self):
mock_tool_call = MagicMock()
mock_tool_call.id = "call_team"
mock_tool_call.function.name = "get_team_usage_data"
mock_tool_call.function.arguments = json.dumps(
{
"start_date": "2025-01-01",
"end_date": "2025-01-31",
}
)
mock_first_response = MagicMock()
mock_first_response.choices = [MagicMock()]
mock_first_response.choices[0].message.tool_calls = [mock_tool_call]
mock_first_response.choices[0].message.model_dump.return_value = {
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_team",
"type": "function",
"function": {
"name": "get_team_usage_data",
"arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31"}',
},
}
],
}
async def mock_stream():
chunk = MagicMock()
chunk.choices = [MagicMock()]
chunk.choices[0].delta.content = "Engineering is the top team."
yield chunk
with patch(
"litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm"
) as mock_litellm, patch(
"litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat._fetch_team_usage_data",
new_callable=AsyncMock,
) as mock_fetch:
mock_litellm.acompletion = AsyncMock(
side_effect=[
mock_first_response,
mock_stream(),
]
)
mock_fetch.return_value = SAMPLE_TEAM_RESPONSE
events = []
async for event in stream_usage_ai_chat(
messages=[{"role": "user", "content": "Which team spends the most?"}],
model="gpt-4o-mini",
is_admin=True,
):
events.append(json.loads(event.replace("data: ", "").strip()))
chunk_events = [e for e in events if e["type"] == "chunk"]
assert len(chunk_events) >= 1
assert "Engineering" in chunk_events[0]["content"]
@pytest.mark.asyncio
async def test_stream_handles_error(self):
with patch(
"litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm"
) as mock_litellm:
mock_litellm.acompletion = AsyncMock(side_effect=Exception("LLM error"))
events = []
async for event in stream_usage_ai_chat(
messages=[{"role": "user", "content": "test"}],
):
events.append(json.loads(event.replace("data: ", "").strip()))
error_events = [e for e in events if e["type"] == "error"]
assert len(error_events) == 1
assert "internal error" in error_events[0]["message"].lower()
@pytest.mark.asyncio
async def test_non_admin_enforces_user_id(self):
mock_tool_call = MagicMock()
mock_tool_call.id = "call_456"
mock_tool_call.function.name = "get_usage_data"
mock_tool_call.function.arguments = json.dumps(
{
"start_date": "2025-01-01",
"end_date": "2025-01-31",
"user_id": "other-user",
}
)
mock_first_response = MagicMock()
mock_first_response.choices = [MagicMock()]
mock_first_response.choices[0].message.tool_calls = [mock_tool_call]
mock_first_response.choices[0].message.model_dump.return_value = {
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_456",
"type": "function",
"function": {
"name": "get_usage_data",
"arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31","user_id":"other-user"}',
},
}
],
}
async def mock_stream():
chunk = MagicMock()
chunk.choices = [MagicMock()]
chunk.choices[0].delta.content = "Data."
yield chunk
mock_fetch = AsyncMock(return_value=SAMPLE_AGGREGATED_RESPONSE)
with patch(
"litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm"
) as mock_litellm, patch.dict(
"litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.TOOL_HANDLERS",
{
"get_usage_data": {
"fetch": mock_fetch,
"summarise": _summarise_usage_data,
"label": "global usage data",
}
},
):
mock_litellm.acompletion = AsyncMock(
side_effect=[
mock_first_response,
mock_stream(),
]
)
events = []
async for event in stream_usage_ai_chat(
messages=[{"role": "user", "content": "Show data"}],
model="gpt-4o-mini",
user_id="my-user-id",
is_admin=False,
):
events.append(event)
mock_fetch.assert_called_once_with(
start_date="2025-01-01",
end_date="2025-01-31",
user_id="my-user-id",
)

39
tests/test_litellm/proxy/test_health_check_functions.py
View file

@ -12,6 +12,7 @@ sys.path.insert(0, os.path.abspath("../../.."))
from litellm.proxy.health_endpoints._health_endpoints import (
_aggregate_health_check_results,
_build_model_param_to_info_mapping,
_perform_health_check_and_save,
_save_background_health_checks_to_db,
_save_health_check_results_if_changed,
_save_health_check_to_db,
@ -466,5 +467,43 @@ async def test_get_all_latest_health_checks_without_model_id(mock_prisma):
assert result[0].checked_at == mock_check2.checked_at # Latest
@pytest.mark.asyncio
async def test_perform_health_check_and_save_passes_model_id_to_perform_health_check():
"""Test that _perform_health_check_and_save passes model_id to perform_health_check so health checks run by model id."""
model_list = [
{
"model_name": "gpt-4",
"model_info": {"id": "deployment-abc"},
"litellm_params": {"model": "gpt-4"},
},
]
healthy = [{"model": "gpt-4"}]
unhealthy = []
async def mock_perform_health_check(model_list, model=None, cli_model=None, details=True, model_id=None):
return healthy, unhealthy
with patch(
"litellm.proxy.health_endpoints._health_endpoints.perform_health_check",
side_effect=mock_perform_health_check,
) as mock_perform:
result = await _perform_health_check_and_save(
model_list=model_list,
target_model=None,
cli_model=None,
details=True,
prisma_client=None,
start_time=0.0,
user_id="user-1",
model_id="deployment-abc",
)
mock_perform.assert_called_once()
call_kwargs = mock_perform.call_args[1]
assert call_kwargs["model_id"] == "deployment-abc"
assert result["healthy_count"] == 1
assert result["unhealthy_count"] == 0
if __name__ == "__main__":
pytest.main([__file__])

67
tests/test_litellm/test_router.py
View file

@ -925,6 +925,73 @@ def test_router_get_model_access_groups_team_only_models():
assert list(access_groups.keys()) == ["default-models"]
def test_cached_get_model_group_info():
"""
Test that _cached_get_model_group_info caches results and
invalidates on deployment changes.
"""
from litellm.types.router import Deployment, LiteLLM_Params
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4", "api_key": "fake"},
"model_info": {"tpm": 1000, "rpm": 100},
},
]
)
# First call should compute and cache
result1 = router._cached_get_model_group_info("gpt-4")
assert result1 is not None
assert result1.tpm == 1000
# Second call should hit cache (same object)
result2 = router._cached_get_model_group_info("gpt-4")
assert result1 is result2
# Add a deployment — cache should be invalidated
router.add_deployment(
Deployment(
model_name="gpt-4",
litellm_params=LiteLLM_Params(model="gpt-4", api_key="fake2"),
model_info={"tpm": 2000, "rpm": 200},
)
)
result3 = router._cached_get_model_group_info("gpt-4")
assert result3 is not result2
assert result3 is not None
assert result3.tpm == 3000 # 1000 + 2000
# Delete a deployment — cache should be invalidated
deployment_id = router.model_list[-1]["model_info"]["id"]
router.delete_deployment(id=deployment_id)
result4 = router._cached_get_model_group_info("gpt-4")
assert result4 is not result3
assert result4 is not None
assert result4.tpm == 1000
# set_model_list — cache should be invalidated
router.set_model_list(
[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4", "api_key": "fake"},
"model_info": {"tpm": 5000},
},
]
)
result5 = router._cached_get_model_group_info("gpt-4")
assert result5 is not result4
assert result5 is not None
assert result5.tpm == 5000
# Verify cache still works after invalidation
result6 = router._cached_get_model_group_info("gpt-4")
assert result5 is result6
def test_get_model_access_groups_caching():
"""
Test that get_model_access_groups caches the no-args result

151
ui/litellm-dashboard/package-lock.json generated
View file

@ -1757,29 +1757,6 @@
"url": "https://opencollective.com/libvips"
}
},
"node_modules/@isaacs/balanced-match": {
"version": "4.0.1",
"resolved": "https://registry.npmjs.org/@isaacs/balanced-match/-/balanced-match-4.0.1.tgz",
"integrity": "sha512-yzMTt9lEb8Gv7zRioUilSglI0c0smZ9k5D65677DLWLtWJaXIS3CqcGyUFByYKlnUj6TkjLVs54fBl6+TiGQDQ==",
"dev": true,
"license": "MIT",
"engines": {
"node": "20 || >=22"
}
},
"node_modules/@isaacs/brace-expansion": {
"version": "5.0.1",
"resolved": "https://registry.npmjs.org/@isaacs/brace-expansion/-/brace-expansion-5.0.1.tgz",
"integrity": "sha512-WMz71T1JS624nWj2n2fnYAuPovhv7EUhk69R6i9dsVyzxt5eM3bjwvgk9L+APE1TRscGysAVMANkB0jh0LQZrQ==",
"dev": true,
"license": "MIT",
"dependencies": {
"@isaacs/balanced-match": "^4.0.1"
},
"engines": {
"node": "20 || >=22"
}
},
"node_modules/@istanbuljs/schema": {
"version": "0.1.3",
"resolved": "https://registry.npmjs.org/@istanbuljs/schema/-/schema-0.1.3.tgz",
@ -3696,32 +3673,6 @@
"typescript": ">=4.8.4 <6.0.0"
}
},
"node_modules/@typescript-eslint/typescript-estree/node_modules/brace-expansion": {
"version": "2.0.2",
"resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.0.2.tgz",
"integrity": "sha512-Jt0vHyM+jmUBqojB7E1NIYadt0vI0Qxjxd2TErW94wDz+E2LAm5vKMXXwg6ZZBTHPuUlDgQHKXvjGBdfcF1ZDQ==",
"dev": true,
"license": "MIT",
"dependencies": {
"balanced-match": "^1.0.0"
}
},
"node_modules/@typescript-eslint/typescript-estree/node_modules/minimatch": {
"version": "9.0.5",
"resolved": "https://registry.npmjs.org/minimatch/-/minimatch-9.0.5.tgz",
"integrity": "sha512-G6T0ZX48xgozx7587koeX9Ys2NYy6Gmv//P89sEte9V9whIapMNF4idKxnW2QtCcLiTWlb/wfCabAtAFWhhBow==",
"dev": true,
"license": "ISC",
"dependencies": {
"brace-expansion": "^2.0.1"
},
"engines": {
"node": ">=16 || 14 >=14.17"
},
"funding": {
"url": "https://github.com/sponsors/isaacs"
}
},
"node_modules/@typescript-eslint/utils": {
"version": "8.54.0",
"resolved": "https://registry.npmjs.org/@typescript-eslint/utils/-/utils-8.54.0.tgz",
@ -4749,11 +4700,14 @@
}
},
"node_modules/balanced-match": {
"version": "1.0.2",
"resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-1.0.2.tgz",
"integrity": "sha512-3oSeUO0TMV67hN1AmbXsK4yaqU7tjiHlbxRDZOpH0KW9+CeX4bRAaX0Anxt0tx2MrpRpWwQaPwIlISEJhYU5Pw==",
"version": "4.0.4",
"resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-4.0.4.tgz",
"integrity": "sha512-BLrgEcRTwX2o6gGxGOCNyMvGSp35YofuYzw9h1IMTRmKqttAZZVU67bdb9Pr2vUHA8+j3i2tJfjO6C6+4myGTA==",
"dev": true,
"license": "MIT"
"license": "MIT",
"engines": {
"node": "18 || 20 || >=22"
}
},
"node_modules/baseline-browser-mapping": {
"version": "2.9.19",
@ -4787,14 +4741,16 @@
}
},
"node_modules/brace-expansion": {
"version": "1.1.12",
"resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.12.tgz",
"integrity": "sha512-9T9UjW3r0UW5c1Q7GTwllptXwhvYmEzFhzMfZ9H7FQWt+uZePjZPjBP/W1ZEyZ1twGWom5/56TF4lPcqjnDHcg==",
"version": "5.0.3",
"resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.3.tgz",
"integrity": "sha512-fy6KJm2RawA5RcHkLa1z/ScpBeA762UF9KmZQxwIbDtRJrgLzM10depAiEQ+CXYcoiqW1/m96OAAoke2nE9EeA==",
"dev": true,
"license": "MIT",
"dependencies": {
"balanced-match": "^1.0.0",
"concat-map": "0.0.1"
"balanced-match": "^4.0.2"
},
"engines": {
"node": "18 || 20 || >=22"
}
},
"node_modules/braces": {
@ -5149,13 +5105,6 @@
"integrity": "sha512-VRhuHOLoKYOy4UbilLbUzbYg93XLjv2PncJC50EuTWPA3gaja1UjBsUP/D/9/juV3vQFr6XBEzn9KCAHdUvOHw==",
"license": "MIT"
},
"node_modules/concat-map": {
"version": "0.0.1",
"resolved": "https://registry.npmjs.org/concat-map/-/concat-map-0.0.1.tgz",
"integrity": "sha512-/Srv4dswyQNBfohGpz9o6Yb3Gz3SrUDqBH5rTuhGR7ahtlbYKnVxw2bCFMRljaA7EXHaXZ8wsHdodFvbkhKmqg==",
"dev": true,
"license": "MIT"
},
"node_modules/copy-to-clipboard": {
"version": "3.3.3",
"resolved": "https://registry.npmjs.org/copy-to-clipboard/-/copy-to-clipboard-3.3.3.tgz",
@ -6924,22 +6873,6 @@
"node": ">=10.13.0"
}
},
"node_modules/glob/node_modules/minimatch": {
"version": "10.1.1",
"resolved": "https://registry.npmjs.org/minimatch/-/minimatch-10.1.1.tgz",
"integrity": "sha512-enIvLvRAFZYXJzkCYG5RKmPfrFArdLv+R+lbQ53BmIMLIry74bjKzX6iHAm8WYamJkhSSEabrWN5D97XnKObjQ==",
"dev": true,
"license": "BlueOak-1.0.0",
"dependencies": {
"@isaacs/brace-expansion": "^5.0.0"
},
"engines": {
"node": "20 || >=22"
},
"funding": {
"url": "https://github.com/sponsors/isaacs"
}
},
"node_modules/globals": {
"version": "14.0.0",
"resolved": "https://registry.npmjs.org/globals/-/globals-14.0.0.tgz",
@ -9035,16 +8968,19 @@
}
},
"node_modules/minimatch": {
"version": "3.1.2",
"resolved": "https://registry.npmjs.org/minimatch/-/minimatch-3.1.2.tgz",
"integrity": "sha512-J7p63hRiAjw1NDEww1W7i37+ByIrOWO5XQQAzZ3VOcL0PNybwpfmV/N05zFAzwQ9USyEcX6t3UO+K5aqBQOIHw==",
"version": "10.2.2",
"resolved": "https://registry.npmjs.org/minimatch/-/minimatch-10.2.2.tgz",
"integrity": "sha512-+G4CpNBxa5MprY+04MbgOw1v7So6n5JY166pFi9KfYwT78fxScCeSNQSNzp6dpPSW2rONOps6Ocam1wFhCgoVw==",
"dev": true,
"license": "ISC",
"license": "BlueOak-1.0.0",
"dependencies": {
"brace-expansion": "^1.1.7"
"brace-expansion": "^5.0.2"
},
"engines": {
"node": "*"
"node": "18 || 20 || >=22"
},
"funding": {
"url": "https://github.com/sponsors/isaacs"
}
},
"node_modules/minimist": {
@ -12004,32 +11940,6 @@
"node": ">=18"
}
},
"node_modules/test-exclude/node_modules/brace-expansion": {
"version": "2.0.2",
"resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.0.2.tgz",
"integrity": "sha512-Jt0vHyM+jmUBqojB7E1NIYadt0vI0Qxjxd2TErW94wDz+E2LAm5vKMXXwg6ZZBTHPuUlDgQHKXvjGBdfcF1ZDQ==",
"dev": true,
"license": "MIT",
"dependencies": {
"balanced-match": "^1.0.0"
}
},
"node_modules/test-exclude/node_modules/minimatch": {
"version": "9.0.5",
"resolved": "https://registry.npmjs.org/minimatch/-/minimatch-9.0.5.tgz",
"integrity": "sha512-G6T0ZX48xgozx7587koeX9Ys2NYy6Gmv//P89sEte9V9whIapMNF4idKxnW2QtCcLiTWlb/wfCabAtAFWhhBow==",
"dev": true,
"license": "ISC",
"dependencies": {
"brace-expansion": "^2.0.1"
},
"engines": {
"node": ">=16 || 14 >=14.17"
},
"funding": {
"url": "https://github.com/sponsors/isaacs"
}
},
"node_modules/thenify": {
"version": "3.3.1",
"resolved": "https://registry.npmjs.org/thenify/-/thenify-3.3.1.tgz",
@ -13085,21 +12995,6 @@
"type": "github",
"url": "https://github.com/sponsors/wooorm"
}
},
"node_modules/@next/swc-win32-ia32-msvc": {
"version": "14.2.33",
"resolved": "https://registry.npmjs.org/@next/swc-win32-ia32-msvc/-/swc-win32-ia32-msvc-14.2.33.tgz",
"integrity": "sha512-pc9LpGNKhJ0dXQhZ5QMmYxtARwwmWLpeocFmVG5Z0DzWq5Uf0izcI8tLc+qOpqxO1PWqZ5A7J1blrUIKrIFc7Q==",
"cpu": [
"ia32"
],
"optional": true,
"os": [
"win32"
],
"engines": {
"node": ">= 10"
}
}
}
}

51
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx
View file

@ -1,6 +1,6 @@
/* @vitest-environment jsdom */
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { render } from "@testing-library/react";
import { act, render } from "@testing-library/react";
import { beforeEach, describe, expect, it, vi } from "vitest";
import ModelsAndEndpointsView from "./ModelsAndEndpointsView";
@ -13,6 +13,8 @@ vi.mock("@/components/networking", () => ({
getCallbacksCall: vi.fn().mockResolvedValue({ router_settings: {} }),
setCallbacksCall: vi.fn().mockResolvedValue(undefined),
getUiSettings: vi.fn().mockResolvedValue({ values: {} }),
latestHealthChecksCall: vi.fn().mockResolvedValue({ latest_health_checks: {} }),
getModelCostMapReloadStatus: vi.fn().mockResolvedValue({}),
}));
vi.mock("@/app/(dashboard)/models-and-endpoints/components/ModelAnalyticsTab/ModelAnalyticsTab", () => ({
@ -27,6 +29,14 @@ vi.mock("@/components/add_model/AddModelForm", () => ({
default: () => null,
}));
const mockHealthCheckComponent = vi.fn((_props: { all_models_on_proxy?: string[] }) => null);
vi.mock("@/components/model_dashboard/HealthCheckComponent", () => ({
default: (props: { all_models_on_proxy?: string[] }) => {
mockHealthCheckComponent(props);
return null;
},
}));
vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({
default: () => ({
teams: [],
@ -104,4 +114,43 @@ describe("ModelsAndEndpointsView", () => {
);
expect(await findByText("Model Management", {}, { timeout: 10000 })).toBeInTheDocument();
}, 15000);
it("should pass model IDs (not model names) to HealthCheckComponent as all_models_on_proxy", async () => {
mockHealthCheckComponent.mockClear();
const modelDataWithIds = {
data: [
{ model_name: "gpt-4", model_info: { id: "deployment-id-1" } },
{ model_name: "gpt-4", model_info: { id: "deployment-id-2" } },
],
};
mockUseModelsInfo.mockReturnValue({
data: { data: modelDataWithIds.data },
isLoading: false,
refetch: vi.fn(),
});
const queryClient = createQueryClient();
const { getByRole } = render(
<QueryClientProvider client={queryClient}>
<ModelsAndEndpointsView
token="123"
modelData={{ data: modelDataWithIds.data }}
keys={[]}
setModelData={() => {}}
premiumUser={false}
teams={[]}
/>
</QueryClientProvider>,
);
const healthStatusTab = getByRole("tab", { name: "Health Status" });
await act(async () => {
healthStatusTab.click();
});
expect(mockHealthCheckComponent).toHaveBeenCalled();
const healthCheckProps = mockHealthCheckComponent.mock.calls[0][0];
expect(healthCheckProps.all_models_on_proxy).toEqual(["deployment-id-1", "deployment-id-2"]);
expect(healthCheckProps.all_models_on_proxy).not.toContain("gpt-4");
});
});

9
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx
View file

@ -98,6 +98,13 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, te
return modelDataResponse.data.map((model: any) => model.model_name);
}, [modelDataResponse?.data]);
const allModelIdsOnProxy = useMemo<string[]>(() => {
if (!modelDataResponse?.data) return [];
return modelDataResponse.data
.map((model: any) => model.model_info?.id)
.filter((id: string | undefined): id is string => Boolean(id));
}, [modelDataResponse?.data]);
const getProviderFromModel = (model: string) => {
if (modelCostMapData !== null && modelCostMapData !== undefined) {
if (typeof modelCostMapData == "object" && model in modelCostMapData) {
@ -397,7 +404,7 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, te
<HealthCheckComponent
accessToken={accessToken}
modelData={processedModelData}
all_models_on_proxy={allModelsOnProxy}
all_models_on_proxy={allModelIdsOnProxy}
getDisplayModelName={getDisplayModelName}
setSelectedModelId={setSelectedModelId}
teams={teams}

3
ui/litellm-dashboard/src/app/page.tsx
View file

@ -38,6 +38,7 @@ import Usage from "@/components/usage";
import UserDashboard from "@/components/user_dashboard";
import { AccessGroupsPage } from "@/components/AccessGroups/AccessGroupsPage";
import VectorStoreManagement from "@/components/vector_store_management";
import ToolPolicies from "@/components/ToolPolicies";
import SpendLogsTable from "@/components/view_logs";
import ViewUserDashboard from "@/components/view_users";
import { ThemeProvider } from "@/contexts/ThemeContext";
@ -548,6 +549,8 @@ function CreateKeyPageContent() {
<AccessGroupsPage />
) : page == "vector-stores" ? (
<VectorStoreManagement accessToken={accessToken} userRole={userRole} userID={userID} />
) : page == "tool-policies" ? (
<ToolPolicies accessToken={accessToken} userRole={userRole} />
) : page == "guardrails-monitor" ? (
<GuardrailsMonitorView accessToken={accessToken} />
) : page == "new_usage" ? (

415
ui/litellm-dashboard/src/components/ToolPolicies.tsx Normal file
View file

@ -0,0 +1,415 @@
"use client";
import React, { useCallback, useDeferredValue, useEffect, useState } from "react";
import { Select, Switch, Tooltip } from "antd";
import { Select, Tooltip } from "antd";
import {
Table,
TableHead,
TableHeaderCell,
TableBody,
TableRow,
TableCell,
} from "@tremor/react";
import { TimeCell } from "./view_logs/time_cell";
import { TableHeaderSortDropdown } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown";
import type { SortState } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown";
import FilterComponent, { FilterOption } from "./molecules/filter";
import { fetchToolsList, updateToolPolicy, ToolRow } from "./networking";
const POLICY_OPTIONS = [
{ value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" },
{ value: "blocked", label: "blocked", color: "#991b1b", bg: "#fee2e2", border: "#fca5a5" },
] as const;
type PolicyValue = "trusted" | "blocked";
const policyStyle = (p: string) =>
POLICY_OPTIONS.find((o) => o.value === p) ?? POLICY_OPTIONS[1];
type SortField = "tool_name" | "call_policy" | "team_id" | "key_alias" | "created_at" | "call_count";
interface FilterValues {
[key: string]: string;
}
interface ToolPoliciesProps {
accessToken: string | null;
userRole?: string;
}
const PolicySelect: React.FC<{
value: string;
toolName: string;
saving: boolean;
onChange: (toolName: string, policy: string) => void;
}> = ({ value, toolName, saving, onChange }) => {
const style = policyStyle(value);
return (
<Select
size="small"
value={value}
disabled={saving}
loading={saving}
onChange={(v) => onChange(toolName, v)}
onClick={(e) => e.stopPropagation()}
style={{
minWidth: 110,
fontWeight: 500,
}}
styles={{
selector: {
backgroundColor: style.bg,
borderColor: style.border,
color: style.color,
borderRadius: 999,
fontSize: 11,
fontWeight: 600,
paddingLeft: 8,
paddingRight: 4,
},
}}
popupMatchSelectWidth={false}
options={POLICY_OPTIONS.map((o) => ({
value: o.value,
label: (
<span
style={{
display: "inline-flex",
alignItems: "center",
gap: 6,
fontSize: 12,
fontWeight: 500,
color: o.color,
}}
>
<span
style={{
width: 8,
height: 8,
borderRadius: "50%",
backgroundColor: o.color,
display: "inline-block",
flexShrink: 0,
}}
/>
{o.label}
</span>
),
}))}
/>
);
};
export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
const [tools, setTools] = useState<ToolRow[]>([]);
const [loading, setLoading] = useState(true);
const [isFetching, setIsFetching] = useState(false);
const [error, setError] = useState<string | null>(null);
const [saving, setSaving] = useState<string | null>(null);
const [searchTerm, setSearchTerm] = useState("");
const [sortField, setSortField] = useState<SortField>("created_at");
const [sortOrder, setSortOrder] = useState<"asc" | "desc">("desc");
const [currentPage, setCurrentPage] = useState(1);
const [isLiveTail, setIsLiveTail] = useState(true);
const [activeFilters, setActiveFilters] = useState<FilterValues>({});
const pageSize = 50;
const isFetchingDeferred = useDeferredValue(isFetching);
const isButtonLoading = isFetching || isFetchingDeferred;
const load = useCallback(async () => {
if (!accessToken) return;
setIsFetching(true);
setError(null);
try {
const rows = await fetchToolsList(accessToken);
setTools(rows);
} catch (e: any) {
setError(e.message ?? "Failed to load tools");
} finally {
setIsFetching(false);
setLoading(false);
}
}, [accessToken]);
useEffect(() => { load(); }, [load]);
useEffect(() => {
if (!isLiveTail) return;
const id = setInterval(load, 15000);
return () => clearInterval(id);
}, [isLiveTail, load]);
const handlePolicyChange = async (toolName: string, newPolicy: string) => {
if (!accessToken) return;
setSaving(toolName);
try {
await updateToolPolicy(accessToken, toolName, newPolicy);
setTools((prev) =>
prev.map((t) => (t.tool_name === toolName ? { ...t, call_policy: newPolicy } : t))
);
} catch (e: any) {
alert(`Failed to update policy: ${e.message}`);
} finally {
setSaving(null);
}
};
const handleSortChange = (field: SortField, newState: SortState) => {
if (newState === false) {
setSortField("created_at");
setSortOrder("desc");
} else {
setSortField(field);
setSortOrder(newState);
}
setCurrentPage(1);
};
const handleApplyFilters = (filters: FilterValues) => {
setActiveFilters(filters);
setCurrentPage(1);
};
const handleResetFilters = () => {
setActiveFilters({});
setCurrentPage(1);
};
// Build unique team/key options from loaded data
const teamOptions = Array.from(new Set(tools.map((t) => t.team_id).filter(Boolean))).map(
(v) => ({ label: v as string, value: v as string })
);
const keyAliasOptions = Array.from(new Set(tools.map((t) => t.key_alias).filter(Boolean))).map(
(v) => ({ label: v as string, value: v as string })
);
const filterOptions: FilterOption[] = [
{
name: "Policy",
label: "Policy",
options: POLICY_OPTIONS.map((o) => ({ label: o.label, value: o.value })),
},
{
name: "Team Name",
label: "Team Name",
options: teamOptions,
},
{
name: "Key Name",
label: "Key Name",
options: keyAliasOptions,
},
];
const SortHeader = ({ label, field }: { label: string; field: SortField }) => (
<div className="flex items-center gap-1">
<span>{label}</span>
<TableHeaderSortDropdown
sortState={sortField === field ? sortOrder : false}
onSortChange={(s) => handleSortChange(field, s)}
/>
</div>
);
const filtered = tools.filter((t) => {
if (searchTerm) {
const q = searchTerm.toLowerCase();
const matchesSearch =
t.tool_name.toLowerCase().includes(q) ||
(t.team_id ?? "").toLowerCase().includes(q) ||
(t.key_alias ?? "").toLowerCase().includes(q) ||
(t.key_hash ?? "").toLowerCase().includes(q) ||
t.call_policy.toLowerCase().includes(q);
if (!matchesSearch) return false;
}
if (activeFilters["Policy"] && t.call_policy !== activeFilters["Policy"]) return false;
if (activeFilters["Team Name"] && t.team_id !== activeFilters["Team Name"]) return false;
if (activeFilters["Key Name"] && t.key_alias !== activeFilters["Key Name"]) return false;
return true;
});
const sorted = [...filtered].sort((a, b) => {
const av = (a as any)[sortField] ?? "";
const bv = (b as any)[sortField] ?? "";
if (av < bv) return sortOrder === "desc" ? 1 : -1;
if (av > bv) return sortOrder === "desc" ? -1 : 1;
return 0;
});
const totalPages = Math.max(1, Math.ceil(sorted.length / pageSize));
const paginated = sorted.slice((currentPage - 1) * pageSize, currentPage * pageSize);
return (
<div className="p-6 w-full">
<h1 className="text-2xl font-semibold text-gray-900 mb-6">Tool Policies</h1>
<div className="bg-white rounded-lg shadow w-full max-w-full box-border">
{/* Toolbar */}
<div className="border-b px-6 py-4 w-full max-w-full box-border">
<div className="flex flex-col md:flex-row items-start md:items-center justify-between space-y-4 md:space-y-0 w-full max-w-full box-border">
<div className="flex flex-wrap items-center gap-3">
<div className="relative w-64">
<input
type="text"
placeholder="Search by Tool Name"
className="w-full px-3 py-2 pl-8 border rounded-md text-sm focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-blue-500"
value={searchTerm}
onChange={(e) => { setSearchTerm(e.target.value); setCurrentPage(1); }}
/>
<svg className="absolute left-2.5 top-2.5 h-4 w-4 text-gray-500" fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M21 21l-6-6m2-5a7 7 0 11-14 0 7 7 0 0114 0z" />
</svg>
</div>
<div className="flex items-center gap-2">
<span className="text-sm font-medium text-gray-900">Live Tail</span>
<Switch color="green" checked={isLiveTail} onChange={setIsLiveTail} />
</div>
<button
onClick={load}
disabled={isButtonLoading}
className="flex items-center gap-1.5 px-3 py-2 text-sm border rounded-md hover:bg-gray-50 disabled:opacity-60"
>
<svg className={`w-4 h-4 ${isButtonLoading ? "animate-spin" : ""}`} fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M4 4v5h.582m15.356 2A8.001 8.001 0 004.582 9m0 0H9m11 11v-5h-.581m0 0a8.003 8.003 0 01-15.357-2m15.357 2H15" />
</svg>
{isButtonLoading ? "Fetching" : "Fetch"}
</button>
</div>
<div className="flex items-center gap-4 text-sm text-gray-600 whitespace-nowrap">
<span>
Showing {filtered.length === 0 ? 0 : (currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, filtered.length)} of {filtered.length} results
</span>
<span>Page {currentPage} of {totalPages}</span>
<div className="flex gap-1">
<button onClick={() => setCurrentPage((p) => Math.max(1, p - 1))} disabled={currentPage === 1}
className="px-3 py-1.5 border rounded-md text-sm hover:bg-gray-50 disabled:opacity-40">Previous</button>
<button onClick={() => setCurrentPage((p) => Math.min(totalPages, p + 1))} disabled={currentPage === totalPages}
className="px-3 py-1.5 border rounded-md text-sm hover:bg-gray-50 disabled:opacity-40">Next</button>
</div>
</div>
</div>
{/* Filter row */}
<div className="mt-3">
<FilterComponent
options={filterOptions}
onApplyFilters={handleApplyFilters}
onResetFilters={handleResetFilters}
buttonLabel="Filters"
/>
</div>
</div>
{/* Auto-refresh banner */}
{isLiveTail && (
<div className="bg-green-50 border-b border-green-100 px-6 py-2 flex items-center justify-between">
<span className="text-sm text-green-700">Auto-refreshing every 15 seconds</span>
<button onClick={() => setIsLiveTail(false)} className="text-xs text-green-600 underline">Stop</button>
</div>
)}
{error && (
<div className="mx-6 mt-4 p-3 bg-red-50 border border-red-200 rounded text-sm text-red-700">{error}</div>
)}
{/* Table */}
<Table className="[&_td]:py-0.5 [&_th]:py-1 w-full">
<TableHead>
<TableRow>
<TableHeaderCell className="py-1 h-8"><SortHeader label="Discovered" field="created_at" /></TableHeaderCell>
<TableHeaderCell className="py-1 h-8"><SortHeader label="Tool Name" field="tool_name" /></TableHeaderCell>
<TableHeaderCell className="py-1 h-8"><SortHeader label="Policy" field="call_policy" /></TableHeaderCell>
<TableHeaderCell className="py-1 h-8"><SortHeader label="# Calls" field="call_count" /></TableHeaderCell>
<TableHeaderCell className="py-1 h-8"><SortHeader label="Team Name" field="team_id" /></TableHeaderCell>
<TableHeaderCell className="py-1 h-8">Key Hash</TableHeaderCell>
<TableHeaderCell className="py-1 h-8"><SortHeader label="Key Name" field="key_alias" /></TableHeaderCell>
<TableHeaderCell className="py-1 h-8">Origin</TableHeaderCell>
</TableRow>
</TableHead>
<TableBody>
{loading ? (
<TableRow>
<TableCell colSpan={8} className="h-8 text-center text-gray-500">Loading tools…</TableCell>
</TableRow>
) : paginated.length === 0 ? (
<TableRow>
<TableCell colSpan={8} className="h-8 text-center text-gray-500">
No tools discovered yet. Make a chat completion that returns tool_calls to start auto-discovery.
</TableCell>
</TableRow>
) : (
paginated.map((tool) => (
<TableRow key={tool.tool_id} className="h-8 hover:bg-gray-50">
<TableCell className="py-0.5 max-h-8 overflow-hidden whitespace-nowrap">
<TimeCell utcTime={tool.created_at ?? ""} />
</TableCell>
<TableCell className="py-0.5 max-h-8 overflow-hidden">
<Tooltip title={tool.tool_name}>
<span className="font-mono text-xs max-w-[20ch] truncate block font-medium">
{tool.tool_name}
</span>
</Tooltip>
</TableCell>
<TableCell className="py-0.5 max-h-8">
<PolicySelect
value={tool.call_policy}
toolName={tool.tool_name}
saving={saving === tool.tool_name}
onChange={handlePolicyChange}
/>
</TableCell>
<TableCell className="py-0.5 max-h-8 text-right tabular-nums text-sm font-mono text-gray-700">
{(tool.call_count ?? 0).toLocaleString()}
</TableCell>
<TableCell className="py-0.5 max-h-8 overflow-hidden whitespace-nowrap">
<Tooltip title={tool.team_id ?? "-"}>
<span className="max-w-[15ch] truncate block">{tool.team_id ?? "-"}</span>
</Tooltip>
</TableCell>
<TableCell className="py-0.5 max-h-8 overflow-hidden whitespace-nowrap">
<Tooltip title={tool.key_hash ?? "-"}>
<span className="font-mono max-w-[15ch] truncate block text-blue-600">
{tool.key_hash ?? "-"}
</span>
</Tooltip>
</TableCell>
<TableCell className="py-0.5 max-h-8 overflow-hidden whitespace-nowrap">
<Tooltip title={tool.key_alias ?? "-"}>
<span className="max-w-[15ch] truncate block">{tool.key_alias ?? "-"}</span>
</Tooltip>
</TableCell>
<TableCell className="py-0.5 max-h-8 overflow-hidden whitespace-nowrap">
<Tooltip title={tool.origin ?? "-"}>
<span className="max-w-[15ch] truncate block">{tool.origin ?? "-"}</span>
</Tooltip>
</TableCell>
</TableRow>
))
)}
</TableBody>
</Table>
{/* Bottom pagination (only when > 1 page) */}
{totalPages > 1 && (
<div className="border-t px-6 py-3 flex items-center justify-between text-sm text-gray-600">
<span>Showing {(currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, sorted.length)} of {sorted.length}</span>
<div className="flex gap-1">
<button onClick={() => setCurrentPage((p) => Math.max(1, p - 1))} disabled={currentPage === 1}
className="px-3 py-1.5 border rounded-md hover:bg-gray-50 disabled:opacity-40">Previous</button>
<button onClick={() => setCurrentPage((p) => Math.min(totalPages, p + 1))} disabled={currentPage === totalPages}
className="px-3 py-1.5 border rounded-md hover:bg-gray-50 disabled:opacity-40">Next</button>
</div>
</div>
)}
</div>
</div>
);
};
export default ToolPolicies;

85
ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx Normal file
View file

@ -0,0 +1,85 @@
import { screen } from "@testing-library/react";
import { beforeAll, describe, expect, it, vi } from "vitest";
import { renderWithProviders } from "../../../../tests/test-utils";
import UsageAIChatPanel from "./UsageAIChatPanel";
beforeAll(() => {
if (typeof window !== "undefined" && !window.ResizeObserver) {
window.ResizeObserver = class ResizeObserver {
observe() {}
unobserve() {}
disconnect() {}
} as any;
}
});
vi.mock("../../networking", () => ({
modelHubCall: vi.fn().mockResolvedValue({
data: [
{ model_group: "gpt-4" },
{ model_group: "claude-3-opus" },
],
}),
usageAiChatStream: vi.fn(),
}));
const defaultProps = {
open: true,
onClose: vi.fn(),
accessToken: "test-token",
};
describe("UsageAIChatPanel", () => {
it("should render the panel when open", () => {
renderWithProviders(<UsageAIChatPanel {...defaultProps} />);
expect(screen.getByText("Ask AI")).toBeInTheDocument();
expect(
screen.getByText("Ask about your spend, models, keys, and trends")
).toBeInTheDocument();
});
it("should render model selector", () => {
renderWithProviders(<UsageAIChatPanel {...defaultProps} />);
expect(screen.getByText("Select a model (optional, defaults to gpt-4o-mini)")).toBeInTheDocument();
});
it("should render empty state message when no conversation", () => {
renderWithProviders(<UsageAIChatPanel {...defaultProps} />);
expect(screen.getByText("Ask a question about your usage")).toBeInTheDocument();
});
it("should render the send button", () => {
renderWithProviders(<UsageAIChatPanel {...defaultProps} />);
expect(screen.getByText("Send")).toBeInTheDocument();
});
it("should render input placeholder", () => {
renderWithProviders(<UsageAIChatPanel {...defaultProps} />);
expect(screen.getByPlaceholderText("Ask about your usage...")).toBeInTheDocument();
});
it("should render clear chat button", () => {
renderWithProviders(<UsageAIChatPanel {...defaultProps} />);
expect(screen.getByText("Clear chat")).toBeInTheDocument();
});
it("should have the panel element even when closed (just off-screen)", () => {
renderWithProviders(<UsageAIChatPanel {...defaultProps} open={false} />);
expect(screen.getByTestId("usage-ai-chat-panel")).toBeInTheDocument();
expect(screen.getByTestId("usage-ai-chat-panel")).toHaveClass("translate-x-full");
});
it("should not have translate-x-full class when open", () => {
renderWithProviders(<UsageAIChatPanel {...defaultProps} open={true} />);
expect(screen.getByTestId("usage-ai-chat-panel")).not.toHaveClass("translate-x-full");
expect(screen.getByTestId("usage-ai-chat-panel")).toHaveClass("translate-x-0");
});
});

402
ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx Normal file
View file

@ -0,0 +1,402 @@
import React, { useEffect, useRef, useState } from "react";
import { Button, Select, Input, Spin } from "antd";
import ReactMarkdown from "react-markdown";
import { modelHubCall, usageAiChatStream, UsageAiToolCallEvent } from "../../networking";
const { TextArea } = Input;
interface ToolCallStep {
tool_name: string;
tool_label: string;
arguments: Record<string, string>;
status: "running" | "complete" | "error";
error?: string;
}
interface ChatMessage {
role: "user" | "assistant";
content: string;
toolCalls?: ToolCallStep[];
}
interface UsageAIChatPanelProps {
open: boolean;
onClose: () => void;
accessToken: string | null;
}
const TOOL_ICONS: Record<string, string> = {
get_usage_data: "📊",
get_team_usage_data: "👥",
get_tag_usage_data: "🏷️",
};
const ToolCallDisplay: React.FC<{ step: ToolCallStep }> = ({ step }) => {
const icon = TOOL_ICONS[step.tool_name] || "🔧";
const args = step.arguments;
const dateRange = args.start_date && args.end_date
? `${args.start_date} → ${args.end_date}`
: "";
const filter = args.team_ids || args.tags || args.user_id || "";
return (
<div className="flex items-start gap-2 px-3 py-2 rounded-lg bg-gray-100 border border-gray-200 text-xs">
<span className="flex-shrink-0 mt-0.5">
{step.status === "running" ? (
<Spin size="small" />
) : step.status === "error" ? (
<span className="text-red-500">✗</span>
) : (
<span className="text-green-600">✓</span>
)}
</span>
<div className="min-w-0">
<div className="font-medium text-gray-700">
{icon} {step.tool_label}
</div>
{dateRange && (
<div className="text-gray-500 mt-0.5">{dateRange}</div>
)}
{filter && (
<div className="text-gray-500 mt-0.5">Filter: {filter}</div>
)}
{step.status === "error" && step.error && (
<div className="text-red-600 mt-0.5">{step.error}</div>
)}
</div>
</div>
);
};
const MarkdownContent: React.FC<{ content: string }> = ({ content }) => (
<ReactMarkdown
components={{
p: ({ children }) => <p className="mb-2 last:mb-0">{children}</p>,
strong: ({ children }) => <strong className="font-semibold">{children}</strong>,
ul: ({ children }) => <ul className="list-disc pl-4 mb-2 space-y-0.5">{children}</ul>,
ol: ({ children }) => <ol className="list-decimal pl-4 mb-2 space-y-0.5">{children}</ol>,
li: ({ children }) => <li>{children}</li>,
h1: ({ children }) => <h4 className="font-semibold text-sm mt-2 mb-1">{children}</h4>,
h2: ({ children }) => <h4 className="font-semibold text-sm mt-2 mb-1">{children}</h4>,
h3: ({ children }) => <h4 className="font-semibold text-sm mt-2 mb-1">{children}</h4>,
code: ({ children, className }) => {
const isBlock = className?.includes("language-");
return isBlock ? (
<pre className="bg-gray-100 rounded p-2 my-1 overflow-x-auto text-xs">
<code>{children}</code>
</pre>
) : (
<code className="px-1 py-0.5 rounded bg-gray-100 text-xs font-mono">{children}</code>
);
},
table: ({ children }) => (
<div className="overflow-x-auto my-2">
<table className="text-xs border-collapse w-full">{children}</table>
</div>
),
th: ({ children }) => <th className="border border-gray-200 px-2 py-1 bg-gray-50 font-medium text-left">{children}</th>,
td: ({ children }) => <td className="border border-gray-200 px-2 py-1">{children}</td>,
}}
>
{content}
</ReactMarkdown>
);
const UsageAIChatPanel: React.FC<UsageAIChatPanelProps> = ({
open,
onClose,
accessToken,
}) => {
const [messages, setMessages] = useState<ChatMessage[]>([]);
const [inputText, setInputText] = useState("");
const [isLoading, setIsLoading] = useState(false);
const [selectedModel, setSelectedModel] = useState<string | undefined>(undefined);
const [availableModels, setAvailableModels] = useState<string[]>([]);
const [isLoadingModels, setIsLoadingModels] = useState(false);
const [streamingContent, setStreamingContent] = useState("");
const [statusMessage, setStatusMessage] = useState<string | null>(null);
const [activeToolCalls, setActiveToolCalls] = useState<ToolCallStep[]>([]);
const messagesEndRef = useRef<HTMLDivElement>(null);
const abortControllerRef = useRef<AbortController | null>(null);
useEffect(() => {
if (open && availableModels.length === 0) {
loadModels();
}
}, [open]);
useEffect(() => {
if (typeof messagesEndRef.current?.scrollIntoView === "function") {
messagesEndRef.current.scrollIntoView({ behavior: "smooth" });
}
}, [messages, streamingContent, activeToolCalls, statusMessage]);
const loadModels = async () => {
if (!accessToken) return;
setIsLoadingModels(true);
try {
const fetchedModels = await modelHubCall(accessToken);
if (fetchedModels?.data?.length > 0) {
const models = fetchedModels.data
.map((item: any) => item.model_group as string)
.sort();
setAvailableModels(models);
}
} catch (error) {
console.error("Failed to load models:", error);
} finally {
setIsLoadingModels(false);
}
};
const handleSend = async () => {
if (!accessToken || !inputText.trim() || isLoading) return;
const userMessage: ChatMessage = { role: "user", content: inputText.trim() };
const updatedMessages = [...messages, userMessage];
setMessages(updatedMessages);
setInputText("");
setIsLoading(true);
setStreamingContent("");
setStatusMessage(null);
setActiveToolCalls([]);
const abortController = new AbortController();
abortControllerRef.current = abortController;
let accumulated = "";
const toolCalls: ToolCallStep[] = [];
try {
await usageAiChatStream(
accessToken,
updatedMessages.slice(-20).map((m) => ({ role: m.role, content: m.content })),
selectedModel || "",
(content: string) => {
setStatusMessage(null);
accumulated += content;
setStreamingContent(accumulated);
},
() => {
setStatusMessage(null);
setActiveToolCalls([]);
setMessages((prev) => [
...prev,
{ role: "assistant", content: accumulated, toolCalls: toolCalls.length > 0 ? [...toolCalls] : undefined },
]);
setStreamingContent("");
},
(errorMsg: string) => {
setStatusMessage(null);
setActiveToolCalls([]);
setMessages((prev) => [
...prev,
{ role: "assistant", content: `Error: ${errorMsg}` },
]);
setStreamingContent("");
},
(status: string) => {
setStatusMessage(status);
},
(event: UsageAiToolCallEvent) => {
const idx = toolCalls.findIndex((tc) => tc.tool_name === event.tool_name);
if (idx >= 0) {
toolCalls[idx] = { ...event };
} else {
toolCalls.push({ ...event });
}
setActiveToolCalls([...toolCalls]);
},
abortController.signal,
);
} catch (error: any) {
if (error?.name === "AbortError" || abortController.signal.aborted) {
return;
}
const errorMsg = error?.message || "Failed to get response. Please try again.";
setMessages((prev) => [
...prev,
{ role: "assistant", content: `Error: ${errorMsg}` },
]);
setStreamingContent("");
} finally {
setIsLoading(false);
abortControllerRef.current = null;
}
};
const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
if (e.key === "Enter" && !e.shiftKey) {
e.preventDefault();
handleSend();
}
};
const handleClose = () => {
if (abortControllerRef.current) {
abortControllerRef.current.abort();
}
onClose();
};
const handleClear = () => {
setMessages([]);
setStreamingContent("");
setActiveToolCalls([]);
setStatusMessage(null);
};
return (
<div
data-testid="usage-ai-chat-panel"
className={`fixed top-0 right-0 h-full bg-white border-l border-gray-200 shadow-2xl z-50 flex flex-col transition-transform duration-300 ease-in-out ${
open ? "translate-x-0" : "translate-x-full"
}`}
style={{ width: 420 }}
>
{/* Header */}
<div className="px-5 pt-5 pb-3 border-b border-gray-100 flex-shrink-0">
<div className="flex items-center justify-between mb-1">
<div className="flex items-center gap-2">
<svg className="w-5 h-5 text-blue-600" viewBox="0 0 16 16" fill="currentColor">
<path d="M8 1l1.5 3.5L13 6l-3.5 1.5L8 11 6.5 7.5 3 6l3.5-1.5L8 1zm4 7l.75 1.75L14.5 10.5l-1.75.75L12 13l-.75-1.75L9.5 10.5l1.75-.75L12 8zM4 9l.75 1.75L6.5 11.5l-1.75.75L4 14l-.75-1.75L1.5 11.5l1.75-.75L4 9z" />
</svg>
<h3 className="text-base font-semibold text-gray-900">Ask AI</h3>
</div>
<button
onClick={handleClose}
className="text-gray-400 hover:text-gray-600 transition-colors p-1 rounded-md hover:bg-gray-100"
>
<svg className="w-5 h-5" fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M6 18L18 6M6 6l12 12" />
</svg>
</button>
</div>
<p className="text-xs text-gray-500">
Ask about your spend, models, keys, and trends
</p>
</div>
{/* Model selector */}
<div className="px-5 py-3 border-b border-gray-100 flex-shrink-0">
<Select
placeholder="Select a model (optional, defaults to gpt-4o-mini)"
value={selectedModel}
onChange={(value) => setSelectedModel(value)}
loading={isLoadingModels}
showSearch
allowClear
size="small"
className="w-full"
options={availableModels.map((m) => ({ label: m, value: m }))}
filterOption={(input, option) =>
(option?.label ?? "").toLowerCase().includes(input.toLowerCase())
}
/>
</div>
{/* Chat messages */}
<div className="flex-1 overflow-y-auto p-4 space-y-3 bg-gray-50">
{messages.length === 0 && !streamingContent && !isLoading && (
<div className="flex flex-col items-center justify-center h-full text-gray-400">
<svg className="w-8 h-8 mb-2" fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M8 10h.01M12 10h.01M16 10h.01M9 16H5a2 2 0 01-2-2V6a2 2 0 012-2h14a2 2 0 012 2v8a2 2 0 01-2 2h-5l-5 5v-5z" />
</svg>
<p className="text-sm font-medium">Ask a question about your usage</p>
<p className="text-xs mt-1">e.g. &quot;Which model costs me the most?&quot;</p>
</div>
)}
{messages.map((msg, idx) => (
<div key={idx}>
{msg.role === "user" ? (
<div className="flex justify-end">
<div className="max-w-[88%] rounded-xl px-3.5 py-2 text-sm leading-relaxed bg-blue-600 text-white">
{msg.content}
</div>
</div>
) : (
<div className="space-y-2">
{/* Tool calls for this message */}
{msg.toolCalls && msg.toolCalls.length > 0 && (
<div className="space-y-1.5">
{msg.toolCalls.map((tc, tcIdx) => (
<ToolCallDisplay key={tcIdx} step={tc} />
))}
</div>
)}
{/* Response */}
<div className="max-w-[95%] rounded-xl px-3.5 py-2.5 text-sm leading-relaxed bg-white border border-gray-200 text-gray-800">
<MarkdownContent content={msg.content} />
</div>
</div>
)}
</div>
))}
{/* Active tool calls (in-progress) */}
{isLoading && activeToolCalls.length > 0 && (
<div className="space-y-1.5">
{activeToolCalls.map((tc, idx) => (
<ToolCallDisplay key={idx} step={tc} />
))}
</div>
)}
{/* Status / spinner */}
{isLoading && !streamingContent && (
<div className="flex items-center gap-2 px-3 py-2 text-xs text-gray-500">
<Spin size="small" />
<span className="italic">{statusMessage || "Thinking..."}</span>
</div>
)}
{/* Streaming response */}
{streamingContent && (
<div className="max-w-[95%] rounded-xl px-3.5 py-2.5 text-sm leading-relaxed bg-white border border-gray-200 text-gray-800">
<MarkdownContent content={streamingContent} />
</div>
)}
<div ref={messagesEndRef} />
</div>
{/* Input area */}
<div className="px-4 py-3 border-t border-gray-200 bg-white flex-shrink-0">
<div className="flex gap-2">
<TextArea
value={inputText}
onChange={(e) => setInputText(e.target.value)}
onKeyDown={handleKeyDown}
placeholder="Ask about your usage..."
autoSize={{ minRows: 1, maxRows: 3 }}
className="flex-1"
disabled={isLoading}
/>
<Button
type="primary"
onClick={handleSend}
disabled={!inputText.trim() || isLoading}
loading={isLoading}
>
Send
</Button>
</div>
<div className="flex justify-between items-center mt-2">
<button
onClick={handleClear}
className="text-xs text-gray-400 hover:text-gray-600 transition-colors"
disabled={messages.length === 0}
>
Clear chat
</button>
<span className="text-xs text-gray-400">
Enter to send
</span>
</div>
</div>
</div>
);
};
export default UsageAIChatPanel;

26
ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx
View file

@ -100,6 +100,10 @@ vi.mock("../../EntityUsageExport", () => ({
default: () => <div>Entity Usage Export Modal</div>,
}));
vi.mock("./UsageAIChatPanel", () => ({
default: () => <div data-testid="usage-ai-chat-panel">Usage AI Chat Panel</div>,
}));
vi.mock("@/app/(dashboard)/hooks/customers/useCustomers", () => ({
useCustomers: vi.fn(),
}));
@ -990,6 +994,28 @@ describe("UsagePage", () => {
});
});
describe("Ask AI button", () => {
it("should render Ask AI button in global view", async () => {
renderWithProviders(<UsagePage {...defaultProps} />);
await waitFor(() => {
expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
});
expect(screen.getByText("Ask AI")).toBeInTheDocument();
});
it("should render AI chat panel component", async () => {
renderWithProviders(<UsagePage {...defaultProps} />);
await waitFor(() => {
expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
});
expect(screen.getByTestId("usage-ai-chat-panel")).toBeInTheDocument();
});
});
describe("model view toggle", () => {
it("should show Public Model Name view by default", async () => {
renderWithProviders(<UsagePage {...defaultProps} />);

51
ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx
View file

@ -50,6 +50,7 @@ import EntityUsage, { EntityList } from "./EntityUsage/EntityUsage";
import SpendByProvider from "./EntityUsage/SpendByProvider";
import TopKeyView from "./EntityUsage/TopKeyView";
import { UsageOption, UsageViewSelect } from "./UsageViewSelect/UsageViewSelect";
import UsageAIChatPanel from "./UsageAIChatPanel";
interface UsagePageProps {
teams: Team[];
@ -142,6 +143,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
const [modelViewType, setModelViewType] = useState<"groups" | "individual">("groups");
const [isCloudZeroModalOpen, setIsCloudZeroModalOpen] = useState(false);
const [isGlobalExportModalOpen, setIsGlobalExportModalOpen] = useState(false);
const [isAiChatOpen, setIsAiChatOpen] = useState(false);
const [usageView, setUsageView] = useState<UsageOption>("global");
const [showCredentialBanner, setShowCredentialBanner] = useState(true);
const [topKeysLimit, setTopKeysLimit] = useState<number>(5);
@ -505,21 +507,33 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
<Tab>MCP Server Activity</Tab>
<Tab>Endpoint Activity</Tab>
</TabList>
<Button
onClick={() => setIsGlobalExportModalOpen(true)}
icon={() => (
<svg className="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path
strokeLinecap="round"
strokeLinejoin="round"
strokeWidth={2}
d="M4 16v1a3 3 0 003 3h10a3 3 0 003-3v-1m-4-4l-4 4m0 0l-4-4m4 4V4"
/>
</svg>
)}
>
Export Data
</Button>
<div className="flex items-center gap-2">
<Button
onClick={() => setIsAiChatOpen(true)}
icon={() => (
<svg className="w-4 h-4" viewBox="0 0 16 16" fill="currentColor">
<path d="M8 1l1.5 3.5L13 6l-3.5 1.5L8 11 6.5 7.5 3 6l3.5-1.5L8 1zm4 7l.75 1.75L14.5 10.5l-1.75.75L12 13l-.75-1.75L9.5 10.5l1.75-.75L12 8zM4 9l.75 1.75L6.5 11.5l-1.75.75L4 14l-.75-1.75L1.5 11.5l1.75-.75L4 9z" />
</svg>
)}
>
Ask AI
</Button>
<Button
onClick={() => setIsGlobalExportModalOpen(true)}
icon={() => (
<svg className="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path
strokeLinecap="round"
strokeLinejoin="round"
strokeWidth={2}
d="M4 16v1a3 3 0 003 3h10a3 3 0 003-3v-1m-4-4l-4 4m0 0l-4-4m4 4V4"
/>
</svg>
)}
>
Export Data
</Button>
</div>
</div>
<TabPanels>
{/* Cost Panel */}
@ -925,6 +939,13 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
selectedFilters={[]}
customTitle="Export Usage Data"
/>
{/* AI Chat Panel */}
<UsageAIChatPanel
open={isAiChatOpen}
onClose={() => setIsAiChatOpen(false)}
accessToken={accessToken}
/>
</div>
);
};

4
ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx
View file

@ -453,8 +453,8 @@ it("should open KeyInfoView when clicking on a key ID button", async () => {
// Wait for KeyInfoView to appear - check for unique elements that only exist in KeyInfoView
await waitFor(() => {
expect(screen.getByText("Back to Keys")).toBeInTheDocument();
// KeyInfoView shows "Created:" or "Updated:" which is unique to it
expect(screen.getByText(/Created:|Updated:/)).toBeInTheDocument();
// KeyInfoHeader shows "Created At" metadata label
expect(screen.getByText("Created At")).toBeInTheDocument();
});
// Verify that table-specific elements are no longer visible

625
ui/litellm-dashboard/src/components/agents/add_agent_form.tsx
View file

@ -1,11 +1,24 @@
import React, { useState, useEffect } from "react";
import { Modal, Form, message, Select, Input } from "antd";
import { Modal, Form, message, Select, Input, Steps, Radio, Tag, Divider } from "antd";
import { Button } from "@tremor/react";
import { createAgentCall, getAgentCreateMetadata, AgentCreateInfo } from "../networking";
import { CheckCircleFilled, KeyOutlined, RobotOutlined, AppstoreOutlined } from "@ant-design/icons";
import CreatedKeyDisplay from "../shared/CreatedKeyDisplay";
import {
createAgentCall,
getAgentCreateMetadata,
keyCreateForAgentCall,
keyListCall,
keyUpdateCall,
AgentCreateInfo,
} from "../networking";
import AgentFormFields from "./agent_form_fields";
import DynamicAgentFormFields, { buildDynamicAgentData } from "./dynamic_agent_form_fields";
import { getDefaultFormValues, buildAgentDataFromForm } from "./agent_config";
const { Step } = Steps;
const CUSTOM_AGENT_TYPE = "custom";
interface AddAgentFormProps {
visible: boolean;
onClose: () => void;
@ -20,11 +33,25 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
onSuccess,
}) => {
const [form] = Form.useForm();
const [currentStep, setCurrentStep] = useState(0);
const [isSubmitting, setIsSubmitting] = useState(false);
const [agentType, setAgentType] = useState<string>("a2a");
const [agentTypeMetadata, setAgentTypeMetadata] = useState<AgentCreateInfo[]>([]);
const [loadingMetadata, setLoadingMetadata] = useState(false);
// Step 1: key assignment state
const [keyAssignOption, setKeyAssignOption] = useState<"create_new" | "existing_key" | "skip">("create_new");
const [newKeyName, setNewKeyName] = useState<string>("");
const [newKeyModels, setNewKeyModels] = useState<string[]>([]);
const [existingKeys, setExistingKeys] = useState<any[]>([]);
const [selectedExistingKey, setSelectedExistingKey] = useState<string | null>(null);
const [loadingKeys, setLoadingKeys] = useState(false);
// Step 2: results
const [createdAgentName, setCreatedAgentName] = useState<string>("");
const [createdKeyValue, setCreatedKeyValue] = useState<string | null>(null);
const [assignedKeyAlias, setAssignedKeyAlias] = useState<string | null>(null);
// Fetch agent type metadata on mount
useEffect(() => {
const fetchMetadata = async () => {
@ -41,11 +68,87 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
fetchMetadata();
}, []);
// Fetch existing keys when assign key step becomes active
useEffect(() => {
if (currentStep === 1 && accessToken && existingKeys.length === 0) {
const fetchKeys = async () => {
setLoadingKeys(true);
try {
const result = await keyListCall(accessToken, null, null, null, null, null, 1, 100);
setExistingKeys(result?.keys || []);
} catch (error) {
console.error("Error fetching keys:", error);
} finally {
setLoadingKeys(false);
}
};
fetchKeys();
}
}, [currentStep, accessToken]);
const selectedAgentTypeInfo = agentTypeMetadata.find(
(info) => info.agent_type === agentType
);
const handleSubmit = async (values: any) => {
const handleNext = async () => {
try {
if (currentStep === 0) {
await form.validateFields(["agent_name"]);
const agentName = form.getFieldValue("agent_name");
if (agentName && !newKeyName) {
setNewKeyName(`${agentName}-key`);
}
}
setCurrentStep((s) => s + 1);
} catch {
// validation failed — stay on current step
}
};
const handleBack = () => {
setCurrentStep((s) => Math.max(0, s - 1));
};
const buildAgentData = (values: any) => {
if (agentType === CUSTOM_AGENT_TYPE) {
return {
agent_name: values.agent_name,
agent_card_params: {
protocolVersion: "1.0",
name: values.agent_name,
description: values.description || "",
url: "",
version: "1.0.0",
defaultInputModes: ["text"],
defaultOutputModes: ["text"],
capabilities: { streaming: false },
skills: [],
},
};
} else if (agentType === "a2a") {
return buildAgentDataFromForm(values);
} else if (selectedAgentTypeInfo?.use_a2a_form_fields) {
const agentData = buildAgentDataFromForm(values);
if (selectedAgentTypeInfo.litellm_params_template) {
agentData.litellm_params = {
...agentData.litellm_params,
...selectedAgentTypeInfo.litellm_params_template,
};
}
for (const field of selectedAgentTypeInfo.credential_fields) {
const value = values[field.key];
if (value && field.include_in_litellm_params !== false) {
agentData.litellm_params[field.key] = value;
}
}
return agentData;
} else if (selectedAgentTypeInfo) {
return buildDynamicAgentData(values, selectedAgentTypeInfo);
}
return null;
};
const handleCreateAgent = async () => {
if (!accessToken) {
message.error("No access token available");
return;
@ -53,40 +156,46 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
setIsSubmitting(true);
try {
let agentData: any;
if (agentType === "a2a") {
agentData = buildAgentDataFromForm(values);
} else if (selectedAgentTypeInfo?.use_a2a_form_fields) {
// A2A-compatible agents use the standard A2A form builder
// but need to add litellm_params from the agent type config
agentData = buildAgentDataFromForm(values);
// Merge litellm_params_template
if (selectedAgentTypeInfo.litellm_params_template) {
agentData.litellm_params = {
...agentData.litellm_params,
...selectedAgentTypeInfo.litellm_params_template,
};
}
// Add credential fields to litellm_params
for (const field of selectedAgentTypeInfo.credential_fields) {
const value = values[field.key];
if (value && field.include_in_litellm_params !== false) {
agentData.litellm_params[field.key] = value;
}
}
} else if (selectedAgentTypeInfo) {
agentData = buildDynamicAgentData(values, selectedAgentTypeInfo);
// getFieldsValue(true) returns ALL preserved values including fields from
// unmounted steps; merge with any currently-mounted validated fields.
await form.validateFields();
const values = { ...form.getFieldsValue(true) };
const agentData = buildAgentData(values);
if (!agentData) {
message.error("Failed to build agent data");
setIsSubmitting(false);
return;
}
await createAgentCall(accessToken, agentData);
message.success("Agent created successfully");
form.resetFields();
setAgentType("a2a");
const agentResponse = await createAgentCall(accessToken, agentData);
const agentId: string = agentResponse.agent_id;
const agentName: string = agentResponse.agent_name || values.agent_name || agentId;
setCreatedAgentName(agentName);
if (keyAssignOption === "create_new" && newKeyName) {
const keyResponse = await keyCreateForAgentCall(
accessToken,
agentId,
newKeyName,
newKeyModels,
);
setCreatedKeyValue(keyResponse.key || null);
} else if (keyAssignOption === "existing_key") {
if (!selectedExistingKey) {
message.error("Please select an existing key to assign");
setIsSubmitting(false);
return;
}
await keyUpdateCall(accessToken, {
key: selectedExistingKey,
agent_id: agentId,
});
const keyInfo = existingKeys.find((k) => k.token === selectedExistingKey);
setAssignedKeyAlias(keyInfo?.key_alias || selectedExistingKey.slice(0, 12) + "…");
}
setCurrentStep(2);
onSuccess();
onClose();
} catch (error) {
console.error("Error creating agent:", error);
message.error("Failed to create agent");
@ -95,9 +204,17 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
}
};
const handleCancel = () => {
const handleClose = () => {
form.resetFields();
setAgentType("a2a");
setCurrentStep(0);
setKeyAssignOption("create_new");
setNewKeyName("");
setNewKeyModels([]);
setSelectedExistingKey(null);
setCreatedAgentName("");
setCreatedKeyValue(null);
setAssignedKeyAlias(null);
onClose();
};
@ -106,25 +223,308 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
form.resetFields();
};
// Get the logo for the selected agent type for the header
const selectedLogo = selectedAgentTypeInfo?.logo_url || agentTypeMetadata.find(a => a.agent_type === "a2a")?.logo_url;
const isCustomAgent = agentType === CUSTOM_AGENT_TYPE;
const selectedLogo = isCustomAgent
? null
: selectedAgentTypeInfo?.logo_url ||
agentTypeMetadata.find((a) => a.agent_type === "a2a")?.logo_url;
const renderConfigureStep = () => (
<>
<Form.Item
label={<span className="text-sm font-medium text-gray-700">Agent Type</span>}
required
tooltip="Select the type of agent you want to create"
>
<Select
value={agentType}
onChange={handleAgentTypeChange}
size="large"
style={{ width: "100%" }}
optionLabelProp="label"
dropdownRender={(menu) => (
<>
{menu}
<Divider style={{ margin: "4px 0" }} />
<div className="px-2 py-1">
<div className="text-xs text-gray-400 font-medium mb-1 uppercase tracking-wide px-2">
Not listed?
</div>
<div
className={`flex items-center gap-3 px-2 py-2 rounded cursor-pointer transition-colors ${
agentType === CUSTOM_AGENT_TYPE
? "bg-amber-50"
: "hover:bg-amber-50"
}`}
onClick={() => handleAgentTypeChange(CUSTOM_AGENT_TYPE)}
>
<AppstoreOutlined className="text-amber-600 text-lg" />
<div className="flex-1">
<div className="flex items-center gap-2">
<span className="font-medium text-amber-700">Custom / Other</span>
<Tag color="orange" style={{ fontSize: 10, padding: "0 4px" }}>GENERIC</Tag>
</div>
<div className="text-xs text-amber-600">
For agents that don&apos;t follow a standard protocol — just needs a virtual key
</div>
</div>
</div>
</div>
</>
)}
>
{agentTypeMetadata.map((info) => (
<Select.Option
key={info.agent_type}
value={info.agent_type}
label={
<div className="flex items-center gap-2">
<img src={info.logo_url || ""} alt="" className="w-4 h-4 object-contain" />
<span>{info.agent_type_display_name}</span>
</div>
}
>
<div className="flex items-center gap-3 py-1">
<img
src={info.logo_url || ""}
alt={info.agent_type_display_name}
className="w-5 h-5 object-contain"
/>
<div>
<div className="font-medium">{info.agent_type_display_name}</div>
{info.description && (
<div className="text-xs text-gray-500">{info.description}</div>
)}
</div>
</div>
</Select.Option>
))}
</Select>
</Form.Item>
<div className="mt-4">
{agentType === CUSTOM_AGENT_TYPE ? (
<div className="space-y-4">
<Form.Item
label="Agent Name"
name="agent_name"
rules={[{ required: true, message: "Please enter an agent name" }]}
>
<Input placeholder="e.g. my-custom-agent" />
</Form.Item>
<Form.Item
label="Description"
name="description"
>
<Input.TextArea placeholder="Describe what this agent does…" rows={3} />
</Form.Item>
</div>
) : agentType === "a2a" ? (
<AgentFormFields showAgentName={true} />
) : selectedAgentTypeInfo?.use_a2a_form_fields ? (
<>
<AgentFormFields showAgentName={true} />
{selectedAgentTypeInfo.credential_fields.length > 0 && (
<div className="mt-4 p-4 border border-gray-200 rounded-lg">
<h4 className="text-sm font-medium text-gray-700 mb-3">
{selectedAgentTypeInfo.agent_type_display_name} Settings
</h4>
{selectedAgentTypeInfo.credential_fields.map((field) => (
<Form.Item
key={field.key}
label={field.label}
name={field.key}
rules={
field.required
? [{ required: true, message: `Please enter ${field.label}` }]
: undefined
}
tooltip={field.tooltip}
initialValue={field.default_value}
>
{field.field_type === "password" ? (
<Input.Password placeholder={field.placeholder || ""} />
) : (
<Input placeholder={field.placeholder || ""} />
)}
</Form.Item>
))}
</div>
)}
</>
) : selectedAgentTypeInfo ? (
<DynamicAgentFormFields agentTypeInfo={selectedAgentTypeInfo} />
) : null}
</div>
</>
);
const renderAssignKeyStep = () => {
const agentName = form.getFieldValue("agent_name") || "your-agent";
return (
<div>
{/* Agent name chip */}
<div className="flex justify-center mb-6">
<Tag icon={<RobotOutlined />} color="purple" className="px-3 py-1 text-sm">
{agentName}
</Tag>
</div>
<div className="space-y-3">
{/* Option: Create new key */}
<div
className={`p-4 border-2 rounded-lg cursor-pointer transition-colors ${
keyAssignOption === "create_new"
? "border-indigo-600 bg-indigo-50"
: "border-gray-200 bg-white hover:border-gray-300"
}`}
onClick={() => setKeyAssignOption("create_new")}
>
<div className="flex items-start justify-between">
<div className="flex items-start gap-3 flex-1">
<Radio
value="create_new"
checked={keyAssignOption === "create_new"}
onChange={() => setKeyAssignOption("create_new")}
/>
<div className="flex-1">
<div className="flex items-center gap-2">
<KeyOutlined className="text-indigo-600" />
<span className="font-medium text-gray-900">Create a new key for this agent</span>
</div>
<p className="text-sm text-gray-500 mt-1">
A dedicated key scoped to this agent.
</p>
{keyAssignOption === "create_new" && (
<div className="mt-3 space-y-3" onClick={(e) => e.stopPropagation()}>
<div>
<label className="text-sm text-gray-600 block mb-1">Key Name</label>
<Input
value={newKeyName}
onChange={(e) => setNewKeyName(e.target.value)}
placeholder="e.g. my-agent-key"
/>
</div>
<div>
<label className="text-sm text-gray-600 block mb-1">
Allowed Models <span className="text-gray-400">(optional — leave empty for all models)</span>
</label>
<Select
mode="tags"
style={{ width: "100%" }}
placeholder="e.g. gpt-4o, claude-3-5-sonnet"
value={newKeyModels}
onChange={setNewKeyModels}
tokenSeparators={[","]}
/>
</div>
</div>
)}
</div>
</div>
<Tag color="green">Recommended</Tag>
</div>
</div>
{/* Option: Assign existing key */}
<div
className={`p-4 border-2 rounded-lg cursor-pointer transition-colors ${
keyAssignOption === "existing_key"
? "border-indigo-600 bg-indigo-50"
: "border-gray-200 bg-white hover:border-gray-300"
}`}
onClick={() => setKeyAssignOption("existing_key")}
>
<div className="flex items-start gap-3">
<Radio
value="existing_key"
checked={keyAssignOption === "existing_key"}
onChange={() => setKeyAssignOption("existing_key")}
/>
<div className="flex-1">
<div className="flex items-center gap-2">
<KeyOutlined className="text-gray-500" />
<span className="font-medium text-gray-900">Assign an existing key</span>
</div>
<p className="text-sm text-gray-500 mt-1">
Re-assign a key you already have to this agent.
</p>
{keyAssignOption === "existing_key" && (
<div className="mt-3" onClick={(e) => e.stopPropagation()}>
<Select
showSearch
style={{ width: "100%" }}
placeholder="Search by key name…"
loading={loadingKeys}
value={selectedExistingKey}
onChange={(value) => setSelectedExistingKey(value)}
filterOption={(input, option) =>
(option?.label as string ?? "").toLowerCase().includes(input.toLowerCase())
}
options={existingKeys.map((k) => ({
label: k.key_alias || k.token?.slice(0, 12) + "…",
value: k.token,
}))}
/>
</div>
)}
</div>
</div>
</div>
</div>
<div className="text-center mt-4">
<button
type="button"
className="text-sm text-gray-500 underline hover:text-gray-700"
onClick={() => setKeyAssignOption("skip")}
>
Skip for now — I&apos;ll assign a key later
</button>
</div>
</div>
);
};
const renderReadyStep = () => (
<div className="text-center py-6">
<CheckCircleFilled className="text-5xl text-green-500 mb-4" style={{ fontSize: 48 }} />
<h3 className="text-xl font-semibold text-gray-900 mb-2">Agent Created!</h3>
<div className="flex justify-center mb-4">
<Tag icon={<RobotOutlined />} color="purple" className="px-3 py-1 text-sm">
{createdAgentName}
</Tag>
</div>
{createdKeyValue && (
<div className="mt-4 text-left max-w-md mx-auto">
<CreatedKeyDisplay apiKey={createdKeyValue} />
</div>
)}
{assignedKeyAlias && (
<p className="text-sm text-gray-600 mt-2">
Key <span className="font-medium">{assignedKeyAlias}</span> has been assigned to this agent.
</p>
)}
{!createdKeyValue && !assignedKeyAlias && keyAssignOption === "skip" && (
<p className="text-sm text-gray-500 mt-2">
No key assigned. You can create one from the Virtual Keys page.
</p>
)}
</div>
);
return (
<Modal
title={
<div className="flex items-center space-x-3 pb-4 border-b border-gray-100">
{selectedLogo && (
<img
src={selectedLogo}
alt="Agent"
className="w-6 h-6 object-contain"
/>
{selectedLogo && currentStep < 1 && (
<img src={selectedLogo} alt="Agent" className="w-6 h-6 object-contain" />
)}
<h2 className="text-xl font-semibold text-gray-900">Add New Agent</h2>
</div>
}
open={visible}
onCancel={handleCancel}
onCancel={handleClose}
footer={null}
width={900}
className="top-8"
@ -134,103 +534,60 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
}}
>
<div className="mt-4">
{/* Step indicator */}
<Steps current={currentStep} size="small" className="mb-8">
<Step title="Configure" />
<Step title="Assign Key" />
<Step title="Ready" />
</Steps>
<Form
form={form}
layout="vertical"
onFinish={handleSubmit}
initialValues={agentType === "a2a" ? getDefaultFormValues() : {}}
className="space-y-4"
>
{/* Agent Type Selection */}
<Form.Item
label={<span className="text-sm font-medium text-gray-700">Agent Type</span>}
required
tooltip="Select the type of agent you want to create"
>
<Select
value={agentType}
onChange={handleAgentTypeChange}
size="large"
style={{ width: "100%" }}
optionLabelProp="label"
>
{agentTypeMetadata.map((info) => (
<Select.Option
key={info.agent_type}
value={info.agent_type}
label={
<div className="flex items-center gap-2">
<img src={info.logo_url || ""} alt="" className="w-4 h-4 object-contain" />
<span>{info.agent_type_display_name}</span>
</div>
}
>
<div className="flex items-center gap-3 py-1">
<img
src={info.logo_url || ""}
alt={info.agent_type_display_name}
className="w-5 h-5 object-contain"
/>
<div>
<div className="font-medium">{info.agent_type_display_name}</div>
{info.description && (
<div className="text-xs text-gray-500">{info.description}</div>
)}
</div>
</div>
</Select.Option>
))}
</Select>
</Form.Item>
{/* Conditional Form Fields */}
<div className="mt-6">
{agentType === "a2a" ? (
<AgentFormFields showAgentName={true} />
) : selectedAgentTypeInfo?.use_a2a_form_fields ? (
// A2A-compatible agents (like Pydantic AI) use full A2A form fields
// plus any additional credential fields
<>
<AgentFormFields showAgentName={true} />
{selectedAgentTypeInfo.credential_fields.length > 0 && (
<div className="mt-4 p-4 border border-gray-200 rounded-lg">
<h4 className="text-sm font-medium text-gray-700 mb-3">
{selectedAgentTypeInfo.agent_type_display_name} Settings
</h4>
{selectedAgentTypeInfo.credential_fields.map((field) => (
<Form.Item
key={field.key}
label={field.label}
name={field.key}
rules={field.required ? [{ required: true, message: `Please enter ${field.label}` }] : undefined}
tooltip={field.tooltip}
initialValue={field.default_value}
>
{field.field_type === "password" ? (
<Input.Password placeholder={field.placeholder || ""} />
) : (
<Input placeholder={field.placeholder || ""} />
)}
</Form.Item>
))}
</div>
)}
</>
) : selectedAgentTypeInfo ? (
<DynamicAgentFormFields agentTypeInfo={selectedAgentTypeInfo} />
) : null}
</div>
{/* Footer Buttons */}
<div className="flex items-center justify-end space-x-3 pt-6 border-t border-gray-100 mt-6">
<Button variant="secondary" onClick={handleCancel}>
Cancel
</Button>
<Button variant="primary" loading={isSubmitting}>
{isSubmitting ? "Creating..." : "Create Agent"}
</Button>
</div>
{currentStep === 0 && renderConfigureStep()}
{currentStep === 1 && renderAssignKeyStep()}
{currentStep === 2 && renderReadyStep()}
</Form>
{/* Footer navigation */}
<div className="flex items-center justify-between pt-6 border-t border-gray-100 mt-6">
<div>
{currentStep > 0 && currentStep < 2 && (
<button
type="button"
onClick={handleBack}
className="text-sm text-gray-600 border border-gray-300 rounded px-4 py-2 hover:bg-gray-50"
>
← Back
</button>
)}
</div>
<div className="flex gap-3">
{currentStep < 2 && (
<Button variant="secondary" onClick={handleClose}>
Cancel
</Button>
)}
{currentStep === 0 && (
<Button variant="primary" onClick={handleNext}>
Next →
</Button>
)}
{currentStep === 1 && (
<Button variant="primary" loading={isSubmitting} onClick={handleCreateAgent}>
{isSubmitting ? "Creating..." : "Create Agent →"}
</Button>
)}
{currentStep === 2 && (
<Button variant="primary" onClick={handleClose}>
Done
</Button>
)}
</div>
</div>
</div>
</Modal>
);

10
ui/litellm-dashboard/src/components/agents/agent_config.ts
View file

@ -54,9 +54,9 @@ export const AGENT_FORM_CONFIG: {
name: "url",
label: "URL",
type: "url",
required: true,
required: false,
placeholder: "http://localhost:9999/",
tooltip: "Base URL where the agent is hosted",
tooltip: "Base URL where the agent is hosted (optional)",
},
{
name: "version",
@ -237,9 +237,9 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => {
agent_name: values.agent_name,
agent_card_params: {
protocolVersion: values.protocolVersion || "1.0",
name: values.name,
description: values.description,
url: values.url,
name: values.name || values.agent_name,
description: values.description || "",
url: values.url || "",
version: values.version || "1.0.0",
defaultInputModes: existingAgent?.agent_card_params?.defaultInputModes || ["text"],
defaultOutputModes: existingAgent?.agent_card_params?.defaultOutputModes || ["text"],

16
ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx
View file

@ -10,13 +10,15 @@ const { Panel } = Collapse;
interface AgentFormFieldsProps {
showAgentName?: boolean;
visiblePanels?: string[];
}
/**
* Reusable form fields component for agent forms
* Uses shared configuration from agent_config.ts
*/
const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true }) => {
const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true, visiblePanels }) => {
const shouldShow = (key: string) => !visiblePanels || visiblePanels.includes(key);
return (
<>
{showAgentName && (
@ -32,6 +34,7 @@ const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true
<Collapse defaultActiveKey={['basic']} style={{ marginBottom: 16 }}>
{/* Basic Information */}
{shouldShow(AGENT_FORM_CONFIG.basic.key) && (
<Panel header={`${AGENT_FORM_CONFIG.basic.title} (Required)`} key={AGENT_FORM_CONFIG.basic.key}>
{AGENT_FORM_CONFIG.basic.fields.map((field) => (
<Form.Item
@ -49,8 +52,10 @@ const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true
</Form.Item>
))}
</Panel>
)}
{/* Skills */}
{shouldShow(AGENT_FORM_CONFIG.skills.key) && (
<Panel header={`${AGENT_FORM_CONFIG.skills.title} (Required)`} key={AGENT_FORM_CONFIG.skills.key}>
<Form.List name="skills">
{(fields, { add, remove }) => (
@ -127,8 +132,10 @@ const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true
)}
</Form.List>
</Panel>
)}
{/* Capabilities */}
{shouldShow(AGENT_FORM_CONFIG.capabilities.key) && (
<Panel header={AGENT_FORM_CONFIG.capabilities.title} key={AGENT_FORM_CONFIG.capabilities.key}>
{AGENT_FORM_CONFIG.capabilities.fields.map((field) => (
<Form.Item
@ -141,8 +148,10 @@ const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true
</Form.Item>
))}
</Panel>
)}
{/* Optional Settings */}
{shouldShow(AGENT_FORM_CONFIG.optional.key) && (
<Panel header={AGENT_FORM_CONFIG.optional.title} key={AGENT_FORM_CONFIG.optional.key}>
{AGENT_FORM_CONFIG.optional.fields.map((field) => (
<Form.Item
@ -155,13 +164,17 @@ const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true
</Form.Item>
))}
</Panel>
)}
{/* Cost Configuration */}
{shouldShow(AGENT_FORM_CONFIG.cost.key) && (
<Panel header={AGENT_FORM_CONFIG.cost.title} key={AGENT_FORM_CONFIG.cost.key}>
<CostConfigFields />
</Panel>
)}
{/* LiteLLM Parameters */}
{shouldShow(AGENT_FORM_CONFIG.litellm.key) && (
<Panel header={AGENT_FORM_CONFIG.litellm.title} key={AGENT_FORM_CONFIG.litellm.key}>
{AGENT_FORM_CONFIG.litellm.fields.map((field) => (
<Form.Item
@ -174,6 +187,7 @@ const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true
</Form.Item>
))}
</Panel>
)}
</Collapse>
</>
);

57
ui/litellm-dashboard/src/components/common_components/LabeledField.test.tsx Normal file
View file

@ -0,0 +1,57 @@
import { render, screen } from "@testing-library/react";
import { describe, expect, it } from "vitest";
import LabeledField from "./LabeledField";
describe("LabeledField", () => {
it("should render the label and value", () => {
render(<LabeledField label="User Email" value="test@example.com" />);
expect(screen.getByText("User Email")).toBeInTheDocument();
expect(screen.getByText("test@example.com")).toBeInTheDocument();
});
it("should render the icon when provided", () => {
render(
<LabeledField label="Name" value="Alice" icon={<span data-testid="test-icon" />} />,
);
expect(screen.getByTestId("test-icon")).toBeInTheDocument();
});
it("should show '-' when value is empty", () => {
render(<LabeledField label="User ID" value="" />);
expect(screen.getByText("-")).toBeInTheDocument();
});
it("should show 'Default Proxy Admin' tag when value is default_user_id and defaultUserIdCheck is true", () => {
render(
<LabeledField label="User ID" value="default_user_id" copyable defaultUserIdCheck />,
);
expect(screen.getByText("Default Proxy Admin")).toBeInTheDocument();
expect(screen.queryByText("default_user_id")).not.toBeInTheDocument();
});
it("should show raw value when value is default_user_id but defaultUserIdCheck is false", () => {
render(<LabeledField label="User ID" value="default_user_id" />);
expect(screen.getByText("default_user_id")).toBeInTheDocument();
expect(screen.queryByText("Default Proxy Admin")).not.toBeInTheDocument();
});
it("should not be copyable when value is empty", () => {
const { container } = render(<LabeledField label="User ID" value="" copyable />);
// antd adds a .ant-typography-copy element when copyable; should not be present
expect(container.querySelector(".ant-typography-copy")).not.toBeInTheDocument();
});
it("should not be copyable when value is default_user_id and defaultUserIdCheck is true", () => {
const { container } = render(
<LabeledField label="User ID" value="default_user_id" copyable defaultUserIdCheck />,
);
expect(container.querySelector(".ant-typography-copy")).not.toBeInTheDocument();
});
it("should be copyable when copyable is true and value is present", () => {
const { container } = render(
<LabeledField label="User ID" value="user-123" copyable />,
);
expect(container.querySelector(".ant-typography-copy")).toBeInTheDocument();
});
});

52
ui/litellm-dashboard/src/components/common_components/LabeledField.tsx Normal file
View file

@ -0,0 +1,52 @@
import React from "react";
import { Typography, Space } from "antd";
import DefaultProxyAdminTag from "./DefaultProxyAdminTag";
const { Text } = Typography;
interface LabeledFieldProps {
label: string;
value: string;
icon?: React.ReactNode;
truncate?: boolean;
copyable?: boolean;
defaultUserIdCheck?: boolean;
}
export default function LabeledField({
label,
value,
icon,
truncate = false,
copyable = false,
defaultUserIdCheck = false,
}: LabeledFieldProps) {
const isEmpty = !value;
const isDefaultUser = defaultUserIdCheck && value === "default_user_id";
const displayValue = isEmpty ? "-" : value;
const isCopyable = copyable && !isEmpty && !isDefaultUser;
const valueEl = isDefaultUser ? (
<DefaultProxyAdminTag userId={value} />
) : (
<Text
strong
copyable={isCopyable ? { tooltips: [`Copy ${label}`, "Copied!"] } : false}
ellipsis={truncate}
style={truncate ? { maxWidth: 160, display: "block" } : undefined}
>
{displayValue}
</Text>
);
return (
<div>
<Space size={4}>
<Text type="secondary">{icon}</Text>
<Text type="secondary" style={{ fontSize: 12, textTransform: "uppercase", letterSpacing: "0.05em" }}>
{label}
</Text>
</Space>
<div>{valueEl}</div>
</div>
);
}

6
ui/litellm-dashboard/src/components/leftnav.tsx
View file

@ -134,6 +134,12 @@ const menuGroups: MenuGroup[] = [
label: "Vector Stores",
icon: <DatabaseOutlined />,
},
{
key: "tool-policies",
page: "tool-policies",
label: "Tool Policies",
icon: <SafetyOutlined />,
},
],
},
],

231
ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.test.tsx Normal file
View file

@ -0,0 +1,231 @@
/* @vitest-environment jsdom */
import { act, render, screen } from "@testing-library/react";
import { describe, it, expect, vi, beforeEach } from "vitest";
import HealthCheckComponent from "./HealthCheckComponent";
const mockIndividualModelHealthCheckCall = vi.fn();
const mockLatestHealthChecksCall = vi.fn();
vi.mock("../networking", () => ({
individualModelHealthCheckCall: (...args: unknown[]) => mockIndividualModelHealthCheckCall(...args),
latestHealthChecksCall: (...args: unknown[]) => mockLatestHealthChecksCall(...args),
}));
describe("HealthCheckComponent", () => {
const getDisplayModelName = (model: { model_name?: string }) => model.model_name ?? "";
beforeEach(() => {
vi.clearAllMocks();
mockLatestHealthChecksCall.mockResolvedValue({ latest_health_checks: {} });
mockIndividualModelHealthCheckCall.mockResolvedValue({
healthy_count: 1,
unhealthy_count: 0,
healthy_endpoints: [],
unhealthy_endpoints: [],
});
});
it("should render the health check section", async () => {
const modelData = {
data: [
{
model_name: "gpt-4",
model_info: { id: "deployment-1" },
litellm_model_name: "gpt-4",
},
],
};
await act(async () => {
render(
<HealthCheckComponent
accessToken="token"
modelData={modelData}
all_models_on_proxy={["deployment-1"]}
getDisplayModelName={getDisplayModelName}
/>,
);
});
await act(async () => {
await new Promise((r) => setTimeout(r, 0));
});
expect(screen.getByText("Model Health Status")).toBeInTheDocument();
expect(
screen.getByText("Run health checks on individual models to verify they are working correctly"),
).toBeInTheDocument();
});
it("should call individualModelHealthCheckCall with model id when run health check is triggered", async () => {
const modelData = {
data: [
{
model_name: "gpt-4",
model_info: { id: "deployment-abc-123" },
litellm_model_name: "gpt-4",
},
],
};
render(
<HealthCheckComponent
accessToken="token-123"
modelData={modelData}
all_models_on_proxy={["deployment-abc-123"]}
getDisplayModelName={getDisplayModelName}
/>,
);
const runButtons = screen.getAllByTestId("run-health-check-btn");
expect(runButtons.length).toBeGreaterThanOrEqual(1);
const runButton = runButtons[0];
await act(async () => {
runButton.click();
});
await act(async () => {
await new Promise((r) => setTimeout(r, 50));
});
expect(mockIndividualModelHealthCheckCall).toHaveBeenCalledWith("token-123", "deployment-abc-123");
expect(mockIndividualModelHealthCheckCall).not.toHaveBeenCalledWith("token-123", "gpt-4");
});
describe("latest_health_checks keyed by model id", () => {
it("should show status from latest_health_checks when keys match model ids", async () => {
const modelData = {
data: [
{
model_name: "gpt-4",
model_info: { id: "id-alpha" },
litellm_model_name: "gpt-4",
},
{
model_name: "gpt-4",
model_info: { id: "id-beta" },
litellm_model_name: "gpt-4",
},
],
};
mockLatestHealthChecksCall.mockResolvedValue({
latest_health_checks: {
"id-alpha": {
status: "healthy",
checked_at: "2024-01-15T10:00:00Z",
error_message: null,
},
"id-beta": {
status: "unhealthy",
checked_at: "2024-01-15T10:05:00Z",
error_message: "Connection failed",
},
},
});
await act(async () => {
render(
<HealthCheckComponent
accessToken="token"
modelData={modelData}
all_models_on_proxy={["id-alpha", "id-beta"]}
getDisplayModelName={getDisplayModelName}
/>,
);
});
await act(async () => {
await new Promise((r) => setTimeout(r, 0));
});
expect(mockLatestHealthChecksCall).toHaveBeenCalledWith("token");
const healthyBadges = screen.getAllByText("healthy");
const unhealthyBadges = screen.getAllByText("unhealthy");
expect(healthyBadges.length).toBeGreaterThanOrEqual(1);
expect(unhealthyBadges.length).toBeGreaterThanOrEqual(1);
});
it("should skip latest_health_checks entries whose key is not a known model id", async () => {
const modelData = {
data: [
{
model_name: "gpt-4",
model_info: { id: "current-model-id" },
litellm_model_name: "gpt-4",
},
],
};
mockLatestHealthChecksCall.mockResolvedValue({
latest_health_checks: {
"current-model-id": {
status: "healthy",
checked_at: "2024-01-15T10:00:00Z",
error_message: null,
},
"deleted-or-unknown-id": {
status: "unhealthy",
checked_at: "2024-01-15T10:05:00Z",
error_message: "Stale entry",
},
},
});
await act(async () => {
render(
<HealthCheckComponent
accessToken="token"
modelData={modelData}
all_models_on_proxy={["current-model-id"]}
getDisplayModelName={getDisplayModelName}
/>,
);
});
await act(async () => {
await new Promise((r) => setTimeout(r, 0));
});
expect(screen.getByText("healthy")).toBeInTheDocument();
expect(screen.queryByText("unhealthy")).not.toBeInTheDocument();
});
it("should not apply status when latest_health_checks key is model name not model id", async () => {
const modelData = {
data: [
{
model_name: "gpt-4",
model_info: { id: "model-id-123" },
litellm_model_name: "gpt-4",
},
],
};
mockLatestHealthChecksCall.mockResolvedValue({
latest_health_checks: {
"gpt-4": {
status: "healthy",
checked_at: "2024-01-15T10:00:00Z",
error_message: null,
},
},
});
await act(async () => {
render(
<HealthCheckComponent
accessToken="token"
modelData={modelData}
all_models_on_proxy={["model-id-123"]}
getDisplayModelName={getDisplayModelName}
/>,
);
});
await act(async () => {
await new Promise((r) => setTimeout(r, 0));
});
expect(screen.queryByText("healthy")).not.toBeInTheDocument();
expect(screen.getByText("none")).toBeInTheDocument();
});
});
});

199
ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx
View file

@ -53,65 +53,48 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
const healthTableRef = useRef<TableInstance<any>>(null);
// Initialize health statuses on component mount
// Initialize health statuses on component mount (keyed by model id)
useEffect(() => {
if (!accessToken || !modelData?.data) return;
const initializeHealthStatuses = async () => {
const healthStatusMap: { [key: string]: HealthStatus } = {};
// Initialize all models with default state using model names
// Initialize all models with default state using model ids
modelData.data.forEach((model: any) => {
const modelName = model.model_name;
healthStatusMap[modelName] = {
status: "none",
lastCheck: "None",
lastSuccess: "None",
loading: false,
error: undefined,
fullError: undefined,
successResponse: undefined,
};
const modelId = model.model_info?.id;
if (modelId) {
healthStatusMap[modelId] = {
status: "none",
lastCheck: "None",
lastSuccess: "None",
loading: false,
error: undefined,
fullError: undefined,
successResponse: undefined,
};
}
});
try {
const latestHealthChecks = await latestHealthChecksCall(accessToken);
// Override with actual database data if it exists
// Override with actual database data if it exists (latest_health_checks is keyed by model_id)
if (
latestHealthChecks &&
latestHealthChecks.latest_health_checks &&
typeof latestHealthChecks.latest_health_checks === "object"
) {
Object.entries(latestHealthChecks.latest_health_checks).forEach(([key, checkData]: [string, any]) => {
Object.entries(latestHealthChecks.latest_health_checks).forEach(([modelId, checkData]: [string, any]) => {
if (!checkData) return;
let targetModelName: string | null = null;
// Key is model_id from the backend (guaranteed by DB schema)
const modelExists = modelData.data.some((m: any) => m.model_info?.id === modelId);
if (!modelExists) return;
// The key could be either model_id or model_name, try both approaches
const directModelMatch = modelData.data.find((m: any) => m.model_name === key);
if (directModelMatch) {
targetModelName = directModelMatch.model_name;
} else {
// If not a direct match, treat as model_id and find the corresponding model
const modelByIdMatch = modelData.data.find((m: any) => m.model_info && m.model_info.id === key);
if (modelByIdMatch) {
targetModelName = modelByIdMatch.model_name;
} else {
// Check if checkData contains model_name and use that
if (checkData.model_name) {
const modelByNameInData = modelData.data.find((m: any) => m.model_name === checkData.model_name);
if (modelByNameInData) {
targetModelName = modelByNameInData.model_name;
}
}
}
}
const fullError = checkData.error_message || undefined;
if (targetModelName) {
const fullError = checkData.error_message || undefined;
healthStatusMap[targetModelName] = {
healthStatusMap[modelId] = {
status: checkData.status || "unknown",
lastCheck: checkData.checked_at ? new Date(checkData.checked_at).toLocaleString() : "None",
lastSuccess:
@ -125,7 +108,6 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
fullError: fullError,
successResponse: checkData.status === "healthy" ? checkData : undefined,
};
}
});
}
} catch (healthError) {
@ -246,33 +228,31 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
return cleaned.length > 100 ? cleaned.substring(0, 97) + "..." : cleaned;
};
const runIndividualHealthCheck = async (modelName: string) => {
const runIndividualHealthCheck = async (modelId: string) => {
if (!accessToken) return;
setModelHealthStatuses((prev) => ({
...prev,
[modelName]: {
...prev[modelName],
[modelId]: {
...prev[modelId],
loading: true,
status: "checking",
},
}));
try {
// Run the health check and process the response directly
const response = await individualModelHealthCheckCall(accessToken, modelName);
const response = await individualModelHealthCheckCall(accessToken, modelId);
const currentTime = new Date().toLocaleString();
// Check if there are any unhealthy endpoints (which means this specific model failed)
if (response.unhealthy_count > 0 && response.unhealthy_endpoints && response.unhealthy_endpoints.length > 0) {
const rawError = response.unhealthy_endpoints[0]?.error || "Health check failed";
const errorMessage = extractMeaningfulError(rawError);
setModelHealthStatuses((prev) => ({
...prev,
[modelName]: {
[modelId]: {
status: "unhealthy",
lastCheck: currentTime,
lastSuccess: prev[modelName]?.lastSuccess || "None",
lastSuccess: prev[modelId]?.lastSuccess || "None",
loading: false,
error: errorMessage,
fullError: rawError,
@ -281,7 +261,7 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
} else {
setModelHealthStatuses((prev) => ({
...prev,
[modelName]: {
[modelId]: {
status: "healthy",
lastCheck: currentTime,
lastSuccess: currentTime,
@ -291,41 +271,33 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
}));
}
// Refresh health status from database to get the saved check data including timestamp
try {
const latestHealthChecks = await latestHealthChecksCall(accessToken);
const checkData = latestHealthChecks.latest_health_checks?.[modelId];
// Find the model ID for this model name to look up database data
const model = modelData.data.find((m: any) => m.model_name === modelName);
if (model) {
const modelId = model.model_info.id;
const checkData = latestHealthChecks.latest_health_checks?.[modelId];
if (checkData) {
const fullError = checkData.error_message || undefined;
setModelHealthStatuses((prev) => ({
...prev,
[modelName]: {
status: checkData.status || prev[modelName]?.status || "unknown",
lastCheck: checkData.checked_at
? new Date(checkData.checked_at).toLocaleString()
: prev[modelName]?.lastCheck || "None",
lastSuccess:
checkData.status === "healthy"
? checkData.checked_at
? new Date(checkData.checked_at).toLocaleString()
: prev[modelName]?.lastSuccess || "None"
: prev[modelName]?.lastSuccess || "None",
loading: false,
error: fullError ? extractMeaningfulError(fullError) : prev[modelName]?.error,
fullError: fullError || prev[modelName]?.fullError,
successResponse: checkData.status === "healthy" ? checkData : prev[modelName]?.successResponse,
},
}));
}
if (checkData) {
const fullError = checkData.error_message || undefined;
setModelHealthStatuses((prev) => ({
...prev,
[modelId]: {
status: checkData.status || prev[modelId]?.status || "unknown",
lastCheck: checkData.checked_at
? new Date(checkData.checked_at).toLocaleString()
: prev[modelId]?.lastCheck || "None",
lastSuccess:
checkData.status === "healthy"
? checkData.checked_at
? new Date(checkData.checked_at).toLocaleString()
: prev[modelId]?.lastSuccess || "None"
: prev[modelId]?.lastSuccess || "None",
loading: false,
error: fullError ? extractMeaningfulError(fullError) : prev[modelId]?.error,
fullError: fullError || prev[modelId]?.fullError,
successResponse: checkData.status === "healthy" ? checkData : prev[modelId]?.successResponse,
},
}));
}
} catch (dbError) {
// Ignore database errors - we already have the health check result from the API call
console.debug("Could not fetch updated status from database (non-critical):", dbError);
}
} catch (error) {
@ -334,10 +306,10 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
const errorMessage = extractMeaningfulError(rawError);
setModelHealthStatuses((prev) => ({
...prev,
[modelName]: {
[modelId]: {
status: "unhealthy",
lastCheck: currentTime,
lastSuccess: prev[modelName]?.lastSuccess || "None",
lastSuccess: prev[modelId]?.lastSuccess || "None",
loading: false,
error: errorMessage,
fullError: rawError,
@ -349,11 +321,10 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
const runAllHealthChecks = async () => {
const modelsToCheck = selectedModelsForHealth.length > 0 ? selectedModelsForHealth : all_models_on_proxy;
// Set all models to loading state
const loadingStatuses = modelsToCheck.reduce(
(acc, modelName) => {
acc[modelName] = {
...modelHealthStatuses[modelName],
(acc, modelId) => {
acc[modelId] = {
...modelHealthStatuses[modelId],
loading: true,
status: "checking",
};
@ -364,30 +335,25 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
setModelHealthStatuses((prev) => ({ ...prev, ...loadingStatuses }));
// Store results from individual health checks
const healthCheckResults: { [key: string]: any } = {};
// Run all health checks in parallel and collect results
const healthCheckPromises = modelsToCheck.map(async (modelName) => {
const healthCheckPromises = modelsToCheck.map(async (modelId) => {
if (!accessToken) return;
try {
// Run the health check and store the result
const response = await individualModelHealthCheckCall(accessToken, modelName);
healthCheckResults[modelName] = response;
const response = await individualModelHealthCheckCall(accessToken, modelId);
healthCheckResults[modelId] = response;
// Update status immediately based on response
const currentTime = new Date().toLocaleString();
// Check if there are any unhealthy endpoints (which means this specific model failed)
if (response.unhealthy_count > 0 && response.unhealthy_endpoints && response.unhealthy_endpoints.length > 0) {
const rawError = response.unhealthy_endpoints[0]?.error || "Health check failed";
const errorMessage = extractMeaningfulError(rawError);
setModelHealthStatuses((prev) => ({
...prev,
[modelName]: {
[modelId]: {
status: "unhealthy",
lastCheck: currentTime,
lastSuccess: prev[modelName]?.lastSuccess || "None",
lastSuccess: prev[modelId]?.lastSuccess || "None",
loading: false,
error: errorMessage,
fullError: rawError,
@ -396,7 +362,7 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
} else {
setModelHealthStatuses((prev) => ({
...prev,
[modelName]: {
[modelId]: {
status: "healthy",
lastCheck: currentTime,
lastSuccess: currentTime,
@ -406,17 +372,16 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
}));
}
} catch (error) {
console.error(`Health check failed for ${modelName}:`, error);
// Set error status for failed health checks
console.error(`Health check failed for model id ${modelId}:`, error);
const currentTime = new Date().toLocaleString();
const rawError = error instanceof Error ? error.message : String(error);
const errorMessage = extractMeaningfulError(rawError);
setModelHealthStatuses((prev) => ({
...prev,
[modelName]: {
[modelId]: {
status: "unhealthy",
lastCheck: currentTime,
lastSuccess: prev[modelName]?.lastSuccess || "None",
lastSuccess: prev[modelId]?.lastSuccess || "None",
loading: false,
error: errorMessage,
fullError: rawError,
@ -425,27 +390,21 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
}
});
// Wait for all health checks to complete
await Promise.allSettled(healthCheckPromises);
// Refresh health statuses from database to get the saved check data including timestamps
try {
if (!accessToken) return;
const latestHealthChecks = await latestHealthChecksCall(accessToken);
if (latestHealthChecks.latest_health_checks) {
// Update health statuses from database, which should have the most accurate saved data
Object.entries(latestHealthChecks.latest_health_checks).forEach(([modelId, checkData]: [string, any]) => {
// Find the model name for this model ID
const model = modelData.data.find((m: any) => m.model_info.id === modelId);
if (model && modelsToCheck.includes(model.model_name) && checkData) {
const modelName = model.model_name;
if (modelsToCheck.includes(modelId) && checkData) {
const fullError = checkData.error_message || undefined;
setModelHealthStatuses((prev) => {
const currentStatus = prev[modelName];
const currentStatus = prev[modelId];
return {
...prev,
[modelName]: {
[modelId]: {
status: checkData.status || currentStatus?.status || "unknown",
lastCheck: checkData.checked_at
? new Date(checkData.checked_at).toLocaleString()
@ -468,15 +427,14 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
}
} catch (dbError) {
console.warn("Failed to fetch updated health statuses from database (non-critical):", dbError);
// This is non-critical - we already have the health check results from the API calls
}
};
const handleModelSelection = (modelName: string, checked: boolean) => {
const handleModelSelection = (modelId: string, checked: boolean) => {
if (checked) {
setSelectedModelsForHealth((prev) => [...prev, modelName]);
setSelectedModelsForHealth((prev) => [...prev, modelId]);
} else {
setSelectedModelsForHealth((prev) => prev.filter((name) => name !== modelName));
setSelectedModelsForHealth((prev) => prev.filter((id) => id !== modelId));
setAllModelsSelected(false);
}
};
@ -580,8 +538,9 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
teams,
)}
data={modelData.data.map((model: any) => {
const modelName = model.model_name;
const healthStatus = modelHealthStatuses[modelName] || {
const modelId = model.model_info?.id;
const healthStatus = modelId ? modelHealthStatuses[modelId] : null;
const status = healthStatus || {
status: "none",
lastCheck: "None",
loading: false,
@ -591,12 +550,12 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
model_info: model.model_info,
provider: model.provider,
litellm_model_name: model.litellm_model_name,
health_status: healthStatus.status,
last_check: healthStatus.lastCheck,
last_success: healthStatus.lastSuccess || "None",
health_loading: healthStatus.loading,
health_error: healthStatus.error,
health_full_error: healthStatus.fullError,
health_status: status.status,
last_check: status.lastCheck,
last_success: status.lastSuccess || "None",
health_loading: status.loading,
health_error: status.error,
health_full_error: status.fullError,
};
})}
isLoading={false}

33
ui/litellm-dashboard/src/components/model_dashboard/health_check_columns.tsx
View file

@ -35,9 +35,9 @@ export const healthCheckColumns = (
modelHealthStatuses: { [key: string]: HealthStatus },
selectedModelsForHealth: string[],
allModelsSelected: boolean,
handleModelSelection: (modelName: string, checked: boolean) => void,
handleModelSelection: (modelId: string, checked: boolean) => void,
handleSelectAll: (checked: boolean) => void,
runIndividualHealthCheck: (modelName: string) => void,
runIndividualHealthCheck: (modelId: string) => void,
getStatusBadge: (status: string) => JSX.Element,
getDisplayModelName: (model: any) => string,
showErrorModal?: (modelName: string, cleanedError: string, fullError: string) => void,
@ -62,14 +62,14 @@ export const healthCheckColumns = (
sortingFn: "alphanumeric",
cell: ({ row }) => {
const model = row.original;
const modelName = model.model_name;
const isSelected = selectedModelsForHealth.includes(modelName);
const modelId = model.model_info?.id ?? "";
const isSelected = selectedModelsForHealth.includes(modelId);
return (
<div className="flex items-center gap-2">
<Checkbox
checked={isSelected}
onChange={(e) => handleModelSelection(modelName, e.target.checked)}
onChange={(e) => handleModelSelection(modelId, e.target.checked)}
onClick={(e) => e.stopPropagation()}
/>
<Tooltip title={model.model_info.id}>
@ -169,8 +169,9 @@ export const healthCheckColumns = (
);
}
const modelName = model.model_name;
const hasSuccessResponse = healthStatus.status === "healthy" && modelHealthStatuses[modelName]?.successResponse;
const modelId = model.model_info?.id ?? "";
const displayName = getDisplayModelName(model) || model.model_name;
const hasSuccessResponse = healthStatus.status === "healthy" && modelHealthStatuses[modelId]?.successResponse;
return (
<div className="flex items-center space-x-2">
@ -178,7 +179,7 @@ export const healthCheckColumns = (
{hasSuccessResponse && showSuccessModal && (
<Tooltip title="View response details" placement="top">
<button
onClick={() => showSuccessModal(modelName, modelHealthStatuses[modelName]?.successResponse)}
onClick={() => showSuccessModal(displayName, modelHealthStatuses[modelId]?.successResponse)}
className="p-1 text-green-600 hover:text-green-800 hover:bg-green-50 rounded cursor-pointer transition-colors"
>
<InformationCircleIcon className="h-4 w-4" />
@ -195,8 +196,9 @@ export const healthCheckColumns = (
enableSorting: false,
cell: ({ row }) => {
const model = row.original;
const modelName = model.model_name;
const healthStatus = modelHealthStatuses[modelName];
const modelId = model.model_info?.id ?? "";
const displayName = getDisplayModelName(model) || model.model_name;
const healthStatus = modelHealthStatuses[modelId];
if (!healthStatus?.error) {
return <Text className="text-gray-400 text-sm">No errors</Text>;
@ -215,7 +217,7 @@ export const healthCheckColumns = (
{showErrorModal && fullError !== cleanedError && (
<Tooltip title="View full error details" placement="top">
<button
onClick={() => showErrorModal(modelName, cleanedError, fullError)}
onClick={() => showErrorModal(displayName, cleanedError, fullError)}
className="p-1 text-red-600 hover:text-red-800 hover:bg-red-50 rounded cursor-pointer transition-colors"
>
<InformationCircleIcon className="h-4 w-4" />
@ -294,8 +296,8 @@ export const healthCheckColumns = (
},
cell: ({ row }) => {
const model = row.original;
const modelName = model.model_name;
const healthStatus = modelHealthStatuses[modelName];
const modelId = model.model_info?.id ?? "";
const healthStatus = modelHealthStatuses[modelId];
const lastSuccess = healthStatus?.lastSuccess || "None";
return <Text className="text-gray-600 text-sm">{lastSuccess}</Text>;
@ -306,7 +308,7 @@ export const healthCheckColumns = (
id: "actions",
cell: ({ row }) => {
const model = row.original;
const modelName = model.model_name;
const modelId = model.model_info?.id ?? "";
const hasExistingStatus = model.health_status && model.health_status !== "none";
const tooltipText = model.health_loading
@ -318,6 +320,7 @@ export const healthCheckColumns = (
return (
<Tooltip title={tooltipText} placement="top">
<button
data-testid="run-health-check-btn"
className={`p-2 rounded-md transition-colors ${
model.health_loading
? "text-gray-400 cursor-not-allowed bg-gray-100"
@ -325,7 +328,7 @@ export const healthCheckColumns = (
}`}
onClick={() => {
if (!model.health_loading) {
runIndividualHealthCheck(modelName);
runIndividualHealthCheck(modelId);
}
}}
disabled={model.health_loading}

54
ui/litellm-dashboard/src/components/networking.test.ts
View file

@ -317,3 +317,57 @@ describe("UI config and public endpoints", () => {
expect(configCall).toBeDefined();
});
});
describe("individualModelHealthCheckCall", () => {
const originalFetch = global.fetch;
beforeEach(() => {
vi.clearAllMocks();
});
afterEach(() => {
global.fetch = originalFetch;
});
it("should call /health with model_id query param so health checks run by deployment id", async () => {
const mockFetch = vi.fn().mockResolvedValue({
ok: true,
json: vi.fn().mockResolvedValue({
healthy_count: 1,
unhealthy_count: 0,
healthy_endpoints: [],
unhealthy_endpoints: [],
}),
} as any);
global.fetch = mockFetch as any;
await Networking.individualModelHealthCheckCall("token-123", "deployment-abc-456");
expect(mockFetch).toHaveBeenCalledOnce();
const [url] = mockFetch.mock.calls[0];
const urlStr = typeof url === "string" ? url : (url as Request).url;
expect(urlStr).toContain("health");
const parsed = typeof url === "string" ? new URL(url, "http://example.com") : new URL((url as Request).url);
expect(parsed.searchParams.get("model_id")).toBe("deployment-abc-456");
expect(parsed.searchParams.has("model")).toBe(false);
});
it("should encode model_id in URL", async () => {
const mockFetch = vi.fn().mockResolvedValue({
ok: true,
json: vi.fn().mockResolvedValue({
healthy_count: 0,
unhealthy_count: 0,
healthy_endpoints: [],
unhealthy_endpoints: [],
}),
} as any);
global.fetch = mockFetch as any;
await Networking.individualModelHealthCheckCall("token", "id/with/slashes");
const [url] = mockFetch.mock.calls[0];
const parsed = typeof url === "string" ? new URL(url, "http://example.com") : new URL((url as Request).url);
expect(parsed.searchParams.get("model_id")).toBe("id/with/slashes");
});
});

169
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -930,6 +930,35 @@ export const keyCreateCall = async (
}
};
export const keyCreateForAgentCall = async (
accessToken: string,
agentId: string,
keyAlias: string,
models: string[],
) => {
const url = proxyBaseUrl ? `${proxyBaseUrl}/key/generate` : `/key/generate`;
const response = await fetch(url, {
method: "POST",
headers: {
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
body: JSON.stringify({
agent_id: agentId,
key_alias: keyAlias,
models: models.length > 0 ? models : [],
}),
});
if (!response.ok) {
const errorData = await response.text();
handleError(errorData);
throw new Error("Failed to create key for agent");
}
return response.json();
};
export const userCreateCall = async (
accessToken: string,
userID: string | null,
@ -5054,14 +5083,14 @@ export const healthCheckCall = async (accessToken: string) => {
}
};
export const individualModelHealthCheckCall = async (accessToken: string, modelName: string) => {
export const individualModelHealthCheckCall = async (accessToken: string, modelId: string) => {
/**
* Run health check for a specific model using model name
* Run health check for a specific model using model ID (so each deployment is checked separately).
*/
try {
let url = proxyBaseUrl
? `${proxyBaseUrl}/health?model=${encodeURIComponent(modelName)}`
: `/health?model=${encodeURIComponent(modelName)}`;
? `${proxyBaseUrl}/health?model_id=${encodeURIComponent(modelId)}`
: `/health?model_id=${encodeURIComponent(modelId)}`;
const response = await fetch(url, {
method: "GET",
@ -5081,7 +5110,7 @@ export const individualModelHealthCheckCall = async (accessToken: string, modelN
const data = await response.json();
return data;
} catch (error) {
console.error(`Failed to call /health for model ${modelName}:`, error);
console.error(`Failed to call /health for model id ${modelId}:`, error);
throw error;
}
};
@ -5907,6 +5936,82 @@ export const enrichPolicyTemplateStream = async (
}
};
export interface UsageAiToolCallEvent {
tool_name: string;
tool_label: string;
arguments: Record<string, string>;
status: "running" | "complete" | "error";
error?: string;
}
export const usageAiChatStream = async (
accessToken: string,
messages: { role: string; content: string }[],
model: string,
onChunk: (content: string) => void,
onDone: () => void,
onError?: (error: string) => void,
onStatus?: (message: string) => void,
onToolCall?: (event: UsageAiToolCallEvent) => void,
signal?: AbortSignal,
) => {
const url = proxyBaseUrl
? `${proxyBaseUrl}/usage/ai/chat`
: `/usage/ai/chat`;
const response = await fetch(url, {
method: "POST",
headers: {
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
body: JSON.stringify({ messages, model }),
signal,
});
if (!response.ok) {
const errorData = await response.json();
const errorMessage = deriveErrorMessage(errorData);
handleError(errorMessage);
throw new Error(errorMessage);
}
const reader = response.body?.getReader();
if (!reader) throw new Error("No response body");
const decoder = new TextDecoder();
let buffer = "";
while (true) {
const { done, value } = await reader.read();
if (done) break;
buffer += decoder.decode(value, { stream: true });
const lines = buffer.split("\n");
buffer = lines.pop() || "";
for (const line of lines) {
if (!line.startsWith("data: ")) continue;
try {
const event = JSON.parse(line.slice(6));
if (event.type === "chunk") {
onChunk(event.content);
} else if (event.type === "status") {
onStatus?.(event.message);
} else if (event.type === "tool_call") {
onToolCall?.(event as UsageAiToolCallEvent);
} else if (event.type === "done") {
onDone();
} else if (event.type === "error") {
onError?.(event.message);
}
} catch {
// skip malformed events
}
}
}
};
export const createPolicyCall = async (accessToken: string, policyData: any) => {
try {
const url = proxyBaseUrl ? `${proxyBaseUrl}/policies` : `/policies`;
@ -9854,3 +9959,57 @@ export const checkGdprCompliance = async (
}
return response.json();
};
export interface ToolRow {
tool_id: string;
tool_name: string;
origin?: string;
call_policy: string;
call_count?: number;
assignments?: Record<string, any>;
key_hash?: string;
team_id?: string;
key_alias?: string;
created_at?: string;
updated_at?: string;
created_by?: string;
updated_by?: string;
}
export const fetchToolsList = async (accessToken: string): Promise<ToolRow[]> => {
const url = proxyBaseUrl ? `${proxyBaseUrl}/v1/tool/list` : `/v1/tool/list`;
const response = await fetch(url, {
method: "GET",
headers: {
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
});
if (!response.ok) {
const errorData = await response.text();
throw new Error(errorData);
}
const data = await response.json();
return data.tools ?? [];
};
export const updateToolPolicy = async (
accessToken: string,
toolName: string,
callPolicy: string
): Promise<ToolRow> => {
const url = proxyBaseUrl ? `${proxyBaseUrl}/v1/tool/policy` : `/v1/tool/policy`;
const response = await fetch(url, {
method: "POST",
headers: {
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
body: JSON.stringify({ tool_name: toolName, call_policy: callPolicy }),
});
if (!response.ok) {
const errorData = await response.text();
throw new Error(errorData);
}
return response.json();
};

77
ui/litellm-dashboard/src/components/organisms/create_key_button.tsx
View file

@ -5,7 +5,7 @@ import { formatNumberWithCommas } from "@/utils/dataUtils";
import { InfoCircleOutlined } from "@ant-design/icons";
import { useQueryClient } from "@tanstack/react-query";
import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, TextInput, Title } from "@tremor/react";
import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tooltip } from "antd";
import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd";
import debounce from "lodash/debounce";
import React, { useCallback, useEffect, useState } from "react";
import { CopyToClipboard } from "react-copy-to-clipboard";
@ -29,6 +29,7 @@ import MCPServerSelector from "../mcp_server_management/MCPServerSelector";
import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions";
import NotificationsManager from "../molecules/notifications_manager";
import {
getAgentsList,
getGuardrailsList,
getPoliciesList,
getPossibleUserRoles,
@ -42,6 +43,7 @@ import {
import NumericalInput from "../shared/numerical_input";
import VectorStoreSelector from "../vector_store_management/VectorStoreSelector";
import { simplifyKeyGenerateError } from "./utils";
import CreatedKeyDisplay from "../shared/CreatedKeyDisplay";
const { Option } = Select;
@ -169,6 +171,8 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
const [rotationInterval, setRotationInterval] = useState<string>("30d");
const [routerSettings, setRouterSettings] = useState<RouterSettingsAccordionValue | null>(null);
const [routerSettingsKey, setRouterSettingsKey] = useState<number>(0);
const [agentsList, setAgentsList] = useState<{ agent_id: string; agent_name: string }[]>([]);
const [selectedAgentId, setSelectedAgentId] = useState<string | null>(null);
const handleOk = () => {
setIsModalVisible(false);
form.resetFields();
@ -180,6 +184,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
setRotationInterval("30d");
setRouterSettings(null);
setRouterSettingsKey((prev) => prev + 1);
setSelectedAgentId(null);
};
const handleCancel = () => {
@ -195,6 +200,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
setRotationInterval("30d");
setRouterSettings(null);
setRouterSettingsKey((prev) => prev + 1);
setSelectedAgentId(null);
};
useEffect(() => {
@ -203,6 +209,14 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
}
}, [accessToken, userID, userRole]);
useEffect(() => {
if (accessToken) {
getAgentsList(accessToken)
.then((res) => setAgentsList(res?.agents || []))
.catch(() => setAgentsList([]));
}
}, [accessToken]);
useEffect(() => {
const fetchGuardrails = async () => {
try {
@ -283,6 +297,12 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
if (keyOwner === "you") {
formValues.user_id = userID;
} else if (keyOwner === "agent") {
if (!selectedAgentId) {
message.error("Please select an agent");
return;
}
formValues.agent_id = selectedAgentId;
}
// Handle metadata for all key types
@ -539,6 +559,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
<Radio value="you">You</Radio>
<Radio value="service_account">Service Account</Radio>
{userRole === "Admin" && <Radio value="another_user">Another User</Radio>}
<Radio value="agent">Agent <Tag color="purple">New</Tag></Radio>
</Radio.Group>
</Form.Item>
@ -583,6 +604,32 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
</div>
</Form.Item>
)}
{keyOwner === "agent" && (
<div className="mt-4 p-4 bg-purple-50 border border-purple-200 rounded-md">
<div className="mb-3">
<span className="text-sm font-medium text-gray-700">
Select Agent <span className="text-red-500">*</span>
</span>
</div>
<Select
showSearch
placeholder="Select an agent"
style={{ width: "100%" }}
value={selectedAgentId}
onChange={(value) => setSelectedAgentId(value)}
filterOption={(input, option) =>
(option?.label as string)?.toLowerCase().includes(input.toLowerCase())
}
options={agentsList.map((a) => ({
label: a.agent_name || a.agent_id,
value: a.agent_id,
}))}
/>
<div className="text-xs text-gray-500 mt-2">
This key will be used by the selected agent to make requests to LiteLLM
</div>
</div>
)}
<Form.Item
label={
<span>
@ -1380,35 +1427,9 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
<Modal open={isModalVisible} onOk={handleOk} onCancel={handleCancel} footer={null}>
<Grid numItems={1} className="gap-2 w-full">
<Title>Save your Key</Title>
<Col numColSpan={1}>
<p>
Please save this secret key somewhere safe and accessible. For security reasons,{" "}
<b>you will not be able to view it again</b> through your LiteLLM account. If you lose this secret key,
you will need to generate a new one.
</p>
</Col>
<Col numColSpan={1}>
{apiKey != null ? (
<div>
<Text className="mt-3">Virtual Key:</Text>
<div
style={{
background: "#f8f8f8",
padding: "10px",
borderRadius: "5px",
marginBottom: "10px",
}}
>
<pre style={{ wordWrap: "break-word", whiteSpace: "normal" }}>{apiKey}</pre>
</div>
<CopyToClipboard text={apiKey} onCopy={handleCopy}>
<Button className="mt-3">Copy Virtual Key</Button>
</CopyToClipboard>
{/* <Button className="mt-3" onClick={sendSlackAlert}>
Test Key
</Button> */}
</div>
<CreatedKeyDisplay apiKey={apiKey} />
) : (
<Text>Key being created, this might take 30s</Text>
)}

54
ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.test.tsx Normal file
View file

@ -0,0 +1,54 @@
import { describe, it, expect } from "vitest";
import { render, screen } from "@testing-library/react";
import LatencyBasedConfiguration from "./LatencyBasedConfiguration";
describe("LatencyBasedConfiguration", () => {
it("should render the section heading", () => {
render(<LatencyBasedConfiguration routingStrategyArgs={{}} />);
expect(screen.getByText("Latency-Based Configuration")).toBeInTheDocument();
});
it("should render default params when no args are provided", () => {
render(<LatencyBasedConfiguration routingStrategyArgs={null as any} />);
// Default: ttl=3600, lowest_latency_buffer=0
expect(screen.getByDisplayValue("3600")).toBeInTheDocument();
expect(screen.getByDisplayValue("0")).toBeInTheDocument();
});
it("should render the provided routing strategy args as inputs", () => {
const args = { ttl: 7200, lowest_latency_buffer: 0.1 };
render(<LatencyBasedConfiguration routingStrategyArgs={args} />);
expect(screen.getByDisplayValue("7200")).toBeInTheDocument();
expect(screen.getByDisplayValue("0.1")).toBeInTheDocument();
});
it("should render an input with the correct name attribute for each param", () => {
const args = { ttl: 3600, lowest_latency_buffer: 0 };
render(<LatencyBasedConfiguration routingStrategyArgs={args} />);
expect(screen.getByRole("textbox", { name: /ttl/i })).toBeInTheDocument();
expect(screen.getByRole("textbox", { name: /lowest latency buffer/i })).toBeInTheDocument();
});
it("should display the TTL parameter explanation", () => {
render(<LatencyBasedConfiguration routingStrategyArgs={null as any} />);
expect(
screen.getByText(/sliding window to look back over/i)
).toBeInTheDocument();
});
it("should display the lowest_latency_buffer parameter explanation", () => {
render(<LatencyBasedConfiguration routingStrategyArgs={null as any} />);
expect(
screen.getByText(/shuffle between deployments within this %/i)
).toBeInTheDocument();
});
it("should render object values stringified into the input", () => {
const args = { ttl: { nested: true } };
render(<LatencyBasedConfiguration routingStrategyArgs={args} />);
// HTML input type=text strips newlines, so check that the key/value appears
const input = screen.getByRole("textbox", { name: /ttl/i }) as HTMLInputElement;
expect(input.value).toContain('"nested"');
expect(input.value).toContain('true');
});
});

79
ui/litellm-dashboard/src/components/router_settings/ReliabilityRetriesSection.test.tsx Normal file
View file

@ -0,0 +1,79 @@
import { describe, it, expect } from "vitest";
import { render, screen } from "@testing-library/react";
import ReliabilityRetriesSection from "./ReliabilityRetriesSection";
const baseSettings = {
num_retries: 3,
timeout: 30,
allowed_fails: 2,
fallbacks: ["gpt-3.5"],
context_window_fallbacks: [],
routing_strategy_args: { ttl: 3600 },
routing_strategy: "simple-shuffle",
enable_tag_filtering: false,
};
describe("ReliabilityRetriesSection", () => {
it("should render the section heading", () => {
render(<ReliabilityRetriesSection routerSettings={{}} routerFieldsMetadata={{}} />);
expect(screen.getByText("Reliability & Retries")).toBeInTheDocument();
});
it("should render input fields for non-excluded settings", () => {
render(<ReliabilityRetriesSection routerSettings={baseSettings} routerFieldsMetadata={{}} />);
expect(screen.getByDisplayValue("3")).toBeInTheDocument(); // num_retries
expect(screen.getByDisplayValue("30")).toBeInTheDocument(); // timeout
expect(screen.getByDisplayValue("2")).toBeInTheDocument(); // allowed_fails
});
it("should not render inputs for excluded keys", () => {
render(<ReliabilityRetriesSection routerSettings={baseSettings} routerFieldsMetadata={{}} />);
// Each excluded key must not produce a visible input value
const inputs = screen.queryAllByRole("textbox");
const inputNames = inputs.map((el) => el.getAttribute("name"));
expect(inputNames).not.toContain("fallbacks");
expect(inputNames).not.toContain("context_window_fallbacks");
expect(inputNames).not.toContain("routing_strategy_args");
expect(inputNames).not.toContain("routing_strategy");
expect(inputNames).not.toContain("enable_tag_filtering");
});
it("should use ui_field_name from metadata as the label", () => {
const metadata = {
num_retries: { ui_field_name: "Number of Retries", field_description: "How many times to retry" },
};
render(
<ReliabilityRetriesSection routerSettings={{ num_retries: 3 }} routerFieldsMetadata={metadata} />
);
expect(screen.getByText("Number of Retries")).toBeInTheDocument();
});
it("should fall back to the raw param name when no metadata label is available", () => {
render(
<ReliabilityRetriesSection routerSettings={{ num_retries: 3 }} routerFieldsMetadata={{}} />
);
expect(screen.getByText("num_retries")).toBeInTheDocument();
});
it("should render null values as an empty input", () => {
render(
<ReliabilityRetriesSection routerSettings={{ timeout: null }} routerFieldsMetadata={{}} />
);
const input = screen.getByRole("textbox", { name: /timeout/i }) as HTMLInputElement;
expect(input.value).toBe("");
});
it("should render object values stringified into the input", () => {
const settings = { retry_policy: { "rate-limited": 2 } };
render(<ReliabilityRetriesSection routerSettings={settings} routerFieldsMetadata={{}} />);
// HTML input type=text strips newlines, so check that the key/value appears
const input = screen.getByRole("textbox", { name: /retry_policy/i }) as HTMLInputElement;
expect(input.value).toContain('"rate-limited"');
expect(input.value).toContain('2');
});
it("should render no inputs when routerSettings is empty", () => {
render(<ReliabilityRetriesSection routerSettings={{}} routerFieldsMetadata={{}} />);
expect(screen.queryAllByRole("textbox")).toHaveLength(0);
});
});

133
ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx Normal file
View file

@ -0,0 +1,133 @@
import { describe, it, expect, vi } from "vitest";
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import RouterSettingsForm from "./RouterSettingsForm";
import type { RouterSettingsFormValue } from "./RouterSettingsForm";
// Use the same antd mock as RoutingStrategySelector to keep things consistent
vi.mock("antd", () => ({
Select: Object.assign(
({ value, onChange, children }: any) => (
<select
data-testid="strategy-select"
value={value ?? ""}
onChange={(e) => onChange(e.target.value)}
>
{children}
</select>
),
{
Option: ({ value, children }: any) => (
<option value={value}>{children}</option>
),
}
),
}));
vi.mock("@tremor/react", async (importOriginal) => {
const actual = await importOriginal<typeof import("@tremor/react")>();
return {
...actual,
Switch: ({ checked, onChange }: any) => (
<input
type="checkbox"
role="switch"
checked={checked}
onChange={(e) => onChange(e.target.checked)}
/>
),
};
});
const defaultValue: RouterSettingsFormValue = {
routerSettings: {},
selectedStrategy: null,
enableTagFiltering: false,
};
const baseProps = {
value: defaultValue,
onChange: vi.fn(),
routerFieldsMetadata: {},
availableRoutingStrategies: [],
routingStrategyDescriptions: {},
};
describe("RouterSettingsForm", () => {
it("should render", () => {
render(<RouterSettingsForm {...baseProps} />);
expect(screen.getByText("Routing Settings")).toBeInTheDocument();
});
it("should not show the strategy selector when no strategies are provided", () => {
render(<RouterSettingsForm {...baseProps} />);
expect(screen.queryByTestId("strategy-select")).not.toBeInTheDocument();
});
it("should show the strategy selector when strategies are available", () => {
const props = {
...baseProps,
availableRoutingStrategies: ["simple-shuffle", "latency-based-routing"],
};
render(<RouterSettingsForm {...props} />);
expect(screen.getByTestId("strategy-select")).toBeInTheDocument();
});
it("should not render LatencyBasedConfiguration for non-latency strategies", () => {
const props = {
...baseProps,
value: { ...defaultValue, selectedStrategy: "simple-shuffle" },
availableRoutingStrategies: ["simple-shuffle"],
};
render(<RouterSettingsForm {...props} />);
expect(screen.queryByText("Latency-Based Configuration")).not.toBeInTheDocument();
});
it("should render LatencyBasedConfiguration when strategy is latency-based-routing", () => {
const props = {
...baseProps,
value: {
...defaultValue,
selectedStrategy: "latency-based-routing",
routerSettings: { routing_strategy_args: { ttl: 3600, lowest_latency_buffer: 0 } },
},
availableRoutingStrategies: ["latency-based-routing"],
};
render(<RouterSettingsForm {...props} />);
expect(screen.getByText("Latency-Based Configuration")).toBeInTheDocument();
});
it("should call onChange with the updated strategy when the selector changes", async () => {
const onChange = vi.fn();
const user = userEvent.setup();
const props = {
...baseProps,
onChange,
availableRoutingStrategies: ["simple-shuffle", "latency-based-routing"],
};
render(<RouterSettingsForm {...props} />);
await user.selectOptions(screen.getByTestId("strategy-select"), "latency-based-routing");
expect(onChange).toHaveBeenCalledWith(
expect.objectContaining({ selectedStrategy: "latency-based-routing" })
);
});
it("should call onChange with the updated enableTagFiltering when the toggle changes", async () => {
const onChange = vi.fn();
const user = userEvent.setup();
render(<RouterSettingsForm {...baseProps} onChange={onChange} />);
await user.click(screen.getByRole("switch"));
expect(onChange).toHaveBeenCalledWith(
expect.objectContaining({ enableTagFiltering: true })
);
});
it("should show the Reliability & Retries section", () => {
render(<RouterSettingsForm {...baseProps} />);
expect(screen.getByText("Reliability & Retries")).toBeInTheDocument();
});
});

94
ui/litellm-dashboard/src/components/router_settings/RoutingStrategySelector.test.tsx Normal file
View file

@ -0,0 +1,94 @@
import { describe, it, expect, vi } from "vitest";
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import RoutingStrategySelector from "./RoutingStrategySelector";
// Ant Design's Select is complex to drive in JSDOM; swap it for a plain
// <select> so we can assert options and fire change events normally.
vi.mock("antd", () => ({
Select: Object.assign(
({ value, onChange, children }: any) => (
<div data-testid="ant-select">
<select
data-testid="strategy-select"
value={value ?? ""}
onChange={(e) => onChange(e.target.value)}
>
{children}
</select>
</div>
),
{
Option: ({ value, children }: any) => (
<option value={value}>{children}</option>
),
}
),
}));
const baseProps = {
selectedStrategy: null,
availableStrategies: ["simple-shuffle", "latency-based-routing", "least-busy"],
routingStrategyDescriptions: {
"simple-shuffle": "Randomly pick a deployment",
"latency-based-routing": "Pick the lowest-latency deployment",
},
routerFieldsMetadata: {},
onStrategyChange: vi.fn(),
};
describe("RoutingStrategySelector", () => {
it("should render", () => {
render(<RoutingStrategySelector {...baseProps} />);
expect(screen.getByTestId("ant-select")).toBeInTheDocument();
});
it("should display default label when no metadata is provided", () => {
render(<RoutingStrategySelector {...baseProps} />);
expect(screen.getByText("Routing Strategy")).toBeInTheDocument();
});
it("should display ui_field_name from metadata when provided", () => {
const props = {
...baseProps,
routerFieldsMetadata: {
routing_strategy: {
ui_field_name: "Strategy",
field_description: "How to pick a deployment",
},
},
};
render(<RoutingStrategySelector {...props} />);
expect(screen.getByText("Strategy")).toBeInTheDocument();
expect(screen.getByText("How to pick a deployment")).toBeInTheDocument();
});
it("should render all available strategies as options", () => {
render(<RoutingStrategySelector {...baseProps} />);
expect(screen.getByText("simple-shuffle")).toBeInTheDocument();
expect(screen.getByText("latency-based-routing")).toBeInTheDocument();
expect(screen.getByText("least-busy")).toBeInTheDocument();
});
it("should display strategy descriptions alongside option labels", () => {
render(<RoutingStrategySelector {...baseProps} />);
expect(screen.getByText("Randomly pick a deployment")).toBeInTheDocument();
expect(screen.getByText("Pick the lowest-latency deployment")).toBeInTheDocument();
});
it("should not render a description for a strategy that has none", () => {
render(<RoutingStrategySelector {...baseProps} />);
// "least-busy" has no entry in routingStrategyDescriptions — it still renders without crashing
expect(screen.getByText("least-busy")).toBeInTheDocument();
});
it("should call onStrategyChange with the selected strategy value", async () => {
const onStrategyChange = vi.fn();
const user = userEvent.setup();
render(<RoutingStrategySelector {...baseProps} onStrategyChange={onStrategyChange} />);
await user.selectOptions(screen.getByTestId("strategy-select"), "latency-based-routing");
expect(onStrategyChange).toHaveBeenCalledWith("latency-based-routing");
});
});

113
ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.test.tsx Normal file
View file

@ -0,0 +1,113 @@
import { describe, it, expect, vi } from "vitest";
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import TagFilteringToggle from "./TagFilteringToggle";
// setupTests.ts mocks @tremor/react but leaves Switch as the real implementation.
// Re-mock Switch as a plain checkbox so toggle interactions are trivially testable.
vi.mock("@tremor/react", async (importOriginal) => {
const actual = await importOriginal<typeof import("@tremor/react")>();
return {
...actual,
Switch: ({ checked, onChange, className }: any) => (
<input
type="checkbox"
role="switch"
checked={checked}
onChange={(e) => onChange(e.target.checked)}
className={className}
/>
),
};
});
const baseMetadata = {
enable_tag_filtering: {
ui_field_name: "Tag Filtering",
field_description: "Route requests based on tags",
link: null,
},
};
describe("TagFilteringToggle", () => {
it("should render", () => {
render(
<TagFilteringToggle enabled={false} routerFieldsMetadata={{}} onToggle={vi.fn()} />
);
expect(screen.getByRole("switch")).toBeInTheDocument();
});
it("should display default label when no metadata is provided", () => {
render(
<TagFilteringToggle enabled={false} routerFieldsMetadata={{}} onToggle={vi.fn()} />
);
expect(screen.getByText("Enable Tag Filtering")).toBeInTheDocument();
});
it("should display the label from metadata when provided", () => {
render(
<TagFilteringToggle
enabled={false}
routerFieldsMetadata={baseMetadata}
onToggle={vi.fn()}
/>
);
expect(screen.getByText("Tag Filtering")).toBeInTheDocument();
});
it("should display the description from metadata", () => {
render(
<TagFilteringToggle
enabled={false}
routerFieldsMetadata={baseMetadata}
onToggle={vi.fn()}
/>
);
expect(screen.getByText("Route requests based on tags")).toBeInTheDocument();
});
it("should render a Learn more link when metadata provides one", () => {
const metadata = {
enable_tag_filtering: {
...baseMetadata.enable_tag_filtering,
link: "https://docs.example.com/tag-filtering",
},
};
render(
<TagFilteringToggle enabled={false} routerFieldsMetadata={metadata} onToggle={vi.fn()} />
);
const link = screen.getByRole("link", { name: /learn more/i });
expect(link).toBeInTheDocument();
expect(link).toHaveAttribute("href", "https://docs.example.com/tag-filtering");
});
it("should not render a Learn more link when metadata has no link", () => {
render(
<TagFilteringToggle
enabled={false}
routerFieldsMetadata={baseMetadata}
onToggle={vi.fn()}
/>
);
expect(screen.queryByRole("link", { name: /learn more/i })).not.toBeInTheDocument();
});
it("should reflect the enabled=true state on the switch", () => {
render(
<TagFilteringToggle enabled={true} routerFieldsMetadata={{}} onToggle={vi.fn()} />
);
expect(screen.getByRole("switch")).toBeChecked();
});
it("should call onToggle with the new value when the switch is toggled", async () => {
const onToggle = vi.fn();
const user = userEvent.setup();
render(
<TagFilteringToggle enabled={false} routerFieldsMetadata={{}} onToggle={onToggle} />
);
await user.click(screen.getByRole("switch"));
expect(onToggle).toHaveBeenCalledWith(true);
});
});

176
ui/litellm-dashboard/src/components/router_settings/index.test.tsx Normal file
View file

@ -0,0 +1,176 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils";
import userEvent from "@testing-library/user-event";
import RouterSettings from "./index";
vi.mock("antd", () => ({
Select: Object.assign(
({ value, onChange, children }: any) => (
<select
data-testid="strategy-select"
value={value ?? ""}
onChange={(e) => onChange(e.target.value)}
>
{children}
</select>
),
{
Option: ({ value, children }: any) => (
<option value={value}>{children}</option>
),
}
),
}));
vi.mock("@tremor/react", async (importOriginal) => {
const actual = await importOriginal<typeof import("@tremor/react")>();
return {
...actual,
Switch: ({ checked, onChange }: any) => (
<input
type="checkbox"
role="switch"
checked={checked}
onChange={(e) => onChange(e.target.checked)}
/>
),
};
});
vi.mock("@/components/networking", () => ({
getCallbacksCall: vi.fn(),
getRouterSettingsCall: vi.fn(),
setCallbacksCall: vi.fn(),
}));
import {
getCallbacksCall,
getRouterSettingsCall,
setCallbacksCall,
} from "@/components/networking";
import NotificationsManager from "@/components/molecules/notifications_manager";
const mockCallbacksResponse = {
router_settings: {
routing_strategy: "simple-shuffle",
num_retries: 3,
timeout: 30,
},
};
const mockRouterSettingsResponse = {
fields: [
{
field_name: "routing_strategy",
ui_field_name: "Routing Strategy",
field_description: "How requests are distributed",
options: ["simple-shuffle", "latency-based-routing"],
link: null,
},
{
field_name: "enable_tag_filtering",
ui_field_name: "Tag Filtering",
field_description: "Route by tag",
field_value: false,
link: null,
},
],
routing_strategy_descriptions: {
"simple-shuffle": "Randomly pick a deployment",
"latency-based-routing": "Pick the lowest-latency deployment",
},
};
const defaultProps = {
accessToken: "test-token",
userRole: "Admin",
userID: "user-1",
modelData: null,
};
describe("RouterSettings", () => {
beforeEach(() => {
vi.clearAllMocks();
vi.mocked(getCallbacksCall).mockResolvedValue(mockCallbacksResponse);
vi.mocked(getRouterSettingsCall).mockResolvedValue(mockRouterSettingsResponse);
vi.mocked(setCallbacksCall).mockResolvedValue({});
});
it("should render nothing when accessToken is null", () => {
const { container } = renderWithProviders(
<RouterSettings {...defaultProps} accessToken={null} />
);
expect(container).toBeEmptyDOMElement();
});
it("should render the Save Changes and Reset buttons when authenticated", () => {
renderWithProviders(<RouterSettings {...defaultProps} />);
expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument();
expect(screen.getByRole("button", { name: /reset/i })).toBeInTheDocument();
});
it("should fetch callbacks and router settings on mount", async () => {
renderWithProviders(<RouterSettings {...defaultProps} />);
await waitFor(() => {
expect(getCallbacksCall).toHaveBeenCalledWith("test-token", "user-1", "Admin");
});
expect(getRouterSettingsCall).toHaveBeenCalledWith("test-token");
});
it("should not fetch data when any required prop is missing", () => {
renderWithProviders(
<RouterSettings {...defaultProps} userRole={null} />
);
expect(getCallbacksCall).not.toHaveBeenCalled();
});
it("should render routing strategies loaded from the API", async () => {
renderWithProviders(<RouterSettings {...defaultProps} />);
await waitFor(() => {
expect(screen.getByTestId("strategy-select")).toBeInTheDocument();
});
const select = screen.getByTestId("strategy-select") as HTMLSelectElement;
const optionValues = Array.from(select.options).map((o) => o.value);
expect(optionValues).toContain("simple-shuffle");
expect(optionValues).toContain("latency-based-routing");
});
it("should call setCallbacksCall with updated settings on Save Changes", async () => {
const user = userEvent.setup();
renderWithProviders(<RouterSettings {...defaultProps} />);
// Wait for the strategy select to appear — it only renders after getRouterSettingsCall resolves
await waitFor(() => {
expect(screen.getByTestId("strategy-select")).toBeInTheDocument();
});
await user.click(screen.getByRole("button", { name: /save changes/i }));
expect(setCallbacksCall).toHaveBeenCalledWith(
"test-token",
expect.objectContaining({
router_settings: expect.objectContaining({
routing_strategy: "simple-shuffle",
}),
})
);
});
it("should show a success notification after saving", async () => {
const user = userEvent.setup();
renderWithProviders(<RouterSettings {...defaultProps} />);
// Wait for data to load before interacting
await waitFor(() => {
expect(screen.getByTestId("strategy-select")).toBeInTheDocument();
});
await user.click(screen.getByRole("button", { name: /save changes/i }));
expect(NotificationsManager.success).toHaveBeenCalledWith(
"router settings updated successfully"
);
});
});

53
ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx Normal file
View file

@ -0,0 +1,53 @@
import React, { useState } from "react";
import { CopyToClipboard } from "react-copy-to-clipboard";
import { Button, message } from "antd";
interface CreatedKeyDisplayProps {
apiKey: string;
}
/**
* Shared component for displaying a newly-created virtual key.
* Used on the Virtual Keys page and in the Add Agent wizard.
*/
const CreatedKeyDisplay: React.FC<CreatedKeyDisplayProps> = ({ apiKey }) => {
const [copied, setCopied] = useState(false);
const handleCopy = () => {
setCopied(true);
message.success("Key copied to clipboard");
setTimeout(() => setCopied(false), 2000);
};
return (
<div>
<p className="mb-2">
Please save this secret key somewhere safe and accessible. For security reasons,{" "}
<b>you will not be able to view it again</b> through your LiteLLM account. If you
lose this secret key, you will need to generate a new one.
</p>
<p className="text-sm text-gray-600 mt-3 mb-1">Virtual Key:</p>
<div
style={{
background: "#f8f8f8",
padding: "10px",
borderRadius: "5px",
marginBottom: "10px",
}}
>
<pre style={{ wordWrap: "break-word", whiteSpace: "normal", margin: 0 }}>
{apiKey}
</pre>
</div>
<CopyToClipboard text={apiKey} onCopy={handleCopy}>
<Button type="primary" style={{ marginTop: 12 }}>
{copied ? "Copied!" : "Copy Virtual Key"}
</Button>
</CopyToClipboard>
</div>
);
};
export default CreatedKeyDisplay;

Some files were not shown because too many files have changed in this diff Show more