Merge branch 'main' into litellm_responses_structured_output

This commit is contained in:
Sameer Kankute 2025-09-04 22:35:30 +05:30 committed by GitHub
commit 8d67392e99
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
60 changed files with 4331 additions and 495 deletions

View file

@ -1913,7 +1913,7 @@ jobs:
-e APORIA_API_BASE_1=$APORIA_API_BASE_1 \
-e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \
-e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \
-e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1
-e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \
-e USE_DDTRACE=True \
-e DD_API_KEY=$DD_API_KEY \
-e DD_SITE=$DD_SITE \

View file

@ -0,0 +1,311 @@
#!/usr/bin/env python3
"""
Complete example for Veo video generation through LiteLLM proxy.
This script demonstrates how to:
1. Generate videos using Google's Veo model
2. Poll for completion status
3. Download the generated video file
Requirements:
- LiteLLM proxy running with Google AI Studio pass-through configured
- Google AI Studio API key with Veo access
"""
import json
import os
import time
import requests
from typing import Optional
class VeoVideoGenerator:
"""Complete Veo video generation client using LiteLLM proxy."""
def __init__(self, base_url: str = "http://localhost:4000/gemini/v1beta",
api_key: str = "sk-1234"):
"""
Initialize the Veo video generator.
Args:
base_url: Base URL for the LiteLLM proxy with Gemini pass-through
api_key: API key for LiteLLM proxy authentication
"""
self.base_url = base_url
self.api_key = api_key
self.headers = {
"x-goog-api-key": api_key,
"Content-Type": "application/json"
}
def generate_video(self, prompt: str) -> Optional[str]:
"""
Initiate video generation with Veo.
Args:
prompt: Text description of the video to generate
Returns:
Operation name if successful, None otherwise
"""
print(f"🎬 Generating video with prompt: '{prompt}'")
url = f"{self.base_url}/models/veo-3.0-generate-preview:predictLongRunning"
payload = {
"instances": [{
"prompt": prompt
}]
}
try:
response = requests.post(url, headers=self.headers, json=payload)
response.raise_for_status()
data = response.json()
operation_name = data.get("name")
if operation_name:
print(f"✅ Video generation started: {operation_name}")
return operation_name
else:
print("❌ No operation name returned")
print(f"Response: {json.dumps(data, indent=2)}")
return None
except requests.RequestException as e:
print(f"❌ Failed to start video generation: {e}")
if hasattr(e, 'response') and e.response is not None:
try:
error_data = e.response.json()
print(f"Error details: {json.dumps(error_data, indent=2)}")
except:
print(f"Error response: {e.response.text}")
return None
def wait_for_completion(self, operation_name: str, max_wait_time: int = 600) -> Optional[str]:
"""
Poll operation status until video generation is complete.
Args:
operation_name: Name of the operation to monitor
max_wait_time: Maximum time to wait in seconds (default: 10 minutes)
Returns:
Video URI if successful, None otherwise
"""
print("⏳ Waiting for video generation to complete...")
operation_url = f"{self.base_url}/{operation_name}"
start_time = time.time()
poll_interval = 10 # Start with 10 seconds
while time.time() - start_time < max_wait_time:
try:
print(f"🔍 Polling status... ({int(time.time() - start_time)}s elapsed)")
response = requests.get(operation_url, headers=self.headers)
response.raise_for_status()
data = response.json()
# Check for errors
if "error" in data:
print("❌ Error in video generation:")
print(json.dumps(data["error"], indent=2))
return None
# Check if operation is complete
is_done = data.get("done", False)
if is_done:
print("🎉 Video generation complete!")
try:
# Extract video URI from nested response
video_uri = data["response"]["generateVideoResponse"]["generatedSamples"][0]["video"]["uri"]
print(f"📹 Video URI: {video_uri}")
return video_uri
except KeyError as e:
print(f"❌ Could not extract video URI: {e}")
print("Full response:")
print(json.dumps(data, indent=2))
return None
# Wait before next poll, with exponential backoff
time.sleep(poll_interval)
poll_interval = min(poll_interval * 1.2, 30) # Cap at 30 seconds
except requests.RequestException as e:
print(f"❌ Error polling operation status: {e}")
time.sleep(poll_interval)
print(f"⏰ Timeout after {max_wait_time} seconds")
return None
def download_video(self, video_uri: str, output_filename: str = "generated_video.mp4") -> bool:
"""
Download the generated video file.
Args:
video_uri: URI of the video to download (from Google's response)
output_filename: Local filename to save the video
Returns:
True if download successful, False otherwise
"""
print(f"⬇️ Downloading video...")
print(f"Original URI: {video_uri}")
# Convert Google URI to LiteLLM proxy URI
# Example: files/abc123 -> /gemini/v1beta/files/abc123:download?alt=media
if video_uri.startswith("files/"):
download_path = f"{video_uri}:download?alt=media"
else:
download_path = video_uri
litellm_download_url = f"{self.base_url}/{download_path}"
print(f"Download URL: {litellm_download_url}")
try:
# Download with streaming and redirect handling
response = requests.get(
litellm_download_url,
headers=self.headers,
stream=True,
allow_redirects=True # Handle redirects automatically
)
response.raise_for_status()
# Save video file
with open(output_filename, 'wb') as f:
downloaded_size = 0
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
downloaded_size += len(chunk)
# Progress indicator for large files
if downloaded_size % (1024 * 1024) == 0: # Every MB
print(f"📦 Downloaded {downloaded_size / (1024*1024):.1f} MB...")
# Verify file was created and has content
if os.path.exists(output_filename):
file_size = os.path.getsize(output_filename)
if file_size > 0:
print(f"✅ Video downloaded successfully!")
print(f"📁 Saved as: {output_filename}")
print(f"📏 File size: {file_size / (1024*1024):.2f} MB")
return True
else:
print("❌ Downloaded file is empty")
os.remove(output_filename)
return False
else:
print("❌ File was not created")
return False
except requests.RequestException as e:
print(f"❌ Download failed: {e}")
if hasattr(e, 'response') and e.response is not None:
print(f"Status code: {e.response.status_code}")
print(f"Response headers: {dict(e.response.headers)}")
return False
def generate_and_download(self, prompt: str, output_filename: str = None) -> bool:
"""
Complete workflow: generate video and download it.
Args:
prompt: Text description for video generation
output_filename: Output filename (auto-generated if None)
Returns:
True if successful, False otherwise
"""
# Auto-generate filename if not provided
if output_filename is None:
timestamp = int(time.time())
safe_prompt = "".join(c for c in prompt[:30] if c.isalnum() or c in (' ', '-', '_')).rstrip()
output_filename = f"veo_video_{safe_prompt.replace(' ', '_')}_{timestamp}.mp4"
print("=" * 60)
print("🎬 VEO VIDEO GENERATION WORKFLOW")
print("=" * 60)
# Step 1: Generate video
operation_name = self.generate_video(prompt)
if not operation_name:
return False
# Step 2: Wait for completion
video_uri = self.wait_for_completion(operation_name)
if not video_uri:
return False
# Step 3: Download video
success = self.download_video(video_uri, output_filename)
if success:
print("=" * 60)
print("🎉 SUCCESS! Video generation complete!")
print(f"📁 Video saved as: {output_filename}")
print("=" * 60)
else:
print("=" * 60)
print("❌ FAILED! Video generation or download failed")
print("=" * 60)
return success
def main():
"""
Example usage of the VeoVideoGenerator.
Configure these environment variables:
- LITELLM_BASE_URL: Your LiteLLM proxy URL (default: http://localhost:4000/gemini/v1beta)
- LITELLM_API_KEY: Your LiteLLM API key (default: sk-1234)
"""
# Configuration from environment or defaults
base_url = os.getenv("LITELLM_BASE_URL", "http://localhost:4000/gemini/v1beta")
api_key = os.getenv("LITELLM_API_KEY", "sk-1234")
print("🚀 Starting Veo Video Generation Example")
print(f"📡 Using LiteLLM proxy at: {base_url}")
# Initialize generator
generator = VeoVideoGenerator(base_url=base_url, api_key=api_key)
# Example prompts - try different ones!
example_prompts = [
"A cat playing with a ball of yarn in a sunny garden",
"Ocean waves crashing against rocky cliffs at sunset",
"A bustling city street with people walking and cars passing by",
"A peaceful forest with sunlight filtering through the trees"
]
# Use first example or get from user
prompt = example_prompts[0]
print(f"🎬 Using prompt: '{prompt}'")
# Generate and download video
success = generator.generate_and_download(prompt)
if success:
print("\n✅ Example completed successfully!")
print("💡 Try modifying the prompt in the script for different videos!")
else:
print("\n❌ Example failed!")
print("🔧 Check your LiteLLM proxy configuration and Google AI Studio API key")
# Troubleshooting tips
print("\n🔍 Troubleshooting:")
print("1. Ensure LiteLLM proxy is running with Google AI Studio pass-through")
print("2. Verify your Google AI Studio API key has Veo access")
print("3. Check that your prompt meets Veo's content guidelines")
print("4. Review the LiteLLM proxy logs for detailed error information")
if __name__ == "__main__":
main()

View file

