Merge branch 'main' into litellm_oss_staging_03_11_2026

This commit is contained in:
Cesar Garcia 2026-03-12 10:43:08 -03:00 committed by GitHub
commit 6bd7cd7573
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
125 changed files with 7141 additions and 857 deletions

View file

@ -4337,7 +4337,8 @@ jobs:
name: Check for expected error name: Check for expected error
command: | command: |
if grep -q "Error: P1001: Can't reach database server at" docker_output.log && \ if grep -q "Error: P1001: Can't reach database server at" docker_output.log && \
grep -q "ERROR: Application startup failed. Exiting." docker_output.log; then (grep -q "Database setup failed after multiple retries" docker_output.log || \
grep -q "ERROR: Application startup failed. Exiting." docker_output.log); then
echo "Expected error found. Test passed." echo "Expected error found. Test passed."
else else
echo "Expected error not found. Test failed." echo "Expected error not found. Test failed."

View file

@ -41,3 +41,39 @@ jobs:
git push origin $BRANCH_NAME git push origin $BRANCH_NAME
echo "Successfully created and pushed branch: $BRANCH_NAME" echo "Successfully created and pushed branch: $BRANCH_NAME"
fi fi
create-internal-dev-branch:
runs-on: ubuntu-latest
steps:
- name: Checkout repository
uses: actions/checkout@v3
with:
fetch-depth: 0
- name: Create internal dev branch
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
# Configure Git user
git config user.name "github-actions[bot]"
git config user.email "github-actions[bot]@users.noreply.github.com"
# Generate branch name with MM_DD_YYYY format
BRANCH_NAME="litellm_internal_dev_$(date +'%m_%d_%Y')"
echo "Creating branch: $BRANCH_NAME"
# Fetch all branches
git fetch --all
# Check if the branch already exists
if git show-ref --verify --quiet refs/remotes/origin/$BRANCH_NAME; then
echo "Branch $BRANCH_NAME already exists. Skipping creation."
else
echo "Creating new branch: $BRANCH_NAME"
# Create the new branch from main
git checkout -b $BRANCH_NAME origin/main
# Push the new branch
git push origin $BRANCH_NAME
echo "Successfully created and pushed branch: $BRANCH_NAME"
fi

View file

@ -102,6 +102,22 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components:
### UI / Backend Consistency ### UI / Backend Consistency
- When wiring a new UI entity type to an existing backend endpoint, verify the backend API contract (single value vs. array, required vs. optional params) and ensure the UI controls match — e.g., use a single-select dropdown when the backend accepts a single value, not a multi-select - When wiring a new UI entity type to an existing backend endpoint, verify the backend API contract (single value vs. array, required vs. optional params) and ensure the UI controls match — e.g., use a single-select dropdown when the backend accepts a single value, not a multi-select
### MCP OAuth / OpenAPI Transport Mapping
- `TRANSPORT.OPENAPI` is a UI-only concept. The backend only accepts `"http"`, `"sse"`, or `"stdio"`. Always map it to `"http"` before any API call (including pre-OAuth temp-session calls).
- FastAPI validation errors return `detail` as an array of `{loc, msg, type}` objects. Error extractors must handle: array (map `.msg`), string, nested `{error: string}`, and fallback.
- When an MCP server already has `authorization_url` stored, skip OAuth discovery (`_discovery_metadata`) — the server URL for OpenAPI MCPs is the spec file, not the API base, and fetching it causes timeouts.
- `client_id` should be optional in the `/authorize` endpoint — if the server has a stored `client_id` in credentials, use that. Never require callers to re-supply it.
### MCP Credential Storage
- OAuth credentials and BYOK credentials share the `litellm_mcpusercredentials` table, distinguished by a `"type"` field in the JSON payload (`"oauth2"` vs plain string).
- When deleting OAuth credentials, check type before deleting to avoid accidentally deleting a BYOK credential for the same `(user_id, server_id)` pair.
- Always pass the raw `expires_at` timestamp to the client — never set it to `None` for expired credentials. Let the frontend compute the "Expired" display state from the timestamp.
- Use `RecordNotFoundError` (not bare `except Exception`) when catching "already deleted" in credential delete endpoints.
### Browser Storage Safety (UI)
- Never write LiteLLM access tokens or API keys to `localStorage` — use `sessionStorage` only. `localStorage` survives browser close and is readable by any injected script (XSS).
- Shared utility functions (e.g. `extractErrorMessage`) belong in `src/utils/` — never define them inline in hooks or duplicate them across files.
### Database Migrations ### Database Migrations
- Prisma handles schema migrations - Prisma handles schema migrations
- Migration files auto-generated with `prisma migrate dev` - Migration files auto-generated with `prisma migrate dev`

View file

@ -944,7 +944,7 @@ router_settings:
| QDRANT_URL | Connection URL for Qdrant database | QDRANT_URL | Connection URL for Qdrant database
| QDRANT_VECTOR_SIZE | Vector size for Qdrant operations. Default is 1536 | QDRANT_VECTOR_SIZE | Vector size for Qdrant operations. Default is 1536
| REDIS_CONNECTION_POOL_TIMEOUT | Timeout in seconds for Redis connection pool. Default is 5 | REDIS_CONNECTION_POOL_TIMEOUT | Timeout in seconds for Redis connection pool. Default is 5
| REDIS_CLUSTER_NODES | JSON-formatted list of Redis cluster startup nodes for Redis Cluster mode. Example: '[{"host": "node1", "port": 6379}]' | REDIS_CLUSTER_NODES | JSON-formatted list of Redis cluster startup nodes for Redis Cluster mode. Example: `[{"host": "node1", "port": 6379}]`
| REDIS_HOST | Hostname for Redis server | REDIS_HOST | Hostname for Redis server
| REDIS_PASSWORD | Password for Redis service | REDIS_PASSWORD | Password for Redis service
| REDIS_PORT | Port number for Redis server | REDIS_PORT | Port number for Redis server

View file

@ -309,6 +309,10 @@ Response:
</TabItem> </TabItem>
</Tabs> </Tabs>
## Policy Flow Builder
For conditional execution (e.g., run a second guardrail only if the first fails), use the [Policy Flow Builder](./policy_flow_builder) to define pipelines with per-step pass/fail actions.
## Config Reference ## Config Reference
### `policies` ### `policies`
@ -323,6 +327,7 @@ policies:
remove: [...] remove: [...]
condition: condition:
model: ... model: ...
pipeline: ... # optional; see Policy Flow Builder
``` ```
| Field | Type | Description | | Field | Type | Description |
@ -332,6 +337,7 @@ policies:
| `guardrails.add` | `list[string]` | Guardrails to enable. | | `guardrails.add` | `list[string]` | Guardrails to enable. |
| `guardrails.remove` | `list[string]` | Guardrails to disable (useful with inheritance). | | `guardrails.remove` | `list[string]` | Guardrails to disable (useful with inheritance). |
| `condition.model` | `string` or `list[string]` | Optional. Only apply when model matches. Supports regex. | | `condition.model` | `string` or `list[string]` | Optional. Only apply when model matches. Supports regex. |
| `pipeline` | `object` | Optional. Ordered guardrail execution with per-step actions. See [Policy Flow Builder](./policy_flow_builder). |
### `policy_attachments` ### `policy_attachments`

View file

@ -0,0 +1,219 @@
# Policy Flow Builder
The Policy Flow Builder lets you design guardrail pipelines with **conditional execution**. Instead of running guardrails independently, you chain them into ordered steps and control what happens when each guardrail passes or fails.
Two powerful patterns it enables: **guardrail fallbacks** (try a different guardrail when one fails) and **retrying the same guardrail** (run the same guardrail again if it fails, e.g. to handle transient errors).
## When to use the Flow Builder
| Approach | Use case |
|----------|----------|
| **Simple policy** (`guardrails.add`) | All guardrails run in parallel; any failure blocks the request. |
| **Flow Builder** (pipeline) | Guardrails run in sequence; you choose actions per step (next, block, allow, custom response). |
Use the Flow Builder when you need:
- **Guardrail fallbacks** — use `on_fail: next` to try a different guardrail when one fails (e.g., fast filter → stricter filter)
- **Retrying the same guardrail** — add the same guardrail as multiple steps; if it fails, `on_fail: next` moves to the next step, which can be the same guardrail again (useful for transient API errors or rate limits)
- **Conditional routing** — e.g., if a fast guardrail fails, run a more advanced one instead of blocking immediately
- **Custom responses** — return a specific message when a guardrail fails instead of a generic block
- **Data chaining** — pass modified data (e.g., PII-masked content) from one step to the next
- **Fine-grained control** — different actions on pass vs. fail per step
## Concepts
### Pipeline
A pipeline has:
- **Mode**: `pre_call` (before the LLM) or `post_call` (after the LLM)
- **Steps**: Ordered list of guardrail steps
### Step actions
Each step defines what happens when the guardrail **passes** and when it **fails**:
| Action | Description |
|--------|-------------|
| **Next Step** | Continue to the next guardrail in the pipeline |
| **Allow** | Stop the pipeline and allow the request to proceed |
| **Block** | Stop the pipeline and block the request |
| **Custom Response** | Return a custom message instead of the default block |
### Step options
| Field | Type | Description |
|-------|------|--------------|
| `guardrail` | `string` | Name of the guardrail to run |
| `on_pass` | `string` | Action when guardrail passes: `next`, `allow`, `block`, `modify_response` |
| `on_fail` | `string` | Action when guardrail fails: `next`, `allow`, `block`, `modify_response` |
| `pass_data` | `boolean` | Forward modified request data (e.g., PII-masked) to the next step |
| `modify_response_message` | `string` | Custom message when using `modify_response` action |
## Using the Flow Builder (UI)
1. Go to **Policies** in the LiteLLM Admin UI
2. Click **+ Create New Policy** or **Edit** on an existing policy
3. Select **Flow Builder** (instead of the simple form)
4. Design your flow:
- **Trigger** — Incoming LLM request (runs when the policy matches)
- **Steps** — Add guardrails, set ON PASS and ON FAIL actions per step
- **End** — Request proceeds to the LLM
5. Use the **+** between steps to insert new steps
6. Use the **Test** panel to run sample messages through the pipeline before saving
7. Click **Save** to create or update the policy
## Config (YAML)
Define a pipeline in your policy config:
```yaml showLineNumbers title="config.yaml"
guardrails:
- guardrail_name: pii_masking
litellm_params:
guardrail: presidio
mode: pre_call
- guardrail_name: prompt_injection
litellm_params:
guardrail: lakera
mode: pre_call
policies:
my-pipeline-policy:
description: "PII mask first, then check for prompt injection"
guardrails:
add:
- pii_masking
- prompt_injection
pipeline:
mode: pre_call
steps:
- guardrail: pii_masking
on_pass: next
on_fail: block
pass_data: true
- guardrail: prompt_injection
on_pass: allow
on_fail: block
policy_attachments:
- policy: my-pipeline-policy
scope: "*"
```
## Fallbacks and retries
### Guardrail fallbacks
Use `on_fail: next` to fall back to another guardrail when one fails. Run a lightweight guardrail first; if it fails, escalate to a stricter or different provider:
```yaml
policies:
fallback-policy:
guardrails:
add:
- fast_content_filter
- strict_content_filter
pipeline:
mode: pre_call
steps:
- guardrail: fast_content_filter
on_pass: allow
on_fail: next
- guardrail: strict_content_filter
on_pass: allow
on_fail: block
```
If `fast_content_filter` passes → allow. If it fails → run `strict_content_filter`; pass → allow, fail → block.
### Retrying the same guardrail
Add the same guardrail as multiple steps to retry on failure. Useful for transient errors (API timeouts, rate limits):
```yaml
policies:
retry-policy:
guardrails:
add:
- lakera_prompt_injection
pipeline:
mode: pre_call
steps:
- guardrail: lakera_prompt_injection
on_pass: allow
on_fail: next
- guardrail: lakera_prompt_injection
on_pass: allow
on_fail: block
```
First attempt passes → allow. First attempt fails → retry the same guardrail; second pass → allow, second fail → block.
## Example: Custom response on fail
Return a branded message instead of a generic block:
```yaml
policies:
branded-block-policy:
guardrails:
add:
- pii_detector
pipeline:
mode: pre_call
steps:
- guardrail: pii_detector
on_pass: allow
on_fail: modify_response
modify_response_message: "Your message contains sensitive information. Please remove PII and try again."
```
## Test a pipeline (API)
Test a pipeline with sample messages before attaching it:
```bash
curl -X POST "http://localhost:4000/policies/test-pipeline" \
-H "Authorization: Bearer <your_api_key>" \
-H "Content-Type: application/json" \
-d '{
"pipeline": {
"mode": "pre_call",
"steps": [
{
"guardrail": "pii_masking",
"on_pass": "next",
"on_fail": "block",
"pass_data": true
},
{
"guardrail": "prompt_injection",
"on_pass": "allow",
"on_fail": "block"
}
]
},
"test_messages": [
{"role": "user", "content": "What is 2+2?"},
{"role": "user", "content": "My SSN is 123-45-6789"}
]
}'
```
Response includes per-step outcomes (pass/fail/error), actions taken, and timing.
## Pipeline vs simple policy
When a policy has a `pipeline`, the pipeline defines execution order and actions. The `guardrails.add` list must include all guardrails used in the pipeline steps.
| Policy type | Execution |
|-------------|-----------|
| Simple (`guardrails.add` only) | All guardrails run; any failure blocks |
| Pipeline (`pipeline` present) | Steps run in order; actions control flow |
## Related docs
- [Guardrail Policies](./guardrail_policies) — Policy basics, attachments, inheritance
- [Policy Templates](./policy_templates) — Pre-built policy templates

View file

@ -0,0 +1,143 @@
import Image from '@theme/IdealImage';
# Retool Assist
This guide walks you through connecting [Retool Assist](https://docs.retool.com/apps/guides/assist/) to LiteLLM Proxy. Retool Assist uses AI to generate and edit apps from within the Retool app IDE. Using LiteLLM with Retool Assist allows you to:
- Access 100+ LLMs through Retool Assist
- Track spend and usage, set budget limits per virtual key
- Control which models Retool Assist can access
- Use your own LLM providers via a unified OpenAI-compatible API
<div style={{ maxWidth: '100%', overflow: 'hidden', paddingBottom: '59.52%', position: 'relative', height: 0 }}>
<iframe
style={{ position: 'absolute', top: 0, left: 0, width: '100%', height: '100%', maxWidth: '840px' }}
src="https://www.youtube.com/embed/aN-Iua5dHGg"
frameborder="0"
webkitallowfullscreen
mozallowfullscreen
allowfullscreen
></iframe>
</div>
---
:::info
**Hosted Retool requires a public URL.** Retool Cloud runs on Retool's servers, so `localhost` will not work. You must expose your LiteLLM proxy via ngrok, Cloudflare Tunnel, or by deploying to a cloud provider.
:::
## Quick Reference
| Setting | Value |
|---------|-------|
| Provider Schema | OpenAI |
| Base URL | Your ngrok URL (e.g. `https://abc123.ngrok-free.app`) or deployed proxy URL |
| API Key | Your LiteLLM Virtual Key |
| Model | Public model name from LiteLLM (e.g. `openai/gpt-4o-mini`, `openai/gpt-5.2-2025-12-11`) |
---
## Prerequisites
- LiteLLM Proxy running locally or deployed
- [ngrok](https://ngrok.com/download) (or similar tunnel) for local development with hosted Retool
- A [Retool](https://retool.com) account (Cloud or self-hosted)
## 1. Start LiteLLM Proxy
Set up LiteLLM Proxy following the [Getting Started Guide](https://docs.litellm.ai/docs/proxy/docker_quick_start). Ensure your proxy is running on port 4000.
## 2. Expose LiteLLM with a Public URL
<Image img={require('../../img/ngrok_public_url.gif')} />
Retool Cloud runs on Retool's servers. You must expose your local LiteLLM proxy with a public URL.
### Using ngrok
- Install [ngrok](https://ngrok.com/download)
- In a separate terminal, run:
```bash
ngrok http 4000
```
- Copy the generated HTTPS URL (e.g. `https://abc123.ngrok-free.app`). This is your **Base URL** for Retool.
### Alternative
If you deploy LiteLLM to Railway, Render, Fly.io, or another cloud provider, use that public URL as your Base URL. See the [Deploy guide](https://docs.litellm.ai/docs/proxy/deploy) for details.
## 3. Generate a Virtual Key
<Image img={require('../../img/litellm_virtual_key.gif')} />
Create a virtual key that Retool Assist will use to authenticate with LiteLLM. The key must have access to the models you want to use (e.g. `openai/*` for all OpenAI models).
### Via LiteLLM UI
- Navigate to [http://localhost:4000/ui](http://localhost:4000/ui)
- Go to **Virtual Keys****+ Create New Key**
- Select the models you need (or `openai/*` for all OpenAI models)
- Copy the key
## 4. Add LiteLLM as a Custom Provider in Retool
Inside your Retool dashboard, configure LiteLLM as a custom AI resource:
<Image img={require('../../img/retool_resource_setup.gif')} />
1. Go to **Resources**
2. Under the **AI** category, select **Custom Provider**
3. Fill in the form:
- **Name:** `LiteLLM`
- **Description:** (optional) e.g. `LiteLLM Proxy - 100+ LLMs`
- **Provider Schema:** `OpenAI`
- **Base URL:** Your ngrok-generated URL (e.g. `https://abc123.ngrok-free.app`) or deployed proxy URL—do not add `/v1` unless Retool requires it
- **API Key:** Your LiteLLM virtual key from Step 3
4. **Add model names** from your LiteLLM proxy (e.g. `openai/gpt-4o-mini`, `openai/gpt-5.2-2025-12-11`).
5. Click **Create Resource**
<Image img={require('../../img/retool_llm_setup.gif')} />
## 5. Test the Connection
<Image img={require('../../img/retool_litellm_connection.gif')} />
- Open an app in Retool and enable **Assist** (if not already enabled in your organization)
- Use Assist to generate or edit app elements, it will route requests through LiteLLM
- Use the code option from the Sidebar to add a resource query, select the LiteLLM resource, and run it to test the setup.
- Check the LiteLLM **Logs** section to verify requests and track usage
<Image img={require('../../img/retool_litellm_logs.gif')} />
---
## Troubleshooting
### 401 Unauthorized
- Ensure the **API Key** in Retool matches your LiteLLM virtual key exactly
- Verify the key is not expired or blocked in LiteLLM
### 401 "key not allowed to access model"
Your virtual key is restricted to specific models. Generate a new key with `openai/*` or include the model you need (e.g. `openai/gpt-5.2-2025-12-11`) in the key's allowed models list.
### 500 "api_key client option must be set"
LiteLLM could not use your OpenAI API key to call the provider. Ensure `OPENAI_API_KEY` is set in your LiteLLM environment (e.g. in `.env` or `docker-compose.yml`) when using `openai/*` models.
### localhost does not work
Retool Cloud cannot reach `localhost` it points to Retool's servers. Use ngrok or deploy LiteLLM to a public URL.
---
## Additional Resources
- [Virtual Keys](https://docs.litellm.ai/docs/proxy/virtual_keys) Create and manage API keys
- [Deploy LiteLLM](https://docs.litellm.ai/docs/proxy/deploy) Production deployment options
- [Retool Assist Documentation](https://docs.retool.com/apps/guides/assist/) Configure Assist and prompting guides

Binary file not shown.

After

Width:  |  Height:  |  Size: 16 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 22 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 40 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 18 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 14 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 6.9 MiB

View file

@ -1,5 +1,5 @@
--- ---
title: "[Preview] v1.82.0 - Realtime Guardrails, Projects Management, and 10+ Performance Optimizations" title: "v1.82.0 - Realtime Guardrails, Projects Management, and 10+ Performance Optimizations"
slug: "v1-82-0" slug: "v1-82-0"
date: 2026-02-28T00:00:00 date: 2026-02-28T00:00:00
authors: authors:
@ -26,7 +26,7 @@ import TabItem from '@theme/TabItem';
docker run \ docker run \
-e STORE_MODEL_IN_DB=True \ -e STORE_MODEL_IN_DB=True \
-p 4000:4000 \ -p 4000:4000 \
ghcr.io/berriai/litellm:main-1.82.0 ghcr.io/berriai/litellm:main-1.82.0-stable
``` ```
</TabItem> </TabItem>

View file

@ -100,6 +100,7 @@ const sidebars = {
label: "Policies", label: "Policies",
items: [ items: [
"proxy/guardrails/guardrail_policies", "proxy/guardrails/guardrail_policies",
"proxy/guardrails/policy_flow_builder",
"proxy/guardrails/policy_templates", "proxy/guardrails/policy_templates",
"proxy/guardrails/policy_tags", "proxy/guardrails/policy_tags",
], ],
@ -172,7 +173,8 @@ const sidebars = {
"tutorials/litellm_gemini_cli", "tutorials/litellm_gemini_cli",
"tutorials/google_genai_sdk", "tutorials/google_genai_sdk",
"tutorials/litellm_qwen_code_cli", "tutorials/litellm_qwen_code_cli",
"tutorials/openai_codex" "tutorials/openai_codex",
"tutorials/retool_assist"
] ]
}, },
{ {

Binary file not shown.

View file

@ -1,13 +0,0 @@
-- SkipTransactionBlock
-- Drop invalid indexes left behind by failed CONCURRENTLY builds
DROP INDEX CONCURRENTLY IF EXISTS "LiteLLM_VerificationToken_key_alias_idx";
-- CreateIndex
CREATE INDEX CONCURRENTLY "LiteLLM_VerificationToken_key_alias_idx" ON "LiteLLM_VerificationToken"("key_alias");
-- Drop invalid indexes left behind by failed CONCURRENTLY builds
DROP INDEX CONCURRENTLY IF EXISTS "LiteLLM_SpendLogs_user_startTime_idx";
-- CreateIndex
CREATE INDEX CONCURRENTLY "LiteLLM_SpendLogs_user_startTime_idx" ON "LiteLLM_SpendLogs"("user", "startTime");

View file

@ -0,0 +1,11 @@
-- DropIndex
DROP INDEX "LiteLLM_MCPServerTable_approval_status_idx";
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN "approval_status",
DROP COLUMN "review_notes",
DROP COLUMN "reviewed_at",
DROP COLUMN "source_url",
DROP COLUMN "submitted_at",
DROP COLUMN "submitted_by";

View file

@ -267,6 +267,7 @@ model LiteLLM_ObjectPermissionTable {
vector_stores String[] @default([]) vector_stores String[] @default([])
agents String[] @default([]) agents String[] @default([])
agent_access_groups String[] @default([]) agent_access_groups String[] @default([])
models String[] @default([])
blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission
teams LiteLLM_TeamTable[] teams LiteLLM_TeamTable[]
projects LiteLLM_ProjectTable[] projects LiteLLM_ProjectTable[]
@ -315,6 +316,11 @@ model LiteLLM_MCPServerTable {
is_byok Boolean @default(false) is_byok Boolean @default(false)
byok_description String[] @default([]) byok_description String[] @default([])
byok_api_key_help_url String? byok_api_key_help_url String?
approval_status String @default("approved")
submitted_by String?
submitted_at DateTime?
reviewed_at DateTime?
review_notes String?
} }
// Per-user BYOK credentials for MCP servers // Per-user BYOK credentials for MCP servers
@ -388,9 +394,6 @@ model LiteLLM_VerificationToken {
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3
@@index([budget_reset_at, expires]) @@index([budget_reset_at, expires])
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (...) ORDER BY "public"."LiteLLM_VerificationToken"."key_alias" ASC
@@index([key_alias])
} }
model LiteLLM_JWTKeyMapping { model LiteLLM_JWTKeyMapping {
@ -556,9 +559,6 @@ model LiteLLM_SpendLogs {
@@index([startTime, request_id]) @@index([startTime, request_id])
@@index([end_user]) @@index([end_user])
@@index([session_id]) @@index([session_id])
// SELECT ... FROM "LiteLLM_SpendLogs" WHERE ("startTime" >= $1 AND "startTime" <= $2 AND "user" = $3) GROUP BY ...
@@index([user, startTime])
} }
// View spend, model, api_key per request // View spend, model, api_key per request

View file

@ -1,6 +1,6 @@
[tool.poetry] [tool.poetry]
name = "litellm-proxy-extras" name = "litellm-proxy-extras"
version = "0.4.53" version = "0.4.54"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
authors = ["BerriAI"] authors = ["BerriAI"]
readme = "README.md" readme = "README.md"
@ -22,7 +22,7 @@ requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api" build-backend = "poetry.core.masonry.api"
[tool.commitizen] [tool.commitizen]
version = "0.4.53" version = "0.4.54"
version_files = [ version_files = [
"pyproject.toml:version", "pyproject.toml:version",
"../requirements.txt:litellm-proxy-extras==", "../requirements.txt:litellm-proxy-extras==",

View file

@ -1,6 +1,6 @@
# What is this? # What is this?
## Helper utilities ## Helper utilities
from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union, get_args from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union
import httpx import httpx

View file

@ -1198,6 +1198,9 @@ class AmazonConverseConfig(BaseConfig):
+ supported_config_params + supported_config_params
) )
inference_params.pop("json_mode", None) # used for handling json_schema inference_params.pop("json_mode", None) # used for handling json_schema
# Anthropic-only key. Bedrock expects `outputConfig` (camelCase) and
# will reject `output_config` if it leaks through pass-through routes.
inference_params.pop("output_config", None)
# Extract requestMetadata before processing other parameters # Extract requestMetadata before processing other parameters
request_metadata = inference_params.pop("requestMetadata", None) request_metadata = inference_params.pop("requestMetadata", None)

View file

@ -7719,6 +7719,536 @@ class BaseLLMHTTPHandler:
response=response, response=response,
) )
async def async_vector_store_retrieve_handler(
self,
vector_store_id: str,
vector_store_provider_config: BaseVectorStoreConfig,
custom_llm_provider: str,
litellm_params: GenericLiteLLMParams,
logging_obj: LiteLLMLoggingObj,
extra_headers: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
) -> VectorStoreCreateResponse:
if client is None or not isinstance(client, AsyncHTTPHandler):
async_httpx_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders(custom_llm_provider),
params={"ssl_verify": litellm_params.get("ssl_verify", None)},
)
else:
async_httpx_client = client
headers = vector_store_provider_config.validate_environment(
headers=extra_headers or {}, litellm_params=litellm_params
)
if extra_headers:
headers.update(extra_headers)
api_base = vector_store_provider_config.get_complete_url(
api_base=litellm_params.api_base,
litellm_params=dict(litellm_params),
)
url = f"{api_base}/{vector_store_id}"
logging_obj.pre_call(
input="",
api_key="",
additional_args={
"api_base": api_base,
"headers": headers,
},
)
try:
response = await async_httpx_client.get(
url=url, headers=headers
)
except Exception as e:
raise self._handle_error(e=e, provider_config=vector_store_provider_config)
return vector_store_provider_config.transform_create_vector_store_response(
response=response,
)
def vector_store_retrieve_handler(
self,
vector_store_id: str,
vector_store_provider_config: BaseVectorStoreConfig,
custom_llm_provider: str,
litellm_params: GenericLiteLLMParams,
logging_obj: LiteLLMLoggingObj,
extra_headers: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
_is_async: bool = False,
) -> Union[
VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]
]:
if _is_async:
return self.async_vector_store_retrieve_handler(
vector_store_id=vector_store_id,
vector_store_provider_config=vector_store_provider_config,
litellm_params=litellm_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
extra_body=extra_body,
timeout=timeout,
client=client,
)
if client is None or not isinstance(client, HTTPHandler):
sync_httpx_client = _get_httpx_client(
params={"ssl_verify": litellm_params.get("ssl_verify", None)}
)
else:
sync_httpx_client = client
headers = vector_store_provider_config.validate_environment(
headers=extra_headers or {}, litellm_params=litellm_params
)
if extra_headers:
headers.update(extra_headers)
api_base = vector_store_provider_config.get_complete_url(
api_base=litellm_params.api_base,
litellm_params=dict(litellm_params),
)
url = f"{api_base}/{vector_store_id}"
logging_obj.pre_call(
input="",
api_key="",
additional_args={
"api_base": api_base,
"headers": headers,
},
)
try:
response = sync_httpx_client.get(url=url, headers=headers)
except Exception as e:
raise self._handle_error(e=e, provider_config=vector_store_provider_config)
return vector_store_provider_config.transform_create_vector_store_response(
response=response,
)
async def async_vector_store_list_handler(
self,
after: Optional[str],
before: Optional[str],
limit: Optional[int],
order: Optional[str],
vector_store_provider_config: BaseVectorStoreConfig,
custom_llm_provider: str,
litellm_params: GenericLiteLLMParams,
logging_obj: LiteLLMLoggingObj,
extra_headers: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
):
if client is None or not isinstance(client, AsyncHTTPHandler):
async_httpx_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders(custom_llm_provider),
params={"ssl_verify": litellm_params.get("ssl_verify", None)},
)
else:
async_httpx_client = client
headers = vector_store_provider_config.validate_environment(
headers=extra_headers or {}, litellm_params=litellm_params
)
if extra_headers:
headers.update(extra_headers)
api_base = vector_store_provider_config.get_complete_url(
api_base=litellm_params.api_base,
litellm_params=dict(litellm_params),
)
url = api_base
params = {}
if after is not None:
params["after"] = after
if before is not None:
params["before"] = before
if limit is not None:
params["limit"] = limit
if order is not None:
params["order"] = order
logging_obj.pre_call(
input="",
api_key="",
additional_args={
"api_base": api_base,
"headers": headers,
"params": params,
},
)
try:
response = await async_httpx_client.get(
url=url, headers=headers, params=params
)
except Exception as e:
raise self._handle_error(e=e, provider_config=vector_store_provider_config)
return response.json()
def vector_store_list_handler(
self,
after: Optional[str],
before: Optional[str],
limit: Optional[int],
order: Optional[str],
vector_store_provider_config: BaseVectorStoreConfig,
custom_llm_provider: str,
litellm_params: GenericLiteLLMParams,
logging_obj: LiteLLMLoggingObj,
extra_headers: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
_is_async: bool = False,
):
if _is_async:
return self.async_vector_store_list_handler(
after=after,
before=before,
limit=limit,
order=order,
vector_store_provider_config=vector_store_provider_config,
litellm_params=litellm_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
extra_body=extra_body,
timeout=timeout,
client=client,
)
if client is None or not isinstance(client, HTTPHandler):
sync_httpx_client = _get_httpx_client(
params={"ssl_verify": litellm_params.get("ssl_verify", None)}
)
else:
sync_httpx_client = client
headers = vector_store_provider_config.validate_environment(
headers=extra_headers or {}, litellm_params=litellm_params
)
if extra_headers:
headers.update(extra_headers)
api_base = vector_store_provider_config.get_complete_url(
api_base=litellm_params.api_base,
litellm_params=dict(litellm_params),
)
url = api_base
params = {}
if after is not None:
params["after"] = after
if before is not None:
params["before"] = before
if limit is not None:
params["limit"] = limit
if order is not None:
params["order"] = order
logging_obj.pre_call(
input="",
api_key="",
additional_args={
"api_base": api_base,
"headers": headers,
"params": params,
},
)
try:
response = sync_httpx_client.get(url=url, headers=headers, params=params)
except Exception as e:
raise self._handle_error(e=e, provider_config=vector_store_provider_config)
return response.json()
async def async_vector_store_update_handler(
self,
vector_store_id: str,
vector_store_update_optional_params: VectorStoreCreateOptionalRequestParams,
vector_store_provider_config: BaseVectorStoreConfig,
custom_llm_provider: str,
litellm_params: GenericLiteLLMParams,
logging_obj: LiteLLMLoggingObj,
extra_headers: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
) -> VectorStoreCreateResponse:
if client is None or not isinstance(client, AsyncHTTPHandler):
async_httpx_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders(custom_llm_provider),
params={"ssl_verify": litellm_params.get("ssl_verify", None)},
)
else:
async_httpx_client = client
headers = vector_store_provider_config.validate_environment(
headers=extra_headers or {}, litellm_params=litellm_params
)
if extra_headers:
headers.update(extra_headers)
api_base = vector_store_provider_config.get_complete_url(
api_base=litellm_params.api_base,
litellm_params=dict(litellm_params),
)
url = f"{api_base}/{vector_store_id}"
request_body = dict(vector_store_update_optional_params)
# Clean metadata to only include string values (OpenAI requirement)
if "metadata" in request_body and request_body["metadata"] is not None:
from litellm.utils import add_openai_metadata
request_body["metadata"] = add_openai_metadata(request_body["metadata"])
if extra_body:
request_body.update(extra_body)
logging_obj.pre_call(
input="",
api_key="",
additional_args={
"complete_input_dict": request_body,
"api_base": api_base,
"headers": headers,
},
)
try:
response = await async_httpx_client.post(
url=url, headers=headers, json=request_body, timeout=timeout
)
except Exception as e:
raise self._handle_error(e=e, provider_config=vector_store_provider_config)
return vector_store_provider_config.transform_create_vector_store_response(
response=response,
)
def vector_store_update_handler(
self,
vector_store_id: str,
vector_store_update_optional_params: VectorStoreCreateOptionalRequestParams,
vector_store_provider_config: BaseVectorStoreConfig,
custom_llm_provider: str,
litellm_params: GenericLiteLLMParams,
logging_obj: LiteLLMLoggingObj,
extra_headers: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
_is_async: bool = False,
) -> Union[
VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]
]:
if _is_async:
return self.async_vector_store_update_handler(
vector_store_id=vector_store_id,
vector_store_update_optional_params=vector_store_update_optional_params,
vector_store_provider_config=vector_store_provider_config,
litellm_params=litellm_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
extra_body=extra_body,
timeout=timeout,
client=client,
)
if client is None or not isinstance(client, HTTPHandler):
sync_httpx_client = _get_httpx_client(
params={"ssl_verify": litellm_params.get("ssl_verify", None)}
)
else:
sync_httpx_client = client
headers = vector_store_provider_config.validate_environment(
headers=extra_headers or {}, litellm_params=litellm_params
)
if extra_headers:
headers.update(extra_headers)
api_base = vector_store_provider_config.get_complete_url(
api_base=litellm_params.api_base,
litellm_params=dict(litellm_params),
)
url = f"{api_base}/{vector_store_id}"
request_body = dict(vector_store_update_optional_params)
# Clean metadata to only include string values (OpenAI requirement)
if "metadata" in request_body and request_body["metadata"] is not None:
from litellm.utils import add_openai_metadata
request_body["metadata"] = add_openai_metadata(request_body["metadata"])
if extra_body:
request_body.update(extra_body)
logging_obj.pre_call(
input="",
api_key="",
additional_args={
"complete_input_dict": request_body,
"api_base": api_base,
"headers": headers,
},
)
try:
response = sync_httpx_client.post(
url=url, headers=headers, json=request_body
)
except Exception as e:
raise self._handle_error(e=e, provider_config=vector_store_provider_config)
return vector_store_provider_config.transform_create_vector_store_response(
response=response,
)
async def async_vector_store_delete_handler(
self,
vector_store_id: str,
vector_store_provider_config: BaseVectorStoreConfig,
custom_llm_provider: str,
litellm_params: GenericLiteLLMParams,
logging_obj: LiteLLMLoggingObj,
extra_headers: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
):
if client is None or not isinstance(client, AsyncHTTPHandler):
async_httpx_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders(custom_llm_provider),
params={"ssl_verify": litellm_params.get("ssl_verify", None)},
)
else:
async_httpx_client = client
headers = vector_store_provider_config.validate_environment(
headers=extra_headers or {}, litellm_params=litellm_params
)
if extra_headers:
headers.update(extra_headers)
api_base = vector_store_provider_config.get_complete_url(
api_base=litellm_params.api_base,
litellm_params=dict(litellm_params),
)
url = f"{api_base}/{vector_store_id}"
logging_obj.pre_call(
input="",
api_key="",
additional_args={
"api_base": api_base,
"headers": headers,
},
)
try:
response = await async_httpx_client.delete(
url=url, headers=headers, timeout=timeout
)
except Exception as e:
raise self._handle_error(e=e, provider_config=vector_store_provider_config)
return response.json()
def vector_store_delete_handler(
self,
vector_store_id: str,
vector_store_provider_config: BaseVectorStoreConfig,
custom_llm_provider: str,
litellm_params: GenericLiteLLMParams,
logging_obj: LiteLLMLoggingObj,
extra_headers: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
_is_async: bool = False,
):
if _is_async:
return self.async_vector_store_delete_handler(
vector_store_id=vector_store_id,
vector_store_provider_config=vector_store_provider_config,
litellm_params=litellm_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
extra_body=extra_body,
timeout=timeout,
client=client,
)
if client is None or not isinstance(client, HTTPHandler):
sync_httpx_client = _get_httpx_client(
params={"ssl_verify": litellm_params.get("ssl_verify", None)}
)
else:
sync_httpx_client = client
headers = vector_store_provider_config.validate_environment(
headers=extra_headers or {}, litellm_params=litellm_params
)
if extra_headers:
headers.update(extra_headers)
api_base = vector_store_provider_config.get_complete_url(
api_base=litellm_params.api_base,
litellm_params=dict(litellm_params),
)
url = f"{api_base}/{vector_store_id}"
logging_obj.pre_call(
input="",
api_key="",
additional_args={
"api_base": api_base,
"headers": headers,
},
)
try:
response = sync_httpx_client.delete(url=url, headers=headers)
except Exception as e:
raise self._handle_error(e=e, provider_config=vector_store_provider_config)
return response.json()
##################################################################### #####################################################################
################ Vector Store Files HANDLERS ######################## ################ Vector Store Files HANDLERS ########################
##################################################################### #####################################################################

View file

@ -219,17 +219,32 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
tool_choice: Tool choice in OpenAI format (str or dict) tool_choice: Tool choice in OpenAI format (str or dict)
Returns: Returns:
Tool choice in Snowflake format (always an object) Tool choice in Snowflake format (always an object, never a string)
OpenAI format (string): "auto", "required", "none" OpenAI format (string):
OpenAI format (object): {"type": "function", "function": {"name": "get_weather"}} "auto", "required", "none"
Snowflake format (string values become objects): {"type": "auto"} OpenAI format (dict):
Snowflake format (specific tool): {"type": "tool", "name": ["get_weather"]} {"type": "function", "function": {"name": "get_weather"}}
Snowflake format:
{"type": "auto"} / {"type": "any"} / {"type": "none"}
{"type": "tool", "name": ["get_weather"]}
Snowflake's API (like Anthropic) requires tool_choice as an object
with a "type" field, not as a bare string.
""" """
if isinstance(tool_choice, str): if isinstance(tool_choice, str):
# Snowflake requires object format: {"type": "auto"} not string "auto" # Snowflake requires object format, not string.
return {"type": tool_choice} # Map OpenAI string values to Snowflake object format.
# "required" maps to "any" (Snowflake/Anthropic convention).
_type_map = {
"auto": "auto",
"required": "any",
"none": "none",
}
mapped_type = _type_map.get(tool_choice, tool_choice)
return {"type": mapped_type}
if isinstance(tool_choice, dict): if isinstance(tool_choice, dict):
if tool_choice.get("type") == "function": if tool_choice.get("type") == "function":

View file

@ -522,29 +522,6 @@ def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False):
return parameters return parameters
def _build_vertex_schema_for_gemini_2(parameters: dict) -> dict:
"""
Minimal schema builder for Gemini 2.0+ tool parameters.
Gemini 2.0+ accepts standard JSON Schema natively in tool parameters,
including lowercase types, anyOf with null, and bare {} (TYPE_UNSPECIFIED).
The only transformation needed is resolving $ref/$defs, which Gemini does
NOT support in tool parameters (returns 400).
This avoids the harmful transforms in _build_vertex_schema that break
JsonValue/Any semantics by coercing {} to {"type": "object"}.
"""
valid_schema_fields = set(get_type_hints(Schema).keys())
parameters = dict(parameters) # shallow copy to avoid mutating caller's dict
defs = parameters.pop("$defs", {})
unpack_defs(parameters, defs)
parameters = filter_schema_fields(parameters, valid_schema_fields)
return parameters
def _build_json_schema(parameters: dict) -> dict: def _build_json_schema(parameters: dict) -> dict:
""" """
Build a JSON Schema for use with Gemini's responseJsonSchema parameter. Build a JSON Schema for use with Gemini's responseJsonSchema parameter.

View file

@ -97,7 +97,6 @@ from ..common_utils import (
VertexAIError, VertexAIError,
_build_json_schema, _build_json_schema,
_build_vertex_schema, _build_vertex_schema,
_build_vertex_schema_for_gemini_2,
supports_response_json_schema, supports_response_json_schema,
) )
from ..vertex_llm_base import VertexBase from ..vertex_llm_base import VertexBase
@ -468,7 +467,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
return None return None
def _map_function( # noqa: PLR0915 def _map_function( # noqa: PLR0915
self, value: List[dict], optional_params: dict, model: str = "" self, value: List[dict], optional_params: dict
) -> List[Tools]: ) -> List[Tools]:
""" """
Map OpenAI-style tools/functions to Vertex AI format. Map OpenAI-style tools/functions to Vertex AI format.
@ -511,21 +510,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"parameters" in _openai_function_object "parameters" in _openai_function_object
and _openai_function_object["parameters"] is not None and _openai_function_object["parameters"] is not None
and isinstance(_openai_function_object["parameters"], dict) and isinstance(_openai_function_object["parameters"], dict)
): ): # OPENAI accepts JSON Schema, Google accepts OpenAPI schema.
if supports_response_json_schema(model): _openai_function_object["parameters"] = _build_vertex_schema(
# Gemini 2.0+: minimal transform (resolve $ref only) _openai_function_object["parameters"]
_openai_function_object["parameters"] = ( )
_build_vertex_schema_for_gemini_2(
_openai_function_object["parameters"]
)
)
else:
# Gemini 1.5: full OpenAPI-style transform
_openai_function_object["parameters"] = (
_build_vertex_schema(
_openai_function_object["parameters"]
)
)
openai_function_object = _openai_function_object openai_function_object = _openai_function_object
@ -1060,7 +1048,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
): ):
# Pass optional_params so _map_function can add toolConfig if needed # Pass optional_params so _map_function can add toolConfig if needed
mapped_tools = self._map_function( mapped_tools = self._map_function(
value=value, optional_params=optional_params, model=model value=value, optional_params=optional_params
) )
optional_params = self._add_tools_to_optional_params( optional_params = self._add_tools_to_optional_params(
optional_params, mapped_tools optional_params, mapped_tools

View file

@ -63,6 +63,7 @@ from litellm.utils import exception_type, get_litellm_params, get_optional_param
# Logging is imported lazily when needed to avoid loading litellm_logging at import time # Logging is imported lazily when needed to avoid loading litellm_logging at import time
if TYPE_CHECKING: if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.utils import TokenCountResponse
from litellm.constants import ( from litellm.constants import (
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,

View file

@ -7997,6 +7997,80 @@
"supports_response_schema": true, "supports_response_schema": true,
"supports_tool_choice": true "supports_tool_choice": true
}, },
"black_forest_labs/flux-kontext-pro": {
"litellm_provider": "black_forest_labs",
"mode": "image_edit",
"output_cost_per_image": 0.04,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/edits",
"/v1/images/generations"
]
},
"black_forest_labs/flux-kontext-max": {
"litellm_provider": "black_forest_labs",
"mode": "image_edit",
"output_cost_per_image": 0.08,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/edits",
"/v1/images/generations"
]
},
"black_forest_labs/flux-pro-1.0-fill": {
"litellm_provider": "black_forest_labs",
"mode": "image_edit",
"output_cost_per_image": 0.05,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/edits"
]
},
"black_forest_labs/flux-pro-1.0-expand": {
"litellm_provider": "black_forest_labs",
"mode": "image_edit",
"output_cost_per_image": 0.05,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/edits"
]
},
"black_forest_labs/flux-pro-1.1": {
"litellm_provider": "black_forest_labs",
"mode": "image_generation",
"output_cost_per_image": 0.04,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/generations"
]
},
"black_forest_labs/flux-pro-1.1-ultra": {
"litellm_provider": "black_forest_labs",
"mode": "image_generation",
"output_cost_per_image": 0.06,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/generations"
]
},
"black_forest_labs/flux-dev": {
"litellm_provider": "black_forest_labs",
"mode": "image_generation",
"output_cost_per_image": 0.025,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/generations"
]
},
"black_forest_labs/flux-pro": {
"litellm_provider": "black_forest_labs",
"mode": "image_generation",
"output_cost_per_image": 0.05,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/generations"
]
},
"cerebras/llama-3.3-70b": { "cerebras/llama-3.3-70b": {
"input_cost_per_token": 8.5e-07, "input_cost_per_token": 8.5e-07,
"litellm_provider": "cerebras", "litellm_provider": "cerebras",
@ -14442,7 +14516,7 @@
"input_cost_per_audio_per_second": 0.00016, "input_cost_per_audio_per_second": 0.00016,
"input_cost_per_image": 0.00012, "input_cost_per_image": 0.00012,
"input_cost_per_token": 2e-07, "input_cost_per_token": 2e-07,
"input_cost_per_video_per_second": 0.0237, "input_cost_per_video_per_second": 0.00079,
"litellm_provider": "vertex_ai-embedding-models", "litellm_provider": "vertex_ai-embedding-models",
"max_input_tokens": 8192, "max_input_tokens": 8192,
"max_tokens": 8192, "max_tokens": 8192,
@ -14453,14 +14527,17 @@
"uses_embed_content": true "uses_embed_content": true
}, },
"vertex_ai/gemini-embedding-2-preview": { "vertex_ai/gemini-embedding-2-preview": {
"input_cost_per_token": 1.5e-07, "input_cost_per_audio_per_second": 0.00016,
"input_cost_per_image": 0.00012,
"input_cost_per_token": 2e-07,
"input_cost_per_video_per_second": 0.00079,
"litellm_provider": "vertex_ai", "litellm_provider": "vertex_ai",
"max_input_tokens": 8192, "max_input_tokens": 8192,
"max_tokens": 8192, "max_tokens": 8192,
"mode": "embedding", "mode": "embedding",
"output_cost_per_token": 0, "output_cost_per_token": 0,
"output_vector_size": 3072, "output_vector_size": 3072,
"source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
"supports_multimodal": true, "supports_multimodal": true,
"uses_embed_content": true "uses_embed_content": true
}, },
@ -14502,7 +14579,10 @@
"tpm": 10000000 "tpm": 10000000
}, },
"gemini/gemini-embedding-2-preview": { "gemini/gemini-embedding-2-preview": {
"input_cost_per_token": 1.5e-07, "input_cost_per_audio_per_second": 0.00016,
"input_cost_per_image": 0.00012,
"input_cost_per_token": 2e-07,
"input_cost_per_video_per_second": 0.00079,
"litellm_provider": "gemini", "litellm_provider": "gemini",
"max_input_tokens": 8192, "max_input_tokens": 8192,
"max_tokens": 8192, "max_tokens": 8192,
@ -14510,7 +14590,7 @@
"output_cost_per_token": 0, "output_cost_per_token": 0,
"output_vector_size": 3072, "output_vector_size": 3072,
"rpm": 10000, "rpm": 10000,
"source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", "source": "https://ai.google.dev/gemini-api/docs/pricing",
"supports_multimodal": true, "supports_multimodal": true,
"tpm": 10000000 "tpm": 10000000
}, },

