Merge pull request #18809 from BerriAI/main

merge main
This commit is contained in:
Sameer Kankute 2026-01-08 15:28:13 +05:30 committed by GitHub
commit c05e148488
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
113 changed files with 4688 additions and 639 deletions

View file

@ -1465,7 +1465,7 @@ jobs:
- run:
name: Run core tests
command: |
python -m pytest tests/test_litellm --ignore=tests/test_litellm/proxy --ignore=tests/test_litellm/llms --cov=litellm --cov-report=xml --junitxml=test-results/junit-core.xml --durations=10 -n 16 --maxfail=5 --timeout=300 -vv --log-cli-level=WARNING
python -m pytest tests/test_litellm --ignore=tests/test_litellm/proxy --ignore=tests/test_litellm/llms --ignore=tests/test_litellm/integrations --ignore=tests/test_litellm/litellm_core_utils --cov=litellm --cov-report=xml --junitxml=test-results/junit-core.xml --durations=10 -n 16 --maxfail=5 --timeout=300 -vv --log-cli-level=WARNING
no_output_timeout: 120m
- run:
name: Rename the coverage files
@ -1479,6 +1479,60 @@ jobs:
paths:
- litellm_core_tests_coverage.xml
- litellm_core_tests_coverage
litellm_mapped_tests_litellm_core_utils:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
resource_class: xlarge
steps:
- setup_litellm_test_deps
- run:
name: Run litellm_core_utils tests
command: |
python -m pytest tests/test_litellm/litellm_core_utils --cov=litellm --cov-report=xml --junitxml=test-results/junit-litellm-core-utils.xml --durations=10 -n 16 --maxfail=5 --timeout=300 -vv --log-cli-level=WARNING
no_output_timeout: 120m
- run:
name: Rename the coverage files
command: |
mv coverage.xml litellm_core_utils_tests_coverage.xml
mv .coverage litellm_core_utils_tests_coverage
- store_test_results:
path: test-results
- persist_to_workspace:
root: .
paths:
- litellm_core_utils_tests_coverage.xml
- litellm_core_utils_tests_coverage
litellm_mapped_tests_integrations:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
resource_class: xlarge
steps:
- setup_litellm_test_deps
- run:
name: Run integrations tests
command: |
python -m pytest tests/test_litellm/integrations --cov=litellm --cov-report=xml --junitxml=test-results/junit-integrations.xml --durations=10 -n 16 --maxfail=5 --timeout=300 -vv --log-cli-level=WARNING
no_output_timeout: 120m
- run:
name: Rename the coverage files
command: |
mv coverage.xml litellm_integrations_tests_coverage.xml
mv .coverage litellm_integrations_tests_coverage
- store_test_results:
path: test-results
- persist_to_workspace:
root: .
paths:
- litellm_integrations_tests_coverage.xml
- litellm_integrations_tests_coverage
litellm_mapped_enterprise_tests:
docker:
- image: cimg/python:3.11
@ -1960,6 +2014,7 @@ jobs:
- run: ruff check ./litellm
# - run: python ./tests/documentation_tests/test_general_setting_keys.py
- run: python ./tests/code_coverage_tests/check_licenses.py
- run: python ./tests/code_coverage_tests/check_provider_folders_documented.py
- run: python ./tests/code_coverage_tests/router_code_coverage.py
- run: python ./tests/code_coverage_tests/test_chat_completion_imports.py
- run: python ./tests/code_coverage_tests/info_log_check.py
@ -3871,6 +3926,18 @@ workflows:
only:
- main
- /litellm_.*/
- litellm_mapped_tests_integrations:
filters:
branches:
only:
- main
- /litellm_.*/
- litellm_mapped_tests_litellm_core_utils:
filters:
branches:
only:
- main
- /litellm_.*/
- batches_testing:
filters:
branches:
@ -3919,6 +3986,8 @@ workflows:
- litellm_mapped_tests_proxy
- litellm_mapped_tests_llms
- litellm_mapped_tests_core
- litellm_mapped_tests_integrations
- litellm_mapped_tests_litellm_core_utils
- litellm_mapped_enterprise_tests
- batches_testing
- litellm_utils_testing
@ -3990,6 +4059,8 @@ workflows:
- litellm_mapped_tests_proxy
- litellm_mapped_tests_llms
- litellm_mapped_tests_core
- litellm_mapped_tests_integrations
- litellm_mapped_tests_litellm_core_utils
- litellm_mapped_enterprise_tests
- batches_testing
- litellm_utils_testing

View file

@ -27,6 +27,7 @@ body:
attributes:
label: What part of LiteLLM is this about?
options:
- ''
- "SDK (litellm Python package)"
- "Proxy"
- "UI Dashboard"

View file

@ -27,6 +27,7 @@ body:
attributes:
label: What part of LiteLLM is this about?
options:
- ''
- "SDK (litellm Python package)"
- "Proxy"
- "UI Dashboard"

View file

@ -12,7 +12,7 @@ jobs:
issues: write
steps:
- name: Add SDK label
if: contains(github.event.issue.body, 'SDK (litellm Python package)')
if: contains(github.event.issue.body, 'What part of LiteLLM is this about?\n\nSDK (litellm Python package)')
uses: actions/github-script@v7
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
@ -45,7 +45,7 @@ jobs:
});
- name: Add Proxy label
if: contains(github.event.issue.body, 'Proxy')
if: contains(github.event.issue.body, 'What part of LiteLLM is this about?\n\nProxy')
uses: actions/github-script@v7
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
@ -78,7 +78,7 @@ jobs:
});
- name: Add UI Dashboard label
if: contains(github.event.issue.body, 'UI Dashboard')
if: contains(github.event.issue.body, 'What part of LiteLLM is this about?\n\nUI Dashboard')
uses: actions/github-script@v7
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
@ -111,7 +111,7 @@ jobs:
});
- name: Add Docs label
if: contains(github.event.issue.body, 'Docs')
if: contains(github.event.issue.body, 'What part of LiteLLM is this about?\n\nDocs')
uses: actions/github-script@v7
with:
github-token: ${{ secrets.GITHUB_TOKEN }}

View file

@ -20,7 +20,8 @@ RUN python -m pip install build
COPY . .
# Build Admin UI
RUN chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
# Convert Windows line endings to Unix and make executable
RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
# Build the package
RUN rm -rf dist/* && python -m build
@ -65,12 +66,14 @@ RUN find /usr/lib -type f -path "*/tornado/test/*" -delete && \
find /usr/lib -type d -path "*/tornado/test" -delete
# Install semantic_router and aurelio-sdk using script
RUN chmod +x docker/install_auto_router.sh && ./docker/install_auto_router.sh
# Convert Windows line endings to Unix and make executable
RUN sed -i 's/\r$//' docker/install_auto_router.sh && chmod +x docker/install_auto_router.sh && ./docker/install_auto_router.sh
# Generate prisma client
RUN prisma generate
RUN chmod +x docker/entrypoint.sh
RUN chmod +x docker/prod_entrypoint.sh
# Convert Windows line endings to Unix for entrypoint scripts
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh
RUN sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
EXPOSE 4000/tcp

View file

