merge: resolve conflict with main in tool_management_endpoints.py

Remove unused ToolOutputPolicy import.

Made-with: Cursor
This commit is contained in:
Ishaan Jaffer 2026-03-04 12:29:09 -08:00
commit b4cd5879fb
129 changed files with 12806 additions and 905 deletions

View file

@ -3689,6 +3689,114 @@ jobs:
- store_test_results:
path: test-results
proxy_e2e_azure_batches_tests:
machine:
image: ubuntu-2204:2023.10.1
resource_class: xlarge
working_directory: ~/project
steps:
- checkout
- setup_google_dns
- run:
name: Install Docker CLI
command: |
curl -fsSL https://get.docker.com | sh
sudo usermod -aG docker $USER
docker version
- run:
name: Install Python 3.12
command: |
curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output miniconda.sh
bash miniconda.sh -b -p $HOME/miniconda
export PATH="$HOME/miniconda/bin:$PATH"
conda init bash
source ~/.bashrc
conda create -n myenv python=3.12 -y
conda activate myenv
python --version
- run:
name: Install Poetry
command: |
export PATH="$HOME/miniconda/bin:$PATH"
source $HOME/miniconda/etc/profile.d/conda.sh
conda activate myenv
pip install poetry
- run:
name: Install dockerize
command: |
wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz
sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz
rm dockerize-linux-amd64-v0.6.1.tar.gz
- run:
name: Start PostgreSQL Database
command: |
docker run -d \
--name postgres-db \
-e POSTGRES_USER=llmproxy \
-e POSTGRES_PASSWORD=dbpassword9090 \
-e POSTGRES_DB=litellm \
-p 5432:5432 \
postgres:15
- run:
name: Wait for PostgreSQL to be ready
command: dockerize -wait tcp://localhost:5432 -timeout 1m
- run:
name: Install system dependencies
command: |
sudo apt-get update -y
sudo apt-get install -y libpq-dev
- run:
name: Install Dependencies
command: |
export PATH="$HOME/miniconda/bin:$PATH"
source $HOME/miniconda/etc/profile.d/conda.sh
conda activate myenv
poetry config virtualenvs.in-project true
poetry install --with dev,proxy-dev --extras "proxy"
poetry run pip install psycopg2-binary uvicorn fastapi httpx tenacity
- run:
name: Setup litellm-enterprise
command: |
export PATH="$HOME/miniconda/bin:$PATH"
source $HOME/miniconda/etc/profile.d/conda.sh
conda activate myenv
poetry run pip install --force-reinstall --no-deps -e enterprise/
- run:
name: Generate Prisma client
command: |
export PATH="$HOME/miniconda/bin:$PATH"
source $HOME/miniconda/etc/profile.d/conda.sh
conda activate myenv
poetry run prisma generate --schema litellm/proxy/schema.prisma
- run:
name: Run Prisma migrations
command: |
export PATH="$HOME/miniconda/bin:$PATH"
source $HOME/miniconda/etc/profile.d/conda.sh
conda activate myenv
export DATABASE_URL=postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
cd litellm/proxy
poetry run prisma migrate deploy --schema schema.prisma
cd ../..
- run:
name: Run Azure Batch E2E Tests
command: |
export PATH="$HOME/miniconda/bin:$PATH"
source $HOME/miniconda/etc/profile.d/conda.sh
conda activate myenv
export DATABASE_URL=postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
export USE_LOCAL_LITELLM=true
export USE_MOCK_MODELS=true
export USE_STATE_TRACKER=true
export LITELLM_LOG=DEBUG
poetry run pytest tests/proxy_e2e_azure_batches_tests/test_proxy_e2e_azure_batches.py \
-vv -s -k "test_e2e_managed_batch" \
--tb=short \
--maxfail=3 \
--durations=10 \
--junitxml=test-results/junit.xml
no_output_timeout: 30m
upload-coverage:
docker:
- image: cimg/python:3.9
@ -4458,6 +4566,12 @@ workflows:
only:
- main
- /litellm_.*/
- proxy_e2e_azure_batches_tests:
filters:
branches:
only:
- main
- /litellm_.*/
- llm_translation_testing:
filters:
branches:

View file

@ -32,7 +32,6 @@ jobs:
run: |
poetry lock
poetry install --with dev
poetry run pip install openai==1.100.1
- name: Run Black formatting
run: |

View file

@ -38,7 +38,7 @@ jobs:
poetry run pip install "google-genai==1.22.0"
poetry run pip install "google-cloud-aiplatform>=1.38"
poetry run pip install "fastapi-offline==1.7.3"
poetry run pip install "python-multipart==0.0.22"
poetry run pip install "python-multipart>=0.0.20"
poetry run pip install "openapi-core"
- name: Setup litellm-enterprise as local package
run: |

View file

@ -0,0 +1,90 @@
name: Proxy E2E Azure Batches Tests
on:
pull_request:
branches: [main]
workflow_dispatch:
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
proxy_e2e_azure_batches_tests:
runs-on: ubuntu-latest
timeout-minutes: 30
services:
postgres:
image: postgres:15
env:
POSTGRES_USER: llmproxy
POSTGRES_PASSWORD: dbpassword9090
POSTGRES_DB: litellm
ports:
- 5432:5432
options: >-
--health-cmd pg_isready
--health-interval 10s
--health-timeout 5s
--health-retries 5
steps:
- uses: actions/checkout@v4
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: "3.12"
- name: Install Poetry
uses: snok/install-poetry@v1
- name: Cache Poetry dependencies
uses: actions/cache@v4
with:
path: |
~/.cache/pypoetry
~/.cache/pip
.venv
key: ${{ runner.os }}-poetry-e2e-batches-${{ hashFiles('poetry.lock') }}
restore-keys: |
${{ runner.os }}-poetry-e2e-batches-
${{ runner.os }}-poetry-
- name: Install dependencies
run: |
poetry config virtualenvs.in-project true
poetry install --with dev,proxy-dev --extras "proxy"
poetry run pip install psycopg2-binary uvicorn fastapi httpx tenacity
- name: Setup litellm-enterprise
run: |
poetry run pip install --force-reinstall --no-deps -e enterprise/
- name: Generate Prisma client
run: |
poetry run prisma generate --schema litellm/proxy/schema.prisma
- name: Run Prisma migrations
env:
DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
run: |
cd litellm/proxy
poetry run prisma migrate deploy --schema schema.prisma
cd ../..
- name: Run Azure Batch E2E Tests
env:
DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
USE_LOCAL_LITELLM: "true"
USE_MOCK_MODELS: "true"
USE_STATE_TRACKER: "true"
LITELLM_LOG: DEBUG
run: |
poetry run pytest tests/proxy_e2e_azure_batches_tests/test_proxy_e2e_azure_batches.py \
-vv -s -k "test_e2e_managed_batch" \
--tb=short \
--maxfail=3 \
--durations=10

View file

@ -16,7 +16,7 @@ LiteLLM provides image editing functionality that maps to OpenAI's `/images/edit
| Supported operations | Create image edits | Single and multiple images supported |
| Supported LiteLLM SDK Versions | 1.63.8+ | Gemini support requires 1.79.3+ |
| Supported LiteLLM Proxy Versions | 1.71.1+ | Gemini support requires 1.79.3+ |
| Supported LLM providers | **OpenAI**, **Gemini (Google AI Studio)**, **Vertex AI**, **Stability AI**, **AWS Bedrock (Stability)** | Gemini supports the new `gemini-2.5-flash-image` family. Vertex AI supports both Gemini and Imagen models. Stability AI and Bedrock Stability support various image editing operations. |
| Supported LLM providers | **OpenAI**, **Gemini (Google AI Studio)**, **Vertex AI**, **OpenRouter**, **Stability AI**, **AWS Bedrock (Stability)** | Gemini supports the new `gemini-2.5-flash-image` family. Vertex AI supports both Gemini and Imagen models. OpenRouter routes image edits through chat completions. Stability AI and Bedrock Stability support various image editing operations. |
#### ⚡See all supported models and providers at [models.litellm.ai](https://models.litellm.ai/)
@ -244,6 +244,47 @@ response = litellm.image_edit(
print(response)
```
</TabItem>
<TabItem value="openrouter" label="OpenRouter">
#### Basic Image Edit
```python showLineNumbers title="OpenRouter Image Edit"
import os
from litellm import image_edit
os.environ["OPENROUTER_API_KEY"] = "your-api-key"
response = image_edit(
model="openrouter/google/gemini-2.5-flash-image",
image=open("original_image.png", "rb"),
prompt="Add aurora borealis to the night sky",
)
print(response)
```
#### Multiple Images Edit
```python showLineNumbers title="OpenRouter Multiple Images Edit"
import os
from litellm import image_edit
os.environ["OPENROUTER_API_KEY"] = "your-api-key"
response = image_edit(
model="openrouter/google/gemini-2.5-flash-image",
image=[
open("scene.png", "rb"),
open("style_reference.png", "rb"),
],
prompt="Blend the reference style into the scene",
size="1536x1024", # mapped to aspect_ratio 3:2
quality="high", # mapped to image_size 4K
)
print(response)
```
</TabItem>
</Tabs>
@ -398,6 +439,34 @@ curl -X POST "http://0.0.0.0:4000/v1/images/edits" \
-F "size=1024x1024"
```
</TabItem>
<TabItem value="openrouter" label="OpenRouter">
1. Add the OpenRouter image edit model to your `config.yaml`:
```yaml showLineNumbers title="OpenRouter Proxy Configuration"
model_list:
- model_name: openrouter-image-edit
litellm_params:
model: openrouter/google/gemini-2.5-flash-image
api_key: os.environ/OPENROUTER_API_KEY
```
2. Start the LiteLLM proxy server:
```bash showLineNumbers title="Start LiteLLM Proxy Server"
litellm --config /path/to/config.yaml
```
3. Make an image edit request:
```bash showLineNumbers title="OpenRouter Proxy Image Edit"
curl -X POST "http://0.0.0.0:4000/v1/images/edits" \
-H "Authorization: Bearer <YOUR-LITELLM-KEY>" \
-F "model=openrouter-image-edit" \
-F "image=@original_image.png" \
-F "prompt=Make the sky a vibrant purple sunset" \
-F "size=1024x1024"
```
</TabItem>
</Tabs>

View file

@ -191,6 +191,7 @@ os.environ["OPENAI_BASE_URL"] = "https://your_host/v1" # OPTIONAL
| gpt-5.2 | `response = completion(model="gpt-5.2", messages=messages)` |
| gpt-5.2-2025-12-11 | `response = completion(model="gpt-5.2-2025-12-11", messages=messages)` |
| gpt-5.2-chat-latest | `response = completion(model="gpt-5.2-chat-latest", messages=messages)` |
| gpt-5.3-chat-latest | `response = completion(model="gpt-5.3-chat-latest", messages=messages)` |
| gpt-5.2-pro | `response = completion(model="gpt-5.2-pro", messages=messages)` |
| gpt-5.2-pro-2025-12-11 | `response = completion(model="gpt-5.2-pro-2025-12-11", messages=messages)` |
| gpt-5.1 | `response = completion(model="gpt-5.1", messages=messages)` |

View file

@ -210,3 +210,90 @@ response = image_generation(
# Cost is available in the response metadata
print(f"Request cost: ${response._hidden_params['additional_headers']['llm_provider-x-litellm-response-cost']}")
```
## Image Edit
OpenRouter supports image editing through select models like Google Gemini image models. LiteLLM routes image edit requests to OpenRouter's chat completions endpoint with the source image sent as a base64 data URL and `modalities: ["image", "text"]`.
### Supported Models
| Model | Description |
|-------|-------------|
| `openrouter/google/gemini-2.5-flash-image` | Gemini 2.5 Flash with image editing |
See all available image models on [OpenRouter's model list](https://openrouter.ai/models?modality=image).
### Supported Parameters
| Parameter | OpenRouter Mapping | Notes |
|-----------|--------------------|-------|
| `size` | `image_config.aspect_ratio` | `1024x1024``1:1`, `1536x1024``3:2`, `1024x1536``2:3`, `1792x1024``16:9`, `1024x1792``9:16` |
| `quality` | `image_config.image_size` | `low`/`standard``1K`, `medium``2K`, `high`/`hd``4K` |
| `n` | `n` | Number of images |
:::note
`quality=high` (4K) is only supported by `google/gemini-3-pro-image-preview` and `google/gemini-3.1-flash-image-preview`. The `google/gemini-2.5-flash-image` model supports up to `medium` (2K).
:::
### Usage
```python
from litellm import image_edit
import os
os.environ["OPENROUTER_API_KEY"] = "your-api-key"
# Basic image edit
response = image_edit(
model="openrouter/google/gemini-2.5-flash-image",
image=open("original_image.png", "rb"),
prompt="Make the sky a vibrant purple sunset",
)
print(response)
```
### Advanced Usage with Parameters
```python
from litellm import image_edit
import os
os.environ["OPENROUTER_API_KEY"] = "your-api-key"
# Edit with size and quality parameters
response = image_edit(
model="openrouter/google/gemini-2.5-flash-image",
image=open("photo.png", "rb"),
prompt="Add northern lights to the sky",
size="1536x1024", # Maps to aspect_ratio 3:2
quality="high", # Maps to image_size 4K
)
# Access the edited image
image_data = response.data[0]
if image_data.b64_json:
import base64
with open("edited.png", "wb") as f:
f.write(base64.b64decode(image_data.b64_json))
```
### Multiple Images Edit
```python
from litellm import image_edit
import os
os.environ["OPENROUTER_API_KEY"] = "your-api-key"
response = image_edit(
model="openrouter/google/gemini-2.5-flash-image",
image=[
open("scene.png", "rb"),
open("style_reference.png", "rb"),
],
prompt="Blend the reference style into the scene",
)
print(response)
```

View file

@ -100,6 +100,19 @@ AzureHarmCategories:
n/a
## Important Notes
### Azure Content Safety Character Limit
Both Azure Prompt Shield and Azure Text Moderation have a **10,000 character limit** per request. When text exceeds this limit:
- LiteLLM automatically splits the text into chunks at word boundaries (no words are broken)
- Each chunk is sent separately to the Azure Content Safety API for analysis
- If any chunk is flagged (attack detected or severity threshold exceeded), the entire request is blocked
- If all chunks are safe, the request is allowed to proceed
This applies to both `pre_call` and `post_call` hooks and ensures that long prompts are properly analyzed without breaking words or losing context.
## Further Reading

View file

@ -358,13 +358,13 @@ response = client.chat.completions.create(
}
],
extra_body={
"guardrails": [
"guardrails": {
"aporia-pre-guard": {
"extra_body": {
"success_threshold": 0.9
}
}
]
}
}
)
@ -387,13 +387,13 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
"content": "what llm are you"
}
],
"guardrails": [
"guardrails": {
"aporia-pre-guard": {
"extra_body": {
"success_threshold": 0.9
}
}
]
}
}'
```
</TabItem>
@ -451,7 +451,6 @@ curl -X POST 'http://0.0.0.0:4000/key/generate' \
-H 'Content-Type: application/json' \
-d '{
"guardrails": ["aporia-pre-guard", "aporia-post-guard"]
}
}'
```
@ -465,7 +464,6 @@ curl --location 'http://0.0.0.0:4000/key/update' \
--data '{
"key": "sk-jNm1Zar7XfNdZXp49Z1kSQ",
"guardrails": ["aporia-pre-guard", "aporia-post-guard"]
}
}'
```

View file

@ -14,6 +14,7 @@ Requests to /chat/completions may be bridged here automatically when the provide
| Logging | ✅ | Works across all integrations |
| End-user Tracking | ✅ | |
| Streaming | ✅ | |
| WebSocket Mode | ✅ | Lower-latency persistent connections for all providers |
| Image Generation Streaming | ✅ | Progressive image generation with partial images (1-3) |
| Fallbacks | ✅ | Works between supported models |
| Loadbalancing | ✅ | Works between supported models |
@ -810,6 +811,245 @@ for event in response:
</TabItem>
</Tabs>
## WebSocket Mode
The Responses API supports **WebSocket mode** for lower-latency, persistent connections ideal for agentic workflows. WebSocket mode works with **all LiteLLM providers**, not just those with native WebSocket support.
### Architecture
LiteLLM provides two WebSocket modes:
1. **Native WebSocket**: Direct `wss://` connection to providers that support it (OpenAI, Azure)
2. **Managed WebSocket**: HTTP streaming over WebSocket for all other providers (Anthropic, Gemini, Bedrock, etc.)
The system automatically selects the appropriate mode based on provider capabilities.
### Usage
<Tabs>
<TabItem value="python" label="Python (websocket-client)">
```python showLineNumbers title="WebSocket with Python"
import json
from websocket import create_connection # pip install websocket-client
# Connect to LiteLLM proxy WebSocket endpoint
ws = create_connection(
"ws://localhost:4000/v1/responses?model=gemini-2.5-flash",
header=["Authorization: Bearer sk-1234"]
)
try:
# Send initial message
ws.send(json.dumps({
"type": "response.create",
"model": "gemini-2.5-flash",
"store": True,
"input": [{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "My favorite color is blue."}]
}]
}))
# Collect response events
response_id = None
while True:
event = json.loads(ws.recv())
print(f"Event: {event['type']}")
if event["type"] == "response.completed":
response_id = event["response"]["id"]
break
elif event["type"] == "response.output_text.delta":
print(f"Text: {event.get('delta', '')}", end="", flush=True)
print(f"\nResponse ID: {response_id}")
# Send follow-up with previous_response_id for multi-turn
ws.send(json.dumps({
"type": "response.create",
"model": "gemini-2.5-flash",
"previous_response_id": response_id,
"input": [{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "What is my favorite color?"}]
}]
}))
# Collect follow-up response
while True:
event = json.loads(ws.recv())
if event["type"] == "response.completed":
break
elif event["type"] == "response.output_text.delta":
print(event.get("delta", ""), end="", flush=True)
finally:
ws.close()
```
</TabItem>
<TabItem value="javascript" label="JavaScript (ws)">
```javascript showLineNumbers title="WebSocket with JavaScript"
const WebSocket = require('ws'); // npm install ws
const ws = new WebSocket(
'ws://localhost:4000/v1/responses?model=gemini-2.5-flash',
{
headers: {
'Authorization': 'Bearer sk-1234'
}
}
);
ws.on('open', () => {
// Send initial message
ws.send(JSON.stringify({
type: 'response.create',
model: 'gemini-2.5-flash',
store: true,
input: [{
type: 'message',
role: 'user',
content: [{ type: 'input_text', text: 'My favorite color is blue.' }]
}]
}));
});
let responseId = null;
ws.on('message', (data) => {
const event = JSON.parse(data.toString());
console.log(`Event: ${event.type}`);
if (event.type === 'response.completed') {
responseId = event.response.id;
console.log(`Response ID: ${responseId}`);
// Send follow-up
ws.send(JSON.stringify({
type: 'response.create',
model: 'gemini-2.5-flash',
previous_response_id: responseId,
input: [{
type: 'message',
role: 'user',
content: [{ type: 'input_text', text: 'What is my favorite color?' }]
}]
}));
} else if (event.type === 'response.output_text.delta') {
process.stdout.write(event.delta || '');
}
});
ws.on('error', (error) => {
console.error('WebSocket error:', error);
});
```
</TabItem>
<TabItem value="curl" label="curl (websocat)">
```bash showLineNumbers title="WebSocket with websocat"
# Install websocat: brew install websocat (macOS) or cargo install websocat
# Connect to WebSocket endpoint
websocat "ws://localhost:4000/v1/responses?model=gemini-2.5-flash" \
-H="Authorization: Bearer sk-1234"
# Then send JSON events (paste and press Enter):
{"type":"response.create","model":"gemini-2.5-flash","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"Hello!"}]}]}
# You'll receive streaming events back:
# {"type":"response.created",...}
# {"type":"response.in_progress",...}
# {"type":"response.output_text.delta","delta":"Hello",...}
# {"type":"response.completed",...}
```
</TabItem>
</Tabs>
### Event Types
WebSocket connections receive Server-Sent Events (SSE) formatted as JSON:
| Event Type | Description |
|------------|-------------|
| `response.created` | Response generation started |
| `response.in_progress` | Response is being generated |
| `response.output_item.added` | New output item (message, tool call, etc.) added |
| `response.output_text.delta` | Incremental text chunk |
| `response.output_text.done` | Text output completed |
| `response.content_part.done` | Content part completed |
| `response.output_item.done` | Output item completed |
| `response.completed` | Full response completed successfully |
| `response.failed` | Response generation failed |
| `response.incomplete` | Response incomplete (e.g., max tokens reached) |
| `error` | Error occurred |
### Multi-Turn Conversations
Use `previous_response_id` to maintain conversation context across multiple WebSocket messages:
```python showLineNumbers title="Multi-turn WebSocket Conversation"
# Turn 1
ws.send(json.dumps({
"type": "response.create",
"model": "gemini-2.5-flash",
"store": True, # Required for multi-turn
"input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Hello"}]}]
}))
# ... collect events and get response_id from response.completed event ...
# Turn 2 - reference previous response
ws.send(json.dumps({
"type": "response.create",
"model": "gemini-2.5-flash",
"previous_response_id": response_id, # Links to previous turn
"input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Continue"}]}]
}))
```
### Provider Support
| Provider | WebSocket Mode | Notes |
|----------|----------------|-------|
| OpenAI | Native | Direct `wss://` connection to OpenAI |
| Azure OpenAI | Native | Direct `wss://` connection to Azure |
| Anthropic | Managed | HTTP streaming over WebSocket |
| Google AI Studio (Gemini) | Managed | HTTP streaming over WebSocket |
| Vertex AI | Managed | HTTP streaming over WebSocket |
| AWS Bedrock | Managed | HTTP streaming over WebSocket |
| All other providers | Managed | HTTP streaming over WebSocket |
**Note**: Both native and managed modes provide the same event stream format. The difference is transparent to clients.
### Configuration
No special configuration needed. WebSocket mode is automatically available on the `/v1/responses` endpoint when accessed via WebSocket protocol (`ws://` or `wss://`).
For LiteLLM Proxy, ensure your models are configured normally:
```yaml showLineNumbers title="config.yaml"
model_list:
- model_name: gemini-2.5-flash
litellm_params:
model: gemini/gemini-2.5-flash
api_key: os.environ/GEMINI_API_KEY
- model_name: gpt-4o
litellm_params:
model: openai/gpt-4o
api_key: os.environ/OPENAI_API_KEY
```
Both models will automatically support WebSocket mode at `ws://localhost:4000/v1/responses`.
## Response ID Security
By default, LiteLLM Proxy prevents users from accessing other users' response IDs.

View file

