tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-10 13:03:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into wandb-inference

Browse Source
This commit is contained in:
Anubhav Singh
2025-09-23 00:14:11 +05:30
committed by GitHub
parent 470e1b7d2d fa20abfe48
commit fde4dcb0b8
271 changed files with 61393 additions and 45494 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.circleci/config.yml
+77 -4
View File
@@ -1050,6 +1050,51 @@ jobs:
ls
python -m pytest -vv tests/test_litellm --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit-litellm.xml --durations=10 -n 8
no_output_timeout: 120m
- run:
name: Rename the coverage files
command: |
mv coverage.xml litellm_mapped_tests_coverage.xml
mv .coverage litellm_mapped_tests_coverage
# Store test results
- store_test_results:
path: test-results
- persist_to_workspace:
root: .
paths:
- litellm_mapped_tests_coverage.xml
- litellm_mapped_tests_coverage
litellm_mapped_enterprise_tests:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
steps:
- checkout
- setup_google_dns
- run:
name: Install Dependencies
command: |
python -m pip install --upgrade pip
python -m pip install -r requirements.txt
pip install "pytest-mock==3.12.0"
pip install "pytest==7.3.1"
pip install "pytest-retry==1.6.3"
pip install "pytest-cov==5.0.0"
pip install "pytest-asyncio==0.21.1"
pip install "respx==0.22.0"
pip install "hypercorn==0.17.3"
pip install "pydantic==2.10.2"
pip install "mcp==1.10.1"
pip install "requests-mock>=1.12.1"
pip install "responses==0.25.7"
pip install "pytest-xdist==3.6.1"
pip install "semantic_router==0.1.10"
pip install "fastapi-offline==1.7.3"
- setup_litellm_enterprise_pip
- run:
name: Run enterprise tests
command: |
@@ -1458,6 +1503,7 @@ jobs:
# - run: python ./tests/documentation_tests/test_general_setting_keys.py
- run: python ./tests/code_coverage_tests/check_licenses.py
- run: python ./tests/code_coverage_tests/router_code_coverage.py
- run: python ./tests/code_coverage_tests/test_chat_completion_imports.py
- run: python ./tests/code_coverage_tests/info_log_check.py
- run: python ./tests/code_coverage_tests/test_ban_set_verbose.py
- run: python ./tests/code_coverage_tests/code_qa_check_tests.py
@@ -1778,8 +1824,8 @@ jobs:
docker run -d \
-p 4000:4000 \
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e AZURE_API_KEY=$AZURE_BATCHES_API_KEY \
-e AZURE_API_BASE=$AZURE_BATCHES_API_BASE \
-e AZURE_API_KEY=$AZURE_API_KEY \
-e AZURE_API_BASE=$AZURE_API_BASE \
-e AZURE_API_VERSION="2024-05-01-preview" \
-e REDIS_HOST=$REDIS_HOST \
-e REDIS_PASSWORD=$REDIS_PASSWORD \
@@ -2825,8 +2871,8 @@ jobs:
source "$NVM_DIR/bash_completion"
# Install and use Node version
nvm install v18.17.0
nvm use v18.17.0
nvm install v20
nvm use v20
cd ui/litellm-dashboard
@@ -2879,7 +2925,26 @@ jobs:
name: Install Playwright Browsers
command: |
npx playwright install
- run:
name: Run UI unit tests (Vitest)
command: |
# Use Node 20 (several deps require >=20)
export NVM_DIR="/opt/circleci/.nvm"
source "$NVM_DIR/nvm.sh"
nvm install 20
nvm use 20
cd ui/litellm-dashboard
npm ci || npm install
# CI run, with both LCOV (Codecov) and HTML (artifact you can click)
CI=true npm run test -- --run --coverage \
--coverage.provider=v8 \
--coverage.reporter=lcov \
--coverage.reporter=html \
--coverage.reportsDirectory=coverage/html
- run:
name: Build Docker image
command: docker build -t my-app:latest -f ./docker/Dockerfile.database .
@@ -3155,6 +3220,12 @@ workflows:
only:
- main
- /litellm_.*/
- litellm_mapped_enterprise_tests:
filters:
branches:
only:
- main
- /litellm_.*/
- litellm_mapped_tests:
filters:
branches:
@@ -3199,6 +3270,7 @@ workflows:
- guardrails_testing
- llm_responses_api_testing
- litellm_mapped_tests
- litellm_mapped_enterprise_tests
- batches_testing
- litellm_utils_testing
- pass_through_unit_testing
@@ -3259,6 +3331,7 @@ workflows:
- google_generate_content_endpoint_testing
- llm_responses_api_testing
- litellm_mapped_tests
- litellm_mapped_enterprise_tests
- batches_testing
- litellm_utils_testing
- pass_through_unit_testing
.gitignore
+2 -1
View File
@@ -95,4 +95,5 @@ test.py
litellm_config.yaml
.cursor
.vscode/launch.json
litellm/proxy/to_delete_loadtest_work/*
litellm/proxy/to_delete_loadtest_work/*
update_model_cost_map.py
Dockerfile
-3
View File
@@ -41,9 +41,6 @@ RUN pip uninstall jwt -y
RUN pip uninstall PyJWT -y
RUN pip install PyJWT==2.9.0 --no-cache-dir
# Build Admin UI
RUN chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
# Runtime stage
FROM $LITELLM_RUNTIME_IMAGE AS runtime
README.md
+1 -1
View File
@@ -37,7 +37,7 @@ LiteLLM manages:
- Retry/fallback logic across multiple deployments (e.g. Azure/OpenAI) - [Router](https://docs.litellm.ai/docs/routing)
- Set Budgets & Rate limits per project, api key, model [LiteLLM Proxy Server (LLM Gateway)](https://docs.litellm.ai/docs/simple_proxy)
[**Jump to LiteLLM Proxy (LLM Gateway) Docs**](https://github.com/BerriAI/litellm?tab=readme-ov-file#openai-proxy---docs) <br>
[**Jump to LiteLLM Proxy (LLM Gateway) Docs**](https://github.com/BerriAI/litellm?tab=readme-ov-file#litellm-proxy-server-llm-gateway---docs) <br>
[**Jump to Supported LLM Providers**](https://github.com/BerriAI/litellm?tab=readme-ov-file#supported-providers-docs)
🚨 **Stable Release:** Use docker images with the `-stable` tag. These have undergone 12 hour load tests, before being published. [More information about the release cycle here](https://docs.litellm.ai/docs/proxy/release_cycle)
cookbook/misc/RELEASE_NOTES_GENERATION_INSTRUCTIONS.md
+99 -35
View File
@@ -4,10 +4,10 @@ This document provides comprehensive instructions for AI agents to generate rele
## Required Inputs
1. **Release Version** (e.g., `v1.76.3-stable`)
1. **Release Version** (e.g., `v1.77.3-stable`)
2. **PR Diff/Changelog** - List of PRs with titles and contributors
3. **Previous Version Commit Hash** - To compare model pricing changes
4. **Reference Release Notes** - Previous release notes to follow style/format
4. **Reference Release Notes** - Use recent stable releases (v1.76.3-stable, v1.77.2-stable) as templates for consistent formatting
## Step-by-Step Process
@@ -26,12 +26,12 @@ git diff <previous_commit_hash> HEAD -- model_prices_and_context_window.json
### 2. Release Notes Structure
Follow this exact structure based on `docs/my-website/release_notes/v1.76.1-stable/index.md`:
Follow this exact structure based on recent stable releases (v1.76.3-stable, v1.77.2-stable):
```markdown
---
title: "v1.76.X-stable - [Key Theme]"
slug: "v1-76-X"
title: "v1.77.X-stable - [Key Theme]"
slug: "v1-77-X"
date: YYYY-MM-DDTHH:mm:ss
authors: [standard author block]
hide_table_of_contents: false
@@ -43,23 +43,42 @@ hide_table_of_contents: false
## Key Highlights
[3-5 bullet points of major features]
## Major Changes
[Critical changes users need to know]
## Performance Improvements
[Performance-related changes]
## New Models / Updated Models
[Detailed model tables and provider updates]
#### New Model Support
[Model pricing table]
#### Features
[Provider-specific features organized by provider]
### Bug Fixes
[Provider-specific bug fixes organized by provider]
#### New Provider Support
[New provider integrations]
## LLM API Endpoints
[API-related features and fixes]
#### Features
[API-specific features organized by API type]
#### Bugs
[General bug fixes]
## Management Endpoints / UI
[Admin interface and management changes]
#### Features
[UI and management features]
#### Bugs
[Management-related bug fixes]
## Logging / Guardrail Integrations
[Observability and security features]
#### Features
[Organized by integration provider with proper doc links]
#### Guardrails
[Guardrail-specific features and fixes]
#### New Integration
[Major new integrations]
## Performance / Loadbalancing / Reliability improvements
[Infrastructure improvements]
@@ -86,21 +105,27 @@ hide_table_of_contents: false
**New Models/Updated Models:**
- Extract from model_prices_and_context_window.json diff
- Create tables with: Provider, Model, Context Window, Input Cost, Output Cost, Features
- Group by provider
- Note pricing corrections
- Highlight deprecated models
- **Structure:**
- `#### New Model Support` - pricing table
- `#### Features` - organized by provider with documentation links
- `### Bug Fixes` - provider-specific bug fixes
- `#### New Provider Support` - major new provider integrations
- Group by provider with proper doc links: `**[Provider Name](../../docs/providers/[provider])**`
- Use bullet points under each provider for multiple features
- Separate features from bug fixes clearly
**Provider Features:**
- Group by provider (Gemini, OpenAI, Anthropic, etc.)
- Link to provider docs: `../../docs/providers/[provider_name]`
- Separate features from bug fixes
**API Endpoints:**
- Images API
- Video Generation (if applicable)
- Responses API
- Passthrough endpoints
- General chat completions
**LLM API Endpoints:**
- **Structure:**
- `#### Features` - organized by API type (Responses API, Batch API, etc.)
- `#### Bugs` - general bug fixes under **General** category
- **API Categories:**
- Responses API
- Batch API
- CountTokens API
- Images API
- Video Generation (if applicable)
- General (miscellaneous improvements)
- Use proper documentation links for each API type
**UI/Management:**
- Authentication changes
@@ -108,11 +133,19 @@ hide_table_of_contents: false
- Team management
- Key management
**Integrations:**
- Logging providers (Datadog, Braintrust, etc.)
- Guardrails
- Cost tracking
- Observability
**Logging / Guardrail Integrations:**
- **Structure:**
- `#### Features` - organized by integration provider with proper doc links
- `#### Guardrails` - guardrail-specific features and fixes
- `#### New Integration` - major new integrations
- **Integration Categories:**
- **[DataDog](../../docs/proxy/logging#datadog)** - group all DataDog-related changes
- **[Langfuse](../../docs/proxy/logging#langfuse)** - Langfuse-specific features
- **[Prometheus](../../docs/proxy/logging#prometheus)** - monitoring improvements
- **[PostHog](../../docs/observability/posthog)** - observability integration
- Other logging providers with proper doc links
- Use bullet points under each provider for multiple features
- Separate logging features from guardrails clearly
### 4. Documentation Linking Strategy
@@ -211,10 +244,41 @@ This release has a known issue...
:::
```
**Provider Features:**
**Provider Features (New Models / Updated Models section):**
```markdown
#### Features
- **[Provider Name](../../docs/providers/provider)**
- Feature description - [PR #XXXXX](link)
- Another feature description - [PR #YYYYY](link)
```
**API Features (LLM API Endpoints section):**
```markdown
#### Features
- **[API Name](../../docs/api_path)**
- Feature description - [PR #XXXXX](link)
- Another feature - [PR #YYYYY](link)
- **General**
- Miscellaneous improvements - [PR #ZZZZZ](link)
```
**Integration Features (Logging / Guardrail Integrations section):**
```markdown
#### Features
- **[Integration Name](../../docs/proxy/logging#integration)**
- Feature description - [PR #XXXXX](link)
- Bug fix description - [PR #YYYYY](link)
```
**Bug Fixes Pattern:**
```markdown
### Bug Fixes
- **[Provider/Component Name](../../docs/providers/provider)**
- Bug fix description - [PR #XXXXX](link)
```
### 10. Missing Documentation Check
docs/my-website/docs/completion/provider_specific_params.md
+50
View File
@@ -433,4 +433,54 @@ curl -X POST 'http://0.0.0.0:4000/chat/completions' \
],
"adapater_id": "my-special-adapter-id" # 👈 PROVIDER-SPECIFIC PARAM
}'
## Provider-Specific Metadata Parameters
| Provider | Parameter | Use Case |
|----------|-----------|----------|
| **AWS Bedrock** | `requestMetadata` | Cost attribution, logging |
| **Gemini/Vertex AI** | `labels` | Resource labeling |
| **Anthropic** | `metadata` | User identification |
<Tabs>
<TabItem value="bedrock" label="AWS Bedrock">
```python
import litellm
response = litellm.completion(
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
messages=[{"role": "user", "content": "Hello!"}],
requestMetadata={"cost_center": "engineering"}
)
```
</TabItem>
<TabItem value="gemini" label="Gemini/Vertex AI">
```python
import litellm
response = litellm.completion(
model="vertex_ai/gemini-pro",
messages=[{"role": "user", "content": "Hello!"}],
labels={"environment": "production"}
)
```
</TabItem>
<TabItem value="anthropic" label="Anthropic">
```python
import litellm
response = litellm.completion(
model="anthropic/claude-3-sonnet-20240229",
messages=[{"role": "user", "content": "Hello!"}],
metadata={"user_id": "user123"}
)
```
</TabItem>
</Tabs>
```
docs/my-website/docs/mcp.md
-2
View File
@@ -114,7 +114,6 @@ mcp_servers:
description: "My custom MCP server"
auth_type: "api_key"
auth_value: "abc123"
spec_version: "2025-03-26"
```
**Configuration Options:**
@@ -716,7 +715,6 @@ mcp_servers:
url: https://mcp.deepwiki.com/mcp
transport: "http"
auth_type: "none"
spec_version: "2025-03-26"
access_groups: ["dev_group"]
```
docs/my-website/docs/observability/helicone_integration.md
+232 -48
View File
@@ -1,3 +1,6 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Helicone - OSS LLM Observability Platform
:::tip
@@ -9,9 +12,68 @@ https://github.com/BerriAI/litellm
[Helicone](https://helicone.ai/) is an open source observability platform that proxies your LLM requests and provides key insights into your usage, spend, latency and more.
## Using Helicone with LiteLLM
## Quick Start
LiteLLM provides `success_callbacks` and `failure_callbacks`, allowing you to easily log data to Helicone based on the status of your responses.
<Tabs>
<TabItem value="sdk" label="Python SDK">
Use just 1 line of code to instantly log your responses **across all providers** with Helicone:
```python
import os
from litellm import completion
## Set env variables
os.environ["HELICONE_API_KEY"] = "your-helicone-key"
os.environ["OPENAI_API_KEY"] = "your-openai-key"
# Set callbacks
litellm.success_callback = ["helicone"]
# OpenAI call
response = completion(
model="gpt-4o",
messages=[{"role": "user", "content": "Hi 👋 - I'm OpenAI"}],
)
print(response)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM Proxy">
Add Helicone to your LiteLLM proxy configuration:
```yaml title="config.yaml"
model_list:
- model_name: gpt-4
litellm_params:
model: gpt-4
api_key: os.environ/OPENAI_API_KEY
# Add Helicone callback
litellm_settings:
success_callback: ["helicone"]
# Set Helicone API key
environment_variables:
HELICONE_API_KEY: "your-helicone-key"
```
Start the proxy:
```bash
litellm --config config.yaml
```
</TabItem>
</Tabs>
## Integration Methods
There are two main approaches to integrate Helicone with LiteLLM:
1. **Callbacks**: Log to Helicone while using any provider
2. **Proxy Mode**: Use Helicone as a proxy for advanced features
### Supported LLM Providers
@@ -26,27 +88,16 @@ Helicone can log requests across [various LLM providers](https://docs.helicone.a
- Replicate
- And more
### Integration Methods
## Method 1: Using Callbacks
There are two main approaches to integrate Helicone with LiteLLM:
Log requests to Helicone while using any LLM provider directly.
1. Using callbacks
2. Using Helicone as a proxy
Let's explore each method in detail.
### Approach 1: Use Callbacks
Use just 1 line of code to instantly log your responses **across all providers** with Helicone:
```python
litellm.success_callback = ["helicone"]
```
Complete Code
<Tabs>
<TabItem value="sdk" label="Python SDK">
```python
import os
import litellm
from litellm import completion
## Set env variables
@@ -66,28 +117,78 @@ response = completion(
print(response)
```
### Approach 2: Use Helicone as a proxy
</TabItem>
<TabItem value="proxy" label="LiteLLM Proxy">
```yaml title="config.yaml"
model_list:
- model_name: gpt-4
litellm_params:
model: gpt-4
api_key: os.environ/OPENAI_API_KEY
- model_name: claude-3
litellm_params:
model: anthropic/claude-3-sonnet-20240229
api_key: os.environ/ANTHROPIC_API_KEY
# Add Helicone logging
litellm_settings:
success_callback: ["helicone"]
# Environment variables
environment_variables:
HELICONE_API_KEY: "your-helicone-key"
OPENAI_API_KEY: "your-openai-key"
ANTHROPIC_API_KEY: "your-anthropic-key"
```
Start the proxy:
```bash
litellm --config config.yaml
```
Make requests to your proxy:
```python
import openai
client = openai.OpenAI(
api_key="anything", # proxy doesn't require real API key
base_url="http://localhost:4000"
)
response = client.chat.completions.create(
model="gpt-4", # This gets logged to Helicone
messages=[{"role": "user", "content": "Hello!"}]
)
```
</TabItem>
</Tabs>
## Method 2: Using Helicone as a Proxy
Helicone's proxy provides [advanced functionality](https://docs.helicone.ai/getting-started/proxy-vs-async) like caching, rate limiting, LLM security through [PromptArmor](https://promptarmor.com/) and more.
To use Helicone as a proxy for your LLM requests:
<Tabs>
<TabItem value="sdk" label="Python SDK">
1. Set Helicone as your base URL via: litellm.api_base
2. Pass in Helicone request headers via: litellm.metadata
Complete Code:
Set Helicone as your base URL and pass authentication headers:
```python
import os
import litellm
from litellm import completion
# Configure LiteLLM to use Helicone proxy
litellm.api_base = "https://oai.hconeai.com/v1"
litellm.headers = {
"Helicone-Auth": f"Bearer {os.getenv('HELICONE_API_KEY')}", # Authenticate to send requests to Helicone API
"Helicone-Auth": f"Bearer {os.getenv('HELICONE_API_KEY')}",
}
response = litellm.completion(
# Set your OpenAI API key
os.environ["OPENAI_API_KEY"] = "your-openai-key"
response = completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "How does a court case get to the Supreme Court?"}]
)
@@ -136,36 +237,119 @@ litellm.metadata = {
}
```
### Session Tracking and Tracing
</TabItem>
</Tabs>
## Session Tracking and Tracing
Track multi-step and agentic LLM interactions using session IDs and paths:
```python
litellm.metadata = {
"Helicone-Auth": f"Bearer {os.getenv('HELICONE_API_KEY')}", # Authenticate to send requests to Helicone API
"Helicone-Session-Id": "session-abc-123", # The session ID you want to track
"Helicone-Session-Path": "parent-trace/child-trace", # The path of the session
}
```
- `Helicone-Session-Id`: Use this to specify the unique identifier for the session you want to track. This allows you to group related requests together.
- `Helicone-Session-Path`: This header defines the path of the session, allowing you to represent parent and child traces. For example, "parent/child" represents a child trace of a parent trace.
By using these two headers, you can effectively group and visualize multi-step LLM interactions, gaining insights into complex AI workflows.
### Retry and Fallback Mechanisms
Set up retry mechanisms and fallback options:
<Tabs>
<TabItem value="sdk" label="Python SDK">
```python
import litellm
litellm.api_base = "https://oai.hconeai.com/v1"
litellm.metadata = {
"Helicone-Auth": f"Bearer {os.getenv('HELICONE_API_KEY')}", # Authenticate to send requests to Helicone API
"Helicone-Retry-Enabled": "true", # Enable retry mechanism
"helicone-retry-num": "3", # Set number of retries
"helicone-retry-factor": "2", # Set exponential backoff factor
"Helicone-Fallbacks": '["gpt-3.5-turbo", "gpt-4"]', # Set fallback models
"Helicone-Auth": f"Bearer {os.getenv('HELICONE_API_KEY')}",
"Helicone-Session-Id": "session-abc-123",
"Helicone-Session-Path": "parent-trace/child-trace",
}
response = litellm.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Start a conversation"}]
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM Proxy">
```python
import openai
client = openai.OpenAI(
api_key="anything",
base_url="http://localhost:4000"
)
# First request in session
response1 = client.chat.completions.create(
model="gpt-4",
messages=[{"role": "user", "content": "Hello"}],
extra_headers={
"Helicone-Session-Id": "session-abc-123",
"Helicone-Session-Path": "conversation/greeting"
}
)
# Follow-up request in same session
response2 = client.chat.completions.create(
model="gpt-4",
messages=[{"role": "user", "content": "Tell me more"}],
extra_headers={
"Helicone-Session-Id": "session-abc-123",
"Helicone-Session-Path": "conversation/follow-up"
}
)
```
</TabItem>
</Tabs>
- `Helicone-Session-Id`: Unique identifier for the session to group related requests
- `Helicone-Session-Path`: Hierarchical path to represent parent/child traces (e.g., "parent/child")
## Retry and Fallback Mechanisms
<Tabs>
<TabItem value="sdk" label="Python SDK">
```python
import litellm
litellm.api_base = "https://oai.hconeai.com/v1"
litellm.metadata = {
"Helicone-Auth": f"Bearer {os.getenv('HELICONE_API_KEY')}",
"Helicone-Retry-Enabled": "true",
"helicone-retry-num": "3",
"helicone-retry-factor": "2", # Exponential backoff
"Helicone-Fallbacks": '["gpt-3.5-turbo", "gpt-4"]',
}
response = litellm.completion(
model="gpt-4",
messages=[{"role": "user", "content": "Hello"}]
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM Proxy">
```yaml title="config.yaml"
model_list:
- model_name: gpt-4
litellm_params:
model: gpt-4
api_key: os.environ/OPENAI_API_KEY
api_base: "https://oai.hconeai.com/v1"
default_litellm_params:
headers:
Helicone-Auth: "Bearer ${HELICONE_API_KEY}"
Helicone-Retry-Enabled: "true"
helicone-retry-num: "3"
helicone-retry-factor: "2"
Helicone-Fallbacks: '["gpt-3.5-turbo", "gpt-4"]'
environment_variables:
HELICONE_API_KEY: "your-helicone-key"
OPENAI_API_KEY: "your-openai-key"
```
</TabItem>
</Tabs>
> **Supported Headers** - For a full list of supported Helicone headers and their descriptions, please refer to the [Helicone documentation](https://docs.helicone.ai/getting-started/quick-start).
> By utilizing these headers and metadata options, you can gain deeper insights into your LLM usage, optimize performance, and better manage your AI workflows with Helicone and LiteLLM.
docs/my-website/docs/observability/posthog_integration.md
+216
View File
@@ -0,0 +1,216 @@
# PostHog - Tracking LLM Usage Analytics
## What is PostHog?
PostHog is an open-source product analytics platform that helps you track and analyze how users interact with your product. For LLM applications, PostHog provides specialized AI features to track model usage, performance, and user interactions with your AI features.
## Usage with LiteLLM Proxy (LLM Gateway)
**Step 1**: Create a `config.yaml` file and set `litellm_settings`: `success_callback`
```yaml
model_list:
- model_name: gpt-3.5-turbo
litellm_params:
model: gpt-3.5-turbo
litellm_settings:
success_callback: ["posthog"]
failure_callback: ["posthog"]
```
**Step 2**: Set required environment variables
```shell
export POSTHOG_API_KEY="your-posthog-api-key"
# Optional, defaults to https://app.posthog.com
export POSTHOG_API_URL="https://app.posthog.com" # optional
```
**Step 3**: Start the proxy, make a test request
Start proxy
```shell
litellm --config config.yaml --debug
```
Test Request
```shell
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Content-Type: application/json' \
--data '{
"model": "gpt-3.5-turbo",
"messages": [
{
"role": "user",
"content": "what llm are you"
}
],
"metadata": {
"user_id": "user-123",
"custom_field": "custom_value"
}
}'
```
## Usage with LiteLLM Python SDK
### Quick Start
Use just 2 lines of code, to instantly log your responses **across all providers** with PostHog:
```python
litellm.success_callback = ["posthog"]
litellm.failure_callback = ["posthog"] # logs errors to posthog
```
```python
import litellm
import os
# from PostHog
os.environ["POSTHOG_API_KEY"] = ""
# Optional, defaults to https://app.posthog.com
os.environ["POSTHOG_API_URL"] = "" # optional
# LLM API Keys
os.environ['OPENAI_API_KEY']=""
# set posthog as a callback, litellm will send the data to posthog
litellm.success_callback = ["posthog"]
# openai call
response = litellm.completion(
model="gpt-3.5-turbo",
messages=[
{"role": "user", "content": "Hi - i'm openai"}
],
metadata = {
"user_id": "user-123", # set posthog user ID
}
)
```
### Advanced
#### Set User ID and Custom Metadata
Pass `user_id` in `metadata` to associate events with specific users in PostHog:
**With LiteLLM Python SDK:**
```python
import litellm
litellm.success_callback = ["posthog"]
response = litellm.completion(
model="gpt-3.5-turbo",
messages=[
{"role": "user", "content": "Hello world"}
],
metadata={
"user_id": "user-123", # Add user ID for PostHog tracking
"custom_field": "custom_value" # Add custom metadata
}
)
```
**With LiteLLM Proxy using OpenAI Python SDK:**
```python
import openai
client = openai.OpenAI(
api_key="sk-1234", # Your LiteLLM Proxy API key
base_url="http://0.0.0.0:4000" # Your LiteLLM Proxy URL
)
response = client.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "user", "content": "Hello world"}
],
extra_body={
"metadata": {
"user_id": "user-123", # Add user ID for PostHog tracking
"project_name": "my-project", # Add custom metadata
"environment": "production"
}
}
)
```
#### Disable Logging for Specific Calls
Use the `no-log` flag to prevent logging for specific calls:
```python
import litellm
litellm.success_callback = ["posthog"]
response = litellm.completion(
model="gpt-3.5-turbo",
messages=[
{"role": "user", "content": "This won't be logged"}
],
metadata={"no-log": True}
)
```
## What's Logged to PostHog?
When LiteLLM logs to PostHog, it captures detailed information about your LLM usage:
### For Completion Calls
- **Model Information**: Provider, model name, model parameters
- **Usage Metrics**: Input tokens, output tokens, total cost
- **Performance**: Latency, completion time
- **Content**: Input messages, model responses (respects privacy settings)
- **Metadata**: Custom fields, user ID, trace information
### For Embedding Calls
- **Model Information**: Provider, model name
- **Usage Metrics**: Input tokens, total cost
- **Performance**: Latency
- **Content**: Input text (respects privacy settings)
- **Metadata**: Custom fields, user ID, trace information
### For Errors
- **Error Details**: Error type, error message, stack trace
- **Context**: Model, provider, input that caused the error
- **Timing**: When the error occurred, request duration
## Environment Variables
| Variable | Required | Description |
|----------|----------|-------------|
| `POSTHOG_API_KEY` | Yes | Your PostHog project API key |
| `POSTHOG_API_URL` | No | PostHog API URL (defaults to https://app.posthog.com) |
## Troubleshooting
### 1. Missing API Key
```
Error: POSTHOG_API_KEY is not set
```
Set your PostHog API key:
```python
import os
os.environ["POSTHOG_API_KEY"] = "your-api-key"
```
### 2. Custom PostHog Instance
If you're using a self-hosted PostHog instance:
```python
import os
os.environ["POSTHOG_API_URL"] = "https://your-posthog-instance.com"
```
### 3. Events Not Appearing
- Check that your API key is correct
- Verify network connectivity to PostHog
- Events may take a few minutes to appear in PostHog dashboard
docs/my-website/docs/providers/bedrock.md
+157 -2
View File
@@ -308,6 +308,65 @@ print(response)
</TabItem>
</Tabs>
## Usage - Request Metadata
Attach metadata to Bedrock requests for logging and cost attribution.
<Tabs>
<TabItem value="sdk" label="SDK">
```python
import os
from litellm import completion
os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = ""
response = completion(
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
messages=[{"role": "user", "content": "Hello, how are you?"}],
requestMetadata={
"cost_center": "engineering",
"user_id": "user123"
}
)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
**Set on yaml**
```yaml
model_list:
- model_name: bedrock-claude-v1
litellm_params:
model: bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0
requestMetadata:
cost_center: "engineering"
```
**Set on request**
```python
import openai
client = openai.OpenAI(
api_key="anything",
base_url="http://0.0.0.0:4000"
)
response = client.chat.completions.create(
model="bedrock-claude-v1",
messages=[{"role": "user", "content": "Hello"}],
extra_body={
"requestMetadata": {"cost_center": "engineering"}
}
)
```
</TabItem>
</Tabs>
## Usage - Function Calling / Tool calling
LiteLLM supports tool calling via Bedrock's Converse and Invoke API's.
@@ -889,6 +948,19 @@ curl http://0.0.0.0:4000/v1/chat/completions \
Example of using [Bedrock Guardrails with LiteLLM](https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails-use-converse-api.html)
### Selective Content Moderation with `guarded_text`
LiteLLM supports selective content moderation using the `guarded_text` content type. This allows you to wrap only specific content that should be moderated by Bedrock Guardrails, rather than evaluating the entire conversation.
**How it works:**
- Content with `type: "guarded_text"` gets automatically wrapped in `guardrailConverseContent` blocks
- Only the wrapped content is evaluated by Bedrock Guardrails
- Regular content with `type: "text"` bypasses guardrail evaluation
:::note
If `guarded_text` is not used, the entire conversation history will be sent to the guardrail for evaluation, which can increase latency and costs.
:::
<Tabs>
<TabItem value="sdk" label="LiteLLM SDK">
@@ -915,6 +987,24 @@ response = completion(
"trace": "disabled", # The trace behavior for the guardrail. Can either be "disabled" or "enabled"
},
)
# Selective guardrail usage with guarded_text - only specific content is evaluated
response_guard = completion(
model="anthropic.claude-v2",
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "What is the main topic of this legal document?"},
{"type": "guarded_text", "text": "This document contains sensitive legal information that should be moderated by guardrails."}
]
}
],
guardrailConfig={
"guardrailIdentifier": "gr-abc123",
"guardrailVersion": "DRAFT"
}
)
```
</TabItem>
<TabItem value="proxy" label="Proxy on request">
@@ -993,7 +1083,20 @@ response = client.chat.completions.create(model="bedrock-claude-v1", messages =
temperature=0.7
)
print(response)
# For adding selective guardrail usage with guarded_text
response_guard = client.chat.completions.create(model="bedrock-claude-v1", messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is the main topic of this legal document?"},
{"type": "guarded_text", "text": "This document contains sensitive legal information that should be moderated by guardrails."}
]
}
],
temperature=0.7
)
print(response_guard)
```
</TabItem>
</Tabs>
@@ -1777,6 +1880,7 @@ Here's an example of using a bedrock model with LiteLLM. For a complete list, re
| Mistral 7B Instruct | `completion(model='bedrock/mistral.mistral-7b-instruct-v0:2', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
| 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']` |
## Bedrock Embedding
### API keys
@@ -1798,11 +1902,29 @@ response = embedding(
print(response)
```
#### Titan V2 - encoding_format support
```python
from litellm import embedding
# Float format (default)
response = embedding(
model="bedrock/amazon.titan-embed-text-v2:0",
input=["good morning from litellm"],
encoding_format="float" # Returns float array
)
# Binary format
response = embedding(
model="bedrock/amazon.titan-embed-text-v2:0",
input=["good morning from litellm"],
encoding_format="base64" # Returns base64 encoded binary
)
```
## Supported AWS Bedrock Embedding Models
| Model Name | Usage | Supported Additional OpenAI params |
|----------------------|---------------------------------------------|-----|
| Titan Embeddings V2 | `embedding(model="bedrock/amazon.titan-embed-text-v2:0", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py#L59) |
| Titan Embeddings V2 | `embedding(model="bedrock/amazon.titan-embed-text-v2:0", input=input)` | `dimensions`, `encoding_format` |
| Titan Embeddings - V1 | `embedding(model="bedrock/amazon.titan-embed-text-v1", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py#L53)
| Titan Multimodal Embeddings | `embedding(model="bedrock/amazon.titan-embed-image-v1", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_multimodal_transformation.py#L28) |
| Cohere Embeddings - English | `embedding(model="bedrock/cohere.embed-english-v3", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/cohere_transformation.py#L18)
@@ -1891,6 +2013,39 @@ curl -L -X POST 'http://0.0.0.0:4000/v1/images/generations' \
</TabItem>
</Tabs>
### Using Inference Profiles with Image Generation
For AWS Bedrock Application Inference Profiles with image generation, use the `model_id` parameter to specify the inference profile ARN:
<Tabs>
<TabItem value="sdk" label="SDK">
```python
from litellm import image_generation
response = image_generation(
model="bedrock/amazon.nova-canvas-v1:0",
model_id="arn:aws:bedrock:eu-west-1:000000000000:application-inference-profile/a0a0a0a0a0a0",
prompt="A cute baby sea otter"
)
print(f"response: {response}")
```
</TabItem>
<TabItem value="proxy" label="PROXY">
```yaml
model_list:
- model_name: nova-canvas-inference-profile
litellm_params:
model: bedrock/amazon.nova-canvas-v1:0
model_id: arn:aws:bedrock:eu-west-1:000000000000:application-inference-profile/a0a0a0a0a0a0
aws_region_name: "eu-west-1"
```
</TabItem>
</Tabs>
## Supported AWS Bedrock Image Generation Models
| Model Name | Function Call |
docs/my-website/docs/providers/bedrock_embedding.md
+95
View File
@@ -0,0 +1,95 @@
# Bedrock Embedding
## Supported Embedding Models
| Provider | LiteLLM Route | AWS Documentation |
|----------|---------------|-------------------|
| Amazon Titan | `bedrock/amazon.*` | [Amazon Titan Embeddings](https://docs.aws.amazon.com/bedrock/latest/userguide/titan-embedding-models.html) |
| Cohere | `bedrock/cohere.*` | [Cohere Embeddings](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-embed.html) |
| TwelveLabs | `bedrock/us.twelvelabs.*` | [TwelveLabs](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-twelvelabs.html) |
### API keys
This can be set as env variables or passed as **params to litellm.embedding()**
```python
import os
os.environ["AWS_ACCESS_KEY_ID"] = "" # Access key
os.environ["AWS_SECRET_ACCESS_KEY"] = "" # Secret access key
os.environ["AWS_REGION_NAME"] = "" # us-east-1, us-east-2, us-west-1, us-west-2
```
## Usage
### LiteLLM Python SDK
```python
from litellm import embedding
response = embedding(
model="bedrock/amazon.titan-embed-text-v1",
input=["good morning from litellm"],
)
print(response)
```
### LiteLLM Proxy Server
#### 1. Setup config.yaml
```yaml
model_list:
- model_name: titan-embed-v1
litellm_params:
model: bedrock/amazon.titan-embed-text-v1
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
aws_region_name: us-east-1
- model_name: titan-embed-v2
litellm_params:
model: bedrock/amazon.titan-embed-text-v2:0
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
aws_region_name: us-east-1
```
#### 2. Start Proxy
```bash
litellm --config /path/to/config.yaml
```
#### 3. Use with OpenAI Python SDK
```python
import openai
client = openai.OpenAI(
api_key="anything",
base_url="http://0.0.0.0:4000"
)
response = client.embeddings.create(
input=["good morning from litellm"],
model="titan-embed-v1"
)
print(response)
```
#### 4. Use with LiteLLM Python SDK
```python
import litellm
response = litellm.embedding(
model="titan-embed-v1", # model alias from config.yaml
input=["good morning from litellm"],
api_base="http://0.0.0.0:4000",
api_key="anything"
)
print(response)
```
## Supported AWS Bedrock Embedding Models
| Model Name | Usage | Supported Additional OpenAI params |
|----------------------|---------------------------------------------|-----|
| Titan Embeddings V2 | `embedding(model="bedrock/amazon.titan-embed-text-v2:0", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py#L59) |
| Titan Embeddings - V1 | `embedding(model="bedrock/amazon.titan-embed-text-v1", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py#L53)
| Titan Multimodal Embeddings | `embedding(model="bedrock/amazon.titan-embed-image-v1", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_multimodal_transformation.py#L28) |
| TwelveLabs Marengo Embed 2.7 | `embedding(model="bedrock/us.twelvelabs.marengo-embed-2-7-v1:0", input=input)` | Supports multimodal input (text, video, audio, image) |
| Cohere Embeddings - English | `embedding(model="bedrock/cohere.embed-english-v3", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/cohere_transformation.py#L18)
| Cohere Embeddings - Multilingual | `embedding(model="bedrock/cohere.embed-multilingual-v3", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/cohere_transformation.py#L18)
### Advanced - [Drop Unsupported Params](https://docs.litellm.ai/docs/completion/drop_params#openai-proxy-usage)
### Advanced - [Pass model/provider-specific Params](https://docs.litellm.ai/docs/completion/provider_specific_params#proxy-usage)
docs/my-website/docs/providers/compactifai.md
+2 -2
View File
@@ -4,7 +4,7 @@ import TabItem from '@theme/TabItem';
# CompactifAI
https://docs.compactif.ai/
CompactifAI offers highly compressed versions of leading language models, delivering up to **70% lower inference costs**, **4x throughput gains**, and **low-latency inference** with minimal quality loss (<5%). CompactifAI's OpenAI-compatible API makes integration straightforward, enabling developers to build ultra-efficient, scalable AI applications with superior concurrency and resource efficiency.
CompactifAI offers highly compressed versions of leading language models, delivering up to **70% lower inference costs**, **4x throughput gains**, and **low-latency inference** with minimal quality loss (under 5%). CompactifAI's OpenAI-compatible API makes integration straightforward, enabling developers to build ultra-efficient, scalable AI applications with superior concurrency and resource efficiency.
| Property | Details |
|-------|-------|
@@ -192,7 +192,7 @@ Common model formats:
## Benefits
- **Cost Efficient**: Up to 70% lower inference costs compared to standard models
- **High Performance**: 4x throughput gains with minimal quality loss (<5%)
- **High Performance**: 4x throughput gains with minimal quality loss (under 5%)
- **Low Latency**: Optimized for fast response times
- **Drop-in Replacement**: Full OpenAI API compatibility
- **Scalable**: Superior concurrency and resource efficiency
docs/my-website/docs/providers/vertex.md
+38 -144
View File
@@ -2509,150 +2509,6 @@ print("response from proxy", response)
</TabItem>
</Tabs>
## **Batch APIs**
Just add the following Vertex env vars to your environment.
```bash
# GCS Bucket settings, used to store batch prediction files in
export GCS_BUCKET_NAME = "litellm-testing-bucket" # the bucket you want to store batch prediction files in
export GCS_PATH_SERVICE_ACCOUNT="/path/to/service_account.json" # path to your service account json file
# Vertex /batch endpoint settings, used for LLM API requests
export GOOGLE_APPLICATION_CREDENTIALS="/path/to/service_account.json" # path to your service account json file
export VERTEXAI_LOCATION="us-central1" # can be any vertex location
export VERTEXAI_PROJECT="my-test-project"
```
### Usage
#### 1. Create a file of batch requests for vertex
LiteLLM expects the file to follow the **[OpenAI batches files format](https://platform.openai.com/docs/guides/batch)**
Each `body` in the file should be an **OpenAI API request**
Create a file called `vertex_batch_completions.jsonl` in the current working directory, the `model` should be the Vertex AI model name
```
{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gemini-1.5-flash-001", "messages": [{"role": "system", "content": "You are a helpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}}
{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gemini-1.5-flash-001", "messages": [{"role": "system", "content": "You are an unhelpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}}
```
#### 2. Upload a File of batch requests
For `vertex_ai` litellm will upload the file to the provided `GCS_BUCKET_NAME`
```python
import os
oai_client = OpenAI(
api_key="sk-1234", # litellm proxy API key
base_url="http://localhost:4000" # litellm proxy base url
)
file_name = "vertex_batch_completions.jsonl" #
_current_dir = os.path.dirname(os.path.abspath(__file__))
file_path = os.path.join(_current_dir, file_name)
file_obj = oai_client.files.create(
file=open(file_path, "rb"),
purpose="batch",
extra_body={"custom_llm_provider": "vertex_ai"}, # tell litellm to use vertex_ai for this file upload
)
```
**Expected Response**
```json
{
"id": "gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/d3f198cd-c0d1-436d-9b1e-28e3f282997a",
"bytes": 416,
"created_at": 1733392026,
"filename": "litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/d3f198cd-c0d1-436d-9b1e-28e3f282997a",
"object": "file",
"purpose": "batch",
"status": "uploaded",
"status_details": null
}
```
#### 3. Create a batch
```python
batch_input_file_id = file_obj.id # use `file_obj` from step 2
create_batch_response = oai_client.batches.create(
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id=batch_input_file_id, # example input_file_id = "gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/c2b1b785-252b-448c-b180-033c4c63b3ce"
extra_body={"custom_llm_provider": "vertex_ai"}, # tell litellm to use `vertex_ai` for this batch request
)
```
**Expected Response**
```json
{
"id": "3814889423749775360",
"completion_window": "24hrs",
"created_at": 1733392026,
"endpoint": "",
"input_file_id": "gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/d3f198cd-c0d1-436d-9b1e-28e3f282997a",
"object": "batch",
"status": "validating",
"cancelled_at": null,
"cancelling_at": null,
"completed_at": null,
"error_file_id": null,
"errors": null,
"expired_at": null,
"expires_at": null,
"failed_at": null,
"finalizing_at": null,
"in_progress_at": null,
"metadata": null,
"output_file_id": "gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001",
"request_counts": null
}
```
#### 4. Retrieve a batch
```python
retrieved_batch = oai_client.batches.retrieve(
batch_id=create_batch_response.id,
extra_body={"custom_llm_provider": "vertex_ai"}, # tell litellm to use `vertex_ai` for this batch request
)
```
**Expected Response**
```json
{
"id": "3814889423749775360",
"completion_window": "24hrs",
"created_at": 1736500100,
"endpoint": "",
"input_file_id": "gs://example-bucket-1-litellm/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/7b2e47f5-3dd4-436d-920f-f9155bbdc952",
"object": "batch",
"status": "completed",
"cancelled_at": null,
"cancelling_at": null,
"completed_at": null,
"error_file_id": null,
"errors": null,
"expired_at": null,
"expires_at": null,
"failed_at": null,
"finalizing_at": null,
"in_progress_at": null,
"metadata": null,
"output_file_id": "gs://example-bucket-1-litellm/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001",
"request_counts": null
}
```
## **Fine Tuning APIs**
@@ -2758,6 +2614,44 @@ curl http://localhost:4000/v1/fine_tuning/jobs \
</Tabs>
## Labels
Google enables you to add custom metadata to its `generateContent` and `streamGenerateContent` calls.
This mechanism is useful in Vertex AI because it allows costs and usage tracking over multiple
different applications or users.
### Usage
You can use that feature through LiteLLM by sending `labels` or `metadata` field in your requests.
If the client sets the `labels` field in the request to the LiteLLM,
the LiteLLM will pass the `labels` field to the Vertex AI backend.
If the client sets the `metadata` field in the request to the LiteLLM and the `labels` field is not set,
the LiteLLM will create the `labels` field filled with `metadata` key/value pairs for all string values and
pass it to the Vertex AI backend.
Here is an example JSON request demonstrating the labels usage:
```json
{
"model": "gemini-2.0-flash-lite",
"messages": [
{ "role": "user", "content": "respond in 20 words. who are you?" }
],
"labels": {
"client_app": "acme_comp_financial_app",
"department": "finance",
"project": "acme_ai"
}
}
```
## Extra
### Using `GOOGLE_APPLICATION_CREDENTIALS`
docs/my-website/docs/providers/vertex_batch.md
+264
View File
@@ -0,0 +1,264 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
## **Batch APIs**
Just add the following Vertex env vars to your environment.
```bash
# GCS Bucket settings, used to store batch prediction files in
export GCS_BUCKET_NAME="my-batch-bucket" # the bucket you want to store batch prediction files in
export GCS_PATH_SERVICE_ACCOUNT="/path/to/service_account.json" # path to your service account json file
# Vertex /batch endpoint settings, used for LLM API requests
export GOOGLE_APPLICATION_CREDENTIALS="/path/to/service_account.json" # path to your service account json file
export VERTEXAI_LOCATION="us-central1" # can be any vertex location
export VERTEXAI_PROJECT="my-project"
```
### Usage
Follow this complete workflow: create JSONL file → upload file → create batch → retrieve batch status → get file content
#### 1. Create a JSONL file of batch requests
LiteLLM expects the file to follow the **[OpenAI batches files format](https://platform.openai.com/docs/guides/batch)**.
Each `body` in the file should be an **OpenAI API request**.
Create a file called `batch_requests.jsonl` with your requests:
```jsonl
{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gemini-2.5-flash-lite", "messages": [{"role": "system", "content": "You are a helpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}}
{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gemini-2.5-flash-lite", "messages": [{"role": "system", "content": "You are an unhelpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}}
```
#### 2. Upload the file
Upload your JSONL file. For `vertex_ai`, the file will be stored in your configured GCS bucket provided by `GCS_BUCKET_NAME`.
<Tabs>
<TabItem value="python" label="Python">
```python showLineNumbers title="upload_file.py"
from openai import OpenAI
oai_client = OpenAI(
api_key="sk-1234", # litellm proxy API key
base_url="http://localhost:4000" # litellm proxy base url
)
file_obj = oai_client.files.create(
file=open("batch_requests.jsonl", "rb"),
purpose="batch",
extra_body={"custom_llm_provider": "vertex_ai"}
)
print(f"File uploaded with ID: {file_obj.id}")
```
</TabItem>
<TabItem value="curl" label="Curl">
```bash showLineNumbers title="Upload File"
curl --request POST \
--url http://localhost:4000/v1/files \
--header 'Content-Type: multipart/form-data' \
--form purpose=batch \
--form file=@batch_requests.jsonl \
--form custom_llm_provider=vertex_ai
```
</TabItem>
</Tabs>
**Expected Response:**
```json
{
"id": "gs://my-batch-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash-lite/abc123-def4-5678-9012-34567890abcd",
"bytes": 416,
"created_at": 1758303684,
"filename": "litellm-vertex-files/publishers/google/models/gemini-2.5-flash-lite/abc123-def4-5678-9012-34567890abcd",
"object": "file",
"purpose": "batch",
"status": "uploaded",
"expires_at": null,
"status_details": null
}
```
#### 3. Create a batch
Create a batch job using the uploaded file ID.
<Tabs>
<TabItem value="python" label="Python">
```python showLineNumbers title="create_batch.py"
batch_input_file_id = file_obj.id # from step 2
create_batch_response = oai_client.batches.create(
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id=batch_input_file_id, # e.g. "gs://my-batch-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash-lite/abc123-def4-5678-9012-34567890abcd"
extra_body={"custom_llm_provider": "vertex_ai"}
)
print(f"Batch created with ID: {create_batch_response.id}")
```
</TabItem>
<TabItem value="curl" label="Curl">
```bash showLineNumbers title="Create Batch Request"
curl --request POST \
--url http://localhost:4000/v1/batches \
--header 'Content-Type: application/json' \
--data '{
"input_file_id": "gs://my-batch-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash-lite/abc123-def4-5678-9012-34567890abcd",
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"custom_llm_provider": "vertex_ai"
}'
```
</TabItem>
</Tabs>
**Expected Response:**
```json
{
"id": "7814463557919047680",
"completion_window": "24hrs",
"created_at": 1758328011,
"endpoint": "",
"input_file_id": "gs://my-batch-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash-lite/abc123-def4-5678-9012-34567890abcd",
"object": "batch",
"status": "validating",
"cancelled_at": null,
"cancelling_at": null,
"completed_at": null,
"error_file_id": null,
"errors": null,
"expired_at": null,
"expires_at": null,
"failed_at": null,
"finalizing_at": null,
"in_progress_at": null,
"metadata": null,
"output_file_id": "gs://my-batch-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash-lite",
"request_counts": null,
"usage": null
}
```
#### 4. Retrieve batch status
Check the status of your batch job. The batch will progress through states: `validating` → `in_progress` → `completed`.
<Tabs>
<TabItem value="python" label="Python">
```python showLineNumbers title="retrieve_batch.py"
retrieved_batch = oai_client.batches.retrieve(
batch_id=create_batch_response.id, # Created batch id, e.g. 7814463557919047680
extra_body={"custom_llm_provider": "vertex_ai"}
)
print(f"Batch status: {retrieved_batch.status}")
if retrieved_batch.status == "completed":
print(f"Output file: {retrieved_batch.output_file_id}")
```
</TabItem>
<TabItem value="curl" label="Curl">
```bash showLineNumbers title="Retrieve Batch Status"
curl --request GET \
--url 'http://localhost:4000/batches/7814463557919047680?provider=vertex_ai' \
--header 'Authorization: Bearer sk-1234'
```
</TabItem>
</Tabs>
**Expected Response (when completed):**
```json
{
"id": "7814463557919047680",
"completion_window": "24hrs",
"created_at": 1758328011,
"endpoint": "",
"input_file_id": "gs://my-batch-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash-lite/abc123-def4-5678-9012-34567890abcd",
"object": "batch",
"status": "completed",
"cancelled_at": null,
"cancelling_at": null,
"completed_at": null,
"error_file_id": null,
"errors": null,
"expired_at": null,
"expires_at": null,
"failed_at": null,
"finalizing_at": null,
"in_progress_at": null,
"metadata": null,
"output_file_id": "gs://my-batch-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash-lite/prediction-model-2025-09-19T21:26:51.569037Z/predictions.jsonl",
"request_counts": null,
"usage": null
}
```
#### 5. Get file content
Once the batch is completed, retrieve the results using the `output_file_id` from the batch response.
**Important:** The `output_file_id` must be URL encoded when used in the request path.
<Tabs>
<TabItem value="python" label="Python">
```python showLineNumbers title="get_file_content.py"
import urllib.parse
import json
output_file_id = retrieved_batch.output_file_id
# URL encode the file ID
encoded_file_id = urllib.parse.quote_plus(output_file_id)
# Get file content
file_content = oai_client.files.content(
file_id=encoded_file_id,
extra_body={"custom_llm_provider": "vertex_ai"}
)
# Process the results
for line in file_content.text.strip().split('\n'):
result = json.loads(line)
print(f"Request: {result['request']}")
print(f"Response: {result['response']}")
print("---")
```
</TabItem>
<TabItem value="curl" label="Curl">
```bash showLineNumbers title="Get File Content"
# Note: The file ID must be URL encoded
curl --request GET \
--url 'http://localhost:4000/files/gs%253A%252F%252Fmy-batch-bucket%252Flitellm-vertex-files%252Fpublishers%252Fgoogle%252Fmodels%252Fgemini-2.5-flash-lite%252Fprediction-model-2025-09-19T21%253A26%253A51.569037Z%252Fpredictions.jsonl/content?provider=vertex_ai' \
--header 'Authorization: Bearer sk-1234'
```
</TabItem>
</Tabs>
**Expected Response:**
The response contains JSONL format with one result per line:
```jsonl
{"status":"","processed_time":"2025-09-19T21:29:47.352+00:00","request":{"contents":[{"parts":[{"text":"Hello world!"}],"role":"user"}],"generationConfig":{"max_output_tokens":10},"system_instruction":{"parts":[{"text":"You are a helpful assistant."}]}},"response":{"candidates":[{"avgLogprobs":-0.48079710006713866,"content":{"parts":[{"text":"Hello there! It's nice to meet you"}],"role":"model"},"finishReason":"MAX_TOKENS"}],"createTime":"2025-09-19T21:29:47.484619Z","modelVersion":"gemini-2.5-flash-lite","responseId":"S8vNaIvKHdvshMIP_aOtuAg","usageMetadata":{"candidatesTokenCount":10,"candidatesTokensDetails":[{"modality":"TEXT","tokenCount":10}],"promptTokenCount":9,"promptTokensDetails":[{"modality":"TEXT","tokenCount":9}],"totalTokenCount":19,"trafficType":"ON_DEMAND"}}}
{"status":"","processed_time":"2025-09-19T21:29:47.358+00:00","request":{"contents":[{"parts":[{"text":"Hello world!"}],"role":"user"}],"generationConfig":{"max_output_tokens":10},"system_instruction":{"parts":[{"text":"You are an unhelpful assistant."}]}},"response":{"candidates":[{"avgLogprobs":-0.6168075137668185,"content":{"parts":[{"text":"I am unable to assist with this request."}],"role":"model"},"finishReason":"STOP"}],"createTime":"2025-09-19T21:29:47.470889Z","modelVersion":"gemini-2.5-flash-lite","responseId":"S8vNaOneHISShMIP28nA8QQ","usageMetadata":{"candidatesTokenCount":9,"candidatesTokensDetails":[{"modality":"TEXT","tokenCount":9}],"promptTokenCount":9,"promptTokensDetails":[{"modality":"TEXT","tokenCount":9}],"totalTokenCount":18,"trafficType":"ON_DEMAND"}}}
```
docs/my-website/docs/proxy/budget_reset_and_tz.md
+1
View File
@@ -29,5 +29,6 @@ Common timezone values:
- `US/Pacific` - Pacific Time
- `Europe/London` - UK Time
- `Asia/Kolkata` - Indian Standard Time (IST)
- `Asia/Bangkok` - Indochina Time (ICT)
- `Asia/Tokyo` - Japan Standard Time
- `Australia/Sydney` - Australian Eastern Time
docs/my-website/docs/proxy/config_settings.md
+34
View File
@@ -93,6 +93,8 @@ callback_settings:
general_settings:
completion_model: string
store_prompts_in_spend_logs: boolean
forward_client_headers_to_llm_api: boolean
disable_spend_logs: boolean # turn off writing each transaction to the db
disable_master_key_return: boolean # turn off returning master key on UI (checked on '/user/info' endpoint)
disable_retry_on_max_parallel_request_limit_error: boolean # turn off retries when max parallel request limit is reached
@@ -121,6 +123,35 @@ general_settings:
alerting: ["slack", "email"]
alerting_threshold: 0
use_client_credentials_pass_through_routes: boolean # use client credentials for all pass through routes like "/vertex-ai", /bedrock/. When this is True Virtual Key auth will not be applied on these endpoints
router_settings:
routing_strategy: simple-shuffle # Literal["simple-shuffle", "least-busy", "usage-based-routing","latency-based-routing"], default="simple-shuffle" - RECOMMENDED for best performance
redis_host: <your-redis-host> # string
redis_password: <your-redis-password> # string
redis_port: <your-redis-port> # string
enable_pre_call_checks: true # bool - Before call is made check if a call is within model context window
allowed_fails: 3 # cooldown model if it fails > 1 call in a minute.
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
retry_policy: { # Dict[str, int]: retry policy for different types of exceptions
"AuthenticationErrorRetries": 3,
"TimeoutErrorRetries": 3,
"RateLimitErrorRetries": 3,
"ContentPolicyViolationErrorRetries": 4,
"InternalServerErrorRetries": 4
}
allowed_fails_policy: {
"BadRequestErrorAllowedFails": 1000, # Allow 1000 BadRequestErrors before cooling down a deployment
"AuthenticationErrorAllowedFails": 10, # int
"TimeoutErrorAllowedFails": 12, # int
"RateLimitErrorAllowedFails": 10000, # int
"ContentPolicyViolationErrorAllowedFails": 15, # int
"InternalServerErrorAllowedFails": 20, # int
}
content_policy_fallbacks=[{"claude-2": ["my-fallback-model"]}] # List[Dict[str, List[str]]]: Fallback model for content policy violations
fallbacks=[{"claude-2": ["my-fallback-model"]}] # List[Dict[str, List[str]]]: Fallback model for all errors
```
### litellm_settings - Reference
@@ -659,6 +690,8 @@ router_settings:
| PILLAR_API_KEY | API key for Pillar API Guardrails
| PILLAR_ON_FLAGGED_ACTION | Action to take when content is flagged ('block' or 'monitor')
| POD_NAME | Pod name for the server, this will be [emitted to `datadog` logs](https://docs.litellm.ai/docs/proxy/logging#datadog) as `POD_NAME`
| POSTHOG_API_KEY | API key for PostHog analytics integration
| POSTHOG_API_URL | Base URL for PostHog API (defaults to https://us.i.posthog.com)
| PREDIBASE_API_BASE | Base URL for Predibase API
| PRESIDIO_ANALYZER_API_BASE | Base URL for Presidio Analyzer service
| PRESIDIO_ANONYMIZER_API_BASE | Base URL for Presidio Anonymizer service
@@ -739,3 +772,4 @@ router_settings:
| WEBHOOK_URL | URL for receiving webhooks from external services
| SPEND_LOG_RUN_LOOPS | Constant for setting how many runs of 1000 batch deletes should spend_log_cleanup task run |
| SPEND_LOG_CLEANUP_BATCH_SIZE | Number of logs deleted per batch during cleanup. Default is 1000 |
| COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY | Maximum size for CoroutineChecker in-memory cache. Default is 1000 |
docs/my-website/docs/proxy/deploy.md
+1
View File
@@ -13,6 +13,7 @@ To start using Litellm, run the following commands in a shell:
```bash
# Get the code
curl -O https://raw.githubusercontent.com/BerriAI/litellm/main/docker-compose.yml
curl -O https://raw.githubusercontent.com/BerriAI/litellm/main/prometheus.yml
# Add the master key - you can change this after setup
echo 'LITELLM_MASTER_KEY="sk-1234"' > .env
docs/my-website/docs/proxy/dynamic_rate_limit.md
+241
View File
@@ -0,0 +1,241 @@
# Dynamic TPM/RPM Allocation
Prevent projects from gobbling too much tpm/rpm.
Dynamically allocate TPM/RPM quota to api keys, based on active keys in that minute. [**See Code**](https://github.com/BerriAI/litellm/blob/9bffa9a48e610cc6886fc2dce5c1815aeae2ad46/litellm/proxy/hooks/dynamic_rate_limiter.py#L125)
## Quick Start Usage
1. Setup config.yaml
```yaml showLineNumbers title="config.yaml"
model_list:
- model_name: my-fake-model
litellm_params:
model: gpt-3.5-turbo
api_key: my-fake-key
mock_response: hello-world
tpm: 60
litellm_settings:
callbacks: ["dynamic_rate_limiter_v3"]
general_settings:
master_key: sk-1234 # OR set `LITELLM_MASTER_KEY=".."` in your .env
database_url: postgres://.. # OR set `DATABASE_URL=".."` in your .env
```
2. Start proxy
```bash
litellm --config /path/to/config.yaml
```
3. Test it!
```python showLineNumbers title="test.py"
"""
- Run 2 concurrent teams calling same model
- model has 60 TPM
- Mock response returns 30 total tokens / request
- Each team will only be able to make 1 request per minute
"""
import requests
from openai import OpenAI, RateLimitError
def create_key(api_key: str, base_url: str):
response = requests.post(
url="{}/key/generate".format(base_url),
json={},
headers={
"Authorization": "Bearer {}".format(api_key)
}
)
_response = response.json()
return _response["key"]
key_1 = create_key(api_key="sk-1234", base_url="http://0.0.0.0:4000")
key_2 = create_key(api_key="sk-1234", base_url="http://0.0.0.0:4000")
# call proxy with key 1 - works
openai_client_1 = OpenAI(api_key=key_1, base_url="http://0.0.0.0:4000")
response = openai_client_1.chat.completions.with_raw_response.create(
model="my-fake-model", messages=[{"role": "user", "content": "Hello world!"}],
)
print("Headers for call 1 - {}".format(response.headers))
_response = response.parse()
print("Total tokens for call - {}".format(_response.usage.total_tokens))
# call proxy with key 2 - works
openai_client_2 = OpenAI(api_key=key_2, base_url="http://0.0.0.0:4000")
response = openai_client_2.chat.completions.with_raw_response.create(
model="my-fake-model", messages=[{"role": "user", "content": "Hello world!"}],
)
print("Headers for call 2 - {}".format(response.headers))
_response = response.parse()
print("Total tokens for call - {}".format(_response.usage.total_tokens))
# call proxy with key 2 - fails
try:
openai_client_2.chat.completions.with_raw_response.create(model="my-fake-model", messages=[{"role": "user", "content": "Hey, how's it going?"}])
raise Exception("This should have failed!")
except RateLimitError as e:
print("This was rate limited b/c - {}".format(str(e)))
```
**Expected Response**
```
This was rate limited b/c - Error code: 429 - {'error': {'message': {'error': 'Key=<hashed_token> over available TPM=0. Model TPM=0, Active keys=2'}, 'type': 'None', 'param': 'None', 'code': 429}}
```
## [BETA] Set Priority / Reserve Quota
Reserve TPM/RPM capacity for different environments or use cases. This ensures critical production workloads always have guaranteed capacity, while development or lower-priority tasks use remaining quota.
**Use Cases:**
- Production vs Development environments
- Real-time applications vs batch processing
- Critical services vs experimental features
:::tip
Reserving TPM/RPM on keys based on priority is a premium feature. Please [get an enterprise license](./enterprise.md) for it.
:::
### How Priority Reservation Works
Priority reservation allocates a percentage of your model's total TPM/RPM to specific priority levels. Keys with higher priority get guaranteed access to their reserved quota first.
**Example Scenario:**
- Model has 10 RPM total capacity
- Priority reservation: `{"prod": 0.9, "dev": 0.1}`
- Result: Production keys get 9 RPM guaranteed, Development keys get 1 RPM guaranteed
### Configuration
#### 1. Setup config.yaml
```yaml showLineNumbers title="config.yaml"
model_list:
- model_name: gpt-3.5-turbo
litellm_params:
model: "gpt-3.5-turbo"
api_key: os.environ/OPENAI_API_KEY
rpm: 10 # Total model capacity
litellm_settings:
callbacks: ["dynamic_rate_limiter_v3"]
priority_reservation:
"prod": 0.9 # 90% reserved for production (9 RPM)
"dev": 0.1 # 10% reserved for development (1 RPM)
general_settings:
master_key: sk-1234 # OR set `LITELLM_MASTER_KEY=".."` in your .env
database_url: postgres://.. # OR set `DATABASE_URL=".."` in your.env
```
**Configuration Details:**
`priority_reservation`: Dict[str, float]
- **Key (str)**: Priority level name (can be any string like "prod", "dev", "critical", etc.)
- **Value (float)**: Percentage of total TPM/RPM to reserve (0.0 to 1.0)
- **Note**: Values should sum to 1.0 or less
**Start Proxy**
```bash
litellm --config /path/to/config.yaml
```
#### 2. Create Keys with Priority Levels
**Production Key:**
```bash
curl -X POST 'http://0.0.0.0:4000/key/generate' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{
"metadata": {"priority": "prod"}
}'
```
**Development Key:**
```bash
curl -X POST 'http://0.0.0.0:4000/key/generate' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{
"metadata": {"priority": "dev"}
}'
```
**Expected Response for both:**
```json
{
"key": "sk-...",
"metadata": {"priority": "prod"}, // or "dev"
...
}
```
#### 3. Test Priority Allocation
**Test Production Key (should get 9 RPM):**
```bash
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-prod-key' \
-d '{
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "Hello from prod"}]
}'
```
**Test Development Key (should get 1 RPM):**
```bash
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-dev-key' \
-d '{
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "Hello from dev"}]
}'
```
### Expected Behavior
With the configuration above:
1. **Production keys** can make up to 9 requests per minute
2. **Development keys** can make up to 1 request per minute
3. Production requests are never blocked by development usage
**Rate Limit Error Example:**
```json
{
"error": {
"message": "Key=sk-dev-... over available RPM=0. Model RPM=10, Reserved RPM for priority 'dev'=1, Active keys=1",
"type": "rate_limit_exceeded",
"code": 429
}
}
```
### Demo Video
This video walks through setting up dynamic rate limiting with priority reservation and locust tests to validate the behavior.
<iframe width="840" height="500" src="https://www.loom.com/embed/1b54b93139ee415d959402cc0629f3f7
" frameborder="0" webkitallowfullscreen mozallowfullscreen allowfullscreen></iframe>
docs/my-website/docs/proxy/logging_spec.md
+83
View File
@@ -61,6 +61,11 @@ Inherits from `StandardLoggingUserAPIKeyMetadata` and adds:
| `requester_metadata` | `Optional[dict]` | Additional requester metadata |
| `vector_store_request_metadata` | `Optional[List[StandardLoggingVectorStoreRequest]]` | Vector store request metadata |
| `requester_custom_headers` | Dict[str, str] | Any custom (`x-`) headers sent by the client to the proxy. |
| `prompt_management_metadata` | `Optional[StandardLoggingPromptManagementMetadata]` | Prompt management and versioning metadata |
| `mcp_tool_call_metadata` | `Optional[StandardLoggingMCPToolCall]` | MCP (Model Context Protocol) tool call information and cost tracking |
| `applied_guardrails` | `Optional[List[str]]` | List of applied guardrail names |
| `usage_object` | `Optional[dict]` | Raw usage object from the LLM provider |
| `cold_storage_object_key` | `Optional[str]` | S3/GCS object key for cold storage retrieval |
| `guardrail_information` | `Optional[StandardLoggingGuardrailInformation]` | Guardrail information |
@@ -145,4 +150,82 @@ A literal type with two possible values:
| `duration` | `Optional[float]` | Duration of the guardrail in seconds |
| `masked_entity_count` | `Optional[Dict[str, int]]` | Count of masked entities |
## StandardLoggingPromptManagementMetadata
Used for tracking prompt versioning and management information.
| Field | Type | Description |
|-------|------|-------------|
| `prompt_id` | `str` | **Required**. Unique identifier for the prompt template or version |
| `prompt_variables` | `Optional[dict]` | Variables/parameters used in the prompt template (e.g., `{"user_name": "John", "context": "support"}`) |
| `prompt_integration` | `str` | **Required**. Integration or system managing the prompt (e.g., `"langfuse"`, `"promptlayer"`, `"custom"`) |
## StandardLoggingMCPToolCall
Used to track Model Context Protocol (MCP) tool calls within LiteLLM requests. This provides detailed logging for external tool integrations.
| Field | Type | Description |
|-------|------|-------------|
| `name` | `str` | **Required**. The name of the tool being called (e.g., `"get_weather"`, `"search_database"`) |
| `arguments` | `dict` | **Required**. Arguments passed to the tool as key-value pairs |
| `result` | `Optional[dict]` | The response/result returned by the tool execution (populated by custom logging hooks) |
| `mcp_server_name` | `Optional[str]` | Name of the MCP server that handled the tool call (e.g., `"weather-service"`, `"database-connector"`) |
| `mcp_server_logo_url` | `Optional[str]` | URL for the MCP server's logo (used for UI display in LiteLLM dashboard) |
| `namespaced_tool_name` | `Optional[str]` | Fully qualified tool name including server prefix (e.g., `"deepwiki-mcp/get_page_content"`, `"github-mcp/create_issue"`) |
| `mcp_server_cost_info` | `Optional[MCPServerCostInfo]` | Cost tracking information for the tool call |
### MCPServerCostInfo
Cost tracking structure for MCP server tool calls:
| Field | Type | Description |
|-------|------|-------------|
| `default_cost_per_query` | `Optional[float]` | Default cost in USD for any tool call to this MCP server |
| `tool_name_to_cost_per_query` | `Optional[Dict[str, float]]` | Per-tool cost mapping for granular pricing (e.g., `{"search": 0.01, "create": 0.05}`) |
### Usage
```python
# Basic MCP tool call metadata
mcp_tool_call = {
"name": "search_documents",
"arguments": {
"query": "machine learning tutorials",
"limit": 10,
"filter": "type:pdf"
},
"mcp_server_name": "document-search-service",
"namespaced_tool_name": "docs-mcp/search_documents",
"mcp_server_cost_info": {
"default_cost_per_query": 0.02,
"tool_name_to_cost_per_query": {
"search_documents": 0.02,
"get_document": 0.01
}
}
}
# optional result field (via custom logging hooks)
mcp_tool_call_with_result = {
"name": "search_documents",
"arguments": {
"query": "machine learning tutorials",
"limit": 10,
"filter": "type:pdf"
},
"result": {
"documents": [...],
"total_found": 42,
"search_time_ms": 150
},
"mcp_server_name": "document-search-service",
"namespaced_tool_name": "docs-mcp/search_documents",
"mcp_server_cost_info": {
"default_cost_per_query": 0.02,
"tool_name_to_cost_per_query": {
"search_documents": 0.02,
"get_document": 0.01
}
}
}
```
docs/my-website/docs/proxy/team_budgets.md
-185
View File
@@ -178,188 +178,3 @@ Expect to see this metric on prometheus to track the Remaining Budget for the te
```shell
litellm_remaining_team_budget_metric{team_alias="QA Prod Bot",team_id="de35b29e-6ca8-4f47-b804-2b79d07aa99a"} 9.699999999999992e-06
```
### Dynamic TPM/RPM Allocation
Prevent projects from gobbling too much tpm/rpm.
Dynamically allocate TPM/RPM quota to api keys, based on active keys in that minute. [**See Code**](https://github.com/BerriAI/litellm/blob/9bffa9a48e610cc6886fc2dce5c1815aeae2ad46/litellm/proxy/hooks/dynamic_rate_limiter.py#L125)
1. Setup config.yaml
```yaml
model_list:
- model_name: my-fake-model
litellm_params:
model: gpt-3.5-turbo
api_key: my-fake-key
mock_response: hello-world
tpm: 60
litellm_settings:
callbacks: ["dynamic_rate_limiter"]
general_settings:
master_key: sk-1234 # OR set `LITELLM_MASTER_KEY=".."` in your .env
database_url: postgres://.. # OR set `DATABASE_URL=".."` in your .env
```
2. Start proxy
```bash
litellm --config /path/to/config.yaml
```
3. Test it!
```python
"""
- Run 2 concurrent teams calling same model
- model has 60 TPM
- Mock response returns 30 total tokens / request
- Each team will only be able to make 1 request per minute
"""
import requests
from openai import OpenAI, RateLimitError
def create_key(api_key: str, base_url: str):
response = requests.post(
url="{}/key/generate".format(base_url),
json={},
headers={
"Authorization": "Bearer {}".format(api_key)
}
)
_response = response.json()
return _response["key"]
key_1 = create_key(api_key="sk-1234", base_url="http://0.0.0.0:4000")
key_2 = create_key(api_key="sk-1234", base_url="http://0.0.0.0:4000")
# call proxy with key 1 - works
openai_client_1 = OpenAI(api_key=key_1, base_url="http://0.0.0.0:4000")
response = openai_client_1.chat.completions.with_raw_response.create(
model="my-fake-model", messages=[{"role": "user", "content": "Hello world!"}],
)
print("Headers for call 1 - {}".format(response.headers))
_response = response.parse()
print("Total tokens for call - {}".format(_response.usage.total_tokens))
# call proxy with key 2 - works
openai_client_2 = OpenAI(api_key=key_2, base_url="http://0.0.0.0:4000")
response = openai_client_2.chat.completions.with_raw_response.create(
model="my-fake-model", messages=[{"role": "user", "content": "Hello world!"}],
)
print("Headers for call 2 - {}".format(response.headers))
_response = response.parse()
print("Total tokens for call - {}".format(_response.usage.total_tokens))
# call proxy with key 2 - fails
try:
openai_client_2.chat.completions.with_raw_response.create(model="my-fake-model", messages=[{"role": "user", "content": "Hey, how's it going?"}])
raise Exception("This should have failed!")
except RateLimitError as e:
print("This was rate limited b/c - {}".format(str(e)))
```
**Expected Response**
```
This was rate limited b/c - Error code: 429 - {'error': {'message': {'error': 'Key=<hashed_token> over available TPM=0. Model TPM=0, Active keys=2'}, 'type': 'None', 'param': 'None', 'code': 429}}
```
#### ✨ [BETA] Set Priority / Reserve Quota
Reserve tpm/rpm capacity for projects in prod.
:::tip
Reserving tpm/rpm on keys based on priority is a premium feature. Please [get an enterprise license](./enterprise.md) for it.
:::
1. Setup config.yaml
```yaml
model_list:
- model_name: gpt-3.5-turbo
litellm_params:
model: "gpt-3.5-turbo"
api_key: os.environ/OPENAI_API_KEY
rpm: 100
litellm_settings:
callbacks: ["dynamic_rate_limiter"]
priority_reservation: {"dev": 0, "prod": 1}
general_settings:
master_key: sk-1234 # OR set `LITELLM_MASTER_KEY=".."` in your .env
database_url: postgres://.. # OR set `DATABASE_URL=".."` in your .env
```
priority_reservation:
- Dict[str, float]
- str: can be any string
- float: from 0 to 1. Specify the % of tpm/rpm to reserve for keys of this priority.
**Start Proxy**
```
litellm --config /path/to/config.yaml
```
2. Create a key with that priority
```bash
curl -X POST 'http://0.0.0.0:4000/key/generate' \
-H 'Authorization: Bearer <your-master-key>' \
-H 'Content-Type: application/json' \
-D '{
"metadata": {"priority": "dev"} # 👈 KEY CHANGE
}'
```
**Expected Response**
```
{
...
"key": "sk-.."
}
```
3. Test it!
```bash
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: sk-...' \ # 👈 key from step 2.
-d '{
"model": "gpt-3.5-turbo",
"messages": [
{
"role": "user",
"content": "what llm are you"
}
],
}'
```
**Expected Response**
```
Key=... over available RPM=0. Model RPM=100, Active keys=None
```
docs/my-website/docs/sdk/headers.md
+328
View File
@@ -0,0 +1,328 @@
# SDK Header Support
LiteLLM SDK provides comprehensive support for passing additional headers with API requests. This is essential for enterprise environments using API gateways, service meshes, and multi-tenant architectures.
## Overview
Headers can be passed to LiteLLM in three ways, with the following priority order:
1. **Request-specific headers** (highest priority)
2. **extra_headers parameter**
3. **Global litellm.headers** (lowest priority)
When the same header key is specified in multiple places, the higher priority value will be used.
## Usage Methods
### 1. Global Headers (litellm.headers)
Set headers that will be included in all API requests:
```python
import litellm
# Set global headers for all requests
litellm.headers = {
"X-API-Gateway-Key": "your-gateway-key",
"X-Company-ID": "acme-corp",
"X-Environment": "production"
}
# Now all completion calls will include these headers
response = litellm.completion(
model="claude-3-5-sonnet-latest",
messages=[{"role": "user", "content": "Hello"}]
)
```
### 2. Per-Request Headers (extra_headers)
Pass headers for specific requests using the `extra_headers` parameter:
```python
import litellm
response = litellm.completion(
model="gpt-4",
messages=[{"role": "user", "content": "Hello"}],
extra_headers={
"X-Request-ID": "req-12345",
"X-Tenant-ID": "tenant-abc",
"X-Custom-Auth": "bearer-token-xyz"
}
)
```
### 3. Request Headers (headers parameter)
Use the `headers` parameter for the highest priority header control:
```python
import litellm
response = litellm.completion(
model="claude-3-5-sonnet-latest",
messages=[{"role": "user", "content": "Hello"}],
headers={
"X-Priority-Header": "high-priority-value",
"Authorization": "Bearer custom-token"
}
)
```
### 4. Combining All Methods
You can combine all three methods. Headers will be merged with the priority order:
```python
import litellm
# Global headers (lowest priority)
litellm.headers = {
"X-Company-ID": "acme-corp",
"X-Shared-Header": "global-value"
}
response = litellm.completion(
model="gpt-4",
messages=[{"role": "user", "content": "Hello"}],
extra_headers={
"X-Request-ID": "req-12345",
"X-Shared-Header": "extra-value" # Overrides global
},
headers={
"X-Priority-Header": "important",
"X-Shared-Header": "request-value" # Overrides both global and extra
}
)
# Final headers sent to API:
# {
# "X-Company-ID": "acme-corp", # From global
# "X-Request-ID": "req-12345", # From extra_headers
# "X-Priority-Header": "important", # From headers
# "X-Shared-Header": "request-value" # From headers (highest priority)
# }
```
## Enterprise Use Cases
### API Gateway Integration (Apigee, Kong, AWS API Gateway)
```python
import litellm
# Set up headers for API gateway routing and authentication
litellm.headers = {
"X-API-Gateway-Key": "your-gateway-key",
"X-Route-Version": "v2"
}
# Per-tenant requests
response = litellm.completion(
model="claude-3-5-sonnet-latest",
messages=[{"role": "user", "content": "Analyze this data"}],
extra_headers={
"X-Tenant-ID": "tenant-123",
"X-Department": "engineering"
}
)
```
### Service Mesh (Istio, Linkerd)
```python
import litellm
response = litellm.completion(
model="gpt-4",
messages=[{"role": "user", "content": "Hello"}],
extra_headers={
"X-Trace-ID": "trace-abc-123",
"X-Service-Name": "ai-service",
"X-Version": "1.2.3"
}
)
```
### Multi-Tenant SaaS Applications
```python
import litellm
def make_ai_request(user_id, tenant_id, content):
return litellm.completion(
model="claude-3-5-sonnet-latest",
messages=[{"role": "user", "content": content}],
extra_headers={
"X-User-ID": user_id,
"X-Tenant-ID": tenant_id,
"X-Request-Time": str(int(time.time()))
}
)
# Usage
response = make_ai_request("user-456", "tenant-org-1", "Help me write code")
```
### Request Tracing and Debugging
```python
import litellm
import uuid
def traced_completion(model, messages, **kwargs):
trace_id = str(uuid.uuid4())
return litellm.completion(
model=model,
messages=messages,
extra_headers={
"X-Trace-ID": trace_id,
"X-Debug-Mode": "true",
"X-Source-Service": "my-app"
},
**kwargs
)
# Usage
response = traced_completion(
model="gpt-4",
messages=[{"role": "user", "content": "Debug this issue"}]
)
```
### Custom Authentication
```python
import litellm
def get_custom_auth_token():
# Your custom authentication logic
return "custom-auth-token"
response = litellm.completion(
model="claude-3-5-sonnet-latest",
messages=[{"role": "user", "content": "Hello"}],
headers={
"X-Custom-Auth": get_custom_auth_token(),
"X-Auth-Type": "custom"
}
)
```
## Provider Support
Headers are supported across all LiteLLM providers including:
- **OpenAI** (GPT models)
- **Anthropic** (Claude models)
- **Cohere**
- **Hugging Face**
- **Custom providers**
- **Azure OpenAI**
- **AWS Bedrock**
- **Google Vertex AI**
Each provider will receive your custom headers along with their required authentication and API-specific headers.
## Best Practices
### 1. Use Meaningful Header Names
```python
# Good
extra_headers = {
"X-Request-ID": "req-12345",
"X-Tenant-ID": "org-456"
}
# Avoid
extra_headers = {
"custom1": "value1",
"h2": "value2"
}
```
### 2. Include Tracing Information
```python
extra_headers = {
"X-Trace-ID": trace_id,
"X-Span-ID": span_id,
"X-Service-Name": "ai-service"
}
```
### 3. Handle Sensitive Information Carefully
```python
# Don't log sensitive headers
import os
if os.getenv("ENVIRONMENT") != "production":
extra_headers["X-Debug-User"] = user_id
```
### 4. Use Environment-Specific Headers
```python
import os
environment = os.getenv("ENVIRONMENT", "development")
litellm.headers = {
"X-Environment": environment,
"X-Service-Version": os.getenv("SERVICE_VERSION", "unknown")
}
```
## Troubleshooting
### Headers Not Being Passed
If your headers aren't reaching the API:
1. **Check Header Names**: Ensure header names don't conflict with provider-specific headers
2. **Verify Priority**: Remember that `headers` > `extra_headers` > `litellm.headers`
3. **Test with Logging**: Enable verbose logging to see what headers are being sent
```python
import litellm
# Enable debug logging
litellm.set_verbose = True
response = litellm.completion(
model="gpt-4",
messages=[{"role": "user", "content": "test"}],
extra_headers={"X-Debug": "test"}
)
```
### Gateway or Proxy Issues
If using API gateways or proxies:
1. **Check Gateway Requirements**: Verify required headers for your gateway
2. **Test Direct vs Gateway**: Compare direct API calls vs gateway calls
3. **Validate Header Format**: Some gateways have header format requirements
## Security Considerations
1. **Don't Log Sensitive Headers**: Avoid logging authentication tokens or personal data
2. **Use HTTPS**: Always use secure connections when passing sensitive headers
3. **Validate Header Values**: Sanitize user-provided header values
4. **Rotate Keys**: Regularly rotate any API keys passed in headers
```python
import litellm
import re
def safe_header_value(value):
# Remove potentially dangerous characters
return re.sub(r'[^\w\-.]', '', str(value))
response = litellm.completion(
model="gpt-4",
messages=[{"role": "user", "content": "Hello"}],
extra_headers={
"X-User-ID": safe_header_value(user_id)
}
)
```
docs/my-website/docs/troubleshoot.md
+1 -1
View File
@@ -2,7 +2,7 @@
[Schedule Demo 👋](https://calendly.com/d/4mp-gd3-k5k/berriai-1-1-onboarding-litellm-hosted-version)
[Community Discord 💭](https://discord.gg/wuPM9dRgDw)
[Community Slack 💭](https://join.slack.com/share/enQtOTE0ODczMzk2Nzk4NC01YjUxNjY2YjBlYTFmNDRiZTM3NDFiYTM3MzVkODFiMDVjOGRjMmNmZTZkZTMzOWQzZGQyZWIwYjQ0MWExYmE3)
[Community Slack 💭](https://litellmossslack.slack.com/)
Our numbers 📞 +1 (770) 8783-106 / ‭+1 (412) 618-6238‬
docs/my-website/release_notes/v1.76.0-stable/index.md
+1 -1
View File
@@ -1,5 +1,5 @@
---
title: "[PRE-RELEASE]v1.76.0-stable - RPS Improvements"
title: "v1.76.0-stable - RPS Improvements"
slug: "v1-76-0"
date: 2025-08-23T10:00:00
authors:
docs/my-website/release_notes/v1.77.2-stable/index.md
+7 -7
View File
@@ -1,5 +1,5 @@
---
title: "[Pre-Release] v1.77.2-stable - Bedrock Batches API"
title: "v1.77.2-stable - Bedrock Batches API"
slug: "v1-77-2"
date: 2025-09-13T10:00:00
authors:
@@ -21,22 +21,22 @@ import TabItem from '@theme/TabItem';
## Deploy this version
:::info
This release is not yet live.
:::
<Tabs>
<TabItem value="docker" label="Docker">
``` showLineNumbers title="docker run litellm"
docker run \
-e STORE_MODEL_IN_DB=True \
-p 4000:4000 \
ghcr.io/berriai/litellm:main-v1.77.2-stable
```
</TabItem>
<TabItem value="pip" label="Pip">
``` showLineNumbers title="pip install litellm"
pip install litellm==1.77.2.post1
```
</TabItem>
docs/my-website/release_notes/v1.77.3-stable/index.md
+258
View File
@@ -0,0 +1,258 @@
---
title: "[Preview] v1.77.3-stable - Priority Based Rate Limiting"
slug: "v1-77-3"
date: 2025-09-21T10:00:00
authors:
- name: Krrish Dholakia
title: CEO, LiteLLM
url: https://www.linkedin.com/in/krish-d/
image_url: https://pbs.twimg.com/profile_images/1298587542745358340/DZv3Oj-h_400x400.jpg
- name: Ishaan Jaff
title: CTO, LiteLLM
url: https://www.linkedin.com/in/reffajnaahsi/
image_url: https://pbs.twimg.com/profile_images/1613813310264340481/lz54oEiB_400x400.jpg
hide_table_of_contents: false
---
import Image from '@theme/IdealImage';
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
## Deploy this version
<Tabs>
<TabItem value="docker" label="Docker">
``` showLineNumbers title="docker run litellm"
docker run \
-e STORE_MODEL_IN_DB=True \
-p 4000:4000 \
ghcr.io/berriai/litellm:main-v1.77.3.rc.1
```
</TabItem>
<TabItem value="pip" label="Pip">
``` showLineNumbers title="pip install litellm"
pip install litellm==1.77.3
```
</TabItem>
</Tabs>
---
## Key Highlights
- **+550 RPS Performance Improvements** - Optimizations in request handling and object initialization.
- **Priority Quota Reservation** - Proxy admins can now reserve TPM/RPM capacity for specific keys.
## Priority Quota Reservation
This release adds support for priority quota reservation. This allows **Proxy Admins** to reserve TPM/RPM capacity for keys based on metadata priority levels, ensuring critical production workloads get guaranteed access regardless of development traffic volume.
Get started [here](../../docs/proxy/dynamic_rate_limit#priority-quota-reservation)
<iframe width="700" height="500" src="https://www.loom.com/embed/1b54b93139ee415d959402cc0629f3f7" frameborder="0" webkitallowfullscreen mozallowfullscreen allowfullscreen></iframe>
## New Models / Updated Models
#### New Model Support
| Provider | Model | Context Window | Input ($/1M tokens) | Output ($/1M tokens) | Features |
| -------- | ----- | -------------- | ------------------- | -------------------- | -------- |
| SambaNova | `sambanova/deepseek-v3.1` | 128K | $0.90 | $0.90 | Chat completions |
| SambaNova | `sambanova/gpt-oss-120b` | 128K | $0.72 | $0.72 | Chat completions |
| OVHCloud | Various models | Varies | Contact provider | Contact provider | Chat completions |
| CompactifAI | Various models | Varies | Contact provider | Contact provider | Chat completions |
| TwelveLabs | `twelvelabs/marengo-embed-2.7` | 32K | $0.12 | $0.00 | Embeddings |
#### Features
- **[OVHCloud AI Endpoints](../../docs/providers/ovhcloud)**
- New provider support with comprehensive model catalog - [PR #14494](https://github.com/BerriAI/litellm/pull/14494)
- **[CompactifAI](../../docs/providers/compactifai)**
- New provider integration - [PR #14532](https://github.com/BerriAI/litellm/pull/14532)
- **[SambaNova](../../docs/providers/sambanova)**
- Added DeepSeek v3.1 and GPT-OSS-120B models - [PR #14500](https://github.com/BerriAI/litellm/pull/14500)
- **[Bedrock](../../docs/providers/bedrock)**
- Cross-region inference profile cost calculation - [PR #14566](https://github.com/BerriAI/litellm/pull/14566)
- AWS external ID parameter support for authentication - [PR #14582](https://github.com/BerriAI/litellm/pull/14582)
- CountTokens API implementation - [PR #14557](https://github.com/BerriAI/litellm/pull/14557)
- Titan V2 encoding_format parameter support - [PR #14687](https://github.com/BerriAI/litellm/pull/14687)
- Nova Canvas image generation inference profiles - [PR #14578](https://github.com/BerriAI/litellm/pull/14578)
- Bedrock Batches API - batch processing support with file upload and request transformation - [PR #14618](https://github.com/BerriAI/litellm/pull/14618)
- Bedrock Twelve Labs embedding provider support - [PR #14697](https://github.com/BerriAI/litellm/pull/14697)
- **[Vertex AI](../../docs/providers/vertex)**
- Gemini labels field provider-aware filtering - [PR #14563](https://github.com/BerriAI/litellm/pull/14563)
- Gemini Batch API support - [PR #14733](https://github.com/BerriAI/litellm/pull/14733)
- **[Volcengine](../../docs/providers/volcengine)**
- Fixed thinking parameters when disabled - [PR #14569](https://github.com/BerriAI/litellm/pull/14569)
- **[Cohere](../../docs/providers/cohere)**
- Handle Generate API deprecation, default to chat endpoints - [PR #14676](https://github.com/BerriAI/litellm/pull/14676)
- **[TwelveLabs](../../docs/providers/twelvelabs)**
- Added Marengo Embed 2.7 embedding support - [PR #14674](https://github.com/BerriAI/litellm/pull/14674)
### Bug Fixes
- **[Bedrock](../../docs/providers/bedrock)**
- Empty arguments handling in tool call invocation - [PR #14583](https://github.com/BerriAI/litellm/pull/14583)
- **[Vertex AI](../../docs/providers/vertex)**
- Avoid deepcopy crash with non-pickleables in Gemini/Vertex - [PR #14418](https://github.com/BerriAI/litellm/pull/14418)
- **[XAI](../../docs/providers/xai)**
- Fix unsupported stop parameter for grok-code models - [PR #14565](https://github.com/BerriAI/litellm/pull/14565)
- **[Gemini](../../docs/providers/gemini)**
- Updated error message for Gemini API - [PR #14589](https://github.com/BerriAI/litellm/pull/14589)
- Fixed 2.5 Flash Image Preview model routing - [PR #14715](https://github.com/BerriAI/litellm/pull/14715)
- API key passing for token counting endpoints - [PR #14744](https://github.com/BerriAI/litellm/pull/14744)
#### New Provider Support
- **[OVHCloud AI Endpoints](../../docs/providers/ovhcloud)**
- Complete provider integration with model catalog and authentication - [PR #14494](https://github.com/BerriAI/litellm/pull/14494)
- **[CompactifAI](../../docs/providers/compactifai)**
- New provider support with documentation - [PR #14532](https://github.com/BerriAI/litellm/pull/14532)
---
## LLM API Endpoints
#### Features
- **[/responses](../../docs/response_api)**
- Added cancel endpoint support for non-admin users - [PR #14594](https://github.com/BerriAI/litellm/pull/14594)
- Improved response session handling and cold storage configuration with s3 - [PR #14534](https://github.com/BerriAI/litellm/pull/14534)
- Added OpenAI & Azure /responses/cancel endpoint support - [PR #14561](https://github.com/BerriAI/litellm/pull/14561)
- **General**
- Enhanced rate limit error messages with details - [PR #14736](https://github.com/BerriAI/litellm/pull/14736)
- Middle-truncation for spend log payloads - [PR #14637](https://github.com/BerriAI/litellm/pull/14637)
#### Bugs
- **[/chat/completions](../../docs/completion/input)**
- Fixed completion chat ID handling - [PR #14548](https://github.com/BerriAI/litellm/pull/14548)
- Prevent AttributeError for _get_tags_from_request_kwargs - [PR #14735](https://github.com/BerriAI/litellm/pull/14735)
- **[/responses](../../docs/response_api)**
- Fixed cost calculation - [PR #14675](https://github.com/BerriAI/litellm/pull/14675)
- **General**
- Rate limiter AttributeError fix - [PR #14609](https://github.com/BerriAI/litellm/pull/14609)
---
## Spend Tracking, Budgets and Rate Limiting
- **Responses API Cost Calculation** fix - [PR #14675](https://github.com/BerriAI/litellm/pull/14675)
- **Anthropic Cache Token Pricing** - Separate 1-hour vs 5-minute cache creation costs - [PR #14620](https://github.com/BerriAI/litellm/pull/14620), [PR #14652](https://github.com/BerriAI/litellm/pull/14652)
- **Indochina Time Timezone** support for budget resets - [PR #14666](https://github.com/BerriAI/litellm/pull/14666)
- **Soft Budget Alert Cache Issues** - Resolved soft budget alert cache issues - [PR #14491](https://github.com/BerriAI/litellm/pull/14491)
- **Dynamic Rate Limiter v3** - Priority routing improvements - [PR #14734](https://github.com/BerriAI/litellm/pull/14734)
- **Enhanced Rate Limit Errors** - More detailed error messages - [PR #14736](https://github.com/BerriAI/litellm/pull/14736)
---
## Management Endpoints / UI
#### Features
- **Team Member Service Account Keys** - Allow team members to view keys they create - [PR #14619](https://github.com/BerriAI/litellm/pull/14619)
- **Default Budget for JWT Teams** - Auto-assign budgets to generated teams - [PR #14514](https://github.com/BerriAI/litellm/pull/14514)
- **SSO Access Control Groups** - Enhanced token info endpoint integration - [PR #14738](https://github.com/BerriAI/litellm/pull/14738)
- **Health Test Connect Protection** - Restrict access based on model creation permissions - [PR #14650](https://github.com/BerriAI/litellm/pull/14650)
- **Amazon Bedrock Guardrail Info View** - Enhanced logging visualization - [PR #14696](https://github.com/BerriAI/litellm/pull/14696)
#### Bug Fixes
- **SCIM v2** - Fix group PUSH and PUT operations for non-existent members - [PR #14581](https://github.com/BerriAI/litellm/pull/14581)
- **Guardrail View/Edit/Delete** behavior fixes - [PR #14622](https://github.com/BerriAI/litellm/pull/14622)
- **In-Memory Guardrail** update failures - [PR #14653](https://github.com/BerriAI/litellm/pull/14653)
---
## Logging / Guardrail Integrations
#### Features
- **[DataDog](../../docs/proxy/logging#datadog)**
- Enhanced spend tracking metrics - [PR #14555](https://github.com/BerriAI/litellm/pull/14555)
- Stream support with is_streamed_request parameter - [PR #14673](https://github.com/BerriAI/litellm/pull/14673)
- Fixed tool calls metadata passing - [PR #14531](https://github.com/BerriAI/litellm/pull/14531)
- **[Langfuse](../../docs/proxy/logging#langfuse)**
- Added logging support for Responses API - [PR #14597](https://github.com/BerriAI/litellm/pull/14597)
- **[Langsmith](../../docs/proxy/logging#langsmith)**
- Langsmith Sampling Rate - Key/Team-level tracing configuration - [PR #14740](https://github.com/BerriAI/litellm/pull/14740)
- **[Prometheus](../../docs/proxy/logging#prometheus)**
- Multi-worker support improvements - [PR #14530](https://github.com/BerriAI/litellm/pull/14530)
- User email labels in monitoring - [PR #14520](https://github.com/BerriAI/litellm/pull/14520)
- **[Opik](../../docs/proxy/logging#opik)**
- Fixed timezone issue - [PR #14708](https://github.com/BerriAI/litellm/pull/14708)
### Bug Fixes
- **[S3](../../docs/proxy/logging#s3-buckets)**
- Fixed 404 error when using s3_endpoint_url - [PR #14559](https://github.com/BerriAI/litellm/pull/14559)
#### Guardrails
- **Tool Permission Guardrail** - Fine-grained tool access control - [PR #14519](https://github.com/BerriAI/litellm/pull/14519)
- **Bedrock Guardrails** - Selective guarding support with runtime endpoint configuration - [PR #14575](https://github.com/BerriAI/litellm/pull/14575), [PR #14650](https://github.com/BerriAI/litellm/pull/14650)
- **Default Last Message** in guardrails - [PR #14640](https://github.com/BerriAI/litellm/pull/14640)
- **AWS exceptions handling despite 200 response** - [PR #14658](https://github.com/BerriAI/litellm/pull/14658)
#### New Integration
- **[PostHog](../../docs/observability/posthog)** - Complete observability integration for LiteLLM usage tracking and analytics - [PR #14610](https://github.com/BerriAI/litellm/pull/14610)
---
## MCP Gateway
- **MCP Server Alias Parsing** - Multi-part URL path support - [PR #14558](https://github.com/BerriAI/litellm/pull/14558)
- **MCP Filter Recomputation** - After server deletion - [PR #14542](https://github.com/BerriAI/litellm/pull/14542)
- **MCP Gateway Tools List** improvements - [PR #14695](https://github.com/BerriAI/litellm/pull/14695)
---
## Performance / Loadbalancing / Reliability improvements
- **+500 RPS Performance Boost** when sending the `user` field - [PR #14616](https://github.com/BerriAI/litellm/pull/14616)
- **+50 RPS** by removing iscoroutine from hot path - [PR #14649](https://github.com/BerriAI/litellm/pull/14649)
- **7% reduction** in __init__ overhead - [PR #14689](https://github.com/BerriAI/litellm/pull/14689)
- **Generic Object Pool** implementation for better resource management - [PR #14702](https://github.com/BerriAI/litellm/pull/14702)
---
## General Proxy Improvements
- **Middle-Truncation** for spend log payloads - [PR #14637](https://github.com/BerriAI/litellm/pull/14637)
#### Security
- **Security Update** - Bump aiohttp==3.12.14, fix CVE-2025-53643 - [PR #14638](https://github.com/BerriAI/litellm/pull/14638)
---
## New Contributors
* @luisfucros made their first contribution in [PR #14500](https://github.com/BerriAI/litellm/pull/14500)
* @hanakannzashi made their first contribution in [PR #14548](https://github.com/BerriAI/litellm/pull/14548)
* @eliasto made their first contribution in [PR #14494](https://github.com/BerriAI/litellm/pull/14494)
* @Rasmusafj made their first contribution in [PR #14491](https://github.com/BerriAI/litellm/pull/14491)
* @LingXuanYin made their first contribution in [PR #14569](https://github.com/BerriAI/litellm/pull/14569)
* @ronaldpereira made their first contribution in [PR #14613](https://github.com/BerriAI/litellm/pull/14613)
* @hula-la made their first contribution in [PR #14534](https://github.com/BerriAI/litellm/pull/14534)
* @carlos-marchal-ph made their first contribution in [PR #14610](https://github.com/BerriAI/litellm/pull/14610)
* @akraines made their first contribution in [PR #14637](https://github.com/BerriAI/litellm/pull/14637)
* @mrFranklin made their first contribution in [PR #14708](https://github.com/BerriAI/litellm/pull/14708)
* @tcx4c70 made their first contribution in [PR #14675](https://github.com/BerriAI/litellm/pull/14675)
* @michaeltansg made their first contribution in [PR #14666](https://github.com/BerriAI/litellm/pull/14666)
* @tosi29 made their first contribution in [PR #14725](https://github.com/BerriAI/litellm/pull/14725)
* @gmdfalk made their first contribution in [PR #14735](https://github.com/BerriAI/litellm/pull/14735)
* @FelipeRodriguesGare made their first contribution in [PR #14733](https://github.com/BerriAI/litellm/pull/14733)
* @mritunjaysharma394 made their first contribution in [PR #14678](https://github.com/BerriAI/litellm/pull/14678)
---
## **[Full Changelog](https://github.com/BerriAI/litellm/compare/v1.77.2.rc.1...v1.77.3.rc.1)**
docs/my-website/sidebars.js
+4 -1
View File
@@ -201,7 +201,7 @@ const sidebars = {
{
type: "category",
label: "Budgets + Rate Limits",
items: ["proxy/users", "proxy/temporary_budget_increase", "proxy/rate_limit_tiers", "proxy/team_budgets", "proxy/customers"],
items: ["proxy/users", "proxy/temporary_budget_increase", "proxy/rate_limit_tiers", "proxy/team_budgets", "proxy/dynamic_rate_limit", "proxy/customers"],
},
{
type: "link",
@@ -392,6 +392,7 @@ const sidebars = {
"providers/vertex",
"providers/vertex_partner",
"providers/vertex_image",
"providers/vertex_batch",
]
},
{
@@ -411,6 +412,7 @@ const sidebars = {
label: "Bedrock",
items: [
"providers/bedrock",
"providers/bedrock_embedding",
"providers/bedrock_agents",
"providers/bedrock_batches",
"providers/bedrock_vector_store",
@@ -522,6 +524,7 @@ const sidebars = {
"completion/batching",
"completion/mock_requests",
"completion/reliable_completions",
"proxy/veo_video_generation",
]
},
enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py
+6
View File
@@ -109,6 +109,9 @@ class PagerDutyAlerting(SlackAlerting):
error_llm_provider=error_info.get("llm_provider"),
user_api_key_hash=_meta.get("user_api_key_hash"),
user_api_key_alias=_meta.get("user_api_key_alias"),
user_api_key_spend=_meta.get("user_api_key_spend"),
user_api_key_max_budget=_meta.get("user_api_key_max_budget"),
user_api_key_budget_reset_at=_meta.get("user_api_key_budget_reset_at"),
user_api_key_org_id=_meta.get("user_api_key_org_id"),
user_api_key_team_id=_meta.get("user_api_key_team_id"),
user_api_key_user_id=_meta.get("user_api_key_user_id"),
@@ -191,6 +194,9 @@ class PagerDutyAlerting(SlackAlerting):
error_llm_provider="HangingRequest",
user_api_key_hash=user_api_key_dict.api_key,
user_api_key_alias=user_api_key_dict.key_alias,
user_api_key_spend=user_api_key_dict.spend,
user_api_key_max_budget=user_api_key_dict.max_budget,
user_api_key_budget_reset_at=user_api_key_dict.budget_reset_at.isoformat() if user_api_key_dict.budget_reset_at else None,
user_api_key_org_id=user_api_key_dict.org_id,
user_api_key_team_id=user_api_key_dict.team_id,
user_api_key_user_id=user_api_key_dict.user_id,
enterprise/litellm_enterprise/integrations/prometheus.py
+20 -21
View File
@@ -102,7 +102,9 @@ class PrometheusLogger(CustomLogger):
# "team",
# "team_alias",
# ],
labelnames=self.get_labels_for_metric("litellm_llm_api_time_to_first_token_metric"),
labelnames=self.get_labels_for_metric(
"litellm_llm_api_time_to_first_token_metric"
),
buckets=LATENCY_BUCKETS,
)
@@ -240,14 +242,14 @@ class PrometheusLogger(CustomLogger):
self.litellm_deployment_state = self._gauge_factory(
"litellm_deployment_state",
"LLM Deployment Analytics - The state of the deployment: 0 = healthy, 1 = partial outage, 2 = complete outage",
labelnames=self.get_labels_for_metric("litellm_deployment_state")
labelnames=self.get_labels_for_metric("litellm_deployment_state"),
)
self.litellm_deployment_cooled_down = self._counter_factory(
"litellm_deployment_cooled_down",
"LLM Deployment Analytics - Number of times a deployment has been cooled down by LiteLLM load balancing logic. exception_status is the status of the exception that caused the deployment to be cooled down",
# labelnames=_logged_llm_labels + [EXCEPTION_STATUS],
labelnames=self.get_labels_for_metric("litellm_deployment_cooled_down")
labelnames=self.get_labels_for_metric("litellm_deployment_cooled_down"),
)
self.litellm_deployment_success_responses = self._counter_factory(
@@ -1039,20 +1041,12 @@ class PrometheusLogger(CustomLogger):
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name="litellm_proxy_total_requests_metric"
metric_name="litellm_spend_metric"
),
enum_values=enum_values,
)
self.litellm_spend_metric.labels(
end_user_id,
user_api_key,
user_api_key_alias,
model,
user_api_team,
user_api_team_alias,
user_id,
).inc(response_cost)
self.litellm_spend_metric.labels(**_labels).inc(response_cost)
def _set_virtual_key_rate_limit_metrics(
self,
@@ -2280,7 +2274,9 @@ def get_custom_labels_from_metadata(metadata: dict) -> Dict[str, str]:
return result
def _tag_matches_wildcard_configured_pattern(tags: List[str], configured_tag: str) -> bool:
def _tag_matches_wildcard_configured_pattern(
tags: List[str], configured_tag: str
) -> bool:
"""
Check if any of the request tags matches a wildcard configured pattern
@@ -2305,6 +2301,7 @@ def _tag_matches_wildcard_configured_pattern(tags: List[str], configured_tag: st
import re
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
pattern_router = PatternMatchRouter()
regex_pattern = pattern_router._pattern_to_regex(configured_tag)
return any(re.match(pattern=regex_pattern, string=tag) for tag in tags)
@@ -2313,11 +2310,11 @@ def _tag_matches_wildcard_configured_pattern(tags: List[str], configured_tag: st
def get_custom_labels_from_tags(tags: List[str]) -> Dict[str, str]:
"""
Get custom labels from tags based on admin configuration.
Supports both exact matches and wildcard patterns:
- Exact match: "prod" matches "prod" exactly
- Wildcard pattern: "User-Agent: curl/*" matches "User-Agent: curl/7.68.0"
- Wildcard pattern: "User-Agent: curl/*" matches "User-Agent: curl/7.68.0"
Reuses PatternMatchRouter for wildcard pattern matching.
Returns dict of label_name: "true" if the tag matches the configured tag, "false" otherwise
@@ -2345,17 +2342,19 @@ def get_custom_labels_from_tags(tags: List[str]) -> Dict[str, str]:
for configured_tag in configured_tags:
label_name = _sanitize_prometheus_label_name(f"tag_{configured_tag}")
# Check for exact match first (backwards compatibility)
if configured_tag in tags:
result[label_name] = "true"
continue
# Use PatternMatchRouter for wildcard pattern matching
if "*" in configured_tag and _tag_matches_wildcard_configured_pattern(tags=tags, configured_tag=configured_tag):
if "*" in configured_tag and _tag_matches_wildcard_configured_pattern(
tags=tags, configured_tag=configured_tag
):
result[label_name] = "true"
continue
# No match found
result[label_name] = "false"
enterprise/pyproject.toml
+2 -2
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm-enterprise"
version = "0.1.19"
version = "0.1.20"
description = "Package for LiteLLM Enterprise features"
authors = ["BerriAI"]
readme = "README.md"
@@ -22,7 +22,7 @@ requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "0.1.19"
version = "0.1.20"
version_files = [
"pyproject.toml:version",
"../requirements.txt:litellm-enterprise==",
examples/sdk_headers_example.py
+195
View File
@@ -0,0 +1,195 @@
#!/usr/bin/env python3
"""
Example demonstrating LiteLLM SDK header support for enterprise environments.
This example shows how to use additional headers with API gateways, service meshes,
and multi-tenant architectures.
"""
import litellm
import os
from typing import Dict, Any
def example_global_headers():
"""Example: Set global headers for all requests"""
print("=== Global Headers Example ===")
# Set global headers that will be included in all API requests
litellm.headers = {
"X-API-Gateway-Key": "your-gateway-key-here",
"X-Company-ID": "acme-corp",
"X-Environment": "production"
}
print("Global headers set:", litellm.headers)
# These headers will now be included in all completion calls
# (Note: This example doesn't actually make API calls)
print("Global headers will be included in all subsequent completion() calls")
def example_per_request_headers():
"""Example: Using extra_headers for specific requests"""
print("\n=== Per-Request Headers Example ===")
headers_to_send = {
"X-Request-ID": "req-12345",
"X-Tenant-ID": "tenant-abc",
"X-Custom-Auth": "bearer-token-xyz"
}
print("Per-request headers:", headers_to_send)
# Example of how you would use extra_headers in a real call
# response = litellm.completion(
# model="claude-3-5-sonnet-latest",
# messages=[{"role": "user", "content": "Hello"}],
# extra_headers=headers_to_send
# )
def example_header_priority():
"""Example: Demonstrating header priority and merging"""
print("\n=== Header Priority Example ===")
# Set global headers
litellm.headers = {
"X-Company-ID": "acme-corp",
"X-Shared-Header": "global-value"
}
# Headers that would be sent in a request
extra_headers = {
"X-Request-ID": "req-12345",
"X-Shared-Header": "extra-value" # Overrides global
}
request_headers = {
"X-Priority-Header": "important",
"X-Shared-Header": "request-value" # Overrides both global and extra
}
print("Global headers:", litellm.headers)
print("Extra headers:", extra_headers)
print("Request headers:", request_headers)
print("\nFinal headers would be:")
print(" X-Company-ID: acme-corp (from global)")
print(" X-Request-ID: req-12345 (from extra)")
print(" X-Priority-Header: important (from request)")
print(" X-Shared-Header: request-value (request wins - highest priority)")
def example_enterprise_api_gateway():
"""Example: Enterprise API Gateway scenario"""
print("\n=== Enterprise API Gateway Example ===")
# Simulate enterprise environment with Apigee or similar
gateway_config = {
"X-API-Gateway-Key": os.getenv("API_GATEWAY_KEY", "demo-key"),
"X-Route-Version": "v2",
"X-Rate-Limit-Group": "premium"
}
# Set gateway headers globally
litellm.headers = gateway_config
print("Gateway headers configured:", gateway_config)
# Function to make tenant-specific requests
def make_tenant_request(tenant_id: str, user_id: str, content: str) -> Dict[str, Any]:
"""Make an AI request with tenant-specific headers"""
tenant_headers = {
"X-Tenant-ID": tenant_id,
"X-User-ID": user_id,
"X-Request-Time": "2024-01-01T00:00:00Z",
"X-Service-Name": "ai-assistant"
}
print(f"Making request for tenant {tenant_id}, user {user_id}")
print("Tenant-specific headers:", tenant_headers)
# In a real scenario, this would make the actual API call:
# return litellm.completion(
# model="claude-3-5-sonnet-latest",
# messages=[{"role": "user", "content": content}],
# extra_headers=tenant_headers
# )
# For demo purposes, return mock data
return {"mock": "response", "headers_used": {**gateway_config, **tenant_headers}}
# Example usage
result = make_tenant_request("tenant-123", "user-456", "Analyze this data")
print("Response:", result)
def example_service_mesh():
"""Example: Service mesh integration (Istio, Linkerd)"""
print("\n=== Service Mesh Example ===")
service_mesh_headers = {
"X-Trace-ID": "trace-abc-123",
"X-Span-ID": "span-def-456",
"X-Service-Name": "ai-service",
"X-Version": "1.2.3",
"X-Cluster": "prod-us-west-2"
}
print("Service mesh headers:", service_mesh_headers)
# Example of using these headers for distributed tracing
# response = litellm.completion(
# model="gpt-4",
# messages=[{"role": "user", "content": "Hello"}],
# extra_headers=service_mesh_headers
# )
def example_debugging_and_monitoring():
"""Example: Request debugging and monitoring"""
print("\n=== Debugging and Monitoring Example ===")
import uuid
import time
# Generate unique identifiers for request tracking
trace_id = str(uuid.uuid4())
request_id = f"req-{int(time.time())}"
debug_headers = {
"X-Trace-ID": trace_id,
"X-Request-ID": request_id,
"X-Debug-Mode": "true",
"X-Source-Service": "customer-support-bot",
"X-Request-Priority": "high"
}
print("Debug headers:", debug_headers)
print(f"Trace ID: {trace_id}")
print(f"Request ID: {request_id}")
# These headers help with:
# 1. Distributed tracing across services
# 2. Request correlation in logs
# 3. Debug mode enablement
# 4. Priority-based routing
if __name__ == "__main__":
print("LiteLLM SDK Header Support Examples")
print("=" * 50)
example_global_headers()
example_per_request_headers()
example_header_priority()
example_enterprise_api_gateway()
example_service_mesh()
example_debugging_and_monitoring()
print("\n" + "=" * 50)
print("All examples completed!")
print("\nTo use in your application:")
print("1. Set litellm.headers for global headers")
print("2. Use extra_headers parameter for request-specific headers")
print("3. Use headers parameter for highest priority headers")
print("4. Headers are merged with priority: headers > extra_headers > litellm.headers")
litellm-proxy-extras/dist/litellm_proxy_extras-0.2.19-py3-none-any.whl
Vendored
BIN
View File
Binary file not shown.
litellm-proxy-extras/dist/litellm_proxy_extras-0.2.19.tar.gz
Vendored
BIN
View File
Binary file not shown.
litellm-proxy-extras/litellm_proxy_extras/migrations/20250918083359_drop_spec_version_column_from_mcp_table/migration.sql
+8
View File
@@ -0,0 +1,8 @@
/*
Warnings:
- You are about to drop the column `spec_version` on the `LiteLLM_MCPServerTable` table. All the data in the column will be lost.
*/
-- AlterTable
ALTER TABLE "public"."LiteLLM_MCPServerTable" DROP COLUMN "spec_version";
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
-1
View File
@@ -171,7 +171,6 @@ model LiteLLM_MCPServerTable {
description String?
url String?
transport String @default("sse")
spec_version String @default("2025-03-26")
auth_type String?
created_at DateTime? @default(now()) @map("created_at")
created_by String?
litellm-proxy-extras/pyproject.toml
+2 -2
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm-proxy-extras"
version = "0.2.18"
version = "0.2.19"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
authors = ["BerriAI"]
readme = "README.md"
@@ -22,7 +22,7 @@ requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "0.2.18"
version = "0.2.19"
version_files = [
"pyproject.toml:version",
"../requirements.txt:litellm-proxy-extras==",
litellm/__init__.py
+3 -1
View File
@@ -68,6 +68,7 @@ from litellm.constants import (
bedrock_embedding_models,
known_tokenizer_config,
BEDROCK_INVOKE_PROVIDERS_LITERAL,
BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
BEDROCK_CONVERSE_MODELS,
DEFAULT_MAX_TOKENS,
DEFAULT_SOFT_BUDGET,
@@ -117,6 +118,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"logfire",
"literalai",
"dynamic_rate_limiter",
"dynamic_rate_limiter_v3",
"langsmith",
"prometheus",
"otel",
@@ -148,6 +150,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"vector_store_pre_call_hook",
"dotprompt",
"cloudzero",
"posthog",
]
configured_cold_storage_logger: Optional[
_custom_logger_compatible_callbacks_literal
@@ -1051,7 +1054,6 @@ from .llms.databricks.chat.transformation import DatabricksConfig
from .llms.databricks.embed.transformation import DatabricksEmbeddingConfig
from .llms.predibase.chat.transformation import PredibaseConfig
from .llms.replicate.chat.transformation import ReplicateConfig
from .llms.cohere.completion.transformation import CohereTextConfig as CohereConfig
from .llms.snowflake.chat.transformation import SnowflakeConfig
from .llms.cohere.rerank.transformation import CohereRerankConfig
from .llms.cohere.rerank_v2.transformation import CohereRerankV2Config
litellm/batches/main.py
+169 -107
View File
@@ -340,7 +340,7 @@ def create_batch(
@client
async def aretrieve_batch(
batch_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@@ -378,11 +378,129 @@ async def aretrieve_batch(
except Exception as e:
raise e
def _handle_retrieve_batch_providers_without_provider_config(
batch_id: str,
optional_params: GenericLiteLLMParams,
timeout: Union[float, httpx.Timeout],
litellm_params: dict,
_retrieve_batch_request: RetrieveBatchRequest,
_is_async: bool,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai",
):
api_base: Optional[str] = None
if custom_llm_provider == "openai":
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
api_base = (
optional_params.api_base
or litellm.api_base
or os.getenv("OPENAI_BASE_URL")
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
)
# set API KEY
api_key = (
optional_params.api_key
or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
or litellm.openai_key
or os.getenv("OPENAI_API_KEY")
)
response = openai_batches_instance.retrieve_batch(
_is_async=_is_async,
retrieve_batch_data=_retrieve_batch_request,
api_base=api_base,
api_key=api_key,
organization=organization,
timeout=timeout,
max_retries=optional_params.max_retries,
)
elif custom_llm_provider == "azure":
api_base = (
optional_params.api_base
or litellm.api_base
or get_secret_str("AZURE_API_BASE")
)
api_version = (
optional_params.api_version
or litellm.api_version
or get_secret_str("AZURE_API_VERSION")
)
api_key = (
optional_params.api_key
or litellm.api_key
or litellm.azure_key
or get_secret_str("AZURE_OPENAI_API_KEY")
or get_secret_str("AZURE_API_KEY")
)
extra_body = optional_params.get("extra_body", {})
if extra_body is not None:
extra_body.pop("azure_ad_token", None)
else:
get_secret_str("AZURE_AD_TOKEN") # type: ignore
response = azure_batches_instance.retrieve_batch(
_is_async=_is_async,
api_base=api_base,
api_key=api_key,
api_version=api_version,
timeout=timeout,
max_retries=optional_params.max_retries,
retrieve_batch_data=_retrieve_batch_request,
litellm_params=litellm_params,
)
elif custom_llm_provider == "vertex_ai":
api_base = optional_params.api_base or ""
vertex_ai_project = (
optional_params.vertex_project
or litellm.vertex_project
or get_secret_str("VERTEXAI_PROJECT")
)
vertex_ai_location = (
optional_params.vertex_location
or litellm.vertex_location
or get_secret_str("VERTEXAI_LOCATION")
)
vertex_credentials = optional_params.vertex_credentials or get_secret_str(
"VERTEXAI_CREDENTIALS"
)
response = vertex_ai_batches_instance.retrieve_batch(
_is_async=_is_async,
batch_id=batch_id,
api_base=api_base,
vertex_project=vertex_ai_project,
vertex_location=vertex_ai_location,
vertex_credentials=vertex_credentials,
timeout=timeout,
max_retries=optional_params.max_retries,
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format(
custom_llm_provider
),
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
),
)
return response
@client
def retrieve_batch(
batch_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@@ -430,115 +548,59 @@ def retrieve_batch(
)
_is_async = kwargs.pop("aretrieve_batch", False) is True
api_base: Optional[str] = None
if custom_llm_provider == "openai":
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
api_base = (
optional_params.api_base
or litellm.api_base
or os.getenv("OPENAI_BASE_URL")
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
)
# set API KEY
api_key = (
optional_params.api_key
or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
or litellm.openai_key
or os.getenv("OPENAI_API_KEY")
)
response = openai_batches_instance.retrieve_batch(
_is_async=_is_async,
retrieve_batch_data=_retrieve_batch_request,
api_base=api_base,
api_key=api_key,
organization=organization,
timeout=timeout,
max_retries=optional_params.max_retries,
)
elif custom_llm_provider == "azure":
api_base = (
optional_params.api_base
or litellm.api_base
or get_secret_str("AZURE_API_BASE")
)
api_version = (
optional_params.api_version
or litellm.api_version
or get_secret_str("AZURE_API_VERSION")
)
api_key = (
optional_params.api_key
or litellm.api_key
or litellm.azure_key
or get_secret_str("AZURE_OPENAI_API_KEY")
or get_secret_str("AZURE_API_KEY")
)
extra_body = optional_params.get("extra_body", {})
if extra_body is not None:
extra_body.pop("azure_ad_token", None)
else:
get_secret_str("AZURE_AD_TOKEN") # type: ignore
response = azure_batches_instance.retrieve_batch(
_is_async=_is_async,
api_base=api_base,
api_key=api_key,
api_version=api_version,
timeout=timeout,
max_retries=optional_params.max_retries,
retrieve_batch_data=_retrieve_batch_request,
litellm_params=litellm_params,
)
elif custom_llm_provider == "vertex_ai":
api_base = optional_params.api_base or ""
vertex_ai_project = (
optional_params.vertex_project
or litellm.vertex_project
or get_secret_str("VERTEXAI_PROJECT")
)
vertex_ai_location = (
optional_params.vertex_location
or litellm.vertex_location
or get_secret_str("VERTEXAI_LOCATION")
)
vertex_credentials = optional_params.vertex_credentials or get_secret_str(
"VERTEXAI_CREDENTIALS"
)
response = vertex_ai_batches_instance.retrieve_batch(
_is_async=_is_async,
batch_id=batch_id,
api_base=api_base,
vertex_project=vertex_ai_project,
vertex_location=vertex_ai_location,
vertex_credentials=vertex_credentials,
timeout=timeout,
max_retries=optional_params.max_retries,
client = kwargs.get("client", None)
# Try to use provider config first (for providers like bedrock)
model: Optional[str] = kwargs.get("model", None)
if model is not None:
provider_config = ProviderConfigManager.get_provider_batches_config(
model=model,
provider=LlmProviders(custom_llm_provider),
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format(
custom_llm_provider
),
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
status_code=400,
content="Unsupported provider",
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
provider_config = None
if provider_config is not None:
response = base_llm_http_handler.retrieve_batch(
batch_id=batch_id,
provider_config=provider_config,
litellm_params=litellm_params,
headers=extra_headers or {},
api_base=optional_params.api_base,
api_key=optional_params.api_key,
logging_obj=litellm_logging_obj or LiteLLMLoggingObj(
model=model or "bedrock/unknown",
messages=[],
stream=False,
call_type="batch_retrieve",
start_time=None,
litellm_call_id="batch_retrieve_" + batch_id,
function_id="batch_retrieve",
),
_is_async=_is_async,
client=client
if client is not None
and isinstance(client, (HTTPHandler, AsyncHTTPHandler))
else None,
timeout=timeout,
model=model,
)
return response
return response
#########################################################
# Handle providers without provider config
#########################################################
return _handle_retrieve_batch_providers_without_provider_config(
batch_id=batch_id,
custom_llm_provider=custom_llm_provider,
optional_params=optional_params,
litellm_params=litellm_params,
_retrieve_batch_request=_retrieve_batch_request,
_is_async=_is_async,
timeout=timeout,
)
except Exception as e:
raise e
litellm/caching/redis_cache.py
+2 -1
View File
@@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union, cast
import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.services import ServiceTypes
@@ -138,7 +139,7 @@ class RedisCache(BaseCache):
self.redis_flush_size = redis_flush_size
self.redis_version = "Unknown"
try:
if not inspect.iscoroutinefunction(self.redis_client):
if not coroutine_checker.is_async_callable(self.redis_client):
self.redis_version = self.redis_client.info()["redis_version"] # type: ignore
except Exception:
pass
litellm/constants.py
+14 -1
View File
@@ -179,7 +179,7 @@ NON_LLM_CONNECTION_TIMEOUT = int(
os.getenv("NON_LLM_CONNECTION_TIMEOUT", 15)
) # timeout for adjacent services (e.g. jwt auth)
MAX_EXCEPTION_MESSAGE_LENGTH = int(os.getenv("MAX_EXCEPTION_MESSAGE_LENGTH", 2000))
MAX_STRING_LENGTH_PROMPT_IN_DB = int(os.getenv("MAX_STRING_LENGTH_PROMPT_IN_DB", 1000))
MAX_STRING_LENGTH_PROMPT_IN_DB = int(os.getenv("MAX_STRING_LENGTH_PROMPT_IN_DB", 2048))
BEDROCK_MAX_POLICY_SIZE = int(os.getenv("BEDROCK_MAX_POLICY_SIZE", 75))
REPLICATE_POLLING_DELAY_SECONDS = float(
os.getenv("REPLICATE_POLLING_DELAY_SECONDS", 0.5)
@@ -805,6 +805,12 @@ BEDROCK_INVOKE_PROVIDERS_LITERAL = Literal[
"deepseek_r1",
]
BEDROCK_EMBEDDING_PROVIDERS_LITERAL = Literal[
"cohere",
"amazon",
"twelvelabs",
]
BEDROCK_CONVERSE_MODELS = [
"openai.gpt-oss-20b-1:0",
"openai.gpt-oss-120b-1:0",
@@ -858,6 +864,7 @@ bedrock_embedding_models: set = set(
"amazon.titan-embed-text-v1",
"cohere.embed-english-v3",
"cohere.embed-multilingual-v3",
"twelvelabs.marengo-embed-2-7-v1:0",
]
)
@@ -976,6 +983,7 @@ HEALTH_CHECK_TIMEOUT_SECONDS = int(
os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60)
) # 60 seconds
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME = "litellm-internal-health-check"
LITTELM_CLI_SERVICE_ACCOUNT_NAME = "litellm-cli"
UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard"
LITELLM_PROXY_ADMIN_NAME = "default_user_id"
@@ -1099,3 +1107,8 @@ SENTRY_PII_DENYLIST = [
"SMTP_SENDER_EMAIL",
"TEST_EMAIL_ADDRESS",
]
# CoroutineChecker cache configuration
COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY = int(
os.getenv("COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY", 1000)
)
litellm/experimental_mcp_client/client.py
+28 -31
View File
@@ -19,8 +19,6 @@ from litellm._logging import verbose_logger
from litellm.types.mcp import (
MCPAuth,
MCPAuthType,
MCPSpecVersion,
MCPSpecVersionType,
MCPStdioConfig,
MCPTransport,
MCPTransportType,
@@ -48,7 +46,6 @@ class MCPClient:
auth_value: Optional[str] = None,
timeout: float = 60.0,
stdio_config: Optional[MCPStdioConfig] = None,
protocol_version: MCPSpecVersionType = MCPSpecVersion.jun_2025,
):
self.server_url: str = server_url
self.transport_type: MCPTransport = transport_type
@@ -62,7 +59,6 @@ class MCPClient:
self._session_ctx = None
self._task: Optional[asyncio.Task] = None
self.stdio_config: Optional[MCPStdioConfig] = stdio_config
self.protocol_version: MCPSpecVersionType = protocol_version
# handle the basic auth value if provided
if auth_value:
@@ -84,22 +80,24 @@ class MCPClient:
"""Initialize the transport and session."""
if self._session:
return # Already connected
try:
if self.transport_type == MCPTransport.stdio:
# For stdio transport, use stdio_client with command-line parameters
if not self.stdio_config:
raise ValueError("stdio_config is required for stdio transport")
server_params = StdioServerParameters(
command=self.stdio_config.get("command", ""),
args=self.stdio_config.get("args", []),
env=self.stdio_config.get("env", {})
env=self.stdio_config.get("env", {}),
)
self._transport_ctx = stdio_client(server_params)
self._transport = await self._transport_ctx.__aenter__()
self._session_ctx = ClientSession(self._transport[0], self._transport[1])
self._session_ctx = ClientSession(
self._transport[0], self._transport[1]
)
self._session = await self._session_ctx.__aenter__()
await self._session.initialize()
elif self.transport_type == MCPTransport.sse:
@@ -110,7 +108,9 @@ class MCPClient:
headers=headers,
)
self._transport = await self._transport_ctx.__aenter__()
self._session_ctx = ClientSession(self._transport[0], self._transport[1])
self._session_ctx = ClientSession(
self._transport[0], self._transport[1]
)
self._session = await self._session_ctx.__aenter__()
await self._session.initialize()
else: # http
@@ -121,7 +121,9 @@ class MCPClient:
headers=headers,
)
self._transport = await self._transport_ctx.__aenter__()
self._session_ctx = ClientSession(self._transport[0], self._transport[1])
self._session_ctx = ClientSession(
self._transport[0], self._transport[1]
)
self._session = await self._session_ctx.__aenter__()
await self._session.initialize()
except ValueError as e:
@@ -184,8 +186,10 @@ class MCPClient:
def _get_auth_headers(self) -> dict:
"""Generate authentication headers based on auth type."""
headers = {}
headers = {
"MCP-Protocol-Version": "2025-06-18"
}
if self._mcp_auth_value:
if self.auth_type == MCPAuth.bearer_token:
headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
@@ -196,18 +200,8 @@ class MCPClient:
elif self.auth_type == MCPAuth.authorization:
headers["Authorization"] = self._mcp_auth_value
# Handle protocol version - it might be a string or enum
if hasattr(self.protocol_version, 'value'):
# It's an enum
protocol_version_str = self.protocol_version.value
else:
# It's a string
protocol_version_str = str(self.protocol_version)
headers["MCP-Protocol-Version"] = protocol_version_str
return headers
async def list_tools(self) -> List[MCPTool]:
"""List available tools from the server."""
if not self._session:
@@ -216,7 +210,7 @@ class MCPClient:
except Exception as e:
verbose_logger.warning(f"MCP client connection failed: {str(e)}")
return []
if self._session is None:
verbose_logger.warning("MCP client session is not initialized")
return []
@@ -245,17 +239,20 @@ class MCPClient:
except Exception as e:
verbose_logger.warning(f"MCP client connection failed: {str(e)}")
return MCPCallToolResult(
content=[TextContent(type="text", text=f"{str(e)}")],
isError=True
content=[TextContent(type="text", text=f"{str(e)}")], isError=True
)
if self._session is None:
verbose_logger.warning("MCP client session is not initialized")
return MCPCallToolResult(
content=[TextContent(type="text", text="MCP client session is not initialized")],
content=[
TextContent(
type="text", text="MCP client session is not initialized"
)
],
isError=True,
)
try:
tool_result = await self._session.call_tool(
name=call_tool_request_params.name,
@@ -270,8 +267,8 @@ class MCPClient:
await self.disconnect()
# Return a default error result instead of raising
return MCPCallToolResult(
content=[TextContent(type="text", text=f"{str(e)}")], # Empty content for error case
content=[
TextContent(type="text", text=f"{str(e)}")
], # Empty content for error case
isError=True,
)
litellm/files/main.py
+27 -1
View File
@@ -731,7 +731,7 @@ def file_list(
async def afile_content(
file_id: str,
custom_llm_provider: Literal["openai", "azure"] = "openai",
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -887,6 +887,32 @@ def file_content(
client=client,
litellm_params=litellm_params_dict,
)
elif custom_llm_provider == "vertex_ai":
api_base = optional_params.api_base or ""
vertex_ai_project = (
optional_params.vertex_project
or litellm.vertex_project
or get_secret_str("VERTEXAI_PROJECT")
)
vertex_ai_location = (
optional_params.vertex_location
or litellm.vertex_location
or get_secret_str("VERTEXAI_LOCATION")
)
vertex_credentials = optional_params.vertex_credentials or get_secret_str(
"VERTEXAI_CREDENTIALS"
)
response = vertex_ai_files_instance.file_content(
_is_async=_is_async,
file_content_request=_file_content_request,
api_base=api_base,
vertex_credentials=vertex_credentials,
vertex_project=vertex_ai_project,
vertex_location=vertex_ai_location,
timeout=timeout,
max_retries=optional_params.max_retries,
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'custom_llm_provider'. Supported providers are 'openai', 'azure', 'vertex_ai'.".format(
litellm/integrations/custom_guardrail.py
+4 -1
View File
@@ -357,6 +357,7 @@ class CustomGuardrail(CustomLogger):
end_time: Optional[float] = None,
duration: Optional[float] = None,
masked_entity_count: Optional[Dict[str, int]] = None,
guardrail_provider: Optional[str] = None,
) -> None:
"""
Builds `StandardLoggingGuardrailInformation` and adds it to the request metadata so it can be used for logging to DataDog, Langfuse, etc.
@@ -367,6 +368,7 @@ class CustomGuardrail(CustomLogger):
slg = StandardLoggingGuardrailInformation(
guardrail_name=self.guardrail_name,
guardrail_provider=guardrail_provider,
guardrail_mode=(
GuardrailMode(**self.event_hook.model_dump()) # type: ignore
if isinstance(self.event_hook, Mode)
@@ -487,7 +489,8 @@ class CustomGuardrail(CustomLogger):
"""
Update the guardrails litellm params in memory
"""
pass
for key, value in vars(litellm_params).items():
setattr(self, key, value)
def log_guardrail_information(func):
litellm/integrations/datadog/datadog_llm_obs.py
+111 -2
View File
@@ -148,10 +148,22 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger):
),
),
}
verbose_logger.debug("payload %s", json.dumps(payload, indent=4))
# serialize datetime objects - for budget reset time in spend metrics
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
try:
verbose_logger.debug("payload %s", safe_dumps(payload))
except Exception as debug_error:
verbose_logger.debug(
"payload serialization failed: %s", str(debug_error)
)
json_payload = safe_dumps(payload)
response = await self.async_client.post(
url=self.intake_url,
json=payload,
content=json_payload,
headers={
"DD-API-KEY": self.DD_API_KEY,
"Content-Type": "application/json",
@@ -486,6 +498,7 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger):
"guardrail_information": standard_logging_payload.get(
"guardrail_information", None
),
"is_streamed_request": self._get_stream_value_from_payload(standard_logging_payload),
}
#########################################################
@@ -494,6 +507,12 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger):
latency_metrics = self._get_latency_metrics(standard_logging_payload)
_metadata.update({"latency_metrics": dict(latency_metrics)})
#########################################################
# Add spend metrics to metadata
#########################################################
spend_metrics = self._get_spend_metrics(standard_logging_payload)
_metadata.update({"spend_metrics": dict(spend_metrics)})
## extract tool calls and add to metadata
tool_call_metadata = self._extract_tool_call_metadata(standard_logging_payload)
_metadata.update(tool_call_metadata)
@@ -543,6 +562,96 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger):
return latency_metrics
def _get_stream_value_from_payload(self, standard_logging_payload: StandardLoggingPayload) -> bool:
"""
Extract the stream value from standard logging payload.
The stream field in StandardLoggingPayload is only set to True for completed streaming responses.
For non-streaming requests, it's None. The original stream parameter is in model_parameters.
Returns:
bool: True if this was a streaming request, False otherwise
"""
# Check top-level stream field first (only True for completed streaming)
stream_value = standard_logging_payload.get("stream")
if stream_value is True:
return True
# Fallback to model_parameters.stream for original request parameters
model_params = standard_logging_payload.get("model_parameters", {})
if isinstance(model_params, dict):
stream_value = model_params.get("stream")
if stream_value is True:
return True
# Default to False for non-streaming requests
return False
def _get_spend_metrics(
self, standard_logging_payload: StandardLoggingPayload
) -> DDLLMObsSpendMetrics:
"""
Get the spend metrics from the standard logging payload
"""
spend_metrics: DDLLMObsSpendMetrics = DDLLMObsSpendMetrics()
# send response cost
spend_metrics["response_cost"] = standard_logging_payload.get(
"response_cost", 0.0
)
# Get budget information from metadata
metadata = standard_logging_payload.get("metadata", {})
# API key max budget
user_api_key_max_budget = metadata.get("user_api_key_max_budget")
if user_api_key_max_budget is not None:
spend_metrics["user_api_key_max_budget"] = float(user_api_key_max_budget)
# API key spend
user_api_key_spend = metadata.get("user_api_key_spend")
if user_api_key_spend is not None:
try:
spend_metrics["user_api_key_spend"] = float(user_api_key_spend)
except (ValueError, TypeError):
verbose_logger.debug(
f"Invalid user_api_key_spend value: {user_api_key_spend}"
)
# API key budget reset datetime
user_api_key_budget_reset_at = metadata.get("user_api_key_budget_reset_at")
if user_api_key_budget_reset_at is not None:
try:
from datetime import datetime, timezone
budget_reset_at = None
if isinstance(user_api_key_budget_reset_at, str):
# Handle ISO format strings that might have 'Z' suffix
iso_string = user_api_key_budget_reset_at.replace("Z", "+00:00")
budget_reset_at = datetime.fromisoformat(iso_string)
elif isinstance(user_api_key_budget_reset_at, datetime):
budget_reset_at = user_api_key_budget_reset_at
if budget_reset_at is not None:
# Preserve timezone info if already present
if budget_reset_at.tzinfo is None:
budget_reset_at = budget_reset_at.replace(tzinfo=timezone.utc)
# Convert to ISO string format for JSON serialization
# This prevents circular reference issues and ensures proper timezone representation
iso_string = budget_reset_at.isoformat()
spend_metrics["user_api_key_budget_reset_at"] = iso_string
# Debug logging to verify the conversion
verbose_logger.debug(
f"Converted budget_reset_at to ISO format: {iso_string}"
)
except Exception as e:
verbose_logger.debug(f"Error processing budget reset datetime: {e}")
verbose_logger.debug(f"Original value: {user_api_key_budget_reset_at}")
return spend_metrics
def _process_input_messages_preserving_tool_calls(
self, messages: List[Any]
) -> List[Dict[str, Any]]:
litellm/integrations/langfuse/langfuse.py
+21 -1
View File
@@ -15,7 +15,7 @@ from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
from litellm.secret_managers.main import str_to_bool
from litellm.types.integrations.langfuse import *
from litellm.types.llms.openai import HttpxBinaryResponseContent
from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse
from litellm.types.utils import (
EmbeddingResponse,
ImageResponse,
@@ -196,6 +196,7 @@ class LangFuseLogger:
TranscriptionResponse,
RerankResponse,
HttpxBinaryResponseContent,
ResponsesAPIResponse,
],
start_time: Optional[datetime] = None,
end_time: Optional[datetime] = None,
@@ -305,6 +306,7 @@ class LangFuseLogger:
TranscriptionResponse,
RerankResponse,
HttpxBinaryResponseContent,
ResponsesAPIResponse,
],
prompt: dict,
level: str,
@@ -369,6 +371,11 @@ class LangFuseLogger:
):
input = prompt
output = response_obj.results
elif response_obj is not None and isinstance(
response_obj, litellm.ResponsesAPIResponse
):
input = prompt
output = self._get_responses_api_content_for_langfuse(response_obj)
elif (
kwargs.get("call_type") is not None
and kwargs.get("call_type") == "_arealtime"
@@ -768,6 +775,19 @@ class LangFuseLogger:
else:
return None
@staticmethod
def _get_responses_api_content_for_langfuse(
response_obj: ResponsesAPIResponse,
):
"""
Get the responses API content for Langfuse logging
"""
if hasattr(response_obj, 'output') and response_obj.output:
# ResponsesAPIResponse.output is a list of strings
return response_obj.output
else:
return None
@staticmethod
def _get_langfuse_tags(
standard_logging_object: Optional[StandardLoggingPayload],
litellm/integrations/langsmith.py
+34 -24
View File
@@ -39,6 +39,7 @@ class LangsmithLogger(CustomBatchLogger):
langsmith_api_key: Optional[str] = None,
langsmith_project: Optional[str] = None,
langsmith_base_url: Optional[str] = None,
langsmith_sampling_rate: Optional[float] = None,
**kwargs,
):
self.flush_lock = asyncio.Lock()
@@ -49,7 +50,8 @@ class LangsmithLogger(CustomBatchLogger):
langsmith_base_url=langsmith_base_url,
)
self.sampling_rate: float = (
float(os.getenv("LANGSMITH_SAMPLING_RATE")) # type: ignore
langsmith_sampling_rate
or float(os.getenv("LANGSMITH_SAMPLING_RATE")) # type: ignore
if os.getenv("LANGSMITH_SAMPLING_RATE") is not None
and os.getenv("LANGSMITH_SAMPLING_RATE").strip().isdigit() # type: ignore
else 1.0
@@ -76,26 +78,14 @@ class LangsmithLogger(CustomBatchLogger):
langsmith_base_url: Optional[str] = None,
) -> LangsmithCredentialsObject:
_credentials_api_key = langsmith_api_key or os.getenv("LANGSMITH_API_KEY")
if _credentials_api_key is None:
raise Exception(
"Invalid Langsmith API Key given. _credentials_api_key=None."
)
_credentials_project = (
langsmith_project or os.getenv("LANGSMITH_PROJECT") or "litellm-completion"
)
if _credentials_project is None:
raise Exception(
"Invalid Langsmith API Key given. _credentials_project=None."
)
_credentials_base_url = (
langsmith_base_url
or os.getenv("LANGSMITH_BASE_URL")
or "https://api.smith.langchain.com"
)
if _credentials_base_url is None:
raise Exception(
"Invalid Langsmith API Key given. _credentials_base_url=None."
)
return LangsmithCredentialsObject(
LANGSMITH_API_KEY=_credentials_api_key,
@@ -200,12 +190,7 @@ class LangsmithLogger(CustomBatchLogger):
def log_success_event(self, kwargs, response_obj, start_time, end_time):
try:
sampling_rate = (
float(os.getenv("LANGSMITH_SAMPLING_RATE")) # type: ignore
if os.getenv("LANGSMITH_SAMPLING_RATE") is not None
and os.getenv("LANGSMITH_SAMPLING_RATE").strip().isdigit() # type: ignore
else 1.0
)
sampling_rate = self._get_sampling_rate_to_use_for_request(kwargs=kwargs)
random_sample = random.random()
if random_sample > sampling_rate:
verbose_logger.info(
@@ -219,6 +204,7 @@ class LangsmithLogger(CustomBatchLogger):
kwargs,
response_obj,
)
credentials = self._get_credentials_to_use_for_request(kwargs=kwargs)
data = self._prepare_log_data(
kwargs=kwargs,
@@ -245,7 +231,7 @@ class LangsmithLogger(CustomBatchLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
try:
sampling_rate = self.sampling_rate
sampling_rate = self._get_sampling_rate_to_use_for_request(kwargs=kwargs)
random_sample = random.random()
if random_sample > sampling_rate:
verbose_logger.info(
@@ -286,7 +272,7 @@ class LangsmithLogger(CustomBatchLogger):
)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
sampling_rate = self.sampling_rate
sampling_rate = self._get_sampling_rate_to_use_for_request(kwargs=kwargs)
random_sample = random.random()
if random_sample > sampling_rate:
verbose_logger.info(
@@ -417,6 +403,17 @@ class LangsmithLogger(CustomBatchLogger):
for queue_object in self.log_queue:
credentials = queue_object["credentials"]
# if credential missing, skip - log warning
if (
credentials["LANGSMITH_API_KEY"] is None
or credentials["LANGSMITH_PROJECT"] is None
):
verbose_logger.warning(
"Langsmith Logging - credentials missing - api_key: %s, project: %s",
credentials["LANGSMITH_API_KEY"],
credentials["LANGSMITH_PROJECT"],
)
continue
key = CredentialsKey(
api_key=credentials["LANGSMITH_API_KEY"],
project=credentials["LANGSMITH_PROJECT"],
@@ -432,6 +429,19 @@ class LangsmithLogger(CustomBatchLogger):
return log_queue_by_credentials
def _get_sampling_rate_to_use_for_request(self, kwargs: Dict[str, Any]) -> float:
standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = (
kwargs.get("standard_callback_dynamic_params", None)
)
sampling_rate: float = self.sampling_rate
if standard_callback_dynamic_params is not None:
_sampling_rate = standard_callback_dynamic_params.get(
"langsmith_sampling_rate"
)
if _sampling_rate is not None:
sampling_rate = float(_sampling_rate)
return sampling_rate
def _get_credentials_to_use_for_request(
self, kwargs: Dict[str, Any]
) -> LangsmithCredentialsObject:
@@ -442,9 +452,9 @@ class LangsmithLogger(CustomBatchLogger):
Otherwise, use the default credentials.
"""
standard_callback_dynamic_params: Optional[
StandardCallbackDynamicParams
] = kwargs.get("standard_callback_dynamic_params", None)
standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = (
kwargs.get("standard_callback_dynamic_params", None)
)
if standard_callback_dynamic_params is not None:
credentials = self.get_credentials_from_env(
langsmith_api_key=standard_callback_dynamic_params.get(
litellm/integrations/opik/opik.py
+5 -4
View File
@@ -3,6 +3,7 @@ Opik Logger that logs LLM events to an Opik server
"""
import asyncio
from datetime import timezone
import json
import traceback
from typing import Dict, List
@@ -291,8 +292,8 @@ class OpikLogger(CustomBatchLogger):
"project_name": project_name,
"id": trace_id,
"name": trace_name,
"start_time": start_time.isoformat() + "Z",
"end_time": end_time.isoformat() + "Z",
"start_time": start_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z"),
"end_time": end_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z"),
"input": input_data,
"output": output_data,
"metadata": metadata,
@@ -312,8 +313,8 @@ class OpikLogger(CustomBatchLogger):
"parent_span_id": parent_span_id,
"name": span_name,
"type": "llm",
"start_time": start_time.isoformat() + "Z",
"end_time": end_time.isoformat() + "Z",
"start_time": start_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z"),
"end_time": end_time.astimezone(timezone.utc).isoformat().replace("+00:00", "Z"),
"input": input_data,
"output": output_data,
"metadata": metadata,
litellm/integrations/posthog.py
+333
View File
@@ -0,0 +1,333 @@
"""
PostHog Integration - sends LLM analytics events to PostHog
Follows PostHog's LLM Analytics format: https://posthog.com/docs/llm-analytics/manual-capture
async_log_success_event: stores batch of events in memory and flushes to PostHog
async_log_failure_event: logs failed LLM calls with error information
For batching specific details see CustomBatchLogger class
"""
import asyncio
import os
import uuid
from typing import Any, Dict, Optional
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.integrations.posthog import (
POSTHOG_MAX_BATCH_SIZE,
PostHogEventPayload,
)
from litellm.types.utils import StandardLoggingPayload
class PostHogLogger(CustomBatchLogger):
def __init__(self, **kwargs):
"""
Initializes the PostHog logger, checks if the correct env variables are set
Required environment variables:
`POSTHOG_API_KEY` - your PostHog API key
`POSTHOG_API_URL` - your PostHog API URL (defaults to https://app.posthog.com)
"""
try:
verbose_logger.debug("PostHog: in init posthog logger")
if os.getenv("POSTHOG_API_KEY", None) is None:
raise Exception("POSTHOG_API_KEY is not set, set 'POSTHOG_API_KEY=<>'")
self.async_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
self.sync_client = _get_httpx_client()
self.POSTHOG_API_KEY = os.getenv("POSTHOG_API_KEY")
posthog_api_url = os.getenv("POSTHOG_API_URL", "https://us.i.posthog.com")
self.posthog_host = posthog_api_url.rstrip('/')
self.capture_url = f"{self.posthog_host}/batch/"
self._async_initialized = False
self.flush_lock = None
self.log_queue = []
super().__init__(
**kwargs, flush_lock=None, batch_size=POSTHOG_MAX_BATCH_SIZE
)
except Exception as e:
verbose_logger.exception(
f"PostHog: Got exception on init PostHog client {str(e)}"
)
raise e
def log_success_event(self, kwargs, response_obj, start_time, end_time):
try:
verbose_logger.debug(
"PostHog: Sync logging - Enters logging function for model %s", kwargs
)
event_payload = self.create_posthog_event_payload(kwargs)
headers = {
"Content-Type": "application/json",
}
payload = self._create_posthog_payload([event_payload])
response = self.sync_client.post(
url=self.capture_url,
json=payload,
headers=headers,
)
response.raise_for_status()
if response.status_code != 200:
raise Exception(
f"Response from PostHog API status_code: {response.status_code}, text: {response.text}"
)
verbose_logger.debug("PostHog: Sync event successfully sent")
except Exception as e:
verbose_logger.exception(f"PostHog Sync Layer Error - {str(e)}")
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
try:
verbose_logger.debug(
"PostHog: Async logging - Enters logging function for model %s", kwargs
)
self._ensure_async_setup() # Lazy initialization
await self._log_async_event(kwargs, response_obj, start_time, end_time)
except Exception as e:
verbose_logger.exception(f"PostHog Layer Error - {str(e)}")
pass
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
try:
verbose_logger.debug(
"PostHog: Async logging - Enters logging function for model %s", kwargs
)
self._ensure_async_setup() # Lazy initialization
await self._log_async_event(kwargs, response_obj, start_time, end_time)
except Exception as e:
verbose_logger.exception(f"PostHog Layer Error - {str(e)}")
pass
async def _log_async_event(self, kwargs, response_obj=None, start_time=0.0, end_time=0.0):
# Note: response_obj, start_time, end_time not used - all data comes from kwargs
event_payload = self.create_posthog_event_payload(kwargs)
self.log_queue.append(event_payload)
verbose_logger.debug(
f"PostHog, event added to queue. Will flush in {self.flush_interval} seconds..."
)
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
def create_posthog_event_payload(self, kwargs: Dict[str, Any]) -> PostHogEventPayload:
"""
Helper function to create a PostHog event payload for logging
Args:
kwargs (Dict[str, Any]): request kwargs containing standard_logging_object
Returns:
PostHogEventPayload: defined in types.py
"""
standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get(
"standard_logging_object", None
)
if standard_logging_object is None:
raise ValueError("standard_logging_object not found in kwargs")
call_type = standard_logging_object.get("call_type", "")
event_name = "$ai_embedding" if call_type == "embedding" else "$ai_generation"
properties = self._create_posthog_properties(
standard_logging_object=standard_logging_object,
kwargs=kwargs,
event_name=event_name,
)
distinct_id = self._get_distinct_id(standard_logging_object, kwargs)
return PostHogEventPayload(
event=event_name,
properties=properties,
distinct_id=distinct_id,
)
def _create_posthog_properties(
self,
standard_logging_object: StandardLoggingPayload,
kwargs: Dict[str, Any],
event_name: str,
) -> Dict[str, Any]:
"""Create PostHog properties following LLM Analytics spec"""
properties = {}
# Core model information
properties["$ai_model"] = self._safe_get(standard_logging_object, "model", "")
properties["$ai_provider"] = self._safe_get(standard_logging_object, "custom_llm_provider", "")
# Input/Output data
messages = self._safe_get(standard_logging_object, "messages")
if messages is not None:
properties["$ai_input"] = messages
if event_name == "$ai_generation":
response = self._safe_get(standard_logging_object, "response")
if response is not None:
properties["$ai_output_choices"] = response
# Token information
properties["$ai_input_tokens"] = self._safe_get(standard_logging_object, "prompt_tokens", 0)
if event_name == "$ai_generation":
properties["$ai_output_tokens"] = self._safe_get(standard_logging_object, "completion_tokens", 0)
# Cost and performance
response_cost = self._safe_get(standard_logging_object, "response_cost")
if response_cost is not None:
properties["$ai_total_cost_usd"] = response_cost
properties["$ai_latency"] = self._safe_get(standard_logging_object, "response_time", 0.0)
# Error handling
if self._safe_get(standard_logging_object, "status") == "failure":
properties["$ai_is_error"] = True
error_str = self._safe_get(standard_logging_object, "error_str")
if error_str is not None:
properties["$ai_error"] = error_str
# Add trace properties
self._add_trace_properties(properties, kwargs)
# Add custom metadata fields
self._add_custom_metadata_properties(properties, kwargs)
return properties
def _add_trace_properties(self, properties: Dict[str, Any], kwargs: Dict[str, Any]):
standard_logging_object = self._safe_get(kwargs, "standard_logging_object", {})
trace_id = self._safe_get(standard_logging_object, "trace_id", self._safe_uuid())
properties["$ai_trace_id"] = trace_id
span_id = self._safe_get(standard_logging_object, "id", self._safe_uuid())
properties["$ai_span_id"] = span_id
metadata = self._extract_metadata(kwargs)
parent_id = metadata.get("parent_run_id") or metadata.get("parent_id")
if parent_id:
properties["$ai_parent_id"] = parent_id
def _add_custom_metadata_properties(self, properties: Dict[str, Any], kwargs: Dict[str, Any]):
"""Add custom metadata fields to PostHog properties"""
metadata = self._extract_metadata(kwargs)
if not isinstance(metadata, dict):
return
litellm_internal_fields = {
"endpoint", "caching_groups", "user_api_key_hash", "user_api_key_alias",
"user_api_key_team_id", "user_api_key_user_id", "user_api_key_org_id",
"user_api_key_team_alias", "user_api_key_end_user_id", "user_api_key_user_email",
"user_api_key", "user_api_end_user_max_budget", "litellm_api_version",
"global_max_parallel_requests", "user_api_key_team_max_budget", "user_api_key_team_spend",
"user_api_key_spend", "user_api_key_max_budget", "user_api_key_model_max_budget",
"user_api_key_metadata", "headers", "litellm_parent_otel_span", "requester_ip_address",
"model_group", "model_group_size", "deployment", "model_info", "api_base",
"caching_groups", "hidden_params", "parent_run_id", "parent_id", "user_id"
}
for key, value in metadata.items():
if key not in litellm_internal_fields:
properties[key] = value
def _get_distinct_id(
self, standard_logging_object: StandardLoggingPayload, kwargs: Dict[str, Any]
) -> str:
metadata = self._extract_metadata(kwargs)
user_id = self._safe_get(metadata, "user_id")
if user_id:
return str(user_id)
end_user = self._safe_get(standard_logging_object, "end_user")
if end_user:
return str(end_user)
trace_id = self._safe_get(standard_logging_object, "trace_id")
if trace_id:
return str(trace_id)
return self._safe_uuid()
async def async_send_batch(self):
"""
Sends the in memory logs queue to PostHog API
Raises:
Raises a NON Blocking verbose_logger.exception if an error occurs
"""
try:
if not self.log_queue:
return
verbose_logger.debug(
f"PostHog: Sending batch of {len(self.log_queue)} events"
)
headers = {
"Content-Type": "application/json",
}
payload = self._create_posthog_payload(list(self.log_queue))
response = await self.async_client.post(
url=self.capture_url,
json=payload,
headers=headers,
)
response.raise_for_status()
if response.status_code != 200:
raise Exception(
f"Response from PostHog API status_code: {response.status_code}, text: {response.text}"
)
verbose_logger.debug(
f"PostHog: Batch of {len(self.log_queue)} events successfully sent"
)
except Exception as e:
verbose_logger.exception(f"PostHog Error sending batch API - {str(e)}")
def _ensure_async_setup(self):
if not self._async_initialized:
try:
self.flush_lock = asyncio.Lock()
asyncio.create_task(self.periodic_flush())
self._async_initialized = True
verbose_logger.debug("PostHog: Async components initialized")
except Exception as e:
verbose_logger.error(f"PostHog: Failed to initialize async components: {str(e)}")
raise
def _extract_metadata(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
litellm_params = kwargs.get("litellm_params", {}) or {}
return litellm_params.get("metadata", {}) or {}
def _safe_uuid(self) -> str:
return str(uuid.uuid4())
def _create_posthog_payload(self, events: list) -> Dict[str, Any]:
return {"api_key": self.POSTHOG_API_KEY, "batch": events}
def _safe_get(self, obj: Any, key: str, default: Any = None) -> Any:
if obj is None or not hasattr(obj, 'get'):
return default
return obj.get(key, default)
litellm/litellm_core_utils/cached_imports.py
+56
View File
@@ -0,0 +1,56 @@
"""
Cached imports module for LiteLLM.
This module provides cached import functionality to avoid repeated imports
inside functions that are critical to performance.
"""
from typing import TYPE_CHECKING, Callable, Optional, Type
# Type annotations for cached imports
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.coroutine_checker import CoroutineChecker
# Global cache variables
_LiteLLMLogging: Optional[Type["Logging"]] = None
_coroutine_checker: Optional["CoroutineChecker"] = None
_set_callbacks: Optional[Callable] = None
def get_litellm_logging_class() -> Type["Logging"]:
"""Get the cached LiteLLM Logging class, initializing if needed."""
global _LiteLLMLogging
if _LiteLLMLogging is not None:
return _LiteLLMLogging
from litellm.litellm_core_utils.litellm_logging import Logging
_LiteLLMLogging = Logging
return _LiteLLMLogging
def get_coroutine_checker() -> "CoroutineChecker":
"""Get the cached coroutine checker instance, initializing if needed."""
global _coroutine_checker
if _coroutine_checker is not None:
return _coroutine_checker
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
_coroutine_checker = coroutine_checker
return _coroutine_checker
def get_set_callbacks() -> Callable:
"""Get the cached set_callbacks function, initializing if needed."""
global _set_callbacks
if _set_callbacks is not None:
return _set_callbacks
from litellm.litellm_core_utils.litellm_logging import set_callbacks
_set_callbacks = set_callbacks
return _set_callbacks
def clear_cached_imports() -> None:
"""Clear all cached imports. Useful for testing or memory management."""
global _LiteLLMLogging, _coroutine_checker, _set_callbacks
_LiteLLMLogging = None
_coroutine_checker = None
_set_callbacks = None
litellm/litellm_core_utils/core_helpers.py
+21 -6
View File
@@ -228,9 +228,11 @@ def safe_deep_copy(data):
"""
Safe Deep Copy
The LiteLLM Request has some object that can-not be pickled / deep copied
Use this function to safely deep copy the LiteLLM Request
The LiteLLM request may contain objects that cannot be pickled/deep-copied
(e.g., tracing spans, locks, clients).
This helper deep-copies each top-level key independently; on failure keeps
original ref
"""
import copy
@@ -255,9 +257,22 @@ def safe_deep_copy(data):
"litellm_parent_otel_span"
)
data["litellm_metadata"]["litellm_parent_otel_span"] = "placeholder"
new_data = copy.deepcopy(data)
# Step 2: re-add the litellm_parent_otel_span after doing a deep copy
# Step 2: Per-key deepcopy with fallback
if isinstance(data, dict):
new_data = {}
for k, v in data.items():
try:
new_data[k] = copy.deepcopy(v)
except Exception:
new_data[k] = v
else:
try:
new_data = copy.deepcopy(data)
except Exception:
new_data = data
# Step 3: re-add the litellm_parent_otel_span after doing a deep copy
if isinstance(data, dict) and litellm_parent_otel_span is not None:
if "metadata" in data and "litellm_parent_otel_span" in data["metadata"]:
data["metadata"]["litellm_parent_otel_span"] = litellm_parent_otel_span
@@ -268,4 +283,4 @@ def safe_deep_copy(data):
data["litellm_metadata"][
"litellm_parent_otel_span"
] = litellm_parent_otel_span
return new_data
return new_data
litellm/litellm_core_utils/coroutine_checker.py
+63
View File
@@ -0,0 +1,63 @@
# CoroutineChecker utility for checking if functions/callables are coroutines or coroutine functions
import inspect
from typing import Any
from weakref import WeakKeyDictionary
from litellm.constants import (
COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY,
)
class CoroutineChecker:
"""Utility class for checking coroutine status of functions and callables.
Simple bounded cache using WeakKeyDictionary to avoid memory leaks.
"""
def __init__(self):
self._cache = WeakKeyDictionary()
self._max_size = COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY
def is_async_callable(self, callback: Any) -> bool:
"""Fast, cached check for whether a callback is an async function.
Falls back gracefully if the object cannot be weak-referenced or cached.
2.59x speedup.
"""
# Fast path: check cache first (most common case)
try:
cached = self._cache.get(callback)
if cached is not None:
return cached
except Exception:
pass
# Determine target - optimized path for common cases
target = callback
if not inspect.isfunction(target) and not inspect.ismethod(target):
try:
call_attr = getattr(target, "__call__", None)
if call_attr is not None:
target = call_attr
except Exception:
pass
# Compute result
try:
result = inspect.iscoroutinefunction(target)
except Exception:
result = False
# Cache the result with size enforcement
try:
# Simple size enforcement: clear cache if it gets too large
if len(self._cache) >= self._max_size:
self._cache.clear()
self._cache[callback] = result
except Exception:
pass
return result
# Global instance for backward compatibility and convenience
coroutine_checker = CoroutineChecker()
litellm/litellm_core_utils/custom_logger_registry.py
+4
View File
@@ -33,6 +33,7 @@ from litellm.integrations.mlflow import MlflowLogger
from litellm.integrations.openmeter import OpenMeterLogger
from litellm.integrations.opentelemetry import OpenTelemetry
from litellm.integrations.opik.opik import OpikLogger
from litellm.integrations.posthog import PostHogLogger
try:
from litellm_enterprise.integrations.prometheus import PrometheusLogger
@@ -46,6 +47,7 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i
VectorStorePreCallHook,
)
from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHandler
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import _PROXY_DynamicRateLimitHandlerV3
class CustomLoggerRegistry:
@@ -85,9 +87,11 @@ class CustomLoggerRegistry:
"s3_v2": S3Logger,
"aws_sqs": SQSLogger,
"dynamic_rate_limiter": _PROXY_DynamicRateLimitHandler,
"dynamic_rate_limiter_v3": _PROXY_DynamicRateLimitHandlerV3,
"vector_store_pre_call_hook": VectorStorePreCallHook,
"dotprompt": DotpromptManager,
"cloudzero": CloudZeroLogger,
"posthog": PostHogLogger,
}
try:
litellm/litellm_core_utils/duration_parser.py
+1
View File
@@ -158,6 +158,7 @@ def _setup_timezone(
"US/Eastern": timezone(timedelta(hours=-4)), # EDT
"US/Pacific": timezone(timedelta(hours=-7)), # PDT
"Asia/Kolkata": timezone(timedelta(hours=5, minutes=30)), # IST
"Asia/Bangkok": timezone(timedelta(hours=7)), # ICT (Indochina Time)
"Europe/London": timezone(timedelta(hours=1)), # BST
"UTC": timezone.utc,
}
litellm/litellm_core_utils/exception_mapping_utils.py
+71 -99
View File
@@ -6,6 +6,7 @@ import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.types.utils import LlmProviders
from ..exceptions import (
APIConnectionError,
@@ -556,7 +557,7 @@ def exception_type( # type: ignore # noqa: PLR0915
model=model,
llm_provider="anthropic",
)
elif "overloaded_error" in error_str:
elif "overloaded_error" in error_str or "Overloaded" in error_str:
exception_mapping_worked = True
raise InternalServerError(
message="AnthropicError - {}".format(error_str),
@@ -762,7 +763,7 @@ def exception_type( # type: ignore # noqa: PLR0915
error_str += "XXXXXXX" + '"'
raise AuthenticationError(
message=f"{custom_llm_provider}Exception: Authentication Error - {error_str}",
message=f"{custom_llm_provider.capitalize()}Exception: Authentication Error - {error_str}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
@@ -771,14 +772,14 @@ def exception_type( # type: ignore # noqa: PLR0915
elif "model's maximum context limit" in error_str:
exception_mapping_worked = True
raise ContextWindowExceededError(
message=f"{custom_llm_provider}Exception: Context Window Error - {error_str}",
message=f"{custom_llm_provider.capitalize()}Exception: Context Window Error - {error_str}",
model=model,
llm_provider=custom_llm_provider,
)
elif "token_quota_reached" in error_str:
exception_mapping_worked = True
raise RateLimitError(
message=f"{custom_llm_provider}Exception: Rate Limit Errror - {error_str}",
message=f"{custom_llm_provider.capitalize()}Exception: Rate Limit Errror - {error_str}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
@@ -789,14 +790,14 @@ def exception_type( # type: ignore # noqa: PLR0915
):
exception_mapping_worked = True
raise litellm.InternalServerError(
message=f"{custom_llm_provider}Exception - {original_exception.message}",
message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}",
llm_provider=custom_llm_provider,
model=model,
)
elif "model_no_support_for_function" in error_str:
exception_mapping_worked = True
raise BadRequestError(
message=f"{custom_llm_provider}Exception - Use 'watsonx_text' route instead. IBM WatsonX does not support `/text/chat` endpoint. - {error_str}",
message=f"{custom_llm_provider.capitalize()}Exception - Use 'watsonx_text' route instead. IBM WatsonX does not support `/text/chat` endpoint. - {error_str}",
llm_provider=custom_llm_provider,
model=model,
)
@@ -804,7 +805,7 @@ def exception_type( # type: ignore # noqa: PLR0915
if original_exception.status_code == 500:
exception_mapping_worked = True
raise litellm.InternalServerError(
message=f"{custom_llm_provider}Exception - {original_exception.message}",
message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}",
llm_provider=custom_llm_provider,
model=model,
)
@@ -814,28 +815,28 @@ def exception_type( # type: ignore # noqa: PLR0915
):
exception_mapping_worked = True
raise AuthenticationError(
message=f"{custom_llm_provider}Exception - {original_exception.message}",
message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}",
llm_provider=custom_llm_provider,
model=model,
)
elif original_exception.status_code == 400:
exception_mapping_worked = True
raise BadRequestError(
message=f"{custom_llm_provider}Exception - {original_exception.message}",
message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}",
llm_provider=custom_llm_provider,
model=model,
)
elif original_exception.status_code == 404:
exception_mapping_worked = True
raise NotFoundError(
message=f"{custom_llm_provider}Exception - {original_exception.message}",
message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}",
llm_provider=custom_llm_provider,
model=model,
)
elif original_exception.status_code == 408:
exception_mapping_worked = True
raise Timeout(
message=f"{custom_llm_provider}Exception - {original_exception.message}",
message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}",
model=model,
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
@@ -846,7 +847,7 @@ def exception_type( # type: ignore # noqa: PLR0915
):
exception_mapping_worked = True
raise BadRequestError(
message=f"{custom_llm_provider}Exception - {original_exception.message}",
message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}",
model=model,
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
@@ -854,7 +855,7 @@ def exception_type( # type: ignore # noqa: PLR0915
elif original_exception.status_code == 429:
exception_mapping_worked = True
raise RateLimitError(
message=f"{custom_llm_provider}Exception - {original_exception.message}",
message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}",
model=model,
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
@@ -862,7 +863,7 @@ def exception_type( # type: ignore # noqa: PLR0915
elif original_exception.status_code == 503:
exception_mapping_worked = True
raise ServiceUnavailableError(
message=f"{custom_llm_provider}Exception - {original_exception.message}",
message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}",
model=model,
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
@@ -870,7 +871,7 @@ def exception_type( # type: ignore # noqa: PLR0915
elif original_exception.status_code == 504: # gateway timeout error
exception_mapping_worked = True
raise Timeout(
message=f"{custom_llm_provider}Exception - {original_exception.message}",
message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}",
model=model,
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
@@ -1168,9 +1169,9 @@ def exception_type( # type: ignore # noqa: PLR0915
exception_status_code=original_exception.status_code,
)
elif (
custom_llm_provider == "vertex_ai"
or custom_llm_provider == "vertex_ai_beta"
or custom_llm_provider == "gemini"
custom_llm_provider == LlmProviders.VERTEX_AI
or custom_llm_provider == LlmProviders.VERTEX_AI_BETA
or custom_llm_provider == LlmProviders.GEMINI
):
if (
"Vertex AI API has not been used in project" in error_str
@@ -1178,9 +1179,9 @@ def exception_type( # type: ignore # noqa: PLR0915
):
exception_mapping_worked = True
raise BadRequestError(
message=f"litellm.BadRequestError: VertexAIException - {error_str}",
message=f"litellm.BadRequestError: {custom_llm_provider}Exception - {error_str}",
model=model,
llm_provider="vertex_ai",
llm_provider=custom_llm_provider,
response=httpx.Response(
status_code=400,
request=httpx.Request(
@@ -1193,7 +1194,7 @@ def exception_type( # type: ignore # noqa: PLR0915
if "400 Request payload size exceeds" in error_str:
exception_mapping_worked = True
raise ContextWindowExceededError(
message=f"VertexException - {error_str}",
message=f"{custom_llm_provider.capitalize()}Exception - {error_str}",
model=model,
llm_provider=custom_llm_provider,
)
@@ -1203,9 +1204,9 @@ def exception_type( # type: ignore # noqa: PLR0915
):
exception_mapping_worked = True
raise litellm.InternalServerError(
message=f"litellm.InternalServerError: VertexAIException - {error_str}",
message=f"litellm.InternalServerError: {custom_llm_provider}Exception - {error_str}",
model=model,
llm_provider="vertex_ai",
llm_provider=custom_llm_provider,
response=httpx.Response(
status_code=500,
content=str(original_exception),
@@ -1216,7 +1217,7 @@ def exception_type( # type: ignore # noqa: PLR0915
elif "API key not valid." in error_str:
exception_mapping_worked = True
raise AuthenticationError(
message=f"{custom_llm_provider}Exception - {error_str}",
message=f"{custom_llm_provider.capitalize()}Exception - {error_str}",
model=model,
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
@@ -1224,9 +1225,9 @@ def exception_type( # type: ignore # noqa: PLR0915
elif "403" in error_str:
exception_mapping_worked = True
raise BadRequestError(
message=f"VertexAIException BadRequestError - {error_str}",
message=f"{custom_llm_provider.capitalize()}Exception BadRequestError - {error_str}",
model=model,
llm_provider="vertex_ai",
llm_provider=custom_llm_provider,
response=httpx.Response(
status_code=403,
request=httpx.Request(
@@ -1243,9 +1244,9 @@ def exception_type( # type: ignore # noqa: PLR0915
):
exception_mapping_worked = True
raise ContentPolicyViolationError(
message=f"VertexAIException ContentPolicyViolationError - {error_str}",
message=f"{custom_llm_provider.capitalize()}Exception ContentPolicyViolationError - {error_str}",
model=model,
llm_provider="vertex_ai",
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
response=httpx.Response(
status_code=400,
@@ -1264,9 +1265,9 @@ def exception_type( # type: ignore # noqa: PLR0915
):
exception_mapping_worked = True
raise RateLimitError(
message=f"litellm.RateLimitError: VertexAIException - {error_str}",
message=f"litellm.RateLimitError: {custom_llm_provider}Exception - {error_str}",
model=model,
llm_provider="vertex_ai",
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
response=httpx.Response(
status_code=429,
@@ -1282,18 +1283,18 @@ def exception_type( # type: ignore # noqa: PLR0915
):
exception_mapping_worked = True
raise litellm.InternalServerError(
message=f"litellm.InternalServerError: VertexAIException - {error_str}",
message=f"litellm.InternalServerError: {custom_llm_provider}Exception - {error_str}",
model=model,
llm_provider="vertex_ai",
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
)
if hasattr(original_exception, "status_code"):
if original_exception.status_code == 400:
exception_mapping_worked = True
raise BadRequestError(
message=f"VertexAIException BadRequestError - {error_str}",
message=f"{custom_llm_provider.capitalize()}Exception BadRequestError - {error_str}",
model=model,
llm_provider="vertex_ai",
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
response=httpx.Response(
status_code=400,
@@ -1306,21 +1307,35 @@ def exception_type( # type: ignore # noqa: PLR0915
if original_exception.status_code == 401:
exception_mapping_worked = True
raise AuthenticationError(
message=f"VertexAIException - {original_exception.message}",
message=f"{custom_llm_provider.capitalize()}Exception - {error_str}",
llm_provider=custom_llm_provider,
model=model,
)
if original_exception.status_code == 403:
exception_mapping_worked = True
raise PermissionDeniedError(
message=f"{custom_llm_provider.capitalize()}Exception - {error_str}",
llm_provider=custom_llm_provider,
model=model,
response=httpx.Response(
status_code=403,
request=httpx.Request(
method="POST",
url="https://cloud.google.com/vertex-ai/",
),
),
)
if original_exception.status_code == 404:
exception_mapping_worked = True
raise NotFoundError(
message=f"VertexAIException - {original_exception.message}",
message=f"{custom_llm_provider.capitalize()}Exception - {error_str}",
llm_provider=custom_llm_provider,
model=model,
)
if original_exception.status_code == 408:
exception_mapping_worked = True
raise Timeout(
message=f"VertexAIException - {original_exception.message}",
message=f"{custom_llm_provider.capitalize()}Exception - {error_str}",
llm_provider=custom_llm_provider,
model=model,
)
@@ -1328,9 +1343,9 @@ def exception_type( # type: ignore # noqa: PLR0915
if original_exception.status_code == 429:
exception_mapping_worked = True
raise RateLimitError(
message=f"litellm.RateLimitError: VertexAIException - {error_str}",
message=f"litellm.RateLimitError: {custom_llm_provider.capitalize()}Exception - {error_str}",
model=model,
llm_provider="vertex_ai",
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
response=httpx.Response(
status_code=429,
@@ -1343,9 +1358,9 @@ def exception_type( # type: ignore # noqa: PLR0915
if original_exception.status_code == 500:
exception_mapping_worked = True
raise litellm.InternalServerError(
message=f"VertexAIException InternalServerError - {error_str}",
message=f"{custom_llm_provider.capitalize()}Exception InternalServerError - {error_str}",
model=model,
llm_provider="vertex_ai",
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
response=httpx.Response(
status_code=500,
@@ -1353,71 +1368,20 @@ def exception_type( # type: ignore # noqa: PLR0915
request=httpx.Request(method="completion", url="https://github.com/BerriAI/litellm"), # type: ignore
),
)
if original_exception.status_code == 503:
if original_exception.status_code == 502:
exception_mapping_worked = True
raise ServiceUnavailableError(
message=f"VertexAIException - {original_exception.message}",
raise APIConnectionError(
message=f"{custom_llm_provider.capitalize()}Exception - {error_str}",
llm_provider=custom_llm_provider,
model=model,
)
elif custom_llm_provider == "palm" or custom_llm_provider == "gemini":
if "503 Getting metadata" in error_str:
# auth errors look like this
# 503 Getting metadata from plugin failed with error: Reauthentication is needed. Please run `gcloud auth application-default login` to reauthenticate.
exception_mapping_worked = True
raise BadRequestError(
message="GeminiException - Invalid api key",
model=model,
llm_provider="palm",
response=getattr(original_exception, "response", None),
)
if (
"504 Deadline expired before operation could complete." in error_str
or "504 Deadline Exceeded" in error_str
):
exception_mapping_worked = True
raise Timeout(
message=f"GeminiException - {original_exception.message}",
model=model,
llm_provider="palm",
exception_status_code=original_exception.status_code,
)
if "400 Request payload size exceeds" in error_str:
exception_mapping_worked = True
raise ContextWindowExceededError(
message=f"GeminiException - {error_str}",
model=model,
llm_provider="palm",
response=getattr(original_exception, "response", None),
)
if (
"500 An internal error has occurred." in error_str
or "list index out of range" in error_str
):
exception_mapping_worked = True
raise APIError(
status_code=getattr(original_exception, "status_code", 500),
message=f"GeminiException - {original_exception.message}",
llm_provider="palm",
model=model,
request=httpx.Response(
status_code=429,
request=httpx.Request(
method="POST",
url=" https://cloud.google.com/vertex-ai/",
),
),
)
if hasattr(original_exception, "status_code"):
if original_exception.status_code == 400:
if original_exception.status_code == 503:
exception_mapping_worked = True
raise BadRequestError(
message=f"GeminiException - {error_str}",
raise ServiceUnavailableError(
message=f"{custom_llm_provider.capitalize()}Exception - {error_str}",
llm_provider=custom_llm_provider,
model=model,
llm_provider="palm",
response=getattr(original_exception, "response", None),
)
# Dailed: Error occurred: 400 Request payload size exceeds the limit: 20000 bytes
elif custom_llm_provider == "cloudflare":
if "Authentication error" in error_str:
exception_mapping_worked = True
@@ -1449,6 +1413,14 @@ def exception_type( # type: ignore # noqa: PLR0915
model=model,
response=getattr(original_exception, "response", None),
)
elif "invalid type: parameter" in error_str:
exception_mapping_worked = True
raise BadRequestError(
message=f"CohereException - {original_exception.message}",
llm_provider="cohere",
model=model,
response=getattr(original_exception, "response", None),
)
elif "too many tokens" in error_str:
exception_mapping_worked = True
raise ContextWindowExceededError(
litellm/litellm_core_utils/get_supported_openai_params.py
+1 -3
View File
@@ -94,9 +94,7 @@ def get_supported_openai_params( # noqa: PLR0915
return litellm.VLLMConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "deepseek":
return litellm.DeepSeekChatConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "cohere":
return litellm.CohereConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "cohere_chat":
elif custom_llm_provider == "cohere_chat" or custom_llm_provider == "cohere":
return litellm.CohereChatConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "maritalk":
return litellm.MaritalkConfig().get_supported_openai_params(model=model)
litellm/litellm_core_utils/litellm_logging.py
+223 -181
View File
@@ -138,6 +138,7 @@ from ..integrations.logfire_logger import LogfireLevel, LogfireLogger
from ..integrations.lunary import LunaryLogger
from ..integrations.openmeter import OpenMeterLogger
from ..integrations.opik.opik import OpikLogger
from ..integrations.posthog import PostHogLogger
from ..integrations.prompt_layer import PromptLayerLogger
from ..integrations.s3 import S3Logger
from ..integrations.s3_v2 import S3Logger as S3V2Logger
@@ -193,7 +194,6 @@ _in_memory_loggers: List[Any] = []
sentry_sdk_instance = None
capture_exception = None
add_breadcrumb = None
posthog = None
slack_app = None
alerts_channel = None
heliconeLogger = None
@@ -300,9 +300,9 @@ class Logging(LiteLLMLoggingBaseClass):
self.litellm_trace_id: str = litellm_trace_id or str(uuid.uuid4())
self.function_id = function_id
self.streaming_chunks: List[Any] = [] # for generating complete stream response
self.sync_streaming_chunks: List[Any] = (
[]
) # for generating complete stream response
self.sync_streaming_chunks: List[
Any
] = [] # for generating complete stream response
self.log_raw_request_response = log_raw_request_response
# Initialize dynamic callbacks
@@ -672,24 +672,23 @@ class Logging(LiteLLMLoggingBaseClass):
if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook(
non_default_params
):
self.model_call_details["prompt_integration"] = (
anthropic_cache_control_logger.__class__.__name__
)
self.model_call_details[
"prompt_integration"
] = anthropic_cache_control_logger.__class__.__name__
return anthropic_cache_control_logger
#########################################################
# Vector Store / Knowledge Base hooks
#########################################################
if litellm.vector_store_registry is not None:
vector_store_custom_logger = _init_custom_logger_compatible_class(
logging_integration="vector_store_pre_call_hook",
internal_usage_cache=None,
llm_router=None,
)
self.model_call_details["prompt_integration"] = (
vector_store_custom_logger.__class__.__name__
)
self.model_call_details[
"prompt_integration"
] = vector_store_custom_logger.__class__.__name__
return vector_store_custom_logger
return None
@@ -741,9 +740,9 @@ class Logging(LiteLLMLoggingBaseClass):
model
): # if model name was changes pre-call, overwrite the initial model call name with the new one
self.model_call_details["model"] = model
self.model_call_details["litellm_params"]["api_base"] = (
self._get_masked_api_base(additional_args.get("api_base", ""))
)
self.model_call_details["litellm_params"][
"api_base"
] = self._get_masked_api_base(additional_args.get("api_base", ""))
def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR0915
# Log the exact input to the LLM API
@@ -772,10 +771,10 @@ class Logging(LiteLLMLoggingBaseClass):
try:
# [Non-blocking Extra Debug Information in metadata]
if turn_off_message_logging is True:
_metadata["raw_request"] = (
"redacted by litellm. \
_metadata[
"raw_request"
] = "redacted by litellm. \
'litellm.turn_off_message_logging=True'"
)
else:
curl_command = self._get_request_curl_command(
api_base=additional_args.get("api_base", ""),
@@ -786,32 +785,32 @@ class Logging(LiteLLMLoggingBaseClass):
_metadata["raw_request"] = str(curl_command)
# split up, so it's easier to parse in the UI
self.model_call_details["raw_request_typed_dict"] = (
RawRequestTypedDict(
raw_request_api_base=str(
additional_args.get("api_base") or ""
),
raw_request_body=self._get_raw_request_body(
additional_args.get("complete_input_dict", {})
),
raw_request_headers=self._get_masked_headers(
additional_args.get("headers", {}) or {},
ignore_sensitive_headers=True,
),
error=None,
)
self.model_call_details[
"raw_request_typed_dict"
] = RawRequestTypedDict(
raw_request_api_base=str(
additional_args.get("api_base") or ""
),
raw_request_body=self._get_raw_request_body(
additional_args.get("complete_input_dict", {})
),
raw_request_headers=self._get_masked_headers(
additional_args.get("headers", {}) or {},
ignore_sensitive_headers=True,
),
error=None,
)
except Exception as e:
self.model_call_details["raw_request_typed_dict"] = (
RawRequestTypedDict(
error=str(e),
)
self.model_call_details[
"raw_request_typed_dict"
] = RawRequestTypedDict(
error=str(e),
)
_metadata["raw_request"] = (
"Unable to Log \
_metadata[
"raw_request"
] = "Unable to Log \
raw request: {}".format(
str(e)
)
str(e)
)
if getattr(self, "logger_fn", None) and callable(self.logger_fn):
try:
@@ -1112,13 +1111,13 @@ class Logging(LiteLLMLoggingBaseClass):
for callback in callbacks:
try:
if isinstance(callback, CustomLogger):
response: Optional[MCPPostCallResponseObject] = (
await callback.async_post_mcp_tool_call_hook(
kwargs=kwargs,
response_obj=post_mcp_tool_call_response_obj,
start_time=start_time,
end_time=end_time,
)
response: Optional[
MCPPostCallResponseObject
] = await callback.async_post_mcp_tool_call_hook(
kwargs=kwargs,
response_obj=post_mcp_tool_call_response_obj,
start_time=start_time,
end_time=end_time,
)
######################################################################
# if any of the callbacks modify the response, use the modified response
@@ -1238,9 +1237,9 @@ class Logging(LiteLLMLoggingBaseClass):
verbose_logger.debug(
f"response_cost_failure_debug_information: {debug_info}"
)
self.model_call_details["response_cost_failure_debug_information"] = (
debug_info
)
self.model_call_details[
"response_cost_failure_debug_information"
] = debug_info
return None
try:
@@ -1265,9 +1264,9 @@ class Logging(LiteLLMLoggingBaseClass):
verbose_logger.debug(
f"response_cost_failure_debug_information: {debug_info}"
)
self.model_call_details["response_cost_failure_debug_information"] = (
debug_info
)
self.model_call_details[
"response_cost_failure_debug_information"
] = debug_info
return None
@@ -1411,9 +1410,9 @@ class Logging(LiteLLMLoggingBaseClass):
end_time = datetime.datetime.now()
if self.completion_start_time is None:
self.completion_start_time = end_time
self.model_call_details["completion_start_time"] = (
self.completion_start_time
)
self.model_call_details[
"completion_start_time"
] = self.completion_start_time
self.model_call_details["log_event_type"] = "successful_api_call"
self.model_call_details["end_time"] = end_time
self.model_call_details["cache_hit"] = cache_hit
@@ -1466,39 +1465,39 @@ class Logging(LiteLLMLoggingBaseClass):
"response_cost"
]
else:
self.model_call_details["response_cost"] = (
self._response_cost_calculator(result=logging_result)
)
self.model_call_details[
"response_cost"
] = self._response_cost_calculator(result=logging_result)
## STANDARDIZED LOGGING PAYLOAD
self.model_call_details["standard_logging_object"] = (
get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=logging_result,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
self.model_call_details[
"standard_logging_object"
] = get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=logging_result,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
elif isinstance(result, dict) or isinstance(result, list):
## STANDARDIZED LOGGING PAYLOAD
self.model_call_details["standard_logging_object"] = (
get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=result,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
self.model_call_details[
"standard_logging_object"
] = get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=result,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
elif standard_logging_object is not None:
self.model_call_details["standard_logging_object"] = (
standard_logging_object
)
self.model_call_details[
"standard_logging_object"
] = standard_logging_object
else: # streaming chunks + image gen.
self.model_call_details["response_cost"] = None
@@ -1597,7 +1596,6 @@ class Logging(LiteLLMLoggingBaseClass):
)
if complete_streaming_response is not None:
self.success_handler(result=complete_streaming_response)
return
@@ -1650,23 +1648,23 @@ class Logging(LiteLLMLoggingBaseClass):
verbose_logger.debug(
"Logging Details LiteLLM-Success Call streaming complete"
)
self.model_call_details["complete_streaming_response"] = (
complete_streaming_response
)
self.model_call_details["response_cost"] = (
self._response_cost_calculator(result=complete_streaming_response)
)
self.model_call_details[
"complete_streaming_response"
] = complete_streaming_response
self.model_call_details[
"response_cost"
] = self._response_cost_calculator(result=complete_streaming_response)
## STANDARDIZED LOGGING PAYLOAD
self.model_call_details["standard_logging_object"] = (
get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=complete_streaming_response,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
self.model_call_details[
"standard_logging_object"
] = get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=complete_streaming_response,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
callbacks = self.get_combined_callback_list(
dynamic_success_callbacks=self.dynamic_success_callbacks,
@@ -1994,10 +1992,10 @@ class Logging(LiteLLMLoggingBaseClass):
)
else:
if self.stream and complete_streaming_response:
self.model_call_details["complete_response"] = (
self.model_call_details.get(
"complete_streaming_response", {}
)
self.model_call_details[
"complete_response"
] = self.model_call_details.get(
"complete_streaming_response", {}
)
result = self.model_call_details["complete_response"]
openMeterLogger.log_success_event(
@@ -2036,10 +2034,10 @@ class Logging(LiteLLMLoggingBaseClass):
)
else:
if self.stream and complete_streaming_response:
self.model_call_details["complete_response"] = (
self.model_call_details.get(
"complete_streaming_response", {}
)
self.model_call_details[
"complete_response"
] = self.model_call_details.get(
"complete_streaming_response", {}
)
result = self.model_call_details["complete_response"]
@@ -2141,10 +2139,12 @@ class Logging(LiteLLMLoggingBaseClass):
result.usage = batch_usage
elif not is_base64_unified_file_id: # only run for non-unified file ids
response_cost, batch_usage, batch_models = (
await _handle_completed_batch(
batch=result, custom_llm_provider=self.custom_llm_provider
)
(
response_cost,
batch_usage,
batch_models,
) = await _handle_completed_batch(
batch=result, custom_llm_provider=self.custom_llm_provider
)
result._hidden_params["response_cost"] = response_cost
@@ -2175,9 +2175,9 @@ class Logging(LiteLLMLoggingBaseClass):
if complete_streaming_response is not None:
print_verbose("Async success callbacks: Got a complete streaming response")
self.model_call_details["async_complete_streaming_response"] = (
complete_streaming_response
)
self.model_call_details[
"async_complete_streaming_response"
] = complete_streaming_response
try:
if self.model_call_details.get("cache_hit", False) is True:
@@ -2188,10 +2188,10 @@ class Logging(LiteLLMLoggingBaseClass):
model_call_details=self.model_call_details
)
# base_model defaults to None if not set on model_info
self.model_call_details["response_cost"] = (
self._response_cost_calculator(
result=complete_streaming_response
)
self.model_call_details[
"response_cost"
] = self._response_cost_calculator(
result=complete_streaming_response
)
verbose_logger.debug(
@@ -2204,16 +2204,16 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["response_cost"] = None
## STANDARDIZED LOGGING PAYLOAD
self.model_call_details["standard_logging_object"] = (
get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=complete_streaming_response,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
self.model_call_details[
"standard_logging_object"
] = get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj=complete_streaming_response,
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="success",
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
callbacks = self.get_combined_callback_list(
dynamic_success_callbacks=self.dynamic_async_success_callbacks,
@@ -2426,18 +2426,18 @@ class Logging(LiteLLMLoggingBaseClass):
## STANDARDIZED LOGGING PAYLOAD
self.model_call_details["standard_logging_object"] = (
get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj={},
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="failure",
error_str=str(exception),
original_exception=exception,
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
self.model_call_details[
"standard_logging_object"
] = get_standard_logging_object_payload(
kwargs=self.model_call_details,
init_response_obj={},
start_time=start_time,
end_time=end_time,
logging_obj=self,
status="failure",
error_str=str(exception),
original_exception=exception,
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
return start_time, end_time
@@ -3068,7 +3068,7 @@ def set_callbacks(callback_list, function_id=None): # noqa: PLR0915
"""
Globally sets the callback client
"""
global sentry_sdk_instance, capture_exception, add_breadcrumb, posthog, slack_app, alerts_channel, traceloopLogger, athinaLogger, heliconeLogger, supabaseClient, lunaryLogger, promptLayerLogger, langFuseLogger, customLogger, weightsBiasesLogger, logfireLogger, dynamoLogger, s3Logger, dataDogLogger, prometheusLogger, greenscaleLogger, openMeterLogger, deepevalLogger
global sentry_sdk_instance, capture_exception, add_breadcrumb, slack_app, alerts_channel, traceloopLogger, athinaLogger, heliconeLogger, supabaseClient, lunaryLogger, promptLayerLogger, langFuseLogger, customLogger, weightsBiasesLogger, logfireLogger, dynamoLogger, s3Logger, dataDogLogger, prometheusLogger, greenscaleLogger, openMeterLogger, deepevalLogger
try:
for callback in callback_list:
@@ -3107,19 +3107,6 @@ def set_callbacks(callback_list, function_id=None): # noqa: PLR0915
)
capture_exception = sentry_sdk_instance.capture_exception
add_breadcrumb = sentry_sdk_instance.add_breadcrumb
elif callback == "posthog":
try:
from posthog import Posthog
except ImportError:
print_verbose("Package 'posthog' is missing. Installing it...")
subprocess.check_call(
[sys.executable, "-m", "pip", "install", "posthog"]
)
from posthog import Posthog
posthog = Posthog(
project_api_key=os.environ.get("POSTHOG_API_KEY"),
host=os.environ.get("POSTHOG_API_URL"),
)
elif callback == "slack":
try:
from slack_bolt import App
@@ -3214,6 +3201,14 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_openmeter_logger = OpenMeterLogger()
_in_memory_loggers.append(_openmeter_logger)
return _openmeter_logger # type: ignore
elif logging_integration == "posthog":
for callback in _in_memory_loggers:
if isinstance(callback, PostHogLogger):
return callback # type: ignore
_posthog_logger = PostHogLogger()
_in_memory_loggers.append(_posthog_logger)
return _posthog_logger # type: ignore
elif logging_integration == "braintrust":
from litellm.integrations.braintrust_logging import BraintrustLogger
@@ -3326,9 +3321,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
endpoint=arize_config.endpoint,
)
os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
f"space_id={arize_config.space_key},api_key={arize_config.api_key}"
)
os.environ[
"OTEL_EXPORTER_OTLP_TRACES_HEADERS"
] = f"space_id={arize_config.space_key},api_key={arize_config.api_key}"
for callback in _in_memory_loggers:
if (
isinstance(callback, ArizeLogger)
@@ -3352,9 +3347,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
# auth can be disabled on local deployments of arize phoenix
if arize_phoenix_config.otlp_auth_headers is not None:
os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
arize_phoenix_config.otlp_auth_headers
)
os.environ[
"OTEL_EXPORTER_OTLP_TRACES_HEADERS"
] = arize_phoenix_config.otlp_auth_headers
for callback in _in_memory_loggers:
if (
@@ -3449,6 +3444,30 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
dynamic_rate_limiter_obj.update_variables(llm_router=llm_router)
_in_memory_loggers.append(dynamic_rate_limiter_obj)
return dynamic_rate_limiter_obj # type: ignore
elif logging_integration == "dynamic_rate_limiter_v3":
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
_PROXY_DynamicRateLimitHandlerV3,
)
for callback in _in_memory_loggers:
if isinstance(callback, _PROXY_DynamicRateLimitHandlerV3):
return callback # type: ignore
if internal_usage_cache is None:
raise Exception(
"Internal Error: Cache cannot be empty - internal_usage_cache={}".format(
internal_usage_cache
)
)
dynamic_rate_limiter_obj_v3 = _PROXY_DynamicRateLimitHandlerV3(
internal_usage_cache=internal_usage_cache
)
if llm_router is not None and isinstance(llm_router, litellm.Router):
dynamic_rate_limiter_obj_v3.update_variables(llm_router=llm_router)
_in_memory_loggers.append(dynamic_rate_limiter_obj_v3)
return dynamic_rate_limiter_obj_v3 # type: ignore
elif logging_integration == "langtrace":
if "LANGTRACE_API_KEY" not in os.environ:
raise ValueError("LANGTRACE_API_KEY not found in environment variables")
@@ -3462,9 +3481,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
exporter="otlp_http",
endpoint="https://langtrace.ai/api/trace",
)
os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
f"api_key={os.getenv('LANGTRACE_API_KEY')}"
)
os.environ[
"OTEL_EXPORTER_OTLP_TRACES_HEADERS"
] = f"api_key={os.getenv('LANGTRACE_API_KEY')}"
for callback in _in_memory_loggers:
if (
isinstance(callback, OpenTelemetry)
@@ -3712,6 +3731,14 @@ def get_custom_logger_compatible_class( # noqa: PLR0915
for callback in _in_memory_loggers:
if isinstance(callback, _PROXY_DynamicRateLimitHandler):
return callback # type: ignore
elif logging_integration == "dynamic_rate_limiter_v3":
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
_PROXY_DynamicRateLimitHandlerV3,
)
for callback in _in_memory_loggers:
if isinstance(callback, _PROXY_DynamicRateLimitHandlerV3):
return callback # type: ignore
elif logging_integration == "langtrace":
from litellm.integrations.opentelemetry import OpenTelemetry
@@ -3910,22 +3937,25 @@ class StandardLoggingPayloadSetup:
clean_metadata = StandardLoggingMetadata(
user_api_key_hash=None,
user_api_key_alias=None,
user_api_key_spend=None,
user_api_key_max_budget=None,
user_api_key_budget_reset_at=None,
user_api_key_team_id=None,
user_api_key_org_id=None,
user_api_key_user_id=None,
user_api_key_team_alias=None,
user_api_key_user_email=None,
user_api_key_end_user_id=None,
user_api_key_request_route=None,
spend_logs_metadata=None,
requester_ip_address=None,
requester_metadata=None,
user_api_key_end_user_id=None,
prompt_management_metadata=prompt_management_metadata,
applied_guardrails=applied_guardrails,
mcp_tool_call_metadata=mcp_tool_call_metadata,
vector_store_request_metadata=vector_store_request_metadata,
usage_object=usage_object,
requester_custom_headers=None,
user_api_key_request_route=None,
cold_storage_object_key=None,
)
if isinstance(metadata, dict):
@@ -4114,10 +4144,10 @@ class StandardLoggingPayloadSetup:
for key in StandardLoggingHiddenParams.__annotations__.keys():
if key in hidden_params:
if key == "additional_headers":
clean_hidden_params["additional_headers"] = (
StandardLoggingPayloadSetup.get_additional_headers(
hidden_params[key]
)
clean_hidden_params[
"additional_headers"
] = StandardLoggingPayloadSetup.get_additional_headers(
hidden_params[key]
)
else:
clean_hidden_params[key] = hidden_params[key] # type: ignore
@@ -4150,15 +4180,28 @@ class StandardLoggingPayloadSetup:
from litellm.integrations.s3 import get_s3_object_key
# Only generate object key if cold storage is configured
if litellm.configured_cold_storage_logger is None:
configured_cold_storage_logger = litellm.configured_cold_storage_logger
if configured_cold_storage_logger is None:
return None
try:
# Generate file name in same format as litellm.utils.get_logging_id
s3_file_name = f"time-{start_time.strftime('%H-%M-%S-%f')}_{response_id}"
# Get the actual s3_path from the configured cold storage logger instance
s3_path = "" # default value
# Try to get the actual logger instance from the logger name
try:
custom_logger = litellm.logging_callback_manager.get_active_custom_logger_for_callback_name(configured_cold_storage_logger)
if custom_logger and hasattr(custom_logger, 's3_path') and custom_logger.s3_path:
s3_path = custom_logger.s3_path
except Exception:
# If any error occurs in getting the logger instance, use default empty s3_path
pass
s3_object_key = get_s3_object_key(
s3_path="", # Use empty path as default
s3_path=s3_path, # Use actual s3_path from logger configuration
team_alias_prefix="", # Don't split by team alias for cold storage
start_time=start_time,
s3_file_name=s3_file_name,
@@ -4553,6 +4596,9 @@ def get_standard_logging_metadata(
clean_metadata = StandardLoggingMetadata(
user_api_key_hash=None,
user_api_key_alias=None,
user_api_key_spend=None,
user_api_key_max_budget=None,
user_api_key_budget_reset_at=None,
user_api_key_team_id=None,
user_api_key_org_id=None,
user_api_key_user_id=None,
@@ -4572,14 +4618,10 @@ def get_standard_logging_metadata(
cold_storage_object_key=None,
)
if isinstance(metadata, dict):
# Filter the metadata dictionary to include only the specified keys
clean_metadata = StandardLoggingMetadata(
**{ # type: ignore
key: metadata[key]
for key in StandardLoggingMetadata.__annotations__.keys()
if key in metadata
}
)
# Update the clean_metadata with values from input metadata that match StandardLoggingMetadata fields
for key in StandardLoggingMetadata.__annotations__.keys():
if key in metadata:
clean_metadata[key] = metadata[key] # type: ignore
if metadata.get("user_api_key") is not None:
if is_valid_sha256_hash(str(metadata.get("user_api_key"))):
@@ -4602,9 +4644,9 @@ def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]):
):
for k, v in metadata["user_api_key_metadata"].items():
if k == "logging": # prevent logging user logging keys
cleaned_user_api_key_metadata[k] = (
"scrubbed_by_litellm_for_sensitive_keys"
)
cleaned_user_api_key_metadata[
k
] = "scrubbed_by_litellm_for_sensitive_keys"
else:
cleaned_user_api_key_metadata[k] = v
litellm/litellm_core_utils/llm_cost_calc/utils.py
+301 -123
View File
@@ -1,11 +1,12 @@
# What is this?
## Helper utilities for cost_per_token()
from typing import Any, Literal, Optional, Tuple, cast
from typing import Any, Literal, Optional, Tuple, TypedDict, cast
import litellm
from litellm._logging import verbose_logger
from litellm.types.utils import (
CacheCreationTokenDetails,
CallTypes,
ImageResponse,
ModelInfo,
@@ -113,20 +114,34 @@ def _generic_cost_per_character(
return prompt_cost, completion_cost
def _get_token_base_cost(model_info: ModelInfo, usage: Usage) -> Tuple[float, float, float, float]:
def _get_token_base_cost(
model_info: ModelInfo, usage: Usage
) -> Tuple[float, float, float, float, float]:
"""
Return prompt cost, completion cost, and cache costs for a given model and usage.
If input_tokens > threshold and `input_cost_per_token_above_[x]k_tokens` or `input_cost_per_token_above_[x]_tokens` is set,
then we use the corresponding threshold cost for all token types.
Returns:
Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost)
"""
prompt_base_cost = cast(float, _get_cost_per_unit(model_info, "input_cost_per_token"))
completion_base_cost = cast(float, _get_cost_per_unit(model_info, "output_cost_per_token"))
cache_creation_cost = cast(float, _get_cost_per_unit(model_info, "cache_creation_input_token_cost"))
cache_read_cost = cast(float, _get_cost_per_unit(model_info, "cache_read_input_token_cost"))
prompt_base_cost = cast(
float, _get_cost_per_unit(model_info, "input_cost_per_token")
)
completion_base_cost = cast(
float, _get_cost_per_unit(model_info, "output_cost_per_token")
)
cache_creation_cost = cast(
float, _get_cost_per_unit(model_info, "cache_creation_input_token_cost")
)
cache_creation_cost_above_1hr = cast(
float,
_get_cost_per_unit(model_info, "cache_creation_input_token_cost_above_1hr"),
)
cache_read_cost = cast(
float, _get_cost_per_unit(model_info, "cache_read_input_token_cost")
)
## CHECK IF ABOVE THRESHOLD
threshold: Optional[float] = None
@@ -140,34 +155,57 @@ def _get_token_base_cost(model_info: ModelInfo, usage: Usage) -> Tuple[float, fl
)
if usage.prompt_tokens > threshold:
prompt_base_cost = cast(float, _get_cost_per_unit(model_info, key, prompt_base_cost))
completion_base_cost = cast(float, _get_cost_per_unit(
model_info,
f"output_cost_per_token_above_{threshold_str}_tokens",
completion_base_cost,
))
prompt_base_cost = cast(
float, _get_cost_per_unit(model_info, key, prompt_base_cost)
)
completion_base_cost = cast(
float,
_get_cost_per_unit(
model_info,
f"output_cost_per_token_above_{threshold_str}_tokens",
completion_base_cost,
),
)
# Apply tiered pricing to cache costs
cache_creation_tiered_key = f"cache_creation_input_token_cost_above_{threshold_str}_tokens"
cache_read_tiered_key = f"cache_read_input_token_cost_above_{threshold_str}_tokens"
cache_creation_tiered_key = (
f"cache_creation_input_token_cost_above_{threshold_str}_tokens"
)
cache_read_tiered_key = (
f"cache_read_input_token_cost_above_{threshold_str}_tokens"
)
if cache_creation_tiered_key in model_info:
cache_creation_cost = cast(float, _get_cost_per_unit(
model_info, cache_creation_tiered_key, cache_creation_cost
))
cache_creation_cost = cast(
float,
_get_cost_per_unit(
model_info,
cache_creation_tiered_key,
cache_creation_cost,
),
)
if cache_read_tiered_key in model_info:
cache_read_cost = cast(float, _get_cost_per_unit(
model_info, cache_read_tiered_key, cache_read_cost
))
cache_read_cost = cast(
float,
_get_cost_per_unit(
model_info, cache_read_tiered_key, cache_read_cost
),
)
break
except (IndexError, ValueError):
continue
except Exception:
continue
return prompt_base_cost, completion_base_cost, cache_creation_cost, cache_read_cost
return (
prompt_base_cost,
completion_base_cost,
cache_creation_cost,
cache_creation_cost_above_1hr,
cache_read_cost,
)
def calculate_cost_component(
@@ -195,7 +233,9 @@ def calculate_cost_component(
return 0.0
def _get_cost_per_unit(model_info: ModelInfo, cost_key: str, default_value: Optional[float] = 0.0) -> Optional[float]:
def _get_cost_per_unit(
model_info: ModelInfo, cost_key: str, default_value: Optional[float] = 0.0
) -> Optional[float]:
# Sometimes the cost per unit is a string (e.g.: If a value like "3e-7" was read from the config.yaml)
cost_per_unit = model_info.get(cost_key)
if isinstance(cost_per_unit, float):
@@ -210,7 +250,196 @@ def _get_cost_per_unit(model_info: ModelInfo, cost_key: str, default_value: Opti
f"litellm.litellm_core_utils.llm_cost_calc.utils.py::calculate_cost_per_component(): Exception occured - {cost_per_unit}\nDefaulting to 0.0"
)
return default_value
def calculate_cache_writing_cost(
cache_creation_tokens: int,
cache_creation_token_details: Optional[CacheCreationTokenDetails],
cache_creation_cost_above_1hr: float,
cache_creation_cost: float,
) -> float:
"""
Adjust cost of cache creation tokens based on the cache creation token details.
"""
total_cost: float = 0.0
if cache_creation_token_details is not None:
# get the number of 5m and 1h cache creation tokens
cache_creation_tokens_5m = (
cache_creation_token_details.ephemeral_5m_input_tokens
)
cache_creation_tokens_1h = (
cache_creation_token_details.ephemeral_1h_input_tokens
)
# add the number of 5m and 1h cache creation tokens to the cache creation tokens
total_cost += (
cache_creation_tokens_5m * cache_creation_cost
if cache_creation_tokens_5m is not None
else 0.0
)
total_cost += (
cache_creation_tokens_1h * cache_creation_cost_above_1hr
if cache_creation_tokens_1h is not None
else 0.0
)
else:
total_cost += cache_creation_tokens * cache_creation_cost
return total_cost
class PromptTokensDetailsResult(TypedDict):
cache_hit_tokens: int
cache_creation_tokens: int
cache_creation_token_details: Optional[CacheCreationTokenDetails]
text_tokens: int
audio_tokens: int
character_count: int
image_count: int
video_length_seconds: int
def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
cache_hit_tokens = (
cast(Optional[int], getattr(usage.prompt_tokens_details, "cached_tokens", 0))
or 0
)
cache_creation_tokens = (
cast(
Optional[int],
getattr(usage.prompt_tokens_details, "cache_creation_tokens", 0),
)
or 0
)
cache_creation_token_details = (
cast(
Optional[CacheCreationTokenDetails],
getattr(usage.prompt_tokens_details, "cache_creation_token_details", None),
)
or None
)
text_tokens = (
cast(Optional[int], getattr(usage.prompt_tokens_details, "text_tokens", None))
or 0 # default to prompt tokens, if this field is not set
)
audio_tokens = (
cast(Optional[int], getattr(usage.prompt_tokens_details, "audio_tokens", 0))
or 0
)
character_count = (
cast(
Optional[int],
getattr(usage.prompt_tokens_details, "character_count", 0),
)
or 0
)
image_count = (
cast(Optional[int], getattr(usage.prompt_tokens_details, "image_count", 0)) or 0
)
video_length_seconds = (
cast(
Optional[int],
getattr(usage.prompt_tokens_details, "video_length_seconds", 0),
)
or 0
)
return PromptTokensDetailsResult(
cache_hit_tokens=cache_hit_tokens,
cache_creation_tokens=cache_creation_tokens,
cache_creation_token_details=cache_creation_token_details,
text_tokens=text_tokens,
audio_tokens=audio_tokens,
character_count=character_count,
image_count=image_count,
video_length_seconds=video_length_seconds,
)
class CompletionTokensDetailsResult(TypedDict):
audio_tokens: int
text_tokens: int
reasoning_tokens: int
def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult:
audio_tokens = (
cast(
Optional[int],
getattr(usage.completion_tokens_details, "audio_tokens", 0),
)
or 0
)
text_tokens = (
cast(
Optional[int],
getattr(usage.completion_tokens_details, "text_tokens", None),
)
or 0 # default to completion tokens, if this field is not set
)
reasoning_tokens = (
cast(
Optional[int],
getattr(usage.completion_tokens_details, "reasoning_tokens", 0),
)
or 0
)
return CompletionTokensDetailsResult(
audio_tokens=audio_tokens,
text_tokens=text_tokens,
reasoning_tokens=reasoning_tokens,
)
def _calculate_input_cost(
prompt_tokens_details: PromptTokensDetailsResult,
model_info: ModelInfo,
prompt_base_cost: float,
cache_read_cost: float,
cache_creation_cost: float,
cache_creation_cost_above_1hr: float,
) -> float:
"""
Calculates the input cost for a given model, prompt tokens, and completion tokens.
"""
prompt_cost = float(prompt_tokens_details["text_tokens"]) * prompt_base_cost
### CACHE READ COST - Now uses tiered pricing
prompt_cost += float(prompt_tokens_details["cache_hit_tokens"]) * cache_read_cost
### AUDIO COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_audio_token", prompt_tokens_details["audio_tokens"]
)
### CACHE WRITING COST - Now uses tiered pricing
prompt_cost += calculate_cache_writing_cost(
cache_creation_tokens=prompt_tokens_details["cache_creation_tokens"],
cache_creation_token_details=prompt_tokens_details[
"cache_creation_token_details"
],
cache_creation_cost_above_1hr=cache_creation_cost_above_1hr,
cache_creation_cost=cache_creation_cost,
)
### CHARACTER COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_character", prompt_tokens_details["character_count"]
)
### IMAGE COUNT COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_image", prompt_tokens_details["image_count"]
)
### VIDEO LENGTH COST
prompt_cost += calculate_cost_component(
model_info,
"input_cost_per_video_per_second",
prompt_tokens_details["video_length_seconds"],
)
return prompt_cost
def generic_cost_per_token(
@@ -236,83 +465,45 @@ def generic_cost_per_token(
### Cost of processing (non-cache hit + cache hit) + Cost of cache-writing (cache writing)
prompt_cost = 0.0
### PROCESSING COST
text_tokens = usage.prompt_tokens
cache_hit_tokens = 0
audio_tokens = 0
character_count = 0
image_count = 0
video_length_seconds = 0
prompt_tokens_details = PromptTokensDetailsResult(
cache_hit_tokens=0,
cache_creation_tokens=0,
cache_creation_token_details=None,
text_tokens=usage.prompt_tokens,
audio_tokens=0,
character_count=0,
image_count=0,
video_length_seconds=0,
)
if usage.prompt_tokens_details:
cache_hit_tokens = (
cast(
Optional[int], getattr(usage.prompt_tokens_details, "cached_tokens", 0)
)
or 0
)
text_tokens = (
cast(
Optional[int], getattr(usage.prompt_tokens_details, "text_tokens", None)
)
or 0 # default to prompt tokens, if this field is not set
)
audio_tokens = (
cast(Optional[int], getattr(usage.prompt_tokens_details, "audio_tokens", 0))
or 0
)
character_count = (
cast(
Optional[int],
getattr(usage.prompt_tokens_details, "character_count", 0),
)
or 0
)
image_count = (
cast(Optional[int], getattr(usage.prompt_tokens_details, "image_count", 0))
or 0
)
video_length_seconds = (
cast(
Optional[int],
getattr(usage.prompt_tokens_details, "video_length_seconds", 0),
)
or 0
)
prompt_tokens_details = _parse_prompt_tokens_details(usage)
## EDGE CASE - text tokens not set inside PromptTokensDetails
if text_tokens == 0:
text_tokens = usage.prompt_tokens - cache_hit_tokens - audio_tokens
prompt_base_cost, completion_base_cost, cache_creation_cost, cache_read_cost = _get_token_base_cost(
model_info=model_info, usage=usage
)
if prompt_tokens_details["text_tokens"] == 0:
text_tokens = (
usage.prompt_tokens
- prompt_tokens_details["cache_hit_tokens"]
- prompt_tokens_details["audio_tokens"]
- prompt_tokens_details["cache_creation_tokens"]
)
prompt_tokens_details["text_tokens"] = text_tokens
prompt_cost = float(text_tokens) * prompt_base_cost
(
prompt_base_cost,
completion_base_cost,
cache_creation_cost,
cache_creation_cost_above_1hr,
cache_read_cost,
) = _get_token_base_cost(model_info=model_info, usage=usage)
### CACHE READ COST - Now uses tiered pricing
prompt_cost += float(cache_hit_tokens) * cache_read_cost
### AUDIO COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_audio_token", audio_tokens
)
### CACHE WRITING COST - Now uses tiered pricing
prompt_cost += float(usage._cache_creation_input_tokens or 0) * cache_creation_cost
### CHARACTER COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_character", character_count
)
### IMAGE COUNT COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_image", image_count
)
### VIDEO LENGTH COST
prompt_cost += calculate_cost_component(
model_info, "input_cost_per_video_per_second", video_length_seconds
prompt_cost = _calculate_input_cost(
prompt_tokens_details=prompt_tokens_details,
model_info=model_info,
prompt_base_cost=prompt_base_cost,
cache_read_cost=cache_read_cost,
cache_creation_cost=cache_creation_cost,
cache_creation_cost_above_1hr=cache_creation_cost_above_1hr,
)
## CALCULATE OUTPUT COST
@@ -321,27 +512,10 @@ def generic_cost_per_token(
reasoning_tokens = 0
is_text_tokens_total = False
if usage.completion_tokens_details is not None:
audio_tokens = (
cast(
Optional[int],
getattr(usage.completion_tokens_details, "audio_tokens", 0),
)
or 0
)
text_tokens = (
cast(
Optional[int],
getattr(usage.completion_tokens_details, "text_tokens", None),
)
or 0 # default to completion tokens, if this field is not set
)
reasoning_tokens = (
cast(
Optional[int],
getattr(usage.completion_tokens_details, "reasoning_tokens", 0),
)
or 0
)
completion_tokens_details = _parse_completion_tokens_details(usage)
audio_tokens = completion_tokens_details["audio_tokens"]
text_tokens = completion_tokens_details["text_tokens"]
reasoning_tokens = completion_tokens_details["reasoning_tokens"]
if text_tokens == 0:
text_tokens = usage.completion_tokens
@@ -350,8 +524,12 @@ def generic_cost_per_token(
## TEXT COST
completion_cost = float(text_tokens) * completion_base_cost
_output_cost_per_audio_token = _get_cost_per_unit(model_info, "output_cost_per_audio_token", None)
_output_cost_per_reasoning_token = _get_cost_per_unit(model_info, "output_cost_per_reasoning_token", None)
_output_cost_per_audio_token = _get_cost_per_unit(
model_info, "output_cost_per_audio_token", None
)
_output_cost_per_reasoning_token = _get_cost_per_unit(
model_info, "output_cost_per_reasoning_token", None
)
## AUDIO COST
if not is_text_tokens_total and audio_tokens is not None and audio_tokens > 0:
@@ -397,7 +575,7 @@ class CostCalculatorUtils:
]:
return True
return False
@staticmethod
def route_image_generation_cost_calculator(
model: str,
litellm/litellm_core_utils/object_pooling.py
+137
View File
@@ -0,0 +1,137 @@
"""
Generic object pooling utilities for LiteLLM.
This module provides a flexible object pooling system that can be used
to pool any type of object, reducing memory allocation overhead and
improving performance for frequently created/destroyed objects.
Memory Management Strategy:
- Balanced eviction-based memory control to optimize reuse ratio
- Moderate eviction frequency (300s) to maintain high object reuse
- Conservative eviction weight (0.3) to avoid destroying useful objects
- Lower pre-warm count (5) to reduce initial memory footprint
- Always keeps at least one object available for high availability
- Unlimited pools when maxsize is not specified (eviction controls actual usage)
"""
from typing import Any, Callable, Optional, Type, TypeVar
from pond import Pond, PooledObject, PooledObjectFactory
T = TypeVar('T')
class GenericPooledObjectFactory(PooledObjectFactory):
"""Generic factory class for creating pooled objects of any type."""
def __init__(
self,
object_class: Type[T],
pooled_maxsize: Optional[int] = None, # None = unlimited pool with eviction-based memory control
least_one: bool = True, # Always keep at least one for high concurrency
initializer: Optional[Callable[[T], None]] = None
):
# Only pass maxsize to Pond if user specified it - otherwise let Pond handle unlimited pools
if pooled_maxsize is not None:
super().__init__(pooled_maxsize=pooled_maxsize, least_one=least_one)
else:
super().__init__(least_one=least_one)
self.object_class = object_class
self.initializer = initializer
self._user_maxsize = pooled_maxsize # Store original user preference
def createInstance(self) -> PooledObject:
"""Create a new instance wrapped in a PooledObject."""
# Create a properly initialized instance
obj = self.object_class()
return PooledObject(obj)
def destroy(self, pooled_object: PooledObject):
"""Destroy the pooled object."""
if hasattr(pooled_object.keeped_object, '__dict__'):
pooled_object.keeped_object.__dict__.clear()
del pooled_object
def reset(self, pooled_object: PooledObject, **kwargs: Any) -> PooledObject:
"""Reset the pooled object to a clean state."""
obj = pooled_object.keeped_object
# Reset the object by calling its reset method if it exists
if hasattr(obj, 'reset') and callable(getattr(obj, 'reset')):
obj.reset()
else:
# Fallback: clear all attributes to reset the object
if hasattr(obj, '__dict__'):
obj.__dict__.clear()
return pooled_object
def validate(self, pooled_object: PooledObject) -> bool:
"""Validate if the pooled object is still usable."""
return pooled_object.keeped_object is not None
# Global pond instances
_pools: dict[str, Pond] = {}
def get_object_pool(
pool_name: str,
object_class: Type[T],
pooled_maxsize: Optional[int] = None, # None = unlimited pool with eviction-based memory control
least_one: bool = True, # Always keep at least one
borrowed_timeout: int = 10, # Longer timeout for high concurrency
time_between_eviction_runs: int = 300, # Less frequent eviction to maintain high reuse ratio
eviction_weight: float = 0.3, # Less aggressive eviction for better reuse
prewarm_count: int = 5 # Lower pre-warm count to reduce initial memory usage
) -> Pond:
"""Get or create a global object pool instance with balanced eviction-based memory control.
Memory is controlled through moderate eviction to balance reuse ratio and memory usage:
- Moderate eviction frequency (300s) to maintain high object reuse ratio
- Conservative eviction weight (0.3) to avoid destroying useful objects
- Lower pre-warm count (5) to reduce initial memory footprint
Args:
pool_name: Unique name for the pool
object_class: The class type to pool
pooled_maxsize: Maximum number of objects in the pool (None = truly unlimited)
least_one: Whether to keep at least one object in the pool (default: True)
borrowed_timeout: Timeout for borrowing objects (seconds, default: 10)
time_between_eviction_runs: Time between eviction runs (seconds, default: 300)
eviction_weight: Weight for eviction algorithm (default: 0.3, conservative)
prewarm_count: Number of objects to pre-warm the pool with (default: 5)
Returns:
Pond instance for the specified object type
"""
if pool_name in _pools:
return _pools[pool_name]
# Create new pond
pond = Pond(
borrowed_timeout=borrowed_timeout,
time_between_eviction_runs=time_between_eviction_runs,
thread_daemon=True,
eviction_weight=eviction_weight
)
# Register the factory with user's maxsize preference
factory = GenericPooledObjectFactory(
object_class=object_class,
pooled_maxsize=pooled_maxsize,
least_one=least_one
)
pond.register(factory, name=f"{pool_name}Factory")
# Pre-warm the pool
_prewarm_pool(pond, pool_name, prewarm_count)
_pools[pool_name] = pond
return pond
def _prewarm_pool(pond: Pond, pool_name: str, prewarm_count: int = 20) -> None:
"""Pre-warm the pool with initial objects for high concurrency."""
for _ in range(prewarm_count):
try:
pooled_obj = pond.borrow(name=f"{pool_name}Factory")
pond.recycle(pooled_obj, name=f"{pool_name}Factory")
except Exception:
# If pre-warming fails, just continue
break
litellm/litellm_core_utils/prompt_templates/factory.py
+66 -50
View File
@@ -16,8 +16,8 @@ from litellm import verbose_logger
from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client
from litellm.types.files import get_file_extension_from_mime_type
from litellm.types.llms.anthropic import *
from litellm.types.llms.bedrock import MessageBlock as BedrockMessageBlock
from litellm.types.llms.bedrock import CachePointBlock
from litellm.types.llms.bedrock import MessageBlock as BedrockMessageBlock
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.llms.ollama import OllamaVisionModelObject
from litellm.types.llms.openai import (
@@ -1067,10 +1067,10 @@ def convert_to_gemini_tool_call_invoke(
if tool_calls is not None:
for tool in tool_calls:
if "function" in tool:
gemini_function_call: Optional[VertexFunctionCall] = (
_gemini_tool_call_invoke_helper(
function_call_params=tool["function"]
)
gemini_function_call: Optional[
VertexFunctionCall
] = _gemini_tool_call_invoke_helper(
function_call_params=tool["function"]
)
if gemini_function_call is not None:
_parts_list.append(
@@ -1589,9 +1589,9 @@ def anthropic_messages_pt( # noqa: PLR0915
)
if "cache_control" in _content_element:
_anthropic_content_element["cache_control"] = (
_content_element["cache_control"]
)
_anthropic_content_element[
"cache_control"
] = _content_element["cache_control"]
user_content.append(_anthropic_content_element)
elif m.get("type", "") == "text":
m = cast(ChatCompletionTextObject, m)
@@ -1629,9 +1629,9 @@ def anthropic_messages_pt( # noqa: PLR0915
)
if "cache_control" in _content_element:
_anthropic_content_text_element["cache_control"] = (
_content_element["cache_control"]
)
_anthropic_content_text_element[
"cache_control"
] = _content_element["cache_control"]
user_content.append(_anthropic_content_text_element)
@@ -2482,8 +2482,7 @@ class BedrockImageProcessor:
if is_document:
return BedrockImageProcessor._get_document_format(
mime_type=mime_type,
supported_doc_formats=supported_doc_formats
mime_type=mime_type, supported_doc_formats=supported_doc_formats
)
else:
@@ -2495,12 +2494,9 @@ class BedrockImageProcessor:
f"Unsupported image format: {image_format}. Supported formats: {supported_image_and_video_formats}"
)
return image_format
@staticmethod
def _get_document_format(
mime_type: str,
supported_doc_formats: List[str]
) -> str:
def _get_document_format(mime_type: str, supported_doc_formats: List[str]) -> str:
"""
Get the document format from the mime type
@@ -2519,13 +2515,9 @@ class BedrockImageProcessor:
The document format
"""
valid_extensions: Optional[List[str]] = None
potential_extensions = mimetypes.guess_all_extensions(
mime_type, strict=False
)
potential_extensions = mimetypes.guess_all_extensions(mime_type, strict=False)
valid_extensions = [
ext[1:]
for ext in potential_extensions
if ext[1:] in supported_doc_formats
ext[1:] for ext in potential_extensions if ext[1:] in supported_doc_formats
]
# Fallback to types/files.py if mimetypes doesn't return valid extensions
@@ -2689,10 +2681,12 @@ def _convert_to_bedrock_tool_call_invoke(
)
bedrock_content_block = BedrockContentBlock(toolUse=bedrock_tool)
_parts_list.append(bedrock_content_block)
# Check for cache_control and add a separate cachePoint block
if tool.get("cache_control", None) is not None:
cache_point_block = BedrockContentBlock(cachePoint=CachePointBlock(type="default"))
cache_point_block = BedrockContentBlock(
cachePoint=CachePointBlock(type="default")
)
_parts_list.append(cache_point_block)
return _parts_list
except Exception as e:
@@ -2754,7 +2748,7 @@ def _convert_to_bedrock_tool_call_result(
for content in content_list:
if content["type"] == "text":
content_str += content["text"]
message.get("name", "")
id = str(message.get("tool_call_id", str(uuid.uuid4())))
@@ -2763,7 +2757,7 @@ def _convert_to_bedrock_tool_call_result(
content=[tool_result_content_block],
toolUseId=id,
)
content_block = BedrockContentBlock(toolResult=tool_result)
return content_block
@@ -3128,6 +3122,12 @@ class BedrockConverseMessagesProcessor:
if element["type"] == "text":
_part = BedrockContentBlock(text=element["text"])
_parts.append(_part)
elif element["type"] == "guarded_text":
# Wrap guarded_text in guardContent block
_part = BedrockContentBlock(
guardContent={"text": {"text": element["text"]}}
)
_parts.append(_part)
elif element["type"] == "image_url":
format: Optional[str] = None
if isinstance(element["image_url"], dict):
@@ -3199,26 +3199,29 @@ class BedrockConverseMessagesProcessor:
current_message = messages[msg_i]
tool_call_result = _convert_to_bedrock_tool_call_result(current_message)
tool_content.append(tool_call_result)
# Check if we need to add a separate cachePoint block
has_cache_control = False
# Check for message-level cache_control
if current_message.get("cache_control", None) is not None:
has_cache_control = True
# Check for content-level cache_control in list content
elif isinstance(current_message.get("content"), list):
for content_element in current_message["content"]:
if (isinstance(content_element, dict) and
content_element.get("cache_control", None) is not None):
if (
isinstance(content_element, dict)
and content_element.get("cache_control", None) is not None
):
has_cache_control = True
break
# Add a separate cachePoint block if cache_control is present
if has_cache_control:
cache_point_block = BedrockContentBlock(cachePoint=CachePointBlock(type="default"))
cache_point_block = BedrockContentBlock(
cachePoint=CachePointBlock(type="default")
)
tool_content.append(cache_point_block)
msg_i += 1
if tool_content:
@@ -3299,7 +3302,7 @@ class BedrockConverseMessagesProcessor:
image_url=image_url
)
assistants_parts.append(assistants_part)
# Add cache point block for assistant content elements
# Add cache point block for assistant content elements
_cache_point_block = (
litellm.AmazonConverseConfig()._get_cache_point_block(
message_block=cast(
@@ -3311,8 +3314,12 @@ class BedrockConverseMessagesProcessor:
if _cache_point_block is not None:
assistants_parts.append(_cache_point_block)
assistant_content.extend(assistants_parts)
elif _assistant_content is not None and isinstance(_assistant_content, str):
assistant_content.append(BedrockContentBlock(text=_assistant_content))
elif _assistant_content is not None and isinstance(
_assistant_content, str
):
assistant_content.append(
BedrockContentBlock(text=_assistant_content)
)
# Add cache point block for assistant string content
_cache_point_block = (
litellm.AmazonConverseConfig()._get_cache_point_block(
@@ -3496,6 +3503,12 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915
if element["type"] == "text":
_part = BedrockContentBlock(text=element["text"])
_parts.append(_part)
elif element["type"] == "guarded_text":
# Wrap guarded_text in guardContent block
_part = BedrockContentBlock(
guardContent={"text": {"text": element["text"]}}
)
_parts.append(_part)
elif element["type"] == "image_url":
format: Optional[str] = None
if isinstance(element["image_url"], dict):
@@ -3565,29 +3578,33 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915
while msg_i < len(messages) and messages[msg_i]["role"] == "tool":
tool_call_result = _convert_to_bedrock_tool_call_result(messages[msg_i])
current_message = messages[msg_i]
# Add the tool result first
tool_content.append(tool_call_result)
# Check if we need to add a separate cachePoint block
has_cache_control = False
# Check for message-level cache_control
if current_message.get("cache_control", None) is not None:
has_cache_control = True
# Check for content-level cache_control in list content
elif isinstance(current_message.get("content"), list):
for content_element in current_message["content"]:
if (isinstance(content_element, dict) and
content_element.get("cache_control", None) is not None):
if (
isinstance(content_element, dict)
and content_element.get("cache_control", None) is not None
):
has_cache_control = True
break
# Add a separate cachePoint block if cache_control is present
if has_cache_control:
cache_point_block = BedrockContentBlock(cachePoint=CachePointBlock(type="default"))
cache_point_block = BedrockContentBlock(
cachePoint=CachePointBlock(type="default")
)
tool_content.append(cache_point_block)
msg_i += 1
if tool_content:
# if last message was a 'user' message, then add a blank assistant message (bedrock requires alternating roles)
@@ -3852,10 +3869,9 @@ def function_call_prompt(messages: list, functions: list):
if isinstance(message["content"], str):
message["content"] += f""" {function_prompt}"""
else:
message["content"].append({
"type": "text",
"text": f""" {function_prompt}"""
})
message["content"].append(
{"type": "text", "text": f""" {function_prompt}"""}
)
function_added_to_prompt = True
if function_added_to_prompt is False:
litellm/llms/anthropic/chat/transformation.py
+24 -3
View File
@@ -45,7 +45,10 @@ from litellm.types.llms.openai import (
OpenAIMcpServerTool,
OpenAIWebSearchOptions,
)
from litellm.types.utils import CompletionTokensDetailsWrapper
from litellm.types.utils import (
CacheCreationTokenDetails,
CompletionTokensDetailsWrapper,
)
from litellm.types.utils import Message as LitellmMessage
from litellm.types.utils import PromptTokensDetailsWrapper, ServerToolUse
from litellm.utils import (
@@ -200,8 +203,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
)
_allowed_properties = set(AnthropicInputSchema.__annotations__.keys())
input_schema_filtered = {k: v for k, v in _input_schema.items() if k in _allowed_properties}
input_anthropic_schema: AnthropicInputSchema = AnthropicInputSchema(**input_schema_filtered)
input_schema_filtered = {
k: v for k, v in _input_schema.items() if k in _allowed_properties
}
input_anthropic_schema: AnthropicInputSchema = AnthropicInputSchema(
**input_schema_filtered
)
_tool = AnthropicMessagesTool(
name=tool["function"]["name"],
@@ -816,12 +823,14 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
_usage = usage_object
cache_creation_input_tokens: int = 0
cache_read_input_tokens: int = 0
cache_creation_token_details: Optional[CacheCreationTokenDetails] = None
web_search_requests: Optional[int] = None
if (
"cache_creation_input_tokens" in _usage
and _usage["cache_creation_input_tokens"] is not None
):
cache_creation_input_tokens = _usage["cache_creation_input_tokens"]
prompt_tokens += cache_creation_input_tokens
if (
"cache_read_input_tokens" in _usage
and _usage["cache_read_input_tokens"] is not None
@@ -837,8 +846,20 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
int, _usage["server_tool_use"]["web_search_requests"]
)
if "cache_creation" in _usage and _usage["cache_creation"] is not None:
cache_creation_token_details = CacheCreationTokenDetails(
ephemeral_5m_input_tokens=_usage["cache_creation"].get(
"ephemeral_5m_input_tokens"
),
ephemeral_1h_input_tokens=_usage["cache_creation"].get(
"ephemeral_1h_input_tokens"
),
)
prompt_tokens_details = PromptTokensDetailsWrapper(
cached_tokens=cache_read_input_tokens,
cache_creation_tokens=cache_read_input_tokens,
cache_creation_token_details=cache_creation_token_details,
)
completion_token_details = (
CompletionTokensDetailsWrapper(
litellm/llms/base_llm/batches/transformation.py
+42
View File
@@ -158,6 +158,48 @@ class BaseBatchesConfig(ABC):
"""
pass
@abstractmethod
def transform_retrieve_batch_request(
self,
batch_id: str,
optional_params: dict,
litellm_params: dict,
) -> Union[bytes, str, Dict[str, Any]]:
"""
Transform the batch retrieval request to provider-specific format.
Args:
batch_id: Batch ID to retrieve
optional_params: Optional parameters
litellm_params: LiteLLM parameters
Returns:
Transformed request data
"""
pass
@abstractmethod
def transform_retrieve_batch_response(
self,
model: Optional[str],
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> LiteLLMBatch:
"""
Transform provider-specific batch retrieval response to LiteLLM format.
Args:
model: Model name
raw_response: Raw HTTP response
logging_obj: Logging object
litellm_params: LiteLLM parameters
Returns:
LiteLLM batch object
"""
pass
@abstractmethod
def get_error_class(
self, error_message: str, status_code: int, headers: Union[Dict, Headers]
litellm/llms/bedrock/base_aws_llm.py
+147 -58
View File
@@ -20,7 +20,11 @@ from pydantic import BaseModel
from litellm._logging import verbose_logger
from litellm.caching.caching import DualCache
from litellm.constants import BEDROCK_INVOKE_PROVIDERS_LITERAL, BEDROCK_MAX_POLICY_SIZE
from litellm.constants import (
BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
BEDROCK_INVOKE_PROVIDERS_LITERAL,
BEDROCK_MAX_POLICY_SIZE,
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.secret_managers.main import get_secret, get_secret_str
@@ -189,23 +193,32 @@ class BaseAWSLLM:
# Check if we're in IRSA and trying to assume the same role we already have
current_role_arn = os.getenv("AWS_ROLE_ARN")
web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")
# In IRSA environments, we should skip role assumption if we're already running as the target role
# This is true when:
# 1. We have AWS_ROLE_ARN set (current role)
# 2. We have AWS_WEB_IDENTITY_TOKEN_FILE set (IRSA environment)
# 3. The current role matches the requested role
if (current_role_arn and web_identity_token_file and
current_role_arn == aws_role_name):
verbose_logger.debug("Using IRSA same-role optimization: calling _auth_with_env_vars")
if (
current_role_arn
and web_identity_token_file
and current_role_arn == aws_role_name
):
verbose_logger.debug(
"Using IRSA same-role optimization: calling _auth_with_env_vars"
)
# We're already running as this role via IRSA, no need to assume it again
# Use the default boto3 credentials (which will use the IRSA credentials)
credentials, _cache_ttl = self._auth_with_env_vars()
else:
verbose_logger.debug("Using role assumption: calling _auth_with_aws_role")
verbose_logger.debug(
"Using role assumption: calling _auth_with_aws_role"
)
# If aws_session_name is not provided, generate a default one
if aws_session_name is None:
aws_session_name = f"litellm-session-{int(datetime.now().timestamp())}"
aws_session_name = (
f"litellm-session-{int(datetime.now().timestamp())}"
)
credentials, _cache_ttl = self._auth_with_aws_role(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
@@ -318,6 +331,40 @@ class BaseAWSLLM:
return provider
return None
@staticmethod
def get_bedrock_embedding_provider(
model: str,
) -> Optional[BEDROCK_EMBEDDING_PROVIDERS_LITERAL]:
"""
Helper function to get the bedrock embedding provider from the model
Handles scenarios like:
1. model=cohere.embed-english-v3:0 -> Returns `cohere`
2. model=amazon.titan-embed-text-v1 -> Returns `amazon`
3. model=us.twelvelabs.marengo-embed-2-7-v1:0 -> Returns `twelvelabs`
4. model=twelvelabs.marengo-embed-2-7-v1:0 -> Returns `twelvelabs`
"""
# Handle regional models like us.twelvelabs.marengo-embed-2-7-v1:0
if "." in model:
parts = model.split(".")
# Check if the second part (after potential region) is a known provider
if len(parts) >= 2:
potential_provider = parts[1] # e.g., "twelvelabs" from "us.twelvelabs.marengo-embed-2-7-v1:0"
if potential_provider in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL):
return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, potential_provider)
# Check if the first part is a known provider (standard format)
potential_provider = parts[0] # e.g., "cohere" from "cohere.embed-english-v3:0"
if potential_provider in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL):
return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, potential_provider)
# Fallback: check if any provider name appears in the model string
for provider in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL):
if provider in model:
return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, provider)
return None
def _get_aws_region_name(
self,
optional_params: dict,
@@ -479,55 +526,67 @@ class BaseAWSLLM:
iam_creds = session.get_credentials()
return iam_creds, self._get_default_ttl_for_boto3_credentials()
def _handle_irsa_cross_account(self, irsa_role_arn: str, aws_role_name: str,
aws_session_name: str, region: str, web_identity_token_file: str,
aws_external_id: Optional[str] = None) -> dict:
def _handle_irsa_cross_account(
self,
irsa_role_arn: str,
aws_role_name: str,
aws_session_name: str,
region: str,
web_identity_token_file: str,
aws_external_id: Optional[str] = None,
) -> dict:
"""Handle cross-account role assumption for IRSA."""
import boto3
verbose_logger.debug("Cross-account role assumption detected")
# Read the web identity token
with open(web_identity_token_file, 'r') as f:
with open(web_identity_token_file, "r") as f:
web_identity_token = f.read().strip()
# 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)
# Manually assume the IRSA role with the session name
verbose_logger.debug(f"Manually assuming IRSA role {irsa_role_arn} with session {aws_session_name}")
verbose_logger.debug(
f"Manually assuming IRSA role {irsa_role_arn} with session {aws_session_name}"
)
irsa_response = sts_client.assume_role_with_web_identity(
RoleArn=irsa_role_arn,
RoleSessionName=aws_session_name,
WebIdentityToken=web_identity_token
WebIdentityToken=web_identity_token,
)
# Extract the credentials from the IRSA assumption
irsa_creds = irsa_response["Credentials"]
# Create a new STS client with the IRSA credentials
with tracer.trace("boto3.client(sts) with manual IRSA credentials"):
sts_client_with_creds = boto3.client(
'sts',
"sts",
region_name=region,
aws_access_key_id=irsa_creds["AccessKeyId"],
aws_secret_access_key=irsa_creds["SecretAccessKey"],
aws_session_token=irsa_creds["SessionToken"]
aws_session_token=irsa_creds["SessionToken"],
)
# Get current caller identity for debugging
try:
caller_identity = sts_client_with_creds.get_caller_identity()
verbose_logger.debug(f"Current identity after manual IRSA assumption: {caller_identity.get('Arn', 'unknown')}")
verbose_logger.debug(
f"Current identity after manual IRSA assumption: {caller_identity.get('Arn', 'unknown')}"
)
except Exception as e:
verbose_logger.debug(f"Failed to get caller identity: {e}")
# Now assume the target role
verbose_logger.debug(f"Attempting to assume target role: {aws_role_name} with session: {aws_session_name}")
verbose_logger.debug(
f"Attempting to assume target role: {aws_role_name} with session: {aws_session_name}"
)
assume_role_params = {
"RoleArn": aws_role_name,
"RoleSessionName": aws_session_name
"RoleSessionName": aws_session_name,
}
# Add ExternalId parameter if provided
@@ -536,27 +595,36 @@ class BaseAWSLLM:
return sts_client_with_creds.assume_role(**assume_role_params)
def _handle_irsa_same_account(self, aws_role_name: str, aws_session_name: str, region: str,
aws_external_id: Optional[str] = None) -> dict:
def _handle_irsa_same_account(
self,
aws_role_name: str,
aws_session_name: str,
region: str,
aws_external_id: Optional[str] = None,
) -> dict:
"""Handle same-account role assumption for IRSA."""
import boto3
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)
# Get current caller identity for debugging
try:
caller_identity = sts_client.get_caller_identity()
verbose_logger.debug(f"Current IRSA identity: {caller_identity.get('Arn', 'unknown')}")
verbose_logger.debug(
f"Current IRSA identity: {caller_identity.get('Arn', 'unknown')}"
)
except Exception as e:
verbose_logger.debug(f"Failed to get caller identity: {e}")
# Assume the role
verbose_logger.debug(f"Attempting to assume role: {aws_role_name} with session: {aws_session_name}")
verbose_logger.debug(
f"Attempting to assume role: {aws_role_name} with session: {aws_session_name}"
)
assume_role_params = {
"RoleArn": aws_role_name,
"RoleSessionName": aws_session_name
"RoleSessionName": aws_session_name,
}
# Add ExternalId parameter if provided
@@ -565,20 +633,24 @@ class BaseAWSLLM:
return sts_client.assume_role(**assume_role_params)
def _extract_credentials_and_ttl(self, sts_response: dict) -> Tuple[Credentials, Optional[int]]:
def _extract_credentials_and_ttl(
self, sts_response: dict
) -> Tuple[Credentials, Optional[int]]:
"""Extract credentials and TTL from STS response."""
from botocore.credentials import Credentials
sts_credentials = sts_response["Credentials"]
credentials = Credentials(
access_key=sts_credentials["AccessKeyId"],
secret_key=sts_credentials["SecretAccessKey"],
token=sts_credentials["SessionToken"],
)
expiration_time = sts_credentials["Expiration"]
ttl = int((expiration_time - datetime.now(expiration_time.tzinfo)).total_seconds())
ttl = int(
(expiration_time - datetime.now(expiration_time.tzinfo)).total_seconds()
)
return credentials, ttl
@tracer.wrap()
@@ -600,34 +672,51 @@ class BaseAWSLLM:
# Check if we're in an EKS/IRSA environment
web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")
irsa_role_arn = os.getenv("AWS_ROLE_ARN")
# If we have IRSA environment variables and no explicit credentials,
# we need to use the web identity token flow
if (web_identity_token_file and irsa_role_arn and
aws_access_key_id is None and aws_secret_access_key is None):
if (
web_identity_token_file
and irsa_role_arn
and aws_access_key_id is None
and aws_secret_access_key is None
):
# For cross-account role assumption with specific session names,
# we need to manually assume the IRSA role first with the correct session name
verbose_logger.debug(f"IRSA detected: using web identity token from {web_identity_token_file}")
verbose_logger.debug(
f"IRSA detected: using web identity token from {web_identity_token_file}"
)
try:
# Get region from environment
region = os.getenv("AWS_REGION") or os.getenv("AWS_DEFAULT_REGION") or "us-east-1"
region = (
os.getenv("AWS_REGION")
or os.getenv("AWS_DEFAULT_REGION")
or "us-east-1"
)
# Check if we need to do cross-account role assumption
if aws_role_name != irsa_role_arn:
sts_response = self._handle_irsa_cross_account(
irsa_role_arn, aws_role_name, aws_session_name, region, web_identity_token_file, aws_external_id
irsa_role_arn,
aws_role_name,
aws_session_name,
region,
web_identity_token_file,
aws_external_id,
)
else:
sts_response = self._handle_irsa_same_account(
aws_role_name, aws_session_name, region, aws_external_id
)
return self._extract_credentials_and_ttl(sts_response)
except Exception as e:
verbose_logger.debug(f"Failed to assume role via IRSA: {e}")
if "AccessDenied" in str(e) and "is not authorized to perform: sts:AssumeRole" in str(e):
if "AccessDenied" in str(
e
) and "is not authorized to perform: sts:AssumeRole" in str(e):
# Provide a more helpful error message for trust policy issues
verbose_logger.error(
f"Access denied when trying to assume role {aws_role_name}. "
@@ -636,7 +725,7 @@ class BaseAWSLLM:
)
# Re-raise the exception instead of falling through
raise
# In EKS/IRSA environments, use ambient credentials (no explicit keys needed)
# This allows the web identity token to work automatically
if aws_access_key_id is None and aws_secret_access_key is None:
@@ -653,7 +742,7 @@ class BaseAWSLLM:
assume_role_params = {
"RoleArn": aws_role_name,
"RoleSessionName": aws_session_name
"RoleSessionName": aws_session_name,
}
# Add ExternalId parameter if provided
@@ -782,14 +871,14 @@ class BaseAWSLLM:
)
# Determine proxy_endpoint_url
if env_aws_bedrock_runtime_endpoint and isinstance(
env_aws_bedrock_runtime_endpoint, str
):
proxy_endpoint_url = env_aws_bedrock_runtime_endpoint
elif aws_bedrock_runtime_endpoint is not None and isinstance(
if aws_bedrock_runtime_endpoint is not None and isinstance(
aws_bedrock_runtime_endpoint, str
):
proxy_endpoint_url = aws_bedrock_runtime_endpoint
elif env_aws_bedrock_runtime_endpoint and isinstance(
env_aws_bedrock_runtime_endpoint, str
):
proxy_endpoint_url = env_aws_bedrock_runtime_endpoint
else:
proxy_endpoint_url = endpoint_url
litellm/llms/bedrock/batches/transformation.py
+206 -6
View File
@@ -7,7 +7,6 @@ from httpx import Headers, Response
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.bedrock import (
BedrockBatchJobStatus,
BedrockCreateBatchRequest,
BedrockCreateBatchResponse,
BedrockInputDataConfig,
@@ -200,19 +199,23 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
# Extract information from typed Bedrock response
job_arn = response_data.get("jobArn", "")
status: BedrockBatchJobStatus = response_data.get("status", "Submitted")
status_str: str = str(response_data.get("status", "Submitted"))
# Map Bedrock status to OpenAI-compatible status
status_mapping: Dict[BedrockBatchJobStatus, str] = {
status_mapping: Dict[str, str] = {
"Submitted": "validating",
"Validating": "validating",
"Scheduled": "in_progress",
"InProgress": "in_progress",
"PartiallyCompleted": "completed",
"Completed": "completed",
"Failed": "failed",
"Stopping": "cancelling",
"Stopped": "cancelled"
"Stopped": "cancelled",
"Expired": "expired",
}
openai_status = cast(Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"], status_mapping.get(status, "validating"))
openai_status = cast(Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"], status_mapping.get(status_str, "validating"))
# Get original request data from litellm_params if available
original_request = litellm_params.get("original_batch_request", {})
@@ -229,7 +232,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
output_file_id=None, # Will be populated when job completes
error_file_id=None,
created_at=int(time.time()),
in_progress_at=int(time.time()) if status == "InProgress" else None,
in_progress_at=int(time.time()) if status_str == "InProgress" else None,
expires_at=None,
finalizing_at=None,
completed_at=None,
@@ -241,6 +244,203 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
metadata=original_request.get("metadata", {}),
)
def transform_retrieve_batch_request(
self,
batch_id: str,
optional_params: dict,
litellm_params: dict,
) -> Dict[str, Any]:
"""
Transform batch retrieval request for Bedrock.
Args:
batch_id: Bedrock job ARN
optional_params: Optional parameters
litellm_params: LiteLLM parameters
Returns:
Transformed request data for Bedrock GetModelInvocationJob API
"""
# For Bedrock, batch_id should be the full job ARN
# The GetModelInvocationJob API expects the full ARN as the identifier
if not batch_id.startswith("arn:aws:bedrock:"):
raise ValueError(f"Invalid batch_id format. Expected ARN, got: {batch_id}")
# Extract the job identifier from the ARN - use the full ARN path part
# ARN format: arn:aws:bedrock:region:account:model-invocation-job/job-name
arn_parts = batch_id.split(":")
if len(arn_parts) < 6:
raise ValueError(f"Invalid ARN format: {batch_id}")
region = arn_parts[3]
# arn_parts[5] contains "model-invocation-job/{jobId}"
# Build the endpoint URL for GetModelInvocationJob
# AWS API format: GET /model-invocation-job/{jobIdentifier}
# Use the FULL ARN as jobIdentifier and URL-encode it (includes ':' and '/')
import urllib.parse as _ul
encoded_arn = _ul.quote(batch_id, safe="")
endpoint_url = f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}"
# Use common utility for AWS signing
signed_headers, _ = self.common_utils.sign_aws_request(
service_name="bedrock",
data={}, # GET request has no body
endpoint_url=endpoint_url,
optional_params=optional_params,
method="GET"
)
# Return pre-signed request format
return {
"method": "GET",
"url": endpoint_url,
"headers": signed_headers,
"data": None
}
def _parse_timestamps_and_status(self, response_data, status_str: str):
"""Helper to parse timestamps based on status."""
import datetime
def parse_timestamp(ts_str: Optional[str]) -> Optional[int]:
if not ts_str:
return None
try:
dt = datetime.datetime.fromisoformat(ts_str.replace('Z', '+00:00'))
return int(dt.timestamp())
except Exception:
return None
created_at = parse_timestamp(str(response_data.get("submitTime")) if response_data.get("submitTime") is not None else None)
in_progress_states = {"InProgress", "Validating", "Scheduled"}
in_progress_at = (
parse_timestamp(str(response_data.get("lastModifiedTime")) if response_data.get("lastModifiedTime") is not None else None)
if status_str in in_progress_states
else None
)
completed_at = parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str in {"Completed", "PartiallyCompleted"} else None
failed_at = parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str == "Failed" else None
cancelled_at = parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str == "Stopped" else None
expires_at = parse_timestamp(str(response_data.get("jobExpirationTime")) if response_data.get("jobExpirationTime") is not None else None)
return created_at, in_progress_at, completed_at, failed_at, cancelled_at, expires_at
def _extract_file_configs(self, response_data):
"""Helper to extract input and output file configurations."""
# Extract input file ID
input_file_id = ""
input_data_config = response_data.get("inputDataConfig", {})
if isinstance(input_data_config, dict):
s3_input_config = input_data_config.get("s3InputDataConfig", {})
if isinstance(s3_input_config, dict):
input_file_id = s3_input_config.get("s3Uri", "")
# Extract output file ID
output_file_id = None
output_data_config = response_data.get("outputDataConfig", {})
if isinstance(output_data_config, dict):
s3_output_config = output_data_config.get("s3OutputDataConfig", {})
if isinstance(s3_output_config, dict):
output_file_id = s3_output_config.get("s3Uri", "")
return input_file_id, output_file_id
def _extract_errors_and_metadata(self, response_data, raw_response):
"""Helper to extract errors and enriched metadata."""
# Extract errors
message = response_data.get("message")
errors = None
if message:
from openai.types.batch import Errors
from openai.types.batch_error import BatchError
errors = Errors(
data=[BatchError(message=message, code=str(raw_response.status_code))],
object="list"
)
# Enrich metadata with useful Bedrock fields
enriched_metadata_raw: Dict[str, Any] = {
"jobName": response_data.get("jobName"),
"clientRequestToken": response_data.get("clientRequestToken"),
"modelId": response_data.get("modelId"),
"roleArn": response_data.get("roleArn"),
"timeoutDurationInHours": response_data.get("timeoutDurationInHours"),
"vpcConfig": response_data.get("vpcConfig"),
}
import json as _json
enriched_metadata: Dict[str, str] = {}
for _k, _v in enriched_metadata_raw.items():
if _v is None:
continue
if isinstance(_v, (dict, list)):
try:
enriched_metadata[_k] = _json.dumps(_v)
except Exception:
enriched_metadata[_k] = str(_v)
else:
enriched_metadata[_k] = str(_v)
return errors, enriched_metadata
def transform_retrieve_batch_response(
self,
model: Optional[str],
raw_response: Response,
logging_obj: Any,
litellm_params: dict,
) -> LiteLLMBatch:
"""
Transform Bedrock batch retrieval response to LiteLLM format.
"""
from litellm.types.llms.bedrock import BedrockGetBatchResponse
try:
response_data: BedrockGetBatchResponse = raw_response.json()
except Exception as e:
raise ValueError(f"Failed to parse Bedrock batch response: {e}")
job_arn = response_data.get("jobArn", "")
status_str: str = str(response_data.get("status", "Submitted"))
# Map Bedrock status to OpenAI-compatible status
status_mapping: Dict[str, str] = {
"Submitted": "validating", "Validating": "validating", "Scheduled": "in_progress",
"InProgress": "in_progress", "PartiallyCompleted": "completed", "Completed": "completed",
"Failed": "failed", "Stopping": "cancelling", "Stopped": "cancelled", "Expired": "expired"
}
openai_status = cast(Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"], status_mapping.get(status_str, "validating"))
# Parse timestamps
created_at, in_progress_at, completed_at, failed_at, cancelled_at, expires_at = self._parse_timestamps_and_status(response_data, status_str)
# Extract file configurations
input_file_id, output_file_id = self._extract_file_configs(response_data)
# Extract errors and metadata
errors, enriched_metadata = self._extract_errors_and_metadata(response_data, raw_response)
return LiteLLMBatch(
id=job_arn,
object="batch",
endpoint="/v1/chat/completions",
errors=errors,
input_file_id=input_file_id,
completion_window="24h",
status=openai_status,
output_file_id=output_file_id,
error_file_id=None,
created_at=created_at or int(time.time()),
in_progress_at=in_progress_at,
expires_at=expires_at,
finalizing_at=None,
completed_at=completed_at,
failed_at=failed_at,
expired_at=None,
cancelling_at=None,
cancelled_at=cancelled_at,
request_counts=None,
metadata=enriched_metadata,
)
def get_error_class(
self, error_message: str, status_code: int, headers: Union[Dict, Headers]
) -> BaseLLMException:
litellm/llms/bedrock/chat/converse_transformation.py
+208 -45
View File
@@ -102,6 +102,61 @@ class AmazonConverseConfig(BaseConfig):
"performanceConfig": PerformanceConfigBlock,
}
@staticmethod
def _convert_consecutive_user_messages_to_guarded_text(
messages: List[AllMessageValues], optional_params: dict
) -> List[AllMessageValues]:
"""
Convert consecutive user messages at the end to guarded_text type if guardrailConfig is present
and no guarded_text is already present in those messages.
"""
# Check if guardrailConfig is present
if "guardrailConfig" not in optional_params:
return messages
# Find all consecutive user messages at the end
consecutive_user_message_indices = []
for i in range(len(messages) - 1, -1, -1):
if messages[i].get("role") == "user":
consecutive_user_message_indices.append(i)
else:
break
if not consecutive_user_message_indices:
return messages
# Process each consecutive user message
messages_copy = copy.deepcopy(messages)
for user_message_index in consecutive_user_message_indices:
user_message = messages_copy[user_message_index]
content = user_message.get("content", [])
if isinstance(content, list):
has_guarded_text = any(
isinstance(item, dict) and item.get("type") == "guarded_text"
for item in content
)
if has_guarded_text:
continue # Skip this message if it already has guarded_text
# Convert text elements to guarded_text
new_content = []
for item in content:
if isinstance(item, dict) and item.get("type") == "text":
new_item = {"type": "guarded_text", "text": item["text"]} # type: ignore
new_content.append(new_item)
else:
new_content.append(item)
messages_copy[user_message_index]["content"] = new_content # type: ignore
elif isinstance(content, str):
# If content is a string, convert it to guarded_text
messages_copy[user_message_index]["content"] = [ # type: ignore
{"type": "guarded_text", "text": content} # type: ignore
]
return messages_copy
@classmethod
def get_config(cls):
return {
@@ -120,6 +175,77 @@ class AmazonConverseConfig(BaseConfig):
and v is not None
}
def _validate_request_metadata(self, metadata: dict) -> None:
"""
Validate requestMetadata according to AWS Bedrock Converse API constraints.
Constraints:
- Maximum of 16 items
- Keys: 1-256 characters, pattern [a-zA-Z0-9\\s:_@$#=/+,-.]{1,256}
- Values: 0-256 characters, pattern [a-zA-Z0-9\\s:_@$#=/+,-.]{0,256}
"""
import re
if not isinstance(metadata, dict):
raise litellm.exceptions.BadRequestError(
message="requestMetadata must be a dictionary",
model="bedrock",
llm_provider="bedrock",
)
if len(metadata) > 16:
raise litellm.exceptions.BadRequestError(
message="requestMetadata can contain a maximum of 16 items",
model="bedrock",
llm_provider="bedrock",
)
key_pattern = re.compile(r'^[a-zA-Z0-9\s:_@$#=/+,.-]{1,256}$')
value_pattern = re.compile(r'^[a-zA-Z0-9\s:_@$#=/+,.-]{0,256}$')
for key, value in metadata.items():
if not isinstance(key, str):
raise litellm.exceptions.BadRequestError(
message="requestMetadata keys must be strings",
model="bedrock",
llm_provider="bedrock",
)
if not isinstance(value, str):
raise litellm.exceptions.BadRequestError(
message="requestMetadata values must be strings",
model="bedrock",
llm_provider="bedrock",
)
if len(key) == 0 or len(key) > 256:
raise litellm.exceptions.BadRequestError(
message="requestMetadata key length must be 1-256 characters",
model="bedrock",
llm_provider="bedrock",
)
if len(value) > 256:
raise litellm.exceptions.BadRequestError(
message="requestMetadata value length must be 0-256 characters",
model="bedrock",
llm_provider="bedrock",
)
if not key_pattern.match(key):
raise litellm.exceptions.BadRequestError(
message=f"requestMetadata key '{key}' contains invalid characters. Allowed: [a-zA-Z0-9\\s:_@$#=/+,.-]",
model="bedrock",
llm_provider="bedrock",
)
if not value_pattern.match(value):
raise litellm.exceptions.BadRequestError(
message=f"requestMetadata value '{value}' contains invalid characters. Allowed: [a-zA-Z0-9\\s:_@$#=/+,.-]",
model="bedrock",
llm_provider="bedrock",
)
def get_supported_openai_params(self, model: str) -> List[str]:
from litellm.utils import supports_function_calling
@@ -133,6 +259,7 @@ class AmazonConverseConfig(BaseConfig):
"top_p",
"extra_headers",
"response_format",
"requestMetadata",
]
if (
@@ -442,6 +569,10 @@ class AmazonConverseConfig(BaseConfig):
optional_params["thinking"] = AnthropicConfig._map_reasoning_effort(
value
)
if param == "requestMetadata":
if value is not None and isinstance(value, dict):
self._validate_request_metadata(value) # type: ignore
optional_params["requestMetadata"] = value
# Only update thinking tokens for non-GPT-OSS models
if "gpt-oss" not in model:
@@ -501,7 +632,6 @@ class AmazonConverseConfig(BaseConfig):
)
and not is_thinking_enabled
):
optional_params["tool_choice"] = ToolChoiceValuesBlock(
tool=SpecificToolChoiceBlock(name=RESPONSE_FORMAT_TOOL_NAME)
)
@@ -632,34 +762,8 @@ class AmazonConverseConfig(BaseConfig):
return {}
def _transform_request_helper(
self,
model: str,
system_content_blocks: List[SystemContentBlock],
optional_params: dict,
messages: Optional[List[AllMessageValues]] = None,
headers: Optional[dict] = None,
) -> CommonRequestObject:
## VALIDATE REQUEST
"""
Bedrock doesn't support tool calling without `tools=` param specified.
"""
if (
"tools" not in optional_params
and messages is not None
and has_tool_call_blocks(messages)
):
if litellm.modify_params:
optional_params["tools"] = add_dummy_tool(
custom_llm_provider="bedrock_converse"
)
else:
raise litellm.UnsupportedParamsError(
message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.",
model="",
llm_provider="bedrock",
)
def _prepare_request_params(self, optional_params: dict, model: str) -> tuple[dict, dict, dict]:
"""Prepare and separate request parameters."""
inference_params = copy.deepcopy(optional_params)
supported_converse_params = list(
AmazonConverseConfig.__annotations__.keys()
@@ -673,6 +777,11 @@ class AmazonConverseConfig(BaseConfig):
)
inference_params.pop("json_mode", None) # used for handling json_schema
# Extract requestMetadata before processing other parameters
request_metadata = inference_params.pop("requestMetadata", None)
if request_metadata is not None:
self._validate_request_metadata(request_metadata)
# keep supported params in 'inference_params', and set all model-specific params in 'additional_request_params'
additional_request_params = {
k: v for k, v in inference_params.items() if k not in total_supported_params
@@ -686,9 +795,10 @@ class AmazonConverseConfig(BaseConfig):
self._handle_top_k_value(model, inference_params)
)
original_tools = inference_params.pop("tools", [])
return inference_params, additional_request_params, request_metadata
# Initialize bedrock_tools
def _process_tools_and_beta(self, original_tools: list, model: str, headers: Optional[dict], additional_request_params: dict) -> tuple[List[ToolBlock], list]:
"""Process tools and collect anthropic_beta values."""
bedrock_tools: List[ToolBlock] = []
# Collect anthropic_beta values from user headers
@@ -730,6 +840,44 @@ class AmazonConverseConfig(BaseConfig):
seen.add(beta)
additional_request_params["anthropic_beta"] = unique_betas
return bedrock_tools, anthropic_beta_list
def _transform_request_helper(
self,
model: str,
system_content_blocks: List[SystemContentBlock],
optional_params: dict,
messages: Optional[List[AllMessageValues]] = None,
headers: Optional[dict] = None,
) -> CommonRequestObject:
## VALIDATE REQUEST
"""
Bedrock doesn't support tool calling without `tools=` param specified.
"""
if (
"tools" not in optional_params
and messages is not None
and has_tool_call_blocks(messages)
):
if litellm.modify_params:
optional_params["tools"] = add_dummy_tool(
custom_llm_provider="bedrock_converse"
)
else:
raise litellm.UnsupportedParamsError(
message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.",
model="",
llm_provider="bedrock",
)
# Prepare and separate parameters
inference_params, additional_request_params, request_metadata = self._prepare_request_params(optional_params, model)
original_tools = inference_params.pop("tools", [])
# Process tools and collect beta values
bedrock_tools, anthropic_beta_list = self._process_tools_and_beta(original_tools, model, headers, additional_request_params)
bedrock_tool_config: Optional[ToolConfigBlock] = None
if len(bedrock_tools) > 0:
tool_choice_values: ToolChoiceValuesBlock = inference_params.pop(
@@ -759,6 +907,10 @@ class AmazonConverseConfig(BaseConfig):
if bedrock_tool_config is not None:
data["toolConfig"] = bedrock_tool_config
# Request Metadata (top-level field)
if request_metadata is not None:
data["requestMetadata"] = request_metadata
return data
async def _async_transform_request(
@@ -770,6 +922,11 @@ class AmazonConverseConfig(BaseConfig):
headers: Optional[dict] = None,
) -> RequestObject:
messages, system_content_blocks = self._transform_system_message(messages)
# Convert last user message to guarded_text if guardrailConfig is present
messages = self._convert_consecutive_user_messages_to_guarded_text(
messages, optional_params
)
## TRANSFORMATION ##
_data: CommonRequestObject = self._transform_request_helper(
@@ -822,6 +979,11 @@ class AmazonConverseConfig(BaseConfig):
) -> RequestObject:
messages, system_content_blocks = self._transform_system_message(messages)
# Convert last user message to guarded_text if guardrailConfig is present
messages = self._convert_consecutive_user_messages_to_guarded_text(
messages, optional_params
)
_data: CommonRequestObject = self._transform_request_helper(
model=model,
system_content_blocks=system_content_blocks,
@@ -995,7 +1157,9 @@ class AmazonConverseConfig(BaseConfig):
return message, returned_finish_reason
def _translate_message_content(self, content_blocks: List[ContentBlock]) -> Tuple[
def _translate_message_content(
self, content_blocks: List[ContentBlock]
) -> Tuple[
str,
List[ChatCompletionToolCallChunk],
Optional[List[BedrockConverseReasoningContentBlock]],
@@ -1010,9 +1174,9 @@ class AmazonConverseConfig(BaseConfig):
"""
content_str = ""
tools: List[ChatCompletionToolCallChunk] = []
reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = (
None
)
reasoningContentBlocks: Optional[
List[BedrockConverseReasoningContentBlock]
] = None
for idx, content in enumerate(content_blocks):
"""
- Content is either a tool response or text
@@ -1133,9 +1297,9 @@ class AmazonConverseConfig(BaseConfig):
chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"}
content_str = ""
tools: List[ChatCompletionToolCallChunk] = []
reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = (
None
)
reasoningContentBlocks: Optional[
List[BedrockConverseReasoningContentBlock]
] = None
if message is not None:
(
@@ -1148,12 +1312,12 @@ class AmazonConverseConfig(BaseConfig):
chat_completion_message["provider_specific_fields"] = {
"reasoningContentBlocks": reasoningContentBlocks,
}
chat_completion_message["reasoning_content"] = (
self._transform_reasoning_content(reasoningContentBlocks)
)
chat_completion_message["thinking_blocks"] = (
self._transform_thinking_blocks(reasoningContentBlocks)
)
chat_completion_message[
"reasoning_content"
] = self._transform_reasoning_content(reasoningContentBlocks)
chat_completion_message[
"thinking_blocks"
] = self._transform_thinking_blocks(reasoningContentBlocks)
chat_completion_message["content"] = content_str
if (
json_mode is True
@@ -1171,7 +1335,6 @@ class AmazonConverseConfig(BaseConfig):
# Bedrock returns the response wrapped in a "properties" object
# We need to extract the actual content from this wrapper
try:
response_data = json.loads(json_mode_content_str)
# If Bedrock wrapped the response in "properties", extract the content
litellm/llms/bedrock/common_utils.py
+13 -8
View File
@@ -738,19 +738,24 @@ class CommonBatchFilesUtils:
)
# Prepare the request data
if isinstance(data, dict):
import json
request_data = json.dumps(data)
method_upper = method.upper()
if method_upper == "GET":
# GET requests should be signed with an empty payload
request_data = ""
headers = {}
else:
request_data = data
# Prepare headers
headers = {"Content-Type": "application/json"}
if isinstance(data, dict):
import json
request_data = json.dumps(data)
else:
request_data = data
# Prepare headers for non-GET requests
headers = {"Content-Type": "application/json"}
# Create AWS request and sign it
sigv4 = SigV4Auth(credentials, service_name, aws_region_name)
request = AWSRequest(
method=method.upper(), url=endpoint_url, data=request_data, headers=headers
method=method_upper, url=endpoint_url, data=request_data, headers=headers
)
sigv4.add_auth(request)
prepped = request.prepare()
litellm/llms/bedrock/count_tokens/handler.py
+123
View File
@@ -0,0 +1,123 @@
"""
AWS Bedrock CountTokens API handler.
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.count_tokens.transformation import BedrockCountTokensConfig
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
class BedrockCountTokensHandler(BedrockCountTokensConfig):
"""
Simplified handler for AWS Bedrock CountTokens API requests.
Uses existing LiteLLM infrastructure for authentication and request handling.
"""
async def handle_count_tokens_request(
self,
request_data: Dict[str, Any],
litellm_params: Dict[str, Any],
resolved_model: str,
) -> Dict[str, Any]:
"""
Handle a CountTokens request using existing LiteLLM patterns.
Args:
request_data: The incoming request payload
litellm_params: LiteLLM configuration parameters
resolved_model: The actual model ID resolved from router
Returns:
Dictionary containing token count response
"""
try:
# Validate the request
self.validate_count_tokens_request(request_data)
verbose_logger.debug(
f"Processing CountTokens request for resolved model: {resolved_model}"
)
# Get AWS region using existing LiteLLM function
aws_region_name = self._get_aws_region_name(
optional_params=litellm_params,
model=resolved_model,
model_id=None,
)
verbose_logger.debug(f"Retrieved AWS region: {aws_region_name}")
# Transform request to Bedrock format (supports both Converse and InvokeModel)
bedrock_request = self.transform_anthropic_to_bedrock_count_tokens(
request_data=request_data
)
verbose_logger.debug(f"Transformed request: {bedrock_request}")
# Get endpoint URL using simplified function
endpoint_url = self.get_bedrock_count_tokens_endpoint(
resolved_model, aws_region_name
)
verbose_logger.debug(f"Making request to: {endpoint_url}")
# Use existing _sign_request method from BaseAWSLLM
headers = {"Content-Type": "application/json"}
signed_headers, signed_body = self._sign_request(
service_name="bedrock",
headers=headers,
optional_params=litellm_params,
request_data=bedrock_request,
api_base=endpoint_url,
model=resolved_model,
)
async_client = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)
response = await async_client.post(
endpoint_url,
headers=signed_headers,
data=signed_body,
timeout=30.0,
)
verbose_logger.debug(f"Response status: {response.status_code}")
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}"},
)
bedrock_response = response.json()
verbose_logger.debug(f"Bedrock response: {bedrock_response}")
# Transform response back to expected format
final_response = self.transform_bedrock_response_to_anthropic(
bedrock_response
)
verbose_logger.debug(f"Final response: {final_response}")
return final_response
except HTTPException:
# Re-raise HTTP exceptions as-is
raise
except Exception as e:
verbose_logger.error(f"Error in CountTokens handler: {str(e)}")
raise HTTPException(
status_code=500,
detail={"error": f"CountTokens processing error: {str(e)}"},
)
litellm/llms/bedrock/count_tokens/transformation.py
+213
View File
@@ -0,0 +1,213 @@
"""
AWS Bedrock CountTokens API transformation logic.
This module handles the transformation of requests from Anthropic Messages API format
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
class BedrockCountTokensConfig(BaseAWSLLM):
"""
Configuration and transformation logic for AWS Bedrock CountTokens API.
AWS Bedrock CountTokens API Specification:
- Endpoint: POST /model/{modelId}/count-tokens
- Input formats: 'invokeModel' or 'converse'
- Response: {"inputTokens": <number>}
"""
def _detect_input_type(self, request_data: Dict[str, Any]) -> str:
"""
Detect whether to use 'converse' or 'invokeModel' input format.
Args:
request_data: The original request data
Returns:
'converse' or 'invokeModel'
"""
# If the request has messages in the expected Anthropic format, use converse
if "messages" in request_data and isinstance(request_data["messages"], list):
return "converse"
# For raw text or other formats, use invokeModel
# This handles cases where the input is prompt-based or already in raw Bedrock format
return "invokeModel"
def transform_anthropic_to_bedrock_count_tokens(
self,
request_data: Dict[str, Any],
) -> Dict[str, Any]:
"""
Transform request to Bedrock CountTokens format.
Supports both Converse and InvokeModel input types.
Input (Anthropic format):
{
"model": "claude-3-5-sonnet",
"messages": [{"role": "user", "content": "Hello!"}]
}
Output (Bedrock CountTokens format for Converse):
{
"input": {
"converse": {
"messages": [...],
"system": [...] (if present)
}
}
}
Output (Bedrock CountTokens format for InvokeModel):
{
"input": {
"invokeModel": {
"body": "{...raw model input...}"
}
}
}
"""
input_type = self._detect_input_type(request_data)
if input_type == "converse":
return self._transform_to_converse_format(request_data.get("messages", []))
else:
return self._transform_to_invoke_model_format(request_data)
def _transform_to_converse_format(
self, messages: List[Dict[str, Any]]
) -> Dict[str, Any]:
"""Transform to Converse input format."""
# Extract system messages if present
system_messages = []
user_messages = []
for message in messages:
if message.get("role") == "system":
system_messages.append({"text": message.get("content", "")})
else:
# Transform message content to Bedrock format
transformed_message: Dict[str, Any] = {"role": message.get("role"), "content": []}
# Handle content - ensure it's in the correct array format
content = message.get("content", "")
if isinstance(content, str):
# String content -> convert to text block
transformed_message["content"].append({"text": content})
elif isinstance(content, list):
# Already in blocks format - use as is
transformed_message["content"] = content
user_messages.append(transformed_message)
# Build the converse input format
converse_input = {"messages": user_messages}
# Add system messages if present
if system_messages:
converse_input["system"] = system_messages
# Build the complete request
return {"input": {"converse": converse_input}}
def _transform_to_invoke_model_format(
self, request_data: Dict[str, Any]
) -> Dict[str, Any]:
"""Transform to InvokeModel input format."""
import json
# For InvokeModel, we need to provide the raw body that would be sent to the model
# Remove the 'model' field from the body as it's not part of the model input
body_data = {k: v for k, v in request_data.items() if k != "model"}
return {"input": {"invokeModel": {"body": json.dumps(body_data)}}}
def get_bedrock_count_tokens_endpoint(
self, model: str, aws_region_name: str
) -> str:
"""
Construct the AWS Bedrock CountTokens API endpoint using existing LiteLLM functions.
Args:
model: The resolved model ID from router lookup
aws_region_name: AWS region (e.g., "eu-west-1")
Returns:
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)
# Remove bedrock/ prefix if present
if model_id.startswith("bedrock/"):
model_id = model_id[8:] # Remove "bedrock/" prefix
base_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com"
endpoint = f"{base_url}/model/{model_id}/count-tokens"
return endpoint
def transform_bedrock_response_to_anthropic(
self, bedrock_response: Dict[str, Any]
) -> Dict[str, Any]:
"""
Transform Bedrock CountTokens response to Anthropic format.
Input (Bedrock response):
{
"inputTokens": 123
}
Output (Anthropic format):
{
"input_tokens": 123
}
"""
input_tokens = bedrock_response.get("inputTokens", 0)
return {"input_tokens": input_tokens}
def validate_count_tokens_request(self, request_data: Dict[str, Any]) -> None:
"""
Validate the incoming count tokens request.
Supports both Converse and InvokeModel input formats.
Args:
request_data: The request payload
Raises:
ValueError: If the request is invalid
"""
if not request_data.get("model"):
raise ValueError("model parameter is required")
input_type = self._detect_input_type(request_data)
if input_type == "converse":
# Validate Converse format (messages-based)
messages = request_data.get("messages", [])
if not messages:
raise ValueError("messages parameter is required for Converse input")
if not isinstance(messages, list):
raise ValueError("messages must be a list")
for i, message in enumerate(messages):
if not isinstance(message, dict):
raise ValueError(f"Message {i} must be a dictionary")
if "role" not in message:
raise ValueError(f"Message {i} must have a 'role' field")
if "content" not in message:
raise ValueError(f"Message {i} must have a 'content' field")
else:
# For InvokeModel format, we need at least some content to count tokens
# The content structure varies by model, so we do minimal validation
if len(request_data) <= 1: # Only has 'model' field
raise ValueError("Request must contain content to count tokens")
litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py
+37 -15
View File
@@ -10,7 +10,7 @@ Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-tit
"""
import types
from typing import List, Optional
from typing import List, Optional, Union
from litellm.types.llms.bedrock import (
AmazonTitanV2EmbeddingRequest,
@@ -30,9 +30,7 @@ class AmazonTitanV2Config:
normalize: Optional[bool] = None
dimensions: Optional[int] = None
def __init__(
self, normalize: Optional[bool] = None, dimensions: Optional[int] = None
) -> None:
def __init__(self, normalize: Optional[bool] = None, dimensions: Optional[int] = None) -> None:
locals_ = locals().copy()
for key, value in locals_.items():
if key != "self" and value is not None:
@@ -57,32 +55,56 @@ class AmazonTitanV2Config:
}
def get_supported_openai_params(self) -> List[str]:
return ["dimensions"]
return ["dimensions", "encoding_format"]
def map_openai_params(
self, non_default_params: dict, optional_params: dict
) -> dict:
def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict:
for k, v in non_default_params.items():
if k == "dimensions":
optional_params["dimensions"] = v
elif k == "encoding_format":
# Map OpenAI encoding_format to AWS embeddingTypes
if v == "float":
optional_params["embeddingTypes"] = ["float"]
elif v == "base64":
# base64 maps to binary format in AWS
optional_params["embeddingTypes"] = ["binary"]
else:
# For any other encoding format, default to float
optional_params["embeddingTypes"] = ["float"]
return optional_params
def _transform_request(
self, input: str, inference_params: dict
) -> AmazonTitanV2EmbeddingRequest:
def _transform_request(self, input: str, inference_params: dict) -> AmazonTitanV2EmbeddingRequest:
return AmazonTitanV2EmbeddingRequest(inputText=input, **inference_params) # type: ignore
def _transform_response(
self, response_list: List[dict], model: str
) -> EmbeddingResponse:
def _transform_response(self, response_list: List[dict], model: str) -> EmbeddingResponse:
total_prompt_tokens = 0
transformed_responses: List[Embedding] = []
for index, response in enumerate(response_list):
_parsed_response = AmazonTitanV2EmbeddingResponse(**response) # type: ignore
# According to AWS docs, embeddingsByType is always present
# If binary was requested (encoding_format="base64"), use binary data
# Otherwise, use float data from embeddingsByType or fallback to embedding field
embedding_data: Union[List[float], List[int]]
if ("embeddingsByType" in _parsed_response and
"binary" in _parsed_response["embeddingsByType"]):
# Use binary data if available (for encoding_format="base64")
embedding_data = _parsed_response["embeddingsByType"]["binary"]
elif ("embeddingsByType" in _parsed_response and
"float" in _parsed_response["embeddingsByType"]):
# Use float data from embeddingsByType
embedding_data = _parsed_response["embeddingsByType"]["float"]
elif "embedding" in _parsed_response:
# Fallback to legacy embedding field
embedding_data = _parsed_response["embedding"]
else:
raise ValueError(f"No embedding data found in response: {response}")
transformed_responses.append(
Embedding(
embedding=_parsed_response["embedding"],
embedding=embedding_data,
index=index,
object="embedding",
)
litellm/llms/bedrock/embed/embedding.py
+95 -76
View File
@@ -4,12 +4,13 @@ Handles embedding calls to Bedrock's `/invoke` endpoint
import copy
import json
from typing import Any, Callable, List, Optional, Tuple, Union
import urllib.parse
from typing import Any, Callable, List, Optional, Tuple, Union, get_args
import httpx
import litellm
from litellm.constants import BEDROCK_EMBEDDING_PROVIDERS_LITERAL
from litellm.llms.cohere.embed.handler import embedding as cohere_embedding
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
@@ -18,7 +19,11 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
)
from litellm.secret_managers.main import get_secret
from litellm.types.llms.bedrock import AmazonEmbeddingRequest, CohereEmbeddingRequest
from litellm.types.llms.bedrock import (
AmazonEmbeddingRequest,
CohereEmbeddingRequest,
TwelveLabsMarengoEmbeddingRequest,
)
from litellm.types.utils import EmbeddingResponse
from ..base_aws_llm import BaseAWSLLM
@@ -29,6 +34,7 @@ from .amazon_titan_multimodal_transformation import (
)
from .amazon_titan_v2_transformation import AmazonTitanV2Config
from .cohere_transformation import BedrockCohereEmbeddingConfig
from .twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig
class BedrockEmbedding(BaseAWSLLM):
@@ -145,6 +151,44 @@ class BedrockEmbedding(BaseAWSLLM):
raise BedrockError(status_code=408, message="Timeout error occurred.")
return response.json()
def _transform_response(
self, response_list: List[dict], model: str, provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL
) -> Optional[EmbeddingResponse]:
"""
Transforms the response from the Bedrock embedding provider to the OpenAI format.
"""
returned_response: Optional[EmbeddingResponse] = None
if model == "amazon.titan-embed-image-v1":
returned_response = (
AmazonTitanMultimodalEmbeddingG1Config()._transform_response(
response_list=response_list, model=model
)
)
elif model == "amazon.titan-embed-text-v1":
returned_response = AmazonTitanG1Config()._transform_response(
response_list=response_list, model=model
)
elif model == "amazon.titan-embed-text-v2:0":
returned_response = AmazonTitanV2Config()._transform_response(
response_list=response_list, model=model
)
elif provider == "twelvelabs":
returned_response = TwelveLabsMarengoEmbeddingConfig()._transform_response(
response_list=response_list, model=model
)
##########################################################
# Validate returned response
##########################################################
if returned_response is None:
raise Exception(
"Unable to map model response to known provider format. model={}".format(
model
)
)
return returned_response
def _single_func_embeddings(
self,
@@ -157,6 +201,7 @@ class BedrockEmbedding(BaseAWSLLM):
aws_region_name: str,
model: str,
logging_obj: Any,
provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
api_key: Optional[str] = None,
):
responses: List[dict] = []
@@ -164,16 +209,16 @@ class BedrockEmbedding(BaseAWSLLM):
headers = {"Content-Type": "application/json"}
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
prepped = self.get_request_headers(
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=endpoint_url,
data=json.dumps(data),
headers=headers,
api_key=api_key
)
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=endpoint_url,
data=json.dumps(data),
headers=headers,
api_key=api_key,
)
## LOGGING
logging_obj.pre_call(
@@ -203,32 +248,9 @@ class BedrockEmbedding(BaseAWSLLM):
responses.append(response)
returned_response: Optional[EmbeddingResponse] = None
## TRANSFORM RESPONSE ##
if model == "amazon.titan-embed-image-v1":
returned_response = (
AmazonTitanMultimodalEmbeddingG1Config()._transform_response(
response_list=responses, model=model
)
)
elif model == "amazon.titan-embed-text-v1":
returned_response = AmazonTitanG1Config()._transform_response(
response_list=responses, model=model
)
elif model == "amazon.titan-embed-text-v2:0":
returned_response = AmazonTitanV2Config()._transform_response(
response_list=responses, model=model
)
if returned_response is None:
raise Exception(
"Unable to map model response to known provider format. model={}".format(
model
)
)
return returned_response
return self._transform_response(
response_list=responses, model=model, provider=provider
)
async def _async_single_func_embeddings(
self,
@@ -241,6 +263,7 @@ class BedrockEmbedding(BaseAWSLLM):
aws_region_name: str,
model: str,
logging_obj: Any,
provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
api_key: Optional[str] = None,
):
responses: List[dict] = []
@@ -248,16 +271,16 @@ class BedrockEmbedding(BaseAWSLLM):
headers = {"Content-Type": "application/json"}
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
prepped = self.get_request_headers(
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=endpoint_url,
data=json.dumps(data),
headers=headers,
api_key=api_key,
)
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=endpoint_url,
data=json.dumps(data),
headers=headers,
api_key=api_key,
)
## LOGGING
logging_obj.pre_call(
@@ -286,33 +309,10 @@ class BedrockEmbedding(BaseAWSLLM):
)
responses.append(response)
returned_response: Optional[EmbeddingResponse] = None
## TRANSFORM RESPONSE ##
if model == "amazon.titan-embed-image-v1":
returned_response = (
AmazonTitanMultimodalEmbeddingG1Config()._transform_response(
response_list=responses, model=model
)
)
elif model == "amazon.titan-embed-text-v1":
returned_response = AmazonTitanG1Config()._transform_response(
response_list=responses, model=model
)
elif model == "amazon.titan-embed-text-v2:0":
returned_response = AmazonTitanV2Config()._transform_response(
response_list=responses, model=model
)
if returned_response is None:
raise Exception(
"Unable to map model response to known provider format. model={}".format(
model
)
)
return returned_response
return self._transform_response(
response_list=responses, model=model, provider=provider
)
def embeddings(
self,
@@ -336,7 +336,7 @@ class BedrockEmbedding(BaseAWSLLM):
### TRANSFORMATION ###
unencoded_model_id = (
optional_params.pop("model_id", None) or model
) # default to model if not passed
) # default to model if not passed
modelId = urllib.parse.quote(unencoded_model_id, safe="")
aws_region_name = self._get_aws_region_name(
optional_params=optional_params,
@@ -344,7 +344,12 @@ class BedrockEmbedding(BaseAWSLLM):
model_id=unencoded_model_id,
)
provider = model.split(".")[0]
provider = self.get_bedrock_embedding_provider(model)
if provider is None:
raise Exception(
f"Unable to determine bedrock embedding provider for model: {model}. "
f"Supported providers: {list(get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL))}"
)
inference_params = copy.deepcopy(optional_params)
inference_params = {
k: v
@@ -394,6 +399,15 @@ class BedrockEmbedding(BaseAWSLLM):
)
)
batch_data.append(transformed_request)
elif provider == "twelvelabs":
batch_data = []
for i in input:
twelvelabs_request: (
TwelveLabsMarengoEmbeddingRequest
) = TwelveLabsMarengoEmbeddingConfig()._transform_request(
input=i, inference_params=inference_params
)
batch_data.append(twelvelabs_request)
### SET RUNTIME ENDPOINT ###
endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint(
@@ -422,8 +436,9 @@ class BedrockEmbedding(BaseAWSLLM):
model=model,
logging_obj=logging_obj,
api_key=api_key,
provider=provider,
)
return self._single_func_embeddings(
returned_response = self._single_func_embeddings(
client=(
client
if client is not None and isinstance(client, HTTPHandler)
@@ -438,14 +453,18 @@ class BedrockEmbedding(BaseAWSLLM):
model=model,
logging_obj=logging_obj,
api_key=api_key,
provider=provider,
)
if returned_response is None:
raise Exception("Unable to map Bedrock request to provider")
return returned_response
elif data is None:
raise Exception("Unable to map Bedrock request to provider")
headers = {"Content-Type": "application/json"}
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
prepped = self.get_request_headers(
credentials=credentials,
aws_region_name=aws_region_name,
litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py
+140
View File
@@ -0,0 +1,140 @@
"""
Transformation logic from OpenAI /v1/embeddings format to Bedrock TwelveLabs Marengo /invoke format.
Why separate file? Make it easy to see how transformation works
Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo.html
"""
from typing import List
from litellm.types.llms.bedrock import (
TwelveLabsMarengoEmbeddingRequest,
)
from litellm.types.utils import Embedding, EmbeddingResponse, Usage
from litellm.utils import get_base64_str, is_base64_encoded
class TwelveLabsMarengoEmbeddingConfig:
"""
Reference - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo.html
Supports text and image inputs for Phase 1.
Video and audio support will be added in Phase 2.
"""
def __init__(self) -> None:
pass
def get_supported_openai_params(self) -> List[str]:
return ["encoding_format", "textTruncate", "embeddingOption"]
def map_openai_params(
self, non_default_params: dict, optional_params: dict
) -> dict:
for k, v in non_default_params.items():
if k == "encoding_format":
# TwelveLabs doesn't have encoding_format, but we can map it to embeddingOption
if v == "float":
optional_params["embeddingOption"] = ["visual-text", "visual-image"]
elif k == "textTruncate":
optional_params["textTruncate"] = v
elif k == "embeddingOption":
optional_params["embeddingOption"] = v
return optional_params
def _transform_request(
self, input: str, inference_params: dict
) -> TwelveLabsMarengoEmbeddingRequest:
"""
Transform OpenAI-style input to TwelveLabs Marengo format.
Phase 1: Supports text and image inputs only.
"""
# Check if input is base64 encoded image
is_encoded = is_base64_encoded(input)
if is_encoded:
# Image input
b64_str = get_base64_str(input)
transformed_request = TwelveLabsMarengoEmbeddingRequest(
inputType="image", mediaSource={"base64String": b64_str}
)
else:
# Text input
transformed_request = TwelveLabsMarengoEmbeddingRequest(
inputType="text", inputText=input
)
# Set default textTruncate if not specified
if "textTruncate" not in inference_params:
transformed_request["textTruncate"] = "end"
# Apply any additional inference parameters
for k, v in inference_params.items():
if k not in [
"inputType",
"inputText",
"mediaSource",
]: # Don't override core fields
transformed_request[k] = v # type: ignore
return transformed_request
def _transform_response(
self, response_list: List[dict], model: str
) -> EmbeddingResponse:
"""
Transform TwelveLabs response to OpenAI format.
Handles the actual TwelveLabs response format: {"data": [{"embedding": [...]}]}
"""
embeddings: List[Embedding] = []
total_tokens = 0
for response in response_list:
# TwelveLabs response format has a "data" field containing the embeddings
if "data" in response and isinstance(response["data"], list):
for item in response["data"]:
if "embedding" in item:
# Single embedding response
embedding = Embedding(
embedding=item["embedding"],
index=len(embeddings),
object="embedding",
)
embeddings.append(embedding)
# Estimate token count (rough approximation)
if "inputTextTokenCount" in item:
total_tokens += item["inputTextTokenCount"]
else:
# Rough estimate: 1 token per 4 characters for text, or use embedding size
total_tokens += len(item["embedding"]) // 4
elif "embedding" in response:
# Direct embedding response (fallback for other formats)
embedding = Embedding(
embedding=response["embedding"],
index=len(embeddings),
object="embedding",
)
embeddings.append(embedding)
# Estimate token count (rough approximation)
if "inputTextTokenCount" in response:
total_tokens += response["inputTextTokenCount"]
else:
# Rough estimate: 1 token per 4 characters for text
total_tokens += len(response.get("inputText", "")) // 4
elif "embeddings" in response:
# Multiple embeddings response (from video/audio)
for i, emb in enumerate(response["embeddings"]):
embedding = Embedding(
embedding=emb["embedding"],
index=len(embeddings),
object="embedding",
)
embeddings.append(embedding)
total_tokens += len(emb["embedding"]) // 4 # Rough estimate
usage = Usage(prompt_tokens=total_tokens, total_tokens=total_tokens)
return EmbeddingResponse(data=embeddings, model=model, usage=usage)
litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py
+7 -2
View File
@@ -7,12 +7,12 @@ from litellm.types.llms.bedrock import (
AmazonNovaCanvasColorGuidedGenerationParams,
AmazonNovaCanvasColorGuidedRequest,
AmazonNovaCanvasImageGenerationConfig,
AmazonNovaCanvasInpaintingParams,
AmazonNovaCanvasInpaintingRequest,
AmazonNovaCanvasRequestBase,
AmazonNovaCanvasTextToImageParams,
AmazonNovaCanvasTextToImageRequest,
AmazonNovaCanvasTextToImageResponse,
AmazonNovaCanvasInpaintingParams,
AmazonNovaCanvasInpaintingRequest,
)
from litellm.types.utils import ImageResponse
@@ -67,6 +67,11 @@ class AmazonNovaCanvasConfig:
"""
task_type = optional_params.pop("taskType", "TEXT_IMAGE")
image_generation_config = optional_params.pop("imageGenerationConfig", {})
# Extract model_id parameter to prevent "extraneous key" error from Bedrock API
# Following the same pattern as chat completions and embeddings
unencoded_model_id = optional_params.pop("model_id", None) # noqa: F841
image_generation_config = {**image_generation_config, **optional_params}
if task_type == "TEXT_IMAGE":
text_to_image_params: Dict[str, Any] = image_generation_config.pop(
litellm/llms/bedrock/image/image_handler.py
+11 -1
View File
@@ -233,7 +233,17 @@ class BedrockImageGeneration(BaseAWSLLM):
Returns:
dict: The request body to use for the Bedrock Image Generation API
"""
provider = model.split(".")[0]
# Use the existing ARN-aware provider detection method
bedrock_provider = self.get_bedrock_invoke_provider(model)
if bedrock_provider == "amazon" or bedrock_provider == "nova":
# Handle Amazon Nova Canvas models
provider = "amazon"
elif bedrock_provider == "stability":
provider = "stability"
else:
# Fallback to original logic for backward compatibility
provider = model.split(".")[0]
inference_params = copy.deepcopy(optional_params)
inference_params.pop(
"user", None
litellm/llms/cohere/completion/handler.py
-5
View File
@@ -1,5 +0,0 @@
"""
Cohere /generate API - uses `llm_http_handler.py` to make httpx requests
Request/Response transformation is handled in `transformation.py`
"""
litellm/llms/cohere/completion/transformation.py
-265
View File
@@ -1,265 +0,0 @@
import time
from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, List, Optional, Union
import httpx
import litellm
from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_content_list_to_str,
)
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import Choices, Message, ModelResponse, Usage
from ..common_utils import CohereError
from ..common_utils import ModelResponseIterator as CohereModelResponseIterator
from ..common_utils import validate_environment as cohere_validate_environment
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class CohereTextConfig(BaseConfig):
"""
Reference: https://docs.cohere.com/reference/generate
The class `CohereConfig` provides configuration for the Cohere's API interface. Below are the parameters:
- `num_generations` (integer): Maximum number of generations returned. Default is 1, with a minimum value of 1 and a maximum value of 5.
- `max_tokens` (integer): Maximum number of tokens the model will generate as part of the response. Default value is 20.
- `truncate` (string): Specifies how the API handles inputs longer than maximum token length. Options include NONE, START, END. Default is END.
- `temperature` (number): A non-negative float controlling the randomness in generation. Lower temperatures result in less random generations. Default is 0.75.
- `preset` (string): Identifier of a custom preset, a combination of parameters such as prompt, temperature etc.
- `end_sequences` (array of strings): The generated text gets cut at the beginning of the earliest occurrence of an end sequence, which will be excluded from the text.
- `stop_sequences` (array of strings): The generated text gets cut at the end of the earliest occurrence of a stop sequence, which will be included in the text.
- `k` (integer): Limits generation at each step to top `k` most likely tokens. Default is 0.
- `p` (number): Limits generation at each step to most likely tokens with total probability mass of `p`. Default is 0.
- `frequency_penalty` (number): Reduces repetitiveness of generated tokens. Higher values apply stronger penalties to previously occurred tokens.
- `presence_penalty` (number): Reduces repetitiveness of generated tokens. Similar to frequency_penalty, but this penalty applies equally to all tokens that have already appeared.
- `return_likelihoods` (string): Specifies how and if token likelihoods are returned with the response. Options include GENERATION, ALL and NONE.
- `logit_bias` (object): Used to prevent the model from generating unwanted tokens or to incentivize it to include desired tokens. e.g. {"hello_world": 1233}
"""
num_generations: Optional[int] = None
max_tokens: Optional[int] = None
truncate: Optional[str] = None
temperature: Optional[int] = None
preset: Optional[str] = None
end_sequences: Optional[list] = None
stop_sequences: Optional[list] = None
k: Optional[int] = None
p: Optional[int] = None
frequency_penalty: Optional[int] = None
presence_penalty: Optional[int] = None
return_likelihoods: Optional[str] = None
logit_bias: Optional[dict] = None
def __init__(
self,
num_generations: Optional[int] = None,
max_tokens: Optional[int] = None,
truncate: Optional[str] = None,
temperature: Optional[int] = None,
preset: Optional[str] = None,
end_sequences: Optional[list] = None,
stop_sequences: Optional[list] = None,
k: Optional[int] = None,
p: Optional[int] = None,
frequency_penalty: Optional[int] = None,
presence_penalty: Optional[int] = None,
return_likelihoods: Optional[str] = None,
logit_bias: Optional[dict] = None,
) -> None:
locals_ = locals().copy()
for key, value in locals_.items():
if key != "self" and value is not None:
setattr(self.__class__, key, value)
@classmethod
def get_config(cls):
return super().get_config()
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
return cohere_validate_environment(
headers=headers,
model=model,
messages=messages,
optional_params=optional_params,
api_key=api_key,
)
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BaseLLMException:
return CohereError(status_code=status_code, message=error_message)
def get_supported_openai_params(self, model: str) -> List:
return [
"stream",
"temperature",
"max_tokens",
"logit_bias",
"top_p",
"frequency_penalty",
"presence_penalty",
"stop",
"n",
"extra_headers",
]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
for param, value in non_default_params.items():
if param == "stream":
optional_params["stream"] = value
elif param == "temperature":
optional_params["temperature"] = value
elif param == "max_tokens":
optional_params["max_tokens"] = value
elif param == "n":
optional_params["num_generations"] = value
elif param == "logit_bias":
optional_params["logit_bias"] = value
elif param == "top_p":
optional_params["p"] = value
elif param == "frequency_penalty":
optional_params["frequency_penalty"] = value
elif param == "presence_penalty":
optional_params["presence_penalty"] = value
elif param == "stop":
optional_params["stop_sequences"] = value
return optional_params
def transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
prompt = " ".join(
convert_content_list_to_str(message=message) for message in messages
)
## Load Config
config = litellm.CohereConfig.get_config()
for k, v in config.items():
if (
k not in optional_params
): # completion(top_k=3) > cohere_config(top_k=3) <- allows for dynamic variables to be passed in
optional_params[k] = v
## Handle Tool Calling
if "tools" in optional_params:
_is_function_call = True
tool_calling_system_prompt = self._construct_cohere_tool_for_completion_api(
tools=optional_params["tools"]
)
optional_params["tools"] = tool_calling_system_prompt
data = {
"model": model,
"prompt": prompt,
**optional_params,
}
return data
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:
prompt = " ".join(
convert_content_list_to_str(message=message) for message in messages
)
completion_response = raw_response.json()
choices_list = []
for idx, item in enumerate(completion_response["generations"]):
if len(item["text"]) > 0:
message_obj = Message(content=item["text"])
else:
message_obj = Message(content=None)
choice_obj = Choices(
finish_reason=item["finish_reason"],
index=idx + 1,
message=message_obj,
)
choices_list.append(choice_obj)
model_response.choices = choices_list # type: ignore
## CALCULATING USAGE
prompt_tokens = len(encoding.encode(prompt))
completion_tokens = len(
encoding.encode(model_response["choices"][0]["message"].get("content", ""))
)
model_response.created = int(time.time())
model_response.model = model
usage = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
)
setattr(model_response, "usage", usage)
return model_response
def _construct_cohere_tool_for_completion_api(
self,
tools: Optional[List] = None,
) -> dict:
if tools is None:
tools = []
return {"tools": tools}
def get_model_response_iterator(
self,
streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse],
sync_stream: bool,
json_mode: Optional[bool] = False,
):
return CohereModelResponseIterator(
streaming_response=streaming_response,
sync_stream=sync_stream,
json_mode=json_mode,
)
litellm/llms/custom_httpx/llm_http_handler.py
+172
View File
@@ -2520,6 +2520,95 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params_with_request,
)
def retrieve_batch(
self,
batch_id: str,
litellm_params: dict,
provider_config: "BaseBatchesConfig",
headers: dict,
api_base: Optional[str],
api_key: Optional[str],
logging_obj: "LiteLLMLoggingObj",
_is_async: bool = False,
client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
model: Optional[str] = None,
) -> Union["LiteLLMBatch", Coroutine[Any, Any, "LiteLLMBatch"]]:
"""
Retrieve a batch using provider-specific configuration.
"""
# Transform the request using provider config
transformed_request = provider_config.transform_retrieve_batch_request(
batch_id=batch_id,
optional_params=litellm_params,
litellm_params=litellm_params,
)
if _is_async:
return self.async_retrieve_batch(
transformed_request=transformed_request,
litellm_params=litellm_params,
provider_config=provider_config,
headers=headers,
api_base=api_base,
logging_obj=logging_obj,
client=client,
timeout=timeout,
batch_id=batch_id,
model=model,
)
if client is None or not isinstance(client, HTTPHandler):
sync_httpx_client = _get_httpx_client()
else:
sync_httpx_client = client
try:
if (
isinstance(transformed_request, dict)
and "method" in transformed_request
):
# Handle pre-signed requests (e.g., from Bedrock with AWS auth)
method = transformed_request["method"].lower()
request_kwargs = {
"url": transformed_request["url"],
"headers": transformed_request["headers"],
}
# Only add data for non-GET requests
if method != "get" and transformed_request.get("data") is not None:
request_kwargs["data"] = transformed_request["data"]
batch_response = getattr(sync_httpx_client, method)(**request_kwargs)
elif isinstance(transformed_request, dict) and api_base:
# For other providers that use JSON requests
batch_response = sync_httpx_client.get(
url=api_base,
headers={**headers, "Content-Type": "application/json"},
params=transformed_request,
)
else:
# Handle other request types if needed
if not api_base:
raise ValueError("api_base is required for non-pre-signed requests")
batch_response = sync_httpx_client.get(
url=api_base,
headers=headers,
)
except Exception as e:
verbose_logger.exception(f"Error retrieving batch: {e}")
raise self._handle_error(
e=e,
provider_config=provider_config,
)
return provider_config.transform_retrieve_batch_response(
model=model,
raw_response=batch_response,
logging_obj=logging_obj,
litellm_params=litellm_params,
)
async def async_create_batch(
self,
transformed_request: Union[bytes, str, dict],
@@ -2606,6 +2695,89 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params_with_request,
)
async def async_retrieve_batch(
self,
transformed_request: Union[bytes, str, dict],
litellm_params: dict,
provider_config: "BaseBatchesConfig",
headers: dict,
api_base: Optional[str],
logging_obj: "LiteLLMLoggingObj",
client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
batch_id: Optional[str] = None,
model: Optional[str] = None,
):
"""
Async version of retrieve_batch
"""
if client is None or not isinstance(client, AsyncHTTPHandler):
async_httpx_client = get_async_httpx_client(
llm_provider=provider_config.custom_llm_provider
)
else:
async_httpx_client = client
#########################################################
# Debug Logging
#########################################################
logging_obj.pre_call(
input="",
api_key="",
additional_args={
"complete_input_dict": transformed_request,
"api_base": api_base,
"headers": headers,
"batch_id": batch_id,
},
)
try:
if (
isinstance(transformed_request, dict)
and "method" in transformed_request
):
# Handle pre-signed requests (e.g., from Bedrock with AWS auth)
method = transformed_request["method"].lower()
request_kwargs = {
"url": transformed_request["url"],
"headers": transformed_request["headers"],
}
# Only add data for non-GET requests
if method != "get" and transformed_request.get("data") is not None:
request_kwargs["data"] = transformed_request["data"]
batch_response = await getattr(async_httpx_client, method)(**request_kwargs)
elif isinstance(transformed_request, dict) and api_base:
# For other providers that use JSON requests
batch_response = await async_httpx_client.get(
url=api_base,
headers={**headers, "Content-Type": "application/json"},
params=transformed_request,
)
else:
# Handle other request types if needed
if not api_base:
raise ValueError("api_base is required for non-pre-signed requests")
batch_response = await async_httpx_client.get(
url=api_base,
headers=headers,
)
except Exception as e:
verbose_logger.exception(f"Error retrieving batch: {e}")
raise self._handle_error(
e=e,
provider_config=provider_config,
)
return provider_config.transform_retrieve_batch_response(
model=model,
raw_response=batch_response,
logging_obj=logging_obj,
litellm_params=litellm_params,
)
def cancel_response_api_handler(
self,
response_id: str,
litellm/llms/gemini/google_genai/transformation.py
+2 -1
View File
@@ -121,7 +121,8 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM):
default_headers = {
"Content-Type": "application/json",
}
gemini_api_key = self._get_google_ai_studio_api_key(dict(litellm_params or {}))
# Use the passed api_key first, then fall back to litellm_params and environment
gemini_api_key = api_key or self._get_google_ai_studio_api_key(dict(litellm_params or {}))
if gemini_api_key is not None:
default_headers[self.XGOOGLE_API_KEY] = gemini_api_key
if headers is not None:
litellm/llms/gemini/image_generation/transformation.py
+78 -37
View File
@@ -85,17 +85,25 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
) -> str:
"""
Get the complete url for the request
Google AI API format: https://generativelanguage.googleapis.com/v1beta/models/{model}:predict
Gemini 2.5 Flash Image Preview: :generateContent
Other Imagen models: :predict
"""
complete_url: str = (
api_base
or get_secret_str("GEMINI_API_BASE")
api_base
or get_secret_str("GEMINI_API_BASE")
or self.DEFAULT_BASE_URL
)
complete_url = complete_url.rstrip("/")
complete_url = f"{complete_url}/models/{model}:predict"
# Gemini 2.5 Flash Image Preview uses generateContent endpoint
if "2.5-flash-image-preview" in model:
complete_url = f"{complete_url}/models/{model}:generateContent"
else:
# All other Imagen models use predict endpoint
complete_url = f"{complete_url}/models/{model}:predict"
return complete_url
def validate_environment(
@@ -128,35 +136,52 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
headers: dict,
) -> dict:
"""
Transform the image generation request to Google AI Imagen format
Google AI API format:
Transform the image generation request to Gemini format
For Gemini 2.5 Flash Image Preview, use the standard Gemini format with response_modalities:
{
"instances": [
"contents": [
{
"prompt": "Robot holding a red skateboard"
"parts": [
{"text": "Generate an image of..."}
]
}
],
"parameters": {
"sampleCount": 4,
"aspectRatio": "1:1",
"personGeneration": "allow_adult"
"generationConfig": {
"response_modalities": ["IMAGE", "TEXT"]
}
}
"""
from litellm.types.llms.gemini import (
GeminiImageGenerationInstance,
GeminiImageGenerationParameters,
)
request_body: GeminiImageGenerationRequest = GeminiImageGenerationRequest(
instances=[
GeminiImageGenerationInstance(
prompt=prompt
)
],
parameters=GeminiImageGenerationParameters(**optional_params)
)
return request_body.model_dump(exclude_none=True)
# For Gemini 2.5 Flash Image Preview, use standard Gemini format
if "2.5-flash-image-preview" in model:
request_body: dict = {
"contents": [
{
"parts": [
{"text": prompt}
]
}
],
"generationConfig": {
"response_modalities": ["IMAGE", "TEXT"]
}
}
return request_body
else:
# For other Imagen models, use the original Imagen format
from litellm.types.llms.gemini import (
GeminiImageGenerationInstance,
GeminiImageGenerationParameters,
)
request_body_obj: GeminiImageGenerationRequest = GeminiImageGenerationRequest(
instances=[
GeminiImageGenerationInstance(
prompt=prompt
)
],
parameters=GeminiImageGenerationParameters(**optional_params)
)
return request_body_obj.model_dump(exclude_none=True)
def transform_image_generation_response(
self,
@@ -185,14 +210,30 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
if not model_response.data:
model_response.data = []
# Google AI returns predictions with generated images
predictions = response_data.get("predictions", [])
for prediction in predictions:
# Google AI returns base64 encoded images in the prediction
model_response.data.append(ImageObject(
b64_json=prediction.get("bytesBase64Encoded", None),
url=None, # Google AI returns base64, not URLs
))
# Handle different response formats based on model
if "2.5-flash-image-preview" in model:
# Gemini 2.5 Flash Image Preview returns in candidates format
candidates = response_data.get("candidates", [])
for candidate in candidates:
content = candidate.get("content", {})
parts = content.get("parts", [])
for part in parts:
# Look for inlineData with image
if "inlineData" in part:
inline_data = part["inlineData"]
if "data" in inline_data:
model_response.data.append(ImageObject(
b64_json=inline_data["data"],
url=None,
))
else:
# Original Imagen format - predictions with generated images
predictions = response_data.get("predictions", [])
for prediction in predictions:
# Google AI returns base64 encoded images in the prediction
model_response.data.append(ImageObject(
b64_json=prediction.get("bytesBase64Encoded", None),
url=None, # Google AI returns base64, not URLs
))
return model_response
litellm/llms/oci/chat/transformation.py
+1 -1
View File
@@ -378,7 +378,7 @@ class OCIChatConfig(BaseConfig):
or not oci_compartment_id
):
raise Exception(
"Missing required parameters: oci_user, oci_fingerprint, oci_tenancy, "
"Missing required parameters: oci_user, oci_fingerprint, oci_tenancy, oci_compartment_id "
"and at least one of oci_key or oci_key_file."
)
litellm/llms/vertex_ai/batches/transformation.py
+8 -1
View File
@@ -114,7 +114,14 @@ class VertexAIBatchTransformation:
"""
Gets the output file id from the Vertex AI Batch response
"""
output_file_id: str = ""
output_file_id: str = (
response.get("outputInfo", OutputInfo()).get("gcsOutputDirectory", "")
+ "/predictions.jsonl"
)
if output_file_id != "/predictions.jsonl":
return output_file_id
output_config = response.get("outputConfig")
if output_config is None:
return output_file_id
litellm/llms/vertex_ai/files/handler.py
+141 -2
View File
@@ -1,5 +1,6 @@
import asyncio
from typing import Any, Coroutine, Optional, Union
import urllib.parse
from typing import Any, Coroutine, Optional, Tuple, Union
import httpx
@@ -9,7 +10,12 @@ from litellm.integrations.gcs_bucket.gcs_bucket_base import (
GCSLoggingConfig,
)
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.openai import CreateFileRequest, OpenAIFileObject
from litellm.types.llms.openai import (
CreateFileRequest,
FileContentRequest,
HttpxBinaryResponseContent,
OpenAIFileObject,
)
from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
from .transformation import VertexAIJsonlFilesTransformation
@@ -105,3 +111,136 @@ class VertexAIFilesHandler(GCSBucketBase):
max_retries=max_retries,
)
)
def _extract_bucket_and_object_from_file_id(self, file_id: str) -> Tuple[str, str]:
"""
Extract bucket name and object path from URL-encoded file_id.
Expected format: gs%3A%2F%2Fbucket-name%2Fpath%2Fto%2Ffile
Which decodes to: gs://bucket-name/path/to/file
Returns:
tuple: (bucket_name, url_encoded_object_path)
- bucket_name: "bucket-name"
- url_encoded_object_path: "path%2Fto%2Ffile"
"""
decoded_path = urllib.parse.unquote(file_id)
if decoded_path.startswith("gs://"):
full_path = decoded_path[5:] # Remove 'gs://' prefix
else:
full_path = decoded_path
if "/" in full_path:
bucket_name, object_path = full_path.split("/", 1)
else:
bucket_name = full_path
object_path = ""
encoded_object_path = urllib.parse.quote(object_path, safe="")
return bucket_name, encoded_object_path
async def afile_content(
self,
file_content_request: FileContentRequest,
vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES],
vertex_project: Optional[str],
vertex_location: Optional[str],
timeout: Union[float, httpx.Timeout],
max_retries: Optional[int],
) -> HttpxBinaryResponseContent:
"""
Download file content from GCS bucket for VertexAI files.
Args:
file_content_request: Contains file_id (URL-encoded GCS path)
vertex_credentials: VertexAI credentials
vertex_project: VertexAI project ID
vertex_location: VertexAI location
timeout: Request timeout
max_retries: Max retry attempts
Returns:
HttpxBinaryResponseContent: Binary content wrapped in compatible response format
"""
file_id = file_content_request.get("file_id")
if not file_id:
raise ValueError("file_id is required in file_content_request")
bucket_name, encoded_object_path = self._extract_bucket_and_object_from_file_id(
file_id
)
download_kwargs = {
"standard_callback_dynamic_params": {"gcs_bucket_name": bucket_name}
}
file_content = await self.download_gcs_object(
object_name=encoded_object_path, **download_kwargs
)
if file_content is None:
decoded_path = urllib.parse.unquote(file_id)
raise ValueError(f"Failed to download file from GCS: {decoded_path}")
decoded_path = urllib.parse.unquote(file_id)
mock_response = httpx.Response(
status_code=200,
content=file_content,
headers={"content-type": "application/octet-stream"},
request=httpx.Request(method="GET", url=decoded_path),
)
return HttpxBinaryResponseContent(response=mock_response)
def file_content(
self,
_is_async: bool,
file_content_request: FileContentRequest,
api_base: Optional[str],
vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES],
vertex_project: Optional[str],
vertex_location: Optional[str],
timeout: Union[float, httpx.Timeout],
max_retries: Optional[int],
) -> Union[
HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent]
]:
"""
Download file content from GCS bucket for VertexAI files.
Supports both sync and async operations.
Args:
_is_async: Whether to run asynchronously
file_content_request: Contains file_id (URL-encoded GCS path)
api_base: API base (unused for GCS operations)
vertex_credentials: VertexAI credentials
vertex_project: VertexAI project ID
vertex_location: VertexAI location
timeout: Request timeout
max_retries: Max retry attempts
Returns:
HttpxBinaryResponseContent or Coroutine: Binary content wrapped in compatible response format
"""
if _is_async:
return self.afile_content(
file_content_request=file_content_request,
vertex_credentials=vertex_credentials,
vertex_project=vertex_project,
vertex_location=vertex_location,
timeout=timeout,
max_retries=max_retries,
)
else:
return asyncio.run(
self.afile_content(
file_content_request=file_content_request,
vertex_credentials=vertex_credentials,
vertex_project=vertex_project,
vertex_location=vertex_location,
timeout=timeout,
max_retries=max_retries,
)
)
litellm/llms/vertex_ai/files/transformation.py
+3 -3
View File
@@ -261,10 +261,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
raise ValueError("file is required")
extracted_file_data = extract_file_data(file_data)
extracted_file_data_content = extracted_file_data.get("content")
if extracted_file_data_content is None:
raise ValueError("file content is required")
if FilesAPIUtils.is_batch_jsonl_file(
create_file_data=create_file_data,
extracted_file_data=extracted_file_data,
@@ -283,7 +283,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
openai_jsonl_content
)
)
return json.dumps(vertex_jsonl_content)
return "\n".join(json.dumps(item) for item in vertex_jsonl_content)
elif isinstance(extracted_file_data_content, bytes):
return extracted_file_data_content
else:
litellm/llms/vertex_ai/vertex_llm_base.py
+17 -9
View File
@@ -239,6 +239,7 @@ class VertexBase:
stream=stream,
auth_header=None,
url=default_api_base,
model=model,
)
return api_base
@@ -292,6 +293,7 @@ class VertexBase:
stream: Optional[bool],
auth_header: Optional[str],
url: str,
model: Optional[str] = None,
) -> Tuple[Optional[str], str]:
"""
for cloudflare ai gateway - https://github.com/BerriAI/litellm/issues/4317
@@ -301,7 +303,12 @@ class VertexBase:
"""
if api_base:
if custom_llm_provider == "gemini":
url = "{}:{}".format(api_base, endpoint)
# For Gemini (Google AI Studio), construct the full path like other providers
if model is None:
raise ValueError(
"Model parameter is required for Gemini custom API base URLs"
)
url = "{}/models/{}:{}".format(api_base, model, endpoint)
if gemini_api_key is None:
raise ValueError(
"Missing gemini_api_key, please set `GEMINI_API_KEY`"
@@ -373,6 +380,7 @@ class VertexBase:
endpoint=endpoint,
stream=stream,
url=url,
model=model,
)
def _handle_reauthentication(
@@ -384,19 +392,19 @@ class VertexBase:
) -> Tuple[str, str]:
"""
Handle reauthentication when credentials refresh fails.
This method clears the cached credentials and attempts to reload them once.
It should only be called when "Reauthentication is needed" error occurs.
Args:
credentials: The original credentials
project_id: The project ID
credential_cache_key: The cache key to clear
error: The original error that triggered reauthentication
Returns:
Tuple of (access_token, project_id)
Raises:
The original error if reauthentication fails
"""
@@ -404,11 +412,11 @@ class VertexBase:
f"Handling reauthentication for project_id: {project_id}. "
f"Clearing cache and retrying once."
)
# Clear the cached credentials
if credential_cache_key in self._credentials_project_mapping:
del self._credentials_project_mapping[credential_cache_key]
# Retry once with _retry_reauth=True to prevent infinite recursion
try:
return self.get_access_token(
@@ -438,12 +446,12 @@ class VertexBase:
3. Check if loaded credentials have expired
4. If expired, refresh credentials
5. Return access token and project id
Args:
credentials: The credentials to use for authentication
project_id: The Google Cloud project ID
_retry_reauth: Internal flag to prevent infinite recursion during reauthentication
Returns:
Tuple of (access_token, project_id)
"""
litellm/main.py
+38 -90
View File
@@ -116,6 +116,7 @@ from litellm.utils import (
from ._logging import verbose_logger
from .caching.caching import disable_cache, enable_cache, update_cache
from .litellm_core_utils.core_helpers import safe_deep_copy
from .litellm_core_utils.fallback_utils import (
async_completion_with_fallbacks,
completion_with_fallbacks,
@@ -1003,7 +1004,15 @@ def completion( # type: ignore # noqa: PLR0915
provider_specific_header = cast(
Optional[ProviderSpecificHeader], kwargs.get("provider_specific_header", None)
)
headers = kwargs.get("headers", None) or extra_headers
# Properly merge headers with priority: request headers > extra_headers > global litellm.headers
headers = {}
if litellm.headers is not None and isinstance(litellm.headers, dict):
headers.update(litellm.headers)
if extra_headers is not None and isinstance(extra_headers, dict):
headers.update(extra_headers)
request_headers = kwargs.get("headers", None)
if request_headers is not None and isinstance(request_headers, dict):
headers.update(request_headers)
ensure_alternating_roles: Optional[bool] = kwargs.get(
"ensure_alternating_roles", None
@@ -1014,10 +1023,6 @@ def completion( # type: ignore # noqa: PLR0915
assistant_continue_message: Optional[ChatCompletionAssistantMessage] = kwargs.get(
"assistant_continue_message", None
)
if headers is None:
headers = {}
if extra_headers is not None:
headers.update(extra_headers)
num_retries = kwargs.get(
"num_retries", None
) ## alt. param for 'max_retries'. Use this to pass retries w/ instructor.
@@ -1074,7 +1079,6 @@ def completion( # type: ignore # noqa: PLR0915
prompt_id=prompt_id, non_default_params=non_default_params
)
):
(
model,
messages,
@@ -1427,8 +1431,7 @@ def completion( # type: ignore # noqa: PLR0915
"azure_ad_token_provider", None
)
headers = headers or litellm.headers
# Use the consolidated headers that were already merged at the top of the function
if extra_headers is not None:
optional_params["extra_headers"] = extra_headers
if max_retries is not None:
@@ -1693,8 +1696,7 @@ def completion( # type: ignore # noqa: PLR0915
or get_secret("OPENAI_API_KEY")
)
headers = headers or litellm.headers
# Use the consolidated headers that were already merged at the top of the function
if extra_headers is not None:
optional_params["extra_headers"] = extra_headers
@@ -2032,7 +2034,6 @@ def completion( # type: ignore # noqa: PLR0915
try:
if use_base_llm_http_handler:
response = base_llm_http_handler.completion(
model=model,
messages=messages,
@@ -2395,47 +2396,7 @@ def completion( # type: ignore # noqa: PLR0915
)
return response
response = model_response
elif custom_llm_provider == "cohere":
cohere_key = (
api_key
or litellm.cohere_key
or get_secret("COHERE_API_KEY")
or get_secret("CO_API_KEY")
or litellm.api_key
)
api_base = (
api_base
or litellm.api_base
or get_secret("COHERE_API_BASE")
or "https://api.cohere.ai/v1/generate"
)
headers = headers or litellm.headers or {}
if headers is None:
headers = {}
if extra_headers is not None:
headers.update(extra_headers)
response = base_llm_http_handler.completion(
model=model,
stream=stream,
messages=messages,
acompletion=acompletion,
api_base=api_base,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params,
custom_llm_provider="cohere",
timeout=timeout,
headers=headers,
encoding=encoding,
api_key=cohere_key,
logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements
client=client,
)
elif custom_llm_provider == "cohere_chat":
elif custom_llm_provider == "cohere_chat" or custom_llm_provider == "cohere":
cohere_key = (
api_key
or litellm.cohere_key
@@ -2451,12 +2412,8 @@ def completion( # type: ignore # noqa: PLR0915
or "https://api.cohere.ai/v1/chat"
)
headers = headers or litellm.headers or {}
if headers is None:
headers = {}
if extra_headers is not None:
headers.update(extra_headers)
# Use the consolidated headers that were already merged at the top of the function
# No need for additional merging here as it's already done
response = base_llm_http_handler.completion(
model=model,
@@ -2552,15 +2509,10 @@ def completion( # type: ignore # noqa: PLR0915
)
elif custom_llm_provider == "compactifai":
api_key = (
api_key
or get_secret_str("COMPACTIFAI_API_KEY")
or litellm.api_key
api_key or get_secret_str("COMPACTIFAI_API_KEY") or litellm.api_key
)
api_base = (
api_base
or "https://api.compactif.ai/v1"
)
api_base = api_base or "https://api.compactif.ai/v1"
## COMPLETION CALL
response = base_llm_http_handler.completion(
@@ -2848,8 +2800,7 @@ def completion( # type: ignore # noqa: PLR0915
)
api_base = api_base or litellm.api_base or get_secret("GEMINI_API_BASE")
new_params = deepcopy(optional_params)
new_params = safe_deep_copy(optional_params or {})
response = vertex_chat_completion.completion( # type: ignore
model=model,
messages=messages,
@@ -2893,7 +2844,7 @@ def completion( # type: ignore # noqa: PLR0915
api_base = api_base or litellm.api_base or get_secret("VERTEXAI_API_BASE")
new_params = deepcopy(optional_params)
new_params = safe_deep_copy(optional_params or {})
if vertex_partner_models_chat_completion.is_vertex_partner_model(model):
model_response = vertex_partner_models_chat_completion.completion(
model=model,
@@ -3147,9 +3098,9 @@ def completion( # type: ignore # noqa: PLR0915
"aws_region_name" not in optional_params
or optional_params["aws_region_name"] is None
):
optional_params["aws_region_name"] = (
aws_bedrock_client.meta.region_name
)
optional_params[
"aws_region_name"
] = aws_bedrock_client.meta.region_name
bedrock_route = BedrockModelInfo.get_bedrock_route(model)
if bedrock_route == "converse":
@@ -3491,7 +3442,6 @@ def completion( # type: ignore # noqa: PLR0915
)
raise e
elif custom_llm_provider == "gradient_ai":
api_base = litellm.api_base or api_base
response = base_llm_http_handler.completion(
model=model,
@@ -3851,7 +3801,7 @@ def embedding(
*,
aembedding: Literal[True],
**kwargs,
) -> Coroutine[Any, Any, EmbeddingResponse]:
) -> Coroutine[Any, Any, EmbeddingResponse]:
...
@@ -3877,7 +3827,7 @@ def embedding(
*,
aembedding: Literal[False] = False,
**kwargs,
) -> EmbeddingResponse:
) -> EmbeddingResponse:
...
# fmt: on
@@ -4188,10 +4138,8 @@ def embedding( # noqa: PLR0915
or litellm.api_key
)
if extra_headers is not None and isinstance(extra_headers, dict):
headers = extra_headers
else:
headers = {}
# Use the consolidated headers that were already merged at the top of the function
# No need for additional merging here as it's already done
response = base_llm_http_handler.embedding(
model=model,
@@ -5151,9 +5099,9 @@ def adapter_completion(
new_kwargs = translation_obj.translate_completion_input_params(kwargs=kwargs)
response: Union[ModelResponse, CustomStreamWrapper] = completion(**new_kwargs) # type: ignore
translated_response: Optional[Union[BaseModel, AdapterCompletionStreamWrapper]] = (
None
)
translated_response: Optional[
Union[BaseModel, AdapterCompletionStreamWrapper]
] = None
if isinstance(response, ModelResponse):
translated_response = translation_obj.translate_completion_output_params(
response=response
@@ -6141,9 +6089,9 @@ def stream_chunk_builder( # noqa: PLR0915
]
if len(content_chunks) > 0:
response["choices"][0]["message"]["content"] = (
processor.get_combined_content(content_chunks)
)
response["choices"][0]["message"][
"content"
] = processor.get_combined_content(content_chunks)
thinking_blocks = [
chunk
@@ -6154,9 +6102,9 @@ def stream_chunk_builder( # noqa: PLR0915
]
if len(thinking_blocks) > 0:
response["choices"][0]["message"]["thinking_blocks"] = (
processor.get_combined_thinking_content(thinking_blocks)
)
response["choices"][0]["message"][
"thinking_blocks"
] = processor.get_combined_thinking_content(thinking_blocks)
reasoning_chunks = [
chunk
@@ -6167,9 +6115,9 @@ def stream_chunk_builder( # noqa: PLR0915
]
if len(reasoning_chunks) > 0:
response["choices"][0]["message"]["reasoning_content"] = (
processor.get_combined_reasoning_content(reasoning_chunks)
)
response["choices"][0]["message"][
"reasoning_content"
] = processor.get_combined_reasoning_content(reasoning_chunks)
audio_chunks = [
chunk
litellm/model_prices_and_context_window_backup.json
+20270 -19930
View File
File diff suppressed because it is too large Load Diff
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
+131 -74
View File
@@ -28,12 +28,16 @@ class MCPRequestHandler:
LITELLM_MCP_SERVERS_HEADER_NAME = SpecialHeaders.mcp_servers.value
LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME = SpecialHeaders.mcp_access_groups.value
# MCP Protocol Version header
MCP_PROTOCOL_VERSION_HEADER_NAME = "MCP-Protocol-Version"
@staticmethod
async def process_mcp_request(scope: Scope) -> Tuple[UserAPIKeyAuth, Optional[str], Optional[List[str]], Optional[Dict[str, str]], Optional[str]]:
async def process_mcp_request(
scope: Scope,
) -> Tuple[
UserAPIKeyAuth, Optional[str], Optional[List[str]], Optional[Dict[str, str]]
]:
"""
Process and validate MCP request headers from the ASGI scope.
This includes:
@@ -49,7 +53,6 @@ class MCPRequestHandler:
mcp_auth_header: Optional[str] MCP auth header to be passed to the MCP server (deprecated)
mcp_servers: Optional[List[str]] List of MCP servers and access groups to use
mcp_server_auth_headers: Optional[Dict[str, str]] Server-specific auth headers in format {server_alias: auth_value}
mcp_protocol_version: Optional[str] MCP protocol version from request header
Raises:
HTTPException: If headers are invalid or missing required headers
@@ -58,39 +61,50 @@ class MCPRequestHandler:
litellm_api_key = (
MCPRequestHandler.get_litellm_api_key_from_headers(headers) or ""
)
# Get the old mcp_auth_header for backward compatibility
mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(headers)
# Get the new server-specific auth headers
mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
# Get MCP protocol version from header
mcp_protocol_version = headers.get(MCPRequestHandler.MCP_PROTOCOL_VERSION_HEADER_NAME)
# Get the new server-specific auth headers
mcp_server_auth_headers = (
MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
)
# Parse MCP servers from header
mcp_servers_header = headers.get(MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME)
mcp_servers_header = headers.get(
MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME
)
verbose_logger.debug(f"Raw MCP servers header: {mcp_servers_header}")
mcp_servers = None
if mcp_servers_header is not None:
try:
mcp_servers = [s.strip() for s in mcp_servers_header.split(",") if s.strip()]
mcp_servers = [
s.strip() for s in mcp_servers_header.split(",") if s.strip()
]
verbose_logger.debug(f"Parsed MCP servers: {mcp_servers}")
except Exception as e:
verbose_logger.debug(f"Error parsing mcp_servers header: {e}")
mcp_servers = None
if mcp_servers_header == "" or (mcp_servers is not None and len(mcp_servers) == 0):
if mcp_servers_header == "" or (
mcp_servers is not None and len(mcp_servers) == 0
):
mcp_servers = []
# Create a proper Request object with mock body method to avoid ASGI receive channel issues
request = Request(scope=scope)
async def mock_body():
return b"{}"
request.body = mock_body # type: ignore
validated_user_api_key_auth = await user_api_key_auth(
api_key=litellm_api_key, request=request
)
return validated_user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, mcp_protocol_version
return (
validated_user_api_key_auth,
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
)
@staticmethod
def _get_mcp_auth_header_from_headers(headers: Headers) -> Optional[str]:
@@ -104,10 +118,12 @@ class MCPRequestHandler:
Support this auth: https://docs.litellm.ai/docs/mcp#using-your-mcp-with-client-side-credentials
If you want to use a different header name, you can set the `LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME` in the secret manager or `mcp_client_side_auth_header_name` in the general settings.
DEPRECATED: This method is deprecated in favor of server-specific auth headers using the format x-mcp-{{server_alias}}-{{header_name}} instead.
"""
mcp_client_side_auth_header_name: str = MCPRequestHandler._get_mcp_client_side_auth_header_name()
mcp_client_side_auth_header_name: str = (
MCPRequestHandler._get_mcp_client_side_auth_header_name()
)
auth_header = headers.get(mcp_client_side_auth_header_name)
if auth_header:
verbose_logger.warning(
@@ -115,42 +131,49 @@ class MCPRequestHandler:
f"Please use server-specific auth headers in the format 'x-mcp-{{server_alias}}-{{header_name}}' instead."
)
return auth_header
@staticmethod
def _get_mcp_server_auth_headers_from_headers(headers: Headers) -> Dict[str, str]:
"""
Parse server-specific MCP auth headers from the request headers.
Looks for headers in the format: x-mcp-{server_alias}-{header_name}
Examples:
- x-mcp-github-authorization: Bearer token123
- x-mcp-zapier-x-api-key: api_key_456
- x-mcp-deepwiki-authorization: Basic base64_encoded_creds
Returns:
Dict[str, str]: Mapping of server alias to auth value
"""
server_auth_headers = {}
prefix = "x-mcp-"
for header_name, header_value in headers.items():
if header_name.lower().startswith(prefix):
# Skip the access groups header as it's not a server auth header
if header_name.lower() == MCPRequestHandler.LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME.lower() or header_name.lower() == MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME.lower():
if (
header_name.lower()
== MCPRequestHandler.LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME.lower()
or header_name.lower()
== MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME.lower()
):
continue
# Extract server_alias and header_name from x-mcp-{server_alias}-{header_name}
remaining = header_name[len(prefix):].lower()
if '-' in remaining:
remaining = header_name[len(prefix) :].lower()
if "-" in remaining:
# Split on the last dash to separate server_alias from header_name
parts = remaining.rsplit('-', 1)
parts = remaining.rsplit("-", 1)
if len(parts) == 2:
server_alias, auth_header_name = parts
server_auth_headers[server_alias] = header_value
verbose_logger.debug(f"Found server auth header: {server_alias} -> {auth_header_name}: {header_value[:10]}...")
verbose_logger.debug(
f"Found server auth header: {server_alias} -> {auth_header_name}: {header_value[:10]}..."
)
return server_auth_headers
@staticmethod
def _get_mcp_client_side_auth_header_name() -> str:
"""
@@ -162,13 +185,21 @@ class MCPRequestHandler:
"""
from litellm.proxy.proxy_server import general_settings
from litellm.secret_managers.main import get_secret_str
MCP_CLIENT_SIDE_AUTH_HEADER_NAME: str = MCPRequestHandler.LITELLM_MCP_AUTH_HEADER_NAME
if get_secret_str("LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME") is not None:
MCP_CLIENT_SIDE_AUTH_HEADER_NAME = get_secret_str("LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME") or MCP_CLIENT_SIDE_AUTH_HEADER_NAME
elif general_settings.get("mcp_client_side_auth_header_name") is not None:
MCP_CLIENT_SIDE_AUTH_HEADER_NAME = general_settings.get("mcp_client_side_auth_header_name") or MCP_CLIENT_SIDE_AUTH_HEADER_NAME
return MCP_CLIENT_SIDE_AUTH_HEADER_NAME
MCP_CLIENT_SIDE_AUTH_HEADER_NAME: str = (
MCPRequestHandler.LITELLM_MCP_AUTH_HEADER_NAME
)
if get_secret_str("LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME") is not None:
MCP_CLIENT_SIDE_AUTH_HEADER_NAME = (
get_secret_str("LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME")
or MCP_CLIENT_SIDE_AUTH_HEADER_NAME
)
elif general_settings.get("mcp_client_side_auth_header_name") is not None:
MCP_CLIENT_SIDE_AUTH_HEADER_NAME = (
general_settings.get("mcp_client_side_auth_header_name")
or MCP_CLIENT_SIDE_AUTH_HEADER_NAME
)
return MCP_CLIENT_SIDE_AUTH_HEADER_NAME
@staticmethod
def get_litellm_api_key_from_headers(headers: Headers) -> Optional[str]:
@@ -229,10 +260,14 @@ class MCPRequestHandler:
try:
allowed_mcp_servers: List[str] = []
allowed_mcp_servers_for_key = (
await MCPRequestHandler._get_allowed_mcp_servers_for_key(user_api_key_auth)
await MCPRequestHandler._get_allowed_mcp_servers_for_key(
user_api_key_auth
)
)
allowed_mcp_servers_for_team = (
await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_api_key_auth)
await MCPRequestHandler._get_allowed_mcp_servers_for_team(
user_api_key_auth
)
)
#########################################################
@@ -274,7 +309,9 @@ class MCPRequestHandler:
try:
key_object_permission = (
await prisma_client.db.litellm_objectpermissiontable.find_unique(
where={"object_permission_id": user_api_key_auth.object_permission_id},
where={
"object_permission_id": user_api_key_auth.object_permission_id
},
)
)
if key_object_permission is None:
@@ -282,17 +319,21 @@ class MCPRequestHandler:
# Get direct MCP servers
direct_mcp_servers = key_object_permission.mcp_servers or []
# Get MCP servers from access groups
access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups(
key_object_permission.mcp_access_groups or []
access_group_servers = (
await MCPRequestHandler._get_mcp_servers_from_access_groups(
key_object_permission.mcp_access_groups or []
)
)
# Combine both lists
all_servers = direct_mcp_servers + access_group_servers
return list(set(all_servers))
except Exception as e:
verbose_logger.warning(f"Failed to get allowed MCP servers for key: {str(e)}")
verbose_logger.warning(
f"Failed to get allowed MCP servers for key: {str(e)}"
)
return []
@staticmethod
@@ -318,10 +359,10 @@ class MCPRequestHandler:
return []
try:
team_obj: Optional[LiteLLM_TeamTable] = (
await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": user_api_key_auth.team_id},
)
team_obj: Optional[
LiteLLM_TeamTable
] = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": user_api_key_auth.team_id},
)
if team_obj is None:
verbose_logger.debug("team_obj is None")
@@ -333,21 +374,27 @@ class MCPRequestHandler:
# Get direct MCP servers
direct_mcp_servers = object_permissions.mcp_servers or []
# Get MCP servers from access groups
access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups(
object_permissions.mcp_access_groups or []
access_group_servers = (
await MCPRequestHandler._get_mcp_servers_from_access_groups(
object_permissions.mcp_access_groups or []
)
)
# Combine both lists
all_servers = direct_mcp_servers + access_group_servers
return list(set(all_servers))
except Exception as e:
verbose_logger.warning(f"Failed to get allowed MCP servers for team: {str(e)}")
verbose_logger.warning(
f"Failed to get allowed MCP servers for team: {str(e)}"
)
return []
@staticmethod
def _get_config_server_ids_for_access_groups(config_mcp_servers, access_groups: List[str]) -> Set[str]:
def _get_config_server_ids_for_access_groups(
config_mcp_servers, access_groups: List[str]
) -> Set[str]:
"""
Helper to get server_ids from config-loaded servers that match any of the given access groups.
"""
@@ -359,7 +406,9 @@ class MCPRequestHandler:
return server_ids
@staticmethod
async def _get_db_server_ids_for_access_groups(prisma_client, access_groups: List[str]) -> Set[str]:
async def _get_db_server_ids_for_access_groups(
prisma_client, access_groups: List[str]
) -> Set[str]:
"""
Helper to get server_ids from DB servers that match any of the given access groups.
"""
@@ -367,21 +416,19 @@ class MCPRequestHandler:
if access_groups and prisma_client is not None:
try:
mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many(
where={
"mcp_access_groups": {
"hasSome": access_groups
}
}
where={"mcp_access_groups": {"hasSome": access_groups}}
)
for server in mcp_servers:
server_ids.add(server.server_id)
except Exception as e:
verbose_logger.debug(f"Error getting MCP servers from access groups: {e}")
verbose_logger.debug(
f"Error getting MCP servers from access groups: {e}"
)
return server_ids
@staticmethod
async def _get_mcp_servers_from_access_groups(
access_groups: List[str]
access_groups: List[str],
) -> List[str]:
"""
Resolve MCP access groups to server IDs by querying BOTH the MCP server table (DB) AND config-loaded servers
@@ -390,22 +437,28 @@ class MCPRequestHandler:
try:
# Import here to avoid circular import
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
# Use the new helper for config-loaded servers
server_ids = MCPRequestHandler._get_config_server_ids_for_access_groups(
global_mcp_server_manager.config_mcp_servers, access_groups
)
# Use the new helper for DB servers
db_server_ids = await MCPRequestHandler._get_db_server_ids_for_access_groups(
prisma_client, access_groups
db_server_ids = (
await MCPRequestHandler._get_db_server_ids_for_access_groups(
prisma_client, access_groups
)
)
server_ids.update(db_server_ids)
return list(server_ids)
except Exception as e:
verbose_logger.warning(f"Failed to get MCP servers from access groups: {str(e)}")
verbose_logger.warning(
f"Failed to get MCP servers from access groups: {str(e)}"
)
return []
@staticmethod
@@ -418,8 +471,8 @@ class MCPRequestHandler:
from typing import List
access_groups: List[str] = []
access_groups_for_key = (
await MCPRequestHandler._get_mcp_access_groups_for_key(user_api_key_auth)
access_groups_for_key = await MCPRequestHandler._get_mcp_access_groups_for_key(
user_api_key_auth
)
access_groups_for_team = (
await MCPRequestHandler._get_mcp_access_groups_for_team(user_api_key_auth)
@@ -482,10 +535,10 @@ class MCPRequestHandler:
verbose_logger.debug("prisma_client is None")
return []
team_obj: Optional[LiteLLM_TeamTable] = (
await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": user_api_key_auth.team_id},
)
team_obj: Optional[
LiteLLM_TeamTable
] = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": user_api_key_auth.team_id},
)
if team_obj is None:
verbose_logger.debug("team_obj is None")
@@ -502,10 +555,14 @@ class MCPRequestHandler:
"""
Extract and parse the x-mcp-access-groups header as a list of strings.
"""
mcp_access_groups_header = headers.get(MCPRequestHandler.LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME)
mcp_access_groups_header = headers.get(
MCPRequestHandler.LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME
)
if mcp_access_groups_header is not None:
try:
return [s.strip() for s in mcp_access_groups_header.split(",") if s.strip()]
return [
s.strip() for s in mcp_access_groups_header.split(",") if s.strip()
]
except Exception:
return None
return None
@@ -516,4 +573,4 @@ class MCPRequestHandler:
Extract and parse the x-mcp-access-groups header from an ASGI scope.
"""
headers = MCPRequestHandler._safe_get_headers_from_scope(scope)
return MCPRequestHandler.get_mcp_access_groups_from_headers(headers)
return MCPRequestHandler.get_mcp_access_groups_from_headers(headers)
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+60 -92
View File
@@ -34,8 +34,6 @@ from litellm.proxy._experimental.mcp_server.utils import (
from litellm.proxy._types import (
LiteLLM_MCPServerTable,
MCPAuthType,
MCPSpecVersion,
MCPSpecVersionType,
MCPTransport,
MCPTransportType,
UserAPIKeyAuth,
@@ -70,38 +68,6 @@ def _deserialize_env_dict(env_data: Any) -> Optional[Dict[str, str]]:
return env_data
def _convert_protocol_version_to_enum(
protocol_version: Optional[str | MCPSpecVersionType],
) -> MCPSpecVersionType:
"""
Convert string protocol version to MCPSpecVersion enum.
Args:
protocol_version: String protocol version, enum, or None
Returns:
MCPSpecVersionType: The enum value
"""
if not protocol_version:
return cast(MCPSpecVersionType, MCPSpecVersion.jun_2025)
# If it's already an MCPSpecVersion enum, return it
if isinstance(protocol_version, MCPSpecVersion):
return cast(MCPSpecVersionType, protocol_version)
# If it's a string, try to match it to enum values
if isinstance(protocol_version, str):
for version in MCPSpecVersion:
if version.value == protocol_version:
return cast(MCPSpecVersionType, version)
# If no match found, return default
verbose_logger.warning(
f"Unknown protocol version '{protocol_version}', using default"
)
return cast(MCPSpecVersionType, MCPSpecVersion.jun_2025)
class MCPServerManager:
def __init__(self):
self.registry: Dict[str, MCPServer] = {}
@@ -113,8 +79,7 @@ class MCPServerManager:
"name": "zapier_mcp_server",
"url": "https://actions.zapier.com/mcp/sk-ak-2ew3bofIeQIkNoeKIdXrF1Hhhp/sse"
"transport": "sse",
"auth_type": "api_key",
"spec_version": "2025-03-26"
"auth_type": "api_key"
},
"uuid-2": {
"name": "google_drive_mcp_server",
@@ -223,7 +188,6 @@ class MCPServerManager:
server_name=server_name,
url=server_config.get("url", None) or "",
transport=server_config.get("transport", MCPTransport.http),
spec_version=server_config.get("spec_version", MCPSpecVersion.jun_2025),
auth_type=server_config.get("auth_type", None),
alias=alias,
)
@@ -239,7 +203,6 @@ class MCPServerManager:
env=server_config.get("env", None) or {},
# TODO: utility fn the default values
transport=server_config.get("transport", MCPTransport.http),
spec_version=server_config.get("spec_version", MCPSpecVersion.jun_2025),
auth_type=server_config.get("auth_type", None),
authentication_token=server_config.get(
"authentication_token", server_config.get("auth_value", None)
@@ -287,7 +250,6 @@ class MCPServerManager:
server_name=getattr(mcp_server, "server_name", None),
url=mcp_server.url,
transport=cast(MCPTransportType, mcp_server.transport),
spec_version=_convert_protocol_version_to_enum(mcp_server.spec_version),
auth_type=cast(MCPAuthType, mcp_server.auth_type),
mcp_info=MCPInfo(
server_name=mcp_server.server_name or mcp_server.server_id,
@@ -350,7 +312,6 @@ class MCPServerManager:
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_protocol_version: Optional[str] = None,
) -> List[MCPTool]:
"""
List all tools available across all MCP Servers.
@@ -390,7 +351,6 @@ class MCPServerManager:
tools = await self._get_tools_from_server(
server=server,
mcp_auth_header=server_auth_header,
mcp_protocol_version=mcp_protocol_version,
)
list_tools_result.extend(tools)
verbose_logger.info(
@@ -414,7 +374,6 @@ class MCPServerManager:
self,
server: MCPServer,
mcp_auth_header: Optional[str] = None,
protocol_version: Optional[str] = None,
) -> MCPClient:
"""
Create an MCPClient instance for the given server.
@@ -422,18 +381,12 @@ class MCPServerManager:
Args:
server (MCPServer): The server configuration
mcp_auth_header: MCP auth header to be passed to the MCP server. This is optional and will be used if provided.
protocol_version: Optional MCP protocol version to use. If not provided, uses server's default.
Returns:
MCPClient: Configured MCP client instance
"""
transport = server.transport or MCPTransport.sse
# Convert protocol version string to enum
protocol_version_enum = _convert_protocol_version_to_enum(
protocol_version or server.spec_version
)
# Handle stdio transport
if transport == MCPTransport.stdio:
# For stdio, we need to get the stdio config from the server
@@ -450,7 +403,6 @@ class MCPServerManager:
auth_value=mcp_auth_header or server.authentication_token,
timeout=60.0,
stdio_config=stdio_config,
protocol_version=protocol_version_enum,
)
else:
# For HTTP/SSE transports
@@ -461,14 +413,12 @@ class MCPServerManager:
auth_type=server.auth_type,
auth_value=mcp_auth_header or server.authentication_token,
timeout=60.0,
protocol_version=protocol_version_enum,
)
async def _get_tools_from_server(
self,
server: MCPServer,
mcp_auth_header: Optional[str] = None,
mcp_protocol_version: Optional[str] = None,
) -> List[MCPTool]:
"""
Helper method to get tools from a single MCP server with prefixed names.
@@ -483,22 +433,18 @@ class MCPServerManager:
verbose_logger.debug(f"Connecting to url: {server.url}")
verbose_logger.info(f"_get_tools_from_server for {server.name}...")
protocol_version = (
mcp_protocol_version if mcp_protocol_version else server.spec_version
)
client = None
try:
client = self._create_mcp_client(
server=server,
mcp_auth_header=mcp_auth_header,
protocol_version=protocol_version,
)
tools = await self._fetch_tools_with_timeout(client, server.name)
prefixed_tools = self._create_prefixed_tools(tools, server)
return prefixed_tools
except Exception as e:
@@ -530,7 +476,7 @@ class MCPServerManager:
async def _list_tools_task():
try:
await client.connect()
tools = await client.list_tools()
verbose_logger.debug(f"Tools from {server_name}: {tools}")
return tools
@@ -609,7 +555,6 @@ class MCPServerManager:
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_protocol_version: Optional[str] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> CallToolResult:
"""
@@ -660,32 +605,54 @@ class MCPServerManager:
"arguments": arguments,
"server_name": server_name_from_prefix,
"user_api_key_auth": user_api_key_auth,
"user_api_key_user_id": getattr(user_api_key_auth, 'user_id', None) if user_api_key_auth else None,
"user_api_key_team_id": getattr(user_api_key_auth, 'team_id', None) if user_api_key_auth else None,
"user_api_key_end_user_id": getattr(user_api_key_auth, 'end_user_id', None) if user_api_key_auth else None,
"user_api_key_hash": getattr(user_api_key_auth, 'api_key_hash', None) if user_api_key_auth else None,
"user_api_key_user_id": getattr(user_api_key_auth, "user_id", None)
if user_api_key_auth
else None,
"user_api_key_team_id": getattr(user_api_key_auth, "team_id", None)
if user_api_key_auth
else None,
"user_api_key_end_user_id": getattr(
user_api_key_auth, "end_user_id", None
)
if user_api_key_auth
else None,
"user_api_key_hash": getattr(user_api_key_auth, "api_key_hash", None)
if user_api_key_auth
else None,
}
# Create MCP request object for processing
mcp_request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs(pre_hook_kwargs)
mcp_request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs(
pre_hook_kwargs
)
# Convert to LLM format for existing guardrail compatibility
synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(mcp_request_obj, pre_hook_kwargs)
synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(
mcp_request_obj, pre_hook_kwargs
)
try:
# Use standard pre_call_hook with call_type="mcp_call"
modified_data = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_auth, #type: ignore
user_api_key_dict=user_api_key_auth, # type: ignore
data=synthetic_llm_data,
call_type="mcp_call" #type: ignore
call_type="mcp_call", # type: ignore
)
if modified_data:
# Convert response back to MCP format and apply modifications
modified_kwargs = proxy_logging_obj._convert_mcp_hook_response_to_kwargs(modified_data, pre_hook_kwargs)
modified_kwargs = (
proxy_logging_obj._convert_mcp_hook_response_to_kwargs(
modified_data, pre_hook_kwargs
)
)
if modified_kwargs.get("arguments") != arguments:
arguments = modified_kwargs["arguments"]
except (BlockedPiiEntityError, GuardrailRaisedException, HTTPException) as e:
except (
BlockedPiiEntityError,
GuardrailRaisedException,
HTTPException,
) as e:
# Re-raise guardrail exceptions to properly fail the MCP call
verbose_logger.error(
f"Guardrail blocked MCP tool call pre call: {str(e)}"
@@ -706,11 +673,9 @@ class MCPServerManager:
client = self._create_mcp_client(
server=mcp_server,
mcp_auth_header=server_auth_header,
protocol_version=mcp_protocol_version,
)
async with client:
# Use the original tool name (without prefix) for the actual call
call_tool_params = MCPCallToolRequestParams(
name=original_tool_name,
@@ -721,7 +686,7 @@ class MCPServerManager:
# Create synthetic LLM data for during hook processing
from litellm.types.llms.base import HiddenParams
from litellm.types.mcp import MCPDuringCallRequestObject
request_obj = MCPDuringCallRequestObject(
tool_name=name,
arguments=arguments,
@@ -729,28 +694,29 @@ class MCPServerManager:
start_time=start_time.timestamp() if start_time else None,
hidden_params=HiddenParams(),
)
during_hook_kwargs = {
"name": name,
"arguments": arguments,
"server_name": server_name_from_prefix,
"user_api_key_auth": user_api_key_auth,
}
synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(request_obj, during_hook_kwargs)
synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(
request_obj, during_hook_kwargs
)
during_hook_task = asyncio.create_task(
proxy_logging_obj.during_call_hook(
user_api_key_dict=user_api_key_auth,
data=synthetic_llm_data,
call_type="mcp_call" #type: ignore
call_type="mcp_call", # type: ignore
)
)
tasks.append(during_hook_task)
tasks.append(asyncio.create_task(client.call_tool(call_tool_params)))
try:
mcp_responses = await asyncio.gather(*tasks)
# If proxy_logging_obj is None, the tool call result is at index 0
@@ -839,19 +805,21 @@ class MCPServerManager:
)
verbose_logger.info("Loading MCP servers from database into registry...")
# perform authz check to filter the mcp servers user has access to
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to your proxy"
)
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
for server in db_mcp_servers:
verbose_logger.debug(f"Adding server to registry: {server.server_id} ({server.server_name})")
verbose_logger.debug(
f"Adding server to registry: {server.server_id} ({server.server_name})"
)
self.add_update_server(server)
verbose_logger.info(f"Registry now contains {len(self.get_registry())} servers")
def get_mcp_server_by_id(self, server_id: str) -> Optional[MCPServer]:
@@ -869,7 +837,6 @@ class MCPServerManager:
server_name: str,
url: str,
transport: str,
spec_version: str,
auth_type: Optional[str] = None,
alias: Optional[str] = None,
) -> str:
@@ -885,7 +852,6 @@ class MCPServerManager:
server_name: Name of the server
url: Server URL
transport: Transport type (sse, http, etc.)
spec_version: MCP spec version
auth_type: Authentication type (optional)
alias: Server alias (optional)
@@ -893,7 +859,9 @@ class MCPServerManager:
A deterministic server ID string
"""
# Create a string from all the identifying parameters
params_string = f"{server_name}|{url}|{transport}|{spec_version}|{auth_type or ''}|{alias or ''}"
params_string = (
f"{server_name}|{url}|{transport}|{auth_type or ''}|{alias or ''}"
)
# Generate SHA-256 hash
hash_object = hashlib.sha256(params_string.encode("utf-8"))
@@ -1050,11 +1018,12 @@ class MCPServerManager:
alias=_server_config.alias,
url=_server_config.url,
transport=_server_config.transport,
spec_version=_server_config.spec_version,
auth_type=_server_config.auth_type,
created_at=datetime.datetime.now(),
updated_at=datetime.datetime.now(),
description=_server_config.mcp_info.get("description") if _server_config.mcp_info else None,
description=_server_config.mcp_info.get("description")
if _server_config.mcp_info
else None,
mcp_info=_server_config.mcp_info,
mcp_access_groups=_server_config.access_groups or [],
# Stdio-specific fields
@@ -1111,7 +1080,6 @@ class MCPServerManager:
description=server.description,
url=server.url,
transport=server.transport,
spec_version=server.spec_version,
auth_type=server.auth_type,
created_at=server.created_at,
created_by=server.created_by,
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
+78 -50
View File
@@ -23,7 +23,6 @@ router = APIRouter(
if MCP_AVAILABLE:
from litellm.experimental_mcp_client.client import MCPTool
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
_convert_protocol_version_to_enum,
global_mcp_server_manager,
)
from litellm.proxy._experimental.mcp_server.server import (
@@ -34,18 +33,24 @@ if MCP_AVAILABLE:
########################################################
############ MCP Server REST API Routes #################
def _get_server_auth_header(
server, mcp_server_auth_headers: Optional[Dict[str, str]], mcp_auth_header: Optional[str]
server,
mcp_server_auth_headers: Optional[Dict[str, str]],
mcp_auth_header: Optional[str],
) -> Optional[str]:
"""Helper function to get server-specific auth header with case-insensitive matching."""
if mcp_server_auth_headers and server.alias:
normalized_server_alias = server.alias.lower()
normalized_headers = {k.lower(): v for k, v in mcp_server_auth_headers.items()}
normalized_headers = {
k.lower(): v for k, v in mcp_server_auth_headers.items()
}
server_auth = normalized_headers.get(normalized_server_alias)
if server_auth is not None:
return server_auth
elif mcp_server_auth_headers and server.server_name:
normalized_server_name = server.server_name.lower()
normalized_headers = {k.lower(): v for k, v in mcp_server_auth_headers.items()}
normalized_headers = {
k.lower(): v for k, v in mcp_server_auth_headers.items()
}
server_auth = normalized_headers.get(normalized_server_name)
if server_auth is not None:
return server_auth
@@ -63,12 +68,11 @@ if MCP_AVAILABLE:
for tool in tools
]
async def _get_tools_for_single_server(server, server_auth_header, mcp_protocol_version):
async def _get_tools_for_single_server(server, server_auth_header):
"""Helper function to get tools for a single server."""
tools = await global_mcp_server_manager._get_tools_from_server(
server=server,
mcp_auth_header=server_auth_header,
mcp_protocol_version=mcp_protocol_version,
)
return _create_tool_response_objects(tools, server.mcp_info)
@@ -104,17 +108,20 @@ if MCP_AVAILABLE:
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
try:
# Extract auth headers from request
headers = request.headers
mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(headers)
mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
mcp_protocol_version = headers.get(MCPRequestHandler.MCP_PROTOCOL_VERSION_HEADER_NAME)
mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(
headers
)
mcp_server_auth_headers = (
MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
)
list_tools_result = []
error_message = None
# If server_id is specified, only query that specific server
if server_id:
server = global_mcp_server_manager.get_mcp_server_by_id(server_id)
@@ -122,49 +129,67 @@ if MCP_AVAILABLE:
return {
"tools": [],
"error": "server_not_found",
"message": f"Server with id {server_id} not found"
"message": f"Server with id {server_id} not found",
}
server_auth_header = _get_server_auth_header(server, mcp_server_auth_headers, mcp_auth_header)
server_auth_header = _get_server_auth_header(
server, mcp_server_auth_headers, mcp_auth_header
)
try:
list_tools_result = await _get_tools_for_single_server(server, server_auth_header, mcp_protocol_version)
list_tools_result = await _get_tools_for_single_server(
server, server_auth_header
)
except Exception as e:
verbose_logger.exception(f"Error getting tools from {server.name}: {e}")
verbose_logger.exception(
f"Error getting tools from {server.name}: {e}"
)
return {
"tools": [],
"error": "server_error",
"message": f"Failed to get tools from server {server.name}: {str(e)}"
"message": f"Failed to get tools from server {server.name}: {str(e)}",
}
else:
# Query all servers
errors = []
for server in global_mcp_server_manager.get_registry().values():
server_auth_header = _get_server_auth_header(server, mcp_server_auth_headers, mcp_auth_header)
server_auth_header = _get_server_auth_header(
server, mcp_server_auth_headers, mcp_auth_header
)
try:
tools_result = await _get_tools_for_single_server(server, server_auth_header, mcp_protocol_version)
tools_result = await _get_tools_for_single_server(
server, server_auth_header
)
list_tools_result.extend(tools_result)
except Exception as e:
verbose_logger.exception(f"Error getting tools from {server.name}: {e}")
verbose_logger.exception(
f"Error getting tools from {server.name}: {e}"
)
errors.append(f"{server.name}: {str(e)}")
continue
if errors and not list_tools_result:
error_message = "Failed to get tools from servers: " + "; ".join(errors)
error_message = "Failed to get tools from servers: " + "; ".join(
errors
)
return {
"tools": list_tools_result,
"error": "partial_failure" if error_message else None,
"message": error_message if error_message else "Successfully retrieved tools"
"message": error_message
if error_message
else "Successfully retrieved tools",
}
except Exception as e:
verbose_logger.exception("Unexpected error in list_tool_rest_api: %s", str(e))
verbose_logger.exception(
"Unexpected error in list_tool_rest_api: %s", str(e)
)
return {
"tools": [],
"error": "unexpected_error",
"message": f"An unexpected error occurred: {str(e)}"
"message": f"An unexpected error occurred: {str(e)}",
}
@router.post("/tools/call", dependencies=[Depends(user_api_key_auth)])
@@ -196,9 +221,9 @@ if MCP_AVAILABLE:
detail={
"error": "blocked_pii_entity",
"message": str(e),
"entity_type": getattr(e, 'entity_type', None),
"guardrail_name": getattr(e, 'guardrail_name', None)
}
"entity_type": getattr(e, "entity_type", None),
"guardrail_name": getattr(e, "guardrail_name", None),
},
)
except GuardrailRaisedException as e:
verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {str(e)}")
@@ -207,8 +232,8 @@ if MCP_AVAILABLE:
detail={
"error": "guardrail_violation",
"message": str(e),
"guardrail_name": getattr(e, 'guardrail_name', None)
}
"guardrail_name": getattr(e, "guardrail_name", None),
},
)
except HTTPException as e:
# Re-raise HTTPException as-is to preserve status code and detail
@@ -220,10 +245,10 @@ if MCP_AVAILABLE:
status_code=500,
detail={
"error": "internal_server_error",
"message": f"An unexpected error occurred: {str(e)}"
}
"message": f"An unexpected error occurred: {str(e)}",
},
)
########################################################
# MCP Connection testing routes
# /health -> Test if we can connect to the MCP server
@@ -234,15 +259,15 @@ if MCP_AVAILABLE:
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
NewMCPServerRequest,
)
async def _execute_with_mcp_client(request: NewMCPServerRequest, operation):
"""
Common helper to create MCP client, execute operation, and ensure proper cleanup.
Args:
request: MCP server configuration
operation: Async function that takes a client and returns the operation result
Returns:
Operation result or error response
"""
@@ -254,15 +279,14 @@ if MCP_AVAILABLE:
name=request.alias or request.server_name or "",
url=request.url,
transport=request.transport,
spec_version=_convert_protocol_version_to_enum(request.spec_version),
auth_type=request.auth_type,
mcp_info=request.mcp_info,
),
mcp_auth_header=None,
)
return await operation(client)
except Exception as e:
verbose_logger.error(f"Error in MCP operation: {e}", exc_info=True)
return {"status": "error", "message": "An internal error has occurred."}
@@ -273,6 +297,7 @@ if MCP_AVAILABLE:
await client.disconnect()
except Exception as e:
verbose_logger.warning(f"Error disconnecting MCP client: {e}")
@router.post("/test/connection")
async def test_connection(
request: NewMCPServerRequest,
@@ -280,13 +305,13 @@ if MCP_AVAILABLE:
"""
Test if we can connect to the provided MCP server before adding it
"""
async def _test_connection_operation(client):
await client.connect()
return {"status": "ok"}
return await _execute_with_mcp_client(request, _test_connection_operation)
@router.post("/test/tools/list")
async def test_tools_list(
request: NewMCPServerRequest,
@@ -295,13 +320,16 @@ if MCP_AVAILABLE:
"""
Preview tools available from MCP server before adding it
"""
async def _list_tools_operation(client):
list_tools_result: List[MCPTool] = await client.list_tools()
model_dumped_tools: List[dict] = [tool.model_dump() for tool in list_tools_result]
model_dumped_tools: List[dict] = [
tool.model_dump() for tool in list_tools_result
]
return {
"tools": model_dumped_tools,
"error": None,
"message": "Successfully retrieved tools"
"message": "Successfully retrieved tools",
}
return await _execute_with_mcp_client(request, _list_tools_operation)
litellm/proxy/_experimental/mcp_server/server.py
+155 -71
View File
@@ -130,7 +130,9 @@ if MCP_AVAILABLE:
await _sse_session_manager_cm.__aenter__()
_SESSION_MANAGERS_INITIALIZED = True
verbose_logger.info("MCP Server started with StreamableHTTP and SSE session managers!")
verbose_logger.info(
"MCP Server started with StreamableHTTP and SSE session managers!"
)
async def shutdown_session_managers():
"""Shutdown the session managers."""
@@ -171,11 +173,18 @@ if MCP_AVAILABLE:
"""
try:
# Get user authentication from context variable
user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, mcp_protocol_version = (
get_auth_context()
(
user_api_key_auth,
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
) = get_auth_context()
verbose_logger.debug(
f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}"
)
verbose_logger.debug(
f"MCP list_tools - MCP servers from context: {mcp_servers}"
)
verbose_logger.debug(f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}")
verbose_logger.debug(f"MCP list_tools - MCP servers from context: {mcp_servers}")
verbose_logger.debug(
f"MCP list_tools - MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}"
)
@@ -186,9 +195,10 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
mcp_protocol_version=mcp_protocol_version,
)
verbose_logger.info(f"MCP list_tools - Successfully returned {len(tools)} tools")
verbose_logger.info(
f"MCP list_tools - Successfully returned {len(tools)} tools"
)
return tools
except Exception as e:
verbose_logger.exception(f"Error in list_tools endpoint: {str(e)}")
@@ -220,9 +230,16 @@ if MCP_AVAILABLE:
from litellm.proxy.proxy_server import proxy_config
# Validate arguments
user_api_key_auth, mcp_auth_header, _, mcp_server_auth_headers, mcp_protocol_version = get_auth_context()
(
user_api_key_auth,
mcp_auth_header,
_,
mcp_server_auth_headers,
) = get_auth_context()
verbose_logger.debug(f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}")
verbose_logger.debug(
f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}"
)
try:
# Create a body date for logging
body_data = {"name": name, "arguments": arguments}
@@ -249,17 +266,22 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
mcp_protocol_version=mcp_protocol_version,
**data, # for logging
)
except BlockedPiiEntityError as e:
verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {str(e)}")
# Return error as text content for MCP protocol
return [TextContent(text=f"Error: Blocked PII entity detected - {str(e)}", type="text")]
return [
TextContent(
text=f"Error: Blocked PII entity detected - {str(e)}", type="text"
)
]
except GuardrailRaisedException as e:
verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {str(e)}")
# Return error as text content for MCP protocol
return [TextContent(text=f"Error: Guardrail violation - {str(e)}", type="text")]
return [
TextContent(text=f"Error: Guardrail violation - {str(e)}", type="text")
]
except HTTPException as e:
verbose_logger.error(f"HTTPException in MCP tool call: {str(e)}")
# Return error as text content for MCP protocol
@@ -287,6 +309,7 @@ if MCP_AVAILABLE:
Get the filtered MCP servers from the MCP server names
"""
from typing import Set
filtered_server_ids: Set[str] = set()
# Filter servers based on mcp_servers parameter if provided
if mcp_servers is not None:
@@ -297,7 +320,11 @@ if MCP_AVAILABLE:
server = global_mcp_server_manager.get_mcp_server_by_id(server_id)
if server:
match_list = [s.lower() for s in [server.alias, server.server_name, server_id] if s is not None]
match_list = [
s.lower()
for s in [server.alias, server.server_name, server_id]
if s is not None
]
if server_or_group.lower() in match_list:
filtered_server_ids.add(server_id)
@@ -306,19 +333,23 @@ if MCP_AVAILABLE:
if not server_name_matched:
try:
access_group_server_ids = await MCPRequestHandler._get_mcp_servers_from_access_groups(
[server_or_group]
access_group_server_ids = (
await MCPRequestHandler._get_mcp_servers_from_access_groups(
[server_or_group]
)
)
# Only include servers that the user has access to
for server_id in access_group_server_ids:
if server_id in allowed_mcp_servers:
filtered_server_ids.add(server_id)
except Exception as e:
verbose_logger.debug(f"Could not resolve '{server_or_group}' as access group: {e}")
verbose_logger.debug(
f"Could not resolve '{server_or_group}' as access group: {e}"
)
if filtered_server_ids:
allowed_mcp_servers = list(filtered_server_ids)
return allowed_mcp_servers
async def _get_tools_from_mcp_servers(
@@ -326,7 +357,6 @@ if MCP_AVAILABLE:
mcp_auth_header: Optional[str],
mcp_servers: Optional[List[str]],
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_protocol_version: Optional[str] = None,
) -> List[MCPTool]:
"""
Helper method to fetch tools from MCP servers based on server filtering criteria.
@@ -344,7 +374,9 @@ if MCP_AVAILABLE:
return []
# Get allowed MCP servers based on user permissions
allowed_mcp_servers = await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth)
allowed_mcp_servers = await global_mcp_server_manager.get_allowed_mcp_servers(
user_api_key_auth
)
if mcp_servers is not None:
allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names(
@@ -352,7 +384,6 @@ if MCP_AVAILABLE:
allowed_mcp_servers=allowed_mcp_servers,
)
# Get tools from each allowed server
all_tools = []
for server_id in allowed_mcp_servers:
@@ -375,15 +406,20 @@ if MCP_AVAILABLE:
tools = await global_mcp_server_manager._get_tools_from_server(
server=server,
mcp_auth_header=server_auth_header,
mcp_protocol_version=mcp_protocol_version,
)
all_tools.extend(tools)
verbose_logger.debug(f"Successfully fetched {len(tools)} tools from server {server.name}")
verbose_logger.debug(
f"Successfully fetched {len(tools)} tools from server {server.name}"
)
except Exception as e:
verbose_logger.exception(f"Error getting tools from server {server.name}: {str(e)}")
verbose_logger.exception(
f"Error getting tools from server {server.name}: {str(e)}"
)
# Continue with other servers instead of failing completely
verbose_logger.info(f"Successfully fetched {len(all_tools)} tools total from all MCP servers")
verbose_logger.info(
f"Successfully fetched {len(all_tools)} tools total from all MCP servers"
)
return all_tools
async def _list_mcp_tools(
@@ -391,7 +427,6 @@ if MCP_AVAILABLE:
mcp_auth_header: Optional[str] = None,
mcp_servers: Optional[List[str]] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_protocol_version: Optional[str] = None,
) -> List[MCPTool]:
"""
List all available MCP tools.
@@ -415,11 +450,14 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
mcp_protocol_version=mcp_protocol_version,
)
verbose_logger.debug(f"Successfully fetched {len(managed_tools)} tools from managed MCP servers")
verbose_logger.debug(
f"Successfully fetched {len(managed_tools)} tools from managed MCP servers"
)
except Exception as e:
verbose_logger.exception(f"Error getting tools from managed MCP servers: {str(e)}")
verbose_logger.exception(
f"Error getting tools from managed MCP servers: {str(e)}"
)
# Continue with empty managed tools list instead of failing completely
# Get tools from local registry
@@ -430,10 +468,16 @@ if MCP_AVAILABLE:
# Convert local tools to MCPTool format
for tool in local_tools_raw:
# Convert from litellm.types.mcp_server.tool_registry.MCPTool to mcp.types.Tool
mcp_tool = MCPTool(name=tool.name, description=tool.description, inputSchema=tool.input_schema)
mcp_tool = MCPTool(
name=tool.name,
description=tool.description,
inputSchema=tool.input_schema,
)
local_tools.append(mcp_tool)
except Exception as e:
verbose_logger.exception(f"Error getting tools from local registry: {str(e)}")
verbose_logger.exception(
f"Error getting tools from local registry: {str(e)}"
)
# Continue with empty local tools list instead of failing completely
# Combine all tools
@@ -448,7 +492,6 @@ if MCP_AVAILABLE:
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_protocol_version: Optional[str] = None,
**kwargs: Any,
) -> List[Union[TextContent, ImageContent, EmbeddedResource]]:
"""
@@ -456,35 +499,46 @@ if MCP_AVAILABLE:
"""
start_time = datetime.now()
if arguments is None:
raise HTTPException(status_code=400, detail="Request arguments are required")
raise HTTPException(
status_code=400, detail="Request arguments are required"
)
# Remove prefix from tool name for logging and processing
original_tool_name, server_name_from_prefix = get_server_name_prefix_tool_mcp(name)
standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = _get_standard_logging_mcp_tool_call(
name=original_tool_name, # Use original name for logging
arguments=arguments,
server_name=server_name_from_prefix,
original_tool_name, server_name_from_prefix = get_server_name_prefix_tool_mcp(
name
)
standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = (
_get_standard_logging_mcp_tool_call(
name=original_tool_name, # Use original name for logging
arguments=arguments,
server_name=server_name_from_prefix,
)
)
litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get(
"litellm_logging_obj", None
)
litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None)
if litellm_logging_obj:
litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = standard_logging_mcp_tool_call
litellm_logging_obj.model_call_details[
"mcp_tool_call_metadata"
] = standard_logging_mcp_tool_call
litellm_logging_obj.model = f"MCP: {name}"
# Try managed server tool first (pass the full prefixed name)
# Primary and recommended way to use MCP servers
#########################################################
mcp_server: Optional[MCPServer] = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
mcp_server: Optional[
MCPServer
] = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
if mcp_server:
standard_logging_mcp_tool_call["mcp_server_cost_info"] = (mcp_server.mcp_info or {}).get(
"mcp_server_cost_info"
)
standard_logging_mcp_tool_call["mcp_server_cost_info"] = (
mcp_server.mcp_info or {}
).get("mcp_server_cost_info")
response = await _handle_managed_mcp_tool(
name=name, # Pass the full name (potentially prefixed)
arguments=arguments,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
mcp_protocol_version=mcp_protocol_version,
litellm_logging_obj=litellm_logging_obj,
)
@@ -537,7 +591,6 @@ if MCP_AVAILABLE:
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_protocol_version: Optional[str] = None,
litellm_logging_obj: Optional[Any] = None,
) -> List[Union[TextContent, ImageContent, EmbeddedResource]]:
"""Handle tool execution for managed server tools"""
@@ -577,12 +630,39 @@ if MCP_AVAILABLE:
Get the MCP servers from the path
"""
import re
mcp_servers_from_path: Optional[List[str]] = None
mcp_path_match = re.match(r"^/mcp/([^/]+/[^/]+|[^/]+)(/.*)?$", path)
# Match /mcp/<servers_and_maybe_path>
# Where servers can be comma-separated list of server names
# Server names can contain slashes (e.g., "custom_solutions/user_123")
mcp_path_match = re.match(r"^/mcp/([^?#]+)(?:\?.*)?(?:#.*)?$", path)
if mcp_path_match:
mcp_servers_str = mcp_path_match.group(1)
if mcp_servers_str:
mcp_servers_from_path = [s.strip() for s in mcp_servers_str.split(",") if s.strip()]
servers_and_path = mcp_path_match.group(1)
if servers_and_path:
# Check if it contains commas (comma-separated servers)
if ',' in servers_and_path:
# For comma-separated, look for a path at the end
# Common patterns: /tools, /chat/completions, etc.
path_match = re.search(r'/([^/,]+(?:/[^/,]+)*)$', servers_and_path)
if path_match:
# Path found at the end, remove it from servers
path_part = '/' + path_match.group(1)
servers_part = servers_and_path[:-len(path_part)]
mcp_servers_from_path = [s.strip() for s in servers_part.split(',') if s.strip()]
else:
# No path, just comma-separated servers
mcp_servers_from_path = [s.strip() for s in servers_and_path.split(',') if s.strip()]
else:
# Single server case - use regex approach for server/path separation
# This handles cases like "custom_solutions/user_123/chat/completions"
# where we want to extract "custom_solutions/user_123" as the server name
single_server_match = re.match(r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path)
if single_server_match:
server_name = single_server_match.group(1)
mcp_servers_from_path = [server_name]
else:
mcp_servers_from_path = [servers_and_path]
return mcp_servers_from_path
async def extract_mcp_auth_context(scope, path):
@@ -597,7 +677,6 @@ if MCP_AVAILABLE:
mcp_auth_header,
_,
mcp_server_auth_headers,
mcp_protocol_version,
) = await MCPRequestHandler.process_mcp_request(scope)
mcp_servers = mcp_servers_from_path
else:
@@ -606,11 +685,12 @@ if MCP_AVAILABLE:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
mcp_protocol_version,
) = await MCPRequestHandler.process_mcp_request(scope)
return user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, mcp_protocol_version
return user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers
async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None:
async def handle_streamable_http_mcp(
scope: Scope, receive: Receive, send: Send
) -> None:
"""Handle MCP requests through StreamableHTTP."""
try:
path = scope.get("path", "")
@@ -619,20 +699,19 @@ if MCP_AVAILABLE:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
mcp_protocol_version,
) = await extract_mcp_auth_context(scope, path)
verbose_logger.debug(f"MCP request mcp_servers (header/path): {mcp_servers}")
verbose_logger.debug(
f"MCP request mcp_servers (header/path): {mcp_servers}"
)
verbose_logger.debug(
f"MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}"
)
verbose_logger.debug(f"MCP protocol version: {mcp_protocol_version}")
# Set the auth context variable for easy access in MCP functions
set_auth_context(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
mcp_protocol_version=mcp_protocol_version,
)
# Ensure session managers are initialized
@@ -656,7 +735,9 @@ if MCP_AVAILABLE:
)
await error_response(scope, receive, send)
except Exception as response_error:
verbose_logger.exception(f"Failed to send error response: {response_error}")
verbose_logger.exception(
f"Failed to send error response: {response_error}"
)
# If we can't send a proper response, re-raise the original error
raise e
@@ -669,19 +750,18 @@ if MCP_AVAILABLE:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
mcp_protocol_version,
) = await extract_mcp_auth_context(scope, path)
verbose_logger.debug(f"MCP request mcp_servers (header/path): {mcp_servers}")
verbose_logger.debug(
f"MCP request mcp_servers (header/path): {mcp_servers}"
)
verbose_logger.debug(
f"MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}"
)
verbose_logger.debug(f"MCP protocol version: {mcp_protocol_version}")
set_auth_context(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
mcp_protocol_version=mcp_protocol_version,
)
if not _SESSION_MANAGERS_INITIALIZED:
@@ -703,7 +783,9 @@ if MCP_AVAILABLE:
)
await error_response(scope, receive, send)
except Exception as response_error:
verbose_logger.exception(f"Failed to send error response: {response_error}")
verbose_logger.exception(
f"Failed to send error response: {response_error}"
)
# If we can't send a proper response, re-raise the original error
raise e
@@ -739,7 +821,6 @@ if MCP_AVAILABLE:
mcp_auth_header: Optional[str] = None,
mcp_servers: Optional[List[str]] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_protocol_version: Optional[str] = None,
) -> None:
"""
Set the UserAPIKeyAuth in the auth context variable.
@@ -755,13 +836,17 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
mcp_protocol_version=mcp_protocol_version,
)
auth_context_var.set(auth_user)
def get_auth_context() -> Tuple[
Optional[UserAPIKeyAuth], Optional[str], Optional[List[str]], Optional[Dict[str, str]], Optional[str]
]:
def get_auth_context() -> (
Tuple[
Optional[UserAPIKeyAuth],
Optional[str],
Optional[List[str]],
Optional[Dict[str, str]],
]
):
"""
Get the UserAPIKeyAuth from the auth context variable.
@@ -776,9 +861,8 @@ if MCP_AVAILABLE:
auth_user.mcp_auth_header,
auth_user.mcp_servers,
auth_user.mcp_server_auth_headers,
auth_user.mcp_protocol_version,
)
return None, None, None, None, None
return None, None, None, None
########################################################
############ End of Auth Context Functions #############
litellm/proxy/_experimental/out/_next/static/fhuPj8WYsuMGymIUE7Xgu/_buildManifest.js → litellm/proxy/_experimental/out/_next/static/0oPk2eYtSaTLaPyVixqA8/_buildManifest.js
View File
litellm/proxy/_experimental/out/_next/static/fhuPj8WYsuMGymIUE7Xgu/_ssgManifest.js → litellm/proxy/_experimental/out/_next/static/0oPk2eYtSaTLaPyVixqA8/_ssgManifest.js
View File
litellm/proxy/_experimental/out/_next/static/chunks/117-a0da667066d322b6.js → litellm/proxy/_experimental/out/_next/static/chunks/117-6a9f6a591033d4a0.js
+1 -1
View File
File diff suppressed because one or more lines are too long
litellm/proxy/_experimental/out/_next/static/chunks/154-6f752d9e0a5e497b.js → litellm/proxy/_experimental/out/_next/static/chunks/154-b1f2a106d0e0d77b.js
+1 -1
View File
File diff suppressed because one or more lines are too long
litellm/proxy/_experimental/out/_next/static/chunks/220-8af5927d18414264.js → litellm/proxy/_experimental/out/_next/static/chunks/220-89d73a525e307735.js
+1 -1
View File
File diff suppressed because one or more lines are too long
litellm/proxy/_experimental/out/_next/static/chunks/50-bb8a11a7610535aa.js
-1
View File
File diff suppressed because one or more lines are too long
litellm/proxy/_experimental/out/_next/static/chunks/50-d0da2dd7acce2eb9.js
+1
View File
File diff suppressed because one or more lines are too long
litellm/proxy/_experimental/out/_next/static/chunks/app/layout-df4939b4f2ec37d7.js
+1
View File
@@ -0,0 +1 @@
(self.webpackChunk_N_E=self.webpackChunk_N_E||[]).push([[185],{85210:function(n,e,t){Promise.resolve().then(t.t.bind(t,39974,23)),Promise.resolve().then(t.t.bind(t,2778,23))},2778:function(){},39974:function(n){n.exports={style:{fontFamily:"'__Inter_1c856b', '__Inter_Fallback_1c856b'",fontStyle:"normal"},className:"__className_1c856b"}}},function(n){n.O(0,[919,986,971,117,744],function(){return n(n.s=85210)}),_N_E=n.O()}]);
litellm/proxy/_experimental/out/_next/static/chunks/app/layout-f4acf18888f0aa20.js
-1
View File
@@ -1 +0,0 @@
(self.webpackChunk_N_E=self.webpackChunk_N_E||[]).push([[185],{96443:function(n,e,t){Promise.resolve().then(t.t.bind(t,39974,23)),Promise.resolve().then(t.t.bind(t,2778,23))},2778:function(){},39974:function(n){n.exports={style:{fontFamily:"'__Inter_b0dd8a', '__Inter_Fallback_b0dd8a'",fontStyle:"normal"},className:"__className_b0dd8a"}}},function(n){n.O(0,[919,986,971,117,744],function(){return n(n.s=96443)}),_N_E=n.O()}]);
litellm/proxy/_experimental/out/_next/static/chunks/app/model_hub/page-0dbadf20167b786c.js → litellm/proxy/_experimental/out/_next/static/chunks/app/model_hub/page-13b00ef4a072d920.js
+1 -1
View File
@@ -1 +1 @@
(self.webpackChunk_N_E=self.webpackChunk_N_E||[]).push([[418],{21024:function(e,n,t){Promise.resolve().then(t.bind(t,52829))},52829:function(e,n,t){"use strict";t.r(n),t.d(n,{default:function(){return f}});var u=t(57437),s=t(2265),c=t(99376),r=t(72162);function f(){let e=(0,c.useSearchParams)().get("key"),[n,t]=(0,s.useState)(null);return(0,s.useEffect)(()=>{e&&t(e)},[e]),(0,u.jsx)(r.Z,{accessToken:n})}}},function(e){e.O(0,[50,521,154,162,971,117,744],function(){return e(e.s=21024)}),_N_E=e.O()}]);
(self.webpackChunk_N_E=self.webpackChunk_N_E||[]).push([[418],{67355:function(e,n,t){Promise.resolve().then(t.bind(t,52829))},52829:function(e,n,t){"use strict";t.r(n),t.d(n,{default:function(){return f}});var u=t(57437),s=t(2265),c=t(99376),r=t(72162);function f(){let e=(0,c.useSearchParams)().get("key"),[n,t]=(0,s.useState)(null);return(0,s.useEffect)(()=>{e&&t(e)},[e]),(0,u.jsx)(r.Z,{accessToken:n})}}},function(e){e.O(0,[50,521,154,162,971,117,744],function(){return e(e.s=67355)}),_N_E=e.O()}]);
litellm/proxy/_experimental/out/_next/static/chunks/app/model_hub_table/page-f469bae327299fbb.js → litellm/proxy/_experimental/out/_next/static/chunks/app/model_hub_table/page-304b7041a3fa39f7.js
+1 -1
View File
@@ -1 +1 @@
(self.webpackChunk_N_E=self.webpackChunk_N_E||[]).push([[25],{64563:function(e,n,u){Promise.resolve().then(u.bind(u,22775))},22775:function(e,n,u){"use strict";u.r(n),u.d(n,{default:function(){return f}});var t=u(57437),s=u(2265),r=u(99376),c=u(36172);function f(){let e=(0,r.useSearchParams)().get("key"),[n,u]=(0,s.useState)(null);return(0,s.useEffect)(()=>{e&&u(e)},[e]),(0,t.jsx)(c.Z,{accessToken:n,publicPage:!0,premiumUser:!1,userRole:null})}}},function(e){e.O(0,[50,521,866,154,162,172,971,117,744],function(){return e(e.s=64563)}),_N_E=e.O()}]);
(self.webpackChunk_N_E=self.webpackChunk_N_E||[]).push([[25],{38520:function(e,n,u){Promise.resolve().then(u.bind(u,22775))},22775:function(e,n,u){"use strict";u.r(n),u.d(n,{default:function(){return f}});var t=u(57437),s=u(2265),r=u(99376),c=u(36172);function f(){let e=(0,r.useSearchParams)().get("key"),[n,u]=(0,s.useState)(null);return(0,s.useEffect)(()=>{e&&u(e)},[e]),(0,t.jsx)(c.Z,{accessToken:n,publicPage:!0,premiumUser:!1,userRole:null})}}},function(e){e.O(0,[50,521,866,154,162,172,971,117,744],function(){return e(e.s=38520)}),_N_E=e.O()}]);
litellm/proxy/_experimental/out/_next/static/chunks/app/onboarding/page-7828c2c64e97362a.js → litellm/proxy/_experimental/out/_next/static/chunks/app/onboarding/page-d0d85032bb87ba51.js
+1 -1
View File
File diff suppressed because one or more lines are too long
litellm/proxy/_experimental/out/_next/static/chunks/app/page-338773f18570e0d6.js
-1
View File
File diff suppressed because one or more lines are too long

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