@ -34,47 +34,47 @@ install_ggshield() {
echo "ggshield installed successfully"
}
# Function to run secret detection scans
run_secret_detection() {
echo "Running secret detection scans..."
# # Function to run secret detection scans
# run_secret_detection() {
# echo "Running secret detection scans..."
if ! command -v ggshield &> /dev/null; then
install_ggshield
fi
# if ! command -v ggshield &> /dev/null; then
# install_ggshield
# fi
# Check if GITGUARDIAN_API_KEY is set (required for CI/CD)
if [ -z "$GITGUARDIAN_API_KEY" ]; then
echo "Warning: GITGUARDIAN_API_KEY environment variable is not set."
echo "ggshield requires a GitGuardian API key to scan for secrets."
echo "Please set GITGUARDIAN_API_KEY in your CI/CD environment variables."
exit 1
fi
# # Check if GITGUARDIAN_API_KEY is set (required for CI/CD)
# if [ -z "$GITGUARDIAN_API_KEY" ]; then
# echo "Warning: GITGUARDIAN_API_KEY environment variable is not set."
# echo "ggshield requires a GitGuardian API key to scan for secrets."
# echo "Please set GITGUARDIAN_API_KEY in your CI/CD environment variables."
# exit 1
# fi
echo "Scanning codebase for secrets..."
echo "Note: Large codebases may take several minutes due to API rate limits (50 requests/minute on free plan)"
echo "ggshield will automatically handle rate limits and retry as needed."
echo "Binary files, cache files, and build artifacts are excluded via .gitguardian.yaml"
# echo "Scanning codebase for secrets..."
# echo "Note: Large codebases may take several minutes due to API rate limits (50 requests/minute on free plan)"
# echo "ggshield will automatically handle rate limits and retry as needed."
# echo "Binary files, cache files, and build artifacts are excluded via .gitguardian.yaml"
# Use --recursive for directory scanning and auto-confirm if prompted
# .gitguardian.yaml will automatically exclude binary files, wheel files, etc.
# GITGUARDIAN_API_KEY environment variable will be used for authentication
echo y | ggshield secret scan path . --recursive || {
echo ""
echo "=========================================="
echo "ERROR: Secret Detection Failed"
echo "=========================================="
echo "ggshield has detected secrets in the codebase."
echo "Please review discovered secrets above, revoke any actively used secrets"
echo "from underlying systems and make changes to inject secrets dynamically at runtime."
echo ""
echo "For more information, see: https://docs.gitguardian.com/secrets-detection/"
echo "=========================================="
echo ""
exit 1
}
# # Use --recursive for directory scanning and auto-confirm if prompted
# # .gitguardian.yaml will automatically exclude binary files, wheel files, etc.
# # GITGUARDIAN_API_KEY environment variable will be used for authentication
# echo y | ggshield secret scan path . --recursive || {
# echo ""
# echo "=========================================="
# echo "ERROR: Secret Detection Failed"
# echo "=========================================="
# echo "ggshield has detected secrets in the codebase."
# echo "Please review discovered secrets above, revoke any actively used secrets"
# echo "from underlying systems and make changes to inject secrets dynamically at runtime."
# echo ""
# echo "For more information, see: https://docs.gitguardian.com/secrets-detection/"
# echo "=========================================="
# echo ""
# exit 1
# }
echo "Secret detection scans completed successfully"
}
# echo "Secret detection scans completed successfully"
# }
# Function to run Trivy scans
run_trivy_scans() {
@ -209,8 +209,8 @@ main() {
install_trivy
install_grype
echo "Running secret detection scans..."
run_secret_detection
# echo "Running secret detection scans..."
# run_secret_detection
echo "Running filesystem vulnerability scans..."
run_trivy_scans

View file

@ -8,7 +8,8 @@ WORKDIR /app
COPY config.yaml .
# Make sure your docker/entrypoint.sh is executable
RUN chmod +x docker/entrypoint.sh
# Convert Windows line endings to Unix
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh
# Expose the necessary port
EXPOSE 4000/tcp

View file

@ -46,8 +46,9 @@ COPY --from=builder /wheels/ /wheels/
# Install the built wheel using pip; again using a wildcard if it's the only file
RUN pip install *.whl /wheels/* --no-index --find-links=/wheels/ && rm -f *.whl && rm -rf /wheels
RUN chmod +x docker/entrypoint.sh
RUN chmod +x docker/prod_entrypoint.sh
# Convert Windows line endings to Unix for entrypoint scripts
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh
RUN sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
EXPOSE 4000/tcp

View file

@ -32,8 +32,9 @@ RUN rm -rf /app/litellm/proxy/_experimental/out/* && \
WORKDIR /app
# Make sure your docker/entrypoint.sh is executable
RUN chmod +x docker/entrypoint.sh
RUN chmod +x docker/prod_entrypoint.sh
# Convert Windows line endings to Unix for entrypoint scripts
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh
RUN sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
# Expose the necessary port
EXPOSE 4000/tcp

View file

@ -27,7 +27,8 @@ RUN python -m pip install build
COPY . .
# Build Admin UI
RUN chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
# Convert Windows line endings to Unix and make executable
RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
# Build the package
RUN rm -rf dist/* && python -m build
@ -63,20 +64,23 @@ COPY --from=builder /wheels/ /wheels/
RUN pip install *.whl /wheels/* --no-index --find-links=/wheels/ && rm -f *.whl && rm -rf /wheels
# Install semantic_router and aurelio-sdk using script
RUN chmod +x docker/install_auto_router.sh && ./docker/install_auto_router.sh
# Convert Windows line endings to Unix and make executable
RUN sed -i 's/\r$//' docker/install_auto_router.sh && chmod +x docker/install_auto_router.sh && ./docker/install_auto_router.sh
# ensure pyjwt is used, not jwt
RUN pip uninstall jwt -y
RUN pip uninstall PyJWT -y
RUN pip install PyJWT==2.9.0 --no-cache-dir
# Build Admin UI
RUN chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
# Build Admin UI (runtime stage)
# Convert Windows line endings to Unix and make executable
RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
# Generate prisma client
RUN prisma generate
RUN chmod +x docker/entrypoint.sh
RUN chmod +x docker/prod_entrypoint.sh
# Convert Windows line endings to Unix for entrypoint scripts
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh
RUN sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
EXPOSE 4000/tcp
RUN apk add --no-cache supervisor

View file

@ -40,7 +40,8 @@ COPY enterprise/ ./enterprise/
COPY docker/ ./docker/
# Build Admin UI once
RUN chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
# Convert Windows line endings to Unix and make executable
RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
# Build the package
RUN rm -rf dist/* && python -m build
@ -79,8 +80,12 @@ RUN pip install --no-cache-dir *.whl /wheels/* --no-index --find-links=/wheels/
rm -rf /wheels
# Generate prisma client and set permissions
# Convert Windows line endings to Unix for entrypoint scripts
RUN prisma generate && \
chmod +x docker/entrypoint.sh docker/prod_entrypoint.sh
sed -i 's/\r$//' docker/entrypoint.sh && \
sed -i 's/\r$//' docker/prod_entrypoint.sh && \
chmod +x docker/entrypoint.sh && \
chmod +x docker/prod_entrypoint.sh
EXPOSE 4000/tcp

View file

@ -144,7 +144,10 @@ RUN pip install --no-index --find-links=/wheels/ -r requirements.txt && \
fi
# Permissions, cleanup, and Prisma prep
RUN chmod +x docker/entrypoint.sh docker/prod_entrypoint.sh && \
# Convert Windows line endings to Unix for entrypoint scripts
RUN sed -i 's/\r$//' docker/entrypoint.sh && \
sed -i 's/\r$//' docker/prod_entrypoint.sh && \
chmod +x docker/entrypoint.sh docker/prod_entrypoint.sh && \
mkdir -p /nonexistent /.npm /var/lib/litellm/assets /var/lib/litellm/ui && \
chown -R nobody:nogroup /app /var/lib/litellm/ui /var/lib/litellm/assets /nonexistent /.npm && \
pip uninstall jwt -y || true && \

View file

@ -21,6 +21,7 @@ Looking for how to use Code Interpreter? See the [Code Interpreter Guide](/docs/
| Endpoint | Method | Description |
|----------|--------|-------------|
| `/v1/containers/{container_id}/files` | POST | Upload file to container |
| `/v1/containers/{container_id}/files` | GET | List files in container |
| `/v1/containers/{container_id}/files/{file_id}` | GET | Get file metadata |
| `/v1/containers/{container_id}/files/{file_id}/content` | GET | Download file content |
@ -28,6 +29,45 @@ Looking for how to use Code Interpreter? See the [Code Interpreter Guide](/docs/
## LiteLLM Python SDK
### Upload Container File
Upload files directly to a container session. This is useful when `/chat/completions` or `/responses` sends files to the container but the input file type is limited to PDF. This endpoint lets you work with other file types like CSV, Excel, Python scripts, etc.
```python showLineNumbers title="upload_container_file.py"
from litellm import upload_container_file
# Upload a CSV file
file = upload_container_file(
container_id="cntr_123...",
file=("data.csv", open("data.csv", "rb").read(), "text/csv"),
custom_llm_provider="openai"
)
print(f"Uploaded: {file.id}")
print(f"Path: {file.path}")
```
**Async:**
```python showLineNumbers title="aupload_container_file.py"
from litellm import aupload_container_file
file = await aupload_container_file(
container_id="cntr_123...",
file=("script.py", b"print('hello world')", "text/x-python"),
custom_llm_provider="openai"
)
```
**Supported file formats:**
- CSV (`.csv`)
- Excel (`.xlsx`)
- Python scripts (`.py`)
- JSON (`.json`)
- Markdown (`.md`)
- Text files (`.txt`)
- And more...
### List Container Files
```python showLineNumbers title="list_container_files.py"
@ -103,6 +143,40 @@ print(f"Deleted: {result.deleted}")
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
### Upload File
<Tabs>
<TabItem value="openai-sdk" label="OpenAI SDK">
```python showLineNumbers title="upload_file.py"
from openai import OpenAI
client = OpenAI(
api_key="sk-1234",
base_url="http://localhost:4000"
)
file = client.containers.files.create(
container_id="cntr_123...",
file=open("data.csv", "rb")
)
print(f"Uploaded: {file.id}")
print(f"Path: {file.path}")
```
</TabItem>
<TabItem value="curl" label="curl">
```bash showLineNumbers title="upload_file.sh"
curl "http://localhost:4000/v1/containers/cntr_123.../files" \
-H "Authorization: Bearer sk-1234" \
-F file="@data.csv"
```
</TabItem>
</Tabs>
### List Files
<Tabs>
@ -236,6 +310,13 @@ curl -X DELETE "http://localhost:4000/v1/containers/cntr_123.../files/cfile_456.
## Parameters
### Upload File
| Parameter | Type | Required | Description |
|-----------|------|----------|-------------|
| `container_id` | string | Yes | Container ID |
| `file` | FileTypes | Yes | File to upload. Can be a tuple of (filename, content, content_type), file-like object, or bytes |
### List Files
| Parameter | Type | Required | Description |

View file

@ -8,7 +8,7 @@ import TabItem from '@theme/TabItem';
| Logging | ✅ | Works across all integrations |
| Streaming | ✅ | |
| Loadbalancing | ✅ | Between supported models |
| Supported LLM providers | **All LiteLLM supported providers** | `openai`, `anthropic`, `bedrock`, `vertex_ai`, `gemini`, `azure`, `azure_ai` etc. |
| Supported LLM providers | **All LiteLLM supported CHAT COMPLETION providers** | `openai`, `anthropic`, `bedrock`, `vertex_ai`, `gemini`, `azure`, `azure_ai` etc. |
## **LiteLLM Python SDK Usage**

View file

@ -68,6 +68,7 @@ environment_variables:
ARIZE_API_KEY: "141a****"
ARIZE_ENDPOINT: "https://otlp.arize.com/v1" # OPTIONAL - your custom arize GRPC api endpoint
ARIZE_HTTP_ENDPOINT: "https://otlp.arize.com/v1" # OPTIONAL - your custom arize HTTP api endpoint. Set either this or ARIZE_ENDPOINT or Neither (defaults to https://otlp.arize.com/v1 on grpc)
ARIZE_PROJECT_NAME: "my-litellm-project" # OPTIONAL - sets the arize project name
```
2. Start the proxy

View file

@ -1,7 +1,7 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Azure AI Image Generation
# Azure AI Image Generation (Black Forest Labs - Flux)
Azure AI provides powerful image generation capabilities using FLUX models from Black Forest Labs to create high-quality images from text descriptions.
@ -12,7 +12,7 @@ Azure AI provides powerful image generation capabilities using FLUX models from
| Description | Azure AI Image Generation uses FLUX models to generate high-quality images from text descriptions. |
| Provider Route on LiteLLM | `azure_ai/` |
| Provider Doc | [Azure AI FLUX Models ↗](https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/black-forest-labs-flux-1-kontext-pro-and-flux1-1-pro-now-available-in-azure-ai-f/4434659) |
| Supported Operations | [`/images/generations`](#image-generation) |
| Supported Operations | [`/images/generations`](#image-generation), [`/images/edits`](#image-editing) |
## Setup
@ -33,6 +33,7 @@ Get your API key and endpoint from [Azure AI Studio](https://ai.azure.com/).
|------------|-------------|----------------|
| `azure_ai/FLUX-1.1-pro` | Latest FLUX 1.1 Pro model for high-quality image generation | $0.04 |
| `azure_ai/FLUX.1-Kontext-pro` | FLUX 1 Kontext Pro model with enhanced context understanding | $0.04 |
| `azure_ai/flux.2-pro` | FLUX 2 Pro model for next-generation image generation | $0.04 |
## Image Generation
@ -85,6 +86,32 @@ print(response.data[0].url)
</TabItem>
<TabItem value="flux2" label="FLUX 2 Pro">
```python showLineNumbers title="FLUX 2 Pro Image Generation"
import litellm
import os
# Set your API credentials
os.environ["AZURE_AI_API_KEY"] = "your-api-key-here"
os.environ["AZURE_AI_API_BASE"] = "your-azure-ai-endpoint" # e.g., https://litellm-ci-cd-prod.services.ai.azure.com
# Generate image with FLUX 2 Pro
response = litellm.image_generation(
model="azure_ai/flux.2-pro",
prompt="A photograph of a red fox in an autumn forest",
api_base=os.environ["AZURE_AI_API_BASE"],
api_key=os.environ["AZURE_AI_API_KEY"],
api_version="preview",
size="1024x1024",
n=1
)
print(response.data[0].b64_json) # FLUX 2 returns base64 encoded images
```
</TabItem>
<TabItem value="async" label="Async Usage">
```python showLineNumbers title="Async Image Generation"
@ -165,6 +192,15 @@ model_list:
model_info:
mode: image_generation
- model_name: azure-flux-2-pro
litellm_params:
model: azure_ai/flux.2-pro
api_key: os.environ/AZURE_AI_API_KEY
api_base: os.environ/AZURE_AI_API_BASE
api_version: preview
model_info:
mode: image_generation
general_settings:
master_key: sk-1234
```
@ -239,6 +275,103 @@ curl --location 'http://localhost:4000/v1/images/generations' \
</TabItem>
</Tabs>
## Image Editing
FLUX 2 Pro supports image editing by passing an input image along with a prompt describing the desired modifications.
### Usage - LiteLLM Python SDK
<Tabs>
<TabItem value="basic-edit" label="Basic Image Edit">
```python showLineNumbers title="Basic Image Editing with FLUX 2 Pro"
import litellm
import os
# Set your API credentials
os.environ["AZURE_AI_API_KEY"] = "your-api-key-here"
os.environ["AZURE_AI_API_BASE"] = "your-azure-ai-endpoint" # e.g., https://litellm-ci-cd-prod.services.ai.azure.com
# Edit an existing image
response = litellm.image_edit(
model="azure_ai/flux.2-pro",
prompt="Add a red hat to the subject",
image=open("input_image.png", "rb"),
api_base=os.environ["AZURE_AI_API_BASE"],
api_key=os.environ["AZURE_AI_API_KEY"],
api_version="preview",
)
print(response.data[0].b64_json) # FLUX 2 returns base64 encoded images
```
</TabItem>
<TabItem value="async-edit" label="Async Image Edit">
```python showLineNumbers title="Async Image Editing"
import litellm
import asyncio
import os
async def edit_image():
os.environ["AZURE_AI_API_KEY"] = "your-api-key-here"
os.environ["AZURE_AI_API_BASE"] = "your-azure-ai-endpoint"
response = await litellm.aimage_edit(
model="azure_ai/flux.2-pro",
prompt="Change the background to a sunset beach",
image=open("input_image.png", "rb"),
api_base=os.environ["AZURE_AI_API_BASE"],
api_key=os.environ["AZURE_AI_API_KEY"],
api_version="preview",
)
return response
asyncio.run(edit_image())
```
</TabItem>
</Tabs>
### Usage - LiteLLM Proxy Server
<Tabs>
<TabItem value="curl-edit" label="cURL">
```bash showLineNumbers title="Image Edit via Proxy - cURL"
curl --location 'http://localhost:4000/v1/images/edits' \
--header 'Authorization: Bearer sk-1234' \
--form 'model="azure-flux-2-pro"' \
--form 'prompt="Add sunglasses to the person"' \
--form 'image=@"input_image.png"'
```
</TabItem>
<TabItem value="openai-sdk-edit" label="OpenAI SDK">
```python showLineNumbers title="Image Edit via Proxy - OpenAI SDK"
from openai import OpenAI
client = OpenAI(
base_url="http://localhost:4000",
api_key="sk-1234"
)
response = client.images.edit(
model="azure-flux-2-pro",
prompt="Make the sky more dramatic with storm clouds",
image=open("input_image.png", "rb"),
)
print(response.data[0].b64_json)
```
</TabItem>
</Tabs>
## Supported Parameters
Azure AI Image Generation supports the following OpenAI-compatible parameters:

View file

@ -498,6 +498,7 @@ router_settings:
| DD_VERSION | Version identifier for Datadog logs. Defaults to "unknown"
| DEBUG_OTEL | Enable debug mode for OpenTelemetry
| DEFAULT_ALLOWED_FAILS | Maximum failures allowed before cooling down a model. Default is 3
| DEFAULT_A2A_AGENT_TIMEOUT | Default timeout in seconds for A2A (Agent-to-Agent) protocol requests. Default is 6000
| DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS | Default maximum tokens for Anthropic chat completions. Default is 4096
| DEFAULT_BATCH_SIZE | Default batch size for operations. Default is 512
| DEFAULT_CHUNK_OVERLAP | Default chunk overlap for RAG text splitters. Default is 200

View file

@ -576,10 +576,31 @@ custom_tokenizer:
```yaml
general_settings:
database_connection_pool_limit: 10 # sets connection pool for prisma client to postgres db (default: 10, recommended: 10-20)
database_connection_pool_limit: 10 # sets connection pool per worker for prisma client to postgres db (default: 10, recommended: 10-20)
database_connection_timeout: 60 # sets a 60s timeout for any connection call to the db
```
**How to calculate the right value:**
The connection limit is applied **per worker process**, not per instance. This means if you have multiple workers, each worker will create its own connection pool.
**Formula:**
```
database_connection_pool_limit = MAX_DB_CONNECTIONS ÷ (number_of_instances × number_of_workers_per_instance)
```
**Example:**
- Your database allows a maximum of **100 connections**
- You're running **1 instance** of LiteLLM
- Each instance has **8 workers** (set via `--num_workers 8`)
Calculation: `100 ÷ (1 × 8) = 12.5`
Since you shouldn't use 12.5, round down to **10** to leave a safety buffer. This means:
- Each of the 8 workers will have a connection pool limit of 10
- Total maximum connections: 8 workers × 10 connections = 80 connections
- This stays safely under your database's 100 connection limit
## Extras

View file

@ -19,7 +19,11 @@ general_settings:
master_key: sk-1234 # enter your own master key, ensure it starts with 'sk-'
alerting: ["slack"] # Setup slack alerting - get alerts on LLM exceptions, Budget Alerts, Slow LLM Responses
proxy_batch_write_at: 60 # Batch write spend updates every 60s
database_connection_pool_limit: 10 # limit the number of database connections to = MAX Number of DB Connections/Number of instances of litellm proxy (Around 10-20 is good number)
database_connection_pool_limit: 10 # connection pool limit per worker process. Total connections = limit × workers × instances. Calculate: MAX_DB_CONNECTIONS / (instances × workers). Default: 10.
:::warning
**Multiple instances:** If running multiple LiteLLM instances (e.g., Kubernetes pods), remember each instance multiplies your total connections. Example: 3 instances × 4 workers × 10 connections = 120 total connections.
:::
# OPTIONAL Best Practices
disable_error_logs: True # turn off writing LLM Exceptions to DB
@ -54,8 +58,8 @@ For optimal performance in production, we recommend the following minimum machin
| Resource | Recommended Value |
|----------|------------------|
| CPU | 2 vCPU |
| Memory | 4 GB RAM |
| CPU | 4 vCPU |
| Memory | 8 GB RAM |
These specifications provide:
- Sufficient compute power for handling concurrent requests

View file

@ -5,6 +5,12 @@ import TabItem from '@theme/TabItem';
Use this to loadbalance across Azure + OpenAI.
Supported Providers:
- OpenAI
- Azure
- Google AI Studio (Gemini)
- Vertex AI
## Proxy Usage
### Add model to config

View file

@ -420,14 +420,8 @@ const sidebars = {
],
},
"assistants",
{
type: "category",
label: "/audio",
items: [
"audio_transcription",
"text_to_speech",
]
},
"audio_transcription",
"text_to_speech",
{
type: "category",
label: "/batches",
@ -477,17 +471,13 @@ const sidebars = {
"apply_guardrail",
"bedrock_invoke",
"interactions",
{
type: "category",
label: "/images",
items: [
"image_edits",
"image_generation",
"image_variations",
]
},
"image_edits",
"image_generation",
"image_variations",
"videos",
"vector_store_files",
"vector_stores/create",
"vector_stores/search",
{
type: "category",
label: "/mcp - Model Context Protocol",
@ -531,24 +521,12 @@ const sidebars = {
"proxy/pass_through_guardrails"
]
},
{
type: "category",
label: "/rag",
items: [
"rag_ingest",
"rag_query",
]
},
"rag_ingest",
"rag_query",
"realtime",
"rerank",
{
type: "category",
label: "/responses",
items: [
"response_api",
"response_api_compact",
]
},
"response_api",
"response_api_compact",
{
type: "category",
label: "/search",
@ -566,14 +544,7 @@ const sidebars = {
]
},
"skills",
{
type: "category",
label: "/vector_stores",
items: [
"vector_stores/create",
"vector_stores/search",
]
},
],
},
{

BIN
flux2_test_image.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 172 KiB

Binary file not shown.

Binary file not shown.

Binary file not shown.

Binary file not shown.

View file

@ -0,0 +1,72 @@
-- DropIndex
DROP INDEX "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key";
-- DropIndex
DROP INDEX "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key";
-- DropIndex
DROP INDEX "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key";
-- DropIndex
DROP INDEX "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key";
-- DropIndex
DROP INDEX "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key";
-- DropIndex
DROP INDEX "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key";
-- AlterTable
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN "endpoint" TEXT;
-- AlterTable
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN "endpoint" TEXT;
-- AlterTable
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN "endpoint" TEXT;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN "endpoint" TEXT;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN "endpoint" TEXT;
-- AlterTable
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN "endpoint" TEXT;
-- CreateIndex
CREATE INDEX "LiteLLM_DailyAgentSpend_endpoint_idx" ON "LiteLLM_DailyAgentSpend"("endpoint");
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key" ON "LiteLLM_DailyAgentSpend"("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
-- CreateIndex
CREATE INDEX "LiteLLM_DailyEndUserSpend_endpoint_idx" ON "LiteLLM_DailyEndUserSpend"("endpoint");
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
-- CreateIndex
CREATE INDEX "LiteLLM_DailyOrganizationSpend_endpoint_idx" ON "LiteLLM_DailyOrganizationSpend"("endpoint");
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
-- CreateIndex
CREATE INDEX "LiteLLM_DailyTagSpend_endpoint_idx" ON "LiteLLM_DailyTagSpend"("endpoint");
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
-- CreateIndex
CREATE INDEX "LiteLLM_DailyTeamSpend_endpoint_idx" ON "LiteLLM_DailyTeamSpend"("endpoint");
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
-- CreateIndex
CREATE INDEX "LiteLLM_DailyUserSpend_endpoint_idx" ON "LiteLLM_DailyUserSpend"("endpoint");
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");

View file

@ -0,0 +1,6 @@
-- AlterTable
ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "router_settings" JSONB DEFAULT '{}';
-- AlterTable
ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "router_settings" JSONB DEFAULT '{}';

View file

@ -124,6 +124,7 @@ model LiteLLM_TeamTable {
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
@ -225,6 +226,7 @@ model LiteLLM_VerificationToken {
models String[]
aliases Json @default("{}")
config Json @default("{}")
router_settings Json? @default("{}")
user_id String?
team_id String?
permissions Json @default("{}")
@ -422,6 +424,7 @@ model LiteLLM_DailyUserSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -433,12 +436,13 @@ model LiteLLM_DailyUserSpend {
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([user_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily organization spend metrics per model and key
@ -451,6 +455,7 @@ model LiteLLM_DailyOrganizationSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -462,12 +467,13 @@ model LiteLLM_DailyOrganizationSpend {
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([organization_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily end user (customer) spend metrics per model and key
@ -480,6 +486,7 @@ model LiteLLM_DailyEndUserSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -490,12 +497,13 @@ model LiteLLM_DailyEndUserSpend {
failed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([end_user_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily agent spend metrics per model and key
@ -508,6 +516,7 @@ model LiteLLM_DailyAgentSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -518,12 +527,13 @@ model LiteLLM_DailyAgentSpend {
failed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([agent_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily team spend metrics per model and key
@ -536,6 +546,7 @@ model LiteLLM_DailyTeamSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -547,12 +558,13 @@ model LiteLLM_DailyTeamSpend {
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([team_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily team spend metrics per model and key
@ -566,6 +578,7 @@ model LiteLLM_DailyTagSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -577,12 +590,13 @@ model LiteLLM_DailyTagSpend {
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([tag])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}

View file

@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm-proxy-extras"
version = "0.4.16"
version = "0.4.20"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
authors = ["BerriAI"]
readme = "README.md"
@ -22,7 +22,7 @@ requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "0.4.16"
version = "0.4.20"
version_files = [
"pyproject.toml:version",
"../requirements.txt:litellm-proxy-extras==",

View file

@ -486,6 +486,7 @@ vertex_mistral_models: Set = set()
vertex_openai_models: Set = set()
vertex_minimax_models: Set = set()
vertex_moonshot_models: Set = set()
vertex_zai_models: Set = set()
ai21_models: Set = set()
ai21_chat_models: Set = set()
nlp_cloud_models: Set = set()
@ -557,6 +558,8 @@ stability_models: Set = set()
github_copilot_models: Set = set()
minimax_models: Set = set()
aws_polly_models: Set = set()
gigachat_models: Set = set()
llamagate_models: Set = set()
def is_bedrock_pricing_only_model(key: str) -> bool:
@ -662,6 +665,9 @@ def add_known_models():
elif value.get("litellm_provider") == "vertex_ai-moonshot_models":
key = key.replace("vertex_ai/", "")
vertex_moonshot_models.add(key)
elif value.get("litellm_provider") == "vertex_ai-zai_models":
key = key.replace("vertex_ai/", "")
vertex_zai_models.add(key)
elif value.get("litellm_provider") == "ai21":
if value.get("mode") == "chat":
ai21_chat_models.add(key)
@ -809,6 +815,10 @@ def add_known_models():
minimax_models.add(key)
elif value.get("litellm_provider") == "aws_polly":
aws_polly_models.add(key)
elif value.get("litellm_provider") == "gigachat":
gigachat_models.add(key)
elif value.get("litellm_provider") == "llamagate":
llamagate_models.add(key)
add_known_models()
@ -944,7 +954,8 @@ models_by_provider: dict = {
| vertex_language_models
| vertex_deepseek_models
| vertex_minimax_models
| vertex_moonshot_models,
| vertex_moonshot_models
| vertex_zai_models,
"ai21": ai21_models,
"bedrock": bedrock_models | bedrock_converse_models,
"petals": petals_models,
@ -1015,6 +1026,8 @@ models_by_provider: dict = {
"github_copilot": github_copilot_models,
"minimax": minimax_models,
"aws_polly": aws_polly_models,
"gigachat": gigachat_models,
"llamagate": llamagate_models,
}
# mapping for those models which have larger equivalents

View file

@ -12,6 +12,7 @@ import litellm
from litellm._logging import verbose_logger
from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator
from litellm.a2a_protocol.utils import A2ARequestUtils
from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -494,7 +495,7 @@ async def create_a2a_client(
async def aget_agent_card(
base_url: str,
timeout: float = 60.0,
timeout: float = DEFAULT_A2A_AGENT_TIMEOUT,
extra_headers: Optional[Dict[str, str]] = None,
) -> "AgentCard":
"""

View file

@ -278,6 +278,7 @@ MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB = int(
DEFAULT_MAX_TOKENS_FOR_TRITON = int(os.getenv("DEFAULT_MAX_TOKENS_FOR_TRITON", 2000))
#### Networking settings ####
request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", 6000)) # time in seconds
DEFAULT_A2A_AGENT_TIMEOUT: float = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes
STREAM_SSE_DONE_STRING: str = "[DONE]"
STREAM_SSE_DATA_PREFIX: str = "data: "
### SPEND TRACKING ###

View file

@ -216,6 +216,8 @@ _generated_endpoints = generate_container_endpoints()
# Export generated functions dynamically
list_container_files = _generated_endpoints.get("list_container_files")
alist_container_files = _generated_endpoints.get("alist_container_files")
upload_container_file = _generated_endpoints.get("upload_container_file")
aupload_container_file = _generated_endpoints.get("aupload_container_file")
retrieve_container_file = _generated_endpoints.get("retrieve_container_file")
aretrieve_container_file = _generated_endpoints.get("aretrieve_container_file")
delete_container_file = _generated_endpoints.get("delete_container_file")

View file

@ -9,6 +9,16 @@
"query_params": ["after", "limit", "order"],
"response_type": "ContainerFileListResponse"
},
{
"name": "upload_container_file",
"async_name": "aupload_container_file",
"path": "/containers/{container_id}/files",
"method": "POST",
"path_params": ["container_id"],
"query_params": [],
"response_type": "ContainerFileObject",
"is_multipart": true
},
{
"name": "retrieve_container_file",
"async_name": "aretrieve_container_file",

View file

@ -13,11 +13,13 @@ from litellm.main import base_llm_http_handler
from litellm.types.containers.main import (
ContainerCreateOptionalRequestParams,
ContainerFileListResponse,
ContainerFileObject,
ContainerListOptionalRequestParams,
ContainerListResponse,
ContainerObject,
DeleteContainerResult,
)
from litellm.types.llms.openai import FileTypes
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import CallTypes
from litellm.utils import ProviderConfigManager, client
@ -28,11 +30,13 @@ __all__ = [
"alist_container_files",
"alist_containers",
"aretrieve_container",
"aupload_container_file",
"create_container",
"delete_container",
"list_container_files",
"list_containers",
"retrieve_container",
"upload_container_file",
]
##### Container Create #######################
@ -1011,3 +1015,236 @@ def list_container_files(
extra_kwargs=kwargs,
)
##### Container File Upload #######################
@client
async def aupload_container_file(
container_id: str,
file: FileTypes,
timeout=600, # default to 10 minutes
custom_llm_provider: Literal["openai"] = "openai",
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
**kwargs,
) -> ContainerFileObject:
"""Asynchronously upload a file to a container.
This endpoint allows uploading files directly to a container session,
supporting various file types like CSV, Excel, Python scripts, etc.
Parameters:
- `container_id` (str): The ID of the container to upload the file to
- `file` (FileTypes): The file to upload. Can be:
- A tuple of (filename, content, content_type)
- A tuple of (filename, content)
- A file-like object with read() method
- Bytes
- A string path to a file
- `timeout` (int): Request timeout in seconds
- `custom_llm_provider` (Literal["openai"]): The LLM provider to use
- `extra_headers` (Optional[Dict[str, Any]]): Additional headers
- `extra_query` (Optional[Dict[str, Any]]): Additional query parameters
- `extra_body` (Optional[Dict[str, Any]]): Additional body parameters
- `kwargs` (dict): Additional keyword arguments
Returns:
- `response` (ContainerFileObject): The uploaded file object
Example:
```python
import litellm
# Upload a CSV file
response = await litellm.aupload_container_file(
container_id="container_abc123",
file=("data.csv", open("data.csv", "rb").read(), "text/csv"),
custom_llm_provider="openai",
)
print(response)
```
"""
local_vars = locals()
try:
loop = asyncio.get_event_loop()
kwargs["async_call"] = True
func = partial(
upload_container_file,
container_id=container_id,
file=file,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
extra_query=extra_query,
extra_body=extra_body,
**kwargs,
)
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
init_response = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
response = init_response
return response
except Exception as e:
raise litellm.exception_type(
model="",
custom_llm_provider=custom_llm_provider,
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)
# fmt: off
@overload
def upload_container_file(
container_id: str,
file: FileTypes,
timeout=600,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
custom_llm_provider: Literal["openai"] = "openai",
*,
aupload_container_file: Literal[True],
**kwargs,
) -> Coroutine[Any, Any, ContainerFileObject]:
...
@overload
def upload_container_file(
container_id: str,
file: FileTypes,
timeout=600,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
custom_llm_provider: Literal["openai"] = "openai",
*,
aupload_container_file: Literal[False] = False,
**kwargs,
) -> ContainerFileObject:
...
# fmt: on
@client
def upload_container_file(
container_id: str,
file: FileTypes,
timeout=600, # default to 10 minutes
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
custom_llm_provider: Literal["openai"] = "openai",
extra_headers: Optional[Dict[str, Any]] = None,
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Union[
ContainerFileObject,
Coroutine[Any, Any, ContainerFileObject],
]:
"""Upload a file to a container using the OpenAI Container API.
This endpoint allows uploading files directly to a container session,
supporting various file types like CSV, Excel, Python scripts, JSON, etc.
This is useful when /chat/completions or /responses sends files to the
container but the input file type is limited to PDF. This endpoint lets
you work with other file types.
Currently supports OpenAI
Example:
```python
import litellm
# Upload a CSV file
response = litellm.upload_container_file(
container_id="container_abc123",
file=("data.csv", open("data.csv", "rb").read(), "text/csv"),
custom_llm_provider="openai",
)
print(response)
# Upload a Python script
response = litellm.upload_container_file(
container_id="container_abc123",
file=("script.py", b"print('hello world')", "text/x-python"),
custom_llm_provider="openai",
)
print(response)
```
"""
from litellm.llms.custom_httpx.container_handler import generic_container_handler
local_vars = locals()
try:
litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id")
_is_async = kwargs.pop("async_call", False) is True
# Check for mock response first
mock_response = kwargs.get("mock_response")
if mock_response is not None:
if isinstance(mock_response, str):
mock_response = json.loads(mock_response)
response = ContainerFileObject(**mock_response)
return response
# get llm provider logic
litellm_params = GenericLiteLLMParams(**kwargs)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
)
if container_provider_config is None:
raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}")
# Pre Call logging
litellm_logging_obj.update_environment_variables(
model="",
optional_params={"container_id": container_id},
litellm_params={
"litellm_call_id": litellm_call_id,
},
custom_llm_provider=custom_llm_provider,
)
# Set the correct call type
litellm_logging_obj.call_type = CallTypes.upload_container_file.value
return generic_container_handler.handle(
endpoint_name="upload_container_file",
container_provider_config=container_provider_config,
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
extra_headers=extra_headers,
extra_query=extra_query,
timeout=timeout or DEFAULT_REQUEST_TIMEOUT,
_is_async=_is_async,
container_id=container_id,
file=file,
)
except Exception as e:
raise litellm.exception_type(
model="",
custom_llm_provider=custom_llm_provider,
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)

View file

@ -51,6 +51,7 @@ class ArizeLogger(OpenTelemetry):
space_id = os.environ.get("ARIZE_SPACE_ID")
space_key = os.environ.get("ARIZE_SPACE_KEY")
api_key = os.environ.get("ARIZE_API_KEY")
project_name = os.environ.get("ARIZE_PROJECT_NAME")
grpc_endpoint = os.environ.get("ARIZE_ENDPOINT")
http_endpoint = os.environ.get("ARIZE_HTTP_ENDPOINT")
@ -74,6 +75,7 @@ class ArizeLogger(OpenTelemetry):
api_key=api_key,
protocol=protocol,
endpoint=endpoint,
project_name=project_name,
)
async def async_service_success_hook(

View file

@ -54,38 +54,6 @@ RAW_REQUEST_SPAN_NAME = "raw_gen_ai_request"
LITELLM_REQUEST_SPAN_NAME = "litellm_request"
def _get_litellm_resource():
"""
Create a proper OpenTelemetry Resource that respects OTEL_RESOURCE_ATTRIBUTES
while maintaining backward compatibility with LiteLLM-specific environment variables.
"""
from opentelemetry.sdk.resources import OTELResourceDetector, Resource
# Create base resource attributes with LiteLLM-specific defaults
# These will be overridden by OTEL_RESOURCE_ATTRIBUTES if present
base_attributes: Dict[str, Optional[str]] = {
"service.name": os.getenv("OTEL_SERVICE_NAME", "litellm"),
"deployment.environment": os.getenv("OTEL_ENVIRONMENT_NAME", "production"),
# Fix the model_id to use proper environment variable or default to service name
"model_id": os.getenv(
"OTEL_MODEL_ID", os.getenv("OTEL_SERVICE_NAME", "litellm")
),
}
# Create base resource with LiteLLM-specific defaults
base_resource = Resource.create(base_attributes) # type: ignore
# Create resource from OTEL_RESOURCE_ATTRIBUTES using the detector
otel_resource_detector = OTELResourceDetector()
env_resource = otel_resource_detector.detect()
# Merge the resources: env_resource takes precedence over base_resource
# This ensures OTEL_RESOURCE_ATTRIBUTES overrides LiteLLM defaults
merged_resource = base_resource.merge(env_resource)
return merged_resource
@dataclass
class OpenTelemetryConfig:
exporter: Union[str, SpanExporter] = "console"
@ -93,6 +61,19 @@ class OpenTelemetryConfig:
headers: Optional[str] = None
enable_metrics: bool = False
enable_events: bool = False
service_name: Optional[str] = None
deployment_environment: Optional[str] = None
model_id: Optional[str] = None
def __post_init__(self) -> None:
if not self.service_name:
self.service_name = os.getenv("OTEL_SERVICE_NAME", "litellm")
if not self.deployment_environment:
self.deployment_environment = os.getenv(
"OTEL_ENVIRONMENT_NAME", "production"
)
if not self.model_id:
self.model_id = os.getenv("OTEL_MODEL_ID", self.service_name)
@classmethod
def from_env(cls):
@ -122,6 +103,9 @@ class OpenTelemetryConfig:
os.getenv("LITELLM_OTEL_INTEGRATION_ENABLE_EVENTS", "false").lower()
== "true"
)
service_name = os.getenv("OTEL_SERVICE_NAME", "litellm")
deployment_environment = os.getenv("OTEL_ENVIRONMENT_NAME", "production")
model_id = os.getenv("OTEL_MODEL_ID", service_name)
if exporter == "in_memory":
return cls(exporter=InMemorySpanExporter())
@ -131,6 +115,9 @@ class OpenTelemetryConfig:
headers=headers, # example: OTEL_HEADERS=x-honeycomb-team=B85YgLm96***"
enable_metrics=enable_metrics,
enable_events=enable_events,
service_name=service_name,
deployment_environment=deployment_environment,
model_id=model_id,
)
@ -174,6 +161,22 @@ class OpenTelemetry(CustomLogger):
self._init_logs(logger_provider)
self._init_otel_logger_on_litellm_proxy()
@staticmethod
def _get_litellm_resource(config: OpenTelemetryConfig):
"""Create an OpenTelemetry Resource using config-driven defaults."""
from opentelemetry.sdk.resources import OTELResourceDetector, Resource
base_attributes: Dict[str, Optional[str]] = {
"service.name": config.service_name,
"deployment.environment": config.deployment_environment,
"model_id": config.model_id or config.service_name,
}
base_resource = Resource.create(base_attributes) # type: ignore[arg-type]
otel_resource_detector = OTELResourceDetector()
env_resource = otel_resource_detector.detect()
return base_resource.merge(env_resource)
def _init_otel_logger_on_litellm_proxy(self):
"""
Initializes OpenTelemetry for litellm proxy server
@ -266,7 +269,7 @@ class OpenTelemetry(CustomLogger):
from opentelemetry.trace import SpanKind
def create_tracer_provider():
provider = TracerProvider(resource=_get_litellm_resource())
provider = TracerProvider(resource=self._get_litellm_resource(self.config))
provider.add_span_processor(self._get_span_processor())
return provider
@ -300,7 +303,8 @@ class OpenTelemetry(CustomLogger):
def create_meter_provider():
metric_reader = self._get_metric_reader()
return MeterProvider(
metric_readers=[metric_reader], resource=_get_litellm_resource()
metric_readers=[metric_reader],
resource=self._get_litellm_resource(self.config),
)
meter_provider = self._get_or_create_provider(
@ -355,7 +359,9 @@ class OpenTelemetry(CustomLogger):
from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
def create_logger_provider():
provider = OTLoggerProvider(resource=_get_litellm_resource())
provider = OTLoggerProvider(
resource=self._get_litellm_resource(self.config)
)
log_exporter = self._get_log_exporter()
provider.add_log_record_processor(
BatchLogRecordProcessor(log_exporter) # type: ignore[arg-type]
@ -606,7 +612,7 @@ class OpenTelemetry(CustomLogger):
from opentelemetry.sdk.trace import TracerProvider
# Create a temporary tracer provider with dynamic headers
temp_provider = TracerProvider(resource=_get_litellm_resource())
temp_provider = TracerProvider(resource=self._get_litellm_resource(self.config))
temp_provider.add_span_processor(
self._get_span_processor(dynamic_headers=dynamic_headers)
)
@ -987,9 +993,9 @@ class OpenTelemetry(CustomLogger):
# Get the resource from the logger provider
logger_provider = get_logger_provider()
resource = (
getattr(logger_provider, "_resource", None) or _get_litellm_resource()
)
resource = getattr(
logger_provider, "_resource", None
) or self._get_litellm_resource(self.config)
parent_ctx = span.get_span_context()
provider = (kwargs.get("litellm_params") or {}).get(
@ -1910,7 +1916,9 @@ class OpenTelemetry(CustomLogger):
)
_split_otel_headers = OpenTelemetry._get_headers_dictionary(self.OTEL_HEADERS)
normalized_endpoint = self._normalize_otel_endpoint(self.OTEL_ENDPOINT, "metrics")
normalized_endpoint = self._normalize_otel_endpoint(
self.OTEL_ENDPOINT, "metrics"
)
if self.OTEL_EXPORTER == "console":
exporter = ConsoleMetricExporter()

View file

@ -14,6 +14,7 @@ from typing import (
Literal,
Optional,
Tuple,
Union,
cast,
)
@ -791,6 +792,11 @@ class PrometheusLogger(CustomLogger):
f"standard_logging_object is required, got={standard_logging_payload}"
)
if self._should_skip_metrics_for_invalid_key(
kwargs=kwargs, standard_logging_payload=standard_logging_payload
):
return
model = kwargs.get("model", "")
litellm_params = kwargs.get("litellm_params", {}) or {}
_metadata = litellm_params.get("metadata", {})
@ -1189,11 +1195,17 @@ class PrometheusLogger(CustomLogger):
f"prometheus Logging - Enters failure logging function for kwargs {kwargs}"
)
# unpack kwargs
model = kwargs.get("model", "")
standard_logging_payload: StandardLoggingPayload = kwargs.get(
"standard_logging_object", {}
)
if self._should_skip_metrics_for_invalid_key(
kwargs=kwargs, standard_logging_payload=standard_logging_payload
):
return
model = kwargs.get("model", "")
litellm_params = kwargs.get("litellm_params", {}) or {}
get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking()
@ -1207,7 +1219,6 @@ class PrometheusLogger(CustomLogger):
user_api_team_alias = standard_logging_payload["metadata"][
"user_api_key_team_alias"
]
kwargs.get("exception", None)
try:
self.litellm_llm_api_failed_requests_metric.labels(
@ -1227,6 +1238,139 @@ class PrometheusLogger(CustomLogger):
pass
pass
def _extract_status_code(
self,
kwargs: Optional[dict] = None,
enum_values: Optional[Any] = None,
exception: Optional[Exception] = None,
) -> Optional[int]:
"""
Extract HTTP status code from various input formats for validation.
This is a centralized helper to extract status code from different
callback function signatures. Handles both ProxyException (uses 'code')
and standard exceptions (uses 'status_code').
Args:
kwargs: Dictionary potentially containing 'exception' key
enum_values: Object with 'status_code' attribute
exception: Exception object to extract status code from directly
Returns:
Status code as integer if found, None otherwise
"""
status_code = None
# Try from enum_values first (most common in our callbacks)
if enum_values and hasattr(enum_values, "status_code") and enum_values.status_code:
try:
status_code = int(enum_values.status_code)
except (ValueError, TypeError):
pass
if not status_code and exception:
# ProxyException uses 'code' attribute, other exceptions may use 'status_code'
status_code = getattr(exception, "status_code", None) or getattr(exception, "code", None)
if status_code is not None:
try:
status_code = int(status_code)
except (ValueError, TypeError):
status_code = None
if not status_code and kwargs:
exception_in_kwargs = kwargs.get("exception")
if exception_in_kwargs:
status_code = getattr(exception_in_kwargs, "status_code", None) or getattr(exception_in_kwargs, "code", None)
if status_code is not None:
try:
status_code = int(status_code)
except (ValueError, TypeError):
status_code = None
return status_code
def _is_invalid_api_key_request(
self,
status_code: Optional[int],
exception: Optional[Exception] = None,
) -> bool:
"""
Determine if a request has an invalid API key based on status code and exception.
This method prevents invalid authentication attempts from being recorded in
Prometheus metrics. A 401 status code is the definitive indicator of authentication
failure. Additionally, we check exception messages for authentication error patterns
to catch cases where the exception hasn't been converted to a ProxyException yet.
Args:
status_code: HTTP status code (401 indicates authentication error)
exception: Exception object to check for auth-related error messages
Returns:
True if the request has an invalid API key and metrics should be skipped,
False otherwise
"""
if status_code == 401:
return True
# Handle cases where AssertionError is raised before conversion to ProxyException
if exception is not None:
exception_str = str(exception).lower()
auth_error_patterns = [
"virtual key expected",
"expected to start with 'sk-'",
"authentication error",
"invalid api key",
"api key not valid",
]
if any(pattern in exception_str for pattern in auth_error_patterns):
return True
return False
def _should_skip_metrics_for_invalid_key(
self,
kwargs: Optional[dict] = None,
user_api_key_dict: Optional[Any] = None,
enum_values: Optional[Any] = None,
standard_logging_payload: Optional[Union[dict, StandardLoggingPayload]] = None,
exception: Optional[Exception] = None,
) -> bool:
"""
Determine if Prometheus metrics should be skipped for invalid API key requests.
This is a centralized validation method that extracts status code and exception
information from various callback function signatures and determines if the request
represents an invalid API key attempt that should be filtered from metrics.
Args:
kwargs: Dictionary potentially containing exception and other data
user_api_key_dict: User API key authentication object (currently unused)
enum_values: Object with status_code attribute
standard_logging_payload: Standard logging payload dictionary
exception: Exception object to check directly
Returns:
True if metrics should be skipped (invalid key detected), False otherwise
"""
status_code = self._extract_status_code(
kwargs=kwargs,
enum_values=enum_values,
exception=exception,
)
if exception is None and kwargs:
exception = kwargs.get("exception")
if self._is_invalid_api_key_request(status_code, exception=exception):
verbose_logger.debug(
"Skipping Prometheus metrics for invalid API key request: "
f"status_code={status_code}, exception={type(exception).__name__ if exception else None}"
)
return True
return False
async def async_post_call_failure_hook(
self,
request_data: dict,
@ -1252,6 +1396,14 @@ class PrometheusLogger(CustomLogger):
StandardLoggingPayloadSetup,
)
if self._should_skip_metrics_for_invalid_key(
user_api_key_dict=user_api_key_dict,
exception=original_exception,
):
return
status_code = self._extract_status_code(exception=original_exception)
try:
_tags = StandardLoggingPayloadSetup._get_request_tags(
litellm_params=request_data,
@ -1266,8 +1418,8 @@ class PrometheusLogger(CustomLogger):
team=user_api_key_dict.team_id,
team_alias=user_api_key_dict.team_alias,
requested_model=request_data.get("model", ""),
status_code=str(getattr(original_exception, "status_code", None)),
exception_status=str(getattr(original_exception, "status_code", None)),
status_code=str(status_code),
exception_status=str(status_code),
exception_class=self._get_exception_class_name(original_exception),
tags=_tags,
route=user_api_key_dict.request_route,
@ -1305,6 +1457,11 @@ class PrometheusLogger(CustomLogger):
StandardLoggingPayloadSetup,
)
if self._should_skip_metrics_for_invalid_key(
user_api_key_dict=user_api_key_dict
):
return
enum_values = UserAPIKeyLabelValues(
end_user=user_api_key_dict.end_user_id,
hashed_api_key=user_api_key_dict.api_key,
@ -1360,6 +1517,15 @@ class PrometheusLogger(CustomLogger):
exception = request_kwargs.get("exception", None)
llm_provider = _litellm_params.get("custom_llm_provider", None)
if self._should_skip_metrics_for_invalid_key(
kwargs=request_kwargs,
standard_logging_payload=standard_logging_payload,
):
return
hashed_api_key = standard_logging_payload.get("metadata", {}).get(
"user_api_key_hash"
)
# Create enum_values for the label factory (always create for use in different metrics)
enum_values = UserAPIKeyLabelValues(
@ -1374,9 +1540,7 @@ class PrometheusLogger(CustomLogger):
self._get_exception_class_name(exception) if exception else None
),
requested_model=model_group,
hashed_api_key=standard_logging_payload["metadata"][
"user_api_key_hash"
],
hashed_api_key=hashed_api_key,
api_key_alias=standard_logging_payload["metadata"][
"user_api_key_alias"
],
@ -1441,6 +1605,14 @@ class PrometheusLogger(CustomLogger):
if standard_logging_payload is None:
return
# Skip recording metrics for invalid API key requests
if self._should_skip_metrics_for_invalid_key(
kwargs=request_kwargs,
enum_values=enum_values,
standard_logging_payload=standard_logging_payload,
):
return
api_base = standard_logging_payload["api_base"]
_litellm_params = request_kwargs.get("litellm_params", {}) or {}
_metadata = _litellm_params.get("metadata", {})

View file

@ -3630,6 +3630,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
otel_config = OpenTelemetryConfig(
exporter=arize_config.protocol,
endpoint=arize_config.endpoint,
service_name=arize_config.project_name,
)
os.environ[

View file

@ -2137,6 +2137,14 @@ def anthropic_messages_pt( # noqa: PLR0915
assistant_content.append(
cast(AnthropicMessagesTextParam, _cached_message)
)
# handle server_tool_use blocks (tool search, web search, etc.)
# Pass through as-is since these are Anthropic-native content types
elif m.get("type", "") == "server_tool_use":
assistant_content.append(m) # type: ignore
# handle tool_search_tool_result blocks
# Pass through as-is since these are Anthropic-native content types
elif m.get("type", "") == "tool_search_tool_result":
assistant_content.append(m) # type: ignore
elif (
"content" in assistant_content_block
and isinstance(assistant_content_block["content"], str)
@ -3168,6 +3176,11 @@ def _convert_to_bedrock_tool_call_invoke(
id = tool["id"]
name = tool["function"].get("name", "")
arguments = tool["function"].get("arguments", "")
arguments_dict = json.loads(arguments) if arguments else {}
# Ensure arguments_dict is always a dict (Bedrock requires toolUse.input to be an object)
# When some providers return arguments: '""' (JSON-encoded empty string), json.loads returns ""
if not isinstance(arguments_dict, dict):
arguments_dict = {}
if not arguments or not arguments.strip():
arguments_dict = {}
else:

View file

@ -990,6 +990,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
def create_azure_base_url(
self, azure_client_params: dict, model: Optional[str]
) -> str:
from litellm.llms.azure_ai.image_generation import (
AzureFoundryFluxImageGenerationConfig,
)
api_base: str = azure_client_params.get(
"azure_endpoint", ""
) # "https://example-endpoint.openai.azure.com"
@ -999,6 +1003,15 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
if model is None:
model = ""
# Handle FLUX 2 models on Azure AI which use a different URL pattern
# e.g., /providers/blackforestlabs/v1/flux-2-pro instead of /openai/deployments/{model}/images/generations
if AzureFoundryFluxImageGenerationConfig.is_flux2_model(model):
return AzureFoundryFluxImageGenerationConfig.get_flux2_image_generation_url(
api_base=api_base,
model=model,
api_version=api_version,
)
if "/openai/deployments/" in api_base:
base_url_with_deployment = api_base
else:

View file

@ -1,15 +1,28 @@
from litellm.llms.azure_ai.image_generation.flux_transformation import (
AzureFoundryFluxImageGenerationConfig,
)
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from .flux2_transformation import AzureFoundryFlux2ImageEditConfig
from .transformation import AzureFoundryFluxImageEditConfig
__all__ = ["AzureFoundryFluxImageEditConfig"]
__all__ = ["AzureFoundryFluxImageEditConfig", "AzureFoundryFlux2ImageEditConfig"]
def get_azure_ai_image_edit_config(model: str) -> BaseImageEditConfig:
model = model.lower()
model = model.replace("-", "")
model = model.replace("_", "")
if model == "" or "flux" in model: # empty model is flux
"""
Get the appropriate image edit config for an Azure AI model.
- FLUX 2 models use JSON with base64 image
- FLUX 1 models use multipart/form-data
"""
# Check if it's a FLUX 2 model
if AzureFoundryFluxImageGenerationConfig.is_flux2_model(model):
return AzureFoundryFlux2ImageEditConfig()
# Default to FLUX 1 config for other FLUX models
model_normalized = model.lower().replace("-", "").replace("_", "")
if model_normalized == "" or "flux" in model_normalized:
return AzureFoundryFluxImageEditConfig()
else:
raise ValueError(f"Model {model} is not supported for Azure AI image editing.")
raise ValueError(f"Model {model} is not supported for Azure AI image editing.")

View file

@ -0,0 +1,167 @@
import base64
from io import BufferedReader
from typing import Any, Dict, Optional, Tuple
from httpx._types import RequestFiles
import litellm
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
from litellm.llms.azure_ai.image_generation.flux_transformation import (
AzureFoundryFluxImageGenerationConfig,
)
from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.llms.openai import FileTypes
from litellm.types.router import GenericLiteLLMParams
class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
"""
Azure AI Foundry FLUX 2 image edit config
Supports FLUX 2 models (e.g., flux.2-pro) for image editing.
Uses the same /providers/blackforestlabs/v1/flux-2-pro endpoint as image generation,
with the image passed as base64 in JSON body.
"""
def get_supported_openai_params(self, model: str) -> list:
"""
FLUX 2 supports a subset of OpenAI image edit params
"""
return [
"prompt",
"image",
"model",
"n",
"size",
]
def map_openai_params(
self,
image_edit_optional_params: ImageEditOptionalRequestParams,
model: str,
drop_params: bool,
) -> Dict:
"""
Map OpenAI params to FLUX 2 params.
FLUX 2 uses the same param names as OpenAI for supported params.
"""
mapped_params: Dict[str, Any] = {}
supported_params = self.get_supported_openai_params(model)
for key, value in dict(image_edit_optional_params).items():
if key in supported_params and value is not None:
mapped_params[key] = value
return mapped_params
def use_multipart_form_data(self) -> bool:
"""FLUX 2 uses JSON requests, not multipart/form-data."""
return False
def validate_environment(
self,
headers: dict,
model: str,
api_key: Optional[str] = None,
) -> dict:
"""
Validate Azure AI Foundry environment and set up authentication
"""
api_key = AzureFoundryModelInfo.get_api_key(api_key)
if not api_key:
raise ValueError(
f"Azure AI API key is required for model {model}. Set AZURE_AI_API_KEY environment variable or pass api_key parameter."
)
headers.update(
{
"Api-Key": api_key,
"Content-Type": "application/json",
}
)
return headers
def transform_image_edit_request(
self,
model: str,
prompt: str,
image: FileTypes,
image_edit_optional_request_params: Dict,
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> Tuple[Dict, RequestFiles]:
"""
Transform image edit request for FLUX 2.
FLUX 2 uses the same endpoint for generation and editing,
with the image passed as base64 in the JSON body.
"""
image_b64 = self._convert_image_to_base64(image)
# Build request body with required params
request_body: Dict[str, Any] = {
"prompt": prompt,
"image": image_b64,
"model": model,
}
# Add mapped optional params (already filtered by map_openai_params)
request_body.update(image_edit_optional_request_params)
# Return JSON body and empty files list (FLUX 2 doesn't use multipart)
return request_body, []
def _convert_image_to_base64(self, image: Any) -> str:
"""Convert image file to base64 string"""
# Handle list of images (take first one)
if isinstance(image, list):
if len(image) == 0:
raise ValueError("Empty image list provided")
image = image[0]
if isinstance(image, BufferedReader):
image_bytes = image.read()
image.seek(0) # Reset file pointer for potential reuse
elif isinstance(image, bytes):
image_bytes = image
elif hasattr(image, "read"):
image_bytes = image.read() # type: ignore
else:
raise ValueError(f"Unsupported image type: {type(image)}")
return base64.b64encode(image_bytes).decode("utf-8")
def get_complete_url(
self,
model: str,
api_base: Optional[str],
litellm_params: dict,
) -> str:
"""
Constructs a complete URL for Azure AI Foundry FLUX 2 image edits.
Uses the same /providers/blackforestlabs/v1/flux-2-pro endpoint as image generation.
"""
api_base = AzureFoundryModelInfo.get_api_base(api_base)
if api_base is None:
raise ValueError(
"Azure AI API base is required. Set AZURE_AI_API_BASE environment variable or pass api_base parameter."
)
api_version = (
litellm_params.get("api_version")
or litellm.api_version
or get_secret_str("AZURE_AI_API_VERSION")
or "preview"
)
return AzureFoundryFluxImageGenerationConfig.get_flux2_image_generation_url(
api_base=api_base,
model=model,
api_version=api_version,
)

View file

@ -71,9 +71,11 @@ class AzureFoundryFluxImageEditConfig(OpenAIImageEditConfig):
"Azure AI API base is required. Set AZURE_AI_API_BASE environment variable or pass api_base parameter."
)
api_version = (litellm_params.get("api_version") or litellm.api_version
or get_secret_str("AZURE_AI_API_VERSION")
)
api_version = (
litellm_params.get("api_version")
or litellm.api_version
or get_secret_str("AZURE_AI_API_VERSION")
)
if api_version is None:
# API version is mandatory for Azure AI Foundry
raise ValueError(

View file

@ -1,3 +1,5 @@
from typing import Optional
from litellm.llms.openai.image_generation import GPTImageGenerationConfig
@ -11,4 +13,56 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
From our test suite - following GPTImageGenerationConfig is working for this model
"""
pass
@staticmethod
def get_flux2_image_generation_url(
api_base: Optional[str],
model: str,
api_version: Optional[str],
) -> str:
"""
Constructs the complete URL for Azure AI FLUX 2 image generation.
FLUX 2 models on Azure AI use a different URL pattern than standard Azure OpenAI:
- Standard: /openai/deployments/{model}/images/generations
- FLUX 2: /providers/blackforestlabs/v1/flux-2-pro
Args:
api_base: Base URL (e.g., https://litellm-ci-cd-prod.services.ai.azure.com)
model: Model name (e.g., flux.2-pro)
api_version: API version (e.g., preview)
Returns:
Complete URL for the FLUX 2 image generation endpoint
"""
if api_base is None:
raise ValueError(
"api_base is required for Azure AI FLUX 2 image generation"
)
api_base = api_base.rstrip("/")
api_version = api_version or "preview"
# If the api_base already contains /providers/, it's already a complete path
if "/providers/" in api_base:
if "?" in api_base:
return api_base
return f"{api_base}?api-version={api_version}"
# Construct the FLUX 2 provider path
# Model name flux.2-pro maps to endpoint flux-2-pro
return f"{api_base}/providers/blackforestlabs/v1/flux-2-pro?api-version={api_version}"
@staticmethod
def is_flux2_model(model: str) -> bool:
"""
Check if the model is an Azure AI FLUX 2 model.
Args:
model: Model name (e.g., flux.2-pro, azure_ai/flux.2-pro)
Returns:
True if the model is a FLUX 2 model
"""
model_lower = model.lower().replace(".", "-").replace("_", "-")
return "flux-2" in model_lower or "flux2" in model_lower

View file

@ -88,6 +88,34 @@ def _build_query_params(
return params
def _prepare_multipart_file_upload(
file: Any,
headers: Dict[str, Any],
) -> tuple:
"""
Prepare file and headers for multipart upload.
Returns:
Tuple of (files_dict, headers_without_content_type)
"""
from litellm.litellm_core_utils.prompt_templates.common_utils import (
extract_file_data,
)
extracted = extract_file_data(file)
filename = extracted.get("filename") or "file"
content = extracted.get("content") or b""
content_type = extracted.get("content_type") or "application/octet-stream"
files = {"file": (filename, content, content_type)}
# Remove content-type header - httpx will set it automatically for multipart
headers_copy = headers.copy()
headers_copy.pop("content-type", None)
headers_copy.pop("Content-Type", None)
return files, headers_copy
class GenericContainerHandler:
"""
Generic handler for container file API endpoints.
@ -210,6 +238,7 @@ class GenericContainerHandler:
# Make request
method = endpoint_config["method"].upper()
returns_binary = endpoint_config.get("returns_binary", False)
is_multipart = endpoint_config.get("is_multipart", False)
try:
if method == "GET":
@ -217,7 +246,11 @@ class GenericContainerHandler:
elif method == "DELETE":
response = http_client.delete(url=url, headers=headers, params=query_params)
elif method == "POST":
response = http_client.post(url=url, headers=headers, params=query_params)
if is_multipart and "file" in kwargs:
files, headers = _prepare_multipart_file_upload(kwargs["file"], headers)
response = http_client.post(url=url, headers=headers, params=query_params, files=files)
else:
response = http_client.post(url=url, headers=headers, params=query_params)
else:
raise ValueError(f"Unsupported HTTP method: {method}")
@ -307,6 +340,7 @@ class GenericContainerHandler:
# Make request
method = endpoint_config["method"].upper()
returns_binary = endpoint_config.get("returns_binary", False)
is_multipart = endpoint_config.get("is_multipart", False)
try:
if method == "GET":
@ -314,7 +348,11 @@ class GenericContainerHandler:
elif method == "DELETE":
response = await http_client.delete(url=url, headers=headers, params=query_params)
elif method == "POST":
response = await http_client.post(url=url, headers=headers, params=query_params)
if is_multipart and "file" in kwargs:
files, headers = _prepare_multipart_file_upload(kwargs["file"], headers)
response = await http_client.post(url=url, headers=headers, params=query_params, files=files)
else:
response = await http_client.post(url=url, headers=headers, params=query_params)
else:
raise ValueError(f"Unsupported HTTP method: {method}")

View file

@ -771,9 +771,9 @@ class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator):
return ModelResponseStream(
id=chunk["id"],
object="chat.completion.chunk",
created=chunk["created"],
model=chunk["model"],
choices=chunk["choices"],
created=chunk.get("created"),
model=chunk.get("model"),
choices=chunk.get("choices", []),
)
except Exception as e:
raise e

View file

@ -203,12 +203,8 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig):
litellm_params: dict,
headers: dict,
) -> dict:
supported_params = self.get_supported_openai_params(model)
# Include extra params that passed validation (e.g., thinking_config for Gemini models via allowed_openai_params)
extra_params = [k for k in optional_params if k not in supported_params and k not in {"tools", "model_version"}]
supported_params = supported_params + extra_params
model_params = {
k: v for k, v in optional_params.items() if k in supported_params
k: v for k, v in optional_params.items() if k not in {"tools", "model_version", "deployment_url"}
}
model_version = optional_params.pop("model_version", "latest")

View file

@ -40,6 +40,7 @@ class PartnerModelPrefixes(str, Enum):
GPT_OSS_PREFIX = "openai/gpt-oss-"
MINIMAX_PREFIX = "minimaxai/"
MOONSHOT_PREFIX = "moonshotai/"
ZAI_PREFIX = "zai-org/"
class VertexAIPartnerModels(VertexBase):
@ -66,6 +67,7 @@ class VertexAIPartnerModels(VertexBase):
or model.startswith(PartnerModelPrefixes.GPT_OSS_PREFIX)
or model.startswith(PartnerModelPrefixes.MINIMAX_PREFIX)
or model.startswith(PartnerModelPrefixes.MOONSHOT_PREFIX)
or model.startswith(PartnerModelPrefixes.ZAI_PREFIX)
):
return True
return False
@ -79,6 +81,7 @@ class VertexAIPartnerModels(VertexBase):
PartnerModelPrefixes.GPT_OSS_PREFIX,
PartnerModelPrefixes.MINIMAX_PREFIX,
PartnerModelPrefixes.MOONSHOT_PREFIX,
PartnerModelPrefixes.ZAI_PREFIX,
]
if any(provider in model for provider in OPENAI_LIKE_VERTEX_PROVIDERS):
return True

View file

@ -405,7 +405,23 @@
"supports_video_input": true,
"supports_vision": true
},
"amazon.nova-2-multimodal-embeddings-v1:0": {
"litellm_provider": "bedrock",
"max_input_tokens": 8172,
"max_tokens": 8172,
"mode": "embedding",
"input_cost_per_token": 1.35e-7,
"input_cost_per_image": 6e-5,
"input_cost_per_video_per_second": 0.0007,
"input_cost_per_audio_per_second": 0.00014,
"output_cost_per_token": 0.0,
"output_vector_size": 3072,
"source": "https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/model-catalog/serverless/amazon.nova-2-multimodal-embeddings-v1:0",
"supports_embedding_image_input": true,
"supports_image_input": true,
"supports_video_input": true,
"supports_audio_input": true
},
"amazon.nova-micro-v1:0": {
"input_cost_per_token": 3.5e-08,
"litellm_provider": "bedrock_converse",
@ -4893,6 +4909,15 @@
"/v1/images/generations"
]
},
"azure_ai/flux.2-pro": {
"litellm_provider": "azure_ai",
"mode": "image_generation",
"output_cost_per_image": 0.04,
"source": "https://ai.azure.com/explore/models/flux.2-pro/version/1/registry/azureml-blackforestlabs",
"supported_endpoints": [
"/v1/images/generations"
]
},
"azure_ai/Llama-3.2-11B-Vision-Instruct": {
"input_cost_per_token": 3.7e-07,
"litellm_provider": "azure_ai",
@ -28320,6 +28345,19 @@
"supports_tool_choice": true,
"supports_web_search": true
},
"vertex_ai/zai-org/glm-4.7-maas": {
"input_cost_per_token": 3e-07,
"litellm_provider": "vertex_ai-zai_models",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"vertex_ai/mistral-medium-3": {
"input_cost_per_token": 4e-07,
"litellm_provider": "vertex_ai-mistral_models",
@ -32152,6 +32190,181 @@
"output_cost_per_token": 2e-07,
"litellm_provider": "fireworks_ai",
"mode": "chat"
},
"llamagate/llama-3.1-8b": {
"max_tokens": 8192,
"max_input_tokens": 131072,
"max_output_tokens": 8192,
"input_cost_per_token": 3e-08,
"output_cost_per_token": 5e-08,
"litellm_provider": "llamagate",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": true
},
"llamagate/llama-3.2-3b": {
"max_tokens": 8192,
"max_input_tokens": 131072,
"max_output_tokens": 8192,
"input_cost_per_token": 4e-08,
"output_cost_per_token": 8e-08,
"litellm_provider": "llamagate",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": true
},
"llamagate/mistral-7b-v0.3": {
"max_tokens": 8192,
"max_input_tokens": 32768,
"max_output_tokens": 8192,
"input_cost_per_token": 1e-07,
"output_cost_per_token": 1.5e-07,
"litellm_provider": "llamagate",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": true
},
"llamagate/qwen3-8b": {
"max_tokens": 8192,
"max_input_tokens": 32768,
"max_output_tokens": 8192,
"input_cost_per_token": 4e-08,
"output_cost_per_token": 1.4e-07,
"litellm_provider": "llamagate",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": true
},
"llamagate/dolphin3-8b": {
"max_tokens": 8192,
"max_input_tokens": 128000,
"max_output_tokens": 8192,
"input_cost_per_token": 8e-08,
"output_cost_per_token": 1.5e-07,
"litellm_provider": "llamagate",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": true
},
"llamagate/deepseek-r1-8b": {
"max_tokens": 16384,
"max_input_tokens": 65536,
"max_output_tokens": 16384,
"input_cost_per_token": 1e-07,
"output_cost_per_token": 2e-07,
"litellm_provider": "llamagate",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": true,
"supports_reasoning": true
},
"llamagate/deepseek-r1-7b-qwen": {
"max_tokens": 16384,
"max_input_tokens": 131072,
"max_output_tokens": 16384,
"input_cost_per_token": 8e-08,
"output_cost_per_token": 1.5e-07,
"litellm_provider": "llamagate",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": true,
"supports_reasoning": true
},
"llamagate/openthinker-7b": {
"max_tokens": 8192,
"max_input_tokens": 32768,
"max_output_tokens": 8192,
"input_cost_per_token": 8e-08,
"output_cost_per_token": 1.5e-07,
"litellm_provider": "llamagate",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": true,
"supports_reasoning": true
},
"llamagate/qwen2.5-coder-7b": {
"max_tokens": 8192,
"max_input_tokens": 32768,
"max_output_tokens": 8192,
"input_cost_per_token": 6e-08,
"output_cost_per_token": 1.2e-07,
"litellm_provider": "llamagate",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": true
},
"llamagate/deepseek-coder-6.7b": {
"max_tokens": 4096,
"max_input_tokens": 16384,
"max_output_tokens": 4096,
"input_cost_per_token": 6e-08,
"output_cost_per_token": 1.2e-07,
"litellm_provider": "llamagate",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": true
},
"llamagate/codellama-7b": {
"max_tokens": 4096,
"max_input_tokens": 16384,
"max_output_tokens": 4096,
"input_cost_per_token": 6e-08,
"output_cost_per_token": 1.2e-07,
"litellm_provider": "llamagate",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": true
},
"llamagate/qwen3-vl-8b": {
"max_tokens": 8192,
"max_input_tokens": 32768,
"max_output_tokens": 8192,
"input_cost_per_token": 1.5e-07,
"output_cost_per_token": 5.5e-07,
"litellm_provider": "llamagate",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": true,
"supports_vision": true
},
"llamagate/llava-7b": {
"max_tokens": 2048,
"max_input_tokens": 4096,
"max_output_tokens": 2048,
"input_cost_per_token": 1e-07,
"output_cost_per_token": 2e-07,
"litellm_provider": "llamagate",
"mode": "chat",
"supports_response_schema": true,
"supports_vision": true
},
"llamagate/gemma3-4b": {
"max_tokens": 8192,
"max_input_tokens": 128000,
"max_output_tokens": 8192,
"input_cost_per_token": 3e-08,
"output_cost_per_token": 8e-08,
"litellm_provider": "llamagate",
"mode": "chat",
"supports_function_calling": true,
"supports_response_schema": true,
"supports_vision": true
},
"llamagate/nomic-embed-text": {
"max_tokens": 8192,
"max_input_tokens": 8192,
"input_cost_per_token": 2e-08,
"output_cost_per_token": 0,
"litellm_provider": "llamagate",
"mode": "embedding"
},
"llamagate/qwen3-embedding-8b": {
"max_tokens": 40960,
"max_input_tokens": 40960,
"input_cost_per_token": 2e-08,
"output_cost_per_token": 0,
"litellm_provider": "llamagate",
"mode": "embedding"
}
}

View file

@ -863,6 +863,7 @@ class KeyRequestBase(GenerateRequestBase):
tpm_limit_type: Optional[
Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"]
] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm
router_settings: Optional[UpdateRouterConfig] = None
class LiteLLMKeyType(str, enum.Enum):
@ -918,6 +919,7 @@ class GenerateKeyResponse(KeyRequestBase):
"config",
"permissions",
"model_max_budget",
"router_settings",
]
for field in dict_fields:
value = values.get(field)
@ -1460,6 +1462,7 @@ class TeamBase(LiteLLMPydanticObjectBase):
models: list = []
blocked: bool = False
router_settings: Optional[dict] = None
class NewTeamRequest(TeamBase):
@ -1542,6 +1545,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
model_rpm_limit: Optional[Dict[str, int]] = None
model_tpm_limit: Optional[Dict[str, int]] = None
allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
router_settings: Optional[dict] = None
class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase):
@ -1684,6 +1688,7 @@ class LiteLLM_TeamTable(TeamBase):
"permissions",
"model_max_budget",
"model_aliases",
"router_settings",
]
if isinstance(values, BaseModel):
@ -3735,6 +3740,7 @@ class BaseDailySpendTransaction(TypedDict):
model_group: Optional[str]
mcp_namespaced_tool_name: Optional[str]
custom_llm_provider: Optional[str]
endpoint: Optional[str]
# token count metrics
prompt_tokens: int

View file

@ -0,0 +1,214 @@
# Performance Utilities Documentation
This module provides performance monitoring and profiling functionality for LiteLLM proxy server using `cProfile` and `line_profiler`.
## Table of Contents
- [Line Profiler Usage](#line-profiler-usage)
- [Example 1: Wrapping a function directly](#example-1-wrapping-a-function-directly)
- [Example 2: Wrapping a module function dynamically](#example-2-wrapping-a-module-function-dynamically)
- [Example 3: Manual stats collection](#example-3-manual-stats-collection)
- [Example 4: Analyzing the profile output](#example-4-analyzing-the-profile-output)
- [Example 5: Using in a decorator pattern](#example-5-using-in-a-decorator-pattern)
- [cProfile Usage](#cprofile-usage)
- [Installation](#installation)
- [Notes](#notes)
## Line Profiler Usage
### Example 1: Wrapping a function directly
This is how it's used in `litellm/utils.py` to profile `wrapper_async`:
```python
from litellm.proxy.common_utils.performance_utils import (
register_shutdown_handler,
wrap_function_directly,
)
def client(original_function):
@wraps(original_function)
async def wrapper_async(*args, **kwargs):
# ... function implementation ...
pass
# Wrap the function with line_profiler
wrapper_async = wrap_function_directly(wrapper_async)
# Register shutdown handler to collect stats on server shutdown
register_shutdown_handler(output_file="wrapper_async_line_profile.lprof")
return wrapper_async
```
### Example 2: Wrapping a module function dynamically
```python
import my_module
from litellm.proxy.common_utils.performance_utils import (
wrap_function_with_line_profiler,
register_shutdown_handler,
)
# Wrap a function in a module
wrap_function_with_line_profiler(my_module, "expensive_function")
# Register shutdown handler
register_shutdown_handler(output_file="my_profile.lprof")
# Now all calls to my_module.expensive_function will be profiled
my_module.expensive_function()
```
### Example 3: Manual stats collection
```python
from litellm.proxy.common_utils.performance_utils import (
wrap_function_directly,
collect_line_profiler_stats,
)
def my_function():
# ... implementation ...
pass
# Wrap the function
my_function = wrap_function_directly(my_function)
# Run your code
my_function()
# Collect stats manually (instead of waiting for shutdown)
collect_line_profiler_stats(output_file="manual_profile.lprof")
```
### Example 4: Analyzing the profile output
After running your code, analyze the `.lprof` file:
```bash
# View the profile
python -m line_profiler wrapper_async_line_profile.lprof
# Save to text file
python -m line_profiler wrapper_async_line_profile.lprof > profile_report.txt
```
The output shows:
- **Line #**: Line number in the source file
- **Hits**: Number of times the line was executed
- **Time**: Total time spent on that line (in microseconds)
- **Per Hit**: Average time per execution
- **% Time**: Percentage of total function time
- **Line Contents**: The actual source code
Example output:
```
Timer unit: 1e-06 s
Total time: 3.73697 s
File: litellm/utils.py
Function: client.<locals>.wrapper_async at line 1657
Line # Hits Time Per Hit % Time Line Contents
==============================================================
1657 @wraps(original_function)
1658 async def wrapper_async(*args, **kwargs):
1659 2005 7577.1 3.8 0.2 print_args_passed_to_litellm(...)
1763 2005 1351909.0 674.3 36.2 result = await original_function(*args, **kwargs)
1846 4010 1543688.1 385.0 41.3 update_response_metadata(...)
```
### Example 5: Using in a decorator pattern
```python
from litellm.proxy.common_utils.performance_utils import (
wrap_function_directly,
register_shutdown_handler,
)
def profile_decorator(func):
# Wrap the function
profiled_func = wrap_function_directly(func)
# Register shutdown handler (only once)
if not hasattr(profile_decorator, '_registered'):
register_shutdown_handler(output_file="decorated_functions.lprof")
profile_decorator._registered = True
return profiled_func
@profile_decorator
async def my_async_function():
# This function will be profiled
pass
```
## cProfile Usage
### Example: Using the profile_endpoint decorator
```python
from litellm.proxy.common_utils.performance_utils import profile_endpoint
@profile_endpoint(sampling_rate=0.1) # Profile 10% of requests
async def my_endpoint():
# ... implementation ...
pass
```
The `sampling_rate` parameter controls what percentage of requests are profiled:
- `1.0`: Profile all requests (100%)
- `0.1`: Profile 1 in 10 requests (10%)
- `0.0`: Profile no requests (0%)
## Installation
`line_profiler` must be installed to use the line profiling functionality:
```bash
pip install line_profiler
```
On Windows with Python 3.14+, you may need to install Microsoft Visual C++ Build Tools to compile `line_profiler` from source.
## Notes
- The profiler aggregates stats by source code location, so multiple instances of the same function (e.g., closures) will be profiled together
- Stats are automatically collected on server shutdown via `atexit` handler when using `register_shutdown_handler()`
- You can also manually collect stats using `collect_line_profiler_stats()`
- The line profiler will fail with an `ImportError` if `line_profiler` is not installed (as configured in `litellm/utils.py`)
## API Reference
### `wrap_function_directly(func: Callable) -> Callable`
Wrap a function directly with line_profiler. This is the recommended way to profile functions, especially closures or functions created dynamically.
**Raises:**
- `ImportError`: If line_profiler is not available
- `RuntimeError`: If line_profiler cannot be enabled or function cannot be wrapped
### `wrap_function_with_line_profiler(module: Any, function_name: str) -> bool`
Dynamically wrap a function in a module with line_profiler.
**Returns:** `True` if wrapping was successful, `False` otherwise
### `collect_line_profiler_stats(output_file: Optional[str] = None) -> None`
Collect and save line_profiler statistics. If `output_file` is provided, saves to file. Otherwise, prints to stdout.
### `register_shutdown_handler(output_file: Optional[str] = None) -> None`
Register an `atexit` handler that will automatically save profiling statistics when the Python process exits. Safe to call multiple times (only registers once).
**Default output file:** `line_profile_stats.lprof` if not specified
### `profile_endpoint(sampling_rate: float = 1.0)`
Decorator to sample endpoint hits and save to a profile file using cProfile.
**Args:**
- `sampling_rate`: Rate of requests to profile (0.0 to 1.0)

View file

@ -2,14 +2,19 @@
Performance utilities for LiteLLM proxy server.
This module provides performance monitoring and profiling functionality for endpoint
performance analysis using cProfile with configurable sampling rates.
performance analysis using cProfile with configurable sampling rates, and line_profiler
for line-by-line profiling.
See performance_utils.md for detailed usage examples and documentation.
"""
import asyncio
import atexit
import cProfile
import functools
import threading
from pathlib import Path as PathLib
from typing import Any, Callable, Optional
from litellm._logging import verbose_proxy_logger
@ -20,6 +25,11 @@ _last_profile_file_path = None
_sample_counter = 0
_sample_counter_lock = threading.Lock()
# Global line_profiler state
_line_profiler: Optional[Any] = None
_line_profiler_lock = threading.Lock()
_wrapped_functions: dict[str, Callable] = {} # Store original functions
def _should_sample(profile_sampling_rate: float) -> bool:
"""Determine if current request should be sampled based on sampling rate."""
@ -123,3 +133,156 @@ def profile_endpoint(sampling_rate: float = 1.0):
raise
return sync_wrapper
return decorator
def enable_line_profiler() -> None:
"""Enable line_profiler for dynamic function wrapping.
Raises:
ImportError: If line_profiler is not available
"""
global _line_profiler
from line_profiler import LineProfiler # Will raise ImportError if not available
with _line_profiler_lock:
if _line_profiler is None:
_line_profiler = LineProfiler()
verbose_proxy_logger.info("Line profiler enabled")
def wrap_function_with_line_profiler(module: Any, function_name: str) -> bool:
"""Dynamically wrap a function with line_profiler.
Args:
module: The module containing the function
function_name: Name of the function to wrap
Returns:
True if wrapping was successful, False otherwise
"""
try:
enable_line_profiler() # May raise ImportError if not available
except ImportError:
return False
if _line_profiler is None:
return False
try:
original_function = getattr(module, function_name, None)
if original_function is None:
verbose_proxy_logger.warning(
f"Function {function_name} not found in module {module.__name__}"
)
return False
# Store original function if not already wrapped
if function_name not in _wrapped_functions:
_wrapped_functions[function_name] = original_function
# Wrap with line_profiler
profiled_function = _line_profiler(original_function)
setattr(module, function_name, profiled_function)
verbose_proxy_logger.info(
f"Wrapped {module.__name__}.{function_name} with line_profiler"
)
return True
except Exception as e:
verbose_proxy_logger.error(
f"Error wrapping {function_name} with line_profiler: {e}"
)
return False
def wrap_function_directly(func: Callable) -> Callable:
"""Wrap a function directly with line_profiler.
This is the recommended way to profile functions, especially closures or
functions created dynamically (like wrapper_async in litellm/utils.py).
Args:
func: The function to wrap
Returns:
The wrapped function that will be profiled when called
Raises:
ImportError: If line_profiler is not available
RuntimeError: If line_profiler cannot be enabled or function cannot be wrapped
"""
import warnings
enable_line_profiler() # Will raise ImportError if not available
if _line_profiler is None:
raise RuntimeError("Line profiler was not initialized")
# Suppress warnings about __wrapped__ - we intentionally want to profile the wrapper
with warnings.catch_warnings():
warnings.filterwarnings('ignore', message='.*__wrapped__.*', category=UserWarning)
# Add function to line_profiler and wrap it
_line_profiler.add_function(func)
profiled_function = _line_profiler(func)
verbose_proxy_logger.info(
f"Wrapped function {func.__name__} with line_profiler"
)
return profiled_function
def collect_line_profiler_stats(output_file: Optional[str] = None) -> None:
"""Collect and save line_profiler statistics.
This can be called manually to collect stats at any time, or it's
automatically called on shutdown if register_shutdown_handler() was used.
Args:
output_file: Optional path to save stats. If None, prints to stdout.
"""
global _line_profiler
with _line_profiler_lock:
if _line_profiler is None:
verbose_proxy_logger.debug("Line profiler not enabled, nothing to collect")
return
try:
if output_file:
# Save to file
output_path = PathLib(output_file)
_line_profiler.dump_stats(str(output_path))
verbose_proxy_logger.info(
f"Line profiler stats saved to {output_path}"
)
else:
# Print to stdout
from io import StringIO
stream = StringIO()
_line_profiler.print_stats(stream=stream)
stats_output = stream.getvalue()
verbose_proxy_logger.info("Line profiler stats:\n" + stats_output)
except Exception as e:
verbose_proxy_logger.error(f"Error collecting line profiler stats: {e}")
def register_shutdown_handler(output_file: Optional[str] = None) -> None:
"""Register a shutdown handler to collect line_profiler stats.
This registers an atexit handler that will automatically save profiling
statistics when the Python process exits. Safe to call multiple times
(only registers once).
Args:
output_file: Optional path to save stats on shutdown.
Defaults to 'line_profile_stats.lprof'
"""
if output_file is None:
output_file = "line_profile_stats.lprof"
def shutdown_handler():
collect_line_profiler_stats(output_file=output_file)
atexit.register(shutdown_handler)
verbose_proxy_logger.debug(f"Registered line_profiler shutdown handler for {output_file}")

View file

@ -43,7 +43,7 @@ def _get_container_provider_config(custom_llm_provider: str):
raise ValueError(f"Container API not supported for provider: {custom_llm_provider}")
def _create_handler_for_path_params(path_params: List[str], route_type: str, returns_binary: bool = False):
def _create_handler_for_path_params(path_params: List[str], route_type: str, returns_binary: bool = False, is_multipart: bool = False):
"""
Dynamically create a handler with the correct path parameter signature.
"""
@ -63,6 +63,23 @@ def _create_handler_for_path_params(path_params: List[str], route_type: str, ret
)
return handler_binary_content
# For multipart file upload endpoints
if is_multipart:
async def handler_multipart_upload(
request: Request,
container_id: str,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
return await _process_multipart_upload_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type=route_type,
container_id=container_id,
)
return handler_multipart_upload
# Create handlers for different path parameter combinations
if path_params == ["container_id"]:
async def handler_container_id(
@ -193,6 +210,83 @@ async def _process_binary_request(
raise e
async def _process_multipart_upload_request(
request: Request,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth,
route_type: str,
container_id: str,
):
"""Process multipart file upload requests."""
from litellm.proxy.common_utils.http_parsing_utils import (
convert_upload_files_to_file_data,
get_form_data,
)
from litellm.proxy.proxy_server import (
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
select_data_generator,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
# Parse multipart form data and convert files
form_data = await get_form_data(request)
data = await convert_upload_files_to_file_data(form_data)
if "file" not in data:
from fastapi import HTTPException
raise HTTPException(status_code=400, detail="Missing required 'file' field")
# convert_upload_files_to_file_data returns list of tuples, extract single file
file_list = data["file"]
if isinstance(file_list, list) and len(file_list) > 0:
data["file"] = file_list[0]
data["container_id"] = container_id
custom_llm_provider = (
get_custom_llm_provider_from_request_headers(request=request)
or get_custom_llm_provider_from_request_query(request=request)
or "openai"
)
data["custom_llm_provider"] = custom_llm_provider
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type=route_type, # type: ignore[arg-type]
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=None,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
async def _process_request(
request: Request,
fastapi_response: Response,
@ -272,9 +366,10 @@ def register_container_file_endpoints(router: APIRouter) -> None:
path_params = endpoint_config.get("path_params", [])
route_type = endpoint_config["async_name"]
returns_binary = endpoint_config.get("returns_binary", False)
is_multipart = endpoint_config.get("is_multipart", False)
# Create handler with correct signature for path params
handler = _create_handler_for_path_params(path_params, route_type, returns_binary)
handler = _create_handler_for_path_params(path_params, route_type, returns_binary, is_multipart)
# Register routes
route_method = getattr(router, method)

View file

@ -42,6 +42,7 @@ from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient, ProxyLogging
@ -1205,6 +1206,7 @@ class DBSpendUpdateWriter:
"mcp_namespaced_tool_name"
)
or "",
"endpoint": transaction.get("endpoint") or "",
}
}
@ -1225,6 +1227,7 @@ class DBSpendUpdateWriter:
"custom_llm_provider": transaction.get(
"custom_llm_provider"
),
"endpoint": transaction.get("endpoint"),
"prompt_tokens": transaction["prompt_tokens"],
"completion_tokens": transaction["completion_tokens"],
"spend": transaction["spend"],
@ -1287,6 +1290,9 @@ class DBSpendUpdateWriter:
if entity_type == "tag" and "request_id" in transaction:
update_data["request_id"] = transaction.get("request_id")
# Add endpoint to update_data so existing rows get their endpoint field updated
update_data["endpoint"] = transaction.get("endpoint") or ""
table.upsert(
where=where_clause,
data={
@ -1347,7 +1353,7 @@ class DBSpendUpdateWriter:
entity_type="user",
entity_id_field="user_id",
table_name="litellm_dailyuserspend",
unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name",
unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint",
)
@staticmethod
@ -1368,7 +1374,7 @@ class DBSpendUpdateWriter:
entity_type="team",
entity_id_field="team_id",
table_name="litellm_dailyteamspend",
unique_constraint_name="team_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name",
unique_constraint_name="team_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint",
)
@staticmethod
@ -1389,7 +1395,7 @@ class DBSpendUpdateWriter:
entity_type="org",
entity_id_field="organization_id",
table_name="litellm_dailyorganizationspend",
unique_constraint_name="organization_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name",
unique_constraint_name="organization_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint",
)
@staticmethod
@ -1410,7 +1416,7 @@ class DBSpendUpdateWriter:
entity_type="end_user",
entity_id_field="end_user_id",
table_name="litellm_dailyenduserspend",
unique_constraint_name="end_user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name",
unique_constraint_name="end_user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint",
)
@staticmethod
@ -1431,7 +1437,7 @@ class DBSpendUpdateWriter:
entity_type="agent",
entity_id_field="agent_id",
table_name="litellm_dailyagentspend",
unique_constraint_name="agent_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name",
unique_constraint_name="agent_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint",
)
@staticmethod
@ -1452,7 +1458,7 @@ class DBSpendUpdateWriter:
entity_type="tag",
entity_id_field="tag",
table_name="litellm_dailytagspend",
unique_constraint_name="tag_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name",
unique_constraint_name="tag_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint",
)
async def _common_add_spend_log_transaction_to_daily_transaction(
@ -1513,6 +1519,12 @@ class DBSpendUpdateWriter:
)
return None
try:
# Map call_type to endpoint using ROUTE_ENDPOINT_MAPPING
call_type = payload.get("call_type", None)
endpoint = None
if call_type:
endpoint = ROUTE_ENDPOINT_MAPPING.get(call_type, None)
daily_transaction = BaseDailySpendTransaction(
date=date,
api_key=payload["api_key"],
@ -1520,6 +1532,7 @@ class DBSpendUpdateWriter:
model_group=payload.get("model_group", None),
mcp_namespaced_tool_name=payload.get("mcp_namespaced_tool_name", None),
custom_llm_provider=payload.get("custom_llm_provider", None),
endpoint=endpoint,
prompt_tokens=payload["prompt_tokens"],
completion_tokens=payload["completion_tokens"],
spend=payload["spend"],
@ -1563,7 +1576,8 @@ class DBSpendUpdateWriter:
if base_daily_transaction is None:
return
daily_transaction_key = f"{payload['user']}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}"
endpoint_str = base_daily_transaction.get("endpoint") or ""
daily_transaction_key = f"{payload['user']}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}"
daily_transaction = DailyUserSpendTransaction(
user_id=payload["user"], **base_daily_transaction
)
@ -1595,7 +1609,8 @@ class DBSpendUpdateWriter:
)
return
daily_transaction_key = f"{payload['team_id']}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}"
endpoint_str = base_daily_transaction.get("endpoint") or ""
daily_transaction_key = f"{payload['team_id']}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}"
daily_transaction = DailyTeamSpendTransaction(
team_id=payload["team_id"], **base_daily_transaction
)
@ -1637,7 +1652,8 @@ class DBSpendUpdateWriter:
if base_daily_transaction is None:
return
daily_transaction_key = f"{org_id}_{base_daily_transaction['date']}_{payload_with_org['api_key']}_{payload_with_org['model']}_{payload_with_org['custom_llm_provider']}"
endpoint_str = base_daily_transaction.get("endpoint") or ""
daily_transaction_key = f"{org_id}_{base_daily_transaction['date']}_{payload_with_org['api_key']}_{payload_with_org['model']}_{payload_with_org['custom_llm_provider']}_{endpoint_str}"
daily_transaction = DailyOrganizationSpendTransaction(
organization_id=org_id, **base_daily_transaction
)
@ -1679,7 +1695,8 @@ class DBSpendUpdateWriter:
if base_daily_transaction is None:
return
daily_transaction_key = f"{end_user_id}_{base_daily_transaction['date']}_{payload_with_end_user_id['api_key']}_{payload_with_end_user_id['model']}_{payload_with_end_user_id['custom_llm_provider']}"
endpoint_str = base_daily_transaction.get("endpoint") or ""
daily_transaction_key = f"{end_user_id}_{base_daily_transaction['date']}_{payload_with_end_user_id['api_key']}_{payload_with_end_user_id['model']}_{payload_with_end_user_id['custom_llm_provider']}_{endpoint_str}"
daily_transaction = DailyEndUserSpendTransaction(
end_user_id=end_user_id, **base_daily_transaction
)
@ -1723,7 +1740,8 @@ class DBSpendUpdateWriter:
)
if base_daily_transaction is None:
return
daily_transaction_key = f"{payload['agent_id']}_{base_daily_transaction['date']}_{payload_with_agent_id['api_key']}_{payload_with_agent_id['model']}_{payload_with_agent_id['custom_llm_provider']}"
endpoint_str = base_daily_transaction.get("endpoint") or ""
daily_transaction_key = f"{payload['agent_id']}_{base_daily_transaction['date']}_{payload_with_agent_id['api_key']}_{payload_with_agent_id['model']}_{payload_with_agent_id['custom_llm_provider']}_{endpoint_str}"
daily_transaction = DailyAgentSpendTransaction(
agent_id=payload['agent_id'], **base_daily_transaction
)
@ -1763,7 +1781,8 @@ class DBSpendUpdateWriter:
else:
raise ValueError(f"Invalid request_tags: {payload['request_tags']}")
for tag in request_tags:
daily_transaction_key = f"{tag}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}"
endpoint_str = base_daily_transaction.get("endpoint") or ""
daily_transaction_key = f"{tag}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}"
daily_transaction = DailyTagSpendTransaction(
tag=tag, **base_daily_transaction, request_id=payload["request_id"]
)

View file

@ -227,6 +227,41 @@ def update_breakdown_metrics(
)
)
# Update endpoint breakdown
if record.endpoint:
if record.endpoint not in breakdown.endpoints:
breakdown.endpoints[record.endpoint] = MetricWithMetadata(
metrics=SpendMetrics(),
metadata={},
)
breakdown.endpoints[record.endpoint].metrics = update_metrics(
breakdown.endpoints[record.endpoint].metrics, record
)
# Update API key breakdown for this endpoint
if record.api_key not in breakdown.endpoints[record.endpoint].api_key_breakdown:
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key] = (
KeyMetricWithMetadata(
metrics=SpendMetrics(),
metadata=KeyMetadata(
key_alias=api_key_metadata.get(record.api_key, {}).get(
"key_alias", None
),
team_id=api_key_metadata.get(record.api_key, {}).get(
"team_id", None
),
),
)
)
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics = (
update_metrics(
breakdown.endpoints[record.endpoint]
.api_key_breakdown[record.api_key]
.metrics,
record,
)
)
# Update api key breakdown
if record.api_key not in breakdown.api_keys:
breakdown.api_keys[record.api_key] = KeyMetricWithMetadata(

View file

@ -14,9 +14,10 @@ import copy
import json
import secrets
import traceback
import yaml
from datetime import datetime, timedelta, timezone
from typing import List, Literal, Optional, Tuple, cast
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status
@ -1033,7 +1034,7 @@ async def generate_key_fn(
- auto_rotate: Optional[bool] - Whether this key should be automatically rotated (regenerated)
- rotation_interval: Optional[str] - How often to auto-rotate this key (e.g., '30s', '30m', '30h', '30d'). Required if auto_rotate=True.
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint.
- router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"max_retries": 5}}. IF null or {} then no router settings.
Examples:
@ -1388,6 +1389,10 @@ async def prepare_key_update_data(
if "model_max_budget" in non_default_values:
validate_model_max_budget(non_default_values["model_max_budget"])
# Serialize router_settings to JSON if present
if "router_settings" in non_default_values and non_default_values["router_settings"] is not None:
non_default_values["router_settings"] = safe_dumps(non_default_values["router_settings"])
non_default_values = prepare_metadata_fields(
data=data, non_default_values=non_default_values, existing_metadata=_metadata
)
@ -1489,7 +1494,8 @@ async def update_key_fn(
- auto_rotate: Optional[bool] - Whether this key should be automatically rotated
- rotation_interval: Optional[str] - How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint.
- router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"max_retries": 5}}. IF null or {} then no router settings.
Example:
```bash
curl --location 'http://0.0.0.0:4000/key/update' \
@ -2080,6 +2086,7 @@ async def generate_key_helper_fn( # noqa: PLR0915
object_permission: Optional[LiteLLM_ObjectPermissionBase] = None,
auto_rotate: Optional[bool] = None,
rotation_interval: Optional[str] = None,
router_settings: Optional[dict] = None,
):
from litellm.proxy.proxy_server import premium_user, prisma_client
@ -2114,6 +2121,7 @@ async def generate_key_helper_fn( # noqa: PLR0915
aliases_json = json.dumps(aliases)
config_json = json.dumps(config)
permissions_json = json.dumps(permissions)
router_settings_json = safe_dumps(router_settings) if router_settings is not None else safe_dumps({})
# Add model_rpm_limit and model_tpm_limit to metadata
if model_rpm_limit is not None:
@ -2189,6 +2197,7 @@ async def generate_key_helper_fn( # noqa: PLR0915
"updated_by": updated_by,
"allowed_routes": allowed_routes or [],
"object_permission_id": object_permission_id,
"router_settings": router_settings_json,
}
# Add rotation fields if auto_rotate is enabled
@ -2225,6 +2234,13 @@ async def generate_key_helper_fn( # noqa: PLR0915
saved_token["model_max_budget"] = json.loads(
saved_token["model_max_budget"]
)
router_settings = cast(Optional[dict], saved_token.get("router_settings"))
if router_settings is not None and isinstance(router_settings, str):
try:
saved_token["router_settings"] = yaml.safe_load(router_settings)
except yaml.YAMLError:
# If it's not valid JSON/YAML, keep as is or set to empty dict
saved_token["router_settings"] = {}
if saved_token.get("expires", None) is not None and isinstance(
saved_token["expires"], datetime
@ -2269,6 +2285,15 @@ async def generate_key_helper_fn( # noqa: PLR0915
)
key_data["created_at"] = getattr(create_key_response, "created_at", None)
key_data["updated_at"] = getattr(create_key_response, "updated_at", None)
# Deserialize router_settings from JSON string to dict for response
router_settings_value = key_data.get("router_settings")
if router_settings_value is not None and isinstance(router_settings_value, str):
try:
key_data["router_settings"] = yaml.safe_load(router_settings_value)
except yaml.YAMLError:
# If it's not valid JSON/YAML, keep as is or set to empty dict
key_data["router_settings"] = {}
except Exception as e:
verbose_proxy_logger.error(
"litellm.proxy.proxy_server.generate_key_helper_fn(): Exception occured - {}".format(

View file

@ -100,7 +100,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
TeamMemberAddResult,
UpdateTeamMemberPermissionsRequest,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
router = APIRouter()
@ -696,8 +696,7 @@ async def new_team( # noqa: PLR0915
- allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team.
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint.
- secret_manager_settings: Optional[dict] - Secret manager settings for the team. [Docs](https://docs.litellm.ai/docs/secret_managers/overview)
- router_settings: Optional[UpdateRouterConfig] - team-specific router settings. Example - {"model_group_retry_policy": {"max_retries": 5}}. IF null or {} then no router settings.
Returns:
- team_id: (str) Unique team id - used for tracking spend across multiple keys for same team id.
@ -911,6 +910,12 @@ async def new_team( # noqa: PLR0915
complete_team_data.members_with_roles = []
complete_team_data_dict = complete_team_data.model_dump(exclude_none=True)
# Serialize router_settings to JSON (matching key creation pattern)
router_settings_value = getattr(data, "router_settings", None)
router_settings_json = safe_dumps(router_settings_value) if router_settings_value is not None else safe_dumps({})
complete_team_data_dict["router_settings"] = router_settings_json
complete_team_data_dict = prisma_client.jsonify_team_object(
db_data=complete_team_data_dict
)
@ -1234,7 +1239,7 @@ async def update_team( # noqa: PLR0915
Example - update team TPM Limit
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint.
- secret_manager_settings: Optional[dict] - Secret manager settings for the team. [Docs](https://docs.litellm.ai/docs/secret_managers/overview)
- router_settings: Optional[UpdateRouterConfig] - team-specific router settings. Example - {"model_group_retry_policy": {"max_retries": 5}}. IF null or {} then no router settings.
```
curl --location 'http://0.0.0.0:4000/team/update' \
@ -1396,6 +1401,10 @@ async def update_team( # noqa: PLR0915
if _model_id is not None:
updated_kv["model_id"] = _model_id
# Serialize router_settings to JSON if present (matching key update pattern)
if "router_settings" in updated_kv and updated_kv["router_settings"] is not None:
updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"])
updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv)
team_row: Optional[LiteLLM_TeamTable] = (
await prisma_client.db.litellm_teamtable.update(

View file

@ -2,4 +2,7 @@ model_list:
- model_name: anthropic/*
litellm_params:
model: anthropic/*
- model_name: openai/*
litellm_params:
model: openai/*

View file

@ -38,6 +38,7 @@ ROUTE_ENDPOINT_MAPPING = {
"aretrieve_container": "/containers/{container_id}",
"adelete_container": "/containers/{container_id}",
# Auto-generated container file routes
"aupload_container_file": "/containers/{container_id}/files",
"alist_container_files": "/containers/{container_id}/files",
"aretrieve_container_file": "/containers/{container_id}/files/{file_id}",
"adelete_container_file": "/containers/{container_id}/files/{file_id}",
@ -144,6 +145,7 @@ async def route_request(
"alist_containers",
"aretrieve_container",
"adelete_container",
"aupload_container_file",
"alist_container_files",
"aretrieve_container_file",
"adelete_container_file",
@ -204,6 +206,7 @@ async def route_request(
"alist_containers",
"aretrieve_container",
"adelete_container",
"aupload_container_file",
"alist_container_files",
"aretrieve_container_file",
"adelete_container_file",
@ -287,6 +290,7 @@ async def route_request(
"alist_containers",
"aretrieve_container",
"adelete_container",
"aupload_container_file",
"alist_container_files",
"aretrieve_container_file",
"adelete_container_file",

View file

@ -124,6 +124,7 @@ model LiteLLM_TeamTable {
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
@ -225,6 +226,7 @@ model LiteLLM_VerificationToken {
models String[]
aliases Json @default("{}")
config Json @default("{}")
router_settings Json? @default("{}")
user_id String?
team_id String?
permissions Json @default("{}")
@ -422,6 +424,7 @@ model LiteLLM_DailyUserSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -433,12 +436,13 @@ model LiteLLM_DailyUserSpend {
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([user_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily organization spend metrics per model and key
@ -451,6 +455,7 @@ model LiteLLM_DailyOrganizationSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -462,12 +467,13 @@ model LiteLLM_DailyOrganizationSpend {
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([organization_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily end user (customer) spend metrics per model and key
@ -480,6 +486,7 @@ model LiteLLM_DailyEndUserSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -490,12 +497,13 @@ model LiteLLM_DailyEndUserSpend {
failed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([end_user_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily agent spend metrics per model and key
@ -508,6 +516,7 @@ model LiteLLM_DailyAgentSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -518,12 +527,13 @@ model LiteLLM_DailyAgentSpend {
failed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([agent_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily team spend metrics per model and key
@ -536,6 +546,7 @@ model LiteLLM_DailyTeamSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -547,12 +558,13 @@ model LiteLLM_DailyTeamSpend {
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([team_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily team spend metrics per model and key
@ -566,6 +578,7 @@ model LiteLLM_DailyTagSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -577,12 +590,13 @@ model LiteLLM_DailyTagSpend {
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([tag])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}

View file

@ -4003,6 +4003,8 @@ class Router:
"retrieve_container",
"adelete_container",
"delete_container",
"aupload_container_file",
"upload_container_file",
"alist_container_files",
"list_container_files",
"aretrieve_container_file",
@ -4154,6 +4156,7 @@ class Router:
"alist_containers",
"aretrieve_container",
"adelete_container",
"aupload_container_file",
"alist_container_files",
"aretrieve_container_file",
"adelete_container_file",

View file

@ -14,3 +14,4 @@ class ArizeConfig(BaseModel):
api_key: Optional[str] = None
protocol: Protocol
endpoint: str
project_name: Optional[str] = None

View file

@ -68,6 +68,9 @@ class BreakdownMetrics(BaseModel):
providers: Dict[str, MetricWithMetadata] = Field(
default_factory=dict
) # provider -> {metrics, metadata}
endpoints: Dict[str, MetricWithMetadata] = Field(
default_factory=dict
) # endpoint -> {metrics, metadata}
api_keys: Dict[str, KeyMetricWithMetadata] = Field(
default_factory=dict
) # api_key -> {metrics, metadata}

View file

@ -324,6 +324,8 @@ class CallTypes(str, Enum):
adelete_container = "adelete_container"
list_container_files = "list_container_files"
alist_container_files = "alist_container_files"
upload_container_file = "upload_container_file"
aupload_container_file = "aupload_container_file"
acancel_fine_tuning_job = "acancel_fine_tuning_job"
cancel_fine_tuning_job = "cancel_fine_tuning_job"

View file

@ -4909,6 +4909,15 @@
"/v1/images/generations"
]
},
"azure_ai/flux.2-pro": {
"litellm_provider": "azure_ai",
"mode": "image_generation",
"output_cost_per_image": 0.04,
"source": "https://ai.azure.com/explore/models/flux.2-pro/version/1/registry/azureml-blackforestlabs",
"supported_endpoints": [
"/v1/images/generations"
]
},
"azure_ai/Llama-3.2-11B-Vision-Instruct": {
"input_cost_per_token": 3.7e-07,
"litellm_provider": "azure_ai",
@ -28336,6 +28345,19 @@
"supports_tool_choice": true,
"supports_web_search": true
},
"vertex_ai/zai-org/glm-4.7-maas": {
"input_cost_per_token": 3e-07,
"litellm_provider": "vertex_ai-zai_models",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"vertex_ai/mistral-medium-3": {
"input_cost_per_token": 4e-07,
"litellm_provider": "vertex_ai-mistral_models",

10
poetry.lock generated
View file

@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand.
# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand.
[[package]]
name = "aiofiles"
@ -3081,15 +3081,15 @@ files = [
[[package]]
name = "litellm-proxy-extras"
version = "0.4.16"
version = "0.4.18"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
optional = true
python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8"
groups = ["main"]
markers = "extra == \"proxy\""
files = [
{file = "litellm_proxy_extras-0.4.16-py3-none-any.whl", hash = "sha256:5651e777c7f4c0e87c6722971bca19b8f40f417b08f74001cab2d0a5b1c63a91"},
{file = "litellm_proxy_extras-0.4.16.tar.gz", hash = "sha256:ff1ee4ea119318b471bb71a99d8bc941159d4d2c09bee797dd29768e9504befb"},
{file = "litellm_proxy_extras-0.4.18-py3-none-any.whl", hash = "sha256:c3edee68bf8eb073c6158dcf7df05727dfc829e63c03a617fcb48853d11490df"},
{file = "litellm_proxy_extras-0.4.18.tar.gz", hash = "sha256:898b28e3e74acdc29142906b84787ab05a90e30aa3c0c8aee849915e3a16adb3"},
]
[[package]]
@ -7981,4 +7981,4 @@ utils = ["numpydoc"]
[metadata]
lock-version = "2.1"
python-versions = ">=3.9,<4.0"
content-hash = "7eed2b2c25173a275ac83c55fd901b9b84663b1d7daa54f0e78b30bf1c8f0e3e"
content-hash = "e9fd12b5ccc703ec156d98877452417083e3ac18b5970cb3a58c3bde09d267bb"

View file

@ -20,16 +20,14 @@
"skills": "Supports /skills endpoint",
"interactions": "Supports /interactions endpoint (Google AI Interactions API)",
"a2a_(Agent Gateway)": "Supports /a2a/{agent}/message/send endpoint (A2A Protocol)",
"create_container": "Supports POST /containers endpoint",
"list_containers": "Supports GET /containers endpoint",
"retrieve_container": "Supports GET /containers/{id} endpoint",
"delete_container": "Supports DELETE /containers/{id} endpoint",
"create_container_file": "Supports POST /containers/{id}/files endpoint",
"list_container_files": "Supports GET /containers/{id}/files endpoint",
"retrieve_container_file": "Supports GET /containers/{id}/files/{file_id} endpoint",
"retrieve_container_file_content": "Supports GET /containers/{id}/files/{file_id}/content endpoint",
"delete_container_file": "Supports DELETE /containers/{id}/files/{file_id} endpoint",
"compact": "Supports /responses/compact endpoint"
"container": "Supports OpenAI's /containers endpoint",
"container_file": "Supports OpenAI's /containers/{id}/files endpoint",
"compact": "Supports /responses/compact endpoint",
"files": "Supports /files endpoint for file operations",
"image_edits": "Supports /images/edits endpoint for image editing",
"vector_stores_create": "Supports creating a new vector store via /vector_stores endpoint",
"vector_stores_search": "Supports searching a vector store via /vector_stores/{id}/search endpoint",
"video_generations": "Supports /videos/generations endpoint for video generation"
}
}
},
@ -122,7 +120,8 @@
"rerank": false,
"skills": true,
"a2a": true,
"interactions": true
"interactions": true,
"count_tokens": true
}
},
"anthropic_text": {
@ -211,7 +210,13 @@
"batches": false,
"rerank": true,
"a2a": true,
"interactions": true
"interactions": true,
"bedrock_invoke": true,
"bedrock_converse": true,
"vector_stores_search": true,
"count_tokens": true,
"rag_ingest": true,
"rag_query": true
}
},
"sagemaker": {
@ -263,7 +268,11 @@
"batches": true,
"rerank": false,
"a2a": true,
"interactions": true
"interactions": true,
"vector_stores_search": true,
"assistants": true,
"fine_tuning": true,
"text_completion": true
}
},
"azure_ai": {
@ -275,6 +284,7 @@
"responses": true,
"embeddings": true,
"image_generations": true,
"image_edits": true,
"audio_transcriptions": true,
"audio_speech": true,
"moderations": true,
@ -282,7 +292,9 @@
"rerank": false,
"ocr": true,
"a2a": true,
"interactions": true
"interactions": true,
"vector_stores_create": true,
"vector_stores_search": true
}
},
"azure_ai/doc-intelligence": {
@ -918,29 +930,19 @@
"embeddings": true,
"image_generations": true,
"audio_transcriptions": false,
"audio_speech": false,
"audio_speech": true,
"moderations": false,
"batches": false,
"rerank": false,
"ocr": true,
"a2a": true,
"interactions": true
}
},
"vertex_ai/chirp": {
"display_name": "Google - Vertex AI Chirp3 HD (`vertex_ai/chirp`)",
"url": "https://docs.litellm.ai/docs/providers/vertex_speech",
"endpoints": {
"chat_completions": false,
"messages": false,
"responses": false,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_speech": true,
"moderations": false,
"batches": false,
"rerank": false
"interactions": true,
"vector_stores_search": true,
"count_tokens": true,
"fine_tuning": true,
"rag_ingest": true,
"rag_query": true,
"generateContent": true
}
},
"gemini": {
@ -958,7 +960,12 @@
"batches": false,
"rerank": false,
"interactions": true,
"a2a": true
"a2a": true,
"vector_stores_search": true,
"count_tokens": true,
"rag_ingest": true,
"realtime": true,
"generateContent": true
}
},
"gradient_ai": {
@ -1511,18 +1518,21 @@
"moderations": true,
"batches": true,
"rerank": false,
"create_container": true,
"list_containers": true,
"retrieve_container": true,
"delete_container": true,
"create_container_file": false,
"list_container_files": true,
"retrieve_container_file": true,
"retrieve_container_file_content": true,
"delete_container_file": true,
"container": true,
"compact": true,
"a2a": true,
"interactions": true
"interactions": true,
"vector_store_files": true,
"vector_stores_create": true,
"vector_stores_search": true,
"assistants": true,
"container_files": true,
"fine_tuning": true,
"image_variations": true,
"rag_ingest": true,
"rag_query": true,
"realtime": true,
"text_completion": true
}
},
"openai_like": {
@ -1538,7 +1548,8 @@
"audio_speech": false,
"moderations": false,
"batches": false,
"rerank": false
"rerank": false,
"assistants": true
}
},
"openrouter": {
@ -1897,34 +1908,13 @@
"display_name": "Topaz (`topaz`)",
"url": "https://docs.litellm.ai/docs/providers/topaz",
"endpoints": {
"chat_completions": true,
"messages": true,
"responses": true,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,
"batches": false,
"rerank": false,
"a2a": true,
"interactions": true
"image_variations": true
}
},
"tavily": {
"display_name": "Tavily (`tavily`)",
"url": "https://docs.litellm.ai/docs/search/tavily",
"endpoints": {
"chat_completions": false,
"messages": false,
"responses": false,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,
"batches": false,
"rerank": false,
"search": true
}
},
@ -2137,7 +2127,7 @@
"moderations": false,
"batches": false,
"rerank": false,
"vector_stores": true,
"vector_stores_create": true,
"a2a": true,
"interactions": true
}
@ -2247,6 +2237,340 @@
"a2a": true,
"interactions": true
}
},
"gigachat": {
"display_name": "GigaChat (`gigachat`)",
"url": "https://docs.litellm.ai/docs/providers/gigachat",
"endpoints": {
"chat_completions": true,
"messages": true,
"responses": true,
"embeddings": true
}
},
"google_pse": {
"display_name": "Google PSE (`google_pse`)",
"url": "https://docs.litellm.ai/docs/search/google_pse",
"endpoints": {
"search": true
}
},
"milvus": {
"display_name": "Milvus (`milvus`)",
"url": "https://docs.litellm.ai/docs/providers/milvus_vector_stores",
"endpoints": {
"vector_stores_search": true
}
},
"minimax": {
"display_name": "Minimax (`minimax`)",
"url": "https://docs.litellm.ai/docs/providers/minimax",
"endpoints": {
"chat_completions": true,
"messages": true,
"responses": true
}
},
"pg_vector": {
"display_name": "PG Vector (`pg_vector`)",
"url": "https://docs.litellm.ai/docs/providers/pg_vector",
"endpoints": {
"vector_stores_search": true
}
},
"helicone": {
"display_name": "Helicone (`helicone`)",
"url": "https://docs.litellm.ai/docs/providers/helicone",
"endpoints": {
"chat_completions": true,
"messages": true,
"responses": true
}
},
"llamagate": {
"display_name": "LlamaGate (`llamagate`)",
"url": "https://docs.litellm.ai/docs/providers/llamagate",
"endpoints": {
"chat_completions": true,
"messages": true,
"responses": true
}
},
"xiaomi_mimo": {
"display_name": "Xiaomi Mimo (`xiaomi_mimo`)",
"url": "https://docs.litellm.ai/docs/providers/xiaomi_mimo",
"endpoints": {
"chat_completions": true,
"messages": true,
"responses": true
}
}
},
"endpoints": {
"a2a": {
"docs_label": "a2a",
"display_name": "A2A (Agent-to-Agent) protocol for agent communication",
"leftnav_label": "/a2a",
"provider_json_field": "a2a",
"url": "https://docs.litellm.ai/docs/a2a",
"bridges_to_chat_completion": true
},
"messages": {
"docs_label": "anthropic_unified",
"display_name": "Anthropic /v1/messages API",
"leftnav_label": "/messages",
"provider_json_field": "messages",
"url": "https://docs.litellm.ai/docs/anthropic_unified",
"bridges_to_chat_completion": true
},
"anthropic_count_tokens": {
"docs_label": "anthropic_count_tokens",
"display_name": "Anthropic /v1/messages/count_tokens API",
"leftnav_label": "/count_tokens",
"provider_json_field": "count_tokens",
"url": "https://docs.litellm.ai/docs/anthropic_count_tokens"
},
"apply_guardrail": {
"docs_label": "apply_guardrail",
"display_name": "Unified Apply Guardrail API",
"leftnav_label": "/guardrails/apply_guardrail",
"provider_json_field": "apply_guardrail",
"url": "https://docs.litellm.ai/docs/apply_guardrail"
},
"assistants": {
"docs_label": "assistants",
"display_name": "OpenAI Assistants API",
"leftnav_label": "/assistants",
"provider_json_field": "assistants",
"url": "https://docs.litellm.ai/docs/assistants"
},
"audio_transcription": {
"docs_label": "audio_transcription",
"display_name": "Audio Transcription API",
"leftnav_label": "/audio/transcriptions",
"provider_json_field": "audio_transcriptions",
"url": "https://docs.litellm.ai/docs/audio_transcription"
},
"batches": {
"docs_label": "batches",
"display_name": "Batches API",
"leftnav_label": "/batches",
"provider_json_field": "batches",
"url": "https://docs.litellm.ai/docs/batches"
},
"bedrock_invoke": {
"docs_label": "bedrock_invoke",
"display_name": "Bedrock Invoke API",
"leftnav_label": "/invoke",
"provider_json_field": "bedrock_invoke",
"url": "https://docs.litellm.ai/docs/bedrock_invoke"
},
"bedrock_converse": {
"docs_label": "bedrock_converse",
"display_name": "Bedrock Converse API",
"leftnav_label": "/converse",
"provider_json_field": "bedrock_converse",
"url": "https://docs.litellm.ai/docs/bedrock_converse"
},
"chat_completions": {
"docs_label": "chat_completions",
"display_name": "Chat Completions API",
"leftnav_label": "/chat/completions",
"provider_json_field": "chat_completions",
"url": "https://docs.litellm.ai/docs/chat_completions"
},
"container_files": {
"docs_label": "container_files",
"display_name": "OpenAI Container Files API",
"leftnav_label": "/create/container/files",
"provider_json_field": "container_files",
"url": "https://docs.litellm.ai/docs/container_files"
},
"container": {
"docs_label": "containers",
"display_name": "OpenAI Containers API",
"leftnav_label": "/container",
"provider_json_field": "container",
"url": "https://docs.litellm.ai/docs/containers"
},
"embeddings": {
"docs_label": "embedding/supported_embedding",
"display_name": "Embedding API (OpenAI Format)",
"leftnav_label": "/embeddings",
"provider_json_field": "embeddings",
"url": "https://docs.litellm.ai/docs/embedding/supported_embedding"
},
"files": {
"docs_label": "files",
"display_name": "OpenAI Files API",
"leftnav_label": "/files",
"provider_json_field": "files",
"url": "https://docs.litellm.ai/docs/proxy/litellm_managed_files"
},
"fine_tuning": {
"docs_label": "fine_tuning",
"display_name": "OpenAI Fine-Tuning API",
"leftnav_label": "/fine_tuning",
"provider_json_field": "fine_tuning",
"url": "https://docs.litellm.ai/docs/proxy/managed_finetuning"
},
"generateContent": {
"docs_label": "generateContent",
"display_name": "Google's GenerateContent API",
"leftnav_label": "/generateContent",
"provider_json_field": "generateContent",
"url": "https://docs.litellm.ai/docs/generateContent",
"bridges_to_chat_completion": true
},
"image_edits": {
"docs_label": "image_edits",
"display_name": "OpenAI Images Edits API",
"leftnav_label": "/images/edits",
"provider_json_field": "image_edits",
"url": "https://docs.litellm.ai/docs/image_edits"
},
"image_generations": {
"docs_label": "image_generation",
"display_name": "OpenAI Images Generations API",
"leftnav_label": "/images/generations",
"provider_json_field": "image_generations",
"url": "https://docs.litellm.ai/docs/image_generation"
},
"image_variations": {
"docs_label": "image_variations",
"display_name": "OpenAI Images Variations API",
"leftnav_label": "/images/variations",
"provider_json_field": "image_variations",
"url": "https://docs.litellm.ai/docs/image_variations"
},
"interactions": {
"docs_label": "interactions",
"display_name": "Google Interactions API",
"leftnav_label": "/interactions",
"provider_json_field": "interactions",
"url": "https://docs.litellm.ai/docs/interactions",
"bridges_to_chat_completion": true
},
"mcp": {
"docs_label": "mcp",
"display_name": "Model Context Protocol (MCP)",
"leftnav_label": "/mcp",
"provider_json_field": "mcp",
"url": "https://docs.litellm.ai/docs/mcp"
},
"moderation": {
"docs_label": "moderation",
"display_name": "OpenAI Moderation API",
"leftnav_label": "/moderations",
"provider_json_field": "moderations",
"url": "https://docs.litellm.ai/docs/moderation"
},
"ocr": {
"docs_label": "ocr",
"display_name": "OCR API (Mistral Format)",
"leftnav_label": "/ocr",
"provider_json_field": "ocr",
"url": "https://docs.litellm.ai/docs/ocr"
},
"rag_ingest": {
"docs_label": "rag_ingest",
"display_name": "RAG Ingest API",
"leftnav_label": "/rag/ingest",
"provider_json_field": "rag_ingest",
"url": "https://docs.litellm.ai/docs/rag_ingest"
},
"rag_query": {
"docs_label": "rag_query",
"display_name": "RAG Query API",
"leftnav_label": "/rag/query",
"provider_json_field": "rag_query",
"url": "https://docs.litellm.ai/docs/rag_query"
},
"realtime": {
"docs_label": "realtime",
"display_name": "OpenAI Realtime API",
"leftnav_label": "/realtime",
"provider_json_field": "realtime",
"url": "https://docs.litellm.ai/docs/realtime"
},
"rerank": {
"docs_label": "rerank",
"display_name": "Rerank API (Cohere Format)",
"leftnav_label": "/rerank",
"provider_json_field": "rerank",
"url": "https://docs.litellm.ai/docs/rerank"
},
"responses": {
"docs_label": "response_api",
"display_name": "Responses API (OpenAI Format)",
"leftnav_label": "/responses",
"provider_json_field": "responses",
"url": "https://docs.litellm.ai/docs/response_api",
"bridges_to_chat_completion": true
},
"response_api_compact": {
"docs_label": "response_api_compact",
"display_name": "Responses API (OpenAI Format)",
"leftnav_label": "/responses",
"provider_json_field": "compact",
"url": "https://docs.litellm.ai/docs/response_api"
},
"search": {
"docs_label": "search",
"display_name": "Search API",
"leftnav_label": "/search",
"provider_json_field": "search",
"url": "https://docs.litellm.ai/docs/search"
},
"skills": {
"docs_label": "skills",
"display_name": "Anthropic Skills API",
"leftnav_label": "/skills",
"provider_json_field": "skills",
"url": "https://docs.litellm.ai/docs/skills"
},
"text_completion": {
"docs_label": "text_completion",
"display_name": "Completions API (OpenAI Format)",
"leftnav_label": "/completions",
"provider_json_field": "text_completion",
"url": "https://docs.litellm.ai/docs/text_completion",
"bridges_to_chat_completion": true
},
"text_to_speech": {
"docs_label": "text_to_speech",
"display_name": "Text-to-Speech API (OpenAI Format)",
"leftnav_label": "/audio/speech",
"provider_json_field": "audio_speech",
"url": "https://docs.litellm.ai/docs/text_to_speech"
},
"vector_store_files": {
"docs_label": "vector_store_files",
"display_name": "OpenAI Vector Store Files API",
"leftnav_label": "/vector_stores/files",
"provider_json_field": "vector_store_files",
"url": "https://docs.litellm.ai/docs/vector_store_files"
},
"vector_stores_create": {
"docs_label": "vector_stores_create",
"display_name": "OpenAI Vector Stores Create API",
"leftnav_label": "/vector_stores/create",
"provider_json_field": "vector_stores_create",
"url": "https://docs.litellm.ai/docs/vector_stores/create"
},
"vector_stores_search": {
"docs_label": "vector_stores_search",
"display_name": "OpenAI Vector Stores Search API",
"leftnav_label": "/vector_stores/search",
"provider_json_field": "vector_stores_search",
"url": "https://docs.litellm.ai/docs/vector_stores/search"
},
"videos": {
"docs_label": "videos",
"display_name": "OpenAI Video Generation API",
"leftnav_label": "/videos",
"provider_json_field": "video_generations",
"url": "https://docs.litellm.ai/docs/videos"
}
}
}

View file

@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm"
version = "1.80.11"
version = "1.80.12"
description = "Library to easily interface with LLM API providers"
authors = ["BerriAI"]
license = "MIT"
@ -59,7 +59,7 @@ websockets = {version = "^15.0.1", optional = true}
boto3 = {version = "1.36.0", optional = true}
redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"}
mcp = {version = "^1.21.2", optional = true, python = ">=3.10"}
litellm-proxy-extras = {version = "0.4.16", optional = true}
litellm-proxy-extras = {version = "0.4.20", optional = true}
rich = {version = "13.7.1", optional = true}
litellm-enterprise = {version = "0.1.27", optional = true}
diskcache = {version = "^5.6.1", optional = true}
@ -167,7 +167,7 @@ requires = ["poetry-core", "wheel"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "1.80.11"
version = "1.80.12"
version_files = [
"pyproject.toml:^version"
]

View file

@ -15,7 +15,7 @@ redis==5.2.1 # redis caching
prisma==0.11.0 # for db
nodejs-wheel-binaries==24.12.0 ## required by prisma for migrations, prevents runtime download (updated from nodejs-bin for security fixes)
mangum==0.17.0 # for aws lambda functions
pynacl==1.5.0 # for encrypting keys
pynacl==1.6.2 # for encrypting keys
google-cloud-aiplatform==1.47.0 # for vertex ai calls
google-cloud-iam==2.19.1 # for GCP IAM Redis authentication
google-genai==1.22.0
@ -47,7 +47,7 @@ sentry_sdk==2.21.0 # for sentry error handling
detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests
cryptography==44.0.1
tzdata==2025.1 # IANA time zone database
litellm-proxy-extras==0.4.16 # for proxy extras - e.g. prisma migrations
litellm-proxy-extras==0.4.20 # for proxy extras - e.g. prisma migrations
llm-sandbox==0.3.31 # for skill execution in sandbox
### LITELLM PACKAGE DEPENDENCIES
python-dotenv==1.0.1 # for env
@ -57,7 +57,7 @@ tokenizers==0.20.2 # for calculating usage
click==8.1.7 # for proxy cli
rich==13.7.1 # for litellm proxy cli
jinja2==3.1.6 # for prompt templates
aiohttp==3.12.14 # for network calls
aiohttp==3.13.3 # for network calls
aioboto3==13.4.0 # for async sagemaker calls
tenacity==8.5.0 # for retrying requests, when litellm.num_retries set
pydantic>=2.11,<3 # proxy + openai req. + mcp

View file

@ -124,6 +124,7 @@ model LiteLLM_TeamTable {
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
@ -225,6 +226,7 @@ model LiteLLM_VerificationToken {
models String[]
aliases Json @default("{}")
config Json @default("{}")
router_settings Json? @default("{}")
user_id String?
team_id String?
permissions Json @default("{}")
@ -422,6 +424,7 @@ model LiteLLM_DailyUserSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -433,12 +436,13 @@ model LiteLLM_DailyUserSpend {
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([user_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily organization spend metrics per model and key
@ -451,6 +455,7 @@ model LiteLLM_DailyOrganizationSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -462,12 +467,13 @@ model LiteLLM_DailyOrganizationSpend {
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([organization_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily end user (customer) spend metrics per model and key
@ -480,6 +486,7 @@ model LiteLLM_DailyEndUserSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -490,12 +497,13 @@ model LiteLLM_DailyEndUserSpend {
failed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([end_user_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily agent spend metrics per model and key
@ -508,6 +516,7 @@ model LiteLLM_DailyAgentSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -518,12 +527,13 @@ model LiteLLM_DailyAgentSpend {
failed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([agent_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily team spend metrics per model and key
@ -536,6 +546,7 @@ model LiteLLM_DailyTeamSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -547,12 +558,13 @@ model LiteLLM_DailyTeamSpend {
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([team_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}
// Track daily team spend metrics per model and key
@ -566,6 +578,7 @@ model LiteLLM_DailyTagSpend {
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
@ -577,12 +590,13 @@ model LiteLLM_DailyTagSpend {
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
@@index([tag])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@index([endpoint])
}

BIN
test_image_edit.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 70 B

View file

@ -0,0 +1,379 @@
"""
Code coverage test to ensure all endpoints documented in sidebars.js are defined in provider_endpoints_support.json.
This script:
1. Extracts all endpoint entries from the "Supported Endpoints" section of sidebars.js
2. Validates that each endpoint has a corresponding entry in the "endpoints" object of provider_endpoints_support.json
3. Checks that the "docs_label" field is present in each endpoint definition
"""
import json
import re
import sys
from pathlib import Path
from typing import Dict, List, Set, Tuple
class MissingEndpointDefinitionError(Exception):
"""Raised when endpoints are documented in sidebars.js but missing from provider_endpoints_support.json."""
pass
def get_repo_root() -> Path:
"""Get the repository root directory."""
# Check if litellm directory exists in current working directory
cwd = Path.cwd()
if (cwd / "litellm").exists() and (cwd / "litellm").is_dir():
# We're already at the repo root
return cwd
# Otherwise, navigate up from script location
current = Path(__file__).resolve()
# Navigate up from tests/code_coverage_tests/
return current.parent.parent.parent
def extract_endpoints_from_sidebars() -> Dict[str, str]:
"""
Extract endpoint entries from sidebars.js.
Returns a dict mapping endpoint_key -> label
Only extracts top-level endpoint entries from the "Supported Endpoints" section.
"""
repo_root = get_repo_root()
sidebars_path = repo_root / "docs" / "my-website" / "sidebars.js"
if not sidebars_path.exists():
print(f"❌ ERROR: Could not find sidebars.js at {sidebars_path}")
sys.exit(1)
with open(sidebars_path, "r") as f:
content = f.read()
# Find the Supported Endpoints section
supported_start = content.find('label: "Supported Endpoints"')
if supported_start == -1:
print("⚠️ WARNING: Could not find 'Supported Endpoints' section")
return {}
# Find the items array within this section
items_start = content.find("items: [", supported_start)
if items_start == -1:
print("⚠️ WARNING: Could not find items array in Supported Endpoints")
return {}
# Find the end of this items array
# Look for the closing ], at the same indentation level
items_end = content.find("\n ],\n },\n {", items_start)
if items_end == -1:
items_end = content.find("\n ],\n }", items_start)
section = content[items_start:items_end]
endpoints = {}
# Pattern 1: Categories with labels at the top level (8 spaces indent)
# Example: " {type: "category", label: "/a2a - A2A Agent Gateway""
category_pattern = (
r'^\s{8}\{\s*\n\s{10}type:\s*"category",\s*\n\s{10}label:\s*"([^"]+)"'
)
for match in re.finditer(category_pattern, section, re.MULTILINE):
label = match.group(1)
# Skip utility categories
if "Pass-through" in label or label == "Vertex AI":
continue
endpoint_key = label.split(" - ")[0].strip("/").replace("/", "_")
endpoints[endpoint_key] = label
# Pattern 2: Standalone doc strings at top level (8 spaces indent)
# Example: " "assistants","
standalone_pattern = r'^\s{8}"([a-zA-Z_][a-zA-Z0-9_]*)",?\s*$'
for match in re.finditer(standalone_pattern, section, re.MULTILINE):
doc_id = match.group(1)
endpoints[doc_id] = doc_id
return endpoints
def load_provider_endpoints_file() -> Dict:
"""Load the provider_endpoints_support.json file."""
repo_root = get_repo_root()
file_path = repo_root / "provider_endpoints_support.json"
if not file_path.exists():
print(
f"❌ ERROR: Could not find provider_endpoints_support.json at {file_path}"
)
sys.exit(1)
with open(file_path, "r") as f:
return json.load(f)
def get_defined_endpoints(data: Dict) -> Dict[str, Dict]:
"""Get all endpoint definitions from provider_endpoints_support.json."""
return data.get("endpoints", {})
def normalize_endpoint_key(key: str) -> Set[str]:
"""
Generate variations of an endpoint key for matching.
Examples:
- "a2a" -> {"a2a"}
- "chat_completions" -> {"chat_completions", "chatcompletions"}
- "vector_stores" -> {"vector_stores", "vectorstores"}
"""
variations = {key, key.replace("_", "")}
return variations
def check_provider_endpoint_keys(data: Dict) -> List[str]:
"""
Check that all endpoint keys used in providers are defined in the root endpoints section.
Returns a list of missing endpoint keys.
"""
# Collect all unique endpoint keys used across all providers
provider_endpoint_keys = set()
providers = data.get("providers", {})
for provider_name, provider_data in providers.items():
if "endpoints" in provider_data and isinstance(
provider_data["endpoints"], dict
):
provider_endpoint_keys.update(provider_data["endpoints"].keys())
# Get all endpoint definitions
defined_endpoints = data.get("endpoints", {})
# Collect all provider_json_field values from endpoint definitions
provider_json_fields = set()
for endpoint_key, endpoint_data in defined_endpoints.items():
if isinstance(endpoint_data, dict) and "provider_json_field" in endpoint_data:
provider_json_fields.add(endpoint_data["provider_json_field"])
# Find missing endpoint keys
missing_keys = []
for key in sorted(provider_endpoint_keys):
if key not in provider_json_fields:
missing_keys.append(key)
return missing_keys
def check_unused_endpoints(data: Dict) -> List[Tuple[str, str]]:
"""
Check that all defined endpoints are used by at least one provider.
Returns a list of tuples (endpoint_key, provider_json_field) for unused endpoints.
"""
# Special endpoints that don't need to be used by specific providers
# These are utility/framework endpoints available across the platform
SPECIAL_ENDPOINTS = {
"apply_guardrail", # Guardrail application - works across providers
"mcp", # Model Context Protocol - works across providers
}
# Get all endpoint definitions
defined_endpoints = data.get("endpoints", {})
providers = data.get("providers", {})
# Collect all endpoint keys used by providers
used_keys = set()
for provider_data in providers.values():
if "endpoints" in provider_data and isinstance(
provider_data["endpoints"], dict
):
used_keys.update(provider_data["endpoints"].keys())
# Find unused endpoints (excluding special ones)
unused = []
for endpoint_key, endpoint_data in defined_endpoints.items():
# Skip special endpoints
if endpoint_key in SPECIAL_ENDPOINTS:
continue
if isinstance(endpoint_data, dict) and "provider_json_field" in endpoint_data:
provider_json_field = endpoint_data["provider_json_field"]
# Check if this provider_json_field is used by any provider
if provider_json_field not in used_keys:
unused.append((endpoint_key, provider_json_field))
return sorted(unused)
def main():
"""Main function to validate endpoint coverage."""
print(
"🔍 Checking endpoint coverage between sidebars.js and provider_endpoints_support.json..."
)
has_errors = False
# Load provider_endpoints_support.json
data = load_provider_endpoints_file()
defined_endpoints = get_defined_endpoints(data)
# Test 1: Check that endpoints from sidebars.js have docs_label entries
print("\n📖 Test 1: Checking endpoints from sidebars.js...")
sidebar_endpoints = extract_endpoints_from_sidebars()
print(f"✓ Found {len(sidebar_endpoints)} endpoints in sidebars.js")
print(
f"✓ Found {len(defined_endpoints)} endpoint definitions in provider_endpoints_support.json"
)
# Check for missing endpoints
missing_endpoints = []
# Collect all docs_label values from defined endpoints
defined_docs_labels = set()
for endpoint_data in defined_endpoints.values():
if isinstance(endpoint_data, dict) and "docs_label" in endpoint_data:
defined_docs_labels.add(endpoint_data["docs_label"])
for sidebar_key, sidebar_label in sorted(sidebar_endpoints.items()):
# Generate variations for matching against docs_label
variations = normalize_endpoint_key(sidebar_key)
# Check if any variation exists in docs_label values
if not any(var in defined_docs_labels for var in variations):
missing_endpoints.append((sidebar_key, sidebar_label))
# Report missing endpoints from sidebars
if missing_endpoints:
has_errors = True
error_msg = "\n❌ ERROR: The following endpoints are in sidebars.js but missing from provider_endpoints_support.json:\n"
error_msg += "=" * 70 + "\n"
for key, label in missing_endpoints:
error_msg += f" - {key}\n"
error_msg += f' Label in sidebars.js: "{label}"\n'
error_msg += "\n" + "=" * 70 + "\n"
error_msg += f"\n💡 To fix: Add these {len(missing_endpoints)} endpoint(s) to the 'endpoints' object\n"
error_msg += " in provider_endpoints_support.json\n"
error_msg += "\nExample format:\n"
error_msg += ' "endpoints": {\n'
for key, label in missing_endpoints[:5]:
error_msg += f' "{key}": {{\n'
error_msg += f' "docs_label": "{label}",\n'
error_msg += f' "provider_json_field": "{key}",\n'
error_msg += f' "description": "Description of the {label} endpoint"\n'
error_msg += " },\n"
if len(missing_endpoints) > 5:
error_msg += " ...\n"
error_msg += " }\n"
print(error_msg)
else:
print(
f"✅ All {len(sidebar_endpoints)} endpoints from sidebars.js are defined!"
)
# Test 2: Check that all provider endpoint keys have provider_json_field entries
print("\n📋 Test 2: Checking provider endpoint keys...")
missing_provider_keys = check_provider_endpoint_keys(data)
if missing_provider_keys:
has_errors = True
error_msg = "\n❌ ERROR: The following endpoint keys are used in providers but missing provider_json_field definitions:\n"
error_msg += "=" * 70 + "\n"
for key in missing_provider_keys:
# Find which providers use this key
using_providers = []
for provider_name, provider_data in data.get("providers", {}).items():
if key in provider_data.get("endpoints", {}):
using_providers.append(provider_name)
error_msg += f" - {key}\n"
error_msg += f" Used by {len(using_providers)} provider(s): {', '.join(using_providers[:3])}"
if len(using_providers) > 3:
error_msg += f" and {len(using_providers) - 3} more"
error_msg += "\n"
error_msg += "\n" + "=" * 70 + "\n"
error_msg += f"\n💡 To fix: Add these {len(missing_provider_keys)} endpoint(s) to the 'endpoints' object\n"
error_msg += " in provider_endpoints_support.json with 'provider_json_field' matching the key\n"
error_msg += "\nExample format:\n"
error_msg += ' "endpoints": {\n'
for key in missing_provider_keys[:3]:
error_msg += f' "{key}": {{\n'
error_msg += f' "docs_label": "{key}",\n'
error_msg += f' "provider_json_field": "{key}",\n'
error_msg += f' "description": "Description of the {key} endpoint"\n'
error_msg += " },\n"
if len(missing_provider_keys) > 3:
error_msg += " ...\n"
error_msg += " }\n"
print(error_msg)
else:
print("✅ All provider endpoint keys have provider_json_field definitions!")
# Test 3: Check that all defined endpoints are used by at least one provider
print("\n🔍 Test 3: Checking for unused endpoint definitions...")
unused_endpoints = check_unused_endpoints(data)
if unused_endpoints:
has_errors = True
error_msg = "\n⚠️ WARNING: The following endpoint definitions are not used by any provider:\n"
error_msg += "=" * 70 + "\n"
for endpoint_key, provider_json_field in unused_endpoints:
endpoint_data = defined_endpoints.get(endpoint_key, {})
docs_label = endpoint_data.get("docs_label", "N/A")
error_msg += f" - {endpoint_key}\n"
error_msg += f" provider_json_field: '{provider_json_field}'\n"
error_msg += f" docs_label: '{docs_label}'\n"
error_msg += "\n" + "=" * 70 + "\n"
error_msg += f"\n💡 These {len(unused_endpoints)} endpoint(s) are defined but not used by any provider.\n"
error_msg += " Either:\n"
error_msg += (
" 1. Add the endpoint to relevant providers' 'endpoints' objects, OR\n"
)
error_msg += " 2. Remove the endpoint definition if it's no longer needed\n"
print(error_msg)
else:
print("✅ All endpoint definitions are used by at least one provider!")
# Raise error if any tests failed
if has_errors:
error_summary = []
if missing_endpoints:
error_summary.append(f"{len(missing_endpoints)} endpoints from sidebars.js")
if missing_provider_keys:
error_summary.append(f"{len(missing_provider_keys)} provider endpoint keys")
if unused_endpoints:
error_summary.append(f"{len(unused_endpoints)} unused endpoint definitions")
raise MissingEndpointDefinitionError(
f"Endpoint validation failed: Missing definitions for {' and '.join(error_summary)}"
)
print("\n🎉 All endpoint coverage validations passed!")
return 0
if __name__ == "__main__":
try:
sys.exit(main())
except MissingEndpointDefinitionError as e:
print(f"\n🚨 CRITICAL ERROR: {e}\n")
sys.exit(1)
except Exception as e:
print(f"\n🚨 UNEXPECTED ERROR: {e}\n")
import traceback
traceback.print_exc()
sys.exit(1)

View file

@ -0,0 +1,294 @@
"""
Code coverage test to ensure all provider folders are documented.
This script validates that:
1. Every provider folder in litellm/llms/ has a corresponding entry in provider_endpoints_support.json
2. Every provider in litellm/llms/openai_like/providers.json is documented in provider_endpoints_support.json
"""
import json
import os
import sys
from pathlib import Path
from typing import Dict, List, Set, Tuple
class UndocumentedProviderError(Exception):
"""Raised when providers are found without documentation."""
pass
# Special folders that should be excluded from validation
EXCLUDED_FOLDERS = {
"__pycache__",
"base_llm",
"deprecated_providers",
"custom_httpx",
"pass_through",
"openai_like", # This is a generic handler, not a specific provider
"aiohttp_openai", # Internal implementation detail for async HTTP
}
def get_repo_root() -> Path:
"""Get the repository root directory."""
# Check if litellm directory exists in current working directory
cwd = Path.cwd()
if (cwd / "litellm").exists() and (cwd / "litellm").is_dir():
# We're already at the repo root
return cwd
# Otherwise, navigate up from script location
current = Path(__file__).resolve()
# Navigate up from tests/code_coverage_tests/
return current.parent.parent.parent
def get_llm_provider_folders() -> Set[str]:
"""Get all provider folder names from litellm/llms directory."""
repo_root = get_repo_root()
llms_dir = repo_root / "litellm" / "llms"
if not llms_dir.exists():
print(f"❌ ERROR: Could not find llms directory at {llms_dir}")
sys.exit(1)
folders = set()
for item in llms_dir.iterdir():
if item.is_dir() and item.name not in EXCLUDED_FOLDERS:
folders.add(item.name)
return folders
def load_provider_endpoints_file() -> Dict:
"""Load the provider_endpoints_support.json file."""
repo_root = get_repo_root()
file_path = repo_root / "provider_endpoints_support.json"
if not file_path.exists():
print(
f"❌ ERROR: Could not find provider_endpoints_support.json at {file_path}"
)
sys.exit(1)
with open(file_path, "r") as f:
return json.load(f)
def get_openai_like_providers() -> Set[str]:
"""Get all provider names from litellm/llms/openai_like/providers.json."""
repo_root = get_repo_root()
providers_file = repo_root / "litellm" / "llms" / "openai_like" / "providers.json"
if not providers_file.exists():
print(
f"⚠️ WARNING: Could not find openai_like/providers.json at {providers_file}"
)
return set()
with open(providers_file, "r") as f:
data = json.load(f)
# Return all provider keys from the JSON
return set(data.keys())
def get_documented_providers(data: Dict) -> Set[str]:
"""Get all provider slugs documented in provider_endpoints_support.json."""
providers = data.get("providers", {})
# Get all provider keys, including those with slashes
documented = set()
for provider_key in providers.keys():
# For providers like "azure_ai/doc-intelligence", extract base name
base_name = provider_key.split("/")[0]
documented.add(base_name)
# Also add the full key in case folder name matches exactly
documented.add(provider_key)
return documented
def normalize_provider_name(folder_name: str) -> Set[str]:
"""
Generate possible provider names that might match a folder.
Some folders might have variations in the JSON:
- github_copilot folder -> github_copilot provider
- azure folder -> azure, azure_text, azure_ai providers
"""
variations = {folder_name}
# Add common variations
if "_" in folder_name:
# Try without underscores (though less common)
variations.add(folder_name.replace("_", ""))
return variations
def main():
"""Main function to validate provider documentation."""
print("🔍 Checking that all providers are documented...")
has_errors = False
# Check 1: Provider folders in litellm/llms
print("\n📁 Checking provider folders in litellm/llms/...")
provider_folders = get_llm_provider_folders()
print(f"✓ Found {len(provider_folders)} provider folders")
# Check 2: OpenAI-like providers
print("\n📋 Checking openai_like providers...")
openai_like_providers = get_openai_like_providers()
print(f"✓ Found {len(openai_like_providers)} openai_like providers")
# Load the JSON file
data = load_provider_endpoints_file()
documented_providers = get_documented_providers(data)
print(
f"\n✓ Found {len(data.get('providers', {}))} provider entries in provider_endpoints_support.json"
)
# Check for undocumented folders
undocumented_folders = []
for folder in sorted(provider_folders):
# Check if folder name or any variation is documented
variations = normalize_provider_name(folder)
if not any(var in documented_providers for var in variations):
undocumented_folders.append(folder)
# Check for undocumented openai_like providers
undocumented_openai_like = []
for provider in sorted(openai_like_providers):
# Generate multiple possible variations of the provider name
variations = {
provider, # Original name (e.g., "nano-gpt")
provider.replace(
"-", "_"
), # Replace hyphens with underscores (e.g., "nano_gpt")
provider.replace("-", ""), # Remove hyphens (e.g., "nanogpt")
provider.replace("_", ""), # Remove underscores
}
# Special case mappings for known variations
special_mappings = {
"veniceai": "venice",
"nano-gpt": "nanogpt",
}
if provider in special_mappings:
variations.add(special_mappings[provider])
# Check if any variation is documented
if not any(var in documented_providers for var in variations):
undocumented_openai_like.append(provider)
# Collect all error messages
error_messages: List[str] = []
# Report errors for undocumented folders
if undocumented_folders:
has_errors = True
error_msg = "\n❌ ERROR: The following provider folders are not documented:\n"
error_msg += "=" * 70 + "\n"
for folder in undocumented_folders:
error_msg += f" - litellm/llms/{folder}/\n"
error_msg += "\n" + "=" * 70 + "\n"
error_msg += f"\n💡 To fix: Add entries for these {len(undocumented_folders)} provider(s)\n"
error_msg += (
" in the 'providers' section of provider_endpoints_support.json\n"
)
error_msg += "\nExample format:\n"
error_msg += ' "providers": {\n'
for folder in undocumented_folders[:3]:
error_msg += f' "{folder}": {{\n'
error_msg += f' "display_name": "{folder.replace("_", " ").title()} (`{folder}`)",\n'
error_msg += (
f' "url": "https://docs.litellm.ai/docs/providers/{folder}",\n'
)
error_msg += ' "endpoints": {\n'
error_msg += ' "chat_completions": true,\n'
error_msg += ' "messages": true,\n'
error_msg += ' "responses": true,\n'
error_msg += ' "embeddings": false,\n'
error_msg += " ...\n"
error_msg += " }\n"
error_msg += " },\n"
if len(undocumented_folders) > 3:
error_msg += " ...\n"
error_msg += " }\n"
print(error_msg)
error_messages.append(
f"Found {len(undocumented_folders)} undocumented provider folders: {', '.join(undocumented_folders)}"
)
# Report errors for undocumented openai_like providers
if undocumented_openai_like:
has_errors = True
error_msg = (
"\n❌ ERROR: The following openai_like providers are not documented:\n"
)
error_msg += "=" * 70 + "\n"
for provider in undocumented_openai_like:
error_msg += f" - {provider}\n"
error_msg += "\n" + "=" * 70 + "\n"
error_msg += f"\n💡 To fix: Add entries for these {len(undocumented_openai_like)} provider(s)\n"
error_msg += (
" in the 'providers' section of provider_endpoints_support.json\n"
)
error_msg += "\nExample format:\n"
error_msg += ' "providers": {\n'
for provider in undocumented_openai_like[:3]:
normalized = provider.replace("-", "_")
error_msg += f' "{normalized}": {{\n'
error_msg += f' "display_name": "{provider.replace("-", " ").replace("_", " ").title()} (`{normalized}`)",\n'
error_msg += (
f' "url": "https://docs.litellm.ai/docs/providers/{normalized}",\n'
)
error_msg += ' "endpoints": {\n'
error_msg += ' "chat_completions": true,\n'
error_msg += ' "messages": true,\n'
error_msg += ' "responses": true,\n'
error_msg += ' "embeddings": false,\n'
error_msg += " ...\n"
error_msg += " }\n"
error_msg += " },\n"
if len(undocumented_openai_like) > 3:
error_msg += " ...\n"
error_msg += " }\n"
print(error_msg)
error_messages.append(
f"Found {len(undocumented_openai_like)} undocumented openai_like providers: {', '.join(undocumented_openai_like)}"
)
# Raise exception if there are any errors
if has_errors:
error_summary = " AND ".join(error_messages)
raise UndocumentedProviderError(
f"Provider documentation validation failed: {error_summary}"
)
print(f"\n✅ All {len(provider_folders)} provider folders are documented!")
print(f"✅ All {len(openai_like_providers)} openai_like providers are documented!")
print("\n🎉 All provider documentation checks passed!")
return 0
if __name__ == "__main__":
try:
sys.exit(main())
except UndocumentedProviderError as e:
print(f"\n🚨 CRITICAL ERROR: {e}\n")
sys.exit(1)
except Exception as e:
print(f"\n🚨 UNEXPECTED ERROR: {e}\n")
import traceback
traceback.print_exc()
sys.exit(1)

View file

@ -138,4 +138,5 @@ pondpond: >=1.4.1 # Apache 2.0 License
fastuuid: >=0.13.0 # BSD-3-Clause license
llm-sandbox: >=0.3.31 # MIT License - https://github.com/vndee/llm-sandbox
nodejs-wheel-binaries: >=24.12.0 # MIT license manually verified
grpcio: >=1.69.0 # Apache License 2.0

View file

@ -143,6 +143,23 @@ class TestOpenAIImageEditDallE2(BaseLLMImageEditTest):
}
class TestAzureAIFlux2ImageEdit(BaseLLMImageEditTest):
"""
Concrete implementation of BaseLLMImageEditTest for Azure AI FLUX 2 image edits.
FLUX 2 uses JSON with base64 image instead of multipart/form-data.
"""
def get_base_image_edit_call_args(self) -> dict:
"""Return base call args for Azure AI FLUX 2 image edit"""
return {
"model": "azure_ai/flux.2-pro",
"image": SINGLE_TEST_IMAGE,
"api_base": os.getenv("AZURE_AI_API_BASE", "https://litellm-ci-cd-prod.services.ai.azure.com"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": "preview",
}
@pytest.mark.flaky(retries=3, delay=2)
@pytest.mark.asyncio
async def test_openai_image_edit_litellm_router():
@ -322,14 +339,23 @@ async def test_openai_image_edit_cost_tracking():
litellm.logging_callback_manager._reset_all_callbacks()
litellm.callbacks = [test_custom_logger]
# Mock response for Azure image edit
# Mock response for Azure image edit with usage data for cost tracking
mock_response = {
"created": 1589478378,
"data": [
{
"b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
}
]
],
"usage": {
"total_tokens": 1100,
"input_tokens": 100,
"input_tokens_details": {
"image_tokens": 50,
"text_tokens": 50
},
"output_tokens": 1000
}
}
class MockResponse:
@ -401,14 +427,23 @@ async def test_azure_image_edit_cost_tracking():
litellm.logging_callback_manager._reset_all_callbacks()
litellm.callbacks = [test_custom_logger]
# Mock response for Azure image edit
# Mock response for Azure image edit with usage data for cost tracking
mock_response = {
"created": 1589478378,
"data": [
{
"b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
}
]
],
"usage": {
"total_tokens": 1100,
"input_tokens": 100,
"input_tokens_details": {
"image_tokens": 50,
"text_tokens": 50
},
"output_tokens": 1000
}
}
class MockResponse:

View file

@ -322,10 +322,7 @@ def process_stream_response(res, messages):
return res
@pytest.mark.skipif(
os.environ.get("CIRCLE_OIDC_TOKEN_V2") is None,
reason="Cannot run without being in CircleCI Runner",
)
@pytest.mark.skip(reason="Cannot run without being in CircleCI Runner")
def test_completion_bedrock_claude_aws_session_token(bedrock_session_token_creds):
print("\ncalling bedrock claude with aws_session_token auth")
@ -406,10 +403,7 @@ def test_completion_bedrock_claude_aws_session_token(bedrock_session_token_creds
pytest.fail(f"Error occurred: {e}")
@pytest.mark.skipif(
os.environ.get("CIRCLE_OIDC_TOKEN_V2") is None,
reason="Cannot run without being in CircleCI Runner",
)
@pytest.mark.skip(reason="Cannot run without being in CircleCI Runner")
def test_completion_bedrock_claude_aws_bedrock_client(bedrock_session_token_creds):
print("\ncalling bedrock claude with aws_session_token auth")

View file

@ -71,6 +71,7 @@ def test_get_arize_config(mock_env_vars):
assert config.api_key == "test_api_key"
assert config.endpoint == "https://otlp.arize.com/v1"
assert config.protocol == "otlp_grpc"
assert config.project_name is None
def test_get_arize_config_with_endpoints(mock_env_vars, monkeypatch):
@ -79,10 +80,12 @@ def test_get_arize_config_with_endpoints(mock_env_vars, monkeypatch):
"""
monkeypatch.setenv("ARIZE_ENDPOINT", "grpc://test.endpoint")
monkeypatch.setenv("ARIZE_HTTP_ENDPOINT", "http://test.endpoint")
monkeypatch.setenv("ARIZE_PROJECT_NAME", "custom-project")
config = ArizeLogger.get_arize_config()
assert config.endpoint == "grpc://test.endpoint"
assert config.protocol == "otlp_grpc"
assert config.project_name == "custom-project"
@pytest.mark.skip(

View file

@ -3059,7 +3059,6 @@ def response_format_tests(response: litellm.ModelResponse):
"bedrock/cohere.command-r-plus-v1:0",
"anthropic.claude-3-sonnet-20240229-v1:0",
"mistral.mistral-7b-instruct-v0:2",
# "bedrock/amazon.titan-tg1-large",
"meta.llama3-8b-instruct-v1:0",
],
)
@ -3101,31 +3100,6 @@ async def test_completion_bedrock_httpx_models(sync_mode, model):
pytest.fail(f"An error occurred - {str(e)}")
def test_completion_bedrock_titan_null_response():
try:
# amazon.titan-text-lite-v1 is deprecated, using titan-text-express-v1 instead
response = completion(
model="bedrock/amazon.titan-text-express-v1",
messages=[
{
"role": "user",
"content": "Hello!",
},
{
"role": "assistant",
"content": "Hello! How can I help you?",
},
{
"role": "user",
"content": "What model are you?",
},
],
)
# Add any assertions here to check the response
print(f"response: {response}")
except Exception as e:
pytest.fail(f"An error occurred - {str(e)}")
# test_completion_bedrock_titan()
@ -3916,26 +3890,7 @@ async def test_dynamic_azure_params(stream, sync_mode):
raise e
@pytest.mark.asyncio()
@pytest.mark.flaky(retries=3, delay=1)
async def test_completion_ai21_chat():
litellm.set_verbose = True
try:
response = await litellm.acompletion(
model="ai21_chat/jamba-mini",
user="ishaan",
tool_choice="auto",
seed=123,
messages=[{"role": "user", "content": "what does the document say"}],
documents=[
{
"content": "hello world",
"metadata": {"source": "google", "author": "ishaan"},
}
],
)
except litellm.InternalServerError:
pytest.skip("Model is overloaded")
@pytest.mark.parametrize(

View file

@ -552,36 +552,6 @@ async def test_completion_predibase_streaming(sync_mode):
pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio()
@pytest.mark.flaky(retries=3, delay=1)
async def test_completion_ai21_stream():
litellm.set_verbose = True
response = await litellm.acompletion(
model="ai21_chat/jamba-mini",
user="ishaan",
stream=True,
seed=123,
messages=[{"role": "user", "content": "hi my name is ishaan"}],
)
complete_response = ""
idx = 0
async for init_chunk in response:
chunk, finished = streaming_format_tests(idx, init_chunk)
complete_response += chunk
custom_llm_provider = init_chunk._hidden_params["custom_llm_provider"]
print(f"custom_llm_provider: {custom_llm_provider}")
assert custom_llm_provider == "ai21_chat"
idx += 1
if finished:
assert isinstance(init_chunk.choices[0], litellm.utils.StreamingChoices)
break
if complete_response.strip() == "":
raise Exception("Empty response received")
print(f"complete_response: {complete_response}")
pass
def test_completion_azure_function_calling_stream():
try:
@ -1318,7 +1288,6 @@ async def test_completion_replicate_llama3_streaming(sync_mode):
# ["bedrock/cohere.command-r-plus-v1:0", None],
["anthropic.claude-3-sonnet-20240229-v1:0", None],
# ["mistral.mistral-7b-instruct-v0:2", None],
["bedrock/amazon.titan-tg1-large", None],
# ["meta.llama3-8b-instruct-v1:0", None],
],
)

View file

@ -40,6 +40,7 @@ ignored_keys = [
"metadata.usage_object",
"metadata.cold_storage_object_key",
"metadata.litellm_overhead_time_ms",
"metadata.cost_breakdown",
]

View file

@ -138,64 +138,6 @@ def validate_raw_gen_ai_request_openai_streaming(span):
assert span._attributes[attr] is not None, f"Attribute {attr} has None"
@pytest.mark.parametrize(
"model",
["anthropic/claude-3-opus-20240229"],
)
@pytest.mark.flaky(retries=6, delay=2)
def test_completion_claude_3_function_call_with_otel(model):
litellm.set_verbose = True
litellm.callbacks = [OpenTelemetry(config=OpenTelemetryConfig(exporter=exporter))]
tools = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
},
"required": ["location"],
},
},
}
]
messages = [
{
"role": "user",
"content": "What's the weather like in Boston today in Fahrenheit?",
}
]
try:
# test without max tokens
response = litellm.completion(
model=model,
messages=messages,
tools=tools,
tool_choice={
"type": "function",
"function": {"name": "get_current_weather"},
},
drop_params=True,
)
print("response from LiteLLM", response)
except litellm.InternalServerError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
finally:
# clear in memory exporter
exporter.clear()
@pytest.mark.asyncio
@pytest.mark.parametrize("streaming", [True, False])
@pytest.mark.parametrize("global_redact", [True, False])

View file

@ -21,7 +21,7 @@ class TestAnthropicMessagesEndpoint(BaseAnthropicMessagesTest):
def test_anthropic_messages_to_wildcard_model(self):
client = self.get_client()
response = client.messages.create(
model="anthropic/claude-3-opus-20240229",
model="anthropic/claude-haiku-4-5-20251001",
messages=[{"role": "user", "content": "Hello, world!"}],
max_tokens=100,
)

View file

@ -3504,6 +3504,7 @@ async def test_list_keys(prisma_client):
include_created_by_keys=False,
sort_by=None,
sort_order="desc",
expand=None,
)
print("response=", response)
assert "keys" in response
@ -3528,6 +3529,7 @@ async def test_list_keys(prisma_client):
include_created_by_keys=False,
sort_by=None,
sort_order="desc",
expand=None,
)
print("pagination response=", response)
assert len(response["keys"]) == 2
@ -3568,6 +3570,7 @@ async def test_list_keys(prisma_client):
include_created_by_keys=False,
sort_by=None,
sort_order="desc",
expand=None,
)
print("filtered user_id response=", response)
assert len(response["keys"]) == 1
@ -3589,6 +3592,7 @@ async def test_list_keys(prisma_client):
include_created_by_keys=False,
sort_by=None,
sort_order="desc",
expand=None,
)
assert len(response["keys"]) == 1
assert _key in response["keys"]

View file

@ -73,7 +73,7 @@ def test_routing_strategy_init(model_list):
from litellm.types.router import RoutingStrategy
router = Router(model_list=model_list)
for strategy in RoutingStrategy._member_names_:
for strategy in RoutingStrategy:
router.routing_strategy_init(
routing_strategy=strategy, routing_strategy_args={}
)

View file

@ -134,80 +134,6 @@ class TestContainerAPI:
assert response.id == "cntr_async_123"
assert response.name == "Async Test Container"
def test_list_containers_basic(self):
"""Test basic container listing functionality."""
mock_response = ContainerListResponse(
object="list",
data=[
ContainerObject(
id="cntr_1",
object="container",
created_at=1747857508,
status="running",
expires_after={"anchor": "last_active_at", "minutes": 20},
last_active_at=1747857508,
name="Container 1"
),
ContainerObject(
id="cntr_2",
object="container",
created_at=1747857600,
status="running",
expires_after={"anchor": "last_active_at", "minutes": 15},
last_active_at=1747857600,
name="Container 2"
)
],
first_id="cntr_1",
last_id="cntr_2",
has_more=False
)
with patch('litellm.containers.main.base_llm_http_handler') as mock_handler:
mock_handler.container_list_handler.return_value = mock_response
response = list_containers(
custom_llm_provider="openai"
)
assert isinstance(response, ContainerListResponse)
assert len(response.data) == 2
assert response.data[0].id == "cntr_1"
assert response.data[1].id == "cntr_2"
assert response.has_more == False
def test_list_containers_with_params(self):
"""Test container listing with parameters."""
mock_response = ContainerListResponse(
object="list",
data=[
ContainerObject(
id="cntr_limited",
object="container",
created_at=1747857508,
status="running",
expires_after={"anchor": "last_active_at", "minutes": 20},
last_active_at=1747857508,
name="Limited Container"
)
],
first_id="cntr_limited",
last_id="cntr_limited",
has_more=True
)
with patch('litellm.containers.main.base_llm_http_handler') as mock_handler:
mock_handler.container_list_handler.return_value = mock_response
response = list_containers(
limit=1,
order="desc",
after="cntr_prev",
custom_llm_provider="openai"
)
assert len(response.data) == 1
assert response.has_more == True
@pytest.mark.asyncio
async def test_alist_containers_basic(self):

View file

@ -123,7 +123,8 @@ class TestArizeIntegrationWithProxy:
with patch.dict(os.environ, {
"ARIZE_SPACE_KEY": "test-space-123",
"ARIZE_API_KEY": "test-api-456",
"ARIZE_ENDPOINT": "https://custom.arize.com/v1"
"ARIZE_ENDPOINT": "https://custom.arize.com/v1",
"ARIZE_PROJECT_NAME": "custom-project",
}):
config = ArizeLogger.get_arize_config()
@ -131,13 +132,15 @@ class TestArizeIntegrationWithProxy:
assert config.api_key == "test-api-456"
assert config.endpoint == "https://custom.arize.com/v1"
assert config.protocol == "otlp_grpc"
assert config.project_name == "custom-project"
def test_arize_get_config_defaults(self):
"""Test ArizeLogger.get_arize_config() with default endpoint."""
with patch.dict(os.environ, {
"ARIZE_SPACE_KEY": "test-space-default",
"ARIZE_API_KEY": "test-api-default"
"ARIZE_API_KEY": "test-api-default",
"ARIZE_PROJECT_NAME": "default-project",
}, clear=True):
config = ArizeLogger.get_arize_config()
@ -145,6 +148,7 @@ class TestArizeIntegrationWithProxy:
assert config.api_key == "test-api-default"
assert config.endpoint == "https://otlp.arize.com/v1" # Default endpoint
assert config.protocol == "otlp_grpc" # Default protocol
assert config.project_name == "default-project"
def test_arize_construct_dynamic_headers(self):
"""Test dynamic OTEL headers construction for team/key logging."""
@ -180,4 +184,4 @@ class TestArizeIntegrationWithProxy:
if __name__ == "__main__":
pytest.main([__file__, "-v"])
pytest.main([__file__, "-v"])

View file

@ -1,18 +1,19 @@
import os
import sys
from unittest.mock import MagicMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path
from litellm.integrations.gitlab.gitlab_client import GitLabClient
from litellm.integrations.gitlab.gitlab_prompt_manager import (
GitLabPromptCache,
GitLabPromptManager,
GitLabPromptTemplate,
GitLabTemplateManager,
GitLabPromptCache,
encode_prompt_id,
decode_prompt_id,
encode_prompt_id,
)
# -----------------------
@ -817,22 +818,3 @@ def test_cache_get_by_file_returns_exact_entry(mock_pm_cls, fake_managers):
assert beta and beta["id"] == "nested/beta"
@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabPromptManager")
def test_encode_decode_helpers_roundtrip_in_cache_context(mock_pm_cls, fake_managers):
tm, wrapper = fake_managers
tm._discoverable_ids = ["dir1/dir2/item"]
mock_pm_cls.return_value = wrapper
cache = GitLabPromptCache({"project": "g/s/r", "access_token": "tkn"})
cache.load_all()
encoded = encode_prompt_id("dir1/dir2/item")
assert encoded in cache.list_ids()
# decode → encode → lookup should still work
decoded = decode_prompt_id(encoded)
assert decoded == "dir1/dir2/item"
got = cache.get_by_id(decoded)
assert got is not None
assert got["id"] == "dir1/dir2/item"

View file

@ -9,10 +9,10 @@ from litellm.integrations.opentelemetry import OpenTelemetryConfig
# Try to import OpenTelemetry packages, skip tests if not available
try:
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
InMemorySpanExporter,
)
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
OPENTELEMETRY_AVAILABLE = True
except ImportError:
@ -150,53 +150,6 @@ class TestLevoConfig(unittest.TestCase):
class TestLevoIntegration(unittest.TestCase):
"""Integration tests for LevoLogger."""
@patch.dict(
"os.environ",
{
"LEVOAI_API_KEY": "test-api-key",
"LEVOAI_ORG_ID": "test-org-id",
"LEVOAI_WORKSPACE_ID": "test-workspace-id",
"LEVOAI_COLLECTOR_URL": "https://collector.levo.ai",
},
)
@pytest.mark.skipif(
not OPENTELEMETRY_AVAILABLE, reason="OpenTelemetry packages not installed"
)
@patch(
"litellm.integrations.opentelemetry.OpenTelemetry._init_otel_logger_on_litellm_proxy"
)
def test_levo_logger_instantiation(self, mock_init_proxy):
"""Test that LevoLogger can be instantiated with proper config."""
# Mock the proxy initialization to avoid importing proxy code
mock_init_proxy.return_value = None
config = LevoLogger.get_levo_config()
otel_config = OpenTelemetryConfig(
exporter=config.protocol,
endpoint=config.endpoint,
headers=config.otlp_auth_headers,
)
# Create a tracer provider with in-memory exporter to avoid requiring OTLP packages
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
# Create LevoLogger instance with mocked tracer provider
levo_logger = LevoLogger(
config=otel_config, callback_name="levo", tracer_provider=tracer_provider
)
# Verify it's an instance of OpenTelemetry
self.assertIsInstance(levo_logger, LevoLogger)
# Check it extends OpenTelemetry by checking base classes
from litellm.integrations.opentelemetry import OpenTelemetry
self.assertIsInstance(levo_logger, OpenTelemetry)
# Verify callback_name is set
self.assertEqual(levo_logger.callback_name, "levo")
@patch.dict(
"os.environ",
{

View file

@ -258,6 +258,22 @@ class TestOpenTelemetry(unittest.TestCase):
MODEL = "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0"
HERE = os.path.dirname(__file__)
@patch.dict(os.environ, {}, clear=True)
def test_open_telemetry_config_manual_defaults(self):
"""Manual OpenTelemetryConfig creation should populate default identifiers."""
config = OpenTelemetryConfig(exporter="console", endpoint="http://collector")
self.assertEqual(config.service_name, "litellm")
self.assertEqual(config.deployment_environment, "production")
self.assertEqual(config.model_id, "litellm")
@patch.dict(os.environ, {}, clear=True)
def test_open_telemetry_config_custom_service_name(self):
"""Model ID should inherit provided service name when not explicitly set."""
config = OpenTelemetryConfig(service_name="custom-service", exporter="console")
self.assertEqual(config.service_name, "custom-service")
self.assertEqual(config.deployment_environment, "production")
self.assertEqual(config.model_id, "custom-service")
def wait_for_spans(self, exporter: InMemorySpanExporter, prefix: str):
"""Poll until we see at least one span with an attribute key starting with `prefix`."""
deadline = time.time() + self.POLL_TIMEOUT
@ -504,8 +520,6 @@ class TestOpenTelemetry(unittest.TestCase):
self, mock_detector_cls, mock_resource_create
):
"""Test _get_litellm_resource with default values when no environment variables are set."""
from litellm.integrations.opentelemetry import _get_litellm_resource
# Mock the Resource.create method
mock_base_resource = MagicMock()
mock_resource_create.return_value = mock_base_resource
@ -520,8 +534,8 @@ class TestOpenTelemetry(unittest.TestCase):
mock_merged_resource = MagicMock()
mock_base_resource.merge.return_value = mock_merged_resource
# Call the function
result = _get_litellm_resource()
config = OpenTelemetryConfig()
result = OpenTelemetry._get_litellm_resource(config)
# Verify Resource.create was called with correct default attributes
expected_attributes = {
@ -549,8 +563,6 @@ class TestOpenTelemetry(unittest.TestCase):
self, mock_detector_cls, mock_resource_create
):
"""Test _get_litellm_resource with LiteLLM-specific environment variables."""
from litellm.integrations.opentelemetry import _get_litellm_resource
# Mock the Resource.create method
mock_base_resource = MagicMock()
mock_resource_create.return_value = mock_base_resource
@ -565,8 +577,8 @@ class TestOpenTelemetry(unittest.TestCase):
mock_merged_resource = MagicMock()
mock_base_resource.merge.return_value = mock_merged_resource
# Call the function
result = _get_litellm_resource()
config = OpenTelemetryConfig.from_env()
result = OpenTelemetry._get_litellm_resource(config)
# Verify Resource.create was called with environment variable values
expected_attributes = {
@ -593,8 +605,6 @@ class TestOpenTelemetry(unittest.TestCase):
self, mock_detector_cls, mock_resource_create
):
"""Test _get_litellm_resource with OTEL_RESOURCE_ATTRIBUTES environment variable."""
from litellm.integrations.opentelemetry import _get_litellm_resource
# Mock the Resource.create method to simulate the actual behavior
# In reality, Resource.create() would parse OTEL_RESOURCE_ATTRIBUTES and merge it
mock_base_resource = MagicMock()
@ -610,8 +620,8 @@ class TestOpenTelemetry(unittest.TestCase):
mock_merged_resource = MagicMock()
mock_base_resource.merge.return_value = mock_merged_resource
# Call the function
result = _get_litellm_resource()
config = OpenTelemetryConfig.from_env()
result = OpenTelemetry._get_litellm_resource(config)
# Verify Resource.create was called with the base attributes
# The actual OTEL_RESOURCE_ATTRIBUTES parsing is handled by OpenTelemetry SDK
@ -628,10 +638,8 @@ class TestOpenTelemetry(unittest.TestCase):
@patch.dict(os.environ, {}, clear=True)
def test_get_litellm_resource_integration_with_real_resource(self):
"""Integration test to verify _get_litellm_resource works with actual OpenTelemetry Resource."""
from litellm.integrations.opentelemetry import _get_litellm_resource
# This test uses the real OpenTelemetry Resource.create() method
result = _get_litellm_resource()
config = OpenTelemetryConfig()
result = OpenTelemetry._get_litellm_resource(config)
# Verify the result is a Resource instance
from opentelemetry.sdk.resources import Resource
@ -653,10 +661,8 @@ class TestOpenTelemetry(unittest.TestCase):
)
def test_get_litellm_resource_real_otel_resource_attributes(self):
"""Integration test to verify OTEL_RESOURCE_ATTRIBUTES is properly handled."""
from litellm.integrations.opentelemetry import _get_litellm_resource
# This test uses the real OpenTelemetry Resource.create() method
result = _get_litellm_resource()
config = OpenTelemetryConfig.from_env()
result = OpenTelemetry._get_litellm_resource(config)
print("RESULT", result)
@ -683,10 +689,8 @@ class TestOpenTelemetry(unittest.TestCase):
)
def test_get_litellm_resource_precedence(self):
"""Test that OTEL_SERVICE_NAME takes precedence over OTEL_RESOURCE_ATTRIBUTES according to OpenTelemetry spec."""
from litellm.integrations.opentelemetry import _get_litellm_resource
# This test verifies the OpenTelemetry standard behavior
result = _get_litellm_resource()
config = OpenTelemetryConfig.from_env()
result = OpenTelemetry._get_litellm_resource(config)
# Verify the result is a Resource instance
from opentelemetry.sdk.resources import Resource

View file

@ -0,0 +1,161 @@
"""
Unit tests for Prometheus invalid API key request filtering.
Tests functionality that prevents invalid API key requests (401 status codes)
from being recorded in Prometheus metrics.
"""
import os
import sys
from unittest.mock import Mock, patch
import pytest
from prometheus_client import REGISTRY
sys.path.insert(0, os.path.abspath("../../.."))
from litellm.integrations.prometheus import PrometheusLogger
from litellm.proxy._types import UserAPIKeyAuth
@pytest.fixture(scope="function")
def prometheus_logger():
"""Create a PrometheusLogger instance for testing."""
collectors = list(REGISTRY._collector_to_names.keys())
for collector in collectors:
REGISTRY.unregister(collector)
return PrometheusLogger()
class ExceptionWithCode:
"""Exception-like object with 'code' attribute (ProxyException pattern)."""
def __init__(self, code):
self.code = code
class ExceptionWithStatusCode:
"""Exception-like object with 'status_code' attribute."""
def __init__(self, status_code):
self.status_code = status_code
class TestExtractStatusCode:
"""Test status code extraction from various sources."""
@pytest.mark.parametrize("exception_class,code_value,expected", [
(ExceptionWithCode, "401", 401),
(ExceptionWithStatusCode, 401, 401),
])
def test_extract_from_exception(self, prometheus_logger, exception_class, code_value, expected):
exception = exception_class(code_value)
assert prometheus_logger._extract_status_code(exception=exception) == expected
def test_extract_from_kwargs(self, prometheus_logger):
exception = ExceptionWithCode("401")
assert prometheus_logger._extract_status_code(kwargs={"exception": exception}) == 401
def test_extract_from_enum_values(self, prometheus_logger):
enum_values = Mock(status_code="401")
assert prometheus_logger._extract_status_code(enum_values=enum_values) == 401
class TestInvalidAPIKeyDetection:
"""Test invalid API key request detection logic."""
@pytest.mark.parametrize("status_code,expected", [
(401, True),
(200, False),
(500, False),
(None, False),
])
def test_status_code_detection(self, prometheus_logger, status_code, expected):
assert prometheus_logger._is_invalid_api_key_request(status_code=status_code) == expected
def test_auth_error_message_detection(self, prometheus_logger):
exception = AssertionError("LiteLLM Virtual Key expected. Received=invalid-key-12345, expected to start with 'sk-'.")
assert prometheus_logger._is_invalid_api_key_request(status_code=None, exception=exception) is True
def test_non_auth_exception_not_detected(self, prometheus_logger):
exception = ValueError("Some other error")
assert prometheus_logger._is_invalid_api_key_request(status_code=None, exception=exception) is False
class TestSkipMetricsValidation:
"""Test high-level validation method that orchestrates detection and extraction."""
def test_skip_for_401_exception(self, prometheus_logger):
"""Test full flow: extraction -> detection -> skip decision."""
exception = ExceptionWithCode("401")
assert prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) is True
def test_skip_for_auth_error_message(self, prometheus_logger):
"""Test full flow: exception message -> detection -> skip decision."""
exception = AssertionError("expected to start with 'sk-'")
assert prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) is True
def test_no_skip_for_valid_request(self, prometheus_logger):
assert prometheus_logger._should_skip_metrics_for_invalid_key() is False
class TestAsyncHooks:
"""Test async hook methods skip metrics for invalid API keys."""
@pytest.fixture
def mock_user_api_key(self):
"""Create a mock UserAPIKeyAuth object."""
user_key = Mock(spec=UserAPIKeyAuth)
user_key.api_key = "test-key"
user_key.end_user_id = None
user_key.user_id = None
user_key.user_email = None
user_key.key_alias = None
user_key.team_id = None
user_key.team_alias = None
user_key.request_route = "/test"
return user_key
@pytest.mark.asyncio
async def test_post_call_failure_hook_skips_401(self, prometheus_logger, mock_user_api_key):
exception = ExceptionWithCode("401")
exception.__class__.__name__ = "ProxyException"
with patch.object(prometheus_logger, 'litellm_proxy_failed_requests_metric') as mock_failed, \
patch.object(prometheus_logger, 'litellm_proxy_total_requests_metric') as mock_total:
await prometheus_logger.async_post_call_failure_hook(
request_data={"model": "test-model"},
original_exception=exception,
user_api_key_dict=mock_user_api_key
)
mock_failed.labels.assert_not_called()
mock_total.labels.assert_not_called()
@pytest.mark.asyncio
async def test_log_failure_event_skips_401(self, prometheus_logger):
exception = ExceptionWithCode("401")
kwargs = {
"model": "test-model",
"standard_logging_object": {
"metadata": {
"user_api_key_hash": "test-key",
"user_api_key_user_id": "test-user",
},
"model_group": "test-model",
},
"exception": exception,
"litellm_params": {},
}
with patch.object(prometheus_logger, 'litellm_llm_api_failed_requests_metric') as mock_failed, \
patch.object(prometheus_logger, 'set_llm_deployment_failure_metrics') as mock_deployment:
await prometheus_logger.async_log_failure_event(
kwargs=kwargs,
response_obj=None,
start_time=None,
end_time=None
)
mock_failed.labels.assert_not_called()
mock_deployment.assert_not_called()

View file

@ -1137,3 +1137,94 @@ def test_bedrock_create_bedrock_block_different_document_formats():
assert f"DocumentPDFmessages_" in block["document"]["name"]
assert block["document"]["name"].endswith(f"_{format_type}")
assert block["document"]["format"] == format_type
def test_anthropic_messages_pt_server_tool_use_passthrough():
"""
Test that anthropic_messages_pt passes through server_tool_use and
tool_search_tool_result blocks in assistant message content.
These are Anthropic-native content types used for tool search functionality
that need to be preserved when reconstructing multi-turn conversations.
Fixes: https://github.com/BerriAI/litellm/issues/XXXXX
"""
from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt
messages = [
{
"role": "user",
"content": "I need help with time information."
},
{
"role": "assistant",
"content": [
{
"type": "server_tool_use",
"id": "srvtoolu_01ABC123",
"name": "tool_search_tool_regex",
"input": {"query": ".*time.*"}
},
{
"type": "tool_search_tool_result",
"tool_use_id": "srvtoolu_01ABC123",
"content": {
"type": "tool_search_tool_search_result",
"tool_references": [
{"type": "tool_reference", "tool_name": "get_time"}
]
}
},
{
"type": "text",
"text": "I found the time tool. How can I help you?"
}
],
},
{
"role": "user",
"content": "What's the time in New York?"
},
]
result = anthropic_messages_pt(
messages=messages,
model="claude-sonnet-4-5-20250929",
llm_provider="anthropic",
)
# Verify we have 3 messages (user, assistant, user)
assert len(result) == 3
# Verify the assistant message content
assistant_msg = result[1]
assert assistant_msg["role"] == "assistant"
assert isinstance(assistant_msg["content"], list)
# Find the different content block types
content_types = [block.get("type") for block in assistant_msg["content"]]
# Verify server_tool_use block is preserved
assert "server_tool_use" in content_types
server_tool_use_block = next(
b for b in assistant_msg["content"] if b.get("type") == "server_tool_use"
)
assert server_tool_use_block["id"] == "srvtoolu_01ABC123"
assert server_tool_use_block["name"] == "tool_search_tool_regex"
assert server_tool_use_block["input"] == {"query": ".*time.*"}
# Verify tool_search_tool_result block is preserved
assert "tool_search_tool_result" in content_types
tool_result_block = next(
b for b in assistant_msg["content"] if b.get("type") == "tool_search_tool_result"
)
assert tool_result_block["tool_use_id"] == "srvtoolu_01ABC123"
assert tool_result_block["content"]["type"] == "tool_search_tool_search_result"
assert tool_result_block["content"]["tool_references"][0]["tool_name"] == "get_time"
# Verify text block is also preserved
assert "text" in content_types
text_block = next(
b for b in assistant_msg["content"] if b.get("type") == "text"
)
assert text_block["text"] == "I found the time tool. How can I help you?"

View file

@ -570,6 +570,7 @@ async def test_ensure_initialize_azure_sdk_client_always_used(call_type):
or call_type == CallTypes.acreate_container
or call_type == CallTypes.adelete_container
or call_type == CallTypes.alist_container_files
or call_type == CallTypes.aupload_container_file
):
# Skip container call types as they're not supported for Azure (only OpenAI)
pytest.skip(f"Skipping {call_type.value} because Azure doesn't support container operations")

View file

@ -1175,6 +1175,30 @@ def test_vertex_ai_moonshot_uses_openai_handler():
)
def test_vertex_ai_zai_uses_openai_handler():
"""
Ensure ZAI partner models re-use the OpenAI-format handler.
"""
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
VertexAIPartnerModels,
)
assert VertexAIPartnerModels.should_use_openai_handler(
"zai-org/glm-4.7-maas"
)
def test_vertex_ai_zai_is_partner_model():
"""
Ensure ZAI models are detected as Vertex AI partner models.
"""
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
VertexAIPartnerModels,
)
assert VertexAIPartnerModels.is_vertex_partner_model("zai-org/glm-4.7-maas")
def test_build_vertex_schema_empty_properties():
"""
Test _build_vertex_schema handles empty properties objects correctly.

View file

@ -107,7 +107,7 @@ async def test_update_daily_spend_with_null_entity_id():
entity_type="user",
entity_id_field="user_id",
table_name="litellm_dailyuserspend",
unique_constraint_name="user_id_date_api_key_model_custom_llm_provider",
unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint",
)
# Verify that table.upsert was called
@ -115,12 +115,14 @@ async def test_update_daily_spend_with_null_entity_id():
# Verify the where clause contains null entity_id
call_args = mock_table.upsert.call_args[1]
where_clause = call_args["where"]["user_id_date_api_key_model_custom_llm_provider"]
where_clause = call_args["where"]["user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint"]
assert where_clause["user_id"] is None
assert where_clause["date"] == "2024-01-01"
assert where_clause["api_key"] == "test-api-key"
assert where_clause["model"] == "gpt-4"
assert where_clause["custom_llm_provider"] == "openai"
assert where_clause["mcp_namespaced_tool_name"] == ""
assert where_clause["endpoint"] == ""
# Verify the create data contains null entity_id
create_data = call_args["data"]["create"]
@ -129,6 +131,8 @@ async def test_update_daily_spend_with_null_entity_id():
assert create_data["api_key"] == "test-api-key"
assert create_data["model"] == "gpt-4"
assert create_data["custom_llm_provider"] == "openai"
assert create_data["mcp_namespaced_tool_name"] == ""
assert create_data["endpoint"] is None
assert create_data["prompt_tokens"] == 10
assert create_data["completion_tokens"] == 20
assert create_data["spend"] == 0.1
@ -171,13 +175,14 @@ async def test_update_daily_spend_sorting():
}
upsert_calls.append(call(
where={
"user_id_date_api_key_model_custom_llm_provider": {
"user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint": {
"user_id": f"user{i+11}", # user11 ... user60, sorted order
"date": "2024-01-01",
"api_key": "test-api-key",
"model": "gpt-4",
"custom_llm_provider": "openai",
"mcp_namespaced_tool_name": "",
"endpoint": "",
}
},
data={
@ -189,6 +194,7 @@ async def test_update_daily_spend_sorting():
"model_group": None,
"mcp_namespaced_tool_name": "",
"custom_llm_provider": "openai",
"endpoint": None,
"prompt_tokens": 10,
"completion_tokens": 20,
"spend": 0.1,
@ -203,6 +209,7 @@ async def test_update_daily_spend_sorting():
"api_requests": {"increment": 1},
"successful_requests": {"increment": 1},
"failed_requests": {"increment": 0},
"endpoint": "",
},
},
))
@ -216,7 +223,7 @@ async def test_update_daily_spend_sorting():
entity_type="user",
entity_id_field="user_id",
table_name="litellm_dailyuserspend",
unique_constraint_name="user_id_date_api_key_model_custom_llm_provider",
unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint",
)
# Verify that table.upsert was called
@ -372,7 +379,7 @@ async def test_update_daily_spend_with_none_values_in_sorting_fields():
entity_type="user",
entity_id_field="user_id",
table_name="litellm_dailyuserspend",
unique_constraint_name="user_id_date_api_key_model_custom_llm_provider",
unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint",
)
# Verify that table.upsert was called (should be called 5 times, once for each transaction)
@ -588,7 +595,7 @@ async def test_add_spend_log_transaction_to_daily_org_transaction_injects_org_id
update_dict = call_args["update"]
assert len(update_dict) == 1
for key, transaction in update_dict.items():
assert key == f"{org_id}_2024-01-01_test-key_gpt-4_openai"
assert key == f"{org_id}_2024-01-01_test-key_gpt-4_openai_"
assert transaction["organization_id"] == org_id
assert transaction["date"] == "2024-01-01"
assert transaction["api_key"] == "test-key"
@ -665,7 +672,7 @@ async def test_add_spend_log_transaction_to_daily_end_user_transaction_injects_e
update_dict = call_args["update"]
assert len(update_dict) == 1
for key, transaction in update_dict.items():
assert key == f"{end_user_id}_2024-01-01_test-key_gpt-4_openai"
assert key == f"{end_user_id}_2024-01-01_test-key_gpt-4_openai_"
assert transaction["end_user_id"] == end_user_id
assert transaction["date"] == "2024-01-01"
assert transaction["api_key"] == "test-key"
@ -741,7 +748,7 @@ async def test_add_spend_log_transaction_to_daily_agent_transaction_injects_agen
update_dict = call_args["update"]
assert len(update_dict) == 1
for key, transaction in update_dict.items():
assert key == f"{agent_id}_2024-01-01_test-key_gpt-4_openai"
assert key == f"{agent_id}_2024-01-01_test-key_gpt-4_openai_"
assert transaction["agent_id"] == agent_id
assert transaction["date"] == "2024-01-01"
assert transaction["api_key"] == "test-key"
@ -780,4 +787,55 @@ async def test_add_spend_log_transaction_to_daily_agent_transaction_skips_when_a
prisma_client=mock_prisma,
)
writer.daily_agent_spend_update_queue.add_update.assert_not_called()
writer.daily_agent_spend_update_queue.add_update.assert_not_called()
@pytest.mark.asyncio
async def test_endpoint_field_is_correctly_mapped_from_call_type():
"""
Test that the endpoint field is correctly mapped from call_type using ROUTE_ENDPOINT_MAPPING.
Verifies that when call_type is provided, the endpoint is set in the transaction and included in the key.
"""
writer = DBSpendUpdateWriter()
mock_prisma = MagicMock()
mock_prisma.get_request_status = MagicMock(return_value="success")
payload = {
"request_id": "req-endpoint-test",
"user": "test-user",
"call_type": "acompletion", # Maps to "/chat/completions"
"startTime": "2024-01-01T12:00:00",
"api_key": "test-key",
"model": "gpt-4",
"custom_llm_provider": "openai",
"model_group": "gpt-4-group",
"prompt_tokens": 100,
"completion_tokens": 50,
"spend": 0.15,
"metadata": '{"usage_object": {}}',
}
writer.daily_spend_update_queue.add_update = AsyncMock()
await writer.add_spend_log_transaction_to_daily_user_transaction(
payload=payload,
prisma_client=mock_prisma,
)
writer.daily_spend_update_queue.add_update.assert_called_once()
call_args = writer.daily_spend_update_queue.add_update.call_args[1]
update_dict = call_args["update"]
assert len(update_dict) == 1
for key, transaction in update_dict.items():
# Verify endpoint is included in the key
assert key == f"test-user_2024-01-01_test-key_gpt-4_openai_/chat/completions"
# Verify endpoint is set in the transaction
assert transaction["endpoint"] == "/chat/completions"
assert transaction["user_id"] == "test-user"
assert transaction["date"] == "2024-01-01"
assert transaction["api_key"] == "test-key"
assert transaction["model"] == "gpt-4"
assert transaction["custom_llm_provider"] == "openai"

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