+### OAuth Configuration & Overrides
+
+LiteLLM attempts [OAuth 2.0 Authorization Server Discovery](https://datatracker.ietf.org/doc/html/rfc8414) by default. When you create an MCP server in the UI and set `Authentication: OAuth`, LiteLLM will locate the provider metadata, dynamically register a client, and perform PKCE-based authorization without you providing any additional details.
+
+**Customize the OAuth flow when needed:**
+
+
+
+- **Provide explicit client credentials** – If the MCP provider does not offer dynamic client registration or you prefer to manage the client yourself, fill in `client_id`, `client_secret`, and the desired `scopes`.
+- **Override discovery URLs** – In some environments, LiteLLM might not be able to reach the provider's metadata endpoints. Use the optional `authorization_url`, `token_url`, and `registration_url` fields to point LiteLLM directly to the correct endpoints.
+
+
+
### Static Headers
Sometimes your MCP server needs specific headers on every request. Maybe it's an API key, maybe it's a custom header the server expects. Instead of configuring auth, you can just set them directly.
diff --git a/docs/my-website/docs/observability/generic_api.md b/docs/my-website/docs/observability/generic_api.md
index 2d1a24c317b..93a0762591a 100644
--- a/docs/my-website/docs/observability/generic_api.md
+++ b/docs/my-website/docs/observability/generic_api.md
@@ -47,6 +47,7 @@ callback_settings:
| `endpoint` | string | Yes | HTTP endpoint to send logs to |
| `headers` | dict | No | Custom headers for the request |
| `event_types` | list | No | Filter events: `llm_api_success`, `llm_api_failure`. Defaults to all events. |
+| `log_format` | string | No | Output format: `json_array` (default), `ndjson`, or `single`. Controls how logs are batched and sent. |
## Pre-configured Callbacks
@@ -107,4 +108,62 @@ callback_settings:
flush_interval: 60 # seconds, default: 60
```
+## Log Format Options
+
+Control how logs are formatted and sent to your endpoint.
+
+### JSON Array (Default)
+
+```yaml
+callback_settings:
+ my_api:
+ callback_type: generic_api
+ endpoint: https://your-endpoint.com
+ log_format: json_array # default if not specified
+```
+
+Sends all logs in a batch as a single JSON array `[{log1}, {log2}, ...]`. This is the default behavior and maintains backward compatibility.
+
+**When to use**: Most HTTP endpoints expecting batched JSON data.
+
+### NDJSON (Newline-Delimited JSON)
+
+```yaml
+callback_settings:
+ my_api:
+ callback_type: generic_api
+ endpoint: https://your-endpoint.com
+ log_format: ndjson
+```
+
+Sends logs as newline-delimited JSON (one record per line):
+```
+{log1}
+{log2}
+{log3}
+```
+
+**When to use**: Log aggregation services like Sumo Logic, Splunk, or Datadog that support field extraction on individual records.
+
+**Benefits**:
+- Each log is ingested as a separate message
+- Field Extraction Rules work at ingest time
+- Better parsing and querying performance
+
+### Single
+
+```yaml
+callback_settings:
+ my_api:
+ callback_type: generic_api
+ endpoint: https://your-endpoint.com
+ log_format: single
+```
+
+Sends each log as an individual HTTP request in parallel when the batch is flushed.
+
+**When to use**: Endpoints that expect individual records, or when you need maximum compatibility.
+
+**Note**: This mode sends N HTTP requests per batch (more overhead). Consider using `ndjson` instead if your endpoint supports it.
+
diff --git a/docs/my-website/docs/observability/levo_integration.md b/docs/my-website/docs/observability/levo_integration.md
new file mode 100644
index 00000000000..3e46cf6b921
--- /dev/null
+++ b/docs/my-website/docs/observability/levo_integration.md
@@ -0,0 +1,162 @@
+---
+sidebar_label: Levo AI
+---
+
+import Image from '@theme/IdealImage';
+import Tabs from '@theme/Tabs';
+import TabItem from '@theme/TabItem';
+
+# Levo AI
+
+
+
+
+
+
+
+
+
+
+[Levo](https://levo.ai/) is an AI observability and compliance platform that provides comprehensive monitoring, analysis, and compliance tracking for LLM applications.
+
+## Quick Start
+
+Send all your LLM requests and responses to Levo for monitoring and analysis using LiteLLM's built-in Levo integration.
+
+### What You'll Get
+
+- **Complete visibility** into all LLM API calls across all providers
+- **Request and response data** including prompts, completions, and metadata
+- **Usage and cost tracking** with token counts and cost breakdowns
+- **Error monitoring** and performance metrics
+- **Compliance tracking** for audit and governance
+
+### Setup Steps
+
+**1. Install OpenTelemetry dependencies:**
+
+```bash
+pip install opentelemetry-api opentelemetry-sdk opentelemetry-exporter-otlp-proto-http opentelemetry-exporter-otlp-proto-grpc
+```
+
+**2. Enable Levo callback in your LiteLLM config:**
+
+Add to your `litellm_config.yaml`:
+
+```yaml
+litellm_settings:
+ callbacks: ["levo"]
+```
+
+**3. Configure environment variables:**
+
+[Contact Levo support](mailto:support@levo.ai) to get your collector endpoint URL, API key, organization ID, and workspace ID.
+
+Set these required environment variables:
+
+```bash
+export LEVOAI_API_KEY=""
+export LEVOAI_ORG_ID=""
+export LEVOAI_WORKSPACE_ID=""
+export LEVOAI_COLLECTOR_URL=""
+```
+
+**Note:** The collector URL should be the full endpoint URL provided by Levo support. It will be used exactly as provided.
+
+**4. Start LiteLLM:**
+
+```bash
+litellm --config config.yaml
+```
+
+**5. Make requests - they'll automatically be sent to Levo!**
+
+```bash
+curl --location 'http://0.0.0.0:4000/chat/completions' \
+ --header 'Content-Type: application/json' \
+ --data '{
+ "model": "gpt-3.5-turbo",
+ "messages": [
+ {
+ "role": "user",
+ "content": "Hello, this is a test message"
+ }
+ ]
+ }'
+```
+
+## What Data is Captured
+
+| Feature | Details |
+|---------|---------|
+| **What is logged** | OpenTelemetry Trace Data (OTLP format) |
+| **Events** | Success + Failure |
+| **Format** | OTLP (OpenTelemetry Protocol) |
+| **Headers** | Automatically includes `Authorization: Bearer {LEVOAI_API_KEY}`, `x-levo-organization-id`, and `x-levo-workspace-id` |
+
+## Configuration Reference
+
+### Required Environment Variables
+
+| Variable | Description | Example |
+|----------|-------------|---------|
+| `LEVOAI_API_KEY` | Your Levo API key | `levo_abc123...` |
+| `LEVOAI_ORG_ID` | Your Levo organization ID | `org-123456` |
+| `LEVOAI_WORKSPACE_ID` | Your Levo workspace ID | `workspace-789` |
+| `LEVOAI_COLLECTOR_URL` | Full collector endpoint URL from Levo support | `https://collector.levo.ai/v1/traces` |
+
+### Optional Environment Variables
+
+| Variable | Description | Default |
+|----------|-------------|---------|
+| `LEVOAI_ENV_NAME` | Environment name for tagging traces | `None` |
+
+**Note:** The collector URL is used exactly as provided by Levo support. No path manipulation is performed.
+
+## Troubleshooting
+
+### Not seeing traces in Levo?
+
+1. **Verify Levo callback is enabled**: Check LiteLLM startup logs for `initializing callbacks=['levo']`
+
+2. **Check required environment variables**: Ensure all required variables are set:
+ ```bash
+ echo $LEVOAI_API_KEY
+ echo $LEVOAI_ORG_ID
+ echo $LEVOAI_WORKSPACE_ID
+ echo $LEVOAI_COLLECTOR_URL
+ ```
+
+3. **Verify collector connectivity**: Test if your collector is reachable:
+ ```bash
+ curl /health
+ ```
+
+4. **Check for initialization errors**: Look for errors in LiteLLM startup logs. Common issues:
+ - Missing OpenTelemetry packages: Install with `pip install opentelemetry-api opentelemetry-sdk opentelemetry-exporter-otlp-proto-http opentelemetry-exporter-otlp-proto-grpc`
+ - Missing required environment variables: All four required variables must be set
+ - Invalid collector URL: Ensure the URL is correct and reachable
+
+5. **Enable debug logging**:
+ ```bash
+ export LITELLM_LOG="DEBUG"
+ ```
+
+6. **Wait for async export**: OTLP sends traces asynchronously. Wait 10-15 seconds after making requests before checking Levo.
+
+### Common Errors
+
+**Error: "LEVOAI_COLLECTOR_URL environment variable is required"**
+- Solution: Set the `LEVOAI_COLLECTOR_URL` environment variable with your collector endpoint URL from Levo support.
+
+**Error: "No module named 'opentelemetry'"**
+- Solution: Install OpenTelemetry packages: `pip install opentelemetry-api opentelemetry-sdk opentelemetry-exporter-otlp-proto-http opentelemetry-exporter-otlp-proto-grpc`
+
+## Additional Resources
+
+- [Levo Documentation](https://docs.levo.ai)
+- [OpenTelemetry Specification](https://opentelemetry.io/docs/specs/otel/)
+
+## Need Help?
+
+For issues or questions about the Levo integration with LiteLLM, please [contact Levo support](mailto:support@levo.ai) or open an issue on the [LiteLLM GitHub repository](https://github.com/BerriAI/litellm/issues).
diff --git a/docs/my-website/docs/observability/opentelemetry_integration.md b/docs/my-website/docs/observability/opentelemetry_integration.md
index 2b3cf1313ba..b6eff231620 100644
--- a/docs/my-website/docs/observability/opentelemetry_integration.md
+++ b/docs/my-website/docs/observability/opentelemetry_integration.md
@@ -4,7 +4,7 @@ import TabItem from '@theme/TabItem';
# OpenTelemetry - Tracing LLMs with any observability tool
-OpenTelemetry is a CNCF standard for observability. It connects to any observability tool, such as Jaeger, Zipkin, Datadog, New Relic, Traceloop and others.
+OpenTelemetry is a CNCF standard for observability. It connects to any observability tool, such as Jaeger, Zipkin, Datadog, New Relic, Traceloop, Levo AI and others.
@@ -12,7 +12,9 @@ OpenTelemetry is a CNCF standard for observability. It connects to any observabi
From v1.81.0, the request/response will be set as attributes on the parent "Received Proxy Server Request" span by default. This allows you to see the request/response in the parent span in your observability tool.
-To use the older behavior with nested "litellm_request" spans, set the following environment variable:
+**Note:** When making multiple LLM calls within an external OTEL span context, the last call's attributes will overwrite previous calls' attributes on the parent span.
+
+To use the older behavior with nested "litellm_request" spans (which creates separate spans for each call), set the following environment variable:
```shell
USE_OTEL_LITELLM_REQUEST_SPAN=true
diff --git a/docs/my-website/docs/observability/sumologic_integration.md b/docs/my-website/docs/observability/sumologic_integration.md
index d0894146e4c..c30ee94dad4 100644
--- a/docs/my-website/docs/observability/sumologic_integration.md
+++ b/docs/my-website/docs/observability/sumologic_integration.md
@@ -148,6 +148,51 @@ Example payload:
## Advanced Configuration
+### Log Format
+
+The Sumo Logic integration uses **NDJSON (newline-delimited JSON)** format by default. This format is optimal for Sumo Logic's parsing capabilities and allows Field Extraction Rules to work at ingest time.
+
+#### NDJSON Format
+
+Each log entry is sent as a separate line in the HTTP request:
+```
+{"id":"chatcmpl-1","model":"gpt-3.5-turbo","response_cost":0.0001,...}
+{"id":"chatcmpl-2","model":"gpt-4","response_cost":0.0003,...}
+{"id":"chatcmpl-3","model":"gpt-3.5-turbo","response_cost":0.0001,...}
+```
+
+#### Benefits for Field Extraction Rules (FERs)
+
+With NDJSON format, you can create Field Extraction Rules directly:
+
+```
+_sourceCategory=litellm/logs
+| json field=_raw "model", "response_cost", "user" as model, cost, user
+```
+
+**Before NDJSON** (with JSON array format):
+- Required `parse regex ... multi` workaround
+- FERs couldn't parse at ingest time
+- Query-time parsing impacted dashboard performance
+
+**After NDJSON**:
+- ✅ FERs parse fields at ingest time
+- ✅ No query-time workarounds needed
+- ✅ Better dashboard performance
+- ✅ Simpler query syntax
+
+#### Changing the Log Format (Advanced)
+
+If you need to change the log format (not recommended for Sumo Logic):
+
+```yaml
+callback_settings:
+ sumologic:
+ callback_type: generic_api
+ callback_name: sumologic
+ log_format: json_array # Override to use JSON array instead
+```
+
### Batching Settings
Control how LiteLLM batches logs before sending to Sumo Logic:
diff --git a/docs/my-website/docs/providers/anthropic.md b/docs/my-website/docs/providers/anthropic.md
index bcfb698a0f8..cae8657f1a0 100644
--- a/docs/my-website/docs/providers/anthropic.md
+++ b/docs/my-website/docs/providers/anthropic.md
@@ -444,7 +444,7 @@ Here's what a sample Raw Request from LiteLLM for Anthropic Context Caching look
POST Request Sent from LiteLLM:
curl -X POST \
https://api.anthropic.com/v1/messages \
--H 'accept: application/json' -H 'anthropic-version: 2023-06-01' -H 'content-type: application/json' -H 'x-api-key: sk-...' -H 'anthropic-beta: prompt-caching-2024-07-31' \
+-H 'accept: application/json' -H 'anthropic-version: 2023-06-01' -H 'content-type: application/json' -H 'x-api-key: sk-...' \
-d '{'model': 'claude-3-5-sonnet-20240620', [
{
"role": "user",
@@ -472,6 +472,8 @@ https://api.anthropic.com/v1/messages \
"max_tokens": 10
}'
```
+
+**Note:** Anthropic no longer requires the `anthropic-beta: prompt-caching-2024-07-31` header. Prompt caching now works automatically when you use `cache_control` in your messages.
:::
### Caching - Large Context Caching
diff --git a/docs/my-website/docs/providers/bedrock.md b/docs/my-website/docs/providers/bedrock.md
index 122554fe8a4..f1eed4b4d52 100644
--- a/docs/my-website/docs/providers/bedrock.md
+++ b/docs/my-website/docs/providers/bedrock.md
@@ -2208,6 +2208,53 @@ response = completion(
| `aws_role_name` | `RoleArn` | The Amazon Resource Name (ARN) of the role to assume | [AssumeRole API](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts.html#STS.Client.assume_role) |
| `aws_session_name` | `RoleSessionName` | An identifier for the assumed role session | [AssumeRole API](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts.html#STS.Client.assume_role) |
+### IAM Roles Anywhere (On-Premise / External Workloads)
+
+[IAM Roles Anywhere](https://docs.aws.amazon.com/rolesanywhere/latest/userguide/introduction.html) extends IAM roles to workloads **outside of AWS** (on-premise servers, edge devices, other clouds). It uses the same STS mechanism as regular IAM roles but authenticates via X.509 certificates instead of AWS credentials.
+
+**Setup**: Configure the [AWS Signing Helper](https://docs.aws.amazon.com/rolesanywhere/latest/userguide/credential-helper.html) as a credential process in `~/.aws/config`:
+
+```ini
+[profile litellm-roles-anywhere]
+credential_process = aws_signing_helper credential-process \
+ --certificate /path/to/certificate.pem \
+ --private-key /path/to/private-key.pem \
+ --trust-anchor-arn arn:aws:rolesanywhere:us-east-1:123456789012:trust-anchor/abc123 \
+ --profile-arn arn:aws:rolesanywhere:us-east-1:123456789012:profile/def456 \
+ --role-arn arn:aws:iam::123456789012:role/MyBedrockRole
+```
+
+**Usage**: Reference the profile in LiteLLM:
+
+
+
+
+```python
+from litellm import completion
+
+response = completion(
+ model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
+ messages=[{"role": "user", "content": "Hello!"}],
+ aws_profile_name="litellm-roles-anywhere",
+)
+```
+
+
+
+
+```yaml
+model_list:
+ - model_name: bedrock-claude
+ litellm_params:
+ model: bedrock/anthropic.claude-3-sonnet-20240229-v1:0
+ aws_profile_name: "litellm-roles-anywhere"
+```
+
+
+
+
+See the [IAM Roles Anywhere Getting Started Guide](https://docs.aws.amazon.com/rolesanywhere/latest/userguide/getting-started.html) for trust anchor and profile setup.
+
Make the bedrock completion call
diff --git a/docs/my-website/docs/providers/bedrock_agentcore.md b/docs/my-website/docs/providers/bedrock_agentcore.md
index 43df7f82519..e3e352f7ab6 100644
--- a/docs/my-website/docs/providers/bedrock_agentcore.md
+++ b/docs/my-website/docs/providers/bedrock_agentcore.md
@@ -11,6 +11,12 @@ Call Bedrock AgentCore in the OpenAI Request/Response format.
| Provider Route on LiteLLM | `bedrock/agentcore/{AGENT_RUNTIME_ARN}` |
| Provider Doc | [AWS Bedrock AgentCore ↗](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agentcore_InvokeAgentRuntime.html) |
+:::info
+
+This documentation is for **AgentCore Agents** (agent runtimes). If you want to use AgentCore MCP servers, add them as you would any other MCP server. See the [MCP documentation](https://docs.litellm.ai/docs/mcp) for details.
+
+:::
+
## Quick Start
### Model Format to LiteLLM
diff --git a/docs/my-website/docs/providers/zai.md b/docs/my-website/docs/providers/zai.md
index 5055d0c1cdd..937ccd67680 100644
--- a/docs/my-website/docs/providers/zai.md
+++ b/docs/my-website/docs/providers/zai.md
@@ -19,7 +19,7 @@ import os
os.environ['ZAI_API_KEY'] = ""
response = completion(
- model="zai/glm-4.6",
+ model="zai/glm-4.7",
messages=[
{"role": "user", "content": "hello from litellm"}
],
@@ -34,7 +34,7 @@ import os
os.environ['ZAI_API_KEY'] = ""
response = completion(
- model="zai/glm-4.6",
+ model="zai/glm-4.7",
messages=[
{"role": "user", "content": "hello from litellm"}
],
@@ -51,7 +51,8 @@ We support ALL Z.AI GLM models, just set `zai/` as a prefix when sending complet
| Model Name | Function Call | Notes |
|------------|---------------|-------|
-| glm-4.6 | `completion(model="zai/glm-4.6", messages)` | Latest flagship model, 200K context |
+| glm-4.7 | `completion(model="zai/glm-4.7", messages)` | **Latest flagship**, 200K context, **Reasoning** |
+| glm-4.6 | `completion(model="zai/glm-4.6", messages)` | 200K context |
| glm-4.5 | `completion(model="zai/glm-4.5", messages)` | 128K context |
| glm-4.5v | `completion(model="zai/glm-4.5v", messages)` | Vision model |
| glm-4.5-x | `completion(model="zai/glm-4.5-x", messages)` | Premium tier |
@@ -62,16 +63,17 @@ We support ALL Z.AI GLM models, just set `zai/` as a prefix when sending complet
## Model Pricing
-| Model | Input ($/1M tokens) | Output ($/1M tokens) | Context Window |
-|-------|---------------------|----------------------|----------------|
-| glm-4.6 | $0.60 | $2.20 | 200K |
-| glm-4.5 | $0.60 | $2.20 | 128K |
-| glm-4.5v | $0.60 | $1.80 | 128K |
-| glm-4.5-x | $2.20 | $8.90 | 128K |
-| glm-4.5-air | $0.20 | $1.10 | 128K |
-| glm-4.5-airx | $1.10 | $4.50 | 128K |
-| glm-4-32b-0414-128k | $0.10 | $0.10 | 128K |
-| glm-4.5-flash | **FREE** | **FREE** | 128K |
+| Model | Input ($/1M tokens) | Output ($/1M tokens) | Cached Input ($/1M tokens) | Context Window |
+|-------|---------------------|----------------------|---------------------------|----------------|
+| glm-4.7 | $0.60 | $2.20 | $0.11 | 200K |
+| glm-4.6 | $0.60 | $2.20 | - | 200K |
+| glm-4.5 | $0.60 | $2.20 | - | 128K |
+| glm-4.5v | $0.60 | $1.80 | - | 128K |
+| glm-4.5-x | $2.20 | $8.90 | - | 128K |
+| glm-4.5-air | $0.20 | $1.10 | - | 128K |
+| glm-4.5-airx | $1.10 | $4.50 | - | 128K |
+| glm-4-32b-0414-128k | $0.10 | $0.10 | - | 128K |
+| glm-4.5-flash | **FREE** | **FREE** | - | 128K |
## Using with LiteLLM Proxy
@@ -84,7 +86,7 @@ import os
os.environ['ZAI_API_KEY'] = ""
response = completion(
- model="zai/glm-4.6",
+ model="zai/glm-4.7",
messages=[{"role": "user", "content": "Hello, how are you?"}],
)
@@ -98,9 +100,9 @@ print(response.choices[0].message.content)
```yaml
model_list:
- - model_name: glm-4.6
+ - model_name: glm-4.7
litellm_params:
- model: zai/glm-4.6
+ model: zai/glm-4.7
api_key: os.environ/ZAI_API_KEY
- model_name: glm-4.5-flash # Free tier
litellm_params:
@@ -121,7 +123,7 @@ curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{
- "model": "glm-4.6",
+ "model": "glm-4.7",
"messages": [
{
"role": "user",
diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md
index 343cbd0e53f..87064d442ae 100644
--- a/docs/my-website/docs/proxy/config_settings.md
+++ b/docs/my-website/docs/proxy/config_settings.md
@@ -464,6 +464,9 @@ router_settings:
| DATABASE_USER | Username for database connection
| DATABASE_USERNAME | Alias for database user
| DATABRICKS_API_BASE | Base URL for Databricks API
+| DATABRICKS_CLIENT_ID | Client ID for Databricks OAuth M2M authentication (Service Principal application ID)
+| DATABRICKS_CLIENT_SECRET | Client secret for Databricks OAuth M2M authentication
+| DATABRICKS_USER_AGENT | Custom user agent string for Databricks API requests. Used for partner telemetry attribution
| DAYS_IN_A_MONTH | Days in a month for calculation purposes. Default is 28
| DAYS_IN_A_WEEK | Days in a week for calculation purposes. Default is 7
| DAYS_IN_A_YEAR | Days in a year for calculation purposes. Default is 365
@@ -708,6 +711,7 @@ router_settings:
| LITELLM_SSL_CIPHERS | SSL/TLS cipher configuration for faster handshakes. Controls cipher suite preferences for OpenSSL connections.
| LITELLM_SECRET_AWS_KMS_LITELLM_LICENSE | AWS KMS encrypted license for LiteLLM
| LITELLM_TOKEN | Access token for LiteLLM integration
+| LITELLM_USER_AGENT | Custom user agent string for LiteLLM API requests. Used for partner telemetry attribution
| LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD | If true, prints the standard logging payload to the console - useful for debugging
| LITELM_ENVIRONMENT | Environment for LiteLLM Instance. This is currently only logged to DeepEval to determine the environment for DeepEval integration.
| LOGFIRE_TOKEN | Token for Logfire logging service
diff --git a/docs/my-website/docs/proxy/configs.md b/docs/my-website/docs/proxy/configs.md
index ba4ca190aa9..bc2f6a13362 100644
--- a/docs/my-website/docs/proxy/configs.md
+++ b/docs/my-website/docs/proxy/configs.md
@@ -116,7 +116,7 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
"role": "user",
"content": "what llm are you"
}
- ],
+ ]
}
'
```
diff --git a/docs/my-website/docs/proxy/token_auth.md b/docs/my-website/docs/proxy/token_auth.md
index fe928a596cf..78cd144d56d 100644
--- a/docs/my-website/docs/proxy/token_auth.md
+++ b/docs/my-website/docs/proxy/token_auth.md
@@ -114,6 +114,189 @@ Set `JWT_PUBLIC_KEY_URL` in your environment to a comma-separated list of URLs f
export JWT_PUBLIC_KEY_URL="https://demo.duendesoftware.com/.well-known/openid-configuration/jwks,https://accounts.google.com/.well-known/openid-configuration/jwks"
```
+### Kubernetes ServiceAccount Authentication
+
+Use Kubernetes ServiceAccount tokens to authenticate workloads running in your cluster. This is useful when you want pods to authenticate to LiteLLM using their native Kubernetes identity.
+
+#### Prerequisites
+
+1. Your Kubernetes cluster must have ServiceAccount token projection enabled (default in Kubernetes 1.20+)
+2. Your cluster's OIDC issuer must be accessible (for EKS, GKE, AKS this is automatic)
+
+#### Step 1: Configure the OIDC Discovery URL
+
+Set `JWT_PUBLIC_KEY_URL` to your cluster's OIDC discovery endpoint:
+
+
+
+
+```bash
+# Get your EKS OIDC issuer URL
+aws eks describe-cluster --name --query "cluster.identity.oidc.issuer" --output text
+
+# Set the JWKS URL (append /keys to the issuer URL)
+export JWT_PUBLIC_KEY_URL="https://oidc.eks..amazonaws.com/id//keys"
+```
+
+
+
+
+```bash
+# GKE uses Google's OIDC provider
+export JWT_PUBLIC_KEY_URL="https://container.googleapis.com/v1/projects//locations//clusters//jwks"
+```
+
+
+
+
+```bash
+# Get your AKS OIDC issuer URL
+az aks show --name --resource-group --query "oidcIssuerProfile.issuerUrl" -o tsv
+
+# Set the JWKS URL
+export JWT_PUBLIC_KEY_URL="/openid/v1/jwks"
+```
+
+
+
+
+```bash
+# For self-managed clusters, check your API server's --service-account-issuer flag
+# The JWKS endpoint is typically at:
+export JWT_PUBLIC_KEY_URL="https:///openid/v1/jwks"
+```
+
+
+
+
+#### Step 2: Configure LiteLLM
+
+Configure LiteLLM to extract identity information from Kubernetes ServiceAccount tokens:
+
+```yaml
+general_settings:
+ enable_jwt_auth: True
+ litellm_jwtauth:
+ # Use namespace as team identifier (resolves via team_alias in DB)
+ team_alias_jwt_field: "kubernetes\.io.namespace"
+```
+
+#### Step 3: Create ServiceAccount and Configure Pod
+
+Create a ServiceAccount with an associated secret and configure your pod to use the token:
+
+```yaml
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+ name: my-llm-client
+ namespace: my-app
+---
+apiVersion: v1
+kind: Secret
+metadata:
+ name: my-llm-client-token
+ namespace: my-app
+ annotations:
+ kubernetes.io/service-account.name: my-llm-client
+type: kubernetes.io/service-account-token
+---
+apiVersion: v1
+kind: Pod
+metadata:
+ name: llm-client-pod
+ namespace: my-app
+spec:
+ serviceAccountName: my-llm-client
+ containers:
+ - name: app
+ image: my-app:latest
+ env:
+ - name: LITELLM_TOKEN
+ valueFrom:
+ secretKeyRef:
+ name: my-llm-client-token
+ key: token
+```
+
+Set the expected audience in LiteLLM:
+
+```bash
+export JWT_AUDIENCE="https://kubernetes.default.svc"
+```
+
+#### Step 4: Create Team for Namespace
+
+Create a team in LiteLLM that matches the namespace (using `team_alias`):
+
+```bash
+curl -X POST 'http://0.0.0.0:4000/team/new' \
+-H 'Authorization: Bearer ' \
+-H 'Content-Type: application/json' \
+-d '{
+ "team_alias": "my-app",
+ "team_id": "my-app",
+ "models": ["gpt-4", "claude-sonnet-4-20250514"]
+}'
+```
+
+#### Step 5: Use the Token
+
+From within the pod, the token is available in the `LITELLM_TOKEN` environment variable:
+
+```bash
+# Make a request to LiteLLM using the env var
+curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
+-H 'Content-Type: application/json' \
+-H "Authorization: Bearer $LITELLM_TOKEN" \
+-d '{
+ "model": "gpt-4",
+ "messages": [{"role": "user", "content": "Hello!"}]
+}'
+```
+
+#### Example: ServiceAccount Token Structure
+
+A Kubernetes ServiceAccount token looks like this:
+
+```json
+{
+ "aud": ["litellm-proxy"],
+ "exp": 1234567890,
+ "iat": 1234567890,
+ "iss": "https://oidc.eks.us-west-2.amazonaws.com/id/EXAMPLE",
+ "kubernetes.io": {
+ "namespace": "my-app",
+ "pod": {
+ "name": "llm-client-pod",
+ "uid": "pod-uid"
+ },
+ "serviceaccount": {
+ "name": "my-llm-client",
+ "uid": "sa-uid"
+ }
+ },
+ "nbf": 1234567890,
+ "sub": "system:serviceaccount:my-app:my-llm-client"
+}
+```
+
+#### Advanced: Map Namespace to Team Using Name Resolution
+
+Use the `team_alias_jwt_field` to automatically resolve namespaces to teams:
+
+```yaml
+general_settings:
+ enable_jwt_auth: True
+ litellm_jwtauth:
+ user_id_jwt_field: "sub"
+ # Map the namespace to team_alias in the database
+ team_alias_jwt_field: "kubernetes\.io.namespace"
+ user_id_upsert: true
+```
+
+This way, pods in namespace `production` automatically get associated with the team that has `team_alias: production`.
+
### Set Accepted JWT Scope Names
Change the string in JWT 'scopes', that litellm evaluates to see if a user has admin access.
@@ -183,6 +366,62 @@ litellm_jwtauth:
Now litellm will automatically update the spend for the user/team/org in the db for each call.
+### Resolve by Name (Alias) Instead of ID
+
+Sometimes your JWT token contains human-readable names instead of database IDs. LiteLLM can resolve these names to IDs by looking them up in the database.
+
+**Use Case:** Your IDP provides team/org names in the JWT, but LiteLLM needs the actual database IDs for spend tracking and access control.
+
+```yaml
+general_settings:
+ master_key: sk-1234
+ enable_jwt_auth: True
+ litellm_jwtauth:
+ # Name-based fields (resolved via database lookup)
+ team_alias_jwt_field: "team_alias" # Resolves team by team_alias in DB
+ org_alias_jwt_field: "org_alias" # Resolves org by organization_alias in DB
+```
+
+**Expected JWT:**
+
+```json
+{
+ "sub": "user-123",
+ "team_alias": "engineering-team",
+ "org_alias": "acme-corp"
+}
+```
+
+**How It Works:**
+
+1. LiteLLM extracts the name from the configured JWT field
+2. Looks up the entity in the database by its alias field:
+ - Teams: `team_alias` column in `LiteLLM_TeamTable`
+ - Organizations: `organization_alias` column in `LiteLLM_OrganizationTable`
+3. Uses the resolved ID for spend tracking and access control
+
+**Precedence:** ID fields always take precedence over name fields. If both `team_id_jwt_field` and `team_alias_jwt_field` are configured and both values exist in the JWT, the ID will be used.
+
+```yaml
+# Example: ID takes precedence
+litellm_jwtauth:
+ team_id_jwt_field: "team_id" # Used if present in JWT
+ team_alias_jwt_field: "team_alias" # Fallback if team_id not present
+```
+
+**Nested Fields:** Name fields also support dot notation for nested claims:
+
+```yaml
+litellm_jwtauth:
+ team_alias_jwt_field: "organization.team.name"
+ org_alias_jwt_field: "company.name"
+```
+
+**Important Notes:**
+- The entity (team/org) must already exist in the database with the matching alias
+- Aliases should be unique - if multiple entities share the same alias, an error will be returned
+- Name resolution adds a database lookup, so using IDs directly is slightly more performant
+
### JWT Scopes
Here's what scopes on JWT-Auth tokens look like
diff --git a/docs/my-website/img/levo_logo.png b/docs/my-website/img/levo_logo.png
new file mode 100644
index 00000000000..fdb72470b29
Binary files /dev/null and b/docs/my-website/img/levo_logo.png differ
diff --git a/docs/my-website/img/levo_logo_dark.png b/docs/my-website/img/levo_logo_dark.png
new file mode 100644
index 00000000000..70da632ee90
Binary files /dev/null and b/docs/my-website/img/levo_logo_dark.png differ
diff --git a/docs/my-website/package-lock.json b/docs/my-website/package-lock.json
index 8af06ec1a94..c5f15ebd5f8 100644
--- a/docs/my-website/package-lock.json
+++ b/docs/my-website/package-lock.json
@@ -8904,23 +8904,23 @@
"license": "ISC"
},
"node_modules/body-parser": {
- "version": "1.20.3",
- "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-1.20.3.tgz",
- "integrity": "sha512-7rAxByjUMqQ3/bHJy7D6OGXvx/MMc4IqBn/X0fcM1QUcAItpZrBEYhWGem+tzXH90c+G01ypMcYJBO9Y30203g==",
+ "version": "1.20.4",
+ "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-1.20.4.tgz",
+ "integrity": "sha512-ZTgYYLMOXY9qKU/57FAo8F+HA2dGX7bqGc71txDRC1rS4frdFI5R7NhluHxH6M0YItAP0sHB4uqAOcYKxO6uGA==",
"license": "MIT",
"dependencies": {
- "bytes": "3.1.2",
+ "bytes": "~3.1.2",
"content-type": "~1.0.5",
"debug": "2.6.9",
"depd": "2.0.0",
- "destroy": "1.2.0",
- "http-errors": "2.0.0",
- "iconv-lite": "0.4.24",
- "on-finished": "2.4.1",
- "qs": "6.13.0",
- "raw-body": "2.5.2",
+ "destroy": "~1.2.0",
+ "http-errors": "~2.0.1",
+ "iconv-lite": "~0.4.24",
+ "on-finished": "~2.4.1",
+ "qs": "~6.14.0",
+ "raw-body": "~2.5.3",
"type-is": "~1.6.18",
- "unpipe": "1.0.0"
+ "unpipe": "~1.0.0"
},
"engines": {
"node": ">= 0.8",
@@ -8945,6 +8945,26 @@
"ms": "2.0.0"
}
},
+ "node_modules/body-parser/node_modules/http-errors": {
+ "version": "2.0.1",
+ "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz",
+ "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==",
+ "license": "MIT",
+ "dependencies": {
+ "depd": "~2.0.0",
+ "inherits": "~2.0.4",
+ "setprototypeof": "~1.2.0",
+ "statuses": "~2.0.2",
+ "toidentifier": "~1.0.1"
+ },
+ "engines": {
+ "node": ">= 0.8"
+ },
+ "funding": {
+ "type": "opencollective",
+ "url": "https://opencollective.com/express"
+ }
+ },
"node_modules/body-parser/node_modules/iconv-lite": {
"version": "0.4.24",
"resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.4.24.tgz",
@@ -8957,12 +8977,27 @@
"node": ">=0.10.0"
}
},
+ "node_modules/body-parser/node_modules/inherits": {
+ "version": "2.0.4",
+ "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz",
+ "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==",
+ "license": "ISC"
+ },
"node_modules/body-parser/node_modules/ms": {
"version": "2.0.0",
"resolved": "https://registry.npmjs.org/ms/-/ms-2.0.0.tgz",
"integrity": "sha512-Tpp60P6IUJDTuOq/5Z8cdskzJujfwqfOTkrwIwj7IRISpnkJnT6SyJ4PCPnGMoFjC9ddhal5KVIYtAt97ix05A==",
"license": "MIT"
},
+ "node_modules/body-parser/node_modules/statuses": {
+ "version": "2.0.2",
+ "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz",
+ "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==",
+ "license": "MIT",
+ "engines": {
+ "node": ">= 0.8"
+ }
+ },
"node_modules/bonjour-service": {
"version": "1.3.0",
"resolved": "https://registry.npmjs.org/bonjour-service/-/bonjour-service-1.3.0.tgz",
@@ -11873,39 +11908,39 @@
}
},
"node_modules/express": {
- "version": "4.21.2",
- "resolved": "https://registry.npmjs.org/express/-/express-4.21.2.tgz",
- "integrity": "sha512-28HqgMZAmih1Czt9ny7qr6ek2qddF4FclbMzwhCREB6OFfH+rXAnuNCwo1/wFvrtbgsQDb4kSbX9de9lFbrXnA==",
+ "version": "4.22.1",
+ "resolved": "https://registry.npmjs.org/express/-/express-4.22.1.tgz",
+ "integrity": "sha512-F2X8g9P1X7uCPZMA3MVf9wcTqlyNp7IhH5qPCI0izhaOIYXaW9L535tGA3qmjRzpH+bZczqq7hVKxTR4NWnu+g==",
"license": "MIT",
"dependencies": {
"accepts": "~1.3.8",
"array-flatten": "1.1.1",
- "body-parser": "1.20.3",
- "content-disposition": "0.5.4",
+ "body-parser": "~1.20.3",
+ "content-disposition": "~0.5.4",
"content-type": "~1.0.4",
- "cookie": "0.7.1",
- "cookie-signature": "1.0.6",
+ "cookie": "~0.7.1",
+ "cookie-signature": "~1.0.6",
"debug": "2.6.9",
"depd": "2.0.0",
"encodeurl": "~2.0.0",
"escape-html": "~1.0.3",
"etag": "~1.8.1",
- "finalhandler": "1.3.1",
- "fresh": "0.5.2",
- "http-errors": "2.0.0",
+ "finalhandler": "~1.3.1",
+ "fresh": "~0.5.2",
+ "http-errors": "~2.0.0",
"merge-descriptors": "1.0.3",
"methods": "~1.1.2",
- "on-finished": "2.4.1",
+ "on-finished": "~2.4.1",
"parseurl": "~1.3.3",
- "path-to-regexp": "0.1.12",
+ "path-to-regexp": "~0.1.12",
"proxy-addr": "~2.0.7",
- "qs": "6.13.0",
+ "qs": "~6.14.0",
"range-parser": "~1.2.1",
"safe-buffer": "5.2.1",
- "send": "0.19.0",
- "serve-static": "1.16.2",
+ "send": "~0.19.0",
+ "serve-static": "~1.16.2",
"setprototypeof": "1.2.0",
- "statuses": "2.0.1",
+ "statuses": "~2.0.1",
"type-is": "~1.6.18",
"utils-merge": "1.0.1",
"vary": "~1.1.2"
@@ -19281,12 +19316,12 @@
}
},
"node_modules/qs": {
- "version": "6.13.0",
- "resolved": "https://registry.npmjs.org/qs/-/qs-6.13.0.tgz",
- "integrity": "sha512-+38qI9SOr8tfZ4QmJNplMUxqjbe7LKvvZgWdExBOmd+egZTtjLB67Gu0HRX3u/XOq7UU2Nx6nsjvS16Z9uwfpg==",
+ "version": "6.14.1",
+ "resolved": "https://registry.npmjs.org/qs/-/qs-6.14.1.tgz",
+ "integrity": "sha512-4EK3+xJl8Ts67nLYNwqw/dsFVnCf+qR7RgXSK9jEEm9unao3njwMDdmsdvoKBKHzxd7tCYz5e5M+SnMjdtXGQQ==",
"license": "BSD-3-Clause",
"dependencies": {
- "side-channel": "^1.0.6"
+ "side-channel": "^1.1.0"
},
"engines": {
"node": ">=0.6"
@@ -19362,15 +19397,15 @@
}
},
"node_modules/raw-body": {
- "version": "2.5.2",
- "resolved": "https://registry.npmjs.org/raw-body/-/raw-body-2.5.2.tgz",
- "integrity": "sha512-8zGqypfENjCIqGhgXToC8aB2r7YrBX+AQAfIPs/Mlk+BtPTztOvTS01NRW/3Eh60J+a48lt8qsCzirQ6loCVfA==",
+ "version": "2.5.3",
+ "resolved": "https://registry.npmjs.org/raw-body/-/raw-body-2.5.3.tgz",
+ "integrity": "sha512-s4VSOf6yN0rvbRZGxs8Om5CWj6seneMwK3oDb4lWDH0UPhWcxwOWw5+qk24bxq87szX1ydrwylIOp2uG1ojUpA==",
"license": "MIT",
"dependencies": {
- "bytes": "3.1.2",
- "http-errors": "2.0.0",
- "iconv-lite": "0.4.24",
- "unpipe": "1.0.0"
+ "bytes": "~3.1.2",
+ "http-errors": "~2.0.1",
+ "iconv-lite": "~0.4.24",
+ "unpipe": "~1.0.0"
},
"engines": {
"node": ">= 0.8"
@@ -19385,6 +19420,26 @@
"node": ">= 0.8"
}
},
+ "node_modules/raw-body/node_modules/http-errors": {
+ "version": "2.0.1",
+ "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz",
+ "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==",
+ "license": "MIT",
+ "dependencies": {
+ "depd": "~2.0.0",
+ "inherits": "~2.0.4",
+ "setprototypeof": "~1.2.0",
+ "statuses": "~2.0.2",
+ "toidentifier": "~1.0.1"
+ },
+ "engines": {
+ "node": ">= 0.8"
+ },
+ "funding": {
+ "type": "opencollective",
+ "url": "https://opencollective.com/express"
+ }
+ },
"node_modules/raw-body/node_modules/iconv-lite": {
"version": "0.4.24",
"resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.4.24.tgz",
@@ -19397,6 +19452,21 @@
"node": ">=0.10.0"
}
},
+ "node_modules/raw-body/node_modules/inherits": {
+ "version": "2.0.4",
+ "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz",
+ "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==",
+ "license": "ISC"
+ },
+ "node_modules/raw-body/node_modules/statuses": {
+ "version": "2.0.2",
+ "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz",
+ "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==",
+ "license": "MIT",
+ "engines": {
+ "node": ">= 0.8"
+ }
+ },
"node_modules/rc": {
"version": "1.2.8",
"resolved": "https://registry.npmjs.org/rc/-/rc-1.2.8.tgz",
diff --git a/docs/my-website/src/css/custom.css b/docs/my-website/src/css/custom.css
index 2bc6a4cfdef..9fa4443afc9 100644
--- a/docs/my-website/src/css/custom.css
+++ b/docs/my-website/src/css/custom.css
@@ -28,3 +28,34 @@
--ifm-color-primary-lightest: #4fddbf;
--docusaurus-highlighted-code-line-bg: rgba(0, 0, 0, 0.3);
}
+
+/* Levo logo sizing and theme switching */
+.levo-logo-container {
+ position: relative;
+}
+
+.levo-logo-container img,
+.levo-logo-container picture,
+.levo-logo-container .ideal-image {
+ max-width: 200px !important;
+ width: 200px !important;
+ height: auto !important;
+}
+
+/* Show light logo by default, hide dark logo */
+.levo-logo-dark {
+ display: none !important;
+}
+
+.levo-logo-light {
+ display: block !important;
+}
+
+/* In dark mode, hide light logo and show dark logo */
+[data-theme='dark'] .levo-logo-light {
+ display: none !important;
+}
+
+[data-theme='dark'] .levo-logo-dark {
+ display: block !important;
+}
diff --git a/docs/my-website/src/data/adopters/README.md b/docs/my-website/src/data/adopters/README.md
new file mode 100644
index 00000000000..61a5215f802
--- /dev/null
+++ b/docs/my-website/src/data/adopters/README.md
@@ -0,0 +1,88 @@
+# LiteLLM Adopters
+
+This directory contains data for organizations that use LiteLLM in production.
+
+## Adding Your Organization
+
+We've made it super easy to add your organization! Just follow the steps below.
+
+### Quick Add (Recommended)
+
+**[Edit adopters.json on GitHub →](https://github.com/BerriAI/litellm/edit/main/docs/my-website/src/data/adopters/adopters.json)**
+
+This will open the GitHub editor in your browser where you can:
+
+1. Add your organization's entry to the JSON array
+2. Commit your changes
+3. GitHub will automatically create a pull request for you!
+
+No need to clone the repository or set up a development environment.
+
+### JSON Format
+
+Add your organization to the array in `adopters.json`:
+
+```json
+{
+ "name": "Your Organization Name",
+ "logoUrl": "https://yoursite.com/logo.svg",
+ "url": "https://yourcompany.com",
+ "description": "Brief description of how you use LiteLLM (shown on hover)"
+}
+```
+
+### Fields
+
+- **`name`** (required): Your organization's display name
+- **`logoUrl`** (required): URL to your logo - can be either:
+ - External URL: `https://yoursite.com/logo.svg` (easiest!)
+ - Local path: `/img/adopters/your-logo.svg` (requires uploading logo file)
+- **`url`** (optional): Your organization's website (makes the logo clickable)
+- **`description`** (optional): Brief description shown when users hover over your logo
+
+### Logo Options
+
+#### Option 1: External URL (Easiest)
+
+Simply provide a direct link to your logo hosted anywhere:
+
+```json
+"logoUrl": "https://yourcompany.com/assets/logo.svg"
+```
+
+#### Option 2: Local Logo (Better Performance)
+
+If you prefer to host the logo locally:
+
+1. Add your logo to `docs/my-website/static/img/adopters/your-company.svg`
+2. Reference it as: `"logoUrl": "/img/adopters/your-company.svg"`
+
+**Logo Specifications:**
+
+- **Format**: SVG preferred (PNG also acceptable)
+- **Dimensions**: 240x160px or similar 3:2 ratio recommended
+- **Background**: Transparent or white background works best
+
+### Example
+
+```json
+{
+ "name": "Acme Corporation",
+ "logoUrl": "https://acme.com/logo.svg",
+ "url": "https://acme.com",
+ "description": "Using LiteLLM to route requests across 50+ LLM providers"
+}
+```
+
+### Display Order
+
+Adopters are displayed alphabetically by organization name, so your position will be determined automatically.
+
+### Need Help?
+
+If you have questions about adding your organization:
+
+- Ask in [GitHub Discussions](https://github.com/BerriAI/litellm/discussions)
+- Join our [Discord community](https://discord.com/invite/wuPM9dRgDw)
+
+Thank you for supporting LiteLLM! 🚅
diff --git a/docs/my-website/src/data/adopters/adopters.json b/docs/my-website/src/data/adopters/adopters.json
new file mode 100644
index 00000000000..52319c149e2
--- /dev/null
+++ b/docs/my-website/src/data/adopters/adopters.json
@@ -0,0 +1,8 @@
+[
+ {
+ "name": "Your Logo Here",
+ "logoUrl": "/img/adopters/placeholder-company.svg",
+ "description": "Add your organization to show support for LiteLLM",
+ "url": "https://github.com/BerriAI/litellm/edit/main/docs/my-website/src/data/adopters/adopters.json"
+ }
+]
diff --git a/docs/my-website/src/data/adopters/index.js b/docs/my-website/src/data/adopters/index.js
new file mode 100644
index 00000000000..b1a242dcc33
--- /dev/null
+++ b/docs/my-website/src/data/adopters/index.js
@@ -0,0 +1,23 @@
+import adoptersData from './adopters.json';
+
+/**
+ * @typedef {Object} Adopter
+ * @property {string} name - The organization's display name
+ * @property {string} logoUrl - URL to the organization's logo
+ * @property {string} [url] - The organization's website URL
+ * @property {string} [description] - Brief description shown on hover
+ */
+
+/**
+ * List of organizations using LiteLLM
+ * @type {Adopter[]}
+ */
+export const adopters = adoptersData;
+
+/**
+ * Adopters sorted alphabetically by name
+ * @type {Adopter[]}
+ */
+export const sortedAdopters = [...adopters].sort((a, b) =>
+ a.name.localeCompare(b.name)
+);
diff --git a/docs/my-website/static/img/adopters/placeholder-company.svg b/docs/my-website/static/img/adopters/placeholder-company.svg
new file mode 100644
index 00000000000..937dffc6eaf
--- /dev/null
+++ b/docs/my-website/static/img/adopters/placeholder-company.svg
@@ -0,0 +1,8 @@
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql
new file mode 100644
index 00000000000..8eebb797e2c
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql
@@ -0,0 +1,5 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "authorization_url" TEXT,
+ADD COLUMN "registration_url" TEXT,
+ADD COLUMN "token_url" TEXT;
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
index aac0b5b35de..ea47b6ed03b 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
@@ -208,6 +208,9 @@ model LiteLLM_MCPServerTable {
command String?
args String[] @default([])
env Json? @default("{}")
+ authorization_url String?
+ token_url String?
+ registration_url String?
}
// Generate Tokens for Proxy
diff --git a/litellm/__init__.py b/litellm/__init__.py
index a32d2d3ef90..4f8c57ef442 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -136,6 +136,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"gitlab",
"cloudzero",
"posthog",
+ "levo",
]
cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = None
logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None
@@ -553,6 +554,8 @@ docker_model_runner_models: Set = set()
amazon_nova_models: Set = set()
stability_models: Set = set()
github_copilot_models: Set = set()
+minimax_models: Set = set()
+aws_polly_models: Set = set()
def is_bedrock_pricing_only_model(key: str) -> bool:
@@ -801,6 +804,10 @@ def add_known_models():
stability_models.add(key)
elif value.get("litellm_provider") == "github_copilot":
github_copilot_models.add(key)
+ elif value.get("litellm_provider") == "minimax":
+ minimax_models.add(key)
+ elif value.get("litellm_provider") == "aws_polly":
+ aws_polly_models.add(key)
add_known_models()
@@ -1005,6 +1012,8 @@ models_by_provider: dict = {
"amazon_nova": amazon_nova_models,
"stability": stability_models,
"github_copilot": github_copilot_models,
+ "minimax": minimax_models,
+ "aws_polly": aws_polly_models,
}
# mapping for those models which have larger equivalents
@@ -1049,8 +1058,8 @@ openai_image_generation_models = ["dall-e-2", "dall-e-3"]
openai_video_generation_models = ["sora-2"]
# timeout is lazy-loaded via __getattr__
-from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
-from litellm.litellm_core_utils.core_helpers import remove_index_from_tool_calls
+# get_llm_provider is lazy-loaded via __getattr__
+# remove_index_from_tool_calls is lazy-loaded via __getattr__
# Import KeyManagementSettings here (before utils import) because _key_management_settings
# is accessed during import time in secret_managers/main.py (via dd_tracing -> datadog -> _service_logger -> utils)
@@ -1499,6 +1508,7 @@ if TYPE_CHECKING:
get_first_chars_messages: Callable[..., str]
get_provider_fields: Callable[..., List]
get_valid_models: Callable[..., list]
+ remove_index_from_tool_calls: Callable[..., None]
# Response types - truly lazy loaded only (not in main.py or elsewhere)
ModelResponseListIterator: Type[Any]
@@ -1650,6 +1660,17 @@ def __getattr__(name: str) -> Any:
LoggingCallbackManager = __getattr__("LoggingCallbackManager")
_globals["logging_callback_manager"] = LoggingCallbackManager()
return _globals["logging_callback_manager"]
+
+ # Lazy load _service_logger module
+ if name == "_service_logger":
+ from ._lazy_imports import _get_litellm_globals
+ _globals = _get_litellm_globals()
+ # Check if already cached
+ if "_service_logger" not in _globals:
+ # Import the module lazily
+ import litellm._service_logger
+ _globals["_service_logger"] = litellm._service_logger
+ return _globals["_service_logger"]
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py
index c1b3e1df976..5ea9dd9b321 100644
--- a/litellm/_lazy_imports.py
+++ b/litellm/_lazy_imports.py
@@ -34,6 +34,7 @@ from ._lazy_imports_registry import (
DOTPROMPT_NAMES,
LLM_CONFIG_NAMES,
TYPES_NAMES,
+ LLM_PROVIDER_LOGIC_NAMES,
# Import maps
_UTILS_IMPORT_MAP,
_COST_CALCULATOR_IMPORT_MAP,
@@ -45,6 +46,7 @@ from ._lazy_imports_registry import (
_DOTPROMPT_IMPORT_MAP,
_TYPES_IMPORT_MAP,
_LLM_CONFIGS_IMPORT_MAP,
+ _LLM_PROVIDER_LOGIC_IMPORT_MAP,
)
@@ -181,6 +183,8 @@ def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]:
_LAZY_IMPORT_REGISTRY[name] = _lazy_import_llm_configs
for name in TYPES_NAMES:
_LAZY_IMPORT_REGISTRY[name] = _lazy_import_types
+ for name in LLM_PROVIDER_LOGIC_NAMES:
+ _LAZY_IMPORT_REGISTRY[name] = _lazy_import_llm_provider_logic
return _LAZY_IMPORT_REGISTRY
@@ -297,6 +301,11 @@ def _lazy_import_litellm_logging(name: str) -> Any:
"""Handler for litellm_logging module (Logging, modify_integration)"""
return _generic_lazy_import(name, _LITELLM_LOGGING_IMPORT_MAP, "Litellm logging")
+
+def _lazy_import_llm_provider_logic(name: str) -> Any:
+ """Handler for LLM provider logic functions (get_llm_provider, etc.)"""
+ return _generic_lazy_import(name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic")
+
# ============================================================================
# SPECIAL HANDLERS
# ============================================================================
diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py
index e2f80a14391..93fa8b39af2 100644
--- a/litellm/_lazy_imports_registry.py
+++ b/litellm/_lazy_imports_registry.py
@@ -32,6 +32,7 @@ UTILS_NAMES = (
"ModelResponse", "ModelResponseStream", "EmbeddingResponse", "ImageResponse",
"TranscriptionResponse", "TextCompletionResponse", "get_provider_fields",
"ModelResponseListIterator", "get_valid_models", "timeout",
+ "get_llm_provider", "remove_index_from_tool_calls",
)
# Token counter names that support lazy loading via _lazy_import_token_counter
@@ -287,6 +288,12 @@ TYPES_NAMES = (
# is accessed during import time in secret_managers/main.py
)
+# LLM provider logic names that support lazy loading via _lazy_import_llm_provider_logic
+LLM_PROVIDER_LOGIC_NAMES = (
+ "get_llm_provider",
+ "remove_index_from_tool_calls",
+)
+
# Import maps for registry pattern - reduces repetition
_UTILS_IMPORT_MAP = {
"exception_type": (".utils", "exception_type"),
@@ -330,6 +337,8 @@ _UTILS_IMPORT_MAP = {
"ModelResponseListIterator": (".utils", "ModelResponseListIterator"),
"get_valid_models": (".utils", "get_valid_models"),
"timeout": (".timeout", "timeout"),
+ "get_llm_provider": ("litellm.litellm_core_utils.get_llm_provider_logic", "get_llm_provider"),
+ "remove_index_from_tool_calls": ("litellm.litellm_core_utils.core_helpers", "remove_index_from_tool_calls"),
}
_COST_CALCULATOR_IMPORT_MAP = {
@@ -386,6 +395,11 @@ _TYPES_IMPORT_MAP = {
"LoggingCallbackManager": ("litellm.litellm_core_utils.logging_callback_manager", "LoggingCallbackManager"),
}
+_LLM_PROVIDER_LOGIC_IMPORT_MAP = {
+ "get_llm_provider": ("litellm.litellm_core_utils.get_llm_provider_logic", "get_llm_provider"),
+ "remove_index_from_tool_calls": ("litellm.litellm_core_utils.core_helpers", "remove_index_from_tool_calls"),
+}
+
_LLM_CONFIGS_IMPORT_MAP = {
"AmazonConverseConfig": (".llms.bedrock.chat.converse_transformation", "AmazonConverseConfig"),
"OpenAILikeChatConfig": (".llms.openai_like.chat.handler", "OpenAILikeChatConfig"),
@@ -587,6 +601,7 @@ __all__ = [
"DOTPROMPT_NAMES",
"LLM_CONFIG_NAMES",
"TYPES_NAMES",
+ "LLM_PROVIDER_LOGIC_NAMES",
# Import maps
"_UTILS_IMPORT_MAP",
"_COST_CALCULATOR_IMPORT_MAP",
@@ -598,5 +613,6 @@ __all__ = [
"_DOTPROMPT_IMPORT_MAP",
"_TYPES_IMPORT_MAP",
"_LLM_CONFIGS_IMPORT_MAP",
+ "_LLM_PROVIDER_LOGIC_IMPORT_MAP",
]
diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py
index 9d3f990b1aa..58a52666d38 100644
--- a/litellm/google_genai/adapters/transformation.py
+++ b/litellm/google_genai/adapters/transformation.py
@@ -8,8 +8,10 @@ from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAssistantMessage,
ChatCompletionAssistantToolCall,
+ ChatCompletionImageObject,
ChatCompletionRequest,
ChatCompletionSystemMessage,
+ ChatCompletionTextObject,
ChatCompletionToolCallFunctionChunk,
ChatCompletionToolChoiceValues,
ChatCompletionToolMessage,
@@ -385,13 +387,36 @@ class GoogleGenAIAdapter:
if role == "user":
# Handle user messages with potential function responses
- combined_text = ""
+ content_parts: List[
+ Union[ChatCompletionTextObject, ChatCompletionImageObject]
+ ] = []
tool_messages: List[ChatCompletionToolMessage] = []
for part in parts:
if isinstance(part, dict):
if "text" in part:
- combined_text += part["text"]
+ content_parts.append(
+ cast(
+ ChatCompletionTextObject,
+ {"type": "text", "text": part["text"]},
+ )
+ )
+ elif "inline_data" in part:
+ # Handle Base64 image data
+ inline_data = part["inline_data"]
+ mime_type = inline_data.get("mime_type", "image/jpeg")
+ data = inline_data.get("data", "")
+ content_parts.append(
+ cast(
+ ChatCompletionImageObject,
+ {
+ "type": "image_url",
+ "image_url": {
+ "url": f"data:{mime_type};base64,{data}"
+ },
+ },
+ )
+ )
elif "functionResponse" in part:
# Transform function response to tool message
func_response = part["functionResponse"]
@@ -402,13 +427,33 @@ class GoogleGenAIAdapter:
)
tool_messages.append(tool_message)
elif isinstance(part, str):
- combined_text += part
+ content_parts.append(
+ cast(
+ ChatCompletionTextObject, {"type": "text", "text": part}
+ )
+ )
- # Add user message if there's text content
- if combined_text:
- messages.append(
- ChatCompletionUserMessage(role="user", content=combined_text)
- )
+ # Add user message if there's content
+ if content_parts:
+ # If only one text part, use simple string format for backward compatibility
+ if (
+ len(content_parts) == 1
+ and isinstance(content_parts[0], dict)
+ and content_parts[0].get("type") == "text"
+ ):
+ text_part = cast(ChatCompletionTextObject, content_parts[0])
+ messages.append(
+ ChatCompletionUserMessage(
+ role="user", content=text_part["text"]
+ )
+ )
+ else:
+ # Use multimodal format (array of content parts)
+ messages.append(
+ ChatCompletionUserMessage(
+ role="user", content=content_parts
+ )
+ )
# Add tool messages
messages.extend(tool_messages)
@@ -468,7 +513,6 @@ class GoogleGenAIAdapter:
Dict in Google GenAI generate_content response format
"""
-
# Extract the main response content
choice = response.choices[0] if response.choices else None
if not choice:
diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py
index fe0ce208ee6..6a76b57e7f7 100644
--- a/litellm/integrations/custom_guardrail.py
+++ b/litellm/integrations/custom_guardrail.py
@@ -243,14 +243,14 @@ class CustomGuardrail(CustomLogger):
def _is_valid_response_type(self, result: Any) -> bool:
"""
Check if result is a valid LLMResponseTypes instance.
-
+
Safely handles TypedDict types which don't support isinstance checks.
For non-LiteLLM responses (like passthrough httpx.Response), returns True
to allow them through.
"""
if result is None:
return False
-
+
try:
# Try isinstance check on valid types that support it
response_types = get_args(LLMResponseTypes)
@@ -506,6 +506,7 @@ class CustomGuardrail(CustomLogger):
duration: Optional[float] = None,
masked_entity_count: Optional[Dict[str, int]] = None,
guardrail_provider: Optional[str] = None,
+ event_type: Optional[GuardrailEventHooks] = None,
) -> None:
"""
Builds `StandardLoggingGuardrailInformation` and adds it to the request metadata so it can be used for logging to DataDog, Langfuse, etc.
@@ -514,14 +515,19 @@ class CustomGuardrail(CustomLogger):
guardrail_json_response = str(guardrail_json_response)
from litellm.types.utils import GuardrailMode
+ # Use event_type if provided, otherwise fall back to self.event_hook
+ guardrail_mode: Union[GuardrailEventHooks, GuardrailMode, List[GuardrailEventHooks]]
+ if event_type is not None:
+ guardrail_mode = event_type
+ elif isinstance(self.event_hook, Mode):
+ guardrail_mode = GuardrailMode(**dict(self.event_hook.model_dump())) # type: ignore[typeddict-item]
+ else:
+ guardrail_mode = self.event_hook # type: ignore[assignment]
+
slg = StandardLoggingGuardrailInformation(
guardrail_name=self.guardrail_name,
guardrail_provider=guardrail_provider,
- guardrail_mode=(
- GuardrailMode(**self.event_hook.model_dump()) # type: ignore
- if isinstance(self.event_hook, Mode)
- else self.event_hook
- ),
+ guardrail_mode=guardrail_mode,
guardrail_response=guardrail_json_response,
guardrail_status=guardrail_status,
start_time=start_time,
@@ -589,6 +595,7 @@ class CustomGuardrail(CustomLogger):
start_time: Optional[float] = None,
end_time: Optional[float] = None,
duration: Optional[float] = None,
+ event_type: Optional[GuardrailEventHooks] = None,
):
"""
Add StandardLoggingGuardrailInformation to the request data
@@ -605,6 +612,7 @@ class CustomGuardrail(CustomLogger):
duration=duration,
start_time=start_time,
end_time=end_time,
+ event_type=event_type,
)
return response
@@ -615,6 +623,7 @@ class CustomGuardrail(CustomLogger):
start_time: Optional[float] = None,
end_time: Optional[float] = None,
duration: Optional[float] = None,
+ event_type: Optional[GuardrailEventHooks] = None,
):
"""
Add StandardLoggingGuardrailInformation to the request data
@@ -628,6 +637,7 @@ class CustomGuardrail(CustomLogger):
duration=duration,
start_time=start_time,
end_time=end_time,
+ event_type=event_type,
)
raise e
@@ -712,16 +722,32 @@ def log_guardrail_information(func):
Logs for:
- pre_call
- during_call
- - TODO: log post_call. This is more involved since the logs are sent to DD, s3 before the guardrail is even run
+ - post_call
"""
import asyncio
import functools
+ def _infer_event_type_from_function_name(
+ func_name: str,
+ ) -> Optional[GuardrailEventHooks]:
+ """Infer the actual event type from the function name"""
+ if func_name == "async_pre_call_hook":
+ return GuardrailEventHooks.pre_call
+ elif func_name == "async_moderation_hook":
+ return GuardrailEventHooks.during_call
+ elif func_name in (
+ "async_post_call_success_hook",
+ "async_post_call_streaming_hook",
+ ):
+ return GuardrailEventHooks.post_call
+ return None
+
@functools.wraps(func)
async def async_wrapper(*args, **kwargs):
start_time = datetime.now() # Move start_time inside the wrapper
self: CustomGuardrail = args[0]
request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {}
+ event_type = _infer_event_type_from_function_name(func.__name__)
try:
response = await func(*args, **kwargs)
return self._process_response(
@@ -730,6 +756,7 @@ def log_guardrail_information(func):
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
+ event_type=event_type,
)
except Exception as e:
return self._process_error(
@@ -738,6 +765,7 @@ def log_guardrail_information(func):
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
+ event_type=event_type,
)
@functools.wraps(func)
@@ -745,18 +773,21 @@ def log_guardrail_information(func):
start_time = datetime.now() # Move start_time inside the wrapper
self: CustomGuardrail = args[0]
request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {}
+ event_type = _infer_event_type_from_function_name(func.__name__)
try:
response = func(*args, **kwargs)
return self._process_response(
response=response,
request_data=request_data,
duration=(datetime.now() - start_time).total_seconds(),
+ event_type=event_type,
)
except Exception as e:
return self._process_error(
e=e,
request_data=request_data,
duration=(datetime.now() - start_time).total_seconds(),
+ event_type=event_type,
)
@functools.wraps(func)
diff --git a/litellm/integrations/generic_api/generic_api_callback.py b/litellm/integrations/generic_api/generic_api_callback.py
index 1c8a5b883da..1c62ce9fcc3 100644
--- a/litellm/integrations/generic_api/generic_api_callback.py
+++ b/litellm/integrations/generic_api/generic_api_callback.py
@@ -25,6 +25,7 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.types.utils import StandardLoggingPayload
API_EVENT_TYPES = Literal["llm_api_success", "llm_api_failure"]
+LOG_FORMAT_TYPES = Literal["json_array", "ndjson", "single"]
def load_compatible_callbacks() -> Dict:
@@ -101,6 +102,7 @@ class GenericAPILogger(CustomBatchLogger):
headers: Optional[dict] = None,
event_types: Optional[List[API_EVENT_TYPES]] = None,
callback_name: Optional[str] = None,
+ log_format: Optional[LOG_FORMAT_TYPES] = None,
**kwargs,
):
"""
@@ -111,6 +113,7 @@ class GenericAPILogger(CustomBatchLogger):
headers: Optional[dict] = None,
event_types: Optional[List[API_EVENT_TYPES]] = None,
callback_name: Optional[str] = None - If provided, loads config from generic_api_compatible_callbacks.json
+ log_format: Optional[LOG_FORMAT_TYPES] = None - Format for log output: "json_array" (default), "ndjson", or "single"
"""
#########################################################
# Check if callback_name is provided and load config
@@ -135,6 +138,9 @@ class GenericAPILogger(CustomBatchLogger):
if event_types is None and "event_types" in callback_config:
event_types = callback_config["event_types"]
+
+ if log_format is None and "log_format" in callback_config:
+ log_format = callback_config["log_format"]
else:
verbose_logger.warning(
f"callback_name '{callback_name}' not found in generic_api_compatible_callbacks.json"
@@ -156,8 +162,16 @@ class GenericAPILogger(CustomBatchLogger):
self.endpoint: str = endpoint
self.event_types: Optional[List[API_EVENT_TYPES]] = event_types
self.callback_name: Optional[str] = callback_name
+
+ # Validate and store log_format
+ if log_format is not None and log_format not in ["json_array", "ndjson", "single"]:
+ raise ValueError(
+ f"Invalid log_format: {log_format}. Must be one of: 'json_array', 'ndjson', 'single'"
+ )
+ self.log_format: LOG_FORMAT_TYPES = log_format or "json_array"
+
verbose_logger.debug(
- f"in init GenericAPILogger, callback_name: {self.callback_name}, endpoint {self.endpoint}, headers {self.headers}, event_types: {self.event_types}"
+ f"in init GenericAPILogger, callback_name: {self.callback_name}, endpoint {self.endpoint}, headers {self.headers}, event_types: {self.event_types}, log_format: {self.log_format}"
)
#########################################################
@@ -289,25 +303,65 @@ class GenericAPILogger(CustomBatchLogger):
async def async_send_batch(self):
"""
Sends the batch of messages to Generic API Endpoint
+
+ Supports three formats:
+ - json_array: Sends all logs as a JSON array (default)
+ - ndjson: Sends logs as newline-delimited JSON
+ - single: Sends each log as individual HTTP request in parallel
"""
try:
if not self.log_queue:
return
verbose_logger.debug(
- f"Generic API Logger - about to flush {len(self.log_queue)} events"
+ f"Generic API Logger - about to flush {len(self.log_queue)} events in '{self.log_format}' format"
)
- # make POST request to Generic API Endpoint
- response = await self.async_httpx_client.post(
- url=self.endpoint,
- headers=self.headers,
- data=safe_dumps(self.log_queue),
- )
+ if self.log_format == "single":
+ # Send each log as individual HTTP request in parallel
+ tasks = []
+ for log_entry in self.log_queue:
+ task = self.async_httpx_client.post(
+ url=self.endpoint,
+ headers=self.headers,
+ data=safe_dumps(log_entry),
+ )
+ tasks.append(task)
- verbose_logger.debug(
- f"Generic API Logger - sent batch to {self.endpoint}, status code {response.status_code}"
- )
+ # Execute all requests in parallel
+ responses = await asyncio.gather(*tasks, return_exceptions=True)
+
+ # Log results
+ for idx, result in enumerate(responses):
+ if isinstance(result, Exception):
+ verbose_logger.exception(
+ f"Generic API Logger - Error sending log {idx}: {result}"
+ )
+ else:
+ # result is a Response object
+ verbose_logger.debug(
+ f"Generic API Logger - sent log {idx}, status: {result.status_code}" # type: ignore
+ )
+ else:
+ # Format the payload based on log_format
+ if self.log_format == "json_array":
+ data = safe_dumps(self.log_queue)
+ elif self.log_format == "ndjson":
+ data = "\n".join(safe_dumps(log) for log in self.log_queue)
+ else:
+ raise ValueError(f"Unknown log_format: {self.log_format}")
+
+ # Make POST request
+ response = await self.async_httpx_client.post(
+ url=self.endpoint,
+ headers=self.headers,
+ data=data,
+ )
+
+ verbose_logger.debug(
+ f"Generic API Logger - sent batch to {self.endpoint}, "
+ f"status: {response.status_code}, format: {self.log_format}"
+ )
except Exception as e:
verbose_logger.exception(
diff --git a/litellm/integrations/generic_api/generic_api_compatible_callbacks.json b/litellm/integrations/generic_api/generic_api_compatible_callbacks.json
index 6c8e5fd1b2a..12dc4ae643c 100644
--- a/litellm/integrations/generic_api/generic_api_compatible_callbacks.json
+++ b/litellm/integrations/generic_api/generic_api_compatible_callbacks.json
@@ -22,6 +22,7 @@
"headers": {
"Content-Type": "application/json"
},
- "environment_variables": ["SUMOLOGIC_WEBHOOK_URL"]
+ "environment_variables": ["SUMOLOGIC_WEBHOOK_URL"],
+ "log_format": "ndjson"
}
}
\ No newline at end of file
diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py
index 10347bc7c67..f0f355b4895 100644
--- a/litellm/integrations/langfuse/langfuse.py
+++ b/litellm/integrations/langfuse/langfuse.py
@@ -3,14 +3,27 @@
import os
import traceback
from datetime import datetime
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union, cast
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ Callable,
+ Dict,
+ List,
+ Optional,
+ Tuple,
+ Union,
+ cast,
+)
from packaging.version import Version
import litellm
from litellm._logging import verbose_logger
from litellm.constants import MAX_LANGFUSE_INITIALIZED_CLIENTS
-from litellm.litellm_core_utils.core_helpers import safe_deep_copy
+from litellm.litellm_core_utils.core_helpers import (
+ safe_deep_copy,
+ reconstruct_model_name,
+)
from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
from litellm.secret_managers.main import str_to_bool
@@ -437,12 +450,17 @@ class LangFuseLogger:
)
)
+ custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider"))
+ model_name = reconstruct_model_name(
+ kwargs.get("model", ""), custom_llm_provider, metadata
+ )
+
trace.generation(
CreateGeneration(
name=metadata.get("generation_name", "litellm-completion"),
startTime=start_time,
endTime=end_time,
- model=kwargs["model"],
+ model=model_name,
modelParameters=optional_params,
prompt=input,
completion=output,
@@ -543,7 +561,9 @@ class LangFuseLogger:
# as we want to fall back to litellm_call_id instead for better traceability.
# Note: Users can still explicitly set a UUID trace_id via metadata["trace_id"] (highest priority)
if trace_id is None and standard_logging_object is not None:
- standard_trace_id = cast(Optional[str], standard_logging_object.get("trace_id"))
+ standard_trace_id = cast(
+ Optional[str], standard_logging_object.get("trace_id")
+ )
# Only use standard_logging_object.trace_id if it's not a UUID
# UUIDs are 36 characters with hyphens in format: xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx
# We check for this specific pattern to avoid rejecting valid trace_ids that happen to have hyphens
@@ -575,7 +595,9 @@ class LangFuseLogger:
mask_output = clean_metadata.pop("mask_output", False)
# Look for masking function in the dedicated location first (set by scrub_sensitive_keys_in_metadata)
# Fall back to metadata for backwards compatibility
- masking_function = litellm_params.get("_langfuse_masking_function") or clean_metadata.pop("langfuse_masking_function", None)
+ masking_function = litellm_params.get(
+ "_langfuse_masking_function"
+ ) or clean_metadata.pop("langfuse_masking_function", None)
# Apply custom masking function if provided
if masking_function is not None and callable(masking_function):
@@ -776,12 +798,17 @@ class LangFuseLogger:
if system_fingerprint is not None:
optional_params["system_fingerprint"] = system_fingerprint
+ custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider"))
+ model_name = reconstruct_model_name(
+ kwargs.get("model", ""), custom_llm_provider, metadata
+ )
+
generation_params = {
"name": generation_name,
"id": clean_metadata.pop("generation_id", generation_id),
"start_time": start_time,
"end_time": end_time,
- "model": kwargs["model"],
+ "model": model_name,
"model_parameters": optional_params,
"input": input if not mask_input else "redacted-by-litellm",
"output": output if not mask_output else "redacted-by-litellm",
@@ -918,7 +945,9 @@ class LangFuseLogger:
return Version(self.langfuse_sdk_version) >= Version("2.7.3")
@staticmethod
- def _apply_masking_function(data: Any, masking_function: Callable[[Any], Any]) -> Any:
+ def _apply_masking_function(
+ data: Any, masking_function: Callable[[Any], Any]
+ ) -> Any:
"""
Apply a masking function to data, handling different data types.
diff --git a/litellm/integrations/levo/README.md b/litellm/integrations/levo/README.md
new file mode 100644
index 00000000000..cb18b1dbfb0
--- /dev/null
+++ b/litellm/integrations/levo/README.md
@@ -0,0 +1,125 @@
+# Levo AI Integration
+
+This integration enables sending LLM observability data to Levo AI using OpenTelemetry (OTLP) protocol.
+
+## Overview
+
+The Levo integration extends LiteLLM's OpenTelemetry support to automatically send traces to Levo's collector endpoint with proper authentication and routing headers.
+
+## Features
+
+- **Automatic OTLP Export**: Sends OpenTelemetry traces to Levo collector
+- **Levo-Specific Headers**: Automatically includes `x-levo-organization-id` and `x-levo-workspace-id` for routing
+- **Simple Configuration**: Just use `callbacks: ["levo"]` in your LiteLLM config
+- **Environment-Based Setup**: Configure via environment variables
+
+## Quick Start
+
+### 1. Install Dependencies
+
+```bash
+pip install opentelemetry-api opentelemetry-sdk opentelemetry-exporter-otlp-proto-http opentelemetry-exporter-otlp-proto-grpc
+```
+
+### 2. Configure LiteLLM
+
+Add to your `litellm_config.yaml`:
+
+```yaml
+litellm_settings:
+ callbacks: ["levo"]
+```
+
+### 3. Set Environment Variables
+
+```bash
+export LEVOAI_API_KEY=""
+export LEVOAI_ORG_ID=""
+export LEVOAI_WORKSPACE_ID=""
+export LEVOAI_COLLECTOR_URL=""
+```
+
+### 4. Start LiteLLM
+
+```bash
+litellm --config config.yaml
+```
+
+All LLM requests will now automatically be sent to Levo!
+
+## Configuration
+
+### Required Environment Variables
+
+| Variable | Description |
+|----------|-------------|
+| `LEVOAI_API_KEY` | Your Levo API key for authentication |
+| `LEVOAI_ORG_ID` | Your Levo organization ID for routing |
+| `LEVOAI_WORKSPACE_ID` | Your Levo workspace ID for routing |
+| `LEVOAI_COLLECTOR_URL` | Full collector endpoint URL from Levo support |
+
+### Optional Environment Variables
+
+| Variable | Description | Default |
+|----------|-------------|---------|
+| `LEVOAI_ENV_NAME` | Environment name for tagging traces | `None` |
+
+**Important**: The `LEVOAI_COLLECTOR_URL` is used exactly as provided. No path manipulation is performed.
+
+## How It Works
+
+1. **LevoLogger** extends LiteLLM's `OpenTelemetry` class
+2. **Configuration** is read from environment variables via `get_levo_config()`
+3. **OTLP Headers** are automatically set:
+ - `Authorization: Bearer {LEVOAI_API_KEY}`
+ - `x-levo-organization-id: {LEVOAI_ORG_ID}`
+ - `x-levo-workspace-id: {LEVOAI_WORKSPACE_ID}`
+4. **Traces** are sent to the collector endpoint in OTLP format
+
+## Code Structure
+
+```
+litellm/integrations/levo/
+├── __init__.py # Exports LevoLogger
+├── levo.py # LevoLogger implementation
+└── README.md # This file
+```
+
+### Key Classes
+
+- **LevoLogger**: Extends `OpenTelemetry`, handles Levo-specific configuration
+- **LevoConfig**: Pydantic model for Levo configuration (defined in `levo.py`)
+
+## Testing
+
+See the test files in `tests/test_litellm/integrations/levo/`:
+- `test_levo.py`: Unit tests for configuration
+- `test_levo_integration.py`: Integration tests for callback registration
+
+## Error Handling
+
+The integration validates all required environment variables at initialization:
+- Missing `LEVOAI_API_KEY`: Raises `ValueError` with clear message
+- Missing `LEVOAI_ORG_ID`: Raises `ValueError` with clear message
+- Missing `LEVOAI_WORKSPACE_ID`: Raises `ValueError` with clear message
+- Missing `LEVOAI_COLLECTOR_URL`: Raises `ValueError` with clear message
+
+## Integration with LiteLLM
+
+The Levo callback is registered in:
+- `litellm/litellm_core_utils/custom_logger_registry.py`: Maps `"levo"` to `LevoLogger`
+- `litellm/litellm_core_utils/litellm_logging.py`: Instantiates `LevoLogger` when `callbacks: ["levo"]` is used
+- `litellm/__init__.py`: Added to `_custom_logger_compatible_callbacks_literal`
+
+## Documentation
+
+For detailed documentation, see:
+- [LiteLLM Levo Integration Docs](../../../../docs/my-website/docs/observability/levo_integration.md)
+- [Levo Documentation](https://docs.levo.ai)
+
+## Support
+
+For issues or questions:
+- LiteLLM Issues: https://github.com/BerriAI/litellm/issues
+- Levo Support: support@levo.ai
+
diff --git a/litellm/integrations/levo/__init__.py b/litellm/integrations/levo/__init__.py
new file mode 100644
index 00000000000..7f4f84437d4
--- /dev/null
+++ b/litellm/integrations/levo/__init__.py
@@ -0,0 +1,3 @@
+from litellm.integrations.levo.levo import LevoLogger
+
+__all__ = ["LevoLogger"]
diff --git a/litellm/integrations/levo/levo.py b/litellm/integrations/levo/levo.py
new file mode 100644
index 00000000000..562f2fd9068
--- /dev/null
+++ b/litellm/integrations/levo/levo.py
@@ -0,0 +1,117 @@
+import os
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+from litellm.integrations.opentelemetry import OpenTelemetry
+
+if TYPE_CHECKING:
+ from opentelemetry.trace import Span as _Span
+
+ from litellm.integrations.opentelemetry import OpenTelemetryConfig as _OpenTelemetryConfig
+ from litellm.types.integrations.arize import Protocol as _Protocol
+
+ Protocol = _Protocol
+ OpenTelemetryConfig = _OpenTelemetryConfig
+ Span = Union[_Span, Any]
+else:
+ Protocol = Any
+ OpenTelemetryConfig = Any
+ Span = Any
+
+
+class LevoConfig:
+ """Configuration for Levo OTLP integration."""
+
+ def __init__(
+ self,
+ otlp_auth_headers: Optional[str],
+ protocol: Protocol,
+ endpoint: str,
+ ):
+ self.otlp_auth_headers = otlp_auth_headers
+ self.protocol = protocol
+ self.endpoint = endpoint
+
+
+class LevoLogger(OpenTelemetry):
+ """Levo Logger that extends OpenTelemetry for OTLP integration."""
+
+ @staticmethod
+ def get_levo_config() -> LevoConfig:
+ """
+ Retrieves the Levo configuration based on environment variables.
+
+ Returns:
+ LevoConfig: Configuration object containing Levo OTLP settings.
+
+ Raises:
+ ValueError: If required environment variables are missing.
+ """
+ # Required environment variables
+ api_key = os.environ.get("LEVOAI_API_KEY", None)
+ org_id = os.environ.get("LEVOAI_ORG_ID", None)
+ workspace_id = os.environ.get("LEVOAI_WORKSPACE_ID", None)
+ collector_url = os.environ.get("LEVOAI_COLLECTOR_URL", None)
+
+ # Validate required env vars
+ if not api_key:
+ raise ValueError(
+ "LEVOAI_API_KEY environment variable is required for Levo integration."
+ )
+ if not org_id:
+ raise ValueError(
+ "LEVOAI_ORG_ID environment variable is required for Levo integration."
+ )
+ if not workspace_id:
+ raise ValueError(
+ "LEVOAI_WORKSPACE_ID environment variable is required for Levo integration."
+ )
+ if not collector_url:
+ raise ValueError(
+ "LEVOAI_COLLECTOR_URL environment variable is required for Levo integration. "
+ "Please contact Levo support to get your collector URL."
+ )
+
+ # Use collector URL exactly as provided by the user
+ endpoint = collector_url
+ protocol: Protocol = "otlp_http"
+
+ # Build OTLP headers string
+ # Format: Authorization=Bearer {api_key},x-levo-organization-id={org_id},x-levo-workspace-id={workspace_id}
+ headers_parts = [f"Authorization=Bearer {api_key}"]
+ headers_parts.append(f"x-levo-organization-id={org_id}")
+ headers_parts.append(f"x-levo-workspace-id={workspace_id}")
+
+ otlp_auth_headers = ",".join(headers_parts)
+
+ return LevoConfig(
+ otlp_auth_headers=otlp_auth_headers,
+ protocol=protocol,
+ endpoint=endpoint,
+ )
+
+ async def async_health_check(self):
+ """
+ Health check for Levo integration.
+
+ Returns:
+ dict: Health status with status and message/error_message keys.
+ """
+ try:
+ config = self.get_levo_config()
+
+ if not config.otlp_auth_headers:
+ return {
+ "status": "unhealthy",
+ "error_message": "LEVOAI_API_KEY environment variable not set",
+ }
+
+ return {
+ "status": "healthy",
+ "message": "Levo credentials are configured properly",
+ }
+ except ValueError as e:
+ return {
+ "status": "unhealthy",
+ "error_message": str(e),
+ }
+
diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py
index 93dce578fe1..12e60bc25bb 100644
--- a/litellm/integrations/opentelemetry.py
+++ b/litellm/integrations/opentelemetry.py
@@ -48,6 +48,7 @@ else:
LITELLM_TRACER_NAME = os.getenv("OTEL_TRACER_NAME", "litellm")
LITELLM_METER_NAME = os.getenv("LITELLM_METER_NAME", "litellm")
LITELLM_LOGGER_NAME = os.getenv("LITELLM_LOGGER_NAME", "litellm")
+LITELLM_PROXY_REQUEST_SPAN_NAME = "Received Proxy Server Request"
# Remove the hardcoded LITELLM_RESOURCE dictionary - we'll create it properly later
RAW_REQUEST_SPAN_NAME = "raw_gen_ai_request"
LITELLM_REQUEST_SPAN_NAME = "litellm_request"
@@ -233,14 +234,16 @@ class OpenTelemetry(CustomLogger):
trace.set_tracer_provider(tracer_provider)
else:
# Tracer provider explicitly provided (e.g., for testing)
+ # Do NOT call set_tracer_provider - the caller is responsible for managing global state
+ # If they want it to be global, they've already set it before passing it to us
verbose_logger.debug(
"OpenTelemetry: Using provided TracerProvider: %s",
type(tracer_provider).__name__,
)
- trace.set_tracer_provider(tracer_provider)
- # grab our tracer
- self.tracer = trace.get_tracer(LITELLM_TRACER_NAME)
+ # Grab our tracer from the TracerProvider (not from global context)
+ # This ensures we use the provided TracerProvider (e.g., for testing)
+ self.tracer = tracer_provider.get_tracer(LITELLM_TRACER_NAME)
self.span_kind = SpanKind
def _init_metrics(self, meter_provider):
@@ -527,6 +530,7 @@ class OpenTelemetry(CustomLogger):
# 3. Guardrail span
self._create_guardrail_span(kwargs=kwargs, context=ctx)
+
return response
#########################################################
@@ -557,9 +561,9 @@ class OpenTelemetry(CustomLogger):
def _get_dynamic_otel_headers_from_kwargs(self, kwargs) -> Optional[dict]:
"""Extract dynamic headers from kwargs if available."""
- standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = (
- kwargs.get("standard_callback_dynamic_params")
- )
+ standard_callback_dynamic_params: Optional[
+ StandardCallbackDynamicParams
+ ] = kwargs.get("standard_callback_dynamic_params")
if not standard_callback_dynamic_params:
return None
@@ -607,18 +611,35 @@ class OpenTelemetry(CustomLogger):
)
ctx, parent_span = self._get_span_context(kwargs)
- if get_secret_bool("USE_OTEL_LITELLM_REQUEST_SPAN"):
- primary_span_parent = None
- else:
- primary_span_parent = parent_span
-
- # 1. Primary span
- span = self._start_primary_span(
- kwargs, response_obj, start_time, end_time, ctx, primary_span_parent
+ # Decide whether to create a primary span
+ # Always create if no parent span exists (backward compatibility)
+ # OR if USE_OTEL_LITELLM_REQUEST_SPAN is explicitly enabled
+ should_create_primary_span = parent_span is None or get_secret_bool(
+ "USE_OTEL_LITELLM_REQUEST_SPAN"
)
- # 2. Raw‐request sub-span (if enabled)
- self._maybe_log_raw_request(kwargs, response_obj, start_time, end_time, span)
+ if should_create_primary_span:
+ # Create a new litellm_request span
+ span = self._start_primary_span(
+ kwargs, response_obj, start_time, end_time, ctx
+ )
+ # Raw-request sub-span (if enabled) - child of litellm_request span
+ self._maybe_log_raw_request(
+ kwargs, response_obj, start_time, end_time, span
+ )
+ else:
+ # Do not create primary span (keep hierarchy shallow when parent exists)
+ from opentelemetry.trace import Status, StatusCode
+
+ span = None
+ # Only set attributes if the span is still recording (not closed)
+ # Note: parent_span is guaranteed to be not None here
+ parent_span.set_status(Status(StatusCode.OK))
+ self.set_attributes(parent_span, kwargs, response_obj)
+ # Raw-request as direct child of parent_span
+ self._maybe_log_raw_request(
+ kwargs, response_obj, start_time, end_time, parent_span
+ )
# 3. Guardrail span
self._create_guardrail_span(kwargs=kwargs, context=ctx)
@@ -628,12 +649,18 @@ class OpenTelemetry(CustomLogger):
# 5. Semantic logs.
if self.config.enable_events:
- self._emit_semantic_logs(kwargs, response_obj, span)
+ log_span = span if span is not None else parent_span
+ if log_span is not None:
+ self._emit_semantic_logs(kwargs, response_obj, log_span)
- # 6. End parent span (only if it wasn't reused as the primary span)
- # If parent_span was reused as the primary span, it was already ended in _start_primary_span
- if parent_span is not None and parent_span is not span:
- parent_span.end(end_time=self._to_ns(datetime.now()))
+ # 6. Do NOT end parent span - it should be managed by its creator
+ # External spans (from Langfuse, user code, HTTP headers, global context) must not be closed by LiteLLM
+ # However, proxy-created spans should be closed here
+ if (
+ parent_span is not None
+ and parent_span.name == LITELLM_PROXY_REQUEST_SPAN_NAME
+ ):
+ parent_span.end(end_time=self._to_ns(end_time))
def _start_primary_span(
self,
@@ -642,16 +669,19 @@ class OpenTelemetry(CustomLogger):
start_time,
end_time,
context,
- parent_span: Optional[Span] = None,
):
from opentelemetry.trace import Status, StatusCode
otel_tracer: Tracer = self.get_tracer_to_use_for_request(kwargs)
- span = parent_span or otel_tracer.start_span(
+
+ # Always create a new span
+ # The parent relationship is preserved through the context parameter
+ span = otel_tracer.start_span(
name=self._get_span_name(kwargs),
start_time=self._to_ns(start_time),
context=context,
)
+
span.set_status(Status(StatusCode.OK))
self.set_attributes(span, kwargs, response_obj)
span.end(end_time=self._to_ns(end_time))
@@ -764,10 +794,10 @@ class OpenTelemetry(CustomLogger):
return float(val)
# isinstance(val, str) - parse datetime string (with or without microseconds)
try:
- return datetime.strptime(val, '%Y-%m-%d %H:%M:%S.%f').timestamp()
+ return datetime.strptime(val, "%Y-%m-%d %H:%M:%S.%f").timestamp()
except ValueError:
try:
- return datetime.strptime(val, '%Y-%m-%d %H:%M:%S').timestamp()
+ return datetime.strptime(val, "%Y-%m-%d %H:%M:%S").timestamp()
except ValueError:
return None
@@ -775,23 +805,23 @@ class OpenTelemetry(CustomLogger):
"""Record Time to First Token (TTFT) metric for streaming requests."""
optional_params = kwargs.get("optional_params", {})
is_streaming = optional_params.get("stream", False)
-
+
if not (self._time_to_first_token_histogram and is_streaming):
return
-
+
# Use api_call_start_time for precision (matches Prometheus implementation)
# This excludes LiteLLM overhead and measures pure LLM API latency
api_call_start_time = kwargs.get("api_call_start_time", None)
completion_start_time = kwargs.get("completion_start_time", None)
-
+
if api_call_start_time is not None and completion_start_time is not None:
# Convert to timestamps if needed (handles datetime, float, and string)
api_call_start_ts = self._to_timestamp(api_call_start_time)
completion_start_ts = self._to_timestamp(completion_start_time)
-
+
if api_call_start_ts is None or completion_start_ts is None:
return # Skip recording if conversion failed
-
+
time_to_first_token_seconds = completion_start_ts - api_call_start_ts
self._time_to_first_token_histogram.record(
time_to_first_token_seconds, attributes=common_attrs
@@ -806,38 +836,40 @@ class OpenTelemetry(CustomLogger):
common_attrs: dict,
):
"""Record Time Per Output Token (TPOT) metric.
-
+
Calculated as: generation_time / completion_tokens
- For streaming: uses end_time - completion_start_time (time to generate all tokens after first)
- For non-streaming: uses end_time - api_call_start_time (total generation time)
"""
if not self._time_per_output_token_histogram:
return
-
+
# Get completion tokens from response_obj
completion_tokens = None
if response_obj and (usage := response_obj.get("usage")):
completion_tokens = usage.get("completion_tokens")
-
+
if completion_tokens is None or completion_tokens <= 0:
return
-
+
# Calculate generation time
completion_start_time = kwargs.get("completion_start_time", None)
api_call_start_time = kwargs.get("api_call_start_time", None)
-
+
# Convert end_time to timestamp (handles datetime, float, and string)
end_time_ts = self._to_timestamp(end_time)
if end_time_ts is None:
# Fallback to duration_s if conversion failed
generation_time_seconds = duration_s
if generation_time_seconds > 0:
- time_per_output_token_seconds = generation_time_seconds / completion_tokens
+ time_per_output_token_seconds = (
+ generation_time_seconds / completion_tokens
+ )
self._time_per_output_token_histogram.record(
time_per_output_token_seconds, attributes=common_attrs
)
return
-
+
if completion_start_time is not None:
# Streaming: use completion_start_time (when first token arrived)
# This measures time to generate all tokens after the first one
@@ -858,7 +890,7 @@ class OpenTelemetry(CustomLogger):
else:
# Fallback: use duration_s (already calculated as (end_time - start_time).total_seconds())
generation_time_seconds = duration_s
-
+
if generation_time_seconds > 0:
time_per_output_token_seconds = generation_time_seconds / completion_tokens
self._time_per_output_token_histogram.record(
@@ -872,37 +904,37 @@ class OpenTelemetry(CustomLogger):
common_attrs: dict,
):
"""Record Total Generation Time (response duration) metric.
-
+
Measures pure LLM API generation time: end_time - api_call_start_time
This excludes LiteLLM overhead and measures only the LLM provider's response time.
Works for both streaming and non-streaming requests.
-
+
Mirrors Prometheus's litellm_llm_api_latency_metric.
Uses kwargs.get("end_time") with fallback to parameter for consistency with Prometheus.
"""
if not self._response_duration_histogram:
return
-
+
api_call_start_time = kwargs.get("api_call_start_time", None)
if api_call_start_time is None:
return
-
+
# Use end_time from kwargs if available (matches Prometheus), otherwise use parameter
# For streaming: end_time is when the stream completes (final chunk received)
# For non-streaming: end_time is when the response is received
_end_time = kwargs.get("end_time") or end_time
if _end_time is None:
_end_time = datetime.now()
-
+
# Convert to timestamps if needed (handles datetime, float, and string)
api_call_start_ts = self._to_timestamp(api_call_start_time)
end_time_ts = self._to_timestamp(_end_time)
-
+
if api_call_start_ts is None or end_time_ts is None:
return # Skip recording if conversion failed
-
+
response_duration_seconds = end_time_ts - api_call_start_ts
-
+
if response_duration_seconds > 0:
self._response_duration_histogram.record(
response_duration_seconds, attributes=common_attrs
@@ -1065,26 +1097,49 @@ class OpenTelemetry(CustomLogger):
)
_parent_context, parent_otel_span = self._get_span_context(kwargs)
- # Span 1: Requst sent to litellm SDK
- otel_tracer: Tracer = self.get_tracer_to_use_for_request(kwargs)
- span = otel_tracer.start_span(
- name=self._get_span_name(kwargs),
- start_time=self._to_ns(start_time),
- context=_parent_context,
+ # Decide whether to create a primary span
+ # Always create if no parent span exists (backward compatibility)
+ # OR if USE_OTEL_LITELLM_REQUEST_SPAN is explicitly enabled
+ should_create_primary_span = parent_otel_span is None or get_secret_bool(
+ "USE_OTEL_LITELLM_REQUEST_SPAN"
)
- span.set_status(Status(StatusCode.ERROR))
- self.set_attributes(span, kwargs, response_obj)
- # Record exception information using OTEL standard method
- self._record_exception_on_span(span=span, kwargs=kwargs)
+ if should_create_primary_span:
+ # Span 1: Request sent to litellm SDK
+ otel_tracer: Tracer = self.get_tracer_to_use_for_request(kwargs)
+ span = otel_tracer.start_span(
+ name=self._get_span_name(kwargs),
+ start_time=self._to_ns(start_time),
+ context=_parent_context,
+ )
+ span.set_status(Status(StatusCode.ERROR))
+ self.set_attributes(span, kwargs, response_obj)
- span.end(end_time=self._to_ns(end_time))
+ # Record exception information using OTEL standard method
+ self._record_exception_on_span(span=span, kwargs=kwargs)
+
+ span.end(end_time=self._to_ns(end_time))
+ else:
+ # When parent span exists and USE_OTEL_LITELLM_REQUEST_SPAN=false,
+ # record error on parent span (keeps hierarchy shallow)
+ # Only set attributes if the span is still recording (not closed)
+ # Note: parent_otel_span is guaranteed to be not None here
+ if parent_otel_span.is_recording():
+ parent_otel_span.set_status(Status(StatusCode.ERROR))
+ self.set_attributes(parent_otel_span, kwargs, response_obj)
+ self._record_exception_on_span(span=parent_otel_span, kwargs=kwargs)
# Create span for guardrail information
self._create_guardrail_span(kwargs=kwargs, context=_parent_context)
- if parent_otel_span is not None:
- parent_otel_span.end(end_time=self._to_ns(datetime.now()))
+ # Do NOT end parent span - it should be managed by its creator
+ # External spans (from Langfuse, user code, HTTP headers, global context) must not be closed by LiteLLM
+ # However, proxy-created spans should be closed here
+ if (
+ parent_otel_span is not None
+ and parent_otel_span.name == LITELLM_PROXY_REQUEST_SPAN_NAME
+ ):
+ parent_otel_span.end(end_time=self._to_ns(end_time))
def _record_exception_on_span(self, span: Span, kwargs: dict):
"""
@@ -1263,7 +1318,9 @@ class OpenTelemetry(CustomLogger):
)
return
elif self.callback_name == "weave_otel":
- from litellm.integrations.weave.weave_otel import set_weave_otel_attributes
+ from litellm.integrations.weave.weave_otel import (
+ set_weave_otel_attributes,
+ )
set_weave_otel_attributes(span, kwargs, response_obj)
return
@@ -1994,9 +2051,9 @@ class OpenTelemetry(CustomLogger):
"""
Create a span for the received proxy server request.
"""
-
+
return self.tracer.start_span(
- name="Received Proxy Server Request",
+ name=LITELLM_PROXY_REQUEST_SPAN_NAME,
start_time=self._to_ns(start_time),
context=self.get_traceparent_from_header(headers=headers),
kind=self.span_kind.SERVER,
diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py
index 20f1357a1c8..c01f7481277 100644
--- a/litellm/integrations/prometheus.py
+++ b/litellm/integrations/prometheus.py
@@ -214,7 +214,7 @@ class PrometheusLogger(CustomLogger):
# Remaining Rate Limit for model
self.litellm_remaining_requests_metric = self._gauge_factory(
- "litellm_remaining_requests",
+ "litellm_remaining_requests_metric",
"LLM Deployment Analytics - remaining requests for model, returned from LLM API Provider",
labelnames=self.get_labels_for_metric(
"litellm_remaining_requests_metric"
@@ -222,7 +222,7 @@ class PrometheusLogger(CustomLogger):
)
self.litellm_remaining_tokens_metric = self._gauge_factory(
- "litellm_remaining_tokens",
+ "litellm_remaining_tokens_metric",
"remaining tokens for model, returned from LLM API Provider",
labelnames=self.get_labels_for_metric(
"litellm_remaining_tokens_metric"
diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py
index 9378ca71f54..dadb36f3fd7 100644
--- a/litellm/litellm_core_utils/core_helpers.py
+++ b/litellm/litellm_core_utils/core_helpers.py
@@ -38,18 +38,18 @@ def safe_divide_seconds(
def safe_divide(
- numerator: Union[int, float],
- denominator: Union[int, float],
- default: Union[int, float] = 0
+ numerator: Union[int, float],
+ denominator: Union[int, float],
+ default: Union[int, float] = 0,
) -> Union[int, float]:
"""
Safely divide two numbers, returning a default value if denominator is zero.
-
+
Args:
numerator: The number to divide
denominator: The number to divide by
default: Value to return if denominator is zero (defaults to 0)
-
+
Returns:
The result of numerator/denominator, or default if denominator is zero
"""
@@ -153,7 +153,8 @@ def get_metadata_variable_name_from_kwargs(
- LiteLLM is now moving to using `litellm_metadata` for our metadata
"""
return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
-
+
+
def get_litellm_metadata_from_kwargs(kwargs: dict):
"""
Helper to get litellm metadata from all litellm request kwargs
@@ -176,6 +177,25 @@ def get_litellm_metadata_from_kwargs(kwargs: dict):
return {}
+def reconstruct_model_name(
+ model_name: str,
+ custom_llm_provider: Optional[str],
+ metadata: dict,
+) -> str:
+ """Reconstruct full model name with provider prefix for logging."""
+ # Check if deployment model name from router metadata is available (has original prefix)
+ deployment_model_name = metadata.get("deployment")
+ if deployment_model_name and "/" in deployment_model_name:
+ # Use the deployment model name which preserves the original provider prefix
+ return deployment_model_name
+ elif custom_llm_provider and model_name and "/" not in model_name:
+ # Only add prefix for Bedrock (not for direct Anthropic API)
+ # This ensures Bedrock models get the prefix while direct Anthropic models don't
+ if custom_llm_provider == "bedrock":
+ return f"{custom_llm_provider}/{model_name}"
+ return model_name
+
+
# Helper functions used for OTEL logging
def _get_parent_otel_span_from_kwargs(
kwargs: Optional[dict] = None,
@@ -246,8 +266,8 @@ def safe_deep_copy(data):
Safe Deep Copy
The LiteLLM request may contain objects that cannot be pickled/deep-copied
- (e.g., tracing spans, locks, clients).
-
+ (e.g., tracing spans, locks, clients).
+
This helper deep-copies each top-level key independently; on failure keeps
original ref
"""
@@ -306,23 +326,23 @@ def safe_deep_copy(data):
def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any:
"""
Recursively filter out Exception objects and callable objects from dicts/lists.
-
+
This is a defensive utility to prevent deepcopy failures when exception objects
are accidentally stored in parameter dictionaries (e.g., optional_params).
Also filters callable objects (functions) to prevent JSON serialization errors.
Exceptions and callables should not be stored in params - this function removes them.
-
+
Args:
data: The data structure to filter (dict, list, or any other type)
max_depth: Maximum recursion depth to prevent infinite loops
-
+
Returns:
Filtered data structure with Exception and callable objects removed, or None if the
entire input was an Exception or callable
"""
if max_depth <= 0:
return data
-
+
# Skip exception objects
if isinstance(data, Exception):
return None
@@ -333,7 +353,7 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any:
obj_type_name = type(data).__name__
if obj_type_name in ["Logging", "LiteLLMLoggingObj"]:
return None
-
+
if isinstance(data, dict):
result: dict[str, Any] = {}
for k, v in data.items():
@@ -352,7 +372,9 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any:
result_list: list[Any] = []
for item in data:
# Skip exception and callable items
- if isinstance(item, Exception) or (callable(item) and not isinstance(item, type)):
+ if isinstance(item, Exception) or (
+ callable(item) and not isinstance(item, type)
+ ):
continue
try:
filtered = filter_exceptions_from_params(item, max_depth - 1)
@@ -366,37 +388,35 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any:
return data
-def filter_internal_params(data: dict, additional_internal_params: Optional[set] = None) -> dict:
+def filter_internal_params(
+ data: dict, additional_internal_params: Optional[set] = None
+) -> dict:
"""
Filter out LiteLLM internal parameters that shouldn't be sent to provider APIs.
-
+
This removes internal/MCP-related parameters that are used by LiteLLM internally
but should not be included in API requests to providers.
-
+
Args:
data: Dictionary of parameters to filter
additional_internal_params: Optional set of additional internal parameter names to filter
-
+
Returns:
Filtered dictionary with internal parameters removed
"""
if not isinstance(data, dict):
return data
-
+
# Known internal parameters that should never be sent to provider APIs
internal_params = {
"skip_mcp_handler",
"mcp_handler_context",
"_skip_mcp_handler",
}
-
+
# Add any additional internal params if provided
if additional_internal_params:
internal_params.update(additional_internal_params)
-
+
# Filter out internal parameters
- return {
- k: v
- for k, v in data.items()
- if k not in internal_params
- }
\ No newline at end of file
+ return {k: v for k, v in data.items() if k not in internal_params}
diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py
index fa2ff42e1df..47cbcb8aec9 100644
--- a/litellm/litellm_core_utils/custom_logger_registry.py
+++ b/litellm/litellm_core_utils/custom_logger_registry.py
@@ -76,6 +76,7 @@ class CustomLoggerRegistry:
"arize_phoenix": OpenTelemetry,
"langtrace": OpenTelemetry,
"weave_otel": OpenTelemetry,
+ "levo": OpenTelemetry,
"mlflow": MlflowLogger,
"langfuse": LangfusePromptManagement,
"otel": OpenTelemetry,
diff --git a/litellm/litellm_core_utils/dot_notation_indexing.py b/litellm/litellm_core_utils/dot_notation_indexing.py
index 6e293a4cb77..1e835004e94 100644
--- a/litellm/litellm_core_utils/dot_notation_indexing.py
+++ b/litellm/litellm_core_utils/dot_notation_indexing.py
@@ -9,6 +9,7 @@ Custom implementation with zero external dependencies.
Supported syntax:
- "field" - top-level field
- "parent.child" - nested field
+- "parent\\.with\\.dots.child" - keys containing dots (escape with backslash)
- "array[*]" - all array elements (wildcard)
- "array[0]" - specific array element (index)
- "array[*].field" - field in all array elements
@@ -47,6 +48,9 @@ def get_nested_value(
'value'
>>> get_nested_value(data, "a.b.d", "default")
'default'
+ >>> data = {"kubernetes.io": {"namespace": "default"}}
+ >>> get_nested_value(data, "kubernetes\\.io.namespace")
+ 'default'
"""
if not key_path:
return default
@@ -58,8 +62,11 @@ def get_nested_value(
else key_path
)
- # Split the key path into parts
- parts = key_path.split(".")
+ # Split the key path into parts, respecting escaped dots (\.)
+ # Use a temporary placeholder, split on unescaped dots, then restore
+ placeholder = "\x00"
+ parts = key_path.replace("\\.", placeholder).split(".")
+ parts = [p.replace(placeholder, ".") for p in parts]
# Traverse through the dictionary
current: Any = data
diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py
index 164e2a73e65..b753e9fa8b5 100644
--- a/litellm/litellm_core_utils/get_llm_provider_logic.py
+++ b/litellm/litellm_core_utils/get_llm_provider_logic.py
@@ -229,10 +229,10 @@ def get_llm_provider( # noqa: PLR0915
elif endpoint == "https://api.ai21.com/studio/v1":
custom_llm_provider = "ai21_chat"
dynamic_api_key = get_secret_str("AI21_API_KEY")
- elif endpoint == "https://codestral.mistral.ai/v1":
+ elif endpoint == "codestral.mistral.ai/v1/chat/completions":
custom_llm_provider = "codestral"
dynamic_api_key = get_secret_str("CODESTRAL_API_KEY")
- elif endpoint == "https://codestral.mistral.ai/v1":
+ elif endpoint == "codestral.mistral.ai/v1/fim/completions":
custom_llm_provider = "text-completion-codestral"
dynamic_api_key = get_secret_str("CODESTRAL_API_KEY")
elif endpoint == "app.empower.dev/api/v1":
diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index c0090d1c3e7..cd324935562 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -59,6 +59,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.deepeval.deepeval import DeepEvalLogger
from litellm.integrations.mlflow import MlflowLogger
from litellm.integrations.sqs import SQSLogger
+from litellm.litellm_core_utils.core_helpers import reconstruct_model_name
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
StandardBuiltInToolCostTracking,
@@ -332,9 +333,9 @@ class Logging(LiteLLMLoggingBaseClass):
self.litellm_trace_id: str = litellm_trace_id or str(uuid.uuid4())
self.function_id = function_id
self.streaming_chunks: List[Any] = [] # for generating complete stream response
- self.sync_streaming_chunks: List[Any] = (
- []
- ) # for generating complete stream response
+ self.sync_streaming_chunks: List[
+ Any
+ ] = [] # for generating complete stream response
self.log_raw_request_response = log_raw_request_response
# Initialize dynamic callbacks
@@ -719,9 +720,9 @@ class Logging(LiteLLMLoggingBaseClass):
prompt_spec=prompt_spec,
dynamic_callback_params=dynamic_callback_params,
):
- self.model_call_details["prompt_integration"] = (
- logger.__class__.__name__
- )
+ self.model_call_details[
+ "prompt_integration"
+ ] = logger.__class__.__name__
return logger
except Exception:
# If check fails, continue to next logger
@@ -789,9 +790,9 @@ class Logging(LiteLLMLoggingBaseClass):
if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook(
non_default_params
):
- self.model_call_details["prompt_integration"] = (
- anthropic_cache_control_logger.__class__.__name__
- )
+ self.model_call_details[
+ "prompt_integration"
+ ] = anthropic_cache_control_logger.__class__.__name__
return anthropic_cache_control_logger
#########################################################
@@ -803,9 +804,9 @@ class Logging(LiteLLMLoggingBaseClass):
internal_usage_cache=None,
llm_router=None,
)
- self.model_call_details["prompt_integration"] = (
- vector_store_custom_logger.__class__.__name__
- )
+ self.model_call_details[
+ "prompt_integration"
+ ] = vector_store_custom_logger.__class__.__name__
# Add to global callbacks so post-call hooks are invoked
if (
vector_store_custom_logger
@@ -865,9 +866,9 @@ class Logging(LiteLLMLoggingBaseClass):
model
): # if model name was changes pre-call, overwrite the initial model call name with the new one
self.model_call_details["model"] = model
- self.model_call_details["litellm_params"]["api_base"] = (
- self._get_masked_api_base(additional_args.get("api_base", ""))
- )
+ self.model_call_details["litellm_params"][
+ "api_base"
+ ] = self._get_masked_api_base(additional_args.get("api_base", ""))
def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR0915
# Log the exact input to the LLM API
@@ -896,10 +897,10 @@ class Logging(LiteLLMLoggingBaseClass):
try:
# [Non-blocking Extra Debug Information in metadata]
if turn_off_message_logging is True:
- _metadata["raw_request"] = (
- "redacted by litellm. \
+ _metadata[
+ "raw_request"
+ ] = "redacted by litellm. \
'litellm.turn_off_message_logging=True'"
- )
else:
curl_command = self._get_request_curl_command(
api_base=additional_args.get("api_base", ""),
@@ -910,34 +911,34 @@ class Logging(LiteLLMLoggingBaseClass):
_metadata["raw_request"] = str(curl_command)
# split up, so it's easier to parse in the UI
- self.model_call_details["raw_request_typed_dict"] = (
- RawRequestTypedDict(
- raw_request_api_base=str(
- additional_args.get("api_base") or ""
- ),
- raw_request_body=self._get_raw_request_body(
- additional_args.get("complete_input_dict", {})
- ),
- # NOTE: setting ignore_sensitive_headers to True will cause
- # the Authorization header to be leaked when calls to the health
- # endpoint are made and fail.
- raw_request_headers=self._get_masked_headers(
- additional_args.get("headers", {}) or {},
- ),
- error=None,
- )
+ self.model_call_details[
+ "raw_request_typed_dict"
+ ] = RawRequestTypedDict(
+ raw_request_api_base=str(
+ additional_args.get("api_base") or ""
+ ),
+ raw_request_body=self._get_raw_request_body(
+ additional_args.get("complete_input_dict", {})
+ ),
+ # NOTE: setting ignore_sensitive_headers to True will cause
+ # the Authorization header to be leaked when calls to the health
+ # endpoint are made and fail.
+ raw_request_headers=self._get_masked_headers(
+ additional_args.get("headers", {}) or {},
+ ),
+ error=None,
)
except Exception as e:
- self.model_call_details["raw_request_typed_dict"] = (
- RawRequestTypedDict(
- error=str(e),
- )
+ self.model_call_details[
+ "raw_request_typed_dict"
+ ] = RawRequestTypedDict(
+ error=str(e),
)
- _metadata["raw_request"] = (
- "Unable to Log \
+ _metadata[
+ "raw_request"
+ ] = "Unable to Log \
raw request: {}".format(
- str(e)
- )
+ str(e)
)
if getattr(self, "logger_fn", None) and callable(self.logger_fn):
try:
@@ -1238,13 +1239,13 @@ class Logging(LiteLLMLoggingBaseClass):
for callback in callbacks:
try:
if isinstance(callback, CustomLogger):
- response: Optional[MCPPostCallResponseObject] = (
- await callback.async_post_mcp_tool_call_hook(
- kwargs=kwargs,
- response_obj=post_mcp_tool_call_response_obj,
- start_time=start_time,
- end_time=end_time,
- )
+ response: Optional[
+ MCPPostCallResponseObject
+ ] = await callback.async_post_mcp_tool_call_hook(
+ kwargs=kwargs,
+ response_obj=post_mcp_tool_call_response_obj,
+ start_time=start_time,
+ end_time=end_time,
)
######################################################################
# if any of the callbacks modify the response, use the modified response
@@ -1423,9 +1424,9 @@ class Logging(LiteLLMLoggingBaseClass):
verbose_logger.debug(
f"response_cost_failure_debug_information: {debug_info}"
)
- self.model_call_details["response_cost_failure_debug_information"] = (
- debug_info
- )
+ self.model_call_details[
+ "response_cost_failure_debug_information"
+ ] = debug_info
return None
try:
@@ -1451,9 +1452,9 @@ class Logging(LiteLLMLoggingBaseClass):
verbose_logger.debug(
f"response_cost_failure_debug_information: {debug_info}"
)
- self.model_call_details["response_cost_failure_debug_information"] = (
- debug_info
- )
+ self.model_call_details[
+ "response_cost_failure_debug_information"
+ ] = debug_info
return None
@@ -1603,16 +1604,16 @@ class Logging(LiteLLMLoggingBaseClass):
result=logging_result
)
- self.model_call_details["standard_logging_object"] = (
- get_standard_logging_object_payload(
- kwargs=self.model_call_details,
- init_response_obj=logging_result,
- start_time=start_time,
- end_time=end_time,
- logging_obj=self,
- status="success",
- standard_built_in_tools_params=self.standard_built_in_tools_params,
- )
+ self.model_call_details[
+ "standard_logging_object"
+ ] = get_standard_logging_object_payload(
+ kwargs=self.model_call_details,
+ init_response_obj=logging_result,
+ start_time=start_time,
+ end_time=end_time,
+ logging_obj=self,
+ status="success",
+ standard_built_in_tools_params=self.standard_built_in_tools_params,
)
def _transform_usage_objects(self, result):
@@ -1667,9 +1668,9 @@ class Logging(LiteLLMLoggingBaseClass):
end_time = datetime.datetime.now()
if self.completion_start_time is None:
self.completion_start_time = end_time
- self.model_call_details["completion_start_time"] = (
- self.completion_start_time
- )
+ self.model_call_details[
+ "completion_start_time"
+ ] = self.completion_start_time
self.model_call_details["log_event_type"] = "successful_api_call"
self.model_call_details["end_time"] = end_time
@@ -1706,21 +1707,21 @@ class Logging(LiteLLMLoggingBaseClass):
end_time=end_time,
)
elif isinstance(result, dict) or isinstance(result, list):
- self.model_call_details["standard_logging_object"] = (
- get_standard_logging_object_payload(
- kwargs=self.model_call_details,
- init_response_obj=result,
- start_time=start_time,
- end_time=end_time,
- logging_obj=self,
- status="success",
- standard_built_in_tools_params=self.standard_built_in_tools_params,
- )
+ self.model_call_details[
+ "standard_logging_object"
+ ] = get_standard_logging_object_payload(
+ kwargs=self.model_call_details,
+ init_response_obj=result,
+ start_time=start_time,
+ end_time=end_time,
+ logging_obj=self,
+ status="success",
+ standard_built_in_tools_params=self.standard_built_in_tools_params,
)
elif standard_logging_object is not None:
- self.model_call_details["standard_logging_object"] = (
- standard_logging_object
- )
+ self.model_call_details[
+ "standard_logging_object"
+ ] = standard_logging_object
else:
self.model_call_details["response_cost"] = None
@@ -1870,23 +1871,23 @@ class Logging(LiteLLMLoggingBaseClass):
verbose_logger.debug(
"Logging Details LiteLLM-Success Call streaming complete"
)
- self.model_call_details["complete_streaming_response"] = (
- complete_streaming_response
- )
- self.model_call_details["response_cost"] = (
- self._response_cost_calculator(result=complete_streaming_response)
- )
+ self.model_call_details[
+ "complete_streaming_response"
+ ] = complete_streaming_response
+ self.model_call_details[
+ "response_cost"
+ ] = self._response_cost_calculator(result=complete_streaming_response)
## STANDARDIZED LOGGING PAYLOAD
- self.model_call_details["standard_logging_object"] = (
- get_standard_logging_object_payload(
- kwargs=self.model_call_details,
- init_response_obj=complete_streaming_response,
- start_time=start_time,
- end_time=end_time,
- logging_obj=self,
- status="success",
- standard_built_in_tools_params=self.standard_built_in_tools_params,
- )
+ self.model_call_details[
+ "standard_logging_object"
+ ] = get_standard_logging_object_payload(
+ kwargs=self.model_call_details,
+ init_response_obj=complete_streaming_response,
+ start_time=start_time,
+ end_time=end_time,
+ logging_obj=self,
+ status="success",
+ standard_built_in_tools_params=self.standard_built_in_tools_params,
)
callbacks = self.get_combined_callback_list(
dynamic_success_callbacks=self.dynamic_success_callbacks,
@@ -2214,10 +2215,10 @@ class Logging(LiteLLMLoggingBaseClass):
)
else:
if self.stream and complete_streaming_response:
- self.model_call_details["complete_response"] = (
- self.model_call_details.get(
- "complete_streaming_response", {}
- )
+ self.model_call_details[
+ "complete_response"
+ ] = self.model_call_details.get(
+ "complete_streaming_response", {}
)
result = self.model_call_details["complete_response"]
openMeterLogger.log_success_event(
@@ -2256,10 +2257,10 @@ class Logging(LiteLLMLoggingBaseClass):
)
else:
if self.stream and complete_streaming_response:
- self.model_call_details["complete_response"] = (
- self.model_call_details.get(
- "complete_streaming_response", {}
- )
+ self.model_call_details[
+ "complete_response"
+ ] = self.model_call_details.get(
+ "complete_streaming_response", {}
)
result = self.model_call_details["complete_response"]
@@ -2402,9 +2403,9 @@ class Logging(LiteLLMLoggingBaseClass):
if complete_streaming_response is not None:
print_verbose("Async success callbacks: Got a complete streaming response")
- self.model_call_details["async_complete_streaming_response"] = (
- complete_streaming_response
- )
+ self.model_call_details[
+ "async_complete_streaming_response"
+ ] = complete_streaming_response
try:
if self.model_call_details.get("cache_hit", False) is True:
@@ -2415,10 +2416,10 @@ class Logging(LiteLLMLoggingBaseClass):
model_call_details=self.model_call_details
)
# base_model defaults to None if not set on model_info
- self.model_call_details["response_cost"] = (
- self._response_cost_calculator(
- result=complete_streaming_response
- )
+ self.model_call_details[
+ "response_cost"
+ ] = self._response_cost_calculator(
+ result=complete_streaming_response
)
verbose_logger.debug(
@@ -2431,16 +2432,16 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["response_cost"] = None
## STANDARDIZED LOGGING PAYLOAD
- self.model_call_details["standard_logging_object"] = (
- get_standard_logging_object_payload(
- kwargs=self.model_call_details,
- init_response_obj=complete_streaming_response,
- start_time=start_time,
- end_time=end_time,
- logging_obj=self,
- status="success",
- standard_built_in_tools_params=self.standard_built_in_tools_params,
- )
+ self.model_call_details[
+ "standard_logging_object"
+ ] = get_standard_logging_object_payload(
+ kwargs=self.model_call_details,
+ init_response_obj=complete_streaming_response,
+ start_time=start_time,
+ end_time=end_time,
+ logging_obj=self,
+ status="success",
+ standard_built_in_tools_params=self.standard_built_in_tools_params,
)
callbacks = self.get_combined_callback_list(
dynamic_success_callbacks=self.dynamic_async_success_callbacks,
@@ -2676,18 +2677,18 @@ class Logging(LiteLLMLoggingBaseClass):
## STANDARDIZED LOGGING PAYLOAD
- self.model_call_details["standard_logging_object"] = (
- get_standard_logging_object_payload(
- kwargs=self.model_call_details,
- init_response_obj={},
- start_time=start_time,
- end_time=end_time,
- logging_obj=self,
- status="failure",
- error_str=str(exception),
- original_exception=exception,
- standard_built_in_tools_params=self.standard_built_in_tools_params,
- )
+ self.model_call_details[
+ "standard_logging_object"
+ ] = get_standard_logging_object_payload(
+ kwargs=self.model_call_details,
+ init_response_obj={},
+ start_time=start_time,
+ end_time=end_time,
+ logging_obj=self,
+ status="failure",
+ error_str=str(exception),
+ original_exception=exception,
+ standard_built_in_tools_params=self.standard_built_in_tools_params,
)
return start_time, end_time
@@ -3301,7 +3302,9 @@ class Logging(LiteLLMLoggingBaseClass):
# Deep copy result and add usage
result_copy = result.model_copy(deep=True)
- result_copy.usage = usage.model_dump() if hasattr(usage, "model_dump") else dict(usage)
+ result_copy.usage = (
+ usage.model_dump() if hasattr(usage, "model_dump") else dict(usage)
+ )
return result_copy
@@ -3629,9 +3632,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
endpoint=arize_config.endpoint,
)
- os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
- f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}"
- )
+ os.environ[
+ "OTEL_EXPORTER_OTLP_TRACES_HEADERS"
+ ] = f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}"
for callback in _in_memory_loggers:
if (
isinstance(callback, ArizeLogger)
@@ -3642,7 +3645,6 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_in_memory_loggers.append(_arize_otel_logger)
return _arize_otel_logger # type: ignore
elif logging_integration == "arize_phoenix":
-
from litellm.integrations.opentelemetry import (
OpenTelemetry,
OpenTelemetryConfig,
@@ -3658,13 +3660,13 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "")
# Add openinference.project.name attribute
if existing_attrs:
- os.environ["OTEL_RESOURCE_ATTRIBUTES"] = (
- f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}"
- )
+ os.environ[
+ "OTEL_RESOURCE_ATTRIBUTES"
+ ] = f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}"
else:
- os.environ["OTEL_RESOURCE_ATTRIBUTES"] = (
- f"openinference.project.name={arize_phoenix_config.project_name}"
- )
+ os.environ[
+ "OTEL_RESOURCE_ATTRIBUTES"
+ ] = f"openinference.project.name={arize_phoenix_config.project_name}"
# Set Phoenix project name from environment variable
phoenix_project_name = os.environ.get("PHOENIX_PROJECT_NAME", None)
@@ -3672,19 +3674,19 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "")
# Add openinference.project.name attribute
if existing_attrs:
- os.environ["OTEL_RESOURCE_ATTRIBUTES"] = (
- f"{existing_attrs},openinference.project.name={phoenix_project_name}"
- )
+ os.environ[
+ "OTEL_RESOURCE_ATTRIBUTES"
+ ] = f"{existing_attrs},openinference.project.name={phoenix_project_name}"
else:
- os.environ["OTEL_RESOURCE_ATTRIBUTES"] = (
- f"openinference.project.name={phoenix_project_name}"
- )
+ os.environ[
+ "OTEL_RESOURCE_ATTRIBUTES"
+ ] = f"openinference.project.name={phoenix_project_name}"
# auth can be disabled on local deployments of arize phoenix
if arize_phoenix_config.otlp_auth_headers is not None:
- os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
- arize_phoenix_config.otlp_auth_headers
- )
+ os.environ[
+ "OTEL_EXPORTER_OTLP_TRACES_HEADERS"
+ ] = arize_phoenix_config.otlp_auth_headers
for callback in _in_memory_loggers:
if (
@@ -3697,6 +3699,31 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
)
_in_memory_loggers.append(_arize_phoenix_otel_logger)
return _arize_phoenix_otel_logger # type: ignore
+ elif logging_integration == "levo":
+ from litellm.integrations.levo.levo import LevoLogger
+ from litellm.integrations.opentelemetry import (
+ OpenTelemetry,
+ OpenTelemetryConfig,
+ )
+
+ levo_config = LevoLogger.get_levo_config()
+ otel_config = OpenTelemetryConfig(
+ exporter=levo_config.protocol,
+ endpoint=levo_config.endpoint,
+ headers=levo_config.otlp_auth_headers,
+ )
+
+ # Check if LevoLogger instance already exists
+ for callback in _in_memory_loggers:
+ if (
+ isinstance(callback, LevoLogger)
+ and callback.callback_name == "levo"
+ ):
+ return callback # type: ignore
+
+ _levo_otel_logger = LevoLogger(config=otel_config, callback_name="levo")
+ _in_memory_loggers.append(_levo_otel_logger)
+ return _levo_otel_logger # type: ignore
elif logging_integration == "otel":
from litellm.integrations.opentelemetry import OpenTelemetry
@@ -3816,9 +3843,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
exporter="otlp_http",
endpoint="https://langtrace.ai/api/trace",
)
- os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
- f"api_key={os.getenv('LANGTRACE_API_KEY')}"
- )
+ os.environ[
+ "OTEL_EXPORTER_OTLP_TRACES_HEADERS"
+ ] = f"api_key={os.getenv('LANGTRACE_API_KEY')}"
for callback in _in_memory_loggers:
if (
isinstance(callback, OpenTelemetry)
@@ -4589,10 +4616,10 @@ class StandardLoggingPayloadSetup:
for key in StandardLoggingHiddenParams.__annotations__.keys():
if key in hidden_params:
if key == "additional_headers":
- clean_hidden_params["additional_headers"] = (
- StandardLoggingPayloadSetup.get_additional_headers(
- hidden_params[key]
- )
+ clean_hidden_params[
+ "additional_headers"
+ ] = StandardLoggingPayloadSetup.get_additional_headers(
+ hidden_params[key]
)
else:
clean_hidden_params[key] = hidden_params[key] # type: ignore
@@ -4898,25 +4925,6 @@ def _extract_response_obj_and_hidden_params(
return response_obj, hidden_params
-def _reconstruct_model_name(
- model_name: str,
- custom_llm_provider: Optional[str],
- metadata: dict,
-) -> str:
- """Reconstruct full model name with provider prefix for logging."""
- # Check if deployment model name from router metadata is available (has original prefix)
- deployment_model_name = metadata.get("deployment")
- if deployment_model_name and "/" in deployment_model_name:
- # Use the deployment model name which preserves the original provider prefix
- return deployment_model_name
- elif custom_llm_provider and model_name and "/" not in model_name:
- # Only add prefix for Bedrock (not for direct Anthropic API)
- # This ensures Bedrock models get the prefix while direct Anthropic models don't
- if custom_llm_provider == "bedrock":
- return f"{custom_llm_provider}/{model_name}"
- return model_name
-
-
def get_standard_logging_object_payload(
kwargs: Optional[dict],
init_response_obj: Union[Any, BaseModel, dict],
@@ -5049,7 +5057,7 @@ def get_standard_logging_object_payload(
# This ensures Bedrock models like "us.anthropic.claude-3-5-sonnet-20240620-v1:0"
# are logged as "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0"
custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider"))
- model_name = _reconstruct_model_name(
+ model_name = reconstruct_model_name(
kwargs.get("model", "") or "", custom_llm_provider, metadata
)
@@ -5205,9 +5213,9 @@ def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]):
):
for k, v in metadata["user_api_key_metadata"].items():
if k == "logging": # prevent logging user logging keys
- cleaned_user_api_key_metadata[k] = (
- "scrubbed_by_litellm_for_sensitive_keys"
- )
+ cleaned_user_api_key_metadata[
+ k
+ ] = "scrubbed_by_litellm_for_sensitive_keys"
else:
cleaned_user_api_key_metadata[k] = v
diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py
index 232d9bfc5d1..cbc0763382c 100644
--- a/litellm/litellm_core_utils/llm_cost_calc/utils.py
+++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py
@@ -161,6 +161,15 @@ def _get_token_base_cost(
prompt_base_cost = cast(float, _get_cost_per_unit(model_info, input_cost_key))
completion_base_cost = cast(float, _get_cost_per_unit(model_info, output_cost_key))
+
+ # For image generation models that don't have output_cost_per_token,
+ # use output_cost_per_image_token as the base cost (all output tokens are image tokens)
+ if completion_base_cost == 0.0 or completion_base_cost is None:
+ output_image_cost = _get_cost_per_unit(
+ model_info, "output_cost_per_image_token", None
+ )
+ if output_image_cost is not None:
+ completion_base_cost = cast(float, output_image_cost)
cache_creation_cost = cast(
float, _get_cost_per_unit(model_info, cache_creation_cost_key)
)
@@ -342,6 +351,7 @@ class PromptTokensDetailsResult(TypedDict):
cache_creation_token_details: Optional[CacheCreationTokenDetails]
text_tokens: int
audio_tokens: int
+ image_tokens: int
character_count: int
image_count: int
video_length_seconds: int
@@ -374,6 +384,10 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
cast(Optional[int], getattr(usage.prompt_tokens_details, "audio_tokens", 0))
or 0
)
+ image_tokens = (
+ cast(Optional[int], getattr(usage.prompt_tokens_details, "image_tokens", 0))
+ or 0
+ )
character_count = (
cast(
Optional[int],
@@ -398,6 +412,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
cache_creation_token_details=cache_creation_token_details,
text_tokens=text_tokens,
audio_tokens=audio_tokens,
+ image_tokens=image_tokens,
character_count=character_count,
image_count=image_count,
video_length_seconds=video_length_seconds,
@@ -470,6 +485,11 @@ def _calculate_input_cost(
model_info, "input_cost_per_audio_token", prompt_tokens_details["audio_tokens"]
)
+ ### IMAGE TOKEN COST (for gpt-image-1 and similar models)
+ prompt_cost += calculate_cost_component(
+ model_info, "input_cost_per_image_token", prompt_tokens_details["image_tokens"]
+ )
+
### CACHE WRITING COST - Now uses tiered pricing
prompt_cost += calculate_cache_writing_cost(
cache_creation_tokens=prompt_tokens_details["cache_creation_tokens"],
@@ -533,6 +553,7 @@ def generic_cost_per_token(
cache_creation_token_details=None,
text_tokens=usage.prompt_tokens,
audio_tokens=0,
+ image_tokens=0,
character_count=0,
image_count=0,
video_length_seconds=0,
@@ -583,12 +604,22 @@ def generic_cost_per_token(
reasoning_tokens = completion_tokens_details["reasoning_tokens"]
image_tokens = completion_tokens_details["image_tokens"]
- # Only assume all tokens are text if there's NO breakdown at all
- # If image_tokens, audio_tokens, or reasoning_tokens exist, respect text_tokens=0
+ # Handle text_tokens calculation:
+ # 1. If text_tokens is explicitly provided and > 0, use it
+ # 2. If there's a breakdown (reasoning/audio/image tokens), calculate text_tokens as the remainder
+ # 3. If no breakdown at all, assume all completion_tokens are text_tokens
has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0
- if text_tokens == 0 and not has_token_breakdown:
- text_tokens = usage.completion_tokens
- is_text_tokens_total = True
+ if text_tokens == 0:
+ if has_token_breakdown:
+ # Calculate text tokens as remainder when we have a breakdown
+ # This handles cases like OpenAI's reasoning models where text_tokens isn't provided
+ text_tokens = max(
+ 0, usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens
+ )
+ else:
+ # No breakdown at all, all tokens are text tokens
+ text_tokens = usage.completion_tokens
+ is_text_tokens_total = True
## TEXT COST
completion_cost = float(text_tokens) * completion_base_cost
@@ -782,6 +813,50 @@ class CostCalculatorUtils:
model=model,
image_response=completion_response,
)
+ elif custom_llm_provider == litellm.LlmProviders.OPENAI.value:
+ # Check if this is a gpt-image model (token-based pricing)
+ model_lower = model.lower()
+ if "gpt-image-1" in model_lower:
+ from litellm.llms.openai.image_generation.cost_calculator import (
+ cost_calculator as openai_gpt_image_cost_calculator,
+ )
+
+ return openai_gpt_image_cost_calculator(
+ model=model,
+ image_response=completion_response,
+ custom_llm_provider=custom_llm_provider,
+ )
+ # Fall through to default for DALL-E models
+ return default_image_cost_calculator(
+ model=model,
+ quality=quality,
+ custom_llm_provider=custom_llm_provider,
+ n=n,
+ size=size,
+ optional_params=optional_params,
+ )
+ elif custom_llm_provider == litellm.LlmProviders.AZURE.value:
+ # Check if this is a gpt-image model (token-based pricing)
+ model_lower = model.lower()
+ if "gpt-image-1" in model_lower:
+ from litellm.llms.openai.image_generation.cost_calculator import (
+ cost_calculator as openai_gpt_image_cost_calculator,
+ )
+
+ return openai_gpt_image_cost_calculator(
+ model=model,
+ image_response=completion_response,
+ custom_llm_provider=custom_llm_provider,
+ )
+ # Fall through to default for DALL-E models
+ return default_image_cost_calculator(
+ model=model,
+ quality=quality,
+ custom_llm_provider=custom_llm_provider,
+ n=n,
+ size=size,
+ optional_params=optional_params,
+ )
else:
return default_image_cost_calculator(
model=model,
diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py
index 59d2a8a8dd0..bbe28e3ec2c 100644
--- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py
+++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py
@@ -445,25 +445,43 @@ def convert_to_model_response_object( # noqa: PLR0915
hidden_params["additional_headers"] = additional_headers
### CHECK IF ERROR IN RESPONSE ### - openrouter returns these in the dictionary
+ # Some OpenAI-compatible providers (e.g., Apertis) return empty error objects
+ # even on success. Only raise if the error contains meaningful data.
if (
response_object is not None
and "error" in response_object
and response_object["error"] is not None
):
- error_args = {"status_code": 422, "message": "Error in response object"}
- if isinstance(response_object["error"], dict):
- if "code" in response_object["error"]:
- error_args["status_code"] = response_object["error"]["code"]
- if "message" in response_object["error"]:
- if isinstance(response_object["error"]["message"], dict):
- message_str = json.dumps(response_object["error"]["message"])
- else:
- message_str = str(response_object["error"]["message"])
- error_args["message"] = message_str
- raised_exception = Exception()
- setattr(raised_exception, "status_code", error_args["status_code"])
- setattr(raised_exception, "message", error_args["message"])
- raise raised_exception
+ error_obj = response_object["error"]
+ has_meaningful_error = False
+
+ if isinstance(error_obj, dict):
+ # Check if error dict has non-empty message or non-null code
+ error_message = error_obj.get("message", "")
+ error_code = error_obj.get("code")
+ has_meaningful_error = bool(error_message) or error_code is not None
+ elif isinstance(error_obj, str):
+ # String error is meaningful if non-empty
+ has_meaningful_error = bool(error_obj)
+ else:
+ # Any other truthy value is considered meaningful
+ has_meaningful_error = True
+
+ if has_meaningful_error:
+ error_args = {"status_code": 422, "message": "Error in response object"}
+ if isinstance(error_obj, dict):
+ if "code" in error_obj:
+ error_args["status_code"] = error_obj["code"]
+ if "message" in error_obj:
+ if isinstance(error_obj["message"], dict):
+ message_str = json.dumps(error_obj["message"])
+ else:
+ message_str = str(error_obj["message"])
+ error_args["message"] = message_str
+ raised_exception = Exception()
+ setattr(raised_exception, "status_code", error_args["status_code"])
+ setattr(raised_exception, "message", error_args["message"])
+ raise raised_exception
try:
if response_type == "completion" and (
diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py
index b78484816da..4f76a5bad03 100644
--- a/litellm/litellm_core_utils/logging_callback_manager.py
+++ b/litellm/litellm_core_utils/logging_callback_manager.py
@@ -166,6 +166,7 @@ class LoggingCallbackManager:
endpoint = callback_config.get("endpoint")
headers = callback_config.get("headers")
event_types = callback_config.get("event_types")
+ log_format = callback_config.get("log_format")
if endpoint is None or headers is None:
verbose_logger.warning(
@@ -180,6 +181,7 @@ class LoggingCallbackManager:
and cached_logger.endpoint == endpoint
and cached_logger.headers == headers
and cached_logger.event_types == event_types
+ and cached_logger.log_format == log_format
):
return cached_logger
@@ -187,6 +189,7 @@ class LoggingCallbackManager:
endpoint=endpoint,
headers=headers,
event_types=event_types,
+ log_format=log_format,
)
_generic_api_logger_cache[callback] = new_logger
return new_logger
diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py
index 20b0bc92fb7..13a83956edd 100644
--- a/litellm/litellm_core_utils/logging_worker.py
+++ b/litellm/litellm_core_utils/logging_worker.py
@@ -51,6 +51,7 @@ class LoggingWorker:
self._worker_task: Optional[asyncio.Task] = None
self._running_tasks: set[asyncio.Task] = set()
self._sem: Optional[asyncio.Semaphore] = None
+ self._bound_loop: Optional[asyncio.AbstractEventLoop] = None
self._last_aggressive_clear_time: float = 0.0
self._aggressive_clear_in_progress: bool = False
@@ -58,9 +59,27 @@ class LoggingWorker:
atexit.register(self._flush_on_exit)
def _ensure_queue(self) -> None:
- """Initialize the queue if it doesn't exist."""
+ """Initialize the queue if it doesn't exist or if event loop has changed."""
+ try:
+ current_loop = asyncio.get_running_loop()
+ except RuntimeError:
+ # No running loop, can't initialize
+ return
+
+ # Check if we need to reinitialize due to event loop change
+ if self._queue is not None and self._bound_loop is not current_loop:
+ verbose_logger.debug(
+ "LoggingWorker: Event loop changed, reinitializing queue and worker"
+ )
+ # Clear old state - these are bound to the old loop
+ self._queue = None
+ self._sem = None
+ self._worker_task = None
+ self._running_tasks.clear()
+
if self._queue is None:
self._queue = asyncio.Queue(maxsize=self.max_queue_size)
+ self._bound_loop = current_loop
def start(self) -> None:
"""Start the logging worker. Idempotent - safe to call multiple times."""
@@ -126,7 +145,7 @@ class LoggingWorker:
# Capture the current context when enqueueing
task = LoggingTask(coroutine=coroutine, context=contextvars.copy_context())
-
+
try:
self._queue.put_nowait(task)
except asyncio.QueueFull:
@@ -141,15 +160,15 @@ class LoggingWorker:
"""
if self._aggressive_clear_in_progress:
return False
-
+
try:
loop = asyncio.get_running_loop()
current_time = loop.time()
time_since_last_clear = current_time - self._last_aggressive_clear_time
-
+
if time_since_last_clear < LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS:
return False
-
+
return True
except RuntimeError:
# No event loop running, drop the task
@@ -158,8 +177,8 @@ class LoggingWorker:
def _mark_aggressive_clear_started(self) -> None:
"""
Mark that an aggressive clear operation has started.
-
- Note: This should only be called after _should_start_aggressive_clear()
+
+ Note: This should only be called after _should_start_aggressive_clear()
returns True, which guarantees an event loop exists.
"""
loop = asyncio.get_running_loop()
@@ -171,7 +190,7 @@ class LoggingWorker:
Handle queue full condition by either starting an aggressive clear
or scheduling a delayed retry.
"""
-
+
if self._should_start_aggressive_clear():
self._mark_aggressive_clear_started()
# Schedule clearing as async task so enqueue returns immediately (non-blocking)
@@ -191,7 +210,8 @@ class LoggingWorker:
time_since_last_clear = current_time - self._last_aggressive_clear_time
remaining_cooldown = max(
0.0,
- LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS - time_since_last_clear
+ LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS
+ - time_since_last_clear,
)
# Add a small buffer (10% of cooldown or 50ms, whichever is larger) to ensure
# cooldown has expired and aggressive clear has completed
@@ -212,7 +232,7 @@ class LoggingWorker:
# Check that we have a running event loop (will raise RuntimeError if not)
asyncio.get_running_loop()
delay = self._calculate_retry_delay()
-
+
# Schedule the retry as a background task
asyncio.create_task(self._retry_enqueue_task(task, delay))
except RuntimeError:
@@ -225,11 +245,11 @@ class LoggingWorker:
This is called as a background task from _schedule_delayed_enqueue_retry.
"""
await asyncio.sleep(delay)
-
+
# Try to enqueue the task directly, preserving its original context
if self._queue is None:
return
-
+
try:
self._queue.put_nowait(task)
except asyncio.QueueFull:
@@ -243,15 +263,17 @@ class LoggingWorker:
"""
if self._queue is None:
return []
-
+
# Calculate items based on percentage of queue size
- items_to_extract = (self.max_queue_size * LOGGING_WORKER_CLEAR_PERCENTAGE) // 100
+ items_to_extract = (
+ self.max_queue_size * LOGGING_WORKER_CLEAR_PERCENTAGE
+ ) // 100
# Use actual queue size to avoid unnecessary iterations
actual_size = self._queue.qsize()
if actual_size == 0:
return []
items_to_extract = min(items_to_extract, actual_size)
-
+
# Extract tasks from queue (using list comprehension would require wrapping in try/except)
extracted_tasks = []
for _ in range(items_to_extract):
@@ -259,10 +281,12 @@ class LoggingWorker:
extracted_tasks.append(self._queue.get_nowait())
except asyncio.QueueEmpty:
break
-
+
return extracted_tasks
- async def _aggressively_clear_queue_async(self, new_task: Optional[LoggingTask] = None) -> None:
+ async def _aggressively_clear_queue_async(
+ self, new_task: Optional[LoggingTask] = None
+ ) -> None:
"""
Aggressively clear the queue by extracting and processing items.
This is called when the queue is full to prevent dropping logs.
@@ -271,18 +295,20 @@ class LoggingWorker:
try:
if self._queue is None:
return
-
+
extracted_tasks = self._extract_tasks_from_queue()
-
+
# Add new task to extracted tasks to process directly
if new_task is not None:
extracted_tasks.append(new_task)
-
+
# Process extracted tasks directly
if extracted_tasks:
await self._process_extracted_tasks(extracted_tasks)
except Exception as e:
- verbose_logger.exception(f"LoggingWorker error during aggressive clear: {e}")
+ verbose_logger.exception(
+ f"LoggingWorker error during aggressive clear: {e}"
+ )
finally:
# Always reset the flag even if an error occurs
self._aggressive_clear_in_progress = False
@@ -291,7 +317,7 @@ class LoggingWorker:
"""Process a single task and mark it done."""
if self._queue is None:
return
-
+
try:
await asyncio.wait_for(
task["context"].run(asyncio.create_task, task["coroutine"]),
@@ -310,7 +336,7 @@ class LoggingWorker:
"""
if not tasks or self._queue is None:
return
-
+
# Process all tasks concurrently for maximum speed
await asyncio.gather(*[self._process_single_task(task) for task in tasks])
@@ -361,10 +387,7 @@ class LoggingWorker:
for _ in range(MAX_ITERATIONS_TO_CLEAR_QUEUE):
# Check if we've exceeded the maximum time
- if (
- asyncio.get_event_loop().time() - start_time
- >= MAX_TIME_TO_CLEAR_QUEUE
- ):
+ if asyncio.get_event_loop().time() - start_time >= MAX_TIME_TO_CLEAR_QUEUE:
verbose_logger.warning(
f"clear_queue exceeded max_time of {MAX_TIME_TO_CLEAR_QUEUE}s, stopping early"
)
@@ -381,6 +404,9 @@ class LoggingWorker:
except Exception:
# Suppress errors during cleanup
pass
+ finally:
+ # Clear reference to prevent memory leaks
+ task = None
self._queue.task_done() # If you're using join() elsewhere
except asyncio.QueueEmpty:
break
@@ -410,7 +436,7 @@ class LoggingWorker:
This ensures callbacks queued by async completions are processed
even when the script exits before the worker loop can handle them.
-
+
Note: All logging in this method is wrapped to handle cases where
logging handlers are closed during shutdown.
"""
@@ -423,7 +449,9 @@ class LoggingWorker:
return
queue_size = self._queue.qsize()
- self._safe_log("info", f"[LoggingWorker] atexit: Flushing {queue_size} remaining events...")
+ self._safe_log(
+ "info", f"[LoggingWorker] atexit: Flushing {queue_size} remaining events..."
+ )
# Create a new event loop since the original is closed
loop = asyncio.new_event_loop()
@@ -438,7 +466,7 @@ class LoggingWorker:
if loop.time() - start_time >= MAX_TIME_TO_CLEAR_QUEUE:
self._safe_log(
"warning",
- f"[LoggingWorker] atexit: Reached time limit ({MAX_TIME_TO_CLEAR_QUEUE}s), stopping flush"
+ f"[LoggingWorker] atexit: Reached time limit ({MAX_TIME_TO_CLEAR_QUEUE}s), stopping flush",
)
break
@@ -456,8 +484,14 @@ class LoggingWorker:
except Exception:
# Silent failure to not break user's program
pass
+ finally:
+ # Clear reference to prevent memory leaks
+ task = None
- self._safe_log("info", f"[LoggingWorker] atexit: Successfully flushed {processed} events!")
+ self._safe_log(
+ "info",
+ f"[LoggingWorker] atexit: Successfully flushed {processed} events!",
+ )
finally:
loop.close()
diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py
index ca2a092dbc8..b100b9b516b 100644
--- a/litellm/litellm_core_utils/prompt_templates/common_utils.py
+++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py
@@ -1087,9 +1087,35 @@ def _parse_content_for_reasoning(
return None, message_text
+def _extract_base64_data(image_url: str) -> str:
+ """
+ Extract pure base64 data from an image URL.
+
+ If the URL is a data URL (e.g., "data:image/png;base64,iVBOR..."),
+ extract and return only the base64 data portion.
+ Otherwise, return the original URL unchanged.
+
+ This is needed for providers like Ollama that expect pure base64 data
+ rather than full data URLs.
+
+ Args:
+ image_url: The image URL or data URL to process
+
+ Returns:
+ The base64 data if it's a data URL, otherwise the original URL
+ """
+ if image_url.startswith("data:") and ";base64," in image_url:
+ return image_url.split(";base64,", 1)[1]
+ return image_url
+
+
def extract_images_from_message(message: AllMessageValues) -> List[str]:
"""
- Extract images from a message
+ Extract images from a message.
+
+ For data URLs (e.g., "data:image/png;base64,iVBOR..."), only the base64
+ data portion is extracted. This is required for providers like Ollama
+ that expect pure base64 data rather than full data URLs.
"""
images = []
message_content = message.get("content")
@@ -1098,7 +1124,7 @@ def extract_images_from_message(message: AllMessageValues) -> List[str]:
image_url = m.get("image_url")
if image_url:
if isinstance(image_url, str):
- images.append(image_url)
+ images.append(_extract_base64_data(image_url))
elif isinstance(image_url, dict) and "url" in image_url:
- images.append(image_url["url"])
+ images.append(_extract_base64_data(image_url["url"]))
return images
diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py
index 6cc6c229f56..12570a02de7 100644
--- a/litellm/litellm_core_utils/prompt_templates/factory.py
+++ b/litellm/litellm_core_utils/prompt_templates/factory.py
@@ -930,7 +930,8 @@ def create_anthropic_image_param(
# Check if the image URL is an HTTP/HTTPS URL
if image_url.startswith("http://") or image_url.startswith("https://"):
- # For Bedrock invoke, always convert URLs to base64 (Bedrock invoke doesn't support URLs)
+ # For Bedrock invoke and Vertex AI Anthropic, always convert URLs to base64
+ # as these providers don't support URL sources for images
if is_bedrock_invoke or image_url.startswith("http://"):
base64_url = convert_url_to_base64(url=image_url)
image_chunk = convert_to_anthropic_image_obj(
@@ -1496,9 +1497,10 @@ def convert_to_gemini_tool_call_result(
content_type = content.get("type", "")
if content_type == "text":
content_str += content.get("text", "")
- elif content_type == "input_image":
- # Extract image for inline_data (for Computer Use screenshots)
- image_url = content.get("image_url", "")
+ elif content_type in ("input_image", "image_url"):
+ # Extract image for inline_data (for Computer Use screenshots and tool results)
+ image_url_data = content.get("image_url", "")
+ image_url = image_url_data.get("url", "") if isinstance(image_url_data, dict) else image_url_data
if image_url:
# Convert image to base64 blob format for Gemini
@@ -2022,9 +2024,12 @@ def anthropic_messages_pt( # noqa: PLR0915
"format": image_url_value.get("format"),
}
# Bedrock invoke models have format: invoke/...
+ # Vertex AI Anthropic also doesn't support URL sources for images
is_bedrock_invoke = model.lower().startswith("invoke/")
+ is_vertex_ai = llm_provider.startswith("vertex_ai") if llm_provider else False
+ force_base64 = is_bedrock_invoke or is_vertex_ai
_anthropic_content_element = create_anthropic_image_param(
- image_url_input, format=format, is_bedrock_invoke=is_bedrock_invoke
+ image_url_input, format=format, is_bedrock_invoke=force_base64
)
_content_element = add_cache_control_to_content(
anthropic_content_element=_anthropic_content_element,
diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py
index 15c035ceec8..c73f0b22b4b 100644
--- a/litellm/llms/__init__.py
+++ b/litellm/llms/__init__.py
@@ -45,6 +45,7 @@ def get_cost_for_web_search_request(
return 0.0
elif custom_llm_provider == "xai":
from .xai.cost_calculator import cost_per_web_search_request
+
return cost_per_web_search_request(usage=usage, model_info=model_info)
else:
return None
@@ -110,6 +111,21 @@ def discover_guardrail_translation_mappings() -> (
verbose_logger.error(f"Error processing {module_path}: {e}")
continue
+ try:
+ from litellm.proxy._experimental.mcp_server.guardrail_translation import (
+ guardrail_translation_mappings as mcp_guardrail_translation_mappings,
+ )
+
+ discovered_mappings.update(mcp_guardrail_translation_mappings)
+ verbose_logger.debug(
+ "Loaded MCP guardrail translation mappings: %s",
+ list(mcp_guardrail_translation_mappings.keys()),
+ )
+ except ImportError:
+ verbose_logger.debug(
+ "MCP guardrail translation mappings not available; skipping"
+ )
+
verbose_logger.debug(
f"Discovered {len(discovered_mappings)} guardrail translation mappings: {list(discovered_mappings.keys())}"
)
diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py
index 6bdc17f7979..c71edcdc2d1 100644
--- a/litellm/llms/anthropic/chat/transformation.py
+++ b/litellm/llms/anthropic/chat/transformation.py
@@ -54,7 +54,10 @@ from litellm.types.utils import (
CompletionTokensDetailsWrapper,
)
from litellm.types.utils import Message as LitellmMessage
-from litellm.types.utils import PromptTokensDetailsWrapper, ServerToolUse
+from litellm.types.utils import (
+ PromptTokensDetailsWrapper,
+ ServerToolUse,
+)
from litellm.utils import (
ModelResponse,
Usage,
@@ -204,9 +207,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
) # Relevant issue: https://github.com/BerriAI/litellm/issues/7755
def get_cache_control_headers(self) -> dict:
+ # Anthropic no longer requires the prompt-caching beta header
+ # Prompt caching now works automatically when cache_control is used in messages
+ # Reference: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching
return {
"anthropic-version": "2023-06-01",
- "anthropic-beta": "prompt-caching-2024-07-31",
}
def _map_tool_choice(
@@ -1034,7 +1039,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
anthropic_messages = anthropic_messages_pt(
model=model,
messages=messages,
- llm_provider="anthropic",
+ llm_provider=self.custom_llm_provider or "anthropic",
)
except Exception as e:
raise AnthropicError(
diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py
index 098694f15ae..fcbe9823ed4 100644
--- a/litellm/llms/anthropic/common_utils.py
+++ b/litellm/llms/anthropic/common_utils.py
@@ -12,7 +12,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
)
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
from litellm.llms.base_llm.chat.transformation import BaseLLMException
-from litellm.types.llms.anthropic import AllAnthropicToolsValues, AnthropicMcpServerTool, ANTHROPIC_HOSTED_TOOLS
+from litellm.types.llms.anthropic import (
+ ANTHROPIC_HOSTED_TOOLS,
+ AllAnthropicToolsValues,
+ AnthropicMcpServerTool,
+)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import TokenCountResponse
@@ -273,8 +277,9 @@ class AnthropicModelInfo(BaseLLMModelInfo):
beta_header = self.get_computer_tool_beta_header(computer_tool_used)
betas.append(beta_header)
- if prompt_caching_set:
- betas.append("prompt-caching-2024-07-31")
+ # Anthropic no longer requires the prompt-caching beta header
+ # Prompt caching now works automatically when cache_control is used in messages
+ # Reference: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching
if file_id_used:
betas.append("files-api-2025-04-14")
@@ -305,8 +310,9 @@ class AnthropicModelInfo(BaseLLMModelInfo):
container_with_skills_used: bool = False,
) -> dict:
betas = set()
- if prompt_caching_set:
- betas.add("prompt-caching-2024-07-31")
+ # Anthropic no longer requires the prompt-caching beta header
+ # Prompt caching now works automatically when cache_control is used in messages
+ # Reference: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching
if computer_tool_used:
beta_header = self.get_computer_tool_beta_header(computer_tool_used)
betas.add(beta_header)
diff --git a/litellm/llms/azure_ai/anthropic/messages_transformation.py b/litellm/llms/azure_ai/anthropic/messages_transformation.py
index 73dc84167ab..55818cc07d6 100644
--- a/litellm/llms/azure_ai/anthropic/messages_transformation.py
+++ b/litellm/llms/azure_ai/anthropic/messages_transformation.py
@@ -48,7 +48,12 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig):
headers = BaseAzureLLM._base_validate_azure_environment(
headers=headers, litellm_params=litellm_params_obj
)
-
+
+ # Azure Anthropic uses x-api-key header (not api-key)
+ # Convert api-key to x-api-key if present
+ if "api-key" in headers and "x-api-key" not in headers:
+ headers["x-api-key"] = headers.pop("api-key")
+
# Set anthropic-version header
if "anthropic-version" not in headers:
headers["anthropic-version"] = "2023-06-01"
diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py
index d8692bb6a3a..d3b55b13ad4 100644
--- a/litellm/llms/gemini/google_genai/transformation.py
+++ b/litellm/llms/gemini/google_genai/transformation.py
@@ -312,7 +312,9 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM):
)
request_dict = cast(dict, typed_generate_content_request)
-
+
+ if system_instruction is not None:
+ request_dict["systemInstruction"] = system_instruction
return request_dict
def transform_generate_content_response(
diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py
index 9e6497e66ab..71956158f52 100644
--- a/litellm/llms/ollama/completion/handler.py
+++ b/litellm/llms/ollama/completion/handler.py
@@ -15,7 +15,7 @@ def _prepare_ollama_embedding_payload(
) -> Dict[str, Any]:
data: Dict[str, Any] = {"model": model, "input": prompts}
- special_optional_params = ["truncate", "options", "keep_alive"]
+ special_optional_params = ["truncate", "options", "keep_alive","dimensions"]
for k, v in optional_params.items():
if k in special_optional_params:
diff --git a/litellm/llms/openai/image_generation/cost_calculator.py b/litellm/llms/openai/image_generation/cost_calculator.py
new file mode 100644
index 00000000000..35caaf6e9b1
--- /dev/null
+++ b/litellm/llms/openai/image_generation/cost_calculator.py
@@ -0,0 +1,63 @@
+"""
+Cost calculator for OpenAI image generation models (gpt-image-1, gpt-image-1-mini)
+
+These models use token-based pricing instead of pixel-based pricing like DALL-E.
+"""
+
+from typing import Optional
+
+from litellm import verbose_logger
+from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
+from litellm.responses.utils import ResponseAPILoggingUtils
+from litellm.types.utils import ImageResponse
+
+
+def cost_calculator(
+ model: str,
+ image_response: ImageResponse,
+ custom_llm_provider: Optional[str] = None,
+) -> float:
+ """
+ Calculate cost for OpenAI gpt-image-1 and gpt-image-1-mini models.
+
+ Uses the same usage format as Responses API, so we reuse the helper
+ to transform to chat completion format and use generic_cost_per_token.
+
+ Args:
+ model: The model name (e.g., "gpt-image-1", "gpt-image-1-mini")
+ image_response: The ImageResponse containing usage data
+ custom_llm_provider: Optional provider name
+
+ Returns:
+ float: Total cost in USD
+ """
+ usage = getattr(image_response, "usage", None)
+
+ if usage is None:
+ verbose_logger.debug(
+ f"No usage data available for {model}, cannot calculate token-based cost"
+ )
+ return 0.0
+
+ # Transform ImageUsage to Usage using the existing helper
+ # ImageUsage has the same format as ResponseAPIUsage
+ chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
+ usage
+ )
+
+ # Use generic_cost_per_token for cost calculation
+ prompt_cost, completion_cost = generic_cost_per_token(
+ model=model,
+ usage=chat_usage,
+ custom_llm_provider=custom_llm_provider or "openai",
+ )
+
+ total_cost = prompt_cost + completion_cost
+
+ verbose_logger.debug(
+ f"OpenAI gpt-image cost calculation for {model}: "
+ f"prompt_cost=${prompt_cost:.6f}, completion_cost=${completion_cost:.6f}, "
+ f"total=${total_cost:.6f}"
+ )
+
+ return total_cost
diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py
index 01ceb72c0de..e13abca59f8 100755
--- a/litellm/llms/sap/chat/transformation.py
+++ b/litellm/llms/sap/chat/transformation.py
@@ -203,9 +203,13 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig):
headers: dict,
) -> dict:
supported_params = self.get_supported_openai_params(model)
+ # Include extra params that passed validation (e.g., thinking_config for Gemini models via allowed_openai_params)
+ extra_params = [k for k in optional_params if k not in supported_params and k not in {"tools", "model_version"}]
+ supported_params = supported_params + extra_params
model_params = {
k: v for k, v in optional_params.items() if k in supported_params
}
+
model_version = optional_params.pop("model_version", "latest")
template = []
for message in messages:
diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py
index baa825bfcca..2bbdfa17cde 100644
--- a/litellm/llms/vertex_ai/gemini/transformation.py
+++ b/litellm/llms/vertex_ai/gemini/transformation.py
@@ -383,7 +383,18 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
and isinstance(_message_content, str)
):
assistant_text = _message_content
- assistant_content.append(PartType(text=assistant_text)) # type: ignore
+ # Check if message has thought_signatures in provider_specific_fields
+ provider_specific_fields = assistant_msg.get("provider_specific_fields")
+ thought_signatures = None
+ if provider_specific_fields and isinstance(provider_specific_fields, dict):
+ thought_signatures = provider_specific_fields.get("thought_signatures")
+
+ # If we have thought signatures, add them to the part
+ if thought_signatures and isinstance(thought_signatures, list) and len(thought_signatures) > 0:
+ # Use the first signature for the text part (Gemini expects one signature per part)
+ assistant_content.append(PartType(text=assistant_text, thoughtSignature=thought_signatures[0])) # type: ignore
+ else:
+ assistant_content.append(PartType(text=assistant_text)) # type: ignore
## HANDLE ASSISTANT FUNCTION CALL
if (
diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
index a5cc3dca8c1..ba1788a217f 100644
--- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
@@ -552,24 +552,46 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"Invalid tool={}. Use `litellm.set_verbose` or `litellm --detailed_debug` to see raw request."
)
- # Only include function_declarations if there are actual functions
- _tools = Tools()
+# Build list of Tool objects - each Tool should contain exactly one type
+ # per Vertex AI API spec: "A Tool object should contain exactly one type of Tool"
+ _tools_list: List[Tools] = []
+
+ # Function declarations can be grouped together in one Tool
if gtool_func_declarations:
- _tools["function_declarations"] = gtool_func_declarations
+ func_tool = Tools()
+ func_tool["function_declarations"] = gtool_func_declarations
+ _tools_list.append(func_tool)
+
+ # Each special tool type must be in its own Tool object
if googleSearch is not None:
- _tools[VertexToolName.GOOGLE_SEARCH.value] = googleSearch
+ search_tool = Tools()
+ search_tool[VertexToolName.GOOGLE_SEARCH.value] = googleSearch
+ _tools_list.append(search_tool)
if googleSearchRetrieval is not None:
- _tools[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = googleSearchRetrieval
+ retrieval_tool = Tools()
+ retrieval_tool[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = googleSearchRetrieval
+ _tools_list.append(retrieval_tool)
if enterpriseWebSearch is not None:
- _tools[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = enterpriseWebSearch
+ enterprise_tool = Tools()
+ enterprise_tool[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = enterpriseWebSearch
+ _tools_list.append(enterprise_tool)
if code_execution is not None:
- _tools[VertexToolName.CODE_EXECUTION.value] = code_execution
+ code_tool = Tools()
+ code_tool[VertexToolName.CODE_EXECUTION.value] = code_execution
+ _tools_list.append(code_tool)
if urlContext is not None:
- _tools[VertexToolName.URL_CONTEXT.value] = urlContext
+ url_tool = Tools()
+ url_tool[VertexToolName.URL_CONTEXT.value] = urlContext
+ _tools_list.append(url_tool)
if googleMaps is not None:
- _tools[VertexToolName.GOOGLE_MAPS.value] = googleMaps
+ maps_tool = Tools()
+ maps_tool[VertexToolName.GOOGLE_MAPS.value] = googleMaps
+ _tools_list.append(maps_tool)
if computerUse is not None:
- _tools[VertexToolName.COMPUTER_USE.value] = computerUse
+ computer_tool = Tools()
+ computer_tool[VertexToolName.COMPUTER_USE.value] = computerUse
+ _tools_list.append(computer_tool)
+
# Add retrieval config to toolConfig if googleMaps has location data
if google_maps_retrieval_config is not None:
@@ -579,7 +601,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"retrievalConfig"
] = google_maps_retrieval_config
- return [_tools]
+ return _tools_list
def _map_response_schema(self, value: dict) -> dict:
old_schema = deepcopy(value)
@@ -1210,6 +1232,25 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
thinking_blocks.append(block)
return thinking_blocks
+ def _extract_thought_signatures_from_parts(
+ self, parts: List[HttpxPartType]
+ ) -> Optional[List[str]]:
+ """Extract thoughtSignature values from parts.
+
+ Per Google's docs, thoughtSignature is returned for multi-turn context preservation
+ and can appear on parts even without thought: true (e.g., regular text responses,
+ function calls). This method extracts all thoughtSignature values from parts.
+
+ Returns:
+ List of thoughtSignature strings if any are found, None otherwise
+ """
+ signatures: List[str] = []
+ for part in parts:
+ signature = part.get("thoughtSignature")
+ if signature is not None:
+ signatures.append(signature)
+ return signatures if signatures else None
+
def _extract_image_response_from_parts(
self, parts: List[HttpxPartType]
) -> Optional[List[ImageURLListItem]]:
@@ -1620,6 +1661,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
from litellm.types.utils import Delta, StreamingChoices
annotations = chat_completion_message.get("annotations") # type: ignore
+ provider_specific_fields = chat_completion_message.get("provider_specific_fields") # type: ignore
# create a streaming choice object
choice = StreamingChoices(
finish_reason=VertexGeminiConfig._check_finish_reason(
@@ -1633,6 +1675,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
images=image_response,
function_call=functions,
annotations=annotations, # type: ignore
+ provider_specific_fields=provider_specific_fields,
),
logprobs=chat_completion_logprobs,
enhancements=None,
@@ -1811,6 +1854,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
)
)
+ # Extract thoughtSignatures from parts (can exist without thought: true)
+ thought_signatures = (
+ VertexGeminiConfig()._extract_thought_signatures_from_parts(
+ parts=candidate["content"]["parts"]
+ )
+ )
+
if audio_response is not None:
cast(Dict[str, Any], chat_completion_message)[
"audio"
@@ -1876,6 +1926,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
reasoning_content = "\n".join(reasoning_content_parts)
chat_completion_message["reasoning_content"] = reasoning_content
+ # Store thoughtSignatures in provider_specific_fields
+ if thought_signatures is not None:
+ if "provider_specific_fields" not in chat_completion_message:
+ chat_completion_message["provider_specific_fields"] = {}
+ chat_completion_message["provider_specific_fields"]["thought_signatures"] = thought_signatures # type: ignore
+
if isinstance(model_response, ModelResponseStream):
choice = VertexGeminiConfig._create_streaming_choice(
chat_completion_message=chat_completion_message,
diff --git a/litellm/llms/zai/chat/transformation.py b/litellm/llms/zai/chat/transformation.py
index 47b314d4e0d..4380256f0a4 100644
--- a/litellm/llms/zai/chat/transformation.py
+++ b/litellm/llms/zai/chat/transformation.py
@@ -20,7 +20,7 @@ class ZAIChatConfig(OpenAIGPTConfig):
return api_base, dynamic_api_key
def get_supported_openai_params(self, model: str) -> list:
- return [
+ base_params = [
"max_tokens",
"stream",
"stream_options",
@@ -31,3 +31,12 @@ class ZAIChatConfig(OpenAIGPTConfig):
"tool_choice",
]
+ import litellm
+
+ try:
+ if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider):
+ base_params.append("thinking")
+ except Exception:
+ pass
+
+ return base_params
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 513a4a554e0..d32adf54b5e 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -249,6 +249,30 @@
"/v1/images/generations"
]
},
+ "aiml/google/imagen-4.0-ultra-generate-001": {
+ "litellm_provider": "aiml",
+ "metadata": {
+ "notes": "Imagen 4.0 Ultra Generate API - Photorealistic image generation with precise text rendering"
+ },
+ "mode": "image_generation",
+ "output_cost_per_image": 0.063,
+ "source": "https://docs.aimlapi.com/api-references/image-models/google/imagen-4-ultra-generate",
+ "supported_endpoints": [
+ "/v1/images/generations"
+ ]
+ },
+ "aiml/google/nano-banana-pro": {
+ "litellm_provider": "aiml",
+ "metadata": {
+ "notes": "Gemini 3 Pro Image (Nano Banana Pro) - Advanced text-to-image generation with reasoning and 4K resolution support"
+ },
+ "mode": "image_generation",
+ "output_cost_per_image": 0.1575,
+ "source": "https://docs.aimlapi.com/api-references/image-models/google/gemini-3-pro-image-preview",
+ "supported_endpoints": [
+ "/v1/images/generations"
+ ]
+ },
"amazon.nova-canvas-v1:0": {
"litellm_provider": "bedrock",
"max_input_tokens": 2600,
@@ -3508,6 +3532,40 @@
"supports_service_tier": true,
"supports_vision": true
},
+ "azure/gpt-5.2-chat": {
+ "cache_read_input_token_cost": 1.75e-07,
+ "cache_read_input_token_cost_priority": 3.5e-07,
+ "input_cost_per_token": 1.75e-06,
+ "input_cost_per_token_priority": 3.5e-06,
+ "litellm_provider": "azure",
+ "max_input_tokens": 128000,
+ "max_output_tokens": 16384,
+ "max_tokens": 16384,
+ "mode": "chat",
+ "output_cost_per_token": 1.4e-05,
+ "output_cost_per_token_priority": 2.8e-05,
+ "supported_endpoints": [
+ "/v1/chat/completions",
+ "/v1/responses"
+ ],
+ "supported_modalities": [
+ "text",
+ "image"
+ ],
+ "supported_output_modalities": [
+ "text"
+ ],
+ "supports_function_calling": true,
+ "supports_native_streaming": true,
+ "supports_parallel_function_calling": true,
+ "supports_pdf_input": true,
+ "supports_prompt_caching": true,
+ "supports_reasoning": true,
+ "supports_response_schema": true,
+ "supports_system_messages": true,
+ "supports_tool_choice": true,
+ "supports_vision": true
+ },
"azure/gpt-5.2-chat-2025-12-11": {
"cache_read_input_token_cost": 1.75e-07,
"cache_read_input_token_cost_priority": 3.5e-07,
@@ -3605,12 +3663,16 @@
"supports_web_search": true
},
"azure/gpt-image-1": {
- "input_cost_per_pixel": 4.0054321e-08,
+ "cache_read_input_image_token_cost": 2.5e-06,
+ "cache_read_input_token_cost": 1.25e-06,
+ "input_cost_per_image_token": 1e-05,
+ "input_cost_per_token": 5e-06,
"litellm_provider": "azure",
"mode": "image_generation",
- "output_cost_per_pixel": 0.0,
+ "output_cost_per_image_token": 4e-05,
"supported_endpoints": [
- "/v1/images/generations"
+ "/v1/images/generations",
+ "/v1/images/edits"
]
},
"azure/hd/1024-x-1024/dall-e-3": {
@@ -3713,12 +3775,16 @@
]
},
"azure/gpt-image-1-mini": {
- "input_cost_per_pixel": 8.0566406e-09,
+ "cache_read_input_image_token_cost": 2.5e-07,
+ "cache_read_input_token_cost": 2e-07,
+ "input_cost_per_image_token": 2.5e-06,
+ "input_cost_per_token": 2e-06,
"litellm_provider": "azure",
"mode": "image_generation",
- "output_cost_per_pixel": 0.0,
+ "output_cost_per_image_token": 8e-06,
"supported_endpoints": [
- "/v1/images/generations"
+ "/v1/images/generations",
+ "/v1/images/edits"
]
},
"azure/gpt-image-1.5": {
@@ -10885,13 +10951,13 @@
"supports_tool_choice": true
},
"fireworks_ai/accounts/fireworks/models/deepseek-v3p2": {
- "input_cost_per_token": 1.2e-06,
+ "input_cost_per_token": 5.6e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 163840,
"max_output_tokens": 163840,
"max_tokens": 163840,
"mode": "chat",
- "output_cost_per_token": 1.2e-06,
+ "output_cost_per_token": 1.68e-06,
"source": "https://fireworks.ai/models/fireworks/deepseek-v3p2",
"supports_function_calling": true,
"supports_reasoning": true,
@@ -16922,6 +16988,336 @@
"supports_vision": true,
"supports_pdf_input": true
},
+ "low/1024-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.009,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "low/1024-x-1536/gpt-image-1.5": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "low/1536-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "medium/1024-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.034,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "medium/1024-x-1536/gpt-image-1.5": {
+ "input_cost_per_image": 0.05,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "medium/1536-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.05,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "high/1024-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.133,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "high/1024-x-1536/gpt-image-1.5": {
+ "input_cost_per_image": 0.20,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "high/1536-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.20,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "standard/1024-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.009,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "standard/1024-x-1536/gpt-image-1.5": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "standard/1536-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "1024-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.009,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "1024-x-1536/gpt-image-1.5": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "1536-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "low/1024-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.009,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "low/1024-x-1536/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "low/1536-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "medium/1024-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.034,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "medium/1024-x-1536/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.05,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "medium/1536-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.05,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "high/1024-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.133,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "high/1024-x-1536/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.20,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "high/1536-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.20,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "standard/1024-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.009,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "standard/1024-x-1536/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "standard/1536-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "1024-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.009,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "1024-x-1536/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "1536-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
"gpt-5": {
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_flex": 6.25e-08,
@@ -17683,16 +18079,16 @@
"supports_vision": true
},
"gpt-image-1": {
- "input_cost_per_image": 0.042,
- "input_cost_per_pixel": 4.0054321e-08,
- "input_cost_per_token": 0.000005,
- "input_cost_per_image_token": 0.00001,
+ "cache_read_input_image_token_cost": 2.5e-06,
+ "cache_read_input_token_cost": 1.25e-06,
+ "input_cost_per_image_token": 1e-05,
+ "input_cost_per_token": 5e-06,
"litellm_provider": "openai",
"mode": "image_generation",
- "output_cost_per_pixel": 0.0,
- "output_cost_per_token": 0.00004,
+ "output_cost_per_image_token": 4e-05,
"supported_endpoints": [
- "/v1/images/generations"
+ "/v1/images/generations",
+ "/v1/images/edits"
]
},
"gpt-image-1-mini": {
@@ -18117,6 +18513,18 @@
"supports_response_schema": false,
"supports_tool_choice": true
},
+ "groq/gemma-7b-it": {
+ "input_cost_per_token": 5e-08,
+ "litellm_provider": "groq",
+ "max_input_tokens": 8192,
+ "max_output_tokens": 8192,
+ "max_tokens": 8192,
+ "mode": "chat",
+ "output_cost_per_token": 8e-08,
+ "supports_function_calling": true,
+ "supports_response_schema": false,
+ "supports_tool_choice": true
+ },
"groq/meta-llama/llama-guard-4-12b": {
"input_cost_per_token": 2e-07,
"litellm_provider": "groq",
@@ -29241,6 +29649,20 @@
"supports_vision": true,
"supports_web_search": true
},
+ "zai/glm-4.7": {
+ "cache_creation_input_token_cost": 0,
+ "cache_read_input_token_cost": 1.1e-07,
+ "input_cost_per_token": 6e-07,
+ "output_cost_per_token": 2.2e-06,
+ "litellm_provider": "zai",
+ "max_input_tokens": 200000,
+ "max_output_tokens": 128000,
+ "mode": "chat",
+ "supports_function_calling": true,
+ "supports_reasoning": true,
+ "supports_tool_choice": true,
+ "source": "https://docs.z.ai/guides/overview/pricing"
+ },
"zai/glm-4.6": {
"input_cost_per_token": 6e-07,
"output_cost_per_token": 2.2e-06,
diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/__init__.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/__init__.py
new file mode 100644
index 00000000000..e0fd610e678
--- /dev/null
+++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/__init__.py
@@ -0,0 +1,16 @@
+"""Guardrail translation mapping for MCP tool calls."""
+
+from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import (
+ MCPGuardrailTranslationHandler,
+)
+from litellm.types.utils import CallTypes
+
+# This mapping lives alongside the MCP server implementation because MCP
+# integrations are managed by the proxy subsystem, not litellm.llms providers.
+# Unified guardrails import this module explicitly to register the handler.
+
+guardrail_translation_mappings = {
+ CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler,
+}
+
+__all__ = ["guardrail_translation_mappings", "MCPGuardrailTranslationHandler"]
diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py
new file mode 100644
index 00000000000..8d6d236b884
--- /dev/null
+++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py
@@ -0,0 +1,89 @@
+"""
+MCP Guardrail Handler for Unified Guardrails.
+
+This handler works with the synthetic "messages" payload generated by
+`ProxyLogging._convert_mcp_to_llm_format`, which always produces a single user
+message whose `content` string encodes the MCP tool name and arguments. The
+handler simply feeds that text through the configured guardrail and writes the
+result back onto the message.
+"""
+
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+from litellm._logging import verbose_proxy_logger
+from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
+from litellm.types.utils import GenericGuardrailAPIInputs
+
+if TYPE_CHECKING:
+ from litellm.integrations.custom_guardrail import CustomGuardrail
+ from mcp.types import CallToolResult
+
+
+class MCPGuardrailTranslationHandler(BaseTranslation):
+ """Guardrail translation handler for MCP tool calls."""
+
+ async def process_input_messages(
+ self,
+ data: Dict[str, Any],
+ guardrail_to_apply: "CustomGuardrail",
+ litellm_logging_obj: Optional[Any] = None,
+ ) -> Dict[str, Any]:
+ messages = data.get("messages")
+ if not isinstance(messages, list) or not messages:
+ verbose_proxy_logger.debug("MCP Guardrail: No messages to process")
+ return data
+
+ first_message = messages[0]
+ content: Optional[str] = None
+ if isinstance(first_message, dict):
+ content = first_message.get("content")
+ else:
+ content = getattr(first_message, "content", None)
+
+ if not isinstance(content, str):
+ verbose_proxy_logger.debug(
+ "MCP Guardrail: Message content missing or not a string",
+ )
+ return data
+
+ guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
+ inputs=GenericGuardrailAPIInputs(texts=[content]),
+ request_data=data,
+ input_type="request",
+ logging_obj=litellm_logging_obj,
+ )
+ guardrailed_texts = (
+ guardrailed_inputs.get("texts", []) if guardrailed_inputs else []
+ )
+
+ if guardrailed_texts:
+ new_content = guardrailed_texts[0]
+ if isinstance(first_message, dict):
+ first_message["content"] = new_content
+ else:
+ setattr(first_message, "content", new_content)
+
+ verbose_proxy_logger.debug(
+ "MCP Guardrail: Updated content for tool %s",
+ data.get("mcp_tool_name"),
+ )
+ else:
+ verbose_proxy_logger.debug(
+ "MCP Guardrail: Guardrail returned no text updates for tool %s",
+ data.get("mcp_tool_name"),
+ )
+
+ return data
+
+ async def process_output_response(
+ self,
+ response: "CallToolResult",
+ guardrail_to_apply: "CustomGuardrail",
+ litellm_logging_obj: Optional[Any] = None,
+ user_api_key_dict: Optional[Any] = None,
+ ) -> Any:
+ # Not implemented: MCP guardrail translation never calls this path today.
+ verbose_proxy_logger.debug(
+ "MCP Guardrail: Output processing not implemented for MCP tools",
+ )
+ return response
diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
index 2260649e8b2..d193fc27fb9 100644
--- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
@@ -11,7 +11,7 @@ import datetime
import hashlib
import json
import re
-from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast
+from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, cast
from urllib.parse import urlparse
from fastapi import HTTPException
@@ -30,6 +30,7 @@ from pydantic import AnyUrl
import litellm
from litellm._logging import verbose_logger
+from litellm.types.utils import CallTypes
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.experimental_mcp_client.client import MCPClient
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
@@ -536,9 +537,12 @@ class MCPServerManager:
client_secret=client_secret_value
or getattr(mcp_server, "client_secret", None),
scopes=resolved_scopes,
- authorization_url=getattr(mcp_oauth_metadata, "authorization_url", None),
- token_url=getattr(mcp_oauth_metadata, "token_url", None),
- registration_url=getattr(mcp_oauth_metadata, "registration_url", None),
+ authorization_url=mcp_server.authorization_url
+ or getattr(mcp_oauth_metadata, "authorization_url", None),
+ token_url=mcp_server.token_url
+ or getattr(mcp_oauth_metadata, "token_url", None),
+ registration_url=mcp_server.registration_url
+ or getattr(mcp_oauth_metadata, "registration_url", None),
command=getattr(mcp_server, "command", None),
args=getattr(mcp_server, "args", None) or [],
env=env_dict,
@@ -548,7 +552,7 @@ class MCPServerManager:
)
return new_server
- async def add_update_server(self, mcp_server: LiteLLM_MCPServerTable):
+ async def add_server(self, mcp_server: LiteLLM_MCPServerTable):
try:
if mcp_server.server_id not in self.registry:
new_server = await self.build_mcp_server_from_table(mcp_server)
@@ -559,6 +563,17 @@ class MCPServerManager:
verbose_logger.debug(f"Failed to add MCP server: {str(e)}")
raise e
+ async def update_server(self, mcp_server: LiteLLM_MCPServerTable):
+ try:
+ if mcp_server.server_id in self.registry:
+ new_server = await self.build_mcp_server_from_table(mcp_server)
+ self.registry[mcp_server.server_id] = new_server
+ verbose_logger.debug(f"Updated MCP Server: {new_server.name}")
+
+ except Exception as e:
+ verbose_logger.debug(f"Failed to udpate MCP server: {str(e)}")
+ raise e
+
def get_all_mcp_server_ids(self) -> Set[str]:
"""
Get all MCP server IDs
@@ -1662,11 +1677,11 @@ class MCPServerManager:
)
try:
- # Use standard pre_call_hook with call_type="mcp_call"
+ # Use standard pre_call_hook
modified_data = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_auth, # type: ignore
data=synthetic_llm_data,
- call_type="mcp_call", # type: ignore
+ call_type=CallTypes.call_mcp_tool.value,
)
if modified_data:
# Convert response back to MCP format and apply modifications
@@ -1723,7 +1738,7 @@ class MCPServerManager:
proxy_logging_obj.during_call_hook(
user_api_key_dict=user_api_key_auth,
data=synthetic_llm_data,
- call_type="mcp_call", # type: ignore
+ call_type=CallTypes.call_mcp_tool.value,
)
)
@@ -1879,7 +1894,7 @@ class MCPServerManager:
#########################################################
# Pre MCP Tool Call Hook
# Allow validation and modification of tool calls before execution
- # Using standard pre_call_hook with call_type="mcp_call"
+ # Using standard pre_call_hook
#########################################################
if proxy_logging_obj:
await self.pre_call_tool_check(
@@ -2040,7 +2055,7 @@ class MCPServerManager:
verbose_logger.debug(
f"Adding server to registry: {server.server_id} ({server.server_name})"
)
- await self.add_update_server(server)
+ await self.add_server(server)
verbose_logger.debug(
f"Registry now contains {len(self.get_registry())} servers"
@@ -2127,7 +2142,7 @@ class MCPServerManager:
async def health_check_server(
self, server_id: str, mcp_auth_header: Optional[str] = None
- ) -> Dict[str, Any]:
+ ) -> LiteLLM_MCPServerTable:
"""
Perform a health check on a specific MCP server.
@@ -2138,206 +2153,186 @@ class MCPServerManager:
Returns:
Dict containing health check results
"""
- import time
from datetime import datetime
server = self.get_mcp_server_by_id(server_id)
if not server:
- return {
- "server_id": server_id,
- "server_name": None,
- "status": "unknown",
- "error": "Server not found",
- "last_health_check": datetime.now().isoformat(),
- "response_time_ms": None,
- }
-
- start_time = time.time()
- try:
- # Try to get tools from the server as a health check
- tools = await self._get_tools_from_server(server, mcp_auth_header)
- response_time = (time.time() - start_time) * 1000
-
- return {
- "server_id": server_id,
- "server_name": server.name,
- "status": "healthy",
- "tools_count": len(tools),
- "last_health_check": datetime.now().isoformat(),
- "response_time_ms": round(response_time, 2),
- "error": None,
- }
- except Exception as e:
- response_time = (time.time() - start_time) * 1000
- error_message = str(e)
-
- return {
- "server_id": server_id,
- "server_name": server.name,
- "status": "unhealthy",
- "last_health_check": datetime.now().isoformat(),
- "response_time_ms": round(response_time, 2),
- "error": error_message,
- }
-
- async def health_check_all_servers(
- self, mcp_auth_header: Optional[str] = None
- ) -> Dict[str, Any]:
- """
- Perform health checks on all MCP servers.
-
- Args:
- mcp_auth_header: Optional authentication header for the MCP servers
-
- Returns:
- Dict containing health check results for all servers
- """
- all_servers = self.get_registry()
- results = {}
-
- for server_id, server in all_servers.items():
- results[server_id] = await self.health_check_server(
- server_id, mcp_auth_header
+ verbose_logger.warning(f"MCP Server {server_id} not found")
+ return LiteLLM_MCPServerTable(
+ server_id=server_id,
+ server_name=None,
+ transport=MCPTransport.http, # Default transport for not found servers
+ status="unknown",
+ health_check_error="Server not found",
+ last_health_check=datetime.now(),
)
- return results
+ status: Literal["healthy", "unhealthy", "unknown"] = "unknown"
+ health_check_error = None
- async def health_check_allowed_servers(
- self,
- user_api_key_auth: Optional[UserAPIKeyAuth] = None,
- mcp_auth_header: Optional[str] = None,
- ) -> Dict[str, Any]:
- """
- Perform health checks on all MCP servers that the user has access to.
+ # Check if we should skip health check based on auth configuration
+ should_skip_health_check = False
- Args:
- user_api_key_auth: User authentication info for access control
- mcp_auth_header: Optional authentication header for the MCP servers
+ # Skip if auth_type is oauth2
+ if server.auth_type == MCPAuth.oauth2:
+ should_skip_health_check = True
+ # Skip if auth_type is not none and authentication_token is missing
+ elif (
+ server.auth_type
+ and server.auth_type != MCPAuth.none
+ and not server.authentication_token
+ ):
+ should_skip_health_check = True
- Returns:
- Dict containing health check results for accessible servers
- """
- # Get allowed servers for the user
- allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth)
+ if not should_skip_health_check:
+ extra_headers = {}
+ if server.static_headers:
+ extra_headers.update(server.static_headers)
- # Perform health checks on allowed servers
- results = {}
- for server_id in allowed_server_ids:
- results[server_id] = await self.health_check_server(
- server_id, mcp_auth_header
+ client = self._create_mcp_client(
+ server=server,
+ mcp_auth_header=None,
+ extra_headers=extra_headers,
+ stdio_env=None,
)
- return results
+ try:
+
+ async def _noop(session):
+ return "ok"
+
+ # Add timeout wrapper to prevent hanging
+ await asyncio.wait_for(client.run_with_session(_noop), timeout=10.0)
+ status = "healthy"
+ except asyncio.TimeoutError:
+ health_check_error = "Health check timed out after 10 seconds"
+ status = "unhealthy"
+ except Exception as e:
+ health_check_error = str(e)
+ status = "unhealthy"
+
+ return LiteLLM_MCPServerTable(
+ server_id=server.server_id,
+ server_name=server.server_name,
+ alias=server.alias,
+ description=(
+ server.mcp_info.get("description") if server.mcp_info else None
+ ),
+ url=server.url,
+ transport=server.transport,
+ auth_type=server.auth_type,
+ created_at=datetime.now(),
+ updated_at=datetime.now(),
+ teams=[],
+ mcp_access_groups=server.access_groups or [],
+ allowed_tools=server.allowed_tools or [],
+ extra_headers=server.extra_headers or [],
+ mcp_info=server.mcp_info,
+ static_headers=server.static_headers,
+ status=status,
+ last_health_check=datetime.now(),
+ health_check_error=health_check_error,
+ command=getattr(server, "command", None),
+ args=getattr(server, "args", None) or [],
+ env=getattr(server, "env", None) or {},
+ authorization_url=server.authorization_url,
+ token_url=server.token_url,
+ registration_url=server.registration_url,
+ )
async def get_all_mcp_servers_with_health_and_teams(
self,
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
- include_health: bool = True,
+ server_ids: Optional[List[str]] = None,
) -> List[LiteLLM_MCPServerTable]:
"""
Get all MCP servers that the user has access to, with health status and team information.
Args:
user_api_key_auth: User authentication info for access control
- include_health: Whether to include health check information
+ server_ids: Optional list of server IDs to filter. If provided, only these servers
+ will be checked (subject to access control). If None, all accessible servers are checked.
Returns:
List of MCP server objects with health and team data
"""
- from litellm.proxy._experimental.mcp_server.db import (
- get_all_mcp_servers,
- get_mcp_servers,
- )
- from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
- from litellm.proxy.proxy_server import prisma_client
# Get allowed server IDs
allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth)
- # Get servers from database
+ # Filter by requested server_ids if provided
+ if server_ids:
+ # Only check servers that are both requested AND accessible
+ target_server_ids = [sid for sid in server_ids if sid in allowed_server_ids]
+ else:
+ # Check all accessible servers
+ target_server_ids = allowed_server_ids
+
+ # Run health checks concurrently
+ tasks = [self.health_check_server(server_id) for server_id in target_server_ids]
+ results = await asyncio.gather(*tasks)
+
+ # Filter out None results (servers that were not found)
+ list_mcp_servers = [server for server in results if server is not None]
+
+ return list_mcp_servers
+
+ async def get_all_allowed_mcp_servers(
+ self,
+ user_api_key_auth: Optional[UserAPIKeyAuth] = None,
+ ) -> List[LiteLLM_MCPServerTable]:
+ """
+ Get all MCP servers that the user has access to.
+
+ Args:
+ user_api_key_auth: User authentication info for access control
+
+ Returns:
+ List of MCP server objects without health status
+ """
+ from datetime import datetime
+
+ # Get allowed server IDs
+ allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth)
+
list_mcp_servers: List[LiteLLM_MCPServerTable] = []
- if prisma_client is not None:
- list_mcp_servers = await get_mcp_servers(prisma_client, allowed_server_ids)
- # If admin, also get all servers from database
- if user_api_key_auth and _user_has_admin_view(user_api_key_auth):
- all_mcp_servers = await get_all_mcp_servers(prisma_client)
- for server in all_mcp_servers:
- if server.server_id not in allowed_server_ids:
- list_mcp_servers.append(server)
+ for server_id in allowed_server_ids:
+ server = self.get_mcp_server_by_id(server_id)
+ if not server:
+ verbose_logger.warning(f"MCP Server {server_id} not found in registry")
+ continue
- # Add config.yaml servers
- for _server_id, _server_config in self.config_mcp_servers.items():
- if _server_id in allowed_server_ids:
- list_mcp_servers.append(
- LiteLLM_MCPServerTable(
- **{
- **_server_config.model_dump(),
- "created_at": datetime.datetime.now(),
- "updated_at": datetime.datetime.now(),
- "description": (
- _server_config.mcp_info.get("description")
- if _server_config.mcp_info
- else None
- ),
- "allowed_tools": _server_config.allowed_tools or [],
- "mcp_info": _server_config.mcp_info,
- "mcp_access_groups": _server_config.access_groups or [],
- "extra_headers": _server_config.extra_headers or [],
- "command": getattr(_server_config, "command", None),
- "args": getattr(_server_config, "args", None) or [],
- "env": getattr(_server_config, "env", None) or {},
- }
- )
- )
-
- # Get team information for non-admin users
- server_to_teams_map: Dict[str, List[Dict[str, str]]] = {}
- if (
- user_api_key_auth
- and not _user_has_admin_view(user_api_key_auth)
- and prisma_client is not None
- ):
- teams = await prisma_client.db.litellm_teamtable.find_many(
- include={"object_permission": True}
+ # Build LiteLLM_MCPServerTable without health check
+ mcp_server_table = LiteLLM_MCPServerTable(
+ server_id=server.server_id,
+ server_name=server.server_name,
+ alias=server.alias,
+ description=(
+ server.mcp_info.get("description") if server.mcp_info else None
+ ),
+ url=server.url,
+ transport=server.transport,
+ auth_type=server.auth_type,
+ created_at=datetime.now(),
+ updated_at=datetime.now(),
+ teams=[],
+ mcp_access_groups=server.access_groups or [],
+ allowed_tools=server.allowed_tools or [],
+ extra_headers=server.extra_headers or [],
+ mcp_info=server.mcp_info,
+ static_headers=server.static_headers,
+ status=None, # No health check performed
+ last_health_check=None, # No health check performed
+ health_check_error=None,
+ command=getattr(server, "command", None),
+ args=getattr(server, "args", None) or [],
+ env=getattr(server, "env", None) or {},
+ authorization_url=server.authorization_url,
+ token_url=server.token_url,
+ registration_url=server.registration_url,
)
-
- user_teams = []
- for team in teams:
- if team.members_with_roles:
- for member in team.members_with_roles:
- if (
- "user_id" in member
- and member["user_id"] is not None
- and member["user_id"] == user_api_key_auth.user_id
- ):
- user_teams.append(team)
-
- # Create a mapping of server_id to teams that have access to it
- for team in user_teams:
- if team.object_permission and team.object_permission.mcp_servers:
- for server_id in team.object_permission.mcp_servers:
- if server_id not in server_to_teams_map:
- server_to_teams_map[server_id] = []
- server_to_teams_map[server_id].append(
- {
- "team_id": team.team_id,
- "team_alias": team.team_alias,
- "organization_id": team.organization_id,
- }
- )
-
- ## mark invalid servers w/ reason for being invalid
- valid_server_ids = self.get_all_mcp_server_ids()
- for server in list_mcp_servers:
- if server.server_id not in valid_server_ids:
- server.status = "unhealthy"
- ## try adding server to registry to get error
- try:
- await self.add_update_server(server)
- except Exception as e:
- server.health_check_error = str(e)
- server.health_check_error = "Server is not in in memory registry yet. This could be a temporary sync issue."
+ list_mcp_servers.append(mcp_server_table)
return list_mcp_servers
diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py
index e4969df131e..8cdb381cf39 100644
--- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py
+++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py
@@ -3,7 +3,9 @@ This module is used to generate MCP tools from OpenAPI specs.
"""
import json
+from pathlib import PurePosixPath
from typing import Any, Dict, Optional
+from urllib.parse import quote
import httpx
@@ -17,6 +19,29 @@ BASE_URL = ""
HEADERS: Dict[str, str] = {}
+def _sanitize_path_parameter_value(param_value: Any, param_name: str) -> str:
+ """Ensure path params cannot introduce directory traversal."""
+ if param_value is None:
+ return ""
+
+ value_str = str(param_value)
+ if value_str == "":
+ return ""
+
+ normalized_value = value_str.replace("\\", "/")
+ if "/" in normalized_value:
+ raise ValueError(
+ f"Path parameter '{param_name}' must not contain path separators"
+ )
+
+ if any(part in {".", ".."} for part in PurePosixPath(normalized_value).parts):
+ raise ValueError(
+ f"Path parameter '{param_name}' cannot include '.' or '..' segments"
+ )
+
+ return quote(value_str, safe="")
+
+
def load_openapi_spec(filepath: str) -> Dict[str, Any]:
"""Load OpenAPI specification from JSON file."""
with open(filepath, "r") as f:
@@ -144,12 +169,18 @@ def create_tool_function(
url = base_url + path
# Replace path parameters using original names from OpenAPI spec
+ # Apply path traversal validation and URL encoding
for param_name in path_params:
param_value = kwargs.get(param_name, "")
if param_value:
+ try:
+ # Sanitize and encode path parameter to prevent traversal attacks
+ safe_value = _sanitize_path_parameter_value(param_value, param_name)
+ except ValueError as exc:
+ return "Invalid path parameter: " + str(exc)
# Replace {param_name} or {{param_name}} in URL
- url = url.replace("{" + param_name + "}", str(param_value))
- url = url.replace("{{" + param_name + "}}", str(param_value))
+ url = url.replace("{" + param_name + "}", safe_value)
+ url = url.replace("{{" + param_name + "}}", safe_value)
# Build query params using original parameter names
params: Dict[str, Any] = {}
diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
index 891b52db7af..4c947b99ba3 100644
--- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
+++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py
@@ -1,5 +1,4 @@
import importlib
-import traceback
from typing import Dict, List, Optional, Union
from fastapi import APIRouter, Depends, Query, Request
@@ -347,17 +346,16 @@ if MCP_AVAILABLE:
except Exception as e:
verbose_logger.error(f"Error in MCP operation: {e}", exc_info=True)
- stack_trace = traceback.format_exc()
return {
"status": "error",
- "message": f"An internal error has occurred: {str(e)}",
- "stack_trace": stack_trace,
+ "message": "An internal error has occurred while testing the MCP server.",
}
- @router.post("/test/connection")
+ @router.post("/test/connection", dependencies=[Depends(user_api_key_auth)])
async def test_connection(
request: Request,
new_mcp_server_request: NewMCPServerRequest,
+ user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Test if we can connect to the provided MCP server before adding it
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 570a6bb9f3b..7abfe7a96bc 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -412,7 +412,6 @@ class LiteLLMRoutes(enum.Enum):
agent_routes = [
"/v1/agents",
"/agents",
-
"/a2a/{agent_id}",
"/a2a/{agent_id}/message/send",
"/a2a/{agent_id}/message/stream",
@@ -830,9 +829,9 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase):
allowed_cache_controls: Optional[list] = []
config: Optional[dict] = {}
permissions: Optional[dict] = {}
- model_max_budget: Optional[dict] = (
- {}
- ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {}
+ model_max_budget: Optional[
+ dict
+ ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {}
model_config = ConfigDict(protected_namespaces=())
model_rpm_limit: Optional[dict] = None
@@ -1035,6 +1034,9 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
command: Optional[str] = None
args: List[str] = Field(default_factory=list)
env: Dict[str, str] = Field(default_factory=dict)
+ authorization_url: Optional[str] = None
+ token_url: Optional[str] = None
+ registration_url: Optional[str] = None
@model_validator(mode="before")
@classmethod
@@ -1092,6 +1094,9 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
command: Optional[str] = None
args: List[str] = Field(default_factory=list)
env: Dict[str, str] = Field(default_factory=dict)
+ authorization_url: Optional[str] = None
+ token_url: Optional[str] = None
+ registration_url: Optional[str] = None
@model_validator(mode="before")
@classmethod
@@ -1141,6 +1146,9 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
command: Optional[str] = None
args: List[str] = Field(default_factory=list)
env: Dict[str, str] = Field(default_factory=dict)
+ authorization_url: Optional[str] = None
+ token_url: Optional[str] = None
+ registration_url: Optional[str] = None
class MakeMCPServersPublicRequest(LiteLLMPydanticObjectBase):
@@ -1160,6 +1168,9 @@ class NewSkillRequest(LiteLLMPydanticObjectBase):
file_name: Optional[str] = None # Original filename
file_type: Optional[str] = None # MIME type (e.g., "application/zip")
metadata: Optional[Dict[str, Any]] = None
+ authorization_url: Optional[str] = None
+ token_url: Optional[str] = None
+ registration_url: Optional[str] = None
class UpdateSkillRequest(LiteLLMPydanticObjectBase):
@@ -1347,12 +1358,12 @@ class NewCustomerRequest(BudgetNewRequest):
blocked: bool = False # allow/disallow requests for this end-user
budget_id: Optional[str] = None # give either a budget_id or max_budget
spend: Optional[float] = None
- allowed_model_region: Optional[AllowedModelRegion] = (
- None # require all user requests to use models in this specific region
- )
- default_model: Optional[str] = (
- None # if no equivalent model in allowed region - default all requests to this model
- )
+ allowed_model_region: Optional[
+ AllowedModelRegion
+ ] = None # require all user requests to use models in this specific region
+ default_model: Optional[
+ str
+ ] = None # if no equivalent model in allowed region - default all requests to this model
@model_validator(mode="before")
@classmethod
@@ -1374,12 +1385,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase):
blocked: bool = False # allow/disallow requests for this end-user
max_budget: Optional[float] = None
budget_id: Optional[str] = None # give either a budget_id or max_budget
- allowed_model_region: Optional[AllowedModelRegion] = (
- None # require all user requests to use models in this specific region
- )
- default_model: Optional[str] = (
- None # if no equivalent model in allowed region - default all requests to this model
- )
+ allowed_model_region: Optional[
+ AllowedModelRegion
+ ] = None # require all user requests to use models in this specific region
+ default_model: Optional[
+ str
+ ] = None # if no equivalent model in allowed region - default all requests to this model
class DeleteCustomerRequest(LiteLLMPydanticObjectBase):
@@ -1464,15 +1475,15 @@ class NewTeamRequest(TeamBase):
] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm
model_tpm_limit: Optional[Dict[str, int]] = None
- team_member_budget: Optional[float] = (
- None # allow user to set a budget for all team members
- )
- team_member_rpm_limit: Optional[int] = (
- None # allow user to set RPM limit for all team members
- )
- team_member_tpm_limit: Optional[int] = (
- None # allow user to set TPM limit for all team members
- )
+ team_member_budget: Optional[
+ float
+ ] = None # allow user to set a budget for all team members
+ team_member_rpm_limit: Optional[
+ int
+ ] = None # allow user to set RPM limit for all team members
+ team_member_tpm_limit: Optional[
+ int
+ ] = None # allow user to set TPM limit for all team members
team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m"
allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
@@ -1558,9 +1569,9 @@ class BlockKeyRequest(LiteLLMPydanticObjectBase):
class AddTeamCallback(LiteLLMPydanticObjectBase):
callback_name: str
- callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = (
- "success_and_failure"
- )
+ callback_type: Optional[
+ Literal["success", "failure", "success_and_failure"]
+ ] = "success_and_failure"
callback_vars: Dict[str, str]
@model_validator(mode="before")
@@ -1788,9 +1799,10 @@ class DynamoDBArgs(LiteLLMPydanticObjectBase):
class PassThroughGuardrailSettings(LiteLLMPydanticObjectBase):
"""
Settings for a specific guardrail on a passthrough endpoint.
-
+
Allows field-level targeting for guardrail execution.
"""
+
request_fields: Optional[List[str]] = Field(
default=None,
description="JSONPath expressions for input field targeting (pre_call). Examples: 'query', 'documents[*].text', 'messages[*].content'. If not specified, guardrail runs on entire request payload.",
@@ -1871,9 +1883,9 @@ class ConfigList(LiteLLMPydanticObjectBase):
stored_in_db: Optional[bool]
field_default_value: Any
premium_field: bool = False
- nested_fields: Optional[List[FieldDetail]] = (
- None # For nested dictionary or Pydantic fields
- )
+ nested_fields: Optional[
+ List[FieldDetail]
+ ] = None # For nested dictionary or Pydantic fields
class UserHeaderMapping(LiteLLMPydanticObjectBase):
@@ -2152,6 +2164,7 @@ class UserAPIKeyAuth(
user_rpm_limit: Optional[int] = None
user_email: Optional[str] = None
request_route: Optional[str] = None
+ user: Optional[Any] = None # Expanded user object when expand=user is used
model_config = ConfigDict(arbitrary_types_allowed=True)
@@ -2258,9 +2271,9 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase):
budget_id: Optional[str] = None
created_at: datetime
updated_at: datetime
- user: Optional[Any] = (
- None # You might want to replace 'Any' with a more specific type if available
- )
+ user: Optional[
+ Any
+ ] = None # You might want to replace 'Any' with a more specific type if available
litellm_budget_table: Optional[LiteLLM_BudgetTable] = None
model_config = ConfigDict(protected_namespaces=())
@@ -2702,7 +2715,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
"TRACELOOP_API_KEY",
],
ui_callback_name="Traceloop",
- )
+ )
class SpendLogsMetadata(TypedDict):
@@ -2736,9 +2749,7 @@ class SpendLogsMetadata(TypedDict):
cold_storage_object_key: Optional[
str
] # S3/GCS object key for cold storage retrieval
- litellm_overhead_time_ms: Optional[
- float
- ] # LiteLLM overhead time in milliseconds
+ litellm_overhead_time_ms: Optional[float] # LiteLLM overhead time in milliseconds
cost_breakdown: Optional[
CostBreakdown
] # Detailed cost breakdown (input_cost, output_cost, margin, discount, etc.)
@@ -3222,9 +3233,9 @@ class TeamModelDeleteRequest(BaseModel):
# Organization Member Requests
class OrganizationMemberAddRequest(OrgMemberAddRequest):
organization_id: str
- max_budget_in_organization: Optional[float] = (
- None # Users max budget within the organization
- )
+ max_budget_in_organization: Optional[
+ float
+ ] = None # Users max budget within the organization
class OrganizationMemberDeleteRequest(MemberDeleteRequest):
@@ -3439,9 +3450,9 @@ class ProviderBudgetResponse(LiteLLMPydanticObjectBase):
Maps provider names to their budget configs.
"""
- providers: Dict[str, ProviderBudgetResponseObject] = (
- {}
- ) # Dictionary mapping provider names to their budget configurations
+ providers: Dict[
+ str, ProviderBudgetResponseObject
+ ] = {} # Dictionary mapping provider names to their budget configurations
class ProxyStateVariables(TypedDict):
@@ -3560,8 +3571,16 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
default=None,
description="If no team_id given, default permissions/spend-tracking to this team.s",
)
+ team_alias_jwt_field: Optional[str] = Field(
+ default=None,
+ description="The field in the JWT token that stores the team name/alias. Will be resolved to team_id via database lookup.",
+ )
org_id_jwt_field: Optional[str] = None
+ org_alias_jwt_field: Optional[str] = Field(
+ default=None,
+ description="The field in the JWT token that stores the organization name/alias. Will be resolved to org_id via database lookup.",
+ )
user_id_jwt_field: Optional[str] = None
user_email_jwt_field: Optional[str] = None
user_allowed_email_domain: Optional[str] = None
@@ -3576,9 +3595,9 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
enforce_rbac: bool = False
roles_jwt_field: Optional[str] = None # v2 on role mappings
role_mappings: Optional[List[RoleMapping]] = None
- object_id_jwt_field: Optional[str] = (
- None # can be either user / team, inferred from the role mapping
- )
+ object_id_jwt_field: Optional[
+ str
+ ] = None # can be either user / team, inferred from the role mapping
scope_mappings: Optional[List[ScopeMapping]] = None
enforce_scope_based_access: bool = False
enforce_team_based_model_access: bool = False
@@ -3729,13 +3748,16 @@ class DailyOrganizationSpendTransaction(BaseDailySpendTransaction):
class DailyUserSpendTransaction(BaseDailySpendTransaction):
user_id: str
+
class DailyEndUserSpendTransaction(BaseDailySpendTransaction):
end_user_id: str
+
class DailyTagSpendTransaction(BaseDailySpendTransaction):
request_id: Optional[str]
tag: str
+
class DailyAgentSpendTransaction(BaseDailySpendTransaction):
agent_id: str
@@ -3767,8 +3789,8 @@ class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase):
flat_model_file_ids: List[str]
created_by: Optional[str]
updated_by: Optional[str]
- storage_backend: Optional[str] = None
- storage_url: Optional[str] = None
+ storage_backend: Optional[str] = None
+ storage_url: Optional[str] = None
class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase):
@@ -3799,4 +3821,4 @@ class LiteLLM_ManagedVectorStoresTable(LiteLLMPydanticObjectBase):
class ResponseLiteLLM_ManagedVectorStore(TypedDict, total=False):
- vector_store: LiteLLM_ManagedVectorStoresTable
+ vector_store: LiteLLM_ManagedVectorStoresTable
\ No newline at end of file
diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index e2e90abeb1b..de4973ecc69 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -1372,6 +1372,195 @@ async def get_team_object(
)
+@log_db_metrics
+async def get_team_object_by_alias(
+ team_alias: str,
+ prisma_client: Optional[PrismaClient],
+ user_api_key_cache: DualCache,
+ parent_otel_span: Optional["Span"] = None,
+ proxy_logging_obj: Optional[ProxyLogging] = None,
+) -> LiteLLM_TeamTableCachedObj:
+ """
+ Look up a team by its team_alias (name) in the database.
+
+ Args:
+ team_alias: The team name/alias to look up
+ prisma_client: Database client
+ user_api_key_cache: Cache for storing results
+ parent_otel_span: Optional OpenTelemetry span
+ proxy_logging_obj: Optional proxy logging object
+
+ Returns:
+ LiteLLM_TeamTableCachedObj: The team object if found
+
+ Raises:
+ HTTPException: If team doesn't exist or multiple teams have the same alias
+ """
+ if prisma_client is None:
+ raise Exception(
+ "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys"
+ )
+
+ # Check cache first (keyed by alias)
+ cache_key = "team_alias:{}".format(team_alias)
+
+ cached_team_obj = await _get_team_object_from_cache(
+ key=cache_key,
+ proxy_logging_obj=proxy_logging_obj,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=parent_otel_span,
+ )
+
+ if cached_team_obj is not None:
+ return cached_team_obj
+
+ # Query database by team_alias
+ try:
+ teams = await prisma_client.db.litellm_teamtable.find_many(
+ where={"team_alias": team_alias}
+ )
+
+ if not teams:
+ raise HTTPException(
+ status_code=404,
+ detail={
+ "error": f"Team with alias '{team_alias}' doesn't exist in db. Create team via `/team/new` call."
+ },
+ )
+
+ if len(teams) > 1:
+ raise HTTPException(
+ status_code=400,
+ detail={
+ "error": f"Multiple teams found with alias '{team_alias}'. Please use team_id_jwt_field instead or ensure team aliases are unique."
+ },
+ )
+
+ team = teams[0]
+ team_obj = LiteLLM_TeamTableCachedObj(**team.model_dump())
+
+ # Cache the result by both alias and team_id
+ await user_api_key_cache.async_set_cache(
+ key=cache_key,
+ value=team_obj,
+ ttl=DEFAULT_IN_MEMORY_TTL,
+ )
+ # Also cache by team_id for consistency
+ team_id_cache_key = "team_id:{}".format(team_obj.team_id)
+ await user_api_key_cache.async_set_cache(
+ key=team_id_cache_key,
+ value=team_obj,
+ ttl=DEFAULT_IN_MEMORY_TTL,
+ )
+
+ return team_obj
+
+ except HTTPException:
+ raise
+ except Exception as e:
+ verbose_proxy_logger.exception(
+ "Error looking up team by alias: %s", team_alias
+ )
+ raise HTTPException(
+ status_code=500,
+ detail={
+ "error": f"Error looking up team by alias '{team_alias}': {str(e)}"
+ },
+ )
+
+
+@log_db_metrics
+async def get_org_object_by_alias(
+ org_alias: str,
+ prisma_client: Optional[PrismaClient],
+ user_api_key_cache: DualCache,
+ parent_otel_span: Optional["Span"] = None,
+ proxy_logging_obj: Optional[ProxyLogging] = None,
+) -> Optional[LiteLLM_OrganizationTable]:
+ """
+ Look up an organization by its organization_alias in the database.
+
+ Args:
+ org_alias: The organization name/alias to look up
+ prisma_client: Database client
+ user_api_key_cache: Cache for storing results
+ parent_otel_span: Optional OpenTelemetry span
+ proxy_logging_obj: Optional proxy logging object
+
+ Returns:
+ LiteLLM_OrganizationTable if found, None otherwise
+
+ Raises:
+ HTTPException: If organization not found or multiple orgs have the same alias
+ """
+ if prisma_client is None:
+ raise Exception(
+ "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys"
+ )
+
+ # Check cache first (keyed by alias)
+ cache_key = "org_alias:{}".format(org_alias)
+ cached_org_obj = await user_api_key_cache.async_get_cache(key=cache_key)
+ if cached_org_obj is not None:
+ if isinstance(cached_org_obj, dict):
+ return LiteLLM_OrganizationTable(**cached_org_obj)
+ elif isinstance(cached_org_obj, LiteLLM_OrganizationTable):
+ return cached_org_obj
+
+ # Query database by organization_alias
+ try:
+ orgs = await prisma_client.db.litellm_organizationtable.find_many(
+ where={"organization_alias": org_alias}
+ )
+
+ if not orgs:
+ raise HTTPException(
+ status_code=404,
+ detail={
+ "error": f"Organization with alias '{org_alias}' doesn't exist in db. Create organization via `/organization/new` call."
+ },
+ )
+
+ if len(orgs) > 1:
+ raise HTTPException(
+ status_code=400,
+ detail={
+ "error": f"Multiple organizations found with alias '{org_alias}'. Please use org_id_jwt_field instead or ensure organization aliases are unique."
+ },
+ )
+
+ org = orgs[0]
+ org_obj = LiteLLM_OrganizationTable(**org.model_dump())
+
+ # Cache the result
+ await user_api_key_cache.async_set_cache(
+ key=cache_key,
+ value=org_obj.model_dump(),
+ ttl=DEFAULT_IN_MEMORY_TTL,
+ )
+ # Also cache by org_id for consistency
+ await user_api_key_cache.async_set_cache(
+ key="org_id:{}".format(org_obj.organization_id),
+ value=org_obj.model_dump(),
+ ttl=DEFAULT_IN_MEMORY_TTL,
+ )
+
+ return org_obj
+
+ except HTTPException:
+ raise
+ except Exception as e:
+ verbose_proxy_logger.exception(
+ "Error looking up organization by alias: %s", org_alias
+ )
+ raise HTTPException(
+ status_code=500,
+ detail={
+ "error": f"Error looking up organization by alias '{org_alias}': {str(e)}"
+ },
+ )
+
+
class ExperimentalUIJWTToken:
@staticmethod
def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str:
diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py
index 17ff0de9f7b..33667b5d8d9 100644
--- a/litellm/proxy/auth/handle_jwt.py
+++ b/litellm/proxy/auth/handle_jwt.py
@@ -48,10 +48,12 @@ from .auth_checks import (
get_actual_routes,
get_end_user_object,
get_org_object,
+ get_org_object_by_alias,
get_role_based_models,
get_role_based_routes,
get_team_membership,
get_team_object,
+ get_team_object_by_alias,
get_user_object,
)
@@ -194,10 +196,13 @@ class JWTHandler:
def is_required_team_id(self) -> bool:
"""
Returns:
- - True: if 'team_id_jwt_field' is set
- - False: if not
+ - True: if 'team_id_jwt_field' or 'team_alias_jwt_field' is set
+ - False: if neither is set
"""
- if self.litellm_jwtauth.team_id_jwt_field is None:
+ if (
+ self.litellm_jwtauth.team_id_jwt_field is None
+ and self.litellm_jwtauth.team_alias_jwt_field is None
+ ):
return False
return True
@@ -240,6 +245,31 @@ class JWTHandler:
team_id = default_value
return team_id
+ def get_team_alias(self, token: dict, default_value: Optional[str]) -> Optional[str]:
+ """
+ Extract team name/alias from JWT token using the configured team_alias_jwt_field.
+
+ Args:
+ token: The decoded JWT token dictionary
+ default_value: Default value to return if field not found
+
+ Returns:
+ The team alias from the token, or default_value if not found
+ """
+ try:
+ if self.litellm_jwtauth.team_alias_jwt_field is not None:
+ team_alias = get_nested_value(
+ data=token,
+ key_path=self.litellm_jwtauth.team_alias_jwt_field,
+ default=default_value,
+ )
+ return team_alias
+ else:
+ team_alias = None
+ except KeyError:
+ team_alias = default_value
+ return team_alias
+
def is_upsert_user_id(self, valid_user_email: Optional[bool] = None) -> bool:
"""
Returns:
@@ -383,6 +413,31 @@ class JWTHandler:
org_id = default_value
return org_id
+ def get_org_alias(self, token: dict, default_value: Optional[str]) -> Optional[str]:
+ """
+ Extract organization name/alias from JWT token using the configured org_alias_jwt_field.
+
+ Args:
+ token: The decoded JWT token dictionary
+ default_value: Default value to return if field not found
+
+ Returns:
+ The organization alias from the token, or default_value if not found
+ """
+ try:
+ if self.litellm_jwtauth.org_alias_jwt_field is not None:
+ org_alias = get_nested_value(
+ data=token,
+ key_path=self.litellm_jwtauth.org_alias_jwt_field,
+ default=default_value,
+ )
+ return org_alias
+ else:
+ org_alias = None
+ except KeyError:
+ org_alias = default_value
+ return org_alias
+
def get_scopes(self, token: dict) -> List[str]:
try:
if isinstance(token["scope"], str):
@@ -813,18 +868,14 @@ class JWTAuthManager:
parent_otel_span: Optional[Span],
proxy_logging_obj: ProxyLogging,
) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]:
- """Find and validate specific team ID"""
+ """Find and validate specific team ID from team_id_jwt_field or team_alias_jwt_field"""
individual_team_id = jwt_handler.get_team_id(
token=jwt_valid_token, default_value=None
)
- if not individual_team_id and jwt_handler.is_required_team_id() is True:
- raise Exception(
- f"No team id found in token. Checked team_id field '{jwt_handler.litellm_jwtauth.team_id_jwt_field}'"
- )
-
- ## VALIDATE TEAM OBJECT ###
team_object: Optional[LiteLLM_TeamTable] = None
+
+ # First try to get team by team_id
if individual_team_id:
team_object = await get_team_object(
team_id=individual_team_id,
@@ -834,6 +885,37 @@ class JWTAuthManager:
proxy_logging_obj=proxy_logging_obj,
team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert,
)
+ return individual_team_id, team_object
+
+ # If no team_id found, try to resolve via team_alias_jwt_field
+ team_alias = jwt_handler.get_team_alias(
+ token=jwt_valid_token, default_value=None
+ )
+ if team_alias:
+ verbose_proxy_logger.info(
+ f"JWT Auth: Resolving team by alias: '{team_alias}'"
+ )
+ team_object = await get_team_object_by_alias(
+ team_alias=team_alias,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=parent_otel_span,
+ proxy_logging_obj=proxy_logging_obj,
+ )
+ if team_object:
+ individual_team_id = team_object.team_id
+ verbose_proxy_logger.info(
+ f"JWT Auth: Resolved team_alias='{team_alias}' to team_id='{individual_team_id}'"
+ )
+ return individual_team_id, team_object
+
+ # Check if team is required but not found
+ if jwt_handler.is_required_team_id() is True:
+ team_id_field = jwt_handler.litellm_jwtauth.team_id_jwt_field
+ team_alias_field = jwt_handler.litellm_jwtauth.team_alias_jwt_field
+ raise Exception(
+ f"No team found in token. Checked team_id field '{team_id_field}' and team_alias field '{team_alias_field}'"
+ )
return individual_team_id, team_object
@@ -942,13 +1024,16 @@ class JWTAuthManager:
parent_otel_span: Optional[Span],
proxy_logging_obj: ProxyLogging,
route: str,
+ org_alias: Optional[str] = None,
) -> Tuple[
Optional[LiteLLM_UserTable],
Optional[LiteLLM_OrganizationTable],
- Optional[LiteLLM_EndUserTable],
+ Optional[LiteLLM_EndUserTable],
Optional[LiteLLM_TeamMembership],
]:
- """Get user, org, and end user objects"""
+ """Get user, org, and end user objects. Also resolves org aliases to IDs if configured."""
+
+ # Get org object - first try by ID, then by alias
org_object: Optional[LiteLLM_OrganizationTable] = None
if org_id:
org_object = (
@@ -962,6 +1047,21 @@ class JWTAuthManager:
if org_id
else None
)
+ elif org_alias:
+ verbose_proxy_logger.info(
+ f"JWT Auth: Resolving org by alias: '{org_alias}'"
+ )
+ org_object = await get_org_object_by_alias(
+ org_alias=org_alias,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=parent_otel_span,
+ proxy_logging_obj=proxy_logging_obj,
+ )
+ if org_object:
+ verbose_proxy_logger.info(
+ f"JWT Auth: Resolved org_alias='{org_alias}' to org_id='{org_object.organization_id}'"
+ )
user_object: Optional[LiteLLM_UserTable] = None
if user_id:
@@ -1304,6 +1404,8 @@ class JWTAuthManager:
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
+ # Extract alias fields for resolution (if configured)
+ org_alias = jwt_handler.get_org_alias(token=jwt_valid_token, default_value=None)
# Get other objects
user_object, org_object, end_user_object, team_membership_object = (
@@ -1320,9 +1422,13 @@ class JWTAuthManager:
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
+ org_alias=org_alias,
)
)
+ # Derive org_id from org_object if resolved by alias
+ resolved_org_id = org_object.organization_id if org_object else org_id
+
await JWTAuthManager.sync_user_role_and_teams(
jwt_handler=jwt_handler,
jwt_valid_token=jwt_valid_token,
@@ -1345,10 +1451,9 @@ class JWTAuthManager:
)
# check if user is proxy admin
- if user_object and user_object.user_role == LitellmUserRoles.PROXY_ADMIN:
- is_proxy_admin = True
- else:
- is_proxy_admin = False
+ is_proxy_admin = bool(
+ user_object and user_object.user_role == LitellmUserRoles.PROXY_ADMIN
+ )
return JWTAuthBuilderResult(
is_proxy_admin=is_proxy_admin,
@@ -1356,7 +1461,7 @@ class JWTAuthManager:
team_object=team_object,
user_id=user_id,
user_object=user_object,
- org_id=org_id,
+ org_id=resolved_org_id, # Use resolved org_id (from alias lookup if applicable)
org_object=org_object,
end_user_id=end_user_id,
end_user_object=end_user_object,
diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py
index fb9757ca647..8cc33ce6cdd 100644
--- a/litellm/proxy/auth/login_utils.py
+++ b/litellm/proxy/auth/login_utils.py
@@ -34,6 +34,59 @@ from litellm.secret_managers.main import get_secret_bool
from litellm.types.proxy.ui_sso import ReturnedUITokenObject
+async def expire_previous_ui_session_tokens(
+ user_id: str, prisma_client: Optional[PrismaClient]
+) -> None:
+ """
+ Expire (block) all other valid UI session tokens for a user.
+
+ This prevents accumulation of multiple valid UI session tokens that
+ are supposed to be short-lived test keys. Only affects keys with
+ team_id = "litellm-dashboard" and that haven't expired yet.
+
+ Args:
+ user_id: The user ID whose previous UI session tokens should be expired
+ prisma_client: Database client for performing the update
+ """
+ if prisma_client is None:
+ return
+
+ try:
+ from datetime import datetime, timezone
+
+ current_time = datetime.now(timezone.utc)
+
+ # Find all unblocked AND non-expired UI session tokens for this user
+ ui_session_tokens = await prisma_client.db.litellm_verificationtoken.find_many(
+ where={
+ "user_id": user_id,
+ "team_id": "litellm-dashboard",
+ "OR": [
+ {"blocked": None}, # Tokens that have never been blocked (null)
+ {"blocked": False}, # Tokens explicitly set to not blocked
+ ],
+ "expires": {"gt": current_time}, # Only get tokens that haven't expired
+ }
+ )
+
+ if not ui_session_tokens:
+ return
+
+ # Block all the found tokens
+ tokens_to_block = [token.token for token in ui_session_tokens if token.token]
+
+ if tokens_to_block:
+ await prisma_client.db.litellm_verificationtoken.update_many(
+ where={"token": {"in": tokens_to_block}},
+ data={"blocked": True}
+ )
+
+ except Exception:
+ # Silently fail - don't block login if cleanup fails
+ # This is a best-effort operation
+ pass
+
+
def get_ui_credentials(master_key: Optional[str]) -> tuple[str, str]:
"""
Get UI username and password from environment variables or master key.
@@ -85,7 +138,7 @@ class LoginResult:
self.login_method = login_method
-async def authenticate_user(
+async def authenticate_user( # noqa: PLR0915
username: str,
password: str,
master_key: Optional[str],
@@ -174,6 +227,10 @@ async def authenticate_user(
)
if os.getenv("DATABASE_URL") is not None:
+ # Expire any previous UI session tokens for this user
+ await expire_previous_ui_session_tokens(
+ user_id=key_user_id, prisma_client=prisma_client
+ )
response = await generate_key_helper_fn(
request_type="key",
**{
@@ -260,6 +317,11 @@ async def authenticate_user(
hash_password, _password
):
if os.getenv("DATABASE_URL") is not None:
+ # Expire any previous UI session tokens for this user
+ await expire_previous_ui_session_tokens(
+ user_id=user_id, prisma_client=prisma_client
+ )
+
response = await generate_key_helper_fn(
request_type="key",
**{ # type: ignore
diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py
index 495d4db304c..9b53d9a3a80 100644
--- a/litellm/proxy/auth/user_api_key_auth.py
+++ b/litellm/proxy/auth/user_api_key_auth.py
@@ -551,6 +551,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
valid_token = UserAPIKeyAuth(
api_key=None,
team_id=team_id,
+ team_alias=(
+ team_object.team_alias if team_object is not None else None
+ ),
team_tpm_limit=(
team_object.tpm_limit if team_object is not None else None
),
diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
index 62c997659bd..59032189438 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
@@ -449,6 +449,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
prepared_request.headers,
)
+ event_type = (
+ GuardrailEventHooks.pre_call
+ if source == "INPUT"
+ else GuardrailEventHooks.post_call
+ )
+
try:
httpx_response = await self.async_handler.post(
url=prepared_request.url,
@@ -469,6 +475,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
+ event_type=event_type,
)
# Re-raise the exception to maintain existing behavior
raise
@@ -486,6 +493,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
+ event_type=event_type,
)
#########################################################
if httpx_response.status_code == 200:
@@ -605,10 +613,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
"""
Only raise exception for "BLOCKED" actions, not for "ANONYMIZED" actions.
- If `self.mask_request_content` or `self.mask_response_content` is set to `True`,
+ If `self.mask_request_content` or `self.mask_response_content` is set to `True`,
then use the output from the guardrail to mask the request or response content.
-
- However, even with masking enabled, content with action="BLOCKED" should still
+
+ However, even with masking enabled, content with action="BLOCKED" should still
raise an exception, only content with action="ANONYMIZED" should be masked.
"""
@@ -731,9 +739,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
#########################################################
########## 1. Make the Bedrock API request ##########
#########################################################
- bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = (
- None
- )
+ bedrock_guardrail_response: Optional[
+ Union[BedrockGuardrailResponse, str]
+ ] = None
try:
bedrock_guardrail_response = await self.make_bedrock_api_request(
source="INPUT", messages=filtered_messages, request_data=data
@@ -803,9 +811,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
#########################################################
########## 1. Make the Bedrock API request ##########
#########################################################
- bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = (
- None
- )
+ bedrock_guardrail_response: Optional[
+ Union[BedrockGuardrailResponse, str]
+ ] = None
try:
bedrock_guardrail_response = await self.make_bedrock_api_request(
source="INPUT", messages=filtered_messages, request_data=data
@@ -1296,11 +1304,6 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data=request_data,
)
- if bedrock_response.get("action") == "BLOCKED":
- raise Exception(
- f"Content blocked by Bedrock guardrail: {bedrock_response.get('reason', 'Unknown reason')}"
- )
-
# Apply any masking that was applied by the guardrail
output_list = bedrock_response.get("output")
diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py
index 6915286a2d7..59381149809 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py
@@ -97,6 +97,7 @@ class DynamoAIGuardrails(CustomGuardrail):
async def _call_dynamoai_guardrails(
self,
messages: List[Dict[str, Any]],
+ event_type: GuardrailEventHooks,
text_type: str = "input",
request_data: Optional[dict] = None,
) -> DynamoAIResponse:
@@ -157,6 +158,7 @@ class DynamoAIGuardrails(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
+ event_type=event_type,
)
return response_json
@@ -177,6 +179,7 @@ class DynamoAIGuardrails(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
+ event_type=event_type,
)
raise
@@ -332,6 +335,7 @@ class DynamoAIGuardrails(CustomGuardrail):
messages=_messages,
text_type="input",
request_data=data,
+ event_type=GuardrailEventHooks.pre_call,
)
verbose_proxy_logger.debug(
@@ -380,6 +384,7 @@ class DynamoAIGuardrails(CustomGuardrail):
messages=_messages,
text_type="input",
request_data=data,
+ event_type=GuardrailEventHooks.during_call,
)
verbose_proxy_logger.debug(
@@ -460,6 +465,7 @@ class DynamoAIGuardrails(CustomGuardrail):
messages=dynamoai_messages,
text_type="output",
request_data=data,
+ event_type=GuardrailEventHooks.post_call,
)
verbose_proxy_logger.debug(
diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py
index 55fa17c21e7..2fc05213640 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py
@@ -108,6 +108,7 @@ class IBMGuardrailDetector(CustomGuardrail):
async def _call_detector_server(
self,
contents: List[str],
+ event_type: GuardrailEventHooks,
request_data: Optional[dict] = None,
) -> List[List[IBMDetectorDetection]]:
"""
@@ -142,7 +143,6 @@ class IBMGuardrailDetector(CustomGuardrail):
)
try:
-
response = await self.async_handler.post(
url=self.api_url,
json=payload,
@@ -172,6 +172,7 @@ class IBMGuardrailDetector(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
+ event_type=event_type,
)
return response_json
@@ -192,6 +193,7 @@ class IBMGuardrailDetector(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
+ event_type=event_type,
)
raise
@@ -199,6 +201,7 @@ class IBMGuardrailDetector(CustomGuardrail):
async def _call_orchestrator(
self,
content: str,
+ event_type: GuardrailEventHooks,
request_data: Optional[dict] = None,
) -> List[IBMDetectorDetection]:
"""
@@ -258,6 +261,7 @@ class IBMGuardrailDetector(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
+ event_type=event_type,
)
return response_json.get("detections", [])
@@ -278,6 +282,7 @@ class IBMGuardrailDetector(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
+ event_type=event_type,
)
raise
@@ -472,6 +477,7 @@ class IBMGuardrailDetector(CustomGuardrail):
result = await self._call_detector_server(
contents=contents_to_check,
request_data=data,
+ event_type=GuardrailEventHooks.pre_call,
)
verbose_proxy_logger.debug(
@@ -500,6 +506,7 @@ class IBMGuardrailDetector(CustomGuardrail):
orchestrator_result = await self._call_orchestrator(
content=content,
request_data=data,
+ event_type=GuardrailEventHooks.pre_call,
)
verbose_proxy_logger.debug(
@@ -557,6 +564,7 @@ class IBMGuardrailDetector(CustomGuardrail):
result = await self._call_detector_server(
contents=contents_to_check,
request_data=data,
+ event_type=GuardrailEventHooks.during_call,
)
verbose_proxy_logger.debug(
@@ -585,6 +593,7 @@ class IBMGuardrailDetector(CustomGuardrail):
orchestrator_result = await self._call_orchestrator(
content=content,
request_data=data,
+ event_type=GuardrailEventHooks.during_call,
)
verbose_proxy_logger.debug(
@@ -673,6 +682,7 @@ class IBMGuardrailDetector(CustomGuardrail):
result = await self._call_detector_server(
contents=contents_to_check,
request_data=data,
+ event_type=GuardrailEventHooks.post_call,
)
verbose_proxy_logger.debug(
@@ -702,6 +712,7 @@ class IBMGuardrailDetector(CustomGuardrail):
orchestrator_result = await self._call_orchestrator(
content=content,
request_data=data,
+ event_type=GuardrailEventHooks.post_call,
)
verbose_proxy_logger.debug(
diff --git a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py
index 6d4ed089818..953275acf14 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py
@@ -83,6 +83,7 @@ class JavelinGuardrail(CustomGuardrail):
async def call_javelin_guard(
self,
request: JavelinGuardRequest,
+ event_type: GuardrailEventHooks,
) -> JavelinGuardResponse:
"""
Call the Javelin guard API.
@@ -158,6 +159,7 @@ class JavelinGuardrail(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
+ event_type=event_type,
)
async def async_pre_call_hook(
@@ -208,7 +210,9 @@ class JavelinGuardrail(CustomGuardrail):
config=self.config if self.config else {},
)
- javelin_response = await self.call_javelin_guard(request=javelin_guard_request)
+ javelin_response = await self.call_javelin_guard(
+ request=javelin_guard_request, event_type=GuardrailEventHooks.pre_call
+ )
assessments = javelin_response.get("assessments", [])
reject_prompt = ""
diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
index 6d98866eadf..732331349e0 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
@@ -70,6 +70,7 @@ class LakeraAIGuardrail(CustomGuardrail):
self,
messages: List[AllMessageValues],
request_data: Dict,
+ event_type: GuardrailEventHooks,
) -> Tuple[LakeraAIResponse, Dict]:
"""
Call the Lakera AI v2 guard API.
@@ -128,6 +129,7 @@ class LakeraAIGuardrail(CustomGuardrail):
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
masked_entity_count=masked_entity_count,
+ event_type=event_type,
)
def _mask_pii_in_messages(
@@ -214,6 +216,7 @@ class LakeraAIGuardrail(CustomGuardrail):
lakera_guardrail_response, masked_entity_count = await self.call_v2_guard(
messages=new_messages,
request_data=data,
+ event_type=GuardrailEventHooks.pre_call,
)
#########################################################
@@ -279,6 +282,7 @@ class LakeraAIGuardrail(CustomGuardrail):
lakera_guardrail_response, masked_entity_count = await self.call_v2_guard(
messages=new_messages,
request_data=data,
+ event_type=GuardrailEventHooks.during_call,
)
#########################################################
diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
index 51136c29eca..a12eb2486d2 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
@@ -295,7 +295,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
filters = (
list(filter_results.values())
if isinstance(filter_results, dict)
- else filter_results if isinstance(filter_results, list) else []
+ else filter_results
+ if isinstance(filter_results, list)
+ else []
)
# Prefer sanitized text from deidentifyResult if present
@@ -327,6 +329,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
start_time: Optional[float] = None,
end_time: Optional[float] = None,
duration: Optional[float] = None,
+ event_type: Optional[GuardrailEventHooks] = None,
):
"""
Override to store only the Model Armor API response, not the entire data dict.
@@ -351,6 +354,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
duration=duration,
start_time=start_time,
end_time=end_time,
+ event_type=event_type,
)
return response
diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py
index a0ea90ccf21..1794751a08c 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py
@@ -119,9 +119,7 @@ class NomaGuardrail(CustomGuardrail):
self.api_base = api_base or os.environ.get(
"NOMA_API_BASE", NomaGuardrail._DEFAULT_API_BASE
)
- self.application_id = application_id or os.environ.get(
- "NOMA_APPLICATION_ID"
- )
+ self.application_id = application_id or os.environ.get("NOMA_APPLICATION_ID")
self.default_application_id = "litellm"
if monitor_mode is None:
@@ -163,6 +161,7 @@ class NomaGuardrail(CustomGuardrail):
self,
request_data: dict,
user_auth: UserAPIKeyAuth,
+ event_type: Optional[GuardrailEventHooks] = None,
) -> Optional[str]:
"""Shared logic for processing user message checks"""
start_time = datetime.now()
@@ -213,6 +212,7 @@ class NomaGuardrail(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
+ event_type=event_type,
)
if self.monitor_mode:
@@ -242,6 +242,7 @@ class NomaGuardrail(CustomGuardrail):
request_data: dict,
response: LLMResponse,
user_auth: UserAPIKeyAuth,
+ event_type: Optional[GuardrailEventHooks] = None,
) -> Optional[str]:
"""Shared logic for processing LLM response checks"""
@@ -293,6 +294,7 @@ class NomaGuardrail(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
+ event_type=event_type,
)
if self.monitor_mode:
@@ -578,7 +580,6 @@ class NomaGuardrail(CustomGuardrail):
data: dict,
call_type: CallTypesLiteral,
) -> Optional[Union[Exception, str, dict]]:
-
verbose_proxy_logger.debug("Running Noma pre-call hook")
if (
@@ -602,7 +603,9 @@ class NomaGuardrail(CustomGuardrail):
return data
try:
- return await self._check_user_message(data, user_api_key_dict)
+ return await self._check_user_message(
+ data, user_api_key_dict, GuardrailEventHooks.pre_call
+ )
except NomaBlockedMessage:
# Blocked requests were already logged in _process_user_message_check with "blocked" status
raise
@@ -619,6 +622,7 @@ class NomaGuardrail(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=start_time.timestamp(),
duration=0.0,
+ event_type=GuardrailEventHooks.pre_call,
)
verbose_proxy_logger.error(f"Noma pre-call hook failed: {str(e)}")
@@ -650,7 +654,9 @@ class NomaGuardrail(CustomGuardrail):
return data
try:
- return await self._check_user_message(data, user_api_key_dict)
+ return await self._check_user_message(
+ data, user_api_key_dict, GuardrailEventHooks.during_call
+ )
except NomaBlockedMessage:
# Blocked requests were already logged in _process_user_message_check with "blocked" status
raise
@@ -667,6 +673,7 @@ class NomaGuardrail(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=start_time.timestamp(),
duration=0.0,
+ event_type=GuardrailEventHooks.during_call,
)
verbose_proxy_logger.error(f"Noma moderation hook failed: {str(e)}")
@@ -700,7 +707,9 @@ class NomaGuardrail(CustomGuardrail):
return response
try:
- return await self._check_llm_response(data, response, user_api_key_dict)
+ return await self._check_llm_response(
+ data, response, user_api_key_dict, GuardrailEventHooks.post_call
+ )
except NomaBlockedMessage:
# Blocked requests were already logged in _process_llm_response_check with "blocked" status
raise
@@ -717,6 +726,7 @@ class NomaGuardrail(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=start_time.timestamp(),
duration=0.0,
+ event_type=GuardrailEventHooks.post_call,
)
verbose_proxy_logger.error(f"Noma post-call hook failed: {str(e)}")
@@ -728,9 +738,12 @@ class NomaGuardrail(CustomGuardrail):
self,
request_data: dict,
user_auth: UserAPIKeyAuth,
+ event_type: Optional[GuardrailEventHooks] = None,
) -> Union[Exception, str, dict, None]:
"""Check user message for policy violations"""
- user_message = await self._process_user_message_check(request_data, user_auth)
+ user_message = await self._process_user_message_check(
+ request_data, user_auth, event_type
+ )
if not user_message:
return request_data
@@ -741,10 +754,11 @@ class NomaGuardrail(CustomGuardrail):
request_data: dict,
response: LLMResponse,
user_auth: UserAPIKeyAuth,
+ event_type: Optional[GuardrailEventHooks] = None,
) -> Any:
"""Check LLM response for policy violations"""
content = await self._process_llm_response_check(
- request_data, response, user_auth
+ request_data, response, user_auth, event_type
)
if not content:
return response
@@ -858,7 +872,10 @@ class NomaGuardrail(CustomGuardrail):
if isinstance(assembled_model_response, ModelResponse):
try:
processed_response = await self._check_llm_response(
- request_data, assembled_model_response, user_api_key_dict
+ request_data,
+ assembled_model_response,
+ user_api_key_dict,
+ GuardrailEventHooks.post_call,
)
except NomaBlockedMessage:
raise
diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py
index 88145ae9e47..02e481acddd 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py
@@ -24,6 +24,7 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
+from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import CallTypesLiteral, ModelResponse
if TYPE_CHECKING:
@@ -523,6 +524,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
scan_result: Dict[str, Any],
data: Dict[str, Any],
start_time: datetime,
+ event_type: GuardrailEventHooks,
is_response: bool = False,
) -> Optional[Dict[str, Any]]:
"""Handle API errors with fail-open/fail-closed logic."""
@@ -542,6 +544,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
+ event_type=event_type,
)
if scan_result.get("_always_block"):
@@ -735,7 +738,11 @@ class PanwPrismaAirsHandler(CustomGuardrail):
if scan_result.get("_is_transient") or scan_result.get("_always_block"):
return self._handle_api_error_with_logging(
- scan_result, data, start_time, is_response=False
+ scan_result,
+ data,
+ start_time,
+ is_response=False,
+ event_type=GuardrailEventHooks.pre_call,
)
end_time = datetime.now()
@@ -749,6 +756,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=(end_time - start_time).total_seconds(),
+ event_type=GuardrailEventHooks.pre_call,
)
action = scan_result.get("action", "block")
@@ -872,7 +880,11 @@ class PanwPrismaAirsHandler(CustomGuardrail):
if scan_result.get("_is_transient") or scan_result.get("_always_block"):
self._handle_api_error_with_logging(
- scan_result, data, start_time, is_response=True
+ scan_result,
+ data,
+ start_time,
+ is_response=True,
+ event_type=GuardrailEventHooks.post_call,
)
return response
@@ -887,6 +899,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=(end_time - start_time).total_seconds(),
+ event_type=GuardrailEventHooks.post_call,
)
action = scan_result.get("action", "block")
@@ -1066,7 +1079,11 @@ class PanwPrismaAirsHandler(CustomGuardrail):
if scan_result.get("_is_transient") or scan_result.get("_always_block"):
self._handle_api_error_with_logging(
- scan_result, request_data, start_time, is_response=True
+ scan_result,
+ request_data,
+ start_time,
+ is_response=True,
+ event_type=EventHooks.post_call,
)
for chunk in all_chunks:
yield chunk
@@ -1083,6 +1100,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=(end_time - start_time).total_seconds(),
+ event_type=EventHooks.post_call,
)
# Add guardrail to applied guardrails header for observability
diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
index a1bbf36ac0c..f66341fde5c 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
@@ -28,7 +28,6 @@ class UnifiedLLMGuardrails(CustomLogger):
self,
**kwargs,
):
-
# store kwargs as optional_params
self.optional_params = kwargs
@@ -63,6 +62,9 @@ class UnifiedLLMGuardrails(CustomLogger):
return data
event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call
+ if call_type == CallTypes.call_mcp_tool.value:
+ event_type = GuardrailEventHooks.pre_mcp_call
+
if (
guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type)
is not True
@@ -114,6 +116,9 @@ class UnifiedLLMGuardrails(CustomLogger):
return data
event_type: GuardrailEventHooks = GuardrailEventHooks.during_call
+ if call_type == CallTypes.call_mcp_tool.value:
+ event_type = GuardrailEventHooks.during_mcp_call
+
if (
guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type)
is not True
@@ -128,7 +133,10 @@ class UnifiedLLMGuardrails(CustomLogger):
endpoint_guardrail_translation_mappings = (
load_guardrail_translation_mappings()
)
- if call_type is not None and CallTypes(call_type) not in endpoint_guardrail_translation_mappings:
+ if (
+ call_type is not None
+ and CallTypes(call_type) not in endpoint_guardrail_translation_mappings
+ ):
return data
endpoint_translation = endpoint_guardrail_translation_mappings[
@@ -180,8 +188,8 @@ class UnifiedLLMGuardrails(CustomLogger):
call_type: Optional[CallTypesLiteral] = None
if user_api_key_dict.request_route is not None:
call_types = get_call_types_for_route(user_api_key_dict.request_route)
- if call_types is not None and len(call_types) > 0: # type: ignore
- call_type = call_types[0] # type: ignore
+ if call_types is not None and len(call_types) > 0: # type: ignore
+ call_type = call_types[0] # type: ignore
if call_type is None:
call_type = _infer_call_type(call_type=None, completion_response=response) # type: ignore
@@ -330,7 +338,6 @@ class UnifiedLLMGuardrails(CustomLogger):
# Process chunk based on sampling rate
if chunk_counter % sampling_rate == 0:
-
verbose_proxy_logger.debug(
"Processing streaming chunk %s (sampling_rate=%s) with guardrail %s",
chunk_counter,
diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py
index 2d86f74a41c..e43da32565a 100644
--- a/litellm/proxy/management_endpoints/budget_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py
@@ -55,6 +55,18 @@ async def new_budget(
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
+ # Validate budget values are not negative
+ if budget_obj.max_budget is not None and budget_obj.max_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"max_budget cannot be negative. Received: {budget_obj.max_budget}"}
+ )
+ if budget_obj.soft_budget is not None and budget_obj.soft_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"soft_budget cannot be negative. Received: {budget_obj.soft_budget}"}
+ )
+
# if no budget_reset_at date is set, but a budget_duration is given, then set budget_reset_at initially to the first completed duration interval in future
if budget_obj.budget_reset_at is None and budget_obj.budget_duration is not None:
budget_obj.budget_reset_at = datetime.utcnow() + timedelta(
@@ -107,6 +119,18 @@ async def update_budget(
if budget_obj.budget_id is None:
raise HTTPException(status_code=400, detail={"error": "budget_id is required"})
+ # Validate budget values are not negative
+ if budget_obj.max_budget is not None and budget_obj.max_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"max_budget cannot be negative. Received: {budget_obj.max_budget}"}
+ )
+ if budget_obj.soft_budget is not None and budget_obj.soft_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"soft_budget cannot be negative. Received: {budget_obj.soft_budget}"}
+ )
+
response = await prisma_client.db.litellm_budgettable.update(
where={"budget_id": budget_obj.budget_id},
data={
diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index cc2ac908149..4d72bc86257 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -1069,6 +1069,18 @@ async def generate_key_fn(
verbose_proxy_logger.debug("entered /key/generate")
+ # Validate budget values are not negative
+ if data.max_budget is not None and data.max_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"}
+ )
+ if data.soft_budget is not None and data.soft_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"}
+ )
+
if user_custom_key_generate is not None:
if asyncio.iscoroutinefunction(user_custom_key_generate):
result = await user_custom_key_generate(data) # type: ignore
@@ -1502,6 +1514,13 @@ async def update_key_fn(
)
try:
+ # Validate budget values are not negative
+ if data.max_budget is not None and data.max_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"}
+ )
+
data_json: dict = data.model_dump(exclude_unset=True, exclude_none=True)
key = data_json.pop("key")
@@ -3020,10 +3039,14 @@ async def list_keys(
description="Column to sort by (e.g. 'user_id', 'created_at', 'spend')",
),
sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"),
+ expand: Optional[List[str]] = Query(None, description="Expand related objects (e.g. 'user')"),
) -> KeyListResponseObject:
"""
List all keys for a given user / team / organization.
+ Parameters:
+ expand: Optional[List[str]] - Expand related objects (e.g. 'user' to include user information)
+
Returns:
{
"keys": List[str] or List[UserAPIKeyAuth],
@@ -3031,6 +3054,9 @@ async def list_keys(
"current_page": int,
"total_pages": int,
}
+
+ When expand includes "user", each key object will include a "user" field with the associated user object.
+ Note: When expand=user is specified, full key objects are returned regardless of the return_full_object parameter.
"""
try:
from litellm.proxy.proxy_server import prisma_client
@@ -3080,6 +3106,7 @@ async def list_keys(
include_created_by_keys=include_created_by_keys,
sort_by=sort_by,
sort_order=sort_order,
+ expand=expand,
)
verbose_proxy_logger.debug("Successfully prepared response")
@@ -3215,45 +3242,17 @@ def _validate_sort_params(
return order_by
-async def _list_key_helper(
- prisma_client: PrismaClient,
- page: int,
- size: int,
+def _build_key_filter_conditions(
user_id: Optional[str],
team_id: Optional[str],
organization_id: Optional[str],
key_alias: Optional[str],
key_hash: Optional[str],
- exclude_team_id: Optional[str] = None,
- return_full_object: bool = False,
- admin_team_ids: Optional[
- List[str]
- ] = None, # New parameter for teams where user is admin
- include_created_by_keys: bool = False,
- sort_by: Optional[str] = None,
- sort_order: str = "desc",
-) -> KeyListResponseObject:
- """
- Helper function to list keys
- Args:
- page: int
- size: int
- user_id: Optional[str]
- team_id: Optional[str]
- key_alias: Optional[str]
- exclude_team_id: Optional[str] # exclude a specific team_id
- return_full_object: bool # when true, will return UserAPIKeyAuth objects instead of just the token
- admin_team_ids: Optional[List[str]] # list of team IDs where the user is an admin
-
- Returns:
- KeyListResponseObject
- {
- "keys": List[str] or List[UserAPIKeyAuth], # Updated to reflect possible return types
- "total_count": int,
- "current_page": int,
- "total_pages": int,
- }
- """
+ exclude_team_id: Optional[str],
+ admin_team_ids: Optional[List[str]],
+ include_created_by_keys: bool,
+) -> Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]]:
+ """Build filter conditions for key listing."""
# Prepare filter conditions
where: Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]] = {}
where.update(_get_condition_to_filter_out_ui_session_tokens())
@@ -3294,6 +3293,59 @@ async def _list_key_helper(
where.update(or_conditions[0])
verbose_proxy_logger.debug(f"Filter conditions: {where}")
+ return where
+
+
+async def _list_key_helper(
+ prisma_client: PrismaClient,
+ page: int,
+ size: int,
+ user_id: Optional[str],
+ team_id: Optional[str],
+ organization_id: Optional[str],
+ key_alias: Optional[str],
+ key_hash: Optional[str],
+ exclude_team_id: Optional[str] = None,
+ return_full_object: bool = False,
+ admin_team_ids: Optional[
+ List[str]
+ ] = None, # New parameter for teams where user is admin
+ include_created_by_keys: bool = False,
+ sort_by: Optional[str] = None,
+ sort_order: str = "desc",
+ expand: Optional[List[str]] = None,
+) -> KeyListResponseObject:
+ """
+ Helper function to list keys
+ Args:
+ page: int
+ size: int
+ user_id: Optional[str]
+ team_id: Optional[str]
+ key_alias: Optional[str]
+ exclude_team_id: Optional[str] # exclude a specific team_id
+ return_full_object: bool # when true, will return UserAPIKeyAuth objects instead of just the token
+ admin_team_ids: Optional[List[str]] # list of team IDs where the user is an admin
+
+ Returns:
+ KeyListResponseObject
+ {
+ "keys": List[str] or List[UserAPIKeyAuth], # Updated to reflect possible return types
+ "total_count": int,
+ "current_page": int,
+ "total_pages": int,
+ }
+ """
+ where = _build_key_filter_conditions(
+ user_id=user_id,
+ team_id=team_id,
+ organization_id=organization_id,
+ key_alias=key_alias,
+ key_hash=key_hash,
+ exclude_team_id=exclude_team_id,
+ admin_team_ids=admin_team_ids,
+ include_created_by_keys=include_created_by_keys,
+ )
# Calculate skip for pagination
skip = (page - 1) * size
@@ -3334,13 +3386,28 @@ async def _list_key_helper(
# Calculate total pages
total_pages = -(-total_count // size) # Ceiling division
+ # Fetch user information if expand includes "user"
+ user_map = {}
+ if expand and "user" in expand:
+ user_ids = [key.user_id for key in keys if key.user_id]
+ if user_ids:
+ users = await prisma_client.db.litellm_usertable.find_many(
+ where={"user_id": {"in": list(set(user_ids))}} # Remove duplicates
+ )
+ user_map = {user.user_id: user for user in users}
+
# Prepare response
key_list: List[Union[str, UserAPIKeyAuth]] = []
for key in keys:
key_dict = key.dict()
# Attach object_permission if object_permission_id is set
key_dict = await attach_object_permission_to_dict(key_dict, prisma_client)
- if return_full_object is True:
+
+ # Include user information if expand includes "user"
+ if expand and "user" in expand and key.user_id and key.user_id in user_map:
+ key_dict["user"] = user_map[key.user_id].dict()
+
+ if return_full_object is True or (expand and "user" in expand):
key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object
else:
_token = key_dict.get("token")
diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
index f0eddcc8683..9111f53a517 100644
--- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
@@ -16,7 +16,7 @@ Endpoints here:
import importlib
from dataclasses import dataclass
from datetime import datetime, timedelta
-from typing import Any, Dict, Iterable, List, Optional
+from typing import Any, Dict, Iterable, List, Literal, Optional
from fastapi import (
APIRouter,
@@ -24,6 +24,7 @@ from fastapi import (
Form,
Header,
HTTPException,
+ Query,
Request,
Response,
status,
@@ -208,6 +209,9 @@ if MCP_AVAILABLE:
command=payload.command,
args=payload.args,
env=payload.env,
+ authorization_url=payload.authorization_url,
+ token_url=payload.token_url,
+ registration_url=payload.registration_url,
)
def get_prisma_client_or_throw(message: str):
@@ -296,117 +300,6 @@ if MCP_AVAILABLE:
access_groups_list = sorted(list(access_groups))
return {"access_groups": access_groups_list}
- @router.get(
- "/server/{server_id}/health",
- description="Perform health check on a specific MCP server",
- dependencies=[Depends(user_api_key_auth)],
- )
- async def health_check_mcp_server(
- server_id: str,
- user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
- ):
- """
- Perform a health check on the MCP server specified by the `server_id`
- Parameters:
- - server_id: str - Required. The unique identifier of the mcp server to health check.
- ```
- curl --location 'http://localhost:4000/v1/mcp/server/{server_id}/health' \
- --header 'Authorization: Bearer your_api_key_here'
- ```
- """
- # Check if server exists and user has access
- prisma_client = get_prisma_client_or_throw(
- "Database not connected. Connect a database to your proxy"
- )
-
- # check to see if server exists for all users
- mcp_server = await get_mcp_server(prisma_client, server_id)
- if mcp_server is None:
- raise HTTPException(
- status_code=status.HTTP_404_NOT_FOUND,
- detail={"error": f"MCP Server with id {server_id} not found"},
- )
-
- # Implement authz restriction from requested user
- if not _user_has_admin_view(user_api_key_dict):
- # Perform authz check to filter the mcp servers user has access to
- mcp_server_records = await get_all_mcp_servers_for_user(
- prisma_client, user_api_key_dict
- )
- exists = does_mcp_server_exist(mcp_server_records, server_id)
-
- if not exists:
- raise HTTPException(
- status_code=status.HTTP_403_FORBIDDEN,
- detail={
- "error": f"User does not have permission to access mcp server with id {server_id}. You can only access mcp servers that you have access to."
- },
- )
-
- # Perform health check using server manager
- try:
- health_result = await global_mcp_server_manager.health_check_server(
- server_id
- )
- return health_result
- except Exception as e:
- verbose_proxy_logger.exception(
- f"Error performing health check on MCP server {server_id}: {str(e)}"
- )
- raise HTTPException(
- status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
- detail={"error": f"Error performing health check: {str(e)}"},
- )
-
- @router.get(
- "/server/health",
- description="Perform health check on all accessible MCP servers",
- dependencies=[Depends(user_api_key_auth)],
- )
- async def health_check_all_mcp_servers(
- user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
- ):
- """
- Perform health checks on all MCP servers accessible to the user
- ```
- curl --location 'http://localhost:4000/v1/mcp/server/health' \
- --header 'Authorization: Bearer your_api_key_here'
- ```
- """
- # Use server manager to get health checks for allowed servers
- try:
- all_health_results = (
- await global_mcp_server_manager.health_check_allowed_servers(
- user_api_key_auth=user_api_key_dict
- )
- )
-
- return {
- "total_servers": len(all_health_results),
- "healthy_count": len(
- [r for r in all_health_results.values() if r["status"] == "healthy"]
- ),
- "unhealthy_count": len(
- [
- r
- for r in all_health_results.values()
- if r["status"] == "unhealthy"
- ]
- ),
- "unknown_count": len(
- [r for r in all_health_results.values() if r["status"] == "unknown"]
- ),
- "servers": all_health_results,
- }
- except Exception as e:
- verbose_proxy_logger.exception(
- f"Error performing health checks on MCP servers: {str(e)}"
- )
- raise HTTPException(
- status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
- detail={"error": f"Error performing health checks: {str(e)}"},
- )
-
## FastAPI Routes
@router.get(
"/server",
@@ -429,7 +322,7 @@ if MCP_AVAILABLE:
aggregated_servers: Dict[str, LiteLLM_MCPServerTable] = {}
for auth_context in auth_contexts:
- servers = await global_mcp_server_manager.get_all_mcp_servers_with_health_and_teams(
+ servers = await global_mcp_server_manager.get_all_allowed_mcp_servers(
user_api_key_auth=auth_context
)
for server in servers:
@@ -447,6 +340,56 @@ if MCP_AVAILABLE:
server.mcp_info["is_public"] = True
return redacted_mcp_servers
+ @router.get(
+ "/server/health",
+ description="Health check for MCP servers",
+ dependencies=[Depends(user_api_key_auth)],
+ )
+ async def health_check_servers(
+ server_ids: Optional[List[str]] = Query(
+ None,
+ description="Server IDs to check. If not provided, checks all accessible servers.",
+ ),
+ user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+ ):
+ """
+ Perform health checks on one or more MCP servers.
+
+ Parameters:
+ - server_ids: Optional list of server IDs. If not provided, checks all accessible servers.
+
+ Returns:
+ - Health check results for requested servers
+
+ ```
+ # Check all accessible servers
+ curl --location 'http://localhost:4000/v1/mcp/server/health' \
+ --header 'Authorization: Bearer your_api_key_here'
+
+ # Check specific servers
+ curl --location 'http://localhost:4000/v1/mcp/server/health?server_ids=server-1&server_ids=server-2' \
+ --header 'Authorization: Bearer your_api_key_here'
+ ```
+ """
+ auth_contexts = await build_effective_auth_contexts(user_api_key_dict)
+
+ server_status_map: Dict[
+ str, Optional[Literal["healthy", "unhealthy", "unknown"]]
+ ] = {}
+ for auth_context in auth_contexts:
+ servers = await global_mcp_server_manager.get_all_mcp_servers_with_health_and_teams(
+ user_api_key_auth=auth_context,
+ server_ids=server_ids,
+ )
+ for server in servers:
+ if server.server_id not in server_status_map:
+ server_status_map[server.server_id] = server.status
+
+ return [
+ {"server_id": server_id, "status": status}
+ for server_id, status in server_status_map.items()
+ ]
+
@router.get(
"/server/{server_id}",
description="Returns the mcp server info",
@@ -484,15 +427,11 @@ if MCP_AVAILABLE:
server_id
)
# Update the server object with health check results
- mcp_server.status = health_result.get("status", "unknown")
- mcp_server.last_health_check = (
- datetime.fromisoformat(
- health_result.get("last_health_check", datetime.now().isoformat())
- )
- if health_result.get("last_health_check")
- else None
+ mcp_server.status = (
+ health_result.status if health_result.status else "unknown"
)
- mcp_server.health_check_error = health_result.get("error")
+ mcp_server.last_health_check = health_result.last_health_check
+ mcp_server.health_check_error = health_result.health_check_error
except Exception as e:
verbose_proxy_logger.debug(
f"Error performing health check on server {server_id}: {e}"
@@ -512,7 +451,7 @@ if MCP_AVAILABLE:
exists = does_mcp_server_exist(mcp_server_records, server_id)
if exists:
- await global_mcp_server_manager.add_update_server(mcp_server)
+ await global_mcp_server_manager.add_server(mcp_server)
return _redact_mcp_credentials(mcp_server)
else:
raise HTTPException(
@@ -586,7 +525,7 @@ if MCP_AVAILABLE:
payload,
touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
)
- await global_mcp_server_manager.add_update_server(new_mcp_server)
+ await global_mcp_server_manager.add_server(new_mcp_server)
# Ensure registry is up to date by reloading from database
await global_mcp_server_manager.reload_servers_from_database()
@@ -867,7 +806,7 @@ if MCP_AVAILABLE:
"error": f"MCP Server not found, passed server_id={payload.server_id}"
},
)
- await global_mcp_server_manager.add_update_server(mcp_server_record_updated)
+ await global_mcp_server_manager.update_server(mcp_server_record_updated)
# Ensure registry is up to date by reloading from database
await global_mcp_server_manager.reload_servers_from_database()
diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py
index 99b37c765a8..a088f46d67d 100644
--- a/litellm/proxy/management_endpoints/organization_endpoints.py
+++ b/litellm/proxy/management_endpoints/organization_endpoints.py
@@ -24,8 +24,11 @@ from litellm.proxy.management_endpoints.budget_management_endpoints import (
new_budget,
update_budget,
)
-from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field
-from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
+from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
+from litellm.proxy.management_endpoints.common_utils import (
+ _set_object_metadata_field,
+ _user_has_admin_view,
+)
from litellm.proxy.management_helpers.object_permission_utils import (
handle_update_object_permission_common,
)
@@ -34,11 +37,10 @@ from litellm.proxy.management_helpers.utils import (
management_endpoint_wrapper,
)
from litellm.proxy.utils import PrismaClient
-from litellm.utils import _update_dictionary
from litellm.types.proxy.management_endpoints.common_daily_activity import (
SpendAnalyticsPaginatedResponse,
)
-from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
+from litellm.utils import _update_dictionary
router = APIRouter()
@@ -168,6 +170,18 @@ async def new_organization(
status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}
)
+ # Validate budget values are not negative
+ if data.max_budget is not None and data.max_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"}
+ )
+ if data.soft_budget is not None and data.soft_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"}
+ )
+
user_object_correct_type: Optional[LiteLLM_UserTable] = None
if user_api_key_dict.user_id is not None:
@@ -414,6 +428,18 @@ async def update_organization(
# Create validated data model
data = LiteLLM_OrganizationTableUpdate(**raw_data_with_flat_budget_fields)
+ # Validate budget values are not negative
+ if data.max_budget is not None and data.max_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"}
+ )
+ if data.soft_budget is not None and data.soft_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"}
+ )
+
if data.updated_by is None:
data.updated_by = user_api_key_dict.user_id
diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index c6fab9a73f0..76c607f5c49 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -732,6 +732,18 @@ async def new_team( # noqa: PLR0915
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
+ # Validate budget values are not negative
+ if data.max_budget is not None and data.max_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"}
+ )
+ if data.team_member_budget is not None and data.team_member_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"}
+ )
+
# Check if license is over limit
total_teams = await prisma_client.db.litellm_teamtable.count()
if total_teams and _license_check.is_team_count_over_limit(
@@ -1169,7 +1181,7 @@ def validate_team_org_change(
"/team/update", tags=["team management"], dependencies=[Depends(user_api_key_auth)]
)
@management_endpoint_wrapper
-async def update_team(
+async def update_team( # noqa: PLR0915
data: UpdateTeamRequest,
http_request: Request,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@@ -1254,6 +1266,18 @@ async def update_team(
raise HTTPException(status_code=400, detail={"error": "No team id passed in"})
verbose_proxy_logger.debug("/team/update - %s", data)
+ # Validate budget values are not negative
+ if data.max_budget is not None and data.max_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"}
+ )
+ if data.team_member_budget is not None and data.team_member_budget < 0:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"}
+ )
+
existing_team_row = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": data.team_id}
)
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index f56c0c2b07a..06525e39133 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -4626,7 +4626,7 @@ class ProxyStartupEvent:
verbose_proxy_logger.info("Batch cost check job scheduled successfully")
except Exception as e:
- verbose_proxy_logger.error(f"Failed to setup batch cost checking: {e}")
+ verbose_proxy_logger.debug(f"Failed to setup batch cost checking: {e}")
verbose_proxy_logger.debug(
"Checking batch cost for LiteLLM Managed Files is an Enterprise Feature. Skipping..."
)
@@ -4657,7 +4657,7 @@ class ProxyStartupEvent:
verbose_proxy_logger.info("Responses cost check job scheduled successfully")
except Exception as e:
- verbose_proxy_logger.error(f"Failed to setup responses cost checking: {e}")
+ verbose_proxy_logger.debug(f"Failed to setup responses cost checking: {e}")
verbose_proxy_logger.debug(
"Checking responses cost for LiteLLM Managed Files is an Enterprise Feature. Skipping..."
)
@@ -7322,13 +7322,9 @@ async def model_info_v2(
"""
global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router
- if llm_router is None:
- raise HTTPException(
- status_code=500,
- detail={
- "error": f"No model list passed, models router={llm_router}. You can add a model through the config.yaml or on the LiteLLM Admin UI."
- },
- )
+ # Return empty data array when no models are configured (graceful handling for fresh installs)
+ if llm_router is None or not llm_router.model_list:
+ return {"data": []}
if prisma_client is None:
raise HTTPException(
@@ -8228,14 +8224,9 @@ async def model_group_info(
"""
global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router
- if llm_model_list is None:
- raise HTTPException(
- status_code=500, detail={"error": "LLM Model List not loaded in"}
- )
- if llm_router is None:
- raise HTTPException(
- status_code=500, detail={"error": "LLM Router is not loaded in"}
- )
+ # Return empty data array when no models are configured (graceful handling for fresh installs)
+ if llm_model_list is None or llm_router is None or not llm_model_list:
+ return {"data": []}
from litellm.proxy.utils import get_available_models_for_user
diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json
index 9916bdf6923..ec1c4619527 100644
--- a/litellm/proxy/public_endpoints/provider_create_fields.json
+++ b/litellm/proxy/public_endpoints/provider_create_fields.json
@@ -1680,6 +1680,34 @@
],
"default_model_placeholder": "gpt-3.5-turbo"
},
+ {
+ "provider": "MINIMAX",
+ "provider_display_name": "MiniMax",
+ "litellm_provider": "minimax",
+ "credential_fields": [
+ {
+ "key": "api_key",
+ "label": "API Key",
+ "placeholder": "your-minimax-api-key",
+ "tooltip": "MiniMax API Key from https://platform.minimaxi.com/",
+ "required": true,
+ "field_type": "password",
+ "options": null,
+ "default_value": null
+ },
+ {
+ "key": "api_base",
+ "label": "API Base URL",
+ "placeholder": "https://api.minimax.io/v1",
+ "tooltip": "International: https://api.minimax.io/v1, China: https://api.minimaxi.com/v1",
+ "required": false,
+ "field_type": "text",
+ "options": null,
+ "default_value": "https://api.minimax.io/v1"
+ }
+ ],
+ "default_model_placeholder": "minimax/MiniMax-M2"
+ },
{
"provider": "MOONSHOT",
"provider_display_name": "Moonshot",
@@ -2865,7 +2893,7 @@
"key": "api_base",
"label": "API Base",
"placeholder": null,
- "tooltip": null,
+ "tooltip": "Base URL of your WatsonX instance",
"required": false,
"field_type": "text",
"options": null,
@@ -2875,14 +2903,54 @@
"key": "api_key",
"label": "API Key",
"placeholder": null,
- "tooltip": null,
+ "tooltip": "IBM Cloud API key. Required if not using Token or Zen API Key",
"required": false,
"field_type": "password",
"options": null,
"default_value": null
+ },
+ {
+ "key": "token",
+ "label": "IAM Token",
+ "placeholder": null,
+ "tooltip": "Pre-generated IAM bearer token. Use instead of API Key if you manage tokens externally",
+ "required": false,
+ "field_type": "password",
+ "options": null,
+ "default_value": null
+ },
+ {
+ "key": "zen_api_key",
+ "label": "Zen API Key",
+ "placeholder": null,
+ "tooltip": "Zen API Key for Cloud Pak for Data deployments. Use instead of API Key for on-premises",
+ "required": false,
+ "field_type": "password",
+ "options": null,
+ "default_value": null
+ },
+ {
+ "key": "project_id",
+ "label": "Project ID",
+ "placeholder": null,
+ "tooltip": "Optional: Your Watsonx.ai Project ID",
+ "required": false,
+ "field_type": "text",
+ "options": null,
+ "default_value": null
+ },
+ {
+ "key": "space_id",
+ "label": "Deployment Space ID",
+ "placeholder": null,
+ "tooltip": "Optional: Watsonx.ai Deployment Space ID",
+ "required": false,
+ "field_type": "text",
+ "options": null,
+ "default_value": null
}
],
- "default_model_placeholder": "gpt-3.5-turbo"
+ "default_model_placeholder": "watsonx/ibm/granite-3-3-8b-instruct"
},
{
"provider": "WATSONX_TEXT",
diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma
index aac0b5b35de..ea47b6ed03b 100644
--- a/litellm/proxy/schema.prisma
+++ b/litellm/proxy/schema.prisma
@@ -208,6 +208,9 @@ model LiteLLM_MCPServerTable {
command String?
args String[] @default([])
env Json? @default("{}")
+ authorization_url String?
+ token_url String?
+ registration_url String?
}
// Generate Tokens for Proxy
diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py
index 1861c44b699..1c457d7bf4c 100644
--- a/litellm/proxy/spend_tracking/spend_tracking_utils.py
+++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py
@@ -11,7 +11,10 @@ from pydantic import BaseModel
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB, REDACTED_BY_LITELM_STRING
-from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs
+from litellm.litellm_core_utils.core_helpers import (
+ get_litellm_metadata_from_kwargs,
+ reconstruct_model_name,
+)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload
from litellm.proxy.utils import PrismaClient, hash_token
@@ -100,9 +103,9 @@ def _get_spend_logs_metadata(
clean_metadata["applied_guardrails"] = applied_guardrails
clean_metadata["batch_models"] = batch_models
clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata
- clean_metadata["vector_store_request_metadata"] = (
- _get_vector_store_request_for_spend_logs_payload(vector_store_request_metadata)
- )
+ clean_metadata[
+ "vector_store_request_metadata"
+ ] = _get_vector_store_request_for_spend_logs_payload(vector_store_request_metadata)
clean_metadata["guardrail_information"] = guardrail_information
clean_metadata["usage_object"] = usage_object
clean_metadata["model_map_information"] = model_map_information
@@ -393,6 +396,9 @@ def get_logging_payload( # noqa: PLR0915
# Extract agent_id for A2A requests (set directly on model_call_details)
agent_id: Optional[str] = kwargs.get("agent_id")
+ custom_llm_provider = kwargs.get("custom_llm_provider")
+ raw_model = cast(str, kwargs.get("model") or "")
+ model_name = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {})
try:
payload: SpendLogsPayload = SpendLogsPayload(
@@ -403,7 +409,7 @@ def get_logging_payload( # noqa: PLR0915
startTime=_ensure_datetime_utc(start_time),
endTime=_ensure_datetime_utc(end_time),
completionStartTime=_ensure_datetime_utc(completion_start_time),
- model=kwargs.get("model", "") or "",
+ model=model_name,
user=metadata.get("user_api_key_user_id", "") or "",
team_id=metadata.get("user_api_key_team_id", "") or "",
organization_id=metadata.get("user_api_key_org_id") or "",
@@ -449,7 +455,7 @@ def get_logging_payload( # noqa: PLR0915
# Explicitly clear large intermediate objects to reduce memory pressure
del response_obj_dict, usage, clean_metadata, additional_usage_values
-
+
return payload
except Exception as e:
verbose_proxy_logger.exception(
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index d595db4a2e0..d1a78534dae 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -151,25 +151,25 @@ def _get_email_logger_class():
"""
Determine which email logger class to use based on environment variables.
Priority: SendGrid > Resend > SMTP > BaseEmailLogger (fallback)
-
+
Returns:
The email logger class to use, or None if BaseEmailLogger is not available
"""
if BaseEmailLogger is None:
return None
-
+
# Check for SendGrid API key
if SendGridEmailLogger is not None and os.getenv("SENDGRID_API_KEY"):
return SendGridEmailLogger
-
+
# Check for Resend API key
if ResendEmailLogger is not None and os.getenv("RESEND_API_KEY"):
return ResendEmailLogger
-
+
# Check for SMTP configuration
if SMTPEmailLogger is not None and os.getenv("SMTP_HOST"):
return SMTPEmailLogger
-
+
# Fallback to BaseEmailLogger (though it won't actually send emails)
return BaseEmailLogger
@@ -452,7 +452,6 @@ class ProxyLogging:
litellm.logging_callback_manager.add_litellm_callback(self.service_logging_obj) # type: ignore
for callback in litellm.callbacks:
if isinstance(callback, str):
-
callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( # type: ignore
cast(_custom_logger_compatible_callbacks_literal, callback),
internal_usage_cache=self.internal_usage_cache.dual_cache,
@@ -965,7 +964,7 @@ class ProxyLogging:
# Determine the event type based on call type
event_type = GuardrailEventHooks.pre_call
- if call_type == "mcp_call":
+ if call_type == CallTypes.call_mcp_tool.value:
event_type = GuardrailEventHooks.pre_mcp_call
# Check if the guardrail should run for this request
@@ -1038,7 +1037,6 @@ class ProxyLogging:
data.pop("prompt_id", None)
if custom_logger and prompt_spec is not None:
-
(
model,
messages,
@@ -1261,7 +1259,7 @@ class ProxyLogging:
from litellm.types.guardrails import GuardrailEventHooks
event_type = GuardrailEventHooks.during_call
- if call_type == "mcp_call":
+ if call_type == CallTypes.call_mcp_tool.value:
event_type = GuardrailEventHooks.during_mcp_call
if (
@@ -1270,7 +1268,7 @@ class ProxyLogging:
):
continue
# Convert user_api_key_dict to proper format for async_moderation_hook
- if call_type == "mcp_call":
+ if call_type == CallTypes.call_mcp_tool.value:
user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(
user_api_key_dict
)
@@ -1288,7 +1286,6 @@ class ProxyLogging:
call_type=call_type,
)
else:
-
guardrail_task = callback.async_moderation_hook(
data=data,
user_api_key_dict=user_api_key_auth_dict, # type: ignore
@@ -1337,7 +1334,7 @@ class ProxyLogging:
if self.alerting is None:
# do nothing if alerting is not switched on
return
-
+
if "slack" in self.alerting:
await self.slack_alerting_instance.budget_alerts(
type=type,
@@ -1548,7 +1545,10 @@ class ProxyLogging:
traceback_str=traceback_str,
)
# If callback returned an HTTPException, use it (first one wins)
- if isinstance(hook_result, HTTPException) and transformed_exception is None:
+ if (
+ isinstance(hook_result, HTTPException)
+ and transformed_exception is None
+ ):
transformed_exception = hook_result
except HTTPException as e:
# If callback raised an HTTPException, use it (first one wins)
@@ -1849,7 +1849,6 @@ class ProxyLogging:
current_response = response
for callback in litellm.callbacks:
-
_callback: Optional[CustomLogger] = None
if isinstance(callback, str):
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
@@ -3568,11 +3567,13 @@ class ProxyUpdateSpend:
)
# Atomically read and remove logs to process (protected by lock)
async with prisma_client._spend_log_transactions_lock:
- logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL]
+ logs_to_process = prisma_client.spend_log_transactions[
+ :MAX_LOGS_PER_INTERVAL
+ ]
# Remove the logs we're about to process
- prisma_client.spend_log_transactions = (
- prisma_client.spend_log_transactions[len(logs_to_process):]
- )
+ prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[
+ len(logs_to_process) :
+ ]
start_time = time.time()
try:
for i in range(n_retry_times + 1):
@@ -3675,9 +3676,7 @@ async def update_spend( # noqa: PLR0915
# Check queue size with lock protection
async with prisma_client._spend_log_transactions_lock:
queue_size = len(prisma_client.spend_log_transactions)
- verbose_proxy_logger.debug(
- "Spend Logs transactions: {}".format(queue_size)
- )
+ verbose_proxy_logger.debug("Spend Logs transactions: {}".format(queue_size))
# Process spend log transactions when called directly.
# This keeps backwards compatibility with the old behavior.
@@ -3699,19 +3698,19 @@ async def update_spend_logs_job(
):
"""
Job to process spend_log_transactions queue.
-
+
This job is triggered based on queue size rather than time.
Processes spend log transactions when the queue reaches a threshold.
"""
n_retry_times = 3
-
+
# Check queue size with lock protection
async with prisma_client._spend_log_transactions_lock:
queue_size = len(prisma_client.spend_log_transactions)
-
+
if queue_size == 0:
return
-
+
await ProxyUpdateSpend.update_spend_logs(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
@@ -3728,7 +3727,7 @@ async def _monitor_spend_logs_queue(
"""
Background task that monitors the spend_log_transactions queue size
and triggers processing when the threshold is reached.
-
+
Args:
prisma_client: Prisma client instance
db_writer_client: Optional HTTP handler for external spend logs endpoint
@@ -3738,23 +3737,23 @@ async def _monitor_spend_logs_queue(
SPEND_LOG_QUEUE_POLL_INTERVAL,
SPEND_LOG_QUEUE_SIZE_THRESHOLD,
)
-
+
threshold = SPEND_LOG_QUEUE_SIZE_THRESHOLD
base_interval = SPEND_LOG_QUEUE_POLL_INTERVAL
max_backoff = 30.0 # Maximum backoff interval in seconds
backoff_multiplier = 1.5 # Exponential backoff multiplier
current_interval = base_interval
-
+
verbose_proxy_logger.info(
f"Starting spend logs queue monitor (threshold: {threshold}, poll_interval: {base_interval}s)"
)
-
+
while True:
try:
# Check queue size with lock protection
async with prisma_client._spend_log_transactions_lock:
queue_size = len(prisma_client.spend_log_transactions)
-
+
if queue_size > 0:
if queue_size >= threshold:
verbose_proxy_logger.debug(
@@ -3767,8 +3766,10 @@ async def _monitor_spend_logs_queue(
f"Spend logs queue size ({queue_size}) below threshold ({threshold}), processing with backoff"
)
# Exponential backoff when below threshold but still processing
- current_interval = min(current_interval * backoff_multiplier, max_backoff)
-
+ current_interval = min(
+ current_interval * backoff_multiplier, max_backoff
+ )
+
await update_spend_logs_job(
prisma_client=prisma_client,
db_writer_client=db_writer_client,
@@ -3776,8 +3777,10 @@ async def _monitor_spend_logs_queue(
)
else:
# Exponential backoff when no logs to process
- current_interval = min(current_interval * backoff_multiplier, max_backoff)
-
+ current_interval = min(
+ current_interval * backoff_multiplier, max_backoff
+ )
+
await asyncio.sleep(current_interval)
except Exception as e:
verbose_proxy_logger.error(
@@ -3788,7 +3791,6 @@ async def _monitor_spend_logs_queue(
await asyncio.sleep(current_interval)
-
def _raise_failed_update_spend_exception(
e: Exception, start_time: float, proxy_logging_obj: ProxyLogging
):
diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py
index fe686598141..e26a2477b1b 100644
--- a/litellm/realtime_api/main.py
+++ b/litellm/realtime_api/main.py
@@ -3,7 +3,7 @@
from typing import Any, Optional, cast
import litellm
-from litellm import get_llm_provider
+from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py
index 2eea28f6cc1..9cdcd3894e0 100644
--- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py
+++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py
@@ -205,7 +205,15 @@ class LiteLLM_Proxy_MCP_Handler:
else:
tool_name = getattr(mcp_tool, "name", None)
- if tool_name and tool_name in allowed_tool_names:
+ if not tool_name:
+ continue
+
+ if tool_name in allowed_tool_names:
+ filtered_tools.append(mcp_tool)
+ continue
+
+ unprefixed_name, _ = split_server_prefix_from_name(tool_name)
+ if unprefixed_name in allowed_tool_names:
filtered_tools.append(mcp_tool)
return filtered_tools
diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py
index ad99609e905..a92b5d25a37 100644
--- a/litellm/responses/utils.py
+++ b/litellm/responses/utils.py
@@ -26,7 +26,7 @@ from litellm.types.llms.openai import (
from litellm.types.responses.main import DecodedResponseId
from litellm.types.utils import (
CompletionTokensDetailsWrapper,
- PromptTokensDetails,
+ PromptTokensDetailsWrapper,
SpecialEnums,
Usage,
)
@@ -431,7 +431,12 @@ class ResponseAPILoggingUtils:
def _transform_response_api_usage_to_chat_usage(
usage_input: Optional[Union[dict, ResponseAPIUsage]],
) -> Usage:
- """Tranforms the ResponseAPIUsage object to a Usage object"""
+ """
+ Transforms ResponseAPIUsage or ImageUsage to a Usage object.
+
+ Both have the same spec with input_tokens, output_tokens, and
+ input_tokens_details (text_tokens, image_tokens).
+ """
if usage_input is None:
return Usage(
prompt_tokens=0,
@@ -445,18 +450,19 @@ class ResponseAPILoggingUtils:
)
prompt_tokens: int = response_api_usage.input_tokens or 0
completion_tokens: int = response_api_usage.output_tokens or 0
- prompt_tokens_details: Optional[PromptTokensDetails] = None
+ prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
if response_api_usage.input_tokens_details:
- prompt_tokens_details = PromptTokensDetails(
- cached_tokens=response_api_usage.input_tokens_details.cached_tokens,
- audio_tokens=response_api_usage.input_tokens_details.audio_tokens,
+ prompt_tokens_details = PromptTokensDetailsWrapper(
+ cached_tokens=getattr(response_api_usage.input_tokens_details, "cached_tokens", None),
+ audio_tokens=getattr(response_api_usage.input_tokens_details, "audio_tokens", None),
+ text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None),
+ image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None),
)
completion_tokens_details: Optional[CompletionTokensDetailsWrapper] = None
- if response_api_usage.output_tokens_details:
+ output_tokens_details = getattr(response_api_usage, "output_tokens_details", None)
+ if output_tokens_details:
completion_tokens_details = CompletionTokensDetailsWrapper(
- reasoning_tokens=getattr(
- response_api_usage.output_tokens_details, "reasoning_tokens", None
- )
+ reasoning_tokens=getattr(output_tokens_details, "reasoning_tokens", None)
)
chat_usage = Usage(
diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py
index c6804b1ad4c..69d6ab9b6e2 100644
--- a/litellm/router_utils/pattern_match_deployments.py
+++ b/litellm/router_utils/pattern_match_deployments.py
@@ -7,7 +7,7 @@ import re
from re import Match
from typing import Dict, List, Optional, Tuple
-from litellm import get_llm_provider
+from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm._logging import verbose_router_logger
diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py
index ceeae958a80..c2912558cab 100644
--- a/litellm/types/llms/openai.py
+++ b/litellm/types/llms/openai.py
@@ -1197,6 +1197,39 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject):
# Define private attributes using PrivateAttr
_hidden_params: dict = PrivateAttr(default_factory=dict)
+ @property
+ def output_text(self) -> str:
+ """
+ Convenience property that aggregates all `output_text` items from the `output` list.
+
+ If no `output_text` content blocks exist, then an empty string is returned.
+
+ This matches the OpenAI SDK's Response.output_text behavior.
+ """
+ texts: List[str] = []
+ for output_item in self.output:
+ # Handle both dict and object access patterns
+ if isinstance(output_item, dict):
+ item_type = output_item.get("type")
+ content = output_item.get("content", [])
+ else:
+ item_type = getattr(output_item, "type", None)
+ content = getattr(output_item, "content", [])
+
+ if item_type == "message":
+ for content_item in content:
+ if isinstance(content_item, dict):
+ content_type = content_item.get("type")
+ text = content_item.get("text", "")
+ else:
+ content_type = getattr(content_item, "type", None)
+ text = getattr(content_item, "text", "") or ""
+
+ if content_type == "output_text":
+ texts.append(text)
+
+ return "".join(texts)
+
class ResponsesAPIStreamEvents(str, Enum):
"""
diff --git a/litellm/types/utils.py b/litellm/types/utils.py
index 144e503acdf..3eec67d9d26 100644
--- a/litellm/types/utils.py
+++ b/litellm/types/utils.py
@@ -142,6 +142,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
] # only for vertex ai models
input_cost_per_query: Optional[float] # only for rerank models
input_cost_per_image: Optional[float] # only for vertex ai models
+ input_cost_per_image_token: Optional[float] # for gpt-image-1 and similar models
input_cost_per_audio_per_second: Optional[float] # only for vertex ai models
input_cost_per_video_per_second: Optional[float] # only for vertex ai models
input_cost_per_second: Optional[float] # for OpenAI Speech models
@@ -1300,7 +1301,7 @@ class CacheCreationTokenDetails(BaseModel):
class PromptTokensDetailsWrapper(
PromptTokensDetails
-): # wrapper for older openai versions
+): # extends with image generation fields (text_tokens, image_tokens)
text_tokens: Optional[int] = None
"""Text tokens sent to the model."""
diff --git a/litellm/utils.py b/litellm/utils.py
index 102df5d595e..df0b2317123 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -34,7 +34,6 @@ from inspect import iscoroutine
from io import StringIO
from os.path import abspath, dirname, join
-import aiohttp
import dotenv
import httpx
import openai
@@ -48,21 +47,16 @@ from tiktoken import Encoding
from tokenizers import Tokenizer
import litellm
-import litellm._service_logger # for storing API inputs, outputs, and metadata
+
import litellm.litellm_core_utils
-import litellm.litellm_core_utils.audio_utils.utils
+# audio_utils.utils is lazy-loaded - only imported when needed for transcription calls
import litellm.litellm_core_utils.json_validation_rule
-import litellm.llms
-import litellm.llms.gemini
from litellm._lazy_imports import (
_get_default_encoding,
_get_modified_max_tokens,
_get_token_counter_new,
)
from litellm._uuid import uuid
-from litellm.caching._internal_lru_cache import lru_cache_wrapper
-from litellm.caching.caching import DualCache
-from litellm.caching.caching_handler import CachingHandlerResponse, LLMCachingHandler
from litellm.constants import (
DEFAULT_CHAT_COMPLETION_PARAM_VALUES,
DEFAULT_EMBEDDING_PARAM_VALUES,
@@ -77,90 +71,84 @@ from litellm.constants import (
OPENAI_EMBEDDING_PARAMS,
TOOL_CHOICE_OBJECT_TOKEN_COUNT,
)
-from litellm.integrations.custom_guardrail import CustomGuardrail
-from litellm.integrations.custom_logger import CustomLogger
-from litellm.integrations.vector_store_integrations.base_vector_store import (
- BaseVectorStore,
-)
-# Import cached imports utilities
-from litellm.litellm_core_utils.cached_imports import (
- get_coroutine_checker,
- get_litellm_logging_class,
- get_set_callbacks,
-)
-from litellm.litellm_core_utils.core_helpers import (
- get_litellm_metadata_from_kwargs,
- map_finish_reason,
- process_response_headers,
-)
-from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
-from litellm.litellm_core_utils.dot_notation_indexing import (
- delete_nested_value,
- is_nested_path,
-)
-from litellm.litellm_core_utils.exception_mapping_utils import (
- _get_response_headers,
- exception_type,
- get_error_message,
-)
-from litellm.litellm_core_utils.get_litellm_params import (
- _get_base_model_from_litellm_call_metadata,
- get_litellm_params,
-)
-from litellm.litellm_core_utils.get_llm_provider_logic import (
- _is_non_openai_azure_model,
- get_llm_provider,
-)
-from litellm.litellm_core_utils.get_supported_openai_params import (
- get_supported_openai_params,
-)
-from litellm.litellm_core_utils.llm_request_utils import _ensure_extra_body_is_safe
-from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
- LiteLLMResponseObjectHandler,
- _handle_invalid_parallel_tool_calls,
- convert_to_model_response_object,
- convert_to_streaming_response,
- convert_to_streaming_response_async,
-)
-from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
-from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import (
- get_formatted_prompt,
-)
-from litellm.litellm_core_utils.llm_response_utils.get_headers import (
- get_response_headers,
-)
-from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
- ResponseMetadata,
-)
-from litellm.litellm_core_utils.prompt_templates.common_utils import (
- _parse_content_for_reasoning,
-)
-from litellm.litellm_core_utils.redact_messages import (
- LiteLLMLoggingObject,
- redact_message_input_output_from_logging,
-)
-from litellm.litellm_core_utils.rules import Rules
-from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
-from litellm.llms.base_llm.google_genai.transformation import (
- BaseGoogleGenAIGenerateContentConfig,
-)
-from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig
-from litellm.llms.base_llm.search.transformation import BaseSearchConfig
-from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig
-from litellm.llms.bedrock.common_utils import BedrockModelInfo
-from litellm.llms.cohere.common_utils import CohereModelInfo
-from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
-from litellm.llms.mistral.ocr.transformation import MistralOCRConfig
-from litellm.router_utils.get_retry_from_policy import (
- get_num_retries_from_retry_policy,
- reset_retry_policy,
-)
-from litellm.secret_managers.main import get_secret
-from litellm.types.llms.anthropic import (
- ANTHROPIC_API_ONLY_HEADERS,
- AnthropicThinkingParam,
-)
+
+
+_CachingHandlerResponse = None
+_LLMCachingHandler = None
+_CustomGuardrail = None
+_CustomLogger = None
+
+
+def _get_cached_custom_logger():
+ """
+ Get cached CustomLogger class.
+ Lazy imports on first call to avoid loading custom_logger at import time.
+ Subsequent calls use cached class for better performance.
+ """
+ global _CustomLogger
+ if _CustomLogger is None:
+ from litellm.integrations.custom_logger import CustomLogger
+ _CustomLogger = CustomLogger
+ return _CustomLogger
+
+
+def _get_cached_custom_guardrail():
+ """
+ Get cached CustomGuardrail class.
+ Lazy imports on first call to avoid loading custom_guardrail at import time.
+ Subsequent calls use cached class for better performance.
+ """
+ global _CustomGuardrail
+ if _CustomGuardrail is None:
+ from litellm.integrations.custom_guardrail import CustomGuardrail
+ _CustomGuardrail = CustomGuardrail
+ return _CustomGuardrail
+
+
+def _get_cached_caching_handler_response():
+ """
+ Get cached CachingHandlerResponse class.
+ Lazy imports on first call to avoid loading caching_handler at import time.
+ Subsequent calls use cached class for better performance.
+ """
+ global _CachingHandlerResponse
+ if _CachingHandlerResponse is None:
+ from litellm.caching.caching_handler import CachingHandlerResponse
+ _CachingHandlerResponse = CachingHandlerResponse
+ return _CachingHandlerResponse
+
+
+def _get_cached_llm_caching_handler():
+ """
+ Get cached LLMCachingHandler class.
+ Lazy imports on first call to avoid loading caching_handler at import time.
+ Subsequent calls use cached class for better performance.
+ """
+ global _LLMCachingHandler
+ if _LLMCachingHandler is None:
+ from litellm.caching.caching_handler import LLMCachingHandler
+ _LLMCachingHandler = LLMCachingHandler
+ return _LLMCachingHandler
+
+
+# Cached lazy import for audio_utils.utils
+# Module-level cache to avoid repeated imports while preserving memory benefits
+_audio_utils_module = None
+
+
+def _get_cached_audio_utils():
+ """
+ Get cached audio_utils.utils module.
+ Lazy imports on first call to avoid loading audio_utils.utils at import time.
+ Subsequent calls use cached module for better performance.
+ """
+ global _audio_utils_module
+ if _audio_utils_module is None:
+ import litellm.litellm_core_utils.audio_utils.utils
+ _audio_utils_module = litellm.litellm_core_utils.audio_utils.utils
+ return _audio_utils_module
+
from litellm.types.llms.openai import (
AllMessageValues,
AllPromptValues,
@@ -171,7 +159,6 @@ from litellm.types.llms.openai import (
OpenAITextCompletionUserMessage,
OpenAIWebSearchOptions,
)
-from litellm.types.rerank import RerankResponse
from litellm.types.utils import FileTypes # type: ignore
from litellm.types.utils import (
OPENAI_RESPONSE_HEADERS,
@@ -256,16 +243,7 @@ from typing import (
from openai import OpenAIError as OriginalError
-from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
- update_response_metadata,
-)
-from litellm.litellm_core_utils.thread_pool_executor import executor
-from litellm.llms.base_llm.anthropic_messages.transformation import (
- BaseAnthropicMessagesConfig,
-)
-from litellm.llms.base_llm.audio_transcription.transformation import (
- BaseAudioTranscriptionConfig,
-)
+# These are lazy loaded via __getattr__
from litellm.llms.base_llm.base_utils import (
BaseLLMModelInfo,
type_to_response_format_param,
@@ -274,31 +252,126 @@ from litellm.llms.base_llm.base_utils import (
if TYPE_CHECKING:
# Heavy types that are only needed for type checking; avoid importing
# their modules at runtime during `litellm` import.
+ from litellm.caching.caching_handler import CachingHandlerResponse, LLMCachingHandler
+ from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.base_llm.files.transformation import BaseFilesConfig
from litellm.proxy._types import AllowedModelRegion
+ # Type stubs for lazy-loaded functions to help mypy understand their types
+ # These imports allow mypy to understand the types when these are accessed via __getattr__
+ from litellm.litellm_core_utils.exception_mapping_utils import exception_type
+ from litellm.litellm_core_utils.get_llm_provider_logic import (
+ _is_non_openai_azure_model,
+ get_llm_provider,
+ )
+ from litellm.litellm_core_utils.get_supported_openai_params import (
+ get_supported_openai_params,
+ )
+ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
+ LiteLLMResponseObjectHandler,
+ _handle_invalid_parallel_tool_calls,
+ convert_to_model_response_object,
+ convert_to_streaming_response,
+ convert_to_streaming_response_async,
+ )
+ from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
+ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
+ ResponseMetadata,
+ )
+ from litellm.litellm_core_utils.prompt_templates.common_utils import (
+ _parse_content_for_reasoning,
+ )
+ from litellm.litellm_core_utils.redact_messages import (
+ LiteLLMLoggingObject,
+ redact_message_input_output_from_logging,
+ )
+ from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
+ from litellm.llms.base_llm.google_genai.transformation import (
+ BaseGoogleGenAIGenerateContentConfig,
+ )
+ from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig
+ from litellm.llms.base_llm.search.transformation import BaseSearchConfig
+ from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig
+ from litellm.llms.bedrock.common_utils import BedrockModelInfo
+ from litellm.llms.cohere.common_utils import CohereModelInfo
+ from litellm.llms.mistral.ocr.transformation import MistralOCRConfig
+ # Type stubs for lazy-loaded functions and classes
+ from litellm.litellm_core_utils.cached_imports import (
+ get_coroutine_checker,
+ get_litellm_logging_class,
+ get_set_callbacks,
+ )
+ from litellm.litellm_core_utils.core_helpers import (
+ get_litellm_metadata_from_kwargs,
+ map_finish_reason,
+ process_response_headers,
+ )
+ from litellm.litellm_core_utils.dot_notation_indexing import (
+ delete_nested_value,
+ is_nested_path,
+ )
+ from litellm.litellm_core_utils.get_litellm_params import (
+ _get_base_model_from_litellm_call_metadata,
+ get_litellm_params,
+ )
+ from litellm.litellm_core_utils.llm_request_utils import _ensure_extra_body_is_safe
+ from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import (
+ get_formatted_prompt,
+ )
+ from litellm.litellm_core_utils.llm_response_utils.get_headers import (
+ get_response_headers,
+ )
+ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
+ update_response_metadata,
+ )
+ from litellm.litellm_core_utils.rules import Rules
+ from litellm.litellm_core_utils.thread_pool_executor import executor
+ from litellm.llms.base_llm.anthropic_messages.transformation import (
+ BaseAnthropicMessagesConfig,
+ )
+ from litellm.llms.base_llm.audio_transcription.transformation import (
+ BaseAudioTranscriptionConfig,
+ )
+ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
+ from litellm.router_utils.get_retry_from_policy import (
+ get_num_retries_from_retry_policy,
+ reset_retry_policy,
+ )
+ from litellm.secret_managers.main import get_secret
+ # Type stubs for lazy-loaded config classes and types
+ from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
+ from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
+ from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
+ from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
+ from litellm.llms.base_llm.image_generation.transformation import (
+ BaseImageGenerationConfig,
+ )
+ from litellm.llms.base_llm.image_variations.transformation import (
+ BaseImageVariationConfig,
+ )
+ from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
+ from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
+ from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
+ from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
+ from litellm.llms.base_llm.vector_store_files.transformation import (
+ BaseVectorStoreFilesConfig,
+ )
+ from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
+ from litellm.types.llms.anthropic import (
+ ANTHROPIC_API_ONLY_HEADERS,
+ AnthropicThinkingParam,
+ )
+ from litellm.types.rerank import RerankResponse
+ from litellm.types.llms.openai import (
+ ChatCompletionDeltaToolCallChunk,
+ ChatCompletionToolCallChunk,
+ ChatCompletionToolCallFunctionChunk,
+ )
+ from litellm.types.router import LiteLLM_Params
-from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
from litellm.llms.base_llm.chat.transformation import BaseConfig
from litellm.llms.base_llm.completion.transformation import BaseTextCompletionConfig
-from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
-from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
-from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
-from litellm.llms.base_llm.image_generation.transformation import (
- BaseImageGenerationConfig,
-)
-from litellm.llms.base_llm.image_variations.transformation import (
- BaseImageVariationConfig,
-)
-from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
-from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
-from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.llms.base_llm.skills.transformation import BaseSkillsAPIConfig
-from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
-from litellm.llms.base_llm.vector_store_files.transformation import (
- BaseVectorStoreFilesConfig,
-)
-from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
from ._logging import _is_debugging_on, verbose_logger
from .caching.caching import (
@@ -326,12 +399,6 @@ from .exceptions import (
UnprocessableEntityError,
UnsupportedParamsError,
)
-from .types.llms.openai import (
- ChatCompletionDeltaToolCallChunk,
- ChatCompletionToolCallChunk,
- ChatCompletionToolCallFunctionChunk,
-)
-from .types.router import LiteLLM_Params
if TYPE_CHECKING:
from litellm import MockException
@@ -487,7 +554,7 @@ def _add_custom_logger_callback_to_specific_event(
def _custom_logger_class_exists_in_success_callbacks(
- callback_class: CustomLogger,
+ callback_class: "CustomLogger",
) -> bool:
"""
Returns True if an instance of the custom logger exists in litellm.success_callback or litellm._async_success_callback
@@ -503,7 +570,7 @@ def _custom_logger_class_exists_in_success_callbacks(
def _custom_logger_class_exists_in_failure_callbacks(
- callback_class: CustomLogger,
+ callback_class: "CustomLogger",
) -> bool:
"""
Returns True if an instance of the custom logger exists in litellm.failure_callback or litellm._async_failure_callback
@@ -536,6 +603,7 @@ def get_applied_guardrails(kwargs: Dict[str, Any]) -> List[str]:
request_guardrails = get_request_guardrails(kwargs)
applied_guardrails = []
+ CustomGuardrail = _get_cached_custom_guardrail()
for callback in litellm.callbacks:
if callback is not None and isinstance(callback, CustomGuardrail):
if callback.guardrail_name is not None:
@@ -547,10 +615,22 @@ def get_applied_guardrails(kwargs: Dict[str, Any]) -> List[str]:
return applied_guardrails
+def _get_utils_globals() -> dict:
+ """
+ Get the globals dictionary of the utils module.
+
+ This is where we cache imported attributes so we don't import them twice.
+ """
+ return sys.modules[__name__].__dict__
+
+
def load_credentials_from_list(kwargs: dict):
"""
Updates kwargs with the credentials if credential_name in kwarg
"""
+ # Access CredentialAccessor via module to trigger lazy loading if needed
+ CredentialAccessor = getattr(sys.modules[__name__], 'CredentialAccessor')
+
credential_name = kwargs.get("litellm_credential_name")
if credential_name and litellm.credential_list:
credential_accessor = CredentialAccessor.get_credential_values(credential_name)
@@ -560,7 +640,7 @@ def load_credentials_from_list(kwargs: dict):
def get_dynamic_callbacks(
- dynamic_callbacks: Optional[List[Union[str, Callable, CustomLogger]]],
+ dynamic_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]],
) -> List:
returned_callbacks = litellm.callbacks.copy()
if dynamic_callbacks:
@@ -694,7 +774,7 @@ def function_setup( # noqa: PLR0915
function_id: Optional[str] = kwargs["id"] if "id" in kwargs else None
## DYNAMIC CALLBACKS ##
- dynamic_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = (
+ dynamic_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = (
kwargs.pop("callbacks", None)
)
all_callbacks = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks)
@@ -739,6 +819,7 @@ def function_setup( # noqa: PLR0915
+ litellm.failure_callback
)
)
+ get_set_callbacks = getattr(sys.modules[__name__], 'get_set_callbacks')
get_set_callbacks()(callback_list=callback_list, function_id=function_id)
## ASYNC CALLBACKS
if len(litellm.input_callback) > 0:
@@ -795,16 +876,16 @@ def function_setup( # noqa: PLR0915
litellm.failure_callback.pop(index)
### DYNAMIC CALLBACKS ###
dynamic_success_callbacks: Optional[
- List[Union[str, Callable, CustomLogger]]
+ List[Union[str, Callable, "CustomLogger"]]
] = None
dynamic_async_success_callbacks: Optional[
- List[Union[str, Callable, CustomLogger]]
+ List[Union[str, Callable, "CustomLogger"]]
] = None
dynamic_failure_callbacks: Optional[
- List[Union[str, Callable, CustomLogger]]
+ List[Union[str, Callable, "CustomLogger"]]
] = None
dynamic_async_failure_callbacks: Optional[
- List[Union[str, Callable, CustomLogger]]
+ List[Union[str, Callable, "CustomLogger"]]
] = None
if kwargs.get("success_callback", None) is not None and isinstance(
kwargs["success_callback"], list
@@ -866,6 +947,7 @@ def function_setup( # noqa: PLR0915
elif kwargs.get("messages", None):
messages = kwargs["messages"]
### PRE-CALL RULES ###
+ Rules = getattr(sys.modules[__name__], 'Rules')
if (
Rules.has_pre_call_rules()
and isinstance(messages, list)
@@ -965,7 +1047,9 @@ def function_setup( # noqa: PLR0915
or call_type == CallTypes.transcription.value
):
_file_obj: FileTypes = args[1] if len(args) > 1 else kwargs["file"]
- file_checksum = litellm.litellm_core_utils.audio_utils.utils.get_audio_file_content_hash(
+ # Lazy import audio_utils.utils only when needed for transcription calls
+ audio_utils = _get_cached_audio_utils()
+ file_checksum = audio_utils.get_audio_file_content_hash(
file_obj=_file_obj
)
if "metadata" in kwargs:
@@ -996,6 +1080,7 @@ def function_setup( # noqa: PLR0915
call_type=call_type,
):
stream = True
+ get_litellm_logging_class = getattr(sys.modules[__name__], 'get_litellm_logging_class')
logging_obj = get_litellm_logging_class()( # Victim for object pool
model=model, # type: ignore
messages=messages,
@@ -1079,6 +1164,8 @@ def _get_wrapper_num_retries(
if num_retries is None:
num_retries = litellm.num_retries
if kwargs.get("retry_policy", None):
+ get_num_retries_from_retry_policy = getattr(sys.modules[__name__], 'get_num_retries_from_retry_policy')
+ reset_retry_policy = getattr(sys.modules[__name__], 'reset_retry_policy')
retry_policy_num_retries = get_num_retries_from_retry_policy(
exception=exception,
retry_policy=kwargs.get("retry_policy"),
@@ -1122,6 +1209,7 @@ async def async_pre_call_deployment_hook(kwargs: Dict[str, Any], call_type: str)
modified_kwargs = kwargs.copy()
+ CustomLogger = _get_cached_custom_logger()
for callback in litellm.callbacks:
if isinstance(callback, CustomLogger):
result = await callback.async_pre_call_deployment_hook(
@@ -1144,6 +1232,7 @@ async def async_post_call_success_deployment_hook(
except ValueError:
typed_call_type = None # unknown call type
+ CustomLogger = _get_cached_custom_logger()
for callback in litellm.callbacks:
if isinstance(callback, CustomLogger):
result = await callback.async_post_call_success_deployment_hook(
@@ -1262,6 +1351,7 @@ def post_call_processing(
def client(original_function): # noqa: PLR0915
+ Rules = getattr(sys.modules[__name__], 'Rules')
rules_obj = Rules()
@wraps(original_function)
@@ -1323,7 +1413,8 @@ def client(original_function): # noqa: PLR0915
## LOAD CREDENTIALS
load_credentials_from_list(kwargs)
kwargs["litellm_logging_obj"] = logging_obj
- _llm_caching_handler: LLMCachingHandler = LLMCachingHandler(
+ LLMCachingHandler = _get_cached_llm_caching_handler()
+ _llm_caching_handler: "LLMCachingHandler" = LLMCachingHandler(
original_function=original_function,
request_kwargs=kwargs,
start_time=start_time,
@@ -1378,7 +1469,7 @@ def client(original_function): # noqa: PLR0915
): # allow users to control returning cached responses from the completion function
# checking cache
verbose_logger.debug("INSIDE CHECKING SYNC CACHE")
- caching_handler_response: CachingHandlerResponse = (
+ caching_handler_response: "CachingHandlerResponse" = (
_llm_caching_handler._sync_get_cache(
model=model or "",
original_function=original_function,
@@ -1446,6 +1537,7 @@ def client(original_function): # noqa: PLR0915
)
else:
# RETURN RESULT
+ update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata')
update_response_metadata(
result=result,
logging_obj=logging_obj,
@@ -1489,6 +1581,7 @@ def client(original_function): # noqa: PLR0915
# Copy the current context to propagate it to the background thread
# This is essential for OpenTelemetry span context propagation
ctx = contextvars.copy_context()
+ executor = getattr(sys.modules[__name__], 'executor')
executor.submit(
ctx.run,
logging_obj.success_handler,
@@ -1497,6 +1590,7 @@ def client(original_function): # noqa: PLR0915
end_time,
)
# RETURN RESULT
+ update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata')
update_response_metadata(
result=result,
logging_obj=logging_obj,
@@ -1513,6 +1607,8 @@ def client(original_function): # noqa: PLR0915
kwargs.get("num_retries", None) or litellm.num_retries or None
)
if kwargs.get("retry_policy", None):
+ get_num_retries_from_retry_policy = getattr(sys.modules[__name__], 'get_num_retries_from_retry_policy')
+ reset_retry_policy = getattr(sys.modules[__name__], 'reset_retry_policy')
num_retries = get_num_retries_from_retry_policy(
exception=e,
retry_policy=kwargs.get("retry_policy"),
@@ -1569,7 +1665,8 @@ def client(original_function): # noqa: PLR0915
logging_obj: Optional[LiteLLMLoggingObject] = kwargs.get(
"litellm_logging_obj", None
)
- _llm_caching_handler: LLMCachingHandler = LLMCachingHandler(
+ LLMCachingHandler = _get_cached_llm_caching_handler()
+ _llm_caching_handler: "LLMCachingHandler" = LLMCachingHandler(
original_function=original_function,
request_kwargs=kwargs,
start_time=start_time,
@@ -1608,7 +1705,7 @@ def client(original_function): # noqa: PLR0915
print_verbose(
f"ASYNC kwargs[caching]: {kwargs.get('caching', False)}; litellm.cache: {litellm.cache}; kwargs.get('cache'): {kwargs.get('cache', None)}"
)
- _caching_handler_response: Optional[CachingHandlerResponse] = (
+ _caching_handler_response: "Optional[CachingHandlerResponse]" = (
await _llm_caching_handler._async_get_cache(
model=model or "",
original_function=original_function,
@@ -1683,6 +1780,7 @@ def client(original_function): # noqa: PLR0915
chunks, messages=kwargs.get("messages", None)
)
else:
+ update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata')
update_response_metadata(
result=result,
logging_obj=logging_obj,
@@ -1747,6 +1845,7 @@ def client(original_function): # noqa: PLR0915
end_time=end_time,
)
+ update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata')
update_response_metadata(
result=result,
logging_obj=logging_obj,
@@ -1822,6 +1921,7 @@ def client(original_function): # noqa: PLR0915
setattr(e, "timeout", timeout)
raise e
+ get_coroutine_checker = getattr(sys.modules[__name__], 'get_coroutine_checker')
is_coroutine = get_coroutine_checker().is_async_callable(original_function)
# Return the appropriate wrapper based on the original function type
@@ -2182,6 +2282,7 @@ def supports_response_schema(
"""
## GET LLM PROVIDER ##
try:
+ get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider')
model, custom_llm_provider, _, _ = get_llm_provider(
model=model, custom_llm_provider=custom_llm_provider
)
@@ -2879,6 +2980,9 @@ def get_optional_params_embeddings( # noqa: PLR0915
additional_drop_params: Optional[List[str]] = None,
**kwargs,
):
+ # Lazy load get_supported_openai_params
+ get_supported_openai_params = getattr(sys.modules[__name__], 'get_supported_openai_params')
+
# retrieve all parameters passed to the function
passed_params = locals()
custom_llm_provider = passed_params.pop("custom_llm_provider", None)
@@ -3148,6 +3252,21 @@ def get_optional_params_embeddings( # noqa: PLR0915
drop_params=drop_params if drop_params is not None else False,
)
+ elif custom_llm_provider == "ollama":
+ if 'dimensions' in non_default_params:
+ optional_params['dimensions']=non_default_params.pop('dimensions')
+ if len(non_default_params.keys()) > 0:
+ if (
+ litellm.drop_params is True or drop_params is True
+ ): # drop the unsupported non-default values
+ keys = list(non_default_params.keys())
+ for k in keys:
+ non_default_params.pop(k, None)
+ else:
+ raise UnsupportedParamsError(
+ status_code=500,
+ message=f"Setting {non_default_params} is not supported by {custom_llm_provider}. To drop it from the call, set `litellm.drop_params = True`.",
+ )
elif (
custom_llm_provider != "openai"
and custom_llm_provider != "azure"
@@ -3666,6 +3785,7 @@ def get_optional_params( # noqa: PLR0915
message=f"{custom_llm_provider} does not support parameters: {list(unsupported_params.keys())}, for model={model}. To drop these, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n. \n If you want to use these params dynamically send allowed_openai_params={list(unsupported_params.keys())} in your request.",
)
+ get_supported_openai_params = getattr(sys.modules[__name__], 'get_supported_openai_params')
supported_params = get_supported_openai_params(
model=model, custom_llm_provider=custom_llm_provider
)
@@ -3920,6 +4040,7 @@ def get_optional_params( # noqa: PLR0915
),
)
elif custom_llm_provider == "bedrock":
+ BedrockModelInfo = getattr(sys.modules[__name__], 'BedrockModelInfo')
bedrock_route = BedrockModelInfo.get_bedrock_route(model)
bedrock_base_model = BedrockModelInfo.get_base_model(model)
if bedrock_route == "converse" or bedrock_route == "converse_like":
@@ -4344,6 +4465,8 @@ def get_optional_params( # noqa: PLR0915
# Apply nested drops from additional_drop_params
if additional_drop_params:
+ is_nested_path = getattr(sys.modules[__name__], 'is_nested_path')
+ delete_nested_value = getattr(sys.modules[__name__], 'delete_nested_value')
nested_paths = [p for p in additional_drop_params if is_nested_path(p)]
for path in nested_paths:
optional_params = delete_nested_value(optional_params, path)
@@ -4393,6 +4516,7 @@ def add_provider_specific_params_to_optional_params(
else:
processed_extra_body = initial_extra_body
+ _ensure_extra_body_is_safe = getattr(sys.modules[__name__], '_ensure_extra_body_is_safe')
optional_params["extra_body"] = _ensure_extra_body_is_safe(
extra_body=processed_extra_body
)
@@ -4803,6 +4927,7 @@ def get_max_tokens(model: str) -> Optional[int]:
return litellm.model_cost[model]["max_output_tokens"]
elif "max_tokens" in litellm.model_cost[model]:
return litellm.model_cost[model]["max_tokens"]
+ get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider')
model, custom_llm_provider, _, _ = get_llm_provider(model=model)
if custom_llm_provider == "huggingface":
max_tokens = _get_max_position_embeddings(model_name=model)
@@ -4923,6 +5048,7 @@ def _get_potential_model_names(
if custom_llm_provider is None:
# Get custom_llm_provider
try:
+ get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider')
split_model, custom_llm_provider, _, _ = get_llm_provider(model=model)
except Exception:
split_model = model
@@ -5216,6 +5342,9 @@ def _get_model_info_helper( # noqa: PLR0915
input_cost_per_audio_token=_model_info.get(
"input_cost_per_audio_token", None
),
+ input_cost_per_image_token=_model_info.get(
+ "input_cost_per_image_token", None
+ ),
input_cost_per_token_batches=_model_info.get(
"input_cost_per_token_batches"
),
@@ -5642,6 +5771,7 @@ def validate_environment( # noqa: PLR0915
}
## EXTRACT LLM PROVIDER - if model name provided
try:
+ get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider')
_, custom_llm_provider, _, _ = get_llm_provider(model=model)
except Exception:
custom_llm_provider = None
@@ -6204,6 +6334,7 @@ def register_prompt_template(
complete_model = model
potential_models = [complete_model]
try:
+ get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider')
model = get_llm_provider(model=model)[0]
potential_models.append(model)
except Exception:
@@ -6289,6 +6420,7 @@ class TextCompletionStreamWrapper:
except StopIteration:
raise StopIteration
except Exception as e:
+ exception_type = getattr(sys.modules[__name__], 'exception_type')
raise exception_type(
model=self.model,
custom_llm_provider=self.custom_llm_provider or "",
@@ -6910,14 +7042,14 @@ def _get_base_model_from_metadata(model_call_details=None):
return _base_model
metadata = litellm_params.get("metadata", {})
- base_model_from_metadata = _get_base_model_from_litellm_call_metadata(
- metadata=metadata
- )
+ _get_base_model_from_litellm_call_metadata = getattr(sys.modules[__name__], '_get_base_model_from_litellm_call_metadata')
+ base_model_from_metadata = _get_base_model_from_litellm_call_metadata(metadata=metadata)
if base_model_from_metadata is not None:
return base_model_from_metadata
# Also check litellm_metadata (used by Responses API and other generic API calls)
litellm_metadata = litellm_params.get("litellm_metadata", {})
+ _get_base_model_from_litellm_call_metadata = getattr(sys.modules[__name__], '_get_base_model_from_litellm_call_metadata')
return _get_base_model_from_litellm_call_metadata(metadata=litellm_metadata)
return None
@@ -7329,6 +7461,7 @@ class ProviderConfigManager:
litellm.LlmProviders.COHERE_CHAT == provider
or litellm.LlmProviders.COHERE == provider
):
+ CohereModelInfo = getattr(sys.modules[__name__], 'CohereModelInfo')
route = CohereModelInfo.get_cohere_route(model)
if route == "v2":
return litellm.CohereV2ChatConfig()
@@ -8179,6 +8312,7 @@ class ProviderConfigManager:
return get_vertex_ai_ocr_config(model=model)
+ MistralOCRConfig = getattr(sys.modules[__name__], 'MistralOCRConfig')
PROVIDER_TO_CONFIG_MAP = {
litellm.LlmProviders.MISTRAL: MistralOCRConfig,
}
@@ -8319,6 +8453,7 @@ def get_end_user_id_for_cost_tracking(
service_type: "litellm_logging" or "prometheus" - used to allow prometheus only disable cost tracking.
"""
+ get_litellm_metadata_from_kwargs = getattr(sys.modules[__name__], 'get_litellm_metadata_from_kwargs')
_metadata = cast(
dict, get_litellm_metadata_from_kwargs(dict(litellm_params=litellm_params))
)
@@ -8610,16 +8745,672 @@ def should_run_mock_completion(
return False
-# Re-export encoding from main.py for backward compatibility
-# This allows tests to import: from litellm.utils import encoding
-# We use a lazy import to avoid loading main.py at utils.py import time
-def __getattr__(name: str) -> Any:
+def __getattr__(name: str) -> Any: # noqa: PLR0915
"""Lazy import handler for utils module"""
+ _globals = _get_utils_globals()
+
+ # Lazy load encoding from main.py to avoid heavy tiktoken import
if name == "encoding":
- # Cache it in the module's __dict__ for subsequent accesses
- import sys
-
- from litellm.main import encoding as _encoding
- sys.modules[__name__].__dict__["encoding"] = _encoding
- return _encoding
+ # Check if already cached
+ if "encoding" not in _globals:
+ from litellm.main import encoding as _encoding
+ _globals["encoding"] = _encoding
+ return _globals["encoding"]
+
+ # Lazy load BaseVectorStore to avoid loading it at module import time
+ if name == "BaseVectorStore":
+ # Check if already cached
+ if "BaseVectorStore" not in _globals:
+ from litellm.integrations.vector_store_integrations.base_vector_store import (
+ BaseVectorStore as _BaseVectorStore,
+ )
+ _globals["BaseVectorStore"] = _BaseVectorStore
+ return _globals["BaseVectorStore"]
+
+ # Lazy load CredentialAccessor to avoid loading it at module import time
+ if name == "CredentialAccessor":
+ # Check if already cached
+ if "CredentialAccessor" not in _globals:
+ from litellm.litellm_core_utils.credential_accessor import (
+ CredentialAccessor as _CredentialAccessor,
+ )
+ _globals["CredentialAccessor"] = _CredentialAccessor
+ return _globals["CredentialAccessor"]
+
+ # Lazy load exception_mapping_utils functions to avoid loading at module import time
+ if name == "exception_type":
+ # Check if already cached
+ if "exception_type" not in _globals:
+ from litellm.litellm_core_utils.exception_mapping_utils import (
+ exception_type as _exception_type,
+ )
+ _globals["exception_type"] = _exception_type
+ return _globals["exception_type"]
+
+ if name == "get_error_message":
+ # Check if already cached
+ if "get_error_message" not in _globals:
+ from litellm.litellm_core_utils.exception_mapping_utils import (
+ get_error_message as _get_error_message,
+ )
+ _globals["get_error_message"] = _get_error_message
+ return _globals["get_error_message"]
+
+ if name == "_get_response_headers":
+ # Check if already cached
+ if "_get_response_headers" not in _globals:
+ from litellm.litellm_core_utils.exception_mapping_utils import (
+ _get_response_headers as __get_response_headers,
+ )
+ _globals["_get_response_headers"] = __get_response_headers
+ return _globals["_get_response_headers"]
+
+ # Lazy load get_llm_provider_logic functions to avoid loading at module import time
+ if name == "get_llm_provider":
+ # Check if already cached
+ if "get_llm_provider" not in _globals:
+ from litellm.litellm_core_utils.get_llm_provider_logic import (
+ get_llm_provider as _get_llm_provider,
+ )
+ _globals["get_llm_provider"] = _get_llm_provider
+ return _globals["get_llm_provider"]
+
+ if name == "_is_non_openai_azure_model":
+ # Check if already cached
+ if "_is_non_openai_azure_model" not in _globals:
+ from litellm.litellm_core_utils.get_llm_provider_logic import (
+ _is_non_openai_azure_model as __is_non_openai_azure_model,
+ )
+ _globals["_is_non_openai_azure_model"] = __is_non_openai_azure_model
+ return _globals["_is_non_openai_azure_model"]
+
+ # Lazy load get_supported_openai_params to avoid loading at module import time
+ if name == "get_supported_openai_params":
+ # Check if already cached
+ if "get_supported_openai_params" not in _globals:
+ from litellm.litellm_core_utils.get_supported_openai_params import (
+ get_supported_openai_params as _get_supported_openai_params,
+ )
+ _globals["get_supported_openai_params"] = _get_supported_openai_params
+ return _globals["get_supported_openai_params"]
+
+ # Lazy load convert_dict_to_response functions to avoid loading at module import time
+ if name == "LiteLLMResponseObjectHandler":
+ # Check if already cached
+ if "LiteLLMResponseObjectHandler" not in _globals:
+ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
+ LiteLLMResponseObjectHandler as _LiteLLMResponseObjectHandler,
+ )
+ _globals["LiteLLMResponseObjectHandler"] = _LiteLLMResponseObjectHandler
+ return _globals["LiteLLMResponseObjectHandler"]
+
+ if name == "_handle_invalid_parallel_tool_calls":
+ # Check if already cached
+ if "_handle_invalid_parallel_tool_calls" not in _globals:
+ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
+ _handle_invalid_parallel_tool_calls as __handle_invalid_parallel_tool_calls,
+ )
+ _globals["_handle_invalid_parallel_tool_calls"] = __handle_invalid_parallel_tool_calls
+ return _globals["_handle_invalid_parallel_tool_calls"]
+
+ if name == "convert_to_model_response_object":
+ # Check if already cached
+ if "convert_to_model_response_object" not in _globals:
+ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
+ convert_to_model_response_object as _convert_to_model_response_object,
+ )
+ _globals["convert_to_model_response_object"] = _convert_to_model_response_object
+ return _globals["convert_to_model_response_object"]
+
+ if name == "convert_to_streaming_response":
+ # Check if already cached
+ if "convert_to_streaming_response" not in _globals:
+ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
+ convert_to_streaming_response as _convert_to_streaming_response,
+ )
+ _globals["convert_to_streaming_response"] = _convert_to_streaming_response
+ return _globals["convert_to_streaming_response"]
+
+ if name == "convert_to_streaming_response_async":
+ # Check if already cached
+ if "convert_to_streaming_response_async" not in _globals:
+ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
+ convert_to_streaming_response_async as _convert_to_streaming_response_async,
+ )
+ _globals["convert_to_streaming_response_async"] = _convert_to_streaming_response_async
+ return _globals["convert_to_streaming_response_async"]
+
+ # Lazy load get_api_base to avoid loading at module import time
+ if name == "get_api_base":
+ # Check if already cached
+ if "get_api_base" not in _globals:
+ from litellm.litellm_core_utils.llm_response_utils.get_api_base import (
+ get_api_base as _get_api_base,
+ )
+ _globals["get_api_base"] = _get_api_base
+ return _globals["get_api_base"]
+
+ # Lazy load ResponseMetadata to avoid loading at module import time
+ if name == "ResponseMetadata":
+ # Check if already cached
+ if "ResponseMetadata" not in _globals:
+ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
+ ResponseMetadata as _ResponseMetadata,
+ )
+ _globals["ResponseMetadata"] = _ResponseMetadata
+ return _globals["ResponseMetadata"]
+
+ # Lazy load _parse_content_for_reasoning to avoid loading at module import time
+ if name == "_parse_content_for_reasoning":
+ # Check if already cached
+ if "_parse_content_for_reasoning" not in _globals:
+ from litellm.litellm_core_utils.prompt_templates.common_utils import (
+ _parse_content_for_reasoning as __parse_content_for_reasoning,
+ )
+ _globals["_parse_content_for_reasoning"] = __parse_content_for_reasoning
+ return _globals["_parse_content_for_reasoning"]
+
+ # Lazy load redact_messages to avoid loading at module import time
+ if name == "LiteLLMLoggingObject":
+ # Check if already cached
+ if "LiteLLMLoggingObject" not in _globals:
+ from litellm.litellm_core_utils.redact_messages import (
+ LiteLLMLoggingObject as _LiteLLMLoggingObject,
+ )
+ _globals["LiteLLMLoggingObject"] = _LiteLLMLoggingObject
+ return _globals["LiteLLMLoggingObject"]
+
+ if name == "redact_message_input_output_from_logging":
+ # Check if already cached
+ if "redact_message_input_output_from_logging" not in _globals:
+ from litellm.litellm_core_utils.redact_messages import (
+ redact_message_input_output_from_logging as _redact_message_input_output_from_logging,
+ )
+ _globals["redact_message_input_output_from_logging"] = _redact_message_input_output_from_logging
+ return _globals["redact_message_input_output_from_logging"]
+
+ # Lazy load CustomStreamWrapper to avoid loading at module import time
+ if name == "CustomStreamWrapper":
+ # Check if already cached
+ if "CustomStreamWrapper" not in _globals:
+ from litellm.litellm_core_utils.streaming_handler import (
+ CustomStreamWrapper as _CustomStreamWrapper,
+ )
+ _globals["CustomStreamWrapper"] = _CustomStreamWrapper
+ return _globals["CustomStreamWrapper"]
+
+ # Lazy load BaseGoogleGenAIGenerateContentConfig to avoid loading at module import time
+ if name == "BaseGoogleGenAIGenerateContentConfig":
+ # Check if already cached
+ if "BaseGoogleGenAIGenerateContentConfig" not in _globals:
+ from litellm.llms.base_llm.google_genai.transformation import (
+ BaseGoogleGenAIGenerateContentConfig as _BaseGoogleGenAIGenerateContentConfig,
+ )
+ _globals["BaseGoogleGenAIGenerateContentConfig"] = _BaseGoogleGenAIGenerateContentConfig
+ return _globals["BaseGoogleGenAIGenerateContentConfig"]
+
+ # Lazy load BaseOCRConfig to avoid loading at module import time
+ if name == "BaseOCRConfig":
+ # Check if already cached
+ if "BaseOCRConfig" not in _globals:
+ from litellm.llms.base_llm.ocr.transformation import (
+ BaseOCRConfig as _BaseOCRConfig,
+ )
+ _globals["BaseOCRConfig"] = _BaseOCRConfig
+ return _globals["BaseOCRConfig"]
+
+ # Lazy load BaseSearchConfig to avoid loading at module import time
+ if name == "BaseSearchConfig":
+ # Check if already cached
+ if "BaseSearchConfig" not in _globals:
+ from litellm.llms.base_llm.search.transformation import (
+ BaseSearchConfig as _BaseSearchConfig,
+ )
+ _globals["BaseSearchConfig"] = _BaseSearchConfig
+ return _globals["BaseSearchConfig"]
+
+ # Lazy load BaseTextToSpeechConfig to avoid loading at module import time
+ if name == "BaseTextToSpeechConfig":
+ # Check if already cached
+ if "BaseTextToSpeechConfig" not in _globals:
+ from litellm.llms.base_llm.text_to_speech.transformation import (
+ BaseTextToSpeechConfig as _BaseTextToSpeechConfig,
+ )
+ _globals["BaseTextToSpeechConfig"] = _BaseTextToSpeechConfig
+ return _globals["BaseTextToSpeechConfig"]
+
+ # Lazy load BedrockModelInfo to avoid loading at module import time
+ if name == "BedrockModelInfo":
+ # Check if already cached
+ if "BedrockModelInfo" not in _globals:
+ from litellm.llms.bedrock.common_utils import (
+ BedrockModelInfo as _BedrockModelInfo,
+ )
+ _globals["BedrockModelInfo"] = _BedrockModelInfo
+ return _globals["BedrockModelInfo"]
+
+ # Lazy load CohereModelInfo to avoid loading at module import time
+ if name == "CohereModelInfo":
+ # Check if already cached
+ if "CohereModelInfo" not in _globals:
+ from litellm.llms.cohere.common_utils import (
+ CohereModelInfo as _CohereModelInfo,
+ )
+ _globals["CohereModelInfo"] = _CohereModelInfo
+ return _globals["CohereModelInfo"]
+
+ # Lazy load MistralOCRConfig to avoid loading at module import time
+ if name == "MistralOCRConfig":
+ # Check if already cached
+ if "MistralOCRConfig" not in _globals:
+ from litellm.llms.mistral.ocr.transformation import (
+ MistralOCRConfig as _MistralOCRConfig,
+ )
+ _globals["MistralOCRConfig"] = _MistralOCRConfig
+ return _globals["MistralOCRConfig"]
+
+ # Lazy load Rules to avoid loading at module import time
+ if name == "Rules":
+ # Check if already cached
+ if "Rules" not in _globals:
+ from litellm.litellm_core_utils.rules import Rules as _Rules
+ _globals["Rules"] = _Rules
+ return _globals["Rules"]
+
+ # Lazy load AsyncHTTPHandler and HTTPHandler to avoid loading at module import time
+ if name == "AsyncHTTPHandler":
+ # Check if already cached
+ if "AsyncHTTPHandler" not in _globals:
+ from litellm.llms.custom_httpx.http_handler import (
+ AsyncHTTPHandler as _AsyncHTTPHandler,
+ )
+ _globals["AsyncHTTPHandler"] = _AsyncHTTPHandler
+ return _globals["AsyncHTTPHandler"]
+
+ if name == "HTTPHandler":
+ # Check if already cached
+ if "HTTPHandler" not in _globals:
+ from litellm.llms.custom_httpx.http_handler import (
+ HTTPHandler as _HTTPHandler,
+ )
+ _globals["HTTPHandler"] = _HTTPHandler
+ return _globals["HTTPHandler"]
+
+ # Lazy load get_num_retries_from_retry_policy and reset_retry_policy to avoid loading at module import time
+ if name == "get_num_retries_from_retry_policy":
+ # Check if already cached
+ if "get_num_retries_from_retry_policy" not in _globals:
+ from litellm.router_utils.get_retry_from_policy import (
+ get_num_retries_from_retry_policy as _get_num_retries_from_retry_policy,
+ )
+ _globals["get_num_retries_from_retry_policy"] = _get_num_retries_from_retry_policy
+ return _globals["get_num_retries_from_retry_policy"]
+
+ if name == "reset_retry_policy":
+ # Check if already cached
+ if "reset_retry_policy" not in _globals:
+ from litellm.router_utils.get_retry_from_policy import (
+ reset_retry_policy as _reset_retry_policy,
+ )
+ _globals["reset_retry_policy"] = _reset_retry_policy
+ return _globals["reset_retry_policy"]
+
+ # Lazy load get_secret to avoid loading at module import time
+ if name == "get_secret":
+ # Check if already cached
+ if "get_secret" not in _globals:
+ from litellm.secret_managers.main import get_secret as _get_secret
+ _globals["get_secret"] = _get_secret
+ return _globals["get_secret"]
+
+ # Lazy load cached_imports functions to avoid loading at module import time
+ if name == "get_coroutine_checker":
+ # Check if already cached
+ if "get_coroutine_checker" not in _globals:
+ from litellm.litellm_core_utils.cached_imports import (
+ get_coroutine_checker as _get_coroutine_checker,
+ )
+ _globals["get_coroutine_checker"] = _get_coroutine_checker
+ return _globals["get_coroutine_checker"]
+
+ if name == "get_litellm_logging_class":
+ # Check if already cached
+ if "get_litellm_logging_class" not in _globals:
+ from litellm.litellm_core_utils.cached_imports import (
+ get_litellm_logging_class as _get_litellm_logging_class,
+ )
+ _globals["get_litellm_logging_class"] = _get_litellm_logging_class
+ return _globals["get_litellm_logging_class"]
+
+ if name == "get_set_callbacks":
+ # Check if already cached
+ if "get_set_callbacks" not in _globals:
+ from litellm.litellm_core_utils.cached_imports import (
+ get_set_callbacks as _get_set_callbacks,
+ )
+ _globals["get_set_callbacks"] = _get_set_callbacks
+ return _globals["get_set_callbacks"]
+
+ # Lazy load core_helpers functions to avoid loading at module import time
+ if name == "get_litellm_metadata_from_kwargs":
+ # Check if already cached
+ if "get_litellm_metadata_from_kwargs" not in _globals:
+ from litellm.litellm_core_utils.core_helpers import (
+ get_litellm_metadata_from_kwargs as _get_litellm_metadata_from_kwargs,
+ )
+ _globals["get_litellm_metadata_from_kwargs"] = _get_litellm_metadata_from_kwargs
+ return _globals["get_litellm_metadata_from_kwargs"]
+
+ if name == "map_finish_reason":
+ # Check if already cached
+ if "map_finish_reason" not in _globals:
+ from litellm.litellm_core_utils.core_helpers import (
+ map_finish_reason as _map_finish_reason,
+ )
+ _globals["map_finish_reason"] = _map_finish_reason
+ return _globals["map_finish_reason"]
+
+ if name == "process_response_headers":
+ # Check if already cached
+ if "process_response_headers" not in _globals:
+ from litellm.litellm_core_utils.core_helpers import (
+ process_response_headers as _process_response_headers,
+ )
+ _globals["process_response_headers"] = _process_response_headers
+ return _globals["process_response_headers"]
+
+ # Lazy load dot_notation_indexing functions to avoid loading at module import time
+ if name == "delete_nested_value":
+ # Check if already cached
+ if "delete_nested_value" not in _globals:
+ from litellm.litellm_core_utils.dot_notation_indexing import (
+ delete_nested_value as _delete_nested_value,
+ )
+ _globals["delete_nested_value"] = _delete_nested_value
+ return _globals["delete_nested_value"]
+
+ if name == "is_nested_path":
+ # Check if already cached
+ if "is_nested_path" not in _globals:
+ from litellm.litellm_core_utils.dot_notation_indexing import (
+ is_nested_path as _is_nested_path,
+ )
+ _globals["is_nested_path"] = _is_nested_path
+ return _globals["is_nested_path"]
+
+ # Lazy load get_litellm_params functions to avoid loading at module import time
+ if name == "_get_base_model_from_litellm_call_metadata":
+ # Check if already cached
+ if "_get_base_model_from_litellm_call_metadata" not in _globals:
+ from litellm.litellm_core_utils.get_litellm_params import (
+ _get_base_model_from_litellm_call_metadata as __get_base_model_from_litellm_call_metadata,
+ )
+ _globals["_get_base_model_from_litellm_call_metadata"] = __get_base_model_from_litellm_call_metadata
+ return _globals["_get_base_model_from_litellm_call_metadata"]
+
+ if name == "get_litellm_params":
+ # Check if already cached
+ if "get_litellm_params" not in _globals:
+ from litellm.litellm_core_utils.get_litellm_params import (
+ get_litellm_params as _get_litellm_params,
+ )
+ _globals["get_litellm_params"] = _get_litellm_params
+ return _globals["get_litellm_params"]
+
+ # Lazy load _ensure_extra_body_is_safe to avoid loading at module import time
+ if name == "_ensure_extra_body_is_safe":
+ # Check if already cached
+ if "_ensure_extra_body_is_safe" not in _globals:
+ from litellm.litellm_core_utils.llm_request_utils import (
+ _ensure_extra_body_is_safe as __ensure_extra_body_is_safe,
+ )
+ _globals["_ensure_extra_body_is_safe"] = __ensure_extra_body_is_safe
+ return _globals["_ensure_extra_body_is_safe"]
+
+ # Lazy load get_formatted_prompt to avoid loading at module import time
+ if name == "get_formatted_prompt":
+ # Check if already cached
+ if "get_formatted_prompt" not in _globals:
+ from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import (
+ get_formatted_prompt as _get_formatted_prompt,
+ )
+ _globals["get_formatted_prompt"] = _get_formatted_prompt
+ return _globals["get_formatted_prompt"]
+
+ # Lazy load get_response_headers to avoid loading at module import time
+ if name == "get_response_headers":
+ # Check if already cached
+ if "get_response_headers" not in _globals:
+ from litellm.litellm_core_utils.llm_response_utils.get_headers import (
+ get_response_headers as _get_response_headers,
+ )
+ _globals["get_response_headers"] = _get_response_headers
+ return _globals["get_response_headers"]
+
+ # Lazy load update_response_metadata to avoid loading at module import time
+ if name == "update_response_metadata":
+ # Check if already cached
+ if "update_response_metadata" not in _globals:
+ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
+ update_response_metadata as _update_response_metadata,
+ )
+ _globals["update_response_metadata"] = _update_response_metadata
+ return _globals["update_response_metadata"]
+
+ # Lazy load executor to avoid loading at module import time
+ if name == "executor":
+ # Check if already cached
+ if "executor" not in _globals:
+ from litellm.litellm_core_utils.thread_pool_executor import (
+ executor as _executor,
+ )
+ _globals["executor"] = _executor
+ return _globals["executor"]
+
+ # Lazy load BaseAnthropicMessagesConfig to avoid loading at module import time
+ if name == "BaseAnthropicMessagesConfig":
+ # Check if already cached
+ if "BaseAnthropicMessagesConfig" not in _globals:
+ from litellm.llms.base_llm.anthropic_messages.transformation import (
+ BaseAnthropicMessagesConfig as _BaseAnthropicMessagesConfig,
+ )
+ _globals["BaseAnthropicMessagesConfig"] = _BaseAnthropicMessagesConfig
+ return _globals["BaseAnthropicMessagesConfig"]
+
+ # Lazy load BaseAudioTranscriptionConfig to avoid loading at module import time
+ if name == "BaseAudioTranscriptionConfig":
+ # Check if already cached
+ if "BaseAudioTranscriptionConfig" not in _globals:
+ from litellm.llms.base_llm.audio_transcription.transformation import (
+ BaseAudioTranscriptionConfig as _BaseAudioTranscriptionConfig,
+ )
+ _globals["BaseAudioTranscriptionConfig"] = _BaseAudioTranscriptionConfig
+ return _globals["BaseAudioTranscriptionConfig"]
+
+ # Lazy load BaseBatchesConfig to avoid loading at module import time
+ if name == "BaseBatchesConfig":
+ # Check if already cached
+ if "BaseBatchesConfig" not in _globals:
+ from litellm.llms.base_llm.batches.transformation import (
+ BaseBatchesConfig as _BaseBatchesConfig,
+ )
+ _globals["BaseBatchesConfig"] = _BaseBatchesConfig
+ return _globals["BaseBatchesConfig"]
+
+ # Lazy load BaseContainerConfig to avoid loading at module import time
+ if name == "BaseContainerConfig":
+ # Check if already cached
+ if "BaseContainerConfig" not in _globals:
+ from litellm.llms.base_llm.containers.transformation import (
+ BaseContainerConfig as _BaseContainerConfig,
+ )
+ _globals["BaseContainerConfig"] = _BaseContainerConfig
+ return _globals["BaseContainerConfig"]
+
+ # Lazy load BaseEmbeddingConfig to avoid loading at module import time
+ if name == "BaseEmbeddingConfig":
+ # Check if already cached
+ if "BaseEmbeddingConfig" not in _globals:
+ from litellm.llms.base_llm.embedding.transformation import (
+ BaseEmbeddingConfig as _BaseEmbeddingConfig,
+ )
+ _globals["BaseEmbeddingConfig"] = _BaseEmbeddingConfig
+ return _globals["BaseEmbeddingConfig"]
+
+ # Lazy load BaseImageEditConfig to avoid loading at module import time
+ if name == "BaseImageEditConfig":
+ # Check if already cached
+ if "BaseImageEditConfig" not in _globals:
+ from litellm.llms.base_llm.image_edit.transformation import (
+ BaseImageEditConfig as _BaseImageEditConfig,
+ )
+ _globals["BaseImageEditConfig"] = _BaseImageEditConfig
+ return _globals["BaseImageEditConfig"]
+
+ # Lazy load BaseImageGenerationConfig to avoid loading at module import time
+ if name == "BaseImageGenerationConfig":
+ # Check if already cached
+ if "BaseImageGenerationConfig" not in _globals:
+ from litellm.llms.base_llm.image_generation.transformation import (
+ BaseImageGenerationConfig as _BaseImageGenerationConfig,
+ )
+ _globals["BaseImageGenerationConfig"] = _BaseImageGenerationConfig
+ return _globals["BaseImageGenerationConfig"]
+
+ # Lazy load BaseImageVariationConfig to avoid loading at module import time
+ if name == "BaseImageVariationConfig":
+ # Check if already cached
+ if "BaseImageVariationConfig" not in _globals:
+ from litellm.llms.base_llm.image_variations.transformation import (
+ BaseImageVariationConfig as _BaseImageVariationConfig,
+ )
+ _globals["BaseImageVariationConfig"] = _BaseImageVariationConfig
+ return _globals["BaseImageVariationConfig"]
+
+ # Lazy load BasePassthroughConfig to avoid loading at module import time
+ if name == "BasePassthroughConfig":
+ # Check if already cached
+ if "BasePassthroughConfig" not in _globals:
+ from litellm.llms.base_llm.passthrough.transformation import (
+ BasePassthroughConfig as _BasePassthroughConfig,
+ )
+ _globals["BasePassthroughConfig"] = _BasePassthroughConfig
+ return _globals["BasePassthroughConfig"]
+
+ # Lazy load BaseRealtimeConfig to avoid loading at module import time
+ if name == "BaseRealtimeConfig":
+ # Check if already cached
+ if "BaseRealtimeConfig" not in _globals:
+ from litellm.llms.base_llm.realtime.transformation import (
+ BaseRealtimeConfig as _BaseRealtimeConfig,
+ )
+ _globals["BaseRealtimeConfig"] = _BaseRealtimeConfig
+ return _globals["BaseRealtimeConfig"]
+
+ # Lazy load BaseRerankConfig to avoid loading at module import time
+ if name == "BaseRerankConfig":
+ # Check if already cached
+ if "BaseRerankConfig" not in _globals:
+ from litellm.llms.base_llm.rerank.transformation import (
+ BaseRerankConfig as _BaseRerankConfig,
+ )
+ _globals["BaseRerankConfig"] = _BaseRerankConfig
+ return _globals["BaseRerankConfig"]
+
+ # Lazy load BaseVectorStoreConfig to avoid loading at module import time
+ if name == "BaseVectorStoreConfig":
+ # Check if already cached
+ if "BaseVectorStoreConfig" not in _globals:
+ from litellm.llms.base_llm.vector_store.transformation import (
+ BaseVectorStoreConfig as _BaseVectorStoreConfig,
+ )
+ _globals["BaseVectorStoreConfig"] = _BaseVectorStoreConfig
+ return _globals["BaseVectorStoreConfig"]
+
+ # Lazy load BaseVectorStoreFilesConfig to avoid loading at module import time
+ if name == "BaseVectorStoreFilesConfig":
+ # Check if already cached
+ if "BaseVectorStoreFilesConfig" not in _globals:
+ from litellm.llms.base_llm.vector_store_files.transformation import (
+ BaseVectorStoreFilesConfig as _BaseVectorStoreFilesConfig,
+ )
+ _globals["BaseVectorStoreFilesConfig"] = _BaseVectorStoreFilesConfig
+ return _globals["BaseVectorStoreFilesConfig"]
+
+ # Lazy load BaseVideoConfig to avoid loading at module import time
+ if name == "BaseVideoConfig":
+ # Check if already cached
+ if "BaseVideoConfig" not in _globals:
+ from litellm.llms.base_llm.videos.transformation import (
+ BaseVideoConfig as _BaseVideoConfig,
+ )
+ _globals["BaseVideoConfig"] = _BaseVideoConfig
+ return _globals["BaseVideoConfig"]
+
+ # Lazy load ANTHROPIC_API_ONLY_HEADERS to avoid loading at module import time
+ if name == "ANTHROPIC_API_ONLY_HEADERS":
+ # Check if already cached
+ if "ANTHROPIC_API_ONLY_HEADERS" not in _globals:
+ from litellm.types.llms.anthropic import (
+ ANTHROPIC_API_ONLY_HEADERS as _ANTHROPIC_API_ONLY_HEADERS,
+ )
+ _globals["ANTHROPIC_API_ONLY_HEADERS"] = _ANTHROPIC_API_ONLY_HEADERS
+ return _globals["ANTHROPIC_API_ONLY_HEADERS"]
+
+ # Lazy load AnthropicThinkingParam to avoid loading at module import time
+ if name == "AnthropicThinkingParam":
+ # Check if already cached
+ if "AnthropicThinkingParam" not in _globals:
+ from litellm.types.llms.anthropic import (
+ AnthropicThinkingParam as _AnthropicThinkingParam,
+ )
+ _globals["AnthropicThinkingParam"] = _AnthropicThinkingParam
+ return _globals["AnthropicThinkingParam"]
+
+ # Lazy load RerankResponse to avoid loading at module import time
+ if name == "RerankResponse":
+ # Check if already cached
+ if "RerankResponse" not in _globals:
+ from litellm.types.rerank import RerankResponse as _RerankResponse
+ _globals["RerankResponse"] = _RerankResponse
+ return _globals["RerankResponse"]
+
+ # Lazy load ChatCompletionDeltaToolCallChunk to avoid loading at module import time
+ if name == "ChatCompletionDeltaToolCallChunk":
+ # Check if already cached
+ if "ChatCompletionDeltaToolCallChunk" not in _globals:
+ from litellm.types.llms.openai import (
+ ChatCompletionDeltaToolCallChunk as _ChatCompletionDeltaToolCallChunk,
+ )
+ _globals["ChatCompletionDeltaToolCallChunk"] = _ChatCompletionDeltaToolCallChunk
+ return _globals["ChatCompletionDeltaToolCallChunk"]
+
+ # Lazy load ChatCompletionToolCallChunk to avoid loading at module import time
+ if name == "ChatCompletionToolCallChunk":
+ # Check if already cached
+ if "ChatCompletionToolCallChunk" not in _globals:
+ from litellm.types.llms.openai import (
+ ChatCompletionToolCallChunk as _ChatCompletionToolCallChunk,
+ )
+ _globals["ChatCompletionToolCallChunk"] = _ChatCompletionToolCallChunk
+ return _globals["ChatCompletionToolCallChunk"]
+
+ # Lazy load ChatCompletionToolCallFunctionChunk to avoid loading at module import time
+ if name == "ChatCompletionToolCallFunctionChunk":
+ # Check if already cached
+ if "ChatCompletionToolCallFunctionChunk" not in _globals:
+ from litellm.types.llms.openai import (
+ ChatCompletionToolCallFunctionChunk as _ChatCompletionToolCallFunctionChunk,
+ )
+ _globals["ChatCompletionToolCallFunctionChunk"] = _ChatCompletionToolCallFunctionChunk
+ return _globals["ChatCompletionToolCallFunctionChunk"]
+
+ # Lazy load LiteLLM_Params to avoid loading at module import time
+ if name == "LiteLLM_Params":
+ # Check if already cached
+ if "LiteLLM_Params" not in _globals:
+ from litellm.types.router import LiteLLM_Params as _LiteLLM_Params
+ _globals["LiteLLM_Params"] = _LiteLLM_Params
+ return _globals["LiteLLM_Params"]
+
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index 4651107c5b8..81b4469f24c 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -3663,12 +3663,16 @@
"supports_web_search": true
},
"azure/gpt-image-1": {
- "input_cost_per_pixel": 4.0054321e-08,
+ "cache_read_input_image_token_cost": 2.5e-06,
+ "cache_read_input_token_cost": 1.25e-06,
+ "input_cost_per_image_token": 1e-05,
+ "input_cost_per_token": 5e-06,
"litellm_provider": "azure",
"mode": "image_generation",
- "output_cost_per_pixel": 0.0,
+ "output_cost_per_image_token": 4e-05,
"supported_endpoints": [
- "/v1/images/generations"
+ "/v1/images/generations",
+ "/v1/images/edits"
]
},
"azure/hd/1024-x-1024/dall-e-3": {
@@ -3771,12 +3775,16 @@
]
},
"azure/gpt-image-1-mini": {
- "input_cost_per_pixel": 8.0566406e-09,
+ "cache_read_input_image_token_cost": 2.5e-07,
+ "cache_read_input_token_cost": 2e-07,
+ "input_cost_per_image_token": 2.5e-06,
+ "input_cost_per_token": 2e-06,
"litellm_provider": "azure",
"mode": "image_generation",
- "output_cost_per_pixel": 0.0,
+ "output_cost_per_image_token": 8e-06,
"supported_endpoints": [
- "/v1/images/generations"
+ "/v1/images/generations",
+ "/v1/images/edits"
]
},
"azure/gpt-image-1.5": {
@@ -10943,13 +10951,13 @@
"supports_tool_choice": true
},
"fireworks_ai/accounts/fireworks/models/deepseek-v3p2": {
- "input_cost_per_token": 1.2e-06,
+ "input_cost_per_token": 5.6e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 163840,
"max_output_tokens": 163840,
"max_tokens": 163840,
"mode": "chat",
- "output_cost_per_token": 1.2e-06,
+ "output_cost_per_token": 1.68e-06,
"source": "https://fireworks.ai/models/fireworks/deepseek-v3p2",
"supports_function_calling": true,
"supports_reasoning": true,
@@ -11632,6 +11640,7 @@
"supports_tool_choice": true
},
"gemini-1.5-flash": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_audio_per_second": 2e-06,
"input_cost_per_audio_per_second_above_128k_tokens": 4e-06,
"input_cost_per_character": 1.875e-08,
@@ -11736,6 +11745,7 @@
"supports_vision": true
},
"gemini-1.5-flash-exp-0827": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_audio_per_second": 2e-06,
"input_cost_per_audio_per_second_above_128k_tokens": 4e-06,
"input_cost_per_character": 1.875e-08,
@@ -11770,6 +11780,7 @@
"supports_vision": true
},
"gemini-1.5-flash-preview-0514": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_audio_per_second": 2e-06,
"input_cost_per_audio_per_second_above_128k_tokens": 4e-06,
"input_cost_per_character": 1.875e-08,
@@ -11803,6 +11814,7 @@
"supports_vision": true
},
"gemini-1.5-pro": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_audio_per_second": 3.125e-05,
"input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05,
"input_cost_per_character": 3.125e-07,
@@ -11890,6 +11902,7 @@
"supports_vision": true
},
"gemini-1.5-pro-preview-0215": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_audio_per_second": 3.125e-05,
"input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05,
"input_cost_per_character": 3.125e-07,
@@ -11917,6 +11930,7 @@
"supports_tool_choice": true
},
"gemini-1.5-pro-preview-0409": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_audio_per_second": 3.125e-05,
"input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05,
"input_cost_per_character": 3.125e-07,
@@ -11943,6 +11957,7 @@
"supports_tool_choice": true
},
"gemini-1.5-pro-preview-0514": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_audio_per_second": 3.125e-05,
"input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05,
"input_cost_per_character": 3.125e-07,
@@ -12214,6 +12229,7 @@
"tpm": 250000
},
"gemini-2.0-flash-preview-image-generation": {
+ "deprecation_date": "2025-11-14",
"cache_read_input_token_cost": 2.5e-08,
"input_cost_per_audio_token": 7e-07,
"input_cost_per_token": 1e-07,
@@ -12252,6 +12268,7 @@
"supports_web_search": true
},
"gemini-2.0-flash-thinking-exp": {
+ "deprecation_date": "2025-12-02",
"cache_read_input_token_cost": 0.0,
"input_cost_per_audio_per_second": 0,
"input_cost_per_audio_per_second_above_128k_tokens": 0,
@@ -12300,6 +12317,7 @@
"supports_web_search": true
},
"gemini-2.0-flash-thinking-exp-01-21": {
+ "deprecation_date": "2025-12-02",
"cache_read_input_token_cost": 0.0,
"input_cost_per_audio_per_second": 0,
"input_cost_per_audio_per_second_above_128k_tokens": 0,
@@ -12486,6 +12504,7 @@
"tpm": 8000000
},
"gemini-2.5-flash-image-preview": {
+ "deprecation_date": "2026-01-15",
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
@@ -12796,6 +12815,7 @@
"tpm": 8000000
},
"gemini-2.5-flash-lite-preview-06-17": {
+ "deprecation_date": "2025-11-18",
"cache_read_input_token_cost": 2.5e-08,
"input_cost_per_audio_token": 5e-07,
"input_cost_per_token": 1e-07,
@@ -12885,6 +12905,7 @@
"supports_web_search": true
},
"gemini-2.5-flash-preview-05-20": {
+ "deprecation_date": "2025-11-18",
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
@@ -13156,6 +13177,7 @@
"supports_web_search": true
},
"gemini-2.5-pro-preview-03-25": {
+ "deprecation_date": "2025-12-02",
"cache_read_input_token_cost": 3.125e-07,
"input_cost_per_audio_token": 1.25e-06,
"input_cost_per_token": 1.25e-06,
@@ -13201,6 +13223,7 @@
"supports_web_search": true
},
"gemini-2.5-pro-preview-05-06": {
+ "deprecation_date": "2025-12-02",
"cache_read_input_token_cost": 3.125e-07,
"input_cost_per_audio_token": 1.25e-06,
"input_cost_per_token": 1.25e-06,
@@ -13416,6 +13439,7 @@
"tpm": 10000000
},
"gemini/gemini-1.5-flash": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_token": 7.5e-08,
"input_cost_per_token_above_128k_tokens": 1.5e-07,
"litellm_provider": "gemini",
@@ -13499,6 +13523,7 @@
"tpm": 4000000
},
"gemini/gemini-1.5-flash-8b": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_token": 0,
"input_cost_per_token_above_128k_tokens": 0,
"litellm_provider": "gemini",
@@ -13525,6 +13550,7 @@
"tpm": 4000000
},
"gemini/gemini-1.5-flash-8b-exp-0827": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_token": 0,
"input_cost_per_token_above_128k_tokens": 0,
"litellm_provider": "gemini",
@@ -13550,6 +13576,7 @@
"tpm": 4000000
},
"gemini/gemini-1.5-flash-8b-exp-0924": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_token": 0,
"input_cost_per_token_above_128k_tokens": 0,
"litellm_provider": "gemini",
@@ -13576,6 +13603,7 @@
"tpm": 4000000
},
"gemini/gemini-1.5-flash-exp-0827": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_token": 0,
"input_cost_per_token_above_128k_tokens": 0,
"litellm_provider": "gemini",
@@ -13601,6 +13629,7 @@
"tpm": 4000000
},
"gemini/gemini-1.5-flash-latest": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_token": 7.5e-08,
"input_cost_per_token_above_128k_tokens": 1.5e-07,
"litellm_provider": "gemini",
@@ -13627,6 +13656,7 @@
"tpm": 4000000
},
"gemini/gemini-1.5-pro": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_token": 3.5e-06,
"input_cost_per_token_above_128k_tokens": 7e-06,
"litellm_provider": "gemini",
@@ -13688,6 +13718,7 @@
"tpm": 4000000
},
"gemini/gemini-1.5-pro-exp-0801": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_token": 3.5e-06,
"input_cost_per_token_above_128k_tokens": 7e-06,
"litellm_provider": "gemini",
@@ -13707,6 +13738,7 @@
"tpm": 4000000
},
"gemini/gemini-1.5-pro-exp-0827": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_token": 0,
"input_cost_per_token_above_128k_tokens": 0,
"litellm_provider": "gemini",
@@ -13726,6 +13758,7 @@
"tpm": 4000000
},
"gemini/gemini-1.5-pro-latest": {
+ "deprecation_date": "2025-09-29",
"input_cost_per_token": 3.5e-06,
"input_cost_per_token_above_128k_tokens": 7e-06,
"litellm_provider": "gemini",
@@ -13908,6 +13941,7 @@
"tpm": 4000000
},
"gemini/gemini-2.0-flash-lite-preview-02-05": {
+ "deprecation_date": "2025-12-02",
"cache_read_input_token_cost": 1.875e-08,
"input_cost_per_audio_token": 7.5e-08,
"input_cost_per_token": 7.5e-08,
@@ -13945,6 +13979,7 @@
"tpm": 10000000
},
"gemini/gemini-2.0-flash-live-001": {
+ "deprecation_date": "2025-12-09",
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_audio_token": 2.1e-06,
"input_cost_per_image": 2.1e-06,
@@ -13993,6 +14028,7 @@
"tpm": 250000
},
"gemini/gemini-2.0-flash-preview-image-generation": {
+ "deprecation_date": "2025-11-14",
"cache_read_input_token_cost": 2.5e-08,
"input_cost_per_audio_token": 7e-07,
"input_cost_per_token": 1e-07,
@@ -14032,6 +14068,7 @@
"tpm": 10000000
},
"gemini/gemini-2.0-flash-thinking-exp": {
+ "deprecation_date": "2025-12-02",
"cache_read_input_token_cost": 0.0,
"input_cost_per_audio_per_second": 0,
"input_cost_per_audio_per_second_above_128k_tokens": 0,
@@ -14081,6 +14118,7 @@
"tpm": 4000000
},
"gemini/gemini-2.0-flash-thinking-exp-01-21": {
+ "deprecation_date": "2025-12-02",
"cache_read_input_token_cost": 0.0,
"input_cost_per_audio_per_second": 0,
"input_cost_per_audio_per_second_above_128k_tokens": 0,
@@ -14269,6 +14307,7 @@
"tpm": 8000000
},
"gemini/gemini-2.5-flash-image-preview": {
+ "deprecation_date": "2026-01-15",
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
@@ -14589,6 +14628,7 @@
"tpm": 250000
},
"gemini/gemini-2.5-flash-lite-preview-06-17": {
+ "deprecation_date": "2025-11-18",
"cache_read_input_token_cost": 2.5e-08,
"input_cost_per_audio_token": 5e-07,
"input_cost_per_token": 1e-07,
@@ -14680,6 +14720,7 @@
"tpm": 250000
},
"gemini/gemini-2.5-flash-preview-05-20": {
+ "deprecation_date": "2025-11-18",
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_audio_token": 1e-06,
"input_cost_per_token": 3e-07,
@@ -15026,6 +15067,7 @@
"tpm": 250000
},
"gemini/gemini-2.5-pro-preview-03-25": {
+ "deprecation_date": "2025-12-02",
"cache_read_input_token_cost": 3.125e-07,
"input_cost_per_audio_token": 7e-07,
"input_cost_per_token": 1.25e-06,
@@ -15066,6 +15108,7 @@
"tpm": 10000000
},
"gemini/gemini-2.5-pro-preview-05-06": {
+ "deprecation_date": "2025-12-02",
"cache_read_input_token_cost": 3.125e-07,
"input_cost_per_audio_token": 7e-07,
"input_cost_per_token": 1.25e-06,
@@ -15341,6 +15384,7 @@
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
},
"gemini/imagen-3.0-generate-002": {
+ "deprecation_date": "2025-11-10",
"litellm_provider": "gemini",
"mode": "image_generation",
"output_cost_per_image": 0.04,
@@ -15407,6 +15451,7 @@
]
},
"gemini/veo-3.0-fast-generate-preview": {
+ "deprecation_date": "2025-11-12",
"litellm_provider": "gemini",
"max_input_tokens": 1024,
"max_tokens": 1024,
@@ -15421,6 +15466,7 @@
]
},
"gemini/veo-3.0-generate-preview": {
+ "deprecation_date": "2025-11-12",
"litellm_provider": "gemini",
"max_input_tokens": 1024,
"max_tokens": 1024,
@@ -16980,6 +17026,336 @@
"supports_vision": true,
"supports_pdf_input": true
},
+ "low/1024-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.009,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "low/1024-x-1536/gpt-image-1.5": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "low/1536-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "medium/1024-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.034,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "medium/1024-x-1536/gpt-image-1.5": {
+ "input_cost_per_image": 0.05,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "medium/1536-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.05,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "high/1024-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.133,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "high/1024-x-1536/gpt-image-1.5": {
+ "input_cost_per_image": 0.20,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "high/1536-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.20,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "standard/1024-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.009,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "standard/1024-x-1536/gpt-image-1.5": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "standard/1536-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "1024-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.009,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "1024-x-1536/gpt-image-1.5": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "1536-x-1024/gpt-image-1.5": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "low/1024-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.009,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "low/1024-x-1536/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "low/1536-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "medium/1024-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.034,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "medium/1024-x-1536/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.05,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "medium/1536-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.05,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "high/1024-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.133,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "high/1024-x-1536/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.20,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "high/1536-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.20,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "standard/1024-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.009,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "standard/1024-x-1536/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "standard/1536-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "1024-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.009,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "1024-x-1536/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
+ "1536-x-1024/gpt-image-1.5-2025-12-16": {
+ "input_cost_per_image": 0.013,
+ "litellm_provider": "openai",
+ "mode": "image_generation",
+ "supported_endpoints": [
+ "/v1/images/generations",
+ "/v1/images/edits"
+ ],
+ "supports_vision": true,
+ "supports_pdf_input": true
+ },
"gpt-5": {
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_flex": 6.25e-08,
@@ -17741,16 +18117,16 @@
"supports_vision": true
},
"gpt-image-1": {
- "input_cost_per_image": 0.042,
- "input_cost_per_pixel": 4.0054321e-08,
- "input_cost_per_token": 0.000005,
- "input_cost_per_image_token": 0.00001,
+ "cache_read_input_image_token_cost": 2.5e-06,
+ "cache_read_input_token_cost": 1.25e-06,
+ "input_cost_per_image_token": 1e-05,
+ "input_cost_per_token": 5e-06,
"litellm_provider": "openai",
"mode": "image_generation",
- "output_cost_per_pixel": 0.0,
- "output_cost_per_token": 0.00004,
+ "output_cost_per_image_token": 4e-05,
"supported_endpoints": [
- "/v1/images/generations"
+ "/v1/images/generations",
+ "/v1/images/edits"
]
},
"gpt-image-1-mini": {
@@ -18175,6 +18551,18 @@
"supports_response_schema": false,
"supports_tool_choice": true
},
+ "groq/gemma-7b-it": {
+ "input_cost_per_token": 5e-08,
+ "litellm_provider": "groq",
+ "max_input_tokens": 8192,
+ "max_output_tokens": 8192,
+ "max_tokens": 8192,
+ "mode": "chat",
+ "output_cost_per_token": 8e-08,
+ "supports_function_calling": true,
+ "supports_response_schema": false,
+ "supports_tool_choice": true
+ },
"groq/meta-llama/llama-guard-4-12b": {
"input_cost_per_token": 2e-07,
"litellm_provider": "groq",
@@ -24776,6 +25164,7 @@
"source": "https://docs.mistral.ai/capabilities/code_generation/"
},
"text-embedding-004": {
+ "deprecation_date": "2026-01-14",
"input_cost_per_character": 2.5e-08,
"input_cost_per_token": 1e-07,
"litellm_provider": "vertex_ai-embedding-models",
@@ -27546,6 +27935,7 @@
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
},
"vertex_ai/imagen-3.0-generate-002": {
+ "deprecation_date": "2025-11-10",
"litellm_provider": "vertex_ai-image-models",
"mode": "image_generation",
"output_cost_per_image": 0.04,
@@ -28056,6 +28446,7 @@
]
},
"vertex_ai/veo-3.0-fast-generate-preview": {
+ "deprecation_date": "2025-11-12",
"litellm_provider": "vertex_ai-video-models",
"max_input_tokens": 1024,
"max_tokens": 1024,
@@ -28070,6 +28461,7 @@
]
},
"vertex_ai/veo-3.0-generate-preview": {
+ "deprecation_date": "2025-11-12",
"litellm_provider": "vertex_ai-video-models",
"max_input_tokens": 1024,
"max_tokens": 1024,
@@ -29299,6 +29691,20 @@
"supports_vision": true,
"supports_web_search": true
},
+ "zai/glm-4.7": {
+ "cache_creation_input_token_cost": 0,
+ "cache_read_input_token_cost": 1.1e-07,
+ "input_cost_per_token": 6e-07,
+ "output_cost_per_token": 2.2e-06,
+ "litellm_provider": "zai",
+ "max_input_tokens": 200000,
+ "max_output_tokens": 128000,
+ "mode": "chat",
+ "supports_function_calling": true,
+ "supports_reasoning": true,
+ "supports_tool_choice": true,
+ "source": "https://docs.z.ai/guides/overview/pricing"
+ },
"zai/glm-4.6": {
"input_cost_per_token": 6e-07,
"output_cost_per_token": 2.2e-06,
diff --git a/poetry.lock b/poetry.lock
index ee97c00594c..a0a0f8540e5 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -2273,75 +2273,6 @@ googleapis-common-protos = {version = ">=1.56.0,<2.0.0", extras = ["grpc"]}
grpcio = ">=1.44.0,<2.0.0"
protobuf = ">=3.20.2,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<7.0.0"
-[[package]]
-name = "grpcio"
-version = "1.67.1"
-description = "HTTP/2-based RPC framework"
-optional = false
-python-versions = ">=3.8"
-groups = ["main", "dev", "proxy-dev"]
-markers = "python_version < \"3.14\""
-files = [
- {file = "grpcio-1.67.1-cp310-cp310-linux_armv7l.whl", hash = "sha256:8b0341d66a57f8a3119b77ab32207072be60c9bf79760fa609c5609f2deb1f3f"},
- {file = "grpcio-1.67.1-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:f5a27dddefe0e2357d3e617b9079b4bfdc91341a91565111a21ed6ebbc51b22d"},
- {file = "grpcio-1.67.1-cp310-cp310-manylinux_2_17_aarch64.whl", hash = "sha256:43112046864317498a33bdc4797ae6a268c36345a910de9b9c17159d8346602f"},
- {file = "grpcio-1.67.1-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c9b929f13677b10f63124c1a410994a401cdd85214ad83ab67cc077fc7e480f0"},
- {file = "grpcio-1.67.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e7d1797a8a3845437d327145959a2c0c47c05947c9eef5ff1a4c80e499dcc6fa"},
- {file = "grpcio-1.67.1-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:0489063974d1452436139501bf6b180f63d4977223ee87488fe36858c5725292"},
- {file = "grpcio-1.67.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:9fd042de4a82e3e7aca44008ee2fb5da01b3e5adb316348c21980f7f58adc311"},
- {file = "grpcio-1.67.1-cp310-cp310-win32.whl", hash = "sha256:638354e698fd0c6c76b04540a850bf1db27b4d2515a19fcd5cf645c48d3eb1ed"},
- {file = "grpcio-1.67.1-cp310-cp310-win_amd64.whl", hash = "sha256:608d87d1bdabf9e2868b12338cd38a79969eaf920c89d698ead08f48de9c0f9e"},
- {file = "grpcio-1.67.1-cp311-cp311-linux_armv7l.whl", hash = "sha256:7818c0454027ae3384235a65210bbf5464bd715450e30a3d40385453a85a70cb"},
- {file = "grpcio-1.67.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:ea33986b70f83844cd00814cee4451055cd8cab36f00ac64a31f5bb09b31919e"},
- {file = "grpcio-1.67.1-cp311-cp311-manylinux_2_17_aarch64.whl", hash = "sha256:c7a01337407dd89005527623a4a72c5c8e2894d22bead0895306b23c6695698f"},
- {file = "grpcio-1.67.1-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:80b866f73224b0634f4312a4674c1be21b2b4afa73cb20953cbbb73a6b36c3cc"},
- {file = "grpcio-1.67.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f9fff78ba10d4250bfc07a01bd6254a6d87dc67f9627adece85c0b2ed754fa96"},
- {file = "grpcio-1.67.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:8a23cbcc5bb11ea7dc6163078be36c065db68d915c24f5faa4f872c573bb400f"},
- {file = "grpcio-1.67.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:1a65b503d008f066e994f34f456e0647e5ceb34cfcec5ad180b1b44020ad4970"},
- {file = "grpcio-1.67.1-cp311-cp311-win32.whl", hash = "sha256:e29ca27bec8e163dca0c98084040edec3bc49afd10f18b412f483cc68c712744"},
- {file = "grpcio-1.67.1-cp311-cp311-win_amd64.whl", hash = "sha256:786a5b18544622bfb1e25cc08402bd44ea83edfb04b93798d85dca4d1a0b5be5"},
- {file = "grpcio-1.67.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:267d1745894200e4c604958da5f856da6293f063327cb049a51fe67348e4f953"},
- {file = "grpcio-1.67.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:85f69fdc1d28ce7cff8de3f9c67db2b0ca9ba4449644488c1e0303c146135ddb"},
- {file = "grpcio-1.67.1-cp312-cp312-manylinux_2_17_aarch64.whl", hash = "sha256:f26b0b547eb8d00e195274cdfc63ce64c8fc2d3e2d00b12bf468ece41a0423a0"},
- {file = "grpcio-1.67.1-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4422581cdc628f77302270ff839a44f4c24fdc57887dc2a45b7e53d8fc2376af"},
- {file = "grpcio-1.67.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1d7616d2ded471231c701489190379e0c311ee0a6c756f3c03e6a62b95a7146e"},
- {file = "grpcio-1.67.1-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:8a00efecde9d6fcc3ab00c13f816313c040a28450e5e25739c24f432fc6d3c75"},
- {file = "grpcio-1.67.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:699e964923b70f3101393710793289e42845791ea07565654ada0969522d0a38"},
- {file = "grpcio-1.67.1-cp312-cp312-win32.whl", hash = "sha256:4e7b904484a634a0fff132958dabdb10d63e0927398273917da3ee103e8d1f78"},
- {file = "grpcio-1.67.1-cp312-cp312-win_amd64.whl", hash = "sha256:5721e66a594a6c4204458004852719b38f3d5522082be9061d6510b455c90afc"},
- {file = "grpcio-1.67.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:aa0162e56fd10a5547fac8774c4899fc3e18c1aa4a4759d0ce2cd00d3696ea6b"},
- {file = "grpcio-1.67.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:beee96c8c0b1a75d556fe57b92b58b4347c77a65781ee2ac749d550f2a365dc1"},
- {file = "grpcio-1.67.1-cp313-cp313-manylinux_2_17_aarch64.whl", hash = "sha256:a93deda571a1bf94ec1f6fcda2872dad3ae538700d94dc283c672a3b508ba3af"},
- {file = "grpcio-1.67.1-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0e6f255980afef598a9e64a24efce87b625e3e3c80a45162d111a461a9f92955"},
- {file = "grpcio-1.67.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9e838cad2176ebd5d4a8bb03955138d6589ce9e2ce5d51c3ada34396dbd2dba8"},
- {file = "grpcio-1.67.1-cp313-cp313-musllinux_1_1_i686.whl", hash = "sha256:a6703916c43b1d468d0756c8077b12017a9fcb6a1ef13faf49e67d20d7ebda62"},
- {file = "grpcio-1.67.1-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:917e8d8994eed1d86b907ba2a61b9f0aef27a2155bca6cbb322430fc7135b7bb"},
- {file = "grpcio-1.67.1-cp313-cp313-win32.whl", hash = "sha256:e279330bef1744040db8fc432becc8a727b84f456ab62b744d3fdb83f327e121"},
- {file = "grpcio-1.67.1-cp313-cp313-win_amd64.whl", hash = "sha256:fa0c739ad8b1996bd24823950e3cb5152ae91fca1c09cc791190bf1627ffefba"},
- {file = "grpcio-1.67.1-cp38-cp38-linux_armv7l.whl", hash = "sha256:178f5db771c4f9a9facb2ab37a434c46cb9be1a75e820f187ee3d1e7805c4f65"},
- {file = "grpcio-1.67.1-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:0f3e49c738396e93b7ba9016e153eb09e0778e776df6090c1b8c91877cc1c426"},
- {file = "grpcio-1.67.1-cp38-cp38-manylinux_2_17_aarch64.whl", hash = "sha256:24e8a26dbfc5274d7474c27759b54486b8de23c709d76695237515bc8b5baeab"},
- {file = "grpcio-1.67.1-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3b6c16489326d79ead41689c4b84bc40d522c9a7617219f4ad94bc7f448c5085"},
- {file = "grpcio-1.67.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:60e6a4dcf5af7bbc36fd9f81c9f372e8ae580870a9e4b6eafe948cd334b81cf3"},
- {file = "grpcio-1.67.1-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:95b5f2b857856ed78d72da93cd7d09b6db8ef30102e5e7fe0961fe4d9f7d48e8"},
- {file = "grpcio-1.67.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:b49359977c6ec9f5d0573ea4e0071ad278ef905aa74e420acc73fd28ce39e9ce"},
- {file = "grpcio-1.67.1-cp38-cp38-win32.whl", hash = "sha256:f5b76ff64aaac53fede0cc93abf57894ab2a7362986ba22243d06218b93efe46"},
- {file = "grpcio-1.67.1-cp38-cp38-win_amd64.whl", hash = "sha256:804c6457c3cd3ec04fe6006c739579b8d35c86ae3298ffca8de57b493524b771"},
- {file = "grpcio-1.67.1-cp39-cp39-linux_armv7l.whl", hash = "sha256:a25bdea92b13ff4d7790962190bf6bf5c4639876e01c0f3dda70fc2769616335"},
- {file = "grpcio-1.67.1-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:cdc491ae35a13535fd9196acb5afe1af37c8237df2e54427be3eecda3653127e"},
- {file = "grpcio-1.67.1-cp39-cp39-manylinux_2_17_aarch64.whl", hash = "sha256:85f862069b86a305497e74d0dc43c02de3d1d184fc2c180993aa8aa86fbd19b8"},
- {file = "grpcio-1.67.1-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ec74ef02010186185de82cc594058a3ccd8d86821842bbac9873fd4a2cf8be8d"},
- {file = "grpcio-1.67.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:01f616a964e540638af5130469451cf580ba8c7329f45ca998ab66e0c7dcdb04"},
- {file = "grpcio-1.67.1-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:299b3d8c4f790c6bcca485f9963b4846dd92cf6f1b65d3697145d005c80f9fe8"},
- {file = "grpcio-1.67.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:60336bff760fbb47d7e86165408126f1dded184448e9a4c892189eb7c9d3f90f"},
- {file = "grpcio-1.67.1-cp39-cp39-win32.whl", hash = "sha256:5ed601c4c6008429e3d247ddb367fe8c7259c355757448d7c1ef7bd4a6739e8e"},
- {file = "grpcio-1.67.1-cp39-cp39-win_amd64.whl", hash = "sha256:5db70d32d6703b89912af16d6d45d78406374a8b8ef0d28140351dd0ec610e98"},
- {file = "grpcio-1.67.1.tar.gz", hash = "sha256:3dc2ed4cabea4dc14d5e708c2b426205956077cc5de419b4d4079315017e9732"},
-]
-
-[package.extras]
-protobuf = ["grpcio-tools (>=1.67.1)"]
-
[[package]]
name = "grpcio"
version = "1.76.0"
@@ -2349,7 +2280,6 @@ description = "HTTP/2-based RPC framework"
optional = false
python-versions = ">=3.9"
groups = ["main", "dev", "proxy-dev"]
-markers = "python_version >= \"3.14\""
files = [
{file = "grpcio-1.76.0-cp310-cp310-linux_armv7l.whl", hash = "sha256:65a20de41e85648e00305c1bb09a3598f840422e522277641145a32d42dcefcc"},
{file = "grpcio-1.76.0-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:40ad3afe81676fd9ec6d9d406eda00933f218038433980aa19d401490e46ecde"},
@@ -8051,4 +7981,4 @@ utils = ["numpydoc"]
[metadata]
lock-version = "2.1"
python-versions = ">=3.9,<4.0"
-content-hash = "b010d9da7f5a765670932b78d720aae4fcb819daba050683ee125b4367972419"
+content-hash = "7eed2b2c25173a275ac83c55fd901b9b84663b1d7daa54f0e78b30bf1c8f0e3e"
diff --git a/pyproject.toml b/pyproject.toml
index f929fb94cb0..3b09119a748 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -72,7 +72,7 @@ soundfile = {version = "^0.12.1", optional = true}
# - 1.68.0-1.68.1 has reconnect bug (https://github.com/grpc/grpc/issues/38290)
# - 1.75.0+ has Python 3.14 wheels and bug fix
grpcio = [
- {version = ">=1.62.3,<1.68.0", python = "<3.14"},
+ {version = ">=1.62.3,!=1.68.*,!=1.69.*,!=1.70.*,!=1.71.0,!=1.71.1,!=1.72.0,!=1.72.1,!=1.73.0", python = "<3.14"},
{version = ">=1.75.0", python = ">=3.14"},
]
diff --git a/requirements.txt b/requirements.txt
index 3bc968c8cb8..06a7c17336c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -41,7 +41,7 @@ opentelemetry-api==1.25.0
opentelemetry-sdk==1.25.0
opentelemetry-exporter-otlp==1.25.0
# grpcio: 1.68.0-1.68.1 has reconnect bug (#38290), 1.75+ has Python 3.14 wheels + fix
-grpcio>=1.62.3,<1.68.0; python_version < "3.14"
+grpcio>=1.62.3,!=1.68.*,!=1.69.*,!=1.70.*,!=1.71.0,!=1.71.1,!=1.72.0,!=1.72.1,!=1.73.0; python_version < "3.14"
grpcio>=1.75.0; python_version >= "3.14"
sentry_sdk==2.21.0 # for sentry error handling
detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests
diff --git a/schema.prisma b/schema.prisma
index aac0b5b35de..ea47b6ed03b 100644
--- a/schema.prisma
+++ b/schema.prisma
@@ -208,6 +208,9 @@ model LiteLLM_MCPServerTable {
command String?
args String[] @default([])
env Json? @default("{}")
+ authorization_url String?
+ token_url String?
+ registration_url String?
}
// Generate Tokens for Proxy
diff --git a/tests/code_coverage_tests/memory_test.py b/tests/code_coverage_tests/memory_test.py
new file mode 100644
index 00000000000..1ce93191992
--- /dev/null
+++ b/tests/code_coverage_tests/memory_test.py
@@ -0,0 +1,616 @@
+"""
+Memory Violation Detection Test
+
+Detects bad memory patterns in the LiteLLM codebase that can lead to memory leaks or OOMs.
+
+The detector uses a modular pattern-based system. To add detection for new memory patterns:
+
+1. Create a Pattern subclass implementing get_pattern_name(), visit_assign(), and check_cleanup()
+ - You can extend the Pattern class with additional methods as needed for your detection logic
+2. Add the pattern to MemoryViolationDetector.DEFAULT_PATTERNS
+
+Currently detects:
+- queue.get() / queue.get_nowait() operations where variables aren't set to None
+- Class-level data structures that have add operations during runtime without size limits:
+ * Built-in: list, dict, set
+ * Collections: deque, defaultdict, Counter, OrderedDict, ChainMap
+ * Queues: queue.Queue, asyncio.Queue (if unbounded, i.e., no maxsize parameter)
+ * Heap operations: heapq.heappush(), heapq.heapreplace(), heapq.heappushpop() on class-level lists
+"""
+
+import ast
+import os
+from abc import ABC, abstractmethod
+from typing import List, Dict, Any, Optional, Sequence
+
+
+class Pattern(ABC):
+ """Base class for memory violation detection patterns"""
+
+ @abstractmethod
+ def get_pattern_name(self) -> str:
+ """Return unique identifier for this violation type"""
+ pass
+
+ @abstractmethod
+ def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]:
+ """Detect memory-sensitive operations in assignment. Returns list of {line, var_name, call} dicts."""
+ pass
+
+ @abstractmethod
+ def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt],
+ context: Dict[str, Any]) -> List[Dict[str, Any]]:
+ """Verify variables are set to None. Returns list of violation dicts."""
+ pass
+
+
+class QueueGetPattern(Pattern):
+ """Detects queue.get()/get_nowait() operations that aren't cleared"""
+
+ def get_pattern_name(self) -> str:
+ return "queue_reference_not_cleared"
+
+ def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]:
+ """Detect queue.get() or queue.get_nowait() calls where object name contains 'queue'"""
+ operations = []
+
+ if isinstance(node.value, ast.Call):
+ func = node.value.func
+ if isinstance(func, ast.Attribute) and func.attr in ("get", "get_nowait"):
+ obj_name = context["get_attr_string"](func.value)
+ if "queue" in obj_name.lower() and node.targets and isinstance(node.targets[0], ast.Name):
+ operations.append({
+ "line": node.lineno,
+ "var_name": node.targets[0].id,
+ "call": context["get_call_string"](node.value),
+ })
+
+ return operations
+
+ def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt],
+ context: Dict[str, Any]) -> List[Dict[str, Any]]:
+ """Flag queue variables that aren't set to None"""
+ violations = []
+ is_var_set_to_none = context["is_var_set_to_none"]
+ current_function = context["current_function"]
+ file_path = context["file_path"]
+
+ queue_vars = {op["var_name"]: op["line"] for op in operations}
+
+ for var_name, line_num in queue_vars.items():
+ if not is_var_set_to_none(var_name, function_body):
+ violations.append({
+ "line": line_num,
+ "type": self.get_pattern_name(),
+ "var_name": var_name,
+ "function": current_function,
+ "file_path": file_path,
+ "message": (
+ f"Queue variable '{var_name}' in function "
+ f"'{current_function}' is not set to None after use. "
+ f"If the runtime is overwhelmed, this can cause OOM (Out of Memory) errors."
+ ),
+ })
+
+ return violations
+
+
+class UnboundedDataStructurePattern(Pattern):
+ """Detects class-level data structures (lists, dicts, sets) that can grow unbounded"""
+
+ def get_pattern_name(self) -> str:
+ return "unbounded_data_structure"
+
+ def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]:
+ """Detect list/dict/set creations that are at class level"""
+ operations = []
+
+ # Check if this is a data structure creation
+ is_data_structure = False
+ structure_type = None
+
+ if isinstance(node.value, (ast.List, ast.Dict, ast.Set)):
+ is_data_structure = True
+ if isinstance(node.value, ast.List):
+ structure_type = "list"
+ elif isinstance(node.value, ast.Dict):
+ structure_type = "dict"
+ elif isinstance(node.value, ast.Set):
+ structure_type = "set"
+ elif isinstance(node.value, ast.Call):
+ # Check for list(), dict(), set() calls
+ func = node.value.func
+ if isinstance(func, ast.Name):
+ if func.id in ("list", "dict", "set"):
+ is_data_structure = True
+ structure_type = func.id
+ elif isinstance(func, ast.Attribute):
+ # Handle cases like collections.defaultdict(list), collections.deque(), etc.
+ obj_name = context["get_attr_string"](func.value)
+ attr_name = func.attr
+
+ # Check for collections module data structures
+ if "collections" in obj_name.lower() or "collections" in str(func.value):
+ if attr_name in ("deque", "defaultdict", "Counter", "OrderedDict", "ChainMap"):
+ # For deque, we track it and let size checks determine if it's bounded
+ # (deque with maxlen parameter is bounded, but we detect that via size checks)
+ is_data_structure = True
+ structure_type = attr_name
+ elif attr_name in ("list", "dict", "set"):
+ # collections.defaultdict(list) pattern
+ is_data_structure = True
+ structure_type = "defaultdict" if "defaultdict" in obj_name.lower() else attr_name
+ # Check for queue.Queue, asyncio.Queue (if unbounded)
+ elif "queue" in obj_name.lower() or "asyncio" in obj_name.lower():
+ if attr_name == "Queue":
+ # Check if maxsize is set (bounded queue)
+ has_maxsize = False
+ for keyword in node.value.keywords:
+ if keyword.arg == "maxsize":
+ has_maxsize = True
+ break
+ if not has_maxsize:
+ is_data_structure = True
+ structure_type = "queue"
+ # Direct attribute access like deque(), Counter(), etc.
+ elif attr_name in ("deque", "defaultdict", "Counter", "OrderedDict", "ChainMap"):
+ is_data_structure = True
+ structure_type = attr_name
+
+ if is_data_structure and node.targets and isinstance(node.targets[0], ast.Name):
+ scope = context.get("current_scope", "function")
+ # Only track if it's at class level (not module level)
+ if scope == "class":
+ operations.append({
+ "line": node.lineno,
+ "var_name": node.targets[0].id,
+ "structure_type": structure_type,
+ "scope": scope,
+ "call": context["get_call_string"](node.value) if isinstance(node.value, ast.Call) else f"{structure_type}()",
+ })
+
+ return operations
+
+ def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt],
+ context: Dict[str, Any]) -> List[Dict[str, Any]]:
+ """Flag persistent data structures that have add operations without size limits"""
+ violations = []
+ current_function = context["current_function"]
+ current_scope = context.get("current_scope", "function")
+ file_path = context["file_path"]
+ get_attr_string = context["get_attr_string"]
+
+ # Skip if this is initialization code (module-level, class-level, or __init__ methods)
+ # Only flag operations in regular methods/functions that can be called during runtime
+ is_initialization = (
+ current_scope in ("module", "class") or
+ current_function in ("__init__", "__new__", "__class_init__") or
+ current_function is None # Module-level code
+ )
+
+ if is_initialization:
+ return violations # Don't flag initialization code
+
+ # Track which variables have add operations and size checks
+ var_add_operations = {} # var_name -> list of lines with add operations
+ var_size_checks = {} # var_name -> has size limit check
+
+ # Build a set of variable names to check
+ tracked_vars = {op["var_name"]: op for op in operations}
+
+ # Scan body for operations on these variables
+ for stmt in function_body:
+ for node in ast.walk(stmt):
+ # Check for method calls that add items
+ if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
+ attr_name = node.func.attr
+ obj_name = get_attr_string(node.func.value)
+
+ # Check if this is an add operation on one of our tracked variables
+ for var_name, op in tracked_vars.items():
+ structure_type = op["structure_type"]
+
+ # Match variable name (exact or as attribute)
+ if obj_name == var_name or obj_name.endswith(f".{var_name}") or obj_name.endswith(f"['{var_name}']"):
+ # Check for add operations
+ add_ops = {
+ "list": ["append", "extend", "insert"],
+ "dict": ["update", "setdefault"],
+ "set": ["add", "update"],
+ "deque": ["append", "appendleft", "extend", "extendleft", "insert"],
+ "defaultdict": ["update", "setdefault"],
+ "Counter": ["update"],
+ "OrderedDict": ["update", "setdefault"],
+ "ChainMap": ["new_child"],
+ "queue": ["put", "put_nowait"],
+ }
+
+ if attr_name in add_ops.get(structure_type, []):
+ if var_name not in var_add_operations:
+ var_add_operations[var_name] = []
+ var_add_operations[var_name].append(node.lineno)
+
+ # Check for size limit checks (len() calls, maxsize/maxlen attributes)
+ if (attr_name in ("__len__",) or
+ "maxsize" in attr_name.lower() or
+ "max_size" in attr_name.lower() or
+ attr_name == "maxlen"): # For deque
+ var_size_checks[var_name] = True
+
+ # Check for heapq operations on tracked lists (heapq.heappush, heapq.heappop)
+ if isinstance(node, ast.Call):
+ func = node.func
+ # Check for heapq.heappush(list_var, item) or heapq.heappop(list_var)
+ if isinstance(func, ast.Attribute):
+ func_obj = get_attr_string(func.value)
+ func_name = func.attr
+ # Check if it's a heapq operation
+ if func_obj == "heapq" and func_name in ("heappush", "heapreplace", "heappushpop"):
+ # First argument should be our tracked variable
+ if len(node.args) > 0:
+ arg_name = get_attr_string(node.args[0])
+ for var_name, op in tracked_vars.items():
+ if op["structure_type"] == "list" and (
+ arg_name == var_name or arg_name.endswith(f".{var_name}")
+ ):
+ if var_name not in var_add_operations:
+ var_add_operations[var_name] = []
+ var_add_operations[var_name].append(node.lineno)
+
+ # Check for dict item assignment: dict[key] = value
+ if isinstance(node, ast.Assign):
+ for target in node.targets:
+ if isinstance(target, ast.Subscript):
+ target_name = get_attr_string(target.value)
+ for var_name in tracked_vars:
+ if target_name == var_name or target_name.endswith(f".{var_name}"):
+ if var_name not in var_add_operations:
+ var_add_operations[var_name] = []
+ var_add_operations[var_name].append(node.lineno)
+
+ # Check for augmented assignment: list += [...]
+ if isinstance(node, ast.AugAssign):
+ target_name = get_attr_string(node.target)
+ for var_name in tracked_vars:
+ if target_name == var_name or target_name.endswith(f".{var_name}"):
+ if var_name not in var_add_operations:
+ var_add_operations[var_name] = []
+ var_add_operations[var_name].append(node.lineno)
+
+ # Check for size comparisons in conditionals
+ if isinstance(node, (ast.If, ast.While, ast.Assert)):
+ test = getattr(node, "test", None)
+ if test:
+ for comp_node in ast.walk(test):
+ if isinstance(comp_node, ast.Compare):
+ left_str = get_attr_string(comp_node.left) if hasattr(comp_node, "left") else ""
+ # Check for len() calls
+ if isinstance(comp_node.left, ast.Call):
+ call_func = comp_node.left.func
+ if isinstance(call_func, ast.Name) and call_func.id == "len":
+ if len(comp_node.left.args) > 0:
+ arg_name = get_attr_string(comp_node.left.args[0])
+ for var_name in tracked_vars:
+ if arg_name == var_name or arg_name.endswith(f".{var_name}"):
+ # Check if comparing to a limit
+ for comparator in comp_node.comparators:
+ if isinstance(comparator, ast.Constant):
+ var_size_checks[var_name] = True
+ elif isinstance(comparator, ast.Name):
+ # Could be a constant like MAX_SIZE
+ if "max" in comparator.id.lower() or "limit" in comparator.id.lower():
+ var_size_checks[var_name] = True
+ # Handle deprecated ast.Num for Python < 3.8
+ try:
+ Num = getattr(ast, "Num", None)
+ if Num and isinstance(comparator, Num):
+ var_size_checks[var_name] = True
+ except (AttributeError, TypeError):
+ pass
+ # Check for direct variable comparisons
+ for var_name in tracked_vars:
+ if var_name in left_str:
+ for comparator in comp_node.comparators:
+ if isinstance(comparator, ast.Constant):
+ var_size_checks[var_name] = True
+ # Handle deprecated ast.Num for Python < 3.8
+ try:
+ Num = getattr(ast, "Num", None)
+ if Num and isinstance(comparator, Num):
+ var_size_checks[var_name] = True
+ except (AttributeError, TypeError):
+ pass
+
+ # Flag violations: persistent structures with add operations but no size checks
+ for op in operations:
+ var_name = op["var_name"]
+ structure_type = op["structure_type"]
+
+ if var_name in var_add_operations and var_name not in var_size_checks:
+ violations.append({
+ "line": op["line"],
+ "type": self.get_pattern_name(),
+ "var_name": var_name,
+ "function": current_function or "class-level",
+ "file_path": file_path,
+ "message": (
+ f"Class-level {structure_type} '{var_name}' "
+ f"has add operations (lines {var_add_operations[var_name]}) but no size limit checks. "
+ f"This can lead to unbounded memory growth and OOM errors during runtime."
+ ),
+ })
+
+ return violations
+
+
+class MemoryViolationDetector(ast.NodeVisitor):
+ """AST visitor that detects memory violations using registered patterns"""
+
+ DEFAULT_PATTERNS: List[Pattern] = [QueueGetPattern(), UnboundedDataStructurePattern()]
+
+ def __init__(self, file_path: str, patterns: Optional[Sequence[Pattern]] = None):
+ self.file_path = file_path
+ self.violations: List[Dict[str, Any]] = []
+ self.current_function: Optional[str] = None
+ self.current_scope: str = "module" # Track current scope: module, class, function
+ self.patterns = self.DEFAULT_PATTERNS if patterns is None else patterns
+ self.ast_tree: Optional[ast.Module] = None # Store full AST for module-level checks
+
+ self.pattern_operations: Dict[str, List[Dict[str, Any]]] = {
+ pattern.get_pattern_name(): [] for pattern in self.patterns
+ }
+
+ # Track class-level operations separately (for checking in functions)
+ self.class_level_operations: Dict[str, List[Dict[str, Any]]] = {
+ pattern.get_pattern_name(): [] for pattern in self.patterns
+ }
+
+ self._context = {
+ "get_call_string": self._get_call_string,
+ "get_attr_string": self._get_attr_string,
+ "is_var_set_to_none": self._is_var_set_to_none,
+ "current_function": None,
+ "current_scope": "module",
+ "file_path": file_path,
+ }
+
+ def visit_ClassDef(self, node):
+ """Track class scope"""
+ old_scope = self.current_scope
+ self.current_scope = "class"
+ self._context["current_scope"] = "class"
+
+ self.generic_visit(node)
+
+ self.current_scope = old_scope
+ self._context["current_scope"] = old_scope
+
+ def visit_FunctionDef(self, node):
+ """Track function scope and check cleanup after visiting"""
+ old_function = self.current_function
+ old_scope = self.current_scope
+ self.current_function = node.name
+ self.current_scope = "function"
+ self._context["current_function"] = node.name
+ self._context["current_scope"] = "function"
+
+ for pattern_name in self.pattern_operations:
+ self.pattern_operations[pattern_name] = []
+
+ self.generic_visit(node)
+ self._check_function_cleanup(node)
+
+ self.current_function = old_function
+ self.current_scope = old_scope
+ self._context["current_function"] = old_function
+ self._context["current_scope"] = old_scope
+
+ def visit_AsyncFunctionDef(self, node):
+ """Track async function scope and check cleanup after visiting"""
+ old_function = self.current_function
+ old_scope = self.current_scope
+ self.current_function = node.name
+ self.current_scope = "function"
+ self._context["current_function"] = node.name
+ self._context["current_scope"] = "function"
+
+ for pattern_name in self.pattern_operations:
+ self.pattern_operations[pattern_name] = []
+
+ self.generic_visit(node)
+ self._check_function_cleanup(node)
+
+ self.current_function = old_function
+ self.current_scope = old_scope
+ self._context["current_function"] = old_function
+ self._context["current_scope"] = old_scope
+
+ def visit_Assign(self, node):
+ """Detect memory-sensitive operations in assignments"""
+ for pattern in self.patterns:
+ operations = pattern.visit_assign(node, self._context)
+ # Track function-level operations
+ self.pattern_operations[pattern.get_pattern_name()].extend(operations)
+ # Track class-level operations separately (for checking in functions)
+ for op in operations:
+ if op.get("scope") == "class":
+ self.class_level_operations[pattern.get_pattern_name()].append(op)
+
+ self.generic_visit(node)
+
+ def _check_function_cleanup(self, node):
+ """Check cleanup for all detected operations"""
+ for pattern in self.patterns:
+ operations = self.pattern_operations[pattern.get_pattern_name()]
+ if operations:
+ violations = pattern.check_cleanup(operations, node.body, self._context)
+ self.violations.extend(violations)
+
+ # For UnboundedDataStructurePattern, also check if this function modifies class-level structures
+ if isinstance(pattern, UnboundedDataStructurePattern):
+ class_ops = self.class_level_operations[pattern.get_pattern_name()]
+ if class_ops and self.current_function not in ("__init__", "__new__", "__class_init__", None):
+ # Check if this regular function modifies class-level structures
+ violations = pattern.check_cleanup(class_ops, node.body, self._context)
+ self.violations.extend(violations)
+
+ def _check_module_level_cleanup(self):
+ """Check cleanup for module/class level operations"""
+ # Module-level operations are now checked when visiting functions
+ # This method is kept for potential future use but doesn't need to do anything
+ # since we only want to flag runtime modifications in functions, not initialization code
+ pass
+
+ def _is_var_set_to_none(self, var_name: str, body: List[ast.stmt]) -> bool:
+ """Check if variable is set to None after its initial assignment"""
+ assignment_line = None
+ for stmt in body:
+ for node in ast.walk(stmt):
+ if isinstance(node, ast.Assign):
+ for target in node.targets:
+ if isinstance(target, ast.Name) and target.id == var_name:
+ assignment_line = node.lineno
+ break
+ if assignment_line:
+ break
+ if assignment_line:
+ break
+
+ if not assignment_line:
+ return False
+
+ for stmt in body:
+ for node in ast.walk(stmt):
+ if isinstance(node, ast.Assign):
+ for target in node.targets:
+ if isinstance(target, ast.Name) and target.id == var_name and node.lineno > assignment_line:
+ if isinstance(node.value, ast.Constant) and node.value.value is None:
+ return True
+ try:
+ NameConstant = getattr(ast, "NameConstant", None)
+ if NameConstant and isinstance(node.value, NameConstant):
+ if getattr(node.value, "value", None) is None:
+ return True
+ except (AttributeError, TypeError):
+ pass
+ return False
+
+ def _get_call_string(self, node: ast.Call) -> str:
+ """Get string representation of function call"""
+ try:
+ if hasattr(ast, "unparse"):
+ return ast.unparse(node)
+ elif isinstance(node.func, ast.Attribute):
+ return f"{self._get_attr_string(node.func.value)}.{node.func.attr}()"
+ return str(node)
+ except Exception:
+ return str(node)
+
+ def _get_attr_string(self, node: ast.AST) -> str:
+ """Get string representation of attribute access"""
+ if isinstance(node, ast.Name):
+ return node.id
+ elif isinstance(node, ast.Attribute):
+ return f"{self._get_attr_string(node.value)}.{node.attr}"
+ return str(node)
+
+
+def check_file_for_memory_violations(file_path: str, patterns: Optional[Sequence[Pattern]] = None) -> List[Dict[str, Any]]:
+ """Check a single file for memory violations"""
+ try:
+ with open(file_path, "r", encoding="utf-8") as f:
+ content = f.read()
+
+ if "test" in file_path.lower() or "__pycache__" in file_path:
+ return []
+
+ tree = ast.parse(content, filename=file_path)
+ detector = MemoryViolationDetector(file_path, patterns)
+ detector.ast_tree = tree # Store AST for potential future use
+ detector.visit(tree)
+ # Class-level operations are checked when visiting functions
+ return detector.violations
+ except Exception as e:
+ print(f"Error parsing {file_path}: {e}")
+ return []
+
+
+def check_directory_for_memory_violations(directory_path: str, ignore_patterns: Optional[List[str]] = None,
+ patterns: Optional[Sequence[Pattern]] = None) -> List[Dict[str, Any]]:
+ """Recursively scan directory for memory violations"""
+ if ignore_patterns is None:
+ ignore_patterns = ["__pycache__", ".pyc", "site-packages", "venv", ".venv", "env", ".env", "node_modules", "tests"]
+
+ all_violations = []
+ for root, _dirs, files in os.walk(directory_path):
+ if any(pattern in root for pattern in ignore_patterns):
+ continue
+ for file in files:
+ if file.endswith(".py"):
+ violations = check_file_for_memory_violations(os.path.join(root, file), patterns)
+ all_violations.extend(violations)
+ return all_violations
+
+
+def main():
+ """Run memory violation detection on codebase"""
+ codebase_path = "./litellm"
+
+ print("=" * 80)
+ print("MEMORY VIOLATION DETECTION TEST")
+ print("=" * 80)
+ print(f"Scanning: {codebase_path}")
+ print(f"Active patterns: {', '.join(p.get_pattern_name() for p in MemoryViolationDetector.DEFAULT_PATTERNS)}")
+ print()
+
+ violations = check_directory_for_memory_violations(codebase_path)
+
+ if violations:
+ by_type = {}
+ for v in violations:
+ vtype = v["type"]
+ if vtype not in by_type:
+ by_type[vtype] = []
+ by_type[vtype].append(v)
+
+ print("MEMORY VIOLATIONS FOUND:")
+ print("=" * 80)
+
+ total = len(violations)
+ for vtype, vlist in by_type.items():
+ print(f"\n{vtype.upper().replace('_', ' ')}: {len(vlist)} violation(s)")
+ print("-" * 80)
+ for v in vlist[:10]:
+ print(f" [VIOLATION] {v['file_path'] if 'file_path' in v else 'unknown'}:{v['line']}")
+ print(f" Function: {v['function']}")
+ print(f" Variable: {v['var_name']}")
+ print(f" {v['message']}")
+ print()
+ if len(vlist) > 10:
+ print(f" ... and {len(vlist) - 10} more violations of this type")
+
+ print("=" * 80)
+ print(f"TOTAL VIOLATIONS: {total}")
+ print()
+ print("RECOMMENDATIONS:")
+ print(" 1. Set queue variables to None after use: obj = queue.get(); ...; obj = None")
+ print(" 2. Use bounded queues to prevent unbounded accumulation")
+ print(" 3. Process items faster than they're added, or drain queues periodically")
+ print(" 4. For class-level data structures (lists, dicts, sets) that are modified at runtime:")
+ print(" - Add size limit checks: if len(data) >= MAX_SIZE: ...")
+ print(" - Implement periodic cleanup or use bounded collections")
+ print(" - Consider using collections.deque with maxlen for lists")
+ print("=" * 80)
+
+ first_v = violations[0]
+ raise Exception(
+ f"Found {total} memory violations! "
+ f"First violation: {first_v.get('file_path', 'unknown')}:{first_v['line']} - "
+ f"{first_v['message']}"
+ )
+ else:
+ print("OK No memory violations found!")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py
index 8331738baba..20e6a381e5b 100644
--- a/tests/code_coverage_tests/recursive_detector.py
+++ b/tests/code_coverage_tests/recursive_detector.py
@@ -37,6 +37,7 @@ IGNORE_FUNCTIONS = [
"_split_text", # max depth set.
"_delete_nested_value_custom", # max depth set (bounded by number of path segments).
"filter_exceptions_from_params", # max depth set (default 20) to prevent infinite recursion.
+ "__getattr__", # lazy loading pattern in litellm/__init__.py with proper caching to prevent infinite recursion.
]
diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py
index 7e269f21451..c151150f634 100644
--- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py
+++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py
@@ -903,3 +903,159 @@ def test_convert_to_model_response_object_with_thinking_content():
resp: ModelResponse = convert_to_model_response_object(**args)
assert resp is not None
assert resp.choices[0].message.reasoning_content is not None
+
+
+def test_convert_to_model_response_object_with_empty_error_object():
+ """
+ Test that convert_to_model_response_object handles empty error objects gracefully.
+
+ This is a regression test for issue #18407 where providers like Apertis return
+ empty error objects even on successful responses, causing spurious APIErrors.
+
+ The error object structure:
+ {
+ "error": {
+ "message": "",
+ "type": "",
+ "param": "",
+ "code": null
+ }
+ }
+ """
+ response_object = {
+ "model": "minimax-m2.1",
+ "choices": [
+ {
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": "Hey! I'm doing well, thanks for asking!",
+ },
+ "finish_reason": "stop",
+ }
+ ],
+ "usage": {
+ "prompt_tokens": 49,
+ "completion_tokens": 87,
+ "total_tokens": 136,
+ },
+ "error": {
+ "message": "",
+ "type": "",
+ "param": "",
+ "code": None,
+ },
+ }
+
+ # This should NOT raise an exception
+ result = convert_to_model_response_object(
+ model_response_object=ModelResponse(),
+ response_object=response_object,
+ stream=False,
+ start_time=datetime.now(),
+ end_time=datetime.now(),
+ hidden_params=None,
+ _response_headers=None,
+ convert_tool_call_to_json_mode=False,
+ )
+
+ assert isinstance(result, ModelResponse)
+ assert result.model == "minimax-m2.1"
+ assert len(result.choices) == 1
+ assert result.choices[0].message.content == "Hey! I'm doing well, thanks for asking!"
+
+
+def test_convert_to_model_response_object_with_real_error():
+ """
+ Test that convert_to_model_response_object still raises for real errors.
+
+ Ensures the empty error fix doesn't break legitimate error handling.
+ """
+ response_object = {
+ "error": {
+ "message": "Rate limit exceeded",
+ "type": "rate_limit_error",
+ "param": None,
+ "code": 429,
+ },
+ }
+
+ with pytest.raises(Exception) as exc_info:
+ convert_to_model_response_object(
+ model_response_object=ModelResponse(),
+ response_object=response_object,
+ stream=False,
+ start_time=datetime.now(),
+ end_time=datetime.now(),
+ hidden_params=None,
+ _response_headers=None,
+ convert_tool_call_to_json_mode=False,
+ )
+
+ # The exception should have the error message
+ assert hasattr(exc_info.value, "message")
+ assert "Rate limit exceeded" in str(exc_info.value.message)
+
+
+def test_convert_to_model_response_object_with_empty_dict_error():
+ """
+ Test that convert_to_model_response_object handles completely empty error dict.
+ """
+ response_object = {
+ "model": "test-model",
+ "choices": [
+ {
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": "Hello!",
+ },
+ "finish_reason": "stop",
+ }
+ ],
+ "usage": {
+ "prompt_tokens": 10,
+ "completion_tokens": 5,
+ "total_tokens": 15,
+ },
+ "error": {}, # Completely empty error object
+ }
+
+ # This should NOT raise an exception
+ result = convert_to_model_response_object(
+ model_response_object=ModelResponse(),
+ response_object=response_object,
+ stream=False,
+ start_time=datetime.now(),
+ end_time=datetime.now(),
+ hidden_params=None,
+ _response_headers=None,
+ convert_tool_call_to_json_mode=False,
+ )
+
+ assert isinstance(result, ModelResponse)
+ assert result.choices[0].message.content == "Hello!"
+
+
+def test_convert_to_model_response_object_with_error_code_only():
+ """
+ Test that errors with only a code (no message) are still treated as real errors.
+ """
+ response_object = {
+ "error": {
+ "message": "",
+ "code": 500,
+ },
+ }
+
+ with pytest.raises(Exception):
+ convert_to_model_response_object(
+ model_response_object=ModelResponse(),
+ response_object=response_object,
+ stream=False,
+ start_time=datetime.now(),
+ end_time=datetime.now(),
+ hidden_params=None,
+ _response_headers=None,
+ convert_tool_call_to_json_mode=False,
+ )
diff --git a/tests/local_testing/test_anthropic_prompt_caching.py b/tests/local_testing/test_anthropic_prompt_caching.py
index 0926bd17b70..c8589dd8844 100644
--- a/tests/local_testing/test_anthropic_prompt_caching.py
+++ b/tests/local_testing/test_anthropic_prompt_caching.py
@@ -104,7 +104,6 @@ async def test_litellm_anthropic_prompt_caching_tools():
],
extra_headers={
"anthropic-version": "2023-06-01",
- "anthropic-beta": "prompt-caching-2024-07-31",
},
)
@@ -112,11 +111,12 @@ async def test_litellm_anthropic_prompt_caching_tools():
print("call args=", mock_post.call_args)
expected_url = "https://api.anthropic.com/v1/messages"
+ # Note: anthropic-beta header for prompt-caching is no longer required
+ # Anthropic now supports prompt caching automatically when cache_control is used
expected_headers = {
"accept": "application/json",
"content-type": "application/json",
"anthropic-version": "2023-06-01",
- "anthropic-beta": "prompt-caching-2024-07-31",
"x-api-key": "mock_api_key",
}
@@ -285,7 +285,6 @@ async def test_anthropic_api_prompt_caching_basic():
max_tokens=10,
extra_headers={
"anthropic-version": "2023-06-01",
- "anthropic-beta": "prompt-caching-2024-07-31",
},
)
@@ -356,7 +355,6 @@ async def test_anthropic_api_prompt_caching_basic_with_cache_creation():
max_tokens=10,
extra_headers={
"anthropic-version": "2023-06-01",
- "anthropic-beta": "prompt-caching-2024-07-31",
},
)
@@ -645,7 +643,6 @@ async def test_litellm_anthropic_prompt_caching_system():
],
extra_headers={
"anthropic-version": "2023-06-01",
- "anthropic-beta": "prompt-caching-2024-07-31",
},
)
@@ -657,7 +654,6 @@ async def test_litellm_anthropic_prompt_caching_system():
"accept": "application/json",
"content-type": "application/json",
"anthropic-version": "2023-06-01",
- "anthropic-beta": "prompt-caching-2024-07-31",
"x-api-key": "mock_api_key",
}
diff --git a/tests/local_testing/test_gemini_reasoning_content.py b/tests/local_testing/test_gemini_reasoning_content.py
index 758616b4c94..f1f9c2ab512 100644
--- a/tests/local_testing/test_gemini_reasoning_content.py
+++ b/tests/local_testing/test_gemini_reasoning_content.py
@@ -1,4 +1,5 @@
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
+from litellm.llms.vertex_ai.gemini.transformation import _gemini_convert_messages_with_history
def test_thought_true_creates_thinking_block():
@@ -36,3 +37,108 @@ def test_thought_signature_without_thought_does_not_create_block():
config = VertexGeminiConfig()
thinking_blocks = config._extract_thinking_blocks_from_parts(parts)
assert thinking_blocks == []
+
+
+def test_extract_thought_signatures_from_regular_parts():
+ """
+ Test that thoughtSignatures are extracted from regular text parts (without thought=True).
+ This is the key feature for Gemini 3 multi-turn context preservation.
+ """
+ parts = [{"text": "I am Gemini", "thoughtSignature": "sig-regular-123"}]
+ config = VertexGeminiConfig()
+
+ # Should NOT create thinking block
+ thinking_blocks = config._extract_thinking_blocks_from_parts(parts)
+ assert thinking_blocks == []
+
+ # Should extract thought signature
+ signatures = config._extract_thought_signatures_from_parts(parts)
+ assert signatures is not None
+ assert len(signatures) == 1
+ assert signatures[0] == "sig-regular-123"
+
+
+def test_extract_multiple_thought_signatures():
+ """
+ Test extraction of multiple thoughtSignatures from different parts.
+ """
+ parts = [
+ {"text": "Part 1", "thoughtSignature": "sig-1"},
+ {"text": "Part 2", "thoughtSignature": "sig-2"},
+ {"text": "Part 3"} # No signature
+ ]
+ config = VertexGeminiConfig()
+ signatures = config._extract_thought_signatures_from_parts(parts)
+
+ assert signatures is not None
+ assert len(signatures) == 2
+ assert signatures[0] == "sig-1"
+ assert signatures[1] == "sig-2"
+
+
+def test_round_trip_thought_signature_in_conversation():
+ """
+ Test that thoughtSignatures are properly round-tripped through conversation history.
+ This ensures multi-turn context preservation works correctly.
+ """
+ messages = [
+ {"role": "user", "content": "Hello"},
+ {
+ "role": "assistant",
+ "content": "Hi there",
+ "provider_specific_fields": {
+ "thought_signatures": ["sig-round-trip-abc"]
+ }
+ },
+ {"role": "user", "content": "How are you?"}
+ ]
+
+ gemini_contents = _gemini_convert_messages_with_history(messages)
+
+ # Find the assistant (model) message
+ model_message = None
+ for content in gemini_contents:
+ if content.get("role") == "model":
+ model_message = content
+ break
+
+ assert model_message is not None
+ assert len(model_message["parts"]) >= 1
+
+ # Check that the text part has the thoughtSignature
+ text_part = model_message["parts"][0]
+ assert text_part["text"] == "Hi there"
+ assert "thoughtSignature" in text_part
+ assert text_part["thoughtSignature"] == "sig-round-trip-abc"
+
+
+def test_round_trip_without_thought_signature_still_works():
+ """
+ Test that messages without thoughtSignatures continue to work normally.
+ This ensures backward compatibility.
+ """
+ messages = [
+ {"role": "user", "content": "Hello"},
+ {
+ "role": "assistant",
+ "content": "Hi there"
+ },
+ {"role": "user", "content": "How are you?"}
+ ]
+
+ gemini_contents = _gemini_convert_messages_with_history(messages)
+
+ # Find the assistant (model) message
+ model_message = None
+ for content in gemini_contents:
+ if content.get("role") == "model":
+ model_message = content
+ break
+
+ assert model_message is not None
+ assert len(model_message["parts"]) >= 1
+
+ # Check that the text part works without thoughtSignature
+ text_part = model_message["parts"][0]
+ assert text_part["text"] == "Hi there"
+ assert "thoughtSignature" not in text_part
diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json
index 26b712c1cf2..bd2f06b502e 100644
--- a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json
+++ b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json
@@ -54,7 +54,7 @@
"id": "time-14-13-16-469836_chatcmpl-3803a9e9-aa68-4493-94d9-247f354830d6",
"endTime": "2025-05-26T14:13:16.795438-07:00",
"completionStartTime": "2025-05-26T14:13:16.795438-07:00",
- "model": "anthropic.claude-3-5-sonnet-20240620-v1:0",
+ "model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
"modelParameters": {
"aws_region": "us-east-1"
},
diff --git a/tests/logging_callback_tests/test_generic_api_callback.py b/tests/logging_callback_tests/test_generic_api_callback.py
index 3ddf84f2939..c3e1171e96a 100644
--- a/tests/logging_callback_tests/test_generic_api_callback.py
+++ b/tests/logging_callback_tests/test_generic_api_callback.py
@@ -207,3 +207,253 @@ async def test_generic_api_callback_multiple_logs():
assert (
payload_item["response"]["choices"][0]["message"]["content"] == "hi"
), "Response should be hi"
+
+
+@pytest.mark.asyncio
+async def test_generic_api_callback_ndjson_format():
+ """
+ Test the GenericAPILogger callback with ndjson log format.
+ Validates that logs are sent as newline-delimited JSON.
+ """
+ # Create a mock for the async_httpx_client's post method
+ mock_post = AsyncMock()
+ mock_post.return_value.status_code = 200
+ mock_post.return_value.text = "OK"
+
+ # Set up an endpoint for testing
+ test_endpoint = "https://example.com/api/logs"
+ test_headers = {"Authorization": "Bearer test_token"}
+ os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint
+
+ # Initialize the GenericAPILogger with ndjson format
+ generic_logger = GenericAPILogger(
+ endpoint=test_endpoint,
+ headers=test_headers,
+ flush_interval=1,
+ log_format="ndjson" # Set NDJSON format
+ )
+ generic_logger.async_httpx_client.post = mock_post
+ litellm.callbacks = [generic_logger]
+
+ # Make multiple completion calls to generate multiple logs
+ for i in range(3):
+ response = await litellm.acompletion(
+ model="gpt-4o",
+ messages=[{"role": "user", "content": f"Hello, world! {i}"}],
+ mock_response="hi",
+ user="test_user",
+ )
+
+ # Wait for async flush
+ await asyncio.sleep(3)
+
+ # Assert httpx post was called
+ mock_post.assert_called_once()
+
+ # Get the actual request body from the mock
+ actual_url = mock_post.call_args[1]["url"]
+ assert actual_url == test_endpoint, f"Expected URL {test_endpoint}, got {actual_url}"
+
+ # Get the data sent
+ ndjson_data = mock_post.call_args[1]["data"]
+ print("##########\n")
+ print("ndjson_data:", ndjson_data)
+ print("##########\n")
+
+ # Validate it's NDJSON format (newline-delimited)
+ assert isinstance(ndjson_data, str), "Data should be a string for NDJSON"
+
+ # Split by newlines and parse each line
+ lines = ndjson_data.strip().split("\n")
+ assert len(lines) == 3, f"Expected 3 lines of NDJSON, got {len(lines)}"
+
+ # Validate each line is valid JSON
+ for i, line in enumerate(lines):
+ payload_item = json.loads(line)
+ payload_item = StandardLoggingPayload(**payload_item)
+
+ # Basic assertions
+ assert payload_item["response_cost"] > 0, "Response cost should be greater than 0"
+ assert payload_item["model"] == "gpt-4o", "Model should be gpt-4o"
+ assert payload_item["model_parameters"]["user"] == "test_user", "User should be test_user"
+
+
+@pytest.mark.asyncio
+async def test_generic_api_callback_single_format():
+ """
+ Test the GenericAPILogger callback with single log format.
+ Validates that each log is sent as an individual request in parallel.
+ """
+ # Create a mock for the async_httpx_client's post method
+ mock_post = AsyncMock()
+ mock_post.return_value.status_code = 200
+ mock_post.return_value.text = "OK"
+
+ # Set up an endpoint for testing
+ test_endpoint = "https://example.com/api/logs"
+ test_headers = {"Authorization": "Bearer test_token"}
+ os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint
+
+ # Initialize the GenericAPILogger with single format
+ generic_logger = GenericAPILogger(
+ endpoint=test_endpoint,
+ headers=test_headers,
+ flush_interval=1, # Quick flush to trigger batch send
+ log_format="single" # Set single format
+ )
+ generic_logger.async_httpx_client.post = mock_post
+ litellm.callbacks = [generic_logger]
+
+ # Make 3 completion calls
+ for i in range(3):
+ response = await litellm.acompletion(
+ model="gpt-4o",
+ messages=[{"role": "user", "content": f"Hello, world! {i}"}],
+ mock_response="hi",
+ user="test_user",
+ )
+
+ # Wait for async flush
+ await asyncio.sleep(3)
+
+ # Assert httpx post was called 3 times (once per log in batch)
+ assert mock_post.call_count == 3, f"Expected 3 calls, got {mock_post.call_count}"
+
+ # Validate each call sent a single log object (not an array)
+ for call_idx in range(3):
+ call_args = mock_post.call_args_list[call_idx]
+ json_data = call_args[1]["data"]
+
+ print(f"########## Call {call_idx} ##########")
+ print("json_data:", json_data)
+
+ # Parse and validate - should be a single object, not an array
+ actual_request = json.loads(json_data)
+ assert isinstance(actual_request, dict), f"Call {call_idx}: Expected dict, got {type(actual_request)}"
+
+ # Validate it's a valid StandardLoggingPayload
+ payload_item = StandardLoggingPayload(**actual_request)
+ assert payload_item["response_cost"] > 0, "Response cost should be greater than 0"
+ assert payload_item["model"] == "gpt-4o", "Model should be gpt-4o"
+
+
+@pytest.mark.asyncio
+async def test_generic_api_callback_json_array_format_explicit():
+ """
+ Test the GenericAPILogger callback with explicit json_array format.
+ Validates backward compatibility when explicitly set to json_array.
+ """
+ # Create a mock for the async_httpx_client's post method
+ mock_post = AsyncMock()
+ mock_post.return_value.status_code = 200
+ mock_post.return_value.text = "OK"
+
+ # Set up an endpoint for testing
+ test_endpoint = "https://example.com/api/logs"
+ test_headers = {"Authorization": "Bearer test_token"}
+ os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint
+
+ # Initialize the GenericAPILogger with explicit json_array format
+ generic_logger = GenericAPILogger(
+ endpoint=test_endpoint,
+ headers=test_headers,
+ flush_interval=1,
+ log_format="json_array" # Explicitly set json_array
+ )
+ generic_logger.async_httpx_client.post = mock_post
+ litellm.callbacks = [generic_logger]
+
+ # Make multiple completion calls
+ for i in range(5):
+ response = await litellm.acompletion(
+ model="gpt-4o",
+ messages=[{"role": "user", "content": f"Hello, world! {i}"}],
+ mock_response="hi",
+ user="test_user",
+ )
+
+ # Wait for async flush
+ await asyncio.sleep(3)
+
+ # Assert httpx post was called once (batched)
+ mock_post.assert_called_once()
+
+ # Get the data and validate it's a JSON array
+ json_data = mock_post.call_args[1]["data"]
+ actual_request = json.loads(json_data)
+
+ assert isinstance(actual_request, list), "Request body should be a list (JSON array)"
+ assert len(actual_request) == 5, f"Expected 5 items, got {len(actual_request)}"
+
+ # Validate each item
+ for payload_item in actual_request:
+ payload_item = StandardLoggingPayload(**payload_item)
+ assert payload_item["response_cost"] > 0, "Response cost should be greater than 0"
+ assert payload_item["model"] == "gpt-4o", "Model should be gpt-4o"
+
+
+@pytest.mark.asyncio
+async def test_generic_api_callback_sumologic_uses_ndjson():
+ """
+ Test that the sumologic callback uses ndjson format by default
+ when loaded from generic_api_compatible_callbacks.json
+ """
+ # Create a mock for the async_httpx_client's post method
+ mock_post = AsyncMock()
+ mock_post.return_value.status_code = 200
+ mock_post.return_value.text = "OK"
+
+ # Set environment variable for sumologic
+ os.environ["SUMOLOGIC_WEBHOOK_URL"] = "https://collectors.sumologic.com/receiver/v1/http/test123"
+
+ # Initialize using callback_name (loads from JSON config)
+ generic_logger = GenericAPILogger(
+ callback_name="sumologic",
+ flush_interval=1
+ )
+ generic_logger.async_httpx_client.post = mock_post
+ litellm.callbacks = [generic_logger]
+
+ # Verify the logger has ndjson format
+ assert generic_logger.log_format == "ndjson", "Sumologic should use ndjson format"
+
+ # Make completion calls
+ for i in range(2):
+ await litellm.acompletion(
+ model="gpt-4o",
+ messages=[{"role": "user", "content": f"Test {i}"}],
+ mock_response="response",
+ user="test_user",
+ )
+
+ # Wait for async flush
+ await asyncio.sleep(3)
+
+ # Assert httpx post was called
+ mock_post.assert_called_once()
+
+ # Verify NDJSON format
+ ndjson_data = mock_post.call_args[1]["data"]
+ assert isinstance(ndjson_data, str), "Data should be a string for NDJSON"
+
+ lines = ndjson_data.strip().split("\n")
+ assert len(lines) == 2, f"Expected 2 lines of NDJSON, got {len(lines)}"
+
+ # Each line should be valid JSON
+ for line in lines:
+ json.loads(line) # Will raise if invalid JSON
+
+
+@pytest.mark.asyncio
+async def test_generic_api_callback_invalid_log_format():
+ """
+ Test that invalid log_format values raise a ValueError
+ """
+ test_endpoint = "https://example.com/api/logs"
+ os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint
+
+ with pytest.raises(ValueError, match="Invalid log_format"):
+ GenericAPILogger(
+ endpoint=test_endpoint,
+ log_format="invalid_format" # type: ignore # Intentionally invalid for testing
+ )
diff --git a/tests/logging_callback_tests/test_opentelemetry_unit_tests.py b/tests/logging_callback_tests/test_opentelemetry_unit_tests.py
index c8ceded4cf2..3d0682d9033 100644
--- a/tests/logging_callback_tests/test_opentelemetry_unit_tests.py
+++ b/tests/logging_callback_tests/test_opentelemetry_unit_tests.py
@@ -40,11 +40,15 @@ class TestOpentelemetryUnitTests(BaseLoggingCallbackTest):
@pytest.mark.asyncio
async def test_opentelemetry_integration(self):
"""
- Unit test to confirm the parent otel span is ended.
+ Unit test to confirm external parent otel spans are NOT ended by LiteLLM.
+
+ External spans (passed via metadata) should be managed by their creators,
+ not by LiteLLM. This prevents premature closure of spans from Langfuse,
+ user code, or other external observability tools.
"""
# Reset all callbacks to ensure clean state
litellm.logging_callback_manager._reset_all_callbacks()
-
+
parent_otel_span = MagicMock()
litellm.callbacks = ["otel"]
@@ -57,8 +61,9 @@ class TestOpentelemetryUnitTests(BaseLoggingCallbackTest):
await asyncio.sleep(1)
- # Verify span was ended (may be called multiple times due to callback architecture)
- parent_otel_span.end.assert_called()
+ # Verify external span was NOT ended by LiteLLM
+ # External spans should only be closed by their creators
+ parent_otel_span.end.assert_not_called()
def test_init_tracing_respects_existing_tracer_provider(self):
"""
diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py
index 9242dfc75f4..8fb0e80cc39 100644
--- a/tests/mcp_tests/test_mcp_server.py
+++ b/tests/mcp_tests/test_mcp_server.py
@@ -1057,7 +1057,7 @@ async def test_mcp_server_manager_config_integration_with_database():
)
# Test the add_update_server method (this tests our fix)
- await test_manager.add_update_server(db_server)
+ await test_manager.add_server(db_server)
# Verify the server was added with correct access_groups
registry = test_manager.get_registry()
@@ -1381,7 +1381,7 @@ async def test_add_update_server_with_alias():
mock_mcp_server.token_url = None
# Add server to manager
- await test_manager.add_update_server(mock_mcp_server)
+ await test_manager.add_server(mock_mcp_server)
# Verify server was added with correct name (should use alias)
assert "test-server-123" in test_manager.registry
@@ -1421,7 +1421,7 @@ async def test_add_update_server_without_alias():
mock_mcp_server.token_url = None
# Add server to manager
- await test_manager.add_update_server(mock_mcp_server)
+ await test_manager.add_server(mock_mcp_server)
# Verify server was added with correct name (should use server_name)
assert "test-server-123" in test_manager.registry
@@ -1461,7 +1461,7 @@ async def test_add_update_server_fallback_to_server_id():
mock_mcp_server.token_url = None
# Add server to manager
- await test_manager.add_update_server(mock_mcp_server)
+ await test_manager.add_server(mock_mcp_server)
# Verify server was added with correct name (should use server_id)
assert "test-server-123" in test_manager.registry
diff --git a/tests/old_proxy_tests/tests/test_anthropic_context_caching.py b/tests/old_proxy_tests/tests/test_anthropic_context_caching.py
index 7a153295f35..6b37873df4e 100644
--- a/tests/old_proxy_tests/tests/test_anthropic_context_caching.py
+++ b/tests/old_proxy_tests/tests/test_anthropic_context_caching.py
@@ -30,7 +30,6 @@ response = client.chat.completions.create(
],
extra_headers={
"anthropic-version": "2023-06-01",
- "anthropic-beta": "prompt-caching-2024-07-31",
},
)
diff --git a/tests/store_model_in_db_tests/test_mcp_servers.py b/tests/store_model_in_db_tests/test_mcp_servers.py
index 8197ce2693b..a369cce83c0 100644
--- a/tests/store_model_in_db_tests/test_mcp_servers.py
+++ b/tests/store_model_in_db_tests/test_mcp_servers.py
@@ -139,7 +139,7 @@ async def test_create_mcp_server_direct():
mock_get_prisma.return_value = mock_prisma
# Mock server manager
- mock_manager.add_update_server = mock.AsyncMock()
+ mock_manager.add_server = mock.AsyncMock()
mock_manager.reload_servers_from_database = mock.AsyncMock()
# Set up test data
@@ -195,7 +195,7 @@ async def test_create_mcp_server_direct():
# Verify mocks were called
mock_get_server.assert_called_once_with(mock_prisma, server_id)
mock_create.assert_called_once()
- mock_manager.add_update_server.assert_called_once_with(expected_response)
+ mock_manager.add_server.assert_called_once_with(expected_response)
@pytest.mark.asyncio
@@ -379,7 +379,7 @@ async def test_edit_mcp_server_redacts_credentials():
mock_prisma = mock.Mock()
mock_get_prisma.return_value = mock_prisma
- mock_manager.add_update_server = mock.AsyncMock()
+ mock_manager.update_server = mock.AsyncMock()
mock_manager.reload_servers_from_database = mock.AsyncMock()
server_id = str(uuid.uuid4())
@@ -417,7 +417,7 @@ async def test_edit_mcp_server_redacts_credentials():
mock_validate.assert_called_once()
mock_update.assert_awaited_once()
- mock_manager.add_update_server.assert_called_once_with(updated_server)
+ mock_manager.update_server.assert_called_once_with(updated_server)
mock_manager.reload_servers_from_database.assert_awaited_once()
def test_validate_mcp_server_name_direct():
"""
diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter.py b/tests/test_litellm/google_genai/test_google_genai_adapter.py
index e8882a1acb3..135881ad209 100644
--- a/tests/test_litellm/google_genai/test_google_genai_adapter.py
+++ b/tests/test_litellm/google_genai/test_google_genai_adapter.py
@@ -1197,6 +1197,131 @@ async def test_agenerate_content_x_goog_api_key_header():
# Verify other expected headers
assert headers.get("Content-Type") == "application/json", f"Expected Content-Type application/json, got {headers.get('Content-Type')}"
-
+
print(f"✓ Test passed: x-goog-api-key header correctly set to {api_key_value}")
print(f"✓ All headers: {list(headers.keys())}")
+
+
+def test_inline_data_base64_image_transformation():
+ """Test transformation of Gemini inline_data (Base64 images) to OpenAI format"""
+ from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
+
+ adapter = GoogleGenAIAdapter()
+
+ # Test input with Base64 image
+ model = "gpt-4-vision-preview"
+ contents = {
+ "role": "user",
+ "parts": [
+ {"text": "What's in this image?"},
+ {
+ "inline_data": {
+ "mime_type": "image/jpeg",
+ "data": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
+ }
+ }
+ ]
+ }
+
+ # Transform to completion format
+ completion_request = adapter.translate_generate_content_to_completion(
+ model=model,
+ contents=contents
+ )
+
+ # Verify the transformation
+ assert completion_request["model"] == model
+ assert len(completion_request["messages"]) == 1
+ assert completion_request["messages"][0]["role"] == "user"
+
+ # Verify content is an array (multimodal format)
+ content = completion_request["messages"][0]["content"]
+ assert isinstance(content, list), "Content should be a list for multimodal messages"
+ assert len(content) == 2, "Should have 2 content parts (text + image)"
+
+ # Verify text part
+ text_part = content[0]
+ assert text_part["type"] == "text"
+ assert text_part["text"] == "What's in this image?"
+
+ # Verify image part
+ image_part = content[1]
+ assert image_part["type"] == "image_url"
+ assert "image_url" in image_part
+ assert "url" in image_part["image_url"]
+ assert image_part["image_url"]["url"].startswith("data:image/jpeg;base64,")
+ assert "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" in image_part["image_url"]["url"]
+
+
+def test_inline_data_image_only_transformation():
+ """Test transformation of Gemini inline_data with only image (no text)"""
+ from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
+
+ adapter = GoogleGenAIAdapter()
+
+ # Test input with only Base64 image (no text)
+ model = "gpt-4-vision-preview"
+ contents = {
+ "role": "user",
+ "parts": [
+ {
+ "inline_data": {
+ "mime_type": "image/png",
+ "data": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
+ }
+ }
+ ]
+ }
+
+ # Transform to completion format
+ completion_request = adapter.translate_generate_content_to_completion(
+ model=model,
+ contents=contents
+ )
+
+ # Verify the transformation
+ assert completion_request["model"] == model
+ assert len(completion_request["messages"]) == 1
+ assert completion_request["messages"][0]["role"] == "user"
+
+ # Verify content is an array (multimodal format)
+ content = completion_request["messages"][0]["content"]
+ assert isinstance(content, list), "Content should be a list for multimodal messages"
+ assert len(content) == 1, "Should have 1 content part (image only)"
+
+ # Verify image part
+ image_part = content[0]
+ assert image_part["type"] == "image_url"
+ assert "image_url" in image_part
+ assert "url" in image_part["image_url"]
+ assert image_part["image_url"]["url"].startswith("data:image/png;base64,")
+
+
+def test_inline_data_backward_compatibility_text_only():
+ """Test that pure text messages still use simple string format (backward compatibility)"""
+ from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
+
+ adapter = GoogleGenAIAdapter()
+
+ # Test input with only text (no images)
+ model = "gpt-3.5-turbo"
+ contents = {
+ "role": "user",
+ "parts": [{"text": "Hello, how are you?"}]
+ }
+
+ # Transform to completion format
+ completion_request = adapter.translate_generate_content_to_completion(
+ model=model,
+ contents=contents
+ )
+
+ # Verify the transformation
+ assert completion_request["model"] == model
+ assert len(completion_request["messages"]) == 1
+ assert completion_request["messages"][0]["role"] == "user"
+
+ # Verify content is a simple string (not an array) for backward compatibility
+ content = completion_request["messages"][0]["content"]
+ assert isinstance(content, str), "Content should be a string for text-only messages (backward compatibility)"
+ assert content == "Hello, how are you?"
diff --git a/tests/test_litellm/google_genai/test_google_genai_transformation.py b/tests/test_litellm/google_genai/test_google_genai_transformation.py
index c953a504a38..0d5468cf275 100644
--- a/tests/test_litellm/google_genai/test_google_genai_transformation.py
+++ b/tests/test_litellm/google_genai/test_google_genai_transformation.py
@@ -247,3 +247,122 @@ def test_responses_api_no_reasoning():
# reasoning_effort should not be in result if not provided (filtered out as None)
assert "reasoning_effort" not in result or result.get("reasoning_effort") is None
+
+
+def test_transform_generate_content_request_with_system_instruction():
+ """Test that systemInstruction parameter is properly included in the request"""
+ config = GoogleGenAIConfig()
+
+ system_instruction = {
+ "parts": [{"text": "You are a helpful assistant"}]
+ }
+
+ contents = [
+ {
+ "role": "user",
+ "parts": [{"text": "Hello"}]
+ }
+ ]
+
+ generate_content_config_dict = {
+ "temperature": 1.0,
+ "maxOutputTokens": 100
+ }
+
+ # Call transform_generate_content_request
+ result = config.transform_generate_content_request(
+ model="gemini-3-flash-preview",
+ contents=contents,
+ tools=None,
+ generate_content_config_dict=generate_content_config_dict,
+ system_instruction=system_instruction,
+ )
+
+ # Verify that systemInstruction is in the request
+ assert "systemInstruction" in result, "systemInstruction should be in request body"
+ assert result["systemInstruction"] == system_instruction, "systemInstruction should match input"
+ assert result["model"] == "gemini-3-flash-preview"
+ assert result["contents"] == contents
+
+
+def test_transform_generate_content_request_without_system_instruction():
+ """Test that request works correctly without systemInstruction"""
+ config = GoogleGenAIConfig()
+
+ contents = [
+ {
+ "role": "user",
+ "parts": [{"text": "Hello"}]
+ }
+ ]
+
+ generate_content_config_dict = {
+ "temperature": 1.0
+ }
+
+ # Call transform_generate_content_request without system_instruction
+ result = config.transform_generate_content_request(
+ model="gemini-3-flash-preview",
+ contents=contents,
+ tools=None,
+ generate_content_config_dict=generate_content_config_dict,
+ system_instruction=None,
+ )
+
+ # Verify that systemInstruction is NOT in the request when not provided
+ assert "systemInstruction" not in result, "systemInstruction should not be in request when None"
+ assert result["model"] == "gemini-3-flash-preview"
+ assert result["contents"] == contents
+
+
+def test_transform_generate_content_request_system_instruction_with_tools():
+ """Test that systemInstruction works correctly alongside tools"""
+ config = GoogleGenAIConfig()
+
+ system_instruction = {
+ "parts": [{"text": "You are a helpful assistant that uses tools"}]
+ }
+
+ contents = [
+ {
+ "role": "user",
+ "parts": [{"text": "What's the weather?"}]
+ }
+ ]
+
+ tools = [
+ {
+ "functionDeclarations": [
+ {
+ "name": "get_weather",
+ "description": "Get weather information",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"type": "string"}
+ }
+ }
+ }
+ ]
+ }
+ ]
+
+ generate_content_config_dict = {
+ "temperature": 0.7
+ }
+
+ # Call transform_generate_content_request with both system_instruction and tools
+ result = config.transform_generate_content_request(
+ model="gemini-3-flash-preview",
+ contents=contents,
+ tools=tools,
+ generate_content_config_dict=generate_content_config_dict,
+ system_instruction=system_instruction,
+ )
+
+ # Verify that both systemInstruction and tools are in the request
+ assert "systemInstruction" in result, "systemInstruction should be in request body"
+ assert result["systemInstruction"] == system_instruction
+ assert "tools" in result, "tools should be in request body"
+ assert result["tools"] == tools
+ assert result["model"] == "gemini-3-flash-preview"
diff --git a/tests/test_litellm/integrations/levo/__init__.py b/tests/test_litellm/integrations/levo/__init__.py
new file mode 100644
index 00000000000..1560e78b7b9
--- /dev/null
+++ b/tests/test_litellm/integrations/levo/__init__.py
@@ -0,0 +1 @@
+# Levo integration tests
diff --git a/tests/test_litellm/integrations/levo/test_levo.py b/tests/test_litellm/integrations/levo/test_levo.py
new file mode 100644
index 00000000000..5d042cbc060
--- /dev/null
+++ b/tests/test_litellm/integrations/levo/test_levo.py
@@ -0,0 +1,407 @@
+import unittest
+from unittest.mock import patch
+
+import pytest
+
+from litellm.integrations.levo.levo import LevoConfig, LevoLogger
+from litellm.integrations.opentelemetry import OpenTelemetryConfig
+
+# Try to import OpenTelemetry packages, skip tests if not available
+try:
+ from opentelemetry.sdk.trace import TracerProvider
+ from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
+ InMemorySpanExporter,
+ )
+ from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+ OPENTELEMETRY_AVAILABLE = True
+except ImportError:
+ OPENTELEMETRY_AVAILABLE = False
+
+
+class TestLevoConfig(unittest.TestCase):
+ """Unit tests for LevoLogger configuration."""
+
+ @patch.dict(
+ "os.environ",
+ {
+ "LEVOAI_API_KEY": "test-api-key",
+ "LEVOAI_ORG_ID": "test-org-id",
+ "LEVOAI_WORKSPACE_ID": "test-workspace-id",
+ "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai",
+ },
+ )
+ def test_get_levo_config_with_all_required_vars(self):
+ """Test get_levo_config() with all required environment variables."""
+ config = LevoLogger.get_levo_config()
+
+ # Verify headers include all three values
+ self.assertIn("Authorization=Bearer test-api-key", config.otlp_auth_headers)
+ self.assertIn("x-levo-organization-id=test-org-id", config.otlp_auth_headers)
+ self.assertIn("x-levo-workspace-id=test-workspace-id", config.otlp_auth_headers)
+
+ # Verify endpoint uses provided collector URL exactly as-is
+ self.assertEqual(config.endpoint, "https://collector.levo.ai")
+
+ # Verify protocol is otlp_http
+ self.assertEqual(config.protocol, "otlp_http")
+
+ @patch.dict(
+ "os.environ",
+ {
+ "LEVOAI_API_KEY": "test-api-key",
+ "LEVOAI_ORG_ID": "test-org-id",
+ "LEVOAI_WORKSPACE_ID": "test-workspace-id",
+ "LEVOAI_COLLECTOR_URL": "https://custom.collector.com",
+ },
+ )
+ def test_get_levo_config_with_custom_collector_url(self):
+ """Test get_levo_config() with custom collector URL."""
+ config = LevoLogger.get_levo_config()
+
+ # Verify endpoint uses custom URL exactly as provided
+ self.assertEqual(config.endpoint, "https://custom.collector.com")
+ self.assertEqual(config.protocol, "otlp_http")
+
+ @patch.dict("os.environ", {}, clear=True)
+ def test_get_levo_config_missing_api_key(self):
+ """Test get_levo_config() raises ValueError when LEVOAI_API_KEY is missing."""
+ with pytest.raises(ValueError, match="LEVOAI_API_KEY"):
+ LevoLogger.get_levo_config()
+
+ @patch.dict(
+ "os.environ",
+ {
+ "LEVOAI_API_KEY": "test-api-key",
+ },
+ clear=True,
+ )
+ def test_get_levo_config_missing_org_id(self):
+ """Test get_levo_config() raises ValueError when LEVOAI_ORG_ID is missing."""
+ with pytest.raises(ValueError, match="LEVOAI_ORG_ID"):
+ LevoLogger.get_levo_config()
+
+ @patch.dict(
+ "os.environ",
+ {
+ "LEVOAI_API_KEY": "test-api-key",
+ "LEVOAI_ORG_ID": "test-org-id",
+ },
+ clear=True,
+ )
+ def test_get_levo_config_missing_workspace_id(self):
+ """Test get_levo_config() raises ValueError when LEVOAI_WORKSPACE_ID is missing."""
+ with pytest.raises(ValueError, match="LEVOAI_WORKSPACE_ID"):
+ LevoLogger.get_levo_config()
+
+ @patch.dict(
+ "os.environ",
+ {
+ "LEVOAI_API_KEY": "test-api-key",
+ "LEVOAI_ORG_ID": "test-org-id",
+ "LEVOAI_WORKSPACE_ID": "test-workspace-id",
+ },
+ clear=True,
+ )
+ def test_get_levo_config_missing_collector_url(self):
+ """Test get_levo_config() raises ValueError when LEVOAI_COLLECTOR_URL is missing."""
+ with pytest.raises(ValueError, match="LEVOAI_COLLECTOR_URL"):
+ LevoLogger.get_levo_config()
+
+ @patch.dict(
+ "os.environ",
+ {
+ "LEVOAI_API_KEY": "test-api-key",
+ "LEVOAI_ORG_ID": "test-org-id",
+ "LEVOAI_WORKSPACE_ID": "test-workspace-id",
+ "LEVOAI_COLLECTOR_URL": "http://localhost:4318",
+ },
+ )
+ def test_get_levo_config_with_http_endpoint(self):
+ """Test get_levo_config() with HTTP endpoint."""
+ config = LevoLogger.get_levo_config()
+
+ # Should use HTTP endpoint exactly as provided
+ self.assertEqual(config.endpoint, "http://localhost:4318")
+ self.assertEqual(config.protocol, "otlp_http")
+
+ @patch.dict(
+ "os.environ",
+ {
+ "LEVOAI_API_KEY": "test-api-key",
+ "LEVOAI_ORG_ID": "test-org-id",
+ "LEVOAI_WORKSPACE_ID": "test-workspace-id",
+ "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai",
+ },
+ )
+ def test_levo_config_headers_format(self):
+ """Test that OTLP headers are formatted correctly."""
+ config = LevoLogger.get_levo_config()
+
+ # Verify headers contain all required parts
+ self.assertIn("Authorization=Bearer test-api-key", config.otlp_auth_headers)
+ self.assertIn("x-levo-organization-id=test-org-id", config.otlp_auth_headers)
+ self.assertIn("x-levo-workspace-id=test-workspace-id", config.otlp_auth_headers)
+
+ # Verify headers are comma-separated
+ header_parts = config.otlp_auth_headers.split(",")
+ self.assertEqual(len(header_parts), 3)
+
+
+class TestLevoIntegration(unittest.TestCase):
+ """Integration tests for LevoLogger."""
+
+ @patch.dict(
+ "os.environ",
+ {
+ "LEVOAI_API_KEY": "test-api-key",
+ "LEVOAI_ORG_ID": "test-org-id",
+ "LEVOAI_WORKSPACE_ID": "test-workspace-id",
+ "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai",
+ },
+ )
+ @pytest.mark.skipif(
+ not OPENTELEMETRY_AVAILABLE, reason="OpenTelemetry packages not installed"
+ )
+ @patch(
+ "litellm.integrations.opentelemetry.OpenTelemetry._init_otel_logger_on_litellm_proxy"
+ )
+ def test_levo_logger_instantiation(self, mock_init_proxy):
+ """Test that LevoLogger can be instantiated with proper config."""
+ # Mock the proxy initialization to avoid importing proxy code
+ mock_init_proxy.return_value = None
+
+ config = LevoLogger.get_levo_config()
+ otel_config = OpenTelemetryConfig(
+ exporter=config.protocol,
+ endpoint=config.endpoint,
+ headers=config.otlp_auth_headers,
+ )
+
+ # Create a tracer provider with in-memory exporter to avoid requiring OTLP packages
+ tracer_provider = TracerProvider()
+ tracer_provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
+
+ # Create LevoLogger instance with mocked tracer provider
+ levo_logger = LevoLogger(
+ config=otel_config, callback_name="levo", tracer_provider=tracer_provider
+ )
+
+ # Verify it's an instance of OpenTelemetry
+ self.assertIsInstance(levo_logger, LevoLogger)
+ # Check it extends OpenTelemetry by checking base classes
+ from litellm.integrations.opentelemetry import OpenTelemetry
+
+ self.assertIsInstance(levo_logger, OpenTelemetry)
+
+ # Verify callback_name is set
+ self.assertEqual(levo_logger.callback_name, "levo")
+
+ @patch.dict(
+ "os.environ",
+ {
+ "LEVOAI_API_KEY": "test-api-key",
+ "LEVOAI_ORG_ID": "test-org-id",
+ "LEVOAI_WORKSPACE_ID": "test-workspace-id",
+ "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai",
+ },
+ )
+ @pytest.mark.skipif(
+ not OPENTELEMETRY_AVAILABLE, reason="OpenTelemetry packages not installed"
+ )
+ @patch(
+ "litellm.integrations.opentelemetry.OpenTelemetry._init_otel_logger_on_litellm_proxy"
+ )
+ @pytest.mark.asyncio
+ async def test_levo_logger_health_check_healthy(self, mock_init_proxy):
+ """Test health check returns healthy status when config is valid."""
+ # Mock the proxy initialization to avoid importing proxy code
+ mock_init_proxy.return_value = None
+
+ config = LevoLogger.get_levo_config()
+ otel_config = OpenTelemetryConfig(
+ exporter=config.protocol,
+ endpoint=config.endpoint,
+ headers=config.otlp_auth_headers,
+ )
+
+ # Create tracer provider with in-memory exporter
+ tracer_provider = TracerProvider()
+ tracer_provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
+
+ levo_logger = LevoLogger(
+ config=otel_config, callback_name="levo", tracer_provider=tracer_provider
+ )
+
+ # Run health check
+ result = await levo_logger.async_health_check()
+
+ self.assertEqual(result["status"], "healthy")
+ self.assertIn("message", result)
+
+ @patch.dict("os.environ", {}, clear=True)
+ def test_levo_logger_health_check_unhealthy(self):
+ """Test health check returns unhealthy status when required vars are missing."""
+ # Try to create logger without required env vars
+ # This should fail during config, but we can test health check logic
+ with pytest.raises(ValueError):
+ LevoLogger.get_levo_config()
+
+ @patch.dict(
+ "os.environ",
+ {
+ "LEVOAI_API_KEY": "test-api-key",
+ "LEVOAI_ORG_ID": "test-org-id",
+ "LEVOAI_WORKSPACE_ID": "test-workspace-id",
+ "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai",
+ },
+ )
+ @pytest.mark.skipif(
+ not OPENTELEMETRY_AVAILABLE, reason="OpenTelemetry packages not installed"
+ )
+ @patch(
+ "litellm.integrations.opentelemetry.OpenTelemetry._init_otel_logger_on_litellm_proxy"
+ )
+ def test_levo_logger_callback_name(self, mock_init_proxy):
+ """Test that callback_name is properly set and used."""
+ # Mock the proxy initialization to avoid importing proxy code
+ mock_init_proxy.return_value = None
+
+ config = LevoLogger.get_levo_config()
+ otel_config = OpenTelemetryConfig(
+ exporter=config.protocol,
+ endpoint=config.endpoint,
+ headers=config.otlp_auth_headers,
+ )
+
+ # Create tracer provider with in-memory exporter
+ tracer_provider = TracerProvider()
+ tracer_provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
+
+ levo_logger = LevoLogger(
+ config=otel_config, callback_name="levo", tracer_provider=tracer_provider
+ )
+
+ # Verify callback_name attribute
+ self.assertEqual(levo_logger.callback_name, "levo")
+
+
+@pytest.mark.parametrize(
+ "env_vars, expected_headers_contains, expected_endpoint, expected_protocol",
+ [
+ pytest.param(
+ {
+ "LEVOAI_API_KEY": "test-key",
+ "LEVOAI_ORG_ID": "test-org",
+ "LEVOAI_WORKSPACE_ID": "test-workspace",
+ "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai",
+ },
+ [
+ "Authorization=Bearer test-key",
+ "x-levo-organization-id=test-org",
+ "x-levo-workspace-id=test-workspace",
+ ],
+ "https://collector.levo.ai",
+ "otlp_http",
+ id="collector URL with all required vars",
+ ),
+ pytest.param(
+ {
+ "LEVOAI_API_KEY": "key-123",
+ "LEVOAI_ORG_ID": "org-456",
+ "LEVOAI_WORKSPACE_ID": "workspace-789",
+ "LEVOAI_COLLECTOR_URL": "https://custom.example.com",
+ },
+ [
+ "Authorization=Bearer key-123",
+ "x-levo-organization-id=org-456",
+ "x-levo-workspace-id=workspace-789",
+ ],
+ "https://custom.example.com",
+ "otlp_http",
+ id="custom collector URL",
+ ),
+ pytest.param(
+ {
+ "LEVOAI_API_KEY": "key-123",
+ "LEVOAI_ORG_ID": "org-456",
+ "LEVOAI_WORKSPACE_ID": "workspace-789",
+ "LEVOAI_COLLECTOR_URL": "http://localhost:9999",
+ },
+ ["Authorization=Bearer key-123"],
+ "http://localhost:9999",
+ "otlp_http",
+ id="custom HTTP endpoint",
+ ),
+ ],
+)
+def test_get_levo_config_parametrized(
+ monkeypatch,
+ env_vars,
+ expected_headers_contains,
+ expected_endpoint,
+ expected_protocol,
+):
+ """Parametrized tests for get_levo_config() with various configurations."""
+ # Clear all Levo-related env vars first to ensure clean state
+ for key in [
+ "LEVOAI_API_KEY",
+ "LEVOAI_ORG_ID",
+ "LEVOAI_WORKSPACE_ID",
+ "LEVOAI_COLLECTOR_URL",
+ "LEVOAI_ENV_NAME",
+ ]:
+ monkeypatch.delenv(key, raising=False)
+
+ for key, value in env_vars.items():
+ monkeypatch.setenv(key, value)
+
+ config = LevoLogger.get_levo_config()
+
+ assert isinstance(config, LevoConfig)
+ assert config.endpoint == expected_endpoint
+ assert config.protocol == expected_protocol
+
+ # Verify all expected header parts are present
+ for header_part in expected_headers_contains:
+ assert header_part in config.otlp_auth_headers
+
+
+@pytest.mark.parametrize(
+ "missing_var",
+ [
+ pytest.param("LEVOAI_API_KEY", id="missing API key"),
+ pytest.param("LEVOAI_ORG_ID", id="missing org ID"),
+ pytest.param("LEVOAI_WORKSPACE_ID", id="missing workspace ID"),
+ pytest.param("LEVOAI_COLLECTOR_URL", id="missing collector URL"),
+ ],
+)
+def test_get_levo_config_missing_required_vars(monkeypatch, missing_var):
+ """Test that missing required environment variables raise ValueError."""
+ # Clear all Levo-related env vars
+ for key in [
+ "LEVOAI_API_KEY",
+ "LEVOAI_ORG_ID",
+ "LEVOAI_WORKSPACE_ID",
+ "LEVOAI_COLLECTOR_URL",
+ ]:
+ monkeypatch.delenv(key, raising=False)
+
+ # Set all required vars except the missing one
+ required_vars = {
+ "LEVOAI_API_KEY": "test-key",
+ "LEVOAI_ORG_ID": "test-org",
+ "LEVOAI_WORKSPACE_ID": "test-workspace",
+ "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai",
+ }
+ required_vars.pop(missing_var)
+
+ for key, value in required_vars.items():
+ monkeypatch.setenv(key, value)
+
+ with pytest.raises(ValueError, match=missing_var):
+ LevoLogger.get_levo_config()
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py
index a719d102a7c..a322dfe9a2b 100644
--- a/tests/test_litellm/integrations/test_custom_guardrail.py
+++ b/tests/test_litellm/integrations/test_custom_guardrail.py
@@ -498,7 +498,7 @@ class TestPassthroughCallTypeHandling:
def test_get_pre_call_type_with_allm_passthrough_route(self):
"""
Test that _get_pre_call_type correctly maps allm_passthrough_route.
-
+
This tests Fix #1: allm_passthrough_route was not being handled, causing call_type to be None.
"""
from litellm.proxy.common_request_processing import (
@@ -509,14 +509,14 @@ class TestPassthroughCallTypeHandling:
result = ProxyBaseLLMRequestProcessing._get_pre_call_type(
route_type="allm_passthrough_route"
)
-
+
# Should return allm_passthrough_route, not None
assert result == "allm_passthrough_route"
def test_get_pre_call_type_preserves_standard_mappings(self):
"""
Test that _get_pre_call_type still correctly maps standard route types.
-
+
Ensures Fix #1 didn't break existing functionality.
"""
from litellm.proxy.common_request_processing import (
@@ -536,3 +536,235 @@ class TestPassthroughCallTypeHandling:
ProxyBaseLLMRequestProcessing._get_pre_call_type(route_type="aresponses")
== "responses"
)
+
+
+class TestEventTypeLogging:
+ """Tests for event_type logging in guardrail information."""
+
+ @pytest.mark.asyncio
+ async def test_log_guardrail_information_infers_event_type_from_async_pre_call_hook(
+ self,
+ ):
+ """
+ Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.pre_call
+ from async_pre_call_hook function name.
+ """
+ from litellm.integrations.custom_guardrail import log_guardrail_information
+ from litellm.types.guardrails import GuardrailEventHooks
+
+ class TestGuardrail(CustomGuardrail):
+ def __init__(self):
+ super().__init__(
+ guardrail_name="test_event_type_guardrail",
+ event_hook=[
+ GuardrailEventHooks.pre_call,
+ GuardrailEventHooks.post_call,
+ ],
+ )
+
+ @log_guardrail_information
+ async def async_pre_call_hook(self, data: dict, **kwargs):
+ return {"result": "pre_call_executed"}
+
+ guardrail = TestGuardrail()
+ request_data = {"metadata": {}}
+
+ await guardrail.async_pre_call_hook(data=request_data)
+
+ # Check that the guardrail_mode was set to pre_call (not the full list)
+ logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
+ assert len(logged_info) == 1
+ assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
+
+ @pytest.mark.asyncio
+ async def test_log_guardrail_information_infers_event_type_from_async_post_call_success_hook(
+ self,
+ ):
+ """
+ Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call
+ from async_post_call_success_hook function name.
+ """
+ from litellm.integrations.custom_guardrail import log_guardrail_information
+ from litellm.types.guardrails import GuardrailEventHooks
+
+ class TestGuardrail(CustomGuardrail):
+ def __init__(self):
+ super().__init__(
+ guardrail_name="test_event_type_guardrail",
+ event_hook=[
+ GuardrailEventHooks.pre_call,
+ GuardrailEventHooks.post_call,
+ ],
+ )
+
+ @log_guardrail_information
+ async def async_post_call_success_hook(self, data: dict, **kwargs):
+ return {"result": "post_call_executed"}
+
+ guardrail = TestGuardrail()
+ request_data = {"metadata": {}}
+
+ await guardrail.async_post_call_success_hook(data=request_data)
+
+ # Check that the guardrail_mode was set to post_call (not the full list)
+ logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
+ assert len(logged_info) == 1
+ assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
+
+ @pytest.mark.asyncio
+ async def test_log_guardrail_information_infers_event_type_from_async_moderation_hook(
+ self,
+ ):
+ """
+ Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.during_call
+ from async_moderation_hook function name.
+ """
+ from litellm.integrations.custom_guardrail import log_guardrail_information
+ from litellm.types.guardrails import GuardrailEventHooks
+
+ class TestGuardrail(CustomGuardrail):
+ def __init__(self):
+ super().__init__(
+ guardrail_name="test_event_type_guardrail",
+ event_hook=[
+ GuardrailEventHooks.during_call,
+ GuardrailEventHooks.post_call,
+ ],
+ )
+
+ @log_guardrail_information
+ async def async_moderation_hook(self, data: dict, **kwargs):
+ return {"result": "moderation_executed"}
+
+ guardrail = TestGuardrail()
+ request_data = {"metadata": {}}
+
+ await guardrail.async_moderation_hook(data=request_data)
+
+ # Check that the guardrail_mode was set to during_call (not the full list)
+ logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
+ assert len(logged_info) == 1
+ assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.during_call
+
+ @pytest.mark.asyncio
+ async def test_log_guardrail_information_infers_event_type_from_async_post_call_streaming_hook(
+ self,
+ ):
+ """
+ Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call
+ from async_post_call_streaming_hook function name.
+ """
+ from litellm.integrations.custom_guardrail import log_guardrail_information
+ from litellm.types.guardrails import GuardrailEventHooks
+
+ class TestGuardrail(CustomGuardrail):
+ def __init__(self):
+ super().__init__(
+ guardrail_name="test_event_type_guardrail",
+ event_hook=[
+ GuardrailEventHooks.pre_call,
+ GuardrailEventHooks.post_call,
+ ],
+ )
+
+ @log_guardrail_information
+ async def async_post_call_streaming_hook(self, data: dict, **kwargs):
+ return {"result": "streaming_executed"}
+
+ guardrail = TestGuardrail()
+ request_data = {"metadata": {}}
+
+ await guardrail.async_post_call_streaming_hook(data=request_data)
+
+ # Check that the guardrail_mode was set to post_call (not the full list)
+ logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
+ assert len(logged_info) == 1
+ assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
+
+ @pytest.mark.asyncio
+ async def test_log_guardrail_information_returns_none_for_unknown_function_name(
+ self,
+ ):
+ """
+ Test that log_guardrail_information decorator returns None for event_type
+ when function name doesn't match known patterns, and falls back to self.event_hook.
+ """
+ from litellm.integrations.custom_guardrail import log_guardrail_information
+ from litellm.types.guardrails import GuardrailEventHooks
+
+ class TestGuardrail(CustomGuardrail):
+ def __init__(self):
+ super().__init__(
+ guardrail_name="test_event_type_guardrail",
+ event_hook=GuardrailEventHooks.pre_call,
+ )
+
+ @log_guardrail_information
+ async def some_other_hook(self, data: dict, **kwargs):
+ return {"result": "other_hook_executed"}
+
+ guardrail = TestGuardrail()
+ request_data = {"metadata": {}}
+
+ await guardrail.some_other_hook(data=request_data)
+
+ # Check that the guardrail_mode falls back to self.event_hook
+ logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
+ assert len(logged_info) == 1
+ assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
+
+ def test_add_standard_logging_uses_event_type_over_event_hook(self):
+ """
+ Test that add_standard_logging_guardrail_information_to_request_data
+ prioritizes event_type parameter over self.event_hook.
+ """
+ from litellm.types.guardrails import GuardrailEventHooks
+
+ guardrail = CustomGuardrail(
+ guardrail_name="test_guardrail",
+ event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call],
+ )
+
+ request_data = {"metadata": {}}
+
+ # Call with explicit event_type
+ guardrail.add_standard_logging_guardrail_information_to_request_data(
+ guardrail_json_response={"result": "ok"},
+ request_data=request_data,
+ guardrail_status="success",
+ event_type=GuardrailEventHooks.post_call,
+ )
+
+ # Should use the provided event_type (post_call), not the full event_hook list
+ logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
+ assert len(logged_info) == 1
+ assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
+
+ def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none(
+ self,
+ ):
+ """
+ Test that add_standard_logging_guardrail_information_to_request_data
+ falls back to self.event_hook when event_type is None.
+ """
+ from litellm.types.guardrails import GuardrailEventHooks
+
+ guardrail = CustomGuardrail(
+ guardrail_name="test_guardrail",
+ event_hook=GuardrailEventHooks.pre_call,
+ )
+
+ request_data = {"metadata": {}}
+
+ # Call with event_type=None
+ guardrail.add_standard_logging_guardrail_information_to_request_data(
+ guardrail_json_response={"result": "ok"},
+ request_data=request_data,
+ guardrail_status="success",
+ event_type=None,
+ )
+
+ # Should fall back to self.event_hook
+ logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
+ assert len(logged_info) == 1
+ assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py
index 58fccf5e42d..5d648f601f6 100644
--- a/tests/test_litellm/integrations/test_opentelemetry.py
+++ b/tests/test_litellm/integrations/test_opentelemetry.py
@@ -8,6 +8,7 @@ from unittest.mock import MagicMock, patch
# Adds the grandparent directory to sys.path to allow importing project modules
sys.path.insert(0, os.path.abspath("../.."))
+from opentelemetry import trace
from opentelemetry.sdk._logs import LoggerProvider as OTLoggerProvider
from opentelemetry.sdk._logs.export import InMemoryLogExporter, SimpleLogRecordProcessor
from opentelemetry.sdk.metrics import MeterProvider
@@ -1318,3 +1319,456 @@ class TestOpenTelemetryProtocolSelection(unittest.TestCase):
"http://collector:4317/v1/traces", "logs"
)
self.assertEqual(normalized, "http://collector:4317/v1/logs")
+
+
+class TestOpenTelemetryExternalSpan(unittest.TestCase):
+ """
+ Test suite for external span handling in OpenTelemetry integration.
+
+ These tests verify that LiteLLM correctly handles spans created outside
+ of LiteLLM (e.g., by Langfuse SDK, user application code, or global context)
+ without closing them prematurely.
+
+ Background:
+ - External spans can come from: Langfuse SDK, user code, HTTP traceparent headers, global context
+ - LiteLLM should NEVER close spans it did not create
+ - Bug: LiteLLM was reusing and closing external spans in _start_primary_span
+ """
+
+ HERE = os.path.dirname(__file__)
+
+ def setUp(self):
+ """Set up common test fixtures"""
+ self.span_exporter = InMemorySpanExporter()
+ self.tracer_provider = TracerProvider()
+ self.tracer_provider.add_span_processor(
+ SimpleSpanProcessor(self.span_exporter)
+ )
+
+ # Don't set global tracer provider - instead, get tracers directly from our provider
+ # This avoids "Overriding of current TracerProvider is not allowed" warnings
+
+ # Clear any existing spans
+ self.span_exporter.clear()
+
+ def _create_test_kwargs_and_response(self):
+ """Load test data from JSON files"""
+ with open(
+ os.path.join(self.HERE, "open_telemetry", "data", "captured_kwargs.json")
+ ) as f:
+ kwargs = json.load(f)
+
+ with open(
+ os.path.join(self.HERE, "open_telemetry", "data", "captured_response.json")
+ ) as f:
+ response_obj = json.load(f)
+
+ return kwargs, response_obj
+
+ def _get_spans_by_name(self, name):
+ """Get all spans with the given name"""
+ spans = self.span_exporter.get_finished_spans()
+ return [s for s in spans if s.name == name]
+
+ @patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "false"}, clear=False)
+ def test_external_span_not_closed_with_use_otel_litellm_request_span_false(self):
+ """
+ Test that external spans are not closed when USE_OTEL_LITELLM_REQUEST_SPAN=false (default).
+
+ Expected behavior:
+ - External span remains open (is_recording = True)
+ - raw_gen_ai_request spans are direct children of external span (shallow hierarchy)
+ - No litellm_request span is created
+ - Multiple completions work correctly
+ """
+ # Initialize OpenTelemetry
+ otel = OpenTelemetry(tracer_provider=self.tracer_provider)
+
+ # Load test data
+ kwargs, response_obj = self._create_test_kwargs_and_response()
+
+ # Create external parent span using our test TracerProvider
+ tracer = self.tracer_provider.get_tracer(__name__)
+
+ with tracer.start_as_current_span("external_parent_span") as parent_span:
+ parent_ctx = parent_span.get_span_context()
+ parent_trace_id = parent_ctx.trace_id
+ parent_span_id = parent_ctx.span_id
+
+ self.assertTrue(
+ parent_span.is_recording(),
+ "External span should be recording before completion calls"
+ )
+
+ # First completion call
+ start_time = datetime.utcnow()
+ end_time = start_time + timedelta(seconds=1)
+ otel._handle_success(kwargs, response_obj, start_time, end_time)
+
+ # Verify parent span is still recording
+ self.assertTrue(
+ parent_span.is_recording(),
+ "External span should still be recording after first completion"
+ )
+
+ # Second completion call
+ start_time2 = end_time
+ end_time2 = start_time2 + timedelta(seconds=1)
+ otel._handle_success(kwargs, response_obj, start_time2, end_time2)
+
+ # Verify parent span is still recording
+ self.assertTrue(
+ parent_span.is_recording(),
+ "External span should still be recording after second completion"
+ )
+
+ # After exiting context, verify spans
+ spans = self.span_exporter.get_finished_spans()
+
+ # All spans should have the same trace_id
+ for span in spans:
+ self.assertEqual(
+ span.context.trace_id,
+ parent_trace_id,
+ f"Span {span.name} should have same trace_id as parent"
+ )
+
+ # Should have external_parent_span
+ parent_spans = self._get_spans_by_name("external_parent_span")
+ self.assertEqual(len(parent_spans), 1, "Should have exactly one external_parent_span")
+
+ # Verify LiteLLM set attributes on external parent span
+ parent_span_finished = parent_spans[0]
+ self.assertIsNotNone(
+ parent_span_finished.attributes,
+ "Parent span should have attributes set by LiteLLM"
+ )
+ self.assertIn(
+ "gen_ai.request.model",
+ parent_span_finished.attributes,
+ "Parent span should have model attribute from LiteLLM"
+ )
+
+ # Should have raw_gen_ai_request spans (if message_logging is on)
+ raw_spans = self._get_spans_by_name("raw_gen_ai_request")
+ # Note: May be 0 if message_logging is off, or 2 if on
+
+ # Should NOT have litellm_request spans (USE_OTEL_LITELLM_REQUEST_SPAN=false)
+ litellm_spans = self._get_spans_by_name("litellm_request")
+ self.assertEqual(
+ len(litellm_spans),
+ 0,
+ "Should NOT have litellm_request spans when USE_OTEL_LITELLM_REQUEST_SPAN=false"
+ )
+
+ # Verify raw_gen_ai_request spans are direct children of external span
+ for raw_span in raw_spans:
+ self.assertEqual(
+ raw_span.parent.span_id if raw_span.parent else None,
+ parent_span_id,
+ f"raw_gen_ai_request should be direct child of external_parent_span"
+ )
+
+ @patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "true"}, clear=False)
+ def test_external_span_not_closed_with_use_otel_litellm_request_span_true(self):
+ """
+ Test that external spans are not closed when USE_OTEL_LITELLM_REQUEST_SPAN=true.
+
+ Expected behavior:
+ - External span remains open (is_recording = True)
+ - litellm_request spans are created as children of external span
+ - raw_gen_ai_request spans are children of litellm_request spans
+ - Correct hierarchy: external_parent → litellm_request → raw_gen_ai_request
+ """
+ # Initialize OpenTelemetry
+ otel = OpenTelemetry(tracer_provider=self.tracer_provider)
+
+ # Load test data
+ kwargs, response_obj = self._create_test_kwargs_and_response()
+
+ # Create external parent span using our test TracerProvider
+ tracer = self.tracer_provider.get_tracer(__name__)
+
+ with tracer.start_as_current_span("external_parent_span") as parent_span:
+ parent_ctx = parent_span.get_span_context()
+ parent_trace_id = parent_ctx.trace_id
+ parent_span_id = parent_ctx.span_id
+
+ # First completion call
+ start_time = datetime.utcnow()
+ end_time = start_time + timedelta(seconds=1)
+ otel._handle_success(kwargs, response_obj, start_time, end_time)
+
+ # Verify parent span is still recording
+ self.assertTrue(
+ parent_span.is_recording(),
+ "External span should still be recording after first completion"
+ )
+
+ # Second completion call
+ start_time2 = end_time
+ end_time2 = start_time2 + timedelta(seconds=1)
+ otel._handle_success(kwargs, response_obj, start_time2, end_time2)
+
+ # Verify parent span is still recording
+ self.assertTrue(
+ parent_span.is_recording(),
+ "External span should still be recording after second completion"
+ )
+
+ # After exiting context, verify spans
+ spans = self.span_exporter.get_finished_spans()
+
+ # All spans should have the same trace_id
+ for span in spans:
+ self.assertEqual(
+ span.context.trace_id,
+ parent_trace_id,
+ f"Span {span.name} should have same trace_id as parent"
+ )
+
+ # Should have litellm_request spans (USE_OTEL_LITELLM_REQUEST_SPAN=true)
+ litellm_spans = self._get_spans_by_name("litellm_request")
+ self.assertEqual(
+ len(litellm_spans),
+ 2,
+ "Should have 2 litellm_request spans when USE_OTEL_LITELLM_REQUEST_SPAN=true"
+ )
+
+ # Verify litellm_request spans are children of external span
+ for litellm_span in litellm_spans:
+ self.assertEqual(
+ litellm_span.parent.span_id if litellm_span.parent else None,
+ parent_span_id,
+ "litellm_request should be child of external_parent_span"
+ )
+
+ # Verify raw_gen_ai_request spans (if present) are children of litellm_request
+ raw_spans = self._get_spans_by_name("raw_gen_ai_request")
+ if raw_spans:
+ litellm_span_ids = {s.context.span_id for s in litellm_spans}
+ for raw_span in raw_spans:
+ self.assertIn(
+ raw_span.parent.span_id if raw_span.parent else None,
+ litellm_span_ids,
+ "raw_gen_ai_request should be child of litellm_request"
+ )
+
+ @patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "false"}, clear=False)
+ def test_external_span_with_multiple_completions(self):
+ """
+ Test that multiple completion calls work correctly within external span context.
+
+ Expected behavior:
+ - Both completion calls succeed
+ - All spans belong to the same trace
+ - External span remains open throughout
+ - No errors or warnings about "ended span"
+ """
+ # Initialize OpenTelemetry
+ otel = OpenTelemetry(tracer_provider=self.tracer_provider)
+
+ # Load test data
+ kwargs, response_obj = self._create_test_kwargs_and_response()
+
+ # Create external parent span using our test TracerProvider
+ tracer = self.tracer_provider.get_tracer(__name__)
+
+ with tracer.start_as_current_span("external_parent_span") as parent_span:
+ parent_ctx = parent_span.get_span_context()
+ parent_trace_id = parent_ctx.trace_id
+
+ # Make multiple completion calls
+ for i in range(3):
+ start_time = datetime.utcnow()
+ end_time = start_time + timedelta(seconds=1)
+
+ # This should not raise any exceptions
+ otel._handle_success(kwargs, response_obj, start_time, end_time)
+
+ # Verify parent span is still recording after each call
+ self.assertTrue(
+ parent_span.is_recording(),
+ f"External span should still be recording after completion #{i+1}"
+ )
+
+ # Verify all spans have the same trace_id
+ spans = self.span_exporter.get_finished_spans()
+ for span in spans:
+ self.assertEqual(
+ span.context.trace_id,
+ parent_trace_id,
+ f"All spans should belong to the same trace"
+ )
+
+ # Should have the external parent span
+ parent_spans = self._get_spans_by_name("external_parent_span")
+ self.assertEqual(len(parent_spans), 1, "Should have exactly one external_parent_span")
+
+ # Verify LiteLLM set attributes on external parent span
+ parent_span_finished = parent_spans[0]
+ self.assertIn(
+ "gen_ai.request.model",
+ parent_span_finished.attributes,
+ "Parent span should have model attribute from LiteLLM"
+ )
+
+ @patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "false"}, clear=False)
+ def test_external_span_from_global_context(self):
+ """
+ Test external span detection from global context (Priority 3 in _get_span_context).
+
+ This simulates the case where a span is set in the global context
+ (e.g., by user code or Langfuse SDK) and LiteLLM detects it via
+ trace.get_current_span().
+
+ Expected behavior:
+ - LiteLLM detects the span from global context
+ - External span is not closed
+ - Correct parent-child relationship
+ """
+ # Initialize OpenTelemetry
+ otel = OpenTelemetry(tracer_provider=self.tracer_provider)
+
+ # Load test data
+ kwargs, response_obj = self._create_test_kwargs_and_response()
+
+ # Create external parent span and set it as current using our test TracerProvider
+ tracer = self.tracer_provider.get_tracer(__name__)
+
+ with tracer.start_as_current_span("external_global_span") as parent_span:
+ parent_ctx = parent_span.get_span_context()
+ parent_trace_id = parent_ctx.trace_id
+
+ # Verify the span is in global context
+ current_span = trace.get_current_span()
+ self.assertEqual(current_span, parent_span, "Span should be in global context")
+
+ # Make completion call
+ start_time = datetime.utcnow()
+ end_time = start_time + timedelta(seconds=1)
+ otel._handle_success(kwargs, response_obj, start_time, end_time)
+
+ # Verify parent span is still recording
+ self.assertTrue(
+ parent_span.is_recording(),
+ "External span from global context should not be closed"
+ )
+
+ # Verify trace structure
+ spans = self.span_exporter.get_finished_spans()
+ for span in spans:
+ self.assertEqual(
+ span.context.trace_id,
+ parent_trace_id,
+ "All spans should have the same trace_id"
+ )
+
+ @patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "false"}, clear=False)
+ def test_external_span_hierarchy_preserved(self):
+ """
+ Test that span hierarchy is correctly preserved with external parent.
+
+ Expected behavior:
+ - Parent span IDs are correct
+ - Trace structure matches expected hierarchy
+ - Span names are correct
+ """
+ # Initialize OpenTelemetry
+ otel = OpenTelemetry(tracer_provider=self.tracer_provider)
+ otel.message_logging = True # Enable message logging to get raw_gen_ai_request spans
+
+ # Load test data
+ kwargs, response_obj = self._create_test_kwargs_and_response()
+
+ # Create external parent span using our test TracerProvider
+ tracer = self.tracer_provider.get_tracer(__name__)
+
+ with tracer.start_as_current_span("external_parent_span") as parent_span:
+ parent_span_id = parent_span.get_span_context().span_id
+
+ # Make completion call
+ start_time = datetime.utcnow()
+ end_time = start_time + timedelta(seconds=1)
+ otel._handle_success(kwargs, response_obj, start_time, end_time)
+
+ # Verify hierarchy
+ spans = self.span_exporter.get_finished_spans()
+
+ # Get spans by name
+ parent_spans = self._get_spans_by_name("external_parent_span")
+ raw_spans = self._get_spans_by_name("raw_gen_ai_request")
+
+ self.assertEqual(len(parent_spans), 1, "Should have one parent span")
+
+ # Verify parent-child relationship
+ if raw_spans: # If message_logging is on
+ for raw_span in raw_spans:
+ self.assertEqual(
+ raw_span.parent.span_id if raw_span.parent else None,
+ parent_span_id,
+ "raw_gen_ai_request should be child of external_parent_span"
+ )
+
+ @patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "false"}, clear=False)
+ def test_external_span_not_ended_on_failure(self):
+ """
+ Test that external spans are not closed even on failure.
+
+ Expected behavior:
+ - When _handle_failure is called with external span context
+ - External span remains open (is_recording = True)
+ - Error span is created correctly
+ - External span status is NOT changed by LiteLLM
+ """
+ # Initialize OpenTelemetry
+ otel = OpenTelemetry(tracer_provider=self.tracer_provider)
+
+ # Load test data
+ kwargs, response_obj = self._create_test_kwargs_and_response()
+
+ # Create external parent span using our test TracerProvider
+ tracer = self.tracer_provider.get_tracer(__name__)
+
+ with tracer.start_as_current_span("external_parent_span") as parent_span:
+ parent_ctx = parent_span.get_span_context()
+ parent_trace_id = parent_ctx.trace_id
+
+ # Simulate failure
+ start_time = datetime.utcnow()
+ end_time = start_time + timedelta(seconds=1)
+
+ # Create error response object
+ error_response = {"error": "Test error"}
+
+ # Call _handle_failure
+ otel._handle_failure(kwargs, error_response, start_time, end_time)
+
+ # Verify parent span is still recording
+ self.assertTrue(
+ parent_span.is_recording(),
+ "External span should still be recording even after failure"
+ )
+
+ # Verify trace structure
+ spans = self.span_exporter.get_finished_spans()
+
+ # All spans should have the same trace_id
+ for span in spans:
+ self.assertEqual(
+ span.context.trace_id,
+ parent_trace_id,
+ "All spans should have the same trace_id even on failure"
+ )
+
+ # Should have external_parent_span
+ parent_spans = self._get_spans_by_name("external_parent_span")
+ self.assertEqual(len(parent_spans), 1, "Should have exactly one external_parent_span")
+
+ # Verify LiteLLM set attributes on external parent span even on failure
+ parent_span_finished = parent_spans[0]
+ self.assertIn(
+ "gen_ai.request.model",
+ parent_span_finished.attributes,
+ "Parent span should have model attribute from LiteLLM even on failure"
+ )
diff --git a/tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py b/tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py
new file mode 100644
index 00000000000..9658eff3cc5
--- /dev/null
+++ b/tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py
@@ -0,0 +1,106 @@
+"""
+Unit tests for prometheus metric name consistency
+
+This test ensures that the metric names used when creating Prometheus metrics
+match the names defined in DEFINED_PROMETHEUS_METRICS, so that metric filtering
+configuration works correctly.
+
+Related issue: https://github.com/BerriAI/litellm/issues/18221
+"""
+from typing import get_args
+
+import pytest
+
+
+def test_remaining_requests_metric_name_in_defined_metrics():
+ """
+ Test that litellm_remaining_requests_metric is defined in DEFINED_PROMETHEUS_METRICS.
+
+ The metric name should include the _metric suffix to be consistent with the
+ configuration format users specify in prometheus_metrics_config.
+ """
+ from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS
+
+ defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS)
+ assert (
+ "litellm_remaining_requests_metric" in defined_metrics
+ ), "litellm_remaining_requests_metric should be in DEFINED_PROMETHEUS_METRICS"
+
+
+def test_remaining_tokens_metric_name_in_defined_metrics():
+ """
+ Test that litellm_remaining_tokens_metric is defined in DEFINED_PROMETHEUS_METRICS.
+
+ The metric name should include the _metric suffix to be consistent with the
+ configuration format users specify in prometheus_metrics_config.
+ """
+ from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS
+
+ defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS)
+ assert (
+ "litellm_remaining_tokens_metric" in defined_metrics
+ ), "litellm_remaining_tokens_metric should be in DEFINED_PROMETHEUS_METRICS"
+
+
+def test_prometheus_metric_labels_have_remaining_metrics():
+ """
+ Test that PrometheusMetricLabels has label definitions for remaining metrics.
+
+ This ensures that the labels can be retrieved when creating the metrics.
+ """
+ from litellm.types.integrations.prometheus import PrometheusMetricLabels
+
+ # Test that labels can be retrieved for remaining metrics
+ remaining_requests_labels = PrometheusMetricLabels.get_labels(
+ "litellm_remaining_requests_metric"
+ )
+ remaining_tokens_labels = PrometheusMetricLabels.get_labels(
+ "litellm_remaining_tokens_metric"
+ )
+
+ assert isinstance(
+ remaining_requests_labels, list
+ ), "Labels for litellm_remaining_requests_metric should be a list"
+ assert isinstance(
+ remaining_tokens_labels, list
+ ), "Labels for litellm_remaining_tokens_metric should be a list"
+
+ # These metrics should have api_provider and api_base labels
+ assert (
+ "api_provider" in remaining_requests_labels
+ ), "litellm_remaining_requests_metric should have api_provider label"
+ assert (
+ "api_base" in remaining_requests_labels
+ ), "litellm_remaining_requests_metric should have api_base label"
+ assert (
+ "api_provider" in remaining_tokens_labels
+ ), "litellm_remaining_tokens_metric should have api_provider label"
+ assert (
+ "api_base" in remaining_tokens_labels
+ ), "litellm_remaining_tokens_metric should have api_base label"
+
+
+def test_all_defined_metrics_have_consistent_naming():
+ """
+ Test that all metrics defined in DEFINED_PROMETHEUS_METRICS follow
+ a consistent naming convention.
+
+ This helps prevent similar inconsistencies in the future.
+ """
+ from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS
+
+ defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS)
+
+ for metric_name in defined_metrics:
+ # All metrics should start with 'litellm_'
+ assert metric_name.startswith(
+ "litellm_"
+ ), f"Metric {metric_name} should start with 'litellm_'"
+
+
+if __name__ == "__main__":
+ test_remaining_requests_metric_name_in_defined_metrics()
+ test_remaining_tokens_metric_name_in_defined_metrics()
+ test_prometheus_metric_labels_have_remaining_metrics()
+ test_all_defined_metrics_have_consistent_naming()
+ print("All prometheus metric name consistency tests passed!")
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
index 65e3dbec8bd..5ba78d9eed1 100644
--- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
+++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
@@ -809,3 +809,54 @@ def test_bedrock_anthropic_prompt_caching():
assert completion_cost >= 0
assert round(prompt_cost, 3) == 0.111
assert round(completion_cost, 5) == 0.00820
+
+
+def test_reasoning_tokens_without_text_tokens_gpt5_nano():
+ """
+ Test fix for GitHub issue #18599:
+ https://github.com/BerriAI/litellm/issues/18599
+
+ When OpenAI models (gpt-5-nano, o1, o3) return reasoning_tokens but don't provide
+ text_tokens, LiteLLM should calculate text_tokens as:
+ text_tokens = completion_tokens - reasoning_tokens - audio_tokens - image_tokens
+
+ This ensures ALL completion tokens are billed, not just reasoning tokens.
+ """
+ model = "gpt-5-nano"
+ custom_llm_provider = "openai"
+
+ # Simulate OpenAI gpt-5-nano response where text_tokens is NOT provided
+ # completion_tokens: 977 total
+ # reasoning_tokens: 768
+ # text_tokens: should be calculated as 977 - 768 = 209
+ usage = Usage(
+ prompt_tokens=17,
+ completion_tokens=977,
+ total_tokens=994,
+ completion_tokens_details=CompletionTokensDetailsWrapper(
+ reasoning_tokens=768,
+ audio_tokens=0,
+ # text_tokens NOT provided - this is the key part of the bug
+ ),
+ )
+
+ prompt_cost, completion_cost = generic_cost_per_token(
+ model=model,
+ usage=usage,
+ custom_llm_provider=custom_llm_provider,
+ )
+
+ # gpt-5-nano pricing: $0.05/1M input, $0.40/1M output
+ expected_prompt_cost = 17 * 0.05 / 1_000_000
+ expected_completion_cost = 977 * 0.40 / 1_000_000 # ALL tokens, not just reasoning
+
+ assert abs(prompt_cost - expected_prompt_cost) < 1e-10, \
+ f"Prompt cost incorrect: {prompt_cost} vs {expected_prompt_cost}"
+
+ assert abs(completion_cost - expected_completion_cost) < 1e-10, \
+ f"Completion cost incorrect: {completion_cost} vs {expected_completion_cost}"
+
+ # Verify it's NOT using only reasoning_tokens (the bug)
+ wrong_cost = 768 * 0.40 / 1_000_000 # Only reasoning tokens
+ assert abs(completion_cost - wrong_cost) > 1e-6, \
+ "Bug detected: Cost calculation is using only reasoning_tokens instead of all completion_tokens!"
diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
index 41ac893b4d7..c8fe6efeaa1 100644
--- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
+++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
@@ -497,6 +497,57 @@ def test_convert_gemini_messages():
)
+def test_convert_gemini_tool_call_result_with_image_url():
+ """
+ Test that image_url content type in tool results is handled correctly for Gemini.
+ Fixes: https://github.com/BerriAI/litellm/issues/18187
+ """
+ from litellm.litellm_core_utils.prompt_templates.factory import (
+ convert_to_gemini_tool_call_result,
+ )
+ from litellm.types.llms.openai import ChatCompletionToolMessage
+
+ # Test with string image_url format
+ message_str_format = ChatCompletionToolMessage(
+ role="tool",
+ tool_call_id="call_123",
+ content=[{"type": "image_url", "image_url": "data:image/jpeg;base64,/9j/4AAQ"}],
+ )
+ last_message_with_tool_calls = {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_123",
+ "type": "function",
+ "index": 0,
+ "function": {"name": "get_image", "arguments": "{}"},
+ }
+ ],
+ }
+
+ result = convert_to_gemini_tool_call_result(
+ message=message_str_format,
+ last_message_with_tool_calls=last_message_with_tool_calls,
+ )
+ # Should have inline_data for the image
+ assert isinstance(result, list) and any("inline_data" in p for p in result)
+
+ # Test with dict image_url format (OpenAI standard)
+ message_dict_format = ChatCompletionToolMessage(
+ role="tool",
+ tool_call_id="call_456",
+ content=[{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,/9j/4AAQ"}}],
+ )
+ last_message_with_tool_calls["tool_calls"][0]["id"] = "call_456"
+
+ result2 = convert_to_gemini_tool_call_result(
+ message=message_dict_format,
+ last_message_with_tool_calls=last_message_with_tool_calls,
+ )
+ assert isinstance(result2, list) and any("inline_data" in p for p in result2)
+
+
def test_bedrock_tools_unpack_defs():
"""
Test that the unpack_defs method handles nested $ref inside anyOf items correctly
diff --git a/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py b/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py
new file mode 100644
index 00000000000..1a6ed51afd0
--- /dev/null
+++ b/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py
@@ -0,0 +1,69 @@
+"""
+Unit tests for codestral provider routing.
+
+These tests verify that the chat and FIM endpoints for codestral
+are correctly routed to different providers:
+- Chat endpoint -> codestral provider
+- FIM endpoint -> text-completion-codestral provider
+
+Related issue: https://github.com/BerriAI/litellm/issues/18464
+"""
+import pytest
+
+import litellm
+
+
+class TestCodestralProviderRouting:
+ """Tests for codestral endpoint routing in get_llm_provider"""
+
+ def test_codestral_chat_endpoint_routes_to_codestral_provider(self):
+ """
+ Test that the codestral chat endpoint routes to the 'codestral' provider.
+
+ The chat/completions endpoint should be handled by the codestral provider.
+ """
+ model, custom_llm_provider, _, api_base = litellm.get_llm_provider(
+ model="codestral-latest",
+ api_base="https://codestral.mistral.ai/v1/chat/completions",
+ )
+
+ assert custom_llm_provider == "codestral"
+
+ def test_codestral_fim_endpoint_routes_to_text_completion_provider(self):
+ """
+ Test that the codestral FIM endpoint routes to 'text-completion-codestral'.
+
+ The fim/completions endpoint should be handled by the
+ text-completion-codestral provider for fill-in-the-middle completions.
+ """
+ model, custom_llm_provider, _, api_base = litellm.get_llm_provider(
+ model="codestral-latest",
+ api_base="https://codestral.mistral.ai/v1/fim/completions",
+ )
+
+ assert custom_llm_provider == "text-completion-codestral"
+
+ def test_codestral_endpoints_are_different_providers(self):
+ """
+ Test that chat and FIM endpoints route to different providers.
+
+ This is the core fix for issue #18464 - previously both endpoints
+ would route to 'codestral' due to duplicate conditions.
+ """
+ _, chat_provider, _, _ = litellm.get_llm_provider(
+ model="codestral-latest",
+ api_base="https://codestral.mistral.ai/v1/chat/completions",
+ )
+
+ _, fim_provider, _, _ = litellm.get_llm_provider(
+ model="codestral-latest",
+ api_base="https://codestral.mistral.ai/v1/fim/completions",
+ )
+
+ assert chat_provider != fim_provider
+ assert chat_provider == "codestral"
+ assert fim_provider == "text-completion-codestral"
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py
index 89cc11c40d2..cd9c401143e 100644
--- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py
+++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py
@@ -1,173 +1,45 @@
-import json
-import os
-import sys
-from unittest.mock import MagicMock, patch
+"""Tests for litellm_core_utils.core_helpers module."""
-import pytest
-
-sys.path.insert(
- 0, os.path.abspath("../../..")
-) # Adds the parent directory to the system path
-
-from litellm.litellm_core_utils.core_helpers import (
- get_litellm_metadata_from_kwargs,
- safe_divide,
- safe_deep_copy
-)
+from litellm.litellm_core_utils.core_helpers import reconstruct_model_name
-def test_get_litellm_metadata_from_kwargs():
- kwargs = {
- "litellm_params": {
- "litellm_metadata": {},
- "metadata": {"user_api_key": "1234567890"},
- },
- }
- assert get_litellm_metadata_from_kwargs(kwargs) == {"user_api_key": "1234567890"}
+def test_reconstruct_model_name_prefers_deployment_value():
+ """Ensure deployment metadata wins when reconstructing the model name."""
+ metadata = {"deployment": "vertex_ai/gemini-1.5-flash"}
-def test_add_missing_spend_metadata_to_litellm_metadata():
- litellm_metadata = {"test_key": "test_value"}
- metadata = {"user_api_key_hash_value": "1234567890"}
- kwargs = {
- "litellm_params": {
- "litellm_metadata": litellm_metadata,
- "metadata": metadata,
- },
- }
- assert get_litellm_metadata_from_kwargs(kwargs) == {
- "test_key": "test_value",
- "user_api_key_hash_value": "1234567890",
- }
-
-
-def test_preserve_upstream_non_openai_attributes():
- from litellm.litellm_core_utils.core_helpers import (
- preserve_upstream_non_openai_attributes,
- )
- from litellm.types.utils import ModelResponseStream
-
- model_response = ModelResponseStream(
- id="123",
- object="text_completion",
- created=1715811200,
- model="gpt-3.5-turbo",
+ result = reconstruct_model_name(
+ model_name="gemini-1.5-flash",
+ custom_llm_provider="vertex_ai",
+ metadata=metadata,
)
- setattr(model_response, "test_key", "test_value")
- preserve_upstream_non_openai_attributes(
- model_response=ModelResponseStream(),
- original_chunk=model_response,
+ assert result == "vertex_ai/gemini-1.5-flash"
+
+
+def test_reconstruct_model_name_adds_bedrock_prefix_when_missing():
+ """Bedrock model names without prefixes should gain the provider prefix."""
+
+ metadata = {}
+
+ result = reconstruct_model_name(
+ model_name="us.anthropic.claude-3-sonnet",
+ custom_llm_provider="bedrock",
+ metadata=metadata,
)
- assert model_response.test_key == "test_value"
+ assert result == "bedrock/us.anthropic.claude-3-sonnet"
-def test_safe_divide_basic():
- """Test basic safe division functionality"""
- # Normal division
- result = safe_divide(10, 2)
- assert result == 5.0, f"Expected 5.0, got {result}"
-
- # Division with float
- result = safe_divide(7.5, 2.5)
- assert result == 3.0, f"Expected 3.0, got {result}"
-
- # Division by zero with default
- result = safe_divide(10, 0)
- assert result == 0, f"Expected 0, got {result}"
-
- # Division by zero with custom default
- result = safe_divide(10, 0, default=1)
- assert result == 1, f"Expected 1, got {result}"
-
- # Division by zero with custom default as float
- result = safe_divide(10, 0, default=0.5)
- assert result == 0.5, f"Expected 0.5, got {result}"
+def test_reconstruct_model_name_returns_original_for_other_providers():
+ """Non-Bedrock providers should not prepend anything."""
+ metadata = {}
-def test_safe_divide_edge_cases():
- """Test edge cases for safe division"""
- # Zero numerator
- result = safe_divide(0, 5)
- assert result == 0.0, f"Expected 0.0, got {result}"
-
- # Negative numbers
- result = safe_divide(-10, 2)
- assert result == -5.0, f"Expected -5.0, got {result}"
-
- # Negative denominator
- result = safe_divide(10, -2)
- assert result == -5.0, f"Expected -5.0, got {result}"
-
- # Both negative
- result = safe_divide(-10, -2)
- assert result == 5.0, f"Expected 5.0, got {result}"
-
- # Float division
- result = safe_divide(1, 3)
- assert abs(result - 0.3333333333333333) < 1e-10, f"Expected ~0.333..., got {result}"
+ result = reconstruct_model_name(
+ model_name="claude-3-sonnet",
+ custom_llm_provider="anthropic",
+ metadata=metadata,
+ )
-
-def test_safe_divide_weight_scenario():
- """Test safe division in the context of weight calculations"""
- # Simulate weight calculation scenario
- weights = [3, 7, 0, 2]
- total_weight = sum(weights) # 12
-
- # Normal case
- normalized_weights = [safe_divide(w, total_weight) for w in weights]
- expected = [0.25, 7/12, 0.0, 1/6]
-
- for i, (actual, exp) in enumerate(zip(normalized_weights, expected)):
- assert abs(actual - exp) < 1e-10, f"Weight {i}: Expected {exp}, got {actual}"
-
- # Zero total weight scenario (division by zero)
- zero_weights = [0, 0, 0]
- zero_total = sum(zero_weights) # 0
-
- # Should return default values (0) for all weights
- normalized_zero_weights = [safe_divide(w, zero_total) for w in zero_weights]
- expected_zero = [0, 0, 0]
-
- assert normalized_zero_weights == expected_zero, f"Expected {expected_zero}, got {normalized_zero_weights}"
-
-
-def test_safe_deep_copy_with_non_pickleables_and_span():
- """
- Verify safe_deep_copy:
- - does not crash when non-pickleables are present,
- - preserves structure/keys,
- - deep-copies JSON-y payloads (e.g., messages),
- - keeps non-pickleables by reference,
- - redacts OTEL span in the copy and restores it in the original.
- """
- import threading
- rlock = threading.RLock()
- data = {
- "metadata": {"litellm_parent_otel_span": rlock, "x": 1},
- "messages": [{"role": "user", "content": "hi"}],
- "optional_params": {"handle": rlock},
- "ok": True,
- }
-
- copied = safe_deep_copy(data)
-
- # Structure preserved
- assert set(copied.keys()) == set(data.keys())
-
- # Messages are deep-copied (new object, same content)
- assert copied["messages"] is not data["messages"]
- assert copied["messages"][0] == data["messages"][0]
-
- # Non-pickleable subtree kept by reference (no crash)
- assert copied["optional_params"] is data["optional_params"]
- assert copied["optional_params"]["handle"] is rlock
-
- # OTEL span: redacted in the copy, restored in original
- assert copied["metadata"]["litellm_parent_otel_span"] == "placeholder"
- assert data["metadata"]["litellm_parent_otel_span"] is rlock
-
- # Other simple fields unchanged
- assert copied["ok"] is True
- assert copied["metadata"]["x"] == 1
+ assert result == "claude-3-sonnet"
diff --git a/tests/test_litellm/litellm_core_utils/test_dot_notation_indexing.py b/tests/test_litellm/litellm_core_utils/test_dot_notation_indexing.py
new file mode 100644
index 00000000000..6940a5ea7a5
--- /dev/null
+++ b/tests/test_litellm/litellm_core_utils/test_dot_notation_indexing.py
@@ -0,0 +1,138 @@
+"""
+Tests for litellm.litellm_core_utils.dot_notation_indexing module.
+"""
+
+import pytest
+
+from litellm.litellm_core_utils.dot_notation_indexing import (
+ get_nested_value,
+ delete_nested_value,
+)
+
+
+class TestGetNestedValue:
+ """Tests for get_nested_value function."""
+
+ def test_simple_key(self):
+ """Test accessing a simple top-level key."""
+ data = {"name": "test"}
+ assert get_nested_value(data, "name") == "test"
+
+ def test_nested_key(self):
+ """Test accessing nested keys with dot notation."""
+ data = {"a": {"b": {"c": "value"}}}
+ assert get_nested_value(data, "a.b.c") == "value"
+
+ def test_missing_key_returns_default(self):
+ """Test that missing keys return the default value."""
+ data = {"a": {"b": "value"}}
+ assert get_nested_value(data, "a.b", "default") == "value"
+ assert get_nested_value(data, "a.c", "default") == "default"
+ assert get_nested_value(data, "x.y.z") is None
+
+ def test_empty_key_path(self):
+ """Test that empty key path returns default."""
+ data = {"a": "value"}
+ assert get_nested_value(data, "") is None
+ assert get_nested_value(data, "", "default") == "default"
+
+ def test_metadata_prefix_removal(self):
+ """Test that metadata. prefix is properly removed."""
+ data = {"user": {"email": "test@example.com"}}
+ assert get_nested_value(data, "metadata.user.email") == "test@example.com"
+
+ def test_escaped_dot_in_key(self):
+ """Test accessing keys that contain dots using escape sequence."""
+ data = {"kubernetes.io": {"namespace": "default"}}
+ assert get_nested_value(data, "kubernetes\\.io.namespace") == "default"
+
+ def test_escaped_dot_nested(self):
+ """Test multiple levels with escaped dots."""
+ data = {
+ "kubernetes.io": {
+ "pod.info": {
+ "name": "my-pod"
+ }
+ }
+ }
+ assert get_nested_value(data, "kubernetes\\.io.pod\\.info.name") == "my-pod"
+
+ def test_kubernetes_jwt_example(self):
+ """Test with a realistic Kubernetes JWT structure."""
+ jwt_token = {
+ "aud": ["https://kubernetes.default.svc"],
+ "exp": "1234567890",
+ "iat": "123456789",
+ "iss": "https://oidc.eks.region.amazonaws.com/id/randomstring",
+ "jti": "randomstring",
+ "kubernetes.io": {
+ "namespace": "namespace",
+ "node": {
+ "name": "node-name",
+ "uid": "node-uid"
+ },
+ "pod": {
+ "name": "pod-name",
+ "uid": "pod-uid"
+ },
+ "serviceaccount": {
+ "name": "serviceaccount-name",
+ "uid": "serviceaccount-uid"
+ },
+ "warnafter": 1234567880
+ },
+ "nbf": 123456789,
+ "sub": "system:serviceaccount:namespace:serviceaccount-name"
+ }
+
+ # Test accessing kubernetes.io.namespace
+ assert get_nested_value(jwt_token, "kubernetes\\.io.namespace") == "namespace"
+
+ # Test accessing nested values within kubernetes.io
+ assert get_nested_value(jwt_token, "kubernetes\\.io.pod.name") == "pod-name"
+ assert get_nested_value(jwt_token, "kubernetes\\.io.serviceaccount.name") == "serviceaccount-name"
+
+ # Test accessing regular keys still works
+ assert get_nested_value(jwt_token, "sub") == "system:serviceaccount:namespace:serviceaccount-name"
+
+ def test_mixed_escaped_and_regular_dots(self):
+ """Test path with both escaped dots (in keys) and regular dots (separators)."""
+ data = {
+ "config.v1": {
+ "settings": {
+ "feature.enabled": True
+ }
+ }
+ }
+ assert get_nested_value(data, "config\\.v1.settings.feature\\.enabled") is True
+
+
+class TestDeleteNestedValue:
+ """Tests for delete_nested_value function."""
+
+ def test_delete_simple_key(self):
+ """Test deleting a simple top-level key."""
+ data = {"a": 1, "b": 2}
+ result = delete_nested_value(data, "a")
+ assert result == {"b": 2}
+ # Original should be unchanged
+ assert data == {"a": 1, "b": 2}
+
+ def test_delete_nested_key(self):
+ """Test deleting a nested key."""
+ data = {"a": {"b": {"c": 1, "d": 2}}}
+ result = delete_nested_value(data, "a.b.c")
+ assert result == {"a": {"b": {"d": 2}}}
+
+ def test_delete_array_wildcard(self):
+ """Test deleting a field from all array elements."""
+ data = {"tools": [{"name": "t1", "secret": "s1"}, {"name": "t2", "secret": "s2"}]}
+ result = delete_nested_value(data, "tools[*].secret")
+ assert result == {"tools": [{"name": "t1"}, {"name": "t2"}]}
+
+ def test_delete_array_index(self):
+ """Test deleting a field from a specific array element."""
+ data = {"items": [{"a": 1, "b": 2}, {"a": 3, "b": 4}]}
+ result = delete_nested_value(data, "items[0].b")
+ assert result == {"items": [{"a": 1}, {"a": 3, "b": 4}]}
+
diff --git a/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py b/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py
new file mode 100644
index 00000000000..b17c02d7006
--- /dev/null
+++ b/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py
@@ -0,0 +1,156 @@
+"""
+Unit tests for _extract_base64_data and extract_images_from_message functions.
+
+These tests verify that base64 image data is correctly extracted from data URLs,
+which fixes the Ollama error "illegal base64 data at input byte 4".
+
+Related issue: https://github.com/BerriAI/litellm/issues/18338
+"""
+import pytest
+
+from litellm.litellm_core_utils.prompt_templates.common_utils import (
+ _extract_base64_data,
+ extract_images_from_message,
+)
+
+
+class TestExtractBase64Data:
+ """Tests for _extract_base64_data function"""
+
+ def test_extract_base64_from_png_data_url(self):
+ """Test extracting base64 data from a PNG data URL"""
+ data_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk"
+ expected = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk"
+ assert _extract_base64_data(data_url) == expected
+
+ def test_extract_base64_from_jpeg_data_url(self):
+ """Test extracting base64 data from a JPEG data URL"""
+ data_url = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD"
+ expected = "/9j/4AAQSkZJRgABAQAAAQABAAD"
+ assert _extract_base64_data(data_url) == expected
+
+ def test_extract_base64_from_gif_data_url(self):
+ """Test extracting base64 data from a GIF data URL"""
+ data_url = "data:image/gif;base64,R0lGODlhAQABAIAAAAAAAP"
+ expected = "R0lGODlhAQABAIAAAAAAAP"
+ assert _extract_base64_data(data_url) == expected
+
+ def test_regular_url_unchanged(self):
+ """Test that regular HTTP URLs are returned unchanged"""
+ url = "https://example.com/image.png"
+ assert _extract_base64_data(url) == url
+
+ def test_file_path_unchanged(self):
+ """Test that file paths are returned unchanged"""
+ path = "/path/to/image.png"
+ assert _extract_base64_data(path) == path
+
+ def test_data_url_without_base64_unchanged(self):
+ """Test that data URLs without base64 encoding are returned unchanged"""
+ # This is a data URL with URL encoding, not base64
+ url = "data:text/plain,Hello%20World"
+ assert _extract_base64_data(url) == url
+
+ def test_base64_data_with_special_chars(self):
+ """Test extracting base64 data that contains valid special characters"""
+ # Base64 can contain +, /, and = characters
+ data_url = "data:image/png;base64,abc+def/ghi==="
+ expected = "abc+def/ghi==="
+ assert _extract_base64_data(data_url) == expected
+
+
+class TestExtractImagesFromMessage:
+ """Tests for extract_images_from_message function"""
+
+ def test_extract_from_message_with_data_url_string(self):
+ """Test extracting images when image_url is a string data URL"""
+ message = {
+ "role": "user",
+ "content": [
+ {
+ "type": "image_url",
+ "image_url": "data:image/png;base64,iVBORw0KGgo",
+ }
+ ],
+ }
+ result = extract_images_from_message(message)
+ assert result == ["iVBORw0KGgo"]
+
+ def test_extract_from_message_with_data_url_dict(self):
+ """Test extracting images when image_url is a dict with url key"""
+ message = {
+ "role": "user",
+ "content": [
+ {
+ "type": "image_url",
+ "image_url": {"url": "data:image/png;base64,iVBORw0KGgo"},
+ }
+ ],
+ }
+ result = extract_images_from_message(message)
+ assert result == ["iVBORw0KGgo"]
+
+ def test_extract_from_message_with_regular_url(self):
+ """Test that regular URLs are preserved"""
+ message = {
+ "role": "user",
+ "content": [
+ {
+ "type": "image_url",
+ "image_url": {"url": "https://example.com/image.png"},
+ }
+ ],
+ }
+ result = extract_images_from_message(message)
+ assert result == ["https://example.com/image.png"]
+
+ def test_extract_multiple_images(self):
+ """Test extracting multiple images from a single message"""
+ message = {
+ "role": "user",
+ "content": [
+ {
+ "type": "image_url",
+ "image_url": "data:image/png;base64,image1base64",
+ },
+ {
+ "type": "image_url",
+ "image_url": {"url": "data:image/jpeg;base64,image2base64"},
+ },
+ {
+ "type": "image_url",
+ "image_url": "https://example.com/image3.png",
+ },
+ ],
+ }
+ result = extract_images_from_message(message)
+ assert result == [
+ "image1base64",
+ "image2base64",
+ "https://example.com/image3.png",
+ ]
+
+ def test_empty_content(self):
+ """Test message with empty content"""
+ message = {"role": "user", "content": []}
+ result = extract_images_from_message(message)
+ assert result == []
+
+ def test_no_images_in_content(self):
+ """Test message with content but no images"""
+ message = {
+ "role": "user",
+ "content": [{"type": "text", "text": "Hello world"}],
+ }
+ result = extract_images_from_message(message)
+ assert result == []
+
+ def test_string_content(self):
+ """Test message with string content (no images possible)"""
+ message = {"role": "user", "content": "Hello world"}
+ result = extract_images_from_message(message)
+ assert result == []
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/tests/test_litellm/litellm_core_utils/test_logging_worker.py b/tests/test_litellm/litellm_core_utils/test_logging_worker.py
index 0fe15168467..a44b821db87 100644
--- a/tests/test_litellm/litellm_core_utils/test_logging_worker.py
+++ b/tests/test_litellm/litellm_core_utils/test_logging_worker.py
@@ -323,3 +323,40 @@ class TestLoggingWorker:
await worker.clear_queue()
assert len(processed) >= 4, f"Expected 4+ tasks processed, got {len(processed)}"
+
+ @pytest.mark.asyncio
+ async def test_event_loop_change_handling(self):
+ """Test that LoggingWorker handles event loop changes correctly.
+
+ This tests the fix for GitHub issue #17813 where asyncio.Queue
+ was bound to a different event loop when using multiprocessing.
+ """
+ worker = LoggingWorker(timeout=1.0, max_queue_size=10)
+
+ # Start the worker in the current event loop
+ worker.start()
+
+ # Verify queue was created and bound to current loop
+ assert worker._queue is not None
+ assert worker._bound_loop is not None
+ original_queue = worker._queue
+
+ await worker.stop()
+
+ # Simulate a new event loop by creating a mock scenario
+ # In a real multiprocessing case, asyncio.run() creates a new loop
+ # We test the internal state detection
+
+ # Create a new worker to test the _ensure_queue logic
+ worker2 = LoggingWorker(timeout=1.0, max_queue_size=10)
+ worker2._queue = original_queue # Pretend we have an old queue
+ worker2._bound_loop = None # No bound loop (simulates first call)
+
+ # Calling start should create a new queue since _bound_loop != current
+ worker2.start()
+
+ # The queue should be reinitialized since bound_loop was None
+ assert worker2._queue is not None
+ assert worker2._bound_loop is not None
+
+ await worker2.stop()
diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py
index d78a638fd89..bdced849c7e 100644
--- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py
+++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py
@@ -55,11 +55,12 @@ class TestAzureAnthropicMessagesConfig:
assert isinstance(call_args[1]["litellm_params"], GenericLiteLLMParams)
assert call_args[1]["litellm_params"].api_key == "test-api-key"
assert "anthropic-version" in result
- # api-key header is preserved as-is (no conversion to x-api-key)
- assert "api-key" in result
+ assert "x-api-key" in result
+ assert result["x-api-key"] == "test-api-key"
+ assert "api-key" not in result
- def test_validate_anthropic_messages_environment_preserves_api_key_header(self):
- """Test that api-key header is preserved as-is (Azure handles the header internally)"""
+ def test_validate_anthropic_messages_environment_converts_api_key_to_x_api_key(self):
+ """Test that api-key header is converted to x-api-key"""
config = AzureAnthropicMessagesConfig()
headers = {}
model = "claude-sonnet-4-5"
@@ -79,9 +80,10 @@ class TestAzureAnthropicMessagesConfig:
litellm_params=litellm_params,
)
- # Verify api-key header is preserved as-is
- assert "api-key" in result
- assert result["api-key"] == "test-api-key"
+ # Verify api-key was converted to x-api-key
+ assert "x-api-key" in result
+ assert result["x-api-key"] == "test-api-key"
+ assert "api-key" not in result
def test_validate_anthropic_messages_environment_sets_headers(self):
"""Test that required headers are set"""
@@ -108,8 +110,7 @@ class TestAzureAnthropicMessagesConfig:
assert result["anthropic-version"] == "2023-06-01"
assert "content-type" in result
assert result["content-type"] == "application/json"
- # api-key header is preserved as-is
- assert "api-key" in result
+ assert "x-api-key" in result
def test_get_complete_url_with_base_url(self):
"""Test get_complete_url with base URL"""
diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py
index 2dcec689895..a8ac680908e 100644
--- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py
+++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py
@@ -8,12 +8,13 @@ forward_client_headers_to_llm_api were not being passed to Bedrock rerank provid
import json
import os
import sys
-from unittest.mock import Mock, patch
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
import pytest
sys.path.insert(0, os.path.abspath("../../../../..")) # Adds the parent directory to the system path
import litellm
+from litellm.llms.bedrock.base_aws_llm import Boto3CredentialsInfo
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
# Mock response for Bedrock rerank
@@ -47,6 +48,19 @@ test_documents = [
]
+def create_mock_credentials():
+ """Create mock AWS credentials for testing"""
+ mock_credentials = MagicMock()
+ mock_credentials.access_key = "test-access-key"
+ mock_credentials.secret_key = "test-secret-key"
+ mock_credentials.token = None
+ return Boto3CredentialsInfo(
+ credentials=mock_credentials,
+ aws_region_name="us-east-1",
+ aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
+ )
+
+
@pytest.mark.parametrize(
"model",
[
@@ -73,7 +87,17 @@ def test_bedrock_rerank_header_forwarding_sync(model):
"X-Test-Header": "test-value",
}
- with patch.object(client, "post") as mock_post:
+ # Mock AWS credentials and SigV4 auth
+ mock_credentials_info = create_mock_credentials()
+
+ with patch.object(client, "post") as mock_post, \
+ patch("litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info), \
+ patch("botocore.auth.SigV4Auth") as mock_sigv4:
+
+ # Mock SigV4Auth to not actually sign the request
+ mock_sigv4_instance = MagicMock()
+ mock_sigv4.return_value = mock_sigv4_instance
+
mock_response = Mock()
mock_response.status_code = 200
mock_response.text = json.dumps(bedrock_rerank_response)
@@ -152,9 +176,17 @@ async def test_bedrock_rerank_header_forwarding_async(model):
"X-Test-Header": "test-value",
}
- from unittest.mock import AsyncMock
+ # Mock AWS credentials and SigV4 auth
+ mock_credentials_info = create_mock_credentials()
- with patch.object(client, "post", new_callable=AsyncMock) as mock_post:
+ with patch.object(client, "post", new_callable=AsyncMock) as mock_post, \
+ patch("litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info), \
+ patch("botocore.auth.SigV4Auth") as mock_sigv4:
+
+ # Mock SigV4Auth to not actually sign the request
+ mock_sigv4_instance = MagicMock()
+ mock_sigv4.return_value = mock_sigv4_instance
+
mock_response = AsyncMock()
mock_response.status_code = 200
mock_response.text = json.dumps(bedrock_rerank_response)
@@ -223,7 +255,17 @@ def test_bedrock_rerank_extra_headers_and_headers_merge():
# Explicit extra_headers
explicit_headers = {"X-Explicit-Header": "ExplicitValue"}
- with patch.object(client, "post") as mock_post:
+ # Mock AWS credentials and SigV4 auth
+ mock_credentials_info = create_mock_credentials()
+
+ with patch.object(client, "post") as mock_post, \
+ patch("litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info), \
+ patch("botocore.auth.SigV4Auth") as mock_sigv4:
+
+ # Mock SigV4Auth to not actually sign the request
+ mock_sigv4_instance = MagicMock()
+ mock_sigv4.return_value = mock_sigv4_instance
+
mock_response = Mock()
mock_response.status_code = 200
mock_response.text = json.dumps(bedrock_rerank_response)
diff --git a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py
index b4dc6c68bb0..800066ac5bf 100644
--- a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py
+++ b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py
@@ -122,10 +122,11 @@ class TestRedactSensitiveData:
def test_redact_pat_token(self):
"""Databricks PAT tokens are redacted."""
+ test_token = "dapiTESTTOKENFAKEVALUEFORTESTINGPURPOSESONLY123"
result = DatabricksBase.redact_sensitive_data(
- "Using token dapi_fake_test_token_value"
+ f"Using token {test_token}"
)
- assert "dapi_fake_test_token_value" not in result
+ assert test_token not in result
assert "[REDACTED_PAT]" in result
def test_redact_client_secret(self):
@@ -347,17 +348,27 @@ class TestSDKPartnerTelemetry:
"Authorization": "Bearer token"
}
- with patch(
- "databricks.sdk.WorkspaceClient", return_value=mock_workspace_client
- ):
- with patch("databricks.sdk.useragent.with_partner") as mock_with_partner:
- databricks_base._get_databricks_credentials(
- api_key=None,
- api_base=None,
- headers=None,
- )
+ mock_useragent = MagicMock()
+ # Create a mock databricks.sdk module to simulate the SDK being available
+ # This allows us to test the partner telemetry registration without requiring
+ # the actual databricks-sdk package to be installed
+ mock_sdk_module = MagicMock()
+ mock_sdk_module.WorkspaceClient = MagicMock(return_value=mock_workspace_client)
+ mock_sdk_module.useragent = mock_useragent
+
+ # Mock both databricks and databricks.sdk modules to ensure the import works
+ with patch.dict(sys.modules, {
+ "databricks": MagicMock(),
+ "databricks.sdk": mock_sdk_module
+ }):
+ databricks_base._get_databricks_credentials(
+ api_key=None,
+ api_base=None,
+ headers=None,
+ )
- mock_with_partner.assert_called_once_with("litellm")
+ # Verify that partner telemetry registration was called correctly
+ mock_useragent.with_partner.assert_called_once_with("litellm")
class TestUserAgentFromEnvironment:
@@ -592,19 +603,28 @@ class TestAuthenticationPriority:
"Authorization": "Bearer sdk-token"
}
- with patch(
- "databricks.sdk.WorkspaceClient", return_value=mock_workspace_client
- ):
- with patch("databricks.sdk.useragent.with_partner"):
- api_base, headers = databricks_base.databricks_validate_environment(
- api_key=None,
- api_base=None,
- endpoint_type="chat_completions",
- custom_endpoint=False,
- headers=None,
- )
+ # Create a mock databricks.sdk module to simulate the SDK being available
+ # This allows us to test the SDK fallback authentication without requiring
+ # the actual databricks-sdk package to be installed
+ mock_sdk_module = MagicMock()
+ mock_sdk_module.WorkspaceClient = MagicMock(return_value=mock_workspace_client)
+ mock_sdk_module.useragent = MagicMock()
+
+ # Mock both databricks and databricks.sdk modules to ensure the import works
+ with patch.dict(sys.modules, {
+ "databricks": MagicMock(),
+ "databricks.sdk": mock_sdk_module
+ }):
+ api_base, headers = databricks_base.databricks_validate_environment(
+ api_key=None,
+ api_base=None,
+ endpoint_type="chat_completions",
+ custom_endpoint=False,
+ headers=None,
+ )
- assert "Authorization" in headers
+ # Verify that SDK authentication was used (headers contain Authorization)
+ assert "Authorization" in headers
class TestEndpointURLConstruction:
diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
index 783d85f471b..91a28ee6ec9 100644
--- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
+++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
@@ -2279,3 +2279,194 @@ def test_partial_json_chunk_on_first_chunk():
assert result is None, "Partial first chunk should return None"
assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode"
+
+# ==================== Tool Type Separation Tests ====================
+# These tests verify that each Tool object contains exactly one type per Vertex AI API spec
+# Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1beta1/Tool
+
+
+def test_vertex_ai_multiple_tool_types_separate_objects():
+ """
+ Test that multiple tool types are placed in separate Tool objects.
+
+ This is required by Vertex AI API spec:
+ "A Tool object should contain exactly one type of Tool"
+
+ Related error without this fix:
+ "tools[0].tool_type: one_of 'tool_type' has more than one initialized field:
+ enterprise_web_search, url_context"
+
+ Input:
+ value=[
+ {"enterpriseWebSearch": {}},
+ {"url_context": {}},
+ ]
+
+ Expected Output:
+ tools=[
+ {"enterpriseWebSearch": {}}, # First Tool object
+ {"url_context": {}}, # Second Tool object (separate!)
+ ]
+
+ NOT (incorrect - causes API error):
+ tools=[
+ {"enterpriseWebSearch": {}, "url_context": {}} # Multiple types in one object
+ ]
+ """
+ v = VertexGeminiConfig()
+ optional_params = {}
+
+ tools = v._map_function(
+ value=[
+ {"enterpriseWebSearch": {}},
+ {"url_context": {}},
+ ],
+ optional_params=optional_params
+ )
+
+ # Should have 2 separate Tool objects
+ assert len(tools) == 2, f"Expected 2 separate Tool objects, got {len(tools)}"
+
+ # Each Tool object should contain exactly ONE type
+ tool_types_in_first = [k for k in tools[0].keys()]
+ tool_types_in_second = [k for k in tools[1].keys()]
+
+ assert len(tool_types_in_first) == 1, f"First Tool should have exactly 1 type, got {tool_types_in_first}"
+ assert len(tool_types_in_second) == 1, f"Second Tool should have exactly 1 type, got {tool_types_in_second}"
+
+ # Verify the correct tool types are present
+ assert "enterpriseWebSearch" in tools[0], "First Tool should contain enterpriseWebSearch"
+ assert "url_context" in tools[1], "Second Tool should contain url_context"
+
+
+def test_vertex_ai_function_declarations_with_other_tools_separate():
+ """
+ Test that function declarations and other tool types are in separate Tool objects.
+
+ This ensures that when using both function calling AND special tools like
+ google_search or code_execution, they are properly separated per API spec.
+
+ Input:
+ value=[
+ {"type": "function", "function": {"name": "get_weather", "description": "Get weather"}},
+ {"googleSearch": {}},
+ {"code_execution": {}},
+ ]
+
+ Expected Output:
+ tools=[
+ {"function_declarations": [{"name": "get_weather", "description": "Get weather"}]},
+ {"googleSearch": {}},
+ {"code_execution": {}},
+ ]
+ """
+ v = VertexGeminiConfig()
+ optional_params = {}
+
+ tools = v._map_function(
+ value=[
+ {"type": "function", "function": {"name": "get_weather", "description": "Get weather"}},
+ {"googleSearch": {}},
+ {"code_execution": {}},
+ ],
+ optional_params=optional_params
+ )
+
+ # Should have 3 separate Tool objects
+ assert len(tools) == 3, f"Expected 3 separate Tool objects, got {len(tools)}"
+
+ # Find each tool type
+ func_tool = None
+ search_tool = None
+ code_tool = None
+
+ for tool in tools:
+ if "function_declarations" in tool:
+ func_tool = tool
+ elif "googleSearch" in tool:
+ search_tool = tool
+ elif "code_execution" in tool:
+ code_tool = tool
+
+ # Verify all tools are present and separate
+ assert func_tool is not None, "function_declarations Tool should be present"
+ assert search_tool is not None, "googleSearch Tool should be present"
+ assert code_tool is not None, "code_execution Tool should be present"
+
+ # Verify each Tool has exactly one type
+ assert len(func_tool.keys()) == 1, "function_declarations Tool should have only one key"
+ assert len(search_tool.keys()) == 1, "googleSearch Tool should have only one key"
+ assert len(code_tool.keys()) == 1, "code_execution Tool should have only one key"
+
+ # Verify function declaration content
+ assert func_tool["function_declarations"][0]["name"] == "get_weather"
+
+
+def test_vertex_ai_single_tool_type_still_works():
+ """
+ Test that single tool type usage still works correctly (backward compatibility).
+
+ Input:
+ value=[{"code_execution": {}}]
+
+ Expected Output:
+ tools=[{"code_execution": {}}]
+ """
+ v = VertexGeminiConfig()
+ optional_params = {}
+
+ tools = v._map_function(
+ value=[{"code_execution": {}}],
+ optional_params=optional_params
+ )
+
+ assert len(tools) == 1
+ assert "code_execution" in tools[0]
+ assert tools[0]["code_execution"] == {}
+
+
+def test_vertex_ai_multiple_function_declarations_grouped():
+ """
+ Test that multiple function declarations are grouped in ONE Tool object.
+
+ Function declarations are the exception - they CAN be grouped together
+ in a single Tool object (up to 512 declarations).
+
+ Input:
+ value=[
+ {"type": "function", "function": {"name": "func1", "description": "First function"}},
+ {"type": "function", "function": {"name": "func2", "description": "Second function"}},
+ ]
+
+ Expected Output:
+ tools=[
+ {
+ "function_declarations": [
+ {"name": "func1", "description": "First function"},
+ {"name": "func2", "description": "Second function"},
+ ]
+ }
+ ]
+ """
+ v = VertexGeminiConfig()
+ optional_params = {}
+
+ tools = v._map_function(
+ value=[
+ {"type": "function", "function": {"name": "func1", "description": "First function"}},
+ {"type": "function", "function": {"name": "func2", "description": "Second function"}},
+ ],
+ optional_params=optional_params
+ )
+
+ # Should have only 1 Tool object (function declarations grouped)
+ assert len(tools) == 1, f"Expected 1 Tool object for grouped functions, got {len(tools)}"
+
+ # Should contain function_declarations with 2 functions
+ assert "function_declarations" in tools[0]
+ assert len(tools[0]["function_declarations"]) == 2
+
+ # Verify function names
+ func_names = [f["name"] for f in tools[0]["function_declarations"]]
+ assert "func1" in func_names
+ assert "func2" in func_names
diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py
new file mode 100644
index 00000000000..fca784342d7
--- /dev/null
+++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py
@@ -0,0 +1,179 @@
+"""
+Tests for Vertex AI Anthropic image URL handling.
+
+Issue: https://github.com/BerriAI/litellm/issues/18430
+Vertex AI Anthropic models don't support URL sources for images.
+LiteLLM should convert image URLs to base64 when using Vertex AI Anthropic.
+"""
+import os
+import sys
+from unittest.mock import patch, MagicMock
+
+import pytest
+
+sys.path.insert(
+ 0, os.path.abspath("../../../../../..")
+) # Adds the parent directory to the system path
+
+from litellm.litellm_core_utils.prompt_templates.factory import (
+ anthropic_messages_pt,
+ create_anthropic_image_param,
+)
+
+
+class TestVertexAIAnthropicImageURLHandling:
+ """Test that Vertex AI Anthropic converts image URLs to base64."""
+
+ @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64")
+ def test_vertex_ai_anthropic_converts_https_url_to_base64(
+ self, mock_convert_url: MagicMock
+ ):
+ """
+ Test that HTTPS image URLs are converted to base64 for Vertex AI Anthropic.
+
+ For regular Anthropic, HTTPS URLs are passed through as URL type.
+ For Vertex AI Anthropic, HTTPS URLs should be converted to base64.
+ """
+ mock_convert_url.return_value = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ=="
+
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "Describe this image"},
+ {
+ "type": "image_url",
+ "image_url": {"url": "https://example.com/image.jpg"},
+ },
+ ],
+ }
+ ]
+
+ # For Vertex AI, image URLs should be converted to base64
+ result = anthropic_messages_pt(
+ messages=messages,
+ model="claude-sonnet-4",
+ llm_provider="vertex_ai",
+ )
+
+ # Verify convert_url_to_base64 was called
+ mock_convert_url.assert_called_once_with(url="https://example.com/image.jpg")
+
+ # Check the result has base64 source type
+ user_message = result[0]
+ assert user_message["role"] == "user"
+ image_content = user_message["content"][1]
+ assert image_content["type"] == "image"
+ assert image_content["source"]["type"] == "base64"
+
+ @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64")
+ def test_regular_anthropic_uses_url_type_for_https(
+ self, mock_convert_url: MagicMock
+ ):
+ """
+ Test that regular Anthropic API uses URL type for HTTPS images.
+
+ This confirms the original behavior is preserved for non-Vertex AI.
+ """
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "Describe this image"},
+ {
+ "type": "image_url",
+ "image_url": {"url": "https://example.com/image.jpg"},
+ },
+ ],
+ }
+ ]
+
+ # For regular Anthropic, HTTPS URLs should NOT be converted
+ result = anthropic_messages_pt(
+ messages=messages,
+ model="claude-sonnet-4",
+ llm_provider="anthropic",
+ )
+
+ # convert_url_to_base64 should NOT be called for regular Anthropic with HTTPS
+ mock_convert_url.assert_not_called()
+
+ # Check the result has URL source type
+ user_message = result[0]
+ assert user_message["role"] == "user"
+ image_content = user_message["content"][1]
+ assert image_content["type"] == "image"
+ assert image_content["source"]["type"] == "url"
+ assert image_content["source"]["url"] == "https://example.com/image.jpg"
+
+ @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64")
+ def test_vertex_ai_beta_also_converts_to_base64(
+ self, mock_convert_url: MagicMock
+ ):
+ """
+ Test that vertex_ai_beta provider also converts image URLs to base64.
+ """
+ mock_convert_url.return_value = "data:image/png;base64,iVBORw0KGgo="
+
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "What is in this image?"},
+ {
+ "type": "image_url",
+ "image_url": "https://example.com/photo.png",
+ },
+ ],
+ }
+ ]
+
+ result = anthropic_messages_pt(
+ messages=messages,
+ model="claude-3-opus",
+ llm_provider="vertex_ai_beta",
+ )
+
+ # Verify convert_url_to_base64 was called
+ mock_convert_url.assert_called_once()
+
+ # Check the result has base64 source type
+ user_message = result[0]
+ image_content = user_message["content"][1]
+ assert image_content["source"]["type"] == "base64"
+
+
+class TestCreateAnthropicImageParam:
+ """Test the create_anthropic_image_param function directly."""
+
+ @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64")
+ def test_force_base64_converts_https_url(self, mock_convert_url: MagicMock):
+ """
+ Test that is_bedrock_invoke=True (used for both Bedrock and Vertex AI)
+ forces conversion of HTTPS URLs to base64.
+ """
+ mock_convert_url.return_value = "data:image/jpeg;base64,/9j/4AAQSkZJRg=="
+
+ result = create_anthropic_image_param(
+ image_url_input="https://example.com/image.jpg",
+ format=None,
+ is_bedrock_invoke=True, # This flag is set for both Bedrock and Vertex AI
+ )
+
+ mock_convert_url.assert_called_once_with(url="https://example.com/image.jpg")
+ assert result["source"]["type"] == "base64"
+
+ @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64")
+ def test_no_force_uses_url_type(self, mock_convert_url: MagicMock):
+ """
+ Test that without force, HTTPS URLs use URL type.
+ """
+ result = create_anthropic_image_param(
+ image_url_input="https://example.com/image.jpg",
+ format=None,
+ is_bedrock_invoke=False,
+ )
+
+ mock_convert_url.assert_not_called()
+ assert result["source"]["type"] == "url"
+ assert result["source"]["url"] == "https://example.com/image.jpg"
diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/test_litellm/llms/zai/test_zai_provider.py
index a3d47d666bc..d1e4359d048 100644
--- a/tests/test_litellm/llms/zai/test_zai_provider.py
+++ b/tests/test_litellm/llms/zai/test_zai_provider.py
@@ -1,6 +1,7 @@
"""
Tests for Z.AI (Zhipu AI) provider - GLM models
"""
+
import json
import math
@@ -50,10 +51,12 @@ def test_zai_in_provider_lists():
def test_zai_models_in_model_cost():
"""Test that ZAI models are in the model cost map"""
import os
+
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
zai_models = [
+ "zai/glm-4.7",
"zai/glm-4.6",
"zai/glm-4.5",
"zai/glm-4.5v",
@@ -72,6 +75,7 @@ def test_zai_models_in_model_cost():
def test_zai_glm46_cost_calculation():
"""Test the cost calculation for glm-4.6"""
import os
+
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
@@ -92,6 +96,7 @@ def test_zai_glm46_cost_calculation():
def test_zai_flash_model_is_free():
"""Test that glm-4.5-flash has zero cost"""
import os
+
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
@@ -102,6 +107,38 @@ def test_zai_flash_model_is_free():
assert info["output_cost_per_token"] == 0
+def test_glm47_supports_reasoning():
+ """Test that GLM-4.7 supports reasoning"""
+ import os
+
+ os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
+ litellm.model_cost = litellm.get_model_cost_map(url="")
+
+ key = "zai/glm-4.7"
+ assert key in litellm.model_cost, f"Model {key} not found in model_cost"
+
+ info = litellm.model_cost[key]
+ assert info["supports_reasoning"] is True
+
+
+def test_glm47_cost_calculation():
+ """Test cost calculation for GLM-4.7"""
+ import os
+
+ os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
+ litellm.model_cost = litellm.get_model_cost_map(url="")
+
+ prompt_cost, completion_cost = cost_per_token(
+ model="zai/glm-4.7",
+ prompt_tokens=1000000, # 1M tokens
+ completion_tokens=1000000,
+ )
+
+ # GLM-4.7: $0.6/M input, $2.2/M output (same as GLM-4.6)
+ assert math.isclose(prompt_cost, 0.6, rel_tol=1e-6)
+ assert math.isclose(completion_cost, 2.2, rel_tol=1e-6)
+
+
@pytest.mark.asyncio
async def test_zai_completion_call(respx_mock, zai_response, monkeypatch):
"""Test completion call with zai provider using mocked response"""
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py b/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py
new file mode 100644
index 00000000000..0e150e064c7
--- /dev/null
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py
@@ -0,0 +1,78 @@
+"""Tests for the MCP guardrail translation handler."""
+
+import pytest
+
+from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import (
+ MCPGuardrailTranslationHandler,
+)
+
+
+class MockGuardrail(CustomGuardrail):
+ """Simple guardrail mock that records invocations."""
+
+ def __init__(self, return_texts=None):
+ super().__init__(guardrail_name="mock-mcp-guardrail")
+ self.return_texts = return_texts
+ self.call_count = 0
+ self.last_inputs = None
+
+ async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
+ self.call_count += 1
+ self.last_inputs = inputs
+
+ if self.return_texts is not None:
+ return {"texts": self.return_texts}
+
+ texts = inputs.get("texts", [])
+ return {"texts": [f"{text} [SAFE]" for text in texts]}
+
+
+@pytest.mark.asyncio
+async def test_process_input_messages_updates_content():
+ """Handler should update the synthetic message content when guardrail modifies text."""
+ handler = MCPGuardrailTranslationHandler()
+ guardrail = MockGuardrail()
+
+ original_content = "Tool: weather\nArguments: {'city': 'tokyo'}"
+ data = {
+ "messages": [{"role": "user", "content": original_content}],
+ "mcp_tool_name": "weather",
+ }
+
+ result = await handler.process_input_messages(data, guardrail)
+
+ assert result["messages"][0]["content"].endswith("[SAFE]")
+ assert guardrail.last_inputs == {"texts": [original_content]}
+ assert guardrail.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_process_input_messages_skips_when_no_messages():
+ """Handler should skip guardrail invocation if messages array is missing or empty."""
+ handler = MCPGuardrailTranslationHandler()
+ guardrail = MockGuardrail()
+
+ data = {"mcp_tool_name": "noop"}
+ result = await handler.process_input_messages(data, guardrail)
+
+ assert result == data
+ assert guardrail.call_count == 0
+
+
+@pytest.mark.asyncio
+async def test_process_input_messages_handles_empty_guardrail_result():
+ """Handler should leave content untouched when guardrail returns no text updates."""
+ handler = MCPGuardrailTranslationHandler()
+ guardrail = MockGuardrail(return_texts=[])
+
+ original_content = "Tool: calendar\nArguments: {'date': '2024-12-25'}"
+ data = {
+ "messages": [{"role": "user", "content": original_content}],
+ "mcp_tool_name": "calendar",
+ }
+
+ result = await handler.process_input_messages(data, guardrail)
+
+ assert result["messages"][0]["content"] == original_content
+ assert guardrail.call_count == 1
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py
index 5581070be71..a2425cc659a 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py
@@ -71,28 +71,29 @@ class TestMCPCustomFields:
manager = MCPServerManager()
# Mock database record with custom fields
- mock_server = Mock(spec=LiteLLM_MCPServerTable)
- mock_server.server_id = "test-server-id"
- mock_server.server_name = "Test Server"
- mock_server.description = "A test server"
- mock_server.url = "http://localhost:3000"
- mock_server.transport = "http"
- mock_server.auth_type = MCPAuth.bearer_token
- mock_server.alias = None
- mock_server.mcp_info = {
- "server_name": "Test Server",
- "description": "A test server",
- "custom_db_field": "database_value",
- "metadata": {"source": "database"},
- "version": "1.0.0"
- }
- mock_server.command = None
- mock_server.args = None
- mock_server.env = None
- mock_server.mcp_access_groups = None
+ mock_server = LiteLLM_MCPServerTable(
+ server_id="test-server-id",
+ server_name="Test Server",
+ alias=None,
+ description="A test server",
+ url="http://localhost:3000",
+ transport="http",
+ auth_type=MCPAuth.bearer_token,
+ mcp_info={
+ "server_name": "Test Server",
+ "description": "A test server",
+ "custom_db_field": "database_value",
+ "metadata": {"source": "database"},
+ "version": "1.0.0",
+ },
+ command=None,
+ args=[],
+ env={},
+ mcp_access_groups=[],
+ )
# Add server to manager
- await manager.add_update_server(mock_server)
+ await manager.add_server(mock_server)
# Get the added server
server = manager.get_mcp_server_by_id("test-server-id")
@@ -209,4 +210,4 @@ class TestMCPCustomFields:
# Should use mcp_info description, not config level
assert mcp_info["description"] == "MCP info description"
- assert mcp_info["custom_field"] == "custom_value"
\ No newline at end of file
+ assert mcp_info["custom_field"] == "custom_value"
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
index ff016a1a130..6491e11024a 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
@@ -65,7 +65,7 @@ class TestMCPServerManager:
updated_at=datetime.now(),
)
- await manager.add_update_server(stdio_server)
+ await manager.add_server(stdio_server)
# Verify server was added
assert "stdio-server-1" in manager.registry
@@ -641,37 +641,31 @@ class TestMCPServerManager:
manager = MCPServerManager()
# Mock server
- server = MagicMock()
- server.server_id = "test-server"
- server.name = "test-server"
+ server = MCPServer(
+ server_id="test-server",
+ name="test-server",
+ transport=MCPTransport.http,
+ auth_type=None,
+ authentication_token="test-token",
+ url="http://test-server.com",
+ )
manager.get_mcp_server_by_id = MagicMock(return_value=server)
- # Mock successful _get_tools_from_server
- async def mock_get_tools_from_server(
- server,
- mcp_auth_header=None,
- raw_headers=None,
- ):
- tool1 = MagicMock()
- tool1.name = "tool1"
- tool2 = MagicMock()
- tool2.name = "tool2"
- return [tool1, tool2]
-
- manager._get_tools_from_server = mock_get_tools_from_server
+ # Mock successful client.run_with_session
+ mock_client = AsyncMock()
+ mock_client.run_with_session = AsyncMock(return_value="ok")
+ manager._create_mcp_client = MagicMock(return_value=mock_client)
# Perform health check
result = await manager.health_check_server("test-server")
- # Verify results
- assert result["server_id"] == "test-server"
- assert result["status"] == "healthy"
- assert result["tools_count"] == 2
- assert result["error"] is None
- assert "last_health_check" in result
- assert "response_time_ms" in result
- assert result["response_time_ms"] >= 0 # Allow 0 for very fast mocks
+ # Verify results - result is now LiteLLM_MCPServerTable
+ assert isinstance(result, LiteLLM_MCPServerTable)
+ assert result.server_id == "test-server"
+ assert result.status == "healthy"
+ assert result.health_check_error is None
+ assert result.last_health_check is not None
@pytest.mark.asyncio
async def test_health_check_server_unhealthy(self):
@@ -679,32 +673,33 @@ class TestMCPServerManager:
manager = MCPServerManager()
# Mock server
- server = MagicMock()
- server.server_id = "test-server"
- server.name = "test-server"
+ server = MCPServer(
+ server_id="test-server",
+ name="test-server",
+ transport=MCPTransport.http,
+ auth_type=None,
+ authentication_token="test-token",
+ url="http://test-server.com",
+ )
manager.get_mcp_server_by_id = MagicMock(return_value=server)
- # Mock failed _get_tools_from_server
- async def mock_get_tools_from_server(
- server,
- mcp_auth_header=None,
- raw_headers=None,
- ):
- raise Exception("Connection timeout")
-
- manager._get_tools_from_server = mock_get_tools_from_server
+ # Mock failed client.run_with_session
+ mock_client = AsyncMock()
+ mock_client.run_with_session = AsyncMock(
+ side_effect=Exception("Connection timeout")
+ )
+ manager._create_mcp_client = MagicMock(return_value=mock_client)
# Perform health check
result = await manager.health_check_server("test-server")
# Verify results
- assert result["server_id"] == "test-server"
- assert result["status"] == "unhealthy"
- assert result["error"] == "Connection timeout"
- assert "last_health_check" in result
- assert "response_time_ms" in result
- assert result["response_time_ms"] >= 0 # Allow 0 for very fast mocks
+ assert isinstance(result, LiteLLM_MCPServerTable)
+ assert result.server_id == "test-server"
+ assert result.status == "unhealthy"
+ assert result.health_check_error == "Connection timeout"
+ assert result.last_health_check is not None
@pytest.mark.asyncio
async def test_health_check_server_not_found(self):
@@ -718,104 +713,121 @@ class TestMCPServerManager:
result = await manager.health_check_server("non-existent-server")
# Verify results
- assert result["server_id"] == "non-existent-server"
- assert result["status"] == "unknown"
- assert result["error"] == "Server not found"
- assert result["response_time_ms"] is None
- assert "last_health_check" in result
+ assert isinstance(result, LiteLLM_MCPServerTable)
+ assert result.server_id == "non-existent-server"
+ assert result.server_name is None
+ assert result.status == "unknown"
+ assert result.health_check_error == "Server not found"
+ assert result.last_health_check is not None
@pytest.mark.asyncio
- async def test_health_check_all_servers(self):
- """Test health check for all servers"""
+ async def test_health_check_server_oauth2_skips_check(self):
+ """Test that health check is skipped for OAuth2 servers and returns unknown status"""
manager = MCPServerManager()
- # Mock servers
- server1 = MagicMock()
- server1.server_id = "server1"
- server1.name = "server1"
-
- server2 = MagicMock()
- server2.server_id = "server2"
- server2.name = "server2"
-
- # Mock registry
- manager.registry = {"server1": server1, "server2": server2}
-
- # Mock get_mcp_server_by_id
- def mock_get_server_by_id(server_id):
- if server_id == "server1":
- return server1
- elif server_id == "server2":
- return server2
- return None
-
- manager.get_mcp_server_by_id = mock_get_server_by_id
-
- # Mock _get_tools_from_server with different results
- async def mock_get_tools_from_server(
- server,
- mcp_auth_header=None,
- raw_headers=None,
- ):
- if server.server_id == "server1":
- tool = MagicMock()
- tool.name = "tool1"
- return [tool]
- elif server.server_id == "server2":
- raise Exception("Connection failed")
- return []
-
- manager._get_tools_from_server = mock_get_tools_from_server
-
- # Perform health check for all servers
- result = await manager.health_check_all_servers()
-
- # Verify results
- assert len(result) == 2
- assert "server1" in result
- assert "server2" in result
-
- # Check server1 (healthy)
- assert result["server1"]["status"] == "healthy"
- assert result["server1"]["tools_count"] == 1
- assert result["server1"]["error"] is None
-
- # Check server2 (unhealthy)
- assert result["server2"]["status"] == "unhealthy"
- assert result["server2"]["error"] == "Connection failed"
-
- @pytest.mark.asyncio
- async def test_health_check_server_with_auth_header(self):
- """Test health check with authentication header"""
- manager = MCPServerManager()
-
- # Mock server
- server = MagicMock()
- server.server_id = "test-server"
- server.name = "test-server"
+ # Mock OAuth2 server
+ server = MCPServer(
+ server_id="oauth2-server",
+ name="oauth2-server",
+ transport=MCPTransport.http,
+ auth_type=MCPAuth.oauth2,
+ url="http://oauth2-server.com",
+ )
manager.get_mcp_server_by_id = MagicMock(return_value=server)
- # Mock _get_tools_from_server to verify auth header is passed
- async def mock_get_tools_from_server(
- server,
- mcp_auth_header=None,
- raw_headers=None,
- ):
- assert mcp_auth_header == "test-token"
- tool = MagicMock()
- tool.name = "tool1"
- return [tool]
+ # _create_mcp_client should not be called for OAuth2 servers
+ manager._create_mcp_client = MagicMock()
- manager._get_tools_from_server = mock_get_tools_from_server
+ # Perform health check
+ result = await manager.health_check_server("oauth2-server")
- # Perform health check with auth header
- result = await manager.health_check_server("test-server", "test-token")
+ # Verify that client was not created (health check was skipped)
+ manager._create_mcp_client.assert_not_called()
# Verify results
- assert result["server_id"] == "test-server"
- assert result["status"] == "healthy"
- assert result["tools_count"] == 1
+ assert isinstance(result, LiteLLM_MCPServerTable)
+ assert result.server_id == "oauth2-server"
+ assert result.status == "unknown"
+ assert result.health_check_error is None
+ assert result.last_health_check is not None
+
+ @pytest.mark.asyncio
+ async def test_health_check_server_no_token_skips_check(self):
+ """Test that health check is skipped when auth_type is set but authentication_token is missing"""
+ manager = MCPServerManager()
+
+ # Mock server with auth_type but no authentication_token
+ server = MCPServer(
+ server_id="no-token-server",
+ name="no-token-server",
+ transport=MCPTransport.http,
+ auth_type=MCPAuth.bearer_token,
+ authentication_token=None, # No token
+ url="http://no-token-server.com",
+ )
+
+ manager.get_mcp_server_by_id = MagicMock(return_value=server)
+
+ # _create_mcp_client should not be called
+ manager._create_mcp_client = MagicMock()
+
+ # Perform health check
+ result = await manager.health_check_server("no-token-server")
+
+ # Verify that client was not created (health check was skipped)
+ manager._create_mcp_client.assert_not_called()
+
+ # Verify results
+ assert isinstance(result, LiteLLM_MCPServerTable)
+ assert result.server_id == "no-token-server"
+ assert result.status == "unknown"
+ assert result.health_check_error is None
+ assert result.last_health_check is not None
+
+ @pytest.mark.asyncio
+ async def test_health_check_server_with_static_headers(self):
+ """Test health check with static headers configured"""
+ manager = MCPServerManager()
+
+ # Mock server with static_headers
+ server = MCPServer(
+ server_id="test-server",
+ name="test-server",
+ transport=MCPTransport.http,
+ auth_type=None,
+ authentication_token="test-token",
+ url="http://test-server.com",
+ static_headers={"X-Custom-Header": "custom-value"},
+ )
+
+ manager.get_mcp_server_by_id = MagicMock(return_value=server)
+
+ # Mock successful client
+ mock_client = AsyncMock()
+ mock_client.run_with_session = AsyncMock(return_value="ok")
+
+ # Capture the extra_headers passed to _create_mcp_client
+ captured_extra_headers = None
+
+ def capture_create_mcp_client(server, mcp_auth_header, extra_headers, stdio_env):
+ nonlocal captured_extra_headers
+ captured_extra_headers = extra_headers
+ return mock_client
+
+ manager._create_mcp_client = MagicMock(side_effect=capture_create_mcp_client)
+
+ # Perform health check
+ result = await manager.health_check_server("test-server")
+
+ # Verify static headers were passed
+ assert captured_extra_headers == {"X-Custom-Header": "custom-value"}
+
+ # Verify results
+ assert isinstance(result, LiteLLM_MCPServerTable)
+ assert result.server_id == "test-server"
+ assert result.status == "healthy"
+ assert result.health_check_error is None
@pytest.mark.asyncio
async def test_pre_call_tool_check_allowed_tools_list_allows_tool(self):
@@ -1352,7 +1364,7 @@ class TestMCPServerManager:
"env": {},
},
)
- await manager.add_update_server(server)
+ await manager.add_server(server)
assert server.server_id in manager.get_registry()
@pytest.mark.asyncio
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py
index cb48a940b57..1044ad3ed27 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py
@@ -5,6 +5,8 @@ This test suite ensures that:
1. Parameter names with invalid Python identifiers are handled safely
2. No exec() is used (security)
3. All edge cases (hyphens, dots, keywords, special chars) work correctly
+4. Path traversal attacks are prevented
+5. Path parameters are properly URL encoded
"""
import pytest
@@ -452,3 +454,69 @@ class TestExtractParameters:
if __name__ == "__main__":
pytest.main([__file__, "-v"])
+
+
+class TestPathSecurity:
+ """Test path traversal security and URL encoding."""
+
+ @pytest.mark.asyncio
+ async def test_should_reject_path_traversal_inputs(self):
+ """Test that path traversal attacks (../admin) are rejected."""
+ operation = {
+ "parameters": [
+ {
+ "name": "filename",
+ "in": "path",
+ "required": True,
+ "schema": {"type": "string"},
+ }
+ ]
+ }
+
+ tool_function = create_tool_function(
+ path="/files/{filename}",
+ method="GET",
+ operation=operation,
+ base_url="https://example.com",
+ )
+
+ response = await tool_function(**{"filename": "../admin"})
+
+ assert "Invalid path parameter" in response
+
+ @pytest.mark.asyncio
+ async def test_should_encode_and_request_safe_path_parameters(self):
+ """Test that path parameters are properly URL encoded."""
+ operation = {
+ "parameters": [
+ {
+ "name": "filename",
+ "in": "path",
+ "required": True,
+ "schema": {"type": "string"},
+ }
+ ]
+ }
+
+ tool_function = create_tool_function(
+ path="/files/{filename}",
+ method="GET",
+ operation=operation,
+ base_url="https://example.com",
+ )
+
+ with patch("httpx.AsyncClient") as mock_client:
+ mock_response = AsyncMock()
+ mock_response.text = "dummy-response"
+ mock_client.return_value.__aenter__.return_value.get = AsyncMock(
+ return_value=mock_response
+ )
+
+ response = await tool_function(**{"filename": "report 2024.json"})
+
+ assert response == "dummy-response"
+
+ # Verify URL was properly encoded
+ call_args = mock_client.return_value.__aenter__.return_value.get.call_args
+ url = call_args[0][0]
+ assert url == "https://example.com/files/report%202024.json"
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
index ce38ee59e4d..0c6d0921952 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
@@ -7,6 +7,7 @@ from litellm.proxy._experimental.mcp_server import rest_endpoints
from litellm.proxy._experimental.mcp_server.auth import (
user_api_key_auth_mcp as auth_mcp,
)
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy._types import NewMCPServerRequest, UserAPIKeyAuth
from litellm.types.mcp import MCPAuth
@@ -31,6 +32,60 @@ def _build_request(headers: Optional[Dict[str, str]] = None) -> Request:
return Request(scope, receive=receive)
+def _get_route(path: str, method: str):
+ for route in rest_endpoints.router.routes:
+ if getattr(route, "path", None) == path and method in getattr(
+ route, "methods", set()
+ ):
+ return route
+ raise AssertionError(f"Route {method} {path} not found")
+
+
+def _route_has_dependency(route, dependency) -> bool:
+ if any(
+ getattr(dep, "dependency", None) == dependency
+ for dep in getattr(route, "dependencies", [])
+ ):
+ return True
+ dependant = getattr(route, "dependant", None)
+ if dependant is None:
+ return False
+ return any(getattr(dep, "call", None) == dependency for dep in dependant.dependencies)
+
+
+@pytest.mark.asyncio
+async def test_execute_with_mcp_client_redacts_stack_trace(monkeypatch):
+ def fake_create_client(*args, **kwargs):
+ return object()
+
+ monkeypatch.setattr(
+ rest_endpoints.global_mcp_server_manager,
+ "_create_mcp_client",
+ fake_create_client,
+ )
+
+ async def failing_operation(client):
+ raise RuntimeError("boom")
+
+ payload = NewMCPServerRequest(
+ server_name="example",
+ url="https://example.com",
+ auth_type=MCPAuth.none,
+ )
+
+ result = await rest_endpoints._execute_with_mcp_client(
+ payload, failing_operation
+ )
+
+ assert result["status"] == "error"
+ assert "stack_trace" not in result
+
+
+def test_test_connection_requires_auth_dependency():
+ route = _get_route("/mcp-rest/test/connection", "POST")
+ assert _route_has_dependency(route, user_api_key_auth)
+
+
@pytest.mark.asyncio
async def test_test_tools_list_forwards_mcp_auth_header(monkeypatch):
"""Ensure credential-based auth forwards the auth_value to the MCP client."""
diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py
index 8ecbaced21e..b56d13bb932 100644
--- a/tests/test_litellm/proxy/auth/test_handle_jwt.py
+++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py
@@ -1146,4 +1146,343 @@ async def test_auth_builder_uses_team_from_header_e2e():
)
assert result["team_id"] == "team-2"
- assert result["team_object"] == team_object
\ No newline at end of file
+ assert result["team_object"] == team_object
+
+
+@pytest.mark.asyncio
+async def test_get_team_alias_with_nested_fields():
+ """
+ Test get_team_alias() method with nested JWT fields
+ """
+ from litellm.proxy._types import LiteLLM_JWTAuth
+ from litellm.proxy.auth.handle_jwt import JWTHandler
+
+ jwt_handler = JWTHandler()
+
+ # Test token with nested team name
+ nested_token = {
+ "organization": {
+ "team": {
+ "name": "engineering-team"
+ }
+ },
+ "team_name": "flat-team"
+ }
+
+ # Test nested access
+ jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="organization.team.name")
+ assert jwt_handler.get_team_alias(nested_token, None) == "engineering-team"
+
+ # Test flat access (backward compatibility)
+ jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="team_name")
+ assert jwt_handler.get_team_alias(nested_token, None) == "flat-team"
+
+ # Test missing field returns default
+ jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="nonexistent.field")
+ assert jwt_handler.get_team_alias(nested_token, "default-team") == "default-team"
+
+ # Test with team_alias_jwt_field not configured
+ jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() # team_alias_jwt_field is None
+ assert jwt_handler.get_team_alias(nested_token, "default") is None
+
+
+@pytest.mark.asyncio
+async def test_is_required_team_id_with_team_alias_field():
+ """
+ Test that is_required_team_id() returns True when team_alias_jwt_field is set
+ """
+ from litellm.proxy._types import LiteLLM_JWTAuth
+ from litellm.proxy.auth.handle_jwt import JWTHandler
+
+ jwt_handler = JWTHandler()
+
+ # Neither field set - should return False
+ jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
+ assert jwt_handler.is_required_team_id() is False
+
+ # Only team_id_jwt_field set - should return True
+ jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id")
+ assert jwt_handler.is_required_team_id() is True
+
+ # Only team_alias_jwt_field set - should return True
+ jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="team_name")
+ assert jwt_handler.is_required_team_id() is True
+
+ # Both fields set - should return True
+ jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
+ team_id_jwt_field="team_id",
+ team_alias_jwt_field="team_name"
+ )
+ assert jwt_handler.is_required_team_id() is True
+
+
+@pytest.mark.asyncio
+async def test_find_and_validate_specific_team_id_with_team_alias():
+ """
+ Test that find_and_validate_specific_team_id resolves team by name when team_id is not found
+ """
+ from unittest.mock import MagicMock
+
+ from litellm.caching import DualCache
+ from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_TeamTable
+ from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
+ from litellm.proxy.utils import ProxyLogging
+
+ jwt_handler = JWTHandler()
+ user_api_key_cache = DualCache()
+ proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
+
+ jwt_handler.update_environment(
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ litellm_jwtauth=LiteLLM_JWTAuth(
+ team_alias_jwt_field="team_alias"
+ ),
+ )
+
+ # Token with team name (no team_id)
+ jwt_token = {
+ "sub": "user-1",
+ "team_alias": "my-team"
+ }
+
+ # Mock team object returned by get_team_object_by_alias
+ team_object = LiteLLM_TeamTable(team_id="resolved-team-id", team_alias="my-team")
+
+ with patch(
+ "litellm.proxy.auth.handle_jwt.get_team_object_by_alias",
+ new_callable=AsyncMock
+ ) as mock_get_by_alias:
+ mock_get_by_alias.return_value = team_object
+
+ team_id, result_team = await JWTAuthManager.find_and_validate_specific_team_id(
+ jwt_handler=jwt_handler,
+ jwt_valid_token=jwt_token,
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=None,
+ proxy_logging_obj=proxy_logging_obj,
+ )
+
+ # Should have resolved team_id from team name
+ assert team_id == "resolved-team-id"
+ assert result_team == team_object
+ mock_get_by_alias.assert_called_once_with(
+ team_alias="my-team",
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=None,
+ proxy_logging_obj=proxy_logging_obj,
+ )
+
+
+@pytest.mark.asyncio
+async def test_find_and_validate_team_id_takes_precedence_over_name():
+ """
+ Test that team_id_jwt_field takes precedence over team_alias_jwt_field
+ """
+ from unittest.mock import MagicMock
+
+ from litellm.caching import DualCache
+ from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_TeamTable
+ from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
+ from litellm.proxy.utils import ProxyLogging
+
+ jwt_handler = JWTHandler()
+ user_api_key_cache = DualCache()
+ proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
+
+ jwt_handler.update_environment(
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ litellm_jwtauth=LiteLLM_JWTAuth(
+ team_id_jwt_field="team_id",
+ team_alias_jwt_field="team_alias"
+ ),
+ )
+
+ # Token with both team_id and team name
+ jwt_token = {
+ "sub": "user-1",
+ "team_id": "direct-team-id",
+ "team_alias": "my-team"
+ }
+
+ # Mock team object returned by get_team_object (by ID)
+ team_object = LiteLLM_TeamTable(team_id="direct-team-id")
+
+ with patch(
+ "litellm.proxy.auth.handle_jwt.get_team_object",
+ new_callable=AsyncMock
+ ) as mock_get_by_id, patch(
+ "litellm.proxy.auth.handle_jwt.get_team_object_by_alias",
+ new_callable=AsyncMock
+ ) as mock_get_by_alias:
+ mock_get_by_id.return_value = team_object
+
+ team_id, result_team = await JWTAuthManager.find_and_validate_specific_team_id(
+ jwt_handler=jwt_handler,
+ jwt_valid_token=jwt_token,
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=None,
+ proxy_logging_obj=proxy_logging_obj,
+ )
+
+ # Should use team_id directly, not resolve by name
+ assert team_id == "direct-team-id"
+ assert result_team == team_object
+ mock_get_by_id.assert_called_once()
+ mock_get_by_alias.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_find_and_validate_raises_when_required_team_not_found():
+ """
+ Test that an exception is raised when team is required but neither team_id nor team_name is found
+ """
+ from litellm.caching import DualCache
+ from litellm.proxy._types import LiteLLM_JWTAuth
+ from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
+ from litellm.proxy.utils import ProxyLogging
+
+ jwt_handler = JWTHandler()
+ user_api_key_cache = DualCache()
+ proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
+
+ jwt_handler.update_environment(
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ litellm_jwtauth=LiteLLM_JWTAuth(
+ team_alias_jwt_field="team_alias" # Required, but not in token
+ ),
+ )
+
+ # Token without team info
+ jwt_token = {
+ "sub": "user-1"
+ }
+
+ with pytest.raises(Exception) as exc_info:
+ await JWTAuthManager.find_and_validate_specific_team_id(
+ jwt_handler=jwt_handler,
+ jwt_valid_token=jwt_token,
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=None,
+ proxy_logging_obj=proxy_logging_obj,
+ )
+
+ assert "No team found in token" in str(exc_info.value)
+ assert "team_alias field 'team_alias'" in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_get_org_alias_with_nested_fields():
+ """
+ Test get_org_alias() method with nested JWT fields
+ """
+ from litellm.proxy._types import LiteLLM_JWTAuth
+ from litellm.proxy.auth.handle_jwt import JWTHandler
+
+ jwt_handler = JWTHandler()
+
+ # Test token with nested org name
+ nested_token = {
+ "company": {
+ "organization": {
+ "name": "acme-corp"
+ }
+ },
+ "org_name": "flat-org"
+ }
+
+ # Test nested access
+ jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_alias_jwt_field="company.organization.name")
+ assert jwt_handler.get_org_alias(nested_token, None) == "acme-corp"
+
+ # Test flat access
+ jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_alias_jwt_field="org_name")
+ assert jwt_handler.get_org_alias(nested_token, None) == "flat-org"
+
+ # Test missing field returns default
+ jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_alias_jwt_field="nonexistent.field")
+ assert jwt_handler.get_org_alias(nested_token, "default-org") == "default-org"
+
+ # Test with org_alias_jwt_field not configured
+ jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
+ assert jwt_handler.get_org_alias(nested_token, "default") is None
+
+
+@pytest.mark.asyncio
+async def test_get_objects_resolves_org_by_name():
+ """
+ Test that get_objects resolves organization by name when org_id is not provided
+ """
+ from litellm.caching import DualCache
+ from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_OrganizationTable
+ from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
+ from litellm.proxy.utils import ProxyLogging
+
+ jwt_handler = JWTHandler()
+ user_api_key_cache = DualCache()
+ proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
+
+ jwt_handler.update_environment(
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ litellm_jwtauth=LiteLLM_JWTAuth(
+ org_alias_jwt_field="org_alias"
+ ),
+ )
+
+ # Mock org object returned by get_org_object_by_alias
+ org_object = LiteLLM_OrganizationTable(
+ organization_id="resolved-org-id",
+ organization_alias="my-org",
+ budget_id="budget-1",
+ created_by="admin",
+ updated_by="admin",
+ models=[]
+ )
+
+ with patch(
+ "litellm.proxy.auth.handle_jwt.get_org_object_by_alias",
+ new_callable=AsyncMock
+ ) as mock_get_by_alias:
+ mock_get_by_alias.return_value = org_object
+
+ (
+ result_user_obj,
+ result_org_obj,
+ result_end_user_obj,
+ result_team_membership,
+ ) = await JWTAuthManager.get_objects(
+ user_id=None,
+ user_email=None,
+ org_id=None, # No org_id provided
+ end_user_id=None,
+ team_id=None,
+ valid_user_email=None,
+ jwt_handler=jwt_handler,
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=None,
+ proxy_logging_obj=proxy_logging_obj,
+ route="/chat/completions",
+ org_alias="my-org",
+ )
+
+ # Should resolve org by alias - org_id can be derived from org_object.organization_id
+ assert result_org_obj == org_object
+ assert result_org_obj.organization_id == "resolved-org-id"
+ mock_get_by_alias.assert_called_once_with(
+ org_alias="my-org",
+ prisma_client=None,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=None,
+ proxy_logging_obj=proxy_logging_obj,
+ )
+
+
+
diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py
index 201461dc8b5..c04b8114939 100644
--- a/tests/test_litellm/proxy/auth/test_login_utils.py
+++ b/tests/test_litellm/proxy/auth/test_login_utils.py
@@ -6,6 +6,7 @@ to login_utils.py for better reusability.
"""
import os
+from datetime import datetime, timezone, timedelta
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@@ -21,6 +22,7 @@ from litellm.proxy._types import (
from litellm.proxy.auth.login_utils import (
LoginResult,
authenticate_user,
+ expire_previous_ui_session_tokens,
get_ui_credentials,
)
@@ -282,3 +284,268 @@ async def test_authenticate_user_database_required_for_admin():
finally:
if original_db_url:
os.environ["DATABASE_URL"] = original_db_url
+
+
+@pytest.mark.asyncio
+async def test_expire_previous_ui_session_tokens_none_prisma_client():
+ """Test that function returns early when prisma_client is None"""
+ await expire_previous_ui_session_tokens("test-user", None)
+ # Should not raise any exception
+
+
+@pytest.mark.asyncio
+async def test_expire_previous_ui_session_tokens_only_litellm_dashboard_team():
+ """Test that only tokens with team_id='litellm-dashboard' are expired"""
+ user_id = "test-user"
+ current_time = datetime.now(timezone.utc)
+
+ # Create mock tokens with proper attributes
+ token1 = MagicMock()
+ token1.token = "token1"
+ token1.user_id = user_id
+ token1.team_id = "litellm-dashboard"
+ token1.blocked = None
+ token1.expires = current_time + timedelta(hours=1)
+
+ token2 = MagicMock()
+ token2.token = "token2"
+ token2.user_id = user_id
+ token2.team_id = "other-team"
+ token2.blocked = None
+ token2.expires = current_time + timedelta(hours=1)
+
+ def mock_find_many(**kwargs):
+ """Mock find_many that filters tokens based on query criteria"""
+ where_clause = kwargs.get("where", {})
+ filtered_tokens = []
+
+ for token in [token1, token2]:
+ # Check user_id match
+ if token.user_id != where_clause.get("user_id"):
+ continue
+ # Check team_id match
+ if token.team_id != where_clause.get("team_id"):
+ continue
+ # Check blocked condition (None or False)
+ if token.blocked is not None and token.blocked is not False:
+ continue
+ # Check expires > current_time
+ if token.expires <= where_clause.get("expires", {}).get("gt"):
+ continue
+ filtered_tokens.append(token)
+
+ return filtered_tokens
+
+ mock_prisma_client = MagicMock()
+ mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=mock_find_many)
+ mock_prisma_client.db.litellm_verificationtoken.update_many = AsyncMock()
+
+ await expire_previous_ui_session_tokens(user_id, mock_prisma_client)
+
+ # Should only call update_many with the litellm-dashboard token
+ mock_prisma_client.db.litellm_verificationtoken.update_many.assert_called_once_with(
+ where={"token": {"in": ["token1"]}},
+ data={"blocked": True}
+ )
+
+
+@pytest.mark.asyncio
+async def test_expire_previous_ui_session_tokens_blocks_null_and_false():
+ """Test that tokens with blocked=None and blocked=False are both processed"""
+ user_id = "test-user"
+ current_time = datetime.now(timezone.utc)
+
+ # Create mock tokens with proper attributes
+ token1 = MagicMock()
+ token1.token = "token1"
+ token1.user_id = user_id
+ token1.team_id = "litellm-dashboard"
+ token1.blocked = None
+ token1.expires = current_time + timedelta(hours=1)
+
+ token2 = MagicMock()
+ token2.token = "token2"
+ token2.user_id = user_id
+ token2.team_id = "litellm-dashboard"
+ token2.blocked = False
+ token2.expires = current_time + timedelta(hours=1)
+
+ token3 = MagicMock()
+ token3.token = "token3"
+ token3.user_id = user_id
+ token3.team_id = "litellm-dashboard"
+ token3.blocked = True # This should be ignored
+ token3.expires = current_time + timedelta(hours=1)
+
+ def mock_find_many(**kwargs):
+ """Mock find_many that filters tokens based on query criteria"""
+ where_clause = kwargs.get("where", {})
+ filtered_tokens = []
+
+ for token in [token1, token2, token3]:
+ # Check user_id match
+ if token.user_id != where_clause.get("user_id"):
+ continue
+ # Check team_id match
+ if token.team_id != where_clause.get("team_id"):
+ continue
+ # Check blocked condition (None or False)
+ if token.blocked is not None and token.blocked is not False:
+ continue
+ # Check expires > current_time
+ if token.expires <= where_clause.get("expires", {}).get("gt"):
+ continue
+ filtered_tokens.append(token)
+
+ return filtered_tokens
+
+ mock_prisma_client = MagicMock()
+ mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=mock_find_many)
+ mock_prisma_client.db.litellm_verificationtoken.update_many = AsyncMock()
+
+ await expire_previous_ui_session_tokens(user_id, mock_prisma_client)
+
+ # Should only block token1 and token2 (not token3 which is already blocked)
+ mock_prisma_client.db.litellm_verificationtoken.update_many.assert_called_once_with(
+ where={"token": {"in": ["token1", "token2"]}},
+ data={"blocked": True}
+ )
+
+
+@pytest.mark.asyncio
+async def test_expire_previous_ui_session_tokens_only_non_expired():
+ """Test that only non-expired tokens are processed"""
+ user_id = "test-user"
+ current_time = datetime.now(timezone.utc)
+
+ # Create mock tokens with proper attributes
+ token1 = MagicMock()
+ token1.token = "token1"
+ token1.user_id = user_id
+ token1.team_id = "litellm-dashboard"
+ token1.blocked = None
+ token1.expires = current_time + timedelta(hours=1) # Not expired
+
+ token2 = MagicMock()
+ token2.token = "token2"
+ token2.user_id = user_id
+ token2.team_id = "litellm-dashboard"
+ token2.blocked = None
+ token2.expires = current_time - timedelta(hours=1) # Already expired
+
+ def mock_find_many(**kwargs):
+ """Mock find_many that filters tokens based on query criteria"""
+ where_clause = kwargs.get("where", {})
+ filtered_tokens = []
+
+ for token in [token1, token2]:
+ # Check user_id match
+ if token.user_id != where_clause.get("user_id"):
+ continue
+ # Check team_id match
+ if token.team_id != where_clause.get("team_id"):
+ continue
+ # Check blocked condition (None or False)
+ if token.blocked is not None and token.blocked is not False:
+ continue
+ # Check expires > current_time
+ if token.expires <= where_clause.get("expires", {}).get("gt"):
+ continue
+ filtered_tokens.append(token)
+
+ return filtered_tokens
+
+ mock_prisma_client = MagicMock()
+ mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=mock_find_many)
+ mock_prisma_client.db.litellm_verificationtoken.update_many = AsyncMock()
+
+ await expire_previous_ui_session_tokens(user_id, mock_prisma_client)
+
+ # Should only block the non-expired token
+ mock_prisma_client.db.litellm_verificationtoken.update_many.assert_called_once_with(
+ where={"token": {"in": ["token1"]}},
+ data={"blocked": True}
+ )
+
+
+@pytest.mark.asyncio
+async def test_expire_previous_ui_session_tokens_no_tokens_found():
+ """Test behavior when no valid tokens are found"""
+ user_id = "test-user"
+
+ mock_prisma_client = MagicMock()
+ mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
+ mock_prisma_client.db.litellm_verificationtoken.update_many = AsyncMock()
+
+ await expire_previous_ui_session_tokens(user_id, mock_prisma_client)
+
+ # Should not call update_many when no tokens found
+ mock_prisma_client.db.litellm_verificationtoken.update_many.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_expire_previous_ui_session_tokens_filters_none_token():
+ """Test that tokens with None token value are filtered out"""
+ user_id = "test-user"
+ current_time = datetime.now(timezone.utc)
+
+ # Create mock tokens with proper attributes
+ token1 = MagicMock()
+ token1.token = "token1"
+ token1.user_id = user_id
+ token1.team_id = "litellm-dashboard"
+ token1.blocked = None
+ token1.expires = current_time + timedelta(hours=1)
+
+ token2 = MagicMock()
+ token2.token = None # This should be filtered out in the token collection step
+ token2.user_id = user_id
+ token2.team_id = "litellm-dashboard"
+ token2.blocked = None
+ token2.expires = current_time + timedelta(hours=1)
+
+ def mock_find_many(**kwargs):
+ """Mock find_many that filters tokens based on query criteria"""
+ where_clause = kwargs.get("where", {})
+ filtered_tokens = []
+
+ for token in [token1, token2]:
+ # Check user_id match
+ if token.user_id != where_clause.get("user_id"):
+ continue
+ # Check team_id match
+ if token.team_id != where_clause.get("team_id"):
+ continue
+ # Check blocked condition (None or False)
+ if token.blocked is not None and token.blocked is not False:
+ continue
+ # Check expires > current_time
+ if token.expires <= where_clause.get("expires", {}).get("gt"):
+ continue
+ filtered_tokens.append(token)
+
+ return filtered_tokens
+
+ mock_prisma_client = MagicMock()
+ mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=mock_find_many)
+ mock_prisma_client.db.litellm_verificationtoken.update_many = AsyncMock()
+
+ await expire_previous_ui_session_tokens(user_id, mock_prisma_client)
+
+ # Should only block token1 (token with None value should be filtered out)
+ mock_prisma_client.db.litellm_verificationtoken.update_many.assert_called_once_with(
+ where={"token": {"in": ["token1"]}},
+ data={"blocked": True}
+ )
+
+
+@pytest.mark.asyncio
+async def test_expire_previous_ui_session_tokens_exception_handling():
+ """Test that exceptions during token expiry are silently handled"""
+ user_id = "test-user"
+
+ mock_prisma_client = MagicMock()
+ mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=Exception("Database error"))
+
+ # Should not raise exception despite database error
+ await expire_previous_ui_session_tokens(user_id, mock_prisma_client)
diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
new file mode 100644
index 00000000000..f4aa28d98cc
--- /dev/null
+++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
@@ -0,0 +1,84 @@
+"""Tests for unified guardrail."""
+
+import pytest
+
+from litellm.caching import DualCache
+from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import (
+ MCPGuardrailTranslationHandler,
+)
+from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import unified_guardrail as unified_module
+from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
+ UnifiedLLMGuardrails,
+)
+from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.utils import CallTypes
+
+
+class RecordingGuardrail(CustomGuardrail):
+ """Records the event types it is asked to run for."""
+
+ def __init__(self):
+ super().__init__(guardrail_name="recording-guardrail")
+ self.event_history = []
+
+ def should_run_guardrail(self, data, event_type): # type: ignore[override]
+ self.event_history.append(event_type)
+ return True
+
+ async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
+ return {"texts": inputs.get("texts", [])}
+
+
+@pytest.fixture(autouse=True)
+def _inject_mcp_handler_mapping():
+ """Inject MCP handler mapping so the unified guardrail can run inside tests."""
+ unified_module.endpoint_guardrail_translation_mappings = {
+ CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler,
+ }
+ yield
+ unified_module.endpoint_guardrail_translation_mappings = None
+
+
+@pytest.mark.asyncio
+async def test_pre_call_hook_uses_mcp_event_type():
+ """pre_call hook should swap to GuardrailEventHooks.pre_mcp_call for MCP calls."""
+ handler = UnifiedLLMGuardrails()
+ guardrail = RecordingGuardrail()
+ cache = DualCache()
+
+ data = {
+ "guardrail_to_apply": guardrail,
+ "messages": [{"role": "user", "content": "Tool: test\nArguments: {}"}],
+ "model": "mcp-tool-call",
+ }
+
+ await handler.async_pre_call_hook(
+ user_api_key_dict=None,
+ cache=cache,
+ data=data,
+ call_type=CallTypes.call_mcp_tool.value,
+ )
+
+ assert guardrail.event_history == [GuardrailEventHooks.pre_mcp_call]
+
+
+@pytest.mark.asyncio
+async def test_moderation_hook_uses_mcp_event_type():
+ """moderation hook should request GuardrailEventHooks.during_mcp_call for MCP calls."""
+ handler = UnifiedLLMGuardrails()
+ guardrail = RecordingGuardrail()
+
+ data = {
+ "guardrail_to_apply": guardrail,
+ "messages": [{"role": "user", "content": "Tool: test\nArguments: {}"}],
+ "model": "mcp-tool-call",
+ }
+
+ await handler.async_moderation_hook(
+ data=data,
+ user_api_key_dict=None,
+ call_type=CallTypes.call_mcp_tool.value,
+ )
+
+ assert guardrail.event_history == [GuardrailEventHooks.during_mcp_call]
diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py
index b4dcc33c747..d5c3ecae7d6 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py
@@ -163,3 +163,78 @@ async def test_update_budget_allows_null_max_budget(client_and_mocks):
assert captured_data["max_budget"] is None, "max_budget should be None"
mock_table.update.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_new_budget_negative_max_budget(client_and_mocks):
+ """
+ Test that /budget/new rejects negative max_budget values.
+
+ This prevents the issue where negative budgets would always trigger
+ budget exceeded errors.
+ """
+ client, _, _ = client_and_mocks
+
+ payload = {
+ "budget_id": "budget_negative",
+ "max_budget": -7.0,
+ }
+ resp = client.post("/budget/new", json=payload)
+ assert resp.status_code == 400, resp.text
+
+ detail = resp.json()["detail"]
+ assert "max_budget cannot be negative" in str(detail)
+
+
+@pytest.mark.asyncio
+async def test_new_budget_negative_soft_budget(client_and_mocks):
+ """
+ Test that /budget/new rejects negative soft_budget values.
+ """
+ client, _, _ = client_and_mocks
+
+ payload = {
+ "budget_id": "budget_negative_soft",
+ "soft_budget": -10.0,
+ }
+ resp = client.post("/budget/new", json=payload)
+ assert resp.status_code == 400, resp.text
+
+ detail = resp.json()["detail"]
+ assert "soft_budget cannot be negative" in str(detail)
+
+
+@pytest.mark.asyncio
+async def test_update_budget_negative_max_budget(client_and_mocks):
+ """
+ Test that /budget/update rejects negative max_budget values.
+ """
+ client, _, _ = client_and_mocks
+
+ payload = {
+ "budget_id": "budget_update_negative",
+ "max_budget": -5.0,
+ }
+ resp = client.post("/budget/update", json=payload)
+ assert resp.status_code == 400, resp.text
+
+ detail = resp.json()["detail"]
+ assert "max_budget cannot be negative" in str(detail)
+
+
+@pytest.mark.asyncio
+async def test_update_budget_negative_soft_budget(client_and_mocks):
+ """
+ Test that /budget/update rejects negative soft_budget values.
+ """
+ client, _, _ = client_and_mocks
+
+ payload = {
+ "budget_id": "budget_update_negative_soft",
+ "soft_budget": -15.0,
+ }
+ resp = client.post("/budget/update", json=payload)
+ assert resp.status_code == 400, resp.text
+
+ detail = resp.json()["detail"]
+ assert "soft_budget cannot be negative" in str(detail)
diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
index 648045a7ea6..3e69b4e0caf 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
@@ -3405,3 +3405,161 @@ async def test_can_modify_verification_token_personal_key_no_user_id(monkeypatch
)
assert result is False
+
+
+@pytest.mark.asyncio
+async def test_list_keys_with_expand_user():
+ """
+ Test that expand=user parameter correctly includes user information in the response.
+ """
+ mock_prisma_client = AsyncMock()
+
+ # Create mock keys with user_ids
+ mock_key1 = MagicMock()
+ mock_key1.token = "token1"
+ mock_key1.user_id = "user123"
+ mock_key1.dict.return_value = {
+ "token": "token1",
+ "user_id": "user123",
+ "key_alias": "key1",
+ "models": ["gpt-4"],
+ }
+
+ mock_key2 = MagicMock()
+ mock_key2.token = "token2"
+ mock_key2.user_id = "user456"
+ mock_key2.dict.return_value = {
+ "token": "token2",
+ "user_id": "user456",
+ "key_alias": "key2",
+ "models": ["gpt-3.5-turbo"],
+ }
+
+ mock_find_many_keys = AsyncMock(return_value=[mock_key1, mock_key2])
+ mock_count_keys = AsyncMock(return_value=2)
+
+ # Create mock users
+ mock_user1 = MagicMock()
+ mock_user1.user_id = "user123"
+ mock_user1.user_email = "user1@example.com"
+ mock_user1.dict.return_value = {
+ "user_id": "user123",
+ "user_email": "user1@example.com",
+ "user_alias": "User One",
+ }
+
+ mock_user2 = MagicMock()
+ mock_user2.user_id = "user456"
+ mock_user2.user_email = "user2@example.com"
+ mock_user2.dict.return_value = {
+ "user_id": "user456",
+ "user_email": "user2@example.com",
+ "user_alias": "User Two",
+ }
+
+ mock_find_many_users = AsyncMock(return_value=[mock_user1, mock_user2])
+
+ mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys
+ mock_prisma_client.db.litellm_verificationtoken.count = mock_count_keys
+ mock_prisma_client.db.litellm_usertable.find_many = mock_find_many_users
+
+ args = {
+ "prisma_client": mock_prisma_client,
+ "page": 1,
+ "size": 50,
+ "user_id": None,
+ "team_id": None,
+ "organization_id": None,
+ "key_alias": None,
+ "key_hash": None,
+ "exclude_team_id": None,
+ "return_full_object": False, # This should be overridden by expand=user
+ "admin_team_ids": None,
+ "include_created_by_keys": False,
+ "expand": ["user"], # Test the expand parameter
+ }
+
+ result = await _list_key_helper(**args)
+
+ # Verify that keys were fetched
+ mock_find_many_keys.assert_called_once()
+ mock_count_keys.assert_called_once()
+
+ # Verify that users were fetched
+ # Note: Order doesn't matter for the 'in' query, so we just check that both user_ids are present
+ call_args = mock_find_many_users.call_args
+ assert call_args is not None
+ where_clause = call_args.kwargs["where"]
+ assert "user_id" in where_clause
+ assert "in" in where_clause["user_id"]
+ user_ids_in_query = set(where_clause["user_id"]["in"])
+ assert user_ids_in_query == {"user123", "user456"}
+
+ # Verify response structure
+ assert len(result["keys"]) == 2
+ assert result["total_count"] == 2
+ assert result["current_page"] == 1
+ assert result["total_pages"] == 1
+
+ # Verify that user data is included in the response
+ # Since expand=user is specified, keys should be full objects
+ assert isinstance(result["keys"][0], UserAPIKeyAuth)
+ assert isinstance(result["keys"][1], UserAPIKeyAuth)
+
+ # Verify user data is attached to keys
+ assert result["keys"][0].user == {
+ "user_id": "user123",
+ "user_email": "user1@example.com",
+ "user_alias": "User One",
+ }
+ assert result["keys"][1].user == {
+ "user_id": "user456",
+ "user_email": "user2@example.com",
+ "user_alias": "User Two",
+ }
+
+
+@pytest.mark.asyncio
+async def test_generate_key_negative_max_budget():
+ """
+ Test that GenerateKeyRequest model allows negative max_budget values.
+ Validation is done at API level, not model level.
+
+ This prevents GET requests from breaking when they receive data with negative budgets.
+ """
+ # Should not raise any errors at model level
+ request = GenerateKeyRequest(max_budget=-7.0)
+ assert request.max_budget == -7.0
+
+
+@pytest.mark.asyncio
+async def test_generate_key_negative_soft_budget():
+ """
+ Test that GenerateKeyRequest model allows negative soft_budget values.
+ Validation is done at API level, not model level.
+ """
+ # Should not raise any errors at model level
+ request = GenerateKeyRequest(soft_budget=-10.0)
+ assert request.soft_budget == -10.0
+
+
+@pytest.mark.asyncio
+async def test_generate_key_positive_budgets_accepted():
+ """
+ Test that GenerateKeyRequest accepts positive budget values.
+ """
+ # Should not raise any errors
+ request = GenerateKeyRequest(max_budget=100.0, soft_budget=50.0)
+ assert request.max_budget == 100.0
+ assert request.soft_budget == 50.0
+
+
+@pytest.mark.asyncio
+async def test_update_key_negative_max_budget():
+ """
+ Test that UpdateKeyRequest model allows negative max_budget values.
+ Validation is done at API level, not model level.
+ """
+ # Should not raise any errors at model level
+ request = UpdateKeyRequest(key="test-key", max_budget=-5.0)
+ assert request.max_budget == -5.0
diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
index 61342e8025b..cc268ab9925 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
@@ -169,8 +169,8 @@ class TestListMCPServers:
return_value=["config_server_1", "config_server_2"]
)
- # Mock the new method that returns servers with health and team data
- mock_servers_with_health = [
+ # Mock the new method that returns servers without health check
+ mock_servers = [
generate_mock_mcp_server_db_record(
server_id="config_server_1",
alias="Zapier MCP",
@@ -184,11 +184,11 @@ class TestListMCPServers:
transport="http",
),
]
- mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock(
- return_value=mock_servers_with_health
+ mock_manager.get_all_allowed_mcp_servers = AsyncMock(
+ return_value=mock_servers
)
- for idx, server in enumerate(mock_servers_with_health):
+ for idx, server in enumerate(mock_servers):
server.credentials = {"auth_value": f"secret_{idx}"}
with patch(
@@ -200,6 +200,9 @@ class TestListMCPServers:
), patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=mock_prisma_client,
+ ), patch(
+ "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
+ AsyncMock(return_value=[mock_user_auth]),
):
# Import and call the function
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
@@ -300,8 +303,8 @@ class TestListMCPServers:
]
)
- # Mock the new method that returns servers with health and team data
- mock_servers_with_health = [
+ # Mock the new method that returns servers without health check
+ mock_servers = [
db_server_1,
db_server_2,
generate_mock_mcp_server_db_record(
@@ -317,11 +320,11 @@ class TestListMCPServers:
transport="http",
),
]
- mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock(
- return_value=mock_servers_with_health
+ mock_manager.get_all_allowed_mcp_servers = AsyncMock(
+ return_value=mock_servers
)
- for idx, server in enumerate(mock_servers_with_health):
+ for idx, server in enumerate(mock_servers):
server.credentials = {"auth_value": f"secret_{idx}"}
with patch(
@@ -333,6 +336,9 @@ class TestListMCPServers:
), patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=mock_prisma_client,
+ ), patch(
+ "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
+ AsyncMock(return_value=[mock_user_auth]),
):
# Import and call the function
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
@@ -425,8 +431,8 @@ class TestListMCPServers:
return_value=["db_server_allowed", "config_server_allowed"]
)
- # Mock the new method that returns servers with health and team data
- mock_servers_with_health = [
+ # Mock the new method that returns servers without health check
+ mock_servers = [
db_server_allowed,
generate_mock_mcp_server_db_record(
server_id="config_server_allowed",
@@ -434,11 +440,11 @@ class TestListMCPServers:
url="https://actions.zapier.com/mcp/sse",
),
]
- mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock(
- return_value=mock_servers_with_health
+ mock_manager.get_all_allowed_mcp_servers = AsyncMock(
+ return_value=mock_servers
)
- for idx, server in enumerate(mock_servers_with_health):
+ for idx, server in enumerate(mock_servers):
server.credentials = {"auth_value": f"secret_{idx}"}
with patch(
@@ -450,6 +456,9 @@ class TestListMCPServers:
), patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=mock_prisma_client,
+ ), patch(
+ "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
+ AsyncMock(return_value=[mock_user_auth]),
):
# Import and call the function
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
@@ -486,11 +495,14 @@ class TestListMCPServers:
mock_server.credentials = {"auth_value": "top-secret"}
mock_prisma_client = MagicMock()
- mock_health_result = {
- "status": "healthy",
- "last_health_check": datetime.now().isoformat(),
- "error": None,
- }
+
+ # Mock health check result as LiteLLM_MCPServerTable
+ mock_health_result = generate_mock_mcp_server_db_record(
+ server_id="server-1", alias="Server 1"
+ )
+ mock_health_result.status = "healthy"
+ mock_health_result.last_health_check = datetime.now()
+ mock_health_result.health_check_error = None
mock_user_auth = generate_mock_user_api_key_auth(
user_role=LitellmUserRoles.PROXY_ADMIN
@@ -531,11 +543,14 @@ class TestListMCPServers:
delattr(mock_server, "credentials")
mock_prisma_client = MagicMock()
- mock_health_result = {
- "status": "healthy",
- "last_health_check": datetime.now().isoformat(),
- "error": None,
- }
+
+ # Mock health check result as LiteLLM_MCPServerTable
+ mock_health_result = generate_mock_mcp_server_db_record(
+ server_id="server-2", alias="Server 2"
+ )
+ mock_health_result.status = "healthy"
+ mock_health_result.last_health_check = datetime.now()
+ mock_health_result.health_check_error = None
mock_user_auth = generate_mock_user_api_key_auth(
user_role=LitellmUserRoles.PROXY_ADMIN
@@ -568,296 +583,6 @@ class TestListMCPServers:
assert result.status == "healthy"
-class TestMCPHealthCheckEndpoints:
- """Test MCP health check endpoints"""
-
- @pytest.mark.asyncio
- async def test_health_check_mcp_server_success(self):
- """Test successful health check for a specific MCP server"""
- # Mock server
- mock_server = generate_mock_mcp_server_db_record(
- server_id="test-server", alias="Test Server"
- )
-
- # Mock dependencies
- mock_prisma_client = MagicMock()
-
- # Mock global MCP server manager
- mock_manager = MagicMock()
- mock_manager.health_check_server = AsyncMock(
- return_value={
- "server_id": "test-server",
- "server_name": "Test Server",
- "status": "healthy",
- "tools_count": 3,
- "last_health_check": "2024-01-01T12:00:00",
- "response_time_ms": 150.5,
- "error": None,
- }
- )
-
- mock_user_auth = generate_mock_user_api_key_auth(
- user_role=LitellmUserRoles.PROXY_ADMIN
- )
-
- with patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
- return_value=mock_prisma_client,
- ), patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view",
- return_value=True,
- ), patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
- mock_manager,
- ), patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
- AsyncMock(return_value=mock_server),
- ):
- # Import and call the function
- from litellm.proxy.management_endpoints.mcp_management_endpoints import (
- health_check_mcp_server,
- )
-
- result = await health_check_mcp_server(
- server_id="test-server", user_api_key_dict=mock_user_auth
- )
-
- # Verify results
- assert result["server_id"] == "test-server"
- assert result["server_name"] == "Test Server"
- assert result["status"] == "healthy"
- assert result["tools_count"] == 3
- assert result["response_time_ms"] == 150.5
- assert result["error"] is None
-
- @pytest.mark.asyncio
- async def test_health_check_mcp_server_not_found(self):
- """Test health check for a server that doesn't exist"""
- # Mock dependencies
- mock_prisma_client = MagicMock()
-
- mock_user_auth = generate_mock_user_api_key_auth(
- user_role=LitellmUserRoles.PROXY_ADMIN
- )
-
- with patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
- return_value=mock_prisma_client,
- ), patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
- AsyncMock(return_value=None),
- ):
- # Import and call the function
- from litellm.proxy.management_endpoints.mcp_management_endpoints import (
- health_check_mcp_server,
- )
-
- # Should raise HTTPException
- with pytest.raises(Exception) as exc_info:
- await health_check_mcp_server(
- server_id="non-existent-server", user_api_key_dict=mock_user_auth
- )
-
- assert "not found" in str(exc_info.value)
-
- @pytest.mark.asyncio
- async def test_health_check_mcp_server_unauthorized(self):
- """Test health check for a server user doesn't have access to"""
- # Mock server
- mock_server = generate_mock_mcp_server_db_record(
- server_id="test-server", alias="Test Server"
- )
-
- # Mock dependencies
- mock_prisma_client = MagicMock()
-
- mock_user_auth = generate_mock_user_api_key_auth(
- user_role=LitellmUserRoles.INTERNAL_USER # Non-admin user
- )
-
- # Mock user doesn't have access to this server
- mock_user_servers = []
-
- with patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
- return_value=mock_prisma_client,
- ), patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view",
- return_value=False,
- ), patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers_for_user",
- return_value=mock_user_servers,
- ), patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
- AsyncMock(return_value=mock_server),
- ):
- # Import and call the function
- from litellm.proxy.management_endpoints.mcp_management_endpoints import (
- health_check_mcp_server,
- )
-
- # Should raise HTTPException
- with pytest.raises(Exception) as exc_info:
- await health_check_mcp_server(
- server_id="test-server", user_api_key_dict=mock_user_auth
- )
-
- assert "permission" in str(exc_info.value)
-
- @pytest.mark.asyncio
- async def test_health_check_all_mcp_servers(self):
- """Test health check for all accessible MCP servers"""
- # Mock team records
- team_records = [
- generate_mock_team_record(
- team_id="team1",
- team_alias="Team 1",
- organization_id="org1",
- mcp_servers=["server1", "server2"],
- )
- ]
-
- # Mock DB servers
- db_servers = [
- generate_mock_mcp_server_db_record(server_id="server1"),
- generate_mock_mcp_server_db_record(server_id="server2"),
- ]
-
- # Mock dependencies
- mock_prisma_client = MagicMock()
- mock_prisma_client = setup_mock_prisma_client(
- mock_prisma_client=mock_prisma_client,
- team_records=team_records,
- mcp_servers=db_servers,
- )
-
- # Mock global MCP server manager
- mock_manager = MagicMock()
- mock_manager.health_check_allowed_servers = AsyncMock(
- return_value={
- "server1": {
- "server_id": "server1",
- "server_name": "Test DB Server",
- "status": "healthy",
- "tools_count": 2,
- "last_health_check": "2024-01-01T12:00:00",
- "response_time_ms": 100.0,
- "error": None,
- },
- "server2": {
- "server_id": "server2",
- "server_name": "Test DB Server",
- "status": "unhealthy",
- "last_health_check": "2024-01-01T12:00:00",
- "response_time_ms": 5000.0,
- "error": "Connection timeout",
- },
- }
- )
- mock_manager.get_allowed_mcp_servers = AsyncMock(
- return_value=["server1", "server2"]
- )
-
- mock_user_auth = generate_mock_user_api_key_auth(
- user_role=LitellmUserRoles.INTERNAL_USER
- )
-
- with patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
- return_value=mock_prisma_client,
- ), patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view",
- return_value=False,
- ), patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
- mock_manager,
- ):
- # Import and call the function
- from litellm.proxy.management_endpoints.mcp_management_endpoints import (
- health_check_all_mcp_servers,
- )
-
- result = await health_check_all_mcp_servers(
- user_api_key_dict=mock_user_auth
- )
-
- # Verify results
- assert result["total_servers"] == 2
- assert result["healthy_count"] == 1
- assert result["unhealthy_count"] == 1
- assert result["unknown_count"] == 0
- assert "server1" in result["servers"]
- assert "server2" in result["servers"]
-
- # Check individual server results
- assert result["servers"]["server1"]["status"] == "healthy"
- assert result["servers"]["server1"]["tools_count"] == 2
- assert result["servers"]["server1"]["server_name"] == "Test DB Server"
- assert result["servers"]["server2"]["status"] == "unhealthy"
- assert result["servers"]["server2"]["error"] == "Connection timeout"
- assert result["servers"]["server2"]["server_name"] == "Test DB Server"
-
- @pytest.mark.asyncio
- async def test_fetch_all_mcp_servers_with_health_status(self):
- """Test that fetch_all_mcp_servers includes health check status"""
- # Mock server with health status
- mock_server = generate_mock_mcp_server_db_record(
- server_id="test-server", alias="Test Server"
- )
- # Add health status to the mock server
- mock_server.status = "healthy"
- mock_server.last_health_check = datetime.now()
- mock_server.health_check_error = None
-
- # Mock dependencies
- mock_prisma_client = MagicMock()
- mock_prisma_client = setup_mock_prisma_client(
- mock_prisma_client=mock_prisma_client,
- team_records=[],
- mcp_servers=[], # Don't add servers here since we're mocking get_all_mcp_servers
- )
-
- # Mock global MCP server manager
- mock_manager = MagicMock()
- mock_manager.config_mcp_servers = {}
- mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=[])
- mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock(
- return_value=[mock_server]
- )
-
- mock_server.credentials = {"auth_value": "secret"}
-
- mock_user_auth = generate_mock_user_api_key_auth(
- user_role=LitellmUserRoles.PROXY_ADMIN
- )
-
- with patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
- return_value=mock_prisma_client,
- ), patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view",
- return_value=True,
- ), patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
- mock_manager,
- ):
- # Import and call the function
- from litellm.proxy.management_endpoints.mcp_management_endpoints import (
- fetch_all_mcp_servers,
- )
-
- result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth)
-
- # Verify health check status is included
- assert len(result) == 1
- server = result[0]
- assert server.server_id == "test-server"
- assert server.status == "healthy"
- assert server.last_health_check is not None
- assert server.health_check_error is None
- assert server.credentials is None
-
-
class TestTemporaryMCPSessionEndpoints:
def test_inherit_credentials_from_existing_server(self):
payload = NewMCPServerRequest(
@@ -1170,7 +895,6 @@ class TestTemporaryMCPSessionEndpoints:
fallback_client_id="server-1",
)
-
class TestUpdateMCPServer:
"""Test suite for update MCP server functionality"""
@@ -1233,7 +957,7 @@ class TestUpdateMCPServer:
"litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server",
AsyncMock(return_value=updated_server),
) as update_mock, patch(
- "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.add_update_server",
+ "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.add_server",
AsyncMock(),
), patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.reload_servers_from_database",
@@ -1260,3 +984,163 @@ class TestUpdateMCPServer:
# Verify the result includes extra_headers
assert result.extra_headers == ["X-Custom-Header", "X-Another-Header"]
assert result.alias == "Updated Test Server"
+
+
+class TestHealthCheckServers:
+ """Test suite for health check servers endpoint"""
+
+ @pytest.mark.asyncio
+ async def test_health_check_all_servers(self):
+ """
+ Test health check for all accessible servers
+
+ Scenario: User has access to 2 servers, checks all
+ Expected: Returns health status for both servers
+ """
+ from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+ health_check_servers,
+ )
+
+ # Mock user auth
+ mock_user_auth = generate_mock_user_api_key_auth()
+
+ # Mock health check results
+ mock_health_result_1 = generate_mock_mcp_server_db_record(
+ server_id="server-1",
+ alias="Server 1",
+ url="https://server1.example.com",
+ )
+ mock_health_result_1.status = "healthy"
+ mock_health_result_1.last_health_check = datetime.now()
+ mock_health_result_1.health_check_error = None
+
+ mock_health_result_2 = generate_mock_mcp_server_db_record(
+ server_id="server-2",
+ alias="Server 2",
+ url="https://server2.example.com",
+ )
+ mock_health_result_2.status = "unhealthy"
+ mock_health_result_2.last_health_check = datetime.now()
+ mock_health_result_2.health_check_error = "Connection timeout"
+
+ # Mock manager
+ mock_manager = MagicMock()
+ mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock(
+ return_value=[mock_health_result_1, mock_health_result_2]
+ )
+
+ with patch(
+ "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
+ mock_manager,
+ ), patch(
+ "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
+ AsyncMock(return_value=[mock_user_auth]),
+ ):
+ result = await health_check_servers(
+ server_ids=None,
+ user_api_key_dict=mock_user_auth,
+ )
+
+ # Verify results
+ assert len(result) == 2
+ assert result[0]["server_id"] == "server-1"
+ assert result[0]["status"] == "healthy"
+ assert result[1]["server_id"] == "server-2"
+ assert result[1]["status"] == "unhealthy"
+
+ @pytest.mark.asyncio
+ async def test_health_check_specific_servers(self):
+ """
+ Test health check for specific servers
+
+ Scenario: User requests health check for specific server IDs
+ Expected: Returns health status only for requested servers
+ """
+ from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+ health_check_servers,
+ )
+
+ # Mock user auth
+ mock_user_auth = generate_mock_user_api_key_auth()
+
+ # Mock health check result
+ mock_health_result = generate_mock_mcp_server_db_record(
+ server_id="server-1",
+ alias="Server 1",
+ url="https://server1.example.com",
+ )
+ mock_health_result.status = "healthy"
+ mock_health_result.last_health_check = datetime.now()
+ mock_health_result.health_check_error = None
+
+ # Mock manager
+ mock_manager = MagicMock()
+ mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock(
+ return_value=[mock_health_result]
+ )
+
+ with patch(
+ "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
+ mock_manager,
+ ), patch(
+ "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
+ AsyncMock(return_value=[mock_user_auth]),
+ ):
+ result = await health_check_servers(
+ server_ids=["server-1"],
+ user_api_key_dict=mock_user_auth,
+ )
+
+ # Verify results
+ assert len(result) == 1
+ assert result[0]["server_id"] == "server-1"
+ assert result[0]["status"] == "healthy"
+
+ @pytest.mark.asyncio
+ async def test_health_check_unauthorized_servers(self):
+ """
+ Test health check with unauthorized servers
+
+ Scenario: User requests health check for servers they don't have access to
+ Expected: Only checks accessible servers, unauthorized servers are filtered out
+ """
+ from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+ health_check_servers,
+ )
+
+ # Mock user auth
+ mock_user_auth = generate_mock_user_api_key_auth()
+
+ # Mock health check result for authorized server
+ mock_health_result = generate_mock_mcp_server_db_record(
+ server_id="server-1",
+ alias="Server 1",
+ url="https://server1.example.com",
+ )
+ mock_health_result.status = "healthy"
+ mock_health_result.last_health_check = datetime.now()
+ mock_health_result.health_check_error = None
+
+ # Mock manager - server_ids filter is applied inside get_all_mcp_servers_with_health_and_teams
+ # So it only returns servers the user has access to
+ mock_manager = MagicMock()
+ mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock(
+ return_value=[mock_health_result] # Only server-1 is returned (accessible)
+ )
+
+ with patch(
+ "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
+ mock_manager,
+ ), patch(
+ "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
+ AsyncMock(return_value=[mock_user_auth]),
+ ):
+ result = await health_check_servers(
+ server_ids=["server-1", "server-unauthorized"],
+ user_api_key_dict=mock_user_auth,
+ )
+
+ # Verify results - only accessible server is returned
+ assert len(result) == 1
+ assert result[0]["server_id"] == "server-1"
+ assert result[0]["status"] == "healthy"
diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
index 83b4fc35a0d..6cf8f745e07 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
@@ -3958,3 +3958,74 @@ async def test_update_team_guardrails_with_org_id():
assert "include" in first_call_kwargs
assert "teams" in first_call_kwargs["include"]
assert first_call_kwargs["include"]["teams"] is True
+
+
+@pytest.mark.asyncio
+async def test_new_team_negative_max_budget():
+ """
+ Test that NewTeamRequest model allows negative max_budget values.
+ Validation is done at API level, not model level.
+
+ This prevents GET requests from breaking when they receive data with negative budgets.
+ """
+ from litellm.proxy._types import NewTeamRequest
+
+ # Should not raise any errors at model level
+ request = NewTeamRequest(team_alias="test-team", max_budget=-7.0)
+ assert request.max_budget == -7.0
+
+
+@pytest.mark.asyncio
+async def test_new_team_negative_team_member_budget():
+ """
+ Test that NewTeamRequest model allows negative team_member_budget values.
+ Validation is done at API level, not model level.
+ """
+ from litellm.proxy._types import NewTeamRequest
+
+ # Should not raise any errors at model level
+ request = NewTeamRequest(team_alias="test-team", team_member_budget=-10.0)
+ assert request.team_member_budget == -10.0
+
+
+@pytest.mark.asyncio
+async def test_update_team_negative_max_budget():
+ """
+ Test that UpdateTeamRequest model allows negative max_budget values.
+ Validation is done at API level, not model level.
+ """
+ from litellm.proxy._types import UpdateTeamRequest
+
+ # Should not raise any errors at model level
+ request = UpdateTeamRequest(team_id="test-team-id", max_budget=-5.0)
+ assert request.max_budget == -5.0
+
+
+@pytest.mark.asyncio
+async def test_update_team_negative_team_member_budget():
+ """
+ Test that UpdateTeamRequest model allows negative team_member_budget values.
+ Validation is done at API level, not model level.
+ """
+ from litellm.proxy._types import UpdateTeamRequest
+
+ # Should not raise any errors at model level
+ request = UpdateTeamRequest(team_id="test-team-id", team_member_budget=-15.0)
+ assert request.team_member_budget == -15.0
+
+
+@pytest.mark.asyncio
+async def test_new_team_positive_budgets_accepted():
+ """
+ Test that NewTeamRequest accepts positive budget values.
+ """
+ from litellm.proxy._types import NewTeamRequest
+
+ # Should not raise any errors
+ request = NewTeamRequest(
+ team_alias="test-team",
+ max_budget=100.0,
+ team_member_budget=50.0
+ )
+ assert request.max_budget == 100.0
+ assert request.team_member_budget == 50.0
diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
index f6a9b5bddb3..213297fc80a 100644
--- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
+++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
@@ -78,3 +78,26 @@ def test_get_litellm_model_cost_map_returns_cost_map():
# Check for common cost fields that should be present
assert "input_cost_per_token" in sample_model_data or "output_cost_per_token" in sample_model_data
+
+def test_watsonx_provider_fields():
+ """Test that Watsonx provider has all required credential fields including multiple auth options."""
+ app = FastAPI()
+ app.include_router(router)
+ client = TestClient(app)
+
+ response = client.get("/public/providers/fields")
+ providers = response.json()
+
+ watsonx = next((p for p in providers if p["provider"] == "WATSONX"), None)
+ assert watsonx is not None
+
+ field_keys = [f["key"] for f in watsonx["credential_fields"]]
+ # Core fields
+ assert "api_base" in field_keys
+ assert "project_id" in field_keys
+ assert "space_id" in field_keys
+ # Multiple auth methods supported
+ assert "api_key" in field_keys
+ assert "token" in field_keys
+ assert "zen_api_key" in field_keys
+
diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
index 5e3652c6d9d..56bba39e6c3 100644
--- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
+++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
@@ -202,6 +202,7 @@ ignored_keys = [
"metadata.cold_storage_object_key",
"metadata.additional_usage_values.prompt_tokens_details.cache_creation_tokens",
"metadata.litellm_overhead_time_ms",
+ "metadata.cost_breakdown",
]
MODEL_LIST = [
diff --git a/tests/test_litellm/proxy/test_empty_model_list.py b/tests/test_litellm/proxy/test_empty_model_list.py
new file mode 100644
index 00000000000..6b3e59d3194
--- /dev/null
+++ b/tests/test_litellm/proxy/test_empty_model_list.py
@@ -0,0 +1,155 @@
+"""
+Tests for graceful handling of empty model list scenarios.
+
+These tests verify that /v2/model/info and /model_group/info endpoints
+return empty data arrays instead of 500 errors when no models are configured.
+"""
+
+import os
+import sys
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from fastapi.testclient import TestClient
+
+sys.path.insert(
+ 0, os.path.abspath("../../..")
+) # Adds the parent directory to the system-path
+
+from litellm.proxy.proxy_server import app
+
+
+@pytest.fixture
+def client():
+ """Create a test client for the FastAPI app."""
+ return TestClient(app)
+
+
+class TestEmptyModelListHandling:
+ """Test suite for empty model list scenarios."""
+
+ def test_v2_model_info_returns_empty_data_when_router_is_none(
+ self, client, monkeypatch
+ ):
+ """
+ Test that /v2/model/info returns {"data": []} instead of 500
+ when llm_router is None.
+ """
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_model_list", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
+
+ with patch(
+ "litellm.proxy.auth.user_api_key_auth.user_api_key_auth",
+ return_value=MagicMock(
+ user_id="test-user",
+ team_id=None,
+ team_models=[],
+ models=[],
+ user_role="proxy_admin",
+ ),
+ ):
+ response = client.get(
+ "/v2/model/info",
+ headers={"Authorization": "Bearer sk-test"},
+ )
+
+ assert response.status_code == 200
+ assert response.json() == {"data": []}
+
+ def test_v2_model_info_returns_empty_data_when_model_list_empty(
+ self, client, monkeypatch
+ ):
+ """
+ Test that /v2/model/info returns {"data": []} instead of 500
+ when llm_router exists but model_list is empty.
+ """
+ mock_router = MagicMock()
+ mock_router.model_list = []
+
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router)
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_model_list", [])
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
+
+ with patch(
+ "litellm.proxy.auth.user_api_key_auth.user_api_key_auth",
+ return_value=MagicMock(
+ user_id="test-user",
+ team_id=None,
+ team_models=[],
+ models=[],
+ user_role="proxy_admin",
+ ),
+ ):
+ response = client.get(
+ "/v2/model/info",
+ headers={"Authorization": "Bearer sk-test"},
+ )
+
+ assert response.status_code == 200
+ assert response.json() == {"data": []}
+
+ def test_model_group_info_returns_empty_data_when_model_list_none(
+ self, client, monkeypatch
+ ):
+ """
+ Test that /model_group/info returns {"data": []} instead of 500
+ when llm_model_list is None.
+ """
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_model_list", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
+
+ with patch(
+ "litellm.proxy.auth.user_api_key_auth.user_api_key_auth",
+ return_value=MagicMock(
+ user_id="test-user",
+ team_id=None,
+ team_models=[],
+ models=[],
+ user_role="proxy_admin",
+ ),
+ ):
+ response = client.get(
+ "/model_group/info",
+ headers={"Authorization": "Bearer sk-test"},
+ )
+
+ assert response.status_code == 200
+ assert response.json() == {"data": []}
+
+ def test_model_group_info_returns_empty_data_when_model_list_empty(
+ self, client, monkeypatch
+ ):
+ """
+ Test that /model_group/info returns {"data": []} instead of 500
+ when llm_model_list is empty.
+ """
+ mock_router = MagicMock()
+ mock_router.model_list = []
+
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router)
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_model_list", [])
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
+
+ with patch(
+ "litellm.proxy.auth.user_api_key_auth.user_api_key_auth",
+ return_value=MagicMock(
+ user_id="test-user",
+ team_id=None,
+ team_models=[],
+ models=[],
+ user_role="proxy_admin",
+ ),
+ ):
+ response = client.get(
+ "/model_group/info",
+ headers={"Authorization": "Bearer sk-test"},
+ )
+
+ assert response.status_code == 200
+ assert response.json() == {"data": []}
diff --git a/tests/test_litellm/test_gpt_image_cost_calculator.py b/tests/test_litellm/test_gpt_image_cost_calculator.py
new file mode 100644
index 00000000000..0a2a62b6c97
--- /dev/null
+++ b/tests/test_litellm/test_gpt_image_cost_calculator.py
@@ -0,0 +1,239 @@
+"""
+Tests for OpenAI gpt-image-1 cost calculator
+
+This tests the fix for GitHub issue #13847:
+https://github.com/BerriAI/litellm/issues/13847
+
+gpt-image-1 uses token-based pricing:
+- Text Input: $5.00/1M tokens
+- Image Input: $10.00/1M tokens
+- Image Output: $40.00/1M tokens
+"""
+
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath("../.."))
+
+import pytest
+
+import litellm
+from litellm.types.utils import (
+ ImageResponse,
+ ImageObject,
+ ImageUsage,
+ ImageUsageInputTokensDetails,
+)
+
+
+class TestGPTImageCostCalculator:
+ """Test the OpenAI gpt-image-1 cost calculator"""
+
+ def test_gpt_image_1_cost_with_text_only(self):
+ """Test cost calculation with only text input tokens"""
+ from litellm.llms.openai.image_generation.cost_calculator import cost_calculator
+
+ usage = ImageUsage(
+ input_tokens=100,
+ output_tokens=5000,
+ total_tokens=5100,
+ input_tokens_details=ImageUsageInputTokensDetails(
+ text_tokens=100,
+ image_tokens=0,
+ ),
+ )
+
+ image_response = ImageResponse(
+ created=1234567890,
+ data=[ImageObject(url="http://example.com/image.jpg")],
+ )
+ image_response.usage = usage
+
+ cost = cost_calculator(
+ model="gpt-image-1",
+ image_response=image_response,
+ custom_llm_provider="openai",
+ )
+
+ # Expected cost:
+ # Text input: 100 * $5/1M = 0.0005
+ # Image output: 5000 * $40/1M = 0.2
+ # Total: 0.2005
+ expected_cost = 0.0005 + 0.2
+ assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
+
+ def test_gpt_image_1_cost_with_image_input(self):
+ """Test cost calculation with both text and image input tokens (for edits)"""
+ from litellm.llms.openai.image_generation.cost_calculator import cost_calculator
+
+ usage = ImageUsage(
+ input_tokens=600,
+ output_tokens=5000,
+ total_tokens=5600,
+ input_tokens_details=ImageUsageInputTokensDetails(
+ text_tokens=100,
+ image_tokens=500,
+ ),
+ )
+
+ image_response = ImageResponse(
+ created=1234567890,
+ data=[ImageObject(url="http://example.com/image.jpg")],
+ )
+ image_response.usage = usage
+
+ cost = cost_calculator(
+ model="gpt-image-1",
+ image_response=image_response,
+ custom_llm_provider="openai",
+ )
+
+ # Expected cost:
+ # Text input: 100 * $5/1M = 0.0005
+ # Image input: 500 * $10/1M = 0.005
+ # Image output: 5000 * $40/1M = 0.2
+ # Total: 0.2055
+ expected_cost = 0.0005 + 0.005 + 0.2
+ assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
+
+ def test_gpt_image_1_mini_cost(self):
+ """Test cost calculation for gpt-image-1-mini model"""
+ from litellm.llms.openai.image_generation.cost_calculator import cost_calculator
+
+ usage = ImageUsage(
+ input_tokens=100,
+ output_tokens=5000,
+ total_tokens=5100,
+ input_tokens_details=ImageUsageInputTokensDetails(
+ text_tokens=100,
+ image_tokens=0,
+ ),
+ )
+
+ image_response = ImageResponse(
+ created=1234567890,
+ data=[ImageObject(url="http://example.com/image.jpg")],
+ )
+ image_response.usage = usage
+
+ cost = cost_calculator(
+ model="gpt-image-1-mini",
+ image_response=image_response,
+ custom_llm_provider="openai",
+ )
+
+ # Expected cost for gpt-image-1-mini:
+ # Text input: 100 * $2/1M = 0.0002
+ # Image output: 5000 * $8/1M = 0.04
+ # Total: 0.0402
+ expected_cost = 0.0002 + 0.04
+ assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
+
+ def test_gpt_image_1_cost_no_usage(self):
+ """Test that cost returns 0 when no usage data is available"""
+ from litellm.llms.openai.image_generation.cost_calculator import cost_calculator
+
+ image_response = ImageResponse(
+ created=1234567890,
+ data=[ImageObject(url="http://example.com/image.jpg")],
+ )
+
+ cost = cost_calculator(
+ model="gpt-image-1",
+ image_response=image_response,
+ custom_llm_provider="openai",
+ )
+
+ assert cost == 0.0
+
+
+class TestGPTImageCostRouting:
+ """Test that gpt-image models are properly routed to the token-based calculator"""
+
+ def test_openai_gpt_image_routes_to_token_calculator(self):
+ """Test that OpenAI gpt-image-1 routes to token-based calculator"""
+ from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils
+
+ usage = ImageUsage(
+ input_tokens=100,
+ output_tokens=5000,
+ total_tokens=5100,
+ input_tokens_details=ImageUsageInputTokensDetails(
+ text_tokens=100,
+ image_tokens=0,
+ ),
+ )
+
+ image_response = ImageResponse(
+ created=1234567890,
+ data=[ImageObject(url="http://example.com/image.jpg")],
+ )
+ image_response.usage = usage
+
+ cost = CostCalculatorUtils.route_image_generation_cost_calculator(
+ model="gpt-image-1",
+ completion_response=image_response,
+ custom_llm_provider="openai",
+ )
+
+ expected_cost = 0.0005 + 0.2
+ assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
+
+ def test_openai_dalle_routes_to_pixel_calculator(self):
+ """Test that OpenAI DALL-E still routes to pixel-based calculator"""
+ from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils
+
+ image_response = ImageResponse(
+ created=1234567890,
+ data=[ImageObject(url="http://example.com/image.jpg")],
+ )
+ image_response.size = "1024x1024"
+ image_response.quality = "standard"
+
+ cost = CostCalculatorUtils.route_image_generation_cost_calculator(
+ model="dall-e-3",
+ completion_response=image_response,
+ custom_llm_provider="openai",
+ size="1024x1024",
+ quality="standard",
+ n=1,
+ )
+
+ assert cost >= 0
+
+
+class TestCompletionCostIntegration:
+ """Test the full completion_cost integration for gpt-image-1"""
+
+ def test_completion_cost_gpt_image_1(self):
+ """Test completion_cost correctly calculates gpt-image-1 costs"""
+ usage = ImageUsage(
+ input_tokens=100,
+ output_tokens=5000,
+ total_tokens=5100,
+ input_tokens_details=ImageUsageInputTokensDetails(
+ text_tokens=100,
+ image_tokens=0,
+ ),
+ )
+
+ image_response = ImageResponse(
+ created=1234567890,
+ data=[ImageObject(url="http://example.com/image.jpg")],
+ )
+ image_response.usage = usage
+ image_response._hidden_params = {"custom_llm_provider": "openai"}
+
+ cost = litellm.completion_cost(
+ completion_response=image_response,
+ model="gpt-image-1",
+ call_type="image_generation",
+ custom_llm_provider="openai",
+ )
+
+ expected_cost = 0.0005 + 0.2
+ assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/tests/test_litellm/test_lazy_imports.py b/tests/test_litellm/test_lazy_imports.py
index 0eaedaab601..7de99495b79 100644
--- a/tests/test_litellm/test_lazy_imports.py
+++ b/tests/test_litellm/test_lazy_imports.py
@@ -33,6 +33,8 @@ from litellm._lazy_imports import (
_lazy_import_llm_configs,
TYPES_NAMES,
_lazy_import_types,
+ LLM_PROVIDER_LOGIC_NAMES,
+ _lazy_import_llm_provider_logic,
)
@@ -218,6 +220,9 @@ def test_unknown_attribute_raises_error():
with pytest.raises(AttributeError):
_lazy_import_types("unknown")
+ with pytest.raises(AttributeError):
+ _lazy_import_llm_provider_logic("unknown")
+
def test_llm_config_lazy_imports():
"""Test that LLM config classes can be lazy imported."""
@@ -246,3 +251,16 @@ def test_types_lazy_imports():
_verify_only_requested_name_imported(name, TYPES_NAMES)
+
+def test_llm_provider_logic_lazy_imports():
+ """Test that LLM provider logic functions can be lazy imported."""
+ for name in LLM_PROVIDER_LOGIC_NAMES:
+ _clear_names_from_globals(LLM_PROVIDER_LOGIC_NAMES)
+
+ func = _lazy_import_llm_provider_logic(name)
+ assert func is not None
+ assert callable(func)
+ assert name in litellm.__dict__
+
+ _verify_only_requested_name_imported(name, LLM_PROVIDER_LOGIC_NAMES)
+
diff --git a/tests/test_litellm/types/llms/test_types_llms_openai.py b/tests/test_litellm/types/llms/test_types_llms_openai.py
index 05dec06d469..87cc9586665 100644
--- a/tests/test_litellm/types/llms/test_types_llms_openai.py
+++ b/tests/test_litellm/types/llms/test_types_llms_openai.py
@@ -35,3 +35,137 @@ def test_output_item_added_event():
assert event.sequence_number == 4
assert event.output_index == 1
assert event.item is None
+
+
+class TestResponsesAPIResponseOutputText:
+ """Tests for the output_text property on ResponsesAPIResponse"""
+
+ def test_output_text_with_single_message(self):
+ """Test output_text with a single message containing text output"""
+ from litellm.types.llms.openai import ResponsesAPIResponse
+
+ response = ResponsesAPIResponse(
+ id="resp_123",
+ created_at=1234567890,
+ output=[
+ {
+ "type": "message",
+ "id": "msg_123",
+ "status": "completed",
+ "role": "assistant",
+ "content": [
+ {
+ "type": "output_text",
+ "text": "Hello, world!",
+ }
+ ],
+ }
+ ],
+ )
+
+ assert response.output_text == "Hello, world!"
+
+ def test_output_text_with_multiple_messages(self):
+ """Test output_text with multiple messages aggregates all text"""
+ from litellm.types.llms.openai import ResponsesAPIResponse
+
+ response = ResponsesAPIResponse(
+ id="resp_123",
+ created_at=1234567890,
+ output=[
+ {
+ "type": "message",
+ "id": "msg_1",
+ "status": "completed",
+ "role": "assistant",
+ "content": [
+ {
+ "type": "output_text",
+ "text": "First part. ",
+ }
+ ],
+ },
+ {
+ "type": "message",
+ "id": "msg_2",
+ "status": "completed",
+ "role": "assistant",
+ "content": [
+ {
+ "type": "output_text",
+ "text": "Second part.",
+ }
+ ],
+ },
+ ],
+ )
+
+ assert response.output_text == "First part. Second part."
+
+ def test_output_text_with_no_text_content(self):
+ """Test output_text returns empty string when no output_text content exists"""
+ from litellm.types.llms.openai import ResponsesAPIResponse
+
+ response = ResponsesAPIResponse(
+ id="resp_123",
+ created_at=1234567890,
+ output=[
+ {
+ "type": "function_call",
+ "id": "call_123",
+ "status": "completed",
+ "name": "get_weather",
+ "arguments": "{}",
+ }
+ ],
+ )
+
+ assert response.output_text == ""
+
+ def test_output_text_with_mixed_content(self):
+ """Test output_text only aggregates output_text type content"""
+ from litellm.types.llms.openai import ResponsesAPIResponse
+
+ response = ResponsesAPIResponse(
+ id="resp_123",
+ created_at=1234567890,
+ output=[
+ {
+ "type": "message",
+ "id": "msg_1",
+ "status": "completed",
+ "role": "assistant",
+ "content": [
+ {
+ "type": "output_text",
+ "text": "The weather is sunny. ",
+ },
+ {
+ "type": "refusal",
+ "refusal": "I cannot do that.",
+ },
+ ],
+ },
+ {
+ "type": "function_call",
+ "id": "call_123",
+ "status": "completed",
+ "name": "get_weather",
+ "arguments": "{}",
+ },
+ ],
+ )
+
+ assert response.output_text == "The weather is sunny. "
+
+ def test_output_text_with_empty_output(self):
+ """Test output_text returns empty string with empty output list"""
+ from litellm.types.llms.openai import ResponsesAPIResponse
+
+ response = ResponsesAPIResponse(
+ id="resp_123",
+ created_at=1234567890,
+ output=[],
+ )
+
+ assert response.output_text == ""
diff --git a/ui/litellm-dashboard/e2e_tests/constants.ts b/ui/litellm-dashboard/e2e_tests/constants.ts
new file mode 100644
index 00000000000..b07bd68fcf1
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/constants.ts
@@ -0,0 +1 @@
+export const ADMIN_STORAGE_PATH = "admin.storageState.json";
diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/roles.ts b/ui/litellm-dashboard/e2e_tests/fixtures/roles.ts
new file mode 100644
index 00000000000..913230ad44b
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/fixtures/roles.ts
@@ -0,0 +1,6 @@
+export enum Role {
+ ProxyAdmin = "proxy_admin",
+ ProxyAdminViewer = "proxy_admin_viewer",
+ InternalUser = "internal_user",
+ InternalUserViewer = "internal_user_viewer",
+}
diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/users.ts b/ui/litellm-dashboard/e2e_tests/fixtures/users.ts
new file mode 100644
index 00000000000..d1f1eab00e5
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/fixtures/users.ts
@@ -0,0 +1,10 @@
+import { Role } from "./roles";
+
+const isCI = !!process.env.CI;
+
+export const users = {
+ [Role.ProxyAdmin]: {
+ email: "admin",
+ password: isCI ? "gm" : "sk-1234",
+ },
+};
diff --git a/ui/litellm-dashboard/e2e_tests/globalSetup.ts b/ui/litellm-dashboard/e2e_tests/globalSetup.ts
new file mode 100644
index 00000000000..a725c58f35b
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/globalSetup.ts
@@ -0,0 +1,18 @@
+import { chromium } from "@playwright/test";
+import { users } from "./fixtures/users";
+import { Role } from "./fixtures/roles";
+
+async function globalSetup() {
+ const browser = await chromium.launch();
+ const page = await browser.newPage();
+ await page.goto("http://localhost:4000/ui/login");
+ await page.getByPlaceholder("Enter your username").fill(users[Role.ProxyAdmin].email);
+ await page.getByPlaceholder("Enter your password").fill(users[Role.ProxyAdmin].password);
+ const loginButton = page.getByRole("button", { name: "Login" });
+ await loginButton.click();
+ await page.waitForSelector("text=AI Gateway");
+ await page.context().storageState({ path: "admin.storageState.json" });
+ await browser.close();
+}
+
+export default globalSetup;
diff --git a/ui/litellm-dashboard/e2e_tests/playwright.config.ts b/ui/litellm-dashboard/e2e_tests/playwright.config.ts
new file mode 100644
index 00000000000..329bb7f7afc
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/playwright.config.ts
@@ -0,0 +1,48 @@
+import { defineConfig, devices } from "@playwright/test";
+
+/**
+ * See https://playwright.dev/docs/test-configuration.
+ */
+export default defineConfig({
+ testDir: ".",
+ testMatch: ["**/*.spec.ts", "**/*.setup.ts"],
+ testIgnore: ["**/*.test.*"],
+ /* Run tests in files in parallel */
+ fullyParallel: true,
+ /* Fail the build on CI if you accidentally left test.only in the source code. */
+ forbidOnly: !!process.env.CI,
+ /* Retry on CI only */
+ retries: process.env.CI ? 2 : 0,
+ /* Opt out of parallel tests on CI. */
+ workers: process.env.CI ? 1 : undefined,
+ /* Reporter to use. See https://playwright.dev/docs/test-reporters */
+ reporter: "html",
+ /* Shared settings for all the projects below. See https://playwright.dev/docs/api/class-testoptions. */
+ use: {
+ /* Base URL to use in actions like `await page.goto('/')`. */
+ baseURL: "http://localhost:4000",
+
+ /* Collect trace when retrying the failed test. See https://playwright.dev/docs/trace-viewer */
+ trace: "on-first-retry",
+ },
+
+ /* Configure projects for major browsers */
+ projects: [
+ {
+ name: "chromium",
+ use: { ...devices["Desktop Chrome"] },
+ },
+
+ {
+ name: "firefox",
+ use: { ...devices["Desktop Firefox"] },
+ },
+ ],
+
+ /* Timeout settings */
+ timeout: 4 * 60 * 1000,
+ expect: {
+ timeout: 10 * 1000,
+ },
+ globalSetup: require.resolve("./globalSetup"),
+});
diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts
new file mode 100644
index 00000000000..d8cc26f8642
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts
@@ -0,0 +1,11 @@
+import { test, expect } from "@playwright/test";
+
+test.describe("Authentication Checks", () => {
+ test("should redirect unauthenticated user from a protected page", async ({ page }) => {
+ const protectedPageUrl = "http://localhost:4000/ui?page=llm-playground";
+ const expectedRedirectUrl = "http://localhost:4000/ui/login/";
+ await page.goto(protectedPageUrl, { waitUntil: "domcontentloaded" });
+ await expect(page).toHaveURL(expectedRedirectUrl);
+ await expect(page.getByRole("heading", { name: "Login" })).toBeVisible();
+ });
+});
diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts
new file mode 100644
index 00000000000..5ac977ff0c8
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts
@@ -0,0 +1,13 @@
+import { expect, test } from "@playwright/test";
+import { users } from "../../fixtures/users";
+import { Role } from "../../fixtures/roles";
+
+test("user can log in", async ({ page }) => {
+ await page.goto("http://localhost:4000/ui/login");
+ await page.getByPlaceholder("Enter your username").fill(users[Role.ProxyAdmin].email);
+ await page.getByPlaceholder("Enter your password").fill(users[Role.ProxyAdmin].password);
+ const loginButton = page.getByRole("button", { name: "Login" });
+ await expect(loginButton).toBeEnabled();
+ await loginButton.click();
+ await expect(page.getByText("AI Gateway")).toBeVisible();
+});
diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts
new file mode 100644
index 00000000000..5fa11a98ef6
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts
@@ -0,0 +1,23 @@
+import { test, expect } from "@playwright/test";
+import { ADMIN_STORAGE_PATH } from "../../constants";
+
+test.describe("Add Model", () => {
+ test.use({ storageState: ADMIN_STORAGE_PATH });
+
+ test("Able to see all models for a specific provider in the model dropdown", async ({ page }) => {
+ await page.goto("http://localhost:4000/ui");
+
+ await page.getByText("Models + Endpoints").click();
+ await page.getByRole("tab", { name: "Add Model" }).click();
+
+ const providerInputDropdown = page.getByRole("combobox", { name: /Provider/i });
+ await providerInputDropdown.fill("Anthropic");
+ await page.waitForTimeout(1000);
+ await providerInputDropdown.press("Enter");
+ await page.waitForTimeout(1000);
+
+ const providerModelsDropdown = page.locator(".ant-select-selection-overflow").first();
+ await providerModelsDropdown.click();
+ await expect(page.getByTitle("claude-haiku-4-5", { exact: true })).toBeVisible();
+ });
+});
diff --git a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts
new file mode 100644
index 00000000000..6801f891e87
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts
@@ -0,0 +1,35 @@
+import test, { expect } from "@playwright/test";
+import { Role } from "../../fixtures/roles";
+import { ADMIN_STORAGE_PATH } from "../../constants";
+
+const sidebarButtons = {
+ [Role.ProxyAdmin]: [
+ "Virtual Keys",
+ "Playground",
+ "Models",
+ "Usage",
+ "Teams",
+ "Internal User",
+ "Settings",
+ "Experimental",
+ "API Reference",
+ "AI Hub",
+ ],
+};
+
+const roles = [{ role: Role.ProxyAdmin, storage: ADMIN_STORAGE_PATH }];
+
+for (const { role, storage } of roles) {
+ test.describe(`${role} sidebar`, () => {
+ test.use({ storageState: storage });
+
+ test("can see and navigate all sidebar buttons", async ({ page }) => {
+ await page.goto("http://localhost:4000/ui");
+ for (const button of sidebarButtons[role as keyof typeof sidebarButtons]) {
+ const tab = page.getByRole("menuitem", { name: button });
+ await expect(tab).toBeVisible();
+ await tab.click();
+ }
+ });
+ });
+}
diff --git a/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts
new file mode 100644
index 00000000000..5873bb3125c
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts
@@ -0,0 +1,90 @@
+import { test, expect, Page } from "@playwright/test";
+test.describe("Internal Users Search", () => {
+ test.use({ storageState: "admin.storageState.json" });
+
+ async function goToInternalUsers(page: Page) {
+ await page.goto("http://localhost:4000/ui");
+
+ const tab = page.getByRole("menuitem", { name: "Internal User" });
+ await expect(tab).toBeVisible();
+ await tab.click();
+
+ await expect(page.locator("tbody tr").first()).toBeVisible();
+ await expect(page.locator(".ant-skeleton")).toHaveCount(0);
+ }
+
+ test("can search users by email", async ({ page }) => {
+ await goToInternalUsers(page);
+
+ const rows = page.locator("tbody tr");
+ const searchInput = page.getByPlaceholder("Search by email...");
+
+ await expect(searchInput).toBeVisible();
+
+ // Ensure initial data is loaded
+ const initialCount = await rows.count();
+ expect(initialCount).toBeGreaterThan(0);
+
+ // 🔹 Apply filter + wait for backend response
+ await Promise.all([
+ page.waitForResponse(
+ (res) =>
+ res.url().includes("/user/list") &&
+ res.url().includes("user_email=test%40") && // encoded "test@"
+ res.status() === 200,
+ ),
+ searchInput.fill("test@"),
+ ]);
+ await page.waitForTimeout(5000);
+ const filteredCount = await rows.count();
+ await expect(filteredCount).toBeLessThan(initialCount);
+
+ // 🔹 Clear filter + wait for unfiltered request
+ await Promise.all([
+ page.waitForResponse(
+ (res) => res.url().includes("/user/list") && !res.url().includes("user_email=") && res.status() === 200,
+ ),
+ searchInput.clear(),
+ ]);
+
+ const resetCount = await rows.count();
+ await expect(resetCount).toBe(initialCount);
+ });
+
+ test("can filter users by user ID and SSO ID", async ({ page }) => {
+ await goToInternalUsers(page);
+ const rows = page.locator("tbody tr");
+
+ // Ensure initial data is loaded
+ const initialCount = await rows.count();
+ expect(initialCount).toBeGreaterThan(0);
+
+ const filtersButton = page.getByRole("button", {
+ name: "Filters",
+ exact: true,
+ });
+ await filtersButton.click();
+
+ const userIdInput = page.getByPlaceholder("Filter by User ID");
+ const ssoIdInput = page.getByPlaceholder("Filter by SSO ID");
+ await Promise.all([
+ page.waitForResponse(
+ (res) => res.url().includes("/user/list") && res.url().includes("user_ids=user") && res.status() === 200,
+ ),
+ userIdInput.fill("user"),
+ ]);
+
+ await Promise.all([
+ page.waitForResponse(
+ (res) =>
+ res.url().includes("/user/list") &&
+ res.url().includes("user_ids=user") &&
+ res.url().includes("sso_user_ids=sso") &&
+ res.status() === 200,
+ ),
+ ssoIdInput.fill("sso"),
+ ]);
+ const combinedFilteredCount = await rows.count();
+ await expect(combinedFilteredCount).toBeLessThan(initialCount);
+ });
+});
diff --git a/ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts
new file mode 100644
index 00000000000..980c7233e42
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts
@@ -0,0 +1,52 @@
+import { test, expect, Page } from "@playwright/test";
+
+test.describe("Internal Users Page", () => {
+ test.use({ storageState: "admin.storageState.json" });
+
+ async function goToInternalUsers(page: Page) {
+ await page.goto("http://localhost:4000/ui");
+
+ const internalUserTab = page.getByRole("menuitem", { name: "Internal User" });
+ await expect(internalUserTab).toBeVisible();
+ await internalUserTab.click();
+
+ const firstRow = page.locator("tbody tr").first();
+ await expect(firstRow).toBeVisible();
+ await expect(page.locator(".ant-skeleton")).toHaveCount(0);
+ }
+
+ test("renders internal users table correctly", async ({ page }) => {
+ await goToInternalUsers(page);
+
+ const rows = page.locator("tbody tr");
+ const rowCount = await rows.count();
+ expect(rowCount).toBeGreaterThan(0);
+
+ const userIdHeader = page.getByRole("columnheader", { name: "User ID" });
+ await expect(userIdHeader).toBeVisible();
+
+ const virtualKeysHeader = page.getByRole("columnheader", { name: "Virtual Keys" });
+ await expect(virtualKeysHeader).toBeVisible();
+ });
+
+ test("pagination controls work correctly", async ({ page }) => {
+ await goToInternalUsers(page);
+
+ const paginationInfo = page.locator(".text-sm.text-gray-700");
+ const prevButton = page.getByRole("button", { name: "Previous" });
+ const nextButton = page.getByRole("button", { name: "Next" });
+
+ const infoText = (await paginationInfo.textContent()) || "";
+
+ // On first page, Previous should be disabled
+ if (infoText.includes("1 -")) {
+ await expect(prevButton).toBeDisabled();
+ }
+
+ // Check if there are more pages
+ const hasMorePages = infoText.includes("of") && !infoText.endsWith("25 of 25");
+ if (hasMorePages) {
+ await expect(nextButton).toBeEnabled();
+ }
+ });
+});
diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json
index b65c35e78d8..0e9675f40a1 100644
--- a/ui/litellm-dashboard/package-lock.json
+++ b/ui/litellm-dashboard/package-lock.json
@@ -39,6 +39,7 @@
"uuid": "^11.1.0"
},
"devDependencies": {
+ "@playwright/test": "^1.57.0",
"@tailwindcss/forms": "^0.5.7",
"@testing-library/dom": "^10.4.1",
"@testing-library/jest-dom": "^6.8.0",
@@ -91,7 +92,6 @@
"version": "5.2.0",
"resolved": "https://registry.npmjs.org/@alloc/quick-lru/-/quick-lru-5.2.0.tgz",
"integrity": "sha512-UrcABB+4bUrFABwbluTIBErXwvbsU/V7TZWfmbgJfbkwiBuziS9gxdODUyuiecfdGQ85jglMW6juS3+z5TsKLw==",
- "dev": true,
"license": "MIT",
"engines": {
"node": ">=10"
@@ -325,6 +325,7 @@
"resolved": "https://registry.npmjs.org/@babel/core/-/core-7.28.5.tgz",
"integrity": "sha512-e7jT4DxYvIDLk1ZHmU/m/mB19rex9sv0c2ftBtjSBv+kVM/902eh0fINUzD7UwLLNR+jU585GxUJ8/EBfAM5fw==",
"license": "MIT",
+ "peer": true,
"dependencies": {
"@babel/code-frame": "^7.27.1",
"@babel/generator": "^7.28.5",
@@ -2186,6 +2187,7 @@
}
],
"license": "MIT",
+ "peer": true,
"engines": {
"node": ">=18"
},
@@ -2228,6 +2230,7 @@
}
],
"license": "MIT",
+ "peer": true,
"engines": {
"node": ">=18"
}
@@ -2337,6 +2340,7 @@
"resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-7.1.0.tgz",
"integrity": "sha512-8sLjZwK0R+JlxlYcTuVnyT2v+htpdrjDOKuMcOVdYjt52Lh8hWRYpxBPoKx/Zg+bcjc3wx6fmQevMmUztS/ccA==",
"license": "MIT",
+ "peer": true,
"dependencies": {
"cssesc": "^3.0.0",
"util-deprecate": "^1.0.2"
@@ -2758,6 +2762,7 @@
"resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-7.1.0.tgz",
"integrity": "sha512-8sLjZwK0R+JlxlYcTuVnyT2v+htpdrjDOKuMcOVdYjt52Lh8hWRYpxBPoKx/Zg+bcjc3wx6fmQevMmUztS/ccA==",
"license": "MIT",
+ "peer": true,
"dependencies": {
"cssesc": "^3.0.0",
"util-deprecate": "^1.0.2"
@@ -3541,6 +3546,40 @@
"react-dom": "*"
}
},
+ "node_modules/@docusaurus/plugin-content-docs": {
+ "version": "3.9.2",
+ "resolved": "https://registry.npmjs.org/@docusaurus/plugin-content-docs/-/plugin-content-docs-3.9.2.tgz",
+ "integrity": "sha512-C5wZsGuKTY8jEYsqdxhhFOe1ZDjH0uIYJ9T/jebHwkyxqnr4wW0jTkB72OMqNjsoQRcb0JN3PcSeTwFlVgzCZg==",
+ "license": "MIT",
+ "peer": true,
+ "dependencies": {
+ "@docusaurus/core": "3.9.2",
+ "@docusaurus/logger": "3.9.2",
+ "@docusaurus/mdx-loader": "3.9.2",
+ "@docusaurus/module-type-aliases": "3.9.2",
+ "@docusaurus/theme-common": "3.9.2",
+ "@docusaurus/types": "3.9.2",
+ "@docusaurus/utils": "3.9.2",
+ "@docusaurus/utils-common": "3.9.2",
+ "@docusaurus/utils-validation": "3.9.2",
+ "@types/react-router-config": "^5.0.7",
+ "combine-promises": "^1.1.0",
+ "fs-extra": "^11.1.1",
+ "js-yaml": "^4.1.0",
+ "lodash": "^4.17.21",
+ "schema-dts": "^1.1.2",
+ "tslib": "^2.6.0",
+ "utility-types": "^3.10.0",
+ "webpack": "^5.88.1"
+ },
+ "engines": {
+ "node": ">=20.0"
+ },
+ "peerDependencies": {
+ "react": "^18.0.0 || ^19.0.0",
+ "react-dom": "^18.0.0 || ^19.0.0"
+ }
+ },
"node_modules/@docusaurus/theme-common": {
"version": "3.9.2",
"resolved": "https://registry.npmjs.org/@docusaurus/theme-common/-/theme-common-3.9.2.tgz",
@@ -4700,6 +4739,24 @@
"url": "https://opencollective.com/unified"
}
},
+ "node_modules/@mdx-js/react": {
+ "version": "3.1.1",
+ "resolved": "https://registry.npmjs.org/@mdx-js/react/-/react-3.1.1.tgz",
+ "integrity": "sha512-f++rKLQgUVYDAtECQ6fn/is15GkEH9+nZPM3MS0RcxVqoTfawHvDlSCH7JbMhAM6uJ32v3eXLvLmLvjGu7PTQw==",
+ "license": "MIT",
+ "peer": true,
+ "dependencies": {
+ "@types/mdx": "^2.0.0"
+ },
+ "funding": {
+ "type": "opencollective",
+ "url": "https://opencollective.com/unified"
+ },
+ "peerDependencies": {
+ "@types/react": ">=16",
+ "react": ">=16"
+ }
+ },
"node_modules/@mermaid-js/parser": {
"version": "0.6.3",
"resolved": "https://registry.npmjs.org/@mermaid-js/parser/-/parser-0.6.3.tgz",
@@ -4927,6 +4984,23 @@
"node": ">=12.4.0"
}
},
+ "node_modules/@playwright/test": {
+ "version": "1.57.0",
+ "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.57.0.tgz",
+ "integrity": "sha512-6TyEnHgd6SArQO8UO2OMTxshln3QMWBtPGrOCgs3wVEmQmwyuNtB10IZMfmYDE0riwNR1cu4q+pPcxMVtaG3TA==",
+ "devOptional": true,
+ "license": "Apache-2.0",
+ "peer": true,
+ "dependencies": {
+ "playwright": "1.57.0"
+ },
+ "bin": {
+ "playwright": "cli.js"
+ },
+ "engines": {
+ "node": ">=18"
+ }
+ },
"node_modules/@pnpm/config.env-replace": {
"version": "1.1.0",
"resolved": "https://registry.npmjs.org/@pnpm/config.env-replace/-/config.env-replace-1.1.0.tgz",
@@ -5773,6 +5847,7 @@
"integrity": "sha512-o4PXJQidqJl82ckFaXUeoAW+XysPLauYI43Abki5hABd853iMhitooc6znOnczgbTYmEP6U6/y1ZyKAIsvMKGg==",
"dev": true,
"license": "MIT",
+ "peer": true,
"dependencies": {
"@babel/code-frame": "^7.10.4",
"@babel/runtime": "^7.12.5",
@@ -6566,6 +6641,7 @@
"resolved": "https://registry.npmjs.org/@types/react/-/react-18.2.48.tgz",
"integrity": "sha512-qboRCl6Ie70DQQG9hhNREz81jqC1cs9EVNcjQ1AU+jH6NFfSAhVVbrrY/+nSF+Bsk4AOwm9Qa61InvMCyV+H3w==",
"license": "MIT",
+ "peer": true,
"dependencies": {
"@types/prop-types": "*",
"@types/scheduler": "*",
@@ -6588,6 +6664,7 @@
"integrity": "sha512-MEe3UeoENYVFXzoXEWsvcpg6ZvlrFNlOQ7EOsvhI3CfAXwzPfO8Qwuxd40nepsYKqyyVQnTdEfv68q91yLcKrQ==",
"dev": true,
"license": "MIT",
+ "peer": true,
"peerDependencies": {
"@types/react": "^18.0.0"
}
@@ -6784,6 +6861,7 @@
"integrity": "sha512-lJi3PfxVmo0AkEY93ecfN+r8SofEqZNGByvHAI3GBLrvt1Cw6H5k1IM02nSzu0RfUafr2EvFSw0wAsZgubNplQ==",
"dev": true,
"license": "MIT",
+ "peer": true,
"dependencies": {
"@typescript-eslint/scope-manager": "8.47.0",
"@typescript-eslint/types": "8.47.0",
@@ -7445,6 +7523,7 @@
"integrity": "sha512-hGISOaP18plkzbWEcP/QvtRW1xDXF2+96HbEX6byqQhAUbiS5oH6/9JwW+QsQCIYON2bI6QZBF+2PvOmrRZ9wA==",
"dev": true,
"license": "MIT",
+ "peer": true,
"dependencies": {
"@vitest/utils": "3.2.4",
"fflate": "^0.8.2",
@@ -7673,6 +7752,7 @@
"resolved": "https://registry.npmjs.org/acorn/-/acorn-8.15.0.tgz",
"integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==",
"license": "MIT",
+ "peer": true,
"bin": {
"acorn": "bin/acorn"
},
@@ -7762,6 +7842,7 @@
"resolved": "https://registry.npmjs.org/ajv/-/ajv-6.12.6.tgz",
"integrity": "sha512-j3fVLgvTo527anyYyJOGTYJbG+vnnQYvE0m5mmkc1TK+nxAppkCLMIL0aZ4dblVCNoGShhm+kzE4ZUykBoMg4g==",
"license": "MIT",
+ "peer": true,
"dependencies": {
"fast-deep-equal": "^3.1.1",
"fast-json-stable-stringify": "^2.0.0",
@@ -7982,7 +8063,6 @@
"version": "1.3.0",
"resolved": "https://registry.npmjs.org/any-promise/-/any-promise-1.3.0.tgz",
"integrity": "sha512-7UvmKalWRt1wgjL1RrGxoSJW/0QZFIegpeGvZG9kjp8vrRu55XTHbwnqq2GpXm9uLbcuhxm3IqX9OB4MZR1b2A==",
- "dev": true,
"license": "MIT"
},
"node_modules/anymatch": {
@@ -8002,7 +8082,6 @@
"version": "5.0.2",
"resolved": "https://registry.npmjs.org/arg/-/arg-5.0.2.tgz",
"integrity": "sha512-PYjyFOLKQ9y57JvQ6QLo8dAgNqswh8M1RMJYdQduT6xbWSgK36P/Z/v+p888pM69jMMfS8Xd8F6I1kQ/I9HUGg==",
- "dev": true,
"license": "MIT"
},
"node_modules/argparse": {
@@ -8479,23 +8558,23 @@
}
},
"node_modules/body-parser": {
- "version": "1.20.3",
- "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-1.20.3.tgz",
- "integrity": "sha512-7rAxByjUMqQ3/bHJy7D6OGXvx/MMc4IqBn/X0fcM1QUcAItpZrBEYhWGem+tzXH90c+G01ypMcYJBO9Y30203g==",
+ "version": "1.20.4",
+ "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-1.20.4.tgz",
+ "integrity": "sha512-ZTgYYLMOXY9qKU/57FAo8F+HA2dGX7bqGc71txDRC1rS4frdFI5R7NhluHxH6M0YItAP0sHB4uqAOcYKxO6uGA==",
"license": "MIT",
"dependencies": {
- "bytes": "3.1.2",
+ "bytes": "~3.1.2",
"content-type": "~1.0.5",
"debug": "2.6.9",
"depd": "2.0.0",
- "destroy": "1.2.0",
- "http-errors": "2.0.0",
- "iconv-lite": "0.4.24",
- "on-finished": "2.4.1",
- "qs": "6.13.0",
- "raw-body": "2.5.2",
+ "destroy": "~1.2.0",
+ "http-errors": "~2.0.1",
+ "iconv-lite": "~0.4.24",
+ "on-finished": "~2.4.1",
+ "qs": "~6.14.0",
+ "raw-body": "~2.5.3",
"type-is": "~1.6.18",
- "unpipe": "1.0.0"
+ "unpipe": "~1.0.0"
},
"engines": {
"node": ">= 0.8",
@@ -8520,6 +8599,26 @@
"ms": "2.0.0"
}
},
+ "node_modules/body-parser/node_modules/http-errors": {
+ "version": "2.0.1",
+ "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz",
+ "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==",
+ "license": "MIT",
+ "dependencies": {
+ "depd": "~2.0.0",
+ "inherits": "~2.0.4",
+ "setprototypeof": "~1.2.0",
+ "statuses": "~2.0.2",
+ "toidentifier": "~1.0.1"
+ },
+ "engines": {
+ "node": ">= 0.8"
+ },
+ "funding": {
+ "type": "opencollective",
+ "url": "https://opencollective.com/express"
+ }
+ },
"node_modules/body-parser/node_modules/iconv-lite": {
"version": "0.4.24",
"resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.4.24.tgz",
@@ -8538,6 +8637,15 @@
"integrity": "sha512-Tpp60P6IUJDTuOq/5Z8cdskzJujfwqfOTkrwIwj7IRISpnkJnT6SyJ4PCPnGMoFjC9ddhal5KVIYtAt97ix05A==",
"license": "MIT"
},
+ "node_modules/body-parser/node_modules/statuses": {
+ "version": "2.0.2",
+ "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz",
+ "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==",
+ "license": "MIT",
+ "engines": {
+ "node": ">= 0.8"
+ }
+ },
"node_modules/bonjour-service": {
"version": "1.3.0",
"resolved": "https://registry.npmjs.org/bonjour-service/-/bonjour-service-1.3.0.tgz",
@@ -8617,6 +8725,7 @@
}
],
"license": "MIT",
+ "peer": true,
"dependencies": {
"baseline-browser-mapping": "^2.8.25",
"caniuse-lite": "^1.0.30001754",
@@ -8797,7 +8906,6 @@
"version": "2.0.1",
"resolved": "https://registry.npmjs.org/camelcase-css/-/camelcase-css-2.0.1.tgz",
"integrity": "sha512-QOSvevhslijgYwRx6Rv7zKdMF8lbRmx+uQGx2+vDc+KI/eBnsy9kit5aj23AgGu3pa4t9AgwbnXWqS+iOY+2aA==",
- "dev": true,
"license": "MIT",
"engines": {
"node": ">= 6"
@@ -8942,6 +9050,7 @@
"resolved": "https://registry.npmjs.org/chevrotain/-/chevrotain-11.0.3.tgz",
"integrity": "sha512-ci2iJH6LeIkvP9eJW6gpueU8cnZhv85ELY8w8WiFtNjMHA5ad6pQLaJo9mEly/9qUyCpvqX8/POVUTf18/HFdw==",
"license": "Apache-2.0",
+ "peer": true,
"dependencies": {
"@chevrotain/cst-dts-gen": "11.0.3",
"@chevrotain/gast": "11.0.3",
@@ -9670,6 +9779,7 @@
"resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-7.1.0.tgz",
"integrity": "sha512-8sLjZwK0R+JlxlYcTuVnyT2v+htpdrjDOKuMcOVdYjt52Lh8hWRYpxBPoKx/Zg+bcjc3wx6fmQevMmUztS/ccA==",
"license": "MIT",
+ "peer": true,
"dependencies": {
"cssesc": "^3.0.0",
"util-deprecate": "^1.0.2"
@@ -10032,6 +10142,7 @@
"resolved": "https://registry.npmjs.org/cytoscape/-/cytoscape-3.33.1.tgz",
"integrity": "sha512-iJc4TwyANnOGR1OmWhsS9ayRS3s+XQ185FmuHObThD+5AeJCakAAbWv8KimMTt08xCCLNgneQwFp+JRJOr9qGQ==",
"license": "MIT",
+ "peer": true,
"engines": {
"node": ">=0.10"
}
@@ -10441,6 +10552,7 @@
"resolved": "https://registry.npmjs.org/d3-selection/-/d3-selection-3.0.0.tgz",
"integrity": "sha512-fmTRWbNMmsmWq6xJV8D19U/gw/bwrHfNXxrIN+HfZgnzqTHp9jOmKMhsTUjXOJnZOdZY9Q28y4yebKzqDKlxlQ==",
"license": "ISC",
+ "peer": true,
"engines": {
"node": ">=12"
}
@@ -10615,6 +10727,7 @@
"resolved": "https://registry.npmjs.org/date-fns/-/date-fns-3.6.0.tgz",
"integrity": "sha512-fRHTG8g/Gif+kSh50gaGEdToemgfj74aRX3swtiouboip5JDLAyDE9F11nHMIcvOaXeOC6D7SpNhi7uFyB7Uww==",
"license": "MIT",
+ "peer": true,
"funding": {
"type": "github",
"url": "https://github.com/sponsors/kossnocorp"
@@ -10894,7 +11007,6 @@
"version": "1.2.2",
"resolved": "https://registry.npmjs.org/didyoumean/-/didyoumean-1.2.2.tgz",
"integrity": "sha512-gxtyfqMg7GKyhQmb056K7M3xszy/myH8w+B4RT+QXBQsvAOdc3XymqDDPHx1BgPgsdAA5SIifona89YtRATDzw==",
- "dev": true,
"license": "Apache-2.0"
},
"node_modules/dir-glob": {
@@ -10913,7 +11025,6 @@
"version": "1.1.3",
"resolved": "https://registry.npmjs.org/dlv/-/dlv-1.1.3.tgz",
"integrity": "sha512-+HlytyjlPKnIG8XuRG8WvmBP8xs8P71y+SKKS6ZXWoEgLuePxtDoUEiH7WkdePWrQ5JBpE6aoVqfZfJUQkjXwA==",
- "dev": true,
"license": "MIT"
},
"node_modules/dns-packet": {
@@ -11494,6 +11605,7 @@
"deprecated": "This version is no longer supported. Please see https://eslint.org/version-support for other options.",
"dev": true,
"license": "MIT",
+ "peer": true,
"dependencies": {
"@eslint-community/eslint-utils": "^4.2.0",
"@eslint-community/regexpp": "^4.6.1",
@@ -11679,6 +11791,7 @@
"integrity": "sha512-whOE1HFo/qJDyX4SnXzP4N6zOWn79WhnCUY/iDR0mPfQZO8wcYE4JClzI2oZrhBnnMUCBCHZhO6VQyoBU95mZA==",
"dev": true,
"license": "MIT",
+ "peer": true,
"dependencies": {
"@rtsao/scc": "^1.1.0",
"array-includes": "^3.1.9",
@@ -12181,39 +12294,39 @@
}
},
"node_modules/express": {
- "version": "4.21.2",
- "resolved": "https://registry.npmjs.org/express/-/express-4.21.2.tgz",
- "integrity": "sha512-28HqgMZAmih1Czt9ny7qr6ek2qddF4FclbMzwhCREB6OFfH+rXAnuNCwo1/wFvrtbgsQDb4kSbX9de9lFbrXnA==",
+ "version": "4.22.1",
+ "resolved": "https://registry.npmjs.org/express/-/express-4.22.1.tgz",
+ "integrity": "sha512-F2X8g9P1X7uCPZMA3MVf9wcTqlyNp7IhH5qPCI0izhaOIYXaW9L535tGA3qmjRzpH+bZczqq7hVKxTR4NWnu+g==",
"license": "MIT",
"dependencies": {
"accepts": "~1.3.8",
"array-flatten": "1.1.1",
- "body-parser": "1.20.3",
- "content-disposition": "0.5.4",
+ "body-parser": "~1.20.3",
+ "content-disposition": "~0.5.4",
"content-type": "~1.0.4",
- "cookie": "0.7.1",
- "cookie-signature": "1.0.6",
+ "cookie": "~0.7.1",
+ "cookie-signature": "~1.0.6",
"debug": "2.6.9",
"depd": "2.0.0",
"encodeurl": "~2.0.0",
"escape-html": "~1.0.3",
"etag": "~1.8.1",
- "finalhandler": "1.3.1",
- "fresh": "0.5.2",
- "http-errors": "2.0.0",
+ "finalhandler": "~1.3.1",
+ "fresh": "~0.5.2",
+ "http-errors": "~2.0.0",
"merge-descriptors": "1.0.3",
"methods": "~1.1.2",
- "on-finished": "2.4.1",
+ "on-finished": "~2.4.1",
"parseurl": "~1.3.3",
- "path-to-regexp": "0.1.12",
+ "path-to-regexp": "~0.1.12",
"proxy-addr": "~2.0.7",
- "qs": "6.13.0",
+ "qs": "~6.14.0",
"range-parser": "~1.2.1",
"safe-buffer": "5.2.1",
- "send": "0.19.0",
- "serve-static": "1.16.2",
+ "send": "~0.19.0",
+ "serve-static": "~1.16.2",
"setprototypeof": "1.2.0",
- "statuses": "2.0.1",
+ "statuses": "~2.0.1",
"type-is": "~1.6.18",
"utils-merge": "1.0.1",
"vary": "~1.1.2"
@@ -14942,6 +15055,7 @@
"integrity": "sha512-454TI39PeRDW1LgpyLPyURtB4Zx1tklSr6+OFOipsxGUH1WMTvk6C65JQdrj455+DP2uJ1+veBEHTGFKWVLFoA==",
"dev": true,
"license": "MIT",
+ "peer": true,
"dependencies": {
"@acemir/cssom": "^0.9.23",
"@asamuzakjp/dom-selector": "^6.7.4",
@@ -18069,6 +18183,7 @@
"resolved": "https://registry.npmjs.org/moment/-/moment-2.30.1.tgz",
"integrity": "sha512-uEmtNhbDOrWPFS+hdjFCBfy9f2YoyzRpwcl+DqpC6taX21FzsTLQVbMV/W7PzNSX6x/bhC1zA3c2UQ5NzH6how==",
"license": "MIT",
+ "peer": true,
"engines": {
"node": "*"
}
@@ -18105,7 +18220,6 @@
"version": "2.7.0",
"resolved": "https://registry.npmjs.org/mz/-/mz-2.7.0.tgz",
"integrity": "sha512-z81GNO7nnYMEhrGh9LeymoE4+Yr0Wn5McHIZMK5cfQCl+NDX08sCZgUc9/6MHni9IWuFLm1Z3HTCXu2z9fN62Q==",
- "dev": true,
"license": "MIT",
"dependencies": {
"any-promise": "^1.0.0",
@@ -18454,7 +18568,6 @@
"version": "3.0.0",
"resolved": "https://registry.npmjs.org/object-hash/-/object-hash-3.0.0.tgz",
"integrity": "sha512-RSn9F68PjH9HqtltsSnqYC1XXoWe9Bju5+213R98cNGttag9q9yAOTzdbsqvIa7aNm5WffBZFpWYr2aWrklWAw==",
- "dev": true,
"license": "MIT",
"engines": {
"node": ">= 6"
@@ -19095,7 +19208,6 @@
"version": "2.3.0",
"resolved": "https://registry.npmjs.org/pify/-/pify-2.3.0.tgz",
"integrity": "sha512-udgsAY+fTnvv7kI7aaxbqwWNb0AHiB0qBO89PZKPkoTmGOgdbrHDKD+0B2X4uTfJ/FT1R09r9gTsjUjNJotuog==",
- "dev": true,
"license": "MIT",
"engines": {
"node": ">=0.10.0"
@@ -19105,7 +19217,6 @@
"version": "4.0.7",
"resolved": "https://registry.npmjs.org/pirates/-/pirates-4.0.7.tgz",
"integrity": "sha512-TfySrs/5nm8fQJDcBDuUng3VOUKsd7S+zqvbOTiGXHfxX4wK31ard+hoNuvkicM/2YFzlpDgABOevKSsB4G/FA==",
- "dev": true,
"license": "MIT",
"engines": {
"node": ">= 6"
@@ -19219,6 +19330,53 @@
"pathe": "^2.0.3"
}
},
+ "node_modules/playwright": {
+ "version": "1.57.0",
+ "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.57.0.tgz",
+ "integrity": "sha512-ilYQj1s8sr2ppEJ2YVadYBN0Mb3mdo9J0wQ+UuDhzYqURwSoW4n1Xs5vs7ORwgDGmyEh33tRMeS8KhdkMoLXQw==",
+ "devOptional": true,
+ "license": "Apache-2.0",
+ "dependencies": {
+ "playwright-core": "1.57.0"
+ },
+ "bin": {
+ "playwright": "cli.js"
+ },
+ "engines": {
+ "node": ">=18"
+ },
+ "optionalDependencies": {
+ "fsevents": "2.3.2"
+ }
+ },
+ "node_modules/playwright-core": {
+ "version": "1.57.0",
+ "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.57.0.tgz",
+ "integrity": "sha512-agTcKlMw/mjBWOnD6kFZttAAGHgi/Nw0CZ2o6JqWSbMlI219lAFLZZCyqByTsvVAJq5XA5H8cA6PrvBRpBWEuQ==",
+ "devOptional": true,
+ "license": "Apache-2.0",
+ "bin": {
+ "playwright-core": "cli.js"
+ },
+ "engines": {
+ "node": ">=18"
+ }
+ },
+ "node_modules/playwright/node_modules/fsevents": {
+ "version": "2.3.2",
+ "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz",
+ "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==",
+ "dev": true,
+ "hasInstallScript": true,
+ "license": "MIT",
+ "optional": true,
+ "os": [
+ "darwin"
+ ],
+ "engines": {
+ "node": "^8.16.0 || ^10.6.0 || >=11.0.0"
+ }
+ },
"node_modules/points-on-curve": {
"version": "0.2.0",
"resolved": "https://registry.npmjs.org/points-on-curve/-/points-on-curve-0.2.0.tgz",
@@ -19264,6 +19422,7 @@
}
],
"license": "MIT",
+ "peer": true,
"dependencies": {
"nanoid": "^3.3.11",
"picocolors": "^1.1.1",
@@ -19820,7 +19979,6 @@
"version": "15.1.0",
"resolved": "https://registry.npmjs.org/postcss-import/-/postcss-import-15.1.0.tgz",
"integrity": "sha512-hpr+J05B2FVYUAXHeK1YyI267J/dDDhMU6B6civm8hSY1jYJnBXxzKDKDswzJmtLHryrjhnDjqqp/49t8FALew==",
- "dev": true,
"license": "MIT",
"dependencies": {
"postcss-value-parser": "^4.0.0",
@@ -19838,7 +19996,6 @@
"version": "4.1.0",
"resolved": "https://registry.npmjs.org/postcss-js/-/postcss-js-4.1.0.tgz",
"integrity": "sha512-oIAOTqgIo7q2EOwbhb8UalYePMvYoIeRY2YKntdpFQXNosSu3vLrniGgmH9OKs/qAkfoj5oB3le/7mINW1LCfw==",
- "dev": true,
"funding": [
{
"type": "opencollective",
@@ -19893,7 +20050,6 @@
"version": "6.0.1",
"resolved": "https://registry.npmjs.org/postcss-load-config/-/postcss-load-config-6.0.1.tgz",
"integrity": "sha512-oPtTM4oerL+UXmx+93ytZVN82RrlY/wPUV8IeDxFrzIjXOLF1pN+EmKPLbubvKHT2HC20xXsCAH2Z+CKV6Oz/g==",
- "dev": true,
"funding": [
{
"type": "opencollective",
@@ -20182,7 +20338,6 @@
"version": "6.2.0",
"resolved": "https://registry.npmjs.org/postcss-nested/-/postcss-nested-6.2.0.tgz",
"integrity": "sha512-HQbt28KulC5AJzG+cZtj9kvKB93CFCdLvog1WFLf1D+xmMvPGlBstkpTEZfK5+AN9hfJocyBFCNiqyS48bpgzQ==",
- "dev": true,
"funding": [
{
"type": "opencollective",
@@ -20280,6 +20435,7 @@
"resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-7.1.0.tgz",
"integrity": "sha512-8sLjZwK0R+JlxlYcTuVnyT2v+htpdrjDOKuMcOVdYjt52Lh8hWRYpxBPoKx/Zg+bcjc3wx6fmQevMmUztS/ccA==",
"license": "MIT",
+ "peer": true,
"dependencies": {
"cssesc": "^3.0.0",
"util-deprecate": "^1.0.2"
@@ -21011,12 +21167,12 @@
}
},
"node_modules/qs": {
- "version": "6.13.0",
- "resolved": "https://registry.npmjs.org/qs/-/qs-6.13.0.tgz",
- "integrity": "sha512-+38qI9SOr8tfZ4QmJNplMUxqjbe7LKvvZgWdExBOmd+egZTtjLB67Gu0HRX3u/XOq7UU2Nx6nsjvS16Z9uwfpg==",
+ "version": "6.14.1",
+ "resolved": "https://registry.npmjs.org/qs/-/qs-6.14.1.tgz",
+ "integrity": "sha512-4EK3+xJl8Ts67nLYNwqw/dsFVnCf+qR7RgXSK9jEEm9unao3njwMDdmsdvoKBKHzxd7tCYz5e5M+SnMjdtXGQQ==",
"license": "BSD-3-Clause",
"dependencies": {
- "side-channel": "^1.0.6"
+ "side-channel": "^1.1.0"
},
"engines": {
"node": ">=0.6"
@@ -21092,15 +21248,15 @@
}
},
"node_modules/raw-body": {
- "version": "2.5.2",
- "resolved": "https://registry.npmjs.org/raw-body/-/raw-body-2.5.2.tgz",
- "integrity": "sha512-8zGqypfENjCIqGhgXToC8aB2r7YrBX+AQAfIPs/Mlk+BtPTztOvTS01NRW/3Eh60J+a48lt8qsCzirQ6loCVfA==",
+ "version": "2.5.3",
+ "resolved": "https://registry.npmjs.org/raw-body/-/raw-body-2.5.3.tgz",
+ "integrity": "sha512-s4VSOf6yN0rvbRZGxs8Om5CWj6seneMwK3oDb4lWDH0UPhWcxwOWw5+qk24bxq87szX1ydrwylIOp2uG1ojUpA==",
"license": "MIT",
"dependencies": {
- "bytes": "3.1.2",
- "http-errors": "2.0.0",
- "iconv-lite": "0.4.24",
- "unpipe": "1.0.0"
+ "bytes": "~3.1.2",
+ "http-errors": "~2.0.1",
+ "iconv-lite": "~0.4.24",
+ "unpipe": "~1.0.0"
},
"engines": {
"node": ">= 0.8"
@@ -21115,6 +21271,26 @@
"node": ">= 0.8"
}
},
+ "node_modules/raw-body/node_modules/http-errors": {
+ "version": "2.0.1",
+ "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz",
+ "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==",
+ "license": "MIT",
+ "dependencies": {
+ "depd": "~2.0.0",
+ "inherits": "~2.0.4",
+ "setprototypeof": "~1.2.0",
+ "statuses": "~2.0.2",
+ "toidentifier": "~1.0.1"
+ },
+ "engines": {
+ "node": ">= 0.8"
+ },
+ "funding": {
+ "type": "opencollective",
+ "url": "https://opencollective.com/express"
+ }
+ },
"node_modules/raw-body/node_modules/iconv-lite": {
"version": "0.4.24",
"resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.4.24.tgz",
@@ -21127,6 +21303,15 @@
"node": ">=0.10.0"
}
},
+ "node_modules/raw-body/node_modules/statuses": {
+ "version": "2.0.2",
+ "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz",
+ "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==",
+ "license": "MIT",
+ "engines": {
+ "node": ">= 0.8"
+ }
+ },
"node_modules/rc": {
"version": "1.2.8",
"resolved": "https://registry.npmjs.org/rc/-/rc-1.2.8.tgz",
@@ -21774,6 +21959,7 @@
"resolved": "https://registry.npmjs.org/react/-/react-18.3.1.tgz",
"integrity": "sha512-wS+hAgJShR0KhEvPJArfuPVN1+Hz1t0Y6n5jLrGQbkb4urgPE/0Rve+1kMB1v/oWgHgm4WIcV+i7F2pTVj+2iQ==",
"license": "MIT",
+ "peer": true,
"dependencies": {
"loose-envify": "^1.1.0"
},
@@ -21813,6 +21999,7 @@
"resolved": "https://registry.npmjs.org/react-dom/-/react-dom-18.3.1.tgz",
"integrity": "sha512-5m4nQKp+rZRb09LNH59GM4BxTh9251/ylbKIbpe7TpGxfJ+9kv6BLkLBXIjjspbgbnIBNqlI23tRnTWT0snUIw==",
"license": "MIT",
+ "peer": true,
"dependencies": {
"loose-envify": "^1.1.0",
"scheduler": "^0.23.2"
@@ -21870,6 +22057,7 @@
"resolved": "https://registry.npmjs.org/@docusaurus/react-loadable/-/react-loadable-6.0.0.tgz",
"integrity": "sha512-YMMxTUQV/QFSnbgrP3tjDzLHRg7vsbMn8e9HAa8o/1iXoiomo48b7sk/kkmWEuWNDPJVlKSJRB6Y2fHqdJk+SQ==",
"license": "MIT",
+ "peer": true,
"dependencies": {
"@types/react": "*"
},
@@ -21935,6 +22123,7 @@
"resolved": "https://registry.npmjs.org/react-router/-/react-router-5.3.4.tgz",
"integrity": "sha512-Ys9K+ppnJah3QuaRiLxk+jDWOR1MekYQrlytiXxC1RyfbdsZkS5pvKAzCCr031xHixZwpnsYNT5xysdFHQaYsA==",
"license": "MIT",
+ "peer": true,
"dependencies": {
"@babel/runtime": "^7.12.13",
"history": "^4.9.0",
@@ -22049,7 +22238,6 @@
"version": "1.0.0",
"resolved": "https://registry.npmjs.org/read-cache/-/read-cache-1.0.0.tgz",
"integrity": "sha512-Owdv/Ft7IjOgm/i0xvNDZ1LrRANRfew4b2prF3OWMQLxLfu3bS8FVhCsrSCMK4lR56Y9ya+AThoTpDCTxCmpRA==",
- "dev": true,
"license": "MIT",
"dependencies": {
"pify": "^2.3.0"
@@ -22969,6 +23157,12 @@
"loose-envify": "^1.1.0"
}
},
+ "node_modules/schema-dts": {
+ "version": "1.1.5",
+ "resolved": "https://registry.npmjs.org/schema-dts/-/schema-dts-1.1.5.tgz",
+ "integrity": "sha512-RJr9EaCmsLzBX2NDiO5Z3ux2BVosNZN5jo0gWgsyKvxKIUL5R3swNvoorulAeL9kLB0iTSX7V6aokhla2m7xbg==",
+ "license": "Apache-2.0"
+ },
"node_modules/schema-utils": {
"version": "4.3.3",
"resolved": "https://registry.npmjs.org/schema-utils/-/schema-utils-4.3.3.tgz",
@@ -22993,6 +23187,7 @@
"resolved": "https://registry.npmjs.org/ajv/-/ajv-8.17.1.tgz",
"integrity": "sha512-B/gBuNg5SiMTrPkC+A2+cW0RszwxYmn6VYxB/inlBStS5nx6xHIt/ehKRhIMhqusl7a8LjQoZnjCs5vhwxOQ1g==",
"license": "MIT",
+ "peer": true,
"dependencies": {
"fast-deep-equal": "^3.1.3",
"fast-uri": "^3.0.1",
@@ -24033,7 +24228,6 @@
"version": "3.35.1",
"resolved": "https://registry.npmjs.org/sucrase/-/sucrase-3.35.1.tgz",
"integrity": "sha512-DhuTmvZWux4H1UOnWMB3sk0sbaCVOoQZjv8u1rDoTV0HTdGem9hkAZtl4JZy8P2z4Bg0nT+YMeOFyVr4zcG5Tw==",
- "dev": true,
"license": "MIT",
"dependencies": {
"@jridgewell/gen-mapping": "^0.3.2",
@@ -24056,7 +24250,6 @@
"version": "4.1.1",
"resolved": "https://registry.npmjs.org/commander/-/commander-4.1.1.tgz",
"integrity": "sha512-NOKm8xhkzAjzFx8B2v5OAHT+u5pRQc2UCa2Vq9jYL/31o2wi9mxBA7LIFs3sV5VSC49z6pEhfbMULvShKj26WA==",
- "dev": true,
"license": "MIT",
"engines": {
"node": ">= 6"
@@ -24225,8 +24418,8 @@
"version": "3.4.18",
"resolved": "https://registry.npmjs.org/tailwindcss/-/tailwindcss-3.4.18.tgz",
"integrity": "sha512-6A2rnmW5xZMdw11LYjhcI5846rt9pbLSabY5XPxo+XWdxwZaFEn47Go4NzFiHu9sNNmr/kXivP1vStfvMaK1GQ==",
- "dev": true,
"license": "MIT",
+ "peer": true,
"dependencies": {
"@alloc/quick-lru": "^5.2.0",
"arg": "^5.0.2",
@@ -24263,7 +24456,6 @@
"version": "6.0.2",
"resolved": "https://registry.npmjs.org/glob-parent/-/glob-parent-6.0.2.tgz",
"integrity": "sha512-XxwI8EOhVQgWp6iDL+3b0r86f4d6AX6zSU55HfB4ydCEuXLXc5FcYeOu+nnGftS4TEju/11rt4KJPTMgbfmv4A==",
- "dev": true,
"license": "ISC",
"dependencies": {
"is-glob": "^4.0.3"
@@ -24424,7 +24616,6 @@
"version": "3.3.1",
"resolved": "https://registry.npmjs.org/thenify/-/thenify-3.3.1.tgz",
"integrity": "sha512-RVZSIV5IG10Hk3enotrhvz0T9em6cyHBLkH/YAZuKqd8hRkKhSfCGIcP2KUY0EPxndzANBmNllzWPwak+bheSw==",
- "dev": true,
"license": "MIT",
"dependencies": {
"any-promise": "^1.0.0"
@@ -24434,7 +24625,6 @@
"version": "1.6.0",
"resolved": "https://registry.npmjs.org/thenify-all/-/thenify-all-1.6.0.tgz",
"integrity": "sha512-RNxQH/qI8/t3thXJDwcstUO4zeqo64+Uy/+sNVRBx4Xn2OX+OZ9oP+iJnNFqplFra2ZUVeKCSa2oVWi3T4uVmA==",
- "dev": true,
"license": "MIT",
"dependencies": {
"thenify": ">= 3.1.0 < 4"
@@ -24506,7 +24696,6 @@
"version": "0.2.15",
"resolved": "https://registry.npmjs.org/tinyglobby/-/tinyglobby-0.2.15.tgz",
"integrity": "sha512-j2Zq4NyQYG5XMST4cbs02Ak8iJUdxRM0XI5QyxXuZOzKOINmWurp3smXu3y5wDcJrptwpSjgXHzIQxR0omXljQ==",
- "dev": true,
"license": "MIT",
"dependencies": {
"fdir": "^6.5.0",
@@ -24523,7 +24712,6 @@
"version": "6.5.0",
"resolved": "https://registry.npmjs.org/fdir/-/fdir-6.5.0.tgz",
"integrity": "sha512-tIbYtZbucOs0BRGqPJkshJUYdL+SDH7dVM8gjy+ERp3WAUjLEFJE+02kanyHtwjWOnwrKYBiwAmM0p4kLJAnXg==",
- "dev": true,
"license": "MIT",
"engines": {
"node": ">=12.0.0"
@@ -24541,8 +24729,8 @@
"version": "4.0.3",
"resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz",
"integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==",
- "dev": true,
"license": "MIT",
+ "peer": true,
"engines": {
"node": ">=12"
},
@@ -24723,7 +24911,6 @@
"version": "0.1.13",
"resolved": "https://registry.npmjs.org/ts-interface-checker/-/ts-interface-checker-0.1.13.tgz",
"integrity": "sha512-Y/arvbn+rrz3JCKl9C4kVNfTfSm2/mEp5FSz5EsZSANGPSlQrpRI5M4PKF+mJnE52jOO90PnPSc3Ur3bTQw0gA==",
- "dev": true,
"license": "Apache-2.0"
},
"node_modules/tsconfig-paths": {
@@ -24756,7 +24943,8 @@
"version": "2.8.1",
"resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz",
"integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==",
- "license": "0BSD"
+ "license": "0BSD",
+ "peer": true
},
"node_modules/type-check": {
"version": "0.4.0",
@@ -24887,8 +25075,9 @@
"version": "5.3.3",
"resolved": "https://registry.npmjs.org/typescript/-/typescript-5.3.3.tgz",
"integrity": "sha512-pXWcraxM0uxAS+tN0AG/BF2TyqmHO014Z070UsJ+pFvYuRSq8KH8DmWpnbXe0pEPDHXZV3FcAbJkijJ5oNEnWw==",
- "dev": true,
+ "devOptional": true,
"license": "Apache-2.0",
+ "peer": true,
"bin": {
"tsc": "bin/tsc",
"tsserver": "bin/tsserver"
@@ -25431,6 +25620,7 @@
"integrity": "sha512-NL8jTlbo0Tn4dUEXEsUg8KeyG/Lkmc4Fnzb8JXN/Ykm9G4HNImjtABMJgkQoVjOBN/j2WAwDTRytdqJbZsah7w==",
"dev": true,
"license": "MIT",
+ "peer": true,
"dependencies": {
"esbuild": "^0.25.0",
"fdir": "^6.5.0",
@@ -25547,6 +25737,7 @@
"integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==",
"dev": true,
"license": "MIT",
+ "peer": true,
"engines": {
"node": ">=12"
},
@@ -25560,6 +25751,7 @@
"integrity": "sha512-LUCP5ev3GURDysTWiP47wRRUpLKMOfPh+yKTx3kVIEiu5KOMeqzpnYNsKyOoVrULivR8tLcks4+lga33Whn90A==",
"dev": true,
"license": "MIT",
+ "peer": true,
"dependencies": {
"@types/chai": "^5.2.2",
"@vitest/expect": "3.2.4",
@@ -25765,6 +25957,7 @@
"resolved": "https://registry.npmjs.org/webpack/-/webpack-5.103.0.tgz",
"integrity": "sha512-HU1JOuV1OavsZ+mfigY0j8d1TgQgbZ6M+J75zDkpEAwYeXjWSqrGJtgnPblJjd/mAyTNQ7ygw0MiKOn6etz8yw==",
"license": "MIT",
+ "peer": true,
"dependencies": {
"@types/eslint-scope": "^3.7.7",
"@types/estree": "^1.0.8",
diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json
index ce42d0ba41a..0ecba62140f 100644
--- a/ui/litellm-dashboard/package.json
+++ b/ui/litellm-dashboard/package.json
@@ -9,8 +9,11 @@
"lint": "next lint",
"test": "vitest",
"test:watch": "vitest -w",
+ "test:coverage": "vitest run --coverage",
"format": "prettier --write .",
- "format:check": "prettier --check ."
+ "format:check": "prettier --check .",
+ "e2e": "playwright test --config e2e_tests/playwright.config.ts",
+ "e2e:ui": "playwright test --ui --config e2e_tests/playwright.config.ts"
},
"dependencies": {
"@anthropic-ai/sdk": "^0.54.0",
@@ -44,6 +47,7 @@
"uuid": "^11.1.0"
},
"devDependencies": {
+ "@playwright/test": "^1.57.0",
"@tailwindcss/forms": "^0.5.7",
"@testing-library/dom": "^10.4.1",
"@testing-library/jest-dom": "^6.8.0",
diff --git a/ui/litellm-dashboard/public/assets/logos/minimax.svg b/ui/litellm-dashboard/public/assets/logos/minimax.svg
new file mode 100644
index 00000000000..59b741bbcb7
--- /dev/null
+++ b/ui/litellm-dashboard/public/assets/logos/minimax.svg
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.test.ts
new file mode 100644
index 00000000000..44fc6a96836
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.test.ts
@@ -0,0 +1,332 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import React, { ReactNode } from "react";
+import { useAgents } from "./useAgents";
+import { getAgentsList } from "@/components/networking";
+import type { AgentsResponse, Agent } from "@/components/agents/types";
+
+// Mock the networking function
+vi.mock("@/components/networking", () => ({
+ getAgentsList: vi.fn(),
+}));
+
+// Mock useAuthorized hook - we can override this in individual tests
+const mockUseAuthorized = vi.fn();
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+ default: () => mockUseAuthorized(),
+}));
+
+// Import actual roles instead of mocking them
+
+// Mock data
+const mockAgents: Agent[] = [
+ {
+ agent_id: "agent-1",
+ agent_name: "Test Agent 1",
+ litellm_params: {
+ model: "gpt-3.5-turbo",
+ api_key: "test-key-1",
+ },
+ agent_card_params: {
+ description: "A test agent for unit testing",
+ },
+ created_at: "2024-01-01T00:00:00Z",
+ updated_at: "2024-01-01T00:00:00Z",
+ created_by: "user-1",
+ updated_by: "user-1",
+ },
+ {
+ agent_id: "agent-2",
+ agent_name: "Test Agent 2",
+ litellm_params: {
+ model: "claude-3",
+ api_key: "test-key-2",
+ },
+ agent_card_params: {
+ description: "Another test agent",
+ },
+ created_at: "2024-01-01T00:00:00Z",
+ updated_at: "2024-01-01T00:00:00Z",
+ created_by: "user-2",
+ updated_by: "user-2",
+ },
+];
+
+const mockAgentsResponse: AgentsResponse = {
+ agents: mockAgents,
+};
+
+describe("useAgents", () => {
+ let queryClient: QueryClient;
+
+ beforeEach(() => {
+ queryClient = new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ },
+ },
+ });
+
+ // Reset all mocks
+ vi.clearAllMocks();
+
+ // Set default mock for useAuthorized (enabled state)
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userRole: "Admin",
+ userId: "test-user-id",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+ });
+
+ const wrapper = ({ children }: { children: ReactNode }) =>
+ React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+ it("should return agents data when query is successful", async () => {
+ // Mock successful API call
+ (getAgentsList as any).mockResolvedValue(mockAgentsResponse);
+
+ const { result } = renderHook(() => useAgents(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+ expect(result.current.data).toBeUndefined();
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(mockAgentsResponse);
+ expect(result.current.error).toBeNull();
+ expect(getAgentsList).toHaveBeenCalledWith("test-access-token");
+ expect(getAgentsList).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle error when getAgentsList fails", async () => {
+ const errorMessage = "Failed to fetch agents";
+ const testError = new Error(errorMessage);
+
+ // Mock failed API call
+ (getAgentsList as any).mockRejectedValue(testError);
+
+ const { result } = renderHook(() => useAgents(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(testError);
+ expect(result.current.data).toBeUndefined();
+ expect(getAgentsList).toHaveBeenCalledWith("test-access-token");
+ expect(getAgentsList).toHaveBeenCalledTimes(1);
+ });
+
+ it("should not execute query when accessToken is missing", async () => {
+ // Mock missing accessToken
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userRole: "Admin",
+ userId: "test-user-id",
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useAgents(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(getAgentsList).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when userRole is not an admin role", async () => {
+ // Mock non-admin userRole
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userRole: "member", // Not in all_admin_roles
+ userId: "test-user-id",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useAgents(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(getAgentsList).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when userRole is null", async () => {
+ // Mock null userRole
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userRole: null,
+ userId: "test-user-id",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useAgents(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(getAgentsList).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when userRole is empty string", async () => {
+ // Mock empty string userRole
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userRole: "",
+ userId: "test-user-id",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useAgents(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(getAgentsList).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when both accessToken and userRole are missing", async () => {
+ // Mock both auth values missing
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userRole: null,
+ userId: "test-user-id",
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useAgents(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(getAgentsList).not.toHaveBeenCalled();
+ });
+
+ it("should execute query when accessToken is present and userRole is Admin", async () => {
+ // Mock successful API call
+ (getAgentsList as any).mockResolvedValue(mockAgentsResponse);
+
+ // Ensure auth values are set (already done in beforeEach)
+ const { result } = renderHook(() => useAgents(), { wrapper });
+
+ // Wait for query to execute
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ expect(getAgentsList).toHaveBeenCalledWith("test-access-token");
+ expect(getAgentsList).toHaveBeenCalledTimes(1);
+ });
+
+ it("should execute query when accessToken is present and userRole is proxy_admin", async () => {
+ // Mock successful API call
+ (getAgentsList as any).mockResolvedValue(mockAgentsResponse);
+
+ // Mock proxy_admin role
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userRole: "proxy_admin",
+ userId: "test-user-id",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useAgents(), { wrapper });
+
+ // Wait for query to execute
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ expect(getAgentsList).toHaveBeenCalledWith("test-access-token");
+ expect(getAgentsList).toHaveBeenCalledTimes(1);
+ });
+
+ it("should return empty agents array when API returns empty data", async () => {
+ // Mock API returning empty agents array
+ (getAgentsList as any).mockResolvedValue({ agents: [] });
+
+ const { result } = renderHook(() => useAgents(), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual({ agents: [] });
+ expect(getAgentsList).toHaveBeenCalledWith("test-access-token");
+ });
+
+ it("should handle network timeout error", async () => {
+ const timeoutError = new Error("Network timeout");
+
+ // Mock network timeout
+ (getAgentsList as any).mockRejectedValue(timeoutError);
+
+ const { result } = renderHook(() => useAgents(), { wrapper });
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(timeoutError);
+ expect(result.current.data).toBeUndefined();
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.test.ts
new file mode 100644
index 00000000000..ee903628f08
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.test.ts
@@ -0,0 +1,194 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import React, { ReactNode } from "react";
+import { useCredentials } from "./useCredentials";
+import { credentialListCall, CredentialsResponse, CredentialItem } from "@/components/networking";
+
+// Mock the networking function
+vi.mock("@/components/networking", () => ({
+ credentialListCall: vi.fn(),
+}));
+
+// Mock useAuthorized hook - we can override this in individual tests
+const mockUseAuthorized = vi.fn();
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+ default: () => mockUseAuthorized(),
+}));
+
+// Mock data
+const mockCredentialItems: CredentialItem[] = [
+ {
+ credential_name: "openai-api-key",
+ credential_values: { api_key: "sk-test123" },
+ credential_info: {
+ custom_llm_provider: "openai",
+ description: "OpenAI API Key for GPT models",
+ required: true,
+ },
+ },
+ {
+ credential_name: "anthropic-api-key",
+ credential_values: { api_key: "sk-ant-test456" },
+ credential_info: {
+ custom_llm_provider: "anthropic",
+ description: "Anthropic API Key for Claude models",
+ required: true,
+ },
+ },
+];
+
+const mockCredentialsResponse: CredentialsResponse = {
+ credentials: mockCredentialItems,
+};
+
+describe("useCredentials", () => {
+ let queryClient: QueryClient;
+
+ beforeEach(() => {
+ queryClient = new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ },
+ },
+ });
+
+ // Reset all mocks
+ vi.clearAllMocks();
+
+ // Set default mock for useAuthorized (enabled state)
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userRole: "Admin",
+ userId: "test-user-id",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+ });
+
+ const wrapper = ({ children }: { children: ReactNode }) =>
+ React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+ it("should return credentials data when query is successful", async () => {
+ // Mock successful API call
+ (credentialListCall as any).mockResolvedValue(mockCredentialsResponse);
+
+ const { result } = renderHook(() => useCredentials(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+ expect(result.current.data).toBeUndefined();
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(mockCredentialsResponse);
+ expect(result.current.error).toBeNull();
+ expect(credentialListCall).toHaveBeenCalledWith("test-access-token");
+ expect(credentialListCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle error when credentialListCall fails", async () => {
+ const errorMessage = "Failed to fetch credentials";
+ const testError = new Error(errorMessage);
+
+ // Mock failed API call
+ (credentialListCall as any).mockRejectedValue(testError);
+
+ const { result } = renderHook(() => useCredentials(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(testError);
+ expect(result.current.data).toBeUndefined();
+ expect(credentialListCall).toHaveBeenCalledWith("test-access-token");
+ expect(credentialListCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should not execute query when accessToken is missing", async () => {
+ // Mock missing accessToken
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userRole: "Admin",
+ userId: "test-user-id",
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useCredentials(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(credentialListCall).not.toHaveBeenCalled();
+ });
+
+ it("should return empty credentials array when API returns empty data", async () => {
+ // Mock API returning empty credentials array
+ (credentialListCall as any).mockResolvedValue({ credentials: [] });
+
+ const { result } = renderHook(() => useCredentials(), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual({ credentials: [] });
+ expect(credentialListCall).toHaveBeenCalledWith("test-access-token");
+ });
+
+ it("should handle network timeout error", async () => {
+ const timeoutError = new Error("Network timeout");
+
+ // Mock network timeout
+ (credentialListCall as any).mockRejectedValue(timeoutError);
+
+ const { result } = renderHook(() => useCredentials(), { wrapper });
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(timeoutError);
+ expect(result.current.data).toBeUndefined();
+ });
+
+ it("should execute query when accessToken is present", async () => {
+ // Mock successful API call
+ (credentialListCall as any).mockResolvedValue(mockCredentialsResponse);
+
+ // Ensure auth values are set (already done in beforeEach)
+ const { result } = renderHook(() => useCredentials(), { wrapper });
+
+ // Wait for query to execute
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ expect(credentialListCall).toHaveBeenCalledWith("test-access-token");
+ expect(credentialListCall).toHaveBeenCalledTimes(1);
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.test.ts
new file mode 100644
index 00000000000..716d6f75399
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.test.ts
@@ -0,0 +1,334 @@
+import { allEndUsersCall } from "@/components/networking";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import { renderHook, waitFor } from "@testing-library/react";
+import React, { ReactNode } from "react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import type { Customer, CustomersResponse } from "./useCustomers";
+import { useCustomers } from "./useCustomers";
+
+// Mock the networking function
+vi.mock("@/components/networking", () => ({
+ allEndUsersCall: vi.fn(),
+}));
+
+// Mock useAuthorized hook - we can override this in individual tests
+const mockUseAuthorized = vi.fn();
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+ default: () => mockUseAuthorized(),
+}));
+
+// Import actual roles instead of mocking them
+
+// Mock data
+const mockCustomers: Customer[] = [
+ {
+ user_id: "customer-1",
+ alias: "Test Customer 1",
+ spend: 150.5,
+ blocked: false,
+ allowed_model_region: "us-east-1",
+ default_model: "gpt-3.5-turbo",
+ budget_id: "budget-1",
+ litellm_budget_table: {
+ budget_id: "budget-1",
+ max_budget: 1000,
+ soft_budget: 800,
+ max_parallel_requests: 10,
+ tpm_limit: 1000,
+ rpm_limit: 100,
+ model_max_budget: { "gpt-4": 500 },
+ budget_duration: "monthly",
+ budget_reset_at: "2024-02-01T00:00:00Z",
+ created_at: "2024-01-01T00:00:00Z",
+ created_by: "admin-1",
+ updated_at: "2024-01-01T00:00:00Z",
+ updated_by: "admin-1",
+ },
+ },
+ {
+ user_id: "customer-2",
+ alias: null,
+ spend: 0,
+ blocked: true,
+ allowed_model_region: null,
+ default_model: null,
+ budget_id: null,
+ litellm_budget_table: null,
+ },
+];
+
+const mockCustomersResponse: CustomersResponse = mockCustomers;
+
+describe("useCustomers", () => {
+ let queryClient: QueryClient;
+
+ beforeEach(() => {
+ queryClient = new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ },
+ },
+ });
+
+ // Reset all mocks
+ vi.clearAllMocks();
+
+ // Set default mock for useAuthorized (enabled state)
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userRole: "Admin",
+ userId: "test-user-id",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+ });
+
+ const wrapper = ({ children }: { children: ReactNode }) =>
+ React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+ it("should return customers data when query is successful", async () => {
+ // Mock successful API call
+ (allEndUsersCall as any).mockResolvedValue(mockCustomersResponse);
+
+ const { result } = renderHook(() => useCustomers(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+ expect(result.current.data).toBeUndefined();
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(mockCustomersResponse);
+ expect(result.current.error).toBeNull();
+ expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token");
+ expect(allEndUsersCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle error when allEndUsersCall fails", async () => {
+ const errorMessage = "Failed to fetch customers";
+ const testError = new Error(errorMessage);
+
+ // Mock failed API call
+ (allEndUsersCall as any).mockRejectedValue(testError);
+
+ const { result } = renderHook(() => useCustomers(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(testError);
+ expect(result.current.data).toBeUndefined();
+ expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token");
+ expect(allEndUsersCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should not execute query when accessToken is missing", async () => {
+ // Mock missing accessToken
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userRole: "Admin",
+ userId: "test-user-id",
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useCustomers(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(allEndUsersCall).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when userRole is not an admin role", async () => {
+ // Mock non-admin userRole
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userRole: "member", // Not in all_admin_roles
+ userId: "test-user-id",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useCustomers(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(allEndUsersCall).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when userRole is null", async () => {
+ // Mock null userRole
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userRole: null,
+ userId: "test-user-id",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useCustomers(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(allEndUsersCall).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when userRole is empty string", async () => {
+ // Mock empty string userRole
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userRole: "",
+ userId: "test-user-id",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useCustomers(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(allEndUsersCall).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when both accessToken and userRole are missing", async () => {
+ // Mock both auth values missing
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userRole: null,
+ userId: "test-user-id",
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useCustomers(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(allEndUsersCall).not.toHaveBeenCalled();
+ });
+
+ it("should execute query when accessToken is present and userRole is Admin", async () => {
+ // Mock successful API call
+ (allEndUsersCall as any).mockResolvedValue(mockCustomersResponse);
+
+ // Ensure auth values are set (already done in beforeEach)
+ const { result } = renderHook(() => useCustomers(), { wrapper });
+
+ // Wait for query to execute
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token");
+ expect(allEndUsersCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should execute query when accessToken is present and userRole is proxy_admin", async () => {
+ // Mock successful API call
+ (allEndUsersCall as any).mockResolvedValue(mockCustomersResponse);
+
+ // Mock proxy_admin role
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userRole: "proxy_admin",
+ userId: "test-user-id",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useCustomers(), { wrapper });
+
+ // Wait for query to execute
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token");
+ expect(allEndUsersCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should return empty customers array when API returns empty data", async () => {
+ // Mock API returning empty customers array
+ (allEndUsersCall as any).mockResolvedValue([]);
+
+ const { result } = renderHook(() => useCustomers(), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual([]);
+ expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token");
+ });
+
+ it("should handle network timeout error", async () => {
+ const timeoutError = new Error("Network timeout");
+
+ // Mock network timeout
+ (allEndUsersCall as any).mockRejectedValue(timeoutError);
+
+ const { result } = renderHook(() => useCustomers(), { wrapper });
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(timeoutError);
+ expect(result.current.data).toBeUndefined();
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts
new file mode 100644
index 00000000000..d9e96a5308c
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts
@@ -0,0 +1,273 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import React, { ReactNode } from "react";
+import { useGuardrails } from "./useGuardrails";
+import { getGuardrailsList } from "@/components/networking";
+
+// Mock the networking function
+vi.mock("@/components/networking", () => ({
+ getGuardrailsList: vi.fn(),
+}));
+
+// Mock useAuthorized hook - we can override this in individual tests
+const mockUseAuthorized = vi.fn();
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+ default: () => mockUseAuthorized(),
+}));
+
+// Mock data
+const mockGuardrailsResponse = {
+ guardrails: [
+ { guardrail_name: "content-safety" },
+ { guardrail_name: "toxicity-filter" },
+ { guardrail_name: "pii-detection" },
+ ],
+};
+
+const expectedGuardrailNames = ["content-safety", "toxicity-filter", "pii-detection"];
+
+describe("useGuardrails", () => {
+ let queryClient: QueryClient;
+
+ beforeEach(() => {
+ queryClient = new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ },
+ },
+ });
+
+ // Reset all mocks
+ vi.clearAllMocks();
+
+ // Set default mock for useAuthorized (enabled state)
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: "test-user-id",
+ userRole: "Admin",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+ });
+
+ const wrapper = ({ children }: { children: ReactNode }) =>
+ React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+ it("should return guardrail names when query is successful", async () => {
+ // Mock successful API call
+ (getGuardrailsList as any).mockResolvedValue(mockGuardrailsResponse);
+
+ const { result } = renderHook(() => useGuardrails(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+ expect(result.current.data).toBeUndefined();
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(expectedGuardrailNames);
+ expect(result.current.error).toBeNull();
+ expect(getGuardrailsList).toHaveBeenCalledWith("test-access-token");
+ expect(getGuardrailsList).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle error when getGuardrailsList fails", async () => {
+ const errorMessage = "Failed to fetch guardrails";
+ const testError = new Error(errorMessage);
+
+ // Mock failed API call
+ (getGuardrailsList as any).mockRejectedValue(testError);
+
+ const { result } = renderHook(() => useGuardrails(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(testError);
+ expect(result.current.data).toBeUndefined();
+ expect(getGuardrailsList).toHaveBeenCalledWith("test-access-token");
+ expect(getGuardrailsList).toHaveBeenCalledTimes(1);
+ });
+
+ it("should not execute query when accessToken is missing", async () => {
+ // Mock missing accessToken
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userId: "test-user-id",
+ userRole: "Admin",
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useGuardrails(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(getGuardrailsList).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when userId is missing", async () => {
+ // Mock missing userId
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: null,
+ userRole: "Admin",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useGuardrails(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(getGuardrailsList).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when userRole is missing", async () => {
+ // Mock missing userRole
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: "test-user-id",
+ userRole: null,
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useGuardrails(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(getGuardrailsList).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when all auth values are missing", async () => {
+ // Mock all auth values missing
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userId: null,
+ userRole: null,
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useGuardrails(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(getGuardrailsList).not.toHaveBeenCalled();
+ });
+
+ it("should execute query when all auth values are present", async () => {
+ // Mock successful API call
+ (getGuardrailsList as any).mockResolvedValue(mockGuardrailsResponse);
+
+ // Ensure all auth values are present (already set in beforeEach)
+ const { result } = renderHook(() => useGuardrails(), { wrapper });
+
+ // Wait for query to execute
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ expect(getGuardrailsList).toHaveBeenCalledWith("test-access-token");
+ expect(getGuardrailsList).toHaveBeenCalledTimes(1);
+ });
+
+ it("should return empty array when API returns empty guardrails", async () => {
+ // Mock API returning empty guardrails array
+ (getGuardrailsList as any).mockResolvedValue({ guardrails: [] });
+
+ const { result } = renderHook(() => useGuardrails(), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual([]);
+ expect(getGuardrailsList).toHaveBeenCalledWith("test-access-token");
+ });
+
+ it("should handle network timeout error", async () => {
+ const timeoutError = new Error("Network timeout");
+
+ // Mock network timeout
+ (getGuardrailsList as any).mockRejectedValue(timeoutError);
+
+ const { result } = renderHook(() => useGuardrails(), { wrapper });
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(timeoutError);
+ expect(result.current.data).toBeUndefined();
+ });
+
+ it("should correctly transform guardrail objects to names array", async () => {
+ const customGuardrailsResponse = {
+ guardrails: [{ guardrail_name: "custom-guardrail-1" }, { guardrail_name: "custom-guardrail-2" }],
+ };
+ const expectedNames = ["custom-guardrail-1", "custom-guardrail-2"];
+
+ // Mock API call with custom data
+ (getGuardrailsList as any).mockResolvedValue(customGuardrailsResponse);
+
+ const { result } = renderHook(() => useGuardrails(), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(expectedNames);
+ expect(result.current.data).toHaveLength(2);
+ expect(result.current.data).toContain("custom-guardrail-1");
+ expect(result.current.data).toContain("custom-guardrail-2");
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts
new file mode 100644
index 00000000000..9786b7fa359
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts
@@ -0,0 +1,18 @@
+import { useQuery, UseQueryResult } from "@tanstack/react-query";
+import { createQueryKeys } from "../common/queryKeysFactory";
+import { getGuardrailsList } from "@/components/networking";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+
+const guardrailKeys = createQueryKeys("guardrails");
+
+export const useGuardrails = (): UseQueryResult => {
+ const { accessToken, userId, userRole } = useAuthorized();
+ return useQuery({
+ queryKey: guardrailKeys.list({}),
+ queryFn: async () => {
+ const response = await getGuardrailsList(accessToken!);
+ return response.guardrails.map((g: { guardrail_name: string }) => g.guardrail_name);
+ },
+ enabled: Boolean(accessToken && userId && userRole),
+ });
+};
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts
new file mode 100644
index 00000000000..db394b9f7f8
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts
@@ -0,0 +1,27 @@
+import { getProxyBaseUrl } from "@/components/networking";
+import { useQuery, UseQueryResult } from "@tanstack/react-query";
+import { createQueryKeys } from "../common/queryKeysFactory";
+
+const healthReadinessKeys = createQueryKeys("healthReadiness");
+
+interface HealthReadinessResponse {
+ litellm_version?: string;
+ [key: string]: any;
+}
+
+const fetchHealthReadiness = async (): Promise => {
+ const baseUrl = getProxyBaseUrl();
+ const response = await fetch(`${baseUrl}/health/readiness`);
+ if (!response.ok) {
+ throw new Error(`Failed to fetch health readiness: ${response.statusText}`);
+ }
+ return response.json();
+};
+
+export const useHealthReadiness = (): UseQueryResult => {
+ return useQuery({
+ queryKey: healthReadinessKeys.detail("readiness"),
+ queryFn: fetchHealthReadiness,
+ staleTime: 5 * 60 * 1000, // 5 minutes
+ });
+};
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts
new file mode 100644
index 00000000000..c4ffb7041aa
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts
@@ -0,0 +1,362 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import React, { ReactNode } from "react";
+import { useKeys } from "./useKeys";
+import { keyListCall } from "@/components/networking";
+import type { KeyResponse } from "@/components/key_team_helpers/key_list";
+
+// Mock the networking function
+vi.mock("@/components/networking", () => ({
+ keyListCall: vi.fn(),
+}));
+
+// Mock useAuthorized hook - we can override this in individual tests
+const mockUseAuthorized = vi.fn();
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+ default: () => mockUseAuthorized(),
+}));
+
+// Mock data
+const mockKeys: KeyResponse[] = [
+ {
+ token: "sk-test-key-1",
+ token_id: "key-1",
+ key_name: "Test Key 1",
+ key_alias: "test-key-1",
+ spend: 10.5,
+ max_budget: 100,
+ expires: "2024-12-31T23:59:59Z",
+ models: ["gpt-3.5-turbo"],
+ aliases: {},
+ config: {},
+ user_id: "user-1",
+ team_id: null,
+ max_parallel_requests: 10,
+ metadata: {},
+ tpm_limit: 1000,
+ rpm_limit: 100,
+ duration: "30d",
+ budget_duration: "1mo",
+ budget_reset_at: "2024-02-01T00:00:00Z",
+ allowed_cache_controls: [],
+ allowed_routes: [],
+ permissions: {},
+ model_spend: { "gpt-3.5-turbo": 10.5 },
+ model_max_budget: { "gpt-3.5-turbo": 100 },
+ soft_budget_cooldown: false,
+ blocked: false,
+ litellm_budget_table: {},
+ organization_id: null,
+ created_at: "2024-01-01T00:00:00Z",
+ updated_at: "2024-01-01T00:00:00Z",
+ team_spend: 0,
+ team_alias: "",
+ team_tpm_limit: 0,
+ team_rpm_limit: 0,
+ team_max_budget: 0,
+ team_models: [],
+ team_blocked: false,
+ soft_budget: 0,
+ team_model_aliases: {},
+ team_member_spend: 0,
+ team_metadata: {},
+ end_user_id: "",
+ end_user_tpm_limit: 0,
+ end_user_rpm_limit: 0,
+ end_user_max_budget: 0,
+ last_refreshed_at: 0,
+ api_key: "",
+ user_role: "user",
+ rpm_limit_per_model: {},
+ tpm_limit_per_model: {},
+ user_tpm_limit: 0,
+ user_rpm_limit: 0,
+ user_email: "",
+ },
+ {
+ token: "sk-test-key-2",
+ token_id: "key-2",
+ key_name: "Test Key 2",
+ key_alias: "test-key-2",
+ spend: 25.0,
+ max_budget: 200,
+ expires: "2024-12-31T23:59:59Z",
+ models: ["claude-3"],
+ aliases: {},
+ config: {},
+ user_id: "user-2",
+ team_id: "team-1",
+ max_parallel_requests: 5,
+ metadata: {},
+ tpm_limit: 500,
+ rpm_limit: 50,
+ duration: "30d",
+ budget_duration: "1mo",
+ budget_reset_at: "2024-02-01T00:00:00Z",
+ allowed_cache_controls: [],
+ allowed_routes: [],
+ permissions: {},
+ model_spend: { "claude-3": 25.0 },
+ model_max_budget: { "claude-3": 200 },
+ soft_budget_cooldown: false,
+ blocked: false,
+ litellm_budget_table: {},
+ organization_id: null,
+ created_at: "2024-01-01T00:00:00Z",
+ updated_at: "2024-01-01T00:00:00Z",
+ team_spend: 0,
+ team_alias: "test-team",
+ team_tpm_limit: 1000,
+ team_rpm_limit: 100,
+ team_max_budget: 500,
+ team_models: ["claude-3"],
+ team_blocked: false,
+ soft_budget: 0,
+ team_model_aliases: {},
+ team_member_spend: 0,
+ team_metadata: {},
+ end_user_id: "",
+ end_user_tpm_limit: 0,
+ end_user_rpm_limit: 0,
+ end_user_max_budget: 0,
+ last_refreshed_at: 0,
+ api_key: "",
+ user_role: "user",
+ rpm_limit_per_model: {},
+ tpm_limit_per_model: {},
+ user_tpm_limit: 0,
+ user_rpm_limit: 0,
+ user_email: "",
+ },
+];
+
+const mockKeysResponse = {
+ keys: mockKeys,
+ total_count: 2,
+ current_page: 1,
+ total_pages: 1,
+};
+
+describe("useKeys", () => {
+ let queryClient: QueryClient;
+
+ beforeEach(() => {
+ queryClient = new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ },
+ },
+ });
+
+ // Reset all mocks
+ vi.clearAllMocks();
+
+ // Set default mock for useAuthorized (enabled state)
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userRole: "Admin",
+ userId: "test-user-id",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+ });
+
+ const wrapper = ({ children }: { children: ReactNode }) =>
+ React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+ it("should return keys data when query is successful", async () => {
+ // Mock successful API call
+ (keyListCall as any).mockResolvedValue(mockKeysResponse);
+
+ const { result } = renderHook(() => useKeys(1, 10), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+ expect(result.current.data).toBeUndefined();
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(mockKeysResponse);
+ expect(result.current.error).toBeNull();
+ expect(keyListCall).toHaveBeenCalledWith(
+ "test-access-token",
+ null, // organizationID
+ null, // teamID
+ null, // selectedKeyAlias
+ null, // userID
+ null, // keyHash
+ 1, // page
+ 10, // pageSize
+ );
+ expect(keyListCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle error when keyListCall fails", async () => {
+ const errorMessage = "Failed to fetch keys";
+ const testError = new Error(errorMessage);
+
+ // Mock failed API call
+ (keyListCall as any).mockRejectedValue(testError);
+
+ const { result } = renderHook(() => useKeys(1, 10), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(testError);
+ expect(result.current.data).toBeUndefined();
+ expect(keyListCall).toHaveBeenCalledWith(
+ "test-access-token",
+ null, // organizationID
+ null, // teamID
+ null, // selectedKeyAlias
+ null, // userID
+ null, // keyHash
+ 1, // page
+ 10, // pageSize
+ );
+ expect(keyListCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should not execute query when accessToken is missing", async () => {
+ // Mock missing accessToken
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userRole: "Admin",
+ userId: "test-user-id",
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useKeys(1, 10), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(keyListCall).not.toHaveBeenCalled();
+ });
+
+ it("should pass correct page and pageSize parameters to the API", async () => {
+ // Mock successful API call
+ (keyListCall as any).mockResolvedValue(mockKeysResponse);
+
+ const page = 2;
+ const pageSize = 20;
+
+ const { result } = renderHook(() => useKeys(page, pageSize), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ expect(keyListCall).toHaveBeenCalledWith(
+ "test-access-token",
+ null, // organizationID
+ null, // teamID
+ null, // selectedKeyAlias
+ null, // userID
+ null, // keyHash
+ page, // page
+ pageSize, // pageSize
+ );
+ });
+
+ it("should return empty keys array when API returns empty data", async () => {
+ // Mock API returning empty keys array
+ const emptyResponse = {
+ keys: [],
+ total_count: 0,
+ current_page: 1,
+ total_pages: 0,
+ };
+ (keyListCall as any).mockResolvedValue(emptyResponse);
+
+ const { result } = renderHook(() => useKeys(1, 10), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(emptyResponse);
+ expect(keyListCall).toHaveBeenCalledWith(
+ "test-access-token",
+ null, // organizationID
+ null, // teamID
+ null, // selectedKeyAlias
+ null, // userID
+ null, // keyHash
+ 1, // page
+ 10, // pageSize
+ );
+ });
+
+ it("should handle network timeout error", async () => {
+ const timeoutError = new Error("Network timeout");
+
+ // Mock network timeout
+ (keyListCall as any).mockRejectedValue(timeoutError);
+
+ const { result } = renderHook(() => useKeys(1, 10), { wrapper });
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(timeoutError);
+ expect(result.current.data).toBeUndefined();
+ });
+
+ it("should handle pagination correctly", async () => {
+ const paginatedResponse = {
+ keys: [mockKeys[0]], // Only first key
+ total_count: 15,
+ current_page: 2,
+ total_pages: 2,
+ };
+ (keyListCall as any).mockResolvedValue(paginatedResponse);
+
+ const { result } = renderHook(() => useKeys(2, 10), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ expect(result.current.data).toEqual(paginatedResponse);
+ expect(keyListCall).toHaveBeenCalledWith(
+ "test-access-token",
+ null, // organizationID
+ null, // teamID
+ null, // selectedKeyAlias
+ null, // userID
+ null, // keyHash
+ 2, // page
+ 10, // pageSize
+ );
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts
new file mode 100644
index 00000000000..8ae4d76ff5d
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts
@@ -0,0 +1,36 @@
+import { keepPreviousData, useQuery, UseQueryResult } from "@tanstack/react-query";
+import { createQueryKeys } from "../common/queryKeysFactory";
+import { keyListCall } from "@/components/networking";
+import { KeyResponse } from "@/components/key_team_helpers/key_list";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+
+const keyKeys = createQueryKeys("keys");
+
+export interface KeysResponse {
+ keys: KeyResponse[];
+ total_count: number;
+ current_page: number;
+ total_pages: number;
+}
+
+export const useKeys = (page: number, pageSize: number): UseQueryResult => {
+ const { accessToken } = useAuthorized();
+
+ return useQuery({
+ queryKey: keyKeys.list({ page, limit: pageSize }),
+ queryFn: async () =>
+ await keyListCall(
+ accessToken!,
+ null, // organizationID
+ null, // teamID
+ null, // selectedKeyAlias
+ null, // userID
+ null, // keyHash
+ page,
+ pageSize,
+ ),
+ enabled: Boolean(accessToken),
+ staleTime: 30000, // 30 seconds
+ placeholderData: keepPreviousData,
+ });
+};
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts
new file mode 100644
index 00000000000..be910acf7e4
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts
@@ -0,0 +1,127 @@
+/* @vitest-environment jsdom */
+import React from "react";
+import { renderHook, waitFor } from "@testing-library/react";
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import { useMCPServerHealth } from "./useMCPServerHealth";
+import * as networking from "@/components/networking";
+
+// Mock the networking module
+vi.mock("@/components/networking", () => ({
+ fetchMCPServerHealth: vi.fn(),
+}));
+
+// Mock useAuthorized hook
+vi.mock("../useAuthorized", () => ({
+ default: vi.fn(() => ({
+ accessToken: "test-token-123",
+ })),
+}));
+
+const createQueryClient = () =>
+ new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ gcTime: 0,
+ },
+ },
+ });
+
+const wrapper = ({ children }: { children: React.ReactNode }) => {
+ const queryClient = createQueryClient();
+ return React.createElement(QueryClientProvider, { client: queryClient }, children);
+};
+
+describe("useMCPServerHealth", () => {
+ beforeEach(() => {
+ vi.clearAllMocks();
+ });
+
+ it("should fetch health status for given server IDs", async () => {
+ const mockHealthStatuses = [
+ { server_id: "server-1", status: "healthy" },
+ { server_id: "server-2", status: "unhealthy" },
+ ];
+
+ vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses);
+
+ const { result } = renderHook(() => useMCPServerHealth(["server-1", "server-2"]), {
+ wrapper,
+ });
+
+ await waitFor(() => {
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", ["server-1", "server-2"]);
+ expect(result.current.data).toEqual(mockHealthStatuses);
+ });
+
+ it("should fetch health status for all servers when no server IDs provided", async () => {
+ const mockHealthStatuses = [
+ { server_id: "server-1", status: "healthy" },
+ { server_id: "server-2", status: "healthy" },
+ { server_id: "server-3", status: "unhealthy" },
+ ];
+
+ vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses);
+
+ const { result } = renderHook(() => useMCPServerHealth(), {
+ wrapper,
+ });
+
+ await waitFor(() => {
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", undefined);
+ expect(result.current.data).toEqual(mockHealthStatuses);
+ });
+
+ it("should handle empty server list", async () => {
+ vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]);
+
+ const { result } = renderHook(() => useMCPServerHealth([]), {
+ wrapper,
+ });
+
+ await waitFor(() => {
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", []);
+ expect(result.current.data).toEqual([]);
+ });
+
+ it("should handle errors when fetching health status", async () => {
+ const mockError = new Error("Failed to fetch health status");
+ vi.mocked(networking.fetchMCPServerHealth).mockRejectedValue(mockError);
+
+ const { result } = renderHook(() => useMCPServerHealth(["server-1"]), {
+ wrapper,
+ });
+
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(mockError);
+ });
+
+ it("should not fetch when accessToken is not available", async () => {
+ // Mock useAuthorized to return no token
+ const useAuthorizedModule = await import("../useAuthorized");
+ vi.mocked(useAuthorizedModule.default).mockReturnValue({
+ accessToken: null,
+ } as any);
+
+ const { result } = renderHook(() => useMCPServerHealth(["server-1"]), {
+ wrapper,
+ });
+
+ // Should remain in idle state since query is not enabled
+ expect(result.current.status).toBe("pending");
+ expect(networking.fetchMCPServerHealth).not.toHaveBeenCalled();
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts
new file mode 100644
index 00000000000..95d7f3bcee0
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts
@@ -0,0 +1,22 @@
+import { useQuery } from "@tanstack/react-query";
+import { createQueryKeys } from "../common/queryKeysFactory";
+import { fetchMCPServerHealth } from "@/components/networking";
+import useAuthorized from "../useAuthorized";
+
+const mcpServerHealthKeys = createQueryKeys("mcpServerHealth");
+
+interface MCPServerHealth {
+ server_id: string;
+ status: string;
+}
+
+export const useMCPServerHealth = (serverIds?: string[]) => {
+ const { accessToken } = useAuthorized();
+ return useQuery({
+ queryKey: [...mcpServerHealthKeys.lists(), { serverIds }],
+ queryFn: async () => await fetchMCPServerHealth(accessToken!, serverIds),
+ enabled: !!accessToken,
+ // Refetch health status every 30 seconds to keep it up to date
+ refetchInterval: 30000,
+ });
+};
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.test.ts
new file mode 100644
index 00000000000..f79ca33bc5d
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.test.ts
@@ -0,0 +1,144 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import React, { ReactNode } from "react";
+import { useModelCostMap } from "./useModelCostMap";
+import { modelCostMap } from "@/components/networking";
+
+// Mock the networking function
+vi.mock("@/components/networking", () => ({
+ modelCostMap: vi.fn(),
+}));
+
+// Mock data
+const mockModelCostData: Record = {
+ "gpt-3.5-turbo": {
+ litellm_provider: "openai",
+ input_cost_per_token: 0.0015,
+ output_cost_per_token: 0.002,
+ },
+ "claude-3-sonnet-20240229": {
+ litellm_provider: "anthropic",
+ input_cost_per_token: 0.003,
+ output_cost_per_token: 0.015,
+ },
+};
+
+describe("useModelCostMap", () => {
+ let queryClient: QueryClient;
+
+ beforeEach(() => {
+ queryClient = new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ },
+ },
+ });
+
+ // Reset all mocks
+ vi.clearAllMocks();
+ });
+
+ const wrapper = ({ children }: { children: ReactNode }) =>
+ React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+ it("should return model cost map data when query is successful", async () => {
+ // Mock successful API call
+ (modelCostMap as any).mockResolvedValue(mockModelCostData);
+
+ const { result } = renderHook(() => useModelCostMap(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+ expect(result.current.data).toBeUndefined();
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(mockModelCostData);
+ expect(result.current.error).toBeNull();
+ expect(modelCostMap).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle error when modelCostMap fails", async () => {
+ const errorMessage = "Failed to fetch model cost map";
+ const testError = new Error(errorMessage);
+
+ // Mock failed API call
+ (modelCostMap as any).mockRejectedValue(testError);
+
+ const { result } = renderHook(() => useModelCostMap(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(testError);
+ expect(result.current.data).toBeUndefined();
+ expect(modelCostMap).toHaveBeenCalledTimes(1);
+ });
+
+ it("should return empty object when API returns empty data", async () => {
+ // Mock API returning empty object
+ (modelCostMap as any).mockResolvedValue({});
+
+ const { result } = renderHook(() => useModelCostMap(), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual({});
+ expect(modelCostMap).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle network timeout error", async () => {
+ const timeoutError = new Error("Network timeout");
+
+ // Mock network timeout
+ (modelCostMap as any).mockRejectedValue(timeoutError);
+
+ const { result } = renderHook(() => useModelCostMap(), { wrapper });
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(timeoutError);
+ expect(result.current.data).toBeUndefined();
+ });
+
+ it("should have correct query configuration", async () => {
+ // Mock successful API call
+ (modelCostMap as any).mockResolvedValue(mockModelCostData);
+
+ const { result } = renderHook(() => useModelCostMap(), { wrapper });
+
+ // Wait for query to complete
+ await waitFor(() => {
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ // Verify the query was called
+ expect(modelCostMap).toHaveBeenCalledTimes(1);
+
+ // The hook should have the expected properties from useQuery
+ expect(result.current).toHaveProperty("data");
+ expect(result.current).toHaveProperty("isLoading");
+ expect(result.current).toHaveProperty("isError");
+ expect(result.current).toHaveProperty("isSuccess");
+ expect(result.current).toHaveProperty("error");
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.ts
new file mode 100644
index 00000000000..2d82eedf25c
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.ts
@@ -0,0 +1,14 @@
+import { modelCostMap } from "@/components/networking";
+import { useQuery } from "@tanstack/react-query";
+import { createQueryKeys } from "../common/queryKeysFactory";
+
+const modelCostMapKeys = createQueryKeys("modelCostMap");
+
+export const useModelCostMap = () => {
+ return useQuery>({
+ queryKey: modelCostMapKeys.list({}),
+ queryFn: async () => await modelCostMap(),
+ staleTime: 60 * 1000, // 1 minute
+ gcTime: 60 * 1000, // 1 minute
+ });
+};
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts
new file mode 100644
index 00000000000..66c005f37c4
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts
@@ -0,0 +1,282 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import React, { ReactNode } from "react";
+import { useOrganizations } from "./useOrganizations";
+import { organizationListCall } from "@/components/networking";
+import type { Organization } from "@/components/networking";
+
+// Mock the networking function
+vi.mock("@/components/networking", () => ({
+ organizationListCall: vi.fn(),
+}));
+
+// Mock useAuthorized hook - we can override this in individual tests
+const mockUseAuthorized = vi.fn();
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+ default: () => mockUseAuthorized(),
+}));
+
+// Mock data
+const mockOrganizations: Organization[] = [
+ {
+ organization_id: "org-1",
+ organization_alias: "Test Organization 1",
+ budget_id: "budget-1",
+ metadata: {},
+ models: ["gpt-3.5-turbo", "gpt-4"],
+ spend: 100.5,
+ model_spend: { "gpt-3.5-turbo": 50.25, "gpt-4": 50.25 },
+ created_at: "2024-01-01T00:00:00Z",
+ created_by: "user-1",
+ updated_at: "2024-01-01T00:00:00Z",
+ updated_by: "user-1",
+ litellm_budget_table: null,
+ teams: null,
+ users: null,
+ members: [
+ { user_id: "user-1", user_role: "admin" },
+ { user_id: "user-2", user_role: "member" },
+ ],
+ },
+ {
+ organization_id: "org-2",
+ organization_alias: "Test Organization 2",
+ budget_id: "budget-2",
+ metadata: {},
+ models: ["claude-3"],
+ spend: 250.75,
+ model_spend: { "claude-3": 250.75 },
+ created_at: "2024-01-01T00:00:00Z",
+ created_by: "user-3",
+ updated_at: "2024-01-01T00:00:00Z",
+ updated_by: "user-3",
+ litellm_budget_table: null,
+ teams: null,
+ users: null,
+ members: [{ user_id: "user-3", user_role: "admin" }],
+ },
+];
+
+describe("useOrganizations", () => {
+ let queryClient: QueryClient;
+
+ beforeEach(() => {
+ queryClient = new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ },
+ },
+ });
+
+ // Reset all mocks
+ vi.clearAllMocks();
+
+ // Set default mock for useAuthorized (enabled state)
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: "test-user-id",
+ userRole: "Admin",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+ });
+
+ const wrapper = ({ children }: { children: ReactNode }) =>
+ React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+ it("should return organizations data when query is successful", async () => {
+ // Mock successful API call
+ (organizationListCall as any).mockResolvedValue(mockOrganizations);
+
+ const { result } = renderHook(() => useOrganizations(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+ expect(result.current.data).toBeUndefined();
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(mockOrganizations);
+ expect(result.current.error).toBeNull();
+ expect(organizationListCall).toHaveBeenCalledWith("test-access-token");
+ expect(organizationListCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle error when organizationListCall fails", async () => {
+ const errorMessage = "Failed to fetch organizations";
+ const testError = new Error(errorMessage);
+
+ // Mock failed API call
+ (organizationListCall as any).mockRejectedValue(testError);
+
+ const { result } = renderHook(() => useOrganizations(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(testError);
+ expect(result.current.data).toBeUndefined();
+ expect(organizationListCall).toHaveBeenCalledWith("test-access-token");
+ expect(organizationListCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should not execute query when accessToken is missing", async () => {
+ // Mock missing accessToken
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userId: "test-user-id",
+ userRole: "Admin",
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useOrganizations(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(organizationListCall).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when userId is missing", async () => {
+ // Mock missing userId
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: null,
+ userRole: "Admin",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useOrganizations(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(organizationListCall).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when userRole is missing", async () => {
+ // Mock missing userRole
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: "test-user-id",
+ userRole: null,
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useOrganizations(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(organizationListCall).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when all auth values are missing", async () => {
+ // Mock all auth values missing
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userId: null,
+ userRole: null,
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useOrganizations(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(organizationListCall).not.toHaveBeenCalled();
+ });
+
+ it("should execute query when all auth values are present", async () => {
+ // Mock successful API call
+ (organizationListCall as any).mockResolvedValue(mockOrganizations);
+
+ // Ensure all auth values are present (already set in beforeEach)
+ const { result } = renderHook(() => useOrganizations(), { wrapper });
+
+ // Wait for query to execute
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ expect(organizationListCall).toHaveBeenCalledWith("test-access-token");
+ expect(organizationListCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should return empty array when API returns empty data", async () => {
+ // Mock API returning empty array
+ (organizationListCall as any).mockResolvedValue([]);
+
+ const { result } = renderHook(() => useOrganizations(), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual([]);
+ expect(organizationListCall).toHaveBeenCalledWith("test-access-token");
+ });
+
+ it("should handle network timeout error", async () => {
+ const timeoutError = new Error("Network timeout");
+
+ // Mock network timeout
+ (organizationListCall as any).mockRejectedValue(timeoutError);
+
+ const { result } = renderHook(() => useOrganizations(), { wrapper });
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(timeoutError);
+ expect(result.current.data).toBeUndefined();
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts
index 57c9c057652..27a946d112a 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts
@@ -6,8 +6,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
const organizationKeys = createQueryKeys("organizations");
export const useOrganizations = (): UseQueryResult => {
- const { accessToken } = useAuthorized();
- const { userId, userRole } = useAuthorized();
+ const { accessToken, userId, userRole } = useAuthorized();
return useQuery({
queryKey: organizationKeys.list({}),
queryFn: async () => await organizationListCall(accessToken!),
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/providers/useProviderFields.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/providers/useProviderFields.test.ts
new file mode 100644
index 00000000000..33242e0452f
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/providers/useProviderFields.test.ts
@@ -0,0 +1,182 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import React, { ReactNode } from "react";
+import { useProviderFields } from "./useProviderFields";
+import { getProviderCreateMetadata } from "@/components/networking";
+import type { ProviderCreateInfo } from "@/components/networking";
+
+// Mock the networking function
+vi.mock("@/components/networking", () => ({
+ getProviderCreateMetadata: vi.fn(),
+}));
+
+// Mock data
+const mockProviderFields: ProviderCreateInfo[] = [
+ {
+ provider: "OpenAI",
+ provider_display_name: "OpenAI",
+ litellm_provider: "openai",
+ default_model_placeholder: "gpt-3.5-turbo",
+ credential_fields: [],
+ },
+ {
+ provider: "Anthropic",
+ provider_display_name: "Anthropic",
+ litellm_provider: "anthropic",
+ default_model_placeholder: "claude-3-sonnet-20240229",
+ credential_fields: [],
+ },
+ {
+ provider: "Azure",
+ provider_display_name: "Azure OpenAI",
+ litellm_provider: "azure",
+ default_model_placeholder: "gpt-35-turbo",
+ credential_fields: [],
+ },
+];
+
+describe("useProviderFields", () => {
+ let queryClient: QueryClient;
+
+ beforeEach(() => {
+ queryClient = new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ },
+ },
+ });
+
+ // Reset all mocks
+ vi.clearAllMocks();
+ });
+
+ const wrapper = ({ children }: { children: ReactNode }) =>
+ React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+ it("should return provider fields data when query is successful", async () => {
+ // Mock successful API call
+ (getProviderCreateMetadata as any).mockResolvedValue(mockProviderFields);
+
+ const { result } = renderHook(() => useProviderFields(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+ expect(result.current.data).toBeUndefined();
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(mockProviderFields);
+ expect(result.current.error).toBeNull();
+ expect(getProviderCreateMetadata).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle error when getProviderCreateMetadata fails", async () => {
+ const errorMessage = "Failed to fetch provider fields";
+ const testError = new Error(errorMessage);
+
+ // Mock failed API call
+ (getProviderCreateMetadata as any).mockRejectedValue(testError);
+
+ const { result } = renderHook(() => useProviderFields(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(testError);
+ expect(result.current.data).toBeUndefined();
+ expect(getProviderCreateMetadata).toHaveBeenCalledTimes(1);
+ });
+
+ it("should return empty array when API returns empty data", async () => {
+ // Mock API returning empty array
+ (getProviderCreateMetadata as any).mockResolvedValue([]);
+
+ const { result } = renderHook(() => useProviderFields(), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual([]);
+ expect(getProviderCreateMetadata).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle network timeout error", async () => {
+ const timeoutError = new Error("Network timeout");
+
+ // Mock network timeout
+ (getProviderCreateMetadata as any).mockRejectedValue(timeoutError);
+
+ const { result } = renderHook(() => useProviderFields(), { wrapper });
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(timeoutError);
+ expect(result.current.data).toBeUndefined();
+ });
+
+ it("should have correct query configuration", async () => {
+ // Mock successful API call
+ (getProviderCreateMetadata as any).mockResolvedValue(mockProviderFields);
+
+ const { result } = renderHook(() => useProviderFields(), { wrapper });
+
+ // Wait for query to complete
+ await waitFor(() => {
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ // Verify the query was called
+ expect(getProviderCreateMetadata).toHaveBeenCalledTimes(1);
+
+ // The hook should have the expected properties from useQuery
+ expect(result.current).toHaveProperty("data");
+ expect(result.current).toHaveProperty("isLoading");
+ expect(result.current).toHaveProperty("isError");
+ expect(result.current).toHaveProperty("isSuccess");
+ expect(result.current).toHaveProperty("error");
+ });
+
+ it("should return provider fields with populated credential fields", async () => {
+ const mockFieldsWithCredentials: ProviderCreateInfo[] = [
+ {
+ provider: "TestProvider",
+ provider_display_name: "Test Provider",
+ litellm_provider: "test",
+ default_model_placeholder: "test-model",
+ credential_fields: [], // Keeping empty as per existing test patterns
+ },
+ ];
+
+ // Mock successful API call with provider that has credential fields
+ (getProviderCreateMetadata as any).mockResolvedValue(mockFieldsWithCredentials);
+
+ const { result } = renderHook(() => useProviderFields(), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(mockFieldsWithCredentials);
+ expect(result.current.data?.[0].provider).toBe("TestProvider");
+ expect(result.current.data?.[0].litellm_provider).toBe("test");
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useEditSSOSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useEditSSOSettings.ts
new file mode 100644
index 00000000000..69e52d0ff25
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useEditSSOSettings.ts
@@ -0,0 +1,38 @@
+import { useMutation, UseMutationResult } from "@tanstack/react-query";
+import { updateSSOSettings } from "@/components/networking";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+
+export interface EditSSOSettingsParams {
+ google_client_id?: string | null;
+ google_client_secret?: string | null;
+ microsoft_client_id?: string | null;
+ microsoft_client_secret?: string | null;
+ microsoft_tenant?: string | null;
+ generic_client_id?: string | null;
+ generic_client_secret?: string | null;
+ generic_authorization_endpoint?: string | null;
+ generic_token_endpoint?: string | null;
+ generic_userinfo_endpoint?: string | null;
+ proxy_base_url?: string | null;
+ user_email?: string | null;
+ sso_provider?: string | null;
+ role_mappings?: any;
+ [key: string]: any;
+}
+
+export interface EditSSOSettingsResponse {
+ [key: string]: any;
+}
+
+export const useEditSSOSettings = (): UseMutationResult => {
+ const { accessToken } = useAuthorized();
+
+ return useMutation({
+ mutationFn: async (params: EditSSOSettingsParams) => {
+ if (!accessToken) {
+ throw new Error("Access token is required");
+ }
+ return await updateSSOSettings(accessToken, params);
+ },
+ });
+};
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts
new file mode 100644
index 00000000000..3e09c3c2ca8
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts
@@ -0,0 +1,54 @@
+import { useQuery, UseQueryResult } from "@tanstack/react-query";
+import { createQueryKeys } from "../common/queryKeysFactory";
+import { getSSOSettings } from "@/components/networking";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+
+export interface SSOFieldSchema {
+ description: string;
+ properties: {
+ [key: string]: {
+ description: string;
+ type: string;
+ };
+ };
+}
+
+export interface SSOSettingsValues {
+ google_client_id: string | null;
+ google_client_secret: string | null;
+ microsoft_client_id: string | null;
+ microsoft_client_secret: string | null;
+ microsoft_tenant: string | null;
+ generic_client_id: string | null;
+ generic_client_secret: string | null;
+ generic_authorization_endpoint: string | null;
+ generic_token_endpoint: string | null;
+ generic_userinfo_endpoint: string | null;
+ proxy_base_url: string | null;
+ user_email: string | null;
+ ui_access_mode: string | null;
+ role_mappings: {
+ provider: string;
+ group_claim: string;
+ default_role: "internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer";
+ roles: {
+ [key: string]: string[];
+ };
+ };
+}
+
+export interface SSOSettingsResponse {
+ values: SSOSettingsValues;
+ field_schema: SSOFieldSchema;
+}
+
+const ssoKeys = createQueryKeys("sso");
+
+export const useSSOSettings = (): UseQueryResult => {
+ const { accessToken, userId, userRole } = useAuthorized();
+ return useQuery({
+ queryKey: ssoKeys.detail("settings"),
+ queryFn: async () => await getSSOSettings(accessToken!),
+ enabled: Boolean(accessToken && userId && userRole),
+ });
+};
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.test.ts
new file mode 100644
index 00000000000..a1751339568
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.test.ts
@@ -0,0 +1,283 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import React, { ReactNode } from "react";
+import { useTags } from "./useTags";
+import { tagListCall } from "@/components/networking";
+import type { TagListResponse } from "@/components/tag_management/types";
+
+// Mock the networking function
+vi.mock("@/components/networking", () => ({
+ tagListCall: vi.fn(),
+}));
+
+// Mock useAuthorized hook - we can override this in individual tests
+const mockUseAuthorized = vi.fn();
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+ default: () => mockUseAuthorized(),
+}));
+
+// Mock data
+const mockTags: TagListResponse = {
+ "tag-1": {
+ name: "tag-1",
+ description: "Test tag 1 description",
+ models: ["gpt-3.5-turbo", "gpt-4"],
+ model_info: { "gpt-3.5-turbo": "GPT-3.5 Turbo", "gpt-4": "GPT-4" },
+ created_at: "2024-01-01T00:00:00Z",
+ updated_at: "2024-01-01T00:00:00Z",
+ created_by: "user-1",
+ updated_by: "user-1",
+ litellm_budget_table: {
+ max_budget: 1000,
+ soft_budget: 800,
+ tpm_limit: 100000,
+ rpm_limit: 1000,
+ max_parallel_requests: 10,
+ budget_duration: "monthly",
+ model_max_budget: { "gpt-3.5-turbo": 500, "gpt-4": 500 },
+ },
+ },
+ "tag-2": {
+ name: "tag-2",
+ description: "Test tag 2 description",
+ models: ["claude-3"],
+ model_info: { "claude-3": "Claude 3" },
+ created_at: "2024-01-02T00:00:00Z",
+ updated_at: "2024-01-02T00:00:00Z",
+ created_by: "user-2",
+ updated_by: "user-2",
+ litellm_budget_table: {
+ max_budget: 2000,
+ soft_budget: 1500,
+ tpm_limit: 200000,
+ rpm_limit: 2000,
+ max_parallel_requests: 20,
+ budget_duration: "monthly",
+ model_max_budget: { "claude-3": 2000 },
+ },
+ },
+};
+
+describe("useTags", () => {
+ let queryClient: QueryClient;
+
+ beforeEach(() => {
+ queryClient = new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ },
+ },
+ });
+
+ // Reset all mocks
+ vi.clearAllMocks();
+
+ // Set default mock for useAuthorized (enabled state)
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: "test-user-id",
+ userRole: "Admin",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+ });
+
+ const wrapper = ({ children }: { children: ReactNode }) =>
+ React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+ it("should return tags data when query is successful", async () => {
+ // Mock successful API call
+ (tagListCall as any).mockResolvedValue(mockTags);
+
+ const { result } = renderHook(() => useTags(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+ expect(result.current.data).toBeUndefined();
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(mockTags);
+ expect(result.current.error).toBeNull();
+ expect(tagListCall).toHaveBeenCalledWith("test-access-token");
+ expect(tagListCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle error when tagListCall fails", async () => {
+ const errorMessage = "Failed to fetch tags";
+ const testError = new Error(errorMessage);
+
+ // Mock failed API call
+ (tagListCall as any).mockRejectedValue(testError);
+
+ const { result } = renderHook(() => useTags(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(testError);
+ expect(result.current.data).toBeUndefined();
+ expect(tagListCall).toHaveBeenCalledWith("test-access-token");
+ expect(tagListCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should not execute query when accessToken is missing", async () => {
+ // Mock missing accessToken
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userId: "test-user-id",
+ userRole: "Admin",
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useTags(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(tagListCall).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when userId is missing", async () => {
+ // Mock missing userId
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: null,
+ userRole: "Admin",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useTags(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(tagListCall).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when userRole is missing", async () => {
+ // Mock missing userRole
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: "test-user-id",
+ userRole: null,
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useTags(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(tagListCall).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when all auth values are missing", async () => {
+ // Mock all auth values missing
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userId: null,
+ userRole: null,
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useTags(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(tagListCall).not.toHaveBeenCalled();
+ });
+
+ it("should execute query when all auth values are present", async () => {
+ // Mock successful API call
+ (tagListCall as any).mockResolvedValue(mockTags);
+
+ // Ensure all auth values are present (already set in beforeEach)
+ const { result } = renderHook(() => useTags(), { wrapper });
+
+ // Wait for query to execute
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ expect(tagListCall).toHaveBeenCalledWith("test-access-token");
+ expect(tagListCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should return empty object when API returns empty data", async () => {
+ // Mock API returning empty object
+ (tagListCall as any).mockResolvedValue({});
+
+ const { result } = renderHook(() => useTags(), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual({});
+ expect(tagListCall).toHaveBeenCalledWith("test-access-token");
+ });
+
+ it("should handle network timeout error", async () => {
+ const timeoutError = new Error("Network timeout");
+
+ // Mock network timeout
+ (tagListCall as any).mockRejectedValue(timeoutError);
+
+ const { result } = renderHook(() => useTags(), { wrapper });
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(timeoutError);
+ expect(result.current.data).toBeUndefined();
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.ts
new file mode 100644
index 00000000000..8f82502a74c
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.ts
@@ -0,0 +1,16 @@
+import { useQuery, UseQueryResult } from "@tanstack/react-query";
+import { createQueryKeys } from "../common/queryKeysFactory";
+import { tagListCall } from "@/components/networking";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+import { TagListResponse } from "@/components/tag_management/types";
+
+const tagKeys = createQueryKeys("tags");
+
+export const useTags = (): UseQueryResult => {
+ const { accessToken, userId, userRole } = useAuthorized();
+ return useQuery({
+ queryKey: tagKeys.list({}),
+ queryFn: async () => await tagListCall(accessToken!),
+ enabled: Boolean(accessToken && userId && userRole),
+ });
+};
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts
new file mode 100644
index 00000000000..91ffbcfafa2
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts
@@ -0,0 +1,275 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import React, { ReactNode } from "react";
+import { useTeams } from "./useTeams";
+import { fetchTeams } from "@/app/(dashboard)/networking";
+import type { Team } from "@/components/key_team_helpers/key_list";
+
+// Mock the networking function
+vi.mock("@/app/(dashboard)/networking", () => ({
+ fetchTeams: vi.fn(),
+}));
+
+// Mock useAuthorized hook - we can override this in individual tests
+const mockUseAuthorized = vi.fn();
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+ default: () => mockUseAuthorized(),
+}));
+
+// Mock data
+const mockTeams: Team[] = [
+ {
+ team_id: "team-1",
+ team_alias: "Test Team 1",
+ models: ["gpt-3.5-turbo", "claude-3"],
+ max_budget: 100.0,
+ budget_duration: "monthly",
+ tpm_limit: 1000,
+ rpm_limit: 100,
+ organization_id: "org-1",
+ created_at: "2024-01-01T00:00:00Z",
+ keys: [],
+ members_with_roles: [],
+ },
+ {
+ team_id: "team-2",
+ team_alias: "Test Team 2",
+ models: ["gpt-4"],
+ max_budget: 200.0,
+ budget_duration: "monthly",
+ tpm_limit: 2000,
+ rpm_limit: 200,
+ organization_id: "org-1",
+ created_at: "2024-01-02T00:00:00Z",
+ keys: [],
+ members_with_roles: [],
+ },
+];
+
+describe("useTeams", () => {
+ let queryClient: QueryClient;
+
+ beforeEach(() => {
+ queryClient = new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ },
+ },
+ });
+
+ // Reset all mocks
+ vi.clearAllMocks();
+
+ // Set default mock for useAuthorized (enabled state)
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: "test-user-id",
+ userRole: "Admin",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+ });
+
+ const wrapper = ({ children }: { children: ReactNode }) =>
+ React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+ it("should return teams data when query is successful", async () => {
+ // Mock successful API call
+ (fetchTeams as any).mockResolvedValue(mockTeams);
+
+ const { result } = renderHook(() => useTeams(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+ expect(result.current.data).toBeUndefined();
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(mockTeams);
+ expect(result.current.error).toBeNull();
+ expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", null);
+ expect(fetchTeams).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle error when fetchTeams fails", async () => {
+ const errorMessage = "Failed to fetch teams";
+ const testError = new Error(errorMessage);
+
+ // Mock failed API call
+ (fetchTeams as any).mockRejectedValue(testError);
+
+ const { result } = renderHook(() => useTeams(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(testError);
+ expect(result.current.data).toBeUndefined();
+ expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", null);
+ expect(fetchTeams).toHaveBeenCalledTimes(1);
+ });
+
+ it("should not execute query when accessToken is missing", async () => {
+ // Mock missing accessToken
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userId: "test-user-id",
+ userRole: "Admin",
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useTeams(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(fetchTeams).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when accessToken is empty string", async () => {
+ // Mock empty string accessToken
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "",
+ userId: "test-user-id",
+ userRole: "Admin",
+ token: "",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useTeams(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(fetchTeams).not.toHaveBeenCalled();
+ });
+
+ it("should execute query when accessToken is present", async () => {
+ // Mock successful API call
+ (fetchTeams as any).mockResolvedValue(mockTeams);
+
+ // Ensure auth values are set (already done in beforeEach)
+ const { result } = renderHook(() => useTeams(), { wrapper });
+
+ // Wait for query to execute
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", null);
+ expect(fetchTeams).toHaveBeenCalledTimes(1);
+ });
+
+ it("should return empty teams array when API returns empty data", async () => {
+ // Mock API returning empty teams array
+ (fetchTeams as any).mockResolvedValue([]);
+
+ const { result } = renderHook(() => useTeams(), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual([]);
+ expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", null);
+ });
+
+ it("should handle network timeout error", async () => {
+ const timeoutError = new Error("Network timeout");
+
+ // Mock network timeout
+ (fetchTeams as any).mockRejectedValue(timeoutError);
+
+ const { result } = renderHook(() => useTeams(), { wrapper });
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(timeoutError);
+ expect(result.current.data).toBeUndefined();
+ });
+
+ it("should pass userId and userRole to fetchTeams", async () => {
+ // Mock successful API call
+ (fetchTeams as any).mockResolvedValue(mockTeams);
+
+ // Mock specific userId and userRole
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: "custom-user-id",
+ userRole: "member",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useTeams(), { wrapper });
+
+ // Wait for query to execute
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "custom-user-id", "member", null);
+ });
+
+ it("should handle null userId", async () => {
+ // Mock successful API call
+ (fetchTeams as any).mockResolvedValue(mockTeams);
+
+ // Mock null userId
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: null,
+ userRole: "Admin",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useTeams(), { wrapper });
+
+ // Wait for query to execute
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ expect(fetchTeams).toHaveBeenCalledWith("test-access-token", null, "Admin", null);
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts
new file mode 100644
index 00000000000..6429aeafb5a
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts
@@ -0,0 +1,169 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import React, { ReactNode } from "react";
+import { useUIConfig } from "./useUIConfig";
+import { getUiConfig, LiteLLMWellKnownUiConfig } from "@/components/networking";
+
+// Mock the networking function
+vi.mock("@/components/networking", () => ({
+ getUiConfig: vi.fn(),
+}));
+
+// Mock the queryKeysFactory - we'll mock the specific return value
+vi.mock("../common/queryKeysFactory", () => ({
+ createQueryKeys: vi.fn((resource: string) => ({
+ all: [resource],
+ lists: () => [resource, "list"],
+ list: (params?: any) => [resource, "list", { params }],
+ details: () => [resource, "detail"],
+ detail: (uid: string) => [resource, "detail", uid],
+ })),
+}));
+
+// Mock data
+const mockUIConfig: LiteLLMWellKnownUiConfig = {
+ server_root_path: "/api",
+ proxy_base_url: "https://proxy.example.com",
+ auto_redirect_to_sso: true,
+ admin_ui_disabled: false,
+};
+
+describe("useUIConfig", () => {
+ let queryClient: QueryClient;
+
+ beforeEach(() => {
+ queryClient = new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ },
+ },
+ });
+
+ // Reset all mocks
+ vi.clearAllMocks();
+ });
+
+ const wrapper = ({ children }: { children: ReactNode }) =>
+ React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+ it("should return UI config data when query is successful", async () => {
+ // Mock successful API call
+ (getUiConfig as any).mockResolvedValue(mockUIConfig);
+
+ const { result } = renderHook(() => useUIConfig(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+ expect(result.current.data).toBeUndefined();
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(mockUIConfig);
+ expect(result.current.error).toBeNull();
+ expect(getUiConfig).toHaveBeenCalledWith();
+ expect(getUiConfig).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle error when getUiConfig fails", async () => {
+ const errorMessage = "Failed to fetch UI config";
+ const testError = new Error(errorMessage);
+
+ // Mock failed API call
+ (getUiConfig as any).mockRejectedValue(testError);
+
+ const { result } = renderHook(() => useUIConfig(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(testError);
+ expect(result.current.data).toBeUndefined();
+ expect(getUiConfig).toHaveBeenCalledWith();
+ expect(getUiConfig).toHaveBeenCalledTimes(1);
+ });
+
+ it("should return different UI config data correctly", async () => {
+ const alternativeUIConfig: LiteLLMWellKnownUiConfig = {
+ server_root_path: "/v1",
+ proxy_base_url: null,
+ auto_redirect_to_sso: false,
+ admin_ui_disabled: true,
+ };
+
+ // Mock successful API call with different data
+ (getUiConfig as any).mockResolvedValue(alternativeUIConfig);
+
+ const { result } = renderHook(() => useUIConfig(), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(alternativeUIConfig);
+ expect(result.current.error).toBeNull();
+ });
+
+ it("should handle network timeout error", async () => {
+ const timeoutError = new Error("Network timeout");
+
+ // Mock network timeout
+ (getUiConfig as any).mockRejectedValue(timeoutError);
+
+ const { result } = renderHook(() => useUIConfig(), { wrapper });
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(timeoutError);
+ expect(result.current.data).toBeUndefined();
+ });
+
+ it("should handle malformed response error", async () => {
+ const malformedError = new Error("Invalid JSON response");
+
+ // Mock malformed response
+ (getUiConfig as any).mockRejectedValue(malformedError);
+
+ const { result } = renderHook(() => useUIConfig(), { wrapper });
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(malformedError);
+ expect(result.current.data).toBeUndefined();
+ });
+
+ it("should use correct query key structure", async () => {
+ // Mock successful API call
+ (getUiConfig as any).mockResolvedValue(mockUIConfig);
+
+ const { result } = renderHook(() => useUIConfig(), { wrapper });
+
+ // Wait for query to execute
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ // The query key should be generated by createQueryKeys("uiConfig").list({})
+ // Based on our mock, this should be ["uiConfig", "list", {}]
+ expect(getUiConfig).toHaveBeenCalledTimes(1);
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.test.ts
new file mode 100644
index 00000000000..785f003d2f8
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.test.ts
@@ -0,0 +1,185 @@
+import { getUiSettings } from "@/components/networking";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import { renderHook, waitFor } from "@testing-library/react";
+import React, { ReactNode } from "react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { useUISettings } from "./useUISettings";
+
+// Mock the networking function
+vi.mock("@/components/networking", () => ({
+ getUiSettings: vi.fn(),
+}));
+
+// Mock useAuthorized hook - we can override this in individual tests
+const mockUseAuthorized = vi.fn();
+vi.mock("../useAuthorized", () => ({
+ default: () => mockUseAuthorized(),
+}));
+
+// Mock data
+const mockUISettings: Record = {
+ theme: "dark",
+ language: "en",
+ notifications: true,
+ dashboard_layout: "compact",
+ api_keys_visible: false,
+};
+
+describe("useUISettings", () => {
+ let queryClient: QueryClient;
+
+ beforeEach(() => {
+ queryClient = new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ },
+ },
+ });
+
+ // Reset all mocks
+ vi.clearAllMocks();
+
+ // Set default mock for useAuthorized (enabled state)
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userRole: "Admin",
+ userId: "test-user-id",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+ });
+
+ const wrapper = ({ children }: { children: ReactNode }) =>
+ React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+ it("should return UI settings data when query is successful", async () => {
+ // Mock successful API call
+ (getUiSettings as any).mockResolvedValue(mockUISettings);
+
+ const { result } = renderHook(() => useUISettings(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+ expect(result.current.data).toBeUndefined();
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(mockUISettings);
+ expect(result.current.error).toBeNull();
+ expect(getUiSettings).toHaveBeenCalledWith("test-access-token");
+ expect(getUiSettings).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle error when getUiSettings fails", async () => {
+ const errorMessage = "Failed to fetch UI settings";
+ const testError = new Error(errorMessage);
+
+ // Mock failed API call
+ (getUiSettings as any).mockRejectedValue(testError);
+
+ const { result } = renderHook(() => useUISettings(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(testError);
+ expect(result.current.data).toBeUndefined();
+ expect(getUiSettings).toHaveBeenCalledWith("test-access-token");
+ expect(getUiSettings).toHaveBeenCalledTimes(1);
+ });
+
+ it("should not execute query when accessToken is missing", async () => {
+ // Mock missing accessToken
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userRole: "Admin",
+ userId: "test-user-id",
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useUISettings(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(getUiSettings).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when accessToken is empty string", async () => {
+ // Mock empty accessToken
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "",
+ userRole: "Admin",
+ userId: "test-user-id",
+ token: "",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useUISettings(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(getUiSettings).not.toHaveBeenCalled();
+ });
+
+ it("should return empty object when API returns empty settings", async () => {
+ // Mock API returning empty object
+ (getUiSettings as any).mockResolvedValue({});
+
+ const { result } = renderHook(() => useUISettings(), { wrapper });
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual({});
+ expect(getUiSettings).toHaveBeenCalledWith("test-access-token");
+ });
+
+ it("should handle network timeout error", async () => {
+ const timeoutError = new Error("Network timeout");
+
+ // Mock network timeout
+ (getUiSettings as any).mockRejectedValue(timeoutError);
+
+ const { result } = renderHook(() => useUISettings(), { wrapper });
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(timeoutError);
+ expect(result.current.data).toBeUndefined();
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts
index 823c0067b5c..46a0254d0db 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts
@@ -1,10 +1,12 @@
import { getUiSettings } from "@/components/networking";
import { useQuery } from "@tanstack/react-query";
import { createQueryKeys } from "../common/queryKeysFactory";
+import useAuthorized from "../useAuthorized";
const uiSettingsKeys = createQueryKeys("uiSettings");
-export const useUISettings = (accessToken: string) => {
+export const useUISettings = () => {
+ const { accessToken } = useAuthorized();
return useQuery>({
queryKey: uiSettingsKeys.list({}),
queryFn: async () => await getUiSettings(accessToken),
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowNewBadge.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowNewBadge.ts
new file mode 100644
index 00000000000..d0a618e27ba
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowNewBadge.ts
@@ -0,0 +1,35 @@
+// hooks/useDisableShowNewBadge.ts
+import { useSyncExternalStore } from "react";
+import { getLocalStorageItem } from "@/utils/localStorageUtils";
+import { LOCAL_STORAGE_EVENT } from "@/utils/localStorageUtils";
+
+function subscribe(callback: () => void) {
+ const onStorage = (e: StorageEvent) => {
+ if (e.key === "disableShowNewBadge") {
+ callback();
+ }
+ };
+
+ const onCustom = (e: Event) => {
+ const { key } = (e as CustomEvent).detail;
+ if (key === "disableShowNewBadge") {
+ callback();
+ }
+ };
+
+ window.addEventListener("storage", onStorage);
+ window.addEventListener(LOCAL_STORAGE_EVENT, onCustom);
+
+ return () => {
+ window.removeEventListener("storage", onStorage);
+ window.removeEventListener(LOCAL_STORAGE_EVENT, onCustom);
+ };
+}
+
+function getSnapshot() {
+ return getLocalStorageItem("disableShowNewBadge") === "true";
+}
+
+export function useDisableShowNewBadge() {
+ return useSyncExternalStore(subscribe, getSnapshot);
+}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts
new file mode 100644
index 00000000000..a392a940f98
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts
@@ -0,0 +1,253 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderHook, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import React, { ReactNode } from "react";
+import { useCurrentUser } from "./useCurrentUser";
+import { userInfoCall } from "@/components/networking";
+import type { UserInfo } from "@/components/view_users/types";
+
+// Mock the networking function
+vi.mock("@/components/networking", () => ({
+ userInfoCall: vi.fn(),
+}));
+
+// Mock the queryKeysFactory - we'll mock the specific return value
+vi.mock("../common/queryKeysFactory", () => ({
+ createQueryKeys: vi.fn((resource: string) => ({
+ all: [resource],
+ lists: () => [resource, "list"],
+ list: (params?: any) => [resource, "list", { params }],
+ details: () => [resource, "detail"],
+ detail: (uid: string) => [resource, "detail", uid],
+ })),
+}));
+
+// Mock useAuthorized hook - we can override this in individual tests
+const mockUseAuthorized = vi.fn();
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+ default: () => mockUseAuthorized(),
+}));
+
+// Mock data - response from userInfoCall should have user_info property
+const mockUserInfoResponse = {
+ user_info: {
+ user_id: "test-user-id",
+ user_email: "test@example.com",
+ user_alias: "Test User",
+ user_role: "Admin",
+ spend: 150.75,
+ max_budget: 1000.0,
+ key_count: 5,
+ created_at: "2024-01-01T00:00:00Z",
+ updated_at: "2024-01-01T00:00:00Z",
+ sso_user_id: null,
+ budget_duration: "monthly",
+ } as UserInfo,
+};
+
+describe("useCurrentUser", () => {
+ let queryClient: QueryClient;
+
+ beforeEach(() => {
+ queryClient = new QueryClient({
+ defaultOptions: {
+ queries: {
+ retry: false,
+ },
+ },
+ });
+
+ // Reset all mocks
+ vi.clearAllMocks();
+
+ // Set default mock for useAuthorized (enabled state)
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: "test-user-id",
+ userRole: "Admin",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+ });
+
+ const wrapper = ({ children }: { children: ReactNode }) =>
+ React.createElement(QueryClientProvider, { client: queryClient }, children);
+
+ it("should return user info data when query is successful", async () => {
+ // Mock successful API call
+ (userInfoCall as any).mockResolvedValue(mockUserInfoResponse);
+
+ const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+ expect(result.current.data).toBeUndefined();
+
+ // Wait for success
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isSuccess).toBe(true);
+ });
+
+ expect(result.current.data).toEqual(mockUserInfoResponse.user_info);
+ expect(result.current.error).toBeNull();
+ expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null);
+ expect(userInfoCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle error when userInfoCall fails", async () => {
+ const errorMessage = "Failed to fetch user info";
+ const testError = new Error(errorMessage);
+
+ // Mock failed API call
+ (userInfoCall as any).mockRejectedValue(testError);
+
+ const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+ // Initially loading
+ expect(result.current.isLoading).toBe(true);
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(testError);
+ expect(result.current.data).toBeUndefined();
+ expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null);
+ expect(userInfoCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should not execute query when accessToken is missing", async () => {
+ // Mock missing accessToken
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userId: "test-user-id",
+ userRole: "Admin",
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(userInfoCall).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when userId is missing", async () => {
+ // Mock missing userId
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: null,
+ userRole: "Admin",
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(userInfoCall).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when userRole is missing", async () => {
+ // Mock missing userRole
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "test-access-token",
+ userId: "test-user-id",
+ userRole: null,
+ token: "test-token",
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(userInfoCall).not.toHaveBeenCalled();
+ });
+
+ it("should not execute query when all auth values are missing", async () => {
+ // Mock all auth values missing
+ mockUseAuthorized.mockReturnValue({
+ accessToken: null,
+ userId: null,
+ userRole: null,
+ token: null,
+ userEmail: "test@example.com",
+ premiumUser: false,
+ disabledPersonalKeyCreation: null,
+ showSSOBanner: false,
+ });
+
+ const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+ // Query should not execute
+ expect(result.current.isLoading).toBe(false);
+ expect(result.current.data).toBeUndefined();
+ expect(result.current.isFetched).toBe(false);
+
+ // API should not be called
+ expect(userInfoCall).not.toHaveBeenCalled();
+ });
+
+ it("should execute query when all auth values are present", async () => {
+ // Mock successful API call
+ (userInfoCall as any).mockResolvedValue(mockUserInfoResponse);
+
+ // Ensure all auth values are present (already set in beforeEach)
+ const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+ // Wait for query to execute
+ await waitFor(() => {
+ expect(result.current.isLoading).toBe(false);
+ });
+
+ expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null);
+ expect(userInfoCall).toHaveBeenCalledTimes(1);
+ });
+
+ it("should handle network timeout error", async () => {
+ const timeoutError = new Error("Network timeout");
+
+ // Mock network timeout
+ (userInfoCall as any).mockRejectedValue(timeoutError);
+
+ const { result } = renderHook(() => useCurrentUser(), { wrapper });
+
+ // Wait for error
+ await waitFor(() => {
+ expect(result.current.isError).toBe(true);
+ });
+
+ expect(result.current.error).toEqual(timeoutError);
+ expect(result.current.data).toBeUndefined();
+ });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts
new file mode 100644
index 00000000000..f4028ada0dc
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts
@@ -0,0 +1,19 @@
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+import { UserInfo, userInfoCall } from "@/components/networking";
+import { useQuery, UseQueryResult } from "@tanstack/react-query";
+import { createQueryKeys } from "../common/queryKeysFactory";
+
+const userKeys = createQueryKeys("users");
+
+export const useCurrentUser = (): UseQueryResult => {
+ const { accessToken, userId, userRole } = useAuthorized();
+ return useQuery({
+ queryKey: userKeys.detail(userId!),
+ queryFn: async () => {
+ const data = await userInfoCall(accessToken!, userId!, userRole!, false, null, null);
+ console.log(`userInfo: ${JSON.stringify(data)}`);
+ return data.user_info;
+ },
+ enabled: Boolean(accessToken && userId && userRole),
+ });
+};
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/model-hub/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/model-hub/page.tsx
index 86967b660fd..c37a935976b 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/model-hub/page.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/model-hub/page.tsx
@@ -1,6 +1,6 @@
"use client";
-import ModelHubTable from "@/components/model_hub_table";
+import ModelHubTable from "@/components/AIHub/ModelHubTable";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
const ModelHubPage = () => {
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx
index 8dc7d1ff3d6..1e8eabaea2e 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx
@@ -9,34 +9,24 @@ vi.mock("@/components/networking", () => ({
credentialListCall: vi.fn().mockResolvedValue({ credentials: [] }),
modelInfoCall: vi.fn().mockResolvedValue({ data: [] }),
modelCostMap: vi.fn().mockResolvedValue({}),
- modelMetricsCall: vi.fn().mockResolvedValue({ data: [], all_api_bases: [] }),
- streamingModelMetricsCall: vi.fn().mockResolvedValue({ data: [], all_api_bases: [] }),
- modelExceptionsCall: vi.fn().mockResolvedValue({ data: [], exception_types: [] }),
- modelMetricsSlowResponsesCall: vi.fn().mockResolvedValue([]),
+ getPassThroughEndpointsCall: vi.fn().mockResolvedValue({ endpoints: {} }),
getCallbacksCall: vi.fn().mockResolvedValue({ router_settings: {} }),
setCallbacksCall: vi.fn().mockResolvedValue(undefined),
- modelSettingsCall: vi.fn().mockResolvedValue([]),
- adminGlobalActivityExceptions: vi.fn().mockResolvedValue({ sum_num_rate_limit_exceptions: 0, daily_data: [] }),
- adminGlobalActivityExceptionsPerDeployment: vi.fn().mockResolvedValue([]),
- allEndUsersCall: vi.fn().mockResolvedValue([]),
- latestHealthChecksCall: vi.fn().mockResolvedValue({ latest_health_checks: {} }),
- getPassThroughEndpointsCall: vi.fn().mockResolvedValue({ endpoints: {} }),
- getGuardrailsList: vi.fn().mockResolvedValue([]),
- tagListCall: vi.fn().mockResolvedValue([]),
- modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }),
- modelHubCall: vi.fn().mockResolvedValue({ data: [] }),
- getModelCostMapReloadStatus: vi.fn().mockResolvedValue({
- scheduled: false,
- interval_hours: null,
- last_run: null,
- next_run: null,
- }),
+ getUiSettings: vi.fn().mockResolvedValue({ values: {} }),
}));
vi.mock("@/app/(dashboard)/models-and-endpoints/components/ModelAnalyticsTab/ModelAnalyticsTab", () => ({
default: () => null,
}));
+vi.mock("@/components/add_model/add_auto_router_tab", () => ({
+ default: () => null,
+}));
+
+vi.mock("@/components/add_model/AddModelForm", () => ({
+ default: () => null,
+}));
+
vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({
default: () => ({
teams: [],
@@ -54,6 +44,16 @@ vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({
useUISettings: () => mockUseUISettings(),
}));
+const mockUseModelCostMap = vi.fn();
+vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({
+ useModelCostMap: () => mockUseModelCostMap(),
+}));
+
+const mockUseAuthorized = vi.fn();
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+ default: () => mockUseAuthorized(),
+}));
+
const createQueryClient = () =>
new QueryClient({
defaultOptions: { queries: { retry: false, gcTime: 0 } },
@@ -69,6 +69,17 @@ describe("ModelsAndEndpointsView", () => {
mockUseUISettings.mockReturnValue({
data: { values: {} },
});
+ mockUseModelCostMap.mockReturnValue({
+ data: {},
+ isLoading: false,
+ error: null,
+ });
+ mockUseAuthorized.mockReturnValue({
+ accessToken: "123",
+ token: "123",
+ userRole: "Admin",
+ userId: "123",
+ });
// eslint-disable-next-line @typescript-eslint/no-explicit-any
(global as any).ResizeObserver = class {
observe() {}
@@ -82,10 +93,7 @@ describe("ModelsAndEndpointsView", () => {
const { findByText } = render(
{}}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx
index bf001c62126..65418179595 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx
@@ -1,52 +1,35 @@
-import { useQueryClient } from "@tanstack/react-query";
-import { Col, Grid, Text } from "@tremor/react";
-import React, { useEffect, useRef, useState } from "react";
-
-import { handleAddModelSubmit } from "@/components/add_model/handle_add_model_submit";
-
import { useCredentials } from "@/app/(dashboard)/hooks/credentials/useCredentials";
+import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap";
import { useModelsInfo } from "@/app/(dashboard)/hooks/models/useModels";
-import { Team } from "@/components/key_team_helpers/key_list";
-import CredentialsPanel from "@/components/model_add/credentials";
-import {
- adminGlobalActivityExceptions,
- adminGlobalActivityExceptionsPerDeployment,
- allEndUsersCall,
- getCallbacksCall,
- modelCostMap,
- modelExceptionsCall,
- modelMetricsCall,
- modelMetricsSlowResponsesCall,
- modelSettingsCall,
- setCallbacksCall,
- streamingModelMetricsCall,
-} from "@/components/networking";
-import { Providers, getPlaceholder, getProviderModels } from "@/components/provider_info_helpers";
-import { getDisplayModelName } from "@/components/view_model/model_name_display";
-import { RefreshIcon } from "@heroicons/react/outline";
-import { DateRangePickerValue, Icon, Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react";
-import type { UploadProps } from "antd";
-import { Form, Typography } from "antd";
-import AddModelTab from "../../../components/add_model/add_model_tab";
-import ModelInfoView from "../../../components/model_info_view";
-import TeamInfoView from "../../../components/team/team_info";
-
+import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings";
import AllModelsTab from "@/app/(dashboard)/models-and-endpoints/components/AllModelsTab";
-import ModelAnalyticsTab from "@/app/(dashboard)/models-and-endpoints/components/ModelAnalyticsTab/ModelAnalyticsTab";
import ModelRetrySettingsTab from "@/app/(dashboard)/models-and-endpoints/components/ModelRetrySettingsTab";
import PriceDataManagementTab from "@/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab";
-import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings";
+import { handleAddModelSubmit } from "@/components/add_model/handle_add_model_submit";
+import { Team } from "@/components/key_team_helpers/key_list";
+import CredentialsPanel from "@/components/model_add/credentials";
+import { getCallbacksCall, setCallbacksCall } from "@/components/networking";
+import { Providers, getPlaceholder, getProviderModels } from "@/components/provider_info_helpers";
+import { getDisplayModelName } from "@/components/view_model/model_name_display";
+import { transformModelData } from "./utils/modelDataTransformer";
import { all_admin_roles, internalUserRoles, isProxyAdminRole, isUserTeamAdminForAnyTeam } from "@/utils/roles";
+import { RefreshIcon } from "@heroicons/react/outline";
+import { useQueryClient } from "@tanstack/react-query";
+import { Col, Grid, Icon, Tab, TabGroup, TabList, TabPanel, TabPanels, Text } from "@tremor/react";
+import type { UploadProps } from "antd";
+import { Form, Typography } from "antd";
+import React, { useEffect, useMemo, useState } from "react";
+import AddModelTab from "../../../components/add_model/add_model_tab";
import HealthCheckComponent from "../../../components/model_dashboard/HealthCheckComponent";
import ModelGroupAliasSettings from "../../../components/model_group_alias_settings";
+import ModelInfoView from "../../../components/model_info_view";
import NotificationsManager from "../../../components/molecules/notifications_manager";
import PassThroughSettings from "../../../components/pass_through_settings";
+import TeamInfoView from "../../../components/team/team_info";
+import useAuthorized from "../hooks/useAuthorized";
interface ModelDashboardProps {
- accessToken: string | null;
token: string | null;
- userRole: string | null;
- userID: string | null;
modelData: any;
keys: any[] | null;
setModelData: any;
@@ -62,100 +45,70 @@ interface GlobalRetryPolicyObject {
[retryPolicyKey: string]: number;
}
-interface GlobalExceptionActivityData {
- sum_num_rate_limit_exceptions: number;
- daily_data: { date: string; num_rate_limit_exceptions: number }[];
-}
-
-//["OpenAI", "Azure OpenAI", "Anthropic", "Gemini (Google AI Studio)", "Amazon Bedrock", "OpenAI-Compatible Endpoints (Groq, Together AI, Mistral AI, etc.)"]
-
-interface ProviderFields {
- field_name: string;
- field_type: string;
- field_description: string;
- field_value: string;
-}
-
-interface ProviderSettings {
- name: string;
- fields: ProviderFields[];
-}
-
-const ModelsAndEndpointsView: React.FC = ({
- accessToken,
- token,
- userRole,
- userID,
- modelData = { data: [] },
- keys,
- setModelData,
- premiumUser,
- teams,
-}) => {
+const ModelsAndEndpointsView: React.FC = ({ premiumUser, teams }) => {
+ const { accessToken, token, userRole, userId: userID } = useAuthorized();
const [addModelForm] = Form.useForm();
- const [modelMap, setModelMap] = useState(null);
const [lastRefreshed, setLastRefreshed] = useState("");
-
- const [providerModels, setProviderModels] = useState>([]); // Explicitly typing providerModels as a string array
-
- const [providerSettings, setProviderSettings] = useState([]);
+ const [providerModels, setProviderModels] = useState>([]);
const [selectedProvider, setSelectedProvider] = useState(Providers.Anthropic);
- const [editModalVisible, setEditModalVisible] = useState(false);
-
- const [selectedModel, setSelectedModel] = useState(null);
- const [availableModelGroups, setAvailableModelGroups] = useState>([]);
- const [availableModelAccessGroups, setAvailableModelAccessGroups] = useState>([]);
const [selectedModelGroup, setSelectedModelGroup] = useState(null);
- const [modelMetrics, setModelMetrics] = useState([]);
- const [modelMetricsCategories, setModelMetricsCategories] = useState([]);
- const [streamingModelMetrics, setStreamingModelMetrics] = useState([]);
- const [streamingModelMetricsCategories, setStreamingModelMetricsCategories] = useState([]);
- const [modelExceptions, setModelExceptions] = useState([]);
- const [allExceptions, setAllExceptions] = useState([]);
- const [slowResponsesData, setSlowResponsesData] = useState([]);
- const [dateValue, setDateValue] = useState({
- from: new Date(Date.now() - 7 * 24 * 60 * 60 * 1000),
- to: new Date(),
- });
const [modelGroupRetryPolicy, setModelGroupRetryPolicy] = useState(null);
const [globalRetryPolicy, setGlobalRetryPolicy] = useState(null);
const [defaultRetry, setDefaultRetry] = useState(0);
-
- const [globalExceptionData, setGlobalExceptionData] = useState(
- {} as GlobalExceptionActivityData,
- );
- const [globalExceptionPerDeployment, setGlobalExceptionPerDeployment] = useState([]);
-
- const [showAdvancedFilters, setShowAdvancedFilters] = useState(false);
- const [selectedAPIKey, setSelectedAPIKey] = useState(null);
- const [selectedCustomer, setSelectedCustomer] = useState(null);
-
- const [allEndUsers, setAllEndUsers] = useState([]);
-
- // Model Group Alias state
const [modelGroupAlias, setModelGroupAlias] = useState<{ [key: string]: string }>({});
-
- // Add state for advanced settings visibility
const [showAdvancedSettings, setShowAdvancedSettings] = useState(false);
-
- // Add these state variables
const [selectedModelId, setSelectedModelId] = useState(null);
- const [editModel, setEditModel] = useState(false);
-
const [selectedTeamId, setSelectedTeamId] = useState(null);
- const [selectedTeam, setSelectedTeam] = useState(null);
-
- const [isDropdownOpen, setIsDropdownOpen] = useState(false);
- const dropdownRef = useRef(null);
-
const [selectedTabIndex, setSelectedTabIndex] = useState(0);
const queryClient = useQueryClient();
const { data: modelDataResponse, isLoading: isLoadingModels, refetch: refetchModels } = useModelsInfo();
- const { data: credentialsResponse } = useCredentials();
+ const { data: modelCostMapData, isLoading: isLoadingModelCostMap } = useModelCostMap();
+ const { data: credentialsResponse, isLoading: isLoadingCredentials } = useCredentials();
const credentialsList = credentialsResponse?.credentials || [];
- const { data: uiSettings } = useUISettings(accessToken || "");
+ const { data: uiSettings, isLoading: isLoadingUISettings } = useUISettings();
+
+ const availableModelGroups = useMemo(() => {
+ if (!modelDataResponse?.data) return [];
+ const allModelGroups = new Set();
+ for (const model of modelDataResponse.data) {
+ allModelGroups.add(model.model_name);
+ }
+ return Array.from(allModelGroups).sort();
+ }, [modelDataResponse?.data]);
+
+ const availableModelAccessGroups = useMemo(() => {
+ if (!modelDataResponse?.data) return [];
+ const allModelAccessGroups = new Set();
+ for (const model of modelDataResponse.data) {
+ const modelInfo = model.model_info;
+ if (modelInfo?.access_groups) {
+ for (const group of modelInfo.access_groups) {
+ allModelAccessGroups.add(group);
+ }
+ }
+ }
+ return Array.from(allModelAccessGroups);
+ }, [modelDataResponse?.data]);
+
+ const allModelsOnProxy = useMemo(() => {
+ return modelDataResponse?.data?.map((model: any) => model.model_name);
+ }, [modelDataResponse?.data]);
+
+ const getProviderFromModel = (model: string) => {
+ if (modelCostMapData !== null && modelCostMapData !== undefined) {
+ if (typeof modelCostMapData == "object" && model in modelCostMapData) {
+ return modelCostMapData[model]["litellm_provider"];
+ }
+ }
+ return "openai";
+ };
+
+ const processedModelData = useMemo(() => {
+ if (!modelDataResponse?.data) return { data: [] };
+ return transformModelData(modelDataResponse, getProviderFromModel);
+ }, [modelDataResponse?.data, getProviderFromModel]);
const isProxyAdmin = userRole && isProxyAdminRole(userRole);
const isInternalUser = userRole && internalUserRoles.includes(userRole);
@@ -166,21 +119,10 @@ const ModelsAndEndpointsView: React.FC = ({
const shouldHideAddModelTab = !isProxyAdmin && (addModelDisabledForInternalUsers || !isUserTeamAdmin);
const setProviderModelsFn = (provider: Providers) => {
- const _providerModels = getProviderModels(provider, modelMap);
+ const _providerModels = getProviderModels(provider, modelCostMapData);
setProviderModels(_providerModels);
};
- useEffect(() => {
- const handleClickOutside = (event: MouseEvent) => {
- if (dropdownRef.current && !dropdownRef.current.contains(event.target as Node)) {
- setIsDropdownOpen(false);
- }
- };
-
- document.addEventListener("mousedown", handleClickOutside);
- return () => document.removeEventListener("mousedown", handleClickOutside);
- }, []);
-
const uploadProps: UploadProps = {
name: "file",
accept: ".json",
@@ -196,7 +138,6 @@ const ModelsAndEndpointsView: React.FC = ({
};
reader.readAsText(file);
}
- // Prevent upload
return false;
},
onChange(info) {
@@ -209,10 +150,8 @@ const ModelsAndEndpointsView: React.FC = ({
};
const handleRefreshClick = () => {
- // Update the 'lastRefreshed' state to the current date and time
const currentDate = new Date();
setLastRefreshed(currentDate.toLocaleString());
- // Invalidate and refetch models data using React Query
queryClient.invalidateQueries({ queryKey: ["models", "list"] });
refetchModels();
};
@@ -228,7 +167,6 @@ const ModelsAndEndpointsView: React.FC = ({
};
if (selectedModelGroup === "global") {
- // Only update global retry policy
if (globalRetryPolicy) {
payload.router_settings.retry_policy = globalRetryPolicy;
}
@@ -252,114 +190,6 @@ const ModelsAndEndpointsView: React.FC = ({
}
const fetchData = async () => {
try {
- setModelData(modelDataResponse);
- const _providerSettings = await modelSettingsCall(accessToken);
- if (_providerSettings) {
- setProviderSettings(_providerSettings);
- }
-
- // loop through modelDataResponse and get all`model_name` values
- let all_model_groups: Set = new Set();
- for (let i = 0; i < modelDataResponse.data.length; i++) {
- const model = modelDataResponse.data[i];
- all_model_groups.add(model.model_name);
- }
- let _array_model_groups = Array.from(all_model_groups);
- // sort _array_model_groups alphabetically
- _array_model_groups = _array_model_groups.sort();
-
- setAvailableModelGroups(_array_model_groups);
-
- let all_model_access_groups: Set = new Set();
- for (let i = 0; i < modelDataResponse.data.length; i++) {
- const model = modelDataResponse.data[i];
- let model_info: any | null = model.model_info;
- if (model_info) {
- let access_groups = model_info.access_groups;
- if (access_groups) {
- for (let j = 0; j < access_groups.length; j++) {
- all_model_access_groups.add(access_groups[j]);
- }
- }
- }
- }
-
- setAvailableModelAccessGroups(Array.from(all_model_access_groups));
-
- let _initial_model_group = "all";
- if (_array_model_groups.length > 0) {
- _initial_model_group = _array_model_groups[_array_model_groups.length - 1];
- }
-
- const modelMetricsResponse = await modelMetricsCall(
- accessToken,
- userID,
- userRole,
- _initial_model_group,
- dateValue.from?.toISOString(),
- dateValue.to?.toISOString(),
- selectedAPIKey?.token,
- selectedCustomer,
- );
-
- setModelMetrics(modelMetricsResponse.data);
- setModelMetricsCategories(modelMetricsResponse.all_api_bases);
-
- const streamingModelMetricsResponse = await streamingModelMetricsCall(
- accessToken,
- _initial_model_group,
- dateValue.from?.toISOString(),
- dateValue.to?.toISOString(),
- );
-
- // Assuming modelMetricsResponse now contains the metric data for the specified model group
- setStreamingModelMetrics(streamingModelMetricsResponse.data);
- setStreamingModelMetricsCategories(streamingModelMetricsResponse.all_api_bases);
-
- const modelExceptionsResponse = await modelExceptionsCall(
- accessToken,
- userID,
- userRole,
- _initial_model_group,
- dateValue.from?.toISOString(),
- dateValue.to?.toISOString(),
- selectedAPIKey?.token,
- selectedCustomer,
- );
- setModelExceptions(modelExceptionsResponse.data);
- setAllExceptions(modelExceptionsResponse.exception_types);
-
- const slowResponses = await modelMetricsSlowResponsesCall(
- accessToken,
- userID,
- userRole,
- _initial_model_group,
- dateValue.from?.toISOString(),
- dateValue.to?.toISOString(),
- selectedAPIKey?.token,
- selectedCustomer,
- );
-
- const dailyExceptions = await adminGlobalActivityExceptions(
- accessToken,
- dateValue.from?.toISOString().split("T")[0],
- dateValue.to?.toISOString().split("T")[0],
- _initial_model_group,
- );
-
- setGlobalExceptionData(dailyExceptions);
-
- const dailyExceptionsPerDeplyment = await adminGlobalActivityExceptionsPerDeployment(
- accessToken,
- dateValue.from?.toISOString().split("T")[0],
- dateValue.to?.toISOString().split("T")[0],
- _initial_model_group,
- );
-
- setGlobalExceptionPerDeployment(dailyExceptionsPerDeplyment);
- setSlowResponsesData(slowResponses);
- let all_end_users_data = await allEndUsersCall(accessToken);
- setAllEndUsers(all_end_users_data?.map((u: any) => u.user_id));
const routerSettingsInfo = await getCallbacksCall(accessToken, userID, userRole);
let router_settings = routerSettingsInfo.router_settings;
@@ -370,7 +200,6 @@ const ModelsAndEndpointsView: React.FC = ({
setGlobalRetryPolicy(router_settings.retry_policy);
setDefaultRetry(default_retries);
- // Set model group alias
const model_group_alias = router_settings.model_group_alias || {};
setModelGroupAlias(model_group_alias);
} catch (error) {
@@ -381,110 +210,9 @@ const ModelsAndEndpointsView: React.FC = ({
if (accessToken && token && userRole && userID && modelDataResponse) {
fetchData();
}
-
- const fetchModelMap = async () => {
- const data = await modelCostMap();
- console.log(`received model cost map data: ${Object.keys(data)}`);
- setModelMap(data);
- };
- if (modelMap == null) {
- fetchModelMap();
- }
}, [accessToken, token, userRole, userID, modelDataResponse]);
- if (!modelData || isLoadingModels) {
- return
+
+ No SSO Configuration Found
+
+ Configure Single Sign-On (SSO) to enable seamless authentication for your team members using your identity
+ provider.
+
+
+
+ {budgetList
+ .slice() // Creates a shallow copy to avoid mutating the original array
+ .sort((a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime()) // Sort by updated_at in descending order
+ .map((value: budgetItem, index: number) => (
+
+ {value.budget_id}
+ {value.max_budget ? value.max_budget : "n/a"}
+ {value.tpm_limit ? value.tpm_limit : "n/a"}
+ {value.rpm_limit ? value.rpm_limit : "n/a"}
+ handleEditCall(value)}
+ dataTestId="edit-budget-button"
+ />
+ handleDeleteClick(value)}
+ dataTestId="delete-budget-button"
+ />
+
+ ))}
+
+
+
+
+
+
+
+
+ How to use budget id
+
+
+ Assign Budget to Customer
+ Test it (Curl)
+ Test it (OpenAI SDK)
+
+
+
+ {CREATE_END_USER_CURL_COMMAND}
+
+
+ {CHAT_COMPLETIONS_CURL_COMMAND}
+
+
+ {OPENAI_SDK_PYTHON_CODE}
+
+
+
+