scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge pull request #23163 from BerriAI/litellm_oss_staging_03_04_2026

Browse source 
Litellm oss staging 03 04 2026
This commit is contained in:
Sameer Kankute 2026-03-11 18:46:51 +05:30 • committed by GitHub
commit 2343149f2d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
97 changed files with 5896 additions and 123 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
deploy/charts/litellm-helm/templates/_helpers.tpl
View file

@ -61,6 +61,20 @@ Create the name of the service account to use
{{- end }}
{{- end }}
{{/*
Create the service account name used by migration jobs.
When Helm hooks are enabled, pre-install/pre-upgrade hooks run before normal resources.
If this chart is creating the ServiceAccount, it is not yet available for the hook job,
so fall back to "default" (or an explicit override) to avoid a cyclic dependency.
*/}}
{{- define "litellm.migrationServiceAccountName" -}}
{{- if and .Values.migrationJob.hooks.helm.enabled .Values.serviceAccount.create }}
{{- default "default" .Values.migrationJob.serviceAccountName }}
{{- else }}
{{- include "litellm.serviceAccountName" . }}
{{- end }}
{{- end }}
{{/*
Get redis service name
*/}}

2
deploy/charts/litellm-helm/templates/migrations-job.yaml
View file

@ -34,7 +34,7 @@ spec:
imagePullSecrets:
{{- toYaml . | nindent 8 }}
{{- end }}
serviceAccountName: {{ include "litellm.serviceAccountName" . }}
serviceAccountName: {{ include "litellm.migrationServiceAccountName" . }}
{{- with .Values.migrationJob.extraInitContainers }}
initContainers:
{{- toYaml . | nindent 8 }}

65
deploy/charts/litellm-helm/tests/migrations-job_tests.yaml
View file

@ -124,4 +124,67 @@ tests:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
name: DATABASE_URL
- it: should use default service account for helm hooks when serviceAccount.create is true
template: migrations-job.yaml
set:
migrationJob:
enabled: true
hooks:
helm:
enabled: true
serviceAccount:
create: true
asserts:
- equal:
path: spec.template.spec.serviceAccountName
value: default
- it: should use migrationJob.serviceAccountName override for helm hooks when serviceAccount.create is true
template: migrations-job.yaml
set:
migrationJob:
enabled: true
serviceAccountName: migration-sa
hooks:
helm:
enabled: true
serviceAccount:
create: true
asserts:
- equal:
path: spec.template.spec.serviceAccountName
value: migration-sa
- it: should use chart service account when helm hooks are disabled
template: migrations-job.yaml
set:
migrationJob:
enabled: true
hooks:
helm:
enabled: false
serviceAccount:
create: true
name: my-custom-sa
asserts:
- equal:
path: spec.template.spec.serviceAccountName
value: my-custom-sa
- it: should use pre-existing service account when helm hooks are enabled but serviceAccount.create is false
template: migrations-job.yaml
set:
migrationJob:
enabled: true
hooks:
helm:
enabled: true
serviceAccount:
create: false
name: pre-existing-sa
asserts:
- equal:
path: spec.template.spec.serviceAccountName
value: pre-existing-sa

4
deploy/charts/litellm-helm/values.yaml
View file

@ -307,6 +307,10 @@ migrationJob:
retries: 3 # Number of retries for the Job in case of failure
backoffLimit: 4 # Backoff limit for Job restarts
disableSchemaUpdate: false # Skip schema migrations for specific environments. When True, the job will exit with code 0.
# Optional service account for the migration job.
# Only used when migrationJob.hooks.helm.enabled=true and serviceAccount.create=true.
# In that case, pre-install/pre-upgrade hooks run before normal resources, so this defaults to "default".
serviceAccountName: ""
annotations: {}
ttlSecondsAfterFinished: 120
resources: {}

1
docs/my-website/docs/anthropic_count_tokens.md
View file

