tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-20 20:19:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into litellm_oss_staging_01_27_2026

Browse Source
This commit is contained in:
Sameer Kankute
2026-01-27 17:11:15 +05:30
committed by GitHub
parent 13313ac2be ea0a264a3c
commit c834d7d1fe
125 changed files with 8481 additions and 776 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.circleci/requirements.txt
+2 -1
View File
@@ -16,4 +16,5 @@ uvloop==0.21.0
mcp==1.25.0 # for MCP server
semantic_router==0.1.10 # for auto-routing with litellm
fastuuid==0.12.0
responses==0.25.7 # for proxy client tests
responses==0.25.7 # for proxy client tests
pytest-retry==1.6.3 # for automatic test retries
AGENTS.md
+5 -3
View File
@@ -51,12 +51,14 @@ LiteLLM is a unified interface for 100+ LLMs that:
### MAKING CODE CHANGES FOR THE UI (IGNORE FOR BACKEND)
1. **Use Common Components as much as possible**:
1. **Tremor is DEPRECATED, do not use Tremor components in new features/changes**
- The only exception is the Tremor Table component and its required Tremor Table sub components.
2. **Use Common Components as much as possible**:
- These are usually defined in the `common_components` directory
- Use these components as much as possible and avoid building new components unless needed
- Tremor components are deprecated; prefer using Ant Design (AntD) as much as possible
2. **Testing**:
3. **Testing**:
- The codebase uses **Vitest** and **React Testing Library**
- **Query Priority Order**: Use query methods in this order: `getByRole`, `getByLabelText`, `getByPlaceholderText`, `getByText`, `getByTestId`
- **Always use `screen`** instead of destructuring from `render()` (e.g., use `screen.getByText()` not `getByText`)
Dockerfile
+2 -2
View File
@@ -69,8 +69,8 @@ RUN find /usr/lib -type f -path "*/tornado/test/*" -delete && \
# Convert Windows line endings to Unix and make executable
RUN sed -i 's/\r$//' docker/install_auto_router.sh && chmod +x docker/install_auto_router.sh && ./docker/install_auto_router.sh
# Generate prisma client
RUN prisma generate
# Generate prisma client using the correct schema
RUN prisma generate --schema=./litellm/proxy/schema.prisma
# Convert Windows line endings to Unix for entrypoint scripts
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh
RUN sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
README.md
+1
View File
@@ -267,6 +267,7 @@ Support for more providers. Missing a provider or LLM Platform, raise a [feature
<td><img height="60" alt="Greptile" src="https://github.com/user-attachments/assets/0be4bd8a-7cfa-48d3-9090-f415fe948280" /></td>
<td><img height="60" alt="OpenHands" src="https://github.com/user-attachments/assets/a6150c4c-149e-4cae-888b-8b92be6e003f" /></td>
<td><h2>Netflix</h2></td>
<td><img height="60" alt="OpenAI Agents SDK" src="https://github.com/user-attachments/assets/c02f7be0-8c2e-4d27-aea7-7c024bfaebc0" /></td>
</tr>
</table>
docs/my-website/docs/a2a.md
+115 -1
View File
@@ -68,7 +68,7 @@ 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/) to invoke agents through LiteLLM.
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
@@ -193,6 +193,120 @@ The logs show:
style={{width: '100%', display: 'block', margin: '2rem auto'}}
/>
## Forwarding LiteLLM Context Headers
When LiteLLM invokes your A2A agent, it sends special headers that enable:
- **Trace Grouping**: All LLM calls from the same agent execution appear under one trace
- **Agent Spend Tracking**: Costs are attributed to the specific agent
| Header | Purpose |
|--------|---------|
| `X-LiteLLM-Trace-Id` | Links all LLM calls to the same execution flow |
| `X-LiteLLM-Agent-Id` | Attributes spend to the correct agent |
To enable these features, your A2A server must **forward these headers** to any LLM calls it makes back to LiteLLM.
### Implementation Steps
**Step 1: Extract headers from incoming A2A request**
```python def get_litellm_headers(request) -> dict:
"""Extract X-LiteLLM-* headers from incoming A2A request."""
all_headers = request.call_context.state.get('headers', {})
return {
k: v for k, v in all_headers.items()
if k.lower().startswith('x-litellm-')
}
```
**Step 2: Forward headers to your LLM calls**
Pass the extracted headers when making calls back to LiteLLM:
<Tabs>
<TabItem value="openai" label="OpenAI SDK" default>
```python from openai import OpenAI
headers = get_litellm_headers(request)
client = OpenAI(
api_key="sk-your-litellm-key",
base_url="http://localhost:4000",
default_headers=headers, # Forward headers
)
response = client.chat.completions.create(
model="gpt-4o",
messages=[{"role": "user", "content": "Hello"}]
)
```
</TabItem>
<TabItem value="langchain" label="LangChain">
```python
from langchain_openai import ChatOpenAI
headers = get_litellm_headers(request)
llm = ChatOpenAI(
model="gpt-4o",
openai_api_key="sk-your-litellm-key",
base_url="http://localhost:4000",
default_headers=headers, # Forward headers
)
```
</TabItem>
<TabItem value="litellm" label="LiteLLM SDK">
```python
import litellm
headers = get_litellm_headers(request)
response = litellm.completion(
model="gpt-4o",
messages=[{"role": "user", "content": "Hello"}],
api_base="http://localhost:4000",
extra_headers=headers, # Forward headers
)
```
</TabItem>
<TabItem value="requests" label="HTTP (requests/httpx)">
```python
import httpx
headers = get_litellm_headers(request)
headers["Authorization"] = "Bearer sk-your-litellm-key"
response = httpx.post(
"http://localhost:4000/v1/chat/completions",
headers=headers,
json={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]}
)
```
</TabItem>
</Tabs>
### Result
With header forwarding enabled, you'll see:
**Trace Grouping in Langfuse:**
<Image
img={require('../img/a2a_trace_grouping.png')}
style={{width: '80%', display: 'block', margin: '0', borderRadius: '8px'}}
/>
**Agent Spend Attribution:**
<Image
img={require('../img/a2a_agent_spend.png')}
style={{width: '80%', display: 'block', margin: '0', borderRadius: '8px'}}
/>
## API Reference
### Endpoint
docs/my-website/docs/observability/datadog.md
+50 -2
View File
@@ -7,6 +7,7 @@ import TabItem from '@theme/TabItem';
LiteLLM Supports logging to the following Datdog Integrations:
- `datadog` [Datadog Logs](https://docs.datadoghq.com/logs/)
- `datadog_llm_observability` [Datadog LLM Observability](https://www.datadoghq.com/product/llm-observability/)
- `datadog_cost_management` [Datadog Cloud Cost Management](#datadog-cloud-cost-management)
- `ddtrace-run` [Datadog Tracing](#datadog-tracing)
## Datadog Logs
@@ -73,7 +74,7 @@ Send logs through a local DataDog agent (useful for containerized environments):
```shell
LITELLM_DD_AGENT_HOST="localhost" # hostname or IP of DataDog agent
LITELLM_DD_AGENT_PORT="10518" # [OPTIONAL] port of DataDog agent (default: 10518)
DD_API_KEY="5f2d0f310***********" # [OPTIONAL] your datadog API Key (agent handles auth)
DD_API_KEY="5f2d0f310***********" # [OPTIONAL] your datadog API Key (Agent handles auth for Logs. REQUIRED for LLM Observability)
DD_SOURCE="litellm_dev" # [OPTIONAL] your datadog source
```
@@ -84,6 +85,9 @@ When `LITELLM_DD_AGENT_HOST` is set, logs are sent to the agent instead of direc
**Note:** We use `LITELLM_DD_AGENT_HOST` instead of `DD_AGENT_HOST` to avoid conflicts with `ddtrace` which automatically sets `DD_AGENT_HOST` for APM tracing.
> [!IMPORTANT]
> **Datadog LLM Observability**: `DD_API_KEY` is **REQUIRED** even when using the Datadog Agent (`LITELLM_DD_AGENT_HOST`). The agent acts as a proxy but the API key header is mandatory for the LLM Observability endpoint.
**Step 3**: Start the proxy, make a test request
Start proxy
@@ -161,6 +165,50 @@ On the Datadog LLM Observability page, you should see that both input messages a
<Image img={require('../../img/dd_llm_obs.png')} />
## Datadog Cloud Cost Management
| Feature | Details |
|---------|---------|
| **What is logged** | Aggregated LLM Costs (FOCUS format) |
| **Events** | Periodic Uploads of Aggregated Cost Data |
| **Product Link** | [Datadog Cloud Cost Management](https://docs.datadoghq.com/cost_management/) |
We will use the `--config` to set `litellm.callbacks = ["datadog_cost_management"]`. This will periodically upload aggregated LLM cost data to Datadog.
**Step 1**: Create a `config.yaml` file and set `litellm_settings`: `success_callback`
```yaml
model_list:
- model_name: gpt-3.5-turbo
litellm_params:
model: gpt-3.5-turbo
litellm_settings:
callbacks: ["datadog_cost_management"]
```
**Step 2**: Set Required env variables
```shell
DD_API_KEY="your-api-key"
DD_APP_KEY="your-app-key" # REQUIRED for Cost Management
DD_SITE="us5.datadoghq.com"
```
**Step 3**: Start the proxy
```shell
litellm --config config.yaml
```
**How it works**
* LiteLLM aggregates costs in-memory by Provider, Model, Date, and Tags.
* Requires `DD_APP_KEY` for the Custom Costs API.
* Costs are uploaded periodically (flushed).
### Datadog Tracing
Use `ddtrace-run` to enable [Datadog Tracing](https://ddtrace.readthedocs.io/en/stable/installation_quickstart.html) on litellm proxy
@@ -203,5 +251,5 @@ LiteLLM supports customizing the following Datadog environment variables
| `POD_NAME` | Pod name tag (useful for Kubernetes deployments) | "unknown" | ❌ No |
\* **Required when using Direct API** (default): `DD_API_KEY` and `DD_SITE` are required
\* **Optional when using DataDog Agent**: Set `LITELLM_DD_AGENT_HOST` to use agent mode; `DD_API_KEY` and `DD_SITE` are not required
\* **Optional when using DataDog Agent**: Set `LITELLM_DD_AGENT_HOST` to use agent mode; `DD_API_KEY` and `DD_SITE` are not required for **Datadog Logs**. (**Note: `DD_API_KEY` IS REQUIRED for Datadog LLM Observability**)
docs/my-website/docs/providers/vercel_ai_gateway.md
+34 -2
View File
@@ -11,7 +11,7 @@ import TabItem from '@theme/TabItem';
| Provider Route on LiteLLM | `vercel_ai_gateway/` |
| Link to Provider Doc | [Vercel AI Gateway Documentation ↗](https://vercel.com/docs/ai-gateway) |
| Base URL | `https://ai-gateway.vercel.sh/v1` |
| Supported Operations | `/chat/completions`, `/models` |
| Supported Operations | `/chat/completions`, `/embeddings`, `/models` |
<br />
<br />
@@ -73,7 +73,7 @@ messages = [{"content": "Hello, how are you?", "role": "user"}]
# Vercel AI Gateway call with streaming
response = completion(
model="vercel_ai_gateway/openai/gpt-4o",
model="vercel_ai_gateway/openai/gpt-4o",
messages=messages,
stream=True
)
@@ -82,6 +82,33 @@ for chunk in response:
print(chunk)
```
### Embeddings
```python showLineNumbers title="Vercel AI Gateway Embeddings"
import os
from litellm import embedding
os.environ["VERCEL_AI_GATEWAY_API_KEY"] = "your-api-key"
# Vercel AI Gateway embedding call
response = embedding(
model="vercel_ai_gateway/openai/text-embedding-3-small",
input="Hello world"
)
print(response.data[0]["embedding"][:5]) # Print first 5 dimensions
```
You can also specify the `dimensions` parameter:
```python showLineNumbers title="Vercel AI Gateway Embeddings with Dimensions"
response = embedding(
model="vercel_ai_gateway/openai/text-embedding-3-small",
input=["Hello world", "Goodbye world"],
dimensions=768
)
```
## Usage - LiteLLM Proxy
Add the following to your LiteLLM Proxy configuration file:
@@ -97,6 +124,11 @@ model_list:
litellm_params:
model: vercel_ai_gateway/anthropic/claude-4-sonnet
api_key: os.environ/VERCEL_AI_GATEWAY_API_KEY
- model_name: text-embedding-3-small-gateway
litellm_params:
model: vercel_ai_gateway/openai/text-embedding-3-small
api_key: os.environ/VERCEL_AI_GATEWAY_API_KEY
```
Start your LiteLLM Proxy server:
docs/my-website/docs/proxy/guardrails/onyx_security.md
+3
View File
@@ -128,6 +128,7 @@ guardrails:
mode: ["pre_call", "post_call", "during_call"] # Run at multiple stages
api_key: os.environ/ONYX_API_KEY
api_base: os.environ/ONYX_API_BASE
timeout: 10.0 # Optional, defaults to 10 seconds
```
### Required Parameters
@@ -137,6 +138,7 @@ guardrails:
### Optional Parameters
- **`api_base`**: Onyx API base URL (defaults to `https://ai-guard.onyx.security`)
- **`timeout`**: Request timeout in seconds (defaults to `10.0`)
## Environment Variables
@@ -145,4 +147,5 @@ You can set these environment variables instead of hardcoding values in your con
```shell
export ONYX_API_KEY="your-api-key-here"
export ONYX_API_BASE="https://ai-guard.onyx.security" # Optional
export ONYX_TIMEOUT=10 # Optional, timeout in seconds
```
docs/my-website/img/a2a_agent_spend.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 184 KiB

docs/my-website/img/a2a_trace_grouping.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 388 KiB

docs/my-website/release_notes/v1.81.3-stable/index.md
+423
View File
@@ -0,0 +1,423 @@
---
title: "v1.81.3-stable - Performance - 25% CPU Usage Reduction"
slug: "v1-81-3"
date: 2026-01-26T10:00:00
authors:
- 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
hide_table_of_contents: false
---
import Image from '@theme/IdealImage';
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
## Deploy this version
<Tabs>
<TabItem value="docker" label="Docker">
``` showLineNumbers title="docker run litellm"
docker run \
-e STORE_MODEL_IN_DB=True \
-p 4000:4000 \
docker.litellm.ai/berriai/litellm:v1.81.3.rc.2
```
</TabItem>
<TabItem value="pip" label="Pip">
``` showLineNumbers title="pip install litellm"
pip install litellm==1.81.3.rc.2
```
</TabItem>
</Tabs>
---
## New Models / Updated Models
### New Model Support
| Provider | Model | Context Window | Input ($/1M tokens) | Output ($/1M tokens) | Deprecation Date |
| -------- | ----- | -------------- | ------------------- | -------------------- | ---------------- |
| OpenAI | `gpt-audio`, `gpt-audio-2025-08-28` | 128K | $32/1M audio tokens, $2.5/1M text tokens | $64/1M audio tokens, $10/1M text tokens | - |
| OpenAI | `gpt-audio-mini`, `gpt-audio-mini-2025-08-28` | 128K | $10/1M audio tokens, $0.6/1M text tokens | $20/1M audio tokens, $2.4/1M text tokens | - |
| Deepinfra, Vertex AI, Google AI Studio, OpenRouter, Vercel AI Gateway | `gemini-2.0-flash-001`, `gemini-2.0-flash` | - | - | - | 2026-03-31 |
| Groq | `openai/gpt-oss-120b` | 131K | 0.075/1M cache read | 0.6/1M output tokens | - |
| Groq | `groq/openai/gpt-oss-20b` | 131K | 0.0375/1M cache read, $0.075/1M text tokens | 0.3/1M output tokens | - |
| Vertex AI | `gemini-2.5-computer-use-preview-10-2025` | 128K | $1.25 | $10 | - |
| Azure AI | `claude-haiku-4-5` | $1.25/1M cache read, $2/1M cache read above 1 hr, $0.1/1M text tokens | $5/1M output tokens | - |
| Azure AI | `claude-sonnet-4-5` | $3.75/1M cache read, $6/1M cache read above 1 hr, $3/1M text tokens | $15/1M output tokens | - |
| Azure AI | `claude-opus-4-5` | $6.25/1M cache read, $10/1M cache read above 1 hr, $0.5/1M text tokens | $25/1M output tokens | - |
| Azure AI | `claude-opus-4-1` | $18.75/1M cache read, $30/1M cache read above 1 hr, $1.5/1M text tokens | $75/1M output tokens | - |
### Features
- **[OpenAI](../../docs/providers/openai)**
- Add gpt-audio and gpt-audio-mini models to pricing - [PR #19509](https://github.com/BerriAI/litellm/pull/19509)
- correct audio token costs for gpt-4o-audio-preview models - [PR #19500](https://github.com/BerriAI/litellm/pull/19500)
- Limit stop sequence as per openai spec (ensures JetBrains IDE compatibility) - [PR #19562](https://github.com/BerriAI/litellm/pull/19562)
- **[VertexAI](../../docs/providers/vertex)**
- Docs - Google Workload Identity Federation (WIF) support - [PR #19320](https://github.com/BerriAI/litellm/pull/19320)
- **[Agentcore](../../docs/providers/bedrock_agentcore)**
- Fixes streaming issues with AWS Bedrock AgentCore where responses would stop after the first chunk, particularly affecting OAuth-enabled agents - [PR #17141](https://github.com/BerriAI/litellm/pull/17141)
- **[Chatgpt](../../docs/providers/chatgpt)**
- Adds support for calling chatgpt subscription via LiteLLM - [PR #19030](https://github.com/BerriAI/litellm/pull/19030)
- Adds responses API bridge support for chatgpt subscription provider - [PR #19030](https://github.com/BerriAI/litellm/pull/19030)
- **[Bedrock](../../docs/providers/bedrock)**
- support for output format for bedrock invoke via v1/messages - [PR #19560](https://github.com/BerriAI/litellm/pull/19560)
- **[Azure](../../docs/providers/azure/azure)**
- Add support for Azure OpenAI v1 API - [PR #19313](https://github.com/BerriAI/litellm/pull/19313)
- preserve content_policy_violation details for images (#19328) - [PR #19372](https://github.com/BerriAI/litellm/pull/19372)
- Support OpenAI-format nested tool definitions for Responses API - [PR #19526](https://github.com/BerriAI/litellm/pull/19526)
- **Gemini([Vertex AI](../../docs/providers/vertex), [Google AI Studio](../../docs/providers/gemini))**
- use responseJsonSchema for Gemini 2.0+ models - [PR #19314](https://github.com/BerriAI/litellm/pull/19314)
- **[Volcengine](../../docs/providers/volcano)**
- Support Volcengine responses api - [PR #18508](https://github.com/BerriAI/litellm/pull/18508)
- **[Anthropic](../../docs/providers/anthropic)**
- Add Support for calling Claude Code Max subscriptions via LiteLLM - [PR #19453](https://github.com/BerriAI/litellm/pull/19453)
- Add Structured output for /v1/messages with Anthropic API, Azure Anthropic API, Bedrock Converse - [PR #19545](https://github.com/BerriAI/litellm/pull/19545)
- **[Brave Search](../../docs/search/brave)**
- New Search provider - [PR #19433](https://github.com/BerriAI/litellm/pull/19433)
- **Sarvam ai**
- Add support for new sarvam models - [PR #19479](https://github.com/BerriAI/litellm/pull/19479)
- **[GMI](../../docs/providers/gmi)**
- add GMI Cloud provider support - [PR #19376](https://github.com/BerriAI/litellm/pull/19376)
### Bug Fixes
- **[Anthropic](../../docs/providers/anthropic)**
- Fix anthropic-beta sent client side being overridden instead of appended to - [PR #19343](https://github.com/BerriAI/litellm/pull/19343)
- Filter out unsupported fields from JSON schema for Anthropic's output_format API - [PR #19482](https://github.com/BerriAI/litellm/pull/19482)
- **[Bedrock](../../docs/providers/bedrock)**
- Expose stability models via /image_edits endpoint and ensure proper request transformation - [PR #19323](https://github.com/BerriAI/litellm/pull/19323)
- Claude Code x Bedrock Invoke fails with advanced-tool-use-2025-11-20 - [PR #19373](https://github.com/BerriAI/litellm/pull/19373)
- deduplicate tool calls in assistant history - [PR #19324](https://github.com/BerriAI/litellm/pull/19324)
- fix: correct us.anthropic.claude-opus-4-5 In-region pricing - [PR #19310](https://github.com/BerriAI/litellm/pull/19310)
- Fix request validation errors when using Claude 4 via bedrock invoke - [PR #19381](https://github.com/BerriAI/litellm/pull/19381)
- Handle thinking with tool calls for Claude 4 models - [PR #19506](https://github.com/BerriAI/litellm/pull/19506)
- correct streaming choice index for tool calls - [PR #19506](https://github.com/BerriAI/litellm/pull/19506)
- **[Ollama](../../docs/providers/ollama)**
- Fix tool call errors due with improved message extraction - [PR #19369](https://github.com/BerriAI/litellm/pull/19369)
- **[VertexAI](../../docs/providers/vertex)**
- Removed optional vertex_count_tokens_location param before request is sent to vertex - [PR #19359](https://github.com/BerriAI/litellm/pull/19359)
- **Gemini([Vertex AI](../../docs/providers/vertex), [Google AI Studio](../../docs/providers/gemini))**
- Supports setting media_resolution and fps parameters on each video file, when using Gemini video understanding - [PR #19273](https://github.com/BerriAI/litellm/pull/19273)
- handle reasoning_effort as dict from OpenAI Agents SDK - [PR #19419](https://github.com/BerriAI/litellm/pull/19419)
- add file content support in tool results - [PR #19416](https://github.com/BerriAI/litellm/pull/19416)
- **[Azure](../../docs/providers/azure_ai)**
- Fix Azure AI costs for Anthropic models - [PR #19530](https://github.com/BerriAI/litellm/pull/19530)
- **[Giga Chat](../../docs/providers/gigachat)**
- Add tool choice mapping - [PR #19645](https://github.com/BerriAI/litellm/pull/19645)
---
## AI API Endpoints (LLMs, MCP, Agents)
### Features
- **[Files API](../../docs/files_endpoints)**
- Add managed files support when load_balancing is True - [PR #19338](https://github.com/BerriAI/litellm/pull/19338)
- **[Claude Plugin Marketplace](../../docs/tutorials/claude_code_plugin_marketplace)**
- Add self hosted Claude Code Plugin Marketplace - [PR #19378](https://github.com/BerriAI/litellm/pull/19378)
- **[MCP](../../docs/mcp)**
- Add MCP Protocol version 2025-11-25 support - [PR #19379](https://github.com/BerriAI/litellm/pull/19379)
- Log MCP tool calls and list tools in the LiteLLM Spend Logs table for easier debugging - [PR #19469](https://github.com/BerriAI/litellm/pull/19469)
- **[Vertex AI](../../docs/providers/vertex)**
- Ensure only anthropic betas are forwarded down to LLM API (by default) - [PR #19542](https://github.com/BerriAI/litellm/pull/19542)
- Allow overriding to support forwarding incoming headers are forwarded down to target - [PR #19524](https://github.com/BerriAI/litellm/pull/19524)
- **[Chat/Completions](../../docs/completion/input)**
- Add MCP tools response to chat completions - [PR #19552](https://github.com/BerriAI/litellm/pull/19552)
- Add custom vertex ai finish reasons to the output - [PR #19558](https://github.com/BerriAI/litellm/pull/19558)
- Return MCP execution in /chat/completions before model output during streaming - [PR #19623](https://github.com/BerriAI/litellm/pull/19623)
### Bugs
- **[Responses API](../../docs/response_api)**
- Fix duplicate messages during MCP streaming tool execution - [PR #19317](https://github.com/BerriAI/litellm/pull/19317)
- Fix pickle error when using OpenAI's Responses API with stream=True and tool_choice of type allowed_tools (an OpenAI-native parameter) - [PR #17205](https://github.com/BerriAI/litellm/pull/17205)
- stream tool call events for non-openai models - [PR #19368](https://github.com/BerriAI/litellm/pull/19368)
- preserve tool output ordering for gemini in responses bridge - [PR #19360](https://github.com/BerriAI/litellm/pull/19360)
- Add ID caching to prevent ID mismatch text-start and text-delta - [PR #19390](https://github.com/BerriAI/litellm/pull/19390)
- Include output_item, reasoning_summary_Text_done and reasoning_summary_part_done events for non-openai models - [PR #19472](https://github.com/BerriAI/litellm/pull/19472)
- **[Chat/Completions](../../docs/completion/input)**
- fix: drop_params not dropping prompt_cache_key for non-OpenAI providers - [PR #19346](https://github.com/BerriAI/litellm/pull/19346)
- **[Realtime API](../../docs/realtime)**
- disable SSL for ws:// WebSocket connections - [PR #19345](https://github.com/BerriAI/litellm/pull/19345)
- **[Generate Content](../../docs/generateContent)**
- Log actual user input when google genai/vertex endpoints are called client-side - [PR #19156](https://github.com/BerriAI/litellm/pull/19156)
- **[/messages/count_tokens Anthropic Token Counting](../../docs/anthropic_count_tokens)**
- ensure it works for Anthropic, Azure AI Anthropic on AI Gateway - [PR #19432](https://github.com/BerriAI/litellm/pull/19432)
- **[MCP](../../docs/mcp)**
- forward static_headers to MCP servers - [PR #19366](https://github.com/BerriAI/litellm/pull/19366)
- **[Batch API](../../docs/batches)**
- Fix: generation config empty for batch - [PR #19556](https://github.com/BerriAI/litellm/pull/19556)
- **[Pass Through Endpoints](../../docs/proxy/pass_through)**
- Always reupdate registry - [PR #19420](https://github.com/BerriAI/litellm/pull/19420)
---
## Management Endpoints / UI
### Features
- **Cost Estimator**
- Fix model dropdown - [PR #19529](https://github.com/BerriAI/litellm/pull/19529)
- **Claude Code Plugins**
- Allow Adding Claude Code Plugins via UI - [PR #19387](https://github.com/BerriAI/litellm/pull/19387)
- **Guardrails**
- New Policy management UI - [PR #19668](https://github.com/BerriAI/litellm/pull/19668)
- Allow adding policies on Keys/Teams + Viewing on Info panels - [PR #19688](https://github.com/BerriAI/litellm/pull/19688)
- **General**
- respects custom authentication header override - [PR #19276](https://github.com/BerriAI/litellm/pull/19276)
- **Playground**
- Button to Fill Custom API Base - [PR #19440](https://github.com/BerriAI/litellm/pull/19440)
- display mcp output on the play ground - [PR #19553](https://github.com/BerriAI/litellm/pull/19553)
- **Models**
- Paginate /v2/models/info - [PR #19521](https://github.com/BerriAI/litellm/pull/19521)
- All Model Tab Pagination - [PR #19525](https://github.com/BerriAI/litellm/pull/19525)
- Adding Optional scope Param to /models - [PR #19539](https://github.com/BerriAI/litellm/pull/19539)
- Model Search - [PR #19622](https://github.com/BerriAI/litellm/pull/19622)
- Filter by Model ID and Team ID - [PR #19713](https://github.com/BerriAI/litellm/pull/19713)
- **MCP Servers**
- MCP Tools Tab Resetting to Overview - [PR #19468](https://github.com/BerriAI/litellm/pull/19468)
- **Organizations**
- Prevent org admin from creating a new user with proxy_admin permissions - [PR #19296](https://github.com/BerriAI/litellm/pull/19296)
- Edit Page: Reusable Model Select - [PR #19601](https://github.com/BerriAI/litellm/pull/19601)
- **Teams**
- Reusable Model Select - [PR #19543](https://github.com/BerriAI/litellm/pull/19543)
- [Fix] Team Update with Organization having All Proxy Models - [PR #19604](https://github.com/BerriAI/litellm/pull/19604)
- **Logs**
- Include tool arguments in spend logs table - [PR #19640](https://github.com/BerriAI/litellm/pull/19640)
- **Fallbacks / Loadbalancing**
- New fallbacks modal - [PR #19673](https://github.com/BerriAI/litellm/pull/19673)
- Set fallbacks/loadbalancing by team/key - [PR #19686](https://github.com/BerriAI/litellm/pull/19686)
### Bugs
- **Playground**
- increase model selector width in playground Compare view - [PR #19423](https://github.com/BerriAI/litellm/pull/19423)
- **Virtual Keys**
- Sorting Shows Incorrect Entries - [PR #19534](https://github.com/BerriAI/litellm/pull/19534)
- **General**
- UI 404 error when SERVER_ROOT_PATH is set - [PR #19467](https://github.com/BerriAI/litellm/pull/19467)
- Redirect to ui/login on expired JWT - [PR #19687](https://github.com/BerriAI/litellm/pull/19687)
- **SSO**
- Fix SSO user roles not updating for existing users - [PR #19621](https://github.com/BerriAI/litellm/pull/19621)
- **Guardrails**
- ensure guardrail patterns persist on edit and mode toggle - [PR #19265](https://github.com/BerriAI/litellm/pull/19265)
---
## AI Integrations
### Logging
- **General Logging**
- prevent printing duplicate StandardLoggingPayload logs - [PR #19325](https://github.com/BerriAI/litellm/pull/19325)
- Fix: log duplication when json_logs is enabled - [PR #19705](https://github.com/BerriAI/litellm/pull/19705)
- **Langfuse OTEL**
- ignore service logs and fix callback shadowing - [PR #19298](https://github.com/BerriAI/litellm/pull/19298)
- **Langfuse**
- Send litellm_trace_id - [PR #19528](https://github.com/BerriAI/litellm/pull/19528)
- Add Langfuse mock mode for testing without API calls - [PR #19676](https://github.com/BerriAI/litellm/pull/19676)
- **GCS Bucket**
- prevent unbounded queue growth due to slow API calls - [PR #19297](https://github.com/BerriAI/litellm/pull/19297)
- Add GCS mock mode for testing without API calls - [PR #19683](https://github.com/BerriAI/litellm/pull/19683)
- **Responses API Logging**
- Fix pydantic serialization error - [PR #19486](https://github.com/BerriAI/litellm/pull/19486)
- **Arize Phoenix**
- add openinference span kinds to arize phoenix - [PR #19267](https://github.com/BerriAI/litellm/pull/19267)
- **Prometheus**
- Added new prometheus metrics for user count and team count - [PR #19520](https://github.com/BerriAI/litellm/pull/19520)
### Guardrails
- **Bedrock Guardrails**
- Ensure post_call guardrail checks input+output - [PR #19151](https://github.com/BerriAI/litellm/pull/19151)
- **Prompt Security**
- fixing prompt-security's guardrail implementation - [PR #19374](https://github.com/BerriAI/litellm/pull/19374)
- **Presidio**
- Fixes crash in Presidio Guardrail when running in background threads (logging_hook) - [PR #19714](https://github.com/BerriAI/litellm/pull/19714)
- **Pillar Security**
- Migrate Pillar Security to Generic Guardrail API - [PR #19364](https://github.com/BerriAI/litellm/pull/19364)
- **Policy Engine**
- New LiteLLM Policy engine - create policies to manage guardrails, conditions - permissions per Key, Team - [PR #19612](https://github.com/BerriAI/litellm/pull/19612)
- **General**
- add case-insensitive support for guardrail mode and actions - [PR #19480](https://github.com/BerriAI/litellm/pull/19480)
### Prompt Management
- **General**
- fix prompt info lookup and delete using correct IDs - [PR #19358](https://github.com/BerriAI/litellm/pull/19358)
### Secret Manager
- **AWS Secret Manager**
- ensure auto-rotation updates existing AWS secret instead of creating new one - [PR #19455](https://github.com/BerriAI/litellm/pull/19455)
- **Hashicorp Vault**
- Ensure key rotations work with Vault - [PR #19634](https://github.com/BerriAI/litellm/pull/19634)
---
## Spend Tracking, Budgets and Rate Limiting
- **Pricing Updates**
- Add openai/dall-e base pricing entries - [PR #19133](https://github.com/BerriAI/litellm/pull/19133)
- Add `input_cost_per_video_per_second` in ModelInfoBase - [PR #19398](https://github.com/BerriAI/litellm/pull/19398)
---
## Performance / Loadbalancing / Reliability improvements
- **General**
- Fix date overflow/division by zero in proxy utils - [PR #19527](https://github.com/BerriAI/litellm/pull/19527)
- Fix in-flight request termination on SIGTERM when health-check runs in a separate process - [PR #19427](https://github.com/BerriAI/litellm/pull/19427)
- Fix Pass through routes to work with server root path - [PR #19383](https://github.com/BerriAI/litellm/pull/19383)
- Fix logging error for stop iteration - [PR #19649](https://github.com/BerriAI/litellm/pull/19649)
- prevent retrying 4xx client errors - [PR #19275](https://github.com/BerriAI/litellm/pull/19275)
- add better error handling for misconfig on health check - [PR #19441](https://github.com/BerriAI/litellm/pull/19441)
- **Router**
- Fix Azure RPM calculation formula - [PR #19513](https://github.com/BerriAI/litellm/pull/19513)
- Persist scheduler request queue to redis - [PR #19304](https://github.com/BerriAI/litellm/pull/19304)
- pass search_tools to Router during DB-triggered initialization - [PR #19388](https://github.com/BerriAI/litellm/pull/19388)
- Fixed PromptCachingCache to correctly handle messages where cache_control is a sibling key of string content - [PR #19266](https://github.com/BerriAI/litellm/pull/19266)
- **Memory Leaks/OOM**
- prevent OOM with nested $defs in tool schemas - [PR #19112](https://github.com/BerriAI/litellm/pull/19112)
- fix: HTTP client memory leaks in Presidio, OpenAI, and Gemini - [PR #19190](https://github.com/BerriAI/litellm/pull/19190)
- **Non root**
- fix logfile and pidfile of supervisor for non root environment - [PR #17267](https://github.com/BerriAI/litellm/pull/17267)
- resolve Read-only file system error in non-root images - [PR #19449](https://github.com/BerriAI/litellm/pull/19449)
- **Dockerfile**
- Redis Semantic Caching - add missing redisvl dependency to requirements.txt - [PR #19417](https://github.com/BerriAI/litellm/pull/19417)
- Bump OTEL versions to support a2a dependency - resolves modulenotfounderror for Microsoft Agents by @Harshit28j in #18991
- **DB**
- Handle PostgreSQL cached plan errors during rolling deployments - [PR #19424](https://github.com/BerriAI/litellm/pull/19424)
- **Timeouts**
- Fix: total timeout is not respected - [PR #19389](https://github.com/BerriAI/litellm/pull/19389)
- **SDK**
- Field-Existence Checks to Type Classes to Prevent Attribute Errors - [PR #18321](https://github.com/BerriAI/litellm/pull/18321)
- add google-cloud-aiplatform as optional dependency with clear error message - [PR #19437](https://github.com/BerriAI/litellm/pull/19437)
- Make grpc dependency optional - [PR #19447](https://github.com/BerriAI/litellm/pull/19447)
- Add support for retry policies - [PR #19645](https://github.com/BerriAI/litellm/pull/19645)
- **Performance**
- Cut chat_completion latency by ~21% by reducing pre-call processing time - [PR #19535](https://github.com/BerriAI/litellm/pull/19535)
- Optimize strip_trailing_slash with O(1) index check - [PR #19679](https://github.com/BerriAI/litellm/pull/19679)
- Optimize use_custom_pricing_for_model with set intersection - [PR #19677](https://github.com/BerriAI/litellm/pull/19677)
- perf: skip pattern_router.route() for non-wildcard models - [PR #19664](https://github.com/BerriAI/litellm/pull/19664)
- perf: Add LRU caching to get_model_info for faster cost lookups - [PR #19606](https://github.com/BerriAI/litellm/pull/19606)
---
## General Proxy Improvements
### Doc Improvements
- new tutorial for adding MCPs to Cursor via LiteLLM - [PR #19317](https://github.com/BerriAI/litellm/pull/19317)
- fix vertex_region to vertex_location in Vertex AI pass-through docs - [PR #19380](https://github.com/BerriAI/litellm/pull/19380)
- clarify Gemini and Vertex AI model prefix in json file - [PR #19443](https://github.com/BerriAI/litellm/pull/19443)
- update Claude Code integration guides - [PR #19415](https://github.com/BerriAI/litellm/pull/19415)
- adjust opencode tutorial - [PR #19605](https://github.com/BerriAI/litellm/pull/19605)
- add spend-queue-troubleshooting docs - [PR #19659](https://github.com/BerriAI/litellm/pull/19659)
- docs: add litellm-enterprise requirement for managed files - [PR #19689](https://github.com/BerriAI/litellm/pull/19689)
### Helm
- Add support for keda in helm chart - [PR #19337](https://github.com/BerriAI/litellm/pull/19337)
- sync Helm chart version with LiteLLM release version - [PR #19438](https://github.com/BerriAI/litellm/pull/19438)
- Enable PreStop hook configuration in values.yaml - [PR #19613](https://github.com/BerriAI/litellm/pull/19613)
### General
- Add health check scripts and parallel execution support - [PR #19295](https://github.com/BerriAI/litellm/pull/19295)
---
## New Contributors
* @dushyantzz made their first contribution in [PR #19158](https://github.com/BerriAI/litellm/pull/19158)
* @obod-mpw made their first contribution in [PR #19133](https://github.com/BerriAI/litellm/pull/19133)
* @msexxeta made their first contribution in [PR #19030](https://github.com/BerriAI/litellm/pull/19030)
* @rsicart made their first contribution in [PR #19337](https://github.com/BerriAI/litellm/pull/19337)
* @cluebbehusen made their first contribution in [PR #19311](https://github.com/BerriAI/litellm/pull/19311)
* @Lucky-Lodhi2004 made their first contribution in [PR #19315](https://github.com/BerriAI/litellm/pull/19315)
* @binbandit made their first contribution in [PR #19324](https://github.com/BerriAI/litellm/pull/19324)
* @flex-myeonghyeon made their first contribution in [PR #19381](https://github.com/BerriAI/litellm/pull/19381)
* @Lrakotoson made their first contribution in [PR #18321](https://github.com/BerriAI/litellm/pull/18321)
* @bensi94 made their first contribution in [PR #18787](https://github.com/BerriAI/litellm/pull/18787)
* @victorigualada made their first contribution in [PR #19368](https://github.com/BerriAI/litellm/pull/19368)
* @VedantMadane made their first contribution in #19266
* @stiyyagura0901 made their first contribution in #19276
* @kamilio made their first contribution in [PR #19447](https://github.com/BerriAI/litellm/pull/19447)
* @jonathansampson made their first contribution in [PR #19433](https://github.com/BerriAI/litellm/pull/19433)
* @rynecarbone made their first contribution in [PR #19416](https://github.com/BerriAI/litellm/pull/19416)
* @jayy-77 made their first contribution in #19366
* @davida-ps made their first contribution in [PR #19374](https://github.com/BerriAI/litellm/pull/19374)
* @joaodinissf made their first contribution in [PR #19506](https://github.com/BerriAI/litellm/pull/19506)
* @ecao310 made their first contribution in [PR #19520](https://github.com/BerriAI/litellm/pull/19520)
* @mpcusack-altos made their first contribution in [PR #19577](https://github.com/BerriAI/litellm/pull/19577)
* @milan-berri made their first contribution in [PR #19602](https://github.com/BerriAI/litellm/pull/19602)
* @xqe2011 made their first contribution in #19621
---
## Full Changelog
**[View complete changelog on GitHub](https://github.com/BerriAI/litellm/releases/tag/v1.81.3.rc)**
litellm/a2a_protocol/main.py
+11 -2
View File
@@ -9,7 +9,7 @@ import datetime
from typing import TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, Optional, Union
import litellm
from litellm._logging import verbose_logger
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator
from litellm.a2a_protocol.utils import A2ARequestUtils
from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT
@@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.types.agents import LiteLLMSendMessageResponse
from litellm.utils import client
import uuid
if TYPE_CHECKING:
from a2a.client import A2AClient as A2AClientType
@@ -225,7 +226,11 @@ async def asend_message(
raise ValueError(
"Either a2a_client or api_base is required for standard A2A flow"
)
a2a_client = await create_a2a_client(base_url=api_base)
trace_id = str(uuid.uuid4())
extra_headers = {"X-LiteLLM-Trace-Id": trace_id}
if agent_id:
extra_headers["X-LiteLLM-Agent-Id"] = agent_id
a2a_client = await create_a2a_client(base_url=api_base, extra_headers=extra_headers)
# Type assertion: a2a_client is guaranteed to be non-None here
assert a2a_client is not None
@@ -490,6 +495,10 @@ async def create_a2a_client(
)
httpx_client = http_handler.client
if extra_headers:
httpx_client.headers.update(extra_headers)
verbose_proxy_logger.debug(f"A2A client created with extra_headers={extra_headers}")
# Resolve agent card
resolver = A2ACardResolver(
httpx_client=httpx_client,
litellm/experimental_mcp_client/client.py
+65 -65
View File
@@ -4,13 +4,17 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
import asyncio
import base64
from typing import Any, Awaitable, Callable, Dict, List, Optional, TypeVar, Union
from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple, TypeVar, Union
import httpx
from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import streamable_http_client
try:
from mcp.client.streamable_http import streamable_http_client # type: ignore
except ImportError:
streamable_http_client = None
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
from mcp.types import CallToolResult as MCPCallToolResult
from mcp.types import (
@@ -76,104 +80,100 @@ class MCPClient:
def _create_transport_context(
self,
) -> tuple[Any, Optional[httpx.AsyncClient]]:
"""Create the appropriate transport context based on transport type."""
) -> Tuple[Any, Optional[httpx.AsyncClient]]:
"""
Create the appropriate transport context based on transport type.
Returns:
Tuple of (transport_context, http_client).
http_client is only set for HTTP transport and needs cleanup.
"""
http_client: Optional[httpx.AsyncClient] = None
if self.transport_type == MCPTransport.stdio:
if not self.stdio_config:
raise ValueError("stdio_config is required for stdio transport")
server_params = StdioServerParameters(
command=self.stdio_config.get("command", ""),
args=self.stdio_config.get("args", []),
env=self.stdio_config.get("env", {}),
)
transport_ctx = stdio_client(server_params)
elif self.transport_type == MCPTransport.sse:
return stdio_client(server_params), None
if self.transport_type == MCPTransport.sse:
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
transport_ctx = sse_client(
return sse_client(
url=self.server_url,
timeout=self.timeout,
headers=headers,
httpx_client_factory=httpx_client_factory,
)
else:
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
verbose_logger.debug(
"litellm headers for streamable_http_client: %s", headers
)
http_client = httpx_client_factory(
headers=headers,
timeout=httpx.Timeout(self.timeout),
)
transport_ctx = streamable_http_client(
url=self.server_url,
http_client=http_client,
)
if transport_ctx is None:
raise RuntimeError("Failed to create transport context")
), None
# HTTP transport (default)
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
verbose_logger.debug(
"litellm headers for streamable_http_client: %s", headers
)
http_client = httpx_client_factory(
headers=headers,
timeout=httpx.Timeout(self.timeout),
)
transport_ctx = streamable_http_client(
url=self.server_url,
http_client=http_client,
)
return transport_ctx, http_client
async def _execute_session_operation(
self,
transport_ctx: Any,
operation: Callable[[ClientSession], Awaitable[TSessionResult]],
) -> TSessionResult:
"""
Execute an operation within a transport and session context.
Handles entering/exiting contexts and running the operation.
"""
transport = await transport_ctx.__aenter__()
try:
read_stream, write_stream = transport[0], transport[1]
session_ctx = ClientSession(read_stream, write_stream)
session = await session_ctx.__aenter__()
try:
await session.initialize()
return await operation(session)
finally:
try:
await session_ctx.__aexit__(None, None, None)
except BaseException as e:
verbose_logger.debug(f"Error during session context exit: {e}")
finally:
try:
await transport_ctx.__aexit__(None, None, None)
except BaseException as e:
verbose_logger.debug(f"Error during transport context exit: {e}")
async def run_with_session(
self, operation: Callable[[ClientSession], Awaitable[TSessionResult]]
) -> TSessionResult:
"""Open a session, run the provided coroutine, and clean up."""
transport_ctx = None
http_client: Optional[httpx.AsyncClient] = None
session_ctx = None
try:
transport_ctx, http_client = self._create_transport_context()
# Enter transport context
transport = await transport_ctx.__aenter__()
try:
read_stream, write_stream = transport[0], transport[1]
session_ctx = ClientSession(read_stream, write_stream)
# Enter session context
session = await session_ctx.__aenter__()
try:
await session.initialize()
result = await operation(session)
return result
finally:
# Ensure session context is properly exited
if session_ctx is not None:
try:
await session_ctx.__aexit__(None, None, None)
except Exception as e:
verbose_logger.debug(
f"Error during session context exit: {e}"
)
finally:
# Ensure transport context is properly exited
if transport_ctx is not None:
try:
await transport_ctx.__aexit__(None, None, None)
except Exception as e:
verbose_logger.debug(
f"Error during transport context exit: {e}"
)
return await self._execute_session_operation(transport_ctx, operation)
except Exception:
verbose_logger.warning(
"MCP client run_with_session failed for %s", self.server_url or "stdio"
)
raise
finally:
# Always clean up http_client if it was created
if http_client is not None:
try:
await http_client.aclose()
except Exception as e:
verbose_logger.debug(
f"Error during http_client cleanup: {e}"
)
except BaseException as e:
verbose_logger.debug(f"Error during http_client cleanup: {e}")
def update_auth_value(self, mcp_auth_value: Union[str, Dict[str, str]]):
"""
litellm/integrations/callback_configs.json
+28 -1
View File
@@ -83,6 +83,33 @@
},
"description": "Datadog Logging Integration"
},
{
"id": "datadog_cost_management",
"displayName": "Datadog Cost Management",
"logo": "datadog.png",
"supports_key_team_logging": false,
"dynamic_params": {
"dd_api_key": {
"type": "password",
"ui_name": "API Key",
"description": "Datadog API key for authentication",
"required": true
},
"dd_app_key": {
"type": "password",
"ui_name": "App Key",
"description": "Datadog Application Key for Cloud Cost Management",
"required": true
},
"dd_site": {
"type": "text",
"ui_name": "Site",
"description": "Datadog site URL (e.g., us5.datadoghq.com)",
"required": true
}
},
"description": "Datadog Cloud Cost Management Integration"
},
{
"id": "lago",
"displayName": "Lago",
@@ -407,4 +434,4 @@
},
"description": "SQS Queue (AWS) Logging Integration"
}
]
]
litellm/integrations/custom_guardrail.py
+14 -2
View File
@@ -516,7 +516,9 @@ class CustomGuardrail(CustomLogger):
from litellm.types.utils import GuardrailMode
# Use event_type if provided, otherwise fall back to self.event_hook
guardrail_mode: Union[GuardrailEventHooks, GuardrailMode, List[GuardrailEventHooks]]
guardrail_mode: Union[
GuardrailEventHooks, GuardrailMode, List[GuardrailEventHooks]
]
if event_type is not None:
guardrail_mode = event_type
elif isinstance(self.event_hook, Mode):
@@ -524,11 +526,21 @@ class CustomGuardrail(CustomLogger):
else:
guardrail_mode = self.event_hook # type: ignore[assignment]
from litellm.litellm_core_utils.core_helpers import (
filter_exceptions_from_params,
)
# Sanitize the response to ensure it's JSON serializable and free of circular refs
# This prevents RecursionErrors in downstream loggers (Langfuse, Datadog, etc.)
clean_guardrail_response = filter_exceptions_from_params(
guardrail_json_response
)
slg = StandardLoggingGuardrailInformation(
guardrail_name=self.guardrail_name,
guardrail_provider=guardrail_provider,
guardrail_mode=guardrail_mode,
guardrail_response=guardrail_json_response,
guardrail_response=clean_guardrail_response,
guardrail_status=guardrail_status,
start_time=start_time,
end_time=end_time,
litellm/integrations/datadog/datadog.py
+4 -13
View File
@@ -32,6 +32,7 @@ from litellm.integrations.datadog.datadog_handler import (
get_datadog_service,
get_datadog_source,
get_datadog_tags,
get_datadog_base_url_from_env,
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.llms.custom_httpx.http_handler import (
@@ -100,7 +101,9 @@ class DataDogLogger(
self._configure_dd_direct_api()
# Optional override for testing
self._apply_dd_base_url_override()
dd_base_url = get_datadog_base_url_from_env()
if dd_base_url:
self.intake_url = f"{dd_base_url}/api/v2/logs"
self.sync_client = _get_httpx_client()
asyncio.create_task(self.periodic_flush())
self.flush_lock = asyncio.Lock()
@@ -159,18 +162,6 @@ class DataDogLogger(
self.DD_API_KEY = os.getenv("DD_API_KEY")
self.intake_url = f"https://http-intake.logs.{os.getenv('DD_SITE')}/api/v2/logs"
def _apply_dd_base_url_override(self) -> None:
"""
Apply base URL override for testing purposes
"""
dd_base_url: Optional[str] = (
os.getenv("_DATADOG_BASE_URL")
or os.getenv("DATADOG_BASE_URL")
or os.getenv("DD_BASE_URL")
)
if dd_base_url is not None:
self.intake_url = f"{dd_base_url}/api/v2/logs"
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
"""
Async Log success events to Datadog
litellm/integrations/datadog/datadog_cost_management.py
+204
View File
@@ -0,0 +1,204 @@
import asyncio
import os
import time
from datetime import datetime
from typing import Dict, List, Optional, Tuple
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.integrations.datadog_cost_management import (
DatadogFOCUSCostEntry,
)
from litellm.types.utils import StandardLoggingPayload
class DatadogCostManagementLogger(CustomBatchLogger):
def __init__(self, **kwargs):
self.dd_api_key = os.getenv("DD_API_KEY")
self.dd_app_key = os.getenv("DD_APP_KEY")
self.dd_site = os.getenv("DD_SITE", "datadoghq.com")
if not self.dd_api_key or not self.dd_app_key:
verbose_logger.warning(
"Datadog Cost Management: DD_API_KEY and DD_APP_KEY are required. Integration will not work."
)
self.upload_url = f"https://api.{self.dd_site}/api/v2/cost/custom_costs"
self.async_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
# Initialize lock and start periodic flush task
self.flush_lock = asyncio.Lock()
asyncio.create_task(self.periodic_flush())
# Check if flush_lock is already in kwargs to avoid double passing (unlikely but safe)
if "flush_lock" not in kwargs:
kwargs["flush_lock"] = self.flush_lock
super().__init__(**kwargs)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
try:
standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get(
"standard_logging_object", None
)
if standard_logging_object is None:
return
# Only log if there is a cost associated
if standard_logging_object.get("response_cost", 0) > 0:
self.log_queue.append(standard_logging_object)
if len(self.log_queue) >= self.batch_size:
await self.async_send_batch()
except Exception as e:
verbose_logger.exception(
f"Datadog Cost Management: Error in async_log_success_event: {str(e)}"
)
async def async_send_batch(self):
if not self.log_queue:
return
try:
# Aggregate costs from the batch
aggregated_entries = self._aggregate_costs(self.log_queue)
if not aggregated_entries:
return
# Send to Datadog
await self._upload_to_datadog(aggregated_entries)
# Clear queue only on success (or if we decide to drop on failure)
# CustomBatchLogger clears queue in flush_queue, so we just process here
except Exception as e:
verbose_logger.exception(
f"Datadog Cost Management: Error in async_send_batch: {str(e)}"
)
def _aggregate_costs(
self, logs: List[StandardLoggingPayload]
) -> List[DatadogFOCUSCostEntry]:
"""
Aggregates costs by Provider, Model, and Date.
Returns a list of DatadogFOCUSCostEntry.
"""
aggregator: Dict[Tuple[str, str, str, Tuple[Tuple[str, str], ...]], DatadogFOCUSCostEntry] = {}
for log in logs:
try:
# Extract keys for aggregation
provider = log.get("custom_llm_provider") or "unknown"
model = log.get("model") or "unknown"
cost = log.get("response_cost", 0)
if cost == 0:
continue
# Get date strings (FOCUS format requires specific keys, but for aggregation we group by Day)
# UTC date
# We interpret "ChargePeriod" as the day of the request.
ts = log.get("startTime") or time.time()
dt = datetime.fromtimestamp(ts)
date_str = dt.strftime("%Y-%m-%d")
# ChargePeriodStart and End
# If we want daily granularity, end date is usually same day or next day?
# Datadog Custom Costs usually expects periods.
# "ChargePeriodStart": "2023-01-01", "ChargePeriodEnd": "2023-12-31" in example.
# If we send daily, we can say Start=Date, End=Date.
# Grouping Key: Provider + Model + Date + Tags?
# For simplicity, let's aggregate by Provider + Model + Date first.
# If we handle tags, we need to include them in the key.
tags = self._extract_tags(log)
tags_key = tuple(sorted(tags.items())) if tags else ()
key = (provider, model, date_str, tags_key)
if key not in aggregator:
aggregator[key] = {
"ProviderName": provider,
"ChargeDescription": f"LLM Usage for {model}",
"ChargePeriodStart": date_str,
"ChargePeriodEnd": date_str,
"BilledCost": 0.0,
"BillingCurrency": "USD",
"Tags": tags if tags else None,
}
aggregator[key]["BilledCost"] += cost
except Exception as e:
verbose_logger.warning(
f"Error processing log for cost aggregation: {e}"
)
continue
return list(aggregator.values())
def _extract_tags(self, log: StandardLoggingPayload) -> Dict[str, str]:
from litellm.integrations.datadog.datadog_handler import (
get_datadog_env,
get_datadog_hostname,
get_datadog_pod_name,
get_datadog_service,
)
tags = {
"env": get_datadog_env(),
"service": get_datadog_service(),
"host": get_datadog_hostname(),
"pod_name": get_datadog_pod_name(),
}
# Add metadata as tags
metadata = log.get("metadata", {})
if metadata:
# Add user info
if "user_api_key_alias" in metadata:
tags["user"] = str(metadata["user_api_key_alias"])
if "user_api_key_team_alias" in metadata:
tags["team"] = str(metadata["user_api_key_team_alias"])
# model_group is not in StandardLoggingMetadata TypedDict, so we need to access it via dict.get()
model_group = metadata.get("model_group") # type: ignore[misc]
if model_group:
tags["model_group"] = str(model_group)
return tags
async def _upload_to_datadog(self, payload: List[Dict]):
if not self.dd_api_key or not self.dd_app_key:
return
headers = {
"Content-Type": "application/json",
"DD-API-KEY": self.dd_api_key,
"DD-APPLICATION-KEY": self.dd_app_key,
}
# The API endpoint expects a list of objects directly in the body (file content behavior)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
data_json = safe_dumps(payload)
response = await self.async_client.put(
self.upload_url, content=data_json, headers=headers
)
response.raise_for_status()
verbose_logger.debug(
f"Datadog Cost Management: Uploaded {len(payload)} cost entries. Status: {response.status_code}"
)
litellm/integrations/datadog/datadog_handler.py
+8
View File
@@ -20,6 +20,14 @@ def get_datadog_hostname() -> str:
return os.getenv("HOSTNAME", "")
def get_datadog_base_url_from_env() -> Optional[str]:
"""
Get base URL override from common DD_BASE_URL env var.
This is useful for testing or custom endpoints.
"""
return os.getenv("DD_BASE_URL")
def get_datadog_env() -> str:
return os.getenv("DD_ENV", "unknown")
litellm/integrations/datadog/datadog_llm_obs.py
+48 -16
View File
@@ -21,6 +21,7 @@ from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.integrations.datadog.datadog_handler import (
get_datadog_service,
get_datadog_tags,
get_datadog_base_url_from_env,
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.prompt_templates.common_utils import (
@@ -43,24 +44,22 @@ class DataDogLLMObsLogger(CustomBatchLogger):
def __init__(self, **kwargs):
try:
verbose_logger.debug("DataDogLLMObs: Initializing logger")
if os.getenv("DD_API_KEY", None) is None:
raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>'")
if os.getenv("DD_SITE", None) is None:
raise Exception(
"DD_SITE is not set, set 'DD_SITE=<>', example sit = `us5.datadoghq.com`"
)
# Configure DataDog endpoint (Agent or Direct API)
# Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST
dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST")
self.async_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
self.DD_API_KEY = os.getenv("DD_API_KEY")
self.DD_SITE = os.getenv("DD_SITE")
self.intake_url = (
f"https://api.{self.DD_SITE}/api/intake/llm-obs/v1/trace/spans"
)
# testing base url
dd_base_url = os.getenv("DD_BASE_URL")
if dd_agent_host:
self._configure_dd_agent(dd_agent_host=dd_agent_host)
else:
self._configure_dd_direct_api()
# Optional override for testing
dd_base_url = get_datadog_base_url_from_env()
if dd_base_url:
self.intake_url = f"{dd_base_url}/api/intake/llm-obs/v1/trace/spans"
@@ -78,6 +77,38 @@ class DataDogLLMObsLogger(CustomBatchLogger):
verbose_logger.exception(f"DataDogLLMObs: Error initializing - {str(e)}")
raise e
def _configure_dd_agent(self, dd_agent_host: str):
"""
Configure the Datadog logger to send traces to the Agent.
"""
# When using the Agent, LLM Observability Intake does NOT require the API Key
# Reference: https://docs.datadoghq.com/llm_observability/setup/sdk/#agent-setup
# Use specific port for LLM Obs (Trace Agent) to avoid conflict with Logs Agent (10518)
agent_port = os.getenv("LITELLM_DD_LLM_OBS_PORT", "8126")
self.DD_SITE = "localhost" # Not used for URL construction in agent mode
self.intake_url = (
f"http://{dd_agent_host}:{agent_port}/api/intake/llm-obs/v1/trace/spans"
)
verbose_logger.debug(f"DataDogLLMObs: Using DD Agent at {self.intake_url}")
def _configure_dd_direct_api(self):
"""
Configure the Datadog logger to send traces directly to the Datadog API.
"""
if not self.DD_API_KEY:
raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>'")
self.DD_SITE = os.getenv("DD_SITE")
if not self.DD_SITE:
raise Exception(
"DD_SITE is not set, set 'DD_SITE=<>', example site = `us5.datadoghq.com`"
)
self.intake_url = (
f"https://api.{self.DD_SITE}/api/intake/llm-obs/v1/trace/spans"
)
def _get_datadog_llm_obs_params(self) -> Dict:
"""
Get the datadog_llm_observability_params from litellm.datadog_llm_observability_params
@@ -164,13 +195,14 @@ class DataDogLLMObsLogger(CustomBatchLogger):
json_payload = safe_dumps(payload)
headers = {"Content-Type": "application/json"}
if self.DD_API_KEY:
headers["DD-API-KEY"] = self.DD_API_KEY
response = await self.async_client.post(
url=self.intake_url,
content=json_payload,
headers={
"DD-API-KEY": self.DD_API_KEY,
"Content-Type": "application/json",
},
headers=headers,
)
if response.status_code != 202:
litellm/integrations/langfuse/langfuse.py
+7 -7
View File
@@ -23,6 +23,7 @@ from litellm.constants import MAX_LANGFUSE_INITIALIZED_CLIENTS
from litellm.litellm_core_utils.core_helpers import (
safe_deep_copy,
reconstruct_model_name,
filter_exceptions_from_params,
)
from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info
from litellm.integrations.langfuse.langfuse_mock_client import (
@@ -75,9 +76,8 @@ def _extract_cache_read_input_tokens(usage_obj) -> int:
# Check prompt_tokens_details.cached_tokens (used by Gemini and other providers)
if hasattr(usage_obj, "prompt_tokens_details"):
prompt_tokens_details = getattr(usage_obj, "prompt_tokens_details", None)
if (
prompt_tokens_details is not None
and hasattr(prompt_tokens_details, "cached_tokens")
if prompt_tokens_details is not None and hasattr(
prompt_tokens_details, "cached_tokens"
):
cached_tokens = getattr(prompt_tokens_details, "cached_tokens", None)
if (
@@ -540,7 +540,6 @@ class LangFuseLogger:
verbose_logger.debug("Langfuse Layer Logging - logging to langfuse v2")
try:
metadata = metadata or {}
standard_logging_object: Optional[StandardLoggingPayload] = cast(
Optional[StandardLoggingPayload],
kwargs.get("standard_logging_object", None),
@@ -706,9 +705,10 @@ class LangFuseLogger:
clean_metadata["litellm_response_cost"] = cost
if standard_logging_object is not None:
clean_metadata["hidden_params"] = standard_logging_object[
"hidden_params"
]
hidden_params = standard_logging_object.get("hidden_params", {})
clean_metadata["hidden_params"] = filter_exceptions_from_params(
hidden_params
)
if (
litellm.langfuse_default_tags is not None
litellm/integrations/prometheus.py
+105 -16
View File
@@ -229,14 +229,18 @@ class PrometheusLogger(CustomLogger):
self.litellm_remaining_api_key_requests_for_model = self._gauge_factory(
"litellm_remaining_api_key_requests_for_model",
"Remaining Requests API Key can make for model (model based rpm limit on key)",
labelnames=["hashed_api_key", "api_key_alias", "model"],
labelnames=self.get_labels_for_metric(
"litellm_remaining_api_key_requests_for_model"
),
)
# Remaining MODEL TPM limit for API Key
self.litellm_remaining_api_key_tokens_for_model = self._gauge_factory(
"litellm_remaining_api_key_tokens_for_model",
"Remaining Tokens API Key can make for model (model based tpm limit on key)",
labelnames=["hashed_api_key", "api_key_alias", "model"],
labelnames=self.get_labels_for_metric(
"litellm_remaining_api_key_tokens_for_model"
),
)
########################################
@@ -312,6 +316,18 @@ class PrometheusLogger(CustomLogger):
labelnames=self.get_labels_for_metric("litellm_deployment_state"),
)
self.litellm_deployment_tpm_limit = self._gauge_factory(
"litellm_deployment_tpm_limit",
"Deployment TPM limit found in config",
labelnames=self.get_labels_for_metric("litellm_deployment_tpm_limit"),
)
self.litellm_deployment_rpm_limit = self._gauge_factory(
"litellm_deployment_rpm_limit",
"Deployment RPM limit found in config",
labelnames=self.get_labels_for_metric("litellm_deployment_rpm_limit"),
)
self.litellm_deployment_cooled_down = self._counter_factory(
"litellm_deployment_cooled_down",
"LLM Deployment Analytics - Number of times a deployment has been cooled down by LiteLLM load balancing logic. exception_status is the status of the exception that caused the deployment to be cooled down",
@@ -373,15 +389,9 @@ class PrometheusLogger(CustomLogger):
self.litellm_llm_api_failed_requests_metric = self._counter_factory(
name="litellm_llm_api_failed_requests_metric",
documentation="deprecated - use litellm_proxy_failed_requests_metric",
labelnames=[
"end_user",
"hashed_api_key",
"api_key_alias",
"model",
"team",
"team_alias",
"user",
],
labelnames=self.get_labels_for_metric(
"litellm_llm_api_failed_requests_metric"
),
)
self.litellm_requests_metric = self._counter_factory(
@@ -954,6 +964,8 @@ class PrometheusLogger(CustomLogger):
route=standard_logging_payload["metadata"].get(
"user_api_key_request_route"
),
client_ip=standard_logging_payload["metadata"].get("requester_ip_address"),
user_agent=standard_logging_payload["metadata"].get("user_agent"),
)
if (
@@ -1011,6 +1023,7 @@ class PrometheusLogger(CustomLogger):
user_api_key_alias=user_api_key_alias,
kwargs=kwargs,
metadata=_metadata,
model_id=enum_values.model_id,
)
# set latency metrics
@@ -1234,6 +1247,7 @@ class PrometheusLogger(CustomLogger):
user_api_key_alias: Optional[str],
kwargs: dict,
metadata: dict,
model_id: Optional[str] = None,
):
from litellm.proxy.common_utils.callback_utils import (
get_model_group_from_litellm_kwargs,
@@ -1255,11 +1269,11 @@ class PrometheusLogger(CustomLogger):
)
self.litellm_remaining_api_key_requests_for_model.labels(
user_api_key, user_api_key_alias, model_group
user_api_key, user_api_key_alias, model_group, model_id
).set(remaining_requests)
self.litellm_remaining_api_key_tokens_for_model.labels(
user_api_key, user_api_key_alias, model_group
user_api_key, user_api_key_alias, model_group, model_id
).set(remaining_tokens)
def _set_latency_metrics(
@@ -1387,6 +1401,7 @@ class PrometheusLogger(CustomLogger):
user_api_team,
user_api_team_alias,
user_id,
standard_logging_payload.get("model_id", ""),
).inc()
self.set_llm_deployment_failure_metrics(kwargs)
except Exception as e:
@@ -1575,6 +1590,10 @@ class PrometheusLogger(CustomLogger):
litellm_params=request_data,
proxy_server_request=request_data.get("proxy_server_request", {}),
)
_metadata = request_data.get("metadata", {}) or {}
model_id = _metadata.get("model_info", {}).get("id") or request_data.get(
"model_info", {}
).get("id")
enum_values = UserAPIKeyLabelValues(
end_user=user_api_key_dict.end_user_id,
user=user_api_key_dict.user_id,
@@ -1589,6 +1608,9 @@ class PrometheusLogger(CustomLogger):
exception_class=self._get_exception_class_name(original_exception),
tags=_tags,
route=user_api_key_dict.request_route,
client_ip=_metadata.get("requester_ip_address"),
user_agent=_metadata.get("user_agent"),
model_id=model_id,
)
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
@@ -1628,6 +1650,7 @@ class PrometheusLogger(CustomLogger):
):
return
_metadata = data.get("metadata", {}) or {}
enum_values = UserAPIKeyLabelValues(
end_user=user_api_key_dict.end_user_id,
hashed_api_key=user_api_key_dict.api_key,
@@ -1643,6 +1666,8 @@ class PrometheusLogger(CustomLogger):
litellm_params=data,
proxy_server_request=data.get("proxy_server_request", {}),
),
client_ip=_metadata.get("requester_ip_address"),
user_agent=_metadata.get("user_agent"),
)
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
@@ -1715,6 +1740,10 @@ class PrometheusLogger(CustomLogger):
"user_api_key_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"),
)
"""
@@ -1752,6 +1781,49 @@ class PrometheusLogger(CustomLogger):
)
)
def _set_deployment_tpm_rpm_limit_metrics(
self,
model_info: dict,
litellm_params: dict,
litellm_model_name: Optional[str],
model_id: Optional[str],
api_base: Optional[str],
llm_provider: Optional[str],
):
"""
Set the deployment TPM and RPM limits metrics
"""
tpm = model_info.get("tpm") or litellm_params.get("tpm")
rpm = model_info.get("rpm") or litellm_params.get("rpm")
if tpm is not None:
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name="litellm_deployment_tpm_limit"
),
enum_values=UserAPIKeyLabelValues(
litellm_model_name=litellm_model_name,
model_id=model_id,
api_base=api_base,
api_provider=llm_provider,
),
)
self.litellm_deployment_tpm_limit.labels(**_labels).set(tpm)
if rpm is not None:
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name="litellm_deployment_rpm_limit"
),
enum_values=UserAPIKeyLabelValues(
litellm_model_name=litellm_model_name,
model_id=model_id,
api_base=api_base,
api_provider=llm_provider,
),
)
self.litellm_deployment_rpm_limit.labels(**_labels).set(rpm)
def set_llm_deployment_success_metrics(
self,
request_kwargs: dict,
@@ -1785,6 +1857,16 @@ class PrometheusLogger(CustomLogger):
_model_info = _metadata.get("model_info") or {}
model_id = _model_info.get("id", None)
if _model_info or _litellm_params:
self._set_deployment_tpm_rpm_limit_metrics(
model_info=_model_info,
litellm_params=_litellm_params,
litellm_model_name=litellm_model_name,
model_id=model_id,
api_base=api_base,
llm_provider=llm_provider,
)
remaining_requests: Optional[int] = None
remaining_tokens: Optional[int] = None
if additional_headers := standard_logging_payload["hidden_params"][
@@ -2262,7 +2344,10 @@ class PrometheusLogger(CustomLogger):
async def fetch_keys(
page_size: int, page: int
) -> Tuple[List[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]], Optional[int]]:
) -> Tuple[
List[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]],
Optional[int],
]:
key_list_response = await _list_key_helper(
prisma_client=prisma_client,
page=page,
@@ -2378,12 +2463,16 @@ class PrometheusLogger(CustomLogger):
# Get total user count
total_users = await prisma_client.db.litellm_usertable.count()
self.litellm_total_users_metric.set(total_users)
verbose_logger.debug(f"Prometheus: set litellm_total_users to {total_users}")
verbose_logger.debug(
f"Prometheus: set litellm_total_users to {total_users}"
)
# Get total team count
total_teams = await prisma_client.db.litellm_teamtable.count()
self.litellm_teams_count_metric.set(total_teams)
verbose_logger.debug(f"Prometheus: set litellm_teams_count to {total_teams}")
verbose_logger.debug(
f"Prometheus: set litellm_teams_count to {total_teams}"
)
except Exception as e:
verbose_logger.exception(
f"Error initializing user/team count metrics: {str(e)}"
litellm/litellm_core_utils/core_helpers.py
+2 -2
View File
@@ -351,9 +351,9 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any:
# Skip callable objects (functions, methods, lambdas) but not classes (type objects)
if callable(data) and not isinstance(data, type):
return None
# Skip known non-serializable object types (Logging, etc.)
# Skip known non-serializable object types (Logging, Router, etc.)
obj_type_name = type(data).__name__
if obj_type_name in ["Logging", "LiteLLMLoggingObj"]:
if obj_type_name in ["Logging", "LiteLLMLoggingObj", "Router"]:
return None
if isinstance(data, dict):
litellm/litellm_core_utils/get_litellm_params.py
+7 -2
View File
@@ -93,8 +93,11 @@ def get_litellm_params(
"text_completion": text_completion,
"azure_ad_token_provider": azure_ad_token_provider,
"user_continue_message": user_continue_message,
"base_model": base_model or (
_get_base_model_from_litellm_call_metadata(metadata=metadata) if metadata else None
"base_model": base_model
or (
_get_base_model_from_litellm_call_metadata(metadata=metadata)
if metadata
else None
),
"litellm_trace_id": litellm_trace_id,
"litellm_session_id": litellm_session_id,
@@ -139,5 +142,7 @@ def get_litellm_params(
"aws_sts_endpoint": kwargs.get("aws_sts_endpoint"),
"aws_external_id": kwargs.get("aws_external_id"),
"aws_bedrock_runtime_endpoint": kwargs.get("aws_bedrock_runtime_endpoint"),
"tpm": kwargs.get("tpm"),
"rpm": kwargs.get("rpm"),
}
return litellm_params
litellm/litellm_core_utils/litellm_logging.py
+11 -2
View File
@@ -335,7 +335,9 @@ class Logging(LiteLLMLoggingBaseClass):
self.start_time = start_time # log the call start time
self.call_type = call_type
self.litellm_call_id = litellm_call_id
self.litellm_trace_id: str = litellm_trace_id if litellm_trace_id else str(uuid.uuid4())
self.litellm_trace_id: str = (
litellm_trace_id if litellm_trace_id else str(uuid.uuid4())
)
self.function_id = function_id
self.streaming_chunks: List[Any] = [] # for generating complete stream response
self.sync_streaming_chunks: List[
@@ -544,7 +546,10 @@ class Logging(LiteLLMLoggingBaseClass):
if "stream_options" in additional_params:
self.stream_options = additional_params["stream_options"]
## check if custom pricing set ##
if any(litellm_params.get(key) is not None for key in _CUSTOM_PRICING_KEYS & litellm_params.keys()):
if any(
litellm_params.get(key) is not None
for key in _CUSTOM_PRICING_KEYS & litellm_params.keys()
):
self.custom_pricing = True
if "custom_llm_provider" in self.model_call_details:
@@ -3299,6 +3304,7 @@ def _get_masked_values(
"token",
"key",
"secret",
"vertex_credentials",
]
return {
k: (
@@ -4453,6 +4459,7 @@ class StandardLoggingPayloadSetup:
user_api_key_request_route=None,
spend_logs_metadata=None,
requester_ip_address=None,
user_agent=None,
requester_metadata=None,
prompt_management_metadata=prompt_management_metadata,
applied_guardrails=applied_guardrails,
@@ -5138,6 +5145,7 @@ def get_standard_logging_object_payload(
model_group=_model_group,
model_id=_model_id,
requester_ip_address=clean_metadata.get("requester_ip_address", None),
user_agent=clean_metadata.get("user_agent", None),
messages=StandardLoggingPayloadSetup.append_system_prompt_messages(
kwargs=kwargs, messages=kwargs.get("messages")
),
@@ -5203,6 +5211,7 @@ def get_standard_logging_metadata(
user_api_key_team_alias=None,
spend_logs_metadata=None,
requester_ip_address=None,
user_agent=None,
requester_metadata=None,
user_api_key_end_user_id=None,
prompt_management_metadata=None,
litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py
+18
View File
@@ -21,11 +21,13 @@ from litellm.types.utils import (
ChatCompletionMessageToolCall,
ChatCompletionRedactedThinkingBlock,
Choices,
CompletionTokensDetailsWrapper,
Delta,
EmbeddingResponse,
Function,
HiddenParams,
ImageResponse,
PromptTokensDetailsWrapper,
)
from litellm.types.utils import Logprobs as TextCompletionLogprobs
from litellm.types.utils import (
@@ -304,6 +306,22 @@ class LiteLLMResponseObjectHandler:
"text_tokens": 0,
}
# Map Responses API naming to Chat Completions API naming for cost calculator
if usage.get("prompt_tokens") is None:
usage["prompt_tokens"] = usage.get("input_tokens", 0)
if usage.get("completion_tokens") is None:
usage["completion_tokens"] = usage.get("output_tokens", 0)
# Convert dicts to wrapper objects so getattr() works in cost calculation
if isinstance(usage.get("input_tokens_details"), dict):
usage["prompt_tokens_details"] = PromptTokensDetailsWrapper(
**usage["input_tokens_details"]
)
if isinstance(usage.get("output_tokens_details"), dict):
usage["completion_tokens_details"] = CompletionTokensDetailsWrapper(
**usage["output_tokens_details"]
)
if model_response_object is None:
model_response_object = ImageResponse(**response_object)
return model_response_object
litellm/litellm_core_utils/prompt_templates/factory.py
+2 -1
View File
@@ -4408,7 +4408,7 @@ def _bedrock_tools_pt(tools: List) -> List[BedrockToolBlock]:
]
"""
"""
Bedrock toolConfig looks like:
Bedrock toolConfig looks like:
"tools": [
{
"toolSpec": {
@@ -4436,6 +4436,7 @@ def _bedrock_tools_pt(tools: List) -> List[BedrockToolBlock]:
tool_block_list: List[BedrockToolBlock] = []
for tool in tools:
# Handle regular function tools
parameters = tool.get("function", {}).get(
"parameters", {"type": "object", "properties": {}}
)
litellm/llms/anthropic/chat/guardrail_translation/handler.py
+13 -1
View File
@@ -110,6 +110,10 @@ class AnthropicMessagesHandler(BaseTranslation):
inputs["tools"] = tools_to_check
if structured_messages:
inputs["structured_messages"] = structured_messages
# Include model information if available
model = data.get("model")
if model:
inputs["model"] = model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=data,
@@ -309,6 +313,14 @@ class AnthropicMessagesHandler(BaseTranslation):
inputs["images"] = images_to_check
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check
# Include model information from the response if available
response_model = None
if isinstance(response, dict):
response_model = response.get("model")
elif hasattr(response, "model"):
response_model = getattr(response, "model", None)
if response_model:
inputs["model"] = response_model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
@@ -552,7 +564,7 @@ class AnthropicMessagesHandler(BaseTranslation):
response_content = response.get("content", [])
else:
response_content = getattr(response, "content", None) or []
if not response_content:
return False
for content_block in response_content:
litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py
+64 -42
View File
@@ -170,21 +170,33 @@ class LiteLLMAnthropicMessagesAdapter:
def _add_cache_control_if_applicable(
self,
source: Dict[str, Any],
target: Dict[str, Any],
source: Any,
target: Any,
model: Optional[str],
) -> None:
"""
Extract cache_control from source and add to target if it should be preserved.
This method accepts Any type to support both regular dicts and TypedDict objects.
TypedDict objects (like ChatCompletionTextObject, ChatCompletionImageObject, etc.)
are dicts at runtime but have specific types at type-check time. Using Any allows
this method to work with both while maintaining runtime correctness.
Args:
source: Dict containing potential cache_control field
target: Dict to add cache_control to
source: Dict or TypedDict containing potential cache_control field
target: Dict or TypedDict to add cache_control to
model: Model name to check if cache_control should be preserved
"""
cache_control = source.get("cache_control")
# TypedDict objects are dicts at runtime, so .get() works
cache_control = source.get("cache_control") if isinstance(source, dict) else getattr(source, "cache_control", None)
if cache_control and model and self.is_anthropic_claude_model(model):
target["cache_control"] = cache_control
# TypedDict objects support dict operations at runtime
# Use type ignore consistent with codebase pattern (see anthropic/chat/transformation.py:432)
if isinstance(target, dict):
target["cache_control"] = cache_control # type: ignore[typeddict-item]
else:
# Fallback for non-dict objects (shouldn't happen in practice)
cast(Dict[str, Any], target)["cache_control"] = cache_control
def translatable_anthropic_params(self) -> List:
"""
@@ -220,7 +232,7 @@ class LiteLLMAnthropicMessagesAdapter:
elif message_content and isinstance(message_content, list):
for content in message_content:
if content.get("type") == "text":
text_obj: Dict[str, Any] = ChatCompletionTextObject(
text_obj = ChatCompletionTextObject(
type="text", text=content.get("text", "")
)
self._add_cache_control_if_applicable(content, text_obj, model)
@@ -236,7 +248,7 @@ class LiteLLMAnthropicMessagesAdapter:
image_url_obj = ChatCompletionImageUrlObject(
url=openai_image_url
)
image_obj: Dict[str, Any] = ChatCompletionImageObject(
image_obj = ChatCompletionImageObject(
type="image_url", image_url=image_url_obj
)
self._add_cache_control_if_applicable(content, image_obj, model)
@@ -245,21 +257,21 @@ class LiteLLMAnthropicMessagesAdapter:
# Convert Anthropic document format (PDF, etc.) to OpenAI format
source = content.get("source", {})
openai_image_url = (
self._translate_anthropic_image_to_openai(source)
self._translate_anthropic_image_to_openai(cast(dict, source))
)
if openai_image_url:
image_url_obj = ChatCompletionImageUrlObject(
url=openai_image_url
)
doc_obj: Dict[str, Any] = ChatCompletionImageObject(
doc_obj = ChatCompletionImageObject(
type="image_url", image_url=image_url_obj
)
self._add_cache_control_if_applicable(content, doc_obj, model)
new_user_content_list.append(doc_obj) # type: ignore
elif content.get("type") == "tool_result":
if "content" not in content:
tool_result: Dict[str, Any] = ChatCompletionToolMessage(
tool_result = ChatCompletionToolMessage(
role="tool",
tool_call_id=content.get("tool_use_id", ""),
content="",
@@ -383,7 +395,7 @@ class LiteLLMAnthropicMessagesAdapter:
assistant_message_str: Optional[str] = None
assistant_content_list: List[Dict[str, Any]] = [] # For content blocks with cache_control
has_cache_control_in_text = False
tool_calls: List[Dict[str, Any]] = []
tool_calls: List[ChatCompletionAssistantToolCall] = []
thinking_blocks: List[
Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]
] = []
@@ -427,7 +439,7 @@ class LiteLLMAnthropicMessagesAdapter:
provider_specific_fields
)
tool_call: Dict[str, Any] = ChatCompletionAssistantToolCall(
tool_call = ChatCompletionAssistantToolCall(
id=content.get("id", ""),
type="function",
function=function_chunk,
@@ -603,7 +615,7 @@ class LiteLLMAnthropicMessagesAdapter:
for k, v in tool.items():
if k not in mapped_tool_params: # pass additional computer kwargs
function_chunk.setdefault("parameters", {}).update({k: v})
tool_param: Dict[str, Any] = ChatCompletionToolParam(type="function", function=function_chunk)
tool_param = ChatCompletionToolParam(type="function", function=function_chunk)
self._add_cache_control_if_applicable(tool, tool_param, model)
new_tools.append(tool_param) # type: ignore[arg-type]
@@ -645,6 +657,41 @@ class LiteLLMAnthropicMessagesAdapter:
},
}
def _add_system_message_to_messages(
self,
new_messages: List[AllMessageValues],
anthropic_message_request: AnthropicMessagesRequest,
) -> None:
"""Add system message to messages list if present in request."""
if "system" not in anthropic_message_request:
return
system_content = anthropic_message_request["system"]
if not system_content:
return
# Handle system as string or array of content blocks
if isinstance(system_content, str):
new_messages.insert(
0,
ChatCompletionSystemMessage(role="system", content=system_content),
)
elif isinstance(system_content, list):
# Convert Anthropic system content blocks to OpenAI format
openai_system_content: List[Dict[str, Any]] = []
model_name = anthropic_message_request.get("model", "")
for block in system_content:
if isinstance(block, dict) and block.get("type") == "text":
text_block: Dict[str, Any] = {
"type": "text",
"text": block.get("text", ""),
}
self._add_cache_control_if_applicable(block, text_block, model_name)
openai_system_content.append(text_block)
if openai_system_content:
new_messages.insert(
0,
ChatCompletionSystemMessage(role="system", content=openai_system_content), # type: ignore
)
def translate_anthropic_to_openai(
self, anthropic_message_request: AnthropicMessagesRequest
) -> ChatCompletionRequest:
@@ -673,32 +720,7 @@ class LiteLLMAnthropicMessagesAdapter:
model=anthropic_message_request.get("model"),
)
## ADD SYSTEM MESSAGE TO MESSAGES
if "system" in anthropic_message_request:
system_content = anthropic_message_request["system"]
if system_content:
# Handle system as string or array of content blocks
if isinstance(system_content, str):
new_messages.insert(
0,
ChatCompletionSystemMessage(role="system", content=system_content),
)
elif isinstance(system_content, list):
# Convert Anthropic system content blocks to OpenAI format
openai_system_content: List[Dict[str, Any]] = []
model_name = anthropic_message_request.get("model", "")
for block in system_content:
if isinstance(block, dict) and block.get("type") == "text":
text_block: Dict[str, Any] = {
"type": "text",
"text": block.get("text", ""),
}
self._add_cache_control_if_applicable(block, text_block, model_name)
openai_system_content.append(text_block)
if openai_system_content:
new_messages.insert(
0,
ChatCompletionSystemMessage(role="system", content=openai_system_content), # type: ignore
)
self._add_system_message_to_messages(new_messages, anthropic_message_request)
new_kwargs: ChatCompletionRequest = {
"model": anthropic_message_request["model"],
@@ -902,7 +924,7 @@ class LiteLLMAnthropicMessagesAdapter:
)
# extract usage
usage: Usage = getattr(response, "usage")
anthropic_usage: Dict[str, Any] = AnthropicUsage(
anthropic_usage = AnthropicUsage(
input_tokens=usage.prompt_tokens or 0,
output_tokens=usage.completion_tokens or 0,
)
@@ -1055,7 +1077,7 @@ class LiteLLMAnthropicMessagesAdapter:
else:
litellm_usage_chunk = None
if litellm_usage_chunk is not None:
usage_delta: Dict[str, Any] = UsageDelta(
usage_delta = UsageDelta(
input_tokens=litellm_usage_chunk.prompt_tokens or 0,
output_tokens=litellm_usage_chunk.completion_tokens or 0,
)
litellm/llms/bedrock/chat/converse_transformation.py
+67 -6
View File
@@ -298,6 +298,39 @@ class AmazonConverseConfig(BaseConfig):
# Check if the model is specifically Nova Lite 2
return "nova-2-lite" in model_without_region
def _map_web_search_options(
self,
web_search_options: dict,
model: str
) -> Optional[BedrockToolBlock]:
"""
Map web_search_options to Nova grounding systemTool.
Nova grounding (web search) is only supported on Amazon Nova models.
Returns None for non-Nova models.
Args:
web_search_options: The web_search_options dict from the request
model: The model identifier string
Returns:
BedrockToolBlock with systemTool for Nova models, None otherwise
Reference: https://docs.aws.amazon.com/nova/latest/userguide/grounding.html
"""
# Only Nova models support nova_grounding
# Model strings can be like: "amazon.nova-pro-v1:0", "us.amazon.nova-pro-v1:0", etc.
if "nova" not in model.lower():
verbose_logger.debug(
f"web_search_options passed but model {model} is not a Nova model. "
"Nova grounding is only supported on Amazon Nova models."
)
return None
# Nova doesn't support search_context_size or user_location params
# (unlike Anthropic), so we just enable grounding with no options
return BedrockToolBlock(systemTool={"name": "nova_grounding"})
def _transform_reasoning_effort_to_reasoning_config(
self, reasoning_effort: str
) -> dict:
@@ -438,6 +471,10 @@ class AmazonConverseConfig(BaseConfig):
):
supported_params.append("tools")
# Nova models support web_search_options (mapped to nova_grounding systemTool)
if base_model.startswith("amazon.nova"):
supported_params.append("web_search_options")
if litellm.utils.supports_tool_choice(
model=model, custom_llm_provider=self.custom_llm_provider
) or litellm.utils.supports_tool_choice(
@@ -730,6 +767,13 @@ class AmazonConverseConfig(BaseConfig):
if bedrock_tier in ("default", "flex", "priority"):
optional_params["serviceTier"] = {"type": bedrock_tier}
if param == "web_search_options" and value and isinstance(value, dict):
grounding_tool = self._map_web_search_options(value, model)
if grounding_tool is not None:
optional_params = self._add_tools_to_optional_params(
optional_params=optional_params, tools=[grounding_tool]
)
# Only update thinking tokens for non-GPT-OSS models and non-Nova-Lite-2 models
# Nova Lite 2 handles token budgeting differently through reasoningConfig
if "gpt-oss" not in model and not self._is_nova_lite_2_model(model):
@@ -1388,20 +1432,23 @@ class AmazonConverseConfig(BaseConfig):
str,
List[ChatCompletionToolCallChunk],
Optional[List[BedrockConverseReasoningContentBlock]],
Optional[List[CitationsContentBlock]],
]:
"""
Translate the message content to a string and a list of tool calls and reasoning content blocks
Translate the message content to a string and a list of tool calls, reasoning content blocks, and citations.
Returns:
content_str: str
tools: List[ChatCompletionToolCallChunk]
reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]]
citationsContentBlocks: Optional[List[CitationsContentBlock]] - Citations from Nova grounding
"""
content_str = ""
tools: List[ChatCompletionToolCallChunk] = []
reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = (
None
)
citationsContentBlocks: Optional[List[CitationsContentBlock]] = None
for idx, content in enumerate(content_blocks):
"""
- Content is either a tool response or text
@@ -1446,10 +1493,15 @@ class AmazonConverseConfig(BaseConfig):
if reasoningContentBlocks is None:
reasoningContentBlocks = []
reasoningContentBlocks.append(content["reasoningContent"])
# Handle Nova grounding citations content
if "citationsContent" in content:
if citationsContentBlocks is None:
citationsContentBlocks = []
citationsContentBlocks.append(content["citationsContent"])
return content_str, tools, reasoningContentBlocks
return content_str, tools, reasoningContentBlocks, citationsContentBlocks
def _transform_response(
def _transform_response( # noqa: PLR0915
self,
model: str,
response: httpx.Response,
@@ -1525,18 +1577,27 @@ class AmazonConverseConfig(BaseConfig):
reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = (
None
)
citationsContentBlocks: Optional[List[CitationsContentBlock]] = None
if message is not None:
(
content_str,
tools,
reasoningContentBlocks,
citationsContentBlocks,
) = self._translate_message_content(message["content"])
# Initialize provider_specific_fields if we have any special content blocks
provider_specific_fields: dict = {}
if reasoningContentBlocks is not None:
provider_specific_fields["reasoningContentBlocks"] = reasoningContentBlocks
if citationsContentBlocks is not None:
provider_specific_fields["citationsContent"] = citationsContentBlocks
if provider_specific_fields:
chat_completion_message["provider_specific_fields"] = provider_specific_fields
if reasoningContentBlocks is not None:
chat_completion_message["provider_specific_fields"] = {
"reasoningContentBlocks": reasoningContentBlocks,
}
chat_completion_message["reasoning_content"] = (
self._transform_reasoning_content(reasoningContentBlocks)
)
litellm/llms/bedrock/chat/invoke_handler.py
+5
View File
@@ -1476,6 +1476,11 @@ class AWSEventStreamDecoder:
reasoning_content = (
"" # set to non-empty string to ensure consistency with Anthropic
)
elif "citationsContent" in delta_obj:
# Handle Nova grounding citations in streaming responses
provider_specific_fields = {
"citationsContent": delta_obj["citationsContent"],
}
return (
text,
tool_use,
litellm/llms/cohere/rerank/guardrail_translation/handler.py
+7 -1
View File
@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
@@ -49,8 +50,13 @@ class CohereRerankHandler(BaseTranslation):
# Process query only
query = data.get("query")
if query is not None and isinstance(query, str):
inputs = GenericGuardrailAPIInputs(texts=[query])
# Include model information if available
model = data.get("model")
if model:
inputs["model"] = model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [query]},
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
litellm/llms/openai/chat/guardrail_translation/handler.py
+14
View File
@@ -87,6 +87,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
tools = data.get("tools")
if tools:
inputs["tools"] = tools
# Include model information if available
model = data.get("model")
if model:
inputs["model"] = model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
@@ -297,6 +301,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
inputs["images"] = images_to_check
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check # type: ignore
# Include model information from the response if available
if hasattr(response, "model") and response.model:
inputs["model"] = response.model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
@@ -417,6 +424,13 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
inputs["images"] = images_to_check
# Include model information from the first response if available
if (
responses_so_far
and hasattr(responses_so_far[0], "model")
and responses_so_far[0].model
):
inputs["model"] = responses_so_far[0].model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=request_data,
litellm/llms/openai/completion/guardrail_translation/handler.py
+18 -3
View File
@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
@@ -53,8 +54,13 @@ class OpenAITextCompletionHandler(BaseTranslation):
if isinstance(prompt, str):
# Single string prompt
inputs = GenericGuardrailAPIInputs(texts=[prompt])
# Include model information if available
model = data.get("model")
if model:
inputs["model"] = model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [prompt]},
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
@@ -80,8 +86,13 @@ class OpenAITextCompletionHandler(BaseTranslation):
text_indices.append(idx)
if texts_to_check:
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
# Include model information if available
model = data.get("model")
if model:
inputs["model"] = model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": texts_to_check},
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
@@ -154,8 +165,12 @@ class OpenAITextCompletionHandler(BaseTranslation):
if user_metadata:
request_data["litellm_metadata"] = user_metadata
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
# Include model information from the response if available
if hasattr(response, "model") and response.model:
inputs["model"] = response.model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": texts_to_check},
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
litellm/llms/openai/image_generation/cost_calculator.py
+13 -7
View File
@@ -8,8 +8,7 @@ from typing import Optional
from litellm import verbose_logger
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
from litellm.responses.utils import ResponseAPILoggingUtils
from litellm.types.utils import ImageResponse
from litellm.types.utils import ImageResponse, Usage
def cost_calculator(
@@ -39,11 +38,18 @@ def cost_calculator(
)
return 0.0
# Transform ImageUsage to Usage using the existing helper
# ImageUsage has the same format as ResponseAPIUsage
chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
usage
)
# If usage is already a Usage object with completion_tokens_details set,
# use it directly (it was already transformed in convert_to_image_response)
if isinstance(usage, Usage) and usage.completion_tokens_details is not None:
chat_usage = usage
else:
# Transform ImageUsage to Usage using the existing helper
# ImageUsage has the same format as ResponseAPIUsage
from litellm.responses.utils import ResponseAPILoggingUtils
chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
usage
)
# Use generic_cost_per_token for cost calculation
prompt_cost, completion_cost = generic_cost_per_token(
litellm/llms/openai/image_generation/guardrail_translation/handler.py
+7 -1
View File
@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
@@ -52,8 +53,13 @@ class OpenAIImageGenerationHandler(BaseTranslation):
# Apply guardrail to the prompt
if isinstance(prompt, str):
inputs = GenericGuardrailAPIInputs(texts=[prompt])
# Include model information if available
model = data.get("model")
if model:
inputs["model"] = model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [prompt]},
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
litellm/llms/openai/responses/guardrail_translation/handler.py
+36 -7
View File
@@ -105,6 +105,10 @@ class OpenAIResponsesHandler(BaseTranslation):
inputs["tools"] = tools_to_check
if structured_messages:
inputs["structured_messages"] = structured_messages # type: ignore
# Include model information if available
model = data.get("model")
if model:
inputs["model"] = model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
@@ -150,6 +154,10 @@ class OpenAIResponsesHandler(BaseTranslation):
inputs["tools"] = tools_to_check
if structured_messages:
inputs["structured_messages"] = structured_messages # type: ignore
# Include model information if available
model = data.get("model")
if model:
inputs["model"] = model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=data,
@@ -344,6 +352,14 @@ class OpenAIResponsesHandler(BaseTranslation):
inputs["images"] = images_to_check
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check
# Include model information from the response if available
response_model = None
if isinstance(response, dict):
response_model = response.get("model")
elif hasattr(response, "model"):
response_model = getattr(response, "model", None)
if response_model:
inputs["model"] = response_model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
@@ -388,12 +404,15 @@ class OpenAIResponsesHandler(BaseTranslation):
tool_calls = model_response_stream.choices[0].delta.tool_calls
if tool_calls:
inputs = GenericGuardrailAPIInputs()
inputs["tool_calls"] = cast(
List[ChatCompletionToolCallChunk], tool_calls
)
# Include model information if available
if hasattr(model_response_stream, "model") and model_response_stream.model:
inputs["model"] = model_response_stream.model
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={
"tool_calls": cast(
List[ChatCompletionToolCallChunk], tool_calls
)
},
inputs=inputs,
request_data={},
input_type="response",
logging_obj=litellm_logging_obj,
@@ -417,7 +436,11 @@ class OpenAIResponsesHandler(BaseTranslation):
guardrail_inputs["tool_calls"] = cast(
List[ChatCompletionToolCallChunk], tool_calls
)
if 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={},
@@ -429,8 +452,14 @@ class OpenAIResponsesHandler(BaseTranslation):
# tool_calls = model_response_stream.choices[0].tool_calls
# convert openai response to model response
string_so_far = self.get_streaming_string_so_far(responses_so_far)
inputs = GenericGuardrailAPIInputs(texts=[string_so_far])
# Try to get model from the final chunk if available
if isinstance(final_chunk, dict):
response_model = final_chunk.get("response", {}).get("model") if isinstance(final_chunk.get("response"), dict) else None
if response_model:
inputs["model"] = response_model
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [string_so_far]},
inputs=inputs,
request_data={},
input_type="response",
logging_obj=litellm_logging_obj,
litellm/llms/openai/speech/guardrail_translation/handler.py
+7 -1
View File
@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
@@ -50,8 +51,13 @@ class OpenAITextToSpeechHandler(BaseTranslation):
return data
if isinstance(input_text, str):
inputs = GenericGuardrailAPIInputs(texts=[input_text])
# Include model information if available (voice model)
model = data.get("model")
if model:
inputs["model"] = model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [input_text]},
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
litellm/llms/openai/transcriptions/guardrail_translation/handler.py
+6 -1
View File
@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
@@ -88,8 +89,12 @@ class OpenAIAudioTranscriptionHandler(BaseTranslation):
if user_metadata:
request_data["litellm_metadata"] = user_metadata
inputs = GenericGuardrailAPIInputs(texts=[original_text])
# Include model information from the response if available
if hasattr(response, "model") and response.model:
inputs["model"] = response.model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [original_text]},
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
litellm/llms/pass_through/guardrail_translation/handler.py
+13 -2
View File
@@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, List, Optional
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.proxy._types import PassThroughGuardrailSettings
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
@@ -118,8 +119,13 @@ class PassThroughEndpointHandler(BaseTranslation):
return data
# Apply guardrail (pass-through doesn't modify the text, just checks it)
inputs = GenericGuardrailAPIInputs(texts=[text_to_check])
# Include model information if available
model = data.get("model")
if model:
inputs["model"] = model
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [text_to_check]},
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
@@ -178,8 +184,13 @@ class PassThroughEndpointHandler(BaseTranslation):
request_data["litellm_metadata"] = user_metadata
# Apply guardrail (pass-through doesn't modify the text, just checks it)
inputs = GenericGuardrailAPIInputs(texts=[text_to_check])
# Include model information from the response if available
response_model = response.get("model") if isinstance(response, dict) else None
if response_model:
inputs["model"] = response_model
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [text_to_check]},
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
litellm/llms/vercel_ai_gateway/embedding/__init__.py
View File
litellm/llms/vercel_ai_gateway/embedding/transformation.py
+176
View File
@@ -0,0 +1,176 @@
"""
Vercel AI Gateway Embedding API Configuration.
This module provides the configuration for Vercel AI Gateway's Embedding API.
Vercel AI Gateway is OpenAI-compatible and supports embeddings via the /v1/embeddings endpoint.
Docs: https://vercel.com/docs/ai-gateway/openai-compat/embeddings
"""
from typing import TYPE_CHECKING, Any, Optional
import httpx
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllEmbeddingInputValues
from litellm.types.utils import EmbeddingResponse
from litellm.utils import convert_to_model_response_object
from ..common_utils import VercelAIGatewayException
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class VercelAIGatewayEmbeddingConfig(BaseEmbeddingConfig):
"""
Configuration for Vercel AI Gateway's Embedding API.
Reference: https://vercel.com/docs/ai-gateway/openai-compat/embeddings
"""
def validate_environment(
self,
headers: dict,
model: str,
messages: list,
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
"""
Validate environment and set up headers for Vercel AI Gateway API.
Vercel AI Gateway requires:
- Authorization header with Bearer token (API key or OIDC token)
"""
vercel_headers = {
"Content-Type": "application/json",
}
# Add Authorization header if api_key is provided
if api_key:
vercel_headers["Authorization"] = f"Bearer {api_key}"
# Merge with existing headers (user's extra_headers take priority)
merged_headers = {**vercel_headers, **headers}
return merged_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 URL for Vercel AI Gateway Embedding API endpoint.
"""
if api_base:
api_base = api_base.rstrip("/")
else:
api_base = (
get_secret_str("VERCEL_AI_GATEWAY_API_BASE")
or "https://ai-gateway.vercel.sh/v1"
)
return f"{api_base}/embeddings"
def transform_embedding_request(
self,
model: str,
input: AllEmbeddingInputValues,
optional_params: dict,
headers: dict,
) -> dict:
"""
Transform embedding request to Vercel AI Gateway format (OpenAI-compatible).
"""
# Ensure input is a list
if isinstance(input, str):
input = [input]
# Strip 'vercel_ai_gateway/' prefix if present
if model.startswith("vercel_ai_gateway/"):
model = model.replace("vercel_ai_gateway/", "", 1)
return {
"model": model,
"input": input,
**optional_params,
}
def transform_embedding_response(
self,
model: str,
raw_response: httpx.Response,
model_response: EmbeddingResponse,
logging_obj: LiteLLMLoggingObj,
api_key: Optional[str],
request_data: dict,
optional_params: dict,
litellm_params: dict,
) -> EmbeddingResponse:
"""
Transform embedding response from Vercel AI Gateway format (OpenAI-compatible).
"""
logging_obj.post_call(original_response=raw_response.text)
# Vercel AI Gateway returns standard OpenAI-compatible embedding response
response_json = raw_response.json()
return convert_to_model_response_object(
response_object=response_json,
model_response_object=model_response,
response_type="embedding",
)
def get_supported_openai_params(self, model: str) -> list:
"""
Get list of supported OpenAI parameters for Vercel AI Gateway embeddings.
Vercel AI Gateway supports the standard OpenAI embeddings parameters
and auto-maps 'dimensions' to each provider's expected field.
"""
return [
"timeout",
"dimensions",
"encoding_format",
"user",
]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
"""
Map OpenAI parameters to Vercel AI Gateway format.
"""
for param, value in non_default_params.items():
if param in self.get_supported_openai_params(model):
optional_params[param] = value
return optional_params
def get_error_class(
self, error_message: str, status_code: int, headers: Any
) -> Any:
"""
Get the error class for Vercel AI Gateway errors.
"""
return VercelAIGatewayException(
message=error_message,
status_code=status_code,
headers=headers,
)
litellm/main.py
+49 -19
View File
@@ -148,7 +148,7 @@ from litellm.utils import (
validate_and_fix_openai_messages,
validate_and_fix_openai_tools,
validate_chat_completion_tool_choice,
validate_openai_optional_params
validate_openai_optional_params,
)
from ._logging import verbose_logger
@@ -368,7 +368,7 @@ class AsyncCompletions:
@tracer.wrap()
@client
async def acompletion( # noqa: PLR0915
async def acompletion( # noqa: PLR0915
model: str,
# Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create
messages: List = [],
@@ -603,12 +603,11 @@ async def acompletion( # noqa: PLR0915
if timeout is not None and isinstance(timeout, (int, float)):
timeout_value = float(timeout)
init_response = await asyncio.wait_for(
loop.run_in_executor(None, func_with_context),
timeout=timeout_value
loop.run_in_executor(None, func_with_context), timeout=timeout_value
)
else:
init_response = await loop.run_in_executor(None, func_with_context)
if isinstance(init_response, dict) or isinstance(
init_response, ModelResponse
): ## CACHING SCENARIO
@@ -640,6 +639,7 @@ async def acompletion( # noqa: PLR0915
except asyncio.TimeoutError:
custom_llm_provider = custom_llm_provider or "openai"
from litellm.exceptions import Timeout
raise Timeout(
message=f"Request timed out after {timeout} seconds",
model=model,
@@ -1118,7 +1118,6 @@ def completion( # type: ignore # noqa: PLR0915
# validate optional params
stop = validate_openai_optional_params(stop=stop)
######### unpacking kwargs #####################
args = locals()
@@ -1135,7 +1134,9 @@ def completion( # type: ignore # noqa: PLR0915
# Check if MCP tools are present (following responses pattern)
# Cast tools to Optional[Iterable[ToolParam]] for type checking
tools_for_mcp = cast(Optional[Iterable[ToolParam]], tools)
if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools_for_mcp):
if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(
tools=tools_for_mcp
):
# Return coroutine - acompletion will await it
# completion() can return a coroutine when MCP tools are present, which acompletion() awaits
return acompletion_with_mcp( # type: ignore[return-value]
@@ -1536,6 +1537,8 @@ def completion( # type: ignore # noqa: PLR0915
max_retries=max_retries,
timeout=timeout,
litellm_request_debug=kwargs.get("litellm_request_debug", False),
tpm=kwargs.get("tpm"),
rpm=kwargs.get("rpm"),
)
cast(LiteLLMLoggingObj, logging).update_environment_variables(
model=model,
@@ -2361,11 +2364,7 @@ def completion( # type: ignore # noqa: PLR0915
input=messages, api_key=api_key, original_response=response
)
elif custom_llm_provider == "minimax":
api_key = (
api_key
or get_secret_str("MINIMAX_API_KEY")
or litellm.api_key
)
api_key = api_key or get_secret_str("MINIMAX_API_KEY") or litellm.api_key
api_base = (
api_base
@@ -2413,7 +2412,9 @@ def completion( # type: ignore # noqa: PLR0915
or custom_llm_provider == "wandb"
or custom_llm_provider == "clarifai"
or custom_llm_provider in litellm.openai_compatible_providers
or JSONProviderRegistry.exists(custom_llm_provider) # JSON-configured providers
or JSONProviderRegistry.exists(
custom_llm_provider
) # JSON-configured providers
or "ft:gpt-3.5-turbo" in model # finetune gpt-3.5-turbo
): # allow user to make an openai call with a custom base
# note: if a user sets a custom base - we should ensure this works
@@ -4724,7 +4725,7 @@ def embedding( # noqa: PLR0915
if headers is not None and headers != {}:
optional_params["extra_headers"] = headers
if encoding_format is not None:
optional_params["encoding_format"] = encoding_format
else:
@@ -4866,6 +4867,36 @@ def embedding( # noqa: PLR0915
headers = openrouter_headers
response = base_llm_http_handler.embedding(
model=model,
input=input,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=api_key,
logging_obj=logging,
timeout=timeout,
model_response=EmbeddingResponse(),
optional_params=optional_params,
client=client,
aembedding=aembedding,
litellm_params=litellm_params_dict,
headers=headers,
)
elif custom_llm_provider == "vercel_ai_gateway":
api_base = (
api_base
or litellm.api_base
or get_secret_str("VERCEL_AI_GATEWAY_API_BASE")
or "https://ai-gateway.vercel.sh/v1"
)
api_key = (
api_key
or litellm.api_key
or get_secret_str("VERCEL_AI_GATEWAY_API_KEY")
or get_secret_str("VERCEL_OIDC_TOKEN")
)
response = base_llm_http_handler.embedding(
model=model,
input=input,
@@ -6759,9 +6790,7 @@ def speech( # noqa: PLR0915
if text_to_speech_provider_config is None:
text_to_speech_provider_config = MinimaxTextToSpeechConfig()
minimax_config = cast(
MinimaxTextToSpeechConfig, text_to_speech_provider_config
)
minimax_config = cast(MinimaxTextToSpeechConfig, text_to_speech_provider_config)
if api_base is not None:
litellm_params_dict["api_base"] = api_base
@@ -6901,7 +6930,7 @@ async def ahealth_check(
custom_llm_provider_from_params = model_params.get("custom_llm_provider", None)
api_base_from_params = model_params.get("api_base", None)
api_key_from_params = model_params.get("api_key", None)
model, custom_llm_provider, _, _ = get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider_from_params,
@@ -7275,8 +7304,9 @@ def __getattr__(name: str) -> Any:
_encoding = tiktoken.get_encoding("cl100k_base")
# Cache it in the module's __dict__ for subsequent accesses
import sys
sys.modules[__name__].__dict__["encoding"] = _encoding
global _encoding_cache
_encoding_cache = _encoding
return _encoding
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
+148 -24
View File
@@ -387,25 +387,57 @@ async def callback(code: str, state: str):
1. Try resource_metadata from WWW-Authenticate header (if present)
2. Fall back to path-based well-known URI: /.well-known/oauth-protected-resource/{path}
(
If the resource identifier value contains a path or query component, any terminating slash (/)
following the host component MUST be removed before inserting /.well-known/ and the well-known
URI path suffix between the host component and the path(include root path) and/or query components.
If the resource identifier value contains a path or query component, any terminating slash (/)
following the host component MUST be removed before inserting /.well-known/ and the well-known
URI path suffix between the host component and the path(include root path) and/or query components.
https://datatracker.ietf.org/doc/html/rfc9728#section-3.1)
3. Fall back to root-based well-known URI: /.well-known/oauth-protected-resource
Dual Pattern Support:
- Standard MCP pattern: /mcp/{server_name} (recommended, used by mcp-inspector, VSCode Copilot)
- LiteLLM legacy pattern: /{server_name}/mcp (backward compatibility)
The resource URL returned matches the pattern used in the discovery request.
"""
@router.get(f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}/mcp")
@router.get("/.well-known/oauth-protected-resource")
async def oauth_protected_resource_mcp(
request: Request, mcp_server_name: Optional[str] = None
):
def _build_oauth_protected_resource_response(
request: Request,
mcp_server_name: Optional[str],
use_standard_pattern: bool,
) -> dict:
"""
Build OAuth protected resource response with the appropriate URL pattern.
Args:
request: FastAPI Request object
mcp_server_name: Name of the MCP server
use_standard_pattern: If True, use /mcp/{server_name} pattern;
if False, use /{server_name}/mcp pattern
Returns:
OAuth protected resource metadata dict
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
# Get the correct base URL considering X-Forwarded-* headers
request_base_url = get_request_base_url(request)
mcp_server: Optional[MCPServer] = None
if mcp_server_name:
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name)
# Build resource URL based on the pattern
if mcp_server_name:
if use_standard_pattern:
# Standard MCP pattern: /mcp/{server_name}
resource_url = f"{request_base_url}/mcp/{mcp_server_name}"
else:
# LiteLLM legacy pattern: /{server_name}/mcp
resource_url = f"{request_base_url}/{mcp_server_name}/mcp"
else:
resource_url = f"{request_base_url}/mcp"
return {
"authorization_servers": [
(
@@ -414,14 +446,55 @@ async def oauth_protected_resource_mcp(
else f"{request_base_url}"
)
],
"resource": (
f"{request_base_url}/{mcp_server_name}/mcp"
if mcp_server_name
else f"{request_base_url}/mcp"
), # this is what Claude will call
"resource": resource_url,
"scopes_supported": mcp_server.scopes if mcp_server else [],
}
# Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name}
# This is the pattern expected by standard MCP clients (mcp-inspector, VSCode Copilot)
@router.get(f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}")
async def oauth_protected_resource_mcp_standard(
request: Request, mcp_server_name: str
):
"""
OAuth protected resource discovery endpoint using standard MCP URL pattern.
Standard pattern: /mcp/{server_name}
Discovery path: /.well-known/oauth-protected-resource/mcp/{server_name}
This endpoint is compliant with MCP specification and works with standard
MCP clients like mcp-inspector and VSCode Copilot.
"""
return _build_oauth_protected_resource_response(
request=request,
mcp_server_name=mcp_server_name,
use_standard_pattern=True,
)
# LiteLLM legacy pattern: /.well-known/oauth-protected-resource/{server_name}/mcp
# Kept for backward compatibility with existing deployments
@router.get(f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}/mcp")
@router.get("/.well-known/oauth-protected-resource")
async def oauth_protected_resource_mcp(
request: Request, mcp_server_name: Optional[str] = None
):
"""
OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.
Legacy pattern: /{server_name}/mcp
Discovery path: /.well-known/oauth-protected-resource/{server_name}/mcp
This endpoint is kept for backward compatibility. New integrations should
use the standard MCP pattern (/mcp/{server_name}) instead.
"""
return _build_oauth_protected_resource_response(
request=request,
mcp_server_name=mcp_server_name,
use_standard_pattern=False,
)
"""
https://datatracker.ietf.org/doc/html/rfc8414#section-3.1
RFC 8414: Path-aware OAuth discovery
@@ -430,15 +503,26 @@ async def oauth_protected_resource_mcp(
the well-known URI suffix between the host component and the path(include root path)
component.
"""
@router.get(f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}")
@router.get("/.well-known/oauth-authorization-server")
async def oauth_authorization_server_mcp(
request: Request, mcp_server_name: Optional[str] = None
):
def _build_oauth_authorization_server_response(
request: Request,
mcp_server_name: Optional[str],
) -> dict:
"""
Build OAuth authorization server metadata response.
Args:
request: FastAPI Request object
mcp_server_name: Name of the MCP server
Returns:
OAuth authorization server metadata dict
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
# Get the correct base URL considering X-Forwarded-* headers
request_base_url = get_request_base_url(request)
authorization_endpoint = (
@@ -470,18 +554,58 @@ async def oauth_authorization_server_mcp(
}
# Standard MCP pattern: /.well-known/oauth-authorization-server/mcp/{server_name}
@router.get(f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}")
async def oauth_authorization_server_mcp_standard(
request: Request, mcp_server_name: str
):
"""
OAuth authorization server discovery endpoint using standard MCP URL pattern.
Standard pattern: /mcp/{server_name}
Discovery path: /.well-known/oauth-authorization-server/mcp/{server_name}
"""
return _build_oauth_authorization_server_response(
request=request,
mcp_server_name=mcp_server_name,
)
# LiteLLM legacy pattern and root endpoint
@router.get(f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}")
@router.get("/.well-known/oauth-authorization-server")
async def oauth_authorization_server_mcp(
request: Request, mcp_server_name: Optional[str] = None
):
"""
OAuth authorization server discovery endpoint.
Supports both legacy pattern (/{server_name}) and root endpoint.
"""
return _build_oauth_authorization_server_response(
request=request,
mcp_server_name=mcp_server_name,
)
# Alias for standard OpenID discovery
@router.get("/.well-known/openid-configuration")
async def openid_configuration(request: Request):
return await oauth_authorization_server_mcp(request)
# Additional legacy pattern support
@router.get("/.well-known/oauth-authorization-server/{mcp_server_name}/mcp")
@router.get("/.well-known/oauth-authorization-server")
async def oauth_authorization_server_root(
request: Request, mcp_server_name: Optional[str] = None
async def oauth_authorization_server_legacy(
request: Request, mcp_server_name: str
):
return await oauth_authorization_server_mcp(request, mcp_server_name)
"""
OAuth authorization server discovery for legacy /{server_name}/mcp pattern.
"""
return _build_oauth_authorization_server_response(
request=request,
mcp_server_name=mcp_server_name,
)
@router.post("/{mcp_server_name}/register")
litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py
+6 -1
View File
@@ -46,8 +46,13 @@ class MCPGuardrailTranslationHandler(BaseTranslation):
)
return data
inputs = GenericGuardrailAPIInputs(texts=[content])
# Include model information if available
model = data.get("model")
if model:
inputs["model"] = model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=GenericGuardrailAPIInputs(texts=[content]),
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+24 -7
View File
@@ -63,7 +63,20 @@ from litellm.types.mcp_server.mcp_server_manager import (
MCPOAuthMetadata,
MCPServer,
)
from mcp.shared.tool_name_validation import SEP_986_URL, validate_tool_name
try:
from mcp.shared.tool_name_validation import SEP_986_URL, validate_tool_name # type: ignore
except ImportError:
SEP_986_URL = "https://github.com/modelcontextprotocol/protocol/blob/main/proposals/0001-tool-name-validation.md"
def validate_tool_name(name: str):
from pydantic import BaseModel
class MockResult(BaseModel):
is_valid: bool = True
warnings: list = []
return MockResult()
# Probe includes characters on both sides of the separator to mimic real prefixed tool names.
@@ -90,7 +103,9 @@ def _warn_on_server_name_fields(
if result.is_valid:
return
warning_text = "; ".join(result.warnings) if result.warnings else "Validation failed"
warning_text = (
"; ".join(result.warnings) if result.warnings else "Validation failed"
)
verbose_logger.warning(
"MCP server '%s' has invalid %s '%s': %s",
server_id,
@@ -103,7 +118,6 @@ def _warn_on_server_name_fields(
_warn("server_name", server_name)
def _deserialize_json_dict(data: Any) -> Optional[Dict[str, str]]:
"""
Deserialize optional JSON mappings stored in the database.
@@ -391,10 +405,13 @@ class MCPServerManager:
# Note: `extra_headers` on MCPServer is a List[str] of header names to forward
# from the client request (not available in this OpenAPI tool generation step).
# `static_headers` is a dict of concrete headers to always send.
headers = merge_mcp_headers(
extra_headers=headers,
static_headers=server.static_headers,
) or {}
headers = (
merge_mcp_headers(
extra_headers=headers,
static_headers=server.static_headers,
)
or {}
)
verbose_logger.debug(
f"Using headers for OpenAPI tools (excluding sensitive values): "
litellm/proxy/_experimental/mcp_server/server.py
+5 -1
View File
@@ -73,7 +73,11 @@ if MCP_AVAILABLE:
AuthContextMiddleware,
auth_context_var,
)
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
try:
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
except ImportError:
StreamableHTTPSessionManager = None # type: ignore
from mcp.types import (
CallToolResult,
EmbeddedResource,
litellm/proxy/_types.py
+8 -6
View File
@@ -425,12 +425,12 @@ class LiteLLMRoutes(enum.Enum):
]
google_routes = [
"/v1beta/models/{model_name}:countTokens",
"/v1beta/models/{model_name}:generateContent",
"/v1beta/models/{model_name}:streamGenerateContent",
"/models/{model_name}:countTokens",
"/models/{model_name}:generateContent",
"/models/{model_name}:streamGenerateContent",
"/v1beta/models/{model_name:path}:countTokens",
"/v1beta/models/{model_name:path}:generateContent",
"/v1beta/models/{model_name:path}:streamGenerateContent",
"/models/{model_name:path}:countTokens",
"/models/{model_name:path}:generateContent",
"/models/{model_name:path}:streamGenerateContent",
# Google Interactions API
"/interactions",
"/v1beta/interactions",
@@ -3428,6 +3428,8 @@ class LitellmMetadataFromRequestHeaders(TypedDict, total=False):
"""
spend_logs_metadata: Optional[dict]
agent_id: Optional[str]
trace_id: Optional[str]
class JWTKeyItem(TypedDict, total=False):
litellm/proxy/auth/auth_utils.py
+18 -2
View File
@@ -758,11 +758,27 @@ def get_model_from_request(
if match:
model = match.group(1)
# If still not found, extract model from Google generateContent-style routes.
# These routes put the model in the path and allow "/" inside the model id.
# Examples:
# - /v1beta/models/gemini-2.0-flash:generateContent
# - /v1beta/models/bedrock/claude-sonnet-3.7:generateContent
# - /models/custom/ns/model:streamGenerateContent
if model is None and not route.lower().startswith("/vertex"):
google_match = re.search(r"/(?:v1beta|beta)/models/([^:]+):", route)
if google_match:
model = google_match.group(1)
if model is None and not route.lower().startswith("/vertex"):
google_match = re.search(r"^/models/([^:]+):", route)
if google_match:
model = google_match.group(1)
# If still not found, extract from Vertex AI passthrough route
# Pattern: /vertex_ai/.../models/{model_id}:*
# Example: /vertex_ai/v1/.../models/gemini-1.5-pro:generateContent
if model is None and "/vertex" in route.lower():
vertex_match = re.search(r"/models/([^/:]+)", route)
if model is None and route.lower().startswith("/vertex"):
vertex_match = re.search(r"/models/([^:]+)", route)
if vertex_match:
model = vertex_match.group(1)
litellm/proxy/auth/login_utils.py
+6 -6
View File
@@ -197,9 +197,9 @@ async def authenticate_user( # noqa: PLR0915
- Login with UI_USERNAME and UI_PASSWORD
- Login with Invite Link `user_email` and `password` combination
"""
if secrets.compare_digest(username, ui_username) and secrets.compare_digest(
password, ui_password
):
if secrets.compare_digest(
username.encode("utf-8"), ui_username.encode("utf-8")
) and secrets.compare_digest(password.encode("utf-8"), ui_password.encode("utf-8")):
# Non SSO -> If user is using UI_USERNAME and UI_PASSWORD they are Proxy admin
user_role = LitellmUserRoles.PROXY_ADMIN
user_id = LITELLM_PROXY_ADMIN_NAME
@@ -313,9 +313,9 @@ async def authenticate_user( # noqa: PLR0915
# check if password == _user_row.password
hash_password = hash_token(token=password)
if secrets.compare_digest(password, _password) or secrets.compare_digest(
hash_password, _password
):
if secrets.compare_digest(
password.encode("utf-8"), _password.encode("utf-8")
) or secrets.compare_digest(hash_password.encode("utf-8"), _password.encode("utf-8")):
if os.getenv("DATABASE_URL") is not None:
# Expire any previous UI session tokens for this user
await expire_previous_ui_session_tokens(
litellm/proxy/auth/route_checks.py
+9 -1
View File
@@ -392,7 +392,15 @@ class RouteChecks:
# Ensure route is a string before attempting regex matching
if not isinstance(route, str):
return False
pattern = re.sub(r"\{[^}]+\}", r"[^/]+", pattern)
def _placeholder_to_regex(match: re.Match) -> str:
placeholder = match.group(0).strip("{}")
if placeholder.endswith(":path"):
# allow "/" in the placeholder value, but don't eat the route suffix after ":"
return r"[^:]+"
return r"[^/]+"
pattern = re.sub(r"\{[^}]+\}", _placeholder_to_regex, pattern)
# Anchor the pattern to match the entire string
pattern = f"^{pattern}$"
if re.match(pattern, route):
litellm/proxy/common_utils/callback_utils.py
+27 -18
View File
@@ -274,11 +274,20 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915
WebSearchInterceptionLogger,
)
websearch_interception_obj = WebSearchInterceptionLogger.initialize_from_proxy_config(
litellm_settings=litellm_settings,
callback_specific_params=callback_specific_params,
websearch_interception_obj = (
WebSearchInterceptionLogger.initialize_from_proxy_config(
litellm_settings=litellm_settings,
callback_specific_params=callback_specific_params,
)
)
imported_list.append(websearch_interception_obj)
elif isinstance(callback, str) and callback == "datadog_cost_management":
from litellm.integrations.datadog.datadog_cost_management import (
DatadogCostManagementLogger,
)
datadog_cost_management_obj = DatadogCostManagementLogger()
imported_list.append(datadog_cost_management_obj)
elif isinstance(callback, CustomLogger):
imported_list.append(callback)
else:
@@ -353,17 +362,17 @@ def get_remaining_tokens_and_requests_from_request_data(data: Dict) -> Dict[str,
remaining_requests_variable_name = f"litellm-key-remaining-requests-{model_group}"
remaining_requests = _metadata.get(remaining_requests_variable_name, None)
if remaining_requests:
headers[f"x-litellm-key-remaining-requests-{h11_model_group_name}"] = (
remaining_requests
)
headers[
f"x-litellm-key-remaining-requests-{h11_model_group_name}"
] = remaining_requests
# Remaining Tokens
remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}"
remaining_tokens = _metadata.get(remaining_tokens_variable_name, None)
if remaining_tokens:
headers[f"x-litellm-key-remaining-tokens-{h11_model_group_name}"] = (
remaining_tokens
)
headers[
f"x-litellm-key-remaining-tokens-{h11_model_group_name}"
] = remaining_tokens
return headers
@@ -438,9 +447,9 @@ def add_guardrail_response_to_standard_logging_object(
):
if litellm_logging_obj is None:
return
standard_logging_object: Optional[StandardLoggingPayload] = (
litellm_logging_obj.model_call_details.get("standard_logging_object")
)
standard_logging_object: Optional[
StandardLoggingPayload
] = litellm_logging_obj.model_call_details.get("standard_logging_object")
if standard_logging_object is None:
return
guardrail_information = standard_logging_object.get("guardrail_information", [])
@@ -469,7 +478,9 @@ def get_metadata_variable_name_from_kwargs(
return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
def process_callback(_callback: str, callback_type: str, environment_variables: dict) -> dict:
def process_callback(
_callback: str, callback_type: str, environment_variables: dict
) -> dict:
"""Process a single callback and return its data with environment variables"""
env_vars = CustomLogger.get_callback_env_vars(_callback)
@@ -481,11 +492,9 @@ def process_callback(_callback: str, callback_type: str, environment_variables:
else:
env_vars_dict[_var] = env_variable
return {
"name": _callback,
"variables": env_vars_dict,
"type": callback_type
}
return {"name": _callback, "variables": env_vars_dict, "type": callback_type}
def normalize_callback_names(callbacks: Iterable[Any]) -> List[Any]:
if callbacks is None:
return []
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py
+2
View File
@@ -185,6 +185,7 @@ class GenericGuardrailAPI(CustomGuardrail):
tools = inputs.get("tools")
structured_messages = inputs.get("structured_messages")
tool_calls = inputs.get("tool_calls")
model = inputs.get("model")
# Use provided request_data or create an empty dict
if request_data is None:
@@ -215,6 +216,7 @@ class GenericGuardrailAPI(CustomGuardrail):
tool_calls=tool_calls,
additional_provider_specific_params=additional_params,
input_type=input_type,
model=model,
)
# Prepare headers
litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py
+5 -2
View File
@@ -8,6 +8,7 @@ import os
import uuid
from typing import TYPE_CHECKING, Any, Literal, Optional, Type
import httpx
from fastapi import HTTPException
from litellm._logging import verbose_proxy_logger
@@ -25,10 +26,12 @@ if TYPE_CHECKING:
class OnyxGuardrail(CustomGuardrail):
def __init__(
self, api_base: Optional[str] = None, api_key: Optional[str] = None, **kwargs
self, api_base: Optional[str] = None, api_key: Optional[str] = None, timeout: Optional[float] = 10.0, **kwargs
):
timeout = timeout or int(os.getenv("ONYX_TIMEOUT", 10.0))
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback
llm_provider=httpxSpecialProvider.GuardrailCallback,
params={"timeout": httpx.Timeout(timeout=timeout, connect=5.0)},
)
self.api_base = api_base or os.getenv(
"ONYX_API_BASE",
litellm/proxy/litellm_pre_call_utils.py
+43 -9
View File
@@ -558,6 +558,16 @@ class LiteLLMProxyRequestSetup:
#########################################################################################
# Finally update the requests metadata with the `metadata_from_headers`
#########################################################################################
agent_id_from_header = headers.get("x-litellm-agent-id")
trace_id_from_header = headers.get("x-litellm-trace-id")
if agent_id_from_header:
metadata_from_headers["agent_id"] = agent_id_from_header
verbose_proxy_logger.debug(f"Extracted agent_id from header: {agent_id_from_header}")
if trace_id_from_header:
metadata_from_headers["trace_id"] = trace_id_from_header
verbose_proxy_logger.debug(f"Extracted trace_id from header: {trace_id_from_header}")
if isinstance(data[_metadata_variable_name], dict):
data[_metadata_variable_name].update(metadata_from_headers)
return data
@@ -846,7 +856,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915
# Add headers to metadata for guardrails to access (fixes #17477)
# Guardrails use metadata["headers"] to access request headers (e.g., User-Agent)
if _metadata_variable_name in data and isinstance(data[_metadata_variable_name], dict):
if _metadata_variable_name in data and isinstance(
data[_metadata_variable_name], dict
):
data[_metadata_variable_name]["headers"] = _headers
# check for forwardable headers
@@ -1002,7 +1014,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915
# User spend, budget - used by prometheus.py
# Follow same pattern as team and API key budgets
data[_metadata_variable_name]["user_api_key_user_spend"] = user_api_key_dict.user_spend
data[_metadata_variable_name][
"user_api_key_user_spend"
] = user_api_key_dict.user_spend
data[_metadata_variable_name][
"user_api_key_user_max_budget"
] = user_api_key_dict.user_max_budget
@@ -1029,8 +1043,8 @@ async def add_litellm_data_to_request( # noqa: PLR0915
## [Enterprise Only]
# Add User-IP Address
requester_ip_address = ""
if premium_user is True:
# Only set the IP Address for Enterprise Users
if True: # Always set the IP Address if available
# logic for tracking IP Address
# logic for tracking IP Address
if (
@@ -1050,6 +1064,16 @@ async def add_litellm_data_to_request( # noqa: PLR0915
requester_ip_address = request.client.host
data[_metadata_variable_name]["requester_ip_address"] = requester_ip_address
# Add User-Agent
user_agent = ""
if (
request is not None
and hasattr(request, "headers")
and "user-agent" in request.headers
):
user_agent = request.headers["user-agent"]
data[_metadata_variable_name]["user_agent"] = user_agent
# Check if using tag based routing
tags = LiteLLMProxyRequestSetup.add_request_tag_to_metadata(
llm_router=llm_router,
@@ -1532,7 +1556,9 @@ def add_guardrails_from_policy_engine(
f"policy_count={len(registry.get_all_policies())}"
)
if not registry.is_initialized():
verbose_proxy_logger.debug("Policy engine not initialized, skipping policy matching")
verbose_proxy_logger.debug(
"Policy engine not initialized, skipping policy matching"
)
return
# Build context from request
@@ -1550,13 +1576,17 @@ def add_guardrails_from_policy_engine(
# Get matching policies via attachments
matching_policy_names = PolicyMatcher.get_matching_policies(context=context)
verbose_proxy_logger.debug(f"Policy engine: matched policies via attachments: {matching_policy_names}")
verbose_proxy_logger.debug(
f"Policy engine: matched policies via attachments: {matching_policy_names}"
)
# Combine attachment-based policies with dynamic request body policies
all_policy_names = set(matching_policy_names)
if request_body_policies and isinstance(request_body_policies, list):
all_policy_names.update(request_body_policies)
verbose_proxy_logger.debug(f"Policy engine: added dynamic policies from request body: {request_body_policies}")
verbose_proxy_logger.debug(
f"Policy engine: added dynamic policies from request body: {request_body_policies}"
)
if not all_policy_names:
return
@@ -1567,7 +1597,9 @@ def add_guardrails_from_policy_engine(
context=context,
)
verbose_proxy_logger.debug(f"Policy engine: applied policies (conditions matched): {applied_policy_names}")
verbose_proxy_logger.debug(
f"Policy engine: applied policies (conditions matched): {applied_policy_names}"
)
# Track applied policies in metadata for response headers
for policy_name in applied_policy_names:
@@ -1578,7 +1610,9 @@ def add_guardrails_from_policy_engine(
# Resolve guardrails from matching policies
resolved_guardrails = PolicyResolver.resolve_guardrails_for_context(context=context)
verbose_proxy_logger.debug(f"Policy engine: resolved guardrails: {resolved_guardrails}")
verbose_proxy_logger.debug(
f"Policy engine: resolved guardrails: {resolved_guardrails}"
)
if not resolved_guardrails:
return
litellm/proxy/management_endpoints/mcp_management_endpoints.py
+14 -4
View File
@@ -56,7 +56,19 @@ except ImportError as e:
MCP_AVAILABLE = False
if MCP_AVAILABLE:
from mcp.shared.tool_name_validation import validate_tool_name
try:
from mcp.shared.tool_name_validation import validate_tool_name # type: ignore
except ImportError:
def validate_tool_name(name: str):
from pydantic import BaseModel
class MockResult(BaseModel):
is_valid: bool = True
warnings: list = []
return MockResult()
from litellm.proxy._experimental.mcp_server.db import (
create_mcp_server,
delete_mcp_server,
@@ -122,9 +134,7 @@ if MCP_AVAILABLE:
)
if validation_result.warnings:
error_messages_text = (
error_messages_text
+ "\n"
+ "\n".join(validation_result.warnings)
error_messages_text + "\n" + "\n".join(validation_result.warnings)
)
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
litellm/proxy/management_endpoints/ui_sso.py
+73 -23
View File
@@ -99,24 +99,24 @@ def determine_role_from_groups(
) -> Optional[LitellmUserRoles]:
"""
Determine the highest privilege role for a user based on their groups.
Role hierarchy (highest to lowest):
- proxy_admin
- proxy_admin_viewer
- internal_user
- internal_user_viewer
Args:
user_groups: List of group names from the SSO token
role_mappings: RoleMappings configuration object
Returns:
The highest privilege role found, or default_role if no matches, or None
"""
if not role_mappings.roles:
# No role mappings configured, return default_role
return role_mappings.default_role
# Role hierarchy (highest to lowest)
role_hierarchy = [
LitellmUserRoles.PROXY_ADMIN,
@@ -124,20 +124,22 @@ def determine_role_from_groups(
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
]
# Convert user_groups to a set for efficient lookup
user_groups_set = set(user_groups) if isinstance(user_groups, list) else set()
# Find the highest privilege role the user belongs to
for role in role_hierarchy:
if role in role_mappings.roles:
role_groups = role_mappings.roles[role]
if isinstance(role_groups, list) and user_groups_set.intersection(set(role_groups)):
if isinstance(role_groups, list) and user_groups_set.intersection(
set(role_groups)
):
verbose_proxy_logger.debug(
f"User groups {user_groups} matched role '{role.value}' via groups: {role_groups}"
)
return role
# No matching groups found, return default_role
verbose_proxy_logger.debug(
f"User groups {user_groups} did not match any role mappings, using default_role: {role_mappings.default_role}"
@@ -326,9 +328,7 @@ def generic_response_convertor(
"GENERIC_USER_PROVIDER_ATTRIBUTE", "provider"
)
generic_user_role_attribute_name = os.getenv(
"GENERIC_USER_ROLE_ATTRIBUTE", "role"
)
generic_user_role_attribute_name = os.getenv("GENERIC_USER_ROLE_ATTRIBUTE", "role")
verbose_proxy_logger.debug(
f" generic_user_id_attribute_name: {generic_user_id_attribute_name}\n generic_user_email_attribute_name: {generic_user_email_attribute_name}"
@@ -345,12 +345,15 @@ def generic_response_convertor(
# Determine user role based on role_mappings if available
# Only apply role_mappings for GENERIC SSO provider
user_role: Optional[LitellmUserRoles] = None
if role_mappings is not None and role_mappings.provider.lower() in ["generic", "okta"]:
if role_mappings is not None and role_mappings.provider.lower() in [
"generic",
"okta",
]:
# Use role_mappings to determine role from groups
group_claim = role_mappings.group_claim
user_groups_raw: Any = get_nested_value(response, group_claim)
# Handle different formats: could be a list, string (comma-separated), or single value
user_groups: List[str] = []
if isinstance(user_groups_raw, list):
@@ -361,7 +364,7 @@ def generic_response_convertor(
elif user_groups_raw is not None:
# Single value
user_groups = [str(user_groups_raw)]
if user_groups:
user_role = determine_role_from_groups(user_groups, role_mappings)
verbose_proxy_logger.debug(
@@ -373,10 +376,12 @@ def generic_response_convertor(
verbose_proxy_logger.debug(
f"No groups found in '{group_claim}', using default_role: {role_mappings.default_role}"
)
# Fallback to existing logic if role_mappings not used
if user_role is None:
user_role_from_sso = get_nested_value(response, generic_user_role_attribute_name)
user_role_from_sso = get_nested_value(
response, generic_user_role_attribute_name
)
if user_role_from_sso is not None:
role = get_litellm_user_role(user_role_from_sso)
if role is not None:
@@ -399,7 +404,9 @@ def generic_response_convertor(
)
def _setup_generic_sso_env_vars(generic_client_id: str, redirect_url: str) -> Tuple[str, List[str], str, str, str, bool]:
def _setup_generic_sso_env_vars(
generic_client_id: str, redirect_url: str
) -> Tuple[str, List[str], str, str, str, bool]:
"""Setup and validate Generic SSO environment variables."""
generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None)
generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(" ")
@@ -492,7 +499,43 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]:
verbose_proxy_logger.debug(
f"Could not load role_mappings from database: {e}. Continuing with existing role logic."
)
generic_role_mappings = os.getenv("GENERIC_ROLE_MAPPINGS_ROLES", None)
generic_role_mappings_group_claim = os.getenv(
"GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", None
)
generic_role_mappoings_default_role = os.getenv(
"GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", None
)
if generic_role_mappings is not None:
verbose_proxy_logger.debug(
"Found role_mappings for generic provider in environment variables"
)
import ast
try:
generic_user_role_mappings_data: Dict[
LitellmUserRoles, List[str]
] = ast.literal_eval(generic_role_mappings)
if isinstance(generic_user_role_mappings_data, dict):
from litellm.types.proxy.management_endpoints.ui_sso import (
RoleMappings,
)
role_mappings_data = {
"provider": "generic",
"group_claim": generic_role_mappings_group_claim,
"default_role": generic_role_mappoings_default_role,
"roles": generic_user_role_mappings_data,
}
role_mappings = RoleMappings(**role_mappings_data)
verbose_proxy_logger.debug(
f"Loaded role_mappings from environments for provider '{role_mappings.provider}'."
)
return role_mappings
except TypeError as e:
verbose_proxy_logger.warning(f"Error decoding role mappings from environment variables: {e}. Continuing with existing role logic.")
return role_mappings
@@ -529,7 +572,7 @@ async def get_generic_sso_response(
# Get role_mappings from SSO settings if available
role_mappings = await _setup_role_mappings()
def response_convertor(response, client):
nonlocal received_response # return for user debugging
received_response = response
@@ -1217,20 +1260,24 @@ async def insert_sso_user(
role_mappings_configured = False
try:
from litellm.proxy.utils import get_prisma_client_or_throw
prisma_client = get_prisma_client_or_throw(
"Prisma client is None, connect a database to your proxy"
)
# Get SSO config from dedicated table
sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique(
where={"id": "sso_config"}
)
if sso_db_record and sso_db_record.sso_settings:
sso_settings_dict = dict(sso_db_record.sso_settings)
role_mappings_data = sso_settings_dict.get("role_mappings")
role_mappings_configured = role_mappings_data is not None
generic_user_role_mappings = os.getenv("GENERIC_USER_ROLE_MAPPINGS", None)
if generic_user_role_mappings is not None:
role_mappings_configured = True
except Exception as e:
# If we can't check role_mappings, continue with existing logic
verbose_proxy_logger.debug(
@@ -1240,7 +1287,10 @@ async def insert_sso_user(
# Apply default_internal_user_params
if litellm.default_internal_user_params:
# If role_mappings is configured and user_role is already set from SSO, preserve it
if role_mappings_configured and user_defined_values.get("user_role") is not None:
if (
role_mappings_configured
and user_defined_values.get("user_role") is not None
):
# Preserve the SSO-extracted role, but apply other defaults
preserved_role = user_defined_values.get("user_role")
user_defined_values.update(litellm.default_internal_user_params) # type: ignore
litellm/proxy/proxy_server.py
+81 -14
View File
@@ -5420,6 +5420,24 @@ async def chat_completion( # noqa: PLR0915
global general_settings, user_debug, proxy_logging_obj, llm_model_list
global user_temperature, user_request_timeout, user_max_tokens, user_api_base
data = await _read_request_body(request=request)
if user_api_key_dict is not None:
if data.get("metadata") is None:
data["metadata"] = {}
if (
hasattr(user_api_key_dict, "user_id")
and user_api_key_dict.user_id is not None
):
data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id
if (
hasattr(user_api_key_dict, "team_id")
and user_api_key_dict.team_id is not None
):
data["metadata"]["user_api_key_team_id"] = user_api_key_dict.team_id
if (
hasattr(user_api_key_dict, "org_id")
and user_api_key_dict.org_id is not None
):
data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id
base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data)
try:
result = await base_llm_response_processor.base_process_llm_request(
@@ -5571,6 +5589,24 @@ async def completion( # noqa: PLR0915
data = {}
try:
data = await _read_request_body(request=request)
if user_api_key_dict is not None:
if data.get("metadata") is None:
data["metadata"] = {}
if (
hasattr(user_api_key_dict, "user_id")
and user_api_key_dict.user_id is not None
):
data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id
if (
hasattr(user_api_key_dict, "team_id")
and user_api_key_dict.team_id is not None
):
data["metadata"]["user_api_key_team_id"] = user_api_key_dict.team_id
if (
hasattr(user_api_key_dict, "org_id")
and user_api_key_dict.org_id is not None
):
data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id
base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data)
return await base_llm_response_processor.base_process_llm_request(
request=request,
@@ -5790,6 +5826,25 @@ async def embeddings( # noqa: PLR0915
)
data["input"] = input_list
if user_api_key_dict is not None:
if data.get("metadata") is None:
data["metadata"] = {}
if (
hasattr(user_api_key_dict, "user_id")
and user_api_key_dict.user_id is not None
):
data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id
if (
hasattr(user_api_key_dict, "team_id")
and user_api_key_dict.team_id is not None
):
data["metadata"]["user_api_key_team_id"] = user_api_key_dict.team_id
if (
hasattr(user_api_key_dict, "org_id")
and user_api_key_dict.org_id is not None
):
data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id
# Use unified request processor (same as chat/completions and responses)
base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data)
@@ -9645,7 +9700,7 @@ def get_logo_url():
@app.get("/get_image", include_in_schema=False)
def get_image():
async def get_image():
"""Get logo to show on admin UI"""
# get current_dir
@@ -9664,25 +9719,37 @@ def get_image():
if is_non_root and not os.path.exists(default_logo):
default_logo = default_site_logo
cache_dir = assets_dir if is_non_root else current_dir
cache_path = os.path.join(cache_dir, "cached_logo.jpg")
# [OPTIMIZATION] Check if the cached image exists first
if os.path.exists(cache_path):
return FileResponse(cache_path, media_type="image/jpeg")
logo_path = os.getenv("UI_LOGO_PATH", default_logo)
verbose_proxy_logger.debug("Reading logo from path: %s", logo_path)
# Check if the logo path is an HTTP/HTTPS URL
if logo_path.startswith(("http://", "https://")):
# Download the image and cache it
client = HTTPHandler()
response = client.get(logo_path)
if response.status_code == 200:
# Save the image to a local file
cache_dir = assets_dir if is_non_root else current_dir
cache_path = os.path.join(cache_dir, "cached_logo.jpg")
with open(cache_path, "wb") as f:
f.write(response.content)
try:
# Download the image and cache it
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
# Return the cached image as a FileResponse
return FileResponse(cache_path, media_type="image/jpeg")
else:
# Handle the case when the image cannot be downloaded
async_client = AsyncHTTPHandler(timeout=5.0)
response = await async_client.get(logo_path)
if response.status_code == 200:
# Save the image to a local file
with open(cache_path, "wb") as f:
f.write(response.content)
# Return the cached image as a FileResponse
return FileResponse(cache_path, media_type="image/jpeg")
else:
# Handle the case when the image cannot be downloaded
return FileResponse(default_logo, media_type="image/jpeg")
except Exception as e:
# Handle any exceptions during the download (e.g., timeout, connection error)
verbose_proxy_logger.debug(f"Error downloading logo from {logo_path}: {e}")
return FileResponse(default_logo, media_type="image/jpeg")
else:
# Return the local image file if the logo path is not an HTTP/HTTPS URL
litellm/proxy/rag_endpoints/endpoints.py
+198
View File
@@ -26,6 +26,184 @@ from litellm.proxy.common_utils.http_parsing_utils import (
router = APIRouter()
def _build_file_metadata_entry(
response: Any,
file_data: Optional[Tuple[str, bytes, str]] = None,
file_url: Optional[str] = None,
) -> Dict[str, Any]:
"""
Build a file metadata entry for storing in vector_store_metadata.
Args:
response: The response from litellm.aingest containing file_id
file_data: Optional tuple of (filename, content, content_type)
file_url: Optional URL if file was ingested from URL
Returns:
Dictionary with file metadata (file_id, filename, file_url, ingested_at, etc.)
"""
from datetime import datetime, timezone
# Extract file_id from response
file_id = None
if hasattr(response, "get"):
file_id = response.get("file_id")
elif hasattr(response, "file_id"):
file_id = response.file_id
# Extract file information from file_data tuple
filename = None
file_size = None
content_type = None
if file_data:
filename = file_data[0]
file_size = len(file_data[1]) if len(file_data) > 1 else None
content_type = file_data[2] if len(file_data) > 2 else None
# Build file metadata entry
file_entry = {
"file_id": file_id,
"filename": filename,
"file_url": file_url,
"ingested_at": datetime.now(timezone.utc).isoformat(),
}
# Add optional fields if available
if file_size is not None:
file_entry["file_size"] = file_size
if content_type is not None:
file_entry["content_type"] = content_type
return file_entry
async def _save_vector_store_to_db_from_rag_ingest(
response: Any,
ingest_options: Dict[str, Any],
prisma_client,
user_api_key_dict: UserAPIKeyAuth,
file_data: Optional[Tuple[str, bytes, str]] = None,
file_url: Optional[str] = None,
) -> None:
"""
Helper function to save a newly created vector store from RAG ingest to the database.
This function:
- Extracts vector store ID and config from the ingest response
- Checks if the vector store already exists in the database
- Creates a new database entry if it doesn't exist
- Adds the vector store to the registry
Args:
response: The response from litellm.aingest()
ingest_options: The ingest options containing vector store config
prisma_client: The Prisma database client
user_api_key_dict: User API key authentication info
"""
from litellm.proxy.vector_store_endpoints.management_endpoints import (
create_vector_store_in_db,
)
# Handle both dict and object responses
if hasattr(response, "get"):
vector_store_id = response.get("vector_store_id")
elif hasattr(response, "vector_store_id"):
vector_store_id = response.vector_store_id
else:
verbose_proxy_logger.warning(
f"Unable to extract vector_store_id from response type: {type(response)}"
)
return
if vector_store_id is None or not isinstance(vector_store_id, str):
verbose_proxy_logger.warning(
"Vector store ID is None or not a string, skipping database save"
)
return
vector_store_config = ingest_options.get("vector_store", {})
custom_llm_provider = vector_store_config.get("custom_llm_provider")
# Extract litellm_vector_store_params for custom name and description
litellm_vector_store_params = ingest_options.get("litellm_vector_store_params", {})
custom_vector_store_name = litellm_vector_store_params.get("vector_store_name")
custom_vector_store_description = litellm_vector_store_params.get("vector_store_description")
# Build file metadata entry using helper
file_entry = _build_file_metadata_entry(
response=response,
file_data=file_data,
file_url=file_url,
)
try:
# Check if vector store already exists in database
existing_vector_store = (
await prisma_client.db.litellm_managedvectorstorestable.find_unique(
where={"vector_store_id": vector_store_id}
)
)
# Only create if it doesn't exist
if existing_vector_store is None:
verbose_proxy_logger.info(
f"Saving newly created vector store {vector_store_id} to database"
)
# Initialize metadata with first file
initial_metadata = {
"ingested_files": [file_entry]
}
# Use custom name if provided, otherwise default
vector_store_name = custom_vector_store_name or f"RAG Vector Store - {vector_store_id[:8]}"
vector_store_description = custom_vector_store_description or "Created via RAG ingest endpoint"
await create_vector_store_in_db(
vector_store_id=vector_store_id,
custom_llm_provider=custom_llm_provider or "openai",
prisma_client=prisma_client,
vector_store_name=vector_store_name,
vector_store_description=vector_store_description,
vector_store_metadata=initial_metadata,
)
verbose_proxy_logger.info(
f"Vector store {vector_store_id} saved to database successfully"
)
else:
verbose_proxy_logger.info(
f"Vector store {vector_store_id} already exists, appending file to metadata"
)
# Update existing vector store with new file
existing_metadata = existing_vector_store.vector_store_metadata or {}
if isinstance(existing_metadata, str):
import json
existing_metadata = json.loads(existing_metadata)
ingested_files = existing_metadata.get("ingested_files", [])
ingested_files.append(file_entry)
existing_metadata["ingested_files"] = ingested_files
# Update the vector store
from litellm.proxy.utils import safe_dumps
await prisma_client.db.litellm_managedvectorstorestable.update(
where={"vector_store_id": vector_store_id},
data={"vector_store_metadata": safe_dumps(existing_metadata)}
)
verbose_proxy_logger.info(
f"Added file {file_entry.get('filename') or file_entry.get('file_url', 'Unknown')} to vector store {vector_store_id} metadata"
)
except Exception as db_error:
# Log the error but don't fail the request since ingestion succeeded
verbose_proxy_logger.exception(
f"Failed to save vector store {vector_store_id} to database: {db_error}"
)
async def parse_rag_ingest_request(
request: Request,
) -> Tuple[Dict[str, Any], Optional[Tuple[str, bytes, str]], Optional[str], Optional[str]]:
@@ -158,6 +336,7 @@ async def rag_ingest(
add_litellm_data_to_request,
general_settings,
llm_router,
prisma_client,
proxy_config,
version,
)
@@ -189,6 +368,25 @@ async def rag_ingest(
**request_data,
)
# Save vector store to database if it was newly created and prisma_client is available
verbose_proxy_logger.debug(
f"RAG Ingest - Checking database save conditions: prisma_client={prisma_client is not None}, response={response is not None}, response_type={type(response)}"
)
if prisma_client is not None and response is not None:
await _save_vector_store_to_db_from_rag_ingest(
response=response,
ingest_options=ingest_options,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
file_data=file_data,
file_url=file_url,
)
else:
verbose_proxy_logger.warning(
f"Skipping database save: prisma_client={prisma_client is not None}, response={response is not None}"
)
return response
except HTTPException:
litellm/proxy/spend_tracking/spend_tracking_utils.py
+1 -1
View File
@@ -396,7 +396,7 @@ def get_logging_payload( # noqa: PLR0915
)
# Extract agent_id for A2A requests (set directly on model_call_details)
agent_id: Optional[str] = kwargs.get("agent_id")
agent_id: Optional[str] = kwargs.get("agent_id") or metadata.get("agent_id")
custom_llm_provider = kwargs.get("custom_llm_provider")
raw_model = cast(str, kwargs.get("model") or "")
model_name = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {})
litellm/proxy/vector_store_endpoints/management_endpoints.py
+127 -58
View File
@@ -133,6 +133,112 @@ async def _resolve_embedding_config_from_db(
return None
########################################################
# Helper Functions
########################################################
async def create_vector_store_in_db(
vector_store_id: str,
custom_llm_provider: str,
prisma_client,
vector_store_name: Optional[str] = None,
vector_store_description: Optional[str] = None,
vector_store_metadata: Optional[Dict] = None,
litellm_params: Optional[Dict] = None,
litellm_credential_name: Optional[str] = None,
) -> LiteLLM_ManagedVectorStore:
"""
Helper function to create a vector store in the database.
This function handles:
- Checking if vector store already exists
- Creating the vector store in the database
- Adding it to the vector store registry
Returns:
LiteLLM_ManagedVectorStore: The created vector store object
Raises:
HTTPException: If vector store already exists or database error occurs
"""
from litellm.types.router import GenericLiteLLMParams
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database not connected")
# Check if vector store already exists
existing_vector_store = (
await prisma_client.db.litellm_managedvectorstorestable.find_unique(
where={"vector_store_id": vector_store_id}
)
)
if existing_vector_store is not None:
raise HTTPException(
status_code=400,
detail=f"Vector store with ID {vector_store_id} already exists",
)
# Prepare data for database
data_to_create: Dict[str, Any] = {
"vector_store_id": vector_store_id,
"custom_llm_provider": custom_llm_provider,
}
if vector_store_name is not None:
data_to_create["vector_store_name"] = vector_store_name
if vector_store_description is not None:
data_to_create["vector_store_description"] = vector_store_description
if vector_store_metadata is not None:
data_to_create["vector_store_metadata"] = safe_dumps(vector_store_metadata)
if litellm_credential_name is not None:
data_to_create["litellm_credential_name"] = litellm_credential_name
# Handle litellm_params - always provide at least an empty dict
if litellm_params:
# Auto-resolve embedding config if embedding model is provided but config is not
embedding_model = litellm_params.get("litellm_embedding_model")
if embedding_model and not litellm_params.get("litellm_embedding_config"):
resolved_config = await _resolve_embedding_config_from_db(
embedding_model=embedding_model,
prisma_client=prisma_client
)
if resolved_config:
litellm_params["litellm_embedding_config"] = resolved_config
verbose_proxy_logger.info(
f"Auto-resolved embedding config for model {embedding_model}"
)
litellm_params_dict = GenericLiteLLMParams(
**litellm_params
).model_dump(exclude_none=True)
data_to_create["litellm_params"] = safe_dumps(litellm_params_dict)
else:
# Provide empty dict if no litellm_params provided
data_to_create["litellm_params"] = safe_dumps({})
# Create in database
_new_vector_store = (
await prisma_client.db.litellm_managedvectorstorestable.create(
data=data_to_create
)
)
new_vector_store: LiteLLM_ManagedVectorStore = LiteLLM_ManagedVectorStore(
**_new_vector_store.model_dump()
)
# Add vector store to registry
if litellm.vector_store_registry is not None:
litellm.vector_store_registry.add_vector_store_to_registry(
vector_store=new_vector_store
)
verbose_proxy_logger.info(
f"Vector store {vector_store_id} created in database successfully"
)
return new_vector_store
########################################################
# Management Endpoints
########################################################
@@ -156,71 +262,34 @@ async def new_vector_store(
- vector_store_metadata: Optional[Dict] - Additional metadata for the vector store
"""
from litellm.proxy.proxy_server import prisma_client
from litellm.types.router import GenericLiteLLMParams
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database not connected")
try:
# Check if vector store already exists
existing_vector_store = (
await prisma_client.db.litellm_managedvectorstorestable.find_unique(
where={"vector_store_id": vector_store.get("vector_store_id")}
)
)
if existing_vector_store is not None:
vector_store_id = vector_store.get("vector_store_id")
custom_llm_provider = vector_store.get("custom_llm_provider")
if not vector_store_id or not custom_llm_provider:
raise HTTPException(
status_code=400,
detail=f"Vector store with ID {vector_store.get('vector_store_id')} already exists",
)
if vector_store.get("vector_store_metadata") is not None:
vector_store["vector_store_metadata"] = safe_dumps(
vector_store.get("vector_store_metadata")
)
# Safely handle JSON serialization of litellm_params
litellm_params_json: Optional[str] = None
_input_litellm_params: dict = vector_store.get("litellm_params", {}) or {}
if _input_litellm_params is not None:
# Auto-resolve embedding config if embedding model is provided but config is not
embedding_model = _input_litellm_params.get("litellm_embedding_model")
if embedding_model and not _input_litellm_params.get("litellm_embedding_config"):
resolved_config = await _resolve_embedding_config_from_db(
embedding_model=embedding_model,
prisma_client=prisma_client
)
if resolved_config:
_input_litellm_params["litellm_embedding_config"] = resolved_config
verbose_proxy_logger.info(
f"Auto-resolved embedding config for model {embedding_model}"
)
litellm_params_dict = GenericLiteLLMParams(
**_input_litellm_params
).model_dump(exclude_none=True)
litellm_params_json = safe_dumps(litellm_params_dict)
del vector_store["litellm_params"]
_new_vector_store = (
await prisma_client.db.litellm_managedvectorstorestable.create(
data={
**vector_store,
"litellm_params": litellm_params_json,
}
detail="vector_store_id and custom_llm_provider are required"
)
# Extract and validate metadata
metadata = vector_store.get("vector_store_metadata")
validated_metadata: Optional[Dict] = None
if metadata is not None and isinstance(metadata, dict):
validated_metadata = metadata
new_vector_store = await create_vector_store_in_db(
vector_store_id=vector_store_id,
custom_llm_provider=custom_llm_provider,
prisma_client=prisma_client,
vector_store_name=vector_store.get("vector_store_name"),
vector_store_description=vector_store.get("vector_store_description"),
vector_store_metadata=validated_metadata,
litellm_params=vector_store.get("litellm_params"),
litellm_credential_name=vector_store.get("litellm_credential_name"),
)
new_vector_store: LiteLLM_ManagedVectorStore = LiteLLM_ManagedVectorStore(
**_new_vector_store.model_dump()
)
# Add vector store to registry
if litellm.vector_store_registry is not None:
litellm.vector_store_registry.add_vector_store_to_registry(
vector_store=new_vector_store
)
return {
"status": "success",
"message": f"Vector store {vector_store.get('vector_store_id')} created successfully",
litellm/rag/main.py
+19 -6
View File
@@ -198,6 +198,10 @@ async def _execute_query_pipeline(
"""
Execute the RAG query pipeline.
"""
# Extract router from kwargs - use it for completion if available
# to properly resolve virtual model names
router: Optional["Router"] = kwargs.pop("router", None)
# 1. Extract query from last user message
query_text = RAGQuery.extract_query_from_messages(messages)
if not query_text:
@@ -233,12 +237,21 @@ async def _execute_query_pipeline(
context_message = RAGQuery.build_context_message(context_chunks)
modified_messages = messages[:-1] + [context_message] + [messages[-1]]
response = await litellm.acompletion(
model=model,
messages=modified_messages,
stream=stream,
**kwargs,
)
# Use router if available to properly resolve virtual model names
if router is not None:
response = await router.acompletion(
model=model,
messages=modified_messages,
stream=stream,
**kwargs,
)
else:
response = await litellm.acompletion(
model=model,
messages=modified_messages,
stream=stream,
**kwargs,
)
# 5. Attach search results to response
if not stream and isinstance(response, ModelResponse):
litellm/router.py
+5 -2
View File
@@ -1707,8 +1707,11 @@ class Router:
litellm_params = deployment.get("litellm_params", {})
dep_num_retries = litellm_params.get("num_retries")
if dep_num_retries is not None and isinstance(dep_num_retries, int):
exception.num_retries = dep_num_retries # type: ignore
if dep_num_retries is not None:
try:
exception.num_retries = int(dep_num_retries) # type: ignore # Handle both int and str
except (ValueError, TypeError):
pass # Skip if value can't be converted to int
def _update_kwargs_with_default_litellm_params(
self, kwargs: dict, metadata_variable_name: Optional[str] = "metadata"
litellm/types/integrations/datadog_cost_management.py
+27
View File
@@ -0,0 +1,27 @@
from typing import Dict, Optional, TypedDict
from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams
class DatadogCostManagementInitParams(StandardCustomLoggerInitParams):
"""
Init params for Datadog Cost Management
"""
datadog_cost_management_params: Optional[Dict] = None
class DatadogFOCUSCostEntry(TypedDict):
"""
Represents a single cost line item in the FOCUS format.
Ref: https://focus.finops.org/#specification
"""
ProviderName: str
ChargeDescription: str
ChargePeriodStart: str
ChargePeriodEnd: str
BilledCost: float
BillingCurrency: str
Tags: Optional[Dict[str, str]]
litellm/types/integrations/prometheus.py
+93
View File
@@ -150,6 +150,9 @@ class UserAPIKeyLabelNames(Enum):
FALLBACK_MODEL = "fallback_model"
ROUTE = "route"
MODEL_GROUP = "model_group"
CLIENT_IP = "client_ip"
USER_AGENT = "user_agent"
CALLBACK_NAME = "callback_name"
DEFINED_PROMETHEUS_METRICS = Literal[
@@ -196,6 +199,12 @@ DEFINED_PROMETHEUS_METRICS = Literal[
"litellm_cache_hits_metric",
"litellm_cache_misses_metric",
"litellm_cached_tokens_metric",
"litellm_deployment_tpm_limit",
"litellm_deployment_rpm_limit",
"litellm_remaining_api_key_requests_for_model",
"litellm_remaining_api_key_tokens_for_model",
"litellm_llm_api_failed_requests_metric",
"litellm_callback_logging_failures_metric",
]
@@ -209,6 +218,7 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.REQUESTED_MODEL.value,
UserAPIKeyLabelNames.END_USER.value,
UserAPIKeyLabelNames.USER.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
litellm_llm_api_time_to_first_token_metric = [
@@ -244,6 +254,7 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.TEAM_ALIAS.value,
UserAPIKeyLabelNames.USER.value,
UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
# Guardrail metrics - these use custom labels (guardrail_name, status, error_type, hook_type)
@@ -263,6 +274,8 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.STATUS_CODE.value,
UserAPIKeyLabelNames.USER_EMAIL.value,
UserAPIKeyLabelNames.ROUTE.value,
UserAPIKeyLabelNames.CLIENT_IP.value,
UserAPIKeyLabelNames.USER_AGENT.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
@@ -278,6 +291,8 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.EXCEPTION_STATUS.value,
UserAPIKeyLabelNames.EXCEPTION_CLASS.value,
UserAPIKeyLabelNames.ROUTE.value,
UserAPIKeyLabelNames.CLIENT_IP.value,
UserAPIKeyLabelNames.USER_AGENT.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
@@ -299,6 +314,7 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value,
UserAPIKeyLabelNames.API_KEY_HASH.value,
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
litellm_remaining_requests_metric = [
@@ -308,6 +324,7 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value,
UserAPIKeyLabelNames.API_KEY_HASH.value,
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
litellm_remaining_tokens_metric = [
@@ -317,6 +334,7 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value,
UserAPIKeyLabelNames.API_KEY_HASH.value,
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
litellm_requests_metric = [
@@ -328,6 +346,9 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.TEAM_ALIAS.value,
UserAPIKeyLabelNames.USER.value,
UserAPIKeyLabelNames.USER_EMAIL.value,
UserAPIKeyLabelNames.CLIENT_IP.value,
UserAPIKeyLabelNames.USER_AGENT.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
litellm_spend_metric = [
@@ -339,6 +360,9 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.TEAM_ALIAS.value,
UserAPIKeyLabelNames.USER.value,
UserAPIKeyLabelNames.USER_EMAIL.value,
UserAPIKeyLabelNames.CLIENT_IP.value,
UserAPIKeyLabelNames.USER_AGENT.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
litellm_input_tokens_metric = [
@@ -351,6 +375,7 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.USER.value,
UserAPIKeyLabelNames.USER_EMAIL.value,
UserAPIKeyLabelNames.REQUESTED_MODEL.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
litellm_total_tokens_metric = [
@@ -363,6 +388,7 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.USER.value,
UserAPIKeyLabelNames.USER_EMAIL.value,
UserAPIKeyLabelNames.REQUESTED_MODEL.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
litellm_output_tokens_metric = [
@@ -375,6 +401,7 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.USER.value,
UserAPIKeyLabelNames.USER_EMAIL.value,
UserAPIKeyLabelNames.REQUESTED_MODEL.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
litellm_deployment_state = [
@@ -384,6 +411,15 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.API_PROVIDER.value,
]
litellm_deployment_tpm_limit = [
UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value,
UserAPIKeyLabelNames.MODEL_ID.value,
UserAPIKeyLabelNames.API_BASE.value,
UserAPIKeyLabelNames.API_PROVIDER.value,
]
litellm_deployment_rpm_limit = litellm_deployment_tpm_limit
litellm_deployment_cooled_down = [
UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value,
UserAPIKeyLabelNames.MODEL_ID.value,
@@ -401,6 +437,7 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.TEAM_ALIAS.value,
UserAPIKeyLabelNames.EXCEPTION_STATUS.value,
UserAPIKeyLabelNames.EXCEPTION_CLASS.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
litellm_deployment_failed_fallbacks = litellm_deployment_successful_fallbacks
@@ -443,6 +480,26 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.USER.value,
]
litellm_user_budget_remaining_hours_metric = [
UserAPIKeyLabelNames.USER.value,
]
litellm_remaining_api_key_requests_for_model = [
UserAPIKeyLabelNames.API_KEY_HASH.value,
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value,
]
litellm_remaining_api_key_tokens_for_model = [
UserAPIKeyLabelNames.API_KEY_HASH.value,
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value,
]
litellm_callback_logging_failures_metric = [
UserAPIKeyLabelNames.CALLBACK_NAME.value,
]
# Add deployment metrics
litellm_deployment_failure_responses = [
UserAPIKeyLabelNames.REQUESTED_MODEL.value,
@@ -456,6 +513,8 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
UserAPIKeyLabelNames.TEAM.value,
UserAPIKeyLabelNames.TEAM_ALIAS.value,
UserAPIKeyLabelNames.CLIENT_IP.value,
UserAPIKeyLabelNames.USER_AGENT.value,
]
litellm_deployment_total_requests = [
@@ -468,10 +527,37 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
UserAPIKeyLabelNames.TEAM.value,
UserAPIKeyLabelNames.TEAM_ALIAS.value,
UserAPIKeyLabelNames.CLIENT_IP.value,
UserAPIKeyLabelNames.USER_AGENT.value,
]
litellm_deployment_success_responses = litellm_deployment_total_requests
litellm_remaining_api_key_requests_for_model = [
UserAPIKeyLabelNames.API_KEY_HASH.value,
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
litellm_remaining_api_key_tokens_for_model = [
UserAPIKeyLabelNames.API_KEY_HASH.value,
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
litellm_llm_api_failed_requests_metric = [
UserAPIKeyLabelNames.END_USER.value,
UserAPIKeyLabelNames.API_KEY_HASH.value,
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value,
UserAPIKeyLabelNames.TEAM.value,
UserAPIKeyLabelNames.TEAM_ALIAS.value,
UserAPIKeyLabelNames.USER.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
# Buffer monitoring metrics - these typically don't need additional labels
litellm_pod_lock_manager_size: List[str] = []
@@ -492,6 +578,7 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.TEAM_ALIAS.value,
UserAPIKeyLabelNames.END_USER.value,
UserAPIKeyLabelNames.USER.value,
UserAPIKeyLabelNames.MODEL_ID.value,
]
litellm_cache_hits_metric = _cache_metric_labels
@@ -584,6 +671,12 @@ class UserAPIKeyLabelValues(BaseModel):
route: Annotated[
Optional[str], Field(..., alias=UserAPIKeyLabelNames.ROUTE.value)
] = None
client_ip: Annotated[
Optional[str], Field(..., alias=UserAPIKeyLabelNames.CLIENT_IP.value)
] = None
user_agent: Annotated[
Optional[str], Field(..., alias=UserAPIKeyLabelNames.USER_AGENT.value)
] = None
class PrometheusMetricsConfig(BaseModel):
litellm/types/llms/bedrock.py
+85
View File
@@ -93,6 +93,67 @@ class GuardrailConverseContentBlock(TypedDict, total=False):
text: GuardrailConverseTextBlock
class CitationWebLocationBlock(TypedDict, total=False):
"""
Web location block for Nova grounding citations.
Contains the URL and domain from web search results.
Reference: https://docs.aws.amazon.com/nova/latest/userguide/grounding.html
"""
url: str
domain: str
class CitationLocationBlock(TypedDict, total=False):
"""
Location block containing the web location for a citation.
"""
web: CitationWebLocationBlock
class CitationReferenceBlock(TypedDict, total=False):
"""
Citation reference block containing a single citation with its location.
Each citation contains:
- location.web.url: The URL of the source
- location.web.domain: The domain of the source
"""
location: CitationLocationBlock
class CitationsContentBlock(TypedDict, total=False):
"""
Citations content block returned by Nova grounding (web search) tool.
When Nova grounding is enabled via systemTool, the model may return
citationsContent blocks containing web search citation references.
Reference: https://docs.aws.amazon.com/nova/latest/userguide/grounding.html
Example response structure:
{
"citationsContent": {
"citations": [
{
"location": {
"web": {
"url": "https://example.com/article",
"domain": "example.com"
}
}
}
]
}
}
"""
citations: List[CitationReferenceBlock]
class ContentBlock(TypedDict, total=False):
text: str
image: ImageBlock
@@ -103,6 +164,7 @@ class ContentBlock(TypedDict, total=False):
cachePoint: CachePointBlock
reasoningContent: BedrockConverseReasoningContentBlock
guardContent: GuardrailConverseContentBlock
citationsContent: CitationsContentBlock
class MessageBlock(TypedDict):
@@ -159,8 +221,24 @@ class ToolSpecBlock(TypedDict, total=False):
description: str
class SystemToolBlock(TypedDict, total=False):
"""
System tool block for Nova grounding and other built-in tools.
Example:
{
"systemTool": {
"name": "nova_grounding"
}
}
"""
name: Required[str]
class ToolBlock(TypedDict, total=False):
toolSpec: Optional[ToolSpecBlock]
systemTool: Optional[SystemToolBlock]
cachePoint: Optional[CachePointBlock]
@@ -210,11 +288,13 @@ class ContentBlockStartEvent(TypedDict, total=False):
class ContentBlockDeltaEvent(TypedDict, total=False):
"""
Either 'text' or 'toolUse' will be specified for Converse API streaming response.
May also include 'citationsContent' when Nova grounding is enabled.
"""
text: str
toolUse: ToolBlockDeltaEvent
reasoningContent: BedrockConverseReasoningContentBlockDelta
citationsContent: CitationsContentBlock
class PerformanceConfigBlock(TypedDict):
@@ -879,3 +959,8 @@ class BedrockGetBatchResponse(TypedDict, total=False):
outputDataConfig: BedrockOutputDataConfig
timeoutDurationInHours: Optional[int]
clientRequestToken: Optional[str]
class BedrockToolBlock(TypedDict, total=False):
toolSpec: Optional[ToolSpecBlock]
systemTool: Optional[SystemToolBlock] # For Nova grounding
cachePoint: Optional[CachePointBlock]
litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py
+9 -8
View File
@@ -51,19 +51,20 @@ class GenericGuardrailAPIRequest(BaseModel):
"""Request model for the Generic Guardrail API"""
input_type: Literal["request", "response"]
litellm_call_id: Optional[str] # the call id of the individual LLM call
litellm_call_id: Optional[str] = None # the call id of the individual LLM call
litellm_trace_id: Optional[
str
] # the trace id of the LLM call - useful if there are multiple LLM calls for the same conversation
structured_messages: Optional[List[AllMessageValues]]
images: Optional[List[str]]
tools: Optional[List[ChatCompletionToolParam]]
texts: Optional[List[str]]
] = None # the trace id of the LLM call - useful if there are multiple LLM calls for the same conversation
structured_messages: Optional[List[AllMessageValues]] = None
images: Optional[List[str]] = None
tools: Optional[List[ChatCompletionToolParam]] = None
texts: Optional[List[str]] = None
request_data: GenericGuardrailAPIMetadata
additional_provider_specific_params: Optional[Dict[str, Any]]
additional_provider_specific_params: Optional[Dict[str, Any]] = None
tool_calls: Optional[
Union[List[ChatCompletionToolCallChunk], List[ChatCompletionMessageToolCall]]
]
] = None
model: Optional[str] = None # the model being used for the LLM call
class GenericGuardrailAPIResponse:
litellm/types/proxy/guardrails/guardrail_hooks/onyx.py
+5
View File
@@ -16,6 +16,11 @@ class OnyxGuardrailConfigModel(GuardrailConfigModel):
description="The API key for the Onyx Guard server. If not provided, the `ONYX_API_KEY` environment variable is checked.",
)
timeout: Optional[float] = Field(
default=None,
description="The timeout for the Onyx Guard server in seconds. If not provided, the `ONYX_TIMEOUT` environment variable is checked.",
)
@staticmethod
def ui_friendly_name() -> str:
return "Onyx Guardrail"
litellm/types/utils.py
+15 -11
View File
@@ -3,25 +3,26 @@ import time
from enum import Enum
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Union
from aiohttp import FormData
from openai._models import BaseModel as OpenAIObject
from openai.types.audio.transcription_create_params import FileTypes # type: ignore
from openai.types.chat.chat_completion import ChatCompletion
from openai.types.audio.transcription_create_params import FileTypes as FileTypes # type: ignore
from openai.types.chat.chat_completion import ChatCompletion as ChatCompletion
from openai.types.completion_usage import (
CompletionTokensDetails,
CompletionUsage,
PromptTokensDetails,
)
from openai.types.moderation import (
Categories,
CategoryAppliedInputTypes,
CategoryScores,
Categories as Categories,
CategoryAppliedInputTypes as CategoryAppliedInputTypes,
CategoryScores as CategoryScores,
)
from openai.types.moderation_create_response import (
Moderation as Moderation,
ModerationCreateResponse as ModerationCreateResponse,
)
from openai.types.moderation_create_response import Moderation, ModerationCreateResponse
from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator
from typing_extensions import Callable, Dict, Required, TypedDict, override
from typing_extensions import Required, TypedDict
import litellm
from litellm._uuid import uuid
from litellm.types.llms.base import (
BaseLiteLLMOpenAIResponseObject,
@@ -52,7 +53,7 @@ from .llms.openai import (
ResponsesAPIResponse,
WebSearchOptions,
)
from .rerank import RerankResponse
from .rerank import RerankResponse as RerankResponse
if TYPE_CHECKING:
from .vector_stores import VectorStoreSearchResponse
@@ -1411,7 +1412,7 @@ class Usage(SafeAttributeModel, CompletionUsage):
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
"""Breakdown of tokens used in the prompt."""
def __init__(
def __init__( # noqa: PLR0915
self,
prompt_tokens: Optional[int] = None,
completion_tokens: Optional[int] = None,
@@ -2501,6 +2502,7 @@ class StandardLoggingMetadata(StandardLoggingUserAPIKeyMetadata):
dict
] # special param to log k,v pairs to spendlogs for a call
requester_ip_address: Optional[str]
user_agent: Optional[str]
requester_metadata: Optional[dict]
requester_custom_headers: Optional[
Dict[str, str]
@@ -2686,6 +2688,7 @@ class StandardLoggingPayload(TypedDict):
request_tags: list
end_user: Optional[str]
requester_ip_address: Optional[str]
user_agent: Optional[str]
messages: Optional[Union[str, list, dict]]
response: Optional[Union[str, list, dict]]
error_str: Optional[str]
@@ -3428,3 +3431,4 @@ class GenericGuardrailAPIInputs(TypedDict, total=False):
structured_messages: List[
AllMessageValues
] # structured messages sent to the LLM - indicates if text is from system or user
model: Optional[str] # the model being used for the LLM call
litellm/utils.py
+6
View File
@@ -8053,6 +8053,12 @@ class ProviderConfigManager:
)
return OpenrouterEmbeddingConfig()
elif litellm.LlmProviders.VERCEL_AI_GATEWAY == provider:
from litellm.llms.vercel_ai_gateway.embedding.transformation import (
VercelAIGatewayEmbeddingConfig,
)
return VercelAIGatewayEmbeddingConfig()
elif litellm.LlmProviders.GIGACHAT == provider:
return litellm.GigaChatEmbeddingConfig()
elif litellm.LlmProviders.SAGEMAKER == provider:
model_prices_and_context_window.json
+42
View File
@@ -10231,6 +10231,48 @@
"mode": "completion",
"output_cost_per_token": 5e-07
},
"deepseek-v3-2-251201": {
"input_cost_per_token": 0.0,
"litellm_provider": "volcengine",
"max_input_tokens": 98304,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 0.0,
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"glm-4-7-251222": {
"input_cost_per_token": 0.0,
"litellm_provider": "volcengine",
"max_input_tokens": 204800,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 0.0,
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"kimi-k2-thinking-251104": {
"input_cost_per_token": 0.0,
"litellm_provider": "volcengine",
"max_input_tokens": 229376,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 0.0,
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"doubao-embedding": {
"input_cost_per_token": 0.0,
"litellm_provider": "volcengine",
pyproject.toml
+3
View File
@@ -145,6 +145,7 @@ mypy = "^1.0"
pytest = "^7.4.3"
pytest-mock = "^3.12.0"
pytest-asyncio = "^0.21.1"
pytest-retry = "^1.6.3"
requests-mock = "^1.12.1"
responses = "^0.25.7"
respx = "^0.22.0"
@@ -183,6 +184,8 @@ plugins = "pydantic.mypy"
[tool.pytest.ini_options]
asyncio_mode = "auto"
retries = 20
retry_delay = 5
markers = [
"asyncio: mark test as an asyncio test",
"limit_leaks: mark test with memory limit for leak detection (e.g., '40 MB')",
tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
+1108
View File
File diff suppressed because it is too large Load Diff
tests/llm_translation/test_bedrock_completion.py
+285
View File
@@ -3954,3 +3954,288 @@ def test_bedrock_openai_error_handling():
assert exc_info.value.status_code == 422
print("✓ Error handling works correctly")
# ============================================================================
# Nova Grounding (web_search_options) Unit Tests (Mocked)
# ============================================================================
def test_bedrock_nova_grounding_web_search_options_non_streaming():
"""
Unit test for Nova grounding using web_search_options parameter (non-streaming).
This test mocks the HTTP call to verify:
1. web_search_options is correctly mapped to systemTool for Nova models
2. The request structure is correct
Related: https://docs.aws.amazon.com/nova/latest/userguide/grounding.html
"""
from unittest.mock import patch, MagicMock
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
messages = [
{
"role": "user",
"content": "What is the current population of Tokyo, Japan?",
}
]
with patch.object(client, "post") as mock_post:
try:
completion(
model="us.amazon.nova-pro-v1:0", # No bedrock/ prefix when using api_base
messages=messages,
web_search_options={}, # Enables Nova grounding
max_tokens=500,
client=client,
api_base="https://bedrock-runtime.us-east-1.amazonaws.com",
)
except Exception:
pass # Expected - we're just checking the request structure
# Verify the request was made correctly
if mock_post.called:
request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}"))
print(f"Request body: {json.dumps(request_body, indent=2)}")
# Verify toolConfig is present with systemTool
assert "toolConfig" in request_body, "toolConfig should be in request"
tool_config = request_body["toolConfig"]
assert "tools" in tool_config, "tools should be in toolConfig"
# Find the systemTool for nova_grounding
system_tool_found = False
for tool in tool_config["tools"]:
if "systemTool" in tool:
assert tool["systemTool"]["name"] == "nova_grounding"
system_tool_found = True
break
assert system_tool_found, "systemTool with nova_grounding should be present"
print(f"✓ web_search_options correctly transformed to systemTool (non-streaming)")
def test_bedrock_nova_grounding_with_function_tools():
"""
Unit test for Nova grounding combined with regular function tools.
This tests the scenario where users want both web grounding AND
custom function calling capabilities.
"""
from unittest.mock import patch
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
# Regular function tool
tools = [
{
"type": "function",
"function": {
"name": "get_stock_price",
"description": "Get the current stock price for a given ticker symbol",
"parameters": {
"type": "object",
"properties": {
"ticker": {
"type": "string",
"description": "The stock ticker symbol, e.g. AAPL, GOOGL",
}
},
"required": ["ticker"],
},
},
}
]
messages = [
{
"role": "user",
"content": "What is the current market cap of Apple Inc?",
}
]
with patch.object(client, "post") as mock_post:
try:
completion(
model="us.amazon.nova-pro-v1:0", # No bedrock/ prefix when using api_base
messages=messages,
tools=tools,
web_search_options={}, # Also enable web grounding
max_tokens=500,
client=client,
api_base="https://bedrock-runtime.us-east-1.amazonaws.com",
)
except Exception:
pass # Expected - we're just checking the request structure
# Verify the request was made correctly
if mock_post.called:
request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}"))
print(f"Request body: {json.dumps(request_body, indent=2)}")
# Verify toolConfig has both function tool and systemTool
assert "toolConfig" in request_body, "toolConfig should be in request"
tool_config = request_body["toolConfig"]
assert "tools" in tool_config, "tools should be in toolConfig"
tools_in_request = tool_config["tools"]
# Should have both the function tool and the systemTool
function_tool_found = False
system_tool_found = False
for tool in tools_in_request:
if "toolSpec" in tool:
assert tool["toolSpec"]["name"] == "get_stock_price"
function_tool_found = True
if "systemTool" in tool:
assert tool["systemTool"]["name"] == "nova_grounding"
system_tool_found = True
assert function_tool_found, "Function tool (get_stock_price) should be present"
assert system_tool_found, "systemTool (nova_grounding) should be present"
print(f"✓ Both function tools and web_search_options correctly combined")
@pytest.mark.asyncio
async def test_bedrock_nova_grounding_async():
"""
Async unit test for Nova grounding via web_search_options.
This test verifies the request transformation for async calls.
"""
from unittest.mock import patch, AsyncMock
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
client = AsyncHTTPHandler()
messages = [
{
"role": "user",
"content": "What is the weather forecast for New York City today?",
}
]
with patch.object(client, "post", new=AsyncMock()) as mock_post:
try:
await litellm.acompletion(
model="us.amazon.nova-pro-v1:0", # No bedrock/ prefix when using api_base
messages=messages,
web_search_options={},
max_tokens=500,
client=client,
api_base="https://bedrock-runtime.us-east-1.amazonaws.com",
)
except Exception:
pass # Expected - we're just checking the request structure
# Verify the request was made correctly
if mock_post.called:
request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}"))
print(f"Request body: {json.dumps(request_body, indent=2)}")
# Verify toolConfig is present with systemTool
assert "toolConfig" in request_body, "toolConfig should be in request"
tool_config = request_body["toolConfig"]
assert "tools" in tool_config, "tools should be in toolConfig"
# Find the systemTool for nova_grounding
system_tool_found = False
for tool in tool_config["tools"]:
if "systemTool" in tool:
assert tool["systemTool"]["name"] == "nova_grounding"
system_tool_found = True
break
assert system_tool_found, "systemTool with nova_grounding should be present"
print(f"✓ Async web_search_options correctly transformed to systemTool")
def test_bedrock_nova_web_search_options_ignored_for_non_nova():
"""
Test that web_search_options is ignored for non-Nova Bedrock models.
Nova grounding is only supported on Nova models. For other models,
the parameter should be silently ignored.
"""
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
config = AmazonConverseConfig()
# Should return None for non-Nova models
result = config._map_web_search_options({}, "anthropic.claude-3-sonnet-v1")
assert result is None
result = config._map_web_search_options({}, "amazon.titan-text-express-v1")
assert result is None
# Should return systemTool for Nova models
result = config._map_web_search_options({}, "amazon.nova-pro-v1:0")
assert result is not None
system_tool = result.get("systemTool")
assert system_tool is not None
assert system_tool["name"] == "nova_grounding"
result2 = config._map_web_search_options({}, "us.amazon.nova-premier-v1:0")
assert result2 is not None
system_tool2 = result2.get("systemTool")
assert system_tool2 is not None
assert system_tool2["name"] == "nova_grounding"
def test_bedrock_nova_grounding_request_transformation():
"""
Unit test to verify that web_search_options transforms to systemTool in the request.
"""
from unittest.mock import patch, MagicMock
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
messages = [{"role": "user", "content": "What is the population of Tokyo?"}]
with patch.object(client, "post") as mock_post:
mock_post.return_value = MagicMock(
status_code=200,
json=lambda: {
"output": {"message": {"role": "assistant", "content": [{"text": "Test"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 10, "outputTokens": 5}
}
)
try:
response = completion(
model="bedrock/us.amazon.nova-pro-v1:0",
messages=messages,
web_search_options={},
max_tokens=100,
client=client,
)
except Exception:
pass # Expected - we're just checking the request
if mock_post.called:
request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}"))
print(f"Request body: {json.dumps(request_body, indent=2)}")
# Verify toolConfig is present with systemTool
assert "toolConfig" in request_body, "toolConfig should be in request"
tool_config = request_body["toolConfig"]
assert "tools" in tool_config, "tools should be in toolConfig"
tools_in_request = tool_config["tools"]
# Find the systemTool
system_tool_found = False
for tool in tools_in_request:
if "systemTool" in tool:
assert tool["systemTool"]["name"] == "nova_grounding"
system_tool_found = True
break
assert system_tool_found, "systemTool with nova_grounding should be present"
print("✓ web_search_options correctly transformed to systemTool")
tests/mcp_tests/test_aresponses_api_with_mcp.py
+4 -2
View File
@@ -683,9 +683,11 @@ async def test_streaming_responses_api_with_mcp_tools(
Return the user the result of request 2
"""
# Skip test if ANTHROPIC_API_KEY is not set for anthropic models
if "anthropic" in model.lower() and not os.getenv("ANTHROPIC_API_KEY"):
# Skip test if required API keys are not set
if ("anthropic" in model.lower() or "claude" in model.lower()) and not os.getenv("ANTHROPIC_API_KEY"):
pytest.skip("ANTHROPIC_API_KEY not set, skipping anthropic model test")
if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv("OPENAI_API_KEY"):
pytest.skip("OPENAI_API_KEY not set, skipping openai model test")
from unittest.mock import AsyncMock, patch
tests/mcp_tests/test_mcp_server.py
+18 -8
View File
@@ -1105,14 +1105,24 @@ async def test_mcp_server_manager_config_integration_with_database():
test_manager.get_allowed_mcp_servers = mock_get_allowed_servers
# Mock _create_mcp_client to return a client that completes immediately
# This avoids network calls while preserving the actual conversion logic
def mock_create_mcp_client(*args, **kwargs):
mock_client = MagicMock()
mock_client.run_with_session = AsyncMock(return_value="ok")
return mock_client
test_manager._create_mcp_client = mock_create_mcp_client
# Mock health_check_server to avoid real network calls that timeout
async def mock_health_check(server_id: str, mcp_auth_header=None):
server = test_manager.get_mcp_server_by_id(server_id)
if not server:
return None
return LiteLLM_MCPServerTable(
server_id=server_id,
server_name=server.name,
url=server.url,
transport=server.transport,
description=server.mcp_info.get("description") if server.mcp_info else None,
mcp_access_groups=server.access_groups,
status="healthy",
last_health_check=datetime.datetime.now(),
mcp_info=server.mcp_info,
)
test_manager.health_check_server = mock_health_check
# Test the method (this tests our second fix)
servers_list = await test_manager.get_all_mcp_servers_with_health_and_teams(
tests/proxy_unit_tests/test_get_image.py
+89
View File
@@ -0,0 +1,89 @@
import os
import sys
from unittest import mock
# Standard path insertion
sys.path.insert(0, os.path.abspath("../.."))
import pytest
import httpx
from litellm.proxy.proxy_server import app
@pytest.mark.asyncio
async def test_get_image_error_handling():
"""
Test that get_image handles network errors gracefully and doesn't hang.
"""
# Set an unreachable URL
os.environ["UI_LOGO_PATH"] = "http://invalid-url-12345.com/logo.jpg"
# Clear cache
parent_dir = os.path.dirname(
os.path.dirname(
app.__file__
if hasattr(app, "__file__")
else "litellm/proxy/proxy_server.py"
)
)
cache_path = os.path.join(parent_dir, "proxy", "cached_logo.jpg")
if os.path.exists(cache_path):
os.remove(cache_path)
# Mock AsyncHTTPHandler to simulate a timeout or connection error
with mock.patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get"
) as mock_get:
mock_get.side_effect = httpx.ConnectError("Network is unreachable")
async with httpx.AsyncClient(
transport=httpx.ASGITransport(app=app), base_url="http://testserver"
) as ac:
response = await ac.get("/get_image")
assert response.status_code == 200
assert response.headers["content-type"] == "image/jpeg"
@pytest.mark.asyncio
async def test_get_image_cache_logic():
"""
Test that once cached, get_image doesn't hit the network.
"""
os.environ["UI_LOGO_PATH"] = "http://example.com/logo.jpg"
# Clear cache
parent_dir = os.path.dirname(
os.path.dirname(
app.__file__
if hasattr(app, "__file__")
else "litellm/proxy/proxy_server.py"
)
)
cache_path = os.path.join(parent_dir, "proxy", "cached_logo.jpg")
if os.path.exists(cache_path):
os.remove(cache_path)
# Mock response
mock_response = mock.Mock()
mock_response.status_code = 200
mock_response.content = b"fake image data"
with mock.patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get"
) as mock_get:
mock_get.return_value = mock_response
async with httpx.AsyncClient(
transport=httpx.ASGITransport(app=app), base_url="http://testserver"
) as ac:
# First call - should hit download logic
response1 = await ac.get("/get_image")
assert response1.status_code == 200
assert mock_get.call_count == 1
# Second call - should hit cache
response2 = await ac.get("/get_image")
assert response2.status_code == 200
# If cache works, mock_get shouldn't be called again
assert mock_get.call_count == 1
tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py
+22
View File
@@ -2,7 +2,9 @@ import os
import sys
import unittest.mock as mock
import httpx
import pytest
import respx
from httpx import Response
sys.path.insert(0, os.path.abspath("../../.."))
@@ -11,6 +13,8 @@ from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import (
ResendEmailLogger,
)
# Test file for Resend email integration
@pytest.fixture
def mock_env_vars():
@@ -37,7 +41,13 @@ def mock_httpx_client():
@pytest.mark.asyncio
@respx.mock
async def test_send_email_success(mock_env_vars, mock_httpx_client):
# Block all HTTP requests at network level to prevent real API calls
respx.post("https://api.resend.com/emails").mock(
return_value=httpx.Response(200, json={"id": "test_email_id"})
)
# Initialize the logger
logger = ResendEmailLogger()
@@ -71,7 +81,13 @@ async def test_send_email_success(mock_env_vars, mock_httpx_client):
@pytest.mark.asyncio
@respx.mock
async def test_send_email_missing_api_key(mock_httpx_client):
# Block all HTTP requests at network level to prevent real API calls
respx.post("https://api.resend.com/emails").mock(
return_value=httpx.Response(200, json={"id": "test_email_id"})
)
# Remove the API key from environment before initializing logger
original_key = os.environ.pop("RESEND_API_KEY", None)
@@ -109,7 +125,13 @@ async def test_send_email_missing_api_key(mock_httpx_client):
@pytest.mark.asyncio
@respx.mock
async def test_send_email_multiple_recipients(mock_env_vars, mock_httpx_client):
# Block all HTTP requests at network level to prevent real API calls
respx.post("https://api.resend.com/emails").mock(
return_value=httpx.Response(200, json={"id": "test_email_id"})
)
# Initialize the logger
logger = ResendEmailLogger()
tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py
+8
View File
@@ -2,7 +2,9 @@ import os
import sys
import unittest.mock as mock
import httpx
import pytest
import respx
from httpx import Response
sys.path.insert(0, os.path.abspath("../../.."))
@@ -101,7 +103,13 @@ async def test_send_email_missing_api_key():
@pytest.mark.asyncio
@respx.mock
async def test_send_email_multiple_recipients(mock_env_vars, mock_httpx_client):
# Block all HTTP requests at network level to prevent real API calls
respx.post("https://api.sendgrid.com/v3/mail/send").mock(
return_value=httpx.Response(202, text="accepted")
)
logger = SendGridEmailLogger()
from_email = "test@example.com"
tests/test_litellm/integrations/datadog/test_datadog_cost_management.py
+169
View File
@@ -0,0 +1,169 @@
import os
import time
from unittest.mock import AsyncMock
import pytest
from httpx import Response
from litellm.integrations.datadog.datadog_cost_management import (
DatadogCostManagementLogger,
)
from litellm.types.utils import StandardLoggingPayload
@pytest.fixture
def clean_env():
# Save original env
original_api_key = os.environ.get("DD_API_KEY")
original_app_key = os.environ.get("DD_APP_KEY")
original_site = os.environ.get("DD_SITE")
# Set test env
os.environ["DD_API_KEY"] = "test_api_key"
os.environ["DD_APP_KEY"] = "test_app_key"
os.environ["DD_SITE"] = "test.datadoghq.com"
yield
# Restore original env
if original_api_key:
os.environ["DD_API_KEY"] = original_api_key
else:
del os.environ["DD_API_KEY"]
if original_app_key:
os.environ["DD_APP_KEY"] = original_app_key
else:
del os.environ["DD_APP_KEY"]
if original_site:
os.environ["DD_SITE"] = original_site
else:
del os.environ["DD_SITE"]
@pytest.mark.asyncio
async def test_init(clean_env):
"""
Test initialization sets up clients and url correctly
"""
logger = DatadogCostManagementLogger()
assert logger.dd_api_key == "test_api_key"
assert logger.dd_app_key == "test_app_key"
assert (
logger.upload_url == "https://api.test.datadoghq.com/api/v2/cost/custom_costs"
)
@pytest.mark.asyncio
async def test_aggregate_costs(clean_env):
"""
Test that costs are correctly aggregated by provider, model, and date
"""
logger = DatadogCostManagementLogger()
# Mock some log payloads
now = time.time()
day_str = time.strftime("%Y-%m-%d", time.localtime(now))
logs = [
StandardLoggingPayload(
custom_llm_provider="openai",
model="gpt-4",
response_cost=0.01,
startTime=now,
metadata={"user_api_key_team_alias": "team-a"},
),
StandardLoggingPayload(
custom_llm_provider="openai",
model="gpt-4",
response_cost=0.02,
startTime=now,
metadata={"user_api_key_team_alias": "team-a"},
),
StandardLoggingPayload(
custom_llm_provider="anthropic",
model="claude-3",
response_cost=0.05,
startTime=now,
),
]
aggregated = logger._aggregate_costs(logs)
assert len(aggregated) == 2
# Check OpenAI entry
openai_entry = next(e for e in aggregated if e["ProviderName"] == "openai")
assert openai_entry["BilledCost"] == 0.03
assert openai_entry["ChargeDescription"] == "LLM Usage for gpt-4"
assert openai_entry["ChargePeriodStart"] == day_str
assert openai_entry["Tags"]["team"] == "team-a"
assert "env" in openai_entry["Tags"]
assert "service" in openai_entry["Tags"]
# Check Anthropic entry
anthropic_entry = next(e for e in aggregated if e["ProviderName"] == "anthropic")
assert anthropic_entry["BilledCost"] == 0.05
@pytest.mark.asyncio
async def test_async_log_success_event(clean_env):
"""
Test that logs are added to queue
"""
logger = DatadogCostManagementLogger(batch_size=10)
await logger.async_log_success_event(
kwargs={"standard_logging_object": {"response_cost": 0.01}},
response_obj={},
start_time=time.time(),
end_time=time.time(),
)
assert len(logger.log_queue) == 1
assert logger.log_queue[0]["response_cost"] == 0.01
# Test zero cost ignored
await logger.async_log_success_event(
kwargs={"standard_logging_object": {"response_cost": 0.0}},
response_obj={},
start_time=time.time(),
end_time=time.time(),
)
assert len(logger.log_queue) == 1
@pytest.mark.asyncio
async def test_async_send_batch(clean_env):
"""
Test that batch is aggregated and uploaded
"""
logger = DatadogCostManagementLogger()
logger.async_client = AsyncMock()
logger.async_client.put.return_value = Response(202, json={"status": "ok"})
# Add logs directly to queue
logger.log_queue = [
StandardLoggingPayload(
custom_llm_provider="openai",
model="gpt-4",
response_cost=0.01,
startTime=time.time(),
)
]
await logger.async_send_batch()
# Verify API called
assert logger.async_client.put.called
call_args = logger.async_client.put.call_args
assert call_args[0][0] == "https://api.test.datadoghq.com/api/v2/cost/custom_costs"
import json
# Use call_args.kwargs['content']
content = json.loads(call_args[1]["content"])
assert content[0]["ProviderName"] == "openai"
assert content[0]["BilledCost"] == 0.01
tests/test_litellm/integrations/datadog/test_datadog_llm_obs_agent.py
+62
View File
@@ -0,0 +1,62 @@
import os
from unittest.mock import patch
from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger
def test_datadog_llm_obs_agent_configuration():
"""
Test that DataDog LLM Obs logger correctly configures agent endpoint.
"""
test_env = {
"LITELLM_DD_AGENT_HOST": "localhost",
"LITELLM_DD_LLM_OBS_PORT": "10518",
"DD_API_KEY": "test-api-key", # Optional, but checking if it's preserved
}
# Ensure DD_SITE is NOT set to verify we don't need it in agent mode
with patch.dict(os.environ, test_env, clear=True):
with patch("asyncio.create_task"): # Prevent periodic flush task from running
dd_logger = DataDogLLMObsLogger()
expected_url = "http://localhost:10518/api/intake/llm-obs/v1/trace/spans"
assert dd_logger.intake_url == expected_url
assert dd_logger.DD_API_KEY == "test-api-key"
def test_datadog_llm_obs_agent_no_api_key_ok():
"""
Test that agent mode works WITHOUT DD_API_KEY (agent handles auth).
"""
test_env = {
"LITELLM_DD_AGENT_HOST": "localhost",
# No DD_API_KEY
}
with patch.dict(os.environ, test_env, clear=True):
with patch("asyncio.create_task"):
# Should NOT raise exception anymore
dd_logger = DataDogLLMObsLogger()
assert dd_logger.DD_API_KEY is None
# Default port is 8126 if not set
expected_url = "http://localhost:8126/api/intake/llm-obs/v1/trace/spans"
assert dd_logger.intake_url == expected_url
def test_datadog_llm_obs_direct_api_configuration():
"""
Test that direct API configuration still works as expected.
"""
test_env = {
"DD_API_KEY": "direct-api-key",
"DD_SITE": "us5.datadoghq.com",
}
with patch.dict(os.environ, test_env, clear=True):
with patch("asyncio.create_task"):
dd_logger = DataDogLLMObsLogger()
expected_url = "https://api.us5.datadoghq.com/api/intake/llm-obs/v1/trace/spans"
assert dd_logger.intake_url == expected_url
assert dd_logger.DD_API_KEY == "direct-api-key"
tests/test_litellm/integrations/test_custom_guardrail_recursion.py
+73
View File
@@ -0,0 +1,73 @@
import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.types.guardrails import GuardrailEventHooks
import json
class TestCustomGuardrailRecursion:
"""
Specific tests for the circular reference / RecursionError fix in logging.
"""
def test_log_guardrail_information_handles_circular_references(self):
"""
Test that add_standard_logging method sanitizes input data containing circular references
instead of crashing.
This reproduces the Langfuse crash scenario:
Request -> Metadata -> GuardrailResponse -> DebugContext -> Request
"""
guardrail = CustomGuardrail(
guardrail_name="recursion_test_guardrail",
event_hook=GuardrailEventHooks.pre_call,
)
# 1. Setup Circular Data
request_data = {"user_id": "test_recursive_user"}
metadata = {"session_id": "123"}
request_data["metadata"] = metadata
# Create the danger: Guardrail Response holding a reference back to request_data
dirty_response = {
"flagged": False,
"debug_context": request_data, # <--- ACCESS TO ROOT (Circular Ref)
}
# 2. Invoke the logging method
# If the fix is working, this will NOT raise RecursionError
try:
guardrail.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=dirty_response,
request_data=request_data,
guardrail_status="success",
start_time=1.0,
end_time=2.0,
duration=1.0,
masked_entity_count={},
event_type=GuardrailEventHooks.pre_call,
)
except RecursionError:
pytest.fail(
"RecursionError raised! The cyclic reference sanitization failed."
)
# 3. Verify the data stored is safe
stored_info = request_data["metadata"][
"standard_logging_guardrail_information"
][0]
stored_response = stored_info["guardrail_response"]
# Check that we can dump it to JSON without crashing (Ultimate proof)
try:
json.dumps(stored_response)
except Exception as e:
pytest.fail(f"Stored data is not JSON serializable: {e}")
# Check content - keys should be preserved but recursion broken
assert "debug_context" in stored_response
debug_context = stored_response["debug_context"]
# In a sanitized copy, the nested metadata should be a copy, not the original live dict
assert debug_context["user_id"] == "test_recursive_user"
# The 'metadata' inside 'debug_context' would be where recursion stops or is filtered
assert "metadata" in debug_context
tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py
+203
View File
@@ -0,0 +1,203 @@
import pytest
from unittest.mock import MagicMock, patch
from litellm.integrations.prometheus import PrometheusLogger
from litellm.types.integrations.prometheus import (
UserAPIKeyLabelValues,
)
from litellm.proxy._types import UserAPIKeyAuth
@pytest.mark.asyncio
async def test_async_post_call_failure_hook_includes_client_ip_user_agent():
"""
Test that async_post_call_failure_hook includes client_ip and user_agent in UserAPIKeyLabelValues
"""
# Mocking
# Mocking
with patch(
"litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None
):
logger = PrometheusLogger()
# Initialize attributes manually as __init__ is mocked
logger.litellm_proxy_failed_requests_metric = MagicMock()
logger.litellm_proxy_total_requests_metric = MagicMock()
logger.get_labels_for_metric = MagicMock(
return_value=["client_ip", "user_agent"]
)
request_data = {
"model": "gpt-4",
"metadata": {
"requester_ip_address": "127.0.0.1",
"user_agent": "test-agent",
},
}
user_api_key_dict = UserAPIKeyAuth(token="test_token")
original_exception = Exception("Test exception")
# Mock prometheus_label_factory to inspect arguments
with patch(
"litellm.integrations.prometheus.prometheus_label_factory"
) as mock_label_factory:
mock_label_factory.return_value = {}
await logger.async_post_call_failure_hook(
request_data=request_data,
original_exception=original_exception,
user_api_key_dict=user_api_key_dict,
)
# Verification
assert mock_label_factory.call_count >= 1
# Check calls
calls = mock_label_factory.call_args_list
found = False
for call in calls:
kwargs = call.kwargs
enum_values = kwargs.get("enum_values")
if isinstance(enum_values, UserAPIKeyLabelValues):
if (
enum_values.client_ip == "127.0.0.1"
and enum_values.user_agent == "test-agent"
):
found = True
break
assert (
found
), "UserAPIKeyLabelValues should contain client_ip='127.0.0.1' and user_agent='test-agent'"
@pytest.mark.asyncio
async def test_async_post_call_success_hook_includes_client_ip_user_agent():
"""
Test that async_post_call_success_hook includes client_ip and user_agent in UserAPIKeyLabelValues
"""
# Mocking
# Mocking
with patch(
"litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None
):
logger = PrometheusLogger()
logger.litellm_proxy_total_requests_metric = MagicMock()
logger.get_labels_for_metric = MagicMock(
return_value=["client_ip", "user_agent"]
)
data = {
"model": "gpt-4",
"metadata": {
"requester_ip_address": "192.168.1.1",
"user_agent": "success-agent",
},
}
user_api_key_dict = UserAPIKeyAuth(token="test_token")
response = MagicMock()
# Mock prometheus_label_factory to inspect arguments
with patch(
"litellm.integrations.prometheus.prometheus_label_factory"
) as mock_label_factory:
mock_label_factory.return_value = {}
await logger.async_post_call_success_hook(
data=data,
user_api_key_dict=user_api_key_dict,
response=response,
)
# Verification
assert mock_label_factory.call_count >= 1
# Check calls
calls = mock_label_factory.call_args_list
found = False
for call in calls:
kwargs = call.kwargs
enum_values = kwargs.get("enum_values")
if isinstance(enum_values, UserAPIKeyLabelValues):
if (
enum_values.client_ip == "192.168.1.1"
and enum_values.user_agent == "success-agent"
):
found = True
break
assert (
found
), "UserAPIKeyLabelValues should contain client_ip='192.168.1.1' and user_agent='success-agent'"
def test_set_llm_deployment_failure_metrics_includes_client_ip_user_agent():
"""
Test that set_llm_deployment_failure_metrics includes client_ip and user_agent in UserAPIKeyLabelValues
"""
# Mocking
# Mocking
with patch(
"litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None
):
logger = PrometheusLogger()
logger.litellm_deployment_failure_responses = MagicMock()
logger.litellm_deployment_total_requests = MagicMock()
logger.get_labels_for_metric = MagicMock(
return_value=["client_ip", "user_agent"]
)
logger.set_deployment_partial_outage = MagicMock()
request_kwargs = {
"model": "gpt-4",
"standard_logging_object": {
"metadata": {
"requester_ip_address": "10.0.0.1",
"user_agent": "failure-deployment",
"user_api_key_team_id": "team_1",
"user_api_key_team_alias": "team_alias_1",
"user_api_key_alias": "key_alias_1",
},
"model_group": "group_1",
"api_base": "http://api.base",
"model_id": "model_1",
},
"litellm_params": {},
"exception": Exception("Deployment failure"),
}
# Mock prometheus_label_factory to inspect arguments
with patch(
"litellm.integrations.prometheus.prometheus_label_factory"
) as mock_label_factory:
mock_label_factory.return_value = {}
logger.set_llm_deployment_failure_metrics(request_kwargs=request_kwargs)
# Verification
assert mock_label_factory.call_count >= 1
# Check calls
calls = mock_label_factory.call_args_list
found = False
for call in calls:
kwargs = call.kwargs
enum_values = kwargs.get("enum_values")
if isinstance(enum_values, UserAPIKeyLabelValues):
if (
enum_values.client_ip == "10.0.0.1"
and enum_values.user_agent == "failure-deployment"
):
found = True
break
assert (
found
), "UserAPIKeyLabelValues should contain client_ip='10.0.0.1' and user_agent='failure-deployment'"
if __name__ == "__main__":
import asyncio
asyncio.run(test_async_post_call_failure_hook_includes_client_ip_user_agent())
asyncio.run(test_async_post_call_success_hook_includes_client_ip_user_agent())
test_set_llm_deployment_failure_metrics_includes_client_ip_user_agent()
print("✅ All client_ip and user_agent tests passed!")
tests/test_litellm/integrations/test_prometheus_labels.py
+105 -44
View File
@@ -26,15 +26,49 @@ def test_user_email_in_required_metrics():
"litellm_input_tokens_metric",
"litellm_output_tokens_metric",
"litellm_requests_metric",
"litellm_spend_metric"
"litellm_spend_metric",
]
for metric_name in metrics_with_user_email:
labels = PrometheusMetricLabels.get_labels(metric_name)
assert user_email_label in labels, f"Metric {metric_name} should contain user_email label"
assert (
user_email_label in labels
), f"Metric {metric_name} should contain user_email label"
print(f"✅ {metric_name} contains user_email label")
def test_model_id_in_required_metrics():
"""
Test that model_id label is present in all the metrics that should have it
"""
model_id_label = UserAPIKeyLabelNames.MODEL_ID.value
# Metrics that should have model_id
metrics_with_model_id = [
"litellm_proxy_total_requests_metric",
"litellm_proxy_failed_requests_metric",
"litellm_input_tokens_metric",
"litellm_output_tokens_metric",
"litellm_requests_metric",
"litellm_spend_metric",
"litellm_llm_api_latency_metric",
"litellm_remaining_requests_metric",
"litellm_deployment_successful_fallbacks",
"litellm_cache_hits_metric",
"litellm_cache_misses_metric",
"litellm_remaining_api_key_requests_for_model",
"litellm_remaining_api_key_tokens_for_model",
"litellm_llm_api_failed_requests_metric",
]
for metric_name in metrics_with_model_id:
labels = PrometheusMetricLabels.get_labels(metric_name)
assert (
model_id_label in labels
), f"Metric {metric_name} should contain model_id label"
print(f"✅ {metric_name} contains model_id label")
def test_user_email_label_exists():
"""Test that the USER_EMAIL label is properly defined"""
assert UserAPIKeyLabelNames.USER_EMAIL.value == "user_email"
@@ -52,12 +86,14 @@ def test_prometheus_metric_labels_structure():
"litellm_proxy_failed_requests_metric",
"litellm_input_tokens_metric",
"litellm_output_tokens_metric",
"litellm_spend_metric"
"litellm_spend_metric",
]
for metric_name in test_metrics:
# Check metric is in DEFINED_PROMETHEUS_METRICS
assert metric_name in get_args(DEFINED_PROMETHEUS_METRICS), f"{metric_name} should be in DEFINED_PROMETHEUS_METRICS"
assert metric_name in get_args(
DEFINED_PROMETHEUS_METRICS
), f"{metric_name} should be in DEFINED_PROMETHEUS_METRICS"
# Check labels can be retrieved
labels = PrometheusMetricLabels.get_labels(metric_name)
@@ -98,11 +134,11 @@ def test_route_normalization_for_responses_api():
"""
Test that route normalization prevents high cardinality in Prometheus metrics
for the /v1/responses/{response_id} endpoint.
Issue: https://github.com/BerriAI/litellm/issues/XXXX
Each unique response ID was creating a separate metric line, causing the
/metrics endpoint to grow to ~30MB and take ~40 seconds to respond.
Fix: Routes are normalized to collapse dynamic IDs into placeholders.
"""
from litellm.proxy.auth.auth_utils import normalize_request_route
@@ -115,43 +151,53 @@ def test_route_normalization_for_responses_api():
("/v1/responses/resp_abc123", "/v1/responses/{response_id}"),
("/v1/responses/litellm_poll_xyz", "/v1/responses/{response_id}"),
]
for original, expected in responses_routes:
normalized = normalize_request_route(original)
assert normalized == expected, \
f"Failed: {original} -> {normalized} (expected {expected})"
assert (
normalized == expected
), f"Failed: {original} -> {normalized} (expected {expected})"
# Verify cardinality reduction
unique_normalized = set(normalize_request_route(route) for route, _ in responses_routes)
assert len(unique_normalized) == 1, \
f"Expected 1 unique normalized route, got {len(unique_normalized)}: {unique_normalized}"
print(f"✅ Responses API routes: {len(responses_routes)} different IDs normalized to 1 metric label")
unique_normalized = set(
normalize_request_route(route) for route, _ in responses_routes
)
assert (
len(unique_normalized) == 1
), f"Expected 1 unique normalized route, got {len(unique_normalized)}: {unique_normalized}"
print(
f"✅ Responses API routes: {len(responses_routes)} different IDs normalized to 1 metric label"
)
def test_route_normalization_for_sub_routes():
"""Test that sub-routes like /cancel and /input_items are normalized correctly"""
from litellm.proxy.auth.auth_utils import normalize_request_route
sub_routes = [
("/v1/responses/id1/cancel", "/v1/responses/{response_id}/cancel"),
("/v1/responses/id2/cancel", "/v1/responses/{response_id}/cancel"),
("/v1/responses/id3/input_items", "/v1/responses/{response_id}/input_items"),
("/openai/v1/responses/id4/input_items", "/openai/v1/responses/{response_id}/input_items"),
(
"/openai/v1/responses/id4/input_items",
"/openai/v1/responses/{response_id}/input_items",
),
]
for original, expected in sub_routes:
normalized = normalize_request_route(original)
assert normalized == expected, \
f"Failed: {original} -> {normalized} (expected {expected})"
assert (
normalized == expected
), f"Failed: {original} -> {normalized} (expected {expected})"
print("✅ Sub-routes normalized correctly")
def test_route_normalization_preserves_static_routes():
"""Test that static routes are not affected by normalization"""
from litellm.proxy.auth.auth_utils import normalize_request_route
static_routes = [
"/chat/completions",
"/v1/chat/completions",
@@ -161,46 +207,47 @@ def test_route_normalization_preserves_static_routes():
"/v1/models",
"/v1/responses", # List endpoint without ID
]
for route in static_routes:
normalized = normalize_request_route(route)
assert normalized == route, \
f"Static route should not be modified: {route} -> {normalized}"
assert (
normalized == route
), f"Static route should not be modified: {route} -> {normalized}"
print(f"✅ {len(static_routes)} static routes preserved")
def test_route_normalization_other_dynamic_apis():
"""Test normalization for other OpenAI-compatible APIs with dynamic IDs"""
from litellm.proxy.auth.auth_utils import normalize_request_route
test_cases = [
# Threads API
("/v1/threads/thread_123", "/v1/threads/{thread_id}"),
("/v1/threads/thread_abc/messages", "/v1/threads/{thread_id}/messages"),
("/v1/threads/thread_abc/runs/run_123", "/v1/threads/{thread_id}/runs/{run_id}"),
(
"/v1/threads/thread_abc/runs/run_123",
"/v1/threads/{thread_id}/runs/{run_id}",
),
# Vector Stores API
("/v1/vector_stores/vs_123", "/v1/vector_stores/{vector_store_id}"),
("/v1/vector_stores/vs_123/files", "/v1/vector_stores/{vector_store_id}/files"),
# Assistants API
("/v1/assistants/asst_123", "/v1/assistants/{assistant_id}"),
# Files API
("/v1/files/file_123", "/v1/files/{file_id}"),
("/v1/files/file_123/content", "/v1/files/{file_id}/content"),
# Batches API
("/v1/batches/batch_123", "/v1/batches/{batch_id}"),
("/v1/batches/batch_123/cancel", "/v1/batches/{batch_id}/cancel"),
]
for original, expected in test_cases:
normalized = normalize_request_route(original)
assert normalized == expected, \
f"Failed: {original} -> {normalized} (expected {expected})"
assert (
normalized == expected
), f"Failed: {original} -> {normalized} (expected {expected})"
print(f"✅ {len(test_cases)} other API routes normalized correctly")
@@ -219,25 +266,39 @@ def test_prometheus_metrics_use_normalized_routes():
# Create a mock PrometheusLogger
prometheus_logger = MagicMock()
prometheus_logger.get_labels_for_metric = PrometheusLogger.get_labels_for_metric.__get__(prometheus_logger)
prometheus_logger.get_labels_for_metric = (
PrometheusLogger.get_labels_for_metric.__get__(prometheus_logger)
)
# Test with a normalized route
enum_values = UserAPIKeyLabelValues(
route="/v1/responses/{response_id}", # Normalized route
status_code="200",
requested_model="gpt-4",
)
labels = prometheus_label_factory(
supported_enum_labels=prometheus_logger.get_labels_for_metric(
metric_name="litellm_proxy_total_requests_metric"
),
enum_values=enum_values,
)
# Verify the route is normalized in labels
assert labels["route"] == "/v1/responses/{response_id}", \
f"Expected normalized route in labels, got: {labels.get('route')}"
assert (
labels["route"] == "/v1/responses/{response_id}"
), f"Expected normalized route in labels, got: {labels.get('route')}"
print("✅ Prometheus metrics use normalized routes in labels")
if __name__ == "__main__":
test_user_email_in_required_metrics()
test_user_email_label_exists()
test_prometheus_metric_labels_structure()
test_route_normalization_for_responses_api()
test_route_normalization_for_sub_routes()
test_route_normalization_preserves_static_routes()
test_route_normalization_other_dynamic_apis()
test_prometheus_metrics_use_normalized_routes()
print("\n✅ All prometheus label tests passed!")
tests/test_litellm/integrations/test_prometheus_missing_metrics.py
+77
View File
@@ -0,0 +1,77 @@
"""
Unit tests for the new Prometheus metrics that were previously missing from validation.
Tests for:
- litellm_remaining_api_key_requests_for_model
- litellm_remaining_api_key_tokens_for_model
- litellm_callback_logging_failures_metric
"""
from typing import get_args
from litellm.types.integrations.prometheus import (
DEFINED_PROMETHEUS_METRICS,
PrometheusMetricLabels,
UserAPIKeyLabelNames,
)
def test_new_metrics_in_defined_metrics():
"""
Test that the new metrics are present in DEFINED_PROMETHEUS_METRICS.
"""
defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS)
new_metrics = [
"litellm_remaining_api_key_requests_for_model",
"litellm_remaining_api_key_tokens_for_model",
"litellm_callback_logging_failures_metric",
]
for metric in new_metrics:
assert (
metric in defined_metrics
), f"{metric} should be in DEFINED_PROMETHEUS_METRICS"
def test_new_metrics_have_correct_labels():
"""
Test that the new metrics have the correct labels defined.
"""
# Test API Key limits metrics labels
api_key_metrics = [
"litellm_remaining_api_key_requests_for_model",
"litellm_remaining_api_key_tokens_for_model",
]
expected_api_key_labels = [
UserAPIKeyLabelNames.API_KEY_HASH.value,
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value,
]
for metric in api_key_metrics:
labels = PrometheusMetricLabels.get_labels(metric)
for expected_label in expected_api_key_labels:
assert (
expected_label in labels
), f"{metric} should have label {expected_label}"
# Test Callback failure metric labels
callback_metric = "litellm_callback_logging_failures_metric"
callback_labels = PrometheusMetricLabels.get_labels(callback_metric)
assert (
UserAPIKeyLabelNames.CALLBACK_NAME.value in callback_labels
), f"{callback_metric} should have label {UserAPIKeyLabelNames.CALLBACK_NAME.value}"
def test_callback_name_label_definition():
"""
Test that CALLBACK_NAME is defined correctly in UserAPIKeyLabelNames.
"""
assert UserAPIKeyLabelNames.CALLBACK_NAME.value == "callback_name"
if __name__ == "__main__":
test_new_metrics_in_defined_metrics()
test_new_metrics_have_correct_labels()
test_callback_name_label_definition()
tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
+75 -8
View File
@@ -1138,6 +1138,73 @@ def test_bedrock_create_bedrock_block_different_document_formats():
assert block["document"]["name"].endswith(f"_{format_type}")
assert block["document"]["format"] == format_type
def test_bedrock_nova_web_search_options_mapping():
"""
Test that web_search_options is correctly mapped to Nova grounding.
This follows the LiteLLM pattern for web search where:
- Vertex AI maps web_search_options to {"googleSearch": {}}
- Anthropic maps web_search_options to {"type": "web_search_20250305", ...}
- Nova should map web_search_options to {"systemTool": {"name": "nova_grounding"}}
"""
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
config = AmazonConverseConfig()
# Test basic mapping for Nova model
result = config._map_web_search_options({}, "amazon.nova-pro-v1:0")
assert result is not None
system_tool = result.get("systemTool")
assert system_tool is not None
assert system_tool["name"] == "nova_grounding"
# Test with search_context_size (should be ignored for Nova)
result2 = config._map_web_search_options(
{"search_context_size": "high"},
"us.amazon.nova-premier-v1:0"
)
assert result2 is not None
system_tool2 = result2.get("systemTool")
assert system_tool2 is not None
assert system_tool2["name"] == "nova_grounding"
# Nova doesn't support search_context_size, so it's just ignored
def test_bedrock_tools_pt_does_not_handle_system_tool():
"""
Verify that _bedrock_tools_pt does NOT handle system_tool format.
System tools (nova_grounding) should be added via web_search_options,
not via the tools parameter directly.
"""
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_tools_pt
# Regular function tools should still work
tools = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"}
},
"required": ["location"]
}
}
}
]
result = _bedrock_tools_pt(tools=tools)
assert len(result) == 1
tool_spec = result[0].get("toolSpec")
assert tool_spec is not None
assert tool_spec["name"] == "get_weather"
def test_convert_to_anthropic_tool_result_image_with_cache_control():
"""
@@ -1305,12 +1372,12 @@ def test_convert_to_anthropic_tool_result_image_url_as_http():
assert result["content"][0]["cache_control"]["type"] == "ephemeral"
def test_anthropic_messages_pt_server_tool_use_passthrough():
"""
Test that anthropic_messages_pt passes through server_tool_use and
Test that anthropic_messages_pt passes through server_tool_use and
tool_search_tool_result blocks in assistant message content.
These are Anthropic-native content types used for tool search functionality
that need to be preserved when reconstructing multi-turn conversations.
Fixes: https://github.com/BerriAI/litellm/issues/XXXXX
"""
from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt
@@ -1359,15 +1426,15 @@ def test_anthropic_messages_pt_server_tool_use_passthrough():
# Verify we have 3 messages (user, assistant, user)
assert len(result) == 3
# Verify the assistant message content
assistant_msg = result[1]
assert assistant_msg["role"] == "assistant"
assert isinstance(assistant_msg["content"], list)
# Find the different content block types
content_types = [block.get("type") for block in assistant_msg["content"]]
# Verify server_tool_use block is preserved
assert "server_tool_use" in content_types
server_tool_use_block = next(
@@ -1376,7 +1443,7 @@ def test_anthropic_messages_pt_server_tool_use_passthrough():
assert server_tool_use_block["id"] == "srvtoolu_01ABC123"
assert server_tool_use_block["name"] == "tool_search_tool_regex"
assert server_tool_use_block["input"] == {"query": ".*time.*"}
# Verify tool_search_tool_result block is preserved
assert "tool_search_tool_result" in content_types
tool_result_block = next(
@@ -1385,7 +1452,7 @@ def test_anthropic_messages_pt_server_tool_use_passthrough():
assert tool_result_block["tool_use_id"] == "srvtoolu_01ABC123"
assert tool_result_block["content"]["type"] == "tool_search_tool_search_result"
assert tool_result_block["content"]["tool_references"][0]["tool_name"] == "get_time"
# Verify text block is also preserved
assert "text" in content_types
text_block = next(
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
+2
View File
@@ -787,11 +787,13 @@ def test_get_masked_values():
"presidio_ad_hoc_recognizers": None,
"aws_bedrock_runtime_endpoint": None,
"presidio_anonymizer_api_base": None,
"vertex_credentials": "{sensitive_api_key}",
}
masked_values = _get_masked_values(
sensitive_object, unmasked_length=4, number_of_asterisks=4
)
assert masked_values["presidio_anonymizer_api_base"] is None
assert masked_values["vertex_credentials"] == "{s****y}"
@pytest.mark.asyncio
tests/test_litellm/llms/vercel_ai_gateway/embedding/__init__.py
View File
tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py
+218
View File
@@ -0,0 +1,218 @@
import os
import sys
from unittest.mock import MagicMock, patch
import httpx
import pytest
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
from litellm.llms.vercel_ai_gateway.embedding.transformation import (
VercelAIGatewayEmbeddingConfig,
)
from litellm.llms.vercel_ai_gateway.common_utils import VercelAIGatewayException
from litellm.types.utils import EmbeddingResponse
def test_vercel_ai_gateway_embedding_get_complete_url():
"""Test URL generation for embeddings endpoint"""
config = VercelAIGatewayEmbeddingConfig()
# Test with default API base
url = config.get_complete_url(
api_base=None,
api_key=None,
model="openai/text-embedding-3-small",
optional_params={},
litellm_params={},
)
assert url == "https://ai-gateway.vercel.sh/v1/embeddings"
# Test with custom API base
url = config.get_complete_url(
api_base="https://custom.vercel.sh/v1",
api_key=None,
model="openai/text-embedding-3-small",
optional_params={},
litellm_params={},
)
assert url == "https://custom.vercel.sh/v1/embeddings"
# Test with trailing slash
url = config.get_complete_url(
api_base="https://custom.vercel.sh/v1/",
api_key=None,
model="openai/text-embedding-3-small",
optional_params={},
litellm_params={},
)
assert url == "https://custom.vercel.sh/v1/embeddings"
def test_vercel_ai_gateway_embedding_transform_request():
"""Test request transformation for embeddings"""
config = VercelAIGatewayEmbeddingConfig()
# Test with string input
request = config.transform_embedding_request(
model="openai/text-embedding-3-small",
input="Hello world",
optional_params={},
headers={},
)
assert request["model"] == "openai/text-embedding-3-small"
assert request["input"] == ["Hello world"]
# Test with list input
request = config.transform_embedding_request(
model="openai/text-embedding-3-small",
input=["Hello", "World"],
optional_params={},
headers={},
)
assert request["model"] == "openai/text-embedding-3-small"
assert request["input"] == ["Hello", "World"]
# Test stripping vercel_ai_gateway/ prefix
request = config.transform_embedding_request(
model="vercel_ai_gateway/openai/text-embedding-3-small",
input="Hello",
optional_params={},
headers={},
)
assert request["model"] == "openai/text-embedding-3-small"
def test_vercel_ai_gateway_embedding_transform_request_with_dimensions():
"""Test request transformation with dimensions parameter"""
config = VercelAIGatewayEmbeddingConfig()
request = config.transform_embedding_request(
model="openai/text-embedding-3-small",
input="Hello world",
optional_params={"dimensions": 768},
headers={},
)
assert request["model"] == "openai/text-embedding-3-small"
assert request["input"] == ["Hello world"]
assert request["dimensions"] == 768
def test_vercel_ai_gateway_embedding_validate_environment():
"""Test header validation and setup"""
config = VercelAIGatewayEmbeddingConfig()
headers = config.validate_environment(
headers={},
model="openai/text-embedding-3-small",
messages=[],
optional_params={},
litellm_params={},
api_key="test_key",
)
assert headers["Content-Type"] == "application/json"
assert headers["Authorization"] == "Bearer test_key"
# Test with existing headers (should merge)
headers = config.validate_environment(
headers={"X-Custom": "value"},
model="openai/text-embedding-3-small",
messages=[],
optional_params={},
litellm_params={},
api_key="test_key",
)
assert headers["X-Custom"] == "value"
assert headers["Authorization"] == "Bearer test_key"
def test_vercel_ai_gateway_embedding_get_supported_params():
"""Test supported OpenAI parameters"""
config = VercelAIGatewayEmbeddingConfig()
supported = config.get_supported_openai_params("openai/text-embedding-3-small")
assert "dimensions" in supported
assert "encoding_format" in supported
assert "timeout" in supported
assert "user" in supported
def test_vercel_ai_gateway_embedding_map_openai_params():
"""Test OpenAI parameter mapping"""
config = VercelAIGatewayEmbeddingConfig()
optional_params = config.map_openai_params(
non_default_params={"dimensions": 768, "encoding_format": "float"},
optional_params={},
model="openai/text-embedding-3-small",
drop_params=False,
)
assert optional_params["dimensions"] == 768
assert optional_params["encoding_format"] == "float"
def test_vercel_ai_gateway_embedding_error_class():
"""Test error class creation"""
config = VercelAIGatewayEmbeddingConfig()
error = config.get_error_class(
error_message="Test error",
status_code=400,
headers={"Content-Type": "application/json"},
)
assert isinstance(error, VercelAIGatewayException)
assert error.message == "Test error"
assert error.status_code == 400
def test_vercel_ai_gateway_embedding_transform_response():
"""Test response transformation"""
config = VercelAIGatewayEmbeddingConfig()
mock_response = MagicMock(spec=httpx.Response)
mock_response.text = '{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2,0.3]}],"model":"openai/text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}'
mock_response.json.return_value = {
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}],
"model": "openai/text-embedding-3-small",
"usage": {"prompt_tokens": 2, "total_tokens": 2},
}
mock_logging = MagicMock()
response = config.transform_embedding_response(
model="openai/text-embedding-3-small",
raw_response=mock_response,
model_response=EmbeddingResponse(),
logging_obj=mock_logging,
api_key="test_key",
request_data={},
optional_params={},
litellm_params={},
)
assert response is not None
mock_logging.post_call.assert_called_once()
def test_vercel_ai_gateway_embedding_env_vars():
"""Test environment variable handling"""
config = VercelAIGatewayEmbeddingConfig()
with patch.dict(
os.environ,
{
"VERCEL_AI_GATEWAY_API_BASE": "https://env.vercel.sh/v1",
},
):
url = config.get_complete_url(
api_base=None,
api_key=None,
model="openai/text-embedding-3-small",
optional_params={},
litellm_params={},
)
assert url == "https://env.vercel.sh/v1/embeddings"
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
+1 -1
View File
@@ -1885,7 +1885,7 @@ class TestMCPServerManager:
# Create mock client that tracks call_tool usage
mock_client = AsyncMock()
async def mock_call_tool(params):
async def mock_call_tool(params, host_progress_callback=None):
# Return a mock CallToolResult
result = MagicMock(spec=CallToolResult)
result.content = [{"type": "text", "text": "Tool executed successfully"}]
tests/test_litellm/proxy/auth/test_auth_utils.py
+23
View File
@@ -8,6 +8,7 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import (
_get_customer_id_from_standard_headers,
get_end_user_id_from_request_body,
get_model_from_request,
get_key_model_rpm_limit,
get_key_model_tpm_limit,
)
@@ -186,3 +187,25 @@ class TestGetEndUserIdFromRequestBodyWithStandardHeaders:
request_body=request_body, request_headers=headers
)
assert result == "body-user"
def test_get_model_from_request_supports_google_model_names_with_slashes():
assert (
get_model_from_request(
request_data={},
route="/v1beta/models/bedrock/claude-sonnet-3.7:generateContent",
)
== "bedrock/claude-sonnet-3.7"
)
assert (
get_model_from_request(
request_data={},
route="/models/hosted_vllm/gpt-oss-20b:generateContent",
)
== "hosted_vllm/gpt-oss-20b"
)
def test_get_model_from_request_vertex_passthrough_still_works():
route = "/vertex_ai/v1/projects/p/locations/l/publishers/google/models/gemini-1.5-pro:generateContent"
assert get_model_from_request(request_data={}, route=route) == "gemini-1.5-pro"
tests/test_litellm/proxy/auth/test_login_utils.py
+130
View File
@@ -626,3 +626,133 @@ async def test_expire_previous_ui_session_tokens_exception_handling():
# Should not raise exception despite database error
await expire_previous_ui_session_tokens(user_id, mock_prisma_client)
@pytest.mark.asyncio
async def test_authenticate_user_admin_login_with_non_ascii_characters():
"""Test admin login with non-ASCII characters in password (issue #19559)"""
master_key = "sk-1234"
ui_username = "admin£test"
ui_password = "sk-1234£pass"
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
with patch.dict(
os.environ,
{
"UI_USERNAME": ui_username,
"UI_PASSWORD": ui_password,
"DATABASE_URL": "postgresql://test:test@localhost/test",
},
):
with patch(
"litellm.proxy.auth.login_utils.generate_key_helper_fn",
new_callable=AsyncMock,
) as mock_generate_key:
mock_generate_key.return_value = {
"token": "test-token-123",
"user_id": LITELLM_PROXY_ADMIN_NAME,
}
with patch(
"litellm.proxy.auth.login_utils.user_update",
new_callable=AsyncMock,
return_value=None,
) as mock_user_update:
with patch(
"litellm.proxy.auth.login_utils.get_secret_bool",
return_value=False,
):
result = await authenticate_user(
username=ui_username,
password=ui_password,
master_key=master_key,
prisma_client=mock_prisma_client,
)
assert isinstance(result, LoginResult)
assert result.user_id == LITELLM_PROXY_ADMIN_NAME
assert result.key == "test-token-123"
assert result.user_role == LitellmUserRoles.PROXY_ADMIN
def test_authenticate_user_non_ascii_direct_comparison():
"""Test that non-ASCII characters can be compared directly (unit test for fix)"""
import secrets
# This test verifies the fix handles non-ASCII by encoding to bytes
username = "admin£test"
password = "pass£word"
# This would fail without encoding:
# secrets.compare_digest(username, username) # TypeError!
# But works with the fix:
result = secrets.compare_digest(
username.encode("utf-8"), username.encode("utf-8")
)
assert result is True
# And correctly returns False for different passwords
result = secrets.compare_digest(
password.encode("utf-8"), "different£pass".encode("utf-8")
)
assert result is False
@pytest.mark.asyncio
async def test_authenticate_user_database_login_with_non_ascii_password():
"""Test database user login with non-ASCII characters in password (issue #19559)"""
master_key = "sk-1234"
user_email = "test@example.com"
password_with_special_char = "correct£password"
hashed_password = hash_token(token=password_with_special_char)
mock_user = MagicMock()
mock_user.user_id = "test-user-123"
mock_user.user_email = user_email
mock_user.password = hashed_password
mock_user.user_role = LitellmUserRoles.INTERNAL_USER
def mock_find_first(**kwargs):
where = kwargs.get("where", {})
user_email_filter = where.get("user_email", {})
if str(user_email_filter.get("equals", "")).lower() == user_email.lower():
return mock_user
return None
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(
side_effect=mock_find_first
)
with patch.dict(
os.environ,
{
"DATABASE_URL": "postgresql://test:test@localhost/test",
"UI_USERNAME": "admin",
"UI_PASSWORD": "admin-password",
},
):
with patch(
"litellm.proxy.auth.login_utils.expire_previous_ui_session_tokens",
new_callable=AsyncMock,
return_value=None,
):
with patch(
"litellm.proxy.auth.login_utils.generate_key_helper_fn",
new_callable=AsyncMock,
) as mock_generate_key:
mock_generate_key.return_value = {"token": "token-123"}
result = await authenticate_user(
username=user_email,
password=password_with_special_char,
master_key=master_key,
prisma_client=mock_prisma_client,
)
assert isinstance(result, LoginResult)
assert result.user_id == "test-user-123"
assert result.user_email == user_email
tests/test_litellm/proxy/auth/test_route_checks.py
+4
View File
@@ -161,9 +161,11 @@ def test_virtual_key_llm_api_route_includes_passthrough_prefix(route):
[
"/v1beta/models/gemini-2.5-flash:countTokens",
"/v1beta/models/gemini-2.0-flash:generateContent",
"/v1beta/models/bedrock/claude-sonnet-3.7:generateContent",
"/v1beta/models/gemini-1.5-pro:streamGenerateContent",
"/models/gemini-2.5-flash:countTokens",
"/models/gemini-2.0-flash:generateContent",
"/models/bedrock/claude-sonnet-3.7:generateContent",
"/models/gemini-1.5-pro:streamGenerateContent",
],
)
@@ -187,9 +189,11 @@ def test_virtual_key_llm_api_routes_allows_google_routes(route):
"/v1beta/models/google-gemini-2-5-pro-code-reviewer-k8s:generateContent",
"/v1beta/models/gemini-2.5-flash-exp:countTokens",
"/v1beta/models/custom-model-name-123:streamGenerateContent",
"/v1beta/models/bedrock/claude-sonnet-3.7:generateContent",
"/models/google-gemini-2-5-pro-code-reviewer-k8s:generateContent",
"/models/gemini-2.5-flash-exp:countTokens",
"/models/custom-model-name-123:streamGenerateContent",
"/models/bedrock/claude-sonnet-3.7:generateContent",
],
)
def test_google_routes_with_dynamic_model_names_recognized_as_llm_api_route(route):
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py
+56
View File
@@ -549,6 +549,62 @@ class TestAdditionalParams:
)
class TestModelParameter:
"""Test model parameter handling in guardrail requests"""
@pytest.mark.asyncio
async def test_model_passed_from_inputs(
self, generic_guardrail, mock_request_data_input
):
"""Test that model is passed to the API when provided in inputs"""
mock_response = MagicMock()
mock_response.json.return_value = {
"action": "NONE",
"texts": ["test"],
}
mock_response.raise_for_status = MagicMock()
with patch.object(
generic_guardrail.async_handler, "post", return_value=mock_response
) as mock_post:
await generic_guardrail.apply_guardrail(
inputs={"texts": ["test"], "model": "gpt-4"},
request_data=mock_request_data_input,
input_type="request",
)
# Verify API was called with model
call_args = mock_post.call_args
json_payload = call_args.kwargs["json"]
assert json_payload["model"] == "gpt-4"
@pytest.mark.asyncio
async def test_model_none_when_not_provided(
self, generic_guardrail, mock_request_data_input
):
"""Test that model is None when not provided in inputs"""
mock_response = MagicMock()
mock_response.json.return_value = {
"action": "NONE",
"texts": ["test"],
}
mock_response.raise_for_status = MagicMock()
with patch.object(
generic_guardrail.async_handler, "post", return_value=mock_response
) as mock_post:
await generic_guardrail.apply_guardrail(
inputs={"texts": ["test"]}, # No model in inputs
request_data=mock_request_data_input,
input_type="request",
)
# Verify API was called with model=None
call_args = mock_post.call_args
json_payload = call_args.kwargs["json"]
assert json_payload["model"] is None
class TestErrorHandling:
"""Test error handling scenarios"""
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py
+300 -2
View File
@@ -3,6 +3,7 @@ import sys
import uuid
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from fastapi import HTTPException
from httpx import Request, Response
@@ -47,20 +48,129 @@ def test_onyx_guard_config():
del os.environ["ONYX_API_KEY"]
def test_onyx_guard_with_custom_timeout_from_kwargs():
"""Test Onyx guard instantiation with custom timeout passed via kwargs."""
# Set environment variables for testing
os.environ["ONYX_API_BASE"] = "https://test.onyx.security"
os.environ["ONYX_API_KEY"] = "test-api-key"
with patch(
"litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client"
) as mock_get_client:
mock_get_client.return_value = MagicMock()
# Simulate how guardrail is instantiated from config with timeout
guardrail = OnyxGuardrail(
guardrail_name="onyx-guard-custom-timeout",
event_hook="pre_call",
default_on=True,
timeout=45.0,
)
# Verify the client was initialized with custom timeout
mock_get_client.assert_called()
call_kwargs = mock_get_client.call_args.kwargs
timeout_param = call_kwargs["params"]["timeout"]
assert timeout_param.read == 45.0
assert timeout_param.connect == 5.0
# Clean up
if "ONYX_API_BASE" in os.environ:
del os.environ["ONYX_API_BASE"]
if "ONYX_API_KEY" in os.environ:
del os.environ["ONYX_API_KEY"]
def test_onyx_guard_with_timeout_none_uses_env_var():
"""Test Onyx guard with timeout=None uses ONYX_TIMEOUT env var.
When timeout=None is passed (as it would be from config model with default None),
the ONYX_TIMEOUT environment variable should be used.
"""
# Set environment variables for testing
os.environ["ONYX_API_BASE"] = "https://test.onyx.security"
os.environ["ONYX_API_KEY"] = "test-api-key"
os.environ["ONYX_TIMEOUT"] = "60"
with patch(
"litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client"
) as mock_get_client:
mock_get_client.return_value = MagicMock()
# Pass timeout=None to simulate config model behavior
guardrail = OnyxGuardrail(
guardrail_name="onyx-guard-env-timeout",
event_hook="pre_call",
default_on=True,
timeout=None, # This triggers env var lookup
)
# Verify the client was initialized with timeout from env var
mock_get_client.assert_called()
call_kwargs = mock_get_client.call_args.kwargs
timeout_param = call_kwargs["params"]["timeout"]
assert timeout_param.read == 60.0
assert timeout_param.connect == 5.0
# Clean up
if "ONYX_API_BASE" in os.environ:
del os.environ["ONYX_API_BASE"]
if "ONYX_API_KEY" in os.environ:
del os.environ["ONYX_API_KEY"]
if "ONYX_TIMEOUT" in os.environ:
del os.environ["ONYX_TIMEOUT"]
def test_onyx_guard_with_timeout_none_defaults_to_10():
"""Test Onyx guard with timeout=None and no env var defaults to 10 seconds."""
# Set environment variables for testing
os.environ["ONYX_API_BASE"] = "https://test.onyx.security"
os.environ["ONYX_API_KEY"] = "test-api-key"
# Ensure ONYX_TIMEOUT is not set
if "ONYX_TIMEOUT" in os.environ:
del os.environ["ONYX_TIMEOUT"]
with patch(
"litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client"
) as mock_get_client:
mock_get_client.return_value = MagicMock()
# Pass timeout=None with no env var - should default to 10.0
guardrail = OnyxGuardrail(
guardrail_name="onyx-guard-default-timeout",
event_hook="pre_call",
default_on=True,
timeout=None,
)
# Verify the client was initialized with default timeout of 10.0
mock_get_client.assert_called()
call_kwargs = mock_get_client.call_args.kwargs
timeout_param = call_kwargs["params"]["timeout"]
assert timeout_param.read == 10.0
assert timeout_param.connect == 5.0
# Clean up
if "ONYX_API_BASE" in os.environ:
del os.environ["ONYX_API_BASE"]
if "ONYX_API_KEY" in os.environ:
del os.environ["ONYX_API_KEY"]
class TestOnyxGuardrail:
"""Test suite for Onyx Security Guardrail integration."""
def setup_method(self):
"""Setup test environment."""
# Clean up any existing environment variables
for key in ["ONYX_API_BASE", "ONYX_API_KEY"]:
for key in ["ONYX_API_BASE", "ONYX_API_KEY", "ONYX_TIMEOUT"]:
if key in os.environ:
del os.environ[key]
def teardown_method(self):
"""Clean up test environment."""
# Clean up any environment variables set during tests
for key in ["ONYX_API_BASE", "ONYX_API_KEY"]:
for key in ["ONYX_API_BASE", "ONYX_API_KEY", "ONYX_TIMEOUT"]:
if key in os.environ:
del os.environ[key]
@@ -103,6 +213,95 @@ class TestOnyxGuardrail:
):
OnyxGuardrail(guardrail_name="test-guard", event_hook="pre_call")
def test_initialization_with_default_timeout(self):
"""Test that default timeout is 10.0 seconds."""
os.environ["ONYX_API_KEY"] = "test-api-key"
with patch(
"litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client"
) as mock_get_client:
mock_get_client.return_value = MagicMock()
guardrail = OnyxGuardrail(
guardrail_name="test-guard", event_hook="pre_call", default_on=True
)
# Verify the client was initialized with correct timeout
mock_get_client.assert_called_once()
call_kwargs = mock_get_client.call_args.kwargs
timeout_param = call_kwargs["params"]["timeout"]
assert timeout_param.read == 10.0
assert timeout_param.connect == 5.0
def test_initialization_with_custom_timeout_parameter(self):
"""Test initialization with custom timeout parameter."""
os.environ["ONYX_API_KEY"] = "test-api-key"
with patch(
"litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client"
) as mock_get_client:
mock_get_client.return_value = MagicMock()
guardrail = OnyxGuardrail(
guardrail_name="test-guard",
event_hook="pre_call",
default_on=True,
timeout=30.0,
)
# Verify the client was initialized with custom timeout
mock_get_client.assert_called_once()
call_kwargs = mock_get_client.call_args.kwargs
timeout_param = call_kwargs["params"]["timeout"]
assert timeout_param.read == 30.0
assert timeout_param.connect == 5.0
def test_initialization_with_timeout_from_env_var(self):
"""Test initialization with timeout from ONYX_TIMEOUT environment variable.
Note: The env var is only used when timeout=None is explicitly passed,
since the default parameter value is 10.0 (not None).
"""
os.environ["ONYX_API_KEY"] = "test-api-key"
os.environ["ONYX_TIMEOUT"] = "25"
with patch(
"litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client"
) as mock_get_client:
mock_get_client.return_value = MagicMock()
# Must pass timeout=None explicitly to trigger env var lookup
guardrail = OnyxGuardrail(
guardrail_name="test-guard", event_hook="pre_call", default_on=True, timeout=None
)
# Verify the client was initialized with timeout from env var
mock_get_client.assert_called_once()
call_kwargs = mock_get_client.call_args.kwargs
timeout_param = call_kwargs["params"]["timeout"]
assert timeout_param.read == 25.0
assert timeout_param.connect == 5.0
def test_initialization_timeout_parameter_overrides_env_var(self):
"""Test that timeout parameter overrides ONYX_TIMEOUT environment variable."""
os.environ["ONYX_API_KEY"] = "test-api-key"
os.environ["ONYX_TIMEOUT"] = "25"
with patch(
"litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client"
) as mock_get_client:
mock_get_client.return_value = MagicMock()
guardrail = OnyxGuardrail(
guardrail_name="test-guard",
event_hook="pre_call",
default_on=True,
timeout=15.0,
)
# Verify the client was initialized with parameter timeout (not env var)
mock_get_client.assert_called_once()
call_kwargs = mock_get_client.call_args.kwargs
timeout_param = call_kwargs["params"]["timeout"]
assert timeout_param.read == 15.0
assert timeout_param.connect == 5.0
@pytest.mark.asyncio
async def test_apply_guardrail_request_no_violations(self):
"""Test apply_guardrail for request with no violations detected."""
@@ -388,6 +587,105 @@ class TestOnyxGuardrail:
assert result == inputs
@pytest.mark.asyncio
async def test_apply_guardrail_timeout_error_handling(self):
"""Test handling of timeout errors in apply_guardrail (graceful degradation)."""
# Set required API key
os.environ["ONYX_API_KEY"] = "test-api-key"
guardrail = OnyxGuardrail(
guardrail_name="test-guard", event_hook="pre_call", default_on=True, timeout=1.0
)
inputs = GenericGuardrailAPIInputs()
request_data = {
"proxy_server_request": {
"messages": [{"role": "user", "content": "Test message"}],
"model": "gpt-3.5-turbo",
}
}
# Test httpx timeout error
with patch.object(
guardrail.async_handler, "post", side_effect=httpx.TimeoutException("Request timed out")
):
# Should return original inputs on timeout (graceful degradation)
result = await guardrail.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="request",
logging_obj=None,
)
assert result == inputs
@pytest.mark.asyncio
async def test_apply_guardrail_read_timeout_error_handling(self):
"""Test handling of read timeout errors in apply_guardrail."""
# Set required API key
os.environ["ONYX_API_KEY"] = "test-api-key"
guardrail = OnyxGuardrail(
guardrail_name="test-guard", event_hook="pre_call", default_on=True, timeout=5.0
)
inputs = GenericGuardrailAPIInputs()
request_data = {
"proxy_server_request": {
"messages": [{"role": "user", "content": "Test message"}],
"model": "gpt-3.5-turbo",
}
}
# Test httpx ReadTimeout error
with patch.object(
guardrail.async_handler, "post", side_effect=httpx.ReadTimeout("Read timed out")
):
# Should return original inputs on timeout (graceful degradation)
result = await guardrail.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="request",
logging_obj=None,
)
assert result == inputs
@pytest.mark.asyncio
async def test_apply_guardrail_connect_timeout_error_handling(self):
"""Test handling of connect timeout errors in apply_guardrail."""
# Set required API key
os.environ["ONYX_API_KEY"] = "test-api-key"
guardrail = OnyxGuardrail(
guardrail_name="test-guard", event_hook="pre_call", default_on=True, timeout=5.0
)
inputs = GenericGuardrailAPIInputs()
request_data = {
"proxy_server_request": {
"messages": [{"role": "user", "content": "Test message"}],
"model": "gpt-3.5-turbo",
}
}
# Test httpx ConnectTimeout error
with patch.object(
guardrail.async_handler, "post", side_effect=httpx.ConnectTimeout("Connect timed out")
):
# Should return original inputs on timeout (graceful degradation)
result = await guardrail.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="request",
logging_obj=None,
)
assert result == inputs
@pytest.mark.asyncio
async def test_apply_guardrail_no_logging_obj(self):
"""Test apply_guardrail without logging object (uses UUID)."""
tests/test_litellm/proxy/test_chat_completion_metadata.py
+154
View File
@@ -0,0 +1,154 @@
import pytest
from unittest.mock import MagicMock, AsyncMock, patch
from litellm.proxy.proxy_server import chat_completion, completion, embeddings
from litellm.proxy._types import UserAPIKeyAuth
from fastapi import Request, Response
@pytest.mark.asyncio
async def test_chat_completion_metadata_population():
# Setup
request = MagicMock(spec=Request)
# Mock _read_request_body to return a dict
with patch(
"litellm.proxy.proxy_server._read_request_body", new_callable=AsyncMock
) as mock_read_body:
mock_read_body.return_value = {"model": "gpt-3.5-turbo", "messages": []}
user_api_key_dict = UserAPIKeyAuth(
user_id="test_user_id", team_id="test_team_id", org_id="test_org_id"
)
fastapi_response = MagicMock(spec=Response)
# Mock ProxyBaseLLMRequestProcessing
with patch(
"litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing"
) as MockProcessor:
mock_instance = MockProcessor.return_value
mock_instance.base_process_llm_request = AsyncMock(
return_value={"choices": []}
)
# Execute
await chat_completion(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
)
# Verify
# Check if ProxyBaseLLMRequestProcessing was initialized with data containing metadata
call_args = MockProcessor.call_args
assert call_args is not None
data_arg = call_args.kwargs.get("data")
assert data_arg is not None
assert "metadata" in data_arg
assert data_arg["metadata"]["user_api_key_user_id"] == "test_user_id"
assert data_arg["metadata"]["user_api_key_team_id"] == "test_team_id"
assert data_arg["metadata"]["user_api_key_org_id"] == "test_org_id"
@pytest.mark.asyncio
async def test_embedding_metadata_population():
"""
Test that the embedding endpoint correctly populates metadata
from UserAPIKeyAuth.
"""
# Setup
with patch(
"litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing.base_process_llm_request"
):
with patch(
"litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing.__init__",
return_value=None,
) as mock_base_process_init:
# Create a mock UserAPIKeyAuth object
mock_user_auth = MagicMock(spec=UserAPIKeyAuth)
mock_user_auth.user_id = "test_user_id_emb"
mock_user_auth.team_id = "test_team_id_emb"
mock_user_auth.org_id = "test_org_id_emb"
# Create a mock Request object
mock_request = MagicMock(spec=Request)
mock_request.json = AsyncMock(
return_value={"model": "gpt-3.5-turbo", "input": "hello"}
)
# Mock _read_request_body to return our data
with patch(
"litellm.proxy.proxy_server._read_request_body",
new=AsyncMock(
return_value={"model": "gpt-3.5-turbo", "input": "hello"}
),
):
# Call the endpoint function directly
await embeddings(
request=mock_request,
fastapi_response=MagicMock(spec=Response),
user_api_key_dict=mock_user_auth,
)
# Check if ProxyBaseLLMRequestProcessing was initialized with the correct metadata
mock_base_process_init.assert_called_once()
call_args = mock_base_process_init.call_args
# handle both positional and keyword args for data
if "data" in call_args.kwargs:
data_arg = call_args.kwargs["data"]
else:
data_arg = call_args.args[0]
assert (
data_arg["metadata"]["user_api_key_user_id"] == "test_user_id_emb"
)
assert (
data_arg["metadata"]["user_api_key_team_id"] == "test_team_id_emb"
)
assert data_arg["metadata"]["user_api_key_org_id"] == "test_org_id_emb"
@pytest.mark.asyncio
async def test_completion_metadata_population():
# Setup
request = MagicMock(spec=Request)
# Mock _read_request_body to return a dict
with patch(
"litellm.proxy.proxy_server._read_request_body", new_callable=AsyncMock
) as mock_read_body:
mock_read_body.return_value = {
"model": "gpt-3.5-turbo-instruct",
"prompt": "test",
}
user_api_key_dict = UserAPIKeyAuth(
user_id="test_user_id_2", team_id="test_team_id_2", org_id="test_org_id_2"
)
fastapi_response = MagicMock(spec=Response)
# Mock ProxyBaseLLMRequestProcessing
with patch(
"litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing"
) as MockProcessor:
mock_instance = MockProcessor.return_value
mock_instance.base_process_llm_request = AsyncMock(
return_value={"choices": []}
)
# Execute
await completion(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
)
# Verify
call_args = MockProcessor.call_args
assert call_args is not None
data_arg = call_args.kwargs.get("data")
assert data_arg is not None
assert "metadata" in data_arg
assert data_arg["metadata"]["user_api_key_user_id"] == "test_user_id_2"
assert data_arg["metadata"]["user_api_key_team_id"] == "test_team_id_2"
assert data_arg["metadata"]["user_api_key_org_id"] == "test_org_id_2"
tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py
+110
View File
@@ -1227,3 +1227,113 @@ class TestProxySettingEndpoints:
assert retrieved_role_mappings["provider"] == "google"
assert retrieved_role_mappings["group_claim"] == "groups"
assert retrieved_role_mappings["default_role"] == LitellmUserRoles.INTERNAL_USER
def test_setup_role_mappings_custom_logic_with_env_vars(self, monkeypatch):
"""Test the _setup_role_mappings function directly with custom role mapping logic from environment variables"""
import asyncio
import os
from litellm.proxy.management_endpoints.ui_sso import _setup_role_mappings
from litellm.proxy._types import LitellmUserRoles
# Set up environment variables for custom role mappings using valid Python dict format
monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_ROLES", "{'proxy_admin': ['custom-admin-group'], 'internal_user': ['custom-user-group'], 'proxy_admin_viewer': ['custom-viewer-group']}")
monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", "custom-groups")
monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", "internal_user_viewer")
# Debug: Print environment variables
print("GENERIC_ROLE_MAPPINGS_ROLES:", os.getenv("GENERIC_ROLE_MAPPINGS_ROLES"))
print("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM:", os.getenv("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM"))
print("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE:", os.getenv("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE"))
# Run the async function
role_mappings = asyncio.run(_setup_role_mappings())
# Debug: Print result
print("role_mappings result:", role_mappings)
# Verify role_mappings is returned correctly from environment variables
assert role_mappings is not None
assert role_mappings.provider == "generic"
assert role_mappings.group_claim == "custom-groups"
assert role_mappings.default_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
assert role_mappings.roles[LitellmUserRoles.PROXY_ADMIN] == ["custom-admin-group"]
assert role_mappings.roles[LitellmUserRoles.INTERNAL_USER] == ["custom-user-group"]
assert role_mappings.roles[LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY] == ["custom-viewer-group"]
def test_setup_role_mappings_custom_logic_with_no_config(self, monkeypatch):
"""Test the _setup_role_mappings function returns None when no configuration is available"""
import asyncio
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy.management_endpoints.ui_sso import _setup_role_mappings
# Ensure environment variables are not set
monkeypatch.delenv("GENERIC_ROLE_MAPPINGS_ROLES", raising=False)
monkeypatch.delenv("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", raising=False)
monkeypatch.delenv("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", raising=False)
# Mock the prisma client to return None (no database record)
mock_prisma = MagicMock()
mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None)
# Run the async function
role_mappings = asyncio.run(_setup_role_mappings())
# Should return None when no configuration is available
assert role_mappings is None
def test_get_sso_settings_with_env_role_mappings(self, mock_proxy_config, mock_auth, monkeypatch):
import json
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy._types import LitellmUserRoles
monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_ROLES", '{"proxy_admin": ["custom-admin-group"], "internal_user": ["custom-user-group"], "proxy_admin_viewer": ["custom-viewer-group"]}')
monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", "custom-groups")
monkeypatch.setenv("GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", "internal_user_viewer")
mock_prisma = MagicMock()
mock_db_record = MagicMock()
mock_db_record.sso_settings = {
"google_client_id": "test_google_client_id",
"role_mappings": {
"provider": "google",
"group_claim": "db-groups",
"default_role": "proxy_admin",
"roles": {
"proxy_admin": ["db-admin-group"],
},
},
}
mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_db_record)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
from litellm.proxy.proxy_server import proxy_config
monkeypatch.setattr(
proxy_config, "_decrypt_and_set_db_env_variables", lambda environment_variables: environment_variables
)
response = client.get("/get/sso_settings")
assert response.status_code == 200
data = response.json()
values = data["values"]
assert "role_mappings" in values
assert values["role_mappings"] is not None
# The database values shoeld override the environment variables
assert values["role_mappings"]["provider"] == "google"
assert values["role_mappings"]["group_claim"] == "db-groups"
assert values["role_mappings"]["default_role"] == LitellmUserRoles.PROXY_ADMIN
assert values["role_mappings"]["roles"][LitellmUserRoles.PROXY_ADMIN] == ["db-admin-group"]
# Verify that the database was checked but environment variables took priority
mock_prisma.db.litellm_ssoconfig.find_unique.assert_called_once_with(
where={"id": "sso_config"}
)
# Verify other SSO settings are still correctly returned
assert values["google_client_id"] == "test_google_client_id"
# Verify field_schema is still present
assert "field_schema" in data
assert "properties" in data["field_schema"]
assert "role_mappings" in data["field_schema"]["properties"]
tests/test_litellm/secret_managers/test_secret_managers_main.py
+16 -9
View File
@@ -48,7 +48,8 @@ def mock_env():
@patch("litellm.secret_managers.main.oidc_cache")
@patch("litellm.secret_managers.main._get_oidc_http_handler")
def test_oidc_google_success(mock_get_http_handler, mock_oidc_cache):
@patch("httpx.Client") # Prevent any real HTTP connections
def test_oidc_google_success(mock_httpx_client, mock_get_http_handler, mock_oidc_cache):
mock_oidc_cache.get_cache.return_value = None
mock_handler = MockHTTPHandler(timeout=600.0)
mock_get_http_handler.return_value = mock_handler
@@ -144,7 +145,7 @@ def test_oidc_azure_file_success(mock_env, tmp_path):
mock_env["AZURE_FEDERATED_TOKEN_FILE"] = str(token_file)
secret_name = "oidc/azure/azure-audience"
result = get_secret(secret_name)
result = get_secret(secret_name)
assert result == "azure_token"
@@ -156,16 +157,22 @@ def test_oidc_azure_ad_token_success(mock_get_azure_ad_token_provider):
if "AZURE_FEDERATED_TOKEN_FILE" in os.environ:
del os.environ["AZURE_FEDERATED_TOKEN_FILE"]
# Mock the token provider function that gets returned and called
mock_token_provider = Mock(return_value="azure_ad_token")
mock_get_azure_ad_token_provider.return_value = mock_token_provider
secret_name = "oidc/azure/api://azure-audience"
result = get_secret(secret_name)
# Also mock the Azure Identity SDK to prevent any real Azure calls
with patch("azure.identity.get_bearer_token_provider") as mock_bearer:
mock_bearer.return_value = mock_token_provider
secret_name = "oidc/azure/api://azure-audience"
result = get_secret(secret_name)
assert result == "azure_ad_token"
mock_get_azure_ad_token_provider.assert_called_once_with(
azure_scope="api://azure-audience"
)
mock_token_provider.assert_called_once_with()
assert result == "azure_ad_token"
mock_get_azure_ad_token_provider.assert_called_once_with(
azure_scope="api://azure-audience"
)
mock_token_provider.assert_called_once_with()
def test_oidc_file_success(tmp_path):
tests/test_litellm/test_gpt_image_cost_calculator.py
+68
View File
@@ -19,10 +19,13 @@ import pytest
import litellm
from litellm.types.utils import (
CompletionTokensDetailsWrapper,
ImageResponse,
ImageObject,
ImageUsage,
ImageUsageInputTokensDetails,
PromptTokensDetailsWrapper,
Usage,
)
@@ -202,6 +205,71 @@ class TestGPTImageCostRouting:
assert cost >= 0
class TestGPTImage15OutputImageTokens:
"""
Test for GitHub issue #19508:
Image usage calculation does not include image tokens in gpt-image-1.5
gpt-image-1.5 returns output_tokens_details with separate image_tokens and text_tokens,
and these must be correctly included in cost calculation.
"""
def test_gpt_image_15_output_image_tokens_cost(self):
"""
Test that output image tokens are correctly included in cost calculation.
This tests the fix for issue #19508 where output_tokens_details.image_tokens
were not being included in the cost calculation, causing costs to be
underreported (e.g., $0.046 instead of $0.14).
"""
# Simulate gpt-image-1.5 response with output_tokens_details
# This is what the API returns and what convert_to_image_response transforms
usage = Usage(
prompt_tokens=169,
completion_tokens=4599,
total_tokens=4768,
prompt_tokens_details=PromptTokensDetailsWrapper(
text_tokens=169,
image_tokens=0,
),
completion_tokens_details=CompletionTokensDetailsWrapper(
text_tokens=439,
image_tokens=4160,
),
)
image_response = ImageResponse(
created=1234567890,
data=[ImageObject(b64_json="test")],
)
image_response.usage = usage
image_response._hidden_params = {"custom_llm_provider": "openai"}
cost = litellm.completion_cost(
completion_response=image_response,
model="gpt-image-1.5",
call_type="image_generation",
custom_llm_provider="openai",
)
# gpt-image-1.5 pricing:
# - input_cost_per_token: 5e-06 ($5/1M for text input)
# - output_cost_per_token: 1e-05 ($10/1M for text output)
# - output_cost_per_image_token: 3.2e-05 ($32/1M for image output)
#
# Expected cost:
# Input text: 169 * $5/1M = $0.000845
# Output text: 439 * $10/1M = $0.00439
# Output image: 4160 * $32/1M = $0.13312
# Total: $0.138355
expected_cost = 169 * 5e-06 + 439 * 1e-05 + 4160 * 3.2e-05
assert abs(cost - expected_cost) < 1e-6, (
f"Expected {expected_cost}, got {cost}. "
f"Image tokens may not be included in cost calculation."
)
class TestCompletionCostIntegration:
"""Test the full completion_cost integration for gpt-image-1"""
tests/test_litellm/test_router_per_deployment_num_retries.py
+45 -12
View File
@@ -32,17 +32,17 @@ class TestPerDeploymentNumRetries:
)
deployment = router.model_list[0]
# Create a mock exception without num_retries
class MockException(Exception):
pass
exc = MockException("test error")
assert not hasattr(exc, "num_retries") or exc.num_retries is None
# Call the helper
router._set_deployment_num_retries_on_exception(exc, deployment)
# Verify num_retries was set from deployment
assert exc.num_retries == 5
@@ -66,16 +66,16 @@ class TestPerDeploymentNumRetries:
)
deployment = router.model_list[0]
# Create an exception that already has num_retries
class MockException(Exception):
num_retries = 10 # Already set
exc = MockException("test error")
# Call the helper
router._set_deployment_num_retries_on_exception(exc, deployment)
# Verify num_retries was NOT overridden
assert exc.num_retries == 10
@@ -99,15 +99,15 @@ class TestPerDeploymentNumRetries:
)
deployment = router.model_list[0]
class MockException(Exception):
pass
exc = MockException("test error")
# Call the helper
router._set_deployment_num_retries_on_exception(exc, deployment)
# Verify num_retries was not set (deployment has no num_retries)
assert not hasattr(exc, "num_retries") or exc.num_retries is None
@@ -155,3 +155,36 @@ class TestPerDeploymentNumRetries:
kwargs = {}
router._update_kwargs_before_fallbacks(model="test-model", kwargs=kwargs)
assert kwargs["num_retries"] == 7 # Uses global
def test_set_deployment_num_retries_with_string_value(self):
"""
Test that _set_deployment_num_retries_on_exception handles string values
from environment variables correctly.
GitHub Issue: #19481
"""
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "openai/gpt-4",
"api_key": "test-key",
"num_retries": "6", # String value (as from env var)
},
},
],
num_retries=0, # Global setting
)
deployment = router.model_list[0]
class MockException(Exception):
pass
exc = MockException("test error")
# Call the helper
router._set_deployment_num_retries_on_exception(exc, deployment)
# Verify num_retries was converted from string to int
assert exc.num_retries == 6
tests/test_litellm/vector_stores/test_vector_store_registry.py
+4
View File
@@ -119,8 +119,12 @@ def test_add_vector_store_to_registry():
@respx.mock
def test_search_uses_registry_credentials():
"""search() should pull credentials from vector_store_registry when available"""
# Block all HTTP requests at the network level to prevent real API calls
respx.route().mock(return_value=httpx.Response(200, json={"object": "list", "data": []}))
vector_store = LiteLLM_ManagedVectorStore(
vector_store_id="vs1",
custom_llm_provider="bedrock",

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