View file

@ -1,4 +1,6 @@
from datetime import datetime, timezone import base64
import json
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast
from litellm._logging import verbose_proxy_logger from litellm._logging import verbose_proxy_logger
@ -499,7 +501,6 @@ async def store_user_credential(
credential: str, credential: str,
) -> None: ) -> None:
"""Store a user credential for a BYOK MCP server.""" """Store a user credential for a BYOK MCP server."""
import base64
encoded = base64.urlsafe_b64encode(credential.encode()).decode() encoded = base64.urlsafe_b64encode(credential.encode()).decode()
await prisma_client.db.litellm_mcpusercredentials.upsert( await prisma_client.db.litellm_mcpusercredentials.upsert(
@ -521,7 +522,6 @@ async def get_user_credential(
server_id: str, server_id: str,
) -> Optional[str]: ) -> Optional[str]:
"""Return credential for a user+server pair, or None.""" """Return credential for a user+server pair, or None."""
import base64
row = await prisma_client.db.litellm_mcpusercredentials.find_unique( row = await prisma_client.db.litellm_mcpusercredentials.find_unique(
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
@ -563,6 +563,138 @@ async def delete_user_credential(
) )
# ── OAuth2 user-credential helpers ────────────────────────────────────────────
async def store_user_oauth_credential(
prisma_client: PrismaClient,
user_id: str,
server_id: str,
access_token: str,
refresh_token: Optional[str] = None,
expires_in: Optional[int] = None,
scopes: Optional[List[str]] = None,
) -> None:
"""Persist an OAuth2 access token for a user+server pair.
The payload is JSON-serialised and stored base64-encoded in the same
``credential_b64`` column used by BYOK. A ``"type": "oauth2"`` key
differentiates it from plain BYOK API keys.
"""
expires_at: Optional[str] = None
if expires_in is not None:
expires_at = (
datetime.now(timezone.utc) + timedelta(seconds=expires_in)
).isoformat()
payload: Dict[str, Any] = {
"type": "oauth2",
"access_token": access_token,
"connected_at": datetime.now(timezone.utc).isoformat(),
}
if refresh_token:
payload["refresh_token"] = refresh_token
if expires_at:
payload["expires_at"] = expires_at
if scopes:
payload["scopes"] = scopes
# Guard against silently overwriting a BYOK credential with an OAuth token.
# BYOK credentials lack a "type" field (or use a non-"oauth2" type).
existing = await prisma_client.db.litellm_mcpusercredentials.find_unique(
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
)
if existing is not None:
_byok_error = ValueError(
f"A non-OAuth2 credential already exists for user {user_id} "
f"and server {server_id}. Refusing to overwrite."
)
try:
raw = json.loads(base64.urlsafe_b64decode(existing.credential_b64).decode())
except Exception:
# Credential is not base64+JSON — it's a plain-text BYOK key.
raise _byok_error
if raw.get("type") != "oauth2":
raise _byok_error
encoded = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode()
await prisma_client.db.litellm_mcpusercredentials.upsert(
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}},
data={
"create": {
"user_id": user_id,
"server_id": server_id,
"credential_b64": encoded,
},
"update": {"credential_b64": encoded},
},
)
def is_oauth_credential_expired(cred: Dict[str, Any]) -> bool:
"""Return True if the OAuth2 credential's access_token has expired.
Checks the ``expires_at`` ISO-format string stored in the credential payload.
Returns False when ``expires_at`` is absent or unparseable (treat as non-expired).
"""
expires_at = cred.get("expires_at")
if not expires_at:
return False
try:
exp_dt = datetime.fromisoformat(expires_at)
if exp_dt.tzinfo is None:
exp_dt = exp_dt.replace(tzinfo=timezone.utc)
return datetime.now(timezone.utc) > exp_dt
except (ValueError, TypeError):
return False
async def get_user_oauth_credential(
prisma_client: PrismaClient,
user_id: str,
server_id: str,
) -> Optional[Dict[str, Any]]:
"""Return the decoded OAuth2 payload dict for a user+server pair, or None."""
row = await prisma_client.db.litellm_mcpusercredentials.find_unique(
where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}
)
if row is None:
return None
try:
decoded = base64.urlsafe_b64decode(row.credential_b64).decode()
parsed = json.loads(decoded)
if isinstance(parsed, dict) and parsed.get("type") == "oauth2":
return parsed
# Row exists but is a BYOK (plain string), not an OAuth token
return None
except Exception:
return None
async def list_user_oauth_credentials(
prisma_client: PrismaClient,
user_id: str,
) -> List[Dict[str, Any]]:
"""Return all OAuth2 credential payloads for a user, tagged with server_id."""
rows = await prisma_client.db.litellm_mcpusercredentials.find_many(
where={"user_id": user_id}
)
results: List[Dict[str, Any]] = []
for row in rows:
try:
decoded = base64.urlsafe_b64decode(row.credential_b64).decode()
parsed = json.loads(decoded)
if isinstance(parsed, dict) and parsed.get("type") == "oauth2":
parsed["server_id"] = row.server_id
results.append(parsed)
except Exception:
pass # Skip non-OAuth rows (BYOK plain strings)
return results
async def approve_mcp_server( async def approve_mcp_server(
prisma_client: PrismaClient, prisma_client: PrismaClient,
server_id: str, server_id: str,

View file

@ -312,8 +312,8 @@ async def register_client_with_server(
@router.get("/authorize") @router.get("/authorize")
async def authorize( async def authorize(
request: Request, request: Request,
client_id: str,
redirect_uri: str, redirect_uri: str,
client_id: Optional[str] = None,
state: str = "", state: str = "",
mcp_server_name: Optional[str] = None, mcp_server_name: Optional[str] = None,
code_challenge: Optional[str] = None, code_challenge: Optional[str] = None,
@ -326,19 +326,34 @@ async def authorize(
global_mcp_server_manager, global_mcp_server_manager,
) )
lookup_name = mcp_server_name or client_id lookup_name: Optional[str] = mcp_server_name or client_id
client_ip = IPAddressUtils.get_mcp_client_ip(request) client_ip = IPAddressUtils.get_mcp_client_ip(request)
mcp_server = global_mcp_server_manager.get_mcp_server_by_name( mcp_server = (
lookup_name, client_ip=client_ip global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip)
if lookup_name
else None
) )
if mcp_server is None and mcp_server_name is None: if mcp_server is None and mcp_server_name is None:
mcp_server = _resolve_oauth2_server_for_root_endpoints() mcp_server = _resolve_oauth2_server_for_root_endpoints()
if mcp_server is None: if mcp_server is None:
raise HTTPException(status_code=404, detail="MCP server not found") raise HTTPException(status_code=404, detail="MCP server not found")
# Use server's stored client_id when caller doesn't supply one.
# Raise a clear error instead of passing an empty string — an empty
# client_id would silently produce a broken authorization URL.
resolved_client_id: str = mcp_server.client_id or client_id or ""
if not resolved_client_id:
raise HTTPException(
status_code=400,
detail={
"error": "client_id is required but was not supplied and is not "
"stored on the MCP server record. Provide client_id as a query "
"parameter or configure it on the server."
},
)
return await authorize_with_server( return await authorize_with_server(
request=request, request=request,
mcp_server=mcp_server, mcp_server=mcp_server,
client_id=client_id, client_id=resolved_client_id,
redirect_uri=redirect_uri, redirect_uri=redirect_uri,
state=state, state=state,
code_challenge=code_challenge, code_challenge=code_challenge,

View file

@ -1,6 +1,6 @@
import importlib import importlib
from datetime import datetime from datetime import datetime, timezone
from typing import Any, Awaitable, Callable, Dict, List, Optional, Union from typing import Any, Awaitable, Callable, Dict, List, Optional, Set, Union
from fastapi import APIRouter, Depends, HTTPException, Query, Request from fastapi import APIRouter, Depends, HTTPException, Query, Request
@ -69,6 +69,132 @@ if MCP_AVAILABLE:
return server_auth return server_auth
return mcp_auth_header return mcp_auth_header
def _get_oauth2_server_ids(allowed_server_ids: List[str]) -> Set[str]:
"""Return the subset of *allowed_server_ids* whose servers use OAuth2 auth.
Used as a cheap pre-flight check to skip bulk credential fetching when no
OAuth2 servers are involved in the current request.
"""
return {
sid
for sid in allowed_server_ids
if getattr(
global_mcp_server_manager.get_mcp_server_by_id(sid), "auth_type", None
)
== MCPAuth.oauth2
}
async def _get_user_oauth_extra_headers(
server,
user_api_key_dict: UserAPIKeyAuth,
prefetched_creds: Optional[Dict[str, Dict[str, Any]]] = None,
) -> Optional[Dict[str, str]]:
"""
For OAuth2 servers, look up the user's stored access token and return it
as extra_headers {"Authorization": "Bearer <token>"} so that it reaches
the MCP server the same way the admin "Add MCP / Authorize and Fetch" flow does.
Returns None for non-OAuth2 servers or when no credential is stored.
Args:
prefetched_creds: Optional dict keyed by server_id with credential payloads.
When provided, avoids a per-server DB round-trip.
"""
if getattr(server, "auth_type", None) != MCPAuth.oauth2:
return None
user_id = getattr(user_api_key_dict, "user_id", None)
server_id = getattr(server, "server_id", None)
if not user_id or not server_id:
return None
try:
from litellm.proxy._experimental.mcp_server.db import (
get_user_oauth_credential,
is_oauth_credential_expired,
)
if prefetched_creds is not None:
cred = prefetched_creds.get(server_id)
else:
from litellm.proxy.utils import get_prisma_client_or_throw
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to use OAuth2 MCP tools."
)
cred = await get_user_oauth_credential(prisma_client, user_id, server_id)
if cred and cred.get("access_token"):
if is_oauth_credential_expired(cred):
verbose_logger.debug(
f"_get_user_oauth_extra_headers: token expired for "
f"user={user_id} server={server_id}"
)
return None
return {"Authorization": f"Bearer {cred['access_token']}"}
except Exception as e:
verbose_logger.warning(
f"_get_user_oauth_extra_headers: failed to retrieve credential for "
f"user={user_id} server={server_id}: {e}"
)
return None
async def _prefetch_user_oauth_creds(
user_api_key_dict: UserAPIKeyAuth,
) -> Dict[str, Dict[str, Any]]:
"""Fetch all OAuth2 credentials for the user in a single DB query.
Returns a dict keyed by server_id. Used to avoid N+1 DB queries when
iterating over multiple OAuth2 MCP servers.
"""
user_id = getattr(user_api_key_dict, "user_id", None)
if not user_id:
return {}
try:
from litellm.proxy._experimental.mcp_server.db import (
list_user_oauth_credentials,
)
from litellm.proxy.utils import get_prisma_client_or_throw
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to use OAuth2 MCP tools."
)
creds = await list_user_oauth_credentials(prisma_client, user_id)
return {c["server_id"]: c for c in creds if "server_id" in c}
except Exception as e:
verbose_logger.warning(
f"_prefetch_user_oauth_creds: failed to prefetch for user={user_id}: {e}"
)
return {}
async def _get_bulk_user_oauth_headers(
user_api_key_dict: UserAPIKeyAuth,
) -> Dict[str, Dict[str, str]]:
"""
Fetch ALL OAuth2 credentials for the current user in a single DB query and
return a mapping of server_id {"Authorization": "Bearer <token>"}.
This is the batch alternative to calling _get_user_oauth_extra_headers
per-server inside a loop (N+1 DB queries).
"""
user_id = getattr(user_api_key_dict, "user_id", None)
if not user_id:
return {}
try:
from litellm.proxy._experimental.mcp_server.db import (
list_user_oauth_credentials,
)
from litellm.proxy.utils import get_prisma_client_or_throw
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to use OAuth2 MCP tools."
)
creds = await list_user_oauth_credentials(prisma_client, user_id)
return {
c["server_id"]: {"Authorization": f"Bearer {c['access_token']}"}
for c in creds
if c.get("access_token") and c.get("server_id")
}
except Exception:
verbose_logger.debug("Failed to bulk-fetch OAuth credentials", exc_info=True)
return {}
def _create_tool_response_objects(tools, server_mcp_info): def _create_tool_response_objects(tools, server_mcp_info):
"""Helper function to create tool response objects.""" """Helper function to create tool response objects."""
return [ return [
@ -162,11 +288,13 @@ if MCP_AVAILABLE:
server_auth_header, server_auth_header,
raw_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None,
user_api_key_auth: Optional[UserAPIKeyAuth] = None, user_api_key_auth: Optional[UserAPIKeyAuth] = None,
extra_headers: Optional[Dict[str, str]] = None,
): ):
"""Helper function to get tools for a single server.""" """Helper function to get tools for a single server."""
tools = await global_mcp_server_manager._get_tools_from_server( tools = await global_mcp_server_manager._get_tools_from_server(
server=server, server=server,
mcp_auth_header=server_auth_header, mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
add_prefix=False, add_prefix=False,
raw_headers=raw_headers, raw_headers=raw_headers,
) )
@ -295,8 +423,23 @@ if MCP_AVAILABLE:
# If server_id is specified, only query that specific server # If server_id is specified, only query that specific server
if server_id: if server_id:
# Resolve a server name to its UUID if needed (MCPConnectPicker passes
# server_name strings, but allowed_server_ids_set contains UUIDs).
# _name_resolved is kept so the second check can reuse it for accurate
# IP-filter error reporting if the resolved UUID is not in allowed_server_ids.
_name_resolved = None
if server_id not in allowed_server_ids: if server_id not in allowed_server_ids:
_server = global_mcp_server_manager.get_mcp_server_by_id(server_id) _name_resolved = global_mcp_server_manager.get_mcp_server_by_name(server_id)
if _name_resolved is not None and _name_resolved.server_id in set(allowed_server_ids):
server_id = _name_resolved.server_id
if server_id not in allowed_server_ids:
# Try UUID lookup first; fall back to the name-resolved server so that
# IP-filter reporting works correctly even when server_id is a name string.
_server = (
global_mcp_server_manager.get_mcp_server_by_id(server_id)
or _name_resolved
)
if ( if (
_server is not None _server is not None
and _rest_client_ip is not None and _rest_client_ip is not None
@ -334,6 +477,8 @@ if MCP_AVAILABLE:
server_auth_header = _get_server_auth_header( server_auth_header = _get_server_auth_header(
server, mcp_server_auth_headers, mcp_auth_header server, mcp_server_auth_headers, mcp_auth_header
) )
# Single-server request: targeted lookup is more efficient than a bulk fetch.
user_oauth_extra_headers = await _get_user_oauth_extra_headers(server, user_api_key_dict)
try: try:
list_tools_result = await _get_tools_for_single_server( list_tools_result = await _get_tools_for_single_server(
@ -341,6 +486,7 @@ if MCP_AVAILABLE:
server_auth_header, server_auth_header,
raw_headers_from_request, raw_headers_from_request,
user_api_key_dict, user_api_key_dict,
extra_headers=user_oauth_extra_headers,
) )
except Exception as e: except Exception as e:
verbose_logger.exception( verbose_logger.exception(
@ -374,6 +520,14 @@ if MCP_AVAILABLE:
}, },
) )
# Pre-fetch OAuth credentials only when at least one allowed server uses OAuth2,
# to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers.
prefetched_oauth_creds = (
await _prefetch_user_oauth_creds(user_api_key_dict)
if _get_oauth2_server_ids(allowed_server_ids)
else {}
)
# Query all servers the user has access to # Query all servers the user has access to
errors = [] errors = []
for allowed_server_id in allowed_server_ids: for allowed_server_id in allowed_server_ids:
@ -386,6 +540,9 @@ if MCP_AVAILABLE:
server_auth_header = _get_server_auth_header( server_auth_header = _get_server_auth_header(
server, mcp_server_auth_headers, mcp_auth_header server, mcp_server_auth_headers, mcp_auth_header
) )
user_oauth_extra_headers = await _get_user_oauth_extra_headers(
server, user_api_key_dict, prefetched_creds=prefetched_oauth_creds
)
try: try:
tools_result = await _get_tools_for_single_server( tools_result = await _get_tools_for_single_server(
@ -393,6 +550,7 @@ if MCP_AVAILABLE:
server_auth_header, server_auth_header,
raw_headers_from_request, raw_headers_from_request,
user_api_key_dict, user_api_key_dict,
extra_headers=user_oauth_extra_headers,
) )
list_tools_result.extend(tools_result) list_tools_result.extend(tools_result)
except Exception as e: except Exception as e:
@ -509,6 +667,16 @@ if MCP_AVAILABLE:
request, user_api_key_dict, server_id request, user_api_key_dict, server_id
) )
# Look up per-user OAuth headers for this server (mirrors list_tool_rest_api).
user_oauth_extra_headers: Optional[Dict[str, str]] = None
target_server = next(
(s for s in allowed_mcp_servers if s.server_id == server_id), None
)
if target_server is not None:
user_oauth_extra_headers = await _get_user_oauth_extra_headers(
target_server, user_api_key_dict
)
# Call execute_mcp_tool directly (permission checks already done) # Call execute_mcp_tool directly (permission checks already done)
result = await execute_mcp_tool( result = await execute_mcp_tool(
name=tool_name, name=tool_name,
@ -518,7 +686,7 @@ if MCP_AVAILABLE:
user_api_key_auth=data.get("user_api_key_auth"), user_api_key_auth=data.get("user_api_key_auth"),
mcp_auth_header=data.get("mcp_auth_header"), mcp_auth_header=data.get("mcp_auth_header"),
mcp_server_auth_headers=data.get("mcp_server_auth_headers"), mcp_server_auth_headers=data.get("mcp_server_auth_headers"),
oauth2_headers=data.get("oauth2_headers"), oauth2_headers=user_oauth_extra_headers or data.get("oauth2_headers"),
raw_headers=data.get("raw_headers"), raw_headers=data.get("raw_headers"),
litellm_logging_obj=data.get("litellm_logging_obj"), litellm_logging_obj=data.get("litellm_logging_obj"),
) )

View file

