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

Merge remote-tracking branch 'origin' into litellm_org_usage

Browse Source
This commit is contained in:
yuneng-jiang
2025-11-13 15:26:59 -08:00
parent 35eece2b7b 124ba463f8
commit d62dfd4be2
72 changed files with 5563 additions and 791 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/completion/image_generation_chat.md
+2 -2
View File
@@ -224,8 +224,8 @@ asyncio.run(generate_image())
| Provider | Model |
|----------|--------|
| Google AI Studio | `gemini/gemini-2.5-flash-image-preview` |
| Vertex AI | `vertex_ai/gemini-2.5-flash-image-preview` |
| Google AI Studio | `gemini/gemini-2.0-flash-preview-image-generation`, `gemini/gemini-2.5-flash-image-preview` |
| Vertex AI | `vertex_ai/gemini-2.0-flash-preview-image-generation`, `vertex_ai/gemini-2.5-flash-image-preview` |
## Spec
docs/my-website/docs/image_edits.md
+81 -3
View File
@@ -14,9 +14,9 @@ LiteLLM provides image editing functionality that maps to OpenAI's `/images/edit
| Fallbacks | ✅ | Works between supported models |
| Loadbalancing | ✅ | Works between supported models |
| Supported operations | Create image edits | Single and multiple images supported |
| Supported LiteLLM SDK Versions | 1.63.8+ | |
| Supported LiteLLM Proxy Versions | 1.71.1+ | |
| Supported LLM providers | **OpenAI** | Currently only `openai` is 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)** | Gemini supports the new `gemini-2.5-flash-image` family |
#### ⚡️See all supported models and providers at [models.litellm.ai](https://models.litellm.ai/)
@@ -149,6 +149,54 @@ for i, image_data in enumerate(response.data):
print(f"Image {i+1}: {image_data.url}")
```
```
</TabItem>
<TabItem value="gemini" label="Gemini">
#### Basic Image Edit
```python showLineNumbers title="Gemini Image Edit"
import base64
import os
from litellm import image_edit
os.environ["GEMINI_API_KEY"] = "your-api-key"
response = image_edit(
model="gemini/gemini-2.5-flash-image",
image=open("original_image.png", "rb"),
prompt="Add aurora borealis to the night sky",
size="1792x1024", # mapped to aspectRatio=16:9 for Gemini
)
edited_image_bytes = base64.b64decode(response.data[0].b64_json)
with open("edited_image.png", "wb") as f:
f.write(edited_image_bytes)
```
#### Multiple Images Edit
```python showLineNumbers title="Gemini Multiple Images Edit"
import base64
import os
from litellm import image_edit
os.environ["GEMINI_API_KEY"] = "your-api-key"
response = image_edit(
model="gemini/gemini-2.5-flash-image",
image=[
open("scene.png", "rb"),
open("style_reference.png", "rb"),
],
prompt="Blend the reference style into the scene while keeping the subject sharp.",
)
for idx, image_obj in enumerate(response.data):
with open(f"gemini_edit_{idx}.png", "wb") as f:
f.write(base64.b64decode(image_obj.b64_json))
```
</TabItem>
</Tabs>
@@ -224,6 +272,36 @@ curl -X POST "http://localhost:4000/v1/images/edits" \
-F "response_format=url"
```
```
</TabItem>
<TabItem value="gemini" label="Gemini">
1. Add the Gemini image edit model to your `config.yaml`:
```yaml showLineNumbers title="Gemini Proxy Configuration"
model_list:
- model_name: gemini-image-edit
litellm_params:
model: gemini/gemini-2.5-flash-image
api_key: os.environ/GEMINI_API_KEY
```
2. Start the LiteLLM proxy server:
```bash showLineNumbers title="Start LiteLLM Proxy Server"
litellm --config /path/to/config.yaml
```
3. Make an image edit request (Gemini responses are base64-only):
```bash showLineNumbers title="Gemini Proxy Image Edit"
curl -X POST "http://0.0.0.0:4000/v1/images/edits" \
-H "Authorization: Bearer <YOUR-LITELLM-KEY>" \
-F "model=gemini-image-edit" \
-F "image=@original_image.png" \
-F "prompt=Add a warm golden-hour glow to the scene" \
-F "size=1024x1024"
```
</TabItem>
</Tabs>
docs/my-website/docs/providers/gemini.md
+21 -8
View File
@@ -10,7 +10,7 @@ import TabItem from '@theme/TabItem';
| Provider Route on LiteLLM | `gemini/` |
| Provider Doc | [Google AI Studio ↗](https://aistudio.google.com/) |
| API Endpoint for Provider | https://generativelanguage.googleapis.com |
| Supported OpenAI Endpoints | `/chat/completions`, [`/embeddings`](../embedding/supported_embedding#gemini-ai-embedding-models), `/completions`, [`/videos`](./gemini/videos.md) |
| Supported OpenAI Endpoints | `/chat/completions`, [`/embeddings`](../embedding/supported_embedding#gemini-ai-embedding-models), `/completions`, [`/videos`](./gemini/videos.md), [`/images/edits`](../image_edits.md) |
| Pass-through Endpoint | [Supported](../pass_through/google_ai_studio.md) |
<br />
@@ -64,16 +64,21 @@ response = completion(
LiteLLM translates OpenAI's `reasoning_effort` to Gemini's `thinking` parameter. [Code](https://github.com/BerriAI/litellm/blob/620664921902d7a9bfb29897a7b27c1a7ef4ddfb/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py#L362)
Added an additional non-OpenAI standard "disable" value for non-reasoning Gemini requests.
**Cost Optimization:** Use `reasoning_effort="none"` (OpenAI standard) for significant cost savings - up to 96% cheaper. [Google's docs](https://ai.google.dev/gemini-api/docs/openai)
:::info
Note: Reasoning cannot be turned off on Gemini 2.5 Pro models.
:::
**Mapping**
| reasoning_effort | thinking |
| ---------------- | -------- |
| "disable" | "budget_tokens": 0 |
| "low" | "budget_tokens": 1024 |
| "medium" | "budget_tokens": 2048 |
| "high" | "budget_tokens": 4096 |
| reasoning_effort | thinking | Notes |
| ---------------- | -------- | ----- |
| "none" | "budget_tokens": 0, "includeThoughts": false | 💰 **Recommended for cost optimization** - OpenAI-compatible, always 0 |
| "disable" | "budget_tokens": DEFAULT (0), "includeThoughts": false | LiteLLM-specific, configurable via env var |
| "low" | "budget_tokens": 1024 | |
| "medium" | "budget_tokens": 2048 | |
| "high" | "budget_tokens": 4096 | |
<Tabs>
<TabItem value="sdk" label="SDK">
@@ -81,6 +86,14 @@ Added an additional non-OpenAI standard "disable" value for non-reasoning Gemini
```python
from litellm import completion
# Cost-optimized: Use reasoning_effort="none" for best pricing
resp = completion(
model="gemini/gemini-2.0-flash-thinking-exp-01-21",
messages=[{"role": "user", "content": "What is the capital of France?"}],
reasoning_effort="none", # Up to 96% cheaper!
)
# Or use other levels: "low", "medium", "high"
resp = completion(
model="gemini/gemini-2.5-flash-preview-04-17",
messages=[{"role": "user", "content": "What is the capital of France?"}],
docs/my-website/docs/providers/openai.md
+71
View File
@@ -410,6 +410,77 @@ Expected Response:
```
### Advanced: Using `reasoning_effort` with `summary` field
By default, `reasoning_effort` accepts a string value (`"low"`, `"medium"`, `"high"`, `"minimal"`) and only sets the effort level without including a reasoning summary.
To opt-in to the `summary` feature, you can pass `reasoning_effort` as a dictionary. **Note:** The `summary` field requires your OpenAI organization to have verification status. Using `summary` without verification will result in a 400 error from OpenAI.
<Tabs>
<TabItem value="sdk" label="SDK">
```python
# Option 1: String format (default - no summary)
response = litellm.completion(
model="openai/responses/gpt-5-mini",
messages=[{"role": "user", "content": "What is the capital of France?"}],
reasoning_effort="high" # Only sets effort level
)
# Option 2: Dict format (with optional summary - requires org verification)
response = litellm.completion(
model="openai/responses/gpt-5-mini",
messages=[{"role": "user", "content": "What is the capital of France?"}],
reasoning_effort={"effort": "high", "summary": "auto"} # "auto", "detailed", or "concise" (not all supported by all models)
)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
```bash
# Option 1: String format (default - no summary)
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{
"model": "openai/responses/gpt-5-mini",
"messages": [{"role": "user", "content": "What is the capital of France?"}],
"reasoning_effort": "high"
}'
# Option 2: Dict format (with optional summary - requires org verification)
# summary options: "auto", "detailed", or "concise" (not all supported by all models)
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{
"model": "openai/responses/gpt-5-mini",
"messages": [{"role": "user", "content": "What is the capital of France?"}],
"reasoning_effort": {"effort": "high", "summary": "auto"}
}'
```
</TabItem>
</Tabs>
**Summary field options:**
- `"auto"`: System automatically determines the appropriate summary level based on the model
- `"concise"`: Provides a shorter summary (not supported by GPT-5 series models)
- `"detailed"`: Offers a comprehensive reasoning summary
**Note:** GPT-5 series models support `"auto"` and `"detailed"`, but do not support `"concise"`. O-series models (o3-pro, o4-mini, o3) support all three options. Some models like o3-mini and o1 do not support reasoning summaries at all.
**Supported `reasoning_effort` values by model:**
| Model | Default (when not set) | Supported Values |
|-------|----------------------|------------------|
| `gpt-5` | `medium` | `minimal`, `low`, `medium`, `high` |
| `gpt-5-mini` | `medium` | `minimal`, `low`, `medium`, `high` |
| `gpt-5-codex` | `adaptive` | `low`, `medium`, `high` (no `minimal`) |
| `gpt-5-pro` | `high` | `high` only |
**Note:** `gpt-5-pro` only accepts `reasoning_effort="high"`. Other values will return an error. When `reasoning_effort` is not set (None), OpenAI defaults to the value shown in the "Default" column.
See [OpenAI Reasoning documentation](https://platform.openai.com/docs/guides/reasoning) for more details on organization verification requirements.
## OpenAI Chat Completion to Responses API Bridge
Call any Responses API model from OpenAI's `/chat/completions` endpoint.
docs/my-website/docs/providers/runwayml/text-to-speech.md
+244
View File
@@ -0,0 +1,244 @@
# RunwayML - Text-to-Speech
## Overview
| Property | Details |
|-------|-------|
| Description | RunwayML provides high-quality AI-powered text-to-speech with natural-sounding voices |
| Provider Route on LiteLLM | `runwayml/` |
| Supported Operations | [`/audio/speech`](#quick-start) |
| Link to Provider Doc | [RunwayML API ↗](https://docs.dev.runwayml.com/) |
LiteLLM supports RunwayML's text-to-speech API with automatic task polling, allowing you to generate natural-sounding audio from text.
## Quick Start
```python showLineNumbers title="Basic Text-to-Speech"
from litellm import speech
import os
os.environ["RUNWAYML_API_KEY"] = "your-api-key"
response = speech(
model="runwayml/eleven_multilingual_v2",
input="Step right up, ladies and gentlemen! Have you ever wished for a toaster that's not just a toaster but a marvel of modern ingenuity?",
voice="alloy"
)
# Save the audio
with open("output.mp3", "wb") as f:
f.write(response.content)
```
## Authentication
Set your RunwayML API key:
```python showLineNumbers title="Set API Key"
import os
os.environ["RUNWAYML_API_KEY"] = "your-api-key"
```
## Supported Parameters
| Parameter | Type | Required | Description |
|-----------|------|----------|-------------|
| `model` | string | Yes | Model to use (e.g., `runwayml/eleven_multilingual_v2`) |
| `input` | string | Yes | Text to convert to speech |
| `voice` | string or dict | Yes | Voice to use (OpenAI name, RunwayML preset, or voice config) |
## Voice Options
### Using OpenAI Voice Names
OpenAI voice names are automatically mapped to appropriate RunwayML voices:
```python showLineNumbers title="OpenAI Voice Names"
from litellm import speech
# These OpenAI voice names work automatically
response = speech(
model="runwayml/eleven_multilingual_v2",
input="Hello, world!",
voice="alloy" # Maya - neutral, balanced female voice
)
```
**Voice Mappings:**
- `alloy` → Maya (neutral, balanced female voice)
- `echo` → James (male voice)
- `fable` → Bernard (warm, storytelling voice)
- `onyx` → Vincent (deep male voice)
- `nova` → Serene (warm, expressive female voice)
- `shimmer` → Ella (clear, friendly female voice)
### Using RunwayML Preset Voices
You can directly specify any RunwayML preset voice by passing the preset name as a string:
```python showLineNumbers title="RunwayML Preset Names"
from litellm import speech
# Pass the RunwayML voice name as a string
response = speech(
model="runwayml/eleven_multilingual_v2",
input="Hello, world!",
voice="Maya" # LiteLLM automatically formats this for RunwayML
)
# Try different RunwayML voices
response = speech(
model="runwayml/eleven_multilingual_v2",
input="Step right up, ladies and gentlemen!",
voice="Bernard" # Great for storytelling
)
```
**Available RunwayML Voices:**
Maya, Arjun, Serene, Bernard, Billy, Mark, Clint, Mabel, Chad, Leslie, Eleanor, Elias, Elliot, Grungle, Brodie, Sandra, Kirk, Kylie, Lara, Lisa, Malachi, Marlene, Martin, Miriam, Monster, Paula, Pip, Rusty, Ragnar, Xylar, Maggie, Jack, Katie, Noah, James, Rina, Ella, Mariah, Frank, Claudia, Niki, Vincent, Kendrick, Myrna, Tom, Wanda, Benjamin, Kiana, Rachel
:::tip
Simply pass the voice name as a string - LiteLLM automatically handles the internal RunwayML API format conversion.
:::
## Async Usage
```python showLineNumbers title="Async Text-to-Speech"
from litellm import aspeech
import os
import asyncio
os.environ["RUNWAYML_API_KEY"] = "your-api-key"
async def generate_speech():
response = await aspeech(
model="runwayml/eleven_multilingual_v2",
input="This is an asynchronous text-to-speech request.",
voice="nova"
)
with open("output.mp3", "wb") as f:
f.write(response.content)
print("Audio generated successfully!")
asyncio.run(generate_speech())
```
## LiteLLM Proxy Usage
Add RunwayML to your proxy configuration:
```yaml showLineNumbers title="config.yaml"
model_list:
- model_name: runway-tts
litellm_params:
model: runwayml/eleven_multilingual_v2
api_key: os.environ/RUNWAYML_API_KEY
```
Start the proxy:
```bash
litellm --config /path/to/config.yaml
```
Generate speech through the proxy:
```bash showLineNumbers title="Proxy Request"
curl --location 'http://localhost:4000/v1/audio/speech' \
--header 'Content-Type: application/json' \
--header 'x-litellm-api-key: sk-1234' \
--data '{
"model": "runwayml/eleven_multilingual_v2",
"input": "Hello from the LiteLLM proxy!",
"voice": "alloy"
}'
```
With RunwayML-specific voice:
```bash showLineNumbers title="Proxy Request with RunwayML Voice"
curl --location 'http://localhost:4000/v1/audio/speech' \
--header 'Content-Type: application/json' \
--header 'x-litellm-api-key: sk-1234' \
--data '{
"model": "runwayml/eleven_multilingual_v2",
"input": "Hello with a custom RunwayML voice!",
"voice": "Bernard"
}'
```
## Supported Models
| Model | Description |
|-------|-------------|
| `runwayml/eleven_multilingual_v2` | High-quality multilingual text-to-speech |
## Cost Tracking
LiteLLM automatically tracks RunwayML text-to-speech costs:
```python showLineNumbers title="Cost Tracking"
from litellm import speech, completion_cost
response = speech(
model="runwayml/eleven_multilingual_v2",
input="Hello, world!",
voice="alloy"
)
cost = completion_cost(completion_response=response)
print(f"Text-to-speech cost: ${cost}")
```
## Supported Features
| Feature | Supported |
|---------|-----------|
| Text-to-Speech | ✅ |
| Cost Tracking | ✅ |
| Logging | ✅ |
| Fallbacks | ✅ |
| Load Balancing | ✅ |
| 50+ Voice Presets | ✅ |
## How It Works
RunwayML uses an asynchronous task-based API pattern. LiteLLM handles the polling and response transformation automatically.
### Complete Flow Diagram
```mermaid
sequenceDiagram
participant Client
box rgb(200, 220, 255) LiteLLM AI Gateway
participant LiteLLM
end
participant RunwayML as RunwayML API
participant Storage as Audio Storage
Client->>LiteLLM: POST /audio/speech (OpenAI format)
Note over LiteLLM: Transform to RunwayML format<br/>Map voice to preset ID
LiteLLM->>RunwayML: POST v1/text_to_speech
RunwayML-->>LiteLLM: 200 OK + task ID
Note over LiteLLM: Automatic Polling
loop Every 2 seconds
LiteLLM->>RunwayML: GET v1/tasks/{task_id}
RunwayML-->>LiteLLM: Status: RUNNING
end
LiteLLM->>RunwayML: GET v1/tasks/{task_id}
RunwayML-->>LiteLLM: Status: SUCCEEDED + audio URL
LiteLLM->>Storage: GET audio URL
Storage-->>LiteLLM: Audio data (MP3)
Note over LiteLLM: Return audio content
LiteLLM-->>Client: Audio Response (binary)
```
docs/my-website/docs/providers/vertex.md
+47 -509
View File
@@ -1604,6 +1604,53 @@ litellm.vertex_location = "us-central1 # Your Location
| gemini-2.5-flash-preview-09-2025 | `completion('gemini-2.5-flash-preview-09-2025', messages)`, `completion('vertex_ai/gemini-2.5-flash-preview-09-2025', messages)` |
| gemini-2.5-flash-lite-preview-09-2025 | `completion('gemini-2.5-flash-lite-preview-09-2025', messages)`, `completion('vertex_ai/gemini-2.5-flash-lite-preview-09-2025', messages)` |
## Private Service Connect (PSC) Endpoints
LiteLLM supports Vertex AI models deployed to Private Service Connect (PSC) endpoints, allowing you to use custom `api_base` URLs for private deployments.
### Usage
```python
from litellm import completion
# Use PSC endpoint with custom api_base
response = completion(
model="vertex_ai/1234567890", # Numeric endpoint ID
messages=[{"role": "user", "content": "Hello!"}],
api_base="http://10.96.32.8", # Your PSC endpoint
vertex_project="my-project-id",
vertex_location="us-central1"
)
```
**Key Features:**
- Supports both numeric endpoint IDs and custom model names
- Works with both completion and embedding endpoints
- Automatically constructs full PSC URL: `{api_base}/v1/projects/{project}/locations/{location}/endpoints/{model}:{endpoint}`
- Compatible with streaming requests
### Configuration
Add PSC endpoints to your `config.yaml`:
```yaml
model_list:
- model_name: psc-gemini
litellm_params:
model: vertex_ai/1234567890 # Numeric endpoint ID
api_base: "http://10.96.32.8" # Your PSC endpoint
vertex_project: "my-project-id"
vertex_location: "us-central1"
vertex_credentials: "/path/to/service_account.json"
- model_name: psc-embedding
litellm_params:
model: vertex_ai/text-embedding-004
api_base: "http://10.96.32.8" # Your PSC endpoint
vertex_project: "my-project-id"
vertex_location: "us-central1"
vertex_credentials: "/path/to/service_account.json"
```
## Fine-tuned Models
You can call fine-tuned Vertex AI Gemini models through LiteLLM
@@ -2042,515 +2089,6 @@ curl http://0.0.0.0:4000/v1/chat/completions \
| code-gecko@latest| `completion('code-gecko@latest', messages)` |
## **Embedding Models**
#### Usage - Embedding
<Tabs>
<TabItem value="sdk" label="SDK">
```python
import litellm
from litellm import embedding
litellm.vertex_project = "hardy-device-38811" # Your Project ID
litellm.vertex_location = "us-central1" # proj location
response = embedding(
model="vertex_ai/textembedding-gecko",
input=["good morning from litellm"],
)
print(response)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY">
1. Add model to config.yaml
```yaml
model_list:
- model_name: snowflake-arctic-embed-m-long-1731622468876
litellm_params:
model: vertex_ai/<your-model-id>
vertex_project: "adroit-crow-413218"
vertex_location: "us-central1"
vertex_credentials: adroit-crow-413218-a956eef1a2a8.json
litellm_settings:
drop_params: True
```
2. Start Proxy
```
$ litellm --config /path/to/config.yaml
```
3. Make Request using OpenAI Python SDK, Langchain Python SDK
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
response = client.embeddings.create(
model="snowflake-arctic-embed-m-long-1731622468876",
input = ["good morning from litellm", "this is another item"],
)
print(response)
```
</TabItem>
</Tabs>
#### Supported Embedding Models
All models listed [here](https://github.com/BerriAI/litellm/blob/57f37f743886a0249f630a6792d49dffc2c5d9b7/model_prices_and_context_window.json#L835) are supported
| Model Name | Function Call |
|--------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| text-embedding-004 | `embedding(model="vertex_ai/text-embedding-004", input)` |
| text-multilingual-embedding-002 | `embedding(model="vertex_ai/text-multilingual-embedding-002", input)` |
| textembedding-gecko | `embedding(model="vertex_ai/textembedding-gecko", input)` |
| textembedding-gecko-multilingual | `embedding(model="vertex_ai/textembedding-gecko-multilingual", input)` |
| textembedding-gecko-multilingual@001 | `embedding(model="vertex_ai/textembedding-gecko-multilingual@001", input)` |
| textembedding-gecko@001 | `embedding(model="vertex_ai/textembedding-gecko@001", input)` |
| textembedding-gecko@003 | `embedding(model="vertex_ai/textembedding-gecko@003", input)` |
| text-embedding-preview-0409 | `embedding(model="vertex_ai/text-embedding-preview-0409", input)` |
| text-multilingual-embedding-preview-0409 | `embedding(model="vertex_ai/text-multilingual-embedding-preview-0409", input)` |
| Fine-tuned OR Custom Embedding models | `embedding(model="vertex_ai/<your-model-id>", input)` |
### Supported OpenAI (Unified) Params
| [param](../embedding/supported_embedding.md#input-params-for-litellmembedding) | type | [vertex equivalent](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api) |
|-------|-------------|--------------------|
| `input` | **string or List[string]** | `instances` |
| `dimensions` | **int** | `output_dimensionality` |
| `input_type` | **Literal["RETRIEVAL_QUERY","RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", "FACT_VERIFICATION"]** | `task_type` |
#### Usage with OpenAI (Unified) Params
<Tabs>
<TabItem value="sdk" label="SDK">
```python
response = litellm.embedding(
model="vertex_ai/text-embedding-004",
input=["good morning from litellm", "gm"]
input_type = "RETRIEVAL_DOCUMENT",
dimensions=1,
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY">
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
response = client.embeddings.create(
model="text-embedding-004",
input = ["good morning from litellm", "gm"],
dimensions=1,
extra_body = {
"input_type": "RETRIEVAL_QUERY",
}
)
print(response)
```
</TabItem>
</Tabs>
### Supported Vertex Specific Params
| param | type |
|-------|-------------|
| `auto_truncate` | **bool** |
| `task_type` | **Literal["RETRIEVAL_QUERY","RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", "FACT_VERIFICATION"]** |
| `title` | **str** |
#### Usage with Vertex Specific Params (Use `task_type` and `title`)
You can pass any vertex specific params to the embedding model. Just pass them to the embedding function like this:
[Relevant Vertex AI doc with all embedding params](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api#request_body)
<Tabs>
<TabItem value="sdk" label="SDK">
```python
response = litellm.embedding(
model="vertex_ai/text-embedding-004",
input=["good morning from litellm", "gm"]
task_type = "RETRIEVAL_DOCUMENT",
title = "test",
dimensions=1,
auto_truncate=True,
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY">
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
response = client.embeddings.create(
model="text-embedding-004",
input = ["good morning from litellm", "gm"],
dimensions=1,
extra_body = {
"task_type": "RETRIEVAL_QUERY",
"auto_truncate": True,
"title": "test",
}
)
print(response)
```
</TabItem>
</Tabs>
## **Multi-Modal Embeddings**
Known Limitations:
- Only supports 1 image / video / image per request
- Only supports GCS or base64 encoded images / videos
### Usage
<Tabs>
<TabItem value="sdk" label="SDK">
Using GCS Images
```python
response = await litellm.aembedding(
model="vertex_ai/multimodalembedding@001",
input="gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png" # will be sent as a gcs image
)
```
Using base 64 encoded images
```python
response = await litellm.aembedding(
model="vertex_ai/multimodalembedding@001",
input="data:image/jpeg;base64,..." # will be sent as a base64 encoded image
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY (Unified Endpoint)">
1. Add model to config.yaml
```yaml
model_list:
- model_name: multimodalembedding@001
litellm_params:
model: vertex_ai/multimodalembedding@001
vertex_project: "adroit-crow-413218"
vertex_location: "us-central1"
vertex_credentials: adroit-crow-413218-a956eef1a2a8.json
litellm_settings:
drop_params: True
```
2. Start Proxy
```
$ litellm --config /path/to/config.yaml
```
3. Make Request use OpenAI Python SDK, Langchain Python SDK
<Tabs>
<TabItem value="OpenAI SDK" label="OpenAI SDK">
Requests with GCS Image / Video URI
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
# # request sent to model set on litellm proxy, `litellm --model`
response = client.embeddings.create(
model="multimodalembedding@001",
input = "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png",
)
print(response)
```
Requests with base64 encoded images
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
# # request sent to model set on litellm proxy, `litellm --model`
response = client.embeddings.create(
model="multimodalembedding@001",
input = "data:image/jpeg;base64,...",
)
print(response)
```
</TabItem>
<TabItem value="langchain" label="Langchain">
Requests with GCS Image / Video URI
```python
from langchain_openai import OpenAIEmbeddings
embeddings_models = "multimodalembedding@001"
embeddings = OpenAIEmbeddings(
model="multimodalembedding@001",
base_url="http://0.0.0.0:4000",
api_key="sk-1234", # type: ignore
)
query_result = embeddings.embed_query(
"gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"
)
print(query_result)
```
Requests with base64 encoded images
```python
from langchain_openai import OpenAIEmbeddings
embeddings_models = "multimodalembedding@001"
embeddings = OpenAIEmbeddings(
model="multimodalembedding@001",
base_url="http://0.0.0.0:4000",
api_key="sk-1234", # type: ignore
)
query_result = embeddings.embed_query(
"data:image/jpeg;base64,..."
)
print(query_result)
```
</TabItem>
</Tabs>
</TabItem>
<TabItem value="proxy-vtx" label="LiteLLM PROXY (Vertex SDK)">
1. Add model to config.yaml
```yaml
default_vertex_config:
vertex_project: "adroit-crow-413218"
vertex_location: "us-central1"
vertex_credentials: adroit-crow-413218-a956eef1a2a8.json
```
2. Start Proxy
```
$ litellm --config /path/to/config.yaml
```
3. Make Request use OpenAI Python SDK
```python
import vertexai
from vertexai.vision_models import Image, MultiModalEmbeddingModel, Video
from vertexai.vision_models import VideoSegmentConfig
from google.auth.credentials import Credentials
LITELLM_PROXY_API_KEY = "sk-1234"
LITELLM_PROXY_BASE = "http://0.0.0.0:4000/vertex-ai"
import datetime
class CredentialsWrapper(Credentials):
def __init__(self, token=None):
super().__init__()
self.token = token
self.expiry = None # or set to a future date if needed
def refresh(self, request):
pass
def apply(self, headers, token=None):
headers['Authorization'] = f'Bearer {self.token}'
@property
def expired(self):
return False # Always consider the token as non-expired
@property
def valid(self):
return True # Always consider the credentials as valid
credentials = CredentialsWrapper(token=LITELLM_PROXY_API_KEY)
vertexai.init(
project="adroit-crow-413218",
location="us-central1",
api_endpoint=LITELLM_PROXY_BASE,
credentials = credentials,
api_transport="rest",
)
model = MultiModalEmbeddingModel.from_pretrained("multimodalembedding")
image = Image.load_from_file(
"gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"
)
embeddings = model.get_embeddings(
image=image,
contextual_text="Colosseum",
dimension=1408,
)
print(f"Image Embedding: {embeddings.image_embedding}")
print(f"Text Embedding: {embeddings.text_embedding}")
```
</TabItem>
</Tabs>
### Text + Image + Video Embeddings
<Tabs>
<TabItem value="sdk" label="SDK">
Text + Image
```python
response = await litellm.aembedding(
model="vertex_ai/multimodalembedding@001",
input=["hey", "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"] # will be sent as a gcs image
)
```
Text + Video
```python
response = await litellm.aembedding(
model="vertex_ai/multimodalembedding@001",
input=["hey", "gs://my-bucket/embeddings/supermarket-video.mp4"] # will be sent as a gcs image
)
```
Image + Video
```python
response = await litellm.aembedding(
model="vertex_ai/multimodalembedding@001",
input=["gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png", "gs://my-bucket/embeddings/supermarket-video.mp4"] # will be sent as a gcs image
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY (Unified Endpoint)">
1. Add model to config.yaml
```yaml
model_list:
- model_name: multimodalembedding@001
litellm_params:
model: vertex_ai/multimodalembedding@001
vertex_project: "adroit-crow-413218"
vertex_location: "us-central1"
vertex_credentials: adroit-crow-413218-a956eef1a2a8.json
litellm_settings:
drop_params: True
```
2. Start Proxy
```
$ litellm --config /path/to/config.yaml
```
3. Make Request use OpenAI Python SDK, Langchain Python SDK
Text + Image
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
# # request sent to model set on litellm proxy, `litellm --model`
response = client.embeddings.create(
model="multimodalembedding@001",
input = ["hey", "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"],
)
print(response)
```
Text + Video
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
# # request sent to model set on litellm proxy, `litellm --model`
response = client.embeddings.create(
model="multimodalembedding@001",
input = ["hey", "gs://my-bucket/embeddings/supermarket-video.mp4"],
)
print(response)
```
Image + Video
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
# # request sent to model set on litellm proxy, `litellm --model`
response = client.embeddings.create(
model="multimodalembedding@001",
input = ["gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png", "gs://my-bucket/embeddings/supermarket-video.mp4"],
)
print(response)
```
</TabItem>
</Tabs>
## **Gemini TTS (Text-to-Speech) Audio Output**
:::info
docs/my-website/docs/providers/vertex_embedding.md
+587
View File
@@ -0,0 +1,587 @@
import Image from '@theme/IdealImage';
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Vertex AI Embedding
## Usage - Embedding
<Tabs>
<TabItem value="sdk" label="SDK">
```python
import litellm
from litellm import embedding
litellm.vertex_project = "hardy-device-38811" # Your Project ID
litellm.vertex_location = "us-central1" # proj location
response = embedding(
model="vertex_ai/textembedding-gecko",
input=["good morning from litellm"],
)
print(response)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY">
1. Add model to config.yaml
```yaml
model_list:
- model_name: snowflake-arctic-embed-m-long-1731622468876
litellm_params:
model: vertex_ai/<your-model-id>
vertex_project: "adroit-crow-413218"
vertex_location: "us-central1"
vertex_credentials: adroit-crow-413218-a956eef1a2a8.json
litellm_settings:
drop_params: True
```
2. Start Proxy
```
$ litellm --config /path/to/config.yaml
```
3. Make Request using OpenAI Python SDK, Langchain Python SDK
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
response = client.embeddings.create(
model="snowflake-arctic-embed-m-long-1731622468876",
input = ["good morning from litellm", "this is another item"],
)
print(response)
```
</TabItem>
</Tabs>
#### Supported Embedding Models
All models listed [here](https://github.com/BerriAI/litellm/blob/57f37f743886a0249f630a6792d49dffc2c5d9b7/model_prices_and_context_window.json#L835) are supported
| Model Name | Function Call |
|--------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| text-embedding-004 | `embedding(model="vertex_ai/text-embedding-004", input)` |
| text-multilingual-embedding-002 | `embedding(model="vertex_ai/text-multilingual-embedding-002", input)` |
| textembedding-gecko | `embedding(model="vertex_ai/textembedding-gecko", input)` |
| textembedding-gecko-multilingual | `embedding(model="vertex_ai/textembedding-gecko-multilingual", input)` |
| textembedding-gecko-multilingual@001 | `embedding(model="vertex_ai/textembedding-gecko-multilingual@001", input)` |
| textembedding-gecko@001 | `embedding(model="vertex_ai/textembedding-gecko@001", input)` |
| textembedding-gecko@003 | `embedding(model="vertex_ai/textembedding-gecko@003", input)` |
| text-embedding-preview-0409 | `embedding(model="vertex_ai/text-embedding-preview-0409", input)` |
| text-multilingual-embedding-preview-0409 | `embedding(model="vertex_ai/text-multilingual-embedding-preview-0409", input)` |
| Fine-tuned OR Custom Embedding models | `embedding(model="vertex_ai/<your-model-id>", input)` |
### Supported OpenAI (Unified) Params
| [param](../embedding/supported_embedding.md#input-params-for-litellmembedding) | type | [vertex equivalent](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api) |
|-------|-------------|--------------------|
| `input` | **string or List[string]** | `instances` |
| `dimensions` | **int** | `output_dimensionality` |
| `input_type` | **Literal["RETRIEVAL_QUERY","RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", "FACT_VERIFICATION"]** | `task_type` |
#### Usage with OpenAI (Unified) Params
<Tabs>
<TabItem value="sdk" label="SDK">
```python
response = litellm.embedding(
model="vertex_ai/text-embedding-004",
input=["good morning from litellm", "gm"]
input_type = "RETRIEVAL_DOCUMENT",
dimensions=1,
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY">
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
response = client.embeddings.create(
model="text-embedding-004",
input = ["good morning from litellm", "gm"],
dimensions=1,
extra_body = {
"input_type": "RETRIEVAL_QUERY",
}
)
print(response)
```
</TabItem>
</Tabs>
### Supported Vertex Specific Params
| param | type |
|-------|-------------|
| `auto_truncate` | **bool** |
| `task_type` | **Literal["RETRIEVAL_QUERY","RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", "FACT_VERIFICATION"]** |
| `title` | **str** |
#### Usage with Vertex Specific Params (Use `task_type` and `title`)
You can pass any vertex specific params to the embedding model. Just pass them to the embedding function like this:
[Relevant Vertex AI doc with all embedding params](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api#request_body)
<Tabs>
<TabItem value="sdk" label="SDK">
```python
response = litellm.embedding(
model="vertex_ai/text-embedding-004",
input=["good morning from litellm", "gm"]
task_type = "RETRIEVAL_DOCUMENT",
title = "test",
dimensions=1,
auto_truncate=True,
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY">
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
response = client.embeddings.create(
model="text-embedding-004",
input = ["good morning from litellm", "gm"],
dimensions=1,
extra_body = {
"task_type": "RETRIEVAL_QUERY",
"auto_truncate": True,
"title": "test",
}
)
print(response)
```
</TabItem>
</Tabs>
## **BGE Embeddings**
Use BGE (Baidu General Embedding) models deployed on Vertex AI.
### Usage
<Tabs>
<TabItem value="sdk" label="SDK">
```python showLineNumbers title="Using BGE on Vertex AI"
import litellm
response = litellm.embedding(
model="vertex_ai/bge/<your-endpoint-id>",
input=["Hello", "World"],
vertex_project="your-project-id",
vertex_location="your-location"
)
print(response)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY">
1. Add model to config.yaml
```yaml showLineNumbers title="config.yaml"
model_list:
- model_name: bge-embedding
litellm_params:
model: vertex_ai/bge/<your-endpoint-id>
vertex_project: "your-project-id"
vertex_location: "us-central1"
vertex_credentials: your-credentials.json
litellm_settings:
drop_params: True
```
2. Start Proxy
```bash
$ litellm --config /path/to/config.yaml
```
3. Make Request using OpenAI Python SDK
```python showLineNumbers title="Making requests to BGE"
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
response = client.embeddings.create(
model="bge-embedding",
input=["good morning from litellm", "this is another item"]
)
print(response)
```
Using a Private Service Connect (PSC) endpoint
```yaml showLineNumbers title="config.yaml (PSC)"
model_list:
- model_name: bge-small-en-v1.5
litellm_params:
model: vertex_ai/bge/1234567890
api_base: http://10.96.32.8 # Your PSC IP
vertex_project: my-project-id #optional
vertex_location: us-central1 #optional
```
</TabItem>
</Tabs>
## **Multi-Modal Embeddings**
Known Limitations:
- Only supports 1 image / video / image per request
- Only supports GCS or base64 encoded images / videos
### Usage
<Tabs>
<TabItem value="sdk" label="SDK">
Using GCS Images
```python
response = await litellm.aembedding(
model="vertex_ai/multimodalembedding@001",
input="gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png" # will be sent as a gcs image
)
```
Using base 64 encoded images
```python
response = await litellm.aembedding(
model="vertex_ai/multimodalembedding@001",
input="data:image/jpeg;base64,..." # will be sent as a base64 encoded image
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY (Unified Endpoint)">
1. Add model to config.yaml
```yaml
model_list:
- model_name: multimodalembedding@001
litellm_params:
model: vertex_ai/multimodalembedding@001
vertex_project: "adroit-crow-413218"
vertex_location: "us-central1"
vertex_credentials: adroit-crow-413218-a956eef1a2a8.json
litellm_settings:
drop_params: True
```
2. Start Proxy
```
$ litellm --config /path/to/config.yaml
```
3. Make Request use OpenAI Python SDK, Langchain Python SDK
<Tabs>
<TabItem value="OpenAI SDK" label="OpenAI SDK">
Requests with GCS Image / Video URI
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
# # request sent to model set on litellm proxy, `litellm --model`
response = client.embeddings.create(
model="multimodalembedding@001",
input = "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png",
)
print(response)
```
Requests with base64 encoded images
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
# # request sent to model set on litellm proxy, `litellm --model`
response = client.embeddings.create(
model="multimodalembedding@001",
input = "data:image/jpeg;base64,...",
)
print(response)
```
</TabItem>
<TabItem value="langchain" label="Langchain">
Requests with GCS Image / Video URI
```python
from langchain_openai import OpenAIEmbeddings
embeddings_models = "multimodalembedding@001"
embeddings = OpenAIEmbeddings(
model="multimodalembedding@001",
base_url="http://0.0.0.0:4000",
api_key="sk-1234", # type: ignore
)
query_result = embeddings.embed_query(
"gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"
)
print(query_result)
```
Requests with base64 encoded images
```python
from langchain_openai import OpenAIEmbeddings
embeddings_models = "multimodalembedding@001"
embeddings = OpenAIEmbeddings(
model="multimodalembedding@001",
base_url="http://0.0.0.0:4000",
api_key="sk-1234", # type: ignore
)
query_result = embeddings.embed_query(
"data:image/jpeg;base64,..."
)
print(query_result)
```
</TabItem>
</Tabs>
</TabItem>
<TabItem value="proxy-vtx" label="LiteLLM PROXY (Vertex SDK)">
1. Add model to config.yaml
```yaml
default_vertex_config:
vertex_project: "adroit-crow-413218"
vertex_location: "us-central1"
vertex_credentials: adroit-crow-413218-a956eef1a2a8.json
```
2. Start Proxy
```
$ litellm --config /path/to/config.yaml
```
3. Make Request use OpenAI Python SDK
```python
import vertexai
from vertexai.vision_models import Image, MultiModalEmbeddingModel, Video
from vertexai.vision_models import VideoSegmentConfig
from google.auth.credentials import Credentials
LITELLM_PROXY_API_KEY = "sk-1234"
LITELLM_PROXY_BASE = "http://0.0.0.0:4000/vertex-ai"
import datetime
class CredentialsWrapper(Credentials):
def __init__(self, token=None):
super().__init__()
self.token = token
self.expiry = None # or set to a future date if needed
def refresh(self, request):
pass
def apply(self, headers, token=None):
headers['Authorization'] = f'Bearer {self.token}'
@property
def expired(self):
return False # Always consider the token as non-expired
@property
def valid(self):
return True # Always consider the credentials as valid
credentials = CredentialsWrapper(token=LITELLM_PROXY_API_KEY)
vertexai.init(
project="adroit-crow-413218",
location="us-central1",
api_endpoint=LITELLM_PROXY_BASE,
credentials = credentials,
api_transport="rest",
)
model = MultiModalEmbeddingModel.from_pretrained("multimodalembedding")
image = Image.load_from_file(
"gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"
)
embeddings = model.get_embeddings(
image=image,
contextual_text="Colosseum",
dimension=1408,
)
print(f"Image Embedding: {embeddings.image_embedding}")
print(f"Text Embedding: {embeddings.text_embedding}")
```
</TabItem>
</Tabs>
### Text + Image + Video Embeddings
<Tabs>
<TabItem value="sdk" label="SDK">
Text + Image
```python
response = await litellm.aembedding(
model="vertex_ai/multimodalembedding@001",
input=["hey", "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"] # will be sent as a gcs image
)
```
Text + Video
```python
response = await litellm.aembedding(
model="vertex_ai/multimodalembedding@001",
input=["hey", "gs://my-bucket/embeddings/supermarket-video.mp4"] # will be sent as a gcs image
)
```
Image + Video
```python
response = await litellm.aembedding(
model="vertex_ai/multimodalembedding@001",
input=["gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png", "gs://my-bucket/embeddings/supermarket-video.mp4"] # will be sent as a gcs image
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY (Unified Endpoint)">
1. Add model to config.yaml
```yaml
model_list:
- model_name: multimodalembedding@001
litellm_params:
model: vertex_ai/multimodalembedding@001
vertex_project: "adroit-crow-413218"
vertex_location: "us-central1"
vertex_credentials: adroit-crow-413218-a956eef1a2a8.json
litellm_settings:
drop_params: True
```
2. Start Proxy
```
$ litellm --config /path/to/config.yaml
```
3. Make Request use OpenAI Python SDK, Langchain Python SDK
Text + Image
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
# # request sent to model set on litellm proxy, `litellm --model`
response = client.embeddings.create(
model="multimodalembedding@001",
input = ["hey", "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"],
)
print(response)
```
Text + Video
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
# # request sent to model set on litellm proxy, `litellm --model`
response = client.embeddings.create(
model="multimodalembedding@001",
input = ["hey", "gs://my-bucket/embeddings/supermarket-video.mp4"],
)
print(response)
```
Image + Video
```python
import openai
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
# # request sent to model set on litellm proxy, `litellm --model`
response = client.embeddings.create(
model="multimodalembedding@001",
input = ["gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png", "gs://my-bucket/embeddings/supermarket-video.mp4"],
)
print(response)
```
</TabItem>
</Tabs>
docs/my-website/docs/proxy/config_settings.md
+10 -1
View File
@@ -661,6 +661,7 @@ router_settings:
| LITELLM_LICENSE | License key for LiteLLM usage
| LITELLM_LOCAL_MODEL_COST_MAP | Local configuration for model cost mapping in LiteLLM
| LITELLM_LOG | Enable detailed logging for LiteLLM
| LITELLM_MODEL_COST_MAP_URL | URL for fetching model cost map data. Default is https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json
| LITELLM_LOG_FILE | File path to write LiteLLM logs to. When set, logs will be written to both console and the specified file
| LITELLM_LOGGER_NAME | Name for OTEL logger
| LITELLM_METER_NAME | Name for OTEL Meter
@@ -773,10 +774,15 @@ router_settings:
| REPLICATE_POLLING_DELAY_SECONDS | Delay in seconds for Replicate polling operations. Default is 0.5
| REQUEST_TIMEOUT | Timeout in seconds for requests. Default is 6000
| ROUTER_MAX_FALLBACKS | Maximum number of fallbacks for router. Default is 5
| RUNWAYML_DEFAULT_API_VERSION | Default API version for RunwayML service. Default is "2024-11-06"
| RUNWAYML_POLLING_TIMEOUT | Timeout in seconds for RunwayML image generation polling. Default is 600 (10 minutes)
| SECRET_MANAGER_REFRESH_INTERVAL | Refresh interval in seconds for secret manager. Default is 86400 (24 hours)
| SEPARATE_HEALTH_APP | If set to '1', runs health endpoints on a separate ASGI app and port. Default: '0'.
| SEPARATE_HEALTH_PORT | Port for the separate health endpoints app. Only used if SEPARATE_HEALTH_APP=1. Default: 4001.
| SERVER_ROOT_PATH | Root path for the server application
| SEND_USER_API_KEY_ALIAS | Flag to send user API key alias to Zscaler AI Guard. Default is False
| SEND_USER_API_KEY_TEAM_ID | Flag to send user API key team ID to Zscaler AI Guard. Default is False
| SEND_USER_API_KEY_USER_ID | Flag to send user API key user ID to Zscaler AI Guard. Default is False
| SET_VERBOSE | Flag to enable verbose logging
| SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD | Minimum number of requests to consider "reasonable traffic" for single-deployment cooldown logic. Default is 1000
| SLACK_DAILY_REPORT_FREQUENCY | Frequency of daily Slack reports (e.g., daily, weekly)
@@ -824,4 +830,7 @@ router_settings:
| SPEND_LOG_CLEANUP_BATCH_SIZE | Number of logs deleted per batch during cleanup. Default is 1000
| COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY | Maximum size for CoroutineChecker in-memory cache. Default is 1000
| DEFAULT_SHARED_HEALTH_CHECK_TTL | Time-to-live in seconds for cached health check results in shared health check mode. Default is 300 (5 minutes)
| DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL | Time-to-live in seconds for health check lock in shared health check mode. Default is 60 (1 minute)
| DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL | Time-to-live in seconds for health check lock in shared health check mode. Default is 60 (1 minute)
| ZSCALER_AI_GUARD_API_KEY | API key for Zscaler AI Guard service
| ZSCALER_AI_GUARD_POLICY_ID | Policy ID for Zscaler AI Guard guardrails
| ZSCALER_AI_GUARD_URL | Base URL for Zscaler AI Guard API. Default is https://api.us1.zseclipse.net/v1/detection/execute-policy
docs/my-website/docs/proxy/docker_quick_start.md
+139 -1
View File
@@ -2,7 +2,7 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# E2E Tutorial
# Getting Started Tutorial
End-to-End tutorial for LiteLLM Proxy to:
- Add an Azure OpenAI model
@@ -82,6 +82,8 @@ model_list:
### Model List Specification
You can read more about how model resolution works in the [Model Configuration](#understanding-model-configuration) section.
- **`model_name`** (`str`) - This field should contain the name of the model as received.
- **`litellm_params`** (`dict`) [See All LiteLLM Params](https://github.com/BerriAI/litellm/blob/559a6ad826b5daef41565f54f06c739c8c068b28/litellm/types/router.py#L222)
- **`model`** (`str`) - Specifies the model name to be sent to `litellm.acompletion` / `litellm.aembedding`, etc. This is the identifier used by LiteLLM to route to the correct model + provider logic on the backend.
@@ -89,6 +91,10 @@ model_list:
- **`api_base`** (`str`) - The API base for your azure deployment.
- **`api_version`** (`str`) - The API Version to use when calling Azure's OpenAI API. Get the latest Inference API version [here](https://learn.microsoft.com/en-us/azure/ai-services/openai/api-version-deprecation?source=recommendations#latest-preview-api-releases).
---
---
### Useful Links
- [**All Supported LLM API Providers (OpenAI/Bedrock/Vertex/etc.)**](../providers/)
@@ -407,6 +413,138 @@ curl -X POST 'http://0.0.0.0:4000/chat/completions' \
- [Set Budgets / Rate Limits per key/user/teams](./users.md)
- [Dynamic TPM/RPM Limits for keys](./team_budgets.md#dynamic-tpmrpm-allocation)
## Key Concepts
This section explains key concepts on LiteLLM AI Gateway.
### Understanding Model Configuration
For this config.yaml example:
```yaml
model_list:
- model_name: gpt-4o
litellm_params:
model: azure/my_azure_deployment
api_base: os.environ/AZURE_API_BASE
api_key: "os.environ/AZURE_API_KEY"
api_version: "2025-01-01-preview" # [OPTIONAL] litellm uses the latest azure api_version by default
```
**How Model Resolution Works:**
```
Client Request LiteLLM Proxy Provider API
────────────── ──────────────── ─────────────
POST /chat/completions
{ 1. Looks up model_name
"model": "gpt-4o" ──────────▶ in config.yaml
...
} 2. Finds matching entry:
model_name: gpt-4o
3. Extracts litellm_params:
model: azure/my_azure_deployment
api_base: https://...
api_key: sk-...
4. Routes to provider ──▶ Azure OpenAI API
POST /deployments/my_azure_deployment/...
```
**Breaking Down the `model` Parameter under `litellm_params`:**
```yaml
model_list:
- model_name: gpt-4o # What the client calls
litellm_params:
model: azure/my_azure_deployment # <provider>/<model-name>
───── ───────────────────
│ │
│ └─────▶ Model name sent to the provider API
│
└─────────────────▶ Provider that LiteLLM routes to
```
**Visual Breakdown:**
```
model: azure/my_azure_deployment
└─┬─┘ └─────────┬─────────┘
│ │
│ └────▶ The actual model identifier that gets sent to Azure
│ (e.g., your deployment name, or the model name)
│
└──────────────────▶ Tells LiteLLM which provider to use
(azure, openai, anthropic, bedrock, etc.)
```
**Key Concepts:**
- **`model_name`**: The alias your client uses to call the model. This is what you send in your API requests (e.g., `gpt-4o`).
- **`model` (in litellm_params)**: Format is `<provider>/<model-identifier>`
- **Provider** (before `/`): Routes to the correct LLM provider (e.g., `azure`, `openai`, `anthropic`, `bedrock`)
- **Model identifier** (after `/`): The actual model/deployment name sent to that provider's API
**Advanced Configuration Examples:**
For custom OpenAI-compatible endpoints (e.g., vLLM, Ollama, custom deployments):
```yaml
model_list:
- model_name: my-custom-model
litellm_params:
model: openai/nvidia/llama-3.2-nv-embedqa-1b-v2
api_base: http://my-service.svc.cluster.local:8000/v1
api_key: "sk-1234"
```
**Breaking down complex model paths:**
```
model: openai/nvidia/llama-3.2-nv-embedqa-1b-v2
└─┬──┘ └────────────┬────────────────┘
│ │
│ └────▶ Full model string sent to the provider API
│ (in this case: "nvidia/llama-3.2-nv-embedqa-1b-v2")
│
└──────────────────────▶ Provider (openai = OpenAI-compatible API)
```
The key point: Everything after the first `/` is passed as-is to the provider's API.
**Common Patterns:**
```yaml
model_list:
# Azure deployment
- model_name: gpt-4
litellm_params:
model: azure/gpt-4-deployment
api_base: https://my-azure.openai.azure.com
# OpenAI
- model_name: gpt-4
litellm_params:
model: openai/gpt-4
api_key: os.environ/OPENAI_API_KEY
# Custom OpenAI-compatible endpoint
- model_name: my-llama-model
litellm_params:
model: openai/meta/llama-3-8b
api_base: http://my-vllm-server:8000/v1
api_key: "optional-key"
# Bedrock
- model_name: claude-3
litellm_params:
model: bedrock/anthropic.claude-3-sonnet-20240229-v1:0
aws_region_name: us-east-1
```
## Troubleshooting
docs/my-website/sidebars.js
+1
View File
@@ -488,6 +488,7 @@ const sidebars = {
"providers/vertex_ai/videos",
"providers/vertex_partner",
"providers/vertex_self_deployed",
"providers/vertex_embedding",
"providers/vertex_image",
"providers/vertex_batch",
"providers/vertex_ocr",
litellm/cost_calculator.py
+1
View File
@@ -943,6 +943,7 @@ def completion_cost( # noqa: PLR0915
n=n,
size=size,
optional_params=optional_params,
call_type=call_type,
)
elif (
call_type == CallTypes.create_video.value
litellm/integrations/sqs.py
+17
View File
@@ -30,6 +30,7 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.types.utils import StandardLoggingPayload
from .custom_batch_logger import CustomBatchLogger
from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus
_BASE64_INLINE_PATTERN = re.compile(
r"data:(?:application|image|audio|video)/[a-zA-Z0-9.+-]+;base64,[A-Za-z0-9+/=\s]+",
@@ -354,3 +355,19 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM):
response.raise_for_status()
except Exception as e:
verbose_logger.exception(f"Error sending to SQS: {str(e)}")
async def async_health_check(self) -> IntegrationHealthCheckStatus:
"""
Health check for SQS by sending a small test message to the configured queue.
"""
try:
from litellm.litellm_core_utils.litellm_logging import (
create_dummy_standard_logging_payload,
)
# Create a minimal standard logging payload
standard_logging_object: StandardLoggingPayload = create_dummy_standard_logging_payload()
# Attempt to send a single message
await self.async_send_message(standard_logging_object)
return IntegrationHealthCheckStatus(status="healthy", error_message=None)
except Exception as e:
return IntegrationHealthCheckStatus(status="unhealthy", error_message=str(e))
litellm/litellm_core_utils/exception_mapping_utils.py
+10 -5
View File
@@ -3,6 +3,7 @@ import traceback
from typing import Any, Optional
import httpx
import re
import litellm
from litellm._logging import verbose_logger
@@ -45,13 +46,20 @@ class ExceptionCheckers:
if not isinstance(error_str, str):
return False
if "429" in error_str or "rate limit" in error_str.lower():
# Only treat 429 as a rate limit signal when it appears as a standalone token
if re.search(r"\b429\b", error_str):
return True
_error_str_lower = error_str.lower()
# Match "rate limit" (including variations like rate-limit / rate_limit)
if re.search(r"rate[\s_\-]*limit", _error_str_lower):
return True
#######################################
# Mistral API returns this error string
#########################################
if "service tier capacity exceeded" in error_str.lower():
if "service tier capacity exceeded" in _error_str_lower:
return True
return False
@@ -155,9 +163,6 @@ def _get_response_headers(original_exception: Exception) -> Optional[httpx.Heade
return _response_headers
import re
def extract_and_raise_litellm_exception(
response: Optional[Any],
error_str: str,
litellm/litellm_core_utils/llm_cost_calc/utils.py
+13
View File
@@ -640,6 +640,7 @@ class CostCalculatorUtils:
n: Optional[int] = None,
size: Optional[str] = None,
optional_params: Optional[dict] = None,
call_type: Optional[str] = None,
) -> float:
"""
Route the image generation cost calculator based on the custom_llm_provider
@@ -713,6 +714,18 @@ class CostCalculatorUtils:
image_response=completion_response,
)
elif custom_llm_provider == litellm.LlmProviders.GEMINI.value:
if call_type in (
CallTypes.image_edit.value,
CallTypes.aimage_edit.value,
):
from litellm.llms.gemini.image_edit.cost_calculator import (
cost_calculator as gemini_image_edit_cost_calculator,
)
return gemini_image_edit_cost_calculator(
model=model,
image_response=completion_response,
)
from litellm.llms.gemini.image_generation.cost_calculator import (
cost_calculator as gemini_image_cost_calculator,
)
litellm/litellm_core_utils/token_counter.py
+158 -29
View File
@@ -3,7 +3,17 @@
import base64
import io
import struct
from typing import Callable, List, Literal, Optional, Tuple, Union, cast
from typing import (
Any,
Callable,
List,
Literal,
Mapping,
Optional,
Tuple,
Union,
cast,
)
import tiktoken
@@ -20,6 +30,10 @@ from litellm.constants import (
)
from litellm.litellm_core_utils.default_encoding import encoding as default_encoding
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
from litellm.types.llms.anthropic import (
AnthropicMessagesToolResultParam,
AnthropicMessagesToolUseParam,
)
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionNamedToolChoiceParam,
@@ -552,6 +566,131 @@ def _fix_model_name(model: str) -> str:
return "gpt-3.5-turbo"
def _count_image_tokens(
image_url: Any,
use_default_image_token_count: bool,
) -> int:
"""
Count tokens for an image_url content block.
Args:
image_url: The image URL data - can be a string URL or dict with 'url' and 'detail'
use_default_image_token_count: Whether to use default image token counts
Returns:
int: Number of tokens for the image
Raises:
ValueError: If image_url is invalid type or detail value is invalid
"""
if isinstance(image_url, dict):
detail = image_url.get("detail", "auto")
if detail not in ["low", "high", "auto"]:
raise ValueError(
f"Invalid detail value: {detail}. Expected 'low', 'high', or 'auto'."
)
url = image_url.get("url")
if not url:
raise ValueError("Missing required key 'url' in image_url dict.")
return calculate_img_tokens(
data=url,
mode=detail, # type: ignore
use_default_image_token_count=use_default_image_token_count,
)
elif isinstance(image_url, str):
if not image_url.strip():
raise ValueError("Empty image_url string is not valid.")
return calculate_img_tokens(
data=image_url,
mode="auto",
use_default_image_token_count=use_default_image_token_count,
)
else:
raise ValueError(
f"Invalid image_url type: {type(image_url).__name__}. "
"Expected str or dict with 'url' field."
)
def _validate_anthropic_content(content: Mapping[str, Any]) -> type:
"""
Validate and determine which Anthropic TypedDict applies.
Returns the corresponding TypedDict class if recognized, otherwise raises.
"""
content_type = content.get("type")
if not content_type:
raise ValueError("Anthropic content missing required field: 'type'")
mapping = {
"tool_use": AnthropicMessagesToolUseParam,
"tool_result": AnthropicMessagesToolResultParam,
}
expected_cls = mapping.get(content_type)
if expected_cls is None:
raise ValueError(f"Unknown Anthropic content type: '{content_type}'")
missing = [
k for k in getattr(expected_cls, "__required_keys__", set()) if k not in content
]
if missing:
raise ValueError(
f"Missing required fields in {content_type} block: {', '.join(missing)}"
)
return expected_cls
def _count_anthropic_content(
content: Mapping[str, Any],
count_function: TokenCounterFunction,
use_default_image_token_count: bool,
default_token_count: Optional[int],
) -> int:
"""
Count tokens in Anthropic-specific content blocks (tool_use, tool_result, etc.).
Uses TypedDict definitions from litellm.types.llms.anthropic to determine
what fields to count and how to handle nested structures.
Dynamically infers which fields to count based on the TypedDict definition,
avoiding hardcoded field names.
"""
typeddict_cls = _validate_anthropic_content(content)
type_hints = getattr(typeddict_cls, "__annotations__", {})
tokens = 0
# Fields to skip (metadata/identifiers that don't contribute to prompt tokens)
skip_fields = {"type", "id", "tool_use_id", "cache_control", "is_error"}
# Iterate over all fields defined in the TypedDict
for field_name, field_type in type_hints.items():
if field_name in skip_fields:
continue
field_value = content.get(field_name)
if field_value is None:
continue
try:
if isinstance(field_value, str):
tokens += count_function(field_value)
elif isinstance(field_value, list):
tokens += _count_content_list(
count_function,
field_value, # type: ignore
use_default_image_token_count,
default_token_count,
)
elif isinstance(field_value, dict):
tokens += count_function(str(field_value))
except Exception as e:
if default_token_count is not None:
return default_token_count
raise ValueError(f"Error counting field '{field_name}': {e}")
return tokens
def _count_content_list(
count_function: TokenCounterFunction,
content_list: OpenAIMessageContent,
@@ -559,7 +698,7 @@ def _count_content_list(
default_token_count: Optional[int],
) -> int:
"""
Get the number of tokens from a list of content.
Recursively count tokens from a list of content blocks.
"""
try:
num_tokens = 0
@@ -567,42 +706,32 @@ def _count_content_list(
if isinstance(c, str):
num_tokens += count_function(c)
elif c["type"] == "text":
num_tokens += count_function(c["text"])
num_tokens += count_function(c.get("text", ""))
elif c["type"] == "image_url":
if isinstance(c["image_url"], dict):
image_url_dict = c["image_url"]
detail = image_url_dict.get("detail", "auto")
if detail not in ["low", "high", "auto"]:
raise ValueError(
f"Invalid detail value: {detail}. Expected 'low', 'high', or 'auto'."
)
url = image_url_dict.get("url")
num_tokens += calculate_img_tokens(
data=url,
mode=detail, # type: ignore
use_default_image_token_count=use_default_image_token_count,
)
elif isinstance(c["image_url"], str):
image_url_str = c["image_url"]
num_tokens += calculate_img_tokens(
data=image_url_str,
mode="auto",
use_default_image_token_count=use_default_image_token_count,
)
else:
raise ValueError(
f"Invalid image_url type: {type(c['image_url'])}. Expected str or dict."
)
image_url = c.get("image_url")
num_tokens += _count_image_tokens(
image_url, use_default_image_token_count
)
elif c["type"] in ("tool_use", "tool_result"):
num_tokens += _count_anthropic_content(
c,
count_function,
use_default_image_token_count,
default_token_count,
)
else:
raise ValueError(
f"Invalid content type: {type(c)}. Expected str or dict."
f"Invalid content item type: {type(c).__name__}. "
f"Expected str or dict with 'type' field. "
f"Value: {c!r}"
)
return num_tokens
except Exception as e:
if default_token_count is not None:
return default_token_count
raise ValueError(
f"Error getting number of tokens from content list: {e}, default_token_count={default_token_count}"
f"Error getting number of tokens from content list: {e}, "
f"default_token_count={default_token_count}"
)
litellm/llms/bedrock/chat/agentcore/transformation.py
+107 -73
View File
@@ -79,25 +79,25 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
aws_bedrock_runtime_endpoint = optional_params.get(
"aws_bedrock_runtime_endpoint", None
)
# Extract ARN from model string
agent_runtime_arn = self._get_agent_runtime_arn(model)
# Parse ARN to get region
region = self._extract_region_from_arn(agent_runtime_arn)
# Build the base endpoint URL for AgentCore
# Note: We don't use get_runtime_endpoint as AgentCore has its own endpoint structure
if aws_bedrock_runtime_endpoint:
base_url = aws_bedrock_runtime_endpoint
else:
base_url = f"https://bedrock-agentcore.{region}.amazonaws.com"
# Based on boto3 client.invoke_agent_runtime, the path is:
# /runtimes/{URL-ENCODED-ARN}/invocations?qualifier=<value>
encoded_arn = quote(agent_runtime_arn, safe='')
encoded_arn = quote(agent_runtime_arn, safe="")
endpoint_url = f"{base_url}/runtimes/{encoded_arn}/invocations"
# Add qualifier as query parameter if provided
if "qualifier" in optional_params:
endpoint_url = f"{endpoint_url}?qualifier={optional_params['qualifier']}"
@@ -115,6 +115,19 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
stream: Optional[bool] = None,
fake_stream: Optional[bool] = None,
) -> Tuple[dict, Optional[bytes]]:
# Check if api_key (bearer token) is provided for Cognito authentication
jwt_token = optional_params.get("api_key")
if jwt_token:
verbose_logger.debug(
f"AgentCore: Using Bearer token authentication (Cognito/JWT) - token: {jwt_token[:50]}..."
)
headers["Content-Type"] = "application/json"
headers["Authorization"] = f"Bearer {jwt_token}"
# Return headers with bearer token and JSON-encoded body (not SigV4 signed)
return headers, json.dumps(request_data).encode()
# Otherwise, use AWS SigV4 authentication
verbose_logger.debug("AgentCore: Using AWS SigV4 authentication (IAM)")
return self._sign_request(
service_name="bedrock-agentcore",
headers=headers,
@@ -157,16 +170,22 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
"""
session_id = optional_params.get("runtimeSessionId", None)
if session_id:
verbose_logger.debug(f"Using provided runtimeSessionId: {session_id}")
return session_id
# Generate a session ID with 33+ characters
return f"litellm-session-{str(uuid.uuid4())}"
generated_id = f"litellm-session-{str(uuid.uuid4())}"
verbose_logger.debug(f"Generated new session ID: {generated_id}")
return generated_id
def _get_runtime_user_id(self, optional_params: dict) -> Optional[str]:
"""
Get runtime user ID if provided
"""
return optional_params.get("runtimeUserId", None)
user_id = optional_params.get("runtimeUserId", None)
if user_id:
verbose_logger.debug(f"Using provided runtimeUserId: {user_id}")
return user_id
def transform_request(
self,
@@ -188,6 +207,10 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
Returns:
dict: Payload dict containing the prompt
"""
verbose_logger.debug(
f"AgentCore transform_request - optional_params keys: {list(optional_params.keys())}"
)
# Use the last message content as the prompt
prompt = convert_content_list_to_str(messages[-1])
@@ -206,17 +229,18 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
# The request data is the payload dict (will be JSON encoded by the HTTP handler)
# Qualifier will be handled as a query parameter in get_complete_url
verbose_logger.debug(f"PAYLOAD: {payload}")
return payload
def _extract_sse_json(self, line: str) -> Optional[Dict]:
"""Extract and parse JSON from an SSE data line."""
if not line.startswith('data:'):
if not line.startswith("data:"):
return None
json_str = line[5:].strip()
if not json_str:
return None
try:
data = json.loads(json_str)
# Skip non-dict data (some lines contain JSON strings)
@@ -230,11 +254,11 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
event_payload = event_data.get("event")
if not event_payload:
return None
metadata = event_payload.get("metadata")
if metadata and "usage" in metadata:
return metadata["usage"] # type: ignore
return None
def _extract_content_delta(self, event_data: Dict) -> Optional[str]:
@@ -242,11 +266,11 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
event_payload = event_data.get("event")
if not event_payload:
return None
content_block_delta = event_payload.get("contentBlockDelta")
if not content_block_delta:
return None
delta = content_block_delta.get("delta", {})
return delta.get("text")
@@ -258,7 +282,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
content_list = message.get("content", [])
if not isinstance(content_list, list):
return ""
return "".join(
block["text"]
for block in content_list
@@ -270,31 +294,28 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
) -> Optional[Usage]:
"""
Calculate token usage using LiteLLM's token counter.
Args:
model: The model name
messages: Input messages
content: Response content
Returns:
Usage object with calculated tokens, or None if calculation fails
"""
try:
from litellm.utils import token_counter
prompt_tokens = token_counter(model=model, messages=messages)
completion_tokens = token_counter(
model=model,
text=content,
count_response_tokens=True
model=model, text=content, count_response_tokens=True
)
total_tokens = prompt_tokens + completion_tokens
verbose_logger.debug(
f"Calculated usage - prompt: {prompt_tokens}, "
f"completion: {completion_tokens}, total: {total_tokens}"
f"Calculated usage - prompt: {prompt_tokens}, completion: {completion_tokens}, total: {total_tokens}"
)
return Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
@@ -307,7 +328,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
def _parse_json_response(self, response_json: dict) -> AgentCoreParsedResponse:
"""
Parse direct JSON response (non-streaming).
JSON response structure:
{
"result": {
@@ -317,15 +338,15 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
}
"""
result = response_json.get("result", {})
# Extract content using the same helper as SSE parsing
content = self._extract_content_from_message(result) # type: ignore
# JSON responses don't include usage data
return AgentCoreParsedResponse(
content=content,
usage=None,
final_message=result # type: ignore
final_message=result, # type: ignore
)
def _get_parsed_response(
@@ -333,16 +354,16 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
) -> AgentCoreParsedResponse:
"""
Parse AgentCore response based on content type.
Args:
raw_response: Raw HTTP response from AgentCore
Returns:
AgentCoreParsedResponse: Parsed response data
"""
content_type = raw_response.headers.get("content-type", "").lower()
verbose_logger.debug(f"AgentCore response Content-Type: {content_type}")
# Parse response based on content type
if "application/json" in content_type:
# Direct JSON response
@@ -354,64 +375,66 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
# SSE stream response (text/event-stream or default)
verbose_logger.debug("Parsing SSE stream response")
response_text = raw_response.text
verbose_logger.debug(f"AgentCore response (first 500 chars): {response_text[:500]}")
verbose_logger.debug(
f"AgentCore response (first 500 chars): {response_text[:500]}"
)
return self._parse_sse_stream(response_text)
def _parse_sse_stream(self, response_text: str) -> AgentCoreParsedResponse:
"""
Parse Server-Sent Events (SSE) stream format.
Each line starts with 'data:' followed by JSON.
Returns:
AgentCoreParsedResponse: Parsed response with content, usage, and message
"""
final_message: Optional[AgentCoreMessage] = None
usage_data: Optional[AgentCoreUsage] = None
content_blocks: List[str] = []
for line in response_text.strip().split('\n'):
for line in response_text.strip().split("\n"):
line = line.strip()
if not line:
continue
data = self._extract_sse_json(line)
if not data:
continue
verbose_logger.debug(f"SSE event keys: {list(data.keys())}")
# Check for final complete message
if "message" in data and isinstance(data["message"], dict):
final_message = data["message"] # type: ignore
verbose_logger.debug("Found final message")
# Process event data
if "event" in data and isinstance(data["event"], dict):
event_payload = data["event"]
verbose_logger.debug(f"Event payload keys: {list(event_payload.keys())}")
verbose_logger.debug(
f"Event payload keys: {list(event_payload.keys())}"
)
# Extract usage metadata
if usage := self._extract_usage_from_event(data):
usage_data = usage
verbose_logger.debug(f"Found usage data: {usage_data}")
# Collect content deltas
if text := self._extract_content_delta(data):
content_blocks.append(text)
# Build final content
content = (
self._extract_content_from_message(final_message)
if final_message
else "".join(content_blocks)
)
verbose_logger.debug(f"Final usage_data: {usage_data}")
return AgentCoreParsedResponse(
content=content,
usage=usage_data,
final_message=final_message
content=content, usage=usage_data, final_message=final_message
)
def get_streaming_response(
@@ -421,11 +444,11 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
) -> AgentCoreSSEStreamIterator:
"""
Return a streaming iterator for SSE responses.
Args:
model: The model name
raw_response: Raw HTTP response with streaming data
Returns:
AgentCoreSSEStreamIterator: Iterator that yields ModelResponse chunks
"""
@@ -446,7 +469,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
) -> CustomStreamWrapper:
"""
Get a CustomStreamWrapper for synchronous streaming.
This is called when stream=True is passed to completion().
"""
from litellm.llms.custom_httpx.http_handler import (
@@ -454,10 +477,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
_get_httpx_client,
)
from litellm.utils import CustomStreamWrapper
if client is None or not isinstance(client, HTTPHandler):
client = _get_httpx_client(params={})
verbose_logger.debug(f"Making sync streaming request to: {api_base}")
# Make streaming request
response = client.post(
api_base,
@@ -466,22 +491,24 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
stream=True, # THIS IS KEY - tells httpx to not buffer
logging_obj=logging_obj,
)
if response.status_code != 200:
raise BedrockError(
status_code=response.status_code, message=str(response.read())
)
# Create iterator for SSE stream
completion_stream = self.get_streaming_response(model=model, raw_response=response)
completion_stream = self.get_streaming_response(
model=model, raw_response=response
)
streaming_response = CustomStreamWrapper(
completion_stream=completion_stream,
model=model,
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
)
# LOGGING
logging_obj.post_call(
input=messages,
@@ -489,7 +516,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
original_response="first stream response received",
additional_args={"complete_input_dict": data},
)
return streaming_response
async def get_async_custom_stream_wrapper(
@@ -517,7 +544,11 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
from litellm.utils import CustomStreamWrapper
if client is None or not isinstance(client, AsyncHTTPHandler):
client = get_async_httpx_client(llm_provider=cast(Any, "bedrock"), params={})
client = get_async_httpx_client(
llm_provider=cast(Any, "bedrock"), params={}
)
verbose_logger.debug(f"Making async streaming request to: {api_base}")
# Make async streaming request
response = await client.post(
@@ -534,7 +565,9 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
)
# Create iterator for SSE stream
completion_stream = self.get_streaming_response(model=model, raw_response=response)
completion_stream = self.get_streaming_response(
model=model, raw_response=response
)
streaming_response = CustomStreamWrapper(
completion_stream=completion_stream,
@@ -583,29 +616,29 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
"""
Transform the AgentCore response to LiteLLM ModelResponse format.
AgentCore can return either JSON or SSE (Server-Sent Events) stream responses.
Note: For streaming responses, use get_streaming_response() instead.
"""
try:
# Parse the response based on content type (JSON or SSE)
parsed_data = self._get_parsed_response(raw_response)
content = parsed_data["content"]
usage_data = parsed_data["usage"]
verbose_logger.debug(f"Parsed content length: {len(content)}")
verbose_logger.debug(f"Usage data: {usage_data}")
# Create the message
message = Message(content=content, role="assistant")
# Create choices
choice = Choices(finish_reason="stop", index=0, message=message)
# Update model response
model_response.choices = [choice]
model_response.model = model
# Add usage information if available
# Note: AgentCore JSON responses don't include usage data
# SSE responses may include usage in metadata events
@@ -618,11 +651,13 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
setattr(model_response, "usage", usage)
else:
# Calculate token usage using LiteLLM's token counter
verbose_logger.debug("No usage data from AgentCore - calculating tokens")
verbose_logger.debug(
"No usage data from AgentCore - calculating tokens"
)
calculated_usage = self._calculate_usage(model, messages, content)
if calculated_usage:
setattr(model_response, "usage", calculated_usage)
return model_response
except Exception as e:
@@ -658,4 +693,3 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
custom_llm_provider: Optional[str] = None,
) -> bool:
return True
litellm/llms/gemini/image_edit/__init__.py
+11
View File
@@ -0,0 +1,11 @@
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from .transformation import GeminiImageEditConfig
from .cost_calculator import cost_calculator
__all__ = ["GeminiImageEditConfig", "get_gemini_image_edit_config", "cost_calculator"]
def get_gemini_image_edit_config(model: str) -> BaseImageEditConfig:
return GeminiImageEditConfig()
litellm/llms/gemini/image_edit/cost_calculator.py
+35
View File
@@ -0,0 +1,35 @@
"""
Gemini Image Edit Cost Calculator
"""
from typing import Any
import litellm
from litellm.types.utils import ImageResponse
def cost_calculator(
model: str,
image_response: Any,
) -> float:
"""
Gemini image edit cost calculator.
Mirrors image generation pricing: charge per returned image based on
model metadata (`output_cost_per_image`).
"""
model_info = litellm.get_model_info(
model=model,
custom_llm_provider="gemini",
)
output_cost_per_image: float = model_info.get("output_cost_per_image") or 0.0
if not isinstance(image_response, ImageResponse):
raise ValueError(
f"image_response must be of type ImageResponse got type={type(image_response)}"
)
num_images = len(image_response.data or [])
return output_cost_per_image * num_images
litellm/llms/gemini/image_edit/transformation.py
+197
View File
@@ -0,0 +1,197 @@
import base64
from io import BufferedReader, BytesIO
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
import httpx
from httpx._types import RequestFiles
from litellm.images.utils import ImageEditRequestUtils
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, OpenAIImage
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class GeminiImageEditConfig(BaseImageEditConfig):
DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta"
SUPPORTED_PARAMS: List[str] = ["size"]
def get_supported_openai_params(self, model: str) -> List[str]:
return list(self.SUPPORTED_PARAMS)
def map_openai_params(
self,
image_edit_optional_params: ImageEditOptionalRequestParams,
model: str,
drop_params: bool,
) -> Dict[str, Any]:
supported_params = self.get_supported_openai_params(model)
filtered_params = {
key: value
for key, value in image_edit_optional_params.items()
if key in supported_params
}
mapped_params: Dict[str, Any] = {}
if "size" in filtered_params:
mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(
filtered_params["size"] # type: ignore[arg-type]
)
return mapped_params
def validate_environment(
self,
headers: dict,
model: str,
api_key: Optional[str] = None,
) -> dict:
final_api_key: Optional[str] = api_key or get_secret_str("GEMINI_API_KEY")
if not final_api_key:
raise ValueError("GEMINI_API_KEY is not set")
headers["x-goog-api-key"] = final_api_key
headers["Content-Type"] = "application/json"
return headers
def get_complete_url(
self,
model: str,
api_base: Optional[str],
litellm_params: dict,
) -> str:
base_url = api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL
base_url = base_url.rstrip("/")
return f"{base_url}/models/{model}:generateContent"
def transform_image_edit_request( # type: ignore[override]
self,
model: str,
prompt: str,
image: FileTypes,
image_edit_optional_request_params: Dict[str, Any],
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> Tuple[Dict[str, Any], Optional[RequestFiles]]:
inline_parts = self._prepare_inline_image_parts(image)
if not inline_parts:
raise ValueError("Gemini image edit requires at least one image.")
contents = [
{
"parts": inline_parts + [{"text": prompt}],
}
]
request_body: Dict[str, Any] = {"contents": contents}
generation_config: Dict[str, Any] = {}
if "aspectRatio" in image_edit_optional_request_params:
generation_config["aspectRatio"] = image_edit_optional_request_params[
"aspectRatio"
]
if generation_config:
request_body["generationConfig"] = generation_config
empty_files = cast(RequestFiles, [])
return request_body, empty_files
def transform_image_edit_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: Any,
) -> ImageResponse:
model_response = ImageResponse()
try:
response_json = raw_response.json()
except Exception as exc:
raise self.get_error_class(
error_message=f"Error transforming image edit response: {exc}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
candidates = response_json.get("candidates", [])
data_list: List[ImageObject] = []
for candidate in candidates:
content = candidate.get("content", {})
parts = content.get("parts", [])
for part in parts:
inline_data = part.get("inlineData")
if inline_data and inline_data.get("data"):
data_list.append(
ImageObject(
b64_json=inline_data["data"],
url=None,
)
)
model_response.data = cast(List[OpenAIImage], data_list)
return model_response
def _map_size_to_aspect_ratio(self, size: str) -> str:
aspect_ratio_map = {
"1024x1024": "1:1",
"1792x1024": "16:9",
"1024x1792": "9:16",
"1280x896": "4:3",
"896x1280": "3:4",
}
return aspect_ratio_map.get(size, "1:1")
def _prepare_inline_image_parts(
self, image: Union[FileTypes, List[FileTypes]]
) -> List[Dict[str, Any]]:
images: List[FileTypes]
if isinstance(image, list):
images = image
else:
images = [image]
inline_parts: List[Dict[str, Any]] = []
for img in images:
if img is None:
continue
mime_type = ImageEditRequestUtils.get_image_content_type(img)
image_bytes = self._read_all_bytes(img)
inline_parts.append(
{
"inlineData": {
"mimeType": mime_type,
"data": base64.b64encode(image_bytes).decode("utf-8"),
}
}
)
return inline_parts
def _read_all_bytes(self, image: FileTypes) -> bytes:
if isinstance(image, bytes):
return image
if isinstance(image, BytesIO):
current_pos = image.tell()
image.seek(0)
data = image.read()
image.seek(current_pos)
return data
if isinstance(image, BufferedReader):
current_pos = image.tell()
image.seek(0)
data = image.read()
image.seek(current_pos)
return data
raise ValueError("Unsupported image type for Gemini image edit.")
litellm/llms/gemini/image_generation/transformation.py
+11 -6
View File
@@ -21,6 +21,11 @@ else:
LiteLLMLoggingObj = Any
FLASH_IMAGE_PREVIEW_MODEL_IDENTIFIERS = (
"2.0-flash-preview-image",
"2.0-flash-preview-image-generation",
"2.5-flash-image-preview",
)
class GoogleImageGenConfig(BaseImageGenerationConfig):
DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta"
@@ -97,8 +102,8 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
complete_url = complete_url.rstrip("/")
# Gemini 2.5 Flash Image Preview uses generateContent endpoint
if "2.5-flash-image-preview" in model:
# Gemini Flash Image Preview models use generateContent endpoint
if any(identifier in model for identifier in FLASH_IMAGE_PREVIEW_MODEL_IDENTIFIERS):
complete_url = f"{complete_url}/models/{model}:generateContent"
else:
# All other Imagen models use predict endpoint
@@ -152,8 +157,8 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
}
}
"""
# For Gemini 2.5 Flash Image Preview, use standard Gemini format
if "2.5-flash-image-preview" in model:
# For Gemini Flash Image Preview models, use standard Gemini format
if any(identifier in model for identifier in FLASH_IMAGE_PREVIEW_MODEL_IDENTIFIERS):
request_body: dict = {
"contents": [
{
@@ -212,8 +217,8 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
model_response.data = []
# Handle different response formats based on model
if "2.5-flash-image-preview" in model:
# Gemini 2.5 Flash Image Preview returns in candidates format
if any(identifier in model for identifier in FLASH_IMAGE_PREVIEW_MODEL_IDENTIFIERS):
# Gemini Flash Image Preview models return in candidates format
candidates = response_data.get("candidates", [])
for candidate in candidates:
content = candidate.get("content", {})
litellm/llms/runwayml/text_to_speech/__init__.py
+5
View File
@@ -0,0 +1,5 @@
"""RunwayML Text-to-Speech implementation."""
from .transformation import RunwayMLTextToSpeechConfig
__all__ = ["RunwayMLTextToSpeechConfig"]
litellm/llms/runwayml/text_to_speech/transformation.py
+591
View File
@@ -0,0 +1,591 @@
"""
RunwayML Text-to-Speech transformation
Maps OpenAI TTS spec to RunwayML Text-to-Speech API
"""
import asyncio
import time
from typing import TYPE_CHECKING, Any, Coroutine, Dict, Optional, Tuple, Union
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.constants import (
RUNWAYML_DEFAULT_API_VERSION,
RUNWAYML_POLLING_TIMEOUT,
)
from litellm.llms.base_llm.text_to_speech.transformation import (
BaseTextToSpeechConfig,
TextToSpeechRequestData,
)
from litellm.secret_managers.main import get_secret_str
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.llms.openai import HttpxBinaryResponseContent
else:
LiteLLMLoggingObj = Any
HttpxBinaryResponseContent = Any
class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig):
"""
Configuration for RunwayML Text-to-Speech
Reference: https://api.dev.runwayml.com/v1/text_to_speech
"""
DEFAULT_BASE_URL: str = "https://api.dev.runwayml.com"
TTS_ENDPOINT_PATH: str = "v1/text_to_speech"
DEFAULT_MODEL: str = "eleven_multilingual_v2"
DEFAULT_VOICE_TYPE: str = "runway-preset"
DEFAULT_VOICE_PRESET_ID: str = "Bernard"
# Voice mappings from OpenAI voices to RunwayML preset IDs
# OpenAI voices mapped to similar-sounding RunwayML voices
VOICE_MAPPINGS = {
"alloy": "Maya", # Neutral, balanced female voice
"echo": "James", # Male voice
"fable": "Bernard", # Warm, storytelling voice
"onyx": "Vincent", # Deep male voice
"nova": "Serene", # Warm, expressive female voice
"shimmer": "Ella", # Clear, friendly female voice
}
def dispatch_text_to_speech(
self,
model: str,
input: str,
voice: Optional[Union[str, Dict]],
optional_params: Dict,
litellm_params_dict: Dict,
logging_obj: "LiteLLMLoggingObj",
timeout: Union[float, httpx.Timeout],
extra_headers: Optional[Dict[str, Any]],
base_llm_http_handler: Any,
aspeech: bool,
api_base: Optional[str],
api_key: Optional[str],
**kwargs: Any,
) -> Union[
"HttpxBinaryResponseContent",
Coroutine[Any, Any, "HttpxBinaryResponseContent"],
]:
"""
Dispatch method to handle RunwayML TTS requests
This method encapsulates RunwayML-specific credential resolution and parameter handling
Args:
base_llm_http_handler: The BaseLLMHTTPHandler instance from main.py
"""
# Resolve api_base from multiple sources
api_base = (
api_base
or litellm_params_dict.get("api_base")
or litellm.api_base
or get_secret_str("RUNWAYML_API_BASE")
or self.DEFAULT_BASE_URL
)
# Resolve api_key from multiple sources
api_key = (
api_key
or litellm_params_dict.get("api_key")
or litellm.api_key
or get_secret_str("RUNWAYML_API_SECRET")
or get_secret_str("RUNWAYML_API_KEY")
)
# Convert voice to appropriate format
voice_param: Optional[Union[str, Dict]] = voice
if isinstance(voice, str):
# Keep as string, will be processed in map_openai_params
voice_param = voice
elif isinstance(voice, dict):
# Already in dict format, pass through
voice_param = voice
litellm_params_dict.update({
"api_key": api_key,
"api_base": api_base,
})
# Call the text_to_speech_handler
response = base_llm_http_handler.text_to_speech_handler(
model=model,
input=input,
voice=voice_param,
text_to_speech_provider_config=self,
text_to_speech_optional_params=optional_params,
custom_llm_provider="runwayml",
litellm_params=litellm_params_dict,
logging_obj=logging_obj,
timeout=timeout,
extra_headers=extra_headers,
client=None,
_is_async=aspeech,
)
return response
def get_supported_openai_params(self, model: str) -> list:
"""
RunwayML TTS supports these OpenAI parameters
"""
return ["voice"]
def map_openai_params(
self,
model: str,
optional_params: Dict,
voice: Optional[Union[str, Dict]] = None,
drop_params: bool = False,
kwargs: Dict = {},
) -> Tuple[Optional[str], Dict]:
"""
Map OpenAI parameters to RunwayML TTS parameters
Returns:
Tuple of (mapped_voice_string, mapped_params)
Note: Since RunwayML requires voice as a dict, we store it in
mapped_params["runwayml_voice"] and return None for the voice string.
"""
mapped_params = {}
# Map voice parameter to RunwayML format dict
voice_dict: Optional[Dict] = None
if isinstance(voice, str):
# Check if it's an OpenAI voice name that needs mapping
if voice in self.VOICE_MAPPINGS:
preset_id = self.VOICE_MAPPINGS[voice]
voice_dict = {
"type": self.DEFAULT_VOICE_TYPE,
"presetId": preset_id,
}
else:
# Assume it's a RunwayML preset ID
voice_dict = {
"type": self.DEFAULT_VOICE_TYPE,
"presetId": voice,
}
elif isinstance(voice, dict):
# Already in RunwayML format, use as-is
voice_dict = voice
# Store the voice dict in optional_params for later use
if voice_dict is not None:
mapped_params["runwayml_voice"] = voice_dict
# No other OpenAI params are currently supported by RunwayML TTS
# (response_format, speed, etc. are not supported)
# Return None for voice string since RunwayML uses dict format
return None, mapped_params
def validate_environment(
self,
headers: dict,
model: str,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
"""
Validate RunwayML environment and set up authentication headers
"""
validated_headers = headers.copy()
final_api_key = (
api_key
or get_secret_str("RUNWAYML_API_SECRET")
or get_secret_str("RUNWAYML_API_KEY")
)
if not final_api_key:
raise ValueError("RUNWAYML_API_SECRET or RUNWAYML_API_KEY is not set")
validated_headers["Authorization"] = f"Bearer {final_api_key}"
validated_headers["X-Runway-Version"] = RUNWAYML_DEFAULT_API_VERSION
validated_headers["Content-Type"] = "application/json"
return validated_headers
def get_complete_url(
self,
model: str,
api_base: Optional[str],
litellm_params: dict,
) -> str:
"""
Get the complete URL for RunwayML TTS request
"""
complete_url = (
api_base
or get_secret_str("RUNWAYML_API_BASE")
or self.DEFAULT_BASE_URL
)
complete_url = complete_url.rstrip("/")
return f"{complete_url}/{self.TTS_ENDPOINT_PATH}"
@staticmethod
def _check_timeout(start_time: float, timeout_secs: float) -> None:
"""
Check if operation has timed out.
Args:
start_time: Start time of the operation
timeout_secs: Timeout duration in seconds
Raises:
TimeoutError: If operation has exceeded timeout
"""
if time.time() - start_time > timeout_secs:
raise TimeoutError(
f"RunwayML TTS task polling timed out after {timeout_secs} seconds"
)
@staticmethod
def _check_task_status(response_data: Dict[str, Any]) -> str:
"""
Check RunwayML task status from response.
RunwayML statuses: PENDING, RUNNING, SUCCEEDED, FAILED, CANCELLED, THROTTLED
Args:
response_data: JSON response from RunwayML task endpoint
Returns:
Normalized status string: "running", "succeeded", or raises on failure
Raises:
ValueError: If task failed or status is unknown
"""
status = response_data.get("status", "").upper()
verbose_logger.debug(f"RunwayML TTS task status: {status}")
if status == "SUCCEEDED":
return "succeeded"
elif status == "FAILED":
failure_reason = response_data.get("failure", "Unknown error")
failure_code = response_data.get("failureCode", "unknown")
raise ValueError(
f"RunwayML TTS failed: {failure_reason} (code: {failure_code})"
)
elif status == "CANCELLED":
raise ValueError("RunwayML TTS was cancelled")
elif status in ["PENDING", "RUNNING", "THROTTLED"]:
return "running"
else:
raise ValueError(f"Unknown RunwayML task status: {status}")
def _poll_task_sync(
self,
task_id: str,
api_base: str,
headers: Dict[str, str],
timeout_secs: float = 600,
) -> httpx.Response:
"""
Poll RunwayML task until completion (sync).
RunwayML POST returns immediately with a task that has status PENDING/RUNNING.
We need to poll GET /v1/tasks/{task_id} until status is SUCCEEDED or FAILED.
Args:
task_id: The task ID to poll
api_base: Base URL for RunwayML API
headers: Request headers (including auth)
timeout_secs: Total timeout in seconds (default: 600s = 10 minutes)
Returns:
Final response with completed task
"""
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
client = _get_httpx_client()
start_time = time.time()
# Build task status URL
api_base = api_base.rstrip("/")
task_url = f"{api_base}/v1/tasks/{task_id}"
verbose_logger.debug(f"Polling RunwayML TTS task: {task_url}")
while True:
self._check_timeout(start_time=start_time, timeout_secs=timeout_secs)
# Poll the task status
response = client.get(url=task_url, headers=headers)
response.raise_for_status()
response_data = response.json()
# Check task status
status = self._check_task_status(response_data=response_data)
if status == "succeeded":
return response
elif status == "running":
# Wait before polling again (RunwayML recommends 1-2 second intervals)
time.sleep(2)
async def _poll_task_async(
self,
task_id: str,
api_base: str,
headers: Dict[str, str],
timeout_secs: float = 600,
) -> httpx.Response:
"""
Poll RunwayML task until completion (async).
Args:
task_id: The task ID to poll
api_base: Base URL for RunwayML API
headers: Request headers (including auth)
timeout_secs: Total timeout in seconds (default: 600s = 10 minutes)
Returns:
Final response with completed task
"""
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
client = get_async_httpx_client(llm_provider=litellm.LlmProviders.RUNWAYML)
start_time = time.time()
# Build task status URL
api_base = api_base.rstrip("/")
task_url = f"{api_base}/v1/tasks/{task_id}"
verbose_logger.debug(f"Polling RunwayML TTS task (async): {task_url}")
while True:
self._check_timeout(start_time=start_time, timeout_secs=timeout_secs)
# Poll the task status
response = await client.get(url=task_url, headers=headers)
response.raise_for_status()
response_data = response.json()
# Check task status
status = self._check_task_status(response_data=response_data)
if status == "succeeded":
return response
elif status == "running":
# Wait before polling again (RunwayML recommends 1-2 second intervals)
await asyncio.sleep(2)
def transform_text_to_speech_request(
self,
model: str,
input: str,
voice: Optional[Union[str, Dict]],
optional_params: Dict,
litellm_params: Dict,
headers: dict,
) -> TextToSpeechRequestData:
"""
Transform OpenAI TTS request to RunwayML TTS format
RunwayML expects:
- model: The model to use (e.g., 'eleven_multilingual_v2')
- promptText: The text to convert to speech
- voice: Voice configuration object
{
"type": "runway-preset",
"presetId": "Bernard"
}
Returns:
TextToSpeechRequestData: Contains JSON body and headers
"""
# Get voice from optional_params (mapped in map_openai_params)
runwayml_voice = optional_params.get("runwayml_voice")
if runwayml_voice is None:
# Use default voice if not provided
runwayml_voice = {
"type": self.DEFAULT_VOICE_TYPE,
"presetId": self.DEFAULT_VOICE_PRESET_ID,
}
# Build request body
request_body = {
"model": model or self.DEFAULT_MODEL,
"promptText": input,
"voice": runwayml_voice,
}
# Add any other optional parameters (except runwayml_voice which we already used)
for k, v in optional_params.items():
if k not in request_body and k != "runwayml_voice":
request_body[k] = v
return {
"dict_body": request_body,
"headers": headers,
}
def transform_text_to_speech_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
) -> "HttpxBinaryResponseContent":
"""
Transform RunwayML TTS response to standard format
RunwayML returns a task immediately with status PENDING/RUNNING.
We need to poll the task until it completes, then download the audio.
Initial response:
{
"id": "task_123...",
"status": "PENDING" | "RUNNING",
"createdAt": "2025-11-13T..."
}
After polling:
{
"id": "task_123...",
"status": "SUCCEEDED",
"output": ["https://storage.googleapis.com/.../audio.mp3"],
"completedAt": "2025-11-13T..."
}
"""
from litellm.types.llms.openai import HttpxBinaryResponseContent
try:
response_data = raw_response.json()
except Exception as e:
raise self.get_error_class(
error_message=f"Error parsing RunwayML TTS response: {e}",
status_code=raw_response.status_code,
headers=dict(raw_response.headers),
)
verbose_logger.debug("RunwayML TTS starting polling...")
# Get task ID
task_id = response_data.get("id")
if not task_id:
raise ValueError("RunwayML TTS response missing task ID")
# Get headers for polling (need auth)
poll_headers = {
"Authorization": raw_response.request.headers.get("Authorization", ""),
"X-Runway-Version": raw_response.request.headers.get(
"X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION
),
}
# Poll until task completes
polled_response = self._poll_task_sync(
task_id=task_id,
api_base=self.DEFAULT_BASE_URL,
headers=poll_headers,
timeout_secs=RUNWAYML_POLLING_TIMEOUT,
)
# Get the completed task data
task_data = polled_response.json()
verbose_logger.debug("RunwayML TTS polling complete, downloading audio")
# Get audio URL from output
output = task_data.get("output", [])
if not output or not isinstance(output, list) or len(output) == 0:
raise ValueError("RunwayML TTS response missing audio URL in output")
audio_url = output[0]
if not isinstance(audio_url, str):
raise ValueError(f"RunwayML TTS audio URL is not a string: {audio_url}")
# Download the audio file
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
client = _get_httpx_client()
audio_response = client.get(url=audio_url)
audio_response.raise_for_status()
verbose_logger.debug("RunwayML TTS audio downloaded successfully")
# Return the audio data wrapped in HttpxBinaryResponseContent
return HttpxBinaryResponseContent(audio_response)
async def async_transform_text_to_speech_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
) -> "HttpxBinaryResponseContent":
"""
Async transform RunwayML TTS response to standard format
Same as sync version but uses async polling and download
"""
from litellm.types.llms.openai import HttpxBinaryResponseContent
try:
response_data = raw_response.json()
except Exception as e:
raise self.get_error_class(
error_message=f"Error parsing RunwayML TTS response: {e}",
status_code=raw_response.status_code,
headers=dict(raw_response.headers),
)
verbose_logger.debug("RunwayML TTS starting polling (async)...")
# Get task ID
task_id = response_data.get("id")
if not task_id:
raise ValueError("RunwayML TTS response missing task ID")
# Get headers for polling (need auth)
poll_headers = {
"Authorization": raw_response.request.headers.get("Authorization", ""),
"X-Runway-Version": raw_response.request.headers.get(
"X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION
),
}
# Poll until task completes (async)
polled_response = await self._poll_task_async(
task_id=task_id,
api_base=self.DEFAULT_BASE_URL,
headers=poll_headers,
timeout_secs=RUNWAYML_POLLING_TIMEOUT,
)
# Get the completed task data
task_data = polled_response.json()
verbose_logger.debug("RunwayML TTS polling complete (async), downloading audio")
# Get audio URL from output
output = task_data.get("output", [])
if not output or not isinstance(output, list) or len(output) == 0:
raise ValueError("RunwayML TTS response missing audio URL in output")
audio_url = output[0]
if not isinstance(audio_url, str):
raise ValueError(f"RunwayML TTS audio URL is not a string: {audio_url}")
# Download the audio file (async)
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
client = get_async_httpx_client(llm_provider=litellm.LlmProviders.RUNWAYML)
audio_response = await client.get(url=audio_url)
audio_response.raise_for_status()
verbose_logger.debug("RunwayML TTS audio downloaded successfully (async)")
# Return the audio data wrapped in HttpxBinaryResponseContent
return HttpxBinaryResponseContent(audio_response)
litellm/llms/runwayml/videos/transformation.py
+2 -7
View File
@@ -6,6 +6,8 @@ from httpx._types import RequestFiles
import litellm
from litellm.constants import RUNWAYML_DEFAULT_API_VERSION
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
@@ -23,16 +25,9 @@ from litellm.types.videos.utils import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from ...base_llm.chat.transformation import BaseLLMException as _BaseLLMException
from ...base_llm.videos.transformation import BaseVideoConfig as _BaseVideoConfig
LiteLLMLoggingObj = _LiteLLMLoggingObj
BaseVideoConfig = _BaseVideoConfig
BaseLLMException = _BaseLLMException
else:
LiteLLMLoggingObj = Any
BaseVideoConfig = Any
BaseLLMException = Any
class RunwayMLVideoConfig(BaseVideoConfig):
litellm/llms/vertex_ai/batches/handler.py
+8
View File
@@ -61,6 +61,10 @@ class VertexAIBatchPrediction(VertexLLM):
stream=None,
auth_header=None,
url=default_api_base,
model=None,
vertex_project=vertex_project or project_id,
vertex_location=vertex_location or "us-central1",
vertex_api_version="v1",
)
headers = {
@@ -166,6 +170,10 @@ class VertexAIBatchPrediction(VertexLLM):
stream=None,
auth_header=None,
url=default_api_base,
model=None,
vertex_project=vertex_project or project_id,
vertex_location=vertex_location or "us-central1",
vertex_api_version="v1",
)
headers = {
litellm/llms/vertex_ai/common_utils.py
+88 -7
View File
@@ -31,9 +31,11 @@ class VertexAIModelRoute(str, Enum):
PARTNER_MODELS = "partner_models"
GEMINI = "gemini"
GEMMA = "gemma"
BGE = "bge"
MODEL_GARDEN = "model_garden"
NON_GEMINI = "non_gemini"
VERTEX_AI_MODEL_ROUTES = [f"{route.value}/" for route in VertexAIModelRoute]
def get_vertex_ai_model_route(
model: str, litellm_params: Optional[dict] = None
@@ -60,6 +62,9 @@ def get_vertex_ai_model_route(
>>> get_vertex_ai_model_route("openai/gpt-oss-120b")
VertexAIModelRoute.MODEL_GARDEN
>>> get_vertex_ai_model_route("1234567890", {"api_base": "http://10.96.32.8"})
VertexAIModelRoute.GEMINI # Numeric endpoints with api_base use HTTP path
"""
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
VertexAIPartnerModels,
@@ -69,11 +74,20 @@ def get_vertex_ai_model_route(
if litellm_params and litellm_params.get("base_model") is not None:
if "gemini" in litellm_params["base_model"]:
return VertexAIModelRoute.GEMINI
# Check if numeric endpoint ID with custom api_base (PSC endpoint)
# Route to GEMINI (HTTP path) to support PSC endpoints properly
if model.isdigit() and litellm_params and litellm_params.get("api_base"):
return VertexAIModelRoute.GEMINI
# Check for partner models (llama, mistral, claude, etc.)
if VertexAIPartnerModels.is_vertex_partner_model(model=model):
return VertexAIModelRoute.PARTNER_MODELS
# Check for BGE models
if "bge/" in model or "bge" in model.lower():
return VertexAIModelRoute.BGE
# Check for gemma models
if "gemma/" in model:
return VertexAIModelRoute.GEMMA
@@ -136,6 +150,71 @@ all_gemini_url_modes = Literal[
]
def get_vertex_base_model_name(model: str) -> str:
"""
Strip routing prefixes from model name for PSC/endpoint URL construction.
Patterns like "bge/", "gemma/", "openai/" are used for internal routing but
should not appear in the actual endpoint URL. Routing prefixes are derived
from VertexAIModelRoute enum values.
Args:
model: The model name with potential prefix (e.g., "bge/123456", "gemma/gemma-3-12b-it")
Returns:
str: The model name without routing prefix (e.g., "123456", "gemma-3-12b-it")
Examples:
>>> get_vertex_base_model_name("bge/378943383978115072")
"378943383978115072"
>>> get_vertex_base_model_name("gemma/gemma-3-12b-it")
"gemma-3-12b-it"
>>> get_vertex_base_model_name("openai/gpt-oss-120b")
"gpt-oss-120b"
>>> get_vertex_base_model_name("1234567890")
"1234567890"
"""
# Derive routing prefixes from VertexAIModelRoute enum
# Map specific routes to their prefixes (some routes like PARTNER_MODELS, GEMINI don't have prefixes)
for route in VERTEX_AI_MODEL_ROUTES:
if model.startswith(route):
return model.replace(route, "", 1)
return model
def _get_embedding_url(
model: str,
vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_api_version: Literal["v1", "v1beta1"],
) -> Tuple[str, str]:
"""
Get URL for embedding models.
Handles special patterns:
- bge/endpoint_id -> strips to endpoint_id for endpoints/ routing
- numeric model -> routes to endpoints/
- regular model -> routes to publishers/google/models/
"""
endpoint = "predict"
# Strip routing prefixes (bge/, gemma/, etc.) for endpoint URL construction
model = get_vertex_base_model_name(model=model)
url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}"
if model.isdigit():
# https://us-central1-aiplatform.googleapis.com/v1/projects/$PROJECT_ID/locations/us-central1/endpoints/$ENDPOINT_ID:predict
url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}:{endpoint}"
return url, endpoint
def _get_vertex_url(
mode: all_gemini_url_modes,
model: str,
@@ -148,6 +227,7 @@ def _get_vertex_url(
endpoint: Optional[str] = None
model = litellm.VertexGeminiConfig.get_model_for_vertex_ai_url(model=model)
if mode == "chat":
### SET RUNTIME ENDPOINT ###
endpoint = "generateContent"
@@ -172,11 +252,12 @@ def _get_vertex_url(
if stream is True:
url += "?alt=sse"
elif mode == "embedding":
endpoint = "predict"
url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}"
if model.isdigit():
# https://us-central1-aiplatform.googleapis.com/v1/projects/$PROJECT_ID/locations/us-central1/endpoints/$ENDPOINT_ID:predict
url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}:{endpoint}"
return _get_embedding_url(
model=model,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_api_version=vertex_api_version,
)
elif mode == "image_generation":
endpoint = "predict"
url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}"
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
+4
View File
@@ -79,6 +79,10 @@ class ContextCachingEndpoints(VertexBase):
stream=None,
auth_header=auth_header,
url=url,
model=None,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_api_version="v1beta1" if custom_llm_provider == "vertex_ai_beta" else "v1",
)
def check_cache(
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+6 -1
View File
@@ -567,6 +567,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"thinkingBudget": DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET,
"includeThoughts": False,
}
elif reasoning_effort == "none":
return {
"thinkingBudget": 0,
"includeThoughts": False,
}
else:
raise ValueError(f"Invalid reasoning effort: {reasoning_effort}")
@@ -1022,7 +1027,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
if "functionCall" in part:
_function_chunk = ChatCompletionToolCallFunctionChunk(
name=part["functionCall"]["name"],
arguments=json.dumps(part["functionCall"]["args"]),
arguments=json.dumps(part["functionCall"]["args"], ensure_ascii=False),
)
if is_function_call is True:
function = _function_chunk
litellm/llms/vertex_ai/vertex_embeddings/bge.py
+182
View File
@@ -0,0 +1,182 @@
"""
Vertex AI BGE (BAAI General Embedding) Configuration
BGE models deployed on Vertex AI require different input/output format:
- Request: Use "prompt" instead of "content" as the input field
- Response: Embeddings are returned directly as arrays, not wrapped in objects
Model name handling:
- Model names like "bge/endpoint_id" are automatically transformed in common_utils._get_vertex_url()
- This module focuses on request/response transformation only
"""
from typing import List, Optional, Union
from litellm.types.utils import EmbeddingResponse, Usage
from .types import (
EmbeddingParameters,
TaskType,
TextEmbeddingBGEInput,
VertexEmbeddingRequest,
)
class VertexBGEConfig:
"""
Configuration and transformation logic for BGE models on Vertex AI.
BGE (BAAI General Embedding) models use a different request format
where the input field is named "prompt" instead of "content".
Supported model patterns (after provider split in main.py):
- "bge-small-en-v1.5" (model name)
- "bge/204379420394258432" (endpoint ID pattern)
Note: Model name transformation (bge/ -> numeric ID) is handled automatically
in common_utils._get_vertex_url(). This class focuses on request/response format only.
"""
@staticmethod
def is_bge_model(model: str) -> bool:
"""
Check if the model is a BGE (BAAI General Embedding) model.
After provider split in main.py, supports:
- "bge-small-en-v1.5" (model name)
- "bge/204379420394258432" (endpoint ID pattern)
Args:
model: The model name after provider split
Returns:
bool: True if the model is a BGE model
"""
model_lower = model.lower()
# Check for "bge/" prefix (endpoint pattern) or "bge" in model name
return model_lower.startswith("bge/") or "bge" in model_lower
@staticmethod
def transform_request(
input: Union[list, str], optional_params: dict, model: str
) -> VertexEmbeddingRequest:
"""
Transforms an OpenAI request to a Vertex BGE embedding request.
BGE models use "prompt" instead of "content" as the input field.
Args:
input: The input text(s) to embed
optional_params: Optional parameters for the request
model: The model name
Returns:
VertexEmbeddingRequest: The transformed request
"""
vertex_request: VertexEmbeddingRequest = VertexEmbeddingRequest()
vertex_text_embedding_input_list: List[TextEmbeddingBGEInput] = []
task_type: Optional[TaskType] = optional_params.get("task_type")
title = optional_params.get("title")
if isinstance(input, str):
input = [input]
for text in input:
embedding_input = VertexBGEConfig._create_embedding_input(
prompt=text, task_type=task_type, title=title
)
vertex_text_embedding_input_list.append(embedding_input)
vertex_request["instances"] = vertex_text_embedding_input_list
vertex_request["parameters"] = EmbeddingParameters(**optional_params)
return vertex_request
@staticmethod
def _create_embedding_input(
prompt: str,
task_type: Optional[TaskType] = None,
title: Optional[str] = None,
) -> TextEmbeddingBGEInput:
"""
Creates a TextEmbeddingBGEInput object for BGE models.
BGE models use "prompt" instead of "content" as the input field.
Args:
prompt: The prompt to be embedded
task_type: The type of task to be performed
title: The title of the document to be embedded
Returns:
TextEmbeddingBGEInput: A TextEmbeddingBGEInput object
"""
text_embedding_input = TextEmbeddingBGEInput(prompt=prompt)
if task_type is not None:
text_embedding_input["task_type"] = task_type
if title is not None:
text_embedding_input["title"] = title
return text_embedding_input
@staticmethod
def transform_response(
response: dict, model: str, model_response: EmbeddingResponse
) -> EmbeddingResponse:
"""
Transforms a Vertex BGE embedding response to OpenAI format.
BGE models return embeddings directly as arrays in predictions:
{
"predictions": [
[0.002, 0.021, ...],
[0.003, 0.022, ...]
]
}
Args:
response: The raw response from Vertex AI
model: The model name
model_response: The EmbeddingResponse object to populate
Returns:
EmbeddingResponse: The transformed response in OpenAI format
Raises:
KeyError: If response doesn't contain 'predictions'
ValueError: If predictions is not a list or contains invalid data
"""
if "predictions" not in response:
raise KeyError("Response missing 'predictions' field")
_predictions = response["predictions"]
if not isinstance(_predictions, list):
raise ValueError(f"Expected 'predictions' to be a list, got {type(_predictions)}")
embedding_response = []
# BGE models don't return token counts, so we estimate or set to 0
input_tokens = 0
for idx, embedding_values in enumerate(_predictions):
if not isinstance(embedding_values, list):
raise ValueError(
f"Expected embedding at index {idx} to be a list, got {type(embedding_values)}"
)
embedding_response.append(
{
"object": "embedding",
"index": idx,
"embedding": embedding_values,
}
)
model_response.object = "list"
model_response.data = embedding_response
model_response.model = model
usage = Usage(
prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens
)
setattr(model_response, "usage", usage)
return model_response
litellm/llms/vertex_ai/vertex_embeddings/transformation.py
+19 -2
View File
@@ -105,10 +105,16 @@ class VertexAITextEmbeddingConfig(BaseModel):
"""
Transforms an openai request to a vertex embedding request.
"""
# Import here to avoid circular import issues with litellm.__init__
from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig
if model.isdigit():
return self._transform_openai_request_to_fine_tuned_embedding_request(
input, optional_params, model
)
if VertexBGEConfig.is_bge_model(model):
return VertexBGEConfig.transform_request(
input=input, optional_params=optional_params, model=model
)
vertex_request: VertexEmbeddingRequest = VertexEmbeddingRequest()
vertex_text_embedding_input_list: List[TextEmbeddingInput] = []
@@ -167,6 +173,9 @@ class VertexAITextEmbeddingConfig(BaseModel):
vertex_request["parameters"] = TextEmbeddingFineTunedParameters(
**optional_params
)
# Remove 'shared_session' from parameters if present
if vertex_request["parameters"] is not None and "shared_session" in vertex_request["parameters"]:
del vertex_request["parameters"]["shared_session"] # type: ignore[typeddict-item]
return vertex_request
@@ -183,8 +192,8 @@ class VertexAITextEmbeddingConfig(BaseModel):
Args:
content (str): The content to be embedded.
task_type (Optional[TaskType]): The type of task to be performed".
title (Optional[str]): The title of the document to be embedded
task_type (Optional[TaskType]): The type of task to be performed.
title (Optional[str]): The title of the document to be embedded.
Returns:
TextEmbeddingInput: A TextEmbeddingInput object.
@@ -206,6 +215,14 @@ class VertexAITextEmbeddingConfig(BaseModel):
return self._transform_vertex_response_to_openai_for_fine_tuned_models(
response, model, model_response
)
# Import here to avoid circular import issues with litellm.__init__
from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig
if VertexBGEConfig.is_bge_model(model):
return VertexBGEConfig.transform_response(
response=response, model=model, model_response=model_response
)
_predictions = response["predictions"]
litellm/llms/vertex_ai/vertex_embeddings/types.py
+7 -1
View File
@@ -25,6 +25,12 @@ class TextEmbeddingInput(TypedDict, total=False):
title: Optional[str]
class TextEmbeddingBGEInput(TypedDict, total=False):
prompt: str
task_type: Optional[TaskType]
title: Optional[str]
# Fine-tuned models require a different input format
# Ref: https://console.cloud.google.com/vertex-ai/model-garden?hl=en&project=adroit-crow-413218&pageState=(%22galleryStateKey%22:(%22f%22:(%22g%22:%5B%5D,%22o%22:%5B%5D),%22s%22:%22%22))
class TextEmbeddingFineTunedInput(TypedDict, total=False):
@@ -44,7 +50,7 @@ class EmbeddingParameters(TypedDict, total=False):
class VertexEmbeddingRequest(TypedDict, total=False):
instances: Union[List[TextEmbeddingInput], List[TextEmbeddingFineTunedInput]]
instances: Union[List[TextEmbeddingInput], List[TextEmbeddingBGEInput], List[TextEmbeddingFineTunedInput]]
parameters: Optional[Union[EmbeddingParameters, TextEmbeddingFineTunedParameters]]
litellm/llms/vertex_ai/vertex_gemma_models/main.py
+3 -2
View File
@@ -25,7 +25,7 @@ import httpx # type: ignore
from litellm.utils import ModelResponse
from ..common_utils import VertexAIError
from ..common_utils import VertexAIError, get_vertex_base_model_name
from ..vertex_llm_base import VertexBase
@@ -82,7 +82,8 @@ class VertexAIGemmaModels(VertexBase):
message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""",
)
try:
model = model.replace("gemma/", "")
model = get_vertex_base_model_name(model=model)
vertex_httpx_logic = VertexLLM()
access_token, project_id = vertex_httpx_logic._ensure_access_token(
litellm/llms/vertex_ai/vertex_llm_base.py
+49 -3
View File
@@ -19,6 +19,7 @@ from .common_utils import (
_get_gemini_url,
_get_vertex_url,
all_gemini_url_modes,
get_vertex_base_model_name,
is_global_only_vertex_model,
)
@@ -241,6 +242,9 @@ class VertexBase:
auth_header=None,
url=default_api_base,
model=model,
vertex_project=vertex_project or project_id,
vertex_location=vertex_location or "us-central1",
vertex_api_version="v1", # Partner models typically use v1
)
return api_base
@@ -289,9 +293,18 @@ class VertexBase:
auth_header: Optional[str],
url: str,
model: Optional[str] = None,
vertex_project: Optional[str] = None,
vertex_location: Optional[str] = None,
vertex_api_version: Optional[Literal["v1", "v1beta1"]] = None,
) -> Tuple[Optional[str], str]:
"""
for cloudflare ai gateway - https://github.com/BerriAI/litellm/issues/4317
Handles custom api_base for:
1. Gemini (Google AI Studio) - constructs /models/{model}:{endpoint}
2. Vertex AI with standard proxies - constructs {api_base}:{endpoint}
3. Vertex AI with PSC endpoints - constructs full path structure
{api_base}/v1/projects/{project}/locations/{location}/endpoints/{model}:{endpoint}
## Returns
- (auth_header, url) - Tuple[Optional[str], str]
@@ -311,8 +324,37 @@ class VertexBase:
if gemini_api_key is not None:
auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment]
else:
url = "{}:{}".format(api_base, endpoint)
# For Vertex AI
# Check if this is a PSC endpoint or custom deployment
# PSC/custom endpoints need the full path structure
if vertex_project and vertex_location and model:
# Strip routing prefixes (bge/, gemma/, etc.) for endpoint URL construction
model_for_url = get_vertex_base_model_name(model=model)
# Check if model is numeric (endpoint ID) or if api_base doesn't contain googleapis.com
# These are indicators of PSC/custom endpoints
is_psc_or_custom = (
"googleapis.com" not in api_base.lower() or model_for_url.isdigit()
)
if is_psc_or_custom:
# Construct full PSC/custom endpoint URL
# Format: {api_base}/v1/projects/{project}/locations/{location}/endpoints/{model}:{endpoint}
version = vertex_api_version or "v1"
url = "{}/{}/projects/{}/locations/{}/endpoints/{}:{}".format(
api_base.rstrip("/"),
version,
vertex_project,
vertex_location,
model_for_url,
endpoint,
)
else:
# Standard proxy - just append endpoint
url = "{}:{}".format(api_base, endpoint)
else:
# Fallback to simple format if we don't have all parameters
url = "{}:{}".format(api_base, endpoint)
if stream is True:
url = url + "?alt=sse"
return auth_header, url
@@ -339,6 +381,7 @@ class VertexBase:
Returns
token, url
"""
version: Optional[Literal["v1beta1", "v1"]] = None
if custom_llm_provider == "gemini":
url, endpoint = _get_gemini_url(
mode=mode,
@@ -354,7 +397,7 @@ class VertexBase:
)
### SET RUNTIME ENDPOINT ###
version: Literal["v1beta1", "v1"] = (
version = (
"v1beta1" if should_use_v1beta1_features is True else "v1"
)
url, endpoint = _get_vertex_url(
@@ -375,6 +418,9 @@ class VertexBase:
stream=stream,
url=url,
model=model,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_api_version=version,
)
def _handle_reauthentication(
litellm/llms/vertex_ai/vertex_model_garden/main.py
+6 -2
View File
@@ -22,7 +22,7 @@ import httpx # type: ignore
from litellm.utils import ModelResponse
from ..common_utils import VertexAIError
from ..common_utils import VertexAIError, get_vertex_base_model_name
from ..vertex_llm_base import VertexBase
@@ -89,7 +89,7 @@ class VertexAIModelGardenModels(VertexBase):
message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""",
)
try:
model = model.replace("openai/", "")
model = get_vertex_base_model_name(model=model)
vertex_httpx_logic = VertexLLM()
access_token, project_id = vertex_httpx_logic._ensure_access_token(
@@ -123,6 +123,10 @@ class VertexAIModelGardenModels(VertexBase):
stream=stream,
auth_header=None,
url=default_api_base,
model=model,
vertex_project=vertex_project or project_id,
vertex_location=vertex_location or "us-central1",
vertex_api_version="v1beta1",
)
model = ""
return openai_like_chat_completions.completion(
litellm/main.py
+33
View File
@@ -6006,6 +6006,39 @@ def speech( # noqa: PLR0915
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
)
elif custom_llm_provider == "runwayml":
from litellm.llms.runwayml.text_to_speech.transformation import (
RunwayMLTextToSpeechConfig,
)
# RunwayML Text-to-Speech
if text_to_speech_provider_config is None:
raise litellm.BadRequestError(
message="RunwayML Text-to-Speech configuration not found",
model=model,
llm_provider=custom_llm_provider,
)
# Cast to specific RunwayML config type to access dispatch method
runwayml_config = cast(
RunwayMLTextToSpeechConfig, text_to_speech_provider_config
)
response = runwayml_config.dispatch_text_to_speech( # type: ignore
model=model,
input=input,
voice=voice,
optional_params=optional_params,
litellm_params_dict=litellm_params_dict,
logging_obj=logging_obj,
timeout=timeout,
extra_headers=extra_headers,
base_llm_http_handler=base_llm_http_handler,
aspeech=aspeech or False,
api_base=api_base,
api_key=api_key,
**kwargs,
)
if response is None:
raise Exception(
litellm/model_prices_and_context_window_backup.json
+176
View File
@@ -9963,6 +9963,7 @@
"supports_function_calling": false,
"supports_parallel_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": false,
"supports_system_messages": true,
"supports_tool_choice": true,
@@ -11568,6 +11569,7 @@
"supports_audio_output": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@@ -11670,6 +11672,7 @@
"litellm_provider": "vertex_ai-language-models",
"max_audio_length_hours": 8.4,
"max_audio_per_prompt": 1,
"supports_reasoning": false,
"max_images_per_prompt": 3000,
"max_input_tokens": 32768,
"max_output_tokens": 32768,
@@ -13849,6 +13852,113 @@
"supports_service_tier": true,
"supports_vision": true
},
"gpt-5.1": {
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_priority": 2.5e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_priority": 2.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_priority": 2e-05,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_service_tier": true,
"supports_vision": true
},
"gpt-5.1-2025-11-13": {
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_priority": 2.5e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_priority": 2.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_priority": 2e-05,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_service_tier": true,
"supports_vision": true
},
"gpt-5.1-chat-latest": {
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_priority": 2.5e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_priority": 2.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_priority": 2e-05,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_function_calling": false,
"supports_native_streaming": true,
"supports_parallel_function_calling": false,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": false,
"supports_vision": true
},
"gpt-5-pro": {
"input_cost_per_token": 1.5e-05,
"input_cost_per_token_batches": 7.5e-06,
@@ -14048,6 +14158,72 @@
"supports_tool_choice": true,
"supports_vision": true
},
"gpt-5.1-codex": {
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_priority": 2.5e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_priority": 2.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 1e-05,
"output_cost_per_token_priority": 2e-05,
"supported_endpoints": [
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": false,
"supports_tool_choice": true,
"supports_vision": true
},
"gpt-5.1-codex-mini": {
"cache_read_input_token_cost": 2.5e-08,
"cache_read_input_token_cost_priority": 4.5e-08,
"input_cost_per_token": 2.5e-07,
"input_cost_per_token_priority": 4.5e-07,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 2e-06,
"output_cost_per_token_priority": 3.6e-06,
"supported_endpoints": [
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": false,
"supports_tool_choice": true,
"supports_vision": true
},
"gpt-5-mini": {
"cache_read_input_token_cost": 2.5e-08,
"cache_read_input_token_cost_flex": 1.25e-08,
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+11 -9
View File
@@ -478,6 +478,12 @@ class MCPServerManager:
"""
Get the allowed MCP Servers for the user
"""
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
# If admin, get all servers
if user_api_key_auth and _user_has_admin_view(user_api_key_auth):
return list(self.get_registry().keys())
try:
allowed_mcp_servers = await MCPRequestHandler.get_allowed_mcp_servers(
user_api_key_auth
@@ -485,18 +491,14 @@ class MCPServerManager:
verbose_logger.debug(
f"Allowed MCP Servers for user api key auth: {allowed_mcp_servers}"
)
if len(allowed_mcp_servers) > 0:
return allowed_mcp_servers
else:
if len(allowed_mcp_servers) == 0:
verbose_logger.debug(
"No allowed MCP Servers found for user api key auth, returning default registry servers"
"No allowed MCP Servers found for user api key auth."
)
return list(self.get_registry().keys())
return allowed_mcp_servers
except Exception as e:
verbose_logger.warning(
f"Failed to get allowed MCP servers: {str(e)}. Returning default registry servers."
)
return list(self.get_registry().keys())
verbose_logger.warning(f"Failed to get allowed MCP servers: {str(e)}.")
return []
async def get_tools_for_server(self, server_id: str) -> List[MCPTool]:
"""
litellm/proxy/_experimental/mcp_server/server.py
+5 -3
View File
@@ -628,8 +628,10 @@ if MCP_AVAILABLE:
)
## CHECK IF USER IS ALLOWED TO CALL THIS TOOL
allowed_mcp_server_ids = await MCPRequestHandler.get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth,
allowed_mcp_server_ids = (
await global_mcp_server_manager.get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth,
)
)
allowed_mcp_servers = global_mcp_server_manager.get_mcp_servers_from_ids(
@@ -638,7 +640,7 @@ if MCP_AVAILABLE:
allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names(
mcp_servers=mcp_servers,
allowed_mcp_servers=allowed_mcp_servers,
allowed_mcp_servers=allowed_mcp_servers
)
server_name: Optional[str]
litellm/proxy/common_utils/callback_utils.py
+6 -1
View File
@@ -1,4 +1,4 @@
from typing import Any, Dict, List, Literal, Optional
from typing import Any, Dict, List, Literal, Optional, Iterable
import litellm
from litellm import get_secret
@@ -382,3 +382,8 @@ def get_metadata_variable_name_from_kwargs(
- LiteLLM is now moving to using `litellm_metadata` for our metadata
"""
return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
def normalize_callback_names(callbacks: Iterable[Any]) -> List[Any]:
if callbacks is None:
return []
return [c.lower() if isinstance(c, str) else c for c in callbacks]
litellm/proxy/health_endpoints/_health_endpoints.py
+12
View File
@@ -38,6 +38,7 @@ services = Union[
Literal[
"slack_budget_alerts",
"langfuse",
"langfuse_otel",
"slack",
"openmeter",
"webhook",
@@ -46,6 +47,7 @@ services = Union[
"datadog",
"generic_api",
"arize",
"sqs"
],
str,
]
@@ -106,6 +108,7 @@ async def health_services_endpoint( # noqa: PLR0915
"slack_budget_alerts",
"email",
"langfuse",
"langfuse_otel",
"slack",
"openmeter",
"webhook",
@@ -116,6 +119,7 @@ async def health_services_endpoint( # noqa: PLR0915
"datadog",
"generic_api",
"arize",
"sqs"
]:
raise HTTPException(
status_code=400,
@@ -196,6 +200,14 @@ async def health_services_endpoint( # noqa: PLR0915
type="user_budget",
user_info=user_info,
)
elif service == "sqs":
from litellm.integrations.sqs import SQSLogger
sqs_logger = SQSLogger()
response = await sqs_logger.async_health_check()
return {
"status": response["status"],
"message": response["error_message"],
}
if service == "slack" or service == "slack_budget_alerts":
if "slack" in general_settings.get("alerting", []):
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
+1
View File
@@ -521,6 +521,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
"url": str(request.url),
"method": request.method,
"body": copy.copy(_parsed_body), # use copy instead of deepcopy
"headers": request.headers,
},
},
"call_type": "pass_through_endpoint",
litellm/proxy/proxy_server.py
+4 -1
View File
@@ -48,6 +48,8 @@ from litellm.types.utils import (
)
from litellm.utils import load_credentials_from_list
from litellm.proxy.common_utils.callback_utils import normalize_callback_names
if TYPE_CHECKING:
from aiohttp import ClientSession
from opentelemetry.trace import Span as _Span
@@ -9052,9 +9054,10 @@ async def update_config(config_info: ConfigYAML): # noqa: PLR0915
if isinstance(
config["litellm_settings"]["success_callback"], list
) and isinstance(updated_litellm_settings["success_callback"], list):
updated_success_callbacks_normalized = normalize_callback_names(updated_litellm_settings["success_callback"])
combined_success_callback = (
config["litellm_settings"]["success_callback"]
+ updated_litellm_settings["success_callback"]
+ updated_success_callbacks_normalized
)
combined_success_callback = list(set(combined_success_callback))
config["litellm_settings"][
litellm/proxy/spend_tracking/spend_management_endpoints.py
+68
View File
@@ -18,6 +18,10 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
)
from litellm.proxy.utils import handle_exception_on_proxy
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.proxy.management_endpoints.common_utils import (
_is_user_team_admin,
_user_has_admin_view,
)
if TYPE_CHECKING:
from litellm.proxy.proxy_server import PrismaClient
@@ -1749,6 +1753,28 @@ async def ui_view_spend_logs( # noqa: PLR0915
where_conditions["spend"]["gte"] = min_spend
if max_spend is not None:
where_conditions["spend"]["lte"] = max_spend
is_admin_view = _is_admin_view_safe(user_api_key_dict=user_api_key_dict)
if not is_admin_view:
if team_id is not None:
can_view_team = await _can_team_member_view_log(
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
team_id=team_id,
)
if not can_view_team:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={
"error": "Not authorized to view team spend for team_id={}".format(
team_id
)
},
)
where_conditions["team_id"] = team_id
else:
if _can_user_view_spend_log(user_api_key_dict=user_api_key_dict):
where_conditions["user"] = user_api_key_dict.user_id
where_conditions.pop("team_id", None)
# Calculate skip value for pagination
skip = (page - 1) * page_size
@@ -2990,3 +3016,45 @@ def _build_status_filter_condition(status_filter: Optional[str]) -> Dict[str, An
return {"OR": [{"status": {"equals": "success"}}, {"status": None}]}
else:
return {"status": {"equals": status_filter}}
def _is_admin_view_safe(user_api_key_dict: UserAPIKeyAuth) -> bool:
"""
Safely determine if the current user has admin view permissions.
Wraps the underlying check and defaults to False on any exception.
"""
try:
return _user_has_admin_view(user_api_key_dict=user_api_key_dict)
except Exception:
return False
async def _can_team_member_view_log(
prisma_client,
user_api_key_dict: UserAPIKeyAuth,
team_id: Optional[str],
) -> bool:
"""
Check if the requesting user can view spend logs for the given team.
Returns True only if the team exists and the user is a team admin.
"""
if team_id is None:
return False
team_obj = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
if team_obj is None:
return False
return _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj)
def _can_user_view_spend_log(user_api_key_dict: UserAPIKeyAuth) -> bool:
"""
Check if the requesting user can view their own spend logs.
"""
user_role = user_api_key_dict.user_role
user_id = user_api_key_dict.user_id
return user_role in (
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
) and user_id is not None
litellm/utils.py
+10
View File
@@ -7720,6 +7720,10 @@ class ProviderConfigManager:
from litellm.llms.azure_ai.image_edit import get_azure_ai_image_edit_config
return get_azure_ai_image_edit_config(model)
elif LlmProviders.GEMINI == provider:
from litellm.llms.gemini.image_edit import get_gemini_image_edit_config
return get_gemini_image_edit_config(model)
elif LlmProviders.LITELLM_PROXY == provider:
from litellm.llms.litellm_proxy.image_edit.transformation import (
LiteLLMProxyImageEditConfig,
@@ -7807,6 +7811,12 @@ class ProviderConfigManager:
)
return AzureAVATextToSpeechConfig()
elif litellm.LlmProviders.RUNWAYML == provider:
from litellm.llms.runwayml.text_to_speech.transformation import (
RunwayMLTextToSpeechConfig,
)
return RunwayMLTextToSpeechConfig()
return None
@staticmethod
model_prices_and_context_window.json
+176
View File
@@ -9963,6 +9963,7 @@
"supports_function_calling": false,
"supports_parallel_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": false,
"supports_system_messages": true,
"supports_tool_choice": true,
@@ -11568,6 +11569,7 @@
"supports_audio_output": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@@ -11670,6 +11672,7 @@
"litellm_provider": "vertex_ai-language-models",
"max_audio_length_hours": 8.4,
"max_audio_per_prompt": 1,
"supports_reasoning": false,
"max_images_per_prompt": 3000,
"max_input_tokens": 32768,
"max_output_tokens": 32768,
@@ -13849,6 +13852,113 @@
"supports_service_tier": true,
"supports_vision": true
},
"gpt-5.1": {
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_priority": 2.5e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_priority": 2.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_priority": 2e-05,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_service_tier": true,
"supports_vision": true
},
"gpt-5.1-2025-11-13": {
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_priority": 2.5e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_priority": 2.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_priority": 2e-05,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_service_tier": true,
"supports_vision": true
},
"gpt-5.1-chat-latest": {
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_priority": 2.5e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_priority": 2.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_priority": 2e-05,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_function_calling": false,
"supports_native_streaming": true,
"supports_parallel_function_calling": false,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": false,
"supports_vision": true
},
"gpt-5-pro": {
"input_cost_per_token": 1.5e-05,
"input_cost_per_token_batches": 7.5e-06,
@@ -14048,6 +14158,72 @@
"supports_tool_choice": true,
"supports_vision": true
},
"gpt-5.1-codex": {
"cache_read_input_token_cost": 1.25e-07,
"cache_read_input_token_cost_priority": 2.5e-07,
"input_cost_per_token": 1.25e-06,
"input_cost_per_token_priority": 2.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 1e-05,
"output_cost_per_token_priority": 2e-05,
"supported_endpoints": [
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": false,
"supports_tool_choice": true,
"supports_vision": true
},
"gpt-5.1-codex-mini": {
"cache_read_input_token_cost": 2.5e-08,
"cache_read_input_token_cost_priority": 4.5e-08,
"input_cost_per_token": 2.5e-07,
"input_cost_per_token_priority": 4.5e-07,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"output_cost_per_token": 2e-06,
"output_cost_per_token_priority": 3.6e-06,
"supported_endpoints": [
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": false,
"supports_tool_choice": true,
"supports_vision": true
},
"gpt-5-mini": {
"cache_read_input_token_cost": 2.5e-08,
"cache_read_input_token_cost_flex": 1.25e-08,
provider_endpoints_support.json
+1 -1
View File
@@ -1368,7 +1368,7 @@
"embeddings": false,
"image_generations": true,
"audio_transcriptions": false,
"audio_speech": false,
"audio_speech": true,
"moderations": false,
"batches": false,
"rerank": false,
tests/audio_tests/runwayml_speech.mp3
BIN
View File
Binary file not shown.
tests/audio_tests/test_audio_speech.py
+54
View File
@@ -382,6 +382,60 @@ async def test_azure_ava_tts_async():
pytest.fail(f"Test failed with exception: {str(e)}")
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
async def test_runwayml_tts_async():
"""
Test RunwayML Text-to-Speech with real API request.
"""
litellm._turn_on_debug()
api_key = os.getenv("RUNWAYML_API_KEY")
api_base = os.getenv("RUNWAYML_API_BASE")
speech_file_path = Path(__file__).parent / "runwayml_speech.mp3"
try:
response = await litellm.aspeech(
model="runwayml/eleven_multilingual_v2",
voice="Rachel",
input="Yuneng is gone, we miss him so much I hope he has a good coffee",
api_base=api_base,
api_key=api_key,
response_format="mp3",
speed=1.0,
)
# Assert the response is HttpxBinaryResponseContent
from litellm.types.llms.openai import HttpxBinaryResponseContent
assert isinstance(response, HttpxBinaryResponseContent)
# Get the binary content
binary_content = response.content
assert len(binary_content) > 0
# MP3 files start with these magic bytes
# ID3 tag or MPEG sync word
assert binary_content[:3] == b"ID3" or binary_content[:2] == b"\xff\xfb" or binary_content[:2] == b"\xff\xf3"
# Write to file
response.stream_to_file(speech_file_path)
# Verify file was created and has content
assert speech_file_path.exists()
assert speech_file_path.stat().st_size > 0
print(f"Azure TTS audio saved to: {speech_file_path}")
# assert response cost is greater than 0
print("Response cost: ", response._hidden_params["response_cost"])
assert response._hidden_params["response_cost"] > 0
except Exception as e:
pytest.fail(f"Test failed with exception: {str(e)}")
@pytest.mark.asyncio
async def test_azure_ava_tts_with_custom_voice():
"""
tests/llm_translation/test_bedrock_agentcore.py
+241
View File
@@ -126,3 +126,244 @@ def test_bedrock_agentcore_with_custom_params():
assert "prompt" in request_data
assert request_data["prompt"] == "Explain machine learning in simple terms"
def test_bedrock_agentcore_with_runtime_user_id():
"""
Test AgentCore with runtimeUserId parameter
"""
import json
litellm._turn_on_debug()
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
with patch.object(client, "post", return_value=MagicMock()) as mock_post:
try:
response = litellm.completion(
model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC",
messages=[
{
"role": "user",
"content": "Hello",
}
],
runtimeUserId="test-user-123",
client=client,
)
except Exception as e:
print(f"Error: {e}")
mock_post.assert_called_once()
call_kwargs = mock_post.call_args.kwargs
print(f"mock_post.call_args.kwargs: {call_kwargs}")
# Verify headers - user ID should be in header
assert "headers" in call_kwargs
headers = call_kwargs["headers"]
print(f"Headers: {headers}")
assert "X-Amzn-Bedrock-AgentCore-Runtime-User-Id" in headers
assert headers["X-Amzn-Bedrock-AgentCore-Runtime-User-Id"] == "test-user-123"
def test_bedrock_agentcore_with_session_and_user():
"""
Test AgentCore with both runtimeSessionId and runtimeUserId
"""
import json
litellm._turn_on_debug()
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
with patch.object(client, "post", return_value=MagicMock()) as mock_post:
try:
response = litellm.completion(
model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC",
messages=[
{
"role": "user",
"content": "Test message",
}
],
runtimeSessionId="session-abc-123",
runtimeUserId="user-xyz-789",
client=client,
)
except Exception as e:
print(f"Error: {e}")
mock_post.assert_called_once()
call_kwargs = mock_post.call_args.kwargs
print(f"mock_post.call_args.kwargs: {call_kwargs}")
# Verify headers contain both session and user IDs
assert "headers" in call_kwargs
headers = call_kwargs["headers"]
print(f"Headers: {headers}")
assert "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id" in headers
assert headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] == "session-abc-123"
assert "X-Amzn-Bedrock-AgentCore-Runtime-User-Id" in headers
assert headers["X-Amzn-Bedrock-AgentCore-Runtime-User-Id"] == "user-xyz-789"
def test_bedrock_agentcore_with_api_key_bearer_token():
"""
Test AgentCore with api_key parameter for JWT/Bearer token authentication
"""
import json
litellm._turn_on_debug()
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
test_jwt_token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"
with patch.object(client, "post", return_value=MagicMock()) as mock_post:
try:
response = litellm.completion(
model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC",
messages=[
{
"role": "user",
"content": "Test JWT authentication",
}
],
api_key=test_jwt_token,
client=client,
)
except Exception as e:
print(f"Error: {e}")
mock_post.assert_called_once()
call_kwargs = mock_post.call_args.kwargs
print(f"mock_post.call_args.kwargs: {call_kwargs}")
# Verify Authorization header with Bearer token
assert "headers" in call_kwargs
headers = call_kwargs["headers"]
print(f"Headers: {headers}")
assert "Authorization" in headers
assert headers["Authorization"] == f"Bearer {test_jwt_token}"
assert headers["Content-Type"] == "application/json"
# Verify the request body is JSON-encoded (not SigV4 signed)
assert "data" in call_kwargs
request_data = json.loads(call_kwargs["data"])
print(f"Request data: {json.dumps(request_data, indent=2)}")
assert "prompt" in request_data
assert request_data["prompt"] == "Test JWT authentication"
def test_bedrock_agentcore_with_all_parameters():
"""
Test AgentCore with all parameters: api_key, runtimeSessionId, runtimeUserId
"""
import json
litellm._turn_on_debug()
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
test_jwt_token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.signature"
with patch.object(client, "post", return_value=MagicMock()) as mock_post:
try:
response = litellm.completion(
model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC",
messages=[
{
"role": "user",
"content": "Complete test",
}
],
api_key=test_jwt_token,
runtimeSessionId="full-test-session-id",
runtimeUserId="full-test-user-id",
qualifier="LATEST",
client=client,
)
except Exception as e:
print(f"Error: {e}")
mock_post.assert_called_once()
call_kwargs = mock_post.call_args.kwargs
print(f"mock_post.call_args.kwargs: {call_kwargs}")
# Verify URL includes qualifier
assert "url" in call_kwargs
url = call_kwargs["url"]
print(f"URL: {url}")
assert "qualifier=LATEST" in url
# Verify all headers are present
assert "headers" in call_kwargs
headers = call_kwargs["headers"]
print(f"Headers: {headers}")
# Check Bearer token authorization
assert "Authorization" in headers
assert headers["Authorization"] == f"Bearer {test_jwt_token}"
# Check session and user IDs
assert "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id" in headers
assert headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] == "full-test-session-id"
assert "X-Amzn-Bedrock-AgentCore-Runtime-User-Id" in headers
assert headers["X-Amzn-Bedrock-AgentCore-Runtime-User-Id"] == "full-test-user-id"
# Verify JSON body
assert "data" in call_kwargs
request_data = json.loads(call_kwargs["data"])
print(f"Request data: {json.dumps(request_data, indent=2)}")
assert "prompt" in request_data
assert request_data["prompt"] == "Complete test"
def test_bedrock_agentcore_without_api_key_uses_sigv4():
"""
Test that AgentCore uses AWS SigV4 signing when api_key is not provided
"""
import json
litellm._turn_on_debug()
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
with patch.object(client, "post", return_value=MagicMock()) as mock_post:
try:
response = litellm.completion(
model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC",
messages=[
{
"role": "user",
"content": "Test SigV4",
}
],
# No api_key provided - should use SigV4
runtimeSessionId="sigv4-test-session",
client=client,
)
except Exception as e:
print(f"Error: {e}")
mock_post.assert_called_once()
call_kwargs = mock_post.call_args.kwargs
print(f"mock_post.call_args.kwargs: {call_kwargs}")
# Verify headers - should have AWS SigV4 headers, not Bearer token
assert "headers" in call_kwargs
headers = call_kwargs["headers"]
print(f"Headers: {headers}")
# Should NOT have Bearer Authorization when using SigV4
if "Authorization" in headers:
assert not headers["Authorization"].startswith("Bearer ")
# Should have AWS4-HMAC-SHA256 signature
assert "AWS4-HMAC-SHA256" in headers["Authorization"]
# Session ID should still be present
assert "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id" in headers
assert headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] == "sigv4-test-session"
tests/llm_translation/test_gemini.py
+101 -7
View File
@@ -290,10 +290,17 @@ def test_gemini_image_generation():
)
def test_gemini_2_5_flash_image_preview():
@pytest.mark.parametrize(
"model_name",
[
"gemini/gemini-2.5-flash-image-preview",
"gemini/gemini-2.0-flash-preview-image-generation",
],
)
def test_gemini_flash_image_preview_models(model_name: str):
"""
Test for GitHub issue #14120 - gemini-2.5-flash-image-preview model routing fix
Validates that the model correctly routes to image generation instead of chat completion
Validate Gemini Flash image preview models route through image_generation()
and invoke the generateContent endpoint returning inline image data.
"""
from unittest.mock import patch, MagicMock
from litellm.types.utils import ImageResponse, ImageObject
@@ -321,7 +328,7 @@ def test_gemini_2_5_flash_image_preview():
# Test that the function works without throwing the original 400 error
response = litellm.image_generation(
model="gemini/gemini-2.5-flash-image-preview",
model=model_name,
prompt="Generate a simple test image",
api_key="test_api_key",
)
@@ -339,9 +346,9 @@ def test_gemini_2_5_flash_image_preview():
call_args[0][0] if call_args[0] else call_args.kwargs.get("url", "")
)
# Verify it uses generateContent endpoint for gemini-2.5-flash-image-preview (not predict)
# Verify it uses generateContent endpoint for Gemini Flash image preview models (not predict)
assert ":generateContent" in called_url
assert "gemini-2.5-flash-image-preview" in called_url
assert model_name.split("/", 1)[1] in called_url
# Verify request format is Gemini format (not Imagen)
request_data = call_args.kwargs.get("json", {})
@@ -356,7 +363,6 @@ def test_gemini_2_5_flash_image_preview():
"TEXT",
]
def test_gemini_imagen_models_use_predict_endpoint():
"""
Test that Imagen models still use :predict endpoint (not broken by gemini-2.5-flash-image-preview fix)
@@ -1129,3 +1135,91 @@ def test_gemini_embedding():
)
print("response: ", response)
assert response is not None
def test_reasoning_effort_none_mapping():
"""
Test that reasoning_effort='none' correctly maps to thinkingConfig.
Related issue: https://github.com/BerriAI/litellm/issues/16420
"""
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexGeminiConfig,
)
# Test reasoning_effort="none" mapping
result = VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(
reasoning_effort="none",
model="gemini-2.0-flash-thinking-exp-01-21",
)
assert result is not None
assert result["thinkingBudget"] == 0
assert result["includeThoughts"] is False
def test_gemini_function_args_preserve_unicode():
"""
Test for Issue #16533: Gemini function call arguments should preserve non-ASCII characters
https://github.com/BerriAI/litellm/issues/16533
Before fix: "や" becomes "\u3084"
After fix: "や" stays as "や"
"""
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
# Test Japanese characters
parts = [
{
"functionCall": {
"name": "send_message",
"args": {
"message": "やあ", # Japanese "hello"
"recipient": "たけし" # Japanese name
}
}
}
]
function, tools, _ = VertexGeminiConfig._transform_parts(
parts=parts,
cumulative_tool_call_idx=0,
is_function_call=False
)
arguments_str = tools[0]['function']['arguments']
parsed_args = json.loads(arguments_str)
# Verify characters are preserved
assert parsed_args["message"] == "やあ", "Japanese characters should be preserved"
assert parsed_args["recipient"] == "たけし", "Japanese characters should be preserved"
# Verify no Unicode escape sequences in raw string
assert "\\u" not in arguments_str, "Should not contain Unicode escape sequences"
assert "やあ" in arguments_str, "Original Japanese characters should be in the string"
assert "たけし" in arguments_str, "Original Japanese characters should be in the string"
# Test Spanish characters
parts_spanish = [
{
"functionCall": {
"name": "send_message",
"args": {
"message": "¡Hola! ¿Cómo estás?",
"recipient": "José"
}
}
}
]
function, tools, _ = VertexGeminiConfig._transform_parts(
parts=parts_spanish,
cumulative_tool_call_idx=0,
is_function_call=False
)
arguments_str = tools[0]['function']['arguments']
parsed_args = json.loads(arguments_str)
assert parsed_args["message"] == "¡Hola! ¿Cómo estás?"
assert parsed_args["recipient"] == "José"
assert "\\u" not in arguments_str
assert "José" in arguments_str
tests/logging_callback_tests/test_sqs_logger.py
+24
View File
@@ -432,3 +432,27 @@ async def test_strip_base64_recursive_redaction():
s = json.dumps(c).lower()
# allow "[base64_redacted]" but nothing else
assert "base64," not in s, f"Found real base64 blob in: {s}"
@pytest.mark.asyncio
async def test_async_health_check_healthy(monkeypatch):
monkeypatch.setattr("litellm.aws_sqs_callback_params", {})
monkeypatch.setattr(asyncio, "create_task", MagicMock())
logger = SQSLogger(sqs_queue_url="https://example.com", sqs_region_name="us-west-2")
logger.async_send_message = AsyncMock(return_value=None)
result = await logger.async_health_check()
assert result["status"] == "healthy"
assert result.get("error_message") is None
@pytest.mark.asyncio
async def test_async_health_check_unhealthy(monkeypatch):
monkeypatch.setattr("litellm.aws_sqs_callback_params", {})
monkeypatch.setattr(asyncio, "create_task", MagicMock())
logger = SQSLogger(sqs_queue_url="https://example.com", sqs_region_name="us-west-2")
logger.async_send_message = AsyncMock(side_effect=Exception("boom"))
result = await logger.async_health_check()
assert result["status"] == "unhealthy"
assert "boom" in (result.get("error_message") or "")
tests/mcp_tests/test_mcp_server.py
+158 -4
View File
@@ -88,7 +88,15 @@ async def test_mcp_server_manager_https_server():
}
)
tools = await mcp_server_manager.list_tools()
allowed_server_ids = list(mcp_server_manager.get_registry().keys())
assert allowed_server_ids, "Expected registry to contain the configured server"
with patch.object(
mcp_server_manager,
"get_allowed_mcp_servers",
new=AsyncMock(return_value=allowed_server_ids),
):
tools = await mcp_server_manager.list_tools()
print("TOOLS FROM MCP SERVER MANAGER== ", tools)
# Verify tools were returned and properly prefixed
@@ -192,7 +200,15 @@ async def test_mcp_http_transport_list_tools_mock():
)
# Call list_tools
tools = await test_manager.list_tools()
allowed_server_ids = list(test_manager.get_registry().keys())
assert allowed_server_ids, "Expected registry to contain configured server"
with patch.object(
test_manager,
"get_allowed_mcp_servers",
new=AsyncMock(return_value=allowed_server_ids),
):
tools = await test_manager.list_tools()
# Assertions
assert len(tools) == 2
@@ -1530,8 +1546,16 @@ async def test_mcp_protocol_version_passed_to_client():
}
)
# Call list_tools with a specific protocol version from request
await test_manager.list_tools()
allowed_server_ids = list(test_manager.get_registry().keys())
assert allowed_server_ids, "Expected registry to contain configured server"
with patch.object(
test_manager,
"get_allowed_mcp_servers",
new=AsyncMock(return_value=allowed_server_ids),
):
# Call list_tools with a specific protocol version from request
await test_manager.list_tools()
# Verify the client was created with the correct protocol version
mock_client.list_tools.assert_called()
@@ -2436,3 +2460,133 @@ async def test_mcp_server_manager_with_access_groups_integration():
# Should only get servers user has access to
assert len(allowed_servers) >= 0 # At least verify no errors
mock_get_allowed.assert_called_once_with(user_auth)
@pytest.mark.asyncio
async def test_get_allowed_mcp_servers_returns_registry_for_admin():
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
test_manager = MCPServerManager()
test_manager.load_servers_from_config(
{
"alpha_server": {
"url": "https://alpha.server/mcp",
"transport": MCPTransport.http,
},
"beta_server": {
"url": "https://beta.server/mcp",
"transport": MCPTransport.http,
},
}
)
admin_auth = UserAPIKeyAuth(
api_key="admin-key",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
with patch.object(
MCPRequestHandler, "get_allowed_mcp_servers", new_callable=AsyncMock
) as mock_permission_lookup:
allowed_servers = await test_manager.get_allowed_mcp_servers(admin_auth)
assert set(allowed_servers) == set(test_manager.get_registry().keys())
mock_permission_lookup.assert_not_called()
@pytest.mark.asyncio
async def test_get_allowed_mcp_servers_returns_empty_for_non_admin_without_permissions():
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
test_manager = MCPServerManager()
test_manager.load_servers_from_config(
{
"alpha_server": {
"url": "https://alpha.server/mcp",
"transport": MCPTransport.http,
},
"beta_server": {
"url": "https://beta.server/mcp",
"transport": MCPTransport.http,
},
}
)
user_auth = UserAPIKeyAuth(
api_key="user-key",
user_role=LitellmUserRoles.INTERNAL_USER,
)
with patch.object(
MCPRequestHandler, "get_allowed_mcp_servers", new_callable=AsyncMock
) as mock_permission_lookup:
mock_permission_lookup.return_value = []
allowed_servers = await test_manager.get_allowed_mcp_servers(user_auth)
assert allowed_servers == []
mock_permission_lookup.assert_awaited_once()
@pytest.mark.asyncio
async def test_call_mcp_tool_uses_manager_permission_lookup():
from litellm.proxy._experimental.mcp_server.server import (
call_mcp_tool,
global_mcp_server_manager,
)
mock_server = MCPServer(
server_id="server-123",
name="test_server",
alias="test_server",
server_name="test_server",
url="https://test-server.com/mcp",
transport=MCPTransport.http,
mcp_info={"server_name": "test_server"},
)
expected_response = [TextContent(type="text", text="ok")]
with patch.object(
global_mcp_server_manager,
"get_allowed_mcp_servers",
new_callable=AsyncMock,
) as mock_get_allowed, patch.object(
global_mcp_server_manager,
"get_mcp_servers_from_ids",
return_value=[mock_server],
), patch.object(
global_mcp_server_manager,
"_get_mcp_server_from_tool_name",
return_value=mock_server,
) as mock_get_server, patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_tool_registry"
) as mock_tool_registry, patch(
"litellm.proxy._experimental.mcp_server.server._handle_managed_mcp_tool",
new_callable=AsyncMock,
) as mock_handle_managed, patch(
"litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed",
return_value=True,
):
mock_get_allowed.return_value = [mock_server.server_id]
mock_tool_registry.get_tool.return_value = None
mock_handle_managed.return_value = expected_response
result = await call_mcp_tool(
name=f"{mock_server.name}/gmail_send_email",
arguments={"body": "hello"},
mcp_servers=["test_server"],
)
assert result == expected_response
mock_get_allowed.assert_awaited_once()
assert mock_get_server.call_count == 2
assert (
mock_get_server.call_args_list[0][0][0]
== f"{mock_server.name}/gmail_send_email"
)
tests/proxy_unit_tests/test_proxy_server.py
+58
View File
@@ -2401,3 +2401,61 @@ def test_non_root_ui_path_logic(monkeypatch, tmp_path, ui_exists, ui_has_content
error_calls = [call[0][0] for call in mock_logger.error.call_args_list]
assert any("Path exists:" in call for call in error_calls)
assert mock_logger.info.call_count == 0
@pytest.mark.asyncio
async def test_update_config_success_callback_normalization():
"""
Ensure success_callback values are normalized to lowercase when updating config.
This prevents delete_callback (which searches lowercase) from failing on mixed case inputs like 'SQS'.
"""
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._types import ConfigYAML
# Ensure feature is enabled and prisma_client is set
setattr(proxy_server, "store_model_in_db", True)
setattr(proxy_server, "proxy_logging_obj", MagicMock())
class MockPrisma:
def __init__(self):
self.db = MagicMock()
self.db.litellm_config = MagicMock()
self.db.litellm_config.upsert = AsyncMock()
# proxy_server.update_config expects this to be sync returning a dict
def jsonify_object(self, obj):
return obj
setattr(proxy_server, "prisma_client", MockPrisma())
class MockProxyConfig:
def __init__(self):
self.saved_config = None
async def get_config(self):
# Existing config has one lowercase callback already
return {"litellm_settings": {"success_callback": ["langfuse"]}}
async def save_config(self, new_config: dict):
self.saved_config = new_config
async def add_deployment(self, prisma_client=None, proxy_logging_obj=None):
return None
mock_proxy_config = MockProxyConfig()
setattr(proxy_server, "proxy_config", mock_proxy_config)
# Update config with mixed-case callbacks - expect normalization to lowercase
config_update = ConfigYAML(litellm_settings={"success_callback": ["SQS", "sQs"]})
await proxy_server.update_config(config_update)
saved = mock_proxy_config.saved_config
assert saved is not None, "save_config was not called"
callbacks = saved["litellm_settings"]["success_callback"]
# Deduped and normalized
assert "sqs" in callbacks
assert "SQS" not in callbacks
assert "sQs" not in callbacks
# Existing callback should still be present
assert "langfuse" in callbacks
tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py
+14
View File
@@ -57,6 +57,20 @@ def test_is_error_str_context_window_exceeded(error_str, expected):
class TestExceptionCheckers:
"""Test the ExceptionCheckers utility methods"""
def test_is_error_str_rate_limit_ignores_embedded_numbers(self):
"""An arbitrary 429 inside user-provided payload must not trigger rate-limit detection"""
error_str = "Invalid user message={'role': 'user', 'content': [{'text': 'payload429snippet'}]}"
result = ExceptionCheckers.is_error_str_rate_limit(error_str)
assert result is False
def test_is_error_str_rate_limit_detects_true_rate_limit(self):
"""A real rate-limit error string should still be detected"""
error_str = "RateLimitError: OpenAIException - You exceeded your current quota. (status code 429)"
result = ExceptionCheckers.is_error_str_rate_limit(error_str)
assert result is True
def test_is_azure_content_policy_violation_error_with_policy_violation_text(self):
"""Test detection of Azure content policy violation with explicit policy violation text"""
tests/test_litellm/litellm_core_utils/test_token_counter.py
+266
View File
@@ -631,3 +631,269 @@ def test_bad_input_token_counter(model, messages):
messages=messages,
default_token_count=1000,
)
def test_token_counter_with_anthropic_tool_use():
"""
Test that _count_anthropic_content() correctly handles tool_use blocks.
Validates that:
- 'name' field is counted (string)
- 'input' field is counted (dict serialized to string)
- Metadata fields ('type', 'id') are skipped
"""
messages = [
{
"role": "user",
"content": "What's the weather in San Francisco?"
},
{
"role": "assistant",
"content": [
{
"type": "text",
"text": "I'll check the weather for you."
},
{
"type": "tool_use",
"id": "toolu_01234567890", # Should be skipped
"name": "get_weather", # Should be counted
"input": { # Should be counted (serialized)
"location": "San Francisco, CA",
"unit": "fahrenheit"
}
}
]
}
]
tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
assert tokens > 0, f"Expected positive token count, got {tokens}"
# Should count: user message + "I'll check" text + "get_weather" name + input dict
assert tokens > 15, f"Expected reasonable token count for message with tool_use, got {tokens}"
def test_token_counter_with_anthropic_tool_result():
"""
Test that _count_anthropic_content() correctly handles tool_result blocks.
Validates that:
- 'content' field (when string) is counted
- Metadata fields ('type', 'tool_use_id') are skipped
- Full conversation with tool_use → tool_result flow works
"""
messages = [
{
"role": "user",
"content": "What's the weather in San Francisco?"
},
{
"role": "assistant",
"content": [
{
"type": "tool_use",
"id": "toolu_01234567890",
"name": "get_weather",
"input": {
"location": "San Francisco, CA"
}
}
]
},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "toolu_01234567890", # Should be skipped
"content": "The weather in San Francisco is 65°F and sunny." # Should be counted
}
]
}
]
tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
assert tokens > 0, f"Expected positive token count, got {tokens}"
assert tokens > 25, f"Expected reasonable token count for conversation with tool_result, got {tokens}"
def test_token_counter_with_nested_tool_result():
"""
Test that _count_anthropic_content() recursively handles nested content lists.
Validates that:
- tool_result with 'content' as a list (not string) is handled
- Nested content blocks are recursively counted via _count_content_list()
- TypedDict inference correctly identifies list fields
"""
messages = [
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "toolu_01234567890",
"content": [ # Nested list - should recursively count
{
"type": "text",
"text": "The weather in San Francisco is 65°F and sunny."
},
{
"type": "text",
"text": "UV index is moderate."
}
]
}
]
}
]
tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
assert tokens > 0, f"Expected positive token count, got {tokens}"
# Should count both nested text blocks
assert tokens > 15, f"Expected reasonable token count for nested tool_result, got {tokens}"
def test_token_counter_tool_use_and_result_combined():
"""
Test dynamic field inference with multiple tool_use and tool_result blocks.
Validates that:
- Multiple tool_use blocks in same message are handled
- Multiple tool_result blocks in same message are handled
- skip_fields correctly filters metadata across all blocks
- Full realistic conversation flow works end-to-end
"""
messages = [
{
"role": "user",
"content": "What's the weather in San Francisco and New York?"
},
{
"role": "assistant",
"content": [
{
"type": "text",
"text": "I'll check the weather in both cities for you."
},
{
"type": "tool_use",
"id": "toolu_01A",
"name": "get_weather",
"input": {"location": "San Francisco, CA"}
},
{
"type": "tool_use",
"id": "toolu_01B",
"name": "get_weather",
"input": {"location": "New York, NY"}
}
]
},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "toolu_01A",
"content": "San Francisco: 65°F, sunny"
},
{
"type": "tool_result",
"tool_use_id": "toolu_01B",
"content": "New York: 45°F, cloudy"
}
]
},
{
"role": "assistant",
"content": "The weather in San Francisco is 65°F and sunny, while New York is cooler at 45°F and cloudy."
}
]
tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
assert tokens > 0, f"Expected positive token count, got {tokens}"
# Should count all text, tool names, inputs, and results
assert tokens > 60, f"Expected substantial token count for full tool conversation, got {tokens}"
def test_token_counter_with_image_url():
"""
Test that _count_image_tokens() correctly handles image_url content blocks.
Validates that:
- image_url as dict with 'url' and 'detail' is handled
- image_url as string is handled
- 'detail' field validation works ('low', 'high', 'auto')
- calculate_img_tokens is called with correct parameters
"""
# Test with dict format (detail: low)
messages_dict = [
{
"role": "user",
"content": [
{
"type": "text",
"text": "What's in this image?"
},
{
"type": "image_url",
"image_url": {
"url": "https://example.com/image.jpg",
"detail": "low" # Should use low token count (85 base tokens)
}
}
]
}
]
tokens_dict = token_counter(
model="gpt-3.5-turbo",
messages=messages_dict,
use_default_image_token_count=True # Avoid actual HTTP request
)
assert tokens_dict > 0, f"Expected positive token count, got {tokens_dict}"
assert tokens_dict > 85, f"Expected at least base image tokens, got {tokens_dict}"
# Test with string format (defaults to auto/low)
messages_str = [
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": "https://example.com/image.jpg" # String format
}
]
}
]
tokens_str = token_counter(
model="gpt-3.5-turbo",
messages=messages_str,
use_default_image_token_count=True
)
assert tokens_str > 0, f"Expected positive token count for string image_url, got {tokens_str}"
# Test invalid detail value raises error
messages_invalid = [
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {
"url": "https://example.com/image.jpg",
"detail": "invalid" # Should raise ValueError
}
}
]
}
]
try:
token_counter(model="gpt-3.5-turbo", messages=messages_invalid)
assert False, "Expected ValueError for invalid detail value"
except ValueError as e:
assert "Invalid detail value" in str(e), f"Expected detail validation error, got: {e}"
tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py
+149
View File
@@ -0,0 +1,149 @@
import base64
import json
from io import BytesIO
from typing import Dict
from unittest.mock import MagicMock
import httpx
import pytest
from litellm.llms.gemini.image_edit.transformation import GeminiImageEditConfig
class TestGeminiImageEditTransformation:
def setup_method(self) -> None:
self.config = GeminiImageEditConfig()
self.model = "gemini-2.5-flash-image-preview"
self.prompt = "Enhance this photo with a dramatic night sky."
self.logging_obj = MagicMock()
def test_map_openai_params(self) -> None:
optional_params: Dict[str, object] = {
"size": "1792x1024",
"response_format": "b64_json",
"quality": "high",
}
mapped = self.config.map_openai_params(
image_edit_optional_params=optional_params, # type: ignore[arg-type]
model=self.model,
drop_params=False,
)
assert mapped["aspectRatio"] == "16:9"
assert "response_format" not in mapped
assert "quality" not in mapped
def test_transform_image_edit_request(self) -> None:
image_bytes = b"fake_image_data"
image = BytesIO(image_bytes)
optional_params = {
"sampleCount": 2,
"aspectRatio": "16:9",
}
request_body, files = self.config.transform_image_edit_request(
model=self.model,
prompt=self.prompt,
image=[image], # Gemini pipeline passes list of images
image_edit_optional_request_params=optional_params,
litellm_params=MagicMock(),
headers={},
)
assert files == []
parts = request_body["contents"][0]["parts"]
assert parts[-1]["text"] == self.prompt
inline_data = parts[0]["inlineData"]
assert inline_data["mimeType"] == "image/png"
assert base64.b64decode(inline_data["data"]) == image_bytes
generation_config = request_body["generationConfig"]
assert generation_config["aspectRatio"] == "16:9"
def test_transform_image_edit_request_multiple_images(self) -> None:
image_one = BytesIO(b"image_one")
image_two = BytesIO(b"image_two")
request_body, files = self.config.transform_image_edit_request(
model=self.model,
prompt=self.prompt,
image=[image_one, image_two],
image_edit_optional_request_params={},
litellm_params=MagicMock(),
headers={},
)
assert files == []
parts = request_body["contents"][0]["parts"]
assert len(parts) == 3 # two images + text prompt
assert parts[-1]["text"] == self.prompt
assert base64.b64decode(parts[0]["inlineData"]["data"]) == b"image_one"
assert base64.b64decode(parts[1]["inlineData"]["data"]) == b"image_two"
def test_transform_image_edit_response(self) -> None:
response_payload = {
"candidates": [
{
"content": {
"parts": [
{
"inlineData": {
"mimeType": "image/png",
"data": base64.b64encode(b"image-one").decode("utf-8"),
}
}
]
}
},
{
"content": {
"parts": [
{
"inlineData": {
"mimeType": "image/png",
"data": base64.b64encode(b"image-two").decode("utf-8"),
}
}
]
}
},
]
}
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.return_value = response_payload
mock_response.status_code = 200
mock_response.headers = {}
image_response = self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
assert image_response.data is not None
assert len(image_response.data) == 2
assert image_response.data[0].b64_json == base64.b64encode(b"image-one").decode(
"utf-8"
)
assert image_response.data[1].b64_json == base64.b64encode(b"image-two").decode(
"utf-8"
)
def test_transform_image_edit_request_without_image_raises(self) -> None:
optional_params = {}
with pytest.raises(ValueError):
self.config.transform_image_edit_request(
model=self.model,
prompt=self.prompt,
image=[],
image_edit_optional_request_params=optional_params,
litellm_params=MagicMock(),
headers={},
)
tests/test_litellm/llms/runwayml/test_text_to_speech_transformation.py
+67
View File
@@ -0,0 +1,67 @@
"""
Test RunwayML text-to-speech transformation
"""
import os
import sys
sys.path.insert(0, os.path.abspath("../../../.."))
from litellm.llms.runwayml.text_to_speech.transformation import (
RunwayMLTextToSpeechConfig,
)
def test_openai_voice_mapping_to_runwayml():
"""
Test that OpenAI voice names are correctly mapped to RunwayML preset IDs
"""
config = RunwayMLTextToSpeechConfig()
# Test OpenAI voice mappings
openai_to_runway = {
"alloy": "Maya",
"echo": "James",
"fable": "Bernard",
"onyx": "Vincent",
"nova": "Serene",
"shimmer": "Ella",
}
for openai_voice, expected_runway_voice in openai_to_runway.items():
mapped_voice, mapped_params = config.map_openai_params(
model="eleven_multilingual_v2",
optional_params={},
voice=openai_voice,
drop_params=False,
kwargs={},
)
assert mapped_voice is None
assert "runwayml_voice" in mapped_params
assert mapped_params["runwayml_voice"]["type"] == "runway-preset"
assert mapped_params["runwayml_voice"]["presetId"] == expected_runway_voice
def test_runwayml_native_voice_passthrough():
"""
Test that RunwayML native voice names are passed through correctly as-is
"""
config = RunwayMLTextToSpeechConfig()
# Test various RunwayML native voices
runway_voices = ["Bernard", "Maya", "Arjun", "Serene", "Chad"]
for runway_voice in runway_voices:
mapped_voice, mapped_params = config.map_openai_params(
model="eleven_multilingual_v2",
optional_params={},
voice=runway_voice,
drop_params=False,
kwargs={},
)
assert mapped_voice is None
assert "runwayml_voice" in mapped_params
assert mapped_params["runwayml_voice"]["type"] == "runway-preset"
assert mapped_params["runwayml_voice"]["presetId"] == runway_voice
tests/test_litellm/llms/vertex_ai/test_bge_embedding.py
+251
View File
@@ -0,0 +1,251 @@
"""
Test BGE embeddings with Vertex AI using custom api_base.
This test ensures that BGE embeddings work correctly with Vertex AI
and that the request body is properly formatted.
"""
import json
import os
import sys
from unittest.mock import MagicMock, patch
sys.path.insert(
0, os.path.abspath("../../../..")
)
import pytest
import litellm
from litellm.llms.custom_httpx.http_handler import HTTPHandler
def test_vertex_ai_bge_embedding_with_custom_api_base():
"""
Test Vertex AI BGE embeddings with custom api_base.
This test verifies that when using a BGE model with Vertex AI and
a custom api_base, the request is properly formatted and sent to
the correct endpoint.
"""
client = HTTPHandler()
def mock_auth_token(*args, **kwargs):
return "fake-token", "fake-project"
with patch.object(client, "post") as mock_post, patch(
"litellm.llms.vertex_ai.vertex_embeddings.embedding_handler.VertexEmbedding._ensure_access_token",
side_effect=mock_auth_token
):
mock_response = MagicMock()
mock_response.status_code = 200
# BGE models return embeddings directly as arrays, not wrapped in objects
mock_response.json.return_value = {
"predictions": [
[0.1, 0.2, 0.3, 0.4, 0.5],
[0.6, 0.7, 0.8, 0.9, 1.0]
],
"deployedModelId": "849506872875548672",
"model": "projects/1060139831167/locations/us-central1/models/baai_bge-small-en-v1.5",
"modelDisplayName": "baai_bge-small-en-v1.5",
"modelVersionId": "1"
}
mock_post.return_value = mock_response
response = litellm.embedding(
model="vertex_ai/bge-small-en-v1.5",
input=["Hello", "World"],
api_base="http://10.96.32.8",
client=client
)
mock_post.assert_called_once()
call_args = mock_post.call_args
kwargs = call_args.kwargs if hasattr(call_args, 'kwargs') else call_args[1]
if "url" in kwargs:
api_url_called = kwargs["url"]
elif len(call_args[0]) > 0:
api_url_called = call_args[0][0]
else:
api_url_called = "Unknown"
# Vertex AI may use 'json' or 'data' parameter
if "json" in kwargs:
request_data = kwargs["json"]
elif "data" in kwargs:
request_data = json.loads(kwargs["data"])
else:
request_data = {}
print("\n" + "="*50)
print("Mock Request Body Received:")
print("="*50)
print(json.dumps(request_data, indent=2))
print("="*50)
print(f"API Base: {api_url_called}")
print("="*50 + "\n")
assert "instances" in request_data
assert len(request_data["instances"]) == 2
# BGE models should use "prompt" instead of "content"
assert "prompt" in request_data["instances"][0]
assert request_data["instances"][0]["prompt"] == "Hello"
assert "prompt" in request_data["instances"][1]
assert request_data["instances"][1]["prompt"] == "World"
assert isinstance(response.data, list)
assert len(response.data) == 2
assert "embedding" in response.data[0]
def test_vertex_ai_bge_with_endpoint_id_pattern():
"""
Test BGE with vertex_ai/bge/endpoint_id pattern.
This test verifies that the pattern vertex_ai/bge/204379420394258432
correctly triggers BGE transformations and routes to the endpoint.
"""
client = HTTPHandler()
def mock_auth_token(*args, **kwargs):
return "fake-token", "fake-project"
with patch.object(client, "post") as mock_post, patch(
"litellm.llms.vertex_ai.vertex_embeddings.embedding_handler.VertexEmbedding._ensure_access_token",
side_effect=mock_auth_token
):
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"predictions": [
[0.1, 0.2, 0.3, 0.4, 0.5],
[0.6, 0.7, 0.8, 0.9, 1.0]
],
"deployedModelId": "204379420394258432",
"model": "projects/1060139831167/locations/europe-west4/models/baai_bge-base-en",
"modelDisplayName": "baai_bge-base-en",
"modelVersionId": "1"
}
mock_post.return_value = mock_response
response = litellm.embedding(
model="vertex_ai/bge/204379420394258432",
input=["Hello", "World"],
vertex_project="1060139831167",
vertex_location="europe-west4",
client=client
)
mock_post.assert_called_once()
call_args = mock_post.call_args
kwargs = call_args.kwargs if hasattr(call_args, 'kwargs') else call_args[1]
if "url" in kwargs:
api_url_called = kwargs["url"]
elif len(call_args[0]) > 0:
api_url_called = call_args[0][0]
else:
api_url_called = "Unknown"
# Vertex AI may use 'json' or 'data' parameter
if "json" in kwargs:
request_data = kwargs["json"]
elif "data" in kwargs:
request_data = json.loads(kwargs["data"])
else:
request_data = {}
print("\n" + "="*50)
print("BGE Endpoint Pattern Test:")
print("="*50)
print(f"Model: vertex_ai/bge/204379420394258432")
print(f"API URL: {api_url_called}")
print("Request Body:")
print(json.dumps(request_data, indent=2))
print("="*50 + "\n")
# Verify URL contains the endpoint ID and uses endpoints/ path
assert "204379420394258432" in api_url_called, f"Endpoint ID not in URL: {api_url_called}"
assert "endpoints" in api_url_called, f"Expected 'endpoints' in URL, got: {api_url_called}"
# Verify BGE-specific request format (uses "prompt" not "content")
assert "instances" in request_data
assert "prompt" in request_data["instances"][0]
assert request_data["instances"][0]["prompt"] == "Hello"
# Verify response
assert isinstance(response.data, list)
assert len(response.data) == 2
def test_vertex_ai_bge_psc_endpoint_url_construction():
"""
Test that BGE models with PSC endpoints construct correct URL without bge/ prefix.
Verifies that vertex_ai/bge/378943383978115072 with api_base http://10.128.16.2
constructs URL: http://10.128.16.2/v1/projects/{project}/locations/{location}/endpoints/378943383978115072:predict
The bge/ prefix should be stripped from the endpoint URL.
"""
client = HTTPHandler()
def mock_auth_token(*args, **kwargs):
return "fake-token", "gen-lang-client-0682925754"
with patch.object(client, "post") as mock_post, patch(
"litellm.llms.vertex_ai.vertex_embeddings.embedding_handler.VertexEmbedding._ensure_access_token",
side_effect=mock_auth_token
):
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"predictions": [
[0.1, 0.2, 0.3, 0.4, 0.5]
]
}
mock_post.return_value = mock_response
response = litellm.embedding(
model="vertex_ai/bge/378943383978115072",
input=["The food was delicious and the waiter.."],
api_base="http://10.128.16.2",
vertex_project="gen-lang-client-0682925754",
vertex_location="us-central1",
client=client
)
mock_post.assert_called_once()
call_args = mock_post.call_args
kwargs = call_args.kwargs if hasattr(call_args, 'kwargs') else call_args[1]
if "url" in kwargs:
api_url_called = kwargs["url"]
elif len(call_args[0]) > 0:
api_url_called = call_args[0][0]
else:
api_url_called = "Unknown"
print("\n" + "="*50)
print("PSC Endpoint URL Construction Test:")
print("="*50)
print(f"Model: vertex_ai/bge/378943383978115072")
print(f"API Base: http://10.128.16.2")
print(f"Constructed URL: {api_url_called}")
print("="*50 + "\n")
# Verify the URL is constructed correctly
expected_url = "http://10.128.16.2/v1/projects/gen-lang-client-0682925754/locations/us-central1/endpoints/378943383978115072:predict"
assert api_url_called == expected_url, f"Expected URL: {expected_url}, Got: {api_url_called}"
# Verify bge/ prefix is NOT in the URL
assert "bge/" not in api_url_called, f"URL should not contain 'bge/' prefix: {api_url_called}"
# Verify response works
assert isinstance(response.data, list)
assert len(response.data) == 1
tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py
+111
View File
@@ -0,0 +1,111 @@
"""
Test BGE response transformation validation.
This test verifies that the BGE response transformer properly validates
and handles different response formats.
"""
import os
import sys
sys.path.insert(
0, os.path.abspath("../../../..")
)
import pytest
from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig
from litellm.types.utils import EmbeddingResponse
def test_is_bge_model_detection():
"""
Test BGE model detection for post-provider-split patterns.
After main.py splits the provider, model strings are passed without the provider prefix.
Model name transformation (bge/ -> numeric ID) is handled in common_utils._get_vertex_url().
"""
# Should detect BGE models (after provider split)
assert VertexBGEConfig.is_bge_model("bge-small-en-v1.5") is True
assert VertexBGEConfig.is_bge_model("bge/204379420394258432") is True
assert VertexBGEConfig.is_bge_model("BGE-large-en-v1.5") is True # case insensitive
# Should not detect non-BGE models
assert VertexBGEConfig.is_bge_model("textembedding-gecko") is False
assert VertexBGEConfig.is_bge_model("gemma") is False
assert VertexBGEConfig.is_bge_model("123456789") is False
def test_bge_response_transformation_success():
"""
Test successful BGE response transformation.
Verifies that a valid BGE response is properly transformed
to OpenAI format.
"""
response = {
"predictions": [
[0.1, 0.2, 0.3],
[0.4, 0.5, 0.6]
],
"deployedModelId": "123456",
"model": "projects/test/models/bge-base"
}
model_response = EmbeddingResponse()
result = VertexBGEConfig.transform_response(
response=response,
model="bge-small-en-v1.5",
model_response=model_response
)
assert result.object == "list"
assert len(result.data) == 2
assert result.data[0]["embedding"] == [0.1, 0.2, 0.3]
assert result.data[1]["embedding"] == [0.4, 0.5, 0.6]
assert result.data[0]["index"] == 0
assert result.data[1]["index"] == 1
assert result.model == "bge-small-en-v1.5"
def test_bge_response_missing_predictions():
"""
Test BGE response transformation with missing predictions field.
Verifies that a KeyError is raised when the response doesn't
contain the required 'predictions' field.
"""
response = {
"deployedModelId": "123456",
"model": "projects/test/models/bge-base"
}
model_response = EmbeddingResponse()
with pytest.raises(KeyError, match="Response missing 'predictions' field"):
VertexBGEConfig.transform_response(
response=response,
model="bge-small-en-v1.5",
model_response=model_response
)
def test_bge_response_invalid_predictions_type():
"""
Test BGE response transformation with invalid predictions type.
Verifies that a ValueError is raised when predictions is not a list.
"""
response = {
"predictions": "not-a-list"
}
model_response = EmbeddingResponse()
with pytest.raises(ValueError, match="Expected 'predictions' to be a list"):
VertexBGEConfig.transform_response(
response=response,
model="bge-small-en-v1.5",
model_response=model_response
)
tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py
+258
View File
@@ -0,0 +1,258 @@
"""
Unit tests for Vertex AI Private Service Connect (PSC) endpoint support
Tests that LiteLLM properly constructs URLs when using custom api_base
for PSC endpoints.
"""
import pytest
import sys
import os
# Add the litellm package to the path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../../.."))
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
class TestVertexAIPSCEndpointSupport:
"""Test cases for PSC endpoint URL construction"""
def test_psc_endpoint_url_construction_basic(self):
"""Test basic PSC endpoint URL construction for predict endpoint"""
vertex_base = VertexBase()
psc_api_base = "http://10.96.32.8"
endpoint_id = "1234567890"
project_id = "test-project"
location = "us-central1"
auth_header, url = vertex_base._check_custom_proxy(
api_base=psc_api_base,
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="predict",
stream=False,
auth_header="test-token",
url="", # This will be replaced
model=endpoint_id,
vertex_project=project_id,
vertex_location=location,
vertex_api_version="v1",
)
expected_url = f"{psc_api_base}/v1/projects/{project_id}/locations/{location}/endpoints/{endpoint_id}:predict"
assert (
url == expected_url
), f"Expected {expected_url}, but got {url}"
def test_psc_endpoint_url_construction_with_streaming(self):
"""Test PSC endpoint URL construction with streaming enabled"""
vertex_base = VertexBase()
psc_api_base = "http://10.96.32.8"
endpoint_id = "1234567890"
project_id = "test-project"
location = "us-central1"
auth_header, url = vertex_base._check_custom_proxy(
api_base=psc_api_base,
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="streamGenerateContent",
stream=True,
auth_header="test-token",
url="",
model=endpoint_id,
vertex_project=project_id,
vertex_location=location,
vertex_api_version="v1",
)
expected_url = f"{psc_api_base}/v1/projects/{project_id}/locations/{location}/endpoints/{endpoint_id}:streamGenerateContent?alt=sse"
assert (
url == expected_url
), f"Expected {expected_url}, but got {url}"
def test_psc_endpoint_url_construction_v1beta1(self):
"""Test PSC endpoint URL construction with v1beta1 API version"""
vertex_base = VertexBase()
psc_api_base = "http://10.96.32.8"
endpoint_id = "1234567890"
project_id = "test-project"
location = "us-central1"
auth_header, url = vertex_base._check_custom_proxy(
api_base=psc_api_base,
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="predict",
stream=False,
auth_header="test-token",
url="",
model=endpoint_id,
vertex_project=project_id,
vertex_location=location,
vertex_api_version="v1beta1",
)
expected_url = f"{psc_api_base}/v1beta1/projects/{project_id}/locations/{location}/endpoints/{endpoint_id}:predict"
assert (
url == expected_url
), f"Expected {expected_url}, but got {url}"
def test_psc_endpoint_url_with_https(self):
"""Test PSC endpoint URL construction with HTTPS"""
vertex_base = VertexBase()
psc_api_base = "https://10.96.32.8"
endpoint_id = "1234567890"
project_id = "test-project"
location = "us-central1"
auth_header, url = vertex_base._check_custom_proxy(
api_base=psc_api_base,
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="predict",
stream=False,
auth_header="test-token",
url="",
model=endpoint_id,
vertex_project=project_id,
vertex_location=location,
vertex_api_version="v1",
)
expected_url = f"{psc_api_base}/v1/projects/{project_id}/locations/{location}/endpoints/{endpoint_id}:predict"
assert (
url == expected_url
), f"Expected {expected_url}, but got {url}"
def test_psc_endpoint_with_trailing_slash(self):
"""Test that trailing slashes in api_base are handled correctly"""
vertex_base = VertexBase()
psc_api_base = "http://10.96.32.8/"
endpoint_id = "1234567890"
project_id = "test-project"
location = "us-central1"
auth_header, url = vertex_base._check_custom_proxy(
api_base=psc_api_base,
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="predict",
stream=False,
auth_header="test-token",
url="",
model=endpoint_id,
vertex_project=project_id,
vertex_location=location,
vertex_api_version="v1",
)
# rstrip('/') should remove the trailing slash
expected_url = f"{psc_api_base.rstrip('/')}/v1/projects/{project_id}/locations/{location}/endpoints/{endpoint_id}:predict"
assert (
url == expected_url
), f"Expected {expected_url}, but got {url}"
def test_standard_proxy_with_googleapis(self):
"""Test that standard proxies with googleapis.com in URL use simple format"""
vertex_base = VertexBase()
proxy_api_base = "https://my-proxy.googleapis.com"
endpoint_id = "gemini-pro" # Not numeric
project_id = "test-project"
location = "us-central1"
auth_header, url = vertex_base._check_custom_proxy(
api_base=proxy_api_base,
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="generateContent",
stream=False,
auth_header="test-token",
url="",
model=endpoint_id,
vertex_project=project_id,
vertex_location=location,
vertex_api_version="v1",
)
# Should use simple format: api_base:endpoint
expected_url = f"{proxy_api_base}:generateContent"
assert (
url == expected_url
), f"Expected {expected_url}, but got {url}"
def test_custom_proxy_with_numeric_model(self):
"""Test that numeric model IDs trigger PSC-style URL construction"""
vertex_base = VertexBase()
proxy_api_base = "https://my-custom-proxy.example.com"
endpoint_id = "9876543210" # Numeric endpoint ID
project_id = "test-project"
location = "us-central1"
auth_header, url = vertex_base._check_custom_proxy(
api_base=proxy_api_base,
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="predict",
stream=False,
auth_header="test-token",
url="",
model=endpoint_id,
vertex_project=project_id,
vertex_location=location,
vertex_api_version="v1",
)
# Numeric model should trigger full path construction
expected_url = f"{proxy_api_base}/v1/projects/{project_id}/locations/{location}/endpoints/{endpoint_id}:predict"
assert (
url == expected_url
), f"Expected {expected_url}, but got {url}"
def test_no_api_base_returns_original_url(self):
"""Test that when api_base is None, the original URL is returned"""
vertex_base = VertexBase()
original_url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test/locations/us-central1/publishers/google/models/gemini-pro:generateContent"
auth_header, url = vertex_base._check_custom_proxy(
api_base=None,
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="generateContent",
stream=False,
auth_header="test-token",
url=original_url,
model="gemini-pro",
vertex_project="test-project",
vertex_location="us-central1",
vertex_api_version="v1",
)
# When api_base is None, original URL should be returned unchanged
assert url == original_url, f"Expected {original_url}, but got {url}"
def test_auth_header_preserved(self):
"""Test that auth_header is properly preserved"""
vertex_base = VertexBase()
psc_api_base = "http://10.96.32.8"
test_auth_header = "Bearer test-token-12345"
auth_header, url = vertex_base._check_custom_proxy(
api_base=psc_api_base,
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="predict",
stream=False,
auth_header=test_auth_header,
url="",
model="1234567890",
vertex_project="test-project",
vertex_location="us-central1",
vertex_api_version="v1",
)
assert (
auth_header == test_auth_header
), f"Auth header should be preserved, got {auth_header}"
tests/test_litellm/proxy/common_utils/test_callback_utils.py
+11
View File
@@ -7,6 +7,7 @@ sys.path.insert(
from litellm.proxy.common_utils.callback_utils import (
get_remaining_tokens_and_requests_from_request_data,
normalize_callback_names,
)
@@ -27,3 +28,13 @@ def test_get_remaining_tokens_and_requests_from_request_data():
f"x-litellm-key-remaining-requests-{expected_name}": 100,
f"x-litellm-key-remaining-tokens-{expected_name}": 200,
}
def test_normalize_callback_names_none_returns_empty_list():
assert normalize_callback_names(None) == []
assert normalize_callback_names([]) == []
def test_normalize_callback_names_lowercases_strings():
assert normalize_callback_names(["SQS", "S3", "CUSTOM_CALLBACK"]) == ["sqs", "s3", "custom_callback"]
tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
+30 -1
View File
@@ -3,7 +3,7 @@ import json
import os
import sys
from datetime import datetime, timedelta
from unittest.mock import MagicMock, patch
from unittest.mock import MagicMock, patch, AsyncMock
sys.path.insert(
0, os.path.abspath("../../..")
@@ -16,6 +16,7 @@ from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy.health_endpoints._health_endpoints import (
_db_health_readiness_check,
db_health_cache,
health_services_endpoint,
)
@@ -97,3 +98,31 @@ async def test_db_health_readiness_check_with_error_and_flag_off(prisma_error):
# Verify that the raised exception is the same
assert excinfo.value == prisma_error
@pytest.mark.asyncio
@pytest.mark.parametrize(
"status,error_message",
[
("healthy", ""),
("unhealthy", "queue not reachable"),
],
)
async def test_health_services_endpoint_sqs(status, error_message):
"""
Verify the /health/services SQS branch returns expected status and message
based on SQSLogger.async_health_check().
"""
with patch("litellm.integrations.sqs.SQSLogger") as MockSQSLogger:
mock_instance = MagicMock()
mock_instance.async_health_check = AsyncMock(
return_value={"status": status, "error_message": error_message}
)
MockSQSLogger.return_value = mock_instance
result = await health_services_endpoint(service="sqs")
assert result["status"] == status
assert result["message"] == error_message
mock_instance.async_health_check.assert_awaited_once()
tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
+291
View File
@@ -21,6 +21,169 @@ from litellm.proxy.proxy_server import app, prisma_client
from litellm.proxy.spend_tracking import spend_management_endpoints
from litellm.router import Router
from litellm.types.utils import BudgetConfig
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles, Member
from litellm.proxy.spend_tracking import spend_management_endpoints
import litellm.proxy.proxy_server as ps
@pytest.mark.asyncio
async def test_is_admin_view_safe_true(monkeypatch):
# Force underlying check to return True
monkeypatch.setattr(
spend_management_endpoints, "_user_has_admin_view", lambda user_api_key_dict: True
)
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user")
assert spend_management_endpoints._is_admin_view_safe(auth) is True
@pytest.mark.asyncio
async def test_is_admin_view_safe_false(monkeypatch):
# Force underlying check to return False
monkeypatch.setattr(
spend_management_endpoints, "_user_has_admin_view", lambda user_api_key_dict: False
)
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1")
assert spend_management_endpoints._is_admin_view_safe(auth) is False
@pytest.mark.asyncio
async def test_is_admin_view_safe_exception(monkeypatch):
# Ensure exceptions are swallowed and return False
def raise_err(*args, **kwargs):
raise RuntimeError("boom")
monkeypatch.setattr(spend_management_endpoints, "_user_has_admin_view", raise_err)
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1")
assert spend_management_endpoints._is_admin_view_safe(auth) is False
@pytest.mark.asyncio
async def test_can_team_member_view_log_none_team_id():
# team_id=None should immediately return False
class MockPrisma:
class DB:
class TeamTable:
async def find_unique(self, where: dict):
return None
def __init__(self):
self.litellm_teamtable = self.TeamTable()
def __init__(self):
self.db = self.DB()
prisma = MockPrisma()
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1")
allowed = await spend_management_endpoints._can_team_member_view_log(
prisma, auth, None
)
assert allowed is False
@pytest.mark.asyncio
async def test_can_team_member_view_log_team_not_found(monkeypatch):
# Non-existent team should return False
class MockPrisma:
class DB:
class TeamTable:
async def find_unique(self, where: dict):
return None
def __init__(self):
self.litellm_teamtable = self.TeamTable()
def __init__(self):
self.db = self.DB()
prisma = MockPrisma()
# Even if admin check would return True, no team means False
monkeypatch.setattr(
spend_management_endpoints, "_is_user_team_admin", lambda user_api_key_dict, team_obj: True
)
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1")
allowed = await spend_management_endpoints._can_team_member_view_log(
prisma, auth, "team_x"
)
assert allowed is False
@pytest.mark.asyncio
async def test_can_team_member_view_log_not_admin(monkeypatch):
# Existing team but caller is not a team admin -> False
class MockTeam:
pass
class MockPrisma:
class DB:
class TeamTable:
async def find_unique(self, where: dict):
return MockTeam()
def __init__(self):
self.litellm_teamtable = self.TeamTable()
def __init__(self):
self.db = self.DB()
prisma = MockPrisma()
monkeypatch.setattr(
spend_management_endpoints, "_is_user_team_admin", lambda user_api_key_dict, team_obj: False
)
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1")
allowed = await spend_management_endpoints._can_team_member_view_log(
prisma, auth, "team_x"
)
assert allowed is False
@pytest.mark.asyncio
async def test_can_team_member_view_log_admin(monkeypatch):
# Existing team and caller is team admin -> True
class MockTeam:
pass
class MockPrisma:
class DB:
class TeamTable:
async def find_unique(self, where: dict):
return MockTeam()
def __init__(self):
self.litellm_teamtable = self.TeamTable()
def __init__(self):
self.db = self.DB()
prisma = MockPrisma()
monkeypatch.setattr(
spend_management_endpoints, "_is_user_team_admin", lambda user_api_key_dict, team_obj: True
)
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1")
allowed = await spend_management_endpoints._can_team_member_view_log(
prisma, auth, "team_x"
)
assert allowed is True
def test_can_user_view_spend_log_true_for_internal_user():
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="u1")
assert spend_management_endpoints._can_user_view_spend_log(auth) is True
def test_can_user_view_spend_log_true_for_internal_view_only():
auth = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, user_id="u1"
)
assert spend_management_endpoints._can_user_view_spend_log(auth) is True
def test_can_user_view_spend_log_false_without_user_id():
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id=None)
assert spend_management_endpoints._can_user_view_spend_log(auth) is False
def test_can_user_view_spend_log_false_for_other_roles():
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin")
assert spend_management_endpoints._can_user_view_spend_log(auth) is False
ignored_keys = [
"request_id",
@@ -255,6 +418,134 @@ async def test_ui_view_spend_logs_with_team_id(client, monkeypatch):
assert data["data"][0]["team_id"] == "team1"
@pytest.mark.asyncio
async def test_ui_view_spend_logs_internal_user_scoped_without_user_id(client, monkeypatch):
"""
Internal users should only be able to view their own spend even if user_id is not provided.
"""
# Mock spend logs for 2 users
mock_spend_logs = [
{"id": "log1", "request_id": "req1", "api_key": "sk-test-key", "user": "internal_user_1", "team_id": "team1", "spend": 0.05, "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-3.5-turbo"},
{"id": "log2", "request_id": "req2", "api_key": "sk-test-key", "user": "internal_user_2", "team_id": "team1", "spend": 0.10, "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-4"},
]
# Prisma client mock that filters by "user" where condition
class MockDB:
async def find_many(self, *args, **kwargs):
where = kwargs.get("where", {})
if "user" in where and where["user"] == "internal_user_1":
return [mock_spend_logs[0]]
return mock_spend_logs
async def count(self, *args, **kwargs):
where = kwargs.get("where", {})
if "user" in where and where["user"] == "internal_user_1":
return 1
return len(mock_spend_logs)
class MockPrismaClient:
def __init__(self):
self.db = MockDB()
self.db.litellm_spendlogs = self.db
mock_prisma_client = MockPrismaClient()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
# Override auth dependency to return INTERNAL_USER with specific user_id
# Override using the function reference attached to the running app module
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user_1"
)
try:
start_date = (datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7)).strftime("%Y-%m-%d %H:%M:%S")
end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
# No user_id provided; should auto-scope to authenticated internal user's own id
response = client.get(
"/spend/logs/ui",
params={"start_date": start_date, "end_date": end_date},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
data = response.json()
assert data["total"] == 1
assert len(data["data"]) == 1
assert data["data"][0]["user"] == "internal_user_1"
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_ui_view_spend_logs_team_admin_can_view_team_spend(client, monkeypatch):
"""
Team admins should be able to view team-wide spend when team_id is provided.
"""
# Mock spend logs for two teams
mock_spend_logs = [
{"id": "log1", "request_id": "req1", "api_key": "sk-test-key", "user": "member1", "team_id": "team_admin_team", "spend": 0.05, "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-3.5-turbo"},
{"id": "log2", "request_id": "req2", "api_key": "sk-test-key", "user": "member2", "team_id": "team_other", "spend": 0.10, "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-4"},
]
class MockDB:
async def find_many(self, *args, **kwargs):
where = kwargs.get("where", {})
if "team_id" in where and where["team_id"] == "team_admin_team":
return [mock_spend_logs[0]]
return mock_spend_logs
async def count(self, *args, **kwargs):
where = kwargs.get("where", {})
if "team_id" in where and where["team_id"] == "team_admin_team":
return 1
return len(mock_spend_logs)
class MockPrismaClient:
def __init__(self):
self.db = MockDB()
self.db.litellm_spendlogs = self.db
# Team lookup for RBAC check
class TeamTable:
def __init__(self):
# user "admin_user" is team admin
self.members_with_roles = [Member(user_id="admin_user", role="admin")]
async def find_unique(where: dict):
if where == {"team_id": "team_admin_team"}:
return TeamTable()
return None
self.db.litellm_teamtable = self
self.litellm_teamtable = self
self.find_unique = find_unique
mock_prisma_client = MockPrismaClient()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
# Override auth dependency to return INTERNAL_USER (who is a team admin via team.members_with_roles)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, user_id="admin_user"
)
try:
start_date = (datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7)).strftime("%Y-%m-%d %H:%M:%S")
end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
response = client.get(
"/spend/logs/ui",
params={"team_id": "team_admin_team", "start_date": start_date, "end_date": end_date},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
data = response.json()
assert data["total"] == 1
assert len(data["data"]) == 1
assert data["data"][0]["team_id"] == "team_admin_team"
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio
async def test_ui_view_spend_logs_pagination(client, monkeypatch):
# Create a larger set of mock data for pagination testing
ui/litellm-dashboard/src/components/callback_info_helpers.tsx
+79 -57
View File
@@ -7,7 +7,7 @@ interface CallbackConfig {
description: string;
}
const asset_logos_folder = '/ui/assets/logos/';
const asset_logos_folder = "/ui/assets/logos/";
export const CALLBACK_CONFIGS: CallbackConfig[] = [
{
@@ -16,10 +16,10 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [
logo: `${asset_logos_folder}arize.png`,
supports_key_team_logging: true,
dynamic_params: {
"arize_api_key": "password",
"arize_space_key": "password",
arize_api_key: "password",
arize_space_key: "password",
},
description: "Arize Logging Integration"
description: "Arize Logging Integration",
},
{
id: "braintrust",
@@ -27,10 +27,10 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [
logo: `${asset_logos_folder}braintrust.png`,
supports_key_team_logging: false,
dynamic_params: {
"braintrust_api_key": "password",
"braintrust_project_name": "text"
braintrust_api_key: "password",
braintrust_project_name: "text",
},
description: "Braintrust Logging Integration"
description: "Braintrust Logging Integration",
},
{
id: "custom_callback_api",
@@ -38,10 +38,10 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [
logo: `${asset_logos_folder}custom.svg`,
supports_key_team_logging: true,
dynamic_params: {
"custom_callback_api_url": "text",
"custom_callback_api_headers": "text"
custom_callback_api_url: "text",
custom_callback_api_headers: "text",
},
description: "Custom Callback API Logging Integration"
description: "Custom Callback API Logging Integration",
},
{
id: "datadog",
@@ -49,10 +49,10 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [
logo: `${asset_logos_folder}datadog.png`,
supports_key_team_logging: false,
dynamic_params: {
"dd_api_key": "password",
"dd_site": "text"
dd_api_key: "password",
dd_site: "text",
},
description: "Datadog Logging Integration"
description: "Datadog Logging Integration",
},
{
id: "lago",
@@ -60,10 +60,10 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [
logo: `${asset_logos_folder}lago.svg`,
supports_key_team_logging: false,
dynamic_params: {
"lago_api_url": "text",
"lago_api_key": "password"
lago_api_url: "text",
lago_api_key: "password",
},
description: "Lago Billing Logging Integration"
description: "Lago Billing Logging Integration",
},
{
id: "langfuse",
@@ -71,11 +71,11 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [
logo: `${asset_logos_folder}langfuse.png`,
supports_key_team_logging: true,
dynamic_params: {
"langfuse_public_key": "text",
"langfuse_secret_key": "password",
"langfuse_host": "text"
langfuse_public_key: "text",
langfuse_secret_key: "password",
langfuse_host: "text",
},
description: "Langfuse v2 Logging Integration"
description: "Langfuse v2 Logging Integration",
},
{
id: "langfuse_otel",
@@ -83,11 +83,11 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [
logo: `${asset_logos_folder}langfuse.png`,
supports_key_team_logging: true,
dynamic_params: {
"langfuse_public_key": "text",
"langfuse_secret_key": "password",
"langfuse_host": "text"
langfuse_public_key: "text",
langfuse_secret_key: "password",
langfuse_host: "text",
},
description: "Langfuse v3 OTEL Logging Integration"
description: "Langfuse v3 OTEL Logging Integration",
},
{
id: "langsmith",
@@ -95,12 +95,12 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [
logo: `${asset_logos_folder}langsmith.png`,
supports_key_team_logging: true,
dynamic_params: {
"langsmith_api_key": "password",
"langsmith_project": "text",
"langsmith_base_url": "text",
"langsmith_sampling_rate": "number"
langsmith_api_key: "password",
langsmith_project: "text",
langsmith_base_url: "text",
langsmith_sampling_rate: "number",
},
description: "Langsmith Logging Integration"
description: "Langsmith Logging Integration",
},
{
id: "openmeter",
@@ -108,10 +108,10 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [
logo: `${asset_logos_folder}openmeter.png`,
supports_key_team_logging: false,
dynamic_params: {
"openmeter_api_key": "password",
"openmeter_base_url": "text"
openmeter_api_key: "password",
openmeter_base_url: "text",
},
description: "OpenMeter Logging Integration"
description: "OpenMeter Logging Integration",
},
{
id: "otel",
@@ -119,10 +119,10 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [
logo: `${asset_logos_folder}otel.png`,
supports_key_team_logging: false,
dynamic_params: {
"otel_endpoint": "text",
"otel_headers": "text"
otel_endpoint: "text",
otel_headers: "text",
},
description: "OpenTelemetry Logging Integration"
description: "OpenTelemetry Logging Integration",
},
{
id: "s3",
@@ -130,48 +130,70 @@ export const CALLBACK_CONFIGS: CallbackConfig[] = [
logo: `${asset_logos_folder}aws.svg`,
supports_key_team_logging: false,
dynamic_params: {
"s3_bucket_name": "text",
"aws_access_key_id": "password",
"aws_secret_access_key": "password",
"aws_region": "text"
s3_bucket_name: "text",
aws_access_key_id: "password",
aws_secret_access_key: "password",
aws_region: "text",
},
description: "S3 Bucket (AWS) Logging Integration"
}
description: "S3 Bucket (AWS) Logging Integration",
},
{
id: "SQS",
displayName: "SQS",
logo: `${asset_logos_folder}aws.svg`,
supports_key_team_logging: false,
dynamic_params: {
sqs_queue_url: "text",
aws_access_key_id: "password",
aws_secret_access_key: "password",
aws_region: "text",
},
description: "SQS Queue (AWS) Logging Integration",
},
];
// Create callbackInfo object mapping display names to config objects
export const callbackInfo: Record<string, CallbackConfig> = CALLBACK_CONFIGS.reduce((acc, config) => {
acc[config.displayName] = config;
return acc;
}, {} as Record<string, CallbackConfig>);
export const callbackInfo: Record<string, CallbackConfig> = CALLBACK_CONFIGS.reduce(
(acc, config) => {
acc[config.displayName] = config;
return acc;
},
{} as Record<string, CallbackConfig>,
);
// Create callback_map mapping display names to internal IDs
export const callback_map: Record<string, string> = CALLBACK_CONFIGS.reduce((acc, config) => {
acc[config.displayName] = config.id;
return acc;
}, {} as Record<string, string>);
export const callback_map: Record<string, string> = CALLBACK_CONFIGS.reduce(
(acc, config) => {
acc[config.displayName] = config.id;
return acc;
},
{} as Record<string, string>,
);
// create reverse_callback_map to map internal IDs to display names
export const reverse_callback_map: Record<string, string> = CALLBACK_CONFIGS.reduce((acc, config) => {
acc[config.id] = config.displayName;
return acc;
}, {} as Record<string, string>);
export const reverse_callback_map: Record<string, string> = CALLBACK_CONFIGS.reduce(
(acc, config) => {
acc[config.id] = config.displayName;
return acc;
},
{} as Record<string, string>,
);
// Function to map display names to internal names
export const mapDisplayToInternalNames = (displayNames: string[]): string[] => {
return displayNames.map(name => callback_map[name] || name);
return displayNames.map((name) => callback_map[name] || name);
};
// Function to map internal names to display names
export const mapInternalToDisplayNames = (internalNames: string[]): string[] => {
return internalNames.map(name => reverse_callback_map[name] || name);
return internalNames.map((name) => reverse_callback_map[name] || name);
};
// Utility functions for easy access
export const getCallbackById = (id: string): CallbackConfig | undefined => {
return CALLBACK_CONFIGS.find(callback => callback.id === id);
return CALLBACK_CONFIGS.find((callback) => callback.id === id);
};
export const getCallbackByDisplayName = (displayName: string): CallbackConfig | undefined => {
return CALLBACK_CONFIGS.find(callback => callback.displayName === displayName);
return CALLBACK_CONFIGS.find((callback) => callback.displayName === displayName);
};
ui/litellm-dashboard/src/components/guardrails.tsx
+4 -2
View File
@@ -8,6 +8,7 @@ import { isAdminRole } from "@/utils/roles";
import GuardrailInfoView from "./guardrails/guardrail_info";
import GuardrailTestPlayground from "./guardrails/GuardrailTestPlayground";
import NotificationsManager from "./molecules/notifications_manager";
import { Guardrail, GuardrailDefinitionLocation } from "./guardrails/types";
interface GuardrailsPanelProps {
accessToken: string | null;
@@ -25,14 +26,15 @@ interface GuardrailItem {
guardrail_info: Record<string, any> | null;
created_at?: string;
updated_at?: string;
guardrail_definition_location: GuardrailDefinitionLocation;
}
interface GuardrailsResponse {
guardrails: GuardrailItem[];
guardrails: Guardrail[];
}
const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ accessToken, userRole }) => {
const [guardrailsList, setGuardrailsList] = useState<GuardrailItem[]>([]);
const [guardrailsList, setGuardrailsList] = useState<Guardrail[]>([]);
const [isAddModalVisible, setIsAddModalVisible] = useState(false);
const [isLoading, setIsLoading] = useState(false);
const [isDeleting, setIsDeleting] = useState(false);
ui/litellm-dashboard/src/components/guardrails/guardrail_table.test.tsx
+56
View File
@@ -0,0 +1,56 @@
import GuardrailTable from "./guardrail_table";
import { render } from "@testing-library/react";
import { describe, it, expect } from "vitest";
import { GuardrailDefinitionLocation } from "./types";
describe("GuardrailTable", () => {
it("should render", () => {
const { getByText } = render(
<GuardrailTable
guardrailsList={[]}
isLoading={false}
onDeleteClick={() => {}}
accessToken={null}
onGuardrailUpdated={() => {}}
onGuardrailClick={() => {}}
/>,
);
expect(getByText("Guardrail ID")).toBeInTheDocument();
expect(getByText("Name")).toBeInTheDocument();
expect(getByText("Provider")).toBeInTheDocument();
expect(getByText("Mode")).toBeInTheDocument();
expect(getByText("Default On")).toBeInTheDocument();
expect(getByText("Created At")).toBeInTheDocument();
expect(getByText("Updated At")).toBeInTheDocument();
});
it("should not allow deletion of config guardrails", () => {
const { getByTestId } = render(
<GuardrailTable
guardrailsList={[
{
guardrail_id: "1",
guardrail_name: "Guardrail 1",
litellm_params: { guardrail: "presidio", mode: "pre_call", default_on: true },
guardrail_info: null,
created_at: "2021-01-01",
updated_at: "2021-01-01",
guardrail_definition_location: GuardrailDefinitionLocation.CONFIG,
},
]}
isLoading={false}
onDeleteClick={() => {}}
accessToken={null}
onGuardrailUpdated={() => {}}
onGuardrailClick={() => {}}
/>,
);
const deleteGuardrailButton = getByTestId("config-delete-icon");
expect(deleteGuardrailButton).toBeInTheDocument();
expect(deleteGuardrailButton).toHaveClass("cursor-not-allowed text-gray-400");
expect(deleteGuardrailButton).toHaveAttribute(
"title",
"Config guardrail cannot be deleted on the dashboard. Please delete it from the config file.",
);
});
});
ui/litellm-dashboard/src/components/guardrails/guardrail_table.tsx
+29 -29
View File
@@ -13,24 +13,10 @@ import {
} from "@tanstack/react-table";
import { getGuardrailLogoAndName, guardrail_provider_map } from "./guardrail_info_helpers";
import EditGuardrailForm from "./edit_guardrail_form";
interface GuardrailItem {
guardrail_id?: string;
guardrail_name: string | null;
litellm_params: {
guardrail: string;
mode: string;
default_on: boolean;
pii_entities_config?: { [key: string]: string };
[key: string]: any;
};
guardrail_info: Record<string, any> | null;
created_at?: string;
updated_at?: string;
}
import { Guardrail, GuardrailDefinitionLocation } from "./types";
interface GuardrailTableProps {
guardrailsList: GuardrailItem[];
guardrailsList: Guardrail[];
isLoading: boolean;
onDeleteClick: (guardrailId: string, guardrailName: string) => void;
accessToken: string | null;
@@ -50,7 +36,7 @@ const GuardrailTable: React.FC<GuardrailTableProps> = ({
}) => {
const [sorting, setSorting] = useState<SortingState>([{ id: "created_at", desc: true }]);
const [editModalVisible, setEditModalVisible] = useState(false);
const [selectedGuardrail, setSelectedGuardrail] = useState<GuardrailItem | null>(null);
const [selectedGuardrail, setSelectedGuardrail] = useState<Guardrail | null>(null);
// Format date helper function
const formatDate = (dateString?: string) => {
@@ -59,7 +45,7 @@ const GuardrailTable: React.FC<GuardrailTableProps> = ({
return date.toLocaleString();
};
const handleEditClick = (guardrail: GuardrailItem) => {
const handleEditClick = (guardrail: Guardrail) => {
setSelectedGuardrail(guardrail);
setEditModalVisible(true);
};
@@ -70,7 +56,7 @@ const GuardrailTable: React.FC<GuardrailTableProps> = ({
onGuardrailUpdated();
};
const columns: ColumnDef<GuardrailItem>[] = [
const columns: ColumnDef<Guardrail>[] = [
{
header: "Guardrail ID",
accessorKey: "guardrail_id",
@@ -176,18 +162,32 @@ const GuardrailTable: React.FC<GuardrailTableProps> = ({
header: "",
cell: ({ row }) => {
const guardrail = row.original;
const isConfigGuardrail = guardrail.guardrail_definition_location === GuardrailDefinitionLocation.CONFIG;
return (
<div className="flex space-x-2">
<Icon
icon={TrashIcon}
size="sm"
onClick={() =>
guardrail.guardrail_id &&
onDeleteClick(guardrail.guardrail_id, guardrail.guardrail_name || "Unnamed Guardrail")
}
className="cursor-pointer hover:text-red-500"
tooltip="Delete guardrail"
/>
{isConfigGuardrail ? (
<Tooltip title="Config guardrail cannot be deleted on the dashboard. Please delete it from the config file.">
<Icon
data-testid="config-delete-icon"
icon={TrashIcon}
size="sm"
className="cursor-not-allowed text-gray-400"
title="Config guardrail cannot be deleted on the dashboard. Please delete it from the config file."
aria-label="Delete guardrail (config)"
/>
</Tooltip>
) : (
<Icon
icon={TrashIcon}
size="sm"
onClick={() =>
guardrail.guardrail_id &&
onDeleteClick(guardrail.guardrail_id, guardrail.guardrail_name || "Unnamed Guardrail")
}
className="cursor-pointer hover:text-red-500"
tooltip="Delete guardrail"
/>
)}
</div>
);
},
ui/litellm-dashboard/src/components/guardrails/pii_components.test.tsx
+40
View File
@@ -0,0 +1,40 @@
import { render } from "@testing-library/react";
import { describe, it, expect } from "vitest";
import { CategoryFilter, QuickActions, PiiEntityList } from "./pii_components";
import type { PiiEntityCategory } from "./types";
describe("CategoryFilter", () => {
it("should render", () => {
const emptyCategories: PiiEntityCategory[] = [];
const { getByText } = render(
<CategoryFilter categories={emptyCategories} selectedCategories={[]} onChange={() => {}} />,
);
expect(getByText("Filter by category")).toBeInTheDocument();
});
});
describe("QuickActions", () => {
it("should render", () => {
const { getByText } = render(
<QuickActions onSelectAll={() => {}} onUnselectAll={() => {}} hasSelectedEntities={false} />,
);
expect(getByText("Quick Actions")).toBeInTheDocument();
});
});
describe("PiiEntityList", () => {
it("should render", () => {
const { getByText } = render(
<PiiEntityList
entities={[]}
selectedEntities={[]}
selectedActions={{}}
actions={[]}
onEntitySelect={() => {}}
onActionSelect={() => {}}
entityToCategoryMap={new Map()}
/>,
);
expect(getByText("No PII types match your filter criteria")).toBeInTheDocument();
});
});
ui/litellm-dashboard/src/components/guardrails/pii_components.tsx
+2
View File
@@ -83,6 +83,7 @@ export const QuickActions: React.FC<QuickActionsProps> = ({ onSelectAll, onUnsel
</div>
<Button
type="default"
danger
onClick={onUnselectAll}
disabled={!hasSelectedEntities}
icon={<CloseOutlined />}
@@ -103,6 +104,7 @@ export const QuickActions: React.FC<QuickActionsProps> = ({ onSelectAll, onUnsel
</Button>
<Button
type="default"
danger
onClick={() => onSelectAll("BLOCK")}
className="flex items-center justify-center h-10 border-red-200 hover:border-red-300 hover:text-red-700 bg-red-50 hover:bg-red-100 text-red-600"
block
ui/litellm-dashboard/src/components/guardrails/pii_configuration.test.tsx
+20
View File
@@ -0,0 +1,20 @@
import { render } from "@testing-library/react";
import { describe, it, expect } from "vitest";
import PiiConfiguration from "./pii_configuration";
describe("PiiConfiguration", () => {
it("should render", () => {
const { getByText } = render(
<PiiConfiguration
entities={[]}
actions={[]}
selectedEntities={[]}
selectedActions={{}}
onEntitySelect={() => {}}
onActionSelect={() => {}}
entityCategories={[]}
/>,
);
expect(getByText("Configure PII Protection")).toBeInTheDocument();
});
});
ui/litellm-dashboard/src/components/guardrails/pii_configuration.tsx
+4 -11
View File
@@ -1,7 +1,7 @@
import { Typography } from "antd";
import React, { useState } from "react";
import { Typography, Badge } from "antd";
import { CategoryFilter, PiiEntityList, QuickActions } from "./pii_components";
import { PiiConfigurationProps } from "./types";
import { CategoryFilter, QuickActions, PiiEntityList } from "./pii_components";
const { Title, Text } = Typography;
@@ -57,18 +57,11 @@ const PiiConfiguration: React.FC<PiiConfigurationProps> = ({
<div className="pii-configuration">
<div className="flex justify-between items-center mb-5">
<div className="flex items-center">
<Title level={4} className="mb-0 font-semibold text-gray-800">
<Title level={4} className="!m-0 font-semibold text-gray-800">
Configure PII Protection
</Title>
</div>
<Badge
count={selectedEntities.length}
showZero
style={{ backgroundColor: selectedEntities.length > 0 ? "#4f46e5" : "#d9d9d9" }}
overflowCount={999}
>
<Text className="text-gray-500">{selectedEntities.length} items selected</Text>
</Badge>
<Text className="text-gray-500">{selectedEntities.length} items selected</Text>
</div>
<div className="mb-6">
ui/litellm-dashboard/src/components/guardrails/types.ts
+6
View File
@@ -31,4 +31,10 @@ export interface Guardrail {
guardrail_info: Record<string, any> | null;
created_at?: string;
updated_at?: string;
guardrail_definition_location: GuardrailDefinitionLocation;
}
export enum GuardrailDefinitionLocation {
DB = "db",
CONFIG = "config",
}
ui/litellm-dashboard/src/components/settings.tsx
+3 -3
View File
@@ -565,7 +565,7 @@ const Settings: React.FC<SettingsPageProps> = ({ accessToken, userRole, userID,
<Modal
title="Add Logging Callback"
visible={showAddCallbacksModal}
open={showAddCallbacksModal}
width={800}
onCancel={() => {
setShowAddCallbacksModal(false);
@@ -682,7 +682,7 @@ const Settings: React.FC<SettingsPageProps> = ({ accessToken, userRole, userID,
)}
<div className="flex justify-end space-x-3 pt-6 mt-6 border-t border-gray-200">
<Button
<Button2
onClick={() => {
setShowAddCallbacksModal(false);
setSelectedCallback(null);
@@ -691,7 +691,7 @@ const Settings: React.FC<SettingsPageProps> = ({ accessToken, userRole, userID,
}}
>
Cancel
</Button>
</Button2>
<Button2 htmlType="submit">Add Callback</Button2>
</div>
</Form>