@ -138,6 +138,7 @@ The `/v1/messages/count_tokens` endpoint automatically routes to the appropriate
| Provider | Token Counting Method |
|----------|----------------------|
| Anthropic | [Anthropic Token Counting API](https://docs.anthropic.com/en/docs/build-with-claude/token-counting) |
| OpenAI | [OpenAI Responses API `/input_tokens`](https://platform.openai.com/docs/api-reference/responses/input-tokens) — see [Token Counting](./count_tokens.md) |
| Vertex AI (Claude) | Vertex AI Partner Models Token Counter |
| Bedrock (Claude) | AWS Bedrock CountTokens API |
| Gemini | Google AI Studio countTokens API |

3
docs/my-website/docs/audio_transcription.md
View file

@ -13,7 +13,7 @@ import TabItem from '@theme/TabItem';
| Fallbacks | ✅ | Works between supported models |
| Loadbalancing | ✅ | Works between supported models |
| Guardrails | ✅ | Applies to output transcribed text (non-streaming only) |
| Supported Providers | `openai`, `azure`, `vertex_ai`, `gemini`, `deepgram`, `groq`, `fireworks_ai`, `ovhcloud` | |
| Supported Providers | `openai`, `azure`, `vertex_ai`, `gemini`, `deepgram`, `groq`, `fireworks_ai`, `ovhcloud`, `mistral` | |
## Quick Start
@ -126,6 +126,7 @@ transcript = client.audio.transcriptions.create(
- [Fireworks AI](./providers/fireworks_ai.md#audio-transcription)
- [Groq](./providers/groq.md#speech-to-text---whisper)
- [Deepgram](./providers/deepgram.md)
- [Mistral (Voxtral)](./providers/mistral.md#audio-transcription)
- [OVHcloud AI Endpoints](./providers/ovhcloud.md)
---

189
docs/my-website/docs/count_tokens.md Normal file
View file

@ -0,0 +1,189 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Token Counting
## Overview
LiteLLM provides exact token counting by calling provider-specific token counting APIs. This gives you accurate token counts before sending requests, helping with cost estimation and context window management.
| Feature | Details |
|---------|---------|
| SDK Method | `litellm.acount_tokens()` |
| Proxy Endpoints | `/v1/messages/count_tokens` (Anthropic format), `/v1/responses/input_tokens` (OpenAI format) |
| Fallback | Local tiktoken-based counting for unsupported providers |
## Supported Providers
| Provider | Token Counting API | Format |
|----------|-------------------|--------|
| OpenAI | [Responses API `/input_tokens`](https://platform.openai.com/docs/api-reference/responses/input-tokens) | OpenAI Responses |
| Anthropic | [Messages `/count_tokens`](https://docs.anthropic.com/en/docs/build-with-claude/token-counting) | Anthropic Messages |
| Vertex AI (Claude) | Vertex AI Partner Models Token Counter | Anthropic Messages |
| Bedrock (Claude) | AWS Bedrock CountTokens API | Anthropic Messages |
| Gemini | Google AI Studio countTokens API | Anthropic Messages |
| Vertex AI (Gemini) | Vertex AI countTokens API | Anthropic Messages |
| Other providers | Local tiktoken fallback | N/A |
## SDK Usage
### Basic Usage
```python
import asyncio
import litellm
async def main():
# OpenAI
result = await litellm.acount_tokens(
model="openai/gpt-4o",
messages=[{"role": "user", "content": "Hello, how are you?"}],
)
print(f"Token count: {result.total_tokens}")
print(f"Tokenizer: {result.tokenizer_type}") # "openai_api"
# Anthropic
result = await litellm.acount_tokens(
model="anthropic/claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "Hello, how are you?"}],
)
print(f"Token count: {result.total_tokens}")
print(f"Tokenizer: {result.tokenizer_type}") # "anthropic_api"
asyncio.run(main())
```
### With Tools and System Message
```python
import asyncio
import litellm
async def main():
result = await litellm.acount_tokens(
model="openai/gpt-4o",
messages=[{"role": "user", "content": "What's the weather in Paris?"}],
tools=[{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather for a city",
"parameters": {
"type": "object",
"properties": {"city": {"type": "string"}},
},
},
}],
system="You are a helpful weather assistant.",
)
print(f"Token count (with tools): {result.total_tokens}")
asyncio.run(main())
```
### Response Format
`litellm.acount_tokens()` returns a `TokenCountResponse`:
```python
TokenCountResponse(
total_tokens=15, # Token count
request_model="openai/gpt-4o", # Model requested
model_used="gpt-4o", # Model used for counting
tokenizer_type="openai_api", # "openai_api", "anthropic_api", "local_tokenizer"
original_response={"input_tokens": 15}, # Raw API response
error=False, # True if counting failed
error_message=None, # Error details if failed
)
```
### Fallback Behavior
If a provider doesn't support a token counting API, or if the API key is missing, `acount_tokens()` automatically falls back to local tiktoken-based counting:
```python
# Unsupported provider → automatic fallback
result = await litellm.acount_tokens(
model="together_ai/meta-llama/Llama-3-8b-chat-hf",
messages=[{"role": "user", "content": "Hello"}],
)
print(result.tokenizer_type) # "local_tokenizer"
```
## Proxy Usage
### OpenAI Format — `/v1/responses/input_tokens`
<Tabs>
<TabItem value="curl" label="curl">
```bash
curl -X POST "http://localhost:4000/v1/responses/input_tokens" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-1234" \
-d '{
"model": "gpt-4o",
"input": "Hello, how are you?"
}'
```
</TabItem>
<TabItem value="python" label="Python (httpx)">
```python
import httpx
response = httpx.post(
"http://localhost:4000/v1/responses/input_tokens",
headers={
"Content-Type": "application/json",
"Authorization": "Bearer sk-1234"
},
json={
"model": "gpt-4o",
"input": "Hello, how are you?"
}
)
print(response.json())
# {"input_tokens": 7}
```
</TabItem>
</Tabs>
**Response:**
```json
{"input_tokens": 7}
```
### Anthropic Format — `/v1/messages/count_tokens`
See [Anthropic Token Counting](./anthropic_count_tokens.md) for full documentation.
```bash
curl -X POST "http://localhost:4000/v1/messages/count_tokens" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-1234" \
-d '{
"model": "claude-3-5-sonnet-20241022",
"messages": [
{"role": "user", "content": "Hello, how are you?"}
]
}'
```
## Proxy Configuration
```yaml
model_list:
- model_name: gpt-4o
litellm_params:
model: openai/gpt-4o
api_key: os.environ/OPENAI_API_KEY
- model_name: claude-3-5-sonnet
litellm_params:
model: anthropic/claude-3-5-sonnet-20241022
api_key: os.environ/ANTHROPIC_API_KEY
```

88
docs/my-website/docs/image_edits.md
View file

@ -16,7 +16,7 @@ LiteLLM provides image editing functionality that maps to OpenAI's `/images/edit
| Supported operations | Create image edits | Single and multiple images supported |
| Supported LiteLLM SDK Versions | 1.63.8+ | Gemini support requires 1.79.3+ |
| Supported LiteLLM Proxy Versions | 1.71.1+ | Gemini support requires 1.79.3+ |
| Supported LLM providers | **OpenAI**, **Gemini (Google AI Studio)**, **Vertex AI**, **OpenRouter**, **Stability AI**, **AWS Bedrock (Stability)** | Gemini supports the new `gemini-2.5-flash-image` family. Vertex AI supports both Gemini and Imagen models. OpenRouter routes image edits through chat completions. Stability AI and Bedrock Stability support various image editing operations. |
| Supported LLM providers | **OpenAI**, **Gemini (Google AI Studio)**, **Vertex AI**, **OpenRouter**, **Stability AI**, **AWS Bedrock (Stability)**, **Black Forest Labs** | Gemini supports the new `gemini-2.5-flash-image` family. Vertex AI supports both Gemini and Imagen models. OpenRouter routes image edits through chat completions. Stability AI and Bedrock Stability support various image editing operations. Black Forest Labs supports FLUX Kontext models. |
#### ⚡️See all supported models and providers at [models.litellm.ai](https://models.litellm.ai/)
@ -199,6 +199,63 @@ for idx, image_obj in enumerate(response.data):
</TabItem>
<TabItem value="bfl" label="Black Forest Labs">
#### Basic Image Edit
```python showLineNumbers title="Black Forest Labs Image Edit"
import os
import litellm
os.environ["BFL_API_KEY"] = "your-api-key"
response = litellm.image_edit(
model="black_forest_labs/flux-kontext-pro",
image=open("original_image.png", "rb"),
prompt="Add a green leaf to the scene",
)
print(response.data[0].url)
```
#### Inpainting with Mask
```python showLineNumbers title="Black Forest Labs Inpainting"
import os
import litellm
os.environ["BFL_API_KEY"] = "your-api-key"
# Use flux-pro-1.0-fill for inpainting
response = litellm.image_edit(
model="black_forest_labs/flux-pro-1.0-fill",
image=open("original_image.png", "rb"),
mask=open("mask_image.png", "rb"),
prompt="Replace with a garden",
)
print(response.data[0].url)
```
#### Outpainting (Expand)
```python showLineNumbers title="Black Forest Labs Outpainting"
import os
import litellm
os.environ["BFL_API_KEY"] = "your-api-key"
# Use flux-pro-1.0-expand to extend image borders
response = litellm.image_edit(
model="black_forest_labs/flux-pro-1.0-expand",
image=open("original_image.png", "rb"),
prompt="Continue the scene with mountains",
top=256,
bottom=256,
)
print(response.data[0].url)
```
</TabItem>
<TabItem value="vertex_ai" label="Vertex AI">
#### Basic Image Edit (Gemini)
@ -392,6 +449,35 @@ curl -X POST "http://0.0.0.0:4000/v1/images/edits" \
</TabItem>
<TabItem value="bfl" label="Black Forest Labs">
1. Add Black Forest Labs image edit models to your `config.yaml`:
```yaml showLineNumbers title="Black Forest Labs Proxy Configuration"
model_list:
- model_name: bfl-kontext-pro
litellm_params:
model: black_forest_labs/flux-kontext-pro
api_key: os.environ/BFL_API_KEY
model_info:
mode: image_edit
```
2. Start the LiteLLM proxy server:
```bash showLineNumbers title="Start LiteLLM Proxy Server"
litellm --config /path/to/config.yaml
```
3. Make an image edit request:
```bash showLineNumbers title="Black Forest Labs Proxy Image Edit"
curl -X POST "http://0.0.0.0:4000/v1/images/edits" \
-H "Authorization: Bearer <YOUR-LITELLM-KEY>" \
-F "model=bfl-kontext-pro" \
-F "image=@original_image.png" \
-F "prompt=Add a sunset in the background"
```
</TabItem>
<TabItem value="vertex_ai" label="Vertex AI">
1. Add Vertex AI image edit models to your `config.yaml`:

2
docs/my-website/docs/image_generation.md
View file

@ -15,7 +15,7 @@ import TabItem from '@theme/TabItem';
| Fallbacks | ✅ | Works between supported models |
| Loadbalancing | ✅ | Works between supported models |
| Guardrails | ✅ | Applies to input prompts (non-streaming only) |
| Supported Providers | OpenAI, Azure, Google AI Studio, Vertex AI, AWS Bedrock, Recraft, OpenRouter, Xinference, Nscale | |
| Supported Providers | OpenAI, Azure, Google AI Studio, Vertex AI, AWS Bedrock, Black Forest Labs, Recraft, OpenRouter, Xinference, Nscale | |
## Quick Start

291
docs/my-website/docs/providers/black_forest_labs.md Normal file
View file

@ -0,0 +1,291 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Black Forest Labs Image Generation
Black Forest Labs provides state-of-the-art text-to-image generation using their FLUX models.
## Overview
| Property | Details |
|----------|---------|
| Description | Black Forest Labs FLUX models for high-quality text-to-image generation |
| Provider Route on LiteLLM | `black_forest_labs/` |
| Provider Doc | [Black Forest Labs API ↗](https://docs.bfl.ai/) |
| Supported Operations | [`/images/generations`](#image-generation) |
## Setup
### API Key
```python showLineNumbers
import os
# Set your Black Forest Labs API key
os.environ["BFL_API_KEY"] = "your-api-key-here"
```
Get your API key from [Black Forest Labs](https://blackforestlabs.ai/).
## Supported Models
| Model Name | Description | Price |
|------------|-------------|-------|
| `black_forest_labs/flux-pro-1.1` | Fast & reliable standard generation | $0.04/image |
| `black_forest_labs/flux-pro-1.1-ultra` | Ultra high-resolution (up to 4MP) | $0.06/image |
| `black_forest_labs/flux-dev` | Development/open-source variant | $0.025/image |
| `black_forest_labs/flux-pro` | Original pro model | $0.05/image |
## Image Generation
### Usage - LiteLLM Python SDK
<Tabs>
<TabItem value="basic" label="Basic Usage">
```python showLineNumbers title="Basic Image Generation"
import os
import litellm
# Set your API key
os.environ["BFL_API_KEY"] = "your-api-key-here"
# Generate an image
response = litellm.image_generation(
model="black_forest_labs/flux-pro-1.1",
prompt="A beautiful sunset over the ocean with sailing boats",
)
# BFL returns URLs
print(response.data[0].url)
```
</TabItem>
<TabItem value="async" label="Async Usage">
```python showLineNumbers title="Async Image Generation"
import os
import asyncio
import litellm
# Set your API key
os.environ["BFL_API_KEY"] = "your-api-key-here"
async def generate_image():
response = await litellm.aimage_generation(
model="black_forest_labs/flux-pro-1.1",
prompt="A futuristic city skyline at night",
)
print(response.data[0].url)
# Run the async function
asyncio.run(generate_image())
```
</TabItem>
<TabItem value="size" label="Custom Size">
```python showLineNumbers title="Image Generation with Custom Size"
import os
import litellm
# Set your API key
os.environ["BFL_API_KEY"] = "your-api-key-here"
# Generate with specific dimensions
response = litellm.image_generation(
model="black_forest_labs/flux-pro-1.1",
prompt="A majestic mountain landscape",
size="1792x1024", # Maps to width/height
)
print(response.data[0].url)
```
</TabItem>
<TabItem value="ultra" label="Ultra High-Res">
```python showLineNumbers title="Ultra High Resolution with flux-pro-1.1-ultra"
import os
import litellm
# Set your API key
os.environ["BFL_API_KEY"] = "your-api-key-here"
# Generate ultra high-resolution image
response = litellm.image_generation(
model="black_forest_labs/flux-pro-1.1-ultra",
prompt="Detailed portrait of a fantasy character",
size="2048x2048", # Up to 4MP supported
quality="hd", # Maps to raw=True for natural look
)
print(response.data[0].url)
```
</TabItem>
<TabItem value="advanced" label="Advanced Parameters">
```python showLineNumbers title="Advanced Image Generation with BFL Parameters"
import os
import litellm
# Set your API key
os.environ["BFL_API_KEY"] = "your-api-key-here"
# Generate with BFL-specific parameters
response = litellm.image_generation(
model="black_forest_labs/flux-pro-1.1",
prompt="A cute orange cat sitting on a windowsill",
seed=42, # For reproducible results
output_format="png", # png or jpeg
safety_tolerance=2, # 0-6, higher = more permissive
prompt_upsampling=True, # Enhance prompt for better results
)
print(response.data[0].url)
```
</TabItem>
</Tabs>
### Usage - LiteLLM Proxy Server
#### 1. Configure your config.yaml
```yaml showLineNumbers title="Black Forest Labs Image Generation Configuration"
model_list:
- model_name: flux-pro
litellm_params:
model: black_forest_labs/flux-pro-1.1
api_key: os.environ/BFL_API_KEY
model_info:
mode: image_generation
- model_name: flux-ultra
litellm_params:
model: black_forest_labs/flux-pro-1.1-ultra
api_key: os.environ/BFL_API_KEY
model_info:
mode: image_generation
- model_name: flux-dev
litellm_params:
model: black_forest_labs/flux-dev
api_key: os.environ/BFL_API_KEY
model_info:
mode: image_generation
general_settings:
master_key: sk-1234
```
#### 2. Start LiteLLM Proxy Server
```bash showLineNumbers title="Start LiteLLM Proxy Server"
litellm --config /path/to/config.yaml
# RUNNING on http://0.0.0.0:4000
```
#### 3. Make image generation requests
<Tabs>
<TabItem value="openai-sdk" label="OpenAI SDK">
```python showLineNumbers title="Black Forest Labs via Proxy - OpenAI SDK"
from openai import OpenAI
# Initialize client with your proxy URL
client = OpenAI(
base_url="http://localhost:4000",
api_key="sk-1234"
)
# Generate image with FLUX Pro
response = client.images.generate(
model="flux-pro",
prompt="A beautiful garden with colorful flowers",
size="1024x1024",
)
print(response.data[0].url)
```
</TabItem>
<TabItem value="curl" label="cURL">
```bash showLineNumbers title="Black Forest Labs via Proxy - cURL"
curl -X POST 'http://localhost:4000/v1/images/generations' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{
"model": "flux-pro",
"prompt": "A beautiful garden with colorful flowers",
"size": "1024x1024"
}'
```
</TabItem>
</Tabs>
## Supported Parameters
### OpenAI-Compatible Parameters
| Parameter | Type | Description | Mapping |
|-----------|------|-------------|---------|
| `prompt` | string | Text description of the image to generate | Direct |
| `model` | string | The FLUX model to use | Direct |
| `size` | string | Image dimensions (e.g., `1024x1024`) | Maps to `width` and `height` |
| `n` | integer | Number of images (ultra model only, up to 4) | Maps to `num_images` |
| `quality` | string | `hd` for natural look | Maps to `raw=True` for ultra |
| `response_format` | string | `url` or `b64_json` | Direct |
### Black Forest Labs Specific Parameters
| Parameter | Type | Description | Default |
|-----------|------|-------------|---------|
| `width` | integer | Image width (256-1920, multiples of 16) | 1024 |
| `height` | integer | Image height (256-1920, multiples of 16) | 1024 |
| `aspect_ratio` | string | Alternative to width/height (e.g., `16:9`, `1:1`) | - |
| `seed` | integer | Seed for reproducible results | Random |
| `output_format` | string | Output format: `png` or `jpeg` | `png` |
| `safety_tolerance` | integer | Safety filter tolerance (0-6, higher = more permissive) | 2 |
| `prompt_upsampling` | boolean | Enhance prompt for better results | `false` |
### Ultra Model Specific Parameters
| Parameter | Type | Description | Default |
|-----------|------|-------------|---------|
| `raw` | boolean | Raw mode for more natural, less synthetic look | `false` |
| `num_images` | integer | Number of images to generate (1-4) | 1 |
## How It Works
Black Forest Labs uses a polling-based API:
1. **Submit Request**: LiteLLM sends your prompt to BFL
2. **Get Task ID**: BFL returns a task ID and polling URL
3. **Poll for Result**: LiteLLM automatically polls until the image is ready
4. **Return Result**: The generated image URL is returned
This polling is handled automatically by LiteLLM - you just call `image_generation()` and get the result.
## Getting Started
1. Create an account at [Black Forest Labs](https://blackforestlabs.ai/)
2. Get your API key from the dashboard
3. Set your `BFL_API_KEY` environment variable
4. Use `litellm.image_generation()` with any supported model
## Additional Resources
- [Black Forest Labs Documentation](https://docs.bfl.ai/)
- [Black Forest Labs Image Editing](./black_forest_labs_img_edit.md) - For editing existing images
- [FLUX Model Information](https://blackforestlabs.ai/)

301
docs/my-website/docs/providers/black_forest_labs_img_edit.md Normal file
View file

@ -0,0 +1,301 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Black Forest Labs Image Editing
Black Forest Labs provides powerful image editing capabilities using their FLUX models to modify existing images based on text descriptions.
## Overview
| Property | Details |
|----------|---------|
| Description | Black Forest Labs Image Editing uses FLUX Kontext and other models to modify, inpaint, and expand images based on text prompts. |
| Provider Route on LiteLLM | `black_forest_labs/` |
| Provider Doc | [Black Forest Labs API ↗](https://docs.bfl.ai/) |
| Supported Operations | [`/images/edits`](#image-editing) |
## Setup
### API Key
```python showLineNumbers
import os
# Set your Black Forest Labs API key
os.environ["BFL_API_KEY"] = "your-api-key-here"
```
Get your API key from [Black Forest Labs](https://blackforestlabs.ai/).
## Supported Models
| Model Name | Description | Use Case |
|------------|-------------|----------|
| `black_forest_labs/flux-kontext-pro` | FLUX Kontext Pro - General image editing with prompts | General editing, style transfer |
| `black_forest_labs/flux-kontext-max` | FLUX Kontext Max - Premium quality editing | High-quality edits |
| `black_forest_labs/flux-pro-1.0-fill` | FLUX Pro Fill - Inpainting with mask | Remove/replace objects |
| `black_forest_labs/flux-pro-1.0-expand` | FLUX Pro Expand - Outpainting | Expand image borders |
## Image Editing
### Usage - LiteLLM Python SDK
<Tabs>
<TabItem value="basic-edit" label="Basic Usage">
```python showLineNumbers title="Basic Image Editing"
import os
import litellm
# Set your API key
os.environ["BFL_API_KEY"] = "your-api-key-here"
# Edit an image with a prompt
response = litellm.image_edit(
model="black_forest_labs/flux-kontext-pro",
image=open("path/to/your/image.png", "rb"),
prompt="Add a green leaf to the scene",
)
# BFL returns URLs
print(response.data[0].url)
```
</TabItem>
<TabItem value="async-edit" label="Async Usage">
```python showLineNumbers title="Async Image Editing"
import os
import asyncio
import litellm
# Set your API key
os.environ["BFL_API_KEY"] = "your-api-key-here"
async def edit_image():
response = await litellm.aimage_edit(
model="black_forest_labs/flux-kontext-pro",
image=open("path/to/your/image.png", "rb"),
prompt="Make this image look like a watercolor painting",
)
print(response.data[0].url)
# Run the async function
asyncio.run(edit_image())
```
</TabItem>
<TabItem value="inpainting" label="Inpainting (Fill)">
```python showLineNumbers title="Inpainting with Mask"
import os
import litellm
# Set your API key
os.environ["BFL_API_KEY"] = "your-api-key-here"
# Use flux-pro-1.0-fill for inpainting
response = litellm.image_edit(
model="black_forest_labs/flux-pro-1.0-fill",
image=open("path/to/your/image.png", "rb"),
mask=open("path/to/mask.png", "rb"), # White areas will be edited
prompt="Replace with a beautiful garden",
steps=50, # BFL-specific parameter
guidance=30, # BFL-specific parameter
)
print(response.data[0].url)
```
</TabItem>
<TabItem value="outpainting" label="Outpainting (Expand)">
```python showLineNumbers title="Outpainting - Expand Image Borders"
import os
import litellm
# Set your API key
os.environ["BFL_API_KEY"] = "your-api-key-here"
# Use flux-pro-1.0-expand to extend image borders
response = litellm.image_edit(
model="black_forest_labs/flux-pro-1.0-expand",
image=open("path/to/your/image.png", "rb"),
prompt="Continue the scene with a mountain landscape",
top=256, # Expand 256 pixels at top
bottom=256, # Expand 256 pixels at bottom
left=128, # Expand 128 pixels at left
right=128, # Expand 128 pixels at right
)
print(response.data[0].url)
```
</TabItem>
<TabItem value="advanced" label="Advanced Parameters">
```python showLineNumbers title="Advanced Image Editing with BFL Parameters"
import os
import litellm
# Set your API key
os.environ["BFL_API_KEY"] = "your-api-key-here"
# Edit image with BFL-specific parameters
response = litellm.image_edit(
model="black_forest_labs/flux-kontext-pro",
image=open("path/to/your/image.png", "rb"),
prompt="Transform into cyberpunk style with neon lights",
seed=42, # For reproducible results
output_format="png", # png or jpeg
safety_tolerance=2, # 0-6, higher = more permissive
aspect_ratio="16:9", # Output aspect ratio
)
print(response.data[0].url)
```
</TabItem>
</Tabs>
### Usage - LiteLLM Proxy Server
#### 1. Configure your config.yaml
```yaml showLineNumbers title="Black Forest Labs Image Editing Configuration"
model_list:
- model_name: bfl-kontext-pro
litellm_params:
model: black_forest_labs/flux-kontext-pro
api_key: os.environ/BFL_API_KEY
model_info:
mode: image_edit
- model_name: bfl-kontext-max
litellm_params:
model: black_forest_labs/flux-kontext-max
api_key: os.environ/BFL_API_KEY
model_info:
mode: image_edit
- model_name: bfl-fill
litellm_params:
model: black_forest_labs/flux-pro-1.0-fill
api_key: os.environ/BFL_API_KEY
model_info:
mode: image_edit
- model_name: bfl-expand
litellm_params:
model: black_forest_labs/flux-pro-1.0-expand
api_key: os.environ/BFL_API_KEY
model_info:
mode: image_edit
general_settings:
master_key: sk-1234
```
#### 2. Start LiteLLM Proxy Server
```bash showLineNumbers title="Start LiteLLM Proxy Server"
litellm --config /path/to/config.yaml
# RUNNING on http://0.0.0.0:4000
```
#### 3. Make image editing requests
<Tabs>
<TabItem value="openai-sdk" label="OpenAI SDK">
```python showLineNumbers title="Black Forest Labs via Proxy - OpenAI SDK"
from openai import OpenAI
# Initialize client with your proxy URL
client = OpenAI(
base_url="http://localhost:4000",
api_key="sk-1234"
)
# Edit image with FLUX Kontext Pro
response = client.images.edit(
model="bfl-kontext-pro",
image=open("path/to/your/image.png", "rb"),
prompt="Add magical sparkles and fairy dust",
)
print(response.data[0].url)
```
</TabItem>
<TabItem value="curl" label="cURL">
```bash showLineNumbers title="Black Forest Labs via Proxy - cURL"
curl --location 'http://localhost:4000/v1/images/edits' \
--header 'Authorization: Bearer sk-1234' \
--form 'model="bfl-kontext-pro"' \
--form 'prompt="Add a sunset in the background"' \
--form 'image=@"path/to/your/image.png"'
```
</TabItem>
</Tabs>
## Supported Parameters
### OpenAI-Compatible Parameters
| Parameter | Type | Description | Default |
|-----------|------|-------------|---------|
| `image` | file | The image file to edit | Required |
| `prompt` | string | Text description of the desired changes | Required |
| `model` | string | The FLUX model to use | Required |
| `mask` | file | Mask image for inpainting (flux-pro-1.0-fill) | Optional |
| `n` | integer | Number of images (BFL returns 1 per request) | `1` |
| `size` | string | Maps to aspect_ratio | Optional |
| `response_format` | string | `url` or `b64_json` | `url` |
### Black Forest Labs Specific Parameters
| Parameter | Type | Description | Default | Models |
|-----------|------|-------------|---------|--------|
| `seed` | integer | Seed for reproducible results | Random | All |
| `output_format` | string | Output format: `png` or `jpeg` | `png` | All |
| `safety_tolerance` | integer | Safety filter tolerance (0-6) | 2 | All |
| `aspect_ratio` | string | Output aspect ratio (e.g., `16:9`, `1:1`) | Original | Kontext models |
| `steps` | integer | Number of inference steps | Model default | Fill |
| `guidance` | float | Guidance scale | Model default | Fill |
| `grow_mask` | integer | Pixels to grow mask | 0 | Fill |
| `top` | integer | Pixels to expand at top | 0 | Expand |
| `bottom` | integer | Pixels to expand at bottom | 0 | Expand |
| `left` | integer | Pixels to expand at left | 0 | Expand |
| `right` | integer | Pixels to expand at right | 0 | Expand |
## How It Works
Black Forest Labs uses a polling-based API:
1. **Submit Request**: LiteLLM sends your image and prompt to BFL
2. **Get Task ID**: BFL returns a task ID and polling URL
3. **Poll for Result**: LiteLLM automatically polls until the image is ready
4. **Return Result**: The generated image URL is returned
This polling is handled automatically by LiteLLM - you just call `image_edit()` and get the result.
## Getting Started
1. Create an account at [Black Forest Labs](https://blackforestlabs.ai/)
2. Get your API key from the dashboard
3. Set your `BFL_API_KEY` environment variable
4. Use `litellm.image_edit()` with any supported model
## Additional Resources
- [Black Forest Labs Documentation](https://docs.bfl.ai/)
- [FLUX Model Information](https://blackforestlabs.ai/)

22
docs/my-website/docs/providers/gemini.md
View file

@ -1562,13 +1562,18 @@ LiteLLM Supports the following image types passed in `url`
## Media Resolution Control (Images & Videos)
For Gemini 3+ models, LiteLLM supports per-part media resolution control using OpenAI's `detail` parameter. This allows you to specify different resolution levels for individual images and videos in your request, whether using `image_url` or `file` content types.
LiteLLM supports OpenAI's `detail` parameter for specifying the image resolution when using Gemini models. The behavior differs between Gemini versions:
| Gemini Version | Resolution Control | Behavior |
|----------------|-------------------|----------|
| Gemini 3+ | Per-part | Each image/video can have its own `detail` setting |
| Gemini 2.x (2.0, 2.5) | Global | The highest `detail` from all images is applied globally via `mediaResolution` in `generationConfig` |
**Supported `detail` values:**
- `"low"` - Maps to `media_resolution: "low"` (280 tokens for images, 70 tokens per frame for videos)
- `"medium"` - Maps to `media_resolution: "medium"`
- `"high"` - Maps to `media_resolution: "high"` (1120 tokens for images)
- `"ultra_high"` - Maps to `media_resolution: "ultra_high"`
- `"low"` - Maps to `MEDIA_RESOLUTION_LOW` (280 tokens for images, 70 tokens per frame for videos)
- `"medium"` - Maps to `MEDIA_RESOLUTION_MEDIUM`
- `"high"` - Maps to `MEDIA_RESOLUTION_HIGH` (1120 tokens for images)
- `"ultra_high"` - Maps to `MEDIA_RESOLUTION_ULTRA_HIGH`
- `"auto"` or `None` - Model decides optimal resolution (no `media_resolution` set)
**Usage Examples:**
@ -1605,8 +1610,9 @@ messages = [
}
]
# Works with both Gemini 2.x and 3+
response = completion(
model="gemini/gemini-3-pro-preview",
model="gemini/gemini-2.5-flash", # or gemini-3-pro-preview
messages=messages,
)
```
@ -1647,7 +1653,9 @@ response = completion(
</Tabs>
:::info
**Per-Part Resolution:** Each image or video in your request can have its own `detail` setting, allowing mixed-resolution requests (e.g., a high-res chart alongside a low-res icon). This feature works with both `image_url` and `file` content types, and is only available for Gemini 3+ models.
**Gemini 3+ Per-Part Resolution:** Each image or video can have its own `detail` setting, allowing mixed-resolution requests (e.g., a high-res chart alongside a low-res icon). This works with both `image_url` and `file` content types.
**Gemini 2.x Global Resolution:** When multiple images have different `detail` values, LiteLLM uses the highest resolution found and applies it globally via `mediaResolution` in `generationConfig` (e.g., if one image has `"low"` and another has `"high"`, all images will use `"high"`).
:::
## Video Metadata Control

73
docs/my-website/docs/providers/mistral.md
View file

@ -311,6 +311,79 @@ print(response)
- **Model Compatibility**: Reasoning parameters only work with magistral models
- **Backward Compatibility**: Non-magistral models will ignore reasoning parameters and work normally
## Audio Transcription
Use Mistral's Voxtral models for audio transcription via `litellm.transcription()`.
### SDK Usage
```python
from litellm import transcription
import os
os.environ["MISTRAL_API_KEY"] = ""
audio_file = open("path/to/audio.wav", "rb")
response = transcription(
model="mistral/voxtral-mini-latest",
file=audio_file,
)
print(response.text)
```
### With Optional Parameters
```python
response = transcription(
model="mistral/voxtral-mini-latest",
file=audio_file,
language="en",
temperature=0.0,
response_format="json",
)
```
### Mistral-Specific Parameters
Mistral supports additional parameters beyond the OpenAI-compatible ones:
| Parameter | Type | Description |
|-----------|------|-------------|
| `diarize` | `bool` | Enable speaker diarization |
```python
response = transcription(
model="mistral/voxtral-mini-latest",
file=audio_file,
diarize=True,
)
```
### Usage with LiteLLM Proxy
```yaml
model_list:
- model_name: voxtral
litellm_params:
model: mistral/voxtral-mini-latest
api_key: os.environ/MISTRAL_API_KEY
model_info:
mode: audio_transcription
```
```bash
litellm --config /path/to/config.yaml
```
```bash
curl --location 'http://0.0.0.0:4000/v1/audio/transcriptions' \
--header 'Authorization: Bearer sk-1234' \
--form 'file=@"audio.wav"' \
--form 'model="voxtral"'
```
## Sample Usage - Embedding
```python
from litellm import embedding

2
docs/my-website/sidebars.js
View file

@ -814,6 +814,8 @@ const sidebars = {
"providers/anyscale",
"providers/apertis",
"providers/baseten",
"providers/black_forest_labs",
"providers/black_forest_labs_img_edit",
"providers/bytez",
"providers/cerebras",
"providers/chutes",

5
litellm/__init__.py
View file

@ -578,6 +578,7 @@ v0_models: Set = set()
morph_models: Set = set()
lambda_ai_models: Set = set()
hyperbolic_models: Set = set()
black_forest_labs_models: Set = set()
recraft_models: Set = set()
cometapi_models: Set = set()
oci_models: Set = set()
@ -825,6 +826,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None):
lambda_ai_models.add(key)
elif value.get("litellm_provider") == "hyperbolic":
hyperbolic_models.add(key)
elif value.get("litellm_provider") == "black_forest_labs":
black_forest_labs_models.add(key)
elif value.get("litellm_provider") == "recraft":
recraft_models.add(key)
elif value.get("litellm_provider") == "cometapi":
@ -958,6 +961,7 @@ model_list = list(
| v0_models
| morph_models
| lambda_ai_models
| black_forest_labs_models
| recraft_models
| cometapi_models
| oci_models
@ -1056,6 +1060,7 @@ models_by_provider: dict = {
"morph": morph_models,
"lambda_ai": lambda_ai_models,
"hyperbolic": hyperbolic_models,
"black_forest_labs": black_forest_labs_models,
"recraft": recraft_models,
"cometapi": cometapi_models,
"oci": oci_models,

39
litellm/images/main.py
View file

@ -50,6 +50,10 @@ from litellm.main import (
openai_image_variations,
)
# BFL handlers
from litellm.llms.black_forest_labs.image_edit.handler import bfl_image_edit
from litellm.llms.black_forest_labs.image_generation.handler import bfl_image_generation
###########################################
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
@ -404,7 +408,7 @@ def image_generation( # noqa: PLR0915
litellm.LlmProviders.STABILITY,
litellm.LlmProviders.RUNWAYML,
litellm.LlmProviders.VERTEX_AI,
litellm.LlmProviders.OPENROUTER
litellm.LlmProviders.OPENROUTER,
):
if image_generation_config is None:
raise ValueError(
@ -427,6 +431,22 @@ def image_generation( # noqa: PLR0915
timeout=timeout,
client=client,
)
elif custom_llm_provider == "black_forest_labs":
# Route to BFL-specific handler (polling required)
if model is None:
raise Exception("Model needs to be set for black_forest_labs")
return bfl_image_generation.image_generation(
model=model,
prompt=prompt,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params_dict,
logging_obj=litellm_logging_obj,
timeout=timeout,
extra_headers=extra_headers,
client=client,
aimg_generation=aimg_generation,
)
elif custom_llm_provider == "azure_ai":
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
@ -920,6 +940,23 @@ def image_edit( # noqa: PLR0915
_is_async=_is_async,
client=kwargs.get("client"),
)
elif custom_llm_provider == "black_forest_labs":
# Route to BFL-specific handler (polling required)
if model is None:
raise Exception("Model needs to be set for black_forest_labs")
image_edit_request_params.update(non_default_params)
return bfl_image_edit.image_edit(
model=model,
image=images,
prompt=prompt,
image_edit_optional_request_params=image_edit_request_params,
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
timeout=timeout or DEFAULT_REQUEST_TIMEOUT,
extra_headers=extra_headers,
client=kwargs.get("client"),
aimage_edit=_is_async,
)
# Call the handler with _is_async flag instead of directly calling the async handler
return base_llm_http_handler.image_edit_handler(
model=model,

8
litellm/litellm_core_utils/get_supported_openai_params.py
View file

@ -144,6 +144,14 @@ def get_supported_openai_params( # noqa: PLR0915
return litellm.MistralConfig().get_supported_openai_params(model=model)
elif request_type == "embeddings":
return litellm.MistralEmbeddingConfig().get_supported_openai_params()
elif request_type == "transcription":
from litellm.llms.mistral.audio_transcription.transformation import (
MistralAudioTranscriptionConfig,
)
return MistralAudioTranscriptionConfig().get_supported_openai_params(
model=model
)
elif custom_llm_provider == "text-completion-codestral":
return litellm.CodestralTextCompletionConfig().get_supported_openai_params(
model=model

21
litellm/llms/black_forest_labs/__init__.py Normal file
View file

@ -0,0 +1,21 @@
from .common_utils import (
DEFAULT_API_BASE,
DEFAULT_MAX_POLLING_TIME,
DEFAULT_POLLING_INTERVAL,
IMAGE_EDIT_MODELS,
IMAGE_GENERATION_MODELS,
BlackForestLabsError,
)
from .image_edit import BlackForestLabsImageEditConfig
from .image_generation import BlackForestLabsImageGenerationConfig
__all__ = [
"BlackForestLabsError",
"BlackForestLabsImageEditConfig",
"BlackForestLabsImageGenerationConfig",
"DEFAULT_API_BASE",
"DEFAULT_MAX_POLLING_TIME",
"DEFAULT_POLLING_INTERVAL",
"IMAGE_EDIT_MODELS",
"IMAGE_GENERATION_MODELS",
]

42
litellm/llms/black_forest_labs/common_utils.py Normal file
View file

@ -0,0 +1,42 @@
"""
Black Forest Labs Common Utilities
Common utilities, constants, and error handling for Black Forest Labs API.
"""
from typing import Dict
from litellm.llms.base_llm.chat.transformation import BaseLLMException
class BlackForestLabsError(BaseLLMException):
"""Exception class for Black Forest Labs API errors."""
pass
# API Constants
DEFAULT_API_BASE = "https://api.bfl.ai"
# Polling configuration
DEFAULT_POLLING_INTERVAL = 1.5 # seconds
DEFAULT_MAX_POLLING_TIME = 300 # 5 minutes
# Model to endpoint mapping for image edit
IMAGE_EDIT_MODELS: Dict[str, str] = {
"flux-kontext-pro": "/v1/flux-kontext-pro",
"flux-kontext-max": "/v1/flux-kontext-max",
"flux-pro-1.0-fill": "/v1/flux-pro-1.0-fill",
"flux-pro-1.0-expand": "/v1/flux-pro-1.0-expand",
}
# Model to endpoint mapping for image generation
IMAGE_GENERATION_MODELS: Dict[str, str] = {
"flux-pro-1.1": "/v1/flux-pro-1.1",
"flux-pro-1.1-ultra": "/v1/flux-pro-1.1-ultra",
"flux-dev": "/v1/flux-dev",
"flux-pro": "/v1/flux-pro",
# Kontext models support both text-to-image and image editing
"flux-kontext-pro": "/v1/flux-kontext-pro",
"flux-kontext-max": "/v1/flux-kontext-max",
}

8
litellm/llms/black_forest_labs/image_edit/__init__.py Normal file
View file

@ -0,0 +1,8 @@
from .handler import BlackForestLabsImageEdit, bfl_image_edit
from .transformation import BlackForestLabsImageEditConfig
__all__ = [
"BlackForestLabsImageEditConfig",
"BlackForestLabsImageEdit",
"bfl_image_edit",
]

454
litellm/llms/black_forest_labs/image_edit/handler.py Normal file
View file

@ -0,0 +1,454 @@
"""
Black Forest Labs Image Edit Handler
Handles image edit requests for Black Forest Labs models.
BFL uses an async polling pattern - the initial request returns a task ID,
then we poll until the result is ready.
"""
import asyncio
import time
from typing import Any, Dict, List, Optional, Union
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
_get_httpx_client,
get_async_httpx_client,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import FileTypes, ImageResponse
from ..common_utils import (
DEFAULT_MAX_POLLING_TIME,
DEFAULT_POLLING_INTERVAL,
BlackForestLabsError,
)
from .transformation import BlackForestLabsImageEditConfig
class BlackForestLabsImageEdit:
"""
Black Forest Labs Image Edit handler.
Handles the HTTP requests and polling logic, delegating data transformation
to the BlackForestLabsImageEditConfig class.
"""
def __init__(self):
self.config = BlackForestLabsImageEditConfig()
def image_edit(
self,
model: str,
image: Union[FileTypes, List[FileTypes]],
prompt: Optional[str],
image_edit_optional_request_params: Dict,
litellm_params: Union[GenericLiteLLMParams, Dict],
logging_obj: LiteLLMLoggingObj,
timeout: Optional[Union[float, httpx.Timeout]],
extra_headers: Optional[Dict[str, Any]] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
aimage_edit: bool = False,
) -> Union[ImageResponse, Any]:
"""
Main entry point for image edit requests.
Args:
model: The model to use (e.g., "black_forest_labs/flux-kontext-pro")
image: The image(s) to edit
prompt: The edit instruction
image_edit_optional_request_params: Optional parameters for the request
litellm_params: LiteLLM parameters including api_key, api_base
logging_obj: Logging object
timeout: Request timeout
extra_headers: Additional headers
client: HTTP client to use
aimage_edit: If True, return async coroutine
Returns:
ImageResponse or coroutine if aimage_edit=True
"""
# Handle litellm_params as dict or object
if isinstance(litellm_params, dict):
api_key = litellm_params.get("api_key")
api_base = litellm_params.get("api_base")
litellm_params_dict = litellm_params
else:
api_key = litellm_params.api_key
api_base = litellm_params.api_base
litellm_params_dict = dict(litellm_params)
if aimage_edit:
return self.async_image_edit(
model=model,
image=image,
prompt=prompt,
image_edit_optional_request_params=image_edit_optional_request_params,
litellm_params=litellm_params,
logging_obj=logging_obj,
timeout=timeout,
extra_headers=extra_headers,
client=client if isinstance(client, AsyncHTTPHandler) else None,
)
# Sync version
if client is None or not isinstance(client, HTTPHandler):
sync_client = _get_httpx_client()
else:
sync_client = client
# Validate environment and get headers
headers = self.config.validate_environment(
api_key=api_key,
headers=image_edit_optional_request_params.get("extra_headers", {}) or {},
model=model,
)
if extra_headers:
headers.update(extra_headers)
# Get complete URL
complete_url = self.config.get_complete_url(
model=model,
api_base=api_base,
litellm_params=litellm_params_dict,
)
# Transform request
# Handle image list vs single image
if isinstance(image, list):
if not image:
raise BlackForestLabsError(status_code=400, message="No image provided")
image_input = image[0]
else:
image_input = image
data, _ = self.config.transform_image_edit_request(
model=model,
prompt=prompt or "",
image=image_input,
image_edit_optional_request_params=image_edit_optional_request_params,
litellm_params=litellm_params_dict,
headers=headers,
)
# Logging
logging_obj.pre_call(
input=prompt,
api_key="",
additional_args={
"complete_input_dict": data,
"api_base": complete_url,
"headers": headers,
},
)
# Make initial request
try:
response = sync_client.post(
url=complete_url,
headers=headers,
json=data,
timeout=timeout,
)
except Exception as e:
raise BlackForestLabsError(
status_code=500,
message=f"Request failed: {str(e)}",
)
# Poll for result
final_response = self._poll_for_result_sync(
initial_response=response,
headers=headers,
sync_client=sync_client,
)
# Transform response
return self.config.transform_image_edit_response(
model=model,
raw_response=final_response,
logging_obj=logging_obj,
)
async def async_image_edit(
self,
model: str,
image: Union[FileTypes, List[FileTypes]],
prompt: Optional[str],
image_edit_optional_request_params: Dict,
litellm_params: Union[GenericLiteLLMParams, Dict],
logging_obj: LiteLLMLoggingObj,
timeout: Optional[Union[float, httpx.Timeout]],
extra_headers: Optional[Dict[str, Any]] = None,
client: Optional[AsyncHTTPHandler] = None,
) -> ImageResponse:
"""
Async version of image edit.
"""
# Handle litellm_params as dict or object
if isinstance(litellm_params, dict):
api_key = litellm_params.get("api_key")
api_base = litellm_params.get("api_base")
litellm_params_dict = litellm_params
else:
api_key = litellm_params.api_key
api_base = litellm_params.api_base
litellm_params_dict = dict(litellm_params)
if client is None:
async_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders.BLACK_FOREST_LABS,
)
else:
async_client = client
# Validate environment and get headers
headers = self.config.validate_environment(
api_key=api_key,
headers=image_edit_optional_request_params.get("extra_headers", {}) or {},
model=model,
)
if extra_headers:
headers.update(extra_headers)
# Get complete URL
complete_url = self.config.get_complete_url(
model=model,
api_base=api_base,
litellm_params=litellm_params_dict,
)
# Transform request
if isinstance(image, list):
if not image:
raise BlackForestLabsError(status_code=400, message="No image provided")
image_input = image[0]
else:
image_input = image
data, _ = self.config.transform_image_edit_request(
model=model,
prompt=prompt or "",
image=image_input,
image_edit_optional_request_params=image_edit_optional_request_params,
litellm_params=litellm_params_dict,
headers=headers,
)
# Logging
logging_obj.pre_call(
input=prompt,
api_key="",
additional_args={
"complete_input_dict": data,
"api_base": complete_url,
"headers": headers,
},
)
# Make initial request
try:
response = await async_client.post(
url=complete_url,
headers=headers,
json=data,
timeout=timeout,
)
except Exception as e:
raise BlackForestLabsError(
status_code=500,
message=f"Request failed: {str(e)}",
)
# Poll for result
final_response = await self._poll_for_result_async(
initial_response=response,
headers=headers,
async_client=async_client,
)
# Transform response
return self.config.transform_image_edit_response(
model=model,
raw_response=final_response,
logging_obj=logging_obj,
)
def _poll_for_result_sync(
self,
initial_response: httpx.Response,
headers: dict,
sync_client: HTTPHandler,
max_wait: float = DEFAULT_MAX_POLLING_TIME,
interval: float = DEFAULT_POLLING_INTERVAL,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> httpx.Response:
"""
Poll BFL API until result is ready (sync version).
Args:
initial_response: The initial response containing polling_url
headers: Headers to use for polling (must include x-key)
sync_client: HTTP client
max_wait: Maximum time to wait in seconds
interval: Polling interval in seconds
timeout: Timeout for each individual polling request
Returns:
Final response with completed result
"""
# Validate initial response status code
if initial_response.status_code >= 400:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"BFL initial request failed: {initial_response.text}",
)
# Parse initial response to get polling URL
try:
response_data = initial_response.json()
except Exception as e:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"Error parsing initial response: {e}",
)
# Check for immediate errors
if "errors" in response_data:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"BFL error: {response_data['errors']}",
)
polling_url = response_data.get("polling_url")
if not polling_url:
raise BlackForestLabsError(
status_code=500,
message="No polling_url in BFL response",
)
# Get just the auth header for polling
polling_headers = {"x-key": headers.get("x-key", "")}
start_time = time.time()
verbose_logger.debug(f"BFL starting sync polling at {polling_url}")
while time.time() - start_time < max_wait:
response = sync_client.get(
url=polling_url,
headers=polling_headers,
)
if response.status_code != 200:
raise BlackForestLabsError(
status_code=response.status_code,
message=f"Polling failed: {response.text}",
)
data = response.json()
status = data.get("status")
verbose_logger.debug(f"BFL poll status: {status}")
if status == "Ready":
return response
elif status in ["Error", "Failed", "Content Moderated", "Request Moderated"]:
raise BlackForestLabsError(
status_code=400,
message=f"Image generation failed: {status}",
)
time.sleep(interval)
raise BlackForestLabsError(
status_code=408,
message=f"Polling timed out after {max_wait} seconds",
)
async def _poll_for_result_async(
self,
initial_response: httpx.Response,
headers: dict,
async_client: AsyncHTTPHandler,
max_wait: float = DEFAULT_MAX_POLLING_TIME,
interval: float = DEFAULT_POLLING_INTERVAL,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> httpx.Response:
"""
Poll BFL API until result is ready (async version).
"""
# Validate initial response status code
if initial_response.status_code >= 400:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"BFL initial request failed: {initial_response.text}",
)
# Parse initial response to get polling URL
try:
response_data = initial_response.json()
except Exception as e:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"Error parsing initial response: {e}",
)
# Check for immediate errors
if "errors" in response_data:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"BFL error: {response_data['errors']}",
)
polling_url = response_data.get("polling_url")
if not polling_url:
raise BlackForestLabsError(
status_code=500,
message="No polling_url in BFL response",
)
# Get just the auth header for polling
polling_headers = {"x-key": headers.get("x-key", "")}
start_time = time.time()
verbose_logger.debug(f"BFL starting async polling at {polling_url}")
while time.time() - start_time < max_wait:
response = await async_client.get(
url=polling_url,
headers=polling_headers,
)
if response.status_code != 200:
raise BlackForestLabsError(
status_code=response.status_code,
message=f"Polling failed: {response.text}",
)
data = response.json()
status = data.get("status")
verbose_logger.debug(f"BFL poll status: {status}")
if status == "Ready":
return response
elif status in ["Error", "Failed", "Content Moderated", "Request Moderated"]:
raise BlackForestLabsError(
status_code=400,
message=f"Image generation failed: {status}",
)
await asyncio.sleep(interval)
raise BlackForestLabsError(
status_code=408,
message=f"Polling timed out after {max_wait} seconds",
)
# Singleton instance for use in images/main.py
bfl_image_edit = BlackForestLabsImageEdit()

308
litellm/llms/black_forest_labs/image_edit/transformation.py Normal file
View file

@ -0,0 +1,308 @@
"""
Black Forest Labs Image Edit Configuration
Handles transformation between OpenAI-compatible format and Black Forest Labs API format
for image editing endpoints (flux-kontext-pro, flux-kontext-max, etc.).
API Reference: https://docs.bfl.ai/
"""
import base64
import time
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
import httpx
from httpx._types import RequestFiles
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import FileTypes, ImageObject, ImageResponse
from ..common_utils import (
DEFAULT_API_BASE,
IMAGE_EDIT_MODELS,
BlackForestLabsError,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class BlackForestLabsImageEditConfig(BaseImageEditConfig):
"""
Configuration for Black Forest Labs image editing.
Supports:
- flux-kontext-pro: General image editing with prompts
- flux-kontext-max: Premium quality editing
- flux-pro-1.0-fill: Inpainting with mask
- flux-pro-1.0-expand: Outpainting (expand image borders)
Note: HTTP requests and polling are handled by the handler (handler.py).
This class only handles data transformation.
"""
def get_supported_openai_params(self, model: str) -> List[str]:
"""
Return list of OpenAI params supported by Black Forest Labs.
Note: BFL uses different parameter names, these are mapped in map_openai_params.
"""
return [
"mask",
"seed",
"output_format",
"safety_tolerance",
"prompt_upsampling",
"aspect_ratio",
"steps",
"guidance",
"grow_mask",
"top",
"bottom",
"left",
"right",
]
def map_openai_params(
self,
image_edit_optional_params: ImageEditOptionalRequestParams,
model: str,
drop_params: bool,
) -> Dict:
"""
Map OpenAI parameters to Black Forest Labs parameters.
BFL-specific params are passed through directly.
"""
optional_params: Dict[str, Any] = {}
# Pass through BFL-specific params
bfl_params = [
"seed",
"output_format",
"safety_tolerance",
"prompt_upsampling",
# Kontext-specific
"aspect_ratio",
# Fill/Inpaint-specific
"steps",
"guidance",
"grow_mask",
# Expand-specific
"top",
"bottom",
"left",
"right",
]
# Convert TypedDict to regular dict for access
params_dict = dict(image_edit_optional_params)
for param in bfl_params:
if param in params_dict:
value = params_dict[param]
if value is not None:
optional_params[param] = value
# Set default output format
if "output_format" not in optional_params:
optional_params["output_format"] = "png"
return optional_params
def validate_environment(
self,
headers: dict,
model: str,
api_key: Optional[str] = None,
) -> dict:
"""
Validate environment and set up headers for Black Forest Labs.
BFL uses x-key header for authentication.
"""
final_api_key: Optional[str] = (
api_key
or get_secret_str("BFL_API_KEY")
or get_secret_str("BLACK_FOREST_LABS_API_KEY")
)
if not final_api_key:
raise BlackForestLabsError(
status_code=401,
message="BFL_API_KEY is not set. Please set it via environment variable or pass api_key parameter.",
)
headers["x-key"] = final_api_key
headers["Content-Type"] = "application/json"
headers["Accept"] = "application/json"
return headers
def use_multipart_form_data(self) -> bool:
"""
BFL uses JSON requests, not multipart/form-data.
"""
return False
def _get_model_endpoint(self, model: str) -> str:
"""
Get the API endpoint for a given model.
"""
# Remove provider prefix if present (e.g., "black_forest_labs/flux-kontext-pro")
model_name = model.lower()
if "/" in model_name:
model_name = model_name.split("/")[-1]
# Check if model is in our mapping
if model_name in IMAGE_EDIT_MODELS:
return IMAGE_EDIT_MODELS[model_name]
raise ValueError(
f"Unknown BFL image edit model: {model_name}. "
f"Supported models: {list(IMAGE_EDIT_MODELS.keys())}"
)
def get_complete_url(
self,
model: str,
api_base: Optional[str],
litellm_params: dict,
) -> str:
"""
Get the complete URL for the Black Forest Labs API request.
"""
base_url: str = (
api_base
or get_secret_str("BFL_API_BASE")
or DEFAULT_API_BASE
)
base_url = base_url.rstrip("/")
endpoint = self._get_model_endpoint(model)
return f"{base_url}{endpoint}"
def _read_image_bytes(self, image: Any) -> bytes:
"""Read image bytes from various input types."""
if isinstance(image, bytes):
return image
elif isinstance(image, list):
# If it's a list, take the first image
return self._read_image_bytes(image[0])
elif isinstance(image, str):
if image.startswith(("http://", "https://")):
# Download image from URL
response = httpx.get(image, timeout=60.0)
response.raise_for_status()
return response.content
else:
# Assume it's a file path
with open(image, "rb") as f:
return f.read()
elif hasattr(image, "read"):
# File-like object
pos = getattr(image, "tell", lambda: 0)()
if hasattr(image, "seek"):
image.seek(0)
data = image.read()
if hasattr(image, "seek"):
image.seek(pos)
return data
else:
raise ValueError(
f"Unsupported image type: {type(image)}. "
"Expected bytes, str (URL or file path), or file-like object."
)
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 OpenAI-style request to Black Forest Labs request format.
BFL uses JSON body with base64-encoded images, not multipart/form-data.
"""
# Read and encode image
image_bytes = self._read_image_bytes(image)
b64_image = base64.b64encode(image_bytes).decode("utf-8")
# Build request body
request_body: Dict[str, Any] = {
"prompt": prompt,
"input_image": b64_image,
}
# Add optional params (only BFL-recognized parameters)
bfl_request_params = [
"seed", "output_format", "safety_tolerance", "prompt_upsampling",
"aspect_ratio", "steps", "guidance", "grow_mask",
"top", "bottom", "left", "right",
]
for key, value in image_edit_optional_request_params.items():
if key in bfl_request_params and value is not None:
request_body[key] = value
# Handle mask if provided (for inpainting)
if "mask" in image_edit_optional_request_params:
mask = image_edit_optional_request_params["mask"]
mask_bytes = self._read_image_bytes(mask)
request_body["mask"] = base64.b64encode(mask_bytes).decode("utf-8")
# BFL uses JSON, not multipart - return empty files
return request_body, []
def transform_image_edit_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
) -> ImageResponse:
"""
Transform Black Forest Labs response to OpenAI-compatible ImageResponse.
This is called with the FINAL polled response (after handler does polling).
The response contains: {"status": "Ready", "result": {"sample": "https://..."}}
"""
try:
response_data = raw_response.json()
except Exception as e:
raise BlackForestLabsError(
status_code=raw_response.status_code,
message=f"Error parsing BFL response: {e}",
)
# Get image URL from result
image_url = response_data.get("result", {}).get("sample")
if not image_url:
raise BlackForestLabsError(
status_code=500,
message="No image URL in BFL result",
)
# Build ImageResponse
return ImageResponse(
created=int(time.time()),
data=[ImageObject(url=image_url)],
)
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BlackForestLabsError:
"""Return the appropriate error class for Black Forest Labs."""
return BlackForestLabsError(
status_code=status_code,
message=error_message,
)

12
litellm/llms/black_forest_labs/image_generation/__init__.py Normal file
View file

@ -0,0 +1,12 @@
from .handler import BlackForestLabsImageGeneration, bfl_image_generation
from .transformation import (
BlackForestLabsImageGenerationConfig,
get_black_forest_labs_image_generation_config,
)
__all__ = [
"BlackForestLabsImageGenerationConfig",
"get_black_forest_labs_image_generation_config",
"BlackForestLabsImageGeneration",
"bfl_image_generation",
]

440
litellm/llms/black_forest_labs/image_generation/handler.py Normal file
View file

@ -0,0 +1,440 @@
"""
Black Forest Labs Image Generation Handler
Handles image generation requests for Black Forest Labs models.
BFL uses an async polling pattern - the initial request returns a task ID,
then we poll until the result is ready.
"""
import asyncio
import time
from typing import Any, Dict, Optional, Union
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
_get_httpx_client,
get_async_httpx_client,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ImageResponse
from ..common_utils import (
DEFAULT_MAX_POLLING_TIME,
DEFAULT_POLLING_INTERVAL,
BlackForestLabsError,
)
from .transformation import BlackForestLabsImageGenerationConfig
class BlackForestLabsImageGeneration:
"""
Black Forest Labs Image Generation handler.
Handles the HTTP requests and polling logic, delegating data transformation
to the BlackForestLabsImageGenerationConfig class.
"""
def __init__(self):
self.config = BlackForestLabsImageGenerationConfig()
def image_generation(
self,
model: str,
prompt: str,
model_response: ImageResponse,
optional_params: Dict,
litellm_params: Union[GenericLiteLLMParams, Dict],
logging_obj: LiteLLMLoggingObj,
timeout: Optional[Union[float, httpx.Timeout]],
extra_headers: Optional[Dict[str, Any]] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
aimg_generation: bool = False,
) -> Union[ImageResponse, Any]:
"""
Main entry point for image generation requests.
Args:
model: The model to use (e.g., "black_forest_labs/flux-pro-1.1")
prompt: The text prompt for image generation
model_response: ImageResponse object to populate
optional_params: Optional parameters for the request
litellm_params: LiteLLM parameters including api_key, api_base
logging_obj: Logging object
timeout: Request timeout
extra_headers: Additional headers
client: HTTP client to use
aimg_generation: If True, return async coroutine
Returns:
ImageResponse or coroutine if aimg_generation=True
"""
# Handle litellm_params as dict or object
if isinstance(litellm_params, dict):
api_key = litellm_params.get("api_key")
api_base = litellm_params.get("api_base")
litellm_params_dict = litellm_params
else:
api_key = litellm_params.api_key
api_base = litellm_params.api_base
litellm_params_dict = dict(litellm_params)
if aimg_generation:
return self.async_image_generation(
model=model,
prompt=prompt,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params,
logging_obj=logging_obj,
timeout=timeout,
extra_headers=extra_headers,
client=client if isinstance(client, AsyncHTTPHandler) else None,
)
# Sync version
if client is None or not isinstance(client, HTTPHandler):
sync_client = _get_httpx_client()
else:
sync_client = client
# Validate environment and get headers
headers = self.config.validate_environment(
api_key=api_key,
headers={},
model=model,
messages=[],
optional_params=optional_params,
litellm_params=litellm_params_dict,
)
if extra_headers:
headers.update(extra_headers)
# Get complete URL
complete_url = self.config.get_complete_url(
api_base=api_base,
api_key=api_key,
model=model,
optional_params=optional_params,
litellm_params=litellm_params_dict,
)
# Transform request
data = self.config.transform_image_generation_request(
model=model,
prompt=prompt,
optional_params=optional_params,
litellm_params=litellm_params_dict,
headers=headers,
)
# Logging
logging_obj.pre_call(
input=prompt,
api_key="",
additional_args={
"complete_input_dict": data,
"api_base": complete_url,
"headers": headers,
},
)
# Make initial request
try:
response = sync_client.post(
url=complete_url,
headers=headers,
json=data,
timeout=timeout,
)
except Exception as e:
raise BlackForestLabsError(
status_code=500,
message=f"Request failed: {str(e)}",
)
# Poll for result
final_response = self._poll_for_result_sync(
initial_response=response,
headers=headers,
sync_client=sync_client,
)
# Transform response
return self.config.transform_image_generation_response(
model=model,
raw_response=final_response,
model_response=model_response,
logging_obj=logging_obj,
)
async def async_image_generation(
self,
model: str,
prompt: str,
model_response: ImageResponse,
optional_params: Dict,
litellm_params: Union[GenericLiteLLMParams, Dict],
logging_obj: LiteLLMLoggingObj,
timeout: Optional[Union[float, httpx.Timeout]],
extra_headers: Optional[Dict[str, Any]] = None,
client: Optional[AsyncHTTPHandler] = None,
) -> ImageResponse:
"""
Async version of image generation.
"""
# Handle litellm_params as dict or object
if isinstance(litellm_params, dict):
api_key = litellm_params.get("api_key")
api_base = litellm_params.get("api_base")
litellm_params_dict = litellm_params
else:
api_key = litellm_params.api_key
api_base = litellm_params.api_base
litellm_params_dict = dict(litellm_params)
if client is None:
async_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders.BLACK_FOREST_LABS,
)
else:
async_client = client
# Validate environment and get headers
headers = self.config.validate_environment(
api_key=api_key,
headers={},
model=model,
messages=[],
optional_params=optional_params,
litellm_params=litellm_params_dict,
)
if extra_headers:
headers.update(extra_headers)
# Get complete URL
complete_url = self.config.get_complete_url(
api_base=api_base,
api_key=api_key,
model=model,
optional_params=optional_params,
litellm_params=litellm_params_dict,
)
# Transform request
data = self.config.transform_image_generation_request(
model=model,
prompt=prompt,
optional_params=optional_params,
litellm_params=litellm_params_dict,
headers=headers,
)
# Logging
logging_obj.pre_call(
input=prompt,
api_key="",
additional_args={
"complete_input_dict": data,
"api_base": complete_url,
"headers": headers,
},
)
# Make initial request
try:
response = await async_client.post(
url=complete_url,
headers=headers,
json=data,
timeout=timeout,
)
except Exception as e:
raise BlackForestLabsError(
status_code=500,
message=f"Request failed: {str(e)}",
)
# Poll for result
final_response = await self._poll_for_result_async(
initial_response=response,
headers=headers,
async_client=async_client,
)
# Transform response
return self.config.transform_image_generation_response(
model=model,
raw_response=final_response,
model_response=model_response,
logging_obj=logging_obj,
)
def _poll_for_result_sync(
self,
initial_response: httpx.Response,
headers: dict,
sync_client: HTTPHandler,
max_wait: float = DEFAULT_MAX_POLLING_TIME,
interval: float = DEFAULT_POLLING_INTERVAL,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> httpx.Response:
"""
Poll BFL API until result is ready (sync version).
"""
# Validate initial response status code
if initial_response.status_code >= 400:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"BFL initial request failed: {initial_response.text}",
)
# Parse initial response to get polling URL
try:
response_data = initial_response.json()
except Exception as e:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"Error parsing initial response: {e}",
)
# Check for immediate errors
if "errors" in response_data:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"BFL error: {response_data['errors']}",
)
polling_url = response_data.get("polling_url")
if not polling_url:
raise BlackForestLabsError(
status_code=500,
message="No polling_url in BFL response",
)
# Get just the auth header for polling
polling_headers = {"x-key": headers.get("x-key", "")}
start_time = time.time()
verbose_logger.debug(f"BFL starting sync polling at {polling_url}")
while time.time() - start_time < max_wait:
response = sync_client.get(
url=polling_url,
headers=polling_headers,
)
if response.status_code != 200:
raise BlackForestLabsError(
status_code=response.status_code,
message=f"Polling failed: {response.text}",
)
data = response.json()
status = data.get("status")
verbose_logger.debug(f"BFL poll status: {status}")
if status == "Ready":
return response
elif status in ["Error", "Failed", "Content Moderated", "Request Moderated"]:
raise BlackForestLabsError(
status_code=400,
message=f"Image generation failed: {status}",
)
time.sleep(interval)
raise BlackForestLabsError(
status_code=408,
message=f"Polling timed out after {max_wait} seconds",
)
async def _poll_for_result_async(
self,
initial_response: httpx.Response,
headers: dict,
async_client: AsyncHTTPHandler,
max_wait: float = DEFAULT_MAX_POLLING_TIME,
interval: float = DEFAULT_POLLING_INTERVAL,
timeout: Optional[Union[float, httpx.Timeout]] = None,
) -> httpx.Response:
"""
Poll BFL API until result is ready (async version).
"""
# Validate initial response status code
if initial_response.status_code >= 400:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"BFL initial request failed: {initial_response.text}",
)
# Parse initial response to get polling URL
try:
response_data = initial_response.json()
except Exception as e:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"Error parsing initial response: {e}",
)
# Check for immediate errors
if "errors" in response_data:
raise BlackForestLabsError(
status_code=initial_response.status_code,
message=f"BFL error: {response_data['errors']}",
)
polling_url = response_data.get("polling_url")
if not polling_url:
raise BlackForestLabsError(
status_code=500,
message="No polling_url in BFL response",
)
# Get just the auth header for polling
polling_headers = {"x-key": headers.get("x-key", "")}
start_time = time.time()
verbose_logger.debug(f"BFL starting async polling at {polling_url}")
while time.time() - start_time < max_wait:
response = await async_client.get(
url=polling_url,
headers=polling_headers,
)
if response.status_code != 200:
raise BlackForestLabsError(
status_code=response.status_code,
message=f"Polling failed: {response.text}",
)
data = response.json()
status = data.get("status")
verbose_logger.debug(f"BFL poll status: {status}")
if status == "Ready":
return response
elif status in ["Error", "Failed", "Content Moderated", "Request Moderated"]:
raise BlackForestLabsError(
status_code=400,
message=f"Image generation failed: {status}",
)
await asyncio.sleep(interval)
raise BlackForestLabsError(
status_code=408,
message=f"Polling timed out after {max_wait} seconds",
)
# Singleton instance for use in images/main.py
bfl_image_generation = BlackForestLabsImageGeneration()

324
litellm/llms/black_forest_labs/image_generation/transformation.py Normal file
View file

@ -0,0 +1,324 @@
"""
Black Forest Labs Image Generation Configuration
Handles transformation between OpenAI-compatible format and Black Forest Labs API format
for image generation endpoints (flux-pro-1.1, flux-pro-1.1-ultra, flux-dev, flux-pro).
API Reference: https://docs.bfl.ai/
"""
import time
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
import httpx
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIImageGenerationOptionalParams,
)
from litellm.types.utils import ImageObject, ImageResponse
from ..common_utils import (
DEFAULT_API_BASE,
IMAGE_GENERATION_MODELS,
BlackForestLabsError,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig):
"""
Configuration for Black Forest Labs image generation (text-to-image).
Supports:
- flux-pro-1.1: Fast & reliable standard generation
- flux-pro-1.1-ultra: Ultra high-resolution (up to 4MP)
- flux-dev: Development/open-source variant
- flux-pro: Original pro model
Note: HTTP requests and polling are handled by the handler (handler.py).
This class only handles data transformation.
"""
def get_supported_openai_params(
self, model: str
) -> List[OpenAIImageGenerationOptionalParams]:
"""
Return list of OpenAI params supported by Black Forest Labs.
Note: BFL uses different parameter names, these are mapped in map_openai_params.
"""
return [
"n", # Number of images (BFL returns 1 per request, but ultra supports up to 4)
"size", # Maps to width/height or aspect_ratio
"quality", # Maps to raw mode for ultra
"seed",
"output_format",
"safety_tolerance",
"prompt_upsampling",
"raw",
"num_images",
"image_url",
"image_prompt_strength",
"aspect_ratio",
]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
"""
Map OpenAI parameters to Black Forest Labs parameters.
BFL-specific params are passed through directly.
"""
supported_params = self.get_supported_openai_params(model)
for k, v in non_default_params.items():
if k in optional_params:
continue
if k in supported_params:
# Map OpenAI 'size' to BFL width/height
if k == "size" and v:
self._map_size_param(v, optional_params)
elif k == "n":
if "ultra" in model.lower():
optional_params["num_images"] = v
# non-ultra: silently skip (n=1 is BFL default)
elif k == "quality":
if v == "hd" and "ultra" in model.lower():
optional_params["raw"] = True
# other quality values have no BFL mapping
else:
optional_params[k] = v
elif not drop_params:
raise ValueError(
f"Parameter {k} is not supported for model {model}. "
f"Supported parameters are {supported_params}. "
f"Set drop_params=True to drop unsupported parameters."
)
return optional_params
def _map_size_param(self, size: str, optional_params: dict) -> None:
"""Map OpenAI size parameter to BFL width/height."""
# Common size mappings
size_mapping = {
"1024x1024": (1024, 1024),
"1792x1024": (1792, 1024),
"1024x1792": (1024, 1792),
"512x512": (512, 512),
"256x256": (256, 256),
}
if size in size_mapping:
width, height = size_mapping[size]
optional_params["width"] = width
optional_params["height"] = height
elif "x" in size:
# Parse custom size
try:
width, height = map(int, size.lower().split("x"))
optional_params["width"] = width
optional_params["height"] = height
except ValueError:
raise ValueError(
f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')."
)
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
"""
Validate environment and set up headers for Black Forest Labs.
BFL uses x-key header for authentication.
"""
final_api_key: Optional[str] = (
api_key
or get_secret_str("BFL_API_KEY")
or get_secret_str("BLACK_FOREST_LABS_API_KEY")
)
if not final_api_key:
raise BlackForestLabsError(
status_code=401,
message="BFL_API_KEY is not set. Please set it via environment variable or pass api_key parameter.",
)
headers["x-key"] = final_api_key
headers["Content-Type"] = "application/json"
headers["Accept"] = "application/json"
return headers
def _get_model_endpoint(self, model: str) -> str:
"""
Get the API endpoint for a given model.
"""
# Remove provider prefix if present (e.g., "black_forest_labs/flux-pro-1.1")
model_name = model.lower()
if "/" in model_name:
model_name = model_name.split("/")[-1]
# Check if model is in our mapping
if model_name in IMAGE_GENERATION_MODELS:
return IMAGE_GENERATION_MODELS[model_name]
raise ValueError(
f"Unknown BFL image generation model: {model_name}. "
f"Supported models: {list(IMAGE_GENERATION_MODELS.keys())}"
)
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
"""
Get the complete URL for the Black Forest Labs API request.
"""
base_url: str = (
api_base or get_secret_str("BFL_API_BASE") or DEFAULT_API_BASE
)
base_url = base_url.rstrip("/")
endpoint = self._get_model_endpoint(model)
return f"{base_url}{endpoint}"
def transform_image_generation_request(
self,
model: str,
prompt: str,
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
"""
Transform OpenAI-style request to Black Forest Labs request format.
https://docs.bfl.ai/flux_models/flux_1_1_pro
"""
# Build request body with prompt
request_body: Dict[str, Any] = {
"prompt": prompt,
}
# BFL-specific params that can be passed through
bfl_params = [
"width",
"height",
"aspect_ratio",
"seed",
"output_format",
"safety_tolerance",
"prompt_upsampling",
# Ultra-specific
"raw",
"num_images",
"image_url",
"image_prompt_strength",
]
for param in bfl_params:
if param in optional_params and optional_params[param] is not None:
request_body[param] = optional_params[param]
# Set default output format if not specified
if "output_format" not in request_body:
request_body["output_format"] = "png"
return request_body
def transform_image_generation_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: LiteLLMLoggingObj,
**kwargs,
) -> ImageResponse:
"""
Transform Black Forest Labs response to OpenAI-compatible ImageResponse.
This is called with the FINAL polled response (after handler does polling).
The response contains: {"status": "Ready", "result": {"sample": "https://..."}}
"""
try:
response_data = raw_response.json()
except Exception as e:
raise BlackForestLabsError(
status_code=raw_response.status_code,
message=f"Error parsing BFL response: {e}",
)
result = response_data.get("result", {})
if not model_response.data:
model_response.data = []
# Handle single image (sample) or multiple images
if isinstance(result, dict) and "sample" in result:
model_response.data.append(ImageObject(url=result["sample"]))
elif isinstance(result, list):
# Multiple images returned
for img in result:
if isinstance(img, str):
model_response.data.append(ImageObject(url=img))
elif isinstance(img, dict) and "url" in img:
model_response.data.append(ImageObject(url=img["url"]))
if not model_response.data:
raise BlackForestLabsError(
status_code=500,
message="No image URL in BFL result",
)
model_response.created = int(time.time())
return model_response
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BlackForestLabsError:
"""Return the appropriate error class for Black Forest Labs."""
return BlackForestLabsError(
status_code=status_code,
message=error_message,
)
def get_black_forest_labs_image_generation_config(
model: str,
) -> BlackForestLabsImageGenerationConfig:
"""
Get the appropriate image generation config for a Black Forest Labs model.
Currently returns a single config class, but can be extended
for model-specific configurations if needed.
"""
return BlackForestLabsImageGenerationConfig()

152
litellm/llms/mistral/audio_transcription/transformation.py Normal file
View file

@ -0,0 +1,152 @@
"""
Support for Mistral Voxtral audio transcription via ``/v1/audio/transcriptions``.
API reference: https://docs.mistral.ai/api/#tag/audio/operation/audio_transcriptions_v1_audio_transcriptions_post
"""
from typing import List, Optional, Union
import httpx
from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
from litellm.llms.base_llm.audio_transcription.transformation import (
AudioTranscriptionRequestData,
BaseAudioTranscriptionConfig,
)
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIAudioTranscriptionOptionalParams,
)
from litellm.types.utils import FileTypes, TranscriptionResponse
class MistralAudioTranscriptionException(BaseLLMException):
pass
class MistralAudioTranscriptionConfig(BaseAudioTranscriptionConfig):
def get_supported_openai_params(
self, model: str
) -> List[OpenAIAudioTranscriptionOptionalParams]:
return [
"language",
"temperature",
"timestamp_granularities",
"response_format",
]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
supported_params = self.get_supported_openai_params(model)
for k, v in non_default_params.items():
if k in supported_params:
optional_params[k] = v
return optional_params
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
api_base = (
"https://api.mistral.ai/v1"
if api_base is None
else api_base.rstrip("/")
)
return f"{api_base}/audio/transcriptions"
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BaseLLMException:
return MistralAudioTranscriptionException(
message=error_message,
status_code=status_code,
headers=headers,
)
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
if api_key is None:
api_key = get_secret_str("MISTRAL_API_KEY")
default_headers = {
"Authorization": f"Bearer {api_key}",
"accept": "application/json",
}
default_headers.update(headers or {})
return default_headers
def transform_audio_transcription_request(
self,
model: str,
audio_file: FileTypes,
optional_params: dict,
litellm_params: dict,
) -> AudioTranscriptionRequestData:
processed_audio = process_audio_file(audio_file)
form_fields: dict = {
"model": model,
}
# OpenAI-compatible params
for key in self.get_supported_openai_params(model):
value = optional_params.get(key)
if value is not None:
form_fields[key] = value
# Mistral-specific params (e.g. diarize)
provider_specific_params = self.get_provider_specific_params(
model=model,
optional_params=optional_params,
openai_params=self.get_supported_openai_params(model),
)
for key, value in provider_specific_params.items():
form_fields[key] = str(value).lower() if isinstance(value, bool) else str(value)
files = {
"file": (
processed_audio.filename,
processed_audio.file_content,
processed_audio.content_type,
)
}
return AudioTranscriptionRequestData(data=form_fields, files=files)
def transform_audio_transcription_response(
self,
raw_response: httpx.Response,
) -> TranscriptionResponse:
try:
response_json = raw_response.json()
except Exception:
raise MistralAudioTranscriptionException(
message=raw_response.text,
status_code=raw_response.status_code,
headers=raw_response.headers,
)
text = response_json.get("text") or ""
response = TranscriptionResponse(text=text)
response._hidden_params = response_json
return response

8
litellm/llms/openai/chat/gpt_transformation.py
View file

@ -58,6 +58,7 @@ from ..common_utils import OpenAIError
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.types.llms.openai import ChatCompletionToolParam
LiteLLMLoggingObj = _LiteLLMLoggingObj
@ -758,6 +759,13 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
def get_base_model(model: Optional[str] = None) -> Optional[str]:
return model
def get_token_counter(self) -> Optional["BaseTokenCounter"]:
from litellm.llms.openai.responses.count_tokens.token_counter import (
OpenAITokenCounter,
)
return OpenAITokenCounter()
def get_model_response_iterator(
self,
streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse],

19
litellm/llms/openai/responses/count_tokens/__init__.py Normal file
View file

@ -0,0 +1,19 @@
"""
OpenAI Responses API token counting implementation.
"""
from litellm.llms.openai.responses.count_tokens.handler import (
OpenAICountTokensHandler,
)
from litellm.llms.openai.responses.count_tokens.token_counter import (
OpenAITokenCounter,
)
from litellm.llms.openai.responses.count_tokens.transformation import (
OpenAICountTokensConfig,
)
__all__ = [
"OpenAICountTokensHandler",
"OpenAICountTokensConfig",
"OpenAITokenCounter",
]

105
litellm/llms/openai/responses/count_tokens/handler.py Normal file
View file

@ -0,0 +1,105 @@
"""
OpenAI Responses API token counting handler.
Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint.
"""
import json
from typing import Any, Dict, List, Optional, Union
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.llms.openai.common_utils import OpenAIError
from litellm.llms.openai.responses.count_tokens.transformation import (
OpenAICountTokensConfig,
)
class OpenAICountTokensHandler(OpenAICountTokensConfig):
"""
Handler for OpenAI Responses API token counting requests.
"""
async def handle_count_tokens_request(
self,
model: str,
input: Union[str, List[Any]],
api_key: str,
api_base: Optional[str] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
tools: Optional[List[Dict[str, Any]]] = None,
instructions: Optional[str] = None,
) -> Dict[str, Any]:
"""
Handle a token counting request to OpenAI's Responses API.
Returns:
Dictionary containing {"input_tokens": <number>}
Raises:
OpenAIError: If the API request fails
"""
try:
self.validate_request(model, input)
verbose_logger.debug(
f"Processing OpenAI CountTokens request for model: {model}"
)
request_body = self.transform_request_to_count_tokens(
model=model,
input=input,
tools=tools,
instructions=instructions,
)
endpoint_url = self.get_openai_count_tokens_endpoint(api_base)
verbose_logger.debug(f"Making request to: {endpoint_url}")
headers = self.get_required_headers(api_key)
async_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders.OPENAI
)
request_timeout = timeout if timeout is not None else litellm.request_timeout
response = await async_client.post(
endpoint_url,
headers=headers,
json=request_body,
timeout=request_timeout,
)
verbose_logger.debug(f"Response status: {response.status_code}")
if response.status_code != 200:
error_text = response.text
verbose_logger.error(f"OpenAI API error: {error_text}")
raise OpenAIError(
status_code=response.status_code,
message=error_text,
)
openai_response = response.json()
verbose_logger.debug(f"OpenAI response: {openai_response}")
return openai_response
except OpenAIError:
raise
except httpx.HTTPStatusError as e:
verbose_logger.error(f"HTTP error in CountTokens handler: {str(e)}")
raise OpenAIError(
status_code=e.response.status_code,
message=e.response.text,
)
except (httpx.RequestError, json.JSONDecodeError, ValueError) as e:
verbose_logger.error(f"Error in CountTokens handler: {str(e)}")
raise OpenAIError(
status_code=500,
message=f"CountTokens processing error: {str(e)}",
)

118
litellm/llms/openai/responses/count_tokens/token_counter.py Normal file
View file

@ -0,0 +1,118 @@
"""
OpenAI Token Counter implementation using the Responses API /input_tokens endpoint.
"""
import os
from typing import Any, Dict, List, Optional
from litellm._logging import verbose_logger
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.llms.openai.common_utils import OpenAIError
from litellm.llms.openai.responses.count_tokens.handler import (
OpenAICountTokensHandler,
)
from litellm.llms.openai.responses.count_tokens.transformation import (
OpenAICountTokensConfig,
)
from litellm.types.utils import LlmProviders, TokenCountResponse
# Global handler instance - reuse across all token counting requests
openai_count_tokens_handler = OpenAICountTokensHandler()
class OpenAITokenCounter(BaseTokenCounter):
"""Token counter implementation for OpenAI provider using the Responses API."""
def should_use_token_counting_api(
self,
custom_llm_provider: Optional[str] = None,
) -> bool:
return custom_llm_provider == LlmProviders.OPENAI.value
async def count_tokens(
self,
model_to_use: str,
messages: Optional[List[Dict[str, Any]]],
contents: Optional[List[Dict[str, Any]]],
deployment: Optional[Dict[str, Any]] = None,
request_model: str = "",
tools: Optional[List[Dict[str, Any]]] = None,
system: Optional[Any] = None,
) -> Optional[TokenCountResponse]:
"""
Count tokens using OpenAI's Responses API /input_tokens endpoint.
"""
if not messages:
return None
deployment = deployment or {}
litellm_params = deployment.get("litellm_params", {})
# Get OpenAI API key from deployment config or environment
api_key = litellm_params.get("api_key")
if not api_key:
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
verbose_logger.warning("No OpenAI API key found for token counting")
return None
api_base = litellm_params.get("api_base")
# Convert chat messages to Responses API input format
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(
messages
)
# Use system param if instructions not extracted from messages
if instructions is None and system is not None:
instructions = system if isinstance(system, str) else str(system)
# If no input items were produced (e.g., system-only messages), fall back to local counting
if not input_items:
return None
try:
result = await openai_count_tokens_handler.handle_count_tokens_request(
model=model_to_use,
input=input_items if input_items is not None else [],
api_key=api_key,
api_base=api_base,
tools=tools,
instructions=instructions,
)
if result is not None:
return TokenCountResponse(
total_tokens=result.get("input_tokens", 0),
request_model=request_model,
model_used=model_to_use,
tokenizer_type="openai_api",
original_response=result,
)
except OpenAIError as e:
verbose_logger.warning(
f"OpenAI CountTokens API error: status={e.status_code}, message={e.message}"
)
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="openai_api",
error=True,
error_message=e.message,
status_code=e.status_code,
)
except Exception as e:
verbose_logger.warning(f"Error calling OpenAI CountTokens API: {e}")
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="openai_api",
error=True,
error_message=str(e),
status_code=500,
)
return None

158
litellm/llms/openai/responses/count_tokens/transformation.py Normal file
View file

@ -0,0 +1,158 @@
"""
OpenAI Responses API token counting transformation logic.
This module handles the transformation of requests to OpenAI's /v1/responses/input_tokens endpoint.
"""
from typing import Any, Dict, List, Optional, Union
class OpenAICountTokensConfig:
"""
Configuration and transformation logic for OpenAI Responses API token counting.
OpenAI Responses API Token Counting Specification:
- Endpoint: POST https://api.openai.com/v1/responses/input_tokens
- Response: {"input_tokens": <number>}
"""
def get_openai_count_tokens_endpoint(self, api_base: Optional[str] = None) -> str:
base = api_base or "https://api.openai.com/v1"
base = base.rstrip("/")
return f"{base}/responses/input_tokens"
def transform_request_to_count_tokens(
self,
model: str,
input: Union[str, List[Any]],
tools: Optional[List[Dict[str, Any]]] = None,
instructions: Optional[str] = None,
) -> Dict[str, Any]:
"""
Transform request to OpenAI Responses API token counting format.
The Responses API uses `input` (not `messages`) and `instructions` (not `system`).
"""
request: Dict[str, Any] = {
"model": model,
"input": input,
}
if instructions is not None:
request["instructions"] = instructions
if tools is not None:
request["tools"] = self._transform_tools_for_responses_api(tools)
return request
def get_required_headers(self, api_key: str) -> Dict[str, str]:
return {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}",
}
def validate_request(
self, model: str, input: Union[str, List[Any]]
) -> None:
if not model:
raise ValueError("model parameter is required")
if not input:
raise ValueError("input parameter is required")
@staticmethod
def _transform_tools_for_responses_api(
tools: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
"""
Transform OpenAI chat tools format to Responses API tools format.
Chat format: {"type": "function", "function": {"name": "...", "parameters": {...}}}
Responses format: {"type": "function", "name": "...", "parameters": {...}}
"""
transformed = []
for tool in tools:
if tool.get("type") == "function" and "function" in tool:
func = tool["function"]
item: Dict[str, Any] = {
"type": "function",
"name": func.get("name", ""),
"description": func.get("description", ""),
"parameters": func.get("parameters", {}),
}
if "strict" in func:
item["strict"] = func["strict"]
transformed.append(item)
else:
# Pass through non-function tools (e.g., web_search, file_search)
transformed.append(tool)
return transformed
@staticmethod
def messages_to_responses_input(
messages: List[Dict[str, Any]],
) -> tuple:
"""
Convert standard chat messages format to OpenAI Responses API input format.
Returns:
(input_items, instructions) tuple where instructions is extracted
from system/developer messages.
"""
input_items: List[Dict[str, Any]] = []
instructions_parts: List[str] = []
for msg in messages:
role = msg.get("role", "")
content = msg.get("content") or ""
if role in ("system", "developer"):
# Extract system/developer messages as instructions
if isinstance(content, str):
instructions_parts.append(content)
elif isinstance(content, list):
# Handle content blocks - extract text
text_parts = []
for block in content:
if isinstance(block, dict) and block.get("type") == "text":
text_parts.append(block.get("text", ""))
elif isinstance(block, str):
text_parts.append(block)
instructions_parts.append("\n".join(text_parts))
elif role == "user":
if isinstance(content, list):
# Extract text from content blocks for Responses API
text_parts = []
for block in content:
if isinstance(block, dict) and block.get("type") == "text":
text_parts.append(block.get("text", ""))
elif isinstance(block, str):
text_parts.append(block)
content = "\n".join(text_parts)
input_items.append({"role": "user", "content": content})
elif role == "assistant":
# Map tool_calls to Responses API function_call items
tool_calls = msg.get("tool_calls")
if content:
input_items.append({"role": "assistant", "content": content})
if tool_calls:
for tc in tool_calls:
func = tc.get("function", {})
input_items.append({
"type": "function_call",
"call_id": tc.get("id", ""),
"name": func.get("name", ""),
"arguments": func.get("arguments", ""),
})
elif not content:
input_items.append({"role": "assistant", "content": content})
elif role == "tool":
input_items.append({
"type": "function_call_output",
"call_id": msg.get("tool_call_id", ""),
"output": content if isinstance(content, str) else str(content),
})
instructions = "\n".join(instructions_parts) if instructions_parts else None
return input_items, instructions

31
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
View file

@ -4,13 +4,16 @@ import httpx
import litellm
from litellm.caching.caching import Cache, LiteLLMCacheType
from litellm.constants import MINIMUM_PROMPT_CACHE_TOKEN_COUNT
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
get_async_httpx_client,
)
from litellm._logging import verbose_logger
from litellm.llms.openai.openai import AllMessageValues
from litellm.utils import is_prompt_caching_valid_prompt
from litellm.types.llms.vertex_ai import (
CachedContentListAllResponseBody,
VertexAICachedContentResponseObject,
@ -314,6 +317,20 @@ class ContextCachingEndpoints(VertexBase):
if len(cached_messages) == 0:
return messages, optional_params, None
# Gemini requires a minimum of 1024 tokens for context caching.
# Skip caching if the cached content is too small to avoid API errors.
if not is_prompt_caching_valid_prompt(
model=model,
messages=cached_messages,
custom_llm_provider=custom_llm_provider,
):
verbose_logger.debug(
"Vertex AI context caching: cached content is below minimum token "
"count (%d). Skipping context caching.",
MINIMUM_PROMPT_CACHE_TOKEN_COUNT,
)
return messages, optional_params, None
tools = optional_params.pop("tools", None)
## AUTHORIZATION ##
@ -446,6 +463,20 @@ class ContextCachingEndpoints(VertexBase):
if len(cached_messages) == 0:
return messages, optional_params, None
# Gemini requires a minimum of 1024 tokens for context caching.
# Skip caching if the cached content is too small to avoid API errors.
if not is_prompt_caching_valid_prompt(
model=model,
messages=cached_messages,
custom_llm_provider=custom_llm_provider,
):
verbose_logger.debug(
"Vertex AI context caching: cached content is below minimum token "
"count (%d). Skipping context caching.",
MINIMUM_PROMPT_CACHE_TOKEN_COUNT,
)
return messages, optional_params, None
tools = optional_params.pop("tools", None)
## AUTHORIZATION ##

71
litellm/llms/vertex_ai/gemini/transformation.py
View file

@ -77,6 +77,60 @@ def _convert_detail_to_media_resolution_enum(
return None
def _get_highest_media_resolution(
current: Optional[str], new_detail: Optional[str]
) -> Optional[str]:
"""
Compare two media resolution values and return the highest one.
Resolution hierarchy: ultra_high > high > medium > low > None
"""
resolution_priority = {"ultra_high": 4, "high": 3, "medium": 2, "low": 1}
current_priority = resolution_priority.get(current, 0) if current else 0
new_priority = resolution_priority.get(new_detail, 0) if new_detail else 0
if new_priority > current_priority:
return new_detail
return current
def _extract_max_media_resolution_from_messages(
messages: List[AllMessageValues],
) -> Optional[str]:
"""
Extract the highest media resolution (detail) from image content in messages.
This is used to set the global media_resolution in generation_config for
Gemini 2.x models which don't support per-part media resolution.
Args:
messages: List of messages in OpenAI format
Returns:
The highest detail level found ("high", "low", or None)
"""
max_resolution: Optional[str] = None
for msg in messages:
content = msg.get("content")
if isinstance(content, list):
for item in content:
if not isinstance(item, dict):
continue
detail: Optional[str] = None
if item.get("type") == "image_url":
image_url = item.get("image_url")
if isinstance(image_url, dict):
detail = image_url.get("detail")
elif item.get("type") == "file":
file_obj = item.get("file")
if isinstance(file_obj, dict):
detail = file_obj.get("detail")
if detail:
max_resolution = _get_highest_media_resolution(
max_resolution, detail
)
return max_resolution
def _apply_gemini_3_metadata(
part: PartType,
model: Optional[str],
@ -84,7 +138,7 @@ def _apply_gemini_3_metadata(
video_metadata: Optional[Dict[str, Any]],
) -> PartType:
"""
Apply the unique media_resolution and video_metadata parameters of Gemini 3+
Apply the unique media_resolution and video_metadata parameters of Gemini 3+
"""
if model is None:
return part
@ -547,7 +601,7 @@ def _pop_and_merge_extra_body(data: RequestBody, optional_params: dict) -> None:
data_dict[k] = v
def _transform_request_body(
def _transform_request_body( # noqa: PLR0915
messages: List[AllMessageValues],
model: str,
optional_params: dict,
@ -623,6 +677,19 @@ def _transform_request_body(
generation_config: Optional[GenerationConfig] = GenerationConfig(
**filtered_params
)
# For Gemini 2.x models, add media_resolution to generation_config (global)
# Gemini 3+ supports per-part media_resolution, but 2.x only supports global
# Gemini 1.x does not support mediaResolution at all
if "gemini-2" in model:
max_media_resolution = _extract_max_media_resolution_from_messages(messages)
if max_media_resolution:
media_resolution_value = _convert_detail_to_media_resolution_enum(
max_media_resolution
)
if media_resolution_value and generation_config is not None:
generation_config["mediaResolution"] = media_resolution_value["level"]
data = RequestBody(contents=content)
if system_instructions is not None:
data["system_instruction"] = system_instructions

11
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -2363,7 +2363,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
async def make_call(
client: Optional[AsyncHTTPHandler],
client: Optional[AsyncHTTPHandler], # module-level client
gemini_client: Optional[AsyncHTTPHandler], # if passed by user
api_base: str,
headers: dict,
data: str,
@ -2371,6 +2372,8 @@ async def make_call(
messages: list,
logging_obj,
):
if gemini_client is not None:
client = gemini_client
if client is None:
client = get_async_httpx_client(
llm_provider=litellm.LlmProviders.VERTEX_AI,
@ -2542,7 +2545,11 @@ class VertexLLM(VertexBase):
completion_stream=None,
make_call=partial(
make_call,
client=client,
gemini_client=(
client
if client is not None and isinstance(client, AsyncHTTPHandler)
else None
),
api_base=api_base,
headers=headers,
data=request_body_str,

105
litellm/main.py
View file

@ -7567,6 +7567,111 @@ def stream_chunk_builder( # noqa: PLR0915
)
########## Token Counting API ##########
async def acount_tokens(
model: str,
messages: Optional[List[Dict[str, Any]]] = None,
tools: Optional[List[Dict[str, Any]]] = None,
system: Optional[str] = None,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> "TokenCountResponse":
"""
Count tokens for a given model and messages using provider-specific APIs.
Routes to the appropriate provider's token counting API (OpenAI, Anthropic, etc.)
for exact token counts. Falls back to local tiktoken-based counting for unsupported providers.
Args:
model: The model identifier (e.g., "openai/gpt-4o", "anthropic/claude-3-5-sonnet-20241022")
messages: The messages to count tokens for (standard chat format)
tools: Optional tools/functions to include in token count
system: Optional system message/instructions
api_key: Optional API key (falls back to environment variable)
api_base: Optional custom API base URL
Returns:
TokenCountResponse with total_tokens and metadata
"""
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.types.utils import LlmProviders, TokenCountResponse
from litellm.utils import ProviderConfigManager
# Determine provider from model string
resolved_model, custom_llm_provider, dynamic_api_key, dynamic_api_base = (
get_llm_provider(
model=model,
api_base=api_base,
api_key=api_key,
)
)
# Use dynamic key/base if not explicitly provided
if api_key is None:
api_key = dynamic_api_key
if api_base is None:
api_base = dynamic_api_base
# Build deployment dict for the token counter
deployment: Dict[str, Any] = {
"litellm_params": {
"model": model,
"api_key": api_key,
"api_base": api_base,
}
}
# Try to get provider-specific token counter
try:
llm_provider_enum = LlmProviders(custom_llm_provider)
provider_model_info = ProviderConfigManager.get_provider_model_info(
model=model, provider=llm_provider_enum
)
if provider_model_info is not None:
token_counter_instance = provider_model_info.get_token_counter()
if (
token_counter_instance is not None
and token_counter_instance.should_use_token_counting_api(
custom_llm_provider
)
):
result = await token_counter_instance.count_tokens(
model_to_use=resolved_model,
messages=messages,
contents=None,
deployment=deployment,
request_model=model,
tools=tools,
system=system,
)
if result is not None and not result.error:
return result
except Exception as e:
verbose_logger.debug(
f"Provider token counting failed for model={model}, falling back to local: {e}"
)
# Fallback to local tiktoken-based token counting
fallback_messages = messages or []
if system and fallback_messages:
fallback_messages = [{"role": "system", "content": system}] + fallback_messages
local_count = litellm.token_counter(
model=model,
messages=fallback_messages,
tools=tools,
)
return TokenCountResponse(
total_tokens=local_count,
request_model=model,
model_used=resolved_model,
tokenizer_type="local_tokenizer",
)
# Cache for encoding to avoid repeated __getattr__ calls
_encoding_cache: Optional[Any] = None

0
litellm/proxy/_experimental/out/404.html → litellm/proxy/_experimental/out/404/index.html
View file

0
litellm/proxy/_experimental/out/_not-found.html → litellm/proxy/_experimental/out/_not-found/index.html
View file

0
litellm/proxy/_experimental/out/api-reference.html → litellm/proxy/_experimental/out/api-reference/index.html
View file

0
litellm/proxy/_experimental/out/experimental/api-playground.html → litellm/proxy/_experimental/out/experimental/api-playground/index.html
View file

0
litellm/proxy/_experimental/out/experimental/budgets.html → litellm/proxy/_experimental/out/experimental/budgets/index.html
View file

0
litellm/proxy/_experimental/out/experimental/caching.html → litellm/proxy/_experimental/out/experimental/caching/index.html
View file

0
litellm/proxy/_experimental/out/experimental/claude-code-plugins.html → litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html
View file

0
litellm/proxy/_experimental/out/experimental/old-usage.html → litellm/proxy/_experimental/out/experimental/old-usage/index.html
View file

0
litellm/proxy/_experimental/out/experimental/prompts.html → litellm/proxy/_experimental/out/experimental/prompts/index.html
View file

0
litellm/proxy/_experimental/out/experimental/tag-management.html → litellm/proxy/_experimental/out/experimental/tag-management/index.html
View file

0
litellm/proxy/_experimental/out/guardrails.html → litellm/proxy/_experimental/out/guardrails/index.html
View file

0
litellm/proxy/_experimental/out/login.html → litellm/proxy/_experimental/out/login/index.html
View file

0
litellm/proxy/_experimental/out/logs.html → litellm/proxy/_experimental/out/logs/index.html
View file

0
litellm/proxy/_experimental/out/mcp/oauth/callback.html → litellm/proxy/_experimental/out/mcp/oauth/callback/index.html
View file

0
litellm/proxy/_experimental/out/model-hub.html → litellm/proxy/_experimental/out/model-hub/index.html
View file

0
litellm/proxy/_experimental/out/model_hub.html → litellm/proxy/_experimental/out/model_hub/index.html
View file

0
litellm/proxy/_experimental/out/model_hub_table.html → litellm/proxy/_experimental/out/model_hub_table/index.html
View file

0
litellm/proxy/_experimental/out/models-and-endpoints.html → litellm/proxy/_experimental/out/models-and-endpoints/index.html
View file

0
litellm/proxy/_experimental/out/onboarding.html → litellm/proxy/_experimental/out/onboarding/index.html
View file

0
litellm/proxy/_experimental/out/organizations.html → litellm/proxy/_experimental/out/organizations/index.html
View file

0
litellm/proxy/_experimental/out/playground.html → litellm/proxy/_experimental/out/playground/index.html
View file

0
litellm/proxy/_experimental/out/policies.html → litellm/proxy/_experimental/out/policies/index.html
View file

0
litellm/proxy/_experimental/out/settings/admin-settings.html → litellm/proxy/_experimental/out/settings/admin-settings/index.html
View file

0
litellm/proxy/_experimental/out/settings/logging-and-alerts.html → litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html
View file

0
litellm/proxy/_experimental/out/settings/router-settings.html → litellm/proxy/_experimental/out/settings/router-settings/index.html
View file

0
litellm/proxy/_experimental/out/settings/ui-theme.html → litellm/proxy/_experimental/out/settings/ui-theme/index.html
View file

0
litellm/proxy/_experimental/out/teams.html → litellm/proxy/_experimental/out/teams/index.html
View file

0
litellm/proxy/_experimental/out/test-key.html → litellm/proxy/_experimental/out/test-key/index.html
View file

0
litellm/proxy/_experimental/out/tools/mcp-servers.html → litellm/proxy/_experimental/out/tools/mcp-servers/index.html
View file

0
litellm/proxy/_experimental/out/tools/vector-stores.html → litellm/proxy/_experimental/out/tools/vector-stores/index.html
View file

0
litellm/proxy/_experimental/out/usage.html → litellm/proxy/_experimental/out/usage/index.html
View file

0
litellm/proxy/_experimental/out/users.html → litellm/proxy/_experimental/out/users/index.html
View file

0
litellm/proxy/_experimental/out/virtual-keys.html → litellm/proxy/_experimental/out/virtual-keys/index.html
View file

9
litellm/proxy/management_endpoints/customer_endpoints.py
View file

@ -10,12 +10,14 @@ All /customer management endpoints
"""
#### END-USER/CUSTOMER MANAGEMENT ####
from datetime import datetime, timedelta
from typing import List, Optional
import fastapi
from fastapi import APIRouter, Depends, HTTPException, Request
import litellm
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -161,7 +163,12 @@ def new_budget_request(data: NewCustomerRequest) -> Optional[BudgetNewRequest]:
budget_kv_pairs[field_name] = value
if budget_kv_pairs:
return BudgetNewRequest(**budget_kv_pairs)
budget_request = BudgetNewRequest(**budget_kv_pairs)
if budget_request.budget_reset_at is None and budget_request.budget_duration is not None:
budget_request.budget_reset_at = datetime.utcnow() + timedelta(
seconds=duration_in_seconds(duration=budget_request.budget_duration)
)
return budget_request
return None

1
litellm/proxy/management_endpoints/tool_management_endpoints.py
View file

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

1
litellm/types/llms/vertex_ai.py
View file

@ -214,6 +214,7 @@ class GenerationConfig(TypedDict, total=False):
responseModalities: List[GeminiResponseModalities]
imageConfig: GeminiImageConfig
thinkingConfig: GeminiThinkingConfig
mediaResolution: str
speechConfig: SpeechConfig

8
litellm/types/utils.py
View file

@ -1234,6 +1234,13 @@ class Delta(SafeAttributeModel, OpenAIObject):
annotations: Optional[List[ChatCompletionAnnotation]] = None,
**params,
):
# Map 'reasoning' to 'reasoning_content' for providers that return
# delta.reasoning (e.g., Cerebras, Groq gpt-oss models).
# Must be done before super().__init__ to prevent 'reasoning' from
# leaking as an extra attribute on the parent model.
if reasoning_content is None and "reasoning" in params:
reasoning_content = params.pop("reasoning", None)
super(Delta, self).__init__(**params)
add_provider_specific_fields(self, params.get("provider_specific_fields", {}))
self.content = content
@ -3110,6 +3117,7 @@ class LlmProviders(str, Enum):
GEMINI = "gemini"
AI21 = "ai21"
BASETEN = "baseten"
BLACK_FOREST_LABS = "black_forest_labs"
AZURE = "azure"
AZURE_TEXT = "azure_text"
AZURE_AI = "azure_ai"

27
litellm/utils.py
View file

@ -7516,6 +7516,15 @@ def is_cached_message(message: AllMessageValues) -> bool:
if litellm.disable_anthropic_gemini_context_caching_transform is True:
return False
# Check message-level cache_control (set by cache_control_injection_points hook for string content)
message_level_cache_control = message.get("cache_control")
if (
message_level_cache_control is not None
and isinstance(message_level_cache_control, dict)
and message_level_cache_control.get("type") == "ephemeral"
):
return True
if "content" not in message:
return False
@ -8322,6 +8331,12 @@ class ProviderConfigManager:
)
return OVHCloudAudioTranscriptionConfig()
elif litellm.LlmProviders.MISTRAL == provider:
from litellm.llms.mistral.audio_transcription.transformation import (
MistralAudioTranscriptionConfig,
)
return MistralAudioTranscriptionConfig()
return None
@staticmethod
@ -8720,6 +8735,12 @@ class ProviderConfigManager:
)
return get_runwayml_image_generation_config(model)
elif LlmProviders.BLACK_FOREST_LABS == provider:
from litellm.llms.black_forest_labs.image_generation import (
get_black_forest_labs_image_generation_config,
)
return get_black_forest_labs_image_generation_config(model)
elif LlmProviders.VERTEX_AI == provider:
from litellm.llms.vertex_ai.image_generation import (
get_vertex_ai_image_generation_config,
@ -8805,6 +8826,12 @@ class ProviderConfigManager:
)
return RecraftImageEditConfig()
elif LlmProviders.BLACK_FOREST_LABS == provider:
from litellm.llms.black_forest_labs.image_edit.transformation import (
BlackForestLabsImageEditConfig,
)
return BlackForestLabsImageEditConfig()
elif LlmProviders.AZURE_AI == provider:
from litellm.llms.azure_ai.image_edit import get_azure_ai_image_edit_config

74
model_prices_and_context_window.json
View file

@ -8023,6 +8023,80 @@
"supports_response_schema": true,
"supports_tool_choice": true
},
"black_forest_labs/flux-kontext-pro": {
"litellm_provider": "black_forest_labs",
"mode": "image_edit",
"output_cost_per_image": 0.04,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/edits",
"/v1/images/generations"
]
},
"black_forest_labs/flux-kontext-max": {
"litellm_provider": "black_forest_labs",
"mode": "image_edit",
"output_cost_per_image": 0.08,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/edits",
"/v1/images/generations"
]
},
"black_forest_labs/flux-pro-1.0-fill": {
"litellm_provider": "black_forest_labs",
"mode": "image_edit",
"output_cost_per_image": 0.05,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/edits"
]
},
"black_forest_labs/flux-pro-1.0-expand": {
"litellm_provider": "black_forest_labs",
"mode": "image_edit",
"output_cost_per_image": 0.05,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/edits"
]
},
"black_forest_labs/flux-pro-1.1": {
"litellm_provider": "black_forest_labs",
"mode": "image_generation",
"output_cost_per_image": 0.04,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/generations"
]
},
"black_forest_labs/flux-pro-1.1-ultra": {
"litellm_provider": "black_forest_labs",
"mode": "image_generation",
"output_cost_per_image": 0.06,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/generations"
]
},
"black_forest_labs/flux-dev": {
"litellm_provider": "black_forest_labs",
"mode": "image_generation",
"output_cost_per_image": 0.025,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/generations"
]
},
"black_forest_labs/flux-pro": {
"litellm_provider": "black_forest_labs",
"mode": "image_generation",
"output_cost_per_image": 0.05,
"source": "https://bfl.ai/pricing",
"supported_endpoints": [
"/v1/images/generations"
]
},
"cerebras/llama-3.3-70b": {
"input_cost_per_token": 8.5e-07,
"litellm_provider": "cerebras",

98
tests/test_litellm/integrations/test_anthropic_cache_control_hook.py
View file

@ -905,6 +905,104 @@ async def test_anthropic_cache_control_hook_document_analysis_multiple_pages():
assert cache_control_count == 1, f"Expected exactly 1 cache control point (last item only), found {cache_control_count}. Before fix, this would be 6 (one for each content item)."
def test_gemini_cache_control_injection_points_detected():
"""
Test that cache_control_injection_points work for Gemini models.
Verifies the full flow:
1. The hook injects cache_control markers on string-content messages
2. is_cached_message() detects the injected markers (message-level cache_control)
3. separate_cached_messages() correctly separates the messages
Fixes GitHub issue #18519.
"""
from litellm.llms.vertex_ai.context_caching.transformation import (
separate_cached_messages,
)
from litellm.utils import is_cached_message
hook = AnthropicCacheControlHook()
# Simulate messages as they would appear for a Gemini call with string content
messages: List[AllMessageValues] = [
{
"role": "system",
"content": "You are a helpful assistant that analyzes legal documents.",
},
{
"role": "user",
"content": "What are the key terms?",
},
]
# Simulate what the hook does: inject cache_control on the system message
injection_points = [{"location": "message", "role": "system"}]
# Manually apply the hook's logic for the system message (string content case)
# The hook sets message["cache_control"] = {"type": "ephemeral"} for string content
hook._safe_insert_cache_control_in_message(
message=messages[0],
control={"type": "ephemeral"},
)
# Verify the hook injected message-level cache_control (string content path)
assert messages[0].get("cache_control") == {"type": "ephemeral"}
# Verify is_cached_message detects message-level cache_control
assert is_cached_message(messages[0]) is True
assert is_cached_message(messages[1]) is False
# Verify separate_cached_messages correctly separates them
cached, non_cached = separate_cached_messages(messages)
assert len(cached) == 1
assert cached[0]["role"] == "system"
assert len(non_cached) == 1
assert non_cached[0]["role"] == "user"
def test_gemini_cache_control_injection_list_content_detected():
"""
Test that cache_control_injection_points work for Gemini models
when the message content is a list (not string).
"""
from litellm.llms.vertex_ai.context_caching.transformation import (
separate_cached_messages,
)
from litellm.utils import is_cached_message
hook = AnthropicCacheControlHook()
messages: List[AllMessageValues] = [
{
"role": "system",
"content": [
{"type": "text", "text": "You are a helpful assistant."},
{"type": "text", "text": "Analyze legal documents carefully."},
],
},
{
"role": "user",
"content": "What are the key terms?",
},
]
# Apply the hook's logic for list content - sets cache_control on last item
hook._safe_insert_cache_control_in_message(
message=messages[0],
control={"type": "ephemeral"},
)
# Verify cache_control was set on the last content item
assert messages[0]["content"][-1]["cache_control"] == {"type": "ephemeral"}
# Verify is_cached_message detects content-item-level cache_control
assert is_cached_message(messages[0]) is True
assert is_cached_message(messages[1]) is False
# Verify separate_cached_messages correctly separates them
cached, non_cached = separate_cached_messages(messages)
assert len(cached) == 1
assert len(non_cached) == 1
@pytest.mark.asyncio
async def test_anthropic_cache_control_hook_string_negative_index():
"""

0
tests/test_litellm/llms/black_forest_labs/__init__.py Normal file
View file

0
tests/test_litellm/llms/black_forest_labs/image_edit/__init__.py Normal file
View file

304
tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py Normal file
View file

@ -0,0 +1,304 @@
"""
Unit tests for Black Forest Labs image edit transformation functionality.
Note: Polling tests are now in test_bfl_image_edit_handler.py
since polling logic was moved to the handler.
"""
import base64
import json
import os
import sys
import time
from io import BytesIO
from typing import Dict, List
from unittest.mock import MagicMock, patch
import httpx
import pytest
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
from litellm.llms.black_forest_labs.image_edit.transformation import (
BlackForestLabsImageEditConfig,
)
from litellm.llms.black_forest_labs.common_utils import BlackForestLabsError
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ImageObject, ImageResponse
class TestBlackForestLabsImageEditTransformation:
"""
Unit tests for Black Forest Labs image edit transformation functionality.
"""
def setup_method(self):
"""Set up test fixtures before each test method."""
self.config = BlackForestLabsImageEditConfig()
self.model = "flux-kontext-pro"
self.logging_obj = MagicMock()
self.prompt = "Add a red hat to the person in the image"
def test_get_supported_openai_params(self):
"""Test that supported OpenAI params are returned correctly."""
params = self.config.get_supported_openai_params(self.model)
# BFL image edit supports BFL-specific params passed through directly
assert isinstance(params, list)
assert len(params) > 0
assert "seed" in params
assert "output_format" in params
assert "safety_tolerance" in params
def test_map_openai_params_basic(self):
"""Test mapping of OpenAI params to BFL params."""
optional_params = ImageEditOptionalRequestParams()
result = self.config.map_openai_params(
image_edit_optional_params=optional_params,
model=self.model,
drop_params=False,
)
# Should have default output_format
assert result.get("output_format") == "png"
def test_map_openai_params_with_bfl_specific(self):
"""Test that BFL-specific params are passed through."""
# BFL-specific params are passed as dict keys
optional_params: ImageEditOptionalRequestParams = {
"seed": 42,
"safety_tolerance": 2,
"aspect_ratio": "16:9",
}
result = self.config.map_openai_params(
image_edit_optional_params=optional_params,
model=self.model,
drop_params=False,
)
assert result.get("seed") == 42
assert result.get("safety_tolerance") == 2
assert result.get("aspect_ratio") == "16:9"
assert result.get("output_format") == "png"
def test_validate_environment_with_api_key(self):
"""Test environment validation with provided API key."""
headers = {}
result = self.config.validate_environment(
headers=headers,
model=self.model,
api_key="test-api-key",
)
assert result["x-key"] == "test-api-key"
assert result["Content-Type"] == "application/json"
assert result["Accept"] == "application/json"
def test_validate_environment_missing_api_key(self):
"""Test that missing API key raises error."""
headers = {}
with patch("litellm.llms.black_forest_labs.image_edit.transformation.get_secret_str") as mock_get_secret:
mock_get_secret.return_value = None
with pytest.raises(BlackForestLabsError) as exc_info:
self.config.validate_environment(
headers=headers,
model=self.model,
api_key=None,
)
assert exc_info.value.status_code == 401
assert "BFL_API_KEY is not set" in exc_info.value.message
def test_get_model_endpoint_kontext_pro(self):
"""Test endpoint resolution for flux-kontext-pro."""
endpoint = self.config._get_model_endpoint("flux-kontext-pro")
assert endpoint == "/v1/flux-kontext-pro"
def test_get_model_endpoint_kontext_max(self):
"""Test endpoint resolution for flux-kontext-max."""
endpoint = self.config._get_model_endpoint("flux-kontext-max")
assert endpoint == "/v1/flux-kontext-max"
def test_get_model_endpoint_with_provider_prefix(self):
"""Test endpoint resolution with provider prefix."""
endpoint = self.config._get_model_endpoint("black_forest_labs/flux-kontext-pro")
assert endpoint == "/v1/flux-kontext-pro"
def test_get_model_endpoint_fill(self):
"""Test endpoint resolution for flux-pro-1.0-fill."""
endpoint = self.config._get_model_endpoint("flux-pro-1.0-fill")
assert endpoint == "/v1/flux-pro-1.0-fill"
def test_get_complete_url(self):
"""Test complete URL generation."""
url = self.config.get_complete_url(
model="flux-kontext-pro",
api_base=None,
litellm_params={},
)
assert url == "https://api.bfl.ai/v1/flux-kontext-pro"
def test_get_complete_url_custom_base(self):
"""Test complete URL generation with custom base."""
url = self.config.get_complete_url(
model="flux-kontext-pro",
api_base="https://custom.api.com/",
litellm_params={},
)
assert url == "https://custom.api.com/v1/flux-kontext-pro"
def test_transform_image_edit_request(self):
"""Test request transformation to BFL format."""
image_data = b"fake_image_data"
image = BytesIO(image_data)
image_edit_optional_params = {
"seed": 123,
"output_format": "jpeg",
}
litellm_params = GenericLiteLLMParams()
headers = {}
data, files = self.config.transform_image_edit_request(
model=self.model,
prompt=self.prompt,
image=image,
image_edit_optional_request_params=image_edit_optional_params,
litellm_params=litellm_params,
headers=headers,
)
# Check that data contains the expected parameters
assert data["prompt"] == self.prompt
assert "input_image" in data
# Verify base64 encoding
decoded = base64.b64decode(data["input_image"])
assert decoded == image_data
assert data["seed"] == 123
assert data["output_format"] == "jpeg"
# BFL uses JSON, not multipart - files should be empty
assert files == []
def test_transform_image_edit_request_with_mask(self):
"""Test request transformation with mask for inpainting."""
image_data = b"fake_image_data"
mask_data = b"fake_mask_data"
image = BytesIO(image_data)
image_edit_optional_params = {
"mask": BytesIO(mask_data),
"output_format": "png",
}
litellm_params = GenericLiteLLMParams()
headers = {}
data, files = self.config.transform_image_edit_request(
model="flux-pro-1.0-fill",
prompt=self.prompt,
image=image,
image_edit_optional_request_params=image_edit_optional_params,
litellm_params=litellm_params,
headers=headers,
)
# Check mask is base64 encoded
assert "mask" in data
decoded_mask = base64.b64decode(data["mask"])
assert decoded_mask == mask_data
def test_read_image_bytes_from_bytes(self):
"""Test reading image bytes from bytes input."""
image_data = b"test_image_bytes"
result = self.config._read_image_bytes(image_data)
assert result == image_data
def test_read_image_bytes_from_file_like(self):
"""Test reading image bytes from file-like object."""
image_data = b"test_image_bytes"
image = BytesIO(image_data)
result = self.config._read_image_bytes(image)
assert result == image_data
def test_read_image_bytes_from_list(self):
"""Test reading image bytes from list (takes first)."""
image_data = b"test_image_bytes"
images = [BytesIO(image_data), BytesIO(b"other")]
result = self.config._read_image_bytes(images)
assert result == image_data
def test_transform_image_edit_response_success(self):
"""Test response transformation with final polled response."""
# The response is now the FINAL polled response from handler
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = {
"status": "Ready",
"result": {"sample": "https://example.com/edited_image.png"},
}
mock_response.status_code = 200
result = self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
assert len(result.data) == 1
assert result.data[0].url == "https://example.com/edited_image.png"
def test_transform_image_edit_response_no_image_url(self):
"""Test response transformation when no image URL is present."""
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = {
"status": "Ready",
"result": {},
}
mock_response.status_code = 200
with pytest.raises(BlackForestLabsError, match="No image URL"):
self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
def test_transform_image_edit_response_json_parse_error(self):
"""Test response transformation with JSON parse error."""
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.side_effect = json.JSONDecodeError("error", "doc", 0)
mock_response.status_code = 200
with pytest.raises(BlackForestLabsError, match="Error parsing"):
self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
def test_get_error_class(self):
"""Test that get_error_class returns BlackForestLabsError."""
error = self.config.get_error_class(
error_message="Test error",
status_code=400,
headers={},
)
assert isinstance(error, BlackForestLabsError)
assert error.status_code == 400
assert "Test error" in str(error.message)
def test_use_multipart_form_data_returns_false(self):
"""Test that use_multipart_form_data returns False for BFL."""
assert self.config.use_multipart_form_data() is False

0
tests/test_litellm/llms/black_forest_labs/image_generation/__init__.py Normal file
View file

350
tests/test_litellm/llms/black_forest_labs/image_generation/test_bfl_image_generation_transformation.py Normal file
View file

@ -0,0 +1,350 @@
"""
Unit tests for Black Forest Labs image generation transformation functionality.
Note: Polling tests are now in test_bfl_image_generation_handler.py
since polling logic was moved to the handler.
"""
import json
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
from litellm.llms.black_forest_labs.image_generation.transformation import (
BlackForestLabsImageGenerationConfig,
get_black_forest_labs_image_generation_config,
)
from litellm.llms.black_forest_labs.common_utils import BlackForestLabsError
from litellm.types.utils import ImageObject, ImageResponse
class TestBlackForestLabsImageGenerationTransformation:
"""
Unit tests for Black Forest Labs image generation transformation functionality.
"""
def setup_method(self):
"""Set up test fixtures before each test method."""
self.config = BlackForestLabsImageGenerationConfig()
self.model = "flux-pro-1.1"
self.logging_obj = MagicMock()
self.prompt = "A beautiful sunset over the ocean"
def test_get_supported_openai_params(self):
"""Test that supported OpenAI params are returned correctly."""
params = self.config.get_supported_openai_params(self.model)
assert "n" in params
assert "size" in params
assert "quality" in params
def test_map_openai_params_basic(self):
"""Test mapping of OpenAI params to BFL params."""
non_default_params = {}
optional_params = {}
result = self.config.map_openai_params(
non_default_params, optional_params, self.model, drop_params=False
)
# Empty input should return empty output
assert result == {}
def test_map_openai_params_size_mapping(self):
"""Test that OpenAI size is mapped to BFL width/height."""
non_default_params = {"size": "1024x1024"}
optional_params = {}
result = self.config.map_openai_params(
non_default_params, optional_params, self.model, drop_params=False
)
assert result["width"] == 1024
assert result["height"] == 1024
def test_map_openai_params_size_custom(self):
"""Test custom size parsing."""
non_default_params = {"size": "800x600"}
optional_params = {}
result = self.config.map_openai_params(
non_default_params, optional_params, self.model, drop_params=False
)
assert result["width"] == 800
assert result["height"] == 600
def test_map_openai_params_n_for_ultra(self):
"""Test that n is mapped to num_images for ultra model."""
non_default_params = {"n": 4}
optional_params = {}
result = self.config.map_openai_params(
non_default_params, optional_params, "flux-pro-1.1-ultra", drop_params=False
)
assert result["num_images"] == 4
def test_map_openai_params_quality_hd_for_ultra(self):
"""Test that 'hd' quality maps to raw=True for ultra model."""
non_default_params = {"quality": "hd"}
optional_params = {}
result = self.config.map_openai_params(
non_default_params, optional_params, "flux-pro-1.1-ultra", drop_params=False
)
assert result["raw"] is True
def test_map_openai_params_unsupported_raises(self):
"""Test that unsupported params raise ValueError when drop_params=False."""
non_default_params = {"unsupported_param": "value"}
optional_params = {}
with pytest.raises(ValueError, match="not supported"):
self.config.map_openai_params(
non_default_params, optional_params, self.model, drop_params=False
)
def test_map_openai_params_unsupported_dropped(self):
"""Test that unsupported params are dropped when drop_params=True."""
non_default_params = {"unsupported_param": "value"}
optional_params = {}
result = self.config.map_openai_params(
non_default_params, optional_params, self.model, drop_params=True
)
assert "unsupported_param" not in result
def test_validate_environment_with_api_key(self):
"""Test that validate_environment sets headers correctly."""
headers = {}
result = self.config.validate_environment(
headers=headers,
model=self.model,
messages=[],
optional_params={},
litellm_params={},
api_key="test_api_key",
)
assert result["x-key"] == "test_api_key"
assert result["Content-Type"] == "application/json"
def test_validate_environment_missing_api_key(self):
"""Test that validate_environment raises error when API key is missing."""
headers = {}
with patch(
"litellm.llms.black_forest_labs.image_generation.transformation.get_secret_str",
return_value=None,
):
with pytest.raises(BlackForestLabsError, match="BFL_API_KEY"):
self.config.validate_environment(
headers=headers,
model=self.model,
messages=[],
optional_params={},
litellm_params={},
api_key=None,
)
def test_get_model_endpoint_flux_pro_1_1(self):
"""Test endpoint for flux-pro-1.1 model."""
endpoint = self.config._get_model_endpoint("flux-pro-1.1")
assert endpoint == "/v1/flux-pro-1.1"
def test_get_model_endpoint_flux_pro_1_1_ultra(self):
"""Test endpoint for flux-pro-1.1-ultra model."""
endpoint = self.config._get_model_endpoint("flux-pro-1.1-ultra")
assert endpoint == "/v1/flux-pro-1.1-ultra"
def test_get_model_endpoint_flux_dev(self):
"""Test endpoint for flux-dev model."""
endpoint = self.config._get_model_endpoint("flux-dev")
assert endpoint == "/v1/flux-dev"
def test_get_model_endpoint_flux_pro(self):
"""Test endpoint for flux-pro model."""
endpoint = self.config._get_model_endpoint("flux-pro")
assert endpoint == "/v1/flux-pro"
def test_get_model_endpoint_flux_kontext_pro(self):
"""Test endpoint for flux-kontext-pro model (supports both generation and editing)."""
endpoint = self.config._get_model_endpoint("flux-kontext-pro")
assert endpoint == "/v1/flux-kontext-pro"
def test_get_model_endpoint_flux_kontext_max(self):
"""Test endpoint for flux-kontext-max model (supports both generation and editing)."""
endpoint = self.config._get_model_endpoint("flux-kontext-max")
assert endpoint == "/v1/flux-kontext-max"
def test_get_model_endpoint_unknown_raises(self):
"""Test that unknown models raise ValueError."""
with pytest.raises(ValueError, match="Unknown BFL image generation model"):
self.config._get_model_endpoint("unknown-model")
def test_get_model_endpoint_with_provider_prefix(self):
"""Test that provider prefix is stripped from model name."""
endpoint = self.config._get_model_endpoint("black_forest_labs/flux-pro-1.1")
assert endpoint == "/v1/flux-pro-1.1"
def test_get_complete_url(self):
"""Test URL construction with default base."""
url = self.config.get_complete_url(
api_base=None,
api_key=None,
model="flux-pro-1.1",
optional_params={},
litellm_params={},
)
assert "https://api.bfl.ai/v1/flux-pro-1.1" == url
def test_get_complete_url_custom_base(self):
"""Test URL construction with custom base."""
url = self.config.get_complete_url(
api_base="https://custom.api.com",
api_key=None,
model="flux-pro-1.1",
optional_params={},
litellm_params={},
)
assert "https://custom.api.com/v1/flux-pro-1.1" == url
def test_transform_image_generation_request(self):
"""Test request body transformation."""
request = self.config.transform_image_generation_request(
model=self.model,
prompt=self.prompt,
optional_params={},
litellm_params={},
headers={},
)
assert request["prompt"] == self.prompt
assert request["output_format"] == "png"
def test_transform_image_generation_request_custom_format(self):
"""Test request body with custom output format."""
request = self.config.transform_image_generation_request(
model=self.model,
prompt=self.prompt,
optional_params={"output_format": "jpeg"},
litellm_params={},
headers={},
)
assert request["output_format"] == "jpeg"
def test_transform_image_generation_request_ultra_params(self):
"""Test request body with ultra-specific params."""
request = self.config.transform_image_generation_request(
model="flux-pro-1.1-ultra",
prompt=self.prompt,
optional_params={
"raw": True,
"num_images": 2,
"aspect_ratio": "16:9",
},
litellm_params={},
headers={},
)
assert request["raw"] is True
assert request["num_images"] == 2
assert request["aspect_ratio"] == "16:9"
def test_transform_image_generation_response_success(self):
"""Test response transformation with final polled response."""
# The response is now the FINAL polled response from handler
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = {
"status": "Ready",
"result": {"sample": "https://example.com/image.png"},
}
mock_response.status_code = 200
model_response = ImageResponse(created=0, data=[])
result = self.config.transform_image_generation_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=self.logging_obj,
)
assert len(result.data) == 1
assert result.data[0].url == "https://example.com/image.png"
def test_transform_image_generation_response_multiple_images(self):
"""Test response transformation with multiple images."""
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = {
"status": "Ready",
"result": [
"https://example.com/image1.png",
"https://example.com/image2.png",
],
}
mock_response.status_code = 200
model_response = ImageResponse(created=0, data=[])
result = self.config.transform_image_generation_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=self.logging_obj,
)
assert len(result.data) == 2
assert result.data[0].url == "https://example.com/image1.png"
assert result.data[1].url == "https://example.com/image2.png"
def test_transform_image_generation_response_no_image(self):
"""Test response transformation when no image URL is present."""
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = {
"status": "Ready",
"result": {},
}
mock_response.status_code = 200
model_response = ImageResponse(created=0, data=[])
with pytest.raises(BlackForestLabsError, match="No image URL"):
self.config.transform_image_generation_response(
model=self.model,
raw_response=mock_response,
model_response=model_response,
logging_obj=self.logging_obj,
)
def test_get_error_class(self):
"""Test that get_error_class returns BlackForestLabsError."""
error = self.config.get_error_class(
error_message="Test error",
status_code=400,
headers={},
)
assert isinstance(error, BlackForestLabsError)
assert error.status_code == 400
assert "Test error" in str(error.message)
def test_get_black_forest_labs_image_generation_config(self):
"""Test the factory function."""
config = get_black_forest_labs_image_generation_config("flux-pro-1.1")
assert isinstance(config, BlackForestLabsImageGenerationConfig)

170
tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py Normal file
View file

@ -0,0 +1,170 @@
import os
from typing import Dict
from unittest.mock import MagicMock
import httpx
import litellm
import pytest
from litellm.llms.base_llm.audio_transcription.transformation import (
BaseAudioTranscriptionConfig,
)
from litellm.llms.mistral.audio_transcription.transformation import (
MistralAudioTranscriptionConfig,
)
from litellm.types.utils import TranscriptionResponse
from litellm.utils import ProviderConfigManager
from tests.llm_translation.base_audio_transcription_unit_tests import (
BaseLLMAudioTranscriptionTest,
)
@pytest.mark.skipif(
not os.getenv("MISTRAL_API_KEY"),
reason="MISTRAL_API_KEY not set, skipping Mistral audio transcription tests",
)
class TestMistralAudioTranscription(BaseLLMAudioTranscriptionTest):
def get_base_audio_transcription_call_args(self) -> Dict:
return {
"model": "mistral/voxtral-mini-latest",
}
def get_custom_llm_provider(self) -> litellm.LlmProviders:
return litellm.LlmProviders.MISTRAL
def test_audio_transcription_async(self): # type: ignore[override]
pytest.skip(
"Async audio transcription test for Mistral is skipped in this suite; "
"async test plugins (e.g. pytest-asyncio/anyio) are not configured here."
)
def test_mistral_audio_transcription_config_installed():
"""Ensure Mistral audio transcription config is registered with ProviderConfigManager."""
config = ProviderConfigManager.get_provider_audio_transcription_config(
model="mistral/voxtral-mini-latest",
provider=litellm.LlmProviders.MISTRAL,
)
assert config is not None
assert isinstance(config, BaseAudioTranscriptionConfig)
assert isinstance(config, MistralAudioTranscriptionConfig)
def test_mistral_audio_transcription_get_complete_url():
config = MistralAudioTranscriptionConfig()
url = config.get_complete_url(
api_base=None,
api_key="fake-key",
model="voxtral-mini-latest",
optional_params={},
litellm_params={},
)
assert url == "https://api.mistral.ai/v1/audio/transcriptions"
def test_mistral_audio_transcription_get_complete_url_custom_base():
config = MistralAudioTranscriptionConfig()
url = config.get_complete_url(
api_base="https://custom.api.example.com/v1/",
api_key="fake-key",
model="voxtral-mini-latest",
optional_params={},
litellm_params={},
)
assert url == "https://custom.api.example.com/v1/audio/transcriptions"
def test_mistral_audio_transcription_validate_environment():
config = MistralAudioTranscriptionConfig()
headers = config.validate_environment(
headers={},
model="voxtral-mini-latest",
messages=[],
optional_params={},
litellm_params={},
api_key="test-key-123",
)
assert headers["Authorization"] == "Bearer test-key-123"
assert headers["accept"] == "application/json"
def test_mistral_audio_transcription_supported_params():
config = MistralAudioTranscriptionConfig()
params = config.get_supported_openai_params("voxtral-mini-latest")
assert "language" in params
assert "temperature" in params
assert "response_format" in params
assert "timestamp_granularities" in params
def test_mistral_audio_transcription_request_transform():
config = MistralAudioTranscriptionConfig()
wav_path = os.path.join(
os.path.dirname(__file__), "../../../../..", "tests", "llm_translation", "gettysburg.wav"
)
audio_file = open(wav_path, "rb")
result = config.transform_audio_transcription_request(
model="voxtral-mini-latest",
audio_file=audio_file,
optional_params={"language": "en", "temperature": 0.0},
litellm_params={},
)
audio_file.close()
assert isinstance(result.data, dict)
assert result.data["model"] == "voxtral-mini-latest"
assert result.data["language"] == "en"
assert result.data["temperature"] == 0.0
assert result.files is not None
assert "file" in result.files
def test_mistral_audio_transcription_request_with_diarize():
"""Test that Mistral-specific params like diarize are passed through."""
config = MistralAudioTranscriptionConfig()
wav_path = os.path.join(
os.path.dirname(__file__), "../../../../..", "tests", "llm_translation", "gettysburg.wav"
)
audio_file = open(wav_path, "rb")
result = config.transform_audio_transcription_request(
model="voxtral-mini-latest",
audio_file=audio_file,
optional_params={"diarize": True},
litellm_params={},
)
audio_file.close()
assert isinstance(result.data, dict)
assert result.data["diarize"] == "true"
def test_mistral_audio_transcription_response_transform():
config = MistralAudioTranscriptionConfig()
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = {
"text": "Four score and seven years ago..."
}
response = config.transform_audio_transcription_response(mock_response)
assert isinstance(response, TranscriptionResponse)
assert response.text == "Four score and seven years ago..."
def test_mistral_audio_transcription_response_transform_empty():
config = MistralAudioTranscriptionConfig()
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = {}
response = config.transform_audio_transcription_response(mock_response)
assert isinstance(response, TranscriptionResponse)
assert response.text == ""

202
tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py Normal file
View file

@ -0,0 +1,202 @@
import os
import sys
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
from litellm.llms.openai.responses.count_tokens.transformation import (
OpenAICountTokensConfig,
)
def test_transform_basic_request():
"""Test basic request with model and input."""
config = OpenAICountTokensConfig()
result = config.transform_request_to_count_tokens(
model="gpt-4o",
input="Hello, how are you?",
)
assert result == {
"model": "gpt-4o",
"input": "Hello, how are you?",
}
def test_transform_with_list_input():
"""Test request with list input format."""
config = OpenAICountTokensConfig()
input_items = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there!"},
]
result = config.transform_request_to_count_tokens(
model="gpt-4o",
input=input_items,
)
assert result["model"] == "gpt-4o"
assert result["input"] == input_items
def test_transform_includes_instructions():
"""Test that instructions are included when provided."""
config = OpenAICountTokensConfig()
result = config.transform_request_to_count_tokens(
model="gpt-4o",
input="Hello",
instructions="You are a helpful assistant.",
)
assert result["instructions"] == "You are a helpful assistant."
assert result["model"] == "gpt-4o"
assert result["input"] == "Hello"
def test_transform_includes_tools():
"""Test that tools are included when provided."""
config = OpenAICountTokensConfig()
tools = [
{
"type": "function",
"name": "get_weather",
"description": "Get weather",
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}},
}
]
result = config.transform_request_to_count_tokens(
model="gpt-4o",
input="What's the weather?",
tools=tools,
)
assert result["tools"] == tools
def test_transform_no_instructions_no_tools():
"""Test that None values are not included."""
config = OpenAICountTokensConfig()
result = config.transform_request_to_count_tokens(
model="gpt-4o",
input="Hello",
instructions=None,
tools=None,
)
assert "instructions" not in result
assert "tools" not in result
def test_messages_to_responses_input_basic():
"""Test converting basic chat messages to Responses API input format."""
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there!"},
{"role": "user", "content": "How are you?"},
]
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
assert len(input_items) == 3
assert input_items[0] == {"role": "user", "content": "Hello"}
assert input_items[1] == {"role": "assistant", "content": "Hi there!"}
assert input_items[2] == {"role": "user", "content": "How are you?"}
assert instructions is None
def test_messages_to_responses_input_with_system():
"""Test that system messages are extracted as instructions."""
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "Hello"},
]
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
assert len(input_items) == 1
assert input_items[0] == {"role": "user", "content": "Hello"}
assert instructions == "You are helpful."
def test_messages_to_responses_input_with_developer():
"""Test that developer messages are extracted as instructions."""
messages = [
{"role": "developer", "content": "Be concise."},
{"role": "user", "content": "Hello"},
]
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
assert len(input_items) == 1
assert instructions == "Be concise."
def test_messages_to_responses_input_with_tool():
"""Test that tool messages are converted to function_call_output."""
messages = [
{"role": "user", "content": "What's the weather?"},
{"role": "tool", "content": "72°F", "tool_call_id": "call_123"},
]
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
assert len(input_items) == 2
assert input_items[1] == {
"type": "function_call_output",
"call_id": "call_123",
"output": "72°F",
}
def test_validate_request_valid():
"""Test that valid requests pass validation."""
config = OpenAICountTokensConfig()
config.validate_request(model="gpt-4o", input="Hello")
def test_validate_request_missing_model():
"""Test that missing model raises ValueError."""
config = OpenAICountTokensConfig()
try:
config.validate_request(model="", input="Hello")
assert False, "Should have raised ValueError"
except ValueError as e:
assert "model" in str(e)
def test_validate_request_missing_input():
"""Test that missing input raises ValueError."""
config = OpenAICountTokensConfig()
try:
config.validate_request(model="gpt-4o", input="")
assert False, "Should have raised ValueError"
except ValueError as e:
assert "input" in str(e)
def test_get_endpoint_default():
"""Test default endpoint URL."""
config = OpenAICountTokensConfig()
assert config.get_openai_count_tokens_endpoint() == "https://api.openai.com/v1/responses/input_tokens"
def test_get_endpoint_custom_base():
"""Test custom API base URL."""
config = OpenAICountTokensConfig()
assert config.get_openai_count_tokens_endpoint("https://custom.api.com/v1") == "https://custom.api.com/v1/responses/input_tokens"
def test_get_required_headers():
"""Test required headers include Authorization."""
config = OpenAICountTokensConfig()
headers = config.get_required_headers("sk-test-key")
assert headers["Authorization"] == "Bearer sk-test-key"
assert headers["Content-Type"] == "application/json"

119
tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py
View file

@ -29,6 +29,15 @@ class TestContextCachingEndpoints:
self.mock_client = MagicMock(spec=HTTPHandler)
self.mock_async_client = MagicMock(spec=AsyncHTTPHandler)
# Mock is_prompt_caching_valid_prompt to return True by default.
# This avoids token counting in unit tests. The min-token guard is
# tested explicitly in test_check_and_create_cache_skips_when_below_min_tokens.
self._token_check_patcher = patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.is_prompt_caching_valid_prompt",
return_value=True,
)
self._token_check_patcher.start()
# Sample messages for testing
self.sample_messages = [
{
@ -56,6 +65,10 @@ class TestContextCachingEndpoints:
self.sample_optional_params = {"tools": self.sample_tools.copy()}
def teardown_method(self):
"""Teardown for each test method"""
self._token_check_patcher.stop()
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
@ -787,6 +800,112 @@ class TestContextCachingEndpoints:
# But original tools should still be available for comparison
assert original_tools == self.sample_tools
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
@patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
)
def test_check_and_create_cache_skips_when_below_min_tokens(
self, mock_separate, custom_llm_provider
):
"""Test that context caching is skipped when cached content is below 1024 tokens.
Gemini requires a minimum of 1024 tokens for context caching. If the cached
content is too small, the request should proceed without caching instead of
failing with a Gemini API error.
"""
# Stop the default mock so the real token count check runs
self._token_check_patcher.stop()
short_cached_messages = [
{
"role": "system",
"content": "You are a helpful assistant.",
"cache_control": {"type": "ephemeral"},
}
]
non_cached_messages = [
{"role": "user", "content": "Hello"},
]
all_messages = short_cached_messages + non_cached_messages
mock_separate.return_value = (short_cached_messages, non_cached_messages)
optional_params = self.sample_optional_params.copy()
result = self.context_caching.check_and_create_cache(
messages=all_messages,
optional_params=optional_params,
api_key="test_key",
api_base=None,
model="gemini-1.5-pro",
client=self.mock_client,
timeout=30.0,
logging_obj=self.mock_logging,
cached_content=None,
custom_llm_provider=custom_llm_provider,
vertex_project="test_project",
vertex_location="test_location",
vertex_auth_header="test_token",
)
messages, returned_params, returned_cache = result
assert messages == all_messages
assert returned_cache is None
# Restart the patcher so teardown_method can stop it cleanly
self._token_check_patcher.start()
@pytest.mark.parametrize(
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
)
@patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages"
)
@pytest.mark.asyncio
async def test_async_check_and_create_cache_skips_when_below_min_tokens(
self, mock_separate, custom_llm_provider
):
"""Test that async context caching is skipped when cached content is below 1024 tokens."""
# Stop the default mock so the real token count check runs
self._token_check_patcher.stop()
short_cached_messages = [
{
"role": "system",
"content": "You are a helpful assistant.",
"cache_control": {"type": "ephemeral"},
}
]
non_cached_messages = [
{"role": "user", "content": "Hello"},
]
all_messages = short_cached_messages + non_cached_messages
mock_separate.return_value = (short_cached_messages, non_cached_messages)
optional_params = self.sample_optional_params.copy()
result = await self.context_caching.async_check_and_create_cache(
messages=all_messages,
optional_params=optional_params,
api_key="test_key",
api_base=None,
model="gemini-1.5-pro",
client=self.mock_async_client,
timeout=30.0,
logging_obj=self.mock_logging,
cached_content=None,
custom_llm_provider=custom_llm_provider,
vertex_project="test_project",
vertex_location="test_location",
vertex_auth_header="test_token",
)
messages, returned_params, returned_cache = result
assert messages == all_messages
assert returned_cache is None
# Restart the patcher so teardown_method can stop it cleanly
self._token_check_patcher.start()
class TestCheckCachePagination:
"""Test pagination logic in check_cache and async_check_cache methods."""

390
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py
View file

@ -5,6 +5,8 @@ from litellm.llms.vertex_ai.gemini.transformation import (
_gemini_convert_messages_with_history,
_transform_request_body,
check_if_part_exists_in_parts,
_get_highest_media_resolution,
_extract_max_media_resolution_from_messages,
)
from litellm.types.llms.vertex_ai import BlobType
from litellm.types.utils import Message
@ -616,12 +618,306 @@ def test_dummy_signature_with_function_call_mode():
assert gemini_parts[0]["thoughtSignature"] == expected_dummy
# Tests for media_resolution (detail parameter) handling - Issue #17084
class TestMediaResolution:
"""Tests for media_resolution handling in Gemini 2.x models"""
def test_get_highest_media_resolution_high_wins(self):
"""Test that 'high' resolution takes precedence over 'low'"""
assert _get_highest_media_resolution("low", "high") == "high"
assert _get_highest_media_resolution("high", "low") == "high"
assert _get_highest_media_resolution(None, "high") == "high"
assert _get_highest_media_resolution("high", None) == "high"
def test_get_highest_media_resolution_low_over_none(self):
"""Test that 'low' resolution takes precedence over None"""
assert _get_highest_media_resolution(None, "low") == "low"
assert _get_highest_media_resolution("low", None) == "low"
def test_get_highest_media_resolution_same_values(self):
"""Test handling of same resolution values"""
assert _get_highest_media_resolution("high", "high") == "high"
assert _get_highest_media_resolution("low", "low") == "low"
assert _get_highest_media_resolution(None, None) is None
def test_get_highest_media_resolution_medium(self):
"""Test that 'medium' resolution is correctly ranked between 'low' and 'high'"""
assert _get_highest_media_resolution("low", "medium") == "medium"
assert _get_highest_media_resolution("medium", "low") == "medium"
assert _get_highest_media_resolution("medium", "high") == "high"
assert _get_highest_media_resolution("high", "medium") == "high"
assert _get_highest_media_resolution(None, "medium") == "medium"
assert _get_highest_media_resolution("medium", None) == "medium"
def test_get_highest_media_resolution_ultra_high(self):
"""Test that 'ultra_high' resolution takes precedence over all others"""
assert _get_highest_media_resolution("high", "ultra_high") == "ultra_high"
assert _get_highest_media_resolution("ultra_high", "high") == "ultra_high"
assert _get_highest_media_resolution("medium", "ultra_high") == "ultra_high"
assert _get_highest_media_resolution("low", "ultra_high") == "ultra_high"
assert _get_highest_media_resolution(None, "ultra_high") == "ultra_high"
assert _get_highest_media_resolution("ultra_high", None) == "ultra_high"
def test_extract_max_media_resolution_single_image_high(self):
"""Test extraction of media resolution from single image with detail=high"""
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,abc123", "detail": "high"},
},
],
}
]
assert _extract_max_media_resolution_from_messages(messages) == "high"
def test_extract_max_media_resolution_single_image_low(self):
"""Test extraction of media resolution from single image with detail=low"""
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,abc123", "detail": "low"},
},
],
}
]
assert _extract_max_media_resolution_from_messages(messages) == "low"
def test_extract_max_media_resolution_no_detail(self):
"""Test extraction when no detail parameter is provided"""
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,abc123"},
},
],
}
]
assert _extract_max_media_resolution_from_messages(messages) is None
def test_extract_max_media_resolution_multiple_images_mixed(self):
"""Test that highest resolution is returned when multiple images have different details"""
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "Compare these images"},
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,abc123", "detail": "low"},
},
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,def456", "detail": "high"},
},
],
}
]
assert _extract_max_media_resolution_from_messages(messages) == "high"
def test_extract_max_media_resolution_text_only(self):
"""Test extraction from messages with no images"""
messages = [
{"role": "user", "content": "Hello, how are you?"},
{"role": "assistant", "content": "I'm doing well!"},
]
assert _extract_max_media_resolution_from_messages(messages) is None
def test_transform_request_body_gemini_2x_adds_media_resolution(self):
"""Test that media_resolution is added to generationConfig for Gemini 2.x models"""
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,iVBORw0KGgo=", "detail": "high"},
},
],
}
]
result = _transform_request_body(
messages=messages,
model="gemini-2.5-flash",
optional_params={},
custom_llm_provider="gemini",
litellm_params={},
cached_content=None,
)
assert "generationConfig" in result
assert "mediaResolution" in result["generationConfig"]
assert result["generationConfig"]["mediaResolution"] == "MEDIA_RESOLUTION_HIGH"
def test_transform_request_body_gemini_2x_low_resolution(self):
"""Test that low media_resolution is correctly added for Gemini 2.x"""
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,iVBORw0KGgo=", "detail": "low"},
},
],
}
]
result = _transform_request_body(
messages=messages,
model="gemini-2.5-flash",
optional_params={},
custom_llm_provider="gemini",
litellm_params={},
cached_content=None,
)
assert "generationConfig" in result
assert "mediaResolution" in result["generationConfig"]
assert result["generationConfig"]["mediaResolution"] == "MEDIA_RESOLUTION_LOW"
def test_transform_request_body_gemini_3_no_global_media_resolution(self):
"""Test that Gemini 3 models don't add media_resolution to generationConfig (they use per-part)"""
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,iVBORw0KGgo=", "detail": "high"},
},
],
}
]
result = _transform_request_body(
messages=messages,
model="gemini-3-pro-preview",
optional_params={},
custom_llm_provider="gemini",
litellm_params={},
cached_content=None,
)
# Gemini 3 should NOT have mediaResolution in generationConfig
# (it's handled per-part in the content transformation)
if "generationConfig" in result:
assert "mediaResolution" not in result["generationConfig"]
def test_transform_request_body_no_detail_no_media_resolution(self):
"""Test that no mediaResolution is added when detail is not specified"""
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,iVBORw0KGgo="},
},
],
}
]
result = _transform_request_body(
messages=messages,
model="gemini-2.5-flash",
optional_params={},
custom_llm_provider="gemini",
litellm_params={},
cached_content=None,
)
# When no detail is specified, mediaResolution should not be in generationConfig
if "generationConfig" in result:
assert "mediaResolution" not in result["generationConfig"]
def test_extract_max_media_resolution_file_type_with_detail(self):
"""Test that detail is extracted from file content type, not just image_url"""
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is in this file?"},
{
"type": "file",
"file": {"url": "data:image/png;base64,abc123", "detail": "high"},
},
],
}
]
assert _extract_max_media_resolution_from_messages(messages) == "high"
def test_extract_max_media_resolution_mixed_image_and_file(self):
"""Test that highest detail is returned across both image_url and file types"""
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "Compare these"},
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,abc123", "detail": "low"},
},
{
"type": "file",
"file": {"url": "data:image/png;base64,def456", "detail": "high"},
},
],
}
]
assert _extract_max_media_resolution_from_messages(messages) == "high"
def test_transform_request_body_gemini_1x_no_media_resolution(self):
"""Test that Gemini 1.x models don't get mediaResolution in generationConfig"""
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,iVBORw0KGgo=", "detail": "high"},
},
],
}
]
result = _transform_request_body(
messages=messages,
model="gemini-1.5-pro",
optional_params={},
custom_llm_provider="gemini",
litellm_params={},
cached_content=None,
)
# Gemini 1.x should NOT have mediaResolution (not supported)
if "generationConfig" in result:
assert "mediaResolution" not in result["generationConfig"]
def test_convert_tool_response_with_base64_image():
"""Test tool response with base64 data URI image."""
# Create a small test image (1x1 red pixel PNG)
test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
image_data_uri = f"data:image/png;base64,{test_image_base64}"
# Create tool message with image
tool_message = {
"role": "tool",
@ -637,7 +933,7 @@ def test_convert_tool_response_with_base64_image():
}
]
}
# Mock last message with tool calls
last_message_with_tool_calls = {
"tool_calls": [
@ -650,16 +946,16 @@ def test_convert_tool_response_with_base64_image():
}
]
}
# Convert tool response (returns list when image is present)
result = convert_to_gemini_tool_call_result(
tool_message, last_message_with_tool_calls
)
# Verify results - should be a list with 2 parts (function_response + inline_data)
assert isinstance(result, list), f"Expected list when image present, got {type(result)}"
assert len(result) == 2, f"Expected 2 parts, got {len(result)}"
# Find function_response part and inline_data part
function_response_part = None
inline_data_part = None
@ -668,7 +964,7 @@ def test_convert_tool_response_with_base64_image():
function_response_part = part
elif "inline_data" in part:
inline_data_part = part
# Check function_response exists
assert function_response_part is not None, "Missing function_response part"
function_response = function_response_part["function_response"]
@ -677,7 +973,7 @@ def test_convert_tool_response_with_base64_image():
# Verify JSON response is parsed correctly
assert "url" in function_response["response"]
assert function_response["response"]["url"] == "https://example.com"
# Check inline_data exists
assert inline_data_part is not None, "Missing inline_data part"
inline_data: BlobType = inline_data_part["inline_data"]
@ -693,7 +989,7 @@ def test_convert_tool_response_with_url_image():
# Use a publicly accessible test image URL
test_image_url = "https://via.placeholder.com/1x1.png"
tool_message = {
"role": "tool",
"tool_call_id": "call_test456",
@ -708,7 +1004,7 @@ def test_convert_tool_response_with_url_image():
}
]
}
last_message_with_tool_calls = {
"tool_calls": [
{
@ -720,25 +1016,25 @@ def test_convert_tool_response_with_url_image():
}
]
}
try:
result = convert_to_gemini_tool_call_result(
tool_message, last_message_with_tool_calls
)
# Should be a list with 2 parts when image is present
assert isinstance(result, list), f"Expected list when image present, got {type(result)}"
assert len(result) == 2, f"Expected 2 parts, got {len(result)}"
# Find parts
function_response_part = next(p for p in result if "function_response" in p)
inline_data_part = next(p for p in result if "inline_data" in p)
# Check function_response exists
assert function_response_part is not None, "Missing function_response part"
function_response = function_response_part["function_response"]
assert function_response["name"] == "type_text_at"
# Check inline_data exists (URL should be downloaded and converted)
assert inline_data_part is not None, "Missing inline_data part"
inline_data: BlobType = inline_data_part["inline_data"]
@ -761,7 +1057,7 @@ def test_convert_tool_response_text_only():
}
]
}
last_message_with_tool_calls = {
"tool_calls": [
{
@ -773,14 +1069,14 @@ def test_convert_tool_response_text_only():
}
]
}
result = convert_to_gemini_tool_call_result(
tool_message, last_message_with_tool_calls
)
# Should be a single part (no list) when no image
assert not isinstance(result, list), "Should return single part when no image"
# Check function_response exists
assert "function_response" in result
function_response = result["function_response"]
@ -788,7 +1084,7 @@ def test_convert_tool_response_text_only():
# Verify JSON response is parsed correctly
assert "status" in function_response["response"]
assert function_response["response"]["status"] == "completed"
# Check inline_data does NOT exist (no image provided)
assert "inline_data" not in result
@ -796,12 +1092,12 @@ def test_convert_tool_response_text_only():
def test_file_data_field_order():
"""
Test that file_data fields are in the correct order (mime_type before file_uri).
The Gemini API is sensitive to field order in the file_data object.
This test verifies that mime_type comes before file_uri in both:
1. Dictionary key order
2. JSON serialization
Related issue: Gemini API returns 400 INVALID_ARGUMENT when fields are in wrong order.
"""
import json
@ -811,25 +1107,25 @@ def test_file_data_field_order():
# Test with HTTPS URL and explicit format (audio file)
file_url = "https://generativelanguage.googleapis.com/v1beta/files/test123"
format = "audio/mpeg"
result = _process_gemini_media(image_url=file_url, format=format)
# Verify the result has file_data
assert "file_data" in result
file_data = result["file_data"]
# Verify both fields are present
assert "mime_type" in file_data
assert "file_uri" in file_data
assert file_data["mime_type"] == "audio/mpeg"
assert file_data["file_uri"] == file_url
# Verify field order by checking dictionary keys
# In Python 3.7+, dict maintains insertion order
file_data_keys = list(file_data.keys())
assert file_data_keys.index("mime_type") < file_data_keys.index("file_uri"), \
"mime_type must come before file_uri in the file_data dict"
# Also verify by serializing to JSON string
json_str = json.dumps(file_data)
mime_type_pos = json_str.find('"mime_type"')
@ -846,17 +1142,17 @@ def test_file_data_field_order_gcs_urls():
# Test with GCS URL
gcs_url = "gs://bucket/audio.mp3"
result = _process_gemini_media(image_url=gcs_url)
# Verify the result has file_data
assert "file_data" in result
file_data = result["file_data"]
# Verify both fields are present
assert "mime_type" in file_data
assert "file_uri" in file_data
# Verify field order
file_data_keys = list(file_data.keys())
assert file_data_keys.index("mime_type") < file_data_keys.index("file_uri"), \
@ -866,11 +1162,11 @@ def test_file_data_field_order_gcs_urls():
def test_extract_file_data_with_path_object():
"""
Test that filename is correctly extracted from Path objects for MIME type detection.
When uploading files using Path objects (e.g., Path("speech.mp3")), the filename
must be extracted to enable proper MIME type detection. Without this, files get
uploaded with 'application/octet-stream' instead of the correct MIME type.
Related issue: Files uploaded with wrong MIME type cause Gemini API to reject
requests where the specified format doesn't match the uploaded file's MIME type.
"""
@ -886,23 +1182,23 @@ def test_extract_file_data_with_path_object():
with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as tmp:
tmp.write(b"fake mp3 content")
tmp_path = tmp.name
try:
# Test with Path object
path_obj = Path(tmp_path)
extracted = extract_file_data(path_obj)
# Verify filename was extracted
assert extracted["filename"] is not None
assert extracted["filename"].endswith(".mp3")
# Verify MIME type was correctly detected
assert extracted["content_type"] == "audio/mpeg", \
f"Expected 'audio/mpeg' but got '{extracted['content_type']}'"
# Verify content was read
assert extracted["content"] == b"fake mp3 content"
finally:
# Clean up temporary file
os.unlink(tmp_path)
@ -921,22 +1217,22 @@ def test_extract_file_data_with_string_path():
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
tmp.write(b"fake wav content")
tmp_path = tmp.name
try:
# Test with string path
extracted = extract_file_data(tmp_path)
# Verify filename was extracted
assert extracted["filename"] is not None
assert extracted["filename"].endswith(".wav")
# Verify MIME type was correctly detected (can be audio/wav or audio/x-wav depending on system)
assert extracted["content_type"] in ["audio/wav", "audio/x-wav"], \
f"Expected 'audio/wav' or 'audio/x-wav' but got '{extracted['content_type']}'"
# Verify content was read
assert extracted["content"] == b"fake wav content"
finally:
# Clean up temporary file
os.unlink(tmp_path)
@ -952,9 +1248,9 @@ def test_extract_file_data_with_tuple_format():
filename = "test_audio.mp3"
content = b"test audio content"
content_type = "audio/mpeg"
extracted = extract_file_data((filename, content, content_type))
# Verify all fields are correct
assert extracted["filename"] == filename
assert extracted["content"] == content
@ -974,15 +1270,15 @@ def test_extract_file_data_fallback_to_octet_stream():
with tempfile.NamedTemporaryFile(suffix=".xyz123", delete=False) as tmp:
tmp.write(b"unknown content")
tmp_path = tmp.name
try:
# Test with unknown file type
extracted = extract_file_data(tmp_path)
# Verify filename was extracted
assert extracted["filename"] is not None
assert extracted["filename"].endswith(".xyz123")
# Verify MIME type falls back to octet-stream
assert extracted["content_type"] == "application/octet-stream", \
f"Expected 'application/octet-stream' for unknown type, got '{extracted['content_type']}'"