@ -8,7 +8,7 @@ import contextlib
import time import time
import traceback import traceback
import uuid import uuid
from datetime import datetime from datetime import datetime, timezone
from typing import ( from typing import (
Any, Any,
AsyncIterator, AsyncIterator,
@ -877,6 +877,84 @@ if MCP_AVAILABLE:
return allowed_mcp_servers return allowed_mcp_servers
async def _get_user_oauth_extra_headers_from_db(
server: MCPServer,
user_api_key_auth: Optional[UserAPIKeyAuth],
prefetched_creds: Optional[Dict[str, Dict[str, Any]]] = None,
) -> Optional[Dict[str, str]]:
"""Look up stored OAuth2 token for (user, server) from DB and return as extra_headers dict.
Args:
prefetched_creds: Optional dict keyed by server_id with credential payloads.
When provided, avoids a per-server DB round-trip.
"""
if server.auth_type != MCPAuth.oauth2:
return None
if user_api_key_auth is None:
return None
user_id = getattr(user_api_key_auth, "user_id", None)
server_id = getattr(server, "server_id", None)
if not user_id or not server_id:
return None
try:
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
get_user_oauth_credential,
is_oauth_credential_expired,
)
if prefetched_creds is not None:
cred = prefetched_creds.get(server_id)
else:
from litellm.proxy.utils import ( # noqa: PLC0415
get_prisma_client_or_throw,
)
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to use OAuth2 MCP tools."
)
cred = await get_user_oauth_credential(prisma_client, user_id, server_id)
if cred and cred.get("access_token"):
if is_oauth_credential_expired(cred):
verbose_logger.debug(
f"_get_user_oauth_extra_headers_from_db: token expired for "
f"user={user_id} server={server_id}"
)
return None
return {"Authorization": f"Bearer {cred['access_token']}"}
except Exception as e:
verbose_logger.warning(
f"_get_user_oauth_extra_headers_from_db: failed to retrieve credential for "
f"user={user_id} server={server_id}: {e}"
)
return None
async def _prefetch_oauth_creds_for_user(
user_api_key_auth: Optional[UserAPIKeyAuth],
) -> Dict[str, Dict[str, Any]]:
"""Fetch all OAuth2 credentials for the user in one DB query.
Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops.
"""
user_id = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None
if not user_id:
return {}
try:
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
list_user_oauth_credentials,
)
from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to use OAuth2 MCP tools."
)
creds = await list_user_oauth_credentials(prisma_client, user_id)
return {c["server_id"]: c for c in creds if "server_id" in c}
except Exception as e:
verbose_logger.warning(
f"_prefetch_oauth_creds_for_user: failed to prefetch for user={user_id}: {e}"
)
return {}
def _prepare_mcp_server_headers( def _prepare_mcp_server_headers(
server: MCPServer, server: MCPServer,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]],
@ -1020,6 +1098,18 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers, mcp_servers=mcp_servers,
) )
# Pre-fetch OAuth credentials only when at least one server uses OAuth2,
# to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers.
_has_oauth2_server = any(
getattr(s, "auth_type", None) == MCPAuth.oauth2
for s in allowed_mcp_servers
)
_prefetched_oauth_creds = (
await _prefetch_oauth_creds_for_user(user_api_key_auth)
if _has_oauth2_server
else {}
)
async def _fetch_and_filter_server_tools( async def _fetch_and_filter_server_tools(
server: MCPServer, server: MCPServer,
) -> List[MCPTool]: ) -> List[MCPTool]:
@ -1035,6 +1125,12 @@ if MCP_AVAILABLE:
raw_headers=raw_headers, raw_headers=raw_headers,
) )
# If no OAuth2 token came from request headers, fall back to pre-fetched creds
if extra_headers is None and server.auth_type == MCPAuth.oauth2:
extra_headers = await _get_user_oauth_extra_headers_from_db(
server, user_api_key_auth, prefetched_creds=_prefetched_oauth_creds
)
try: try:
tools = await global_mcp_server_manager._get_tools_from_server( tools = await global_mcp_server_manager._get_tools_from_server(
server=server, server=server,

View file

@ -1,59 +1,40 @@
import enum import enum
import json import json
from datetime import datetime from datetime import datetime
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal,
Optional, Union)
import httpx import httpx
from pydantic import ( from pydantic import (BaseModel, ConfigDict, Field, Json, field_validator,
BaseModel, model_validator)
ConfigDict,
Field,
Json,
field_validator,
model_validator,
)
from typing_extensions import Required, TypedDict from typing_extensions import Required, TypedDict
from litellm._uuid import uuid from litellm._uuid import uuid
from litellm.types.integrations.slack_alerting import AlertType from litellm.types.integrations.slack_alerting import AlertType
from litellm.types.llms.openai import ( from litellm.types.llms.openai import (AllMessageValues, OpenAIFileObject,
AllMessageValues, ResponsesAPIResponse)
OpenAIFileObject, from litellm.types.mcp import (MCPAuthType, MCPCredentials, MCPTransport,
ResponsesAPIResponse, MCPTransportType)
)
from litellm.types.mcp import (
MCPAuthType,
MCPCredentials,
MCPTransport,
MCPTransportType,
)
from litellm.types.mcp_server.mcp_server_manager import MCPInfo from litellm.types.mcp_server.mcp_server_manager import MCPInfo
from litellm.types.router import RouterErrors, UpdateRouterConfig from litellm.types.router import RouterErrors, UpdateRouterConfig
from litellm.types.secret_managers.main import KeyManagementSystem from litellm.types.secret_managers.main import KeyManagementSystem
from litellm.types.utils import ( from litellm.types.utils import (CallTypes, CostBreakdown, EmbeddingResponse,
CallTypes, GenericBudgetConfigType, ImageResponse,
CostBreakdown, LiteLLMBatch, LiteLLMFineTuningJob,
EmbeddingResponse, LiteLLMPydanticObjectBase, ModelResponse,
GenericBudgetConfigType, ProviderField, StandardCallbackDynamicParams,
ImageResponse, StandardLoggingGuardrailInformation,
LiteLLMBatch, StandardLoggingMCPToolCall,
LiteLLMFineTuningJob, StandardLoggingModelInformation,
LiteLLMPydanticObjectBase, StandardLoggingPayloadErrorInformation,
ModelResponse, StandardLoggingPayloadStatus,
ProviderField, StandardLoggingVectorStoreRequest,
StandardCallbackDynamicParams, StandardPassThroughResponseObject,
StandardLoggingGuardrailInformation, TextCompletionResponse)
StandardLoggingMCPToolCall,
StandardLoggingModelInformation,
StandardLoggingPayloadErrorInformation,
StandardLoggingPayloadStatus,
StandardLoggingVectorStoreRequest,
StandardPassThroughResponseObject,
TextCompletionResponse,
)
from litellm.types.videos.main import VideoObject from litellm.types.videos.main import VideoObject
from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type from .types_utils.utils import (get_instance_fn,
validate_custom_validate_return_type)
if TYPE_CHECKING: if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span from opentelemetry.trace import Span as _Span
@ -684,6 +665,8 @@ class LiteLLMRoutes(enum.Enum):
"/team/daily/activity", "/team/daily/activity",
"/tag/daily/activity", "/tag/daily/activity",
"/tag/list", "/tag/list",
"/audit",
"/audit/{id}",
] + info_routes ] + info_routes
# All routes accesible by an Org Admin # All routes accesible by an Org Admin
@ -855,6 +838,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase):
vector_stores: Optional[List[str]] = None vector_stores: Optional[List[str]] = None
agents: Optional[List[str]] = None agents: Optional[List[str]] = None
agent_access_groups: Optional[List[str]] = None agent_access_groups: Optional[List[str]] = None
models: Optional[List[str]] = None
class GenerateRequestBase(LiteLLMPydanticObjectBase): class GenerateRequestBase(LiteLLMPydanticObjectBase):
@ -1287,6 +1271,37 @@ class MCPUserCredentialResponse(LiteLLMPydanticObjectBase):
has_credential: bool has_credential: bool
class MCPOAuthUserCredentialRequest(LiteLLMPydanticObjectBase):
"""Stores a user's OAuth2 token for an OpenAPI MCP server."""
access_token: str
refresh_token: Optional[str] = None
expires_in: Optional[int] = None # seconds until expiry
scopes: Optional[List[str]] = None
class MCPOAuthUserCredentialStatus(LiteLLMPydanticObjectBase):
"""Describes whether the calling user has a stored OAuth credential."""
server_id: str
has_credential: bool
expires_at: Optional[str] = None # ISO-8601
is_expired: bool = False
connected_at: Optional[str] = None # ISO-8601
class MCPUserCredentialListItem(LiteLLMPydanticObjectBase):
"""One entry in the /user-credentials list."""
server_id: str
server_name: Optional[str] = None
alias: Optional[str] = None
credential_type: str # "oauth2" or "byok"
has_credential: bool
expires_at: Optional[str] = None # ISO-8601; None means non-expiring
connected_at: Optional[str] = None # ISO-8601
class RejectMCPServerRequest(LiteLLMPydanticObjectBase): class RejectMCPServerRequest(LiteLLMPydanticObjectBase):
review_notes: Optional[str] = None review_notes: Optional[str] = None
@ -2430,6 +2445,7 @@ class UserAPIKeyAuth(
user_max_budget: Optional[float] = None user_max_budget: Optional[float] = None
request_route: Optional[str] = None request_route: Optional[str] = None
user: Optional[Any] = None # Expanded user object when expand=user is used user: Optional[Any] = None # Expanded user object when expand=user is used
created_by_user: Optional[Any] = None # Expanded created_by user when expand=user is used
end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
model_config = ConfigDict(arbitrary_types_allowed=True) model_config = ConfigDict(arbitrary_types_allowed=True)
@ -2473,7 +2489,8 @@ class UserAPIKeyAuth(
This is used to track number of requests/spend for health check calls. This is used to track number of requests/spend for health check calls.
""" """
from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME from litellm.constants import \
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
return cls( return cls(
api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
@ -2505,7 +2522,8 @@ class UserAPIKeyAuth(
This is used to track actions performed by automated system jobs. This is used to track actions performed by automated system jobs.
""" """
from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME from litellm.constants import \
LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
return cls( return cls(
api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
@ -2911,7 +2929,8 @@ class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase):
@model_validator(mode="after") @model_validator(mode="after")
def mask_api_keys(self): def mask_api_keys(self):
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.litellm_core_utils.sensitive_data_masker import \
SensitiveDataMasker
masker = SensitiveDataMasker(sensitive_patterns={"key"}) masker = SensitiveDataMasker(sensitive_patterns={"key"})

View file

@ -18,8 +18,11 @@ async def get_ui_config():
from litellm.proxy.auth.auth_utils import _has_user_setup_sso from litellm.proxy.auth.auth_utils import _has_user_setup_sso
from litellm.proxy.utils import get_proxy_base_url, get_server_root_path from litellm.proxy.utils import get_proxy_base_url, get_server_root_path
from litellm.proxy.proxy_server import general_settings
auto_redirect_ui_login_to_sso = ( auto_redirect_ui_login_to_sso = (
os.getenv("AUTO_REDIRECT_UI_LOGIN_TO_SSO", "false").lower() == "true" os.getenv("AUTO_REDIRECT_UI_LOGIN_TO_SSO", "false").lower() == "true"
or general_settings.get("auto_redirect_ui_login_to_sso", False) is True
) )
admin_ui_disabled = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true" admin_ui_disabled = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true"

View file

@ -1122,87 +1122,73 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
) )
return response return response
async def async_post_call_streaming_iterator_hook( async def _stream_apply_output_masking(
self, self,
user_api_key_dict: UserAPIKeyAuth,
response: Any, response: Any,
request_data: dict, request_data: dict,
) -> AsyncGenerator[Union[ModelResponseStream, bytes], None]: ) -> AsyncGenerator[Union[ModelResponseStream, bytes], None]:
""" """Apply Presidio masking to streaming output (apply_to_output=True path)."""
Process streaming response chunks to unmask PII tokens when needed.
"""
from litellm.llms.base_llm.base_model_iterator import ( from litellm.llms.base_llm.base_model_iterator import (
convert_model_response_to_streaming, convert_model_response_to_streaming,
) )
from litellm.main import stream_chunk_builder from litellm.main import stream_chunk_builder
from litellm.types.utils import ModelResponse from litellm.types.utils import ModelResponse
# --- Output masking path (apply_to_output=True) --- all_chunks: List[ModelResponseStream] = []
if self.apply_to_output: try:
all_chunks: List[ModelResponseStream] = [] async for chunk in response:
try: if isinstance(chunk, ModelResponseStream):
async for chunk in response: all_chunks.append(chunk)
if isinstance(chunk, ModelResponseStream): elif isinstance(chunk, bytes):
all_chunks.append(chunk) yield chunk # type: ignore[misc]
elif isinstance(chunk, bytes): continue
# Anthropic native SSE: pass through as-is
yield chunk # type: ignore[misc]
continue
if not all_chunks: if not all_chunks:
# All chunks were Anthropic native SSE bytes — output verbose_proxy_logger.warning(
# masking cannot be applied to raw bytes. Log a warning "Presidio apply_to_output: streaming response contained only "
# so operators know PII masking was skipped for this stream. "bytes chunks (Anthropic native SSE). Output PII masking was "
verbose_proxy_logger.warning( "skipped for this response."
"Presidio apply_to_output: streaming response contained only "
"bytes chunks (Anthropic native SSE). Output PII masking was "
"skipped for this response."
)
return
assembled_model_response = stream_chunk_builder(
chunks=all_chunks, messages=request_data.get("messages")
) )
if not isinstance(assembled_model_response, ModelResponse):
for chunk in all_chunks:
yield chunk
return
# Apply Presidio masking on the assembled response
await self._process_response_for_pii(
response=assembled_model_response,
request_data=request_data,
mode="mask",
)
mock_response_stream = convert_model_response_to_streaming(
assembled_model_response
)
yield mock_response_stream
return return
except Exception as e: assembled_model_response = stream_chunk_builder(
verbose_proxy_logger.error( chunks=all_chunks, messages=request_data.get("messages")
f"Error masking streaming PII output: {str(e)}" )
)
# Cannot re-iterate `response` — it's already consumed. if not isinstance(assembled_model_response, ModelResponse):
# If we collected chunks before the error, replay those.
for chunk in all_chunks: for chunk in all_chunks:
yield chunk yield chunk
return return
# --- PII unmasking path (output_parse_pii=True) --- await self._process_response_for_pii(
metadata = (request_data.get("metadata") or {}) if request_data else {} response=assembled_model_response,
pii_tokens = metadata.get("pii_tokens", {}) request_data=request_data,
if not pii_tokens and request_data: mode="mask",
verbose_proxy_logger.debug(
"No pii_tokens in request_data['metadata'] for streaming unmask path"
) )
if not (self.output_parse_pii and pii_tokens):
async for chunk in response: mock_response_stream = convert_model_response_to_streaming(
assembled_model_response
)
yield mock_response_stream
except Exception as e:
verbose_proxy_logger.error(
f"Error masking streaming PII output: {str(e)}"
)
for chunk in all_chunks:
yield chunk yield chunk
return
async def _stream_pii_unmasking(
self,
response: Any,
request_data: dict,
) -> AsyncGenerator[Union[ModelResponseStream, bytes], None]:
"""Apply PII unmasking to streaming output (output_parse_pii=True path)."""
from litellm.llms.base_llm.base_model_iterator import (
convert_model_response_to_streaming,
)
from litellm.main import stream_chunk_builder
from litellm.types.utils import ModelResponse
remaining_chunks: List[ModelResponseStream] = [] remaining_chunks: List[ModelResponseStream] = []
try: try:
@ -1210,7 +1196,6 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
if isinstance(chunk, ModelResponseStream): if isinstance(chunk, ModelResponseStream):
remaining_chunks.append(chunk) remaining_chunks.append(chunk)
elif isinstance(chunk, bytes): elif isinstance(chunk, bytes):
# Anthropic native SSE: pass through as-is
yield chunk # type: ignore[misc] yield chunk # type: ignore[misc]
continue continue
@ -1226,17 +1211,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
yield chunk yield chunk
return return
# --- PRESERVE USAGE METADATA --- self._preserve_usage_from_last_chunk(
# stream_chunk_builder might miss usage if it's only in the last chunk assembled_model_response, remaining_chunks
if ( )
not getattr(assembled_model_response, "usage", None)
) and remaining_chunks:
last_chunk = remaining_chunks[-1]
last_chunk_usage = getattr(last_chunk, "usage", None)
if last_chunk_usage:
setattr(assembled_model_response, "usage", last_chunk_usage)
# Apply PII unmasking to assembled content (unmasking tokens back to original text)
await self._process_response_for_pii( await self._process_response_for_pii(
response=assembled_model_response, response=assembled_model_response,
request_data=request_data, request_data=request_data,
@ -1253,6 +1231,47 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
for chunk in remaining_chunks: for chunk in remaining_chunks:
yield chunk yield chunk
async def async_post_call_streaming_iterator_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
response: Any,
request_data: dict,
) -> AsyncGenerator[Union[ModelResponseStream, bytes], None]:
"""
Process streaming response chunks to unmask PII tokens when needed.
"""
if self.apply_to_output:
async for chunk in self._stream_apply_output_masking(
response, request_data
):
yield chunk
return
metadata = (request_data.get("metadata") or {}) if request_data else {}
pii_tokens = metadata.get("pii_tokens", {})
if not pii_tokens and request_data:
verbose_proxy_logger.debug(
"No pii_tokens in request_data['metadata'] for streaming unmask path"
)
if not (self.output_parse_pii and pii_tokens):
async for chunk in response:
yield chunk
return
async for chunk in self._stream_pii_unmasking(response, request_data):
yield chunk
@staticmethod
def _preserve_usage_from_last_chunk(
assembled_model_response: Any,
chunks: List[Any],
) -> None:
"""Copy usage metadata from the last chunk when stream_chunk_builder misses it."""
if not getattr(assembled_model_response, "usage", None) and chunks:
last_chunk_usage = getattr(chunks[-1], "usage", None)
if last_chunk_usage:
setattr(assembled_model_response, "usage", last_chunk_usage)
def get_presidio_settings_from_request_data( def get_presidio_settings_from_request_data(
self, data: dict self, data: dict
) -> Optional[PresidioPerRequestConfig]: ) -> Optional[PresidioPerRequestConfig]:

View file

@ -1759,6 +1759,34 @@ async def _process_single_key_update(
return updated_key_info return updated_key_info
async def _validate_mcp_servers_for_key_update(
data: "UpdateKeyRequest",
team_obj: Optional["LiteLLM_TeamTableCachedObj"],
existing_key_row: Any,
prisma_client: Any,
user_api_key_cache: Any,
) -> None:
"""Validate MCP servers in object_permission against the effective team."""
effective_team_obj = team_obj
# If team_id isn't being changed, resolve the existing key's team
if effective_team_obj is None and existing_key_row.team_id:
effective_team_obj = await get_team_object(
team_id=existing_key_row.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
check_db_only=True,
)
object_permission_dict = (
data.object_permission.model_dump()
if hasattr(data.object_permission, "model_dump")
else data.object_permission
)
await validate_key_mcp_servers_against_team(
object_permission=object_permission_dict,
team_obj=effective_team_obj,
)
@router.post( @router.post(
"/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)] "/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
) )
@ -1959,23 +1987,12 @@ async def update_key_fn(
# Validate MCP servers in object_permission against the effective team # Validate MCP servers in object_permission against the effective team
if data.object_permission is not None: if data.object_permission is not None:
effective_team_obj = team_obj await _validate_mcp_servers_for_key_update(
# If team_id isn't being changed, resolve the existing key's team data=data,
if effective_team_obj is None and existing_key_row.team_id: team_obj=team_obj,
effective_team_obj = await get_team_object( existing_key_row=existing_key_row,
team_id=existing_key_row.team_id, prisma_client=prisma_client,
prisma_client=prisma_client, user_api_key_cache=user_api_key_cache,
user_api_key_cache=user_api_key_cache,
check_db_only=True,
)
object_permission_dict = (
data.object_permission.model_dump()
if hasattr(data.object_permission, "model_dump")
else data.object_permission
)
await validate_key_mcp_servers_against_team(
object_permission=object_permission_dict,
team_obj=effective_team_obj,
) )
non_default_values = await prepare_key_update_data( non_default_values = await prepare_key_update_data(
@ -4587,9 +4604,11 @@ async def _list_key_helper(
user_map = {} user_map = {}
if expand and "user" in expand: if expand and "user" in expand:
user_ids = [key.user_id for key in keys if key.user_id] user_ids = [key.user_id for key in keys if key.user_id]
if user_ids: created_by_ids = [key.created_by for key in keys if key.created_by]
all_ids = list(set(user_ids + created_by_ids)) # Remove duplicates
if all_ids:
users = await prisma_client.db.litellm_usertable.find_many( users = await prisma_client.db.litellm_usertable.find_many(
where={"user_id": {"in": list(set(user_ids))}} # Remove duplicates where={"user_id": {"in": all_ids}}
) )
user_map = {user.user_id: user for user in users} user_map = {user.user_id: user for user in users}
@ -4607,11 +4626,19 @@ async def _list_key_helper(
key_dict = await attach_object_permission_to_dict(key_dict, prisma_client) key_dict = await attach_object_permission_to_dict(key_dict, prisma_client)
# Include user information if expand includes "user" # Include user information if expand includes "user"
if expand and "user" in expand and key.user_id and key.user_id in user_map: if expand and "user" in expand:
try: if key.user_id and key.user_id in user_map:
key_dict["user"] = user_map[key.user_id].model_dump() try:
except Exception: key_dict["user"] = user_map[key.user_id].model_dump()
key_dict["user"] = user_map[key.user_id].dict() except Exception:
key_dict["user"] = user_map[key.user_id].dict()
if key.created_by and key.created_by in user_map:
created_by_user = user_map[key.created_by]
key_dict["created_by_user"] = {
"user_id": created_by_user.user_id,
"user_email": created_by_user.user_email,
"user_alias": created_by_user.user_alias,
}
if return_full_object is True or (expand and "user" in expand): if return_full_object is True or (expand and "user" in expand):
if use_deleted_table: if use_deleted_table:

View file

@ -36,6 +36,11 @@ from fastapi import (
) )
from fastapi.responses import JSONResponse from fastapi.responses import JSONResponse
try:
from prisma.errors import RecordNotFoundError
except ImportError:
RecordNotFoundError = Exception # type: ignore
import litellm import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._uuid import uuid from litellm._uuid import uuid
@ -84,9 +89,13 @@ if MCP_AVAILABLE:
delete_user_credential, delete_user_credential,
get_all_mcp_servers_for_user, get_all_mcp_servers_for_user,
get_mcp_server, get_mcp_server,
get_mcp_servers,
get_mcp_submissions, get_mcp_submissions,
get_user_oauth_credential,
list_user_oauth_credentials,
reject_mcp_server, reject_mcp_server,
store_user_credential, store_user_credential,
store_user_oauth_credential,
update_mcp_server, update_mcp_server,
) )
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
@ -106,7 +115,10 @@ if MCP_AVAILABLE:
LitellmUserRoles, LitellmUserRoles,
MakeMCPServersPublicRequest, MakeMCPServersPublicRequest,
MCPApprovalStatus, MCPApprovalStatus,
MCPOAuthUserCredentialRequest,
MCPOAuthUserCredentialStatus,
MCPSubmissionsSummary, MCPSubmissionsSummary,
MCPUserCredentialListItem,
MCPUserCredentialRequest, MCPUserCredentialRequest,
MCPUserCredentialResponse, MCPUserCredentialResponse,
NewMCPServerRequest, NewMCPServerRequest,
@ -1045,6 +1057,7 @@ if MCP_AVAILABLE:
response_model=LiteLLM_MCPServerTable, response_model=LiteLLM_MCPServerTable,
) )
async def fetch_mcp_server( async def fetch_mcp_server(
request: Request,
server_id: str, server_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
): ):
@ -1061,8 +1074,30 @@ if MCP_AVAILABLE:
"Database not connected. Connect a database to your proxy" "Database not connected. Connect a database to your proxy"
) )
# check to see if server exists for all users # check to see if server exists (DB first, then registry for config-based servers)
mcp_server = await get_mcp_server(prisma_client, server_id) mcp_server = await get_mcp_server(prisma_client, server_id)
from_db = mcp_server is not None
if mcp_server is None:
# Fallback: check registry (config-based servers) - list endpoint uses get_registry()
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
client_ip = IPAddressUtils.get_mcp_client_ip(request)
registry_server = global_mcp_server_manager.get_mcp_server_by_id(server_id)
if registry_server is not None and not global_mcp_server_manager._is_server_accessible_from_ip(
registry_server, client_ip
):
registry_server = None
if registry_server is None:
# Try lookup by server_name or alias (client may use display name in URL)
registry_server = global_mcp_server_manager.get_mcp_server_by_name(
server_id, client_ip=client_ip
)
if registry_server is not None:
mcp_server = global_mcp_server_manager._build_mcp_server_table(
registry_server
)
if mcp_server is None: if mcp_server is None:
raise HTTPException( raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND, status_code=status.HTTP_404_NOT_FOUND,
@ -1078,10 +1113,17 @@ if MCP_AVAILABLE:
if not is_admin_view: if not is_admin_view:
# Perform authz check BEFORE any health check (avoid side-effects for # Perform authz check BEFORE any health check (avoid side-effects for
# unauthorized callers). # unauthorized callers).
mcp_server_records = await get_all_mcp_servers_for_user( if from_db:
prisma_client, user_api_key_dict 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) )
exists = does_mcp_server_exist(mcp_server_records, server_id)
else:
# Registry/config server: use same access logic as list endpoint
allowed_server_ids = await global_mcp_server_manager.get_allowed_mcp_servers(
user_api_key_dict
)
exists = mcp_server.server_id in allowed_server_ids
if not exists: if not exists:
raise HTTPException( raise HTTPException(
@ -1095,7 +1137,8 @@ if MCP_AVAILABLE:
) )
# At this point caller is authorized to view the server. # At this point caller is authorized to view the server.
await global_mcp_server_manager.add_server(mcp_server) if from_db:
await global_mcp_server_manager.add_server(mcp_server)
# Perform health check on the server using server manager # Perform health check on the server using server manager
try: try:
@ -1269,10 +1312,21 @@ if MCP_AVAILABLE:
def _get_cached_temporary_mcp_server_or_404(server_id: str) -> MCPServer: def _get_cached_temporary_mcp_server_or_404(server_id: str) -> MCPServer:
server = get_cached_temporary_mcp_server(server_id) server = get_cached_temporary_mcp_server(server_id)
if server is None:
# Fall back to real DB/config server (e.g. for the user-side OAuth flow
# which calls these endpoints with a real server_id, not a temp session id).
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
server = (
global_mcp_server_manager.get_mcp_server_by_id(server_id)
or global_mcp_server_manager.get_mcp_server_by_name(server_id)
)
if server is None: if server is None:
raise HTTPException( raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND, status_code=status.HTTP_404_NOT_FOUND,
detail={"error": f"Temporary MCP server {server_id} not found"}, detail={"error": f"MCP server {server_id} not found"},
) )
return server return server
@ -1283,8 +1337,8 @@ if MCP_AVAILABLE:
async def mcp_authorize( async def mcp_authorize(
request: Request, request: Request,
server_id: str, server_id: str,
client_id: str, client_id: Optional[str] = None,
redirect_uri: str, redirect_uri: str = Query(...),
state: str = "", state: str = "",
code_challenge: Optional[str] = None, code_challenge: Optional[str] = None,
code_challenge_method: Optional[str] = None, code_challenge_method: Optional[str] = None,
@ -1292,10 +1346,23 @@ if MCP_AVAILABLE:
scope: Optional[str] = None, scope: Optional[str] = None,
): ):
mcp_server = _get_cached_temporary_mcp_server_or_404(server_id) mcp_server = _get_cached_temporary_mcp_server_or_404(server_id)
# Use the server's stored client_id when the caller doesn't supply one
resolved_client_id = mcp_server.client_id or client_id or ""
if not resolved_client_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={
"error": "missing_client_id",
"message": (
"No client_id available for this MCP server. "
"Either configure the server with a client_id or supply one in the request."
),
},
)
return await authorize_with_server( return await authorize_with_server(
request=request, request=request,
mcp_server=mcp_server, mcp_server=mcp_server,
client_id=client_id, client_id=resolved_client_id,
redirect_uri=redirect_uri, redirect_uri=redirect_uri,
state=state, state=state,
code_challenge=code_challenge, code_challenge=code_challenge,
@ -1314,18 +1381,30 @@ if MCP_AVAILABLE:
grant_type: str = Form(...), grant_type: str = Form(...),
code: Optional[str] = Form(None), code: Optional[str] = Form(None),
redirect_uri: Optional[str] = Form(None), redirect_uri: Optional[str] = Form(None),
client_id: str = Form(...), client_id: Optional[str] = Form(None),
client_secret: Optional[str] = Form(None), client_secret: Optional[str] = Form(None),
code_verifier: Optional[str] = Form(None), code_verifier: Optional[str] = Form(None),
): ):
mcp_server = _get_cached_temporary_mcp_server_or_404(server_id) mcp_server = _get_cached_temporary_mcp_server_or_404(server_id)
resolved_client_id = mcp_server.client_id or client_id or ""
if not resolved_client_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={
"error": "missing_client_id",
"message": (
"No client_id available for this MCP server. "
"Either configure the server with a client_id or supply one in the request."
),
},
)
return await exchange_token_with_server( return await exchange_token_with_server(
request=request, request=request,
mcp_server=mcp_server, mcp_server=mcp_server,
grant_type=grant_type, grant_type=grant_type,
code=code, code=code,
redirect_uri=redirect_uri, redirect_uri=redirect_uri,
client_id=client_id, client_id=resolved_client_id,
client_secret=client_secret, client_secret=client_secret,
code_verifier=code_verifier, code_verifier=code_verifier,
) )
@ -1484,7 +1563,7 @@ if MCP_AVAILABLE:
) )
try: try:
await delete_user_credential(prisma_client, user_id, server_id) await delete_user_credential(prisma_client, user_id, server_id)
except Exception: except RecordNotFoundError:
pass # Already deleted or didn't exist pass # Already deleted or didn't exist
from litellm.proxy._experimental.mcp_server.server import ( from litellm.proxy._experimental.mcp_server.server import (
_invalidate_byok_cred_cache, _invalidate_byok_cred_cache,
@ -1493,6 +1572,182 @@ if MCP_AVAILABLE:
_invalidate_byok_cred_cache(user_id, server_id) _invalidate_byok_cred_cache(user_id, server_id)
return MCPUserCredentialResponse(server_id=server_id, has_credential=False) return MCPUserCredentialResponse(server_id=server_id, has_credential=False)
# ── OAuth2 user-credential endpoints ──────────────────────────────────────
@router.post(
"/server/{server_id}/oauth-user-credential",
description="Store the calling user's OAuth2 token for an OpenAPI MCP server",
dependencies=[Depends(user_api_key_auth)],
response_model=MCPOAuthUserCredentialStatus,
)
@management_endpoint_wrapper
async def store_mcp_oauth_user_credential(
server_id: str,
payload: MCPOAuthUserCredentialRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""Persist the OAuth2 access token obtained by the calling user."""
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to your proxy"
)
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 {server_id} not found"},
)
user_id = user_api_key_dict.user_id or ""
if not user_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={"error": "User ID not found in token"},
)
await store_user_oauth_credential(
prisma_client,
user_id,
server_id,
payload.access_token,
refresh_token=payload.refresh_token,
expires_in=payload.expires_in,
scopes=payload.scopes,
)
# Read back the persisted record so the response reflects the stored
# expires_at rather than recomputing it here (which could diverge by
# milliseconds or if the storage logic ever adds a grace period).
stored = await get_user_oauth_credential(prisma_client, user_id, server_id)
expires_at: Optional[str] = stored.get("expires_at") if stored else None
return MCPOAuthUserCredentialStatus(
server_id=server_id,
has_credential=True,
expires_at=expires_at,
is_expired=False,
)
@router.delete(
"/server/{server_id}/oauth-user-credential",
description="Revoke the calling user's stored OAuth2 token for an MCP server",
dependencies=[Depends(user_api_key_auth)],
response_model=MCPOAuthUserCredentialStatus,
)
@management_endpoint_wrapper
async def delete_mcp_oauth_user_credential(
server_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""Revoke/delete the user's OAuth2 credential."""
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to your proxy"
)
user_id = user_api_key_dict.user_id or ""
if not user_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={"error": "User ID not found in token"},
)
# Only delete if the stored credential is actually an OAuth2 token.
# This prevents accidentally deleting a BYOK credential if one exists
# for the same (user_id, server_id) pair.
cred_to_delete = await get_user_oauth_credential(prisma_client, user_id, server_id)
if cred_to_delete is not None:
try:
await delete_user_credential(prisma_client, user_id, server_id)
except RecordNotFoundError:
pass # Already gone — treat as a successful delete
return MCPOAuthUserCredentialStatus(
server_id=server_id,
has_credential=False,
is_expired=False,
)
@router.get(
"/server/{server_id}/oauth-user-credential/status",
description="Check whether the calling user has a stored OAuth2 credential for this MCP server",
dependencies=[Depends(user_api_key_auth)],
response_model=MCPOAuthUserCredentialStatus,
)
@management_endpoint_wrapper
async def get_mcp_oauth_user_credential_status(
server_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""Return credential status (has_credential, expiry) without exposing the token."""
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to your proxy"
)
user_id = user_api_key_dict.user_id or ""
if not user_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={"error": "User ID not found in token"},
)
cred = await get_user_oauth_credential(prisma_client, user_id, server_id)
if cred is None:
return MCPOAuthUserCredentialStatus(
server_id=server_id, has_credential=False, is_expired=False
)
expires_at: Optional[str] = cred.get("expires_at")
is_expired = False
if expires_at:
try:
exp = datetime.fromisoformat(expires_at)
is_expired = exp < datetime.now(timezone.utc)
except Exception:
pass
return MCPOAuthUserCredentialStatus(
server_id=server_id,
has_credential=True,
expires_at=expires_at,
is_expired=is_expired,
connected_at=cred.get("connected_at"),
)
@router.get(
"/user-credentials",
description="List all OAuth2 MCP credentials stored for the calling user",
dependencies=[Depends(user_api_key_auth)],
response_model=List[MCPUserCredentialListItem],
)
@management_endpoint_wrapper
async def list_mcp_user_credentials(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""Return all servers the calling user has connected via OAuth2."""
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to your proxy"
)
user_id = user_api_key_dict.user_id or ""
if not user_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={"error": "User ID not found in token"},
)
oauth_creds = await list_user_oauth_credentials(prisma_client, user_id)
if not oauth_creds:
return []
# Fetch server metadata for display names — single batch query instead of N+1.
server_ids = [c["server_id"] for c in oauth_creds]
servers = {
srv.server_id: srv
for srv in await get_mcp_servers(prisma_client, server_ids)
}
items: List[MCPUserCredentialListItem] = []
for cred in oauth_creds:
sid = cred["server_id"]
srv = servers.get(sid)
expires_at: Optional[str] = cred.get("expires_at")
items.append(
MCPUserCredentialListItem(
server_id=sid,
server_name=getattr(srv, "server_name", None) if srv else None,
alias=getattr(srv, "alias", None) if srv else None,
credential_type="oauth2",
has_credential=True,
expires_at=expires_at, # always pass the raw timestamp; client computes expiry state
connected_at=cred.get("connected_at"),
)
)
return items
@router.put( @router.put(
"/server", "/server",
description="Allows deleting mcp serves in the db", description="Allows deleting mcp serves in the db",

View file

@ -26,7 +26,6 @@ from litellm.types.tool_management import (
ToolDetailResponse, ToolDetailResponse,
ToolInputPolicy, ToolInputPolicy,
ToolListResponse, ToolListResponse,
ToolOutputPolicy,
ToolPolicyOption, ToolPolicyOption,
ToolPolicyOptionsResponse, ToolPolicyOptionsResponse,
ToolPolicyUpdateRequest, ToolPolicyUpdateRequest,

View file

@ -4,7 +4,7 @@ organizations, teams, and keys.
""" """
import json import json
from typing import Dict, List, Optional, Set, Union from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
from fastapi import HTTPException, status from fastapi import HTTPException, status
@ -12,6 +12,12 @@ from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid from litellm._uuid import uuid
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy.utils import PrismaClient from litellm.proxy.utils import PrismaClient
if TYPE_CHECKING:
from litellm.proxy._types import (
LiteLLM_ObjectPermissionTable,
LiteLLM_TeamTableCachedObj,
)

View file

@ -33,8 +33,8 @@
"icon_url": "https://cdn.simpleicons.org/atlassian", "icon_url": "https://cdn.simpleicons.org/atlassian",
"category": "Developer Tools", "category": "Developer Tools",
"registry_url": "https://registry.modelcontextprotocol.io/servers/com.atlassian%2Fatlassian-mcp-server", "registry_url": "https://registry.modelcontextprotocol.io/servers/com.atlassian%2Fatlassian-mcp-server",
"transport": "sse", "transport": "http",
"url": "https://mcp.atlassian.com/v1/sse", "url": "https://mcp.atlassian.com/v1/mcp",
"env_vars": [] "env_vars": []
}, },
{ {

View file

@ -1133,6 +1133,7 @@ def create_pass_through_route(
fastapi_response: Response, fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
subpath: str = "", # captures sub-paths when include_subpath=True subpath: str = "", # captures sub-paths when include_subpath=True
custom_body: Optional[dict] = None, # caller-supplied body takes precedence over request-parsed body
): ):
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
InitPassThroughEndpointHelpers, InitPassThroughEndpointHelpers,
@ -1208,9 +1209,11 @@ def create_pass_through_route(
) )
if query_params: if query_params:
final_query_params.update(query_params) final_query_params.update(query_params)
# Use the body parsed from the raw request # Caller-supplied custom_body takes precedence over the request-parsed body
final_custom_body: Optional[dict] = None final_custom_body: Optional[dict] = None
if isinstance(custom_body_data, dict): if custom_body is not None:
final_custom_body = custom_body
elif isinstance(custom_body_data, dict):
final_custom_body = custom_body_data final_custom_body = custom_body_data
return await pass_through_request( # type: ignore return await pass_through_request( # type: ignore

View file

@ -267,6 +267,7 @@ model LiteLLM_ObjectPermissionTable {
vector_stores String[] @default([]) vector_stores String[] @default([])
agents String[] @default([]) agents String[] @default([])
agent_access_groups String[] @default([]) agent_access_groups String[] @default([])
models String[] @default([])
blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission
teams LiteLLM_TeamTable[] teams LiteLLM_TeamTable[]
projects LiteLLM_ProjectTable[] projects LiteLLM_ProjectTable[]
@ -397,9 +398,6 @@ model LiteLLM_VerificationToken {
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3
@@index([budget_reset_at, expires]) @@index([budget_reset_at, expires])
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (...) ORDER BY "public"."LiteLLM_VerificationToken"."key_alias" ASC
@@index([key_alias])
} }
model LiteLLM_JWTKeyMapping { model LiteLLM_JWTKeyMapping {
@ -565,9 +563,6 @@ model LiteLLM_SpendLogs {
@@index([startTime, request_id]) @@index([startTime, request_id])
@@index([end_user]) @@index([end_user])
@@index([session_id]) @@index([session_id])
// SELECT ... FROM "LiteLLM_SpendLogs" WHERE ("startTime" >= $1 AND "startTime" <= $2 AND "user" = $3) GROUP BY ...
@@index([user, startTime])
} }
// View spend, model, api_key per request // View spend, model, api_key per request

View file

@ -52,8 +52,8 @@ def _get_max_string_length_prompt_in_db() -> int:
return DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB return DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB
def _is_master_key(api_key: str, _master_key: Optional[str]) -> bool: def _is_master_key(api_key: Optional[str], _master_key: Optional[str]) -> bool:
if _master_key is None: if _master_key is None or api_key is None:
return False return False
## string comparison ## string comparison

View file

@ -451,6 +451,7 @@ async def _update_litellm_setting(
# Update the in-memory settings # Update the in-memory settings
in_memory_var = settings.model_dump(exclude_none=True) in_memory_var = settings.model_dump(exclude_none=True)
setattr(litellm, settings_key, in_memory_var)
# Load existing config # Load existing config
config = await proxy_config.get_config() config = await proxy_config.get_config()
@ -459,7 +460,7 @@ async def _update_litellm_setting(
if "litellm_settings" not in config: if "litellm_settings" not in config:
config["litellm_settings"] = {} config["litellm_settings"] = {}
config["litellm_settings"][settings_key] = settings.model_dump(exclude_none=True) config["litellm_settings"][settings_key] = in_memory_var
# Save the updated config # Save the updated config
await proxy_config.save_config(new_config=config) await proxy_config.save_config(new_config=config)

View file

@ -288,6 +288,268 @@ async def vector_store_create(
) )
@router.get("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
@router.get("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
async def vector_store_retrieve(
request: Request,
vector_store_id: str,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Retrieve a vector store.
API Reference:
https://platform.openai.com/docs/api-reference/vector-stores/retrieve
"""
from litellm.proxy.proxy_server import (
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
select_data_generator,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
data = {"vector_store_id": vector_store_id}
data = _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
)
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="avector_store_retrieve",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=None,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
@router.get("/v1/vector_stores", dependencies=[Depends(user_api_key_auth)])
@router.get("/vector_stores", dependencies=[Depends(user_api_key_auth)])
async def vector_store_list(
request: Request,
fastapi_response: Response,
after: Optional[str] = None,
before: Optional[str] = None,
limit: Optional[int] = 20,
order: Optional[str] = "desc",
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
List vector stores.
API Reference:
https://platform.openai.com/docs/api-reference/vector-stores/list
"""
from litellm.proxy.proxy_server import (
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
select_data_generator,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
data = {}
if after is not None:
data["after"] = after
if before is not None:
data["before"] = before
if limit is not None:
data["limit"] = limit
if order is not None:
data["order"] = order
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="avector_store_list",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=None,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
@router.post("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
@router.post("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
async def vector_store_update(
request: Request,
vector_store_id: str,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Update a vector store.
API Reference:
https://platform.openai.com/docs/api-reference/vector-stores/modify
"""
from litellm.proxy.proxy_server import (
_read_request_body,
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
select_data_generator,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
data = await _read_request_body(request=request)
if "vector_store_id" not in data:
data["vector_store_id"] = vector_store_id
data = _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
)
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="avector_store_update",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=None,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
@router.delete("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
@router.delete("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)])
async def vector_store_delete(
request: Request,
vector_store_id: str,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Delete a vector store.
API Reference:
https://platform.openai.com/docs/api-reference/vector-stores/delete
"""
from litellm.proxy.proxy_server import (
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
select_data_generator,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
data = {"vector_store_id": vector_store_id}
data = _update_request_data_with_litellm_managed_vector_store_registry(
data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
)
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="avector_store_delete",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=None,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
@router.post( @router.post(
"/v1/indexes", "/v1/indexes",
dependencies=[Depends(user_api_key_auth)], dependencies=[Depends(user_api_key_auth)],

View file

@ -24,6 +24,7 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i
LiteLLMResponsesTransformationHandler, LiteLLMResponsesTransformationHandler,
) )
from litellm.constants import request_timeout from litellm.constants import request_timeout
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.prompt_templates.common_utils import (
update_responses_input_with_model_file_ids, update_responses_input_with_model_file_ids,
@ -655,34 +656,37 @@ def responses(
# Native MCP Responses API # Native MCP Responses API
######################################################### #########################################################
if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools): if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools):
return aresponses_api_with_mcp( mcp_call_kwargs = {
input=input, "input": input,
model=model, "model": model,
include=include, "include": include,
instructions=instructions, "instructions": instructions,
max_output_tokens=max_output_tokens, "max_output_tokens": max_output_tokens,
prompt=prompt, "prompt": prompt,
metadata=metadata, "metadata": metadata,
parallel_tool_calls=parallel_tool_calls, "parallel_tool_calls": parallel_tool_calls,
previous_response_id=previous_response_id, "previous_response_id": previous_response_id,
reasoning=reasoning, "reasoning": reasoning,
store=store, "store": store,
background=background, "background": background,
stream=stream, "stream": stream,
temperature=temperature, "temperature": temperature,
text=text, "text": text,
tool_choice=tool_choice, "tool_choice": tool_choice,
tools=tools, "tools": tools,
top_p=top_p, "top_p": top_p,
truncation=truncation, "truncation": truncation,
user=user, "user": user,
extra_headers=extra_headers, "extra_headers": extra_headers,
extra_query=extra_query, "extra_query": extra_query,
extra_body=extra_body, "extra_body": extra_body,
timeout=timeout, "timeout": timeout,
custom_llm_provider=custom_llm_provider, "custom_llm_provider": custom_llm_provider,
**kwargs, **kwargs,
) }
if _is_async:
return aresponses_api_with_mcp(**mcp_call_kwargs)
return run_async_function(aresponses_api_with_mcp, **mcp_call_kwargs)
# get provider config # get provider config
responses_api_provider_config: Optional[ responses_api_provider_config: Optional[

View file

@ -1,3 +1,4 @@
import re
import traceback import traceback
from datetime import datetime from datetime import datetime
from typing import ( from typing import (
@ -43,6 +44,13 @@ ToolParam = Any
LITELLM_PROXY_MCP_SERVER_URL = "litellm_proxy" LITELLM_PROXY_MCP_SERVER_URL = "litellm_proxy"
LITELLM_PROXY_MCP_SERVER_URL_PREFIX = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" LITELLM_PROXY_MCP_SERVER_URL_PREFIX = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/"
# Matches any URL whose path ends with /mcp/<server_name> — covers both root-path
# (http://host:port/mcp/name) and sub-path (http://host/base/mcp/name) proxy deployments.
# A false-positive match (e.g. an external URL that happens to end with /mcp/<name>) results
# in a "server not found" error from the internal gateway, not a silent failure or data leak,
# so this broad pattern is intentional and preferred over anchoring to localhost only.
_PROXY_MCP_PATH_RE = re.compile(r"^https?://.+/mcp/([^/]+)$")
class LiteLLM_Proxy_MCP_Handler: class LiteLLM_Proxy_MCP_Handler:
""" """
@ -54,7 +62,8 @@ class LiteLLM_Proxy_MCP_Handler:
@staticmethod @staticmethod
def _should_use_litellm_mcp_gateway(tools: Optional[Iterable[ToolParam]]) -> bool: def _should_use_litellm_mcp_gateway(tools: Optional[Iterable[ToolParam]]) -> bool:
""" """
Returns True if the user passed a MCP tool with server_url="litellm_proxy" Returns True if any MCP tool should be handled via the litellm proxy MCP gateway.
This includes tools with server_url="litellm_proxy" as well as URLs ending in /mcp/<name>.
""" """
if tools: if tools:
for tool in tools: for tool in tools:
@ -64,6 +73,10 @@ class LiteLLM_Proxy_MCP_Handler:
LITELLM_PROXY_MCP_SERVER_URL LITELLM_PROXY_MCP_SERVER_URL
): ):
return True return True
if isinstance(server_url, str) and _PROXY_MCP_PATH_RE.match(
server_url
):
return True
return False return False
@staticmethod @staticmethod
@ -87,6 +100,18 @@ class LiteLLM_Proxy_MCP_Handler:
LITELLM_PROXY_MCP_SERVER_URL LITELLM_PROXY_MCP_SERVER_URL
): ):
mcp_tools_with_litellm_proxy.append(tool) mcp_tools_with_litellm_proxy.append(tool)
elif isinstance(server_url, str):
# Also intercept URLs like http://localhost:4000/mcp/atlassian_test
# by rewriting them to the internal litellm_proxy format.
m = _PROXY_MCP_PATH_RE.match(server_url)
if m:
rewritten = {
**tool,
"server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{m.group(1)}",
}
mcp_tools_with_litellm_proxy.append(rewritten)
else:
other_tools.append(tool)
else: else:
other_tools.append(tool) other_tools.append(tool)
else: else:

View file

@ -164,7 +164,11 @@ from litellm.types.utils import (
) )
from litellm.types.utils import ModelInfo from litellm.types.utils import ModelInfo
from litellm.types.utils import ModelInfo as ModelMapInfo from litellm.types.utils import ModelInfo as ModelMapInfo
from litellm.types.utils import ModelResponseStream, StandardLoggingPayload, Usage from litellm.types.utils import (
ModelResponseStream,
StandardLoggingPayload,
Usage,
)
from litellm.utils import ( from litellm.utils import (
CustomStreamWrapper, CustomStreamWrapper,
EmbeddingResponse, EmbeddingResponse,
@ -913,7 +917,19 @@ class Router:
def _initialize_vector_store_endpoints(self): def _initialize_vector_store_endpoints(self):
"""Initialize vector store endpoints.""" """Initialize vector store endpoints."""
from litellm.vector_stores.main import asearch, create, search from litellm.vector_stores.main import (
adelete,
alist,
aretrieve,
asearch,
aupdate,
create,
delete,
list,
retrieve,
search,
update,
)
self.avector_store_search = self.factory_function( self.avector_store_search = self.factory_function(
asearch, call_type="avector_store_search" asearch, call_type="avector_store_search"
@ -924,6 +940,30 @@ class Router:
self.vector_store_create = self.factory_function( self.vector_store_create = self.factory_function(
create, call_type="vector_store_create" create, call_type="vector_store_create"
) )
self.avector_store_retrieve = self.factory_function(
aretrieve, call_type="avector_store_retrieve"
)
self.vector_store_retrieve = self.factory_function(
retrieve, call_type="vector_store_retrieve"
)
self.avector_store_list = self.factory_function(
alist, call_type="avector_store_list"
)
self.vector_store_list = self.factory_function(
list, call_type="vector_store_list"
)
self.avector_store_update = self.factory_function(
aupdate, call_type="avector_store_update"
)
self.vector_store_update = self.factory_function(
update, call_type="vector_store_update"
)
self.avector_store_delete = self.factory_function(
adelete, call_type="avector_store_delete"
)
self.vector_store_delete = self.factory_function(
delete, call_type="vector_store_delete"
)
def _initialize_vector_store_file_endpoints(self): def _initialize_vector_store_file_endpoints(self):
"""Initialize vector store file endpoints.""" """Initialize vector store file endpoints."""
@ -4725,6 +4765,10 @@ class Router:
"generate_content_stream", "generate_content_stream",
"avector_store_search", "avector_store_search",
"avector_store_create", "avector_store_create",
"avector_store_retrieve",
"avector_store_list",
"avector_store_update",
"avector_store_delete",
"avector_store_file_create", "avector_store_file_create",
"avector_store_file_list", "avector_store_file_list",
"avector_store_file_retrieve", "avector_store_file_retrieve",
@ -4733,6 +4777,10 @@ class Router:
"avector_store_file_delete", "avector_store_file_delete",
"vector_store_search", "vector_store_search",
"vector_store_create", "vector_store_create",
"vector_store_retrieve",
"vector_store_list",
"vector_store_update",
"vector_store_delete",
"vector_store_file_create", "vector_store_file_create",
"vector_store_file_list", "vector_store_file_list",
"vector_store_file_retrieve", "vector_store_file_retrieve",
@ -4798,6 +4846,10 @@ class Router:
"generate_content_stream", "generate_content_stream",
"vector_store_search", "vector_store_search",
"vector_store_create", "vector_store_create",
"vector_store_retrieve",
"vector_store_list",
"vector_store_update",
"vector_store_delete",
"ocr", "ocr",
"search", "search",
"video_generation", "video_generation",
@ -4946,6 +4998,10 @@ class Router:
elif call_type in ( elif call_type in (
"avector_store_search", "avector_store_search",
"avector_store_create", "avector_store_create",
"avector_store_retrieve",
"avector_store_list",
"avector_store_update",
"avector_store_delete",
): ):
return await self._init_vector_store_api_endpoints( return await self._init_vector_store_api_endpoints(
original_function=original_function, original_function=original_function,

View file

@ -172,6 +172,8 @@ class AgentObjectPermission(TypedDict, total=False):
mcp_servers: Optional[List[str]] mcp_servers: Optional[List[str]]
mcp_access_groups: Optional[List[str]] mcp_access_groups: Optional[List[str]]
mcp_tool_permissions: Optional[Dict[str, List[str]]] mcp_tool_permissions: Optional[Dict[str, List[str]]]
models: Optional[List[str]]
agents: Optional[List[str]]
class AgentConfig(TypedDict, total=False): class AgentConfig(TypedDict, total=False):

View file

@ -479,3 +479,588 @@ def search(
completion_kwargs=local_vars, completion_kwargs=local_vars,
extra_kwargs=kwargs, extra_kwargs=kwargs,
) )
@client
async def aretrieve(
vector_store_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
**kwargs,
) -> VectorStoreCreateResponse:
"""
Async: Retrieve a vector store.
"""
local_vars = locals()
try:
loop = asyncio.get_event_loop()
kwargs["aretrieve"] = True
if custom_llm_provider is None:
custom_llm_provider = "openai"
func = partial(
retrieve,
vector_store_id=vector_store_id,
extra_headers=extra_headers,
extra_query=extra_query,
extra_body=extra_body,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
**kwargs,
)
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
init_response = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
response = init_response
return response
except Exception as e:
raise litellm.exception_type(
model=None,
custom_llm_provider=custom_llm_provider,
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)
@client
def retrieve(
vector_store_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
**kwargs,
) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]:
"""
Retrieve a vector store.
Args:
vector_store_id: The ID of the vector store to retrieve.
Returns:
VectorStoreCreateResponse containing the vector store details.
"""
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("aretrieve", False) is True
litellm_params = GenericLiteLLMParams(**kwargs)
if custom_llm_provider is None:
custom_llm_provider = "openai"
if "/" in custom_llm_provider:
api_type, custom_llm_provider, _, _ = get_llm_provider(
model=custom_llm_provider,
custom_llm_provider=None,
litellm_params=None,
)
else:
api_type = None
custom_llm_provider = custom_llm_provider
vector_store_provider_config = (
ProviderConfigManager.get_provider_vector_stores_config(
provider=litellm.LlmProviders(custom_llm_provider),
api_type=api_type,
)
)
if vector_store_provider_config is None:
raise ValueError(
f"Vector store retrieve is not supported for {custom_llm_provider}"
)
litellm_logging_obj.update_environment_variables(
model=None,
optional_params={"vector_store_id": vector_store_id},
litellm_params={"litellm_call_id": litellm_call_id},
custom_llm_provider=custom_llm_provider,
)
response = base_llm_http_handler.vector_store_retrieve_handler(
vector_store_id=vector_store_id,
vector_store_provider_config=vector_store_provider_config,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
extra_headers=extra_headers,
extra_body=extra_body,
timeout=timeout or request_timeout,
_is_async=_is_async,
client=kwargs.get("client"),
)
return response
except Exception as e:
raise litellm.exception_type(
model=None,
custom_llm_provider=custom_llm_provider,
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)
@client
async def alist(
after: Optional[str] = None,
before: Optional[str] = None,
limit: Optional[int] = 20,
order: Optional[str] = "desc",
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
**kwargs,
):
"""
Async: List vector stores.
"""
local_vars = locals()
try:
loop = asyncio.get_event_loop()
kwargs["alist"] = True
if custom_llm_provider is None:
custom_llm_provider = "openai"
func = partial(
list,
after=after,
before=before,
limit=limit,
order=order,
extra_headers=extra_headers,
extra_query=extra_query,
extra_body=extra_body,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
**kwargs,
)
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
init_response = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
response = init_response
return response
except Exception as e:
raise litellm.exception_type(
model=None,
custom_llm_provider=custom_llm_provider,
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)
@client
def list(
after: Optional[str] = None,
before: Optional[str] = None,
limit: Optional[int] = 20,
order: Optional[str] = "desc",
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
**kwargs,
):
"""
List vector stores.
Args:
after: A cursor for use in pagination.
before: A cursor for use in pagination.
limit: A limit on the number of objects to be returned.
order: Sort order by the created_at timestamp.
Returns:
List of vector stores.
"""
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("alist", False) is True
litellm_params = GenericLiteLLMParams(**kwargs)
if custom_llm_provider is None:
custom_llm_provider = "openai"
if "/" in custom_llm_provider:
api_type, custom_llm_provider, _, _ = get_llm_provider(
model=custom_llm_provider,
custom_llm_provider=None,
litellm_params=None,
)
else:
api_type = None
custom_llm_provider = custom_llm_provider
vector_store_provider_config = (
ProviderConfigManager.get_provider_vector_stores_config(
provider=litellm.LlmProviders(custom_llm_provider),
api_type=api_type,
)
)
if vector_store_provider_config is None:
raise ValueError(
f"Vector store list is not supported for {custom_llm_provider}"
)
litellm_logging_obj.update_environment_variables(
model=None,
optional_params={
"after": after,
"before": before,
"limit": limit,
"order": order,
},
litellm_params={"litellm_call_id": litellm_call_id},
custom_llm_provider=custom_llm_provider,
)
response = base_llm_http_handler.vector_store_list_handler(
after=after,
before=before,
limit=limit,
order=order,
vector_store_provider_config=vector_store_provider_config,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
extra_headers=extra_headers,
extra_body=extra_body,
timeout=timeout or request_timeout,
_is_async=_is_async,
client=kwargs.get("client"),
)
return response
except Exception as e:
raise litellm.exception_type(
model=None,
custom_llm_provider=custom_llm_provider,
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)
@client
async def aupdate(
vector_store_id: str,
name: Optional[str] = None,
expires_after: Optional[Dict] = None,
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
**kwargs,
) -> VectorStoreCreateResponse:
"""
Async: Update a vector store.
"""
local_vars = locals()
try:
loop = asyncio.get_event_loop()
kwargs["aupdate"] = True
if custom_llm_provider is None:
custom_llm_provider = "openai"
func = partial(
update,
vector_store_id=vector_store_id,
name=name,
expires_after=expires_after,
metadata=metadata,
extra_headers=extra_headers,
extra_query=extra_query,
extra_body=extra_body,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
**kwargs,
)
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
init_response = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
response = init_response
return response
except Exception as e:
raise litellm.exception_type(
model=None,
custom_llm_provider=custom_llm_provider,
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)
@client
def update(
vector_store_id: str,
name: Optional[str] = None,
expires_after: Optional[Dict] = None,
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
**kwargs,
) -> Union[VectorStoreCreateResponse, Coroutine[Any, Any, VectorStoreCreateResponse]]:
"""
Update a vector store.
Args:
vector_store_id: The ID of the vector store to update.
name: The name of the vector store.
expires_after: The expiration policy for the vector store.
metadata: Set of 16 key-value pairs that can be attached to an object.
Returns:
VectorStoreCreateResponse containing the updated vector store details.
"""
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("aupdate", False) is True
litellm_params = GenericLiteLLMParams(**kwargs)
if custom_llm_provider is None:
custom_llm_provider = "openai"
if "/" in custom_llm_provider:
api_type, custom_llm_provider, _, _ = get_llm_provider(
model=custom_llm_provider,
custom_llm_provider=None,
litellm_params=None,
)
else:
api_type = None
custom_llm_provider = custom_llm_provider
vector_store_provider_config = (
ProviderConfigManager.get_provider_vector_stores_config(
provider=litellm.LlmProviders(custom_llm_provider),
api_type=api_type,
)
)
if vector_store_provider_config is None:
raise ValueError(
f"Vector store update is not supported for {custom_llm_provider}"
)
local_vars.update(kwargs)
vector_store_update_optional_params: VectorStoreCreateOptionalRequestParams = (
VectorStoreRequestUtils.get_requested_vector_store_create_optional_param(
local_vars
)
)
litellm_logging_obj.update_environment_variables(
model=None,
optional_params={
"vector_store_id": vector_store_id,
"name": name,
**vector_store_update_optional_params,
},
litellm_params={"litellm_call_id": litellm_call_id},
custom_llm_provider=custom_llm_provider,
)
response = base_llm_http_handler.vector_store_update_handler(
vector_store_id=vector_store_id,
vector_store_update_optional_params=vector_store_update_optional_params,
vector_store_provider_config=vector_store_provider_config,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
extra_headers=extra_headers,
extra_body=extra_body,
timeout=timeout or request_timeout,
_is_async=_is_async,
client=kwargs.get("client"),
)
return response
except Exception as e:
raise litellm.exception_type(
model=None,
custom_llm_provider=custom_llm_provider,
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)
@client
async def adelete(
vector_store_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
**kwargs,
):
"""
Async: Delete a vector store.
"""
local_vars = locals()
try:
loop = asyncio.get_event_loop()
kwargs["adelete"] = True
if custom_llm_provider is None:
custom_llm_provider = "openai"
func = partial(
delete,
vector_store_id=vector_store_id,
extra_headers=extra_headers,
extra_query=extra_query,
extra_body=extra_body,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
**kwargs,
)
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
init_response = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
response = init_response
return response
except Exception as e:
raise litellm.exception_type(
model=None,
custom_llm_provider=custom_llm_provider,
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)
@client
def delete(
vector_store_id: str,
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
**kwargs,
):
"""
Delete a vector store.
Args:
vector_store_id: The ID of the vector store to delete.
Returns:
Deletion confirmation response.
"""
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("adelete", False) is True
litellm_params = GenericLiteLLMParams(**kwargs)
if custom_llm_provider is None:
custom_llm_provider = "openai"
if "/" in custom_llm_provider:
api_type, custom_llm_provider, _, _ = get_llm_provider(
model=custom_llm_provider,
custom_llm_provider=None,
litellm_params=None,
)
else:
api_type = None
custom_llm_provider = custom_llm_provider
vector_store_provider_config = (
ProviderConfigManager.get_provider_vector_stores_config(
provider=litellm.LlmProviders(custom_llm_provider),
api_type=api_type,
)
)
if vector_store_provider_config is None:
raise ValueError(
f"Vector store delete is not supported for {custom_llm_provider}"
)
litellm_logging_obj.update_environment_variables(
model=None,
optional_params={"vector_store_id": vector_store_id},
litellm_params={"litellm_call_id": litellm_call_id},
custom_llm_provider=custom_llm_provider,
)
response = base_llm_http_handler.vector_store_delete_handler(
vector_store_id=vector_store_id,
vector_store_provider_config=vector_store_provider_config,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
extra_headers=extra_headers,
extra_body=extra_body,
timeout=timeout or request_timeout,
_is_async=_is_async,
client=kwargs.get("client"),
)
return response
except Exception as e:
raise litellm.exception_type(
model=None,
custom_llm_provider=custom_llm_provider,
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)

View file

@ -14516,7 +14516,7 @@
"input_cost_per_audio_per_second": 0.00016, "input_cost_per_audio_per_second": 0.00016,
"input_cost_per_image": 0.00012, "input_cost_per_image": 0.00012,
"input_cost_per_token": 2e-07, "input_cost_per_token": 2e-07,
"input_cost_per_video_per_second": 0.0237, "input_cost_per_video_per_second": 0.00079,
"litellm_provider": "vertex_ai-embedding-models", "litellm_provider": "vertex_ai-embedding-models",
"max_input_tokens": 8192, "max_input_tokens": 8192,
"max_tokens": 8192, "max_tokens": 8192,
@ -14527,14 +14527,17 @@
"uses_embed_content": true "uses_embed_content": true
}, },
"vertex_ai/gemini-embedding-2-preview": { "vertex_ai/gemini-embedding-2-preview": {
"input_cost_per_token": 1.5e-07, "input_cost_per_audio_per_second": 0.00016,
"input_cost_per_image": 0.00012,
"input_cost_per_token": 2e-07,
"input_cost_per_video_per_second": 0.00079,
"litellm_provider": "vertex_ai", "litellm_provider": "vertex_ai",
"max_input_tokens": 8192, "max_input_tokens": 8192,
"max_tokens": 8192, "max_tokens": 8192,
"mode": "embedding", "mode": "embedding",
"output_cost_per_token": 0, "output_cost_per_token": 0,
"output_vector_size": 3072, "output_vector_size": 3072,
"source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
"supports_multimodal": true, "supports_multimodal": true,
"uses_embed_content": true "uses_embed_content": true
}, },
@ -14576,7 +14579,10 @@
"tpm": 10000000 "tpm": 10000000
}, },
"gemini/gemini-embedding-2-preview": { "gemini/gemini-embedding-2-preview": {
"input_cost_per_token": 1.5e-07, "input_cost_per_audio_per_second": 0.00016,
"input_cost_per_image": 0.00012,
"input_cost_per_token": 2e-07,
"input_cost_per_video_per_second": 0.00079,
"litellm_provider": "gemini", "litellm_provider": "gemini",
"max_input_tokens": 8192, "max_input_tokens": 8192,
"max_tokens": 8192, "max_tokens": 8192,
@ -14584,7 +14590,7 @@
"output_cost_per_token": 0, "output_cost_per_token": 0,
"output_vector_size": 3072, "output_vector_size": 3072,
"rpm": 10000, "rpm": 10000,
"source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", "source": "https://ai.google.dev/gemini-api/docs/pricing",
"supports_multimodal": true, "supports_multimodal": true,
"tpm": 10000000 "tpm": 10000000
}, },

8
poetry.lock generated
View file

@ -3222,15 +3222,15 @@ files = [
[[package]] [[package]]
name = "litellm-proxy-extras" name = "litellm-proxy-extras"
version = "0.4.53" version = "0.4.54"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
optional = true optional = true
python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8" python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8"
groups = ["main"] groups = ["main"]
markers = "extra == \"proxy\"" markers = "extra == \"proxy\""
files = [ files = [
{file = "litellm_proxy_extras-0.4.53-py3-none-any.whl", hash = "sha256:9224c667144774b6119e4de9b4b2d52fafc58442e6db317785c43b2d833665d6"}, {file = "litellm_proxy_extras-0.4.54-py3-none-any.whl", hash = "sha256:6621cf529f7f3647eb2dd0d2c417d91db8c7a05c3c592bef251887a122928837"},
{file = "litellm_proxy_extras-0.4.53.tar.gz", hash = "sha256:22c53fa8890d93d4a0d24171726e4e2bba8be6fef4838317cb74284fa9d27f70"}, {file = "litellm_proxy_extras-0.4.54.tar.gz", hash = "sha256:2c777ecdf39901c4007ade4466eb6398985ed4000afe3fc2cac997e1169e8cee"},
] ]
[[package]] [[package]]
@ -8002,4 +8002,4 @@ utils = ["numpydoc"]
[metadata] [metadata]
lock-version = "2.1" lock-version = "2.1"
python-versions = ">=3.9,<4.0" python-versions = ">=3.9,<4.0"
content-hash = "3036cfcdc06fb4293e248a2edd9c32a7afe6846920167527e247b2aefd74cfa6" content-hash = "5ed0af4e3644bc7b5a02b8bfc8b3eda15c014b43aa6da7a9a97a9b070fba5366"

View file

@ -2518,6 +2518,39 @@
"search": true, "search": true,
"a2a": false "a2a": false
} }
},
"black_forest_labs": {
"display_name": "Black Forest Labs (`black_forest_labs`)",
"url": "https://docs.litellm.ai/docs/providers/black_forest_labs",
"endpoints": {
"chat_completions": false,
"messages": false,
"responses": false,
"embeddings": false,
"image_generations": true,
"image_edits": true,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,
"batches": false,
"rerank": false
}
},
"charity_engine": {
"display_name": "Charity Engine (`charity_engine`)",
"url": "https://docs.litellm.ai/docs/providers/charity_engine",
"endpoints": {
"chat_completions": true,
"messages": true,
"responses": true,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,
"batches": false,
"rerank": false
}
} }
}, },
"endpoints": { "endpoints": {

View file

@ -61,7 +61,7 @@ boto3 = { version = "^1.40.76", optional = true }
redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"} redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"}
mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"} mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"}
a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"} a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"}
litellm-proxy-extras = {version = "^0.4.53", optional = true} litellm-proxy-extras = {version = "^0.4.54", optional = true}
rich = {version = "^13.7.1", optional = true} rich = {version = "^13.7.1", optional = true}
litellm-enterprise = {version = "^0.1.33", optional = true} litellm-enterprise = {version = "^0.1.33", optional = true}
diskcache = {version = "^5.6.1", optional = true} diskcache = {version = "^5.6.1", optional = true}

View file

@ -57,7 +57,7 @@ grpcio>=1.75.0; python_version >= "3.14"
sentry_sdk==2.21.0 # for sentry error handling sentry_sdk==2.21.0 # for sentry error handling
detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests
tzdata==2025.1 # IANA time zone database tzdata==2025.1 # IANA time zone database
litellm-proxy-extras==0.4.53 # for proxy extras - e.g. prisma migrations litellm-proxy-extras==0.4.54 # for proxy extras - e.g. prisma migrations
llm-sandbox==0.3.31 # for skill execution in sandbox llm-sandbox==0.3.31 # for skill execution in sandbox
### LITELLM PACKAGE DEPENDENCIES ### LITELLM PACKAGE DEPENDENCIES
python-dotenv==1.0.1 # for env python-dotenv==1.0.1 # for env

View file

@ -267,6 +267,7 @@ model LiteLLM_ObjectPermissionTable {
vector_stores String[] @default([]) vector_stores String[] @default([])
agents String[] @default([]) agents String[] @default([])
agent_access_groups String[] @default([]) agent_access_groups String[] @default([])
models String[] @default([])
blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission
teams LiteLLM_TeamTable[] teams LiteLLM_TeamTable[]
projects LiteLLM_ProjectTable[] projects LiteLLM_ProjectTable[]
@ -388,9 +389,6 @@ model LiteLLM_VerificationToken {
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3
@@index([budget_reset_at, expires]) @@index([budget_reset_at, expires])
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (...) ORDER BY "public"."LiteLLM_VerificationToken"."key_alias" ASC
@@index([key_alias])
} }
model LiteLLM_JWTKeyMapping { model LiteLLM_JWTKeyMapping {
@ -556,9 +554,6 @@ model LiteLLM_SpendLogs {
@@index([startTime, request_id]) @@index([startTime, request_id])
@@index([end_user]) @@index([end_user])
@@index([session_id]) @@index([session_id])
// SELECT ... FROM "LiteLLM_SpendLogs" WHERE ("startTime" >= $1 AND "startTime" <= $2 AND "user" = $3) GROUP BY ...
@@index([user, startTime])
} }
// View spend, model, api_key per request // View spend, model, api_key per request

View file

@ -93,6 +93,8 @@ class BaseImageGenTest(ABC):
except Exception as e: except Exception as e:
if "Your task failed as a result of our safety system." in str(e): if "Your task failed as a result of our safety system." in str(e):
pass pass
elif "ModelDeprecated" in str(e):
pass # Azure model deployment has been deprecated - skip
else: else:
pytest.fail(f"An exception occurred - {str(e)}") pytest.fail(f"An exception occurred - {str(e)}")

View file

@ -271,7 +271,7 @@ def test_trimming_should_not_change_original_messages():
assert messages == messages_copy assert messages == messages_copy
@pytest.mark.parametrize("model", ["gpt-4-0125-preview", "claude-3-opus-20240229"]) @pytest.mark.parametrize("model", ["gpt-4-0125-preview", "claude-sonnet-4-6"])
def test_trimming_with_model_cost_max_input_tokens(model): def test_trimming_with_model_cost_max_input_tokens(model):
messages = [ messages = [
{"role": "system", "content": "This is a normal system message"}, {"role": "system", "content": "This is a normal system message"},
@ -521,7 +521,7 @@ def test_function_to_dict():
("gpt-3.5-turbo", True), ("gpt-3.5-turbo", True),
("azure/gpt-4-1106-preview", True), ("azure/gpt-4-1106-preview", True),
("groq/gemma-7b-it", True), ("groq/gemma-7b-it", True),
("gemini/gemini-1.5-flash", True), ("gemini/gemini-2.5-flash", True),
], ],
) )
def test_supports_function_calling(model, expected_bool): def test_supports_function_calling(model, expected_bool):
@ -1062,8 +1062,8 @@ def test_parse_content_for_reasoning(content, expected_reasoning, expected_conte
@pytest.mark.parametrize( @pytest.mark.parametrize(
"model, expected_bool", "model, expected_bool",
[ [
("vertex_ai/gemini-1.5-pro", True), ("vertex_ai/gemini-2.5-pro", True),
("gemini/gemini-1.5-pro", True), ("gemini/gemini-2.5-pro", True),
("predibase/llama3-8b-instruct", True), ("predibase/llama3-8b-instruct", True),
("databricks/databricks-meta-llama-3-1-70b-instruct", True), ("databricks/databricks-meta-llama-3-1-70b-instruct", True),
("gpt-3.5-turbo", False), ("gpt-3.5-turbo", False),
@ -1074,7 +1074,7 @@ def test_supports_response_schema(model, expected_bool):
""" """
Unit tests for 'supports_response_schema' helper function. Unit tests for 'supports_response_schema' helper function.
Should be true for gemini-1.5-pro on google ai studio / vertex ai AND predibase models Should be true for gemini-2.5-pro on google ai studio / vertex ai AND predibase models
Should be false otherwise Should be false otherwise
""" """
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
@ -1093,7 +1093,7 @@ def test_supports_response_schema(model, expected_bool):
("gpt-3.5-turbo", True), ("gpt-3.5-turbo", True),
("gpt-4", True), ("gpt-4", True),
("command-nightly", False), ("command-nightly", False),
("gemini-pro", True), ("gemini-2.5-pro", True),
], ],
) )
def test_supports_function_calling_v2(model, expected_bool): def test_supports_function_calling_v2(model, expected_bool):
@ -1109,10 +1109,10 @@ def test_supports_function_calling_v2(model, expected_bool):
@pytest.mark.parametrize( @pytest.mark.parametrize(
"model, expected_bool", "model, expected_bool",
[ [
("gpt-4-vision-preview", True), ("gpt-4o", True),
("gpt-3.5-turbo", False), ("gpt-3.5-turbo", False),
("claude-3-opus-20240229", True), ("claude-sonnet-4-6", True),
("gemini-pro-vision", True), ("gemini-2.5-flash", True),
("command-nightly", False), ("command-nightly", False),
], ],
) )
@ -1727,7 +1727,7 @@ def test_supports_vision_gemini():
litellm.model_cost = litellm.get_model_cost_map(url="") litellm.model_cost = litellm.get_model_cost_map(url="")
from litellm.utils import supports_vision from litellm.utils import supports_vision
assert supports_vision("gemini-1.5-pro") is True assert supports_vision("gemini-2.5-pro") is True
def test_pick_cheapest_chat_model_from_llm_provider(): def test_pick_cheapest_chat_model_from_llm_provider():

View file

@ -271,7 +271,7 @@ def test_gemini_context_caching_separate_messages():
def test_gemini_image_generation(): def test_gemini_image_generation():
# litellm._turn_on_debug() # litellm._turn_on_debug()
response = completion( response = completion(
model="gemini/gemini-2.0-flash-exp-image-generation", model="gemini/gemini-2.5-flash-image-preview",
messages=[{"role": "user", "content": "Generate an image of a cat"}], messages=[{"role": "user", "content": "Generate an image of a cat"}],
modalities=["image", "text"], modalities=["image", "text"],
) )

View file

@ -1012,8 +1012,8 @@ def test_completion_cost_azure_common_deployment_name():
@pytest.mark.parametrize( @pytest.mark.parametrize(
"model, custom_llm_provider", "model, custom_llm_provider",
[ [
("claude-3-5-sonnet-20240620", "anthropic"), ("claude-sonnet-4-6", "anthropic"),
("gemini/gemini-1.5-flash-001", "gemini"), ("claude-haiku-4-5", "anthropic"),
], ],
) )
def test_completion_cost_prompt_caching(model, custom_llm_provider): def test_completion_cost_prompt_caching(model, custom_llm_provider):

View file

@ -1085,10 +1085,15 @@ def test_standard_logging_payload(model, turn_off_message_logging):
if turn_off_message_logging: if turn_off_message_logging:
print("checks redacted-by-litellm") print("checks redacted-by-litellm")
assert "redacted-by-litellm" == slobject["messages"][0]["content"] assert "redacted-by-litellm" == slobject["messages"][0]["content"]
# response is a full ModelResponse dict (choices format) since d84e5e381acf
response = slobject["response"] response = slobject["response"]
assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" if "choices" in response:
assert response["choices"][0]["message"].get("audio") is None assert (
response["choices"][0]["message"]["content"]
== "redacted-by-litellm"
)
assert response["choices"][0]["message"].get("audio") is None
else:
assert response["text"] == "redacted-by-litellm"
@pytest.mark.parametrize( @pytest.mark.parametrize(
@ -1188,10 +1193,15 @@ def test_standard_logging_payload_audio(turn_off_message_logging, stream):
if turn_off_message_logging: if turn_off_message_logging:
print("checks redacted-by-litellm") print("checks redacted-by-litellm")
assert "redacted-by-litellm" == slobject["messages"][0]["content"] assert "redacted-by-litellm" == slobject["messages"][0]["content"]
# response is a full ModelResponse dict (choices format) since d84e5e381acf
response = slobject["response"] response = slobject["response"]
assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" if "choices" in response:
assert response["choices"][0]["message"].get("audio") is None assert (
response["choices"][0]["message"]["content"]
== "redacted-by-litellm"
)
assert response["choices"][0]["message"].get("audio") is None
else:
assert response["text"] == "redacted-by-litellm"
@pytest.mark.skip(reason="Works locally. Flaky on ci/cd") @pytest.mark.skip(reason="Works locally. Flaky on ci/cd")

View file

@ -927,7 +927,7 @@ def test_anthropic_tool_calling_exception():
] ]
try: try:
litellm.completion( litellm.completion(
model="claude-3-5-sonnet-20240620", model="claude-haiku-4-5-20251001",
messages=[{"role": "user", "content": "Hey, how's it going?"}], messages=[{"role": "user", "content": "Hey, how's it going?"}],
tools=tools, tools=tools,
) )

View file

@ -1456,6 +1456,7 @@ async def test_add_update_server_with_alias():
mock_mcp_server.authorization_url = None mock_mcp_server.authorization_url = None
mock_mcp_server.registration_url = None mock_mcp_server.registration_url = None
mock_mcp_server.token_url = None mock_mcp_server.token_url = None
mock_mcp_server.oauth2_flow = None
# Additional fields used by build_mcp_server_from_table # Additional fields used by build_mcp_server_from_table
mock_mcp_server.extra_headers = None mock_mcp_server.extra_headers = None
mock_mcp_server.allow_all_keys = False mock_mcp_server.allow_all_keys = False
@ -1511,6 +1512,7 @@ async def test_add_update_server_without_alias():
mock_mcp_server.authorization_url = None mock_mcp_server.authorization_url = None
mock_mcp_server.registration_url = None mock_mcp_server.registration_url = None
mock_mcp_server.token_url = None mock_mcp_server.token_url = None
mock_mcp_server.oauth2_flow = None
# Additional fields used by build_mcp_server_from_table # Additional fields used by build_mcp_server_from_table
mock_mcp_server.extra_headers = None mock_mcp_server.extra_headers = None
mock_mcp_server.allow_all_keys = False mock_mcp_server.allow_all_keys = False
@ -1566,6 +1568,7 @@ async def test_add_update_server_fallback_to_server_id():
mock_mcp_server.authorization_url = None mock_mcp_server.authorization_url = None
mock_mcp_server.registration_url = None mock_mcp_server.registration_url = None
mock_mcp_server.token_url = None mock_mcp_server.token_url = None
mock_mcp_server.oauth2_flow = None
# Additional fields used by build_mcp_server_from_table - set explicitly # Additional fields used by build_mcp_server_from_table - set explicitly
# to avoid MagicMock objects being passed to Pydantic MCPServer constructor # to avoid MagicMock objects being passed to Pydantic MCPServer constructor
mock_mcp_server.extra_headers = None mock_mcp_server.extra_headers = None
@ -1823,6 +1826,7 @@ async def test_get_tools_for_single_server():
mock_manager._get_tools_from_server.assert_called_once_with( mock_manager._get_tools_from_server.assert_called_once_with(
server=mock_server, server=mock_server,
mcp_auth_header="Bearer test_token", mcp_auth_header="Bearer test_token",
extra_headers=None,
add_prefix=False, add_prefix=False,
raw_headers=None, raw_headers=None,
) )

View file

@ -42,7 +42,9 @@ from litellm.types.utils import CacheCreationTokenDetails, Usage
def test_reasoning_tokens_no_price_set(): def test_reasoning_tokens_no_price_set():
model = "o1-mini" # Use o1 - o1-mini was deprecated/renamed; o1 has same reasoning-token semantics
# (no separate output_cost_per_reasoning_token, so all completion tokens use output_cost_per_token)
model = "o1"
custom_llm_provider = "openai" custom_llm_provider = "openai"
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="") litellm.model_cost = litellm.get_model_cost_map(url="")
@ -266,7 +268,8 @@ def test_image_tokens_fallback_to_base_cost():
def test_generic_cost_per_token_above_200k_tokens(): def test_generic_cost_per_token_above_200k_tokens():
model = "gemini-2.5-pro-exp-03-25" # gemini-2.5-pro-exp-03-25 was removed; gemini-2.5-pro has same above-200k pricing
model = "gemini-2.5-pro"
custom_llm_provider = "vertex_ai" custom_llm_provider = "vertex_ai"
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="") litellm.model_cost = litellm.get_model_cost_map(url="")

View file

@ -121,17 +121,20 @@ def test_get_cost_for_built_in_tools_file_search():
def test_get_cost_for_anthropic_web_search(): def test_get_cost_for_anthropic_web_search():
""" """
Test that the cost for a web search is 0.00 when no response object is provided Test that Anthropic web search cost is tracked when usage.server_tool_use.web_search_requests
is set. Use claude-3-7-sonnet-20250219 (has search_context_cost_per_query) and
custom_llm_provider=anthropic so get_cost_for_anthropic_web_search is invoked.
""" """
from litellm.types.utils import ServerToolUse, Usage from litellm.types.utils import ServerToolUse, Usage
model = "claude-3-7-sonnet-latest" model = "claude-3-7-sonnet-20250219"
usage = Usage(server_tool_use=ServerToolUse(web_search_requests=1)) usage = Usage(server_tool_use=ServerToolUse(web_search_requests=1))
cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools(
model=model, model=model,
usage=usage, usage=usage,
response_object=None, response_object=None,
standard_built_in_tools_params=None, standard_built_in_tools_params=None,
custom_llm_provider="anthropic",
) )
assert cost > 0.0 assert cost > 0.0

View file

@ -219,6 +219,7 @@ class TestAgentCoreStreamingJsonFallback:
messages=[{"role": "user", "content": "test"}], messages=[{"role": "user", "content": "test"}],
stream=True, stream=True,
client=client, client=client,
api_key="test-jwt-token",
) )
# Collect content across all chunks # Collect content across all chunks
@ -257,6 +258,7 @@ class TestAgentCoreStreamingJsonFallback:
messages=[{"role": "user", "content": "test"}], messages=[{"role": "user", "content": "test"}],
stream=True, stream=True,
client=client, client=client,
api_key="test-jwt-token",
) )
# Collect content across all chunks # Collect content across all chunks
@ -289,6 +291,7 @@ class TestAgentCoreStreamingJsonFallback:
messages=[{"role": "user", "content": "test"}], messages=[{"role": "user", "content": "test"}],
stream=True, stream=True,
client=client, client=client,
api_key="test-jwt-token",
) )
async def test_async_streaming_malformed_json_raises_error(self): async def test_async_streaming_malformed_json_raises_error(self):
@ -316,4 +319,5 @@ class TestAgentCoreStreamingJsonFallback:
messages=[{"role": "user", "content": "test"}], messages=[{"role": "user", "content": "test"}],
stream=True, stream=True,
client=client, client=client,
api_key="test-jwt-token",
) )