@ -4,9 +4,14 @@
liteLLM provides `input_callbacks`, `success_callbacks` and `failure_callbacks`, making it easy for you to send data to a particular provider depending on the status of your responses.
:::tip
**New to LiteLLM Callbacks?** Check out our comprehensive [Callback Management Guide](./callback_management.md) to understand when to use different callback hooks like `async_log_success_event` vs `async_post_call_success_hook`.
:::
liteLLM supports:
- [Custom Callback Functions](https://docs.litellm.ai/docs/observability/custom_callback)
- [Callback Management Guide](./callback_management.md) - **Comprehensive guide for choosing the right hooks**
- [Lunary](https://lunary.ai/docs)
- [Langfuse](https://langfuse.com/docs)
- [LangSmith](https://www.langchain.com/langsmith)

View file

@ -4,7 +4,6 @@
**For PROXY** [Go Here](../proxy/logging.md#custom-callback-class-async)
:::
## Callback Class
You can create a custom callback class to precisely log events as they occur in litellm.
@ -57,6 +56,17 @@ def async completion():
asyncio.run(completion())
```
## Common Hooks
- `async_log_success_event` - Log successful API calls
- `async_log_failure_event` - Log failed API calls
- `log_pre_api_call` - Log before API call
- `log_post_api_call` - Log after API call
**Proxy-only hooks** (only work with LiteLLM Proxy):
- `async_post_call_success_hook` - Access user data + modify responses
- `async_pre_call_hook` - Modify requests before sending
## Callback Functions
If you just want to log on a specific event (e.g. on input) - you can use callback functions.
@ -174,260 +184,87 @@ async def test_chat_openai():
asyncio.run(test_chat_openai())
```
:::info
## What's Available in kwargs?
We're actively trying to expand this to other event types. [Tell us if you need this!](https://github.com/BerriAI/litellm/issues/1007)
:::
## What's in kwargs?
Notice we pass in a kwargs argument to custom callback.
```python
def custom_callback(
kwargs, # kwargs to completion
completion_response, # response from completion
start_time, end_time # start/end time
):
# Your custom code here
print("LITELLM: in custom callback function")
print("kwargs", kwargs)
print("completion_response", completion_response)
print("start_time", start_time)
print("end_time", end_time)
```
This is a dictionary containing all the model-call details (the params we receive, the values we send to the http endpoint, the response we receive, stacktrace in case of errors, etc.).
This is all logged in the [model_call_details via our Logger](https://github.com/BerriAI/litellm/blob/fc757dc1b47d2eb9d0ea47d6ad224955b705059d/litellm/utils.py#L246).
Here's exactly what you can expect in the kwargs dictionary:
```shell
### DEFAULT PARAMS ###
"model": self.model,
"messages": self.messages,
"optional_params": self.optional_params, # model-specific params passed in
"litellm_params": self.litellm_params, # litellm-specific params passed in (e.g. metadata passed to completion call)
"start_time": self.start_time, # datetime object of when call was started
### PRE-API CALL PARAMS ### (check via kwargs["log_event_type"]="pre_api_call")
"input" = input # the exact prompt sent to the LLM API
"api_key" = api_key # the api key used for that LLM API
"additional_args" = additional_args # any additional details for that API call (e.g. contains optional params sent)
### POST-API CALL PARAMS ### (check via kwargs["log_event_type"]="post_api_call")
"original_response" = original_response # the original http response received (saved via response.text)
### ON-SUCCESS PARAMS ### (check via kwargs["log_event_type"]="successful_api_call")
"complete_streaming_response" = complete_streaming_response # the complete streamed response (only set if `completion(..stream=True)`)
"end_time" = end_time # datetime object of when call was completed
### ON-FAILURE PARAMS ### (check via kwargs["log_event_type"]="failed_api_call")
"exception" = exception # the Exception raised
"traceback_exception" = traceback_exception # the traceback generated via `traceback.format_exc()`
"end_time" = end_time # datetime object of when call was completed
```
### Cache hits
Cache hits are logged in success events as `kwarg["cache_hit"]`.
Here's an example of accessing it:
```python
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm import completion, acompletion, Cache
class MyCustomHandler(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
print(f"On Success")
print(f"Value of Cache hit: {kwargs['cache_hit']"})
async def test_async_completion_azure_caching():
customHandler_caching = MyCustomHandler()
litellm.cache = Cache(type="redis", host=os.environ['REDIS_HOST'], port=os.environ['REDIS_PORT'], password=os.environ['REDIS_PASSWORD'])
litellm.callbacks = [customHandler_caching]
unique_time = time.time()
response1 = await litellm.acompletion(model="azure/chatgpt-v-2",
messages=[{
"role": "user",
"content": f"Hi 👋 - i'm async azure {unique_time}"
}],
caching=True)
await asyncio.sleep(1)
print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}")
response2 = await litellm.acompletion(model="azure/chatgpt-v-2",
messages=[{
"role": "user",
"content": f"Hi 👋 - i'm async azure {unique_time}"
}],
caching=True)
await asyncio.sleep(1) # success callbacks are done in parallel
print(f"customHandler_caching.states post-cache hit: {customHandler_caching.states}")
assert len(customHandler_caching.errors) == 0
assert len(customHandler_caching.states) == 4 # pre, post, success, success
```
### Get complete streaming response
LiteLLM will pass you the complete streaming response in the final streaming chunk as part of the kwargs for your custom callback function.
The kwargs dictionary contains all the details about your API call:
```python
# litellm.set_verbose = False
def custom_callback(
kwargs, # kwargs to completion
completion_response, # response from completion
start_time, end_time # start/end time
):
# print(f"streaming response: {completion_response}")
if "complete_streaming_response" in kwargs:
print(f"Complete Streaming Response: {kwargs['complete_streaming_response']}")
# Assign the custom callback function
litellm.success_callback = [custom_callback]
response = completion(model="claude-instant-1", messages=messages, stream=True)
for idx, chunk in enumerate(response):
pass
```
### Log additional metadata
LiteLLM accepts a metadata dictionary in the completion call. You can pass additional metadata into your completion call via `completion(..., metadata={"key": "value"})`.
Since this is a [litellm-specific param](https://github.com/BerriAI/litellm/blob/b6a015404eed8a0fa701e98f4581604629300ee3/litellm/main.py#L235), it's accessible via kwargs["litellm_params"]
```python
from litellm import completion
import os, litellm
## set ENV variables
os.environ["OPENAI_API_KEY"] = "your-api-key"
messages = [{ "content": "Hello, how are you?","role": "user"}]
def custom_callback(
kwargs, # kwargs to completion
completion_response, # response from completion
start_time, end_time # start/end time
):
print(kwargs["litellm_params"]["metadata"])
def custom_callback(kwargs, completion_response, start_time, end_time):
# Access common data
model = kwargs.get("model")
messages = kwargs.get("messages", [])
cost = kwargs.get("response_cost", 0)
cache_hit = kwargs.get("cache_hit", False)
# Assign the custom callback function
litellm.success_callback = [custom_callback]
response = litellm.completion(model="gpt-3.5-turbo", messages=messages, metadata={"hello": "world"})
# Access metadata you passed in
metadata = kwargs.get("litellm_params", {}).get("metadata", {})
```
## Examples
**Key fields in kwargs:**
- `model` - The model name
- `messages` - Input messages
- `response_cost` - Calculated cost
- `cache_hit` - Whether response was cached
- `litellm_params.metadata` - Your custom metadata
### Custom Callback to track costs for Streaming + Non-Streaming
By default, the response cost is accessible in the logging object via `kwargs["response_cost"]` on success (sync + async)
## Practical Examples
### Track API Costs
```python
def track_cost_callback(kwargs, completion_response, start_time, end_time):
cost = kwargs["response_cost"] # litellm calculates this for you
print(f"Request cost: ${cost}")
# Step 1. Write your custom callback function
def track_cost_callback(
kwargs, # kwargs to completion
completion_response, # response from completion
start_time, end_time # start/end time
):
try:
response_cost = kwargs["response_cost"] # litellm calculates response cost for you
print("regular response_cost", response_cost)
except:
pass
# Step 2. Assign the custom callback function
litellm.success_callback = [track_cost_callback]
# Step 3. Make litellm.completion call
response = completion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
"content": "Hi 👋 - i'm openai"
}
]
)
print(response)
response = completion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello"}])
```
### Custom Callback to log transformed Input to LLMs
### Log Inputs to LLMs
```python
def get_transformed_inputs(
kwargs,
):
def get_transformed_inputs(kwargs):
params_to_model = kwargs["additional_args"]["complete_input_dict"]
print("params to model", params_to_model)
litellm.input_callback = [get_transformed_inputs]
def test_chat_openai():
try:
response = completion(model="claude-2",
messages=[{
"role": "user",
"content": "Hi 👋 - i'm openai"
}])
print(response)
except Exception as e:
print(e)
pass
response = completion(model="claude-2", messages=[{"role": "user", "content": "Hello"}])
```
#### Output
```shell
params to model {'model': 'claude-2', 'prompt': "\n\nHuman: Hi 👋 - i'm openai\n\nAssistant: ", 'max_tokens_to_sample': 256}
### Send to External Service
```python
import requests
def send_to_analytics(kwargs, completion_response, start_time, end_time):
data = {
"model": kwargs.get("model"),
"cost": kwargs.get("response_cost", 0),
"duration": (end_time - start_time).total_seconds()
}
requests.post("https://your-analytics.com/api", json=data)
litellm.success_callback = [send_to_analytics]
```
### Custom Callback to write to Mixpanel
## Common Issues
### Callback Not Called
Make sure you:
1. Register callbacks correctly: `litellm.callbacks = [MyHandler()]`
2. Use the right hook names (check spelling)
3. Don't use proxy-only hooks in library mode
### Performance Issues
- Use async hooks for I/O operations
- Don't block in callback functions
- Handle exceptions properly:
```python
import mixpanel
import litellm
from litellm import completion
def custom_callback(
kwargs, # kwargs to completion
completion_response, # response from completion
start_time, end_time # start/end time
):
# Your custom code here
mixpanel.track("LLM Response", {"llm_response": completion_response})
# Assign the custom callback function
litellm.success_callback = [custom_callback]
response = completion(
model="gpt-3.5-turbo",
messages=[
{
"role": "user",
"content": "Hi 👋 - i'm openai"
}
]
)
print(response)
class SafeHandler(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
try:
await external_service(response_obj)
except Exception as e:
print(f"Callback error: {e}") # Log but don't break the flow
```

View file

@ -230,6 +230,13 @@ curl -X POST "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5
```
## **Example 4: Video Generation with Veo**
Generate videos using Google's Veo model through LiteLLM pass-through routes.
[**→ Complete Veo Video Generation Guide**](../proxy/veo_video_generation.md)
## Advanced
Pre-requisites

View file

@ -11,3 +11,43 @@ These endpoints are useful for 2 scenarios:
## How is your request handled?
The request is passed through to the provider's endpoint. The response is then passed back to the client. **No translation is done.**
### Request Forwarding Process
1. **Request Reception**: LiteLLM receives your request at `/provider/endpoint`
2. **Authentication**: Your LiteLLM API key is validated and mapped to the provider's API key
3. **Request Transformation**: Request is reformatted for the target provider's API
4. **Forwarding**: Request is sent to the actual provider endpoint
5. **Response Handling**: Provider response is returned directly to you
### Authentication Flow
```mermaid
graph LR
A[Client Request] --> B[LiteLLM Proxy]
B --> C[Validate LiteLLM API Key]
C --> D[Map to Provider API Key]
D --> E[Forward to Provider]
E --> F[Return Response]
```
**Key Points:**
- Use your **LiteLLM API key** in requests, not the provider's key
- LiteLLM handles the provider authentication internally
- Same authentication works across all passthrough endpoints
### Error Handling
**Provider Errors**: Forwarded directly to you with original error codes and messages
**LiteLLM Errors**:
- `401`: Invalid LiteLLM API key
- `404`: Provider or endpoint not supported
- `500`: Internal routing/forwarding errors
### Benefits
- **Unified Authentication**: One API key for all providers
- **Centralized Logging**: All requests logged through LiteLLM
- **Cost Tracking**: Usage tracked across all endpoints
- **Access Control**: Same permissions apply to passthrough endpoints

View file

@ -6,6 +6,10 @@ import Image from '@theme/IdealImage';
- Reject data before making llm api calls / before returning the response
- Enforce 'user' param for all openai endpoint calls
:::tip
**Understanding Callback Hooks?** Check out our [Callback Management Guide](../observability/callback_management.md) to understand the differences between proxy-specific hooks like `async_pre_call_hook` and general logging hooks like `async_log_success_event`.
:::
See a complete example with our [parallel request rate limiter](https://github.com/BerriAI/litellm/blob/main/litellm/proxy/hooks/parallel_request_limiter.py)
## Quick Start

View file

@ -335,12 +335,15 @@ router_settings:
| ANTHROPIC_API_KEY | API key for Anthropic service
| ANTHROPIC_API_BASE | Base URL for Anthropic API. Default is https://api.anthropic.com
| AWS_ACCESS_KEY_ID | Access Key ID for AWS services
| AWS_BATCH_ROLE_ARN | ARN of the AWS IAM role for batch operations
| AWS_DEFAULT_REGION | Default AWS region for service interactions when AWS_REGION is not set
| AWS_PROFILE_NAME | AWS CLI profile name to be used
| AWS_REGION | AWS region for service interactions (takes precedence over AWS_DEFAULT_REGION)
| AWS_REGION_NAME | Default AWS region for service interactions
| AWS_ROLE_ARN | ARN of the AWS IAM role to assume for authentication
| AWS_ROLE_NAME | Role name for AWS IAM usage
| AWS_S3_BUCKET_NAME | Name of the AWS S3 bucket for file operations
| AWS_S3_OUTPUT_BUCKET_NAME | Name of the AWS S3 output bucket for batch operations
| AWS_SECRET_ACCESS_KEY | Secret Access Key for AWS services
| AWS_SESSION_NAME | Name for AWS session
| AWS_WEB_IDENTITY_TOKEN | Web identity token for AWS

View file

@ -505,11 +505,11 @@ litellm_settings:
### Disable user-agent tracking
You can disable user-agent tracking by setting `litellm_settings.disable_user_agent_tracking` to `true`.
You can disable user-agent tracking by setting `litellm_settings.disable_add_user_agent_to_request_tags` to `true`.
```yaml
litellm_settings:
disable_user_agent_tracking: true
disable_add_user_agent_to_request_tags: true
```
## ✨ (Enterprise) Generate Spend Reports

View file

@ -13,6 +13,23 @@ For more details on routing strategies / params, see [Routing](../routing.md)
:::
## How Load Balancing Works
LiteLLM automatically distributes requests across multiple deployments of the same model using its built-in router. the proxy routes traffic to optimize performance and reliability.
"simple-shuffle" routing strategy is used by default
### Routing Strategies
| Strategy | Description | When to Use |
|----------|-------------|-------------|
| **simple-shuffle** (recommended) | Randomly distributes requests | General purpose, good for even load distribution |
| **least-busy** | Routes to deployment with fewest active requests | High concurrency scenarios |
| **usage-based-routing** (bad for perf) | Routes to deployment with lowest current usage (RPM/TPM) | When you want to respect rate limits evenly |
| **latency-based-routing** | Routes to fastest responding deployment | Latency-critical applications |
| **cost-based-routing** | Routes to deployment with lowest cost | Cost-sensitive applications |
## Quick Start - Load Balancing
#### Step 1 - Set deployments on config
@ -106,49 +123,13 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
]
}'
```
</TabItem>
<TabItem value="langchain" label="Langchain">
```python
from langchain.chat_models import ChatOpenAI
from langchain.prompts.chat import (
ChatPromptTemplate,
HumanMessagePromptTemplate,
SystemMessagePromptTemplate,
)
from langchain.schema import HumanMessage, SystemMessage
import os
os.environ["OPENAI_API_KEY"] = "anything"
chat = ChatOpenAI(
openai_api_base="http://0.0.0.0:4000",
model="gpt-3.5-turbo",
)
messages = [
SystemMessage(
content="You are a helpful assistant that im using to make a test request to."
),
HumanMessage(
content="test from litellm. tell me why it's amazing in 1 sentence"
),
]
response = chat(messages)
print(response)
```
</TabItem>
</Tabs>
### Test - Loadbalancing
In this request, the following will occur:
1. A rate limit exception will be raised
2. LiteLLM proxy will retry the request on the model group (default is 3).
2. LiteLLM proxy will retry the request on the model group (default retries are 3).
```bash
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
@ -256,4 +237,16 @@ model_group_alias: Optional[Dict[str, Union[str, RouterModelGroupAliasItem]]] =
class RouterModelGroupAliasItem(TypedDict):
model: str
hidden: bool # if 'True', don't return on `/v1/models`, `/v1/model/info`, `/v1/model_group/info`
```
```
### When You'll See Load Balancing in Action
**Immediate Effects:**
- Different deployments serve subsequent requests (visible in logs)
- Better response times during high traffic
**Observable Benefits:**
- **Higher throughput**: More requests handled simultaneously across deployments
- **Improved reliability**: If one deployment fails, traffic automatically routes to healthy ones
- **Better resource utilization**: Load spread evenly across all available deployments

View file

@ -63,7 +63,7 @@ Use this for for tracking per [user, key, team, etc.](virtual_keys)
| Metric Name | Description |
|----------------------|--------------------------------------|
| `litellm_spend_metric` | Total Spend, per `"user", "key", "model", "team", "end-user"` |
| `litellm_spend_metric` | Total Spend, per `"end_user", "hashed_api_key", "api_key_alias", "model", "team", "team_alias", "user"` |
| `litellm_total_tokens_metric` | input + output tokens per `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "model"` |
| `litellm_input_tokens_metric` | input tokens per `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "model"` |
| `litellm_output_tokens_metric` | output tokens per `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "model"` |
@ -73,9 +73,9 @@ Use this for for tracking per [user, key, team, etc.](virtual_keys)
| Metric Name | Description |
|----------------------|--------------------------------------|
| `litellm_team_max_budget_metric` | Max Budget for Team Labels: `"team_id", "team_alias"`|
| `litellm_remaining_team_budget_metric` | Remaining Budget for Team (A team created on LiteLLM) Labels: `"team_id", "team_alias"`|
| `litellm_team_budget_remaining_hours_metric` | Hours before the team budget is reset Labels: `"team_id", "team_alias"`|
| `litellm_team_max_budget_metric` | Max Budget for Team Labels: `"team", "team_alias"`|
| `litellm_remaining_team_budget_metric` | Remaining Budget for Team (A team created on LiteLLM) Labels: `"team", "team_alias"`|
| `litellm_team_budget_remaining_hours_metric` | Hours before the team budget is reset Labels: `"team", "team_alias"`|
### Virtual Key - Budget
@ -119,8 +119,8 @@ Use this to track overall LiteLLM Proxy usage.
| Metric Name | Description |
|----------------------|--------------------------------------|
| `litellm_proxy_failed_requests_metric` | Total number of failed responses from proxy - the client did not get a success response from litellm proxy. Labels: `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "exception_status", "exception_class"` |
| `litellm_proxy_total_requests_metric` | Total number of requests made to the proxy server - track number of client side requests. Labels: `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "status_code"` |
| `litellm_proxy_failed_requests_metric` | Total number of failed responses from proxy - the client did not get a success response from litellm proxy. Labels: `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "exception_status", "exception_class", "route"` |
| `litellm_proxy_total_requests_metric` | Total number of requests made to the proxy server - track number of client side requests. Labels: `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "status_code", "user_email", "route"` |
## LLM Provider Metrics
@ -155,7 +155,7 @@ Use this for LLM API Error monitoring and tracking remaining rate limits and tok
| Metric Name | Description |
|----------------------|--------------------------------------|
| `litellm_remaining_requests_metric` | Track `x-ratelimit-remaining-requests` returned from LLM API Deployment. Labels: `"model_group", "api_provider", "api_base", "litellm_model_name", "hashed_api_key", "api_key_alias"` |
| `litellm_remaining_tokens` | Track `x-ratelimit-remaining-tokens` return from LLM API Deployment. Labels: `"model_group", "api_provider", "api_base", "litellm_model_name", "hashed_api_key", "api_key_alias"` |
| `litellm_remaining_tokens_metric` | Track `x-ratelimit-remaining-tokens` return from LLM API Deployment. Labels: `"model_group", "api_provider", "api_base", "litellm_model_name", "hashed_api_key", "api_key_alias"` |
### Deployment State
| Metric Name | Description |
@ -167,16 +167,22 @@ Use this for LLM API Error monitoring and tracking remaining rate limits and tok
| Metric Name | Description |
|----------------------|--------------------------------------|
| `litellm_deployment_cooled_down` | Number of times a deployment has been cooled down by LiteLLM load balancing logic. Labels: `"litellm_model_name", "model_id", "api_base", "api_provider", "exception_status"` |
| `litellm_deployment_cooled_down` | Number of times a deployment has been cooled down by LiteLLM load balancing logic. Labels: `"litellm_model_name", "model_id", "api_base", "api_provider"` |
| `litellm_deployment_successful_fallbacks` | Number of successful fallback requests from primary model -> fallback model. Labels: `"requested_model", "fallback_model", "hashed_api_key", "api_key_alias", "team", "team_alias", "exception_status", "exception_class"` |
| `litellm_deployment_failed_fallbacks` | Number of failed fallback requests from primary model -> fallback model. Labels: `"requested_model", "fallback_model", "hashed_api_key", "api_key_alias", "team", "team_alias", "exception_status", "exception_class"` |
## Request Counting Metrics
| Metric Name | Description |
|----------------------|--------------------------------------|
| `litellm_requests_metric` | Total number of requests tracked per endpoint. Labels: `"end_user", "hashed_api_key", "api_key_alias", "model", "team", "team_alias", "user", "user_email"` |
## Request Latency Metrics
| Metric Name | Description |
|----------------------|--------------------------------------|
| `litellm_request_total_latency_metric` | Total latency (seconds) for a request to LiteLLM Proxy Server - tracked for labels "end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "model" |
| `litellm_overhead_latency_metric` | Latency overhead (seconds) added by LiteLLM processing - tracked for labels "end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "model" |
| `litellm_overhead_latency_metric` | Latency overhead (seconds) added by LiteLLM processing - tracked for labels "model_group", "api_provider", "api_base", "litellm_model_name", "hashed_api_key", "api_key_alias" |
| `litellm_llm_api_latency_metric` | Latency (seconds) for just the LLM API call - tracked for labels "model", "hashed_api_key", "api_key_alias", "team", "team_alias", "requested_model", "end_user", "user" |
| `litellm_llm_api_time_to_first_token_metric` | Time to first token for LLM API call - tracked for labels `model`, `hashed_api_key`, `api_key_alias`, `team`, `team_alias` [Note: only emitted for streaming requests] |
@ -486,7 +492,6 @@ Here is a screenshot of the metrics you can monitor with the LiteLLM Grafana Das
| Metric Name | Description |
|----------------------|--------------------------------------|
| `litellm_llm_api_failed_requests_metric` | **deprecated** use `litellm_proxy_failed_requests_metric` |
| `litellm_requests_metric` | **deprecated** use `litellm_proxy_total_requests_metric` |

View file

@ -0,0 +1,163 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Veo Video Generation with Google AI Studio
Generate videos using Google's Veo model through LiteLLM's pass-through endpoints.
## Quick Start
LiteLLM allows you to use Google AI Studio's Veo video generation API through pass-through routes with zero configuration.
### 1. Add Google AI Studio API Key to your environment
```bash
export GEMINI_API_KEY="your_google_ai_studio_api_key"
```
### 2. Start LiteLLM Proxy
```bash
litellm
# RUNNING on http://0.0.0.0:4000
```
### 3. Generate Video
<Tabs>
<TabItem value="python" label="Python">
```python
import requests
import time
import json
# Configuration
BASE_URL = "http://localhost:4000/gemini/v1beta"
API_KEY = "anything" # Use "anything" as the key
headers = {
"x-goog-api-key": API_KEY,
"Content-Type": "application/json"
}
# Step 1: Initiate video generation
def generate_video(prompt):
url = f"{BASE_URL}/models/veo-3.0-generate-preview:predictLongRunning"
payload = {
"instances": [{
"prompt": prompt
}]
}
response = requests.post(url, headers=headers, json=payload)
response.raise_for_status()
data = response.json()
return data.get("name") # Operation name
# Step 2: Poll for completion
def wait_for_completion(operation_name):
operation_url = f"{BASE_URL}/{operation_name}"
while True:
response = requests.get(operation_url, headers=headers)
response.raise_for_status()
data = response.json()
if data.get("done", False):
# Extract video URI
video_uri = data["response"]["generateVideoResponse"]["generatedSamples"][0]["video"]["uri"]
return video_uri
time.sleep(10) # Wait 10 seconds before next poll
# Step 3: Download video
def download_video(video_uri, filename="generated_video.mp4"):
# Replace Google URL with LiteLLM proxy URL
litellm_url = video_uri.replace(
"https://generativelanguage.googleapis.com/v1beta",
BASE_URL
)
response = requests.get(litellm_url, headers=headers, stream=True)
response.raise_for_status()
with open(filename, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
return filename
# Complete workflow
prompt = "A cat playing with a ball of yarn in a sunny garden"
print("Generating video...")
operation_name = generate_video(prompt)
print("Waiting for completion...")
video_uri = wait_for_completion(operation_name)
print("Downloading video...")
filename = download_video(video_uri)
print(f"Video saved as: {filename}")
```
</TabItem>
<TabItem value="curl" label="Curl">
```bash
# Step 1: Initiate video generation
curl -X POST "http://localhost:4000/gemini/v1beta/models/veo-3.0-generate-preview:predictLongRunning" \
-H "x-goog-api-key: anything" \
-H "Content-Type: application/json" \
-d '{
"instances": [{
"prompt": "A cat playing with a ball of yarn in a sunny garden"
}]
}'
# Response will include operation name:
# {"name": "operations/generate_12345"}
# Step 2: Poll for completion
curl -X GET "http://localhost:4000/gemini/v1beta/operations/generate_12345" \
-H "x-goog-api-key: anything"
# Step 3: Download video (when done=true)
curl -X GET "http://localhost:4000/gemini/v1beta/files/VIDEO_ID:download?alt=media" \
-H "x-goog-api-key: anything" \
--output generated_video.mp4
```
</TabItem>
</Tabs>
## Complete Example
For a full working example with error handling and logging, see our [Veo Video Generation Cookbook](https://github.com/BerriAI/litellm/blob/main/cookbook/veo_video_generation.py).
## How It Works
1. **Video Generation Request**: Send a prompt to Veo's `predictLongRunning` endpoint
2. **Operation Polling**: Monitor the long-running operation until completion
3. **File Download**: Download the generated video through LiteLLM's pass-through with automatic redirect handling
LiteLLM handles:
- ✅ Authentication with Google AI Studio
- ✅ Request routing and proxying
- ✅ Automatic redirect handling for file downloads
## Configuration Options
### Environment Variables
```bash
export GEMINI_API_KEY="your_google_ai_studio_api_key"
```

View file

@ -20,6 +20,7 @@ Supported Providers:
- Vertex AI (`vertex_ai/`)
- Perplexity (`perplexity/`)
- Mistral AI (Magistral models) (`mistral/`)
- Groq (`groq/`)
LiteLLM will standardize the `reasoning_content` in the response and `thinking_blocks` in the assistant message.

View file

@ -450,6 +450,7 @@ vertex_vision_models: Set = set()
vertex_chat_models: Set = set()
vertex_code_chat_models: Set = set()
vertex_ai_image_models: Set = set()
vertex_ai_video_models: Set = set()
vertex_text_models: Set = set()
vertex_code_text_models: Set = set()
vertex_embedding_models: Set = set()
@ -605,6 +606,9 @@ def add_known_models():
elif value.get("litellm_provider") == "vertex_ai-image-models":
key = key.replace("vertex_ai/", "")
vertex_ai_image_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-video-models":
key = key.replace("vertex_ai/", "")
vertex_ai_video_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-openai_models":
key = key.replace("vertex_ai/", "")
vertex_openai_models.add(key)

View file

@ -14,13 +14,15 @@ import asyncio
import contextvars
import os
from functools import partial
from typing import Any, Coroutine, Dict, Literal, Optional, Union
from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast
import httpx
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.azure.batches.handler import AzureBatchesAPI
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.llms.openai.openai import OpenAIBatchesAPI
from litellm.llms.vertex_ai.batches.handler import VertexAIBatchPrediction
from litellm.secret_managers.main import get_secret_str
@ -31,13 +33,19 @@ from litellm.types.llms.openai import (
RetrieveBatchRequest,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LiteLLMBatch
from litellm.utils import client, get_litellm_params, supports_httpx_timeout
from litellm.types.utils import LiteLLMBatch, LlmProviders
from litellm.utils import (
ProviderConfigManager,
client,
get_litellm_params,
supports_httpx_timeout,
)
####### ENVIRONMENT VARIABLES ###################
openai_batches_instance = OpenAIBatchesAPI()
azure_batches_instance = AzureBatchesAPI()
vertex_ai_batches_instance = VertexAIBatchPrediction(gcs_bucket_name="")
base_llm_http_handler = BaseLLMHTTPHandler()
#################################################
@ -46,7 +54,7 @@ async def acreate_batch(
completion_window: Literal["24h"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
input_file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@ -94,7 +102,7 @@ def create_batch(
completion_window: Literal["24h"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
input_file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@ -111,8 +119,8 @@ def create_batch(
proxy_server_request = kwargs.get("proxy_server_request", None)
model_info = kwargs.get("model_info", None)
_is_async = kwargs.pop("acreate_batch", False) is True
litellm_params = get_litellm_params(**kwargs)
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj", None)
litellm_params = dict(GenericLiteLLMParams(**kwargs))
litellm_logging_obj: LiteLLMLoggingObj = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None))
### TIMEOUT LOGIC ###
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
litellm_logging_obj.update_environment_variables(
@ -142,6 +150,7 @@ def create_batch(
timeout = float(timeout) # type: ignore
elif timeout is None:
timeout = 600.0
_create_batch_request = CreateBatchRequest(
completion_window=completion_window,
@ -151,6 +160,27 @@ def create_batch(
extra_headers=extra_headers,
extra_body=extra_body,
)
provider_config = ProviderConfigManager.get_provider_batches_config(
model="",
provider=LlmProviders(custom_llm_provider),
)
if provider_config is not None:
response = base_llm_http_handler.create_batch(
provider_config=provider_config,
litellm_params=litellm_params,
create_batch_data=_create_batch_request,
headers=extra_headers or {},
api_base=optional_params.api_base,
api_key=optional_params.api_key,
logging_obj=litellm_logging_obj,
_is_async=_is_async,
client=client
if client is not None
and isinstance(client, (HTTPHandler, AsyncHTTPHandler))
else None,
timeout=timeout,
)
return response
api_base: Optional[str] = None
if custom_llm_provider == "openai":
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -322,20 +352,21 @@ def retrieve_batch(
"""
try:
optional_params = GenericLiteLLMParams(**kwargs)
litellm_logging_obj: 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(
custom_llm_provider=custom_llm_provider,
**kwargs,
)
litellm_logging_obj.update_environment_variables(
model=None,
user=None,
optional_params=optional_params.model_dump(),
litellm_params=litellm_params,
custom_llm_provider=custom_llm_provider,
)
if litellm_logging_obj is not None:
litellm_logging_obj.update_environment_variables(
model=None,
user=None,
optional_params=optional_params.model_dump(),
litellm_params=litellm_params,
custom_llm_provider=custom_llm_provider,
)
if (
timeout is not None

View file

@ -50,7 +50,7 @@ vertex_ai_files_instance = VertexAIFilesHandler()
async def acreate_file(
file: FileTypes,
purpose: Literal["assistants", "batch", "fine-tune"],
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@ -94,7 +94,7 @@ async def acreate_file(
def create_file(
file: FileTypes,
purpose: Literal["assistants", "batch", "fine-tune"],
custom_llm_provider: Optional[Literal["openai", "azure", "vertex_ai"]] = None,
custom_llm_provider: Optional[Literal["openai", "azure", "vertex_ai", "bedrock"]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@ -109,7 +109,7 @@ def create_file(
try:
_is_async = kwargs.pop("acreate_file", False) is True
optional_params = GenericLiteLLMParams(**kwargs)
litellm_params_dict = get_litellm_params(**kwargs)
litellm_params_dict = dict(**kwargs)
logging_obj = cast(
Optional[LiteLLMLoggingObj], kwargs.get("litellm_logging_obj")
)

View file

@ -320,6 +320,7 @@ def get_llm_provider( # noqa: PLR0915
or model in litellm.vertex_embedding_models
or model in litellm.vertex_vision_models
or model in litellm.vertex_ai_image_models
or model in litellm.vertex_ai_video_models
):
custom_llm_provider = "vertex_ai"
## ai21

View file

@ -1,5 +1,6 @@
from .anthropic_messages.transformation import BaseAnthropicMessagesConfig
from .audio_transcription.transformation import BaseAudioTranscriptionConfig
from .batches.transformation import BaseBatchesConfig
from .chat.transformation import BaseConfig
from .embedding.transformation import BaseEmbeddingConfig
from .image_edit.transformation import BaseImageEditConfig
@ -12,4 +13,5 @@ __all__ = [
"BaseAnthropicMessagesConfig",
"BaseEmbeddingConfig",
"BaseImageEditConfig",
"BaseBatchesConfig",
]

View file

@ -0,0 +1,176 @@
import types
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
import httpx
from httpx import Headers
from litellm.types.llms.openai import (
AllMessageValues,
CreateBatchRequest,
)
from litellm.types.utils import LiteLLMBatch, LlmProviders
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from ..chat.transformation import BaseLLMException as _BaseLLMException
LiteLLMLoggingObj = _LiteLLMLoggingObj
BaseLLMException = _BaseLLMException
else:
LiteLLMLoggingObj = Any
BaseLLMException = Any
class BaseBatchesConfig(ABC):
"""
Abstract base class for batch processing configurations across different LLM providers.
This class defines the interface that all provider-specific batch configurations
must implement to work with LiteLLM's unified batch processing system.
"""
def __init__(self):
pass
@property
@abstractmethod
def custom_llm_provider(self) -> LlmProviders:
"""Return the LLM provider type for this configuration."""
pass
@classmethod
def get_config(cls):
"""Get configuration dictionary for this class."""
return {
k: v
for k, v in cls.__dict__.items()
if not k.startswith("__")
and not k.startswith("_abc")
and not isinstance(
v,
(
types.FunctionType,
types.BuiltinFunctionType,
classmethod,
staticmethod,
),
)
and v is not None
}
@abstractmethod
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
"""
Validate and prepare environment-specific headers and parameters.
Args:
headers: HTTP headers dictionary
model: Model name
messages: List of messages
optional_params: Optional parameters
litellm_params: LiteLLM parameters
api_key: API key
api_base: API base URL
Returns:
Updated headers dictionary
"""
pass
@abstractmethod
def get_complete_batch_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: Dict,
litellm_params: Dict,
data: CreateBatchRequest,
) -> str:
"""
Get the complete URL for batch creation request.
Args:
api_base: Base API URL
api_key: API key
model: Model name
optional_params: Optional parameters
litellm_params: LiteLLM parameters
data: Batch creation request data
Returns:
Complete URL for the batch request
"""
pass
@abstractmethod
def transform_create_batch_request(
self,
model: str,
create_batch_data: CreateBatchRequest,
optional_params: dict,
litellm_params: dict,
) -> Union[bytes, str, Dict[str, Any]]:
"""
Transform the batch creation request to provider-specific format.
Args:
model: Model name
create_batch_data: Batch creation request data
optional_params: Optional parameters
litellm_params: LiteLLM parameters
Returns:
Transformed request data
"""
pass
@abstractmethod
def transform_create_batch_response(
self,
model: Optional[str],
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> LiteLLMBatch:
"""
Transform provider-specific batch response to LiteLLM format.
Args:
model: Model name
raw_response: Raw HTTP response
logging_obj: Logging object
litellm_params: LiteLLM parameters
Returns:
LiteLLM batch object
"""
pass
@abstractmethod
def get_error_class(
self, error_message: str, status_code: int, headers: Union[Dict, Headers]
) -> "BaseLLMException":
"""
Get the appropriate error class for this provider.
Args:
error_message: Error message
status_code: HTTP status code
headers: Response headers
Returns:
Provider-specific exception class
"""
pass

View file

@ -35,6 +35,16 @@ class BaseFilesConfig(BaseConfig):
def custom_llm_provider(self) -> LlmProviders:
pass
@property
def file_upload_http_method(self) -> str:
"""
HTTP method to use for file uploads.
Override this in provider configs if they need different methods.
Default is POST (used by most providers like OpenAI, Anthropic).
S3-based providers like Bedrock should return "PUT".
"""
return "POST"
@abstractmethod
def get_supported_openai_params(
self, model: str

View file

@ -0,0 +1,254 @@
import os
import time
from typing import Any, Dict, List, Literal, Optional, Union, cast
from httpx import Headers, Response
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.bedrock import (
BedrockBatchJobStatus,
BedrockCreateBatchRequest,
BedrockCreateBatchResponse,
BedrockInputDataConfig,
BedrockOutputDataConfig,
BedrockS3InputDataConfig,
BedrockS3OutputDataConfig,
)
from litellm.types.llms.openai import (
AllMessageValues,
CreateBatchRequest,
)
from litellm.types.utils import LiteLLMBatch, LlmProviders
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import CommonBatchFilesUtils
class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
"""
Config for Bedrock Batches - handles batch job creation and management for Bedrock
"""
def __init__(self):
super().__init__()
self.common_utils = CommonBatchFilesUtils()
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.BEDROCK
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
"""
Validate and prepare environment for Bedrock batch requests.
AWS credentials are handled by BaseAWSLLM.
"""
# Add any Bedrock-specific headers if needed
return headers
def get_complete_batch_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: Dict,
litellm_params: Dict,
data: CreateBatchRequest,
) -> str:
"""
Get the complete URL for Bedrock batch creation.
Bedrock batch jobs are created via the model invocation job API.
"""
aws_region_name = self._get_aws_region_name(optional_params, model)
# Bedrock model invocation job endpoint
# Format: https://bedrock.{region}.amazonaws.com/model-invocation-job
bedrock_endpoint = f"https://bedrock.{aws_region_name}.amazonaws.com/model-invocation-job"
return bedrock_endpoint
def transform_create_batch_request(
self,
model: str,
create_batch_data: CreateBatchRequest,
optional_params: dict,
litellm_params: dict,
) -> Dict[str, Any]:
"""
Transform the batch creation request to Bedrock format.
Bedrock batch inference requires:
- modelId: The Bedrock model ID
- jobName: Unique name for the batch job
- inputDataConfig: Configuration for input data (S3 location)
- outputDataConfig: Configuration for output data (S3 location)
- roleArn: IAM role ARN for the batch job
"""
# Get required parameters
input_file_id = create_batch_data.get("input_file_id")
if not input_file_id:
raise ValueError("input_file_id is required for Bedrock batch creation")
# Extract S3 information from file ID using common utility
input_bucket, input_key = self.common_utils.parse_s3_uri(input_file_id)
# Get output S3 configuration
output_bucket = litellm_params.get("s3_output_bucket_name") or os.getenv("AWS_S3_OUTPUT_BUCKET_NAME")
if not output_bucket:
# Use same bucket as input if no output bucket specified
output_bucket = input_bucket
# Get IAM role ARN
role_arn = (
litellm_params.get("aws_batch_role_arn")
or optional_params.get("aws_batch_role_arn")
or os.getenv("AWS_BATCH_ROLE_ARN")
)
if not role_arn:
raise ValueError(
"AWS IAM role ARN is required for Bedrock batch jobs. "
"Set 'aws_batch_role_arn' in litellm_params or AWS_BATCH_ROLE_ARN env var"
)
# Get the actual Bedrock model ID using common utility
bedrock_model_id = self.common_utils.extract_model_from_s3_file_path(input_file_id, optional_params)
if not bedrock_model_id:
raise ValueError("Could not determine Bedrock model ID. Ensure the model is specified in the input file or passed as a parameter.")
# Generate job name with the correct model ID using common utility
job_name = self.common_utils.generate_unique_job_name(bedrock_model_id, prefix="litellm")
output_key = f"litellm-batch-outputs/{job_name}/"
# Build input data config
input_data_config: BedrockInputDataConfig = {
"s3InputDataConfig": BedrockS3InputDataConfig(
s3Uri=f"s3://{input_bucket}/{input_key}"
)
}
# Build output data config
output_data_config: BedrockOutputDataConfig = {
"s3OutputDataConfig": BedrockS3OutputDataConfig(
s3Uri=f"s3://{output_bucket}/{output_key}"
)
}
# Create Bedrock batch request with proper typing
bedrock_request: BedrockCreateBatchRequest = {
"modelId": bedrock_model_id,
"jobName": job_name,
"inputDataConfig": input_data_config,
"outputDataConfig": output_data_config,
"roleArn": role_arn
}
# Add optional parameters if provided
completion_window = create_batch_data.get("completion_window")
if completion_window:
# Map OpenAI completion window to Bedrock timeout
# OpenAI uses "24h", Bedrock expects timeout in hours
if completion_window == "24h":
bedrock_request["timeoutDurationInHours"] = 24
# For Bedrock, we need to return a pre-signed request with AWS auth headers
# Use common utility for AWS signing
endpoint_url = f"https://bedrock.{self._get_aws_region_name(optional_params, model)}.amazonaws.com/model-invocation-job"
signed_headers, signed_data = self.common_utils.sign_aws_request(
service_name="bedrock",
data=bedrock_request,
endpoint_url=endpoint_url,
optional_params=optional_params,
method="POST"
)
# Return a pre-signed request format that the HTTP handler can use
return {
"method": "POST",
"url": endpoint_url,
"headers": signed_headers,
"data": signed_data.decode('utf-8')
}
def transform_create_batch_response(
self,
model: Optional[str],
raw_response: Response,
logging_obj: Any,
litellm_params: dict,
) -> LiteLLMBatch:
"""
Transform Bedrock batch creation response to LiteLLM format.
"""
try:
response_data: BedrockCreateBatchResponse = raw_response.json()
except Exception as e:
raise ValueError(f"Failed to parse Bedrock batch response: {e}")
# Extract information from typed Bedrock response
job_arn = response_data.get("jobArn", "")
status: BedrockBatchJobStatus = response_data.get("status", "Submitted")
# Map Bedrock status to OpenAI-compatible status
status_mapping: Dict[BedrockBatchJobStatus, str] = {
"Submitted": "validating",
"InProgress": "in_progress",
"Completed": "completed",
"Failed": "failed",
"Stopping": "cancelling",
"Stopped": "cancelled"
}
openai_status = cast(Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"], status_mapping.get(status, "validating"))
# Get original request data from litellm_params if available
original_request = litellm_params.get("original_batch_request", {})
# Create LiteLLM batch object
return LiteLLMBatch(
id=job_arn, # Use ARN as the batch ID
object="batch",
endpoint=original_request.get("endpoint", "/v1/chat/completions"),
errors=None,
input_file_id=original_request.get("input_file_id", ""),
completion_window=original_request.get("completion_window", "24h"),
status=openai_status,
output_file_id=None, # Will be populated when job completes
error_file_id=None,
created_at=int(time.time()),
in_progress_at=int(time.time()) if status == "InProgress" else None,
expires_at=None,
finalizing_at=None,
completed_at=None,
failed_at=None,
expired_at=None,
cancelling_at=None,
cancelled_at=None,
request_counts=None,
metadata=original_request.get("metadata", {}),
)
def get_error_class(
self, error_message: str, status_code: int, headers: Union[Dict, Headers]
) -> BaseLLMException:
"""
Get Bedrock-specific error class using common utility.
"""
return self.common_utils.get_error_class(error_message, status_code, headers)

View file

@ -6,6 +6,9 @@ import json
import os
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union
if TYPE_CHECKING:
from litellm.types.llms.bedrock import BedrockCreateBatchRequest
import httpx
import litellm
@ -608,3 +611,218 @@ def get_anthropic_beta_from_headers(headers: dict) -> List[str]:
# Split comma-separated values and strip whitespace
return [beta.strip() for beta in anthropic_beta_header.split(",")]
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")
"""
if model.startswith("bedrock/"):
return model[8:] # Remove "bedrock/" prefix
return model
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:
"""
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"])
# 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
model_part = object_key[22:] # Remove "litellm-bedrock-files-"
# Find the last dash before the UUID
parts = model_part.split("-")
if len(parts) > 1:
# Reconstruct model name (everything except the last UUID part and .jsonl)
model_name = "-".join(parts[:-1])
if model_name.endswith(".jsonl"):
model_name = model_name[:-6] # Remove .jsonl
return model_name
except Exception:
pass
# Fallback to default model
return "anthropic.claude-3-5-sonnet-20240620-v1:0"
def sign_aws_request(
self,
service_name: str,
data: Union[str, dict, "BedrockCreateBatchRequest"],
endpoint_url: str,
optional_params: dict,
method: str = "POST",
) -> 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)
"""
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
# Get AWS credentials using existing methods
aws_region_name = self._base_aws._get_aws_region_name(
optional_params=optional_params, model=""
)
credentials = self._base_aws.get_credentials(
aws_access_key_id=optional_params.get("aws_access_key_id"),
aws_secret_access_key=optional_params.get("aws_secret_access_key"),
aws_session_token=optional_params.get("aws_session_token"),
aws_region_name=aws_region_name,
aws_session_name=optional_params.get("aws_session_name"),
aws_profile_name=optional_params.get("aws_profile_name"),
aws_role_name=optional_params.get("aws_role_name"),
aws_web_identity_token=optional_params.get("aws_web_identity_token"),
aws_sts_endpoint=optional_params.get("aws_sts_endpoint"),
)
# Prepare the request data
if isinstance(data, dict):
import json
request_data = json.dumps(data)
else:
request_data = data
# Prepare headers
headers = {"Content-Type": "application/json"}
# Create AWS request and sign it
sigv4 = SigV4Auth(credentials, service_name, aws_region_name)
request = AWSRequest(
method=method.upper(), url=endpoint_url, data=request_data, headers=headers
)
sigv4.add_auth(request)
prepped = request.prepare()
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)
"""
import fastuuid as 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,
optional_params: dict,
bucket_env_var: str = "AWS_S3_BUCKET_NAME",
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)
"""
import time
import uuid
# Get bucket name
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")
# 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(
self, error_message: str, status_code: int, headers: Union[Dict, httpx.Headers]
) -> BaseLLMException:
"""
Get Bedrock-specific error class.
"""
return BedrockError(
status_code=status_code,
message=error_message,
headers=headers
)

View file

@ -0,0 +1,607 @@
import json
import os
import time
import uuid
from typing import Any, Dict, List, Optional, Tuple, Union
from httpx import Headers, Response
from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.files.transformation import (
BaseFilesConfig,
LiteLLMLoggingObj,
)
from litellm.types.llms.openai import (
AllMessageValues,
CreateFileRequest,
FileTypes,
OpenAICreateFileRequestOptionalParams,
OpenAIFileObject,
PathLike,
)
from litellm.types.utils import ExtractedFileData, LlmProviders
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockError
class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
"""
Config for Bedrock Files - handles S3 uploads for Bedrock batch processing
"""
def __init__(self):
self.jsonl_transformation = BedrockJsonlFilesTransformation()
super().__init__()
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.BEDROCK
@property
def file_upload_http_method(self) -> str:
"""
Bedrock files are uploaded to S3, which requires PUT requests
"""
return "PUT"
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
# No additional headers needed for S3 uploads - AWS credentials handled by BaseAWSLLM
return headers
def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str:
"""
Helper to extract content from various OpenAI file types and return as string.
Handles:
- Direct content (str, bytes, IO[bytes])
- Tuple formats: (filename, content, [content_type], [headers])
- PathLike objects
"""
content: Union[str, bytes] = b""
# Extract file content from tuple if necessary
if isinstance(openai_file_content, tuple):
# Take the second element which is always the file content
file_content = openai_file_content[1]
else:
file_content = openai_file_content
# Handle different file content types
if isinstance(file_content, str):
# String content can be used directly
content = file_content
elif isinstance(file_content, bytes):
# Bytes content can be decoded
content = file_content
elif isinstance(file_content, PathLike): # PathLike
with open(str(file_content), "rb") as f:
content = f.read()
elif hasattr(file_content, "read"): # IO[bytes]
# File-like objects need to be read
content = file_content.read()
# Ensure content is string
if isinstance(content, bytes):
content = content.decode("utf-8")
return content
def _get_s3_object_name_from_batch_jsonl(
self,
openai_jsonl_content: List[Dict[str, Any]],
) -> str:
"""
Gets a unique S3 object name for the Bedrock batch processing job
named as: litellm-bedrock-files/{model}/{uuid}
"""
_model = openai_jsonl_content[0].get("body", {}).get("model", "")
# Remove bedrock/ prefix if present
if _model.startswith("bedrock/"):
_model = _model[8:]
object_name = f"litellm-bedrock-files-{_model}-{uuid.uuid4()}.jsonl"
return object_name
def get_object_name(
self, extracted_file_data: ExtractedFileData, purpose: str
) -> str:
"""
Get the object name for the request
"""
extracted_file_data_content = extracted_file_data.get("content")
if extracted_file_data_content is None:
raise ValueError("file content is required")
if purpose == "batch":
## 1. If jsonl, check if there's a model name
file_content = self._get_content_from_openai_file(
extracted_file_data_content
)
# Split into lines and parse each line as JSON
openai_jsonl_content = [
json.loads(line) for line in file_content.splitlines() if line.strip()
]
if len(openai_jsonl_content) > 0:
return self._get_s3_object_name_from_batch_jsonl(openai_jsonl_content)
## 2. If not jsonl, return the filename
filename = extracted_file_data.get("filename")
if filename:
return filename
## 3. If no file name, return timestamp
return str(int(time.time()))
def get_complete_file_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: Dict,
litellm_params: Dict,
data: CreateFileRequest,
) -> str:
"""
Get the complete S3 URL for the file upload request
"""
bucket_name = litellm_params.get("s3_bucket_name") or os.getenv("AWS_S3_BUCKET_NAME")
if not bucket_name:
raise ValueError("S3 bucket_name is required. Set 's3_bucket_name' in litellm_params or AWS_S3_BUCKET_NAME env var")
aws_region_name = self._get_aws_region_name(optional_params, model)
file_data = data.get("file")
purpose = data.get("purpose")
if file_data is None:
raise ValueError("file is required")
if purpose is None:
raise ValueError("purpose is required")
extracted_file_data = extract_file_data(file_data)
object_name = self.get_object_name(extracted_file_data, purpose)
# S3 endpoint URL format
s3_endpoint_url = optional_params.get("s3_endpoint_url") or f"https://s3.{aws_region_name}.amazonaws.com"
return f"{s3_endpoint_url}/{bucket_name}/{object_name}"
def get_supported_openai_params(
self, model: str
) -> List[OpenAICreateFileRequestOptionalParams]:
return []
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
return optional_params
def _get_bedrock_provider_from_model(self, model: str) -> Optional[str]:
"""
Extract provider from Bedrock model name
"""
if model.startswith("anthropic."):
return "anthropic"
elif model.startswith("cohere."):
return "cohere"
elif model.startswith("meta.") or model.startswith("llama"):
return "meta"
elif model.startswith("mistral."):
return "mistral"
elif model.startswith("ai21."):
return "ai21"
elif model.startswith("amazon."):
return "amazon"
else:
return None
def _map_openai_to_bedrock_params(
self,
openai_request_body: Dict[str, Any],
provider: Optional[str] = None,
) -> Dict[str, Any]:
"""
Transform OpenAI request body to Bedrock-compatible modelInput parameters using existing transformation logic
"""
_model = openai_request_body.get("model", "")
messages = openai_request_body.get("messages", [])
# Use existing Anthropic transformation logic for Anthropic models
if provider == "anthropic":
from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
AmazonAnthropicClaudeConfig,
)
anthropic_config = AmazonAnthropicClaudeConfig()
# Extract optional params (everything except model and messages)
optional_params = {k: v for k, v in openai_request_body.items() if k not in ["model", "messages"]}
# Transform using existing Anthropic logic
bedrock_params = anthropic_config.transform_request(
model=_model,
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
)
return bedrock_params
else:
# For other providers, use basic mapping
bedrock_params = {
"messages": messages,
**{k: v for k, v in openai_request_body.items() if k not in ["model", "messages"]}
}
return bedrock_params
def _transform_openai_jsonl_content_to_bedrock_jsonl_content(
self, openai_jsonl_content: List[Dict[str, Any]]
) -> List[Dict[str, Any]]:
"""
Transforms OpenAI JSONL content to Bedrock batch format
Bedrock batch format: { "recordId": "alphanumeric string", "modelInput": {JSON body} }
Example:
{
"recordId": "CALL0000001",
"modelInput": {
"anthropic_version": "bedrock-2023-05-31",
"max_tokens": 1024,
"messages": [
{
"role": "user",
"content": [{"type": "text", "text": "Hello"}]
}
]
}
}
"""
bedrock_jsonl_content = []
for idx, _openai_jsonl_content in enumerate(openai_jsonl_content):
# Extract the request body from OpenAI format
openai_body = _openai_jsonl_content.get("body", {})
model = openai_body.get("model", "")
# Determine provider from model name
provider = self._get_bedrock_provider_from_model(model)
# Transform to Bedrock modelInput format
model_input = self._map_openai_to_bedrock_params(
openai_request_body=openai_body,
provider=provider
)
# Create Bedrock batch record
record_id = _openai_jsonl_content.get("custom_id", f"CALL{str(idx).zfill(7)}")
bedrock_record = {
"recordId": record_id,
"modelInput": model_input
}
bedrock_jsonl_content.append(bedrock_record)
return bedrock_jsonl_content
def transform_create_file_request(
self,
model: str,
create_file_data: CreateFileRequest,
optional_params: dict,
litellm_params: dict,
) -> Union[bytes, str, dict]:
"""
Transform file request and return a pre-signed request for S3.
This keeps the HTTP handler clean by doing all the signing here.
"""
file_data = create_file_data.get("file")
if file_data is None:
raise ValueError("file is required")
extracted_file_data = extract_file_data(file_data)
extracted_file_data_content = extracted_file_data.get("content")
# Get and transform the file content
if (
create_file_data.get("purpose") == "batch"
and extracted_file_data.get("content_type") == "application/jsonl"
and extracted_file_data_content is not None
):
## Transform JSONL content to Bedrock format
original_file_content = self._get_content_from_openai_file(
extracted_file_data_content
)
openai_jsonl_content = [
json.loads(line) for line in original_file_content.splitlines() if line.strip()
]
bedrock_jsonl_content = (
self._transform_openai_jsonl_content_to_bedrock_jsonl_content(
openai_jsonl_content
)
)
file_content = "\n".join(json.dumps(item) for item in bedrock_jsonl_content)
elif isinstance(extracted_file_data_content, bytes):
file_content = extracted_file_data_content.decode('utf-8')
elif isinstance(extracted_file_data_content, str):
file_content = extracted_file_data_content
else:
raise ValueError("Unsupported file content type")
# Get the S3 URL for upload
api_base = self.get_complete_file_url(
api_base=None,
api_key=None,
model=model,
optional_params=optional_params,
litellm_params=litellm_params,
data=create_file_data,
)
# Sign the request and return a pre-signed request object
signed_headers, signed_body = self._sign_s3_request(
content=file_content,
api_base=api_base,
optional_params=optional_params,
)
# Return a dict that tells the HTTP handler exactly what to do
return {
"method": "PUT",
"url": api_base,
"headers": signed_headers,
"data": signed_body or file_content,
}
def _sign_s3_request(
self,
content: str,
api_base: str,
optional_params: dict,
) -> Tuple[dict, str]:
"""
Sign S3 PUT request using the same proven logic as S3Logger.
Reuses the exact pattern from litellm/integrations/s3_v2.py
"""
try:
import hashlib
import requests
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
# Get AWS credentials using existing methods
aws_region_name = self._get_aws_region_name(
optional_params=optional_params, model=""
)
credentials = self.get_credentials(
aws_access_key_id=optional_params.get("aws_access_key_id"),
aws_secret_access_key=optional_params.get("aws_secret_access_key"),
aws_session_token=optional_params.get("aws_session_token"),
aws_region_name=aws_region_name,
aws_session_name=optional_params.get("aws_session_name"),
aws_profile_name=optional_params.get("aws_profile_name"),
aws_role_name=optional_params.get("aws_role_name"),
aws_web_identity_token=optional_params.get("aws_web_identity_token"),
aws_sts_endpoint=optional_params.get("aws_sts_endpoint"),
)
# Calculate SHA256 hash of the content (REQUIRED for S3)
content_hash = hashlib.sha256(content.encode("utf-8")).hexdigest()
# Prepare headers with required S3 headers (same as s3_v2.py)
request_headers = {
"Content-Type": "application/json", # JSONL files are JSON content
"x-amz-content-sha256": content_hash, # REQUIRED by S3
"Content-Language": "en",
"Cache-Control": "private, immutable, max-age=31536000, s-maxage=0",
}
# Use requests.Request to prepare the request (same pattern as s3_v2.py)
req = requests.Request("PUT", api_base, data=content, headers=request_headers)
prepped = req.prepare()
# Sign the request with S3 service
aws_request = AWSRequest(
method=prepped.method,
url=prepped.url,
data=prepped.body,
headers=prepped.headers,
)
# Get region name for non-LLM API calls (same as s3_v2.py)
signing_region = self.get_aws_region_name_for_non_llm_api_calls(
aws_region_name=aws_region_name
)
SigV4Auth(credentials, "s3", signing_region).add_auth(aws_request)
# Return signed headers and body
signed_body = aws_request.body
if isinstance(signed_body, bytes):
signed_body = signed_body.decode('utf-8')
elif signed_body is None:
signed_body = content # Fallback to original content
return dict(aws_request.headers), signed_body
def transform_create_file_response(
self,
model: Optional[str],
raw_response: Response,
logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> OpenAIFileObject:
"""
Transform S3 File upload response into OpenAI-style FileObject
"""
# For S3 uploads, we typically get an ETag and other metadata
response_headers = raw_response.headers
# Extract S3 object information from the response
# S3 PUT object returns ETag and other metadata in headers
content_length = response_headers.get("Content-Length", "0")
# Extract bucket and key from the request URL or litellm_params
bucket_name = litellm_params.get("s3_bucket_name") or os.getenv("AWS_S3_BUCKET_NAME")
# Generate file ID in S3 format
object_key = getattr(logging_obj, 'object_key', None) or f"file-{int(time.time())}"
file_id = f"s3://{bucket_name}/{object_key}"
# Extract filename from object key
filename = object_key.split("/")[-1] if "/" in object_key else object_key
return OpenAIFileObject(
purpose="batch", # Default purpose for Bedrock files
id=file_id,
filename=filename,
created_at=int(time.time()), # Current timestamp
status="uploaded",
bytes=int(content_length) if content_length.isdigit() else 0,
object="file",
)
def get_error_class(
self, error_message: str, status_code: int, headers: Union[Dict, Headers]
) -> BaseLLMException:
return BedrockError(
status_code=status_code, message=error_message, headers=headers
)
class BedrockJsonlFilesTransformation:
"""
Transforms OpenAI /v1/files/* requests to Bedrock S3 file uploads for batch processing
"""
def transform_openai_file_content_to_bedrock_file_content(
self, openai_file_content: Optional[FileTypes] = None
) -> Tuple[str, str]:
"""
Transforms OpenAI FileContentRequest to Bedrock S3 file format
"""
if openai_file_content is None:
raise ValueError("contents of file are None")
# Read the content of the file
file_content = self._get_content_from_openai_file(openai_file_content)
# Split into lines and parse each line as JSON
openai_jsonl_content = [
json.loads(line) for line in file_content.splitlines() if line.strip()
]
bedrock_jsonl_content = (
self._transform_openai_jsonl_content_to_bedrock_jsonl_content(
openai_jsonl_content
)
)
bedrock_jsonl_string = "\n".join(
json.dumps(item) for item in bedrock_jsonl_content
)
object_name = self._get_s3_object_name(
openai_jsonl_content=openai_jsonl_content
)
return bedrock_jsonl_string, object_name
def _transform_openai_jsonl_content_to_bedrock_jsonl_content(
self, openai_jsonl_content: List[Dict[str, Any]]
):
"""
Delegate to the main BedrockFilesConfig transformation method
"""
config = BedrockFilesConfig()
return config._transform_openai_jsonl_content_to_bedrock_jsonl_content(openai_jsonl_content)
def _get_s3_object_name(
self,
openai_jsonl_content: List[Dict[str, Any]],
) -> str:
"""
Gets a unique S3 object name for the Bedrock batch processing job
named as: litellm-bedrock-files-{model}-{uuid}
"""
_model = openai_jsonl_content[0].get("body", {}).get("model", "")
# Remove bedrock/ prefix if present
if _model.startswith("bedrock/"):
_model = _model[8:]
object_name = f"litellm-bedrock-files-{_model}-{uuid.uuid4()}.jsonl"
return object_name
def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str:
"""
Helper to extract content from various OpenAI file types and return as string.
Handles:
- Direct content (str, bytes, IO[bytes])
- Tuple formats: (filename, content, [content_type], [headers])
- PathLike objects
"""
content: Union[str, bytes] = b""
# Extract file content from tuple if necessary
if isinstance(openai_file_content, tuple):
# Take the second element which is always the file content
file_content = openai_file_content[1]
else:
file_content = openai_file_content
# Handle different file content types
if isinstance(file_content, str):
# String content can be used directly
content = file_content
elif isinstance(file_content, bytes):
# Bytes content can be decoded
content = file_content
elif isinstance(file_content, PathLike): # PathLike
with open(str(file_content), "rb") as f:
content = f.read()
elif hasattr(file_content, "read"): # IO[bytes]
# File-like objects need to be read
content = file_content.read()
# Ensure content is string
if isinstance(content, bytes):
content = content.decode("utf-8")
return content
def transform_s3_bucket_response_to_openai_file_object(
self, create_file_data: CreateFileRequest, s3_upload_response: Dict[str, Any]
) -> OpenAIFileObject:
"""
Transforms S3 Bucket upload file response to OpenAI FileObject
"""
# S3 response typically contains ETag, key, etc.
object_key = s3_upload_response.get("Key", "")
bucket_name = s3_upload_response.get("Bucket", "")
# Extract filename from object key
filename = object_key.split("/")[-1] if "/" in object_key else object_key
return OpenAIFileObject(
purpose=create_file_data.get("purpose", "batch"),
id=f"s3://{bucket_name}/{object_key}",
filename=filename,
created_at=int(time.time()), # Current timestamp
status="uploaded",
bytes=s3_upload_response.get("ContentLength", 0),
object="file",
)

View file

@ -212,6 +212,7 @@ class AsyncHTTPHandler:
verify=ssl_config,
cert=cert,
headers=headers,
follow_redirects=True,
)
async def close(self):
@ -687,6 +688,7 @@ class HTTPHandler:
verify=ssl_config,
cert=cert,
headers=headers,
follow_redirects=True,
)
else:
self.client = client

View file

@ -28,6 +28,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import (
BaseAudioTranscriptionConfig,
)
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
from litellm.llms.base_llm.chat.transformation import BaseConfig
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
from litellm.llms.base_llm.files.transformation import BaseFilesConfig
@ -58,6 +59,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
from litellm.types.llms.openai import (
CreateBatchRequest,
CreateFileRequest,
OpenAIFileObject,
ResponseInputParam,
@ -66,7 +68,12 @@ from litellm.types.llms.openai import (
from litellm.types.rerank import OptionalRerankParams, RerankResponse
from litellm.types.responses.main import DeleteResponseResult
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import EmbeddingResponse, FileTypes, TranscriptionResponse
from litellm.types.utils import (
EmbeddingResponse,
FileTypes,
LiteLLMBatch,
TranscriptionResponse,
)
from litellm.types.vector_stores import (
VectorStoreCreateOptionalRequestParams,
VectorStoreCreateResponse,
@ -2212,15 +2219,38 @@ class BaseLLMHTTPHandler:
else:
sync_httpx_client = client
if isinstance(transformed_request, str) or isinstance(
transformed_request, bytes
):
upload_response = sync_httpx_client.post(
url=api_base,
headers=headers,
data=transformed_request,
if isinstance(transformed_request, dict) and "method" in transformed_request:
# Handle pre-signed requests (e.g., from Bedrock S3 uploads)
upload_response = getattr(sync_httpx_client, transformed_request["method"].lower())(
url=transformed_request["url"],
headers=transformed_request["headers"],
data=transformed_request["data"],
timeout=timeout,
)
elif isinstance(transformed_request, str) or isinstance(
transformed_request, bytes
):
# Handle traditional file uploads
# Ensure transformed_request is a string for httpx compatibility
if isinstance(transformed_request, bytes):
transformed_request = transformed_request.decode('utf-8')
# Use the HTTP method specified by the provider config
http_method = provider_config.file_upload_http_method.upper()
if http_method == "PUT":
upload_response = sync_httpx_client.put(
url=api_base,
headers=headers,
data=transformed_request,
timeout=timeout,
)
else: # Default to POST
upload_response = sync_httpx_client.post(
url=api_base,
headers=headers,
data=transformed_request,
timeout=timeout,
)
else:
try:
# Step 1: Initial request to get upload URL
@ -2280,16 +2310,52 @@ class BaseLLMHTTPHandler:
)
else:
async_httpx_client = client
#########################################################
# Debug Logging
#########################################################
logging_obj.pre_call(
input="",
api_key="",
additional_args={
"complete_input_dict": transformed_request,
"api_base": api_base,
"headers": headers,
},
)
if isinstance(transformed_request, str) or isinstance(
transformed_request, bytes
):
upload_response = await async_httpx_client.post(
url=api_base,
headers=headers,
data=transformed_request,
if isinstance(transformed_request, dict) and "method" in transformed_request:
# Handle pre-signed requests (e.g., from Bedrock S3 uploads)
upload_response = await getattr(async_httpx_client, transformed_request["method"].lower())(
url=transformed_request["url"],
headers=transformed_request["headers"],
data=transformed_request["data"],
timeout=timeout,
)
elif isinstance(transformed_request, str) or isinstance(
transformed_request, bytes
):
# Handle traditional file uploads
# Ensure transformed_request is a string for httpx compatibility
if isinstance(transformed_request, bytes):
transformed_request = transformed_request.decode('utf-8')
# Use the HTTP method specified by the provider config
http_method = provider_config.file_upload_http_method.upper()
if http_method == "PUT":
upload_response = await async_httpx_client.put(
url=api_base,
headers=headers,
data=transformed_request,
timeout=timeout,
)
else: # Default to POST
upload_response = await async_httpx_client.post(
url=api_base,
headers=headers,
data=transformed_request,
timeout=timeout,
)
else:
try:
# Step 1: Initial request to get upload URL
@ -2330,6 +2396,188 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params,
)
def create_batch(
self,
create_batch_data: "CreateBatchRequest",
litellm_params: dict,
provider_config: "BaseBatchesConfig",
headers: dict,
api_base: Optional[str],
api_key: Optional[str],
logging_obj: "LiteLLMLoggingObj",
_is_async: bool = False,
client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> Union["LiteLLMBatch", Coroutine[Any, Any, "LiteLLMBatch"]]:
"""
Creates a batch using provider-specific batch creation process
"""
# get config from model, custom llm provider
headers = provider_config.validate_environment(
api_key=api_key,
headers=headers,
model="",
messages=[],
optional_params={},
litellm_params=litellm_params,
)
api_base = provider_config.get_complete_batch_url(
api_base=api_base,
api_key=api_key,
model="",
optional_params={},
litellm_params=litellm_params,
data=create_batch_data,
)
if api_base is None:
raise ValueError("api_base is required for create_batch")
# Get the transformed request data
transformed_request = provider_config.transform_create_batch_request(
model="",
create_batch_data=create_batch_data,
litellm_params=litellm_params,
optional_params={},
)
if _is_async:
return self.async_create_batch(
transformed_request=transformed_request,
litellm_params=litellm_params,
provider_config=provider_config,
headers=headers,
api_base=api_base,
logging_obj=logging_obj,
client=client,
timeout=timeout,
create_batch_data=create_batch_data,
)
if client is None or not isinstance(client, HTTPHandler):
sync_httpx_client = _get_httpx_client()
else:
sync_httpx_client = client
try:
if isinstance(transformed_request, dict) and "method" in transformed_request:
# Handle pre-signed requests (e.g., from Bedrock with AWS auth)
batch_response = getattr(sync_httpx_client, transformed_request["method"].lower())(
url=transformed_request["url"],
headers=transformed_request["headers"],
data=transformed_request["data"],
timeout=timeout,
)
elif isinstance(transformed_request, dict):
# For other providers that use JSON requests
batch_response = sync_httpx_client.post(
url=api_base,
headers={**headers, "Content-Type": "application/json"},
json=transformed_request,
timeout=timeout,
)
else:
# Handle other request types if needed
batch_response = sync_httpx_client.post(
url=api_base,
headers=headers,
data=transformed_request,
timeout=timeout,
)
except Exception as e:
verbose_logger.exception(f"Error creating batch: {e}")
raise self._handle_error(
e=e,
provider_config=provider_config,
)
# Store original request for response transformation
litellm_params_with_request = {**litellm_params, "original_batch_request": create_batch_data}
return provider_config.transform_create_batch_response(
model=None,
raw_response=batch_response,
logging_obj=logging_obj,
litellm_params=litellm_params_with_request,
)
async def async_create_batch(
self,
transformed_request: Union[bytes, str, dict],
litellm_params: dict,
provider_config: "BaseBatchesConfig",
headers: dict,
api_base: str,
logging_obj: "LiteLLMLoggingObj",
client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
create_batch_data: Optional["CreateBatchRequest"] = None,
):
"""
Async version of create_batch
"""
if client is None or not isinstance(client, AsyncHTTPHandler):
async_httpx_client = get_async_httpx_client(
llm_provider=provider_config.custom_llm_provider
)
else:
async_httpx_client = client
#########################################################
# Debug Logging
#########################################################
logging_obj.pre_call(
input="",
api_key="",
additional_args={
"complete_input_dict": transformed_request,
"api_base": api_base,
"headers": headers,
},
)
try:
if isinstance(transformed_request, dict) and "method" in transformed_request:
# Handle pre-signed requests (e.g., from Bedrock with AWS auth)
batch_response = await getattr(async_httpx_client, transformed_request["method"].lower())(
url=transformed_request["url"],
headers=transformed_request["headers"],
data=transformed_request["data"],
timeout=timeout,
)
elif isinstance(transformed_request, dict):
# For other providers that use JSON requests
batch_response = await async_httpx_client.post(
url=api_base,
headers={**headers, "Content-Type": "application/json"},
json=transformed_request,
timeout=timeout,
)
else:
# Handle other request types if needed
batch_response = await async_httpx_client.post(
url=api_base,
headers=headers,
data=transformed_request,
timeout=timeout,
)
except Exception as e:
verbose_logger.exception(f"Error creating batch: {e}")
raise self._handle_error(
e=e,
provider_config=provider_config,
)
# Store original request for response transformation (for async version)
litellm_params_with_request = {**litellm_params, "original_batch_request": create_batch_data or {}}
return provider_config.transform_create_batch_response(
model=None,
raw_response=batch_response,
logging_obj=logging_obj,
litellm_params=litellm_params_with_request,
)
def list_files(self):
"""
Lists all files
@ -2381,6 +2629,7 @@ class BaseLLMHTTPHandler:
BaseVectorStoreConfig,
BaseGoogleGenAIGenerateContentConfig,
BaseAnthropicMessagesConfig,
BaseBatchesConfig,
"BasePassthroughConfig",
],
):

View file

@ -6,6 +6,8 @@ from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, cast,
import httpx
from pydantic import BaseModel
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
@ -55,6 +57,10 @@ class GroqChatConfig(OpenAILikeChatConfig):
if key != "self" and value is not None:
setattr(self.__class__, key, value)
@property
def custom_llm_provider(self) -> Optional[str]:
return "groq"
@classmethod
def get_config(cls):
return super().get_config()
@ -65,6 +71,15 @@ class GroqChatConfig(OpenAILikeChatConfig):
base_params.remove("max_retries")
except ValueError:
pass
try:
if litellm.supports_reasoning(
model=model, custom_llm_provider=self.custom_llm_provider
):
base_params.append("reasoning_effort")
except Exception as e:
verbose_logger.debug(f"Error checking if model supports reasoning: {e}")
return base_params
@overload

View file

@ -1,5 +1,6 @@
# What is this?
## API Handler for calling Vertex AI Partner Models
from enum import Enum
from typing import Callable, Optional, Union
import httpx # type: ignore
@ -27,6 +28,16 @@ class VertexAIError(Exception):
self.message
) # Call the base class constructor with the parameters it needs
class PartnerModelPrefixes(str, Enum):
META_PREFIX = "meta/"
DEEPSEEK_PREFIX = "deepseek-ai"
MISTRAL_PREFIX = "mistral"
CODERESTAL_PREFIX = "codestral"
JAMBA_PREFIX = "jamba"
CLAUDE_PREFIX = "claude"
QWEN_PREFIX = "qwen"
GPT_OSS_PREFIX = "openai/gpt-oss-"
class VertexAIPartnerModels(VertexBase):
def __init__(self) -> None:
@ -42,14 +53,14 @@ class VertexAIPartnerModels(VertexBase):
bool: True if the model string is a Vertex AI Partner Model, False otherwise
"""
if (
model.startswith("meta/")
or model.startswith("deepseek-ai")
or model.startswith("mistral")
or model.startswith("codestral")
or model.startswith("jamba")
or model.startswith("claude")
or model.startswith("qwen")
or model.startswith("openai")
model.startswith(PartnerModelPrefixes.META_PREFIX)
or model.startswith(PartnerModelPrefixes.DEEPSEEK_PREFIX)
or model.startswith(PartnerModelPrefixes.MISTRAL_PREFIX)
or model.startswith(PartnerModelPrefixes.CODERESTAL_PREFIX)
or model.startswith(PartnerModelPrefixes.JAMBA_PREFIX)
or model.startswith(PartnerModelPrefixes.CLAUDE_PREFIX)
or model.startswith(PartnerModelPrefixes.QWEN_PREFIX)
or model.startswith(PartnerModelPrefixes.GPT_OSS_PREFIX)
):
return True
return False
@ -58,9 +69,9 @@ class VertexAIPartnerModels(VertexBase):
def should_use_openai_handler(model: str):
OPENAI_LIKE_VERTEX_PROVIDERS = [
"llama",
"deepseek-ai",
"qwen",
"openai",
PartnerModelPrefixes.DEEPSEEK_PREFIX,
PartnerModelPrefixes.QWEN_PREFIX,
PartnerModelPrefixes.GPT_OSS_PREFIX,
]
if any(provider in model for provider in OPENAI_LIKE_VERTEX_PROVIDERS):
return True

View file

@ -356,7 +356,7 @@ async def acompletion(
logprobs: Optional[bool] = None,
top_logprobs: Optional[int] = None,
deployment_id=None,
reasoning_effort: Optional[Literal["minimal", "low", "medium", "high"]] = None,
reasoning_effort: Optional[Literal["none", "minimal", "low", "medium", "high", "default"]] = None,
safety_identifier: Optional[str] = None,
# set api_base, api_version, api_key
base_url: Optional[str] = None,
@ -897,7 +897,7 @@ def completion( # type: ignore # noqa: PLR0915
logit_bias: Optional[dict] = None,
user: Optional[str] = None,
# openai v1.0+ new params
reasoning_effort: Optional[Literal["minimal", "low", "medium", "high"]] = None,
reasoning_effort: Optional[Literal["none", "minimal", "low", "medium", "high", "default"]] = None,
response_format: Optional[Union[dict, Type[BaseModel]]] = None,
seed: Optional[int] = None,
tools: Optional[List] = None,

View file

@ -6157,21 +6157,7 @@
"supports_tool_choice": true,
"source": "https://inference-docs.cerebras.ai/support/pricing"
},
"cerebras/openai/gpt-oss-20b": {
"max_tokens": 32768,
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"input_cost_per_token": 7e-08,
"output_cost_per_token": 3e-07,
"litellm_provider": "cerebras",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"source": "https://inference-docs.cerebras.ai/support/pricing"
},
"cerebras/openai/gpt-oss-120b": {
"max_tokens": 32768,
"max_input_tokens": 131072,
@ -9498,6 +9484,48 @@
"source": "https://aistudio.google.com",
"supports_tool_choice": true
},
"gemini/veo-3.0-generate-preview": {
"max_tokens": 1024,
"max_input_tokens": 1024,
"output_cost_per_second": 0.75,
"litellm_provider": "gemini",
"mode": "video_generation",
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"video"
],
"source": "https://ai.google.dev/gemini-api/docs/video"
},
"gemini/veo-3.0-fast-generate-preview": {
"max_tokens": 1024,
"max_input_tokens": 1024,
"output_cost_per_second": 0.40,
"litellm_provider": "gemini",
"mode": "video_generation",
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"video"
],
"source": "https://ai.google.dev/gemini-api/docs/video"
},
"gemini/veo-2.0-generate-001": {
"max_tokens": 1024,
"max_input_tokens": 1024,
"output_cost_per_second": 0.35,
"litellm_provider": "gemini",
"mode": "video_generation",
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"video"
],
"source": "https://ai.google.dev/gemini-api/docs/video"
},
"vertex_ai/claude-opus-4-1": {
"max_tokens": 4096,
"max_input_tokens": 200000,
@ -10315,6 +10343,48 @@
"mode": "image_generation",
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
},
"vertex_ai/veo-3.0-generate-preview": {
"max_tokens": 1024,
"max_input_tokens": 1024,
"output_cost_per_second": 0.75,
"litellm_provider": "vertex_ai-video-models",
"mode": "video_generation",
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"video"
],
"source": "https://ai.google.dev/gemini-api/docs/video"
},
"vertex_ai/veo-3.0-fast-generate-preview": {
"max_tokens": 1024,
"max_input_tokens": 1024,
"output_cost_per_second": 0.40,
"litellm_provider": "vertex_ai-video-models",
"mode": "video_generation",
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"video"
],
"source": "https://ai.google.dev/gemini-api/docs/video"
},
"vertex_ai/veo-2.0-generate-001": {
"max_tokens": 1024,
"max_input_tokens": 1024,
"output_cost_per_second": 0.35,
"litellm_provider": "vertex_ai-video-models",
"mode": "video_generation",
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"video"
],
"source": "https://ai.google.dev/gemini-api/docs/video"
},
"text-embedding-004": {
"max_tokens": 2048,
"max_input_tokens": 2048,

View file

@ -388,7 +388,11 @@ class LiteLLMRoutes(enum.Enum):
]
# NOTE: ROUTES ONLY FOR MASTER KEY - only the Master Key should be able to Reset Spend
master_key_only_routes = ["/global/spend/reset"]
master_key_only_routes = [
"/global/spend/reset",
"/memory-usage-in-mem-cache",
"/memory-usage-in-mem-cache-items",
]
key_management_routes = [
KeyManagementRoutes.KEY_GENERATE,

View file

@ -5,10 +5,12 @@ import os
import tracemalloc
from collections import Counter
from fastapi import APIRouter
from fastapi import APIRouter, Depends
from litellm import get_secret_str
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
router = APIRouter()
@ -84,7 +86,9 @@ if os.environ.get("LITELLM_PROFILE", "false").lower() == "true":
@router.get("/memory-usage-in-mem-cache", include_in_schema=False)
async def memory_usage_in_mem_cache():
async def memory_usage_in_mem_cache(
_: UserAPIKeyAuth = Depends(user_api_key_auth),
):
# returns the size of all in-memory caches on the proxy server
"""
1. user_api_key_cache
@ -121,7 +125,9 @@ async def memory_usage_in_mem_cache():
@router.get("/memory-usage-in-mem-cache-items", include_in_schema=False)
async def memory_usage_in_mem_cache_items():
async def memory_usage_in_mem_cache_items(
_: UserAPIKeyAuth = Depends(user_api_key_auth),
):
# returns the size of all in-memory caches on the proxy server
"""
1. user_api_key_cache

View file

@ -68,6 +68,32 @@ end
return results
"""
TOKEN_INCREMENT_SCRIPT = """
local results = {}
-- Process each key/increment_value/ttl triplet
for i = 1, #KEYS do
local key = KEYS[i]
local increment_value = tonumber(ARGV[i * 2 - 1])
local ttl_seconds = tonumber(ARGV[i * 2])
-- Increment the value
local new_value = redis.call('INCRBYFLOAT', key, increment_value)
-- Handle TTL: only set expire if ttl_seconds > 0 and key has no current TTL
-- ttl_seconds can be 0 (no TTL) or positive (set TTL)
if ttl_seconds and ttl_seconds > 0 then
local current_ttl = redis.call('TTL', key)
if current_ttl == -1 then
redis.call('EXPIRE', key, ttl_seconds)
end
end
table.insert(results, new_value)
end
return results
"""
class RateLimitDescriptorRateLimitObject(TypedDict, total=False):
requests_per_unit: Optional[int]
@ -109,8 +135,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
BATCH_RATE_LIMITER_SCRIPT
)
)
self.token_increment_script = (
self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
TOKEN_INCREMENT_SCRIPT
)
)
else:
self.batch_rate_limiter_script = None
self.token_increment_script = None
self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60))
@ -567,6 +599,62 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return pipeline_operations
async def async_increment_tokens_with_ttl_preservation(
self,
pipeline_operations: List["RedisPipelineIncrementOperation"],
parent_otel_span: Optional[Span] = None,
) -> None:
"""
Increment token counters using Lua script to preserve existing TTL.
This prevents TTL reset on every token increment.
"""
if not pipeline_operations:
return
# Check if script is available
if self.token_increment_script is None:
verbose_proxy_logger.debug("TTL preservation script not available, using regular pipeline")
await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline(
increment_list=pipeline_operations,
litellm_parent_otel_span=parent_otel_span,
)
return
try:
# Use Lua script for all operations
keys = []
args = []
for op in pipeline_operations:
# Convert None TTL to 0 for Lua script
ttl_value = op["ttl"] if op["ttl"] is not None else 0
verbose_proxy_logger.debug(
f"Executing TTL-preserving increment for key={op['key']}, "
f"increment={op['increment_value']}, ttl={ttl_value}"
)
keys.append(op["key"])
args.extend([op["increment_value"], ttl_value])
await self.token_increment_script(
keys=keys,
args=args,
)
verbose_proxy_logger.debug(
f"Successfully executed TTL-preserving increment for {len(pipeline_operations)} keys"
)
except Exception as e:
verbose_proxy_logger.warning(
f"TTL preservation failed, falling back to regular pipeline: {str(e)}"
)
# Fallback to regular pipeline on error
await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline(
increment_list=pipeline_operations,
litellm_parent_otel_span=parent_otel_span,
)
def get_rate_limit_type(self) -> Literal["output", "input", "total"]:
from litellm.proxy.proxy_server import general_settings
@ -713,9 +801,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# Execute all increments in a single pipeline
if pipeline_operations:
await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline(
increment_list=pipeline_operations,
litellm_parent_otel_span=litellm_parent_otel_span,
await self.async_increment_tokens_with_ttl_preservation(
pipeline_operations=pipeline_operations,
parent_otel_span=litellm_parent_otel_span,
)
except Exception as e:

View file

@ -677,11 +677,6 @@ async def add_litellm_data_to_request( # noqa: PLR0915
from litellm.proxy.proxy_server import llm_router, premium_user
from litellm.types.proxy.litellm_pre_call_utils import SecretFields
safe_add_api_version_from_query_params(data, request)
_metadata_variable_name = _get_metadata_variable_name(request)
if data.get(_metadata_variable_name, None) is None:
data[_metadata_variable_name] = {}
_headers = clean_headers(
request.headers,
@ -692,6 +687,24 @@ async def add_litellm_data_to_request( # noqa: PLR0915
),
)
##########################################################
# Init - Proxy Server Request
# we do this as soon as entering so we track the original request
##########################################################
data["proxy_server_request"] = {
"url": str(request.url),
"method": request.method,
"headers": _headers,
"body": copy.copy(data), # use copy instead of deepcopy
}
safe_add_api_version_from_query_params(data, request)
_metadata_variable_name = _get_metadata_variable_name(request)
if data.get(_metadata_variable_name, None) is None:
data[_metadata_variable_name] = {}
data.update(
LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call(
headers=_headers,
@ -721,13 +734,6 @@ async def add_litellm_data_to_request( # noqa: PLR0915
if "user" not in data:
data["user"] = user
# Include original request and headers in the data
data["proxy_server_request"] = {
"url": str(request.url),
"method": request.method,
"headers": _headers,
"body": copy.copy(data), # use copy instead of deepcopy
}
data["secret_fields"] = SecretFields(raw_headers=dict(request.headers))

View file

@ -1566,14 +1566,12 @@ async def generate_key_helper_fn( # noqa: PLR0915
if duration is None: # allow tokens that never expire
expires = None
else:
duration_s = duration_in_seconds(duration=duration)
expires = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
expires = get_budget_reset_time(budget_duration=duration)
if key_budget_duration is None: # one-time budget
key_reset_at = None
else:
duration_s = duration_in_seconds(duration=key_budget_duration)
key_reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
key_reset_at = get_budget_reset_time(budget_duration=key_budget_duration)
if budget_duration is None: # one-time budget
reset_at = None

View file

@ -121,15 +121,16 @@ class ScimTransformations:
if isinstance(team, dict):
team = LiteLLM_TeamTable(**team)
# Get team members
# Get team members with proper display names
scim_members: List[SCIMMember] = []
for member in team.members_with_roles or []:
if isinstance(member, dict):
member = Member(**member)
scim_members.append(
SCIMMember(
value=ScimTransformations._get_scim_member_value(member),
display=member.user_email,
display=ScimTransformations._get_scim_member_display(member),
)
)
@ -151,6 +152,24 @@ class ScimTransformations:
@staticmethod
def _get_scim_member_value(member: Member) -> str:
if member.user_email:
"""
Get the SCIM member value. Use user_email if available, otherwise use user_id.
SCIM member value should be the unique identifier for the user.
"""
if hasattr(member, "user_email") and member.user_email:
return member.user_email
elif hasattr(member, "user_id"):
return member.user_id or ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE
return ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE
@staticmethod
def _get_scim_member_display(member: Member) -> str:
"""
Get the SCIM member display. Use user_email if available, otherwise use user_id.
SCIM member display should be the display name for the user.
"""
if hasattr(member, "user_email") and member.user_email:
return member.user_email
elif hasattr(member, "user_id"):
return member.user_id or ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE
return ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE

View file

@ -22,6 +22,7 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import (
LiteLLM_TeamTable,
LiteLLM_UserTable,
LitellmUserRoles,
Member,
@ -237,6 +238,23 @@ async def _handle_team_membership_changes(user_id: str, existing_teams: List[str
)
async def _get_team_member_user_ids_from_team(team: LiteLLM_TeamTable) -> List[str]:
"""
Get the IDs of the members from a team.
Use one source of truth for the member IDs: team.members_with_roles
"""
member_user_ids: List[str] = []
for member in team.members_with_roles or []:
if hasattr(member, "user_id") and member.user_id is not None:
member_user_ids.append(member.user_id)
elif isinstance(member, dict) and "user_id" in member:
user_id = member.get("user_id")
if user_id is not None:
member_user_ids.append(user_id)
return member_user_ids
# Dependency to set the correct SCIM Content-Type
async def set_scim_content_type(response: Response):
"""Sets the Content-Type header to application/scim+json"""
@ -253,6 +271,12 @@ async def set_scim_content_type(response: Response):
)
async def get_service_provider_config(request: Request):
"""Return SCIM Service Provider Configuration."""
verbose_proxy_logger.debug(
"SCIM ServiceProviderConfig request: method=%s url=%s headers=%s",
request.method,
request.url,
dict(request.headers),
)
meta = {
"resourceType": "ServiceProviderConfig",
"location": str(request.url),
@ -275,6 +299,12 @@ async def get_users(
"""
Get a list of users according to SCIM v2 protocol
"""
verbose_proxy_logger.debug(
"SCIM GET USERS request: startIndex=%s count=%s filter=%s",
startIndex,
count,
filter,
)
try:
prisma_client = await _get_prisma_client_or_raise_exception()
# Parse filter if provided (basic support)
@ -334,6 +364,7 @@ async def get_user(
"""
Get a single user by ID according to SCIM v2 protocol
"""
verbose_proxy_logger.debug("SCIM GET USER request for user_id=%s", user_id)
try:
user = await _check_user_exists(user_id)
@ -357,7 +388,9 @@ async def create_user(
Create a user according to SCIM v2 protocol
"""
try:
verbose_proxy_logger.debug("SCIM CREATE USER request: %s", user)
verbose_proxy_logger.debug(
"SCIM CREATE USER request: %s", user.model_dump()
)
prisma_client = await _get_prisma_client_or_raise_exception()
# Extract data from SCIM user
@ -435,7 +468,11 @@ async def update_user(
"""
Update a user according to SCIM v2 protocol (full replacement)
"""
verbose_proxy_logger.debug("SCIM PUT USER request: %s", user)
verbose_proxy_logger.debug(
"SCIM PUT USER request for user_id=%s: %s",
user_id,
user.model_dump(),
)
try:
prisma_client = await _get_prisma_client_or_raise_exception()
@ -497,6 +534,9 @@ async def delete_user(
"""
Delete a user according to SCIM v2 protocol
"""
verbose_proxy_logger.debug(
"SCIM DELETE USER request for user_id=%s", user_id
)
try:
prisma_client = await _get_prisma_client_or_raise_exception()
existing_user = await _check_user_exists(user_id)
@ -691,7 +731,11 @@ async def patch_user(
"""
Patch a user according to SCIM v2 protocol
"""
verbose_proxy_logger.debug("SCIM PATCH USER request: %s", patch_ops)
verbose_proxy_logger.debug(
"SCIM PATCH USER request for user_id=%s: %s",
user_id,
patch_ops.model_dump(),
)
try:
prisma_client = await _get_prisma_client_or_raise_exception()
@ -744,6 +788,12 @@ async def get_groups(
"""
Get a list of groups according to SCIM v2 protocol
"""
verbose_proxy_logger.debug(
"SCIM GET GROUPS request: startIndex=%s count=%s filter=%s",
startIndex,
count,
filter,
)
try:
prisma_client = await _get_prisma_client_or_raise_exception()
# Parse filter if provided (basic support)
@ -814,6 +864,9 @@ async def get_group(
"""
Get a single group by ID according to SCIM v2 protocol
"""
verbose_proxy_logger.debug(
"SCIM GET GROUP request for group_id=%s", group_id
)
try:
team = await _check_team_exists(group_id)
@ -839,6 +892,10 @@ async def create_group(
"""
Create a group according to SCIM v2 protocol
"""
verbose_proxy_logger.debug(
"SCIM CREATE GROUP request: %s",
group.model_dump(),
)
try:
prisma_client = await _get_prisma_client_or_raise_exception()
@ -892,78 +949,51 @@ async def update_group(
"""
Update a group according to SCIM v2 protocol
"""
verbose_proxy_logger.debug(
"SCIM PUT GROUP request for group_id=%s: %s",
group_id,
group.model_dump(),
)
try:
prisma_client = await _get_prisma_client_or_raise_exception()
existing_team = await _check_team_exists(group_id)
# Extract valid member IDs
member_ids = await _extract_group_member_ids(group)
verbose_proxy_logger.debug(f"SCIM PUT GROUP member_ids: {member_ids}")
# Update team in database
# Prepare update data
existing_metadata = existing_team.metadata if existing_team.metadata else {}
updated_metadata = {**existing_metadata, "scim_data": group.model_dump()}
update_data = {
"team_alias": group.displayName,
"metadata": safe_dumps(updated_metadata),
}
# Update team in database
updated_team = await prisma_client.db.litellm_teamtable.update(
where={"team_id": group_id},
data={
"team_alias": group.displayName,
"members": member_ids,
"metadata": safe_dumps(updated_metadata),
},
data=update_data,
)
# Handle user-team relationships
current_members = existing_team.members or []
# Add new members to team
for member_id in member_ids:
if member_id not in current_members:
user = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": member_id}
)
if user:
current_user_teams = user.teams or []
if group_id not in current_user_teams:
await prisma_client.db.litellm_usertable.update(
where={"user_id": member_id},
data={"teams": {"push": group_id}},
)
# Remove former members from team
for member_id in current_members:
if member_id not in member_ids:
user = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": member_id}
)
if user:
current_user_teams = user.teams or []
if group_id in current_user_teams:
new_teams = [t for t in current_user_teams if t != group_id]
await prisma_client.db.litellm_usertable.update(
where={"user_id": member_id}, data={"teams": new_teams}
)
# Get updated members for response
members = await _get_team_members_display(member_ids)
team_created_at = (
updated_team.created_at.isoformat() if updated_team.created_at else None
)
team_updated_at = (
updated_team.updated_at.isoformat() if updated_team.updated_at else None
# Handle user-team relationship changes using the same approach as patch_group
current_members = set(await _get_team_member_user_ids_from_team(existing_team))
verbose_proxy_logger.debug(f"SCIM PUT GROUP current_members: {current_members}")
final_members = set(member_ids)
verbose_proxy_logger.debug(f"SCIM PUT GROUP final_members: {final_members}")
await _handle_group_membership_changes(
group_id=group_id,
current_members=current_members,
final_members=final_members,
)
return SCIMGroup(
schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"],
id=group_id,
displayName=updated_team.team_alias or group_id,
members=members,
meta={
"resourceType": "Group",
"created": team_created_at,
"lastModified": team_updated_at,
},
# Convert to SCIM format and return
scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(
updated_team
)
return scim_group
except Exception as e:
raise handle_exception_on_proxy(e)
@ -980,6 +1010,9 @@ async def delete_group(
"""
Delete a group according to SCIM v2 protocol
"""
verbose_proxy_logger.debug(
"SCIM DELETE GROUP request for group_id=%s", group_id
)
try:
prisma_client = await _get_prisma_client_or_raise_exception()
existing_team = await _check_team_exists(group_id)
@ -1135,7 +1168,11 @@ async def patch_group(
"""
Patch a group according to SCIM v2 protocol
"""
verbose_proxy_logger.debug("SCIM PATCH GROUP request: %s", patch_ops)
verbose_proxy_logger.debug(
"SCIM PATCH GROUP request for group_id=%s: %s",
group_id,
patch_ops.model_dump(),
)
try:
prisma_client = await _get_prisma_client_or_raise_exception()
@ -1147,7 +1184,7 @@ async def patch_group(
)
# Track current members for comparison
current_members = set(existing_team.members or [])
current_members = set(await _get_team_member_user_ids_from_team(existing_team))
# Apply updates to the database
updated_team = await _apply_group_patch_updates(

View file

@ -0,0 +1,383 @@
"""
OpenAI Passthrough Logging Handler
Handles cost tracking and logging for OpenAI passthrough endpoints, specifically /chat/completions.
"""
from datetime import datetime
from typing import List, Optional, Union
from urllib.parse import urlparse
import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import (
get_standard_logging_object_payload,
)
from litellm.llms.openai.openai import OpenAIConfig
from litellm.llms.openai.openai import OpenAIConfig as OpenAIConfigType
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import (
BasePassthroughLoggingHandler,
)
from litellm.proxy.pass_through_endpoints.success_handler import (
PassThroughEndpointLogging,
)
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
EndpointType,
PassthroughStandardLoggingPayload,
)
from litellm.types.utils import LlmProviders
from litellm.utils import ModelResponse, TextCompletionResponse
class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
"""
OpenAI-specific passthrough logging handler that provides cost tracking for /chat/completions endpoints.
"""
@property
def llm_provider_name(self) -> LlmProviders:
return LlmProviders.OPENAI
@staticmethod
def get_provider_config(model: str) -> OpenAIConfigType:
"""Get OpenAI provider configuration for the given model."""
return OpenAIConfig()
@staticmethod
def is_openai_chat_completions_route(url_route: str) -> bool:
"""Check if the URL route is an OpenAI chat completions endpoint."""
if not url_route:
return False
parsed_url = urlparse(url_route)
return bool(
parsed_url.hostname
and (
"api.openai.com" in parsed_url.hostname
or "openai.azure.com" in parsed_url.hostname
)
and "/v1/chat/completions" in parsed_url.path
)
@staticmethod
def _get_user_from_metadata(
passthrough_logging_payload: PassthroughStandardLoggingPayload,
) -> Optional[str]:
"""Extract user information from passthrough logging payload."""
request_body = passthrough_logging_payload.get("request_body")
if request_body:
return request_body.get("user")
return None
@staticmethod
def openai_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:
"""
Handle OpenAI passthrough logging with cost tracking for chat completions.
"""
# Only handle chat completions endpoints
if not OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(
url_route
):
# For non-chat-completions endpoints, use the base handler without cost tracking
base_handler = OpenAIPassthroughLoggingHandler()
return base_handler.passthrough_chat_handler(
httpx_response=httpx_response,
response_body=response_body,
logging_obj=logging_obj,
url_route=url_route,
result=result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
request_body=request_body,
**kwargs,
)
# Extract model from request or response
model = request_body.get("model", response_body.get("model", ""))
if not model:
verbose_proxy_logger.warning(
"No model found in request or response for OpenAI passthrough cost tracking"
)
base_handler = OpenAIPassthroughLoggingHandler()
return base_handler.passthrough_chat_handler(
httpx_response=httpx_response,
response_body=response_body,
logging_obj=logging_obj,
url_route=url_route,
result=result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
request_body=request_body,
**kwargs,
)
try:
# Transform the response to LiteLLM format for cost calculation
provider_config = OpenAIPassthroughLoggingHandler.get_provider_config(
model=model
)
litellm_model_response: ModelResponse = provider_config.transform_response(
raw_response=httpx_response,
model_response=litellm.ModelResponse(),
model=model,
messages=request_body.get("messages", []),
logging_obj=logging_obj,
optional_params=request_body.get("optional_params", {}),
api_key="",
request_data=request_body,
encoding=litellm.encoding,
json_mode=request_body.get("response_format", {}).get("type")
== "json_object",
litellm_params={},
)
# Calculate cost using LiteLLM's cost calculator
response_cost = litellm.completion_cost(
completion_response=litellm_model_response,
model=model,
custom_llm_provider="openai",
)
# Update kwargs with cost information
kwargs["response_cost"] = response_cost
kwargs["model"] = model
kwargs["custom_llm_provider"] = "openai"
# Extract user information for tracking
passthrough_logging_payload: Optional[
PassthroughStandardLoggingPayload
] = kwargs.get("passthrough_logging_payload")
if passthrough_logging_payload:
user = OpenAIPassthroughLoggingHandler._get_user_from_metadata(
passthrough_logging_payload=passthrough_logging_payload,
)
if user:
kwargs.setdefault("litellm_params", {})
kwargs["litellm_params"].update(
{"proxy_server_request": {"body": {"user": user}}}
)
# Create standard logging object
get_standard_logging_object_payload(
kwargs=kwargs,
init_response_obj=litellm_model_response,
start_time=start_time,
end_time=end_time,
logging_obj=logging_obj,
status="success",
)
# Update logging object with cost information
logging_obj.model_call_details["model"] = model
logging_obj.model_call_details["custom_llm_provider"] = "openai"
logging_obj.model_call_details["response_cost"] = response_cost
verbose_proxy_logger.debug(
f"OpenAI passthrough cost tracking - Model: {model}, Cost: ${response_cost:.6f}"
)
return {
"result": litellm_model_response,
"kwargs": kwargs,
}
except Exception as e:
verbose_proxy_logger.error(
f"Error in OpenAI passthrough cost tracking: {str(e)}"
)
# Fall back to base handler without cost tracking
base_handler = OpenAIPassthroughLoggingHandler()
return base_handler.passthrough_chat_handler(
httpx_response=httpx_response,
response_body=response_body,
logging_obj=logging_obj,
url_route=url_route,
result=result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
request_body=request_body,
**kwargs,
)
def _build_complete_streaming_response(
self,
all_chunks: list,
litellm_logging_obj: LiteLLMLoggingObj,
model: str,
) -> Optional[Union[ModelResponse, TextCompletionResponse]]:
"""
Builds complete response from raw chunks for OpenAI streaming responses.
- Converts str chunks to generic chunks
- Converts generic chunks to litellm chunks (OpenAI format)
- Builds complete response from litellm chunks
"""
try:
# OpenAI's response iterator to parse chunks
from litellm.llms.openai.openai import OpenAIChatCompletionResponseIterator
openai_iterator = OpenAIChatCompletionResponseIterator(
streaming_response=None,
sync_stream=False,
)
all_openai_chunks = []
for chunk_str in all_chunks:
try:
# Parse the string chunk using the base iterator's string parser
from litellm.llms.base_llm.base_model_iterator import (
BaseModelResponseIterator,
)
# Convert string chunk to dict
stripped_json_chunk = (
BaseModelResponseIterator._string_to_dict_parser(
str_line=chunk_str
)
)
if stripped_json_chunk:
# Parse the chunk using OpenAI's chunk parser
transformed_chunk = openai_iterator.chunk_parser(
chunk=stripped_json_chunk
)
if transformed_chunk is not None:
all_openai_chunks.append(transformed_chunk)
except (StopIteration, StopAsyncIteration, Exception) as e:
verbose_proxy_logger.debug(f"Error parsing streaming chunk: {e}")
continue
if not all_openai_chunks:
verbose_proxy_logger.warning(
"No valid chunks found in streaming response"
)
return None
# Build complete response from chunks
complete_streaming_response = litellm.stream_chunk_builder(
chunks=all_openai_chunks
)
return complete_streaming_response
except Exception as e:
verbose_proxy_logger.error(
f"Error building complete streaming response: {str(e)}"
)
return None
@staticmethod
def _handle_logging_openai_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],
end_time: datetime,
) -> PassThroughEndpointLoggingTypedDict:
"""
Handle logging for collected OpenAI streaming chunks with cost tracking.
"""
try:
# Extract model from request body
model = request_body.get("model", "gpt-4o")
# Build complete response from chunks using our streaming handler
handler = OpenAIPassthroughLoggingHandler()
complete_response = handler._build_complete_streaming_response(
all_chunks=all_chunks,
litellm_logging_obj=litellm_logging_obj,
model=model,
)
if complete_response is None:
verbose_proxy_logger.warning(
"Failed to build complete response from OpenAI streaming chunks"
)
return {
"result": None,
"kwargs": {},
}
# Calculate cost using LiteLLM's cost calculator
response_cost = litellm.completion_cost(
completion_response=complete_response,
model=model,
custom_llm_provider="openai",
)
# Prepare kwargs for logging
kwargs = {
"response_cost": response_cost,
"model": model,
"custom_llm_provider": "openai",
}
# Extract user information for tracking
passthrough_logging_payload: Optional[
PassthroughStandardLoggingPayload
] = litellm_logging_obj.model_call_details.get(
"passthrough_logging_payload"
)
if passthrough_logging_payload:
user = OpenAIPassthroughLoggingHandler._get_user_from_metadata(
passthrough_logging_payload=passthrough_logging_payload,
)
if user:
kwargs.setdefault("litellm_params", {})
kwargs["litellm_params"].update(
{"proxy_server_request": {"body": {"user": user}}}
)
# Create standard logging object
get_standard_logging_object_payload(
kwargs=kwargs,
init_response_obj=complete_response,
start_time=start_time,
end_time=end_time,
logging_obj=litellm_logging_obj,
status="success",
)
# Update logging object with cost information
litellm_logging_obj.model_call_details["model"] = model
litellm_logging_obj.model_call_details["custom_llm_provider"] = "openai"
litellm_logging_obj.model_call_details["response_cost"] = response_cost
verbose_proxy_logger.debug(
f"OpenAI streaming passthrough cost tracking - Model: {model}, Cost: ${response_cost:.6f}"
)
return {
"result": complete_response,
"kwargs": kwargs,
}
except Exception as e:
verbose_proxy_logger.error(
f"Error in OpenAI streaming passthrough cost tracking: {str(e)}"
)
return {
"result": None,
"kwargs": {},
}

View file

@ -314,6 +314,12 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
return EndpointType.VERTEX_AI
elif parsed_url.hostname == "api.anthropic.com":
return EndpointType.ANTHROPIC
elif (
parsed_url.hostname == "api.openai.com"
or parsed_url.hostname == "openai.azure.com"
or (parsed_url.hostname and "openai.com" in parsed_url.hostname)
):
return EndpointType.OPENAI
return EndpointType.GENERIC
@staticmethod
@ -415,10 +421,10 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
for field_name, field_value in form_data.items():
if isinstance(field_value, (StarletteUploadFile, UploadFile)):
files[field_name] = (
await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file(
upload_file=field_value
)
files[
field_name
] = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file(
upload_file=field_value
)
else:
form_data_dict[field_name] = field_value
@ -497,9 +503,9 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
"passthrough_logging_payload": passthrough_logging_payload,
}
logging_obj.model_call_details["passthrough_logging_payload"] = (
passthrough_logging_payload
)
logging_obj.model_call_details[
"passthrough_logging_payload"
] = passthrough_logging_payload
return kwargs
@ -531,10 +537,10 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
subpath = subpath[1:]
return base_target + subpath
@staticmethod
def _update_stream_param_based_on_request_body(
parsed_body: dict,
parsed_body: dict,
stream: Optional[bool] = None,
) -> Optional[bool]:
"""
@ -699,9 +705,11 @@ async def pass_through_request( # noqa: PLR0915
"headers": headers,
},
)
stream = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body(
parsed_body=_parsed_body,
stream=stream,
stream = (
HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body(
parsed_body=_parsed_body,
stream=stream,
)
)
if stream:

View file

@ -14,6 +14,9 @@ from litellm.types.utils import StandardPassThroughResponseObject
from .llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
)
from .llm_provider_handlers.openai_passthrough_logging_handler import (
OpenAIPassthroughLoggingHandler,
)
from .llm_provider_handlers.vertex_passthrough_logging_handler import (
VertexPassthroughLoggingHandler,
)
@ -78,6 +81,7 @@ class PassThroughStreamingHandler:
Supported endpoint types:
- Anthropic
- Vertex AI
- OpenAI
"""
all_chunks = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(
raw_bytes
@ -119,6 +123,23 @@ class PassThroughStreamingHandler:
vertex_passthrough_logging_handler_result["result"]
)
kwargs = vertex_passthrough_logging_handler_result["kwargs"]
elif endpoint_type == EndpointType.OPENAI:
openai_passthrough_logging_handler_result = (
OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks(
litellm_logging_obj=litellm_logging_obj,
passthrough_success_handler_obj=passthrough_success_handler_obj,
url_route=url_route,
request_body=request_body,
endpoint_type=endpoint_type,
start_time=start_time,
all_chunks=all_chunks,
end_time=end_time,
)
)
standard_logging_response_object = (
openai_passthrough_logging_handler_result["result"]
)
kwargs = openai_passthrough_logging_handler_result["kwargs"]
if standard_logging_response_object is None:
standard_logging_response_object = StandardPassThroughResponseObject(

View file

@ -162,9 +162,32 @@ class PassThroughEndpointLogging:
cohere_passthrough_logging_handler_result["result"]
)
kwargs = cohere_passthrough_logging_handler_result["kwargs"]
return_dict["standard_logging_response_object"] = (
standard_logging_response_object
)
elif self.is_openai_route(url_route):
from .llm_provider_handlers.openai_passthrough_logging_handler import (
OpenAIPassthroughLoggingHandler,
)
openai_passthrough_logging_handler_result = (
OpenAIPassthroughLoggingHandler.openai_passthrough_handler(
httpx_response=httpx_response,
response_body=response_body or {},
logging_obj=logging_obj,
url_route=url_route,
result=result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
request_body=request_body,
**kwargs,
)
)
standard_logging_response_object = (
openai_passthrough_logging_handler_result["result"]
)
kwargs = openai_passthrough_logging_handler_result["kwargs"]
return_dict[
"standard_logging_response_object"
] = standard_logging_response_object
return_dict["kwargs"] = kwargs
return return_dict
@ -185,9 +208,9 @@ class PassThroughEndpointLogging:
standard_logging_response_object: Optional[
PassThroughEndpointLoggingResultValues
] = None
logging_obj.model_call_details["passthrough_logging_payload"] = (
passthrough_logging_payload
)
logging_obj.model_call_details[
"passthrough_logging_payload"
] = passthrough_logging_payload
if self.is_assemblyai_route(url_route):
if (
AssemblyAIPassthroughLoggingHandler._should_log_request(
@ -286,6 +309,16 @@ class PassThroughEndpointLogging:
return True
return False
def is_openai_route(self, url_route: str):
"""Check if the URL route is an OpenAI API route."""
if not url_route:
return False
parsed_url = urlparse(url_route)
return parsed_url.hostname and (
"api.openai.com" in parsed_url.hostname
or "openai.azure.com" in parsed_url.hostname
)
def _set_cost_per_request(
self,
logging_obj: LiteLLMLoggingObj,
@ -305,8 +338,8 @@ class PassThroughEndpointLogging:
kwargs["response_cost"] = passthrough_logging_payload.get(
"cost_per_request"
)
logging_obj.model_call_details["response_cost"] = (
passthrough_logging_payload.get("cost_per_request")
)
logging_obj.model_call_details[
"response_cost"
] = passthrough_logging_payload.get("cost_per_request")
return kwargs

View file

@ -1029,4 +1029,4 @@ def list_input_items(
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)
)

View file

@ -379,4 +379,4 @@ class ResponseAPILoggingUtils:
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
)
)

View file

@ -573,3 +573,84 @@ class AmazonDeepSeekR1StreamingResponse(TypedDict):
generation_token_count: int
stop_reason: Optional[str]
prompt_token_count: int
################ Bedrock Batch Types #################
class BedrockS3InputDataConfig(TypedDict):
"""S3 input data configuration for Bedrock batch jobs."""
s3Uri: str
class BedrockInputDataConfig(TypedDict):
"""Input data configuration for Bedrock batch jobs."""
s3InputDataConfig: BedrockS3InputDataConfig
class BedrockS3OutputDataConfig(TypedDict):
"""S3 output data configuration for Bedrock batch jobs."""
s3Uri: str
class BedrockOutputDataConfig(TypedDict):
"""Output data configuration for Bedrock batch jobs."""
s3OutputDataConfig: BedrockS3OutputDataConfig
class BedrockCreateBatchRequest(TypedDict, total=False):
"""
Request structure for creating a Bedrock batch inference job.
Reference: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_CreateModelInvocationJob.html
"""
jobName: str
roleArn: str
modelId: str
inputDataConfig: BedrockInputDataConfig
outputDataConfig: BedrockOutputDataConfig
timeoutDurationInHours: Optional[int]
clientRequestToken: Optional[str]
tags: Optional[List[dict]]
BedrockBatchJobStatus = Literal[
"Submitted",
"InProgress",
"Completed",
"Failed",
"Stopping",
"Stopped"
]
class BedrockCreateBatchResponse(TypedDict):
"""
Response structure from creating a Bedrock batch inference job.
Reference: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_CreateModelInvocationJob.html
"""
jobArn: str
jobName: str
status: BedrockBatchJobStatus
class BedrockGetBatchResponse(TypedDict, total=False):
"""
Response structure from getting a Bedrock batch inference job.
Reference: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_GetModelInvocationJob.html
"""
jobArn: str
jobName: str
modelId: str
roleArn: str
status: BedrockBatchJobStatus
message: Optional[str]
submitTime: Optional[str]
lastModifiedTime: Optional[str]
endTime: Optional[str]
inputDataConfig: BedrockInputDataConfig
outputDataConfig: BedrockOutputDataConfig
timeoutDurationInHours: Optional[int]
clientRequestToken: Optional[str]

View file

@ -5,6 +5,7 @@ from typing import Optional, TypedDict
class EndpointType(str, Enum):
VERTEX_AI = "vertex-ai"
ANTHROPIC = "anthropic"
OPENAI = "openai"
GENERIC = "generic"

View file

@ -216,6 +216,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
auto_router_default_model: Optional[str] = None
auto_router_embedding_model: Optional[str] = None
# Batch/File API Params
s3_bucket_name: Optional[str] = None
gcs_bucket_name: Optional[str] = None
def __init__(
self,
custom_llm_provider: Optional[str] = None,
@ -265,6 +269,9 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
auto_router_config: Optional[str] = None,
auto_router_default_model: Optional[str] = None,
auto_router_embedding_model: Optional[str] = None,
# Batch/File API Params
s3_bucket_name: Optional[str] = None,
gcs_bucket_name: Optional[str] = None,
**params,
):
args = locals()

View file

@ -234,6 +234,7 @@ from litellm.llms.base_llm.base_utils import (
BaseLLMModelInfo,
type_to_response_format_param,
)
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
from litellm.llms.base_llm.chat.transformation import BaseConfig
from litellm.llms.base_llm.completion.transformation import BaseTextCompletionConfig
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
@ -2395,7 +2396,7 @@ def _should_drop_param(k, additional_drop_params) -> bool:
def _get_non_default_params(
passed_params: dict, default_params: dict, additional_drop_params: Optional[bool]
passed_params: dict, default_params: dict, additional_drop_params: Optional[list]
) -> dict:
non_default_params = {}
for k, v in passed_params.items():
@ -2509,7 +2510,7 @@ def get_optional_params_image_gen(
user: Optional[str] = None,
input_fidelity: Optional[str] = None,
custom_llm_provider: Optional[str] = None,
additional_drop_params: Optional[bool] = None,
additional_drop_params: Optional[list] = None,
provider_config: Optional[BaseImageGenerationConfig] = None,
drop_params: Optional[bool] = None,
**kwargs,
@ -2628,9 +2629,20 @@ def get_optional_params_image_gen(
) # Default to square if size not recognized
optional_params["aspectRatio"] = aspect_ratio
for k in passed_params.keys():
if k not in default_params.keys():
optional_params[k] = passed_params[k]
openai_params: list[str] = list(default_params.keys())
if provider_config is not None:
supported_params = provider_config.get_supported_openai_params(
model=model or ""
)
openai_params = list(supported_params)
optional_params = add_provider_specific_params_to_optional_params(
optional_params=optional_params,
passed_params=passed_params,
custom_llm_provider=custom_llm_provider or "",
openai_params=openai_params,
additional_drop_params=additional_drop_params,
)
return optional_params
@ -7287,6 +7299,20 @@ class ProviderConfigManager:
from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig
return VertexAIFilesConfig()
elif LlmProviders.BEDROCK == provider:
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
return BedrockFilesConfig()
return None
@staticmethod
def get_provider_batches_config(
model: str,
provider: LlmProviders,
) -> Optional[BaseBatchesConfig]:
if LlmProviders.BEDROCK == provider:
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
return BedrockBatchesConfig()
return None
@staticmethod

View file

@ -6157,21 +6157,7 @@
"supports_tool_choice": true,
"source": "https://inference-docs.cerebras.ai/support/pricing"
},
"cerebras/openai/gpt-oss-20b": {
"max_tokens": 32768,
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"input_cost_per_token": 7e-08,
"output_cost_per_token": 3e-07,
"litellm_provider": "cerebras",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"source": "https://inference-docs.cerebras.ai/support/pricing"
},
"cerebras/openai/gpt-oss-120b": {
"max_tokens": 32768,
"max_input_tokens": 131072,
@ -9498,6 +9484,48 @@
"source": "https://aistudio.google.com",
"supports_tool_choice": true
},
"gemini/veo-3.0-generate-preview": {
"max_tokens": 1024,
"max_input_tokens": 1024,
"output_cost_per_second": 0.75,
"litellm_provider": "gemini",
"mode": "video_generation",
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"video"
],
"source": "https://ai.google.dev/gemini-api/docs/video"
},
"gemini/veo-3.0-fast-generate-preview": {
"max_tokens": 1024,
"max_input_tokens": 1024,
"output_cost_per_second": 0.40,
"litellm_provider": "gemini",
"mode": "video_generation",
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"video"
],
"source": "https://ai.google.dev/gemini-api/docs/video"
},
"gemini/veo-2.0-generate-001": {
"max_tokens": 1024,
"max_input_tokens": 1024,
"output_cost_per_second": 0.35,
"litellm_provider": "gemini",
"mode": "video_generation",
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"video"
],
"source": "https://ai.google.dev/gemini-api/docs/video"
},
"vertex_ai/claude-opus-4-1": {
"max_tokens": 4096,
"max_input_tokens": 200000,
@ -10315,6 +10343,48 @@
"mode": "image_generation",
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
},
"vertex_ai/veo-3.0-generate-preview": {
"max_tokens": 1024,
"max_input_tokens": 1024,
"output_cost_per_second": 0.75,
"litellm_provider": "vertex_ai-video-models",
"mode": "video_generation",
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"video"
],
"source": "https://ai.google.dev/gemini-api/docs/video"
},
"vertex_ai/veo-3.0-fast-generate-preview": {
"max_tokens": 1024,
"max_input_tokens": 1024,
"output_cost_per_second": 0.40,
"litellm_provider": "vertex_ai-video-models",
"mode": "video_generation",
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"video"
],
"source": "https://ai.google.dev/gemini-api/docs/video"
},
"vertex_ai/veo-2.0-generate-001": {
"max_tokens": 1024,
"max_input_tokens": 1024,
"output_cost_per_second": 0.35,
"litellm_provider": "vertex_ai-video-models",
"mode": "video_generation",
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"video"
],
"source": "https://ai.google.dev/gemini-api/docs/video"
},
"text-embedding-004": {
"max_tokens": 2048,
"max_input_tokens": 2048,

View file

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

View file

@ -0,0 +1,3 @@
{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0", "messages": [{"role": "system", "content": "You are a helpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}}
{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0", "messages": [{"role": "system", "content": "You are an unhelpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}}

View file

@ -0,0 +1,70 @@
# What is this?
## Unit Tests for OpenAI Batches API
import asyncio
import json
import os
import sys
import traceback
import tempfile
from dotenv import load_dotenv
load_dotenv()
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system-path
import pytest
from typing import Optional
import litellm
@pytest.mark.asyncio()
async def test_async_create_file():
"""
1. Create File for Batch completion
2. Create Batch Request
3. Retrieve the specific batch
"""
litellm._turn_on_debug()
print("Testing async create batch")
file_name = "bedrock_batch_completions.jsonl"
_current_dir = os.path.dirname(os.path.abspath(__file__))
file_path = os.path.join(_current_dir, file_name)
file_obj = await litellm.acreate_file(
file=open(file_path, "rb"),
purpose="batch",
custom_llm_provider="bedrock",
s3_bucket_name="litellm-proxy",
)
@pytest.mark.asyncio()
async def test_async_file_and_batch():
"""
Test file retrieval
"""
litellm._turn_on_debug()
file_name = "bedrock_batch_completions.jsonl"
_current_dir = os.path.dirname(os.path.abspath(__file__))
file_path = os.path.join(_current_dir, file_name)
file_obj = await litellm.acreate_file(
file=open(file_path, "rb"),
purpose="batch",
custom_llm_provider="bedrock",
s3_bucket_name="litellm-proxy",
)
print("CREATED FILE RESPONSE=", file_obj)
# create batch
create_batch_response = await litellm.acreate_batch(
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id=file_obj.id,
metadata={"key1": "value1", "key2": "value2"},
custom_llm_provider="bedrock",
aws_batch_role_arn="arn:aws:iam::888602223428:role/service-role/AmazonBedrockExecutionRoleForAgents_BB9HNW6V4CV"
)
print("CREATED BATCH RESPONSE=", create_batch_response)

View file

@ -528,7 +528,6 @@ class BaseResponsesAPITest(ABC):
# Validate final response structure
validate_responses_api_response(final_response, final_chunk=True)
assert final_response.output is not None
assert len(final_response.output) > 0
@pytest.mark.asyncio
async def test_text_format_to_text_conversion(self):

View file

@ -1,6 +1,16 @@
from base_llm_unit_tests import BaseLLMChatTest
import os
import sys
import pytest
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
from base_llm_unit_tests import BaseLLMChatTest
from litellm.llms.groq.chat.transformation import GroqChatConfig
class TestGroq(BaseLLMChatTest):
def get_base_completion_call_args(self) -> dict:
return {
@ -10,3 +20,9 @@ class TestGroq(BaseLLMChatTest):
def test_tool_call_no_arguments(self, tool_call_no_arguments):
"""Test that tool calls with no arguments is translated correctly. Relevant issue: https://github.com/BerriAI/litellm/issues/6833"""
pass
@pytest.mark.parametrize("model", ["groq/qwen/qwen3-32b", "groq/openai/gpt-oss-20b", "groq/openai/gpt-oss-120b"])
def test_reasoning_effort_in_supported_params(self, model):
"""Test that reasoning_effort is in the list of supported parameters for Groq"""
supported_params = GroqChatConfig().get_supported_openai_params(model=model)
assert "reasoning_effort" in supported_params

View file

@ -839,7 +839,6 @@ from test_completion import response_format_tests
"model,region",
[
("vertex_ai/mistral-large-2411", "us-central1"),
("vertex_ai/mistral-nemo@2407", "us-central1"),
("vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas", "us-south1"),
("vertex_ai/openai/gpt-oss-20b-maas", "us-central1"),
],

View file

@ -934,3 +934,204 @@ async def test_team_member_rate_limits_v3():
assert team_member_descriptor["value"] == f"{_team_id}:{_user_id}", "Team member value should combine team_id and user_id"
assert team_member_descriptor["rate_limit"]["requests_per_unit"] == 10, "Team member RPM limit should be set"
assert team_member_descriptor["rate_limit"]["tokens_per_unit"] == 1000, "Team member TPM limit should be set"
@pytest.mark.asyncio
async def test_async_increment_tokens_with_ttl_preservation():
"""
Test TTL preservation functionality for token increment operations.
This test verifies that:
1. Keys are created with proper TTL on first increment
2. TTL is preserved on subsequent increments (not reset)
3. Both TTL and non-TTL operations work correctly in the same call
Environment variables required:
- REDIS_HOST: Redis server hostname
- REDIS_PORT: Redis server port
- REDIS_PASSWORD: Redis password (optional)
Test scenario:
1. First call: Create keys with TTL=60s and TTL=None
2. Wait 2 seconds
3. Second call: Increment same keys
4. Verify TTL decreased but wasn't reset to 60s
"""
import os
import time
from litellm.caching.redis_cache import RedisCache
from litellm.types.caching import RedisPipelineIncrementOperation
# Skip test if Redis environment variables are not set
redis_host = os.getenv("REDIS_HOST")
redis_port = os.getenv("REDIS_PORT")
redis_password = os.getenv("REDIS_PASSWORD")
if not redis_host or not redis_port:
pytest.skip("Redis environment variables (REDIS_HOST, REDIS_PORT) not set")
# Setup Redis cache
redis_cache = RedisCache(
host=redis_host,
port=int(redis_port),
password=redis_password,
)
local_cache = DualCache(redis_cache=redis_cache)
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
# Verify Redis connection is working
try:
await redis_cache.ping()
except Exception as e:
pytest.skip(f"Redis connection failed: {str(e)}")
# Test keys
test_key_with_ttl = "test_ttl_preservation:with_ttl"
test_key_without_ttl = "test_ttl_preservation:without_ttl"
try:
# Clean up any existing test keys
try:
await redis_cache.async_delete_cache(test_key_with_ttl)
await redis_cache.async_delete_cache(test_key_without_ttl)
except Exception:
# Keys might not exist, ignore cleanup errors
pass
# First increment: Create operations with mixed TTL scenarios
pipeline_operations_first = [
RedisPipelineIncrementOperation(
key=test_key_with_ttl,
increment_value=10.0,
ttl=60
),
RedisPipelineIncrementOperation(
key=test_key_without_ttl,
increment_value=5.0,
ttl=None # No TTL
)
]
# Execute first increment
await parallel_request_handler.async_increment_tokens_with_ttl_preservation(
pipeline_operations=pipeline_operations_first
)
# Verify keys exist and check initial TTL
ttl_after_first = await redis_cache.async_get_ttl(test_key_with_ttl)
value_after_first_with_ttl = await redis_cache.async_get_cache(test_key_with_ttl)
value_after_first_without_ttl = await redis_cache.async_get_cache(test_key_without_ttl)
assert value_after_first_with_ttl == 10.0, "First increment should set value to 10.0"
assert value_after_first_without_ttl == 5.0, "First increment should set value to 5.0"
assert ttl_after_first is not None and ttl_after_first > 0, "Key with TTL should have positive TTL after first increment"
assert ttl_after_first <= 60, "TTL should not exceed the set value"
# Check TTL for key without TTL (should be None, meaning no expiry)
ttl_no_ttl_key = await redis_cache.async_get_ttl(test_key_without_ttl)
assert ttl_no_ttl_key is None, "Key without TTL should have no expiry (None from async_get_ttl)"
# Wait a moment to ensure TTL decreases
await asyncio.sleep(2)
# Second increment: Same operations to test TTL preservation
pipeline_operations_second = [
RedisPipelineIncrementOperation(
key=test_key_with_ttl,
increment_value=15.0,
ttl=60 # Same TTL value
),
RedisPipelineIncrementOperation(
key=test_key_without_ttl,
increment_value=7.0,
ttl=None # No TTL
)
]
# Execute second increment
await parallel_request_handler.async_increment_tokens_with_ttl_preservation(
pipeline_operations=pipeline_operations_second
)
# Verify TTL preservation and value updates
ttl_after_second = await redis_cache.async_get_ttl(test_key_with_ttl)
value_after_second_with_ttl = await redis_cache.async_get_cache(test_key_with_ttl)
value_after_second_without_ttl = await redis_cache.async_get_cache(test_key_without_ttl)
assert value_after_second_with_ttl == 25.0, "Second increment should update value to 25.0"
assert value_after_second_without_ttl == 12.0, "Second increment should update value to 12.0"
# Critical test: TTL should be preserved (not reset to 60)
assert ttl_after_second is not None, "TTL should still exist"
assert ttl_after_second < ttl_after_first, "TTL should have decreased (not been reset)"
assert ttl_after_second > 0, "TTL should still be positive"
# TTL should not be close to the original 60 seconds (proving it wasn't reset)
assert ttl_after_second < 59, "TTL should be significantly less than original, proving preservation"
# Key without TTL should still have no expiry
ttl_no_ttl_key_after_second = await redis_cache.async_get_ttl(test_key_without_ttl)
assert ttl_no_ttl_key_after_second is None, "Key without TTL should still have no expiry"
finally:
# Clean up test keys
try:
await redis_cache.async_delete_cache(test_key_with_ttl)
await redis_cache.async_delete_cache(test_key_without_ttl)
except Exception:
# Ignore cleanup errors
pass
# Properly close Redis connections to prevent warnings
try:
await redis_cache.disconnect()
except Exception:
# Ignore disconnect errors
pass
@pytest.mark.asyncio
async def test_async_increment_tokens_fallback_behavior():
"""
Test fallback behavior when Lua script is not available.
"""
from litellm.types.caching import RedisPipelineIncrementOperation
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
# Mock the token_increment_script to None to simulate unavailable script
parallel_request_handler.token_increment_script = None
# Mock the fallback method
fallback_called = False
original_method = parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline
async def mock_fallback(*args, **kwargs):
nonlocal fallback_called
fallback_called = True
return await original_method(*args, **kwargs)
parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = mock_fallback
# Test operations
pipeline_operations = [
RedisPipelineIncrementOperation(
key="test_fallback_key",
increment_value=10.0,
ttl=60
)
]
# Execute increment
await parallel_request_handler.async_increment_tokens_with_ttl_preservation(
pipeline_operations=pipeline_operations
)
# Verify fallback was called
assert fallback_called, "Fallback method should be called when Lua script is not available"

View file

@ -224,10 +224,10 @@ class TestScimTransformations:
result = ScimTransformations._get_scim_member_value(member_with_email)
assert result == member_with_email.user_email
# Member without email
# Member without email should fall back to user_id
member_without_email = Member(user_id="user-456", user_email=None, role="user")
result = ScimTransformations._get_scim_member_value(member_without_email)
assert result == ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE
assert result == member_without_email.user_id
class TestSCIMPatchOperations:

View file

@ -10,10 +10,13 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import (
create_user,
get_service_provider_config,
patch_user,
update_group,
update_user,
)
from litellm.types.proxy.management_endpoints.scim_v2 import (
SCIMFeature,
SCIMGroup,
SCIMMember,
SCIMPatchOp,
SCIMPatchOperation,
SCIMServiceProviderConfig,
@ -678,4 +681,233 @@ async def test_update_group_metadata_serialization_issue(mocker):
parsed_metadata = json.loads(metadata)
assert "existing_key" in parsed_metadata
assert "scim_data" in parsed_metadata
assert parsed_metadata["existing_key"] == "existing_value"
@pytest.mark.asyncio
async def test_team_membership_management(mocker):
"""
Test that team membership changes work correctly:
- Adding members to team
- Removing members from team
- members_with_roles is used as source of truth
"""
from litellm.proxy._types import Member
from litellm.proxy.management_endpoints.scim.scim_v2 import (
_get_team_member_user_ids_from_team,
_handle_group_membership_changes,
patch_team_membership,
)
# Mock team with members_with_roles as source of truth
mock_team = mocker.MagicMock()
mock_team.members_with_roles = [
Member(user_id="user1", role="user"),
Member(user_id="user2", role="user")
]
mock_team.members = ["user1", "user2", "user3"] # This should be ignored
# Test that members_with_roles is source of truth
member_ids = await _get_team_member_user_ids_from_team(mock_team)
assert set(member_ids) == {"user1", "user2"}
assert "user3" not in member_ids # Should not be included even though in members
# Mock patch_team_membership function
mock_patch_team_membership = mocker.patch(
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
AsyncMock()
)
# Test adding and removing members
group_id = "test-group-id"
current_members = {"user1", "user2"}
final_members = {"user2", "user3", "user4"} # Remove user1, add user3 and user4
await _handle_group_membership_changes(
group_id=group_id,
current_members=current_members,
final_members=final_members
)
# Verify patch_team_membership was called correctly
assert mock_patch_team_membership.call_count == 3
# Check calls for adding members
add_calls = [call for call in mock_patch_team_membership.call_args_list
if call[1]["teams_ids_to_add_user_to"] == [group_id]]
assert len(add_calls) == 2 # user3 and user4
add_user_ids = {call[1]["user_id"] for call in add_calls}
assert add_user_ids == {"user3", "user4"}
# Check calls for removing members
remove_calls = [call for call in mock_patch_team_membership.call_args_list
if call[1]["teams_ids_to_remove_user_from"] == [group_id]]
assert len(remove_calls) == 1 # user1
remove_user_ids = {call[1]["user_id"] for call in remove_calls}
assert remove_user_ids == {"user1"}
# Verify all calls have correct structure
for call in mock_patch_team_membership.call_args_list:
assert "user_id" in call[1]
assert "teams_ids_to_add_user_to" in call[1]
assert "teams_ids_to_remove_user_from" in call[1]
# Each call should either add OR remove, not both
add_teams = call[1]["teams_ids_to_add_user_to"]
remove_teams = call[1]["teams_ids_to_remove_user_from"]
assert (len(add_teams) > 0) != (len(remove_teams) > 0) # XOR - one should be empty
@pytest.mark.asyncio
async def test_update_group_e2e(mocker):
"""
End-to-end test for update_group endpoint:
- Updates group metadata (displayName)
- Handles complete member replacement (add/remove members)
- Verifies members_with_roles is updated as source of truth
- Tests the full flow from SCIM request to database updates
"""
from litellm.proxy._types import LiteLLM_TeamTable, Member
from litellm.proxy.management_endpoints.scim.scim_transformations import (
ScimTransformations,
)
from litellm.proxy.utils import safe_dumps
# Setup test data
group_id = "test-team-123"
# Mock existing team in database
existing_team = LiteLLM_TeamTable(
team_id=group_id,
team_alias="Old Team Name",
members=["user1", "user2"], # This should be ignored
members_with_roles=[
Member(user_id="user1", role="user"),
Member(user_id="user2", role="user")
],
metadata={"existing_key": "existing_value"}
)
# Mock updated SCIM group request
scim_group_update = SCIMGroup(
schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"],
id=group_id,
displayName="Updated Team Name",
members=[
SCIMMember(value="user2", display="User Two"), # Keep user2
SCIMMember(value="user3", display="User Three"), # Add user3
SCIMMember(value="user4", display="User Four") # Add user4
]
)
# Mock prisma client
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
# Mock database operations
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team)
# Mock the updated team that gets returned from database
updated_team = LiteLLM_TeamTable(
team_id=group_id,
team_alias="Updated Team Name",
members=["user2", "user3", "user4"],
members_with_roles=[
Member(user_id="user2", role="user"),
Member(user_id="user3", role="user"),
Member(user_id="user4", role="user")
],
metadata={
"existing_key": "existing_value",
"scim_data": scim_group_update.model_dump()
}
)
mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=updated_team)
# Mock user validation (all users exist)
mock_user = mocker.MagicMock()
mock_user.user_id = "test-user"
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user)
# Mock dependencies
mocker.patch(
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
AsyncMock(return_value=mock_prisma_client)
)
# Mock patch_team_membership to track membership changes
mock_patch_team_membership = mocker.patch(
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
AsyncMock()
)
# Mock SCIM transformation
expected_scim_response = SCIMGroup(
schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"],
id=group_id,
displayName="Updated Team Name",
members=[
SCIMMember(value="user2", display="user2"),
SCIMMember(value="user3", display="user3"),
SCIMMember(value="user4", display="user4")
]
)
mocker.patch.object(
ScimTransformations,
"transform_litellm_team_to_scim_group",
AsyncMock(return_value=expected_scim_response)
)
# Execute the update_group function
result = await update_group(group_id=group_id, group=scim_group_update)
# Verify database update was called with correct data
mock_prisma_client.db.litellm_teamtable.update.assert_called_once()
update_call_args = mock_prisma_client.db.litellm_teamtable.update.call_args
# Check the update parameters
assert update_call_args[1]["where"]["team_id"] == group_id
update_data = update_call_args[1]["data"]
assert update_data["team_alias"] == "Updated Team Name"
# Verify metadata includes both existing data and SCIM data
metadata_str = update_data["metadata"]
import json
metadata = json.loads(metadata_str)
assert metadata["existing_key"] == "existing_value"
assert "scim_data" in metadata
assert metadata["scim_data"]["displayName"] == "Updated Team Name"
# Verify team membership changes were handled correctly
assert mock_patch_team_membership.call_count == 3 # Remove user1, add user3, add user4
# Check membership changes
call_args_list = mock_patch_team_membership.call_args_list
# Find remove operation (user1)
remove_calls = [call for call in call_args_list
if call[1]["teams_ids_to_remove_user_from"] == [group_id]]
assert len(remove_calls) == 1
assert remove_calls[0][1]["user_id"] == "user1"
assert remove_calls[0][1]["teams_ids_to_add_user_to"] == []
# Find add operations (user3, user4)
add_calls = [call for call in call_args_list
if call[1]["teams_ids_to_add_user_to"] == [group_id]]
assert len(add_calls) == 2
add_user_ids = {call[1]["user_id"] for call in add_calls}
assert add_user_ids == {"user3", "user4"}
# Verify all add calls have empty remove lists
for call in add_calls:
assert call[1]["teams_ids_to_remove_user_from"] == []
# Verify the response
assert result.id == group_id
assert result.displayName == "Updated Team Name"
assert len(result.members) == 3
# Verify SCIM transformation was called with updated team
ScimTransformations.transform_litellm_team_to_scim_group.assert_called_once_with(updated_team)

View file

@ -118,9 +118,9 @@ async def test_key_token_handling(monkeypatch):
@pytest.mark.asyncio
async def test_budget_reset_at_first_of_month(monkeypatch):
async def test_budget_reset_and_expires_at_first_of_month(monkeypatch):
"""
Test that when budget_duration is "1mo", budget_reset_at is set to first of next month
Test that when budget_duration, duration, and key_budget_duration are "1mo", budget_reset_at and expires are set to first of next month
"""
mock_prisma_client = AsyncMock()
mock_insert_data = AsyncMock(
@ -152,10 +152,12 @@ async def test_budget_reset_at_first_of_month(monkeypatch):
# Use monkeypatch to set the prisma_client
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
# Test key generation with budget_duration="1mo"
# Test key generation with budget_duration="1mo", duration="1mo", key_budget_duration="1mo"
response = await generate_key_helper_fn(
request_type="user",
budget_duration="1mo",
duration="1mo",
key_budget_duration="1mo",
user_id="test_user",
)
@ -171,17 +173,17 @@ async def test_budget_reset_at_first_of_month(monkeypatch):
expected_month = now.month + 1
expected_year = now.year
# Parse the response date
response_date = response["budget_reset_at"]
# Verify budget_reset_at is set to first of next month
assert (
response_date.year == expected_year
), f"Expected year {expected_year}, got {response_date.year}"
assert (
response_date.month == expected_month
), f"Expected month {expected_month}, got {response_date.month}"
assert response_date.day == 1, f"Expected day 1, got {response_date.day}"
# Verify budget_reset_at, expires is set to first of next month
for key in ["budget_reset_at", "expires"]:
response_date = response.get(key)
assert response_date is not None, f"{key} not found in response"
assert (
response_date.year == expected_year
), f"Expected year {expected_year}, got {response_date.year} for {key}"
assert (
response_date.month == expected_month
), f"Expected month {expected_month}, got {response_date.month} for {key}"
assert response_date.day == 1, f"Expected day 1, got {response_date.day} for {key}"
@pytest.mark.asyncio

View file

@ -0,0 +1,451 @@
import json
import os
import sys
from datetime import datetime
from typing import Any, Dict, List
from unittest.mock import MagicMock, patch, AsyncMock
import pytest
import httpx
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
OpenAIPassthroughLoggingHandler,
)
from litellm.proxy.pass_through_endpoints.success_handler import (
PassThroughEndpointLogging,
)
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
PassthroughStandardLoggingPayload,
)
class TestOpenAIPassthroughLoggingHandler:
"""Test the OpenAI passthrough logging handler for cost tracking."""
def setup_method(self):
"""Set up test fixtures"""
self.start_time = datetime.now()
self.end_time = datetime.now()
self.handler = OpenAIPassthroughLoggingHandler()
# Mock OpenAI chat completions response
self.mock_openai_response = {
"id": "chatcmpl-123",
"object": "chat.completion",
"created": 1677652288,
"model": "gpt-4o-2024-08-06",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello! How can I help you today?"
},
"finish_reason": "stop"
}
],
"usage": {
"prompt_tokens": 20,
"completion_tokens": 15,
"total_tokens": 35
}
}
def _create_mock_logging_obj(self) -> LiteLLMLoggingObj:
"""Create a mock logging object"""
mock_logging_obj = MagicMock()
mock_logging_obj.model_call_details = {}
return mock_logging_obj
def _create_mock_httpx_response(self, response_data: dict = None) -> httpx.Response:
"""Create a mock httpx response"""
if response_data is None:
response_data = self.mock_openai_response
mock_response = MagicMock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.text = json.dumps(response_data)
mock_response.json.return_value = response_data
mock_response.headers = {"content-type": "application/json"}
return mock_response
def _create_passthrough_logging_payload(self, user: str = "test_user") -> PassthroughStandardLoggingPayload:
"""Create a mock passthrough logging payload"""
return PassthroughStandardLoggingPayload(
url="https://api.openai.com/v1/chat/completions",
request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]},
request_method="POST",
)
def test_llm_provider_name(self):
"""Test that the handler returns the correct provider name"""
assert self.handler.llm_provider_name == "openai"
def test_get_provider_config(self):
"""Test that the handler returns an OpenAI config"""
config = OpenAIPassthroughLoggingHandler.get_provider_config(model="gpt-4o")
assert config is not None
# Verify it's an OpenAI config by checking if it has the expected methods
assert hasattr(config, 'transform_response')
def test_is_openai_chat_completions_route(self):
"""Test OpenAI chat completions route detection"""
# Positive cases
assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.openai.com/v1/chat/completions") == True
assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://openai.azure.com/v1/chat/completions") == True
# Negative cases
assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.openai.com/v1/models") == False
assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("http://localhost:4000/openai/v1/chat/completions") == False
assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.anthropic.com/v1/messages") == False
assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("") == False
@patch('litellm.completion_cost')
@patch('litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload')
def test_openai_passthrough_handler_success(self, mock_get_standard_logging, mock_completion_cost):
"""Test successful cost tracking for OpenAI chat completions"""
# Arrange
mock_completion_cost.return_value = 0.000045
mock_get_standard_logging.return_value = {"test": "logging_payload"}
mock_httpx_response = self._create_mock_httpx_response()
mock_logging_obj = self._create_mock_logging_obj()
passthrough_payload = self._create_passthrough_logging_payload()
kwargs = {
"passthrough_logging_payload": passthrough_payload,
"model": "gpt-4o",
}
# Act
result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler(
httpx_response=mock_httpx_response,
response_body=self.mock_openai_response,
logging_obj=mock_logging_obj,
url_route="https://api.openai.com/v1/chat/completions",
result="",
start_time=self.start_time,
end_time=self.end_time,
cache_hit=False,
request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]},
**kwargs
)
# Assert
assert result is not None
assert "result" in result
assert "kwargs" in result
assert result["kwargs"]["response_cost"] == 0.000045
assert result["kwargs"]["model"] == "gpt-4o"
assert result["kwargs"]["custom_llm_provider"] == "openai"
# Verify cost calculation was called
mock_completion_cost.assert_called_once()
# Verify logging object was updated
assert mock_logging_obj.model_call_details["response_cost"] == 0.000045
assert mock_logging_obj.model_call_details["model"] == "gpt-4o"
assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai"
@patch('litellm.completion_cost')
def test_openai_passthrough_handler_non_chat_completions(self, mock_completion_cost):
"""Test that non-chat-completions routes fall back to base handler"""
# Arrange
mock_httpx_response = self._create_mock_httpx_response()
mock_logging_obj = self._create_mock_logging_obj()
passthrough_payload = self._create_passthrough_logging_payload()
kwargs = {
"passthrough_logging_payload": passthrough_payload,
"model": "gpt-4o",
}
# Act - Use a non-chat-completions route
result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler(
httpx_response=mock_httpx_response,
response_body={"id": "file-123", "object": "file"},
logging_obj=mock_logging_obj,
url_route="https://api.openai.com/v1/files",
result="",
start_time=self.start_time,
end_time=self.end_time,
cache_hit=False,
request_body={"purpose": "fine-tune"},
**kwargs
)
# Assert - Should fall back to base handler for non-chat-completions
assert result is not None
assert "result" in result
assert "kwargs" in result
# Cost calculation may be called by the base handler fallback
# The important thing is that our specific OpenAI handler logic didn't run
@patch('litellm.completion_cost')
@patch('litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload')
def test_openai_passthrough_handler_with_user_tracking(self, mock_get_standard_logging, mock_completion_cost):
"""Test cost tracking with user information"""
# Arrange
mock_completion_cost.return_value = 0.000123
mock_get_standard_logging.return_value = {"test": "logging_payload"}
mock_httpx_response = self._create_mock_httpx_response()
mock_logging_obj = self._create_mock_logging_obj()
# Create payload with user information
passthrough_payload = PassthroughStandardLoggingPayload(
url="https://api.openai.com/v1/chat/completions",
request_body={
"model": "gpt-4o",
"messages": [{"role": "user", "content": "Hello"}],
"user": "test_user_123"
},
request_method="POST",
)
kwargs = {
"passthrough_logging_payload": passthrough_payload,
"model": "gpt-4o",
}
# Act
result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler(
httpx_response=mock_httpx_response,
response_body=self.mock_openai_response,
logging_obj=mock_logging_obj,
url_route="https://api.openai.com/v1/chat/completions",
result="",
start_time=self.start_time,
end_time=self.end_time,
cache_hit=False,
request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}], "user": "test_user_123"},
**kwargs
)
# Assert
assert result is not None
assert "result" in result
assert "kwargs" in result
assert result["kwargs"]["response_cost"] == 0.000123
# Verify user information is included in litellm_params
assert "litellm_params" in result["kwargs"]
assert "proxy_server_request" in result["kwargs"]["litellm_params"]
assert "body" in result["kwargs"]["litellm_params"]["proxy_server_request"]
assert result["kwargs"]["litellm_params"]["proxy_server_request"]["body"]["user"] == "test_user_123"
@patch('litellm.completion_cost')
def test_openai_passthrough_handler_cost_calculation_error(self, mock_completion_cost):
"""Test error handling in cost calculation"""
# Arrange
mock_completion_cost.side_effect = Exception("Cost calculation failed")
mock_httpx_response = self._create_mock_httpx_response()
mock_logging_obj = self._create_mock_logging_obj()
passthrough_payload = self._create_passthrough_logging_payload()
kwargs = {
"passthrough_logging_payload": passthrough_payload,
"model": "gpt-4o",
}
# Act
result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler(
httpx_response=mock_httpx_response,
response_body=self.mock_openai_response,
logging_obj=mock_logging_obj,
url_route="https://api.openai.com/v1/chat/completions",
result="",
start_time=self.start_time,
end_time=self.end_time,
cache_hit=False,
request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]},
**kwargs
)
# Assert - Should fall back to base handler when cost calculation fails
assert result is not None
assert "result" in result
assert "kwargs" in result
def test_build_complete_streaming_response(self):
"""Test the streaming response builder (placeholder implementation)"""
# This is a placeholder method that returns None for now
result = self.handler._build_complete_streaming_response(
all_chunks=["chunk1", "chunk2"],
litellm_logging_obj=self._create_mock_logging_obj(),
model="gpt-4o",
)
assert result is None # Placeholder implementation
@patch('litellm.completion_cost')
@patch('litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload')
def test_different_models_cost_tracking(self, mock_get_standard_logging, mock_completion_cost):
"""Test cost tracking for different OpenAI models"""
# Arrange
mock_get_standard_logging.return_value = {"test": "logging_payload"}
test_cases = [
("gpt-4o", 0.000045),
("gpt-4o-mini", 0.000015),
("gpt-3.5-turbo", 0.000002),
]
for model, expected_cost in test_cases:
mock_completion_cost.return_value = expected_cost
mock_httpx_response = self._create_mock_httpx_response()
mock_httpx_response.json.return_value = {
**self.mock_openai_response,
"model": model
}
mock_logging_obj = self._create_mock_logging_obj()
passthrough_payload = self._create_passthrough_logging_payload()
kwargs = {
"passthrough_logging_payload": passthrough_payload,
"model": model,
}
# Act
result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler(
httpx_response=mock_httpx_response,
response_body={**self.mock_openai_response, "model": model},
logging_obj=mock_logging_obj,
url_route="https://api.openai.com/v1/chat/completions",
result="",
start_time=self.start_time,
end_time=self.end_time,
cache_hit=False,
request_body={"model": model, "messages": [{"role": "user", "content": "Hello"}]},
**kwargs
)
# Assert
assert result is not None
assert "result" in result
assert "kwargs" in result
assert result["kwargs"]["response_cost"] == expected_cost
assert result["kwargs"]["model"] == model
assert result["kwargs"]["custom_llm_provider"] == "openai"
def test_static_methods(self):
"""Test that static methods work correctly"""
# Test static method calls
assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.openai.com/v1/chat/completions") == True
assert OpenAIPassthroughLoggingHandler.get_provider_config("gpt-4o") is not None
class TestOpenAIPassthroughIntegration:
"""Integration tests for OpenAI passthrough cost tracking"""
def setup_method(self):
"""Set up test fixtures"""
self.handler = PassThroughEndpointLogging()
def test_is_openai_route_detection(self):
"""Test OpenAI route detection in the main success handler"""
# Positive cases
assert self.handler.is_openai_route("https://api.openai.com/v1/chat/completions") == True
assert self.handler.is_openai_route("https://openai.azure.com/v1/chat/completions") == True
assert self.handler.is_openai_route("https://api.openai.com/v1/models") == True
# Negative cases
assert self.handler.is_openai_route("http://localhost:4000/openai/v1/chat/completions") == False
assert self.handler.is_openai_route("https://api.anthropic.com/v1/messages") == False
assert self.handler.is_openai_route("https://api.assemblyai.com/v2/transcript") == False
assert self.handler.is_openai_route("") == False
@patch('litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler')
@pytest.mark.asyncio
async def test_success_handler_calls_openai_handler(self, mock_openai_handler):
"""Test that the success handler calls our OpenAI handler for OpenAI routes"""
# Arrange
mock_openai_handler.return_value = {
"result": {"id": "chatcmpl-123"},
"kwargs": {
"response_cost": 0.000045,
"model": "gpt-4o",
"custom_llm_provider": "openai"
}
}
mock_httpx_response = MagicMock(spec=httpx.Response)
mock_httpx_response.text = '{"id": "chatcmpl-123", "choices": [{"message": {"content": "Hello"}}]}'
mock_logging_obj = AsyncMock()
mock_logging_obj.model_call_details = {}
mock_logging_obj.async_success_handler = AsyncMock()
passthrough_payload = PassthroughStandardLoggingPayload(
url="https://api.openai.com/v1/chat/completions",
request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]},
request_method="POST",
)
# Act
result = await self.handler.pass_through_async_success_handler(
httpx_response=mock_httpx_response,
response_body={"id": "chatcmpl-123", "choices": [{"message": {"content": "Hello"}}]},
logging_obj=mock_logging_obj,
url_route="https://api.openai.com/v1/chat/completions",
result="",
start_time=datetime.now(),
end_time=datetime.now(),
cache_hit=False,
request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]},
passthrough_logging_payload=passthrough_payload,
)
# Assert
mock_openai_handler.assert_called_once()
# The success handler returns None on success, which is expected
assert result is None
@pytest.mark.asyncio
async def test_success_handler_falls_back_for_non_openai_routes(self):
"""Test that non-OpenAI routes don't call our handler"""
# Arrange
mock_httpx_response = MagicMock(spec=httpx.Response)
mock_httpx_response.text = '{"status": "success"}'
mock_httpx_response.headers = {"content-type": "application/json"}
mock_logging_obj = MagicMock()
mock_logging_obj.model_call_details = {}
passthrough_payload = PassthroughStandardLoggingPayload(
url="https://api.anthropic.com/v1/messages",
request_body={"model": "claude-3-sonnet", "messages": [{"role": "user", "content": "Hello"}]},
request_method="POST",
)
# Mock the _handle_logging method to capture calls
self.handler._handle_logging = AsyncMock()
# Act
result = await self.handler.pass_through_async_success_handler(
httpx_response=mock_httpx_response,
response_body={"status": "success"},
logging_obj=mock_logging_obj,
url_route="https://api.anthropic.com/v1/messages",
result="",
start_time=datetime.now(),
end_time=datetime.now(),
cache_hit=False,
request_body={"model": "claude-3-sonnet", "messages": [{"role": "user", "content": "Hello"}]},
passthrough_logging_payload=passthrough_payload,
)
# Assert - Should call the base handler, not our OpenAI handler
self.handler._handle_logging.assert_called_once()
if __name__ == "__main__":
pytest.main([__file__])

