mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
merge: resolve conflict with main in tool_management_endpoints.py
Remove unused ToolOutputPolicy import. Made-with: Cursor
This commit is contained in:
commit
b4cd5879fb
129 changed files with 12806 additions and 905 deletions
|
|
@ -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:
|
||||
|
|
|
|||
1
.github/workflows/test-linting.yml
vendored
1
.github/workflows/test-linting.yml
vendored
|
|
@ -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: |
|
||||
|
|
|
|||
2
.github/workflows/test-litellm.yml
vendored
2
.github/workflows/test-litellm.yml
vendored
|
|
@ -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: |
|
||||
|
|
|
|||
90
.github/workflows/test-proxy-e2e-azure-batches.yml
vendored
Normal file
90
.github/workflows/test-proxy-e2e-azure-batches.yml
vendored
Normal 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
|
||||
|
||||
|
|
@ -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>
|
||||
|
||||
|
|
|
|||
|
|
@ -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)` |
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
```
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
197
docs/my-website/docs/search/searchapi.md
Normal file
197
docs/my-website/docs/search/searchapi.md
Normal 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
|
||||
|
|
@ -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");
|
||||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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 ##########
|
||||
#########################################################
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
11
litellm/llms/openrouter/image_edit/__init__.py
Normal file
11
litellm/llms/openrouter/image_edit/__init__.py
Normal 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()
|
||||
367
litellm/llms/openrouter/image_edit/transformation.py
Normal file
367
litellm/llms/openrouter/image_edit/transformation.py
Normal 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.")
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
1
litellm/llms/searchapi/__init__.py
Normal file
1
litellm/llms/searchapi/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
"""SearchAPI.io integration for LiteLLM."""
|
||||
4
litellm/llms/searchapi/search/__init__.py
Normal file
4
litellm/llms/searchapi/search/__init__.py
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
"""SearchAPI.io search integration for LiteLLM."""
|
||||
from litellm.llms.searchapi.search.transformation import SearchAPIConfig
|
||||
|
||||
__all__ = ["SearchAPIConfig"]
|
||||
232
litellm/llms/searchapi/search/transformation.py
Normal file
232
litellm/llms/searchapi/search/transformation.py
Normal 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",
|
||||
)
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
60
litellm/proxy/dd_span_tagger.py
Normal file
60
litellm/proxy/dd_span_tagger.py
Normal 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,
|
||||
)
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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"]]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -26,7 +26,6 @@ from litellm.types.tool_management import (
|
|||
ToolDetailResponse,
|
||||
ToolInputPolicy,
|
||||
ToolListResponse,
|
||||
ToolOutputPolicy,
|
||||
ToolPolicyOption,
|
||||
ToolPolicyOptionsResponse,
|
||||
ToolPolicyUpdateRequest,
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
80
poetry.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ SEARCH_PROVIDERS = [
|
|||
"searxng",
|
||||
"linkup",
|
||||
"duckduckgo",
|
||||
"searchapi",
|
||||
]
|
||||
|
||||
ALLOWED_FILES_IN_LLMS_FOLDER = [
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
)
|
||||
|
|
|
|||
136
tests/litellm/litellm_core_utils/test_json_schema_validation.py
Normal file
136
tests/litellm/litellm_core_utils/test_json_schema_validation.py
Normal 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
|
||||
362
tests/litellm/proxy/test_batch_x_litellm_model_encoding.py
Normal file
362
tests/litellm/proxy/test_batch_x_litellm_model_encoding.py
Normal 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
|
||||
|
|
@ -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=[
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
0
tests/proxy_e2e_azure_batches_tests/__init__.py
Normal file
0
tests/proxy_e2e_azure_batches_tests/__init__.py
Normal file
494
tests/proxy_e2e_azure_batches_tests/base_integration_test.py
Normal file
494
tests/proxy_e2e_azure_batches_tests/base_integration_test.py
Normal 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
|
||||
311
tests/proxy_e2e_azure_batches_tests/conftest.py
Normal file
311
tests/proxy_e2e_azure_batches_tests/conftest.py
Normal 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()
|
||||
0
tests/proxy_e2e_azure_batches_tests/fixtures/__init__.py
Normal file
0
tests/proxy_e2e_azure_batches_tests/fixtures/__init__.py
Normal file
56
tests/proxy_e2e_azure_batches_tests/fixtures/config.yml
Normal file
56
tests/proxy_e2e_azure_batches_tests/fixtures/config.yml
Normal 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
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
from .server import create_mock_azure_batch_server
|
||||
|
||||
__all__ = ["create_mock_azure_batch_server"]
|
||||
|
|
@ -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()},
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
@ -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},
|
||||
}
|
||||
|
|
@ -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,
|
||||
}
|
||||
|
|
@ -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}
|
||||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
41
tests/proxy_e2e_azure_batches_tests/test_fixtures_smoke.py
Normal file
41
tests/proxy_e2e_azure_batches_tests/test_fixtures_smoke.py
Normal 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")
|
||||
1085
tests/proxy_e2e_azure_batches_tests/test_managed_files_base.py
Normal file
1085
tests/proxy_e2e_azure_batches_tests/test_managed_files_base.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -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.",
|
||||
)
|
||||
119
tests/proxy_e2e_azure_batches_tests/validate_e2e_setup.py
Normal file
119
tests/proxy_e2e_azure_batches_tests/validate_e2e_setup.py
Normal 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())
|
||||
270
tests/search_tests/test_searchapi_search.py
Normal file
270
tests/search_tests/test_searchapi_search.py
Normal 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"])
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
|
|
|
|||
|
|
@ -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
Loading…
Add table
Reference in a new issue