scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge branch 'litellm_oss_staging_02_03_2026' into litellm_dev_02_02_2026_p2

Browse source
This commit is contained in:
Krish Dholakia 2026-02-03 19:57:15 -08:00 • committed by GitHub
commit c05cb5c27f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
147 changed files with 10392 additions and 2317 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
ci_cd/security_scans.sh
View file

@ -154,6 +154,7 @@ run_grype_scans() {
"CVE-2025-15367" # No fix available yet
"CVE-2025-12781" # No fix available yet
"CVE-2025-11468" # No fix available yet
"CVE-2026-1299" # Python 3.13 email module header injection - not applicable, LiteLLM doesn't use BytesGenerator for email serialization
)
# Build JSON array of allowlisted CVE IDs for jq

92
docs/my-website/blog/sub_millisecond_proxy_overhead/index.md Normal file
View file

@ -0,0 +1,92 @@
---
slug: sub-millisecond-proxy-overhead
title: "Achieving Sub-Millisecond Proxy Overhead"
date: 2026-02-02T10:00:00
authors:
- name: Alexsander Hamir
title: "Performance Engineer, LiteLLM"
url: https://www.linkedin.com/in/alexsander-baptista/
image_url: https://github.com/AlexsanderHamir.png
- name: Krrish Dholakia
title: "CEO, LiteLLM"
url: https://www.linkedin.com/in/krish-d/
image_url: https://pbs.twimg.com/profile_images/1298587542745358340/DZv3Oj-h_400x400.jpg
- name: Ishaan Jaff
title: "CTO, LiteLLM"
url: https://www.linkedin.com/in/reffajnaahsi/
image_url: https://pbs.twimg.com/profile_images/1613813310264340481/lz54oEiB_400x400.jpg
description: "Our Q1 performance target and architectural direction for achieving sub-millisecond proxy overhead on modest hardware."
tags: [performance, architecture]
hide_table_of_contents: false
---
![Sidecar architecture: Python control plane vs. sidecar hot path](https://raw.githubusercontent.com/AlexsanderHamir/assets/main/Screenshot%202026-02-02%20172554.png)
# Achieving Sub-Millisecond Proxy Overhead
## Introduction
Our Q1 performance target is to aggressively move toward sub-millisecond proxy overhead on a single instance with 4 CPUs and 8 GB of RAM, and to continue pushing that boundary over time. Our broader goal is to make LiteLLM inexpensive to deploy, lightweight, and fast. This post outlines the architectural direction behind that effort.
Proxy overhead refers to the latency introduced by LiteLLM itself, independent of the upstream provider.
To measure it, we run the same workload directly against the provider and through LiteLLM at identical QPS (for example, 1,000 QPS) and compare the latency delta. To reduce noise, the load generator, LiteLLM, and a mock LLM endpoint all run on the same machine, ensuring the difference reflects proxy overhead rather than network latency.
---
## Where We're Coming From
Under the same benchmark originally conducted by [TensorZero](https://www.tensorzero.com/docs/gateway/benchmarks), LiteLLM previously failed at around 1,000 QPS.
That is no longer the case. Today, LiteLLM can be stress-tested at 1,000 QPS with no failures and can scale up to 5,000 QPS without failures on a 4-CPU, 8-GB RAM single instance setup.
This establishes a more up to date baseline and provides useful context as we continue working on proxy overhead and overall performance.
---
## Design Choice
Achieving sub-millisecond proxy overhead with a Python-based system requires being deliberate about where work happens.
Python is a strong fit for flexibility and extensibility: provider abstraction, configuration-driven routing, and a rich callback ecosystem. These are areas where development velocity and correctness matter more than raw throughput.
At higher request rates, however, certain classes of work become expensive when executed inside the Python process on every request. Rather than rewriting LiteLLM or introducing complex deployment requirements, we adopt an optional **sidecar architecture**.
This architectural change is how we intend to make LiteLLM **permanently fast**. While it supports our near-term performance targets, it is a long-term investment.
Python continues to own:
- Request validation and normalization
- Model and provider selection
- Callbacks and integrations
The sidecar owns **performance-critical execution**, such as:
- Efficient request forwarding
- Connection reuse and pooling
- Enforcing timeouts and limits
- Aggregating high-frequency metrics
This separation allows each component to focus on what it does best: Python acts as the control plane, while the sidecar handles the hot path.
---
### Why the Sidecar Is Optional
The sidecar is intentionally **optional**.
This allows us to ship it incrementally, validate it under real-world workloads, and avoid making it a hard dependency before it is fully battle-tested across all LiteLLM features.
Just as importantly, this ensures that self-hosting LiteLLM remains simple. The sidecar is bundled and started automatically, requires no additional infrastructure, and can be disabled entirely. From a user's perspective, LiteLLM continues to behave like a single service.
As of today, the sidecar is an optimization, not a requirement.
---
## Conclusion
Sub-millisecond proxy overhead is not achieved through a single optimization, but through architectural changes.
By keeping Python focused on orchestration and extensibility, and offloading performance-critical execution to a sidecar, we establish a foundation for making LiteLLM **permanently fast over time**—even on modest hardware such as a 1-CPU, 2-GB RAM instance, while keeping deployment and self-hosting simple.
This work extends beyond Q1, and we will continue sharing benchmarks and updates as the architecture evolves.

113
docs/my-website/docs/a2a.md
View file

@ -68,116 +68,9 @@ Follow [this guide, to add your pydantic ai agent to LiteLLM Agent Gateway](./pr
## Invoking your Agents
Use the [A2A Python SDK](https://pypi.org/project/a2a-sdk) to invoke agents through LiteLLM.
This example shows how to:
1. **List available agents** - Query `/v1/agents` to see which agents your key can access
2. **Select an agent** - Pick an agent from the list
3. **Invoke via A2A** - Use the A2A protocol to send messages to the agent
```python showLineNumbers title="invoke_a2a_agent.py"
from uuid import uuid4
import httpx
import asyncio
from a2a.client import A2ACardResolver, A2AClient
from a2a.types import MessageSendParams, SendMessageRequest
# === CONFIGURE THESE ===
LITELLM_BASE_URL = "http://localhost:4000" # Your LiteLLM proxy URL
LITELLM_VIRTUAL_KEY = "sk-1234" # Your LiteLLM Virtual Key
# =======================
async def main():
headers = {"Authorization": f"Bearer {LITELLM_VIRTUAL_KEY}"}
async with httpx.AsyncClient(headers=headers) as client:
# Step 1: List available agents
response = await client.get(f"{LITELLM_BASE_URL}/v1/agents")
agents = response.json()
print("Available agents:")
for agent in agents:
print(f" - {agent['agent_name']} (ID: {agent['agent_id']})")
if not agents:
print("No agents available for this key")
return
# Step 2: Select an agent and invoke it
selected_agent = agents[0]
agent_id = selected_agent["agent_id"]
agent_name = selected_agent["agent_name"]
print(f"\nInvoking: {agent_name}")
# Step 3: Use A2A protocol to invoke the agent
base_url = f"{LITELLM_BASE_URL}/a2a/{agent_id}"
resolver = A2ACardResolver(httpx_client=client, base_url=base_url)
agent_card = await resolver.get_agent_card()
a2a_client = A2AClient(httpx_client=client, agent_card=agent_card)
request = SendMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
message={
"role": "user",
"parts": [{"kind": "text", "text": "Hello, what can you do?"}],
"messageId": uuid4().hex,
}
),
)
response = await a2a_client.send_message(request)
print(f"Response: {response.model_dump(mode='json', exclude_none=True, indent=4)}")
if __name__ == "__main__":
asyncio.run(main())
```
### Streaming Responses
For streaming responses, use `send_message_streaming`:
```python showLineNumbers title="invoke_a2a_agent_streaming.py"
from uuid import uuid4
import httpx
import asyncio
from a2a.client import A2ACardResolver, A2AClient
from a2a.types import MessageSendParams, SendStreamingMessageRequest
# === CONFIGURE THESE ===
LITELLM_BASE_URL = "http://localhost:4000" # Your LiteLLM proxy URL
LITELLM_VIRTUAL_KEY = "sk-1234" # Your LiteLLM Virtual Key
LITELLM_AGENT_NAME = "ij-local" # Agent name registered in LiteLLM
# =======================
async def main():
base_url = f"{LITELLM_BASE_URL}/a2a/{LITELLM_AGENT_NAME}"
headers = {"Authorization": f"Bearer {LITELLM_VIRTUAL_KEY}"}
async with httpx.AsyncClient(headers=headers) as httpx_client:
# Resolve agent card and create client
resolver = A2ACardResolver(httpx_client=httpx_client, base_url=base_url)
agent_card = await resolver.get_agent_card()
client = A2AClient(httpx_client=httpx_client, agent_card=agent_card)
# Send a streaming message
request = SendStreamingMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
message={
"role": "user",
"parts": [{"kind": "text", "text": "Hello, what can you do?"}],
"messageId": uuid4().hex,
}
),
)
# Stream the response
async for chunk in client.send_message_streaming(request):
print(chunk.model_dump(mode="json", exclude_none=True))
if __name__ == "__main__":
asyncio.run(main())
```
See the [Invoking A2A Agents](./a2a_invoking_agents) guide to learn how to call your agents using:
- **A2A SDK** - Native A2A protocol with full support for tasks and artifacts
- **OpenAI SDK** - Familiar `/chat/completions` interface with `a2a/` model prefix
## Tracking Agent Logs

280
docs/my-website/docs/a2a_invoking_agents.md Normal file
View file

@ -0,0 +1,280 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Invoking A2A Agents
Learn how to invoke A2A agents through LiteLLM using different methods.
:::tip Deploy Your Own A2A Agent
Want to test with your own agent? Deploy this template A2A agent powered by Google Gemini:
[**shin-bot-litellm/a2a-gemini-agent**](https://github.com/shin-bot-litellm/a2a-gemini-agent) - Simple deployable A2A agent with streaming support
:::
## A2A SDK
Use the [A2A Python SDK](https://pypi.org/project/a2a-sdk) to invoke agents through LiteLLM using the A2A protocol.
### Non-Streaming
This example shows how to:
1. **List available agents** - Query `/v1/agents` to see which agents your key can access
2. **Select an agent** - Pick an agent from the list
3. **Invoke via A2A** - Use the A2A protocol to send messages to the agent
```python showLineNumbers title="invoke_a2a_agent.py"
from uuid import uuid4
import httpx
import asyncio
from a2a.client import A2ACardResolver, A2AClient
from a2a.types import MessageSendParams, SendMessageRequest
# === CONFIGURE THESE ===
LITELLM_BASE_URL = "http://localhost:4000" # Your LiteLLM proxy URL
LITELLM_VIRTUAL_KEY = "sk-1234" # Your LiteLLM Virtual Key
# =======================
async def main():
headers = {"Authorization": f"Bearer {LITELLM_VIRTUAL_KEY}"}
async with httpx.AsyncClient(headers=headers) as client:
# Step 1: List available agents
response = await client.get(f"{LITELLM_BASE_URL}/v1/agents")
agents = response.json()
print("Available agents:")
for agent in agents:
print(f" - {agent['agent_name']} (ID: {agent['agent_id']})")
if not agents:
print("No agents available for this key")
return
# Step 2: Select an agent and invoke it
selected_agent = agents[0]
agent_id = selected_agent["agent_id"]
agent_name = selected_agent["agent_name"]
print(f"\nInvoking: {agent_name}")
# Step 3: Use A2A protocol to invoke the agent
base_url = f"{LITELLM_BASE_URL}/a2a/{agent_id}"
resolver = A2ACardResolver(httpx_client=client, base_url=base_url)
agent_card = await resolver.get_agent_card()
a2a_client = A2AClient(httpx_client=client, agent_card=agent_card)
request = SendMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
message={
"role": "user",
"parts": [{"kind": "text", "text": "Hello, what can you do?"}],
"messageId": uuid4().hex,
}
),
)
response = await a2a_client.send_message(request)
print(f"Response: {response.model_dump(mode='json', exclude_none=True, indent=4)}")
if __name__ == "__main__":
asyncio.run(main())
```
### Streaming
For streaming responses, use `send_message_streaming`:
```python showLineNumbers title="invoke_a2a_agent_streaming.py"
from uuid import uuid4
import httpx
import asyncio
from a2a.client import A2ACardResolver, A2AClient
from a2a.types import MessageSendParams, SendStreamingMessageRequest
# === CONFIGURE THESE ===
LITELLM_BASE_URL = "http://localhost:4000" # Your LiteLLM proxy URL
LITELLM_VIRTUAL_KEY = "sk-1234" # Your LiteLLM Virtual Key
LITELLM_AGENT_NAME = "ij-local" # Agent name registered in LiteLLM
# =======================
async def main():
base_url = f"{LITELLM_BASE_URL}/a2a/{LITELLM_AGENT_NAME}"
headers = {"Authorization": f"Bearer {LITELLM_VIRTUAL_KEY}"}
async with httpx.AsyncClient(headers=headers) as httpx_client:
# Resolve agent card and create client
resolver = A2ACardResolver(httpx_client=httpx_client, base_url=base_url)
agent_card = await resolver.get_agent_card()
client = A2AClient(httpx_client=httpx_client, agent_card=agent_card)
# Send a streaming message
request = SendStreamingMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
message={
"role": "user",
"parts": [{"kind": "text", "text": "Tell me a long story"}],
"messageId": uuid4().hex,
}
),
)
# Stream the response
async for chunk in client.send_message_streaming(request):
print(chunk.model_dump(mode="json", exclude_none=True))
if __name__ == "__main__":
asyncio.run(main())
```
## /chat/completions API (OpenAI SDK)
You can also invoke A2A agents using the familiar OpenAI SDK by using the `a2a/` model prefix.
### Non-Streaming
<Tabs>
<TabItem value="python" label="Python" default>
```python showLineNumbers title="openai_non_streaming.py"
import openai
client = openai.OpenAI(
api_key="sk-1234", # Your LiteLLM Virtual Key
base_url="http://localhost:4000" # Your LiteLLM proxy URL
)
response = client.chat.completions.create(
model="a2a/my-agent", # Use a2a/ prefix with your agent name
messages=[
{"role": "user", "content": "Hello, what can you do?"}
]
)
print(response.choices[0].message.content)
```
</TabItem>
<TabItem value="typescript" label="TypeScript">
```typescript showLineNumbers title="openai_non_streaming.ts"
import OpenAI from 'openai';
const client = new OpenAI({
apiKey: 'sk-1234', // Your LiteLLM Virtual Key
baseURL: 'http://localhost:4000' // Your LiteLLM proxy URL
});
const response = await client.chat.completions.create({
model: 'a2a/my-agent', // Use a2a/ prefix with your agent name
messages: [
{ role: 'user', content: 'Hello, what can you do?' }
]
});
console.log(response.choices[0].message.content);
```
</TabItem>
<TabItem value="curl" label="cURL">
```bash showLineNumbers title="curl_non_streaming.sh"
curl -X POST http://localhost:4000/v1/chat/completions \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"model": "a2a/my-agent",
"messages": [
{"role": "user", "content": "Hello, what can you do?"}
]
}'
```
</TabItem>
</Tabs>
### Streaming
<Tabs>
<TabItem value="python" label="Python" default>
```python showLineNumbers title="openai_streaming.py"
import openai
client = openai.OpenAI(
api_key="sk-1234", # Your LiteLLM Virtual Key
base_url="http://localhost:4000" # Your LiteLLM proxy URL
)
stream = client.chat.completions.create(
model="a2a/my-agent", # Use a2a/ prefix with your agent name
messages=[
{"role": "user", "content": "Tell me a long story"}
],
stream=True
)
for chunk in stream:
if chunk.choices[0].delta.content:
print(chunk.choices[0].delta.content, end="", flush=True)
```
</TabItem>
<TabItem value="typescript" label="TypeScript">
```typescript showLineNumbers title="openai_streaming.ts"
import OpenAI from 'openai';
const client = new OpenAI({
apiKey: 'sk-1234', // Your LiteLLM Virtual Key
baseURL: 'http://localhost:4000' // Your LiteLLM proxy URL
});
const stream = await client.chat.completions.create({
model: 'a2a/my-agent', // Use a2a/ prefix with your agent name
messages: [
{ role: 'user', content: 'Tell me a long story' }
],
stream: true
});
for await (const chunk of stream) {
const content = chunk.choices[0]?.delta?.content;
if (content) {
process.stdout.write(content);
}
}
```
</TabItem>
<TabItem value="curl" label="cURL">
```bash showLineNumbers title="curl_streaming.sh"
curl -X POST http://localhost:4000/v1/chat/completions \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"model": "a2a/my-agent",
"messages": [
{"role": "user", "content": "Tell me a long story"}
],
"stream": true
}'
```
</TabItem>
</Tabs>
## Key Differences
| Method | Use Case | Advantages |
|--------|----------|------------|
| **A2A SDK** | Native A2A protocol integration | • Full A2A protocol support<br/>• Access to task states and artifacts<br/>• Context management |
| **OpenAI SDK** | Familiar OpenAI-style interface | • Drop-in replacement for OpenAI calls<br/>• Easier migration from LLM to agent workflows<br/>• Works with existing OpenAI tooling |
:::tip Model Prefix
When using the OpenAI SDK, always prefix your agent name with `a2a/` (e.g., `a2a/my-agent`) to route requests to the A2A agent instead of an LLM provider.
:::

158
docs/my-website/docs/mcp_semantic_filter.md Normal file
View file

@ -0,0 +1,158 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# MCP Semantic Tool Filter
Automatically filter MCP tools by semantic relevance. When you have many MCP tools registered, LiteLLM semantically matches the user's query against tool descriptions and sends only the most relevant tools to the LLM.
## How It Works
Tool search shifts tool selection from a prompt-engineering problem to a retrieval problem. Instead of injecting a large static list of tools into every prompt, the semantic filter:
1. Builds a semantic index of all available MCP tools on startup
2. On each request, semantically matches the user's query against tool descriptions
3. Returns only the top-K most relevant tools to the LLM
This approach improves context efficiency, increases reliability by reducing tool confusion, and enables scalability to ecosystems with hundreds or thousands of MCP tools.
```mermaid
sequenceDiagram
participant Client
participant LiteLLM as LiteLLM Proxy
participant SemanticFilter as Semantic Filter
participant MCP as MCP Registry
participant LLM as LLM Provider
Note over LiteLLM,MCP: Startup: Build Semantic Index
LiteLLM->>MCP: Fetch all registered MCP tools
MCP->>LiteLLM: Return all tools (e.g., 50 tools)
LiteLLM->>SemanticFilter: Build semantic router with embeddings
SemanticFilter->>LLM: Generate embeddings for tool descriptions
LLM->>SemanticFilter: Return embeddings
Note over SemanticFilter: Index ready for fast lookup
Note over Client,LLM: Request: Semantic Tool Filtering
Client->>LiteLLM: POST /v1/responses with MCP tools
LiteLLM->>SemanticFilter: Expand MCP references (50 tools available)
SemanticFilter->>SemanticFilter: Extract user query from request
SemanticFilter->>LLM: Generate query embedding
LLM->>SemanticFilter: Return query embedding
SemanticFilter->>SemanticFilter: Match query against tool embeddings
SemanticFilter->>LiteLLM: Return top-K tools (e.g., 3 most relevant)
LiteLLM->>LLM: Forward request with filtered tools (3 tools)
LLM->>LiteLLM: Return response
LiteLLM->>Client: Response with headers<br/>x-litellm-semantic-filter: 50->3<br/>x-litellm-semantic-filter-tools: tool1,tool2,tool3
```
## Configuration
Enable semantic filtering in your LiteLLM config:
```yaml title="config.yaml" showLineNumbers
litellm_settings:
mcp_semantic_tool_filter:
enabled: true
embedding_model: "text-embedding-3-small" # Model for semantic matching
top_k: 5 # Max tools to return
similarity_threshold: 0.3 # Min similarity score
```
**Configuration Options:**
- `enabled` - Enable/disable semantic filtering (default: `false`)
- `embedding_model` - Model for generating embeddings (default: `"text-embedding-3-small"`)
- `top_k` - Maximum number of tools to return (default: `10`)
- `similarity_threshold` - Minimum similarity score for matches (default: `0.3`)
## Usage
Use MCP tools normally with the Responses API or Chat Completions. The semantic filter runs automatically:
<Tabs>
<TabItem value="responses" label="Responses API">
```bash title="Responses API with Semantic Filtering" showLineNumbers
curl --location 'http://localhost:4000/v1/responses' \
--header 'Content-Type: application/json' \
--header "Authorization: Bearer sk-1234" \
--data '{
"model": "gpt-4o",
"input": [
{
"role": "user",
"content": "give me TLDR of what BerriAI/litellm repo is about",
"type": "message"
}
],
"tools": [
{
"type": "mcp",
"server_url": "litellm_proxy",
"require_approval": "never"
}
],
"tool_choice": "required"
}'
```
</TabItem>
<TabItem value="chat" label="Chat Completions">
```bash title="Chat Completions with Semantic Filtering" showLineNumbers
curl --location 'http://localhost:4000/v1/chat/completions' \
--header 'Content-Type: application/json' \
--header "Authorization: Bearer sk-1234" \
--data '{
"model": "gpt-4o",
"messages": [
{"role": "user", "content": "Search Wikipedia for LiteLLM"}
],
"tools": [
{
"type": "mcp",
"server_url": "litellm_proxy"
}
]
}'
```
</TabItem>
</Tabs>
## Response Headers
The semantic filter adds diagnostic headers to every response:
```
x-litellm-semantic-filter: 10->3
x-litellm-semantic-filter-tools: wikipedia-fetch,github-search,slack-post
```
- **`x-litellm-semantic-filter`** - Shows before→after tool count (e.g., `10->3` means 10 tools were filtered down to 3)
- **`x-litellm-semantic-filter-tools`** - CSV list of the filtered tool names (max 150 chars, clipped with `...` if longer)
These headers help you understand which tools were selected for each request and verify the filter is working correctly.
## Example
If you have 50 MCP tools registered and make a request asking about Wikipedia, the semantic filter will:
1. Semantically match your query `"Search Wikipedia for LiteLLM"` against all 50 tool descriptions
2. Select the top 5 most relevant tools (e.g., `wikipedia-fetch`, `wikipedia-search`, etc.)
3. Pass only those 5 tools to the LLM
4. Add headers showing `x-litellm-semantic-filter: 50->5`
This dramatically reduces prompt size while ensuring the LLM has access to the right tools for the task.
## Performance
The semantic filter is optimized for production:
- Router builds once on startup (no per-request overhead)
- Semantic matching typically takes under 50ms
- Fails gracefully - returns all tools if filtering fails
- No impact on latency for requests without MCP tools
## Related
- [MCP Overview](./mcp.md) - Learn about MCP in LiteLLM
- [MCP Permission Management](./mcp_control.md) - Control tool access by key/team
- [Using MCP](./mcp_usage.md) - Complete MCP usage guide

2
docs/my-website/docs/providers/bedrock.md
View file

@ -9,7 +9,7 @@ ALL Bedrock models (Anthropic, Meta, Deepseek, Mistral, Amazon, etc.) are Suppor
| Description | Amazon Bedrock is a fully managed service that offers a choice of high-performing foundation models (FMs). |
| Provider Route on LiteLLM | `bedrock/`, [`bedrock/converse/`](#set-converse--invoke-route), [`bedrock/invoke/`](#set-invoke-route), [`bedrock/converse_like/`](#calling-via-internal-proxy), [`bedrock/llama/`](#deepseek-not-r1), [`bedrock/deepseek_r1/`](#deepseek-r1), [`bedrock/qwen3/`](#qwen3-imported-models), [`bedrock/qwen2/`](./bedrock_imported.md#qwen2-imported-models), [`bedrock/openai/`](./bedrock_imported.md#openai-compatible-imported-models-qwen-25-vl-etc), [`bedrock/moonshot`](./bedrock_imported.md#moonshot-kimi-k2-thinking) |
| Provider Doc | [Amazon Bedrock ↗](https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html) |
| Supported OpenAI Endpoints | `/chat/completions`, `/completions`, `/embeddings`, `/images/generations` |
| Supported OpenAI Endpoints | `/chat/completions`, `/completions`, `/embeddings`, `/images/generations`, `/v1/realtime`|
| Rerank Endpoint | `/rerank` |
| Pass-through Endpoint | [Supported](../pass_through/bedrock.md) |

6
docs/my-website/docs/tutorials/bedrock_realtime_with_audio.md → docs/my-website/docs/providers/bedrock_realtime_with_audio.md
View file

@ -1,8 +1,4 @@
# Call Bedrock Nova Sonic Realtime API with Audio Input/Output
:::info
Requires LiteLLM Proxy v1.70.1+
:::
# Bedrock Realtime API
## Overview

35
docs/my-website/docs/providers/github_copilot.md
View file

@ -35,11 +35,10 @@ from litellm import completion
response = completion(
model="github_copilot/gpt-4",
messages=[{"role": "user", "content": "Write a Python function to calculate fibonacci numbers"}],
extra_headers={
"editor-version": "vscode/1.85.1",
"Copilot-Integration-Id": "vscode-chat"
}
messages=[
{"role": "system", "content": "You are a helpful coding assistant"},
{"role": "user", "content": "Write a Python function to calculate fibonacci numbers"}
]
)
print(response)
```
@ -50,11 +49,7 @@ from litellm import completion
stream = completion(
model="github_copilot/gpt-4",
messages=[{"role": "user", "content": "Explain async/await in Python"}],
stream=True,
extra_headers={
"editor-version": "vscode/1.85.1",
"Copilot-Integration-Id": "vscode-chat"
}
stream=True
)
for chunk in stream:
@ -134,11 +129,7 @@ client = OpenAI(
# Non-streaming response
response = client.chat.completions.create(
model="github_copilot/gpt-4",
messages=[{"role": "user", "content": "How do I optimize this SQL query?"}],
extra_headers={
"editor-version": "vscode/1.85.1",
"Copilot-Integration-Id": "vscode-chat"
}
messages=[{"role": "user", "content": "How do I optimize this SQL query?"}]
)
print(response.choices[0].message.content)
@ -156,11 +147,7 @@ response = litellm.completion(
model="litellm_proxy/github_copilot/gpt-4",
messages=[{"role": "user", "content": "Review this code for bugs"}],
api_base="http://localhost:4000",
api_key="your-proxy-api-key",
extra_headers={
"editor-version": "vscode/1.85.1",
"Copilot-Integration-Id": "vscode-chat"
}
api_key="your-proxy-api-key"
)
print(response.choices[0].message.content)
@ -174,8 +161,6 @@ print(response.choices[0].message.content)
curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer your-proxy-api-key" \
-H "editor-version: vscode/1.85.1" \
-H "Copilot-Integration-Id: vscode-chat" \
-d '{
"model": "github_copilot/gpt-4",
"messages": [{"role": "user", "content": "Explain this error message"}]
@ -211,9 +196,11 @@ export GITHUB_COPILOT_API_KEY_FILE="api-key.json"
### Headers
GitHub Copilot supports various editor-specific headers:
LiteLLM automatically injects the required GitHub Copilot headers (simulating VSCode). You don't need to specify them manually.
```python showLineNumbers title="Common Headers"
If you want to override the defaults (e.g., to simulate a different editor), you can use `extra_headers`:
```python showLineNumbers title="Custom Headers (Optional)"
extra_headers = {
"editor-version": "vscode/1.85.1", # Editor version
"editor-plugin-version": "copilot/1.155.0", # Plugin version

8
docs/my-website/docs/proxy/config_settings.md
View file

@ -94,7 +94,7 @@ litellm_settings:
# /chat/completions, /completions, /embeddings, /audio/transcriptions
mode: default_off # if default_off, you need to opt in to caching on a per call basis
ttl: 600 # ttl for caching
disable_copilot_system_to_assistant: False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior.
disable_copilot_system_to_assistant: False # DEPRECATED - GitHub Copilot API supports system prompts.
callback_settings:
otel:
@ -197,7 +197,7 @@ router_settings:
| disable_add_transform_inline_image_block | boolean | For Fireworks AI models - if true, turns off the auto-add of `#transform=inline` to the url of the image_url, if the model is not a vision model. |
| disable_hf_tokenizer_download | boolean | If true, it defaults to using the openai tokenizer for all models (including huggingface models). |
| enable_json_schema_validation | boolean | If true, enables json schema validation for all requests. |
| disable_copilot_system_to_assistant | boolean | If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior. Useful for tools (like Claude Code) that send system messages, which Copilot does not support. |
| disable_copilot_system_to_assistant | boolean | **DEPRECATED** - GitHub Copilot API supports system prompts. |
### general_settings - Reference
@ -545,6 +545,9 @@ router_settings:
| DEFAULT_MAX_TOKENS | Default maximum tokens for LLM calls. Default is 4096
| DEFAULT_MAX_TOKENS_FOR_TRITON | Default maximum tokens for Triton models. Default is 2000
| DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE | Default maximum size for redis batch cache. Default is 1000
| DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL | Default embedding model for MCP semantic tool filtering. Default is "text-embedding-3-small"
| DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD | Default similarity threshold for MCP semantic tool filtering. Default is 0.3
| DEFAULT_MCP_SEMANTIC_FILTER_TOP_K | Default number of top results to return for MCP semantic tool filtering. Default is 10
| DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT | Default token count for mock response completions. Default is 20
| DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT | Default token count for mock response prompts. Default is 10
| DEFAULT_MODEL_CREATED_AT_TIME | Default creation timestamp for models. Default is 1677610602
@ -802,6 +805,7 @@ router_settings:
| MAXIMUM_TRACEBACK_LINES_TO_LOG | Maximum number of lines to log in traceback in LiteLLM Logs UI. Default is 100
| MAX_RETRY_DELAY | Maximum delay in seconds for retrying requests. Default is 8.0
| MAX_LANGFUSE_INITIALIZED_CLIENTS | Maximum number of Langfuse clients to initialize on proxy. Default is 50. This is set since langfuse initializes 1 thread everytime a client is initialized. We've had an incident in the past where we reached 100% cpu utilization because Langfuse was initialized several times.
| MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH | Maximum header length for MCP semantic filter tools. Default is 150
| MIN_NON_ZERO_TEMPERATURE | Minimum non-zero temperature value. Default is 0.0001
| MINIMUM_PROMPT_CACHE_TOKEN_COUNT | Minimum token count for caching a prompt. Default is 1024
| MISTRAL_API_BASE | Base URL for Mistral API. Default is https://api.mistral.ai

197
docs/my-website/docs/proxy/guardrails/grayswan.md
View file

@ -13,20 +13,26 @@ Cygnal returns a `violation` score between `0` and `1` (higher means more likely
### 1. Obtain Credentials
1. Create a Gray Swan account and generate a Cygnal API key.
1. Log in to our Gray Swan platform and generate a Cygnal API key.
For existing customers, you should already have access to our [platform](https://platform.grayswan.ai).
For new users, please register at this [page](https://hubs.ly/Q03-sX1J0) and we are more than happy to give you an onboarding!
2. Configure environment variables for the LiteLLM proxy host:
```bash
export GRAYSWAN_API_KEY="your-grayswan-key"
export GRAYSWAN_API_BASE="https://api.grayswan.ai"
```
```bash
export GRAYSWAN_API_KEY="your-grayswan-key"
export GRAYSWAN_API_BASE="https://api.grayswan.ai"
```
### 2. Configure `config.yaml`
Add a guardrail entry that references the Gray Swan integration. Below is a balanced example that monitors both input and output but only blocks once the violation score reaches the configured threshold.
Add a guardrail entry that references the Gray Swan integration. Below is our recommmended settings.
```yaml
model_list:
model_list: # this part is a standard litellm configuration for reference
- model_name: openai/gpt-4.1-mini
litellm_params:
model: openai/gpt-4.1-mini
@ -40,13 +46,14 @@ guardrails:
api_key: os.environ/GRAYSWAN_API_KEY
api_base: os.environ/GRAYSWAN_API_BASE # optional
optional_params:
on_flagged_action: monitor # or "block"
on_flagged_action: passthrough # or "block" or "monitor"
violation_threshold: 0.5 # score >= threshold is flagged
reasoning_mode: hybrid # off | hybrid | thinking
categories:
safety: "Detect jailbreaks and policy violations"
policy_id: "your-cygnal-policy-id"
policy_id: "your-cygnal-policy-id" # Optional: Your Cygnal policy ID. Defaults to a content safety policy if empty.
streaming_end_of_stream_only: true # For streaming API, only send the assembled message to Cygnal (post_call only). Defaults to false.
default_on: true
guardrail_timeout: 30 # Defaults to 30 seconds. Change accordingly.
fail_open: true # Defaults to true; set to false to propagate guardrail errors.
general_settings:
master_key: "your-litellm-master-key"
@ -65,13 +72,13 @@ litellm --config config.yaml --port 4000
## Choosing Guardrail Modes
Gray Swan can run during `pre_call`, `during_call`, and `post_call` stages. Combine modes based on your latency and coverage requirements.
Gray Swan can run during `pre_call`, `during_call`, and `post_call` stages. Combine modes based on your latency and coverage requirements.
| Mode | When it Runs | Protects | Typical Use Case |
|--------------|-------------------|-----------------------|------------------|
| `pre_call` | Before LLM call | User input only | Block prompt injection before it reaches the model |
| `during_call`| Parallel to call | User input only | Low-latency monitoring without blocking |
| `post_call` | After response | Full conversation | Scan output for policy violations, leaked secrets, or IPI |
| `post_call` | After response | Model Outputs | Scan output for policy violations, leaked secrets, or IPI |
When using `during_call` with `on_flagged_action: block` or `on_flagged_action: passthrough`:
@ -81,87 +88,110 @@ When using `during_call` with `on_flagged_action: block` or `on_flagged_action:
- The guardrail exception prevents the response from reaching the user, but **does not cancel the running LLM task**
- This means you pay full LLM costs while returning an error/passthrough message to the user
**Recommendation:** For cost-sensitive applications, use `pre_call` and `post_call` instead of `during_call` for blocking or passthrough modes. Reserve `during_call` for `monitor` mode where you want low-latency logging without impacting the user experience.
**Recommendation:** Use `pre_call` and `post_call` instead of `during_call` for `passthrough` (or `block`) `on_flagged_action` (see our recommended configuration above). Reserve `during_call` for `monitor` mode ONLY when you want low-latency logging without impacting the user experience.
<Tabs>
<TabItem value="monitor" label="Monitor Only">
---
```yaml
guardrails:
- guardrail_name: "cygnal-monitor-only"
litellm_params:
guardrail: grayswan
mode: "during_call"
api_key: os.environ/GRAYSWAN_API_KEY
optional_params:
on_flagged_action: monitor
violation_threshold: 0.6
default_on: true
## Work with Claude Code
Follow the official litellm [guide](https://docs.litellm.ai/docs/tutorials/claude_responses_api) on setting up Claude Code with litellm, with the guardrail part mentioned above added to your litellm configuration. Cygnal natively supports coding agent policies defense. Define your own policy or use the provided coding policies on the platform. The example config we show above is also the recommended setup for Claude Code (with the `policy_id` replaced with an appropriate one).
---
## Per-request overrides via `extra_body`
You can override parts of the Gray Swan guardrail configuration on a per-request basis by passing `litellm_metadata.guardrails[*].grayswan.extra_body`.
`extra_body` is merged into the Cygnal request body and takes precedence over specific fields from `config.yaml`, which are `policy_id`, `violation_threshold`, and `reasoning_mode`.
If you include a `metadata` field inside `extra_body`, it is forwarded to the Cygnal API as-is under the request body's `metadata` field.
Example:
```bash
curl -X POST "http://0.0.0.0:4000/v1/messages?beta=true" \
-H "Authorization: Bearer token" \
-H "Content-Type: application/json" \
-d '{
"model": "openrouter/anthropic/claude-sonnet-4.5",
"messages": [{"role": "user", "content": "hello"}],
"litellm_metadata": {
"guardrails": [
{
"cygnal-monitor": {
"extra_body": {
"policy_id": "specific policy id you want to use",
"metadata": {
"user": "health-check"
}
}
}
}
]
}
}'
```
Best for visibility without blocking. Alerts are logged via LiteLLM’s standard logging callbacks.
OpenAI client:
</TabItem>
<TabItem value="block-input" label="Block Input">
```python
from openai import OpenAI
```yaml
guardrails:
- guardrail_name: "cygnal-block-input"
litellm_params:
guardrail: grayswan
mode: "pre_call"
api_key: os.environ/GRAYSWAN_API_KEY
optional_params:
on_flagged_action: block
violation_threshold: 0.4
categories:
pii: "Detect sensitive data"
default_on: true
client = OpenAI(api_key="anything", base_url="http://0.0.0.0:4000")
resp = client.responses.create(
model="openrouter/anthropic/claude-sonnet-4.5",
input="hello",
extra_body={
"litellm_metadata": {
"guardrails": [
{
"cygnal-monitor": {
"extra_body": {
"policy_id": "69038214e5cdb6befc5e991e",
"metadata": {"trace_id": "trace-123"},
}
}
}
]
}
},
)
```
Stops malicious or sensitive prompts before any tokens are generated.
Anthropic client:
</TabItem>
<TabItem value="full-coverage" label="Full Coverage">
```python
from anthropic import Anthropic
```yaml
guardrails:
- guardrail_name: "cygnal-full-coverage"
litellm_params:
guardrail: grayswan
mode: [pre_call, post_call]
api_key: os.environ/GRAYSWAN_API_KEY
optional_params:
on_flagged_action: block
violation_threshold: 0.5
reasoning_mode: thinking
policy_id: "policy-id-from-grayswan"
default_on: true
client = Anthropic(api_key="anything", base_url="http://0.0.0.0:4000")
resp = client.messages.create(
model="openrouter/anthropic/claude-sonnet-4.5",
max_tokens=256,
messages=[{"role": "user", "content": "hello"}],
extra_body={
"litellm_metadata": {
"guardrails": [
{
"cygnal-monitor": {
"extra_body": {
"policy_id": "69038214e5cdb6befc5e991e",
"metadata": {"trace_id": "trace-123"},
}
}
}
]
}
},
)
```
Provides the strongest enforcement by inspecting both prompts and responses.
Notes:
</TabItem>
<TabItem value="passthrough" label="Passthrough Mode">
```yaml
guardrails:
- guardrail_name: "cygnal-passthrough"
litellm_params:
guardrail: grayswan
mode: [pre_call, post_call]
api_key: os.environ/GRAYSWAN_API_KEY
optional_params:
on_flagged_action: passthrough
violation_threshold: 0.5
default_on: true
```
Allows requests to proceed without raising a 400 error when content is flagged. Instead of blocking, the model response content is replaced with a detailed violation message including violation score, violated rules, and detection flags (mutation, IPI). **Supported Response Formats:** OpenAI chat/text completions, Anthropic Messages API. Other response types (embeddings, images, etc.) will log a warning and return unchanged.
</TabItem>
</Tabs>
- The guardrail name (for example, `cygnal-monitor`) must match the `guardrail_name` in `config.yaml`.
- Per-request guardrail overrides may require a premium license, depending on your proxy settings.
---
@ -170,9 +200,14 @@ Allows requests to proceed without raising a 400 error when content is flagged.
| Parameter | Type | Description |
|---------------------------------------|-----------------|-------------|
| `api_key` | string | Gray Swan Cygnal API key. Reads from `GRAYSWAN_API_KEY` if omitted. |
| `api_base` | string | Override for the Gray Swan API base URL. Defaults to `https://api.grayswan.ai` or `GRAYSWAN_API_BASE`. |
| `mode` | string or list | Guardrail stages (`pre_call`, `during_call`, `post_call`). |
| `optional_params.on_flagged_action` | string | `monitor` (log only), `block` (raise `HTTPException`), or `passthrough` (replace response content with violation message, no 400 error). |
| `.optional_params.violation_threshold`| number (0-1) | Scores at or above this value are considered violations. |
| `optional_params.violation_threshold` | number (0-1) | Scores at or above this value are considered violations. |
| `optional_params.reasoning_mode` | string | `off`, `hybrid`, or `thinking`. Enables Cygnal's reasoning capabilities. |
| `optional_params.categories` | object | Map of custom category names to descriptions. |
| `optional_params.policy_id` | string | Gray Swan policy identifier. |
| `guardrail_timeout` | number | Timeout in seconds for the Cygnal request. Defaults to 30. |
| `fail_open` | boolean | If true, errors contacting Cygnal are logged and the request proceeds; if false, errors propagate. Defaults to treu. |
| `streaming_end_of_stream_only` | boolean | For streaming `post_call`, only send the final assembled response to Cygnal. Defaults to false. |
| `default_on` | boolean | Run the guardrail on every request by default. |

34
docs/my-website/docs/proxy/ui_logs.md
View file

@ -37,6 +37,40 @@ general_settings:
<Image img={require('../../img/ui_request_logs_content.png')}/>
## Tracing Tools
View which tools were provided and called in your completion requests.
<Image img={require('../../img/ui_tools.png')}/>
**Example:** Make a completion request with tools:
```bash
curl -X POST 'http://localhost:4000/chat/completions' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{
"model": "gpt-4",
"messages": [{"role": "user", "content": "What is the weather?"}],
"tools": [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"}
}
}
}
}
]
}'
```
Check the Logs page to see all tools provided and which ones were called.
## Stop storing Error Logs in DB

BIN
docs/my-website/img/ui_tools.png Normal file
View file

Binary file not shown.
Side by side

After

Width:  |  Height:  |  Size: 420 KiB

3
docs/my-website/sidebars.js
View file

@ -470,6 +470,7 @@ const sidebars = {
label: "/a2a - A2A Agent Gateway",
items: [
"a2a",
"a2a_invoking_agents",
"a2a_cost_tracking",
"a2a_agent_permissions"
],
@ -539,6 +540,7 @@ const sidebars = {
items: [
"mcp",
"mcp_usage",
"mcp_semantic_filter",
"mcp_control",
"mcp_cost",
"mcp_guardrail",
@ -717,6 +719,7 @@ const sidebars = {
"providers/bedrock_agents",
"providers/bedrock_writer",
"providers/bedrock_batches",
"providers/bedrock_realtime_with_audio",
"providers/aws_polly",
"providers/bedrock_vector_store",
]

BIN
enterprise/dist/litellm_enterprise-0.1.29-py3-none-any.whl vendored Normal file
View file

Binary file not shown.

BIN
enterprise/dist/litellm_enterprise-0.1.29.tar.gz vendored Normal file
View file

Binary file not shown.

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm-enterprise"
version = "0.1.28"
version = "0.1.29"
description = "Package for LiteLLM Enterprise features"
authors = ["BerriAI"]
readme = "README.md"
@ -22,7 +22,7 @@ requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "0.1.28"
version = "0.1.29"
version_files = [
"pyproject.toml:version",
"../requirements.txt:litellm-enterprise==",

8
litellm-proxy-extras/litellm_proxy_extras/migrations/20260129103648_add_verificationtoken_indexes/migration.sql
View file

@ -1,8 +0,0 @@
-- CreateIndex
CREATE INDEX "LiteLLM_VerificationToken_user_id_team_id_idx" ON "LiteLLM_VerificationToken"("user_id", "team_id");
-- CreateIndex
CREATE INDEX "LiteLLM_VerificationToken_team_id_idx" ON "LiteLLM_VerificationToken"("team_id");
-- CreateIndex
CREATE INDEX "LiteLLM_VerificationToken_budget_reset_at_expires_idx" ON "LiteLLM_VerificationToken"("budget_reset_at", "expires");

8
litellm-proxy-extras/litellm_proxy_extras/migrations/20260202204523_add_allow_team_guardrail_config/migration.sql Normal file
View file

@ -0,0 +1,8 @@
-- AlterTable
ALTER TABLE "LiteLLM_GuardrailsTable" ADD COLUMN "team_id" TEXT;
-- AlterTable
ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "allow_team_guardrail_config" BOOLEAN NOT NULL DEFAULT false;
-- AlterTable
ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN "allow_team_guardrail_config" BOOLEAN NOT NULL DEFAULT false;

10
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -305,16 +305,6 @@ model LiteLLM_VerificationToken {
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
// SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub"
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2
@@index([user_id, team_id])
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."team_id" = $1 OFFSET $2
@@index([team_id])
// 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])
}
// Audit table for deleted keys - preserves spend and key information for historical tracking

3
litellm/__init__.py
View file

@ -261,6 +261,8 @@ extra_spend_tag_headers: Optional[List[str]] = None
in_memory_llm_clients_cache: "LLMClientCache"
safe_memory_mode: bool = False
enable_azure_ad_token_refresh: Optional[bool] = False
# Proxy Authentication - auto-obtain/refresh OAuth2/JWT tokens for LiteLLM Proxy
proxy_auth: Optional[Any] = None
### DEFAULT AZURE API VERSION ###
AZURE_DEFAULT_API_VERSION = "2025-02-01-preview" # this is updated to the latest
### DEFAULT WATSONX API VERSION ###
@ -1378,6 +1380,7 @@ if TYPE_CHECKING:
from .llms.topaz.image_variations.transformation import TopazImageVariationConfig as TopazImageVariationConfig
from litellm.llms.openai.completion.transformation import OpenAITextCompletionConfig as OpenAITextCompletionConfig
from .llms.groq.chat.transformation import GroqChatConfig as GroqChatConfig
from .llms.a2a.chat.transformation import A2AConfig as A2AConfig
from .llms.voyage.embedding.transformation import VoyageEmbeddingConfig as VoyageEmbeddingConfig
from .llms.voyage.embedding.transformation_contextual import VoyageContextualEmbeddingConfig as VoyageContextualEmbeddingConfig
from .llms.infinity.embedding.transformation import InfinityEmbeddingConfig as InfinityEmbeddingConfig

2
litellm/_lazy_imports_registry.py
View file

@ -213,6 +213,7 @@ LLM_CONFIG_NAMES = (
"TopazImageVariationConfig",
"OpenAITextCompletionConfig",
"GroqChatConfig",
"A2AConfig",
"GenAIHubOrchestrationConfig",
"VoyageEmbeddingConfig",
"VoyageContextualEmbeddingConfig",
@ -850,6 +851,7 @@ _LLM_CONFIGS_IMPORT_MAP = {
"OpenAITextCompletionConfig",
),
"GroqChatConfig": (".llms.groq.chat.transformation", "GroqChatConfig"),
"A2AConfig": (".llms.a2a.chat.transformation", "A2AConfig"),
"GenAIHubOrchestrationConfig": (
".llms.sap.chat.transformation",
"GenAIHubOrchestrationConfig",

3
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -329,6 +329,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
else:
request_data[key] = value
if headers:
request_data["extra_headers"] = headers
return request_data
@staticmethod

14
litellm/constants.py
View file

@ -67,6 +67,20 @@ DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET = int(
os.getenv("DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET", 0)
)
# MCP Semantic Tool Filter Defaults
DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL = str(
os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL", "text-embedding-3-small")
)
DEFAULT_MCP_SEMANTIC_FILTER_TOP_K = int(
os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_TOP_K", 10)
)
DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD = float(
os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD", 0.3)
)
MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH = int(
os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150)
)
# Gemini model-specific minimal thinking budget constants
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH = int(
os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH", 1)

9
litellm/integrations/custom_guardrail.py
View file

@ -475,11 +475,18 @@ class CustomGuardrail(CustomLogger):
guardrail_config: DynamicGuardrailParams = DynamicGuardrailParams(
**guardrail[self.guardrail_name]
)
extra_body = guardrail_config.get("extra_body", {})
if self._validate_premium_user() is not True:
if isinstance(extra_body, dict) and extra_body:
verbose_logger.warning(
"Guardrail %s: ignoring dynamic extra_body keys %s because premium_user is False",
self.guardrail_name,
list(extra_body.keys()),
)
return {}
# Return the extra_body if it exists, otherwise empty dict
return guardrail_config.get("extra_body", {})
return extra_body
return {}

199
litellm/integrations/prometheus.py
View file

@ -1478,35 +1478,19 @@ class PrometheusLogger(CustomLogger):
"""
Determine if a request has an invalid API key based on status code and exception.
This method prevents invalid authentication attempts from being recorded in
Prometheus metrics. A 401 status code is the definitive indicator of authentication
failure. Additionally, we check exception messages for authentication error patterns
to catch cases where the exception hasn't been converted to a ProxyException yet.
Returns True only when we truly cannot record useful metrics (e.g. missing required
data). We no longer skip 401/invalid-key requests - all requests including
authentication failures and bad requests must be tracked for debugging, security
auditing, abuse detection, and capacity planning.
Args:
status_code: HTTP status code (401 indicates authentication error)
exception: Exception object to check for auth-related error messages
status_code: HTTP status code
exception: Exception object (unused, kept for API compatibility)
Returns:
True if the request has an invalid API key and metrics should be skipped,
True if metrics should be skipped (currently always False - track all requests),
False otherwise
"""
if status_code == 401:
return True
# Handle cases where AssertionError is raised before conversion to ProxyException
if exception is not None:
exception_str = str(exception).lower()
auth_error_patterns = [
"virtual key expected",
"expected to start with 'sk-'",
"authentication error",
"invalid api key",
"api key not valid",
]
if any(pattern in exception_str for pattern in auth_error_patterns):
return True
return False
def _should_skip_metrics_for_invalid_key(
@ -1683,6 +1667,108 @@ class PrometheusLogger(CustomLogger):
)
pass
def _safe_get(self, obj: Any, key: str, default: Any = None) -> Any:
"""Get value from dict or Pydantic model."""
if obj is None:
return default
if isinstance(obj, dict):
return obj.get(key, default)
return getattr(obj, key, default)
def _extract_deployment_failure_label_values(
self, request_kwargs: dict
) -> Dict[str, Optional[str]]:
"""
Extract label values for deployment failure metrics from all available
sources in request_kwargs. Falls back to litellm_params metadata and
user_api_key_auth when standard_logging_payload has None values.
"""
standard_logging_payload = (
request_kwargs.get("standard_logging_object", {}) or {}
)
_litellm_params = request_kwargs.get("litellm_params", {}) or {}
_metadata_raw = self._safe_get(standard_logging_payload, "metadata") or {}
if isinstance(_metadata_raw, dict):
_metadata = _metadata_raw
else:
_metadata = {
"user_api_key_alias": getattr(
_metadata_raw, "user_api_key_alias", None
),
"user_api_key_team_id": getattr(
_metadata_raw, "user_api_key_team_id", None
),
"user_api_key_team_alias": getattr(
_metadata_raw, "user_api_key_team_alias", None
),
"user_api_key_hash": getattr(_metadata_raw, "user_api_key_hash", None),
"requester_ip_address": getattr(
_metadata_raw, "requester_ip_address", None
),
"user_agent": getattr(_metadata_raw, "user_agent", None),
}
_litellm_params_metadata = _litellm_params.get("metadata", {}) or {}
# Extract user_api_key_auth if present (proxy injects this, skipped in merge)
user_api_key_auth = _litellm_params_metadata.get("user_api_key_auth")
def _get_api_key_alias() -> Optional[str]:
val = _metadata.get("user_api_key_alias")
if val is not None:
return val
val = _litellm_params_metadata.get("user_api_key_alias")
if val is not None:
return val
if user_api_key_auth is not None:
return getattr(user_api_key_auth, "key_alias", None)
return None
def _get_team_id() -> Optional[str]:
val = _metadata.get("user_api_key_team_id")
if val is not None:
return val
val = _litellm_params_metadata.get("user_api_key_team_id")
if val is not None:
return val
if user_api_key_auth is not None:
return getattr(user_api_key_auth, "team_id", None)
return None
def _get_team_alias() -> Optional[str]:
val = _metadata.get("user_api_key_team_alias")
if val is not None:
return val
val = _litellm_params_metadata.get("user_api_key_team_alias")
if val is not None:
return val
if user_api_key_auth is not None:
return getattr(user_api_key_auth, "team_alias", None)
return None
def _get_hashed_api_key() -> Optional[str]:
val = _metadata.get("user_api_key_hash")
if val is not None:
return val
val = _litellm_params_metadata.get("user_api_key_hash")
if val is not None:
return val
if user_api_key_auth is not None:
return getattr(user_api_key_auth, "api_key", None) or getattr(
user_api_key_auth, "api_key_hash", None
)
return None
return {
"api_key_alias": _get_api_key_alias(),
"team": _get_team_id(),
"team_alias": _get_team_alias(),
"hashed_api_key": _get_hashed_api_key(),
"client_ip": _metadata.get("requester_ip_address")
or _litellm_params_metadata.get("requester_ip_address"),
"user_agent": _metadata.get("user_agent")
or _litellm_params_metadata.get("user_agent"),
}
def set_llm_deployment_failure_metrics(self, request_kwargs: dict):
"""
Sets Failure metrics when an LLM API call fails
@ -1707,6 +1793,21 @@ class PrometheusLogger(CustomLogger):
model_id = standard_logging_payload.get("model_id", None)
exception = request_kwargs.get("exception", None)
# Fallback: model_id from litellm_metadata.model_info
if model_id is None:
_model_info = (
(_litellm_params.get("litellm_metadata") or {}).get("model_info")
or (_litellm_params.get("metadata") or {}).get("model_info")
or {}
)
model_id = _model_info.get("id")
# Fallback: model_group from litellm_metadata
if model_group is None:
model_group = (_litellm_params.get("litellm_metadata") or {}).get(
"model_group"
) or (_litellm_params.get("metadata") or {}).get("model_group")
llm_provider = _litellm_params.get("custom_llm_provider", None)
if self._should_skip_metrics_for_invalid_key(
@ -1714,9 +1815,37 @@ class PrometheusLogger(CustomLogger):
standard_logging_payload=standard_logging_payload,
):
return
hashed_api_key = standard_logging_payload.get("metadata", {}).get(
# Extract context labels from all available sources (fix for None labels)
fallback_values = self._extract_deployment_failure_label_values(
request_kwargs
)
_metadata = standard_logging_payload.get("metadata", {}) or {}
hashed_api_key = fallback_values.get("hashed_api_key") or _metadata.get(
"user_api_key_hash"
)
api_key_alias = fallback_values.get("api_key_alias") or _metadata.get(
"user_api_key_alias"
)
team = fallback_values.get("team") or _metadata.get("user_api_key_team_id")
team_alias = fallback_values.get("team_alias") or _metadata.get(
"user_api_key_team_alias"
)
client_ip = fallback_values.get("client_ip") or _metadata.get(
"requester_ip_address"
)
user_agent = fallback_values.get("user_agent") or _metadata.get(
"user_agent"
)
# exception_status: prefer status_code, fallback to exception class for known types
exception_status = None
if exception is not None:
exception_status = str(getattr(exception, "status_code", None))
if exception_status == "None" or not exception_status:
code = getattr(exception, "code", None)
if code is not None:
exception_status = str(code)
# Create enum_values for the label factory (always create for use in different metrics)
enum_values = UserAPIKeyLabelValues(
@ -1724,26 +1853,18 @@ class PrometheusLogger(CustomLogger):
model_id=model_id,
api_base=api_base,
api_provider=llm_provider,
exception_status=(
str(getattr(exception, "status_code", None)) if exception else None
),
exception_status=exception_status,
exception_class=(
self._get_exception_class_name(exception) if exception else None
),
requested_model=model_group,
requested_model=model_group or litellm_model_name,
hashed_api_key=hashed_api_key,
api_key_alias=standard_logging_payload["metadata"][
"user_api_key_alias"
],
team=standard_logging_payload["metadata"]["user_api_key_team_id"],
team_alias=standard_logging_payload["metadata"][
"user_api_key_team_alias"
],
api_key_alias=api_key_alias,
team=team,
team_alias=team_alias,
tags=standard_logging_payload.get("request_tags", []),
client_ip=standard_logging_payload["metadata"].get(
"requester_ip_address"
),
user_agent=standard_logging_payload["metadata"].get("user_agent"),
client_ip=client_ip,
user_agent=user_agent,
)
"""

6
litellm/llms/a2a/__init__.py Normal file
View file

@ -0,0 +1,6 @@
"""
A2A (Agent-to-Agent) Protocol Provider for LiteLLM
"""
from .chat.transformation import A2AConfig
__all__ = ["A2AConfig"]

6
litellm/llms/a2a/chat/__init__.py Normal file
View file

@ -0,0 +1,6 @@
"""
A2A Chat Completion Implementation
"""
from .transformation import A2AConfig
__all__ = ["A2AConfig"]

103
litellm/llms/a2a/chat/streaming_iterator.py Normal file
View file

@ -0,0 +1,103 @@
"""
A2A Streaming Response Iterator
"""
from typing import Optional, Union
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
from ..common_utils import extract_text_from_a2a_response
class A2AModelResponseIterator(BaseModelResponseIterator):
"""
Iterator for parsing A2A streaming responses.
Converts A2A JSON-RPC streaming chunks to OpenAI-compatible format.
"""
def __init__(
self,
streaming_response,
sync_stream: bool,
json_mode: Optional[bool] = False,
model: str = "a2a/agent",
):
super().__init__(
streaming_response=streaming_response,
sync_stream=sync_stream,
json_mode=json_mode,
)
self.model = model
def chunk_parser(self, chunk: dict) -> Union[GenericStreamingChunk, ModelResponseStream]:
"""
Parse A2A streaming chunk to OpenAI format.
A2A chunk format:
{
"jsonrpc": "2.0",
"id": "request-id",
"result": {
"message": {
"parts": [{"kind": "text", "text": "content"}]
}
}
}
Or for tasks:
{
"jsonrpc": "2.0",
"result": {
"kind": "task",
"status": {"state": "running"},
"artifacts": [{"parts": [{"kind": "text", "text": "content"}]}]
}
}
"""
try:
# Extract text from A2A response
text = extract_text_from_a2a_response(chunk)
# Determine finish reason
finish_reason = self._get_finish_reason(chunk)
# Return generic streaming chunk
return GenericStreamingChunk(
text=text,
is_finished=bool(finish_reason),
finish_reason=finish_reason or "",
usage=None,
index=0,
tool_use=None,
)
except Exception:
# Return empty chunk on parse error
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=0,
tool_use=None,
)
def _get_finish_reason(self, chunk: dict) -> Optional[str]:
"""Extract finish reason from A2A chunk"""
result = chunk.get("result", {})
# Check for task completion
if isinstance(result, dict):
status = result.get("status", {})
if isinstance(status, dict):
state = status.get("state")
if state == "completed":
return "stop"
elif state == "failed":
return "stop" # Map failed state to 'stop' (valid finish_reason)
# Check for [DONE] marker
if chunk.get("done") is True:
return "stop"
return None

370
litellm/llms/a2a/chat/transformation.py Normal file
View file

@ -0,0 +1,370 @@
"""
A2A Protocol Transformation for LiteLLM
"""
import uuid
from typing import Any, Dict, Iterator, List, Optional, Union
import httpx
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import Choices, Message, ModelResponse
from ..common_utils import (
A2AError,
convert_messages_to_prompt,
extract_text_from_a2a_response,
)
from .streaming_iterator import A2AModelResponseIterator
class A2AConfig(BaseConfig):
"""
Configuration for A2A (Agent-to-Agent) Protocol.
Handles transformation between OpenAI and A2A JSON-RPC 2.0 formats.
"""
@staticmethod
def resolve_agent_config_from_registry(
model: str,
api_base: Optional[str],
api_key: Optional[str],
headers: Optional[Dict[str, Any]],
optional_params: Dict[str, Any],
) -> tuple[Optional[str], Optional[str], Optional[Dict[str, Any]]]:
"""
Resolve agent configuration from registry if model format is "a2a/<agent-name>".
Extracts agent name from model string and looks up configuration in the
agent registry (if available in proxy context).
Args:
model: Model string (e.g., "a2a/my-agent")
api_base: Explicit api_base (takes precedence over registry)
api_key: Explicit api_key (takes precedence over registry)
headers: Explicit headers (takes precedence over registry)
optional_params: Dict to merge additional litellm_params into
Returns:
Tuple of (api_base, api_key, headers) with registry values filled in
"""
# Extract agent name from model (e.g., "a2a/my-agent" -> "my-agent")
agent_name = model.split("/", 1)[1] if "/" in model else None
# Only lookup if agent name exists and some config is missing
if not agent_name or (api_base is not None and api_key is not None and headers is not None):
return api_base, api_key, headers
# Try registry lookup (only available in proxy context)
try:
from litellm.proxy.agent_endpoints.agent_registry import (
global_agent_registry,
)
agent = global_agent_registry.get_agent_by_name(agent_name)
if agent:
# Get api_base from agent card URL
if api_base is None and agent.agent_card_params:
api_base = agent.agent_card_params.get("url")
# Get api_key, headers, and other params from litellm_params
if agent.litellm_params:
if api_key is None:
api_key = agent.litellm_params.get("api_key")
if headers is None:
agent_headers = agent.litellm_params.get("headers")
if agent_headers:
headers = agent_headers
# Merge other litellm_params (timeout, max_retries, etc.)
for key, value in agent.litellm_params.items():
if key not in ["api_key", "api_base", "headers", "model"] and key not in optional_params:
optional_params[key] = value
except ImportError:
pass # Registry not available (not running in proxy context)
return api_base, api_key, headers
def get_supported_openai_params(self, model: str) -> List[str]:
"""Return list of supported OpenAI parameters"""
return [
"stream",
"temperature",
"max_tokens",
"top_p",
]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
"""
Map OpenAI parameters to A2A parameters.
For A2A protocol, we need to map the stream parameter so
transform_request can determine which JSON-RPC method to use.
"""
# Map stream parameter
for param, value in non_default_params.items():
if param == "stream" and value is True:
optional_params["stream"] = value
return optional_params
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
"""
Validate environment and set headers for A2A requests.
Args:
headers: Request headers dict
model: Model name
messages: Messages list
optional_params: Optional parameters
litellm_params: LiteLLM parameters
api_key: API key (optional for A2A)
api_base: API base URL
Returns:
Updated headers dict
"""
# Ensure Content-Type is set to application/json for JSON-RPC 2.0
if "content-type" not in headers and "Content-Type" not in headers:
headers["Content-Type"] = "application/json"
# Add Authorization header if API key is provided
if api_key is not None:
headers["Authorization"] = f"Bearer {api_key}"
return headers
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
"""
Get the complete A2A agent endpoint URL.
A2A agents use JSON-RPC 2.0 at the base URL, not specific paths.
The method (message/send or message/stream) is specified in the
JSON-RPC request body, not in the URL.
Args:
api_base: Base URL of the A2A agent (e.g., "http://0.0.0.0:9999")
api_key: API key (not used for URL construction)
model: Model name (not used for A2A, agent determined by api_base)
optional_params: Optional parameters
litellm_params: LiteLLM parameters
stream: Whether this is a streaming request (affects JSON-RPC method)
Returns:
Complete URL for the A2A endpoint (base URL)
"""
if api_base is None:
raise ValueError("api_base is required for A2A provider")
# A2A uses JSON-RPC 2.0 at the base URL
# Remove trailing slash for consistency
return api_base.rstrip("/")
def transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
"""
Transform OpenAI request to A2A JSON-RPC 2.0 format.
Args:
model: Model name
messages: List of OpenAI messages
optional_params: Optional parameters
litellm_params: LiteLLM parameters
headers: Request headers
Returns:
A2A JSON-RPC 2.0 request dict
"""
# Generate request ID
request_id = str(uuid.uuid4())
if not messages:
raise ValueError("At least one message is required for A2A completion")
# Convert all messages to maintain conversation history
# Use helper to format conversation with role prefixes
full_context = convert_messages_to_prompt(messages)
# Create single A2A message with full conversation context
a2a_message = {
"role": "user",
"parts": [{"kind": "text", "text": full_context}],
"messageId": str(uuid.uuid4()),
}
# Build JSON-RPC 2.0 request
# For A2A protocol, the method is "message/send" for non-streaming
# and "message/stream" for streaming
stream = optional_params.get("stream", False)
method = "message/stream" if stream else "message/send"
request_data = {
"jsonrpc": "2.0",
"id": request_id,
"method": method,
"params": {
"message": a2a_message
}
}
return request_data
def transform_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ModelResponse,
logging_obj: Any,
request_data: dict,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ModelResponse:
"""
Transform A2A JSON-RPC 2.0 response to OpenAI format.
Args:
model: Model name
raw_response: HTTP response from A2A agent
model_response: Model response object to populate
logging_obj: Logging object
request_data: Original request data
messages: Original messages
optional_params: Optional parameters
litellm_params: LiteLLM parameters
encoding: Encoding object
api_key: API key
json_mode: JSON mode flag
Returns:
Populated ModelResponse object
"""
try:
response_json = raw_response.json()
except Exception as e:
raise A2AError(
status_code=raw_response.status_code,
message=f"Failed to parse A2A response: {str(e)}",
headers=dict(raw_response.headers),
)
# Check for JSON-RPC error
if "error" in response_json:
error = response_json["error"]
raise A2AError(
status_code=raw_response.status_code,
message=f"A2A error: {error.get('message', 'Unknown error')}",
headers=dict(raw_response.headers),
)
# Extract text from A2A response
text = extract_text_from_a2a_response(response_json)
# Populate model response
model_response.choices = [
Choices(
finish_reason="stop",
index=0,
message=Message(
content=text,
role="assistant",
),
)
]
# Set model
model_response.model = model
# Set ID from response
model_response.id = response_json.get("id", str(uuid.uuid4()))
return model_response
def get_model_response_iterator(
self,
streaming_response: Union[Iterator, Any],
sync_stream: bool,
json_mode: Optional[bool] = False,
) -> BaseModelResponseIterator:
"""
Get streaming iterator for A2A responses.
Args:
streaming_response: Streaming response iterator
sync_stream: Whether this is a sync stream
json_mode: JSON mode flag
Returns:
A2A streaming iterator
"""
return A2AModelResponseIterator(
streaming_response=streaming_response,
sync_stream=sync_stream,
json_mode=json_mode,
)
def _openai_message_to_a2a_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
"""
Convert OpenAI message to A2A message format.
Args:
message: OpenAI message dict
Returns:
A2A message dict
"""
content = message.get("content", "")
role = message.get("role", "user")
return {
"role": role,
"parts": [{"kind": "text", "text": str(content)}],
"messageId": str(uuid.uuid4()),
}
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BaseLLMException:
"""Return appropriate error class for A2A errors"""
# Convert headers to dict if needed
headers_dict = dict(headers) if isinstance(headers, httpx.Headers) else headers
return A2AError(
status_code=status_code,
message=error_message,
headers=headers_dict,
)

152
litellm/llms/a2a/common_utils.py Normal file
View file

@ -0,0 +1,152 @@
"""
Common utilities for A2A (Agent-to-Agent) Protocol
"""
from typing import Any, Dict, List
from pydantic import BaseModel
from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_content_list_to_str,
)
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.openai import AllMessageValues
class A2AError(BaseLLMException):
"""Base exception for A2A protocol errors"""
def __init__(
self,
status_code: int,
message: str,
headers: Dict[str, Any] = {},
):
super().__init__(
status_code=status_code,
message=message,
headers=headers,
)
def convert_messages_to_prompt(messages: List[AllMessageValues]) -> str:
"""
Convert OpenAI messages to a single prompt string for A2A agent.
Formats each message as "{role}: {content}" and joins with newlines
to preserve conversation history. Handles both string and list content.
Args:
messages: List of OpenAI-format messages
Returns:
Formatted prompt string with full conversation context
"""
conversation_parts = []
for msg in messages:
# Use LiteLLM's helper to extract text from content (handles both str and list)
content_text = convert_content_list_to_str(message=msg)
# Get role
if isinstance(msg, BaseModel):
role = msg.model_dump().get("role", "user")
elif isinstance(msg, dict):
role = msg.get("role", "user")
else:
role = dict(msg).get("role", "user") # type: ignore
if content_text:
conversation_parts.append(f"{role}: {content_text}")
return "\n".join(conversation_parts)
def extract_text_from_a2a_message(
message: Dict[str, Any], depth: int = 0, max_depth: int = 10
) -> str:
"""
Extract text content from A2A message parts.
Args:
message: A2A message dict with 'parts' containing text parts
depth: Current recursion depth (internal use)
max_depth: Maximum recursion depth to prevent infinite loops
Returns:
Concatenated text from all text parts
"""
if message is None or depth >= max_depth:
return ""
parts = message.get("parts", [])
text_parts: List[str] = []
for part in parts:
if part.get("kind") == "text":
text_parts.append(part.get("text", ""))
# Handle nested parts if they exist
elif "parts" in part:
nested_text = extract_text_from_a2a_message(part, depth + 1, max_depth)
if nested_text:
text_parts.append(nested_text)
return " ".join(text_parts)
def extract_text_from_a2a_response(
response_dict: Dict[str, Any], max_depth: int = 10
) -> str:
"""
Extract text content from A2A response result.
Args:
response_dict: A2A response dict with 'result' containing message
max_depth: Maximum recursion depth to prevent infinite loops
Returns:
Text from response message parts
"""
result = response_dict.get("result", {})
if not isinstance(result, dict):
return ""
# A2A response can have different formats:
# 1. Direct message: {"result": {"kind": "message", "parts": [...]}}
# 2. Nested message: {"result": {"message": {"parts": [...]}}}
# 3. Task with artifacts: {"result": {"kind": "task", "artifacts": [{"parts": [...]}]}}
# 4. Task with status message: {"result": {"kind": "task", "status": {"message": {"parts": [...]}}}}
# 5. Streaming artifact-update: {"result": {"kind": "artifact-update", "artifact": {"parts": [...]}}}
# Check if result itself has parts (direct message)
if "parts" in result:
return extract_text_from_a2a_message(result, depth=0, max_depth=max_depth)
# Check for nested message
message = result.get("message")
if message:
return extract_text_from_a2a_message(message, depth=0, max_depth=max_depth)
# Check for streaming artifact-update (singular artifact)
artifact = result.get("artifact")
if artifact and isinstance(artifact, dict):
return extract_text_from_a2a_message(
artifact, depth=0, max_depth=max_depth
)
# Check for task status message (common in Gemini A2A agents)
status = result.get("status", {})
if isinstance(status, dict):
status_message = status.get("message")
if status_message:
return extract_text_from_a2a_message(
status_message, depth=0, max_depth=max_depth
)
# Handle task result with artifacts (plural, array)
artifacts = result.get("artifacts", [])
if artifacts and len(artifacts) > 0:
first_artifact = artifacts[0]
return extract_text_from_a2a_message(
first_artifact, depth=0, max_depth=max_depth
)
return ""

72
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -34,6 +34,7 @@ from litellm.types.llms.openai import (
)
from litellm.types.utils import (
ChatCompletionMessageToolCall,
Choices,
GenericGuardrailAPIInputs,
ModelResponse,
)
@ -76,7 +77,8 @@ class AnthropicMessagesHandler(BaseTranslation):
chat_completion_compatible_request, tool_name_mapping = (
LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
anthropic_message_request=cast(AnthropicMessagesRequest, data)
# Use a shallow copy to avoid mutating request data (pop on litellm_metadata).
anthropic_message_request=cast(AnthropicMessagesRequest, data.copy())
)
)
@ -84,9 +86,9 @@ class AnthropicMessagesHandler(BaseTranslation):
texts_to_check: List[str] = []
images_to_check: List[str] = []
tools_to_check: List[ChatCompletionToolParam] = (
chat_completion_compatible_request.get("tools", [])
)
tools_to_check: List[
ChatCompletionToolParam
] = chat_completion_compatible_request.get("tools", [])
task_mappings: List[Tuple[int, Optional[int]]] = []
# Track (message_index, content_index) for each text
# content_index is None for string content, int for list content
@ -282,7 +284,10 @@ class AnthropicMessagesHandler(BaseTranslation):
if hasattr(content_block, "model_dump"):
block_dict = content_block.model_dump()
else:
block_dict = {"type": block_type, "text": getattr(content_block, "text", None)}
block_dict = {
"type": block_type,
"text": getattr(content_block, "text", None),
}
else:
continue
@ -358,30 +363,40 @@ class AnthropicMessagesHandler(BaseTranslation):
"""
has_ended = self._check_streaming_has_ended(responses_so_far)
if has_ended:
# build the model response from the responses_so_far
model_response = cast(
ModelResponse,
AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
all_chunks=responses_so_far,
litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj),
model="",
),
built_response = AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
all_chunks=responses_so_far,
litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj),
model="",
)
tool_calls_list = cast(Optional[List[ChatCompletionMessageToolCall]], model_response.choices[0].message.tool_calls) # type: ignore
string_so_far = model_response.choices[0].message.content # type: ignore
guardrail_inputs = GenericGuardrailAPIInputs()
if string_so_far:
guardrail_inputs["texts"] = [string_so_far]
if tool_calls_list:
guardrail_inputs["tool_calls"] = tool_calls_list
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid
inputs=guardrail_inputs,
request_data={},
input_type="response",
logging_obj=litellm_logging_obj,
)
# Check if model_response is valid and has choices before accessing
if (
built_response is not None
and hasattr(built_response, "choices")
and built_response.choices
):
model_response = cast(ModelResponse, built_response)
first_choice = cast(Choices, model_response.choices[0])
tool_calls_list = cast(
Optional[List[ChatCompletionMessageToolCall]],
first_choice.message.tool_calls,
)
string_so_far = first_choice.message.content
guardrail_inputs = GenericGuardrailAPIInputs()
if string_so_far:
guardrail_inputs["texts"] = [string_so_far]
if tool_calls_list:
guardrail_inputs["tool_calls"] = tool_calls_list
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid
inputs=guardrail_inputs,
request_data={},
input_type="response",
logging_obj=litellm_logging_obj,
)
else:
verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices")
return responses_so_far
string_so_far = self.get_streaming_string_so_far(responses_so_far)
@ -648,7 +663,10 @@ class AnthropicMessagesHandler(BaseTranslation):
if isinstance(content_block, dict):
if content_block.get("type") == "text":
cast(Dict[str, Any], content_block)["text"] = guardrail_response
elif hasattr(content_block, "type") and getattr(content_block, "type", None) == "text":
elif (
hasattr(content_block, "type")
and getattr(content_block, "type", None) == "text"
):
# Update Pydantic object's text attribute
if hasattr(content_block, "text"):
content_block.text = guardrail_response

4
litellm/llms/fireworks_ai/chat/transformation.py
View file

@ -236,6 +236,10 @@ class FireworksAIConfig(OpenAIGPTConfig):
disable_add_transform_inline_image_block=disable_add_transform_inline_image_block,
)
filter_value_from_dict(cast(dict, message), "cache_control")
# Remove fields not permitted by FireworksAI that may cause:
# "Not permitted, field: 'messages[n].provider_specific_fields'"
if isinstance(message, dict) and "provider_specific_fields" in message:
message.pop("provider_specific_fields", None)
return messages

7
litellm/llms/gemini/files/transformation.py
View file

@ -210,7 +210,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
We expect file_id to be the URI (e.g. https://generativelanguage.googleapis.com/v1beta/files/...)
as returned by the upload response.
"""
api_key = litellm_params.get("api_key")
api_key = litellm_params.get("api_key") or self.get_api_key()
if not api_key:
raise ValueError("api_key is required")
@ -222,7 +222,8 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
api_base = api_base.rstrip("/")
url = "{}/v1beta/{}?key={}".format(api_base, file_id, api_key)
return url, {"Content-Type": "application/json"}
# Return empty params dict - API key is already in URL, no query params needed
return url, {}
def transform_retrieve_file_response(
self,
@ -299,7 +300,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
# Extract the file path from full URI
file_name = file_id.split("/v1beta/")[-1]
else:
file_name = file_id
file_name = file_id if file_id.startswith("files/") else f"files/{file_id}"
# Construct the delete URL
url = f"{api_base}/v1beta/{file_name}"

21
litellm/llms/github_copilot/chat/transformation.py
View file

@ -5,7 +5,7 @@ from litellm.llms.openai.openai import OpenAIConfig
from litellm.types.llms.openai import AllMessageValues
from ..authenticator import Authenticator
from ..common_utils import GetAPIKeyError, GITHUB_COPILOT_API_BASE
from ..common_utils import GetAPIKeyError, GITHUB_COPILOT_API_BASE, get_copilot_default_headers
class GithubCopilotConfig(OpenAIConfig):
@ -43,15 +43,8 @@ class GithubCopilotConfig(OpenAIConfig):
messages,
model: str,
):
import litellm
disable_copilot_system_to_assistant = (
litellm.disable_copilot_system_to_assistant
)
if not disable_copilot_system_to_assistant:
for message in messages:
if "role" in message and message["role"] == "system":
cast(Any, message)["role"] = "assistant"
# GitHub Copilot API now supports system prompts for all models (Claude, GPT, etc.)
# No conversion needed - just return messages as-is
return messages
def validate_environment(
@ -69,6 +62,14 @@ class GithubCopilotConfig(OpenAIConfig):
headers, model, messages, optional_params, litellm_params, api_key, api_base
)
# Add Copilot-specific headers (editor-version, user-agent, etc.)
try:
copilot_api_key = self.authenticator.get_api_key()
copilot_headers = get_copilot_default_headers(copilot_api_key)
validated_headers = {**copilot_headers, **validated_headers}
except GetAPIKeyError:
pass # Will be handled later in the request flow
# Add X-Initiator header based on message roles
initiator = self._determine_initiator(messages)
validated_headers["X-Initiator"] = initiator

21
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -21,7 +21,13 @@ from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.main import stream_chunk_builder
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.utils import Choices, GenericGuardrailAPIInputs, ModelResponse, ModelResponseStream, StreamingChoices
from litellm.types.utils import (
Choices,
GenericGuardrailAPIInputs,
ModelResponse,
ModelResponseStream,
StreamingChoices,
)
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
@ -80,9 +86,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check # type: ignore
if messages:
inputs["structured_messages"] = (
messages # pass the openai /chat/completions messages to the guardrail, as-is
)
inputs[
"structured_messages"
] = messages # pass the openai /chat/completions messages to the guardrail, as-is
# Pass tools (function definitions) to the guardrail
tools = data.get("tools")
if tools:
@ -362,14 +368,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
# check if the stream has ended
has_stream_ended = False
for chunk in responses_so_far:
if chunk.choices[0].finish_reason is not None:
if chunk.choices and chunk.choices[0].finish_reason is not None:
has_stream_ended = True
break
if has_stream_ended:
# convert to model response
model_response = cast(
ModelResponse, stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj)
ModelResponse,
stream_chunk_builder(
chunks=responses_so_far, logging_obj=litellm_logging_obj
),
)
# run process_output_response
await self.process_output_response(

74
litellm/llms/openai/common_utils.py
View file

@ -22,7 +22,6 @@ from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_ssl_configuration,
)
from litellm.types.utils import LlmProviders
class OpenAIError(BaseLLMException):
@ -205,67 +204,30 @@ class BaseOpenAILLM:
if litellm.aclient_session is not None:
return litellm.aclient_session
# Use the global cached client system to prevent memory leaks (issue #14540)
# This routes through get_async_httpx_client() which provides TTL-based caching
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
# Get unified SSL configuration
ssl_config = get_ssl_configuration()
try:
# Get SSL config and include in params for proper cache key
ssl_config = get_ssl_configuration()
params = {"ssl_verify": ssl_config} if ssl_config is not None else {}
params["disable_aiohttp_transport"] = litellm.disable_aiohttp_transport
# Get a cached AsyncHTTPHandler which manages the httpx.AsyncClient
cached_handler = get_async_httpx_client(
llm_provider=LlmProviders.OPENAI, # Cache key includes provider
params=params, # Include SSL config in cache key
return httpx.AsyncClient(
verify=ssl_config,
transport=AsyncHTTPHandler._create_async_transport(
ssl_context=ssl_config
if isinstance(ssl_config, ssl.SSLContext)
else None,
ssl_verify=ssl_config if isinstance(ssl_config, bool) else None,
shared_session=shared_session,
)
# Return the underlying httpx client from the handler
return cached_handler.client
except (ImportError, AttributeError, KeyError) as e:
# Fallback to creating a client directly if caching system unavailable
# This preserves backwards compatibility
verbose_logger.debug(
f"Client caching unavailable ({type(e).__name__}), using direct client creation"
)
ssl_config = get_ssl_configuration()
return httpx.AsyncClient(
verify=ssl_config,
transport=AsyncHTTPHandler._create_async_transport(
ssl_context=ssl_config
if isinstance(ssl_config, ssl.SSLContext)
else None,
ssl_verify=ssl_config if isinstance(ssl_config, bool) else None,
shared_session=shared_session,
),
follow_redirects=True,
)
),
follow_redirects=True,
)
@staticmethod
def _get_sync_http_client() -> Optional[httpx.Client]:
if litellm.client_session is not None:
return litellm.client_session
# Use the global cached client system to prevent memory leaks (issue #14540)
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
# Get unified SSL configuration
ssl_config = get_ssl_configuration()
try:
# Get SSL config and include in params for proper cache key
ssl_config = get_ssl_configuration()
params = {"ssl_verify": ssl_config} if ssl_config is not None else None
# Get a cached HTTPHandler which manages the httpx.Client
cached_handler = _get_httpx_client(params=params)
# Return the underlying httpx client from the handler
return cached_handler.client
except (ImportError, AttributeError, KeyError) as e:
# Fallback to creating a client directly if caching system unavailable
verbose_logger.debug(
f"Client caching unavailable ({type(e).__name__}), using direct client creation"
)
ssl_config = get_ssl_configuration()
return httpx.Client(
verify=ssl_config,
follow_redirects=True,
)
return httpx.Client(
verify=ssl_config,
follow_redirects=True,
)

51
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -319,9 +319,7 @@ class OpenAIResponsesHandler(BaseTranslation):
return response
if not response_output:
verbose_proxy_logger.debug(
"OpenAI Responses API: Empty output in response"
)
verbose_proxy_logger.debug("OpenAI Responses API: Empty output in response")
return response
# Step 1: Extract all text content and tool calls from response output
@ -427,27 +425,30 @@ class OpenAIResponsesHandler(BaseTranslation):
handle_raw_dict_callback=None,
)
tool_calls = model_response_choices[0].message.tool_calls
text = model_response_choices[0].message.content
guardrail_inputs = GenericGuardrailAPIInputs()
if text:
guardrail_inputs["texts"] = [text]
if tool_calls:
guardrail_inputs["tool_calls"] = cast(
List[ChatCompletionToolCallChunk], tool_calls
)
# Include model information from the response if available
response_model = final_chunk.get("response", {}).get("model")
if response_model:
guardrail_inputs["model"] = response_model
if tool_calls or text:
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=guardrail_inputs,
request_data={},
input_type="response",
logging_obj=litellm_logging_obj,
)
return responses_so_far
if model_response_choices:
tool_calls = model_response_choices[0].message.tool_calls
text = model_response_choices[0].message.content
guardrail_inputs = GenericGuardrailAPIInputs()
if text:
guardrail_inputs["texts"] = [text]
if tool_calls:
guardrail_inputs["tool_calls"] = cast(
List[ChatCompletionToolCallChunk], tool_calls
)
# Include model information from the response if available
response_model = final_chunk.get("response", {}).get("model")
if response_model:
guardrail_inputs["model"] = response_model
if tool_calls or text:
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=guardrail_inputs,
request_data={},
input_type="response",
logging_obj=litellm_logging_obj,
)
return responses_so_far
else:
verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices")
# model_response_stream = OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(final_chunk)
# tool_calls = model_response_stream.choices[0].tool_calls
# convert openai response to model response
@ -513,11 +514,9 @@ class OpenAIResponsesHandler(BaseTranslation):
# Check if it's an OutputText with text
if isinstance(content_item, OutputText):
if content_item.text:
return True
elif isinstance(content_item, dict):
if content_item.get("text"):
return True
return False

7
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -478,6 +478,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
if "type" in tool and tool["type"] == "computer_use":
computer_use_config = {k: v for k, v in tool.items() if k != "type"}
tool = {VertexToolName.COMPUTER_USE.value: computer_use_config}
# Handle OpenAI-style web_search and web_search_preview tools
# Transform them to Gemini's googleSearch tool
elif "type" in tool and tool["type"] in ("web_search", "web_search_preview"):
verbose_logger.info(
f"Gemini: Transforming OpenAI-style '{tool['type']}' tool to googleSearch"
)
tool = {VertexToolName.GOOGLE_SEARCH.value: {}}
# Handle tools with 'type' field (OpenAI spec compliance) Ignore this field -> https://github.com/BerriAI/litellm/issues/14644#issuecomment-3342061838
elif "type" in tool:
tool = {k: tool[k] for k in tool if k != "type"}

78
litellm/main.py
View file

@ -1199,6 +1199,13 @@ def completion( # type: ignore # noqa: PLR0915
headers = {}
if extra_headers is not None:
headers.update(extra_headers)
# Inject proxy auth headers if configured
if litellm.proxy_auth is not None:
try:
proxy_headers = litellm.proxy_auth.get_auth_headers()
headers.update(proxy_headers)
except Exception as e:
verbose_logger.warning(f"Failed to get proxy auth headers: {e}")
num_retries = kwargs.get(
"num_retries", None
) ## alt. param for 'max_retries'. Use this to pass retries w/ instructor.
@ -2199,6 +2206,48 @@ def completion( # type: ignore # noqa: PLR0915
logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements
client=client,
)
elif custom_llm_provider == "a2a":
# A2A (Agent-to-Agent) Protocol
# Resolve agent configuration from registry if model format is "a2a/<agent-name>"
api_base, api_key, headers = litellm.A2AConfig.resolve_agent_config_from_registry(
model=model,
api_base=api_base,
api_key=api_key,
headers=headers,
optional_params=optional_params,
)
# Fall back to environment variables and defaults
api_base = api_base or litellm.api_base or get_secret_str("A2A_API_BASE")
if api_base is None:
raise Exception(
"api_base is required for A2A provider. "
"Either provide api_base parameter, set A2A_API_BASE environment variable, "
"or register the agent in the proxy with model='a2a/<agent-name>'."
)
headers = headers or litellm.headers
response = base_llm_http_handler.completion(
model=model,
stream=stream,
messages=messages,
acompletion=acompletion,
api_base=api_base,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params,
shared_session=shared_session,
custom_llm_provider=custom_llm_provider,
timeout=timeout,
headers=headers,
encoding=_get_encoding(),
api_key=api_key,
logging_obj=logging,
client=client,
provider_config=provider_config,
)
elif custom_llm_provider == "gigachat":
# GigaChat - Sber AI's LLM (Russia)
api_key = (
@ -2455,6 +2504,20 @@ def completion( # type: ignore # noqa: PLR0915
headers = headers or litellm.headers
# Add GitHub Copilot headers (same as /responses endpoint does)
if custom_llm_provider == "github_copilot":
from litellm.llms.github_copilot.common_utils import (
get_copilot_default_headers,
)
from litellm.llms.github_copilot.authenticator import Authenticator
copilot_auth = Authenticator()
copilot_api_key = copilot_auth.get_api_key()
copilot_headers = get_copilot_default_headers(copilot_api_key)
if extra_headers:
copilot_headers.update(extra_headers)
extra_headers = copilot_headers
if extra_headers is not None:
optional_params["extra_headers"] = extra_headers
@ -3113,8 +3176,8 @@ def completion( # type: ignore # noqa: PLR0915
api_key
or litellm.api_key
or litellm.openrouter_key
or get_secret("OPENROUTER_API_KEY")
or get_secret("OR_API_KEY")
or get_secret_str("OPENROUTER_API_KEY")
or get_secret_str("OR_API_KEY")
)
openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai"
@ -4555,6 +4618,13 @@ def embedding( # noqa: PLR0915
headers = {}
if extra_headers is not None:
headers.update(extra_headers)
# Inject proxy auth headers if configured
if litellm.proxy_auth is not None:
try:
proxy_headers = litellm.proxy_auth.get_auth_headers()
headers.update(proxy_headers)
except Exception as e:
verbose_logger.warning(f"Failed to get proxy auth headers: {e}")
### CUSTOM MODEL COST ###
input_cost_per_token = kwargs.get("input_cost_per_token", None)
output_cost_per_token = kwargs.get("output_cost_per_token", None)
@ -4884,8 +4954,8 @@ def embedding( # noqa: PLR0915
api_key
or litellm.api_key
or litellm.openrouter_key
or get_secret("OPENROUTER_API_KEY")
or get_secret("OR_API_KEY")
or get_secret_str("OPENROUTER_API_KEY")
or get_secret_str("OR_API_KEY")
)
openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai"

28
litellm/model_prices_and_context_window_backup.json
View file

@ -27113,6 +27113,34 @@
"supports_reasoning": true,
"supports_tool_choice": true
},
"together_ai/zai-org/GLM-4.7": {
"input_cost_per_token": 4.5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 200000,
"max_output_tokens": 200000,
"max_tokens": 200000,
"mode": "chat",
"output_cost_per_token": 2e-06,
"source": "https://www.together.ai/models/glm-4-7",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"together_ai/moonshotai/Kimi-K2.5": {
"input_cost_per_token": 5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 256000,
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"output_cost_per_token": 2.8e-06,
"source": "https://www.together.ai/models/kimi-k2-5",
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_reasoning": true
},
"together_ai/moonshotai/Kimi-K2-Instruct-0905": {
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",

3
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -387,6 +387,9 @@ class MCPRequestHandler:
user_api_key_cache,
)
verbose_logger.debug(
f"MCP team permission lookup: team_id={user_api_key_auth.team_id if user_api_key_auth else None}"
)
if not user_api_key_auth or not user_api_key_auth.team_id or not prisma_client:
return None

250
litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py Normal file
View file

@ -0,0 +1,250 @@
"""
Semantic MCP Tool Filtering using semantic-router
Filters MCP tools semantically for /chat/completions and /responses endpoints.
"""
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from litellm._logging import verbose_logger
if TYPE_CHECKING:
from semantic_router.routers import SemanticRouter
from litellm.router import Router
class SemanticMCPToolFilter:
"""Filters MCP tools using semantic similarity to reduce context window size."""
def __init__(
self,
embedding_model: str,
litellm_router_instance: "Router",
top_k: int = 10,
similarity_threshold: float = 0.3,
enabled: bool = True,
):
"""
Initialize the semantic tool filter.
Args:
embedding_model: Model to use for embeddings (e.g., "text-embedding-3-small")
litellm_router_instance: Router instance for embedding generation
top_k: Maximum number of tools to return
similarity_threshold: Minimum similarity score for filtering
enabled: Whether filtering is enabled
"""
self.enabled = enabled
self.top_k = top_k
self.similarity_threshold = similarity_threshold
self.embedding_model = embedding_model
self.router_instance = litellm_router_instance
self.tool_router: Optional["SemanticRouter"] = None
self._tool_map: Dict[str, Any] = {} # MCPTool objects or OpenAI function dicts
async def build_router_from_mcp_registry(self) -> None:
"""Build semantic router from all MCP tools in the registry (no auth checks)."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
try:
# Get all servers from registry without auth checks
registry = global_mcp_server_manager.get_registry()
if not registry:
verbose_logger.warning("MCP registry is empty")
self.tool_router = None
return
# Fetch tools from all servers in parallel
all_tools = []
for server_id, server in registry.items():
try:
tools = await global_mcp_server_manager.get_tools_for_server(server_id)
all_tools.extend(tools)
except Exception as e:
verbose_logger.warning(f"Failed to fetch tools from server {server_id}: {e}")
continue
if not all_tools:
verbose_logger.warning("No MCP tools found in registry")
self.tool_router = None
return
verbose_logger.info(f"Fetched {len(all_tools)} tools from {len(registry)} MCP servers")
self._build_router(all_tools)
except Exception as e:
verbose_logger.error(f"Failed to build router from MCP registry: {e}")
self.tool_router = None
raise
def _extract_tool_info(self, tool) -> tuple[str, str]:
"""Extract name and description from MCP tool or OpenAI function dict."""
name: str
description: str
if isinstance(tool, dict):
# OpenAI function format
name = tool.get("name", "")
description = tool.get("description", name)
else:
# MCPTool object
name = str(tool.name)
description = str(tool.description) if tool.description else str(tool.name)
return name, description
def _build_router(self, tools: List) -> None:
"""Build semantic router with tools (MCPTool objects or OpenAI function dicts)."""
from semantic_router.routers import SemanticRouter
from semantic_router.routers.base import Route
from litellm.router_strategy.auto_router.litellm_encoder import (
LiteLLMRouterEncoder,
)
if not tools:
self.tool_router = None
return
try:
# Convert tools to routes
routes = []
self._tool_map = {}
for tool in tools:
name, description = self._extract_tool_info(tool)
self._tool_map[name] = tool
routes.append(
Route(
name=name,
description=description,
utterances=[description],
score_threshold=self.similarity_threshold,
)
)
self.tool_router = SemanticRouter(
routes=routes,
encoder=LiteLLMRouterEncoder(
litellm_router_instance=self.router_instance,
model_name=self.embedding_model,
score_threshold=self.similarity_threshold,
),
auto_sync="local",
)
verbose_logger.info(
f"Built semantic router with {len(routes)} tools"
)
except Exception as e:
verbose_logger.error(f"Failed to build semantic router: {e}")
self.tool_router = None
raise
async def filter_tools(
self,
query: str,
available_tools: List[Any],
top_k: Optional[int] = None,
) -> List[Any]:
"""
Filter tools semantically based on query.
Args:
query: User query to match against tools
available_tools: Full list of available MCP tools
top_k: Override default top_k (optional)
Returns:
Filtered and ordered list of tools (up to top_k)
"""
# Early returns for cases where we can't/shouldn't filter
if not self.enabled:
return available_tools
if not available_tools:
return available_tools
if not query or not query.strip():
return available_tools
# Router should be built on startup - if not, something went wrong
if self.tool_router is None:
verbose_logger.warning("Router not initialized - was build_router_from_mcp_registry() called on startup?")
return available_tools
# Run semantic filtering
try:
limit = top_k or self.top_k
matches = self.tool_router(text=query, limit=limit)
matched_tool_names = self._extract_tool_names_from_matches(matches)
if not matched_tool_names:
return available_tools
return self._get_tools_by_names(matched_tool_names, available_tools)
except Exception as e:
verbose_logger.error(f"Semantic tool filter failed: {e}", exc_info=True)
return available_tools
def _extract_tool_names_from_matches(self, matches) -> List[str]:
"""Extract tool names from semantic router match results."""
if not matches:
return []
# Handle single match
if hasattr(matches, "name") and matches.name:
return [matches.name]
# Handle list of matches
if isinstance(matches, list):
return [m.name for m in matches if hasattr(m, "name") and m.name]
return []
def _get_tools_by_names(
self, tool_names: List[str], available_tools: List[Any]
) -> List[Any]:
"""Get tools from available_tools by their names, preserving order."""
# Match tools from available_tools (preserves format - dict or MCPTool)
matched_tools = []
for tool in available_tools:
tool_name, _ = self._extract_tool_info(tool)
if tool_name in tool_names:
matched_tools.append(tool)
# Reorder to match semantic router's ordering
tool_map = {self._extract_tool_info(t)[0]: t for t in matched_tools}
return [tool_map[name] for name in tool_names if name in tool_map]
def extract_user_query(self, messages: List[Dict[str, Any]]) -> str:
"""
Extract user query from messages for /chat/completions or /responses.
Args:
messages: List of message dictionaries (from 'messages' or 'input' field)
Returns:
Extracted query string
"""
for msg in reversed(messages):
if msg.get("role") == "user":
content = msg.get("content", "")
if isinstance(content, str):
return content
if isinstance(content, list):
texts = [
block.get("text", "") if isinstance(block, dict) else str(block)
for block in content
if isinstance(block, (dict, str))
]
return " ".join(texts)
return ""

12
litellm/proxy/_types.py
View file

@ -359,7 +359,6 @@ class LiteLLMRoutes(enum.Enum):
"/v1/vector_stores/{vector_store_id}/files/{file_id}/content",
"/vector_store/list",
"/v1/vector_store/list",
# search
"/search",
"/v1/search",
@ -631,6 +630,9 @@ class LiteLLMRoutes(enum.Enum):
"/model/{model_id}/update",
"/prompt/list",
"/prompt/info",
"/guardrails",
"/guardrails/{guardrail_id}",
"/v2/guardrails/list",
] # routes that manage their own allowed/disallowed logic
## Org Admin Routes ##
@ -1482,6 +1484,9 @@ class TeamBase(LiteLLMPydanticObjectBase):
members: list = []
members_with_roles: List[Member] = []
team_member_permissions: Optional[List[str]] = None
allow_team_guardrail_config: Optional[
bool
] = None # if True, team admin can configure guardrails for this team
metadata: Optional[dict] = None
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
@ -1579,6 +1584,9 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
model_tpm_limit: Optional[Dict[str, int]] = None
allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
router_settings: Optional[dict] = None
allow_team_guardrail_config: Optional[
bool
] = None # if True, team admin can configure guardrails for this team
class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase):
@ -3673,7 +3681,7 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
team_id_upsert: bool = False
team_ids_jwt_field: Optional[str] = None
upsert_sso_user_to_team: bool = False
team_allowed_routes: List[str] = ["openai_routes", "info_routes"]
team_allowed_routes: List[str] = ["openai_routes", "info_routes", "mcp_routes"]
team_id_default: Optional[str] = Field(
default=None,
description="If no team_id given, default permissions/spend-tracking to this team.s",

53
litellm/proxy/agent_endpoints/a2a_routing.py Normal file
View file

@ -0,0 +1,53 @@
"""
A2A Agent Routing
Handles routing for A2A agents (models with "a2a/<agent-name>" prefix).
Looks up agents in the registry and injects their API base URL.
"""
from typing import Any, Optional
import litellm
from litellm._logging import verbose_proxy_logger
async def route_a2a_agent_request(data: dict, route_type: str) -> Optional[Any]:
"""
Route A2A agent requests directly to litellm with injected API base.
Returns None if not an A2A request (allows normal routing to continue).
"""
# Import here to avoid circular imports
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.proxy.route_llm_request import (
ROUTE_ENDPOINT_MAPPING,
ProxyModelNotFoundError,
)
model_name = data.get("model", "")
# Check if this is an A2A agent request
if not isinstance(model_name, str) or not model_name.startswith("a2a/"):
return None
# Extract agent name (e.g., "a2a/my-agent" -> "my-agent")
agent_name = model_name[4:]
# Look up agent in registry
agent = global_agent_registry.get_agent_by_name(agent_name)
if agent is None:
verbose_proxy_logger.error(f"[A2A] Agent '{agent_name}' not found in registry")
route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type)
raise ProxyModelNotFoundError(route=route_name, model_name=model_name)
# Get API base URL from agent config
if not agent.agent_card_params or "url" not in agent.agent_card_params:
verbose_proxy_logger.error(f"[A2A] Agent '{agent_name}' has no URL configured")
route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type)
raise ProxyModelNotFoundError(route=route_name, model_name=model_name)
# Inject API base and route to litellm
data["api_base"] = agent.agent_card_params["url"]
verbose_proxy_logger.debug(f"[A2A] Routing {model_name} to {data['api_base']}")
return getattr(litellm, f"{route_type}")(**data)

96
litellm/proxy/agent_endpoints/model_list_helpers.py Normal file
View file

@ -0,0 +1,96 @@
"""
Helper functions for appending A2A agents to model lists.
Used by proxy model endpoints to make agents appear in UI alongside models.
"""
from typing import List
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
ModelGroupInfoProxy,
)
async def append_agents_to_model_group(
model_groups: List[ModelGroupInfoProxy],
user_api_key_dict: UserAPIKeyAuth,
) -> List[ModelGroupInfoProxy]:
"""
Append A2A agents to model groups list for UI display.
Converts agents to model format with "a2a/<agent-name>" naming
so they appear in playground and work with LiteLLM routing.
"""
try:
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
AgentRequestHandler,
)
allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(
user_api_key_auth=user_api_key_dict
)
for agent_id in allowed_agent_ids:
agent = global_agent_registry.get_agent_by_id(agent_id)
if agent is not None:
model_groups.append(
ModelGroupInfoProxy(
model_group=f"a2a/{agent.agent_name}",
mode="chat",
providers=["a2a"],
)
)
except Exception as e:
verbose_proxy_logger.debug(
f"Error appending agents to model_group/info: {e}"
)
return model_groups
async def append_agents_to_model_info(
models: List[dict],
user_api_key_dict: UserAPIKeyAuth,
) -> List[dict]:
"""
Append A2A agents to model info list for UI display.
Converts agents to model format with "a2a/<agent-name>" naming
so they appear in models page and work with LiteLLM routing.
"""
try:
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
AgentRequestHandler,
)
allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(
user_api_key_auth=user_api_key_dict
)
for agent_id in allowed_agent_ids:
agent = global_agent_registry.get_agent_by_id(agent_id)
if agent is not None:
models.append({
"model_name": f"a2a/{agent.agent_name}",
"litellm_params": {
"model": f"a2a/{agent.agent_name}",
"custom_llm_provider": "a2a",
},
"model_info": {
"id": agent.agent_id,
"mode": "chat",
"db_model": True,
"created_by": agent.created_by,
"created_at": agent.created_at,
"updated_at": agent.updated_at,
},
})
except Exception as e:
verbose_proxy_logger.debug(
f"Error appending agents to v2/model/info: {e}"
)
return models

8
litellm/proxy/auth/auth_checks.py
View file

@ -198,7 +198,7 @@ async def common_checks(
message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}",
type=ProxyErrorTypes.team_model_access_denied,
param="model",
code=status.HTTP_401_UNAUTHORIZED,
code=status.HTTP_400_BAD_REQUEST,
)
## 2.1 If user can call model (if personal key)
@ -2056,7 +2056,7 @@ def _can_object_call_model(
object_type=object_type
),
param="model",
code=status.HTTP_401_UNAUTHORIZED,
code=status.HTTP_400_BAD_REQUEST,
)
@ -2157,7 +2157,7 @@ async def can_user_call_model(
message=f"User not allowed to access model. No default model access, only team models allowed. Tried to access {model}",
type=ProxyErrorTypes.key_model_access_denied,
param="model",
code=status.HTTP_401_UNAUTHORIZED,
code=status.HTTP_400_BAD_REQUEST,
)
return _can_object_call_model(
@ -2739,7 +2739,7 @@ def _can_object_call_vector_stores(
object_type
),
param="vector_store",
code=status.HTTP_401_UNAUTHORIZED,
code=status.HTTP_400_BAD_REQUEST,
)
return True

65
litellm/proxy/auth/auth_exception_handler.py
View file

@ -9,7 +9,10 @@ from fastapi import HTTPException, Request, status
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import _get_request_ip_address
from litellm.proxy.auth.auth_utils import (
_get_request_ip_address,
add_client_context_to_request_data,
)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.types.services import ServiceTypes
@ -30,6 +33,7 @@ class UserAPIKeyAuthExceptionHandler:
route: str,
parent_otel_span: Optional[Span],
api_key: str,
valid_token: Optional[UserAPIKeyAuth] = None,
) -> UserAPIKeyAuth:
"""
Handles Connection Errors when reading a Virtual Key from LiteLLM DB
@ -71,30 +75,63 @@ class UserAPIKeyAuthExceptionHandler:
)
else:
# raise the exception to the caller
use_x_forwarded_for = general_settings.get("use_x_forwarded_for", False)
requester_ip = _get_request_ip_address(
request=request,
use_x_forwarded_for=general_settings.get("use_x_forwarded_for", False),
use_x_forwarded_for=use_x_forwarded_for,
)
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - {}\nRequester IP Address:{}".format(
str(e),
requester_ip,
),
extra={"requester_ip": requester_ip},
user_agent = request.headers.get("user-agent", "") if request else ""
# Ensure request_data has client context for callbacks (Prometheus, etc.)
add_client_context_to_request_data(
request=request,
request_data=request_data,
use_x_forwarded_for=use_x_forwarded_for,
)
# Log this exception to OTEL, Datadog etc
user_api_key_dict = UserAPIKeyAuth(
parent_otel_span=parent_otel_span,
api_key=api_key,
request_route=route,
key_name = (
valid_token.key_alias or getattr(valid_token, "key_name", None)
if valid_token
else "<unknown-key>"
)
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - {}\nRequester IP Address:{}\nUser-Agent:{}\nKey Hash:{}\nKey Name:{}".format(
str(e),
requester_ip or "<unknown>",
user_agent or "<unknown>",
api_key or "<unknown>",
key_name,
),
extra={
"requester_ip": requester_ip,
"user_agent": user_agent,
"key_hash": api_key,
"key_name": key_name,
},
)
# Log this exception to OTEL, Datadog etc - use valid_token when available (e.g. model access denied)
if valid_token is not None:
user_api_key_dict = valid_token
else:
user_api_key_dict = UserAPIKeyAuth(
parent_otel_span=parent_otel_span,
api_key=api_key,
request_route=route,
key_alias="<unknown-key>",
)
# Allow callbacks to transform the error response
error_type = ProxyErrorTypes.auth_error
if isinstance(e, ProxyException) and hasattr(e, "type"):
try:
error_type = ProxyErrorTypes(e.type)
except (ValueError, TypeError):
pass
transformed_exception = await proxy_logging_obj.post_call_failure_hook(
request_data=request_data,
original_exception=e,
user_api_key_dict=user_api_key_dict,
error_type=ProxyErrorTypes.auth_error,
error_type=error_type,
route=route,
)
# Use transformed exception if callback returned one, otherwise use original

145
litellm/proxy/auth/auth_utils.py
View file

@ -26,6 +26,28 @@ def _get_request_ip_address(
return client_ip
def add_client_context_to_request_data(
request: Request, request_data: dict, use_x_forwarded_for: bool = False
) -> None:
"""
Add client_ip (requester_ip_address) and User-Agent to request_data metadata early.
Call this at the start of the request pipeline so failures have client context for logging.
"""
if "metadata" not in request_data:
request_data["metadata"] = {}
metadata = request_data["metadata"]
requester_ip = _get_request_ip_address(
request=request, use_x_forwarded_for=use_x_forwarded_for
)
metadata["requester_ip_address"] = requester_ip or ""
user_agent = ""
if hasattr(request, "headers") and "user-agent" in request.headers:
user_agent = request.headers.get("user-agent", "")
metadata["user_agent"] = user_agent
def _check_valid_ip(
allowed_ips: Optional[List[str]],
request: Request,
@ -314,17 +336,17 @@ def get_request_route(request: Request) -> str:
def normalize_request_route(route: str) -> str:
"""
Normalize request routes by replacing dynamic path parameters with placeholders.
This prevents high cardinality in Prometheus metrics by collapsing routes like:
- /v1/responses/1234567890 -> /v1/responses/{response_id}
- /v1/threads/thread_123 -> /v1/threads/{thread_id}
Args:
route: The request route path
Returns:
Normalized route with dynamic parameters replaced by placeholders
Examples:
>>> normalize_request_route("/v1/responses/abc123")
'/v1/responses/{response_id}'
@ -337,58 +359,90 @@ def normalize_request_route(route: str) -> str:
# Format: (regex_pattern, replacement_template)
patterns = [
# Responses API - must come before generic patterns
(r'^(/(?:openai/)?v1/responses)/([^/]+)(/input_items)$', r'\1/{response_id}\3'),
(r'^(/(?:openai/)?v1/responses)/([^/]+)(/cancel)$', r'\1/{response_id}\3'),
(r'^(/(?:openai/)?v1/responses)/([^/]+)$', r'\1/{response_id}'),
(r'^(/responses)/([^/]+)(/input_items)$', r'\1/{response_id}\3'),
(r'^(/responses)/([^/]+)(/cancel)$', r'\1/{response_id}\3'),
(r'^(/responses)/([^/]+)$', r'\1/{response_id}'),
(r"^(/(?:openai/)?v1/responses)/([^/]+)(/input_items)$", r"\1/{response_id}\3"),
(r"^(/(?:openai/)?v1/responses)/([^/]+)(/cancel)$", r"\1/{response_id}\3"),
(r"^(/(?:openai/)?v1/responses)/([^/]+)$", r"\1/{response_id}"),
(r"^(/responses)/([^/]+)(/input_items)$", r"\1/{response_id}\3"),
(r"^(/responses)/([^/]+)(/cancel)$", r"\1/{response_id}\3"),
(r"^(/responses)/([^/]+)$", r"\1/{response_id}"),
# Threads API
(r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)/([^/]+)$', r'\1/{thread_id}\3/{run_id}\5/{step_id}'),
(r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)$', r'\1/{thread_id}\3/{run_id}\5'),
(r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/cancel)$', r'\1/{thread_id}\3/{run_id}\5'),
(r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/submit_tool_outputs)$', r'\1/{thread_id}\3/{run_id}\5'),
(r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)$', r'\1/{thread_id}\3/{run_id}'),
(r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)$', r'\1/{thread_id}\3'),
(r'^(/(?:openai/)?v1/threads)/([^/]+)(/messages)/([^/]+)$', r'\1/{thread_id}\3/{message_id}'),
(r'^(/(?:openai/)?v1/threads)/([^/]+)(/messages)$', r'\1/{thread_id}\3'),
(r'^(/(?:openai/)?v1/threads)/([^/]+)$', r'\1/{thread_id}'),
(
r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)/([^/]+)$",
r"\1/{thread_id}\3/{run_id}\5/{step_id}",
),
(
r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)$",
r"\1/{thread_id}\3/{run_id}\5",
),
(
r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/cancel)$",
r"\1/{thread_id}\3/{run_id}\5",
),
(
r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/submit_tool_outputs)$",
r"\1/{thread_id}\3/{run_id}\5",
),
(
r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)$",
r"\1/{thread_id}\3/{run_id}",
),
(r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)$", r"\1/{thread_id}\3"),
(
r"^(/(?:openai/)?v1/threads)/([^/]+)(/messages)/([^/]+)$",
r"\1/{thread_id}\3/{message_id}",
),
(r"^(/(?:openai/)?v1/threads)/([^/]+)(/messages)$", r"\1/{thread_id}\3"),
(r"^(/(?:openai/)?v1/threads)/([^/]+)$", r"\1/{thread_id}"),
# Vector Stores API
(r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)/([^/]+)$', r'\1/{vector_store_id}\3/{file_id}'),
(r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)$', r'\1/{vector_store_id}\3'),
(r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)/([^/]+)$', r'\1/{vector_store_id}\3/{batch_id}'),
(r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)$', r'\1/{vector_store_id}\3'),
(r'^(/(?:openai/)?v1/vector_stores)/([^/]+)$', r'\1/{vector_store_id}'),
(
r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)/([^/]+)$",
r"\1/{vector_store_id}\3/{file_id}",
),
(
r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)$",
r"\1/{vector_store_id}\3",
),
(
r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)/([^/]+)$",
r"\1/{vector_store_id}\3/{batch_id}",
),
(
r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)$",
r"\1/{vector_store_id}\3",
),
(r"^(/(?:openai/)?v1/vector_stores)/([^/]+)$", r"\1/{vector_store_id}"),
# Assistants API
(r'^(/(?:openai/)?v1/assistants)/([^/]+)$', r'\1/{assistant_id}'),
(r"^(/(?:openai/)?v1/assistants)/([^/]+)$", r"\1/{assistant_id}"),
# Files API
(r'^(/(?:openai/)?v1/files)/([^/]+)(/content)$', r'\1/{file_id}\3'),
(r'^(/(?:openai/)?v1/files)/([^/]+)$', r'\1/{file_id}'),
(r"^(/(?:openai/)?v1/files)/([^/]+)(/content)$", r"\1/{file_id}\3"),
(r"^(/(?:openai/)?v1/files)/([^/]+)$", r"\1/{file_id}"),
# Batches API
(r'^(/(?:openai/)?v1/batches)/([^/]+)(/cancel)$', r'\1/{batch_id}\3'),
(r'^(/(?:openai/)?v1/batches)/([^/]+)$', r'\1/{batch_id}'),
(r"^(/(?:openai/)?v1/batches)/([^/]+)(/cancel)$", r"\1/{batch_id}\3"),
(r"^(/(?:openai/)?v1/batches)/([^/]+)$", r"\1/{batch_id}"),
# Fine-tuning API
(r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/events)$', r'\1/{fine_tuning_job_id}\3'),
(r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/cancel)$', r'\1/{fine_tuning_job_id}\3'),
(r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/checkpoints)$', r'\1/{fine_tuning_job_id}\3'),
(r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)$', r'\1/{fine_tuning_job_id}'),
(
r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/events)$",
r"\1/{fine_tuning_job_id}\3",
),
(
r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/cancel)$",
r"\1/{fine_tuning_job_id}\3",
),
(
r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/checkpoints)$",
r"\1/{fine_tuning_job_id}\3",
),
(r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)$", r"\1/{fine_tuning_job_id}"),
# Models API
(r'^(/(?:openai/)?v1/models)/([^/]+)$', r'\1/{model}'),
(r"^(/(?:openai/)?v1/models)/([^/]+)$", r"\1/{model}"),
]
# Apply patterns in order
for pattern, replacement in patterns:
normalized = re.sub(pattern, replacement, route)
if normalized != route:
return normalized
# Return original route if no pattern matched
return route
@ -644,6 +698,7 @@ def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[str]:
return header_name
return None
def _get_customer_id_from_standard_headers(
request_headers: Optional[dict],
) -> Optional[str]:
@ -679,7 +734,9 @@ def get_end_user_id_from_request_body(
from litellm.proxy.proxy_server import general_settings
# Check 1: Standard customer ID headers (always checked, no configuration required)
customer_id = _get_customer_id_from_standard_headers(request_headers=request_headers)
customer_id = _get_customer_id_from_standard_headers(
request_headers=request_headers
)
if customer_id is not None:
return customer_id

3
litellm/proxy/auth/handle_jwt.py
View file

@ -976,6 +976,9 @@ class JWTAuthManager:
user_route=route,
litellm_proxy_roles=jwt_handler.litellm_jwtauth,
)
verbose_proxy_logger.debug(
f"JWT team route check: team_id={team_id}, route={route}, is_allowed={is_allowed}"
)
if is_allowed:
return team_id, team_object
except Exception:

60
litellm/proxy/auth/login_utils.py
View file

@ -34,6 +34,58 @@ from litellm.secret_managers.main import get_secret_bool
from litellm.types.proxy.ui_sso import ReturnedUITokenObject
async def expire_previous_ui_session_tokens(
user_id: str, prisma_client: Optional[PrismaClient]
) -> None:
"""
Expire (block) all other valid UI session tokens for a user.
This prevents accumulation of multiple valid UI session tokens that
are supposed to be short-lived test keys. Only affects keys with
team_id = "litellm-dashboard" and that haven't expired yet.
Args:
user_id: The user ID whose previous UI session tokens should be expired
prisma_client: Database client for performing the update
"""
if prisma_client is None:
return
try:
from datetime import datetime, timezone
current_time = datetime.now(timezone.utc)
# Find all unblocked AND non-expired UI session tokens for this user
ui_session_tokens = await prisma_client.db.litellm_verificationtoken.find_many(
where={
"user_id": user_id,
"team_id": "litellm-dashboard",
"OR": [
{"blocked": None}, # Tokens that have never been blocked (null)
{"blocked": False}, # Tokens explicitly set to not blocked
],
"expires": {"gt": current_time}, # Only get tokens that haven't expired
}
)
if not ui_session_tokens:
return
# Block all the found tokens
tokens_to_block = [token.token for token in ui_session_tokens if token.token]
if tokens_to_block:
await prisma_client.db.litellm_verificationtoken.update_many(
where={"token": {"in": tokens_to_block}}, data={"blocked": True}
)
except Exception:
# Silently fail - don't block login if cleanup fails
# This is a best-effort operation
pass
def get_ui_credentials(master_key: Optional[str]) -> tuple[str, str]:
"""
Get UI username and password from environment variables or master key.
@ -245,6 +297,11 @@ async def authenticate_user( # noqa: PLR0915
)
user_email = getattr(_user_row, "user_email", "unknown")
_password = getattr(_user_row, "password", "unknown")
user_team_id = getattr(_user_row, "team_id", "litellm-dashboard")
# if user_team_id is None, set it to "litellm-dashboard"
if user_team_id is None:
user_team_id = "litellm-dashboard"
if _password is None:
raise ProxyException(
@ -271,7 +328,7 @@ async def authenticate_user( # noqa: PLR0915
"config": {},
"spend": 0,
"user_id": user_id,
"team_id": "litellm-dashboard",
"team_id": user_team_id,
},
)
else:
@ -340,4 +397,3 @@ def create_ui_token_object(
disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation,
server_root_path=get_server_root_path(),
)

25
litellm/proxy/auth/model_checks.py
View file

@ -64,27 +64,6 @@ def _get_models_from_access_groups(
return all_models
def get_access_groups_from_models(
model_access_groups: Dict[str, List[str]],
models: List[str],
) -> List[str]:
"""
Extract access group names from a models list.
Given a models list like ["gpt-4", "beta-models", "claude-v1"]
and access groups like {"beta-models": ["gpt-5", "gpt-6"]},
returns ["beta-models"].
This is used to pass allowed access groups to the router for filtering
deployments during load balancing (GitHub issue #18333).
"""
access_groups = []
for model in models:
if model in model_access_groups:
access_groups.append(model)
return access_groups
async def get_mcp_server_ids(
user_api_key_dict: UserAPIKeyAuth,
) -> List[str]:
@ -101,6 +80,7 @@ async def get_mcp_server_ids(
# Make a direct SQL query to get just the mcp_servers
try:
result = await prisma_client.db.litellm_objectpermissiontable.find_unique(
where={"object_permission_id": user_api_key_dict.object_permission_id},
)
@ -196,7 +176,6 @@ def get_complete_model_list(
"""
unique_models = []
def append_unique(models):
for model in models:
if model not in unique_models:
@ -209,7 +188,7 @@ def get_complete_model_list(
else:
append_unique(proxy_model_list)
if include_model_access_groups:
append_unique(list(model_access_groups.keys())) # TODO: keys order
append_unique(list(model_access_groups.keys())) # TODO: keys order
if user_model:
append_unique([user_model])

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

@ -42,6 +42,7 @@ from litellm.proxy.auth.auth_checks import (
from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler
from litellm.proxy.auth.auth_utils import (
abbreviate_api_key,
add_client_context_to_request_data,
get_end_user_id_from_request_body,
get_model_from_request,
get_request_route,
@ -247,7 +248,9 @@ async def get_global_proxy_spend(
proxy_logging_obj: ProxyLogging,
) -> Optional[float]:
global_proxy_spend = None
if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget
if (
litellm.max_budget > 0 and prisma_client is not None
): # user set proxy max budget
# Use event-driven coordination to prevent cache stampede
cache_key = "{}:spend".format(litellm_proxy_admin_name)
global_proxy_spend = await _fetch_global_spend_with_event_coordination(
@ -1251,6 +1254,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
route=route,
parent_otel_span=parent_otel_span,
api_key=api_key,
valid_token=valid_token,
)
@ -1278,6 +1282,14 @@ async def user_api_key_auth(
request_data = populate_request_with_path_params(
request_data=request_data, request=request
)
# Capture client context early so failures have IP/User-Agent for logging and metrics
from litellm.proxy.proxy_server import general_settings
add_client_context_to_request_data(
request=request,
request_data=request_data,
use_x_forwarded_for=general_settings.get("use_x_forwarded_for", False),
)
route: str = get_request_route(request=request)
## CHECK IF ROUTE IS ALLOWED

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

@ -11,7 +11,7 @@ from pydantic import BaseModel
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry
from litellm.types.guardrails import (
@ -32,6 +32,8 @@ from litellm.types.guardrails import (
PresidioPresidioConfigModelUserInterface,
SupportedGuardrailIntegrations,
ToolPermissionGuardrailConfigModel,
CreateGuardrailRequest,
UpdateGuardrailRequest,
)
#### GUARDRAILS ENDPOINTS ####
@ -40,6 +42,37 @@ router = APIRouter()
GUARDRAIL_REGISTRY = GuardrailRegistry()
async def _check_team_can_configure_guardrails(
user_api_key_dict: UserAPIKeyAuth,
team_id: Optional[str],
prisma_client: Any,
user_api_key_cache: Any,
) -> None:
"""
If the user is not proxy admin and team_id is set, verify the team has
allow_team_guardrail_config enabled. Raise HTTPException 403 otherwise.
"""
if team_id is None or team_id == "litellm-dashboard":
return
user_role = getattr(user_api_key_dict, "user_role", None)
if user_role == LitellmUserRoles.PROXY_ADMIN:
return
from litellm.proxy.auth.auth_checks import get_team_object
team_table = await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
check_db_only=True,
)
allow_config = getattr(team_table, "allow_team_guardrail_config", False)
if not allow_config:
raise HTTPException(
status_code=403,
detail="Guardrail configuration is not enabled for this team. Contact your administrator to enable it.",
)
def _get_guardrails_list_response(
guardrails_config: List[Dict],
) -> ListGuardrailsResponse:
@ -64,9 +97,19 @@ def _get_guardrails_list_response(
dependencies=[Depends(user_api_key_auth)],
response_model=ListGuardrailsResponse,
)
async def list_guardrails():
@router.get(
"/v2/guardrails/list",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
response_model=ListGuardrailsResponse,
)
async def list_guardrails(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
List the guardrails that are available on the proxy server
List the guardrails that are available in the database using GuardrailRegistry
Now supports both DB-persisted and In-Memory (Config) guardrails.
👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start)
@ -75,59 +118,6 @@ async def list_guardrails():
curl -X GET "http://localhost:4000/guardrails/list" -H "Authorization: Bearer <your_api_key>"
```
Example Response:
```json
{
"guardrails": [
{
"guardrail_name": "bedrock-pre-guard",
"guardrail_info": {
"params": [
{
"name": "toxicity_score",
"type": "float",
"description": "Score between 0-1 indicating content toxicity level"
},
{
"name": "pii_detection",
"type": "boolean"
}
]
}
}
]
}
```
"""
from litellm.proxy.proxy_server import proxy_config
config = proxy_config.config
_guardrails_config = cast(Optional[list[dict]], config.get("guardrails"))
if _guardrails_config is None:
return _get_guardrails_list_response([])
return _get_guardrails_list_response(_guardrails_config)
@router.get(
"/v2/guardrails/list",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
response_model=ListGuardrailsResponse,
)
async def list_guardrails_v2():
"""
List the guardrails that are available in the database using GuardrailRegistry
👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start)
Example Request:
```bash
curl -X GET "http://localhost:4000/v2/guardrails/list" -H "Authorization: Bearer <your_api_key>"
```
Example Response:
```json
{
@ -139,12 +129,14 @@ async def list_guardrails_v2():
"guardrail": "bedrock",
"mode": "pre_call",
"guardrailIdentifier": "ff6ujrregl1q",
"guardrailVersion": "DRAFT",
"guardrailVersion": "1.0",
"default_on": true
},
"guardrail_info": {
"description": "Bedrock content moderation guardrail"
}
"description": "Updated Bedrock content moderation guardrail"
},
"created_at": "2023-11-09T12:34:56.789Z",
"updated_at": "2023-11-09T13:45:12.345Z"
}
]
}
@ -153,12 +145,29 @@ async def list_guardrails_v2():
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
# Check if DB is connected
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
# Fallback to config only if DB is missing (though prisma_client usually exists)
from litellm.proxy.proxy_server import proxy_config
config = proxy_config.config
_guardrails_config = cast(Optional[list[dict]], config.get("guardrails"))
if _guardrails_config is None:
return _get_guardrails_list_response([])
return _get_guardrails_list_response(_guardrails_config)
try:
team_id = getattr(user_api_key_dict, "team_id", None)
user_role = getattr(user_api_key_dict, "user_role", None)
filter_team_id = None
if user_role != LitellmUserRoles.PROXY_ADMIN or (
team_id is not None and team_id != "litellm-dashboard"
):
filter_team_id = team_id
guardrails = await GUARDRAIL_REGISTRY.get_all_guardrails_from_db(
prisma_client=prisma_client
prisma_client=prisma_client, team_id=filter_team_id
)
guardrail_configs: List[GuardrailInfoResponse] = []
@ -173,6 +182,7 @@ async def list_guardrails_v2():
created_at=guardrail.get("created_at"),
updated_at=guardrail.get("updated_at"),
guardrail_definition_location="db",
team_id=guardrail.get("team_id"),
)
)
seen_guardrail_ids.add(guardrail.get("guardrail_id"))
@ -180,6 +190,15 @@ async def list_guardrails_v2():
# get guardrails initialized on litellm config.yaml
in_memory_guardrails = IN_MEMORY_GUARDRAIL_HANDLER.list_in_memory_guardrails()
for guardrail in in_memory_guardrails:
# Check access for in-memory guardrails too
if filter_team_id:
g_team = guardrail.get("team_id")
# If guardrail has a team_id, it must match.
# If guardrail has NO team_id, it is likely a global config guardrail, so we usually allow it.
if g_team and g_team != filter_team_id:
continue
# only add guardrails that are not in DB guardrail list already
if guardrail.get("guardrail_id") not in seen_guardrail_ids:
guardrail_configs.append(
@ -189,6 +208,7 @@ async def list_guardrails_v2():
litellm_params=dict(guardrail.get("litellm_params") or {}),
guardrail_info=dict(guardrail.get("guardrail_info") or {}),
guardrail_definition_location="config",
team_id=guardrail.get("team_id"),
)
)
seen_guardrail_ids.add(guardrail.get("guardrail_id"))
@ -199,16 +219,15 @@ async def list_guardrails_v2():
raise HTTPException(status_code=500, detail=str(e))
class CreateGuardrailRequest(BaseModel):
guardrail: Guardrail
@router.post(
"/guardrails",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def create_guardrail(request: CreateGuardrailRequest):
async def create_guardrail(
request: CreateGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Create a new guardrail
@ -257,14 +276,21 @@ async def create_guardrail(request: CreateGuardrailRequest):
```
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
try:
team_id = getattr(user_api_key_dict, "team_id", None)
await _check_team_can_configure_guardrails(
user_api_key_dict=user_api_key_dict,
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
result = await GUARDRAIL_REGISTRY.add_guardrail_to_db(
guardrail=request.guardrail, prisma_client=prisma_client
guardrail=request.guardrail, prisma_client=prisma_client, team_id=team_id
)
guardrail_name = result.get("guardrail_name", "Unknown")
@ -283,21 +309,23 @@ async def create_guardrail(request: CreateGuardrailRequest):
)
return result
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(f"Error adding guardrail to db: {e}")
raise HTTPException(status_code=500, detail=str(e))
class UpdateGuardrailRequest(BaseModel):
guardrail: Guardrail
@router.put(
"/guardrails/{guardrail_id}",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest):
async def update_guardrail(
guardrail_id: str,
request: UpdateGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Update an existing guardrail
@ -346,7 +374,7 @@ async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest):
```
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@ -362,6 +390,25 @@ async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest):
status_code=404, detail=f"Guardrail with ID {guardrail_id} not found"
)
if existing_guardrail.get("team_id"):
team_id = getattr(user_api_key_dict, "team_id", None)
if getattr(
user_api_key_dict, "user_role", None
) != LitellmUserRoles.PROXY_ADMIN or (
team_id is not None and team_id != "litellm-dashboard"
):
if existing_guardrail.get("team_id") != team_id:
raise HTTPException(
status_code=403,
detail="Not authorized to access this guardrail",
)
await _check_team_can_configure_guardrails(
user_api_key_dict=user_api_key_dict,
team_id=existing_guardrail.get("team_id"),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
result = await GUARDRAIL_REGISTRY.update_guardrail_in_db(
guardrail_id=guardrail_id,
guardrail=request.guardrail,
@ -394,7 +441,9 @@ async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest):
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def delete_guardrail(guardrail_id: str):
async def delete_guardrail(
guardrail_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)
):
"""
Delete a guardrail
@ -414,7 +463,7 @@ async def delete_guardrail(guardrail_id: str):
```
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@ -430,6 +479,25 @@ async def delete_guardrail(guardrail_id: str):
status_code=404, detail=f"Guardrail with ID {guardrail_id} not found"
)
if existing_guardrail.get("team_id"):
team_id = getattr(user_api_key_dict, "team_id", None)
if getattr(
user_api_key_dict, "user_role", None
) != LitellmUserRoles.PROXY_ADMIN or (
team_id is not None and team_id != "litellm-dashboard"
):
if existing_guardrail.get("team_id") != team_id:
raise HTTPException(
status_code=403,
detail="Not authorized to access this guardrail",
)
await _check_team_can_configure_guardrails(
user_api_key_dict=user_api_key_dict,
team_id=existing_guardrail.get("team_id"),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
result = await GUARDRAIL_REGISTRY.delete_guardrail_from_db(
guardrail_id=guardrail_id, prisma_client=prisma_client
)
@ -460,7 +528,11 @@ async def delete_guardrail(guardrail_id: str):
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest):
async def patch_guardrail(
guardrail_id: str,
request: PatchGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Partially update an existing guardrail
@ -507,7 +579,7 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest):
```
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@ -523,6 +595,25 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest):
status_code=404, detail=f"Guardrail with ID {guardrail_id} not found"
)
if existing_guardrail.get("team_id"):
team_id = getattr(user_api_key_dict, "team_id", None)
if getattr(
user_api_key_dict, "user_role", None
) != LitellmUserRoles.PROXY_ADMIN or (
team_id is not None and team_id != "litellm-dashboard"
):
if existing_guardrail.get("team_id") != team_id:
raise HTTPException(
status_code=403,
detail="Not authorized to access this guardrail",
)
await _check_team_can_configure_guardrails(
user_api_key_dict=user_api_key_dict,
team_id=existing_guardrail.get("team_id"),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
# Create updated guardrail object
guardrail_name = (
request.guardrail_name
@ -594,7 +685,10 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest):
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def get_guardrail_info(guardrail_id: str):
async def get_guardrail_info(
guardrail_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get detailed information about a specific guardrail by ID
@ -653,6 +747,19 @@ async def get_guardrail_info(guardrail_id: str):
status_code=404, detail=f"Guardrail with ID {guardrail_id} not found"
)
if result.get("team_id"):
team_id = getattr(user_api_key_dict, "team_id", None)
if getattr(
user_api_key_dict, "user_role", None
) != LitellmUserRoles.PROXY_ADMIN or (
team_id is not None and team_id != "litellm-dashboard"
):
if result.get("team_id") != team_id:
raise HTTPException(
status_code=403,
detail="Not authorized to access this guardrail",
)
litellm_params: Optional[Union[LitellmParams, dict]] = result.get(
"litellm_params"
)
@ -675,6 +782,7 @@ async def get_guardrail_info(guardrail_id: str):
created_at=result.get("created_at"),
updated_at=result.get("updated_at"),
guardrail_definition_location=guardrail_definition_location,
team_id=result.get("team_id"),
)
except HTTPException as e:
raise e
@ -1209,9 +1317,9 @@ async def get_provider_specific_params():
lakera_v2_fields = _get_fields_from_model(LakeraV2GuardrailConfigModel)
tool_permission_fields = _get_fields_from_model(ToolPermissionGuardrailConfigModel)
tool_permission_fields["ui_friendly_name"] = (
ToolPermissionGuardrailConfigModel.ui_friendly_name()
)
tool_permission_fields[
"ui_friendly_name"
] = ToolPermissionGuardrailConfigModel.ui_friendly_name()
# Return the provider-specific parameters
provider_params = {
@ -1519,10 +1627,10 @@ async def apply_guardrail(
from litellm.proxy.utils import handle_exception_on_proxy
try:
active_guardrail: Optional[CustomGuardrail] = (
GUARDRAIL_REGISTRY.get_initialized_guardrail_callback(
guardrail_name=request.guardrail_name
)
active_guardrail: Optional[
CustomGuardrail
] = GUARDRAIL_REGISTRY.get_initialized_guardrail_callback(
guardrail_name=request.guardrail_name
)
if active_guardrail is None:
raise HTTPException(

42
litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py
View file

@ -9,8 +9,10 @@ from fastapi import HTTPException
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -21,6 +23,8 @@ from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
GRAYSWAN_BLOCK_ERROR_MSG = "Blocked by Gray Swan Guardrail"
class GraySwanGuardrailMissingSecrets(Exception):
"""Raised when the Gray Swan API key is missing."""
@ -205,9 +209,13 @@ class GraySwanGuardrail(CustomGuardrail):
# Get dynamic params from request metadata
dynamic_body = self.get_guardrail_dynamic_request_body_params(request_data) or {}
if dynamic_body:
verbose_proxy_logger.debug(
"Gray Swan Guardrail: dynamic extra_body=%s", safe_dumps(dynamic_body)
)
# Prepare and send payload
payload = self._prepare_payload(messages, dynamic_body)
payload = self._prepare_payload(messages, dynamic_body, request_data)
if payload is None:
return inputs
@ -223,6 +231,8 @@ class GraySwanGuardrail(CustomGuardrail):
)
return result
except Exception as exc:
if self._is_grayswan_exception(exc):
raise
end_time = time.time()
status_code = getattr(exc, "status_code", None) or getattr(
exc, "exception_status_code", None
@ -240,8 +250,20 @@ class GraySwanGuardrail(CustomGuardrail):
exc,
)
return inputs
if isinstance(exc, GraySwanGuardrailAPIError):
raise exc
raise GraySwanGuardrailAPIError(str(exc), status_code=status_code) from exc
def _is_grayswan_exception(self, exc: Exception) -> bool:
# Guardrail decision (passthrough) should always propagate,
# regardless of fail_open.
if isinstance(exc, ModifyResponseException):
return True
detail = getattr(exc, "detail", None)
if isinstance(detail, dict):
return detail.get("error") == GRAYSWAN_BLOCK_ERROR_MSG
return False
# ------------------------------------------------------------------
# Legacy Test Interface (for backward compatibility)
# ------------------------------------------------------------------
@ -324,7 +346,7 @@ class GraySwanGuardrail(CustomGuardrail):
raise HTTPException(
status_code=400,
detail={
"error": "Blocked by Gray Swan Guardrail",
"error": GRAYSWAN_BLOCK_ERROR_MSG,
"violation_location": violation_location,
"violation": violation_score,
"violated_rules": violated_rules,
@ -445,7 +467,7 @@ class GraySwanGuardrail(CustomGuardrail):
raise HTTPException(
status_code=400,
detail={
"error": "Blocked by Gray Swan Guardrail",
"error": GRAYSWAN_BLOCK_ERROR_MSG,
"violation_location": violation_location,
"violation": violation_score,
"violated_rules": violated_rules,
@ -494,7 +516,7 @@ class GraySwanGuardrail(CustomGuardrail):
}
def _prepare_payload(
self, messages: List[Dict[str, str]], dynamic_body: dict
self, messages: List[Dict[str, str]], dynamic_body: dict, request_data: dict
) -> Optional[Dict[str, Any]]:
payload: Dict[str, Any] = {"messages": messages}
@ -510,6 +532,18 @@ class GraySwanGuardrail(CustomGuardrail):
if reasoning_mode:
payload["reasoning_mode"] = reasoning_mode
# Pass through arbitrary metadata when provided via dynamic extra_body.
if "metadata" in dynamic_body:
payload["metadata"] = dynamic_body["metadata"]
litellm_metadata = request_data.get("litellm_metadata")
if isinstance(litellm_metadata, dict) and litellm_metadata:
cleaned_litellm_metadata = dict(litellm_metadata)
# cleaned_litellm_metadata.pop("user_api_key_auth", None)
sanitized = safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={})
if isinstance(sanitized, dict) and sanitized:
payload["litellm_metadata"] = sanitized
return payload
def _format_violation_message(

99
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -233,7 +233,10 @@ class GuardrailRegistry:
########### DB management helpers for guardrails ###########
############################################################
async def add_guardrail_to_db(
self, guardrail: Guardrail, prisma_client: PrismaClient
self,
guardrail: Guardrail,
prisma_client: PrismaClient,
team_id: Optional[str] = None,
):
"""
Add a guardrail to the database
@ -248,6 +251,8 @@ class GuardrailRegistry:
litellm_params_dict = (
dict(litellm_params_obj) if litellm_params_obj else {}
)
# Use safe_dumps to store as string, as Prisma client seems to prefer this for now
litellm_params: str = safe_dumps(litellm_params_dict)
guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {}))
@ -257,6 +262,7 @@ class GuardrailRegistry:
"guardrail_name": guardrail_name,
"litellm_params": litellm_params,
"guardrail_info": guardrail_info,
"team_id": team_id,
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
}
@ -265,18 +271,32 @@ class GuardrailRegistry:
# Add guardrail_id to the returned guardrail object
guardrail_dict = dict(guardrail)
guardrail_dict["guardrail_id"] = created_guardrail.guardrail_id
guardrail_dict["team_id"] = team_id
return guardrail_dict
except Exception as e:
raise Exception(f"Error adding guardrail to DB: {str(e)}")
async def delete_guardrail_from_db(
self, guardrail_id: str, prisma_client: PrismaClient
self,
guardrail_id: str,
prisma_client: PrismaClient,
team_id: Optional[str] = None,
):
"""
Delete a guardrail from the database
"""
try:
# Check ownership if team_id is provided
if team_id:
existing_guardrail = (
await prisma_client.db.litellm_guardrailstable.find_unique(
where={"guardrail_id": guardrail_id}
)
)
if not existing_guardrail or existing_guardrail.team_id != team_id:
raise Exception("Guardrail not found or access denied")
# Delete from DB
await prisma_client.db.litellm_guardrailstable.delete(
where={"guardrail_id": guardrail_id}
@ -287,12 +307,26 @@ class GuardrailRegistry:
raise Exception(f"Error deleting guardrail from DB: {str(e)}")
async def update_guardrail_in_db(
self, guardrail_id: str, guardrail: Guardrail, prisma_client: PrismaClient
self,
guardrail_id: str,
guardrail: Guardrail,
prisma_client: PrismaClient,
team_id: Optional[str] = None,
):
"""
Update a guardrail in the database
"""
try:
# Check ownership if team_id is provided
if team_id:
existing_guardrail = (
await prisma_client.db.litellm_guardrailstable.find_unique(
where={"guardrail_id": guardrail_id}
)
)
if not existing_guardrail or existing_guardrail.team_id != team_id:
raise Exception("Guardrail not found or access denied")
guardrail_name = guardrail.get("guardrail_name")
# Properly serialize LitellmParams Pydantic model to dict
litellm_params_obj: Any = guardrail.get("litellm_params", {})
@ -302,6 +336,7 @@ class GuardrailRegistry:
litellm_params_dict = (
dict(litellm_params_obj) if litellm_params_obj else {}
)
litellm_params: str = safe_dumps(litellm_params_dict)
guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {}))
@ -324,27 +359,67 @@ class GuardrailRegistry:
@staticmethod
async def get_all_guardrails_from_db(
prisma_client: PrismaClient,
team_id: Optional[str] = None,
) -> List[Guardrail]:
"""
Get all guardrails from the database
"""
try:
guardrails_from_db = (
await prisma_client.db.litellm_guardrailstable.find_many(
order={"created_at": "desc"},
)
# Normal ORM Fetch
all_guardrails = await prisma_client.db.litellm_guardrailstable.find_many(
order={"created_at": "desc"},
)
# Filter in Python
if team_id:
guardrails_from_db = [g for g in all_guardrails if g.team_id == team_id]
else:
guardrails_from_db = all_guardrails
guardrails: List[Guardrail] = []
for guardrail in guardrails_from_db:
guardrails.append(Guardrail(**(dict(guardrail)))) # type: ignore
try:
# Deep copy to avoid mutating cache/original
g_dict = (
guardrail.dict()
if hasattr(guardrail, "dict")
else dict(guardrail)
)
# Handle litellm_params
params = g_dict.get("litellm_params")
if isinstance(params, str):
import json
try:
g_dict["litellm_params"] = json.loads(params)
except Exception:
g_dict["litellm_params"] = {}
# Handle guardrail_info
info = g_dict.get("guardrail_info")
if isinstance(info, str):
import json
try:
g_dict["guardrail_info"] = json.loads(info)
except Exception:
g_dict["guardrail_info"] = {}
# Construct
guardrails.append(Guardrail(**g_dict)) # type: ignore
except Exception:
continue
return guardrails
except Exception as e:
raise Exception(f"Error getting guardrails from DB: {str(e)}")
async def get_guardrail_by_id_from_db(
self, guardrail_id: str, prisma_client: PrismaClient
self,
guardrail_id: str,
prisma_client: PrismaClient,
team_id: Optional[str] = None,
) -> Optional[Guardrail]:
"""
Get a guardrail by its ID from the database
@ -357,6 +432,11 @@ class GuardrailRegistry:
if not guardrail:
return None
if team_id and guardrail.team_id != team_id:
# Return None if not found or not owned by team
# Alternatively could raise exception, but returning None resembles "not found"
return None
return Guardrail(**(dict(guardrail))) # type: ignore
except Exception as e:
raise Exception(f"Error getting guardrail from DB: {str(e)}")
@ -473,6 +553,7 @@ class InMemoryGuardrailHandler:
guardrail_id=guardrail.get("guardrail_id"),
guardrail_name=guardrail["guardrail_name"],
litellm_params=litellm_params,
team_id=guardrail.get("team_id"),
)
# store references to the guardrail in memory

96
litellm/proxy/hooks/mcp_semantic_filter/ARCHITECTURE.md Normal file
View file

@ -0,0 +1,96 @@
# MCP Semantic Tool Filter Architecture
## Why Filter MCP Tools
When multiple MCP servers are connected, the proxy may expose hundreds of tools. Sending all tools in every request wastes context window tokens and increases cost. The semantic filter keeps only the top-K most relevant tools based on embedding similarity.
```mermaid
sequenceDiagram
participant Client
participant Hook as SemanticToolFilterHook
participant Filter as SemanticMCPToolFilter
participant Router as semantic-router
participant LLM
Client->>Hook: POST /chat/completions
Note over Client,Hook: tools: [100+ MCP tools]
Note over Client,Hook: messages: [{"role": "user", "content": "Get my Jira issues"}]
rect rgb(240, 240, 240)
Note over Hook: 1. Extract User Query
Hook->>Filter: filter_tools("Get my Jira issues", tools)
end
rect rgb(240, 240, 240)
Note over Filter: 2. Convert Tools → Routes
Note over Filter: Tool name + description → Route
end
rect rgb(240, 240, 240)
Note over Filter: 3. Semantic Matching
Filter->>Router: router(query)
Router->>Router: Embeddings + similarity
Router-->>Filter: [top 10 matches]
end
rect rgb(240, 240, 240)
Note over Filter: 4. Return Filtered Tools
Filter-->>Hook: [10 relevant tools]
end
Hook->>LLM: POST /chat/completions
Note over Hook,LLM: tools: [10 Jira-related tools] ← FILTERED
Note over Hook,LLM: messages: [...] ← UNCHANGED
LLM-->>Client: Response (unchanged)
```
## Filter Operations
The hook intercepts requests before they reach the LLM:
| Operation | Description |
|-----------|-------------|
| **Extract query** | Get user message from `messages[-1]` |
| **Convert to Routes** | Transform MCP tools into semantic-router Routes |
| **Semantic match** | Use `semantic-router` to find top-K similar tools |
| **Filter tools** | Replace request `tools` with filtered subset |
## Trigger Conditions
The filter only runs when:
- Call type is `completion` or `acompletion`
- Request contains `tools` field
- Request contains `messages` field
- Filter is enabled in config
## What Does NOT Change
- Request messages
- Response body
- Non-tool parameters
## Integration with semantic-router
Reuses existing LiteLLM infrastructure:
- `semantic-router` - Already an optional dependency
- `LiteLLMRouterEncoder` - Wraps `Router.aembedding()` for embeddings
- `SemanticRouter` - Handles similarity calculation and top-K selection
## Configuration
```yaml
litellm_settings:
mcp_semantic_tool_filter:
enabled: true
embedding_model: "openai/text-embedding-3-small"
top_k: 10
similarity_threshold: 0.3
```
## Error Handling
The filter fails gracefully:
- If filtering fails → Return all tools (no impact on functionality)
- If query extraction fails → Skip filtering
- If no matches found → Return all tools

9
litellm/proxy/hooks/mcp_semantic_filter/__init__.py Normal file
View file

@ -0,0 +1,9 @@
"""
MCP Semantic Tool Filter Hook
Semantic filtering for MCP tools to reduce context window size
and improve tool selection accuracy.
"""
from litellm.proxy.hooks.mcp_semantic_filter.hook import SemanticToolFilterHook
__all__ = ["SemanticToolFilterHook"]

353
litellm/proxy/hooks/mcp_semantic_filter/hook.py Normal file
View file

@ -0,0 +1,353 @@
"""
Semantic Tool Filter Hook
Pre-call hook that filters MCP tools semantically before LLM inference.
Reduces context window size and improves tool selection accuracy.
"""
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL,
DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD,
DEFAULT_MCP_SEMANTIC_FILTER_TOP_K,
)
from litellm.integrations.custom_logger import CustomLogger
if TYPE_CHECKING:
from litellm.caching.caching import DualCache
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
SemanticMCPToolFilter,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.router import Router
class SemanticToolFilterHook(CustomLogger):
"""
Pre-call hook that filters MCP tools semantically.
This hook:
1. Extracts the user query from messages
2. Filters tools based on semantic similarity to the query
3. Returns only the top-k most relevant tools to the LLM
"""
def __init__(self, semantic_filter: "SemanticMCPToolFilter"):
"""
Initialize the hook.
Args:
semantic_filter: SemanticMCPToolFilter instance
"""
super().__init__()
self.filter = semantic_filter
verbose_proxy_logger.debug(
f"Initialized SemanticToolFilterHook with filter: "
f"enabled={semantic_filter.enabled}, top_k={semantic_filter.top_k}"
)
def _should_expand_mcp_tools(self, tools: List[Any]) -> bool:
"""
Check if tools contain MCP references with server_url="litellm_proxy".
Only expands MCP tools pointing to litellm proxy, not external MCP servers.
"""
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)
return LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools)
async def _expand_mcp_tools(
self,
tools: List[Any],
user_api_key_dict: "UserAPIKeyAuth",
) -> List[Dict[str, Any]]:
"""
Expand MCP references to actual tool definitions.
Reuses LiteLLM_Proxy_MCP_Handler._process_mcp_tools_to_openai_format
which internally does: parse -> fetch -> filter -> deduplicate -> transform
"""
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)
# Parse to separate MCP tools from other tools
mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
if not mcp_tools:
return []
# Use single combined method instead of 3 separate calls
# This already handles: fetch -> filter by allowed_tools -> deduplicate -> transform
openai_tools, _ = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_to_openai_format(
user_api_key_auth=user_api_key_dict,
mcp_tools_with_litellm_proxy=mcp_tools
)
# Convert Pydantic models to dicts for compatibility
openai_tools_as_dicts = []
for tool in openai_tools:
if hasattr(tool, "model_dump"):
tool_dict = tool.model_dump(exclude_none=True)
verbose_proxy_logger.debug(f"Converted Pydantic tool to dict: {type(tool).__name__} -> dict with keys: {list(tool_dict.keys())}")
openai_tools_as_dicts.append(tool_dict)
elif hasattr(tool, "dict"):
tool_dict = tool.dict(exclude_none=True)
verbose_proxy_logger.debug(f"Converted Pydantic tool (v1) to dict: {type(tool).__name__} -> dict")
openai_tools_as_dicts.append(tool_dict)
elif isinstance(tool, dict):
verbose_proxy_logger.debug(f"Tool is already a dict with keys: {list(tool.keys())}")
openai_tools_as_dicts.append(tool)
else:
verbose_proxy_logger.warning(f"Tool is unknown type: {type(tool)}, passing as-is")
openai_tools_as_dicts.append(tool)
verbose_proxy_logger.debug(
f"Expanded {len(mcp_tools)} MCP reference(s) to {len(openai_tools_as_dicts)} tools (all as dicts)"
)
return openai_tools_as_dicts
def _get_metadata_variable_name(self, data: dict) -> str:
if "litellm_metadata" in data:
return "litellm_metadata"
return "metadata"
async def async_pre_call_hook(
self,
user_api_key_dict: "UserAPIKeyAuth",
cache: "DualCache",
data: dict,
call_type: str,
) -> Optional[Union[Exception, str, dict]]:
"""
Filter tools before LLM call based on user query.
This hook is called before the LLM request is made. It filters the
tools list to only include semantically relevant tools.
Args:
user_api_key_dict: User authentication
cache: Cache instance
data: Request data containing messages and tools
call_type: Type of call (completion, acompletion, etc.)
Returns:
Modified data dict with filtered tools, or None if no changes
"""
# Only filter endpoints that support tools
if call_type not in ("completion", "acompletion", "aresponses"):
verbose_proxy_logger.debug(
f"Skipping semantic filter for call_type={call_type}"
)
return None
# Check if tools are present
tools = data.get("tools")
if not tools:
verbose_proxy_logger.debug("No tools in request, skipping semantic filter")
return None
original_tool_count = len(tools)
# Check for MCP references (server_url="litellm_proxy") and expand them
if self._should_expand_mcp_tools(tools):
verbose_proxy_logger.debug(
"Detected litellm_proxy MCP references, expanding before semantic filtering"
)
try:
expanded_tools = await self._expand_mcp_tools(
tools, user_api_key_dict
)
if not expanded_tools:
verbose_proxy_logger.warning(
"No tools expanded from MCP references"
)
return None
verbose_proxy_logger.info(
f"Expanded {len(tools)} MCP reference(s) to {len(expanded_tools)} tools"
)
# Update tools for filtering
tools = expanded_tools
original_tool_count = len(tools)
except Exception as e:
verbose_proxy_logger.error(
f"Failed to expand MCP references: {e}", exc_info=True
)
return None
# Check if messages are present (try both "messages" and "input" for responses API)
messages = data.get("messages", [])
if not messages:
messages = data.get("input", [])
if not messages:
verbose_proxy_logger.debug("No messages in request, skipping semantic filter")
return None
# Check if filter is enabled
if not self.filter.enabled:
verbose_proxy_logger.debug("Semantic filter disabled, skipping")
return None
try:
# Extract user query from messages
user_query = self.filter.extract_user_query(messages)
if not user_query:
verbose_proxy_logger.debug("No user query found, skipping semantic filter")
return None
verbose_proxy_logger.debug(
f"Applying semantic filter to {len(tools)} tools "
f"with query: '{user_query[:50]}...'"
)
# Filter tools semantically
filtered_tools = await self.filter.filter_tools(
query=user_query,
available_tools=tools, # type: ignore
)
# Always update tools and emit header (even if count unchanged)
data["tools"] = filtered_tools
# Store filter stats and tool names for response header
filter_stats = f"{original_tool_count}->{len(filtered_tools)}"
tool_names_csv = self._get_tool_names_csv(filtered_tools)
_metadata_variable_name = self._get_metadata_variable_name(data)
data[_metadata_variable_name]["litellm_semantic_filter_stats"] = filter_stats
data[_metadata_variable_name]["litellm_semantic_filter_tools"] = tool_names_csv
verbose_proxy_logger.info(
f"Semantic tool filter: {filter_stats} tools"
)
return data
except Exception as e:
verbose_proxy_logger.warning(
f"Semantic tool filter hook failed: {e}. Proceeding with all tools."
)
return None
async def async_post_call_response_headers_hook(
self,
data: dict,
user_api_key_dict: "UserAPIKeyAuth",
response: Any,
request_headers: Optional[Dict[str, str]] = None,
) -> Optional[Dict[str, str]]:
"""Add semantic filter stats and tool names to response headers."""
from litellm.constants import MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH
_metadata_variable_name = self._get_metadata_variable_name(data)
metadata = data[_metadata_variable_name]
filter_stats = metadata.get("litellm_semantic_filter_stats")
if not filter_stats:
return None
headers = {"x-litellm-semantic-filter": filter_stats}
# Add CSV of filtered tool names (nginx-safe length)
tool_names_csv = metadata.get("litellm_semantic_filter_tools", "")
if tool_names_csv:
if len(tool_names_csv) > MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH:
tool_names_csv = tool_names_csv[:MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH - 3] + "..."
headers["x-litellm-semantic-filter-tools"] = tool_names_csv
return headers
def _get_tool_names_csv(self, tools: List[Any]) -> str:
"""Extract tool names and return as CSV string."""
if not tools:
return ""
tool_names = []
for tool in tools:
name = tool.get("name", "") if isinstance(tool, dict) else getattr(tool, "name", "")
if name:
tool_names.append(name)
return ",".join(tool_names)
@staticmethod
async def initialize_from_config(
config: Optional[Dict[str, Any]],
llm_router: Optional["Router"],
) -> Optional["SemanticToolFilterHook"]:
"""
Initialize semantic tool filter from proxy config.
Args:
config: Proxy configuration dict (litellm_settings.mcp_semantic_tool_filter)
llm_router: LiteLLM router instance for embeddings
Returns:
SemanticToolFilterHook instance if enabled, None otherwise
"""
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
SemanticMCPToolFilter,
)
if not config or not config.get("enabled", False):
verbose_proxy_logger.debug("Semantic tool filter not enabled in config")
return None
if llm_router is None:
verbose_proxy_logger.warning(
"Cannot initialize semantic filter: llm_router is None"
)
return None
try:
embedding_model = config.get(
"embedding_model", DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL
)
top_k = config.get("top_k", DEFAULT_MCP_SEMANTIC_FILTER_TOP_K)
similarity_threshold = config.get(
"similarity_threshold", DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD
)
semantic_filter = SemanticMCPToolFilter(
embedding_model=embedding_model,
litellm_router_instance=llm_router,
top_k=top_k,
similarity_threshold=similarity_threshold,
enabled=True,
)
# Build router from MCP registry on startup
await semantic_filter.build_router_from_mcp_registry()
hook = SemanticToolFilterHook(semantic_filter)
verbose_proxy_logger.info(
f"✅ MCP Semantic Tool Filter enabled: "
f"embedding_model={embedding_model}, top_k={top_k}, "
f"similarity_threshold={similarity_threshold}"
)
return hook
except ImportError as e:
verbose_proxy_logger.warning(
f"semantic-router not installed. Install with: "
f"pip install 'litellm[semantic-router]'. Error: {e}"
)
return None
except Exception as e:
verbose_proxy_logger.exception(
f"Failed to initialize MCP semantic tool filter: {e}"
)
return None

31
litellm/proxy/litellm_pre_call_utils.py
View file

@ -1021,37 +1021,6 @@ async def add_litellm_data_to_request( # noqa: PLR0915
"user_api_key_user_max_budget"
] = user_api_key_dict.user_max_budget
# Extract allowed access groups for router filtering (GitHub issue #18333)
# This allows the router to filter deployments based on key's and team's access groups
# NOTE: We keep key and team access groups SEPARATE because a key doesn't always
# inherit all team access groups (per maintainer feedback).
if llm_router is not None:
from litellm.proxy.auth.model_checks import get_access_groups_from_models
model_access_groups = llm_router.get_model_access_groups()
# Key-level access groups (from user_api_key_dict.models)
key_models = list(user_api_key_dict.models) if user_api_key_dict.models else []
key_allowed_access_groups = get_access_groups_from_models(
model_access_groups=model_access_groups, models=key_models
)
if key_allowed_access_groups:
data[_metadata_variable_name][
"user_api_key_allowed_access_groups"
] = key_allowed_access_groups
# Team-level access groups (from user_api_key_dict.team_models)
team_models = (
list(user_api_key_dict.team_models) if user_api_key_dict.team_models else []
)
team_allowed_access_groups = get_access_groups_from_models(
model_access_groups=model_access_groups, models=team_models
)
if team_allowed_access_groups:
data[_metadata_variable_name][
"user_api_key_team_allowed_access_groups"
] = team_allowed_access_groups
data[_metadata_variable_name]["user_api_key_metadata"] = user_api_key_dict.metadata
_headers = dict(request.headers)
_headers.pop(

5
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -813,9 +813,12 @@ def _update_internal_user_params(
data_json: dict, data: Union[UpdateUserRequest, UpdateUserRequestNoUserIDorEmail]
) -> dict:
non_default_values = {}
fields_set = data.fields_set() if hasattr(data, 'fields_set') else set()
for k, v in data_json.items():
if k == "max_budget":
non_default_values[k] = v
if "max_budget" in fields_set:
non_default_values[k] = v
elif (
v is not None
and v

279
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -753,12 +753,16 @@ async def new_team( # noqa: PLR0915
if data.max_budget is not None and data.max_budget < 0:
raise HTTPException(
status_code=400,
detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"}
detail={
"error": f"max_budget cannot be negative. Received: {data.max_budget}"
},
)
if data.team_member_budget is not None and data.team_member_budget < 0:
raise HTTPException(
status_code=400,
detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"}
detail={
"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"
},
)
# Check if license is over limit
@ -918,12 +922,16 @@ async def new_team( # noqa: PLR0915
complete_team_data.members_with_roles = []
complete_team_data_dict = complete_team_data.model_dump(exclude_none=True)
# Serialize router_settings to JSON (matching key creation pattern)
router_settings_value = getattr(data, "router_settings", None)
router_settings_json = safe_dumps(router_settings_value) if router_settings_value is not None else safe_dumps({})
router_settings_json = (
safe_dumps(router_settings_value)
if router_settings_value is not None
else safe_dumps({})
)
complete_team_data_dict["router_settings"] = router_settings_json
complete_team_data_dict = prisma_client.jsonify_team_object(
db_data=complete_team_data_dict
)
@ -1099,7 +1107,9 @@ async def fetch_and_validate_organization(
validate_team_org_change(
team=LiteLLM_TeamTable(**existing_team_row.model_dump()),
organization=LiteLLM_OrganizationTableWithMembers(**organization_row.model_dump()),
organization=LiteLLM_OrganizationTableWithMembers(
**organization_row.model_dump()
),
llm_router=llm_router,
)
@ -1107,7 +1117,9 @@ async def fetch_and_validate_organization(
def validate_team_org_change(
team: LiteLLM_TeamTable, organization: LiteLLM_OrganizationTableWithMembers, llm_router: Router
team: LiteLLM_TeamTable,
organization: LiteLLM_OrganizationTableWithMembers,
llm_router: Router,
) -> bool:
"""
Validate that a team can be moved to an organization.
@ -1158,7 +1170,9 @@ def validate_team_org_change(
# Check if the team's user_id is a member of the org
team_members = [m.user_id for m in team.members_with_roles]
org_members = [m.user_id for m in organization.members] if organization.members else []
org_members = (
[m.user_id for m in organization.members] if organization.members else []
)
not_in_org = [
m
for m in team_members
@ -1204,7 +1218,7 @@ def validate_team_org_change(
"/team/update", tags=["team management"], dependencies=[Depends(user_api_key_auth)]
)
@management_endpoint_wrapper
async def update_team( # noqa: PLR0915
async def update_team( # noqa: PLR0915
data: UpdateTeamRequest,
http_request: Request,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@ -1288,19 +1302,25 @@ async def update_team( # noqa: PLR0915
)
if data.team_id is None:
raise HTTPException(status_code=400, detail={"error": "No team id passed in"})
raise HTTPException(
status_code=400, detail={"error": "No team id passed in"}
)
verbose_proxy_logger.debug("/team/update - %s", data)
# Validate budget values are not negative
if data.max_budget is not None and data.max_budget < 0:
raise HTTPException(
status_code=400,
detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"}
detail={
"error": f"max_budget cannot be negative. Received: {data.max_budget}"
},
)
if data.team_member_budget is not None and data.team_member_budget < 0:
raise HTTPException(
status_code=400,
detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"}
detail={
"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"
},
)
existing_team_row = await prisma_client.db.litellm_teamtable.find_unique(
@ -1367,6 +1387,22 @@ async def update_team( # noqa: PLR0915
updated_kv = data.json(exclude_unset=True)
# Only proxy admin can change allow_team_guardrail_config
if (
"allow_team_guardrail_config" in updated_kv
and user_api_key_dict.user_role
not in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN.value,
)
):
raise HTTPException(
status_code=403,
detail={
"error": "Only proxy admin can change 'Allow team to configure guardrails'. Contact your administrator."
},
)
# Check budget_duration and budget_reset_at
_set_budget_reset_at(data, updated_kv)
@ -1411,16 +1447,19 @@ async def update_team( # noqa: PLR0915
updated_kv["model_id"] = _model_id
# Serialize router_settings to JSON if present (matching key update pattern)
if "router_settings" in updated_kv and updated_kv["router_settings"] is not None:
if (
"router_settings" in updated_kv
and updated_kv["router_settings"] is not None
):
updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"])
updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv)
team_row: Optional[LiteLLM_TeamTable] = (
await prisma_client.db.litellm_teamtable.update(
where={"team_id": data.team_id},
data=updated_kv,
include={"litellm_model_table": True}, # type: ignore
)
team_row: Optional[
LiteLLM_TeamTable
] = await prisma_client.db.litellm_teamtable.update(
where={"team_id": data.team_id},
data=updated_kv,
include={"litellm_model_table": True}, # type: ignore
)
if team_row is None or team_row.team_id is None:
@ -1429,7 +1468,9 @@ async def update_team( # noqa: PLR0915
detail={"error": "Team doesn't exist. Got={}".format(team_row)},
)
verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id)
verbose_proxy_logger.info(
"Successfully updated team - %s, info", team_row.team_id
)
await _cache_team_object(
team_id=team_row.team_id,
team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()),
@ -1771,113 +1812,6 @@ async def _add_team_members_to_team(
return updated_team, updated_users, updated_team_memberships
async def _validate_and_populate_member_user_info(
member: Member,
prisma_client: PrismaClient,
) -> Member:
"""
Validate and populate user_email/user_id for a member.
Logic:
1. If both user_email and user_id are provided, verify they belong to the same user (use user_email as source of truth)
2. If only user_email is provided, populate user_id from DB
3. If only user_id is provided, populate user_email from DB (if user exists)
4. If only user_id is provided and doesn't exist, allow it to pass with user_email as None (will be upserted later)
5. If user_email and user_id mismatch, throw error
Returns a Member with user_email and user_id populated (user_email may be None if only user_id provided and user doesn't exist).
"""
if member.user_email is None and member.user_id is None:
raise HTTPException(
status_code=400,
detail={"error": "Either user_id or user_email must be provided"},
)
# Case 1: Both user_email and user_id provided - verify they match
if member.user_email is not None and member.user_id is not None:
# Use user_email as source of truth
# Check for multiple users with same email first
users_by_email = await prisma_client.get_data(
key_val={"user_email": member.user_email},
table_name="user",
query_type="find_all",
)
if users_by_email is None or (
isinstance(users_by_email, list) and len(users_by_email) == 0
):
# User doesn't exist yet - this is fine, will be created later
return member
if isinstance(users_by_email, list) and len(users_by_email) > 1:
raise HTTPException(
status_code=400,
detail={
"error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead."
},
)
# Get the single user
user_by_email = users_by_email[0]
# Verify the user_id matches
if user_by_email.user_id != member.user_id:
raise HTTPException(
status_code=400,
detail={
"error": f"user_email '{member.user_email}' and user_id '{member.user_id}' do not belong to the same user."
},
)
# Both match, return as is
return member
# Case 2: Only user_email provided - populate user_id from DB
if member.user_email is not None and member.user_id is None:
user_by_email = await prisma_client.db.litellm_usertable.find_first(
where={"user_email": {"equals": member.user_email, "mode": "insensitive"}}
)
if user_by_email is None:
# User doesn't exist yet - this is fine, will be created later
return member
# Check for multiple users with same email
users_by_email = await prisma_client.get_data(
key_val={"user_email": member.user_email},
table_name="user",
query_type="find_all",
)
if users_by_email and isinstance(users_by_email, list) and len(users_by_email) > 1:
raise HTTPException(
status_code=400,
detail={
"error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead."
},
)
# Populate user_id
member.user_id = user_by_email.user_id
return member
# Case 3: Only user_id provided - populate user_email from DB if user exists
if member.user_id is not None and member.user_email is None:
user_by_id = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": member.user_id}
)
if user_by_id is None:
# User doesn't exist yet - allow it to pass with user_email as None
# Will be upserted later with just user_id and null email
return member
# Populate user_email
member.user_email = user_by_id.user_email
return member
return member
@router.post(
"/team/member_add",
tags=["team management"],
@ -1953,27 +1887,16 @@ async def team_member_add(
complete_team_data=complete_team_data,
)
# Validate and populate user_email/user_id for members before processing
if isinstance(data.member, Member):
await _validate_and_populate_member_user_info(
member=data.member,
prisma_client=prisma_client,
)
elif isinstance(data.member, List):
for m in data.member:
await _validate_and_populate_member_user_info(
member=m,
prisma_client=prisma_client,
)
updated_team, updated_users, updated_team_memberships = (
await _add_team_members_to_team(
data=data,
complete_team_data=complete_team_data,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
(
updated_team,
updated_users,
updated_team_memberships,
) = await _add_team_members_to_team(
data=data,
complete_team_data=complete_team_data,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
# Check if updated_team is None
@ -2152,15 +2075,15 @@ async def team_member_delete(
)
# Fetch keys before deletion to persist them
keys_to_delete: List[LiteLLM_VerificationToken] = (
await prisma_client.db.litellm_verificationtoken.find_many(
where={
"user_id": {"in": list(user_ids_to_delete)},
"team_id": data.team_id,
}
)
keys_to_delete: List[
LiteLLM_VerificationToken
] = await prisma_client.db.litellm_verificationtoken.find_many(
where={
"user_id": {"in": list(user_ids_to_delete)},
"team_id": data.team_id,
}
)
if keys_to_delete:
await _persist_deleted_verification_tokens(
keys=keys_to_delete,
@ -2539,10 +2462,10 @@ async def delete_team(
team_rows: List[LiteLLM_TeamTable] = []
for team_id in data.team_ids:
try:
team_row_base: Optional[BaseModel] = (
await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
team_row_base: Optional[
BaseModel
] = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
if team_row_base is None:
raise Exception
@ -2601,10 +2524,10 @@ async def delete_team(
_persist_deleted_verification_tokens,
)
keys_to_delete: List[LiteLLM_VerificationToken] = (
await prisma_client.db.litellm_verificationtoken.find_many(
where={"team_id": {"in": data.team_ids}}
)
keys_to_delete: List[
LiteLLM_VerificationToken
] = await prisma_client.db.litellm_verificationtoken.find_many(
where={"team_id": {"in": data.team_ids}}
)
if keys_to_delete:
@ -2643,7 +2566,6 @@ async def delete_team(
return deleted_teams
def _transform_teams_to_deleted_records(
teams: List[LiteLLM_TeamTable],
user_api_key_dict: UserAPIKeyAuth,
@ -2666,7 +2588,13 @@ def _transform_teams_to_deleted_records(
)
record = deleted_record.model_dump()
for json_field in ["members_with_roles", "metadata", "model_spend", "model_max_budget", "router_settings"]:
for json_field in [
"members_with_roles",
"metadata",
"model_spend",
"model_max_budget",
"router_settings",
]:
if json_field in record and record[json_field] is not None:
record[json_field] = json.dumps(record[json_field])
@ -2685,9 +2613,7 @@ async def _save_deleted_team_records(
"""Save deleted team records to the database."""
if not records:
return
await prisma_client.db.litellm_deletedteamtable.create_many(
data=records
)
await prisma_client.db.litellm_deletedteamtable.create_many(data=records)
async def _persist_deleted_team_records(
@ -2707,6 +2633,7 @@ async def _persist_deleted_team_records(
prisma_client=prisma_client,
)
def validate_membership(
user_api_key_dict: UserAPIKeyAuth, team_table: LiteLLM_TeamTable
):
@ -2830,11 +2757,11 @@ async def team_info(
)
try:
team_info: Optional[BaseModel] = (
await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id},
include={"object_permission": True},
)
team_info: Optional[
BaseModel
] = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id},
include={"object_permission": True},
)
if team_info is None:
raise Exception
@ -3297,7 +3224,9 @@ async def list_team_v2(
order=order_by if order_by else {"created_at": "desc"}, # Default sort
)
# Get total count for pagination
total_count = await prisma_client.db.litellm_teamtable.count(where=where_conditions)
total_count = await prisma_client.db.litellm_teamtable.count(
where=where_conditions
)
# Calculate total pages
total_pages = -(-total_count // page_size) # Ceiling division

39
litellm/proxy/proxy_config.yaml
View file

@ -1,4 +1,14 @@
model_list:
- model_name: gpt-4o
litellm_params:
model: openai/gpt-4o
api_key: os.environ/OPENAI_API_KEY
- model_name: text-embedding-3-small
litellm_params:
model: openai/text-embedding-3-small
api_key: os.environ/OPENAI_API_KEY
- model_name: bedrock-claude-sonnet-3.5
litellm_params:
model: "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0"
@ -22,4 +32,31 @@ model_list:
- model_name: bedrock-nova-premier
litellm_params:
model: "bedrock/us.amazon.nova-premier-v1:0"
aws_region_name: "us-east-1"
aws_region_name: "us-east-1"
# MCP Server Configuration
mcp_servers:
# Wikipedia MCP - reliable and works without external deps
wikipedia:
transport: "stdio"
command: "uvx"
args: ["mcp-server-fetch"]
description: "Fetch web pages and Wikipedia content"
deepwiki:
transport: "http"
url: "https://mcp.deepwiki.com/mcp"
# General Settings
general_settings:
master_key: sk-1234
store_model_in_db: false
# LiteLLM Settings
litellm_settings:
# Enable MCP Semantic Tool Filter
mcp_semantic_tool_filter:
enabled: true
embedding_model: "text-embedding-3-small"
top_k: 5
similarity_threshold: 0.3

62
litellm/proxy/proxy_server.py
View file

@ -239,6 +239,10 @@ from litellm.proxy._types import *
from litellm.proxy.agent_endpoints.a2a_endpoints import router as a2a_router
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
from litellm.proxy.agent_endpoints.endpoints import router as agent_endpoints_router
from litellm.proxy.agent_endpoints.model_list_helpers import (
append_agents_to_model_group,
append_agents_to_model_info,
)
from litellm.proxy.analytics_endpoints.analytics_endpoints import (
router as analytics_router,
)
@ -793,6 +797,21 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915
redis_usage_cache=redis_usage_cache,
)
## SEMANTIC TOOL FILTER ##
# Read litellm_settings from config for semantic filter initialization
try:
verbose_proxy_logger.debug("About to initialize semantic tool filter")
_config = proxy_config.get_config_state()
_litellm_settings = _config.get("litellm_settings", {})
verbose_proxy_logger.debug(f"litellm_settings keys = {list(_litellm_settings.keys())}")
await ProxyStartupEvent._initialize_semantic_tool_filter(
llm_router=llm_router,
litellm_settings=_litellm_settings,
)
verbose_proxy_logger.debug("After semantic tool filter initialization")
except Exception as e:
verbose_proxy_logger.error(f"Semantic filter init failed: {e}", exc_info=True)
## JWT AUTH ##
ProxyStartupEvent._initialize_jwt_auth(
general_settings=general_settings,
@ -4742,6 +4761,34 @@ class ProxyStartupEvent:
llm_router=llm_router, redis_usage_cache=redis_usage_cache
)
@classmethod
async def _initialize_semantic_tool_filter(
cls,
llm_router: Optional[Router],
litellm_settings: Dict[str, Any],
):
"""Initialize MCP semantic tool filter if configured"""
from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook
verbose_proxy_logger.info(
f"Initializing semantic tool filter: llm_router={llm_router is not None}, "
f"litellm_settings keys={list(litellm_settings.keys())}"
)
mcp_semantic_filter_config = litellm_settings.get("mcp_semantic_tool_filter", None)
verbose_proxy_logger.debug(f"Semantic filter config: {mcp_semantic_filter_config}")
hook = await SemanticToolFilterHook.initialize_from_config(
config=mcp_semantic_filter_config,
llm_router=llm_router,
)
if hook:
verbose_proxy_logger.debug("✅ Semantic tool filter hook registered")
litellm.logging_callback_manager.add_litellm_callback(hook)
else:
verbose_proxy_logger.warning("❌ Semantic tool filter hook not initialized")
@classmethod
def _initialize_jwt_auth(
cls,
@ -8573,6 +8620,15 @@ async def model_info_v2(
)
verbose_proxy_logger.debug("all_models: %s", all_models)
# Append A2A agents to models list
all_models = await append_agents_to_model_info(
models=all_models,
user_api_key_dict=user_api_key_dict,
)
# Update total count to include agents
search_total_count = len(all_models)
return _paginate_models_response(
all_models=all_models,
@ -9413,6 +9469,12 @@ async def model_group_info(
model_groups: List[ModelGroupInfoProxy] = _get_model_group_info(
llm_router=llm_router, all_models_str=all_models_str, model_group=model_group
)
# Append A2A agents to model groups
model_groups = await append_agents_to_model_group(
model_groups=model_groups,
user_api_key_dict=user_api_key_dict,
)
return {"data": model_groups}

11
litellm/proxy/route_llm_request.py
View file

@ -12,6 +12,11 @@ else:
LitellmRouter = Any
def _is_a2a_agent_model(model_name: Any) -> bool:
"""Check if the model name is for an A2A agent (a2a/ prefix)."""
return isinstance(model_name, str) and model_name.startswith("a2a/")
ROUTE_ENDPOINT_MAPPING = {
"acompletion": "/chat/completions",
"atext_completion": "/completions",
@ -322,6 +327,12 @@ async def route_request(
except Exception:
# If router fails (e.g., model not found in router), fall back to direct call
return getattr(litellm, f"{route_type}")(**data)
elif _is_a2a_agent_model(data.get("model", "")):
from litellm.proxy.agent_endpoints.a2a_routing import (
route_a2a_agent_request,
)
return await route_a2a_agent_request(data, route_type)
elif user_model is not None:
return getattr(litellm, f"{route_type}")(**data)

15
litellm/proxy/schema.prisma
View file

@ -128,6 +128,7 @@ model LiteLLM_TeamTable {
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
policies String[] @default([])
allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team
model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id])
@ -159,8 +160,9 @@ model LiteLLM_DeletedTeamTable {
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
policies String[] @default([])
allow_team_guardrail_config Boolean @default(false)
model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases
// Original timestamps from team creation/updates
created_at DateTime? @map("created_at")
updated_at DateTime? @map("updated_at")
@ -305,16 +307,6 @@ model LiteLLM_VerificationToken {
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
// SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub"
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2
@@index([user_id, team_id])
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."team_id" = $1 OFFSET $2
@@index([team_id])
// 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])
}
// Audit table for deleted keys - preserves spend and key information for historical tracking
@ -784,6 +776,7 @@ model LiteLLM_GuardrailsTable {
guardrail_name String @unique
litellm_params Json
guardrail_info Json?
team_id String?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
}

10
litellm/proxy/utils.py
View file

@ -1633,8 +1633,16 @@ class ProxyLogging:
if RouteChecks.is_llm_api_route(route) is not True:
return False
proxy_only_error_types = {
ProxyErrorTypes.auth_error,
ProxyErrorTypes.key_model_access_denied,
ProxyErrorTypes.team_model_access_denied,
ProxyErrorTypes.user_model_access_denied,
ProxyErrorTypes.org_model_access_denied,
ProxyErrorTypes.token_not_found_in_db,
}
return isinstance(original_exception, HTTPException) or (
error_type == ProxyErrorTypes.auth_error
error_type in proxy_only_error_types
)
async def _handle_logging_proxy_only_error(

30
litellm/proxy_auth/__init__.py Normal file
View file

@ -0,0 +1,30 @@
"""
Proxy Authentication module for LiteLLM SDK.
This module provides OAuth2/JWT token management for authenticating
with LiteLLM Proxy or any OAuth2-protected endpoint.
Usage:
from litellm.proxy_auth import AzureADCredential, ProxyAuthHandler
litellm.proxy_auth = ProxyAuthHandler(
credential=AzureADCredential(),
scope="api://my-proxy/.default"
)
"""
from .credentials import (
AccessToken,
TokenCredential,
AzureADCredential,
GenericOAuth2Credential,
ProxyAuthHandler,
)
__all__ = [
"AccessToken",
"TokenCredential",
"AzureADCredential",
"GenericOAuth2Credential",
"ProxyAuthHandler",
]

240
litellm/proxy_auth/credentials.py Normal file
View file

@ -0,0 +1,240 @@
"""
Credential providers for proxy authentication.
This module provides a provider-agnostic interface for obtaining OAuth2/JWT tokens.
It follows the same TokenCredential protocol used by Azure SDK.
"""
import time
from dataclasses import dataclass
from typing import Any, Optional, Protocol, runtime_checkable
@dataclass
class AccessToken:
"""
Represents an OAuth2 access token with expiration.
This matches the structure used by azure.core.credentials.AccessToken.
Attributes:
token: The access token string (typically a JWT).
expires_on: Unix timestamp when the token expires.
"""
token: str
expires_on: int
@runtime_checkable
class TokenCredential(Protocol):
"""
Protocol for credential providers.
This matches the azure.core.credentials.TokenCredential interface,
allowing any Azure SDK credential to be used directly.
Any class implementing get_token(scope) -> AccessToken can be used.
"""
def get_token(self, scope: str) -> AccessToken:
"""
Get an access token for the specified scope.
Args:
scope: The OAuth2 scope to request (e.g., "api://my-app/.default")
Returns:
AccessToken with the token string and expiration timestamp.
"""
...
class AzureADCredential:
"""
Wrapper for Azure Identity credentials.
This wraps any azure-identity credential (DefaultAzureCredential,
ClientSecretCredential, ManagedIdentityCredential, etc.) and converts
the token to our AccessToken format.
If no credential is provided, it will use DefaultAzureCredential
which tries multiple authentication methods automatically.
Example:
# Use default credential chain (env vars, managed identity, CLI, etc.)
cred = AzureADCredential()
# Or provide a specific credential
from azure.identity import ClientSecretCredential
azure_cred = ClientSecretCredential(tenant_id, client_id, client_secret)
cred = AzureADCredential(credential=azure_cred)
"""
def __init__(self, credential: Optional[Any] = None):
"""
Initialize with an optional Azure credential.
Args:
credential: An azure-identity credential object. If None,
DefaultAzureCredential will be used on first token request.
"""
self._credential = credential
self._initialized = credential is not None
def get_token(self, scope: str) -> AccessToken:
"""
Get an access token from Azure AD.
Args:
scope: The OAuth2 scope (e.g., "api://my-app/.default")
Returns:
AccessToken with the JWT and expiration.
Raises:
ImportError: If azure-identity is not installed.
"""
if not self._initialized:
try:
from azure.identity import DefaultAzureCredential
self._credential = DefaultAzureCredential()
self._initialized = True
except ImportError:
raise ImportError(
"azure-identity is required for AzureADCredential. "
"Install it with: pip install azure-identity"
)
result = self._credential.get_token(scope)
return AccessToken(token=result.token, expires_on=result.expires_on)
class GenericOAuth2Credential:
"""
Generic OAuth2 client credentials flow.
This works with any OAuth2 provider (Okta, Auth0, Keycloak, etc.)
that supports the client_credentials grant type.
Example:
cred = GenericOAuth2Credential(
client_id="my-client-id",
client_secret="my-client-secret",
token_url="https://my-idp.com/oauth2/token"
)
"""
def __init__(self, client_id: str, client_secret: str, token_url: str):
"""
Initialize OAuth2 client credentials.
Args:
client_id: OAuth2 client ID
client_secret: OAuth2 client secret
token_url: Token endpoint URL (e.g., "https://idp.com/oauth2/token")
"""
self.client_id = client_id
self.client_secret = client_secret
self.token_url = token_url
self._cached_token: Optional[AccessToken] = None
def get_token(self, scope: str) -> AccessToken:
"""
Get an access token using OAuth2 client credentials flow.
Tokens are cached and reused until they expire (with 60s buffer).
Args:
scope: The OAuth2 scope to request
Returns:
AccessToken with the token and expiration.
"""
# Return cached token if still valid (with 60s buffer)
if self._cached_token and self._cached_token.expires_on > time.time() + 60:
return self._cached_token
import httpx
response = httpx.post(
self.token_url,
data={
"grant_type": "client_credentials",
"client_id": self.client_id,
"client_secret": self.client_secret,
"scope": scope,
},
)
response.raise_for_status()
data = response.json()
self._cached_token = AccessToken(
token=data["access_token"],
expires_on=int(time.time()) + data.get("expires_in", 3600),
)
return self._cached_token
class ProxyAuthHandler:
"""
Manages OAuth2/JWT token lifecycle for proxy authentication.
This handler:
- Obtains tokens from the configured credential provider
- Caches tokens to avoid unnecessary requests
- Automatically refreshes tokens before they expire (60s buffer)
- Generates Authorization headers for HTTP requests
Set this as litellm.proxy_auth to automatically inject auth headers
into all requests to your LiteLLM Proxy.
Example:
import litellm
from litellm.proxy_auth import AzureADCredential, ProxyAuthHandler
litellm.proxy_auth = ProxyAuthHandler(
credential=AzureADCredential(),
scope="api://my-litellm-proxy/.default"
)
litellm.api_base = "https://my-proxy.example.com"
# Auth headers are now automatically injected
response = litellm.completion(model="gpt-4", messages=[...])
"""
def __init__(self, credential: TokenCredential, scope: str):
"""
Initialize the proxy auth handler.
Args:
credential: A TokenCredential implementation (AzureADCredential,
GenericOAuth2Credential, or any custom implementation)
scope: The OAuth2 scope to request tokens for
"""
self.credential = credential
self.scope = scope
self._cached_token: Optional[AccessToken] = None
def get_token(self) -> AccessToken:
"""
Get a valid access token, refreshing if necessary.
Returns:
AccessToken that is valid for at least 60 more seconds.
"""
# Refresh if no token or token expires within 60 seconds
if not self._cached_token or self._cached_token.expires_on <= time.time() + 60:
self._cached_token = self.credential.get_token(self.scope)
return self._cached_token
def get_auth_headers(self) -> dict:
"""
Get HTTP headers for authentication.
Returns:
Dict with Authorization header containing Bearer token.
"""
token = self.get_token()
return {"Authorization": f"Bearer {token.token}"}

12
litellm/router.py
View file

@ -88,7 +88,6 @@ from litellm.router_utils.clientside_credential_handler import (
is_clientside_credential,
)
from litellm.router_utils.common_utils import (
filter_deployments_by_access_groups,
filter_team_based_models,
filter_web_search_deployments,
)
@ -8088,17 +8087,10 @@ class Router:
request_kwargs=request_kwargs,
)
verbose_router_logger.debug(f"healthy_deployments after web search filter: {healthy_deployments}")
# Filter by allowed access groups (GitHub issue #18333)
# This prevents cross-team load balancing when teams have models with same name in different access groups
healthy_deployments = filter_deployments_by_access_groups(
healthy_deployments=healthy_deployments,
request_kwargs=request_kwargs,
verbose_router_logger.debug(
f"healthy_deployments after web search filter: {healthy_deployments}"
)
verbose_router_logger.debug(f"healthy_deployments after access group filter: {healthy_deployments}")
if isinstance(healthy_deployments, dict):
return healthy_deployments

79
litellm/router_utils/common_utils.py
View file

@ -75,7 +75,6 @@ def filter_team_based_models(
if deployment.get("model_info", {}).get("id") not in ids_to_remove
]
def _deployment_supports_web_search(deployment: Dict) -> bool:
"""
Check if a deployment supports web search.
@ -113,7 +112,7 @@ def filter_web_search_deployments(
is_web_search_request = False
tools = request_kwargs.get("tools") or []
for tool in tools:
# These are the two websearch tools for OpenAI / Azure.
# These are the two websearch tools for OpenAI / Azure.
if tool.get("type") == "web_search" or tool.get("type") == "web_search_preview":
is_web_search_request = True
break
@ -122,82 +121,8 @@ def filter_web_search_deployments(
return healthy_deployments
# Filter out deployments that don't support web search
final_deployments = [
d for d in healthy_deployments if _deployment_supports_web_search(d)
]
final_deployments = [d for d in healthy_deployments if _deployment_supports_web_search(d)]
if len(healthy_deployments) > 0 and len(final_deployments) == 0:
verbose_logger.warning("No deployments support web search for request")
return final_deployments
def filter_deployments_by_access_groups(
healthy_deployments: Union[List[Dict], Dict],
request_kwargs: Optional[Dict] = None,
) -> Union[List[Dict], Dict]:
"""
Filter deployments to only include those matching the user's allowed access groups.
Reads from TWO separate metadata fields (per maintainer feedback):
- `user_api_key_allowed_access_groups`: Access groups from the API Key's models.
- `user_api_key_team_allowed_access_groups`: Access groups from the Team's models.
A deployment is included if its access_groups overlap with EITHER the key's
or the team's allowed access groups. Deployments with no access_groups are
always included (not restricted).
This prevents cross-team load balancing when multiple teams have models with
the same name but in different access groups (GitHub issue #18333).
"""
if request_kwargs is None:
return healthy_deployments
if isinstance(healthy_deployments, dict):
return healthy_deployments
metadata = request_kwargs.get("metadata") or {}
litellm_metadata = request_kwargs.get("litellm_metadata") or {}
# Gather key-level allowed access groups
key_allowed_access_groups = (
metadata.get("user_api_key_allowed_access_groups")
or litellm_metadata.get("user_api_key_allowed_access_groups")
or []
)
# Gather team-level allowed access groups
team_allowed_access_groups = (
metadata.get("user_api_key_team_allowed_access_groups")
or litellm_metadata.get("user_api_key_team_allowed_access_groups")
or []
)
# Combine both for the final allowed set
combined_allowed_access_groups = list(key_allowed_access_groups) + list(
team_allowed_access_groups
)
# If no access groups specified from either source, return all deployments (backwards compatible)
if not combined_allowed_access_groups:
return healthy_deployments
allowed_set = set(combined_allowed_access_groups)
filtered = []
for deployment in healthy_deployments:
model_info = deployment.get("model_info") or {}
deployment_access_groups = model_info.get("access_groups") or []
# If deployment has no access groups, include it (not restricted)
if not deployment_access_groups:
filtered.append(deployment)
continue
# Include if any of deployment's groups overlap with allowed groups
if set(deployment_access_groups) & allowed_set:
filtered.append(deployment)
if len(healthy_deployments) > 0 and len(filtered) == 0:
verbose_logger.warning(
f"No deployments match allowed access groups {combined_allowed_access_groups}"
)
return filtered

12
litellm/router_utils/fallback_event_handlers.py
View file

@ -113,16 +113,8 @@ async def run_async_fallback(
The most recent exception if all fallback model groups fail.
"""
### BASE CASE ### MAX FALLBACK DEPTH REACHED
if fallback_depth >= max_fallbacks:
raise original_exception
### CHECK IF MODEL GROUP LIST EXHAUSTED
if original_model_group in fallback_model_group:
fallback_group_length = len(fallback_model_group) - 1
else:
fallback_group_length = len(fallback_model_group)
if fallback_depth >= fallback_group_length:
### BASE CASE ### MAX FALLBACK DEPTH REACHED
if fallback_depth >= max_fallbacks:
raise original_exception
error_from_fallbacks = original_exception

10
litellm/types/guardrails.py
View file

@ -731,6 +731,7 @@ class Guardrail(TypedDict, total=False):
guardrail_name: Required[str]
litellm_params: Required[LitellmParams]
guardrail_info: Optional[Dict]
team_id: Optional[str]
created_at: Optional[datetime]
updated_at: Optional[datetime]
@ -762,6 +763,7 @@ class GuardrailInfoResponse(BaseModel):
guardrail_name: str
litellm_params: Optional[BaseLitellmParams] = None
guardrail_info: Optional[Dict] = None
team_id: Optional[str] = None
created_at: Optional[datetime] = None
updated_at: Optional[datetime] = None
guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = (
@ -808,3 +810,11 @@ class PatchGuardrailRequest(BaseModel):
guardrail_name: Optional[str] = None
litellm_params: Optional[BaseLitellmParams] = None
guardrail_info: Optional[Dict[str, Any]] = None
class CreateGuardrailRequest(BaseModel):
guardrail: Guardrail
class UpdateGuardrailRequest(BaseModel):
guardrail: Guardrail

1
litellm/types/utils.py
View file

@ -3029,6 +3029,7 @@ class LlmProviders(str, Enum):
MISTRAL = "mistral"
MILVUS = "milvus"
GROQ = "groq"
A2A = "a2a"
GIGACHAT = "gigachat"
NVIDIA_NIM = "nvidia_nim"
CEREBRAS = "cerebras"

8
litellm/utils.py
View file

@ -1453,6 +1453,10 @@ def client(original_function): # noqa: PLR0915
logging_obj, kwargs = function_setup(
original_function.__name__, rules_obj, start_time, *args, **kwargs
)
# Type assertion: logging_obj is guaranteed to be non-None after function_setup
assert logging_obj is not None, "logging_obj should not be None after function_setup"
## LOAD CREDENTIALS
load_credentials_from_list(kwargs)
kwargs["litellm_logging_obj"] = logging_obj
@ -1771,6 +1775,9 @@ def client(original_function): # noqa: PLR0915
logging_obj, kwargs = function_setup(
original_function.__name__, rules_obj, start_time, *args, **kwargs
)
# Type assertion: logging_obj is guaranteed to be non-None after function_setup
assert logging_obj is not None, "logging_obj should not be None after function_setup"
modified_kwargs = await async_pre_call_deployment_hook(kwargs, call_type)
if modified_kwargs is not None:
@ -7799,6 +7806,7 @@ class ProviderConfigManager:
# Simple provider mappings (no model parameter needed)
LlmProviders.DEEPSEEK: (lambda: litellm.DeepSeekChatConfig(), False),
LlmProviders.GROQ: (lambda: litellm.GroqChatConfig(), False),
LlmProviders.A2A: (lambda: litellm.A2AConfig(), False),
LlmProviders.BYTEZ: (lambda: litellm.BytezChatConfig(), False),
LlmProviders.DATABRICKS: (lambda: litellm.DatabricksConfig(), False),
LlmProviders.XAI: (lambda: litellm.XAIChatConfig(), False),

361
model_prices_and_context_window.json
View file

@ -27113,6 +27113,34 @@
"supports_reasoning": true,
"supports_tool_choice": true
},
"together_ai/zai-org/GLM-4.7": {
"input_cost_per_token": 4.5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 200000,
"max_output_tokens": 200000,
"max_tokens": 200000,
"mode": "chat",
"output_cost_per_token": 2e-06,
"source": "https://www.together.ai/models/glm-4-7",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"together_ai/moonshotai/Kimi-K2.5": {
"input_cost_per_token": 5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 256000,
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"output_cost_per_token": 2.8e-06,
"source": "https://www.together.ai/models/kimi-k2-5",
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_reasoning": true
},
"together_ai/moonshotai/Kimi-K2-Instruct-0905": {
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",
@ -27829,7 +27857,9 @@
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 3e-07
"output_cost_per_token": 3e-07,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/alibaba/qwen3-coder": {
"input_cost_per_token": 4e-07,
@ -27838,7 +27868,9 @@
"max_output_tokens": 66536,
"max_tokens": 66536,
"mode": "chat",
"output_cost_per_token": 1.6e-06
"output_cost_per_token": 1.6e-06,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/amazon/nova-lite": {
"input_cost_per_token": 6e-08,
@ -27847,7 +27879,10 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.4e-07
"output_cost_per_token": 2.4e-07,
"supports_vision": true,
"supports_function_calling": true,
"supports_response_schema": true
},
"vercel_ai_gateway/amazon/nova-micro": {
"input_cost_per_token": 3.5e-08,
@ -27856,7 +27891,9 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 1.4e-07
"output_cost_per_token": 1.4e-07,
"supports_function_calling": true,
"supports_response_schema": true
},
"vercel_ai_gateway/amazon/nova-pro": {
"input_cost_per_token": 8e-07,
@ -27865,7 +27902,10 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 3.2e-06
"output_cost_per_token": 3.2e-06,
"supports_vision": true,
"supports_function_calling": true,
"supports_response_schema": true
},
"vercel_ai_gateway/amazon/titan-embed-text-v2": {
"input_cost_per_token": 2e-08,
@ -27885,7 +27925,11 @@
"max_output_tokens": 4096,
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 1.25e-06
"output_cost_per_token": 1.25e-06,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/anthropic/claude-3-opus": {
"cache_creation_input_token_cost": 1.875e-05,
@ -27896,7 +27940,11 @@
"max_output_tokens": 4096,
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 7.5e-05
"output_cost_per_token": 7.5e-05,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/anthropic/claude-3.5-haiku": {
"cache_creation_input_token_cost": 1e-06,
@ -27907,7 +27955,11 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 4e-06
"output_cost_per_token": 4e-06,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/anthropic/claude-3.5-sonnet": {
"cache_creation_input_token_cost": 3.75e-06,
@ -27918,7 +27970,11 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 1.5e-05
"output_cost_per_token": 1.5e-05,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/anthropic/claude-3.7-sonnet": {
"cache_creation_input_token_cost": 3.75e-06,
@ -27929,7 +27985,11 @@
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 1.5e-05
"output_cost_per_token": 1.5e-05,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/anthropic/claude-4-opus": {
"cache_creation_input_token_cost": 1.875e-05,
@ -27940,7 +28000,11 @@
"max_output_tokens": 32000,
"max_tokens": 32000,
"mode": "chat",
"output_cost_per_token": 7.5e-05
"output_cost_per_token": 7.5e-05,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/anthropic/claude-4-sonnet": {
"cache_creation_input_token_cost": 3.75e-06,
@ -27951,7 +28015,9 @@
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 1.5e-05
"output_cost_per_token": 1.5e-05,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/cohere/command-a": {
"input_cost_per_token": 2.5e-06,
@ -27960,7 +28026,9 @@
"max_output_tokens": 8000,
"max_tokens": 8000,
"mode": "chat",
"output_cost_per_token": 1e-05
"output_cost_per_token": 1e-05,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/cohere/command-r": {
"input_cost_per_token": 1.5e-07,
@ -27969,7 +28037,9 @@
"max_output_tokens": 4096,
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 6e-07
"output_cost_per_token": 6e-07,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/cohere/command-r-plus": {
"input_cost_per_token": 2.5e-06,
@ -27978,7 +28048,9 @@
"max_output_tokens": 4096,
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 1e-05
"output_cost_per_token": 1e-05,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/cohere/embed-v4.0": {
"input_cost_per_token": 1.2e-07,
@ -27996,7 +28068,8 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2.19e-06
"output_cost_per_token": 2.19e-06,
"supports_tool_choice": true
},
"vercel_ai_gateway/deepseek/deepseek-r1-distill-llama-70b": {
"input_cost_per_token": 7.5e-07,
@ -28005,7 +28078,10 @@
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 9.9e-07
"output_cost_per_token": 9.9e-07,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/deepseek/deepseek-v3": {
"input_cost_per_token": 9e-07,
@ -28014,7 +28090,8 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 9e-07
"output_cost_per_token": 9e-07,
"supports_tool_choice": true
},
"vercel_ai_gateway/google/gemini-2.0-flash": {
"deprecation_date": "2026-03-31",
@ -28024,7 +28101,11 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 6e-07
"output_cost_per_token": 6e-07,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/google/gemini-2.0-flash-lite": {
"deprecation_date": "2026-03-31",
@ -28034,7 +28115,11 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 3e-07
"output_cost_per_token": 3e-07,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/google/gemini-2.5-flash": {
"input_cost_per_token": 3e-07,
@ -28043,7 +28128,11 @@
"max_output_tokens": 65536,
"max_tokens": 65536,
"mode": "chat",
"output_cost_per_token": 2.5e-06
"output_cost_per_token": 2.5e-06,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/google/gemini-2.5-pro": {
"input_cost_per_token": 2.5e-06,
@ -28052,7 +28141,11 @@
"max_output_tokens": 65536,
"max_tokens": 65536,
"mode": "chat",
"output_cost_per_token": 1e-05
"output_cost_per_token": 1e-05,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/google/gemini-embedding-001": {
"input_cost_per_token": 1.5e-07,
@ -28070,7 +28163,10 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 2e-07
"output_cost_per_token": 2e-07,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/google/text-embedding-005": {
"input_cost_per_token": 2.5e-08,
@ -28106,7 +28202,8 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 7.9e-07
"output_cost_per_token": 7.9e-07,
"supports_tool_choice": true
},
"vercel_ai_gateway/meta/llama-3-8b": {
"input_cost_per_token": 5e-08,
@ -28115,7 +28212,8 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 8e-08
"output_cost_per_token": 8e-08,
"supports_tool_choice": true
},
"vercel_ai_gateway/meta/llama-3.1-70b": {
"input_cost_per_token": 7.2e-07,
@ -28124,7 +28222,8 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 7.2e-07
"output_cost_per_token": 7.2e-07,
"supports_tool_choice": true
},
"vercel_ai_gateway/meta/llama-3.1-8b": {
"input_cost_per_token": 5e-08,
@ -28133,7 +28232,9 @@
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 8e-08
"output_cost_per_token": 8e-08,
"supports_function_calling": true,
"supports_response_schema": true
},
"vercel_ai_gateway/meta/llama-3.2-11b": {
"input_cost_per_token": 1.6e-07,
@ -28142,7 +28243,10 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 1.6e-07
"output_cost_per_token": 1.6e-07,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/meta/llama-3.2-1b": {
"input_cost_per_token": 1e-07,
@ -28160,7 +28264,9 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 1.5e-07
"output_cost_per_token": 1.5e-07,
"supports_function_calling": true,
"supports_response_schema": true
},
"vercel_ai_gateway/meta/llama-3.2-90b": {
"input_cost_per_token": 7.2e-07,
@ -28169,7 +28275,10 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 7.2e-07
"output_cost_per_token": 7.2e-07,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/meta/llama-3.3-70b": {
"input_cost_per_token": 7.2e-07,
@ -28178,7 +28287,9 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 7.2e-07
"output_cost_per_token": 7.2e-07,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/meta/llama-4-maverick": {
"input_cost_per_token": 2e-07,
@ -28187,7 +28298,8 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 6e-07
"output_cost_per_token": 6e-07,
"supports_tool_choice": true
},
"vercel_ai_gateway/meta/llama-4-scout": {
"input_cost_per_token": 1e-07,
@ -28196,7 +28308,10 @@
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 3e-07
"output_cost_per_token": 3e-07,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/mistral/codestral": {
"input_cost_per_token": 3e-07,
@ -28205,7 +28320,9 @@
"max_output_tokens": 4000,
"max_tokens": 4000,
"mode": "chat",
"output_cost_per_token": 9e-07
"output_cost_per_token": 9e-07,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/mistral/codestral-embed": {
"input_cost_per_token": 1.5e-07,
@ -28223,7 +28340,10 @@
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.8e-07
"output_cost_per_token": 2.8e-07,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/mistral/magistral-medium": {
"input_cost_per_token": 2e-06,
@ -28232,7 +28352,10 @@
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 5e-06
"output_cost_per_token": 5e-06,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/mistral/magistral-small": {
"input_cost_per_token": 5e-07,
@ -28241,7 +28364,8 @@
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 1.5e-06
"output_cost_per_token": 1.5e-06,
"supports_function_calling": true
},
"vercel_ai_gateway/mistral/ministral-3b": {
"input_cost_per_token": 4e-08,
@ -28250,7 +28374,9 @@
"max_output_tokens": 4000,
"max_tokens": 4000,
"mode": "chat",
"output_cost_per_token": 4e-08
"output_cost_per_token": 4e-08,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/mistral/ministral-8b": {
"input_cost_per_token": 1e-07,
@ -28259,7 +28385,10 @@
"max_output_tokens": 4000,
"max_tokens": 4000,
"mode": "chat",
"output_cost_per_token": 1e-07
"output_cost_per_token": 1e-07,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/mistral/mistral-embed": {
"input_cost_per_token": 1e-07,
@ -28277,7 +28406,9 @@
"max_output_tokens": 4000,
"max_tokens": 4000,
"mode": "chat",
"output_cost_per_token": 6e-06
"output_cost_per_token": 6e-06,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/mistral/mistral-saba-24b": {
"input_cost_per_token": 7.9e-07,
@ -28295,7 +28426,10 @@
"max_output_tokens": 4000,
"max_tokens": 4000,
"mode": "chat",
"output_cost_per_token": 3e-07
"output_cost_per_token": 3e-07,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/mistral/mixtral-8x22b-instruct": {
"input_cost_per_token": 1.2e-06,
@ -28304,7 +28438,8 @@
"max_output_tokens": 2048,
"max_tokens": 2048,
"mode": "chat",
"output_cost_per_token": 1.2e-06
"output_cost_per_token": 1.2e-06,
"supports_function_calling": true
},
"vercel_ai_gateway/mistral/pixtral-12b": {
"input_cost_per_token": 1.5e-07,
@ -28313,7 +28448,11 @@
"max_output_tokens": 4000,
"max_tokens": 4000,
"mode": "chat",
"output_cost_per_token": 1.5e-07
"output_cost_per_token": 1.5e-07,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/mistral/pixtral-large": {
"input_cost_per_token": 2e-06,
@ -28322,7 +28461,11 @@
"max_output_tokens": 4000,
"max_tokens": 4000,
"mode": "chat",
"output_cost_per_token": 6e-06
"output_cost_per_token": 6e-06,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/moonshotai/kimi-k2": {
"input_cost_per_token": 5.5e-07,
@ -28331,7 +28474,9 @@
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 2.2e-06
"output_cost_per_token": 2.2e-06,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/morph/morph-v3-fast": {
"input_cost_per_token": 8e-07,
@ -28358,7 +28503,9 @@
"max_output_tokens": 4096,
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 1.5e-06
"output_cost_per_token": 1.5e-06,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/openai/gpt-3.5-turbo-instruct": {
"input_cost_per_token": 1.5e-06,
@ -28376,7 +28523,10 @@
"max_output_tokens": 4096,
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 3e-05
"output_cost_per_token": 3e-05,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/openai/gpt-4.1": {
"cache_creation_input_token_cost": 0.0,
@ -28387,7 +28537,11 @@
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 8e-06
"output_cost_per_token": 8e-06,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/openai/gpt-4.1-mini": {
"cache_creation_input_token_cost": 0.0,
@ -28398,7 +28552,11 @@
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 1.6e-06
"output_cost_per_token": 1.6e-06,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/openai/gpt-4.1-nano": {
"cache_creation_input_token_cost": 0.0,
@ -28409,7 +28567,11 @@
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 4e-07
"output_cost_per_token": 4e-07,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/openai/gpt-4o": {
"cache_creation_input_token_cost": 0.0,
@ -28420,7 +28582,11 @@
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1e-05
"output_cost_per_token": 1e-05,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/openai/gpt-4o-mini": {
"cache_creation_input_token_cost": 0.0,
@ -28431,7 +28597,11 @@
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 6e-07
"output_cost_per_token": 6e-07,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/openai/o1": {
"cache_creation_input_token_cost": 0.0,
@ -28442,7 +28612,11 @@
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 6e-05
"output_cost_per_token": 6e-05,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/openai/o3": {
"cache_creation_input_token_cost": 0.0,
@ -28453,7 +28627,11 @@
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 8e-06
"output_cost_per_token": 8e-06,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/openai/o3-mini": {
"cache_creation_input_token_cost": 0.0,
@ -28464,7 +28642,10 @@
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.4e-06
"output_cost_per_token": 4.4e-06,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/openai/o4-mini": {
"cache_creation_input_token_cost": 0.0,
@ -28475,7 +28656,11 @@
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.4e-06
"output_cost_per_token": 4.4e-06,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true
},
"vercel_ai_gateway/openai/text-embedding-3-large": {
"input_cost_per_token": 1.3e-07,
@ -28547,7 +28732,10 @@
"max_output_tokens": 32000,
"max_tokens": 32000,
"mode": "chat",
"output_cost_per_token": 1.5e-05
"output_cost_per_token": 1.5e-05,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/vercel/v0-1.5-md": {
"input_cost_per_token": 3e-06,
@ -28556,7 +28744,10 @@
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 1.5e-05
"output_cost_per_token": 1.5e-05,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/xai/grok-2": {
"input_cost_per_token": 2e-06,
@ -28565,7 +28756,9 @@
"max_output_tokens": 4000,
"max_tokens": 4000,
"mode": "chat",
"output_cost_per_token": 1e-05
"output_cost_per_token": 1e-05,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/xai/grok-2-vision": {
"input_cost_per_token": 2e-06,
@ -28574,7 +28767,10 @@
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 1e-05
"output_cost_per_token": 1e-05,
"supports_vision": true,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/xai/grok-3": {
"input_cost_per_token": 3e-06,
@ -28583,7 +28779,9 @@
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.5e-05
"output_cost_per_token": 1.5e-05,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/xai/grok-3-fast": {
"input_cost_per_token": 5e-06,
@ -28592,7 +28790,8 @@
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 2.5e-05
"output_cost_per_token": 2.5e-05,
"supports_function_calling": true
},
"vercel_ai_gateway/xai/grok-3-mini": {
"input_cost_per_token": 3e-07,
@ -28601,7 +28800,9 @@
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 5e-07
"output_cost_per_token": 5e-07,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/xai/grok-3-mini-fast": {
"input_cost_per_token": 6e-07,
@ -28610,7 +28811,9 @@
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4e-06
"output_cost_per_token": 4e-06,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/xai/grok-4": {
"input_cost_per_token": 3e-06,
@ -28619,7 +28822,9 @@
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"output_cost_per_token": 1.5e-05
"output_cost_per_token": 1.5e-05,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/zai/glm-4.5": {
"input_cost_per_token": 6e-07,
@ -28628,7 +28833,9 @@
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 2.2e-06
"output_cost_per_token": 2.2e-06,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/zai/glm-4.5-air": {
"input_cost_per_token": 2e-07,
@ -28637,7 +28844,9 @@
"max_output_tokens": 96000,
"max_tokens": 96000,
"mode": "chat",
"output_cost_per_token": 1.1e-06
"output_cost_per_token": 1.1e-06,
"supports_function_calling": true,
"supports_tool_choice": true
},
"vercel_ai_gateway/zai/glm-4.6": {
"litellm_provider": "vercel_ai_gateway",
@ -29799,7 +30008,9 @@
"mode": "chat",
"output_cost_per_token": 1e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
"supported_regions": ["global"],
"supported_regions": [
"global"
],
"supports_function_calling": true,
"supports_tool_choice": true
},
@ -29812,7 +30023,9 @@
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
"supported_regions": ["global"],
"supported_regions": [
"global"
],
"supports_function_calling": true,
"supports_tool_choice": true
},
@ -29825,7 +30038,9 @@
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
"supported_regions": ["global"],
"supported_regions": [
"global"
],
"supports_function_calling": true,
"supports_tool_choice": true
},
@ -29838,7 +30053,9 @@
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
"supported_regions": ["global"],
"supported_regions": [
"global"
],
"supports_function_calling": true,
"supports_tool_choice": true
},
@ -34787,4 +35004,4 @@
"output_cost_per_token": 0,
"supports_reasoning": true
}
}
}

20
poetry.lock generated
View file

@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 2.1.4 and should not be changed by hand.
# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand.
[[package]]
name = "a2a-sdk"
@ -5704,24 +5704,6 @@ pytest = ">=7.0.0"
[package.extras]
dev = ["black", "flake8", "isort", "mypy"]
[[package]]
name = "pytest-retry"
version = "1.7.0"
description = "Adds the ability to retry flaky tests in CI environments"
optional = false
python-versions = ">=3.9"
groups = ["dev"]
files = [
{file = "pytest_retry-1.7.0-py3-none-any.whl", hash = "sha256:a2dac85b79a4e2375943f1429479c65beb6c69553e7dae6b8332be47a60954f4"},
{file = "pytest_retry-1.7.0.tar.gz", hash = "sha256:f8d52339f01e949df47c11ba9ee8d5b362f5824dff580d3870ec9ae0057df80f"},
]
[package.dependencies]
pytest = ">=7.0.0"
[package.extras]
dev = ["black", "flake8", "isort", "mypy"]
[[package]]
name = "python-dateutil"
version = "2.9.0.post0"

17
provider_endpoints_support.json
View file

@ -32,6 +32,23 @@
}
},
"providers": {
"a2a": {
"display_name": "A2A (Agent-to-Agent) (`a2a`)",
"url": "https://docs.litellm.ai/docs/providers/a2a",
"endpoints": {
"chat_completions": true,
"messages": false,
"responses": false,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,
"batches": false,
"rerank": false,
"a2a": false
}
},
"abliteration": {
"display_name": "Abliteration (`abliteration`)",
"url": "https://docs.litellm.ai/docs/providers/abliteration",

4
pyproject.toml
View file

@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm"
version = "1.81.6"
version = "1.81.7"
description = "Library to easily interface with LLM API providers"
authors = ["BerriAI"]
license = "MIT"
@ -174,7 +174,7 @@ requires = ["poetry-core", "wheel"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "1.81.6"
version = "1.81.7"
version_files = [
"pyproject.toml:^version"
]

2
requirements.txt
View file

@ -73,4 +73,4 @@ pypdf>=6.6.2 # for PDF text extraction in RAG ingestion
########################
# LITELLM ENTERPRISE DEPENDENCIES
########################
litellm-enterprise==0.1.28
litellm-enterprise==0.1.29

15
schema.prisma
View file

@ -129,6 +129,7 @@ model LiteLLM_TeamTable {
team_member_permissions String[] @default([])
policies String[] @default([])
model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id])
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
@ -160,7 +161,8 @@ model LiteLLM_DeletedTeamTable {
team_member_permissions String[] @default([])
policies String[] @default([])
model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases
allow_team_guardrail_config Boolean @default(false)
// Original timestamps from team creation/updates
created_at DateTime? @map("created_at")
updated_at DateTime? @map("updated_at")
@ -305,16 +307,6 @@ model LiteLLM_VerificationToken {
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
// SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub"
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2
@@index([user_id, team_id])
// SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."team_id" = $1 OFFSET $2
@@index([team_id])
// 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])
}
// Audit table for deleted keys - preserves spend and key information for historical tracking
@ -784,6 +776,7 @@ model LiteLLM_GuardrailsTable {
guardrail_name String @unique
litellm_params Json
guardrail_info Json?
team_id String?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
}

1
tests/code_coverage_tests/recursive_detector.py
View file

@ -41,6 +41,7 @@ IGNORE_FUNCTIONS = [
"__getattr__", # lazy loading pattern in litellm/__init__.py with proper caching to prevent infinite recursion.
"_validate_inheritance_chain", # max depth set (default 100) to prevent infinite recursion in policy inheritance validation.
"_basic_json_schema_validate", # max depth set.
"extract_text_from_a2a_message", # max depth set (default 10) to prevent infinite recursion in A2A message parsing.
]

144
tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py
View file

@ -1,4 +1,3 @@
import io
import os
import sys
@ -10,13 +9,10 @@ from datetime import datetime, timedelta, timezone
from unittest.mock import MagicMock, call, patch
import pytest
from prometheus_client import REGISTRY, CollectorRegistry
from prometheus_client import REGISTRY
import litellm
from litellm import completion
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.types.utils import (
StandardLoggingHiddenParams,
StandardLoggingMetadata,
@ -37,7 +33,6 @@ from litellm.proxy._types import UserAPIKeyAuth
verbose_logger.setLevel(logging.DEBUG)
litellm.set_verbose = True
import time
@pytest.fixture
@ -293,7 +288,6 @@ async def test_increment_remaining_budget_metrics(prometheus_logger):
) as mock_get_team, patch(
"litellm.proxy.auth.auth_checks.get_key_object"
) as mock_get_key:
mock_get_team.return_value = MagicMock(budget_reset_at=future_reset_time_team)
mock_get_key.return_value = MagicMock(budget_reset_at=future_reset_time_key)
@ -648,25 +642,16 @@ async def test_async_log_failure_event(prometheus_logger):
)
# litellm_llm_api_failed_requests_metric incremented
"""
Expected metrics
end_user_id,
user_api_key,
user_api_key_alias,
model,
user_api_team,
user_api_team_alias,
user_id,
"""
# Labels: end_user, api_key_hash, api_key_alias, model, team, team_alias, user, model_id
prometheus_logger.litellm_llm_api_failed_requests_metric.labels.assert_called_once_with(
None,
None, # end_user_id
"test_hash",
"test_alias",
"gpt-3.5-turbo",
"test_team",
"test_team_alias",
"test_user",
"model-123",
"model-123", # model_id from standard_logging_payload
)
prometheus_logger.litellm_llm_api_failed_requests_metric.labels().inc.assert_called_once()
@ -678,38 +663,54 @@ async def test_async_log_failure_event(prometheus_logger):
api_provider="openai",
)
# deployment failure responses incremented
prometheus_logger.litellm_deployment_failure_responses.labels.assert_called_once_with(
litellm_model_name="gpt-3.5-turbo",
model_id="model-123",
api_base="https://api.openai.com",
api_provider="openai",
exception_status="None",
exception_class="Exception",
requested_model="openai-gpt", # passed in standard logging payload
hashed_api_key="test_hash",
api_key_alias="test_alias",
team="test_team",
team_alias="test_team_alias",
client_ip="127.0.0.1", # from standard logging payload
user_agent=None,
# deployment failure responses incremented - verify key labels are populated
prometheus_logger.litellm_deployment_failure_responses.labels.assert_called_once()
actual_failure_labels = (
prometheus_logger.litellm_deployment_failure_responses.labels.call_args.kwargs
)
expected_failure_labels = {
"litellm_model_name": "gpt-3.5-turbo",
"model_id": "model-123",
"api_base": "https://api.openai.com",
"api_provider": "openai",
"exception_class": "Exception",
"requested_model": "openai-gpt",
"hashed_api_key": "test_hash",
"api_key_alias": "test_alias",
"team": "test_team",
"team_alias": "test_team_alias",
}
for key, expected_val in expected_failure_labels.items():
assert key in actual_failure_labels, f"Missing label {key}"
assert (
actual_failure_labels[key] == expected_val
), f"Label {key}: expected {expected_val!r}, got {actual_failure_labels[key]!r}"
assert actual_failure_labels.get("exception_status") in ("None", None)
assert actual_failure_labels.get("client_ip") == "127.0.0.1"
prometheus_logger.litellm_deployment_failure_responses.labels().inc.assert_called_once()
# deployment total requests incremented
prometheus_logger.litellm_deployment_total_requests.labels.assert_called_once_with(
litellm_model_name="gpt-3.5-turbo",
model_id="model-123",
api_base="https://api.openai.com",
api_provider="openai",
requested_model="openai-gpt", # passed in standard logging payload
hashed_api_key="test_hash",
api_key_alias="test_alias",
team="test_team",
team_alias="test_team_alias",
client_ip="127.0.0.1", # from standard logging payload
user_agent=None,
# deployment total requests incremented - verify key labels are populated
prometheus_logger.litellm_deployment_total_requests.labels.assert_called_once()
actual_total_labels = (
prometheus_logger.litellm_deployment_total_requests.labels.call_args.kwargs
)
expected_total_labels = {
"litellm_model_name": "gpt-3.5-turbo",
"model_id": "model-123",
"api_base": "https://api.openai.com",
"api_provider": "openai",
"requested_model": "openai-gpt",
"hashed_api_key": "test_hash",
"api_key_alias": "test_alias",
"team": "test_team",
"team_alias": "test_team_alias",
}
for key, expected_val in expected_total_labels.items():
assert key in actual_total_labels, f"Missing label {key}"
assert (
actual_total_labels[key] == expected_val
), f"Label {key}: expected {expected_val!r}, got {actual_total_labels[key]!r}"
assert actual_total_labels.get("client_ip") == "127.0.0.1"
prometheus_logger.litellm_deployment_total_requests.labels().inc.assert_called_once()
@ -1095,7 +1096,7 @@ def test_increment_deployment_cooled_down(prometheus_logger):
import inspect
method_sig = inspect.signature(prometheus_logger.increment_deployment_cooled_down)
expected_label_count = len([p for p in method_sig.parameters.keys() if p != 'self'])
expected_label_count = len([p for p in method_sig.parameters.keys() if p != "self"])
mock_chain = MagicMock()
@ -1103,11 +1104,15 @@ def test_increment_deployment_cooled_down(prometheus_logger):
"""Validate label count matches metric definition"""
total = len(label_values) + len(label_kwargs)
if total != expected_label_count:
raise ValueError(f"Incorrect label count: expected {expected_label_count}, got {total}")
raise ValueError(
f"Incorrect label count: expected {expected_label_count}, got {total}"
)
return mock_chain
prometheus_logger.litellm_deployment_cooled_down = MagicMock()
prometheus_logger.litellm_deployment_cooled_down.labels = MagicMock(side_effect=validating_labels)
prometheus_logger.litellm_deployment_cooled_down.labels = MagicMock(
side_effect=validating_labels
)
prometheus_logger.increment_deployment_cooled_down(
litellm_model_name="gpt-3.5-turbo",
@ -1179,8 +1184,12 @@ def test_get_custom_labels_from_top_level_metadata(monkeypatch):
metadata = {
"requester_ip_address": "10.48.203.20", # Top-level field
"user_api_key_alias": "TestAlias", # Top-level field
"requester_metadata": {"nested_field": "nested_value"}, # Nested dict (excluded)
"user_api_key_auth_metadata": {"another_nested": "value"}, # Nested dict (excluded)
"requester_metadata": {
"nested_field": "nested_value"
}, # Nested dict (excluded)
"user_api_key_auth_metadata": {
"another_nested": "value"
}, # Nested dict (excluded)
}
result = get_custom_labels_from_metadata(metadata)
assert result == {
@ -1217,7 +1226,9 @@ def test_get_custom_labels_from_top_level_and_nested_metadata(monkeypatch):
}
async def test_async_log_success_event_with_top_level_metadata(prometheus_logger, monkeypatch):
async def test_async_log_success_event_with_top_level_metadata(
prometheus_logger, monkeypatch
):
"""
Test that async_log_success_event correctly extracts custom labels from top-level metadata
fields like requester_ip_address, not just from nested dictionaries.
@ -1231,7 +1242,9 @@ async def test_async_log_success_event_with_top_level_metadata(prometheus_logger
standard_logging_object = create_standard_logging_payload()
standard_logging_object["metadata"]["requester_ip_address"] = "10.48.203.20"
standard_logging_object["metadata"]["requester_metadata"] = {} # Empty nested dict
standard_logging_object["metadata"]["user_api_key_auth_metadata"] = {} # Empty nested dict
standard_logging_object["metadata"][
"user_api_key_auth_metadata"
] = {} # Empty nested dict
kwargs = {
"model": "gpt-3.5-turbo",
@ -1273,7 +1286,9 @@ async def test_async_log_success_event_with_top_level_metadata(prometheus_logger
prometheus_logger.litellm_remaining_user_budget_metric = create_mock_metric()
prometheus_logger.litellm_user_max_budget_metric = create_mock_metric()
prometheus_logger.litellm_user_budget_remaining_hours_metric = create_mock_metric()
prometheus_logger.litellm_remaining_api_key_requests_for_model = create_mock_metric()
prometheus_logger.litellm_remaining_api_key_requests_for_model = (
create_mock_metric()
)
prometheus_logger.litellm_remaining_api_key_tokens_for_model = create_mock_metric()
prometheus_logger.litellm_llm_api_time_to_first_token_metric = create_mock_metric()
prometheus_logger.litellm_llm_api_latency_metric = create_mock_metric()
@ -1302,7 +1317,7 @@ async def test_async_log_success_event_with_top_level_metadata(prometheus_logger
# This confirms that the custom label extraction logic ran without errors
assert prometheus_logger.litellm_requests_metric.labels.called
assert prometheus_logger.litellm_spend_metric.labels.called
# Verify that the labels() method was called with some arguments (either positional or keyword)
# This ensures the custom label extraction happened and didn't cause a "Incorrect label names" error
call_args = prometheus_logger.litellm_requests_metric.labels.call_args
@ -1494,7 +1509,6 @@ async def test_initialize_remaining_budget_metrics(prometheus_logger):
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
"litellm.proxy.management_endpoints.team_endpoints.get_paginated_teams"
) as mock_get_teams:
# Create mock team data with proper datetime objects for budget_reset_at
future_reset = datetime.now() + timedelta(hours=24) # Reset 24 hours from now
mock_teams = [
@ -1592,21 +1606,22 @@ async def test_initialize_remaining_budget_metrics_exception_handling(
) as mock_get_teams, patch(
"litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper"
) as mock_list_keys:
# Make get_paginated_teams raise an exception
mock_get_teams.side_effect = Exception("Database error")
mock_list_keys.side_effect = Exception("Key listing error")
# Mock prisma_client structure to raise an exception for user budget metrics
# The code accesses prisma_client.db.litellm_usertable.find_many and count
mock_usertable = MagicMock()
mock_usertable.find_many = MagicMock(side_effect=Exception("User database error"))
mock_usertable.find_many = MagicMock(
side_effect=Exception("User database error")
)
mock_usertable.count = MagicMock(side_effect=Exception("User count error"))
# Mock litellm_teamtable to raise an exception for team count metrics
mock_teamtable = MagicMock()
mock_teamtable.count = MagicMock(side_effect=Exception("Team count error"))
mock_db = MagicMock()
mock_db.litellm_usertable = mock_usertable
mock_db.litellm_teamtable = mock_teamtable
@ -1661,7 +1676,6 @@ async def test_initialize_api_key_budget_metrics(prometheus_logger):
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
"litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper"
) as mock_list_keys:
# Create mock key data with proper datetime objects for budget_reset_at
future_reset = datetime.now() + timedelta(hours=24) # Reset 24 hours from now
key1 = UserAPIKeyAuth(
@ -1916,7 +1930,6 @@ def test_prometheus_label_factory_with_custom_tags(monkeypatch):
Test that prometheus_label_factory correctly handles custom tags
"""
from litellm.integrations.prometheus import (
get_custom_labels_from_tags,
prometheus_label_factory,
)
from litellm.types.integrations.prometheus import UserAPIKeyLabelValues
@ -1954,7 +1967,6 @@ def test_prometheus_label_factory_with_no_custom_tags(monkeypatch):
Test that prometheus_label_factory works when no custom tags are configured
"""
from litellm.integrations.prometheus import (
get_custom_labels_from_tags,
prometheus_label_factory,
)
from litellm.types.integrations.prometheus import UserAPIKeyLabelValues
@ -2179,9 +2191,7 @@ async def test_prometheus_token_metrics_with_prometheus_config():
All three metrics should be properly incremented when making a successful completion request.
"""
from prometheus_client import CollectorRegistry, Counter
import litellm
from litellm.types.integrations.prometheus import PrometheusMetricsConfig
# Clear registry before test

204
tests/litellm/test_proxy_auth.py Normal file
View file

@ -0,0 +1,204 @@
"""
Unit tests for litellm.proxy_auth module.
Tests the OAuth2/JWT token management for LiteLLM Proxy authentication.
"""
import time
from unittest.mock import Mock, patch
import pytest
from litellm.proxy_auth import (
AccessToken,
AzureADCredential,
GenericOAuth2Credential,
ProxyAuthHandler,
)
class TestAccessToken:
"""Tests for AccessToken dataclass."""
def test_access_token_creation(self):
"""Test AccessToken can be created with required fields."""
token = AccessToken(token="test-token", expires_on=1234567890)
assert token.token == "test-token"
assert token.expires_on == 1234567890
def test_access_token_equality(self):
"""Test AccessToken equality comparison."""
token1 = AccessToken(token="test", expires_on=123)
token2 = AccessToken(token="test", expires_on=123)
assert token1 == token2
class MockCredential:
"""Mock credential for testing."""
def __init__(self, expires_in_seconds: int = 3600):
self.call_count = 0
self.expires_in = expires_in_seconds
def get_token(self, scope: str) -> AccessToken:
self.call_count += 1
return AccessToken(
token=f"mock-token-{self.call_count}",
expires_on=int(time.time()) + self.expires_in,
)
class TestProxyAuthHandler:
"""Tests for ProxyAuthHandler."""
def test_get_auth_headers_returns_bearer_token(self):
"""Test that get_auth_headers returns correct Authorization header."""
cred = MockCredential()
handler = ProxyAuthHandler(credential=cred, scope="test-scope")
headers = handler.get_auth_headers()
assert "Authorization" in headers
assert headers["Authorization"].startswith("Bearer ")
assert "mock-token-1" in headers["Authorization"]
def test_token_caching(self):
"""Test that tokens are cached and not re-requested."""
cred = MockCredential(expires_in_seconds=3600) # Long expiry
handler = ProxyAuthHandler(credential=cred, scope="test-scope")
# Multiple calls should only request token once
handler.get_auth_headers()
handler.get_auth_headers()
handler.get_auth_headers()
assert cred.call_count == 1
def test_token_refresh_when_about_to_expire(self):
"""Test that tokens are refreshed when about to expire (within 60s buffer)."""
cred = MockCredential(expires_in_seconds=30) # Expires in 30s (< 60s buffer)
handler = ProxyAuthHandler(credential=cred, scope="test-scope")
# First call gets token
handler.get_auth_headers()
# Second call should refresh because token expires within 60s buffer
handler.get_auth_headers()
assert cred.call_count == 2
def test_get_token_method(self):
"""Test the get_token method returns AccessToken."""
cred = MockCredential()
handler = ProxyAuthHandler(credential=cred, scope="test-scope")
token = handler.get_token()
assert isinstance(token, AccessToken)
assert token.token == "mock-token-1"
class TestAzureADCredential:
"""Tests for AzureADCredential."""
def test_lazy_initialization(self):
"""Test that azure-identity is not imported until get_token is called."""
# This should not raise ImportError even if azure-identity is not installed
cred = AzureADCredential(credential=None)
# _initialized should be False until get_token is called
assert cred._initialized is False
def test_wraps_azure_credential(self):
"""Test that AzureADCredential wraps an azure-identity credential."""
# Mock Azure credential
mock_azure_cred = Mock()
mock_azure_cred.get_token.return_value = Mock(
token="azure-token", expires_on=9999999999
)
cred = AzureADCredential(credential=mock_azure_cred)
token = cred.get_token("https://graph.microsoft.com/.default")
assert token.token == "azure-token"
assert token.expires_on == 9999999999
mock_azure_cred.get_token.assert_called_once_with(
"https://graph.microsoft.com/.default"
)
class TestGenericOAuth2Credential:
"""Tests for GenericOAuth2Credential."""
def test_token_request(self):
"""Test that GenericOAuth2Credential makes correct OAuth2 request."""
with patch("httpx.post") as mock_post:
mock_response = Mock()
mock_response.json.return_value = {
"access_token": "oauth2-token",
"expires_in": 3600,
}
mock_response.raise_for_status = Mock()
mock_post.return_value = mock_response
cred = GenericOAuth2Credential(
client_id="test-client",
client_secret="test-secret",
token_url="https://example.com/oauth2/token",
)
token = cred.get_token("test-scope")
assert token.token == "oauth2-token"
mock_post.assert_called_once()
call_kwargs = mock_post.call_args
assert call_kwargs[1]["data"]["grant_type"] == "client_credentials"
assert call_kwargs[1]["data"]["client_id"] == "test-client"
assert call_kwargs[1]["data"]["client_secret"] == "test-secret"
assert call_kwargs[1]["data"]["scope"] == "test-scope"
def test_token_caching(self):
"""Test that GenericOAuth2Credential caches tokens."""
with patch("httpx.post") as mock_post:
mock_response = Mock()
mock_response.json.return_value = {
"access_token": "oauth2-token",
"expires_in": 3600,
}
mock_response.raise_for_status = Mock()
mock_post.return_value = mock_response
cred = GenericOAuth2Credential(
client_id="test-client",
client_secret="test-secret",
token_url="https://example.com/oauth2/token",
)
# Multiple calls should only make one HTTP request
cred.get_token("test-scope")
cred.get_token("test-scope")
cred.get_token("test-scope")
assert mock_post.call_count == 1
class TestLiteLLMIntegration:
"""Tests for integration with litellm module."""
def test_proxy_auth_variable_exists(self):
"""Test that litellm.proxy_auth variable exists."""
import litellm
# Should be None by default
assert hasattr(litellm, "proxy_auth")
def test_proxy_auth_can_be_set(self):
"""Test that litellm.proxy_auth can be set to a ProxyAuthHandler."""
import litellm
original_value = litellm.proxy_auth
try:
cred = MockCredential()
handler = ProxyAuthHandler(credential=cred, scope="test")
litellm.proxy_auth = handler
assert litellm.proxy_auth is handler
finally:
litellm.proxy_auth = original_value

132
tests/llm_translation/test_a2a.py Normal file
View file

@ -0,0 +1,132 @@
"""
Minimal E2E tests for A2A (Agent-to-Agent) Protocol provider.
Tests validate that the endpoint is reachable and can handle both
streaming and non-streaming requests.
"""
import os
import sys
import pytest
sys.path.insert(0, os.path.abspath("../.."))
import litellm
@pytest.mark.asyncio
async def test_a2a_completion_async_non_streaming():
"""
Test A2A provider with async non-streaming request.
Minimal test to validate endpoint reachability.
Note: Requires an A2A agent running at http://0.0.0.0:9999
Set A2A_API_BASE environment variable to use a different endpoint.
"""
api_base = os.environ.get("A2A_API_BASE", "http://0.0.0.0:9999")
try:
response = await litellm.acompletion(
model="a2a/test-agent",
messages=[{"role": "user", "content": "Hello"}],
api_base=api_base,
stream=False,
)
print(f"Response: {response}")
assert response is not None, "Expected non-None response"
print(f"✅ Async non-streaming test passed")
except litellm.exceptions.APIConnectionError as e:
pytest.skip(f"A2A agent not reachable at {api_base}: {e}")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio
async def test_a2a_completion_async_streaming():
"""
Test A2A provider with async streaming request.
Minimal test to validate streaming endpoint reachability.
"""
api_base = os.environ.get("A2A_API_BASE", "http://0.0.0.0:9999")
try:
response = await litellm.acompletion(
model="a2a/test-agent",
messages=[{"role": "user", "content": "Hello"}],
api_base=api_base,
stream=True,
)
chunks = []
async for chunk in response: # type: ignore
chunks.append(chunk)
print(f"Chunk: {chunk}")
assert len(chunks) > 0, "Expected at least one chunk in streaming response"
print(f"✅ Async streaming test passed: received {len(chunks)} chunks")
except litellm.exceptions.APIConnectionError as e:
pytest.skip(f"A2A agent not reachable at {api_base}: {e}")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_a2a_completion_sync():
"""
Test A2A provider with synchronous non-streaming request.
Minimal test to validate sync endpoint reachability.
"""
api_base = os.environ.get("A2A_API_BASE", "http://0.0.0.0:9999")
try:
response = litellm.completion(
model="a2a/test-agent",
messages=[{"role": "user", "content": "Hello"}],
api_base=api_base,
stream=False,
)
print(f"Response: {response}")
assert response is not None, "Expected non-None response"
print(f"✅ Sync non-streaming test passed")
except litellm.exceptions.APIConnectionError as e:
pytest.skip(f"A2A agent not reachable at {api_base}: {e}")
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_a2a_completion_sync_streaming():
"""
Test A2A provider with synchronous streaming request.
Minimal test to validate sync streaming endpoint reachability.
"""
api_base = os.environ.get("A2A_API_BASE", "http://0.0.0.0:9999")
try:
response = litellm.completion(
model="a2a/test-agent",
messages=[{"role": "user", "content": "Hello"}],
api_base=api_base,
stream=True,
)
chunks = []
for chunk in response: # type: ignore
chunks.append(chunk)
print(f"Chunk: {chunk}")
assert len(chunks) > 0, "Expected at least one chunk in streaming response"
print(f"✅ Sync streaming test passed: received {len(chunks)} chunks")
except litellm.exceptions.APIConnectionError as e:
pytest.skip(f"A2A agent not reachable at {api_base}: {e}")
except Exception as e:
pytest.fail(f"Error occurred: {e}")

17
tests/llm_translation/test_gemini.py
View file

@ -1435,3 +1435,20 @@ def test_gemini_image_size_limit_exceeded():
error_message = str(excinfo.value)
assert "Image size" in error_message
assert "exceeds maximum allowed size" in error_message
@pytest.mark.asyncio
async def test_gemini_openai_web_search_tool_to_google_search():
"""
Test that OpenAI-style web_search tools are transformed to Gemini's googleSearch.
When passing {"type": "web_search"} or {"type": "web_search_preview"} to Gemini,
these should be transformed to googleSearch, not silently ignored.
"""
response = await litellm.acompletion(
model="gemini/gemini-2.5-flash",
messages=[{"role": "user", "content": "What is the capital of France?"}],
tools=[{"type": "web_search"}],
)
print("response: ", response.model_dump_json(indent=4))
assert hasattr(response, "vertex_ai_grounding_metadata")
assert getattr(response, "vertex_ai_grounding_metadata") is not None

89
tests/mcp_tests/test_semantic_tool_filter_e2e.py Normal file
View file

@ -0,0 +1,89 @@
"""
End-to-end test for MCP Semantic Tool Filtering
"""
import asyncio
import os
import sys
from unittest.mock import Mock
import pytest
sys.path.insert(0, os.path.abspath("../.."))
from mcp.types import Tool as MCPTool
# Check if semantic-router is available
try:
import semantic_router
SEMANTIC_ROUTER_AVAILABLE = True
except ImportError:
SEMANTIC_ROUTER_AVAILABLE = False
@pytest.mark.asyncio
@pytest.mark.skipif(
not SEMANTIC_ROUTER_AVAILABLE,
reason="semantic-router not installed. Install with: pip install 'litellm[semantic-router]'"
)
async def test_e2e_semantic_filter():
"""E2E: Load router/filter and verify hook filters tools."""
from litellm import Router
from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
SemanticMCPToolFilter,
)
# Create router and filter
router = Router(
model_list=[{
"model_name": "text-embedding-3-small",
"litellm_params": {"model": "openai/text-embedding-3-small"},
}]
)
filter_instance = SemanticMCPToolFilter(
embedding_model="text-embedding-3-small",
litellm_router_instance=router,
top_k=3,
enabled=True,
)
# Create 10 tools
tools = [
MCPTool(name="gmail_send", description="Send an email via Gmail", inputSchema={"type": "object"}),
MCPTool(name="calendar_create", description="Create a calendar event", inputSchema={"type": "object"}),
MCPTool(name="file_upload", description="Upload a file", inputSchema={"type": "object"}),
MCPTool(name="web_search", description="Search the web", inputSchema={"type": "object"}),
MCPTool(name="slack_send", description="Send Slack message", inputSchema={"type": "object"}),
MCPTool(name="doc_read", description="Read document", inputSchema={"type": "object"}),
MCPTool(name="db_query", description="Query database", inputSchema={"type": "object"}),
MCPTool(name="api_call", description="Make API call", inputSchema={"type": "object"}),
MCPTool(name="task_create", description="Create task", inputSchema={"type": "object"}),
MCPTool(name="note_add", description="Add note", inputSchema={"type": "object"}),
]
# Build router with test tools
filter_instance._build_router(tools)
hook = SemanticToolFilterHook(filter_instance)
data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Send an email and create a calendar event"}],
"tools": tools,
"metadata": {}, # Initialize metadata dict for hook to store filter stats
}
# Call hook
result = await hook.async_pre_call_hook(
user_api_key_dict=Mock(),
cache=Mock(),
data=data,
call_type="completion",
)
# Single assertion: hook filtered tools
assert result and len(result["tools"]) < len(tools), f"Expected filtered tools, got {len(result['tools'])} tools (original: {len(tools)})"
print(f"✅ E2E test passed: Filtering reduced tools from {len(tools)} to {len(result['tools'])}")
print(f" Filtered tools: {[t.name for t in result['tools']]}")

42
tests/test_fallbacks.py
View file

@ -336,45 +336,3 @@ async def test_chat_completion_bad_and_good_model():
f"Iteration {iteration + 1}: {'✓' if success else '✗'} ({time.time() - start_time:.2f}s)"
)
assert success, "Not all good model requests succeeded"
@pytest.mark.asyncio
async def test_router_fallback_exhaustion():
"""
Test for Bug 19985:
"""
from litellm import Router
import pytest
# Setup: Only ONE fallback model available
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "openai/fake", "api_key": "bad-key"},
},
{
"model_name": "bad-model-1",
"litellm_params": {"model": "azure/fake", "api_key": "bad-key"},
}
]
# max_fallbacks=10 is much larger than the 1 fallback provided in the list
router = Router(
model_list=model_list,
fallbacks=[{"gpt-3.5-turbo": ["bad-model-1"]}],
max_fallbacks=10
)
try:
# This will fail and attempt to fallback
await router.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "test"}]
)
except Exception as e:
# The success criteria is that we DON'T get an IndexError
assert not isinstance(e, IndexError), f"Expected API error, but got IndexError: {e}"
# Also ensure we actually hit a fallback attempt
print(f"Caught expected exception: {type(e).__name__}")

22
tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py
View file

@ -134,3 +134,25 @@ def test_transform_request_with_response_format():
assert result["text"]["format"]["type"] == "json_schema"
assert result["text"]["format"]["name"] == "person_schema"
assert "schema" in result["text"]["format"]
def test_transform_request_includes_extra_headers():
"""Test that transform_request forwards headers as extra_headers for upstream call."""
handler = LiteLLMResponsesTransformationHandler()
messages = [{"role": "user", "content": "Hello"}]
optional_params = {}
litellm_params = {}
class MockLoggingObj:
pass
headers = {"cf-aig-authorization": "secret-token"}
result = handler.transform_request(
model="gpt-5-pro",
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
litellm_logging_obj=MockLoggingObj(),
)
assert result.get("extra_headers") == headers

134
tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py
View file

@ -1,8 +1,8 @@
"""
Unit tests for Prometheus invalid API key request filtering.
Unit tests for Prometheus request tracking.
Tests functionality that prevents invalid API key requests (401 status codes)
from being recorded in Prometheus metrics.
Tests that all requests including 401/invalid-key failures are tracked in metrics
for debugging, security auditing, abuse detection, and capacity planning.
"""
import os
@ -29,12 +29,14 @@ def prometheus_logger():
class ExceptionWithCode:
"""Exception-like object with 'code' attribute (ProxyException pattern)."""
def __init__(self, code):
self.code = code
class ExceptionWithStatusCode:
"""Exception-like object with 'status_code' attribute."""
def __init__(self, status_code):
self.status_code = status_code
@ -42,17 +44,25 @@ class ExceptionWithStatusCode:
class TestExtractStatusCode:
"""Test status code extraction from various sources."""
@pytest.mark.parametrize("exception_class,code_value,expected", [
(ExceptionWithCode, "401", 401),
(ExceptionWithStatusCode, 401, 401),
])
def test_extract_from_exception(self, prometheus_logger, exception_class, code_value, expected):
@pytest.mark.parametrize(
"exception_class,code_value,expected",
[
(ExceptionWithCode, "401", 401),
(ExceptionWithStatusCode, 401, 401),
],
)
def test_extract_from_exception(
self, prometheus_logger, exception_class, code_value, expected
):
exception = exception_class(code_value)
assert prometheus_logger._extract_status_code(exception=exception) == expected
def test_extract_from_kwargs(self, prometheus_logger):
exception = ExceptionWithCode("401")
assert prometheus_logger._extract_status_code(kwargs={"exception": exception}) == 401
assert (
prometheus_logger._extract_status_code(kwargs={"exception": exception})
== 401
)
def test_extract_from_enum_values(self, prometheus_logger):
enum_values = Mock(status_code="401")
@ -60,45 +70,62 @@ class TestExtractStatusCode:
class TestInvalidAPIKeyDetection:
"""Test invalid API key request detection logic."""
"""Test that we no longer skip metrics - all requests are tracked."""
@pytest.mark.parametrize("status_code,expected", [
(401, True),
(200, False),
(500, False),
(None, False),
])
def test_status_code_detection(self, prometheus_logger, status_code, expected):
assert prometheus_logger._is_invalid_api_key_request(status_code=status_code) == expected
@pytest.mark.parametrize("status_code", [401, 200, 500, None])
def test_no_skip_for_any_status_code(self, prometheus_logger, status_code):
"""All status codes are tracked - never skip metrics."""
assert (
prometheus_logger._is_invalid_api_key_request(status_code=status_code)
is False
)
def test_auth_error_message_detection(self, prometheus_logger):
exception = AssertionError("LiteLLM Virtual Key expected. Received=invalid-key-12345, expected to start with 'sk-'.")
assert prometheus_logger._is_invalid_api_key_request(status_code=None, exception=exception) is True
def test_auth_error_message_not_skipped(self, prometheus_logger):
exception = AssertionError(
"LiteLLM Virtual Key expected. Received=invalid-key-12345, expected to start with 'sk-'."
)
assert (
prometheus_logger._is_invalid_api_key_request(
status_code=None, exception=exception
)
is False
)
def test_non_auth_exception_not_detected(self, prometheus_logger):
def test_non_auth_exception_not_skipped(self, prometheus_logger):
exception = ValueError("Some other error")
assert prometheus_logger._is_invalid_api_key_request(status_code=None, exception=exception) is False
assert (
prometheus_logger._is_invalid_api_key_request(
status_code=None, exception=exception
)
is False
)
class TestSkipMetricsValidation:
"""Test high-level validation method that orchestrates detection and extraction."""
"""Test that we never skip metrics - all requests are tracked."""
def test_skip_for_401_exception(self, prometheus_logger):
"""Test full flow: extraction -> detection -> skip decision."""
def test_no_skip_for_401_exception(self, prometheus_logger):
"""401 requests are now tracked for security auditing and abuse detection."""
exception = ExceptionWithCode("401")
assert prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) is True
assert (
prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception)
is False
)
def test_skip_for_auth_error_message(self, prometheus_logger):
"""Test full flow: exception message -> detection -> skip decision."""
def test_no_skip_for_auth_error_message(self, prometheus_logger):
"""Auth error messages are now tracked."""
exception = AssertionError("expected to start with 'sk-'")
assert prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) is True
assert (
prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception)
is False
)
def test_no_skip_for_valid_request(self, prometheus_logger):
assert prometheus_logger._should_skip_metrics_for_invalid_key() is False
class TestAsyncHooks:
"""Test async hook methods skip metrics for invalid API keys."""
"""Test async hook methods record metrics for all requests including 401s."""
@pytest.fixture
def mock_user_api_key(self):
@ -115,24 +142,33 @@ class TestAsyncHooks:
return user_key
@pytest.mark.asyncio
async def test_post_call_failure_hook_skips_401(self, prometheus_logger, mock_user_api_key):
async def test_post_call_failure_hook_records_401(
self, prometheus_logger, mock_user_api_key
):
"""401 failures are now recorded for security auditing and abuse detection."""
exception = ExceptionWithCode("401")
exception.__class__.__name__ = "ProxyException"
with patch.object(prometheus_logger, 'litellm_proxy_failed_requests_metric') as mock_failed, \
patch.object(prometheus_logger, 'litellm_proxy_total_requests_metric') as mock_total:
with patch.object(
prometheus_logger, "litellm_proxy_failed_requests_metric"
) as mock_failed, patch.object(
prometheus_logger, "litellm_proxy_total_requests_metric"
) as mock_total:
mock_failed.labels.return_value = Mock(inc=Mock())
mock_total.labels.return_value = Mock(inc=Mock())
await prometheus_logger.async_post_call_failure_hook(
request_data={"model": "test-model"},
request_data={"model": "test-model", "metadata": {}},
original_exception=exception,
user_api_key_dict=mock_user_api_key
user_api_key_dict=mock_user_api_key,
)
mock_failed.labels.assert_not_called()
mock_total.labels.assert_not_called()
mock_failed.labels.assert_called_once()
mock_total.labels.assert_called_once()
@pytest.mark.asyncio
async def test_log_failure_event_skips_401(self, prometheus_logger):
async def test_log_failure_event_records_401(self, prometheus_logger):
"""401 failures in log_failure_event are now recorded."""
exception = ExceptionWithCode("401")
kwargs = {
"model": "test-model",
@ -140,6 +176,9 @@ class TestAsyncHooks:
"metadata": {
"user_api_key_hash": "test-key",
"user_api_key_user_id": "test-user",
"user_api_key_alias": None,
"user_api_key_team_id": None,
"user_api_key_team_alias": None,
},
"model_group": "test-model",
},
@ -147,15 +186,16 @@ class TestAsyncHooks:
"litellm_params": {},
}
with patch.object(prometheus_logger, 'litellm_llm_api_failed_requests_metric') as mock_failed, \
patch.object(prometheus_logger, 'set_llm_deployment_failure_metrics') as mock_deployment:
with patch.object(
prometheus_logger, "litellm_llm_api_failed_requests_metric"
) as mock_failed, patch.object(
prometheus_logger, "set_llm_deployment_failure_metrics"
) as mock_deployment:
mock_failed.labels.return_value = Mock(inc=Mock())
await prometheus_logger.async_log_failure_event(
kwargs=kwargs,
response_obj=None,
start_time=None,
end_time=None
kwargs=kwargs, response_obj=None, start_time=None, end_time=None
)
mock_failed.labels.assert_not_called()
mock_deployment.assert_not_called()
mock_failed.labels.assert_called_once()
mock_deployment.assert_called_once()

236
tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py Normal file
View file

@ -0,0 +1,236 @@
"""
Unit tests for Anthropic Messages Guardrail Translation Handler
Tests the handler's ability to process streaming output for Anthropic Messages API
with guardrail transformations, specifically testing edge cases with empty choices.
"""
import os
import sys
from typing import Any, List, Literal, Optional
from unittest.mock import MagicMock, patch
import pytest
sys.path.insert(
0, os.path.abspath("../../../../../../..")
) # Adds the parent directory to the system path
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.anthropic.chat.guardrail_translation.handler import (
AnthropicMessagesHandler,
)
from litellm.types.utils import GenericGuardrailAPIInputs
class MockPassThroughGuardrail(CustomGuardrail):
"""Mock guardrail that passes through without blocking - for testing streaming fallback behavior"""
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
"""Simply return inputs unchanged"""
return inputs
class MockDynamicGuardrail(CustomGuardrail):
"""Mock guardrail that records dynamic params from request metadata."""
def __init__(self, guardrail_name: str):
super().__init__(guardrail_name=guardrail_name)
self.dynamic_params: Optional[dict] = None
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
self.dynamic_params = self.get_guardrail_dynamic_request_body_params(
request_data
)
return inputs
class TestAnthropicMessagesHandlerStreamingOutputProcessing:
"""Test streaming output processing functionality"""
@pytest.mark.asyncio
async def test_process_output_streaming_response_empty_model_response(self):
"""Test that streaming response with None model_response doesn't raise error
This test verifies the fix for the bug where accessing model_response.choices[0]
would raise an error when _build_complete_streaming_response returns None.
"""
handler = AnthropicMessagesHandler()
guardrail = MockPassThroughGuardrail(guardrail_name="test")
# Mock _check_streaming_has_ended to return True (stream ended)
# and _build_complete_streaming_response to return None
with patch.object(
handler, "_check_streaming_has_ended", return_value=True
), patch(
"litellm.llms.anthropic.chat.guardrail_translation.handler.AnthropicPassthroughLoggingHandler._build_complete_streaming_response",
return_value=None,
):
responses_so_far = [b"data: some chunk"]
# This should not raise an error
result = await handler.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail,
litellm_logging_obj=MagicMock(),
)
# Should return the responses unchanged
assert result == responses_so_far
class TestAnthropicMessagesHandlerInputProcessing:
"""Test input processing preserves litellm_metadata for dynamic guardrails."""
@pytest.mark.asyncio
async def test_process_input_messages_preserves_litellm_metadata_guardrails(self):
handler = AnthropicMessagesHandler()
guardrail = MockDynamicGuardrail(guardrail_name="cygnal-monitor")
data = {
"model": "claude-3-5-sonnet-20241022",
"messages": [{"role": "user", "content": "hello"}],
"litellm_metadata": {
"guardrails": [
{
"cygnal-monitor": {
"extra_body": {"policy_id": "policy-123"}
}
}
]
},
}
with patch("litellm.proxy.proxy_server.premium_user", True):
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert data.get("litellm_metadata", {}).get("guardrails")
assert guardrail.dynamic_params == {"policy_id": "policy-123"}
@pytest.mark.asyncio
async def test_process_output_streaming_response_empty_choices(self):
"""Test that streaming response with empty choices doesn't raise IndexError
This test verifies the fix for the bug where accessing model_response.choices[0]
would raise IndexError when the response has an empty choices list.
"""
from litellm.types.utils import ModelResponse
handler = AnthropicMessagesHandler()
guardrail = MockPassThroughGuardrail(guardrail_name="test")
# Create a mock response with empty choices
mock_response = ModelResponse(
id="msg_123",
created=1234567890,
model="claude-3",
object="chat.completion",
choices=[], # Empty choices
)
# Mock _check_streaming_has_ended to return True (stream ended)
# and _build_complete_streaming_response to return the mock response
with patch.object(
handler, "_check_streaming_has_ended", return_value=True
), patch(
"litellm.llms.anthropic.chat.guardrail_translation.handler.AnthropicPassthroughLoggingHandler._build_complete_streaming_response",
return_value=mock_response,
):
responses_so_far = [b"data: some chunk"]
# This should not raise IndexError
result = await handler.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail,
litellm_logging_obj=MagicMock(),
)
# Should return the responses unchanged
assert result == responses_so_far
@pytest.mark.asyncio
async def test_process_output_streaming_response_with_valid_choices(self):
"""Test that streaming response with valid choices still works correctly"""
from litellm.types.utils import Choices, Message, ModelResponse
handler = AnthropicMessagesHandler()
guardrail = MockPassThroughGuardrail(guardrail_name="test")
# Create a mock response with valid choices
mock_response = ModelResponse(
id="msg_123",
created=1234567890,
model="claude-3",
object="chat.completion",
choices=[
Choices(
finish_reason="stop",
index=0,
message=Message(
content="Hello world",
role="assistant",
),
)
],
)
# Mock _check_streaming_has_ended to return True (stream ended)
# and _build_complete_streaming_response to return the mock response
with patch.object(
handler, "_check_streaming_has_ended", return_value=True
), patch(
"litellm.llms.anthropic.chat.guardrail_translation.handler.AnthropicPassthroughLoggingHandler._build_complete_streaming_response",
return_value=mock_response,
):
responses_so_far = [b"data: some chunk"]
# This should process successfully
result = await handler.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail,
litellm_logging_obj=MagicMock(),
)
# Should return the responses
assert result == responses_so_far
@pytest.mark.asyncio
async def test_process_output_streaming_response_stream_not_ended(self):
"""Test that streaming response falls back to text processing when stream hasn't ended"""
handler = AnthropicMessagesHandler()
guardrail = MockPassThroughGuardrail(guardrail_name="test")
# Mock _check_streaming_has_ended to return False (stream not ended)
with patch.object(
handler, "_check_streaming_has_ended", return_value=False
), patch.object(
handler, "get_streaming_string_so_far", return_value="partial text"
):
responses_so_far = [b"data: some chunk"]
# This should process successfully using text-based guardrail
result = await handler.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail,
litellm_logging_obj=MagicMock(),
)
# Should return the responses
assert result == responses_so_far
if __name__ == "__main__":
# Run the tests
pytest.main([__file__, "-v"])

29
tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
View file

@ -108,3 +108,32 @@ def test_get_supported_openai_params_reasoning_effort():
"fireworks_ai/accounts/fireworks/models/llama-v3-70b-instruct"
)
assert "reasoning_effort" not in unsupported_params
def test_transform_messages_helper_removes_provider_specific_fields():
"""
Test that _transform_messages_helper removes provider_specific_fields from messages.
"""
config = FireworksAIConfig()
# Simulated messages, as dicts, including provider_specific_fields
messages = [
{
"role": "user",
"content": "Hello!",
"provider_specific_fields": {"extra": "should be removed"},
},
{
"role": "assistant",
"content": "Hi there!",
"provider_specific_fields": {"more": "remove this"},
},
{
"role": "user",
"content": "How are you?",
# no provider_specific_fields
}
]
# Call helper
out = config._transform_messages_helper(messages, model="fireworks/test", litellm_params={})
for msg in out:
assert "provider_specific_fields" not in msg

1
tests/test_litellm/llms/gemini/files/__init__.py Normal file
View file

@ -0,0 +1 @@
"""Tests for Gemini files functionality"""

298
tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py Normal file
View file

@ -0,0 +1,298 @@
"""
Test Google AI Studio (Gemini) files transformation functionality
"""
import os
import pytest
from unittest.mock import Mock, patch
import httpx
from litellm.llms.gemini.files.transformation import GoogleAIStudioFilesHandler
from litellm.types.llms.openai import OpenAIFileObject
class TestGoogleAIStudioFilesTransformation:
"""Test Google AI Studio files transformation"""
def setup_method(self):
"""Setup test method"""
self.handler = GoogleAIStudioFilesHandler()
def test_transform_retrieve_file_request_with_full_uri(self):
"""
Test that transform_retrieve_file_request returns empty params dict
to avoid 'Content-Type' query parameter error
Regression test for: https://github.com/BerriAI/litellm/issues/XXX
When retrieving a file, the API was incorrectly trying to pass Content-Type
as a query parameter, which Gemini API rejected.
"""
file_id = "https://generativelanguage.googleapis.com/v1beta/files/test123"
litellm_params = {"api_key": "test-api-key"}
url, params = self.handler.transform_retrieve_file_request(
file_id=file_id,
optional_params={},
litellm_params=litellm_params,
)
# Verify URL is constructed correctly with API key
assert "key=test-api-key" in url
assert file_id in url
# CRITICAL: params should be empty dict, not contain Content-Type or any other params
# These would be incorrectly interpreted as query parameters
assert params == {}, f"Expected empty params dict, got: {params}"
assert "Content-Type" not in params, "Content-Type should not be in query params"
def test_transform_retrieve_file_request_with_file_name_only(self):
"""
Test that transform_retrieve_file_request handles file_id without full URI
"""
file_id = "files/test123"
litellm_params = {"api_key": "test-api-key"}
url, params = self.handler.transform_retrieve_file_request(
file_id=file_id,
optional_params={},
litellm_params=litellm_params,
)
# Verify URL is constructed correctly
assert "generativelanguage.googleapis.com" in url
assert file_id in url
assert "key=test-api-key" in url
# CRITICAL: params should be empty dict
assert params == {}, f"Expected empty params dict, got: {params}"
assert "Content-Type" not in params, "Content-Type should not be in query params"
@patch.dict('os.environ', {}, clear=True)
@patch('litellm.llms.gemini.common_utils.get_secret_str', return_value=None)
def test_transform_retrieve_file_request_missing_api_key(self, mock_get_secret):
"""Test that transform_retrieve_file_request raises error when API key is missing"""
file_id = "files/test123"
litellm_params = {}
with pytest.raises(ValueError, match="api_key is required"):
self.handler.transform_retrieve_file_request(
file_id=file_id,
optional_params={},
litellm_params=litellm_params,
)
def test_transform_retrieve_file_response_success(self):
"""Test successful transformation of Gemini file retrieval response"""
# Mock response data from Gemini API
mock_response_data = {
"name": "files/test123",
"displayName": "test_file.pdf",
"mimeType": "application/pdf",
"sizeBytes": "1024",
"createTime": "2024-01-15T10:30:00.123456Z",
"updateTime": "2024-01-15T10:30:00.123456Z",
"expirationTime": "2024-01-17T10:30:00.123456Z",
"sha256Hash": "abcd1234",
"uri": "https://generativelanguage.googleapis.com/v1beta/files/test123",
"state": "ACTIVE",
}
# Create mock httpx response
mock_response = Mock(spec=httpx.Response)
mock_response.json.return_value = mock_response_data
# Create mock logging object
mock_logging_obj = Mock()
# Transform response
result = self.handler.transform_retrieve_file_response(
raw_response=mock_response,
logging_obj=mock_logging_obj,
litellm_params={},
)
# Verify transformation
assert isinstance(result, OpenAIFileObject)
assert result.id == mock_response_data["uri"]
assert result.filename == mock_response_data["displayName"]
assert result.bytes == int(mock_response_data["sizeBytes"])
assert result.object == "file"
assert result.purpose == "user_data"
assert result.status == "processed" # ACTIVE state maps to processed
assert result.status_details is None
def test_transform_retrieve_file_response_failed_state(self):
"""Test transformation of Gemini file retrieval response with FAILED state"""
mock_response_data = {
"name": "files/test123",
"displayName": "test_file.pdf",
"mimeType": "application/pdf",
"sizeBytes": "1024",
"createTime": "2024-01-15T10:30:00.123456Z",
"uri": "https://generativelanguage.googleapis.com/v1beta/files/test123",
"state": "FAILED",
"error": {"message": "Upload failed", "code": "INTERNAL"},
}
mock_response = Mock(spec=httpx.Response)
mock_response.json.return_value = mock_response_data
mock_logging_obj = Mock()
result = self.handler.transform_retrieve_file_response(
raw_response=mock_response,
logging_obj=mock_logging_obj,
litellm_params={},
)
# Verify error state handling
assert result.status == "error"
assert result.status_details is not None
assert "message" in result.status_details
def test_transform_retrieve_file_response_processing_state(self):
"""Test transformation of Gemini file retrieval response with PROCESSING state"""
mock_response_data = {
"name": "files/test123",
"displayName": "test_file.pdf",
"mimeType": "application/pdf",
"sizeBytes": "1024",
"createTime": "2024-01-15T10:30:00.123456Z",
"uri": "https://generativelanguage.googleapis.com/v1beta/files/test123",
"state": "PROCESSING",
}
mock_response = Mock(spec=httpx.Response)
mock_response.json.return_value = mock_response_data
mock_logging_obj = Mock()
result = self.handler.transform_retrieve_file_response(
raw_response=mock_response,
logging_obj=mock_logging_obj,
litellm_params={},
)
# PROCESSING state should map to "uploaded" status
assert result.status == "uploaded"
def test_transform_retrieve_file_response_missing_createTime(self):
"""
Test that transform_retrieve_file_response raises proper error when createTime is missing
This tests the error scenario that occurs when API returns an error response
without the expected file metadata fields.
"""
# Mock error response from Gemini API (missing createTime)
mock_response_data = {
"error": {
"code": 400,
"message": "Invalid request",
"status": "INVALID_ARGUMENT",
}
}
mock_response = Mock(spec=httpx.Response)
mock_response.json.return_value = mock_response_data
mock_logging_obj = Mock()
# Should raise ValueError with helpful message
with pytest.raises(ValueError, match="Error parsing file retrieve response"):
self.handler.transform_retrieve_file_response(
raw_response=mock_response,
logging_obj=mock_logging_obj,
litellm_params={},
)
def test_validate_environment(self):
"""Test that validate_environment properly adds API key to headers"""
headers = {}
api_key = "test-gemini-api-key"
result_headers = self.handler.validate_environment(
headers=headers,
model="gemini-pro",
messages=[],
optional_params={},
litellm_params={},
api_key=api_key,
)
# Verify API key is added to headers
assert "x-goog-api-key" in result_headers
assert result_headers["x-goog-api-key"] == api_key
@patch.dict('os.environ', {}, clear=True)
@patch('litellm.llms.gemini.common_utils.get_secret_str', return_value=None)
def test_validate_environment_missing_api_key(self, mock_get_secret):
"""Test that validate_environment raises error when API key is missing"""
headers = {}
with pytest.raises(
ValueError, match="GEMINI_API_KEY is required for Google AI Studio file operations"
):
self.handler.validate_environment(
headers=headers,
model="gemini-pro",
messages=[],
optional_params={},
litellm_params={},
api_key=None,
)
def test_get_complete_url(self):
"""Test that get_complete_url constructs proper upload URL"""
api_base = "https://generativelanguage.googleapis.com"
api_key = "test-api-key"
url = self.handler.get_complete_url(
api_base=api_base,
api_key=api_key,
model="gemini-pro",
optional_params={},
litellm_params={},
)
# Verify URL structure
assert api_base in url
assert "upload/v1beta/files" in url
assert f"key={api_key}" in url
def test_transform_delete_file_request_with_full_uri(self):
"""Test delete file request transformation with full URI"""
file_id = "https://generativelanguage.googleapis.com/v1beta/files/test123"
litellm_params = {
"api_key": "test-api-key",
"api_base": "https://generativelanguage.googleapis.com",
}
url, params = self.handler.transform_delete_file_request(
file_id=file_id,
optional_params={},
litellm_params=litellm_params,
)
# Verify URL extraction
assert "files/test123" in url
assert "generativelanguage.googleapis.com" in url
# Params should be empty (API key goes in header via validate_environment)
assert params == {}
def test_transform_delete_file_request_with_file_name_only(self):
"""Test delete file request transformation with file name only"""
file_id = "files/test123"
litellm_params = {
"api_key": "test-api-key",
"api_base": "https://generativelanguage.googleapis.com",
}
url, params = self.handler.transform_delete_file_request(
file_id=file_id,
optional_params={},
litellm_params=litellm_params,
)
# Verify URL construction
assert file_id in url
assert "generativelanguage.googleapis.com" in url
assert params == {}

148
tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py
View file

@ -733,6 +733,154 @@ class TestOpenAIChatCompletionsHandlerToolCallsOutput:
assert response.choices[0].finish_reason == "tool_calls"
class MockPassThroughGuardrail(CustomGuardrail):
"""Mock guardrail that passes through without blocking - for testing streaming fallback behavior"""
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
"""Simply return inputs unchanged"""
return inputs
class TestOpenAIChatCompletionsHandlerStreamingOutput:
"""Test streaming output processing functionality"""
@pytest.mark.asyncio
async def test_process_output_streaming_response_empty_choices(self):
"""Test that streaming response with empty choices doesn't raise IndexError
This test verifies the fix for the bug where accessing chunk.choices[0]
would raise IndexError when a streaming chunk has an empty choices list.
"""
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
handler = OpenAIChatCompletionsHandler()
guardrail = MockPassThroughGuardrail(guardrail_name="test")
# Create a streaming chunk with empty choices
chunk_with_empty_choices = ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[], # Empty choices - this was causing the IndexError
)
responses_so_far = [chunk_with_empty_choices]
# This should not raise IndexError
result = await handler.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
# Should return the responses unchanged
assert result == responses_so_far
@pytest.mark.asyncio
async def test_process_output_streaming_response_with_valid_choices(self):
"""Test that streaming response with valid choices still works correctly"""
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
handler = OpenAIChatCompletionsHandler()
guardrail = MockPassThroughGuardrail(guardrail_name="test")
# Create streaming chunks with valid choices
chunk1 = ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[
StreamingChoices(
index=0,
delta=Delta(content="Hello"),
finish_reason=None,
)
],
)
chunk2 = ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[
StreamingChoices(
index=0,
delta=Delta(content=" world"),
finish_reason="stop",
)
],
)
responses_so_far = [chunk1, chunk2]
# This should process successfully
result = await handler.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
# Should return the responses
assert result == responses_so_far
@pytest.mark.asyncio
async def test_process_output_streaming_response_mixed_empty_and_valid_choices_no_finish(self):
"""Test streaming response with mix of empty and valid choices chunks (stream not finished)
This tests the has_stream_ended check when iterating through chunks with mixed choices.
The stream hasn't finished yet (no finish_reason), so it won't trigger stream_chunk_builder.
"""
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
handler = OpenAIChatCompletionsHandler()
guardrail = MockPassThroughGuardrail(guardrail_name="test")
# Mix of chunks - some with empty choices, some with valid choices
# Stream hasn't finished (no finish_reason)
chunk_empty = ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[],
)
chunk_valid = ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[
StreamingChoices(
index=0,
delta=Delta(content="Hello"),
finish_reason=None, # Stream not finished
)
],
)
responses_so_far = [chunk_empty, chunk_valid]
# This should not raise IndexError when checking has_stream_ended
result = await handler.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
# Should return the responses
assert result == responses_so_far
if __name__ == "__main__":
# Run the tests
pytest.main([__file__, "-v"])

178
tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
View file

@ -817,3 +817,181 @@ class TestOpenAIResponsesHandlerToolCallExtraction:
assert task_mappings[0] == (0, 0)
assert task_mappings[1] == (0, 1)
assert task_mappings[2] == (0, 2)
class MockPassThroughGuardrail(CustomGuardrail):
"""Mock guardrail that passes through without blocking - for testing streaming fallback behavior"""
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
"""Simply return inputs unchanged"""
return inputs
class TestOpenAIResponsesHandlerStreamingOutputProcessing:
"""Test streaming output processing functionality"""
@pytest.mark.asyncio
async def test_process_output_streaming_response_empty_output(self):
"""Test that streaming response with empty output doesn't raise IndexError
This test verifies the fix for the bug where accessing model_response_choices[0]
would raise IndexError when the response.completed event has an empty output array.
"""
handler = OpenAIResponsesHandler()
guardrail = MockPassThroughGuardrail(guardrail_name="test")
# Simulate a response.completed streaming event with empty output
responses_so_far = [
{
"type": "response.completed",
"response": {
"id": "resp_123",
"output": [], # Empty output - this was causing the IndexError
"status": "completed",
},
}
]
# This should not raise IndexError
result = await handler.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
# Should return the responses unchanged
assert result == responses_so_far
@pytest.mark.asyncio
async def test_process_output_streaming_response_missing_output_key(self):
"""Test that streaming response with missing output key doesn't raise IndexError
This test verifies the handler gracefully handles when the response dict
doesn't contain an 'output' key at all.
"""
handler = OpenAIResponsesHandler()
guardrail = MockPassThroughGuardrail(guardrail_name="test")
# Simulate a response.completed streaming event with missing output key
responses_so_far = [
{
"type": "response.completed",
"response": {
"id": "resp_123",
"status": "completed",
# No 'output' key - get() will return []
},
}
]
# This should not raise IndexError
result = await handler.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
# Should return the responses unchanged
assert result == responses_so_far
@pytest.mark.asyncio
async def test_process_output_streaming_response_unrecognized_output_type(self):
"""Test that streaming response with unrecognized output types doesn't raise IndexError
This test verifies the handler gracefully handles when output items are of
unrecognized types that _convert_response_output_to_choices skips over.
"""
handler = OpenAIResponsesHandler()
guardrail = MockPassThroughGuardrail(guardrail_name="test")
# Simulate a response.completed streaming event with unrecognized output type
responses_so_far = [
{
"type": "response.completed",
"response": {
"id": "resp_123",
"output": [
{
"type": "unknown_type", # Unrecognized type
"id": "item_123",
"data": "some data",
}
],
"status": "completed",
},
}
]
# This should not raise IndexError
result = await handler.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
# Should return the responses unchanged
assert result == responses_so_far
@pytest.mark.asyncio
async def test_process_output_streaming_response_with_valid_output(self):
"""Test that streaming response with valid output still works correctly"""
handler = OpenAIResponsesHandler()
guardrail = MockPassThroughGuardrail(guardrail_name="test")
# Simulate a response.completed streaming event with valid message output
responses_so_far = [
{
"type": "response.created",
"response": {"id": "resp_123"},
},
{
"type": "response.output_item.added",
"item": {"type": "message", "id": "msg_123"},
},
{
"type": "response.content_part.added",
"part": {"type": "output_text", "text": ""},
},
{
"type": "response.output_text.delta",
"delta": "Hello",
},
{
"type": "response.output_text.delta",
"delta": " world",
},
{
"type": "response.completed",
"response": {
"id": "resp_123",
"output": [
{
"type": "message",
"id": "msg_123",
"status": "completed",
"role": "assistant",
"content": [
{"type": "output_text", "text": "Hello world"},
],
}
],
"status": "completed",
},
},
]
# This should process successfully
result = await handler.process_output_streaming_response(
responses_so_far=responses_so_far,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
# Should return the responses
assert result == responses_so_far

46
tests/test_litellm/llms/test_lifecycle_fix.py Normal file
View file

@ -0,0 +1,46 @@
"""
Verifies that the httpx client used by AsyncOpenAI is NOT closed
when AsyncHTTPHandler instances are garbage collected.
"""
import asyncio
import gc
import httpx
from litellm.llms.openai.common_utils import BaseOpenAILLM
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
async def test_httpx_client_not_closed_by_handler_gc():
"""
Before the fix: _get_async_http_client() returned handler.client,
so when handler was GC'd its __del__ closed the client.
After the fix: returns a standalone httpx.AsyncClient, no handler involved.
"""
# Get the client the same way AsyncOpenAI would
client = BaseOpenAILLM._get_async_http_client()
assert isinstance(client, httpx.AsyncClient)
# Simulate what the old code did: create an AsyncHTTPHandler and GC it
handler = AsyncHTTPHandler()
handler_client = handler.client
del handler
gc.collect()
# The client from _get_async_http_client should still be open
# because it's NOT tied to any AsyncHTTPHandler
assert not client.is_closed, "Client was closed prematurely!"
# Verify it can actually send (build a request without sending)
try:
req = client.build_request("GET", "https://example.com")
print("PASS: Client is still usable after handler GC")
except RuntimeError as e:
if "closed" in str(e):
print(f"FAIL: {e}")
raise
raise
await client.aclose()
print("All checks passed!")
asyncio.run(test_httpx_client_not_closed_by_handler_gc())

105
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -2663,6 +2663,111 @@ def test_vertex_ai_single_tool_type_still_works():
assert tools[0]["code_execution"] == {}
def test_vertex_ai_openai_web_search_tool_transformation():
"""
Test that OpenAI-style web_search and web_search_preview tools are transformed to googleSearch.
This fixes the issue where passing OpenAI-style web search tools like:
{"type": "web_search"} or {"type": "web_search_preview"}
would be silently ignored (the request succeeds but grounding is not applied).
The fix transforms these to Gemini's googleSearch tool.
Input:
value=[{"type": "web_search"}]
Expected Output:
tools=[{"googleSearch": {}}]
"""
v = VertexGeminiConfig()
optional_params = {}
# Test web_search transformation
tools = v._map_function(
value=[{"type": "web_search"}],
optional_params=optional_params
)
assert len(tools) == 1, f"Expected 1 Tool object, got {len(tools)}"
assert "googleSearch" in tools[0], f"Expected googleSearch in tool, got {tools[0].keys()}"
assert tools[0]["googleSearch"] == {}, f"Expected empty googleSearch config, got {tools[0]['googleSearch']}"
def test_vertex_ai_openai_web_search_preview_tool_transformation():
"""
Test that OpenAI-style web_search_preview tool is transformed to googleSearch.
Input:
value=[{"type": "web_search_preview"}]
Expected Output:
tools=[{"googleSearch": {}}]
"""
v = VertexGeminiConfig()
optional_params = {}
# Test web_search_preview transformation
tools = v._map_function(
value=[{"type": "web_search_preview"}],
optional_params=optional_params
)
assert len(tools) == 1, f"Expected 1 Tool object, got {len(tools)}"
assert "googleSearch" in tools[0], f"Expected googleSearch in tool, got {tools[0].keys()}"
assert tools[0]["googleSearch"] == {}, f"Expected empty googleSearch config, got {tools[0]['googleSearch']}"
def test_vertex_ai_openai_web_search_with_function_tools():
"""
Test that OpenAI-style web_search tool works alongside function tools.
Input:
value=[
{"type": "web_search"},
{"type": "function", "function": {"name": "get_weather", "description": "Get weather"}},
]
Expected Output:
tools=[
{"googleSearch": {}},
{"function_declarations": [{"name": "get_weather", "description": "Get weather"}]},
]
"""
v = VertexGeminiConfig()
optional_params = {}
tools = v._map_function(
value=[
{"type": "web_search"},
{"type": "function", "function": {"name": "get_weather", "description": "Get weather"}},
],
optional_params=optional_params
)
# Should have 2 separate Tool objects
assert len(tools) == 2, f"Expected 2 Tool objects, got {len(tools)}"
# Find each tool type
search_tool = None
func_tool = None
for tool in tools:
if "googleSearch" in tool:
search_tool = tool
elif "function_declarations" in tool:
func_tool = tool
# Verify both tools are present
assert search_tool is not None, "googleSearch Tool should be present"
assert func_tool is not None, "function_declarations Tool should be present"
# Verify googleSearch is empty config
assert search_tool["googleSearch"] == {}
# Verify function declaration content
assert func_tool["function_declarations"][0]["name"] == "get_weather"
def test_vertex_ai_multiple_function_declarations_grouped():
"""
Test that multiple function declarations are grouped in ONE Tool object.

480
tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_enforcement.py Normal file
View file

@ -0,0 +1,480 @@
"""
Test to verify Team MCP permissions are enforced when using JWT authentication.
Scenario:
1. Team "ABC" exists with models configured and MCPs assigned
2. User JWT has team "ABC" in groups (via team_ids_jwt_field)
3. Call MCP list endpoint
4. EXPECTED: Team MCP permissions should be enforced
"""
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy._types import (
LiteLLM_JWTAuth,
LiteLLM_TeamTable,
LiteLLM_ObjectPermissionTable,
UserAPIKeyAuth,
)
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
@pytest.mark.asyncio
async def test_reproduce_jwt_mcp_enforcement_issue(monkeypatch):
"""
Reproduce the bug where Team MCP permissions are NOT enforced when using JWT.
Setup:
- Team "ABC" has models ["gpt-4"] and MCPs ["mcp-server-1"] assigned
- JWT has team "ABC" in groups field
- User calls MCP list endpoint (no model requested)
Expected: team_id should be set to "ABC" so MCP permissions are enforced
Actual (BUG): team_id is None because route check fails for MCP routes
"""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
# Setup mock router
router = Router(model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}])
import sys
import types
proxy_server_module = types.ModuleType("proxy_server")
proxy_server_module.llm_router = router
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
# Team "ABC" has models configured AND MCPs assigned
team_with_mcp = LiteLLM_TeamTable(
team_id="ABC",
models=["gpt-4"], # Team HAS models
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="perm-123",
mcp_servers=["mcp-server-1"], # Team has MCPs assigned
),
)
async def mock_get_team_object(*args, **kwargs):
team_id = kwargs.get("team_id") or args[0]
if team_id == "ABC":
return team_with_mcp
return None
monkeypatch.setattr(
"litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object
)
# Setup JWT handler with team_ids_jwt_field (groups)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
team_ids_jwt_field="groups", # Use groups field for teams
# NOTE: team_allowed_routes defaults to ["openai_routes", "info_routes"]
# which does NOT include "mcp_routes"
)
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
# Simulate JWT payload with team in groups
jwt_token = {
"sub": "user-123",
"groups": ["ABC"], # Team "ABC" is in groups
"scope": "",
}
# Mock auth_jwt to return our token
with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt:
mock_auth_jwt.return_value = jwt_token
# Call auth_builder for MCP route (like /mcp/tools/list)
result = await JWTAuthManager.auth_builder(
api_key="test-jwt-token",
jwt_handler=jwt_handler,
request_data={}, # No model in request (MCP endpoint)
general_settings={},
route="/mcp/tools/list", # MCP route
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
# THIS IS THE BUG: team_id should be "ABC" but it's None!
print(f"Result team_id: {result['team_id']}")
print(f"Result team_object: {result['team_object']}")
# The test should FAIL if the bug exists (team_id is None)
# If the fix is applied, team_id should be "ABC"
assert result["team_id"] == "ABC", (
f"BUG: team_id should be 'ABC' but got '{result['team_id']}'. "
f"This happens because default team_allowed_routes does not include 'mcp_routes', "
f"so allowed_routes_check() fails and the team is skipped in find_team_with_model_access()."
)
@pytest.mark.asyncio
async def test_verify_mcp_routes_in_default_team_allowed_routes():
"""
Verify that mcp_routes IS in the default team_allowed_routes.
This is required for team MCP permissions to work with JWT auth.
"""
default_jwt_auth = LiteLLM_JWTAuth()
print(f"Default team_allowed_routes: {default_jwt_auth.team_allowed_routes}")
# mcp_routes must be in defaults for team MCP permissions to work
assert "mcp_routes" in default_jwt_auth.team_allowed_routes, (
"mcp_routes must be in default team_allowed_routes for JWT MCP enforcement to work"
)
@pytest.mark.asyncio
async def test_mcp_route_check_passes_for_team():
"""
Verify that allowed_routes_check returns True for MCP routes with default settings.
This is required for teams to access MCP endpoints with JWT auth.
"""
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.auth.auth_checks import allowed_routes_check
jwt_auth = LiteLLM_JWTAuth() # Use defaults
# Check if MCP route is allowed for TEAM role
is_allowed = allowed_routes_check(
user_role=LitellmUserRoles.TEAM,
user_route="/mcp/tools/list",
litellm_proxy_roles=jwt_auth,
)
print(f"Is /mcp/tools/list allowed for TEAM with defaults? {is_allowed}")
# MCP routes should be allowed by default for teams
assert is_allowed is True, (
"MCP routes must be allowed by default for teams for JWT MCP enforcement to work"
)
@pytest.mark.asyncio
async def test_e2e_jwt_team_mcp_permissions_enforced(monkeypatch):
"""
End-to-end test verifying that team MCP permissions are properly enforced
when using JWT authentication with teams in groups.
This test verifies the complete flow:
1. JWT token contains team "ABC" in groups field
2. Team "ABC" exists with MCP servers ["mcp-server-1", "mcp-server-2"] assigned
3. JWT auth properly sets team_id on UserAPIKeyAuth
4. MCPRequestHandler.get_allowed_mcp_servers() returns team's MCP servers
"""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
# Setup mock router
router = Router(model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}])
import sys
import types
proxy_server_module = types.ModuleType("proxy_server")
proxy_server_module.llm_router = router
proxy_server_module.prisma_client = MagicMock() # Mock prisma client
proxy_server_module.user_api_key_cache = DualCache()
proxy_server_module.proxy_logging_obj = MagicMock()
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
# Team "ABC" has MCP servers assigned via object_permission
team_mcp_servers = ["mcp-server-1", "mcp-server-2"]
team_object_permission = LiteLLM_ObjectPermissionTable(
object_permission_id="perm-abc-123",
mcp_servers=team_mcp_servers,
mcp_access_groups=[],
vector_stores=[],
)
team_with_mcp = LiteLLM_TeamTable(
team_id="ABC",
models=["gpt-4"],
object_permission=team_object_permission,
object_permission_id="perm-abc-123",
)
async def mock_get_team_object(*args, **kwargs):
team_id = kwargs.get("team_id") or (args[0] if args else None)
if team_id == "ABC":
return team_with_mcp
return None
monkeypatch.setattr(
"litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object
)
monkeypatch.setattr(
"litellm.proxy.auth.auth_checks.get_team_object", mock_get_team_object
)
# Setup JWT handler with team_ids_jwt_field (groups)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
team_ids_jwt_field="groups",
)
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
# Simulate JWT payload with team in groups
jwt_token = {
"sub": "user-123",
"groups": ["ABC"],
"scope": "",
}
# Step 1: Verify JWT auth returns correct team_id
with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt:
mock_auth_jwt.return_value = jwt_token
result = await JWTAuthManager.auth_builder(
api_key="test-jwt-token",
jwt_handler=jwt_handler,
request_data={},
general_settings={},
route="/mcp/tools/list",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
# Verify team_id is set correctly
assert result["team_id"] == "ABC", f"Expected team_id='ABC', got '{result['team_id']}'"
assert result["team_object"] is not None, "team_object should not be None"
# Step 2: Create UserAPIKeyAuth with the team_id from JWT auth
user_api_key_auth = UserAPIKeyAuth(
api_key=None,
team_id=result["team_id"],
user_id=result["user_id"],
)
# Step 3: Verify MCPRequestHandler returns team's MCP servers
# Mock _get_team_object_permission to return our team's object_permission
with patch.object(
MCPRequestHandler, "_get_team_object_permission"
) as mock_get_team_perm:
mock_get_team_perm.return_value = team_object_permission
# Mock _get_allowed_mcp_servers_for_key to return empty (no key-level permissions)
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_key"
) as mock_key_servers:
mock_key_servers.return_value = []
# Mock _get_mcp_servers_from_access_groups to return empty
with patch.object(
MCPRequestHandler, "_get_mcp_servers_from_access_groups"
) as mock_access_groups:
mock_access_groups.return_value = []
allowed_servers = await MCPRequestHandler.get_allowed_mcp_servers(
user_api_key_auth
)
print(f"Allowed MCP servers: {allowed_servers}")
# Verify team's MCP servers are returned
assert set(allowed_servers) == set(team_mcp_servers), (
f"Expected team MCP servers {team_mcp_servers}, got {allowed_servers}"
)
@pytest.mark.asyncio
async def test_e2e_jwt_without_team_no_mcp_servers(monkeypatch):
"""
End-to-end test verifying that when JWT has no teams, no MCP servers are returned.
This ensures:
1. JWT token with no groups returns no team_id
2. MCPRequestHandler.get_allowed_mcp_servers() returns empty list
"""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
# Setup mock router
router = Router(model_list=[])
import sys
import types
proxy_server_module = types.ModuleType("proxy_server")
proxy_server_module.llm_router = router
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
async def mock_get_team_object(*args, **kwargs):
return None
monkeypatch.setattr(
"litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object
)
# Setup JWT handler
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
team_ids_jwt_field="groups",
)
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
# JWT payload with empty groups
jwt_token = {
"sub": "user-123",
"groups": [], # No teams
"scope": "",
}
with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt:
mock_auth_jwt.return_value = jwt_token
result = await JWTAuthManager.auth_builder(
api_key="test-jwt-token",
jwt_handler=jwt_handler,
request_data={},
general_settings={},
route="/mcp/tools/list",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
# Verify no team_id is set
assert result["team_id"] is None, f"Expected team_id=None, got '{result['team_id']}'"
# Create UserAPIKeyAuth without team_id
user_api_key_auth = UserAPIKeyAuth(
api_key=None,
team_id=None,
user_id=result["user_id"],
)
# Verify no MCP servers are returned when there's no team
allowed_servers = await MCPRequestHandler._get_allowed_mcp_servers_for_team(
user_api_key_auth
)
assert allowed_servers == [], f"Expected empty list, got {allowed_servers}"
@pytest.mark.asyncio
async def test_e2e_jwt_team_mcp_key_intersection(monkeypatch):
"""
End-to-end test verifying MCP permission intersection between key and team.
Scenario:
- Team has MCP servers: ["server-1", "server-2", "server-3"]
- Key has MCP servers: ["server-2", "server-4"]
- Result should be intersection: ["server-2"]
"""
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
# Setup mock router
router = Router(model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}])
import sys
import types
proxy_server_module = types.ModuleType("proxy_server")
proxy_server_module.llm_router = router
proxy_server_module.prisma_client = MagicMock()
proxy_server_module.user_api_key_cache = DualCache()
proxy_server_module.proxy_logging_obj = MagicMock()
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
# Team MCP servers
team_mcp_servers = ["server-1", "server-2", "server-3"]
team_object_permission = LiteLLM_ObjectPermissionTable(
object_permission_id="team-perm",
mcp_servers=team_mcp_servers,
)
team_with_mcp = LiteLLM_TeamTable(
team_id="TEAM-X",
models=["gpt-4"],
object_permission=team_object_permission,
)
# Key MCP servers
key_mcp_servers = ["server-2", "server-4"]
key_object_permission = LiteLLM_ObjectPermissionTable(
object_permission_id="key-perm",
mcp_servers=key_mcp_servers,
)
async def mock_get_team_object(*args, **kwargs):
team_id = kwargs.get("team_id") or (args[0] if args else None)
if team_id == "TEAM-X":
return team_with_mcp
return None
monkeypatch.setattr(
"litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups")
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
jwt_token = {"sub": "user-123", "groups": ["TEAM-X"], "scope": ""}
with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt:
mock_auth_jwt.return_value = jwt_token
result = await JWTAuthManager.auth_builder(
api_key="test-jwt-token",
jwt_handler=jwt_handler,
request_data={},
general_settings={},
route="/mcp/tools/list",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
assert result["team_id"] == "TEAM-X"
user_api_key_auth = UserAPIKeyAuth(
api_key=None,
team_id=result["team_id"],
user_id=result["user_id"],
object_permission=key_object_permission, # Key has its own permissions
)
# Mock the helper methods to return our test data
with patch.object(
MCPRequestHandler, "_get_team_object_permission"
) as mock_team_perm:
mock_team_perm.return_value = team_object_permission
with patch.object(
MCPRequestHandler, "_get_key_object_permission"
) as mock_key_perm:
mock_key_perm.return_value = key_object_permission
with patch.object(
MCPRequestHandler, "_get_mcp_servers_from_access_groups"
) as mock_access_groups:
mock_access_groups.return_value = []
allowed_servers = await MCPRequestHandler.get_allowed_mcp_servers(
user_api_key_auth
)
# Should be intersection: only server-2 is in both
expected = ["server-2"]
assert sorted(allowed_servers) == sorted(expected), (
f"Expected intersection {expected}, got {allowed_servers}"
)

277
tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_simple.py Normal file
View file

@ -0,0 +1,277 @@
"""
Simple test to validate MCP permissions are enforced when calling MCP routes with JWT.
"""
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy._types import (
LiteLLM_JWTAuth,
LiteLLM_TeamTable,
LiteLLM_ObjectPermissionTable,
UserAPIKeyAuth,
)
@pytest.mark.asyncio
async def test_simple_jwt_mcp_permissions_enforced():
"""
Simple test: Call MCP route with JWT, verify team's MCP servers are returned.
Setup:
- Team "my-team" has MCP servers: ["github-mcp", "slack-mcp"]
- JWT user belongs to "my-team"
Expected: Only ["github-mcp", "slack-mcp"] should be allowed
"""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
# 1. Create a user authenticated via JWT with team_id set
user_auth = UserAPIKeyAuth(
api_key=None, # JWT auth doesn't have api_key
user_id="jwt-user-123",
team_id="my-team", # This is set by JWT auth when team is in groups
)
# 2. Team's MCP permissions
team_mcp_servers = ["github-mcp", "slack-mcp"]
team_object_permission = LiteLLM_ObjectPermissionTable(
object_permission_id="perm-123",
mcp_servers=team_mcp_servers,
)
# 3. Mock the team permission lookup
with patch.object(
MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock
) as mock_team_perm:
mock_team_perm.return_value = team_object_permission
# Mock key permissions (empty - user has no key-level MCP permissions)
with patch.object(
MCPRequestHandler, "_get_key_object_permission", new_callable=AsyncMock
) as mock_key_perm:
mock_key_perm.return_value = None
# Mock access groups (empty)
with patch.object(
MCPRequestHandler, "_get_mcp_servers_from_access_groups", new_callable=AsyncMock
) as mock_access_groups:
mock_access_groups.return_value = []
# 4. Call get_allowed_mcp_servers - this is what MCP routes use
allowed = await MCPRequestHandler.get_allowed_mcp_servers(user_auth)
# 5. Verify only team's MCP servers are returned
assert sorted(allowed) == sorted(team_mcp_servers), (
f"Expected {team_mcp_servers}, got {allowed}"
)
# Verify team permission was looked up
mock_team_perm.assert_called_once_with(user_auth)
@pytest.mark.asyncio
async def test_simple_jwt_no_team_no_mcp_servers():
"""
Simple test: JWT user with no team should get no MCP servers.
"""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
# User with no team_id (JWT didn't have teams in groups)
user_auth = UserAPIKeyAuth(
api_key=None,
user_id="jwt-user-no-team",
team_id=None, # No team
)
# _get_allowed_mcp_servers_for_team returns [] when team_id is None
allowed = await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_auth)
assert allowed == [], f"Expected [], got {allowed}"
@pytest.mark.asyncio
async def test_simple_jwt_team_id_required_for_mcp_permissions():
"""
Simple test: Verify that team_id must be set for team MCP permissions to work.
This is the key insight - if JWT auth doesn't set team_id,
team MCP permissions won't be enforced.
"""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
# Case 1: team_id is set -> team permissions should be checked
user_with_team = UserAPIKeyAuth(
api_key=None,
user_id="user-1",
team_id="team-abc",
)
team_mcp_servers = ["server-1", "server-2"]
team_perm = LiteLLM_ObjectPermissionTable(
object_permission_id="perm-1",
mcp_servers=team_mcp_servers,
)
with patch.object(
MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock
) as mock_perm:
mock_perm.return_value = team_perm
with patch.object(
MCPRequestHandler, "_get_mcp_servers_from_access_groups", new_callable=AsyncMock
) as mock_groups:
mock_groups.return_value = []
result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_with_team)
assert sorted(result) == sorted(team_mcp_servers)
mock_perm.assert_called_once() # Permission WAS checked
# Case 2: team_id is None -> team permissions NOT checked
user_without_team = UserAPIKeyAuth(
api_key=None,
user_id="user-2",
team_id=None,
)
result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_without_team)
assert result == [] # No permissions returned
@pytest.mark.asyncio
async def test_jwt_auth_sets_team_id_for_mcp_route():
"""
Test that JWT auth properly sets team_id when accessing MCP routes.
This is the critical test - when user calls /mcp/tools/list with JWT,
the team_id from JWT groups must be set on UserAPIKeyAuth.
"""
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
# Setup
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
team_ids_jwt_field="groups", # Teams come from "groups" field in JWT
)
# Team exists with models
team = LiteLLM_TeamTable(
team_id="team-from-jwt",
models=["gpt-4"],
)
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
# Mock JWT token with team in groups
jwt_payload = {
"sub": "user-123",
"groups": ["team-from-jwt"],
"scope": "",
}
with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth:
mock_auth.return_value = jwt_payload
with patch(
"litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock
) as mock_get_team:
mock_get_team.return_value = team
# Simulate calling MCP route
result = await JWTAuthManager.auth_builder(
api_key="jwt-token",
jwt_handler=jwt_handler,
request_data={},
general_settings={},
route="/mcp/tools/list", # MCP route
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
# THE KEY ASSERTION: team_id must be set
assert result["team_id"] == "team-from-jwt", (
f"team_id should be 'team-from-jwt' but got '{result['team_id']}'. "
"This means JWT auth is not properly setting team_id for MCP routes!"
)
@pytest.mark.asyncio
async def test_mcp_route_without_model_still_returns_team_id():
"""
Test that MCP routes (which don't specify a model) still get team_id assigned.
Key insight: MCP routes don't require a model in the request, but the JWT auth
flow must still assign a team_id so that team MCP permissions are enforced.
The flow is:
1. JWT token contains team in "groups" field
2. find_team_with_model_access() is called with requested_model=None
3. Since `not requested_model` is True, model check passes
4. Route check passes because "mcp_routes" is in team_allowed_routes
5. team_id is returned and set on UserAPIKeyAuth
"""
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
# Setup
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
team_ids_jwt_field="groups",
)
# Team exists - note: models is a list (can be empty or have values)
# The key is that when no model is requested, model check is skipped
team = LiteLLM_TeamTable(
team_id="my-team",
models=["gpt-4", "gpt-3.5-turbo"], # Team has models, but MCP request won't specify one
)
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
# JWT with team in groups
jwt_payload = {
"sub": "user-abc",
"groups": ["my-team"],
"scope": "",
}
with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth:
mock_auth.return_value = jwt_payload
with patch(
"litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock
) as mock_get_team:
mock_get_team.return_value = team
# Call MCP route with NO MODEL in request_data
result = await JWTAuthManager.auth_builder(
api_key="jwt-token",
jwt_handler=jwt_handler,
request_data={}, # <-- NO MODEL SPECIFIED
general_settings={},
route="/mcp/tools/list", # MCP route
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
# Team ID must still be set even though no model was requested
assert result["team_id"] == "my-team", (
f"Expected team_id='my-team' but got '{result['team_id']}'. "
"MCP routes without model should still get team_id from JWT!"
)

394
tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py Normal file
View file

@ -0,0 +1,394 @@
"""
Unit tests for MCP Semantic Tool Filtering
Tests the core filtering logic that takes a long list of tools and returns
an ordered set of top K tools based on semantic similarity.
"""
import asyncio
import os
import sys
from unittest.mock import AsyncMock, Mock, patch
import pytest
sys.path.insert(0, os.path.abspath("../.."))
from mcp.types import Tool as MCPTool
@pytest.mark.asyncio
async def test_semantic_filter_basic_filtering():
"""
Test that the semantic filter correctly filters tools based on query.
Given: 10 email/calendar tools
When: Query is "send an email"
Then: Email tools should rank higher than calendar tools
"""
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
SemanticMCPToolFilter,
)
# Create mock tools - mix of email and calendar tools
tools = [
MCPTool(name="gmail_send", description="Send an email via Gmail", inputSchema={"type": "object"}),
MCPTool(name="outlook_send", description="Send an email via Outlook", inputSchema={"type": "object"}),
MCPTool(name="calendar_create", description="Create a calendar event", inputSchema={"type": "object"}),
MCPTool(name="calendar_update", description="Update a calendar event", inputSchema={"type": "object"}),
MCPTool(name="email_read", description="Read emails from inbox", inputSchema={"type": "object"}),
MCPTool(name="email_delete", description="Delete an email", inputSchema={"type": "object"}),
MCPTool(name="calendar_delete", description="Delete a calendar event", inputSchema={"type": "object"}),
MCPTool(name="email_search", description="Search for emails", inputSchema={"type": "object"}),
MCPTool(name="calendar_list", description="List calendar events", inputSchema={"type": "object"}),
MCPTool(name="email_forward", description="Forward an email to someone", inputSchema={"type": "object"}),
]
# Mock router that returns mock embeddings
from litellm.types.utils import Embedding, EmbeddingResponse
mock_router = Mock()
def mock_embedding_sync(*args, **kwargs):
return EmbeddingResponse(
data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")],
model="text-embedding-3-small",
object="list",
usage={"prompt_tokens": 10, "total_tokens": 10}
)
async def mock_embedding_async(*args, **kwargs):
return mock_embedding_sync()
mock_router.embedding = mock_embedding_sync
mock_router.aembedding = mock_embedding_async
# Create filter
filter_instance = SemanticMCPToolFilter(
embedding_model="text-embedding-3-small",
litellm_router_instance=mock_router,
top_k=3,
similarity_threshold=0.3,
enabled=True,
)
# Build router with the tools before filtering
filter_instance._build_router(tools)
# Filter tools with email-related query
filtered = await filter_instance.filter_tools(
query="send an email to john@example.com",
available_tools=tools,
)
# Assertions - validate filtering mechanics work
assert len(filtered) <= 3, f"Should return at most 3 tools (top_k), got {len(filtered)}"
assert len(filtered) > 0, "Should return at least some tools"
assert len(filtered) < len(tools), f"Should filter down from {len(tools)} tools, got {len(filtered)}"
# Validate tools are actual MCPTool objects
for tool in filtered:
assert hasattr(tool, 'name'), "Filtered result should be MCPTool with name"
assert hasattr(tool, 'description'), "Filtered result should be MCPTool with description"
filtered_names = [t.name for t in filtered]
print(f"✅ Successfully filtered {len(tools)} tools down to top {len(filtered)}: {filtered_names}")
print(f" Filter respects top_k parameter correctly")
@pytest.mark.asyncio
async def test_semantic_filter_top_k_limiting():
"""
Test that the filter respects top_k parameter.
Given: 20 tools
When: top_k=5
Then: Should return at most 5 tools
"""
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
SemanticMCPToolFilter,
)
# Create 20 tools
tools = [
MCPTool(name=f"tool_{i}", description=f"Tool number {i} for testing", inputSchema={"type": "object"})
for i in range(20)
]
# Mock router
from litellm.types.utils import Embedding, EmbeddingResponse
mock_router = Mock()
def mock_embedding_sync(*args, **kwargs):
return EmbeddingResponse(
data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")],
model="text-embedding-3-small",
object="list",
usage={"prompt_tokens": 10, "total_tokens": 10}
)
async def mock_embedding_async(*args, **kwargs):
return mock_embedding_sync()
mock_router.embedding = mock_embedding_sync
mock_router.aembedding = mock_embedding_async
# Create filter with top_k=5
filter_instance = SemanticMCPToolFilter(
embedding_model="text-embedding-3-small",
litellm_router_instance=mock_router,
top_k=5,
similarity_threshold=0.3,
enabled=True,
)
# Build router with the tools before filtering
filter_instance._build_router(tools)
# Filter tools
filtered = await filter_instance.filter_tools(
query="test query",
available_tools=tools,
)
# Should return at most 5 tools
assert len(filtered) <= 5, f"Expected at most 5 tools, got {len(filtered)}"
print(f"Returned {len(filtered)} tools out of {len(tools)} (top_k=5)")
@pytest.mark.asyncio
async def test_semantic_filter_disabled():
"""
Test that when filter is disabled, all tools are returned.
"""
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
SemanticMCPToolFilter,
)
tools = [
MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"})
for i in range(10)
]
mock_router = Mock()
# Create disabled filter
filter_instance = SemanticMCPToolFilter(
embedding_model="text-embedding-3-small",
litellm_router_instance=mock_router,
top_k=3,
similarity_threshold=0.3,
enabled=False, # Disabled
)
# Filter tools
filtered = await filter_instance.filter_tools(
query="test query",
available_tools=tools,
)
# Should return all tools when disabled
assert len(filtered) == len(tools), f"Expected all {len(tools)} tools, got {len(filtered)}"
@pytest.mark.asyncio
async def test_semantic_filter_empty_tools():
"""
Test that filter handles empty tool list gracefully.
"""
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
SemanticMCPToolFilter,
)
mock_router = Mock()
filter_instance = SemanticMCPToolFilter(
embedding_model="text-embedding-3-small",
litellm_router_instance=mock_router,
top_k=3,
similarity_threshold=0.3,
enabled=True,
)
# Filter empty list
filtered = await filter_instance.filter_tools(
query="test query",
available_tools=[],
)
assert len(filtered) == 0, "Should return empty list for empty input"
@pytest.mark.asyncio
async def test_semantic_filter_extract_user_query():
"""
Test that user query extraction works correctly from messages.
"""
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
SemanticMCPToolFilter,
)
mock_router = Mock()
filter_instance = SemanticMCPToolFilter(
embedding_model="text-embedding-3-small",
litellm_router_instance=mock_router,
top_k=3,
similarity_threshold=0.3,
enabled=True,
)
# Test string content
messages = [
{"role": "system", "content": "You are a helpful assistant"},
{"role": "user", "content": "Send an email to john@example.com"},
]
query = filter_instance.extract_user_query(messages)
assert query == "Send an email to john@example.com"
# Test list content blocks
messages_with_blocks = [
{"role": "user", "content": [
{"type": "text", "text": "Hello, "},
{"type": "text", "text": "send email please"},
]},
]
query2 = filter_instance.extract_user_query(messages_with_blocks)
assert "Hello" in query2 and "send email" in query2
# Test no user messages
messages_no_user = [
{"role": "system", "content": "System message only"},
]
query3 = filter_instance.extract_user_query(messages_no_user)
assert query3 == ""
@pytest.mark.asyncio
async def test_semantic_filter_hook_triggers_on_completion():
"""
Test that the hook triggers for completion requests with tools.
"""
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
SemanticMCPToolFilter,
)
from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook
from litellm.types.utils import Embedding, EmbeddingResponse
# Create mock filter
mock_router = Mock()
def mock_embedding_sync(*args, **kwargs):
return EmbeddingResponse(
data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")],
model="text-embedding-3-small",
object="list",
usage={"prompt_tokens": 10, "total_tokens": 10}
)
async def mock_embedding_async(*args, **kwargs):
return mock_embedding_sync()
mock_router.embedding = mock_embedding_sync
mock_router.aembedding = mock_embedding_async
filter_instance = SemanticMCPToolFilter(
embedding_model="text-embedding-3-small",
litellm_router_instance=mock_router,
top_k=3,
similarity_threshold=0.3,
enabled=True,
)
# Prepare data - completion request with tools
tools = [
MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"})
for i in range(10)
]
# Build router with the tools before filtering
filter_instance._build_router(tools)
# Create hook
hook = SemanticToolFilterHook(filter_instance)
data = {
"model": "gpt-4",
"messages": [
{"role": "user", "content": "Send an email"}
],
"tools": tools,
"metadata": {}, # Hook needs metadata field to store filter stats
}
# Mock user API key dict and cache
mock_user_api_key_dict = Mock()
mock_cache = Mock()
# Call hook
result = await hook.async_pre_call_hook(
user_api_key_dict=mock_user_api_key_dict,
cache=mock_cache,
data=data,
call_type="completion",
)
# Assertions
assert result is not None, "Hook should return modified data"
assert "tools" in result, "Result should contain tools"
assert len(result["tools"]) < len(tools), f"Hook should filter tools, got {len(result['tools'])}/{len(tools)}"
print(f"✅ Hook triggered correctly: {len(tools)} -> {len(result['tools'])} tools")
@pytest.mark.asyncio
async def test_semantic_filter_hook_skips_no_tools():
"""
Test that the hook does NOT trigger when there are no tools.
"""
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
SemanticMCPToolFilter,
)
from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook
# Create mock filter
mock_router = Mock()
filter_instance = SemanticMCPToolFilter(
embedding_model="text-embedding-3-small",
litellm_router_instance=mock_router,
top_k=3,
similarity_threshold=0.3,
enabled=True,
)
# Create hook
hook = SemanticToolFilterHook(filter_instance)
# Prepare data - completion without tools
data = {
"model": "gpt-4",
"messages": [
{"role": "user", "content": "Hello"}
],
}
# Mock user API key dict and cache
mock_user_api_key_dict = Mock()
mock_cache = Mock()
# Call hook
result = await hook.async_pre_call_hook(
user_api_key_dict=mock_user_api_key_dict,
cache=mock_cache,
data=data,
call_type="completion",
)
# Should return None (no modification)
assert result is None, "Hook should skip requests without tools"
print("✅ Hook correctly skips requests without tools")

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