View file

@ -2616,11 +2616,11 @@ def test_empty_assistant_message_handling():
empty or whitespace-only content with a placeholder to prevent AWS Bedrock empty or whitespace-only content with a placeholder to prevent AWS Bedrock
Converse API 400 Bad Request errors. Converse API 400 Bad Request errors.
""" """
# Import the litellm module that factory.py uses to ensure we patch the correct reference
import litellm.litellm_core_utils.prompt_templates.factory as factory_module
from litellm.litellm_core_utils.prompt_templates.factory import ( from litellm.litellm_core_utils.prompt_templates.factory import (
_bedrock_converse_messages_pt, _bedrock_converse_messages_pt,
) )
# Import the litellm module that factory.py uses to ensure we patch the correct reference
import litellm.litellm_core_utils.prompt_templates.factory as factory_module
# Test case 1: Empty string content - test with modify_params=True to prevent merging # Test case 1: Empty string content - test with modify_params=True to prevent merging
messages = [ messages = [
@ -3135,7 +3135,12 @@ def test_native_structured_output_no_fake_stream():
def test_transform_request_with_output_config(): def test_transform_request_with_output_config():
"""Test that outputConfig flows through _transform_request_helper into the final request.""" """Test that outputConfig flows through _transform_request_helper into the final request."""
from litellm.types.llms.bedrock import OutputConfigBlock, OutputFormat, OutputFormatStructure, JsonSchemaDefinition from litellm.types.llms.bedrock import (
JsonSchemaDefinition,
OutputConfigBlock,
OutputFormat,
OutputFormatStructure,
)
config = AmazonConverseConfig() config = AmazonConverseConfig()
@ -3170,6 +3175,29 @@ def test_transform_request_with_output_config():
assert result["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["name"] == "TestSchema" assert result["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["name"] == "TestSchema"
def test_transform_request_strips_anthropic_output_config():
"""
output_config is Anthropic-specific and must never be forwarded to Bedrock.
"""
config = AmazonConverseConfig()
messages = [{"role": "user", "content": "hello"}]
result = config._transform_request(
model="us.amazon.nova-pro-v1:0",
messages=messages,
optional_params={
"maxTokens": 64,
"output_config": {"effort": "low"},
},
litellm_params={},
headers={},
)
assert "outputConfig" not in result
additional_fields = result.get("additionalModelRequestFields", {})
assert "output_config" not in additional_fields
def test_transform_response_native_structured_output(): def test_transform_response_native_structured_output():
"""Test response handling when model returns JSON as text content (native structured output).""" """Test response handling when model returns JSON as text content (native structured output)."""
response_json = { response_json = {

View file

@ -107,12 +107,19 @@ class TestSnowflakeToolTransformation:
""" """
Test that string tool_choice values are transformed to Snowflake object format. Test that string tool_choice values are transformed to Snowflake object format.
Snowflake requires tool_choice to be an object, not a string. Snowflake's API (like Anthropic) requires tool_choice as an object
Ref: https://docs.snowflake.com/en/developer-guide/snowflake-rest-api/reference/cortex-inference#post--api-v2-cortex-inference-complete-req-body-schema with a "type" field, not as a bare string. OpenAI's "required" maps
to Snowflake's "any".
""" """
config = SnowflakeConfig() config = SnowflakeConfig()
for value in ["auto", "required", "none"]: expected_mappings = {
"auto": {"type": "auto"},
"required": {"type": "any"},
"none": {"type": "none"},
}
for value, expected in expected_mappings.items():
optional_params = {"tool_choice": value} optional_params = {"tool_choice": value}
transformed_request = config.transform_request( transformed_request = config.transform_request(
@ -123,8 +130,10 @@ class TestSnowflakeToolTransformation:
headers={}, headers={},
) )
# Snowflake requires object format: {"type": "auto"} not string "auto" assert transformed_request["tool_choice"] == expected, (
assert transformed_request["tool_choice"] == {"type": value} f"tool_choice='{value}' should be transformed to {expected}, "
f"got {transformed_request['tool_choice']}"
)
def test_transform_response_with_tool_calls(self): def test_transform_response_with_tool_calls(self):
""" """

View file

@ -11,7 +11,6 @@ sys.path.insert(
) # Adds the parent directory to the system path ) # Adds the parent directory to the system path
from litellm.llms.vertex_ai.common_utils import ( from litellm.llms.vertex_ai.common_utils import (
_build_vertex_schema_for_gemini_2,
_get_vertex_url, _get_vertex_url,
convert_anyof_null_to_nullable, convert_anyof_null_to_nullable,
get_vertex_location_from_url, get_vertex_location_from_url,
@ -1383,93 +1382,3 @@ def test_add_object_type_does_not_add_type_when_anyof_present():
# Verify type was not added (anyOf handles the type) # Verify type was not added (anyOf handles the type)
assert "type" not in input_schema, "type should not be added when anyOf is present" assert "type" not in input_schema, "type should not be added when anyOf is present"
class TestBuildVertexSchemaForGemini2:
"""Tests for _build_vertex_schema_for_gemini_2 — minimal transform for Gemini 2.0+ tools."""
def test_jsonvalue_standalone_preserved(self):
"""JsonValue (bare {}) should NOT be coerced to {"type": "object"}."""
schema = {
"type": "object",
"properties": {
"name": {"type": "string"},
"value": {},
},
"required": ["name", "value"],
}
result = _build_vertex_schema_for_gemini_2(schema)
assert result["properties"]["value"] == {}
def test_optional_jsonvalue_anyof_preserved(self):
"""Optional[JsonValue] anyOf with null should be preserved, not converted to nullable."""
schema = {
"type": "object",
"properties": {
"name": {"type": "string"},
"value": {
"anyOf": [
{"type": "array", "items": {}},
{},
{"type": "null"},
]
},
},
"required": ["name"],
}
result = _build_vertex_schema_for_gemini_2(schema)
value_schema = result["properties"]["value"]
assert "anyOf" in value_schema
assert len(value_schema["anyOf"]) == 3
assert {"type": "null"} in value_schema["anyOf"]
assert {} in value_schema["anyOf"]
def test_ref_defs_resolved(self):
"""$ref/$defs should be resolved since Gemini doesn't support them in tool params."""
schema = {
"type": "object",
"properties": {
"value": {"$ref": "#/$defs/JsonValue"},
},
"$defs": {"JsonValue": {}},
}
result = _build_vertex_schema_for_gemini_2(schema)
assert "$ref" not in result["properties"]["value"]
assert "$defs" not in result
assert result["properties"]["value"] == {}
def test_unsupported_fields_stripped(self):
"""Fields not in Vertex Schema TypedDict should be removed."""
schema = {
"type": "object",
"properties": {
"name": {"type": "string", "additionalProperties": False},
},
"additionalProperties": False,
"$schema": "http://json-schema.org/draft-07/schema#",
}
result = _build_vertex_schema_for_gemini_2(schema)
assert "additionalProperties" not in result
assert "$schema" not in result
def test_no_type_coercion(self):
"""Schemas without type should NOT have type: object added."""
schema = {
"type": "object",
"properties": {
"data": {"description": "Any data"},
},
}
result = _build_vertex_schema_for_gemini_2(schema)
assert "type" not in result["properties"]["data"]
def test_items_empty_preserved(self):
"""items: {} should NOT be coerced to items: {"type": "object"}."""
schema = {
"type": "object",
"properties": {
"values": {"type": "array", "items": {}},
},
}
result = _build_vertex_schema_for_gemini_2(schema)
assert result["properties"]["values"]["items"] == {}

View file

@ -2093,3 +2093,83 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab
assert spend_meta["tool_count_total"] == 1 assert spend_meta["tool_count_total"] == 1
assert spend_meta["allowed_server_count"] == 1 assert spend_meta["allowed_server_count"] == 1
assert spend_meta["per_server_tool_counts"]["server_a"] == 1 assert spend_meta["per_server_tool_counts"]["server_a"] == 1
@pytest.mark.asyncio
async def test_get_tools_from_mcp_servers_injects_stored_oauth2_token():
"""
When _get_tools_from_mcp_servers is called for an OAuth2 MCP server and no
oauth2_headers are provided in the request (e.g. a /responses API call from a
chat UI), the per-user stored token must be fetched from the DB and passed as
extra_headers to _get_tools_from_server.
The implementation pre-fetches all user credentials in a single bulk query
(_prefetch_oauth_creds_for_user) to avoid N+1 queries in the gather loop.
This covers the bug where OAuth2 MCP tools were always empty in the /responses
API because the stored credential was never injected.
"""
try:
from litellm.proxy._experimental.mcp_server.server import (
_get_tools_from_mcp_servers,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp import MCPAuth
except ImportError:
pytest.skip("MCP server not available")
STORED_TOKEN = "atlassian-oauth-access-token-xyz"
SERVER_ID = "srv-oauth2-id"
USER_ID = "user-123"
user_auth = UserAPIKeyAuth(api_key="test-key", user_id=USER_ID)
oauth2_server = MagicMock(name="atlassian_server")
oauth2_server.name = "atlassian_test"
oauth2_server.alias = "atlassian_test"
oauth2_server.server_name = "atlassian_test"
oauth2_server.server_id = SERVER_ID
oauth2_server.auth_type = MCPAuth.oauth2
oauth2_server.extra_headers = None
# Simulate the DB returning a valid credential for this user+server
prefetched_creds = {SERVER_ID: {"access_token": STORED_TOKEN, "server_id": SERVER_ID}}
tool_1 = MagicMock()
tool_1.name = "atlassian_test-search"
with patch(
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
new=AsyncMock(return_value=[oauth2_server]),
), patch(
# Patch the bulk prefetch so no real DB connection is needed
"litellm.proxy._experimental.mcp_server.server._prefetch_oauth_creds_for_user",
new=AsyncMock(return_value=prefetched_creds),
) as mock_prefetch, patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
) as mock_manager, patch(
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools",
side_effect=lambda tools, _server: tools,
), patch(
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions",
new=AsyncMock(side_effect=lambda tools, **_: tools),
):
mock_manager._get_tools_from_server = AsyncMock(return_value=[tool_1])
tools = await _get_tools_from_mcp_servers(
user_api_key_auth=user_auth,
mcp_auth_header=None,
mcp_servers=["atlassian_test"],
mcp_server_auth_headers=None,
oauth2_headers=None, # No token from request — must fall back to DB
)
# Bulk credential prefetch was called once (not once per server)
mock_prefetch.assert_awaited_once_with(user_auth)
# The stored token was forwarded to the MCP transport layer as extra_headers
mock_manager._get_tools_from_server.assert_awaited_once()
call_kwargs = mock_manager._get_tools_from_server.await_args.kwargs
assert call_kwargs["extra_headers"] == {"Authorization": f"Bearer {STORED_TOKEN}"}
assert tools == [tool_1]

View file

@ -484,7 +484,7 @@ class TestListToolsRestAPI:
captured = {"called": False} captured = {"called": False}
async def fake_get_tools( async def fake_get_tools(
server, server_auth_header, raw_headers=None, user_api_key_auth=None server, server_auth_header, raw_headers=None, user_api_key_auth=None, extra_headers=None
): ):
captured["called"] = True captured["called"] = True
captured["server"] = server captured["server"] = server
@ -529,6 +529,175 @@ class TestListToolsRestAPI:
assert result["error"] is None assert result["error"] is None
assert result["message"] == "Successfully retrieved tools" assert result["message"] == "Successfully retrieved tools"
async def test_name_resolution_finds_server_by_uuid(self, monkeypatch):
"""When server_id is a name string, it should be resolved to its UUID
and used for the tools lookup when the UUID is in allowed_server_ids."""
from litellm.proxy._experimental.mcp_server.server import MCPServer
from litellm.types.mcp import MCPTransport
stub_server = MCPServer(
server_id="uuid-abc-123",
name="my-server",
transport=MCPTransport.sse,
)
stub_server.alias = "my-server"
stub_server.server_name = "my-server"
stub_server.available_on_public_internet = True
stub_server.allowed_tools = None
stub_server.mcp_info = {"server_name": "my-server"}
async def fake_contexts(user_api_key_auth):
return [user_api_key_auth]
# Allowed list contains the UUID, not the name
async def fake_get_allowed_mcp_servers(*args, **kwargs):
return ["uuid-abc-123"]
captured = {"called": False, "server_arg": None}
async def fake_get_tools(server, server_auth_header, raw_headers=None, user_api_key_auth=None, extra_headers=None):
captured["called"] = True
captured["server_arg"] = server
return ["tool-x"]
monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager, "get_allowed_mcp_servers",
fake_get_allowed_mcp_servers, raising=False,
)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_name",
lambda name: stub_server if name == "my-server" else None,
raising=False,
)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_id",
lambda sid: stub_server if sid == "uuid-abc-123" else None,
raising=False,
)
monkeypatch.setattr(rest_endpoints, "_get_tools_for_single_server", fake_get_tools, raising=False)
request = _build_request(path="/mcp-rest/tools/list", method="GET")
result = await rest_endpoints.list_tool_rest_api(
request,
server_id="my-server", # pass name, not UUID
user_api_key_dict=UserAPIKeyAuth(),
)
assert captured["called"] is True
assert captured["server_arg"] is stub_server
assert result["tools"] == ["tool-x"]
assert result["error"] is None
async def test_name_not_in_allowed_returns_access_denied(self, monkeypatch):
"""When name resolves to a server whose UUID is NOT in allowed_server_ids,
the result should be an access_denied error (not a crash or silent pass)."""
from litellm.proxy._experimental.mcp_server.server import MCPServer
from litellm.types.mcp import MCPTransport
stub_server = MCPServer(
server_id="uuid-xyz-999",
name="restricted-server",
transport=MCPTransport.sse,
)
stub_server.available_on_public_internet = True
async def fake_contexts(user_api_key_auth):
return [user_api_key_auth]
# No allowed servers for this key
async def fake_get_allowed_mcp_servers(*args, **kwargs):
return []
monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager, "get_allowed_mcp_servers",
fake_get_allowed_mcp_servers, raising=False,
)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_name",
lambda name: stub_server if name == "restricted-server" else None,
raising=False,
)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_id",
lambda sid: stub_server if sid == "uuid-xyz-999" else None,
raising=False,
)
request = _build_request(path="/mcp-rest/tools/list", method="GET")
result = await rest_endpoints.list_tool_rest_api(
request,
server_id="restricted-server",
user_api_key_dict=UserAPIKeyAuth(),
)
assert result["tools"] == []
assert result["error"] == "unexpected_error"
assert "access_denied" in result["message"]
async def test_oauth2_user_token_injected_for_single_server(self, monkeypatch):
"""For a single-server OAuth2 request, _get_user_oauth_extra_headers is called
and the returned headers are forwarded to _get_tools_for_single_server."""
from litellm.proxy._experimental.mcp_server.server import MCPServer
from litellm.types.mcp import MCPTransport
stub_server = MCPServer(
server_id="oauth-server-id",
name="oauth-server",
transport=MCPTransport.sse,
)
stub_server.alias = "oauth-server"
stub_server.server_name = "oauth-server"
stub_server.available_on_public_internet = True
stub_server.allowed_tools = None
stub_server.mcp_info = {"server_name": "oauth-server"}
stub_server.auth_type = MCPAuth.oauth2
async def fake_contexts(user_api_key_auth):
return [user_api_key_auth]
async def fake_get_allowed_mcp_servers(*args, **kwargs):
return ["oauth-server-id"]
oauth_headers = {"Authorization": "Bearer user-oauth-token"}
async def fake_get_user_oauth_extra_headers(server, user_api_key_dict, prefetched_creds=None):
return oauth_headers
captured = {}
async def fake_get_tools(server, server_auth_header, raw_headers=None, user_api_key_auth=None, extra_headers=None):
captured["server"] = server
captured["auth_header"] = server_auth_header
return ["oauth-tool"]
monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager, "get_allowed_mcp_servers",
fake_get_allowed_mcp_servers, raising=False,
)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager, "get_mcp_server_by_id",
lambda sid: stub_server if sid == "oauth-server-id" else None,
raising=False,
)
monkeypatch.setattr(
rest_endpoints, "_get_user_oauth_extra_headers",
fake_get_user_oauth_extra_headers, raising=False,
)
monkeypatch.setattr(rest_endpoints, "_get_tools_for_single_server", fake_get_tools, raising=False)
request = _build_request(path="/mcp-rest/tools/list", method="GET")
result = await rest_endpoints.list_tool_rest_api(
request,
server_id="oauth-server-id",
user_api_key_dict=UserAPIKeyAuth(user_id="user-123"),
)
assert result["tools"] == ["oauth-tool"]
assert result["error"] is None
class TestCallToolRestAPI: class TestCallToolRestAPI:
pytestmark = pytest.mark.asyncio pytestmark = pytest.mark.asyncio

