tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-18 06:26:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into litellm_dev_09_12_2025_p1

Browse Source
This commit is contained in:
Krish Dholakia
2025-09-13 10:10:30 -07:00
committed by GitHub
parent 9ee9745781 ec32a0a7d7
commit 269515e525
56 changed files with 11730 additions and 429 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.circleci/config.yml
+1
View File
@@ -671,6 +671,7 @@ jobs:
pip install mypy
pip install "google-generativeai==0.3.2"
pip install "google-cloud-aiplatform==1.43.0"
pip install "google-genai==1.22.0"
pip install pyarrow
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
cookbook/litellm_proxy_server/batch_api/bedrock/bedrock.py
+25
View File
@@ -0,0 +1,25 @@
from openai import OpenAI
client = OpenAI(
base_url="http://0.0.0.0:4000",
api_key="sk-1234",
)
BEDROCK_BATCH_MODEL = "bedrock/batch-anthropic.claude-3-5-sonnet-20240620-v1:0"
# Upload file
batch_input_file = client.files.create(
file=open("./bedrock_batch_completions.jsonl", "rb"),
purpose="batch",
extra_body={"target_model_names": BEDROCK_BATCH_MODEL}
)
print(batch_input_file)
# Create batch
batch = client.batches.create(
input_file_id=batch_input_file.id,
endpoint="/v1/chat/completions",
completion_window="24h",
metadata={"description": "Test batch job"},
)
print(batch)
cookbook/litellm_proxy_server/batch_api/bedrock/bedrock_batch_completions.jsonl
+128
View File
@@ -0,0 +1,128 @@
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
docs/my-website/docs/batches.md
+2 -1
View File
@@ -7,7 +7,7 @@ Covers Batches, Files
| Feature | Supported | Notes |
|-------|-------|-------|
| Supported Providers | OpenAI, Azure, Vertex | - |
| Supported Providers | OpenAI, Azure, Vertex, Bedrock | - |
| ✨ Cost Tracking | ✅ | LiteLLM Enterprise only |
| Logging | ✅ | Works across all logging integrations |
@@ -178,6 +178,7 @@ print("list_batches_response=", list_batches_response)
### [Azure OpenAI](./providers/azure#azure-batches-api)
### [OpenAI](#quick-start)
### [Vertex AI](./providers/vertex#batch-apis)
### [Bedrock](./providers/bedrock_batches)
## How Cost Tracking for Batches API Works
docs/my-website/docs/providers/bedrock_batches.md
+180
View File
@@ -0,0 +1,180 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Bedrock Batches
Use Amazon Bedrock Batch Inference API through LiteLLM.
| Property | Details |
|----------|---------|
| Description | Amazon Bedrock Batch Inference allows you to run inference on large datasets asynchronously |
| Provider Doc | [AWS Bedrock Batch Inference ↗](https://docs.aws.amazon.com/bedrock/latest/userguide/batch-inference.html) |
## Overview
Use this to:
- Run batch inference on large datasets with Bedrock models
- Control batch model access by key/user/team (same as chat completion models)
- Manage S3 storage for batch input/output files
## (Proxy Admin) Usage
Here's how to give developers access to your Bedrock Batch models.
### 1. Setup config.yaml
- Specify `mode: batch` for each model: Allows developers to know this is a batch model
- Configure S3 bucket and AWS credentials for batch operations
```yaml showLineNumbers title="litellm_config.yaml"
model_list:
- model_name: "bedrock-batch-claude"
litellm_params:
model: bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0
#########################################################
########## batch specific params ########################
s3_bucket_name: litellm-proxy
s3_region_name: us-west-2
s3_access_key_id: os.environ/AWS_ACCESS_KEY_ID
s3_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
aws_batch_role_arn: arn:aws:iam::888602223428:role/service-role/AmazonBedrockExecutionRoleForAgents_BB9HNW6V4CV
model_info:
mode: batch # 👈 SPECIFY MODE AS BATCH, to tell user this is a batch model
```
**Required Parameters:**
| Parameter | Description |
|-----------|-------------|
| `s3_bucket_name` | S3 bucket for batch input/output files |
| `s3_region_name` | AWS region for S3 bucket |
| `s3_access_key_id` | AWS access key for S3 bucket |
| `s3_secret_access_key` | AWS secret key for S3 bucket |
| `aws_batch_role_arn` | IAM role ARN for Bedrock batch operations. Bedrock Batch APIs require an IAM role ARN to be set. |
| `mode: batch` | Indicates to LiteLLM this is a batch model |
### 2. Create Virtual Key
```bash showLineNumbers title="create_virtual_key.sh"
curl -L -X POST 'https://{PROXY_BASE_URL}/key/generate' \
-H 'Authorization: Bearer ${PROXY_API_KEY}' \
-H 'Content-Type: application/json' \
-d '{"models": ["bedrock-batch-claude"]}'
```
You can now use the virtual key to access the batch models (See Developer flow).
## (Developer) Usage
Here's how to create a LiteLLM managed file and execute Bedrock Batch CRUD operations with the file.
### 1. Create request.jsonl
- Check models available via `/model_group/info`
- See all models with `mode: batch`
- Set `model` in .jsonl to the model from `/model_group/info`
```json showLineNumbers title="bedrock_batch_completions.jsonl"
{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "bedrock-batch-claude", "messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Hello world!"}], "max_tokens": 1000}}
{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "bedrock-batch-claude", "messages": [{"role": "system", "content": "You are an unhelpful assistant."}, {"role": "user", "content": "Hello world!"}], "max_tokens": 1000}}
```
Expectation:
- LiteLLM translates this to the bedrock deployment specific value (e.g. `bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0`)
### 2. Upload File
Specify `target_model_names: "<model-name>"` to enable LiteLLM managed files and request validation.
model-name should be the same as the model-name in the request.jsonl
<Tabs>
<TabItem value="python" label="Python">
```python showLineNumbers title="bedrock_batch.py"
from openai import OpenAI
client = OpenAI(
base_url="http://0.0.0.0:4000",
api_key="sk-1234",
)
# Upload file
batch_input_file = client.files.create(
file=open("./bedrock_batch_completions.jsonl", "rb"), # {"model": "bedrock-batch-claude"} <-> {"model": "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0"}
purpose="batch",
extra_body={"target_model_names": "bedrock-batch-claude"}
)
print(batch_input_file)
```
</TabItem>
<TabItem value="curl" label="Curl">
```bash showLineNumbers title="Upload File"
curl http://localhost:4000/v1/files \
-H "Authorization: Bearer sk-1234" \
-F purpose="batch" \
-F file="@bedrock_batch_completions.jsonl" \
-F extra_body='{"target_model_names": "bedrock-batch-claude"}'
```
</TabItem>
</Tabs>
**Where is the file written?**:
The file is written to S3 bucket specified in your config and prepared for Bedrock batch inference.
### 3. Create the batch
<Tabs>
<TabItem value="python" label="Python">
```python showLineNumbers title="bedrock_batch.py"
...
# Create batch
batch = client.batches.create(
input_file_id=batch_input_file.id,
endpoint="/v1/chat/completions",
completion_window="24h",
metadata={"description": "Test batch job"},
)
print(batch)
```
</TabItem>
<TabItem value="curl" label="Curl">
```bash showLineNumbers title="Create Batch Request"
curl http://localhost:4000/v1/batches \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"input_file_id": "file-abc123",
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"metadata": {"description": "Test batch job"}
}'
```
</TabItem>
</Tabs>
## FAQ
### Where are my files written?
When a `target_model_names` is specified, the file is written to the S3 bucket configured in your Bedrock batch model configuration.
### What models are supported?
LiteLLM only supports Bedrock Anthropic Models for Batch API. If you want other bedrock models file an issue [here](https://github.com/BerriAI/litellm/issues/new/choose).
## Further Reading
- [AWS Bedrock Batch Inference Documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/batch-inference.html)
- [LiteLLM Managed Batches](../proxy/managed_batches)
- [LiteLLM Authentication to Bedrock](https://docs.litellm.ai/docs/providers/bedrock#boto3---authentication)
docs/my-website/docs/proxy/admin_ui_sso.md
+4
View File
@@ -4,6 +4,10 @@ import TabItem from '@theme/TabItem';
# ✨ SSO for Admin UI
:::info
From v1.76.0, SSO is now Free for up to 5 users.
:::
:::info
✨ SSO is on LiteLLM Enterprise
docs/my-website/docs/proxy/debugging.md
+53 -15
View File
@@ -11,13 +11,13 @@ The proxy also supports json logs. [See here](#json-logs)
**via cli**
```bash
```bash showLineNumbers
$ litellm --debug
```
**via env**
```python
```python showLineNumbers
os.environ["LITELLM_LOG"] = "INFO"
```
@@ -25,25 +25,25 @@ os.environ["LITELLM_LOG"] = "INFO"
**via cli**
```bash
```bash showLineNumbers
$ litellm --detailed_debug
```
**via env**
```python
```python showLineNumbers
os.environ["LITELLM_LOG"] = "DEBUG"
```
### Debug Logs
Run the proxy with `--detailed_debug` to view detailed debug logs
```shell
```shell showLineNumbers
litellm --config /path/to/config.yaml --detailed_debug
```
When making requests you should see the POST request sent by LiteLLM to the LLM on the Terminal output
```shell
```shell showLineNumbers
POST Request Sent from LiteLLM:
curl -X POST \
https://api.openai.com/v1/chat/completions \
@@ -51,25 +51,63 @@ https://api.openai.com/v1/chat/completions \
-d '{"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "this is a test request, write a short poem"}]}'
```
## Debug single request
Pass in `litellm_request_debug=True` in the request body
```bash showLineNumbers
curl -L -X POST 'http://0.0.0.0:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{
"model":"fake-openai-endpoint",
"messages": [{"role": "user","content": "How many r in the word strawberry?"}],
"litellm_request_debug": true
}'
```
This will emit the raw request sent by LiteLLM to the API Provider and raw response received from the API Provider for **just** this request in the logs.
```bash showLineNumbers
INFO: Uvicorn running on http://0.0.0.0:4000 (Press CTRL+C to quit)
20:14:06 - LiteLLM:WARNING: litellm_logging.py:938 -
POST Request Sent from LiteLLM:
curl -X POST \
https://exampleopenaiendpoint-production.up.railway.app/chat/completions \
-H 'Authorization: Be****ey' -H 'Content-Type: application/json' \
-d '{'model': 'fake', 'messages': [{'role': 'user', 'content': 'How many r in the word strawberry?'}], 'stream': False}'
20:14:06 - LiteLLM:WARNING: litellm_logging.py:1015 - RAW RESPONSE:
{"id":"chatcmpl-817fc08f0d6c451485d571dab39b26a1","object":"chat.completion","created":1677652288,"model":"gpt-3.5-turbo-0301","system_fingerprint":"fp_44709d6fcb","choices":[{"index":0,"message":{"role":"assistant","content":"\n\nHello there, how may I assist you today?"},"logprobs":null,"finish_reason":"stop"}],"usage":{"prompt_tokens":9,"completion_tokens":12,"total_tokens":21}}
INFO: 127.0.0.1:56155 - "POST /chat/completions HTTP/1.1" 200 OK
```
## JSON LOGS
Set `JSON_LOGS="True"` in your env:
```bash
```bash showLineNumbers
export JSON_LOGS="True"
```
**OR**
Set `json_logs: true` in your yaml:
```yaml
```yaml showLineNumbers
litellm_settings:
json_logs: true
```
Start proxy
```bash
```bash showLineNumbers
$ litellm
```
@@ -80,7 +118,7 @@ The proxy will now all logs in json format.
Turn off fastapi's default 'INFO' logs
1. Turn on 'json logs'
```yaml
```yaml showLineNumbers
litellm_settings:
json_logs: true
```
@@ -89,20 +127,20 @@ litellm_settings:
Only get logs if an error occurs.
```bash
```bash showLineNumbers
LITELLM_LOG="ERROR"
```
3. Start proxy
```bash
```bash showLineNumbers
$ litellm
```
Expected Output:
```bash
```bash showLineNumbers
# no info statements
```
@@ -119,14 +157,14 @@ This can be caused due to all your models hitting rate limit errors, causing the
How to control this?
- Adjust the cooldown time
```yaml
```yaml showLineNumbers
router_settings:
cooldown_time: 0 # 👈 KEY CHANGE
```
- Disable Cooldowns [NOT RECOMMENDED]
```yaml
```yaml showLineNumbers
router_settings:
disable_cooldowns: True
```
docs/my-website/docs/proxy/forward_client_headers.md
+212
View File
@@ -0,0 +1,212 @@
# Forward Client Headers to LLM API
Control which model groups can forward client headers to the underlying LLM provider APIs.
## Overview
By default, LiteLLM does not forward client headers to LLM provider APIs for security reasons. However, you can selectively enable header forwarding for specific model groups using the `forward_client_headers_to_llm_api` setting.
## Configuration
## Enable Globally
```yaml
general_settings:
forward_client_headers_to_llm_api: true
```
## Enable for a Model Group
Add the `forward_client_headers_to_llm_api` setting under `model_group_settings` in your configuration:
```yaml
model_list:
- model_name: gpt-4o-mini
litellm_params:
model: openai/gpt-4o-mini
api_key: "your-api-key"
- model_name: "wildcard-models/*"
litellm_params:
model: "openai/*"
api_key: "your-api-key"
litellm_settings:
model_group_settings:
forward_client_headers_to_llm_api:
- gpt-4o-mini
- wildcard-models/*
```
## Supported Model Patterns
The configuration supports various model matching patterns:
### 1. Exact Model Names
```yaml
forward_client_headers_to_llm_api:
- gpt-4o-mini
- claude-3-sonnet
```
### 2. Wildcard Patterns
```yaml
forward_client_headers_to_llm_api:
- "openai/*" # All OpenAI models
- "anthropic/*" # All Anthropic models
- "wildcard-group/*" # All models in wildcard-group
```
### 3. Team Model Aliases
If your team has model aliases configured, the forwarding will work with both the original model name and the alias.
## Forwarded Headers
When enabled for a model group, LiteLLM forwards the following types of headers:
### Custom Headers (x- prefix)
- Any header starting with `x-` (except `x-stainless-*` which can cause OpenAI SDK issues)
- Examples: `x-custom-header`, `x-request-id`, `x-trace-id`
### Provider-Specific Headers
- **Anthropic**: `anthropic-beta` headers
- **OpenAI**: `openai-organization` (when enabled via `forward_openai_org_id: true`)
### User Information Headers (Optional)
When `add_user_information_to_llm_headers` is enabled, LiteLLM adds:
- `x-litellm-user-id`
- `x-litellm-org-id`
- Other user metadata as `x-litellm-*` headers
## Security Considerations
⚠️ **Important Security Notes:**
1. **Sensitive Data**: Only enable header forwarding for trusted model groups, as headers may contain sensitive information
2. **API Keys**: Never include API keys or secrets in forwarded headers
3. **PII**: Be cautious about forwarding headers that might contain personally identifiable information
4. **Provider Limits**: Some providers have restrictions on custom headers
## Example Use Cases
### 1. Request Tracing
Forward tracing headers to track requests across your system:
```bash
curl -X POST "https://your-proxy.com/v1/chat/completions" \
-H "Authorization: Bearer your-key" \
-H "x-trace-id: abc123" \
-H "x-request-source: mobile-app" \
-d '{
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "Hello"}]
}'
```
### 2. Custom Metadata
Pass custom metadata to your LLM provider:
```bash
curl -X POST "https://your-proxy.com/v1/chat/completions" \
-H "Authorization: Bearer your-key" \
-H "x-customer-id: customer-123" \
-H "x-environment: production" \
-d '{
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "Hello"}]
}'
```
### 3. Anthropic Beta Features
Enable beta features for Anthropic models:
```bash
curl -X POST "https://your-proxy.com/v1/chat/completions" \
-H "Authorization: Bearer your-key" \
-H "anthropic-beta: tools-2024-04-04" \
-d '{
"model": "claude-3-sonnet",
"messages": [{"role": "user", "content": "Hello"}]
}'
```
## Complete Configuration Example
```yaml
model_list:
# Fixed model with header forwarding
- model_name: byok-fixed-gpt-4o-mini
litellm_params:
model: openai/gpt-4o-mini
api_base: "https://your-openai-endpoint.com"
api_key: "your-api-key"
# Wildcard model group with header forwarding
- model_name: "byok-wildcard/*"
litellm_params:
model: "openai/*"
api_base: "https://your-openai-endpoint.com"
api_key: "your-api-key"
# Standard model without header forwarding
- model_name: standard-gpt-4
litellm_params:
model: openai/gpt-4
api_key: "your-api-key"
litellm_settings:
# Enable user info headers globally (optional)
add_user_information_to_llm_headers: true
model_group_settings:
forward_client_headers_to_llm_api:
- byok-fixed-gpt-4o-mini
- byok-wildcard/*
# Note: standard-gpt-4 is NOT included, so no headers forwarded
general_settings:
# Enable OpenAI organization header forwarding (optional)
forward_openai_org_id: true
```
## Testing Header Forwarding
To test if headers are being forwarded:
1. **Enable Debug Logging**: Set `set_verbose: true` in your config
2. **Check Provider Logs**: Monitor your LLM provider's request logs
3. **Use Webhook Sites**: For testing, you can use webhook.site URLs as api_base to see forwarded headers
## Troubleshooting
### Headers Not Being Forwarded
1. **Check Model Name**: Ensure the model name in your request matches the configuration
2. **Verify Pattern Matching**: Wildcard patterns must match exactly
3. **Review Logs**: Enable verbose logging to see header processing
### Provider Errors
1. **Invalid Headers**: Some providers reject unknown headers
2. **Header Limits**: Providers may have limits on header count/size
3. **Authentication**: Ensure forwarded headers don't conflict with authentication
## Related Features
- [Request Headers](./request_headers.md) - Complete list of supported request headers
- [Response Headers](./response_headers.md) - Headers returned by LiteLLM
- [Team Model Aliases](./team_model_add.md) - Configure model aliases for teams
- [Model Access Control](./model_access.md) - Control which users can access which models
## API Reference
The header forwarding is controlled by the `ModelGroupSettings` configuration:
```python
class ModelGroupSettings(BaseModel):
forward_client_headers_to_llm_api: Optional[List[str]] = None
```
Where each string in the list can be:
- An exact model name (e.g., `"gpt-4o-mini"`)
- A wildcard pattern (e.g., `"openai/*"`)
- A model group name (e.g., `"my-model-group/*"`)
docs/my-website/docs/proxy/guardrails/noma_security.md
+17
View File
@@ -135,6 +135,7 @@ guardrails:
# application_id: "my-app"
# monitor_mode: false
# block_failures: true
# anonymize_input: false
```
### Required Parameters
@@ -147,6 +148,7 @@ guardrails:
- **`application_id`**: Your application identifier (defaults to `"litellm"`)
- **`monitor_mode`**: If `true`, logs violations without blocking (defaults to `false`)
- **`block_failures`**: If `true`, blocks requests when guardrail API failures occur (defaults to `true`)
- **`anonymize_input`**: If `true`, replaces sensitive content with anonymized version (defaults to `false`)
## Environment Variables
@@ -158,6 +160,7 @@ export NOMA_API_BASE="https://api.noma.security/" # Optional
export NOMA_APPLICATION_ID="my-app" # Optional
export NOMA_MONITOR_MODE="false" # Optional
export NOMA_BLOCK_FAILURES="true" # Optional
export NOMA_ANONYMIZE_INPUT="false" # Optional
```
## Advanced Configuration
@@ -190,6 +193,20 @@ guardrails:
block_failures: false # Allow requests to proceed if guardrail API fails
```
### Content Anonymization
Enable anonymization to replace sensitive content instead of blocking:
```yaml
guardrails:
- guardrail_name: "noma-anonymize"
litellm_params:
guardrail: noma
mode: "pre_call"
api_key: os.environ/NOMA_API_KEY
anonymize_input: true # Replace sensitive data with anonymized version
```
### Multiple Guardrails
Apply different configurations for input and output:
docs/my-website/docs/proxy/request_headers.md
+9 -1
View File
@@ -2,6 +2,10 @@
Special headers that are supported by LiteLLM.
## Header Forwarding
By default, LiteLLM does not forward client headers to LLM provider APIs. However, you can selectively enable header forwarding for specific model groups. [Learn more about configuring header forwarding](./forward_client_headers.md).
## LiteLLM Headers
`x-litellm-timeout` Optional[float]: The timeout for the request in seconds.
@@ -21,11 +25,15 @@ Special headers that are supported by LiteLLM.
`anthropic-version` Optional[str]: The version of the Anthropic API to use.
`anthropic-beta` Optional[str]: The beta version of the Anthropic API to use.
- For `/v1/messages` endpoint, this will always be forward the header to the underlying model.
- For `/chat/completions` endpoint, this will only be forwarded if `forward_client_headers_to_llm_api` is true.
- For `/chat/completions` endpoint, this will only be forwarded if the model is configured in `forward_client_headers_to_llm_api`. [Learn more](./forward_client_headers.md)
## OpenAI Headers
`openai-organization` Optional[str]: The organization to use for the OpenAI API. (currently needs to be enabled via `general_settings::forward_openai_org_id: true`)
## Custom Headers
Custom headers starting with `x-` can be forwarded to LLM provider APIs when the model is configured in `forward_client_headers_to_llm_api`. [Learn more about header forwarding configuration](./forward_client_headers.md).
docs/my-website/sidebars.js
+2
View File
@@ -141,6 +141,7 @@ const sidebars = {
"proxy/clientside_auth",
"proxy/request_headers",
"proxy/response_headers",
"proxy/forward_client_headers",
"proxy/model_discovery",
],
},
@@ -410,6 +411,7 @@ const sidebars = {
items: [
"providers/bedrock",
"providers/bedrock_agents",
"providers/bedrock_batches",
"providers/bedrock_vector_store",
]
},
litellm/batches/main.py
+61 -18
View File
@@ -19,6 +19,7 @@ from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast
import httpx
import litellm
from litellm._logging import verbose_logger
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
@@ -38,6 +39,7 @@ from litellm.utils import (
ProviderConfigManager,
client,
get_litellm_params,
get_llm_provider,
supports_httpx_timeout,
)
@@ -49,6 +51,45 @@ base_llm_http_handler = BaseLLMHTTPHandler()
#################################################
def _resolve_timeout(
optional_params: GenericLiteLLMParams,
kwargs: Dict[str, Any],
custom_llm_provider: str,
default_timeout: float = 600.0,
) -> float:
"""
Resolve timeout value from various sources and handle httpx.Timeout objects.
Args:
optional_params: GenericLiteLLMParams object containing timeout
kwargs: Additional kwargs that may contain request_timeout
custom_llm_provider: Provider name for httpx timeout support check
default_timeout: Default timeout value to use
Returns:
Resolved timeout as float
"""
timeout = optional_params.timeout or kwargs.get("request_timeout", default_timeout) or default_timeout
# Handle httpx.Timeout objects
if isinstance(timeout, httpx.Timeout):
if supports_httpx_timeout(custom_llm_provider) is False:
# Extract read timeout for providers that don't support httpx.Timeout
read_timeout = timeout.read or default_timeout
return float(read_timeout)
else:
# For providers that support httpx.Timeout, we still need to return a float
# This case might need to be handled differently based on the actual use case
return float(timeout.read or default_timeout)
# Handle None case
if timeout is None:
return float(default_timeout)
# Handle numeric values (int, float, string representations)
return float(timeout)
@client
async def acreate_batch(
completion_window: Literal["24h"],
@@ -118,13 +159,23 @@ def create_batch(
litellm_call_id = kwargs.get("litellm_call_id", None)
proxy_server_request = kwargs.get("proxy_server_request", None)
model_info = kwargs.get("model_info", None)
model: Optional[str] = kwargs.get("model", None)
try:
if model is not None:
model, _, _, _ = get_llm_provider(
model=model,
custom_llm_provider=None,
)
except Exception as e:
verbose_logger.exception(f"litellm.batches.main.py::create_batch() - Error inferring custom_llm_provider - {str(e)}")
_is_async = kwargs.pop("acreate_batch", False) is True
litellm_params = dict(GenericLiteLLMParams(**kwargs))
litellm_logging_obj: LiteLLMLoggingObj = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None))
### TIMEOUT LOGIC ###
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
timeout = _resolve_timeout(optional_params, kwargs, custom_llm_provider)
litellm_logging_obj.update_environment_variables(
model=None,
model=model,
user=None,
optional_params=optional_params.model_dump(),
litellm_params={
@@ -138,18 +189,6 @@ def create_batch(
},
custom_llm_provider=custom_llm_provider,
)
if (
timeout is not None
and isinstance(timeout, httpx.Timeout)
and supports_httpx_timeout(custom_llm_provider) is False
):
read_timeout = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
elif timeout is None:
timeout = 600.0
_create_batch_request = CreateBatchRequest(
@@ -160,10 +199,13 @@ 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 model is not None:
provider_config = ProviderConfigManager.get_provider_batches_config(
model=model,
provider=LlmProviders(custom_llm_provider),
)
else:
provider_config = None
if provider_config is not None:
response = base_llm_http_handler.create_batch(
provider_config=provider_config,
@@ -179,6 +221,7 @@ def create_batch(
and isinstance(client, (HTTPHandler, AsyncHTTPHandler))
else None,
timeout=timeout,
model=model,
)
return response
api_base: Optional[str] = None
litellm/files/utils.py
+27
View File
@@ -0,0 +1,27 @@
from typing import Optional
from litellm.types.llms.openai import CreateFileRequest
from litellm.types.utils import ExtractedFileData
class FilesAPIUtils:
"""
Utils for files API interface on litellm
"""
@staticmethod
def is_batch_jsonl_file(create_file_data: CreateFileRequest, extracted_file_data: ExtractedFileData) -> bool:
"""
Check if the file is a batch jsonl file
"""
return (
create_file_data.get("purpose") == "batch"
and FilesAPIUtils.valid_content_type(extracted_file_data.get("content_type"))
and extracted_file_data.get("content") is not None
)
@staticmethod
def valid_content_type(content_type: Optional[str]) -> bool:
"""
Check if the content type is valid
"""
return content_type in set(["application/jsonl", "application/octet-stream"])
litellm/litellm_core_utils/get_litellm_params.py
+2
View File
@@ -62,6 +62,7 @@ def get_litellm_params(
use_litellm_proxy: Optional[bool] = None,
api_version: Optional[str] = None,
max_retries: Optional[int] = None,
litellm_request_debug: Optional[bool] = None,
**kwargs,
) -> dict:
litellm_params = {
@@ -118,5 +119,6 @@ def get_litellm_params(
"vertex_credentials": kwargs.get("vertex_credentials"),
"vertex_project": kwargs.get("vertex_project"),
"use_litellm_proxy": use_litellm_proxy,
"litellm_request_debug": litellm_request_debug,
}
return litellm_params
litellm/litellm_core_utils/litellm_logging.py
+41 -15
View File
@@ -245,6 +245,7 @@ class Logging(LiteLLMLoggingBaseClass):
global supabaseClient, promptLayerLogger, weightsBiasesLogger, logfireLogger, capture_exception, add_breadcrumb, lunaryLogger, logfireLogger, prometheusLogger, slack_app
custom_pricing: bool = False
stream_options = None
litellm_request_debug: bool = False
def __init__(
self,
@@ -470,6 +471,7 @@ class Logging(LiteLLMLoggingBaseClass):
**self.litellm_params,
**scrub_sensitive_keys_in_metadata(litellm_params),
}
self.litellm_request_debug = litellm_params.get("litellm_request_debug", False)
self.logger_fn = litellm_params.get("logger_fn", None)
verbose_logger.debug(f"self.optional_params: {self.optional_params}")
@@ -907,13 +909,19 @@ class Logging(LiteLLMLoggingBaseClass):
Prints the RAW curl command sent from LiteLLM
"""
if _is_debugging_on():
if _is_debugging_on() or self.litellm_request_debug:
if json_logs:
masked_headers = self._get_masked_headers(headers)
verbose_logger.debug(
"POST Request Sent from LiteLLM",
extra={"api_base": {api_base}, **masked_headers},
)
if self.litellm_request_debug:
verbose_logger.warning( # .warning ensures this shows up in all environments
"POST Request Sent from LiteLLM",
extra={"api_base": {api_base}, **masked_headers},
)
else:
verbose_logger.debug(
"POST Request Sent from LiteLLM",
extra={"api_base": {api_base}, **masked_headers},
)
else:
headers = additional_args.get("headers", {})
if headers is None:
@@ -926,7 +934,12 @@ class Logging(LiteLLMLoggingBaseClass):
additional_args=additional_args,
data=data,
)
verbose_logger.debug(f"\033[92m{curl_command}\033[0m\n")
if self.litellm_request_debug:
verbose_logger.warning(
f"\033[92m{curl_command}\033[0m\n"
) # .warning ensures this shows up in all environments
else:
verbose_logger.debug(f"\033[92m{curl_command}\033[0m\n")
def _get_request_body(self, data: dict) -> str:
return str(data)
@@ -983,8 +996,14 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["additional_args"] = additional_args
self.model_call_details["log_event_type"] = "post_api_call"
if self.litellm_request_debug:
attr = "warning"
else:
attr = "debug"
if json_logs:
verbose_logger.debug(
callattr = getattr(verbose_logger, attr)
callattr(
"RAW RESPONSE:\n{}\n\n".format(
self.model_call_details.get(
"original_response", self.model_call_details
@@ -992,7 +1011,8 @@ class Logging(LiteLLMLoggingBaseClass):
),
)
else:
print_verbose(
callattr = getattr(verbose_logger, attr)
callattr(
"RAW RESPONSE:\n{}\n\n".format(
self.model_call_details.get(
"original_response", self.model_call_details
@@ -1714,12 +1734,16 @@ class Logging(LiteLLMLoggingBaseClass):
response_obj=result,
start_time=start_time,
end_time=end_time,
litellm_call_id=current_call_id
if (
current_call_id := litellm_params.get("litellm_call_id")
)
is not None
else str(uuid.uuid4()),
litellm_call_id=(
current_call_id
if (
current_call_id := litellm_params.get(
"litellm_call_id"
)
)
is not None
else str(uuid.uuid4())
),
print_verbose=print_verbose,
)
if callback == "wandb" and weightsBiasesLogger is not None:
@@ -3367,6 +3391,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
return galileo_logger # type: ignore
elif logging_integration == "cloudzero":
from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
for callback in _in_memory_loggers:
if isinstance(callback, CloudZeroLogger):
return callback # type: ignore
@@ -3594,6 +3619,7 @@ def get_custom_logger_compatible_class( # noqa: PLR0915
return callback
elif logging_integration == "cloudzero":
from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
for callback in _in_memory_loggers:
if isinstance(callback, CloudZeroLogger):
return callback
@@ -4504,7 +4530,7 @@ def get_standard_logging_object_payload(
def emit_standard_logging_payload(payload: StandardLoggingPayload):
if os.getenv("LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD"):
print(json.dumps(payload, indent=4)) # noqa
print(json.dumps(payload, indent=4)) # noqa
def get_standard_logging_metadata(
litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py
+70 -63
View File
@@ -28,10 +28,6 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
TextBlock,
)
def __init__(self, completion_stream: Any, model: str):
super().__init__(completion_stream)
self.model = model
sent_first_chunk: bool = False
sent_content_block_start: bool = False
sent_content_block_finish: bool = False
@@ -39,6 +35,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
sent_last_message: bool = False
holding_chunk: Optional[Any] = None
holding_stop_reason_chunk: Optional[Any] = None
queued_usage_chunk: bool = False
current_content_block_index: int = 0
current_content_block_start: ContentBlockContentBlockDict = TextBlock(
type="text",
@@ -47,6 +44,10 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
pending_new_content_block: bool = False
chunk_queue: deque = deque() # Queue for buffering multiple chunks
def __init__(self, completion_stream: Any, model: str):
super().__init__(completion_stream)
self.model = model
def __next__(self):
from .transformation import LiteLLMAnthropicMessagesAdapter
@@ -217,77 +218,83 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
# Queue the merged chunk and reset
self.chunk_queue.append(merged_chunk)
self.queued_usage_chunk = True
self.holding_stop_reason_chunk = None
return self.chunk_queue.popleft()
# Check if this processed chunk has a stop_reason - hold it for next chunk
if should_start_new_block and not self.sent_content_block_finish:
# Queue the sequence: content_block_stop -> content_block_start -> current_chunk
if not self.queued_usage_chunk:
if should_start_new_block and not self.sent_content_block_finish:
# Queue the sequence: content_block_stop -> content_block_start -> current_chunk
# 1. Stop current content block
self.chunk_queue.append(
{
"type": "content_block_stop",
"index": max(self.current_content_block_index - 1, 0),
}
)
# 1. Stop current content block
self.chunk_queue.append(
{
"type": "content_block_stop",
"index": max(self.current_content_block_index - 1, 0),
}
)
# 2. Start new content block
self.chunk_queue.append(
{
"type": "content_block_start",
"index": self.current_content_block_index,
"content_block": self.current_content_block_start,
}
)
# 2. Start new content block
self.chunk_queue.append(
{
"type": "content_block_start",
"index": self.current_content_block_index,
"content_block": self.current_content_block_start,
}
)
# 3. Queue the current chunk (don't lose it!)
self.chunk_queue.append(processed_chunk)
# Reset state for new block
self.sent_content_block_finish = False
# Return the first queued item
return self.chunk_queue.popleft()
if (
processed_chunk["type"] == "message_delta"
and self.sent_content_block_finish is False
):
# Queue both the content_block_stop and the holding chunk
self.chunk_queue.append(
{
"type": "content_block_stop",
"index": self.current_content_block_index,
}
)
self.sent_content_block_finish = True
if processed_chunk.get("delta", {}).get("stop_reason") is not None:
self.holding_stop_reason_chunk = processed_chunk
else:
# 3. Queue the current chunk (don't lose it!)
self.chunk_queue.append(processed_chunk)
return self.chunk_queue.popleft()
elif self.holding_chunk is not None:
# Queue both chunks
self.chunk_queue.append(self.holding_chunk)
self.chunk_queue.append(processed_chunk)
self.holding_chunk = None
return self.chunk_queue.popleft()
else:
# Queue the current chunk
self.chunk_queue.append(processed_chunk)
return self.chunk_queue.popleft()
# Reset state for new block
self.sent_content_block_finish = False
# Return the first queued item
return self.chunk_queue.popleft()
if (
processed_chunk["type"] == "message_delta"
and self.sent_content_block_finish is False
):
# Queue both the content_block_stop and the holding chunk
self.chunk_queue.append(
{
"type": "content_block_stop",
"index": self.current_content_block_index,
}
)
self.sent_content_block_finish = True
if (
processed_chunk.get("delta", {}).get("stop_reason")
is not None
):
self.holding_stop_reason_chunk = processed_chunk
else:
self.chunk_queue.append(processed_chunk)
return self.chunk_queue.popleft()
elif self.holding_chunk is not None:
# Queue both chunks
self.chunk_queue.append(self.holding_chunk)
self.chunk_queue.append(processed_chunk)
self.holding_chunk = None
return self.chunk_queue.popleft()
else:
# Queue the current chunk
self.chunk_queue.append(processed_chunk)
return self.chunk_queue.popleft()
# Handle any remaining held chunks after stream ends
if self.holding_stop_reason_chunk is not None:
self.chunk_queue.append(self.holding_stop_reason_chunk)
self.holding_stop_reason_chunk = None
if not self.queued_usage_chunk:
if self.holding_stop_reason_chunk is not None:
self.chunk_queue.append(self.holding_stop_reason_chunk)
self.holding_stop_reason_chunk = None
if self.holding_chunk is not None:
self.chunk_queue.append(self.holding_chunk)
self.holding_chunk = None
if self.holding_chunk is not None:
self.chunk_queue.append(self.holding_chunk)
self.holding_chunk = None
if not self.sent_last_message:
self.sent_last_message = True
litellm/llms/bedrock/batches/transformation.py
+5 -7
View File
@@ -124,15 +124,13 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
"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.")
if not model:
raise ValueError("Could not determine Bedrock model ID. Please pass `model` in your request body.")
# 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")
job_name = self.common_utils.generate_unique_job_name(model, prefix="litellm")
output_key = f"litellm-batch-outputs/{job_name}/"
# Build input data config
@@ -151,7 +149,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
# Create Bedrock batch request with proper typing
bedrock_request: BedrockCreateBatchRequest = {
"modelId": bedrock_model_id,
"modelId": model,
"jobName": job_name,
"inputDataConfig": input_data_config,
"outputDataConfig": output_data_config,
litellm/llms/bedrock/files/transformation.py
+92 -37
View File
@@ -6,6 +6,8 @@ from typing import Any, Dict, List, Optional, Tuple, Union
from httpx import Headers, Response
from litellm._logging import verbose_logger
from litellm.files.utils import FilesAPIUtils
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 (
@@ -21,6 +23,7 @@ from litellm.types.llms.openai import (
PathLike,
)
from litellm.types.utils import ExtractedFileData, LlmProviders
from litellm.utils import get_llm_provider
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockError
@@ -111,6 +114,10 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
# Remove bedrock/ prefix if present
if _model.startswith("bedrock/"):
_model = _model[8:]
# Replace colons with hyphens for Bedrock S3 URI compliance
_model = _model.replace(":", "-")
object_name = f"litellm-bedrock-files-{_model}-{uuid.uuid4()}.jsonl"
return object_name
@@ -191,24 +198,6 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
) -> 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,
@@ -218,11 +207,12 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
"""
Transform OpenAI request body to Bedrock-compatible modelInput parameters using existing transformation logic
"""
from litellm.types.utils import LlmProviders
_model = openai_request_body.get("model", "")
messages = openai_request_body.get("messages", [])
# Use existing Anthropic transformation logic for Anthropic models
if provider == "anthropic":
if provider == LlmProviders.ANTHROPIC:
from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
AmazonAnthropicClaudeConfig,
)
@@ -231,16 +221,22 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
# 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"]}
mapped_params = anthropic_config.map_openai_params(
non_default_params={},
optional_params=optional_params,
model=_model,
drop_params=False
)
# Transform using existing Anthropic logic
bedrock_params = anthropic_config.transform_request(
model=_model,
messages=messages,
optional_params=optional_params,
optional_params=mapped_params,
litellm_params={},
headers={}
)
return bedrock_params
else:
# For other providers, use basic mapping
@@ -278,9 +274,17 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
# Extract the request body from OpenAI format
openai_body = _openai_jsonl_content.get("body", {})
model = openai_body.get("model", "")
try:
model, _, _, _ = get_llm_provider(
model=model,
custom_llm_provider=None,
)
except Exception as e:
verbose_logger.exception(f"litellm.llms.bedrock.files.transformation.py::_transform_openai_jsonl_content_to_bedrock_jsonl_content() - Error inferring custom_llm_provider - {str(e)}")
# Determine provider from model name
provider = self._get_bedrock_provider_from_model(model)
provider = self.get_bedrock_invoke_provider(model)
# Transform to Bedrock modelInput format
model_input = self._map_openai_to_bedrock_params(
@@ -315,11 +319,13 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
extracted_file_data = extract_file_data(file_data)
extracted_file_data_content = extracted_file_data.get("content")
if extracted_file_data_content is None:
raise ValueError("file content is required")
# 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
if FilesAPIUtils.is_batch_jsonl_file(
create_file_data=create_file_data,
extracted_file_data=extracted_file_data,
):
## Transform JSONL content to Bedrock format
original_file_content = self._get_content_from_openai_file(
@@ -357,6 +363,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
api_base=api_base,
optional_params=optional_params,
)
litellm_params["upload_url"] = api_base
# Return a dict that tells the HTTP handler exactly what to do
return {
@@ -440,6 +448,56 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
return dict(aws_request.headers), signed_body
def _convert_https_url_to_s3_uri(self, https_url: str) -> tuple[str, str]:
"""
Convert HTTPS S3 URL to s3:// URI format.
Args:
https_url: HTTPS S3 URL (e.g., "https://s3.us-west-2.amazonaws.com/bucket/key")
Returns:
Tuple of (s3_uri, filename)
Example:
Input: "https://s3.us-west-2.amazonaws.com/litellm-proxy/file.jsonl"
Output: ("s3://litellm-proxy/file.jsonl", "file.jsonl")
"""
import re
# Match HTTPS S3 URL patterns
# Pattern 1: https://s3.region.amazonaws.com/bucket/key
# Pattern 2: https://bucket.s3.region.amazonaws.com/key
pattern1 = r"https://s3\.([^.]+)\.amazonaws\.com/([^/]+)/(.+)"
pattern2 = r"https://([^.]+)\.s3\.([^.]+)\.amazonaws\.com/(.+)"
match1 = re.match(pattern1, https_url)
match2 = re.match(pattern2, https_url)
if match1:
# Pattern: https://s3.region.amazonaws.com/bucket/key
region, bucket, key = match1.groups()
s3_uri = f"s3://{bucket}/{key}"
elif match2:
# Pattern: https://bucket.s3.region.amazonaws.com/key
bucket, region, key = match2.groups()
s3_uri = f"s3://{bucket}/{key}"
else:
# Fallback: try to extract bucket and key from URL path
from urllib.parse import urlparse
parsed = urlparse(https_url)
path_parts = parsed.path.lstrip('/').split('/', 1)
if len(path_parts) >= 2:
bucket, key = path_parts[0], path_parts[1]
s3_uri = f"s3://{bucket}/{key}"
else:
raise ValueError(f"Unable to parse S3 URL: {https_url}")
# Extract filename from key
filename = key.split("/")[-1] if "/" in key else key
return s3_uri, filename
def transform_create_file_response(
self,
model: Optional[str],
@@ -452,21 +510,18 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
"""
# 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
# Use the actual upload URL that was used for the S3 upload
upload_url = litellm_params.get("upload_url")
file_id: str = ""
filename: str = ""
if upload_url:
# Convert HTTPS S3 URL to s3:// URI format
file_id, filename = self._convert_https_url_to_s3_uri(upload_url)
return OpenAIFileObject(
purpose="batch", # Default purpose for Bedrock files
id=file_id,
litellm/llms/custom_httpx/llm_http_handler.py
+16 -7
View File
@@ -2200,7 +2200,6 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params,
optional_params={},
)
if _is_async:
return self.async_create_file(
transformed_request=transformed_request,
@@ -2217,6 +2216,7 @@ class BaseLLMHTTPHandler:
sync_httpx_client = _get_httpx_client()
else:
sync_httpx_client = client
if isinstance(transformed_request, dict) and "method" in transformed_request:
# Handle pre-signed requests (e.g., from Bedrock S3 uploads)
@@ -2284,11 +2284,15 @@ class BaseLLMHTTPHandler:
provider_config=provider_config,
)
# Store the upload URL in litellm_params for the transformation method
litellm_params_with_url = dict(litellm_params)
litellm_params_with_url["upload_url"] = api_base
return provider_config.transform_create_file_response(
model=None,
raw_response=upload_response,
logging_obj=logging_obj,
litellm_params=litellm_params,
litellm_params=litellm_params_with_url,
)
async def async_create_file(
@@ -2411,15 +2415,19 @@ class BaseLLMHTTPHandler:
_is_async: bool = False,
client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
model: Optional[str] = None,
) -> Union["LiteLLMBatch", Coroutine[Any, Any, "LiteLLMBatch"]]:
"""
Creates a batch using provider-specific batch creation process
"""
# get config from model, custom llm provider
if model is None:
raise ValueError("model is required for create_batch")
headers = provider_config.validate_environment(
api_key=api_key,
headers=headers,
model="",
model=model,
messages=[],
optional_params={},
litellm_params=litellm_params,
@@ -2428,7 +2436,7 @@ class BaseLLMHTTPHandler:
api_base = provider_config.get_complete_batch_url(
api_base=api_base,
api_key=api_key,
model="",
model=model,
optional_params={},
litellm_params=litellm_params,
data=create_batch_data,
@@ -2438,7 +2446,7 @@ class BaseLLMHTTPHandler:
# Get the transformed request data
transformed_request = provider_config.transform_create_batch_request(
model="",
model=model,
create_batch_data=create_batch_data,
litellm_params=litellm_params,
optional_params={},
@@ -2506,7 +2514,7 @@ class BaseLLMHTTPHandler:
}
return provider_config.transform_create_batch_response(
model=None,
model=model,
raw_response=batch_response,
logging_obj=logging_obj,
litellm_params=litellm_params_with_request,
@@ -2523,6 +2531,7 @@ class BaseLLMHTTPHandler:
client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
create_batch_data: Optional["CreateBatchRequest"] = None,
model: Optional[str] = None,
):
"""
Async version of create_batch
@@ -2591,7 +2600,7 @@ class BaseLLMHTTPHandler:
}
return provider_config.transform_create_batch_response(
model=None,
model=model,
raw_response=batch_response,
logging_obj=logging_obj,
litellm_params=litellm_params_with_request,
litellm/llms/lm_studio/chat/transformation.py
+2 -2
View File
@@ -15,8 +15,8 @@ class LMStudioChatConfig(OpenAIGPTConfig):
) -> Tuple[Optional[str], Optional[str]]:
api_base = api_base or get_secret_str("LM_STUDIO_API_BASE") # type: ignore
dynamic_api_key = (
api_key or get_secret_str("LM_STUDIO_API_KEY") or " "
) # vllm does not require an api key
api_key or get_secret_str("LM_STUDIO_API_KEY") or "fake-api-key"
) # LM Studio does not require an api key, but OpenAI client requires non-None value
return api_base, dynamic_api_key
def map_openai_params(
litellm/llms/vertex_ai/files/transformation.py
+8 -4
View File
@@ -6,6 +6,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
from httpx import Headers, Response
from litellm.files.utils import FilesAPIUtils
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 (
@@ -260,10 +261,13 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
raise ValueError("file is required")
extracted_file_data = extract_file_data(file_data)
extracted_file_data_content = extracted_file_data.get("content")
if (
create_file_data.get("purpose") == "batch"
and extracted_file_data.get("content_type") == "application/jsonl"
and extracted_file_data_content is not None
if extracted_file_data_content is None:
raise ValueError("file content is required")
if FilesAPIUtils.is_batch_jsonl_file(
create_file_data=create_file_data,
extracted_file_data=extracted_file_data,
):
## 1. If jsonl, check if there's a model name
file_content = self._get_content_from_openai_file(
litellm/main.py
+1
View File
@@ -1333,6 +1333,7 @@ def completion( # type: ignore # noqa: PLR0915
azure_scope=kwargs.get("azure_scope"),
max_retries=max_retries,
timeout=timeout,
litellm_request_debug=kwargs.get("litellm_request_debug", False),
)
cast(LiteLLMLoggingObj, logging).update_environment_variables(
model=model,
litellm/proxy/_new_secret_config.yaml
+7 -12
View File
@@ -1,17 +1,12 @@
model_list:
- model_name: fake-openai-endpoint
- model_name: byok-fixed-gpt-4o-mini
litellm_params:
model: openai/fake
api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/
- model_name: wildcard_models/*
model: openai/gpt-4o-mini
api_base: "https://webhook.site/2f385e05-00aa-402b-86d1-efc9261471a5"
api_key: dummy
- model_name: "byok-wildcard/*"
litellm_params:
model: openai/*
- model_name: hosted_vllm/*
- model_name: xai-grok-3
litellm_params:
model: hosted_vllm/*
api_base: https://webhook.site/6fbe498e-88b5-4a5f-8f07-edb9806c1937
api_key: fake-key
- model_name: deepseek-r1-5b
litellm_params:
model: ollama_chat/deepseek-r1:1.5b
model: xai/grok-3
litellm/proxy/auth/auth_checks.py
-2
View File
@@ -1211,7 +1211,6 @@ def _check_model_access_helper(
models: List[str],
team_model_aliases: Optional[Dict[str, str]] = None,
team_id: Optional[str] = None,
object_type: Literal["user", "team", "key", "org"] = "user",
) -> bool:
## check if model in allowed model names
from collections import defaultdict
@@ -1316,7 +1315,6 @@ def _can_object_call_model(
models=models,
team_model_aliases=team_model_aliases,
team_id=team_id,
object_type=object_type,
):
return True
litellm/proxy/example_config_yaml/oai_misc_config.yaml
+13
View File
@@ -18,6 +18,19 @@ model_list:
litellm_params:
model: "groq/*"
api_key: os.environ/GROQ_API_KEY
- model_name: bedrock/batch-anthropic.claude-3-5-sonnet-20240620-v1:0
litellm_params:
model: bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0
#########################################################
########## batch specific params ########################
s3_bucket_name: litellm-proxy
s3_region_name: us-west-2
s3_access_key_id: os.environ/AWS_ACCESS_KEY_ID
s3_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
aws_batch_role_arn: arn:aws:iam::888602223428:role/service-role/AmazonBedrockExecutionRoleForAgents_BB9HNW6V4CV
model_info:
mode: batch
litellm_settings:
# set_verbose: True # Uncomment this if you want to see verbose logs; not recommended in production
drop_params: True
litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py
+1
View File
@@ -18,6 +18,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
application_id=litellm_params.application_id,
monitor_mode=litellm_params.monitor_mode,
block_failures=litellm_params.block_failures,
anonymize_input=litellm_params.anonymize_input,
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
)
litellm/proxy/guardrails/guardrail_hooks/noma/noma.py
+384 -55
View File
@@ -5,9 +5,10 @@
#
# +-------------------------------------------------------------+
import asyncio
import copy
import os
from typing import Any, Dict, Literal, Optional, Union
from typing import Any, Dict, Final, Literal, Optional, Union
from urllib.parse import urljoin
from fastapi import HTTPException
@@ -24,6 +25,15 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import EmbeddingResponse, ImageResponse
# Constants
USER_ROLE: Final[Literal["user"]] = "user"
ASSISTANT_ROLE: Final[Literal["assistant"]] = "assistant"
SENSITIVE_DATA_DETECTOR_KEYS: Final[list[str]] = ["sensitiveData", "dataDetector"]
# Type aliases
MessageRole = Literal["user", "assistant"]
LLMResponse = Union[Any, ModelResponse, EmbeddingResponse, ImageResponse]
class NomaBlockedMessage(HTTPException):
"""Exception raised when Noma guardrail blocks a message"""
@@ -77,6 +87,7 @@ class NomaBlockedMessage(HTTPException):
"allowedTopics",
"bannedTopics",
"topicGuardrails",
"topicDetector", # Mock name for tests
] and isinstance(value, dict):
filtered_topics = {}
for topic, topic_result in value.items():
@@ -86,7 +97,7 @@ class NomaBlockedMessage(HTTPException):
if filtered_topics:
result[key] = filtered_topics
elif key == "sensitiveData" and isinstance(value, dict):
elif key in SENSITIVE_DATA_DETECTOR_KEYS and isinstance(value, dict):
filtered_sensitive = {}
for data_type, data_result in value.items():
if self._is_result_true(data_result):
@@ -135,6 +146,7 @@ class NomaGuardrail(CustomGuardrail):
application_id: Optional[str] = None,
monitor_mode: Optional[bool] = None,
block_failures: Optional[bool] = None,
anonymize_input: Optional[bool] = None,
**kwargs,
):
self.async_handler = get_async_httpx_client(
@@ -162,8 +174,326 @@ class NomaGuardrail(CustomGuardrail):
else:
self.block_failures = block_failures
if anonymize_input is None:
self.anonymize_input = (
os.environ.get("NOMA_ANONYMIZE_INPUT", "false").lower() == "true"
)
else:
self.anonymize_input = anonymize_input
super().__init__(**kwargs)
def _create_background_noma_check(
self,
coro,
) -> None:
"""Create a background task for Noma API calls without blocking the main flow"""
try:
asyncio.create_task(coro)
except Exception as e:
verbose_proxy_logger.error(
f"Failed to create background Noma task: {str(e)}"
)
async def _process_user_message_check(
self,
request_data: dict,
user_auth: UserAPIKeyAuth,
) -> Optional[str]:
"""Shared logic for processing user message checks"""
extra_data = self.get_guardrail_dynamic_request_body_params(request_data)
user_message = await self._extract_user_message(request_data)
if not user_message:
return None
payload = {"request": {"text": user_message}}
response_json = await self._call_noma_api(
payload=payload,
llm_request_id=None,
request_data=request_data,
user_auth=user_auth,
extra_data=extra_data,
)
if self.monitor_mode:
await self._handle_verdict_background(
USER_ROLE, user_message, response_json
)
return user_message
# Check if we should anonymize content
if self._should_anonymize(response_json, USER_ROLE):
anonymized_content = self._extract_anonymized_content(
response_json, USER_ROLE
)
if anonymized_content:
# Replace the user message content with anonymized version
self._replace_user_message_content(request_data, anonymized_content)
verbose_proxy_logger.debug(
f"Noma guardrail anonymized user message: {anonymized_content}"
)
return anonymized_content
await self._check_verdict(USER_ROLE, user_message, response_json)
return user_message
async def _process_llm_response_check(
self,
request_data: dict,
response: LLMResponse,
user_auth: UserAPIKeyAuth,
) -> Optional[str]:
"""Shared logic for processing LLM response checks"""
extra_data = self.get_guardrail_dynamic_request_body_params(request_data)
if not isinstance(response, litellm.ModelResponse):
return None
content = None
for choice in response.choices:
if isinstance(choice, litellm.Choices) and choice.message.content:
content = choice.message.content
break
if not content or not isinstance(content, str):
return None
payload = {"response": {"text": content}}
response_json = await self._call_noma_api(
payload=payload,
llm_request_id=response.id,
request_data=request_data,
user_auth=user_auth,
extra_data=extra_data,
)
if self.monitor_mode:
await self._handle_verdict_background(
ASSISTANT_ROLE, content, response_json
)
return content
# Check if we should anonymize content
if self._should_anonymize(response_json, ASSISTANT_ROLE):
anonymized_content = self._extract_anonymized_content(
response_json, ASSISTANT_ROLE
)
if anonymized_content:
# Replace the LLM response content with anonymized version
self._replace_llm_response_content(response, anonymized_content)
verbose_proxy_logger.debug(
f"Noma guardrail anonymized LLM response: {anonymized_content}"
)
return anonymized_content
await self._check_verdict(ASSISTANT_ROLE, content, response_json)
return content
def _should_only_sensitive_data_failed(self, classification_obj: dict) -> bool:
"""
Check if only sensitive data detectors (PII, PCI, secrets) have result=true in the classification.
Args:
classification_obj: The prompt or response classification object from Noma API
Returns:
True if only sensitiveData detectors have result=true, False otherwise
"""
if not classification_obj:
return False
# Track which detectors have result=true (detected violations)
failed_detectors = []
sensitive_data_detected = False
for key, value in classification_obj.items():
if key in SENSITIVE_DATA_DETECTOR_KEYS and isinstance(value, dict):
# Check if any sensitive data detector has result=true
for data_type, data_result in value.items():
if self._is_result_true(data_result):
sensitive_data_detected = True
# Don't add to failed_detectors as we want to allow these
elif isinstance(value, dict) and "result" in value:
# Check other detectors - these should NOT have result=true
if self._is_result_true(value):
failed_detectors.append(key)
elif isinstance(value, dict):
# Handle nested detectors
for nested_key, nested_value in value.items():
if self._is_result_true(nested_value):
failed_detectors.append(f"{key}.{nested_key}")
# Return True only if sensitive data was detected AND no other detectors have result=true
return sensitive_data_detected and len(failed_detectors) == 0
def _extract_anonymized_content(
self, response_json: dict, message_type: MessageRole
) -> Optional[str]:
"""
Extract anonymized content from Noma API response.
Args:
response_json: The full response from Noma API
message_type: Either 'user' or 'assistant' to determine which content to extract
Returns:
The anonymized content string if available, None otherwise
"""
original_response = response_json.get("originalResponse", {})
if message_type == USER_ROLE:
prompt_data = original_response.get("prompt", {})
anonymized_data = prompt_data.get("anonymizedContent", {})
return anonymized_data.get("anonymized")
elif message_type == ASSISTANT_ROLE:
response_data = original_response.get("response", {})
anonymized_data = response_data.get("anonymizedContent", {})
return anonymized_data.get("anonymized")
return None
def _should_anonymize(self, response_json: dict, message_type: MessageRole) -> bool:
"""
Determine if content should be anonymized based on Noma API response.
Logic:
- If verdict=True: Content is safe, anonymize if anonymized version exists
- If verdict=False: Check if only sensitiveData detectors have result=True
- If yes: Anonymize
- If no: Block (other violations detected)
Args:
response_json: The full response from Noma API
message_type: Either 'user' or 'assistant' to determine which classification to check
Returns:
True if content should be anonymized, False if it should be blocked
"""
# Only anonymize in blocking mode when anonymize_input is enabled
if self.monitor_mode or not self.anonymize_input:
return False
verdict = response_json.get("verdict", True)
# If verdict is True, anonymize (content is considered safe)
if verdict:
return True
# If verdict is False, check if only sensitive data detectors have result=True
original_response = response_json.get("originalResponse", {})
if message_type == USER_ROLE:
classification_obj = original_response.get("prompt", {})
elif message_type == ASSISTANT_ROLE:
classification_obj = original_response.get("response", {})
else:
return False
# Anonymize only if solely sensitive data (PII/PCI/secrets) was detected
return self._should_only_sensitive_data_failed(classification_obj)
def _is_result_true(self, result_obj: Optional[Dict[str, Any]]) -> bool:
"""
Check if a result object has a "result" field that is True.
Args:
result_obj: A dictionary that may contain a "result" field
Returns:
True if the "result" field exists and is True, False otherwise
"""
if not result_obj or not isinstance(result_obj, dict):
return False
return result_obj.get("result") is True
def _replace_user_message_content(
self, request_data: dict, anonymized_content: str
):
"""
Replace the user message content in request data with anonymized version.
Args:
request_data: The original request data
anonymized_content: The anonymized content to replace with
"""
messages = request_data.get("messages", [])
if not messages:
return
# Find and replace the last user message
for i in range(len(messages) - 1, -1, -1):
if messages[i].get("role") == USER_ROLE:
messages[i]["content"] = anonymized_content
break
def _replace_llm_response_content(
self, response: LLMResponse, anonymized_content: str
):
"""
Replace the LLM response content with anonymized version.
Args:
response: The original LLM response
anonymized_content: The anonymized content to replace with
"""
if not isinstance(response, litellm.ModelResponse):
return
# Replace content in all choices
for choice in response.choices:
if isinstance(choice, litellm.Choices) and choice.message.content:
choice.message.content = anonymized_content
async def _check_user_message_background(
self,
request_data: dict,
user_auth: UserAPIKeyAuth,
) -> None:
"""Check user message in background for monitor mode - non-blocking"""
try:
await self._process_user_message_check(request_data, user_auth)
except Exception as e:
verbose_proxy_logger.error(
f"Noma background user message check failed: {str(e)}"
)
async def _check_llm_response_background(
self,
request_data: dict,
response: LLMResponse,
user_auth: UserAPIKeyAuth,
) -> None:
"""Check LLM response in background for monitor mode - non-blocking"""
try:
await self._process_llm_response_check(request_data, response, user_auth)
except Exception as e:
verbose_proxy_logger.error(
f"Noma background response check failed: {str(e)}"
)
async def _handle_verdict_background(
self,
type: MessageRole,
message: str,
response_json: dict,
) -> None:
"""Handle verdict from Noma API in background - logging only, never blocks"""
try:
if not response_json.get("verdict", True):
msg = f"Noma guardrail blocked {type} message: {message}"
verbose_proxy_logger.warning(msg)
else:
msg = f"Noma guardrail allowed {type} message: {message}"
verbose_proxy_logger.info(msg)
except Exception as e:
verbose_proxy_logger.error(
f"Noma background verdict handling failed: {str(e)}"
)
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
@@ -191,6 +521,18 @@ class NomaGuardrail(CustomGuardrail):
):
return data
# In monitor mode, run Noma check in background and return immediately
if self.monitor_mode:
try:
self._create_background_noma_check(
self._check_user_message_background(data, user_api_key_dict)
)
except Exception as e:
verbose_proxy_logger.error(
f"Failed to start background Noma pre-call check: {str(e)}"
)
return data
try:
return await self._check_user_message(data, user_api_key_dict)
except NomaBlockedMessage:
@@ -198,7 +540,7 @@ class NomaGuardrail(CustomGuardrail):
except Exception as e:
verbose_proxy_logger.error(f"Noma pre-call hook failed: {str(e)}")
if self.block_failures and not self.monitor_mode:
if self.block_failures:
raise
return data
@@ -220,6 +562,18 @@ class NomaGuardrail(CustomGuardrail):
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
return data
# In monitor mode, run Noma check in background and return immediately
if self.monitor_mode:
try:
self._create_background_noma_check(
self._check_user_message_background(data, user_api_key_dict)
)
except Exception as e:
verbose_proxy_logger.error(
f"Failed to start background Noma moderation check: {str(e)}"
)
return data
try:
return await self._check_user_message(data, user_api_key_dict)
except NomaBlockedMessage:
@@ -227,7 +581,7 @@ class NomaGuardrail(CustomGuardrail):
except Exception as e:
verbose_proxy_logger.error(f"Noma moderation hook failed: {str(e)}")
if self.block_failures and not self.monitor_mode:
if self.block_failures:
raise
return data
@@ -235,19 +589,33 @@ class NomaGuardrail(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
response: Union[Any, ModelResponse, EmbeddingResponse, ImageResponse],
response: LLMResponse,
):
event_type: GuardrailEventHooks = GuardrailEventHooks.post_call
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
return response
# In monitor mode, run Noma check in background and return immediately
if self.monitor_mode:
try:
self._create_background_noma_check(
self._check_llm_response_background(
data, response, user_api_key_dict
)
)
except Exception as e:
verbose_proxy_logger.error(
f"Failed to start background Noma post-call check: {str(e)}"
)
return response
try:
return await self._check_llm_response(data, response, user_api_key_dict)
except NomaBlockedMessage:
raise
except Exception as e:
verbose_proxy_logger.error(f"Noma post-call hook failed: {str(e)}")
if self.block_failures and not self.monitor_mode:
if self.block_failures:
raise
return response
@@ -257,55 +625,24 @@ class NomaGuardrail(CustomGuardrail):
user_auth: UserAPIKeyAuth,
) -> Union[Exception, str, dict, None]:
"""Check user message for policy violations"""
extra_data = self.get_guardrail_dynamic_request_body_params(request_data)
user_message = await self._extract_user_message(request_data)
user_message = await self._process_user_message_check(request_data, user_auth)
if not user_message:
return request_data
payload = {"request": {"text": user_message}}
response_json = await self._call_noma_api(
payload=payload,
llm_request_id=None,
request_data=request_data,
user_auth=user_auth,
extra_data=extra_data,
)
await self._check_verdict("user", user_message, response_json)
return request_data
async def _check_llm_response(
self,
request_data: dict,
response: Union[Any, ModelResponse, EmbeddingResponse, ImageResponse],
response: LLMResponse,
user_auth: UserAPIKeyAuth,
) -> Union[Exception, ModelResponse, Any]:
"""Check LLM response for policy violations"""
extra_data = self.get_guardrail_dynamic_request_body_params(request_data)
if not isinstance(response, litellm.ModelResponse):
return response
content = None
for choice in response.choices:
if isinstance(choice, litellm.Choices) and choice.message.content:
content = choice.message.content
break
if not content or not isinstance(content, str):
return response
payload = {"response": {"text": content}}
response_json = await self._call_noma_api(
payload=payload,
llm_request_id=response.id,
request_data=request_data,
user_auth=user_auth,
extra_data=extra_data,
content = await self._process_llm_response_check(
request_data, response, user_auth
)
await self._check_verdict("assistant", content, response_json)
if not content:
return response
return response
@@ -316,7 +653,7 @@ class NomaGuardrail(CustomGuardrail):
return None
# Get the last user message
user_messages = [msg for msg in messages if msg.get("role") == "user"]
user_messages = [msg for msg in messages if msg.get("role") == USER_ROLE]
if not user_messages:
return None
@@ -371,7 +708,7 @@ class NomaGuardrail(CustomGuardrail):
async def _check_verdict(
self,
type: Literal["user", "assistant"],
type: MessageRole,
message: str,
response_json: dict,
) -> None:
@@ -379,11 +716,7 @@ class NomaGuardrail(CustomGuardrail):
Check the verdict from the Noma API and raise an exception if needed
"""
if not response_json.get("verdict", True):
msg = str.format(
"Noma guardrail blocked {type} message: {message}",
type=type,
message=message,
)
msg = f"Noma guardrail blocked {type} message: {message}"
if self.monitor_mode:
verbose_proxy_logger.warning(msg)
@@ -392,11 +725,7 @@ class NomaGuardrail(CustomGuardrail):
original_response = response_json.get("originalResponse", {})
raise NomaBlockedMessage(original_response)
else:
msg = str.format(
"Noma guardrail allowed {type} message: {message}",
type=type,
message=message,
)
msg = f"Noma guardrail allowed {type} message: {message}"
if self.monitor_mode:
verbose_proxy_logger.info(msg)
else:
litellm/proxy/hooks/__init__.py
+3 -3
View File
@@ -17,13 +17,13 @@ except ImportError:
# List of all available hooks that can be enabled
PROXY_HOOKS = {
"max_budget_limiter": _PROXY_MaxBudgetLimiter,
"parallel_request_limiter": _PROXY_MaxParallelRequestsHandler,
"parallel_request_limiter": _PROXY_MaxParallelRequestsHandler_v3,
"cache_control_check": _PROXY_CacheControlCheck,
}
## FEATURE FLAG HOOKS ##
if os.getenv("EXPERIMENTAL_MULTI_INSTANCE_RATE_LIMITING", "false").lower() == "true":
PROXY_HOOKS["parallel_request_limiter"] = _PROXY_MaxParallelRequestsHandler_v3
if os.getenv("LEGACY_MULTI_INSTANCE_RATE_LIMITING", "false").lower() == "true":
PROXY_HOOKS["parallel_request_limiter"] = _PROXY_MaxParallelRequestsHandler
### update PROXY_HOOKS with ENTERPRISE_PROXY_HOOKS ###
litellm/proxy/hooks/parallel_request_limiter_v3.py
+26 -18
View File
@@ -6,6 +6,7 @@ This is currently in development and not yet ready for production.
import os
from datetime import datetime
from math import floor
from typing import (
TYPE_CHECKING,
Any,
@@ -17,7 +18,7 @@ from typing import (
Union,
cast,
)
from math import floor
from fastapi import HTTPException
from litellm import DualCache
@@ -95,6 +96,7 @@ end
return results
"""
class RateLimitDescriptorRateLimitObject(TypedDict, total=False):
requests_per_unit: Optional[int]
tokens_per_unit: Optional[int]
@@ -266,7 +268,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if current_limit is None or rate_limit_type is None:
continue
if counter_value is not None and int(counter_value) + 1 > current_limit:
if counter_value is not None and int(counter_value) > current_limit:
overall_code = "OVER_LIMIT"
item_code = "OVER_LIMIT"
@@ -480,10 +482,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
},
)
)
# Team Member rate limits
if user_api_key_dict.user_id and (user_api_key_dict.team_member_rpm_limit is not None or user_api_key_dict.team_member_tpm_limit is not None):
team_member_value = f"{user_api_key_dict.team_id}:{user_api_key_dict.user_id}"
if user_api_key_dict.user_id and (
user_api_key_dict.team_member_rpm_limit is not None
or user_api_key_dict.team_member_tpm_limit is not None
):
team_member_value = (
f"{user_api_key_dict.team_id}:{user_api_key_dict.user_id}"
)
descriptors.append(
RateLimitDescriptor(
key="team_member",
@@ -557,13 +564,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# Find which descriptor hit the limit
for i, status in enumerate(response["statuses"]):
if status["code"] == "OVER_LIMIT":
descriptor = descriptors[floor(i/2)]
descriptor = descriptors[floor(i / 2)]
raise HTTPException(
status_code=429,
detail=f"Rate limit exceeded for {descriptor['key']}: {descriptor['value']}. Remaining: {status['limit_remaining']}",
headers={
"retry-after": str(self.window_size),
"rate_limit_type": str(status["rate_limit_type"])
"rate_limit_type": str(status["rate_limit_type"]),
}, # Retry after 1 minute
)
@@ -613,7 +620,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# Check if script is available
if self.token_increment_script is None:
verbose_proxy_logger.debug("TTL preservation script not available, using regular pipeline")
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,
@@ -628,7 +637,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
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}"
@@ -693,16 +702,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
# Get metadata from kwargs
user_api_key = kwargs["litellm_params"]["metadata"].get("user_api_key")
user_api_key_user_id = kwargs["litellm_params"]["metadata"].get(
"user_api_key_user_id"
litellm_metadata = kwargs["litellm_params"]["metadata"]
if litellm_metadata is None:
return
user_api_key = litellm_metadata.get("user_api_key")
user_api_key_user_id = litellm_metadata.get("user_api_key_user_id")
user_api_key_team_id = litellm_metadata.get("user_api_key_team_id")
user_api_key_end_user_id = kwargs.get("user") or litellm_metadata.get(
"user_api_key_end_user_id"
)
user_api_key_team_id = kwargs["litellm_params"]["metadata"].get(
"user_api_key_team_id"
)
user_api_key_end_user_id = kwargs.get("user") or kwargs["litellm_params"][
"metadata"
].get("user_api_key_end_user_id")
model_group = get_model_group_from_litellm_kwargs(kwargs)
# Get total tokens from response
litellm/proxy/litellm_pre_call_utils.py
+42 -18
View File
@@ -14,10 +14,10 @@ from litellm.proxy._types import (
AddTeamCallback,
CommonProxyErrors,
LitellmDataForBackendLLMCall,
LitellmUserRoles,
SpecialHeaders,
TeamCallbackMetadata,
UserAPIKeyAuth,
LitellmUserRoles,
)
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.router import Router
@@ -272,7 +272,7 @@ class LiteLLMProxyRequestSetup:
if timeout_header is not None:
return float(timeout_header)
return None
@staticmethod
def _get_stream_timeout_from_request(headers: dict) -> Optional[float]:
"""
@@ -292,13 +292,14 @@ class LiteLLMProxyRequestSetup:
if num_retries_header is not None:
return int(num_retries_header)
return None
@staticmethod
def _get_spend_logs_metadata_from_request_headers(headers: dict) -> Optional[dict]:
"""
Get the `spend_logs_metadata` from the request headers.
"""
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
spend_logs_metadata_header = headers.get("x-litellm-spend-logs-metadata", None)
if spend_logs_metadata_header is not None:
return safe_json_loads(spend_logs_metadata_header)
@@ -337,16 +338,24 @@ class LiteLLMProxyRequestSetup:
return None
@staticmethod
def add_internal_user_from_user_mapping(general_settings: Optional[Dict], user_api_key_dict: UserAPIKeyAuth, headers: dict) -> UserAPIKeyAuth:
def add_internal_user_from_user_mapping(
general_settings: Optional[Dict],
user_api_key_dict: UserAPIKeyAuth,
headers: dict,
) -> UserAPIKeyAuth:
if general_settings is None:
return user_api_key_dict
user_header_mapping = general_settings.get("user_header_mappings")
if not user_header_mapping:
return user_api_key_dict
header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(user_header_mapping)
header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(
user_header_mapping
)
if not header_name:
return user_api_key_dict
header_value = LiteLLMProxyRequestSetup._get_case_insensitive_header(headers, header_name)
header_value = LiteLLMProxyRequestSetup._get_case_insensitive_header(
headers, header_name
)
if header_value:
user_api_key_dict.user_id = header_value
return user_api_key_dict
@@ -429,15 +438,25 @@ class LiteLLMProxyRequestSetup:
"""
Add headers to the LLM call by model group
"""
from litellm.proxy.auth.auth_checks import _check_model_access_helper
from litellm.proxy.proxy_server import llm_router
data_model = data.get("model")
if (
data_model is not None
and litellm.model_group_settings is not None
and litellm.model_group_settings.forward_client_headers_to_llm_api
is not None
and data_model
in litellm.model_group_settings.forward_client_headers_to_llm_api
and _check_model_access_helper(
model=data_model,
llm_router=llm_router,
models=litellm.model_group_settings.forward_client_headers_to_llm_api,
team_model_aliases=user_api_key_dict.team_model_aliases,
team_id=user_api_key_dict.team_id,
) # handles aliases, wildcards, etc.
):
_headers = LiteLLMProxyRequestSetup.add_headers_to_llm_call(
headers, user_api_key_dict
)
@@ -497,8 +516,10 @@ class LiteLLMProxyRequestSetup:
timeout = LiteLLMProxyRequestSetup._get_timeout_from_request(headers)
if timeout is not None:
data["timeout"] = timeout
stream_timeout = LiteLLMProxyRequestSetup._get_stream_timeout_from_request(headers)
stream_timeout = LiteLLMProxyRequestSetup._get_stream_timeout_from_request(
headers
)
if stream_timeout is not None:
data["stream_timeout"] = stream_timeout
@@ -507,7 +528,7 @@ class LiteLLMProxyRequestSetup:
data["num_retries"] = num_retries
return data
@staticmethod
def add_litellm_metadata_from_request_headers(
headers: dict,
@@ -520,11 +541,16 @@ class LiteLLMProxyRequestSetup:
Relevant issue: https://github.com/BerriAI/litellm/issues/14008
"""
from litellm.proxy._types import LitellmMetadataFromRequestHeaders
metadata_from_headers = LitellmMetadataFromRequestHeaders()
spend_logs_metadata = LiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headers(headers)
spend_logs_metadata = (
LiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headers(
headers
)
)
if spend_logs_metadata is not None:
metadata_from_headers["spend_logs_metadata"] = spend_logs_metadata
#########################################################################################
# Finally update the requests metadata with the `metadata_from_headers`
#########################################################################################
@@ -714,7 +740,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
_headers = clean_headers(
request.headers,
litellm_key_header_name=(
@@ -740,8 +765,6 @@ async def add_litellm_data_to_request( # noqa: PLR0915
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,
@@ -763,7 +786,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915
data=data, headers=_headers, user_api_key_dict=user_api_key_dict
)
user_api_key_dict = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(general_settings, user_api_key_dict, _headers)
user_api_key_dict = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(
general_settings, user_api_key_dict, _headers
)
# Parse user info from headers
user = LiteLLMProxyRequestSetup.get_user_from_headers(_headers, general_settings)
@@ -773,7 +798,6 @@ async def add_litellm_data_to_request( # noqa: PLR0915
if "user" not in data:
data["user"] = user
data["secret_fields"] = SecretFields(raw_headers=dict(request.headers))
## Dynamic api version (Azure OpenAI endpoints) ##
litellm/proxy/management_endpoints/key_management_endpoints.py
+4 -1
View File
@@ -2899,7 +2899,10 @@ async def unblock_key(
param="key",
code=status.HTTP_400_BAD_REQUEST,
)
hashed_token = hash_token(token=data.key)
if data.key.startswith("sk-"):
hashed_token = hash_token(token=data.key)
else:
hashed_token = data.key
if litellm.store_audit_logs is True:
# make an audit log for key update
litellm/proxy/proxy_config.yaml
+11 -16
View File
@@ -1,18 +1,13 @@
model_list:
- model_name: db-openai-endpoint
- model_name: bedrock/batch-anthropic.claude-3-5-sonnet-20240620-v1:0
litellm_params:
model: openai/*
api_base: https://exampleopenaiendpoint-production-0ee2.up.railway.app/
- model_name: bedrock/*
litellm_params:
model: bedrock/*
- model_name: openai/*
litellm_params:
model: openai/*
- model_name: dashscope/*
litellm_params:
model: dashscope/*
litellm_settings:
callbacks: ["cloudzero"]
model: bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0
#########################################################
########## batch specific params ########################
s3_bucket_name: litellm-proxy
s3_region_name: us-west-2
s3_access_key_id: os.environ/AWS_ACCESS_KEY_ID
s3_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
aws_batch_role_arn: arn:aws:iam::888602223428:role/service-role/AmazonBedrockExecutionRoleForAgents_BB9HNW6V4CV
model_info:
mode: batch
litellm/proxy/spend_tracking/spend_management_endpoints.py
+9
View File
@@ -1660,6 +1660,9 @@ async def ui_view_spend_logs( # noqa: PLR0915
model: Optional[str] = fastapi.Query(
default=None, description="Filter logs by model"
),
key_alias: Optional[str] = fastapi.Query(
default=None, description="Filter logs by key alias"
),
):
"""
View spend logs for UI with pagination support
@@ -1727,6 +1730,12 @@ async def ui_view_spend_logs( # noqa: PLR0915
if model is not None:
where_conditions["model"] = model
if key_alias is not None:
where_conditions["metadata"] = {
"path": ["user_api_key_alias"],
"string_contains": key_alias
}
if min_spend is not None or max_spend is not None:
where_conditions["spend"] = {}
litellm/router.py
+1 -1
View File
@@ -3046,7 +3046,7 @@ class Router:
from litellm.router_utils.common_utils import add_model_file_id_mappings
verbose_router_logger.debug(
f"Inside _atext_completion()- model: {model}; kwargs: {kwargs}"
f"Inside _acreate_file()- model: {model}; kwargs: {kwargs}"
)
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
healthy_deployments = await self.async_get_healthy_deployments(
litellm/types/guardrails.py
+8 -3
View File
@@ -42,6 +42,7 @@ class SupportedGuardrailIntegrations(Enum):
OPENAI_MODERATION = "openai_moderation"
NOMA = "noma"
class Role(Enum):
SYSTEM = "system"
ASSISTANT = "assistant"
@@ -312,7 +313,6 @@ class BedrockGuardrailConfigModel(BaseModel):
)
class LakeraV2GuardrailConfigModel(BaseModel):
"""Configuration parameters for the Lakera AI v2 guardrail"""
@@ -375,6 +375,10 @@ class NomaGuardrailConfigModel(BaseModel):
default=None,
description="If True, blocks requests on API failures. Defaults to True if not provided",
)
anonymize_input: Optional[bool] = Field(
default=None,
description="If True, replaces sensitive content with anonymized version when only PII/PCI/secrets are detected. Only applies in blocking mode. Defaults to False if not provided",
)
class BaseLitellmParams(BaseModel): # works for new and patch update guardrails
@@ -425,7 +429,8 @@ class BaseLitellmParams(BaseModel): # works for new and patch update guardrails
)
model: Optional[str] = Field(
default=None, description="Optional field if guardrail requires a 'model' parameter"
default=None,
description="Optional field if guardrail requires a 'model' parameter",
)
# Model Armor params
@@ -446,7 +451,7 @@ class BaseLitellmParams(BaseModel): # works for new and patch update guardrails
default=True,
description="Whether to fail the request if Model Armor encounters an error",
)
model_config = ConfigDict(extra="allow", protected_namespaces=())
litellm/types/utils.py
+2 -1
View File
@@ -1996,7 +1996,7 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False):
]
guardrail_request: Optional[dict]
guardrail_response: Optional[Union[dict, str, List[dict]]]
guardrail_status: Literal["success", "failure","blocked"]
guardrail_status: Literal["success", "failure", "blocked"]
start_time: Optional[float]
end_time: Optional[float]
duration: Optional[float]
@@ -2124,6 +2124,7 @@ all_litellm_params = [
"metadata",
"litellm_metadata",
"litellm_trace_id",
"litellm_request_debug",
"guardrails",
"tags",
"acompletion",
package-lock.json
Generated
+7975 -4
View File
File diff suppressed because it is too large Load Diff
package.json
+4 -1
View File
@@ -5,6 +5,9 @@
"react-copy-to-clipboard": "^5.1.0"
},
"devDependencies": {
"@types/react-copy-to-clipboard": "^5.0.7"
"@testing-library/jest-dom": "^6.8.0",
"@testing-library/react": "^14.3.1",
"@types/react-copy-to-clipboard": "^5.0.7",
"jest": "^29.7.0"
}
}
tests/batches_tests/bedrock_batch_completions.jsonl
+126 -1
View File
@@ -1,3 +1,128 @@
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
{"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}}
tests/batches_tests/test_bedrock_files_and_batches.py
+50
View File
@@ -18,6 +18,8 @@ sys.path.insert(
import pytest
from typing import Optional
import litellm
from unittest.mock import patch, MagicMock
import httpx
@pytest.mark.asyncio()
@@ -64,7 +66,55 @@ async def test_async_file_and_batch():
input_file_id=file_obj.id,
metadata={"key1": "value1", "key2": "value2"},
custom_llm_provider="bedrock",
#########################################################
# bedrock specific params
#########################################################
model="us.anthropic.claude-3-5-sonnet-20240620-v1:0",
aws_batch_role_arn="arn:aws:iam::888602223428:role/service-role/AmazonBedrockExecutionRoleForAgents_BB9HNW6V4CV"
)
print("CREATED BATCH RESPONSE=", create_batch_response)
@pytest.mark.asyncio()
async def test_mock_bedrock_file_url_mapping():
"""
Simple test to capture PUT URL and validate mapping to file ID.
"""
print("Testing Bedrock file URL mapping")
captured_put_url = None
async def mock_async_create_file(transformed_request, **kwargs):
nonlocal captured_put_url
# Capture PUT URL from transformed request
if isinstance(transformed_request, dict) and "url" in transformed_request:
captured_put_url = transformed_request["url"]
# Call the real method to get actual response
from litellm.files.main import base_llm_http_handler
return await base_llm_http_handler.__class__.async_create_file(
base_llm_http_handler, transformed_request, **kwargs
)
with patch('litellm.files.main.base_llm_http_handler.async_create_file', side_effect=mock_async_create_file):
file_obj = await litellm.acreate_file(
file=open(os.path.join(os.path.dirname(__file__), "bedrock_batch_completions.jsonl"), "rb"),
purpose="batch",
custom_llm_provider="bedrock",
s3_bucket_name="litellm-proxy",
)
print(f"PUT URL: {captured_put_url}")
print(f"File ID: {file_obj.id}")
# Validate URL was captured and response is correct
assert captured_put_url is not None
assert file_obj.id.startswith("s3://")
# Verify mapping
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
bedrock_config = BedrockFilesConfig()
expected_s3_uri, _ = bedrock_config._convert_https_url_to_s3_uri(captured_put_url)
assert file_obj.id == expected_s3_uri
tests/llm_translation/test_gemini.py
+1 -1
View File
@@ -21,7 +21,7 @@ import json
class TestGoogleAIStudioGemini(BaseLLMChatTest):
def get_base_completion_call_args(self) -> dict:
return {"model": "gemini/gemini-2.0-flash"}
return {"model": "gemini/gemini-2.5-flash"}
def get_base_completion_call_args_with_reasoning_model(self) -> dict:
return {"model": "gemini/gemini-2.5-flash"}
tests/local_testing/test_pass_through_endpoints.py
+157 -17
View File
@@ -1,5 +1,7 @@
import os
import sys
import uuid
from functools import partial
from typing import Optional
import pytest
@@ -146,24 +148,36 @@ async def test_pass_through_endpoint_rerank(client):
@pytest.mark.parametrize(
"auth, rpm_limit, expected_error_code",
[(True, 0, 429), (True, 1, 200), (False, 0, 200)],
"auth, rpm_limit, requests_to_make, expected_status_codes, num_users",
[
# Single user tests
(True, 0, 1, [429], 1),
(True, 1, 1, [200], 1),
(True, 1, 2, [200, 429], 1),
(True, 2, 4, [200, 200, 429, 429], 1),
(True, 3, 4, [200, 200, 200, 429], 1),
(True, 4, 4, [200, 200, 200, 200], 1),
(False, 0, 1, [200], 1),
(False, 0, 4, [200, 200, 200, 200], 1),
# Multiple user tests (same parameters as single user)
(True, 0, 1, [429], 2),
(True, 1, 1, [200], 2),
(True, 1, 2, [200, 429], 2),
(True, 2, 4, [200, 200, 429, 429], 2),
(True, 3, 4, [200, 200, 200, 429], 2),
(True, 4, 4, [200, 200, 200, 200], 2),
(False, 0, 1, [200], 2),
(False, 0, 4, [200, 200, 200, 200], 2),
],
)
@pytest.mark.asyncio
async def test_pass_through_endpoint_rpm_limit(
client, auth, expected_error_code, rpm_limit
client, auth, rpm_limit, requests_to_make, expected_status_codes, num_users
):
import litellm
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.proxy_server import ProxyLogging, hash_token, user_api_key_cache
mock_api_key = "sk-my-test-key"
cache_value = UserAPIKeyAuth(token=hash_token(mock_api_key), rpm_limit=rpm_limit)
_cohere_api_key = os.environ.get("COHERE_API_KEY")
user_api_key_cache.set_cache(key=hash_token(mock_api_key), value=cache_value)
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
proxy_logging_obj._init_litellm_callbacks()
@@ -173,6 +187,7 @@ async def test_pass_through_endpoint_rpm_limit(
setattr(litellm.proxy.proxy_server, "proxy_logging_obj", proxy_logging_obj)
# Define a pass-through endpoint
_cohere_api_key = os.environ.get("COHERE_API_KEY")
pass_through_endpoints = [
{
"path": "/v1/rerank",
@@ -190,6 +205,13 @@ async def test_pass_through_endpoint_rpm_limit(
general_settings.update({"pass_through_endpoints": pass_through_endpoints})
setattr(litellm.proxy.proxy_server, "general_settings", general_settings)
# Setup API keys and cache
mock_api_keys = [f"sk-test-{uuid.uuid4().hex}" for _ in range(num_users)]
for mock_api_key in mock_api_keys:
cache_value = UserAPIKeyAuth(token=hash_token(mock_api_key), rpm_limit=rpm_limit)
user_api_key_cache.set_cache(key=hash_token(mock_api_key), value=cache_value)
_json_data = {
"model": "rerank-english-v3.0",
"query": "What is the capital of the United States?",
@@ -200,16 +222,134 @@ async def test_pass_through_endpoint_rpm_limit(
}
# Make a request to the pass-through endpoint
response = client.post(
"/v1/rerank",
json=_json_data,
headers={"Authorization": "Bearer {}".format(mock_api_key)},
)
tasks = []
for mock_api_key in mock_api_keys:
for _ in range(requests_to_make):
task = asyncio.get_running_loop().run_in_executor(
None,
partial(
client.post,
"/v1/rerank",
json=_json_data,
headers={"Authorization": "Bearer {}".format(mock_api_key)},
),
)
tasks.append(task)
responses = await asyncio.gather(*tasks)
if num_users == 1:
status_codes = sorted([response.status_code for response in responses])
assert status_codes == sorted(expected_status_codes)
else:
first_user_responses = responses[requests_to_make:]
second_user_responses = responses[:requests_to_make]
first_user_status_codes = sorted([response.status_code for response in first_user_responses])
second_user_status_codes = sorted([response.status_code for response in second_user_responses])
expected_status_codes.sort()
assert first_user_status_codes == expected_status_codes
assert second_user_status_codes == expected_status_codes
print("JSON response: ", _json_data)
# Assert the response
assert response.status_code == expected_error_code
@pytest.mark.parametrize(
"auth, rpm_limit, requests_to_make, expected_status_codes",
[
# Multiple user tests (same parameters as single user)
(True, 0, 1, [429]),
(True, 1, 1, [200]),
(True, 1, 2, [200, 429]),
(True, 2, 4, [200, 200, 429, 429]),
(True, 3, 4, [200, 200, 200, 429]),
(True, 4, 4, [200, 200, 200, 200]),
(False, 0, 1, [200]),
(False, 0, 4, [200, 200, 200, 200]),
],
)
@pytest.mark.asyncio
async def test_pass_through_endpoint_sequential_rpm_limit(
client, auth, rpm_limit, requests_to_make, expected_status_codes
):
import litellm
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.proxy_server import ProxyLogging, hash_token, user_api_key_cache
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
proxy_logging_obj._init_litellm_callbacks()
setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
setattr(litellm.proxy.proxy_server, "prisma_client", "FAKE-VAR")
setattr(litellm.proxy.proxy_server, "proxy_logging_obj", proxy_logging_obj)
# Define a pass-through endpoint
_cohere_api_key = os.environ.get("COHERE_API_KEY")
pass_through_endpoints = [
{
"path": "/v1/rerank",
"target": "https://api.cohere.com/v1/rerank",
"auth": auth,
"headers": {"Authorization": f"bearer {_cohere_api_key}"},
}
]
# Initialize the pass-through endpoint
await initialize_pass_through_endpoints(pass_through_endpoints)
general_settings: Optional[dict] = (
getattr(litellm.proxy.proxy_server, "general_settings", {}) or {}
)
general_settings.update({"pass_through_endpoints": pass_through_endpoints})
setattr(litellm.proxy.proxy_server, "general_settings", general_settings)
# Setup API keys and cache
mock_api_keys = [f"sk-test-{uuid.uuid4().hex}" for _ in range(2)]
for mock_api_key in mock_api_keys:
cache_value = UserAPIKeyAuth(token=hash_token(mock_api_key), rpm_limit=rpm_limit)
user_api_key_cache.set_cache(key=hash_token(mock_api_key), value=cache_value)
_json_data = {
"model": "rerank-english-v3.0",
"query": "What is the capital of the United States?",
"top_n": 3,
"documents": [
"Carson City is the capital city of the American state of Nevada."
],
}
# Make a request to the pass-through endpoint
first_user_responses = []
second_user_responses = []
for _ in range(requests_to_make):
requests = []
for mock_api_key in mock_api_keys:
task = asyncio.get_running_loop().run_in_executor(
None,
partial(
client.post,
"/v1/rerank",
json=_json_data,
headers={"Authorization": "Bearer {}".format(mock_api_key)},
),
)
requests.append(task)
first_user_response, second_user_response = await asyncio.gather(*requests)
first_user_responses.append(first_user_response)
second_user_responses.append(second_user_response)
first_user_status_codes = sorted([response.status_code for response in first_user_responses])
second_user_status_codes = sorted([response.status_code for response in second_user_responses])
expected_status_codes.sort()
assert first_user_status_codes == expected_status_codes
assert second_user_status_codes == expected_status_codes
print("JSON response: ", _json_data)
@pytest.mark.parametrize(
tests/openai_endpoints_tests/bedrock_batch_completions.jsonl
+2
View File
@@ -0,0 +1,2 @@
{"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}}
tests/openai_endpoints_tests/test_bedrock_batches_api.py
+37
View File
@@ -0,0 +1,37 @@
from openai import OpenAI
import pytest
client = OpenAI(
base_url="http://0.0.0.0:4000",
api_key="sk-1234",
)
BEDROCK_BATCH_MODEL = "bedrock/batch-anthropic.claude-3-5-sonnet-20240620-v1:0"
@pytest.mark.asyncio
async def test_bedrock_batches_api():
"""
Test bedrock batches api
E2E Test Creating a File and a Batch on Bedrock
"""
# Upload file
batch_input_file = client.files.create(
file=open("tests/openai_endpoints_tests/bedrock_batch_completions.jsonl", "rb"),
purpose="batch",
extra_body={"target_model_names": BEDROCK_BATCH_MODEL}
)
print(batch_input_file)
# Create batch
batch = client.batches.create(
input_file_id=batch_input_file.id,
endpoint="/v1/chat/completions",
completion_window="24h",
metadata={"description": "Test batch job"},
)
print(batch)
assert batch.id is not None
tests/proxy_admin_ui_tests/ui_unit_tests/log_filter_logic.test.tsx
+87
View File
@@ -0,0 +1,87 @@
import { uiSpendLogsCall } from '../../../ui/litellm-dashboard/src/components/networking';
// Mock the networking module
jest.mock('../../../ui/litellm-dashboard/src/components/networking', () => ({
uiSpendLogsCall: jest.fn(),
}));
const mockUiSpendLogsCall = uiSpendLogsCall as jest.MockedFunction<typeof uiSpendLogsCall>;
describe('Key Alias Filtering Integration Test', () => {
beforeEach(() => {
jest.clearAllMocks();
});
it('should call API with correct key_alias parameter', async () => {
// Mock API response with both success and failure logs
const mockResponse = {
data: [
{ request_id: 'req-1', status: 'success', metadata: { user_api_key_alias: 'test-key' } },
{ request_id: 'req-2', status: 'failure', metadata: { user_api_key_alias: 'test-key' } }
],
total: 2,
page: 1,
page_size: 50,
total_pages: 1
};
mockUiSpendLogsCall.mockResolvedValueOnce(mockResponse);
// Simulate the API call that would happen when filtering by key alias
const result = await uiSpendLogsCall(
'test-token',
undefined,
undefined,
undefined,
'2024-01-15 09:00:00',
'2024-01-15 11:00:00',
1,
50,
undefined,
undefined,
undefined,
undefined,
'test-key-alias' // key_alias - this is the fix
);
// Verify the API was called correctly
expect(mockUiSpendLogsCall).toHaveBeenCalledWith(
'test-token',
undefined,
undefined,
undefined,
'2024-01-15 09:00:00',
'2024-01-15 11:00:00',
1,
50,
undefined,
undefined,
undefined,
undefined,
'test-key-alias' // The key assertion - this parameter should be passed through
);
// Verify response contains both success and failure logs
expect(result.data).toHaveLength(2);
expect(result.data[0].status).toBe('success');
expect(result.data[1].status).toBe('failure');
});
it('should pass undefined for empty key alias', async () => {
mockUiSpendLogsCall.mockResolvedValueOnce({ data: [], total: 0, page: 1, page_size: 50, total_pages: 0 });
await uiSpendLogsCall(
'test-token', undefined, undefined, undefined,
'2024-01-15 09:00:00', '2024-01-15 11:00:00',
1, 50, undefined, undefined, undefined, undefined,
undefined // Empty string should become undefined
);
expect(mockUiSpendLogsCall).toHaveBeenCalledWith(
'test-token', undefined, undefined, undefined,
'2024-01-15 09:00:00', '2024-01-15 11:00:00',
1, 50, undefined, undefined, undefined, undefined,
undefined // Should be undefined for empty key alias
);
});
});
tests/proxy_admin_ui_tests/ui_unit_tests/package-lock.json
Generated
+91 -12
View File
@@ -8,13 +8,14 @@
"name": "ui-unit-tests",
"version": "1.0.0",
"dependencies": {
"antd": "^5.0.0",
"@ant-design/icons": "^5.0.0",
"antd": "^5.12.5",
"react": "^18.2.0",
"react-dom": "^18.2.0"
},
"devDependencies": {
"@testing-library/jest-dom": "^6.0.0",
"@testing-library/react": "^14.0.0",
"@types/antd": "^1.0.0",
"@types/jest": "^29.5.0",
"@types/react": "^18.2.0",
"@types/react-dom": "^18.2.0",
@@ -25,6 +26,13 @@
"typescript": "^5.0.0"
}
},
"node_modules/@adobe/css-tools": {
"version": "4.4.4",
"resolved": "https://registry.npmjs.org/@adobe/css-tools/-/css-tools-4.4.4.tgz",
"integrity": "sha512-Elp+iwUx5rN5+Y8xLt5/GRoG20WGoDCQ/1Fb+1LiGtvwbDavuSk0jhD/eZdckHAuzcDzccnkv+rEjyWfRx18gg==",
"dev": true,
"license": "MIT"
},
"node_modules/@ampproject/remapping": {
"version": "2.3.0",
"resolved": "https://registry.npmjs.org/@ampproject/remapping/-/remapping-2.3.0.tgz",
@@ -1167,6 +1175,33 @@
"node": ">=14"
}
},
"node_modules/@testing-library/jest-dom": {
"version": "6.8.0",
"resolved": "https://registry.npmjs.org/@testing-library/jest-dom/-/jest-dom-6.8.0.tgz",
"integrity": "sha512-WgXcWzVM6idy5JaftTVC8Vs83NKRmGJz4Hqs4oyOuO2J4r/y79vvKZsb+CaGyCSEbUPI6OsewfPd0G1A0/TUZQ==",
"dev": true,
"license": "MIT",
"dependencies": {
"@adobe/css-tools": "^4.4.0",
"aria-query": "^5.0.0",
"css.escape": "^1.5.1",
"dom-accessibility-api": "^0.6.3",
"picocolors": "^1.1.1",
"redent": "^3.0.0"
},
"engines": {
"node": ">=14",
"npm": ">=6",
"yarn": ">=1"
}
},
"node_modules/@testing-library/jest-dom/node_modules/dom-accessibility-api": {
"version": "0.6.3",
"resolved": "https://registry.npmjs.org/dom-accessibility-api/-/dom-accessibility-api-0.6.3.tgz",
"integrity": "sha512-7ZgogeTnjuHbo+ct10G9Ffp0mif17idi0IyWNVA/wcwcm7NPOD/WEHVP3n7n3MhXqxoIYm8d6MuZohYWIZ4T3w==",
"dev": true,
"license": "MIT"
},
"node_modules/@testing-library/react": {
"version": "14.3.1",
"resolved": "https://registry.npmjs.org/@testing-library/react/-/react-14.3.1.tgz",
@@ -1194,16 +1229,6 @@
"node": ">= 10"
}
},
"node_modules/@types/antd": {
"version": "1.0.4",
"resolved": "https://registry.npmjs.org/@types/antd/-/antd-1.0.4.tgz",
"integrity": "sha512-gp4PGQckP1kNjj2H6juhjKIVwkpXwCIyIvOlwp2DC6geuhVpDHEEB5gwH4hJabVgBAFtrjBPJ58VIRV9VV9W2g==",
"deprecated": "This is a stub types definition. antd provides its own type definitions, so you do not need this installed.",
"dev": true,
"dependencies": {
"antd": "*"
}
},
"node_modules/@types/aria-query": {
"version": "5.0.4",
"resolved": "https://registry.npmjs.org/@types/aria-query/-/aria-query-5.0.4.tgz",
@@ -2081,6 +2106,13 @@
"node": ">= 8"
}
},
"node_modules/css.escape": {
"version": "1.5.1",
"resolved": "https://registry.npmjs.org/css.escape/-/css.escape-1.5.1.tgz",
"integrity": "sha512-YUifsXXuknHlUsmlgyY0PKzgPOr7/FjCePfHNt0jxm83wHZi44VDMQ7/fGNkjY3/jV1MC+1CmZbaHzugyeRtpg==",
"dev": true,
"license": "MIT"
},
"node_modules/cssom": {
"version": "0.5.0",
"resolved": "https://registry.npmjs.org/cssom/-/cssom-0.5.0.tgz",
@@ -2974,6 +3006,16 @@
"node": ">=0.8.19"
}
},
"node_modules/indent-string": {
"version": "4.0.0",
"resolved": "https://registry.npmjs.org/indent-string/-/indent-string-4.0.0.tgz",
"integrity": "sha512-EdDDZu4A2OyIK7Lr/2zG+w5jmbuk1DVBnEwREQvBzspBJkCEbRa8GxU1lghYcaGJCnRWibjDXlq779X1/y5xwg==",
"dev": true,
"license": "MIT",
"engines": {
"node": ">=8"
}
},
"node_modules/inflight": {
"version": "1.0.6",
"resolved": "https://registry.npmjs.org/inflight/-/inflight-1.0.6.tgz",
@@ -4559,6 +4601,16 @@
"node": ">=6"
}
},
"node_modules/min-indent": {
"version": "1.0.1",
"resolved": "https://registry.npmjs.org/min-indent/-/min-indent-1.0.1.tgz",
"integrity": "sha512-I9jwMn07Sy/IwOj3zVkVik2JTvgpaykDZEigL6Rx6N9LbMywwUSMtxET+7lVoDLLd3O3IXwJwvuuns8UB/HeAg==",
"dev": true,
"license": "MIT",
"engines": {
"node": ">=4"
}
},
"node_modules/minimatch": {
"version": "3.1.2",
"resolved": "https://registry.npmjs.org/minimatch/-/minimatch-3.1.2.tgz",
@@ -5552,6 +5604,20 @@
"integrity": "sha512-w2GsyukL62IJnlaff/nRegPQR94C/XXamvMWmSHRJ4y7Ts/4ocGRmTHvOs8PSE6pB3dWOrD/nueuU5sduBsQ4w==",
"dev": true
},
"node_modules/redent": {
"version": "3.0.0",
"resolved": "https://registry.npmjs.org/redent/-/redent-3.0.0.tgz",
"integrity": "sha512-6tDA8g98We0zd0GvVeMT9arEOnTw9qM03L9cJXaCjrip1OO764RDBLBfrB4cwzNGDj5OA5ioymC9GkizgWJDUg==",
"dev": true,
"license": "MIT",
"dependencies": {
"indent-string": "^4.0.0",
"strip-indent": "^3.0.0"
},
"engines": {
"node": ">=8"
}
},
"node_modules/regenerator-runtime": {
"version": "0.14.1",
"resolved": "https://registry.npmjs.org/regenerator-runtime/-/regenerator-runtime-0.14.1.tgz",
@@ -5965,6 +6031,19 @@
"node": ">=6"
}
},
"node_modules/strip-indent": {
"version": "3.0.0",
"resolved": "https://registry.npmjs.org/strip-indent/-/strip-indent-3.0.0.tgz",
"integrity": "sha512-laJTa3Jb+VQpaC6DseHhF7dXVqHTfJPCRDaEbid/drOhgitgYku/letMUqOXFoWV0zIIUbjpdH2t+tYj4bQMRQ==",
"dev": true,
"license": "MIT",
"dependencies": {
"min-indent": "^1.0.0"
},
"engines": {
"node": ">=8"
}
},
"node_modules/strip-json-comments": {
"version": "3.1.1",
"resolved": "https://registry.npmjs.org/strip-json-comments/-/strip-json-comments-3.1.1.tgz",
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
+7 -4
View File
@@ -473,13 +473,16 @@ async def test_logging_opentelemetry_context_propagation():
Test that OpenTelemtry context propagation works with async completion.
"""
import asyncio
import litellm
from litellm.integrations.custom_logger import CustomLogger
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
InMemorySpanExporter,
)
import litellm
from litellm.integrations.custom_logger import CustomLogger
provider = TracerProvider()
exporter = InMemorySpanExporter()
@@ -490,7 +493,7 @@ async def test_logging_opentelemetry_context_propagation():
class MockOpenTelemetryLogger(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
span = tracer.start_span(start_time=start_time.timestamp() * 1e9, name="async_log_success_event")
span.end(end_time=end_time)
span.end(end_time=end_time.timestamp() * 1e9)
mock_logging_obj = MockOpenTelemetryLogger()
tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_content_after_stop_reason.py
+343
View File
@@ -0,0 +1,343 @@
"""
Test for AnthropicStreamWrapper handling content blocks that exist after message_delta with stop_reason and usage.
This tests the scenario where a streaming response includes:
1. Initial content blocks
2. A message_delta chunk with stop_reason and usage
3. Additional content blocks after the stop_reason
The wrapper should properly handle this by:
- Holding the stop_reason chunk until usage is available
- Merging usage into the stop_reason chunk
- Properly managing content_block_stop/start events for subsequent content
"""
import os
import sys
from typing import List
import pytest
sys.path.insert(0, os.path.abspath("../../../../.."))
from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import (
AnthropicStreamWrapper,
)
from litellm.types.utils import Delta, ModelResponse, StreamingChoices, Usage
class MockCompletionStreamWithContentAfterStopReason:
"""Mock stream that simulates content blocks existing after message_delta with stop_reason and usage."""
def __init__(self):
self.responses = [
# Initial text content
ModelResponse(
stream=True,
choices=[
StreamingChoices(
delta=Delta(content="Hello"), index=0, finish_reason=None
)
],
),
ModelResponse(
stream=True,
choices=[
StreamingChoices(
delta=Delta(content=" world"), index=0, finish_reason=None
)
],
),
# Message delta with stop_reason AND usage (this is how it actually comes from the API)
ModelResponse(
stream=True,
choices=[
StreamingChoices(
delta=Delta(content=""), index=0, finish_reason="stop"
)
],
usage=Usage(prompt_tokens=230, completion_tokens=65, total_tokens=295),
),
# Additional content after the stop_reason - this simulates the scenario
# where there might be additional content blocks after the main response
ModelResponse(
stream=True,
choices=[
StreamingChoices(
delta=Delta(content=" Additional content"),
index=0,
finish_reason=None,
)
],
),
]
self.index = 0
def __iter__(self):
return self
def __next__(self):
if self.index >= len(self.responses):
raise StopIteration
response = self.responses[self.index]
self.index += 1
return response
def __aiter__(self):
return self
async def __anext__(self):
if self.index >= len(self.responses):
raise StopAsyncIteration
response = self.responses[self.index]
self.index += 1
return response
def test_anthropic_stream_wrapper_content_after_stop_reason():
"""Test that AnthropicStreamWrapper properly handles content blocks after message_delta with stop_reason."""
wrapper = AnthropicStreamWrapper(
completion_stream=MockCompletionStreamWithContentAfterStopReason(),
model="claude-3",
)
chunks = []
chunk_types = []
# Collect all chunks
for chunk in wrapper:
chunks.append(chunk)
chunk_types.append(chunk.get("type"))
# Verify the expected sequence of chunk types
expected_types = [
"message_start", # Initial message start
"content_block_start", # Start of first content block
"content_block_delta", # "Hello"
"content_block_delta", # " world"
"content_block_stop", # End of first content block due to stop_reason
"message_delta", # Stop reason with merged usage
"message_stop", # Final message stop
]
print(f"Actual chunk types: {chunk_types}")
print(f"Expected chunk types: {expected_types}")
# Verify we have the expected number of chunks
assert len(chunk_types) >= len(
expected_types
), f"Expected at least {len(expected_types)} chunks, got {len(chunk_types)}"
# Verify key chunk types are present
assert "message_start" in chunk_types
assert "content_block_start" in chunk_types
assert "content_block_delta" in chunk_types
assert "content_block_stop" in chunk_types
assert "message_delta" in chunk_types
assert "message_stop" in chunk_types
# Find the message_delta chunk with stop_reason
message_delta_chunk = None
for chunk in chunks:
if chunk.get("type") == "message_delta":
message_delta_chunk = chunk
break
assert message_delta_chunk is not None, "message_delta chunk not found"
# Verify that the message_delta chunk has both stop_reason and usage
delta = message_delta_chunk.get("delta", {})
usage = message_delta_chunk.get("usage", {})
assert (
delta.get("stop_reason") == "end_turn"
), f"Expected stop_reason 'end_turn', got {delta.get('stop_reason')}"
assert (
usage.get("input_tokens") == 230
), f"Expected input_tokens 230, got {usage.get('input_tokens')}"
assert (
usage.get("output_tokens") == 65
), f"Expected output_tokens 65, got {usage.get('output_tokens')}"
# Verify content_block_stop comes before message_delta
content_block_stop_index = None
message_delta_index = None
for i, chunk_type in enumerate(chunk_types):
if chunk_type == "content_block_stop" and content_block_stop_index is None:
content_block_stop_index = i
elif chunk_type == "message_delta":
message_delta_index = i
assert content_block_stop_index is not None, "content_block_stop not found"
assert message_delta_index is not None, "message_delta not found"
assert (
content_block_stop_index < message_delta_index
), "content_block_stop should come before message_delta"
@pytest.mark.asyncio
async def test_async_anthropic_stream_wrapper_content_after_stop_reason():
"""Test async version of AnthropicStreamWrapper handling content blocks after message_delta with stop_reason."""
wrapper = AnthropicStreamWrapper(
completion_stream=MockCompletionStreamWithContentAfterStopReason(),
model="claude-3",
)
chunks = []
chunk_types = []
# Collect all chunks asynchronously
async for chunk in wrapper:
chunks.append(chunk)
chunk_types.append(chunk.get("type"))
print(f"Async - Actual chunk types: {chunk_types}")
# Verify key chunk types are present
assert "message_start" in chunk_types
assert "content_block_start" in chunk_types
assert "content_block_delta" in chunk_types
assert "content_block_stop" in chunk_types
assert "message_delta" in chunk_types
assert "message_stop" in chunk_types
# Find the message_delta chunk with stop_reason
message_delta_chunk = None
for chunk in chunks:
if chunk.get("type") == "message_delta":
message_delta_chunk = chunk
break
assert message_delta_chunk is not None, "message_delta chunk not found"
# Verify that the message_delta chunk has both stop_reason and usage
delta = message_delta_chunk.get("delta", {})
usage = message_delta_chunk.get("usage", {})
assert (
delta.get("stop_reason") == "end_turn"
), f"Expected stop_reason 'end_turn', got {delta.get('stop_reason')}"
assert (
usage.get("input_tokens") == 230
), f"Expected input_tokens 230, got {usage.get('input_tokens')}"
assert (
usage.get("output_tokens") == 65
), f"Expected output_tokens 65, got {usage.get('output_tokens')}"
def test_usage_merging_behavior():
"""Test that usage information is properly merged with stop_reason chunk."""
wrapper = AnthropicStreamWrapper(
completion_stream=MockCompletionStreamWithContentAfterStopReason(),
model="claude-3",
)
# Process chunks and look specifically for the usage merging behavior
chunks = []
for chunk in wrapper:
chunks.append(chunk)
# If this is a message_delta with stop_reason, verify it has usage
if (
chunk.get("type") == "message_delta"
and chunk.get("delta", {}).get("stop_reason") is not None
):
usage = chunk.get("usage", {})
assert (
usage.get("input_tokens") is not None
), "Usage should be merged with stop_reason chunk"
assert (
usage.get("output_tokens") is not None
), "Usage should be merged with stop_reason chunk"
break
def test_sse_wrapper_with_content_after_stop_reason():
"""Test SSE wrapper formatting for the content after stop_reason scenario."""
wrapper = AnthropicStreamWrapper(
completion_stream=MockCompletionStreamWithContentAfterStopReason(),
model="claude-3",
)
# Get SSE formatted chunks
sse_chunks = []
for chunk in wrapper.anthropic_sse_wrapper():
sse_chunks.append(chunk)
if len(sse_chunks) >= 10: # Limit to avoid infinite loops in tests
break
# Verify all chunks are properly formatted as bytes
for chunk in sse_chunks:
assert isinstance(chunk, bytes), "SSE chunks should be bytes"
# Decode and verify SSE format
chunk_str = chunk.decode("utf-8")
lines = chunk_str.split("\n")
# Should have event and data lines
assert any(
line.startswith("event: ") for line in lines
), f"Missing event line in: {chunk_str}"
assert any(
line.startswith("data: ") for line in lines
), f"Missing data line in: {chunk_str}"
@pytest.mark.asyncio
async def test_async_sse_wrapper_with_content_after_stop_reason():
"""Test async SSE wrapper formatting for the content after stop_reason scenario."""
wrapper = AnthropicStreamWrapper(
completion_stream=MockCompletionStreamWithContentAfterStopReason(),
model="claude-3",
)
# Get SSE formatted chunks asynchronously
sse_chunks = []
async for chunk in wrapper.async_anthropic_sse_wrapper():
sse_chunks.append(chunk)
if len(sse_chunks) >= 10: # Limit to avoid infinite loops in tests
break
# Verify all chunks are properly formatted as bytes
for chunk in sse_chunks:
assert isinstance(chunk, bytes), "Async SSE chunks should be bytes"
# Decode and verify SSE format
chunk_str = chunk.decode("utf-8")
lines = chunk_str.split("\n")
# Should have event and data lines
assert any(
line.startswith("event: ") for line in lines
), f"Missing event line in: {chunk_str}"
assert any(
line.startswith("data: ") for line in lines
), f"Missing data line in: {chunk_str}"
if __name__ == "__main__":
# Run a quick test
test_anthropic_stream_wrapper_content_after_stop_reason()
print("✅ Sync test passed")
import asyncio
asyncio.run(test_async_anthropic_stream_wrapper_content_after_stop_reason())
print("✅ Async test passed")
test_usage_merging_behavior()
print("✅ Usage merging test passed")
test_sse_wrapper_with_content_after_stop_reason()
print("✅ SSE wrapper test passed")
asyncio.run(test_async_sse_wrapper_with_content_after_stop_reason())
print("✅ Async SSE wrapper test passed")
print("🎉 All tests passed!")
tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_completions.jsonl
+2
View File
@@ -0,0 +1,2 @@
{"recordId": "request-1", "modelInput": {"messages": [{"role": "user", "content": [{"type": "text", "text": "Hello world!"}]}], "max_tokens": 10, "system": [{"type": "text", "text": "You are a helpful assistant."}], "anthropic_version": "bedrock-2023-05-31"}}
{"recordId": "request-2", "modelInput": {"messages": [{"role": "user", "content": [{"type": "text", "text": "Hello world!"}]}], "max_tokens": 10, "system": [{"type": "text", "text": "You are an unhelpful assistant."}], "anthropic_version": "bedrock-2023-05-31"}}
tests/test_litellm/llms/bedrock/files/input_batch_completions.jsonl
+2
View File
@@ -0,0 +1,2 @@
{"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}}
tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py
+90
View File
@@ -0,0 +1,90 @@
"""
Test bedrock files transformation functionality
"""
import json
import os
from typing import Any, Dict, List
import pytest
from litellm.llms.bedrock.files.transformation import BedrockJsonlFilesTransformation
class TestBedrockFilesTransformation:
"""Test bedrock files transformation"""
def test_transform_openai_jsonl_content_to_bedrock_jsonl_content(self):
"""
Test transformation of OpenAI JSONL format to Bedrock batch format.
Validates that the transformation correctly converts OpenAI batch completion
format to Bedrock's expected batch format with proper recordId and modelInput structure.
"""
# Initialize the transformation class
transformation = BedrockJsonlFilesTransformation()
# Load input JSONL file
input_file_path = os.path.join(
os.path.dirname(__file__),
"input_batch_completions.jsonl"
)
# Read and parse the JSONL content
openai_jsonl_content = []
with open(input_file_path, 'r') as f:
for line in f:
if line.strip():
openai_jsonl_content.append(json.loads(line))
# Transform the content
bedrock_jsonl_content = transformation._transform_openai_jsonl_content_to_bedrock_jsonl_content(
openai_jsonl_content=openai_jsonl_content
)
# Print the transformation results for validation
print("\n=== INPUT (OpenAI format) ===")
for i, content in enumerate(openai_jsonl_content):
print(f"Record {i+1}:")
print(json.dumps(content, indent=2))
print()
print("\n=== OUTPUT (Bedrock format) ===")
for i, content in enumerate(bedrock_jsonl_content):
print(f"Record {i+1}:")
print(json.dumps(content, indent=2))
print()
# Basic validation
assert len(bedrock_jsonl_content) == len(openai_jsonl_content), "Should have same number of records"
# Check structure of transformed records
for i, record in enumerate(bedrock_jsonl_content):
assert "recordId" in record, f"Record {i+1} should have recordId"
assert "modelInput" in record, f"Record {i+1} should have modelInput"
# Check recordId matches custom_id from input
expected_custom_id = openai_jsonl_content[i].get("custom_id")
assert record["recordId"] == expected_custom_id, f"Record {i+1} recordId should match custom_id"
# Check modelInput has expected structure
model_input = record["modelInput"]
assert isinstance(model_input, dict), f"Record {i+1} modelInput should be a dictionary"
# For Anthropic models, should have anthropic_version and messages
if "anthropic.claude" in openai_jsonl_content[i]["body"]["model"]:
assert "anthropic_version" in model_input, f"Record {i+1} should have anthropic_version"
assert "messages" in model_input, f"Record {i+1} should have messages"
assert "max_tokens" in model_input, f"Record {i+1} should have max_tokens"
# Write expected output to file for reference
expected_output_path = os.path.join(
os.path.dirname(__file__),
"expected_bedrock_batch_completions.jsonl"
)
with open(expected_output_path, 'w') as f:
for record in bedrock_jsonl_content:
f.write(json.dumps(record) + '\n')
print(f"\n=== Expected output written to: {expected_output_path} ===")
tests/test_litellm/llms/lm_studio/test_lm_studio_chat_transformation.py
+30
View File
@@ -1,5 +1,6 @@
import os
import sys
from unittest.mock import patch
from pydantic import BaseModel
@@ -53,3 +54,32 @@ class TestLMStudioChatConfigResponseFormat:
assert mapped_schema["properties"] == schema["properties"]
opt_schema = optional_params["response_format"]["json_schema"]["schema"]
assert opt_schema["properties"] == schema["properties"]
def test_lm_studio_get_openai_compatible_provider_info():
"""Test provider info retrieval"""
config = LMStudioChatConfig()
# Test default behavior (no API key provided)
_, api_key = config._get_openai_compatible_provider_info(None, None)
assert api_key == "fake-api-key"
# Test explicit API key
_, api_key = config._get_openai_compatible_provider_info(None, "test-key")
assert api_key == "test-key"
def test_lm_studio_get_openai_compatible_provider_info_with_env():
"""Test provider info retrieval with environment variables."""
config = LMStudioChatConfig()
with patch.dict(
"os.environ",
{
"LM_STUDIO_API_BASE": "http://localhost:1234/v1",
"LM_STUDIO_API_KEY": "env_api_key",
},
):
api_base, api_key = config._get_openai_compatible_provider_info(None, None)
assert api_base == "http://localhost:1234/v1"
assert api_key == "env_api_key"
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py
+1035 -25
View File
File diff suppressed because it is too large Load Diff
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
+202 -43
View File
@@ -183,7 +183,9 @@ async def test_budget_reset_and_expires_at_first_of_month(monkeypatch):
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}"
assert (
response_date.day == 1
), f"Expected day 1, got {response_date.day} for {key}"
@pytest.mark.asyncio
@@ -507,7 +509,6 @@ def test_get_new_token_with_invalid_key():
assert "New key must start with 'sk-'" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_generate_service_account_requires_team_id():
with pytest.raises(HTTPException):
@@ -529,11 +530,12 @@ async def test_generate_service_account_works_with_team_id():
from unittest.mock import patch
# Mock the database and router dependencies from proxy_server
with patch('litellm.proxy.proxy_server.prisma_client') as mock_prisma, \
patch('litellm.proxy.proxy_server.llm_router') as mock_router, \
patch('litellm.proxy.proxy_server.premium_user', False), \
patch('litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn') as mock_generate_key:
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
"litellm.proxy.proxy_server.llm_router"
) as mock_router, patch("litellm.proxy.proxy_server.premium_user", False), patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn"
) as mock_generate_key:
# Configure mocks
mock_prisma.return_value = AsyncMock()
mock_router.return_value = None
@@ -542,9 +544,9 @@ async def test_generate_service_account_works_with_team_id():
"key": "sk-test-key",
"expires": None,
"user_id": "test-user",
"team_id": "IJ"
"team_id": "IJ",
}
# This should not raise an exception since team_id is provided
await _common_key_generation_helper(
data=GenerateKeyRequest(
@@ -559,7 +561,6 @@ async def test_generate_service_account_works_with_team_id():
)
@pytest.mark.asyncio
async def test_update_service_account_requires_team_id():
data = UpdateKeyRequest(key="sk-1", metadata={"service_account_id": "sa"})
@@ -571,7 +572,9 @@ async def test_update_service_account_requires_team_id():
@pytest.mark.asyncio
async def test_update_service_account_works_with_team_id():
data = UpdateKeyRequest(key="sk-1", metadata={"service_account_id": "sa"}, team_id="IJ")
data = UpdateKeyRequest(
key="sk-1", metadata={"service_account_id": "sa"}, team_id="IJ"
)
existing_key = LiteLLM_VerificationToken(token="hashed")
await prepare_key_update_data(data=data, existing_key_row=existing_key)
@@ -580,22 +583,22 @@ async def test_update_service_account_works_with_team_id():
@pytest.mark.asyncio
async def test_validate_team_id_used_in_service_account_request_requires_team_id():
"""
Test that validate_team_id_used_in_service_account_request raises HTTPException
Test that validate_team_id_used_in_service_account_request raises HTTPException
when team_id is None for service account key generation.
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
validate_team_id_used_in_service_account_request,
)
mock_prisma_client = AsyncMock()
# Test that HTTPException is raised when team_id is None
with pytest.raises(HTTPException) as exc_info:
await validate_team_id_used_in_service_account_request(
team_id=None,
prisma_client=mock_prisma_client,
)
assert exc_info.value.status_code == 400
assert "team_id is required for service account keys" in str(exc_info.value.detail)
@@ -603,7 +606,7 @@ async def test_validate_team_id_used_in_service_account_request_requires_team_id
@pytest.mark.asyncio
async def test_validate_team_id_used_in_service_account_request_requires_prisma_client():
"""
Test that validate_team_id_used_in_service_account_request raises HTTPException
Test that validate_team_id_used_in_service_account_request raises HTTPException
when prisma_client is None for service account key generation.
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
@@ -616,78 +619,76 @@ async def test_validate_team_id_used_in_service_account_request_requires_prisma_
team_id="test-team-id",
prisma_client=None,
)
assert exc_info.value.status_code == 400
assert "prisma_client is required for service account keys" in str(exc_info.value.detail)
assert "prisma_client is required for service account keys" in str(
exc_info.value.detail
)
@pytest.mark.asyncio
async def test_validate_team_id_used_in_service_account_request_checks_team_exists():
"""
Test that validate_team_id_used_in_service_account_request validates that
Test that validate_team_id_used_in_service_account_request validates that
the team_id exists in the database for service account key generation.
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
validate_team_id_used_in_service_account_request,
)
mock_prisma_client = AsyncMock()
# Mock the database query to return None (team doesn't exist)
mock_find_unique = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique
# Test that HTTPException is raised when team doesn't exist in DB
with pytest.raises(HTTPException) as exc_info:
await validate_team_id_used_in_service_account_request(
team_id="non-existent-team-id",
prisma_client=mock_prisma_client,
)
assert exc_info.value.status_code == 400
assert "team_id does not exist in the database" in str(exc_info.value.detail)
# Verify the database was queried with the correct parameters
mock_find_unique.assert_called_once_with(
where={"team_id": "non-existent-team-id"}
)
mock_find_unique.assert_called_once_with(where={"team_id": "non-existent-team-id"})
@pytest.mark.asyncio
async def test_validate_team_id_used_in_service_account_request_success():
"""
Test that validate_team_id_used_in_service_account_request returns True
Test that validate_team_id_used_in_service_account_request returns True
when team_id exists in the database for service account key generation.
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
validate_team_id_used_in_service_account_request,
)
mock_prisma_client = AsyncMock()
# Mock the database query to return a team object (team exists)
mock_team = {"team_id": "existing-team-id", "team_name": "Test Team"}
mock_find_unique = AsyncMock(return_value=mock_team)
mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique
# Test that function returns True when team exists
result = await validate_team_id_used_in_service_account_request(
team_id="existing-team-id",
prisma_client=mock_prisma_client,
)
assert result is True
# Verify the database was queried with the correct parameters
mock_find_unique.assert_called_once_with(
where={"team_id": "existing-team-id"}
)
mock_find_unique.assert_called_once_with(where={"team_id": "existing-team-id"})
@pytest.mark.asyncio
async def test_generate_service_account_key_endpoint_validation():
"""
Test that the /key/service-account/generate endpoint properly validates
Test that the /key/service-account/generate endpoint properly validates
team_id requirement and team existence in database.
"""
from unittest.mock import patch
@@ -705,16 +706,16 @@ async def test_generate_service_account_key_endpoint_validation():
),
litellm_changed_by=None,
)
assert exc_info.value.status_code == 400
assert "team_id is required for service account keys" in str(exc_info.value.detail)
# Test case 2: Team doesn't exist in database
with patch('litellm.proxy.proxy_server.prisma_client') as mock_prisma:
# Test case 2: Team doesn't exist in database
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
# Mock team not found
mock_find_unique = AsyncMock(return_value=None)
mock_prisma.db.litellm_teamtable.find_unique = mock_find_unique
with pytest.raises(HTTPException) as exc_info:
await generate_service_account_key_fn(
data=GenerateKeyRequest(team_id="non-existent-team"),
@@ -723,7 +724,165 @@ async def test_generate_service_account_key_endpoint_validation():
),
litellm_changed_by=None,
)
assert exc_info.value.status_code == 400
assert "team_id does not exist in the database" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_unblock_key_supports_both_sk_and_hashed_tokens(monkeypatch):
"""
Test that the unblock_key endpoint correctly handles both sk- prefixed tokens
and hashed tokens by properly converting sk- tokens to hashed format before
database operations.
"""
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy._types import BlockKeyRequest
from litellm.proxy.management_endpoints.key_management_endpoints import unblock_key
# Mock dependencies
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
mock_proxy_logging_obj = MagicMock()
# Use a proper 64-character hex hash for testing
test_hashed_token = (
"a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
)
# Mock the key record that will be returned from database
mock_key_record = MagicMock()
mock_key_record.token = test_hashed_token
mock_key_record.blocked = False
mock_key_record.model_dump_json.return_value = (
f'{{"token": "{test_hashed_token}", "blocked": false}}'
)
# Mock database operations
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=mock_key_record
)
mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
return_value=mock_key_record
)
# Mock get_key_object and _cache_key_object functions
mock_key_object = MagicMock()
mock_key_object.blocked = True # Initially blocked
# Mock hash_token function
def mock_hash_token(token):
if token == "sk-test123456789":
return test_hashed_token
return token
# Apply monkeypatch
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
)
monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)
monkeypatch.setattr(
"litellm.store_audit_logs", False
) # Disable audit logs for simpler test
# Mock get_key_object and _cache_key_object
async def mock_get_key_object(**kwargs):
return mock_key_object
async def mock_cache_key_object(**kwargs):
pass
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_key_object",
mock_get_key_object,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints._cache_key_object",
mock_cache_key_object,
)
# Create mock request and user auth
mock_request = MagicMock()
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
)
# Test Case 1: Using sk- prefixed token
sk_token_request = BlockKeyRequest(key="sk-test123456789")
result = await unblock_key(
data=sk_token_request,
http_request=mock_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
# Verify that the database update was called with hashed token
mock_prisma_client.db.litellm_verificationtoken.update.assert_called_with(
where={"token": test_hashed_token}, data={"blocked": False}
)
assert result == mock_key_record
assert mock_key_object.blocked == False # Should be updated to unblocked
# Reset mocks for second test
mock_prisma_client.db.litellm_verificationtoken.update.reset_mock()
mock_key_object.blocked = True # Reset to blocked state
# Test Case 2: Using already hashed token
hashed_token_request = BlockKeyRequest(key=test_hashed_token)
result = await unblock_key(
data=hashed_token_request,
http_request=mock_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
# Verify that the database update was called with the same hashed token
mock_prisma_client.db.litellm_verificationtoken.update.assert_called_with(
where={"token": test_hashed_token}, data={"blocked": False}
)
assert result == mock_key_record
assert mock_key_object.blocked == False # Should be updated to unblocked
@pytest.mark.asyncio
async def test_unblock_key_invalid_key_format(monkeypatch):
"""
Test that unblock_key properly validates key format and raises appropriate errors
for invalid keys.
"""
from litellm.proxy._types import BlockKeyRequest
from litellm.proxy.management_endpoints.key_management_endpoints import unblock_key
from litellm.proxy.utils import ProxyException
# Mock prisma_client to avoid DB connection error
mock_prisma_client = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
# Mock request and user auth
mock_request = MagicMock()
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
)
# Test with invalid key format
invalid_key_request = BlockKeyRequest(key="invalid-key-format")
with pytest.raises(ProxyException) as exc_info:
await unblock_key(
data=invalid_key_request,
http_request=mock_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
assert exc_info.value.code == "400"
assert "Invalid key format" in str(exc_info.value.message)
ui/litellm-dashboard/src/components/networking.tsx
+3 -1
View File
@@ -2633,7 +2633,8 @@ export const uiSpendLogsCall = async (
user_id?: string,
end_user?: string,
status_filter?: string,
model?: string
model?: string,
keyAlias?: string
) => {
try {
// Construct base URL
@@ -2652,6 +2653,7 @@ export const uiSpendLogsCall = async (
if (end_user) queryParams.append("end_user", end_user);
if (status_filter) queryParams.append("status_filter", status_filter);
if (model) queryParams.append("model", model);
if (keyAlias) queryParams.append("key_alias", keyAlias);
// Append query parameters to URL if any exist
const queryString = queryParams.toString();
if (queryString) {
ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx
+17 -20
View File
@@ -60,6 +60,7 @@ export function useLogFilterLogic({
const performSearch = useCallback(async (filters: LogFilterState, page = 1) => {
if (!accessToken) return;
console.log("Filters being sent to API:", filters);
const currentTimestamp = Date.now();
lastSearchTimestamp.current = currentTimestamp;
@@ -81,7 +82,8 @@ export function useLogFilterLogic({
filters[FILTER_KEYS.USER_ID] || undefined,
filters[FILTER_KEYS.END_USER] || undefined,
filters[FILTER_KEYS.STATUS] || undefined,
filters[FILTER_KEYS.MODEL] || undefined
filters[FILTER_KEYS.MODEL] || undefined,
filters[FILTER_KEYS.KEY_ALIAS] || undefined
);
if (currentTimestamp === lastSearchTimestamp.current && response.data) {
@@ -123,6 +125,19 @@ export function useLogFilterLogic({
});
return;
}
// Only do client-side filtering if no backend filters are active
const hasBackendFilters =
filters[FILTER_KEYS.KEY_ALIAS] ||
filters[FILTER_KEYS.KEY_HASH] ||
filters[FILTER_KEYS.REQUEST_ID] ||
filters[FILTER_KEYS.USER_ID] ||
filters[FILTER_KEYS.END_USER];
if (hasBackendFilters) {
// Backend is handling filtering, don't override the results
return;
}
let filteredData = [...logs.data];
@@ -148,7 +163,7 @@ export function useLogFilterLogic({
log => log.model === filters[FILTER_KEYS.MODEL]
);
}
if (filters[FILTER_KEYS.KEY_HASH]) {
filteredData = filteredData.filter(
log => log.api_key === filters[FILTER_KEYS.KEY_HASH]
@@ -161,24 +176,6 @@ export function useLogFilterLogic({
);
}
// Add key alias filtering
if (filters[FILTER_KEYS.KEY_ALIAS]) {
// We need to fetch the key info to get the key hash for the selected alias
try {
// Get the key hash for the selected alias
const selectedKey = filters[FILTER_KEYS.KEY_ALIAS]
if (selectedKey) {
// Filter logs by the key hash
filteredData = filteredData.filter(
log => log.metadata?.user_api_key_alias === selectedKey
);
}
} catch (error) {
console.error("Error fetching key info for alias:", error);
}
}
const newFilteredLogs: PaginatedResponse = {
data: filteredData,
total: logs.total,