67
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -3724,3 +3724,70 @@ def test_vertex_ai_usage_metadata_video_tokens_with_caching():
assert result.prompt_tokens_details.text_tokens == 9
assert result.prompt_tokens_details.audio_tokens == 200
def test_async_streaming_uses_custom_client():
"""
Test that user-specified async client is correctly passed to make_call
for async streaming calls.
Fixes: https://github.com/BerriAI/litellm/issues/17148
"""
from functools import partial
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
make_call,
)
# Create a mock async client
mock_client = MagicMock(spec=AsyncHTTPHandler)
# Create a partial function like the code does in async_streaming
partial_make_call = partial(
make_call,
gemini_client=mock_client,
api_base="https://example.com",
headers={},
data="{}",
model="gemini-pro",
messages=[],
logging_obj=MagicMock(),
)
# Verify that gemini_client is in the partial's keywords
assert "gemini_client" in partial_make_call.keywords
assert partial_make_call.keywords["gemini_client"] is mock_client
def test_sync_streaming_uses_custom_client():
"""
Test that user-specified sync client is correctly passed to make_sync_call
for sync streaming calls.
This verifies the existing behavior that we want to match for async.
"""
from functools import partial
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
make_sync_call,
)
# Create a mock sync client
mock_client = MagicMock(spec=HTTPHandler)
# Create a partial function like the code does in sync streaming
partial_make_sync_call = partial(
make_sync_call,
gemini_client=mock_client,
api_base="https://example.com",
headers={},
data="{}",
model="gemini-pro",
messages=[],
logging_obj=MagicMock(),
)
# Verify that gemini_client is in the partial's keywords
assert "gemini_client" in partial_make_sync_call.keywords
assert partial_make_sync_call.keywords["gemini_client"] is mock_client

