diff --git a/docs/my-website/docs/completion/image_generation_chat.md b/docs/my-website/docs/completion/image_generation_chat.md index 98b718ef4c..5538b7f8ff 100644 --- a/docs/my-website/docs/completion/image_generation_chat.md +++ b/docs/my-website/docs/completion/image_generation_chat.md @@ -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 diff --git a/docs/my-website/docs/image_edits.md b/docs/my-website/docs/image_edits.md index 84dddd5e4a..9a53da510f 100644 --- a/docs/my-website/docs/image_edits.md +++ b/docs/my-website/docs/image_edits.md @@ -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}") ``` +``` + + + + + +#### 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)) +``` + @@ -224,6 +272,36 @@ curl -X POST "http://localhost:4000/v1/images/edits" \ -F "response_format=url" ``` +``` + + + + + +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 " \ + -F "model=gemini-image-edit" \ + -F "image=@original_image.png" \ + -F "prompt=Add a warm golden-hour glow to the scene" \ + -F "size=1024x1024" +``` + diff --git a/docs/my-website/docs/providers/gemini.md b/docs/my-website/docs/providers/gemini.md index 31d3a491f4..c5014fc2ff 100644 --- a/docs/my-website/docs/providers/gemini.md +++ b/docs/my-website/docs/providers/gemini.md @@ -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) |
@@ -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 | | @@ -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?"}], diff --git a/docs/my-website/docs/providers/openai.md b/docs/my-website/docs/providers/openai.md index f9831c6d8b..51ebc881d2 100644 --- a/docs/my-website/docs/providers/openai.md +++ b/docs/my-website/docs/providers/openai.md @@ -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. + + + +```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) +) +``` + + + +```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"} +}' +``` + + + +**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. diff --git a/docs/my-website/docs/providers/runwayml/text-to-speech.md b/docs/my-website/docs/providers/runwayml/text-to-speech.md new file mode 100644 index 0000000000..020269863c --- /dev/null +++ b/docs/my-website/docs/providers/runwayml/text-to-speech.md @@ -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
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) +``` + diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 874b637e4d..4d7e85f388 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -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 - - - - -```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) -``` - - - - - -1. Add model to config.yaml -```yaml -model_list: - - model_name: snowflake-arctic-embed-m-long-1731622468876 - litellm_params: - model: vertex_ai/ - 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) -``` - - - - - -#### 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/", 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 - - - - - -```python -response = litellm.embedding( - model="vertex_ai/text-embedding-004", - input=["good morning from litellm", "gm"] - input_type = "RETRIEVAL_DOCUMENT", - dimensions=1, -) -``` - - - - -```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) -``` - - - - -### 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) - - - - -```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, -) -``` - - - - -```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) -``` - - - -## **Multi-Modal Embeddings** - - -Known Limitations: -- Only supports 1 image / video / image per request -- Only supports GCS or base64 encoded images / videos - -### Usage - - - - -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 -) -``` - - - - -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 - - - - - - -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) -``` - - - - - -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) - -``` - - - - - - - - - -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}") -``` - - - - - -### Text + Image + Video Embeddings - - - - -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 -) -``` - - - - - -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) -``` - - - - - ## **Gemini TTS (Text-to-Speech) Audio Output** :::info diff --git a/docs/my-website/docs/providers/vertex_embedding.md b/docs/my-website/docs/providers/vertex_embedding.md new file mode 100644 index 0000000000..5656ade337 --- /dev/null +++ b/docs/my-website/docs/providers/vertex_embedding.md @@ -0,0 +1,587 @@ +import Image from '@theme/IdealImage'; +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Vertex AI Embedding + +## Usage - Embedding + + + + +```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) +``` + + + + + +1. Add model to config.yaml +```yaml +model_list: + - model_name: snowflake-arctic-embed-m-long-1731622468876 + litellm_params: + model: vertex_ai/ + 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) +``` + + + + + +#### 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/", 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 + + + + + +```python +response = litellm.embedding( + model="vertex_ai/text-embedding-004", + input=["good morning from litellm", "gm"] + input_type = "RETRIEVAL_DOCUMENT", + dimensions=1, +) +``` + + + + +```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) +``` + + + + +### 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) + + + + +```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, +) +``` + + + + +```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) +``` + + + +## **BGE Embeddings** + +Use BGE (Baidu General Embedding) models deployed on Vertex AI. + +### Usage + + + + +```python showLineNumbers title="Using BGE on Vertex AI" +import litellm + +response = litellm.embedding( + model="vertex_ai/bge/", + input=["Hello", "World"], + vertex_project="your-project-id", + vertex_location="your-location" +) + +print(response) +``` + + + + + +1. Add model to config.yaml +```yaml showLineNumbers title="config.yaml" +model_list: + - model_name: bge-embedding + litellm_params: + model: vertex_ai/bge/ + 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 +``` + + + + +## **Multi-Modal Embeddings** + + +Known Limitations: +- Only supports 1 image / video / image per request +- Only supports GCS or base64 encoded images / videos + +### Usage + + + + +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 +) +``` + + + + +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 + + + + + + +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) +``` + + + + + +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) + +``` + + + + + + + + + +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}") +``` + + + + + +### Text + Image + Video Embeddings + + + + +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 +) +``` + + + + + +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) +``` + + + \ No newline at end of file diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 31aa38c033..4d02d5729b 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -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) \ No newline at end of file +| 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 \ No newline at end of file diff --git a/docs/my-website/docs/proxy/docker_quick_start.md b/docs/my-website/docs/proxy/docker_quick_start.md index 7e380e8308..d82a0b01d1 100644 --- a/docs/my-website/docs/proxy/docker_quick_start.md +++ b/docs/my-website/docs/proxy/docker_quick_start.md @@ -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 # / + ───── ─────────────────── + │ │ + │ └─────▶ 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** (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 diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 0b6c31cc1f..809e0bc0b7 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -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", diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index d4a4c441eb..d1c7ede655 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -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 diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index b353c3670f..97a4c5723d 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -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)) diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 1a43ff2e17..f8c786daef 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -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, diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index b55065352d..eff5376e49 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -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, ) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index fab2c1e76e..a21ebd56f6 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -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}" ) diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 677bd91f98..7c65cad94d 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -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= - 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 - diff --git a/litellm/llms/gemini/image_edit/__init__.py b/litellm/llms/gemini/image_edit/__init__.py new file mode 100644 index 0000000000..6181015b81 --- /dev/null +++ b/litellm/llms/gemini/image_edit/__init__.py @@ -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() + diff --git a/litellm/llms/gemini/image_edit/cost_calculator.py b/litellm/llms/gemini/image_edit/cost_calculator.py new file mode 100644 index 0000000000..31f35345d8 --- /dev/null +++ b/litellm/llms/gemini/image_edit/cost_calculator.py @@ -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 + diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py new file mode 100644 index 0000000000..830c58a006 --- /dev/null +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -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.") \ No newline at end of file diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index f136bd0a40..d47759d0e8 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -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", {}) diff --git a/litellm/llms/runwayml/text_to_speech/__init__.py b/litellm/llms/runwayml/text_to_speech/__init__.py new file mode 100644 index 0000000000..491e8449e0 --- /dev/null +++ b/litellm/llms/runwayml/text_to_speech/__init__.py @@ -0,0 +1,5 @@ +"""RunwayML Text-to-Speech implementation.""" +from .transformation import RunwayMLTextToSpeechConfig + +__all__ = ["RunwayMLTextToSpeechConfig"] + diff --git a/litellm/llms/runwayml/text_to_speech/transformation.py b/litellm/llms/runwayml/text_to_speech/transformation.py new file mode 100644 index 0000000000..ac926beb22 --- /dev/null +++ b/litellm/llms/runwayml/text_to_speech/transformation.py @@ -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) + diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index c45a1cd60d..651acff6fc 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -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): diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 7932881f48..b40f0a72a5 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -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 = { diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 2c53457736..02b0f79280 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -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}" diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 70b068b5a4..dabc620a6d 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -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( diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index b8370d5fef..cbd8cf320c 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -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 diff --git a/litellm/llms/vertex_ai/vertex_embeddings/bge.py b/litellm/llms/vertex_ai/vertex_embeddings/bge.py new file mode 100644 index 0000000000..2eff0ba96d --- /dev/null +++ b/litellm/llms/vertex_ai/vertex_embeddings/bge.py @@ -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 + diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index 97af558041..5a3a4a7188 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -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"] diff --git a/litellm/llms/vertex_ai/vertex_embeddings/types.py b/litellm/llms/vertex_ai/vertex_embeddings/types.py index 7f85ea46f3..fa9794d79a 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/types.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/types.py @@ -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]] diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/main.py b/litellm/llms/vertex_ai/vertex_gemma_models/main.py index 8203b285eb..41bd6b5431 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/main.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/main.py @@ -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( diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 9ddbc461a7..ce50bf311e 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -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( diff --git a/litellm/llms/vertex_ai/vertex_model_garden/main.py b/litellm/llms/vertex_ai/vertex_model_garden/main.py index 1c57096734..fe7d0862e0 100644 --- a/litellm/llms/vertex_ai/vertex_model_garden/main.py +++ b/litellm/llms/vertex_ai/vertex_model_garden/main.py @@ -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( diff --git a/litellm/main.py b/litellm/main.py index 2ad444a9a2..23ffa90e2d 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index aa7774c43e..fa36e2d608 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7cd6546e9c..aefbbc8d4a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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]: """ diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 6f4dafb5fe..ce29d2d32e 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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] diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index fb7ada8ab1..4373f84292 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -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] \ No newline at end of file diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 1af6c9cf28..2226e19090 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -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", []): diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 2f101b13d8..3eee47f201 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -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", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f06b9e70f4..2077955fc5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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"][ diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 51ece5dd84..b60acd1428 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -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 diff --git a/litellm/utils.py b/litellm/utils.py index 4b9c1d9051..ae3de67374 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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 diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index aa7774c43e..fa36e2d608 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 6c3be897f4..4d9218d609 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1368,7 +1368,7 @@ "embeddings": false, "image_generations": true, "audio_transcriptions": false, - "audio_speech": false, + "audio_speech": true, "moderations": false, "batches": false, "rerank": false, diff --git a/tests/audio_tests/runwayml_speech.mp3 b/tests/audio_tests/runwayml_speech.mp3 new file mode 100644 index 0000000000..22aee2084b Binary files /dev/null and b/tests/audio_tests/runwayml_speech.mp3 differ diff --git a/tests/audio_tests/test_audio_speech.py b/tests/audio_tests/test_audio_speech.py index bf886a8d5c..409e3544fc 100644 --- a/tests/audio_tests/test_audio_speech.py +++ b/tests/audio_tests/test_audio_speech.py @@ -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(): """ diff --git a/tests/llm_translation/test_bedrock_agentcore.py b/tests/llm_translation/test_bedrock_agentcore.py index 6211fdbcdb..6dc6215a5e 100644 --- a/tests/llm_translation/test_bedrock_agentcore.py +++ b/tests/llm_translation/test_bedrock_agentcore.py @@ -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" + diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index f4100d9e8a..1065509dd4 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -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 diff --git a/tests/logging_callback_tests/test_sqs_logger.py b/tests/logging_callback_tests/test_sqs_logger.py index ed96334228..8a05ac6d0c 100644 --- a/tests/logging_callback_tests/test_sqs_logger.py +++ b/tests/logging_callback_tests/test_sqs_logger.py @@ -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 "") diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index fbdada9306..8391e8f77e 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -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" + ) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index c41ab391df..8b9c515234 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -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 diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 2487479a8c..d7da759287 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -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""" diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 5d17ea3dc3..8cd623267b 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -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}" + diff --git a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py new file mode 100644 index 0000000000..2732bf1595 --- /dev/null +++ b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py @@ -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={}, + ) + diff --git a/tests/test_litellm/llms/runwayml/test_text_to_speech_transformation.py b/tests/test_litellm/llms/runwayml/test_text_to_speech_transformation.py new file mode 100644 index 0000000000..277a47a03b --- /dev/null +++ b/tests/test_litellm/llms/runwayml/test_text_to_speech_transformation.py @@ -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 + diff --git a/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py new file mode 100644 index 0000000000..156ab95184 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py @@ -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 + + diff --git a/tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py b/tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py new file mode 100644 index 0000000000..20150501ad --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py @@ -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 + ) + diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py new file mode 100644 index 0000000000..46f365094c --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py @@ -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}" + diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index b9ed4b9b50..42290a024b 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -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"] + diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 2502b4e34e..6939a19b7e 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -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() + diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index c9c602e1cc..3a9229b3c4 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -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 diff --git a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx index 22f7989a13..7abe255acc 100644 --- a/ui/litellm-dashboard/src/components/callback_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/callback_info_helpers.tsx @@ -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 = CALLBACK_CONFIGS.reduce((acc, config) => { - acc[config.displayName] = config; - return acc; -}, {} as Record); +export const callbackInfo: Record = CALLBACK_CONFIGS.reduce( + (acc, config) => { + acc[config.displayName] = config; + return acc; + }, + {} as Record, +); // Create callback_map mapping display names to internal IDs -export const callback_map: Record = CALLBACK_CONFIGS.reduce((acc, config) => { - acc[config.displayName] = config.id; - return acc; -}, {} as Record); +export const callback_map: Record = CALLBACK_CONFIGS.reduce( + (acc, config) => { + acc[config.displayName] = config.id; + return acc; + }, + {} as Record, +); // create reverse_callback_map to map internal IDs to display names -export const reverse_callback_map: Record = CALLBACK_CONFIGS.reduce((acc, config) => { - acc[config.id] = config.displayName; - return acc; -}, {} as Record); +export const reverse_callback_map: Record = CALLBACK_CONFIGS.reduce( + (acc, config) => { + acc[config.id] = config.displayName; + return acc; + }, + {} as Record, +); // 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); }; diff --git a/ui/litellm-dashboard/src/components/guardrails.tsx b/ui/litellm-dashboard/src/components/guardrails.tsx index 23aec34c40..3861545f8c 100644 --- a/ui/litellm-dashboard/src/components/guardrails.tsx +++ b/ui/litellm-dashboard/src/components/guardrails.tsx @@ -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 | null; created_at?: string; updated_at?: string; + guardrail_definition_location: GuardrailDefinitionLocation; } interface GuardrailsResponse { - guardrails: GuardrailItem[]; + guardrails: Guardrail[]; } const GuardrailsPanel: React.FC = ({ accessToken, userRole }) => { - const [guardrailsList, setGuardrailsList] = useState([]); + const [guardrailsList, setGuardrailsList] = useState([]); const [isAddModalVisible, setIsAddModalVisible] = useState(false); const [isLoading, setIsLoading] = useState(false); const [isDeleting, setIsDeleting] = useState(false); diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_table.test.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_table.test.tsx new file mode 100644 index 0000000000..dcc6767735 --- /dev/null +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_table.test.tsx @@ -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( + {}} + 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( + {}} + 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.", + ); + }); +}); diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_table.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_table.tsx index 4857cdd31e..6a2476960f 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_table.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_table.tsx @@ -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 | 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 = ({ }) => { const [sorting, setSorting] = useState([{ id: "created_at", desc: true }]); const [editModalVisible, setEditModalVisible] = useState(false); - const [selectedGuardrail, setSelectedGuardrail] = useState(null); + const [selectedGuardrail, setSelectedGuardrail] = useState(null); // Format date helper function const formatDate = (dateString?: string) => { @@ -59,7 +45,7 @@ const GuardrailTable: React.FC = ({ return date.toLocaleString(); }; - const handleEditClick = (guardrail: GuardrailItem) => { + const handleEditClick = (guardrail: Guardrail) => { setSelectedGuardrail(guardrail); setEditModalVisible(true); }; @@ -70,7 +56,7 @@ const GuardrailTable: React.FC = ({ onGuardrailUpdated(); }; - const columns: ColumnDef[] = [ + const columns: ColumnDef[] = [ { header: "Guardrail ID", accessorKey: "guardrail_id", @@ -176,18 +162,32 @@ const GuardrailTable: React.FC = ({ header: "", cell: ({ row }) => { const guardrail = row.original; + const isConfigGuardrail = guardrail.guardrail_definition_location === GuardrailDefinitionLocation.CONFIG; return (
- - guardrail.guardrail_id && - onDeleteClick(guardrail.guardrail_id, guardrail.guardrail_name || "Unnamed Guardrail") - } - className="cursor-pointer hover:text-red-500" - tooltip="Delete guardrail" - /> + {isConfigGuardrail ? ( + + + + ) : ( + + guardrail.guardrail_id && + onDeleteClick(guardrail.guardrail_id, guardrail.guardrail_name || "Unnamed Guardrail") + } + className="cursor-pointer hover:text-red-500" + tooltip="Delete guardrail" + /> + )}
); }, diff --git a/ui/litellm-dashboard/src/components/guardrails/pii_components.test.tsx b/ui/litellm-dashboard/src/components/guardrails/pii_components.test.tsx new file mode 100644 index 0000000000..3accbf13bd --- /dev/null +++ b/ui/litellm-dashboard/src/components/guardrails/pii_components.test.tsx @@ -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( + {}} />, + ); + expect(getByText("Filter by category")).toBeInTheDocument(); + }); +}); + +describe("QuickActions", () => { + it("should render", () => { + const { getByText } = render( + {}} onUnselectAll={() => {}} hasSelectedEntities={false} />, + ); + expect(getByText("Quick Actions")).toBeInTheDocument(); + }); +}); + +describe("PiiEntityList", () => { + it("should render", () => { + const { getByText } = render( + {}} + onActionSelect={() => {}} + entityToCategoryMap={new Map()} + />, + ); + expect(getByText("No PII types match your filter criteria")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/guardrails/pii_components.tsx b/ui/litellm-dashboard/src/components/guardrails/pii_components.tsx index 1b744cad31..e3b3926b99 100644 --- a/ui/litellm-dashboard/src/components/guardrails/pii_components.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/pii_components.tsx @@ -83,6 +83,7 @@ export const QuickActions: React.FC = ({ onSelectAll, onUnsel + Add Callback