View file

@ -1245,3 +1245,58 @@ async def test_delete_pass_through_endpoint_empty_list():
# Verify the exception
assert exc_info.value.status_code == 400
assert "no pass-through endpoints setup" in str(exc_info.value.detail).lower()
@pytest.mark.asyncio
async def test_pass_through_with_httpbin_redirect():
"""
Integration test using httpbin.org redirect endpoint to test real redirect handling.
This tests the actual redirect handling capability end-to-end using the full pass_through_request function.
"""
from unittest.mock import MagicMock
from fastapi import Request
from starlette.datastructures import Headers, QueryParams
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
pass_through_request,
)
# Create mock request
mock_request = MagicMock(spec=Request)
mock_request.method = "GET"
mock_request.headers = Headers({})
mock_request.query_params = QueryParams("")
# Mock the body method to return empty bytes for GET request
async def mock_body():
return b""
mock_request.body = mock_body
# Mock user API key dict
mock_user_api_key_dict = MagicMock()
try:
# Test with httpbin.org redirect endpoint
# This will redirect to httpbin.org/get
response = await pass_through_request(
request=mock_request,
target="https://httpbin.org/redirect/1",
custom_headers={},
user_api_key_dict=mock_user_api_key_dict
)
# Should get the final response (200) from /get endpoint, not the redirect (302)
assert response.status_code == 200
# The response should be from the /get endpoint
response_content = response.body.decode('utf-8')
# httpbin.org/get returns JSON with info about the request
assert '"url": "https://httpbin.org/get"' in response_content
print("GOT A Response from HTTPBIN=", response_content)
except Exception as e:
# If httpbin.org is not accessible, skip the test
import pytest
pytest.skip(f"Could not reach httpbin.org for integration test: {e}")

View file

@ -549,6 +549,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"completion",
"embedding",
"image_generation",
"video_generation",
"moderation",
"rerank",
"responses",
@ -636,7 +637,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"type": "array",
"items": {
"type": "string",
"enum": ["text", "image", "audio", "code"],
"enum": ["text", "image", "audio", "code", "video"],
},
},
"supports_native_streaming": {"type": "boolean"},