39
tests/test_litellm/proxy/management_endpoints/test_customer_budget.py
View file

@ -6,18 +6,25 @@ Tests customer update functionality related to budget management:
- Creating new budgets for customers with proper field validation
- Budget creation with required metadata fields
- Proper database relationship handling
- Budget initialization on customer creation
"""
from datetime import datetime, timedelta
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy._types import (
LiteLLM_BudgetTable,
LiteLLM_EndUserTable,
NewCustomerRequest,
UpdateCustomerRequest,
)
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.management_endpoints.customer_endpoints import update_end_user
from litellm.proxy.management_endpoints.customer_endpoints import (
new_budget_request,
update_end_user,
)
@pytest.fixture
@ -340,4 +347,32 @@ async def test_update_customer_with_budget_id_and_creation_fields(
# The update data should contain budget_id from the created budget, not the original budget_id
update_data = call_args[1]['data']
assert update_data['budget_id'] == "new-budget-combo" # From created budget
assert update_data['budget_id'] == "new-budget-combo" # From created budget
def test_new_budget_request_sets_budget_reset_at_when_duration_provided():
"""
Test that new_budget_request auto-populates budget_reset_at when
budget_duration is provided but budget_reset_at is not.
Without this fix, budgets created via /customer/new with a budget_duration
but no budget_reset_at would have budget_reset_at=NULL in the DB, causing
the ResetBudgetJob to immediately pick them up and zero out enduser spend.
"""
data = NewCustomerRequest(
user_id="test-user",
max_budget=10.0,
budget_duration="30d",
)
before = datetime.utcnow()
result = new_budget_request(data)
after = datetime.utcnow()
assert result is not None
assert result.budget_reset_at is not None
assert result.budget_duration == "30d"
expected_min = before + timedelta(days=30)
expected_max = after + timedelta(days=30)
assert expected_min <= result.budget_reset_at <= expected_max

160
tests/test_litellm/test_count_tokens_public_api.py Normal file
View file

@ -0,0 +1,160 @@
"""
Tests for litellm.acount_tokens() public API.
"""
import asyncio
import os
import sys
from unittest.mock import AsyncMock, patch
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm.types.utils import TokenCountResponse
def test_acount_tokens_routes_to_openai():
"""Test that acount_tokens routes to OpenAI token counter for openai/ models."""
with patch(
"litellm.llms.openai.responses.count_tokens.token_counter.openai_count_tokens_handler.handle_count_tokens_request",
new_callable=AsyncMock,
return_value={"input_tokens": 15},
):
result = asyncio.run(
litellm.acount_tokens(
model="openai/gpt-4o",
messages=[{"role": "user", "content": "Hello, how are you?"}],
api_key="sk-test-key",
)
)
assert result.total_tokens == 15
assert result.tokenizer_type == "openai_api"
assert result.request_model == "openai/gpt-4o"
def test_acount_tokens_routes_to_anthropic():
"""Test that acount_tokens routes to Anthropic token counter for anthropic/ models."""
with patch(
"litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler.handle_count_tokens_request",
new_callable=AsyncMock,
return_value={"input_tokens": 20},
):
result = asyncio.run(
litellm.acount_tokens(
model="anthropic/claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "Hello Claude!"}],
api_key="sk-ant-test-key",
)
)
assert result.total_tokens == 20
assert result.tokenizer_type == "anthropic_api"
assert result.request_model == "anthropic/claude-3-5-sonnet-20241022"
def test_acount_tokens_fallback_to_local():
"""Test that unsupported providers fall back to local tiktoken counting."""
result = asyncio.run(
litellm.acount_tokens(
model="together_ai/meta-llama/Llama-3-8b-chat-hf",
messages=[{"role": "user", "content": "Hello"}],
)
)
assert result.total_tokens > 0
assert result.tokenizer_type == "local_tokenizer"
def test_acount_tokens_with_tools():
"""Test that tools are passed through to the token counter."""
tools = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather info",
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}},
},
}
]
with patch(
"litellm.llms.openai.responses.count_tokens.token_counter.openai_count_tokens_handler.handle_count_tokens_request",
new_callable=AsyncMock,
return_value={"input_tokens": 30},
) as mock_handler:
result = asyncio.run(
litellm.acount_tokens(
model="openai/gpt-4o",
messages=[{"role": "user", "content": "What's the weather?"}],
tools=tools,
api_key="sk-test-key",
)
)
assert result.total_tokens == 30
mock_handler.assert_called_once()
call_kwargs = mock_handler.call_args
assert call_kwargs.kwargs.get("tools") == tools
def test_acount_tokens_with_system():
"""Test that system messages are passed through."""
with patch(
"litellm.llms.openai.responses.count_tokens.token_counter.openai_count_tokens_handler.handle_count_tokens_request",
new_callable=AsyncMock,
return_value={"input_tokens": 25},
):
result = asyncio.run(
litellm.acount_tokens(
model="openai/gpt-4o",
messages=[{"role": "user", "content": "Hello"}],
system="You are a helpful assistant.",
api_key="sk-test-key",
)
)
assert result.total_tokens == 25
def test_acount_tokens_api_error_falls_back():
"""Test that API errors in token counting return error response."""
from litellm.llms.openai.common_utils import OpenAIError
with patch(
"litellm.llms.openai.responses.count_tokens.token_counter.openai_count_tokens_handler.handle_count_tokens_request",
new_callable=AsyncMock,
side_effect=OpenAIError(status_code=401, message="Invalid API key"),
):
result = asyncio.run(
litellm.acount_tokens(
model="openai/gpt-4o",
messages=[{"role": "user", "content": "Hello"}],
api_key="sk-bad-key",
)
)
# Should fall back to local tokenizer when provider API errors
assert result.error is False
assert result.tokenizer_type == "local_tokenizer"
assert result.total_tokens > 0
def test_acount_tokens_no_api_key_falls_back():
"""Test that missing API key falls back to local counting."""
env_backup = os.environ.pop("OPENAI_API_KEY", None)
try:
result = asyncio.run(
litellm.acount_tokens(
model="openai/gpt-4o",
messages=[{"role": "user", "content": "Hello"}],
)
)
# Should fall back to local tokenizer since no API key
assert result.total_tokens > 0
assert result.tokenizer_type == "local_tokenizer"
finally:
if env_backup:
os.environ["OPENAI_API_KEY"] = env_backup

32
tests/test_litellm/test_utils.py
View file

@ -2944,6 +2944,38 @@ class TestIsCachedMessage:
message = {"role": "user", "content": []}
assert is_cached_message(message) is False
def test_message_level_cache_control_returns_true(self):
"""Message with string content and message-level cache_control should return True.
This is the format injected by the cache_control_injection_points hook
when the message content is a string (common for system messages).
Fixes GitHub issue #18519 - Gemini models ignoring cache_control_injection_points.
"""
message = {
"role": "system",
"content": "You are a helpful assistant.",
"cache_control": {"type": "ephemeral"},
}
assert is_cached_message(message) is True
def test_message_level_cache_control_wrong_type_returns_false(self):
"""Message-level cache_control with non-ephemeral type should return False."""
message = {
"role": "system",
"content": "You are a helpful assistant.",
"cache_control": {"type": "permanent"},
}
assert is_cached_message(message) is False
def test_message_level_cache_control_non_dict_returns_false(self):
"""Message-level cache_control that's not a dict should return False."""
message = {
"role": "system",
"content": "You are a helpful assistant.",
"cache_control": "ephemeral",
}
assert is_cached_message(message) is False
@pytest.mark.asyncio
class TestProxyLoggingBudgetAlerts:

27
tests/test_litellm/types/test_types_utils.py
View file

@ -223,3 +223,30 @@ def test_chat_completion_token_logprob_invalid_top_logprobs_rejected():
logprob=-0.31725305,
top_logprobs="invalid_string",
)
def test_delta_maps_reasoning_to_reasoning_content():
"""
Test that Delta maps 'reasoning' field to 'reasoning_content'.
Providers like Cerebras and Groq return delta.reasoning for gpt-oss models,
but LiteLLM expects delta.reasoning_content.
"""
from litellm.types.utils import Delta
# When provider sends 'reasoning' (e.g., Cerebras gpt-oss streaming)
delta = Delta(content=None, role="assistant", reasoning="thinking step by step")
assert delta.reasoning_content == "thinking step by step"
assert not hasattr(delta, "reasoning"), "reasoning should not leak as an extra attribute"
# When provider sends 'reasoning_content' directly (e.g., NIM), it still works
delta2 = Delta(content="hello", reasoning_content="direct reasoning")
assert delta2.reasoning_content == "direct reasoning"
# When both are present, reasoning_content takes precedence
delta3 = Delta(reasoning_content="from_rc", reasoning="from_r")
assert delta3.reasoning_content == "from_rc"
# When neither is present, reasoning_content is not set (OpenAI spec)
delta4 = Delta(content="hello")
assert not hasattr(delta4, "reasoning_content")