View file

@ -976,6 +976,43 @@ def test_proxy_admin_viewer_can_access_global_spend_tags():
) )
@pytest.mark.parametrize("route", ["/audit", "/audit/some-log-id"])
def test_proxy_admin_viewer_can_access_audit_logs(route):
"""
Test that proxy_admin_viewer can access /audit endpoints.
Admin viewers should be able to view audit logs since these are read-only.
"""
user_obj = LiteLLM_UserTable(
user_id="viewer_user",
user_email="viewer@example.com",
user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
)
valid_token = UserAPIKeyAuth(
user_id="viewer_user",
user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
)
request = MagicMock(spec=Request)
request.query_params = {}
try:
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
route=route,
request=request,
valid_token=valid_token,
request_data={},
)
except Exception as e:
pytest.fail(
f"proxy_admin_viewer should be able to access {route} route. Got error: {str(e)}"
)
class TestModelsRouteExemptFromDisableLLMEndpoints: class TestModelsRouteExemptFromDisableLLMEndpoints:
""" """
Test that /models and /v1/models are exempt from DISABLE_LLM_API_ENDPOINTS. Test that /models and /v1/models are exempt from DISABLE_LLM_API_ENDPOINTS.

View file

@ -175,6 +175,46 @@ def test_ui_discovery_endpoints_both_routes_return_same_data():
assert response1.json() == response2.json() assert response1.json() == response2.json()
def test_ui_discovery_endpoints_with_auto_redirect_via_general_settings():
"""When auto_redirect_ui_login_to_sso is set in general_settings (config.yaml), it should be honored."""
app = FastAPI()
app.include_router(router)
client = TestClient(app)
with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \
patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \
patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=True), \
patch("litellm.proxy.proxy_server.general_settings", {"auto_redirect_ui_login_to_sso": True}), \
patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False):
os.environ.pop("AUTO_REDIRECT_UI_LOGIN_TO_SSO", None)
response = client.get("/.well-known/litellm-ui-config")
assert response.status_code == 200
data = response.json()
assert data["auto_redirect_to_sso"] is True
assert data["sso_configured"] is True
def test_ui_discovery_endpoints_with_auto_redirect_env_var_overrides_general_settings():
"""Env var and general_settings should both work — either being true enables the feature."""
app = FastAPI()
app.include_router(router)
client = TestClient(app)
with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \
patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \
patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=True), \
patch("litellm.proxy.proxy_server.general_settings", {"auto_redirect_ui_login_to_sso": False}), \
patch.dict(os.environ, {"AUTO_REDIRECT_UI_LOGIN_TO_SSO": "true", "DISABLE_ADMIN_UI": "false"}, clear=False):
response = client.get("/.well-known/litellm-ui-config")
assert response.status_code == 200
data = response.json()
assert data["auto_redirect_to_sso"] is True
def test_ui_discovery_endpoints_with_admin_ui_disabled(): def test_ui_discovery_endpoints_with_admin_ui_disabled():
app = FastAPI() app = FastAPI()
app.include_router(router) app.include_router(router)

View file

@ -155,6 +155,103 @@ async def test_create_user_uses_default_internal_user_params_role(mocker, monkey
assert called_args.user_role == LitellmUserRoles.PROXY_ADMIN assert called_args.user_role == LitellmUserRoles.PROXY_ADMIN
@pytest.mark.asyncio
async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeypatch):
"""
Default user role set to 'Internal User' via UI,
but SCIM-created users get 'Internal Viewer' instead.
The UI saves the setting via _update_litellm_setting, which should update
litellm.default_internal_user_params in memory. Then SCIM create_user
should read that in-memory value and assign the correct role.
This test simulates the full flow:
1. Start with default_internal_user_params = None
2. Call update_internal_user_settings (the UI endpoint) to set role to INTERNAL_USER
3. Create a user via SCIM
4. Assert the user gets INTERNAL_USER (not INTERNAL_USER_VIEW_ONLY)
"""
from litellm.proxy._types import DefaultInternalUserParams
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
_update_litellm_setting,
)
# Step 1: Start with no default params (fresh proxy state)
monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False)
# Step 2: Simulate the UI saving "Internal User (Create/Delete/View)" as default role
# Mock the proxy_config and store_model_in_db that _update_litellm_setting needs
mock_proxy_config = mocker.MagicMock()
mock_proxy_config.get_config = AsyncMock(return_value={"litellm_settings": {}})
mock_proxy_config.save_config = AsyncMock()
mocker.patch(
"litellm.proxy.proxy_server.proxy_config",
mock_proxy_config,
)
mocker.patch(
"litellm.proxy.proxy_server.store_model_in_db",
True,
)
import litellm
settings = DefaultInternalUserParams(
user_role=LitellmUserRoles.INTERNAL_USER,
)
await _update_litellm_setting(
settings=settings,
settings_key="default_internal_user_params",
in_memory_var=litellm.default_internal_user_params,
success_message="ok",
)
# Verify the in-memory variable was actually updated
assert litellm.default_internal_user_params is not None, (
"BUG: _update_litellm_setting did not update litellm.default_internal_user_params in memory. "
"The local variable reassignment (in_memory_var = ...) doesn't propagate back."
)
assert litellm.default_internal_user_params.get("user_role") == LitellmUserRoles.INTERNAL_USER
# Step 3: Create a user via SCIM
scim_user = SCIMUser(
schemas=["urn:ietf:params:scim:schemas:core:2.0:User"],
userName="idontexist@krakentest.tech",
emails=[SCIMUserEmail(value="idontexist@krakentest.tech")],
)
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
mocker.patch(
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
AsyncMock(return_value=mock_prisma_client),
)
new_user_mock = mocker.patch(
"litellm.proxy.management_endpoints.scim.scim_v2.new_user",
AsyncMock(return_value=NewUserRequest(user_id="idontexist@krakentest.tech")),
)
mocker.patch(
"litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user",
AsyncMock(return_value=scim_user),
)
await create_user(user=scim_user)
# Step 4: Verify the user got INTERNAL_USER, not INTERNAL_USER_VIEW_ONLY
called_args = new_user_mock.call_args.kwargs["data"]
assert called_args.user_role == LitellmUserRoles.INTERNAL_USER, (
f"BUG: SCIM created user with role {called_args.user_role} instead of "
f"{LitellmUserRoles.INTERNAL_USER}. The default_internal_user_params "
f"in-memory variable was not updated by _update_litellm_setting."
)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_handle_existing_user_by_email_no_email(mocker): async def test_handle_existing_user_by_email_no_email(mocker):
"""Should return None when new_user_request has no email""" """Should return None when new_user_request has no email"""

View file

@ -4003,6 +4003,7 @@ async def test_list_keys_with_expand_user():
mock_key1 = MagicMock() mock_key1 = MagicMock()
mock_key1.token = "token1" mock_key1.token = "token1"
mock_key1.user_id = "user123" mock_key1.user_id = "user123"
mock_key1.created_by = None
# Set up model_dump() to raise AttributeError so it falls back to dict() # Set up model_dump() to raise AttributeError so it falls back to dict()
mock_key1.model_dump = MagicMock(side_effect=AttributeError("model_dump not available")) mock_key1.model_dump = MagicMock(side_effect=AttributeError("model_dump not available"))
mock_key1.dict = MagicMock(return_value=key1_dict) mock_key1.dict = MagicMock(return_value=key1_dict)
@ -4016,6 +4017,7 @@ async def test_list_keys_with_expand_user():
mock_key2 = MagicMock() mock_key2 = MagicMock()
mock_key2.token = "token2" mock_key2.token = "token2"
mock_key2.user_id = "user456" mock_key2.user_id = "user456"
mock_key2.created_by = None
# Set up model_dump() to raise AttributeError so it falls back to dict() # Set up model_dump() to raise AttributeError so it falls back to dict()
mock_key2.model_dump = MagicMock(side_effect=AttributeError("model_dump not available")) mock_key2.model_dump = MagicMock(side_effect=AttributeError("model_dump not available"))
mock_key2.dict = MagicMock(return_value=key2_dict) mock_key2.dict = MagicMock(return_value=key2_dict)
@ -4120,6 +4122,98 @@ async def test_list_keys_with_expand_user():
} }
@pytest.mark.asyncio
async def test_list_keys_with_expand_user_includes_created_by_user():
"""
Test that expand=user also resolves created_by to a user object.
"""
mock_prisma_client = AsyncMock()
# Key created by user789 but owned by user123
key1_dict = {
"token": "token1",
"user_id": "user123",
"created_by": "user789",
"key_alias": "key1",
"models": ["gpt-4"],
}
mock_key1 = MagicMock()
mock_key1.token = "token1"
mock_key1.user_id = "user123"
mock_key1.created_by = "user789"
mock_key1.model_dump = MagicMock(return_value=key1_dict)
mock_find_many_keys = AsyncMock(return_value=[mock_key1])
mock_count_keys = AsyncMock(return_value=1)
# Create mock users for both user_id and created_by
mock_user_owner = MagicMock()
mock_user_owner.user_id = "user123"
mock_user_owner.user_email = "owner@example.com"
mock_user_owner.user_alias = "Owner"
mock_user_owner.model_dump = MagicMock(return_value={
"user_id": "user123",
"user_email": "owner@example.com",
"user_alias": "Owner",
})
mock_user_creator = MagicMock()
mock_user_creator.user_id = "user789"
mock_user_creator.user_email = "creator@example.com"
mock_user_creator.user_alias = "Creator"
mock_find_many_users = AsyncMock(return_value=[mock_user_owner, mock_user_creator])
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
async def mock_attach_object_permission(d, _):
return d
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.attach_object_permission_to_dict",
side_effect=mock_attach_object_permission,
):
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,
"admin_team_ids": None,
"include_created_by_keys": False,
"expand": ["user"],
}
result = await _list_key_helper(**args)
# Verify that the user lookup included both user_id and created_by
call_args = mock_find_many_users.call_args
user_ids_in_query = set(call_args.kwargs["where"]["user_id"]["in"])
assert user_ids_in_query == {"user123", "user789"}
# Verify created_by_user is attached
key_result = result["keys"][0]
assert key_result.created_by_user == {
"user_id": "user789",
"user_email": "creator@example.com",
"user_alias": "Creator",
}
# Verify user (owner) is also still attached
assert key_result.user == {
"user_id": "user123",
"user_email": "owner@example.com",
"user_alias": "Owner",
}
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_list_keys_with_status_deleted(): async def test_list_keys_with_status_deleted():
""" """

View file

