scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 154

Merge branch 'main' into litellm_dev_09_30_2025_p1

Browse source
This commit is contained in:
Krish Dholakia 2025-10-04 14:15:55 -07:00 • committed by GitHub
commit bed6c79bed
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
220 changed files with 16777 additions and 2615 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
.circleci/config.yml
View file

@ -616,6 +616,24 @@ jobs:
wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz
sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz
rm dockerize-linux-amd64-v0.6.1.tar.gz
- run:
name: Start PostgreSQL Database
command: |
docker run -d \
--name postgres-db \
-e POSTGRES_USER=postgres \
-e POSTGRES_PASSWORD=postgres \
-e POSTGRES_DB=circle_test \
-p 5432:5432 \
postgres:14
- run:
name: Wait for PostgreSQL to be ready
command: dockerize -wait tcp://localhost:5432 -timeout 1m
- run:
name: Set DATABASE_URL environment variable
command: |
echo 'export DATABASE_URL="postgresql://postgres:postgres@localhost:5432/circle_test"' >> $BASH_ENV
source $BASH_ENV
- run:
name: Run Security Scans
command: |

2
README.md
View file

@ -273,7 +273,7 @@ echo 'LITELLM_SALT_KEY="sk-1234"' >> .env
source .env
# Start
docker-compose up
docker compose up
```

8
docker/README.md
View file

@ -28,7 +28,7 @@ Replace `your-secret-key` with a strong, randomly generated secret.
Once you have set the `MASTER_KEY`, you can build and run the containers using the following command:
```bash
docker-compose up -d --build
docker compose up -d --build
```
This command will:
@ -42,13 +42,13 @@ This command will:
You can check the status of the running containers with the following command:
```bash
docker-compose ps
docker compose ps
```
To view the logs of the `litellm` container, run:
```bash
docker-compose logs -f litellm
docker compose logs -f litellm
```
### 4. Stopping the Application
@ -56,7 +56,7 @@ docker-compose logs -f litellm
To stop the running containers, use the following command:
```bash
docker-compose down
docker compose down
```
## Troubleshooting

2
docs/my-website/docs/adding_provider/new_rerank_provider.md
View file

@ -17,7 +17,7 @@ class YourProviderRerankConfig(BaseRerankConfig):
# ... other supported params
]
def transform_rerank_request(self, model: str, optional_rerank_params: OptionalRerankParams, headers: dict) -> dict:
def transform_rerank_request(self, model: str, optional_rerank_params: Dict, headers: dict) -> dict:
# Transform request to RerankRequest spec
return rerank_request.model_dump(exclude_none=True)

5
docs/my-website/docs/contributing.md
View file

@ -13,9 +13,6 @@ git clone https://github.com/BerriAI/litellm.git
Tell the proxy where the UI is located
```bash
export PROXY_BASE_URL="http://localhost:3000/"
### ALSO ### - set the basic env variables
DATABASE_URL = "postgresql://<user>:<password>@<host>:<port>/<dbname>"
LITELLM_MASTER_KEY = "sk-1234"
STORE_MODEL_IN_DB = "True"
@ -30,7 +27,7 @@ python3 proxy_cli.py --config /path/to/config.yaml --port 4000
Set the mode as development (this will assume the proxy is running on localhost:4000)
```bash
export NODE_ENV="development"
npm install # install dependencies
```
```bash

52
docs/my-website/docs/embedding/supported_embedding.md
View file

@ -266,7 +266,59 @@ print(response)
| Titan Embeddings - G1 | `embedding(model="amazon.titan-embed-text-v1", input=input)` |
| Cohere Embeddings - English | `embedding(model="cohere.embed-english-v3", input=input)` |
| Cohere Embeddings - Multilingual | `embedding(model="cohere.embed-multilingual-v3", input=input)` |
| TwelveLabs Marengo (Async) | `embedding(model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0", input=input, input_type="text")` | [Async Invoke Docs](../providers/bedrock_embedding#async-invoke-embedding) |
## TwelveLabs Bedrock Embedding Models
TwelveLabs Marengo models support multimodal embeddings (text, image, video, audio) and require the `input_type` parameter to specify the input format.
### Usage
```python
from litellm import embedding
import os
# Set AWS credentials
os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = "us-east-1"
# Text embedding
response = embedding(
model="bedrock/us.twelvelabs.marengo-embed-2-7-v1:0",
input=["Hello world from LiteLLM!"],
input_type="text" # Required parameter
)
# Image embedding (base64)
response = embedding(
model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0",
input=["data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..."],
input_type="image", # Required parameter
output_s3_uri="s3://your-bucket/async-invoke-output/"
)
# Video embedding (S3 URL)
response = embedding(
model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0",
input=["s3://your-bucket/video.mp4"],
input_type="video", # Required parameter
output_s3_uri="s3://your-bucket/async-invoke-output/"
)
```
### Required Parameters
| Parameter | Description | Values |
|-----------|-------------|--------|
| `input_type` | Type of input content | `"text"`, `"image"`, `"video"`, `"audio"` |
### Supported Models
| Model Name | Function Call | Notes |
|------------|---------------|-------|
| TwelveLabs Marengo 2.7 (Sync) | `embedding(model="bedrock/us.twelvelabs.marengo-embed-2-7-v1:0", input=input, input_type="text")` | Text embeddings only |
| TwelveLabs Marengo 2.7 (Async) | `embedding(model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0", input=input, input_type="text/image/video/audio")` | All input types, requires `output_s3_uri` |
## Cohere Embedding Models
https://docs.cohere.com/reference/embed

19
docs/my-website/docs/mcp.md
View file

@ -1045,6 +1045,25 @@ curl --location 'http://localhost:4000/github_mcp/mcp' \
---
## MCP Oauth
LiteLLM v 1.77.6 added support for OAuth 2.0 Client Credentials for MCP servers.
This configuration is currently available on the config.yaml, with UI support coming soon.
```yaml
mcp_servers:
github_mcp:
url: "https://api.githubcopilot.com/mcp"
auth_type: oauth2
authorization_url: https://github.com/login/oauth/authorize
token_url: https://github.com/login/oauth/access_token
client_id: os.environ/GITHUB_OAUTH_CLIENT_ID
client_secret: os.environ/GITHUB_OAUTH_CLIENT_SECRET
scopes: ["public_repo", "user:email"]
```
## Using your MCP with client side credentials
Use this if you want to pass a client side authentication token to LiteLLM to then pass to your MCP to auth to your MCP.

49
docs/my-website/docs/pass_through/vertex_ai.md
View file

@ -15,10 +15,11 @@ Pass-through endpoints for Vertex AI - call provider-specific endpoint, in nativ
## Supported Endpoints
LiteLLM supports 2 vertex ai passthrough routes:
LiteLLM supports 3 vertex ai passthrough routes:
1. `/vertex_ai` → routes to `https://{vertex_location}-aiplatform.googleapis.com/`
2. `/vertex_ai/discovery` → routes to [`https://discoveryengine.googleapis.com`](https://discoveryengine.googleapis.com/)
3. `/vertex_ai/live` → upgrades to the Vertex AI Live API WebSocket (`google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent`)
## How to use
@ -170,6 +171,50 @@ generateContent();
</Tabs>
## Vertex AI Live API WebSocket
LiteLLM can now proxy the Vertex AI Live API to help you experiment with streaming audio/text from Gemini Live models without exposing Google credentials to clients.
- Configure default Vertex credentials via `default_vertex_config` or environment variables (see examples above).
- Connect to `wss://<PROXY_URL>/vertex_ai/live`. LiteLLM will exchange your saved credentials for a short-lived access token and forward messages bidirectionally.
- Optional query params `vertex_project`, `vertex_location`, and `model` let you override defaults for multi-project setups or global-only models.
```python title="client.py"
import asyncio
import json
from websockets.asyncio.client import connect
async def main() -> None:
headers = {
"x-litellm-api-key": "Bearer sk-your-litellm-key",
"Content-Type": "application/json",
}
async with connect(
"ws://localhost:4000/vertex_ai/live",
additional_headers=headers,
) as ws:
await ws.send(
json.dumps(
{
"setup": {
"model": "projects/your-project/locations/us-central1/publishers/google/models/gemini-2.0-flash-live-preview-04-09",
"generation_config": {"response_modalities": ["TEXT"]},
}
}
)
)
async for message in ws:
print("server:", message)
if __name__ == "__main__":
asyncio.run(main())
```
## Quick Start
Let's call the Vertex AI [`/generateContent` endpoint](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference)
@ -415,4 +460,4 @@ generateContent();
```
</TabItem>
</Tabs>
</Tabs>

284
docs/my-website/docs/pass_through/vertex_ai_live_websocket.md Normal file
View file

@ -0,0 +1,284 @@
# Vertex AI Live API WebSocket Passthrough
LiteLLM now supports WebSocket passthrough for the Vertex AI Live API, enabling real-time bidirectional communication with Gemini models.
## Overview
The Vertex AI Live API WebSocket passthrough allows you to:
- Connect to Vertex AI Live API through LiteLLM proxy
- Use existing Vertex AI authentication methods
- Pass through all WebSocket messages bidirectionally
- Support text, audio, video, and multimodal interactions
- Track costs automatically for all usage types
## Configuration
### Environment Variables
Set the following environment variables for Vertex AI authentication:
```bash
# Required
DEFAULT_VERTEXAI_PROJECT=your-project-id
DEFAULT_VERTEXAI_LOCATION=us-central1
# Optional - use one of these for authentication
DEFAULT_GOOGLE_APPLICATION_CREDENTIALS=/path/to/service-account.json
# OR run: gcloud auth application-default login
```
### Configuration File
Alternatively, configure in your `config.yaml`:
```yaml
litellm_settings:
default_vertex_config:
vertex_project: "your-project-id"
vertex_location: "us-central1"
vertex_credentials: "os.environ/GOOGLE_APPLICATION_CREDENTIALS"
```
## Usage
### WebSocket Endpoints
- `ws://your-proxy-host/v1/vertex-ai/live`
- `ws://your-proxy-host/vertex-ai/live`
### Query Parameters
- `project_id` (optional): Google Cloud project ID (can be set in config)
- `location` (optional): Vertex AI location (can be set in config, default: us-central1)
### Example Connection
```javascript
// If project_id and location are set in config, you can connect without query params
const ws = new WebSocket('ws://localhost:4000/v1/vertex-ai/live');
// Or specify them explicitly
const ws = new WebSocket('ws://localhost:4000/v1/vertex-ai/live?project_id=your-project-id&location=us-central1');
```
## Cost Tracking
The WebSocket passthrough automatically tracks costs for all usage types based on the [Vertex AI pricing](https://cloud.google.com/vertex-ai/generative-ai/pricing#model-optimizer-pricing):
### Supported Cost Tracking
- **Text**: Character-based or token-based pricing depending on model
- **Audio**: Per-second pricing for audio input/output
- **Video**: Per-second pricing for video input
- **Images**: Per-image pricing for image input
### Cost Calculation
Costs are calculated using the same methods as other Vertex AI models in LiteLLM:
- Uses `cost_per_character` for Gemini models
- Uses `cost_per_token` for partner models (Claude, Llama, etc.)
- Includes audio, video, and image costs when applicable
### Cost Logging
Costs are automatically logged to:
- LiteLLM proxy logs
- Database (if configured)
- Spend tracking system
- Admin dashboard
Example log output:
```
Vertex AI Live WebSocket session cost: $0.001234 (input: $0.000800, output: $0.000434) tokens: 150, characters: 1200, duration: 45.2s
```
## API Reference
### Setup Message
Send this message first to initialize the session:
```json
{
"setup": {
"model": "projects/your-project-id/locations/us-central1/publishers/google/models/gemini-2.0-flash-live-preview-04-09",
"generation_config": {
"response_modalities": ["TEXT"]
}
}
}
```
### Text Input
```json
{
"client_content": {
"turns": [
{
"role": "user",
"parts": [{"text": "Hello! How are you?"}]
}
],
"turn_complete": true
}
}
```
### Audio Input
```json
{
"realtime_input": {
"media_chunks": [
{
"data": "base64-encoded-audio-data",
"mime_type": "audio/pcm"
}
]
}
}
```
## Supported Features
### Response Modalities
- **TEXT**: Text responses
- **AUDIO**: Audio responses with voice synthesis
### Tools
- **Function Calling**: Define and use custom functions
- **Code Execution**: Execute Python code
- **Google Search**: Search the web
- **Voice Activity Detection**: Detect when user is speaking
### Advanced Features
- **Audio Transcription**: Transcribe input and output audio
- **Proactive Audio**: Model responds only when relevant
- **Affective Dialog**: Understand emotional expressions
## Examples
### Python Client
```python
import asyncio
import json
import websockets
async def chat_with_gemini():
uri = "ws://localhost:4000/v1/vertex-ai/live?project_id=your-project-id"
async with websockets.connect(uri) as websocket:
# Setup
setup = {
"setup": {
"model": "projects/your-project-id/locations/us-central1/publishers/google/models/gemini-2.0-flash-live-preview-04-09",
"generation_config": {"response_modalities": ["TEXT"]}
}
}
await websocket.send(json.dumps(setup))
# Wait for setup response
response = await websocket.recv()
print(f"Setup: {response}")
# Send message
message = {
"client_content": {
"turns": [{"role": "user", "parts": [{"text": "Hello!"}]}],
"turn_complete": True
}
}
await websocket.send(json.dumps(message))
# Receive response
async for response in websocket:
print(f"Response: {response}")
# Check if turn is complete
data = json.loads(response)
if data.get("serverContent", {}).get("turnComplete"):
break
asyncio.run(chat_with_gemini())
```
### JavaScript Client
```javascript
const ws = new WebSocket('ws://localhost:4000/v1/vertex-ai/live?project_id=your-project-id');
ws.onopen = function() {
// Send setup
const setup = {
setup: {
model: "projects/your-project-id/locations/us-central1/publishers/google/models/gemini-2.0-flash-live-preview-04-09",
generation_config: { response_modalities: ["TEXT"] }
}
};
ws.send(JSON.stringify(setup));
};
ws.onmessage = function(event) {
const data = JSON.parse(event.data);
console.log('Received:', data);
// Check if setup is complete
if (data.setupComplete) {
// Send a message
const message = {
client_content: {
turns: [{ role: "user", parts: [{ text: "Hello!" }] }],
turn_complete: true
}
};
ws.send(JSON.stringify(message));
}
};
```
## Error Handling
The WebSocket connection may close with these codes:
- `4001`: Vertex AI credentials not configured
- `4002`: Project ID not provided
- `1011`: Internal server error
## Authentication
The WebSocket passthrough uses the same authentication as other LiteLLM endpoints:
1. **API Key**: Pass `Authorization: Bearer your-api-key` header
2. **Vertex AI Credentials**: Set environment variables or config file
## Limitations
- Requires valid Google Cloud project with Vertex AI API enabled
- WebSocket connections are not persistent across server restarts
- Rate limits apply based on your Google Cloud quotas
## Troubleshooting
### Common Issues
1. **Authentication Error**: Ensure Vertex AI credentials are properly configured
2. **Project Not Found**: Verify the project ID exists and has Vertex AI enabled
3. **Connection Refused**: Check that the LiteLLM proxy server is running
### Debug Mode
Enable debug logging to see detailed connection information:
```bash
export LITELLM_LOG=DEBUG
```
## Related Documentation
- [Vertex AI Live API Reference](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/multimodal-live)
- [LiteLLM Proxy Configuration](../proxy/)
- [Vertex AI Passthrough Endpoints](./vertex_ai.md)

7
docs/my-website/docs/projects/Railtracks.md Normal file
View file

@ -0,0 +1,7 @@
# Railtracks
`Railtracks` is an open-source agentic framework that helps developers build resilient agentic systems offering local and remote monitoring tools.
- [Github](https://github.com/RailtownAI/railtracks)
- [Docs](https://railtownai.github.io/railtracks/)
- [Railtracks](https://railtracks.org/)

176
docs/my-website/docs/providers/bedrock_embedding.md
View file

@ -8,6 +8,182 @@
| 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) |
## Async Invoke Support
LiteLLM supports AWS Bedrock's async-invoke feature for embedding models that require asynchronous processing, particularly useful for large media files (video, audio) or when you need to process embeddings in the background.
### Supported Models
| Provider | Async Invoke Route | Use Case |
|----------|-------------------|----------|
| TwelveLabs Marengo | `bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0` | Video, audio, image, and text embeddings |
### Required Parameters
When using async-invoke, you must provide:
| Parameter | Description | Required |
|-----------|-------------|----------|
| `output_s3_uri` | S3 URI where the embedding results will be stored | ✅ Yes |
| `input_type` | Type of input: `"text"`, `"image"`, `"video"`, or `"audio"` | ✅ Yes |
| `aws_region_name` | AWS region for the request | ✅ Yes |
### Usage
#### Basic Async Invoke
```python
from litellm import embedding
# Text embedding with async-invoke
response = embedding(
model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0",
input=["Hello world from LiteLLM async invoke!"],
aws_region_name="us-east-1",
input_type="text",
output_s3_uri="s3://your-bucket/async-invoke-output/"
)
print(f"Job submitted! Invocation ARN: {response._hidden_params._invocation_arn}")
```
#### Video/Audio Embedding
```python
# Video embedding (requires async-invoke)
response = embedding(
model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0",
input=["s3://your-bucket/video.mp4"], # S3 URL for video
aws_region_name="us-east-1",
input_type="video",
output_s3_uri="s3://your-bucket/async-invoke-output/"
)
print(f"Video embedding job submitted! ARN: {response._hidden_params._invocation_arn}")
```
#### Image Embedding with Base64
```python
import base64
# Load and encode image
with open("image.jpg", "rb") as img_file:
img_data = base64.b64encode(img_file.read()).decode('utf-8')
img_base64 = f"data:image/jpeg;base64,{img_data}"
response = embedding(
model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0",
input=[img_base64],
aws_region_name="us-east-1",
input_type="image",
output_s3_uri="s3://your-bucket/async-invoke-output/"
)
```
### Retrieving Job Information
#### Getting Job ID and Invocation ARN
The async-invoke response includes the invocation ARN in the hidden parameters:
```python
response = embedding(
model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0",
input=["Hello world"],
aws_region_name="us-east-1",
input_type="text",
output_s3_uri="s3://your-bucket/async-invoke-output/"
)
# Access invocation ARN
invocation_arn = response._hidden_params._invocation_arn
print(f"Invocation ARN: {invocation_arn}")
# Extract job ID from ARN (last part after the last slash)
job_id = invocation_arn.split("/")[-1]
print(f"Job ID: {job_id}")
```
#### Checking Job Status
Use LiteLLM's `retrieve_batch` function to check if your job is still processing:
```python
from litellm import retrieve_batch
def check_async_job_status(invocation_arn, aws_region_name="us-east-1"):
"""Check the status of an async invoke job using LiteLLM batch API"""
try:
response = retrieve_batch(
batch_id=invocation_arn,
custom_llm_provider="bedrock",
aws_region_name=aws_region_name
)
return response
except Exception as e:
print(f"Error checking job status: {e}")
return None
# Check status
status = check_async_job_status(invocation_arn, "us-east-1")
if status:
print(f"Job Status: {status.status}")
print(f"Output Location: {status.output_file_id}")
```
**Note:** The actual embedding results are stored in S3. The `output_file_id` from the batch status can be used to locate the results file in your S3 bucket.
### Error Handling
#### Common Errors
| Error | Cause | Solution |
|-------|-------|----------|
| `ValueError: output_s3_uri cannot be empty` | Missing S3 output URI | Provide a valid S3 URI |
| `ValueError: Input type 'video' requires async_invoke route` | Using video/audio without async-invoke | Use `bedrock/async_invoke/` model prefix |
| `ValueError: input_type is required` | Missing input type parameter | Specify `input_type` parameter |
#### Example Error Handling
```python
try:
response = embedding(
model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0",
input=["Hello world"],
aws_region_name="us-east-1",
input_type="text",
output_s3_uri="s3://your-bucket/output/" # Required for async-invoke
)
print("Job submitted successfully!")
except ValueError as e:
if "output_s3_uri cannot be empty" in str(e):
print("Error: Please provide a valid S3 output URI")
elif "requires async_invoke route" in str(e):
print("Error: Use async_invoke model for video/audio inputs")
else:
print(f"Error: {e}")
except Exception as e:
print(f"Unexpected error: {e}")
```
### Best Practices
1. **Use async-invoke for large files**: Video and audio files are better processed asynchronously
2. **Use LiteLLM batch API**: Use `retrieve_batch()` instead of direct Bedrock API calls for status checking
3. **Monitor job status**: Check job status periodically using the batch API to know when results are ready
4. **Handle errors gracefully**: Implement proper error handling for network issues and job failures
5. **Set appropriate timeouts**: Consider the processing time for large files
6. **Use S3 for large inputs**: For video/audio, use S3 URLs instead of base64 encoding
### Limitations
- Async-invoke is currently only supported for TwelveLabs Marengo models
- Results are stored in S3 and must be retrieved separately using the output file ID
- Job status checking requires using LiteLLM's `retrieve_batch()` function
- No built-in polling mechanism in LiteLLM (must implement your own status checking loop)
### API keys
This can be set as env variables or passed as **params to litellm.embedding()**
```python

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

@ -186,3 +186,6 @@ print("Available models:", [model['id'] for model in models.get('data', [])])
## Support
For more information regarding Lemonade please go to to the [Lemonade website](https://lemonade-server.ai/) or [Lemonade repository](https://github.com/lemonade-sdk/lemonade).
</TabItem>
</Tabs>

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

@ -15,8 +15,8 @@ https://docs.api.nvidia.com/nim/reference/
| Description | Nvidia NIM is a platform that provides a simple API for deploying and using AI models. LiteLLM supports all models from [Nvidia NIM](https://developer.nvidia.com/nim/) |
| Provider Route on LiteLLM | `nvidia_nim/` |
| Provider Doc | [Nvidia NIM Docs ↗](https://developer.nvidia.com/nim/) |
| API Endpoint for Provider | https://integrate.api.nvidia.com/v1/ |
| Supported OpenAI Endpoints | `/chat/completions`, `/completions`, `/responses`, `/embeddings` |
| API Endpoint for Provider | https://integrate.api.nvidia.com/v1/ (chat/embeddings), https://ai.api.nvidia.com/v1/ (rerank) |
| Supported OpenAI Endpoints | `/chat/completions`, `/completions`, `/responses`, `/embeddings`, `/rerank` |
## API Key
```python

261
docs/my-website/docs/providers/nvidia_nim_rerank.md Normal file
View file

@ -0,0 +1,261 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Nvidia NIM - Rerank
Use Nvidia NIM Rerank models through LiteLLM.
| Property | Details |
|----------|---------|
| Description | Nvidia NIM provides high-performance reranking models for semantic search and retrieval-augmented generation (RAG) |
| Provider Doc | [Nvidia NIM Rerank API ↗](https://docs.api.nvidia.com/nim/reference/nvidia-llama-3_2-nv-rerankqa-1b-v2-infer) |
| Supported Endpoint | `/rerank` |
## Overview
Nvidia NIM rerank models help you:
- Reorder search results by relevance to a query
- Improve RAG (Retrieval-Augmented Generation) accuracy
- Filter and rank large document sets efficiently
**Supported Models:**
- All Nvidia NIM rerank models on their platform
:::tip
See the full list of LiteLLM supported Nvidia NIM rerank models on [Nvidia NIM](https://models.litellm.ai)
:::
## Usage
### LiteLLM Python SDK
<Tabs>
<TabItem value="llama-1b" label="LLaMa 1B Model">
```python
import litellm
import os
os.environ['NVIDIA_NIM_API_KEY'] = "nvapi-..."
response = litellm.rerank(
model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
query="What is the GPU memory bandwidth of H100 SXM?",
documents=[
"The Hopper GPU is paired with the Grace CPU using NVIDIA's ultra-fast chip-to-chip interconnect, delivering 900GB/s of bandwidth.",
"A100 provides up to 20X higher performance over the prior generation.",
"Accelerated servers with H100 deliver 3 terabytes per second (TB/s) of memory bandwidth per GPU."
],
top_n=3,
)
print(response)
```
</TabItem>
<TabItem value="mistral-4b" label="Mistral 4B Model">
```python
import litellm
import os
os.environ['NVIDIA_NIM_API_KEY'] = "nvapi-..."
response = litellm.rerank(
model="nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3",
query="What is the GPU memory bandwidth of H100 SXM?",
documents=[
"The Hopper GPU is paired with the Grace CPU using NVIDIA's ultra-fast chip-to-chip interconnect, delivering 900GB/s of bandwidth.",
"A100 provides up to 20X higher performance over the prior generation.",
"Accelerated servers with H100 deliver 3 terabytes per second (TB/s) of memory bandwidth per GPU."
],
top_n=3,
)
print(response)
```
</TabItem>
</Tabs>
**Response:**
```json
{
"results": [
{
"index": 2,
"relevance_score": 6.828125,
"document": {
"text": "Accelerated servers with H100 deliver 3 terabytes per second (TB/s) of memory bandwidth per GPU."
}
},
{
"index": 0,
"relevance_score": -1.564453125,
"document": {
"text": "The Hopper GPU is paired with the Grace CPU using NVIDIA's ultra-fast chip-to-chip interconnect, delivering 900GB/s of bandwidth."
}
}
]
}
```
## Usage with LiteLLM Proxy
### 1. Setup Config
Add Nvidia NIM rerank models to your proxy configuration:
```yaml
model_list:
- model_name: nvidia-rerank
litellm_params:
model: nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2
api_key: os.environ/NVIDIA_NIM_API_KEY
```
### 2. Start Proxy
```bash
litellm --config /path/to/config.yaml
```
### 3. Make Rerank Requests
```bash
curl -X POST http://0.0.0.0:4000/rerank \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"model": "nvidia-rerank",
"query": "What is the GPU memory bandwidth of H100?",
"documents": [
"H100 delivers 3TB/s memory bandwidth",
"A100 has 2TB/s memory bandwidth",
"V100 offers 900GB/s memory bandwidth"
],
"top_n": 2
}'
```
## API Parameters
### Required Parameters
| Parameter | Type | Description |
|-----------|------|-------------|
| `model` | string | The Nvidia NIM rerank model name with `nvidia_nim/` prefix |
| `query` | string | The search query to rank documents against |
| `documents` | array | List of documents to rank (1-1000 documents) |
### Optional Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `top_n` | integer | All documents | Number of top-ranked documents to return |
### Nvidia-Specific Parameters
**`truncate`**: Controls how text is truncated if it exceeds the model's context window
- `"NONE"`: No truncation (request may fail if too long)
- `"END"`: Truncate from the end of the text
```python
response = litellm.rerank(
model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
query="GPU performance",
documents=["High performance computing", "Fast GPU processing"],
top_n=2,
truncate="END", # Nvidia-specific parameter
)
```
## Authentication
Set your Nvidia NIM API key:
<Tabs>
<TabItem value="env" label="Environment Variable">
```bash
export NVIDIA_NIM_API_KEY="nvapi-..."
```
</TabItem>
<TabItem value="python" label="Python">
```python
import os
os.environ['NVIDIA_NIM_API_KEY'] = "nvapi-..."
# Or pass directly
response = litellm.rerank(
model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
query="test",
documents=["doc1"],
api_key="nvapi-...",
)
```
</TabItem>
</Tabs>
## API Endpoint
The rerank endpoint uses a different base URL than chat/embeddings:
- **Chat/Embeddings:** `https://integrate.api.nvidia.com/v1/`
- **Rerank:** `https://ai.api.nvidia.com/v1/`
LiteLLM automatically uses the correct endpoint for rerank requests.
### Custom API Base URL
You can override the default base URL in several ways:
**Option 1: Environment Variable**
```bash
export NVIDIA_NIM_API_BASE="https://your-custom-endpoint.com"
```
**Option 2: Pass as parameter**
```python
response = litellm.rerank(
model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
query="test",
documents=["doc1"],
api_base="https://your-custom-endpoint.com",
)
```
**Option 3: Full URL (including model path)**
If you have the complete endpoint URL, you can pass it directly:
```python
response = litellm.rerank(
model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
query="test",
documents=["doc1"],
api_base="https://your-custom-endpoint.com/v1/retrieval/nvidia/llama-3_2-nv-rerankqa-1b-v2/reranking",
)
```
LiteLLM will detect the full URL (by checking for `/retrieval/` in the path) and use it as-is.
### How do I get an API key?
Get your Nvidia NIM API key from [Nvidia's website](https://developer.nvidia.com/nim/).
## Related Documentation
- [Nvidia NIM - Main Documentation](./nvidia_nim)
- [Nvidia NIM Chat Completions](./nvidia_nim#sample-usage)
- [LiteLLM Rerank Endpoint](../rerank)
- [Nvidia NIM Official Docs ↗](https://docs.api.nvidia.com/nim/reference/)

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

@ -44,6 +44,7 @@ response = completion(
oci_user=<your_oci_user>,
oci_fingerprint=<your_oci_fingerprint>,
oci_tenancy=<your_oci_tenancy>,
oci_serving_mode="ON_DEMAND", # Optional, default is "ON_DEMAND". Other option is "DEDICATED"
# Provide either the private key string OR the path to the key file:
# Option 1: pass the private key as a string
oci_key=<string_with_content_of_oci_key>,
@ -71,6 +72,7 @@ response = completion(
oci_user=<your_oci_user>,
oci_fingerprint=<your_oci_fingerprint>,
oci_tenancy=<your_oci_tenancy>,
oci_serving_mode="ON_DEMAND", # Optional, default is "ON_DEMAND". Other option is "DEDICATED"
# Provide either the private key string OR the path to the key file:
# Option 1: pass the private key as a string
oci_key=<string_with_content_of_oci_key>,

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

@ -621,6 +621,163 @@ curl -X POST 'http://0.0.0.0:4000/chat/completions' \
#### **Google Maps**
Use Google Maps to provide location-based context to your Gemini models.
[**Relevant Vertex AI Docs**](https://ai.google.dev/gemini-api/docs/grounding#google-maps)
<Tabs>
<TabItem value="sdk" label="SDK">
**Basic Usage - Enable Widget Only**
```python showLineNumbers
from litellm import completion
## SETUP ENVIRONMENT
# !gcloud auth application-default login - run this to add vertex credentials to your env
tools = [{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}] # 👈 ADD GOOGLE MAPS
resp = litellm.completion(
model="vertex_ai/gemini-2.0-flash",
messages=[{"role": "user", "content": "What restaurants are nearby?"}],
tools=tools,
)
print(resp)
```
**With Location Data**
You can specify a location to ground the model's responses with location-specific information:
```python showLineNumbers
from litellm import completion
## SETUP ENVIRONMENT
# !gcloud auth application-default login - run this to add vertex credentials to your env
tools = [{
"googleMaps": {
"enableWidget": "ENABLE_WIDGET",
"latitude": 37.7749, # San Francisco latitude
"longitude": -122.4194, # San Francisco longitude
"languageCode": "en_US" # Optional: language for results
}
}] # 👈 ADD GOOGLE MAPS WITH LOCATION
resp = litellm.completion(
model="vertex_ai/gemini-2.0-flash",
messages=[{"role": "user", "content": "What restaurants are nearby?"}],
tools=tools,
)
print(resp)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
<Tabs>
<TabItem value="openai" label="OpenAI Python SDK">
**Basic Usage - Enable Widget Only**
```python showLineNumbers
from openai import OpenAI
client = OpenAI(
api_key="sk-1234", # pass litellm proxy key, if you're using virtual keys
base_url="http://0.0.0.0:4000/v1/" # point to litellm proxy
)
response = client.chat.completions.create(
model="gemini-2.0-flash",
messages=[{"role": "user", "content": "What restaurants are nearby?"}],
tools=[{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}],
)
print(response)
```
**With Location Data**
```python showLineNumbers
from openai import OpenAI
client = OpenAI(
api_key="sk-1234", # pass litellm proxy key, if you're using virtual keys
base_url="http://0.0.0.0:4000/v1/" # point to litellm proxy
)
response = client.chat.completions.create(
model="gemini-2.0-flash",
messages=[{"role": "user", "content": "What restaurants are nearby?"}],
tools=[{
"googleMaps": {
"enableWidget": "ENABLE_WIDGET",
"latitude": 37.7749, # San Francisco latitude
"longitude": -122.4194, # San Francisco longitude
"languageCode": "en_US" # Optional: language for results
}
}],
)
print(response)
```
</TabItem>
<TabItem value="curl" label="cURL">
**Basic Usage - Enable Widget Only**
```bash showLineNumbers
curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-1234" \
-d '{
"model": "gemini-2.0-flash",
"messages": [
{"role": "user", "content": "What restaurants are nearby?"}
],
"tools": [
{
"googleMaps": {"enableWidget": "ENABLE_WIDGET"}
}
]
}'
```
**With Location Data**
```bash showLineNumbers
curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-1234" \
-d '{
"model": "gemini-2.0-flash",
"messages": [
{"role": "user", "content": "What restaurants are nearby?"}
],
"tools": [
{
"googleMaps": {
"enableWidget": "ENABLE_WIDGET",
"latitude": 37.7749,
"longitude": -122.4194,
"languageCode": "en_US"
}
}
]
}'
```
</TabItem>
</Tabs>
</TabItem>
</Tabs>
#### **Moving from Vertex AI SDK to LiteLLM (GROUNDING)**

4
docs/my-website/docs/proxy/cost_tracking.md
View file

@ -8,6 +8,10 @@ Track spend for keys, users, and teams across 100+ LLMs.
LiteLLM automatically tracks spend for all known models. See our [model cost map](https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json)
:::tip Keep Pricing Data Updated
[Sync model pricing data from GitHub](../sync_models_github.md) to ensure accurate cost tracking.
:::
### How to Track Spend with LiteLLM
**Step 1**

21
docs/my-website/docs/proxy/deploy.md
View file

@ -27,7 +27,7 @@ echo 'LITELLM_SALT_KEY="sk-1234"' >> .env
source .env
# Start
docker-compose up
docker compose up
```
@ -715,6 +715,25 @@ docker run ghcr.io/berriai/litellm:main-stable
```
### Restart Workers After N Requests
Use this to mitigate memory growth by recycling workers after a fixed number of requests. When set, each worker restarts after completing the specified number of requests. Defaults to disabled when unset.
Usage Examples:
```shell showLineNumbers title="docker run (CLI flag)"
docker run ghcr.io/berriai/litellm:main-stable \
--max_requests_before_restart 10000
```
Or set via environment variable:
```shell showLineNumbers title="Environment Variable"
export MAX_REQUESTS_BEFORE_RESTART=10000
docker run ghcr.io/berriai/litellm:main-stable
```
### 5. config.yaml file on s3, GCS Bucket Object/url
Use this if you cannot mount a config file on your deployment service (example - AWS Fargate, Railway etc)

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

@ -55,7 +55,7 @@ echo 'LITELLM_SALT_KEY="sk-1234"' >> .env
source .env
# Start
docker-compose up
docker compose up
```
</TabItem>

91
docs/my-website/docs/proxy/logging_spec.md
View file

@ -14,6 +14,7 @@ Found under `kwargs["standard_logging_object"]`. This is a standard payload, log
| `cost_breakdown` | `Optional[CostBreakdown]` | Detailed cost breakdown object |
| `response_cost_failure_debug_info` | `StandardLoggingModelCostFailureDebugInformation` | Debug information if cost tracking fails |
| `status` | `StandardLoggingPayloadStatus` | Status of the payload |
| `status_fields` | `StandardLoggingPayloadStatusFields` | Typed status fields for easy filtering and analytics |
| `total_tokens` | `int` | Total number of tokens |
| `prompt_tokens` | `int` | Number of prompt tokens |
| `completion_tokens` | `int` | Number of completion tokens |
@ -162,17 +163,89 @@ A literal type with two possible values:
## StandardLoggingGuardrailInformation
| Field | Type | Description |
|-----------------------|------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| `guardrail_name` | `Optional[str]` | Guardrail name |
| `guardrail_provider` | `Optional[str]` | Guardrail provider |
| `guardrail_mode` | `Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks]]]` | Guardrail mode |
| `guardrail_request` | `Optional[dict]` | Guardrail request |
| `guardrail_response` | `Optional[Union[dict, str, List[dict]]]` | Guardrail response |
| `guardrail_status` | `Literal["success", "failure", "blocked"]` | Guardrail execution status: `success` = no violations detected, `blocked` = content blocked/modified due to policy violations, `failure` = technical error or API failure |
| `start_time` | `Optional[float]` | Start time of the guardrail |
| `end_time` | `Optional[float]` | End time of the guardrail |
| `duration` | `Optional[float]` | Duration of the guardrail in seconds |
| `masked_entity_count` | `Optional[Dict[str, int]]` | Count of masked entities |
## StandardLoggingPayloadStatusFields
Typed status fields for easy filtering and analytics.
| Field | Type | Description |
|-------|------|-------------|
| `guardrail_name` | `Optional[str]` | Guardrail name |
| `guardrail_mode` | `Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks]]]` | Guardrail mode |
| `guardrail_request` | `Optional[dict]` | Guardrail request |
| `guardrail_response` | `Optional[Union[dict, str, List[dict]]]` | Guardrail response |
| `guardrail_status` | `Literal["success", "failure"]` | Guardrail status |
| `start_time` | `Optional[float]` | Start time of the guardrail |
| `end_time` | `Optional[float]` | End time of the guardrail |
| `duration` | `Optional[float]` | Duration of the guardrail in seconds |
| `masked_entity_count` | `Optional[Dict[str, int]]` | Count of masked entities |
| `llm_api_status` | `StandardLoggingPayloadStatus` | Status of the LLM API call: `"success"` if completed successfully, `"failure"` if errored |
| `guardrail_status` | `GuardrailStatus` | Status of guardrail execution (see below) |
### StandardLoggingPayloadStatus
A literal type with two possible values:
- `"success"` - The LLM API request completed successfully
- `"failure"` - The LLM API request failed
### GuardrailStatus
A literal type with four possible values:
- `"success"` - Guardrail ran and allowed content through (no violations detected)
- `"guardrail_intervened"` - Guardrail blocked or modified content due to policy violations
- `"guardrail_failed_to_respond"` - Guardrail had a technical failure or API error
- `"not_run"` - No guardrail was executed for this request
### Usage Examples
Filter logs for requests where guardrails intervened:
```json
{
"status_fields": {
"guardrail_status": "guardrail_intervened"
}
}
```
Find guardrail technical failures:
```json
{
"status_fields": {
"guardrail_status": "guardrail_failed_to_respond"
}
}
```
Get successful LLM requests:
```json
{
"status_fields": {
"llm_api_status": "success"
}
}
```
Find requests where guardrails ran successfully without intervention:
```json
{
"status_fields": {
"guardrail_status": "success",
"llm_api_status": "success"
}
}
```
Find requests where no guardrail was run:
```json
{
"status_fields": {
"guardrail_status": "not_run"
}
}
```
## StandardLoggingPromptManagementMetadata

4
docs/my-website/docs/proxy/model_management.md
View file

@ -19,6 +19,10 @@ model_list:
Retrieve detailed information about each model listed in the `/model/info` endpoint, including descriptions from the `config.yaml` file, and additional model info (e.g. max tokens, cost per input token, etc.) pulled from the model_info you set and the [litellm model cost map](https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json). Sensitive details like API keys are excluded for security purposes.
:::tip Sync Model Data
Keep your model pricing data up to date by [syncing models from GitHub](../sync_models_github.md).
:::
<Tabs
defaultValue="curl"
values={[

75
docs/my-website/docs/proxy/native_litellm_prompt.md
View file

@ -9,7 +9,7 @@ Store prompts as `.prompt` files in your repository and use them directly with L
- **File System**: Store `.prompt` files locally
- **BitBucket**: Store `.prompt` files in BitBucket repositories with team-based access control
- **Gitlab**: Store `.prompt` files in Gitlab repositories with team-based access control
## Quick Start
<Tabs>
@ -90,6 +90,51 @@ response = litellm.completion(
```
</TabItem>
<TabItem value="gitlab" label="GITLAB">
**1. Create a .prompt file in a gitlab repo**
Create `prompts/hello.prompt` in your gitlab repository:
```yaml
---
model: gpt-4
temperature: 0.7
---
System: You are a helpful assistant.
User: {{user_message}}
```
**2. Configure Gitlab access**
```python
import litellm
# Configure gitlab access
gitlab_config = {
"workspace": "your-workspace",
"repository": "your-repo",
"access_token": "your-access-token",
"branch": "main"
}
# Set global gitlab configuration
litellm.set_global_gitlab_config(gitlab_config)
```
**3. Use with LiteLLM**
```python
response = litellm.completion(
model="gitlab/gpt-4",
prompt_id="hello",
prompt_variables={"user_message": "What is the capital of France?"}
)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
**1. Create a .prompt file**
@ -124,6 +169,12 @@ litellm_settings:
repository: "your-repo"
access_token: "your-access-token"
branch: "main"
# Or use Gitlab for team-based prompt management
global_gitlab_config:
workspace: "your-workspace"
repository: "your-repo"
access_token: "your-access-token"
branch: "main"
```
**3. Start the proxy**
@ -213,6 +264,14 @@ prompt_variables: Optional[dict] # optional - variables for template rendering
bitbucket_config: Optional[dict] # optional - BitBucket configuration (if not set globally)
```
**Gitlab:**
```
model: gitlab/<base_model> # required (e.g., gitlab/gpt-4)
prompt_id: str # required - the .prompt filename without extension
prompt_variables: Optional[dict] # optional - variables for template rendering
gitlab_config: Optional[dict] # optional - Gitlab configuration (if not set globally)
```
**Example API calls:**
```python
@ -235,4 +294,18 @@ response = litellm.completion(
"access_token": "your-token"
}
)
# Gitlab integration
response = litellm.completion(
model="gitlab/gpt-4",
prompt_id="hello",
prompt_variables={"user_message": "Hello world"},
gitlab_config={
"project": "a/b/<repo_name>",
"access_token": "your-access-token",
"base_url": "gitlab url",
"prompts_path": "src/prompts", # folder to point to, defaults to root
"branch":"main" # optional, defaults to main
}
)
```

12
docs/my-website/docs/proxy/pass_through.md
View file

@ -243,6 +243,18 @@ curl --location 'http://0.0.0.0:4000/v1/messages' \
}'
```
---
## Tutorial - Add Azure OpenAI Assistants API as a Pass Through Endpoint
In this video, we'll add the Azure OpenAI Assistants API as a pass through endpoint to LiteLLM Proxy.
<iframe width="840" height="500" src="https://www.loom.com/embed/12965cb299d24fc0bd7b6b413ab6d0ad" frameborder="0" webkitallowfullscreen mozallowfullscreen allowfullscreen></iframe>
<br/>
<br/>
---
## Troubleshooting

10
docs/my-website/docs/proxy/prod.md
View file

@ -71,6 +71,16 @@ Use this Docker `CMD`. This will start the proxy with 1 Uvicorn Async Worker
CMD ["--port", "4000", "--config", "./proxy_server_config.yaml"]
```
> Optional: If you observe gradual memory growth under sustained load, consider recycling workers after a fixed number of requests to mitigate leaks. Set this via CLI or environment variable:
```shell
# CLI
CMD ["--port", "4000", "--config", "./proxy_server_config.yaml", "--max_requests_before_restart", "10000"]
# or ENV (for deployment manifests / containers)
export MAX_REQUESTS_BEFORE_RESTART=10000
```
## 4. Use Redis 'port','host', 'password'. NOT 'redis_url'

61
docs/my-website/docs/proxy/sync_models_github.md Normal file
View file

@ -0,0 +1,61 @@
# Syncing Models to GitHub model_context_window
Sync model pricing data from GitHub's `model_prices_and_context_window.json` file outside of the LiteLLM UI.
> **📹 Video Tutorial**: [Watch how to sync models via the Admin UI](https://www.loom.com/share/ba41acc1882d41b284bbddbb0e9c27ce?sid=bdae351e-2026-4e39-932b-fcb185ff612c)
## Quick Start
**Manual sync:**
```bash
curl -X POST "https://your-proxy-url/reload/model_cost_map" \
-H "Authorization: Bearer YOUR_ADMIN_TOKEN" \
-H "Content-Type: application/json"
```
**Automatic sync every 6 hours:**
```bash
curl -X POST "https://your-proxy-url/schedule/model_cost_map_reload?hours=6" \
-H "Authorization: Bearer YOUR_ADMIN_TOKEN" \
-H "Content-Type: application/json"
```
## API Endpoints
| Endpoint | Method | Description |
|----------|--------|-------------|
| `/reload/model_cost_map` | POST | Manual sync |
| `/schedule/model_cost_map_reload?hours={hours}` | POST | Schedule periodic sync |
| `/schedule/model_cost_map_reload` | DELETE | Cancel scheduled sync |
| `/schedule/model_cost_map_reload/status` | GET | Check sync status |
**Authentication:** Requires admin role or master key
## Python Example
```python
import requests
def sync_models(proxy_url, admin_token):
response = requests.post(
f"{proxy_url}/reload/model_cost_map",
headers={"Authorization": f"Bearer {admin_token}"}
)
return response.json()
# Usage
result = sync_models("https://your-proxy-url", "your-admin-token")
print(result['message'])
```
## Configuration
**Custom model cost map URL:**
```bash
export LITELLM_MODEL_COST_MAP_URL="https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
```
**Use local model cost map:**
```bash
export LITELLM_LOCAL_MODEL_COST_MAP=True
```

14
docs/my-website/docs/proxy/ui.md
View file

@ -54,6 +54,20 @@ Allow others to create/delete their own keys.
[**Go Here**](./self_serve.md)
## Model Management
The Admin UI provides comprehensive model management capabilities:
- **Add Models**: Add new models through the UI without restarting the proxy
- **Model Hub**: Make models public for developers to discover available models
- **Price Data Sync**: Keep model pricing data up to date by syncing from GitHub
For detailed information on model management, see [Model Management](./model_management.md).
:::tip Sync Model Pricing Data
[Sync model pricing data from GitHub](./sync_models_github.md) to keep your model cost information current.
:::
## Disable Admin UI
Set `DISABLE_ADMIN_UI="True"` in your environment to disable the Admin UI.

25
docs/my-website/release_notes/v1.75.5-stable/index.md
View file

@ -51,6 +51,31 @@ pip install litellm==1.75.5.post2
- **Digital Ocean's Gradient AI** - New LLM provider for calling models on Digital Ocean's Gradient AI platform.
### 54% RPS Improvement
Throughput increased by 54% (1,040 → 1,602 RPS, aggregated) per instance while maintaining a 40 ms median overhead. The improvement comes from fixing major O(n²) inefficiencies in the router, primarily caused by repeated use of in statements inside loops over large arrays. Tests were run with a database-only setup (no cache hits). As a result, p95 latency improved by 30% (2,700 → 1,900 ms), enhancing overall stability and scalability under heavy load.
---
### Test Setup
All benchmarks were executed using Locust with 1,000 concurrent users and a ramp-up of 500. The environment was configured to stress the routing layer and eliminate caching as a variable.
**System Specs**
- **CPU:** 8 vCPUs
- **Memory:** 32 GB RAM
**Configuration (config.yaml)**
View the complete configuration: [gist.github.com/AlexsanderHamir/config.yaml](https://gist.github.com/AlexsanderHamir/53f7d554a5d2afcf2c4edb5b6be68ff4)
**Load Script (no_cache_hits.py)**
View the complete load testing script: [gist.github.com/AlexsanderHamir/no_cache_hits.py](https://gist.github.com/AlexsanderHamir/42c33d7a4dc7a57f56a78b560dee3a42)
---
### Risk of Upgrade
If you build the proxy from the pip package, you should hold off on upgrading. This version makes `prisma migrate deploy` our default for managing the DB. This is safer, as it doesn't reset the DB, but it requires a manual `prisma generate` step.

298
docs/my-website/release_notes/v1.77.7-stable/index.md Normal file
View file

@ -0,0 +1,298 @@
---
title: "[Preview] v1.77.7-stable - Claude Sonnet 4.5"
slug: "v1-77-7"
date: 2025-10-04T10: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
- name: Alexsander Hamir
title: Backend Performance Engineer
url: https://www.linkedin.com/in/alexsander-baptista/
image_url: https://media.licdn.com/dms/image/v2/D5603AQGXnziu4kqNCQ/profile-displayphoto-crop_800_800/B56ZkxEcuOKEAI-/0/1757464874550?e=1762387200&v=beta&t=9SNXLsWhx8OnYPAMQ9fqAr02oevDYEAL2vMYg2f9ieg
- name: Achintya Srivastava
title: Fullstack Engineer
url: https://www.linkedin.com/in/achintya-rajan/
image_url: https://media.licdn.com/dms/image/v2/D5603AQGdkEeyJTdljw/profile-displayphoto-shrink_800_800/profile-displayphoto-shrink_800_800/0/1716271140869?e=1762387200&v=beta&t=9gOoLPeqR2E5z3KSX61EUj3HVZXmgo87vhVuSHeffjc
- name: Sameer Kankute
title: Backend Engineer (LLM Translation)
url: https://www.linkedin.com/in/sameer-kankute/
image_url: https://media.licdn.com/dms/image/v2/D4D03AQHB_loQYd5gjg/profile-displayphoto-shrink_800_800/profile-displayphoto-shrink_800_800/0/1719137160975?e=1762387200&v=beta&t=0jbuX-f4eSnDxBY3olI6meuYr-LMbObhFmFbRcKF5mY
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:v1.77.7.rc.1
```
</TabItem>
<TabItem value="pip" label="Pip">
``` showLineNumbers title="pip install litellm"
pip install litellm==1.77.7.rc.1
```
</TabItem>
</Tabs>
---
## Key Highlights
- **Dynamic Rate Limiter v3** - Automatically maximizes throughput when capacity is available (< 80% saturation) by allowing lower-priority requests to use unused capacity, then switches to fair priority-based allocation under high load (≥ 80%) to prevent blocking
- **Major Performance Improvements** - Router optimization reducing P99 latency by 62.5%, cache improvements from O(n*log(n)) to O(log(n))
- **Claude Sonnet 4.5** - Support for Anthropic's new Claude Sonnet 4.5 model family with 200K+ context and tiered pricing
- **MCP Gateway Enhancements** - Fine-grained tool control, server permissions, and forwardable headers
- **AMD Lemonade & Nvidia NIM** - New provider support for AMD Lemonade and Nvidia NIM Rerank
- **GitLab Prompt Management** - GitLab-based prompt management integration
## New Models / Updated Models
#### New Model Support
| Provider | Model | Context Window | Input ($/1M tokens) | Output ($/1M tokens) | Features |
| -------- | ----- | -------------- | ------------------- | -------------------- | -------- |
| Anthropic | `claude-sonnet-4-5` | 200K | $3.00 | $15.00 | Chat, reasoning, vision, function calling, prompt caching |
| Anthropic | `claude-sonnet-4-5-20250929` | 200K | $3.00 | $15.00 | Chat, reasoning, vision, function calling, prompt caching |
| Bedrock | `eu.anthropic.claude-sonnet-4-5-20250929-v1:0` | 200K | $3.00 | $15.00 | Chat, reasoning, vision, function calling, prompt caching |
| Azure AI | `azure_ai/grok-4` | 131K | $5.50 | $27.50 | Chat, reasoning, function calling, web search |
| Azure AI | `azure_ai/grok-4-fast-reasoning` | 131K | $5.80 | $2,900.00 | Chat, reasoning, function calling, web search |
| Azure AI | `azure_ai/grok-4-fast-non-reasoning` | 131K | $5.00 | $2,500.00 | Chat, function calling, web search |
| Azure AI | `azure_ai/grok-code-fast-1` | 131K | $3.50 | $17.50 | Chat, function calling, web search |
| Groq | `groq/moonshotai/kimi-k2-instruct-0905` | Context varies | Pricing varies | Pricing varies | Chat, function calling |
| Ollama | Ollama Cloud models | Varies | Free | Free | Self-hosted models via Ollama Cloud |
#### Features
- **[Anthropic](../../docs/providers/anthropic)**
- Add new claude-sonnet-4-5 model family with tiered pricing above 200K tokens - [PR #15041](https://github.com/BerriAI/litellm/pull/15041)
- Add anthropic/claude-sonnet-4-5 to model price json with prompt caching support - [PR #15049](https://github.com/BerriAI/litellm/pull/15049)
- Add 200K prices for Sonnet 4.5 - [PR #15140](https://github.com/BerriAI/litellm/pull/15140)
- Add cost tracking for /v1/messages in streaming response - [PR #15102](https://github.com/BerriAI/litellm/pull/15102)
- Add /v1/messages/count_tokens to Anthropic routes for non-admin user access - [PR #15034](https://github.com/BerriAI/litellm/pull/15034)
- **[Gemini](../../docs/providers/gemini)**
- Ignore type param for gemini tools - [PR #15022](https://github.com/BerriAI/litellm/pull/15022)
- **[Vertex AI](../../docs/providers/vertex)**
- Add LiteLLM Overhead metric for VertexAI - [PR #15040](https://github.com/BerriAI/litellm/pull/15040)
- Support googlemap grounding in vertex ai - [PR #15179](https://github.com/BerriAI/litellm/pull/15179)
- **[Azure](../../docs/providers/azure)**
- Add azure_ai grok-4 model family - [PR #15137](https://github.com/BerriAI/litellm/pull/15137)
- Use the `extra_query` parameter for GET requests in Azure Batch - [PR #14997](https://github.com/BerriAI/litellm/pull/14997)
- Use extra_query for download results (Batch API) - [PR #15025](https://github.com/BerriAI/litellm/pull/15025)
- Add support for Azure AD token-based authorization - [PR #14813](https://github.com/BerriAI/litellm/pull/14813)
- **[Ollama](../../docs/providers/ollama)**
- Add ollama cloud models - [PR #15008](https://github.com/BerriAI/litellm/pull/15008)
- **[Groq](../../docs/providers/groq)**
- Add groq/moonshotai/kimi-k2-instruct-0905 - [PR #15079](https://github.com/BerriAI/litellm/pull/15079)
- **[OpenAI](../../docs/providers/openai)**
- Add support for GPT 5 codex models - [PR #14841](https://github.com/BerriAI/litellm/pull/14841)
- **[DeepInfra](../../docs/providers/deepinfra)**
- Update DeepInfra model data refresh with latest pricing - [PR #14939](https://github.com/BerriAI/litellm/pull/14939)
- **[Bedrock](../../docs/providers/bedrock)**
- Add JP Cross-Region Inference - [PR #15188](https://github.com/BerriAI/litellm/pull/15188)
- Add "eu.anthropic.claude-sonnet-4-5-20250929-v1:0" - [PR #15181](https://github.com/BerriAI/litellm/pull/15181)
- Add twelvelabs bedrock Async Invoke Support - [PR #14871](https://github.com/BerriAI/litellm/pull/14871)
- **[Nvidia NIM](../../docs/providers/nvidia_nim)**
- Add Nvidia NIM Rerank Support - [PR #15152](https://github.com/BerriAI/litellm/pull/15152)
### Bug Fixes
- **[VLLM](../../docs/providers/vllm)**
- Fix response_format bug in hosted vllm audio_transcription - [PR #15010](https://github.com/BerriAI/litellm/pull/15010)
- Fix passthrough of atranscription into kwargs going to upstream provider - [PR #15005](https://github.com/BerriAI/litellm/pull/15005)
- **[OCI](../../docs/providers/oci)**
- Fix OCI Generative AI Integration when using Proxy - [PR #15072](https://github.com/BerriAI/litellm/pull/15072)
- **General**
- Fix: Authorization header to use correct "Bearer" capitalization - [PR #14764](https://github.com/BerriAI/litellm/pull/14764)
- Bug fix: gpt-5-chat-latest has incorrect max_input_tokens value - [PR #15116](https://github.com/BerriAI/litellm/pull/15116)
- Update request handling for original exceptions - [PR #15013](https://github.com/BerriAI/litellm/pull/15013)
#### New Provider Support
- **[AMD Lemonade](../../docs/providers/lemonade)**
- Add AMD Lemonade provider support - [PR #14840](https://github.com/BerriAI/litellm/pull/14840)
---
## LLM API Endpoints
#### Features
- **[Responses API](../../docs/response_api)**
- Return Cost for Responses API Streaming requests - [PR #15053](https://github.com/BerriAI/litellm/pull/15053)
- **[/generateContent](../../docs/providers/gemini)**
- Add full support for native Gemini API translation - [PR #15029](https://github.com/BerriAI/litellm/pull/15029)
- **Passthrough Gemini Routes**
- Add Gemini generateContent passthrough cost tracking - [PR #15014](https://github.com/BerriAI/litellm/pull/15014)
- Add streamGenerateContent cost tracking in passthrough - [PR #15199](https://github.com/BerriAI/litellm/pull/15199)
- **Passthrough Vertex AI Routes**
- Add cost tracking for Vertex AI Passthrough `/predict` endpoint - [PR #15019](https://github.com/BerriAI/litellm/pull/15019)
- Add cost tracking for Vertex AI Live API WebSocket Passthrough - [PR #14956](https://github.com/BerriAI/litellm/pull/14956)
- **General**
- Preserve Whitespace Characters in Model Response Streams - [PR #15160](https://github.com/BerriAI/litellm/pull/15160)
- Add provider name to payload specification - [PR #15130](https://github.com/BerriAI/litellm/pull/15130)
- Ensure query params are forwarded from origin url to downstream request - [PR #15087](https://github.com/BerriAI/litellm/pull/15087)
---
## Management Endpoints / UI
#### Features
- **Virtual Keys**
- Fix Session Token Cookie Infinite Logout Loop - [PR #15146](https://github.com/BerriAI/litellm/pull/15146)
- Ensure LLM_API_KEYs can access pass through routes - [PR #15115](https://github.com/BerriAI/litellm/pull/15115)
- **Models + Endpoints**
- Ensure OCI secret fields not shared on /models and /v1/models endpoints - [PR #15085](https://github.com/BerriAI/litellm/pull/15085)
- Add snowflake on UI - [PR #15083](https://github.com/BerriAI/litellm/pull/15083)
- Make UI theme settings publicly accessible for custom branding - [PR #15074](https://github.com/BerriAI/litellm/pull/15074)
- **Admin Settings**
- Ensure OTEL settings are saved in DB after set on UI - [PR #15118](https://github.com/BerriAI/litellm/pull/15118)
- Top api key tags - [PR #15151](https://github.com/BerriAI/litellm/pull/15151), [PR #15156](https://github.com/BerriAI/litellm/pull/15156)
#### Bugs
- **Dashboard** - Fix LiteLLM model name fallback in dashboard overview - [PR #14998](https://github.com/BerriAI/litellm/pull/14998)
---
## Logging / Guardrail / Prompt Management Integrations
#### Features
- **[OpenTelemetry](../../docs/observability/otel)**
- Use generation_name for span naming in logging method - [PR #14799](https://github.com/BerriAI/litellm/pull/14799)
- **[Langfuse](../../docs/proxy/logging#langfuse)**
- Handle non-serializable objects in Langfuse logging - [PR #15148](https://github.com/BerriAI/litellm/pull/15148)
- Set usage_details.total in langfuse integration - [PR #15015](https://github.com/BerriAI/litellm/pull/15015)
#### Guardrails
- **[Javelin](../../docs/proxy/guardrails)**
- Add Javelin standalone guardrails integration for LiteLLM Proxy - [PR #14983](https://github.com/BerriAI/litellm/pull/14983)
- Add logging for important status fields in guardrails - [PR #15090](https://github.com/BerriAI/litellm/pull/15090)
- Don't run post_call guardrail if no text returned from Bedrock - [PR #15106](https://github.com/BerriAI/litellm/pull/15106)
#### Prompt Management
- **[GitLab](../../docs/proxy/prompt_management)**
- GitLab based Prompt manager - [PR #14988](https://github.com/BerriAI/litellm/pull/14988)
---
## Spend Tracking, Budgets and Rate Limiting
- **Cost Tracking**
- Proxy: end user cost tracking in the responses API - [PR #15124](https://github.com/BerriAI/litellm/pull/15124)
- **Parallel Request Limiter v3**
- Use well known redis cluster hashing algorithm - [PR #15052](https://github.com/BerriAI/litellm/pull/15052)
- Fixes to dynamic rate limiter v3 - add saturation detection - [PR #15119](https://github.com/BerriAI/litellm/pull/15119)
- Dynamic Rate Limiter v3 - fixes for detecting saturation + fixes for post saturation behavior - [PR #15192](https://github.com/BerriAI/litellm/pull/15192)
- **Teams**
- Add model specific tpm/rpm limits to teams on LiteLLM - [PR #15044](https://github.com/BerriAI/litellm/pull/15044)
---
## MCP Gateway
- **Server Configuration**
- Specify forwardable headers, specify allowed/disallowed tools for MCP servers - [PR #15002](https://github.com/BerriAI/litellm/pull/15002)
- Enforce server permissions on call tools - [PR #15044](https://github.com/BerriAI/litellm/pull/15044)
- MCP Gateway Fine-grained Tools Addition - [PR #15153](https://github.com/BerriAI/litellm/pull/15153)
- **Bug Fixes**
- Remove servername prefix mcp tools tests - [PR #14986](https://github.com/BerriAI/litellm/pull/14986)
- Resolve regression with duplicate Mcp-Protocol-Version header - [PR #15050](https://github.com/BerriAI/litellm/pull/15050)
- Fix test_mcp_server.py - [PR #15183](https://github.com/BerriAI/litellm/pull/15183)
---
## Performance / Loadbalancing / Reliability improvements
- **Router Optimizations**
- **+62.5% P99 Latency Improvement** - Remove router inefficiencies (from O(M*N) to O(1)) - [PR #15046](https://github.com/BerriAI/litellm/pull/15046)
- Remove hasattr checks in Router - [PR #15082](https://github.com/BerriAI/litellm/pull/15082)
- Remove Double Lookups - [PR #15084](https://github.com/BerriAI/litellm/pull/15084)
- Optimize _filter_cooldown_deployments from O(n×m + k×n) to O(n) - [PR #15091](https://github.com/BerriAI/litellm/pull/15091)
- Optimize unhealthy deployment filtering in retry path (O(n*m) → O(n+m)) - [PR #15110](https://github.com/BerriAI/litellm/pull/15110)
- **Cache Optimizations**
- Reduce complexity of InMemoryCache.evict_cache from O(n*log(n)) to O(log(n)) - [PR #15000](https://github.com/BerriAI/litellm/pull/15000)
- Avoiding expensive operations when cache isn't available - [PR #15182](https://github.com/BerriAI/litellm/pull/15182)
- **Worker Management**
- Add proxy CLI option to recycle workers after N requests - [PR #15007](https://github.com/BerriAI/litellm/pull/15007)
- **Metrics & Monitoring**
- LiteLLM Overhead metric tracking - Add support for tracking litellm overhead on cache hits - [PR #15045](https://github.com/BerriAI/litellm/pull/15045)
---
## Documentation Updates
- **Provider Documentation**
- Update litellm docs from latest release - [PR #15004](https://github.com/BerriAI/litellm/pull/15004)
- Add missing api_key parameter - [PR #15058](https://github.com/BerriAI/litellm/pull/15058)
- **General Documentation**
- Use docker compose instead of docker-compose - [PR #15024](https://github.com/BerriAI/litellm/pull/15024)
- Add railtracks to projects that are using litellm - [PR #15144](https://github.com/BerriAI/litellm/pull/15144)
- Perf: Last week improvement - [PR #15193](https://github.com/BerriAI/litellm/pull/15193)
- Sync models GitHub documentation with Loom video and cross-reference - [PR #15191](https://github.com/BerriAI/litellm/pull/15191)
---
## Security Fixes
- **JWT Token Security** - Don't log JWT SSO token on .info() log - [PR #15145](https://github.com/BerriAI/litellm/pull/15145)
---
## New Contributors
* @herve-ves made their first contribution in [PR #14998](https://github.com/BerriAI/litellm/pull/14998)
* @wenxi-onyx made their first contribution in [PR #15008](https://github.com/BerriAI/litellm/pull/15008)
* @jpetrucciani made their first contribution in [PR #15005](https://github.com/BerriAI/litellm/pull/15005)
* @abhijitjavelin made their first contribution in [PR #14983](https://github.com/BerriAI/litellm/pull/14983)
* @ZeroClover made their first contribution in [PR #15039](https://github.com/BerriAI/litellm/pull/15039)
* @cedarm made their first contribution in [PR #15043](https://github.com/BerriAI/litellm/pull/15043)
* @Isydmr made their first contribution in [PR #15025](https://github.com/BerriAI/litellm/pull/15025)
* @serializer made their first contribution in [PR #15013](https://github.com/BerriAI/litellm/pull/15013)
* @eddierichter-amd made their first contribution in [PR #14840](https://github.com/BerriAI/litellm/pull/14840)
* @malags made their first contribution in [PR #15000](https://github.com/BerriAI/litellm/pull/15000)
* @henryhwang made their first contribution in [PR #15029](https://github.com/BerriAI/litellm/pull/15029)
* @plafleur made their first contribution in [PR #15111](https://github.com/BerriAI/litellm/pull/15111)
* @tyler-liner made their first contribution in [PR #14799](https://github.com/BerriAI/litellm/pull/14799)
* @Amir-R25 made their first contribution in [PR #15144](https://github.com/BerriAI/litellm/pull/15144)
* @georg-wolflein made their first contribution in [PR #15124](https://github.com/BerriAI/litellm/pull/15124)
* @niharm made their first contribution in [PR #15140](https://github.com/BerriAI/litellm/pull/15140)
* @anthony-liner made their first contribution in [PR #15015](https://github.com/BerriAI/litellm/pull/15015)
* @rishiganesh2002 made their first contribution in [PR #15153](https://github.com/BerriAI/litellm/pull/15153)
* @danielaskdd made their first contribution in [PR #15160](https://github.com/BerriAI/litellm/pull/15160)
* @JVenberg made their first contribution in [PR #15146](https://github.com/BerriAI/litellm/pull/15146)
* @speglich made their first contribution in [PR #15072](https://github.com/BerriAI/litellm/pull/15072)
* @daily-kim made their first contribution in [PR #14764](https://github.com/BerriAI/litellm/pull/14764)
---
## **[Full Changelog](https://github.com/BerriAI/litellm/compare/v1.77.5.rc.4...v1.77.7.rc.1)**

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

@ -458,7 +458,14 @@ const sidebars = {
"providers/deepgram",
"providers/watsonx",
"providers/predibase",
"providers/nvidia_nim",
{
type: "category",
label: "Nvidia NIM",
items: [
"providers/nvidia_nim",
"providers/nvidia_nim_rerank",
]
},
{ type: "doc", id: "providers/nscale", label: "Nscale (EU Sovereign)" },
"providers/xai",
"providers/moonshot",
@ -699,7 +706,8 @@ const sidebars = {
"projects/llm_cord",
"projects/pgai",
"projects/GPTLocalhost",
"projects/HolmesGPT"
"projects/HolmesGPT",
"projects/Railtracks",
],
},
"extras/code_quality",

BIN
litellm-proxy-extras/dist/litellm_proxy_extras-0.2.23-py3-none-any.whl vendored Normal file
View file

Binary file not shown.

BIN
litellm-proxy-extras/dist/litellm_proxy_extras-0.2.23.tar.gz vendored Normal file
View file

Binary file not shown.

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20251004123655_mcp_allowed_tools_column/migration.sql Normal file
View file

@ -0,0 +1,3 @@
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "allowed_tools" TEXT[] DEFAULT ARRAY[]::TEXT[];

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

@ -178,6 +178,7 @@ model LiteLLM_MCPServerTable {
updated_by String?
mcp_info Json? @default("{}")
mcp_access_groups String[]
allowed_tools String[] @default([])
// Health check status
status String? @default("unknown")
last_health_check DateTime?

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm-proxy-extras"
version = "0.2.22"
version = "0.2.23"
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.22"
version = "0.2.23"
version_files = [
"pyproject.toml:version",
"../requirements.txt:litellm-proxy-extras==",

133
litellm/__init__.py
View file

@ -152,6 +152,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"vector_store_pre_call_hook",
"dotprompt",
"bitbucket",
"gitlab",
"cloudzero",
"posthog",
]
@ -172,22 +173,22 @@ prometheus_initialize_budget_metrics: Optional[bool] = False
require_auth_for_metrics_endpoint: Optional[bool] = False
argilla_batch_size: Optional[int] = None
datadog_use_v1: Optional[bool] = False # if you want to use v1 datadog logged payload.
gcs_pub_sub_use_v1: Optional[bool] = (
False # if you want to use v1 gcs pubsub logged payload
)
generic_api_use_v1: Optional[bool] = (
False # if you want to use v1 generic api logged payload
)
gcs_pub_sub_use_v1: Optional[
bool
] = False # if you want to use v1 gcs pubsub logged payload
generic_api_use_v1: Optional[
bool
] = False # if you want to use v1 generic api logged payload
argilla_transformation_object: Optional[Dict[str, Any]] = None
_async_input_callback: List[Union[str, Callable, CustomLogger]] = (
[]
) # internal variable - async custom callbacks are routed here.
_async_success_callback: List[Union[str, Callable, CustomLogger]] = (
[]
) # internal variable - async custom callbacks are routed here.
_async_failure_callback: List[Union[str, Callable, CustomLogger]] = (
[]
) # internal variable - async custom callbacks are routed here.
_async_input_callback: List[
Union[str, Callable, CustomLogger]
] = [] # internal variable - async custom callbacks are routed here.
_async_success_callback: List[
Union[str, Callable, CustomLogger]
] = [] # internal variable - async custom callbacks are routed here.
_async_failure_callback: List[
Union[str, Callable, CustomLogger]
] = [] # internal variable - async custom callbacks are routed here.
pre_call_rules: List[Callable] = []
post_call_rules: List[Callable] = []
turn_off_message_logging: Optional[bool] = False
@ -195,18 +196,18 @@ log_raw_request_response: bool = False
redact_messages_in_exceptions: Optional[bool] = False
redact_user_api_key_info: Optional[bool] = False
filter_invalid_headers: Optional[bool] = False
add_user_information_to_llm_headers: Optional[bool] = (
None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers
)
add_user_information_to_llm_headers: Optional[
bool
] = None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers
store_audit_logs = False # Enterprise feature, allow users to see audit logs
### end of callbacks #############
email: Optional[str] = (
None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
)
token: Optional[str] = (
None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
)
email: Optional[
str
] = None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
token: Optional[
str
] = None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
telemetry = True
max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults
drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False))
@ -289,7 +290,7 @@ banned_keywords_list: Optional[Union[str, List]] = None
llm_guard_mode: Literal["all", "key-specific", "request-specific"] = "all"
guardrail_name_config_map: Dict[str, GuardrailItem] = {}
include_cost_in_streaming_usage: bool = False
### PROMPTS ###
### PROMPTS ####
from litellm.types.prompts.init_prompts import PromptSpec
prompt_name_config_map: Dict[str, PromptSpec] = {}
@ -307,24 +308,20 @@ enable_loadbalancing_on_batch_endpoints: Optional[bool] = None
enable_caching_on_provider_specific_optional_params: bool = (
False # feature-flag for caching on optional params - e.g. 'top_k'
)
caching: bool = (
False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
)
caching_with_models: bool = (
False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
)
cache: Optional[Cache] = (
None # cache object <- use this - https://docs.litellm.ai/docs/caching
)
caching: bool = False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
caching_with_models: bool = False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
cache: Optional[
Cache
] = None # cache object <- use this - https://docs.litellm.ai/docs/caching
default_in_memory_ttl: Optional[float] = None
default_redis_ttl: Optional[float] = None
default_redis_batch_cache_expiry: Optional[float] = None
model_alias_map: Dict[str, str] = {}
model_group_settings: Optional["ModelGroupSettings"] = None
max_budget: float = 0.0 # set the max budget across all providers
budget_duration: Optional[str] = (
None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
)
budget_duration: Optional[
str
] = None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
default_soft_budget: float = (
DEFAULT_SOFT_BUDGET # by default all litellm proxy keys have a soft budget of 50.0
)
@ -333,15 +330,11 @@ forward_traceparent_to_llm_provider: bool = False
_current_cost = 0.0 # private variable, used if max budget is set
error_logs: Dict = {}
add_function_to_prompt: bool = (
False # if function calling not supported by api, append function call details to system prompt
)
add_function_to_prompt: bool = False # if function calling not supported by api, append function call details to system prompt
client_session: Optional[httpx.Client] = None
aclient_session: Optional[httpx.AsyncClient] = None
model_fallbacks: Optional[List] = None # Deprecated for 'litellm.fallbacks'
model_cost_map_url: str = (
"https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
)
model_cost_map_url: str = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
suppress_debug_info = False
dynamodb_table_name: Optional[str] = None
s3_callback_params: Optional[Dict] = None
@ -371,12 +364,10 @@ prometheus_metrics_config: Optional[List] = None
disable_add_prefix_to_prompt: bool = (
False # used by anthropic, to disable adding prefix to prompt
)
disable_copilot_system_to_assistant: bool = (
False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior.
)
disable_copilot_system_to_assistant: bool = False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior.
public_model_groups: Optional[List[str]] = None
public_model_groups_links: Dict[str, str] = {}
#### REQUEST PRIORITIZATION ######
#### REQUEST PRIORITIZATION #######
priority_reservation: Optional[Dict[str, float]] = None
priority_reservation_settings: "PriorityReservationSettings" = (
PriorityReservationSettings()
@ -384,17 +375,13 @@ priority_reservation_settings: "PriorityReservationSettings" = (
######## Networking Settings ########
use_aiohttp_transport: bool = (
True # Older variable, aiohttp is now the default. use disable_aiohttp_transport instead.
)
use_aiohttp_transport: bool = True # Older variable, aiohttp is now the default. use disable_aiohttp_transport instead.
aiohttp_trust_env: bool = False # set to true to use HTTP_ Proxy settings
disable_aiohttp_transport: bool = False # Set this to true to use httpx instead
disable_aiohttp_trust_env: bool = (
False # When False, aiohttp will respect HTTP(S)_PROXY env vars
)
force_ipv4: bool = (
False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6.
)
force_ipv4: bool = False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6.
module_level_aclient = AsyncHTTPHandler(
timeout=request_timeout, client_alias="module level aclient"
)
@ -408,13 +395,13 @@ fallbacks: Optional[List] = None
context_window_fallbacks: Optional[List] = None
content_policy_fallbacks: Optional[List] = None
allowed_fails: int = 3
num_retries_per_request: Optional[int] = (
None # for the request overall (incl. fallbacks + model retries)
)
num_retries_per_request: Optional[
int
] = None # for the request overall (incl. fallbacks + model retries)
####### SECRET MANAGERS #####################
secret_manager_client: Optional[Any] = (
None # list of instantiated key management clients - e.g. azure kv, infisical, etc.
)
secret_manager_client: Optional[
Any
] = None # list of instantiated key management clients - e.g. azure kv, infisical, etc.
_google_kms_resource_name: Optional[str] = None
_key_management_system: Optional[KeyManagementSystem] = None
_key_management_settings: KeyManagementSettings = KeyManagementSettings()
@ -510,6 +497,7 @@ azure_text_models: Set = set()
anyscale_models: Set = set()
cerebras_models: Set = set()
galadriel_models: Set = set()
nvidia_nim_models: Set = set()
sambanova_models: Set = set()
sambanova_embedding_models: Set = set()
novita_models: Set = set()
@ -704,6 +692,8 @@ def add_known_models():
cerebras_models.add(key)
elif value.get("litellm_provider") == "galadriel":
galadriel_models.add(key)
elif value.get("litellm_provider") == "nvidia_nim":
nvidia_nim_models.add(key)
elif value.get("litellm_provider") == "sambanova":
sambanova_models.add(key)
elif value.get("litellm_provider") == "sambanova-embedding-models":
@ -831,6 +821,7 @@ model_list = list(
| anyscale_models
| cerebras_models
| galadriel_models
| nvidia_nim_models
| sambanova_models
| azure_text_models
| novita_models
@ -914,6 +905,7 @@ models_by_provider: dict = {
"anyscale": anyscale_models,
"cerebras": cerebras_models,
"galadriel": galadriel_models,
"nvidia_nim": nvidia_nim_models,
"sambanova": sambanova_models | sambanova_embedding_models,
"novita": novita_models,
"nebius": nebius_models | nebius_embedding_models,
@ -1074,6 +1066,7 @@ from .llms.azure_ai.rerank.transformation import AzureAIRerankConfig
from .llms.infinity.rerank.transformation import InfinityRerankConfig
from .llms.jina_ai.rerank.transformation import JinaAIRerankConfig
from .llms.deepinfra.rerank.transformation import DeepinfraRerankConfig
from .llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig
from .llms.clarifai.chat.transformation import ClarifaiConfig
from .llms.ai21.chat.transformation import AI21ChatConfig, AI21ChatConfig as AI21Config
from .llms.meta_llama.chat.transformation import LlamaAPIConfig
@ -1174,6 +1167,7 @@ from .llms.bedrock.embed.amazon_titan_v2_transformation import (
)
from .llms.cohere.chat.transformation import CohereChatConfig
from .llms.bedrock.embed.cohere_transformation import BedrockCohereEmbeddingConfig
from .llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig
from .llms.openai.openai import OpenAIConfig, MistralEmbeddingConfig
from .llms.openai.image_variations.transformation import OpenAIImageVariationConfig
from .llms.deepinfra.chat.transformation import DeepInfraConfig
@ -1349,12 +1343,12 @@ from .types.llms.custom_llm import CustomLLMItem
from .types.utils import GenericStreamingChunk
custom_provider_map: List[CustomLLMItem] = []
_custom_providers: List[str] = (
[]
) # internal helper util, used to track names of custom providers
disable_hf_tokenizer_download: Optional[bool] = (
None # disable huggingface tokenizer download. Defaults to openai clk100
)
_custom_providers: List[
str
] = [] # internal helper util, used to track names of custom providers
disable_hf_tokenizer_download: Optional[
bool
] = None # disable huggingface tokenizer download. Defaults to openai clk100
global_disable_no_log_param: bool = False
### CLI UTILITIES ###
@ -1362,6 +1356,7 @@ from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_k
### PASSTHROUGH ###
from .passthrough import allm_passthrough_route, llm_passthrough_route
from .google_genai import agenerate_content
### GLOBAL CONFIG ###
global_bitbucket_config: Optional[Dict[str, Any]] = None
@ -1371,3 +1366,11 @@ def set_global_bitbucket_config(config: Dict[str, Any]) -> None:
"""Set global BitBucket configuration for prompt management."""
global global_bitbucket_config
global_bitbucket_config = config
### GLOBAL CONFIG ###
global_gitlab_config: Optional[Dict[str, Any]] = None
def set_global_gitlab_config(config: Dict[str, Any]) -> None:
"""Set global BitBucket configuration for prompt management."""
global global_gitlab_config
global_gitlab_config = config

144
litellm/batches/main.py
View file

@ -59,18 +59,22 @@ def _resolve_timeout(
) -> float:
"""
Resolve timeout value from various sources and handle httpx.Timeout objects.
Args:
optional_params: GenericLiteLLMParams object containing timeout
kwargs: Additional kwargs that may contain request_timeout
custom_llm_provider: Provider name for httpx timeout support check
default_timeout: Default timeout value to use
Returns:
Resolved timeout as float
"""
timeout = optional_params.timeout or kwargs.get("request_timeout", default_timeout) or default_timeout
timeout = (
optional_params.timeout
or kwargs.get("request_timeout", default_timeout)
or default_timeout
)
# Handle httpx.Timeout objects
if isinstance(timeout, httpx.Timeout):
if supports_httpx_timeout(custom_llm_provider) is False:
@ -81,11 +85,11 @@ def _resolve_timeout(
# For providers that support httpx.Timeout, we still need to return a float
# This case might need to be handled differently based on the actual use case
return float(timeout.read or default_timeout)
# Handle None case
if timeout is None:
return float(default_timeout)
# Handle numeric values (int, float, string representations)
return float(timeout)
@ -163,15 +167,19 @@ def create_batch(
try:
if model is not None:
model, _, _, _ = get_llm_provider(
model=model,
custom_llm_provider=None,
)
model=model,
custom_llm_provider=None,
)
except Exception as e:
verbose_logger.exception(f"litellm.batches.main.py::create_batch() - Error inferring custom_llm_provider - {str(e)}")
verbose_logger.exception(
f"litellm.batches.main.py::create_batch() - Error inferring custom_llm_provider - {str(e)}"
)
_is_async = kwargs.pop("acreate_batch", False) is True
litellm_params = dict(GenericLiteLLMParams(**kwargs))
litellm_logging_obj: LiteLLMLoggingObj = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None))
litellm_logging_obj: LiteLLMLoggingObj = cast(
LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None)
)
### TIMEOUT LOGIC ###
timeout = _resolve_timeout(optional_params, kwargs, custom_llm_provider)
litellm_logging_obj.update_environment_variables(
@ -189,7 +197,6 @@ def create_batch(
},
custom_llm_provider=custom_llm_provider,
)
_create_batch_request = CreateBatchRequest(
completion_window=completion_window,
@ -378,6 +385,7 @@ async def aretrieve_batch(
except Exception as e:
raise e
def _handle_retrieve_batch_providers_without_provider_config(
batch_id: str,
optional_params: GenericLiteLLMParams,
@ -497,6 +505,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
)
return response
@client
def retrieve_batch(
batch_id: str,
@ -513,7 +522,9 @@ def retrieve_batch(
"""
try:
optional_params = GenericLiteLLMParams(**kwargs)
litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None)
litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get(
"litellm_logging_obj", None
)
### TIMEOUT LOGIC ###
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
litellm_params = get_litellm_params(
@ -549,7 +560,26 @@ def retrieve_batch(
_is_async = kwargs.pop("aretrieve_batch", False) is True
client = kwargs.get("client", None)
# Check if this is an async invoke ARN (different from regular batch ARN)
# Async invoke ARNs have format: arn:aws(-[^:]+)?:bedrock:[a-z0-9-]{1,20}:[0-9]{12}:async-invoke/[a-z0-9]{12}
if (
batch_id.startswith("arn:aws")
and ":bedrock:" in batch_id
and ":async-invoke/" in batch_id
):
# Handle async invoke status check
# Remove aws_region_name from kwargs to avoid duplicate parameter
async_kwargs = kwargs.copy()
async_kwargs.pop("aws_region_name", None)
return _handle_async_invoke_status(
batch_id=batch_id,
aws_region_name=kwargs.get("aws_region_name", "us-east-1"),
logging_obj=litellm_logging_obj,
**async_kwargs,
)
# Try to use provider config first (for providers like bedrock)
model: Optional[str] = kwargs.get("model", None)
if model is not None:
@ -559,7 +589,7 @@ def retrieve_batch(
)
else:
provider_config = None
if provider_config is not None:
response = base_llm_http_handler.retrieve_batch(
batch_id=batch_id,
@ -568,7 +598,8 @@ def retrieve_batch(
headers=extra_headers or {},
api_base=optional_params.api_base,
api_key=optional_params.api_key,
logging_obj=litellm_logging_obj or LiteLLMLoggingObj(
logging_obj=litellm_logging_obj
or LiteLLMLoggingObj(
model=model or "bedrock/unknown",
messages=[],
stream=False,
@ -586,7 +617,6 @@ def retrieve_batch(
model=model,
)
return response
#########################################################
# Handle providers without provider config
@ -600,7 +630,7 @@ def retrieve_batch(
_is_async=_is_async,
timeout=timeout,
)
except Exception as e:
raise e
@ -933,3 +963,79 @@ def cancel_batch(
return response
except Exception as e:
raise e
def _handle_async_invoke_status(
batch_id: str, aws_region_name: str, logging_obj=None, **kwargs
) -> "LiteLLMBatch":
"""
Handle async invoke status check for AWS Bedrock.
Args:
batch_id: The async invoke ARN
aws_region_name: AWS region name
**kwargs: Additional parameters
Returns:
dict: Status information including status, output_file_id (S3 URL), etc.
"""
import asyncio
from litellm.llms.bedrock.embed.embedding import BedrockEmbedding
async def _async_get_status():
# Create embedding handler instance
embedding_handler = BedrockEmbedding()
# Get the status of the async invoke job
status_response = await embedding_handler._get_async_invoke_status(
invocation_arn=batch_id,
aws_region_name=aws_region_name,
logging_obj=logging_obj,
**kwargs,
)
# Transform response to a LiteLLMBatch object
from litellm.types.utils import LiteLLMBatch
result = LiteLLMBatch(
id=status_response["invocationArn"],
object="batch",
status=status_response["status"],
created_at=status_response["submitTime"],
in_progress_at=status_response["lastModifiedTime"],
completed_at=status_response.get("endTime"),
failed_at=status_response.get("endTime")
if status_response["status"] == "failed"
else None,
request_counts={
"total": 1,
"completed": 1 if status_response["status"] == "completed" else 0,
"failed": 1 if status_response["status"] == "failed" else 0,
},
metadata={
"output_file_id": status_response["outputDataConfig"][
"s3OutputDataConfig"
]["s3Uri"],
"failure_message": status_response.get("failureMessage"),
"model_arn": status_response["modelArn"],
},
)
return result
# Since this function is called from within an async context via run_in_executor,
# we need to create a new event loop in a thread to avoid conflicts
import concurrent.futures
def run_in_thread():
new_loop = asyncio.new_event_loop()
asyncio.set_event_loop(new_loop)
try:
return new_loop.run_until_complete(_async_get_status())
finally:
new_loop.close()
with concurrent.futures.ThreadPoolExecutor() as executor:
future = executor.submit(run_in_thread)
return future.result()

85
litellm/caching/caching_handler.py
View file

@ -14,6 +14,7 @@ It utilizes the (RedisCache, s3Cache, RedisSemanticCache, QdrantSemanticCache, I
In each method it will call the appropriate method from caching.py
"""
import time
import asyncio
import datetime
import inspect
@ -57,10 +58,16 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.utils import CustomStreamWrapper
else:
LiteLLMLoggingObj = Any
CustomStreamWrapper = Any
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.litellm_core_utils.core_helpers import (
_get_parent_otel_span_from_kwargs,
)
class CachingHandlerResponse(BaseModel):
@ -112,7 +119,7 @@ class LLMCachingHandler:
call_type: str,
kwargs: Dict[str, Any],
args: Optional[Tuple[Any, ...]] = None,
) -> CachingHandlerResponse:
) -> Optional[CachingHandlerResponse]:
"""
Internal method to get from the cache.
Handles different call types (embeddings, chat/completions, text_completion, transcription)
@ -133,32 +140,27 @@ class LLMCachingHandler:
Raises:
None
"""
from litellm.litellm_core_utils.core_helpers import (
_get_parent_otel_span_from_kwargs,
)
from litellm.utils import CustomStreamWrapper
kwargs = kwargs.copy()
args = args or ()
#########################################################
# Init cache timing metrics
#########################################################
cache_check_start_time = datetime.datetime.now()
cache_check_end_time = None
#########################################################
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
kwargs["parent_otel_span"] = parent_otel_span
final_embedding_cached_response: Optional[EmbeddingResponse] = None
embedding_all_elements_cache_hit: bool = False
cached_result: Optional[Any] = None
# Check if caching should be performed BEFORE doing expensive operations
if (
(kwargs.get("caching", None) is None and litellm.cache is not None)
or kwargs.get("caching", False) is True
) and (
kwargs.get("cache", {}).get("no-cache", False) is not True
): # allow users to control returning cached responses from the completion function
args = args or ()
final_embedding_cached_response: Optional[EmbeddingResponse] = None
embedding_all_elements_cache_hit: bool = False
cached_result: Optional[Any] = None
kwargs = kwargs.copy()
#########################################################
# Init cache timing metrics
#########################################################
cache_check_start_time = time.perf_counter()
cache_check_end_time: Optional[float] = None
#########################################################
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
kwargs["parent_otel_span"] = parent_otel_span
if litellm.cache is not None and self._is_call_type_supported_by_cache(
original_function=original_function
):
@ -168,7 +170,7 @@ class LLMCachingHandler:
kwargs=kwargs,
args=args,
)
cache_check_end_time = datetime.datetime.now()
cache_check_end_time = time.perf_counter()
if cached_result is not None and not isinstance(cached_result, list):
verbose_logger.debug("Cache Hit!")
@ -180,7 +182,7 @@ class LLMCachingHandler:
api_base=kwargs.get("api_base", None),
api_key=kwargs.get("api_key", None),
)
cache_duration_ms = (cache_check_end_time - cache_check_start_time).total_seconds() * 1000
cache_duration_ms = (cache_check_end_time - cache_check_start_time) * 1000
self._update_litellm_logging_obj_environment(
logging_obj=logging_obj,
model=model,
@ -245,11 +247,14 @@ class LLMCachingHandler:
final_embedding_cached_response=final_embedding_cached_response,
embedding_all_elements_cache_hit=embedding_all_elements_cache_hit,
)
verbose_logger.debug(f"CACHE RESULT: {cached_result}")
return CachingHandlerResponse(
cached_result=cached_result,
final_embedding_cached_response=final_embedding_cached_response,
)
verbose_logger.debug(f"CACHE RESULT: {cached_result}")
return CachingHandlerResponse(
cached_result=cached_result,
final_embedding_cached_response=final_embedding_cached_response,
)
# Caching disabled - return None to indicate no caching attempted
return None
def _sync_get_cache(
self,
@ -263,18 +268,22 @@ class LLMCachingHandler:
) -> CachingHandlerResponse:
from litellm.utils import CustomStreamWrapper
args = args or ()
new_kwargs = kwargs.copy()
new_kwargs.update(
convert_args_to_kwargs(
self.original_function,
args,
)
)
cached_result: Optional[Any] = None
# Check if caching should be performed BEFORE doing expensive kwargs copy
if litellm.cache is not None and self._is_call_type_supported_by_cache(
original_function=original_function
):
args = args or ()
# Now that we confirmed caching will happen, prepare kwargs
new_kwargs = kwargs.copy()
new_kwargs.update(
convert_args_to_kwargs(
self.original_function,
args,
)
)
print_verbose("Checking Sync Cache")
cached_result = litellm.cache.get_cache(**new_kwargs)
if cached_result is not None:

39
litellm/caching/in_memory_cache.py
View file

@ -11,6 +11,7 @@ Has 4 methods:
import json
import sys
import time
import heapq
from typing import TYPE_CHECKING, Any, List, Optional
if TYPE_CHECKING:
@ -46,6 +47,7 @@ class InMemoryCache(BaseCache):
# in-memory cache
self.cache_dict: dict = {}
self.ttl_dict: dict = {}
self.expiration_heap: list[tuple[float, str]] = []
def check_value_size(self, value: Any):
"""
@ -114,19 +116,27 @@ class InMemoryCache(BaseCache):
"""
current_time = time.time()
# Step 1: Remove expired items
expired_keys = [key for key, ttl in self.ttl_dict.items() if current_time > ttl]
for key in expired_keys:
self._remove_key(key)
# Step 2: If cache is still full, evict items with earliest expiration times
if len(self.cache_dict) >= self.max_size_in_memory:
# Sort by expiration time (earliest first) and evict until we're under the limit
items_by_expiration = sorted(self.ttl_dict.items(), key=lambda x: x[1])
keys_to_evict = items_by_expiration[:len(self.cache_dict) - self.max_size_in_memory + 1]
for key, _ in keys_to_evict:
# Step 1: Remove expired or outdated items
while self.expiration_heap:
expiration_time, key = self.expiration_heap[0]
# Case 1: Heap entry is outdated
if expiration_time != self.ttl_dict.get(key):
heapq.heappop(self.expiration_heap)
# Case 2: Entry is valid but expired
elif expiration_time <= current_time:
heapq.heappop(self.expiration_heap)
self._remove_key(key)
else:
# Case 3: Entry is valid and not expired
break
# Step 2: Evict if cache is still full
while len(self.cache_dict) >= self.max_size_in_memory:
expiration_time, key = heapq.heappop(self.expiration_heap)
# Skip if key was removed or updated
if self.ttl_dict.get(key) == expiration_time:
self._remove_key(key)
# de-reference the removed item
@ -150,7 +160,7 @@ class InMemoryCache(BaseCache):
# Handle the edge case where max_size_in_memory is 0
if self.max_size_in_memory == 0:
return # Don't cache anything if max size is 0
if len(self.cache_dict) >= self.max_size_in_memory:
# only evict when cache is full
self.evict_cache()
@ -161,8 +171,10 @@ class InMemoryCache(BaseCache):
if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl
if "ttl" in kwargs and kwargs["ttl"] is not None:
self.ttl_dict[key] = time.time() + float(kwargs["ttl"])
heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key))
else:
self.ttl_dict[key] = time.time() + self.default_ttl
heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key))
async def async_set_cache(self, key, value, **kwargs):
self.set_cache(key=key, value=value, **kwargs)
@ -253,6 +265,7 @@ class InMemoryCache(BaseCache):
def flush_cache(self):
self.cache_dict.clear()
self.ttl_dict.clear()
self.expiration_heap.clear()
async def disconnect(self):
pass

46
litellm/google_genai/adapters/handler.py
View file

@ -72,15 +72,24 @@ class GenerateContentToCompletionHandler:
completion_response = await litellm.acompletion(**completion_kwargs)
if stream:
# Transform streaming completion response to generate_content format
transformed_stream = (
GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming(
# Check if completion_response is actually a stream or a ModelResponse
# This can happen in error cases or when stream is not properly supported
if not hasattr(completion_response, "__aiter__"):
# If it's not a stream, treat it as a regular response
generate_content_response = (
GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
cast(ModelResponse, completion_response)
)
)
return generate_content_response
else:
# Transform streaming completion response to generate_content format
transformed_stream = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming(
completion_response
)
)
if transformed_stream is not None:
return transformed_stream
raise ValueError("Failed to transform streaming response")
if transformed_stream is not None:
return transformed_stream
raise ValueError("Failed to transform streaming response")
else:
# Transform completion response back to generate_content format
generate_content_response = (
@ -136,15 +145,24 @@ class GenerateContentToCompletionHandler:
completion_response = litellm.completion(**completion_kwargs)
if stream:
# Transform streaming completion response to generate_content format
transformed_stream = (
GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming(
# Check if completion_response is actually a stream or a ModelResponse
# This can happen in error cases or when stream is not properly supported
if not hasattr(completion_response, "__iter__"):
# If it's not a stream, treat it as a regular response
generate_content_response = (
GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content(
cast(ModelResponse, completion_response)
)
)
return generate_content_response
else:
# Transform streaming completion response to generate_content format
transformed_stream = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming(
completion_response
)
)
if transformed_stream is not None:
return transformed_stream
raise ValueError("Failed to transform streaming response")
if transformed_stream is not None:
return transformed_stream
raise ValueError("Failed to transform streaming response")
else:
# Transform completion response back to generate_content format
generate_content_response = (

300
litellm/google_genai/adapters/transformation.py
View file

@ -1,12 +1,15 @@
import json
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union, cast
from litellm import verbose_logger
from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAssistantMessage,
ChatCompletionAssistantToolCall,
ChatCompletionRequest,
ChatCompletionSystemMessage,
ChatCompletionToolCallFunctionChunk,
ChatCompletionToolChoiceValues,
ChatCompletionToolMessage,
@ -36,43 +39,103 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
def __init__(self, completion_stream: Any):
self.sent_first_chunk = False
self.accumulated_tool_calls = {}
self._returned_response = False
super().__init__(completion_stream)
def __next__(self):
try:
if not hasattr(self.completion_stream, "__iter__"):
if self._returned_response:
raise StopIteration
self._returned_response = True
return GoogleGenAIAdapter().translate_completion_to_generate_content(
self.completion_stream
)
for chunk in self.completion_stream:
if chunk == "None" or chunk is None:
continue
# Transform OpenAI streaming chunk to Google GenAI format
transformed_chunk = GoogleGenAIAdapter().translate_streaming_completion_to_generate_content(
chunk, self
)
if transformed_chunk: # Only return non-empty chunks
if transformed_chunk:
return transformed_chunk
raise StopIteration
except StopIteration:
raise StopIteration
raise
except Exception:
raise StopIteration
async def __anext__(self):
try:
if not hasattr(self.completion_stream, "__aiter__"):
if self._returned_response:
raise StopAsyncIteration
self._returned_response = True
return GoogleGenAIAdapter().translate_completion_to_generate_content(
self.completion_stream
)
async for chunk in self.completion_stream:
if chunk == "None" or chunk is None:
continue
# Transform OpenAI streaming chunk to Google GenAI format
transformed_chunk = GoogleGenAIAdapter().translate_streaming_completion_to_generate_content(
chunk, self
)
if transformed_chunk: # Only return non-empty chunks
if transformed_chunk:
return transformed_chunk
# After the stream is exhausted, check for any remaining accumulated tool calls
if self.accumulated_tool_calls:
try:
parts = []
for (
tool_call_index,
tool_call_data,
) in self.accumulated_tool_calls.items():
try:
# For tool calls with no arguments, accumulated_args will be "", which is not valid JSON.
# We default to an empty JSON object in this case.
parsed_args = json.loads(
tool_call_data["arguments"] or "{}"
)
function_call_part = {
"functionCall": {
"name": tool_call_data["name"]
or "undefined_tool_name",
"args": parsed_args,
}
}
parts.append(function_call_part)
except json.JSONDecodeError:
# This can happen if the stream is abruptly cut off mid-argument string.
verbose_logger.warning(
f"Could not parse tool call arguments at end of stream for index {tool_call_index}. "
f"Name: {tool_call_data['name']}. "
f"Partial args: {tool_call_data['arguments']}"
)
pass
if parts:
final_chunk = {
"candidates": [
{
"content": {"parts": parts, "role": "model"},
"finishReason": "STOP",
"index": 0,
"safetyRatings": [],
}
]
}
return final_chunk
finally:
# Ensure the accumulator is always cleared to prevent memory leaks
self.accumulated_tool_calls.clear()
raise StopAsyncIteration
except StopAsyncIteration:
raise StopAsyncIteration
raise
except Exception:
raise StopAsyncIteration
@ -107,9 +170,14 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
payload = f"data: {json.dumps(transformed_chunk)}\n\n"
yield payload.encode()
else:
raise ValueError(f"Invalid chunk 1: {chunk}")
# For empty chunks, continue to next iteration
continue
else:
raise ValueError(f"Invalid chunk 2: {chunk}")
# For other chunk types, yield them directly
if hasattr(chunk, "encode"):
yield chunk.encode()
else:
yield str(chunk).encode()
class GoogleGenAIAdapter:
@ -133,12 +201,19 @@ class GoogleGenAIAdapter:
model: The model name
contents: Generate content contents (can be list or single dict)
config: Optional config parameters
**kwargs: Additional parameters
**kwargs: Additional parameters from the original request
Returns:
Dict in OpenAI format
"""
# Extract top-level fields from kwargs
system_instruction = kwargs.get("systemInstruction") or kwargs.get(
"system_instruction"
)
tools = kwargs.get("tools")
tool_config = kwargs.get("toolConfig") or kwargs.get("tool_config")
# Normalize contents to list format
if isinstance(contents, dict):
contents_list = [contents]
@ -146,7 +221,9 @@ class GoogleGenAIAdapter:
contents_list = contents
# Transform contents to OpenAI messages format
messages = self._transform_contents_to_messages(contents_list)
messages = self._transform_contents_to_messages(
contents_list, system_instruction=system_instruction
)
# Create base request as dict (which is compatible with ChatCompletionRequest)
completion_request: ChatCompletionRequest = {
@ -182,20 +259,19 @@ class GoogleGenAIAdapter:
completion_request["stop"] = config["stopSequences"]
# Handle tools transformation
if "tools" in kwargs:
tools = kwargs["tools"]
if tools:
# Check if tools are already in OpenAI format or Google GenAI format
if isinstance(tools, list) and len(tools) > 0:
# Tools are in Google GenAI format, transform them
openai_tools = self._transform_google_genai_tools_to_openai(tools)
if openai_tools:
completion_request["tools"] = openai_tools
# Handle tool_config (tool choice)
if "tool_config" in kwargs:
if tool_config:
tool_choice = self._transform_google_genai_tool_config_to_openai(
kwargs["tool_config"]
tool_config
)
if tool_choice:
completion_request["tool_choice"] = tool_choice
@ -235,7 +311,8 @@ class GoogleGenAIAdapter:
return completion_request_dict
def translate_completion_output_params_streaming(
self, completion_stream: Any
self,
completion_stream: Any,
) -> Union[AsyncIterator[bytes], None]:
"""Transform streaming completion output to Google GenAI format"""
google_genai_wrapper = GoogleGenAIStreamWrapper(
@ -245,7 +322,8 @@ class GoogleGenAIAdapter:
return google_genai_wrapper.async_google_genai_sse_wrapper()
def _transform_google_genai_tools_to_openai(
self, tools: List[Dict[str, Any]]
self,
tools: List[Dict[str, Any]],
) -> List[ChatCompletionToolParam]:
"""Transform Google GenAI tools to OpenAI tools format"""
openai_tools: List[Dict[str, Any]] = []
@ -259,8 +337,8 @@ class GoogleGenAIAdapter:
if "description" in func_decl:
function_chunk["description"] = func_decl["description"]
if "parameters" in func_decl:
function_chunk["parameters"] = func_decl["parameters"]
if "parametersJsonSchema" in func_decl:
function_chunk["parameters"] = func_decl["parametersJsonSchema"]
openai_tool = {"type": "function", "function": function_chunk}
openai_tools.append(openai_tool)
@ -271,7 +349,8 @@ class GoogleGenAIAdapter:
return cast(List[ChatCompletionToolParam], normalized_tools)
def _transform_google_genai_tool_config_to_openai(
self, tool_config: Dict[str, Any]
self,
tool_config: Dict[str, Any],
) -> Optional[ChatCompletionToolChoiceValues]:
"""Transform Google GenAI tool_config to OpenAI tool_choice"""
function_calling_config = tool_config.get("functionCallingConfig", {})
@ -283,11 +362,23 @@ class GoogleGenAIAdapter:
return cast(ChatCompletionToolChoiceValues, tool_choice)
def _transform_contents_to_messages(
self, contents: List[Dict[str, Any]]
self,
contents: List[Dict[str, Any]],
system_instruction: Optional[Dict[str, Any]] = None,
) -> List[AllMessageValues]:
"""Transform Google GenAI contents to OpenAI messages format"""
messages: List[AllMessageValues] = []
# Handle system instruction
if system_instruction:
system_parts = system_instruction.get("parts", [])
if system_parts and "text" in system_parts[0]:
messages.append(
ChatCompletionSystemMessage(
role="system", content=system_parts[0]["text"]
)
)
for content in contents:
role = content.get("role", "user")
parts = content.get("parts", [])
@ -364,7 +455,8 @@ class GoogleGenAIAdapter:
return messages
def translate_completion_to_generate_content(
self, response: ModelResponse
self,
response: ModelResponse,
) -> Dict[str, Any]:
"""
Transform litellm completion response to Google GenAI generate_content format
@ -376,6 +468,7 @@ class GoogleGenAIAdapter:
Dict in Google GenAI generate_content response format
"""
# Extract the main response content
choice = response.choices[0] if response.choices else None
if not choice:
@ -388,12 +481,6 @@ class GoogleGenAIAdapter:
"Invalid completion response: no message found in choice"
)
parts = self._transform_openai_message_to_google_genai_parts(choice.message)
elif isinstance(choice, StreamingChoices):
if not choice.delta:
raise ValueError(
"Invalid completion response: no delta found in streaming choice"
)
parts = self._transform_openai_delta_to_google_genai_parts(choice.delta)
else:
# Fallback for generic choice objects
message_content = getattr(choice, "message", {}).get(
@ -438,7 +525,7 @@ class GoogleGenAIAdapter:
self,
response: Union[ModelResponse, ModelResponseStream],
wrapper: GoogleGenAIStreamWrapper,
) -> Dict[str, Any]:
) -> Optional[Dict[str, Any]]:
"""
Transform streaming litellm completion chunk to Google GenAI generate_content format
@ -454,7 +541,7 @@ class GoogleGenAIAdapter:
choice = response.choices[0] if response.choices else None
if not choice:
# Return empty chunk if no choices
return {}
return None
# Handle streaming choice
if isinstance(choice, StreamingChoices):
@ -473,7 +560,7 @@ class GoogleGenAIAdapter:
# Only create response chunk if we have parts or it's the final chunk
if not parts and not finish_reason:
return {}
return None
# Create Google GenAI streaming format response
streaming_chunk: Dict[str, Any] = {
@ -515,7 +602,8 @@ class GoogleGenAIAdapter:
return streaming_chunk
def _transform_openai_message_to_google_genai_parts(
self, message: Any
self,
message: Any,
) -> List[Dict[str, Any]]:
"""Transform OpenAI message to Google GenAI parts format"""
parts: List[Dict[str, Any]] = []
@ -537,112 +625,94 @@ class GoogleGenAIAdapter:
except json.JSONDecodeError:
args = {}
function_call_part = {
"functionCall": {"name": tool_call.function.name, "args": args}
}
parts.append(function_call_part)
return parts if parts else [{"text": ""}]
def _transform_openai_delta_to_google_genai_parts(
self, delta: Any
) -> List[Dict[str, Any]]:
"""Transform OpenAI delta to Google GenAI parts format for streaming"""
parts: List[Dict[str, Any]] = []
# Add text content if present
if hasattr(delta, "content") and delta.content:
parts.append({"text": delta.content})
# Add tool calls if present (for streaming tool calls)
if hasattr(delta, "tool_calls") and delta.tool_calls:
for tool_call in delta.tool_calls:
if hasattr(tool_call, "function") and tool_call.function:
# For streaming, we might get partial function arguments
args_str = getattr(tool_call.function, "arguments", "") or ""
try:
args = json.loads(args_str) if args_str else {}
except json.JSONDecodeError:
# For partial JSON in streaming, return as text for now
args = {"partial": args_str}
function_call_part = {
"functionCall": {
"name": getattr(tool_call.function, "name", "") or "",
"name": tool_call.function.name or "undefined_tool_name",
"args": args,
}
}
parts.append(function_call_part)
return parts
return parts if parts else [{"text": ""}]
def _transform_openai_delta_to_google_genai_parts_with_accumulation(
self, delta: Any, wrapper: GoogleGenAIStreamWrapper
) -> List[Dict[str, Any]]:
"""Transform OpenAI delta to Google GenAI parts format with tool call accumulation"""
"""Transforms OpenAI delta to Google GenAI parts, accumulating streaming tool calls."""
# 1. Initialize wrapper state if it doesn't exist
if not hasattr(wrapper, "accumulated_tool_calls"):
wrapper.accumulated_tool_calls = {}
parts: List[Dict[str, Any]] = []
# Add text content if present
if hasattr(delta, "content") and delta.content:
parts.append({"text": delta.content})
# Handle tool calls with accumulation for streaming
if hasattr(delta, "tool_calls") and delta.tool_calls:
for tool_call in delta.tool_calls:
if hasattr(tool_call, "function") and tool_call.function:
tool_call_id = getattr(tool_call, "id", "") or "call_unknown"
function_name = getattr(tool_call.function, "name", "") or ""
args_str = getattr(tool_call.function, "arguments", "") or ""
# 2. Ensure tool_calls is iterable
tool_calls = delta.tool_calls or []
# Initialize accumulation for this tool call if not exists
if tool_call_id not in wrapper.accumulated_tool_calls:
wrapper.accumulated_tool_calls[tool_call_id] = {
"name": "",
"arguments": "",
"complete": False,
}
for tool_call in tool_calls:
if not hasattr(tool_call, "function"):
continue
# Accumulate function name if provided
if function_name:
wrapper.accumulated_tool_calls[tool_call_id][
"name"
] = function_name
# 3. Use `index` as the primary key for accumulation
tool_call_index = getattr(tool_call, "index", None)
if tool_call_index is None:
continue # Index is essential for tracking streaming tool calls
# Accumulate arguments if provided
if args_str:
wrapper.accumulated_tool_calls[tool_call_id][
"arguments"
] += args_str
# Initialize accumulator for this index if it's new
if tool_call_index not in wrapper.accumulated_tool_calls:
wrapper.accumulated_tool_calls[tool_call_index] = {
"name": "",
"arguments": "",
}
# Try to parse the accumulated arguments as JSON
accumulated_args = wrapper.accumulated_tool_calls[tool_call_id][
"arguments"
]
try:
if accumulated_args:
parsed_args = json.loads(accumulated_args)
# JSON is valid, mark as complete and create function call part
wrapper.accumulated_tool_calls[tool_call_id][
"complete"
] = True
# Accumulate name and arguments
function_name = getattr(tool_call.function, "name", None)
args_chunk = getattr(tool_call.function, "arguments", None)
function_call_part = {
"functionCall": {
"name": wrapper.accumulated_tool_calls[
tool_call_id
]["name"],
"args": parsed_args,
}
}
parts.append(function_call_part)
# Optimization: Skip chunks that have no new data
if not function_name and not args_chunk:
verbose_logger.debug(
f"Skipping empty tool call chunk for index: {tool_call_index}"
)
continue
# Clean up completed tool call
del wrapper.accumulated_tool_calls[tool_call_id]
if function_name:
wrapper.accumulated_tool_calls[tool_call_index]["name"] = function_name
except json.JSONDecodeError:
# JSON is still incomplete, continue accumulating
# Don't add to parts yet
pass
if args_chunk:
wrapper.accumulated_tool_calls[tool_call_index][
"arguments"
] += args_chunk
# Attempt to parse and emit a complete tool call
accumulated_data = wrapper.accumulated_tool_calls[tool_call_index]
accumulated_name = accumulated_data["name"]
accumulated_args = accumulated_data["arguments"]
# 5. Attempt to parse arguments even if name hasn't arrived.
try:
# Attempt to parse the accumulated arguments string
parsed_args = json.loads(accumulated_args)
# If parsing succeeds, but we don't have a name yet, wait.
# The part will be created by a later chunk that brings the name.
if accumulated_name:
# If successful, create the part and clean up
function_call_part = {
"functionCall": {"name": accumulated_name, "args": parsed_args}
}
parts.append(function_call_part)
# Remove the completed tool call from the accumulator
del wrapper.accumulated_tool_calls[tool_call_index]
except json.JSONDecodeError:
# The JSON for arguments is still incomplete.
# We will continue to accumulate and wait for more chunks.
pass
return parts

31
litellm/google_genai/main.py
View file

@ -85,7 +85,6 @@ class GenerateContentHelper:
contents: GenerateContentContentListUnionDict,
config: Optional[GenerateContentConfigDict] = None,
custom_llm_provider: Optional[str] = None,
stream: bool = False,
tools: Optional[ToolConfigDict] = None,
**kwargs,
) -> GenerateContentSetupResult:
@ -97,8 +96,7 @@ class GenerateContentHelper:
contents: The content to generate from
config: Optional configuration
custom_llm_provider: Optional custom LLM provider
stream: Whether this is a streaming call
local_vars: Local variables from the calling function
tools: Optional tools
**kwargs: Additional keyword arguments
Returns:
@ -114,7 +112,7 @@ class GenerateContentHelper:
## MOCK RESPONSE LOGIC (only for non-streaming)
if (
not stream
not kwargs.get("stream", False)
and litellm_params.mock_response
and isinstance(litellm_params.mock_response, str)
):
@ -289,7 +287,7 @@ def generate_content(
"""
local_vars = locals()
try:
_is_async = kwargs.pop("agenerate_content", False) is True
_is_async = kwargs.pop("agenerate_content", False)
# Handle generationConfig parameter from kwargs for backward compatibility
if "generationConfig" in kwargs and config is None:
@ -309,7 +307,6 @@ def generate_content(
contents=contents,
config=config,
custom_llm_provider=custom_llm_provider,
stream=False,
tools=tools,
**kwargs,
)
@ -321,7 +318,7 @@ def generate_content(
model=model,
contents=contents, # type: ignore
config=setup_result.generate_content_config_dict,
stream=False,
tools=tools,
_is_async=_is_async,
litellm_params=setup_result.litellm_params,
**kwargs,
@ -342,7 +339,6 @@ def generate_content(
timeout=timeout or request_timeout,
_is_async=_is_async,
client=kwargs.get("client"),
stream=False,
litellm_metadata=kwargs.get("litellm_metadata", {}),
)
@ -391,15 +387,12 @@ async def agenerate_content_stream(
# Setup the call
setup_result = GenerateContentHelper.setup_generate_content_call(
**{
"model": model,
"contents": contents,
"config": config,
"custom_llm_provider": custom_llm_provider,
"stream": True,
"tools": tools,
**kwargs,
}
model=model,
contents=contents,
config=config,
custom_llm_provider=custom_llm_provider,
tools=tools,
**kwargs,
)
# Check if we should use the adapter (when provider config is None)
@ -411,6 +404,7 @@ async def agenerate_content_stream(
contents=contents, # type: ignore
config=setup_result.generate_content_config_dict,
litellm_params=setup_result.litellm_params,
tools=tools,
stream=True,
**kwargs,
)
@ -479,7 +473,6 @@ def generate_content_stream(
contents=contents,
config=config,
custom_llm_provider=custom_llm_provider,
stream=True,
tools=tools,
**kwargs,
)
@ -491,9 +484,9 @@ def generate_content_stream(
model=model,
contents=contents, # type: ignore
config=setup_result.generate_content_config_dict,
stream=True,
_is_async=_is_async,
litellm_params=setup_result.litellm_params,
stream=True,
**kwargs,
)

7
litellm/integrations/custom_guardrail.py
View file

@ -1,5 +1,5 @@
from datetime import datetime
from typing import Any, Dict, List, Literal, Optional, Type, Union, get_args
from typing import Any, Dict, List, Optional, Type, Union, get_args
from litellm._logging import verbose_logger
from litellm.caching import DualCache
@ -14,6 +14,7 @@ from litellm.types.guardrails import (
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
from litellm.types.utils import (
CallTypes,
GuardrailStatus,
LLMResponseTypes,
StandardLoggingGuardrailInformation,
)
@ -352,7 +353,7 @@ class CustomGuardrail(CustomLogger):
self,
guardrail_json_response: Union[Exception, str, dict, List[dict]],
request_data: dict,
guardrail_status: Literal["success", "failure", "blocked"],
guardrail_status: GuardrailStatus,
start_time: Optional[float] = None,
end_time: Optional[float] = None,
duration: Optional[float] = None,
@ -460,7 +461,7 @@ class CustomGuardrail(CustomLogger):
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=e,
request_data=request_data,
guardrail_status="failure",
guardrail_status="guardrail_failed_to_respond",
duration=duration,
start_time=start_time,
end_time=end_time,

317
litellm/integrations/gitlab/README.md Normal file
View file

@ -0,0 +1,317 @@
# LiteLLM gitlab Prompt Management
A powerful prompt management system for LiteLLM that fetches `.prompt` files from gitlab repositories. This enables team-based prompt management with gitlab's built-in access control and version control capabilities.
## Features
- **🏢 Team-based access control**: Leverage gitlab's workspace and repository permissions
- **📁 Repository-based prompt storage**: Store prompts in gitlab repositories
- **🔐 Multiple authentication methods**: Support for access tokens and basic auth
- **🎯 YAML frontmatter**: Define model, parameters, and schemas in file headers
- **🔧 Handlebars templating**: Use `{{variable}}` syntax with Jinja2 backend
- **✅ Input validation**: Automatic validation against defined schemas
- **🔗 LiteLLM integration**: Works seamlessly with `litellm.completion()`
- **💬 Smart message parsing**: Converts prompts to proper chat messages
- **⚙️ Parameter extraction**: Automatically applies model settings from prompts
## Quick Start
### 1. Set up gitlab Repository
Create a repository in your gitlab workspace and add `.prompt` files:
```
your-repo/
├── prompts/
│ ├── chat_assistant.prompt
│ ├── code_reviewer.prompt
│ └── data_analyst.prompt
```
### 2. Create a `.prompt` file
Create a file called `prompts/chat_assistant.prompt`:
```yaml
---
model: gpt-4
temperature: 0.7
max_tokens: 150
input:
schema:
user_message: string
system_context?: string
---
{% if system_context %}System: {{system_context}}
{% endif %}User: {{user_message}}
```
### 3. Configure gitlab Access
#### Option A: Access Token (Recommended)
```python
import litellm
# Configure gitlab access
gitlab_config = {
"project": "a/b/<repo_name>",
"access_token": "your-access-token",
"base_url": "gitlab url",
"prompts_path": "src/prompts", # folder to point to, defaults to root
"branch":"main" # optional, defaults to main
}
# Set global gitlab configuration
litellm.set_global_gitlab_config(gitlab_config)
```
#### Option B: Basic Authentication
```python
import litellm
# Configure gitlab access with basic auth
gitlab_config = {
"project": "a/b/<repo_name>",
"base_url": "base url",
"access_token": "your-app-password", # Use app password for basic auth
"branch": "main",
"prompts_path": "src/prompts", # folder to point to, defaults to root
}
litellm.set_global_gitlab_config(gitlab_config)
```
### 4. Use with LiteLLM
```python
# Use with completion - the model prefix 'gitlab/' tells LiteLLM to use gitlab prompt management
response = litellm.completion(
model="gitlab/gpt-4", # The actual model comes from the .prompt file
prompt_id="prompts/chat_assistant", # Location of the prompt file
prompt_variables={
"user_message": "What is machine learning?",
"system_context": "You are a helpful AI tutor."
},
# Any additional messages will be appended after the prompt
messages=[{"role": "user", "content": "Please explain it simply."}]
)
print(response.choices[0].message.content)
```
## Proxy Server Configuration
### 1. Create a `.prompt` file
Create `prompts/hello.prompt`:
```yaml
---
model: gpt-4
temperature: 0.7
---
System: You are a helpful assistant.
User: {{user_message}}
```
### 2. Setup config.yaml
```yaml
model_list:
- model_name: my-gitlab-model
litellm_params:
model: gitlab/gpt-4
prompt_id: "prompts/hello"
api_key: os.environ/OPENAI_API_KEY
litellm_settings:
global_gitlab_config:
workspace: "your-workspace"
repository: "your-repo"
access_token: "your-access-token"
branch: "main"
```
### 3. Start the proxy
```bash
litellm --config config.yaml --detailed_debug
```
### 4. Test it!
```bash
curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{
"model": "my-gitlab-model",
"messages": [{"role": "user", "content": "IGNORED"}],
"prompt_variables": {
"user_message": "What is the capital of France?"
}
}'
```
## Prompt File Format
### Basic Structure
```yaml
---
# Model configuration
model: gpt-4
temperature: 0.7
max_tokens: 500
# Input schema (optional)
input:
schema:
user_message: string
system_context?: string
---
System: You are a helpful {{role}} assistant.
User: {{user_message}}
```
### Advanced Features
**Multi-role conversations:**
```yaml
---
model: gpt-4
temperature: 0.3
---
System: You are a helpful coding assistant.
User: {{user_question}}
```
**Dynamic model selection:**
```yaml
---
model: "{{preferred_model}}" # Model can be a variable
temperature: 0.7
---
System: You are a helpful assistant specialized in {{domain}}.
User: {{user_message}}
```
## Team-Based Access Control
gitlab's built-in permission system provides team-based access control:
1. **Workspace-level permissions**: Control access to entire workspaces
2. **Repository-level permissions**: Control access to specific repositories
3. **Branch-level permissions**: Control access to specific branches
4. **User and group management**: Manage team members and their access levels
### Setting up Team Access
1. **Create workspaces for each team**:
```
team-a-prompts/
team-b-prompts/
team-c-prompts/
```
2. **Configure repository permissions**:
- Grant read access to team members
- Grant write access to prompt maintainers
- Use branch protection rules for production prompts
3. **Use different access tokens**:
- Each team can have their own access token
- Tokens can be scoped to specific repositories
- Use app passwords for additional security
## API Reference
### gitlab Configuration
```python
gitlab_config = {
"workspace": str, # Required: gitlab workspace name
"repository": str, # Required: Repository name
"access_token": str, # Required: gitlab access token or app password
"branch": str, # Optional: Branch to fetch from (default: "main")
"base_url": str, # Optional: Custom gitlab API URL
"auth_method": str, # Optional: "token" or "basic" (default: "token")
"username": str, # Optional: Username for basic auth
"base_url" : str # Optional: Incase where the base url is not https://api.gitlab.org/2.0
}
```
### LiteLLM Integration
```python
response = litellm.completion(
model="gitlab/<base_model>", # required (e.g., gitlab/gpt-4)
prompt_id=str, # required - the .prompt filename without extension
prompt_variables=dict, # optional - variables for template rendering
gitlab_config=dict, # optional - gitlab configuration (if not set globally)
messages=list, # optional - additional messages
)
```
## Error Handling
The gitlab integration provides detailed error messages for common issues:
- **Authentication errors**: Invalid access tokens or credentials
- **Permission errors**: Insufficient access to workspace/repository
- **File not found**: Missing .prompt files
- **Network errors**: Connection issues with gitlab API
## Security Considerations
1. **Access Token Security**: Store access tokens securely using environment variables or secret management systems
2. **Repository Permissions**: Use gitlab's permission system to control access
3. **Branch Protection**: Protect main branches from unauthorized changes
4. **Audit Logging**: gitlab provides audit logs for all repository access
## Troubleshooting
### Common Issues
1. **"Access denied" errors**: Check your gitlab permissions for the workspace and repository
2. **"Authentication failed" errors**: Verify your access token or credentials
3. **"File not found" errors**: Ensure the .prompt file exists in the specified branch
4. **Template rendering errors**: Check your Handlebars syntax in the .prompt file
### Debug Mode
Enable debug logging to troubleshoot issues:
```python
import litellm
litellm.set_verbose = True
# Your gitlab prompt calls will now show detailed logs
response = litellm.completion(
model="gitlab/gpt-4",
prompt_id="your_prompt",
prompt_variables={"key": "value"}
)
```
## Migration from File-Based Prompts
If you're currently using file-based prompts with the dotprompt integration, you can easily migrate to gitlab:
1. **Upload your .prompt files** to a gitlab repository
2. **Update your configuration** to use gitlab instead of local files
3. **Set up team access** using gitlab's permission system
4. **Update your code** to use `gitlab/` model prefix instead of `dotprompt/`
This provides better collaboration, version control, and team-based access control for your prompts.

95
litellm/integrations/gitlab/__init__.py Normal file
View file

@ -0,0 +1,95 @@
from typing import TYPE_CHECKING, Optional, Dict, Any
if TYPE_CHECKING:
from .gitlab_prompt_manager import GitLabPromptManager
from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec
from litellm.integrations.custom_prompt_management import CustomPromptManagement
from litellm.types.prompts.init_prompts import SupportedPromptIntegrations
from litellm.integrations.custom_prompt_management import CustomPromptManagement
from litellm.types.prompts.init_prompts import PromptSpec, PromptLiteLLMParams
from .gitlab_prompt_manager import GitLabPromptManager
# Global instances
global_gitlab_config: Optional[dict] = None
def set_global_gitlab_config(config: dict) -> None:
"""
Set the global BitBucket configuration for prompt management.
Args:
config: Dictionary containing BitBucket configuration
- workspace: BitBucket workspace name
- repository: Repository name
- access_token: BitBucket access token
- branch: Branch to fetch prompts from (default: main)
"""
import litellm
litellm.global_gitlab_config = config # type: ignore
def prompt_initializer(
litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec"
) -> "CustomPromptManagement":
"""
Initialize a prompt from a BitBucket repository.
"""
gitlab_config = getattr(litellm_params, "gitlab_config", None)
prompt_id = getattr(litellm_params, "prompt_id", None)
if not gitlab_config:
raise ValueError(
"bitbucket_config is required for BitBucket prompt integration"
)
try:
bitbucket_prompt_manager = GitLabPromptManager(
gitlab_config=gitlab_config,
prompt_id=prompt_id,
)
return bitbucket_prompt_manager
except Exception as e:
raise e
def _gitlab_prompt_initializer(
litellm_params: PromptLiteLLMParams,
prompt: PromptSpec,
) -> CustomPromptManagement:
"""
Build a GitLab-backed prompt manager for this prompt.
Expected fields on litellm_params:
- prompt_integration="gitlab" (handled by the caller)
- gitlab_config: Dict[str, Any] (project/access_token/branch/prompts_path/etc.)
- git_ref (optional): per-prompt tag/branch/SHA override
"""
# You can store arbitrary integration-specific config on PromptLiteLLMParams.
# If your dataclass doesn't have these attributes, add them or put inside
# `litellm_params.extra` and pull them from there.
gitlab_config: Dict[str, Any] = getattr(litellm_params, "gitlab_config", None) or {}
git_ref: Optional[str] = getattr(litellm_params, "git_ref", None)
if not gitlab_config:
raise ValueError("gitlab_config is required for gitlab prompt integration")
# prompt.prompt_id can map to a file path under prompts_path (e.g. "chat/greet/hi")
return GitLabPromptManager(
gitlab_config=gitlab_config,
prompt_id=prompt.prompt_id,
ref=git_ref,
)
prompt_initializer_registry = {
SupportedPromptIntegrations.GITLAB.value: _gitlab_prompt_initializer,
}
# Export public API
__all__ = [
"GitLabPromptManager",
"set_global_gitlab_config",
"global_gitlab_config",
]

285
litellm/integrations/gitlab/gitlab_client.py Normal file
View file

@ -0,0 +1,285 @@
"""
GitLab API client for fetching files from GitLab repositories.
Now supports selecting a tag via `config["tag"]`; falls back to branch ("main").
"""
import base64
from typing import Any, Dict, List, Optional
from urllib.parse import quote
from litellm.llms.custom_httpx.http_handler import HTTPHandler
class GitLabClient:
"""
Client for interacting with the GitLab API to fetch files.
Supports:
- Authentication with personal/access tokens or OAuth bearer tokens
- Fetching file contents from repositories (raw endpoint with JSON fallback)
- Namespace/project path or numeric project ID addressing
- Ref selection via tag (preferred) or branch (default "main")
- Directory listing via the repository tree API
"""
def __init__(self, config: Dict[str, Any]):
"""
Initialize the GitLab client.
Args:
config: Dictionary containing:
- project: Project path ("group/subgroup/repo") or numeric project ID (str|int) [required]
- access_token: GitLab personal/access token or OAuth token [required] (str)
- auth_method: 'token' (default; sends Private-Token) or 'oauth' (Authorization: Bearer)
- tag: Tag name to fetch from (takes precedence over branch if provided)
- branch: Branch to fetch from (default: "main")
- base_url: Base GitLab API URL (default: "https://gitlab.com/api/v4")
"""
project = config.get("project")
access_token = config.get("access_token")
if project is None or access_token is None:
raise ValueError("project and access_token are required")
self.project: str | int = project
self.access_token: str = str(access_token)
self.auth_method = config.get("auth_method", "token") # 'token' or 'oauth'
self.branch = config.get("branch", None)
if not self.branch:
self.branch = 'main'
self.tag = config.get("tag")
self.base_url = config.get("base_url", "https://gitlab.com/api/v4")
if not all([self.project, self.access_token]):
raise ValueError("project and access_token are required")
# Effective ref: prefer tag if provided, else branch ("main")
self.ref = str(self.tag or self.branch)
# Build headers
self.headers = {
"Accept": "application/json",
"Content-Type": "application/json",
}
if self.auth_method == "oauth":
self.headers["Authorization"] = f"Bearer {self.access_token}"
else:
# Default GitLab token header
self.headers["Private-Token"] = self.access_token
# Project identifier must be URL-encoded (slashes become %2F)
self._project_enc = quote(str(self.project), safe="")
# HTTP handler
self.http_handler = HTTPHandler()
# ------------------------
# Core helpers
# ------------------------
def _file_raw_url(self, file_path: str, *, ref: Optional[str] = None) -> str:
file_enc = quote(file_path, safe="")
ref_q = quote(ref or self.ref, safe="")
return f"{self.base_url}/projects/{self._project_enc}/repository/files/{file_enc}/raw?ref={ref_q}"
def _file_json_url(self, file_path: str, *, ref: Optional[str] = None) -> str:
file_enc = quote(file_path, safe="")
ref_q = quote(ref or self.ref, safe="")
return f"{self.base_url}/projects/{self._project_enc}/repository/files/{file_enc}?ref={ref_q}"
def _tree_url(self, directory_path: str = "", recursive: bool = False, *, ref: Optional[str] = None) -> str:
path_q = f"&path={quote(directory_path, safe='')}" if directory_path else ""
rec_q = "&recursive=true" if recursive else ""
ref_q = quote(ref or self.ref, safe="")
return f"{self.base_url}/projects/{self._project_enc}/repository/tree?ref={ref_q}{path_q}{rec_q}"
# ------------------------
# Public API
# ------------------------
def set_ref(self, ref: str) -> None:
"""Override the default ref (tag/branch) for subsequent calls."""
if not ref:
raise ValueError("ref must be a non-empty string")
self.ref = ref
def get_file_content(self, file_path: str, *, ref: Optional[str] = None) -> Optional[str]:
"""
Fetch the content of a file from the GitLab repository at the given ref
(tag, branch, or commit SHA). If `ref` is None, uses self.ref.
Strategy:
1) Try the RAW endpoint (returns bytes of the file)
2) Fallback to the JSON endpoint (returns base64-encoded content)
Returns:
File content as UTF-8 string, or None if file not found.
"""
raw_url = self._file_raw_url(file_path, ref=ref)
try:
resp = self.http_handler.get(raw_url, headers=self.headers)
if resp.status_code == 404:
# Fallback to JSON endpoint
return self._get_file_content_via_json(file_path, ref=ref)
resp.raise_for_status()
ctype = (resp.headers.get("content-type") or "").lower()
if ctype.startswith("text/") or "charset=" in ctype or ctype.startswith("application/json"):
return resp.text
try:
return resp.content.decode("utf-8")
except Exception:
return resp.content.decode("utf-8", errors="replace")
except Exception as e:
status = getattr(getattr(e, "response", None), "status_code", None)
if status == 404:
return None
if status == 403:
raise Exception(
f"Access denied to file '{file_path}'. Check your GitLab permissions for project '{self.project}'."
)
if status == 401:
raise Exception("Authentication failed. Check your GitLab token and auth_method.")
raise Exception(f"Failed to fetch file '{file_path}': {e}")
def _get_file_content_via_json(self, file_path: str, *, ref: Optional[str] = None) -> Optional[str]:
"""
Fallback for get_file_content(): use the JSON file API which returns base64 content.
"""
json_url = self._file_json_url(file_path, ref=ref)
try:
resp = self.http_handler.get(json_url, headers=self.headers)
if resp.status_code == 404:
return None
resp.raise_for_status()
data = resp.json()
content = data.get("content")
encoding = data.get("encoding", "")
if content and encoding == "base64":
try:
return base64.b64decode(content).decode("utf-8")
except Exception:
return base64.b64decode(content).decode("utf-8", errors="replace")
return content
except Exception as e:
status = getattr(getattr(e, "response", None), "status_code", None)
if status == 404:
return None
if status == 403:
raise Exception(
f"Access denied to file '{file_path}'. Check your GitLab permissions for project '{self.project}'."
)
if status == 401:
raise Exception("Authentication failed. Check your GitLab token and auth_method.")
raise Exception(f"Failed to fetch file '{file_path}' via JSON endpoint: {e}")
def list_files(
self,
directory_path: str = "",
file_extension: str = ".prompt",
recursive: bool = False,
*,
ref: Optional[str] = None,
) -> List[str]:
"""
List files in a directory with a specific extension using the repository tree API.
Args:
directory_path: Directory path in the repository (empty for repo root)
file_extension: File extension to filter by (default: .prompt)
recursive: If True, traverses subdirectories
ref: Optional override (tag/branch/SHA). Defaults to self.ref.
Returns:
List of file paths (relative to repo root)
"""
url = self._tree_url(directory_path, recursive=recursive, ref=ref)
try:
resp = self.http_handler.get(url, headers=self.headers)
if resp.status_code == 404:
return []
resp.raise_for_status()
data = resp.json() or []
files: List[str] = []
for item in data:
if item.get("type") == "blob":
file_path = item.get("path", "")
if not file_extension or file_path.endswith(file_extension):
files.append(file_path)
return files
except Exception as e:
status = getattr(getattr(e, "response", None), "status_code", None)
if status == 404:
return []
if status == 403:
raise Exception(
f"Access denied to directory '{directory_path}'. Check your GitLab permissions for project '{self.project}'."
)
if status == 401:
raise Exception("Authentication failed. Check your GitLab token and auth_method.")
raise Exception(f"Failed to list files in '{directory_path}': {e}")
def get_repository_info(self) -> Dict[str, Any]:
"""Get information about the project/repository."""
url = f"{self.base_url}/projects/{self._project_enc}"
try:
resp = self.http_handler.get(url, headers=self.headers)
resp.raise_for_status()
return resp.json()
except Exception as e:
raise Exception(f"Failed to get repository info: {e}")
def test_connection(self) -> bool:
"""Test the connection to the GitLab project."""
try:
self.get_repository_info()
return True
except Exception:
return False
def get_branches(self) -> List[Dict[str, Any]]:
"""Get list of branches in the repository."""
url = f"{self.base_url}/projects/{self._project_enc}/repository/branches"
try:
resp = self.http_handler.get(url, headers=self.headers)
resp.raise_for_status()
data = resp.json()
return data if isinstance(data, list) else []
except Exception as e:
raise Exception(f"Failed to get branches: {e}")
def get_file_metadata(self, file_path: str, *, ref: Optional[str] = None) -> Optional[Dict[str, Any]]:
"""
Get minimal metadata about a file via RAW endpoint headers at a given ref.
Args:
file_path: Path to the file in the repository.
ref: Optional override (tag/branch/SHA). Defaults to self.ref.
"""
url = self._file_raw_url(file_path, ref=ref)
try:
headers = dict(self.headers)
headers["Range"] = "bytes=0-0"
resp = self.http_handler.get(url, headers=headers)
if resp.status_code == 404:
return None
resp.raise_for_status()
return {
"content_type": resp.headers.get("content-type"),
"content_length": resp.headers.get("content-length"),
"last_modified": resp.headers.get("last-modified"),
}
except Exception as e:
status = getattr(getattr(e, "response", None), "status_code", None)
if status == 404:
return None
raise Exception(f"Failed to get file metadata for '{file_path}': {e}")
def close(self):
"""Close the HTTP handler to free resources."""
if hasattr(self, "http_handler"):
self.http_handler.close()

488
litellm/integrations/gitlab/gitlab_prompt_manager.py Normal file
View file

@ -0,0 +1,488 @@
"""
GitLab prompt manager with configurable prompts folder.
"""
from typing import Any, Dict, List, Optional, Tuple, Union
from jinja2 import DictLoader, Environment, select_autoescape
from litellm.integrations.custom_prompt_management import CustomPromptManagement
from litellm.integrations.prompt_management_base import (
PromptManagementBase,
PromptManagementClient,
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import StandardCallbackDynamicParams
from litellm.integrations.gitlab.gitlab_client import GitLabClient
class GitLabPromptTemplate:
def __init__(
self,
template_id: str,
content: str,
metadata: Dict[str, Any],
model: Optional[str] = None,
):
self.template_id = template_id
self.content = content
self.metadata = metadata
self.model = model or metadata.get("model")
self.temperature = metadata.get("temperature")
self.max_tokens = metadata.get("max_tokens")
self.input_schema = metadata.get("input", {}).get("schema", {})
self.optional_params = {
k: v for k, v in metadata.items() if k not in ["model", "input", "content"]
}
def __repr__(self):
return f"GitLabPromptTemplate(id='{self.template_id}', model='{self.model}')"
class GitLabTemplateManager:
"""
Manager for loading and rendering .prompt files from GitLab repositories.
New: supports `prompts_path` (or `folder`) in gitlab_config to scope where prompts live.
"""
def __init__(
self,
gitlab_config: Dict[str, Any],
prompt_id: Optional[str] = None,
ref: Optional[str] = None,
gitlab_client: Optional[GitLabClient] = None
):
self.gitlab_config = dict(gitlab_config)
self.prompt_id = prompt_id
self.prompts: Dict[str, GitLabPromptTemplate] = {}
self.gitlab_client = gitlab_client or GitLabClient(self.gitlab_config)
if ref:
self.gitlab_client.set_ref(ref)
# Folder inside repo to look for prompts (e.g., "prompts" or "prompts/chat")
self.prompts_path: str = (
self.gitlab_config.get("prompts_path")
or self.gitlab_config.get("folder")
or ""
).strip("/")
self.jinja_env = Environment(
loader=DictLoader({}),
autoescape=select_autoescape(["html", "xml"]),
variable_start_string="{{",
variable_end_string="}}",
block_start_string="{%",
block_end_string="%}",
comment_start_string="{#",
comment_end_string="#}",
)
if self.prompt_id:
self._load_prompt_from_gitlab(self.prompt_id)
# ---------- path helpers ----------
def _id_to_repo_path(self, prompt_id: str) -> str:
"""Map a prompt_id to a repo path (respects prompts_path and adds .prompt)."""
if self.prompts_path:
return f"{self.prompts_path}/{prompt_id}.prompt"
return f"{prompt_id}.prompt"
def _repo_path_to_id(self, repo_path: str) -> str:
"""
Map a repo path like 'prompts/chat/greeting.prompt' to an ID relative
to prompts_path without the extension (e.g., 'chat/greeting').
"""
path = repo_path.strip("/")
if self.prompts_path and path.startswith(self.prompts_path.strip("/") + "/"):
path = path[len(self.prompts_path.strip("/")) + 1 :]
if path.endswith(".prompt"):
path = path[: -len(".prompt")]
return path
# ---------- loading ----------
def _load_prompt_from_gitlab(self, prompt_id: str, *, ref: Optional[str] = None) -> None:
"""Load a specific .prompt file from GitLab (scoped under prompts_path if set)."""
try:
file_path = self._id_to_repo_path(prompt_id)
prompt_content = self.gitlab_client.get_file_content(file_path, ref=ref)
if prompt_content:
template = self._parse_prompt_file(prompt_content, prompt_id)
self.prompts[prompt_id] = template
except Exception as e:
raise Exception(f"Failed to load prompt '{prompt_id}' from GitLab: {e}")
def load_all_prompts(self, *, recursive: bool = True) -> List[str]:
"""
Eagerly load all .prompt files from prompts_path. Returns loaded IDs.
"""
files = self.list_templates(recursive=recursive) # reuse logic
loaded: List[str] = []
for pid in files:
if pid not in self.prompts:
self._load_prompt_from_gitlab(pid)
loaded.append(pid)
return loaded
# ---------- parsing & rendering ----------
def _parse_prompt_file(
self, content: str, prompt_id: str
) -> GitLabPromptTemplate:
if content.startswith("---"):
parts = content.split("---", 2)
if len(parts) >= 3:
frontmatter_str = parts[1].strip()
template_content = parts[2].strip()
else:
frontmatter_str = ""
template_content = content
else:
frontmatter_str = ""
template_content = content
metadata: Dict[str, Any] = {}
if frontmatter_str:
try:
import yaml
metadata = yaml.safe_load(frontmatter_str) or {}
except ImportError:
metadata = self._parse_yaml_basic(frontmatter_str)
except Exception:
metadata = {}
return GitLabPromptTemplate(
template_id=prompt_id,
content=template_content,
metadata=metadata,
)
def _parse_yaml_basic(self, yaml_str: str) -> Dict[str, Any]:
result: Dict[str, Any] = {}
for line in yaml_str.split("\n"):
line = line.strip()
if ":" in line and not line.startswith("#"):
key, value = line.split(":", 1)
key = key.strip()
value = value.strip()
if value.lower() in ["true", "false"]:
result[key] = value.lower() == "true"
elif value.isdigit():
result[key] = int(value)
elif value.replace(".", "").isdigit():
try:
result[key] = float(value)
except Exception:
result[key] = value
else:
result[key] = value.strip("\"'")
return result
def render_template(
self, template_id: str, variables: Optional[Dict[str, Any]] = None
) -> str:
if template_id not in self.prompts:
raise ValueError(f"Template '{template_id}' not found")
template = self.prompts[template_id]
jinja_template = self.jinja_env.from_string(template.content)
return jinja_template.render(**(variables or {}))
def get_template(self, template_id: str) -> Optional[GitLabPromptTemplate]:
return self.prompts.get(template_id)
def list_templates(self, *, recursive: bool = True) -> List[str]:
"""
List available prompt IDs discovered under prompts_path (no extension, relative to prompts_path).
"""
"""
List available prompt IDs under prompts_path (no extension).
Compatible with both list_files signatures:
- list_files(directory_path=..., file_extension=..., recursive=...)
- list_files(path=..., ref=None, recursive=...)
"""
# First try the "new" signature (directory_path/file_extension)
try:
files = self.gitlab_client.list_files(
directory_path=self.prompts_path,
file_extension=".prompt",
recursive=recursive,
)
base = self.prompts_path.strip("/")
out: List[str] = []
for p in files or []:
path = str(p).strip("/")
if base and not path.startswith(base + "/"):
# if the client returns extra files outside the folder, skip them
continue
if not path.endswith(".prompt"):
continue
out.append(self._repo_path_to_id(path))
return out
except TypeError:
# Fallback to the "classic" signature
raw = self.gitlab_client.list_files(
directory_path=self.prompts_path or "",
ref=None,
recursive=recursive,
)
# Classic returns GitLab tree entries; filter *.prompt blobs
files = []
for f in (raw or []):
if isinstance(f, dict) and f.get("type") == "blob" and str(f.get("path", "")).endswith(".prompt") and 'path' in f:
files.append(f['path'])
return [self._repo_path_to_id(p) for p in files]
class GitLabPromptManager(CustomPromptManagement):
"""
GitLab prompt manager with folder support.
Example config:
gitlab_config = {
"project": "group/subgroup/repo",
"access_token": "glpat_***",
"tag": "v1.2.3", # optional; takes precedence
"branch": "main", # default fallback
"prompts_path": "prompts/chat" # <--- NEW
}
"""
def __init__(
self,
gitlab_config: Dict[str, Any],
prompt_id: Optional[str] = None,
ref: Optional[str] = None, # tag/branch/SHA override
gitlab_client: Optional[GitLabClient] = None
):
self.gitlab_config = gitlab_config
self.prompt_id = prompt_id
self._prompt_manager: Optional[GitLabTemplateManager] = None
self._ref_override = ref
self._injected_gitlab_client = gitlab_client
if self.prompt_id:
self._prompt_manager = GitLabTemplateManager(
gitlab_config=self.gitlab_config,
prompt_id=self.prompt_id,
ref=self._ref_override,
)
@property
def integration_name(self) -> str:
return "gitlab"
@property
def prompt_manager(self) -> GitLabTemplateManager:
if self._prompt_manager is None:
self._prompt_manager = GitLabTemplateManager(
gitlab_config=self.gitlab_config,
prompt_id=self.prompt_id,
ref=self._ref_override,
gitlab_client=self._injected_gitlab_client
)
return self._prompt_manager
def get_prompt_template(
self,
prompt_id: str,
prompt_variables: Optional[Dict[str, Any]] = None,
*,
ref: Optional[str] = None,
) -> Tuple[str, Dict[str, Any]]:
if prompt_id not in self.prompt_manager.prompts:
self.prompt_manager._load_prompt_from_gitlab(prompt_id, ref=ref)
template = self.prompt_manager.get_template(prompt_id)
if not template:
raise ValueError(f"Prompt template '{prompt_id}' not found")
rendered_prompt = self.prompt_manager.render_template(
prompt_id, prompt_variables or {}
)
metadata = {
"model": template.model,
"temperature": template.temperature,
"max_tokens": template.max_tokens,
**template.optional_params,
}
return rendered_prompt, metadata
def pre_call_hook(
self,
user_id: Optional[str],
messages: List[AllMessageValues],
function_call: Optional[Union[Dict[str, Any], str]] = None,
litellm_params: Optional[Dict[str, Any]] = None,
prompt_id: Optional[str] = None,
prompt_variables: Optional[Dict[str, Any]] = None,
prompt_version: Optional[str] = None,
**kwargs,
) -> Tuple[List[AllMessageValues], Optional[Dict[str, Any]]]:
if not prompt_id:
return messages, litellm_params
try:
# Precedence: explicit prompt_version → per-call git_ref kwarg → manager override → config default
git_ref = prompt_version or kwargs.get("git_ref") or self._ref_override
rendered_prompt, prompt_metadata = self.get_prompt_template(
prompt_id, prompt_variables, ref=git_ref
)
parsed_messages = self._parse_prompt_to_messages(rendered_prompt)
if parsed_messages:
final_messages: List[AllMessageValues] = parsed_messages
else:
final_messages = [{"role": "user", "content": rendered_prompt}] + messages # type: ignore
if litellm_params is None:
litellm_params = {}
if prompt_metadata.get("model"):
litellm_params["model"] = prompt_metadata["model"]
for param in ["temperature", "max_tokens", "top_p", "frequency_penalty", "presence_penalty"]:
if param in prompt_metadata:
litellm_params[param] = prompt_metadata[param]
return final_messages, litellm_params
except Exception as e:
import litellm
litellm._logging.verbose_proxy_logger.error(f"Error in GitLab prompt pre_call_hook: {e}")
return messages, litellm_params
def _parse_prompt_to_messages(self, prompt_content: str) -> List[AllMessageValues]:
messages: List[AllMessageValues] = []
lines = prompt_content.strip().split("\n")
current_role: Optional[str] = None
current_content: List[str] = []
for raw in lines:
line = raw.strip()
if not line:
continue
low = line.lower()
if low.startswith("system:"):
if current_role and current_content:
messages.append({"role": current_role, "content": "\n".join(current_content).strip()}) # type: ignore
current_role = "system"
current_content = [line[7:].strip()]
elif low.startswith("user:"):
if current_role and current_content:
messages.append({"role": current_role, "content": "\n".join(current_content).strip()}) # type: ignore
current_role = "user"
current_content = [line[5:].strip()]
elif low.startswith("assistant:"):
if current_role and current_content:
messages.append({"role": current_role, "content": "\n".join(current_content).strip()}) # type: ignore
current_role = "assistant"
current_content = [line[10:].strip()]
else:
current_content.append(line)
if current_role and current_content:
messages.append({"role": current_role, "content": "\n".join(current_content).strip()}) # type: ignore
if not messages and prompt_content.strip():
messages = [{"role": "user", "content": prompt_content.strip()}] # type: ignore
return messages
def post_call_hook(
self,
user_id: Optional[str],
response: Any,
input_messages: List[AllMessageValues],
function_call: Optional[Union[Dict[str, Any], str]] = None,
litellm_params: Optional[Dict[str, Any]] = None,
prompt_id: Optional[str] = None,
prompt_variables: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Any:
return response
def get_available_prompts(self) -> List[str]:
"""
Return prompt IDs. Prefer already-loaded templates in memory to avoid
unnecessary network calls (and to make tests deterministic).
"""
ids = set(self.prompt_manager.prompts.keys())
try:
ids.update(self.prompt_manager.list_templates())
except Exception:
# If GitLab list fails (auth, network), still return what we've loaded.
pass
return sorted(ids)
def reload_prompts(self) -> None:
if self.prompt_id:
self._prompt_manager = None
_ = self.prompt_manager # trigger re-init/load
def should_run_prompt_management(
self,
prompt_id: str,
dynamic_callback_params: StandardCallbackDynamicParams,
) -> bool:
return True
def _compile_prompt_helper(
self,
prompt_id: str,
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str] = None,
prompt_version: Optional[int] = None,
) -> PromptManagementClient:
try:
if prompt_id not in self.prompt_manager.prompts:
git_ref = getattr(dynamic_callback_params, "extra", {}).get("git_ref") if hasattr(dynamic_callback_params, "extra") else None
self.prompt_manager._load_prompt_from_gitlab(prompt_id, ref=git_ref)
rendered_prompt, prompt_metadata = self.get_prompt_template(
prompt_id, prompt_variables
)
messages = self._parse_prompt_to_messages(rendered_prompt)
template_model = prompt_metadata.get("model")
optional_params: Dict[str, Any] = {}
for param in ["temperature", "max_tokens", "top_p", "frequency_penalty", "presence_penalty"]:
if param in prompt_metadata:
optional_params[param] = prompt_metadata[param]
return PromptManagementClient(
prompt_id=prompt_id,
prompt_template=messages,
prompt_template_model=template_model,
prompt_template_optional_params=optional_params,
completed_messages=None,
)
except Exception as e:
raise ValueError(f"Error compiling prompt '{prompt_id}': {e}")
def get_chat_completion_prompt(
self,
model: str,
messages: List[AllMessageValues],
non_default_params: dict,
prompt_id: Optional[str],
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str] = None,
prompt_version: Optional[int] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
return PromptManagementBase.get_chat_completion_prompt(
self,
model,
messages,
non_default_params,
prompt_id,
prompt_variables,
dynamic_callback_params,
prompt_label,
prompt_version,
)

5
litellm/integrations/langfuse/langfuse.py
View file

@ -1,6 +1,5 @@
#### What this does ####
# On success, logs events to Langfuse
import copy
import os
import traceback
from datetime import datetime
@ -11,6 +10,7 @@ from packaging.version import Version
import litellm
from litellm._logging import verbose_logger
from litellm.constants import MAX_LANGFUSE_INITIALIZED_CLIENTS
from litellm.litellm_core_utils.core_helpers import safe_deep_copy
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
@ -222,7 +222,7 @@ class LangFuseLogger:
litellm_params.get("metadata", {}) or {}
) # if litellm_params['metadata'] == None
metadata = self.add_metadata_from_header(litellm_params, metadata)
optional_params = copy.deepcopy(kwargs.get("optional_params", {}))
optional_params = safe_deep_copy(kwargs.get("optional_params", {}))
prompt = {"messages": kwargs.get("messages")}
@ -690,6 +690,7 @@ class LangFuseLogger:
}
usage_details = LangfuseUsageDetails(input=_usage_obj.prompt_tokens,
output=_usage_obj.completion_tokens,
total=_usage_obj.total_tokens,
cache_creation_input_tokens=_usage_obj.get('cache_creation_input_tokens', 0),
cache_read_input_tokens=_usage_obj.get('cache_read_input_tokens', 0))

99
litellm/integrations/opentelemetry.py
View file

@ -575,9 +575,16 @@ class OpenTelemetry(CustomLogger):
if litellm.turn_off_message_logging or not self.message_logging:
return
litellm_params = kwargs.get("litellm_params", {})
metadata = litellm_params.get("metadata") or {}
generation_name = metadata.get("generation_name")
raw_span_name = generation_name if generation_name else RAW_REQUEST_SPAN_NAME
otel_tracer: Tracer = self.get_tracer_to_use_for_request(kwargs)
raw_span = otel_tracer.start_span(
name=RAW_REQUEST_SPAN_NAME,
name=raw_span_name,
start_time=self._to_ns(start_time),
context=trace.set_span_in_context(parent_span),
)
@ -645,7 +652,7 @@ class OpenTelemetry(CustomLogger):
if not self.config.enable_events:
return
from opentelemetry._logs import get_logger, LogRecord
from opentelemetry._logs import LogRecord, get_logger
otel_logger = get_logger(LITELLM_LOGGER_NAME)
parent_ctx = span.get_span_context()
@ -1115,56 +1122,68 @@ class OpenTelemetry(CustomLogger):
span.set_attribute(key, primitive_value)
def set_raw_request_attributes(self, span: Span, kwargs, response_obj):
kwargs.get("optional_params", {})
litellm_params = kwargs.get("litellm_params", {}) or {}
custom_llm_provider = litellm_params.get("custom_llm_provider", "Unknown")
try:
kwargs.get("optional_params", {})
litellm_params = kwargs.get("litellm_params", {}) or {}
custom_llm_provider = litellm_params.get("custom_llm_provider", "Unknown")
_raw_response = kwargs.get("original_response")
_additional_args = kwargs.get("additional_args", {}) or {}
complete_input_dict = _additional_args.get("complete_input_dict")
#############################################
########## LLM Request Attributes ###########
#############################################
_raw_response = kwargs.get("original_response")
_additional_args = kwargs.get("additional_args", {}) or {}
complete_input_dict = _additional_args.get("complete_input_dict")
#############################################
########## LLM Request Attributes ###########
#############################################
# OTEL Attributes for the RAW Request to https://docs.anthropic.com/en/api/messages
if complete_input_dict and isinstance(complete_input_dict, dict):
for param, val in complete_input_dict.items():
self.safe_set_attribute(
span=span, key=f"llm.{custom_llm_provider}.{param}", value=val
)
# OTEL Attributes for the RAW Request to https://docs.anthropic.com/en/api/messages
if complete_input_dict and isinstance(complete_input_dict, dict):
for param, val in complete_input_dict.items():
self.safe_set_attribute(
span=span, key=f"llm.{custom_llm_provider}.{param}", value=val
)
#############################################
########## LLM Response Attributes ##########
#############################################
if _raw_response and isinstance(_raw_response, str):
# cast sr -> dict
import json
#############################################
########## LLM Response Attributes ##########
#############################################
if _raw_response and isinstance(_raw_response, str):
# cast sr -> dict
import json
try:
_raw_response = json.loads(_raw_response)
for param, val in _raw_response.items():
self.safe_set_attribute(
span=span,
key=f"llm.{custom_llm_provider}.{param}",
value=val,
)
except json.JSONDecodeError:
verbose_logger.debug(
"litellm.integrations.opentelemetry.py::set_raw_request_attributes() - raw_response not json string - {}".format(
_raw_response
)
)
try:
_raw_response = json.loads(_raw_response)
for param, val in _raw_response.items():
self.safe_set_attribute(
span=span,
key=f"llm.{custom_llm_provider}.{param}",
value=val,
key=f"llm.{custom_llm_provider}.stringified_raw_response",
value=_raw_response,
)
except json.JSONDecodeError:
verbose_logger.debug(
"litellm.integrations.opentelemetry.py::set_raw_request_attributes() - raw_response not json string - {}".format(
_raw_response
)
)
self.safe_set_attribute(
span=span,
key=f"llm.{custom_llm_provider}.stringified_raw_response",
value=_raw_response,
)
except Exception as e:
verbose_logger.exception(
"OpenTelemetry logging error in set_raw_request_attributes %s", str(e)
)
def _to_ns(self, dt):
return int(dt.timestamp() * 1e9)
def _get_span_name(self, kwargs):
litellm_params = kwargs.get("litellm_params", {})
metadata = litellm_params.get("metadata") or {}
generation_name = metadata.get("generation_name")
if generation_name:
return generation_name
return LITELLM_REQUEST_SPAN_NAME
def get_traceparent_from_header(self, headers):

2
litellm/litellm_core_utils/custom_logger_registry.py
View file

@ -16,6 +16,7 @@ from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheCont
from litellm.integrations.argilla import ArgillaLogger
from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger
from litellm.integrations.bitbucket import BitBucketPromptManager
from litellm.integrations.gitlab import GitLabPromptManager
from litellm.integrations.braintrust_logging import BraintrustLogger
from litellm.integrations.datadog.datadog import DataDogLogger
from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger
@ -92,6 +93,7 @@ class CustomLoggerRegistry:
"vector_store_pre_call_hook": VectorStorePreCallHook,
"dotprompt": DotpromptManager,
"bitbucket": BitBucketPromptManager,
"gitlab": GitLabPromptManager,
"cloudzero": CloudZeroLogger,
"posthog": PostHogLogger,
}

72
litellm/litellm_core_utils/litellm_logging.py
View file

@ -89,6 +89,7 @@ from litellm.types.utils import (
CostResponseTypes,
DynamicPromptManagementParamLiteral,
EmbeddingResponse,
GuardrailStatus,
ImageResponse,
LiteLLMBatch,
LiteLLMLoggingBaseClass,
@ -107,6 +108,7 @@ from litellm.types.utils import (
StandardLoggingPayload,
StandardLoggingPayloadErrorInformation,
StandardLoggingPayloadStatus,
StandardLoggingPayloadStatusFields,
StandardLoggingPromptManagementMetadata,
StandardLoggingVectorStoreRequest,
TextCompletionResponse,
@ -3667,6 +3669,25 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
bitbucket_logger = BitBucketPromptManager(bitbucket_config=bitbucket_config)
_in_memory_loggers.append(bitbucket_logger)
return bitbucket_logger # type: ignore
elif logging_integration == "gitlab":
from litellm.integrations.gitlab.gitlab_prompt_manager import (
GitLabPromptManager,
)
for callback in _in_memory_loggers:
if isinstance(callback, GitLabPromptManager):
return callback
# Get global BitBucket config
gitlab_config = getattr(litellm, "global_gitlab_config", None)
if gitlab_config is None:
raise ValueError(
"Gitlab configuration not found. Please set litellm.global_gitlab_config first."
)
gitlab_logger = GitLabPromptManager(gitlab_config=gitlab_config)
_in_memory_loggers.append(gitlab_logger)
return gitlab_logger # type: ignore
return None
except Exception as e:
verbose_logger.exception(
@ -4426,6 +4447,51 @@ class StandardLoggingPayloadSetup:
return request_tags
def _get_status_fields(
status: StandardLoggingPayloadStatus,
guardrail_information: Optional[dict],
error_str: Optional[str]
) -> "StandardLoggingPayloadStatusFields":
"""
Determine status fields based on request status and guardrail information.
Args:
status: Overall request status ("success" or "failure")
guardrail_information: Guardrail information from metadata
error_str: Error string if any
Returns:
StandardLoggingPayloadStatusFields with llm_api_status and guardrail_status
"""
# Mapping for legacy guardrail status values to new GuardrailStatus values
GUARDRAIL_STATUS_MAP: Dict[str, GuardrailStatus] = {
"success": "success",
"blocked": "guardrail_intervened", # legacy
"guardrail_intervened": "guardrail_intervened", # direct
"failure": "guardrail_failed_to_respond", # legacy
"guardrail_failed_to_respond": "guardrail_failed_to_respond", # direct
"not_run": "not_run"
}
# Set LLM API status
llm_api_status: StandardLoggingPayloadStatus = status
#########################################################
# Map - guardrail_information.guardrail_status to guardrail_status
#########################################################
guardrail_status: GuardrailStatus = "not_run"
if guardrail_information and isinstance(guardrail_information, dict):
raw_status = guardrail_information.get("guardrail_status", "not_run")
guardrail_status = GUARDRAIL_STATUS_MAP.get(raw_status, "not_run")
return StandardLoggingPayloadStatusFields(
llm_api_status=llm_api_status,
guardrail_status=guardrail_status
)
def get_standard_logging_object_payload(
kwargs: Optional[dict],
init_response_obj: Union[Any, BaseModel, dict],
@ -4535,7 +4601,6 @@ def get_standard_logging_object_payload(
start_time=start_time,
response_id=id,
)
_request_body = proxy_server_request.get("body", {})
end_user_id = clean_metadata["user_api_key_end_user_id"] or _request_body.get(
"user", None
@ -4591,6 +4656,11 @@ def get_standard_logging_object_payload(
cache_hit=cache_hit,
stream=stream,
status=status,
status_fields=_get_status_fields(
status=status,
guardrail_information=metadata.get("standard_logging_guardrail_information", None),
error_str=error_str
),
custom_llm_provider=cast(Optional[str], kwargs.get("custom_llm_provider")),
saved_cache_cost=saved_cache_cost,
startTime=start_time_float,

4
litellm/litellm_core_utils/model_response_utils.py
View file

@ -84,7 +84,9 @@ def _has_meaningful_content(value: Any) -> bool:
return False
if isinstance(value, str):
return len(value.strip()) > 0
# Don't strip whitespace - preserve all content including newlines, spaces, etc.
# Even pure whitespace characters like '\n' or ' ' are meaningful content
return len(value) > 0
if isinstance(value, (list, dict)):
return len(value) > 0

11
litellm/litellm_core_utils/sensitive_data_masker.py
View file

@ -21,6 +21,8 @@ class SensitiveDataMasker:
"access",
"private",
"certificate",
"fingerprint",
"tenancy",
}
self.visible_prefix = visible_prefix
@ -42,7 +44,14 @@ class SensitiveDataMasker:
def is_sensitive_key(self, key: str) -> bool:
key_lower = str(key).lower()
result = any(pattern in key_lower for pattern in self.sensitive_patterns)
# Split on underscores and check if any segment matches the pattern
# This avoids false positives like "max_tokens" matching "token"
# but still catches "api_key", "access_token", etc.
key_segments = key_lower.replace('-', '_').split('_')
result = any(
pattern in key_segments
for pattern in self.sensitive_patterns
)
return result
def mask_dict(

8
litellm/llms/azure/azure.py
View file

@ -1117,6 +1117,14 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
status_code=422, message="max retries must be an int"
)
if api_key is None and azure_ad_token_provider is not None:
azure_ad_token = azure_ad_token_provider()
if azure_ad_token:
headers.pop(
"api-key", None
)
headers["Authorization"] = f"Bearer {azure_ad_token}"
# init AzureOpenAI Client
azure_client_params: Dict[str, Any] = self.initialize_azure_sdk_client(
litellm_params=litellm_params or {},

6
litellm/llms/base_llm/rerank/transformation.py
View file

@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
import httpx
from litellm.types.rerank import OptionalRerankParams, RerankBilledUnits, RerankResponse
from litellm.types.rerank import RerankBilledUnits, RerankResponse
from litellm.types.utils import ModelInfo
from ..chat.transformation import BaseLLMException
@ -30,7 +30,7 @@ class BaseRerankConfig(ABC):
def transform_rerank_request(
self,
model: str,
optional_rerank_params: OptionalRerankParams,
optional_rerank_params: Dict,
headers: dict,
) -> dict:
return {}
@ -78,7 +78,7 @@ class BaseRerankConfig(ABC):
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
) -> OptionalRerankParams:
) -> Dict:
pass
def get_error_class(

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

@ -440,27 +440,30 @@ class BedrockModelInfo(BaseLLMModelInfo):
"""
Abbreviations of regions AWS Bedrock supports for cross region inference
"""
return ["us", "eu", "apac"]
return ["us", "eu", "apac", "jp"]
@staticmethod
def get_bedrock_route(
model: str,
) -> Literal["converse", "invoke", "converse_like", "agent"]:
) -> Literal["converse", "invoke", "converse_like", "agent", "async_invoke"]:
"""
Get the bedrock route for the given model.
"""
route_mappings: Dict[str, Literal["invoke", "converse_like", "converse", "agent"]] = {
route_mappings: Dict[
str, Literal["invoke", "converse_like", "converse", "agent", "async_invoke"]
] = {
"invoke/": "invoke",
"converse_like/": "converse_like",
"converse_like/": "converse_like",
"converse/": "converse",
"agent/": "agent"
"agent/": "agent",
"async_invoke/": "async_invoke",
}
# Check explicit routes first
for prefix, route_type in route_mappings.items():
if prefix in model:
return route_type
base_model = BedrockModelInfo.get_base_model(model)
alt_model = BedrockModelInfo.get_non_litellm_routing_model_name(model=model)
if (
@ -469,38 +472,46 @@ class BedrockModelInfo(BaseLLMModelInfo):
):
return "converse"
return "invoke"
@staticmethod
def _explicit_converse_route(model: str) -> bool:
"""
Check if the model is an explicit converse route.
"""
return "converse/" in model
@staticmethod
def _explicit_invoke_route(model: str) -> bool:
"""
Check if the model is an explicit invoke route.
"""
return "invoke/" in model
@staticmethod
def _explicit_agent_route(model: str) -> bool:
"""
Check if the model is an explicit agent route.
"""
return "agent/" in model
@staticmethod
def _explicit_converse_like_route(model: str) -> bool:
"""
Check if the model is an explicit converse like route.
"""
return "converse_like/" in model
@staticmethod
def get_bedrock_provider_config_for_messages_api(model: str) -> Optional[BaseAnthropicMessagesConfig]:
def _explicit_async_invoke_route(model: str) -> bool:
"""
Check if the model is an explicit async invoke route.
"""
return "async_invoke/" in model
@staticmethod
def get_bedrock_provider_config_for_messages_api(
model: str,
) -> Optional[BaseAnthropicMessagesConfig]:
"""
Get the bedrock provider config for the given model.
@ -513,19 +524,20 @@ class BedrockModelInfo(BaseLLMModelInfo):
# Converse routes should go through litellm.completion()
if BedrockModelInfo._explicit_converse_route(model):
return None
#########################################################
# This goes through litellm.AmazonAnthropicClaude3MessagesConfig()
# Since bedrock Invoke supports Native Anthropic Messages API
#########################################################
if "claude" in model:
return litellm.AmazonAnthropicClaudeMessagesConfig()
#########################################################
# These routes will go through litellm.completion()
#########################################################
return None
class BedrockEventStreamDecoderBase:
"""
Base class for event stream decoding for Bedrock
@ -595,20 +607,20 @@ def get_anthropic_beta_from_headers(headers: dict) -> List[str]:
"""
Extract anthropic-beta header values and convert them to a list.
Supports comma-separated values from user headers.
Used by both converse and invoke transformations for consistent handling
of anthropic-beta headers that should be passed to AWS Bedrock.
Args:
headers (dict): Request headers dictionary
Returns:
List[str]: List of anthropic beta feature strings, empty list if no header
"""
anthropic_beta_header = headers.get("anthropic-beta")
if not anthropic_beta_header:
return []
# Split comma-separated values and strip whitespace
return [beta.strip() for beta in anthropic_beta_header.split(",")]
@ -618,19 +630,20 @@ class CommonBatchFilesUtils:
Common utilities for Bedrock batch and file operations.
Provides shared functionality to reduce code duplication between batches and files.
"""
def __init__(self):
# Import here to avoid circular imports
from .base_aws_llm import BaseAWSLLM
self._base_aws = BaseAWSLLM()
def get_bedrock_model_id_from_litellm_model(self, model: str) -> str:
"""
Extract the actual Bedrock model ID from LiteLLM model name.
Args:
model: LiteLLM model name (e.g., "bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
Returns:
Bedrock model ID (e.g., "anthropic.claude-3-sonnet-20240229-v1:0")
"""
@ -641,41 +654,45 @@ class CommonBatchFilesUtils:
def parse_s3_uri(self, s3_uri: str) -> tuple:
"""
Parse S3 URI into bucket and key components.
Args:
s3_uri: S3 URI (e.g., "s3://bucket/key/path")
Returns:
Tuple of (bucket, key)
Raises:
ValueError: If URI format is invalid
"""
if not s3_uri.startswith("s3://"):
raise ValueError(f"Invalid S3 URI format: {s3_uri}")
s3_parts = s3_uri[5:].split("/", 1) # Remove "s3://" and split on first "/"
if len(s3_parts) != 2:
raise ValueError(f"Invalid S3 URI format: {s3_uri}")
return s3_parts[0], s3_parts[1] # bucket, key
def extract_model_from_s3_file_path(self, s3_uri: str, optional_params: dict) -> str:
def extract_model_from_s3_file_path(
self, s3_uri: str, optional_params: dict
) -> str:
"""
Extract model ID from S3 file path.
The Bedrock file transformation creates S3 objects with the model name embedded:
Format: s3://bucket/litellm-bedrock-files-{model}-{uuid}.jsonl
"""
# Check if model is provided in optional_params first
if "model" in optional_params and optional_params["model"]:
return self.get_bedrock_model_id_from_litellm_model(optional_params["model"])
return self.get_bedrock_model_id_from_litellm_model(
optional_params["model"]
)
# Extract model from S3 URI path
# Expected format: s3://bucket/litellm-bedrock-files-{model}-{uuid}.jsonl
try:
bucket, object_key = self.parse_s3_uri(s3_uri)
# Extract model from object key if it follows our naming pattern
if object_key.startswith("litellm-bedrock-files-"):
# Remove prefix and suffix to get model part
@ -690,7 +707,7 @@ class CommonBatchFilesUtils:
return model_name
except Exception:
pass
# Fallback to default model
return "anthropic.claude-3-5-sonnet-20240620-v1:0"
@ -704,14 +721,14 @@ class CommonBatchFilesUtils:
) -> tuple:
"""
Sign AWS request using Signature Version 4.
Args:
service_name: AWS service name ("bedrock" or "s3")
data: Request data (string or dict)
endpoint_url: Full endpoint URL
optional_params: Optional parameters containing AWS credentials
method: HTTP method (default: POST)
Returns:
Tuple of (signed_headers, signed_data)
"""
@ -736,7 +753,7 @@ class CommonBatchFilesUtils:
aws_web_identity_token=optional_params.get("aws_web_identity_token"),
aws_sts_endpoint=optional_params.get("aws_sts_endpoint"),
)
# Prepare the request data
method_upper = method.upper()
if method_upper == "GET":
@ -746,12 +763,13 @@ class CommonBatchFilesUtils:
else:
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(
@ -759,45 +777,51 @@ class CommonBatchFilesUtils:
)
sigv4.add_auth(request)
prepped = request.prepare()
return dict(prepped.headers), request_data.encode('utf-8') if isinstance(request_data, str) else request_data
return (
dict(prepped.headers),
request_data.encode("utf-8")
if isinstance(request_data, str)
else request_data,
)
def generate_unique_job_name(self, model: str, prefix: str = "litellm") -> str:
"""
Generate a unique job name for AWS services.
AWS services often have length limits, so this creates a concise name.
Args:
model: Model name to include in the job name
prefix: Prefix for the job name
Returns:
Unique job name (≤ 63 characters for Bedrock compatibility)
"""
from litellm._uuid import uuid
unique_id = str(uuid.uuid4())[:8]
# Format: {prefix}-batch-{model}-{uuid}
# Example: litellm-batch-claude-266c398e
job_name = f"{prefix}-batch-{unique_id}"
return job_name
def get_s3_bucket_and_key_from_config(
self,
litellm_params: dict,
self,
litellm_params: dict,
optional_params: dict,
bucket_env_var: str = "AWS_S3_BUCKET_NAME",
key_prefix: str = "litellm"
key_prefix: str = "litellm",
) -> tuple:
"""
Get S3 bucket and generate a unique key from configuration.
Args:
litellm_params: LiteLLM parameters
optional_params: Optional parameters
bucket_env_var: Environment variable name for bucket
key_prefix: Prefix for the S3 key
Returns:
Tuple of (bucket_name, object_key)
"""
@ -806,18 +830,20 @@ class CommonBatchFilesUtils:
# Get bucket name
bucket_name = (
litellm_params.get("s3_bucket_name")
litellm_params.get("s3_bucket_name")
or optional_params.get("s3_bucket_name")
or os.getenv(bucket_env_var)
)
if not bucket_name:
raise ValueError(f"S3 bucket name is required. Set 's3_bucket_name' parameter or {bucket_env_var} env var")
raise ValueError(
f"S3 bucket name is required. Set 's3_bucket_name' parameter or {bucket_env_var} env var"
)
# Generate unique object key
timestamp = int(time.time())
unique_id = str(uuid.uuid4())[:8]
object_key = f"{key_prefix}-{timestamp}-{unique_id}"
return bucket_name, object_key
def get_error_class(
@ -827,7 +853,5 @@ class CommonBatchFilesUtils:
Get Bedrock-specific error class.
"""
return BedrockError(
status_code=status_code,
message=error_message,
headers=headers
status_code=status_code, message=error_message, headers=headers
)

225
litellm/llms/bedrock/embed/embedding.py
View file

@ -22,9 +22,8 @@ from litellm.secret_managers.main import get_secret
from litellm.types.llms.bedrock import (
AmazonEmbeddingRequest,
CohereEmbeddingRequest,
TwelveLabsMarengoEmbeddingRequest,
)
from litellm.types.utils import EmbeddingResponse
from litellm.types.utils import EmbeddingResponse, LlmProviders
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockError
@ -77,7 +76,7 @@ class BedrockEmbedding(BaseAWSLLM):
if aws_region_name is None:
aws_region_name = "us-west-2"
credentials: Credentials = self.get_credentials(
credentials: Credentials = self.get_credentials( # type: ignore
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
@ -151,35 +150,80 @@ 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
self,
response_list: List[dict],
model: str,
provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
is_async_invoke: Optional[bool] = False,
) -> 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(
# Handle async invoke responses (single response with invocationArn)
if (
is_async_invoke
and len(response_list) == 1
and "invocationArn" in response_list[0]
):
if provider == "twelvelabs":
returned_response = (
TwelveLabsMarengoEmbeddingConfig()._transform_async_invoke_response(
response=response_list[0], model=model
)
)
else:
# For other providers, create a generic async response
invocation_arn = response_list[0].get("invocationArn", "")
from litellm.types.utils import Embedding, Usage
embedding = Embedding(
embedding=[],
index=0,
object="embedding", # Must be literal "embedding"
)
usage = Usage(prompt_tokens=0, total_tokens=0)
# Create hidden params with job ID
from litellm.types.llms.base import HiddenParams
hidden_params = HiddenParams()
setattr(hidden_params, "_invocation_arn", invocation_arn)
returned_response = EmbeddingResponse(
data=[embedding],
model=model,
usage=usage,
hidden_params=hidden_params,
)
else:
# Handle regular invoke responses
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-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
)
##########################################################
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:
@ -203,6 +247,7 @@ class BedrockEmbedding(BaseAWSLLM):
logging_obj: Any,
provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
api_key: Optional[str] = None,
is_async_invoke: Optional[bool] = False,
):
responses: List[dict] = []
for data in batch_data:
@ -210,7 +255,7 @@ class BedrockEmbedding(BaseAWSLLM):
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
prepped = self.get_request_headers(
prepped = self.get_request_headers( # type: ignore # type: ignore
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
@ -249,7 +294,10 @@ class BedrockEmbedding(BaseAWSLLM):
responses.append(response)
return self._transform_response(
response_list=responses, model=model, provider=provider
response_list=responses,
model=model,
provider=provider,
is_async_invoke=is_async_invoke,
)
async def _async_single_func_embeddings(
@ -265,6 +313,7 @@ class BedrockEmbedding(BaseAWSLLM):
logging_obj: Any,
provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
api_key: Optional[str] = None,
is_async_invoke: Optional[bool] = False,
):
responses: List[dict] = []
for data in batch_data:
@ -272,7 +321,7 @@ class BedrockEmbedding(BaseAWSLLM):
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
prepped = self.get_request_headers(
prepped = self.get_request_headers( # type: ignore # type: ignore
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
@ -311,7 +360,10 @@ class BedrockEmbedding(BaseAWSLLM):
responses.append(response)
## TRANSFORM RESPONSE ##
return self._transform_response(
response_list=responses, model=model, provider=provider
response_list=responses,
model=model,
provider=provider,
is_async_invoke=is_async_invoke,
)
def embeddings(
@ -343,7 +395,10 @@ class BedrockEmbedding(BaseAWSLLM):
model=model,
model_id=unencoded_model_id,
)
# Check async invoke needs to be used
has_async_invoke = "async_invoke/" in model
if has_async_invoke:
model = model.replace("async_invoke/", "", 1)
provider = self.get_bedrock_embedding_provider(model)
if provider is None:
raise Exception(
@ -402,10 +457,14 @@ class BedrockEmbedding(BaseAWSLLM):
elif provider == "twelvelabs":
batch_data = []
for i in input:
twelvelabs_request: (
TwelveLabsMarengoEmbeddingRequest
) = TwelveLabsMarengoEmbeddingConfig()._transform_request(
input=i, inference_params=inference_params
twelvelabs_request = (
TwelveLabsMarengoEmbeddingConfig()._transform_request(
input=i,
inference_params=inference_params,
async_invoke_route=has_async_invoke,
model_id=modelId,
output_s3_uri=inference_params.get("output_s3_uri"),
)
)
batch_data.append(twelvelabs_request)
@ -417,7 +476,10 @@ class BedrockEmbedding(BaseAWSLLM):
),
aws_region_name=aws_region_name,
)
endpoint_url = f"{endpoint_url}/model/{modelId}/invoke"
if has_async_invoke:
endpoint_url = f"{endpoint_url}/async-invoke"
else:
endpoint_url = f"{endpoint_url}/model/{modelId}/invoke"
if batch_data is not None:
if aembedding:
@ -437,6 +499,7 @@ class BedrockEmbedding(BaseAWSLLM):
logging_obj=logging_obj,
api_key=api_key,
provider=provider,
is_async_invoke=has_async_invoke,
)
returned_response = self._single_func_embeddings(
client=(
@ -454,6 +517,7 @@ class BedrockEmbedding(BaseAWSLLM):
logging_obj=logging_obj,
api_key=api_key,
provider=provider,
is_async_invoke=has_async_invoke,
)
if returned_response is None:
raise Exception("Unable to map Bedrock request to provider")
@ -465,7 +529,7 @@ class BedrockEmbedding(BaseAWSLLM):
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
prepped = self.get_request_headers(
prepped = self.get_request_headers( # type: ignore
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
@ -491,3 +555,94 @@ class BedrockEmbedding(BaseAWSLLM):
client=client,
headers=prepped.headers, # type: ignore
)
async def _get_async_invoke_status(
self, invocation_arn: str, aws_region_name: str, logging_obj=None, **kwargs
) -> dict:
"""
Get the status of an async invoke job using the GetAsyncInvoke operation.
Args:
invocation_arn: The invocation ARN from the async invoke response
aws_region_name: AWS region name
**kwargs: Additional parameters (credentials, etc.)
Returns:
dict: Status response from AWS Bedrock
"""
# Get AWS credentials using the same method as other Bedrock methods
credentials, _ = self._load_credentials(kwargs)
# Get the runtime endpoint
endpoint_url, _ = self.get_runtime_endpoint(
api_base=None,
aws_bedrock_runtime_endpoint=kwargs.get("aws_bedrock_runtime_endpoint"),
aws_region_name=aws_region_name,
)
# Construct the status check URL
status_url = f"{endpoint_url}/async-invoke/{invocation_arn}"
# Prepare headers
headers = {"Content-Type": "application/json"}
# Get AWS signed headers
prepped = self.get_request_headers( # type: ignore
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=None,
endpoint_url=status_url,
data="", # GET request, no body
headers=headers,
api_key=None,
)
# LOGGING
if logging_obj is not None:
# Create custom curl command for GET request
masked_headers = logging_obj._get_masked_headers(prepped.headers)
formatted_headers = " ".join(
[f"-H '{k}: {v}'" for k, v in masked_headers.items()]
)
custom_curl = "\n\nGET Request Sent from LiteLLM:\n"
custom_curl += "curl -X GET \\\n"
custom_curl += f"{prepped.url} \\\n"
custom_curl += f"{formatted_headers}\n"
logging_obj.pre_call(
input=invocation_arn,
api_key="",
additional_args={
"complete_input_dict": {"invocation_arn": invocation_arn},
"api_base": prepped.url,
"headers": prepped.headers,
"request_str": custom_curl, # Override with custom GET curl command
},
)
# Make the GET request
client = get_async_httpx_client(llm_provider=LlmProviders.BEDROCK)
response = await client.get(
url=prepped.url,
headers=prepped.headers,
)
# LOGGING
if logging_obj is not None:
logging_obj.post_call(
input=invocation_arn,
api_key="",
original_response=response,
additional_args={
"complete_input_dict": {"invocation_arn": invocation_arn}
},
)
# Parse response
if response.status_code == 200:
return response.json()
else:
raise Exception(
f"Failed to get async invoke status: {response.status_code} - {response.text}"
)

207
litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py
View file

@ -1,33 +1,47 @@
"""
Transformation logic from OpenAI /v1/embeddings format to Bedrock TwelveLabs Marengo /invoke format.
Transformation logic from OpenAI /v1/embeddings format to Bedrock TwelveLabs Marengo /invoke and /async-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 typing import List, Optional, Union, cast
from litellm.types.llms.bedrock import (
TWELVELABS_EMBEDDING_INPUT_TYPES,
TwelveLabsAsyncInvokeRequest,
TwelveLabsMarengoEmbeddingRequest,
TwelveLabsOutputDataConfig,
TwelveLabsS3Location,
TwelveLabsS3OutputDataConfig,
)
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.
Supports text, image, video, and audio inputs.
- InvokeModel: text and image inputs
- StartAsyncInvoke: video, audio, image, and text inputs
"""
def __init__(self) -> None:
pass
def get_supported_openai_params(self) -> List[str]:
return ["encoding_format", "textTruncate", "embeddingOption"]
return [
"encoding_format",
"textTruncate",
"embeddingOption",
"startSec",
"lengthSec",
"useFixedLengthSec",
"minClipSec",
"input_type",
]
def map_openai_params(
self, non_default_params: dict, optional_params: dict
@ -41,45 +55,142 @@ class TwelveLabsMarengoEmbeddingConfig:
optional_params["textTruncate"] = v
elif k == "embeddingOption":
optional_params["embeddingOption"] = v
elif k == "input_type":
# Map input_type to inputType for Bedrock
optional_params["inputType"] = v
elif k in ["startSec", "lengthSec", "useFixedLengthSec", "minClipSec"]:
optional_params[k] = v
return optional_params
def _extract_bucket_owner_from_params(self, inference_params: dict) -> str:
"""
Extract bucket owner from inference parameters.
"""
return inference_params.get("bucketOwner", "")
def _is_s3_url(self, input: str) -> bool:
"""Check if input is an S3 URL."""
return input.startswith("s3://")
def _transform_request(
self, input: str, inference_params: dict
) -> TwelveLabsMarengoEmbeddingRequest:
self,
input: str,
inference_params: dict,
async_invoke_route: bool = False,
model_id: Optional[str] = None,
output_s3_uri: Optional[str] = None,
) -> Union[TwelveLabsMarengoEmbeddingRequest, TwelveLabsAsyncInvokeRequest]:
"""
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)
Transform OpenAI-style input to TwelveLabs Marengo format/async-invoke format.
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
Supports:
- Text inputs (for both invoke and async-invoke)
- Image inputs (for both invoke and async-invoke)
- Video inputs (async-invoke only)
- Audio inputs (async-invoke only)
- S3 URLs for all media types (async-invoke only)
"""
# Get input_type or default to "text"
input_type = cast(
TWELVELABS_EMBEDDING_INPUT_TYPES,
inference_params.get("inputType") or inference_params.get("input_type") or "text"
)
# Validate that async-invoke is used for video/audio
if input_type in ["video", "audio"] and not async_invoke_route:
raise ValueError(
f"Input type '{input_type}' requires async_invoke route. "
f"Use model format: 'bedrock/async_invoke/model_id'"
)
transformed_request: TwelveLabsMarengoEmbeddingRequest = {
"inputType": input_type
}
if input_type == "text":
transformed_request["inputText"] = input
# Set default textTruncate if not specified
if "textTruncate" not in inference_params:
transformed_request["textTruncate"] = "end"
elif input_type in ["image", "video", "audio"]:
if self._is_s3_url(input):
# S3 URL input
s3_location: TwelveLabsS3Location = {"uri": input}
bucket_owner = self._extract_bucket_owner_from_params(inference_params)
if bucket_owner:
s3_location["bucketOwner"] = bucket_owner
transformed_request["mediaSource"] = {"s3Location": s3_location}
else:
# Base64 encoded input
if input.startswith("data:"):
# Extract base64 data from data URL
b64_str = input.split(",", 1)[1] if "," in input else input
else:
# Direct base64 string
from litellm.utils import get_base64_str
b64_str = get_base64_str(input)
transformed_request["mediaSource"] = {"base64String": b64_str}
# Apply any additional inference parameters
for k, v in inference_params.items():
if k not in [
"inputType",
"input_type", # Exclude both camelCase and snake_case
"inputText",
"mediaSource",
"bucketOwner", # Don't include bucketOwner in the request
]: # Don't override core fields
transformed_request[k] = v # type: ignore
# If async invoke route, wrap in the async invoke format
if async_invoke_route and model_id:
return self._wrap_async_invoke_request(
model_input=transformed_request,
model_id=model_id,
output_s3_uri=output_s3_uri,
)
return transformed_request
def _wrap_async_invoke_request(
self,
model_input: TwelveLabsMarengoEmbeddingRequest,
model_id: str,
output_s3_uri: Optional[str] = None,
) -> TwelveLabsAsyncInvokeRequest:
"""
Wrap the transformed request in the correct AWS Bedrock async invoke format.
Args:
model_input: The transformed TwelveLabs Marengo embedding request
model_id: The model identifier (without async_invoke prefix)
output_s3_uri: Optional S3 URI for output data config
Returns:
TwelveLabsAsyncInvokeRequest: The wrapped async invoke request
"""
import urllib.parse
# Clean the model ID
unquoted_model_id = urllib.parse.unquote(model_id)
if unquoted_model_id.startswith("async_invoke/"):
unquoted_model_id = unquoted_model_id.replace("async_invoke/", "")
# Validate that the S3 URI is not empty
if not output_s3_uri or output_s3_uri.strip() == "":
raise ValueError("output_s3_uri cannot be empty for async invoke requests")
return TwelveLabsAsyncInvokeRequest(
modelId=unquoted_model_id,
modelInput=model_input,
outputDataConfig=TwelveLabsOutputDataConfig(
s3OutputDataConfig=TwelveLabsS3OutputDataConfig(s3Uri=output_s3_uri)
),
)
def _transform_response(
self, response_list: List[dict], model: str
) -> EmbeddingResponse:
@ -138,3 +249,53 @@ class TwelveLabsMarengoEmbeddingConfig:
usage = Usage(prompt_tokens=total_tokens, total_tokens=total_tokens)
return EmbeddingResponse(data=embeddings, model=model, usage=usage)
def _transform_async_invoke_response(
self, response: dict, model: str
) -> EmbeddingResponse:
"""
Transform async invoke response (invocation ARN) to OpenAI format.
AWS async invoke returns:
{
"invocationArn": "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123"
}
We transform this to a job-like embedding response:
{
"object": "list",
"data": [
{
"object": "embedding_job_id:1234567890",
"embedding": [],
"index": 0
}
],
"model": "model",
"usage": {}
}
"""
invocation_arn = response.get("invocationArn", "")
# Create a placeholder embedding object for the job
embedding = Embedding(
embedding=[], # Empty embedding for async jobs
index=0,
object="embedding",
)
# Create usage object (empty for async jobs)
usage = Usage(prompt_tokens=0, total_tokens=0)
# Create hidden params with job ID
from litellm.types.llms.base import HiddenParams
hidden_params = HiddenParams()
setattr(hidden_params, "_invocation_arn", invocation_arn)
return EmbeddingResponse(
data=[embedding],
model=model,
usage=usage,
hidden_params=hidden_params,
)

4
litellm/llms/cohere/common_utils.py
View file

@ -31,7 +31,7 @@ def validate_environment(
"Request-Source": "unspecified:litellm",
"accept": "application/json",
"content-type": "application/json",
"Authorization": "bearer $CO_API_KEY"
"Authorization": "Bearer $CO_API_KEY"
}
"""
headers.update(
@ -42,7 +42,7 @@ def validate_environment(
}
)
if api_key:
headers["Authorization"] = f"bearer {api_key}"
headers["Authorization"] = f"Bearer {api_key}"
return headers

12
litellm/llms/cohere/rerank/transformation.py
View file

@ -1,8 +1,8 @@
from typing import Any, Dict, List, Optional, Union
import httpx
import litellm
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
@ -52,20 +52,20 @@ class CohereRerankConfig(BaseRerankConfig):
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
) -> OptionalRerankParams:
) -> Dict:
"""
Map Cohere rerank params
No mapping required - returns all supported params
"""
return OptionalRerankParams(
return dict(OptionalRerankParams(
query=query,
documents=documents,
top_n=top_n,
rank_fields=rank_fields,
return_documents=return_documents,
max_chunks_per_doc=max_chunks_per_doc,
)
))
def validate_environment(
self,
@ -86,7 +86,7 @@ class CohereRerankConfig(BaseRerankConfig):
)
default_headers = {
"Authorization": f"bearer {api_key}",
"Authorization": f"Bearer {api_key}",
"accept": "application/json",
"content-type": "application/json",
}
@ -101,7 +101,7 @@ class CohereRerankConfig(BaseRerankConfig):
def transform_rerank_request(
self,
model: str,
optional_rerank_params: OptionalRerankParams,
optional_rerank_params: Dict,
headers: dict,
) -> dict:
if "query" not in optional_rerank_params:

8
litellm/llms/cohere/rerank_v2/transformation.py
View file

@ -44,25 +44,25 @@ class CohereRerankV2Config(CohereRerankConfig):
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
) -> OptionalRerankParams:
) -> Dict:
"""
Map Cohere rerank params
No mapping required - returns all supported params
"""
return OptionalRerankParams(
return dict(OptionalRerankParams(
query=query,
documents=documents,
top_n=top_n,
rank_fields=rank_fields,
return_documents=return_documents,
max_tokens_per_doc=max_tokens_per_doc,
)
))
def transform_rerank_request(
self,
model: str,
optional_rerank_params: OptionalRerankParams,
optional_rerank_params: Dict,
headers: dict,
) -> dict:
if "query" not in optional_rerank_params:

4
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -65,7 +65,7 @@ from litellm.types.llms.openai import (
ResponseInputParam,
ResponsesAPIResponse,
)
from litellm.types.rerank import OptionalRerankParams, RerankResponse
from litellm.types.rerank import RerankResponse
from litellm.types.responses.main import DeleteResponseResult
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import (
@ -893,7 +893,7 @@ class BaseLLMHTTPHandler:
custom_llm_provider: str,
logging_obj: LiteLLMLoggingObj,
provider_config: BaseRerankConfig,
optional_rerank_params: OptionalRerankParams,
optional_rerank_params: Dict,
timeout: Optional[Union[float, httpx.Timeout]],
model_response: RerankResponse,
_is_async: bool = False,

6
litellm/llms/deepinfra/rerank/transformation.py
View file

@ -2,11 +2,11 @@
Translate between Cohere's `/rerank` format and Deepinfra's `/rerank` format.
"""
from litellm._uuid import uuid
from typing import Any, Dict, List, Optional, Union
import httpx
from litellm._uuid import uuid
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.rerank.transformation import (
BaseLLMException,
@ -98,7 +98,7 @@ class DeepinfraRerankConfig(BaseRerankConfig):
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
) -> OptionalRerankParams:
) -> Dict:
# Start with the basic parameters
optional_rerank_params = {}
if query:
@ -124,7 +124,7 @@ class DeepinfraRerankConfig(BaseRerankConfig):
def transform_rerank_request(
self,
model: str,
optional_rerank_params: OptionalRerankParams,
optional_rerank_params: Dict,
headers: dict,
) -> dict:
# Convert OptionalRerankParams to dict as expected by parent class

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

@ -3,10 +3,10 @@ This file contains the transformation logic for the Gemini realtime API.
"""
import json
from litellm._uuid import uuid
from typing import Any, Dict, List, Optional, Union, cast
from litellm import verbose_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
@ -186,9 +186,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
)
vertex_gemini_config = VertexGeminiConfig()
vertex_gemini_config._map_function(value)
optional_params["generationConfig"]["tools"] = (
vertex_gemini_config._map_function(value)
vertex_gemini_config._map_function(
value=value, optional_params=optional_params
)
)
elif key == "input_audio_transcription" and value is not None:
optional_params["inputAudioTranscription"] = {}

27
litellm/llms/hosted_vllm/rerank/transformation.py
View file

@ -2,27 +2,26 @@
Transformation logic for Hosted VLLM rerank
"""
from litellm._uuid import uuid
from typing import Any, Dict, List, Optional, Union
import httpx
from litellm._uuid import uuid
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.rerank import (
OptionalRerankParams,
RerankBilledUnits,
RerankRequest,
RerankResponse,
RerankResponseDocument,
RerankResponseMeta,
RerankResponseResult,
RerankTokens,
OptionalRerankParams,
RerankRequest,
)
import httpx
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.secret_managers.main import get_secret_str
class HostedVLLMRerankError(BaseLLMException):
def __init__(
@ -72,20 +71,20 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
) -> OptionalRerankParams:
) -> Dict:
"""
Map parameters for Hosted VLLM rerank
"""
if max_chunks_per_doc is not None:
raise ValueError("Hosted VLLM does not support max_chunks_per_doc")
return OptionalRerankParams(
return dict(OptionalRerankParams(
query=query,
documents=documents,
top_n=top_n,
rank_fields=rank_fields,
return_documents=return_documents,
)
))
def validate_environment(
self,
@ -112,7 +111,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
def transform_rerank_request(
self,
model: str,
optional_rerank_params: OptionalRerankParams,
optional_rerank_params: Dict,
headers: dict,
) -> dict:
if "query" not in optional_rerank_params:

4
litellm/llms/huggingface/rerank/transformation.py
View file

@ -1,11 +1,11 @@
import os
from litellm._uuid import uuid
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
import httpx
from typing_extensions import TypedDict
import litellm
from litellm._uuid import uuid
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.secret_managers.main import get_secret_str
@ -95,7 +95,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig):
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
) -> OptionalRerankParams:
) -> Dict:
optional_rerank_params = {}
if non_default_params is not None:
for k, v in non_default_params.items():

2
litellm/llms/infinity/rerank/transformation.py
View file

@ -49,7 +49,7 @@ class InfinityRerankConfig(CohereRerankConfig):
)
default_headers = {
"Authorization": f"bearer {api_key}",
"Authorization": f"Bearer {api_key}",
"accept": "application/json",
"content-type": "application/json",
}

29
litellm/llms/jina_ai/rerank/transformation.py
View file

@ -6,11 +6,11 @@ Why separate file? Make it easy to see how transformation works
Docs - https://jina.ai/reranker
"""
from litellm._uuid import uuid
from typing import Any, Dict, List, Optional, Tuple, Union
from httpx import URL, Response
from litellm._uuid import uuid
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.types.rerank import (
@ -45,15 +45,15 @@ class JinaAIRerankConfig(BaseRerankConfig):
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
) -> OptionalRerankParams:
) -> Dict:
optional_params = {}
supported_params = self.get_supported_cohere_rerank_params(model)
for k, v in non_default_params.items():
if k in supported_params:
optional_params[k] = v
return OptionalRerankParams(
return dict(OptionalRerankParams(
**optional_params,
)
))
def get_complete_url(self, api_base: Optional[str], model: str) -> str:
base_path = "/v1/rerank"
@ -67,7 +67,7 @@ class JinaAIRerankConfig(BaseRerankConfig):
return cleaned_base
def transform_rerank_request(
self, model: str, optional_rerank_params: OptionalRerankParams, headers: Dict
self, model: str, optional_rerank_params: Dict, headers: Dict
) -> Dict:
return {"model": model, **optional_rerank_params}
@ -98,9 +98,26 @@ class JinaAIRerankConfig(BaseRerankConfig):
if _results is None:
raise ValueError(f"No results found in the response={_json_response}")
# Transform Jina AI's response format to match LiteLLM's expected format
# Jina AI returns: {"index": 0, "relevance_score": 0.72, "document": "hello"}
# LiteLLM expects: {"index": 0, "relevance_score": 0.72, "document": {"text": "hello"}}
transformed_results = []
for result in _results:
transformed_result = {
"index": result["index"],
"relevance_score": result["relevance_score"],
}
# Convert document from string to dict format if it exists
if "document" in result and isinstance(result["document"], str):
transformed_result["document"] = {"text": result["document"]}
elif "document" in result:
# If it's already a dict, keep it as is
transformed_result["document"] = result["document"]
transformed_results.append(transformed_result)
return RerankResponse(
id=_json_response.get("id") or str(uuid.uuid4()),
results=_results, # type: ignore
results=transformed_results, # type: ignore
meta=rerank_meta,
) # Return response

325
litellm/llms/nvidia_nim/rerank/transformation.py Normal file
View file

@ -0,0 +1,325 @@
from typing import Any, Dict, List, Literal, Optional, Union
import httpx
from typing_extensions import Required, TypedDict
import litellm
from litellm._uuid import uuid
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.rerank import (
RerankBilledUnits,
RerankResponse,
RerankResponseMeta,
RerankResponseResult,
)
class NvidiaNimQueryObject(TypedDict):
text: Required[str]
class NvidiaNimPassageObject(TypedDict):
text: Required[str]
class NvidiaNimRerankRequest(TypedDict, total=False):
model: Required[str]
query: Required[NvidiaNimQueryObject]
passages: Required[List[NvidiaNimPassageObject]]
truncate: Literal["NONE", "END"]
top_k: int
class NvidiaNimRankingResult(TypedDict):
index: Required[int]
logit: Required[float]
class NvidiaNimRerankResponse(TypedDict):
rankings: Required[List[NvidiaNimRankingResult]]
class NvidiaNimRerankConfig(BaseRerankConfig):
"""
Reference: https://docs.api.nvidia.com/nim/reference/nvidia-llama-3_2-nv-rerankqa-1b-v2-infer
Nvidia NIM rerank API uses a different format:
- query is an object with 'text' field
- documents are called 'passages' and have 'text' field
"""
DEFAULT_NIM_RERANK_API_BASE = "https://ai.api.nvidia.com"
def __init__(self) -> None:
pass
def get_complete_url(self, api_base: Optional[str], model: str) -> str:
"""
Construct the Nvidia NIM rerank URL.
Format: {api_base}/v1/retrieval/{model}/reranking
If the user provides a full URL (e.g., {api_base}/v1/retrieval/{model}/reranking),
it will be used as-is.
"""
if not api_base:
api_base = self.DEFAULT_NIM_RERANK_API_BASE
api_base = api_base.rstrip("/")
# Check if user already provided the full URL with /retrieval/ path
if "/retrieval/" in api_base:
return api_base
# Ensure we don't have duplicate /v1
if api_base.endswith("/v1"):
api_base = api_base[:-3]
return f"{api_base}/v1/retrieval/{model}/reranking"
def get_supported_cohere_rerank_params(self, model: str) -> list:
"""
Nvidia NIM supports these rerank parameters.
"""
return [
"query",
"documents",
"top_n",
]
def map_cohere_rerank_params(
self,
non_default_params: Optional[dict],
model: str,
drop_params: bool,
query: str,
documents: List[Union[str, Dict[str, Any]]],
custom_llm_provider: Optional[str] = None,
top_n: Optional[int] = None,
rank_fields: Optional[List[str]] = None,
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
) -> Dict:
"""
Map Cohere/OpenAI rerank params to Nvidia NIM format.
Parameter mapping:
- top_n (Cohere) -> top_k (Nvidia)
Nvidia NIM specific params (passed through as-is from non_default_params):
- truncate: How to truncate input if too long (NONE, END)
"""
optional_nvidia_nim_rerank_params: Dict[str, Any] = {
"query": query,
"documents": documents,
}
# Map Cohere's top_n to Nvidia's top_k
if top_n is not None:
optional_nvidia_nim_rerank_params["top_k"] = top_n
# Pass through Nvidia-specific params from non_default_params
if non_default_params:
optional_nvidia_nim_rerank_params.update(non_default_params)
return dict(optional_nvidia_nim_rerank_params)
def validate_environment(
self,
headers: dict,
model: str,
api_key: Optional[str] = None,
) -> dict:
"""
Validate that the Nvidia NIM API key is present.
"""
if api_key is None:
api_key = (
get_secret_str("NVIDIA_NIM_API_KEY")
or litellm.api_key
)
if api_key is None:
raise ValueError(
"Nvidia NIM API key is required. Please set 'NVIDIA_NIM_API_KEY' in your environment"
)
default_headers = {
"Authorization": f"Bearer {api_key}",
"accept": "application/json",
"content-type": "application/json",
}
# If 'Authorization' is provided in headers, it overrides the default
if "Authorization" in headers:
default_headers["Authorization"] = headers["Authorization"]
# Merge other headers, overriding any default ones except Authorization
return {**default_headers, **headers}
def transform_rerank_request(
self,
model: str,
optional_rerank_params: Dict,
headers: dict,
) -> dict:
"""
Transform request to Nvidia NIM format.
Nvidia NIM expects:
- query as {text: "..."}
- documents as passages: [{text: "..."}, ...]
- Optional: truncate (NONE or END), top_k
Note: optional_rerank_params may contain provider-specific params like 'top_k' and 'truncate'
that aren't in the OptionalRerankParams TypedDict but are passed through at runtime.
The mapping from Cohere's 'top_n' to Nvidia's 'top_k' already happened in map_cohere_rerank_params.
"""
if "query" not in optional_rerank_params:
raise ValueError("query is required for Nvidia NIM rerank")
if "documents" not in optional_rerank_params:
raise ValueError("documents is required for Nvidia NIM rerank")
query = optional_rerank_params["query"]
documents = optional_rerank_params["documents"]
# Transform query to object format
query_obj: NvidiaNimQueryObject = {"text": query}
# Transform documents to passages format
passages: List[NvidiaNimPassageObject] = []
for doc in documents:
if isinstance(doc, str):
passages.append({"text": doc})
elif isinstance(doc, dict):
# If document is already a dict, check if it has 'text' field
if "text" in doc:
passages.append({"text": doc["text"]})
else:
# Otherwise, stringify the dict
import json
passages.append({"text": json.dumps(doc)})
else:
passages.append({"text": str(doc)})
# Note: URL path uses underscores (llama-3_2) but JSON body uses periods (llama-3.2)
# Convert underscores back to periods for the model field in request body
model_for_body = model.replace("_", ".")
# Build request using TypedDict
request_data: NvidiaNimRerankRequest = {
"model": model_for_body,
"query": query_obj,
"passages": passages,
}
# Add optional top_k parameter if provided (already mapped from top_n in map_cohere_rerank_params)
if "top_k" in optional_rerank_params and optional_rerank_params.get("top_k") is not None: # type: ignore
request_data["top_k"] = optional_rerank_params.get("top_k") # type: ignore
# Add Nvidia-specific truncate parameter if provided
# This is passed through from non_default_params, not in base OptionalRerankParams
if "truncate" in optional_rerank_params and optional_rerank_params.get("truncate") is not None: # type: ignore
truncate_value = optional_rerank_params.get("truncate") # type: ignore
if truncate_value in ["NONE", "END"]:
request_data["truncate"] = truncate_value # type: ignore
return dict(request_data)
def transform_rerank_response(
self,
model: str,
raw_response: httpx.Response,
model_response: RerankResponse,
logging_obj: LiteLLMLoggingObj,
api_key: Optional[str] = None,
request_data: dict = {},
optional_params: dict = {},
litellm_params: dict = {},
) -> RerankResponse:
"""
Transform Nvidia NIM rerank response to LiteLLM format.
Nvidia NIM returns (NvidiaNimRerankResponse):
{
"rankings": [
{
"index": 0,
"logit": 0.123
}
]
}
LiteLLM expects (RerankResponse):
{
"results": [
{
"index": 0,
"relevance_score": 0.123,
"document": {"text": "..."} # optional
}
]
}
"""
try:
raw_response_json = raw_response.json()
except Exception:
raise BaseLLMException(
status_code=raw_response.status_code,
message=raw_response.text,
headers=raw_response.headers,
)
# Parse as NvidiaNimRerankResponse
nvidia_response: NvidiaNimRerankResponse = raw_response_json
# Transform Nvidia NIM response to LiteLLM format
results: List[RerankResponseResult] = []
rankings = nvidia_response.get("rankings", [])
# Get original documents from request if we need to include them
original_passages: List[NvidiaNimPassageObject] = request_data.get("passages", [])
for ranking in rankings:
result_item: RerankResponseResult = {
"index": ranking["index"],
"relevance_score": ranking["logit"],
}
# Include document if it was in the original request
index: int = ranking["index"]
if index < len(original_passages):
result_item["document"] = {"text": original_passages[index]["text"]} # type: ignore
results.append(result_item)
# Construct metadata with billed_units
# Nvidia NIM uses "usage" field with "total_tokens"
usage = raw_response_json.get("usage", {})
total_tokens = usage.get("total_tokens", 0)
billed_units: RerankBilledUnits = {
"total_tokens": total_tokens if total_tokens > 0 else len(results)
}
meta: RerankResponseMeta = {
"billed_units": billed_units
}
return RerankResponse(
id=raw_response_json.get("id") or str(uuid.uuid4()),
results=results,
meta=meta,
)
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BaseLLMException:
return BaseLLMException(
status_code=status_code,
message=error_message,
headers=headers,
)

26
litellm/llms/oci/chat/transformation.py
View file

@ -207,9 +207,9 @@ class OCIChatConfig(BaseConfig):
alias = open_ai_to_oci_param_map.get(key)
if alias is False:
if drop_params:
# Workaround for mypy issue
if drop_params or litellm.drop_params:
continue
raise Exception(f"param `{key}` is not supported on OCI")
if alias is None:
@ -450,12 +450,26 @@ class OCIChatConfig(BaseConfig):
"Cohere models are not yet supported in the litellm OCI chat completion endpoint. Use the Cohere API directly."
)
else:
data = OCICompletionPayload(
compartmentId=oci_compartment_id,
servingMode=OCIServingMode(
oci_serving_mode = optional_params.get("oci_serving_mode", "ON_DEMAND")
if oci_serving_mode not in ["ON_DEMAND", "DEDICATED"]:
raise Exception(
"kwarg `oci_serving_mode` must be either 'ON_DEMAND' or 'DEDICATED'"
)
if oci_serving_mode == "DEDICATED":
servingMode = OCIServingMode(
servingType="DEDICATED",
endpointId=model,
)
else:
servingMode = OCIServingMode(
servingType="ON_DEMAND",
modelId=model,
),
)
data = OCICompletionPayload(
compartmentId=oci_compartment_id,
servingMode=servingMode,
chatRequest=OCIChatRequestPayload(
apiFormat=vendor.value,
messages=adapt_messages_to_generic_oci_standard(messages),

14
litellm/llms/openai/chat/gpt_5_transformation.py
View file

@ -8,18 +8,24 @@ from .gpt_transformation import OpenAIGPTConfig
class OpenAIGPT5Config(OpenAIGPTConfig):
"""Configuration for gpt-5 models.
"""Configuration for gpt-5 models including GPT-5-Codex variants.
Handles OpenAI API quirks for the gpt-5 series like:
- Mapping ``max_tokens`` -> ``max_completion_tokens``.
- Dropping unsupported ``temperature`` values when requested.
- Support for GPT-5-Codex models optimized for code generation.
"""
@classmethod
def is_model_gpt_5_model(cls, model: str) -> bool:
return "gpt-5" in model
@classmethod
def is_model_gpt_5_codex_model(cls, model: str) -> bool:
"""Check if the model is specifically a GPT-5 Codex variant."""
return "gpt-5-codex" in model
def get_supported_openai_params(self, model: str) -> list:
from litellm.utils import supports_tool_choice
@ -38,7 +44,9 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
]
return [
param for param in base_gpt_series_params if param not in non_supported_params
param
for param in base_gpt_series_params
if param not in non_supported_params
]
def map_openai_params(
@ -67,7 +75,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
else:
raise litellm.utils.UnsupportedParamsError(
message=(
"gpt-5 models don't support temperature={}. Only temperature=1 is supported. To drop unsupported params set `litellm.drop_params = True`"
"gpt-5 models (including gpt-5-codex) don't support temperature={}. Only temperature=1 is supported. To drop unsupported params set `litellm.drop_params = True`"
).format(temperature_value),
status_code=400,
)

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

@ -68,6 +68,7 @@ from litellm.types.llms.vertex_ai import (
ToolConfig,
Tools,
UsageMetadata,
VertexToolName,
)
from litellm.types.utils import (
ChatCompletionAudioResponse,
@ -276,42 +277,106 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"""
return Tools(googleSearch={})
def _map_function(self, value: List[dict]) -> List[Tools]: # noqa: PLR0915
def _extract_google_maps_retrieval_config(
self, google_maps_config: dict
) -> Tuple[dict, Optional[dict]]:
"""
Extract location configuration from googleMaps tool for Vertex AI toolConfig.
Supports two interface styles:
1. Nested (recommended): {"enableWidget": "...", "retrievalConfig": {"latitude": ..., "longitude": ...}}
2. Flat (backward compat): {"enableWidget": "...", "latitude": ..., "longitude": ...}
Args:
google_maps_config: The googleMaps tool configuration from LiteLLM
Returns:
Tuple of (cleaned_google_maps_config, retrieval_config):
- cleaned_google_maps_config: googleMaps config without location fields
- retrieval_config: Location config for toolConfig.retrievalConfig or None
"""
retrieval_config = None
latitude = google_maps_config.get("latitude")
longitude = google_maps_config.get("longitude")
language_code = google_maps_config.get("languageCode")
if latitude is not None and longitude is not None:
retrieval_config = {
"latLng": {
"latitude": latitude,
"longitude": longitude,
}
}
if language_code is not None:
retrieval_config["languageCode"] = language_code
# Remove location fields from tool definition
cleaned_config = {
k: v
for k, v in google_maps_config.items()
if k not in ["latitude", "longitude", "languageCode"]
}
return cleaned_config, retrieval_config
def get_tool_value(
self,
tool: dict,
tool_name: str
) -> Optional[dict]:
"""
Helper function to get tool value handling both camelCase and underscore_case variants
Args:
tool (dict): The tool dictionary
tool_name (str): The base tool name (e.g. "codeExecution")
Returns:
Optional[dict]: The tool value if found, None otherwise
"""
# Convert camelCase to underscore_case
underscore_name = "".join(
["_" + c.lower() if c.isupper() else c for c in tool_name]
).lstrip("_")
# Try both camelCase and underscore_case variants
if tool.get(tool_name) is not None:
return tool.get(tool_name)
elif tool.get(underscore_name) is not None:
return tool.get(underscore_name)
else:
return None
def _map_function( # noqa: PLR0915
self, value: List[dict], optional_params: dict
) -> List[Tools]:
"""
Map OpenAI-style tools/functions to Vertex AI format.
Args:
value: List of tool definitions
optional_params: Request-scoped parameters to store retrieval config
Returns:
List of mapped tools in Vertex AI format
Side effects:
May add 'toolConfig' with 'retrievalConfig' to optional_params if
googleMaps tools contain location data
"""
gtool_func_declarations = []
googleSearch: Optional[dict] = None
googleSearchRetrieval: Optional[dict] = None
enterpriseWebSearch: Optional[dict] = None
urlContext: Optional[dict] = None
code_execution: Optional[dict] = None
googleMaps: Optional[dict] = None
google_maps_retrieval_config: Optional[dict] = None
# remove 'additionalProperties' from tools
value = _remove_additional_properties(value)
# remove 'strict' from tools
value = _remove_strict_from_schema(value)
def get_tool_value(tool: dict, tool_name: str) -> Optional[dict]:
"""
Helper function to get tool value handling both camelCase and underscore_case variants
Args:
tool (dict): The tool dictionary
tool_name (str): The base tool name (e.g. "codeExecution")
Returns:
Optional[dict]: The tool value if found, None otherwise
"""
# Convert camelCase to underscore_case
underscore_name = "".join(
["_" + c.lower() if c.isupper() else c for c in tool_name]
).lstrip("_")
# Try both camelCase and underscore_case variants
if tool.get(tool_name) is not None:
return tool.get(tool_name)
elif tool.get(underscore_name) is not None:
return tool.get(underscore_name)
else:
return None
for tool in value:
openai_function_object: Optional[
ChatCompletionToolParamFunctionChunk
@ -341,17 +406,27 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
tool_name = list(tool.keys())[0] if len(tool.keys()) == 1 else None
if tool_name and (
tool_name == "codeExecution" or tool_name == "code_execution"
tool_name == "codeExecution" or tool_name == VertexToolName.CODE_EXECUTION.value
): # code_execution maintained for backwards compatibility
code_execution = get_tool_value(tool, "codeExecution")
elif tool_name and tool_name == "googleSearch":
googleSearch = get_tool_value(tool, "googleSearch")
elif tool_name and tool_name == "googleSearchRetrieval":
googleSearchRetrieval = get_tool_value(tool, "googleSearchRetrieval")
elif tool_name and tool_name == "enterpriseWebSearch":
enterpriseWebSearch = get_tool_value(tool, "enterpriseWebSearch")
elif tool_name and tool_name == "urlContext":
urlContext = get_tool_value(tool, "urlContext")
code_execution = self.get_tool_value(tool, "codeExecution")
elif tool_name and tool_name == VertexToolName.GOOGLE_SEARCH.value:
googleSearch = self.get_tool_value(tool, VertexToolName.GOOGLE_SEARCH.value)
elif tool_name and tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value:
googleSearchRetrieval = self.get_tool_value(tool, VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value)
elif tool_name and tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value:
enterpriseWebSearch = self.get_tool_value(tool, VertexToolName.ENTERPRISE_WEB_SEARCH.value)
elif tool_name and (tool_name == VertexToolName.URL_CONTEXT.value or tool_name == "urlContext"):
urlContext = self.get_tool_value(tool, tool_name)
elif tool_name and (
tool_name == VertexToolName.GOOGLE_MAPS.value or tool_name == "google_maps"
):
google_maps_value = self.get_tool_value(tool, VertexToolName.GOOGLE_MAPS.value)
# Extract and transform location configuration for toolConfig
if google_maps_value is not None:
googleMaps, google_maps_retrieval_config = self._extract_google_maps_retrieval_config(
google_maps_config=google_maps_value
)
elif openai_function_object is not None:
gtool_func_declaration = FunctionDeclaration(
name=openai_function_object["name"],
@ -373,19 +448,29 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"Invalid tool={}. Use `litellm.set_verbose` or `litellm --detailed_debug` to see raw request."
)
_tools = Tools(
function_declarations=gtool_func_declarations,
)
# Only include function_declarations if there are actual functions
_tools = Tools()
if gtool_func_declarations:
_tools["function_declarations"] = gtool_func_declarations
if googleSearch is not None:
_tools["googleSearch"] = googleSearch
_tools[VertexToolName.GOOGLE_SEARCH.value] = googleSearch
if googleSearchRetrieval is not None:
_tools["googleSearchRetrieval"] = googleSearchRetrieval
_tools[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = googleSearchRetrieval
if enterpriseWebSearch is not None:
_tools["enterpriseWebSearch"] = enterpriseWebSearch
_tools[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = enterpriseWebSearch
if code_execution is not None:
_tools["code_execution"] = code_execution
_tools[VertexToolName.CODE_EXECUTION.value] = code_execution
if urlContext is not None:
_tools["url_context"] = urlContext
_tools[VertexToolName.URL_CONTEXT.value] = urlContext
if googleMaps is not None:
_tools[VertexToolName.GOOGLE_MAPS.value] = googleMaps
# Add retrieval config to toolConfig if googleMaps has location data
if google_maps_retrieval_config is not None:
if "toolConfig" not in optional_params:
optional_params["toolConfig"] = {}
optional_params["toolConfig"]["retrievalConfig"] = google_maps_retrieval_config
return [_tools]
def _map_response_schema(self, value: dict) -> dict:
@ -606,8 +691,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
and isinstance(value, list)
and value
):
# Pass optional_params so _map_function can add toolConfig if needed
mapped_tools = self._map_function(
value=value, optional_params=optional_params
)
optional_params = self._add_tools_to_optional_params(
optional_params, self._map_function(value=value)
optional_params, mapped_tools
)
elif param == "tool_choice" and (
isinstance(value, str) or isinstance(value, dict)

16
litellm/main.py
View file

@ -24,6 +24,7 @@ from functools import partial
from typing import (
TYPE_CHECKING,
Any,
AsyncIterator,
Callable,
Coroutine,
Dict,
@ -5170,6 +5171,21 @@ async def aadapter_completion(
except Exception as e:
raise e
async def aadapter_generate_content(
**kwargs,
) -> Union[Dict[str, Any], AsyncIterator[bytes]]:
from litellm.google_genai.adapters.handler import (
GenerateContentToCompletionHandler,
)
coro = cast(
Coroutine[Any, Any, Union[Dict[str, Any], AsyncIterator[bytes]]],
GenerateContentToCompletionHandler.generate_content_handler(
**kwargs, _is_async=True
),
)
return await coro
def adapter_completion(
*, adapter_id: str, **kwargs

773
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

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

@ -294,6 +294,9 @@ class MCPRequestHandler:
) -> List[str]:
"""
Get list of allowed MCP servers for the given user/key based on permissions
Returns:
List[str]: List of allowed MCP servers by server id
"""
from typing import List
@ -330,11 +333,30 @@ class MCPRequestHandler:
verbose_logger.warning(f"Failed to get allowed MCP servers: {str(e)}")
return []
@staticmethod
def is_tool_allowed(
allowed_mcp_servers: List[str],
server_name: str,
) -> bool:
"""
Check if the tool is allowed for the given user/key based on permissions
"""
if len(allowed_mcp_servers) == 0:
return True
elif server_name in allowed_mcp_servers:
return True
return False
@staticmethod
async def _get_allowed_mcp_servers_for_key(
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
) -> List[str]:
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.auth.auth_checks import get_object_permission
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if user_api_key_auth is None:
return []
@ -347,12 +369,12 @@ class MCPRequestHandler:
return []
try:
key_object_permission = (
await prisma_client.db.litellm_objectpermissiontable.find_unique(
where={
"object_permission_id": user_api_key_auth.object_permission_id
},
)
key_object_permission = await get_object_permission(
object_permission_id=user_api_key_auth.object_permission_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_auth.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if key_object_permission is None:
return []
@ -386,7 +408,12 @@ class MCPRequestHandler:
first we check if the team has a object_permission_id attached
- if it does then we look up the object_permission for the team
"""
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if user_api_key_auth is None:
return []
@ -399,10 +426,12 @@ 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 get_team_object(
team_id=user_api_key_auth.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_auth.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if team_obj is None:
verbose_logger.debug("team_obj is None")
@ -534,7 +563,12 @@ class MCPRequestHandler:
async def _get_mcp_access_groups_for_key(
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
) -> List[str]:
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.auth.auth_checks import get_object_permission
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if user_api_key_auth is None:
return []
@ -546,15 +580,21 @@ class MCPRequestHandler:
verbose_logger.debug("prisma_client is None")
return []
key_object_permission = (
await prisma_client.db.litellm_objectpermissiontable.find_unique(
where={"object_permission_id": user_api_key_auth.object_permission_id},
try:
key_object_permission = await get_object_permission(
object_permission_id=user_api_key_auth.object_permission_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_auth.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
)
if key_object_permission is None:
return []
if key_object_permission is None:
return []
return key_object_permission.mcp_access_groups or []
return key_object_permission.mcp_access_groups or []
except Exception as e:
verbose_logger.warning(f"Failed to get MCP access groups for key: {str(e)}")
return []
@staticmethod
async def _get_mcp_access_groups_for_team(
@ -563,7 +603,12 @@ class MCPRequestHandler:
"""
Get MCP access groups for the team
"""
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if user_api_key_auth is None:
return []
@ -575,20 +620,28 @@ 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},
try:
team_obj: Optional[LiteLLM_TeamTable] = await get_team_object(
team_id=user_api_key_auth.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_auth.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
)
if team_obj is None:
verbose_logger.debug("team_obj is None")
return []
if team_obj is None:
verbose_logger.debug("team_obj is None")
return []
object_permissions = team_obj.object_permission
if object_permissions is None:
return []
object_permissions = team_obj.object_permission
if object_permissions is None:
return []
return object_permissions.mcp_access_groups or []
return object_permissions.mcp_access_groups or []
except Exception as e:
verbose_logger.warning(
f"Failed to get MCP access groups for team: {str(e)}"
)
return []
@staticmethod
def get_mcp_access_groups_from_headers(headers: Headers) -> Optional[List[str]]:

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

@ -919,6 +919,14 @@ class MCPServerManager:
return server
return None
def get_mcp_server_names_from_ids(self, server_ids: List[str]) -> List[str]:
server_names = []
registry = self.get_registry()
for server in registry.values():
if server.server_id in server_ids:
server_names.append(server.name)
return server_names
def get_mcp_server_by_name(self, server_name: str) -> Optional[MCPServer]:
"""
Get the MCP Server from the server name
@ -1189,6 +1197,11 @@ class MCPServerManager:
if server.mcp_access_groups is not None
else []
),
allowed_tools=(
server.allowed_tools
if server.allowed_tools is not None
else []
),
mcp_info=server.mcp_info,
teams=cast(
List[Dict[str, str | None]],

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

@ -28,6 +28,7 @@ if MCP_AVAILABLE:
from litellm.proxy._experimental.mcp_server.server import (
ListMCPToolsRestAPIResponseObject,
call_mcp_tool,
filter_tools_by_allowed_tools,
)
########################################################
@ -75,6 +76,12 @@ if MCP_AVAILABLE:
mcp_auth_header=server_auth_header,
add_prefix=False,
)
# Filter tools based on allowed_tools configuration
# Only filter if allowed_tools is explicitly configured (not None and not empty)
if server.allowed_tools is not None and len(server.allowed_tools) > 0:
tools = filter_tools_by_allowed_tools(tools, server)
return _create_tool_response_objects(tools, server.mcp_info)
########################################################

76
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -361,22 +361,66 @@ if MCP_AVAILABLE:
return allowed_mcp_servers
def _tool_name_matches(tool_name: str, filter_list: List[str]) -> bool:
"""
Check if a tool name matches any name in the filter list.
Checks both the full tool name and unprefixed version (without server prefix).
This allows users to configure simple tool names regardless of prefixing.
Args:
tool_name: The tool name to check (may be prefixed like "server-tool_name")
filter_list: List of tool names to match against
Returns:
True if the tool name (prefixed or unprefixed) is in the filter list
"""
from litellm.proxy._experimental.mcp_server.utils import (
get_server_name_prefix_tool_mcp,
)
# Check if the full name is in the list
if tool_name in filter_list:
return True
# Check if the unprefixed name is in the list
unprefixed_name, _ = get_server_name_prefix_tool_mcp(tool_name)
return unprefixed_name in filter_list
def filter_tools_by_allowed_tools(
tools: List[MCPTool],
mcp_server: MCPServer,
) -> List[MCPTool]:
"""
Filter tools by allowed tools
Filter tools by allowed/disallowed tools configuration.
If allowed_tools is set, only tools in that list are returned.
If disallowed_tools is set, tools in that list are excluded.
Tool names are matched with and without server prefixes for flexibility.
Args:
tools: List of tools to filter
mcp_server: Server configuration with allowed_tools/disallowed_tools
Returns:
Filtered list of tools
"""
tools_to_return = tools
# Filter by allowed_tools (whitelist)
if mcp_server.allowed_tools:
tools_to_return = [
tool for tool in tools if tool.name in mcp_server.allowed_tools
tool for tool in tools
if _tool_name_matches(tool.name, mcp_server.allowed_tools)
]
# Filter by disallowed_tools (blacklist)
if mcp_server.disallowed_tools:
tools_to_return = [
tool for tool in tools if tool.name not in mcp_server.disallowed_tools
tool for tool in tools_to_return
if not _tool_name_matches(tool.name, mcp_server.disallowed_tools)
]
return tools_to_return
async def _get_tools_from_mcp_servers(
@ -453,9 +497,12 @@ if MCP_AVAILABLE:
extra_headers=extra_headers,
add_prefix=add_prefix,
)
all_tools.extend(filter_tools_by_allowed_tools(tools, server))
filtered_tools = filter_tools_by_allowed_tools(tools, server)
all_tools.extend(filtered_tools)
verbose_logger.debug(
f"Successfully fetched {len(tools)} tools from server {server.name}"
f"Successfully fetched {len(tools)} tools from server {server.name}, {len(filtered_tools)} after filtering"
)
except Exception as e:
verbose_logger.exception(
@ -560,6 +607,25 @@ if MCP_AVAILABLE:
name
)
## CHECK IF USER IS ALLOWED TO CALL THIS TOOL
allowed_mcp_server_ids = await MCPRequestHandler.get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth,
)
allowed_mcp_servers = global_mcp_server_manager.get_mcp_server_names_from_ids(
allowed_mcp_server_ids
)
if not MCPRequestHandler.is_tool_allowed(
allowed_mcp_servers=allowed_mcp_servers,
server_name=server_name_from_prefix,
):
raise HTTPException(
status_code=403,
detail=f"User not allowed to call this tool. Allowed MCP servers: {allowed_mcp_servers}",
)
standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = (
_get_standard_logging_mcp_tool_call(
name=original_tool_name, # Use original name for logging

2
litellm/proxy/_types.py
View file

@ -917,6 +917,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
url: Optional[str] = None
mcp_info: Optional[MCPInfo] = None
mcp_access_groups: List[str] = Field(default_factory=list)
allowed_tools: Optional[List[str]] = None
# Stdio-specific fields
command: Optional[str] = None
args: List[str] = Field(default_factory=list)
@ -985,6 +986,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
updated_by: Optional[str] = None
teams: List[Dict[str, Optional[str]]] = Field(default_factory=list)
mcp_access_groups: List[str] = Field(default_factory=list)
allowed_tools: List[str] = Field(default_factory=list)
mcp_info: Optional[MCPInfo] = None
# Health check status
status: Optional[str] = Field(

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

@ -41,12 +41,12 @@ from litellm.proxy._types import (
LiteLLM_UserTable,
LiteLLMRoutes,
LitellmUserRoles,
NewTeamRequest,
ProxyErrorTypes,
ProxyException,
RoleBasedPermissions,
SpecialModelNames,
UserAPIKeyAuth,
NewTeamRequest,
)
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.route_llm_request import route_request
@ -474,7 +474,7 @@ async def get_end_user_object(
return return_obj
# else, check db
try:
try:
response = await prisma_client.db.litellm_endusertable.find_unique(
where={"user_id": end_user_id},
include={"litellm_budget_table": True},
@ -817,7 +817,9 @@ async def _cache_management_object(
):
await user_api_key_cache.async_set_cache(
key=key, value=value, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
key=key,
value=value,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
@ -892,7 +894,9 @@ async def _get_team_db_check(
system_admin_user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
created_team_dict = await new_team(
data=new_team_data, http_request=mock_request, user_api_key_dict=system_admin_user
data=new_team_data,
http_request=mock_request,
user_api_key_dict=system_admin_user,
)
response = LiteLLM_TeamTable(**created_team_dict)
return response
@ -1166,6 +1170,54 @@ async def get_key_object(
return _response
@log_db_metrics
async def get_object_permission(
object_permission_id: str,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
parent_otel_span: Optional[Span] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> Optional[LiteLLM_ObjectPermissionTable]:
"""
- Check if object permission id in proxy ObjectPermissionTable
- if valid, return LiteLLM_ObjectPermissionTable object
- if not, then raise an error
"""
if prisma_client is None:
raise Exception(
"No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys"
)
# check if in cache
key = "object_permission_id:{}".format(object_permission_id)
cached_obj_permission = await user_api_key_cache.async_get_cache(key=key)
if cached_obj_permission is not None:
if isinstance(cached_obj_permission, dict):
return LiteLLM_ObjectPermissionTable(**cached_obj_permission)
elif isinstance(cached_obj_permission, LiteLLM_ObjectPermissionTable):
return cached_obj_permission
# else, check db
try:
response = await prisma_client.db.litellm_objectpermissiontable.find_unique(
where={"object_permission_id": object_permission_id}
)
if response is None:
return None
# save the object permission to cache
await user_api_key_cache.async_set_cache(
key=key,
value=response.model_dump(),
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
return LiteLLM_ObjectPermissionTable(**response.dict())
except Exception:
return None
@log_db_metrics
async def get_org_object(
org_id: str,

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

@ -417,6 +417,12 @@ def bytes_to_mb(bytes_value: int):
def get_key_model_rpm_limit(
user_api_key_dict: UserAPIKeyAuth,
) -> Optional[Dict[str, int]]:
"""
Get the model rpm limit for a given api key
- check key metadata
- check key model max budget
- check team metadata
"""
if user_api_key_dict.metadata:
if "model_rpm_limit" in user_api_key_dict.metadata:
return user_api_key_dict.metadata["model_rpm_limit"]
@ -426,7 +432,9 @@ def get_key_model_rpm_limit(
if "rpm_limit" in budget and budget["rpm_limit"] is not None:
model_rpm_limit[model] = budget["rpm_limit"]
return model_rpm_limit
elif user_api_key_dict.team_metadata:
if "model_rpm_limit" in user_api_key_dict.team_metadata:
return user_api_key_dict.team_metadata["model_rpm_limit"]
return None
@ -439,7 +447,9 @@ def get_key_model_tpm_limit(
elif user_api_key_dict.model_max_budget:
if "tpm_limit" in user_api_key_dict.model_max_budget:
return user_api_key_dict.model_max_budget["tpm_limit"]
elif user_api_key_dict.team_metadata:
if "model_tpm_limit" in user_api_key_dict.team_metadata:
return user_api_key_dict.team_metadata["model_tpm_limit"]
return None
@ -473,6 +483,7 @@ def _has_user_setup_sso():
return sso_setup
def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[str]:
"""Return the header_name mapped to CUSTOMER role, if any (dict-based)."""
if not user_id_mapping:
@ -522,7 +533,11 @@ def get_end_user_id_from_request_body(
for header_name, header_value in request_headers.items():
if header_name.lower() == custom_header_name_to_check.lower():
user_id_from_header = header_value
user_id_str = str(user_id_from_header) if user_id_from_header is not None else ""
user_id_str = (
str(user_id_from_header)
if user_id_from_header is not None
else ""
)
if user_id_str.strip():
return user_id_str

12
litellm/proxy/auth/route_checks.py
View file

@ -62,12 +62,22 @@ class RouteChecks:
for allowed_route in valid_token.allowed_routes
):
for allowed_route in valid_token.allowed_routes:
if allowed_route in LiteLLMRoutes._member_names_:
if allowed_route in LiteLLMRoutes._member_names_:
if RouteChecks.check_route_access(
route=route,
allowed_routes=LiteLLMRoutes._member_map_[allowed_route].value,
):
return True
################################################
# For llm_api_routes, also check registered pass-through endpoints
################################################
if allowed_route == "llm_api_routes":
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
InitPassThroughEndpointHelpers,
)
if InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route):
return True
# check if wildcard pattern is allowed
for allowed_route in valid_token.allowed_routes:

156
litellm/proxy/common_request_processing.py
View file

@ -38,6 +38,7 @@ from litellm.proxy.common_utils.callback_utils import (
from litellm.proxy.route_llm_request import route_request
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
from litellm.types.utils import ServerToolUse
if TYPE_CHECKING:
from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig
@ -46,6 +47,7 @@ if TYPE_CHECKING:
else:
ProxyConfig = Any
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
from litellm.types.utils import ModelResponse, ModelResponseStream, Usage
async def _parse_event_data_for_error(event_line: Union[str, bytes]) -> Optional[int]:
@ -379,6 +381,7 @@ class ProxyBaseLLMRequestProcessing:
user_api_base: Optional[str] = None,
version: Optional[str] = None,
is_streaming_request: Optional[bool] = False,
contents: Optional[list] = None, # Add contents parameter
) -> Any:
"""
Common request processing logic for both chat completions and responses API endpoints
@ -417,6 +420,10 @@ class ProxyBaseLLMRequestProcessing:
)
)
# Pass contents if provided
if contents:
self.data["contents"] = contents
### ROUTE THE REQUEST ###
# Do not change this - it should be a constant time fetch - ALWAYS
llm_call = await route_request(
@ -729,7 +736,6 @@ class ProxyBaseLLMRequestProcessing:
"""
Anthropic /messages and Google /generateContent streaming data generator require SSE events
"""
from litellm.types.utils import ModelResponse, ModelResponseStream
verbose_proxy_logger.debug("inside generator")
try:
@ -754,6 +760,10 @@ class ProxyBaseLLMRequestProcessing:
response_str = litellm.get_response_string(response_obj=chunk)
str_so_far += response_str
# Inject cost into Anthropic-style SSE usage for /v1/messages for any provider
model_name = request_data.get("model", "")
chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, model_name)
# Format chunk using helper function
yield ProxyBaseLLMRequestProcessing.return_sse_chunk(chunk)
except Exception as e:
@ -785,3 +795,147 @@ class ProxyBaseLLMRequestProcessing:
)
error_returned = json.dumps({"error": proxy_exception.to_dict()})
yield f"{STREAM_SSE_DATA_PREFIX}{error_returned}\n\n"
@staticmethod
def _process_chunk_with_cost_injection(chunk: Any, model_name: str) -> Any:
"""
Process a streaming chunk and inject cost information if enabled.
Args:
chunk: The streaming chunk (dict, str, bytes, or bytearray)
model_name: Model name for cost calculation
Returns:
The processed chunk with cost information injected if applicable
"""
if not getattr(litellm, "include_cost_in_streaming_usage", False):
return chunk
try:
if isinstance(chunk, dict):
maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(chunk, model_name)
if maybe_modified is not None:
return maybe_modified
elif isinstance(chunk, (bytes, bytearray)):
# Decode to str, inject, and rebuild as bytes
try:
s = chunk.decode("utf-8", errors="ignore")
maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(s, model_name)
if maybe_mod is not None:
return (maybe_mod + ("" if maybe_mod.endswith("\n\n") else "\n\n")).encode("utf-8")
except Exception:
pass
elif isinstance(chunk, str):
# Try to parse SSE frame and inject cost into the data line
maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(chunk, model_name)
if maybe_mod is not None:
# Ensure trailing frame separator
return maybe_mod if maybe_mod.endswith("\n\n") else (maybe_mod + "\n\n")
except Exception:
# Never break streaming on optional cost injection
pass
return chunk
@staticmethod
def _inject_cost_into_sse_frame_str(frame_str: str, model_name: str) -> Optional[str]:
"""
Inject cost information into an SSE frame string by modifying the JSON in the 'data:' line.
Args:
frame_str: SSE frame string that may contain multiple lines
model_name: Model name for cost calculation
Returns:
Modified SSE frame string with cost injected, or None if no modification needed
"""
try:
# Split preserving lines
lines = frame_str.split("\n")
for idx, ln in enumerate(lines):
stripped_ln = ln.strip()
if stripped_ln.startswith("data:"):
json_part = stripped_ln.split("data:", 1)[1].strip()
if json_part and json_part != "[DONE]":
obj = json.loads(json_part)
maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(obj, model_name)
if maybe_modified is not None:
# Replace just this line with updated JSON using safe_dumps
lines[idx] = f"data: {safe_dumps(maybe_modified)}"
return "\n".join(lines)
return None
except Exception:
return None
@staticmethod
def _inject_cost_into_usage_dict(obj: dict, model_name: str) -> Optional[dict]:
"""
Inject cost information into a usage dictionary for message_delta events.
Args:
obj: Dictionary containing the SSE event data
model_name: Model name for cost calculation
Returns:
Modified dictionary with cost injected, or None if no modification needed
"""
if (
obj.get("type") == "message_delta"
and isinstance(obj.get("usage"), dict)
):
_usage = obj["usage"]
prompt_tokens = int(_usage.get("input_tokens", 0) or 0)
completion_tokens = int(_usage.get("output_tokens", 0) or 0)
total_tokens = int(
_usage.get("total_tokens", prompt_tokens + completion_tokens)
or (prompt_tokens + completion_tokens)
)
# Extract additional usage fields
cache_creation_input_tokens = _usage.get("cache_creation_input_tokens")
cache_read_input_tokens = _usage.get("cache_read_input_tokens")
web_search_requests = _usage.get("web_search_requests")
completion_tokens_details = _usage.get("completion_tokens_details")
prompt_tokens_details = _usage.get("prompt_tokens_details")
usage_kwargs: dict[str, Any] = {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": total_tokens,
}
# Add optional named parameters
if completion_tokens_details is not None:
usage_kwargs["completion_tokens_details"] = completion_tokens_details
if prompt_tokens_details is not None:
usage_kwargs["prompt_tokens_details"] = prompt_tokens_details
# Handle web_search_requests by wrapping in ServerToolUse
if web_search_requests is not None:
usage_kwargs["server_tool_use"] = ServerToolUse(
web_search_requests=web_search_requests
)
# Add cache-related fields to **params (handled by Usage.__init__)
if cache_creation_input_tokens is not None:
usage_kwargs["cache_creation_input_tokens"] = cache_creation_input_tokens
if cache_read_input_tokens is not None:
usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens
_mr = ModelResponse(
usage=Usage(**usage_kwargs)
)
try:
cost_val = litellm.completion_cost(
completion_response=_mr,
model=model_name,
)
except Exception:
cost_val = None
if cost_val is not None:
obj.setdefault("usage", {})["cost"] = cost_val
return obj
return None

5
litellm/proxy/common_utils/openai_endpoint_utils.py
View file

@ -6,8 +6,11 @@ from typing import Optional
from fastapi import Request
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
SENSITIVE_DATA_MASKER = SensitiveDataMasker()
def remove_sensitive_info_from_deployment(deployment_dict: dict) -> dict:
"""
@ -25,6 +28,8 @@ def remove_sensitive_info_from_deployment(deployment_dict: dict) -> dict:
deployment_dict["litellm_params"].pop("aws_access_key_id", None)
deployment_dict["litellm_params"].pop("aws_secret_access_key", None)
deployment_dict["litellm_params"] = SENSITIVE_DATA_MASKER.mask_dict(deployment_dict["litellm_params"])
return deployment_dict

15
litellm/proxy/example_config_yaml/pass_through_config.yaml
View file

@ -24,6 +24,19 @@ model_list:
litellm_params:
model: anthropic/*
api_key: os.environ/ANTHROPIC_API_KEY
- model_name: openai/*
litellm_params:
model: openai/*
api_key: os.environ/OPENAI_API_KEY
general_settings:
master_key: sk-1234
custom_auth: custom_auth_basic.user_api_key_auth
custom_auth: custom_auth_basic.user_api_key_auth
pass_through_endpoints:
- path: "/azure-config-passthrough"
target: os.environ/AZURE_API_BASE_PASSHROUGH
include_subpath: true
headers:
Authorization: os.environ/AZURE_API_KEY_PASSHROUGH
litellm_settings:
include_cost_in_streaming_usage: true

151
litellm/proxy/google_endpoints/endpoints.py
View file

@ -1,8 +1,10 @@
from fastapi import APIRouter, Depends, Request, Response
from fastapi import APIRouter, Depends, Request, Response, HTTPException
from fastapi.responses import StreamingResponse
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
router = APIRouter(
@ -10,140 +12,63 @@ router = APIRouter(
)
@router.post("/v1beta/models/{model_name}:generateContent", dependencies=[Depends(user_api_key_auth)])
@router.post("/models/{model_name}:generateContent", dependencies=[Depends(user_api_key_auth)])
@router.post(
"/v1beta/models/{model_name}:generateContent",
dependencies=[Depends(user_api_key_auth)],
)
@router.post(
"/models/{model_name}:generateContent", dependencies=[Depends(user_api_key_auth)]
)
async def google_generate_content(
request: Request,
model_name: str,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Not Implemented, this is a placeholder for the google genai generateContent endpoint.
"""
from litellm.proxy.proxy_server import (
_read_request_body,
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
select_data_generator,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
from litellm.proxy.proxy_server import llm_router
data = await _read_request_body(request=request)
if "model" not in data:
data["model"] = model_name
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="agenerate_content",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=None,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
# call router
if llm_router is None:
raise HTTPException(status_code=500, detail="Router not initialized")
response = await llm_router.agenerate_content(**data)
return response
class GoogleAIStudioDataGenerator:
"""
Ensures SSE data generator is used for Google AI Studio streaming responses
Thin wrapper around ProxyBaseLLMRequestProcessing.async_sse_data_generator
"""
@staticmethod
def _select_data_generator(response, user_api_key_dict, request_data):
from litellm.proxy.proxy_server import proxy_logging_obj
return ProxyBaseLLMRequestProcessing.async_sse_data_generator(
response=response,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
proxy_logging_obj=proxy_logging_obj,
)
@router.post("/v1beta/models/{model_name}:streamGenerateContent", dependencies=[Depends(user_api_key_auth)])
@router.post("/models/{model_name}:streamGenerateContent", dependencies=[Depends(user_api_key_auth)])
@router.post(
"/v1beta/models/{model_name}:streamGenerateContent",
dependencies=[Depends(user_api_key_auth)],
)
@router.post(
"/models/{model_name}:streamGenerateContent",
dependencies=[Depends(user_api_key_auth)],
)
async def google_stream_generate_content(
request: Request,
model_name: str,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Not Implemented, this is a placeholder for the google genai streamGenerateContent endpoint.
"""
from litellm.proxy.proxy_server import (
_read_request_body,
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
from litellm.proxy.proxy_server import llm_router
data = await _read_request_body(request=request)
if "model" not in data:
data["model"] = model_name
data["stream"] = True # enforce streaming for this endpoint
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="agenerate_content_stream",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=GoogleAIStudioDataGenerator._select_data_generator,
model=None,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
is_streaming_request=True,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
# call router
if llm_router is None:
raise HTTPException(status_code=500, detail="Router not initialized")
response = await llm_router.agenerate_content(**data)
# Check if response is an async iterator (streaming response)
if hasattr(response, "__aiter__"):
return StreamingResponse(response, media_type="text/event-stream")
return response
@router.post(
@ -171,13 +96,13 @@ async def google_count_tokens(request: Request, model_name: str):
}
```
"""
from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.proxy_server import token_counter as internal_token_counter
from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
data = await _read_request_body(request=request)
contents = data.get("contents", [])
#Create TokenCountRequest for the internal endpoint
# Create TokenCountRequest for the internal endpoint
from litellm.proxy._types import TokenCountRequest
# Translate contents to openai format messages using the adapter

67
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -41,6 +41,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
)
from litellm.types.utils import (
Choices,
GuardrailStatus,
ModelResponse,
ModelResponseStream,
StreamingChoices,
@ -361,11 +362,30 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
prepared_request.headers,
)
httpx_response = await self.async_handler.post(
url=prepared_request.url,
data=prepared_request.body, # type: ignore
headers=prepared_request.headers, # type: ignore
)
try:
httpx_response = await self.async_handler.post(
url=prepared_request.url,
data=prepared_request.body, # type: ignore
headers=prepared_request.headers, # type: ignore
)
except Exception as e:
# Endpoint down, timeout, or other HTTP/network errors
verbose_proxy_logger.error(
"Bedrock AI: failed to make guardrail request: %s", str(e)
)
# Add guardrail information with failure status
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response={"error": str(e)},
request_data=request_data or {},
guardrail_status="guardrail_failed_to_respond",
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
)
# Re-raise the exception to maintain existing behavior
raise
#########################################################
# Add guardrail information to request trace
#########################################################
@ -437,15 +457,30 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
def _get_bedrock_guardrail_response_status(
self, response: httpx.Response
) -> Literal["success", "failure"]:
) -> GuardrailStatus:
"""
Get the status of the bedrock guardrail response.
Returns:
"success": Content allowed through with no violations
"guardrail_intervened": Content blocked due to policy violations
"guardrail_failed_to_respond": Technical error or API failure
"""
if response.status_code == 200:
if self._check_bedrock_response_for_exception(response):
return "failure"
return "guardrail_failed_to_respond"
# Check if the guardrail would block content
try:
_json_response = response.json()
bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response)
if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response):
return "guardrail_intervened"
except Exception:
pass
return "success"
return "failure"
return "guardrail_failed_to_respond"
def _get_http_exception_for_blocked_guardrail(
self, response: BedrockGuardrailResponse
@ -692,6 +727,22 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
return
# Check if the ModelResponse has text content in its choices
# to avoid sending empty content to Bedrock (e.g., during tool calls)
if isinstance(response, litellm.ModelResponse):
has_text_content = False
for choice in response.choices:
if isinstance(choice, litellm.Choices):
if choice.message.content and isinstance(choice.message.content, str):
has_text_content = True
break
if not has_text_content:
verbose_proxy_logger.warning(
"Bedrock AI: not running guardrail. No output text in response"
)
return
#########################################################
########## 1. Make parallel Bedrock API requests ##########
#########################################################

18
litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py
View file

@ -1,5 +1,7 @@
from datetime import datetime
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union, Type
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Type, Union
from fastapi import HTTPException
import litellm
from litellm._logging import verbose_proxy_logger
@ -12,11 +14,11 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.secret_managers.main import get_secret_str
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.guardrails.guardrail_hooks.javelin import (
JavelinGuardInput,
JavelinGuardRequest,
JavelinGuardResponse,
JavelinGuardInput,
)
from fastapi import HTTPException
from litellm.types.utils import GuardrailStatus
if TYPE_CHECKING:
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
@ -95,7 +97,7 @@ class JavelinGuardrail(CustomGuardrail):
if self.application:
headers["x-javelin-application"] = self.application
status: Literal["success", "failure", "blocked"] = "failure"
status: GuardrailStatus = "guardrail_failed_to_respond"
javelin_response: Optional[JavelinGuardResponse] = None
exception_str = ""
@ -122,7 +124,7 @@ class JavelinGuardrail(CustomGuardrail):
status = "success"
return javelin_response
except Exception as e:
status = "failure"
status = "guardrail_failed_to_respond"
exception_str = str(e)
return {"assessments": []}
finally:
@ -178,12 +180,12 @@ class JavelinGuardrail(CustomGuardrail):
"""
Pre-call hook for the Javelin guardrail.
"""
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_last_user_message,
)
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
)
verbose_proxy_logger.debug("Javelin Guardrail: pre_call_hook")
verbose_proxy_logger.debug("Javelin Guardrail: Request data: %s", data)

5
litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
View file

@ -20,6 +20,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import (
LakeraAIRequest,
LakeraAIResponse,
)
from litellm.types.utils import GuardrailStatus
class LakeraAIGuardrail(CustomGuardrail):
@ -70,7 +71,7 @@ class LakeraAIGuardrail(CustomGuardrail):
"""
Call the Lakera AI v2 guard API.
"""
status: Literal["success", "failure"] = "success"
status: GuardrailStatus = "success"
exception_str: str = ""
start_time: datetime = datetime.now()
lakera_response: Optional[LakeraAIResponse] = None
@ -99,7 +100,7 @@ class LakeraAIGuardrail(CustomGuardrail):
lakera_response = LakeraAIResponse(**response.json())
return lakera_response, masked_entity_count
except Exception as e:
status = "failure"
status = "guardrail_failed_to_respond"
exception_str = str(e)
raise e
finally:

5
litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
View file

@ -30,6 +30,7 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import (
Choices,
GuardrailStatus,
ModelResponse,
ModelResponseStream,
)
@ -329,14 +330,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
guardrail_response = metadata.get("_model_armor_response", {})
# Determine status – default to "success" but prefer the explicit value if present.
guardrail_status: Literal["success", "failure", "blocked"] = metadata.get(
guardrail_status: GuardrailStatus = metadata.get(
"_model_armor_status", "success"
) # type: ignore
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=guardrail_response,
request_data=request_data,
guardrail_status=guardrail_status, # type: ignore
guardrail_status=guardrail_status,
duration=duration,
start_time=start_time,
end_time=end_time,

119
litellm/proxy/guardrails/guardrail_hooks/noma/noma.py
View file

@ -8,7 +8,8 @@
import asyncio
import copy
import os
from typing import Any, Dict, Final, Literal, Optional, Union, Type, TYPE_CHECKING
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, Final, Literal, Optional, Type, Union
from urllib.parse import urljoin
from fastapi import HTTPException
@ -23,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import EmbeddingResponse, ImageResponse
from litellm.types.utils import EmbeddingResponse, GuardrailStatus, ImageResponse
# Constants
USER_ROLE: Final[Literal["user"]] = "user"
@ -204,6 +205,7 @@ class NomaGuardrail(CustomGuardrail):
user_auth: UserAPIKeyAuth,
) -> Optional[str]:
"""Shared logic for processing user message checks"""
start_time = datetime.now()
extra_data = self.get_guardrail_dynamic_request_body_params(request_data)
user_message = await self._extract_user_message(request_data)
@ -218,6 +220,23 @@ class NomaGuardrail(CustomGuardrail):
user_auth=user_auth,
extra_data=extra_data,
)
end_time = datetime.now()
duration = (end_time - start_time).total_seconds()
# Determine guardrail status based on response
guardrail_status = self._determine_guardrail_status(response_json)
# Always log guardrail information for consistency
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
guardrail_json_response=response_json,
request_data=request_data,
guardrail_status=guardrail_status,
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
)
if self.monitor_mode:
await self._handle_verdict_background(
@ -248,6 +267,8 @@ class NomaGuardrail(CustomGuardrail):
user_auth: UserAPIKeyAuth,
) -> Optional[str]:
"""Shared logic for processing LLM response checks"""
start_time = datetime.now()
extra_data = self.get_guardrail_dynamic_request_body_params(request_data)
if not isinstance(response, litellm.ModelResponse):
@ -271,6 +292,23 @@ class NomaGuardrail(CustomGuardrail):
user_auth=user_auth,
extra_data=extra_data,
)
end_time = datetime.now()
duration = (end_time - start_time).total_seconds()
# Determine guardrail status based on response
guardrail_status = self._determine_guardrail_status(response_json)
# Always log guardrail information for consistency
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
guardrail_json_response=response_json,
request_data=request_data,
guardrail_status=guardrail_status,
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
)
if self.monitor_mode:
await self._handle_verdict_background(
@ -294,6 +332,41 @@ class NomaGuardrail(CustomGuardrail):
await self._check_verdict(ASSISTANT_ROLE, content, response_json)
return content
def _determine_guardrail_status(self, response_json: dict) -> GuardrailStatus:
"""
Determine the guardrail status based on NOMA API response.
Args:
response_json: Response from NOMA API
Returns:
"success": Content allowed through with no violations
"guardrail_intervened": Content blocked due to policy violations
"guardrail_failed_to_respond": Technical error or API failure
"""
try:
# Check if we got a valid response structure
if not isinstance(response_json, dict):
return "guardrail_failed_to_respond"
# Get the verdict from the response
verdict = response_json.get("verdict", True)
# If verdict is True, content is allowed
if verdict is True:
return "success"
# If verdict is False, content is blocked/flagged
if verdict is False:
return "guardrail_intervened"
# If verdict is missing or invalid, treat as failure
return "guardrail_failed_to_respond"
except Exception as e:
verbose_proxy_logger.error(f"Error determining NOMA guardrail status: {str(e)}")
return "guardrail_failed_to_respond"
def _should_only_sensitive_data_failed(self, classification_obj: dict) -> bool:
"""
Check if only sensitive data detectors (PII, PCI, secrets) have result=true in the classification.
@ -539,8 +612,22 @@ class NomaGuardrail(CustomGuardrail):
try:
return await self._check_user_message(data, user_api_key_dict)
except NomaBlockedMessage:
# Blocked requests were already logged in _process_user_message_check with "blocked" status
raise
except Exception as e:
# Log technical failures
from datetime import datetime
start_time = datetime.now()
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
guardrail_json_response=str(e),
request_data=data,
guardrail_status="guardrail_failed_to_respond",
start_time=start_time.timestamp(),
end_time=start_time.timestamp(),
duration=0.0,
)
verbose_proxy_logger.error(f"Noma pre-call hook failed: {str(e)}")
if self.block_failures:
@ -580,8 +667,22 @@ class NomaGuardrail(CustomGuardrail):
try:
return await self._check_user_message(data, user_api_key_dict)
except NomaBlockedMessage:
# Blocked requests were already logged in _process_user_message_check with "blocked" status
raise
except Exception as e:
# Log technical failures
from datetime import datetime
start_time = datetime.now()
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
guardrail_json_response=str(e),
request_data=data,
guardrail_status="guardrail_failed_to_respond",
start_time=start_time.timestamp(),
end_time=start_time.timestamp(),
duration=0.0,
)
verbose_proxy_logger.error(f"Noma moderation hook failed: {str(e)}")
if self.block_failures:
@ -615,8 +716,22 @@ class NomaGuardrail(CustomGuardrail):
try:
return await self._check_llm_response(data, response, user_api_key_dict)
except NomaBlockedMessage:
# Blocked requests were already logged in _process_llm_response_check with "blocked" status
raise
except Exception as e:
# Log technical failures
from datetime import datetime
start_time = datetime.now()
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
guardrail_json_response=str(e),
request_data=data,
guardrail_status="guardrail_failed_to_respond",
start_time=start_time.timestamp(),
end_time=start_time.timestamp(),
duration=0.0,
)
verbose_proxy_logger.error(f"Noma post-call hook failed: {str(e)}")
if self.block_failures:
raise

8
litellm/proxy/guardrails/guardrail_hooks/presidio.py
View file

@ -10,14 +10,12 @@
import asyncio
import json
from litellm._uuid import uuid
from datetime import datetime
from typing import (
Any,
AsyncGenerator,
Dict,
List,
Literal,
Optional,
Tuple,
Union,
@ -29,6 +27,7 @@ import aiohttp
import litellm # noqa: E401
from litellm import get_secret
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.caching.caching import DualCache
from litellm.exceptions import BlockedPiiEntityError
from litellm.integrations.custom_guardrail import CustomGuardrail
@ -45,6 +44,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.presidio import (
PresidioAnalyzeResponseItem,
)
from litellm.types.utils import CallTypes as LitellmCallTypes
from litellm.types.utils import GuardrailStatus
from litellm.utils import (
EmbeddingResponse,
ImageResponse,
@ -324,7 +324,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
"""
start_time = datetime.now()
analyze_results: Optional[Union[List[PresidioAnalyzeResponseItem], Dict]] = None
status: Literal["success", "failure"] = "success"
status: GuardrailStatus = "success"
masked_entity_count: Dict[str, int] = {}
exception_str: str = ""
try:
@ -356,7 +356,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
)
return redacted_text["text"]
except Exception as e:
status = "failure"
status = "guardrail_failed_to_respond"
exception_str = str(e)
raise e
finally:

170
litellm/proxy/hooks/README.dynamic_rate_limiter_v3.md Normal file
View file

@ -0,0 +1,170 @@
# Dynamic Rate Limiter v3 - Saturation-Aware Priority-Based Rate Limiting
## Overview
The v3 dynamic rate limiter implements saturation-aware rate limiting with priority-based allocation. It balances resource efficiency (allowing unused capacity to be borrowed) with fairness guarantees (enforcing priorities during high load).
**Key Behavior:**
- When system is under 80% capacity: Generous mode - allows priority borrowing
- When system is at/above 80% capacity: Strict mode - enforces normalized priority limits
## How It Works
### Flow Diagram
```
┌─────────────────────────────────────────────────────────────┐
│ Incoming Request │
└────────────────────────┬────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ 1. Check Model Saturation │
│ - Query v3 limiter's Redis counters │
│ - Calculate: current_usage / capacity │
│ - Returns: 0.0 (empty) to 1.0+ (saturated) │
└────────────────────────┬────────────────────────────────────┘
│
▼
┌────────┴────────┐
│ Saturation? │
└────────┬────────┘
│
┌───────────────┴───────────────┐
│ │
▼ ▼
< 80% (Generous) >= 80% (Strict)
│ │
▼ ▼
┌─────────────────────┐ ┌─────────────────────┐
│ Generous Mode │ │ Strict Mode │
│ │ │ │
│ - Enforce model- │ │ - Normalize │
│ wide capacity │ │ priority weights │
│ - No priority │ │ (if over 1.0) │
│ restrictions │ │ │
│ - Allows borrowing │ │ - Create priority- │
│ │ │ specific │
│ - First-come- │ │ descriptors │
│ first-served │ │ │
│ until capacity │ │ - Enforce strict │
│ │ │ limits per │
│ │ │ priority │
└──────────┬──────────┘ └──────────┬──────────┘
│ │
│ ▼
│ ┌──────────────────────┐
│ │ Track model usage │
│ │ for future │
│ │ saturation checks │
│ └──────────┬───────────┘
│ │
└───────────────┬───────────────┘
│
▼
┌──────────────┐
│ v3 Limiter │
│ Check │
└──────┬───────┘
│
┌───────────────┴───────────────┐
│ │
▼ ▼
OVER_LIMIT OK
│ │
▼ ▼
Return 429 Error Allow Request
```
## Configuration
### Priority Reservation
Set priority weights in your proxy configuration:
```python
litellm.priority_reservation = {
"premium": 0.75, # 75% of capacity
"standard": 0.25 # 25% of capacity
}
```
### Priority Reservation Settings
Configure saturation-aware behavior:
```python
litellm.priority_reservation_settings = PriorityReservationSettings(
default_priority=0.5, # Default weight for users without explicit priority
saturation_threshold=0.80, # 80% - threshold for strict mode enforcement
tracking_multiplier=10 # 10x - multiplier for non-blocking tracking in strict mode
)
```
**Settings:**
- `default_priority` (default: 0.5) - Priority weight for users without explicit priority metadata
- `saturation_threshold` (default: 0.80) - Saturation level (0.0-1.0) at which strict priority enforcement begins
- `tracking_multiplier` (default: 10) - Multiplier for model-wide tracking limits in strict mode
### User Priority Assignment
Set priority in user metadata:
```python
user_api_key_dict.metadata = {"priority": "premium"}
```
## Priority Weight Normalization
If priorities sum to > 1.0, they are automatically normalized:
```
Input: {key_a: 0.60, key_b: 0.80} = 1.40 total
Output: {key_a: 0.43, key_b: 0.57} = 1.00 total
```
This ensures total allocation never exceeds model capacity.
## Implementation Details
### Saturation Detection
- Queries v3 limiter's Redis counters for model-wide usage
- Checks both RPM and TPM, returns higher saturation value
- Non-blocking reads (doesn't increment counters)
### Mode Selection
**Generous Mode (< 80% saturation):**
- Creates single model-wide descriptor
- Enforces total capacity only
- Allows any priority to use available capacity
- Prevents over-subscription via model-wide limit
**Strict Mode (>= 80% saturation):**
- Creates priority-specific descriptors with normalized weights
- Each priority gets its reserved allocation
- Tracks model-wide usage separately (non-blocking, 10x multiplier)
- Ensures fairness under load
Test scenarios covered:
1. No rate limiting when under capacity
2. Priority queue behavior during saturation
3. Spillover capacity for default keys
4. Over-allocated priorities with normalization
5. Default priority value handling
### `_PROXY_DynamicRateLimitHandlerV3`
Main handler class inheriting from `CustomLogger`.
**Key Methods:**
- `async_pre_call_hook()` - Main entry point, routes to generous/strict mode
- `_check_model_saturation()` - Queries Redis for current usage
- `_handle_generous_mode()` - Enforces model-wide capacity only
- `_handle_strict_mode()` - Enforces normalized priority limits
- `_normalize_priority_weights()` - Handles over-allocation
- `_create_priority_based_descriptors()` - Creates rate limit descriptors

475
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
View file

@ -1,9 +1,9 @@
"""
Dynamic rate limiter v3
Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting
"""
import os
from typing import List, Literal, Optional, Union
from typing import Dict, List, Literal, Optional, Union
from fastapi import HTTPException
@ -24,12 +24,22 @@ from litellm.types.router import ModelGroupInfo
class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
"""
Simple validation version that uses v3 parallel request limiter for priority-based rate limiting.
Saturation-aware priority-based rate limiter using v3 infrastructure.
Key differences from original:
1. Uses v3 limiter's sliding window approach instead of per-minute cache buckets
2. Leverages Redis Lua scripts for atomic operations under high traffic
3. Creates priority-specific rate limit descriptors
Key features:
1. Model capacity ALWAYS enforced at 100% (prevents over-allocation)
2. Priority usage tracked from first request (accurate accounting)
3. Priority limits only enforced when saturated >= threshold
4. Three-phase checking prevents partial counter increments
5. Reuses v3 limiter's Redis-based tracking (multi-instance safe)
How it works:
- Phase 1: Read-only check of ALL limits (no increments)
- Phase 2: Decide enforcement based on saturation
- Phase 3: Increment counters only if request allowed
- When under-saturated: priorities can borrow unused capacity (generous)
- When saturated: strict priority-based limits enforced (fair)
- Uses v3 limiter's atomic Lua scripts for race-free increments
"""
def __init__(self, internal_usage_cache: DualCache):
self.internal_usage_cache = InternalUsageCache(dual_cache=internal_usage_cache)
@ -57,6 +67,147 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
weight = litellm.priority_reservation[priority]
return weight
def _normalize_priority_weights(self) -> Dict[str, float]:
"""
Normalize priority weights if they sum to > 1.0
Handles over-allocation: {key_a: 0.60, key_b: 0.80} -> {key_a: 0.43, key_b: 0.57}
"""
if litellm.priority_reservation is None:
return {}
weights = dict(litellm.priority_reservation)
total_weight = sum(weights.values())
if total_weight > 1.0:
normalized = {k: v / total_weight for k, v in weights.items()}
verbose_proxy_logger.debug(
f"Normalized over-allocated priorities: {weights} -> {normalized}"
)
return normalized
return weights
def _get_priority_allocation(
self,
model: str,
priority: Optional[str],
normalized_weights: Dict[str, float],
) -> tuple[float, str]:
"""
Get priority weight and pool key for a given priority.
For explicit priorities: returns specific allocation and unique pool key
For default priority: returns default allocation and shared pool key
Args:
model: Model name
priority: Priority level (None for default)
normalized_weights: Pre-computed normalized weights
Returns:
tuple: (priority_weight, priority_key)
"""
# Check if this key has an explicit priority in litellm.priority_reservation
has_explicit_priority = (
priority is not None
and litellm.priority_reservation is not None
and priority in litellm.priority_reservation
)
if has_explicit_priority and priority is not None:
# Explicit priority: get its specific allocation
priority_weight = normalized_weights.get(priority, self._get_priority_weight(priority))
# Use unique key per priority level
priority_key = f"{model}:{priority}"
else:
# No explicit priority: share the default_priority pool with ALL other default keys
priority_weight = litellm.priority_reservation_settings.default_priority
# Use shared key for all default-priority requests
priority_key = f"{model}:default_pool"
return priority_weight, priority_key
async def _check_model_saturation(
self,
model: str,
model_group_info: ModelGroupInfo,
) -> float:
"""
Check current saturation by directly querying v3 limiter's cache keys.
Reuses v3 limiter's Redis-based tracking (works across multiple instances).
Reads counters WITHOUT incrementing them.
Returns:
float: Saturation ratio (0.0 = empty, 1.0 = at capacity, >1.0 = over)
"""
try:
max_saturation = 0.0
# Query RPM saturation
if model_group_info.rpm is not None and model_group_info.rpm > 0:
# Use v3 limiter's key format: {key:value}:rate_limit_type
counter_key = self.v3_limiter.create_rate_limit_keys(
key="model_saturation_check",
value=model,
rate_limit_type="requests",
)
# Query cache for current counter value
counter_value = await self.internal_usage_cache.async_get_cache(
key=counter_key,
litellm_parent_otel_span=None,
local_only=False, # Check Redis too
)
if counter_value is not None:
current_requests = int(counter_value)
rpm_saturation = current_requests / model_group_info.rpm
max_saturation = max(max_saturation, rpm_saturation)
verbose_proxy_logger.debug(
f"Model {model} RPM: {current_requests}/{model_group_info.rpm} "
f"({rpm_saturation:.1%})"
)
# Query TPM saturation
if model_group_info.tpm is not None and model_group_info.tpm > 0:
counter_key = self.v3_limiter.create_rate_limit_keys(
key="model_saturation_check",
value=model,
rate_limit_type="tokens",
)
counter_value = await self.internal_usage_cache.async_get_cache(
key=counter_key,
litellm_parent_otel_span=None,
local_only=False,
)
if counter_value is not None:
current_tokens = float(counter_value)
tpm_saturation = current_tokens / model_group_info.tpm
max_saturation = max(max_saturation, tpm_saturation)
verbose_proxy_logger.debug(
f"Model {model} TPM: {current_tokens}/{model_group_info.tpm} "
f"({tpm_saturation:.1%})"
)
verbose_proxy_logger.debug(
f"Model {model} overall saturation: {max_saturation:.1%}"
)
return max_saturation
except Exception as e:
verbose_proxy_logger.error(
f"Error checking saturation for {model}: {str(e)}"
)
# Fail open: assume not saturated on error
return 0.0
def _create_priority_based_descriptors(
self,
model: str,
@ -64,11 +215,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
priority: Optional[str],
) -> List[RateLimitDescriptor]:
"""
Create rate limit descriptors based on priority and model group limits.
Create rate limit descriptors with normalized priority weights.
This is the key change: instead of calculating dynamic quotas based on active projects,
we create descriptors with priority-adjusted limits and let the v3 limiter handle
the actual rate limiting with its sliding window approach.
Uses normalized weights to handle over-allocation scenarios.
For explicit priorities: each priority gets its own pool (e.g., prod gets 75%)
For default priority: ALL keys without explicit priority share ONE pool (e.g., all share 25%)
"""
descriptors: List[RateLimitDescriptor] = []
@ -79,23 +231,22 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
if model_group_info is None:
return descriptors
# Get priority weight
priority_weight = self._get_priority_weight(priority)
# Create priority-specific rate limits
# Use model:priority as the key to separate different priority levels
priority_key = f"{model}:{priority or 'default'}"
# Get normalized priority weight and pool key
normalized_weights = self._normalize_priority_weights()
priority_weight, priority_key = self._get_priority_allocation(
model=model,
priority=priority,
normalized_weights=normalized_weights,
)
rate_limit_config: RateLimitDescriptorRateLimitObject = {}
# Apply priority weight to model limits
if model_group_info.tpm is not None:
# Reserve portion of TPM based on priority
reserved_tpm = int(model_group_info.tpm * priority_weight)
rate_limit_config["tokens_per_unit"] = reserved_tpm
if model_group_info.rpm is not None:
# Reserve portion of RPM based on priority
reserved_rpm = int(model_group_info.rpm * priority_weight)
rate_limit_config["requests_per_unit"] = reserved_rpm
@ -112,6 +263,187 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
return descriptors
def _create_model_tracking_descriptor(
self,
model: str,
model_group_info: ModelGroupInfo,
high_limit_multiplier: int = 1,
) -> RateLimitDescriptor:
"""
Create a descriptor for tracking model-wide usage.
Args:
model: Model name
model_group_info: Model configuration with RPM/TPM limits
high_limit_multiplier: Multiplier for limits (use >1 for tracking-only)
Returns:
Rate limit descriptor for model-wide tracking
"""
return RateLimitDescriptor(
key="model_saturation_check",
value=model,
rate_limit={
"requests_per_unit": (
model_group_info.rpm * high_limit_multiplier
if model_group_info.rpm else None
),
"tokens_per_unit": (
model_group_info.tpm * high_limit_multiplier
if model_group_info.tpm else None
),
"window_size": self.v3_limiter.window_size,
},
)
async def _check_rate_limits(
self,
model: str,
model_group_info: ModelGroupInfo,
user_api_key_dict: UserAPIKeyAuth,
key_priority: Optional[str],
saturation: float,
data: dict,
) -> None:
"""
Check rate limits using THREE-PHASE approach to prevent partial increments.
Phase 1: Read-only check of ALL limits (no increments)
Phase 2: Decide which limits to enforce based on saturation
Phase 3: Increment ALL counters atomically (model + priority)
This prevents the bug where:
- Model counter increments in stage 1
- Priority check fails in stage 2
- Request blocked but model counter already incremented
Key behaviors:
- All checks performed first (read-only)
- Only increment counters if request will be allowed
- Model capacity: Always enforced at 100%
- Priority limits: Only enforced when saturated >= threshold
- Both counters tracked from first request (accurate accounting)
Args:
model: Model name
model_group_info: Model configuration
user_api_key_dict: User authentication info
key_priority: User's priority level
saturation: Current saturation level
data: Request data dictionary
Raises:
HTTPException: If any limit is exceeded
"""
import json
saturation_threshold = litellm.priority_reservation_settings.saturation_threshold
should_enforce_priority = saturation >= saturation_threshold
# Build ALL descriptors upfront
descriptors_to_check: List[RateLimitDescriptor] = []
# Model-wide descriptor (always enforce)
model_wide_descriptor = self._create_model_tracking_descriptor(
model=model,
model_group_info=model_group_info,
high_limit_multiplier=1,
)
descriptors_to_check.append(model_wide_descriptor)
# Priority descriptors (always track, conditionally enforce)
priority_descriptors = self._create_priority_based_descriptors(
model=model,
user_api_key_dict=user_api_key_dict,
priority=key_priority,
)
if priority_descriptors:
descriptors_to_check.extend(priority_descriptors)
# PHASE 1: Read-only check of ALL limits (no increments)
check_response = await self.v3_limiter.should_rate_limit(
descriptors=descriptors_to_check,
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=True, # CRITICAL: Don't increment counters yet
)
verbose_proxy_logger.debug(f"Read-only check: {json.dumps(check_response, indent=2)}")
# PHASE 2: Decide which limits to enforce
if check_response["overall_code"] == "OVER_LIMIT":
for status in check_response["statuses"]:
if status["code"] == "OVER_LIMIT":
descriptor_key = status["descriptor_key"]
# Model-wide limit exceeded (ALWAYS enforce)
if descriptor_key == "model_saturation_check":
raise HTTPException(
status_code=429,
detail={
"error": f"Model capacity reached for {model}. "
f"Priority: {key_priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}"
},
headers={
"retry-after": str(self.v3_limiter.window_size),
"rate_limit_type": str(status["rate_limit_type"]),
"x-litellm-priority": key_priority or "default",
},
)
# Priority limit exceeded (ONLY enforce when saturated)
elif descriptor_key == "priority_model" and should_enforce_priority:
verbose_proxy_logger.debug(
f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, "
f"priority: {key_priority}"
)
raise HTTPException(
status_code=429,
detail={
"error": f"Priority-based rate limit exceeded. "
f"Priority: {key_priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}, "
f"Model saturation: {saturation:.1%}"
},
headers={
"retry-after": str(self.v3_limiter.window_size),
"rate_limit_type": str(status["rate_limit_type"]),
"x-litellm-priority": key_priority or "default",
"x-litellm-saturation": f"{saturation:.2%}",
},
)
# PHASE 3: Increment counters separately to avoid early-exit issues
# Model counter must ALWAYS increment, but priority counter might be over limit
# If we increment them together, v3_limiter's in-memory check will exit early
# and skip incrementing the model counter
# Step 3a: Increment model-wide counter (always)
model_increment_response = await self.v3_limiter.should_rate_limit(
descriptors=[model_wide_descriptor],
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=False,
)
# Step 3b: Increment priority counter (may be over limit, but we still track it)
if priority_descriptors:
priority_increment_response = await self.v3_limiter.should_rate_limit(
descriptors=priority_descriptors,
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=False,
)
# Combine responses for post-call hook
combined_response = {
"overall_code": model_increment_response["overall_code"],
"statuses": model_increment_response["statuses"] + priority_increment_response["statuses"]
}
data["litellm_proxy_rate_limit_response"] = combined_response
else:
data["litellm_proxy_rate_limit_response"] = model_increment_response
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -130,60 +462,87 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
],
) -> Optional[Union[Exception, str, dict]]:
"""
Pre-call hook using v3 limiter for priority-based rate limiting.
Saturation-aware pre-call hook for priority-based rate limiting.
Flow:
1. Check current saturation level
2. THREE-PHASE rate limit check:
- PHASE 1: Read-only check of ALL limits (no increments)
- PHASE 2: Decide which limits to enforce based on saturation
- PHASE 3: Increment ALL counters atomically if request allowed
This three-phase approach ensures:
- Model capacity is NEVER exceeded (always enforced at 100%)
- Priority usage tracked from first request (accurate metrics)
- Counters only increment when request will be allowed (prevents phantom usage)
- When under-saturated: priorities can borrow unused capacity (generous)
- When saturated: fair allocation based on normalized priority weights (strict)
Example with 100 RPM model, 60% priority allocation, 80% threshold:
- Saturation < 80%: Priority can use up to 100 RPM (model limit enforced only)
- Saturation >= 80%: Priority limited to 60 RPM (both limits enforced)
Prevents bugs where:
- Model counter increments but priority check fails → model over-capacity
- Priority counter increments but not enforced → inaccurate metrics
Args:
user_api_key_dict: User authentication and metadata
cache: Dual cache instance
data: Request data containing model name
call_type: Type of API call being made
Returns:
None if request is allowed, otherwise raises HTTPException
"""
if "model" not in data:
return None
model = data["model"]
key_priority: Optional[str] = user_api_key_dict.metadata.get("priority", None)
# Create priority-based descriptors
descriptors = self._create_priority_based_descriptors(
model=data["model"],
user_api_key_dict=user_api_key_dict,
priority=key_priority,
# Get model configuration
model_group_info: Optional[ModelGroupInfo] = self.llm_router.get_model_group_info(
model_group=model
)
if not descriptors:
verbose_proxy_logger.debug("No rate limit descriptors created, allowing request")
if model_group_info is None:
verbose_proxy_logger.debug(f"No model group info for {model}, allowing request")
return None
try:
# Use v3 limiter to check rate limits
response = await self.v3_limiter.should_rate_limit(
descriptors=descriptors,
parent_otel_span=user_api_key_dict.parent_otel_span,
# STEP 1: Check current saturation level
saturation = await self._check_model_saturation(model, model_group_info)
saturation_threshold = litellm.priority_reservation_settings.saturation_threshold
verbose_proxy_logger.debug(
f"[Dynamic Rate Limiter] Model={model}, Saturation={saturation:.1%}, "
f"Threshold={saturation_threshold:.1%}, Priority={key_priority}"
)
if response["overall_code"] == "OVER_LIMIT":
# Find which descriptor hit the limit
for status in response["statuses"]:
if status["code"] == "OVER_LIMIT":
raise HTTPException(
status_code=429,
detail={
"error": f"Priority-based rate limit exceeded for {status['descriptor_key']}. "
f"Priority: {key_priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}"
},
headers={
"retry-after": str(self.v3_limiter.window_size),
"rate_limit_type": str(status["rate_limit_type"]),
"x-litellm-priority": key_priority or "default",
},
)
else:
# Store response for post-call hook
data["litellm_proxy_rate_limit_response"] = response
data["litellm_model_saturation"] = saturation
# STEP 2: Check rate limits in THREE phases
# Phase 1: Read-only check of ALL limits (no increments)
# Phase 2: Decide which limits to enforce (based on saturation)
# Phase 3: Increment ALL counters only if request will be allowed
# This prevents partial increments and ensures accurate tracking
await self._check_rate_limits(
model=model,
model_group_info=model_group_info,
user_api_key_dict=user_api_key_dict,
key_priority=key_priority,
saturation=saturation,
data=data,
)
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(
f"Error in dynamic rate limiter v3 pre-call hook: {str(e)}"
verbose_proxy_logger.error(
f"Error in dynamic rate limiter: {str(e)}, allowing request"
)
# Allow request to proceed on unexpected errors
# Fail open on unexpected errors
return None
return None

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

@ -413,6 +413,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
Check if any of the rate limit descriptors should be rate limited.
Returns a RateLimitResponse with the overall code and status for each descriptor.
Uses batch operations for Redis to improve performance.
Args:
descriptors: List of rate limit descriptors to check
parent_otel_span: Optional OpenTelemetry span for tracing
read_only: If True, only check limits without incrementing counters
"""
now = datetime.now().timestamp()
@ -485,8 +490,22 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if rate_limit_response["overall_code"] == "OVER_LIMIT":
return rate_limit_response
## IF under limit, check Redis
if self.batch_rate_limiter_script is not None:
## IF under limit in-memory, check Redis
if read_only:
# READ-ONLY MODE: Just read current values without incrementing
cache_values = await self.internal_usage_cache.async_batch_get_cache(
keys=keys_to_fetch,
parent_otel_span=parent_otel_span,
local_only=False, # Check Redis too
)
# For keys that don't exist yet, set them to 0
if cache_values is None:
cache_values = []
for _ in keys_to_fetch:
cache_values.append(str(now_int) if _.endswith(":window") else 0)
elif self.batch_rate_limiter_script is not None:
# NORMAL MODE: Increment counters in Redis
# Group keys by hash tag for Redis cluster compatibility
cache_values = await self._execute_redis_batch_rate_limiter_script(
keys_to_fetch=keys_to_fetch,
@ -514,6 +533,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
local_only=True,
)
else:
# NORMAL MODE: In-memory sliding window (no Redis)
cache_values = await self.in_memory_cache_sliding_window(
keys=keys_to_fetch,
now_int=now_int,

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

@ -5,7 +5,6 @@
Endpoints here:
- GET `/v1/mcp/server` - Returns all of the configured mcp servers in the db filtered by requestor's access
- GET `/v1/mcp/server/{server_id}` - Returns the the specific mcp server in the db given `server_id` filtered by requestor's access
- GET `/v1/mcp/server/{server_id}/tools` - Get all the tools from the mcp server specified by the `server_id`
- POST `/v1/mcp/server` - Add a new external mcp server.
- PUT `/v1/mcp/server` - Edits an existing mcp server.
- DELETE `/v1/mcp/server/{server_id}` - Deletes the mcp server given `server_id`.

1
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -1594,7 +1594,6 @@ class SSOAuthenticationHandler:
master_key or "",
algorithm="HS256",
)
verbose_proxy_logger.info(f"user_id: {user_id}; jwt_token: {jwt_token}")
if user_id is not None and isinstance(user_id, str):
litellm_dashboard_ui += "?login=success"
verbose_proxy_logger.info(f"Redirecting to {litellm_dashboard_ui}")

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

@ -6,12 +6,14 @@ Provider-specific Pass-Through Endpoints
Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc.
"""
import json
import os
from typing import Optional, cast
import httpx
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket
from fastapi.responses import StreamingResponse
from starlette.websockets import WebSocketState
import litellm
from litellm._logging import verbose_proxy_logger
@ -19,7 +21,9 @@ from litellm.constants import BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.proxy._types import *
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.auth.user_api_key_auth import (
user_api_key_auth,
)
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
get_form_data,
@ -28,6 +32,8 @@ from litellm.proxy.common_utils.http_parsing_utils import (
from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
create_pass_through_route,
create_websocket_passthrough_route,
websocket_passthrough_request,
)
from litellm.proxy.utils import is_known_model
from litellm.secret_managers.main import get_secret_str
@ -51,9 +57,7 @@ def create_request_copy(request: Request):
}
def is_passthrough_request_using_router_model(
request_body: dict, llm_router: Optional[litellm.Router]
) -> bool:
def is_passthrough_request_using_router_model(request_body: dict, llm_router: Optional[litellm.Router]) -> bool:
"""
Returns True if the model is in the llm_router model names
"""
@ -89,16 +93,12 @@ async def llm_passthrough_factory_proxy_route(
model=None,
)
if provider_config is None:
raise HTTPException(
status_code=404, detail=f"Provider {custom_llm_provider} not found"
)
raise HTTPException(status_code=404, detail=f"Provider {custom_llm_provider} not found")
base_target_url = provider_config.get_api_base()
if base_target_url is None:
raise HTTPException(
status_code=404, detail=f"Provider {custom_llm_provider} api base not found"
)
raise HTTPException(status_code=404, detail=f"Provider {custom_llm_provider} api base not found")
encoded_endpoint = httpx.URL(endpoint).path
@ -143,7 +143,7 @@ async def llm_passthrough_factory_proxy_route(
_request_body = await request.json()
else:
_request_body = await get_form_data(request)
if _request_body.get("stream"):
is_streaming_request = True
@ -177,17 +177,11 @@ async def gemini_proxy_route(
[Docs](https://docs.litellm.ai/docs/pass_through/google_ai_studio)
"""
## CHECK FOR LITELLM API KEY IN THE QUERY PARAMS - ?..key=LITELLM_API_KEY
google_ai_studio_api_key = request.query_params.get("key") or request.headers.get(
"x-goog-api-key"
)
google_ai_studio_api_key = request.query_params.get("key") or request.headers.get("x-goog-api-key")
user_api_key_dict = await user_api_key_auth(
request=request, api_key=f"Bearer {google_ai_studio_api_key}"
)
user_api_key_dict = await user_api_key_auth(request=request, api_key=f"Bearer {google_ai_studio_api_key}")
base_target_url = (
os.getenv("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com"
)
base_target_url = os.getenv("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com"
encoded_endpoint = httpx.URL(endpoint).path
# Ensure endpoint starts with '/' for proper URL construction
@ -220,6 +214,7 @@ async def gemini_proxy_route(
endpoint_func = create_pass_through_route(
endpoint=endpoint,
target=str(updated_url),
custom_llm_provider="gemini",
) # dynamically construct pass-through endpoint based on incoming path
received_value = await endpoint_func(
request,
@ -304,9 +299,7 @@ async def vllm_proxy_route(
from litellm.proxy.proxy_server import llm_router
request_body = await get_request_body(request)
is_router_model = is_passthrough_request_using_router_model(
request_body, llm_router
)
is_router_model = is_passthrough_request_using_router_model(request_body, llm_router)
is_streaming_request = is_passthrough_request_streaming(request_body)
if is_router_model and llm_router:
result = cast(
@ -321,11 +314,7 @@ async def vllm_proxy_route(
content=None,
data=None,
files=None,
json=(
request_body
if request.headers.get("content-type") == "application/json"
else None
),
json=(request_body if request.headers.get("content-type") == "application/json" else None),
params=None,
headers=None,
cookies=None,
@ -503,9 +492,7 @@ async def handle_bedrock_count_tokens(
# Extract model from request body
model = request_body.get("model")
if not model:
raise HTTPException(
status_code=400, detail={"error": "Model is required in request body"}
)
raise HTTPException(status_code=400, detail={"error": "Model is required in request body"})
# Get model parameters from router
litellm_params = {"user_api_key_dict": user_api_key_dict}
@ -544,9 +531,7 @@ async def handle_bedrock_count_tokens(
raise
except Exception as e:
verbose_proxy_logger.error(f"Error in handle_bedrock_count_tokens: {str(e)}")
raise HTTPException(
status_code=500, detail={"error": f"CountTokens processing error: {str(e)}"}
)
raise HTTPException(status_code=500, detail={"error": f"CountTokens processing error: {str(e)}"})
async def bedrock_llm_proxy_route(
@ -598,8 +583,7 @@ async def bedrock_llm_proxy_route(
raise HTTPException(
status_code=400,
detail={
"error": "Model missing from endpoint. Expected format: /model/<Model>/<endpoint>. Got: "
+ endpoint,
"error": "Model missing from endpoint. Expected format: /model/<Model>/<endpoint>. Got: " + endpoint,
},
)
@ -663,9 +647,7 @@ async def bedrock_proxy_route(
aws_region_name = litellm.utils.get_secret(secret_name="AWS_REGION_NAME")
if _is_bedrock_agent_runtime_route(endpoint=endpoint): # handle bedrock agents
base_target_url = (
f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com"
)
base_target_url = f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com"
else:
return await bedrock_llm_proxy_route(
endpoint=endpoint,
@ -686,7 +668,8 @@ async def bedrock_proxy_route(
# Add or update query parameters
from litellm.llms.bedrock.chat import BedrockConverseLLM
credentials: Credentials = BedrockConverseLLM().get_credentials()
bedrock_llm = BedrockConverseLLM()
credentials: Credentials = bedrock_llm.get_credentials() # type: ignore
sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name)
headers = {"Content-Type": "application/json"}
# Assuming the body contains JSON data, parse it
@ -694,9 +677,7 @@ async def bedrock_proxy_route(
data = await request.json()
except Exception as e:
raise HTTPException(status_code=400, detail={"error": e})
_request = AWSRequest(
method="POST", url=str(updated_url), data=json.dumps(data), headers=headers
)
_request = AWSRequest(method="POST", url=str(updated_url), data=json.dumps(data), headers=headers)
sigv4.add_auth(_request)
prepped = _request.prepare()
@ -757,14 +738,8 @@ async def assemblyai_proxy_route(
[Docs](https://api.assemblyai.com)
"""
# Set base URL based on the route
assembly_region = AssemblyAIPassthroughLoggingHandler._get_assembly_region_from_url(
url=str(request.url)
)
base_target_url = (
AssemblyAIPassthroughLoggingHandler._get_assembly_base_url_from_region(
region=assembly_region
)
)
assembly_region = AssemblyAIPassthroughLoggingHandler._get_assembly_region_from_url(url=str(request.url))
base_target_url = AssemblyAIPassthroughLoggingHandler._get_assembly_base_url_from_region(region=assembly_region)
encoded_endpoint = httpx.URL(endpoint).path
# Ensure endpoint starts with '/' for proper URL construction
if not encoded_endpoint.startswith("/"):
@ -822,18 +797,14 @@ async def azure_proxy_route(
"""
base_target_url = get_secret_str(secret_name="AZURE_API_BASE")
if base_target_url is None:
raise Exception(
"Required 'AZURE_API_BASE' in environment to make pass-through calls to Azure."
)
raise Exception("Required 'AZURE_API_BASE' in environment to make pass-through calls to Azure.")
# Add or update query parameters
azure_api_key = passthrough_endpoint_router.get_credentials(
custom_llm_provider=litellm.LlmProviders.AZURE.value,
region_name=None,
)
if azure_api_key is None:
raise Exception(
"Required 'AZURE_API_KEY' in environment to make pass-through calls to Azure."
)
raise Exception("Required 'AZURE_API_KEY' in environment to make pass-through calls to Azure.")
return await BaseOpenAIPassThroughHandler._base_openai_pass_through_handler(
endpoint=endpoint,
@ -857,9 +828,7 @@ class BaseVertexAIPassThroughHandler(ABC):
@staticmethod
@abstractmethod
def update_base_target_url_with_credential_location(
base_target_url: str, vertex_location: Optional[str]
) -> str:
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str:
pass
@ -869,9 +838,7 @@ class VertexAIDiscoveryPassThroughHandler(BaseVertexAIPassThroughHandler):
return "https://discoveryengine.googleapis.com/"
@staticmethod
def update_base_target_url_with_credential_location(
base_target_url: str, vertex_location: Optional[str]
) -> str:
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str:
return base_target_url
@ -881,9 +848,7 @@ class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler):
return get_vertex_base_url(vertex_location)
@staticmethod
def update_base_target_url_with_credential_location(
base_target_url: str, vertex_location: Optional[str]
) -> str:
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str:
return get_vertex_base_url(vertex_location)
@ -949,18 +914,14 @@ async def _base_vertex_proxy_route(
location=vertex_location,
)
base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(
vertex_location
)
base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
headers_passed_through = False
# Use headers from the incoming request if no vertex credentials are found
if vertex_credentials is None or vertex_credentials.vertex_project is None:
headers = dict(request.headers) or {}
headers_passed_through = True
verbose_proxy_logger.debug(
"default_vertex_config not set, incoming request headers %s", headers
)
verbose_proxy_logger.debug("default_vertex_config not set, incoming request headers %s", headers)
headers.pop("content-length", None)
headers.pop("host", None)
else:
@ -1126,9 +1087,7 @@ async def openai_proxy_route(
region_name=None,
)
if openai_api_key is None:
raise Exception(
"Required 'OPENAI_API_KEY' in environment to make pass-through calls to OpenAI."
)
raise Exception("Required 'OPENAI_API_KEY' in environment to make pass-through calls to OpenAI.")
return await BaseOpenAIPassThroughHandler._base_openai_pass_through_handler(
endpoint=endpoint,
@ -1174,9 +1133,7 @@ class BaseOpenAIPassThroughHandler:
endpoint_func = create_pass_through_route(
endpoint=endpoint,
target=str(updated_url),
custom_headers=BaseOpenAIPassThroughHandler._assemble_headers(
api_key=api_key, request=request
),
custom_headers=BaseOpenAIPassThroughHandler._assemble_headers(api_key=api_key, request=request),
) # dynamically construct pass-through endpoint based on incoming path
received_value = await endpoint_func(
request,
@ -1193,10 +1150,7 @@ class BaseOpenAIPassThroughHandler:
"""
Appends the OpenAI-Beta header to the headers if the request is an OpenAI Assistants API request
"""
if (
RouteChecks._is_assistants_api_request(request) is True
and "OpenAI-Beta" not in headers
):
if RouteChecks._is_assistants_api_request(request) is True and "OpenAI-Beta" not in headers:
headers["OpenAI-Beta"] = "assistants=v2"
return headers
@ -1212,9 +1166,7 @@ class BaseOpenAIPassThroughHandler:
)
@staticmethod
def _join_url_paths(
base_url: httpx.URL, path: str, custom_llm_provider: litellm.LlmProviders
) -> str:
def _join_url_paths(base_url: httpx.URL, path: str, custom_llm_provider: litellm.LlmProviders) -> str:
"""
Properly joins a base URL with a path, preserving any existing path in the base URL.
"""
@ -1230,13 +1182,192 @@ class BaseOpenAIPassThroughHandler:
joined_path_str = str(base_url.copy_with(path=full_path))
# Apply OpenAI-specific path handling for both branches
if (
custom_llm_provider == litellm.LlmProviders.OPENAI
and "/v1/" not in joined_path_str
):
if custom_llm_provider == litellm.LlmProviders.OPENAI and "/v1/" not in joined_path_str:
# Insert v1 after api.openai.com for OpenAI requests
joined_path_str = joined_path_str.replace(
"api.openai.com/", "api.openai.com/v1/"
)
joined_path_str = joined_path_str.replace("api.openai.com/", "api.openai.com/v1/")
return joined_path_str
async def vertex_ai_live_websocket_passthrough(
websocket: WebSocket,
model: Optional[str] = None,
vertex_project: Optional[str] = None,
vertex_location: Optional[str] = None,
user_api_key_dict: Optional[UserAPIKeyAuth] = None,
):
"""
Vertex AI Live API WebSocket Pass-through Function
This function provides WebSocket passthrough functionality for Vertex AI Live API,
allowing real-time communication with Google's Live API service.
Note: This function should be registered in proxy_server.py using:
app.websocket("/vertex_ai/live")(vertex_ai_live_websocket_passthrough)
"""
from litellm.proxy.proxy_server import proxy_logging_obj
_ = user_api_key_dict # passthrough route already authenticated; avoid lint warnings
await websocket.accept()
incoming_headers = dict(websocket.headers)
vertex_credentials_config = passthrough_endpoint_router.get_vertex_credentials(
project_id=vertex_project,
location=vertex_location,
)
if vertex_credentials_config is None:
# Attempt to load defaults from environment/config if not already initialised
passthrough_endpoint_router.set_default_vertex_config()
vertex_credentials_config = passthrough_endpoint_router.get_vertex_credentials(
project_id=vertex_project,
location=vertex_location,
)
resolved_project = vertex_project
resolved_location: Optional[str] = vertex_location
credentials_value: Optional[str] = None
if vertex_credentials_config is not None:
resolved_project = resolved_project or vertex_credentials_config.vertex_project
temp_location = (
resolved_location or vertex_credentials_config.vertex_location
)
# Ensure resolved_location is a string
if isinstance(temp_location, dict):
resolved_location = str(temp_location)
elif temp_location is not None:
resolved_location = str(temp_location)
else:
resolved_location = None
credentials_value = str(vertex_credentials_config.vertex_credentials) if vertex_credentials_config.vertex_credentials is not None else None
try:
resolved_location = resolved_location or (
vertex_llm_base.get_default_vertex_location()
)
if model:
resolved_location = vertex_llm_base.get_vertex_region(
vertex_region=resolved_location,
model=model,
)
(
access_token,
resolved_project,
) = await vertex_llm_base._ensure_access_token_async(
credentials=credentials_value,
project_id=resolved_project,
custom_llm_provider="vertex_ai_beta",
)
except Exception as e:
verbose_proxy_logger.exception(
"Failed to prepare Vertex AI credentials for live passthrough"
)
# Log the authentication failure using proxy_logging_obj
if proxy_logging_obj and user_api_key_dict:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data={},
)
if websocket.client_state != WebSocketState.DISCONNECTED:
await websocket.close(code=1011, reason="Vertex AI authentication failed")
return
host_location = resolved_location or vertex_llm_base.get_default_vertex_location()
host = (
"aiplatform.googleapis.com"
if host_location == "global"
else f"{host_location}-aiplatform.googleapis.com"
)
service_url = (
f"wss://{host}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
)
upstream_headers = {
"Authorization": f"Bearer {access_token}",
"Content-Type": "application/json",
}
if resolved_project:
upstream_headers["x-goog-user-project"] = resolved_project
# Forward any custom x-goog-* headers provided by the caller if we haven't overridden them
for header_name, header_value in incoming_headers.items():
lower_header = header_name.lower()
if lower_header.startswith("x-goog-") and header_name not in upstream_headers:
upstream_headers[header_name] = header_value
# Use the new WebSocket passthrough pattern
if user_api_key_dict is None:
raise ValueError("user_api_key_dict is required for WebSocket passthrough")
return await websocket_passthrough_request(
websocket=websocket,
target=service_url,
custom_headers=upstream_headers,
user_api_key_dict=user_api_key_dict,
forward_headers=False,
endpoint="/vertex_ai/live",
accept_websocket=False,
)
def create_vertex_ai_live_websocket_endpoint():
"""
Create a Vertex AI Live WebSocket endpoint using the new passthrough pattern.
This demonstrates how to use the create_websocket_passthrough_route function
for a provider-specific WebSocket endpoint.
"""
# This would be used like:
# endpoint_func = create_vertex_ai_live_websocket_endpoint()
# app.websocket("/vertex_ai/live")(endpoint_func)
# For now, we'll keep the existing implementation since it has
# provider-specific logic for Vertex AI credentials and headers
return vertex_ai_live_websocket_passthrough
def create_generic_websocket_passthrough_endpoint(
provider: str,
target_url: str,
custom_headers: Optional[dict] = None,
forward_headers: bool = False,
cost_per_request: Optional[float] = None,
):
"""
Create a generic WebSocket passthrough endpoint for any provider.
This demonstrates the new WebSocket passthrough pattern that's similar to
the HTTP create_pass_through_route function.
Args:
provider: The provider name (e.g., "anthropic", "cohere")
target_url: The target WebSocket URL
custom_headers: Custom headers to include
forward_headers: Whether to forward incoming headers
Returns:
A WebSocket endpoint function that can be registered with app.websocket()
Example usage:
# Create a WebSocket endpoint for Anthropic
anthropic_ws_func = create_generic_websocket_passthrough_endpoint(
provider="anthropic",
target_url="wss://api.anthropic.com/v1/ws",
custom_headers={"x-api-key": "your-api-key"},
forward_headers=True
)
# Register it in proxy_server.py
app.websocket("/anthropic/ws")(anthropic_ws_func)
"""
return create_websocket_passthrough_route(
endpoint=f"/{provider}/ws",
target=target_url,
custom_headers=custom_headers,
_forward_headers=forward_headers,
cost_per_request=cost_per_request,
)

5
litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
View file

@ -20,7 +20,7 @@ from litellm.types.utils import ModelResponse, TextCompletionResponse
if TYPE_CHECKING:
from ..success_handler import PassThroughEndpointLogging
from ..types import EndpointType
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
else:
PassThroughEndpointLogging = Any
EndpointType = Any
@ -228,6 +228,7 @@ class AnthropicPassthroughLoggingHandler:
except (StopIteration, StopAsyncIteration):
break
complete_streaming_response = litellm.stream_chunk_builder(
chunks=all_openai_chunks
chunks=all_openai_chunks,
logging_obj=litellm_logging_obj,
)
return complete_streaming_response

265
litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py Normal file
View file

@ -0,0 +1,265 @@
import json
import re
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator as GeminiModelResponseIterator,
)
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.types.utils import (
ModelResponse,
TextCompletionResponse,
)
if TYPE_CHECKING:
from ..success_handler import PassThroughEndpointLogging
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
else:
PassThroughEndpointLogging = Any
EndpointType = Any
class GeminiPassthroughLoggingHandler:
@staticmethod
def gemini_passthrough_handler(
httpx_response: httpx.Response,
response_body: dict,
logging_obj: LiteLLMLoggingObj,
url_route: str,
result: str,
start_time: datetime,
end_time: datetime,
cache_hit: bool,
request_body: dict,
**kwargs,
) -> PassThroughEndpointLoggingTypedDict:
if "generateContent" in url_route:
model = GeminiPassthroughLoggingHandler.extract_model_from_url(url_route)
# Use Gemini config for transformation
instance_of_gemini_llm = litellm.GoogleAIStudioGeminiConfig()
litellm_model_response: ModelResponse = instance_of_gemini_llm.transform_response(
model=model,
messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}],
raw_response=httpx_response,
model_response=litellm.ModelResponse(),
logging_obj=logging_obj,
optional_params={},
litellm_params={},
api_key="",
request_data={},
encoding=litellm.encoding,
)
kwargs = GeminiPassthroughLoggingHandler._create_gemini_response_logging_payload_for_generate_content(
litellm_model_response=litellm_model_response,
model=model,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
logging_obj=logging_obj,
custom_llm_provider="gemini",
)
return {
"result": litellm_model_response,
"kwargs": kwargs,
}
else:
return {
"result": None,
"kwargs": kwargs,
}
@staticmethod
def _handle_logging_gemini_collected_chunks(
litellm_logging_obj: LiteLLMLoggingObj,
passthrough_success_handler_obj: PassThroughEndpointLogging,
url_route: str,
request_body: dict,
endpoint_type: EndpointType,
start_time: datetime,
all_chunks: List[str],
model: Optional[str],
end_time: datetime,
) -> PassThroughEndpointLoggingTypedDict:
"""
Takes raw chunks from Gemini passthrough endpoint and logs them in litellm callbacks
- Builds complete response from chunks
- Creates standard logging object
- Logs in litellm callbacks
"""
kwargs: Dict[str, Any] = {}
model = model or GeminiPassthroughLoggingHandler.extract_model_from_url(url_route)
complete_streaming_response = GeminiPassthroughLoggingHandler._build_complete_streaming_response(
all_chunks=all_chunks,
litellm_logging_obj=litellm_logging_obj,
model=model,
url_route=url_route,
)
if complete_streaming_response is None:
verbose_proxy_logger.error(
"Unable to build complete streaming response for Gemini passthrough endpoint, not logging..."
)
return {
"result": None,
"kwargs": kwargs,
}
kwargs = GeminiPassthroughLoggingHandler._create_gemini_response_logging_payload_for_generate_content(
litellm_model_response=complete_streaming_response,
model=model,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
logging_obj=litellm_logging_obj,
custom_llm_provider="gemini",
)
return {
"result": complete_streaming_response,
"kwargs": kwargs,
}
@staticmethod
def _parse_gemini_streaming_json(
all_chunks: List[str],
gemini_iterator: GeminiModelResponseIterator,
) -> List[Any]:
"""
Parse Gemini streaming chunks from fragmented JSON strings.
Gemini's streaming format sends a single JSON array that may be split across
multiple lines. This method:
1. Joins all fragmented string chunks into complete JSON
2. Parses the JSON array/object
3. Transforms each item using Gemini's chunk_parser
Args:
all_chunks: Raw string chunks from the streaming response
gemini_iterator: GeminiModelResponseIterator instance for parsing
Returns:
List of parsed chunks in OpenAI format, or empty list if parsing fails
"""
parsed_chunks = []
verbose_proxy_logger.debug(f"Gemini streaming: Processing {len(all_chunks)} raw chunks")
# Gemini streaming response is a single JSON array that may be split across lines
# Join all chunks back together to reconstruct the complete JSON
combined_chunk = "".join(all_chunks)
# Parse the combined JSON string
try:
dict_chunk = json.loads(combined_chunk)
verbose_proxy_logger.debug(f"Parsed JSON object: {type(dict_chunk)}")
# Gemini returns an array of response objects
if isinstance(dict_chunk, list):
for item in dict_chunk:
try:
# Call chunk_parser directly with the dict, not _common_chunk_parsing_logic
parsed_chunk = gemini_iterator.chunk_parser(chunk=item)
if parsed_chunk is not None:
parsed_chunks.append(parsed_chunk)
except Exception as e:
verbose_proxy_logger.error(f"Error parsing Gemini chunk item: {e}", exc_info=True)
continue
else:
# Single object response
parsed_chunk = gemini_iterator.chunk_parser(chunk=dict_chunk)
if parsed_chunk is not None:
parsed_chunks.append(parsed_chunk)
except json.JSONDecodeError as e:
verbose_proxy_logger.error(f"Failed to parse Gemini streaming response as JSON: {e}")
return []
verbose_proxy_logger.debug(f"Total parsed chunks: {len(parsed_chunks)}")
return parsed_chunks
@staticmethod
def _build_complete_streaming_response(
all_chunks: List[str],
litellm_logging_obj: LiteLLMLoggingObj,
model: str,
url_route: str,
) -> Optional[Union[ModelResponse, TextCompletionResponse]]:
if "generateContent" not in url_route and "streamGenerateContent" not in url_route:
return None
gemini_iterator: Any = GeminiModelResponseIterator(
streaming_response=None,
sync_stream=False,
logging_obj=litellm_logging_obj,
)
# Parse the streaming chunks
parsed_chunks = GeminiPassthroughLoggingHandler._parse_gemini_streaming_json(
all_chunks=all_chunks,
gemini_iterator=gemini_iterator,
)
all_openai_chunks = []
for parsed_chunk in parsed_chunks:
if parsed_chunk is None:
continue
all_openai_chunks.append(parsed_chunk)
complete_streaming_response = litellm.stream_chunk_builder(
chunks=all_openai_chunks,
logging_obj=litellm_logging_obj,
)
return complete_streaming_response
@staticmethod
def extract_model_from_url(url: str) -> str:
pattern = r"/models/([^:]+)"
match = re.search(pattern, url)
if match:
return match.group(1)
return "unknown"
@staticmethod
def _create_gemini_response_logging_payload_for_generate_content(
litellm_model_response: Union[ModelResponse, TextCompletionResponse],
model: str,
kwargs: dict,
start_time: datetime,
end_time: datetime,
logging_obj: LiteLLMLoggingObj,
custom_llm_provider: str,
):
"""
Create the standard logging object for Gemini passthrough generateContent (streaming and non-streaming)
"""
response_cost = litellm.completion_cost(
completion_response=litellm_model_response,
model=model,
custom_llm_provider=custom_llm_provider,
)
kwargs["response_cost"] = response_cost
kwargs["model"] = model
kwargs["custom_llm_provider"] = custom_llm_provider
# pretty print standard logging object
verbose_proxy_logger.debug("kwargs= %s", json.dumps(kwargs, indent=4))
# set litellm_call_id to logging response object
litellm_model_response.id = logging_obj.litellm_call_id
logging_obj.model = litellm_model_response.model or model
logging_obj.model_call_details["model"] = logging_obj.model
logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider
logging_obj.model_call_details["response_cost"] = response_cost
return kwargs

398
litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py Normal file
View file

@ -0,0 +1,398 @@
"""
Vertex AI Live API WebSocket Passthrough Logging Handler
Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough endpoints.
Supports different modalities: text, audio, video, and web search.
"""
from datetime import datetime
from typing import Any, Dict, List, Optional
from litellm._logging import verbose_proxy_logger
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import (
BasePassthroughLoggingHandler,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
PassThroughEndpointLoggingTypedDict,
)
from litellm.types.utils import LlmProviders, ModelResponse, Usage
from litellm.utils import get_model_info
class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
"""
Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough.
Supports:
- Text tokens (input/output)
- Audio tokens (input/output)
- Video tokens (input/output)
- Web search requests
- Tool use tokens
"""
def _build_complete_streaming_response(self, *args, **kwargs):
"""Not applicable for WebSocket passthrough."""
return None
def get_provider_config(self, model: str):
"""Return Vertex AI provider configuration."""
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexGeminiConfig,
)
return VertexGeminiConfig()
@property
def llm_provider_name(self) -> LlmProviders:
"""Return the LLM provider name."""
return LlmProviders.VERTEX_AI
@staticmethod
def _extract_usage_metadata_from_websocket_messages(
websocket_messages: List[Dict],
) -> Optional[Dict]:
"""
Extract and aggregate usage metadata from a list of WebSocket messages.
Args:
websocket_messages: List of WebSocket messages from the Live API
Returns:
Dictionary containing aggregated usage metadata, or None if not found
"""
all_usage_metadata = []
# Collect all usage metadata messages
for message in websocket_messages:
if isinstance(message, dict) and "usageMetadata" in message:
all_usage_metadata.append(message["usageMetadata"])
if not all_usage_metadata:
return None
# If only one usage metadata, return it as-is
if len(all_usage_metadata) == 1:
return all_usage_metadata[0]
# Aggregate multiple usage metadata messages
aggregated: Dict[str, Any] = {
"promptTokenCount": 0,
"candidatesTokenCount": 0,
"totalTokenCount": 0,
"promptTokensDetails": [],
"candidatesTokensDetails": [],
}
# Aggregate token counts
for usage in all_usage_metadata:
aggregated["promptTokenCount"] += usage.get("promptTokenCount", 0)
aggregated["candidatesTokenCount"] += usage.get("candidatesTokenCount", 0)
aggregated["totalTokenCount"] += usage.get("totalTokenCount", 0)
# Aggregate token details by modality
modality_totals = {}
for usage in all_usage_metadata:
# Process prompt tokens details
for detail in usage.get("promptTokensDetails", []):
modality = detail.get("modality", "TEXT")
token_count = detail.get("tokenCount", 0)
if modality not in modality_totals:
modality_totals[modality] = {"prompt": 0, "candidate": 0}
modality_totals[modality]["prompt"] += token_count
# Process candidate tokens details
for detail in usage.get("candidatesTokensDetails", []):
modality = detail.get("modality", "TEXT")
token_count = detail.get("tokenCount", 0)
if modality not in modality_totals:
modality_totals[modality] = {"prompt": 0, "candidate": 0}
modality_totals[modality]["candidate"] += token_count
# Convert aggregated modality totals back to details format
for modality, totals in modality_totals.items():
if totals["prompt"] > 0:
aggregated["promptTokensDetails"].append(
{"modality": modality, "tokenCount": totals["prompt"]}
)
if totals["candidate"] > 0:
aggregated["candidatesTokensDetails"].append(
{"modality": modality, "tokenCount": totals["candidate"]}
)
# Add any additional fields from the first usage metadata
first_usage = all_usage_metadata[0]
for key, value in first_usage.items():
if key not in aggregated:
aggregated[key] = value
return aggregated
@staticmethod
def _calculate_live_api_cost(
model: str,
usage_metadata: Dict,
custom_llm_provider: str = "vertex_ai",
) -> float:
"""
Calculate cost for Vertex AI Live API based on usage metadata.
Args:
model: The model name (e.g., "gemini-2.0-flash-live-preview-04-09")
usage_metadata: Usage metadata from the Live API response
custom_llm_provider: The LLM provider (default: "vertex_ai")
Returns:
Total cost in USD
"""
try:
# Get model pricing information
model_info = get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
verbose_proxy_logger.debug(
f"Vertex AI Live API model info for '{model}': {model_info}"
)
# Check if pricing info is available
if not model_info or not model_info.get("input_cost_per_token"):
verbose_proxy_logger.error(
f"No pricing info found for {model} in local model pricing database"
)
return 0.0
total_cost = 0.0
# Extract token counts from usage metadata
prompt_token_count = usage_metadata.get("promptTokenCount", 0)
candidates_token_count = usage_metadata.get("candidatesTokenCount", 0)
# Calculate base text token costs
input_cost_per_token = model_info.get("input_cost_per_token", 0.0)
output_cost_per_token = model_info.get("output_cost_per_token", 0.0)
total_cost += prompt_token_count * input_cost_per_token
total_cost += candidates_token_count * output_cost_per_token
# Handle modality-specific costs if present
prompt_tokens_details = usage_metadata.get("promptTokensDetails", [])
candidates_tokens_details = usage_metadata.get(
"candidatesTokensDetails", []
)
# Process prompt tokens by modality
for detail in prompt_tokens_details:
modality = detail.get("modality", "TEXT")
token_count = detail.get("tokenCount", 0)
if modality == "AUDIO":
audio_cost_per_token = model_info.get(
"input_cost_per_audio_token", 0.0
)
total_cost += token_count * audio_cost_per_token
elif modality == "VIDEO":
# Video tokens are typically per second, but we'll treat as per token for now
video_cost_per_token = model_info.get(
"input_cost_per_video_per_second", 0.0
)
total_cost += token_count * video_cost_per_token
# TEXT tokens are already handled above
# Process candidate tokens by modality
for detail in candidates_tokens_details:
modality = detail.get("modality", "TEXT")
token_count = detail.get("tokenCount", 0)
if modality == "AUDIO":
audio_cost_per_token = model_info.get(
"output_cost_per_audio_token", 0.0
)
total_cost += token_count * audio_cost_per_token
elif modality == "VIDEO":
# Video tokens are typically per second, but we'll treat as per token for now
video_cost_per_token = model_info.get(
"output_cost_per_video_per_second", 0.0
)
total_cost += token_count * video_cost_per_token
# TEXT tokens are already handled above
# Handle web search costs if present
tool_use_prompt_token_count = usage_metadata.get(
"toolUsePromptTokenCount", 0
)
if tool_use_prompt_token_count > 0:
# Web search typically has a fixed cost per request
web_search_cost = model_info.get("web_search_cost_per_request", 0.0)
if isinstance(web_search_cost, (int, float)) and web_search_cost > 0:
total_cost += web_search_cost
else:
# Fallback to token-based pricing for tool use
total_cost += tool_use_prompt_token_count * input_cost_per_token
verbose_proxy_logger.debug(
f"Vertex AI Live API cost calculation - Model: {model}, "
f"Prompt tokens: {prompt_token_count}, "
f"Candidate tokens: {candidates_token_count}, "
f"Total cost: ${total_cost:.6f}"
)
return total_cost
except Exception as e:
verbose_proxy_logger.error(
f"Error calculating Vertex AI Live API cost: {e}"
)
return 0.0
@staticmethod
def _create_usage_object_from_metadata(
usage_metadata: Dict,
model: str,
) -> Usage:
"""
Create a LiteLLM Usage object from Live API usage metadata.
Args:
usage_metadata: Usage metadata from the Live API response
model: The model name
Returns:
LiteLLM Usage object
"""
prompt_tokens = usage_metadata.get("promptTokenCount", 0)
completion_tokens = usage_metadata.get("candidatesTokenCount", 0)
total_tokens = usage_metadata.get("totalTokenCount", 0)
# Create modality-specific token details if available
prompt_tokens_details = usage_metadata.get("promptTokensDetails", [])
candidates_tokens_details = usage_metadata.get("candidatesTokensDetails", [])
# Extract text tokens from details
text_prompt_tokens = 0
text_completion_tokens = 0
for detail in prompt_tokens_details:
if detail.get("modality") == "TEXT":
text_prompt_tokens = detail.get("tokenCount", 0)
break
for detail in candidates_tokens_details:
if detail.get("modality") == "TEXT":
text_completion_tokens = detail.get("tokenCount", 0)
break
# If no text tokens found in details, use total counts
if text_prompt_tokens == 0:
text_prompt_tokens = prompt_tokens
if text_completion_tokens == 0:
text_completion_tokens = completion_tokens
return Usage(
prompt_tokens=text_prompt_tokens,
completion_tokens=text_completion_tokens,
total_tokens=total_tokens,
)
def vertex_ai_live_passthrough_handler(
self,
websocket_messages: List[Dict],
logging_obj,
url_route: str,
start_time: datetime,
end_time: datetime,
request_body: dict,
**kwargs,
) -> PassThroughEndpointLoggingTypedDict:
"""
Handle cost tracking and logging for Vertex AI Live API WebSocket passthrough.
Args:
websocket_messages: List of WebSocket messages from the Live API
logging_obj: LiteLLM logging object
url_route: The URL route that was called
start_time: Request start time
end_time: Request end time
request_body: The original request body
**kwargs: Additional keyword arguments
Returns:
Dictionary containing the result and kwargs for logging
"""
try:
# Extract model from request body or kwargs
model = kwargs.get("model", "gemini-2.0-flash-live-preview-04-09")
custom_llm_provider = kwargs.get("custom_llm_provider", "vertex_ai")
verbose_proxy_logger.debug(
f"Vertex AI Live API model: {model}, custom_llm_provider: {custom_llm_provider}"
)
# Extract usage metadata from WebSocket messages
usage_metadata = self._extract_usage_metadata_from_websocket_messages(
websocket_messages
)
if not usage_metadata:
verbose_proxy_logger.warning(
"No usage metadata found in Vertex AI Live API WebSocket messages"
)
return {
"result": None,
"kwargs": kwargs,
}
# Calculate cost using Live API specific pricing
response_cost = self._calculate_live_api_cost(
model=model,
usage_metadata=usage_metadata,
custom_llm_provider=custom_llm_provider,
)
# Create Usage object for standard LiteLLM logging
usage = self._create_usage_object_from_metadata(
usage_metadata=usage_metadata,
model=model,
)
# Create a mock ModelResponse for standard logging
litellm_model_response = ModelResponse(
id=f"vertex-ai-live-{start_time.timestamp()}",
object="chat.completion",
created=int(start_time.timestamp()),
model=model,
usage=usage,
choices=[],
)
# Update kwargs with cost information
kwargs["response_cost"] = response_cost
kwargs["model"] = model
kwargs["custom_llm_provider"] = custom_llm_provider
# Safely log the model name: only allow known safe formats, redact otherwise.
import re
allowed_pattern = re.compile(r"^[A-Za-z0-9._\-:]+$")
safe_model = model if isinstance(model, str) and allowed_pattern.match(model) else "[REDACTED]"
verbose_proxy_logger.debug(
f"Vertex AI Live API passthrough cost tracking - "
f"Model: {safe_model}, Cost: ${response_cost:.6f}, "
f"Prompt tokens: {usage.prompt_tokens}, "
f"Completion tokens: {usage.completion_tokens}"
)
return {
"result": litellm_model_response,
"kwargs": kwargs,
}
except Exception as e:
verbose_proxy_logger.error(
f"Error in Vertex AI Live API passthrough handler: {e}"
)
return {
"result": None,
"kwargs": kwargs,
}

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