@ -2,7 +2,7 @@
| Feature | Supported |
|---------|-----------|
| Supported Providers | `perplexity`, `tavily`, `parallel_ai`, `exa_ai`, `brave`, `google_pse`, `dataforseo`, `firecrawl`, `searxng`, `linkup` |
| Supported Providers | `perplexity`, `tavily`, `parallel_ai`, `exa_ai`, `brave`, `google_pse`, `dataforseo`, `firecrawl`, `searxng`, `linkup`, `duckduckgo`, `searchapi` |
| Cost Tracking | ✅ |
| Logging | ✅ |
| Load Balancing | ❌ |
@ -210,7 +210,7 @@ See the [official Perplexity Search documentation](https://docs.perplexity.ai/ap
| Parameter | Type | Required | Description |
|-----------|------|----------|-------------|
| `query` | string or array | Yes | Search query. Can be a single string or array of strings |
| `search_provider` | string | Yes (SDK) | The search provider to use: `"perplexity"`, `"tavily"`, `"parallel_ai"`, `"exa_ai"`, `"brave"`, `"google_pse"`, `"dataforseo"`, `"firecrawl"`, `"searxng"`, or `"linkup"` |
| `search_provider` | string | Yes (SDK) | The search provider to use: `"perplexity"`, `"tavily"`, `"parallel_ai"`, `"exa_ai"`, `"brave"`, `"google_pse"`, `"dataforseo"`, `"firecrawl"`, `"searxng"`, `"linkup"`, `"duckduckgo"`, or `"searchapi"` |
| `search_tool_name` | string | Yes (Proxy) | Name of the search tool configured in `config.yaml` |
| `max_results` | integer | No | Maximum number of results to return (1-20). Default: 10 |
| `search_domain_filter` | array | No | List of domains to filter results (max 20 domains) |
@ -276,7 +276,8 @@ The response follows Perplexity's search format with the following structure:
| Firecrawl | `FIRECRAWL_API_KEY` | `firecrawl` |
| SearXNG | `SEARXNG_API_BASE` (required) | `searxng` |
| Linkup | `LINKUP_API_KEY` | `linkup` |
| DuckDuckGo | `DUCKDUCKGO_API_BASE` | `duckduckgo` |
| DuckDuckGo | `DUCKDUCKGO_API_BASE` | `duckduckgo` |
| SearchAPI.io | `SEARCHAPI_API_KEY` | `searchapi` |
See the individual provider documentation for detailed setup instructions and provider-specific parameters.

View file

@ -0,0 +1,197 @@
# SearchAPI.io (Google Search)
Get started by creating a free API key via https://www.searchapi.io/.
SearchAPI.io provides access to Google Search results with a simple API. It supports all Google Search parameters including location, language, time filters, and more.
For complete documentation on all supported parameters, visit https://www.searchapi.io/docs/google.
## LiteLLM Python SDK
```python showLineNumbers title="SearchAPI.io Search"
import os
from litellm import search
os.environ["SEARCHAPI_API_KEY"] = "your-api-key"
response = search(
query="latest AI developments",
search_provider="searchapi",
max_results=10
)
# Access search results
for result in response.results:
print(f"{result.title}: {result.url}")
print(f"Snippet: {result.snippet}\n")
```
### Advanced Usage with SearchAPI.io Parameters
SearchAPI.io supports many Google Search-specific parameters:
```python showLineNumbers title="Advanced SearchAPI.io Parameters"
import os
from litellm import search
os.environ["SEARCHAPI_API_KEY"] = "your-api-key"
response = search(
query="machine learning research",
search_provider="searchapi",
max_results=10,
# Unified parameters
country="US",
search_domain_filter=["arxiv.org", "nature.com"],
# SearchAPI.io specific parameters
gl="us", # Country code
hl="en", # Interface language
time_period="last_month", # Time filter
safe="active", # SafeSearch
device="desktop", # Device type
location="New York" # Geographic location
)
```
## LiteLLM AI Gateway
### 1. Setup config.yaml
```yaml showLineNumbers title="config.yaml"
model_list:
- model_name: gpt-4
litellm_params:
model: gpt-4
api_key: os.environ/OPENAI_API_KEY
search_tools:
- search_tool_name: google-search
litellm_params:
search_provider: searchapi
api_key: os.environ/SEARCHAPI_API_KEY
```
### 2. Start the proxy
```bash
litellm --config /path/to/config.yaml
# RUNNING on http://0.0.0.0:4000
```
### 3. Test the search endpoint
```bash showLineNumbers title="Test Request"
curl http://0.0.0.0:4000/v1/search/google-search \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"query": "latest AI developments",
"max_results": 10,
"country": "US"
}'
```
## SearchAPI.io Specific Parameters
SearchAPI.io supports many Google Search parameters. Here are some commonly used ones:
| Parameter | Type | Description |
|-----------|------|-------------|
| `gl` | string | Country code (e.g., 'us', 'uk', 'de') |
| `hl` | string | Interface language (e.g., 'en', 'es', 'fr') |
| `location` | string | Geographic location (e.g., 'New York', 'London') |
| `device` | string | Device type: 'desktop', 'mobile', 'tablet' |
| `time_period` | string | Time filter: 'last_hour', 'last_day', 'last_week', 'last_month', 'last_year' |
| `time_period_min` | string | Start date (MM/DD/YYYY) |
| `time_period_max` | string | End date (MM/DD/YYYY) |
| `safe` | string | SafeSearch: 'active' or 'off' |
| `lr` | string | Language restriction (e.g., 'lang_en', 'lang_es') |
| `cr` | string | Country restriction |
| `page` | integer | Page number for pagination |
### Example with Time Filters
```python showLineNumbers title="Search with Time Filter"
response = search(
query="AI breakthroughs",
search_provider="searchapi",
max_results=10,
time_period="last_month"
)
```
### Example with Custom Date Range
```python showLineNumbers title="Search with Custom Date Range"
response = search(
query="AI research papers",
search_provider="searchapi",
max_results=10,
time_period_min="01/01/2024",
time_period_max="03/01/2024"
)
```
### Example with Location
```python showLineNumbers title="Search with Location"
response = search(
query="AI conferences",
search_provider="searchapi",
max_results=10,
location="San Francisco",
gl="us"
)
```
## Response Format
SearchAPI.io returns results in the standard LiteLLM search format:
```json
{
"object": "search",
"results": [
{
"title": "Latest AI Developments",
"url": "https://example.com/ai-news",
"snippet": "Recent breakthroughs in artificial intelligence...",
"date": "2024-01-15"
}
]
}
```
## Rate Limits
SearchAPI.io has different rate limits based on your plan:
- Free tier: 100 requests/month
- Paid plans: Higher limits available
Check your current usage at https://www.searchapi.io/dashboard.
## Error Handling
```python showLineNumbers title="Error Handling"
from litellm import search
import os
os.environ["SEARCHAPI_API_KEY"] = "your-api-key"
try:
response = search(
query="test query",
search_provider="searchapi",
max_results=10
)
print(f"Found {len(response.results)} results")
except Exception as e:
print(f"Search failed: {str(e)}")
```
## Additional Resources
- SearchAPI.io Documentation: https://www.searchapi.io/docs
- API Dashboard: https://www.searchapi.io/dashboard
- Pricing: https://www.searchapi.io/pricing

View file

@ -0,0 +1,20 @@
-- Rename call_policy to input_policy
ALTER TABLE "LiteLLM_ToolTable" RENAME COLUMN "call_policy" TO "input_policy";
-- Add output_policy column
ALTER TABLE "LiteLLM_ToolTable" ADD COLUMN "output_policy" TEXT NOT NULL DEFAULT 'untrusted';
-- Add user_agent column
ALTER TABLE "LiteLLM_ToolTable" ADD COLUMN "user_agent" TEXT;
-- Add last_used_at column
ALTER TABLE "LiteLLM_ToolTable" ADD COLUMN "last_used_at" TIMESTAMP(3);
-- Drop old index on call_policy
DROP INDEX IF EXISTS "LiteLLM_ToolTable_call_policy_idx";
-- CreateIndex
CREATE INDEX "LiteLLM_ToolTable_input_policy_idx" ON "LiteLLM_ToolTable"("input_policy");
-- CreateIndex
CREATE INDEX "LiteLLM_ToolTable_output_policy_idx" ON "LiteLLM_ToolTable"("output_policy");

View file

@ -1246,6 +1246,7 @@ from .ocr.main import *
from .rag.main import *
from .search.main import *
from .realtime_api.main import _arealtime
from .responses.main import _aresponses_websocket
from .fine_tuning.main import *
from .files.main import *
from .vector_store_files.main import (

View file

@ -271,8 +271,13 @@ async def asend_message(
card_url = getattr(agent_card, "url", None) if agent_card else None
context_id = trace_id or str(uuid.uuid4())
if request.params.message.context_id is None:
request.params.message.context_id = context_id
message = request.params.message
if isinstance(message, dict):
if message.get("context_id") is None:
message["context_id"] = context_id
else:
if getattr(message, "context_id", None) is None:
message.context_id = context_id
# Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL
a2a_response = None

View file

@ -1,14 +1,10 @@
import json
import time
from typing import Any, List, Literal, Optional, Tuple
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.types.llms.openai import Batch
from litellm.types.utils import CallTypes, ModelInfo, ModelResponse, Usage
from litellm.types.utils import CallTypes, ModelInfo, Usage
from litellm.utils import token_counter

View file

@ -7,7 +7,6 @@ https://platform.openai.com/docs/api-reference/files
import asyncio
import contextvars
import os
import time
import uuid as uuid_module
from functools import partial
@ -20,10 +19,12 @@ from litellm import get_secret_str
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.files.handler import AnthropicFilesHandler
from litellm.llms.azure.common_utils import get_azure_credentials
from litellm.llms.azure.files.handler import AzureOpenAIFilesAPI
from litellm.llms.bedrock.files.handler import BedrockFilesHandler
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.llms.openai.common_utils import get_openai_credentials
from litellm.llms.openai.openai import FileDeleted, FileObject, OpenAIFilesAPI
from litellm.llms.vertex_ai.files.handler import VertexAIFilesHandler
from litellm.types.llms.openai import (
@ -185,95 +186,36 @@ def create_file(
timeout=timeout,
)
elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
api_base = (
optional_params.api_base
or litellm.api_base
or os.getenv("OPENAI_BASE_URL")
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
openai_creds = get_openai_credentials(
api_base=optional_params.api_base,
api_key=optional_params.api_key,
organization=optional_params.organization,
)
organization = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
)
# set API KEY
api_key = (
optional_params.api_key
or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
or litellm.openai_key
or os.getenv("OPENAI_API_KEY")
)
response = openai_files_instance.create_file(
_is_async=_is_async,
api_base=api_base,
api_key=api_key,
api_base=openai_creds.api_base,
api_key=openai_creds.api_key,
timeout=timeout,
max_retries=optional_params.max_retries,
organization=organization,
organization=openai_creds.organization,
create_file_data=_create_file_request,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore
api_version = (
optional_params.api_version
or litellm.api_version
or get_secret_str("AZURE_API_VERSION")
) # type: ignore
api_key = (
optional_params.api_key
or litellm.api_key
or litellm.azure_key
or get_secret_str("AZURE_OPENAI_API_KEY")
or get_secret_str("AZURE_API_KEY")
) # type: ignore
extra_body = optional_params.get("extra_body", {})
if extra_body is not None:
extra_body.pop("azure_ad_token", None)
else:
get_secret_str("AZURE_AD_TOKEN") # type: ignore
azure_creds = get_azure_credentials(
api_base=optional_params.api_base,
api_key=optional_params.api_key,
api_version=optional_params.api_version,
)
response = azure_files_instance.create_file(
_is_async=_is_async,
api_base=api_base,
api_key=api_key,
api_version=api_version,
api_base=azure_creds.api_base,
api_key=azure_creds.api_key,
api_version=azure_creds.api_version,
timeout=timeout,
max_retries=optional_params.max_retries,
create_file_data=_create_file_request,
litellm_params=litellm_params_dict,
)
elif custom_llm_provider == "vertex_ai":
api_base = optional_params.api_base or ""
vertex_ai_project = (
optional_params.vertex_project
or litellm.vertex_project
or get_secret_str("VERTEXAI_PROJECT")
)
vertex_ai_location = (
optional_params.vertex_location
or litellm.vertex_location
or get_secret_str("VERTEXAI_LOCATION")
)
vertex_credentials = optional_params.vertex_credentials or get_secret_str(
"VERTEXAI_CREDENTIALS"
)
response = vertex_ai_files_instance.create_file(
_is_async=_is_async,
api_base=api_base,
vertex_project=vertex_ai_project,
vertex_location=vertex_ai_location,
vertex_credentials=vertex_credentials,
timeout=timeout,
max_retries=optional_params.max_retries,
create_file_data=_create_file_request,
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_file'. Only ['openai', 'azure', 'vertex_ai', 'manus'] are supported.".format(
@ -367,64 +309,31 @@ def file_retrieve(
_is_async = kwargs.pop("is_async", False) is True
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
api_base = (
optional_params.api_base
or litellm.api_base
or os.getenv("OPENAI_BASE_URL")
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
openai_creds = get_openai_credentials(
api_base=optional_params.api_base,
api_key=optional_params.api_key,
organization=optional_params.organization,
)
organization = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
)
# set API KEY
api_key = (
optional_params.api_key
or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
or litellm.openai_key
or os.getenv("OPENAI_API_KEY")
)
response = openai_files_instance.retrieve_file(
file_id=file_id,
_is_async=_is_async,
api_base=api_base,
api_key=api_key,
api_base=openai_creds.api_base,
api_key=openai_creds.api_key,
timeout=timeout,
max_retries=optional_params.max_retries,
organization=organization,
organization=openai_creds.organization,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore
api_version = (
optional_params.api_version
or litellm.api_version
or get_secret_str("AZURE_API_VERSION")
) # type: ignore
api_key = (
optional_params.api_key
or litellm.api_key
or litellm.azure_key
or get_secret_str("AZURE_OPENAI_API_KEY")
or get_secret_str("AZURE_API_KEY")
) # type: ignore
extra_body = optional_params.get("extra_body", {})
if extra_body is not None:
extra_body.pop("azure_ad_token", None)
else:
get_secret_str("AZURE_AD_TOKEN") # type: ignore
azure_creds = get_azure_credentials(
api_base=optional_params.api_base,
api_key=optional_params.api_key,
api_version=optional_params.api_version,
)
response = azure_files_instance.retrieve_file(
_is_async=_is_async,
api_base=api_base,
api_key=api_key,
api_version=api_version,
api_base=azure_creds.api_base,
api_key=azure_creds.api_key,
api_version=azure_creds.api_version,
timeout=timeout,
max_retries=optional_params.max_retries,
file_id=file_id,
@ -576,63 +485,31 @@ def file_delete(
timeout = 600.0
_is_async = kwargs.pop("is_async", False) is True
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
api_base = (
optional_params.api_base
or litellm.api_base
or os.getenv("OPENAI_BASE_URL")
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
)
# set API KEY
api_key = (
optional_params.api_key
or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
or litellm.openai_key
or os.getenv("OPENAI_API_KEY")
openai_creds = get_openai_credentials(
api_base=optional_params.api_base,
api_key=optional_params.api_key,
organization=optional_params.organization,
)
response = openai_files_instance.delete_file(
file_id=file_id,
_is_async=_is_async,
api_base=api_base,
api_key=api_key,
api_base=openai_creds.api_base,
api_key=openai_creds.api_key,
timeout=timeout,
max_retries=optional_params.max_retries,
organization=organization,
organization=openai_creds.organization,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore
api_version = (
optional_params.api_version
or litellm.api_version
or get_secret_str("AZURE_API_VERSION")
) # type: ignore
api_key = (
optional_params.api_key
or litellm.api_key
or litellm.azure_key
or get_secret_str("AZURE_OPENAI_API_KEY")
or get_secret_str("AZURE_API_KEY")
) # type: ignore
extra_body = optional_params.get("extra_body", {})
if extra_body is not None:
extra_body.pop("azure_ad_token", None)
else:
get_secret_str("AZURE_AD_TOKEN") # type: ignore
azure_creds = get_azure_credentials(
api_base=optional_params.api_base,
api_key=optional_params.api_key,
api_version=optional_params.api_version,
)
response = azure_files_instance.delete_file(
_is_async=_is_async,
api_base=api_base,
api_key=api_key,
api_version=api_version,
api_base=azure_creds.api_base,
api_key=azure_creds.api_key,
api_version=azure_creds.api_version,
timeout=timeout,
max_retries=optional_params.max_retries,
file_id=file_id,
@ -815,64 +692,31 @@ def file_list(
)
return response
elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
api_base = (
optional_params.api_base
or litellm.api_base
or os.getenv("OPENAI_BASE_URL")
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
openai_creds = get_openai_credentials(
api_base=optional_params.api_base,
api_key=optional_params.api_key,
organization=optional_params.organization,
)
organization = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
)
# set API KEY
api_key = (
optional_params.api_key
or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
or litellm.openai_key
or os.getenv("OPENAI_API_KEY")
)
response = openai_files_instance.list_files(
purpose=purpose,
_is_async=_is_async,
api_base=api_base,
api_key=api_key,
api_base=openai_creds.api_base,
api_key=openai_creds.api_key,
timeout=timeout,
max_retries=optional_params.max_retries,
organization=organization,
organization=openai_creds.organization,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore
api_version = (
optional_params.api_version
or litellm.api_version
or get_secret_str("AZURE_API_VERSION")
) # type: ignore
api_key = (
optional_params.api_key
or litellm.api_key
or litellm.azure_key
or get_secret_str("AZURE_OPENAI_API_KEY")
or get_secret_str("AZURE_API_KEY")
) # type: ignore
extra_body = optional_params.get("extra_body", {})
if extra_body is not None:
extra_body.pop("azure_ad_token", None)
else:
get_secret_str("AZURE_AD_TOKEN") # type: ignore
azure_creds = get_azure_credentials(
api_base=optional_params.api_base,
api_key=optional_params.api_key,
api_version=optional_params.api_version,
)
response = azure_files_instance.list_files(
_is_async=_is_async,
api_base=api_base,
api_key=api_key,
api_version=api_version,
api_base=azure_creds.api_base,
api_key=azure_creds.api_key,
api_version=azure_creds.api_version,
timeout=timeout,
max_retries=optional_params.max_retries,
purpose=purpose,
@ -1003,64 +847,31 @@ def file_content(
return response
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
api_base = (
optional_params.api_base
or litellm.api_base
or os.getenv("OPENAI_BASE_URL")
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
openai_creds = get_openai_credentials(
api_base=optional_params.api_base,
api_key=optional_params.api_key,
organization=optional_params.organization,
)
organization = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
)
# set API KEY
api_key = (
optional_params.api_key
or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
or litellm.openai_key
or os.getenv("OPENAI_API_KEY")
)
response = openai_files_instance.file_content(
_is_async=_is_async,
file_content_request=_file_content_request,
api_base=api_base,
api_key=api_key,
api_base=openai_creds.api_base,
api_key=openai_creds.api_key,
timeout=timeout,
max_retries=optional_params.max_retries,
organization=organization,
organization=openai_creds.organization,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore
api_version = (
optional_params.api_version
or litellm.api_version
or get_secret_str("AZURE_API_VERSION")
) # type: ignore
api_key = (
optional_params.api_key
or litellm.api_key
or litellm.azure_key
or get_secret_str("AZURE_OPENAI_API_KEY")
or get_secret_str("AZURE_API_KEY")
) # type: ignore
extra_body = optional_params.get("extra_body", {})
if extra_body is not None:
extra_body.pop("azure_ad_token", None)
else:
get_secret_str("AZURE_AD_TOKEN") # type: ignore
azure_creds = get_azure_credentials(
api_base=optional_params.api_base,
api_key=optional_params.api_key,
api_version=optional_params.api_version,
)
response = azure_files_instance.file_content(
_is_async=_is_async,
api_base=api_base,
api_key=api_key,
api_version=api_version,
api_base=azure_creds.api_base,
api_key=azure_creds.api_key,
api_version=azure_creds.api_version,
timeout=timeout,
max_retries=optional_params.max_retries,
file_content_request=_file_content_request,

View file

@ -167,12 +167,12 @@ class HeliconeLogger:
if "claude" in model and not is_vertex_ai:
url = f"{self.api_base}/anthropic/v1/log"
provider_url = "https://api.anthropic.com/v1/messages"
elif "gemini" in model:
url = f"{self.api_base}/custom/v1/log"
provider_url = "https://generativelanguage.googleapis.com/v1beta"
elif is_vertex_ai:
url = f"{self.api_base}/custom/v1/log"
provider_url = "https://aiplatform.googleapis.com/v1"
elif "gemini" in model:
url = f"{self.api_base}/custom/v1/log"
provider_url = "https://generativelanguage.googleapis.com/v1beta"
headers = {
"Authorization": f"Bearer {self.key}",
"Content-Type": "application/json",

View file

@ -1,6 +1,6 @@
import json
import os
from typing import Any, Callable, Dict, Literal, Optional, Union, cast
from typing import Any, Callable, Dict, Literal, NamedTuple, Optional, Union, cast
import httpx
from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
@ -789,3 +789,39 @@ class BaseAzureLLM(BaseOpenAILLM):
return param_value
return os.getenv(env_var_key)
class AzureCredentials(NamedTuple):
api_base: Optional[str]
api_key: Optional[str]
api_version: Optional[str]
def get_azure_credentials(
api_base: Optional[str] = None,
api_key: Optional[str] = None,
api_version: Optional[str] = None,
) -> AzureCredentials:
"""Resolve Azure credentials from params, litellm globals, and env vars."""
resolved_api_base = (
api_base
or litellm.api_base
or get_secret_str("AZURE_API_BASE")
)
resolved_api_version = (
api_version
or litellm.api_version
or get_secret_str("AZURE_API_VERSION")
)
resolved_api_key = (
api_key
or litellm.api_key
or litellm.azure_key
or get_secret_str("AZURE_OPENAI_API_KEY")
or get_secret_str("AZURE_API_KEY")
)
return AzureCredentials(
api_base=resolved_api_base,
api_key=resolved_api_key,
api_version=resolved_api_version,
)

View file

@ -218,6 +218,18 @@ class BaseResponsesAPIConfig(ABC):
"""Returns True if litellm should fake a stream for the given model and stream value"""
return False
def supports_native_websocket(self) -> bool:
"""
Returns True if the provider has a native WebSocket endpoint for Responses API.
Providers with native websocket support can connect directly to wss:// endpoints.
Providers without native support will use the ManagedResponsesWebSocketHandler
which makes HTTP streaming calls and forwards events over the websocket.
Default: False (use managed websocket handler)
"""
return False
#########################################################
########## CANCEL RESPONSE API TRANSFORMATION ##########
#########################################################

View file

@ -1,14 +1,14 @@
import json
from typing import Any, Optional
from litellm.exceptions import AuthenticationError
from litellm.constants import STREAM_SSE_DONE_STRING
from litellm.exceptions import AuthenticationError
from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.llms.openai.common_utils import OpenAIError
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
_safe_convert_created_field,
)
from litellm.llms.openai.common_utils import OpenAIError
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.types.llms.openai import (
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
@ -200,3 +200,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
api_base = api_base or self.authenticator.get_api_base() or CHATGPT_API_BASE
api_base = api_base.rstrip("/")
return f"{api_base}/responses"
def supports_native_websocket(self) -> bool:
"""ChatGPT does not support native WebSocket for Responses API"""
return False

View file

@ -69,6 +69,7 @@ from litellm.responses.streaming_iterator import (
BaseResponsesAPIStreamingIterator,
MockResponsesAPIStreamingIterator,
ResponsesAPIStreamingIterator,
ResponsesWebSocketStreaming,
SyncResponsesAPIStreamingIterator,
)
from litellm.types.containers.main import (
@ -4731,6 +4732,123 @@ class BaseLLMHTTPHandler:
f"Unexpected error while closing WebSocket: {close_error}"
)
async def async_responses_websocket(
self,
model: str,
websocket: Any,
logging_obj: LiteLLMLoggingObj,
responses_api_provider_config: Optional[BaseResponsesAPIConfig],
api_base: Optional[str] = None,
api_key: Optional[str] = None,
timeout: Optional[float] = None,
user_api_key_dict: Optional[Any] = None,
litellm_metadata: Optional[Dict[str, Any]] = None,
custom_llm_provider: Optional[str] = None,
**kwargs: Any,
):
"""
Handles Responses API WebSocket mode.
For providers with native websocket support (OpenAI, Azure):
- Opens a persistent WebSocket to the provider's /v1/responses endpoint
- Proxies response.create events bidirectionally for lower-latency agentic workflows
For providers without native websocket support (all others):
- Uses ManagedResponsesWebSocketHandler which makes HTTP streaming calls
- Forwards events over the websocket connection
"""
if responses_api_provider_config is None or not responses_api_provider_config.supports_native_websocket():
from litellm.responses.streaming_iterator import (
ManagedResponsesWebSocketHandler,
)
handler = ManagedResponsesWebSocketHandler(
websocket=websocket,
model=model,
logging_obj=logging_obj,
user_api_key_dict=user_api_key_dict,
litellm_metadata=litellm_metadata,
api_key=api_key,
api_base=api_base,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
**kwargs,
)
await handler.run()
return
import websockets
from websockets.asyncio.client import ClientConnection
litellm_params = GenericLiteLLMParams()
headers = responses_api_provider_config.validate_environment(
headers={},
model=model,
litellm_params=litellm_params,
)
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
http_url = responses_api_provider_config.get_complete_url(
api_base=api_base,
litellm_params={},
)
ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://")
try:
ssl_context = get_shared_realtime_ssl_context()
if ws_url.startswith("wss://") and ssl_context is False:
ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
ssl_context.check_hostname = False
ssl_context.verify_mode = ssl.CERT_NONE
logging_obj.pre_call(
input=None,
api_key=api_key or "",
additional_args={
"api_base": ws_url,
"headers": headers,
"complete_input_dict": {"mode": "responses_websocket"},
},
)
async with websockets.connect( # type: ignore
ws_url,
additional_headers=headers,
max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES,
ssl=ssl_context,
) as backend_ws:
_request_data: Dict[str, Any] = {}
if litellm_metadata:
_request_data["litellm_metadata"] = litellm_metadata
streaming = ResponsesWebSocketStreaming(
websocket=websocket,
backend_ws=cast(ClientConnection, backend_ws),
logging_obj=logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=_request_data,
)
await streaming.bidirectional_forward()
except websockets.exceptions.InvalidStatusCode as e: # type: ignore
verbose_logger.exception(f"Error connecting to responses WS backend: {e}")
await websocket.close(code=e.status_code, reason=str(e))
except Exception as e:
verbose_logger.exception(f"Error in responses WS: {e}")
try:
await websocket.close(
code=1011, reason=f"Internal server error: {str(e)}"
)
except RuntimeError as close_error:
if "already completed" in str(close_error) or "websocket.close" in str(
close_error
):
pass
else:
raise Exception(
f"Unexpected error while closing WebSocket: {close_error}"
)
def image_edit_handler(
self,
model: str,

View file

@ -98,3 +98,7 @@ class DatabricksResponsesAPIConfig(DatabricksBase, OpenAIResponsesAPIConfig):
litellm_params=litellm_params,
headers=headers,
)
def supports_native_websocket(self) -> bool:
"""Databricks does not support native WebSocket for Responses API"""
return False

View file

@ -22,8 +22,8 @@ from litellm.types.utils import LlmProviders
from ..authenticator import Authenticator
from ..common_utils import (
GetAPIKeyError,
GITHUB_COPILOT_API_BASE,
GetAPIKeyError,
get_copilot_default_headers,
)
@ -329,3 +329,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig):
)
return False
def supports_native_websocket(self) -> bool:
"""GitHub Copilot does not support native WebSocket for Responses API"""
return False

View file

@ -69,3 +69,7 @@ class HostedVLLMResponsesAPIConfig(OpenAIResponsesAPIConfig):
return f"{api_base}/responses"
return f"{api_base}/v1/responses"
def supports_native_websocket(self) -> bool:
"""Hosted vLLM does not support native WebSocket for Responses API"""
return False

View file

@ -46,3 +46,7 @@ class LiteLLMProxyResponsesAPIConfig(OpenAIResponsesAPIConfig):
api_base = api_base.rstrip("/")
return f"{api_base}/responses"
def supports_native_websocket(self) -> bool:
"""LiteLLM Proxy does not support native WebSocket for Responses API"""
return False

View file

@ -247,6 +247,10 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig):
response._hidden_params["headers"] = raw_response_headers
return response
def supports_native_websocket(self) -> bool:
"""Manus does not support native WebSocket for Responses API"""
return False
def transform_get_response_api_request(
self,
response_id: str,

View file

@ -5,8 +5,9 @@ Common helpers / utils across al OpenAI endpoints
import hashlib
import inspect
import json
import os
import ssl
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union
from typing import TYPE_CHECKING, Any, Dict, List, Literal, NamedTuple, Optional, Tuple, Union
import httpx
import openai
@ -244,3 +245,39 @@ class BaseOpenAILLM:
)
class OpenAICredentials(NamedTuple):
api_base: str
api_key: Optional[str]
organization: Optional[str]
def get_openai_credentials(
api_base: Optional[str] = None,
api_key: Optional[str] = None,
organization: Optional[str] = None,
) -> OpenAICredentials:
"""Resolve OpenAI credentials from params, litellm globals, and env vars."""
resolved_api_base = (
api_base
or litellm.api_base
or os.getenv("OPENAI_BASE_URL")
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
resolved_organization = (
organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
or None
)
resolved_api_key = (
api_key
or litellm.api_key
or litellm.openai_key
or os.getenv("OPENAI_API_KEY")
)
return OpenAICredentials(
api_base=resolved_api_base,
api_key=resolved_api_key,
organization=resolved_organization,
)

View file

@ -344,6 +344,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
)
return False
def supports_native_websocket(self) -> bool:
"""OpenAI supports native WebSocket for Responses API"""
return True
#########################################################
########## DELETE RESPONSE API TRANSFORMATION ##############
#########################################################
@ -524,7 +528,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
OpenAI API expects the following request
- POST /v1/responses/compact
"""
url = f"{api_base}/compact"
# Preserve query params (e.g., api-version) while appending /compact.
parsed_url = httpx.URL(api_base)
compact_path = parsed_url.path.rstrip("/") + "/compact"
url = str(parsed_url.copy_with(path=compact_path))
input = self._validate_input_param(input)
data = dict(

View file

@ -0,0 +1,11 @@
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from .transformation import OpenRouterImageEditConfig
__all__ = [
"OpenRouterImageEditConfig",
]
def get_openrouter_image_edit_config(model: str) -> BaseImageEditConfig:
return OpenRouterImageEditConfig()

View file

@ -0,0 +1,367 @@
"""
OpenRouter Image Edit Support
OpenRouter provides image editing through chat completion endpoints.
The source image is sent as a base64 data URL in the message content,
and the response contains edited images in the message's images array.
Request format:
{
"model": "google/gemini-2.5-flash-image",
"messages": [{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}},
{"type": "text", "text": "Edit this image by..."}
]
}],
"modalities": ["image", "text"]
}
Response format:
{
"choices": [{
"message": {
"content": "Here is the edited image.",
"role": "assistant",
"images": [{
"image_url": {"url": "data:image/png;base64,..."},
"type": "image_url"
}]
}
}],
"usage": {
"completion_tokens": 1299,
"prompt_tokens": 300,
"total_tokens": 1599,
"completion_tokens_details": {"image_tokens": 1290},
"cost": 0.0387243
}
}
"""
import base64
from io import BufferedReader, BytesIO
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
import httpx
from httpx._types import RequestFiles
import litellm
from litellm.images.utils import ImageEditRequestUtils
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.llms.openrouter.common_utils import OpenRouterException
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import FileTypes, ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class OpenRouterImageEditConfig(BaseImageEditConfig):
"""
Configuration for OpenRouter image editing via chat completions.
OpenRouter uses the chat completions endpoint for image editing.
The source image is sent as a base64 data URL in the message content,
and the response contains edited images in the message's images array.
"""
def get_supported_openai_params(self, model: str) -> list:
return ["size", "quality", "n"]
def map_openai_params(
self,
image_edit_optional_params: ImageEditOptionalRequestParams,
model: str,
drop_params: bool,
) -> Dict:
supported_params = self.get_supported_openai_params(model)
mapped_params: Dict[str, Any] = {}
for key, value in image_edit_optional_params.items():
if key in supported_params:
if key == "size":
if "image_config" not in mapped_params:
mapped_params["image_config"] = {}
mapped_params["image_config"]["aspect_ratio"] = self._map_size_to_aspect_ratio(value)
elif key == "quality":
image_size = self._map_quality_to_image_size(value)
if image_size:
if "image_config" not in mapped_params:
mapped_params["image_config"] = {}
mapped_params["image_config"]["image_size"] = image_size
else:
mapped_params[key] = value
return mapped_params
def validate_environment(
self,
headers: dict,
model: str,
api_key: Optional[str] = None,
) -> dict:
api_key = (
api_key
or litellm.api_key
or get_secret_str("OPENROUTER_API_KEY")
)
if not api_key:
raise ValueError("OPENROUTER_API_KEY is not set")
headers.update(
{
"Authorization": f"Bearer {api_key}",
}
)
return headers
def use_multipart_form_data(self) -> bool:
"""OpenRouter uses JSON requests, not multipart/form-data."""
return False
def get_complete_url(
self,
model: str,
api_base: Optional[str],
litellm_params: dict,
) -> str:
base_url = api_base or get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1"
base_url = base_url.rstrip("/")
if not base_url.endswith("/chat/completions"):
return f"{base_url}/chat/completions"
return base_url
def transform_image_edit_request(
self,
model: str,
prompt: Optional[str],
image: Optional[FileTypes],
image_edit_optional_request_params: Dict,
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> Tuple[Dict, RequestFiles]:
content_parts: List[Dict[str, Any]] = []
# Add source image(s) as base64 data URLs
if image is not None:
images = image if isinstance(image, list) else [image]
for img in images:
if img is None:
continue
mime_type = ImageEditRequestUtils.get_image_content_type(img)
image_bytes = self._read_image_bytes(img)
b64_data = base64.b64encode(image_bytes).decode("utf-8")
content_parts.append(
{
"type": "image_url",
"image_url": {
"url": f"data:{mime_type};base64,{b64_data}"
},
}
)
# Add the text prompt
if prompt:
content_parts.append({"type": "text", "text": prompt})
request_body: Dict[str, Any] = {
"model": model,
"messages": [
{
"role": "user",
"content": content_parts,
}
],
"modalities": ["image", "text"],
}
# Add mapped optional params (image_config, n, etc.)
for key, value in image_edit_optional_request_params.items():
if key not in ("model", "messages", "modalities"):
request_body[key] = value
empty_files = cast(RequestFiles, [])
return request_body, empty_files
def transform_image_edit_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
) -> ImageResponse:
try:
response_json = raw_response.json()
except Exception as e:
raise OpenRouterException(
message=f"Error parsing OpenRouter response: {str(e)}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
model_response = ImageResponse()
model_response.data = []
try:
choices = response_json.get("choices", [])
for choice in choices:
message = choice.get("message", {})
images = message.get("images", [])
for image_data in images:
image_url_obj = image_data.get("image_url", {})
image_url = image_url_obj.get("url")
if image_url:
if image_url.startswith("data:"):
# Extract base64 data from data URL
parts = image_url.split(",", 1)
b64_data = parts[1] if len(parts) > 1 else None
model_response.data.append(
ImageObject(
b64_json=b64_data,
url=None,
revised_prompt=None,
)
)
else:
model_response.data.append(
ImageObject(
b64_json=None,
url=image_url,
revised_prompt=None,
)
)
except Exception as e:
raise OpenRouterException(
message=f"Error transforming OpenRouter image edit response: {str(e)}",
status_code=500,
headers={},
)
self._set_usage_and_cost(model_response, response_json, model)
return model_response
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BaseLLMException:
return OpenRouterException(
message=error_message,
status_code=status_code,
headers=headers,
)
# Private helper methods
def _map_size_to_aspect_ratio(self, size: str) -> str:
"""
Map OpenAI size format to OpenRouter aspect_ratio format.
Uses the same mapping as image generation since OpenRouter
handles both through the same chat completions endpoint.
"""
size_to_aspect_ratio = {
"256x256": "1:1",
"512x512": "1:1",
"1024x1024": "1:1",
"1536x1024": "3:2",
"1792x1024": "16:9",
"1024x1536": "2:3",
"1024x1792": "9:16",
"auto": "1:1",
}
return size_to_aspect_ratio.get(size, "1:1")
def _map_quality_to_image_size(self, quality: str) -> Optional[str]:
"""
Map OpenAI quality to OpenRouter image_size format.
Uses the same mapping as image generation since OpenRouter
handles both through the same chat completions endpoint.
"""
quality_to_image_size = {
"low": "1K",
"standard": "1K",
"medium": "2K",
"high": "4K",
"hd": "4K",
"auto": "1K",
}
return quality_to_image_size.get(quality)
def _set_usage_and_cost(
self,
model_response: ImageResponse,
response_json: dict,
model: str,
) -> None:
"""Extract and set usage and cost information from OpenRouter response."""
usage_data = response_json.get("usage", {})
if usage_data:
prompt_tokens = usage_data.get("prompt_tokens", 0)
total_tokens = usage_data.get("total_tokens", 0)
completion_tokens_details = usage_data.get("completion_tokens_details", {})
image_tokens = completion_tokens_details.get("image_tokens", 0)
# For image edit, input may include image tokens
input_image_tokens = 0
prompt_tokens_details = usage_data.get("prompt_tokens_details", {})
if prompt_tokens_details:
input_image_tokens = prompt_tokens_details.get("image_tokens", 0)
model_response.usage = ImageUsage(
input_tokens=prompt_tokens,
input_tokens_details=ImageUsageInputTokensDetails(
image_tokens=input_image_tokens,
text_tokens=prompt_tokens - input_image_tokens,
),
output_tokens=image_tokens,
total_tokens=total_tokens,
)
cost = usage_data.get("cost")
if cost is not None:
if not hasattr(model_response, "_hidden_params"):
model_response._hidden_params = {}
if "additional_headers" not in model_response._hidden_params:
model_response._hidden_params["additional_headers"] = {}
model_response._hidden_params["additional_headers"][
"llm_provider-x-litellm-response-cost"
] = float(cost)
cost_details = usage_data.get("cost_details", {})
if cost_details:
if "response_cost_details" not in model_response._hidden_params:
model_response._hidden_params["response_cost_details"] = {}
model_response._hidden_params["response_cost_details"].update(cost_details)
model_response._hidden_params["model"] = response_json.get("model", model)
def _read_image_bytes(self, image: FileTypes) -> bytes:
"""Read raw bytes from various image input types."""
if isinstance(image, bytes):
return image
if isinstance(image, BytesIO):
current_pos = image.tell()
image.seek(0)
data = image.read()
image.seek(current_pos)
return data
if isinstance(image, BufferedReader):
current_pos = image.tell()
image.seek(0)
data = image.read()
image.seek(current_pos)
return data
raise ValueError("Unsupported image type for OpenRouter image edit.")

View file

@ -75,3 +75,7 @@ class OpenRouterResponsesAPIConfig(OpenAIResponsesAPIConfig):
api_base = api_base.rstrip("/")
return f"{api_base}/responses"
def supports_native_websocket(self) -> bool:
"""OpenRouter does not support native WebSocket for Responses API"""
return False

View file

@ -490,3 +490,7 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig):
verbose_logger.debug("Failed to transform Perplexity cost object: %s", e)
return chunk
def supports_native_websocket(self) -> bool:
"""Perplexity does not support native WebSocket for Responses API"""
return False

View file

@ -0,0 +1 @@
"""SearchAPI.io integration for LiteLLM."""

View file

@ -0,0 +1,4 @@
"""SearchAPI.io search integration for LiteLLM."""
from litellm.llms.searchapi.search.transformation import SearchAPIConfig
__all__ = ["SearchAPIConfig"]

View file

@ -0,0 +1,232 @@
"""
Calls SearchAPI.io's Google Search API endpoint.
SearchAPI.io API Reference: https://www.searchapi.io/docs/google
"""
from typing import Dict, List, Literal, Optional, TypedDict, Union
from urllib.parse import urlencode
import httpx
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.search.transformation import (
BaseSearchConfig,
SearchResponse,
SearchResult,
)
from litellm.secret_managers.main import get_secret_str
class _SearchAPIRequestRequired(TypedDict):
"""Required fields for SearchAPI.io request."""
engine: str # Required - search engine (e.g., 'google')
q: str # Required - search query
class SearchAPIRequest(_SearchAPIRequestRequired, total=False):
"""
SearchAPI.io request format for Google Search.
Based on: https://www.searchapi.io/docs/google
"""
kgmid: str # Optional - Knowledge Graph identifier
device: str # Optional - device type ('desktop', 'mobile', 'tablet')
location: str # Optional - geographic location
uule: str # Optional - Google-encoded location
google_domain: str # Optional - Google domain (deprecated)
gl: str # Optional - country code (e.g., 'us', 'uk')
hl: str # Optional - interface language (e.g., 'en', 'es')
lr: str # Optional - language restriction (e.g., 'lang_en')
cr: str # Optional - country restriction
nfpr: int # Optional - exclude auto-corrected results (0 or 1)
filter: int # Optional - duplicate/host crowding filter (0 or 1)
safe: str # Optional - SafeSearch ('active', 'off')
time_period: str # Optional - time period ('last_hour', 'last_day', 'last_week', 'last_month', 'last_year')
time_period_min: str # Optional - start date (MM/DD/YYYY)
time_period_max: str # Optional - end date (MM/DD/YYYY)
num: int # Optional - number of results (phased out by Google, constant 10)
page: int # Optional - page number for pagination
optimization_strategy: str # Optional - 'performance' or 'ads'
class SearchAPIConfig(BaseSearchConfig):
SEARCHAPI_API_BASE = "https://www.searchapi.io/api/v1/search"
@staticmethod
def ui_friendly_name() -> str:
return "SearchAPI.io (Google Search)"
def get_http_method(self) -> Literal["GET", "POST"]:
"""
SearchAPI.io uses GET requests for search.
"""
return "GET"
def validate_environment(
self,
headers: Dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
**kwargs,
) -> Dict:
"""
Validate environment and return headers.
"""
api_key = api_key or get_secret_str("SEARCHAPI_API_KEY")
if not api_key:
raise ValueError(
"SEARCHAPI_API_KEY is not set. Set `SEARCHAPI_API_KEY` environment variable."
)
headers["Content-Type"] = "application/json"
return headers
def get_complete_url(
self,
api_base: Optional[str],
optional_params: dict,
data: Optional[Union[Dict, List[Dict]]] = None,
**kwargs,
) -> str:
"""
Get complete URL for Search endpoint with query parameters.
SearchAPI.io uses GET requests and includes api_key in query params.
"""
api_base = api_base or get_secret_str("SEARCHAPI_API_BASE") or self.SEARCHAPI_API_BASE
# Build query parameters from the transformed request body
if data and isinstance(data, dict) and "_searchapi_params" in data:
params = data["_searchapi_params"]
query_string = urlencode(params, doseq=True)
return f"{api_base}?{query_string}"
return api_base
def transform_search_request(
self,
query: Union[str, List[str]],
optional_params: dict,
api_key: Optional[str] = None,
search_engine_id: Optional[str] = None,
**kwargs,
) -> Dict:
"""
Transform Search request to SearchAPI.io format.
Transforms unified spec parameters:
- query q
- max_results num (limited to 10 by Google)
- search_domain_filter q (append site: filters)
- country gl
Args:
query: Search query (string or list of strings)
optional_params: Optional parameters for the request
api_key: API key for authentication
Returns:
Dict with typed request data following SearchAPI.io spec
"""
if isinstance(query, list):
query = " ".join(query)
# Get API key from parameter or environment
api_key = api_key or get_secret_str("SEARCHAPI_API_KEY")
if not api_key:
raise ValueError(
"SEARCHAPI_API_KEY is not set. Set `SEARCHAPI_API_KEY` environment variable."
)
request_data: SearchAPIRequest = {
"engine": "google",
"q": query,
}
# Add API key to request
result_data = dict(request_data)
result_data["api_key"] = api_key
# Transform unified spec parameters to SearchAPI.io format
if "max_results" in optional_params:
# Google now returns constant 10 results, but we can still set num
num_results = min(optional_params["max_results"], 10)
result_data["num"] = num_results
if "search_domain_filter" in optional_params:
# Convert to multiple "site:domain" clauses
domains = optional_params["search_domain_filter"]
if isinstance(domains, list) and len(domains) > 0:
result_data["q"] = self._append_domain_filters(
result_data["q"], domains
)
if "country" in optional_params:
# Map to gl parameter
result_data["gl"] = optional_params["country"].lower()
# Pass through all other SearchAPI.io-specific parameters
for param, value in optional_params.items():
if (
param not in self.get_supported_perplexity_optional_params()
and param not in result_data
):
result_data[param] = value
# Store params in special key for URL building (GET request)
return {
"_searchapi_params": result_data,
}
@staticmethod
def _append_domain_filters(query: str, domains: List[str]) -> str:
"""
Add site: filters to restrict search to specific domains.
"""
domain_clauses = [f"site:{domain}" for domain in domains]
domain_query = " OR ".join(domain_clauses)
return f"({query}) AND ({domain_query})"
def transform_search_response(
self,
raw_response: httpx.Response,
logging_obj: Optional[LiteLLMLoggingObj],
**kwargs,
) -> SearchResponse:
"""
Transform SearchAPI.io response to LiteLLM unified SearchResponse format.
SearchAPI.io LiteLLM mappings:
- organic_results[].title SearchResult.title
- organic_results[].link SearchResult.url
- organic_results[].snippet SearchResult.snippet
- organic_results[].date SearchResult.date
"""
response_json = raw_response.json()
# Transform results to SearchResult objects
results: List[SearchResult] = []
# Process organic results
for result in response_json.get("organic_results", []):
title = result.get("title", "")
url = result.get("link", "")
snippet = result.get("snippet", "")
date = result.get("date") # SearchAPI.io provides date in some results
search_result = SearchResult(
title=title,
url=url,
snippet=snippet,
date=date,
last_updated=None, # SearchAPI.io doesn't provide last_updated
)
results.append(search_result)
return SearchResponse(
results=results,
object="search",
)

View file

@ -571,14 +571,38 @@ def _filter_anyof_fields(schema_dict: Dict[str, Any]) -> Dict[str, Any]:
return schema_dict
def _is_any_type_schema(schema: dict) -> bool:
"""
Detect schemas that represent "any JSON value" (no type constraints).
In JSON Schema, an empty schema {} means "any value is valid".
Schemas with only metadata keys (title, description, default, examples)
but no type-constraining keywords also represent "any type".
Gemini's Schema proto uses TYPE_UNSPECIFIED (0) as default,
so omitting the type field is valid and means "any type".
"""
type_constraining_keys = {
"type",
"properties",
"items",
"anyOf",
"oneOf",
"allOf",
"enum",
"required",
"$ref",
"$schema",
}
return not any(key in type_constraining_keys for key in schema.keys())
def process_items(schema, depth=0):
if depth > DEFAULT_MAX_RECURSE_DEPTH:
raise ValueError(
f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema. Please check the schema for excessive nesting."
)
if isinstance(schema, dict):
if "items" in schema and schema["items"] == {}:
schema["items"] = {"type": "object"}
for key, value in schema.items():
if isinstance(value, dict):
process_items(value, depth + 1)
@ -677,9 +701,8 @@ def convert_anyof_null_to_nullable(schema, depth=0):
# remove null type
anyof.remove(atype)
contains_null = True
elif "type" not in atype and len(atype) == 0:
# Handle empty object case
atype["type"] = "object"
elif isinstance(atype, dict) and _is_any_type_schema(atype):
pass # preserve "any type" semantics — don't coerce to object
if len(anyof) == 0:
# Edge case: response schema with only null type present is invalid in Vertex AI
@ -714,7 +737,8 @@ def add_object_type(schema):
# Gemini requires all function parameters to be type OBJECT
# Handle case where schema has no properties and no type (e.g. tools with no arguments)
if "type" not in schema and "anyOf" not in schema and "oneOf" not in schema and "allOf" not in schema:
schema["type"] = "object"
if not _is_any_type_schema(schema):
schema["type"] = "object"
properties = schema.get("properties", None)
if properties is not None:

View file

@ -2905,6 +2905,7 @@ class ModelResponseIterator:
self.logging_obj = logging_obj
self.is_function_call = check_is_function_call(logging_obj)
self.cumulative_tool_call_index: int = 0
self.has_seen_tool_calls: bool = False
def chunk_parser(self, chunk: dict) -> Optional["ModelResponseStream"]:
try:
@ -2943,6 +2944,40 @@ class ModelResponseIterator:
cumulative_tool_call_index=self.cumulative_tool_call_index,
)
# Track whether tool_calls have been seen across streaming chunks.
# Gemini sends tool_calls and finishReason in separate chunks,
# so we need to remember if earlier chunks contained tool_calls
# to correctly set finish_reason="tool_calls" per the OpenAI spec.
if not self.has_seen_tool_calls:
for choice in model_response.choices:
if hasattr(choice, "delta") and choice.delta and choice.delta.tool_calls:
self.has_seen_tool_calls = True
break
# Handle final chunk with finishReason but no content.
# _process_candidates skips candidates without "content",
# so the finish_reason from the final chunk is lost.
if not model_response.choices and _candidates:
from litellm.types.utils import Delta, StreamingChoices
for candidate in _candidates:
finish_reason_str = candidate.get("finishReason")
if finish_reason_str is not None:
if self.has_seen_tool_calls:
mapped_finish_reason = "tool_calls"
else:
mapped_finish_reason = VertexGeminiConfig._check_finish_reason(
None, finish_reason_str
)
choice = StreamingChoices(
finish_reason=mapped_finish_reason,
index=candidate.get("index", 0),
delta=Delta(content=None, role=None),
logprobs=None,
enhancements=None,
)
model_response.choices.append(choice)
setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) # type: ignore
setattr(model_response, "vertex_ai_url_context_metadata", url_context_metadata) # type: ignore
setattr(model_response, "vertex_ai_safety_ratings", safety_ratings) # type: ignore

View file

@ -16,16 +16,17 @@ from pydantic import fields as pyd_fields
import litellm
from litellm._logging import verbose_logger
from litellm.types.llms.openai import ResponseInputParam, ResponsesAPIStreamingResponse
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
_safe_convert_created_field,
)
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
ResponseInputParam,
ResponsesAPIOptionalRequestParams,
ResponsesAPIResponse,
ResponsesAPIStreamingResponse,
)
from litellm.types.responses.main import DeleteResponseResult
from litellm.types.router import GenericLiteLLMParams
@ -555,3 +556,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig):
# Fall back to the first candidate
return candidates[0]
def supports_native_websocket(self) -> bool:
"""VolcEngine does not support native WebSocket for Responses API"""
return False

View file

@ -252,3 +252,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
return f"{api_base}/responses"
def supports_native_websocket(self) -> bool:
"""XAI does not support native WebSocket for Responses API"""
return False

View file

@ -107,6 +107,7 @@ from litellm.realtime_api.main import _realtime_health_check
from litellm.secret_managers.main import get_secret_bool, get_secret_str
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import (
CustomPricingLiteLLMParams,
ModelResponseStream,
RawRequestTypedDict,
StreamingChoices,
@ -418,6 +419,8 @@ async def acompletion( # noqa: PLR0915
web_search_options: Optional[OpenAIWebSearchOptions] = None,
# Session management
shared_session: Optional["ClientSession"] = None,
# Per-request JSON schema validation (overrides litellm.enable_json_schema_validation)
enable_json_schema_validation: Optional[bool] = None,
**kwargs,
) -> Union[ModelResponse, CustomStreamWrapper]:
"""
@ -562,6 +565,7 @@ async def acompletion( # noqa: PLR0915
"thinking": thinking,
"web_search_options": web_search_options,
"shared_session": shared_session,
"enable_json_schema_validation": enable_json_schema_validation,
}
if custom_llm_provider is None:
_, custom_llm_provider, _, _ = get_llm_provider(
@ -996,6 +1000,32 @@ def _drop_input_examples_from_tools(
return cleaned_tools
def _build_custom_pricing_entry(
custom_llm_provider: str,
kwargs: dict,
model_info: Optional[dict] = None,
) -> dict:
"""Build a complete model cost entry from kwargs and model_info.
Collects all CustomPricingLiteLLMParams fields present in kwargs and
merges metadata from model_info (mode, supports_prompt_caching, max_tokens)
so that register_model() receives the full pricing configuration.
"""
entry: dict = {"litellm_provider": custom_llm_provider}
for field_name in CustomPricingLiteLLMParams.model_fields:
value = kwargs.get(field_name)
if value is not None:
entry[field_name] = value
if model_info and isinstance(model_info, dict):
for key in ("mode", "supports_prompt_caching", "max_tokens"):
if key in model_info and model_info[key] is not None:
entry.setdefault(key, model_info[key])
return entry
@tracer.wrap()
@client
def completion( # type: ignore # noqa: PLR0915
@ -1047,6 +1077,8 @@ def completion( # type: ignore # noqa: PLR0915
thinking: Optional[AnthropicThinkingParam] = None,
# Session management
shared_session: Optional["ClientSession"] = None,
# Per-request JSON schema validation (overrides litellm.enable_json_schema_validation)
enable_json_schema_validation: Optional[bool] = None,
**kwargs,
) -> Union[ModelResponse, CustomStreamWrapper]:
"""
@ -1167,6 +1199,7 @@ def completion( # type: ignore # noqa: PLR0915
thinking=thinking,
web_search_options=web_search_options,
shared_session=shared_session,
enable_json_schema_validation=enable_json_schema_validation,
**kwargs,
)
api_base = kwargs.get("api_base", None)
@ -1351,27 +1384,16 @@ def completion( # type: ignore # noqa: PLR0915
timeout = float(timeout) # type: ignore
### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ###
if input_cost_per_token is not None and output_cost_per_token is not None:
if (
input_cost_per_token is not None and output_cost_per_token is not None
) or input_cost_per_second is not None:
litellm.register_model(
{
f"{custom_llm_provider}/{model}": {
"input_cost_per_token": input_cost_per_token,
"output_cost_per_token": output_cost_per_token,
"litellm_provider": custom_llm_provider,
}
}
)
elif (
input_cost_per_second is not None
): # time based pricing just needs cost in place
output_cost_per_second = output_cost_per_second
litellm.register_model(
{
f"{custom_llm_provider}/{model}": {
"input_cost_per_second": input_cost_per_second,
"output_cost_per_second": output_cost_per_second,
"litellm_provider": custom_llm_provider,
}
f"{custom_llm_provider}/{model}": _build_custom_pricing_entry(
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
model_info=model_info,
)
}
)
### BUILD CUSTOM PROMPT TEMPLATE -- IF GIVEN ###
@ -4644,7 +4666,6 @@ def embedding( # noqa: PLR0915
input_cost_per_token = kwargs.get("input_cost_per_token", None)
output_cost_per_token = kwargs.get("output_cost_per_token", None)
input_cost_per_second = kwargs.get("input_cost_per_second", None)
output_cost_per_second = kwargs.get("output_cost_per_second", None)
openai_params = [
"user",
"dimensions",
@ -4694,25 +4715,16 @@ def embedding( # noqa: PLR0915
)
### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ###
if input_cost_per_token is not None and output_cost_per_token is not None:
if (
input_cost_per_token is not None and output_cost_per_token is not None
) or input_cost_per_second is not None:
litellm.register_model(
{
f"{custom_llm_provider}/{model}": {
"input_cost_per_token": input_cost_per_token,
"output_cost_per_token": output_cost_per_token,
"litellm_provider": custom_llm_provider,
}
}
)
if input_cost_per_second is not None: # time based pricing just needs cost in place
output_cost_per_second = output_cost_per_second or 0.0
litellm.register_model(
{
f"{custom_llm_provider}/{model}": {
"input_cost_per_second": input_cost_per_second,
"output_cost_per_second": output_cost_per_second,
"litellm_provider": custom_llm_provider,
}
f"{custom_llm_provider}/{model}": _build_custom_pricing_entry(
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
model_info=kwargs.get("model_info"),
)
}
)

View file

@ -20716,6 +20716,40 @@
"supports_tool_choice": true,
"supports_vision": true
},
"gpt-5.3-chat-latest": {
"cache_read_input_token_cost": 1.75e-07,
"cache_read_input_token_cost_priority": 3.5e-07,
"input_cost_per_token": 1.75e-06,
"input_cost_per_token_priority": 3.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1.4e-05,
"output_cost_per_token_priority": 2.8e-05,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
},
"gpt-5.2-pro": {
"input_cost_per_token": 2.1e-05,
"litellm_provider": "openai",

View file

@ -279,7 +279,6 @@ async def common_checks( # noqa: PLR0915
request: Request,
skip_budget_checks: bool = False,
project_object: Optional[LiteLLM_ProjectTableCachedObj] = None,
skip_route_check: bool = False,
) -> bool:
"""
Common checks across jwt + key-based auth.
@ -499,21 +498,18 @@ async def common_checks( # noqa: PLR0915
user_object=user_object, route=route, request_body=request_body
)
if not skip_route_check:
token_team = getattr(valid_token, "team_id", None)
token_type: Literal["ui", "api"] = (
"ui"
if token_team is not None and token_team == "litellm-dashboard"
else "api"
)
_is_route_allowed = _is_allowed_route(
route=route,
token_type=token_type,
user_obj=user_object,
request=request,
request_data=request_body,
valid_token=valid_token,
)
token_team = getattr(valid_token, "team_id", None)
token_type: Literal["ui", "api"] = (
"ui" if token_team is not None and token_team == "litellm-dashboard" else "api"
)
_is_route_allowed = _is_allowed_route(
route=route,
token_type=token_type,
user_obj=user_object,
request=request,
request_data=request_body,
valid_token=valid_token,
)
# 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store
await vector_store_access_check(

View file

@ -1752,21 +1752,21 @@ async def _run_post_custom_auth_checks(
if _project_obj is not None:
valid_token.project_metadata = _project_obj.metadata
_ = await common_checks(
request=request,
request_body=request_data,
team_object=_team_obj,
user_object=user_object,
end_user_object=end_user_object,
general_settings=general_settings,
global_proxy_spend=None,
route=route,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
skip_budget_checks=False,
project_object=_project_obj,
skip_route_check=True,
)
if general_settings.get("custom_auth_run_common_checks", False):
_ = await common_checks(
request=request,
request_body=request_data,
team_object=_team_obj,
user_object=user_object,
end_user_object=end_user_object,
general_settings=general_settings,
global_proxy_spend=None,
route=route,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
skip_budget_checks=False,
project_object=_project_obj,
)
return valid_token

View file

@ -23,6 +23,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
decode_model_from_file_id,
encode_batch_response_ids,
encode_file_id_with_model,
get_batch_from_database,
get_credentials_for_model,
@ -127,12 +128,22 @@ async def create_batch( # noqa: PLR0915
if enforced_batch_expiry is not None:
if "anchor" not in enforced_batch_expiry or "seconds" not in enforced_batch_expiry:
raise HTTPException(
status_code=400,
status_code=500,
detail={
"error": "enforced_batch_output_expires_after must contain 'anchor' and 'seconds' keys",
"error": "Server configuration error: team metadata field 'enforced_batch_output_expires_after' is malformed - must contain 'anchor' and 'seconds' keys. Contact your team or proxy admin to fix this setting.",
},
)
_create_batch_data["output_expires_after"] = enforced_batch_expiry
if enforced_batch_expiry["anchor"] != "created_at":
raise HTTPException(
status_code=500,
detail={
"error": f"Server configuration error: team metadata field 'enforced_batch_output_expires_after' has invalid anchor '{enforced_batch_expiry['anchor']}' - must be 'created_at'. Contact your team or proxy admin to fix this setting.",
},
)
_create_batch_data["output_expires_after"] = {
"anchor": "created_at",
"seconds": int(enforced_batch_expiry["seconds"]),
}
input_file_id = _create_batch_data.get("input_file_id", None)
unified_file_id: Union[str, Literal[False]] = False
@ -258,7 +269,9 @@ async def create_batch( # noqa: PLR0915
custom_llm_provider=credentials["custom_llm_provider"],
**_create_batch_data # type: ignore
)
encode_batch_response_ids(response, model=model_param)
verbose_proxy_logger.debug(f"Created batch using model: {model_param}")
else:
# SCENARIO 3: Fallback to custom_llm_provider (uses env variables)
@ -456,8 +469,9 @@ async def retrieve_batch( # noqa: PLR0915
custom_llm_provider=credentials["custom_llm_provider"],
**data # type: ignore
)
encode_batch_response_ids(response, model=model_from_id)
verbose_proxy_logger.debug(
f"Retrieved batch using model: {model_from_id}, original_id: {original_batch_id}"
)
@ -649,7 +663,13 @@ async def list_batches(
limit=limit,
**data # type: ignore
)
# Encode batch IDs in the list response so clients can use
# them for retrieve/cancel/file downloads through the proxy.
if response and hasattr(response, "data") and response.data:
for batch in response.data:
encode_batch_response_ids(batch, model=model_param)
verbose_proxy_logger.debug(f"Listed batches using model: {model_param}")
# SCENARIO 2 (alternative): target_model_names based routing
@ -825,7 +845,9 @@ async def cancel_batch(
custom_llm_provider=credentials["custom_llm_provider"],
**data # type: ignore
)
encode_batch_response_ids(response, model=model_from_id)
verbose_proxy_logger.debug(
f"Cancelled batch using model: {model_from_id}, original_id: {original_batch_id}"
)

View file

@ -29,7 +29,7 @@ from litellm.constants import (
MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG,
STREAM_SSE_DATA_PREFIX,
)
from litellm.litellm_core_utils.dd_tracing import set_active_span_tag, tracer
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.llm_response_utils.get_headers import (
get_response_headers,
@ -41,6 +41,7 @@ from litellm.proxy.common_utils.callback_utils import (
get_logging_caching_headers,
get_remaining_tokens_and_requests_from_request_data,
)
from litellm.proxy.dd_span_tagger import DDSpanTagger
from litellm.proxy.route_llm_request import route_request
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
@ -245,26 +246,6 @@ async def create_response(
)
def _add_dd_apm_tags_for_litellm_call_id(litellm_call_id: Optional[str]) -> None:
"""
Attach LiteLLM call id to the active Datadog APM span.
This enables searching APM traces by LiteLLM call id returned in
`x-litellm-call-id`.
"""
if not litellm_call_id:
return
try:
set_active_span_tag("litellm.call_id", str(litellm_call_id))
except Exception:
# Tagging is best-effort and should never impact request processing.
verbose_proxy_logger.debug(
"Failed to tag active ddtrace span with litellm.call_id",
exc_info=True,
)
def _override_openai_response_model(
*,
response_obj: Any,
@ -518,6 +499,7 @@ class ProxyBaseLLMRequestProcessing:
"aembedding",
"aresponses",
"_arealtime",
"_aresponses_websocket",
"aget_responses",
"adelete_responses",
"acancel_responses",
@ -662,7 +644,11 @@ class ProxyBaseLLMRequestProcessing:
self.data["litellm_call_id"] = request.headers.get(
"x-litellm-call-id", str(uuid.uuid4())
)
_add_dd_apm_tags_for_litellm_call_id(self.data.get("litellm_call_id"))
DDSpanTagger.tag_call_id(self.data.get("litellm_call_id"))
DDSpanTagger.tag_request(
user_api_key_dict=user_api_key_dict,
requested_model=self.data.get("model"),
)
### AUTO STREAM USAGE TRACKING ###
# If always_include_stream_usage is enabled and this is a streaming request

View file

@ -0,0 +1,60 @@
from typing import Optional
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.dd_tracing import set_active_span_tag
from litellm.proxy._types import UserAPIKeyAuth
class DDSpanTagger:
"""Best-effort helpers for tagging the active Datadog APM span with LiteLLM request metadata."""
@staticmethod
def tag_call_id(litellm_call_id: Optional[str]) -> None:
"""
Attach LiteLLM call id to the active Datadog APM span.
This enables searching APM traces by LiteLLM call id returned in
`x-litellm-call-id`.
"""
if not litellm_call_id:
return
try:
set_active_span_tag("litellm.call_id", str(litellm_call_id))
except Exception:
verbose_proxy_logger.debug(
"Failed to tag active ddtrace span with litellm.call_id",
exc_info=True,
)
@staticmethod
def tag_request(
user_api_key_dict: UserAPIKeyAuth,
requested_model: Optional[str],
) -> None:
"""
Attach key and model tags to the active Datadog APM span.
Tags set (all best-effort, skipped when value is absent):
- ``litellm.key_alias`` human-readable alias for the API key
- ``litellm.key_hash`` hashed API key (safe to log; never the raw secret)
- ``litellm.requested_model`` model name as sent by the client
Use cases:
- Trace all requests from a specific user/key: filter by ``litellm.key_alias`` or
``litellm.key_hash``.
- Trace all requests for a specific model: filter by ``litellm.requested_model``.
Note: key_alias / key_hash are not available for unauthenticated (e.g. 401) requests.
"""
try:
if user_api_key_dict.key_alias:
set_active_span_tag("litellm.key_alias", str(user_api_key_dict.key_alias))
if user_api_key_dict.token:
set_active_span_tag("litellm.key_hash", str(user_api_key_dict.token))
if requested_model:
set_active_span_tag("litellm.requested_model", str(requested_model))
except Exception:
verbose_proxy_logger.debug(
"Failed to tag active ddtrace span with key/model tags",
exc_info=True,
)

View file

@ -1,14 +1,127 @@
from typing import TYPE_CHECKING, List, Optional
import re
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from litellm._logging import verbose_proxy_logger
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
if TYPE_CHECKING:
from litellm.types.llms.openai import AllMessageValues
# Azure Content Safety APIs have a 10,000 character limit per request.
AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH = 10000
class AzureGuardrailBase:
"""
Base class for Azure guardrails.
Provides shared initialisation (API credentials, HTTP client) and
utilities (text splitting, authenticated POST) used by all Azure
Content Safety guardrails.
"""
def __init__(
self,
api_key: str,
api_base: str,
**kwargs: Any,
):
# Forward remaining kwargs to the next class in the MRO
# (typically CustomGuardrail).
super().__init__(**kwargs)
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback
)
self.api_key = api_key
self.api_base = api_base
self.api_version: str = kwargs.get("api_version") or "2024-09-01"
async def _post_to_content_safety(
self, endpoint_path: str, request_body: Dict[str, Any]
) -> Dict[str, Any]:
"""POST to an Azure Content Safety endpoint with standard auth headers.
Args:
endpoint_path: The API action, e.g. ``"text:shieldPrompt"`` or
``"text:analyze"``.
request_body: JSON-serialisable request payload.
Returns:
Parsed JSON response dict.
"""
url = f"{self.api_base}/contentsafety/{endpoint_path}?api-version={self.api_version}"
headers = {
"Ocp-Apim-Subscription-Key": self.api_key,
"Content-Type": "application/json",
}
verbose_proxy_logger.debug(
"Azure Content Safety request [%s]: %s", endpoint_path, request_body
)
response = await self.async_handler.post(
url=url,
headers=headers,
json=request_body,
)
response_json: Dict[str, Any] = response.json()
verbose_proxy_logger.debug(
"Azure Content Safety response [%s]: %s", endpoint_path, response_json
)
return response_json
@staticmethod
def split_text_by_words(text: str, max_length: int) -> List[str]:
"""
Split text into chunks at word boundaries without breaking words.
Always returns at least one chunk. Short text ( max_length) is
returned as a single-element list so callers can use a uniform
loop without branching on length.
Args:
text: The text to split
max_length: Maximum character length of each chunk
Returns:
List of text chunks, each not exceeding max_length
"""
if len(text) <= max_length:
return [text]
# Tokenize into alternating non-whitespace and whitespace runs so
# that original newlines, tabs, and multiple spaces are preserved
# within each chunk.
tokens = re.findall(r"\S+|\s+", text)
chunks: List[str] = []
current_chunk = ""
for token in tokens:
# Would appending this token exceed the limit?
if len(current_chunk) + len(token) <= max_length:
current_chunk += token
else:
# Flush whatever we have accumulated so far
if current_chunk:
chunks.append(current_chunk)
current_chunk = ""
# Force-split any single token longer than max_length
while len(token) > max_length:
chunks.append(token[:max_length])
token = token[max_length:]
current_chunk = token
if current_chunk:
chunks.append(current_chunk)
return chunks
def get_user_prompt(self, messages: List["AllMessageValues"]) -> Optional[str]:
"""
Get the last consecutive block of messages from the user.

View file

@ -3,7 +3,7 @@
Azure Prompt Shield Native Guardrail Integrationfor LiteLLM
"""
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, cast
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, cast
from fastapi import HTTPException
@ -12,10 +12,7 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.utils import CallTypesLiteral
from .base import AzureGuardrailBase
@ -26,7 +23,6 @@ if TYPE_CHECKING:
AzurePromptShieldGuardrailResponse,
)
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
from litellm.types.utils import ModelResponse
class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrail):
@ -53,25 +49,19 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
"""Initialize Azure Prompt Shield guardrail handler."""
from litellm.types.guardrails import GuardrailEventHooks
# Initialize parent CustomGuardrail
supported_event_hooks = [
GuardrailEventHooks.pre_call,
GuardrailEventHooks.during_call,
]
# AzureGuardrailBase.__init__ stores api_key, api_base, api_version,
# async_handler and forwards the rest to CustomGuardrail.
super().__init__(
api_key=api_key,
api_base=api_base,
guardrail_name=guardrail_name,
supported_event_hooks=supported_event_hooks,
**kwargs,
)
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback
)
# Store configuration
self.api_key = api_key
self.api_base = api_base
self.api_version = kwargs.get("api_version") or "2024-09-01"
verbose_proxy_logger.debug(
f"Initialized Azure Prompt Shield Guardrail: {guardrail_name}"
@ -82,31 +72,50 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
) -> "AzurePromptShieldGuardrailResponse":
"""
Make a request to the Azure Prompt Shield API.
Long prompts are automatically split at word boundaries into chunks
that respect the Azure Content Safety 10 000-character limit. Each
chunk is analysed independently; an attack in *any* chunk raises
an HTTPException immediately.
"""
from .base import AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH
from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_prompt_shield import (
AzurePromptShieldGuardrailRequestBody,
AzurePromptShieldGuardrailResponse,
)
request_body = AzurePromptShieldGuardrailRequestBody(
documents=[], userPrompt=user_prompt
)
verbose_proxy_logger.debug(
"Azure Prompt Shield guard request: %s", request_body
)
response = await self.async_handler.post(
url=f"{self.api_base}/contentsafety/text:shieldPrompt?api-version={self.api_version}",
headers={
"Ocp-Apim-Subscription-Key": self.api_key,
"Content-Type": "application/json",
},
json=cast(dict, request_body),
chunks = self.split_text_by_words(
user_prompt, AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH
)
verbose_proxy_logger.debug(
"Azure Prompt Shield guard response: %s", response.json()
)
return AzurePromptShieldGuardrailResponse(**response.json()) # type: ignore
last_response: Optional[AzurePromptShieldGuardrailResponse] = None
for chunk in chunks:
request_body = AzurePromptShieldGuardrailRequestBody(
documents=[], userPrompt=chunk
)
response_json = await self._post_to_content_safety(
"text:shieldPrompt", cast(dict, request_body)
)
last_response = AzurePromptShieldGuardrailResponse(**response_json)
if last_response["userPromptAnalysis"].get("attackDetected"):
verbose_proxy_logger.warning(
"Azure Prompt Shield: Attack detected in chunk of length %d",
len(chunk),
)
raise HTTPException(
status_code=400,
detail={
"error": "Violated Azure Prompt Shield guardrail policy",
"detection_message": f"Attack detected: {last_response['userPromptAnalysis']}",
},
)
# chunks is always non-empty (split_text_by_words guarantees ≥1 element)
assert last_response is not None
return last_response
@log_guardrail_information
async def async_pre_call_hook(
@ -114,17 +123,7 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
user_api_key_dict: "UserAPIKeyAuth",
cache: Any,
data: Dict[str, Any],
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
],
call_type: CallTypesLiteral,
) -> Optional[Dict[str, Any]]:
"""
Pre-call hook to scan user prompts before sending to LLM.
@ -138,7 +137,7 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
new_messages: Optional[List[AllMessageValues]] = data.get("messages")
if new_messages is None:
verbose_proxy_logger.warning(
"Lakera AI: not running guardrail. No messages in data"
"Azure Prompt Shield: not running guardrail. No messages in data"
)
return data
user_prompt = self.get_user_prompt(new_messages)
@ -147,40 +146,13 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
verbose_proxy_logger.debug(
f"Azure Prompt Shield: User prompt: {user_prompt}"
)
azure_prompt_shield_response = await self.async_make_request(
await self.async_make_request(
user_prompt=user_prompt,
)
if azure_prompt_shield_response["userPromptAnalysis"].get("attackDetected"):
verbose_proxy_logger.warning("Azure Prompt Shield: Attack detected")
raise HTTPException(
status_code=400,
detail={
"error": "Violated Azure Prompt Shield guardrail policy",
"detection_message": f"Attack detected: {azure_prompt_shield_response['userPromptAnalysis']}",
},
)
else:
verbose_proxy_logger.warning("Azure Prompt Shield: No user prompt found")
return None
@log_guardrail_information
async def async_post_call_hook(
self,
data: Dict[str, Any],
user_api_key_dict: "UserAPIKeyAuth",
response: "ModelResponse",
) -> "ModelResponse":
"""
Post-call hook to scan LLM responses before returning to user.
Raises HTTPException if response should be blocked.
"""
verbose_proxy_logger.debug(
"Azure Prompt Shield: Running post-call response scan"
)
return response
@staticmethod
def get_config_model() -> Optional[Type["GuardrailConfigModel"]]:
"""

View file

@ -12,16 +12,12 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import CallTypesLiteral
from .base import AzureGuardrailBase
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import AllMessageValues
from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_text_moderation import (
AzureTextModerationGuardrailResponse,
@ -32,15 +28,15 @@ if TYPE_CHECKING:
class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardrail):
"""
LiteLLM Built-in Guardrail for Azure Content Safety Guardrail (Prompt Shield).
LiteLLM Built-in Guardrail for Azure Content Safety (Text Moderation).
This guardrail scans prompts and responses using the Azure Prompt Shield API to detect
malicious content, injection attempts, and policy violations.
This guardrail scans prompts and responses using the Azure Text Moderation API to detect
malicious content and policy violations based on severity thresholds.
Configuration:
guardrail_name: Name of the guardrail instance
api_key: Azure Prompt Shield API key
api_base: Azure Prompt Shield API endpoint
api_key: Azure Text Moderation API key
api_base: Azure Text Moderation API endpoint
default_on: Whether to enable by default
"""
@ -56,23 +52,19 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
**kwargs,
):
"""Initialize Azure Text Moderation guardrail handler."""
# Initialize parent CustomGuardrail
from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_text_moderation import (
AzureTextModerationRequestBodyOptionalParams,
)
# AzureGuardrailBase.__init__ stores api_key, api_base, api_version,
# async_handler and forwards the rest to CustomGuardrail.
super().__init__(
api_key=api_key,
api_base=api_base,
guardrail_name=guardrail_name,
**kwargs,
)
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback
)
# Store configuration
self.api_key = api_key
self.api_base = api_base
self.api_version = kwargs.get("api_version") or "2024-09-01"
self.optional_params_request_body: (
AzureTextModerationRequestBodyOptionalParams
) = {
@ -96,7 +88,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
self.severity_threshold_by_category = severity_threshold_by_category
verbose_proxy_logger.info(
f"Initialized Azure Prompt Shield Guardrail: {guardrail_name}"
f"Initialized Azure Text Moderation Guardrail: {guardrail_name}"
)
@staticmethod
@ -111,34 +103,53 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
self, text: str
) -> "AzureTextModerationGuardrailResponse":
"""
Make a request to the Azure Prompt Shield API.
Make a request to the Azure Text Moderation API.
Long texts are automatically split at word boundaries into chunks
that respect the Azure Content Safety 10 000-character limit. Each
chunk is analysed independently; a severity-threshold violation in
*any* chunk raises an HTTPException immediately.
"""
from .base import AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH
from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_text_moderation import (
AzureTextModerationGuardrailRequestBody,
AzureTextModerationGuardrailResponse,
)
request_body = AzureTextModerationGuardrailRequestBody(
text=text,
**self.optional_params_request_body,
)
verbose_proxy_logger.debug(
"Azure Text Moderation guard request: %s", request_body
chunks = self.split_text_by_words(
text, AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH
)
response = await self.async_handler.post(
url=f"{self.api_base}/contentsafety/text:analyze?api-version={self.api_version}",
headers={
"Ocp-Apim-Subscription-Key": self.api_key,
"Content-Type": "application/json",
},
json=cast(dict, request_body),
)
last_response: Optional[AzureTextModerationGuardrailResponse] = None
verbose_proxy_logger.debug(
"Azure Text Moderation guard response: %s", response.json()
)
return AzureTextModerationGuardrailResponse(**response.json()) # type: ignore
for chunk in chunks:
request_body = AzureTextModerationGuardrailRequestBody(
text=chunk,
**self.optional_params_request_body,
)
response_json = await self._post_to_content_safety(
"text:analyze", cast(dict, request_body)
)
chunk_response = AzureTextModerationGuardrailResponse(**response_json)
# For multi-chunk texts the callers only see the final response,
# so we must check every intermediate chunk here to avoid silently
# swallowing a violation that appears in an earlier chunk.
try:
self.check_severity_threshold(response=chunk_response)
except HTTPException:
verbose_proxy_logger.warning(
"Azure Text Moderation: Violation detected in chunk of length %d",
len(chunk),
)
raise
last_response = chunk_response
# chunks is always non-empty (split_text_by_words guarantees ≥1 element)
assert last_response is not None
return last_response
def check_severity_threshold(
self, response: "AzureTextModerationGuardrailResponse"
@ -207,17 +218,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
user_api_key_dict: "UserAPIKeyAuth",
cache: Any,
data: Dict[str, Any],
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
],
call_type: CallTypesLiteral,
) -> Optional[Dict[str, Any]]:
"""
Pre-call hook to scan user prompts before sending to LLM.
@ -225,13 +226,13 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
Raises HTTPException if content should be blocked.
"""
verbose_proxy_logger.info(
"Azure Prompt Shield: Running pre-call prompt scan, on call_type: %s",
"Azure Text Moderation: Running pre-call prompt scan, on call_type: %s",
call_type,
)
new_messages: Optional[List[AllMessageValues]] = data.get("messages")
if new_messages is None:
verbose_proxy_logger.warning(
"Lakera AI: not running guardrail. No messages in data"
"Azure Text Moderation: not running guardrail. No messages in data"
)
return data
user_prompt = self.get_user_prompt(new_messages)
@ -240,10 +241,9 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
verbose_proxy_logger.info(
f"Azure Text Moderation: User prompt: {user_prompt}"
)
azure_text_moderation_response = await self.async_make_request(
await self.async_make_request(
text=user_prompt,
)
self.check_severity_threshold(response=azure_text_moderation_response)
else:
verbose_proxy_logger.warning("Azure Text Moderation: No text found")
return None
@ -262,10 +262,9 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
and isinstance(response.choices[0], Choices)
):
content = response.choices[0].message.content or ""
azure_text_moderation_response = await self.async_make_request(
await self.async_make_request(
text=content,
)
self.check_severity_threshold(response=azure_text_moderation_response)
return response
async def async_post_call_streaming_hook(
@ -273,10 +272,9 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr
) -> Any:
try:
if response is not None and len(response) > 0:
azure_text_moderation_response = await self.async_make_request(
await self.async_make_request(
text=response,
)
self.check_severity_threshold(response=azure_text_moderation_response)
return response
except HTTPException as e:
import json

View file

@ -188,6 +188,7 @@ class ResponsesIDSecurity(CustomLogger):
self,
response: BaseLiteLLMOpenAIResponseObject,
user_api_key_dict: "UserAPIKeyAuth",
request_cache: Optional[dict[str, str]] = None,
) -> BaseLiteLLMOpenAIResponseObject:
# encrypt the response id using the symmetric key
# encrypt the response id, and encode the user id and response id in base64
@ -211,31 +212,41 @@ class ResponsesIDSecurity(CustomLogger):
and isinstance(response_id, str)
and response_id.startswith("resp_")
):
encrypted_response_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format(
response_id,
user_api_key_dict.user_id or "",
user_api_key_dict.team_id or "",
)
# Check request-scoped cache first (for streaming consistency)
if request_cache is not None and response_id in request_cache:
setattr(response, "id", request_cache[response_id])
else:
encrypted_response_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format(
response_id,
user_api_key_dict.user_id or "",
user_api_key_dict.team_id or "",
)
encoded_user_id_and_response_id = encrypt_value_helper(
value=encrypted_response_id
)
setattr(
response, "id", f"resp_{encoded_user_id_and_response_id}"
) # maintain the 'resp_' prefix for the responses api response id
encoded_user_id_and_response_id = encrypt_value_helper(
value=encrypted_response_id
)
encrypted_id = f"resp_{encoded_user_id_and_response_id}"
if request_cache is not None:
request_cache[response_id] = encrypted_id
setattr(response, "id", encrypted_id)
elif response_obj and isinstance(response_obj, ResponsesAPIResponse):
encrypted_response_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format(
response_obj.id,
user_api_key_dict.user_id or "",
user_api_key_dict.team_id or "",
)
encoded_user_id_and_response_id = encrypt_value_helper(
value=encrypted_response_id
)
setattr(
response_obj, "id", f"resp_{encoded_user_id_and_response_id}"
) # maintain the 'resp_' prefix for the responses api response id
# Check request-scoped cache first (for streaming consistency)
if request_cache is not None and response_obj.id in request_cache:
setattr(response_obj, "id", request_cache[response_obj.id])
else:
encrypted_response_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format(
response_obj.id,
user_api_key_dict.user_id or "",
user_api_key_dict.team_id or "",
)
encoded_user_id_and_response_id = encrypt_value_helper(
value=encrypted_response_id
)
encrypted_id = f"resp_{encoded_user_id_and_response_id}"
if request_cache is not None:
request_cache[response_obj.id] = encrypted_id
setattr(response_obj, "id", encrypted_id)
setattr(response, "response", response_obj)
return response
@ -258,7 +269,7 @@ class ResponsesIDSecurity(CustomLogger):
if isinstance(response, ResponsesAPIResponse):
response = cast(
ResponsesAPIResponse,
self._encrypt_response_id(response, user_api_key_dict),
self._encrypt_response_id(response, user_api_key_dict, request_cache=None),
)
return response
@ -267,6 +278,9 @@ class ResponsesIDSecurity(CustomLogger):
) -> AsyncGenerator[BaseLiteLLMOpenAIResponseObject, None]:
from litellm.proxy.proxy_server import general_settings
# Create a request-scoped cache for consistent encryption across streaming chunks.
request_encryption_cache: dict[str, str] = {}
async for chunk in response:
if (
isinstance(chunk, BaseLiteLLMOpenAIResponseObject)
@ -274,5 +288,5 @@ class ResponsesIDSecurity(CustomLogger):
== "/v1/responses" # only encrypt the response id for the responses api
and not general_settings.get("disable_responses_id_security", False)
):
chunk = self._encrypt_response_id(chunk, user_api_key_dict)
chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache)
yield chunk

View file

@ -26,7 +26,6 @@ from litellm.types.tool_management import (
ToolDetailResponse,
ToolInputPolicy,
ToolListResponse,
ToolOutputPolicy,
ToolPolicyOption,
ToolPolicyOptionsResponse,
ToolPolicyUpdateRequest,

View file

@ -129,6 +129,22 @@ def encode_file_id_with_model(
return f"{prefix}{encoded_b64}"
def encode_batch_response_ids(response, model: str) -> None:
"""Encode all IDs in a batch response with model routing info (in-place)."""
if not response or not hasattr(response, "id") or not response.id:
return
response.id = encode_file_id_with_model(
file_id=response.id, model=model, id_type="batch"
)
for attr in ("output_file_id", "error_file_id", "input_file_id"):
if hasattr(response, attr) and getattr(response, attr):
setattr(
response,
attr,
encode_file_id_with_model(file_id=getattr(response, attr), model=model),
)
def decode_model_from_file_id(encoded_id: str) -> Optional[str]:
"""
Extract model name from an encoded file/batch ID.

View file

@ -460,21 +460,21 @@ async def create_file( # noqa: PLR0915
if enforced_file_expiry is not None:
if "anchor" not in enforced_file_expiry or "seconds" not in enforced_file_expiry:
raise HTTPException(
status_code=400,
status_code=500,
detail={
"error": "enforced_file_expires_after must contain 'anchor' and 'seconds' keys",
"error": "Server configuration error: team metadata field 'enforced_file_expires_after' is malformed - must contain 'anchor' and 'seconds' keys. Contact your team or proxy admin to fix this setting.",
},
)
if enforced_file_expiry["anchor"] != "created_at":
raise HTTPException(
status_code=400,
status_code=500,
detail={
"error": f"enforced_file_expires_after anchor must be 'created_at', got '{enforced_file_expiry['anchor']}'",
"error": f"Server configuration error: team metadata field 'enforced_file_expires_after' has invalid anchor '{enforced_file_expiry['anchor']}' - must be 'created_at'. Contact your team or proxy admin to fix this setting.",
},
)
expires_after = FileExpiresAfter(
anchor="created_at",
seconds=enforced_file_expiry["seconds"],
seconds=int(enforced_file_expiry["seconds"]),
)
verbose_proxy_logger.debug(

View file

@ -1,14 +1,21 @@
import asyncio
import json
import time
from typing import Any, AsyncIterator, Optional, cast
from typing import Any, AsyncIterator, Dict, Optional, cast
from uuid import uuid4
import fastapi
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from starlette.websockets import WebSocket
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.auth.user_api_key_auth import (
UserAPIKeyAuth,
user_api_key_auth,
user_api_key_auth_websocket,
)
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
from litellm.types.responses.main import DeleteResponseResult
@ -904,3 +911,121 @@ async def cancel_response(
proxy_logging_obj=proxy_logging_obj,
version=version,
)
@router.websocket("/v1/responses")
@router.websocket("/responses")
async def responses_websocket_endpoint(
websocket: WebSocket,
model: str = fastapi.Query(
..., description="The model to use for the responses WebSocket session."
),
user_api_key_dict=Depends(user_api_key_auth_websocket),
):
"""
Responses API WebSocket mode endpoint.
Keeps a persistent WebSocket connection for response.create events,
enabling lower-latency agentic workflows with many tool-call round trips.
See: https://developers.openai.com/api/docs/guides/websocket-mode/
"""
from litellm.proxy.proxy_server import (
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
from litellm.proxy.route_llm_request import route_request
# Accept the WebSocket handshake
requested_protocols = [
p.strip()
for p in (websocket.headers.get("sec-websocket-protocol") or "").split(",")
if p.strip()
]
accept_kwargs: dict = {}
if requested_protocols:
accept_kwargs["subprotocol"] = requested_protocols[0]
await websocket.accept(**accept_kwargs)
data: Dict[str, Any] = {
"model": model,
"websocket": websocket,
}
# Construct a synthetic Request for pre-call processing
headers_list = list(websocket.scope.get("headers") or [])
scope: Dict[str, Any] = {
"type": "http",
"method": "POST",
"path": "/v1/responses",
"headers": headers_list,
}
request = Request(scope=scope)
request._url = websocket.url
async def return_body():
return f'{{"model": "{model}"}}'.encode()
request.body = return_body # type: ignore
# Phase 1: pre-call processing (auth, guardrails, rate limits)
base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data)
try:
(
data,
litellm_logging_obj,
) = await base_llm_response_processor.common_processing_pre_call_logic(
request=request,
general_settings=general_settings,
user_api_key_dict=user_api_key_dict,
version=version,
proxy_logging_obj=proxy_logging_obj,
proxy_config=proxy_config,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
model=model,
route_type="_aresponses_websocket",
)
except Exception as e:
verbose_proxy_logger.exception("Responses WebSocket pre-call error")
try:
await websocket.send_text(
json.dumps(
{
"type": "error",
"error": {
"type": "pre_call_error",
"message": str(e),
},
}
)
)
except Exception:
pass
await websocket.close(code=1011, reason="Pre-call error")
return
# Phase 2: route to upstream provider
try:
data["user_api_key_dict"] = user_api_key_dict
llm_call = await route_request(
data=data,
route_type="_aresponses_websocket",
llm_router=llm_router,
user_model=user_model,
)
await llm_call
except Exception:
verbose_proxy_logger.exception("Responses WebSocket error")
await websocket.close(code=1011, reason="Internal server error")

View file

@ -42,6 +42,7 @@ ROUTE_ENDPOINT_MAPPING = {
"amoderation": "/moderations",
"arerank": "/rerank",
"aresponses": "/responses",
"_aresponses_websocket": "/responses",
"alist_input_items": "/responses/{response_id}/input_items",
"aimage_edit": "/images/edits",
"acancel_responses": "/responses/{response_id}/cancel",
@ -163,6 +164,7 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin
"acreate_response_reply",
"alist_input_items",
"_arealtime", # private function for realtime API
"_aresponses_websocket", # private function for responses WebSocket mode
"aimage_edit",
"agenerate_content",
"agenerate_content_stream",

View file

@ -344,8 +344,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._pending_tool_events.append(item_done_event)
def _default_response_created_event_data(self) -> dict:
# Use cached response ID if available, otherwise generate a new one
if self._cached_response_id is None:
self._cached_response_id = f"resp_{str(uuid.uuid4())}"
response_created_event_data = {
"id": f"resp_{str(uuid.uuid4())}",
"id": self._cached_response_id,
"object": "response",
"created_at": int(time.time()),
"status": "in_progress",
@ -1074,6 +1078,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
responses_api_request=self.responses_api_request,
)
# Use the cached response ID to ensure consistency across all events
if self._cached_response_id:
responses_api_response.id = self._cached_response_id
# Encode the response ID to match non-streaming behavior
encoded_response = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id(
responses_api_response=responses_api_response,

View file

@ -51,6 +51,8 @@ if TYPE_CHECKING:
from litellm.types.llms.openai import ResponseText # type: ignore
else:
ResponseText = str # Fallback for ResponseText import
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.secret_managers.main import get_secret_str
from litellm.types.responses.main import *
from litellm.types.router import GenericLiteLLMParams
from litellm.utils import ProviderConfigManager, client
@ -182,8 +184,6 @@ async def aresponses_api_with_mcp(
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None
secret_fields = kwargs.get("secret_fields")
if secret_fields and isinstance(secret_fields, dict):
from litellm.responses.utils import ResponsesAPIRequestUtils
mcp_auth_header, mcp_server_auth_headers, _, _ = (
ResponsesAPIRequestUtils.extract_mcp_headers_from_request(
secret_fields=secret_fields, tools=tools
@ -1662,3 +1662,100 @@ def compact_responses(
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)
# ---------------------------------------------------------------------------
# Responses API WebSocket mode
# ---------------------------------------------------------------------------
def _build_litellm_metadata_for_ws(kwargs: dict) -> dict:
metadata: dict = {**(kwargs.get("litellm_metadata") or {})}
guardrails = (
(kwargs.get("metadata") or {}).get("guardrails")
or kwargs.get("guardrails")
or []
)
if guardrails:
metadata["guardrails"] = guardrails
return metadata
@client
async def _aresponses_websocket(
model: str,
websocket: Any,
api_base: Optional[str] = None,
api_key: Optional[str] = None,
timeout: Optional[float] = None,
**kwargs,
):
"""
Private function to handle the Responses API WebSocket mode.
For PROXY use only.
Resolves the LLM provider from ``model``, looks up the matching
``BaseResponsesAPIConfig``, and hands off to
``BaseLLMHTTPHandler.async_responses_websocket``.
"""
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
user = kwargs.get("user", None)
litellm_params = GenericLiteLLMParams(**kwargs)
litellm_params_dict = get_litellm_params(**kwargs)
model, _custom_llm_provider, dynamic_api_key, dynamic_api_base = (
litellm.get_llm_provider(
model=model,
api_base=api_base,
api_key=api_key,
)
)
litellm_logging_obj.update_environment_variables(
model=model,
user=user,
optional_params={},
litellm_params=litellm_params_dict,
custom_llm_provider=_custom_llm_provider,
)
responses_api_provider_config: Optional[BaseResponsesAPIConfig] = None
if _custom_llm_provider is not None:
responses_api_provider_config = (
ProviderConfigManager.get_provider_responses_api_config(
model=model,
provider=litellm.LlmProviders(_custom_llm_provider),
)
)
resolved_api_base = (
dynamic_api_base
or litellm_params.api_base
or litellm.api_base
or None
)
resolved_api_key = (
dynamic_api_key
or litellm_params.api_key
or litellm.api_key
or litellm.openai_key
or get_secret_str("OPENAI_API_KEY")
)
# Extract params that we're passing explicitly to avoid duplicates in **kwargs
remaining_kwargs = {k: v for k, v in kwargs.items() if k not in {"user_api_key_dict", "litellm_metadata"}}
await base_llm_http_handler.async_responses_websocket(
model=model,
websocket=websocket,
logging_obj=litellm_logging_obj,
responses_api_provider_config=responses_api_provider_config,
api_base=resolved_api_base,
api_key=resolved_api_key,
timeout=timeout,
user_api_key_dict=kwargs.get("user_api_key_dict"),
litellm_metadata=_build_litellm_metadata_for_ws(kwargs),
custom_llm_provider=_custom_llm_provider,
**remaining_kwargs,
)

View file

@ -269,7 +269,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
self.should_auto_execute = self._should_auto_execute_tools()
# Streaming state management
self.phase = "mcp_discovery" # mcp_discovery -> initial_response -> tool_execution -> follow_up_response -> finished
self.phase = "initial_response" # initial_response -> mcp_discovery -> tool_execution -> follow_up_response -> finished
self.finished = False
# Event queues and generation flags
@ -305,6 +305,12 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
# Mark as async iterator
self.is_async = True
# Track if we've emitted initial OpenAI lifecycle events
self.initial_events_emitted = False
# Cache the response ID to ensure consistency across all events
self._cached_response_id: Optional[str] = None
def _extract_mcp_headers_from_params(self) -> None:
"""Extract MCP headers from original request params to pass to tool calls"""
@ -388,38 +394,51 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
async def __anext__(self) -> ResponsesAPIStreamingResponse:
"""
Phase-based streaming:
1. mcp_discovery - Emit MCP discovery events
2. initial_response - Stream the first LLM response
3. tool_execution - Emit tool execution events
4. follow_up_response - Stream the follow-up response
5. finished - End iteration
1. initial_response - Stream the first LLM response (includes response.created, response.in_progress, response.output_item.added)
2. mcp_discovery - Emit MCP discovery events (after response.output_item.added)
3. continue_initial_response - Continue streaming the initial response content
4. tool_execution - Emit tool execution events
5. follow_up_response - Stream the follow-up response
6. finished - End iteration
"""
# Phase 1: MCP Discovery Events
if self.phase == "mcp_discovery":
# Generate MCP discovery events if not already done
# MCP discovery events are already generated and available
# Emit MCP discovery events
if self.mcp_discovery_events:
return self.mcp_discovery_events.pop(0)
# All MCP discovery events emitted, move to next phase
verbose_logger.debug(
"MCP discovery phase complete, transitioning to initial_response"
)
self.phase = "initial_response"
await self._create_initial_response_iterator()
# Fall through to process the initial response immediately
# Phase 2: Initial Response Stream
# Phase 1: Initial Response Stream (emit standard OpenAI events first)
if self.phase == "initial_response":
# Create the initial response iterator if not already created
if self.base_iterator is None:
await self._create_initial_response_iterator()
if self.base_iterator is None:
# LLM call failed — still emit MCP discovery events before finishing
if self.mcp_discovery_events:
self.phase = "mcp_discovery"
else:
self.phase = "finished"
raise StopAsyncIteration
if self.base_iterator:
# Check if base_iterator is actually iterable
if hasattr(self.base_iterator, "__anext__"):
try:
chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined]
# Capture the response ID from the first event to ensure consistency
if self._cached_response_id is None and hasattr(chunk, 'response'):
response_obj = getattr(chunk, 'response', None)
if response_obj and hasattr(response_obj, 'id'):
self._cached_response_id = response_obj.id
verbose_logger.debug(f"Cached response ID: {self._cached_response_id}")
# After emitting response.output_item.added, transition to MCP discovery
# Check if this is the output_item.added event
if not self.initial_events_emitted and hasattr(chunk, 'type'):
chunk_type = getattr(chunk, 'type', None)
if chunk_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED:
self.initial_events_emitted = True
# Transition to MCP discovery phase after returning this chunk
self.phase = "mcp_discovery"
return chunk
# If auto-execution is enabled, check for completed responses
if self.should_auto_execute and self._is_response_completed(
chunk
@ -454,7 +473,28 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
self.phase = "finished"
raise StopAsyncIteration
# Phase 3: Tool Execution Events
# Phase 2: MCP Discovery Events (after response.output_item.added)
if self.phase == "mcp_discovery":
# Emit MCP discovery events
if self.mcp_discovery_events:
return self.mcp_discovery_events.pop(0)
self.phase = "continue_initial_response"
# Fall through to continue processing the initial response
# Phase 3: Continue Initial Response (after MCP discovery events)
if self.phase == "continue_initial_response":
try:
return await self._process_base_iterator_chunk()
except StopAsyncIteration:
# Initial response ended, move to next phase
if self.should_auto_execute and self.collected_response:
self.phase = "tool_execution"
await self._generate_tool_execution_events()
else:
self.phase = "finished"
raise
# Phase 4: Tool Execution Events
if self.phase == "tool_execution":
# Emit any queued tool execution events
if self.tool_execution_events:
@ -464,7 +504,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
self.phase = "follow_up_response"
await self._create_follow_up_iterator()
# Phase 4: Follow-up Response Stream
# Phase 5: Follow-up Response Stream
if self.phase == "follow_up_response":
if self.follow_up_iterator:
try:
@ -476,7 +516,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
self.phase = "finished"
raise StopAsyncIteration
# Phase 5: Finished
# Phase 6: Finished
if self.phase == "finished":
raise StopAsyncIteration
@ -491,6 +531,35 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
getattr(chunk, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
)
async def _process_base_iterator_chunk(self) -> ResponsesAPIStreamingResponse:
"""
Process a chunk from the base iterator with response ID consistency enforcement.
"""
if not self.base_iterator or not hasattr(self.base_iterator, "__anext__"):
raise StopAsyncIteration
chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined]
# Ensure response ID consistency - update chunk if needed
if self._cached_response_id and hasattr(chunk, 'response'):
response_obj = getattr(chunk, 'response', None)
if response_obj and hasattr(response_obj, 'id'):
if response_obj.id != self._cached_response_id:
verbose_logger.debug(f"Updating response ID from {response_obj.id} to {self._cached_response_id}")
response_obj.id = self._cached_response_id
# If auto-execution is enabled, check for completed responses
if self.should_auto_execute and self._is_response_completed(chunk):
# Collect the response for tool execution
response_obj = getattr(chunk, "response", None)
if isinstance(response_obj, ResponsesAPIResponse):
self.collected_response = response_obj
# Move to tool execution phase after emitting this chunk
self.phase = "tool_execution"
await self._generate_tool_execution_events()
return chunk
async def _create_initial_response_iterator(self) -> None:
"""Create the initial response iterator by making the first LLM call"""
try:
@ -540,7 +609,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
traceback.print_exc()
self.base_iterator = None
self.phase = "finished"
# Don't set phase to "finished" here — let __anext__ emit any
# pre-generated MCP discovery events before ending the iteration.
async def _generate_tool_execution_events(self) -> None:
"""Generate tool execution events and execute tools"""

View file

@ -3,7 +3,7 @@ import json
import time
import traceback
from datetime import datetime
from typing import Any, Dict, Optional
from typing import Any, Dict, List, Optional
import httpx
@ -26,6 +26,7 @@ from litellm.types.llms.openai import (
OutputTextDeltaEvent,
ResponseAPIUsage,
ResponseCompletedEvent,
ResponsesAPIRequestParams,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
ResponsesAPIStreamingResponse,
@ -682,3 +683,602 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
for c in getattr(out_item, "content", []):
out += c.text
return out
# ---------------------------------------------------------------------------
# WebSocket mode streaming (bidirectional forwarding)
# ---------------------------------------------------------------------------
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.thread_pool_executor import executor as _ws_executor
RESPONSES_WS_LOGGED_EVENT_TYPES = [
"response.created",
"response.completed",
"response.failed",
"response.incomplete",
"error",
]
class ResponsesWebSocketStreaming:
"""
Manages bidirectional WebSocket forwarding for the Responses API
WebSocket mode (wss://.../v1/responses).
Unlike the Realtime API, the Responses API WebSocket mode:
- Uses response.create as the client-to-server event
- Streams back the same events as the HTTP streaming Responses API
- Supports previous_response_id for incremental continuation
- Supports generate: false for warmup
- One response at a time per connection (sequential, no multiplexing)
"""
def __init__(
self,
websocket: Any,
backend_ws: Any,
logging_obj: LiteLLMLoggingObj,
user_api_key_dict: Optional[Any] = None,
request_data: Optional[Dict] = None,
):
self.websocket = websocket
self.backend_ws = backend_ws
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.request_data: Dict = request_data or {}
self.messages: list[Dict] = []
self.input_messages: list[Dict[str, str]] = []
def _should_store_event(self, event_obj: dict) -> bool:
return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES
def _store_event(self, event: Any) -> None:
if isinstance(event, bytes):
event = event.decode("utf-8")
if isinstance(event, str):
try:
event_obj = json.loads(event)
except (json.JSONDecodeError, TypeError):
return
else:
event_obj = event
if self._should_store_event(event_obj):
self.messages.append(event_obj)
def _collect_input_from_client_event(self, message: Any) -> None:
"""Extract user input content from response.create for logging."""
try:
if isinstance(message, str):
msg_obj = json.loads(message)
elif isinstance(message, dict):
msg_obj = message
else:
return
if msg_obj.get("type") != "response.create":
return
input_items = msg_obj.get("input", [])
if isinstance(input_items, str):
self.input_messages.append({"role": "user", "content": input_items})
return
if isinstance(input_items, list):
for item in input_items:
if not isinstance(item, dict):
continue
if item.get("type") == "message" and item.get("role") == "user":
content = item.get("content", [])
if isinstance(content, str):
self.input_messages.append(
{"role": "user", "content": content}
)
elif isinstance(content, list):
for c in content:
if (
isinstance(c, dict)
and c.get("type") == "input_text"
):
text = c.get("text", "")
if text:
self.input_messages.append(
{"role": "user", "content": text}
)
except (json.JSONDecodeError, AttributeError, TypeError):
pass
def _store_input(self, message: Any) -> None:
self._collect_input_from_client_event(message)
if self.logging_obj:
self.logging_obj.pre_call(input=message, api_key="")
async def _log_messages(self) -> None:
if not self.logging_obj:
return
if self.input_messages:
self.logging_obj.model_call_details["messages"] = self.input_messages
if self.messages:
asyncio.create_task(
self.logging_obj.async_success_handler(self.messages)
)
_ws_executor.submit(self.logging_obj.success_handler, self.messages)
async def backend_to_client(self) -> None:
"""Forward events from backend WebSocket to the client."""
import websockets
try:
while True:
try:
raw_response = await self.backend_ws.recv(decode=False) # type: ignore[union-attr]
except TypeError:
raw_response = await self.backend_ws.recv() # type: ignore[union-attr, assignment]
if isinstance(raw_response, bytes):
response_str = raw_response.decode("utf-8")
else:
response_str = raw_response
self._store_event(response_str)
await self.websocket.send_text(response_str)
except websockets.exceptions.ConnectionClosed as e: # type: ignore
verbose_logger.debug(
"Responses WS backend connection closed: %s", e
)
except Exception as e:
verbose_logger.exception(
"Error in responses WS backend_to_client: %s", e
)
finally:
await self._log_messages()
async def client_to_backend(self) -> None:
"""Forward response.create events from client to backend."""
try:
while True:
message = await self.websocket.receive_text()
self._store_input(message)
self._store_event(message)
await self.backend_ws.send(message) # type: ignore[union-attr]
except Exception as e:
verbose_logger.debug("Responses WS client_to_backend ended: %s", e)
async def bidirectional_forward(self) -> None:
"""Run both forwarding directions concurrently."""
forward_task = asyncio.create_task(self.backend_to_client())
try:
await self.client_to_backend()
except Exception:
pass
finally:
if not forward_task.done():
forward_task.cancel()
try:
await forward_task
except asyncio.CancelledError:
pass
try:
await self.backend_ws.close()
except Exception:
pass
# ---------------------------------------------------------------------------
# Managed WebSocket mode (HTTP-backed, provider-agnostic)
# ---------------------------------------------------------------------------
_RESPONSE_CREATE_PARAMS: frozenset = (
ResponsesAPIRequestParams.__required_keys__ | ResponsesAPIRequestParams.__optional_keys__
)
_MANAGED_WS_SKIP_KWARGS: frozenset = frozenset(
{
"litellm_logging_obj",
"litellm_call_id",
"aresponses",
"_aresponses_websocket",
"user_api_key_dict",
}
)
class ManagedResponsesWebSocketHandler:
"""
Handles Responses API WebSocket mode for providers that do not expose a
native ``wss://`` responses endpoint.
Instead of proxying to a provider WebSocket, this handler:
- Listens for ``response.create`` events from the client
- Makes HTTP streaming calls via ``litellm.aresponses(stream=True)``
- Serialises and forwards every streaming event back over the WebSocket
- Supports ``previous_response_id`` for multi-turn conversations via
in-memory session tracking (avoids async DB-write timing issues)
- Supports sequential requests over a single persistent connection
This makes every provider that LiteLLM can reach over HTTP available on
the WebSocket transport without any provider-specific changes.
"""
def __init__(
self,
websocket: Any,
model: str,
logging_obj: "LiteLLMLoggingObj",
user_api_key_dict: Optional[Any] = None,
litellm_metadata: Optional[Dict[str, Any]] = None,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
timeout: Optional[float] = None,
custom_llm_provider: Optional[str] = None,
**kwargs: Any,
) -> None:
self.websocket = websocket
self.model = model
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.litellm_metadata: Dict[str, Any] = litellm_metadata or {}
self.api_key = api_key
self.api_base = api_base
self.timeout = timeout
self.custom_llm_provider = custom_llm_provider
# Carry through safe pass-through kwargs (e.g. extra_headers)
self.extra_kwargs: Dict[str, Any] = {
k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS
}
# In-memory session history: response_id → full accumulated message list.
# Keyed by the DECODED (pre-encoding) response ID from response.completed.
# This avoids the async DB-write race condition where spend logs haven't
# been committed yet when the next response.create arrives.
self._session_history: Dict[str, List[Dict[str, Any]]] = {}
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
@staticmethod
def _serialize_chunk(chunk: Any) -> Optional[str]:
"""Serialize a streaming chunk to a JSON string for WebSocket transmission."""
try:
if hasattr(chunk, "model_dump_json"):
return chunk.model_dump_json(exclude_none=True)
if hasattr(chunk, "model_dump"):
return json.dumps(chunk.model_dump(exclude_none=True), default=str)
if isinstance(chunk, dict):
return json.dumps(chunk, default=str)
return json.dumps(str(chunk))
except Exception as exc:
verbose_logger.debug("ManagedResponsesWS: failed to serialize chunk: %s", exc)
return None
async def _send_error(self, message: str, error_type: str = "server_error") -> None:
try:
await self.websocket.send_text(
json.dumps({"type": "error", "error": {"type": error_type, "message": message}})
)
except Exception:
pass
def _get_history_messages(self, previous_response_id: str) -> List[Dict[str, Any]]:
"""
Return accumulated message history for *previous_response_id*.
The key is the *decoded* response ID (the raw provider response ID before
LiteLLM base64-encodes it into the ``resp_...`` format).
"""
decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(
previous_response_id
)
raw_id = decoded.get("response_id", previous_response_id)
return list(self._session_history.get(raw_id, []))
def _store_history(self, response_id: str, messages: List[Dict[str, Any]]) -> None:
"""
Store the complete accumulated message history for *response_id*.
Replaces any prior value callers are responsible for passing the full
history (prior turns + current input + new output).
"""
self._session_history[response_id] = messages
@staticmethod
def _extract_response_id(completed_event: Dict[str, Any]) -> Optional[str]:
"""
Pull the raw (decoded) response ID out of a ``response.completed`` event.
Returns *None* if the event doesn't contain a usable ID.
"""
resp_obj = completed_event.get("response", {})
encoded_id: Optional[str] = resp_obj.get("id") if isinstance(resp_obj, dict) else None
if not encoded_id:
return None
decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(encoded_id)
return decoded.get("response_id", encoded_id)
@staticmethod
def _extract_output_messages(completed_event: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Convert the output items in a ``response.completed`` event into
Responses API message dicts suitable for the next turn's ``input``.
"""
resp_obj = completed_event.get("response", {})
if not isinstance(resp_obj, dict):
return []
messages: List[Dict[str, Any]] = []
for item in resp_obj.get("output", []) or []:
if not isinstance(item, dict):
continue
item_type = item.get("type")
role = item.get("role", "assistant")
if item_type == "message":
content_parts = item.get("content") or []
text_parts = [
p.get("text", "")
for p in content_parts
if isinstance(p, dict) and p.get("type") in ("output_text", "text")
]
text = "".join(text_parts)
if text:
messages.append({"type": "message", "role": role, "content": [{"type": "output_text", "text": text}]})
elif item_type == "function_call":
messages.append(item)
return messages
@staticmethod
def _input_to_messages(input_val: Any) -> List[Dict[str, Any]]:
"""
Normalise the ``input`` field of a ``response.create`` event to a list
of Responses API message dicts.
"""
if isinstance(input_val, str):
return [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": input_val}]}]
if isinstance(input_val, list):
return [item for item in input_val if isinstance(item, dict)]
return []
# ------------------------------------------------------------------
# _process_response_create sub-methods
# ------------------------------------------------------------------
async def _parse_message(self, raw_message: str) -> Optional[Dict[str, Any]]:
"""Parse raw WS text; return the message dict or None (JSON error / ignored type)."""
try:
msg_obj = json.loads(raw_message)
except json.JSONDecodeError:
await self._send_error("Invalid JSON in response.create event", "invalid_request_error")
return None
if msg_obj.get("type") != "response.create":
# Silently ignore non-response.create messages (e.g. warmup pings)
return None
return msg_obj
@staticmethod
def _build_base_call_kwargs(msg_obj: Dict[str, Any]) -> Dict[str, Any]:
"""
Extract Responses API params from the event, handling both wire formats:
Nested: {"type": "response.create", "response": {"input": [...], ...}}
Flat: {"type": "response.create", "input": [...], "model": "...", ...}
"""
nested = msg_obj.get("response")
response_params: Dict[str, Any] = (
nested
if isinstance(nested, dict) and nested
else {k: v for k, v in msg_obj.items() if k != "type"}
)
return {
param: response_params[param]
for param in _RESPONSE_CREATE_PARAMS
if param in response_params and response_params[param] is not None
}
def _apply_history(
self,
call_kwargs: Dict[str, Any],
previous_response_id: Optional[str],
current_messages: List[Dict[str, Any]],
prior_history: List[Dict[str, Any]],
) -> None:
"""Prepend in-memory turn history, or fall back to DB-based reconstruction."""
if not previous_response_id:
return
if prior_history:
call_kwargs["input"] = prior_history + current_messages
verbose_logger.debug(
"ManagedResponsesWS: prepended %d history messages for previous_response_id=%s",
len(prior_history),
previous_response_id,
)
else:
verbose_logger.debug(
"ManagedResponsesWS: no in-memory history for previous_response_id=%s; "
"falling back to DB-based session reconstruction",
previous_response_id,
)
# Fall back to DB-based session reconstruction (may work for
# cross-connection multi-turn when spend logs are committed)
call_kwargs["previous_response_id"] = previous_response_id
def _inject_credentials(
self, call_kwargs: Dict[str, Any], event_model: Optional[str]
) -> None:
"""Inject connection-level credentials and metadata into call_kwargs."""
if self.api_key is not None:
call_kwargs["api_key"] = self.api_key
if self.api_base is not None:
call_kwargs["api_base"] = self.api_base
if self.timeout is not None:
call_kwargs["timeout"] = self.timeout
# Only propagate custom_llm_provider when no per-request model override exists.
# If the payload specifies a different model, let litellm re-resolve the
# provider so we don't accidentally force the wrong backend.
if self.custom_llm_provider is not None and not event_model:
call_kwargs["custom_llm_provider"] = self.custom_llm_provider
if self.litellm_metadata:
call_kwargs["litellm_metadata"] = dict(self.litellm_metadata)
@staticmethod
def _update_proxy_request(call_kwargs: Dict[str, Any], model: str) -> None:
"""Update proxy_server_request body so spend logs record the full request."""
proxy_server_request = (call_kwargs.get("litellm_metadata") or {}).get(
"proxy_server_request"
) or {}
if not isinstance(proxy_server_request, dict):
return
body = dict(proxy_server_request.get("body") or {})
body["input"] = call_kwargs.get("input")
body["store"] = call_kwargs.get("store")
body["model"] = model
for k in ("tools", "tool_choice", "instructions", "metadata"):
if k in call_kwargs and call_kwargs[k] is not None:
body[k] = call_kwargs[k]
proxy_server_request = {**proxy_server_request, "body": body}
if "litellm_metadata" not in call_kwargs:
call_kwargs["litellm_metadata"] = {}
call_kwargs["litellm_metadata"]["proxy_server_request"] = proxy_server_request
call_kwargs.setdefault("litellm_params", {})
call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request
async def _stream_and_forward(
self, model: str, call_kwargs: Dict[str, Any]
) -> Optional[Dict[str, Any]]:
"""
Stream ``litellm.aresponses`` and forward every chunk over the WebSocket.
Captures the ``response.completed`` event type from the chunk object
directly (before serialization) to avoid a redundant JSON round-trip on
every chunk. Returns the completed event dict, or ``None``.
"""
completed_event: Optional[Dict[str, Any]] = None
stream_response = await litellm.aresponses(model=model, **call_kwargs)
async for chunk in stream_response: # type: ignore[union-attr]
if chunk is None:
continue
# Read type from the object before serializing to avoid double JSON parse
chunk_type = getattr(chunk, "type", None) or (
chunk.get("type") if isinstance(chunk, dict) else None
)
serialized = self._serialize_chunk(chunk)
if serialized is None:
continue
if chunk_type == "response.completed" and completed_event is None:
try:
completed_event = json.loads(serialized)
except Exception:
pass
try:
await self.websocket.send_text(serialized)
except Exception as send_exc:
verbose_logger.debug(
"ManagedResponsesWS: error sending chunk to client: %s", send_exc
)
return completed_event # Client disconnected
return completed_event
def _save_turn_history(
self,
completed_event: Optional[Dict[str, Any]],
prior_history: List[Dict[str, Any]],
current_messages: List[Dict[str, Any]],
) -> None:
"""Store this turn in in-memory history for future previous_response_id lookups."""
if completed_event is None:
return
new_response_id = self._extract_response_id(completed_event)
if not new_response_id:
return
output_msgs = self._extract_output_messages(completed_event)
all_messages = prior_history + current_messages + output_msgs
self._store_history(new_response_id, all_messages)
verbose_logger.debug(
"ManagedResponsesWS: stored %d messages for response_id=%s",
len(all_messages),
new_response_id,
)
# ------------------------------------------------------------------
# Core request handler
# ------------------------------------------------------------------
async def _process_response_create(self, raw_message: str) -> None:
"""
Parse one ``response.create`` event, call ``litellm.aresponses(stream=True)``,
and forward every streaming event to the client.
Multi-turn support via in-memory session history
------------------------------------------------
When ``previous_response_id`` is present in the event:
1. Look up the accumulated message history in ``self._session_history``
(keyed by the decoded provider response ID).
2. Prepend those messages to the current ``input`` so the model has full
conversation context.
3. After the stream completes, extract the new response ID and output
messages from ``response.completed`` and store them in
``self._session_history`` for the next turn.
This in-memory approach avoids the async DB-write race condition that
occurs when spend logs haven't been committed by the time the second
``response.create`` arrives over the same WebSocket connection.
"""
msg_obj = await self._parse_message(raw_message)
if msg_obj is None:
return
call_kwargs = self._build_base_call_kwargs(msg_obj)
call_kwargs["stream"] = True
event_model: Optional[str] = call_kwargs.pop("model", None)
model = event_model or self.model
previous_response_id: Optional[str] = call_kwargs.pop("previous_response_id", None)
current_messages = self._input_to_messages(call_kwargs.get("input"))
# Fetch history once; reused in both _apply_history and _save_turn_history
prior_history = (
self._get_history_messages(previous_response_id)
if previous_response_id
else []
)
self._apply_history(call_kwargs, previous_response_id, current_messages, prior_history)
self._inject_credentials(call_kwargs, event_model)
self._update_proxy_request(call_kwargs, model)
call_kwargs.update(self.extra_kwargs)
try:
completed_event = await self._stream_and_forward(model, call_kwargs)
except Exception as exc:
verbose_logger.exception(
"ManagedResponsesWS: error processing response.create: %s", exc
)
await self._send_error(str(exc))
return
self._save_turn_history(completed_event, prior_history, current_messages)
# ------------------------------------------------------------------
# Main entry point
# ------------------------------------------------------------------
async def run(self) -> None:
"""
Main loop: accept ``response.create`` events sequentially and handle
each one before waiting for the next message.
"""
try:
while True:
try:
message = await self.websocket.receive_text()
except Exception as exc:
verbose_logger.debug(
"ManagedResponsesWS: client disconnected: %s", exc
)
break
await self._process_response_create(message)
except Exception as exc:
verbose_logger.exception("ManagedResponsesWS: unexpected error: %s", exc)
await self._send_error(f"Internal server error: {exc}")

View file

@ -115,9 +115,6 @@ from litellm.router_utils.handle_error import (
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
EncryptedContentAffinityCheck,
)
from litellm.router_utils.pre_call_checks.model_rate_limit_check import (
ModelRateLimitingCheck,
)
@ -885,6 +882,9 @@ class Router:
self._arealtime = self.factory_function(
litellm._arealtime, call_type="_arealtime"
)
self._aresponses_websocket = self.factory_function(
litellm._aresponses_websocket, call_type="_aresponses_websocket"
)
self.acreate_fine_tuning_job = self.factory_function(
litellm.acreate_fine_tuning_job, call_type="acreate_fine_tuning_job"
)
@ -1848,7 +1848,7 @@ class Router:
finally:
if hasattr(model_response, "close"):
try:
model_response.close()
model_response.close() # type: ignore[reportAttributeAccessIssue]
except BaseException as close_err:
verbose_router_logger.debug(
"stream_with_fallbacks: error closing model_response: %s",
@ -4683,6 +4683,7 @@ class Router:
"afile_delete",
"afile_content",
"_arealtime",
"_aresponses_websocket",
"acreate_fine_tuning_job",
"acancel_fine_tuning_job",
"alist_fine_tuning_jobs",
@ -4855,6 +4856,7 @@ class Router:
"anthropic_messages",
"aresponses",
"_arealtime",
"_aresponses_websocket",
"acreate_fine_tuning_job",
"acancel_fine_tuning_job",
"alist_fine_tuning_jobs",
@ -7100,6 +7102,17 @@ class Router:
model_group_name=model_id
)
# If still not found, check for wildcard pattern matches
if deployment is None:
potential_wildcard_models = self.pattern_router.route(model_id) or []
if potential_wildcard_models:
# Use the first matching wildcard deployment
deployment_dict = potential_wildcard_models[0]
if isinstance(deployment_dict, dict):
deployment = Deployment(**deployment_dict)
elif isinstance(deployment_dict, Deployment):
deployment = deployment_dict
if deployment is None:
return None

View file

@ -291,6 +291,7 @@ class CallTypes(str, Enum):
search = "search"
asearch = "asearch"
arealtime = "_arealtime"
aresponses_websocket = "_aresponses_websocket"
create_batch = "create_batch"
acreate_batch = "acreate_batch"
aretrieve_batch = "aretrieve_batch"
@ -3026,6 +3027,7 @@ all_litellm_params = (
"shared_session",
"search_tool_name",
"order",
"enable_json_schema_validation",
]
+ list(StandardCallbackDynamicParams.__annotations__.keys())
+ list(CustomPricingLiteLLMParams.model_fields.keys())
@ -3236,6 +3238,7 @@ class SearchProviders(str, Enum):
SEARXNG = "searxng"
LINKUP = "linkup"
DUCKDUCKGO = "duckduckgo"
SEARCHAPI = "searchapi"
# Create a set of all search provider values for quick lookup

View file

@ -1319,7 +1319,18 @@ def post_call_processing(
### POST-CALL RULES ###
rules_obj.post_call_rules(input=model_response, model=model)
### JSON SCHEMA VALIDATION ###
if litellm.enable_json_schema_validation is True:
# Per-request flag takes priority over global flag
_per_request_validation = (
optional_params.get("enable_json_schema_validation")
if optional_params is not None
else None
)
_enable_json_schema_validation = (
_per_request_validation
if _per_request_validation is not None
else litellm.enable_json_schema_validation
)
if _enable_json_schema_validation is True:
try:
if (
optional_params is not None
@ -8794,6 +8805,12 @@ class ProviderConfigManager:
)
return BedrockStabilityImageEditConfig()
elif LlmProviders.OPENROUTER == provider:
from litellm.llms.openrouter.image_edit import (
get_openrouter_image_edit_config,
)
return get_openrouter_image_edit_config(model)
return None
@staticmethod
@ -8844,6 +8861,7 @@ class ProviderConfigManager:
ParallelAISearchConfig,
)
from litellm.llms.perplexity.search.transformation import PerplexitySearchConfig
from litellm.llms.searchapi.search.transformation import SearchAPIConfig
from litellm.llms.searxng.search.transformation import SearXNGSearchConfig
from litellm.llms.tavily.search.transformation import TavilySearchConfig
@ -8859,6 +8877,7 @@ class ProviderConfigManager:
SearchProviders.SEARXNG: SearXNGSearchConfig,
SearchProviders.LINKUP: LinkupSearchConfig,
SearchProviders.DUCKDUCKGO: DuckDuckGoSearchConfig,
SearchProviders.SEARCHAPI: SearchAPIConfig,
}
config_class = PROVIDER_TO_CONFIG_MAP.get(provider, None)
if config_class is None:

View file

@ -20716,6 +20716,40 @@
"supports_tool_choice": true,
"supports_vision": true
},
"gpt-5.3-chat-latest": {
"cache_read_input_token_cost": 1.75e-07,
"cache_read_input_token_cost_priority": 3.5e-07,
"input_cost_per_token": 1.75e-06,
"input_cost_per_token_priority": 3.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1.4e-05,
"output_cost_per_token_priority": 2.8e-05,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
},
"gpt-5.2-pro": {
"input_cost_per_token": 2.1e-05,
"litellm_provider": "openai",

80
poetry.lock generated
View file

@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand.
# This file is automatically @generated by Poetry 2.3.2 and should not be changed by hand.
[[package]]
name = "a2a-sdk"
@ -7,11 +7,11 @@ description = "A2A Python SDK"
optional = false
python-versions = ">=3.10"
groups = ["main", "proxy-dev"]
markers = "python_version >= \"3.10\""
files = [
{file = "a2a_sdk-0.3.22-py3-none-any.whl", hash = "sha256:b98701135bb90b0ff85d35f31533b6b7a299bf810658c1c65f3814a6c15ea385"},
{file = "a2a_sdk-0.3.22.tar.gz", hash = "sha256:77a5694bfc4f26679c11b70c7f1062522206d430b34bc1215cfbb1eba67b7e7d"},
]
markers = {main = "python_version >= \"3.10\" and extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
[package.dependencies]
google-api-core = ">=1.26.0"
@ -385,6 +385,7 @@ files = [
{file = "azure_core-1.36.0-py3-none-any.whl", hash = "sha256:fee9923a3a753e94a259563429f3644aaf05c486d45b1215d098115102d91d3b"},
{file = "azure_core-1.36.0.tar.gz", hash = "sha256:22e5605e6d0bf1d229726af56d9e92bc37b6e726b141a18be0b4d424131741b7"},
]
markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""}
[package.dependencies]
requests = ">=2.21.0"
@ -405,6 +406,7 @@ files = [
{file = "azure_identity-1.25.1-py3-none-any.whl", hash = "sha256:e9edd720af03dff020223cd269fa3a61e8f345ea75443858273bcb44844ab651"},
{file = "azure_identity-1.25.1.tar.gz", hash = "sha256:87ca8328883de6036443e1c37b40e8dc8fb74898240f61071e09d2e369361456"},
]
markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""}
[package.dependencies]
azure-core = ">=1.31.0"
@ -598,7 +600,7 @@ files = [
{file = "cachetools-6.2.2-py3-none-any.whl", hash = "sha256:6c09c98183bf58560c97b2abfcedcbaf6a896a490f534b031b661d3723b45ace"},
{file = "cachetools-6.2.2.tar.gz", hash = "sha256:8e6d266b25e539df852251cfd6f990b4bc3a141db73b939058d809ebd2590fc6"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
[[package]]
name = "certifi"
@ -705,7 +707,7 @@ files = [
{file = "cffi-2.0.0-cp39-cp39-win_amd64.whl", hash = "sha256:b882b3df248017dba09d6b16defe9b5c407fe32fc7c65a9c69798e6175601be9"},
{file = "cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529"},
]
markers = {main = "platform_python_implementation != \"PyPy\" or extra == \"proxy\"", dev = "platform_python_implementation != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\""}
markers = {main = "(platform_python_implementation != \"PyPy\" or extra == \"proxy\") and (python_version >= \"3.10\" or extra == \"proxy\" or extra == \"extra-proxy\") and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"mlflow\")", dev = "platform_python_implementation != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\""}
[package.dependencies]
pycparser = {version = "*", markers = "implementation_name != \"PyPy\""}
@ -1055,6 +1057,7 @@ files = [
{file = "cryptography-43.0.3-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:2ce6fae5bdad59577b44e4dfed356944fbf1d925269114c28be377692643b4ff"},
{file = "cryptography-43.0.3.tar.gz", hash = "sha256:315b9001266a492a6ff443b61238f956b214dbec9910a081ba5b6646a055a805"},
]
markers = {main = "python_version >= \"3.10\" and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"mlflow\") or extra == \"proxy\" or extra == \"extra-proxy\""}
[package.dependencies]
cffi = {version = ">=1.12", markers = "platform_python_implementation != \"PyPy\""}
@ -1837,11 +1840,11 @@ description = "Google API client core library"
optional = false
python-versions = ">=3.7"
groups = ["main", "proxy-dev"]
markers = "python_version >= \"3.14\""
files = [
{file = "google_api_core-2.25.2-py3-none-any.whl", hash = "sha256:e9a8f62d363dc8424a8497f4c2a47d6bcda6c16514c935629c257ab5d10210e7"},
{file = "google_api_core-2.25.2.tar.gz", hash = "sha256:1c63aa6af0d0d5e37966f157a77f9396d820fba59f9e43e9415bc3dc5baff300"},
]
markers = {main = "python_version >= \"3.14\" and (extra == \"extra-proxy\" or extra == \"google\")", proxy-dev = "python_version >= \"3.14\""}
[package.dependencies]
google-auth = ">=2.14.1,<3.0.0"
@ -1869,7 +1872,7 @@ files = [
{file = "google_api_core-2.28.1-py3-none-any.whl", hash = "sha256:4021b0f8ceb77a6fb4de6fde4502cecab45062e66ff4f2895169e0b35bc9466c"},
{file = "google_api_core-2.28.1.tar.gz", hash = "sha256:2b405df02d68e68ce0fbc138559e6036559e685159d148ae5861013dc201baf8"},
]
markers = {main = "(python_version >= \"3.10\" or extra == \"google\" or extra == \"extra-proxy\") and python_version < \"3.14\"", proxy-dev = "python_version >= \"3.10\" and python_version < \"3.14\""}
markers = {main = "python_version < \"3.14\" and (extra == \"extra-proxy\" or extra == \"google\")", proxy-dev = "python_version >= \"3.10\" and python_version < \"3.14\""}
[package.dependencies]
google-auth = ">=2.14.1,<3.0.0"
@ -1906,7 +1909,7 @@ files = [
{file = "google_auth-2.43.0-py2.py3-none-any.whl", hash = "sha256:af628ba6fa493f75c7e9dbe9373d148ca9f4399b5ea29976519e0a3848eddd16"},
{file = "google_auth-2.43.0.tar.gz", hash = "sha256:88228eee5fc21b62a1b5fe773ca15e67778cb07dc8363adcb4a8827b52d81483"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
[package.dependencies]
cachetools = ">=2.0.0,<7.0"
@ -2078,11 +2081,11 @@ files = [
]
[package.dependencies]
google-api-core = {version = ">=1.34.1,<2.0.dev0 || >=2.11.dev0,<3.0.0dev", extras = ["grpc"]}
google-auth = ">=2.14.1,<2.24.0 || >2.24.0,<2.25.0 || >2.25.0,<3.0.0dev"
grpc-google-iam-v1 = ">=0.12.4,<1.0.0dev"
proto-plus = ">=1.22.3,<2.0.0dev"
protobuf = ">=3.20.2,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<6.0.0dev"
google-api-core = {version = ">=1.34.1,<2.0.dev0 || >=2.11.dev0,<3.0.0.dev0", extras = ["grpc"]}
google-auth = ">=2.14.1,<2.24.0 || >2.24.0,<2.25.0 || >2.25.0,<3.0.0.dev0"
grpc-google-iam-v1 = ">=0.12.4,<1.0.0.dev0"
proto-plus = ">=1.22.3,<2.0.0.dev0"
protobuf = ">=3.20.2,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<6.0.0.dev0"
[[package]]
name = "google-cloud-resource-manager"
@ -2264,7 +2267,7 @@ files = [
{file = "googleapis_common_protos-1.72.0-py3-none-any.whl", hash = "sha256:4299c5a82d5ae1a9702ada957347726b167f9f8d1fc352477702a1e851ff4038"},
{file = "googleapis_common_protos-1.72.0.tar.gz", hash = "sha256:e55a601c1b32b52d7a3e65f43563e2aa61bcd737998ee672ac9b951cd49319f5"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\") or extra == \"google\" or extra == \"extra-proxy\""}
[package.dependencies]
grpcio = {version = ">=1.44.0,<2.0.0", optional = true, markers = "extra == \"grpc\""}
@ -2673,11 +2676,11 @@ description = "Consume Server-Sent Event (SSE) messages with HTTPX."
optional = false
python-versions = ">=3.9"
groups = ["main", "proxy-dev"]
markers = "python_version >= \"3.10\""
files = [
{file = "httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc"},
{file = "httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d"},
]
markers = {main = "python_version >= \"3.10\" and (extra == \"proxy\" or extra == \"extra-proxy\")", proxy-dev = "python_version >= \"3.10\""}
[[package]]
name = "huey"
@ -3042,7 +3045,7 @@ files = [
[package.dependencies]
attrs = ">=22.2.0"
jsonschema-specifications = ">=2023.03.6"
jsonschema-specifications = ">=2023.3.6"
referencing = ">=0.28.4"
rpds-py = ">=0.7.1"
@ -3713,6 +3716,7 @@ files = [
{file = "msal-1.34.0-py3-none-any.whl", hash = "sha256:f669b1644e4950115da7a176441b0e13ec2975c29528d8b9e81316023676d6e1"},
{file = "msal-1.34.0.tar.gz", hash = "sha256:76ba83b716ea5a6d75b0279c0ac353a0e05b820ca1f6682c0eb7f45190c43c2f"},
]
markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""}
[package.dependencies]
cryptography = ">=2.5,<49"
@ -3733,6 +3737,7 @@ files = [
{file = "msal_extensions-1.3.1-py3-none-any.whl", hash = "sha256:96d3de4d034504e969ac5e85bae8106c8373b5c6568e4c8fa7af2eca9dbe6bca"},
{file = "msal_extensions-1.3.1.tar.gz", hash = "sha256:c5b0fd10f65ef62b5f1d62f4251d51cbcaf003fcedae8c91b040a488614be1a4"},
]
markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""}
[package.dependencies]
msal = ">=1.29,<2"
@ -3983,6 +3988,7 @@ files = [
{file = "nodeenv-1.9.1-py2.py3-none-any.whl", hash = "sha256:ba11c9782d29c27c70ffbdda2d7415098754709be8a7056d79a737cd901155c9"},
{file = "nodeenv-1.9.1.tar.gz", hash = "sha256:6ec12890a2dab7946721edbfbcd91f3319c6ccc9aec47be7c7e6b7011ee6645f"},
]
markers = {main = "extra == \"extra-proxy\""}
[[package]]
name = "numpy"
@ -4105,7 +4111,7 @@ files = [
{file = "opentelemetry_api-1.39.1-py3-none-any.whl", hash = "sha256:2edd8463432a7f8443edce90972169b195e7d6a05500cd29e6d13898187c9950"},
{file = "opentelemetry_api-1.39.1.tar.gz", hash = "sha256:fbde8c80e1b937a2c61f20347e91c0c18a1940cecf012d62e65a7caf08967c9c"},
]
markers = {main = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""}
[package.dependencies]
importlib-metadata = ">=6.0,<8.8.0"
@ -4220,7 +4226,7 @@ files = [
{file = "opentelemetry_sdk-1.39.1-py3-none-any.whl", hash = "sha256:4d5482c478513ecb0a5d938dcc61394e647066e0cc2676bee9f3af3f3f45f01c"},
{file = "opentelemetry_sdk-1.39.1.tar.gz", hash = "sha256:cf4d4563caf7bff906c9f7967e2be22d0d6b349b908be0d90fb21c8e9c995cc6"},
]
markers = {main = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""}
[package.dependencies]
opentelemetry-api = "1.39.1"
@ -4238,7 +4244,7 @@ files = [
{file = "opentelemetry_semantic_conventions-0.60b1-py3-none-any.whl", hash = "sha256:9fa8c8b0c110da289809292b0591220d3a7b53c1526a23021e977d68597893fb"},
{file = "opentelemetry_semantic_conventions-0.60b1.tar.gz", hash = "sha256:87c228b5a0669b748c76d76df6c364c369c28f1c465e50f661e39737e84bc953"},
]
markers = {main = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""}
[package.dependencies]
opentelemetry-api = "1.39.1"
@ -4722,6 +4728,7 @@ files = [
{file = "prisma-0.11.0-py3-none-any.whl", hash = "sha256:22bb869e59a2968b99f3483bb417717273ffbc569fd1e9ceed95e5614cbaf53a"},
{file = "prisma-0.11.0.tar.gz", hash = "sha256:3f2f2fd2361e1ec5ff655f2a04c7860c2f2a5bc4c91f78ca9c5c6349735bf693"},
]
markers = {main = "extra == \"extra-proxy\""}
[package.dependencies]
click = ">=7.1.2"
@ -4895,7 +4902,7 @@ files = [
{file = "proto_plus-1.26.1-py3-none-any.whl", hash = "sha256:13285478c2dcf2abb829db158e1047e2f1e8d63a077d94263c2b88b043c75a66"},
{file = "proto_plus-1.26.1.tar.gz", hash = "sha256:21a515a4c4c0088a773899e23c7bbade3d18f9c66c73edd4c7ee3816bc96a012"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
markers = {main = "extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
[package.dependencies]
protobuf = ">=3.19.0,<7.0.0"
@ -4923,7 +4930,7 @@ files = [
{file = "protobuf-5.29.5-py3-none-any.whl", hash = "sha256:6cf42630262c59b2d8de33954443d94b746c952b01434fc58a417fdbd2e84bd5"},
{file = "protobuf-5.29.5.tar.gz", hash = "sha256:bc1463bafd4b0929216c35f437a8e28731a2b7fe3d98bb77a600efced5a15c84"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\""}
[[package]]
name = "psutil"
@ -5083,7 +5090,7 @@ files = [
{file = "pyasn1-0.6.1-py3-none-any.whl", hash = "sha256:0d632f46f2ba09143da3a8afe9e33fb6f92fa2320ab7e886e2d0f7672af84629"},
{file = "pyasn1-0.6.1.tar.gz", hash = "sha256:6f580d2bdd84365380830acf45550f2511469f673cb4a5ae3857a3170128b034"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
[[package]]
name = "pyasn1-modules"
@ -5096,7 +5103,7 @@ files = [
{file = "pyasn1_modules-0.4.2-py3-none-any.whl", hash = "sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a"},
{file = "pyasn1_modules-0.4.2.tar.gz", hash = "sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
[package.dependencies]
pyasn1 = ">=0.6.1,<0.7.0"
@ -5124,7 +5131,7 @@ files = [
{file = "pycparser-2.23-py3-none-any.whl", hash = "sha256:e5c6e8d3fbad53479cab09ac03729e0a9faf2bee3db8208a550daf5af81a5934"},
{file = "pycparser-2.23.tar.gz", hash = "sha256:78816d4f24add8f10a06d6f05b4d424ad9e96cfebf68a4ddc99c65c0720d00c2"},
]
markers = {main = "implementation_name != \"PyPy\" and (platform_python_implementation != \"PyPy\" or extra == \"proxy\")", dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\""}
markers = {main = "implementation_name != \"PyPy\" and (platform_python_implementation != \"PyPy\" or extra == \"proxy\") and (python_version >= \"3.10\" or extra == \"proxy\" or extra == \"extra-proxy\") and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"mlflow\")", dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\""}
[[package]]
name = "pydantic"
@ -5347,6 +5354,7 @@ files = [
{file = "PyJWT-2.10.1-py3-none-any.whl", hash = "sha256:dcdd193e30abefd5debf142f9adfcdd2b58004e644f25406ffaebd50bd98dacb"},
{file = "pyjwt-2.10.1.tar.gz", hash = "sha256:3cc5772eb20009233caf06e9d8a0577824723b44e6648ee0a2aedb6cf9381953"},
]
markers = {main = "(python_version <= \"3.13\" or extra == \"proxy\" or extra == \"extra-proxy\") and (extra == \"extra-proxy\" or extra == \"proxy\")"}
[package.dependencies]
cryptography = {version = ">=3.4.0", optional = true, markers = "extra == \"crypto\""}
@ -5599,6 +5607,19 @@ files = [
[package.extras]
cli = ["click (>=5.0)"]
[[package]]
name = "python-multipart"
version = "0.0.20"
description = "A streaming multipart parser for Python"
optional = true
python-versions = ">=3.8"
groups = ["main"]
markers = "python_version == \"3.9\" and extra == \"proxy\""
files = [
{file = "python_multipart-0.0.20-py3-none-any.whl", hash = "sha256:8a62d3a8335e06589fe01f2a3e178cdcc632f3fbe0d492ad9ee0ec35aab1f104"},
{file = "python_multipart-0.0.20.tar.gz", hash = "sha256:8dd0cab45b8e23064ae09147625994d090fa46f5b0d1e13af944c331a7fa9d13"},
]
[[package]]
name = "python-multipart"
version = "0.0.22"
@ -6276,7 +6297,7 @@ files = [
{file = "rsa-4.9.1-py3-none-any.whl", hash = "sha256:68635866661c6836b8d39430f97a996acbd61bfa49406748ea243539fe239762"},
{file = "rsa-4.9.1.tar.gz", hash = "sha256:e7bdbfdb5497da4c07dfd35530e1a902659db6ff241e39d9953cad06ebd0ae75"},
]
markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""}
markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""}
[package.dependencies]
pyasn1 = ">=0.1.3"
@ -6322,10 +6343,10 @@ files = [
]
[package.dependencies]
botocore = ">=1.37.4,<2.0a.0"
botocore = ">=1.37.4,<2.0a0"
[package.extras]
crt = ["botocore[crt] (>=1.37.4,<2.0a.0)"]
crt = ["botocore[crt] (>=1.37.4,<2.0a0)"]
[[package]]
name = "scikit-learn"
@ -6478,9 +6499,9 @@ tornado = ">=6.4.2,<7"
urllib3 = ">=1.26,<3"
[package.extras]
all = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)", "cohere (>=5.9.4,<6.00)", "dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\"", "google-cloud-aiplatform (>=1.45.0,<2)", "ipykernel (>=6.25.0,<7)", "llama-cpp-python (>=0.2.28,<0.2.86) ; python_version < \"3.13\"", "mistralai (>=0.0.12,<0.1.0)", "mypy (>=1.7.1,<2)", "ollama (>=0.1.7)", "pillow (>=10.2.0,<11.0.0) ; python_version < \"3.13\"", "pinecone[asyncio] (>=7.0.0,<8.0.0)", "psycopg[binary] (>=3.1.0,<4)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "qdrant-client (>=1.11.1,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "sentence-transformers (>=5.0.0) ; python_version < \"3.13\"", "tokenizers (>=0.19) ; python_version < \"3.13\"", "torch (>=2.6.0) ; python_version < \"3.13\"", "torchvision (>=0.17.0) ; python_version < \"3.13\"", "transformers (>=4.36.2) ; python_version < \"3.13\"", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"]
all = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)", "cohere (>=5.9.4,<6.0)", "dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\"", "google-cloud-aiplatform (>=1.45.0,<2)", "ipykernel (>=6.25.0,<7)", "llama-cpp-python (>=0.2.28,<0.2.86) ; python_version < \"3.13\"", "mistralai (>=0.0.12,<0.1.0)", "mypy (>=1.7.1,<2)", "ollama (>=0.1.7)", "pillow (>=10.2.0,<11.0.0) ; python_version < \"3.13\"", "pinecone[asyncio] (>=7.0.0,<8.0.0)", "psycopg[binary] (>=3.1.0,<4)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "qdrant-client (>=1.11.1,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "sentence-transformers (>=5.0.0) ; python_version < \"3.13\"", "tokenizers (>=0.19) ; python_version < \"3.13\"", "torch (>=2.6.0) ; python_version < \"3.13\"", "torchvision (>=0.17.0) ; python_version < \"3.13\"", "transformers (>=4.36.2) ; python_version < \"3.13\"", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"]
bedrock = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)"]
cohere = ["cohere (>=5.9.4,<6.00)"]
cohere = ["cohere (>=5.9.4,<6.0)"]
dev = ["dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "ipykernel (>=6.25.0,<7)", "mypy (>=1.7.1,<2)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"]
docs = ["pydoc-markdown (>=4.8.2) ; python_version < \"3.12\""]
fastembed = ["fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\""]
@ -7208,6 +7229,7 @@ files = [
{file = "tomlkit-0.13.3-py3-none-any.whl", hash = "sha256:c89c649d79ee40629a9fda55f8ace8c6a1b42deb912b2a8fd8d942ddadb606b0"},
{file = "tomlkit-0.13.3.tar.gz", hash = "sha256:430cf247ee57df2b94ee3fbe588e71d362a941ebb545dec29b53961d61add2a1"},
]
markers = {main = "extra == \"extra-proxy\""}
[[package]]
name = "tornado"
@ -7980,4 +8002,4 @@ utils = ["numpydoc"]
[metadata]
lock-version = "2.1"
python-versions = ">=3.9,<4.0"
content-hash = "5ae4b43dfe73be01d71f757227eb22245d18c06b5b4d5989b014500f400f1ee9"
content-hash = "70ec9abe5b06e7e81a2d76305cb950eea79692ae40321bac3285dc63fcbcf059"

View file

@ -45,7 +45,7 @@ orjson = {version = "^3.9.7", optional = true}
apscheduler = {version = "^3.10.4", optional = true}
fastapi-sso = { version = "^0.16.0", optional = true }
PyJWT = { version = "^2.10.1", optional = true, python = ">=3.9" }
python-multipart = { version = "^0.0.22", optional = true, python = ">=3.10"}
python-multipart = { version = ">=0.0.20", optional = true}
cryptography = {version = "*", optional = true}
prisma = {version = "0.11.0", optional = true}
azure-identity = {version = "^1.15.0", optional = true, python = ">=3.9"}

View file

@ -41,7 +41,7 @@ polars==1.31.0 # for data processing
apscheduler==3.10.4 # for resetting budget in background
fastapi-sso==0.19.0 # admin UI, SSO
pyjwt[crypto]==2.10.1 ; python_version >= "3.9"
python-multipart==0.0.22 # admin UI
python-multipart>=0.0.20 # admin UI
jaraco.context>=6.1.0
azure-ai-contentsafety==1.0.0 # for azure content safety
azure-identity==1.16.1 ; python_version >= "3.9" # for azure content safety

View file

@ -16,3 +16,4 @@ exclude = ["litellm/types/*", "litellm/__init__.py", "litellm/proxy/example_conf
"litellm/proxy/utils.py" = ["F401", "PLR0915"]
"litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py" = ["PLR0915"]
"litellm/proxy/guardrails/guardrail_hooks/guardrail_benchmarks/test_eval.py" = ["PLR0915"]
"litellm/responses/streaming_iterator.py" = ["PLR0915"]

View file

@ -17,6 +17,7 @@ SEARCH_PROVIDERS = [
"searxng",
"linkup",
"duckduckgo",
"searchapi",
]
ALLOWED_FILES_IN_LLMS_FOLDER = [

View file

@ -114,7 +114,7 @@ apscheduler: >=3.10.4 # Unknown license
fastapi-sso: >=0.16.0 # Unknown license
filelock: >=3.20.0 # Unlicense (public domain) - https://unlicense.org / https://github.com/tox-dev/filelock
pyjwt: >=2.9.0 # Unknown license
python-multipart: >=0.0.18 # Unknown license
python-multipart: >=0.0.20 # Unknown license
pillow: >=11.0.0 # Unknown license
azure-ai-contentsafety: >=1.0.0 # Unknown license
azure-identity: >=1.16.1 # Unknown license

View file

@ -20,7 +20,7 @@
"apscheduler:3.10.4": "MIT",
"fastapi-sso:0.16.0": "MIT",
"pyjwt:2.9.0": "MIT",
"python-multipart:0.0.22": "Apache-2.0",
"python-multipart:0.0.20": "Apache-2.0",
"Pillow:11.0.0": "MIT-CMU",
"azure-ai-contentsafety:1.0.0": "MIT License",
"azure-identity:1.16.1": "MIT License",

View file

@ -62,3 +62,74 @@ def test_helicone_vertex_ai_via_custom_llm_provider():
for model, custom_llm_provider in test_cases:
is_vertex_ai = custom_llm_provider == "vertex_ai" or model.startswith("vertex_ai/")
assert is_vertex_ai, f"{model} with custom_llm_provider={custom_llm_provider} should be recognized as vertex_ai"
def test_helicone_vertex_gemini_gets_vertex_provider_url():
"""
Test that vertex_ai/gemini-* models route to aiplatform.googleapis.com,
not generativelanguage.googleapis.com.
This verifies the branch ordering fix: is_vertex_ai must be checked
before "gemini" in model, otherwise vertex gemini models get the wrong
provider_url.
"""
from unittest.mock import MagicMock, patch
from litellm.integrations.helicone import HeliconeLogger
logger = HeliconeLogger()
captured = {}
def mock_post(url, **kwargs):
captured["url"] = url
captured["data"] = kwargs.get("json", {})
mock_resp = MagicMock()
mock_resp.status_code = 200
return mock_resp
test_cases = [
# (model, custom_llm_provider, expected_provider_url)
(
"vertex_ai/gemini-1.5-pro",
"",
"https://aiplatform.googleapis.com/v1",
),
(
"gemini-2.0-flash",
"vertex_ai",
"https://aiplatform.googleapis.com/v1",
),
(
"gemini-1.5-flash",
"",
"https://generativelanguage.googleapis.com/v1beta",
),
]
for model, custom_llm_provider, expected_url in test_cases:
captured.clear()
mock_client = MagicMock()
mock_client.post = mock_post
with patch("litellm.module_level_client", mock_client):
logger.log_success(
model=model,
messages=[{"role": "user", "content": "test"}],
response_obj={"choices": [{"message": {"content": "hi"}}]},
start_time=MagicMock(),
end_time=MagicMock(),
print_verbose=lambda *args, **kwargs: None,
kwargs={
"litellm_params": {
"custom_llm_provider": custom_llm_provider,
"metadata": {},
},
},
)
assert "data" in captured, f"No request captured for {model}"
actual_url = captured["data"]["providerRequest"]["url"]
assert actual_url == expected_url, (
f"Model {model} (provider={custom_llm_provider!r}): "
f"expected provider_url={expected_url}, got {actual_url}"
)

View file

@ -0,0 +1,136 @@
"""
Tests for per-request enable_json_schema_validation parameter.
Ensures the per-request flag overrides the global litellm.enable_json_schema_validation,
making JSON schema validation thread-safe for concurrent usage.
Related issue: https://github.com/BerriAI/litellm/issues/XXXX
"""
import json
import pytest
import litellm
from litellm.types.utils import ModelResponse
from litellm.utils import Rules, post_call_processing
def _make_response(content: dict) -> ModelResponse:
"""Create a ModelResponse with the given content as JSON string."""
response = ModelResponse()
response.choices[0].message.content = json.dumps(content)
return response
def _mock_completion():
"""Mock function with __name__ == 'completion' for post_call_processing."""
pass
_mock_completion.__name__ = "completion"
# Schema that requires 'title' (string) and 'rating' (integer)
STRICT_SCHEMA = {
"type": "json_schema",
"json_schema": {
"name": "MovieReview",
"schema": {
"type": "object",
"properties": {
"title": {"type": "string"},
"rating": {"type": "integer"},
},
"required": ["title", "rating"],
},
},
}
INVALID_CONTENT = {"name": "test", "age": 25} # Does NOT match the schema
VALID_CONTENT = {"title": "Inception", "rating": 9} # Matches the schema
@pytest.fixture(autouse=True)
def _reset_global_flag():
"""Reset the global flag before and after each test."""
original = litellm.enable_json_schema_validation
litellm.enable_json_schema_validation = False
yield
litellm.enable_json_schema_validation = original
class TestPerRequestJsonSchemaValidation:
"""Test that per-request enable_json_schema_validation overrides the global flag."""
def test_global_off_no_per_request_skips_validation(self):
"""Global OFF + no per-request flag -> no validation (default behavior)."""
litellm.enable_json_schema_validation = False
# Should NOT raise even though response doesn't match schema
post_call_processing(
_make_response(INVALID_CONTENT),
"test-model",
{"response_format": STRICT_SCHEMA},
_mock_completion,
Rules(),
)
def test_per_request_on_overrides_global_off(self):
"""Global OFF + per-request ON -> validation runs and catches invalid response."""
litellm.enable_json_schema_validation = False
with pytest.raises(litellm.JSONSchemaValidationError):
post_call_processing(
_make_response(INVALID_CONTENT),
"test-model",
{
"response_format": STRICT_SCHEMA,
"enable_json_schema_validation": True,
},
_mock_completion,
Rules(),
)
def test_per_request_off_overrides_global_on(self):
"""Global ON + per-request OFF -> validation skipped (per-request wins)."""
litellm.enable_json_schema_validation = True
# Should NOT raise because per-request says False
post_call_processing(
_make_response(INVALID_CONTENT),
"test-model",
{
"response_format": STRICT_SCHEMA,
"enable_json_schema_validation": False,
},
_mock_completion,
Rules(),
)
def test_global_on_no_per_request_validates(self):
"""Global ON + no per-request flag -> validation runs (backward compatible)."""
litellm.enable_json_schema_validation = True
with pytest.raises(litellm.JSONSchemaValidationError):
post_call_processing(
_make_response(INVALID_CONTENT),
"test-model",
{"response_format": STRICT_SCHEMA},
_mock_completion,
Rules(),
)
def test_valid_response_passes_with_per_request_on(self):
"""Per-request ON + valid response -> no error raised."""
post_call_processing(
_make_response(VALID_CONTENT),
"test-model",
{
"response_format": STRICT_SCHEMA,
"enable_json_schema_validation": True,
},
_mock_completion,
Rules(),
)
def test_per_request_flag_is_in_all_litellm_params(self):
"""Ensure the param is registered so it doesn't leak to provider APIs."""
from litellm.types.utils import all_litellm_params
assert "enable_json_schema_validation" in all_litellm_params

View file

@ -0,0 +1,362 @@
"""
Unit tests for batch ID encoding when x-litellm-model header is used.
Verifies that create_batch encodes response IDs with model info so that
retrieve_batch can route back to the correct provider/credentials.
"""
from typing import Optional
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import litellm
from litellm.proxy.openai_files_endpoints.common_utils import (
decode_model_from_file_id,
get_original_file_id,
)
from litellm.types.utils import LiteLLMBatch
def _make_mock_request(headers: dict) -> MagicMock:
"""Create a mock FastAPI Request with the given headers."""
mock_request = MagicMock()
mock_request.headers = headers
mock_request.query_params = {}
mock_request.url = MagicMock()
mock_request.url.port = 4000
mock_request.method = "POST"
mock_request.url.path = "/v1/batches"
return mock_request
def _make_batch_response(
batch_id: str = "batch_abc123",
input_file_id: str = "file-input456",
output_file_id: Optional[str] = None,
error_file_id: Optional[str] = None,
status: str = "validating",
) -> LiteLLMBatch:
"""Create a mock LiteLLMBatch response from a provider."""
return LiteLLMBatch(
id=batch_id,
object="batch",
status=status,
endpoint="/v1/chat/completions",
input_file_id=input_file_id,
completion_window="24h",
created_at=1234567890,
output_file_id=output_file_id,
error_file_id=error_file_id,
)
@pytest.mark.asyncio
async def test_create_batch_with_x_litellm_model_encodes_batch_id():
"""
When x-litellm-model header is provided, create_batch should encode the
response batch_id with model info so retrieve_batch can route correctly.
"""
from litellm.proxy.batches_endpoints.endpoints import create_batch
model_name = "my-vllm-model"
raw_batch_id = "batch_abc123"
mock_response = _make_batch_response(batch_id=raw_batch_id)
mock_request = _make_mock_request(headers={"x-litellm-model": model_name})
mock_fastapi_response = MagicMock()
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.parent_otel_span = None
mock_user_api_key_dict.user_id = "test_user"
mock_credentials = {
"api_key": "sk-test",
"api_base": "http://vllm:8000",
"custom_llm_provider": "openai",
}
with (
patch(
"litellm.proxy.batches_endpoints.endpoints._read_request_body",
new=AsyncMock(
return_value={"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"}
),
),
patch(
"litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing"
) as mock_processor_cls,
patch(
"litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model",
return_value=mock_credentials,
),
patch(
"litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials",
),
patch(
"litellm.acreate_batch",
new=AsyncMock(return_value=mock_response),
),
patch(
"litellm.proxy.batches_endpoints.endpoints.is_known_model",
return_value=False,
),
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.llm_router", MagicMock()),
patch("litellm.proxy.proxy_server.proxy_config", MagicMock()),
patch("litellm.proxy.proxy_server.version", "1.0.0"),
patch(
"litellm.proxy.proxy_server.proxy_logging_obj",
MagicMock(
post_call_success_hook=AsyncMock(return_value=mock_response),
update_request_status=AsyncMock(),
),
),
):
# Setup the mock processor to return data and logging obj
mock_processor = MagicMock()
mock_processor.common_processing_pre_call_logic = AsyncMock(
return_value=(
{"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"},
MagicMock(),
)
)
mock_processor_cls.return_value = mock_processor
response = await create_batch(
request=mock_request,
fastapi_response=mock_fastapi_response,
provider=None,
user_api_key_dict=mock_user_api_key_dict,
)
# The batch_id should be encoded with model info
assert response.id != raw_batch_id, (
f"Expected batch_id to be encoded, but got raw ID: {response.id}"
)
assert response.id.startswith("batch_"), (
f"Encoded batch_id should keep batch_ prefix, got: {response.id}"
)
# Should be decodable back to the original
decoded_model = decode_model_from_file_id(response.id)
assert decoded_model == model_name, (
f"Expected model '{model_name}' from decoded batch_id, got: {decoded_model}"
)
original_id = get_original_file_id(response.id)
assert original_id == raw_batch_id, (
f"Expected original ID '{raw_batch_id}', got: {original_id}"
)
@pytest.mark.asyncio
async def test_create_batch_with_x_litellm_model_encodes_output_and_error_file_ids():
"""
When a completed batch is returned with output_file_id and error_file_id,
these should also be encoded with model info.
"""
from litellm.proxy.batches_endpoints.endpoints import create_batch
model_name = "my-vllm-model"
raw_output_file = "file-output789"
raw_error_file = "file-error012"
mock_response = _make_batch_response(
batch_id="batch_abc123",
output_file_id=raw_output_file,
error_file_id=raw_error_file,
status="completed",
)
mock_request = _make_mock_request(headers={"x-litellm-model": model_name})
mock_fastapi_response = MagicMock()
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.parent_otel_span = None
mock_user_api_key_dict.user_id = "test_user"
mock_credentials = {
"api_key": "sk-test",
"api_base": "http://vllm:8000",
"custom_llm_provider": "openai",
}
with (
patch(
"litellm.proxy.batches_endpoints.endpoints._read_request_body",
new=AsyncMock(
return_value={"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"}
),
),
patch(
"litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing"
) as mock_processor_cls,
patch(
"litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model",
return_value=mock_credentials,
),
patch(
"litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials",
),
patch(
"litellm.acreate_batch",
new=AsyncMock(return_value=mock_response),
),
patch(
"litellm.proxy.batches_endpoints.endpoints.is_known_model",
return_value=False,
),
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.llm_router", MagicMock()),
patch("litellm.proxy.proxy_server.proxy_config", MagicMock()),
patch("litellm.proxy.proxy_server.version", "1.0.0"),
patch(
"litellm.proxy.proxy_server.proxy_logging_obj",
MagicMock(
post_call_success_hook=AsyncMock(return_value=mock_response),
update_request_status=AsyncMock(),
),
),
):
mock_processor = MagicMock()
mock_processor.common_processing_pre_call_logic = AsyncMock(
return_value=(
{"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"},
MagicMock(),
)
)
mock_processor_cls.return_value = mock_processor
response = await create_batch(
request=mock_request,
fastapi_response=mock_fastapi_response,
provider=None,
user_api_key_dict=mock_user_api_key_dict,
)
# output_file_id should be encoded
assert decode_model_from_file_id(response.output_file_id) == model_name
assert get_original_file_id(response.output_file_id) == raw_output_file
# error_file_id should be encoded
assert decode_model_from_file_id(response.error_file_id) == model_name
assert get_original_file_id(response.error_file_id) == raw_error_file
@pytest.mark.asyncio
async def test_create_batch_without_x_litellm_model_returns_raw_ids():
"""
Without x-litellm-model header, create_batch should NOT encode batch IDs
(falls through to Scenario 3 / custom_llm_provider fallback).
"""
from litellm.proxy.batches_endpoints.endpoints import create_batch
raw_batch_id = "batch_abc123"
mock_response = _make_batch_response(batch_id=raw_batch_id)
mock_request = _make_mock_request(headers={})
mock_fastapi_response = MagicMock()
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.parent_otel_span = None
mock_user_api_key_dict.user_id = "test_user"
with (
patch(
"litellm.proxy.batches_endpoints.endpoints._read_request_body",
new=AsyncMock(
return_value={"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"}
),
),
patch(
"litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing"
) as mock_processor_cls,
patch(
"litellm.acreate_batch",
new=AsyncMock(return_value=mock_response),
),
patch(
"litellm.proxy.batches_endpoints.endpoints.is_known_model",
return_value=False,
),
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.proxy_config", MagicMock()),
patch("litellm.proxy.proxy_server.version", "1.0.0"),
patch(
"litellm.proxy.proxy_server.proxy_logging_obj",
MagicMock(
post_call_success_hook=AsyncMock(return_value=mock_response),
update_request_status=AsyncMock(),
),
),
):
mock_processor = MagicMock()
mock_processor.common_processing_pre_call_logic = AsyncMock(
return_value=(
{"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"},
MagicMock(),
)
)
mock_processor_cls.return_value = mock_processor
response = await create_batch(
request=mock_request,
fastapi_response=mock_fastapi_response,
provider=None,
user_api_key_dict=mock_user_api_key_dict,
)
# Without x-litellm-model, the batch_id should remain raw
assert response.id == raw_batch_id
assert decode_model_from_file_id(response.id) is None
class TestBatchIdRoundTripWithRetrieve:
"""
Tests that batch IDs encoded during create_batch can be decoded
correctly during retrieve_batch (Scenario 1: model_from_id).
"""
def test_encoded_batch_id_is_decoded_for_retrieve(self):
"""
Simulates the full round-trip: create encodes the ID,
retrieve decodes it to get the model and original batch_id.
"""
from litellm.proxy.openai_files_endpoints.common_utils import (
encode_file_id_with_model,
)
model_name = "my-vllm-model"
raw_batch_id = "batch_vllm_12345"
# What create_batch does:
encoded_id = encode_file_id_with_model(
file_id=raw_batch_id, model=model_name, id_type="batch"
)
# What retrieve_batch does:
decoded_model = decode_model_from_file_id(encoded_id)
original_id = get_original_file_id(encoded_id)
assert decoded_model == model_name
assert original_id == raw_batch_id
def test_vllm_style_batch_id_roundtrip(self):
"""
VLLM may return batch IDs in various formats.
Verify round-trip works for common patterns.
"""
from litellm.proxy.openai_files_endpoints.common_utils import (
encode_file_id_with_model,
)
test_cases = [
("batch_abc123", "vllm-llama3"),
("batch_67890", "openai/llama-3-8b"),
("batch_some-uuid-here", "my-custom-vllm"),
]
for raw_id, model in test_cases:
encoded = encode_file_id_with_model(
file_id=raw_id, model=model, id_type="batch"
)
assert encoded.startswith("batch_")
assert decode_model_from_file_id(encoded) == model
assert get_original_file_id(encoded) == raw_id

View file

@ -502,6 +502,53 @@ def test_router_get_deployment_credentials_with_provider():
assert credentials3 is None
def test_router_get_deployment_credentials_with_provider_wildcard():
"""
Test that get_deployment_credentials_with_provider handles wildcard patterns.
When a model like openai/gpt-4o is requested and the config has openai/*,
the method should resolve the wildcard pattern and return credentials.
"""
router = Router(
model_list=[
{
"model_name": "openai/*",
"litellm_params": {
"model": "openai/*",
"api_key": "sk-wildcard-123",
"api_base": "https://api.openai.com/v1",
},
"model_info": {"id": "openai-wildcard-deployment"},
},
{
"model_name": "anthropic/*",
"litellm_params": {
"model": "anthropic/*",
"api_key": "sk-ant-wildcard-456",
},
"model_info": {"id": "anthropic-wildcard-deployment"},
},
]
)
# Test wildcard pattern matching for OpenAI
credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-4o")
assert credentials is not None
assert credentials["api_key"] == "sk-wildcard-123"
assert credentials["custom_llm_provider"] == "openai"
assert credentials["api_base"] == "https://api.openai.com/v1"
# Test wildcard pattern matching for Anthropic
credentials2 = router.get_deployment_credentials_with_provider(model_id="anthropic/claude-3-opus")
assert credentials2 is not None
assert credentials2["api_key"] == "sk-ant-wildcard-456"
assert credentials2["custom_llm_provider"] == "anthropic"
# Test with non-matching model
credentials3 = router.get_deployment_credentials_with_provider(model_id="vertex_ai/gemini-pro")
assert credentials3 is None
def test_router_get_deployment_model_info():
router = Router(
model_list=[

View file

@ -10,7 +10,7 @@ sys.path.insert(0, os.path.abspath("../../.."))
# Import required modules
import litellm
from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler
from litellm.types.llms.openai import ResponsesAPIResponse, OpenAIMcpServerTool, ToolParam
from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStreamingResponse, OpenAIMcpServerTool, ToolParam
class MockUserAPIKeyAuth:
@ -542,193 +542,211 @@ async def test_mcp_allowed_tools_filtering():
async def test_streaming_mcp_events_validation():
"""
Test that MCP streaming events are properly emitted when using streaming with MCP tools.
This test validates:
1. MCP discovery events are emitted first
2. Regular streaming response events follow
3. Tool execution events are emitted when tools are auto-executed
"""
from unittest.mock import AsyncMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.types.llms.openai import ResponsesAPIStreamEvents
print("🧪 Testing MCP streaming events...")
# Mock MCP tools that would be returned from the manager
mock_mcp_tools = [
type('MCPTool', (), {
'name': 'search_repo',
'description': 'Search BerriAI/litellm repository for information',
'inputSchema': {
"type": "object",
"properties": {
"query": {"type": "string", "description": "Search query"}
type(
"MCPTool",
(),
{
"name": "search_repo",
"description": "Search BerriAI/litellm repository for information",
"inputSchema": {
"type": "object",
"properties": {
"query": {"type": "string", "description": "Search query"}
},
"required": ["query"],
},
"required": ["query"]
}
})(),
type('MCPTool', (), {
'name': 'get_repo_info',
'description': 'Get repository information',
'inputSchema': {
"type": "object",
"properties": {
"repo_name": {"type": "string", "description": "Repository name"}
},
)(),
type(
"MCPTool",
(),
{
"name": "get_repo_info",
"description": "Get repository information",
"inputSchema": {
"type": "object",
"properties": {
"repo_name": {
"type": "string",
"description": "Repository name",
}
},
"required": ["repo_name"],
},
"required": ["repo_name"]
}
})()
},
)(),
]
# Mock the MCP operations
with patch.object(LiteLLM_Proxy_MCP_Handler, '_get_mcp_tools_from_manager', new_callable=AsyncMock) as mock_get_tools, \
patch.object(LiteLLM_Proxy_MCP_Handler, '_execute_tool_calls', new_callable=AsyncMock) as mock_execute_tools:
# Build fake streaming chunks that the inner aresponses() call would yield
fake_response_obj = MagicMock(spec=ResponsesAPIResponse)
fake_response_obj.id = "resp_fake_123"
fake_response_obj.output = []
fake_created_chunk = MagicMock(spec=ResponsesAPIStreamingResponse)
fake_created_chunk.type = ResponsesAPIStreamEvents.RESPONSE_CREATED
fake_created_chunk.response = fake_response_obj
fake_in_progress_chunk = MagicMock(spec=ResponsesAPIStreamingResponse)
fake_in_progress_chunk.type = ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS
fake_in_progress_chunk.response = fake_response_obj
fake_output_item_added_chunk = MagicMock(spec=ResponsesAPIStreamingResponse)
fake_output_item_added_chunk.type = ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED
fake_output_item_added_chunk.response = fake_response_obj
fake_completed_chunk = MagicMock(spec=ResponsesAPIStreamingResponse)
fake_completed_chunk.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED
fake_completed_chunk.response = fake_response_obj
# Create a fake async iterator for the inner LLM streaming call
class FakeAsyncIterator:
def __init__(self, chunks):
self._chunks = list(chunks)
self._index = 0
def __aiter__(self):
return self
async def __anext__(self):
if self._index >= len(self._chunks):
raise StopAsyncIteration
chunk = self._chunks[self._index]
self._index += 1
return chunk
fake_stream = FakeAsyncIterator(
[
fake_created_chunk,
fake_in_progress_chunk,
fake_output_item_added_chunk,
fake_completed_chunk,
]
)
# Mock the MCP operations and the inner aresponses call
with patch.object(
LiteLLM_Proxy_MCP_Handler,
"_get_mcp_tools_from_manager",
new_callable=AsyncMock,
) as mock_get_tools, patch.object(
LiteLLM_Proxy_MCP_Handler,
"_execute_tool_calls",
new_callable=AsyncMock,
) as mock_execute_tools, patch(
"litellm.responses.main.aresponses",
new_callable=AsyncMock,
return_value=fake_stream,
):
# Setup MCP mocks
mock_get_tools.return_value = (mock_mcp_tools, ["test_server"])
def mock_execute_tool_calls_side_effect(tool_calls, user_api_key_auth):
async def mock_execute_tool_calls_side_effect(
tool_server_map, tool_calls, user_api_key_auth, **kwargs
):
"""Mock tool execution with realistic results"""
results = []
for tool_call in tool_calls:
call_id = None
if isinstance(tool_call, dict):
call_id = tool_call.get("call_id") or tool_call.get("id")
elif hasattr(tool_call, 'call_id'):
elif hasattr(tool_call, "call_id"):
call_id = tool_call.call_id
elif hasattr(tool_call, 'id'):
elif hasattr(tool_call, "id"):
call_id = tool_call.id
if call_id:
results.append({
"tool_call_id": call_id,
"result": "LiteLLM is a unified interface for 100+ LLMs that provides consistent OpenAI-format output and includes proxy server capabilities."
})
results.append(
{
"tool_call_id": call_id,
"result": "LiteLLM is a unified interface for 100+ LLMs.",
}
)
return results
mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect
# Configure MCP tool with streaming and auto-execution
mcp_tool_config = {
"type": "mcp",
"server_url": "litellm_proxy/mcp/test_server",
"require_approval": "never" # This enables auto-execution
"server_url": "litellm_proxy/mcp/test_server",
"require_approval": "never", # This enables auto-execution
}
print("📞 Making streaming request with MCP tools...")
# Make streaming request with MCP tools
response = await litellm.aresponses(
model="gpt-4o-mini", # Use cheaper model for testing
model="gpt-4o-mini",
tools=[mcp_tool_config],
tool_choice="required",
input=[{
"role": "user",
"type": "message",
"content": "What is LiteLLM? Give me a brief overview."
}],
stream=True
input=[
{
"role": "user",
"type": "message",
"content": "What is LiteLLM? Give me a brief overview.",
}
],
stream=True,
)
print(f"📋 Response type: {type(response)}")
assert hasattr(response, '__aiter__'), "Response should be async iterable for streaming"
assert hasattr(
response, "__aiter__"
), "Response should be async iterable for streaming"
# Collect all streaming events
events = []
event_types = []
mcp_discovery_events = []
mcp_execution_events = []
regular_events = []
print("🔄 Collecting streaming events...")
try:
async for chunk in response:
events.append(chunk)
event_type = getattr(chunk, 'type', 'unknown')
event_types.append(event_type)
# Categorize events
if event_type in [
ResponsesAPIStreamEvents.MCP_TOOLS_DISCOVERY_STARTED,
ResponsesAPIStreamEvents.MCP_TOOLS_DISCOVERY_COMPLETED
]:
mcp_discovery_events.append(chunk)
elif event_type in [
ResponsesAPIStreamEvents.MCP_TOOL_EXECUTION_STARTED,
ResponsesAPIStreamEvents.MCP_TOOL_EXECUTION_COMPLETED
]:
mcp_execution_events.append(chunk)
else:
regular_events.append(chunk)
print(f"📦 Event: {event_type}")
# Print MCP-specific event details
if hasattr(chunk, 'mcp_servers'):
print(f" 🔧 MCP Servers: {chunk.mcp_servers}")
elif hasattr(chunk, 'mcp_tools'):
print(f" 🛠️ MCP Tools: {len(chunk.mcp_tools)} tools discovered")
elif hasattr(chunk, 'tool_name'):
print(f" ⚙️ Tool: {chunk.tool_name}")
if hasattr(chunk, 'result'):
print(f" ✅ Result: {chunk.result[:100]}...")
except Exception as e:
print(f"❌ Error during streaming: {e}")
# Continue with validation of events collected so far
print(f"\n📊 Event Summary:")
print(f" Total events: {len(events)}")
print(f" MCP discovery events: {len(mcp_discovery_events)}")
print(f" MCP execution events: {len(mcp_execution_events)}")
print(f" Regular streaming events: {len(regular_events)}")
print(f" Event types: {set(event_types)}")
# Validate MCP discovery events
if mcp_discovery_events:
print("✅ MCP discovery events found!")
# Check for discovery started event
started_events = [e for e in mcp_discovery_events if e.type == ResponsesAPIStreamEvents.MCP_TOOLS_DISCOVERY_STARTED]
if started_events:
print(f" 🚀 Discovery started events: {len(started_events)}")
started_event = started_events[0]
if hasattr(started_event, 'mcp_servers'):
print(f" 📡 MCP servers: {started_event.mcp_servers}")
# Check for discovery completed event
completed_events = [e for e in mcp_discovery_events if e.type == ResponsesAPIStreamEvents.MCP_TOOLS_DISCOVERY_COMPLETED]
if completed_events:
print(f" 🏁 Discovery completed events: {len(completed_events)}")
completed_event = completed_events[0]
if hasattr(completed_event, 'mcp_tools'):
print(f" 🔧 Tools discovered: {len(completed_event.mcp_tools)}")
else:
print("⚠️ No MCP discovery events found")
# Validate MCP execution events (if auto-execution occurred)
if mcp_execution_events:
print("✅ MCP tool execution events found!")
execution_started = [e for e in mcp_execution_events if e.type == ResponsesAPIStreamEvents.MCP_TOOL_EXECUTION_STARTED]
execution_completed = [e for e in mcp_execution_events if e.type == ResponsesAPIStreamEvents.MCP_TOOL_EXECUTION_COMPLETED]
print(f" 🚀 Execution started events: {len(execution_started)}")
print(f" 🏁 Execution completed events: {len(execution_completed)}")
# Validate that we got some form of streaming response
async for chunk in response:
events.append(chunk)
event_type = getattr(chunk, "type", "unknown")
event_types.append(event_type)
# Categorize events
if event_type in [
ResponsesAPIStreamEvents.MCP_LIST_TOOLS_IN_PROGRESS,
ResponsesAPIStreamEvents.MCP_LIST_TOOLS_COMPLETED,
]:
mcp_discovery_events.append(chunk)
else:
regular_events.append(chunk)
# Validate that we got streaming events
assert len(events) > 0, "Should have received at least some streaming events"
# Validate MCP discovery events were emitted
assert (
len(mcp_discovery_events) > 0
), "Should have received MCP discovery events"
# Check that discovery events come before regular content events
first_discovery_idx = next(
i
for i, e in enumerate(events)
if getattr(e, "type", None)
in [
ResponsesAPIStreamEvents.MCP_LIST_TOOLS_IN_PROGRESS,
ResponsesAPIStreamEvents.MCP_LIST_TOOLS_COMPLETED,
]
)
# The output_item.added event triggers the transition to MCP discovery,
# so discovery events should appear after it in the stream
assert first_discovery_idx > 0, "MCP discovery events should follow the initial output_item.added event"
# Verify MCP mocks were called
assert mock_get_tools.called, "MCP tools should have been fetched"
print("✅ MCP tool fetching was called")
print("🎉 MCP streaming events validation completed!")
return {
'total_events': len(events),
'mcp_discovery_events': len(mcp_discovery_events),
'mcp_execution_events': len(mcp_execution_events),
'regular_events': len(regular_events),
'event_types': list(set(event_types))
}
@pytest.mark.asyncio
@ -1250,4 +1268,147 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e():
}
@pytest.mark.asyncio
@pytest.mark.parametrize("model", ["gpt-4o-mini"])
async def test_streaming_mcp_event_order_and_response_id_consistency(
model: str, caplog: pytest.LogCaptureFixture
):
"""
Test that:
1. Streaming events are emitted in correct order (response.created, response.in_progress, response.output_item.added before MCP events)
2. All response lifecycle events share the same response ID within a cycle
"""
if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv("OPENAI_API_KEY"):
pytest.skip("OPENAI_API_KEY not set, skipping openai model test")
from unittest.mock import AsyncMock, patch
mock_mcp_tools = [
type('MCPTool', (), {
'name': 'get_weather',
'description': 'Get weather for a city',
'inputSchema': {
"type": "object",
"properties": {
"city": {"type": "string", "description": "City name"}
},
"required": ["city"]
}
})()
]
with caplog.at_level(logging.ERROR):
with patch.object(
LiteLLM_Proxy_MCP_Handler,
'_get_mcp_tools_from_manager',
new_callable=AsyncMock,
) as mock_get_tools, patch.object(
LiteLLM_Proxy_MCP_Handler,
'_execute_tool_calls',
new_callable=AsyncMock,
) as mock_execute_tools:
mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"])
def mock_execute_side_effect(tool_calls, user_api_key_auth, **kwargs):
results = []
for tool_call in tool_calls:
call_id = None
if isinstance(tool_call, dict):
call_id = tool_call.get("call_id") or tool_call.get("id")
elif hasattr(tool_call, 'call_id'):
call_id = tool_call.call_id
elif hasattr(tool_call, 'id'):
call_id = tool_call.id
if call_id:
results.append({
"tool_call_id": call_id,
"result": "Sunny, 72°F",
})
return results
mock_execute_tools.side_effect = mock_execute_side_effect
mcp_tool_config = cast(Any, {
"type": "mcp",
"server_url": "litellm_proxy",
"require_approval": "never",
})
response = await litellm.aresponses(
model=model,
tools=[mcp_tool_config],
input=[{
"role": "user",
"type": "message",
"content": "What's the weather in San Francisco?"
}],
stream=True,
)
events = []
async for chunk in response:
events.append(chunk)
assert len(events) > 0, "Should receive streaming events"
created_idx = next((i for i, e in enumerate(events) if getattr(e, 'type', None) == 'response.created'), None)
in_progress_idx = next((i for i, e in enumerate(events) if getattr(e, 'type', None) == 'response.in_progress'), None)
output_item_added_idx = next((i for i, e in enumerate(events) if getattr(e, 'type', None) == 'response.output_item.added'), None)
mcp_in_progress_idx = next((i for i, e in enumerate(events) if 'mcp_list_tools.in_progress' in str(getattr(e, 'type', ''))), None)
completed_idx = next((i for i, e in enumerate(events) if getattr(e, 'type', None) == 'response.completed'), None)
assert created_idx is not None, "response.created event should be present"
assert in_progress_idx is not None, "response.in_progress event should be present"
assert output_item_added_idx is not None, "response.output_item.added event should be present"
assert created_idx < in_progress_idx, "response.created should come before response.in_progress"
assert in_progress_idx < output_item_added_idx, "response.in_progress should come before response.output_item.added"
if mcp_in_progress_idx is not None:
assert output_item_added_idx < mcp_in_progress_idx, "response.output_item.added should come before response.mcp_list_tools.in_progress"
response_ids = []
for i, event in enumerate(events):
event_type = getattr(event, 'type', None)
if hasattr(event, 'response'):
response_obj = getattr(event, 'response', None)
if response_obj and hasattr(response_obj, 'id'):
event_type_value = event_type.value if hasattr(event_type, 'value') else str(event_type)
if any(x in event_type_value for x in ['response.created', 'response.in_progress', 'response.completed']):
response_ids.append((i, event_type_value, response_obj.id))
assert len(response_ids) >= 2, f"Should have at least 2 response lifecycle events. Found {len(response_ids)}"
cycles = []
current_cycle = []
current_id = None
for idx, event_type, resp_id in response_ids:
if current_id is None or resp_id == current_id:
current_cycle.append((idx, event_type, resp_id))
current_id = resp_id
else:
if current_cycle:
cycles.append(current_cycle)
current_cycle = [(idx, event_type, resp_id)]
current_id = resp_id
if current_cycle:
cycles.append(current_cycle)
for cycle_num, cycle in enumerate(cycles):
cycle_ids = set(resp_id for _, _, resp_id in cycle)
assert len(cycle_ids) == 1, f"Cycle {cycle_num + 1} should have consistent response ID. Found {len(cycle_ids)} unique IDs"
assert completed_idx is not None, "response.completed event should be present"
lite_errors = [
record for record in caplog.records
if record.levelno >= logging.ERROR
and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage())
]
assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join(
record.getMessage() for record in lite_errors
)

View file

@ -0,0 +1,239 @@
"""
E2E tests for OpenAI Responses API WebSocket mode through the LiteLLM proxy.
Connects to ws://0.0.0.0:4000/v1/responses, sends response.create events,
and validates the streamed response events.
Requires:
- Proxy running: python -m litellm.proxy.proxy_cli --config <config> --port 4000
- Model configured in proxy (e.g. gpt-4o-mini)
See: https://developers.openai.com/api/docs/guides/websocket-mode/
"""
import asyncio
import json
import os
import httpx
import pytest
# ── Configuration ─────────────────────────────────────────────────────────────
PROXY_BASE_URL = os.environ.get("LITELLM_PROXY_BASE_URL", "ws://0.0.0.0:4000")
PROXY_MASTER_KEY = os.environ.get("LITELLM_PROXY_KEY", "sk-1234")
PROXY_MODEL = os.environ.get("LITELLM_PROXY_RESPONSES_MODEL", "gpt-4o-mini")
# ──────────────────────────────────────────────────────────────────────────────
def _generate_key() -> str:
"""Generate a key for testing via proxy key/generate endpoint."""
url = "http://0.0.0.0:4000/key/generate"
headers = {
"Authorization": f"Bearer {PROXY_MASTER_KEY}",
"Content-Type": "application/json",
}
response = httpx.post(url, headers=headers, json={}, timeout=10)
if response.status_code != 200:
raise Exception(
f"Key generation failed with status: {response.status_code}. "
"Is the proxy running?"
)
return response.json()["key"]
def _assert_basic_response(events: list[dict], label: str = "") -> None:
"""Assert that events contain response.created, response.completed, and usage."""
prefix = f"[{label}] " if label else ""
types = [e.get("type") for e in events]
assert len(events) > 0, f"{prefix}no events received"
assert "response.created" in types, f"{prefix}missing response.created, got: {types}"
assert "response.completed" in types, (
f"{prefix}missing response.completed, got: {types}"
)
completed = next(e for e in events if e.get("type") == "response.completed")
resp = completed.get("response", {})
assert resp.get("status") == "completed", (
f"{prefix}status != completed: {resp.get('status')}"
)
usage = resp.get("usage", {})
assert usage.get("input_tokens", 0) > 0, f"{prefix}input_tokens=0"
assert usage.get("output_tokens", 0) > 0, f"{prefix}output_tokens=0"
streaming_types = {
"response.output_item.added",
"response.content_part.added",
"response.output_text.delta",
"response.output_item.done",
}
found = streaming_types & set(types)
assert found, f"{prefix}no streaming delta events found, got: {types}"
@pytest.mark.asyncio
async def test_responses_websocket_proxy_basic():
"""
Sends a simple response.create event to the proxy WebSocket endpoint
and validates response.created, response.completed, and streaming events.
"""
try:
import websockets
except ImportError:
pytest.skip("websockets not installed")
try:
key = _generate_key()
except Exception as e:
pytest.skip(
f"Proxy not available or key generation failed: {e}. "
"Start proxy: python -m litellm.proxy.proxy_cli --config <config> --port 4000"
)
url = f"{PROXY_BASE_URL}/v1/responses?model={PROXY_MODEL}"
headers = {"Authorization": f"Bearer {key}"}
events: list[dict] = []
try:
async with websockets.connect(
url, additional_headers=headers, open_timeout=5
) as ws:
payload = {
"type": "response.create",
"model": PROXY_MODEL,
"store": False,
"input": [
{
"type": "message",
"role": "user",
"content": [
{"type": "input_text", "text": "Say hello in one word."}
],
}
],
"tools": [],
}
await ws.send(json.dumps(payload))
for _ in range(50):
msg = await asyncio.wait_for(ws.recv(), timeout=15)
event = json.loads(msg)
events.append(event)
if event.get("type") in (
"response.completed",
"response.failed",
"error",
):
break
except Exception as e:
pytest.fail(
f"WebSocket connection failed: {e}. "
"Ensure proxy is running and model is configured."
)
_assert_basic_response(events, "proxy-basic")
@pytest.mark.asyncio
async def test_responses_websocket_proxy_multi_turn():
"""
Sends two sequential response.create events with previous_response_id
to validate multi-turn conversation over a single WebSocket.
"""
try:
import websockets
except ImportError:
pytest.skip("websockets not installed")
try:
key = _generate_key()
except Exception as e:
pytest.skip(
f"Proxy not available or key generation failed: {e}. "
"Start proxy: python -m litellm.proxy.proxy_cli --config <config> --port 4000"
)
url = f"{PROXY_BASE_URL}/v1/responses?model={PROXY_MODEL}"
headers = {"Authorization": f"Bearer {key}"}
all_events: list[dict] = []
completed: list[dict] = []
first_id = None
try:
async with websockets.connect(
url, additional_headers=headers, open_timeout=5
) as ws:
# Turn 1
await ws.send(
json.dumps(
{
"type": "response.create",
"model": PROXY_MODEL,
"store": True,
"input": [
{
"type": "message",
"role": "user",
"content": [
{
"type": "input_text",
"text": "Remember the number 7. Just say OK.",
}
],
}
],
}
)
)
for _ in range(50):
msg = await asyncio.wait_for(ws.recv(), timeout=15)
event = json.loads(msg)
all_events.append(event)
if event.get("type") == "response.completed":
completed.append(event)
first_id = event.get("response", {}).get("id")
break
if event.get("type") in ("response.failed", "error"):
break
assert first_id, "Turn 1 never completed"
# Turn 2
await ws.send(
json.dumps(
{
"type": "response.create",
"model": PROXY_MODEL,
"store": True,
"previous_response_id": first_id,
"input": [
{
"type": "message",
"role": "user",
"content": [
{
"type": "input_text",
"text": "What number did I tell you to remember?",
}
],
}
],
}
)
)
for _ in range(50):
msg = await asyncio.wait_for(ws.recv(), timeout=15)
event = json.loads(msg)
all_events.append(event)
if event.get("type") == "response.completed":
completed.append(event)
break
if event.get("type") in ("response.failed", "error"):
break
except Exception as e:
pytest.fail(
f"WebSocket multi-turn failed: {e}. "
"Ensure proxy is running and model is configured."
)
assert len(completed) >= 2, (
f"Expected 2 response.completed events, got {len(completed)}"
)
assert completed[1].get("response", {}).get("status") == "completed"

View file

@ -0,0 +1,494 @@
"""Base class for LiteLLM integration tests.
Supports both local (mock) and remote testing modes via environment variables:
- USE_LOCAL_LITELLM: When "true", uses local LiteLLM at localhost:4000 (default: false)
- USE_MOCK_MODELS: When "true", uses mock model names (default: false)
- LITELLM_API_KEY: API key for remote LiteLLM (required when USE_LOCAL_LITELLM=false)
- LITELLM_BASE_URL: Base URL for remote LiteLLM (required when USE_LOCAL_LITELLM=false)
"""
import enum
import os
import time
import uuid
from abc import ABC
from collections import defaultdict
from typing import Any, Callable, Dict, List, Tuple, Union
import httpx
import openai
import pytest
import requests
from urllib3.exceptions import InsecureRequestWarning
requests.packages.urllib3.disable_warnings(category=InsecureRequestWarning)
LOCAL_LITELLM_BASE_URL = "http://localhost:4000"
LOCAL_MOCK_SERVER_URL = "http://localhost:8090"
if "USE_LOCAL_LITELLM" not in os.environ:
os.environ["USE_LOCAL_LITELLM"] = "true"
if "USE_MOCK_MODELS" not in os.environ:
os.environ["USE_MOCK_MODELS"] = "true"
if "USE_STATE_TRACKER" not in os.environ:
os.environ["USE_STATE_TRACKER"] = "true"
if "DATABASE_URL" not in os.environ:
os.environ["DATABASE_URL"] = "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm"
def use_local_litellm() -> bool:
return os.environ.get("USE_LOCAL_LITELLM", "false").lower() == "true"
def use_remote_litellm() -> bool:
return not use_local_litellm()
def use_mock_models() -> bool:
return os.environ.get("USE_MOCK_MODELS", "false").lower() == "true"
def get_local_litellm_base_url() -> str:
return LOCAL_LITELLM_BASE_URL
def get_remote_litellm_base_url() -> str:
return os.environ.get("LITELLM_BASE_URL", "").rstrip("/")
def get_litellm_base_url() -> str:
if use_local_litellm():
return get_local_litellm_base_url()
return get_remote_litellm_base_url()
def get_litellm_api_key() -> str:
if use_local_litellm():
return "sk-1234"
return os.environ.get("LITELLM_API_KEY", "")
def get_mock_server_base_url() -> str:
return LOCAL_MOCK_SERVER_URL
def get_responses_model_name() -> str:
if use_mock_models():
return "openai-fake-gpt-4o"
return "gpt-4o-mini-2024-07-18"
def model_id(param) -> str:
"""Generate a test ID from a model name or tuple containing model name.
Handles both:
- String: "gpt-4o-mini" -> "gpt_4o_mini"
- Tuple: ("gpt-4o", "openai/gpt-4o") -> "gpt_4o"
"""
if isinstance(param, tuple):
name = param[0]
else:
name = param
return name.replace("-", "_").replace(".", "_")
def generate_test_id(
params: Tuple[str, ...],
test_name: str = "test",
) -> str:
"""Generate test ID from model parameters tuple.
Handles two tuple formats:
- 6 elements: (provider, deployment, model_name, api_version, action, reason)
- 7 elements: (provider, deployment, model_name, api_version, model_id, action, reason)
Uses model_id (position 4) if 7 elements, otherwise model_name (position 2).
"""
provider = params[0]
deployment = params[1]
api_version = params[3]
if len(params) == 7:
identifier = params[4] # model_id
else:
identifier = params[2] # model_name
test_id = "/".join([provider, deployment, api_version, identifier, test_name])
return test_id.replace("-", "_").replace(".", "_")
class ModelTestAction(enum.Enum):
NOT_APPLICABLE = 1
SKIP = 2
RUN = 3
WARN_ON_FAIL = 4
def applicable(self) -> bool:
return self.value != ModelTestAction.NOT_APPLICABLE.value
class BaseLiteLLMIntegrationTest(ABC):
"""Base class for all LiteLLM integration tests.
Supports both local/mock and remote testing based on environment variables.
"""
@staticmethod
def get_api_key() -> str:
return get_litellm_api_key()
@staticmethod
def get_base_url() -> str:
return get_litellm_base_url()
@staticmethod
def get_ca_bundle_path() -> str:
current_dir = os.path.dirname(os.path.abspath(__file__))
# change if needed
@classmethod
def _get_ssl_verify_setting(cls) -> Union[bool, str]:
"""Get the appropriate SSL verification setting based on mode.
Returns path string (not SSLContext) for compatibility with both
requests and httpx libraries.
"""
if use_local_litellm():
return False
ca_bundle_path = cls.get_ca_bundle_path()
if os.path.exists(ca_bundle_path):
return ca_bundle_path
return True
@classmethod
def setup_class(cls):
cls.api_key = cls.get_api_key()
cls.base_url = cls.get_base_url()
if not cls.api_key:
pytest.fail(
"API key is not available. Set LITELLM_API_KEY or USE_LOCAL_LITELLM=true",
)
if not cls.base_url:
pytest.fail(
"Base URL is not available. Set LITELLM_BASE_URL or USE_LOCAL_LITELLM=true",
)
verify_setting = cls._get_ssl_verify_setting()
if use_remote_litellm() and isinstance(verify_setting, str):
os.environ["REQUESTS_CA_BUNDLE"] = verify_setting
os.environ["CURL_CA_BUNDLE"] = verify_setting
print(f"Using CA bundle: {verify_setting}")
cls.openai_client = openai.OpenAI(
base_url=cls.base_url,
api_key=cls.api_key,
http_client=httpx.Client(verify=verify_setting),
)
@classmethod
def make_request(
cls,
method: str,
endpoint: str,
timeout_secs: int,
**kwargs,
) -> requests.Response:
headers = kwargs.get("headers", {})
headers["Authorization"] = f"Bearer {cls.api_key}"
kwargs["headers"] = headers
kwargs.setdefault("timeout", timeout_secs)
kwargs.setdefault("verify", cls._get_ssl_verify_setting())
url = f"{cls.base_url}{endpoint}"
return requests.request(method, url, **kwargs)
@staticmethod
def generate_request_id() -> str:
return f"req-{uuid.uuid4().hex[:8]}"
@staticmethod
def get_timeout_secs(model_name: str) -> int:
model_lower = model_name.lower()
slow_models = ["gpt-5", "gpt_5", "o1", "claude-opus", "claude_opus", "o3", "o4"]
if any(slow_model in model_lower for slow_model in slow_models):
return 300
return 60
@staticmethod
def generate_unique_filename(extension: str = "txt") -> str:
return f"test_{time.time()}.{extension}"
@staticmethod
def extract_model_params(model_data: Dict[str, Any]) -> Tuple[str, str, str, str]:
"""Extract standardized parameters from model data."""
model_name = model_data.get("model_name", "")
model_info = model_data.get("model_info", {})
provider = model_info.get("litellm_provider", "unknown")
litellm_params = model_data.get("litellm_params", {})
if provider == "azure":
api_base = litellm_params.get("api_base", "unknown")
if api_base != "unknown" and "//" in api_base:
domain_name = api_base.split("//")[1]
deployment = domain_name.split(".")[0]
else:
deployment = "unknown"
api_version = litellm_params.get("api_version", "unknown")
elif provider in ["bedrock", "bedrock_converse"]:
deployment = litellm_params.get("aws_region_name", "unknown")
api_version = "unknown"
else:
deployment = "unknown"
api_version = "unknown"
return provider, deployment, model_name, api_version
@classmethod
def _fetch_all_models_from_litellm(cls) -> List[Dict[str, Any]]:
base_url = cls.get_base_url()
api_key = cls.get_api_key()
if not api_key or not base_url:
return []
verify_setting = cls._get_ssl_verify_setting()
response = requests.get(
f"{base_url}/model/info",
headers={"Authorization": f"Bearer {api_key}"},
verify=verify_setting,
timeout=30,
)
if response.status_code != 200:
raise RuntimeError(
f"Failed to fetch all models from {base_url}. Response code: {response.status_code}",
)
data = response.json()
return data.get("data", [])
@classmethod
def _fetch_all_approved_models(cls) -> List[Dict[str, Any]]:
return cls._fetch_all_models_from_litellm()
@classmethod
def build_model_test_params(
cls,
should_skip_model: Callable[
[str, str, str, str, Dict[str, Any]],
Tuple["ModelTestAction", str],
],
include_model_id: bool = False,
include_load_balanced: bool = False,
) -> List[Tuple[str, ...]]:
"""Build test parameters from all approved models.
Args:
should_skip_model: Callback that determines if a model should be skipped.
Signature: (provider, deployment, model_name, api_version, model_info) -> (action, reason)
include_model_id: If True, includes model_id in tuple (7 elements), else 6 elements.
include_load_balanced: If True, adds extra tests for load-balanced model groups.
Returns:
List of tuples with model test parameters.
- 6-element: (provider, deployment, model_name, api_version, action, reason)
- 7-element: (provider, deployment, model_name, api_version, model_id, action, reason)
"""
models = cls._fetch_all_approved_models()
test_params: List[Tuple[str, ...]] = []
models_by_model_name: Dict[str, List[Tuple[str, ...]]] = defaultdict(list)
for model_data in models:
model_info = model_data.get("model_info", {}) or {}
provider, deployment, model_name, api_version = cls.extract_model_params(
model_data,
)
model_test_action, model_test_action_reason = should_skip_model(
provider,
deployment,
model_name,
api_version,
model_info,
)
if model_test_action.applicable():
if include_model_id:
model_id = str(model_info.get("id"))
params_tuple: Tuple[str, ...] = (
provider,
deployment,
model_name,
api_version,
model_id,
model_test_action,
model_test_action_reason,
)
else:
params_tuple = (
provider,
deployment,
model_name,
api_version,
model_test_action,
model_test_action_reason,
)
test_params.append(params_tuple)
if include_load_balanced:
models_by_model_name[model_name].append(params_tuple)
if include_load_balanced and include_model_id:
for load_balanced_model_name, deployments in models_by_model_name.items():
if len(deployments) <= 1:
continue
first_deployment = deployments[0]
test_params.append(
(
first_deployment[0], # provider
"load_balanced",
load_balanced_model_name,
"load_balanced",
load_balanced_model_name, # model_id = model_name for LB
first_deployment[5], # model_test_action
first_deployment[6], # model_test_action_reason
),
)
return test_params
class UserKeyTestMixin:
"""Mixin for tests that need to create users and API keys."""
allowed_routes: list[str] = []
_base_url: str = None
_master_api_key: str = None
admin_client: httpx.Client = None
@classmethod
def setup_admin_client(cls):
cls._base_url = get_litellm_base_url()
cls._master_api_key = get_litellm_api_key()
verify_setting = (
False
if use_local_litellm()
else BaseLiteLLMIntegrationTest._get_ssl_verify_setting()
)
cls.admin_client = httpx.Client(base_url=cls._base_url, verify=verify_setting)
@classmethod
def teardown_admin_client(cls):
if cls.admin_client:
cls.admin_client.close()
@staticmethod
def unique_suffix() -> str:
return f"{time.strftime('%Y%m%d%H%M%S')}{int(time.time() * 1000) % 1000:03d}"
@classmethod
def create_user_and_key(cls, user_suffix: str) -> tuple[str, str, str]:
user_email = f"test-user-{user_suffix}-{cls.unique_suffix()}@test.com"
user_response = cls.admin_client.post(
"/user/new",
json={
"user_email": user_email,
"user_alias": user_email,
"user_role": "internal_user",
"auto_create_key": "false",
},
headers={
"Authorization": f"Bearer {cls._master_api_key}",
"Content-Type": "application/json",
},
timeout=30,
)
assert user_response.status_code == 200, (
f"Failed to create user: {user_response.status_code} - {user_response.text}"
)
user_id = user_response.json().get("user_id")
key_alias = user_email.replace("@", "-at-").replace(".", "-")
key_response = cls.admin_client.post(
"/key/generate",
json={
"user_id": user_id,
"key_alias": key_alias,
"allowed_routes": cls.allowed_routes,
},
headers={
"Authorization": f"Bearer {cls._master_api_key}",
"Content-Type": "application/json",
},
timeout=30,
)
assert key_response.status_code == 200, (
f"Failed to create key: {key_response.status_code} - {key_response.text}"
)
api_key = key_response.json().get("key")
print(f"Created user {user_email}")
return user_id, api_key, user_email
@classmethod
def create_user_key_and_client(
cls,
user_suffix: str,
) -> tuple[str, str, str, openai.OpenAI]:
user_id, api_key, user_email = cls.create_user_and_key(user_suffix)
verify_setting = (
False
if use_local_litellm()
else BaseLiteLLMIntegrationTest._get_ssl_verify_setting()
)
client = openai.OpenAI(
base_url=cls._base_url,
api_key=api_key,
http_client=httpx.Client(verify=verify_setting),
)
return user_id, api_key, user_email, client
@classmethod
def create_key_and_client(
cls,
user_id: str,
key_suffix: str,
) -> tuple[str, openai.OpenAI]:
key_alias = f"additional-key-{key_suffix}-{cls.unique_suffix()}"
key_response = cls.admin_client.post(
"/key/generate",
json={
"user_id": user_id,
"key_alias": key_alias,
"allowed_routes": cls.allowed_routes,
},
headers={
"Authorization": f"Bearer {cls._master_api_key}",
"Content-Type": "application/json",
},
timeout=30,
)
assert key_response.status_code == 200, (
f"Failed to create additional key: {key_response.status_code} - {key_response.text}"
)
api_key = key_response.json().get("key")
verify_setting = (
False
if use_local_litellm()
else BaseLiteLLMIntegrationTest._get_ssl_verify_setting()
)
client = openai.OpenAI(
base_url=cls._base_url,
api_key=api_key,
http_client=httpx.Client(verify=verify_setting),
)
print(f"Created additional key for user {user_id}")
return api_key, client

View file

@ -0,0 +1,311 @@
"""
Pytest configuration for Azure Batch E2E Tests.
This conftest manages:
1. Mock Azure Batch server (FastAPI on port 8090)
2. LiteLLM proxy server (port 4000)
3. PostgreSQL database setup
"""
import asyncio
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Generator
import httpx
import pytest
_test_dir = Path(__file__).parent
sys.path.insert(0, str(_test_dir.parent.parent)) # litellm root
sys.path.insert(0, str(_test_dir)) # test directory for local imports
LOG_DIR = _test_dir
def pytest_configure(config):
"""Ensure test directory is in Python path before collection."""
test_dir = Path(__file__).parent
if str(test_dir) not in sys.path:
sys.path.insert(0, str(test_dir))
MOCK_SERVER_PORT = 8090
MOCK_SERVER_URL = f"http://localhost:{MOCK_SERVER_PORT}"
LITELLM_PROXY_PORT = 4000
LITELLM_PROXY_URL = f"http://localhost:{LITELLM_PROXY_PORT}"
DATABASE_URL = "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm"
def kill_process_on_port(port: int) -> None:
"""Kill any process using the specified port."""
try:
result = subprocess.run(
["lsof", "-ti", f":{port}"],
capture_output=True,
text=True,
timeout=5,
)
if result.stdout.strip():
pids = result.stdout.strip().split("\n")
for pid in pids:
try:
subprocess.run(["kill", "-9", pid.strip()], timeout=5)
except Exception:
pass
time.sleep(1)
except Exception:
pass
def wait_for_server(url: str, max_attempts: int = 30, delay: float = 1.0) -> bool:
"""Wait for a server to become available at url/health.
Any HTTP response (including 401) means the server is up.
Only connection errors count as "not ready yet".
"""
for attempt in range(max_attempts):
try:
response = httpx.get(f"{url}/health", timeout=2.0)
return True
except (httpx.ConnectError, httpx.TimeoutException, httpx.NetworkError):
pass
except Exception:
pass
if attempt < max_attempts - 1:
time.sleep(delay)
return False
def _read_log_tail(log_path: Path, max_lines: int = 80) -> str:
"""Read the last N lines of a log file, returning empty string if not found."""
if not log_path.exists():
return "(log file not found)"
try:
text = log_path.read_text()
lines = text.strip().splitlines()
if len(lines) > max_lines:
return f"... ({len(lines) - max_lines} lines truncated) ...\n" + "\n".join(
lines[-max_lines:]
)
return text
except Exception as e:
return f"(error reading log: {e})"
def _check_process_alive(process: subprocess.Popen, label: str, log_path: Path):
"""Check if a subprocess crashed immediately after starting.
Raises pytest.fail with log output if the process has already exited.
"""
time.sleep(1)
exit_code = process.poll()
if exit_code is not None:
log_output = _read_log_tail(log_path)
pytest.fail(
f"{label} exited immediately with code {exit_code}.\n"
f"--- {label} log ({log_path}) ---\n{log_output}\n"
f"--- end log ---"
)
def setup_database() -> bool:
"""Ensure PostgreSQL database exists and is accessible."""
try:
import psycopg2
conn = psycopg2.connect(
host="localhost",
port=5432,
database="litellm",
user="llmproxy",
password="dbpassword9090",
connect_timeout=5,
)
conn.close()
return True
except ImportError:
print("WARNING: psycopg2 not installed — cannot verify database")
return False
except Exception:
return False
@pytest.fixture(scope="session")
def mock_azure_server() -> Generator[str, None, None]:
"""Start mock Azure batch server as a subprocess."""
print(f"\n{'=' * 60}")
print("Setting up Mock Azure Batch Server")
print(f"{'=' * 60}")
kill_process_on_port(MOCK_SERVER_PORT)
runner_script = Path(__file__).parent / "fixtures" / "run_mock_server.py"
runner_script.write_text(
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
from fixtures.mock_azure_batch_server import create_mock_azure_batch_server
import uvicorn
if __name__ == "__main__":
app = create_mock_azure_batch_server()
uvicorn.run(app, host="0.0.0.0", port=8090, log_level="info", access_log=False)
"""
)
mock_log = LOG_DIR / "mock_server.log"
log_file = open(mock_log, "w")
print(f"Starting mock server on port {MOCK_SERVER_PORT}...")
print(f"Log file: {mock_log}")
process = subprocess.Popen(
[sys.executable, str(runner_script)],
stdout=log_file,
stderr=subprocess.STDOUT,
cwd=Path(__file__).parent,
)
_check_process_alive(process, "Mock server", mock_log)
if not wait_for_server(MOCK_SERVER_URL, max_attempts=30, delay=1.0):
log_output = _read_log_tail(mock_log)
exit_code = process.poll()
process.terminate()
try:
process.wait(timeout=5)
except subprocess.TimeoutExpired:
process.kill()
process.wait()
log_file.close()
pytest.fail(
f"Mock server failed to start on port {MOCK_SERVER_PORT} "
f"(process exit_code={exit_code}).\n"
f"--- mock server log ---\n{log_output}\n--- end log ---\n"
f"Hint: ensure 'uvicorn' and 'fastapi' are installed."
)
print(f"Mock Azure server ready at {MOCK_SERVER_URL}")
yield MOCK_SERVER_URL
print("\nShutting down mock server...")
try:
process.terminate()
process.wait(timeout=5)
except subprocess.TimeoutExpired:
process.kill()
process.wait()
log_file.close()
print("Mock server stopped")
@pytest.fixture(scope="session")
def litellm_proxy_server(mock_azure_server: str) -> Generator[str, None, None]:
"""Start LiteLLM proxy server for the test session."""
print(f"\n{'=' * 60}")
print("Setting up LiteLLM Proxy Server")
print(f"{'=' * 60}")
if not setup_database():
pytest.skip(
"PostgreSQL database not available at localhost:5432. "
"Start PostgreSQL and create a 'litellm' database:\n"
" docker run -d --name litellm-db -p 5432:5432 "
'-e POSTGRES_USER=llmproxy -e POSTGRES_PASSWORD=dbpassword9090 '
"-e POSTGRES_DB=litellm postgres:15\n"
"Then run: prisma db push --schema=litellm/proxy/schema.prisma"
)
print("Database connection verified")
config_path = Path(__file__).parent / "fixtures" / "config.yml"
if not config_path.exists():
pytest.fail(f"Config file not found: {config_path}")
print("Config file found")
kill_process_on_port(LITELLM_PROXY_PORT)
os.environ["MOCK_SERVER_URL_V1"] = f"{mock_azure_server}/v1"
os.environ["MOCK_SERVER_URL_OPENAI_V1"] = f"{mock_azure_server}/openai/v1"
os.environ["DATABASE_URL"] = DATABASE_URL
os.environ["USE_LOCAL_LITELLM"] = "true"
os.environ["USE_MOCK_MODELS"] = "true"
os.environ["USE_STATE_TRACKER"] = "true"
os.environ["PROXY_BATCH_POLLING_INTERVAL"] = "10"
print("Environment configured")
print(f"Starting LiteLLM proxy on port {LITELLM_PROXY_PORT}...")
litellm_root = Path(__file__).parent.parent.parent
cmd = [
sys.executable,
"-m",
"litellm.proxy.proxy_cli",
"--config",
str(config_path),
"--port",
str(LITELLM_PROXY_PORT),
"--detailed_debug",
]
proxy_log = LOG_DIR / "proxy_server.log"
log_file = open(proxy_log, "w")
print(f"Log file: {proxy_log}")
process = subprocess.Popen(
cmd,
stdout=log_file,
stderr=subprocess.STDOUT,
env=os.environ.copy(),
cwd=litellm_root,
)
_check_process_alive(process, "LiteLLM proxy", proxy_log)
if not wait_for_server(LITELLM_PROXY_URL, max_attempts=60, delay=1.0):
log_output = _read_log_tail(proxy_log)
exit_code = process.poll()
process.terminate()
try:
process.wait(timeout=5)
except subprocess.TimeoutExpired:
process.kill()
process.wait()
log_file.close()
pytest.fail(
f"LiteLLM proxy failed to start on port {LITELLM_PROXY_PORT} "
f"(process exit_code={exit_code}).\n"
f"--- proxy log (last 80 lines) ---\n{log_output}\n--- end log ---\n"
f"Hints:\n"
f" 1. Ensure Prisma client is generated: "
f"cd {litellm_root} && prisma generate --schema=litellm/proxy/schema.prisma\n"
f" 2. Ensure DB migrations are applied: "
f"prisma db push --schema=litellm/proxy/schema.prisma\n"
f" 3. Check the full log at: {proxy_log}"
)
print(f"LiteLLM proxy ready at {LITELLM_PROXY_URL}")
yield LITELLM_PROXY_URL
print("\nShutting down LiteLLM proxy...")
try:
process.terminate()
process.wait(timeout=10)
except subprocess.TimeoutExpired:
process.kill()
process.wait()
log_file.close()
print("LiteLLM proxy stopped")
@pytest.fixture(scope="session")
def event_loop():
"""Provide an event loop for async tests."""
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
yield loop
loop.close()

View file

@ -0,0 +1,56 @@
model_list:
- model_name: openai-fake-gpt-3.5-turbo
litellm_params:
model: openai/openai-fake-gpt-3.5-turbo
api_base: os.environ/MOCK_SERVER_URL_V1
api_key: fake-key
- model_name: openai-fake-gpt-4
litellm_params:
model: openai/openai-fake-gpt-4
api_base: os.environ/MOCK_SERVER_URL_V1
api_key: fake-key
- model_name: openai-fake-gpt-4o
litellm_params:
model: openai/openai-fake-gpt-4o
api_base: os.environ/MOCK_SERVER_URL_V1
api_key: fake-key
- model_name: fake-text-embedding-3-small
litellm_params:
model: openai/fake-text-embedding-3-small
api_base: os.environ/MOCK_SERVER_URL_V1
api_key: fake-key
- model_name: o3-mini-batch-2025-01-31
litellm_params:
model: openai/o3-mini-batch-2025-01-31
api_base: os.environ/MOCK_SERVER_URL_OPENAI_V1
api_key: fake-key
model_info:
mode: batch
- model_name: azure-fake-gpt-5-batch-2025-08-07
litellm_params:
api_base: http://0.0.0.0:8090
api_key: fake-key
api_version: 2025-03-01-preview
base_model: azure/gpt-5
model: azure/gpt-5-mini
custom_llm_provider: azure
general_settings:
master_key: sk-1234
database_url: os.environ/DATABASE_URL
proxy_batch_polling_interval: 10
litellm_settings:
drop_params: true
set_verbose: true
json_logs: true
# S3 callback for batch completion logging (points to mock server)
callbacks: ["s3_v2"]
s3_callback_params:
s3_bucket_name: litellm-test-bucket
s3_region_name: us-east-1
s3_endpoint_url: http://0.0.0.0:8090
s3_aws_access_key_id: fake-key
s3_aws_secret_access_key: fake-secret
s3_use_ssl: false
s3_verify: false

View file

@ -0,0 +1,3 @@
from .server import create_mock_azure_batch_server
__all__ = ["create_mock_azure_batch_server"]

View file

@ -0,0 +1,517 @@
import asyncio
import io
import json
import logging
import time
import uuid
from typing import Dict, List, Optional
from fastapi import FastAPI, HTTPException, Query, Request, UploadFile
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class FileObject(BaseModel):
id: str
object: str = "file"
bytes: int
created_at: int
filename: str
purpose: str
status: str = "processed"
status_details: Optional[str] = None
expires_at: Optional[int] = None
class BatchObject(BaseModel):
id: str
object: str = "batch"
endpoint: str
errors: Optional[Dict] = None
input_file_id: str
completion_window: str
status: str
output_file_id: Optional[str] = None
error_file_id: Optional[str] = None
created_at: int
in_progress_at: Optional[int] = None
expires_at: Optional[int] = None
finalizing_at: Optional[int] = None
completed_at: Optional[int] = None
failed_at: Optional[int] = None
expired_at: Optional[int] = None
cancelling_at: Optional[int] = None
cancelled_at: Optional[int] = None
request_counts: Optional[Dict[str, int]] = None
metadata: Optional[Dict] = None
class BatchListResponse(BaseModel):
object: str = "list"
data: List[Dict]
first_id: Optional[str] = None
last_id: Optional[str] = None
has_more: bool = False
file_storage: Dict[str, Dict] = {}
batch_storage: Dict[str, BatchObject] = {}
batch_results: Dict[str, List[Dict]] = {}
PROCESSING_DELAY_SECONDS = float(1)
VALIDATING_DELAY_SECONDS = float(3)
async def process_batch(batch_id: str):
logger.info(f"Starting batch processing for {batch_id}")
try:
batch = batch_storage[batch_id]
await asyncio.sleep(VALIDATING_DELAY_SECONDS)
batch.status = "in_progress"
batch.in_progress_at = int(time.time())
logger.info(f"Batch {batch_id} status: in_progress")
await process_batch_requests(batch_id)
await asyncio.sleep(PROCESSING_DELAY_SECONDS)
batch.status = "finalizing"
batch.finalizing_at = int(time.time())
logger.info(f"Batch {batch_id} status: finalizing")
await asyncio.sleep(PROCESSING_DELAY_SECONDS)
await create_output_file(batch_id)
batch.status = "completed"
batch.completed_at = int(time.time())
logger.info(f"Batch {batch_id} status: completed")
except Exception as e:
logger.error(f"Batch {batch_id} failed: {e}")
batch = batch_storage[batch_id]
batch.status = "failed"
batch.failed_at = int(time.time())
batch.errors = {
"object": "list",
"data": [{"code": "processing_error", "message": str(e)}],
}
async def process_batch_requests(batch_id: str):
batch = batch_storage[batch_id]
input_file = file_storage[batch.input_file_id]
requests = []
for line in input_file["content"].split("\n"):
if line.strip():
try:
requests.append(json.loads(line))
except json.JSONDecodeError as e:
logger.warning(f"Invalid JSON line in batch {batch_id}: {e}")
logger.info(f"Batch {batch_id} has {len(requests)} requests")
results = []
failed_count = 0
for req in requests:
result = await process_single_request(req)
if result.get("error"):
failed_count += 1
results.append(result)
batch_results[batch_id] = results
batch.request_counts = {
"total": len(requests),
"completed": len(results) - failed_count,
"failed": failed_count,
}
async def process_single_request(request_data: Dict) -> Dict:
custom_id = request_data.get("custom_id")
url = request_data.get("url", "/v1/chat/completions")
body = request_data.get("body", {})
if "/chat/completions" in url:
response_body = {
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": int(time.time()),
"model": body.get("model", "gpt-4o"),
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "Mock batch response."},
"finish_reason": "stop",
},
],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
}
status_code = 200
else:
response_body = {"error": {"message": f"Unsupported endpoint: {url}"}}
status_code = 400
return {
"id": f"batch_req_{uuid.uuid4().hex[:12]}",
"custom_id": custom_id,
"response": {
"status_code": status_code,
"request_id": f"req_{uuid.uuid4().hex[:12]}",
"body": response_body,
},
"error": None,
}
async def create_output_file(batch_id: str):
results = batch_results.get(batch_id, [])
output_lines = [json.dumps(result) for result in results]
output_content = "\n".join(output_lines)
output_file_id = f"file-batch-output-{uuid.uuid4().hex[:12]}"
file_storage[output_file_id] = {
"content": output_content,
"filename": f"batch_output_{batch_id}.jsonl",
"purpose": "batch_output",
"bytes": len(output_content.encode()),
"created_at": int(time.time()),
}
batch = batch_storage[batch_id]
batch.output_file_id = output_file_id
logger.info(f"Created output file {output_file_id} for batch {batch_id}")
def validate_batch_input(content: str) -> tuple[bool, str, List[Dict]]:
requests = []
custom_ids = set()
lines = content.strip().split("\n")
if not lines or all(not line.strip() for line in lines):
return False, "empty_batch", []
for line_num, line in enumerate(lines, 1):
if not line.strip():
continue
try:
req = json.loads(line)
except json.JSONDecodeError:
return False, "invalid_json_line", []
for field in ["custom_id", "method", "url", "body"]:
if field not in req:
return False, "invalid_request", []
if req["custom_id"] in custom_ids:
return False, "duplicate_custom_id", []
custom_ids.add(req["custom_id"])
requests.append(req)
if len(requests) > 100000:
return False, "too_many_tasks", []
return True, "", requests
def setup_batch_routes(app: FastAPI):
# Files endpoints (OpenAI and Azure paths)
@app.post("/openai/v1/files")
@app.post("/openai/files")
@app.post("/v1/files")
@app.post("/files")
async def create_file(request: Request):
form = await request.form()
logger.info(f"File upload form fields: {list(form.keys())}")
file: UploadFile = form.get("file")
purpose: str = form.get("purpose", "batch")
if not file:
raise HTTPException(status_code=400, detail="No file provided")
logger.info(f"Uploading file: {file.filename}, purpose: {purpose}")
content = await file.read()
content_str = content.decode("utf-8")
file_id = f"file-{uuid.uuid4().hex[:24]}"
created_at = int(time.time())
expires_at = None
expires_after_seconds = form.get("expires_after[seconds]")
if expires_after_seconds:
try:
seconds = int(expires_after_seconds)
logger.info(f"expires_after[seconds] = {seconds}")
if seconds < 259200 or seconds > 2592000:
raise HTTPException(
status_code=400,
detail={
"error": {
"code": "invalidPayload",
"message": "Value for Seconds must be between 259200 and 2592000.",
},
},
)
expires_at = created_at + seconds
logger.info(f"Calculated expires_at: {expires_at}")
except ValueError as e:
logger.warning(f"Failed to parse expires_after[seconds]: {e}")
file_storage[file_id] = {
"content": content_str,
"filename": file.filename or "batch_input.jsonl",
"purpose": purpose,
"bytes": len(content),
"created_at": created_at,
"expires_at": expires_at,
}
logger.info(f"Created file {file_id}, expires_at={expires_at}")
return FileObject(
id=file_id,
bytes=len(content),
created_at=created_at,
filename=file.filename or "batch_input.jsonl",
purpose=purpose,
expires_at=expires_at,
).model_dump()
@app.get("/openai/v1/files/{file_id}")
@app.get("/openai/files/{file_id}")
@app.get("/v1/files/{file_id}")
@app.get("/files/{file_id}")
async def get_file(file_id: str):
logger.info(f"Getting file: {file_id}")
if file_id not in file_storage:
raise HTTPException(status_code=404, detail="File not found")
file_data = file_storage[file_id]
return FileObject(
id=file_id,
bytes=file_data["bytes"],
created_at=file_data["created_at"],
filename=file_data["filename"],
purpose=file_data["purpose"],
expires_at=file_data.get("expires_at"),
).model_dump()
@app.get("/openai/v1/files/{file_id}/content")
@app.get("/openai/files/{file_id}/content")
@app.get("/v1/files/{file_id}/content")
@app.get("/files/{file_id}/content")
async def get_file_content(file_id: str):
logger.info(f"Getting file content: {file_id}")
if file_id not in file_storage:
raise HTTPException(status_code=404, detail="File not found")
file_data = file_storage[file_id]
content = file_data["content"]
return StreamingResponse(
io.StringIO(content),
media_type="application/octet-stream",
headers={
"Content-Disposition": f"attachment; filename={file_data['filename']}",
},
)
@app.delete("/openai/v1/files/{file_id}")
@app.delete("/openai/files/{file_id}")
@app.delete("/v1/files/{file_id}")
@app.delete("/files/{file_id}")
async def delete_file(file_id: str):
logger.info(f"Deleting file: {file_id}")
if file_id not in file_storage:
raise HTTPException(status_code=404, detail="File not found")
del file_storage[file_id]
return {"id": file_id, "object": "file", "deleted": True}
@app.get("/openai/v1/files")
@app.get("/openai/files")
@app.get("/v1/files")
@app.get("/files")
async def list_files(
purpose: Optional[str] = None,
limit: int = Query(10000, le=10000),
):
logger.info(f"Listing files, purpose: {purpose}, limit: {limit}")
files = []
for file_id, file_data in file_storage.items():
if purpose is None or file_data.get("purpose") == purpose:
files.append(
FileObject(
id=file_id,
bytes=file_data["bytes"],
created_at=file_data["created_at"],
filename=file_data["filename"],
purpose=file_data["purpose"],
expires_at=file_data.get("expires_at"),
).model_dump(),
)
return {"object": "list", "data": files[:limit]}
# Batches endpoints (OpenAI and Azure paths)
@app.post("/openai/v1/batches")
@app.post("/openai/batches")
@app.post("/v1/batches")
@app.post("/batches")
async def create_batch(request_data: dict):
input_file_id = request_data.get("input_file_id")
endpoint = request_data.get("endpoint", "/v1/chat/completions")
completion_window = request_data.get("completion_window", "24h")
metadata = request_data.get("metadata", {})
output_expires_after = request_data.get("output_expires_after")
logger.info(
f"Creating batch with input_file: {input_file_id}, endpoint: {endpoint}, output_expires_after: {output_expires_after}",
)
if not input_file_id or input_file_id not in file_storage:
raise HTTPException(status_code=400, detail="Input file not found")
input_file = file_storage[input_file_id]
is_valid, error_code, _ = validate_batch_input(input_file["content"])
if not is_valid:
raise HTTPException(
status_code=400,
detail={
"error": {
"code": error_code,
"message": f"Validation failed: {error_code}",
},
},
)
batch_id = f"batch_{uuid.uuid4()}"
created_at = int(time.time())
if output_expires_after:
seconds = (
output_expires_after.get("seconds", 0)
if isinstance(output_expires_after, dict)
else 0
)
expires_at = created_at + seconds
logger.info(
f"Using output_expires_after: {seconds}s, expires_at: {expires_at}",
)
elif completion_window == "24h":
expires_at = created_at + (24 * 60 * 60)
else:
expires_at = created_at + (24 * 60 * 60)
batch = BatchObject(
id=batch_id,
endpoint=endpoint,
input_file_id=input_file_id,
completion_window=completion_window,
status="validating",
created_at=created_at,
expires_at=expires_at,
request_counts={"total": 0, "completed": 0, "failed": 0},
metadata=metadata,
)
batch_storage[batch_id] = batch
logger.info(f"Created batch {batch_id}")
asyncio.create_task(process_batch(batch_id))
return batch.model_dump()
@app.get("/openai/v1/batches/{batch_id}")
@app.get("/openai/batches/{batch_id}")
@app.get("/v1/batches/{batch_id}")
@app.get("/batches/{batch_id}")
async def get_batch(batch_id: str):
logger.info(f"Getting batch: {batch_id}")
if batch_id not in batch_storage:
raise HTTPException(status_code=404, detail="Batch not found")
return batch_storage[batch_id].model_dump()
@app.get("/openai/v1/batches")
@app.get("/openai/batches")
@app.get("/v1/batches")
@app.get("/batches")
async def list_batches(
after: Optional[str] = Query(None),
limit: int = Query(20, le=100),
):
logger.info(f"Listing batches, after: {after}, limit: {limit}")
batches = list(batch_storage.values())
batches.sort(key=lambda x: x.created_at, reverse=True)
if after:
after_index = next((i for i, b in enumerate(batches) if b.id == after), -1)
if after_index >= 0:
batches = batches[after_index + 1 :]
batches = batches[:limit]
return BatchListResponse(
data=[batch.model_dump() for batch in batches],
first_id=batches[0].id if batches else None,
last_id=batches[-1].id if batches else None,
has_more=len(batches) == limit,
).model_dump()
@app.post("/openai/v1/batches/{batch_id}/cancel")
@app.post("/openai/batches/{batch_id}/cancel")
@app.post("/v1/batches/{batch_id}/cancel")
@app.post("/batches/{batch_id}/cancel")
async def cancel_batch(batch_id: str):
logger.info(f"Cancelling batch: {batch_id}")
if batch_id not in batch_storage:
raise HTTPException(status_code=404, detail="Batch not found")
batch = batch_storage[batch_id]
if batch.status in ["completed", "failed", "cancelled", "expired"]:
raise HTTPException(
status_code=400,
detail=f"Cannot cancel batch in {batch.status} status",
)
batch.status = "cancelled"
batch.cancelled_at = int(time.time())
logger.info(f"Batch {batch_id} cancelled")
return batch.model_dump()
# Debug endpoints
@app.get("/debug/batches")
async def debug_list_batches():
return {
"batches": {
batch_id: batch.model_dump()
for batch_id, batch in batch_storage.items()
},
"files": {
file_id: {k: v for k, v in data.items() if k != "content"}
for file_id, data in file_storage.items()
},
}
@app.post("/reset")
@app.post("/debug/clear")
async def reset_all():
file_storage.clear()
batch_storage.clear()
batch_results.clear()
logger.info("All data cleared")
return {"message": "All data cleared"}
@app.get("/debug/status")
async def debug_status():
return {
"files_count": len(file_storage),
"batches_count": len(batch_storage),
"batch_statuses": {bid: b.status for bid, b in batch_storage.items()},
}

View file

@ -0,0 +1,124 @@
import json
import time
import uuid
from datetime import datetime
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
def get_request_details(request: Request, body: dict = None) -> str:
details = {
"method": request.method,
"url": str(request.url),
"path": request.url.path,
"headers": dict(request.headers),
"query_params": dict(request.query_params),
}
return json.dumps(details, indent=2)
def data_generator(response_details: str, model: str):
response_id = uuid.uuid4().hex
content = response_details
chunk_size = 50
for i in range(0, len(content), chunk_size):
text_chunk = content[i : i + chunk_size]
chunk = {
"id": f"chatcmpl-{response_id}",
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": model,
"choices": [{"index": 0, "delta": {"content": text_chunk}}],
}
yield f"data: {json.dumps(chunk)}\n\n"
final_chunk = {
"id": f"chatcmpl-{response_id}",
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": model,
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
}
yield f"data: {json.dumps(final_chunk)}\n\n"
yield "data: [DONE]\n\n"
def setup_chat_routes(app: FastAPI):
@app.post("/chat/completions")
@app.post("/v1/chat/completions")
@app.post("/openai/deployments/{model:path}/chat/completions")
async def completion(request: Request):
data = await request.json()
model = data.get("model", "unknown")
request_details = get_request_details(request, data)
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
response_details = f"Request:{request_details}, Canned Response:{timestamp}"
if data.get("stream"):
return StreamingResponse(
content=data_generator(response_details, model),
media_type="text/event-stream",
)
else:
response_id = uuid.uuid4().hex
response = {
"id": f"chatcmpl-{response_id}",
"object": "chat.completion",
"created": int(time.time()),
"model": model,
"system_fingerprint": "fp_mock_server",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": response_details,
},
"logprobs": None,
"finish_reason": "stop",
},
],
"usage": {
"prompt_tokens": 9,
"completion_tokens": 12,
"total_tokens": 21,
},
}
return response
@app.post("/completions")
@app.post("/v1/completions")
async def text_completion(request: Request):
data = await request.json()
model = data.get("model", "unknown")
request_details = get_request_details(request, data)
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
response_details = f"Request:{request_details}, Canned Response:{timestamp}"
if data.get("stream"):
return StreamingResponse(
content=data_generator(response_details, model),
media_type="text/event-stream",
)
else:
response = {
"id": f"cmpl-{uuid.uuid4().hex}",
"choices": [
{
"finish_reason": "stop",
"index": 0,
"logprobs": None,
"text": response_details,
},
],
"created": int(time.time()),
"model": model,
"object": "text_completion",
"system_fingerprint": None,
"usage": {
"completion_tokens": 16,
"prompt_tokens": 10,
"total_tokens": 26,
},
}
return response

View file

@ -0,0 +1,23 @@
from fastapi import FastAPI, Request
def setup_embeddings_routes(app: FastAPI):
@app.post("/embeddings")
@app.post("/v1/embeddings")
@app.post("/openai/deployments/{model:path}/embeddings")
async def embeddings(request: Request):
data = await request.json()
model = data.get("model", "unknown")
_small_embedding = [
-0.006929283495992422,
-0.005336422007530928,
-4.547132266452536e-05,
-0.024047505110502243,
]
big_embedding = _small_embedding * 100
return {
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": big_embedding}],
"model": model,
"usage": {"prompt_tokens": 5, "total_tokens": 5},
}

View file

@ -0,0 +1,170 @@
import json
import re
import time
import uuid
from datetime import datetime
from typing import Any
from fastapi import FastAPI, Request, HTTPException
# Header to identify which model/deployment this request targets (simulates Azure model-specific encryption).
# When set, the mock validates that encrypted_content in input was produced by this model.
MOCK_AZURE_MODEL_HEADER = "X-Mock-Azure-Model"
# Prefix we use in mock encrypted_content: gAAA_model_<model_id>_<32hex uuid>
# Model id can contain underscores (e.g. gpt-5.1-codex-openai-2).
ENCRYPTED_CONTENT_MODEL_PREFIX = re.compile(r"^gAAA_model_(.+)_[0-9a-f]{32}$")
def _extract_model_from_encrypted_content(encrypted: str) -> str | None:
"""Extract model id from our mock encrypted_content format, or None if not our format."""
if not isinstance(encrypted, str) or not encrypted.startswith("gAAA"):
return None
m = ENCRYPTED_CONTENT_MODEL_PREFIX.match(encrypted)
return m.group(1) if m else None
def _collect_encrypted_contents(obj, out: list[str]) -> None:
"""Recursively collect all encrypted_content string values from input structure."""
if isinstance(obj, dict):
if "encrypted_content" in obj and obj["encrypted_content"]:
out.append(obj["encrypted_content"])
for v in obj.values():
_collect_encrypted_contents(v, out)
elif isinstance(obj, list):
for item in obj:
_collect_encrypted_contents(item, out)
def _validate_encrypted_content_model(request_model: str | None, input_data: Any) -> str | None:
"""
If request_model is set, check that all encrypted_content in input was produced by this model.
Returns error message if validation fails, else None.
Content with our format (gAAA_model_<id>_) must match request_model.
"""
if not request_model:
return None
encrypted_values: list[str] = []
_collect_encrypted_contents(input_data, encrypted_values)
for enc in encrypted_values:
content_model = _extract_model_from_encrypted_content(enc)
if content_model is not None and content_model != request_model:
err = enc[:50] + "..." if len(enc) > 50 else enc
return f"The encrypted content {err} could not be verified."
return None
def get_request_details(request: Request, body: dict = None) -> str:
details = {
"method": request.method,
"url": str(request.url),
"path": request.url.path,
"headers": dict(request.headers),
"query_params": dict(request.query_params),
}
return json.dumps(details, indent=2)
def setup_responses_routes(app: FastAPI):
@app.post("/responses")
@app.post("/v1/responses")
@app.post("/openai/responses")
async def responses_api(request: Request):
data = await request.json()
model = data.get("model", "unknown")
# Simulate Azure: encrypted content from one model cannot be verified by another.
input_data = data.get("input")
err_msg = _validate_encrypted_content_model(model, input_data)
if err_msg is not None:
raise HTTPException(
status_code=400,
detail={
"error": {
"message": err_msg,
"type": "invalid_request_error",
"param": None,
"code": "invalid_encrypted_content",
}
},
)
request_details = get_request_details(request, data)
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
response_details = f"Request:{request_details}, Canned Response:{timestamp}"
response_id = uuid.uuid4().hex
message_id = f"msg_{uuid.uuid4().hex[:34]}"
reasoning_id = f"rs_{uuid.uuid4().hex[:34]}"
output_items: list[dict[str, Any]] = [
{
"id": message_id,
"content": [
{
"annotations": [],
"text": response_details,
"type": "output_text",
"logprobs": [],
},
],
"role": "assistant",
"status": "completed",
"type": "message",
},
]
if model:
output_items.append(
{
"id": reasoning_id,
"type": "reasoning",
"status": "completed",
"encrypted_content": f"gAAA_model_{model}_{uuid.uuid4().hex}",
}
)
return {
"id": f"resp_{response_id}",
"created_at": int(time.time()),
"error": None,
"incomplete_details": None,
"instructions": None,
"metadata": {},
"model": model,
"object": "response",
"output": output_items,
"parallel_tool_calls": True,
"temperature": data.get("temperature", 1.0),
"tool_choice": data.get("tool_choice", "auto"),
"tools": data.get("tools", []),
"top_p": data.get("top_p", 1.0),
"max_output_tokens": data.get("max_output_tokens"),
"previous_response_id": None,
"reasoning": {"effort": None, "summary": None},
"status": "completed",
"text": {"format": {"type": "text"}, "verbosity": "medium"},
"truncation": "disabled",
"usage": {
"input_tokens": 11,
"input_tokens_details": {
"audio_tokens": None,
"cached_tokens": 0,
"text_tokens": None,
},
"output_tokens": 19,
"output_tokens_details": {"reasoning_tokens": 0, "text_tokens": None},
"total_tokens": 30,
"cost": None,
},
"user": None,
"store": True,
"background": False,
"content_filters": None,
"max_tool_calls": None,
"prompt_cache_key": None,
"safety_identifier": None,
"service_tier": "default",
"top_logprobs": 0,
}

View file

@ -0,0 +1,98 @@
"""
Mock S3 callback receiver for testing LiteLLM S3 callbacks.
This module provides S3-compatible endpoints that capture callback data
sent by LiteLLM's s3_v2 callback handler after batch completion.
"""
import json
import logging
import time
from typing import Any, Dict, List, Optional
from fastapi import FastAPI, Request
from pydantic import BaseModel
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class S3CallbackRecord(BaseModel):
key: str
bucket: str
content: Dict[str, Any]
timestamp: int
content_type: Optional[str] = None
callback_storage: List[S3CallbackRecord] = []
def setup_s3_callback_routes(app: FastAPI):
@app.put("/{bucket}/{key:path}")
async def s3_put_object(bucket: str, key: str, request: Request):
content_type = request.headers.get("content-type", "application/json")
body = await request.body()
try:
content = json.loads(body.decode("utf-8"))
except (json.JSONDecodeError, UnicodeDecodeError):
content = {"raw": body.decode("utf-8", errors="replace")}
record = S3CallbackRecord(
key=key,
bucket=bucket,
content=content,
timestamp=int(time.time()),
content_type=content_type,
)
callback_storage.append(record)
logger.info(f"S3 callback received: bucket={bucket}, key={key}")
logger.debug(f"Callback content: {json.dumps(content, indent=2)[:500]}")
return {
"ETag": f'"{hash(body)}"',
"VersionId": None,
}
@app.get("/mock-s3/callbacks")
async def list_callbacks(
bucket: Optional[str] = None,
key_prefix: Optional[str] = None,
limit: int = 100,
):
results = callback_storage
if bucket:
results = [r for r in results if r.bucket == bucket]
if key_prefix:
results = [r for r in results if r.key.startswith(key_prefix)]
return {
"count": len(results),
"callbacks": [r.model_dump() for r in results[-limit:]],
}
@app.get("/mock-s3/callbacks/count")
async def count_callbacks(bucket: Optional[str] = None):
if bucket:
count = sum(1 for r in callback_storage if r.bucket == bucket)
else:
count = len(callback_storage)
return {"count": count}
@app.get("/mock-s3/callbacks/latest")
async def get_latest_callback():
if not callback_storage:
return {"callback": None}
return {"callback": callback_storage[-1].model_dump()}
@app.delete("/mock-s3/callbacks")
async def clear_callbacks():
count = len(callback_storage)
callback_storage.clear()
logger.info(f"Cleared {count} S3 callbacks")
return {"cleared": count}

View file

@ -0,0 +1,33 @@
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from .mock_azure_batch import setup_batch_routes
from .mock_chat import setup_chat_routes
from .mock_embeddings import setup_embeddings_routes
from .mock_responses import setup_responses_routes
from .mock_s3_callback import setup_s3_callback_routes
def create_mock_azure_batch_server() -> FastAPI:
"""Create a FastAPI app that mocks Azure Batch API and S3 callbacks."""
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.get("/health")
async def health():
return {"status": "ok"}
setup_chat_routes(app)
setup_responses_routes(app)
setup_embeddings_routes(app)
setup_batch_routes(app)
setup_s3_callback_routes(app)
return app

View file

@ -0,0 +1,12 @@
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
from fixtures.mock_azure_batch_server import create_mock_azure_batch_server
import uvicorn
if __name__ == "__main__":
app = create_mock_azure_batch_server()
uvicorn.run(app, host="0.0.0.0", port=8090, log_level="info", access_log=False)

View file

@ -0,0 +1,41 @@
"""
Smoke test to verify fixtures start and stop correctly.
Run this first to ensure the infrastructure works before running full E2E tests.
"""
import httpx
import pytest
pytestmark = pytest.mark.usefixtures("mock_azure_server", "litellm_proxy_server")
def test_mock_server_health(mock_azure_server):
"""Verify mock Azure server is running and healthy."""
response = httpx.get(f"{mock_azure_server}/health", timeout=5.0)
assert response.status_code == 200
assert response.json() == {"status": "ok"}
print(f"✓ Mock Azure server is healthy at {mock_azure_server}")
def test_litellm_proxy_health(litellm_proxy_server):
"""Verify LiteLLM proxy is running and healthy."""
response = httpx.get(f"{litellm_proxy_server}/health", timeout=5.0)
assert response.status_code == 200
print(f"✓ LiteLLM proxy is healthy at {litellm_proxy_server}")
def test_litellm_proxy_model_list(litellm_proxy_server):
"""Verify LiteLLM proxy can list models."""
response = httpx.get(
f"{litellm_proxy_server}/v1/models",
headers={"Authorization": "Bearer sk-1234"},
timeout=5.0,
)
assert response.status_code == 200
data = response.json()
assert "data" in data
models = [m["id"] for m in data["data"]]
print(f"✓ LiteLLM proxy has {len(models)} models configured")
assert "azure-fake-gpt-5-batch-2025-08-07" in models
print(f"✓ Azure batch model is configured")

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,323 @@
import base64
import os
import sys
import time
import warnings
import httpx
import openai
import pytest
from tenacity import RetryError
sys.path.insert(0, os.path.abspath("../.."))
from base_integration_test import (
get_mock_server_base_url,
model_id,
use_mock_models,
UserKeyTestMixin,
)
from test_managed_files_base import (
ManagedFilesBase,
MIN_EXPIRY_SECONDS,
get_batch_model_names,
)
MANAGED_FILE_ID_PREFIX = "litellm_proxy"
pytestmark = [
pytest.mark.usefixtures("mock_azure_server", "litellm_proxy_server"),
pytest.mark.skipif(
os.environ.get("SKIP_E2E_TESTS", "false").lower() == "true",
reason="E2E tests disabled via SKIP_E2E_TESTS env var"
),
]
def is_managed_id(file_id: str) -> bool:
"""Check if a file ID is a base64-encoded LiteLLM managed/unified ID."""
try:
padded = file_id + "=" * (-len(file_id) % 4)
decoded = base64.urlsafe_b64decode(padded).decode()
return decoded.startswith(MANAGED_FILE_ID_PREFIX)
except Exception:
return False
def assert_managed_id(file_id: str, label: str):
assert is_managed_id(file_id), f"{label} should be a managed ID, got raw: {file_id}"
def wip_features_enabled() -> bool:
return os.environ.get("WIP_FEATURES", "").lower() == "true"
class TestManagedFilesAPI(ManagedFilesBase, UserKeyTestMixin):
@classmethod
def setup_class(cls):
super().setup_class()
cls.setup_admin_client()
@classmethod
def teardown_class(cls):
cls.teardown_admin_client()
@pytest.fixture(autouse=True)
def setup_test(self):
print(
f"\nBase URL: {self.base_url}, Using mock models: {use_mock_models()}",
)
self.clear_s3_callbacks()
user_id, api_key, user_email, client = self.create_user_key_and_client(
"e2e-batch",
)
self.test_user_id = user_id
self.openai_client = client
print(f"Using user {user_email} (id={user_id})")
def _create_and_verify_batch_input_file(self, tmp_path, model_name):
request_file = self.create_batch_request_file_on_disk(tmp_path, model_name)
print("Creating batch input file...")
batch_input_file = self.create_batch_input_file(
self.openai_client,
request_file,
MIN_EXPIRY_SECONDS,
target_model_names=model_name,
)
print(f"Created batch input file: {self.shorten_id(batch_input_file.id)}")
assert_managed_id(batch_input_file.id, "batch_input_file.id")
print("Retrieving batch input file metadata...")
metadata = self.openai_client.files.retrieve(batch_input_file.id)
assert_managed_id(metadata.id, "files.retrieve(input).id")
assert metadata.id == batch_input_file.id, (
f"Input file ID mismatch: retrieve returned '{metadata.id}' but expected '{batch_input_file.id}'"
)
assert metadata.object == "file"
assert metadata.bytes > 0, "bytes not set"
assert metadata.filename == "modified_file.jsonl"
assert metadata.purpose == "batch"
assert metadata.status in ["uploaded", "processed", "error"]
assert metadata.created_at > 0
if wip_features_enabled():
assert metadata.expires_at > 0, "expires_at not set"
self.print_file_metadata(metadata, "Input file")
return batch_input_file
def _create_and_verify_batch(self, input_file_id):
print("\nCreating batch...")
batch = self.create_batch(
self.openai_client,
input_file_id,
MIN_EXPIRY_SECONDS,
)
print(f"Created batch: {self.shorten_id(batch.id)}")
assert batch.id, "No batch ID returned"
assert_managed_id(batch.id, "batch.id")
assert_managed_id(batch.input_file_id, "batch.input_file_id")
assert batch.input_file_id == input_file_id, "batch.input_file_id mismatch"
assert batch.status in ["validating", "in_progress", "finalizing", "completed"]
if not batch.expires_at:
warnings.warn("batch expires_at not set")
else:
assert batch.expires_at > 0
if not batch.endpoint:
warnings.warn("batch.endpoint empty - Azure API quirk, not a bug")
else:
assert batch.endpoint == "/v1/chat/completions"
assert batch.completion_window == "24h"
assert batch.created_at > 0
self.print_batch_metadata(batch)
return batch
def _list_batches(self, batch_id, model_name):
if not wip_features_enabled():
return
print("\nListing batches...")
try:
batches_list = self.wait_for_batch_list(
model_name,
max_seconds=30,
wait_seconds=5,
)
batch_ids = [b.id for b in (batches_list.data if batches_list else [])]
if batch_id not in batch_ids:
warnings.warn(
f"Batch {batch_id} not found in list. "
f"batches.list returns raw IDs, not encoded IDs. raw IDs: {batch_ids}",
)
except openai.APIError as e:
pytest.fail(f"batches.list() failed: {e}")
def _wait_for_batch_completion(self, batch_id, tracker):
print(f"\nWaiting for batch {self.shorten_id(batch_id)} to complete...")
try:
batch_response = self.wait_for_batch_state(
self.openai_client,
batch_id,
"completed",
max_seconds=25 * 60,
wait_seconds=15,
state_tracker=tracker,
)
except RetryError:
tracker.print_state("Timeout waiting for batch completion")
raise TimeoutError("Timed out waiting for batch to be in state: completed")
assert_managed_id(batch_response.id, "batch_response.id")
assert batch_response.id == batch_id, (
f"batch_response.id mismatch: got '{batch_response.id}' but expected '{batch_id}'"
)
assert_managed_id(batch_response.input_file_id, "batch_response.input_file_id")
assert_managed_id(
batch_response.output_file_id,
"batch_response.output_file_id",
)
return batch_response
def _get_and_verify_batch_output(self, output_file_id):
print("\nRetrieving batch output file metadata...")
metadata = self.openai_client.files.retrieve(output_file_id)
assert_managed_id(metadata.id, "files.retrieve(output_file_id).id")
assert metadata.id == output_file_id, (
f"Output file ID mismatch: retrieve returned '{metadata.id}' but expected '{output_file_id}'"
)
assert metadata.object == "file"
assert metadata.bytes > 0, "bytes not set"
assert metadata.filename, "filename not set"
assert metadata.purpose in ["batch_output", "batch"]
assert metadata.created_at > 0
self.print_file_metadata(metadata, "Output file")
print("\nFetching batch output file content...")
content = self.openai_client.files.content(output_file_id)
assert content.text, "No batch file content returned"
assert len(content.text) > 0, "Batch file content is empty"
print(f"Output file content ({len(content.text)} bytes):")
for line in content.text.strip().split("\n")[:3]:
print(f"\t{line}")
return metadata
def _delete_file(self, file_id, label, max_retries=6, retry_delay=10):
print(f"\nDeleting {label}: {self.shorten_id(file_id)}")
for attempt in range(max_retries):
try:
self.openai_client.files.delete(file_id)
return
except openai.BadRequestError as e:
if "batch_processed" in str(e) and attempt < max_retries - 1:
print(
f" File still referenced by unprocessed batch, "
f"retrying in {retry_delay}s ({attempt + 1}/{max_retries})"
)
time.sleep(retry_delay)
else:
pytest.fail(f"files.delete({label}) failed: {e}")
except openai.APIError as e:
pytest.fail(f"files.delete({label}) failed: {e}")
def _verify_file_deleted(self, file_id, label):
print(f"Verifying {label} is deleted...")
try:
self.openai_client.files.content(file_id)
assert False, f"{label} {file_id} still accessible after deletion"
except openai.NotFoundError:
print(f"{label} correctly not accessible after deletion")
# ------------------------------------------------------------------
# Tests
# ------------------------------------------------------------------
@pytest.mark.parametrize(
"model_name",
get_batch_model_names(),
ids=model_id,
)
def test_e2e_managed_batch(self, tmp_path, model_name):
print(
f"\n\nStarting test with base_url={self.base_url} and model_name={model_name}\n",
)
self.reset_mock_server()
tracker = self.create_state_tracker()
batch_input_file = self._create_and_verify_batch_input_file(
tmp_path,
model_name,
)
tracker.set_file_id(batch_input_file.id)
tracker.print_state("After creating batch input file")
batch = self._create_and_verify_batch(batch_input_file.id)
tracker.set_batch_id(batch.id)
tracker.print_state("After creating batch")
self._list_batches(batch.id, model_name)
batch_response = self._wait_for_batch_completion(batch.id, tracker)
tracker.print_state("After batch completed")
self._get_and_verify_batch_output(batch_response.output_file_id)
tracker.print_state("After retrieving output file")
tracker.print_state("Final state after cleanup")
tracker.wait_and_print_s3_callbacks()
tracker.assert_batch_cost_callback()
self._delete_file(batch_input_file.id, "input file")
self._delete_file(batch_response.output_file_id, "output file")
self._verify_file_deleted(batch_input_file.id, "input file")
self._verify_file_deleted(batch_response.output_file_id, "output file")
def cleanup_batches_in_database(self):
import psycopg2
print("Cleaning up stale batch records from database...")
try:
conn = psycopg2.connect(
host="localhost",
port=5432,
database="litellm",
user="llmproxy",
password="dbpassword9090",
)
with conn.cursor() as cur:
cur.execute("""
DELETE FROM "LiteLLM_ManagedObjectTable"
WHERE file_purpose = 'batch' AND status = 'validating'
""")
deleted = cur.rowcount
conn.commit()
if deleted > 0:
print(f"Deleted {deleted} stale batch records")
conn.close()
except Exception as e:
print(f"Warning: Could not clean up database: {e}")
def clear_s3_callbacks(self):
clear_response = httpx.delete(f"{get_mock_server_base_url()}/mock-s3/callbacks")
assert clear_response.status_code == 200, (
f"Failed to clear callbacks: {clear_response.text}"
)
return clear_response.json()
@pytest.mark.skipif(
True,
reason="Skipping managed files test till managed files feature is available",
)
@pytest.mark.parametrize(
"model_name",
get_batch_model_names(),
ids=model_id,
)
def test_error_files(self, tmp_path, model_name):
raise NotImplementedError(
"To implement. Fail a batch and retrieve the error file.",
)

View file

@ -0,0 +1,119 @@
#!/usr/bin/env python
"""
Validation script for Azure Batch E2E test setup.
Run this before running the actual tests to verify all components are accessible.
"""
import os
import sys
from pathlib import Path
sys.path.insert(0, os.path.abspath("../.."))
def check_imports():
"""Verify all required imports work."""
print("Checking imports...")
try:
from base_integration_test import (
get_mock_server_base_url,
get_litellm_base_url,
get_litellm_api_key,
)
print(" ✓ base_integration_test imports OK")
from test_managed_files_base import ManagedFilesBase, get_batch_model_names
print(" ✓ test_managed_files_base imports OK")
from fixtures.mock_azure_batch_server import create_mock_azure_batch_server
print(" ✓ mock_azure_batch_server imports OK")
import httpx
import openai
import psycopg2
import uvicorn
print(" ✓ All external dependencies OK")
return True
except ImportError as e:
print(f" ✗ Import error: {e}")
return False
def check_config_file():
"""Verify config file exists."""
print("\nChecking config file...")
config_path = Path(__file__).parent / "fixtures" / "config.yml"
if config_path.exists():
print(f" ✓ Config file found: {config_path}")
return True
else:
print(f" ✗ Config file not found: {config_path}")
return False
def check_database():
"""Verify database connection."""
print("\nChecking database connection...")
try:
import psycopg2
conn = psycopg2.connect(
host="localhost",
port=5432,
database="litellm",
user="llmproxy",
password="dbpassword9090",
)
conn.close()
print(" ✓ Database connection OK")
return True
except Exception as e:
print(f" ✗ Database connection failed: {e}")
print(" Start PostgreSQL with:")
print(" docker run --name litellm-postgres -e POSTGRES_USER=llmproxy \\")
print(" -e POSTGRES_PASSWORD=dbpassword9090 -e POSTGRES_DB=litellm \\")
print(" -p 5432:5432 -d postgres:15")
return False
def check_ports():
"""Check if required ports are available."""
print("\nChecking ports...")
import socket
for port, name in [(4000, "LiteLLM Proxy"), (8090, "Mock Server")]:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
try:
s.bind(("localhost", port))
print(f" ✓ Port {port} ({name}) is available")
except OSError:
print(f" ⚠ Port {port} ({name}) is in use (will reuse if healthy)")
return True
def main():
print("=" * 70)
print("Azure Batch E2E Test Setup Validation")
print("=" * 70)
checks = [
check_imports(),
check_config_file(),
check_database(),
check_ports(),
]
print("\n" + "=" * 70)
if all(checks):
print("✓ All checks passed! Ready to run E2E tests.")
print("\nRun tests with:")
print(" cd litellm")
print(" export DATABASE_URL='postgresql://llmproxy:dbpassword9090@localhost:5432/litellm'")
print(" poetry run pytest tests/proxy_e2e_azure_batches_tests/test_proxy_e2e_azure_batches.py -vv")
return 0
else:
print("✗ Some checks failed. Please fix the issues above.")
return 1
if __name__ == "__main__":
sys.exit(main())

View file

@ -0,0 +1,270 @@
"""
Tests for SearchAPI.io (Google Search) integration.
Tests the SearchAPI.io search provider implementation including:
- Request transformation
- Response transformation
- Parameter mapping
- Error handling
"""
import json
import os
import sys
from unittest.mock import MagicMock, Mock, patch
import httpx
import pytest
sys.path.insert(
0, os.path.abspath("../..")
)
from litellm.llms.searchapi.search.transformation import SearchAPIConfig
from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
class TestSearchAPIConfig:
"""Test SearchAPI.io configuration and transformations."""
def test_ui_friendly_name(self):
"""Test that UI friendly name is returned correctly."""
config = SearchAPIConfig()
assert config.ui_friendly_name() == "SearchAPI.io (Google Search)"
def test_get_http_method(self):
"""Test that HTTP method is GET."""
config = SearchAPIConfig()
assert config.get_http_method() == "GET"
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
def test_validate_environment_with_api_key(self, mock_get_secret):
"""Test environment validation with API key."""
mock_get_secret.return_value = "test_api_key"
config = SearchAPIConfig()
headers = {}
result = config.validate_environment(headers, api_key="test_api_key")
assert result["Content-Type"] == "application/json"
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
def test_validate_environment_without_api_key(self, mock_get_secret):
"""Test environment validation without API key raises error."""
mock_get_secret.return_value = None
config = SearchAPIConfig()
headers = {}
with pytest.raises(ValueError, match="SEARCHAPI_API_KEY is not set"):
config.validate_environment(headers)
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
def test_transform_search_request_basic(self, mock_get_secret):
"""Test basic search request transformation."""
mock_get_secret.return_value = "test_api_key"
config = SearchAPIConfig()
result = config.transform_search_request(
query="test query",
optional_params={},
api_key="test_api_key"
)
assert "_searchapi_params" in result
params = result["_searchapi_params"]
assert params["engine"] == "google"
assert params["q"] == "test query"
assert params["api_key"] == "test_api_key"
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
def test_transform_search_request_with_max_results(self, mock_get_secret):
"""Test search request transformation with max_results parameter."""
mock_get_secret.return_value = "test_api_key"
config = SearchAPIConfig()
result = config.transform_search_request(
query="test query",
optional_params={"max_results": 5},
api_key="test_api_key"
)
params = result["_searchapi_params"]
assert params["num"] == 5
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
def test_transform_search_request_with_country(self, mock_get_secret):
"""Test search request transformation with country parameter."""
mock_get_secret.return_value = "test_api_key"
config = SearchAPIConfig()
result = config.transform_search_request(
query="test query",
optional_params={"country": "US"},
api_key="test_api_key"
)
params = result["_searchapi_params"]
assert params["gl"] == "us"
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
def test_transform_search_request_with_domain_filter(self, mock_get_secret):
"""Test search request transformation with domain filter."""
mock_get_secret.return_value = "test_api_key"
config = SearchAPIConfig()
result = config.transform_search_request(
query="test query",
optional_params={"search_domain_filter": ["example.com", "test.com"]},
api_key="test_api_key"
)
params = result["_searchapi_params"]
assert "site:example.com" in params["q"]
assert "site:test.com" in params["q"]
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
def test_transform_search_request_with_list_query(self, mock_get_secret):
"""Test search request transformation with list query."""
mock_get_secret.return_value = "test_api_key"
config = SearchAPIConfig()
result = config.transform_search_request(
query=["test", "query"],
optional_params={},
api_key="test_api_key"
)
params = result["_searchapi_params"]
assert params["q"] == "test query"
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
def test_get_complete_url(self, mock_get_secret):
"""Test URL construction with query parameters."""
mock_get_secret.return_value = None
config = SearchAPIConfig()
data = {
"_searchapi_params": {
"engine": "google",
"q": "test query",
"api_key": "test_key"
}
}
url = config.get_complete_url(
api_base=None,
optional_params={},
data=data
)
assert "https://www.searchapi.io/api/v1/search?" in url
assert "engine=google" in url
assert "q=test+query" in url
assert "api_key=test_key" in url
def test_transform_search_response(self):
"""Test search response transformation."""
config = SearchAPIConfig()
# Mock response
mock_response = Mock(spec=httpx.Response)
mock_response.json.return_value = {
"organic_results": [
{
"title": "Test Result 1",
"link": "https://example.com/1",
"snippet": "This is a test snippet 1",
"date": "2024-01-01"
},
{
"title": "Test Result 2",
"link": "https://example.com/2",
"snippet": "This is a test snippet 2"
}
]
}
result = config.transform_search_response(
raw_response=mock_response,
logging_obj=None
)
assert isinstance(result, SearchResponse)
assert result.object == "search"
assert len(result.results) == 2
# Check first result
assert result.results[0].title == "Test Result 1"
assert result.results[0].url == "https://example.com/1"
assert result.results[0].snippet == "This is a test snippet 1"
assert result.results[0].date == "2024-01-01"
assert result.results[0].last_updated is None
# Check second result
assert result.results[1].title == "Test Result 2"
assert result.results[1].url == "https://example.com/2"
assert result.results[1].snippet == "This is a test snippet 2"
assert result.results[1].date is None
def test_transform_search_response_empty(self):
"""Test search response transformation with no results."""
config = SearchAPIConfig()
mock_response = Mock(spec=httpx.Response)
mock_response.json.return_value = {
"organic_results": []
}
result = config.transform_search_response(
raw_response=mock_response,
logging_obj=None
)
assert isinstance(result, SearchResponse)
assert len(result.results) == 0
def test_append_domain_filters(self):
"""Test domain filter appending logic."""
config = SearchAPIConfig()
query = "test query"
domains = ["example.com", "test.com"]
result = config._append_domain_filters(query, domains)
assert "(test query)" in result
assert "site:example.com" in result
assert "site:test.com" in result
assert "OR" in result
assert "AND" in result
@pytest.mark.skipif(
os.environ.get("SEARCHAPI_API_KEY") is None,
reason="SEARCHAPI_API_KEY not set in environment"
)
class TestSearchAPIIntegration:
"""Integration tests for SearchAPI.io (requires API key)."""
def test_real_search_request(self):
"""
Test a real search request to SearchAPI.io.
This test is skipped if SEARCHAPI_API_KEY is not set.
"""
import litellm
response = litellm.search(
query="Python programming",
search_provider="searchapi",
max_results=5
)
assert response is not None
assert hasattr(response, "results")
assert len(response.results) > 0
assert all(hasattr(r, "title") for r in response.results)
assert all(hasattr(r, "url") for r in response.results)
assert all(hasattr(r, "snippet") for r in response.results)
if __name__ == "__main__":
pytest.main([__file__, "-v"])

View file

@ -425,6 +425,7 @@ def test_select_azure_base_url_called(setup_mocks):
"add_message",
"arun_thread_stream",
"aresponses",
"aresponses_websocket",
"alist_input_items",
"acreate_fine_tuning_job",
"acancel_fine_tuning_job",

View file

@ -686,6 +686,48 @@ class TestTransformListInputItemsRequest:
# Assert
assert "include" not in params # Empty list should not be included
def test_openai_transform_compact_response_api_request_query_params_preserved(self):
"""Test compact URL construction preserves query params and appends path."""
# Setup
azure_style_api_base = (
"https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview"
)
# Execute
url, data = self.openai_config.transform_compact_response_api_request(
model="gpt-5.2-codex",
input="hello",
response_api_optional_request_params={},
api_base=azure_style_api_base,
litellm_params=self.litellm_params,
headers=self.headers,
)
# Assert
assert (
url
== "https://test.openai.azure.com/openai/responses/compact?api-version=2024-05-01-preview"
)
assert data["model"] == "gpt-5.2-codex"
assert data["input"] == "hello"
def test_openai_transform_compact_response_api_request_path_without_query(self):
"""Test compact URL construction for base URL without query params."""
# Execute
url, data = self.openai_config.transform_compact_response_api_request(
model="gpt-4o",
input="hello",
response_api_optional_request_params={},
api_base="https://api.openai.com/v1/responses",
litellm_params=self.litellm_params,
headers=self.headers,
)
# Assert
assert url == "https://api.openai.com/v1/responses/compact"
assert data["model"] == "gpt-4o"
assert data["input"] == "hello"
def test_azure_transform_list_input_items_request_minimal(self):
"""Test Azure implementation with minimal parameters"""
# Setup
@ -1239,4 +1281,4 @@ class TestPhaseParameter:
assert validated[0]["phase"] == "commentary"
assert validated[1]["phase"] == "final_answer"
assert "phase" not in validated[2]
assert "phase" not in validated[2]

View file

@ -269,6 +269,7 @@ def test_gpt5_1_model_detection(gpt5_config: OpenAIGPT5Config):
assert gpt5_config.is_model_gpt_5_1_model("gpt-5.2-2025-12-11")
assert not gpt5_config.is_model_gpt_5_1_model("gpt-5.2-chat")
assert not gpt5_config.is_model_gpt_5_1_model("gpt-5.2-chat-latest")
assert not gpt5_config.is_model_gpt_5_1_model("gpt-5.3-chat-latest")
assert not gpt5_config.is_model_gpt_5_1_model("gpt-5.2-pro")
assert not gpt5_config.is_model_gpt_5_1_model("gpt-5")
assert not gpt5_config.is_model_gpt_5_1_model("gpt-5-mini")
@ -402,7 +403,7 @@ def test_gpt5_2_chat_temperature_restricted(config: OpenAIConfig):
Regression test for https://github.com/BerriAI/litellm/issues/21911
"""
# gpt-5.2-chat should reject non-1 temperature when drop_params=False
for model in ["gpt-5.2-chat", "gpt-5.2-chat-latest"]:
for model in ["gpt-5.2-chat", "gpt-5.2-chat-latest", "gpt-5.3-chat-latest"]:
with pytest.raises(litellm.utils.UnsupportedParamsError):
config.map_openai_params(
non_default_params={"temperature": 0.7},

View file

@ -0,0 +1,540 @@
import base64
import json
import os
import sys
from io import BytesIO
from unittest.mock import MagicMock, patch
import httpx
import pytest
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
from litellm.llms.openrouter.common_utils import OpenRouterException
from litellm.llms.openrouter.image_edit.transformation import (
OpenRouterImageEditConfig,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ImageResponse
class TestOpenRouterImageEditTransformation:
def setup_method(self):
"""Set up test fixtures before each test method."""
self.config = OpenRouterImageEditConfig()
self.model = "google/gemini-2.5-flash-image"
self.logging_obj = MagicMock()
self.sample_image_bytes = b"\x89PNG\r\n\x1a\n" + b"\x00" * 100
def test_get_supported_openai_params(self):
"""Test that get_supported_openai_params returns correct parameters."""
supported_params = self.config.get_supported_openai_params(self.model)
assert "size" in supported_params
assert "quality" in supported_params
assert "n" in supported_params
assert len(supported_params) == 3
def test_use_multipart_form_data_returns_false(self):
"""Test that OpenRouter uses JSON, not multipart/form-data."""
assert self.config.use_multipart_form_data() is False
# Parameter mapping tests
def test_map_openai_params_size(self):
"""Test that size is mapped to image_config.aspect_ratio."""
result = self.config.map_openai_params(
image_edit_optional_params={"size": "1024x1024"},
model=self.model,
drop_params=False,
)
assert "image_config" in result
assert result["image_config"]["aspect_ratio"] == "1:1"
def test_map_openai_params_quality(self):
"""Test that quality is mapped to image_config.image_size."""
result = self.config.map_openai_params(
image_edit_optional_params={"quality": "high"},
model=self.model,
drop_params=False,
)
assert "image_config" in result
assert result["image_config"]["image_size"] == "4K"
def test_map_openai_params_size_and_quality(self):
"""Test that both size and quality are mapped correctly."""
result = self.config.map_openai_params(
image_edit_optional_params={"size": "1792x1024", "quality": "hd"},
model=self.model,
drop_params=False,
)
assert result["image_config"]["aspect_ratio"] == "16:9"
assert result["image_config"]["image_size"] == "4K"
def test_map_openai_params_n_passthrough(self):
"""Test that n parameter is passed through directly."""
result = self.config.map_openai_params(
image_edit_optional_params={"n": 2},
model=self.model,
drop_params=False,
)
assert result["n"] == 2
def test_map_openai_params_unknown_quality_ignored(self):
"""Test that unknown quality values produce no image_size mapping."""
result = self.config.map_openai_params(
image_edit_optional_params={"quality": "unknown_value"},
model=self.model,
drop_params=False,
)
assert "image_config" not in result
# Size-to-aspect-ratio mapping tests
def test_map_size_to_aspect_ratio_square(self):
"""Test mapping square sizes to 1:1 aspect ratio."""
assert self.config._map_size_to_aspect_ratio("256x256") == "1:1"
assert self.config._map_size_to_aspect_ratio("512x512") == "1:1"
assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1"
def test_map_size_to_aspect_ratio_landscape(self):
"""Test mapping landscape sizes to correct aspect ratios."""
assert self.config._map_size_to_aspect_ratio("1536x1024") == "3:2"
assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9"
def test_map_size_to_aspect_ratio_portrait(self):
"""Test mapping portrait sizes to correct aspect ratios."""
assert self.config._map_size_to_aspect_ratio("1024x1536") == "2:3"
assert self.config._map_size_to_aspect_ratio("1024x1792") == "9:16"
def test_map_size_to_aspect_ratio_unknown_defaults_to_1_1(self):
"""Test that unknown size defaults to 1:1."""
assert self.config._map_size_to_aspect_ratio("999x999") == "1:1"
# Quality-to-image-size mapping tests
def test_map_quality_to_image_size(self):
"""Test quality to image size mappings."""
assert self.config._map_quality_to_image_size("low") == "1K"
assert self.config._map_quality_to_image_size("standard") == "1K"
assert self.config._map_quality_to_image_size("auto") == "1K"
assert self.config._map_quality_to_image_size("medium") == "2K"
assert self.config._map_quality_to_image_size("high") == "4K"
assert self.config._map_quality_to_image_size("hd") == "4K"
def test_map_quality_to_image_size_unknown_returns_none(self):
"""Test that unknown quality returns None."""
assert self.config._map_quality_to_image_size("unknown") is None
# URL tests
def test_get_complete_url_default(self):
"""Test that default URL is OpenRouter chat completions endpoint."""
result = self.config.get_complete_url(
model=self.model,
api_base=None,
litellm_params={},
)
assert result == "https://openrouter.ai/api/v1/chat/completions"
def test_get_complete_url_with_custom_base(self):
"""Test that custom api_base gets /chat/completions appended."""
result = self.config.get_complete_url(
model=self.model,
api_base="https://custom.openrouter.ai/api/v1",
litellm_params={},
)
assert result == "https://custom.openrouter.ai/api/v1/chat/completions"
def test_get_complete_url_with_complete_base(self):
"""Test that api_base already ending in /chat/completions is not duplicated."""
url = "https://custom.openrouter.ai/api/v1/chat/completions"
result = self.config.get_complete_url(
model=self.model,
api_base=url,
litellm_params={},
)
assert result == url
# Validate environment tests
@patch("litellm.llms.openrouter.image_edit.transformation.get_secret_str")
def test_validate_environment_with_api_key(self, mock_get_secret):
"""Test that validate_environment sets authorization header with provided key."""
headers = {}
result = self.config.validate_environment(
headers=headers,
model=self.model,
api_key="test_api_key",
)
assert result["Authorization"] == "Bearer test_api_key"
mock_get_secret.assert_not_called()
@patch("litellm.llms.openrouter.image_edit.transformation.get_secret_str")
def test_validate_environment_with_secret_key(self, mock_get_secret):
"""Test that validate_environment falls back to secret key."""
mock_get_secret.return_value = "secret_api_key"
headers = {}
result = self.config.validate_environment(
headers=headers,
model=self.model,
api_key=None,
)
assert result["Authorization"] == "Bearer secret_api_key"
@patch("litellm.llms.openrouter.image_edit.transformation.litellm")
@patch("litellm.llms.openrouter.image_edit.transformation.get_secret_str")
def test_validate_environment_missing_api_key_raises(self, mock_get_secret, mock_litellm):
"""Test that validate_environment raises ValueError when no API key is available."""
mock_get_secret.return_value = None
mock_litellm.api_key = None
with pytest.raises(ValueError, match="OPENROUTER_API_KEY is not set"):
self.config.validate_environment(
headers={},
model=self.model,
api_key=None,
)
# Request transformation tests
def test_transform_image_edit_request_basic(self):
"""Test basic request transformation with image and prompt."""
data, files = self.config.transform_image_edit_request(
model=self.model,
prompt="Add a sunset to this image",
image=self.sample_image_bytes,
image_edit_optional_request_params={},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert data["model"] == self.model
assert data["modalities"] == ["image", "text"]
assert len(data["messages"]) == 1
assert data["messages"][0]["role"] == "user"
content = data["messages"][0]["content"]
assert len(content) == 2
# First content part should be the image
assert content[0]["type"] == "image_url"
assert content[0]["image_url"]["url"].startswith("data:image/png;base64,")
# Second content part should be the text prompt
assert content[1]["type"] == "text"
assert content[1]["text"] == "Add a sunset to this image"
# Files should be empty (JSON mode)
assert list(files) == []
def test_transform_image_edit_request_with_bytesio(self):
"""Test request transformation with BytesIO image input."""
image = BytesIO(self.sample_image_bytes)
data, files = self.config.transform_image_edit_request(
model=self.model,
prompt="Edit this",
image=image,
image_edit_optional_request_params={},
litellm_params=GenericLiteLLMParams(),
headers={},
)
content = data["messages"][0]["content"]
assert content[0]["type"] == "image_url"
assert content[0]["image_url"]["url"].startswith("data:image/png;base64,")
def test_transform_image_edit_request_with_multiple_images(self):
"""Test request transformation with a list of images."""
images = [self.sample_image_bytes, self.sample_image_bytes]
data, files = self.config.transform_image_edit_request(
model=self.model,
prompt="Combine these images",
image=images,
image_edit_optional_request_params={},
litellm_params=GenericLiteLLMParams(),
headers={},
)
content = data["messages"][0]["content"]
# Two image parts + one text part
assert len(content) == 3
assert content[0]["type"] == "image_url"
assert content[1]["type"] == "image_url"
assert content[2]["type"] == "text"
def test_transform_image_edit_request_with_optional_params(self):
"""Test that optional params are included in request body."""
data, files = self.config.transform_image_edit_request(
model=self.model,
prompt="Edit this",
image=self.sample_image_bytes,
image_edit_optional_request_params={
"image_config": {"aspect_ratio": "16:9"},
"n": 2,
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert data["image_config"]["aspect_ratio"] == "16:9"
assert data["n"] == 2
def test_transform_image_edit_request_base64_encoding(self):
"""Test that image bytes are correctly base64-encoded in the request."""
raw_bytes = b"test_image_data"
expected_b64 = base64.b64encode(raw_bytes).decode("utf-8")
data, _ = self.config.transform_image_edit_request(
model=self.model,
prompt="Edit",
image=raw_bytes,
image_edit_optional_request_params={},
litellm_params=GenericLiteLLMParams(),
headers={},
)
image_url = data["messages"][0]["content"][0]["image_url"]["url"]
# Extract the base64 part after the data URL prefix
b64_part = image_url.split(",", 1)[1]
assert b64_part == expected_b64
def test_transform_image_edit_request_no_prompt(self):
"""Test request transformation with no prompt (image-only)."""
data, _ = self.config.transform_image_edit_request(
model=self.model,
prompt=None,
image=self.sample_image_bytes,
image_edit_optional_request_params={},
litellm_params=GenericLiteLLMParams(),
headers={},
)
content = data["messages"][0]["content"]
# Only image, no text part
assert len(content) == 1
assert content[0]["type"] == "image_url"
# Response transformation tests
def test_transform_image_edit_response_with_base64(self):
"""Test response transformation with base64 image data."""
response_data = {
"choices": [{
"message": {
"content": "Here is the edited image.",
"role": "assistant",
"images": [{
"image_url": {"url": "data:image/png;base64,iVBORw0KGgoAAAANS"},
"type": "image_url"
}]
}
}],
"usage": {
"prompt_tokens": 300,
"completion_tokens": 1299,
"total_tokens": 1599,
"completion_tokens_details": {"image_tokens": 1290},
"cost": 0.05
},
"model": self.model
}
mock_response = MagicMock()
mock_response.json.return_value = response_data
mock_response.status_code = 200
mock_response.headers = {}
result = self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
assert len(result.data) == 1
assert result.data[0].b64_json == "iVBORw0KGgoAAAANS"
assert result.data[0].url is None
def test_transform_image_edit_response_with_url(self):
"""Test response transformation with URL image data."""
response_data = {
"choices": [{
"message": {
"content": "Edited.",
"role": "assistant",
"images": [{
"image_url": {"url": "https://example.com/edited.png"},
"type": "image_url"
}]
}
}],
"usage": {"prompt_tokens": 10, "total_tokens": 1310},
"model": self.model
}
mock_response = MagicMock()
mock_response.json.return_value = response_data
mock_response.status_code = 200
mock_response.headers = {}
result = self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
assert len(result.data) == 1
assert result.data[0].url == "https://example.com/edited.png"
assert result.data[0].b64_json is None
def test_transform_image_edit_response_usage_and_cost(self):
"""Test that usage and cost are correctly extracted from response."""
response_data = {
"choices": [{
"message": {
"content": "Edited.",
"role": "assistant",
"images": [{
"image_url": {"url": "data:image/png;base64,abc123"},
"type": "image_url"
}]
}
}],
"usage": {
"prompt_tokens": 300,
"completion_tokens": 1299,
"total_tokens": 1599,
"completion_tokens_details": {"image_tokens": 1290},
"prompt_tokens_details": {"image_tokens": 258},
"cost": 0.05,
"cost_details": {"input_cost": 0.01, "output_cost": 0.04}
},
"model": self.model
}
mock_response = MagicMock()
mock_response.json.return_value = response_data
mock_response.status_code = 200
mock_response.headers = {}
result = self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
# Check usage
assert result.usage is not None
assert result.usage.input_tokens == 300
assert result.usage.output_tokens == 1290
assert result.usage.total_tokens == 1599
assert result.usage.input_tokens_details.image_tokens == 258
assert result.usage.input_tokens_details.text_tokens == 42
# Check cost
assert result._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == 0.05
# Check cost details
assert result._hidden_params["response_cost_details"]["input_cost"] == 0.01
assert result._hidden_params["response_cost_details"]["output_cost"] == 0.04
# Check model
assert result._hidden_params["model"] == self.model
def test_transform_image_edit_response_multiple_images(self):
"""Test response transformation with multiple output images."""
response_data = {
"choices": [{
"message": {
"content": "Here are your edits.",
"role": "assistant",
"images": [
{
"image_url": {"url": "data:image/png;base64,img1data"},
"type": "image_url"
},
{
"image_url": {"url": "data:image/png;base64,img2data"},
"type": "image_url"
}
]
}
}],
"usage": {"prompt_tokens": 300, "total_tokens": 2600},
"model": self.model
}
mock_response = MagicMock()
mock_response.json.return_value = response_data
mock_response.status_code = 200
mock_response.headers = {}
result = self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
assert len(result.data) == 2
assert result.data[0].b64_json == "img1data"
assert result.data[1].b64_json == "img2data"
def test_transform_image_edit_response_json_error(self):
"""Test that invalid JSON response raises OpenRouterException."""
mock_response = MagicMock()
mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "", 0)
mock_response.status_code = 500
mock_response.headers = {}
with pytest.raises(OpenRouterException) as exc_info:
self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
assert "Error parsing OpenRouter response" in str(exc_info.value)
assert exc_info.value.status_code == 500
def test_get_error_class(self):
"""Test that get_error_class returns OpenRouterException."""
error = self.config.get_error_class(
error_message="Test error",
status_code=400,
headers={"Content-Type": "application/json"},
)
assert isinstance(error, OpenRouterException)
assert error.status_code == 400
# Read image bytes tests
def test_read_image_bytes_from_bytes(self):
"""Test reading bytes directly."""
result = self.config._read_image_bytes(b"raw_bytes")
assert result == b"raw_bytes"
def test_read_image_bytes_from_bytesio(self):
"""Test reading bytes from BytesIO."""
bio = BytesIO(b"bytesio_data")
bio.seek(5) # Move position to test seek reset
result = self.config._read_image_bytes(bio)
assert result == b"bytesio_data"
assert bio.tell() == 5 # Position should be restored
def test_read_image_bytes_unsupported_type(self):
"""Test that unsupported image type raises ValueError."""
with pytest.raises(ValueError, match="Unsupported image type"):
self.config._read_image_bytes("not_an_image") # type: ignore

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