scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Merge remote-tracking branch 'origin' into litellm_ui_endpoint_usage

Browse source
This commit is contained in:
yuneng-jiang 2026-01-10 09:08:38 -08:00
commit 0d90baf3ac
192 changed files with 16145 additions and 2582 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

64
.circleci/config.yml
View file

@ -2038,6 +2038,39 @@ jobs:
- run: python ./tests/code_coverage_tests/memory_test.py
- run: helm lint ./deploy/charts/litellm-helm
memory_leak_tests:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
resource_class: large
steps:
- setup_litellm_test_deps
- run:
name: Install Memory Test Dependencies
command: |
pip install "psutil>=5.9.0"
pip install "fastapi>=0.100.0"
pip install "httpx>=0.24.0"
pip install "uvicorn>=0.23.0"
- run:
name: Run Linear Memory Growth Tests
command: |
echo "Running memory leak tests individually to avoid baseline drift..."
echo "Running test_memory_baseline_1k..."
python -m pytest tests/load_tests/test_linear_memory_growth.py::test_memory_baseline_1k -v -s --tb=short
echo "Running test_memory_baseline_2k..."
python -m pytest tests/load_tests/test_linear_memory_growth.py::test_memory_baseline_2k -v -s --tb=short
echo "Running test_memory_baseline_4k..."
python -m pytest tests/load_tests/test_linear_memory_growth.py::test_memory_baseline_4k -v -s --tb=short
echo "Running test_memory_baseline_10k..."
python -m pytest tests/load_tests/test_linear_memory_growth.py::test_memory_baseline_10k -v -s --tb=short
echo "Running test_memory_baseline_30k..."
python -m pytest tests/load_tests/test_linear_memory_growth.py::test_memory_baseline_30k -v -s --tb=short
no_output_timeout: 60m
db_migration_disable_update_check:
machine:
image: ubuntu-2204:2023.10.1
@ -3557,12 +3590,34 @@ jobs:
name: Install Playwright Browsers
command: |
npx playwright install
- run:
name: Install Neon CLI
command: |
npm i -g neonctl
- run:
name: Create Neon branch
command: |
export EXPIRES_AT=$(date -u -d "+3 hours" +"%Y-%m-%dT%H:%M:%SZ")
echo "Expires at: $EXPIRES_AT"
neon branches create \
--project-id $NEON_PROJECT_ID \
--name preview/commit-${CIRCLE_SHA1:0:7} \
--expires-at $EXPIRES_AT \
--parent br-fancy-paper-ad1olsb3 \
--api-key $NEON_API_KEY || true
- run:
name: Run Docker container
command: |
E2E_UI_TEST_DATABASE_URL=$(neon connection-string \
--project-id $NEON_PROJECT_ID \
--api-key $NEON_API_KEY \
--branch preview/commit-${CIRCLE_SHA1:0:7} \
--database-name yuneng-trial-db \
--role neondb_owner)
echo $E2E_UI_TEST_DATABASE_URL
docker run -d \
-p 4000:4000 \
-e DATABASE_URL=$SMALL_DATABASE_URL \
-e DATABASE_URL=$E2E_UI_TEST_DATABASE_URL \
-e LITELLM_MASTER_KEY="sk-1234" \
-e OPENAI_API_KEY=$OPENAI_API_KEY \
-e UI_USERNAME="admin" \
@ -3765,6 +3820,12 @@ workflows:
only:
- main
- /litellm_.*/
- memory_leak_tests:
filters:
branches:
only:
- main
- /litellm_.*/
- ui_build:
filters:
branches:
@ -3792,6 +3853,7 @@ workflows:
- main
- /litellm_.*/
- e2e_ui_testing:
context: e2e_ui_tests
requires:
- ui_build
- build_docker_database_image

174
.github/workflows/label-component.yml vendored
View file