@ -76,6 +76,15 @@ def generate_mock_mcp_server_config_record(
) )
def _make_mock_request(ip: str = "127.0.0.1"):
"""Create a mock Request for fetch_mcp_server tests (IP used for access control)."""
req = MagicMock()
req.client = MagicMock()
req.client.host = ip
req.headers = {}
return req
def generate_mock_user_api_key_auth( def generate_mock_user_api_key_auth(
user_role: LitellmUserRoles = LitellmUserRoles.PROXY_ADMIN, user_role: LitellmUserRoles = LitellmUserRoles.PROXY_ADMIN,
user_id: str = "test_user_id", user_id: str = "test_user_id",
@ -735,7 +744,9 @@ class TestListMCPServers:
) )
result = await fetch_mcp_server( result = await fetch_mcp_server(
server_id="server-1", user_api_key_dict=mock_user_auth request=_make_mock_request(),
server_id="server-1",
user_api_key_dict=mock_user_auth,
) )
assert result.server_id == "server-1" assert result.server_id == "server-1"
@ -788,7 +799,9 @@ class TestListMCPServers:
) )
result = await fetch_mcp_server( result = await fetch_mcp_server(
server_id="server-2", user_api_key_dict=mock_user_auth request=_make_mock_request(),
server_id="server-2",
user_api_key_dict=mock_user_auth,
) )
assert result.server_id == "server-2" assert result.server_id == "server-2"
@ -796,6 +809,290 @@ class TestListMCPServers:
assert not hasattr(result, "credentials") assert not hasattr(result, "credentials")
assert result.status == "healthy" assert result.status == "healthy"
@pytest.mark.asyncio
async def test_fetch_single_mcp_server_from_registry_config_based(self):
"""
Test that fetch_mcp_server finds config-based servers when not in DB.
Config servers appear in list via get_registry() but were 404 on fetch.
"""
config_server = generate_mock_mcp_server_config_record(
server_id="serper_custom_dev",
name="Serper MCP",
url="https://serper.example.com/mcp",
transport="http",
)
mock_health_result = generate_mock_mcp_server_db_record(
server_id="serper_custom_dev", alias="Serper MCP"
)
mock_health_result.status = "healthy"
mock_health_result.last_health_check = datetime.now()
mock_health_result.health_check_error = None
mock_manager = MagicMock()
mock_manager.get_mcp_server_by_id = MagicMock(
side_effect=lambda sid: config_server if sid == "serper_custom_dev" else None
)
mock_manager.get_mcp_server_by_name = MagicMock(return_value=None)
mock_manager._build_mcp_server_table = MagicMock(
return_value=generate_mock_mcp_server_db_record(
server_id="serper_custom_dev",
alias="Serper MCP",
url="https://serper.example.com/mcp",
transport="http",
)
)
mock_manager.get_allowed_mcp_servers = AsyncMock(
return_value=["serper_custom_dev"]
)
mock_manager.health_check_server = AsyncMock(return_value=mock_health_result)
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=MagicMock(),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
AsyncMock(return_value=None),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view",
return_value=True,
),
):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
fetch_mcp_server,
)
result = await fetch_mcp_server(
request=_make_mock_request(),
server_id="serper_custom_dev",
user_api_key_dict=mock_user_auth,
)
assert result.server_id == "serper_custom_dev"
assert result.status == "healthy"
mock_manager.get_mcp_server_by_id.assert_called_with("serper_custom_dev")
mock_manager._build_mcp_server_table.assert_called_once()
@pytest.mark.asyncio
async def test_fetch_single_mcp_server_from_registry_by_name_passes_client_ip(self):
"""
When lookup by server_id fails, fallback to get_mcp_server_by_name.
Verify client_ip is passed for IP-based access control (security).
"""
config_server = generate_mock_mcp_server_config_record(
server_id="serper_custom_dev",
name="Serper MCP",
url="https://serper.example.com/mcp",
transport="http",
)
mock_manager = MagicMock()
mock_manager.get_mcp_server_by_id = MagicMock(return_value=None)
mock_manager.get_mcp_server_by_name = MagicMock(return_value=config_server)
mock_manager._build_mcp_server_table = MagicMock(
return_value=generate_mock_mcp_server_db_record(
server_id="serper_custom_dev",
alias="Serper MCP",
url="https://serper.example.com/mcp",
transport="http",
)
)
mock_manager.get_allowed_mcp_servers = AsyncMock(
return_value=["serper_custom_dev"]
)
mock_manager.health_check_server = AsyncMock(
return_value=generate_mock_mcp_server_db_record(
server_id="serper_custom_dev", alias="Serper MCP"
)
)
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=MagicMock(),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
AsyncMock(return_value=None),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view",
return_value=True,
),
):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
fetch_mcp_server,
)
result = await fetch_mcp_server(
request=_make_mock_request(ip="192.168.1.100"),
server_id="Serper MCP",
user_api_key_dict=mock_user_auth,
)
assert result.server_id == "serper_custom_dev"
mock_manager.get_mcp_server_by_id.assert_called_with("Serper MCP")
mock_manager.get_mcp_server_by_name.assert_called_once_with(
"Serper MCP", client_ip="192.168.1.100"
)
@pytest.mark.asyncio
async def test_fetch_single_mcp_server_from_registry_non_admin_denied(self):
"""
Non-admin user: config server NOT in allowed_server_ids -> 403.
"""
config_server = generate_mock_mcp_server_config_record(
server_id="restricted_server",
name="Restricted MCP",
url="https://restricted.example.com/mcp",
transport="http",
)
mock_manager = MagicMock()
mock_manager.get_mcp_server_by_id = MagicMock(
side_effect=lambda sid: config_server if sid == "restricted_server" else None
)
mock_manager.get_mcp_server_by_name = MagicMock(return_value=None)
mock_manager._build_mcp_server_table = MagicMock(
return_value=generate_mock_mcp_server_db_record(
server_id="restricted_server",
alias="Restricted MCP",
url="https://restricted.example.com/mcp",
transport="http",
)
)
mock_manager.get_allowed_mcp_servers = AsyncMock(
return_value=["other_server"] # restricted_server NOT in list
)
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=MagicMock(),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
AsyncMock(return_value=None),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view",
return_value=False,
),
):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
fetch_mcp_server,
)
with pytest.raises(HTTPException) as exc_info:
await fetch_mcp_server(
request=_make_mock_request(),
server_id="restricted_server",
user_api_key_dict=mock_user_auth,
)
assert exc_info.value.status_code == 403
mock_manager.get_allowed_mcp_servers.assert_called_once_with(mock_user_auth)
@pytest.mark.asyncio
async def test_fetch_single_mcp_server_from_registry_non_admin_granted(self):
"""
Non-admin user: config server IS in allowed_server_ids -> 200.
"""
config_server = generate_mock_mcp_server_config_record(
server_id="allowed_config_server",
name="Allowed MCP",
url="https://allowed.example.com/mcp",
transport="http",
)
mock_health_result = generate_mock_mcp_server_db_record(
server_id="allowed_config_server", alias="Allowed MCP"
)
mock_health_result.status = "healthy"
mock_health_result.last_health_check = datetime.now()
mock_health_result.health_check_error = None
mock_manager = MagicMock()
mock_manager.get_mcp_server_by_id = MagicMock(
side_effect=lambda sid: config_server if sid == "allowed_config_server" else None
)
mock_manager.get_mcp_server_by_name = MagicMock(return_value=None)
mock_manager._build_mcp_server_table = MagicMock(
return_value=generate_mock_mcp_server_db_record(
server_id="allowed_config_server",
alias="Allowed MCP",
url="https://allowed.example.com/mcp",
transport="http",
)
)
mock_manager.get_allowed_mcp_servers = AsyncMock(
return_value=["allowed_config_server"]
)
mock_manager.health_check_server = AsyncMock(return_value=mock_health_result)
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=MagicMock(),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
AsyncMock(return_value=None),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view",
return_value=False,
),
):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
fetch_mcp_server,
)
result = await fetch_mcp_server(
request=_make_mock_request(),
server_id="allowed_config_server",
user_api_key_dict=mock_user_auth,
)
assert result.server_id == "allowed_config_server"
assert result.status == "healthy"
mock_manager.get_allowed_mcp_servers.assert_called_once_with(mock_user_auth)
class TestTeamScopedMCPServerAccess: class TestTeamScopedMCPServerAccess:
"""Tests for cross-team information disclosure and restricted key bypass fixes.""" """Tests for cross-team information disclosure and restricted key bypass fixes."""
@ -961,6 +1258,11 @@ class TestTemporaryMCPSessionEndpoints:
existing_server.client_id = "client-123" existing_server.client_id = "client-123"
existing_server.client_secret = "secret-xyz" existing_server.client_secret = "secret-xyz"
existing_server.scopes = ["scope:a", "scope:b"] existing_server.scopes = ["scope:a", "scope:b"]
existing_server.aws_access_key_id = None
existing_server.aws_secret_access_key = None
existing_server.aws_session_token = None
existing_server.aws_region_name = None
existing_server.aws_service_name = None
mock_manager = MagicMock() mock_manager = MagicMock()
mock_manager.get_mcp_server_by_id.return_value = existing_server mock_manager.get_mcp_server_by_id.return_value = existing_server
@ -1068,6 +1370,11 @@ class TestTemporaryMCPSessionEndpoints:
client_id="client-id", client_id="client-id",
client_secret="client-secret", client_secret="client-secret",
scopes=["scope1"], scopes=["scope1"],
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_region_name=None,
aws_service_name=None,
) )
built_server = generate_mock_mcp_server_config_record(server_id="temp-server") built_server = generate_mock_mcp_server_config_record(server_id="temp-server")
mock_manager = MagicMock() mock_manager = MagicMock()
@ -2005,3 +2312,183 @@ class TestValidateMCPRequiredFields:
_validate_mcp_required_fields(payload) _validate_mcp_required_fields(payload)
assert exc_info.value.status_code == 500 assert exc_info.value.status_code == 500
assert "source_Url" in str(exc_info.value.detail) assert "source_Url" in str(exc_info.value.detail)
# ── OAuth user credential endpoint unit tests ──────────────────────────────────
def _make_user_auth(user_id: str = "user-abc") -> "UserAPIKeyAuth":
return UserAPIKeyAuth(
api_key="sk-test",
user_id=user_id,
user_role=LitellmUserRoles.INTERNAL_USER,
)
def _make_prisma_client():
"""Return a minimal mock PrismaClient accepted by get_prisma_client_or_throw."""
client = MagicMock()
client.db = MagicMock()
return client
@pytest.mark.asyncio
async def test_store_mcp_oauth_user_credential_returns_status():
"""store_mcp_oauth_user_credential persists the token and echoes back status."""
from litellm.proxy._types import (
MCPOAuthUserCredentialRequest,
MCPOAuthUserCredentialStatus,
)
if not mgmt_endpoints.MCP_AVAILABLE:
pytest.skip("MCP module not installed")
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
store_mcp_oauth_user_credential,
)
server_id = "srv-1"
user_id = "user-123"
stored_payload = {
"type": "oauth2",
"access_token": "tok",
"expires_at": "2099-01-01T00:00:00+00:00",
"connected_at": "2026-01-01T00:00:00+00:00",
"server_id": server_id,
}
mock_prisma = _make_prisma_client()
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=mock_prisma,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
new=AsyncMock(return_value=generate_mock_mcp_server_db_record(server_id=server_id)),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.store_user_oauth_credential",
new=AsyncMock(return_value=None),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_user_oauth_credential",
new=AsyncMock(return_value=stored_payload),
),
):
result = await store_mcp_oauth_user_credential(
server_id=server_id,
payload=MCPOAuthUserCredentialRequest(
access_token="tok",
expires_in=3600,
),
user_api_key_dict=_make_user_auth(user_id),
)
assert isinstance(result, MCPOAuthUserCredentialStatus)
assert result.has_credential is True
assert result.server_id == server_id
# expires_at should come from the stored record, not be recomputed
assert result.expires_at == "2099-01-01T00:00:00+00:00"
@pytest.mark.asyncio
async def test_delete_mcp_oauth_user_credential_only_deletes_oauth():
"""delete_mcp_oauth_user_credential only deletes OAuth2 credentials, not BYOK."""
from litellm.proxy._types import MCPOAuthUserCredentialStatus
if not mgmt_endpoints.MCP_AVAILABLE:
pytest.skip("MCP module not installed")
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
delete_mcp_oauth_user_credential,
)
server_id = "srv-2"
user_id = "user-456"
delete_mock = AsyncMock(return_value=None)
# When get_user_oauth_credential returns None (no OAuth cred), delete should NOT be called.
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_user_oauth_credential",
new=AsyncMock(return_value=None),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.delete_user_credential",
new=delete_mock,
),
):
result = await delete_mcp_oauth_user_credential(
server_id=server_id,
user_api_key_dict=_make_user_auth(user_id),
)
delete_mock.assert_not_called()
assert isinstance(result, MCPOAuthUserCredentialStatus)
assert result.has_credential is False
@pytest.mark.asyncio
async def test_list_mcp_user_credentials_batch_server_fetch():
"""list_mcp_user_credentials uses a single batch DB call, not N+1 queries."""
from litellm.proxy._types import MCPUserCredentialListItem
if not mgmt_endpoints.MCP_AVAILABLE:
pytest.skip("MCP module not installed")
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
list_mcp_user_credentials,
)
user_id = "user-789"
server_id = "srv-3"
stored_creds = [
{
"type": "oauth2",
"access_token": "tok",
"expires_at": "2099-01-01T00:00:00+00:00",
"connected_at": "2026-01-01T00:00:00+00:00",
"server_id": server_id,
}
]
mock_server = generate_mock_mcp_server_db_record(server_id=server_id, alias="My Server")
# get_mcp_servers (batch) should be called once; get_mcp_server (single) must not be called.
batch_mock = AsyncMock(return_value=[mock_server])
single_mock = AsyncMock(return_value=mock_server)
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.list_user_oauth_credentials",
new=AsyncMock(return_value=stored_creds),
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_servers",
new=batch_mock,
),
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
new=single_mock,
),
):
result = await list_mcp_user_credentials(
user_api_key_dict=_make_user_auth(user_id),
)
batch_mock.assert_called_once()
single_mock.assert_not_called()
assert len(result) == 1
assert isinstance(result[0], MCPUserCredentialListItem)
assert result[0].server_id == server_id
assert result[0].alias == "My Server"
# expires_at should always be the raw timestamp (not set to None when expired)
assert result[0].expires_at == "2099-01-01T00:00:00+00:00"

View file

@ -242,7 +242,11 @@ class TestGeminiPassthroughLoggingHandler:
assert "kwargs" in result assert "kwargs" in result
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_pass_through_success_handler_gemini_routing(self): @patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.gemini_passthrough_logging_handler.litellm.completion_cost",
return_value=0.000050,
)
async def test_pass_through_success_handler_gemini_routing(self, mock_completion_cost):
"""Test that the success handler correctly routes Gemini requests to the Gemini handler""" """Test that the success handler correctly routes Gemini requests to the Gemini handler"""
handler = PassThroughEndpointLogging() handler = PassThroughEndpointLogging()
@ -263,7 +267,7 @@ class TestGeminiPassthroughLoggingHandler:
httpx_response=mock_response, httpx_response=mock_response,
response_body=self.mock_gemini_response, response_body=self.mock_gemini_response,
logging_obj=mock_logging_obj, logging_obj=mock_logging_obj,
url_route="https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent", url_route="https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent",
result="", result="",
start_time=self.start_time, start_time=self.start_time,
end_time=self.end_time, end_time=self.end_time,
@ -277,15 +281,15 @@ class TestGeminiPassthroughLoggingHandler:
assert result is None assert result is None
# Verify that the logging object has the cost set (from Gemini handler) # Verify that the logging object has the cost set (from Gemini handler)
assert mock_logging_obj.model_call_details["response_cost"] is not None assert mock_logging_obj.model_call_details["response_cost"] == 0.000050
assert mock_logging_obj.model_call_details["model"] == "gemini-1.5-flash" assert mock_logging_obj.model_call_details["model"] == "gemini-2.0-flash"
assert mock_logging_obj.model_call_details["custom_llm_provider"] == "gemini" assert mock_logging_obj.model_call_details["custom_llm_provider"] == "gemini"
# Verify that _handle_logging was called with the correct kwargs # Verify that _handle_logging was called with the correct kwargs
handler._handle_logging.assert_called_once() handler._handle_logging.assert_called_once()
call_kwargs = handler._handle_logging.call_args[1] call_kwargs = handler._handle_logging.call_args[1]
assert call_kwargs["response_cost"] is not None assert call_kwargs["response_cost"] == 0.000050
assert call_kwargs["model"] == "gemini-1.5-flash" assert call_kwargs["model"] == "gemini-2.0-flash"
assert call_kwargs["custom_llm_provider"] == "gemini" assert call_kwargs["custom_llm_provider"] == "gemini"
@patch("litellm.completion_cost") @patch("litellm.completion_cost")

View file

@ -243,7 +243,7 @@ class TestVertexAIBatchPassthroughHandler:
] ]
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
vertex_ai_batch_responses, model_name="gemini-1.5-flash-001" vertex_ai_batch_responses, model_name="gemini-2.0-flash-001"
) )
assert usage.total_tokens == 15 assert usage.total_tokens == 15

View file

@ -29,6 +29,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
_get_response_for_spend_logs_payload, _get_response_for_spend_logs_payload,
_get_spend_logs_metadata, _get_spend_logs_metadata,
_get_vector_store_request_for_spend_logs_payload, _get_vector_store_request_for_spend_logs_payload,
_is_master_key,
_sanitize_request_body_for_spend_logs_payload, _sanitize_request_body_for_spend_logs_payload,
_should_store_prompts_and_responses_in_spend_logs, _should_store_prompts_and_responses_in_spend_logs,
get_logging_payload, get_logging_payload,
@ -1440,3 +1441,30 @@ def test_get_logging_payload_includes_request_duration_ms():
) )
assert payload["request_duration_ms"] == 3000 assert payload["request_duration_ms"] == 3000
class TestIsMasterKey:
"""Tests for _is_master_key handling None inputs without raising TypeError."""
def test_none_api_key_returns_false(self):
"""Regression: _is_master_key(None, 'sk-master') should return False, not raise TypeError."""
assert _is_master_key(api_key=None, _master_key="sk-master-key") is False
def test_none_master_key_returns_false(self):
assert _is_master_key(api_key="sk-some-key", _master_key=None) is False
def test_both_none_returns_false(self):
assert _is_master_key(api_key=None, _master_key=None) is False
def test_matching_key_returns_true(self):
assert _is_master_key(api_key="sk-master", _master_key="sk-master") is True
def test_non_matching_key_returns_false(self):
assert _is_master_key(api_key="sk-other", _master_key="sk-master") is False
def test_hashed_key_returns_true(self):
from litellm.proxy.utils import hash_token
master = "sk-master-key-123"
hashed = hash_token(master)
assert _is_master_key(api_key=hashed, _master_key=master) is True

View file

@ -970,12 +970,12 @@ def test_cost_discount_vertex_ai():
# Save original config # Save original config
original_discount_config = litellm.cost_discount_config.copy() original_discount_config = litellm.cost_discount_config.copy()
# Create mock response # Create mock response (use a model that exists in model_prices_and_context_window.json)
response = ModelResponse( response = ModelResponse(
id="test-id", id="test-id",
choices=[], choices=[],
created=1234567890, created=1234567890,
model="gemini-pro", model="gemini-3-pro-preview",
object="chat.completion", object="chat.completion",
usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
) )
@ -984,7 +984,7 @@ def test_cost_discount_vertex_ai():
litellm.cost_discount_config = {} litellm.cost_discount_config = {}
cost_without_discount = completion_cost( cost_without_discount = completion_cost(
completion_response=response, completion_response=response,
model="vertex_ai/gemini-pro", model="vertex_ai/gemini-3-pro-preview",
custom_llm_provider="vertex_ai", custom_llm_provider="vertex_ai",
) )
@ -994,7 +994,7 @@ def test_cost_discount_vertex_ai():
# Calculate cost with discount # Calculate cost with discount
cost_with_discount = completion_cost( cost_with_discount = completion_cost(
completion_response=response, completion_response=response,
model="vertex_ai/gemini-pro", model="vertex_ai/gemini-3-pro-preview",
custom_llm_provider="vertex_ai", custom_llm_provider="vertex_ai",
) )

View file

@ -90,7 +90,7 @@ def test_supports_function_calling_github_openai_alias():
def test_supports_function_calling_github_anthropic_alias(): def test_supports_function_calling_github_anthropic_alias():
assert ( assert (
litellm.utils.supports_function_calling( litellm.utils.supports_function_calling(
model="github/claude-3-5-sonnet-latest" model="github/claude-3-7-sonnet-20250219"
) )
is True is True
) )
@ -384,14 +384,14 @@ def test_all_model_configs():
assert ( assert (
"max_completion_tokens" "max_completion_tokens"
in VertexAIAnthropicConfig().get_supported_openai_params( in VertexAIAnthropicConfig().get_supported_openai_params(
model="claude-3-5-sonnet-20240620" model="claude-sonnet-4-6"
) )
) )
assert VertexAIAnthropicConfig().map_openai_params( assert VertexAIAnthropicConfig().map_openai_params(
non_default_params={"max_completion_tokens": 10}, non_default_params={"max_completion_tokens": 10},
optional_params={}, optional_params={},
model="claude-3-5-sonnet-20240620", model="claude-sonnet-4-6",
drop_params=False, drop_params=False,
) == {"max_tokens": 10} ) == {"max_tokens": 10}
@ -619,6 +619,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"input_cost_per_image_above_128k_tokens": {"type": "number"}, "input_cost_per_image_above_128k_tokens": {"type": "number"},
"input_cost_per_image_token": {"type": "number"}, "input_cost_per_image_token": {"type": "number"},
"input_cost_per_token_above_200k_tokens": {"type": "number"}, "input_cost_per_token_above_200k_tokens": {"type": "number"},
"input_cost_per_token_above_256k_tokens": {"type": "number"},
"input_cost_per_token_above_272k_tokens": {"type": "number"}, "input_cost_per_token_above_272k_tokens": {"type": "number"},
"cache_read_input_token_cost_flex": {"type": "number"}, "cache_read_input_token_cost_flex": {"type": "number"},
"cache_read_input_token_cost_priority": {"type": "number"}, "cache_read_input_token_cost_priority": {"type": "number"},
@ -700,6 +701,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"output_cost_per_token": {"type": "number"}, "output_cost_per_token": {"type": "number"},
"output_cost_per_token_above_128k_tokens": {"type": "number"}, "output_cost_per_token_above_128k_tokens": {"type": "number"},
"output_cost_per_token_above_200k_tokens": {"type": "number"}, "output_cost_per_token_above_200k_tokens": {"type": "number"},
"output_cost_per_token_above_256k_tokens": {"type": "number"},
"output_cost_per_token_above_272k_tokens": {"type": "number"}, "output_cost_per_token_above_272k_tokens": {"type": "number"},
"output_cost_per_image_above_1024_and_1024_pixels": {"type": "number"}, "output_cost_per_image_above_1024_and_1024_pixels": {"type": "number"},
"output_cost_per_image_above_1024_and_1024_pixels_and_premium_image": { "output_cost_per_image_above_1024_and_1024_pixels_and_premium_image": {
@ -738,6 +740,8 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"supports_vision": {"type": "boolean"}, "supports_vision": {"type": "boolean"},
"supports_web_search": {"type": "boolean"}, "supports_web_search": {"type": "boolean"},
"supports_url_context": {"type": "boolean"}, "supports_url_context": {"type": "boolean"},
"supports_multimodal": {"type": "boolean"},
"uses_embed_content": {"type": "boolean"},
"supports_reasoning": {"type": "boolean"}, "supports_reasoning": {"type": "boolean"},
"supports_none_reasoning_effort": {"type": "boolean"}, "supports_none_reasoning_effort": {"type": "boolean"},
"supports_xhigh_reasoning_effort": {"type": "boolean"}, "supports_xhigh_reasoning_effort": {"type": "boolean"},
@ -1136,7 +1140,7 @@ def test_get_model_info_shows_supports_computer_use():
[ [
("gpt-3.5-turbo", "openai"), ("gpt-3.5-turbo", "openai"),
("anthropic.claude-3-7-sonnet-20250219-v1:0", "bedrock"), ("anthropic.claude-3-7-sonnet-20250219-v1:0", "bedrock"),
("gemini-1.5-pro", "vertex_ai"), ("gemini-2.5-pro", "vertex_ai"),
], ],
) )
def test_pre_process_non_default_params(model, custom_llm_provider): def test_pre_process_non_default_params(model, custom_llm_provider):
@ -1225,14 +1229,14 @@ class TestProxyFunctionCalling:
), ),
# Anthropic models (Claude supports function calling) # Anthropic models (Claude supports function calling)
( (
"claude-3-5-sonnet-20240620", "claude-sonnet-4-6",
"litellm_proxy/claude-3-5-sonnet-20240620", "litellm_proxy/claude-sonnet-4-6",
True, True,
), ),
# Google models # Google models
("gemini-pro", "litellm_proxy/gemini-pro", True), ("gemini-2.5-pro", "litellm_proxy/gemini-2.5-pro", True),
("gemini/gemini-1.5-pro", "litellm_proxy/gemini/gemini-1.5-pro", True), ("gemini/gemini-2.5-pro", "litellm_proxy/gemini/gemini-2.5-pro", True),
("gemini/gemini-1.5-flash", "litellm_proxy/gemini/gemini-1.5-flash", True), ("gemini/gemini-2.5-flash", "litellm_proxy/gemini/gemini-2.5-flash", True),
# Groq models (mixed support) # Groq models (mixed support)
("groq/gemma-7b-it", "litellm_proxy/groq/gemma-7b-it", True), ("groq/gemma-7b-it", "litellm_proxy/groq/gemma-7b-it", True),
( (
@ -1437,8 +1441,8 @@ class TestProxyFunctionCalling:
("litellm_proxy/gpt-3.5-turbo", True), ("litellm_proxy/gpt-3.5-turbo", True),
("litellm_proxy/gpt-4", True), ("litellm_proxy/gpt-4", True),
("litellm_proxy/gpt-4o", True), ("litellm_proxy/gpt-4o", True),
("litellm_proxy/claude-3-5-sonnet-20240620", True), ("litellm_proxy/claude-sonnet-4-6", True),
("litellm_proxy/gemini/gemini-1.5-pro", True), ("litellm_proxy/gemini/gemini-2.5-pro", True),
# Test proxy models that should not support function calling # Test proxy models that should not support function calling
("litellm_proxy/command-nightly", False), ("litellm_proxy/command-nightly", False),
("litellm_proxy/anthropic.claude-instant-v1", False), ("litellm_proxy/anthropic.claude-instant-v1", False),
@ -1483,8 +1487,8 @@ class TestProxyFunctionCalling:
[ [
"litellm_proxy/gpt-3.5-turbo", "litellm_proxy/gpt-3.5-turbo",
"litellm_proxy/gpt-4", "litellm_proxy/gpt-4",
"litellm_proxy/claude-3-5-sonnet-20240620", "litellm_proxy/claude-sonnet-4-6",
"litellm_proxy/gemini/gemini-1.5-pro", "litellm_proxy/gemini/gemini-2.5-pro",
], ],
) )
def test_proxy_model_with_custom_llm_provider_none(self, model_name): def test_proxy_model_with_custom_llm_provider_none(self, model_name):

View file

@ -0,0 +1,364 @@
"""
Comprehensive test for new vector store endpoints: retrieve, list, update, delete
Tests both basic functionality and complex scenarios including target_model_names
"""
import asyncio
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm.proxy._types import UserAPIKeyAuth
@pytest.mark.asyncio
async def test_vector_store_retrieve_basic():
"""Test basic vector store retrieve functionality."""
router = litellm.Router(model_list=[])
mock_response = {
"id": "vs_test123",
"object": "vector_store",
"created_at": 1699061776,
"name": "Test Vector Store",
"file_counts": {
"in_progress": 0,
"completed": 5,
"failed": 0,
"cancelled": 0,
"total": 5,
},
"status": "completed",
"usage_bytes": 12345,
}
with patch(
"litellm.vector_stores.main.aretrieve",
new=AsyncMock(return_value=mock_response),
) as mock_retrieve:
result = await router.avector_store_retrieve(
vector_store_id="vs_test123",
custom_llm_provider="openai",
)
assert result["id"] == "vs_test123"
assert result["object"] == "vector_store"
assert result["status"] == "completed"
mock_retrieve.assert_called_once()
@pytest.mark.asyncio
async def test_vector_store_list_basic():
"""Test basic vector store list functionality."""
router = litellm.Router(model_list=[])
mock_response = {
"object": "list",
"data": [
{
"id": "vs_test1",
"object": "vector_store",
"created_at": 1699061776,
"name": "Store 1",
},
{
"id": "vs_test2",
"object": "vector_store",
"created_at": 1699061777,
"name": "Store 2",
},
],
"first_id": "vs_test1",
"last_id": "vs_test2",
"has_more": False,
}
with patch(
"litellm.vector_stores.main.alist",
new=AsyncMock(return_value=mock_response),
) as mock_list:
result = await router.avector_store_list(
limit=20,
order="desc",
custom_llm_provider="openai",
)
assert result["object"] == "list"
assert len(result["data"]) == 2
assert result["data"][0]["id"] == "vs_test1"
mock_list.assert_called_once()
@pytest.mark.asyncio
async def test_vector_store_update_basic():
"""Test basic vector store update functionality."""
router = litellm.Router(model_list=[])
mock_response = {
"id": "vs_test123",
"object": "vector_store",
"created_at": 1699061776,
"name": "Updated Name",
"metadata": {"key": "value"},
"status": "completed",
}
with patch(
"litellm.vector_stores.main.aupdate",
new=AsyncMock(return_value=mock_response),
) as mock_update:
result = await router.avector_store_update(
vector_store_id="vs_test123",
name="Updated Name",
metadata={"key": "value"},
custom_llm_provider="openai",
)
assert result["id"] == "vs_test123"
assert result["name"] == "Updated Name"
assert result["metadata"]["key"] == "value"
mock_update.assert_called_once()
@pytest.mark.asyncio
async def test_vector_store_delete_basic():
"""Test basic vector store delete functionality."""
router = litellm.Router(model_list=[])
mock_response = {
"id": "vs_test123",
"object": "vector_store.deleted",
"deleted": True,
}
with patch(
"litellm.vector_stores.main.adelete",
new=AsyncMock(return_value=mock_response),
) as mock_delete:
result = await router.avector_store_delete(
vector_store_id="vs_test123",
custom_llm_provider="openai",
)
assert result["id"] == "vs_test123"
assert result["deleted"] is True
assert result["object"] == "vector_store.deleted"
mock_delete.assert_called_once()
@pytest.mark.asyncio
async def test_async_vector_store_retrieve():
"""Test async vector store retrieve."""
router = litellm.Router(model_list=[])
mock_response = {
"id": "vs_async123",
"object": "vector_store",
"name": "Async Test Store",
}
with patch(
"litellm.vector_stores.main.aretrieve",
new=AsyncMock(return_value=mock_response),
) as mock_aretrieve:
result = await router.avector_store_retrieve(
vector_store_id="vs_async123",
custom_llm_provider="openai",
)
assert result["id"] == "vs_async123"
mock_aretrieve.assert_called_once()
@pytest.mark.asyncio
async def test_async_vector_store_list():
"""Test async vector store list."""
router = litellm.Router(model_list=[])
mock_response = {
"object": "list",
"data": [{"id": "vs_1"}, {"id": "vs_2"}],
}
with patch(
"litellm.vector_stores.main.alist",
new=AsyncMock(return_value=mock_response),
) as mock_alist:
result = await router.avector_store_list(
limit=10,
custom_llm_provider="openai",
)
assert len(result["data"]) == 2
mock_alist.assert_called_once()
@pytest.mark.asyncio
async def test_async_vector_store_update():
"""Test async vector store update."""
router = litellm.Router(model_list=[])
mock_response = {
"id": "vs_async123",
"name": "Updated Async Name",
}
with patch(
"litellm.vector_stores.main.aupdate",
new=AsyncMock(return_value=mock_response),
) as mock_aupdate:
result = await router.avector_store_update(
vector_store_id="vs_async123",
name="Updated Async Name",
custom_llm_provider="openai",
)
assert result["name"] == "Updated Async Name"
mock_aupdate.assert_called_once()
@pytest.mark.asyncio
async def test_async_vector_store_delete():
"""Test async vector store delete."""
router = litellm.Router(model_list=[])
mock_response = {
"id": "vs_async123",
"deleted": True,
}
with patch(
"litellm.vector_stores.main.adelete",
new=AsyncMock(return_value=mock_response),
) as mock_adelete:
result = await router.avector_store_delete(
vector_store_id="vs_async123",
custom_llm_provider="openai",
)
assert result["deleted"] is True
mock_adelete.assert_called_once()
@pytest.mark.asyncio
async def test_vector_store_list_with_pagination():
"""Test vector store list with pagination parameters."""
router = litellm.Router(model_list=[])
mock_response = {
"object": "list",
"data": [{"id": f"vs_{i}"} for i in range(5)],
"has_more": True,
"first_id": "vs_0",
"last_id": "vs_4",
}
with patch(
"litellm.vector_stores.main.list",
return_value=mock_response,
) as mock_list:
result = router.vector_store_list(
limit=5,
after="vs_previous",
order="asc",
custom_llm_provider="openai",
)
assert result["has_more"] is True
assert len(result["data"]) == 5
# Verify pagination params were passed
call_kwargs = mock_list.call_args.kwargs
assert call_kwargs["limit"] == 5
assert call_kwargs["after"] == "vs_previous"
assert call_kwargs["order"] == "asc"
@pytest.mark.asyncio
async def test_vector_store_update_with_expires_after():
"""Test vector store update with expiration policy."""
router = litellm.Router(model_list=[])
expires_after = {
"anchor": "last_active_at",
"days": 7,
}
mock_response = {
"id": "vs_test123",
"expires_after": expires_after,
"expires_at": 1699668576,
}
with patch(
"litellm.vector_stores.main.update",
return_value=mock_response,
) as mock_update:
result = router.vector_store_update(
vector_store_id="vs_test123",
expires_after=expires_after,
custom_llm_provider="openai",
)
assert result["expires_after"]["days"] == 7
assert result["expires_at"] is not None
call_kwargs = mock_update.call_args.kwargs
assert call_kwargs["expires_after"] == expires_after
def test_router_initializes_new_endpoints():
"""Test that router properly initializes the new vector store endpoints."""
router = litellm.Router(model_list=[])
# Verify all new endpoints are initialized
assert hasattr(router, "vector_store_retrieve")
assert hasattr(router, "avector_store_retrieve")
assert hasattr(router, "vector_store_list")
assert hasattr(router, "avector_store_list")
assert hasattr(router, "vector_store_update")
assert hasattr(router, "avector_store_update")
assert hasattr(router, "vector_store_delete")
assert hasattr(router, "avector_store_delete")
# Verify they are callable
assert callable(router.vector_store_retrieve)
assert callable(router.avector_store_retrieve)
assert callable(router.vector_store_list)
assert callable(router.avector_store_list)
assert callable(router.vector_store_update)
assert callable(router.avector_store_update)
assert callable(router.vector_store_delete)
assert callable(router.avector_store_delete)
if __name__ == "__main__":
# Run basic smoke tests
print("Running smoke tests for new vector store endpoints...")
# Test router initialization
print("✓ Testing router initialization...")
test_router_initializes_new_endpoints()
print("✓ Router initialization successful")
# Test basic sync operations
print("✓ Testing basic sync operations...")
asyncio.run(test_vector_store_retrieve_basic())
asyncio.run(test_vector_store_list_basic())
asyncio.run(test_vector_store_update_basic())
asyncio.run(test_vector_store_delete_basic())
print("✓ Basic sync operations successful")
# Test async operations
print("✓ Testing async operations...")
asyncio.run(test_async_vector_store_retrieve())
asyncio.run(test_async_vector_store_list())
asyncio.run(test_async_vector_store_update())
asyncio.run(test_async_vector_store_delete())
print("✓ Async operations successful")
print("\n✅ All smoke tests passed!")

View file

@ -46,7 +46,7 @@
"@testing-library/user-event": "^14.6.1", "@testing-library/user-event": "^14.6.1",
"@types/babel__traverse": "^7.28.0", "@types/babel__traverse": "^7.28.0",
"@types/lodash": "^4.17.15", "@types/lodash": "^4.17.15",
"@types/node": "^20", "@types/node": "20.19.37",
"@types/react": "18.2.48", "@types/react": "18.2.48",
"@types/react-copy-to-clipboard": "^5.0.7", "@types/react-copy-to-clipboard": "^5.0.7",
"@types/react-dom": "^18", "@types/react-dom": "^18",
@ -65,7 +65,7 @@
"postcss": "^8.4.33", "postcss": "^8.4.33",
"prettier": "3.2.5", "prettier": "3.2.5",
"tailwindcss": "^3.4.1", "tailwindcss": "^3.4.1",
"typescript": "^5.3.3", "typescript": "5.9.3",
"vite": "^7.1.11", "vite": "^7.1.11",
"vitest": "^3.2.4" "vitest": "^3.2.4"
}, },
@ -3394,9 +3394,9 @@
"license": "MIT" "license": "MIT"
}, },
"node_modules/@types/node": { "node_modules/@types/node": {
"version": "20.19.30", "version": "20.19.37",
"resolved": "https://registry.npmjs.org/@types/node/-/node-20.19.30.tgz", "resolved": "https://registry.npmjs.org/@types/node/-/node-20.19.37.tgz",
"integrity": "sha512-WJtwWJu7UdlvzEAUm484QNg5eAoq5QR08KDNx7g45Usrs2NtOPiX8ugDqmKdXkyL03rBqU5dYNYVQetEpBHq2g==", "integrity": "sha512-8kzdPJ3FsNsVIurqBs7oodNnCEVbni9yUEkaHbgptDACOPW04jimGagZ51E6+lXUwJjgnBw+hyko/lkFWCldqw==",
"license": "MIT", "license": "MIT",
"dependencies": { "dependencies": {
"undici-types": "~6.21.0" "undici-types": "~6.21.0"
@ -12555,9 +12555,9 @@
} }
}, },
"node_modules/typescript": { "node_modules/typescript": {
"version": "5.3.3", "version": "5.9.3",
"resolved": "https://registry.npmjs.org/typescript/-/typescript-5.3.3.tgz", "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.9.3.tgz",
"integrity": "sha512-pXWcraxM0uxAS+tN0AG/BF2TyqmHO014Z070UsJ+pFvYuRSq8KH8DmWpnbXe0pEPDHXZV3FcAbJkijJ5oNEnWw==", "integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==",
"devOptional": true, "devOptional": true,
"license": "Apache-2.0", "license": "Apache-2.0",
"bin": { "bin": {
@ -13279,21 +13279,6 @@
"type": "github", "type": "github",
"url": "https://github.com/sponsors/wooorm" "url": "https://github.com/sponsors/wooorm"
} }
},
"node_modules/@next/swc-win32-ia32-msvc": {
"version": "14.2.33",
"resolved": "https://registry.npmjs.org/@next/swc-win32-ia32-msvc/-/swc-win32-ia32-msvc-14.2.33.tgz",
"integrity": "sha512-pc9LpGNKhJ0dXQhZ5QMmYxtARwwmWLpeocFmVG5Z0DzWq5Uf0izcI8tLc+qOpqxO1PWqZ5A7J1blrUIKrIFc7Q==",
"cpu": [
"ia32"
],
"optional": true,
"os": [
"win32"
],
"engines": {
"node": ">= 10"
}
} }
} }
} }