98
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx
View file

@ -1,9 +1,27 @@
/* @vitest-environment jsdom */
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { act, render } from "@testing-library/react";
import { act, fireEvent, render } from "@testing-library/react";
import { beforeEach, describe, expect, it, vi } from "vitest";
import ModelsAndEndpointsView from "./ModelsAndEndpointsView";
// Mock localStorage
const localStorageMock = (() => {
let store: Record<string, string> = {};
return {
getItem: (key: string) => store[key] || null,
setItem: (key: string, value: string) => {
store[key] = value;
},
removeItem: (key: string) => {
delete store[key];
},
clear: () => {
store = {};
},
};
})();
Object.defineProperty(window, "localStorage", { value: localStorageMock });
// Minimal stubs to avoid Next.js router and network usage during render
vi.mock("@/components/networking", () => ({
credentialListCall: vi.fn().mockResolvedValue({ credentials: [] }),
@ -115,6 +133,84 @@ describe("ModelsAndEndpointsView", () => {
expect(await findByText("Model Management", {}, { timeout: 10000 })).toBeInTheDocument();
}, 15000);
it("should show Missing provider banner by default", async () => {
localStorageMock.clear();
const queryClient = createQueryClient();
const { findByText } = render(
<QueryClientProvider client={queryClient}>
<ModelsAndEndpointsView
token="123"
modelData={{ data: [] }}
keys={[]}
setModelData={() => {}}
premiumUser={false}
teams={[]}
/>
</QueryClientProvider>,
);
expect(await findByText("Missing a provider?", {}, { timeout: 10000 })).toBeInTheDocument();
}, 15000);
it("should hide Missing provider banner when dismiss button is clicked and persist to localStorage", async () => {
localStorageMock.clear();
const queryClient = createQueryClient();
const { findByText, queryByText, container } = render(
<QueryClientProvider client={queryClient}>
<ModelsAndEndpointsView
token="123"
modelData={{ data: [] }}
keys={[]}
setModelData={() => {}}
premiumUser={false}
teams={[]}
/>
</QueryClientProvider>,
);
// Wait for banner to appear
expect(await findByText("Missing a provider?", {}, { timeout: 10000 })).toBeInTheDocument();
// Find and click dismiss button (X button)
const dismissButton = container.querySelector('button[aria-label="Dismiss banner"]');
expect(dismissButton).not.toBeNull();
fireEvent.click(dismissButton!);
// Banner should be hidden
expect(queryByText("Missing a provider?")).not.toBeInTheDocument();
// LocalStorage should be updated
expect(localStorageMock.getItem("hideMissingProviderBanner")).toBe("true");
}, 15000);
it("should show compact Request Provider button when banner is dismissed", async () => {
// Set localStorage to hide banner
localStorageMock.setItem("hideMissingProviderBanner", "true");
const queryClient = createQueryClient();
const { findByText, queryByText } = render(
<QueryClientProvider client={queryClient}>
<ModelsAndEndpointsView
token="123"
modelData={{ data: [] }}
keys={[]}
setModelData={() => {}}
premiumUser={false}
teams={[]}
/>
</QueryClientProvider>,
);
// Wait for component to render
await findByText("Model Management", {}, { timeout: 10000 });
// Banner should not be visible
expect(queryByText("Missing a provider?")).not.toBeInTheDocument();
// Compact Request Provider button should be visible in header
const requestProviderLinks = document.querySelectorAll('a[href="https://models.litellm.ai/?request=true"]');
// There should be a compact button when banner is hidden
expect(requestProviderLinks.length).toBeGreaterThan(0);
}, 15000);
it("should pass model IDs (not model names) to HealthCheckComponent as all_models_on_proxy", async () => {
mockHealthCheckComponent.mockClear();
const modelDataWithIds = {

116
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx
View file

@ -15,7 +15,7 @@ import { transformModelData } from "./utils/modelDataTransformer";
import { all_admin_roles, internalUserRoles, isProxyAdminRole, isUserTeamAdminForAnyTeam } from "@/utils/roles";
import { RefreshIcon } from "@heroicons/react/outline";
import { useQueryClient } from "@tanstack/react-query";
import { Col, Grid, Icon, Tab, TabGroup, TabList, TabPanel, TabPanels, Text } from "@tremor/react";
import { Col, Grid, Icon, Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react";
import type { UploadProps } from "antd";
import { Form, Typography } from "antd";
import { PlusCircleOutlined } from "@ant-design/icons";
@ -62,6 +62,12 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, te
const [selectedModelId, setSelectedModelId] = useState<string | null>(null);
const [selectedTeamId, setSelectedTeamId] = useState<string | null>(null);
const [selectedTabIndex, setSelectedTabIndex] = useState(0);
const [showMissingProviderBanner, setShowMissingProviderBanner] = useState(() => {
if (typeof window !== "undefined") {
return localStorage.getItem("hideMissingProviderBanner") !== "true";
}
return true;
});
const queryClient = useQueryClient();
const { data: modelDataResponse, isLoading: isLoadingModels, refetch: refetchModels } = useModelsInfo();
@ -160,7 +166,7 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, te
const handleRefreshClick = () => {
const currentDate = new Date();
setLastRefreshed(currentDate.toLocaleString());
setLastRefreshed(currentDate.toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' }));
queryClient.invalidateQueries({ queryKey: ["models", "list"] });
refetchModels();
};
@ -282,43 +288,75 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, te
<p className="text-sm text-gray-600">Add and manage models for the proxy</p>
)}
</div>
{!showMissingProviderBanner && (
<a
href="https://models.litellm.ai/?request=true"
target="_blank"
rel="noopener noreferrer"
className="inline-flex items-center gap-1.5 px-3 py-1.5 text-xs font-medium text-[#6366f1] hover:text-[#5558e3] border border-[#6366f1] hover:border-[#5558e3] rounded-lg transition-colors"
>
<PlusCircleOutlined style={{ fontSize: "12px" }} />
Request Provider
</a>
)}
</div>
{/* Missing Provider Banner */}
<div className="mb-4 px-4 py-3 bg-blue-50 rounded-lg border border-blue-100 flex items-center gap-4">
<div className="flex-shrink-0 w-10 h-10 bg-white rounded-full flex items-center justify-center border border-blue-200">
<PlusCircleOutlined style={{ fontSize: "18px", color: "#6366f1" }} />
</div>
<div className="flex-1 min-w-0">
<h4 className="text-gray-900 font-semibold text-sm m-0">Missing a provider?</h4>
<p className="text-gray-500 text-xs m-0 mt-0.5">
The LiteLLM engineering team is constantly adding support for new LLM models, providers, endpoints. If
you don&apos;t see the one you need, let us know and we&apos;ll prioritize it.
</p>
</div>
<a
href="https://models.litellm.ai/?request=true"
target="_blank"
rel="noopener noreferrer"
className="flex-shrink-0 inline-flex items-center gap-2 px-4 py-2 bg-[#6366f1] hover:bg-[#5558e3] text-white text-sm font-medium rounded-lg transition-colors"
>
Request Provider
<svg
xmlns="http://www.w3.org/2000/svg"
className="h-4 w-4"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
strokeWidth={2}
{showMissingProviderBanner && (
<div className="mb-4 px-4 py-3 bg-blue-50 rounded-lg border border-blue-100 flex items-center gap-4">
<div className="flex-shrink-0 w-10 h-10 bg-white rounded-full flex items-center justify-center border border-blue-200">
<PlusCircleOutlined style={{ fontSize: "18px", color: "#6366f1" }} />
</div>
<div className="flex-1 min-w-0">
<h4 className="text-gray-900 font-semibold text-sm m-0">Missing a provider?</h4>
<p className="text-gray-500 text-xs m-0 mt-0.5">
The LiteLLM engineering team is constantly adding support for new LLM models, providers, endpoints. If
you don&apos;t see the one you need, let us know and we&apos;ll prioritize it.
</p>
</div>
<a
href="https://models.litellm.ai/?request=true"
target="_blank"
rel="noopener noreferrer"
className="flex-shrink-0 inline-flex items-center gap-2 px-4 py-2 bg-[#6366f1] hover:bg-[#5558e3] text-white text-sm font-medium rounded-lg transition-colors"
>
<path
strokeLinecap="round"
strokeLinejoin="round"
d="M10 6H6a2 2 0 00-2 2v10a2 2 0 002 2h10a2 2 0 002-2v-4M14 4h6m0 0v6m0-6L10 14"
/>
</svg>
</a>
</div>
Request Provider
<svg
xmlns="http://www.w3.org/2000/svg"
className="h-4 w-4"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
strokeWidth={2}
>
<path
strokeLinecap="round"
strokeLinejoin="round"
d="M10 6H6a2 2 0 00-2 2v10a2 2 0 002 2h10a2 2 0 002-2v-4M14 4h6m0 0v6m0-6L10 14"
/>
</svg>
</a>
<button
onClick={() => {
setShowMissingProviderBanner(false);
localStorage.setItem("hideMissingProviderBanner", "true");
}}
className="flex-shrink-0 p-1 text-gray-400 hover:text-gray-600 hover:bg-gray-100 rounded-full transition-colors"
aria-label="Dismiss banner"
>
<svg
xmlns="http://www.w3.org/2000/svg"
className="h-5 w-5"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
strokeWidth={2}
>
<path strokeLinecap="round" strokeLinejoin="round" d="M6 18L18 6M6 6l12 12" />
</svg>
</button>
</div>
)}
{selectedModelId && !isLoading ? (
<ModelInfoView
modelId={selectedModelId}
@ -348,13 +386,13 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, te
{all_admin_roles.includes(userRole) && <Tab>Price Data Reload</Tab>}
</div>
<div className="flex items-center space-x-2">
{lastRefreshed && <Text>Last Refreshed: {lastRefreshed}</Text>}
<div className="flex items-center space-x-2 self-center">
{lastRefreshed && <span className="text-xs text-gray-500">Last Refreshed: {lastRefreshed}</span>}
<Icon
icon={RefreshIcon} // Modify as necessary for correct icon name
icon={RefreshIcon}
variant="shadow"
size="xs"
className="self-center"
className="cursor-pointer"
onClick={handleRefreshClick}
/>
</div>

122
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx
View file

@ -1,8 +1,31 @@
import * as useAuthorizedModule from "@/app/(dashboard)/hooks/useAuthorized";
import { renderWithProviders, screen, waitFor } from "../../../../../tests/test-utils";
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
import { renderWithProviders } from "../../../../../tests/test-utils";
import { beforeEach, describe, expect, it, vi } from "vitest";
import AllModelsTab from "./AllModelsTab";
// Mock modelDeleteCall
const mockModelDeleteCall = vi.fn().mockResolvedValue({});
vi.mock("@/components/networking", () => ({
modelDeleteCall: (...args: any[]) => mockModelDeleteCall(...args),
}));
// Mock NotificationsManager
vi.mock("@/components/molecules/notifications_manager", () => ({
default: {
success: vi.fn(),
fromBackend: vi.fn(),
},
}));
// Mock react-query
const mockInvalidateQueries = vi.fn();
vi.mock("@tanstack/react-query", () => ({
useQueryClient: () => ({
invalidateQueries: mockInvalidateQueries,
}),
}));
// Mock the useModelsInfo hook
const mockUseModelsInfo = vi.fn(() => ({
data: { data: [], total_count: 0, current_page: 1, total_pages: 1, size: 50 },
@ -493,4 +516,101 @@ describe("AllModelsTab", () => {
const previousButton = screen.getByRole("button", { name: /previous/i });
expect(previousButton).toBeDisabled();
});
it("should pass setDeleteModalModelId to columns for delete functionality", async () => {
// This test verifies that the delete modal setter is passed to columns
// The actual modal rendering is handled by DeleteResourceModal component
mockUseTeams.mockReturnValue({
data: [],
isLoading: false,
error: null,
refetch: vi.fn(),
});
mockUseModelCostMap.mockReturnValue(
createModelCostMapMock({
"gpt-4-delete-test": { litellm_provider: "openai" },
}),
);
const modelData = createPaginatedModelData([
{
model_name: "gpt-4-delete-test",
litellm_model_name: "gpt-4-delete-test",
provider: "openai",
model_info: {
id: "model-to-delete",
db_model: true,
direct_access: true,
access_via_team_ids: [],
access_groups: [],
created_by: "user-123",
created_at: "2024-01-01",
updated_at: "2024-01-01",
},
},
], 1, 1, 1, 50);
mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null, refetch: vi.fn() });
render(<AllModelsTab {...defaultProps} />);
await waitFor(() => {
expect(screen.getByText("gpt-4-delete-test")).toBeInTheDocument();
});
// Verify the DB Model badge is shown (indicating it can be deleted)
expect(screen.getByText("DB Model")).toBeInTheDocument();
});
it("should render clickable model ID that calls setSelectedModelId", async () => {
mockUseTeams.mockReturnValue({
data: [],
isLoading: false,
error: null,
refetch: vi.fn(),
});
mockUseModelCostMap.mockReturnValue(
createModelCostMapMock({
"gpt-4-clickable": { litellm_provider: "openai" },
}),
);
const modelData = createPaginatedModelData([
{
model_name: "gpt-4-clickable",
litellm_model_name: "gpt-4-clickable",
provider: "openai",
model_info: {
id: "clickable-model-id",
db_model: true,
direct_access: true,
access_via_team_ids: [],
access_groups: [],
created_by: "user-123",
created_at: "2024-01-01",
updated_at: "2024-01-01",
},
},
], 1, 1, 1, 50);
mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null, refetch: vi.fn() });
render(<AllModelsTab {...defaultProps} />);
await waitFor(() => {
expect(screen.getByText("gpt-4-clickable")).toBeInTheDocument();
});
// Click on the Model ID cell which should call setSelectedModelId
const modelIdCell = screen.getByText("clickable-model-id");
expect(modelIdCell).toBeInTheDocument();
fireEvent.click(modelIdCell);
await waitFor(() => {
expect(mockSetSelectedModelId).toHaveBeenCalledWith("clickable-model-id");
});
});
});

65
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx
View file

@ -5,8 +5,12 @@ import { Team } from "@/components/key_team_helpers/key_list";
import { AllModelsDataTable } from "@/components/model_dashboard/all_models_table";
import { columns } from "@/components/molecules/models/columns";
import { getDisplayModelName } from "@/components/view_model/model_name_display";
import DeleteResourceModal from "@/components/common_components/DeleteResourceModal";
import NotificationsManager from "@/components/molecules/notifications_manager";
import { modelDeleteCall } from "@/components/networking";
import { InfoCircleOutlined, SettingOutlined } from "@ant-design/icons";
import { PaginationState, SortingState } from "@tanstack/react-table";
import { useQueryClient } from "@tanstack/react-query";
import { Grid, TabPanel } from "@tremor/react";
import { Badge, Button, Select, Skeleton, Space, Typography } from "antd";
import ModelSettingsModal from "@/components/model_dashboard/ModelSettingsModal/ModelSettingsModal";
@ -35,8 +39,9 @@ const AllModelsTab = ({
setSelectedTeamId,
}: AllModelsTabProps) => {
const { data: modelCostMapData, isLoading: isLoadingModelCostMap } = useModelCostMap();
const { userId, userRole, premiumUser } = useAuthorized();
const { accessToken, userId, userRole, premiumUser } = useAuthorized();
const { data: teams, isLoading: isLoadingTeams } = useTeams();
const queryClient = useQueryClient();
const [modelNameSearch, setModelNameSearch] = useState<string>("");
const [debouncedSearch, setDebouncedSearch] = useState<string>("");
@ -95,7 +100,7 @@ const AllModelsTab = ({
return sort.desc ? "desc" : "asc";
}, [sorting]);
const { data: rawModelData, isLoading: isLoadingModelsInfo } = useModelsInfo(
const { data: rawModelData, isLoading: isLoadingModelsInfo, refetch: refetchModels } = useModelsInfo(
currentPage,
pageSize,
debouncedSearch || undefined,
@ -120,6 +125,9 @@ const AllModelsTab = ({
return transformModelData(rawModelData, getProviderFromModel);
}, [rawModelData, modelCostMapData]);
const [deleteModalModelId, setDeleteModalModelId] = useState<string | null>(null);
const [deleteLoading, setDeleteLoading] = useState(false);
// Get pagination metadata from the response
const paginationMeta = useMemo(() => {
if (!rawModelData) {
@ -190,6 +198,28 @@ const AllModelsTab = ({
setSorting([]);
};
const modelToDelete = useMemo(() => {
if (!deleteModalModelId || !modelData?.data) return null;
return modelData.data.find((model: any) => model.model_info.id === deleteModalModelId);
}, [deleteModalModelId, modelData]);
const handleDeleteModel = async () => {
if (!accessToken || !deleteModalModelId) return;
try {
setDeleteLoading(true);
await modelDeleteCall(accessToken, deleteModalModelId);
NotificationsManager.success("Model deleted successfully");
queryClient.invalidateQueries({ queryKey: ["models", "list"] });
refetchModels();
} catch (error) {
console.error("Error deleting model:", error);
NotificationsManager.fromBackend(error);
} finally {
setDeleteLoading(false);
setDeleteModalModelId(null);
}
};
return (
<TabPanel>
<Grid>
@ -504,6 +534,7 @@ const AllModelsTab = ({
() => { },
expandedRows,
setExpandedRows,
setDeleteModalModelId,
)}
data={filteredData}
isLoading={isLoadingModelsInfo}
@ -512,10 +543,40 @@ const AllModelsTab = ({
pagination={pagination}
onPaginationChange={setPagination}
enablePagination={true}
onRowClick={(model: any) => setSelectedModelId(model.model_info.id)}
/>
</div>
</div>
</Grid>
<DeleteResourceModal
isOpen={!!deleteModalModelId}
title="Delete Model"
alertMessage="This action cannot be undone."
message="Are you sure you want to delete this model?"
resourceInformationTitle="Model Information"
resourceInformation={modelToDelete ? [
{
label: "Model Name",
value: modelToDelete.model_name || "Not Set",
},
{
label: "LiteLLM Model Name",
value: modelToDelete.litellm_model_name || "Not Set",
},
{
label: "Provider",
value: modelToDelete.provider || "Not Set",
},
{
label: "Created By",
value: modelToDelete.model_info?.created_by || "Not Set",
},
] : []}
onCancel={() => setDeleteModalModelId(null)}
onOk={handleDeleteModel}
confirmLoading={deleteLoading}
/>
<ModelSettingsModal
isVisible={isModelSettingsModalVisible}
onCancel={() => setIsModelSettingsModalVisible(false)}

8
ui/litellm-dashboard/src/components/model_dashboard/all_models_table.tsx
View file

@ -30,6 +30,7 @@ interface AllModelsDataTableProps<TData, TValue> {
pagination?: PaginationState;
onPaginationChange?: OnChangeFn<PaginationState>;
enablePagination?: boolean;
onRowClick?: (row: TData) => void;
}
export function AllModelsDataTable<TData, TValue>({
@ -41,6 +42,7 @@ export function AllModelsDataTable<TData, TValue>({
pagination,
onPaginationChange,
enablePagination = false,
onRowClick,
}: AllModelsDataTableProps<TData, TValue>) {
const [columnResizeMode] = React.useState<ColumnResizeMode>("onChange");
const [columnSizing, setColumnSizing] = React.useState({});
@ -174,7 +176,11 @@ export function AllModelsDataTable<TData, TValue>({
</TableRow>
) : tableInstance.getRowModel().rows.length > 0 ? (
tableInstance.getRowModel().rows.map((row) => (
<TableRow key={row.id}>
<TableRow
key={row.id}
className={onRowClick ? "cursor-pointer hover:bg-gray-50" : ""}
onClick={() => onRowClick?.(row.original)}
>
{row.getVisibleCells().map((cell) => (
<TableCell
key={cell.id}

8
ui/litellm-dashboard/src/components/model_dashboard/table.tsx
View file

@ -30,6 +30,7 @@ interface ModelDataTableProps<TData, TValue> {
pagination?: PaginationState;
onPaginationChange?: OnChangeFn<PaginationState>;
enablePagination?: boolean;
onRowClick?: (row: TData) => void;
}
export function ModelDataTable<TData, TValue>({
@ -40,6 +41,7 @@ export function ModelDataTable<TData, TValue>({
pagination,
onPaginationChange,
enablePagination = false,
onRowClick,
}: ModelDataTableProps<TData, TValue>) {
const [sorting, setSorting] = React.useState<SortingState>(defaultSorting);
const [columnResizeMode] = React.useState<ColumnResizeMode>("onChange");
@ -164,7 +166,11 @@ export function ModelDataTable<TData, TValue>({
</TableRow>
) : tableInstance.getRowModel().rows.length > 0 ? (
tableInstance.getRowModel().rows.map((row) => (
<TableRow key={row.id}>
<TableRow
key={row.id}
onClick={() => onRowClick?.(row.original)}
className={onRowClick ? "cursor-pointer hover:bg-gray-50" : ""}
>
{row.getVisibleCells().map((cell) => (
<TableCell
key={cell.id}

18
ui/litellm-dashboard/src/components/molecules/models/columns.tsx
View file

@ -52,6 +52,7 @@ export const columns = (
handleRefreshClick: () => void,
expandedRows: Set<string>,
setExpandedRows: (expandedRows: Set<string>) => void,
onDeleteClick?: (modelId: string) => void,
): ColumnDef<ModelData>[] => [
{
header: () => <span className="text-sm font-semibold">Model ID</span>,
@ -67,7 +68,10 @@ export const columns = (
ellipsis
className="text-blue-500 bg-blue-50 hover:bg-blue-100 text-xs cursor-pointer w-full block"
style={{ fontSize: 14, padding: '1px 8px' }}
onClick={() => setSelectedModelId(model.model_info.id)}
onClick={(e) => {
e.stopPropagation();
setSelectedModelId(model.model_info.id);
}}
>
{model.model_info.id}
</Text>
@ -297,7 +301,10 @@ export const columns = (
size="xs"
variant="light"
className="font-mono text-blue-500 bg-blue-50 hover:bg-blue-100 text-xs font-normal px-2 py-0.5 text-left overflow-hidden truncate w-full"
onClick={() => setSelectedTeamId(model.model_info.team_id)}
onClick={(e: React.MouseEvent) => {
e.stopPropagation();
setSelectedTeamId(model.model_info.team_id);
}}
>
{model.model_info.team_id.slice(0, 7)}...
</Button>
@ -409,9 +416,10 @@ export const columns = (
<Icon
icon={TrashIcon}
size="sm"
onClick={() => {
if (canEditModel) {
setSelectedModelId(model.model_info.id);
onClick={(e) => {
e.stopPropagation();
if (canEditModel && onDeleteClick) {
onDeleteClick(model.model_info.id);
}
}}
className={!canEditModel ? "opacity-50 cursor-not-allowed" : "cursor-pointer hover:text-red-600"}

10
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -9038,9 +9038,7 @@ export const updateUiSettings = async (accessToken: string, settings: Record<str
};
// ============================================================
// Claude Code Marketplace Networking Functions
// ============================================================
/**
* Get public marketplace catalog (no authentication required)
@ -9386,7 +9384,7 @@ export interface ToolPolicyOption {
description: string;
}
interface ToolPolicyOptionsResponse {
export interface ToolPolicyOptionsResponse {
input_policies: ToolPolicyOption[];
output_policies: ToolPolicyOption[];
}
@ -9439,12 +9437,12 @@ export interface ToolPolicyOverrideRow {
updated_at?: string;
}
interface ToolDetailResponse {
export interface ToolDetailResponse {
tool: ToolRow;
overrides: ToolPolicyOverrideRow[];
}
interface ToolUsageLogEntry {
export interface ToolUsageLogEntry {
id: string;
timestamp: string;
model?: string | null;
@ -9453,7 +9451,7 @@ interface ToolUsageLogEntry {
input_snippet?: string | null;
}
interface ToolUsageLogsResponse {
export interface ToolUsageLogsResponse {
logs: ToolUsageLogEntry[];
total: number;
page: number;