@ -11,134 +11,72 @@ jobs:
permissions:
issues: write
steps:
- name: Add SDK label
if: contains(github.event.issue.body, 'What part of LiteLLM is this about?\n\nSDK (litellm Python package)')
- name: Add component labels
uses: actions/github-script@v7
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
const labelName = 'sdk';
try {
await github.rest.issues.getLabel({
owner: context.repo.owner,
repo: context.repo.repo,
name: labelName
});
} catch (error) {
if (error.status === 404) {
await github.rest.issues.createLabel({
owner: context.repo.owner,
repo: context.repo.repo,
name: labelName,
color: '0E7C86',
description: 'Issues related to the litellm Python SDK'
});
} else {
throw error;
}
}
await github.rest.issues.addLabels({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: context.issue.number,
labels: [labelName]
});
const body = context.payload.issue.body;
if (!body) return;
- name: Add Proxy label
if: contains(github.event.issue.body, 'What part of LiteLLM is this about?\n\nProxy')
uses: actions/github-script@v7
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
const labelName = 'proxy';
try {
await github.rest.issues.getLabel({
owner: context.repo.owner,
repo: context.repo.repo,
name: labelName
});
} catch (error) {
if (error.status === 404) {
await github.rest.issues.createLabel({
owner: context.repo.owner,
repo: context.repo.repo,
name: labelName,
color: '5319E7',
description: 'Issues related to the LiteLLM Proxy'
});
} else {
throw error;
// Define component mappings with regex patterns that handle flexible whitespace
const components = [
{
pattern: /What part of LiteLLM is this about\?\s*SDK \(litellm Python package\)/,
label: 'sdk',
color: '0E7C86',
description: 'Issues related to the litellm Python SDK'
},
{
pattern: /What part of LiteLLM is this about\?\s*Proxy/,
label: 'proxy',
color: '5319E7',
description: 'Issues related to the LiteLLM Proxy'
},
{
pattern: /What part of LiteLLM is this about\?\s*UI Dashboard/,
label: 'ui-dashboard',
color: 'D876E3',
description: 'Issues related to the LiteLLM UI Dashboard'
},
{
pattern: /What part of LiteLLM is this about\?\s*Docs/,
label: 'docs',
color: 'FBCA04',
description: 'Issues related to LiteLLM documentation'
}
}
await github.rest.issues.addLabels({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: context.issue.number,
labels: [labelName]
});
];
- name: Add UI Dashboard label
if: contains(github.event.issue.body, 'What part of LiteLLM is this about?\n\nUI Dashboard')
uses: actions/github-script@v7
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
const labelName = 'ui-dashboard';
try {
await github.rest.issues.getLabel({
owner: context.repo.owner,
repo: context.repo.repo,
name: labelName
});
} catch (error) {
if (error.status === 404) {
await github.rest.issues.createLabel({
owner: context.repo.owner,
repo: context.repo.repo,
name: labelName,
color: 'D876E3',
description: 'Issues related to the LiteLLM UI Dashboard'
});
} else {
throw error;
}
}
await github.rest.issues.addLabels({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: context.issue.number,
labels: [labelName]
});
// Find matching component
for (const component of components) {
if (component.pattern.test(body)) {
// Ensure label exists
try {
await github.rest.issues.getLabel({
owner: context.repo.owner,
repo: context.repo.repo,
name: component.label
});
} catch (error) {
if (error.status === 404) {
await github.rest.issues.createLabel({
owner: context.repo.owner,
repo: context.repo.repo,
name: component.label,
color: component.color,
description: component.description
});
}
}
- name: Add Docs label
if: contains(github.event.issue.body, 'What part of LiteLLM is this about?\n\nDocs')
uses: actions/github-script@v7
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
const labelName = 'docs';
try {
await github.rest.issues.getLabel({
owner: context.repo.owner,
repo: context.repo.repo,
name: labelName
});
} catch (error) {
if (error.status === 404) {
await github.rest.issues.createLabel({
// Add label to issue
await github.rest.issues.addLabels({
owner: context.repo.owner,
repo: context.repo.repo,
name: labelName,
color: 'FBCA04',
description: 'Issues related to LiteLLM documentation'
issue_number: context.issue.number,
labels: [component.label]
});
} else {
throw error;
break;
}
}
await github.rest.issues.addLabels({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: context.issue.number,
labels: [labelName]
});

1
.github/workflows/publish-migrations.yml vendored
View file

@ -13,6 +13,7 @@ on:
jobs:
publish-migrations:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
services:
postgres:

2
README.md
View file

@ -262,6 +262,7 @@ Support for more providers. Missing a provider or LLM Platform, raise a [feature
| Provider | `/chat/completions` | `/messages` | `/responses` | `/embeddings` | `/image/generations` | `/audio/transcriptions` | `/audio/speech` | `/moderations` | `/batches` | `/rerank` |
|-------------------------------------------------------------------------------------|---------------------|-------------|--------------|---------------|----------------------|-------------------------|-----------------|----------------|-----------|-----------|
| [Abliteration (`abliteration`)](https://docs.litellm.ai/docs/providers/abliteration) | ✅ | | | | | | | | | |
| [AI/ML API (`aiml`)](https://docs.litellm.ai/docs/providers/aiml) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | |
| [AI21 (`ai21`)](https://docs.litellm.ai/docs/providers/ai21) | ✅ | ✅ | ✅ | | | | | | | |
| [AI21 Chat (`ai21_chat`)](https://docs.litellm.ai/docs/providers/ai21) | ✅ | ✅ | ✅ | | | | | | | |
@ -455,4 +456,3 @@ All these checks must pass before your PR can be merged.
<img src="https://contrib.rocks/image?repo=BerriAI/litellm" />
</a>

13
docs/my-website/docs/mcp_control.md
View file

@ -649,3 +649,16 @@ general_settings:
```
This is useful when you want discoverability for MCP offerings without granting additional execution privileges.
## Publish MCP Registry
If you want other systems—for example external agent frameworks such as MCP-capable IDEs running outside your network—to automatically discover the MCP servers hosted on LiteLLM, you can expose a Model Context Protocol Registry endpoint. This registry lists the built-in LiteLLM MCP server and every server you have configured, using the [official MCP Registry spec](https://github.com/modelcontextprotocol/registry).
1. Set `enable_mcp_registry: true` under `general_settings` in your proxy config (or DB settings) and restart the proxy.
2. LiteLLM will serve the registry at `GET /v1/mcp/registry.json`.
3. Each entry points to either `/mcp` (built-in server) or `/{mcp_server_name}/mcp` for your custom servers, so clients can connect directly using the advertised Streamable HTTP URL.
:::note Permissions still apply
The registry only advertises server URLs. Actual access control is still enforced by LiteLLM when the client connects to `/mcp` or `/{server}/mcp`, so publishing the registry does not bypass per-key permissions.
:::

93
docs/my-website/docs/observability/focus.md Normal file
View file

@ -0,0 +1,93 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Focus Export (Experimental)
:::caution Experimental feature
Focus Format export is under active development and currently considered experimental.
Interfaces, schema mappings, and configuration options may change as we iterate based on user feedback.
Please treat this integration as a preview and report any issues or suggestions to help us stabilize and improve the workflow.
:::
LiteLLM can emit usage data in the [FinOps FOCUS format](https://focus.finops.org/focus-specification/v1-2/) and push artifacts (for example Parquet files) to destinations such as Amazon S3. This enables downstream cost-analysis tooling to ingest a standardised dataset directly from LiteLLM.
LiteLLM currently conforms to the FinOps FOCUS v1.2 specification when emitting this dataset.
## Overview
| Property | Details |
|----------|---------|
| Destination | Export LiteLLM usage data in FOCUS format to managed storage (currently S3) |
| Callback name | `focus` |
| Supported operations | Automatic scheduled export |
| Data format | FOCUS Normalised Dataset (Parquet) |
## Environment Variables
### Common settings
| Variable | Required | Description |
|----------|----------|-------------|
| `FOCUS_PROVIDER` | No | Destination provider (defaults to `s3`). |
| `FOCUS_FORMAT` | No | Output format (currently only `parquet`). |
| `FOCUS_FREQUENCY` | No | Export cadence. Prefer `hourly` or `daily` for production; `interval` is intended for short test loops. Defaults to `hourly`. |
| `FOCUS_CRON_OFFSET` | No | Minute offset used for hourly/daily cron triggers. Defaults to `5`. |
| `FOCUS_INTERVAL_SECONDS` | No | Interval (seconds) when `FOCUS_FREQUENCY="interval"`. |
| `FOCUS_PREFIX` | No | Object key prefix/folder. Defaults to `focus_exports`. |
### S3 destination
| Variable | Required | Description |
|----------|----------|-------------|
| `FOCUS_S3_BUCKET_NAME` | Yes | Destination bucket for exported files. |
| `FOCUS_S3_REGION_NAME` | No | AWS region for the bucket. |
| `FOCUS_S3_ENDPOINT_URL` | No | Custom endpoint (useful for S3-compatible storage). |
| `FOCUS_S3_ACCESS_KEY` | Yes | AWS access key for uploads. |
| `FOCUS_S3_SECRET_KEY` | Yes | AWS secret key for uploads. |
| `FOCUS_S3_SESSION_TOKEN` | No | AWS session token if using temporary credentials. |
## Setup via Config
### Configure environment variables
```bash
export FOCUS_PROVIDER="s3"
export FOCUS_PREFIX="focus_exports"
# S3 example
export FOCUS_S3_BUCKET_NAME="my-litellm-focus-bucket"
export FOCUS_S3_REGION_NAME="us-east-1"
export FOCUS_S3_ACCESS_KEY="AKIA..."
export FOCUS_S3_SECRET_KEY="..."
```
### Update LiteLLM config
```yaml
model_list:
- model_name: gpt-4o
litellm_params:
model: openai/gpt-4o
api_key: sk-your-key
litellm_settings:
callbacks: ["focus"]
```
### Start the proxy
```bash
litellm --config /path/to/config.yaml
```
During boot LiteLLM registers the Focus logger and a background job that runs according to the configured frequency.
## Planned Enhancements
- Add "Setup on UI" flow alongside the current configuration-based setup.
- Add GCS / Azure Blob to the Destination options.
- Support CSV output alongside Parquet.
## Related Links
- [Focus](https://focus.finops.org/)

122
docs/my-website/docs/observability/qualifire_integration.md Normal file
View file

@ -0,0 +1,122 @@
import Image from '@theme/IdealImage';
# Qualifire - LLM Evaluation, Guardrails & Observability
[Qualifire](https://qualifire.ai/) provides real-time Agentic evaluations, guardrails and observability for production AI applications.
**Key Features:**
- **Evaluation** - Systematically assess AI behavior to detect hallucinations, jailbreaks, policy breaches, and other vulnerabilities
- **Guardrails** - Real-time interventions to prevent risks like brand damage, data leaks, and compliance breaches
- **Observability** - Complete tracing and logging for RAG pipelines, chatbots, and AI agents
- **Prompt Management** - Centralized prompt management with versioning and no-code studio
:::tip
Looking for Qualifire Guardrails? Check out the [Qualifire Guardrails Integration](../proxy/guardrails/qualifire.md) for real-time content moderation, prompt injection detection, PII checks, and more.
:::
## Pre-Requisites
1. Create an account on [Qualifire](https://app.qualifire.ai/)
2. Get your API key and webhook URL from the Qualifire dashboard
```bash
pip install litellm
```
## Quick Start
Use just 2 lines of code to instantly log your responses **across all providers** with Qualifire.
```python
litellm.callbacks = ["qualifire_eval"]
```
```python
import litellm
import os
# Set Qualifire credentials
os.environ["QUALIFIRE_API_KEY"] = "your-qualifire-api-key"
os.environ["QUALIFIRE_WEBHOOK_URL"] = "https://your-qualifire-webhook-url"
# LLM API Keys
os.environ['OPENAI_API_KEY'] = "your-openai-api-key"
# Set qualifire_eval as a callback & LiteLLM will send the data to Qualifire
litellm.callbacks = ["qualifire_eval"]
# OpenAI call
response = litellm.completion(
model="gpt-5",
messages=[
{"role": "user", "content": "Hi 👋 - i'm openai"}
]
)
```
## Using with LiteLLM Proxy
1. Setup config.yaml
```yaml
model_list:
- model_name: gpt-4o
litellm_params:
model: openai/gpt-4o
api_key: os.environ/OPENAI_API_KEY
litellm_settings:
callbacks: ["qualifire_eval"]
general_settings:
master_key: "sk-1234"
environment_variables:
QUALIFIRE_API_KEY: "your-qualifire-api-key"
QUALIFIRE_WEBHOOK_URL: "https://app.qualifire.ai/api/v1/webhooks/evaluations"
```
2. Start the proxy
```bash
litellm --config config.yaml
```
3. Test it!
```bash
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{ "model": "gpt-4o", "messages": [{"role": "user", "content": "Hi 👋 - i'm openai"}]}'
```
## Environment Variables
| Variable | Description |
| ----------------------- | ------------------------------------------------------ |
| `QUALIFIRE_API_KEY` | Your Qualifire API key for authentication |
| `QUALIFIRE_WEBHOOK_URL` | The Qualifire webhook endpoint URL from your dashboard |
## What Gets Logged?
The [LiteLLM Standard Logging Payload](https://docs.litellm.ai/docs/proxy/logging_spec) is sent to your Qualifire endpoint on each successful LLM API call.
This includes:
- Request messages and parameters
- Response content and metadata
- Token usage statistics
- Latency metrics
- Model information
- Cost data
Once data is in Qualifire, you can:
- Run evaluations to detect hallucinations, toxicity, and policy violations
- Set up guardrails to block or modify responses in real-time
- View traces across your entire AI pipeline
- Track performance and quality metrics over time

109
docs/my-website/docs/providers/abliteration.md Normal file
View file

@ -0,0 +1,109 @@
# Abliteration
## Overview
| Property | Details |
|-------|-------|
| Description | Abliteration provides an OpenAI-compatible `/chat/completions` endpoint. |
| Provider Route on LiteLLM | `abliteration/` |
| Link to Provider Doc | [Abliteration](https://abliteration.ai) |
| Base URL | `https://api.abliteration.ai/v1` |
| Supported Operations | [`/chat/completions`](#sample-usage) |
<br />
## Required Variables
```python showLineNumbers title="Environment Variables"
os.environ["ABLITERATION_API_KEY"] = "" # your Abliteration API key
```
## Sample Usage
```python showLineNumbers title="Abliteration Completion"
import os
from litellm import completion
os.environ["ABLITERATION_API_KEY"] = ""
response = completion(
model="abliteration/abliterated-model",
messages=[{"role": "user", "content": "Hello from LiteLLM"}],
)
print(response)
```
## Sample Usage - Streaming
```python showLineNumbers title="Abliteration Streaming Completion"
import os
from litellm import completion
os.environ["ABLITERATION_API_KEY"] = ""
response = completion(
model="abliteration/abliterated-model",
messages=[{"role": "user", "content": "Stream a short reply"}],
stream=True,
)
for chunk in response:
print(chunk)
```
## Usage with LiteLLM Proxy Server
1. Add the model to your proxy config:
```yaml showLineNumbers title="config.yaml"
model_list:
- model_name: abliteration-chat
litellm_params:
model: abliteration/abliterated-model
api_key: os.environ/ABLITERATION_API_KEY
```
2. Start the proxy:
```bash
litellm --config /path/to/config.yaml
```
## Direct API Usage (Bearer Token)
Use the environment variable as a Bearer token against the OpenAI-compatible endpoint:
`https://api.abliteration.ai/v1/chat/completions`.
```bash showLineNumbers title="cURL"
export ABLITERATION_API_KEY=""
curl https://api.abliteration.ai/v1/chat/completions \
-H "Authorization: Bearer ${ABLITERATION_API_KEY}" \
-H "Content-Type: application/json" \
-d '{
"model": "abliterated-model",
"messages": [{"role": "user", "content": "Hello from Abliteration"}]
}'
```
```python showLineNumbers title="Python (requests)"
import os
import requests
api_key = os.environ["ABLITERATION_API_KEY"]
response = requests.post(
"https://api.abliteration.ai/v1/chat/completions",
headers={
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
},
json={
"model": "abliterated-model",
"messages": [{"role": "user", "content": "Hello from Abliteration"}],
},
timeout=60,
)
print(response.json())
```

4
docs/my-website/docs/providers/anthropic.md
View file

@ -1692,9 +1692,9 @@ Assistant:
```
## Usage - PDF
## Usage - PDF
Pass base64 encoded PDF files to Anthropic models using the `image_url` field.
Pass base64 encoded PDF files to Anthropic models using the `file` content type with a `file_data` field.
<Tabs>
<TabItem value="sdk" label="SDK">

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

@ -7,7 +7,7 @@ ALL Bedrock models (Anthropic, Meta, Deepseek, Mistral, Amazon, etc.) are Suppor
| Property | Details |
|-------|-------|
| Description | Amazon Bedrock is a fully managed service that offers a choice of high-performing foundation models (FMs). |
| Provider Route on LiteLLM | `bedrock/`, [`bedrock/converse/`](#set-converse--invoke-route), [`bedrock/invoke/`](#set-invoke-route), [`bedrock/converse_like/`](#calling-via-internal-proxy), [`bedrock/llama/`](#deepseek-not-r1), [`bedrock/deepseek_r1/`](#deepseek-r1), [`bedrock/qwen3/`](#qwen3-imported-models), [`bedrock/qwen2/`](./bedrock_imported.md#qwen2-imported-models), [`bedrock/openai/`](./bedrock_imported.md#openai-compatible-imported-models-qwen-25-vl-etc) |
| Provider Route on LiteLLM | `bedrock/`, [`bedrock/converse/`](#set-converse--invoke-route), [`bedrock/invoke/`](#set-invoke-route), [`bedrock/converse_like/`](#calling-via-internal-proxy), [`bedrock/llama/`](#deepseek-not-r1), [`bedrock/deepseek_r1/`](#deepseek-r1), [`bedrock/qwen3/`](#qwen3-imported-models), [`bedrock/qwen2/`](./bedrock_imported.md#qwen2-imported-models), [`bedrock/openai/`](./bedrock_imported.md#openai-compatible-imported-models-qwen-25-vl-etc), [`bedrock/moonshot`](./bedrock_imported.md#moonshot-kimi-k2-thinking) |
| Provider Doc | [Amazon Bedrock ↗](https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html) |
| Supported OpenAI Endpoints | `/chat/completions`, `/completions`, `/embeddings`, `/images/generations` |
| Rerank Endpoint | `/rerank` |
@ -1941,6 +1941,7 @@ Here's an example of using a bedrock model with LiteLLM. For a complete list, re
| Mixtral 8x7B Instruct | `completion(model='bedrock/mistral.mixtral-8x7b-instruct-v0:1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| TwelveLabs Pegasus 1.2 (US) | `completion(model='bedrock/us.twelvelabs.pegasus-1-2-v1:0', messages=messages, mediaSource={...})` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| TwelveLabs Pegasus 1.2 (EU) | `completion(model='bedrock/eu.twelvelabs.pegasus-1-2-v1:0', messages=messages, mediaSource={...})` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| Moonshot Kimi K2 Thinking | `completion(model='bedrock/moonshot.kimi-k2-thinking', messages=messages)` or `completion(model='bedrock/invoke/moonshot.kimi-k2-thinking', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
## Bedrock Embedding

178
docs/my-website/docs/providers/bedrock_imported.md
View file

@ -431,4 +431,180 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
"max_tokens": 300,
"temperature": 0.5
}'
```
```
### Moonshot Kimi K2 Thinking
Moonshot AI's Kimi K2 Thinking model is now available on Amazon Bedrock. This model features advanced reasoning capabilities with automatic reasoning content extraction.
| Property | Details |
|----------|---------|
| Provider Route | `bedrock/moonshot.kimi-k2-thinking`, `bedrock/invoke/moonshot.kimi-k2-thinking` |
| Provider Documentation | [AWS Bedrock Moonshot Announcement ↗](https://aws.amazon.com/about-aws/whats-new/2025/12/amazon-bedrock-fully-managed-open-weight-models/) |
| Supported Parameters | `temperature`, `max_tokens`, `top_p`, `stream`, `tools`, `tool_choice` |
| Special Features | Reasoning content extraction, Tool calling |
#### Supported Features
- **Reasoning Content Extraction**: Automatically extracts `<reasoning>` tags and returns them as `reasoning_content` (similar to OpenAI's o1 models)
- **Tool Calling**: Full support for function/tool calling with tool responses
- **Streaming**: Both streaming and non-streaming responses
- **System Messages**: System message support
#### Basic Usage
<Tabs>
<TabItem value="sdk" label="SDK">
```python title="Moonshot Kimi K2 SDK Usage" showLineNumbers
from litellm import completion
import os
os.environ["AWS_ACCESS_KEY_ID"] = "your-aws-access-key"
os.environ["AWS_SECRET_ACCESS_KEY"] = "your-aws-secret-key"
os.environ["AWS_REGION_NAME"] = "us-west-2" # or your preferred region
# Basic completion
response = completion(
model="bedrock/moonshot.kimi-k2-thinking", # or bedrock/invoke/moonshot.kimi-k2-thinking
messages=[
{"role": "user", "content": "What is 2+2? Think step by step."}
],
temperature=0.7,
max_tokens=200
)
print(response.choices[0].message.content)
# Access reasoning content if present
if response.choices[0].message.reasoning_content:
print("Reasoning:", response.choices[0].message.reasoning_content)
```
</TabItem>
<TabItem value="proxy" label="Proxy">
**1. Add to config**
```yaml title="config.yaml" showLineNumbers
model_list:
- model_name: kimi-k2
litellm_params:
model: bedrock/moonshot.kimi-k2-thinking
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
aws_region_name: us-west-2
```
**2. Start proxy**
```bash title="Start LiteLLM Proxy" showLineNumbers
litellm --config /path/to/config.yaml
# RUNNING at http://0.0.0.0:4000
```
**3. Test it!**
```bash title="Test Kimi K2 via Proxy" showLineNumbers
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"model": "kimi-k2",
"messages": [
{
"role": "user",
"content": "What is 2+2? Think step by step."
}
],
"temperature": 0.7,
"max_tokens": 200
}'
```
</TabItem>
</Tabs>
#### Tool Calling Example
```python title="Kimi K2 with Tool Calling" showLineNumbers
from litellm import completion
import os
os.environ["AWS_ACCESS_KEY_ID"] = "your-aws-access-key"
os.environ["AWS_SECRET_ACCESS_KEY"] = "your-aws-secret-key"
os.environ["AWS_REGION_NAME"] = "us-west-2"
# Tool calling example
response = completion(
model="bedrock/moonshot.kimi-k2-thinking",
messages=[
{"role": "user", "content": "What's the weather in Tokyo?"}
],
tools=[
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather in a location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city name"
}
},
"required": ["location"]
}
}
}
]
)
if response.choices[0].message.tool_calls:
tool_call = response.choices[0].message.tool_calls[0]
print(f"Tool called: {tool_call.function.name}")
print(f"Arguments: {tool_call.function.arguments}")
```
#### Streaming Example
```python title="Kimi K2 Streaming" showLineNumbers
from litellm import completion
import os
os.environ["AWS_ACCESS_KEY_ID"] = "your-aws-access-key"
os.environ["AWS_SECRET_ACCESS_KEY"] = "your-aws-secret-key"
os.environ["AWS_REGION_NAME"] = "us-west-2"
response = completion(
model="bedrock/moonshot.kimi-k2-thinking",
messages=[
{"role": "user", "content": "Explain quantum computing in simple terms."}
],
stream=True,
temperature=0.7
)
for chunk in response:
if chunk.choices[0].delta.content:
print(chunk.choices[0].delta.content, end="")
# Check for reasoning content in streaming
if hasattr(chunk.choices[0].delta, 'reasoning_content') and chunk.choices[0].delta.reasoning_content:
print(f"\n[Reasoning: {chunk.choices[0].delta.reasoning_content}]")
```
#### Supported Parameters
| Parameter | Type | Description | Supported |
|-----------|------|-------------|-----------|
| `temperature` | float (0-1) | Controls randomness in output | ✅ |
| `max_tokens` | integer | Maximum tokens to generate | ✅ |
| `top_p` | float | Nucleus sampling parameter | ✅ |
| `stream` | boolean | Enable streaming responses | ✅ |
| `tools` | array | Tool/function definitions | ✅ |
| `tool_choice` | string/object | Tool choice specification | ✅ |
| `stop` | array | Stop sequences | ❌ (Not supported on Bedrock) |

194
docs/my-website/docs/providers/manus.md Normal file
View file

@ -0,0 +1,194 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Manus
Use Manus AI agents through LiteLLM's OpenAI-compatible Responses API.
| Property | Details |
|----------|---------|
| Description | Manus is an AI agent platform for complex reasoning tasks, document analysis, and multi-step workflows with asynchronous task execution. |
| Provider Route on LiteLLM | `manus/{agent_profile}` |
| Supported Operations | `/responses` (Responses API) |
| Provider Doc | [Manus API ↗](https://open.manus.im/docs/openai-compatibility) |
## Model Format
```shell
manus/{agent_profile}
```
**Examples:**
- `manus/manus-1.6` - General purpose agent
- `manus/manus-1.6-lite` - Lightweight agent for simple tasks
- `manus/manus-1.6-max` - Advanced agent for complex analysis
## LiteLLM Python SDK
```python showLineNumbers title="Basic Usage"
import litellm
import os
import time
# Set API key
os.environ["MANUS_API_KEY"] = "your-manus-api-key"
# Create task
response = litellm.responses(
model="manus/manus-1.6",
input="What's the capital of France?",
)
print(f"Task ID: {response.id}")
print(f"Status: {response.status}") # "running"
# Poll until complete
task_id = response.id
while response.status == "running":
time.sleep(5)
response = litellm.get_response(
response_id=task_id,
custom_llm_provider="manus",
)
print(f"Status: {response.status}")
# Get results
if response.status == "completed":
for message in response.output:
if message.role == "assistant":
print(message.content[0].text)
```
## LiteLLM AI Gateway
### Setup
```yaml showLineNumbers title="config.yaml"
model_list:
- model_name: manus-agent
litellm_params:
model: manus/manus-1.6
api_key: os.environ/MANUS_API_KEY
```
```bash title="Start Proxy"
litellm --config config.yaml
```
### Usage
<Tabs>
<TabItem value="curl" label="cURL">
```bash showLineNumbers title="Create Task"
# Create task
curl -X POST http://localhost:4000/responses \
-H "Authorization: Bearer your-proxy-key" \
-H "Content-Type: application/json" \
-d '{
"model": "manus-agent",
"input": "What is the capital of France?"
}'
# Response
{
"id": "task_abc123",
"status": "running",
"metadata": {
"task_url": "https://manus.im/app/task_abc123"
}
}
```
```bash showLineNumbers title="Poll for Completion"
# Check status (repeat until status is "completed")
curl http://localhost:4000/responses/task_abc123 \
-H "Authorization: Bearer your-proxy-key"
# When completed
{
"id": "task_abc123",
"status": "completed",
"output": [
{
"role": "user",
"content": [{"text": "What is the capital of France?"}]
},
{
"role": "assistant",
"content": [{"text": "The capital of France is Paris."}]
}
]
}
```
</TabItem>
<TabItem value="openai" label="OpenAI SDK">
```python showLineNumbers title="Create Task and Poll"
import openai
import time
client = openai.OpenAI(
base_url="http://localhost:4000",
api_key="your-proxy-key"
)
# Create task
response = client.responses.create(
model="manus-agent",
input="What is the capital of France?"
)
print(f"Task ID: {response.id}")
print(f"Status: {response.status}") # "running"
# Poll until complete
task_id = response.id
while response.status == "running":
time.sleep(5)
response = client.responses.retrieve(response_id=task_id)
print(f"Status: {response.status}")
# Get results
if response.status == "completed":
for message in response.output:
if message.role == "assistant":
print(message.content[0].text)
```
</TabItem>
</Tabs>
## How It Works
Manus operates as an **asynchronous agent API**:
1. **Create Task**: When you call `litellm.responses()`, Manus creates a task and returns immediately with `status: "running"`
2. **Task Executes**: The agent works on your request in the background
3. **Poll for Completion**: You must repeatedly call `litellm.get_response()` or `client.responses.retrieve()` until the status changes to `"completed"`
4. **Get Results**: Once completed, the `output` field contains the full conversation
**Task Statuses:**
- `running` - Agent is actively working
- `pending` - Agent is waiting for input
- `completed` - Task finished successfully
- `error` - Task failed
:::tip Production Usage
For production applications, use [webhooks](https://open.manus.im/docs/webhooks) instead of polling to get notified when tasks complete.
:::
## Supported Parameters
| Parameter | Supported | Notes |
|-----------|-----------|-------|
| `input` | ✅ | Text, images, or structured content |
| `stream` | ✅ | Fake streaming (task runs async) |
| `max_output_tokens` | ✅ | Limits response length |
| `previous_response_id` | ✅ | For multi-turn conversations |
## Related Documentation
- [LiteLLM Responses API](/docs/response_api)
- [Manus OpenAI Compatibility](https://open.manus.im/docs/openai-compatibility)

45
docs/my-website/docs/providers/vertex.md
View file

@ -35,6 +35,8 @@ import json
# !gcloud auth application-default login - run this to add vertex credentials to your env
## OR ##
file_path = 'path/to/vertex_ai_service_account.json'
## OR ##
export VERTEXAI_API_KEY="your-api-key"
# Load the JSON file
with open(file_path, 'r') as file:
@ -47,7 +49,7 @@ vertex_credentials_json = json.dumps(vertex_credentials)
response = completion(
model="vertex_ai/gemini-2.5-pro",
messages=[{ "content": "Hello, how are you?","role": "user"}],
vertex_credentials=vertex_credentials_json
vertex_credentials=vertex_credentials_json # Can remove this is added VERTEXAI_API_KEY in env
)
```
@ -1329,15 +1331,41 @@ Here's how to use Vertex AI with the LiteLLM Proxy Server
## Authentication - vertex_project, vertex_location, etc.
LiteLLM supports two authentication methods for Vertex AI:
1. **API Key Authentication** (Recommended for getting started)
2. **Service Account Credentials** (Recommended for production)
Set your vertex credentials via:
- dynamic params
OR
- env vars
### **Authentication Method 1:
### **Dynamic Params**
The simplest way to authenticate with Vertex AI. You can set:
- `api_key` (str) - Your Vertex AI API key
You can set:
**Environment Variables:**
```bash
export VERTEXAI_API_KEY="your-api-key"
```
**Or pass as parameters:**
```python
from litellm import completion
response = completion(
model="vertex_ai/gemini-2.0-flash-exp",
messages=[{"role": "user", "content": "Hello!"}],
api_key="your-vertex-api-key",
)
```
### **Authentication Method 2: Service Account Credentials**
For production environments with fine-grained access control. You can set:
- `vertex_credentials` (str) - can be a json string or filepath to your vertex ai service account.json
- `vertex_location` (str) - place where vertex model is deployed (us-central1, asia-southeast1, etc.). Some models support the global location, please see [Vertex AI documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#supported_models)
- `vertex_project` Optional[str] - use if vertex project different from the one in vertex_credentials
@ -1392,7 +1420,16 @@ model_list:
### **Environment Variables**
You can set:
#### For API Key Authentication:
- `VERTEXAI_API_KEY` or `VERTEX_API_KEY` - Your Vertex AI API key
```bash
export VERTEXAI_API_KEY="your-vertex-api-key"
```
#### For Service Account Authentication:
- `GOOGLE_APPLICATION_CREDENTIALS` - store the filepath for your service_account.json in here (used by vertex sdk directly).
- VERTEXAI_LOCATION - place where vertex model is deployed (us-central1, asia-southeast1, etc.)
- VERTEXAI_PROJECT - Optional[str] - use if vertex project different from the one in vertex_credentials

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

@ -146,6 +146,7 @@ router_settings:
cooldown_time: 30 # (in seconds) how long to cooldown model if fails/min > allowed_fails
disable_cooldowns: True # bool - Disable cooldowns for all models
enable_tag_filtering: True # bool - Use tag based routing for requests
tag_filtering_match_any: True # bool - Tag matching behavior (only when enable_tag_filtering=true). `true`: match if deployment has ANY requested tag; `false`: match only if deployment has ALL requested tags
retry_policy: { # Dict[str, int]: retry policy for different types of exceptions
"AuthenticationErrorRetries": 3,
"TimeoutErrorRetries": 3,
@ -293,6 +294,7 @@ router_settings:
cooldown_time: 30 # (in seconds) how long to cooldown model if fails/min > allowed_fails
disable_cooldowns: True # bool - Disable cooldowns for all models
enable_tag_filtering: True # bool - Use tag based routing for requests
tag_filtering_match_any: True # bool - Tag matching behavior (only when enable_tag_filtering=true). `true`: match if deployment has ANY requested tag; `false`: match only if deployment has ALL requested tags
retry_policy: { # Dict[str, int]: retry policy for different types of exceptions
"AuthenticationErrorRetries": 3,
"TimeoutErrorRetries": 3,
@ -322,6 +324,7 @@ router_settings:
| content_policy_fallbacks | array of objects | Specifies fallback models for content policy violations. [More information here](reliability) |
| fallbacks | array of objects | Specifies fallback models for all types of errors. [More information here](reliability) |
| enable_tag_filtering | boolean | If true, uses tag based routing for requests [Tag Based Routing](tag_routing) |
| tag_filtering_match_any | boolean | Tag matching behavior (only when enable_tag_filtering=true). `true`: match if deployment has ANY requested tag; `false`: match only if deployment has ALL requested tags |
| cooldown_time | integer | The duration (in seconds) to cooldown a model if it exceeds the allowed failures. |
| disable_cooldowns | boolean | If true, disables cooldowns for all models. [More information here](reliability) |
| retry_policy | object | Specifies the number of retries for different types of exceptions. [More information here](reliability) |
@ -578,6 +581,18 @@ router_settings:
| FIREWORKS_AI_56_B_MOE | Size parameter for Fireworks AI 56B MOE model. Default is 56
| FIREWORKS_AI_80_B | Size parameter for Fireworks AI 80B model. Default is 80
| FIREWORKS_AI_176_B_MOE | Size parameter for Fireworks AI 176B MOE model. Default is 176
| FOCUS_PROVIDER | Destination provider for Focus exports (e.g., `s3`). Defaults to `s3`.
| FOCUS_FORMAT | Output format for Focus exports. Defaults to `parquet`.
| FOCUS_FREQUENCY | Frequency for scheduled Focus exports (`hourly`, `daily`, or `interval`). Defaults to `hourly`.
| FOCUS_CRON_OFFSET | Minute offset used when scheduling hourly/daily Focus exports. Defaults to `5` minutes.
| FOCUS_INTERVAL_SECONDS | Interval (in seconds) for Focus exports when `frequency` is `interval`.
| FOCUS_PREFIX | Object key prefix (or folder) used when uploading Focus export files. Defaults to `focus_exports`.
| FOCUS_S3_BUCKET_NAME | S3 bucket to upload Focus export files when using the S3 destination.
| FOCUS_S3_REGION_NAME | AWS region for the Focus export S3 bucket.
| FOCUS_S3_ENDPOINT_URL | Custom endpoint for the Focus export S3 client (optional; useful for S3-compatible storage).
| FOCUS_S3_ACCESS_KEY | AWS access key ID used by the Focus export S3 client.
| FOCUS_S3_SECRET_KEY | AWS secret access key used by the Focus export S3 client.
| FOCUS_S3_SESSION_TOKEN | AWS session token used by the Focus export S3 client (optional).
| FUNCTION_DEFINITION_TOKEN_COUNT | Token count for function definitions. Default is 9
| GALILEO_BASE_URL | Base URL for Galileo platform
| GALILEO_PASSWORD | Password for Galileo authentication

41
docs/my-website/docs/proxy/guardrails/qualifire.md
View file

@ -8,13 +8,7 @@ Use [Qualifire](https://qualifire.ai) to evaluate LLM outputs for quality, safet
## Quick Start
### 1. Install the Qualifire SDK
```bash
pip install qualifire
```
### 2. Define Guardrails on your LiteLLM config.yaml
### 1. Define Guardrails on your LiteLLM config.yaml
Define your guardrails under the `guardrails` section:
@ -61,13 +55,13 @@ guardrails:
- `post_call` Run **after** LLM call, on **input & output**
- `during_call` Run **during** LLM call, on **input**. Same as `pre_call` but runs in parallel as LLM call. Response not returned until guardrail check completes
### 3. Start LiteLLM Gateway
### 2. Start LiteLLM Gateway
```shell
litellm --config config.yaml --detailed_debug
```
### 4. Test request
### 3. Test request
**[Langchain, OpenAI SDK Usage Examples](../proxy/user_keys#request-format)**
@ -142,7 +136,7 @@ guardrails:
evaluation_id: eval_abc123 # Your evaluation ID from Qualifire dashboard
```
When `evaluation_id` is provided, LiteLLM will use `invoke_evaluation()` instead of `evaluate()`, running the pre-configured evaluation from your dashboard.
When `evaluation_id` is provided, LiteLLM will use the invoke evaluation API endpoint instead of the evaluate endpoint, running the pre-configured evaluation from your dashboard.
## Available Checks
@ -213,19 +207,19 @@ guardrails:
### Parameter Reference
| Parameter | Type | Default | Description |
| ------------------------------ | ----------- | --------------------------- | -------------------------------------------------------- |
| `api_key` | `str` | `QUALIFIRE_API_KEY` env var | Your Qualifire API key |
| `api_base` | `str` | `None` | Custom API base URL (optional) |
| `evaluation_id` | `str` | `None` | Pre-configured evaluation ID from Qualifire dashboard |
| `prompt_injections` | `bool` | `true` (if no other checks) | Enable prompt injection detection |
| `hallucinations_check` | `bool` | `None` | Enable hallucination detection |
| `grounding_check` | `bool` | `None` | Enable grounding verification |
| `pii_check` | `bool` | `None` | Enable PII detection |
| `content_moderation_check` | `bool` | `None` | Enable content moderation |
| `tool_selection_quality_check` | `bool` | `None` | Enable tool selection quality check |
| `assertions` | `List[str]` | `None` | Custom assertions to validate |
| `on_flagged` | `str` | `"block"` | Action when content is flagged: `"block"` or `"monitor"` |
| Parameter | Type | Default | Description |
| ------------------------------ | ----------- | ---------------------------- | -------------------------------------------------------- |
| `api_key` | `str` | `QUALIFIRE_API_KEY` env var | Your Qualifire API key |
| `api_base` | `str` | `https://proxy.qualifire.ai` | Custom API base URL (optional) |
| `evaluation_id` | `str` | `None` | Pre-configured evaluation ID from Qualifire dashboard |
| `prompt_injections` | `bool` | `true` (if no other checks) | Enable prompt injection detection |
| `hallucinations_check` | `bool` | `None` | Enable hallucination detection |
| `grounding_check` | `bool` | `None` | Enable grounding verification |
| `pii_check` | `bool` | `None` | Enable PII detection |
| `content_moderation_check` | `bool` | `None` | Enable content moderation |
| `tool_selection_quality_check` | `bool` | `None` | Enable tool selection quality check |
| `assertions` | `List[str]` | `None` | Custom assertions to validate |
| `on_flagged` | `str` | `"block"` | Action when content is flagged: `"block"` or `"monitor"` |
### Default Behavior
@ -261,4 +255,3 @@ This evaluates whether the LLM selected the appropriate tools and provided corre
- [Qualifire Documentation](https://docs.qualifire.ai)
- [Qualifire Dashboard](https://app.qualifire.ai)
- [Qualifire Python SDK](https://github.com/qualifire-dev/qualifire-python-sdk)

2
docs/my-website/docs/proxy/logging.md
View file

@ -67,7 +67,7 @@ Set `litellm.turn_off_message_logging=True` This will prevent the messages and r
<TabItem value="global" label="Global">
**1. Setup config.yaml **
**1. Setup config.yaml**
```yaml
model_list:
- model_name: gpt-3.5-turbo

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

@ -55,6 +55,7 @@ const sidebars = {
"proxy/guardrails/test_playground",
"proxy/guardrails/litellm_content_filter",
...[
"proxy/guardrails/qualifire",
"proxy/guardrails/aim_security",
"proxy/guardrails/onyx_security",
"proxy/guardrails/aporia_api",
@ -653,12 +654,13 @@ const sidebars = {
"providers/bedrock_writer",
"providers/bedrock_batches",
"providers/aws_polly",
"providers/bedrock_vector_store",
]
},
"providers/litellm_proxy",
"providers/ai21",
"providers/aiml",
"providers/bedrock_vector_store",
]
},
"providers/litellm_proxy",
"providers/abliteration",
"providers/ai21",
"providers/aiml",
"providers/aleph_alpha",
"providers/amazon_nova",
"providers/anyscale",
@ -710,6 +712,7 @@ const sidebars = {
"providers/llamafile",
"providers/llamagate",
"providers/lm_studio",
"providers/manus",
"providers/meta_llama",
"providers/milvus_vector_stores",
"providers/mistral",

9
litellm-proxy-extras/litellm_proxy_extras/migrations/20260108_add_user_email_lower_idx/migration.sql Normal file
View file

@ -0,0 +1,9 @@
-- CreateIndex
-- Fixes performance issue in _check_duplicate_user_email function
-- by enabling fast case-insensitive email lookups.
--
-- Without this index, queries with mode: "insensitive" cause full table scans.
-- With this index, PostgreSQL can use an Index Scan for O(log n) performance.
--
-- Related: GitHub Issue #18411
CREATE INDEX "LiteLLM_UserTable_user_email_lower_idx" ON "LiteLLM_UserTable"(LOWER("user_email"));

10
litellm/__init__.py
View file

@ -134,6 +134,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"bitbucket",
"gitlab",
"cloudzero",
"focus",
"posthog",
"levo",
]
@ -486,6 +487,7 @@ vertex_mistral_models: Set = set()
vertex_openai_models: Set = set()
vertex_minimax_models: Set = set()
vertex_moonshot_models: Set = set()
vertex_zai_models: Set = set()
ai21_models: Set = set()
ai21_chat_models: Set = set()
nlp_cloud_models: Set = set()
@ -664,6 +666,9 @@ def add_known_models():
elif value.get("litellm_provider") == "vertex_ai-moonshot_models":
key = key.replace("vertex_ai/", "")
vertex_moonshot_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-zai_models":
key = key.replace("vertex_ai/", "")
vertex_zai_models.add(key)
elif value.get("litellm_provider") == "ai21":
if value.get("mode") == "chat":
ai21_chat_models.add(key)
@ -950,7 +955,8 @@ models_by_provider: dict = {
| vertex_language_models
| vertex_deepseek_models
| vertex_minimax_models
| vertex_moonshot_models,
| vertex_moonshot_models
| vertex_zai_models,
"ai21": ai21_models,
"bedrock": bedrock_models | bedrock_converse_models,
"petals": petals_models,
@ -1338,6 +1344,7 @@ if TYPE_CHECKING:
from .llms.bedrock.chat.invoke_transformations.amazon_llama_transformation import AmazonLlamaConfig as AmazonLlamaConfig
from .llms.bedrock.chat.invoke_transformations.amazon_deepseek_transformation import AmazonDeepSeekR1Config as AmazonDeepSeekR1Config
from .llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation import AmazonMistralConfig as AmazonMistralConfig
from .llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import AmazonMoonshotConfig as AmazonMoonshotConfig
from .llms.bedrock.chat.invoke_transformations.amazon_titan_transformation import AmazonTitanConfig as AmazonTitanConfig
from .llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation import AmazonTwelveLabsPegasusConfig as AmazonTwelveLabsPegasusConfig
from .llms.bedrock.chat.invoke_transformations.base_invoke_transformation import AmazonInvokeConfig as AmazonInvokeConfig
@ -1367,6 +1374,7 @@ if TYPE_CHECKING:
from .llms.azure.responses.o_series_transformation import AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig
from .llms.xai.responses.transformation import XAIResponsesAPIConfig as XAIResponsesAPIConfig
from .llms.litellm_proxy.responses.transformation import LiteLLMProxyResponsesAPIConfig as LiteLLMProxyResponsesAPIConfig
from .llms.manus.responses.transformation import ManusResponsesAPIConfig as ManusResponsesAPIConfig
from .llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig
from .llms.openai.chat.o_series_transformation import OpenAIOSeriesConfig as OpenAIOSeriesConfig, OpenAIOSeriesConfig as OpenAIO1Config
from .llms.anthropic.skills.transformation import AnthropicSkillsConfig as AnthropicSkillsConfig

4
litellm/_lazy_imports_registry.py
View file

@ -165,6 +165,7 @@ LLM_CONFIG_NAMES = (
"AmazonLlamaConfig",
"AmazonDeepSeekR1Config",
"AmazonMistralConfig",
"AmazonMoonshotConfig",
"AmazonTitanConfig",
"AmazonTwelveLabsPegasusConfig",
"AmazonInvokeConfig",
@ -252,6 +253,7 @@ LLM_CONFIG_NAMES = (
"IBMWatsonXAudioTranscriptionConfig",
"GithubCopilotConfig",
"GithubCopilotResponsesAPIConfig",
"ManusResponsesAPIConfig",
"GithubCopilotEmbeddingConfig",
"NebiusConfig",
"WandbConfig",
@ -556,6 +558,7 @@ _LLM_CONFIGS_IMPORT_MAP = {
"AmazonLlamaConfig": (".llms.bedrock.chat.invoke_transformations.amazon_llama_transformation", "AmazonLlamaConfig"),
"AmazonDeepSeekR1Config": (".llms.bedrock.chat.invoke_transformations.amazon_deepseek_transformation", "AmazonDeepSeekR1Config"),
"AmazonMistralConfig": (".llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation", "AmazonMistralConfig"),
"AmazonMoonshotConfig": (".llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation", "AmazonMoonshotConfig"),
"AmazonTitanConfig": (".llms.bedrock.chat.invoke_transformations.amazon_titan_transformation", "AmazonTitanConfig"),
"AmazonTwelveLabsPegasusConfig": (".llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation", "AmazonTwelveLabsPegasusConfig"),
"AmazonInvokeConfig": (".llms.bedrock.chat.invoke_transformations.base_invoke_transformation", "AmazonInvokeConfig"),
@ -588,6 +591,7 @@ _LLM_CONFIGS_IMPORT_MAP = {
"AzureOpenAIOSeriesResponsesAPIConfig": (".llms.azure.responses.o_series_transformation", "AzureOpenAIOSeriesResponsesAPIConfig"),
"XAIResponsesAPIConfig": (".llms.xai.responses.transformation", "XAIResponsesAPIConfig"),
"LiteLLMProxyResponsesAPIConfig": (".llms.litellm_proxy.responses.transformation", "LiteLLMProxyResponsesAPIConfig"),
"ManusResponsesAPIConfig": (".llms.manus.responses.transformation", "ManusResponsesAPIConfig"),
"GoogleAIStudioInteractionsConfig": (".llms.gemini.interactions.transformation", "GoogleAIStudioInteractionsConfig"),
"OpenAIOSeriesConfig": (".llms.openai.chat.o_series_transformation", "OpenAIOSeriesConfig"),
"AnthropicSkillsConfig": (".llms.anthropic.skills.transformation", "AnthropicSkillsConfig"),

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

@ -31,6 +31,7 @@ from litellm.llms.base_llm.bridges.completion_transformation import (
CompletionTransformationBridge,
)
from litellm.types.llms.openai import (
ChatCompletionAnnotation,
ChatCompletionToolParamFunctionChunk,
Reasoning,
ResponsesAPIOptionalRequestParams,
@ -90,9 +91,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
content_type = content_item.get("type")
if content_type == "output_text":
response_text = content_item.get("text", "")
# Extract annotations from content if present
annotations = LiteLLMResponsesTransformationHandler._convert_annotations_to_chat_format(
content_item.get("annotations", None)
)
msg = Message(
role=item.get("role", "assistant"),
content=response_text if response_text else "",
annotations=annotations,
)
choice = Choices(message=msg, finish_reason="stop", index=index)
return choice, index + 1
@ -364,10 +370,16 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
elif isinstance(item, ResponseOutputMessage):
for content in item.content:
response_text = getattr(content, "text", "")
# Extract annotations from content if present
raw_annotations = getattr(content, "annotations", None)
annotations = LiteLLMResponsesTransformationHandler._convert_annotations_to_chat_format(
raw_annotations
)
msg = Message(
role=item.role,
content=response_text if response_text else "",
reasoning_content=reasoning_content,
annotations=annotations,
)
choices.append(
@ -763,6 +775,42 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
return {"format": {"type": "text"}}
return None
@staticmethod
def _convert_annotations_to_chat_format(
annotations: Optional[List[Any]],
) -> Optional[List["ChatCompletionAnnotation"]]:
"""
Convert annotations from Responses API to Chat Completions format.
Annotations are already in compatible format between both APIs,
so we just need to convert Pydantic models to dicts.
"""
if not annotations:
return None
result: List[ChatCompletionAnnotation] = []
for annotation in annotations:
try:
# Convert Pydantic models to dicts (handles both v1 and v2)
if hasattr(annotation, "model_dump"):
annotation_dict = annotation.model_dump()
elif hasattr(annotation, "dict"):
annotation_dict = annotation.dict()
elif isinstance(annotation, dict):
annotation_dict = annotation
else:
# Skip unsupported annotation types
verbose_logger.debug(f"Skipping unsupported annotation type: {type(annotation)}")
continue
result.append(annotation_dict) # type: ignore
except Exception as e:
# Skip malformed annotations
verbose_logger.debug(f"Skipping malformed annotation: {annotation}, error: {e}")
continue
return result if result else None
def _map_responses_status_to_finish_reason(self, status: Optional[str]) -> str:
"""Map responses API status to chat completion finish_reason"""

1
litellm/constants.py
View file

@ -909,6 +909,7 @@ BEDROCK_INVOKE_PROVIDERS_LITERAL = Literal[
"twelvelabs",
"openai",
"stability",
"moonshot",
]
BEDROCK_EMBEDDING_PROVIDERS_LITERAL = Literal[

9
litellm/google_genai/main.py
View file

@ -130,6 +130,9 @@ class GenerateContentHelper:
api_key=litellm_params.api_key,
)
if litellm_params.custom_llm_provider is None:
litellm_params.custom_llm_provider = custom_llm_provider
# get provider config
generate_content_provider_config: Optional[
BaseGoogleGenAIGenerateContentConfig
@ -407,6 +410,9 @@ async def agenerate_content_stream(
# Check if we should use the adapter (when provider config is None)
if setup_result.generate_content_provider_config is None:
if "stream" in kwargs:
kwargs.pop("stream", None)
# Use the adapter to convert to completion format
return (
await GenerateContentToCompletionHandler.async_generate_content_handler(
@ -490,6 +496,9 @@ def generate_content_stream(
# Check if we should use the adapter (when provider config is None)
if setup_result.generate_content_provider_config is None:
if "stream" in kwargs:
kwargs.pop("stream", None)
# Use the adapter to convert to completion format
return GenerateContentToCompletionHandler.generate_content_handler(
model=model,

35
litellm/integrations/braintrust_logging.py
View file

@ -225,10 +225,13 @@ class BraintrustLogger(CustomLogger):
"id": litellm_call_id,
"input": prompt["messages"],
"metadata": standard_logging_object,
"tags": tags,
"span_attributes": {"name": span_name, "type": "llm"},
}
# Braintrust cannot specify 'tags' for non-root spans
if dynamic_metadata.get("root_span_id") is None:
request_data["tags"] = tags
# Only add those that are not None (or falsy)
for key, value in span_attributes.items():
if value:
@ -351,14 +354,37 @@ class BraintrustLogger(CustomLogger):
# Allow metadata override for span name
span_name = dynamic_metadata.get("span_name", "Chat Completion")
# Span parents is a special case
span_parents = dynamic_metadata.get("span_parents")
# Convert comma-separated string to list if present
if span_parents:
span_parents = [s.strip() for s in span_parents.split(",") if s.strip()]
# Add optional span attributes only if present
span_attributes = {
"span_id": dynamic_metadata.get("span_id"),
"root_span_id": dynamic_metadata.get("root_span_id"),
"span_parents": span_parents,
}
request_data = {
"id": litellm_call_id,
"input": prompt["messages"],
"output": output,
"metadata": standard_logging_object,
"tags": tags,
"span_attributes": {"name": span_name, "type": "llm"},
}
# Braintrust cannot specify 'tags' for non-root spans
if dynamic_metadata.get("root_span_id") is None:
request_data["tags"] = tags
# Only add those that are not None (or falsy)
for key, value in span_attributes.items():
if value:
request_data[key] = value
if choices is not None:
request_data["output"] = [choice.dict() for choice in choices]
else:
@ -367,9 +393,6 @@ class BraintrustLogger(CustomLogger):
if metrics is not None:
request_data["metrics"] = metrics
if metrics is not None:
request_data["metrics"] = metrics
try:
await self.global_braintrust_http_handler.post(
url=f"{self.api_base}/project_logs/{project_id}/insert",

182
litellm/integrations/cloudzero/database.py
View file

@ -19,7 +19,7 @@
"""Database connection and data extraction for LiteLLM."""
from datetime import datetime
from typing import Any, Dict, Optional
from typing import Any, Optional, List
import polars as pl
@ -46,19 +46,9 @@ class LiteLLMDatabase:
"""Retrieve usage data from LiteLLM daily user spend table."""
client = self._ensure_prisma_client()
# Build WHERE clause for time filtering
where_conditions = []
if start_time_utc:
where_conditions.append(f"dus.updated_at >= '{start_time_utc.isoformat()}'")
if end_time_utc:
where_conditions.append(f"dus.updated_at <= '{end_time_utc.isoformat()}'")
where_clause = ""
if where_conditions:
where_clause = "WHERE " + " AND ".join(where_conditions)
# Query to get user spend data with team information
query = f"""
# Query to get user spend data with team information. Use parameter binding to
# avoid SQL injection from user-supplied timestamps or limits.
query = """
SELECT
dus.id,
dus.date,
@ -85,163 +75,27 @@ class LiteLLMDatabase:
LEFT JOIN "LiteLLM_VerificationToken" vt ON dus.api_key = vt.token
LEFT JOIN "LiteLLM_TeamTable" tt ON vt.team_id = tt.team_id
LEFT JOIN "LiteLLM_UserTable" ut ON dus.user_id = ut.user_id
{where_clause}
WHERE ($1::timestamptz IS NULL OR dus.updated_at >= $1::timestamptz)
AND ($2::timestamptz IS NULL OR dus.updated_at <= $2::timestamptz)
ORDER BY dus.date DESC, dus.created_at DESC
"""
if limit:
query += f" LIMIT {limit}"
params: List[Any] = [
start_time_utc,
end_time_utc,
]
if limit is not None:
try:
params.append(int(limit))
except (TypeError, ValueError):
raise ValueError("limit must be an integer")
query += " LIMIT $3"
try:
db_response = await client.db.query_raw(query)
db_response = await client.db.query_raw(query, *params)
# Convert the response to polars DataFrame with full schema inference
# This prevents schema mismatch errors when data types vary across rows
return pl.DataFrame(db_response, infer_schema_length=None)
except Exception as e:
raise Exception(f"Error retrieving usage data: {str(e)}")
async def get_table_info(self) -> Dict[str, Any]:
"""Get information about the daily user spend table."""
client = self._ensure_prisma_client()
try:
# Get row count from user spend table
user_count = await self._get_table_row_count("LiteLLM_DailyUserSpend")
# Get column structure from user spend table
query = """
SELECT column_name, data_type, is_nullable
FROM information_schema.columns
WHERE table_name = 'LiteLLM_DailyUserSpend'
ORDER BY ordinal_position;
"""
columns_response = await client.db.query_raw(query)
return {
"columns": columns_response,
"row_count": user_count,
"table_name": "LiteLLM_DailyUserSpend",
}
except Exception as e:
raise Exception(f"Error getting table info: {str(e)}")
async def _get_table_row_count(self, table_name: str) -> int:
"""Get row count from specified table."""
client = self._ensure_prisma_client()
try:
query = f'SELECT COUNT(*) as count FROM "{table_name}"'
response = await client.db.query_raw(query)
if response and len(response) > 0:
return response[0].get("count", 0)
return 0
except Exception:
return 0
async def discover_all_tables(self) -> Dict[str, Any]:
"""Discover all tables in the LiteLLM database and their schemas."""
client = self._ensure_prisma_client()
try:
# Get all LiteLLM tables
litellm_tables_query = """
SELECT table_name
FROM information_schema.tables
WHERE table_schema = 'public'
AND table_name LIKE 'LiteLLM_%'
ORDER BY table_name;
"""
tables_response = await client.db.query_raw(litellm_tables_query)
table_names = [row["table_name"] for row in tables_response]
# Get detailed schema for each table
tables_info = {}
for table_name in table_names:
# Get column information
columns_query = """
SELECT
column_name,
data_type,
is_nullable,
column_default,
character_maximum_length,
numeric_precision,
numeric_scale,
ordinal_position
FROM information_schema.columns
WHERE table_name = $1
AND table_schema = 'public'
ORDER BY ordinal_position;
"""
columns_response = await client.db.query_raw(columns_query, table_name)
# Get primary key information
pk_query = """
SELECT a.attname
FROM pg_index i
JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
WHERE i.indrelid = $1::regclass AND i.indisprimary;
"""
pk_response = await client.db.query_raw(pk_query, f'"{table_name}"')
primary_keys = (
[row["attname"] for row in pk_response] if pk_response else []
)
# Get foreign key information
fk_query = """
SELECT
tc.constraint_name,
kcu.column_name,
ccu.table_name AS foreign_table_name,
ccu.column_name AS foreign_column_name
FROM information_schema.table_constraints AS tc
JOIN information_schema.key_column_usage AS kcu
ON tc.constraint_name = kcu.constraint_name
JOIN information_schema.constraint_column_usage AS ccu
ON ccu.constraint_name = tc.constraint_name
WHERE tc.constraint_type = 'FOREIGN KEY'
AND tc.table_name = $1;
"""
fk_response = await client.db.query_raw(fk_query, table_name)
foreign_keys = fk_response if fk_response else []
# Get indexes
indexes_query = """
SELECT
i.relname AS index_name,
array_agg(a.attname ORDER BY a.attnum) AS column_names,
ix.indisunique AS is_unique
FROM pg_class t
JOIN pg_index ix ON t.oid = ix.indrelid
JOIN pg_class i ON i.oid = ix.indexrelid
JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(ix.indkey)
WHERE t.relname = $1
AND t.relkind = 'r'
GROUP BY i.relname, ix.indisunique
ORDER BY i.relname;
"""
indexes_response = await client.db.query_raw(indexes_query, table_name)
indexes = indexes_response if indexes_response else []
# Get row count
try:
row_count = await self._get_table_row_count(table_name)
except Exception:
row_count = 0
tables_info[table_name] = {
"columns": columns_response,
"primary_keys": primary_keys,
"foreign_keys": foreign_keys,
"indexes": indexes,
"row_count": row_count,
}
return {
"tables": tables_info,
"table_count": len(table_names),
"table_names": table_names,
}
except Exception as e:
raise Exception(f"Error discovering tables: {str(e)}")

0
litellm/integrations/focus/__init__.py Normal file
View file

113
litellm/integrations/focus/database.py Normal file
View file

@ -0,0 +1,113 @@
"""Database access helpers for Focus export."""
from __future__ import annotations
from datetime import datetime
from typing import Any, Dict, Optional
import polars as pl
class FocusLiteLLMDatabase:
"""Retrieves LiteLLM usage data for Focus export workflows."""
def _ensure_prisma_client(self):
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise RuntimeError(
"Database not connected. Connect a database to your proxy - "
"https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
)
return prisma_client
async def get_usage_data(
self,
*,
limit: Optional[int] = None,
start_time_utc: Optional[datetime] = None,
end_time_utc: Optional[datetime] = None,
) -> pl.DataFrame:
"""Return usage data for the requested window."""
client = self._ensure_prisma_client()
where_clauses: list[str] = []
query_params: list[Any] = []
placeholder_index = 1
if start_time_utc:
where_clauses.append(f"dus.updated_at >= ${placeholder_index}::timestamptz")
query_params.append(start_time_utc)
placeholder_index += 1
if end_time_utc:
where_clauses.append(f"dus.updated_at <= ${placeholder_index}::timestamptz")
query_params.append(end_time_utc)
placeholder_index += 1
where_clause = ""
if where_clauses:
where_clause = "WHERE " + " AND ".join(where_clauses)
limit_clause = ""
if limit is not None:
try:
limit_value = int(limit)
except (TypeError, ValueError) as exc: # pragma: no cover - defensive guard
raise ValueError("limit must be an integer") from exc
if limit_value < 0:
raise ValueError("limit must be non-negative")
limit_clause = f" LIMIT ${placeholder_index}"
query_params.append(limit_value)
query = f"""
SELECT
dus.id,
dus.date,
dus.user_id,
dus.api_key,
dus.model,
dus.model_group,
dus.custom_llm_provider,
dus.prompt_tokens,
dus.completion_tokens,
dus.spend,
dus.api_requests,
dus.successful_requests,
dus.failed_requests,
dus.cache_creation_input_tokens,
dus.cache_read_input_tokens,
dus.created_at,
dus.updated_at,
vt.team_id,
vt.key_alias as api_key_alias,
tt.team_alias,
ut.user_email as user_email
FROM "LiteLLM_DailyUserSpend" dus
LEFT JOIN "LiteLLM_VerificationToken" vt ON dus.api_key = vt.token
LEFT JOIN "LiteLLM_TeamTable" tt ON vt.team_id = tt.team_id
LEFT JOIN "LiteLLM_UserTable" ut ON dus.user_id = ut.user_id
{where_clause}
ORDER BY dus.date DESC, dus.created_at DESC
{limit_clause}
"""
try:
db_response = await client.db.query_raw(query, *query_params)
return pl.DataFrame(db_response, infer_schema_length=None)
except Exception as exc:
raise RuntimeError(f"Error retrieving usage data: {exc}") from exc
async def get_table_info(self) -> Dict[str, Any]:
"""Return metadata about the spend table for diagnostics."""
client = self._ensure_prisma_client()
info_query = """
SELECT column_name, data_type, is_nullable
FROM information_schema.columns
WHERE table_name = 'LiteLLM_DailyUserSpend'
ORDER BY ordinal_position;
"""
try:
columns_response = await client.db.query_raw(info_query)
return {"columns": columns_response, "table_name": "LiteLLM_DailyUserSpend"}
except Exception as exc:
raise RuntimeError(f"Error getting table info: {exc}") from exc

12
litellm/integrations/focus/destinations/__init__.py Normal file
View file

@ -0,0 +1,12 @@
"""Destination implementations for Focus export."""
from .base import FocusDestination, FocusTimeWindow
from .factory import FocusDestinationFactory
from .s3_destination import FocusS3Destination
__all__ = [
"FocusDestination",
"FocusDestinationFactory",
"FocusTimeWindow",
"FocusS3Destination",
]

30
litellm/integrations/focus/destinations/base.py Normal file
View file

@ -0,0 +1,30 @@
"""Abstract destination interfaces for Focus export."""
from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime
from typing import Protocol
@dataclass(frozen=True)
class FocusTimeWindow:
"""Represents the span of data exported in a single batch."""
start_time: datetime
end_time: datetime
frequency: str
class FocusDestination(Protocol):
"""Protocol for anything that can receive Focus export files."""
async def deliver(
self,
*,
content: bytes,
time_window: FocusTimeWindow,
filename: str,
) -> None:
"""Persist the serialized export for the provided time window."""
...

59
litellm/integrations/focus/destinations/factory.py Normal file
View file

@ -0,0 +1,59 @@
"""Factory helpers for Focus export destinations."""
from __future__ import annotations
import os
from typing import Any, Dict, Optional
from .base import FocusDestination
from .s3_destination import FocusS3Destination
class FocusDestinationFactory:
"""Builds destination instances based on provider/config settings."""
@staticmethod
def create(
*,
provider: str,
prefix: str,
config: Optional[Dict[str, Any]] = None,
) -> FocusDestination:
"""Return a destination implementation for the requested provider."""
provider_lower = provider.lower()
normalized_config = FocusDestinationFactory._resolve_config(
provider=provider_lower, overrides=config or {}
)
if provider_lower == "s3":
return FocusS3Destination(prefix=prefix, config=normalized_config)
raise NotImplementedError(
f"Provider '{provider}' not supported for Focus export"
)
@staticmethod
def _resolve_config(
*,
provider: str,
overrides: Dict[str, Any],
) -> Dict[str, Any]:
if provider == "s3":
resolved = {
"bucket_name": overrides.get("bucket_name")
or os.getenv("FOCUS_S3_BUCKET_NAME"),
"region_name": overrides.get("region_name")
or os.getenv("FOCUS_S3_REGION_NAME"),
"endpoint_url": overrides.get("endpoint_url")
or os.getenv("FOCUS_S3_ENDPOINT_URL"),
"aws_access_key_id": overrides.get("aws_access_key_id")
or os.getenv("FOCUS_S3_ACCESS_KEY"),
"aws_secret_access_key": overrides.get("aws_secret_access_key")
or os.getenv("FOCUS_S3_SECRET_KEY"),
"aws_session_token": overrides.get("aws_session_token")
or os.getenv("FOCUS_S3_SESSION_TOKEN"),
}
if not resolved.get("bucket_name"):
raise ValueError("FOCUS_S3_BUCKET_NAME must be provided for S3 exports")
return {k: v for k, v in resolved.items() if v is not None}
raise NotImplementedError(
f"Provider '{provider}' not supported for Focus export configuration"
)

74
litellm/integrations/focus/destinations/s3_destination.py Normal file
View file

@ -0,0 +1,74 @@
"""S3 destination implementation for Focus export."""
from __future__ import annotations
import asyncio
from datetime import timezone
from typing import Any, Optional
import boto3
from .base import FocusDestination, FocusTimeWindow
class FocusS3Destination(FocusDestination):
"""Handles uploading serialized exports to S3 buckets."""
def __init__(
self,
*,
prefix: str,
config: Optional[dict[str, Any]] = None,
) -> None:
config = config or {}
bucket_name = config.get("bucket_name")
if not bucket_name:
raise ValueError("bucket_name must be provided for S3 destination")
self.bucket_name = bucket_name
self.prefix = prefix.rstrip("/")
self.config = config
async def deliver(
self,
*,
content: bytes,
time_window: FocusTimeWindow,
filename: str,
) -> None:
object_key = self._build_object_key(time_window=time_window, filename=filename)
await asyncio.to_thread(self._upload, content, object_key)
def _build_object_key(self, *, time_window: FocusTimeWindow, filename: str) -> str:
start_utc = time_window.start_time.astimezone(timezone.utc)
date_component = f"date={start_utc.strftime('%Y-%m-%d')}"
parts = [self.prefix, date_component]
if time_window.frequency == "hourly":
parts.append(f"hour={start_utc.strftime('%H')}")
key_prefix = "/".join(filter(None, parts))
return f"{key_prefix}/{filename}" if key_prefix else filename
def _upload(self, content: bytes, object_key: str) -> None:
client_kwargs: dict[str, Any] = {}
region_name = self.config.get("region_name")
if region_name:
client_kwargs["region_name"] = region_name
endpoint_url = self.config.get("endpoint_url")
if endpoint_url:
client_kwargs["endpoint_url"] = endpoint_url
session_kwargs: dict[str, Any] = {}
for key in (
"aws_access_key_id",
"aws_secret_access_key",
"aws_session_token",
):
if self.config.get(key):
session_kwargs[key] = self.config[key]
s3_client = boto3.client("s3", **client_kwargs, **session_kwargs)
s3_client.put_object(
Bucket=self.bucket_name,
Key=object_key,
Body=content,
ContentType="application/octet-stream",
)

124
litellm/integrations/focus/export_engine.py Normal file
View file

@ -0,0 +1,124 @@
"""Core export engine for Focus integrations (heavy dependencies)."""
from __future__ import annotations
from typing import Any, Dict, Optional
import polars as pl
from litellm._logging import verbose_logger
from .database import FocusLiteLLMDatabase
from .destinations import FocusDestinationFactory, FocusTimeWindow
from .serializers import FocusParquetSerializer, FocusSerializer
from .transformer import FocusTransformer
class FocusExportEngine:
"""Engine that fetches, normalizes, and uploads Focus exports."""
def __init__(
self,
*,
provider: str,
export_format: str,
prefix: str,
destination_config: Optional[dict[str, Any]] = None,
) -> None:
self.provider = provider
self.export_format = export_format
self.prefix = prefix
self._destination = FocusDestinationFactory.create(
provider=self.provider,
prefix=self.prefix,
config=destination_config,
)
self._serializer = self._init_serializer()
self._transformer = FocusTransformer()
self._database = FocusLiteLLMDatabase()
def _init_serializer(self) -> FocusSerializer:
if self.export_format != "parquet":
raise NotImplementedError("Only parquet export supported currently")
return FocusParquetSerializer()
async def dry_run_export_usage_data(self, limit: Optional[int]) -> Dict[str, Any]:
data = await self._database.get_usage_data(limit=limit)
normalized = self._transformer.transform(data)
usage_sample = data.head(min(50, len(data))).to_dicts()
normalized_sample = normalized.head(min(50, len(normalized))).to_dicts()
summary = {
"total_records": len(normalized),
"total_spend": self._sum_column(normalized, "spend"),
"total_tokens": self._sum_column(normalized, "total_tokens"),
"unique_teams": self._count_unique(normalized, "team_id"),
"unique_models": self._count_unique(normalized, "model"),
}
return {
"usage_data": usage_sample,
"normalized_data": normalized_sample,
"summary": summary,
}
async def export_window(
self,
*,
window: FocusTimeWindow,
limit: Optional[int],
) -> None:
data = await self._database.get_usage_data(
limit=limit,
start_time_utc=window.start_time,
end_time_utc=window.end_time,
)
if data.is_empty():
verbose_logger.debug("Focus export: no usage data for window %s", window)
return
normalized = self._transformer.transform(data)
if normalized.is_empty():
verbose_logger.debug(
"Focus export: normalized data empty for window %s", window
)
return
await self._serialize_and_upload(normalized, window)
async def _serialize_and_upload(
self, frame: pl.DataFrame, window: FocusTimeWindow
) -> None:
payload = self._serializer.serialize(frame)
if not payload:
verbose_logger.debug("Focus export: serializer returned empty payload")
return
await self._destination.deliver(
content=payload,
time_window=window,
filename=self._build_filename(),
)
def _build_filename(self) -> str:
if not self._serializer.extension:
raise ValueError("Serializer must declare a file extension")
return f"usage.{self._serializer.extension}"
@staticmethod
def _sum_column(frame: pl.DataFrame, column: str) -> float:
if frame.is_empty() or column not in frame.columns:
return 0.0
value = frame.select(pl.col(column).sum().alias("sum")).row(0)[0]
if value is None:
return 0.0
return float(value)
@staticmethod
def _count_unique(frame: pl.DataFrame, column: str) -> int:
if frame.is_empty() or column not in frame.columns:
return 0
value = frame.select(pl.col(column).n_unique().alias("unique")).row(0)[0]
if value is None:
return 0
return int(value)

211
litellm/integrations/focus/focus_logger.py Normal file
View file

@ -0,0 +1,211 @@
"""Focus export logger orchestrating DB pull/transform/upload."""
from __future__ import annotations
import os
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
import litellm
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from .destinations import FocusTimeWindow
if TYPE_CHECKING:
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from .export_engine import FocusExportEngine
else:
AsyncIOScheduler = Any
FOCUS_USAGE_DATA_JOB_NAME = "focus_export_usage_data"
DEFAULT_DRY_RUN_LIMIT = 500
class FocusLogger(CustomLogger):
"""Coordinates Focus export jobs across transformer/serializer/destination layers."""
def __init__(
self,
*,
provider: Optional[str] = None,
export_format: Optional[str] = None,
frequency: Optional[str] = None,
cron_offset_minute: Optional[int] = None,
interval_seconds: Optional[int] = None,
prefix: Optional[str] = None,
destination_config: Optional[dict[str, Any]] = None,
**kwargs: Any,
) -> None:
super().__init__(**kwargs)
self.provider = (provider or os.getenv("FOCUS_PROVIDER") or "s3").lower()
self.export_format = (
export_format or os.getenv("FOCUS_FORMAT") or "parquet"
).lower()
self.frequency = (frequency or os.getenv("FOCUS_FREQUENCY") or "hourly").lower()
self.cron_offset_minute = (
cron_offset_minute
if cron_offset_minute is not None
else int(os.getenv("FOCUS_CRON_OFFSET", "5"))
)
raw_interval = (
interval_seconds
if interval_seconds is not None
else os.getenv("FOCUS_INTERVAL_SECONDS")
)
self.interval_seconds = int(raw_interval) if raw_interval is not None else None
env_prefix = os.getenv("FOCUS_PREFIX")
self.prefix: str = (
prefix if prefix is not None else (env_prefix if env_prefix else "focus_exports")
)
self._destination_config = destination_config
self._engine: Optional["FocusExportEngine"] = None
def _ensure_engine(self) -> "FocusExportEngine":
"""Instantiate the heavy export engine lazily."""
if self._engine is None:
from .export_engine import FocusExportEngine
self._engine = FocusExportEngine(
provider=self.provider,
export_format=self.export_format,
prefix=self.prefix,
destination_config=self._destination_config,
)
return self._engine
async def export_usage_data(
self,
*,
limit: Optional[int] = None,
start_time_utc: Optional[datetime] = None,
end_time_utc: Optional[datetime] = None,
) -> None:
"""Public hook to trigger export immediately."""
if bool(start_time_utc) ^ bool(end_time_utc):
raise ValueError(
"start_time_utc and end_time_utc must be provided together"
)
if start_time_utc and end_time_utc:
window = FocusTimeWindow(
start_time=start_time_utc,
end_time=end_time_utc,
frequency=self.frequency,
)
else:
window = self._compute_time_window(datetime.now(timezone.utc))
await self._export_window(window=window, limit=limit)
async def dry_run_export_usage_data(
self, limit: Optional[int] = DEFAULT_DRY_RUN_LIMIT
) -> dict[str, Any]:
"""Return transformed data without uploading."""
engine = self._ensure_engine()
return await engine.dry_run_export_usage_data(limit=limit)
async def initialize_focus_export_job(self) -> None:
"""Entry point for scheduler jobs to run export cycle with locking."""
from litellm.proxy.proxy_server import proxy_logging_obj
pod_lock_manager = None
if proxy_logging_obj is not None:
writer = getattr(proxy_logging_obj, "db_spend_update_writer", None)
if writer is not None:
pod_lock_manager = getattr(writer, "pod_lock_manager", None)
if pod_lock_manager and pod_lock_manager.redis_cache:
acquired = await pod_lock_manager.acquire_lock(
cronjob_id=FOCUS_USAGE_DATA_JOB_NAME
)
if not acquired:
verbose_logger.debug("Focus export: unable to acquire pod lock")
return
try:
await self._run_scheduled_export()
finally:
await pod_lock_manager.release_lock(
cronjob_id=FOCUS_USAGE_DATA_JOB_NAME
)
else:
await self._run_scheduled_export()
@staticmethod
async def init_focus_export_background_job(
scheduler: AsyncIOScheduler,
) -> None:
"""Register the export cron/interval job with the provided scheduler."""
focus_loggers: List[
CustomLogger
] = litellm.logging_callback_manager.get_custom_loggers_for_type(
callback_type=FocusLogger
)
if not focus_loggers:
verbose_logger.debug(
"No Focus export logger registered; skipping scheduler"
)
return
focus_logger = cast(FocusLogger, focus_loggers[0])
trigger_kwargs = focus_logger._build_scheduler_trigger()
scheduler.add_job(
focus_logger.initialize_focus_export_job,
**trigger_kwargs,
)
def _build_scheduler_trigger(self) -> Dict[str, Any]:
"""Return scheduler configuration for the selected frequency."""
if self.frequency == "interval":
seconds = self.interval_seconds or 60
return {"trigger": "interval", "seconds": seconds}
if self.frequency == "hourly":
minute = max(0, min(59, self.cron_offset_minute))
return {"trigger": "cron", "minute": minute, "second": 0}
if self.frequency == "daily":
total_minutes = max(0, self.cron_offset_minute)
hour = min(23, total_minutes // 60)
minute = min(59, total_minutes % 60)
return {"trigger": "cron", "hour": hour, "minute": minute, "second": 0}
raise ValueError(f"Unsupported frequency: {self.frequency}")
async def _run_scheduled_export(self) -> None:
"""Execute the scheduled export for the configured window."""
window = self._compute_time_window(datetime.now(timezone.utc))
await self._export_window(window=window, limit=None)
async def _export_window(
self,
*,
window: FocusTimeWindow,
limit: Optional[int],
) -> None:
engine = self._ensure_engine()
await engine.export_window(window=window, limit=limit)
def _compute_time_window(self, now: datetime) -> FocusTimeWindow:
"""Derive the time window to export based on configured frequency."""
now_utc = now.astimezone(timezone.utc)
if self.frequency == "hourly":
end_time = now_utc.replace(minute=0, second=0, microsecond=0)
start_time = end_time - timedelta(hours=1)
elif self.frequency == "daily":
end_time = now_utc.replace(hour=0, minute=0, second=0, microsecond=0)
start_time = end_time - timedelta(days=1)
elif self.frequency == "interval":
interval = timedelta(seconds=self.interval_seconds or 60)
end_time = now_utc
start_time = end_time - interval
else:
raise ValueError(f"Unsupported frequency: {self.frequency}")
return FocusTimeWindow(
start_time=start_time,
end_time=end_time,
frequency=self.frequency,
)
__all__ = ["FocusLogger"]

50
litellm/integrations/focus/schema.py Normal file
View file

@ -0,0 +1,50 @@
"""Schema definitions for Focus export data."""
from __future__ import annotations
import polars as pl
# see: https://focus.finops.org/focus-specification/v1-2/
FOCUS_NORMALIZED_SCHEMA = pl.Schema(
[
("BilledCost", pl.Decimal(18, 6)),
("BillingAccountId", pl.String),
("BillingAccountName", pl.String),
("BillingCurrency", pl.String),
("BillingPeriodStart", pl.Datetime(time_unit="us")),
("BillingPeriodEnd", pl.Datetime(time_unit="us")),
("ChargeCategory", pl.String),
("ChargeClass", pl.String),
("ChargeDescription", pl.String),
("ChargeFrequency", pl.String),
("ChargePeriodStart", pl.Datetime(time_unit="us")),
("ChargePeriodEnd", pl.Datetime(time_unit="us")),
("ConsumedQuantity", pl.Decimal(18, 6)),
("ConsumedUnit", pl.String),
("ContractedCost", pl.Decimal(18, 6)),
("ContractedUnitPrice", pl.Decimal(18, 6)),
("EffectiveCost", pl.Decimal(18, 6)),
("InvoiceIssuerName", pl.String),
("ListCost", pl.Decimal(18, 6)),
("ListUnitPrice", pl.Decimal(18, 6)),
("PricingCategory", pl.String),
("PricingQuantity", pl.Decimal(18, 6)),
("PricingUnit", pl.String),
("ProviderName", pl.String),
("PublisherName", pl.String),
("RegionId", pl.String),
("RegionName", pl.String),
("ResourceId", pl.String),
("ResourceName", pl.String),
("ResourceType", pl.String),
("ServiceCategory", pl.String),
("ServiceSubcategory", pl.String),
("ServiceName", pl.String),
("SubAccountId", pl.String),
("SubAccountName", pl.String),
("SubAccountType", pl.String),
("Tags", pl.Object),
]
)
__all__ = ["FOCUS_NORMALIZED_SCHEMA"]

6
litellm/integrations/focus/serializers/__init__.py Normal file
View file

@ -0,0 +1,6 @@
"""Serializer package exports for Focus integration."""
from .base import FocusSerializer
from .parquet import FocusParquetSerializer
__all__ = ["FocusSerializer", "FocusParquetSerializer"]

18
litellm/integrations/focus/serializers/base.py Normal file
View file

@ -0,0 +1,18 @@
"""Serializer abstractions for Focus export."""
from __future__ import annotations
from abc import ABC, abstractmethod
import polars as pl
class FocusSerializer(ABC):
"""Base serializer turning Focus frames into bytes."""
extension: str = ""
@abstractmethod
def serialize(self, frame: pl.DataFrame) -> bytes:
"""Convert the normalized Focus frame into the chosen format."""
raise NotImplementedError

22
litellm/integrations/focus/serializers/parquet.py Normal file
View file

@ -0,0 +1,22 @@
"""Parquet serializer for Focus export."""
from __future__ import annotations
import io
import polars as pl
from .base import FocusSerializer
class FocusParquetSerializer(FocusSerializer):
"""Serialize normalized Focus frames to Parquet bytes."""
extension = "parquet"
def serialize(self, frame: pl.DataFrame) -> bytes:
"""Encode the provided frame as a parquet payload."""
target = frame if not frame.is_empty() else pl.DataFrame(schema=frame.schema)
buffer = io.BytesIO()
target.write_parquet(buffer, compression="snappy")
return buffer.getvalue()

90
litellm/integrations/focus/transformer.py Normal file
View file

@ -0,0 +1,90 @@
"""Focus export data transformer."""
from __future__ import annotations
from datetime import timedelta
import polars as pl
from .schema import FOCUS_NORMALIZED_SCHEMA
class FocusTransformer:
"""Transforms LiteLLM DB rows into Focus-compatible schema."""
schema = FOCUS_NORMALIZED_SCHEMA
def transform(self, frame: pl.DataFrame) -> pl.DataFrame:
"""Return a normalized frame expected by downstream serializers."""
if frame.is_empty():
return pl.DataFrame(schema=self.schema)
# derive period start/end from usage date
frame = frame.with_columns(
pl.col("date")
.cast(pl.Utf8)
.str.strptime(pl.Datetime(time_unit="us"), format="%Y-%m-%d", strict=False)
.alias("usage_date"),
)
frame = frame.with_columns(
pl.col("usage_date").alias("ChargePeriodStart"),
(pl.col("usage_date") + timedelta(days=1)).alias("ChargePeriodEnd"),
)
def fmt(col):
return col.dt.strftime("%Y-%m-%dT%H:%M:%SZ")
DEC = pl.Decimal(18, 6)
def dec(col):
return col.cast(DEC)
none_str = pl.lit(None, dtype=pl.Utf8)
none_dec = pl.lit(None, dtype=pl.Decimal(18, 6))
return frame.select(
dec(pl.col("spend").fill_null(0.0)).alias("BilledCost"),
pl.col("api_key").cast(pl.String).alias("BillingAccountId"),
pl.col("api_key_alias").cast(pl.String).alias("BillingAccountName"),
pl.lit("API Key").alias("BillingAccountType"),
pl.lit("USD").alias("BillingCurrency"),
fmt(pl.col("ChargePeriodEnd")).alias("BillingPeriodEnd"),
fmt(pl.col("ChargePeriodStart")).alias("BillingPeriodStart"),
pl.lit("Usage").alias("ChargeCategory"),
none_str.alias("ChargeClass"),
pl.col("model").cast(pl.String).alias("ChargeDescription"),
pl.lit("Usage-Based").alias("ChargeFrequency"),
fmt(pl.col("ChargePeriodEnd")).alias("ChargePeriodEnd"),
fmt(pl.col("ChargePeriodStart")).alias("ChargePeriodStart"),
dec(pl.lit(1.0)).alias("ConsumedQuantity"),
pl.lit("Requests").alias("ConsumedUnit"),
dec(pl.col("spend").fill_null(0.0)).alias("ContractedCost"),
none_str.alias("ContractedUnitPrice"),
dec(pl.col("spend").fill_null(0.0)).alias("EffectiveCost"),
pl.col("custom_llm_provider").cast(pl.String).alias("InvoiceIssuerName"),
none_str.alias("InvoiceId"),
dec(pl.col("spend").fill_null(0.0)).alias("ListCost"),
none_dec.alias("ListUnitPrice"),
none_str.alias("AvailabilityZone"),
pl.lit("USD").alias("PricingCurrency"),
none_str.alias("PricingCategory"),
dec(pl.lit(1.0)).alias("PricingQuantity"),
none_dec.alias("PricingCurrencyContractedUnitPrice"),
dec(pl.col("spend").fill_null(0.0)).alias("PricingCurrencyEffectiveCost"),
none_dec.alias("PricingCurrencyListUnitPrice"),
pl.lit("Requests").alias("PricingUnit"),
pl.col("custom_llm_provider").cast(pl.String).alias("ProviderName"),
pl.col("custom_llm_provider").cast(pl.String).alias("PublisherName"),
none_str.alias("RegionId"),
none_str.alias("RegionName"),
pl.col("model").cast(pl.String).alias("ResourceId"),
pl.col("model").cast(pl.String).alias("ResourceName"),
pl.col("model").cast(pl.String).alias("ResourceType"),
pl.lit("AI and Machine Learning").alias("ServiceCategory"),
pl.lit("Generative AI").alias("ServiceSubcategory"),
pl.col("model_group").cast(pl.String).alias("ServiceName"),
pl.col("team_id").cast(pl.String).alias("SubAccountId"),
pl.col("team_alias").cast(pl.String).alias("SubAccountName"),
none_str.alias("SubAccountType"),
none_str.alias("Tags"),
)

59
litellm/integrations/generic_api/generic_api_compatible_callbacks.json
View file

@ -1,28 +1,37 @@
{
"sample_callback": {
"event_types": ["llm_api_success", "llm_api_failure"],
"endpoint": "{{environment_variables.SAMPLE_CALLBACK_URL}}",
"headers": {
"Content-Type": "application/json",
"Authorization": "Bearer {{environment_variables.SAMPLE_CALLBACK_API_KEY}}"
},
"environment_variables": ["SAMPLE_CALLBACK_URL", "SAMPLE_CALLBACK_API_KEY"]
"sample_callback": {
"event_types": ["llm_api_success", "llm_api_failure"],
"endpoint": "{{environment_variables.SAMPLE_CALLBACK_URL}}",
"headers": {
"Content-Type": "application/json",
"Authorization": "Bearer {{environment_variables.SAMPLE_CALLBACK_API_KEY}}"
},
"rubrik": {
"event_types": ["llm_api_success"],
"endpoint": "{{environment_variables.RUBRIK_WEBHOOK_URL}}",
"headers": {
"Content-Type": "application/json",
"Authorization": "Bearer {{environment_variables.RUBRIK_API_KEY}}"
},
"environment_variables": ["RUBRIK_API_KEY", "RUBRIK_WEBHOOK_URL"]
"environment_variables": ["SAMPLE_CALLBACK_URL", "SAMPLE_CALLBACK_API_KEY"]
},
"rubrik": {
"event_types": ["llm_api_success"],
"endpoint": "{{environment_variables.RUBRIK_WEBHOOK_URL}}",
"headers": {
"Content-Type": "application/json",
"Authorization": "Bearer {{environment_variables.RUBRIK_API_KEY}}"
},
"sumologic": {
"endpoint": "{{environment_variables.SUMOLOGIC_WEBHOOK_URL}}",
"headers": {
"Content-Type": "application/json"
},
"environment_variables": ["SUMOLOGIC_WEBHOOK_URL"],
"log_format": "ndjson"
}
}
"environment_variables": ["RUBRIK_API_KEY", "RUBRIK_WEBHOOK_URL"]
},
"sumologic": {
"endpoint": "{{environment_variables.SUMOLOGIC_WEBHOOK_URL}}",
"headers": {
"Content-Type": "application/json"
},
"environment_variables": ["SUMOLOGIC_WEBHOOK_URL"],
"log_format": "ndjson"
},
"qualifire_eval": {
"event_types": ["llm_api_success"],
"endpoint": "{{environment_variables.QUALIFIRE_WEBHOOK_URL}}",
"headers": {
"Content-Type": "application/json",
"X-Qualifire-API-Key": "{{environment_variables.QUALIFIRE_API_KEY}}"
},
"environment_variables": ["QUALIFIRE_API_KEY", "QUALIFIRE_WEBHOOK_URL"]
}
}

378
litellm/integrations/prometheus.py
View file

@ -14,6 +14,7 @@ from typing import (
Literal,
Optional,
Tuple,
Union,
cast,
)
@ -44,6 +45,7 @@ def _get_cached_end_user_id_for_cost_tracking():
global _get_end_user_id_for_cost_tracking
if _get_end_user_id_for_cost_tracking is None:
from litellm.utils import get_end_user_id_for_cost_tracking
_get_end_user_id_for_cost_tracking = get_end_user_id_for_cost_tracking
return _get_end_user_id_for_cost_tracking
@ -237,6 +239,36 @@ class PrometheusLogger(CustomLogger):
),
buckets=LATENCY_BUCKETS,
)
# Request queue time metric
self.litellm_request_queue_time_metric = self._histogram_factory(
"litellm_request_queue_time_seconds",
"Time spent in request queue before processing starts (seconds)",
labelnames=self.get_labels_for_metric(
"litellm_request_queue_time_seconds"
),
buckets=LATENCY_BUCKETS,
)
# Guardrail metrics
self.litellm_guardrail_latency_metric = self._histogram_factory(
"litellm_guardrail_latency_seconds",
"Latency (seconds) for guardrail execution",
labelnames=["guardrail_name", "status", "error_type", "hook_type"],
buckets=LATENCY_BUCKETS,
)
self.litellm_guardrail_errors_total = self._counter_factory(
"litellm_guardrail_errors_total",
"Total number of errors encountered during guardrail execution",
labelnames=["guardrail_name", "error_type", "hook_type"],
)
self.litellm_guardrail_requests_total = self._counter_factory(
"litellm_guardrail_requests_total",
"Total number of guardrail invocations",
labelnames=["guardrail_name", "status", "hook_type"],
)
# llm api provider budget metrics
self.litellm_provider_remaining_budget_metric = self._gauge_factory(
"litellm_provider_remaining_budget_metric",
@ -329,6 +361,25 @@ class PrometheusLogger(CustomLogger):
labelnames=self.get_labels_for_metric("litellm_requests_metric"),
)
# Cache metrics
self.litellm_cache_hits_metric = self._counter_factory(
name="litellm_cache_hits_metric",
documentation="Total number of LiteLLM cache hits",
labelnames=self.get_labels_for_metric("litellm_cache_hits_metric"),
)
self.litellm_cache_misses_metric = self._counter_factory(
name="litellm_cache_misses_metric",
documentation="Total number of LiteLLM cache misses",
labelnames=self.get_labels_for_metric("litellm_cache_misses_metric"),
)
self.litellm_cached_tokens_metric = self._counter_factory(
name="litellm_cached_tokens_metric",
documentation="Total tokens served from LiteLLM cache",
labelnames=self.get_labels_for_metric("litellm_cached_tokens_metric"),
)
except Exception as e:
print_verbose(f"Got exception on init prometheus client {str(e)}")
raise e
@ -791,11 +842,16 @@ class PrometheusLogger(CustomLogger):
f"standard_logging_object is required, got={standard_logging_payload}"
)
if self._should_skip_metrics_for_invalid_key(
kwargs=kwargs, standard_logging_payload=standard_logging_payload
):
return
model = kwargs.get("model", "")
litellm_params = kwargs.get("litellm_params", {}) or {}
_metadata = litellm_params.get("metadata", {})
get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking()
end_user_id = get_end_user_id_for_cost_tracking(
litellm_params, service_type="prometheus"
)
@ -815,7 +871,7 @@ class PrometheusLogger(CustomLogger):
user_api_key_auth_metadata: Optional[dict] = standard_logging_payload[
"metadata"
].get("user_api_key_auth_metadata")
# Include top-level metadata fields (excluding nested dictionaries)
# This allows accessing fields like requester_ip_address from top-level metadata
top_level_metadata = standard_logging_payload.get("metadata", {})
@ -826,7 +882,7 @@ class PrometheusLogger(CustomLogger):
for k, v in top_level_metadata.items()
if not isinstance(v, dict) # Exclude nested dicts to avoid conflicts
}
combined_metadata: Dict[str, Any] = {
**top_level_fields, # Include top-level fields first
**(_requester_metadata if _requester_metadata else {}),
@ -945,6 +1001,12 @@ class PrometheusLogger(CustomLogger):
kwargs, start_time, end_time, enum_values, output_tokens
)
# cache metrics
self._increment_cache_metrics(
standard_logging_payload=standard_logging_payload, # type: ignore
enum_values=enum_values,
)
if (
standard_logging_payload["stream"] is True
): # log successful streaming requests from logging event hook.
@ -1014,6 +1076,54 @@ class PrometheusLogger(CustomLogger):
standard_logging_payload["completion_tokens"]
)
def _increment_cache_metrics(
self,
standard_logging_payload: StandardLoggingPayload,
enum_values: UserAPIKeyLabelValues,
):
"""
Increment cache-related Prometheus metrics based on cache hit/miss status.
Args:
standard_logging_payload: Contains cache_hit field (True/False/None)
enum_values: Label values for Prometheus metrics
"""
cache_hit = standard_logging_payload.get("cache_hit")
# Only track if cache_hit has a definite value (True or False)
if cache_hit is None:
return
if cache_hit is True:
# Increment cache hits counter
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name="litellm_cache_hits_metric"
),
enum_values=enum_values,
)
self.litellm_cache_hits_metric.labels(**_labels).inc()
# Increment cached tokens counter
total_tokens = standard_logging_payload.get("total_tokens", 0)
if total_tokens > 0:
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name="litellm_cached_tokens_metric"
),
enum_values=enum_values,
)
self.litellm_cached_tokens_metric.labels(**_labels).inc(total_tokens)
else:
# cache_hit is False - increment cache misses counter
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name="litellm_cache_misses_metric"
),
enum_values=enum_values,
)
self.litellm_cache_misses_metric.labels(**_labels).inc()
async def _increment_remaining_budget_metrics(
self,
user_api_team: Optional[str],
@ -1182,6 +1292,22 @@ class PrometheusLogger(CustomLogger):
total_time_seconds
)
# request queue time (time from arrival to processing start)
_litellm_params = kwargs.get("litellm_params", {}) or {}
queue_time_seconds = _litellm_params.get("metadata", {}).get(
"queue_time_seconds"
)
if queue_time_seconds is not None and queue_time_seconds >= 0:
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name="litellm_request_queue_time_seconds"
),
enum_values=enum_values,
)
self.litellm_request_queue_time_metric.labels(**_labels).observe(
queue_time_seconds
)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
from litellm.types.utils import StandardLoggingPayload
@ -1189,14 +1315,20 @@ class PrometheusLogger(CustomLogger):
f"prometheus Logging - Enters failure logging function for kwargs {kwargs}"
)
# unpack kwargs
model = kwargs.get("model", "")
standard_logging_payload: StandardLoggingPayload = kwargs.get(
"standard_logging_object", {}
)
if self._should_skip_metrics_for_invalid_key(
kwargs=kwargs, standard_logging_payload=standard_logging_payload
):
return
model = kwargs.get("model", "")
litellm_params = kwargs.get("litellm_params", {}) or {}
get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking()
end_user_id = get_end_user_id_for_cost_tracking(
litellm_params, service_type="prometheus"
)
@ -1207,7 +1339,6 @@ class PrometheusLogger(CustomLogger):
user_api_team_alias = standard_logging_payload["metadata"][
"user_api_key_team_alias"
]
kwargs.get("exception", None)
try:
self.litellm_llm_api_failed_requests_metric.labels(
@ -1227,6 +1358,139 @@ class PrometheusLogger(CustomLogger):
pass
pass
def _extract_status_code(
self,
kwargs: Optional[dict] = None,
enum_values: Optional[Any] = None,
exception: Optional[Exception] = None,
) -> Optional[int]:
"""
Extract HTTP status code from various input formats for validation.
This is a centralized helper to extract status code from different
callback function signatures. Handles both ProxyException (uses 'code')
and standard exceptions (uses 'status_code').
Args:
kwargs: Dictionary potentially containing 'exception' key
enum_values: Object with 'status_code' attribute
exception: Exception object to extract status code from directly
Returns:
Status code as integer if found, None otherwise
"""
status_code = None
# Try from enum_values first (most common in our callbacks)
if enum_values and hasattr(enum_values, "status_code") and enum_values.status_code:
try:
status_code = int(enum_values.status_code)
except (ValueError, TypeError):
pass
if not status_code and exception:
# ProxyException uses 'code' attribute, other exceptions may use 'status_code'
status_code = getattr(exception, "status_code", None) or getattr(exception, "code", None)
if status_code is not None:
try:
status_code = int(status_code)
except (ValueError, TypeError):
status_code = None
if not status_code and kwargs:
exception_in_kwargs = kwargs.get("exception")
if exception_in_kwargs:
status_code = getattr(exception_in_kwargs, "status_code", None) or getattr(exception_in_kwargs, "code", None)
if status_code is not None:
try:
status_code = int(status_code)
except (ValueError, TypeError):
status_code = None
return status_code
def _is_invalid_api_key_request(
self,
status_code: Optional[int],
exception: Optional[Exception] = None,
) -> bool:
"""
Determine if a request has an invalid API key based on status code and exception.
This method prevents invalid authentication attempts from being recorded in
Prometheus metrics. A 401 status code is the definitive indicator of authentication
failure. Additionally, we check exception messages for authentication error patterns
to catch cases where the exception hasn't been converted to a ProxyException yet.
Args:
status_code: HTTP status code (401 indicates authentication error)
exception: Exception object to check for auth-related error messages
Returns:
True if the request has an invalid API key and metrics should be skipped,
False otherwise
"""
if status_code == 401:
return True
# Handle cases where AssertionError is raised before conversion to ProxyException
if exception is not None:
exception_str = str(exception).lower()
auth_error_patterns = [
"virtual key expected",
"expected to start with 'sk-'",
"authentication error",
"invalid api key",
"api key not valid",
]
if any(pattern in exception_str for pattern in auth_error_patterns):
return True
return False
def _should_skip_metrics_for_invalid_key(
self,
kwargs: Optional[dict] = None,
user_api_key_dict: Optional[Any] = None,
enum_values: Optional[Any] = None,
standard_logging_payload: Optional[Union[dict, StandardLoggingPayload]] = None,
exception: Optional[Exception] = None,
) -> bool:
"""
Determine if Prometheus metrics should be skipped for invalid API key requests.
This is a centralized validation method that extracts status code and exception
information from various callback function signatures and determines if the request
represents an invalid API key attempt that should be filtered from metrics.
Args:
kwargs: Dictionary potentially containing exception and other data
user_api_key_dict: User API key authentication object (currently unused)
enum_values: Object with status_code attribute
standard_logging_payload: Standard logging payload dictionary
exception: Exception object to check directly
Returns:
True if metrics should be skipped (invalid key detected), False otherwise
"""
status_code = self._extract_status_code(
kwargs=kwargs,
enum_values=enum_values,
exception=exception,
)
if exception is None and kwargs:
exception = kwargs.get("exception")
if self._is_invalid_api_key_request(status_code, exception=exception):
verbose_logger.debug(
"Skipping Prometheus metrics for invalid API key request: "
f"status_code={status_code}, exception={type(exception).__name__ if exception else None}"
)
return True
return False
async def async_post_call_failure_hook(
self,
request_data: dict,
@ -1252,6 +1516,14 @@ class PrometheusLogger(CustomLogger):
StandardLoggingPayloadSetup,
)
if self._should_skip_metrics_for_invalid_key(
user_api_key_dict=user_api_key_dict,
exception=original_exception,
):
return
status_code = self._extract_status_code(exception=original_exception)
try:
_tags = StandardLoggingPayloadSetup._get_request_tags(
litellm_params=request_data,
@ -1266,8 +1538,8 @@ class PrometheusLogger(CustomLogger):
team=user_api_key_dict.team_id,
team_alias=user_api_key_dict.team_alias,
requested_model=request_data.get("model", ""),
status_code=str(getattr(original_exception, "status_code", None)),
exception_status=str(getattr(original_exception, "status_code", None)),
status_code=str(status_code),
exception_status=str(status_code),
exception_class=self._get_exception_class_name(original_exception),
tags=_tags,
route=user_api_key_dict.request_route,
@ -1305,6 +1577,11 @@ class PrometheusLogger(CustomLogger):
StandardLoggingPayloadSetup,
)
if self._should_skip_metrics_for_invalid_key(
user_api_key_dict=user_api_key_dict
):
return
enum_values = UserAPIKeyLabelValues(
end_user=user_api_key_dict.end_user_id,
hashed_api_key=user_api_key_dict.api_key,
@ -1360,6 +1637,15 @@ class PrometheusLogger(CustomLogger):
exception = request_kwargs.get("exception", None)
llm_provider = _litellm_params.get("custom_llm_provider", None)
if self._should_skip_metrics_for_invalid_key(
kwargs=request_kwargs,
standard_logging_payload=standard_logging_payload,
):
return
hashed_api_key = standard_logging_payload.get("metadata", {}).get(
"user_api_key_hash"
)
# Create enum_values for the label factory (always create for use in different metrics)
enum_values = UserAPIKeyLabelValues(
@ -1374,9 +1660,7 @@ class PrometheusLogger(CustomLogger):
self._get_exception_class_name(exception) if exception else None
),
requested_model=model_group,
hashed_api_key=standard_logging_payload["metadata"][
"user_api_key_hash"
],
hashed_api_key=hashed_api_key,
api_key_alias=standard_logging_payload["metadata"][
"user_api_key_alias"
],
@ -1398,7 +1682,6 @@ class PrometheusLogger(CustomLogger):
api_provider=llm_provider or "",
)
if exception is not None:
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name="litellm_deployment_failure_responses"
@ -1431,16 +1714,23 @@ class PrometheusLogger(CustomLogger):
enum_values: UserAPIKeyLabelValues,
output_tokens: float = 1.0,
):
try:
verbose_logger.debug("setting remaining tokens requests metric")
standard_logging_payload: Optional[StandardLoggingPayload] = (
request_kwargs.get("standard_logging_object")
)
standard_logging_payload: Optional[
StandardLoggingPayload
] = request_kwargs.get("standard_logging_object")
if standard_logging_payload is None:
return
# Skip recording metrics for invalid API key requests
if self._should_skip_metrics_for_invalid_key(
kwargs=request_kwargs,
enum_values=enum_values,
standard_logging_payload=standard_logging_payload,
):
return
api_base = standard_logging_payload["api_base"]
_litellm_params = request_kwargs.get("litellm_params", {}) or {}
_metadata = _litellm_params.get("metadata", {})
@ -1571,6 +1861,50 @@ class PrometheusLogger(CustomLogger):
)
return
def _record_guardrail_metrics(
self,
guardrail_name: str,
latency_seconds: float,
status: str,
error_type: Optional[str],
hook_type: str,
):
"""
Record guardrail metrics for prometheus.
Args:
guardrail_name: Name of the guardrail
latency_seconds: Execution latency in seconds
status: "success" or "error"
error_type: Type of error if any, None otherwise
hook_type: "pre_call", "during_call", or "post_call"
"""
try:
# Record latency
self.litellm_guardrail_latency_metric.labels(
guardrail_name=guardrail_name,
status=status,
error_type=error_type or "none",
hook_type=hook_type,
).observe(latency_seconds)
# Record request count
self.litellm_guardrail_requests_total.labels(
guardrail_name=guardrail_name,
status=status,
hook_type=hook_type,
).inc()
# Record error count if there was an error
if status == "error" and error_type:
self.litellm_guardrail_errors_total.labels(
guardrail_name=guardrail_name,
error_type=error_type,
hook_type=hook_type,
).inc()
except Exception as e:
verbose_logger.debug(f"Error recording guardrail metrics: {str(e)}")
@staticmethod
def _get_exception_class_name(exception: Exception) -> str:
exception_class_name = ""
@ -2208,10 +2542,10 @@ class PrometheusLogger(CustomLogger):
from litellm.constants import PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES
from litellm.integrations.custom_logger import CustomLogger
prometheus_loggers: List[CustomLogger] = (
litellm.logging_callback_manager.get_custom_loggers_for_type(
callback_type=PrometheusLogger
)
prometheus_loggers: List[
CustomLogger
] = litellm.logging_callback_manager.get_custom_loggers_for_type(
callback_type=PrometheusLogger
)
# we need to get the initialized prometheus logger instance(s) and call logger.initialize_remaining_budget_metrics() on them
verbose_logger.debug("found %s prometheus loggers", len(prometheus_loggers))
@ -2283,7 +2617,7 @@ def prometheus_label_factory(
if UserAPIKeyLabelNames.END_USER.value in filtered_labels:
get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking()
filtered_labels["end_user"] = get_end_user_id_for_cost_tracking(
litellm_params={"user_api_key_end_user_id": enum_values.end_user},
service_type="prometheus",

2
litellm/litellm_core_utils/custom_logger_registry.py
View file

@ -18,6 +18,7 @@ from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLog
from litellm.integrations.bitbucket import BitBucketPromptManager
from litellm.integrations.braintrust_logging import BraintrustLogger
from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
from litellm.integrations.focus.focus_logger import FocusLogger
from litellm.integrations.datadog.datadog import DataDogLogger
from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger
from litellm.integrations.deepeval import DeepEvalLogger
@ -93,6 +94,7 @@ class CustomLoggerRegistry:
"bitbucket": BitBucketPromptManager,
"gitlab": GitLabPromptManager,
"cloudzero": CloudZeroLogger,
"focus": FocusLogger,
"posthog": PostHogLogger,
}

8
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -913,6 +913,14 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915
or "http://localhost:2024"
)
dynamic_api_key = api_key or get_secret_str("LANGGRAPH_API_KEY")
elif custom_llm_provider == "manus":
# Manus is OpenAI compatible for responses API
api_base = (
api_base
or get_secret_str("MANUS_API_BASE")
or "https://api.manus.im"
)
dynamic_api_key = api_key or get_secret_str("MANUS_API_KEY")
if api_base is not None and not isinstance(api_base, str):
raise Exception("api base needs to be a string. api_base={}".format(api_base))

21
litellm/litellm_core_utils/litellm_logging.py
View file

@ -3756,6 +3756,15 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
cloudzero_logger = CloudZeroLogger()
_in_memory_loggers.append(cloudzero_logger)
return cloudzero_logger # type: ignore
elif logging_integration == "focus":
from litellm.integrations.focus.focus_logger import FocusLogger
for callback in _in_memory_loggers:
if isinstance(callback, FocusLogger):
return callback # type: ignore
focus_logger = FocusLogger()
_in_memory_loggers.append(focus_logger)
return focus_logger # type: ignore
elif logging_integration == "deepeval":
for callback in _in_memory_loggers:
if isinstance(callback, DeepEvalLogger):
@ -4076,6 +4085,12 @@ def get_custom_logger_compatible_class( # noqa: PLR0915
for callback in _in_memory_loggers:
if isinstance(callback, CloudZeroLogger):
return callback
elif logging_integration == "focus":
from litellm.integrations.focus.focus_logger import FocusLogger
for callback in _in_memory_loggers:
if isinstance(callback, FocusLogger):
return callback
elif logging_integration == "deepeval":
for callback in _in_memory_loggers:
if isinstance(callback, DeepEvalLogger):
@ -4800,7 +4815,7 @@ class StandardLoggingPayloadSetup:
"""
Extract additional header tags for spend tracking based on config.
"""
extra_headers: List[str] = litellm.extra_spend_tag_headers or []
extra_headers: List[str] = getattr(litellm, "extra_spend_tag_headers", None) or []
if not extra_headers:
return None
@ -4824,9 +4839,9 @@ class StandardLoggingPayloadSetup:
metadata = litellm_params.get("metadata") or {}
litellm_metadata = litellm_params.get("litellm_metadata") or {}
if metadata.get("tags", []):
request_tags = metadata.get("tags", [])
request_tags = metadata.get("tags", []).copy()
elif litellm_metadata.get("tags", []):
request_tags = litellm_metadata.get("tags", [])
request_tags = litellm_metadata.get("tags", []).copy()
else:
request_tags = []
user_agent_tags = StandardLoggingPayloadSetup._get_user_agent_tags(

17
litellm/litellm_core_utils/prompt_templates/common_utils.py
View file

@ -6,6 +6,7 @@ import io
import mimetypes
import re
from os import PathLike
from pathlib import Path
from typing import (
TYPE_CHECKING,
Any,
@ -533,6 +534,12 @@ def extract_file_data(file_data: FileTypes) -> ExtractedFileData:
# Convert content to bytes
if isinstance(file_content, (str, PathLike)):
# If it's a path, open and read the file
# Extract filename from path if not already set
if filename is None:
if isinstance(file_content, PathLike):
filename = Path(file_content).name
else:
filename = Path(str(file_content)).name
with open(file_content, "rb") as f:
content = f.read()
elif isinstance(file_content, io.IOBase):
@ -550,11 +557,11 @@ def extract_file_data(file_data: FileTypes) -> ExtractedFileData:
# Use provided content type or guess based on filename
if not content_type:
content_type = (
mimetypes.guess_type(filename)[0]
if filename
else "application/octet-stream"
)
if filename:
guessed_type = mimetypes.guess_type(filename)[0]
content_type = guessed_type if guessed_type else "application/octet-stream"
else:
content_type = "application/octet-stream"
return ExtractedFileData(
filename=filename,

7
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -1645,9 +1645,12 @@ def convert_to_anthropic_tool_result(
)
elif content["type"] == "image_url":
format = content["image_url"].get("format") if isinstance(content["image_url"], dict) else None
anthropic_content_list.append(
create_anthropic_image_param(content["image_url"], format=format)
_anthropic_image_param = create_anthropic_image_param(content["image_url"], format=format)
_anthropic_image_param = add_cache_control_to_content(
anthropic_content_element=_anthropic_image_param,
original_content_element=content,
)
anthropic_content_list.append(_anthropic_image_param)
anthropic_content = anthropic_content_list
anthropic_tool_result: Optional[AnthropicMessagesToolResultParam] = None

67
litellm/litellm_core_utils/sensitive_data_masker.py
View file

@ -1,4 +1,5 @@
from typing import Any, Dict, Optional, Set
from collections.abc import Mapping
from typing import Any, Dict, List, Optional, Set
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER
@ -17,6 +18,7 @@ class SensitiveDataMasker:
"key",
"token",
"auth",
"authorization",
"credential",
"access",
"private",
@ -42,22 +44,52 @@ class SensitiveDataMasker:
else:
return f"{value_str[:self.visible_prefix]}{self.mask_char * masked_length}{value_str[-self.visible_suffix:]}"
def is_sensitive_key(self, key: str, excluded_keys: Optional[Set[str]] = None) -> bool:
def is_sensitive_key(
self, key: str, excluded_keys: Optional[Set[str]] = None
) -> bool:
# Check if key is in excluded_keys first (exact match)
if excluded_keys and key in excluded_keys:
return False
key_lower = str(key).lower()
# Split on underscores and check if any segment matches the pattern
# Split on underscores/hyphens and check if any segment matches the pattern
# This avoids false positives like "max_tokens" matching "token"
# but still catches "api_key", "access_token", etc.
key_segments = key_lower.replace('-', '_').split('_')
result = any(
pattern in key_segments
for pattern in self.sensitive_patterns
)
key_segments = key_lower.replace("-", "_").split("_")
result = any(pattern in key_segments for pattern in self.sensitive_patterns)
return result
def _mask_sequence(
self,
values: List[Any],
depth: int,
max_depth: int,
excluded_keys: Optional[Set[str]],
key_is_sensitive: bool,
) -> List[Any]:
masked_items: List[Any] = []
if depth >= max_depth:
return values
for item in values:
if isinstance(item, Mapping):
masked_items.append(
self.mask_dict(dict(item), depth + 1, max_depth, excluded_keys)
)
elif isinstance(item, list):
masked_items.append(
self._mask_sequence(
item, depth + 1, max_depth, excluded_keys, key_is_sensitive
)
)
elif key_is_sensitive and isinstance(item, str):
masked_items.append(self._mask_value(item))
else:
masked_items.append(
item if isinstance(item, (int, float, bool, str, list)) else str(item)
)
return masked_items
def mask_dict(
self,
data: Dict[str, Any],
@ -71,11 +103,20 @@ class SensitiveDataMasker:
masked_data: Dict[str, Any] = {}
for k, v in data.items():
try:
if isinstance(v, dict):
masked_data[k] = self.mask_dict(v, depth + 1, max_depth, excluded_keys)
key_is_sensitive = self.is_sensitive_key(k, excluded_keys)
if isinstance(v, Mapping):
masked_data[k] = self.mask_dict(
dict(v), depth + 1, max_depth, excluded_keys
)
elif isinstance(v, list):
masked_data[k] = self._mask_sequence(
v, depth + 1, max_depth, excluded_keys, key_is_sensitive
)
elif hasattr(v, "__dict__") and not isinstance(v, type):
masked_data[k] = self.mask_dict(vars(v), depth + 1, max_depth, excluded_keys)
elif self.is_sensitive_key(k, excluded_keys):
masked_data[k] = self.mask_dict(
vars(v), depth + 1, max_depth, excluded_keys
)
elif key_is_sensitive:
str_value = str(v) if v is not None else ""
masked_data[k] = self._mask_value(str_value)
else:

15
litellm/llms/anthropic/chat/transformation.py
View file

@ -1265,14 +1265,15 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
cache_creation_tokens=cache_creation_input_tokens,
cache_creation_token_details=cache_creation_token_details,
)
completion_token_details = (
CompletionTokensDetailsWrapper(
reasoning_tokens=token_counter(
text=reasoning_content, count_response_tokens=True
)
)
# Always populate completion_token_details, not just when there's reasoning_content
reasoning_tokens = (
token_counter(text=reasoning_content, count_response_tokens=True)
if reasoning_content
else None
else 0
)
completion_token_details = CompletionTokensDetailsWrapper(
reasoning_tokens=reasoning_tokens if reasoning_tokens > 0 else None,
text_tokens=completion_tokens - reasoning_tokens if reasoning_tokens > 0 else completion_tokens,
)
total_tokens = prompt_tokens + completion_tokens

19
litellm/llms/azure/chat/gpt_5_transformation.py
View file

@ -25,7 +25,24 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
return "gpt-5" in model or "gpt5_series" in model
def get_supported_openai_params(self, model: str) -> List[str]:
return OpenAIGPT5Config.get_supported_openai_params(self, model=model)
"""Get supported parameters for Azure OpenAI GPT-5 models.
Azure OpenAI GPT-5.2 models support logprobs, unlike OpenAI's GPT-5.
This overrides the parent class to add logprobs support back for gpt-5.2.
Reference:
- Tested with Azure OpenAI GPT-5.2 (api-version: 2025-01-01-preview)
- Azure returns logprobs successfully despite Microsoft's general
documentation stating reasoning models don't support it.
"""
params = OpenAIGPT5Config.get_supported_openai_params(self, model=model)
# Only gpt-5.2 has been verified to support logprobs on Azure
if self.is_model_gpt_5_2_model(model):
azure_supported_params = ["logprobs", "top_logprobs"]
params.extend(azure_supported_params)
return params
def map_openai_params(
self,

72
litellm/llms/bedrock/base_aws_llm.py
View file

@ -74,6 +74,41 @@ class BaseAWSLLM:
"aws_external_id",
]
def _get_ssl_verify(self):
"""
Get SSL verification setting for boto3 clients.
This ensures that custom CA certificates are properly used for all AWS API calls,
including STS and Bedrock services.
Returns:
Union[bool, str]: SSL verification setting - False to disable, True to enable,
or a string path to a CA bundle file
"""
import litellm
from litellm.secret_managers.main import str_to_bool
# Check environment variable first (highest priority)
ssl_verify = os.getenv("SSL_VERIFY", litellm.ssl_verify)
# Convert string "False"/"True" to boolean
if isinstance(ssl_verify, str):
# Check if it's a file path
if os.path.exists(ssl_verify):
return ssl_verify
# Otherwise try to convert to boolean
ssl_verify_bool = str_to_bool(ssl_verify)
if ssl_verify_bool is not None:
ssl_verify = ssl_verify_bool
# Check SSL_CERT_FILE environment variable for custom CA bundle
if ssl_verify is True or ssl_verify == "True":
ssl_cert_file = os.getenv("SSL_CERT_FILE")
if ssl_cert_file and os.path.exists(ssl_cert_file):
return ssl_cert_file
return ssl_verify
def get_cache_key(self, credential_args: Dict[str, Optional[str]]) -> str:
"""
Generate a unique cache key based on the credential arguments.
@ -314,6 +349,12 @@ class BaseAWSLLM:
if model.startswith("invoke/"):
model = model.replace("invoke/", "", 1)
# Special case: Check for "nova" in model name first (before "amazon")
# This handles amazon.nova-* models which would otherwise match "amazon" (Titan)
if "nova" in model.lower():
if "nova" in get_args(BEDROCK_INVOKE_PROVIDERS_LITERAL):
return cast(BEDROCK_INVOKE_PROVIDERS_LITERAL, "nova")
_split_model = model.split(".")[0]
if _split_model in get_args(BEDROCK_INVOKE_PROVIDERS_LITERAL):
return cast(BEDROCK_INVOKE_PROVIDERS_LITERAL, _split_model)
@ -323,13 +364,9 @@ class BaseAWSLLM:
if provider is not None:
return provider
# check if provider == "nova"
if "nova" in model:
return "nova"
else:
for provider in get_args(BEDROCK_INVOKE_PROVIDERS_LITERAL):
if provider in model:
return provider
for provider in get_args(BEDROCK_INVOKE_PROVIDERS_LITERAL):
if provider in model:
return provider
return None
@staticmethod
@ -364,11 +401,15 @@ class BaseAWSLLM:
elif provider == "qwen3" and "qwen3/" in model_id:
model_id = BaseAWSLLM._get_model_id_from_model_with_spec(
model_id, spec="qwen3"
)
)
elif provider == "stability" and "stability/" in model_id:
model_id = BaseAWSLLM._get_model_id_from_model_with_spec(
model_id, spec="stability"
)
elif provider == "moonshot" and "moonshot/" in model_id:
model_id = BaseAWSLLM._get_model_id_from_model_with_spec(
model_id, spec="moonshot"
)
return model_id
@staticmethod
@ -412,7 +453,7 @@ class BaseAWSLLM:
if "nova" in model.lower():
if "nova" in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL):
return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, "nova")
# Handle regional models like us.twelvelabs.marengo-embed-2-7-v1:0
if "." in model:
parts = model.split(".")
@ -563,6 +604,7 @@ class BaseAWSLLM:
"sts",
region_name=aws_region_name,
endpoint_url=sts_endpoint,
verify=self._get_ssl_verify(),
)
# https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html
@ -619,7 +661,7 @@ class BaseAWSLLM:
# Create an STS client without credentials
with tracer.trace("boto3.client(sts) for manual IRSA"):
sts_client = boto3.client("sts", region_name=region)
sts_client = boto3.client("sts", region_name=region, verify=self._get_ssl_verify())
# Manually assume the IRSA role with the session name
verbose_logger.debug(
@ -642,6 +684,7 @@ class BaseAWSLLM:
aws_access_key_id=irsa_creds["AccessKeyId"],
aws_secret_access_key=irsa_creds["SecretAccessKey"],
aws_session_token=irsa_creds["SessionToken"],
verify=self._get_ssl_verify(),
)
# Get current caller identity for debugging
@ -680,7 +723,7 @@ class BaseAWSLLM:
verbose_logger.debug("Same account role assumption, using automatic IRSA")
with tracer.trace("boto3.client(sts) with automatic IRSA"):
sts_client = boto3.client("sts", region_name=region)
sts_client = boto3.client("sts", region_name=region, verify=self._get_ssl_verify())
# Get current caller identity for debugging
try:
@ -803,7 +846,7 @@ class BaseAWSLLM:
# This allows the web identity token to work automatically
if aws_access_key_id is None and aws_secret_access_key is None:
with tracer.trace("boto3.client(sts)"):
sts_client = boto3.client("sts")
sts_client = boto3.client("sts", verify=self._get_ssl_verify())
else:
with tracer.trace("boto3.client(sts)"):
sts_client = boto3.client(
@ -811,6 +854,7 @@ class BaseAWSLLM:
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
verify=self._get_ssl_verify(),
)
assume_role_params = {
@ -958,7 +1002,9 @@ class BaseAWSLLM:
return endpoint_url, proxy_endpoint_url
def _select_default_endpoint_url(
self, endpoint_type: Optional[Literal["runtime", "agent", "agentcore"]], aws_region_name: str
self,
endpoint_type: Optional[Literal["runtime", "agent", "agentcore"]],
aws_region_name: str,
) -> str:
"""
Select the default endpoint url based on the endpoint type

256
litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py Normal file
View file

@ -0,0 +1,256 @@
"""
Transformation for Bedrock Moonshot AI (Kimi K2) models.
Supports the Kimi K2 Thinking model available on Amazon Bedrock.
Model format: bedrock/moonshot.kimi-k2-thinking-v1:0
Reference: https://aws.amazon.com/about-aws/whats-new/2025/12/amazon-bedrock-fully-managed-open-weight-models/
"""
from typing import TYPE_CHECKING, Any, List, Optional, Union
import re
import httpx
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
AmazonInvokeConfig,
)
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.moonshot.chat.transformation import MoonshotChatConfig
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import Choices
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from litellm.types.utils import ModelResponse
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig):
"""
Configuration for Bedrock Moonshot AI (Kimi K2) models.
Reference:
https://aws.amazon.com/about-aws/whats-new/2025/12/amazon-bedrock-fully-managed-open-weight-models/
https://platform.moonshot.ai/docs/api/chat
Supported Params for the Amazon / Moonshot models:
- `max_tokens` (integer) max tokens
- `temperature` (float) temperature for model (0-1 for Moonshot)
- `top_p` (float) top p for model
- `stream` (bool) whether to stream responses
- `tools` (list) tool definitions (supported on kimi-k2-thinking)
- `tool_choice` (str|dict) tool choice specification (supported on kimi-k2-thinking)
NOT Supported on Bedrock:
- `stop` sequences (Bedrock doesn't support stopSequences field for this model)
Note: The kimi-k2-thinking model DOES support tool calls, unlike kimi-thinking-preview.
"""
def __init__(self, **kwargs):
AmazonInvokeConfig.__init__(self, **kwargs)
MoonshotChatConfig.__init__(self, **kwargs)
@property
def custom_llm_provider(self) -> Optional[str]:
return "bedrock"
def _get_model_id(self, model: str) -> str:
"""
Extract the actual model ID from the LiteLLM model name.
Removes routing prefixes like:
- bedrock/invoke/moonshot.kimi-k2-thinking -> moonshot.kimi-k2-thinking
- invoke/moonshot.kimi-k2-thinking -> moonshot.kimi-k2-thinking
- moonshot.kimi-k2-thinking -> moonshot.kimi-k2-thinking
"""
# Remove bedrock/ prefix if present
if model.startswith("bedrock/"):
model = model[8:]
# Remove invoke/ prefix if present
if model.startswith("invoke/"):
model = model[7:]
# Remove any provider prefix (e.g., moonshot/)
if "/" in model and not model.startswith("arn:"):
parts = model.split("/", 1)
if len(parts) == 2:
model = parts[1]
return model
def get_supported_openai_params(self, model: str) -> List[str]:
"""
Get the supported OpenAI params for Moonshot AI models on Bedrock.
Bedrock-specific limitations:
- stopSequences field is not supported on Bedrock (unlike native Moonshot API)
- functions parameter is not supported (use tools instead)
- tool_choice doesn't support "required" value
Note: kimi-k2-thinking DOES support tool calls (unlike kimi-thinking-preview)
The parent MoonshotChatConfig class handles the kimi-thinking-preview exclusion.
"""
excluded_params: List[str] = ["functions", "stop"] # Bedrock doesn't support stopSequences
base_openai_params = super(MoonshotChatConfig, self).get_supported_openai_params(model=model)
final_params: List[str] = []
for param in base_openai_params:
if param not in excluded_params:
final_params.append(param)
return final_params
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
"""
Map OpenAI parameters to Moonshot AI parameters for Bedrock.
Handles Moonshot AI specific limitations:
- tool_choice doesn't support "required" value
- Temperature <0.3 limitation for n>1
- Temperature range is [0, 1] (not [0, 2] like OpenAI)
"""
return MoonshotChatConfig.map_openai_params(
self,
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=drop_params,
)
def transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
"""
Transform the request for Bedrock Moonshot AI models.
Uses the Moonshot transformation logic which handles:
- Converting content lists to strings (Moonshot doesn't support list format)
- Adding tool_choice="required" message if needed
- Temperature and parameter validation
"""
# Filter out AWS credentials using the existing method from BaseAWSLLM
self._get_boto_credentials_from_optional_params(optional_params, model)
# Strip routing prefixes to get the actual model ID
clean_model_id = self._get_model_id(model)
# Use Moonshot's transform_request which handles message transformation
# and tool_choice="required" workaround
return MoonshotChatConfig.transform_request(
self,
model=clean_model_id,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
def _extract_reasoning_from_content(self, content: str) -> tuple[Optional[str], str]:
"""
Extract reasoning content from <reasoning> tags in the response.
Moonshot AI's Kimi K2 Thinking model returns reasoning in <reasoning> tags.
This method extracts that content and returns it separately.
Args:
content: The full content string from the API response
Returns:
tuple: (reasoning_content, main_content)
"""
if not content:
return None, content
# Match <reasoning>...</reasoning> tags
reasoning_match = re.match(
r"<reasoning>(.*?)</reasoning>\s*(.*)",
content,
re.DOTALL
)
if reasoning_match:
reasoning_content = reasoning_match.group(1).strip()
main_content = reasoning_match.group(2).strip()
return reasoning_content, main_content
return None, content
def transform_response(
self,
model: str,
raw_response: httpx.Response,
model_response: "ModelResponse",
logging_obj: LiteLLMLoggingObj,
request_data: dict,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> "ModelResponse":
"""
Transform the response from Bedrock Moonshot AI models.
Moonshot AI uses OpenAI-compatible response format, but returns reasoning
content in <reasoning> tags. This method:
1. Calls parent class transformation
2. Extracts reasoning content from <reasoning> tags
3. Sets reasoning_content on the message object
"""
# First, get the standard transformation
model_response = MoonshotChatConfig.transform_response(
self,
model=model,
raw_response=raw_response,
model_response=model_response,
logging_obj=logging_obj,
request_data=request_data,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
api_key=api_key,
json_mode=json_mode,
)
# Extract reasoning content from <reasoning> tags
if model_response.choices and len(model_response.choices) > 0:
for choice in model_response.choices:
# Only process Choices (not StreamingChoices) which have message attribute
if isinstance(choice, Choices) and choice.message and choice.message.content:
reasoning_content, main_content = self._extract_reasoning_from_content(
choice.message.content
)
if reasoning_content:
# Set the reasoning_content field
choice.message.reasoning_content = reasoning_content
# Update the main content without reasoning tags
choice.message.content = main_content
return model_response
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BedrockError:
"""Return the appropriate error class for Bedrock."""
return BedrockError(status_code=status_code, message=error_message)

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

@ -524,6 +524,12 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
if model.startswith("invoke/"):
model = model.replace("invoke/", "", 1)
# Special case: Check for "nova" in model name first (before "amazon")
# This handles amazon.nova-* models which would otherwise match "amazon" (Titan)
if "nova" in model.lower():
if "nova" in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL):
return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, "nova")
_split_model = model.split(".")[0]
if _split_model in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL):
return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, _split_model)
@ -533,10 +539,6 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
if provider is not None:
return provider
# check if provider == "nova"
if "nova" in model:
return "nova"
for provider in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL):
if provider in model:
return provider

161
litellm/llms/bedrock/common_utils.py
View file

@ -15,7 +15,7 @@ import litellm
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret
@ -178,7 +178,26 @@ def init_bedrock_client(
) = params_to_check
# SSL certificates (a.k.a CA bundle) used to verify the identity of requested hosts.
# Use the same logic as BaseAWSLLM._get_ssl_verify() for consistency
from litellm.secret_managers.main import str_to_bool
ssl_verify = os.getenv("SSL_VERIFY", litellm.ssl_verify)
# Convert string "False"/"True" to boolean
if isinstance(ssl_verify, str):
# Check if it's a file path
if os.path.exists(ssl_verify):
pass # Keep the file path
else:
# Otherwise try to convert to boolean
ssl_verify_bool = str_to_bool(ssl_verify)
if ssl_verify_bool is not None:
ssl_verify = ssl_verify_bool
# Check SSL_CERT_FILE environment variable for custom CA bundle
if ssl_verify is True or ssl_verify == "True":
ssl_cert_file = os.getenv("SSL_CERT_FILE")
if ssl_cert_file and os.path.exists(ssl_cert_file):
ssl_verify = ssl_cert_file
### SET REGION NAME
if region_name:
@ -229,7 +248,7 @@ def init_bedrock_client(
status_code=401,
)
sts_client = boto3.client("sts")
sts_client = boto3.client("sts", verify=ssl_verify)
# https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html
# https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html
@ -359,6 +378,70 @@ def get_bedrock_tool_name(response_tool_name: str) -> str:
return response_tool_name
# Cache the global regions list at module level
_BEDROCK_GLOBAL_REGIONS: Optional[List[str]] = None
def _get_all_bedrock_regions() -> List[str]:
"""Get all Bedrock regions, cached at module level."""
global _BEDROCK_GLOBAL_REGIONS
if _BEDROCK_GLOBAL_REGIONS is None:
_BEDROCK_GLOBAL_REGIONS = AmazonBedrockGlobalConfig().get_all_regions()
return _BEDROCK_GLOBAL_REGIONS
def get_bedrock_cross_region_inference_regions() -> List[str]:
"""Abbreviations of regions AWS Bedrock supports for cross region inference."""
return ["global", "us", "eu", "apac", "jp", "au", "us-gov"]
def extract_model_name_from_bedrock_arn(model: str) -> str:
"""
Extract the model name from an AWS Bedrock ARN.
Returns the string after the last '/' if 'arn' is in the input string.
"""
if "arn" in model.lower():
return model.split("/")[-1]
return model
def strip_bedrock_routing_prefix(model: str) -> str:
"""Strip LiteLLM routing prefixes from model name."""
for prefix in ["bedrock/", "converse/", "invoke/", "openai/"]:
if model.startswith(prefix):
model = model.split("/", 1)[1]
return model
def get_bedrock_base_model(model: str) -> str:
"""
Get the base model from the given model name.
Handle model names like:
- "us.meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1"
- "bedrock/converse/model" -> "model"
"""
model = strip_bedrock_routing_prefix(model)
model = extract_model_name_from_bedrock_arn(model)
potential_region = model.split(".", 1)[0]
alt_potential_region = model.split("/", 1)[0]
if potential_region in get_bedrock_cross_region_inference_regions():
return model.split(".", 1)[1]
elif (
alt_potential_region in _get_all_bedrock_regions()
and len(model.split("/", 1)) > 1
):
return model.split("/", 1)[1]
return model
# Import after standalone functions to avoid circular imports
from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter
class BedrockModelInfo(BaseLLMModelInfo):
global_config = AmazonBedrockGlobalConfig()
all_global_regions = global_config.get_all_regions()
@ -394,76 +477,34 @@ class BedrockModelInfo(BaseLLMModelInfo):
) -> List[str]:
return []
@staticmethod
def extract_model_name_from_arn(model: str) -> str:
def get_token_counter(self) -> Optional[BaseTokenCounter]:
"""
Extract the model name from an AWS Bedrock ARN.
Returns the string after the last '/' if 'arn' is in the input string.
Args:
arn (str): The ARN string to parse
Factory method to create a Bedrock token counter.
Returns:
str: The extracted model name if 'arn' is in the string,
otherwise returns the original string
BedrockTokenCounter instance for this provider.
"""
if "arn" in model.lower():
return model.split("/")[-1]
return model
return BedrockTokenCounter()
@staticmethod
def extract_model_name_from_arn(model: str) -> str:
"""Wrapper for standalone function. See extract_model_name_from_bedrock_arn()."""
return extract_model_name_from_bedrock_arn(model)
@staticmethod
def get_non_litellm_routing_model_name(model: str) -> str:
if model.startswith("bedrock/"):
model = model.split("/", 1)[1]
if model.startswith("converse/"):
model = model.split("/", 1)[1]
if model.startswith("invoke/"):
model = model.split("/", 1)[1]
if model.startswith("openai/"):
model = model.split("/", 1)[1]
return model
"""Wrapper for standalone function. See strip_bedrock_routing_prefix()."""
return strip_bedrock_routing_prefix(model)
@staticmethod
def get_base_model(model: str) -> str:
"""
Get the base model from the given model name.
Handle model names like - "us.meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1"
AND "meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1"
"""
model = BedrockModelInfo.get_non_litellm_routing_model_name(model=model)
model = BedrockModelInfo.extract_model_name_from_arn(model)
potential_region = model.split(".", 1)[0]
alt_potential_region = model.split("/", 1)[
0
] # in model cost map we store regional information like `/us-west-2/bedrock-model`
if (
potential_region
in BedrockModelInfo._supported_cross_region_inference_region()
):
return model.split(".", 1)[1]
elif (
alt_potential_region in BedrockModelInfo.all_global_regions
and len(model.split("/", 1)) > 1
):
return model.split("/", 1)[1]
return model
"""Wrapper for standalone function. See get_bedrock_base_model()."""
return get_bedrock_base_model(model)
@staticmethod
def _supported_cross_region_inference_region() -> List[str]:
"""
Abbreviations of regions AWS Bedrock supports for cross region inference
"""
return ["global", "us", "eu", "apac", "jp", "au", "us-gov"]
"""Wrapper for standalone function. See get_bedrock_cross_region_inference_regions()."""
return get_bedrock_cross_region_inference_regions()
@staticmethod
def get_bedrock_route(
@ -629,6 +670,8 @@ def get_bedrock_chat_config(model: str):
return litellm.AmazonCohereConfig()
elif bedrock_invoke_provider == "mistral":
return litellm.AmazonMistralConfig()
elif bedrock_invoke_provider == "moonshot":
return litellm.AmazonMoonshotConfig()
elif bedrock_invoke_provider == "deepseek_r1":
return litellm.AmazonDeepSeekR1Config()
elif bedrock_invoke_provider == "nova":

87
litellm/llms/bedrock/count_tokens/bedrock_token_counter.py Normal file
View file

@ -0,0 +1,87 @@
"""
Bedrock Token Counter implementation using the CountTokens API.
"""
from typing import Any, Dict, List, Optional
from litellm._logging import verbose_logger
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.llms.bedrock.common_utils import get_bedrock_base_model
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
from litellm.types.utils import LlmProviders, TokenCountResponse
class BedrockTokenCounter(BaseTokenCounter):
"""Token counter implementation for AWS Bedrock provider using the CountTokens API."""
def should_use_token_counting_api(
self,
custom_llm_provider: Optional[str] = None,
) -> bool:
"""
Returns True if we should use the Bedrock CountTokens API for token counting.
"""
return custom_llm_provider == LlmProviders.BEDROCK.value
async def count_tokens(
self,
model_to_use: str,
messages: Optional[List[Dict[str, Any]]],
contents: Optional[List[Dict[str, Any]]],
deployment: Optional[Dict[str, Any]] = None,
request_model: str = "",
) -> Optional[TokenCountResponse]:
"""
Count tokens using AWS Bedrock's CountTokens API.
This method calls the existing BedrockCountTokensHandler to make an API call
to Bedrock's token counting endpoint, bypassing the local tiktoken-based counting.
Args:
model_to_use: The model identifier
messages: The messages to count tokens for
contents: Alternative content format (not used for Bedrock)
deployment: Deployment configuration containing litellm_params
request_model: The original request model name
Returns:
TokenCountResponse with token count, or None if counting fails
"""
if not messages:
return None
deployment = deployment or {}
litellm_params = deployment.get("litellm_params", {})
# Build request data in the format expected by BedrockCountTokensHandler
request_data = {
"model": model_to_use,
"messages": messages,
}
# Get the resolved model (strip prefixes like bedrock/, converse/, etc.)
resolved_model = get_bedrock_base_model(model_to_use)
try:
handler = BedrockCountTokensHandler()
result = await handler.handle_count_tokens_request(
request_data=request_data,
litellm_params=litellm_params,
resolved_model=resolved_model,
)
# Transform response to TokenCountResponse
if result is not None:
return TokenCountResponse(
total_tokens=result.get("input_tokens", 0),
request_model=request_model,
model_used=model_to_use,
tokenizer_type="bedrock_api",
original_response=result,
)
except Exception as e:
verbose_logger.warning(
f"Error calling Bedrock CountTokens API: {e}, falling back to default tokenizer"
)
return None

20
litellm/llms/bedrock/count_tokens/handler.py
View file

@ -6,10 +6,9 @@ Simplified handler leveraging existing LiteLLM Bedrock infrastructure.
from typing import Any, Dict
from fastapi import HTTPException
import litellm
from litellm._logging import verbose_logger
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
@ -70,6 +69,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
verbose_logger.debug(f"Making request to: {endpoint_url}")
# Use existing _sign_request method from BaseAWSLLM
# Extract api_key for bearer token auth if provided
api_key = litellm_params.get("api_key", None)
headers = {"Content-Type": "application/json"}
signed_headers, signed_body = self._sign_request(
service_name="bedrock",
@ -78,6 +79,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
request_data=bedrock_request,
api_base=endpoint_url,
model=resolved_model,
api_key=api_key,
)
async_client = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)
@ -94,9 +96,9 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
if response.status_code != 200:
error_text = response.text
verbose_logger.error(f"AWS Bedrock error: {error_text}")
raise HTTPException(
status_code=400,
detail={"error": f"AWS Bedrock error: {error_text}"},
raise BedrockError(
status_code=response.status_code,
message=f"AWS Bedrock error: {error_text}",
)
bedrock_response = response.json()
@ -112,12 +114,12 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
return final_response
except HTTPException:
# Re-raise HTTP exceptions as-is
except BedrockError:
# Re-raise Bedrock exceptions as-is
raise
except Exception as e:
verbose_logger.error(f"Error in CountTokens handler: {str(e)}")
raise HTTPException(
raise BedrockError(
status_code=500,
detail={"error": f"CountTokens processing error: {str(e)}"},
message=f"CountTokens processing error: {str(e)}",
)

4
litellm/llms/bedrock/count_tokens/transformation.py
View file

@ -8,7 +8,7 @@ to AWS Bedrock's CountTokens API format and vice versa.
from typing import Any, Dict, List
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import BedrockModelInfo
from litellm.llms.bedrock.common_utils import get_bedrock_base_model
class BedrockCountTokensConfig(BaseAWSLLM):
@ -141,7 +141,7 @@ class BedrockCountTokensConfig(BaseAWSLLM):
Complete endpoint URL for CountTokens API
"""
# Use existing LiteLLM function to get the base model ID (removes region prefix)
model_id = BedrockModelInfo.get_base_model(model)
model_id = get_bedrock_base_model(model)
# Remove bedrock/ prefix if present
if model_id.startswith("bedrock/"):

1
litellm/llms/bedrock/files/handler.py
View file

@ -142,6 +142,7 @@ class BedrockFilesHandler(BaseAWSLLM):
aws_secret_access_key=credentials.secret_key,
aws_session_token=credentials.token,
region_name=aws_region_name,
verify=self._get_ssl_verify(),
)
# Download file from S3

9
litellm/llms/bedrock/passthrough/transformation.py
View file

@ -34,11 +34,12 @@ class BedrockPassthroughConfig(
litellm_params: dict,
) -> Tuple["URL", str]:
optional_params = litellm_params.copy()
model_id = optional_params.get("model_id", None)
aws_region_name = self._get_aws_region_name(
optional_params=optional_params,
model=model,
model_id=None,
model_id=model_id,
)
aws_bedrock_runtime_endpoint = optional_params.get("aws_bedrock_runtime_endpoint")
@ -49,6 +50,12 @@ class BedrockPassthroughConfig(
endpoint_type="runtime",
)
# If model_id is provided (e.g., Application Inference Profile ARN), use it in the endpoint
# instead of the translated model name
if model_id is not None:
# Replace the model name in the endpoint with the model_id
import re
endpoint = re.sub(r'model/[^/]+/', f'model/{model_id}/', endpoint)
return self.format_url(endpoint, endpoint_url, request_query_params or {}), endpoint_url
def sign_request(

77
litellm/llms/deepinfra/chat/transformation.py
View file

@ -1,9 +1,11 @@
from typing import Optional, Tuple, Union
import json
from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, cast, overload
import litellm
from litellm.constants import MIN_NON_ZERO_TEMPERATURE
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
class DeepInfraConfig(OpenAIGPTConfig):
@ -117,6 +119,79 @@ class DeepInfraConfig(OpenAIGPTConfig):
optional_params[param] = value
return optional_params
def _transform_tool_message_content(self, messages: List[AllMessageValues]) -> List[AllMessageValues]:
"""
Transform tool message content from array to string format for DeepInfra compatibility.
DeepInfra requires tool message content to be a string, not an array.
This method converts tool message content from array format to string format.
Example transformation:
- Input: {"role": "tool", "content": [{"type": "text", "text": "20"}]}
- Output: {"role": "tool", "content": "20"}
Or if content is complex:
- Input: {"role": "tool", "content": [{"type": "text", "text": "result"}]}
- Output: {"role": "tool", "content": "[{\"type\": \"text\", \"text\": \"result\"}]"}
"""
for message in messages:
if message.get("role") == "tool":
content = message.get("content")
# If content is a list/array, convert it to string
if isinstance(content, list):
# Check if it's a simple single text item
if (
len(content) == 1
and isinstance(content[0], dict)
and content[0].get("type") == "text"
and "text" in content[0]
):
# Extract just the text value for simple cases
message["content"] = content[0]["text"]
else:
# For complex content, serialize the entire array as JSON string
message["content"] = json.dumps(content)
return messages
@overload
def _transform_messages(
self, messages: List[AllMessageValues], model: str, is_async: Literal[True]
) -> Coroutine[Any, Any, List[AllMessageValues]]:
...
@overload
def _transform_messages(
self, messages: List[AllMessageValues], model: str, is_async: Literal[False] = False
) -> List[AllMessageValues]:
...
def _transform_messages(
self, messages: List[AllMessageValues], model: str, is_async: bool = False
) -> Union[List[AllMessageValues], Coroutine[Any, Any, List[AllMessageValues]]]:
"""
Transform messages for DeepInfra compatibility.
Handles both sync and async transformations.
"""
if is_async:
# For async case, create an async function that awaits parent and applies our transformation
async def _async_transform():
# Call parent with is_async=True (literal) for async case
parent_result = super(DeepInfraConfig, self)._transform_messages(
messages=messages, model=model, is_async=cast(Literal[True], True)
)
transformed_messages = await parent_result
return self._transform_tool_message_content(transformed_messages)
return _async_transform()
else:
# Call parent with is_async=False (literal) for sync case
parent_result = super()._transform_messages(
messages=messages, model=model, is_async=cast(Literal[False], False)
)
# For sync case, parent_result is already the transformed messages
return self._transform_tool_message_content(parent_result)
def _get_openai_compatible_provider_info(
self, api_base: Optional[str], api_key: Optional[str]
) -> Tuple[Optional[str], Optional[str]]:

9
litellm/llms/gemini/common_utils.py
View file

@ -150,6 +150,15 @@ def get_api_key_from_env() -> Optional[str]:
return get_secret_str("GOOGLE_API_KEY") or get_secret_str("GEMINI_API_KEY")
def get_vertex_api_key_from_env() -> Optional[str]:
"""
Get API key from environment for Vertex AI.
Checks VERTEXAI_API_KEY and VERTEX_API_KEY environment variables.
This allows using Vertex AI with API keys instead of service account credentials.
"""
return get_secret_str("VERTEXAI_API_KEY") or get_secret_str("VERTEX_API_KEY")
class GoogleAIStudioTokenCounter(BaseTokenCounter):
"""Token counter implementation for Google AI Studio provider."""
def should_use_token_counting_api(

2
litellm/llms/manus/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# Manus provider implementation

2
litellm/llms/manus/responses/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# Manus Responses API implementation

308
litellm/llms/manus/responses/transformation.py Normal file
View file

@ -0,0 +1,308 @@
from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
_safe_convert_created_field,
)
from litellm.llms.openai.common_utils import OpenAIError
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
ResponseAPIUsage,
ResponseInputParam,
ResponsesAPIResponse,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
MANUS_API_BASE = "https://api.manus.im"
class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig):
"""
Configuration for Manus API's Responses API.
Manus API is OpenAI-compatible but has some differences:
- API key passed via `API_KEY` header (not `Authorization: Bearer`)
- Model format: `manus/{agent_profile}` (e.g., `manus/manus-1.6`)
- Requires `extra_body` with `task_mode: "agent"` and `agent_profile`
Reference: https://open.manus.im/docs/openai-compatibility
"""
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.MANUS
def should_fake_stream(
self,
model: Optional[str],
stream: Optional[bool],
custom_llm_provider: Optional[str] = None,
) -> bool:
"""
Manus API doesn't support real-time streaming.
It returns a task that runs asynchronously.
We fake streaming by converting the response into streaming events.
"""
return stream is True
def _extract_agent_profile(self, model: str) -> str:
"""
Extract agent profile from model name.
Model format: `manus/{agent_profile}`
Examples: `manus/manus-1.6`, `manus/manus-1.6-lite`, `manus/manus-1.6-max`
Returns:
str: The agent profile (e.g., "manus-1.6")
"""
if "/" in model:
return model.split("/", 1)[1]
# If no slash, assume the model name itself is the agent profile
return model
def validate_environment(
self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]
) -> dict:
"""
Validate environment and set up headers for Manus API.
Manus uses `API_KEY` header instead of `Authorization: Bearer`.
"""
litellm_params = litellm_params or GenericLiteLLMParams()
api_key = (
litellm_params.api_key
or litellm.api_key
or get_secret_str("MANUS_API_KEY")
)
if not api_key:
raise ValueError(
"Manus API key is required. Set MANUS_API_KEY environment variable or pass api_key parameter."
)
# Manus uses API_KEY header, not Authorization: Bearer
headers.update(
{
"API_KEY": api_key,
}
)
return headers
def get_complete_url(
self,
api_base: Optional[str],
litellm_params: dict,
) -> str:
"""
Get the complete URL for Manus Responses API endpoint.
Returns:
str: The full URL for the Manus /v1/responses endpoint
"""
api_base = (
api_base
or litellm.api_base
or get_secret_str("MANUS_API_BASE")
or MANUS_API_BASE
)
# Remove trailing slashes
api_base = api_base.rstrip("/")
# Manus API uses /v1/responses endpoint (OpenAI-compatible)
if api_base.endswith("/v1"):
return f"{api_base}/responses"
return f"{api_base}/v1/responses"
def transform_responses_api_request(
self,
model: str,
input: Union[str, ResponseInputParam],
response_api_optional_request_params: Dict,
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> Dict:
"""
Transform the request for Manus API.
Manus requires:
- `task_mode: "agent"` in the request body
- `agent_profile` extracted from model name in the request body
"""
# First, get the base OpenAI request
base_request = super().transform_responses_api_request(
model=model,
input=input,
response_api_optional_request_params=response_api_optional_request_params,
litellm_params=litellm_params,
headers=headers,
)
# Extract agent profile from model name
agent_profile = self._extract_agent_profile(model=model)
# Add Manus-specific parameters directly to the request body
# These will be sent as part of the request
base_request["task_mode"] = "agent"
base_request["agent_profile"] = agent_profile
# Merge any existing extra_body into the request
extra_body = response_api_optional_request_params.get("extra_body", {}) or {}
if extra_body:
base_request.update(extra_body)
# Avoid logging potentially sensitive agent_profile value
verbose_logger.debug("Manus: Using task_mode=agent")
return base_request
def transform_response_api_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
) -> ResponsesAPIResponse:
"""
Transform Manus API response to OpenAI-compatible format.
Manus uses camelCase (createdAt) instead of snake_case (created_at).
"""
try:
logging_obj.post_call(
original_response=raw_response.text,
additional_args={"complete_input_dict": {}},
)
raw_response_json = raw_response.json()
# Manus uses camelCase "createdAt" instead of snake_case "created_at"
if "createdAt" in raw_response_json and "created_at" not in raw_response_json:
raw_response_json["created_at"] = _safe_convert_created_field(
raw_response_json["createdAt"]
)
# Ensure created_at is set
if "created_at" in raw_response_json:
raw_response_json["created_at"] = _safe_convert_created_field(
raw_response_json["created_at"]
)
except Exception:
raise OpenAIError(
message=raw_response.text, status_code=raw_response.status_code
)
raw_response_headers = dict(raw_response.headers)
processed_headers = process_response_headers(raw_response_headers)
# Ensure reasoning is an empty dict if not present, OpenAI SDK does not allow None
if "reasoning" not in raw_response_json or raw_response_json.get("reasoning") is None:
raw_response_json["reasoning"] = {}
if "text" not in raw_response_json or raw_response_json.get("text") is None:
raw_response_json["text"] = {}
if "output" not in raw_response_json or raw_response_json.get("output") is None:
raw_response_json["output"] = []
# Ensure usage is present with default values if not provided
if "usage" not in raw_response_json or raw_response_json.get("usage") is None:
raw_response_json["usage"] = ResponseAPIUsage(
input_tokens=0,
output_tokens=0,
total_tokens=0,
)
try:
response = ResponsesAPIResponse(**raw_response_json)
except Exception:
verbose_logger.debug(
f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct"
)
response = ResponsesAPIResponse.model_construct(**raw_response_json)
# Store processed headers in additional_headers so they get returned to the client
response._hidden_params["additional_headers"] = processed_headers
response._hidden_params["headers"] = raw_response_headers
return response
def transform_get_response_api_request(
self,
response_id: str,
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> Tuple[str, Dict]:
"""
Transform the get response API request into a URL and data.
Manus API follows OpenAI-compatible format:
- GET /v1/responses/{response_id}
Reference: https://open.manus.im/docs/openai-compatibility
"""
url = f"{api_base}/{response_id}"
data: Dict = {}
return url, data
def transform_get_response_api_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
) -> ResponsesAPIResponse:
"""
Transform Manus API GET response to OpenAI-compatible format.
Manus uses camelCase (createdAt) instead of snake_case (created_at).
Same transformation as transform_response_api_response.
"""
try:
logging_obj.post_call(
original_response=raw_response.text,
additional_args={"complete_input_dict": {}},
)
raw_response_json = raw_response.json()
# Manus uses camelCase "createdAt" instead of snake_case "created_at"
if "createdAt" in raw_response_json and "created_at" not in raw_response_json:
raw_response_json["created_at"] = _safe_convert_created_field(
raw_response_json["createdAt"]
)
# Ensure created_at is set
if "created_at" in raw_response_json:
raw_response_json["created_at"] = _safe_convert_created_field(
raw_response_json["created_at"]
)
except Exception:
raise OpenAIError(
message=raw_response.text, status_code=raw_response.status_code
)
raw_response_headers = dict(raw_response.headers)
processed_headers = process_response_headers(raw_response_headers)
try:
response = ResponsesAPIResponse(**raw_response_json)
except Exception:
verbose_logger.debug(
f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct"
)
response = ResponsesAPIResponse.model_construct(**raw_response_json)
# Store processed headers in additional_headers so they get returned to the client
response._hidden_params["additional_headers"] = processed_headers
response._hidden_params["headers"] = raw_response_headers
return response

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

@ -771,9 +771,9 @@ class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator):
return ModelResponseStream(
id=chunk["id"],
object="chat.completion.chunk",
created=chunk["created"],
model=chunk["model"],
choices=chunk["choices"],
created=chunk.get("created"),
model=chunk.get("model"),
choices=chunk.get("choices", []),
)
except Exception as e:
raise e

4
litellm/llms/openai_like/providers.json
View file

@ -61,6 +61,10 @@
"max_completion_tokens": "max_tokens"
}
},
"abliteration": {
"base_url": "https://api.abliteration.ai/v1",
"api_key_env": "ABLITERATION_API_KEY"
},
"llamagate": {
"base_url": "https://api.llamagate.dev/v1",
"api_key_env": "LLAMAGATE_API_KEY",

182
litellm/llms/openrouter/embedding/transformation.py Normal file
View file

@ -0,0 +1,182 @@
"""
OpenRouter Embedding API Configuration.
This module provides the configuration for OpenRouter's Embedding API.
OpenRouter is OpenAI-compatible and supports embeddings via the /v1/embeddings endpoint.
Docs: https://openrouter.ai/docs
"""
from typing import TYPE_CHECKING, Any, Optional
import httpx
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
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 OpenRouterException
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class OpenrouterEmbeddingConfig(BaseEmbeddingConfig):
"""
Configuration for OpenRouter's Embedding API.
Reference: https://openrouter.ai/docs
"""
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 OpenRouter API.
OpenRouter requires:
- Authorization header with Bearer token
- HTTP-Referer header (site URL)
- X-Title header (app name)
"""
from litellm import get_secret
# Get OpenRouter-specific headers
openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai"
openrouter_app_name = get_secret("OR_APP_NAME") or "liteLLM"
openrouter_headers = {
"HTTP-Referer": openrouter_site_url,
"X-Title": openrouter_app_name,
"Content-Type": "application/json",
}
# Add Authorization header if api_key is provided
if api_key:
openrouter_headers["Authorization"] = f"Bearer {api_key}"
# Merge with existing headers (user's extra_headers take priority)
merged_headers = {**openrouter_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 OpenRouter Embedding API endpoint.
"""
# api_base is already set to https://openrouter.ai/api/v1 in main.py
# Remove trailing slashes
if api_base:
api_base = api_base.rstrip("/")
else:
api_base = "https://openrouter.ai/api/v1"
# Return the embeddings endpoint
return f"{api_base}/embeddings"
def transform_embedding_request(
self,
model: str,
input: AllEmbeddingInputValues,
optional_params: dict,
headers: dict,
) -> dict:
"""
Transform embedding request to OpenRouter format (OpenAI-compatible).
"""
# Ensure input is a list
if isinstance(input, str):
input = [input]
# OpenRouter expects the full model name (e.g., google/gemini-embedding-001)
# Strip 'openrouter/' prefix if present
if model.startswith("openrouter/"):
model = model.replace("openrouter/", "", 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 OpenRouter format (OpenAI-compatible).
"""
logging_obj.post_call(original_response=raw_response.text)
# OpenRouter 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 OpenRouter embeddings.
"""
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 OpenRouter 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 OpenRouter errors.
"""
return OpenRouterException(
message=error_message,
status_code=status_code,
headers=headers,
)

2
litellm/llms/vertex_ai/gemini/transformation.py
View file

@ -115,7 +115,7 @@ def _process_gemini_image(
and (image_type := format or _get_image_mime_type_from_url(image_url))
is not None
):
file_data = FileDataType(file_uri=image_url, mime_type=image_type)
file_data = FileDataType(mime_type=image_type, file_uri=image_url)
part = {"file_data": file_data}
if media_resolution_enum is not None and model is not None:

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

@ -480,20 +480,21 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
or tool_name == VertexToolName.CODE_EXECUTION.value
): # code_execution maintained for backwards compatibility
code_execution = self.get_tool_value(tool, "codeExecution")
elif tool_name and tool_name == VertexToolName.GOOGLE_SEARCH.value:
googleSearch = self.get_tool_value(
tool, VertexToolName.GOOGLE_SEARCH.value
)
elif (
tool_name and tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value
elif tool_name and (
tool_name == VertexToolName.GOOGLE_SEARCH.value
or tool_name == "google_search"
):
googleSearchRetrieval = self.get_tool_value(
tool, VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value
)
elif tool_name and tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value:
enterpriseWebSearch = self.get_tool_value(
tool, VertexToolName.ENTERPRISE_WEB_SEARCH.value
)
googleSearch = self.get_tool_value(tool, tool_name)
elif tool_name and (
tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value
or tool_name == "google_search_retrieval"
):
googleSearchRetrieval = self.get_tool_value(tool, tool_name)
elif tool_name and (
tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value
or tool_name == "enterprise_web_search"
):
enterpriseWebSearch = self.get_tool_value(tool, tool_name)
elif tool_name and (
tool_name == VertexToolName.URL_CONTEXT.value
or tool_name == "urlContext"
@ -1811,6 +1812,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
functions: Optional[ChatCompletionToolCallFunctionChunk] = None
thinking_blocks: Optional[List[ChatCompletionThinkingBlock]] = None
reasoning_content: Optional[str] = None
thought_signatures: Optional[Any] = None
for idx, candidate in enumerate(_candidates):
if "content" not in candidate:

3
litellm/llms/vertex_ai/vertex_ai_partner_models/main.py
View file

@ -40,6 +40,7 @@ class PartnerModelPrefixes(str, Enum):
GPT_OSS_PREFIX = "openai/gpt-oss-"
MINIMAX_PREFIX = "minimaxai/"
MOONSHOT_PREFIX = "moonshotai/"
ZAI_PREFIX = "zai-org/"
class VertexAIPartnerModels(VertexBase):
@ -66,6 +67,7 @@ class VertexAIPartnerModels(VertexBase):
or model.startswith(PartnerModelPrefixes.GPT_OSS_PREFIX)
or model.startswith(PartnerModelPrefixes.MINIMAX_PREFIX)
or model.startswith(PartnerModelPrefixes.MOONSHOT_PREFIX)
or model.startswith(PartnerModelPrefixes.ZAI_PREFIX)
):
return True
return False
@ -79,6 +81,7 @@ class VertexAIPartnerModels(VertexBase):
PartnerModelPrefixes.GPT_OSS_PREFIX,
PartnerModelPrefixes.MINIMAX_PREFIX,
PartnerModelPrefixes.MOONSHOT_PREFIX,
PartnerModelPrefixes.ZAI_PREFIX,
]
if any(provider in model for provider in OPENAI_LIKE_VERTEX_PROVIDERS):
return True

40
litellm/llms/vertex_ai/vertex_llm_base.py
View file

@ -388,6 +388,10 @@ class VertexBase:
Internal function. Returns the token and url for the call.
Handles logic if it's google ai studio vs. vertex ai.
For Vertex AI:
- If gemini_api_key is provided, use API key authentication (x-goog-api-key header)
- Otherwise, use service account credentials (OAuth2 Bearer token)
Returns
token, url
@ -400,7 +404,7 @@ class VertexBase:
stream=stream,
gemini_api_key=gemini_api_key,
)
auth_header = None # this field is not used for gemin
auth_header = None # this field is not used for gemini
else:
vertex_location = self.get_vertex_region(
vertex_region=vertex_location,
@ -409,14 +413,32 @@ class VertexBase:
### SET RUNTIME ENDPOINT ###
version = "v1beta1" if should_use_v1beta1_features is True else "v1"
url, endpoint = _get_vertex_url(
mode=mode,
model=model,
stream=stream,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_api_version=version,
)
# Check if using API key authentication for Vertex AI
if gemini_api_key and not vertex_credentials:
# When using API key with Vertex AI, use the Google AI Studio endpoint
# This is because Vertex AI API keys work with generativelanguage.googleapis.com
verbose_logger.debug(
f"Using Vertex AI API key authentication for model: {model} - routing to Google AI Studio endpoint"
)
url, endpoint = _get_gemini_url(
mode=mode,
model=model,
stream=stream,
gemini_api_key=gemini_api_key,
)
# API key is already included in the URL by _get_gemini_url
auth_header = None
else:
# Use OAuth2 Bearer token authentication (traditional Vertex AI)
url, endpoint = _get_vertex_url(
mode=mode,
model=model,
stream=stream,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_api_version=version,
)
return self._check_custom_proxy(
api_base=api_base,

48
litellm/llms/watsonx/audio_transcription/transformation.py
View file

@ -7,13 +7,14 @@ WatsonX follows the OpenAI spec for audio transcription.
from typing import Any, Dict, List, Optional
import litellm
from httpx import Response
from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIAudioTranscriptionOptionalParams,
)
from litellm.types.llms.watsonx import WatsonXAudioTranscriptionRequestBody
from litellm.types.utils import FileTypes
from litellm.types.utils import FileTypes, TranscriptionResponse
from ...base_llm.audio_transcription.transformation import (
AudioTranscriptionRequestData,
@ -156,3 +157,48 @@ class IBMWatsonXAudioTranscriptionConfig(
url = f"{url}?version={api_version}"
return url
def transform_audio_transcription_response(
self,
raw_response: Response,
) -> TranscriptionResponse:
"""
Transform the audio transcription response from WatsonX.
WatsonX may include a 'model' field in the response, which needs to be
removed before creating the TranscriptionResponse object.
"""
try:
raw_response_json = raw_response.json()
except Exception as e:
raise ValueError(
f"Error transforming response to json: {str(e)}\nResponse: {raw_response.text}"
)
# Extract only valid fields for TranscriptionResponse.__init__()
# TranscriptionResponse only accepts 'text' and 'usage' in __init__()
text = raw_response_json.get("text")
usage = raw_response_json.get("usage")
# Create response with only valid fields
response_kwargs = {}
if text is not None:
response_kwargs["text"] = text
if usage is not None:
response_kwargs["usage"] = usage
if not response_kwargs:
raise ValueError(
"Invalid response format. Received response does not match the expected format. Got: ",
raw_response_json,
)
response = TranscriptionResponse(**response_kwargs)
# Add other fields using dictionary-style assignment (like duration, task, etc.)
# Skip fields that TranscriptionResponse doesn't accept in __init__()
for key, value in raw_response_json.items():
if key not in ["text", "usage", "model"]: # text/usage already set, model should be excluded
response[key] = value
return response

55
litellm/main.py
View file

@ -189,7 +189,7 @@ from .llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from .llms.custom_llm import CustomLLM, custom_chat_llm_router
from .llms.databricks.embed.handler import DatabricksEmbeddingHandler
from .llms.deprecated_providers import aleph_alpha, palm
from .llms.gemini.common_utils import get_api_key_from_env
from .llms.gemini.common_utils import get_api_key_from_env, get_vertex_api_key_from_env
from .llms.groq.chat.handler import GroqChatCompletion
from .llms.heroku.chat.transformation import HerokuChatConfig
from .llms.huggingface.embedding.handler import HuggingFaceEmbedding
@ -3230,6 +3230,12 @@ def completion( # type: ignore # noqa: PLR0915
or get_secret("VERTEXAI_CREDENTIALS")
)
vertex_api_key = (
api_key
or get_vertex_api_key_from_env()
or litellm.api_key
)
api_base = api_base or litellm.api_base or get_secret("VERTEXAI_API_BASE")
new_params = safe_deep_copy(optional_params or {})
@ -3271,7 +3277,7 @@ def completion( # type: ignore # noqa: PLR0915
vertex_location=vertex_ai_location,
vertex_project=vertex_ai_project,
vertex_credentials=vertex_credentials,
gemini_api_key=None,
gemini_api_key=vertex_api_key, # Support for Vertex AI API Key
logging_obj=logging,
acompletion=acompletion,
timeout=timeout,
@ -4701,6 +4707,51 @@ def embedding( # noqa: PLR0915
litellm_params=litellm_params_dict,
headers=headers,
)
elif custom_llm_provider == "openrouter":
api_base = (
api_base
or litellm.api_base
or get_secret_str("OPENROUTER_API_BASE")
or "https://openrouter.ai/api/v1"
)
api_key = (
api_key
or litellm.api_key
or litellm.openrouter_key
or get_secret("OPENROUTER_API_KEY")
or get_secret("OR_API_KEY")
)
openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai"
openrouter_app_name = get_secret("OR_APP_NAME") or "liteLLM"
openrouter_headers = {
"HTTP-Referer": openrouter_site_url,
"X-Title": openrouter_app_name,
}
_headers = headers or litellm.headers
if _headers:
openrouter_headers.update(_headers)
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 == "huggingface":
api_key = (
api_key

13
litellm/model_prices_and_context_window_backup.json
View file

@ -28345,6 +28345,19 @@
"supports_tool_choice": true,
"supports_web_search": true
},
"vertex_ai/zai-org/glm-4.7-maas": {
"input_cost_per_token": 3e-07,
"litellm_provider": "vertex_ai-zai_models",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"vertex_ai/mistral-medium-3": {
"input_cost_per_token": 4e-07,
"litellm_provider": "vertex_ai-mistral_models",

5
litellm/passthrough/main.py
View file

@ -216,6 +216,11 @@ def llm_passthrough_route(
)
litellm_params_dict = get_litellm_params(**kwargs)
# Add model_id to litellm_params if present in kwargs (for Bedrock Application Inference Profiles)
if "model_id" in kwargs:
litellm_params_dict["model_id"] = kwargs["model_id"]
litellm_logging_obj.update_environment_variables(
model=model,
litellm_params=litellm_params_dict,

42
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -551,6 +551,7 @@ class MCPServerManager:
allowed_tools=getattr(mcp_server, "allowed_tools", None),
disallowed_tools=getattr(mcp_server, "disallowed_tools", None),
allow_all_keys=mcp_server.allow_all_keys,
updated_at=getattr(mcp_server, "updated_at", None),
)
return new_server
@ -697,9 +698,7 @@ class MCPServerManager:
results = await asyncio.gather(*tasks)
# Flatten results into single list
list_tools_result: List[MCPTool] = [
tool for tools in results for tool in tools
]
list_tools_result: List[MCPTool] = [tool for tools in results for tool in tools]
verbose_logger.info(
f"Successfully fetched {len(list_tools_result)} tools total from all servers"
@ -2059,7 +2058,8 @@ class MCPServerManager:
return None
async def _add_mcp_servers_from_db_to_in_memory_registry(self):
async def reload_servers_from_database(self):
"""Re-synchronize the in-memory MCP server registry with the database."""
from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
get_prisma_client_or_throw,
@ -2074,15 +2074,34 @@ class MCPServerManager:
db_mcp_servers = await get_all_mcp_servers(prisma_client)
verbose_logger.info(f"Found {len(db_mcp_servers)} MCP servers in database")
# ensure the global_mcp_server_manager is up to date with the db
previous_registry = self.registry
new_registry: Dict[str, MCPServer] = {}
for server in db_mcp_servers:
existing_server = previous_registry.get(server.server_id)
if (
existing_server is not None
and existing_server.updated_at is not None
and server.updated_at is not None
and existing_server.updated_at == server.updated_at
):
# Re-use existing server instance to avoid re-running build_mcp_server_from_table()
# which can perform network discovery for OAuth2 servers.
new_registry[server.server_id] = existing_server
continue
verbose_logger.debug(
f"Adding server to registry: {server.server_id} ({server.server_name})"
f"Building server from DB: {server.server_id} ({server.server_name})"
)
await self.add_server(server)
new_registry[server.server_id] = await self.build_mcp_server_from_table(
server
)
self.registry = new_registry
verbose_logger.debug(
f"Registry now contains {len(self.get_registry())} servers"
"MCP registry refreshed (%s servers in registry)", len(new_registry)
)
def get_mcp_servers_from_ids(self, server_ids: List[str]) -> List[MCPServer]:
@ -2369,13 +2388,6 @@ class MCPServerManager:
servers.append(self._build_mcp_server_table(server))
return servers
async def reload_servers_from_database(self):
"""
Public method to reload all MCP servers from database into registry.
This can be called from management endpoints to ensure registry is up to date.
"""
await self._add_mcp_servers_from_db_to_in_memory_registry()
async def get_all_mcp_servers_with_health_unfiltered(
self, server_ids: Optional[List[str]] = None
) -> List[LiteLLM_MCPServerTable]:

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

@ -218,10 +218,10 @@ if MCP_AVAILABLE:
from fastapi import HTTPException
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.proxy.proxy_server import add_litellm_data_to_request, proxy_config
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
from litellm.proxy.proxy_server import add_litellm_data_to_request, proxy_config
try:
data = await request.json()
@ -252,7 +252,12 @@ if MCP_AVAILABLE:
if mcp_server_auth_headers:
data["mcp_server_auth_headers"] = mcp_server_auth_headers
data["raw_headers"] = raw_headers_from_request
# Extract user_api_key_auth from metadata and add to top level
# call_mcp_tool expects user_api_key_auth as a top-level parameter
if "metadata" in data and "user_api_key_auth" in data["metadata"]:
data["user_api_key_auth"] = data["metadata"]["user_api_key_auth"]
result = await call_mcp_tool(**data)
return result
except BlockedPiiEntityError as e:

3
litellm/proxy/_types.py
View file

@ -1942,7 +1942,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
description="connect to a postgres db - needed for generating temporary keys + tracking spend / key",
)
database_connection_pool_limit: Optional[int] = Field(
100,
10,
description="default connection pool for prisma client connecting to postgres db",
)
database_connection_timeout: Optional[float] = Field(
@ -2104,6 +2104,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
rotation_interval: Optional[str] = None # How often to rotate (e.g., "30d", "90d")
last_rotation_at: Optional[datetime] = None # When this key was last rotated
key_rotation_at: Optional[datetime] = None # When this key should next be rotated
router_settings: Optional[Dict] = None # Router settings for this key (Key > Team > Global precedence)
model_config = ConfigDict(protected_namespaces=())

4
litellm/proxy/anthropic_endpoints/endpoints.py
View file

@ -11,7 +11,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
create_streaming_response,
create_response,
)
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.types.utils import TokenCountResponse
@ -106,7 +106,7 @@ async def anthropic_response( # noqa: PLR0915
)
)
return await create_streaming_response(
return await create_response(
generator=selected_data_generator,
media_type="text/event-stream",
headers={},

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

@ -147,6 +147,7 @@ async def common_checks(
# 3.1. If organization is in budget
await _organization_max_budget_check(
valid_token=valid_token,
team_object=team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
@ -2310,61 +2311,86 @@ async def _team_max_budget_check(
async def _organization_max_budget_check(
valid_token: Optional[UserAPIKeyAuth],
team_object: Optional[LiteLLM_TeamTable],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
proxy_logging_obj: ProxyLogging,
):
"""
Check if the organization is over its max budget.
This function checks the organization budget using:
1. First, tries to use valid_token.org_id (if key has organization_id set)
2. Falls back to team_object.organization_id (if key doesn't have org_id but team does)
This ensures organization budget checks work even when keys don't have organization_id
set directly, as long as their team belongs to an organization.
Raises:
BudgetExceededError if the organization is over its max budget.
Triggers a budget alert if the organization is over its max budget.
"""
# Only check if token has organization info and organization_max_budget is set
if (
valid_token is None
or valid_token.org_id is None
or valid_token.organization_max_budget is None
or valid_token.organization_max_budget <= 0
):
if valid_token is None or prisma_client is None:
return
# Get organization object to check current spend
if prisma_client is not None:
org_table = await get_org_object(
org_id=valid_token.org_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
# Determine organization_id: first try from token, then fallback to team
org_id: Optional[str] = None
if valid_token.org_id is not None:
org_id = valid_token.org_id
elif team_object is not None and team_object.organization_id is not None:
org_id = team_object.organization_id
# If no organization_id found, skip the check
if org_id is None:
return
# Get organization object with budget table to check current spend and max budget
try:
org_table = await prisma_client.db.litellm_organizationtable.find_unique(
where={"organization_id": org_id},
include={"litellm_budget_table": True},
)
except Exception:
# If organization lookup fails, skip the check
return
if org_table is None:
return
# Get max_budget from organization's budget table
org_max_budget: Optional[float] = None
if org_table.litellm_budget_table is not None:
org_max_budget = org_table.litellm_budget_table.max_budget
# Only check if organization has a valid max_budget set
if org_max_budget is None or org_max_budget <= 0:
return
# Check if organization spend exceeds max budget
if org_table.spend >= org_max_budget:
# Trigger budget alert
call_info = CallInfo(
token=valid_token.token,
spend=org_table.spend,
max_budget=org_max_budget,
user_id=valid_token.user_id,
team_id=valid_token.team_id,
team_alias=valid_token.team_alias,
organization_id=org_id,
event_group=Litellm_EntityType.ORGANIZATION,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="organization_budget",
user_info=call_info,
)
)
if (
org_table is not None
and org_table.spend >= valid_token.organization_max_budget
):
# Trigger budget alert
call_info = CallInfo(
token=valid_token.token,
spend=org_table.spend,
max_budget=valid_token.organization_max_budget,
user_id=valid_token.user_id,
team_id=valid_token.team_id,
team_alias=valid_token.team_alias,
organization_id=valid_token.org_id,
event_group=Litellm_EntityType.ORGANIZATION,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="organization_budget",
user_info=call_info,
)
)
raise litellm.BudgetExceededError(
current_cost=org_table.spend,
max_budget=valid_token.organization_max_budget,
message=f"Budget has been exceeded! Organization={valid_token.org_id} Current cost: {org_table.spend}, Max budget: {valid_token.organization_max_budget}",
)
raise litellm.BudgetExceededError(
current_cost=org_table.spend,
max_budget=org_max_budget,
message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_table.spend}, Max budget: {org_max_budget}",
)
async def _tag_max_budget_check(
@ -2601,4 +2627,4 @@ def _can_object_call_vector_stores(
code=status.HTTP_401_UNAUTHORIZED,
)
return True
return True

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

@ -426,38 +426,65 @@ def get_key_model_rpm_limit(
user_api_key_dict: UserAPIKeyAuth,
) -> Optional[Dict[str, int]]:
"""
Get the model rpm limit for a given api key
- check key metadata
- check key model max budget
- check team metadata
Get the model rpm limit for a given api key.
Priority order (returns first found):
1. Key metadata (model_rpm_limit)
2. Key model_max_budget (rpm_limit per model)
3. Team metadata (model_rpm_limit)
"""
# 1. Check key metadata first (takes priority)
if user_api_key_dict.metadata:
if "model_rpm_limit" in user_api_key_dict.metadata:
return user_api_key_dict.metadata["model_rpm_limit"]
elif user_api_key_dict.model_max_budget:
result = user_api_key_dict.metadata.get("model_rpm_limit")
if result:
return result
# 2. Check model_max_budget
if user_api_key_dict.model_max_budget:
model_rpm_limit: Dict[str, Any] = {}
for model, budget in user_api_key_dict.model_max_budget.items():
if "rpm_limit" in budget and budget["rpm_limit"] is not None:
if isinstance(budget, dict) and budget.get("rpm_limit") is not None:
model_rpm_limit[model] = budget["rpm_limit"]
return model_rpm_limit
elif user_api_key_dict.team_metadata:
if "model_rpm_limit" in user_api_key_dict.team_metadata:
return user_api_key_dict.team_metadata["model_rpm_limit"]
if model_rpm_limit:
return model_rpm_limit
# 3. Fallback to team metadata
if user_api_key_dict.team_metadata:
return user_api_key_dict.team_metadata.get("model_rpm_limit")
return None
def get_key_model_tpm_limit(
user_api_key_dict: UserAPIKeyAuth,
) -> Optional[Dict[str, int]]:
"""
Get the model tpm limit for a given api key.
Priority order (returns first found):
1. Key metadata (model_tpm_limit)
2. Key model_max_budget (tpm_limit per model)
3. Team metadata (model_tpm_limit)
"""
# 1. Check key metadata first (takes priority)
if user_api_key_dict.metadata:
if "model_tpm_limit" in user_api_key_dict.metadata:
return user_api_key_dict.metadata["model_tpm_limit"]
elif user_api_key_dict.model_max_budget:
if "tpm_limit" in user_api_key_dict.model_max_budget:
return user_api_key_dict.model_max_budget["tpm_limit"]
elif user_api_key_dict.team_metadata:
if "model_tpm_limit" in user_api_key_dict.team_metadata:
return user_api_key_dict.team_metadata["model_tpm_limit"]
result = user_api_key_dict.metadata.get("model_tpm_limit")
if result:
return result
# 2. Check model_max_budget (iterate per-model like RPM does)
if user_api_key_dict.model_max_budget:
model_tpm_limit: Dict[str, Any] = {}
for model, budget in user_api_key_dict.model_max_budget.items():
if isinstance(budget, dict) and budget.get("tpm_limit") is not None:
model_tpm_limit[model] = budget["tpm_limit"]
if model_tpm_limit:
return model_tpm_limit
# 3. Fallback to team metadata
if user_api_key_dict.team_metadata:
return user_api_key_dict.team_metadata.get("model_tpm_limit")
return None
@ -469,7 +496,8 @@ def get_model_rate_limit_from_metadata(
if getattr(user_api_key_dict, metadata_accessor_key):
return getattr(user_api_key_dict, metadata_accessor_key).get(rate_limit_key)
return None
def get_team_model_rpm_limit(
user_api_key_dict: UserAPIKeyAuth,
) -> Optional[Dict[str, int]]:

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

@ -188,7 +188,7 @@ async def authenticate_user( # noqa: PLR0915
_user_row = cast(
Optional[LiteLLM_UserTable],
await prisma_client.db.litellm_usertable.find_first(
where={"user_email": {"equals": username}}
where={"user_email": {"equals": username, "mode": "insensitive"}}
),
)

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

@ -138,7 +138,7 @@ def _apply_budget_limits_to_end_user_params(
) -> None:
"""
Helper function to apply budget limits to end user parameters.
Args:
end_user_params: Dictionary to update with budget parameters
budget_info: Budget table object containing limits
@ -146,16 +146,14 @@ def _apply_budget_limits_to_end_user_params(
"""
if budget_info.tpm_limit is not None:
end_user_params["end_user_tpm_limit"] = budget_info.tpm_limit
if budget_info.rpm_limit is not None:
end_user_params["end_user_rpm_limit"] = budget_info.rpm_limit
if budget_info.max_budget is not None:
end_user_params["end_user_max_budget"] = budget_info.max_budget
verbose_proxy_logger.debug(
f"Applied budget limits to end user {end_user_id}"
)
verbose_proxy_logger.debug(f"Applied budget limits to end user {end_user_id}")
async def user_api_key_auth_websocket(websocket: WebSocket):
@ -170,12 +168,10 @@ async def user_api_key_auth_websocket(websocket: WebSocket):
model = query_params.get("model")
async def return_body():
return _realtime_request_body(model)
request.body = return_body # type: ignore
request.body = return_body # type: ignore
authorization = websocket.headers.get("authorization")
# If no Authorization header, try the api-key header
@ -586,7 +582,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
if team_membership is not None
else None
),
team_metadata=team_object.metadata if team_object is not None else None,
team_metadata=team_object.metadata
if team_object is not None
else None,
)
# run through common checks
_ = await common_checks(
@ -669,9 +667,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
route=route,
)
if _end_user_object is not None:
end_user_params["allowed_model_region"] = (
_end_user_object.allowed_model_region
)
end_user_params[
"allowed_model_region"
] = _end_user_object.allowed_model_region
if _end_user_object.litellm_budget_table is not None:
_apply_budget_limits_to_end_user_params(
end_user_params=end_user_params,
@ -753,7 +751,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}",
type=ProxyErrorTypes.expired_key,
code=400,
param=api_key,
param=abbreviate_api_key(api_key=api_key),
)
valid_token = update_valid_token_with_end_user_params(
valid_token=valid_token, end_user_params=end_user_params
@ -994,7 +992,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
# Check 3. Check if user is in their team budget
if valid_token.team_member_spend is not None:
if prisma_client is not None:
_cache_key = f"{valid_token.team_id}_{valid_token.user_id}"
@ -1055,7 +1052,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}",
type=ProxyErrorTypes.expired_key,
code=400,
param=api_key,
param=abbreviate_api_key(api_key=api_key),
)
# Check 4. Token Spend is under budget
@ -1216,8 +1213,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
)
@tracer.wrap()
async def user_api_key_auth(
request: Request,

137
litellm/proxy/common_request_processing.py
View file

@ -17,7 +17,7 @@ from typing import (
import httpx
import orjson
from fastapi import HTTPException, Request, status
from fastapi.responses import Response, StreamingResponse
from fastapi.responses import JSONResponse, Response, StreamingResponse
import litellm
from litellm._logging import verbose_proxy_logger
@ -96,16 +96,55 @@ async def _parse_event_data_for_error(event_line: Union[str, bytes]) -> Optional
return None
async def create_streaming_response(
def _extract_error_from_sse_chunk(event_line: Union[str, bytes]) -> dict:
"""
Extract error dictionary from SSE format chunk.
Args:
event_line: SSE format event line, e.g. "data: {"error": {...}}\n\n"
Returns:
Error dictionary in OpenAI API format
"""
event_line = (
event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line
)
# Default error format
default_error = {
"message": "Unknown error",
"type": "internal_server_error",
"param": None,
"code": "500",
}
if event_line.startswith("data: "):
json_str = event_line[len("data: ") :].strip()
if not json_str or json_str == "[DONE]":
return default_error
try:
data = orjson.loads(json_str)
if isinstance(data, dict) and "error" in data:
error_obj = data["error"]
if isinstance(error_obj, dict):
return error_obj
except (orjson.JSONDecodeError, json.JSONDecodeError):
pass
return default_error
async def create_response(
generator: AsyncGenerator[str, None],
media_type: str,
headers: dict,
default_status_code: int = status.HTTP_200_OK,
) -> StreamingResponse:
) -> Union[StreamingResponse, JSONResponse]:
"""
Creates a StreamingResponse by inspecting the first chunk for an error code.
The entire original generator content is streamed, but the HTTP status code
of the response is set based on the first chunk if it's a recognized error.
Create streaming response, checking if the first chunk is an error.
If the first chunk is an error, return a standard JSON error response.
Otherwise, return StreamingResponse and stream all content.
"""
first_chunk_value: Optional[str] = None
final_status_code = default_status_code
@ -124,9 +163,27 @@ async def create_streaming_response(
first_chunk_value
)
if error_code_from_chunk is not None:
# First chunk is an error, stream hasn't really started yet
# Should return standard JSON error response instead of SSE format
final_status_code = error_code_from_chunk
verbose_proxy_logger.debug(
f"Error detected in first stream chunk. Status code set to: {final_status_code}"
f"Error detected in first stream chunk. Returning JSON error response with status code: {final_status_code}"
)
# Parse error content
error_dict = _extract_error_from_sse_chunk(first_chunk_value)
# Consume and close generator (avoid resource leak)
try:
await generator.aclose()
except Exception:
pass
# Return JSON format error response
return JSONResponse(
status_code=final_status_code,
content={"error": error_dict},
headers=headers,
)
except Exception as e:
verbose_proxy_logger.debug(f"Error parsing first chunk value: {e}")
@ -237,7 +294,11 @@ class ProxyBaseLLMRequestProcessing:
if response_cost is not None:
try:
# Convert response_cost to float if it's a string
cost_value = float(response_cost) if isinstance(response_cost, str) else response_cost
cost_value = (
float(response_cost)
if isinstance(response_cost, str)
else response_cost
)
if cost_value > 0:
updated_spend = current_spend + cost_value
except (ValueError, TypeError):
@ -376,6 +437,16 @@ class ProxyBaseLLMRequestProcessing:
) -> Tuple[dict, LiteLLMLoggingObj]:
start_time = datetime.now() # start before calling guardrail hooks
# Calculate request queue time if arrival_time is available
# Use start_time.timestamp() to avoid extra time.time() call for better performance
proxy_server_request = self.data.get("proxy_server_request", {})
arrival_time = proxy_server_request.get("arrival_time")
queue_time_seconds = None
if arrival_time is not None:
# Convert start_time (datetime) to timestamp for calculation
processing_start_time = start_time.timestamp()
queue_time_seconds = processing_start_time - arrival_time
self.data = await add_litellm_data_to_request(
data=self.data,
request=request,
@ -385,6 +456,19 @@ class ProxyBaseLLMRequestProcessing:
proxy_config=proxy_config,
)
# Store queue time in metadata after add_litellm_data_to_request to ensure it's preserved
if queue_time_seconds is not None:
from litellm.proxy.litellm_pre_call_utils import _get_metadata_variable_name
_metadata_variable_name = _get_metadata_variable_name(request)
if _metadata_variable_name not in self.data:
self.data[_metadata_variable_name] = {}
if not isinstance(self.data[_metadata_variable_name], dict):
self.data[_metadata_variable_name] = {}
self.data[_metadata_variable_name][
"queue_time_seconds"
] = queue_time_seconds
self.data["model"] = (
general_settings.get("completion_model", None) # server default
or user_model # model name passed via cli args
@ -440,6 +524,29 @@ class ProxyBaseLLMRequestProcessing:
user_api_key_dict=user_api_key_dict, data=self.data, call_type=route_type # type: ignore
)
# Apply hierarchical router_settings (Key > Team > Global)
if llm_router is not None and proxy_config is not None:
from litellm.proxy.proxy_server import prisma_client
router_settings = await proxy_config._get_hierarchical_router_settings(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
# If router_settings found (from key, team, or global), apply them
# This ensures key/team settings override global settings
if router_settings is not None and router_settings:
# Get model_list from current router
model_list = llm_router.get_model_list()
if model_list is not None:
# Create user_config with model_list and router_settings
# This creates a per-request router with the hierarchical settings
user_config = {
"model_list": model_list,
**router_settings
}
self.data["user_config"] = user_config
if "messages" in self.data and self.data["messages"]:
logging_obj.update_messages(self.data["messages"])
@ -647,7 +754,7 @@ class ProxyBaseLLMRequestProcessing:
proxy_logging_obj=proxy_logging_obj,
)
)
return await create_streaming_response(
return await create_response(
generator=selected_data_generator,
media_type="text/event-stream",
headers=custom_headers,
@ -658,7 +765,7 @@ class ProxyBaseLLMRequestProcessing:
user_api_key_dict=user_api_key_dict,
request_data=self.data,
)
return await create_streaming_response(
return await create_response(
generator=selected_data_generator,
media_type="text/event-stream",
headers=custom_headers,
@ -900,11 +1007,11 @@ class ProxyBaseLLMRequestProcessing:
@staticmethod
def _get_pre_call_type(
route_type: Literal["acompletion", "aembedding", "aresponses", "allm_passthrough_route"],
) -> Literal["completion", "embeddings", "responses", "allm_passthrough_route"]:
) -> Literal["completion", "embedding", "responses", "allm_passthrough_route"]:
if route_type == "acompletion":
return "completion"
elif route_type == "aembedding":
return "embeddings"
return "embedding"
elif route_type == "aresponses":
return "responses"
elif route_type == "allm_passthrough_route":
@ -1155,9 +1262,9 @@ class ProxyBaseLLMRequestProcessing:
# Add cache-related fields to **params (handled by Usage.__init__)
if cache_creation_input_tokens is not None:
usage_kwargs["cache_creation_input_tokens"] = (
cache_creation_input_tokens
)
usage_kwargs[
"cache_creation_input_tokens"
] = cache_creation_input_tokens
if cache_read_input_tokens is not None:
usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens

309
litellm/proxy/db/prisma_client.py
View file

@ -17,50 +17,141 @@ from litellm.secret_managers.main import str_to_bool
class PrismaWrapper:
"""
Wrapper around Prisma client that handles RDS IAM token authentication.
When iam_token_db_auth is enabled, this wrapper:
1. Proactively refreshes IAM tokens before they expire (background task)
2. Falls back to synchronous refresh if a token is found expired
3. Uses proper locking to prevent race conditions during reconnection
RDS IAM tokens are valid for 15 minutes. This wrapper refreshes them
3 minutes before expiration to ensure uninterrupted database connectivity.
"""
# Buffer time in seconds before token expiration to trigger refresh
# Refresh 3 minutes (180 seconds) before the token expires
TOKEN_REFRESH_BUFFER_SECONDS = 180
# Fallback refresh interval if token parsing fails (10 minutes)
FALLBACK_REFRESH_INTERVAL_SECONDS = 600
def __init__(self, original_prisma: Any, iam_token_db_auth: bool):
self._original_prisma = original_prisma
self.iam_token_db_auth = iam_token_db_auth
# Background token refresh task management
self._token_refresh_task: Optional[asyncio.Task] = None
self._reconnection_lock = asyncio.Lock()
self._last_refresh_time: Optional[datetime] = None
def _extract_token_from_db_url(self, db_url: Optional[str]) -> Optional[str]:
"""
Extract the token (password) from the DATABASE_URL.
The token contains the AWS signature with X-Amz-Date and X-Amz-Expires parameters.
Important: We must parse the URL while it's still encoded to preserve structure,
then decode the password portion. Otherwise the '?' in the token breaks URL parsing.
"""
if db_url is None:
return None
try:
# Parse URL while still encoded to preserve structure
parsed = urllib.parse.urlparse(db_url)
if parsed.password:
# Now decode just the password/token
return urllib.parse.unquote(parsed.password)
return None
except Exception:
return None
def _parse_token_expiration(self, token: Optional[str]) -> Optional[datetime]:
"""
Parse the token to extract its expiration time.
Returns the datetime when the token expires, or None if parsing fails.
"""
if token is None:
return None
try:
# Token format: ...?X-Amz-Date=YYYYMMDDTHHMMSSZ&X-Amz-Expires=900&...
if "?" not in token:
return None
query_string = token.split("?", 1)[1]
params = urllib.parse.parse_qs(query_string)
expires_str = params.get("X-Amz-Expires", [None])[0]
date_str = params.get("X-Amz-Date", [None])[0]
if not expires_str or not date_str:
return None
token_created = datetime.strptime(date_str, "%Y%m%dT%H%M%SZ")
expires_in = int(expires_str)
return token_created + timedelta(seconds=expires_in)
except Exception as e:
verbose_proxy_logger.debug(f"Failed to parse token expiration: {e}")
return None
def _calculate_seconds_until_refresh(self) -> float:
"""
Calculate exactly how many seconds until we need to refresh the token.
Uses precise timing: sleeps until (token_expiration - buffer_seconds).
For a 15-minute (900s) token with 180s buffer, this returns ~720s (12 min).
Returns:
Number of seconds to sleep before the next refresh.
Returns 0 if token should be refreshed immediately.
Returns FALLBACK_REFRESH_INTERVAL_SECONDS if parsing fails.
"""
db_url = os.getenv("DATABASE_URL")
token = self._extract_token_from_db_url(db_url)
expiration_time = self._parse_token_expiration(token)
if expiration_time is None:
# If we can't parse the token, use fallback interval
verbose_proxy_logger.debug(
f"Could not parse token expiration, using fallback interval of "
f"{self.FALLBACK_REFRESH_INTERVAL_SECONDS}s"
)
return self.FALLBACK_REFRESH_INTERVAL_SECONDS
# Calculate when we should refresh (expiration - buffer)
refresh_at = expiration_time - timedelta(
seconds=self.TOKEN_REFRESH_BUFFER_SECONDS
)
# How long until refresh time?
now = datetime.utcnow()
seconds_until_refresh = (refresh_at - now).total_seconds()
# If already past refresh time, return 0 (refresh immediately)
return max(0, seconds_until_refresh)
def is_token_expired(self, token_url: Optional[str]) -> bool:
"""Check if the token in the given URL is expired."""
if token_url is None:
return True
# Decode the token URL to handle URL-encoded characters
decoded_url = urllib.parse.unquote(token_url)
# Parse the token URL
parsed_url = urllib.parse.urlparse(decoded_url)
token = self._extract_token_from_db_url(token_url)
expiration_time = self._parse_token_expiration(token)
# Parse the query parameters from the path component (if they exist there)
query_params = urllib.parse.parse_qs(parsed_url.query)
if expiration_time is None:
# If we can't parse the token, assume it's expired to trigger refresh
verbose_proxy_logger.debug(
"Could not parse token expiration, treating as expired"
)
return True
# Get expiration time from the query parameters
expires = query_params.get("X-Amz-Expires", [None])[0]
if expires is None:
raise ValueError("X-Amz-Expires parameter is missing or invalid.")
expires_int = int(expires)
# Get the token's creation time from the X-Amz-Date parameter
token_time_str = query_params.get("X-Amz-Date", [""])[0]
if not token_time_str:
raise ValueError("X-Amz-Date parameter is missing or invalid.")
# Ensure the token time string is parsed correctly
try:
token_time = datetime.strptime(token_time_str, "%Y%m%dT%H%M%SZ")
except ValueError as e:
raise ValueError(f"Invalid X-Amz-Date format: {e}")
# Calculate the expiration time
expiration_time = token_time + timedelta(seconds=expires_int)
# Current time in UTC
current_time = datetime.utcnow()
# Check if the token is expired
return current_time > expiration_time
return datetime.utcnow() > expiration_time
def get_rds_iam_token(self) -> Optional[str]:
"""Generate a new RDS IAM token and update DATABASE_URL."""
if self.iam_token_db_auth:
from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token
@ -74,7 +165,6 @@ class PrismaWrapper:
db_host=db_host, db_port=db_port, db_user=db_user
)
# print(f"token: {token}")
_db_url = f"postgresql://{db_user}:{token}@{db_host}:{db_port}/{db_name}"
if db_schema:
_db_url += f"?schema={db_schema}"
@ -86,6 +176,7 @@ class PrismaWrapper:
async def recreate_prisma_client(
self, new_db_url: str, http_client: Optional[Any] = None
):
"""Disconnect and reconnect the Prisma client with a new database URL."""
from prisma import Prisma # type: ignore
try:
@ -100,21 +191,159 @@ class PrismaWrapper:
await self._original_prisma.connect()
async def start_token_refresh_task(self) -> None:
"""
Start the background token refresh task.
This task proactively refreshes RDS IAM tokens before they expire,
preventing connection failures. Should be called after the initial
Prisma client connection is established.
"""
if not self.iam_token_db_auth:
verbose_proxy_logger.debug(
"IAM token auth not enabled, skipping token refresh task"
)
return
if self._token_refresh_task is not None:
verbose_proxy_logger.debug("Token refresh task already running")
return
self._token_refresh_task = asyncio.create_task(self._token_refresh_loop())
verbose_proxy_logger.info(
"Started RDS IAM token proactive refresh background task"
)
async def stop_token_refresh_task(self) -> None:
"""
Stop the background token refresh task gracefully.
Should be called during application shutdown to clean up resources.
"""
if self._token_refresh_task is None:
return
self._token_refresh_task.cancel()
try:
await self._token_refresh_task
except asyncio.CancelledError:
pass
self._token_refresh_task = None
verbose_proxy_logger.info("Stopped RDS IAM token refresh background task")
async def _token_refresh_loop(self) -> None:
"""
Background loop that proactively refreshes RDS IAM tokens before expiration.
Uses precise timing: calculates the exact sleep duration until the token
needs to be refreshed (expiration - 3 minute buffer), then refreshes.
This is more efficient than polling, requiring only 1 wake-up per token cycle.
"""
verbose_proxy_logger.info(
f"RDS IAM token refresh loop started. "
f"Tokens will be refreshed {self.TOKEN_REFRESH_BUFFER_SECONDS}s before expiration."
)
while True:
try:
# Calculate exactly how long to sleep until next refresh
sleep_seconds = self._calculate_seconds_until_refresh()
if sleep_seconds > 0:
verbose_proxy_logger.info(
f"RDS IAM token refresh scheduled in {sleep_seconds:.0f} seconds "
f"({sleep_seconds / 60:.1f} minutes)"
)
await asyncio.sleep(sleep_seconds)
# Refresh the token
verbose_proxy_logger.info("Proactively refreshing RDS IAM token...")
await self._safe_refresh_token()
except asyncio.CancelledError:
verbose_proxy_logger.info("RDS IAM token refresh loop cancelled")
break
except Exception as e:
verbose_proxy_logger.error(
f"Error in RDS IAM token refresh loop: {e}. "
f"Retrying in {self.FALLBACK_REFRESH_INTERVAL_SECONDS}s..."
)
# On error, wait before retrying to avoid tight error loops
try:
await asyncio.sleep(self.FALLBACK_REFRESH_INTERVAL_SECONDS)
except asyncio.CancelledError:
break
async def _safe_refresh_token(self) -> None:
"""
Refresh the RDS IAM token with proper locking to prevent race conditions.
Uses an asyncio lock to ensure only one refresh operation happens at a time,
preventing multiple concurrent reconnection attempts.
"""
async with self._reconnection_lock:
new_db_url = self.get_rds_iam_token()
if new_db_url:
await self.recreate_prisma_client(new_db_url)
self._last_refresh_time = datetime.utcnow()
verbose_proxy_logger.info(
"RDS IAM token refreshed successfully. New token valid for ~15 minutes."
)
else:
verbose_proxy_logger.error(
"Failed to generate new RDS IAM token during proactive refresh"
)
def __getattr__(self, name: str):
"""
Proxy attribute access to the underlying Prisma client.
If IAM token auth is enabled and the token is expired, this method
provides a synchronous fallback to refresh the token. However, this
should rarely be needed since the background task proactively refreshes
tokens before they expire.
FIXED: Now properly waits for reconnection to complete before returning,
instead of the previous fire-and-forget pattern that caused the bug.
"""
original_attr = getattr(self._original_prisma, name)
if self.iam_token_db_auth:
db_url = os.getenv("DATABASE_URL")
if self.is_token_expired(db_url):
db_url = self.get_rds_iam_token()
loop = asyncio.get_event_loop()
if db_url:
# Check if token is expired (should be rare if background task is running)
if self.is_token_expired(db_url):
verbose_proxy_logger.warning(
"RDS IAM token expired in __getattr__ - proactive refresh may have failed. "
"Triggering synchronous fallback refresh..."
)
new_db_url = self.get_rds_iam_token()
if new_db_url:
loop = asyncio.get_event_loop()
if loop.is_running():
asyncio.run_coroutine_threadsafe(
self.recreate_prisma_client(db_url), loop
# FIXED: Actually wait for the reconnection to complete!
# The previous code used fire-and-forget which caused the bug.
future = asyncio.run_coroutine_threadsafe(
self.recreate_prisma_client(new_db_url), loop
)
try:
# Wait up to 30 seconds for reconnection
future.result(timeout=30)
verbose_proxy_logger.info(
"Synchronous token refresh completed successfully"
)
except Exception as e:
verbose_proxy_logger.error(
f"Failed to refresh token synchronously: {e}"
)
raise
else:
asyncio.run(self.recreate_prisma_client(db_url))
asyncio.run(self.recreate_prisma_client(new_db_url))
# Get the NEW attribute after reconnection
original_attr = getattr(self._original_prisma, name)
else:
raise ValueError("Failed to get RDS IAM token")

244
litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py
View file

@ -5,6 +5,7 @@
# +-------------------------------------------------------------+
# Qualifire - Evaluate LLM outputs for quality, safety, and reliability
import json
import os
from typing import Any, Dict, List, Literal, Optional, Type
@ -15,12 +16,17 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
from litellm.types.utils import GenericGuardrailAPIInputs
GUARDRAIL_NAME = "qualifire"
DEFAULT_QUALIFIRE_API_BASE = "https://proxy.qualifire.ai"
class QualifireGuardrail(CustomGuardrail):
@ -44,7 +50,7 @@ class QualifireGuardrail(CustomGuardrail):
Args:
api_key: API key for Qualifire (or use QUALIFIRE_API_KEY env var)
api_base: Optional custom API base URL
api_base: Optional custom API base URL (defaults to https://api.qualifire.ai)
evaluation_id: Pre-configured evaluation ID from Qualifire dashboard
prompt_injections: Enable prompt injection detection (default if no other checks)
hallucinations_check: Enable hallucination detection
@ -64,6 +70,7 @@ class QualifireGuardrail(CustomGuardrail):
api_base
or get_secret_str("QUALIFIRE_BASE_URL")
or os.environ.get("QUALIFIRE_BASE_URL")
or DEFAULT_QUALIFIRE_API_BASE
)
self.evaluation_id = evaluation_id
self.prompt_injections = prompt_injections
@ -79,7 +86,11 @@ class QualifireGuardrail(CustomGuardrail):
if not self._has_any_check_enabled() and not self.evaluation_id:
self.prompt_injections = True
self._client = None
# Initialize async HTTP client for direct API calls
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback
)
super().__init__(**kwargs)
def _has_any_check_enabled(self) -> bool:
@ -96,43 +107,22 @@ class QualifireGuardrail(CustomGuardrail):
]
)
def _get_client(self):
"""Lazy initialization of Qualifire client."""
if self._client is None:
try:
from qualifire.client import Client
except ImportError:
raise ImportError(
"qualifire package is required for QualifireGuardrail. "
"Install it with: pip install qualifire"
)
client_kwargs: Dict[str, Any] = {}
if self.qualifire_api_key:
client_kwargs["api_key"] = self.qualifire_api_key
if self.qualifire_api_base:
client_kwargs["base_url"] = self.qualifire_api_base
self._client = Client(**client_kwargs)
return self._client
def _convert_messages_to_qualifire_format(
def _convert_messages_to_api_format(
self, messages: List[AllMessageValues]
) -> List[Any]:
) -> List[Dict[str, Any]]:
"""
Convert LiteLLM messages to Qualifire's LLMMessage format.
Convert LiteLLM messages to Qualifire API format.
Supports tool calls for tool_selection_quality_check.
"""
try:
from qualifire.types import LLMMessage, LLMToolCall
except ImportError:
raise ImportError(
"qualifire package is required for QualifireGuardrail. "
"Install it with: pip install qualifire"
)
qualifire_messages = []
Returns a list of dicts matching the API's ModelInvocationCanonicalMessage schema:
{
"role": "user" | "assistant" | "system" | "tool",
"content": "...",
"tool_call_id": "...", # optional
"tool_calls": [{"id": "...", "name": "...", "arguments": {...}}] # optional
}
"""
api_messages = []
for msg in messages:
role = msg.get("role", "user")
content = msg.get("content", "")
@ -147,42 +137,86 @@ class QualifireGuardrail(CustomGuardrail):
text_parts.append(part)
content = "\n".join(text_parts)
llm_message_kwargs: Dict[str, Any] = {
api_message: Dict[str, Any] = {
"role": role,
"content": content if isinstance(content, str) else str(content),
}
# Handle tool_call_id for tool response messages
tool_call_id = msg.get("tool_call_id")
if tool_call_id:
api_message["tool_call_id"] = tool_call_id
# Handle tool calls if present
tool_calls = msg.get("tool_calls")
if tool_calls and isinstance(tool_calls, list):
qualifire_tool_calls = []
api_tool_calls = []
for tc in tool_calls:
if isinstance(tc, dict):
function_info = tc.get("function", {})
# Arguments can be a string (JSON) or dict
args = function_info.get("arguments", {})
if isinstance(args, str):
import json
try:
args = json.loads(args)
except json.JSONDecodeError:
args = {}
qualifire_tool_calls.append(
LLMToolCall(
id=tc.get("id") or "",
name=function_info.get("name") or "",
arguments=args if isinstance(args, dict) else {},
)
api_tool_calls.append(
{
"id": tc.get("id") or "",
"name": function_info.get("name") or "",
"arguments": args if isinstance(args, dict) else {},
}
)
if qualifire_tool_calls:
llm_message_kwargs["tool_calls"] = qualifire_tool_calls
if api_tool_calls:
api_message["tool_calls"] = api_tool_calls
qualifire_messages.append(LLMMessage(**llm_message_kwargs))
api_messages.append(api_message)
return qualifire_messages
return api_messages
def _check_if_flagged(self, result: Any) -> bool:
def _convert_tools_to_api_format(
self, tools: Optional[List[Any]]
) -> Optional[List[Dict[str, Any]]]:
"""
Convert OpenAI-format tools to Qualifire API format.
Returns a list of dicts matching the API's ModelInvocationToolDefinition schema:
{
"name": "...",
"description": "...",
"parameters": {...}
}
"""
if not tools:
return None
api_tools = []
for tool in tools:
if isinstance(tool, dict):
# Handle OpenAI function tool format
if tool.get("type") == "function":
function_def = tool.get("function", {})
api_tools.append(
{
"name": function_def.get("name", ""),
"description": function_def.get("description", ""),
"parameters": function_def.get("parameters", {}),
}
)
# Handle direct tool format
elif "name" in tool:
api_tools.append(
{
"name": tool.get("name", ""),
"description": tool.get("description", ""),
"parameters": tool.get("parameters", {}),
}
)
return api_tools if api_tools else None
def _check_if_flagged(self, result: Dict[str, Any]) -> bool:
"""
Check if the Qualifire evaluation result indicates flagged content.
@ -190,65 +224,53 @@ class QualifireGuardrail(CustomGuardrail):
A high score (close to 100) indicates GOOD content, low score indicates problems.
"""
# Check evaluation results for any flagged items
evaluation_results = getattr(result, "evaluationResults", None) or []
if isinstance(result, dict):
evaluation_results = result.get("evaluationResults", []) or []
evaluation_results = result.get("evaluationResults", []) or []
for eval_result in evaluation_results:
results: List[Any] = []
if isinstance(eval_result, dict):
results = eval_result.get("results", []) or []
else:
results = getattr(eval_result, "results", []) or []
results = eval_result.get("results", []) or []
for r in results:
flagged = (
r.get("flagged")
if isinstance(r, dict)
else getattr(r, "flagged", False)
)
if flagged:
if r.get("flagged"):
return True
return False
def _build_evaluate_kwargs(
def _build_evaluate_payload(
self,
qualifire_messages: List[Any],
api_messages: List[Dict[str, Any]],
output: Optional[str],
assertions: Optional[List[str]],
available_tools: Optional[List[Any]],
available_tools: Optional[List[Dict[str, Any]]],
) -> Dict[str, Any]:
"""Build kwargs dictionary for the evaluate call."""
kwargs: Dict[str, Any] = {"messages": qualifire_messages}
"""Build payload dictionary for the /api/evaluation/evaluate endpoint."""
payload: Dict[str, Any] = {"messages": api_messages}
if output is not None:
kwargs["output"] = output
payload["output"] = output
# Add enabled checks
if self.prompt_injections:
kwargs["prompt_injections"] = True
payload["prompt_injections"] = True
if self.hallucinations_check:
kwargs["hallucinations_check"] = True
payload["hallucinations_check"] = True
if self.grounding_check:
kwargs["grounding_check"] = True
payload["grounding_check"] = True
if self.pii_check:
kwargs["pii_check"] = True
payload["pii_check"] = True
if self.content_moderation_check:
kwargs["content_moderation_check"] = True
payload["content_moderation_check"] = True
if self.tool_selection_quality_check:
# Only enable tool_selection_quality_check if available_tools is provided
if available_tools:
kwargs["tool_selection_quality_check"] = True
kwargs["available_tools"] = available_tools
payload["tool_selection_quality_check"] = True
payload["available_tools"] = available_tools
else:
verbose_proxy_logger.debug(
"Qualifire Guardrail: tool_selection_quality_check enabled but no available_tools provided, skipping this check"
)
if assertions:
kwargs["assertions"] = assertions
payload["assertions"] = assertions
return kwargs
return payload
async def _run_qualifire_check(
self,
@ -274,11 +296,17 @@ class QualifireGuardrail(CustomGuardrail):
assertions = dynamic_params.get("assertions") or self.assertions
on_flagged = dynamic_params.get("on_flagged") or self.on_flagged
try:
client = self._get_client()
qualifire_messages = self._convert_messages_to_qualifire_format(messages)
# Prepare headers
headers = {
"X-Qualifire-API-Key": self.qualifire_api_key or "",
"Content-Type": "application/json",
}
# Use invoke_evaluation if evaluation_id is provided
try:
# Convert messages to API format
api_messages = self._convert_messages_to_api_format(messages)
# Use invoke endpoint if evaluation_id is provided
if evaluation_id:
# For invoke_evaluation, we need to extract input/output
input_text = ""
@ -291,25 +319,47 @@ class QualifireGuardrail(CustomGuardrail):
input_text = content
break
result = client.invoke_evaluation(
evaluation_id=evaluation_id,
input=input_text,
output=output or "",
)
payload = {
"evaluation_id": evaluation_id,
"input": input_text,
"output": output or "",
"messages": api_messages,
}
# Convert tools if provided
api_tools = self._convert_tools_to_api_format(available_tools)
if api_tools:
payload["available_tools"] = api_tools
url = f"{self.qualifire_api_base}/api/evaluation/invoke"
else:
# Use evaluate with individual checks
kwargs = self._build_evaluate_kwargs(
qualifire_messages=qualifire_messages,
# Use evaluate endpoint with individual checks
api_tools = self._convert_tools_to_api_format(available_tools)
payload = self._build_evaluate_payload(
api_messages=api_messages,
output=output,
assertions=assertions,
available_tools=available_tools,
available_tools=api_tools,
)
result = client.evaluate(**kwargs)
url = f"{self.qualifire_api_base}/api/evaluation/evaluate"
# Convert result to dict for logging
verbose_proxy_logger.debug(
f"Qualifire Guardrail: Making request to {url}"
)
# Make the API request
response = await self.async_handler.post(
url=url,
headers=headers,
json=payload,
)
response.raise_for_status()
result = response.json()
# Extract response info for logging
qualifire_response = {
"score": getattr(result, "score", None),
"status": getattr(result, "status", None),
"score": result.get("score"),
"status": result.get("status"),
}
verbose_proxy_logger.debug(

4
litellm/proxy/hooks/litellm_skills/main.py
View file

@ -336,8 +336,8 @@ class SkillsInjectionHook(CustomLogger):
)
# Check if code execution is enabled for this request
litellm_metadata = request_data.get("litellm_metadata", {})
metadata = request_data.get("metadata", {})
litellm_metadata = request_data.get("litellm_metadata") or {}
metadata = request_data.get("metadata") or {}
code_exec_enabled = (
litellm_metadata.get("_litellm_code_execution_enabled") or

84
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -167,7 +167,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self.token_increment_script = None
self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60))
# Batch rate limiter (lazy loaded)
self._batch_rate_limiter: Optional[Any] = None
@ -1013,7 +1013,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
# Fail safe: enforce limits if we can't check
return True
def get_rate_limiter_for_call_type(self, call_type: str) -> Optional[Any]:
"""Get the rate limiter for the call type."""
if call_type == "acreate_batch":
@ -1095,9 +1095,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
now = self._get_current_time().timestamp()
reset_time = now + self.window_size
reset_time_formatted = datetime.fromtimestamp(
reset_time
).strftime("%Y-%m-%d %H:%M:%S UTC")
reset_time_formatted = datetime.fromtimestamp(reset_time).strftime(
"%Y-%m-%d %H:%M:%S UTC"
)
remaining_display = max(0, status["limit_remaining"])
rate_limit_type = status["rate_limit_type"]
@ -1137,7 +1137,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# Check if the call type has a specific rate limiter
# eg. for Batch APIs we need to use the batch rate limiter to read the input file and count the tokens and requests
#########################################################
call_type_specific_rate_limiter = self.get_rate_limiter_for_call_type(call_type=call_type)
call_type_specific_rate_limiter = self.get_rate_limiter_for_call_type(
call_type=call_type
)
if call_type_specific_rate_limiter:
return await call_type_specific_rate_limiter.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
@ -1233,26 +1235,58 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return pipeline_operations
def _get_total_tokens_from_usage(self, usage: Any | None, rate_limit_type: Literal["output", "input", "total"]) -> int:
# Get total tokens from response
def _get_total_tokens_from_usage(
self, usage: Any | None, rate_limit_type: Literal["output", "input", "total"]
) -> int:
"""
Get total tokens from response usage for rate limiting.
For 'input' and 'total' rate limit types, cached tokens are excluded
because providers like AWS Bedrock don't count cached tokens toward
rate limits. This aligns LiteLLM's TPM calculation with provider behavior.
"""
total_tokens = 0
# spot fix for /responses api
cached_tokens = 0
if usage:
if isinstance(usage, Usage):
if rate_limit_type == "output":
total_tokens = usage.completion_tokens
total_tokens = usage.completion_tokens or 0
elif rate_limit_type == "input":
total_tokens = usage.prompt_tokens
total_tokens = usage.prompt_tokens or 0
elif rate_limit_type == "total":
total_tokens = usage.total_tokens
total_tokens = usage.total_tokens or 0
# Get cached tokens to exclude from input/total
if rate_limit_type in ("input", "total"):
if (
hasattr(usage, "prompt_tokens_details")
and usage.prompt_tokens_details is not None
):
cached_tokens = (
getattr(usage.prompt_tokens_details, "cached_tokens", 0)
or 0
)
elif isinstance(usage, dict):
# Responses API usage comes as a dict in ResponsesAPIResponse
# Responses API usage comes as a dict
if rate_limit_type == "output":
total_tokens = usage.get("completion_tokens", 0)
total_tokens = usage.get("completion_tokens", 0) or 0
elif rate_limit_type == "input":
total_tokens = usage.get("prompt_tokens", 0)
total_tokens = usage.get("prompt_tokens", 0) or 0
elif rate_limit_type == "total":
total_tokens = usage.get("total_tokens", 0)
total_tokens = usage.get("total_tokens", 0) or 0
# Get cached tokens from dict
if rate_limit_type in ("input", "total"):
prompt_details = usage.get("prompt_tokens_details") or {}
if isinstance(prompt_details, dict):
cached_tokens = prompt_details.get("cached_tokens", 0) or 0
# Subtract cached tokens for input/total (providers don't count them)
if cached_tokens > 0:
total_tokens = max(0, total_tokens - cached_tokens)
return total_tokens
async def _execute_token_increment_script(
@ -1336,6 +1370,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
def get_rate_limit_type(self) -> Literal["output", "input", "total"]:
from litellm.proxy.proxy_server import general_settings
specified_rate_limit_type = general_settings.get(
"token_rate_limit_type", "total"
)
@ -1381,9 +1416,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
user_api_key_organization_id = standard_logging_metadata.get(
"user_api_key_org_id"
)
user_api_key_end_user_id = kwargs.get("user") or standard_logging_metadata.get(
"user_api_key_end_user_id"
)
user_api_key_end_user_id = kwargs.get(
"user"
) or standard_logging_metadata.get("user_api_key_end_user_id")
model_group = get_model_group_from_litellm_kwargs(kwargs)
# Get total tokens from response
@ -1393,7 +1428,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
response_obj, BaseLiteLLMOpenAIResponseObject
):
_usage = getattr(response_obj, "usage", None)
total_tokens = self._get_total_tokens_from_usage(usage=_usage, rate_limit_type=rate_limit_type)
total_tokens = self._get_total_tokens_from_usage(
usage=_usage, rate_limit_type=rate_limit_type
)
# Create pipeline operations for TPM increments
pipeline_operations: List[RedisPipelineIncrementOperation] = []
@ -1518,9 +1555,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
from litellm.types.caching import RedisPipelineIncrementOperation
try:
litellm_parent_otel_span: Union[
Span, None
] = _get_parent_otel_span_from_kwargs(kwargs)
litellm_parent_otel_span: Union[Span, None] = (
_get_parent_otel_span_from_kwargs(kwargs)
)
# Get metadata from standard_logging_object - this correctly handles both
# 'metadata' and 'litellm_metadata' fields from litellm_params
standard_logging_object = kwargs.get("standard_logging_object") or {}
@ -1555,7 +1592,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
f"Error in rate limit failure event: {str(e)}"
)
async def async_post_call_success_hook(
self, data: dict, user_api_key_dict: UserAPIKeyAuth, response
):

33
litellm/proxy/litellm_pre_call_utils.py
View file

@ -161,7 +161,6 @@ class KeyAndTeamLoggingSettings:
@staticmethod
def get_team_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth):
if (
user_api_key_dict.team_metadata is not None
and "logging" in user_api_key_dict.team_metadata
@ -174,12 +173,12 @@ def _get_dynamic_logging_metadata(
user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig
) -> Optional[TeamCallbackMetadata]:
callback_settings_obj: Optional[TeamCallbackMetadata] = None
key_dynamic_logging_settings: Optional[dict] = (
KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)
)
team_dynamic_logging_settings: Optional[dict] = (
KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict)
)
key_dynamic_logging_settings: Optional[
dict
] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)
team_dynamic_logging_settings: Optional[
dict
] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict)
#########################################################################################
# Key-based callbacks
#########################################################################################
@ -462,7 +461,6 @@ class LiteLLMProxyRequestSetup:
team_id=user_api_key_dict.team_id,
) # handles aliases, wildcards, etc.
):
_headers = LiteLLMProxyRequestSetup.add_headers_to_llm_call(
headers, user_api_key_dict
)
@ -663,11 +661,11 @@ class LiteLLMProxyRequestSetup:
## KEY-LEVEL SPEND LOGS / TAGS
if "tags" in key_metadata and key_metadata["tags"] is not None:
data[_metadata_variable_name]["tags"] = (
LiteLLMProxyRequestSetup._merge_tags(
request_tags=data[_metadata_variable_name].get("tags"),
tags_to_add=key_metadata["tags"],
)
data[_metadata_variable_name][
"tags"
] = LiteLLMProxyRequestSetup._merge_tags(
request_tags=data[_metadata_variable_name].get("tags"),
tags_to_add=key_metadata["tags"],
)
if "disable_global_guardrails" in key_metadata and isinstance(
key_metadata["disable_global_guardrails"], bool
@ -815,11 +813,14 @@ async def add_litellm_data_to_request( # noqa: PLR0915
# Init - Proxy Server Request
# we do this as soon as entering so we track the original request
##########################################################
# Track arrival time for queue time metric
arrival_time = time.time()
data["proxy_server_request"] = {
"url": str(request.url),
"method": request.method,
"headers": _headers,
"body": copy.copy(data), # use copy instead of deepcopy
"arrival_time": arrival_time, # Track when request arrived at proxy
}
safe_add_api_version_from_query_params(data, request)
@ -930,9 +931,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915
data[_metadata_variable_name]["litellm_api_version"] = version
if general_settings is not None:
data[_metadata_variable_name]["global_max_parallel_requests"] = (
general_settings.get("global_max_parallel_requests", None)
)
data[_metadata_variable_name][
"global_max_parallel_requests"
] = general_settings.get("global_max_parallel_requests", None)
### KEY-LEVEL Controls
key_metadata = user_api_key_dict.metadata

109
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -36,12 +36,19 @@ from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import LITELLM_PROXY_ADMIN_NAME
from litellm.proxy._experimental.mcp_server.utils import (
get_server_prefix,
validate_and_normalize_mcp_server_payload,
)
router = APIRouter(prefix="/v1/mcp", tags=["mcp"])
MCP_AVAILABLE: bool = True
TEMPORARY_MCP_SERVER_TTL_SECONDS = 300
DEFAULT_MCP_REGISTRY_VERSION = "1.0.0"
LITELLM_MCP_SERVER_NAME = "litellm-mcp-server"
LITELLM_MCP_SERVER_DESCRIPTION = "MCP Server for LiteLLM"
try:
importlib.import_module("mcp")
except ImportError as e:
@ -57,6 +64,7 @@ if MCP_AVAILABLE:
update_mcp_server,
)
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
get_request_base_url,
authorize_with_server,
exchange_token_with_server,
register_client_with_server,
@ -89,6 +97,66 @@ if MCP_AVAILABLE:
server: MCPServer
expires_at: datetime
def _is_public_registry_enabled() -> bool:
from litellm.proxy.proxy_server import (
general_settings as proxy_general_settings,
)
return bool(proxy_general_settings.get("enable_mcp_registry"))
def _build_registry_remote_url(base_url: str, path: str) -> str:
normalized_base = base_url.rstrip("/")
normalized_path = path if path.startswith("/") else f"/{path}"
return f"{normalized_base}{normalized_path}"
def _build_mcp_registry_server_name(server: MCPServer) -> str:
if server.alias:
return server.alias
if server.server_name:
return server.server_name
return server.server_id
def _build_mcp_registry_entry_for_server(
server: MCPServer, base_url: str
) -> Dict[str, Any]:
server_name = _build_mcp_registry_server_name(server)
title = server_name
description = server_name
version = DEFAULT_MCP_REGISTRY_VERSION
server_prefix = get_server_prefix(server)
if not server_prefix:
raise ValueError("MCP server prefix is missing")
remote_url = _build_registry_remote_url(base_url, f"/{server_prefix}/mcp")
return {
"name": server_name,
"title": title,
"description": description,
"version": version,
"remotes": [
{
"type": "streamable-http",
"url": remote_url,
}
],
}
def _build_builtin_registry_entry(base_url: str) -> Dict[str, Any]:
remote_url = _build_registry_remote_url(base_url, "/mcp")
return {
"name": LITELLM_MCP_SERVER_NAME,
"title": LITELLM_MCP_SERVER_NAME,
"description": LITELLM_MCP_SERVER_DESCRIPTION,
"version": DEFAULT_MCP_REGISTRY_VERSION,
"remotes": [
{
"type": "streamable-http",
"url": remote_url,
}
],
}
_temporary_mcp_servers: Dict[str, _TemporaryMCPServerEntry] = {}
def _prune_expired_temporary_mcp_servers() -> None:
@ -302,15 +370,42 @@ if MCP_AVAILABLE:
access_groups_list = sorted(list(access_groups))
return {"access_groups": access_groups_list}
@router.get(
"/registry.json",
tags=["mcp"],
description="MCP registry endpoint. Spec: https://github.com/modelcontextprotocol/registry",
)
async def get_mcp_registry(request: Request):
if not _is_public_registry_enabled():
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="MCP registry is not enabled",
)
base_url = get_request_base_url(request)
registry_servers: List[Dict[str, Any]] = []
registry_servers.append({"server": _build_builtin_registry_entry(base_url)})
registered_servers = list(global_mcp_server_manager.get_registry().values())
registered_servers.sort(key=_build_mcp_registry_server_name)
for server in registered_servers:
try:
entry = _build_mcp_registry_entry_for_server(server, base_url)
except Exception as e:
verbose_proxy_logger.debug(
f"Skipping MCP server {getattr(server, 'server_id', 'unknown')} in registry: {e}"
)
continue
registry_servers.append({"server": entry})
return {"servers": registry_servers}
## FastAPI Routes
def _get_user_mcp_management_mode() -> UserMCPManagementMode:
proxy_general_settings: dict = {}
try:
from litellm.proxy.proxy_server import (
general_settings as proxy_general_settings,
)
except Exception:
pass
from litellm.proxy.proxy_server import (
general_settings as proxy_general_settings,
)
mode = proxy_general_settings.get("user_mcp_management_mode")
if mode == "view_all":

60
litellm/proxy/management_endpoints/router_settings_endpoints.py
View file

@ -4,6 +4,7 @@ ROUTER SETTINGS MANAGEMENT
Endpoints for accessing router configuration and metadata
GET /router/settings - Get router configuration including available routing strategies
GET /router/fields - Get router settings field definitions without values (for UI rendering)
"""
import inspect
@ -37,6 +38,15 @@ class RouterSettingsResponse(BaseModel):
)
class RouterFieldsResponse(BaseModel):
fields: List[RouterSettingsField] = Field(
description="List of all configurable router settings with metadata (without field values)"
)
routing_strategy_descriptions: Dict[str, str] = Field(
description="Descriptions for each routing strategy option"
)
def _get_routing_strategies_from_router_class() -> List[str]:
"""
Dynamically extract routing strategies from the Router class __init__ method.
@ -120,3 +130,53 @@ async def get_router_settings(
)
raise
@router.get(
"/router/fields",
tags=["Router Settings"],
dependencies=[Depends(user_api_key_auth)],
response_model=RouterFieldsResponse,
)
async def get_router_fields(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get router settings field definitions without values.
Returns only the field metadata (type, description, default, options) without
populating field_value. This is useful for UI components that need to know
what fields to render, but will get the actual values from a different endpoint.
Returns:
- fields: List of all configurable router settings with their metadata (type, description, default, options)
The routing_strategy field includes available options extracted from the Router class
Note: field_value will be None for all fields
- routing_strategy_descriptions: Descriptions for each routing strategy option
"""
try:
# Get available routing strategies dynamically from Router class
available_routing_strategies = _get_routing_strategies_from_router_class()
# Get router settings fields from types file
router_fields = [field.model_copy(deep=True) for field in ROUTER_SETTINGS_FIELDS]
# Populate routing_strategy field with available options
for field in router_fields:
if field.field_name == "routing_strategy":
field.options = available_routing_strategies
break
# Ensure field_value is None for all fields (don't populate values)
for field in router_fields:
field.field_value = None
return RouterFieldsResponse(
fields=router_fields,
routing_strategy_descriptions=ROUTING_STRATEGY_DESCRIPTIONS,
)
except Exception as e:
verbose_proxy_logger.error(
f"Error fetching router fields: {str(e)}"
)
raise

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

@ -776,6 +776,7 @@ async def handle_bedrock_count_tokens(
- /v1/messages/count_tokens
- /v1/messages/count-tokens
"""
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
from litellm.proxy.proxy_server import llm_router
@ -822,6 +823,12 @@ async def handle_bedrock_count_tokens(
return result
except BedrockError as e:
# Convert BedrockError to HTTPException for FastAPI
verbose_proxy_logger.error(f"BedrockError in handle_bedrock_count_tokens: {str(e)}")
raise HTTPException(
status_code=e.status_code, detail={"error": e.message}
)
except HTTPException:
# Re-raise HTTP exceptions as-is
raise

240
litellm/proxy/proxy_server.py
View file

@ -229,7 +229,7 @@ from litellm.proxy.batches_endpoints.endpoints import router as batches_router
from litellm.proxy.caching_routes import router as caching_router
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
create_streaming_response,
create_response,
)
from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy
from litellm.proxy.common_utils.debug_utils import init_verbose_loggers
@ -533,9 +533,9 @@ except ImportError:
server_root_path = os.getenv("SERVER_ROOT_PATH", "")
_license_check = LicenseCheck()
premium_user: bool = _license_check.is_premium()
premium_user_data: Optional["EnterpriseLicenseData"] = (
_license_check.airgapped_license_data
)
premium_user_data: Optional[
"EnterpriseLicenseData"
] = _license_check.airgapped_license_data
global_max_parallel_request_retries_env: Optional[str] = os.getenv(
"LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES"
)
@ -658,7 +658,7 @@ async def _initialize_shared_aiohttp_session():
@asynccontextmanager
async def proxy_startup_event(app: FastAPI):
async def proxy_startup_event(app: FastAPI): # noqa: PLR0915
global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, premium_user, _license_check, proxy_batch_polling_interval, shared_aiohttp_session
import json
@ -788,6 +788,17 @@ async def proxy_startup_event(app: FastAPI):
except Exception as e:
verbose_proxy_logger.error(f"Error closing shared aiohttp session: {e}")
# Shutdown event - stop RDS IAM token refresh background task
if (
prisma_client is not None
and hasattr(prisma_client, "db")
and hasattr(prisma_client.db, "stop_token_refresh_task")
):
try:
await prisma_client.db.stop_token_refresh_task()
except Exception as e:
verbose_proxy_logger.error(f"Error stopping token refresh task: {e}")
await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues]
@ -1083,9 +1094,7 @@ try:
# In non-root Docker, we restructure in /var/lib/litellm/ui.
try:
_restructure_ui_html_files(ui_path)
verbose_proxy_logger.info(
f"Restructured UI directory: {ui_path}"
)
verbose_proxy_logger.info(f"Restructured UI directory: {ui_path}")
except PermissionError as e:
verbose_proxy_logger.exception(
f"Permission error while restructuring UI directory {ui_path}: {e}"
@ -1171,9 +1180,9 @@ master_key: Optional[str] = None
config_agents: Optional[List[AgentConfig]] = None
otel_logging = False
prisma_client: Optional[PrismaClient] = None
shared_aiohttp_session: Optional["ClientSession"] = (
None # Global shared session for connection reuse
)
shared_aiohttp_session: Optional[
"ClientSession"
] = None # Global shared session for connection reuse
user_api_key_cache = DualCache(
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
)
@ -1181,9 +1190,9 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(
dual_cache=user_api_key_cache
)
litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter)
redis_usage_cache: Optional[RedisCache] = (
None # redis cache used for tracking spend, tpm/rpm limits
)
redis_usage_cache: Optional[
RedisCache
] = None # redis cache used for tracking spend, tpm/rpm limits
polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False
polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache
user_custom_auth = None
@ -1522,9 +1531,9 @@ async def update_cache( # noqa: PLR0915
_id = "team_id:{}".format(team_id)
try:
# Fetch the existing cost for the given user
existing_spend_obj: Optional[LiteLLM_TeamTable] = (
await user_api_key_cache.async_get_cache(key=_id)
)
existing_spend_obj: Optional[
LiteLLM_TeamTable
] = await user_api_key_cache.async_get_cache(key=_id)
if existing_spend_obj is None:
# do nothing if team not in api key cache
return
@ -1876,7 +1885,6 @@ class ProxyConfig:
"environment_variables" in config_to_save
and config_to_save["environment_variables"]
):
# decrypt the environment_variables - in case a caller function has already encrypted the environment_variables
decrypted_env_vars = self._decrypt_and_set_db_env_variables(
environment_variables=config_to_save["environment_variables"],
@ -2794,21 +2802,21 @@ class ProxyConfig:
verbose_proxy_logger.debug(f"_alerting_callbacks: {general_settings}")
if _alerting_callbacks is None:
return
# Ensure proxy_logging_obj.alerting is set for all alerting types
_alerting_value = general_settings.get("alerting", None)
verbose_proxy_logger.debug(f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}")
verbose_proxy_logger.debug(
f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}"
)
proxy_logging_obj.update_values(
alerting=_alerting_value,
alerting_threshold=general_settings.get("alerting_threshold", 600),
alert_types=general_settings.get("alert_types", None),
alert_to_webhook_url=general_settings.get(
"alert_to_webhook_url", None
),
alert_to_webhook_url=general_settings.get("alert_to_webhook_url", None),
alerting_args=general_settings.get("alerting_args", None),
redis_cache=redis_usage_cache,
)
for _alert in _alerting_callbacks:
if _alert == "slack":
# [OLD] v0 implementation - already handled by update_values above
@ -3222,6 +3230,84 @@ class ProxyConfig:
decrypted_variables[k] = decrypted_value
return decrypted_variables
async def _get_hierarchical_router_settings(
self,
user_api_key_dict: Optional["UserAPIKeyAuth"],
prisma_client: Optional[PrismaClient],
) -> Optional[dict]:
"""
Get router_settings in priority order: Key > Team > Global
Returns:
dict: Combined router_settings, or None if no settings found
"""
if prisma_client is None:
return None
import json
import yaml
# 1. Try key-level router_settings
if user_api_key_dict is not None:
# Check if router_settings is available on the key object
key_router_settings_value = getattr(user_api_key_dict, "router_settings", None)
if key_router_settings_value is not None:
key_router_settings = None
if isinstance(key_router_settings_value, str):
try:
key_router_settings = yaml.safe_load(key_router_settings_value)
except (yaml.YAMLError, json.JSONDecodeError):
try:
key_router_settings = json.loads(key_router_settings_value)
except json.JSONDecodeError:
pass
elif isinstance(key_router_settings_value, dict):
key_router_settings = key_router_settings_value
# If key has router_settings (non-empty dict), use it
if key_router_settings is not None and isinstance(key_router_settings, dict) and key_router_settings:
return key_router_settings
# 2. Try team-level router_settings
if user_api_key_dict is not None and user_api_key_dict.team_id is not None:
try:
team_obj = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": user_api_key_dict.team_id}
)
if team_obj is not None:
team_router_settings_value = getattr(team_obj, "router_settings", None)
if team_router_settings_value is not None:
team_router_settings = None
if isinstance(team_router_settings_value, str):
try:
team_router_settings = yaml.safe_load(team_router_settings_value)
except (yaml.YAMLError, json.JSONDecodeError):
try:
team_router_settings = json.loads(team_router_settings_value)
except json.JSONDecodeError:
pass
elif isinstance(team_router_settings_value, dict):
team_router_settings = team_router_settings_value
# If team has router_settings (non-empty dict), use it
if team_router_settings is not None and isinstance(team_router_settings, dict) and team_router_settings:
return team_router_settings
except Exception:
# If team lookup fails, continue to global settings
pass
# 3. Try global router_settings
try:
db_router_settings = await prisma_client.db.litellm_config.find_first(
where={"param_name": "router_settings"}
)
if db_router_settings is not None and isinstance(db_router_settings.param_value, dict) and db_router_settings.param_value:
return db_router_settings.param_value
except Exception:
pass
return None
async def _add_router_settings_from_db_config(
self,
config_data: dict,
@ -3279,7 +3365,7 @@ class ProxyConfig:
proxy_logging_obj: ProxyLogging
"""
_general_settings = config_data.get("general_settings", {})
if _general_settings is not None and "alerting" in _general_settings:
if (
general_settings is not None
@ -3294,7 +3380,8 @@ class ProxyConfig:
_merged_alerting = list(_yaml_alerting.union(_db_alerting))
# Preserve order: YAML values first, then DB values
_merged_alerting = list(general_settings["alerting"]) + [
item for item in _general_settings["alerting"]
item
for item in _general_settings["alerting"]
if item not in general_settings["alerting"]
]
verbose_proxy_logger.debug(
@ -3402,8 +3489,8 @@ class ProxyConfig:
def _deep_merge_dicts(dst: dict, src: dict) -> None:
"""
Deep-merge src into dst, skipping None values from src.
On conflicts, src (DB) wins.
Deep-merge src into dst, skipping None values and empty lists from src.
On conflicts, src (DB) wins, but empty lists are treated as "no value" and don't overwrite.
"""
stack = [(dst, src)]
while stack:
@ -3412,6 +3499,9 @@ class ProxyConfig:
if v is None:
# Preserve existing config when DB value is None (matches prior behavior)
continue
# Skip empty lists - treat them as "no value" to preserve file config
if isinstance(v, list) and len(v) == 0:
continue
if isinstance(v, dict) and isinstance(d.get(k), dict):
stack.append((d[k], v))
else:
@ -3605,7 +3695,6 @@ class ProxyConfig:
await self._init_vector_stores_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="vector_store_indexes"):
await self._init_vector_store_indexes_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="mcp"):
@ -3804,10 +3893,10 @@ class ProxyConfig:
)
try:
guardrails_in_db: List[Guardrail] = (
await GuardrailRegistry.get_all_guardrails_from_db(
prisma_client=prisma_client
)
guardrails_in_db: List[
Guardrail
] = await GuardrailRegistry.get_all_guardrails_from_db(
prisma_client=prisma_client
)
verbose_proxy_logger.debug(
"guardrails from the DB %s", str(guardrails_in_db)
@ -3894,7 +3983,7 @@ class ProxyConfig:
)
try:
await global_mcp_server_manager._add_mcp_servers_from_db_to_in_memory_registry()
await global_mcp_server_manager.reload_servers_from_database()
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db - {}".format(
@ -4031,6 +4120,23 @@ class ProxyConfig:
return []
async def _reload_mcp_servers_job():
"""Background job entrypoint for MCP registry refreshes."""
if proxy_config._should_load_db_object(object_type="mcp") is False:
return
try:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
await global_mcp_server_manager.reload_servers_from_database()
except Exception as e:
verbose_proxy_logger.exception(
"Failed to reload MCP servers from database: %s", str(e)
)
proxy_config = ProxyConfig()
@ -4134,9 +4240,9 @@ async def initialize( # noqa: PLR0915
user_api_base = api_base
dynamic_config[user_model]["api_base"] = api_base
if api_version:
os.environ["AZURE_API_VERSION"] = (
api_version # set this for azure - litellm can read this from the env
)
os.environ[
"AZURE_API_VERSION"
] = api_version # set this for azure - litellm can read this from the env
if max_tokens: # model-specific param
dynamic_config[user_model]["max_tokens"] = max_tokens
if temperature: # model-specific param
@ -4568,6 +4674,18 @@ class ProxyStartupEvent:
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
await proxy_config.get_credentials(prisma_client=prisma_client)
from litellm.proxy._experimental.mcp_server.utils import is_mcp_available
if is_mcp_available():
scheduler.add_job(
_reload_mcp_servers_job,
"interval",
seconds=30,
id="reload_mcp_servers_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
await cls._initialize_slack_alerting_jobs(
scheduler=scheduler,
general_settings=general_settings,
@ -4654,10 +4772,14 @@ class ProxyStartupEvent:
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
verbose_proxy_logger.info("Responses cost check job scheduled successfully")
verbose_proxy_logger.info(
"Responses cost check job scheduled successfully"
)
except Exception as e:
verbose_proxy_logger.debug(f"Failed to setup responses cost checking: {e}")
verbose_proxy_logger.debug(
f"Failed to setup responses cost checking: {e}"
)
verbose_proxy_logger.debug(
"Checking responses cost for LiteLLM Managed Files is an Enterprise Feature. Skipping..."
)
@ -4681,8 +4803,9 @@ class ProxyStartupEvent:
"""
Initialize the spend tracking and other background jobs
1. CloudZero Background Job
2. Prometheus Background Job
3. Key Rotation Background Job
2. Focus Background Job
3. Prometheus Background Job
4. Key Rotation Background Job
Args:
scheduler: The scheduler to add the background jobs to
@ -4691,11 +4814,17 @@ class ProxyStartupEvent:
# CloudZero Background Job
########################################################
from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
from litellm.integrations.focus.focus_logger import FocusLogger
from litellm.proxy.spend_tracking.cloudzero_endpoints import is_cloudzero_setup
if await is_cloudzero_setup():
await CloudZeroLogger.init_cloudzero_background_job(scheduler=scheduler)
########################################################
# Focus Background Job
########################################################
await FocusLogger.init_focus_export_background_job(scheduler=scheduler)
########################################################
# Prometheus Background Job
########################################################
@ -4826,6 +4955,14 @@ class ProxyStartupEvent:
await prisma_client.connect()
## Start RDS IAM token refresh background task if enabled ##
# This proactively refreshes IAM tokens before they expire,
# preventing the 15-minute connection failure bug (#16220)
if hasattr(prisma_client, "db") and hasattr(
prisma_client.db, "start_token_refresh_task"
):
await prisma_client.db.start_token_refresh_task()
## Add necessary views to proxy ##
asyncio.create_task(
prisma_client.check_view_exists()
@ -5937,7 +6074,6 @@ async def realtime_websocket_endpoint(
),
user_api_key_dict=Depends(user_api_key_auth_websocket),
):
await websocket.accept()
# Only use explicit parameters, not all query params
@ -6736,7 +6872,7 @@ async def run_thread(
if (
"stream" in data and data["stream"] is True
): # use generate_responses to stream responses
return await create_streaming_response(
return await create_response(
generator=async_assistants_data_generator(
user_api_key_dict=user_api_key_dict,
response=response,
@ -9525,9 +9661,9 @@ async def get_config_list(
hasattr(sub_field_info, "description")
and sub_field_info.description is not None
):
nested_fields[idx].field_description = (
sub_field_info.description
)
nested_fields[
idx
].field_description = sub_field_info.description
idx += 1
_stored_in_db = None
@ -9762,6 +9898,18 @@ async def get_config(): # noqa: PLR0915
_failure_callbacks = _litellm_settings.get("failure_callback", [])
_success_and_failure_callbacks = _litellm_settings.get("callbacks", [])
# Normalize string callbacks to lists
def normalize_callback(callback):
if isinstance(callback, str):
return [callback]
elif callback is None:
return []
return callback
_success_callbacks = normalize_callback(_success_callbacks)
_failure_callbacks = normalize_callback(_failure_callbacks)
_success_and_failure_callbacks = normalize_callback(_success_and_failure_callbacks)
_data_to_return = []
"""
[

91
litellm/proxy/utils.py
View file

@ -131,6 +131,7 @@ else:
unified_guardrail = UnifiedLLMGuardrails()
_anthropic_async_clients = {}
def print_verbose(print_statement):
"""
@ -961,6 +962,7 @@ class ProxyLogging:
Updated data dictionary if guardrail passes, None if guardrail should be skipped
"""
from litellm.types.guardrails import GuardrailEventHooks
from litellm.integrations.prometheus import PrometheusLogger
# Determine the event type based on call type
event_type = GuardrailEventHooks.pre_call
@ -973,30 +975,62 @@ class ProxyLogging:
guardrail_name = callback.guardrail_name
# Check if load balancing should be used
if guardrail_name and self._should_use_guardrail_load_balancing(guardrail_name):
response = await self._execute_guardrail_with_load_balancing(
guardrail_name=guardrail_name,
hook_type="pre_call",
data=data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
)
else:
# Single guardrail - execute directly
response = await self._execute_guardrail_hook(
callback=callback,
hook_type="pre_call",
data=data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
)
# Track timing and errors for prometheus metrics
# Use time.perf_counter() for more accurate duration measurements
guardrail_start_time = time.perf_counter()
status = "success"
error_type = None
# Process the response if one was returned
if response is not None:
data = await self.process_pre_call_hook_response(
response=response, data=data, call_type=call_type
)
try:
# Check if load balancing should be used
if guardrail_name and self._should_use_guardrail_load_balancing(guardrail_name):
response = await self._execute_guardrail_with_load_balancing(
guardrail_name=guardrail_name,
hook_type="pre_call",
data=data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
)
else:
# Single guardrail - execute directly
response = await self._execute_guardrail_hook(
callback=callback,
hook_type="pre_call",
data=data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
)
# Process the response if one was returned
if response is not None:
data = await self.process_pre_call_hook_response(
response=response, data=data, call_type=call_type
)
except Exception as e:
status = "error"
error_type = type(e).__name__
# Re-raise the exception to maintain existing behavior
raise
finally:
# Record prometheus metrics
guardrail_end_time = time.perf_counter()
latency_seconds = guardrail_end_time - guardrail_start_time
# Get guardrail name for metrics (fallback if not set)
metrics_guardrail_name = guardrail_name or getattr(callback, "guardrail_name", callback.__class__.__name__) or "unknown"
# Find PrometheusLogger in callbacks and record metrics
for prom_callback in litellm.callbacks:
if isinstance(prom_callback, PrometheusLogger):
prom_callback._record_guardrail_metrics(
guardrail_name=metrics_guardrail_name,
latency_seconds=latency_seconds,
status=status,
error_type=error_type,
hook_type="pre_call",
)
break
return data
@ -1195,7 +1229,7 @@ class ProxyLogging:
and _callback.__class__.async_pre_call_hook
!= CustomLogger.async_pre_call_hook
):
if call_type == "mcp_call" and user_api_key_dict is None:
if call_type == "call_mcp_tool" and user_api_key_dict is None:
continue
response = await _callback.async_pre_call_hook(
@ -4254,11 +4288,16 @@ async def count_tokens_with_anthropic_api(
if anthropic_api_key and messages:
# Call Anthropic API directly for more accurate token counting
client = anthropic.Anthropic(api_key=anthropic_api_key)
# Use cached client if available to avoid socket exhaustion
if anthropic_api_key not in _anthropic_async_clients:
_anthropic_async_clients[anthropic_api_key] = anthropic.AsyncAnthropic(api_key=anthropic_api_key)
client = _anthropic_async_clients[anthropic_api_key]
# Call with explicit parameters to satisfy type checking
# Type ignore for now since messages come from generic dict input
response = client.beta.messages.count_tokens(
response = await client.beta.messages.count_tokens(
model=model_to_use,
messages=messages, # type: ignore
betas=["token-counting-2024-11-01"],

14
litellm/responses/main.py
View file

@ -577,7 +577,13 @@ def responses(
api_base=litellm_params.api_base,
api_key=litellm_params.api_key,
)
# Use dynamic credentials from get_llm_provider (e.g., when use_litellm_proxy=True)
if dynamic_api_key is not None:
litellm_params.api_key = dynamic_api_key
if dynamic_api_base is not None:
litellm_params.api_base = dynamic_api_base
#########################################################
# Update input with provider-specific file IDs if managed files are used
#########################################################
@ -1483,6 +1489,12 @@ def compact_responses(
api_key=litellm_params.api_key,
)
# Use dynamic credentials from get_llm_provider (e.g., when use_litellm_proxy=True)
if dynamic_api_key is not None:
litellm_params.api_key = dynamic_api_key
if dynamic_api_base is not None:
litellm_params.api_base = dynamic_api_base
if custom_llm_provider is None:
raise ValueError("custom_llm_provider is required but passed as None")

17
litellm/responses/utils.py
View file

@ -443,11 +443,18 @@ class ResponseAPILoggingUtils:
completion_tokens=0,
total_tokens=0,
)
response_api_usage: ResponseAPIUsage = (
ResponseAPIUsage(**usage_input)
if isinstance(usage_input, dict)
else usage_input
)
response_api_usage: ResponseAPIUsage
if isinstance(usage_input, dict):
total_tokens = usage_input.get("total_tokens")
if total_tokens is None:
input_tokens = usage_input.get("input_tokens")
output_tokens = usage_input.get("output_tokens")
if input_tokens is not None and output_tokens is not None:
total_tokens = input_tokens + output_tokens
usage_input["total_tokens"] = total_tokens
response_api_usage = ResponseAPIUsage(**usage_input)
else:
response_api_usage = usage_input
prompt_tokens: int = response_api_usage.input_tokens or 0
completion_tokens: int = response_api_usage.output_tokens or 0
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None

143
litellm/router.py
View file

@ -255,6 +255,7 @@ class Router:
] = {},
enable_pre_call_checks: bool = False,
enable_tag_filtering: bool = False,
tag_filtering_match_any: bool = True,
retry_after: int = 0, # min time to wait before retrying a failed request
retry_policy: Optional[
Union[RetryPolicy, dict]
@ -363,6 +364,7 @@ class Router:
self.debug_level = debug_level
self.enable_pre_call_checks = enable_pre_call_checks
self.enable_tag_filtering = enable_tag_filtering
self.tag_filtering_match_any = tag_filtering_match_any
from litellm._service_logger import ServiceLogging
self.service_logger_obj: ServiceLogging = ServiceLogging()
@ -384,9 +386,9 @@ class Router:
) # names of models under litellm_params. ex. azure/chatgpt-v-2
self.deployment_latency_map = {}
### CACHING ###
cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = (
"local" # default to an in-memory cache
)
cache_type: Literal[
"local", "redis", "redis-semantic", "s3", "disk"
] = "local" # default to an in-memory cache
redis_cache = None
cache_config: Dict[str, Any] = {}
@ -428,9 +430,9 @@ class Router:
self.default_max_parallel_requests = default_max_parallel_requests
self.provider_default_deployment_ids: List[str] = []
self.pattern_router = PatternMatchRouter()
self.team_pattern_routers: Dict[str, PatternMatchRouter] = (
{}
) # {"TEAM_ID": PatternMatchRouter}
self.team_pattern_routers: Dict[
str, PatternMatchRouter
] = {} # {"TEAM_ID": PatternMatchRouter}
self.auto_routers: Dict[str, "AutoRouter"] = {}
# Initialize model_group_alias early since it's used in set_model_list
@ -611,9 +613,9 @@ class Router:
)
)
self.model_group_retry_policy: Optional[Dict[str, RetryPolicy]] = (
model_group_retry_policy
)
self.model_group_retry_policy: Optional[
Dict[str, RetryPolicy]
] = model_group_retry_policy
self.allowed_fails_policy: Optional[AllowedFailsPolicy] = None
if allowed_fails_policy is not None:
@ -720,7 +722,10 @@ class Router:
valid_strategy_strings = ["simple-shuffle"] + [s.value for s in RoutingStrategy]
if routing_strategy is not None:
is_valid_string = isinstance(routing_strategy, str) and routing_strategy in valid_strategy_strings
is_valid_string = (
isinstance(routing_strategy, str)
and routing_strategy in valid_strategy_strings
)
is_valid_enum = isinstance(routing_strategy, RoutingStrategy)
if not is_valid_string and not is_valid_enum:
raise ValueError(
@ -1069,7 +1074,7 @@ class Router:
self.delete_container = self.factory_function(
delete_container, call_type="delete_container"
)
# Auto-register JSON-generated container file endpoints
for name, func in container_file_endpoints.items():
setattr(self, name, self.factory_function(func, call_type=name)) # type: ignore[arg-type]
@ -1498,10 +1503,7 @@ class Router:
async def _acompletion(
self, model: str, messages: List[Dict[str, str]], **kwargs
) -> Union[
ModelResponse,
CustomStreamWrapper,
]:
) -> Union[ModelResponse, CustomStreamWrapper,]:
"""
- Get an available deployment
- call it with a semaphore over the call
@ -3019,7 +3021,9 @@ class Router:
kwargs["original_generic_function"] = original_function
kwargs["original_function"] = self._aguardrail_helper
self._update_kwargs_before_fallbacks(
model=guardrail_name, kwargs=kwargs, metadata_variable_name="litellm_metadata"
model=guardrail_name,
kwargs=kwargs,
metadata_variable_name="litellm_metadata",
)
verbose_router_logger.debug(
f"Inside aguardrail() - guardrail_name: {guardrail_name}; kwargs: {kwargs}"
@ -3312,8 +3316,7 @@ class Router:
kwargs["model"] = model
kwargs["input"] = input
kwargs["original_function"] = self._embedding
kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries)
kwargs.setdefault("metadata", {}).update({"model_group": model})
self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs)
response = self.function_with_fallbacks(**kwargs)
return response
except Exception as e:
@ -3615,9 +3618,9 @@ class Router:
healthy_deployments=healthy_deployments, responses=responses
)
returned_response = cast(OpenAIFileObject, responses[0])
returned_response._hidden_params["model_file_id_mapping"] = (
model_file_id_mapping
)
returned_response._hidden_params[
"model_file_id_mapping"
] = model_file_id_mapping
return returned_response
except Exception as e:
verbose_router_logger.exception(
@ -4364,11 +4367,11 @@ class Router:
if isinstance(e, litellm.ContextWindowExceededError):
if context_window_fallbacks is not None:
context_window_fallback_model_group: Optional[List[str]] = (
self._get_fallback_model_group_from_fallbacks(
fallbacks=context_window_fallbacks,
model_group=model_group,
)
context_window_fallback_model_group: Optional[
List[str]
] = self._get_fallback_model_group_from_fallbacks(
fallbacks=context_window_fallbacks,
model_group=model_group,
)
if context_window_fallback_model_group is None:
raise original_exception
@ -4400,11 +4403,11 @@ class Router:
e.message += "\n{}".format(error_message)
elif isinstance(e, litellm.ContentPolicyViolationError):
if content_policy_fallbacks is not None:
content_policy_fallback_model_group: Optional[List[str]] = (
self._get_fallback_model_group_from_fallbacks(
fallbacks=content_policy_fallbacks,
model_group=model_group,
)
content_policy_fallback_model_group: Optional[
List[str]
] = self._get_fallback_model_group_from_fallbacks(
fallbacks=content_policy_fallbacks,
model_group=model_group,
)
if content_policy_fallback_model_group is None:
raise original_exception
@ -4483,9 +4486,21 @@ class Router:
if hasattr(original_exception, "message"):
# add the available fallbacks to the exception
original_exception.message += ". Received Model Group={}\nAvailable Model Group Fallbacks={}".format( # type: ignore
model_group,
fallback_model_group,
deployment_info = ""
if kwargs is not None:
metadata = kwargs.get('metadata', {})
if metadata and 'deployment' in metadata:
deployment_info = f"\nUsed Deployment: {metadata['deployment']}"
if 'model_info' in metadata:
model_info = metadata['model_info']
if isinstance(model_info, dict):
deployment_info += f"\nDeployment ID: {model_info.get('id', 'unknown')}"
original_exception.message += ( # type: ignore
f". Received Model Group={model_group}"
f"\nAvailable Model Group Fallbacks={fallback_model_group}"
f"{deployment_info}"
f"\n\n💡 Tip: If using wildcard patterns (e.g., 'openai/*'), ensure all matching deployments have credentials with access to this model."
)
if len(fallback_failure_exception_str) > 0:
original_exception.message += ( # type: ignore
@ -5667,26 +5682,26 @@ class Router:
"""
from litellm.router_strategy.auto_router.auto_router import AutoRouter
auto_router_config_path: Optional[str] = (
deployment.litellm_params.auto_router_config_path
)
auto_router_config_path: Optional[
str
] = deployment.litellm_params.auto_router_config_path
auto_router_config: Optional[str] = deployment.litellm_params.auto_router_config
if auto_router_config_path is None and auto_router_config is None:
raise ValueError(
"auto_router_config_path or auto_router_config is required for auto-router deployments. Please set it in the litellm_params"
)
default_model: Optional[str] = (
deployment.litellm_params.auto_router_default_model
)
default_model: Optional[
str
] = deployment.litellm_params.auto_router_default_model
if default_model is None:
raise ValueError(
"auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params"
)
embedding_model: Optional[str] = (
deployment.litellm_params.auto_router_embedding_model
)
embedding_model: Optional[
str
] = deployment.litellm_params.auto_router_embedding_model
if embedding_model is None:
raise ValueError(
"auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params"
@ -6233,9 +6248,9 @@ class Router:
# Add custom_llm_provider
if deployment.litellm_params.custom_llm_provider:
credentials["custom_llm_provider"] = (
deployment.litellm_params.custom_llm_provider
)
credentials[
"custom_llm_provider"
] = deployment.litellm_params.custom_llm_provider
elif "/" in deployment.litellm_params.model:
# Extract provider from "provider/model" format
credentials["custom_llm_provider"] = deployment.litellm_params.model.split(
@ -6929,42 +6944,44 @@ class Router:
"""
return candidate_id in self.model_id_to_deployment_index_map
def resolve_model_name_from_model_id(self, model_id: Optional[str]) -> Optional[str]:
def resolve_model_name_from_model_id(
self, model_id: Optional[str]
) -> Optional[str]:
"""
Resolve model_name from model_id.
This method attempts to find the correct model_name to use with the router
so that litellm_params can be automatically injected from the model config.
Strategy:
1. First, check if model_id directly matches a model_name or deployment ID
2. If not, search through router's model_list to find a match by litellm_params.model
3. Return the model_name if found, None otherwise
Args:
model_id: The model_id extracted from decoded video_id
(could be model_name or litellm_params.model value)
Returns:
model_name if found, None otherwise. If None, the request will fall through
to normal flow using environment variables.
"""
if not model_id:
return None
# Strategy 1: Check if model_id directly matches a model_name or deployment ID
if model_id in self.model_names or self.has_model_id(model_id):
return model_id
# Strategy 2: Search through router's model_list to find by litellm_params.model
all_models = self.get_model_list(model_name=None)
if not all_models:
return None
for deployment in all_models:
litellm_params = deployment.get("litellm_params", {})
actual_model = litellm_params.get("model")
# Match by exact match or by checking if actual_model ends with /model_id or :model_id
# e.g., model_id="veo-2.0-generate-001" matches actual_model="vertex_ai/veo-2.0-generate-001"
matches = (
@ -6972,12 +6989,12 @@ class Router:
or (actual_model and actual_model.endswith(f"/{model_id}"))
or (actual_model and actual_model.endswith(f":{model_id}"))
)
if matches:
model_name = deployment.get("model_name")
if model_name:
return model_name
# No match found
return None
@ -7662,6 +7679,10 @@ class Router:
)
if pattern_deployments:
verbose_router_logger.debug(
f"Pattern match for model='{model}': Found {len(pattern_deployments)} deployments. "
f"Deployment IDs: {[d.get('model_info', {}).get('id', 'unknown') for d in pattern_deployments]}"
)
return model, pattern_deployments
if (
@ -7767,14 +7788,18 @@ class Router:
request_kwargs=request_kwargs,
)
verbose_router_logger.debug(f"healthy_deployments after team filter: {healthy_deployments}")
verbose_router_logger.debug(
f"healthy_deployments after team filter: {healthy_deployments}"
)
healthy_deployments = filter_web_search_deployments(
healthy_deployments=healthy_deployments,
request_kwargs=request_kwargs,
)
verbose_router_logger.debug(f"healthy_deployments after web search filter: {healthy_deployments}")
verbose_router_logger.debug(
f"healthy_deployments after web search filter: {healthy_deployments}"
)
if isinstance(healthy_deployments, dict):
return healthy_deployments

23
litellm/router_strategy/tag_based_routing.py
View file

@ -20,17 +20,28 @@ else:
def is_valid_deployment_tag(
deployment_tags: List[str], request_tags: List[str]
deployment_tags: List[str], request_tags: List[str], match_any: bool = True
) -> bool:
"""
Check if a tag is valid
Check if a tag is valid, the matching can be either any or all based on `match_any` flag
"""
if not request_tags:
return False
if any(tag in deployment_tags for tag in request_tags):
dep_set = set(deployment_tags)
req_set = set(request_tags)
if match_any:
is_valid_deployment = bool(dep_set & req_set)
else:
is_valid_deployment = req_set.issubset(dep_set)
if is_valid_deployment:
verbose_logger.debug(
"adding deployment with tags: %s, request tags: %s",
"adding deployment with tags: %s, request tags: %s for match_any=%s",
deployment_tags,
request_tags,
match_any,
)
return True
return False
@ -68,6 +79,7 @@ async def get_deployments_for_tag(
if metadata_variable_name in request_kwargs:
metadata = request_kwargs[metadata_variable_name]
request_tags = metadata.get("tags")
match_any = llm_router_instance.tag_filtering_match_any
new_healthy_deployments = []
default_deployments = []
@ -76,7 +88,6 @@ async def get_deployments_for_tag(
"get_deployments_for_tag routing: router_keys: %s", request_tags
)
# example this can be router_keys=["free", "custom"]
# get all deployments that have a superset of these router keys
for deployment in healthy_deployments:
deployment_litellm_params = deployment.get("litellm_params")
deployment_tags = deployment_litellm_params.get("tags")
@ -90,7 +101,7 @@ async def get_deployments_for_tag(
if deployment_tags is None:
continue
if is_valid_deployment_tag(deployment_tags, request_tags):
if is_valid_deployment_tag(deployment_tags, request_tags, match_any):
new_healthy_deployments.append(deployment)
if "default" in deployment_tags:

47
litellm/types/integrations/prometheus.py
View file

@ -1,7 +1,7 @@
import re
from dataclasses import dataclass
from enum import Enum
from typing import Dict, List, Literal, Optional, Tuple, Union
from typing import Dict, List, Literal, Optional, Tuple
from pydantic import BaseModel, Field
from typing_extensions import Annotated
@ -185,6 +185,14 @@ DEFINED_PROMETHEUS_METRICS = Literal[
"litellm_redis_daily_spend_update_queue_size",
"litellm_in_memory_spend_update_queue_size",
"litellm_redis_spend_update_queue_size",
"litellm_request_queue_time_seconds",
"litellm_guardrail_latency_seconds",
"litellm_guardrail_errors_total",
"litellm_guardrail_requests_total",
# Cache metrics
"litellm_cache_hits_metric",
"litellm_cache_misses_metric",
"litellm_cached_tokens_metric",
]
@ -219,6 +227,23 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value,
]
litellm_request_queue_time_seconds = [
UserAPIKeyLabelNames.END_USER.value,
UserAPIKeyLabelNames.API_KEY_HASH.value,
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
UserAPIKeyLabelNames.REQUESTED_MODEL.value,
UserAPIKeyLabelNames.TEAM.value,
UserAPIKeyLabelNames.TEAM_ALIAS.value,
UserAPIKeyLabelNames.USER.value,
UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value,
]
# Guardrail metrics - these use custom labels (guardrail_name, status, error_type, hook_type)
# which are not part of UserAPIKeyLabelNames
litellm_guardrail_latency_seconds: List[str] = []
litellm_guardrail_errors_total: List[str] = []
litellm_guardrail_requests_total: List[str] = []
litellm_proxy_total_requests_metric = [
UserAPIKeyLabelNames.END_USER.value,
UserAPIKeyLabelNames.API_KEY_HASH.value,
@ -436,6 +461,21 @@ class PrometheusMetricLabels:
litellm_redis_spend_update_queue_size: List[str] = []
# Cache metrics - track cache hits, misses, and tokens served from cache
_cache_metric_labels = [
UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value,
UserAPIKeyLabelNames.API_KEY_HASH.value,
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
UserAPIKeyLabelNames.TEAM.value,
UserAPIKeyLabelNames.TEAM_ALIAS.value,
UserAPIKeyLabelNames.END_USER.value,
UserAPIKeyLabelNames.USER.value,
]
litellm_cache_hits_metric = _cache_metric_labels
litellm_cache_misses_metric = _cache_metric_labels
litellm_cached_tokens_metric = _cache_metric_labels
@staticmethod
def get_labels(label_name: DEFINED_PROMETHEUS_METRICS) -> List[str]:
default_labels = getattr(PrometheusMetricLabels, label_name)
@ -460,11 +500,6 @@ class PrometheusMetricLabels:
return default_labels + custom_labels
from typing import List, Optional
from pydantic import BaseModel, Field
class UserAPIKeyLabelValues(BaseModel):
end_user: Annotated[
Optional[str], Field(..., alias=UserAPIKeyLabelNames.END_USER.value)

8
litellm/types/management_endpoints/router_settings_endpoints.py
View file

@ -184,6 +184,14 @@ ROUTER_SETTINGS_FIELDS: List[RouterSettingsField] = [
field_default=False,
ui_field_name="Enable Tag Filtering",
link="https://docs.litellm.ai/docs/proxy/tag_routing",
),
RouterSettingsField(
field_name="tag_filtering_match_any",
field_type="Boolean",
field_value=None,
field_description="Match any tag instead of all tags for tag-based routing",
field_default=True,
ui_field_name="Tag Filtering Match Any",
),
RouterSettingsField(
field_name="disable_cooldowns",

2
litellm/types/mcp_server/mcp_server_manager.py
View file

@ -1,3 +1,4 @@
from datetime import datetime
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, ConfigDict
@ -50,4 +51,5 @@ class MCPServer(BaseModel):
env: Optional[Dict[str, str]] = None
access_groups: Optional[List[str]] = None
allow_all_keys: bool = False
updated_at: Optional[datetime] = None
model_config = ConfigDict(arbitrary_types_allowed=True)

2
litellm/types/utils.py
View file

@ -3016,6 +3016,7 @@ class LlmProviders(str, Enum):
AUTO_ROUTER = "auto_router"
VERCEL_AI_GATEWAY = "vercel_ai_gateway"
DOTPROMPT = "dotprompt"
MANUS = "manus"
WANDB = "wandb"
OVHCLOUD = "ovhcloud"
LEMONADE = "lemonade"
@ -3028,6 +3029,7 @@ class LlmProviders(str, Enum):
NANOGPT = "nano-gpt"
POE = "poe"
CHUTES = "chutes"
XIAOMI_MIMO = "xiaomi_mimo"

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