View file

@ -58,7 +58,7 @@
"@testing-library/user-event": "^14.6.1", "@testing-library/user-event": "^14.6.1",
"@types/babel__traverse": "^7.28.0", "@types/babel__traverse": "^7.28.0",
"@types/lodash": "^4.17.15", "@types/lodash": "^4.17.15",
"@types/node": "^20", "@types/node": "20.19.37",
"@types/react": "18.2.48", "@types/react": "18.2.48",
"@types/react-copy-to-clipboard": "^5.0.7", "@types/react-copy-to-clipboard": "^5.0.7",
"@types/react-dom": "^18", "@types/react-dom": "^18",
@ -77,7 +77,7 @@
"postcss": "^8.4.33", "postcss": "^8.4.33",
"prettier": "3.2.5", "prettier": "3.2.5",
"tailwindcss": "^3.4.1", "tailwindcss": "^3.4.1",
"typescript": "^5.3.3", "typescript": "5.9.3",
"vite": "^7.1.11", "vite": "^7.1.11",
"vitest": "^3.2.4" "vitest": "^3.2.4"
}, },

View file

@ -74,7 +74,7 @@ describe("useMCPServers", () => {
expect(result.current.isSuccess).toBe(true); expect(result.current.isSuccess).toBe(true);
}); });
expect(networking.fetchMCPServers).toHaveBeenCalledWith(mockAccessToken); expect(networking.fetchMCPServers).toHaveBeenCalledWith(mockAccessToken, undefined);
expect(result.current.data).toEqual(mockServers); expect(result.current.data).toEqual(mockServers);
}); });

View file

@ -20,11 +20,15 @@ vi.mock("@/components/molecules/notifications_manager", () => ({
// Mock react-query // Mock react-query
const mockInvalidateQueries = vi.fn(); const mockInvalidateQueries = vi.fn();
vi.mock("@tanstack/react-query", () => ({ vi.mock("@tanstack/react-query", async (importOriginal) => {
useQueryClient: () => ({ const actual = await importOriginal() as any;
invalidateQueries: mockInvalidateQueries, return {
}), ...actual,
})); useQueryClient: () => ({
invalidateQueries: mockInvalidateQueries,
}),
};
});
// Mock the useModelsInfo hook // Mock the useModelsInfo hook
const mockUseModelsInfo = vi.fn(() => ({ const mockUseModelsInfo = vi.fn(() => ({
@ -553,7 +557,7 @@ describe("AllModelsTab", () => {
mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null, refetch: vi.fn() }); mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null, refetch: vi.fn() });
render(<AllModelsTab {...defaultProps} />); renderWithProviders(<AllModelsTab {...defaultProps} />);
await waitFor(() => { await waitFor(() => {
expect(screen.getByText("gpt-4-delete-test")).toBeInTheDocument(); expect(screen.getByText("gpt-4-delete-test")).toBeInTheDocument();
@ -597,7 +601,7 @@ describe("AllModelsTab", () => {
mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null, refetch: vi.fn() }); mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null, refetch: vi.fn() });
render(<AllModelsTab {...defaultProps} />); renderWithProviders(<AllModelsTab {...defaultProps} />);
await waitFor(() => { await waitFor(() => {
expect(screen.getByText("gpt-4-clickable")).toBeInTheDocument(); expect(screen.getByText("gpt-4-clickable")).toBeInTheDocument();

View file

@ -3,7 +3,11 @@
import { Suspense, useEffect, useMemo } from "react"; import { Suspense, useEffect, useMemo } from "react";
import { useSearchParams } from "next/navigation"; import { useSearchParams } from "next/navigation";
const RESULT_STORAGE_KEY = "litellm-mcp-oauth-result"; // Written to sessionStorage so both the admin hook (useMcpOAuthFlow) and the
// user hook (useUserMcpOAuthFlow) can pick up the result. Each hook reads
// its own namespace to avoid cross-flow collisions.
const ADMIN_RESULT_KEY = "litellm-mcp-oauth-result";
const USER_RESULT_KEY = "litellm-user-mcp-oauth-result";
const RETURN_URL_STORAGE_KEY = "litellm-mcp-oauth-return-url"; const RETURN_URL_STORAGE_KEY = "litellm-mcp-oauth-return-url";
const resolveDefaultRedirect = () => { const resolveDefaultRedirect = () => {
@ -32,6 +36,10 @@ const McpOAuthCallbackContent = () => {
type: "litellm-mcp-oauth", type: "litellm-mcp-oauth",
code: searchParams.get("code"), code: searchParams.get("code"),
state: searchParams.get("state"), state: searchParams.get("state"),
// Forward OAuth provider error params so the hook can surface the real
// reason (e.g. "access_denied") instead of a generic "code missing" error.
error: searchParams.get("error"),
error_description: searchParams.get("error_description"),
}; };
}, [searchParams]); }, [searchParams]);
@ -41,16 +49,16 @@ const McpOAuthCallbackContent = () => {
} }
try { try {
// Store in both sessionStorage and localStorage for redundancy // Write to both namespace keys (admin and user) so whichever hook is
window.sessionStorage.setItem(RESULT_STORAGE_KEY, JSON.stringify(payload)); // active can consume the result. sessionStorage only — no localStorage.
window.localStorage.setItem(RESULT_STORAGE_KEY, JSON.stringify(payload)); const serialized = JSON.stringify(payload);
window.sessionStorage.setItem(ADMIN_RESULT_KEY, serialized);
window.sessionStorage.setItem(USER_RESULT_KEY, serialized);
} catch (err) { } catch (err) {
// Silently ignore storage errors // Silently ignore storage errors
} }
// Check both sessionStorage and localStorage for return URL const returnUrl = window.sessionStorage.getItem(RETURN_URL_STORAGE_KEY);
const returnUrl = window.sessionStorage.getItem(RETURN_URL_STORAGE_KEY) ||
window.localStorage.getItem(RETURN_URL_STORAGE_KEY);
const destination = returnUrl || resolveDefaultRedirect(); const destination = returnUrl || resolveDefaultRedirect();
window.location.replace(destination); window.location.replace(destination);
}, [payload]); }, [payload]);

View file

@ -547,7 +547,7 @@ function CreateKeyPageContent() {
) : page == "policies" ? ( ) : page == "policies" ? (
<PoliciesPanel accessToken={accessToken} userRole={userRole} /> <PoliciesPanel accessToken={accessToken} userRole={userRole} />
) : page == "agents" ? ( ) : page == "agents" ? (
<AgentsPanel accessToken={accessToken} userRole={userRole} /> <AgentsPanel accessToken={accessToken} userRole={userRole} teams={teams} />
) : page == "prompts" ? ( ) : page == "prompts" ? (
<PromptsPanel accessToken={accessToken} userRole={userRole} /> <PromptsPanel accessToken={accessToken} userRole={userRole} />
) : page == "transform-request" ? ( ) : page == "transform-request" ? (

View file

@ -125,6 +125,208 @@ describe("EntityUsage", () => {
}, },
}; };
const mockAgentSpendData = {
results: [
{
date: "2025-01-01",
metrics: {
spend: 245.8,
api_requests: 3200,
successful_requests: 3100,
failed_requests: 100,
total_tokens: 1250000,
prompt_tokens: 850000,
completion_tokens: 400000,
cache_read_input_tokens: 50000,
cache_creation_input_tokens: 10000,
},
breakdown: {
entities: {
"agent-code-review": {
metrics: {
spend: 120.4,
api_requests: 1500,
successful_requests: 1450,
failed_requests: 50,
total_tokens: 620000,
prompt_tokens: 420000,
completion_tokens: 200000,
cache_read_input_tokens: 30000,
cache_creation_input_tokens: 5000,
},
metadata: { agent_name: "Code Review Agent" },
api_key_breakdown: {},
},
"agent-customer-support": {
metrics: {
spend: 85.2,
api_requests: 1200,
successful_requests: 1170,
failed_requests: 30,
total_tokens: 430000,
prompt_tokens: 290000,
completion_tokens: 140000,
cache_read_input_tokens: 15000,
cache_creation_input_tokens: 3000,
},
metadata: { agent_name: "Customer Support Agent" },
api_key_breakdown: {},
},
"agent-data-analyst": {
metrics: {
spend: 40.2,
api_requests: 500,
successful_requests: 480,
failed_requests: 20,
total_tokens: 200000,
prompt_tokens: 140000,
completion_tokens: 60000,
cache_read_input_tokens: 5000,
cache_creation_input_tokens: 2000,
},
metadata: { agent_name: "Data Analyst Agent" },
api_key_breakdown: {},
},
},
models: {
"gpt-4o": {
metrics: {
spend: 180.0,
api_requests: 2000,
successful_requests: 1950,
failed_requests: 50,
total_tokens: 900000,
prompt_tokens: 600000,
completion_tokens: 300000,
cache_read_input_tokens: 40000,
cache_creation_input_tokens: 8000,
},
metadata: {},
api_key_breakdown: {},
},
"claude-sonnet-4-20250514": {
metrics: {
spend: 65.8,
api_requests: 1200,
successful_requests: 1150,
failed_requests: 50,
total_tokens: 350000,
prompt_tokens: 250000,
completion_tokens: 100000,
cache_read_input_tokens: 10000,
cache_creation_input_tokens: 2000,
},
metadata: {},
api_key_breakdown: {},
},
},
api_keys: {},
providers: {
openai: {
metrics: {
spend: 180.0,
api_requests: 2000,
successful_requests: 1950,
failed_requests: 50,
total_tokens: 900000,
prompt_tokens: 600000,
completion_tokens: 300000,
cache_read_input_tokens: 40000,
cache_creation_input_tokens: 8000,
},
},
anthropic: {
metrics: {
spend: 65.8,
api_requests: 1200,
successful_requests: 1150,
failed_requests: 50,
total_tokens: 350000,
prompt_tokens: 250000,
completion_tokens: 100000,
cache_read_input_tokens: 10000,
cache_creation_input_tokens: 2000,
},
},
},
},
},
{
date: "2025-01-02",
metrics: {
spend: 198.5,
api_requests: 2800,
successful_requests: 2720,
failed_requests: 80,
total_tokens: 980000,
prompt_tokens: 670000,
completion_tokens: 310000,
cache_read_input_tokens: 42000,
cache_creation_input_tokens: 9000,
},
breakdown: {
entities: {
"agent-code-review": {
metrics: {
spend: 95.3,
api_requests: 1300,
successful_requests: 1270,
failed_requests: 30,
total_tokens: 510000,
prompt_tokens: 350000,
completion_tokens: 160000,
cache_read_input_tokens: 25000,
cache_creation_input_tokens: 4000,
},
metadata: { agent_name: "Code Review Agent" },
api_key_breakdown: {},
},
"agent-customer-support": {
metrics: {
spend: 68.7,
api_requests: 1000,
successful_requests: 970,
failed_requests: 30,
total_tokens: 320000,
prompt_tokens: 220000,
completion_tokens: 100000,
cache_read_input_tokens: 12000,
cache_creation_input_tokens: 3000,
},
metadata: { agent_name: "Customer Support Agent" },
api_key_breakdown: {},
},
"agent-data-analyst": {
metrics: {
spend: 34.5,
api_requests: 500,
successful_requests: 480,
failed_requests: 20,
total_tokens: 150000,
prompt_tokens: 100000,
completion_tokens: 50000,
cache_read_input_tokens: 5000,
cache_creation_input_tokens: 2000,
},
metadata: { agent_name: "Data Analyst Agent" },
api_key_breakdown: {},
},
},
models: {},
api_keys: {},
providers: {},
},
},
],
metadata: {
total_spend: 444.3,
total_api_requests: 6000,
total_successful_requests: 5820,
total_failed_requests: 180,
total_tokens: 2230000,
},
};
const defaultProps = { const defaultProps = {
accessToken: "test-token", accessToken: "test-token",
entityType: "tag" as const, entityType: "tag" as const,
@ -153,7 +355,7 @@ describe("EntityUsage", () => {
mockTeamDailyActivityCall.mockResolvedValue(mockSpendData); mockTeamDailyActivityCall.mockResolvedValue(mockSpendData);
mockOrganizationDailyActivityCall.mockResolvedValue(mockSpendData); mockOrganizationDailyActivityCall.mockResolvedValue(mockSpendData);
mockCustomerDailyActivityCall.mockResolvedValue(mockSpendData); mockCustomerDailyActivityCall.mockResolvedValue(mockSpendData);
mockAgentDailyActivityCall.mockResolvedValue(mockSpendData); mockAgentDailyActivityCall.mockResolvedValue(mockAgentSpendData);
mockUserDailyActivityCall.mockResolvedValue(mockSpendData); mockUserDailyActivityCall.mockResolvedValue(mockSpendData);
}); });
@ -231,7 +433,7 @@ describe("EntityUsage", () => {
expect(screen.getByText("Agent Spend Overview")).toBeInTheDocument(); expect(screen.getByText("Agent Spend Overview")).toBeInTheDocument();
await waitFor(() => { await waitFor(() => {
const spendElements = screen.getAllByText("$100.50"); const spendElements = screen.getAllByText("$444.30");
expect(spendElements.length).toBeGreaterThan(0); expect(spendElements.length).toBeGreaterThan(0);
}); });
}); });
@ -385,6 +587,87 @@ describe("EntityUsage", () => {
}); });
}); });
it("should display Agent Activity tab for team entity type", async () => {
render(<EntityUsage {...defaultProps} entityType="team" />);
await waitFor(() => {
expect(mockTeamDailyActivityCall).toHaveBeenCalled();
});
expect(screen.getByText("Agent Activity")).toBeInTheDocument();
});
it("should not display Agent Activity tab for non-team entity types", async () => {
render(<EntityUsage {...defaultProps} entityType="tag" />);
await waitFor(() => {
expect(mockTagDailyActivityCall).toHaveBeenCalled();
});
expect(screen.queryByText("Agent Activity")).not.toBeInTheDocument();
});
it("should display Top Agents Driving Spend card for team entity type", async () => {
render(<EntityUsage {...defaultProps} entityType="team" />);
await waitFor(() => {
expect(mockTeamDailyActivityCall).toHaveBeenCalled();
});
expect(screen.getByText("Top Agents Driving Spend")).toBeInTheDocument();
});
it("should not display Top Agents Driving Spend card for non-team entity types", async () => {
render(<EntityUsage {...defaultProps} entityType="tag" />);
await waitFor(() => {
expect(mockTagDailyActivityCall).toHaveBeenCalled();
});
expect(screen.queryByText("Top Agents Driving Spend")).not.toBeInTheDocument();
});
it("should fetch agent activity data when entity type is team", async () => {
render(<EntityUsage {...defaultProps} entityType="team" />);
await waitFor(() => {
expect(mockAgentDailyActivityCall).toHaveBeenCalledWith(
"test-token",
expect.any(Date),
expect.any(Date),
1,
null,
);
});
});
it("should not fetch agent activity data for non-team entity types", async () => {
render(<EntityUsage {...defaultProps} entityType="tag" />);
await waitFor(() => {
expect(mockTagDailyActivityCall).toHaveBeenCalled();
});
expect(mockAgentDailyActivityCall).not.toHaveBeenCalled();
});
it("should switch to Agent Activity tab for team entity type", async () => {
render(<EntityUsage {...defaultProps} entityType="team" />);
await waitFor(() => {
expect(mockTeamDailyActivityCall).toHaveBeenCalled();
});
const agentActivityTab = screen.getByText("Agent Activity");
act(() => {
fireEvent.click(agentActivityTab);
});
await waitFor(() => {
expect(screen.getAllByText("Activity Metrics").length).toBeGreaterThan(0);
});
});
it("should fallback to entity value when no entityList and no team_alias", async () => { it("should fallback to entity value when no entityList and no team_alias", async () => {
const spendDataWithoutAlias = { const spendDataWithoutAlias = {
...mockSpendData, ...mockSpendData,

View file

@ -100,11 +100,24 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
}); });
const { teams } = useTeams(); const { teams } = useTeams();
const [agentSpendData, setAgentSpendData] = useState<EntitySpendData>({
results: [],
metadata: {
total_spend: 0,
total_api_requests: 0,
total_successful_requests: 0,
total_failed_requests: 0,
total_tokens: 0,
},
});
const modelMetrics = processActivityData(spendData, "models", teams || []); const modelMetrics = processActivityData(spendData, "models", teams || []);
const keyMetrics = processActivityData(spendData, "api_keys", teams || []); const keyMetrics = processActivityData(spendData, "api_keys", teams || []);
const agentMetrics = entityType === "team" ? processActivityData(agentSpendData, "entities", teams || []) : {};
const [selectedTags, setSelectedTags] = useState<string[]>([]); const [selectedTags, setSelectedTags] = useState<string[]>([]);
const [topKeysLimit, setTopKeysLimit] = useState<number>(5); const [topKeysLimit, setTopKeysLimit] = useState<number>(5);
const [topModelsLimit, setTopModelsLimit] = useState<number>(5); const [topModelsLimit, setTopModelsLimit] = useState<number>(5);
const [topAgentsLimit, setTopAgentsLimit] = useState<number>(5);
const fetchSpendData = async () => { const fetchSpendData = async () => {
if (!accessToken || !dateValue.from || !dateValue.to) return; if (!accessToken || !dateValue.from || !dateValue.to) return;
@ -171,8 +184,21 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
} }
}; };
const fetchAgentSpendData = async () => {
if (!accessToken || !dateValue.from || !dateValue.to || entityType !== "team") return;
const startTime = new Date(dateValue.from);
const endTime = new Date(dateValue.to);
try {
const data = await agentDailyActivityCall(accessToken, startTime, endTime, 1, null);
setAgentSpendData(data);
} catch (e) {
console.error("Failed to fetch agent activity data:", e);
}
};
useEffect(() => { useEffect(() => {
fetchSpendData(); fetchSpendData();
fetchAgentSpendData();
}, [accessToken, dateValue, entityId, selectedTags]); }, [accessToken, dateValue, entityId, selectedTags]);
const getTopModels = () => { const getTopModels = () => {
@ -209,6 +235,37 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
.slice(0, topModelsLimit); .slice(0, topModelsLimit);
}; };
const getTopAgents = () => {
const agentSpend: { [key: string]: any } = {};
agentSpendData.results.forEach((day) => {
Object.entries(day.breakdown.entities || {}).forEach(([agentId, data]) => {
if (!agentSpend[agentId]) {
agentSpend[agentId] = {
spend: 0,
requests: 0,
successful_requests: 0,
failed_requests: 0,
tokens: 0,
agent_name: (data.metadata as any)?.agent_name || agentId,
};
}
agentSpend[agentId].spend += data.metrics.spend;
agentSpend[agentId].requests += data.metrics.api_requests;
agentSpend[agentId].successful_requests += data.metrics.successful_requests;
agentSpend[agentId].failed_requests += data.metrics.failed_requests;
agentSpend[agentId].tokens += data.metrics.total_tokens;
});
});
return Object.entries(agentSpend)
.map(([agentId, metrics]) => ({
key: metrics.agent_name,
...metrics,
}))
.sort((a, b) => b.spend - a.spend)
.slice(0, topAgentsLimit);
};
const getTopAPIKeys = () => { const getTopAPIKeys = () => {
console.log("debugTags", { spendData }); console.log("debugTags", { spendData });
const keySpend: { [key: string]: KeyMetricWithMetadata } = {}; const keySpend: { [key: string]: KeyMetricWithMetadata } = {};
@ -408,6 +465,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
<TabList variant="solid" className="mt-1"> <TabList variant="solid" className="mt-1">
<Tab>Cost</Tab> <Tab>Cost</Tab>
<Tab>{entityType === "agent" ? "Request / Token Consumption" : "Model Activity"}</Tab> <Tab>{entityType === "agent" ? "Request / Token Consumption" : "Model Activity"}</Tab>
{entityType === "team" && <Tab>Agent Activity</Tab>}
<Tab>Key Activity</Tab> <Tab>Key Activity</Tab>
<Tab>Endpoint Activity</Tab> <Tab>Endpoint Activity</Tab>
</TabList> </TabList>
@ -621,6 +679,20 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
</Card> </Card>
</Col> </Col>
{/* Top Agents - only for team entity type */}
{entityType === "team" && (
<Col numColSpan={2}>
<Card>
<Title>Top Agents Driving Spend</Title>
<TopModelView
topModels={getTopAgents()}
topModelsLimit={topAgentsLimit}
setTopModelsLimit={setTopAgentsLimit}
/>
</Card>
</Col>
)}
{/* Spend by Provider */} {/* Spend by Provider */}
<Col numColSpan={2}> <Col numColSpan={2}>
<Card> <Card>
@ -696,6 +768,11 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
<TabPanel> <TabPanel>
<ActivityMetrics modelMetrics={modelMetrics} hidePromptCachingMetrics={entityType === "agent"} /> <ActivityMetrics modelMetrics={modelMetrics} hidePromptCachingMetrics={entityType === "agent"} />
</TabPanel> </TabPanel>
{entityType === "team" && (
<TabPanel>
<ActivityMetrics modelMetrics={agentMetrics} />
</TabPanel>
)}
<TabPanel> <TabPanel>
<ActivityMetrics modelMetrics={keyMetrics} hidePromptCachingMetrics={entityType === "agent"} /> <ActivityMetrics modelMetrics={keyMetrics} hidePromptCachingMetrics={entityType === "agent"} />
</TabPanel> </TabPanel>

View file

@ -249,6 +249,8 @@ vi.mock("@ant-design/icons", async () => {
CalendarOutlined: Icon, CalendarOutlined: Icon,
InfoCircleOutlined: Icon, InfoCircleOutlined: Icon,
UserOutlined: Icon, UserOutlined: Icon,
DownOutlined: Icon,
RightOutlined: Icon,
LoadingOutlined, LoadingOutlined,
}; };
}); });

View file

@ -6,7 +6,7 @@
* Works at 1m+ spend logs, by querying an aggregate table instead. * Works at 1m+ spend logs, by querying an aggregate table instead.
*/ */
import { InfoCircleOutlined, LoadingOutlined } from "@ant-design/icons"; import { DownOutlined, InfoCircleOutlined, LoadingOutlined, RightOutlined } from "@ant-design/icons";
import { import {
BarChart, BarChart,
Card, Card,
@ -148,6 +148,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
const [showCredentialBanner, setShowCredentialBanner] = useState(true); const [showCredentialBanner, setShowCredentialBanner] = useState(true);
const [topKeysLimit, setTopKeysLimit] = useState<number>(5); const [topKeysLimit, setTopKeysLimit] = useState<number>(5);
const [topModelsLimit, setTopModelsLimit] = useState<number>(5); const [topModelsLimit, setTopModelsLimit] = useState<number>(5);
const [showTokenBreakdown, setShowTokenBreakdown] = useState(false);
const getAllTags = async () => { const getAllTags = async () => {
if (!accessToken) { if (!accessToken) {
return; return;
@ -176,7 +177,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
const totalSpend = userSpendData.metadata?.total_spend || 0; const totalSpend = userSpendData.metadata?.total_spend || 0;
// Calculate top models from the breakdown data // Calculate top models from the breakdown data
const getTopModels = (limit: number = 5) => { const topModels = useMemo(() => {
const modelSpend: { [key: string]: MetricWithMetadata } = {}; const modelSpend: { [key: string]: MetricWithMetadata } = {};
userSpendData.results.forEach((day) => { userSpendData.results.forEach((day) => {
Object.entries(day.breakdown.models || {}).forEach(([model, metrics]) => { Object.entries(day.breakdown.models || {}).forEach(([model, metrics]) => {
@ -219,10 +220,10 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
tokens: metrics.metrics.total_tokens, tokens: metrics.metrics.total_tokens,
})) }))
.sort((a, b) => b.spend - a.spend) .sort((a, b) => b.spend - a.spend)
.slice(0, limit); .slice(0, topModelsLimit);
}; }, [userSpendData.results, topModelsLimit]);
const getTopModelGroups = (limit: number = 5) => { const topModelGroups = useMemo(() => {
const modelGroupSpend: { [key: string]: MetricWithMetadata } = {}; const modelGroupSpend: { [key: string]: MetricWithMetadata } = {};
userSpendData.results.forEach((day) => { userSpendData.results.forEach((day) => {
Object.entries(day.breakdown.model_groups || {}).forEach(([modelGroup, metrics]) => { Object.entries(day.breakdown.model_groups || {}).forEach(([modelGroup, metrics]) => {
@ -266,16 +267,16 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
tokens: metrics.metrics.total_tokens, tokens: metrics.metrics.total_tokens,
})) }))
.sort((a, b) => b.spend - a.spend) .sort((a, b) => b.spend - a.spend)
.slice(0, limit); .slice(0, topModelsLimit);
}; }, [userSpendData.results, topModelsLimit]);
// Calculate provider spend from the breakdown data // Calculate provider spend from the breakdown data
const getProviderSpend = () => { const providerSpend = useMemo(() => {
const providerSpend: { [key: string]: MetricWithMetadata } = {}; const providerSpendMap: { [key: string]: MetricWithMetadata } = {};
userSpendData.results.forEach((day) => { userSpendData.results.forEach((day) => {
Object.entries(day.breakdown.providers || {}).forEach(([provider, metrics]) => { Object.entries(day.breakdown.providers || {}).forEach(([provider, metrics]) => {
if (!providerSpend[provider]) { if (!providerSpendMap[provider]) {
providerSpend[provider] = { providerSpendMap[provider] = {
metrics: { metrics: {
spend: 0, spend: 0,
prompt_tokens: 0, prompt_tokens: 0,
@ -291,19 +292,19 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
api_key_breakdown: {}, api_key_breakdown: {},
}; };
} }
providerSpend[provider].metrics.spend += metrics.metrics.spend; providerSpendMap[provider].metrics.spend += metrics.metrics.spend;
providerSpend[provider].metrics.prompt_tokens += metrics.metrics.prompt_tokens; providerSpendMap[provider].metrics.prompt_tokens += metrics.metrics.prompt_tokens;
providerSpend[provider].metrics.completion_tokens += metrics.metrics.completion_tokens; providerSpendMap[provider].metrics.completion_tokens += metrics.metrics.completion_tokens;
providerSpend[provider].metrics.total_tokens += metrics.metrics.total_tokens; providerSpendMap[provider].metrics.total_tokens += metrics.metrics.total_tokens;
providerSpend[provider].metrics.api_requests += metrics.metrics.api_requests; providerSpendMap[provider].metrics.api_requests += metrics.metrics.api_requests;
providerSpend[provider].metrics.successful_requests += metrics.metrics.successful_requests || 0; providerSpendMap[provider].metrics.successful_requests += metrics.metrics.successful_requests || 0;
providerSpend[provider].metrics.failed_requests += metrics.metrics.failed_requests || 0; providerSpendMap[provider].metrics.failed_requests += metrics.metrics.failed_requests || 0;
providerSpend[provider].metrics.cache_read_input_tokens += metrics.metrics.cache_read_input_tokens || 0; providerSpendMap[provider].metrics.cache_read_input_tokens += metrics.metrics.cache_read_input_tokens || 0;
providerSpend[provider].metrics.cache_creation_input_tokens += metrics.metrics.cache_creation_input_tokens || 0; providerSpendMap[provider].metrics.cache_creation_input_tokens += metrics.metrics.cache_creation_input_tokens || 0;
}); });
}); });
return Object.entries(providerSpend).map(([provider, metrics]) => ({ return Object.entries(providerSpendMap).map(([provider, metrics]) => ({
provider, provider,
spend: metrics.metrics.spend, spend: metrics.metrics.spend,
requests: metrics.metrics.api_requests, requests: metrics.metrics.api_requests,
@ -311,10 +312,10 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
failed_requests: metrics.metrics.failed_requests, failed_requests: metrics.metrics.failed_requests,
tokens: metrics.metrics.total_tokens, tokens: metrics.metrics.total_tokens,
})); }));
}; }, [userSpendData.results]);
// Calculate top API keys from the breakdown data // Calculate top API keys from the breakdown data
const getTopKeys = (limit: number = 5) => { const topKeys = useMemo(() => {
const keySpend: { [key: string]: KeyMetricWithMetadata } = {}; const keySpend: { [key: string]: KeyMetricWithMetadata } = {};
userSpendData.results.forEach((day) => { userSpendData.results.forEach((day) => {
Object.entries(day.breakdown.api_keys || {}).forEach(([key, metrics]) => { Object.entries(day.breakdown.api_keys || {}).forEach(([key, metrics]) => {
@ -334,7 +335,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
metadata: { metadata: {
key_alias: metrics.metadata.key_alias, key_alias: metrics.metadata.key_alias,
team_id: null, team_id: null,
tags: metrics.metadata.tags || [], // This gets key-level tags tags: metrics.metadata.tags || [],
}, },
}; };
} }
@ -350,18 +351,16 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
}); });
}); });
console.log("debugTags", { keySpend, userSpendData });
return Object.entries(keySpend) return Object.entries(keySpend)
.map(([api_key, metrics]) => ({ .map(([api_key, metrics]) => ({
api_key, api_key,
key_alias: metrics.metadata.key_alias || "-", // Using truncated key as alias key_alias: metrics.metadata.key_alias || "-",
tags: metrics.metadata.tags || [], // This will show key-level tags tags: metrics.metadata.tags || [],
spend: metrics.metrics.spend, spend: metrics.metrics.spend,
})) }))
.sort((a, b) => b.spend - a.spend) .sort((a, b) => b.spend - a.spend)
.slice(0, limit); .slice(0, topKeysLimit);
}; }, [userSpendData.results, topKeysLimit]);
const fetchUserSpendData = useCallback(async () => { const fetchUserSpendData = useCallback(async () => {
if (!accessToken || !dateValue.from || !dateValue.to) return; if (!accessToken || !dateValue.from || !dateValue.to) return;
@ -399,11 +398,15 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
const pageData = await userDailyActivityCall(accessToken, startTime, endTime, page, effectiveUserId); const pageData = await userDailyActivityCall(accessToken, startTime, endTime, page, effectiveUserId);
allResults.push(...pageData.results); allResults.push(...pageData.results);
if (pageData.metadata) { if (pageData.metadata) {
aggregatedMetadata.total_spend += pageData.metadata.total_spend || 0; aggregatedMetadata.total_spend = (aggregatedMetadata.total_spend || 0) + (pageData.metadata.total_spend || 0);
aggregatedMetadata.total_api_requests += pageData.metadata.total_api_requests || 0; aggregatedMetadata.total_api_requests = (aggregatedMetadata.total_api_requests || 0) + (pageData.metadata.total_api_requests || 0);
aggregatedMetadata.total_successful_requests += pageData.metadata.total_successful_requests || 0; aggregatedMetadata.total_successful_requests = (aggregatedMetadata.total_successful_requests || 0) + (pageData.metadata.total_successful_requests || 0);
aggregatedMetadata.total_failed_requests += pageData.metadata.total_failed_requests || 0; aggregatedMetadata.total_failed_requests = (aggregatedMetadata.total_failed_requests || 0) + (pageData.metadata.total_failed_requests || 0);
aggregatedMetadata.total_tokens += pageData.metadata.total_tokens || 0; aggregatedMetadata.total_tokens = (aggregatedMetadata.total_tokens || 0) + (pageData.metadata.total_tokens || 0);
aggregatedMetadata.total_prompt_tokens = (aggregatedMetadata.total_prompt_tokens || 0) + (pageData.metadata.total_prompt_tokens || 0);
aggregatedMetadata.total_completion_tokens = (aggregatedMetadata.total_completion_tokens || 0) + (pageData.metadata.total_completion_tokens || 0);
aggregatedMetadata.total_cache_read_input_tokens = (aggregatedMetadata.total_cache_read_input_tokens || 0) + (pageData.metadata.total_cache_read_input_tokens || 0);
aggregatedMetadata.total_cache_creation_input_tokens = (aggregatedMetadata.total_cache_creation_input_tokens || 0) + (pageData.metadata.total_cache_creation_input_tokens || 0);
} }
} }
@ -440,9 +443,13 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
return () => clearTimeout(timeoutId); return () => clearTimeout(timeoutId);
}, [fetchUserSpendData]); }, [fetchUserSpendData]);
const modelMetrics = processActivityData(userSpendData, "models", teams); const sortedDailyResults = useMemo(
const keyMetrics = processActivityData(userSpendData, "api_keys", teams); () => [...userSpendData.results].sort((a, b) => new Date(a.date).getTime() - new Date(b.date).getTime()),
const mcpServerMetrics = processActivityData(userSpendData, "mcp_servers", teams); [userSpendData.results],
);
const modelMetrics = useMemo(() => processActivityData(userSpendData, "models", teams), [userSpendData, teams]);
const keyMetrics = useMemo(() => processActivityData(userSpendData, "api_keys", teams), [userSpendData, teams]);
const mcpServerMetrics = useMemo(() => processActivityData(userSpendData, "mcp_servers", teams), [userSpendData, teams]);
return ( return (
<div style={{ width: "100%" }} className="p-8 relative"> <div style={{ width: "100%" }} className="p-8 relative">
@ -626,12 +633,6 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
{userSpendData.metadata?.total_failed_requests?.toLocaleString() || 0} {userSpendData.metadata?.total_failed_requests?.toLocaleString() || 0}
</Text> </Text>
</Card> </Card>
<Card>
<Title>Total Tokens</Title>
<Text className="text-2xl font-bold mt-2">
{userSpendData.metadata?.total_tokens?.toLocaleString() || 0}
</Text>
</Card>
<Card> <Card>
<Title>Average Cost per Request</Title> <Title>Average Cost per Request</Title>
<Text className="text-2xl font-bold mt-2"> <Text className="text-2xl font-bold mt-2">
@ -642,7 +643,51 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
)} )}
</Text> </Text>
</Card> </Card>
<Card
className="cursor-pointer hover:bg-gray-50 transition-colors"
onClick={() => setShowTokenBreakdown(!showTokenBreakdown)}
>
<div className="flex items-center gap-2">
<Title>Total Tokens</Title>
{showTokenBreakdown ? (
<DownOutlined className="text-gray-400 text-xs" />
) : (
<RightOutlined className="text-gray-400 text-xs" />
)}
</div>
<Text className="text-2xl font-bold mt-2">
{userSpendData.metadata?.total_tokens?.toLocaleString() || 0}
</Text>
</Card>
</Grid> </Grid>
{showTokenBreakdown && (
<Grid numItems={4} className="gap-4 mt-4">
<Card>
<Title>Input Tokens</Title>
<Text className="text-2xl font-bold mt-2 text-blue-600">
{userSpendData.metadata?.total_prompt_tokens?.toLocaleString() || 0}
</Text>
</Card>
<Card>
<Title>Output Tokens</Title>
<Text className="text-2xl font-bold mt-2 text-cyan-600">
{userSpendData.metadata?.total_completion_tokens?.toLocaleString() || 0}
</Text>
</Card>
<Card>
<Title>Cache Read Tokens</Title>
<Text className="text-2xl font-bold mt-2 text-green-600">
{userSpendData.metadata?.total_cache_read_input_tokens?.toLocaleString() || 0}
</Text>
</Card>
<Card>
<Title>Cache Write Tokens</Title>
<Text className="text-2xl font-bold mt-2 text-purple-600">
{userSpendData.metadata?.total_cache_creation_input_tokens?.toLocaleString() || 0}
</Text>
</Card>
</Grid>
)}
</Card> </Card>
</Col> </Col>
@ -654,9 +699,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
<ChartLoader isDateChanging={isDateChanging} /> <ChartLoader isDateChanging={isDateChanging} />
) : ( ) : (
<BarChart <BarChart
data={[...userSpendData.results].sort( data={sortedDailyResults}
(a, b) => new Date(a.date).getTime() - new Date(b.date).getTime(),
)}
index="date" index="date"
categories={["metrics.spend"]} categories={["metrics.spend"]}
colors={["cyan"]} colors={["cyan"]}
@ -688,7 +731,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
<Card className="h-full"> <Card className="h-full">
<Title>Top Virtual Keys</Title> <Title>Top Virtual Keys</Title>
<TopKeyView <TopKeyView
topKeys={getTopKeys(topKeysLimit)} topKeys={topKeys}
teams={null} teams={null}
topKeysLimit={topKeysLimit} topKeysLimit={topKeysLimit}
setTopKeysLimit={setTopKeysLimit} setTopKeysLimit={setTopKeysLimit}
@ -739,8 +782,8 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
{(() => { {(() => {
const modelData = const modelData =
modelViewType === "groups" modelViewType === "groups"
? getTopModelGroups(topModelsLimit) ? topModelGroups
: getTopModels(topModelsLimit); : topModels;
return ( return (
<BarChart <BarChart
className="mt-4" className="mt-4"
@ -784,7 +827,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
<SpendByProvider <SpendByProvider
loading={loading} loading={loading}
isDateChanging={isDateChanging} isDateChanging={isDateChanging}
providerSpend={getProviderSpend()} providerSpend={providerSpend}
/> />
</Col> </Col>

View file

@ -555,6 +555,94 @@ it("should display 'Default Proxy Admin' for created_by when value is 'default_u
}); });
it("should display created_by_user email in 'Created By' column when available", async () => {
const keyWithCreatedByUser = {
...mockKey,
created_by: "some-uuid-1234",
created_by_user: {
user_id: "some-uuid-1234",
user_email: "creator@example.com",
user_alias: null,
},
};
mockUseFilterLogic.mockReturnValue({
filters: {
"Team ID": "",
"Organization ID": "",
"Key Alias": "",
"User ID": "",
"Sort By": "created_at",
"Sort Order": "desc",
},
filteredKeys: [keyWithCreatedByUser],
allTeams: [mockTeam],
allOrganizations: [mockOrganization],
handleFilterChange: vi.fn(),
handleFilterReset: vi.fn(),
});
const mockProps = {
teams: [mockTeam],
organizations: [mockOrganization],
onSortChange: vi.fn(),
currentSort: {
sortBy: "created_at",
sortOrder: "desc" as const,
},
};
renderWithProviders(<VirtualKeysTable {...mockProps} />);
await waitFor(() => {
expect(screen.getByText("creator@example.com")).toBeInTheDocument();
});
});
it("should display created_by_user alias over email when both available", async () => {
const keyWithCreatedByUser = {
...mockKey,
created_by: "some-uuid-1234",
created_by_user: {
user_id: "some-uuid-1234",
user_email: "creator@example.com",
user_alias: "The Creator",
},
};
mockUseFilterLogic.mockReturnValue({
filters: {
"Team ID": "",
"Organization ID": "",
"Key Alias": "",
"User ID": "",
"Sort By": "created_at",
"Sort Order": "desc",
},
filteredKeys: [keyWithCreatedByUser],
allTeams: [mockTeam],
allOrganizations: [mockOrganization],
handleFilterChange: vi.fn(),
handleFilterReset: vi.fn(),
});
const mockProps = {
teams: [mockTeam],
organizations: [mockOrganization],
onSortChange: vi.fn(),
currentSort: {
sortBy: "created_at",
sortOrder: "desc" as const,
},
};
renderWithProviders(<VirtualKeysTable {...mockProps} />);
await waitFor(() => {
expect(screen.getByText("The Creator")).toBeInTheDocument();
});
});
it("should render table without crashing when models is null", async () => { it("should render table without crashing when models is null", async () => {
const keyWithNullModels = { const keyWithNullModels = {
...mockKey, ...mockKey,

View file

@ -311,25 +311,40 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
cell: (info) => { cell: (info) => {
const userId = info.getValue() as string | null; const userId = info.getValue() as string | null;
if (!userId) return "-"; if (!userId) return "-";
const key = info.row.original;
const createdByUser = key.created_by_user;
const userAlias = createdByUser?.user_alias ?? null;
const userEmail = createdByUser?.user_email ?? null;
const isDefaultAdmin = userId === "default_user_id"; const isDefaultAdmin = userId === "default_user_id";
const displayValue = userAlias || userEmail || userId;
const width = 160; const width = 160;
const popoverContent = ( const popoverContent = (
<div className="flex flex-col gap-2 text-xs min-w-[200px] max-w-[300px]"> <div className="flex flex-col gap-2 text-xs min-w-[200px] max-w-[300px]">
<div className="flex flex-col min-w-0"> {[
<span className="text-gray-400">User ID</span> { label: "User Alias", value: userAlias },
<Typography.Text { label: "User Email", value: userEmail },
className="font-mono text-xs" { label: "User ID", value: userId },
ellipsis={{ tooltip: userId }} ].map(({ label, value }) => (
copyable <div key={label} className="flex flex-col min-w-0">
> <span className="text-gray-400">{label}</span>
{userId} {value ? (
</Typography.Text> <Typography.Text
</div> className="font-mono text-xs"
ellipsis={{ tooltip: value }}
copyable
>
{value}
</Typography.Text>
) : (
<span className="font-mono">-</span>
)}
</div>
))}
</div> </div>
); );
if (isDefaultAdmin) { if (isDefaultAdmin && !userAlias && !userEmail) {
return ( return (
<Popover content={popoverContent} trigger="hover" placement="bottomLeft"> <Popover content={popoverContent} trigger="hover" placement="bottomLeft">
<span className="cursor-default"> <span className="cursor-default">
@ -345,7 +360,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
className="font-mono text-xs truncate block cursor-default" className="font-mono text-xs truncate block cursor-default"
style={{ maxWidth: width, overflow: "hidden" }} style={{ maxWidth: width, overflow: "hidden" }}
> >
{userId} {displayValue}
</span> </span>
</Popover> </Popover>
); );

View file

@ -362,7 +362,7 @@ export const formatKeyLabel = (modelData: KeyMetricWithMetadata, model: string,
// Process data function // Process data function
export const processActivityData = ( export const processActivityData = (
dailyActivity: { results: DailyData[] }, dailyActivity: { results: DailyData[] },
key: "models" | "api_keys" | "mcp_servers", key: "models" | "api_keys" | "mcp_servers" | "entities",
teams: Team[] = [], teams: Team[] = [],
): Record<string, ModelActivityData> => { ): Record<string, ModelActivityData> => {
const modelMetrics: Record<string, ModelActivityData> = {}; const modelMetrics: Record<string, ModelActivityData> = {};
@ -371,7 +371,11 @@ export const processActivityData = (
Object.entries(day.breakdown[key] || {}).forEach(([model, modelData]) => { Object.entries(day.breakdown[key] || {}).forEach(([model, modelData]) => {
if (!modelMetrics[model]) { if (!modelMetrics[model]) {
modelMetrics[model] = { modelMetrics[model] = {
label: key === "api_keys" ? formatKeyLabel(modelData as KeyMetricWithMetadata, model, teams) : model, label: key === "api_keys"
? formatKeyLabel(modelData as KeyMetricWithMetadata, model, teams)
: key === "entities"
? ((modelData as any).metadata?.agent_name || (modelData as any).metadata?.team_alias || model)
: model,
total_requests: 0, total_requests: 0,
total_successful_requests: 0, total_successful_requests: 0,
total_failed_requests: 0, total_failed_requests: 0,

View file

@ -19,19 +19,21 @@ import { isAdminRole } from "@/utils/roles";
import AgentInfoView from "./agents/agent_info"; import AgentInfoView from "./agents/agent_info";
import NotificationsManager from "./molecules/notifications_manager"; import NotificationsManager from "./molecules/notifications_manager";
import { Agent, AgentKeyInfo } from "./agents/types"; import { Agent, AgentKeyInfo } from "./agents/types";
import { Team } from "./key_team_helpers/key_list";
import { formatNumberWithCommas } from "@/utils/dataUtils"; import { formatNumberWithCommas } from "@/utils/dataUtils";
import TableIconActionButton from "./common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; import TableIconActionButton from "./common_components/IconActionButton/TableIconActionButtons/TableIconActionButton";
interface AgentsPanelProps { interface AgentsPanelProps {
accessToken: string | null; accessToken: string | null;
userRole?: string; userRole?: string;
teams?: Team[] | null;
} }
interface AgentsResponse { interface AgentsResponse {
agents: Agent[]; agents: Agent[];
} }
const AgentsPanel: React.FC<AgentsPanelProps> = ({ accessToken, userRole }) => { const AgentsPanel: React.FC<AgentsPanelProps> = ({ accessToken, userRole, teams }) => {
const [agentsList, setAgentsList] = useState<Agent[]>([]); const [agentsList, setAgentsList] = useState<Agent[]>([]);
const [keyInfoMap, setKeyInfoMap] = useState<Record<string, AgentKeyInfo>>({}); const [keyInfoMap, setKeyInfoMap] = useState<Record<string, AgentKeyInfo>>({});
const [isAddModalVisible, setIsAddModalVisible] = useState(false); const [isAddModalVisible, setIsAddModalVisible] = useState(false);
@ -282,6 +284,7 @@ const AgentsPanel: React.FC<AgentsPanelProps> = ({ accessToken, userRole }) => {
onClose={handleCloseModal} onClose={handleCloseModal}
accessToken={accessToken} accessToken={accessToken}
onSuccess={handleSuccess} onSuccess={handleSuccess}
teams={teams}
/> />
{agentToDelete && ( {agentToDelete && (

View file

@ -6,6 +6,7 @@ import CreatedKeyDisplay from "../shared/CreatedKeyDisplay";
import { import {
createAgentCall, createAgentCall,
getAgentCreateMetadata, getAgentCreateMetadata,
getAgentsList,
keyCreateForAgentCall, keyCreateForAgentCall,
keyListCall, keyListCall,
keyUpdateCall, keyUpdateCall,
@ -14,11 +15,14 @@ import {
} from "../networking"; } from "../networking";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
import { Team } from "../key_team_helpers/key_list";
import TeamDropdown from "../common_components/team_dropdown";
import AgentFormFields from "./agent_form_fields"; import AgentFormFields from "./agent_form_fields";
import DynamicAgentFormFields, { buildDynamicAgentData } from "./dynamic_agent_form_fields"; import DynamicAgentFormFields, { buildDynamicAgentData } from "./dynamic_agent_form_fields";
import { getDefaultFormValues, buildAgentDataFromForm } from "./agent_config"; import { getDefaultFormValues, buildAgentDataFromForm } from "./agent_config";
import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; import MCPServerSelector from "../mcp_server_management/MCPServerSelector";
import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions"; import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions";
import GuardrailSelector from "../guardrails/GuardrailSelector";
const { Step } = Steps; const { Step } = Steps;
@ -29,6 +33,7 @@ interface AddAgentFormProps {
onClose: () => void; onClose: () => void;
accessToken: string | null; accessToken: string | null;
onSuccess: () => void; onSuccess: () => void;
teams?: Team[] | null;
} }
const AddAgentForm: React.FC<AddAgentFormProps> = ({ const AddAgentForm: React.FC<AddAgentFormProps> = ({
@ -36,6 +41,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
onClose, onClose,
accessToken, accessToken,
onSuccess, onSuccess,
teams,
}) => { }) => {
const { userId, userRole } = useAuthorized(); const { userId, userRole } = useAuthorized();
const [form] = Form.useForm(); const [form] = Form.useForm();
@ -45,7 +51,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
const [agentTypeMetadata, setAgentTypeMetadata] = useState<AgentCreateInfo[]>([]); const [agentTypeMetadata, setAgentTypeMetadata] = useState<AgentCreateInfo[]>([]);
const [loadingMetadata, setLoadingMetadata] = useState(false); const [loadingMetadata, setLoadingMetadata] = useState(false);
// Step 1: key assignment state // Step 3: key assignment state
const [keyAssignOption, setKeyAssignOption] = useState<"create_new" | "existing_key" | "skip">("create_new"); const [keyAssignOption, setKeyAssignOption] = useState<"create_new" | "existing_key" | "skip">("create_new");
const [newKeyName, setNewKeyName] = useState<string>(""); const [newKeyName, setNewKeyName] = useState<string>("");
const [newKeyModels, setNewKeyModels] = useState<string[]>([]); const [newKeyModels, setNewKeyModels] = useState<string[]>([]);
@ -54,8 +60,10 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
const [loadingKeys, setLoadingKeys] = useState(false); const [loadingKeys, setLoadingKeys] = useState(false);
const [availableModels, setAvailableModels] = useState<string[]>([]); const [availableModels, setAvailableModels] = useState<string[]>([]);
const [loadingModels, setLoadingModels] = useState(false); const [loadingModels, setLoadingModels] = useState(false);
const [availableAgents, setAvailableAgents] = useState<{agent_id: string; agent_name: string}[]>([]);
const [loadingAgents, setLoadingAgents] = useState(false);
// Step 2: results // Step 4: results
const [createdAgentName, setCreatedAgentName] = useState<string>(""); const [createdAgentName, setCreatedAgentName] = useState<string>("");
const [createdKeyValue, setCreatedKeyValue] = useState<string | null>(null); const [createdKeyValue, setCreatedKeyValue] = useState<string | null>(null);
const [assignedKeyAlias, setAssignedKeyAlias] = useState<string | null>(null); const [assignedKeyAlias, setAssignedKeyAlias] = useState<string | null>(null);
@ -82,9 +90,9 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
fetchMetadata(); fetchMetadata();
}, []); }, []);
// Fetch existing keys when assign key step becomes active (step 2) // Fetch existing keys when Agent Management step becomes active (step 3)
useEffect(() => { useEffect(() => {
if (currentStep === 2 && accessToken && existingKeys.length === 0) { if (currentStep === 3 && accessToken && existingKeys.length === 0) {
const fetchKeys = async () => { const fetchKeys = async () => {
setLoadingKeys(true); setLoadingKeys(true);
try { try {
@ -100,9 +108,9 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
} }
}, [currentStep, accessToken]); }, [currentStep, accessToken]);
// Fetch available models when Assign Key step is active (same list as key generation) // Fetch available models when Agent Management step is active (same list as key generation)
useEffect(() => { useEffect(() => {
if (currentStep !== 2 || !accessToken || !userId || !userRole) return; if ((currentStep !== 1 && currentStep !== 3) || !accessToken || !userId || !userRole) return;
let cancelled = false; let cancelled = false;
setLoadingModels(true); setLoadingModels(true);
modelAvailableCall(accessToken, userId, userRole) modelAvailableCall(accessToken, userId, userRole)
@ -125,6 +133,25 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
}; };
}, [currentStep, accessToken, userId, userRole]); }, [currentStep, accessToken, userId, userRole]);
useEffect(() => {
if (currentStep !== 1 || !accessToken) return;
let cancelled = false;
setLoadingAgents(true);
getAgentsList(accessToken)
.then((response) => {
if (cancelled) return;
const agents = response?.agents ?? [];
setAvailableAgents(agents.map((a: any) => ({ agent_id: a.agent_id, agent_name: a.agent_name })));
})
.catch((error) => {
if (!cancelled) console.error("Error fetching agents:", error);
})
.finally(() => {
if (!cancelled) setLoadingAgents(false);
});
return () => { cancelled = true; };
}, [currentStep, accessToken]);
const selectedAgentTypeInfo = agentTypeMetadata.find( const selectedAgentTypeInfo = agentTypeMetadata.find(
(info) => info.agent_type === agentType (info) => info.agent_type === agentType
); );
@ -207,11 +234,14 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
// Build object_permission from MCP Tools step (allowed_mcp_servers_and_groups, mcp_tool_permissions) // Build object_permission from MCP Tools step (allowed_mcp_servers_and_groups, mcp_tool_permissions)
const mcpServersAndGroups = values.allowed_mcp_servers_and_groups; const mcpServersAndGroups = values.allowed_mcp_servers_and_groups;
const mcpToolPermissions = values.mcp_tool_permissions || {}; const mcpToolPermissions = values.mcp_tool_permissions || {};
if ( const entitlementModels = values.entitlement_models || [];
mcpServersAndGroups && const entitlementAgents = values.entitlement_agents || [];
(mcpServersAndGroups.servers?.length > 0 || mcpServersAndGroups.accessGroups?.length > 0) || const hasObjectPermission =
Object.keys(mcpToolPermissions).length > 0 (mcpServersAndGroups?.servers?.length > 0 || mcpServersAndGroups?.accessGroups?.length > 0) ||
) { Object.keys(mcpToolPermissions).length > 0 ||
entitlementModels.length > 0 ||
entitlementAgents.length > 0;
if (hasObjectPermission) {
agentData.object_permission = {}; agentData.object_permission = {};
if (mcpServersAndGroups?.servers?.length > 0) { if (mcpServersAndGroups?.servers?.length > 0) {
agentData.object_permission.mcp_servers = mcpServersAndGroups.servers; agentData.object_permission.mcp_servers = mcpServersAndGroups.servers;
@ -222,6 +252,12 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
if (Object.keys(mcpToolPermissions).length > 0) { if (Object.keys(mcpToolPermissions).length > 0) {
agentData.object_permission.mcp_tool_permissions = mcpToolPermissions; agentData.object_permission.mcp_tool_permissions = mcpToolPermissions;
} }
if (entitlementModels.length > 0) {
agentData.object_permission.models = entitlementModels;
}
if (entitlementAgents.length > 0) {
agentData.object_permission.agents = entitlementAgents;
}
} }
// Wire trace-id flags and budget controls into agent litellm_params (before create call) // Wire trace-id flags and budget controls into agent litellm_params (before create call)
@ -237,6 +273,17 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
} }
} }
const selectedGuardrails = values.guardrails || [];
if (selectedGuardrails.length > 0) {
if (!agentData.litellm_params) agentData.litellm_params = {};
agentData.litellm_params.guardrails = selectedGuardrails;
}
const selectedTeamId = values.team_id || null;
if (selectedTeamId) {
agentData.team_id = selectedTeamId;
}
const agentResponse = await createAgentCall(accessToken, agentData); const agentResponse = await createAgentCall(accessToken, agentData);
const agentId: string = agentResponse.agent_id; const agentId: string = agentResponse.agent_id;
const agentName: string = agentResponse.agent_name || values.agent_name || agentId; const agentName: string = agentResponse.agent_name || values.agent_name || agentId;
@ -248,6 +295,8 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
agentId, agentId,
newKeyName, newKeyName,
newKeyModels, newKeyModels,
undefined,
selectedTeamId,
); );
setCreatedKeyValue(keyResponse.key || null); setCreatedKeyValue(keyResponse.key || null);
} else if (keyAssignOption === "existing_key") { } else if (keyAssignOption === "existing_key") {
@ -264,7 +313,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
setAssignedKeyAlias(keyInfo?.key_alias || selectedExistingKey.slice(0, 12) + "…"); setAssignedKeyAlias(keyInfo?.key_alias || selectedExistingKey.slice(0, 12) + "…");
} }
setCurrentStep(3); setCurrentStep(4);
onSuccess(); onSuccess();
} catch (error) { } catch (error) {
console.error("Error creating agent:", error); console.error("Error creating agent:", error);
@ -293,11 +342,54 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
onClose(); onClose();
}; };
const renderMCPToolsStep = () => ( const renderEntitlementsStep = () => (
<div className="space-y-4"> <div className="space-y-4">
<p className="text-sm text-gray-600"> <p className="text-sm text-gray-600">
Optionally restrict which MCP servers and tools this agent can use. Leave empty to allow all (subject to key/team permissions). Configure which models, agents, and MCP tools this agent is allowed to use. Leave fields empty to allow all (subject to key/team permissions).
</p> </p>
<Form.Item
label={<span className="text-sm font-medium text-gray-700">Allowed Models</span>}
name="entitlement_models"
tooltip="Restrict which models this agent can call. Leave empty to allow all."
>
<Select
mode="tags"
style={{ width: "100%" }}
placeholder={loadingModels ? "Loading models..." : "Select models (leave empty for all)"}
tokenSeparators={[","]}
loading={loadingModels}
showSearch
options={availableModels.map((m) => ({
label: getModelDisplayName(m),
value: m,
}))}
/>
</Form.Item>
<Form.Item
label={<span className="text-sm font-medium text-gray-700">Allowed Agents (Sub-Agents)</span>}
name="entitlement_agents"
tooltip="Restrict which other agents this agent can invoke as sub-agents. Leave empty to allow all."
>
<Select
mode="multiple"
style={{ width: "100%" }}
placeholder={loadingAgents ? "Loading agents..." : "Select agents (leave empty for all)"}
loading={loadingAgents}
showSearch
filterOption={(input, option) =>
(option?.label as string ?? "").toLowerCase().includes(input.toLowerCase())
}
options={availableAgents.map((a) => ({
label: a.agent_name,
value: a.agent_id,
}))}
/>
</Form.Item>
<Divider className="my-2" />
<Form.Item <Form.Item
label={ label={
<span> <span>
@ -338,122 +430,137 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
</div> </div>
)} )}
</Form.Item> </Form.Item>
</div>
);
<Collapse ghost className="mt-6" items={[ const renderObservabilityStep = () => (
{ <div className="space-y-6">
key: "tracing", <div>
label: <span className="text-sm font-medium text-gray-700">Tracing</span>, <h4 className="text-sm font-medium text-gray-700 mb-3">Tracing</h4>
children: ( <div className="space-y-4">
<div className="space-y-4"> <div className="flex items-center justify-between">
<div className="flex items-center justify-between"> <div>
<div> <span className="text-sm font-medium text-gray-700">
<span className="text-sm font-medium text-gray-700"> Require x-litellm-trace-id on calls TO this agent
Require x-litellm-trace-id on calls TO this agent </span>
</span> <p className="text-xs text-gray-500 mt-1">
<p className="text-xs text-gray-500 mt-1"> Only accept this agent being invoked with a trace-id (e.g. when used as a sub-agent).
Only accept this agent being invoked with a trace-id (e.g. when used as a sub-agent).
</p>
</div>
<Switch
checked={requireTraceIdInbound}
onChange={setRequireTraceIdInbound}
/>
</div>
<div className="flex items-center justify-between">
<div>
<span className="text-sm font-medium text-gray-700">
Require x-litellm-trace-id on calls BY this agent
</span>
<p className="text-xs text-gray-500 mt-1">
Requires LLM/MCP calls made by this agent to include x-litellm-trace-id for session tracking.
</p>
</div>
<Switch
checked={requireTraceIdOutbound}
onChange={(checked) => {
setRequireTraceIdOutbound(checked);
if (!checked) {
setMaxIterations(null);
setMaxBudgetPerSession(null);
}
}}
/>
</div>
</div>
),
},
{
key: "budgets_and_rate_limits",
label: <span className="text-sm font-medium text-gray-700">Budgets &amp; Rate Limits</span>,
children: (
<div className="space-y-4">
{!requireTraceIdOutbound && (
<div className="p-3 bg-yellow-50 border border-yellow-200 rounded-lg text-sm text-yellow-800">
Enable &quot;Require x-litellm-trace-id on calls BY this agent&quot; in Tracing to configure budgets and rate limits.
</div>
)}
<div className="text-sm font-medium text-gray-700">Session Budgets</div>
<div className="grid grid-cols-2 gap-4">
<div>
<label className="text-sm text-gray-600 block mb-1">Max Iterations</label>
<InputNumber
className="w-full"
min={1}
placeholder="e.g. 25"
disabled={!requireTraceIdOutbound}
value={maxIterations}
onChange={(val) => setMaxIterations(val)}
/>
<p className="text-xs text-gray-400 mt-1">Hard cap on LLM calls per session</p>
</div>
<div>
<label className="text-sm text-gray-600 block mb-1">Max Budget Per Session ($)</label>
<InputNumber
className="w-full"
min={0.01}
step={0.5}
placeholder="e.g. 5.00"
disabled={!requireTraceIdOutbound}
value={maxBudgetPerSession}
onChange={(val) => setMaxBudgetPerSession(val)}
/>
<p className="text-xs text-gray-400 mt-1">Max spend per trace before returning 429</p>
</div>
</div>
<Divider className="my-2" />
<div className="text-sm font-medium text-gray-700">Agent Rate Limits</div>
<p className="text-xs text-gray-500">
Global rate limits applied across all callers of this agent.
</p> </p>
<div className="grid grid-cols-2 gap-4">
<Form.Item label="TPM Limit" name="tpm_limit" className="mb-0">
<InputNumber className="w-full" min={0} placeholder="e.g. 100000" disabled={!requireTraceIdOutbound} />
</Form.Item>
<Form.Item label="RPM Limit" name="rpm_limit" className="mb-0">
<InputNumber className="w-full" min={0} placeholder="e.g. 100" disabled={!requireTraceIdOutbound} />
</Form.Item>
</div>
<div className="text-sm font-medium text-gray-700 mt-4">Per-Session Rate Limits</div>
<p className="text-xs text-gray-500">
Rate limits per session (x-litellm-trace-id). Each session gets its own counters.
</p>
<div className="grid grid-cols-2 gap-4">
<Form.Item label="Session TPM Limit" name="session_tpm_limit" className="mb-0">
<InputNumber className="w-full" min={0} placeholder="e.g. 10000" disabled={!requireTraceIdOutbound} />
</Form.Item>
<Form.Item label="Session RPM Limit" name="session_rpm_limit" className="mb-0">
<InputNumber className="w-full" min={0} placeholder="e.g. 20" disabled={!requireTraceIdOutbound} />
</Form.Item>
</div>
</div> </div>
), <Switch
}, checked={requireTraceIdInbound}
]} /> onChange={setRequireTraceIdInbound}
/>
</div>
<div className="flex items-center justify-between">
<div>
<span className="text-sm font-medium text-gray-700">
Require x-litellm-trace-id on calls BY this agent
</span>
<p className="text-xs text-gray-500 mt-1">
Requires LLM/MCP calls made by this agent to include x-litellm-trace-id for session tracking.
</p>
</div>
<Switch
checked={requireTraceIdOutbound}
onChange={(checked) => {
setRequireTraceIdOutbound(checked);
if (!checked) {
setMaxIterations(null);
setMaxBudgetPerSession(null);
}
}}
/>
</div>
</div>
</div>
<Divider className="my-0" />
<div>
<h4 className="text-sm font-medium text-gray-700 mb-3">Budgets &amp; Rate Limits</h4>
<div className="space-y-4">
{!requireTraceIdOutbound && (
<div className="p-3 bg-yellow-50 border border-yellow-200 rounded-lg text-sm text-yellow-800">
Enable &quot;Require x-litellm-trace-id on calls BY this agent&quot; in Tracing to configure budgets and rate limits.
</div>
)}
<div className="text-sm font-medium text-gray-700">Session Budgets</div>
<div className="grid grid-cols-2 gap-4">
<div>
<label className="text-sm text-gray-600 block mb-1">Max Iterations</label>
<InputNumber
className="w-full"
min={1}
placeholder="e.g. 25"
disabled={!requireTraceIdOutbound}
value={maxIterations}
onChange={(val) => setMaxIterations(val)}
/>
<p className="text-xs text-gray-400 mt-1">Hard cap on LLM calls per session</p>
</div>
<div>
<label className="text-sm text-gray-600 block mb-1">Max Budget Per Session ($)</label>
<InputNumber
className="w-full"
min={0.01}
step={0.5}
placeholder="e.g. 5.00"
disabled={!requireTraceIdOutbound}
value={maxBudgetPerSession}
onChange={(val) => setMaxBudgetPerSession(val)}
/>
<p className="text-xs text-gray-400 mt-1">Max spend per trace before returning 429</p>
</div>
</div>
<Divider className="my-2" />
<div className="text-sm font-medium text-gray-700">Agent Rate Limits</div>
<p className="text-xs text-gray-500">
Global rate limits applied across all callers of this agent.
</p>
<div className="grid grid-cols-2 gap-4">
<Form.Item label="TPM Limit" name="tpm_limit" className="mb-0">
<InputNumber className="w-full" min={0} placeholder="e.g. 100000" disabled={!requireTraceIdOutbound} />
</Form.Item>
<Form.Item label="RPM Limit" name="rpm_limit" className="mb-0">
<InputNumber className="w-full" min={0} placeholder="e.g. 100" disabled={!requireTraceIdOutbound} />
</Form.Item>
</div>
<div className="text-sm font-medium text-gray-700 mt-4">Per-Session Rate Limits</div>
<p className="text-xs text-gray-500">
Rate limits per session (x-litellm-trace-id). Each session gets its own counters.
</p>
<div className="grid grid-cols-2 gap-4">
<Form.Item label="Session TPM Limit" name="session_tpm_limit" className="mb-0">
<InputNumber className="w-full" min={0} placeholder="e.g. 10000" disabled={!requireTraceIdOutbound} />
</Form.Item>
<Form.Item label="Session RPM Limit" name="session_rpm_limit" className="mb-0">
<InputNumber className="w-full" min={0} placeholder="e.g. 20" disabled={!requireTraceIdOutbound} />
</Form.Item>
</div>
</div>
</div>
<Divider className="my-0" />
<div>
<h4 className="text-sm font-medium text-gray-700 mb-3">Guardrails</h4>
<p className="text-xs text-gray-500 mb-3">
Apply guardrails to this agent. Selected guardrails will run on all calls made by this agent.
</p>
<Form.Item name="guardrails" initialValue={[]}>
<GuardrailSelector
accessToken={accessToken ?? ""}
value={form.getFieldValue("guardrails") ?? []}
onChange={(selected: string[]) => form.setFieldsValue({ guardrails: selected })}
/>
</Form.Item>
</div>
</div> </div>
); );
@ -610,6 +717,19 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
</Tag> </Tag>
</div> </div>
<Form.Item
label={<span className="text-sm font-medium text-gray-700">Assign to Team</span>}
name="team_id"
tooltip="Optionally assign this agent to a team. The agent and its key will belong to the selected team."
>
<TeamDropdown
teams={teams}
loading={!teams}
/>
</Form.Item>
<Divider className="my-4" />
<div className="space-y-3"> <div className="space-y-3">
{/* Option: Create new key */} {/* Option: Create new key */}
<div <div
@ -645,25 +765,6 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
placeholder="e.g. my-agent-key" placeholder="e.g. my-agent-key"
/> />
</div> </div>
<div>
<label className="text-sm text-gray-600 block mb-1">
Allowed Models <span className="text-gray-400">(optional leave empty for all models)</span>
</label>
<Select
mode="tags"
style={{ width: "100%" }}
placeholder={loadingModels ? "Loading models..." : "e.g. gpt-4o, claude-3-5-sonnet"}
value={newKeyModels}
onChange={setNewKeyModels}
tokenSeparators={[","]}
loading={loadingModels}
showSearch
options={availableModels.map((m) => ({
label: getModelDisplayName(m),
value: m,
}))}
/>
</div>
</div> </div>
)} )}
</div> </div>
@ -783,8 +884,9 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
{/* Step indicator */} {/* Step indicator */}
<Steps current={currentStep} size="small" className="mb-8"> <Steps current={currentStep} size="small" className="mb-8">
<Step title="Configure" /> <Step title="Configure" />
<Step title="Agent Settings" /> <Step title="Entitlements" />
<Step title="Assign Key" /> <Step title="Governance" />
<Step title="Agent Management" />
<Step title="Ready" /> <Step title="Ready" />
</Steps> </Steps>
@ -793,21 +895,22 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
layout="vertical" layout="vertical"
initialValues={ initialValues={
agentType === "a2a" agentType === "a2a"
? { ...getDefaultFormValues(), allowed_mcp_servers_and_groups: { servers: [], accessGroups: [] }, mcp_tool_permissions: {} } ? { ...getDefaultFormValues(), allowed_mcp_servers_and_groups: { servers: [], accessGroups: [] }, mcp_tool_permissions: {}, entitlement_models: [], entitlement_agents: [], guardrails: [] }
: { allowed_mcp_servers_and_groups: { servers: [], accessGroups: [] }, mcp_tool_permissions: {} } : { allowed_mcp_servers_and_groups: { servers: [], accessGroups: [] }, mcp_tool_permissions: {}, entitlement_models: [], entitlement_agents: [], guardrails: [] }
} }
className="space-y-4" className="space-y-4"
> >
{currentStep === 0 && renderConfigureStep()} {currentStep === 0 && renderConfigureStep()}
{currentStep === 1 && renderMCPToolsStep()} {currentStep === 1 && renderEntitlementsStep()}
{currentStep === 2 && renderAssignKeyStep()} {currentStep === 2 && renderObservabilityStep()}
{currentStep === 3 && renderReadyStep()} {currentStep === 3 && renderAssignKeyStep()}
{currentStep === 4 && renderReadyStep()}
</Form> </Form>
{/* Footer navigation */} {/* Footer navigation */}
<div className="flex items-center justify-between pt-6 border-t border-gray-100 mt-6"> <div className="flex items-center justify-between pt-6 border-t border-gray-100 mt-6">
<div> <div>
{currentStep > 0 && currentStep < 3 && ( {currentStep > 0 && currentStep < 4 && (
<button <button
type="button" type="button"
onClick={handleBack} onClick={handleBack}
@ -818,7 +921,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
)} )}
</div> </div>
<div className="flex gap-3"> <div className="flex gap-3">
{currentStep < 3 && ( {currentStep < 4 && (
<Button variant="secondary" onClick={handleClose}> <Button variant="secondary" onClick={handleClose}>
Cancel Cancel
</Button> </Button>
@ -834,11 +937,16 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
</Button> </Button>
)} )}
{currentStep === 2 && ( {currentStep === 2 && (
<Button variant="primary" onClick={handleNext}>
Next
</Button>
)}
{currentStep === 3 && (
<Button variant="primary" loading={isSubmitting} onClick={handleCreateAgent}> <Button variant="primary" loading={isSubmitting} onClick={handleCreateAgent}>
{isSubmitting ? "Creating..." : "Create Agent →"} {isSubmitting ? "Creating..." : "Create Agent →"}
</Button> </Button>
)} )}
{currentStep === 3 && ( {currentStep === 4 && (
<Button variant="primary" onClick={handleClose}> <Button variant="primary" onClick={handleClose}>
Done Done
</Button> </Button>

View file

@ -0,0 +1,46 @@
import { describe, it, expect, vi } from "vitest";
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { Tooltip } from "./Tooltip";
vi.mock("@ant-design/icons", () => ({
QuestionCircleOutlined: (props: any) => <span data-testid="question-icon" {...props} />,
}));
describe("Tooltip", () => {
it("should render", () => {
render(<Tooltip content="Help text" />);
expect(screen.getByTestId("question-icon")).toBeInTheDocument();
});
it("should render children instead of the default icon when provided", () => {
render(<Tooltip content="Help text"><button>Info</button></Tooltip>);
expect(screen.getByRole("button", { name: /info/i })).toBeInTheDocument();
expect(screen.queryByTestId("question-icon")).not.toBeInTheDocument();
});
it("should show tooltip content on mouse enter", async () => {
const user = userEvent.setup();
render(<Tooltip content="Help text" />);
await user.hover(screen.getByTestId("question-icon"));
expect(screen.getByText("Help text")).toBeInTheDocument();
});
it("should hide tooltip content on mouse leave", async () => {
const user = userEvent.setup();
render(<Tooltip content="Help text" />);
await user.hover(screen.getByTestId("question-icon"));
expect(screen.getByText("Help text")).toBeInTheDocument();
await user.unhover(screen.getByTestId("question-icon"));
expect(screen.queryByText("Help text")).not.toBeInTheDocument();
});
it("should not show tooltip content before hovering", () => {
render(<Tooltip content="Help text" />);
expect(screen.queryByText("Help text")).not.toBeInTheDocument();
});
});

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