diff --git a/cookbook/ai_coding_tool_guides/claude_code_quickstart/guide.md b/cookbook/ai_coding_tool_guides/claude_code_quickstart/guide.md index ad86c2b7b1..3d6c75498b 100644 --- a/cookbook/ai_coding_tool_guides/claude_code_quickstart/guide.md +++ b/cookbook/ai_coding_tool_guides/claude_code_quickstart/guide.md @@ -97,17 +97,75 @@ export ANTHROPIC_AUTH_TOKEN="$LITELLM_MASTER_KEY" ## Step 5: Use Claude Code -Start Claude Code and it will automatically use your configured models: +### Choosing Your Model + +You have two options for specifying which model Claude Code uses: + +#### Option 1: Command Line / Session Model Selection + +Specify the model directly when starting Claude Code or during a session: ```bash -# Claude Code will use the models configured in your LiteLLM proxy -claude - -# Or specify a model if you have multiple configured +# Specify model at startup claude --model claude-3-5-sonnet-20241022 -claude --model claude-3-5-haiku-20241022 + +# Or change model during a session +/model claude-3-5-haiku-20241022 ``` +This method uses the exact model you specify. + +#### Option 2: Environment Variables + +Configure default models using environment variables: + +```bash +# Tell Claude Code which models to use by default +export ANTHROPIC_DEFAULT_SONNET_MODEL=claude-3-5-sonnet-20241022 +export ANTHROPIC_DEFAULT_HAIKU_MODEL=claude-3-5-haiku-20241022 +export ANTHROPIC_DEFAULT_OPUS_MODEL=claude-opus-3-5-20240229 + +claude # Will use the models specified above +``` + +**Note:** Claude Code may cache the model from a previous session. If environment variables don't take effect, use Option 1 to explicitly set the model. + +**Important:** The `model_name` in your LiteLLM config must match what Claude Code requests (either from env vars or command line). + +### Using 1M Context Window + +Claude Code supports extended context (1 million tokens) using the `[1m]` suffix with Claude 4+ models: + +```bash +# Use Sonnet 4.5 with 1M context (requires quotes for shell) +claude --model 'claude-sonnet-4-5-20250929[1m]' + +# Inside a Claude Code session (no quotes needed) +/model claude-sonnet-4-5-20250929[1m] +``` + +**Important:** When using `--model` with `[1m]` in the shell, you must use quotes to prevent the shell from interpreting the brackets. + +Alternatively, set as default with environment variables: + +```bash +export ANTHROPIC_DEFAULT_SONNET_MODEL='claude-sonnet-4-5-20250929[1m]' +claude +``` + +**How it works:** +- Claude Code strips the `[1m]` suffix before sending to LiteLLM +- Claude Code automatically adds the header `anthropic-beta: context-1m-2025-08-07` +- Your LiteLLM config should **NOT** include `[1m]` in model names + +**Verify 1M context is active:** +```bash +/context +# Should show: 21k/1000k tokens (2%) +``` + +**Pricing:** Models using 1M context have different pricing. Input tokens above 200k are charged at a higher rate. + ## Troubleshooting Common issues and solutions: @@ -123,18 +181,25 @@ Common issues and solutions: - Ensure the `ANTHROPIC_AUTH_TOKEN` matches your LiteLLM master key **Model not found:** -- Ensure the model name in Claude Code matches exactly with your `config.yaml` -- Check LiteLLM logs for detailed error messages +- Check what model Claude Code is requesting in LiteLLM logs +- Ensure your `config.yaml` has a matching `model_name` entry +- If using environment variables, verify they're set: `echo $ANTHROPIC_DEFAULT_SONNET_MODEL` + +**1M context not working (showing 200k instead of 1000k):** +- Verify you're using the `[1m]` suffix: `/model your-model-name[1m]` +- Check LiteLLM logs for the header `context-1m-2025-08-07` in the request +- Ensure your model supports 1M context (only certain Claude models do) +- Your LiteLLM config should **NOT** include `[1m]` in the `model_name` ## Using Multiple Models and Providers -Expand your configuration to support multiple providers and models: +You can configure LiteLLM to route to any supported provider. Here's an example with multiple providers: ```yaml model_list: # OpenAI models - model_name: codex-mini - litellm_params: + litellm_params: model: openai/codex-mini api_key: os.environ/OPENAI_API_KEY api_base: https://api.openai.com/v1 @@ -156,7 +221,7 @@ model_list: litellm_params: model: anthropic/claude-3-5-sonnet-20241022 api_key: os.environ/ANTHROPIC_API_KEY - + - model_name: claude-3-5-haiku-20241022 litellm_params: model: anthropic/claude-3-5-haiku-20241022 @@ -174,19 +239,54 @@ litellm_settings: master_key: os.environ/LITELLM_MASTER_KEY ``` +**Note:** The `model_name` can be anything you choose. Claude Code will request whatever model you specify (via env vars or command line), and LiteLLM will route to the `model` configured in `litellm_params`. + Switch between models seamlessly: ```bash -# Use Claude for complex reasoning -claude --model claude-3-5-sonnet-20241022 +# Use environment variables to set defaults +export ANTHROPIC_DEFAULT_SONNET_MODEL=claude-3-5-sonnet-20241022 +export ANTHROPIC_DEFAULT_HAIKU_MODEL=claude-3-5-haiku-20241022 -# Use Haiku for fast responses -claude --model claude-3-5-haiku-20241022 - -# Use Bedrock deployment -claude --model claude-bedrock +# Or specify directly +claude --model claude-3-5-sonnet-20241022 # Complex reasoning +claude --model claude-3-5-haiku-20241022 # Fast responses +claude --model claude-bedrock # Bedrock deployment ``` +## Default Models Used by Claude Code + +If you **don't** set environment variables, Claude Code uses these default model names: + +| Purpose | Default Model Name (v2.1.14) | +|---------|------------------------------| +| Main model | `claude-sonnet-4-5-20250929` | +| Light tasks (subagents, summaries) | `claude-haiku-4-5-20251001` | +| Planning mode | `claude-opus-4-5-20251101` | + +Your LiteLLM config should include these model names if you want Claude Code to work without setting environment variables: + +```yaml +model_list: + - model_name: claude-sonnet-4-5-20250929 + litellm_params: + # Can be any provider - Anthropic, Bedrock, Vertex AI, etc. + model: anthropic/claude-sonnet-4-5-20250929 + api_key: os.environ/ANTHROPIC_API_KEY + + - model_name: claude-haiku-4-5-20251001 + litellm_params: + model: anthropic/claude-haiku-4-5-20251001 + api_key: os.environ/ANTHROPIC_API_KEY + + - model_name: claude-opus-4-5-20251101 + litellm_params: + model: anthropic/claude-opus-4-5-20251101 + api_key: os.environ/ANTHROPIC_API_KEY +``` + +**Warning:** These default model names may change with new Claude Code versions. Check LiteLLM proxy logs for "model not found" errors to identify what Claude Code is requesting. + ## Additional Resources - [LiteLLM Documentation](https://docs.litellm.ai/) diff --git a/docs/my-website/docs/anthropic_unified.md b/docs/my-website/docs/anthropic_unified/index.md similarity index 100% rename from docs/my-website/docs/anthropic_unified.md rename to docs/my-website/docs/anthropic_unified/index.md diff --git a/docs/my-website/docs/anthropic_unified/structured_output.md b/docs/my-website/docs/anthropic_unified/structured_output.md new file mode 100644 index 0000000000..433f57537d --- /dev/null +++ b/docs/my-website/docs/anthropic_unified/structured_output.md @@ -0,0 +1,237 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Structured Output /v1/messages + +Use LiteLLM to call Anthropic's structured output feature via the `/v1/messages` endpoint. + +## Supported Providers + +| Provider | Supported | Notes | +|----------|-----------|-------| +| Anthropic | ✅ | Native support | +| Azure AI (Anthropic models) | ✅ | Claude models on Azure AI | +| Bedrock (Converse Anthropic models) | ✅ | Claude models via Bedrock Converse API | + +## Usage + +### LiteLLM Proxy Server + + + + +1. Setup config.yaml + +```yaml +model_list: + - model_name: claude-sonnet + litellm_params: + model: anthropic/claude-sonnet-4-5-20250514 + api_key: os.environ/ANTHROPIC_API_KEY +``` + +2. Start proxy + +```bash +litellm --config /path/to/config.yaml +``` + +3. Test it! + +```bash +curl http://localhost:4000/v1/messages \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -H "anthropic-version: 2023-06-01" \ + -d '{ + "model": "claude-sonnet", + "max_tokens": 1024, + "messages": [ + { + "role": "user", + "content": "Extract the key information from this email: John Smith (john@example.com) is interested in our Enterprise plan and wants to schedule a demo for next Tuesday at 2pm." + } + ], + "output_format": { + "type": "json_schema", + "schema": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "email": {"type": "string"}, + "plan_interest": {"type": "string"}, + "demo_requested": {"type": "boolean"} + }, + "required": ["name", "email", "plan_interest", "demo_requested"], + "additionalProperties": false + } + } + }' +``` + + + + + +1. Setup config.yaml + +```yaml +model_list: + - model_name: azure-claude-sonnet + litellm_params: + model: azure_ai/claude-sonnet-4-5-20250514 + api_key: os.environ/AZURE_AI_API_KEY + api_base: https://your-endpoint.inference.ai.azure.com +``` + +2. Start proxy + +```bash +litellm --config /path/to/config.yaml +``` + +3. Test it! + +```bash +curl http://localhost:4000/v1/messages \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -H "anthropic-version: 2023-06-01" \ + -d '{ + "model": "azure-claude-sonnet", + "max_tokens": 1024, + "messages": [ + { + "role": "user", + "content": "Extract the key information from this email: John Smith (john@example.com) is interested in our Enterprise plan and wants to schedule a demo for next Tuesday at 2pm." + } + ], + "output_format": { + "type": "json_schema", + "schema": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "email": {"type": "string"}, + "plan_interest": {"type": "string"}, + "demo_requested": {"type": "boolean"} + }, + "required": ["name", "email", "plan_interest", "demo_requested"], + "additionalProperties": false + } + } + }' +``` + + + + + +1. Setup config.yaml + +```yaml +model_list: + - model_name: bedrock-claude-sonnet + litellm_params: + model: bedrock/anthropic.claude-sonnet-4-5-20250514-v1:0 + aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + aws_region_name: us-west-2 +``` + +2. Start proxy + +```bash +litellm --config /path/to/config.yaml +``` + +3. Test it! + +```bash +curl http://localhost:4000/v1/messages \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -H "anthropic-version: 2023-06-01" \ + -d '{ + "model": "bedrock-claude-sonnet", + "max_tokens": 1024, + "messages": [ + { + "role": "user", + "content": "Extract the key information from this email: John Smith (john@example.com) is interested in our Enterprise plan and wants to schedule a demo for next Tuesday at 2pm." + } + ], + "output_format": { + "type": "json_schema", + "schema": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "email": {"type": "string"}, + "plan_interest": {"type": "string"}, + "demo_requested": {"type": "boolean"} + }, + "required": ["name", "email", "plan_interest", "demo_requested"], + "additionalProperties": false + } + } + }' +``` + + + + +## Example Response + +```json +{ + "id": "msg_01XFDUDYJgAACzvnptvVoYEL", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "text", + "text": "{\"name\":\"John Smith\",\"email\":\"john@example.com\",\"plan_interest\":\"Enterprise\",\"demo_requested\":true}" + } + ], + "model": "claude-sonnet-4-5-20250514", + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 75, + "output_tokens": 28 + } +} +``` + +## Request Format + +### output_format + +The `output_format` parameter specifies the structured output format. + +```json +{ + "output_format": { + "type": "json_schema", + "schema": { + "type": "object", + "properties": { + "field_name": {"type": "string"}, + "another_field": {"type": "integer"} + }, + "required": ["field_name", "another_field"], + "additionalProperties": false + } + } +} +``` + +#### Fields + +- **type** (string): Must be `"json_schema"` +- **schema** (object): A JSON Schema object defining the expected output structure + - **type** (string): The root type, typically `"object"` + - **properties** (object): Defines the fields and their types + - **required** (array): List of required field names + - **additionalProperties** (boolean): Set to `false` to enforce strict schema adherence diff --git a/docs/my-website/docs/pass_through/vertex_ai.md b/docs/my-website/docs/pass_through/vertex_ai.md index adbd06187d..00df6def70 100644 --- a/docs/my-website/docs/pass_through/vertex_ai.md +++ b/docs/my-website/docs/pass_through/vertex_ai.md @@ -461,3 +461,48 @@ generateContent(); + +### Using Anthropic Beta Features on Vertex AI + +When using Anthropic models via Vertex AI passthrough (e.g., Claude on Vertex), you can enable Anthropic beta features like extended context windows. + +The `anthropic-beta` header is automatically forwarded to Vertex AI when calling Anthropic models. + +```bash +curl http://localhost:4000/vertex_ai/v1/projects/${PROJECT_ID}/locations/us-east5/publishers/anthropic/models/claude-3-5-sonnet:rawPredict \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -H "anthropic-beta: context-1m-2025-08-07" \ + -d '{ + "anthropic_version": "vertex-2023-10-16", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 500 + }' +``` + +### Forwarding Custom Headers with `x-pass-` Prefix + +You can forward any custom header to the provider by prefixing it with `x-pass-`. The prefix is stripped before the header is sent to the provider. + +For example: +- `x-pass-anthropic-beta: value` becomes `anthropic-beta: value` +- `x-pass-custom-header: value` becomes `custom-header: value` + +This is useful when you need to send provider-specific headers that aren't in the default allowlist. + +```bash +curl http://localhost:4000/vertex_ai/v1/projects/${PROJECT_ID}/locations/us-east5/publishers/anthropic/models/claude-3-5-sonnet:rawPredict \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -H "x-pass-anthropic-beta: context-1m-2025-08-07" \ + -H "x-pass-custom-feature: enabled" \ + -d '{ + "anthropic_version": "vertex-2023-10-16", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 500 + }' +``` + +:::info +The `x-pass-` prefix works for all LLM pass-through endpoints, not just Vertex AI. +::: diff --git a/docs/my-website/docs/providers/gemini.md b/docs/my-website/docs/providers/gemini.md index 110e3f3f09..23a02f7365 100644 --- a/docs/my-website/docs/providers/gemini.md +++ b/docs/my-website/docs/providers/gemini.md @@ -1558,16 +1558,21 @@ LiteLLM Supports the following image types passed in `url` - Images with direct links - https://storage.googleapis.com/github-repo/img/gemini/intro/landmark3.jpg - Image in local storage - ./localimage.jpeg -## Image Resolution Control (Gemini 3+) +## Media Resolution Control (Images & Videos) -For Gemini 3+ models, LiteLLM supports per-part media resolution control using OpenAI's `detail` parameter. This allows you to specify different resolution levels for individual images in your request. +For Gemini 3+ models, LiteLLM supports per-part media resolution control using OpenAI's `detail` parameter. This allows you to specify different resolution levels for individual images and videos in your request, whether using `image_url` or `file` content types. **Supported `detail` values:** - `"low"` - Maps to `media_resolution: "low"` (280 tokens for images, 70 tokens per frame for videos) +- `"medium"` - Maps to `media_resolution: "medium"` - `"high"` - Maps to `media_resolution: "high"` (1120 tokens for images) +- `"ultra_high"` - Maps to `media_resolution: "ultra_high"` - `"auto"` or `None` - Model decides optimal resolution (no `media_resolution` set) -**Usage Example:** +**Usage Examples:** + + + ```python from litellm import completion @@ -1604,10 +1609,193 @@ response = completion( ) ``` + + + +```python +from litellm import completion + +messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Analyze this video" + }, + { + "type": "file", + "file": { + "file_id": "gs://my-bucket/video.mp4", + "format": "video/mp4", + "detail": "high" # High resolution for detailed video analysis + } + } + ] + } +] + +response = completion( + model="gemini/gemini-3-pro-preview", + messages=messages, +) +``` + + + + :::info -**Per-Part Resolution:** Each image in your request can have its own `detail` setting, allowing mixed-resolution requests (e.g., a high-res chart alongside a low-res icon). This feature is only available for Gemini 3+ models. +**Per-Part Resolution:** Each image or video in your request can have its own `detail` setting, allowing mixed-resolution requests (e.g., a high-res chart alongside a low-res icon). This feature works with both `image_url` and `file` content types, and is only available for Gemini 3+ models. ::: +## Video Metadata Control + +For Gemini 3+ models, LiteLLM supports fine-grained video processing control through the `video_metadata` field. This allows you to specify frame extraction rates and time ranges for video analysis. + +**Supported `video_metadata` parameters:** + +| Parameter | Type | Description | Example | +|-----------|------|-------------|---------| +| `fps` | Number | Frame extraction rate (frames per second) | `5` | +| `start_offset` | String | Start time for video clip processing | `"10s"` | +| `end_offset` | String | End time for video clip processing | `"60s"` | + +:::note +**Field Name Conversion:** LiteLLM automatically converts snake_case field names to camelCase for the Gemini API: +- `start_offset` → `startOffset` +- `end_offset` → `endOffset` +- `fps` remains unchanged +::: + +:::warning +- **Gemini 3+ Only:** This feature is only available for Gemini 3.0 and newer models +- **Video Files Recommended:** While `video_metadata` is designed for video files, error handling for other media types is delegated to the Vertex AI API +- **File Formats Supported:** Works with `gs://`, `https://`, and base64-encoded video files +::: + +**Usage Examples:** + + + + +```python +from litellm import completion + +response = completion( + model="gemini/gemini-3-pro-preview", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze this video clip"}, + { + "type": "file", + "file": { + "file_id": "gs://my-bucket/video.mp4", + "format": "video/mp4", + "video_metadata": { + "fps": 5, # Extract 5 frames per second + "start_offset": "10s", # Start from 10 seconds + "end_offset": "60s" # End at 60 seconds + } + } + } + ] + } + ] +) + +print(response.choices[0].message.content) +``` + + + + +```python +from litellm import completion + +response = completion( + model="gemini/gemini-3-pro-preview", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Provide detailed analysis of this video segment"}, + { + "type": "file", + "file": { + "file_id": "https://example.com/presentation.mp4", + "format": "video/mp4", + "detail": "high", # High resolution for detailed analysis + "video_metadata": { + "fps": 10, # Extract 10 frames per second + "start_offset": "30s", # Start from 30 seconds + "end_offset": "90s" # End at 90 seconds + } + } + } + ] + } + ] +) + +print(response.choices[0].message.content) +``` + + + + +1. Setup config.yaml + +```yaml +model_list: + - model_name: gemini-3-pro + litellm_params: + model: gemini/gemini-3-pro-preview + api_key: os.environ/GEMINI_API_KEY +``` + +2. Start proxy + +```bash +litellm --config /path/to/config.yaml +``` + +3. Make request + +```bash +curl http://0.0.0.0:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer " \ + -d '{ + "model": "gemini-3-pro", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze this video clip"}, + { + "type": "file", + "file": { + "file_id": "gs://my-bucket/video.mp4", + "format": "video/mp4", + "detail": "high", + "video_metadata": { + "fps": 5, + "start_offset": "10s", + "end_offset": "60s" + } + } + } + ] + } + ] + }' +``` + + + + ## Sample Usage ```python import os diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 5647b5292e..63e4dceec0 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -1968,6 +1968,244 @@ assert isinstance( ``` +## Media Resolution Control (Images & Videos) + +For Gemini 3+ models, LiteLLM supports per-part media resolution control using OpenAI's `detail` parameter. This allows you to specify different resolution levels for individual images and videos in your request, whether using `image_url` or `file` content types. + +**Supported `detail` values:** +- `"low"` - Maps to `media_resolution: "low"` (280 tokens for images, 70 tokens per frame for videos) +- `"medium"` - Maps to `media_resolution: "medium"` +- `"high"` - Maps to `media_resolution: "high"` (1120 tokens for images) +- `"ultra_high"` - Maps to `media_resolution: "ultra_high"` +- `"auto"` or `None` - Model decides optimal resolution (no `media_resolution` set) + +**Usage Examples:** + + + + +```python +from litellm import completion + +messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://example.com/chart.png", + "detail": "high" # High resolution for detailed chart analysis + } + }, + { + "type": "text", + "text": "Analyze this chart" + }, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/icon.png", + "detail": "low" # Low resolution for simple icon + } + } + ] + } +] + +response = completion( + model="vertex_ai/gemini-3-pro-preview", + messages=messages, +) +``` + + + + +```python +from litellm import completion + +messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Analyze this video" + }, + { + "type": "file", + "file": { + "file_id": "gs://my-bucket/video.mp4", + "format": "video/mp4", + "detail": "high" # High resolution for detailed video analysis + } + } + ] + } +] + +response = completion( + model="vertex_ai/gemini-3-pro-preview", + messages=messages, +) +``` + + + + +:::info +**Per-Part Resolution:** Each image or video in your request can have its own `detail` setting, allowing mixed-resolution requests (e.g., a high-res chart alongside a low-res icon). This feature works with both `image_url` and `file` content types, and is only available for Gemini 3+ models. +::: + +## Video Metadata Control + +For Gemini 3+ models, LiteLLM supports fine-grained video processing control through the `video_metadata` field. This allows you to specify frame extraction rates and time ranges for video analysis. + +**Supported `video_metadata` parameters:** + +| Parameter | Type | Description | Example | +|-----------|------|-------------|---------| +| `fps` | Number | Frame extraction rate (frames per second) | `5` | +| `start_offset` | String | Start time for video clip processing | `"10s"` | +| `end_offset` | String | End time for video clip processing | `"60s"` | + +:::note +**Field Name Conversion:** LiteLLM automatically converts snake_case field names to camelCase for the Gemini API: +- `start_offset` → `startOffset` +- `end_offset` → `endOffset` +- `fps` remains unchanged +::: + +:::warning +- **Gemini 3+ Only:** This feature is only available for Gemini 3.0 and newer models +- **Video Files Recommended:** While `video_metadata` is designed for video files, error handling for other media types is delegated to the Vertex AI API +- **File Formats Supported:** Works with `gs://`, `https://`, and base64-encoded video files +::: + +**Usage Examples:** + + + + +```python +from litellm import completion + +response = completion( + model="vertex_ai/gemini-3-pro-preview", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze this video clip"}, + { + "type": "file", + "file": { + "file_id": "gs://my-bucket/video.mp4", + "format": "video/mp4", + "video_metadata": { + "fps": 5, # Extract 5 frames per second + "start_offset": "10s", # Start from 10 seconds + "end_offset": "60s" # End at 60 seconds + } + } + } + ] + } + ] +) + +print(response.choices[0].message.content) +``` + + + + +```python +from litellm import completion + +response = completion( + model="vertex_ai/gemini-3-pro-preview", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Provide detailed analysis of this video segment"}, + { + "type": "file", + "file": { + "file_id": "https://example.com/presentation.mp4", + "format": "video/mp4", + "detail": "high", # High resolution for detailed analysis + "video_metadata": { + "fps": 10, # Extract 10 frames per second + "start_offset": "30s", # Start from 30 seconds + "end_offset": "90s" # End at 90 seconds + } + } + } + ] + } + ] +) + +print(response.choices[0].message.content) +``` + + + + +1. Setup config.yaml + +```yaml +model_list: + - model_name: gemini-3-pro + litellm_params: + model: vertex_ai/gemini-3-pro-preview + vertex_project: your-project + vertex_location: us-central1 +``` + +2. Start proxy + +```bash +litellm --config /path/to/config.yaml +``` + +3. Make request + +```bash +curl http://0.0.0.0:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer " \ + -d '{ + "model": "gemini-3-pro", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze this video clip"}, + { + "type": "file", + "file": { + "file_id": "gs://my-bucket/video.mp4", + "format": "video/mp4", + "detail": "high", + "video_metadata": { + "fps": 5, + "start_offset": "10s", + "end_offset": "60s" + } + } + } + ] + } + ] + }' +``` + + + ## Usage - PDF / Videos / Audio etc. Files diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index f76fd21468..67fffc13e4 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -397,6 +397,7 @@ router_settings: | AUDIO_SPEECH_CHUNK_SIZE | Chunk size for audio speech processing. Default is 1024 | ANTHROPIC_API_KEY | API key for Anthropic service | ANTHROPIC_API_BASE | Base URL for Anthropic API. Default is https://api.anthropic.com +| ANTHROPIC_TOKEN_COUNTING_BETA_VERSION | Beta version header for Anthropic token counting API. Default is `token-counting-2024-11-01` | AWS_ACCESS_KEY_ID | Access Key ID for AWS services | AWS_BATCH_ROLE_ARN | ARN of the AWS IAM role for batch operations | AWS_DEFAULT_REGION | Default AWS region for service interactions when AWS_REGION is not set @@ -412,6 +413,8 @@ router_settings: | AWS_WEB_IDENTITY_TOKEN | Web identity token for AWS | AWS_WEB_IDENTITY_TOKEN_FILE | Path to file containing web identity token for AWS | AZURE_API_VERSION | Version of the Azure API being used +| AZURE_AI_API_BASE | Base URL for Azure AI services (e.g., Azure AI Anthropic) +| AZURE_AI_API_KEY | API key for Azure AI services (e.g., Azure AI Anthropic) | AZURE_AUTHORITY_HOST | Azure authority host URL | AZURE_CERTIFICATE_PASSWORD | Password for Azure OpenAI certificate | AZURE_CLIENT_ID | Client ID for Azure services diff --git a/docs/my-website/docs/proxy/custom_pricing.md b/docs/my-website/docs/proxy/custom_pricing.md index f6762f5e45..8f4a4c450f 100644 --- a/docs/my-website/docs/proxy/custom_pricing.md +++ b/docs/my-website/docs/proxy/custom_pricing.md @@ -127,6 +127,28 @@ model_list: base_model: azure/gpt-4-1106-preview ``` +### OpenAI Models with Dated Versions + +`base_model` is also useful when OpenAI returns a dated model name in the response that differs from your configured model name. + +**Example**: You configure custom pricing for `gpt-4o-mini-audio-preview`, but OpenAI returns `gpt-4o-mini-audio-preview-2024-12-17` in the response. Since LiteLLM uses the response model name for pricing lookup, your custom pricing won't be applied. + +**Solution** ✅: Set `base_model` to the key you want LiteLLM to use for pricing lookup. + +```yaml +model_list: + - model_name: my-audio-model + litellm_params: + model: openai/gpt-4o-mini-audio-preview + api_key: os.environ/OPENAI_API_KEY + model_info: + base_model: gpt-4o-mini-audio-preview # 👈 Used for pricing lookup + input_cost_per_token: 0.0000006 + output_cost_per_token: 0.0000024 + input_cost_per_audio_token: 0.00001 + output_cost_per_audio_token: 0.00002 +``` + ## Debugging diff --git a/docs/my-website/docs/tutorials/claude_code_max_subscription.md b/docs/my-website/docs/tutorials/claude_code_max_subscription.md new file mode 100644 index 0000000000..399051d41e --- /dev/null +++ b/docs/my-website/docs/tutorials/claude_code_max_subscription.md @@ -0,0 +1,357 @@ +import Image from '@theme/IdealImage'; +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Using Claude Code Max Subscription + +
+ + +Route Claude Code Max subscription traffic through LiteLLM AI Gateway. +
+ +**Why Claude Code Max over direct API?** +- **Lower costs** — Claude Code Max subscriptions are cheaper for Claude Code power users than per-token API pricing + +**Why route through LiteLLM?** +- **Cost attribution** — Track spend per user, team, or key +- **Budgets & rate limits** — Set spending caps and request limits +- **Guardrails** — Apply content filtering and safety controls to all requests + + + +## Quick Start Video + +Watch the end-to-end walkthrough of setting up Claude Code with LiteLLM Gateway: + + + +## Prerequisites + +- [Claude Code](https://docs.anthropic.com/en/docs/claude-code/overview) installed +- Claude Max subscription +- LiteLLM Gateway running + +## Step 1: Configure LiteLLM Proxy + +Create a `config.yaml` with the critical `forward_client_headers_to_llm_api: true` setting: + +```yaml showLineNumbers title="config.yaml" +model_list: + - model_name: anthropic-claude + litellm_params: + model: anthropic/claude-sonnet-4-20250514 + + - model_name: claude-3-5-sonnet-20241022 + litellm_params: + model: anthropic/claude-3-5-sonnet-20241022 + + - model_name: claude-3-5-haiku-20241022 + litellm_params: + model: anthropic/claude-3-5-haiku-20241022 + +general_settings: + forward_client_headers_to_llm_api: true # Required: forwards OAuth token to Anthropic + +litellm_settings: + master_key: os.environ/LITELLM_MASTER_KEY +``` + +:::info Why `forward_client_headers_to_llm_api`? + +This setting forwards the user's OAuth token (in the `Authorization` header) through LiteLLM to the Anthropic API, enabling per-user authentication with their Max subscription while LiteLLM handles tracking and controls. + +::: + +## Step 2: Start LiteLLM Proxy + +```bash showLineNumbers title="Start LiteLLM Proxy" +litellm --config /path/to/config.yaml + +# RUNNING on http://0.0.0.0:4000 +``` + +## Walkthrough + +### Part 1: Create a Virtual Key in LiteLLM + +Navigate to the LiteLLM Dashboard and create a new virtual key for Claude Code usage. + +#### 1.1 Open Virtual Keys Page + +Navigate to the Virtual Keys section in the LiteLLM Dashboard. + + + +#### 1.2 Click "Create New Key" + + + +#### 1.3 Configure Key Details + +Enter a key name (e.g., `claude-code-test`) and select the models you want to allow access to. + + + +#### 1.4 Select Models + +Choose the Anthropic models that should be accessible via this key (e.g., `anthropic-claude`, `claude-4.5-haiku`). + + + +#### 1.5 Confirm Model Selection + + + +#### 1.6 Create the Key + +Click "Create Key" to generate your virtual key. Copy the generated key value (e.g., `sk-otsclFlEblQ-6D60ua2IZg`). + + + +--- + +### Part 2: Sign into Claude Code Max Plan (Client Side) + +Set up Claude Code environment variables and authenticate with your Max subscription. + +#### 2.1 Set Environment Variables + +Configure Claude Code to use LiteLLM Gateway with your virtual key: + +```bash showLineNumbers title="Configure Claude Code Environment Variables" +export ANTHROPIC_BASE_URL=http://localhost:4000 +export ANTHROPIC_MODEL="anthropic-claude" +export ANTHROPIC_CUSTOM_HEADERS="x-litellm-api-key: Bearer sk-otsclFlEblQ-6D60ua2IZg" +``` + + + +#### Environment Variables Explained + +| Variable | Description | +|----------|-------------| +| `ANTHROPIC_BASE_URL` | Points Claude Code to your LiteLLM Gateway endpoint | +| `ANTHROPIC_MODEL` | The model name configured in your LiteLLM `config.yaml` | +| `ANTHROPIC_CUSTOM_HEADERS` | The `x-litellm-api-key` header for LiteLLM authentication | + +#### 2.2 Launch Claude Code + +Start Claude Code: + +```bash showLineNumbers title="Launch Claude Code" +claude +``` + + + +#### 2.3 Select Login Method + +Choose "Claude account with subscription" (Pro, Max, Team, or Enterprise). + + + +#### 2.4 Authorize in Browser + +Claude Code opens your browser to authenticate. Click "Authorize" to connect your Claude Max account. + + + +#### 2.5 Login Successful + +After authorization, you'll see the login success confirmation. + + + +#### 2.6 Complete Setup + +Press Enter to continue past the security notes and complete the setup. + + + +--- + +### Part 3: Use Claude Code with LiteLLM + +Now you can use Claude Code normally, and all requests will be tracked in LiteLLM. + +#### 3.1 Make a Request in Claude Code + +Start using Claude Code - requests will flow through LiteLLM Gateway. + + + +#### 3.2 View Logs in LiteLLM Dashboard + +Navigate to the Logs page in LiteLLM Dashboard to see all Claude Code requests. + + + +#### 3.3 View Request Details + +Click on a request to see detailed information including tokens, cost, duration, and model used. + + + +The logs show: +- **Key Name**: `claude-code-test` (the virtual key you created) +- **Model**: `anthropic/claude-sonnet-4-20250514` +- **Tokens**: 65012 (64679 prompt + 333 completion) +- **Cost**: $0.249754 +- **Status**: Success + + + +--- + +## How It Works + +LiteLLM Gateway handles two types of authentication: +1. **`x-litellm-api-key`**: Authenticates the request with LiteLLM (usage tracking, budgets, rate limits) +2. **OAuth Token (via `Authorization` header)**: Forwarded to Anthropic API for Claude Max authentication + +```mermaid +sequenceDiagram + participant User as Claude Code User + participant LiteLLM as LiteLLM AI Gateway + participant Anthropic as Anthropic API + + User->>LiteLLM: Request with:
- x-litellm-api-key (LiteLLM auth)
- Authorization: Bearer {oauth_token} + + Note over LiteLLM: 1. Validate x-litellm-api-key
2. Check budgets/rate limits
3. Log request for tracking + + LiteLLM->>Anthropic: Forward request with:
- Authorization: Bearer {oauth_token}
(User's Claude Max OAuth token) + + Note over Anthropic: Authenticate user via
OAuth token from Max plan + + Anthropic-->>LiteLLM: Response + + Note over LiteLLM: Log usage, tokens, cost + + LiteLLM-->>User: Response +``` + +### Header Flow + +| Header | Purpose | Handled By | +|--------|---------|------------| +| `x-litellm-api-key` | LiteLLM Gateway authentication, budget tracking, rate limits | LiteLLM | +| `Authorization: Bearer {oauth_token}` | Claude Max subscription authentication | Anthropic API | + +### Complete Request Flow Example + +Here's what a typical request looks like when Claude Code makes a call through LiteLLM: + +```bash showLineNumbers title="Example Request from Claude Code to LiteLLM" +curl -X POST "http://localhost:4000/v1/messages" \ + -H "x-litellm-api-key: Bearer sk-otsclFlEblQ-6D60ua2IZg" \ + -H "Authorization: Bearer oauth_token_from_max_plan" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "anthropic-claude", + "max_tokens": 1024, + "messages": [{"role": "user", "content": "Hello, Claude!"}] + }' +``` + +LiteLLM then: +1. Validates `x-litellm-api-key` for gateway access +2. Logs the request for usage tracking +3. Forwards the request to Anthropic with the OAuth `Authorization` header (because of `forward_client_headers_to_llm_api: true`) + +## Advanced Configuration + +### Per-Model Header Forwarding + +For more granular control, you can enable header forwarding only for specific models: + +```yaml showLineNumbers title="config.yaml - Per-Model Header Forwarding" +model_list: + - model_name: anthropic-claude + litellm_params: + model: anthropic/claude-sonnet-4-20250514 + + - model_name: claude-3-5-haiku-20241022 + litellm_params: + model: anthropic/claude-3-5-haiku-20241022 + +litellm_settings: + master_key: os.environ/LITELLM_MASTER_KEY + model_group_settings: + forward_client_headers_to_llm_api: + - anthropic-claude + - claude-3-5-haiku-20241022 +``` + +### Budget Controls + +Set up per-user budgets while using Max subscriptions: + +```yaml showLineNumbers title="config.yaml - With Database for Budget Tracking" +model_list: + - model_name: anthropic-claude + litellm_params: + model: anthropic/claude-sonnet-4-20250514 + +general_settings: + forward_client_headers_to_llm_api: true + database_url: "postgresql://..." + +litellm_settings: + master_key: os.environ/LITELLM_MASTER_KEY +``` + +Then create virtual keys with budgets: + +```bash showLineNumbers title="Create Virtual Key with Budget" +curl -X POST "http://localhost:4000/key/generate" \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "key_alias": "developer-1", + "max_budget": 100.00, + "budget_duration": "monthly" + }' +``` + +## Troubleshooting + +### OAuth Token Not Being Forwarded + +**Symptom**: Authentication errors from Anthropic API + +**Solution**: Ensure `forward_client_headers_to_llm_api: true` is set in your config: + +```yaml showLineNumbers title="config.yaml - Enable Header Forwarding" +general_settings: + forward_client_headers_to_llm_api: true +``` + +### LiteLLM Authentication Failing + +**Symptom**: 401 errors from LiteLLM Gateway + +**Solution**: Verify `x-litellm-api-key` header is set correctly in `ANTHROPIC_CUSTOM_HEADERS`: + +```bash showLineNumbers title="Verify Key Info" +curl -X GET "http://localhost:4000/key/info" \ + -H "Authorization: Bearer sk-otsclFlEblQ-6D60ua2IZg" +``` + +### Model Not Found + +**Symptom**: Model not found errors + +**Solution**: Ensure the `ANTHROPIC_MODEL` matches a model name in your config: + +```bash showLineNumbers title="List Available Models" +curl "http://localhost:4000/v1/models" \ + -H "Authorization: Bearer sk-otsclFlEblQ-6D60ua2IZg" +``` + +## Related Documentation + +- [Forward Client Headers](/docs/proxy/forward_client_headers) - Detailed header forwarding configuration +- [Claude Code Quickstart](/docs/tutorials/claude_responses_api) - Basic Claude Code + LiteLLM setup +- [Virtual Keys](/docs/proxy/virtual_keys) - Creating and managing API keys +- [Budgets & Rate Limits](/docs/proxy/users) - Setting up usage controls diff --git a/docs/my-website/docs/tutorials/claude_responses_api.md b/docs/my-website/docs/tutorials/claude_responses_api.md index 6b681d93a8..03ac9935fd 100644 --- a/docs/my-website/docs/tutorials/claude_responses_api.md +++ b/docs/my-website/docs/tutorials/claude_responses_api.md @@ -37,18 +37,22 @@ Create a secure configuration using environment variables: ```yaml model_list: - # Claude models - - model_name: claude-3-5-sonnet-20241022 + # Configure the models you want to use + - model_name: claude-sonnet-4-5-20250929 litellm_params: - model: anthropic/claude-3-5-sonnet-20241022 + model: anthropic/claude-sonnet-4-5-20250929 api_key: os.environ/ANTHROPIC_API_KEY - - - model_name: claude-3-5-haiku-20241022 + + - model_name: claude-haiku-4-5-20251001 litellm_params: - model: anthropic/claude-3-5-haiku-20241022 + model: anthropic/claude-haiku-4-5-20251001 + api_key: os.environ/ANTHROPIC_API_KEY + + - model_name: claude-opus-4-5-20251101 + litellm_params: + model: anthropic/claude-opus-4-5-20251101 api_key: os.environ/ANTHROPIC_API_KEY - litellm_settings: master_key: os.environ/LITELLM_MASTER_KEY ``` @@ -60,6 +64,10 @@ export ANTHROPIC_API_KEY="your-anthropic-api-key" export LITELLM_MASTER_KEY="sk-1234567890" # Generate a secure key ``` +:::tip +Alternatively, you can store `ANTHROPIC_API_KEY` in a `.env` file in your proxy directory. LiteLLM will automatically load it when starting. +::: + ### 2. Start proxy ```bash @@ -111,15 +119,55 @@ export ANTHROPIC_AUTH_TOKEN="$LITELLM_MASTER_KEY" ### 5. Use Claude Code -Start Claude Code and it will automatically use your configured models: +Start Claude Code with the model you want to use: ```bash -# Claude Code will use the models configured in your LiteLLM proxy -claude +# Specify model at startup +claude --model claude-sonnet-4-5-20250929 -# Or specify a model if you have multiple configured -claude --model claude-3-5-sonnet-20241022 -claude --model claude-3-5-haiku-20241022 +# Or specify a different model +claude --model claude-haiku-4-5-20251001 +claude --model claude-opus-4-5-20251101 + +# Or change model during a session +claude +/model claude-sonnet-4-5-20250929 +``` + +Alternatively, set default models with environment variables: + +```bash +export ANTHROPIC_DEFAULT_SONNET_MODEL=claude-sonnet-4-5-20250929 +export ANTHROPIC_DEFAULT_HAIKU_MODEL=claude-haiku-4-5-20251001 +export ANTHROPIC_DEFAULT_OPUS_MODEL=claude-opus-4-5-20251101 +claude +``` + +### Using 1M Context Window + +Claude Code supports extended context (1 million tokens) using the `[1m]` suffix: + +```bash +# Use Sonnet with 1M context (requires quotes in shell) +claude --model 'claude-sonnet-4-5-20250929[1m]' + +# Inside a Claude Code session (no quotes needed) +/model claude-sonnet-4-5-20250929[1m] +``` + +:::warning +**Important:** When using `--model` with `[1m]` in the shell, you must use quotes to prevent the shell from interpreting the brackets. +::: + +**How it works:** +- Claude Code strips the `[1m]` suffix before sending to LiteLLM +- Claude Code automatically adds the header `anthropic-beta: context-1m-2025-08-07` +- Your LiteLLM config should **NOT** include `[1m]` in model names + +**Verify 1M context is active:** +```bash +/context +# Should show: 21k/1000k tokens (2%) ``` Example conversation: @@ -140,6 +188,7 @@ Common issues and solutions: **Model not found:** - Ensure the model name in Claude Code matches exactly with your `config.yaml` +- Use `--model` flag or environment variables to specify the model - Check LiteLLM logs for detailed error messages ## Using Bedrock/Vertex AI/Azure Foundry Models diff --git a/docs/my-website/img/claude_code_max.png b/docs/my-website/img/claude_code_max.png new file mode 100644 index 0000000000..65c9578a45 Binary files /dev/null and b/docs/my-website/img/claude_code_max.png differ diff --git a/docs/my-website/img/claude_code_max/step1.jpeg b/docs/my-website/img/claude_code_max/step1.jpeg new file mode 100644 index 0000000000..6b65d598d3 Binary files /dev/null and b/docs/my-website/img/claude_code_max/step1.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step10.jpeg b/docs/my-website/img/claude_code_max/step10.jpeg new file mode 100644 index 0000000000..326f9b12d1 Binary files /dev/null and b/docs/my-website/img/claude_code_max/step10.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step12.jpeg b/docs/my-website/img/claude_code_max/step12.jpeg new file mode 100644 index 0000000000..97199e9ead Binary files /dev/null and b/docs/my-website/img/claude_code_max/step12.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step13.jpeg b/docs/my-website/img/claude_code_max/step13.jpeg new file mode 100644 index 0000000000..53fd1c9bd5 Binary files /dev/null and b/docs/my-website/img/claude_code_max/step13.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step14.jpeg b/docs/my-website/img/claude_code_max/step14.jpeg new file mode 100644 index 0000000000..5c3e4b05e2 Binary files /dev/null and b/docs/my-website/img/claude_code_max/step14.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step15.jpeg b/docs/my-website/img/claude_code_max/step15.jpeg new file mode 100644 index 0000000000..2c63ba6e75 Binary files /dev/null and b/docs/my-website/img/claude_code_max/step15.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step16.jpeg b/docs/my-website/img/claude_code_max/step16.jpeg new file mode 100644 index 0000000000..7abb53edb8 Binary files /dev/null and b/docs/my-website/img/claude_code_max/step16.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step17.jpeg b/docs/my-website/img/claude_code_max/step17.jpeg new file mode 100644 index 0000000000..a9c352f85e Binary files /dev/null and b/docs/my-website/img/claude_code_max/step17.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step18.jpeg b/docs/my-website/img/claude_code_max/step18.jpeg new file mode 100644 index 0000000000..0177537fef Binary files /dev/null and b/docs/my-website/img/claude_code_max/step18.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step19.jpeg b/docs/my-website/img/claude_code_max/step19.jpeg new file mode 100644 index 0000000000..d84eec24dd Binary files /dev/null and b/docs/my-website/img/claude_code_max/step19.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step2.jpeg b/docs/my-website/img/claude_code_max/step2.jpeg new file mode 100644 index 0000000000..2d7255c73a Binary files /dev/null and b/docs/my-website/img/claude_code_max/step2.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step20.jpeg b/docs/my-website/img/claude_code_max/step20.jpeg new file mode 100644 index 0000000000..3e97cba38c Binary files /dev/null and b/docs/my-website/img/claude_code_max/step20.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step21.jpeg b/docs/my-website/img/claude_code_max/step21.jpeg new file mode 100644 index 0000000000..02387c7666 Binary files /dev/null and b/docs/my-website/img/claude_code_max/step21.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step22.jpeg b/docs/my-website/img/claude_code_max/step22.jpeg new file mode 100644 index 0000000000..7aa920221d Binary files /dev/null and b/docs/my-website/img/claude_code_max/step22.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step23.jpeg b/docs/my-website/img/claude_code_max/step23.jpeg new file mode 100644 index 0000000000..4eb9c62c72 Binary files /dev/null and b/docs/my-website/img/claude_code_max/step23.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step24.jpeg b/docs/my-website/img/claude_code_max/step24.jpeg new file mode 100644 index 0000000000..bb38c2e19a Binary files /dev/null and b/docs/my-website/img/claude_code_max/step24.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step25.jpeg b/docs/my-website/img/claude_code_max/step25.jpeg new file mode 100644 index 0000000000..fb1e095066 Binary files /dev/null and b/docs/my-website/img/claude_code_max/step25.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step26.jpeg b/docs/my-website/img/claude_code_max/step26.jpeg new file mode 100644 index 0000000000..9eb418b9be Binary files /dev/null and b/docs/my-website/img/claude_code_max/step26.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step27.jpeg b/docs/my-website/img/claude_code_max/step27.jpeg new file mode 100644 index 0000000000..b8efb3aeb1 Binary files /dev/null and b/docs/my-website/img/claude_code_max/step27.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step28.jpeg b/docs/my-website/img/claude_code_max/step28.jpeg new file mode 100644 index 0000000000..a2ce52441e Binary files /dev/null and b/docs/my-website/img/claude_code_max/step28.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step3.jpeg b/docs/my-website/img/claude_code_max/step3.jpeg new file mode 100644 index 0000000000..a5f28c8049 Binary files /dev/null and b/docs/my-website/img/claude_code_max/step3.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step4.jpeg b/docs/my-website/img/claude_code_max/step4.jpeg new file mode 100644 index 0000000000..ec9ffa4deb Binary files /dev/null and b/docs/my-website/img/claude_code_max/step4.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step5.jpeg b/docs/my-website/img/claude_code_max/step5.jpeg new file mode 100644 index 0000000000..25d33f27a0 Binary files /dev/null and b/docs/my-website/img/claude_code_max/step5.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step6.jpeg b/docs/my-website/img/claude_code_max/step6.jpeg new file mode 100644 index 0000000000..116f792eac Binary files /dev/null and b/docs/my-website/img/claude_code_max/step6.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step7.jpeg b/docs/my-website/img/claude_code_max/step7.jpeg new file mode 100644 index 0000000000..1a3b232d2b Binary files /dev/null and b/docs/my-website/img/claude_code_max/step7.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step8.jpeg b/docs/my-website/img/claude_code_max/step8.jpeg new file mode 100644 index 0000000000..1a67a135c3 Binary files /dev/null and b/docs/my-website/img/claude_code_max/step8.jpeg differ diff --git a/docs/my-website/img/claude_code_max/step9.jpeg b/docs/my-website/img/claude_code_max/step9.jpeg new file mode 100644 index 0000000000..b95594e617 Binary files /dev/null and b/docs/my-website/img/claude_code_max/step9.jpeg differ diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index d2af237141..ad5019d880 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -121,6 +121,7 @@ const sidebars = { label: "Claude Code", items: [ "tutorials/claude_responses_api", + "tutorials/claude_code_max_subscription", "tutorials/claude_code_customer_tracking", "tutorials/claude_code_websearch", "tutorials/claude_mcp", @@ -516,7 +517,14 @@ const sidebars = { "mcp_troubleshoot", ] }, - "anthropic_unified", + { + type: "category", + label: "/v1/messages", + items: [ + "anthropic_unified/index", + "anthropic_unified/structured_output", + ] + }, "anthropic_count_tokens", "moderation", "ocr", diff --git a/litellm/constants.py b/litellm/constants.py index c98551fb1b..5da0cf1489 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -323,6 +323,9 @@ EMAIL_BUDGET_ALERT_TTL = int(os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60)) EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE = float(os.getenv("EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8)) # 80% of max budget ############### LLM Provider Constants ############### ### ANTHROPIC CONSTANTS ### +ANTHROPIC_TOKEN_COUNTING_BETA_VERSION = os.getenv( + "ANTHROPIC_TOKEN_COUNTING_BETA_VERSION", "token-counting-2024-11-01" +) ANTHROPIC_SKILLS_API_BETA_VERSION = "skills-2025-10-02" ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES = { "low": 1, @@ -1119,6 +1122,20 @@ BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES = [ "generateQuery/", "optimize-prompt/", ] + + +# Headers that are safe to forward from incoming requests to Vertex AI +# Using an allowlist approach for security - only forward headers we explicitly trust +ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS = { + "anthropic-beta", # Required for Anthropic features like extended context windows + "content-type", # Required for request body parsing +} + +# Prefix for headers that should be forwarded to the provider with the prefix stripped +# e.g., 'x-pass-anthropic-beta: value' becomes 'anthropic-beta: value' +# Works for all LLM pass-through endpoints (Vertex AI, Anthropic, Bedrock, etc.) +PASS_THROUGH_HEADER_PREFIX = "x-pass-" + BASE_MCP_ROUTE = "/mcp" BATCH_STATUS_POLL_INTERVAL_SECONDS = int( diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 4603ebab29..aea2c7c0c5 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -16,6 +16,21 @@ import openai from litellm.types.utils import LiteLLMCommonStrings +_MINIMAL_ERROR_RESPONSE: Optional[httpx.Response] = None + + +def _get_minimal_error_response() -> httpx.Response: + """Get a cached minimal httpx.Response object for error cases.""" + global _MINIMAL_ERROR_RESPONSE + if _MINIMAL_ERROR_RESPONSE is None: + _MINIMAL_ERROR_RESPONSE = httpx.Response( + status_code=400, + request=httpx.Request( + method="GET", url="https://litellm.ai" + ), + ) + return _MINIMAL_ERROR_RESPONSE + class AuthenticationError(openai.AuthenticationError): # type: ignore def __init__( @@ -127,16 +142,15 @@ class BadRequestError(openai.BadRequestError): # type: ignore self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - _response_headers = ( - getattr(response, "headers", None) if response is not None else None - ) - self.response = httpx.Response( - status_code=self.status_code, - headers=_response_headers, - request=httpx.Request( - method="GET", url="https://litellm.ai" - ), # mock request object - ) + if ( + response is not None + and isinstance(response, httpx.Response) + and hasattr(response, "request") + and response.request is not None + ): + self.response = response + else: + self.response = _get_minimal_error_response() super().__init__( self.message, response=self.response, body=body ) # Call the base class constructor with the parameters it needs diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 7e62613a7e..8087c17caf 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -593,30 +593,10 @@ class LangFuseLogger: trace_id = clean_metadata.pop("trace_id", None) # Use standard_logging_object.trace_id if available (when trace_id from metadata is None) # This allows standard trace_id to be used when provided in standard_logging_object - # However, we skip standard_logging_object.trace_id if it's a UUID (from litellm_trace_id default), - # as we want to fall back to litellm_call_id instead for better traceability. - # Note: Users can still explicitly set a UUID trace_id via metadata["trace_id"] (highest priority) if trace_id is None and standard_logging_object is not None: - standard_trace_id = cast( + trace_id = cast( Optional[str], standard_logging_object.get("trace_id") ) - # Only use standard_logging_object.trace_id if it's not a UUID - # UUIDs are 36 characters with hyphens in format: xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx - # We check for this specific pattern to avoid rejecting valid trace_ids that happen to have hyphens - # This primarily filters out default litellm_trace_id UUIDs, while still allowing user-provided - # trace_ids via metadata["trace_id"] (which is checked first and not affected by this logic) - if standard_trace_id is not None: - # Check if it's a UUID: 36 chars, 4 hyphens, specific pattern - is_uuid = ( - len(standard_trace_id) == 36 - and standard_trace_id.count("-") == 4 - and standard_trace_id[8] == "-" - and standard_trace_id[13] == "-" - and standard_trace_id[18] == "-" - and standard_trace_id[23] == "-" - ) - if not is_uuid: - trace_id = standard_trace_id # Fallback to litellm_call_id if no trace_id found if trace_id is None: trace_id = litellm_call_id diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 7771f2d342..93d631eb0f 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -988,7 +988,7 @@ class OpenTelemetry(CustomLogger): from opentelemetry._logs import SeverityNumber, get_logger, get_logger_provider try: - from opentelemetry.sdk._logs import LogRecord as SdkLogRecord # OTEL < 1.39.0 + from opentelemetry.sdk._logs import LogRecord as SdkLogRecord # type: ignore[attr-defined] # OTEL < 1.39.0 except ImportError: from opentelemetry.sdk._logs._internal import LogRecord as SdkLogRecord # OTEL >= 1.39.0 diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 0d35cfa314..e290101f8b 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -93,8 +93,9 @@ def get_litellm_params( "text_completion": text_completion, "azure_ad_token_provider": azure_ad_token_provider, "user_continue_message": user_continue_message, - "base_model": base_model - or _get_base_model_from_litellm_call_metadata(metadata=metadata), + "base_model": base_model or ( + _get_base_model_from_litellm_call_metadata(metadata=metadata) if metadata else None + ), "litellm_trace_id": litellm_trace_id, "litellm_session_id": litellm_session_id, "hf_model_name": hf_model_name, diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 807b66faec..718773a1b1 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -1,7 +1,5 @@ from typing import Optional, Tuple -import httpx - import litellm from litellm.constants import REPLICATE_MODEL_NAME_WITH_ID_LENGTH from litellm.llms.openai_like.json_loader import JSONProviderRegistry @@ -453,11 +451,7 @@ def get_llm_provider( # noqa: PLR0915 raise litellm.exceptions.BadRequestError( # type: ignore message=error_str, model=model, - response=httpx.Response( - status_code=400, - content=error_str, - request=httpx.Request(method="completion", url="https://github.com/BerriAI/litellm"), # type: ignore - ), + response=None, llm_provider="", ) if api_base is not None and not isinstance(api_base, str): @@ -481,11 +475,7 @@ def get_llm_provider( # noqa: PLR0915 raise litellm.exceptions.BadRequestError( # type: ignore message=f"GetLLMProvider Exception - {str(e)}\n\noriginal model: {model}", model=model, - response=httpx.Response( - status_code=400, - content=error_str, - request=httpx.Request(method="completion", url="https://github.com/BerriAI/litellm"), # type: ignore - ), + response=None, llm_provider="", ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9cef21bfdf..d3bcfe8200 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -325,12 +325,12 @@ class Logging(LiteLLMLoggingBaseClass): messages = new_messages self.model = model - self.messages = copy.deepcopy(messages) + self.messages = copy.deepcopy(messages) if messages is not None else None self.stream = stream self.start_time = start_time # log the call start time self.call_type = call_type self.litellm_call_id = litellm_call_id - self.litellm_trace_id: str = litellm_trace_id or str(uuid.uuid4()) + self.litellm_trace_id: str = litellm_trace_id if litellm_trace_id else str(uuid.uuid4()) self.function_id = function_id self.streaming_chunks: List[Any] = [] # for generating complete stream response self.sync_streaming_chunks: List[ @@ -1624,15 +1624,7 @@ class Logging(LiteLLMLoggingBaseClass): result.usage ) ) - setattr( - result, - "usage", - ( - transformed_usage.model_dump() - if hasattr(transformed_usage, "model_dump") - else dict(transformed_usage) - ), - ) + setattr(result, "usage", transformed_usage) if ( standard_logging_payload := self.model_call_details.get( "standard_logging_object" diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 3093a37c26..3304759f74 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1571,6 +1571,46 @@ class CustomStreamWrapper: ) return chunk + def _add_mcp_metadata_to_final_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + """ + Add MCP metadata from _hidden_params to the final chunk's delta.provider_specific_fields. + + This method checks if MCP metadata is stored in _hidden_params and adds it to + the chunk's delta.provider_specific_fields, similar to how RAG adds search results. + """ + try: + # Check if MCP metadata should be added to final chunk + if not hasattr(self, "_hidden_params") or not self._hidden_params: + return chunk + + mcp_metadata = self._hidden_params.get("mcp_metadata") + if not mcp_metadata: + return chunk + + # Add MCP metadata to delta.provider_specific_fields + if hasattr(chunk, "choices") and chunk.choices: + for choice in chunk.choices: + if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: + # Get existing provider_specific_fields or create new dict + provider_fields = ( + getattr(choice.delta, "provider_specific_fields", None) or {} + ) + + # Add MCP metadata + if isinstance(mcp_metadata, dict): + provider_fields.update(mcp_metadata) + + # Set the provider_specific_fields + setattr(choice.delta, "provider_specific_fields", provider_fields) + + except Exception as e: + from litellm._logging import verbose_logger + verbose_logger.exception( + f"Error adding MCP metadata to final chunk: {str(e)}" + ) + + return chunk + def cache_streaming_response(self, processed_chunk, cache_hit: bool): """ Caches the streaming response @@ -1712,6 +1752,8 @@ class CustomStreamWrapper: if self.sent_last_chunk is True and self.stream_options is None: usage = calculate_total_usage(chunks=self.chunks) response._hidden_params["usage"] = usage + # Add MCP metadata to final chunk if present + response = self._add_mcp_metadata_to_final_chunk(response) # RETURN RESULT return response @@ -1884,6 +1926,8 @@ class CustomStreamWrapper: processed_chunk ) ) + # Add MCP metadata to final chunk if present (after hooks) + processed_chunk = self._add_mcp_metadata_to_final_chunk(processed_chunk) return processed_chunk raise StopAsyncIteration diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 86378b97d2..82eccee596 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -200,6 +200,68 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return params + @staticmethod + def filter_anthropic_output_schema(schema: Dict[str, Any]) -> Dict[str, Any]: + """ + Filter out unsupported fields from JSON schema for Anthropic's output_format API. + + Anthropic's output_format doesn't support certain JSON schema properties: + - maxItems: Not supported for array types + - minItems: Not supported for array types + + This function recursively removes these unsupported fields while preserving + all other valid schema properties. + + Args: + schema: The JSON schema dictionary to filter + + Returns: + A new dictionary with unsupported fields removed + + Related issue: https://github.com/BerriAI/litellm/issues/19444 + """ + if not isinstance(schema, dict): + return schema + + unsupported_fields = {"maxItems", "minItems"} + + result: Dict[str, Any] = {} + for key, value in schema.items(): + if key in unsupported_fields: + continue + + if key == "properties" and isinstance(value, dict): + result[key] = { + k: AnthropicConfig.filter_anthropic_output_schema(v) + for k, v in value.items() + } + elif key == "items" and isinstance(value, dict): + result[key] = AnthropicConfig.filter_anthropic_output_schema(value) + elif key == "$defs" and isinstance(value, dict): + result[key] = { + k: AnthropicConfig.filter_anthropic_output_schema(v) + for k, v in value.items() + } + elif key == "anyOf" and isinstance(value, list): + result[key] = [ + AnthropicConfig.filter_anthropic_output_schema(item) + for item in value + ] + elif key == "allOf" and isinstance(value, list): + result[key] = [ + AnthropicConfig.filter_anthropic_output_schema(item) + for item in value + ] + elif key == "oneOf" and isinstance(value, list): + result[key] = [ + AnthropicConfig.filter_anthropic_output_schema(item) + for item in value + ] + else: + result[key] = value + + return result + def get_json_schema_from_pydantic_object( self, response_format: Union[Any, Dict, None] ) -> Optional[dict]: @@ -636,9 +698,13 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) if json_schema is None: return None + + # Filter out unsupported fields for Anthropic's output_format API + filtered_schema = self.filter_anthropic_output_schema(json_schema) + return AnthropicOutputSchema( type="json_schema", - schema=json_schema, + schema=filtered_schema, ) def map_response_format_to_anthropic_tool( diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index fcbe9823ed..cb23d21fbc 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -2,7 +2,7 @@ This file contains common utils for anthropic calls. """ -from typing import Any, Dict, List, Optional, Union +from typing import Dict, List, Optional, Union import httpx @@ -14,11 +14,36 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.anthropic import ( ANTHROPIC_HOSTED_TOOLS, + ANTHROPIC_OAUTH_BETA_HEADER, + ANTHROPIC_OAUTH_TOKEN_PREFIX, AllAnthropicToolsValues, AnthropicMcpServerTool, ) from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import TokenCountResponse + + +def optionally_handle_anthropic_oauth( + headers: dict, api_key: Optional[str] +) -> tuple[dict, Optional[str]]: + """ + Handle Anthropic OAuth token detection and header setup. + + If an OAuth token is detected in the Authorization header, extracts it + and sets the required OAuth headers. + + Args: + headers: Request headers dict + api_key: Current API key (may be None) + + Returns: + Tuple of (updated headers, api_key) + """ + auth_header = headers.get("authorization", "") + if auth_header and auth_header.startswith(f"Bearer {ANTHROPIC_OAUTH_TOKEN_PREFIX}"): + api_key = auth_header.replace("Bearer ", "") + headers["anthropic-beta"] = ANTHROPIC_OAUTH_BETA_HEADER + headers["anthropic-dangerous-direct-browser-access"] = "true" + return headers, api_key class AnthropicError(BaseLLMException): @@ -372,6 +397,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> Dict: + # Check for Anthropic OAuth token in headers + headers, api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key) if api_key is None: raise litellm.AuthenticationError( message="Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` in your environment vars", @@ -476,45 +503,11 @@ class AnthropicModelInfo(BaseLLMModelInfo): Returns: AnthropicTokenCounter instance for this provider. """ - return AnthropicTokenCounter() - - -class AnthropicTokenCounter(BaseTokenCounter): - """Token counter implementation for Anthropic provider.""" - - def should_use_token_counting_api( - self, - custom_llm_provider: Optional[str] = None, - ) -> bool: - from litellm.types.utils import LlmProviders - return custom_llm_provider == LlmProviders.ANTHROPIC.value - - async def count_tokens( - self, - model_to_use: str, - messages: Optional[List[Dict[str, Any]]], - contents: Optional[List[Dict[str, Any]]], - deployment: Optional[Dict[str, Any]] = None, - request_model: str = "", - ) -> Optional[TokenCountResponse]: - from litellm.proxy.utils import count_tokens_with_anthropic_api - - result = await count_tokens_with_anthropic_api( - model_to_use=model_to_use, - messages=messages, - deployment=deployment, + from litellm.llms.anthropic.count_tokens.token_counter import ( + AnthropicTokenCounter, ) - - if result is not None: - return TokenCountResponse( - total_tokens=result.get("total_tokens", 0), - request_model=request_model, - model_used=model_to_use, - tokenizer_type=result.get("tokenizer_used", ""), - original_response=result, - ) - - return None + + return AnthropicTokenCounter() def process_anthropic_headers(headers: Union[httpx.Headers, dict]) -> dict: diff --git a/litellm/llms/anthropic/count_tokens/__init__.py b/litellm/llms/anthropic/count_tokens/__init__.py new file mode 100644 index 0000000000..ef46862bda --- /dev/null +++ b/litellm/llms/anthropic/count_tokens/__init__.py @@ -0,0 +1,15 @@ +""" +Anthropic CountTokens API implementation. +""" + +from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler +from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCounter +from litellm.llms.anthropic.count_tokens.transformation import ( + AnthropicCountTokensConfig, +) + +__all__ = [ + "AnthropicCountTokensHandler", + "AnthropicCountTokensConfig", + "AnthropicTokenCounter", +] diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py new file mode 100644 index 0000000000..5b5354228f --- /dev/null +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -0,0 +1,122 @@ +""" +Anthropic CountTokens API handler. + +Uses httpx for HTTP requests instead of the Anthropic SDK. +""" + +from typing import Any, Dict, List, Optional, Union + +import httpx + +import litellm +from litellm._logging import verbose_logger +from litellm.llms.anthropic.common_utils import AnthropicError +from litellm.llms.anthropic.count_tokens.transformation import ( + AnthropicCountTokensConfig, +) +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + + +class AnthropicCountTokensHandler(AnthropicCountTokensConfig): + """ + Handler for Anthropic CountTokens API requests. + + Uses httpx for HTTP requests, following the same pattern as BedrockCountTokensHandler. + """ + + async def handle_count_tokens_request( + self, + model: str, + messages: List[Dict[str, Any]], + api_key: str, + api_base: Optional[str] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + ) -> Dict[str, Any]: + """ + Handle a CountTokens request using httpx. + + Args: + model: The model identifier (e.g., "claude-3-5-sonnet-20241022") + messages: The messages to count tokens for + api_key: The Anthropic API key + api_base: Optional custom API base URL + timeout: Optional timeout for the request (defaults to litellm.request_timeout) + + Returns: + Dictionary containing token count response + + Raises: + AnthropicError: If the API request fails + """ + try: + # Validate the request + self.validate_request(model, messages) + + verbose_logger.debug( + f"Processing Anthropic CountTokens request for model: {model}" + ) + + # Transform request to Anthropic format + request_body = self.transform_request_to_count_tokens( + model=model, + messages=messages, + ) + + verbose_logger.debug(f"Transformed request: {request_body}") + + # Get endpoint URL + endpoint_url = api_base or self.get_anthropic_count_tokens_endpoint() + + verbose_logger.debug(f"Making request to: {endpoint_url}") + + # Get required headers + headers = self.get_required_headers(api_key) + + # Use LiteLLM's async httpx client + async_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders.ANTHROPIC + ) + + # Use provided timeout or fall back to litellm.request_timeout + request_timeout = timeout if timeout is not None else litellm.request_timeout + + response = await async_client.post( + endpoint_url, + headers=headers, + json=request_body, + timeout=request_timeout, + ) + + verbose_logger.debug(f"Response status: {response.status_code}") + + if response.status_code != 200: + error_text = response.text + verbose_logger.error(f"Anthropic API error: {error_text}") + raise AnthropicError( + status_code=response.status_code, + message=error_text, + ) + + anthropic_response = response.json() + + verbose_logger.debug(f"Anthropic response: {anthropic_response}") + + # Return Anthropic response directly - no transformation needed + return anthropic_response + + except AnthropicError: + # Re-raise Anthropic exceptions as-is + raise + except httpx.HTTPStatusError as e: + # HTTP errors - preserve the actual status code + verbose_logger.error(f"HTTP error in CountTokens handler: {str(e)}") + raise AnthropicError( + status_code=e.response.status_code, + message=e.response.text, + ) + except Exception as e: + verbose_logger.error(f"Error in CountTokens handler: {str(e)}") + raise AnthropicError( + status_code=500, + message=f"CountTokens processing error: {str(e)}", + ) diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py new file mode 100644 index 0000000000..266b2794fc --- /dev/null +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -0,0 +1,104 @@ +""" +Anthropic Token Counter implementation using the CountTokens API. +""" + +import os +from typing import Any, Dict, List, Optional + +from litellm._logging import verbose_logger +from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler +from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.types.utils import LlmProviders, TokenCountResponse + +# Global handler instance - reuse across all token counting requests +anthropic_count_tokens_handler = AnthropicCountTokensHandler() + + +class AnthropicTokenCounter(BaseTokenCounter): + """Token counter implementation for Anthropic provider using the CountTokens API.""" + + def should_use_token_counting_api( + self, + custom_llm_provider: Optional[str] = None, + ) -> bool: + return custom_llm_provider == LlmProviders.ANTHROPIC.value + + async def count_tokens( + self, + model_to_use: str, + messages: Optional[List[Dict[str, Any]]], + contents: Optional[List[Dict[str, Any]]], + deployment: Optional[Dict[str, Any]] = None, + request_model: str = "", + ) -> Optional[TokenCountResponse]: + """ + Count tokens using Anthropic's CountTokens API. + + Args: + model_to_use: The model identifier + messages: The messages to count tokens for + contents: Alternative content format (not used for Anthropic) + deployment: Deployment configuration containing litellm_params + request_model: The original request model name + + Returns: + TokenCountResponse with token count, or None if counting fails + """ + from litellm.llms.anthropic.common_utils import AnthropicError + + if not messages: + return None + + deployment = deployment or {} + litellm_params = deployment.get("litellm_params", {}) + + # Get Anthropic API key from deployment config or environment + api_key = litellm_params.get("api_key") + if not api_key: + api_key = os.getenv("ANTHROPIC_API_KEY") + + if not api_key: + verbose_logger.warning("No Anthropic API key found for token counting") + return None + + try: + result = await anthropic_count_tokens_handler.handle_count_tokens_request( + model=model_to_use, + messages=messages, + api_key=api_key, + ) + + if result is not None: + return TokenCountResponse( + total_tokens=result.get("input_tokens", 0), + request_model=request_model, + model_used=model_to_use, + tokenizer_type="anthropic_api", + original_response=result, + ) + except AnthropicError as e: + verbose_logger.warning( + f"Anthropic CountTokens API error: status={e.status_code}, message={e.message}" + ) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="anthropic_api", + error=True, + error_message=e.message, + status_code=e.status_code, + ) + except Exception as e: + verbose_logger.warning(f"Error calling Anthropic CountTokens API: {e}") + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="anthropic_api", + error=True, + error_message=str(e), + status_code=500, + ) + + return None diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py new file mode 100644 index 0000000000..c3ad72436b --- /dev/null +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -0,0 +1,103 @@ +""" +Anthropic CountTokens API transformation logic. + +This module handles the transformation of requests to Anthropic's CountTokens API format. +""" + +from typing import Any, Dict, List + +from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION + + +class AnthropicCountTokensConfig: + """ + Configuration and transformation logic for Anthropic CountTokens API. + + Anthropic CountTokens API Specification: + - Endpoint: POST https://api.anthropic.com/v1/messages/count_tokens + - Beta header required: anthropic-beta: token-counting-2024-11-01 + - Response: {"input_tokens": } + """ + + def get_anthropic_count_tokens_endpoint(self) -> str: + """ + Get the Anthropic CountTokens API endpoint. + + Returns: + The endpoint URL for the CountTokens API + """ + return "https://api.anthropic.com/v1/messages/count_tokens" + + def transform_request_to_count_tokens( + self, + model: str, + messages: List[Dict[str, Any]], + ) -> Dict[str, Any]: + """ + Transform request to Anthropic CountTokens format. + + Input: + { + "model": "claude-3-5-sonnet-20241022", + "messages": [{"role": "user", "content": "Hello!"}] + } + + Output (Anthropic CountTokens format): + { + "model": "claude-3-5-sonnet-20241022", + "messages": [{"role": "user", "content": "Hello!"}] + } + """ + return { + "model": model, + "messages": messages, + } + + def get_required_headers(self, api_key: str) -> Dict[str, str]: + """ + Get the required headers for the CountTokens API. + + Args: + api_key: The Anthropic API key + + Returns: + Dictionary of required headers + """ + return { + "Content-Type": "application/json", + "x-api-key": api_key, + "anthropic-version": "2023-06-01", + "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, + } + + def validate_request( + self, model: str, messages: List[Dict[str, Any]] + ) -> None: + """ + Validate the incoming count tokens request. + + Args: + model: The model name + messages: The messages to count tokens for + + Raises: + ValueError: If the request is invalid + """ + if not model: + raise ValueError("model parameter is required") + + if not messages: + raise ValueError("messages parameter is required") + + if not isinstance(messages, list): + raise ValueError("messages must be a list") + + for i, message in enumerate(messages): + if not isinstance(message, dict): + raise ValueError(f"Message {i} must be a dictionary") + + if "role" not in message: + raise ValueError(f"Message {i} must have a 'role' field") + + if "content" not in message: + raise ValueError(f"Message {i} must have a 'content' field") diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index 795f9a4cd0..8fa7bb7e65 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -45,6 +45,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: tools: Optional[List[Dict]] = None, top_k: Optional[int] = None, top_p: Optional[float] = None, + output_format: Optional[Dict] = None, extra_kwargs: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """Prepare kwargs for litellm.completion/acompletion""" @@ -76,6 +77,8 @@ class LiteLLMMessagesToCompletionTransformationHandler: request_data["top_k"] = top_k if top_p is not None: request_data["top_p"] = top_p + if output_format: + request_data["output_format"] = output_format openai_request = ANTHROPIC_ADAPTER.translate_completion_input_params( request_data @@ -130,6 +133,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: tools: Optional[List[Dict]] = None, top_k: Optional[int] = None, top_p: Optional[float] = None, + output_format: Optional[Dict] = None, **kwargs, ) -> Union[AnthropicMessagesResponse, AsyncIterator]: """Handle non-Anthropic models asynchronously using the adapter""" @@ -148,6 +152,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: tools=tools, top_k=top_k, top_p=top_p, + output_format=output_format, extra_kwargs=kwargs, ) ) @@ -189,6 +194,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: tools: Optional[List[Dict]] = None, top_k: Optional[int] = None, top_p: Optional[float] = None, + output_format: Optional[Dict] = None, _is_async: bool = False, **kwargs, ) -> Union[ @@ -212,6 +218,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: tools=tools, top_k=top_k, top_p=top_p, + output_format=output_format, **kwargs, ) @@ -230,6 +237,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: tools=tools, top_k=top_k, top_p=top_p, + output_format=output_format, extra_kwargs=kwargs, ) ) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 877e47a9ae..1706f045f1 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -172,7 +172,7 @@ class LiteLLMAnthropicMessagesAdapter: """ Which anthropic params, we need to translate to the openai format. """ - return ["messages", "metadata", "system", "tool_choice", "tools", "thinking"] + return ["messages", "metadata", "system", "tool_choice", "tools", "thinking", "output_format"] def translate_anthropic_messages_to_openai( # noqa: PLR0915 self, @@ -554,6 +554,42 @@ class LiteLLMAnthropicMessagesAdapter: return new_tools + def translate_anthropic_output_format_to_openai( + self, output_format: Any + ) -> Optional[Dict[str, Any]]: + """ + Translate Anthropic's output_format to OpenAI's response_format. + + Anthropic output_format: {"type": "json_schema", "schema": {...}} + OpenAI response_format: {"type": "json_schema", "json_schema": {"name": "...", "schema": {...}}} + + Args: + output_format: Anthropic output_format dict with 'type' and 'schema' + + Returns: + OpenAI-compatible response_format dict, or None if invalid + """ + if not isinstance(output_format, dict): + return None + + output_type = output_format.get("type") + if output_type != "json_schema": + return None + + schema = output_format.get("schema") + if not schema: + return None + + # Convert to OpenAI response_format structure + return { + "type": "json_schema", + "json_schema": { + "name": "structured_output", + "schema": schema, + "strict": True, + }, + } + def translate_anthropic_to_openai( self, anthropic_message_request: AnthropicMessagesRequest ) -> ChatCompletionRequest: @@ -636,6 +672,16 @@ class LiteLLMAnthropicMessagesAdapter: if reasoning_effort: new_kwargs["reasoning_effort"] = reasoning_effort + ## CONVERT OUTPUT_FORMAT to RESPONSE_FORMAT + if "output_format" in anthropic_message_request: + output_format = anthropic_message_request["output_format"] + if output_format: + response_format = self.translate_anthropic_output_format_to_openai( + output_format=output_format + ) + if response_format: + new_kwargs["response_format"] = response_format + translatable_params = self.translatable_anthropic_params() for k, v in anthropic_message_request.items(): if k not in translatable_params: # pass remaining params as is diff --git a/litellm/llms/anthropic/experimental_pass_through/architecture.md b/litellm/llms/anthropic/experimental_pass_through/architecture.md new file mode 100644 index 0000000000..b939723513 --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/architecture.md @@ -0,0 +1,51 @@ +# Anthropic Messages Pass-Through Architecture + +## Request Flow + +```mermaid +flowchart TD + A[litellm.anthropic.messages.acreate] --> B{Provider?} + + B -->|anthropic| C[AnthropicMessagesConfig] + B -->|azure_ai| D[AzureAnthropicMessagesConfig] + B -->|bedrock invoke| E[BedrockAnthropicMessagesConfig] + B -->|vertex_ai| F[VertexAnthropicMessagesConfig] + B -->|Other providers| G[LiteLLMAnthropicMessagesAdapter] + + C --> H[Direct Anthropic API] + D --> I[Azure AI Foundry API] + E --> J[Bedrock Invoke API] + F --> K[Vertex AI API] + + G --> L[translate_anthropic_to_openai] + L --> M[litellm.completion] + M --> N[Provider API] + N --> O[translate_openai_response_to_anthropic] + O --> P[Anthropic Response Format] + + H --> P + I --> P + J --> P + K --> P +``` + +## Adapter Flow (Non-Native Providers) + +```mermaid +sequenceDiagram + participant User + participant Handler as anthropic_messages_handler + participant Adapter as LiteLLMAnthropicMessagesAdapter + participant LiteLLM as litellm.completion + participant Provider as Provider API + + User->>Handler: Anthropic Messages Request + Handler->>Adapter: translate_anthropic_to_openai() + Note over Adapter: messages, tools, thinking,
output_format → response_format + Adapter->>LiteLLM: OpenAI Format Request + LiteLLM->>Provider: Provider-specific Request + Provider->>LiteLLM: Provider Response + LiteLLM->>Adapter: OpenAI Format Response + Adapter->>Handler: translate_openai_response_to_anthropic() + Handler->>User: Anthropic Messages Response +``` diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index f67e4c8382..308bf367d0 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -17,7 +17,11 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( from litellm.types.llms.anthropic_tool_search import get_tool_search_beta_header from litellm.types.router import GenericLiteLLMParams -from ...common_utils import AnthropicError, AnthropicModelInfo +from ...common_utils import ( + AnthropicError, + AnthropicModelInfo, + optionally_handle_anthropic_oauth, +) DEFAULT_ANTHROPIC_API_BASE = "https://api.anthropic.com" DEFAULT_ANTHROPIC_API_VERSION = "2023-06-01" @@ -38,6 +42,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): "tool_choice", "thinking", "context_management", + "output_format", # TODO: Add Anthropic `metadata` support # "metadata", ] @@ -68,8 +73,11 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): ) -> Tuple[dict, Optional[str]]: import os + # Check for Anthropic OAuth token in Authorization header + headers, api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key) if api_key is None: api_key = os.getenv("ANTHROPIC_API_KEY") + if "x-api-key" not in headers and api_key: headers["x-api-key"] = api_key if "anthropic-version" not in headers: @@ -162,27 +170,32 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): ) -> dict: """ Auto-inject anthropic-beta headers based on features used. - + Handles: - context_management: adds 'context-management-2025-06-27' - tool_search: adds provider-specific tool search header - + - output_format: adds 'structured-outputs-2025-11-13' + Args: headers: Request headers dict - optional_params: Optional parameters including tools, context_management + optional_params: Optional parameters including tools, context_management, output_format custom_llm_provider: Provider name for looking up correct tool search header """ beta_values: set = set() - + # Get existing beta headers if any existing_beta = headers.get("anthropic-beta") if existing_beta: beta_values.update(b.strip() for b in existing_beta.split(",")) - + # Check for context management if optional_params.get("context_management") is not None: beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value) - + + # Check for structured outputs + if optional_params.get("output_format") is not None: + beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.STRUCTURED_OUTPUT_2025_09_25.value) + # Check for tool search tools tools = optional_params.get("tools") if tools: @@ -191,8 +204,8 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): # Use provider-specific tool search header tool_search_header = get_tool_search_beta_header(custom_llm_provider) beta_values.add(tool_search_header) - + if beta_values: headers["anthropic-beta"] = ",".join(sorted(beta_values)) - + return headers diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/__init__.py b/litellm/llms/azure_ai/anthropic/count_tokens/__init__.py new file mode 100644 index 0000000000..9605d401f8 --- /dev/null +++ b/litellm/llms/azure_ai/anthropic/count_tokens/__init__.py @@ -0,0 +1,19 @@ +""" +Azure AI Anthropic CountTokens API implementation. +""" + +from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( + AzureAIAnthropicCountTokensHandler, +) +from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import ( + AzureAIAnthropicTokenCounter, +) +from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( + AzureAIAnthropicCountTokensConfig, +) + +__all__ = [ + "AzureAIAnthropicCountTokensHandler", + "AzureAIAnthropicCountTokensConfig", + "AzureAIAnthropicTokenCounter", +] diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py new file mode 100644 index 0000000000..52a0bb8bb0 --- /dev/null +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -0,0 +1,127 @@ +""" +Azure AI Anthropic CountTokens API handler. + +Uses httpx for HTTP requests with Azure authentication. +""" + +from typing import Any, Dict, List, Optional, Union + +import httpx + +import litellm +from litellm._logging import verbose_logger +from litellm.llms.anthropic.common_utils import AnthropicError +from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( + AzureAIAnthropicCountTokensConfig, +) +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + + +class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): + """ + Handler for Azure AI Anthropic CountTokens API requests. + + Uses httpx for HTTP requests with Azure authentication. + """ + + async def handle_count_tokens_request( + self, + model: str, + messages: List[Dict[str, Any]], + api_key: str, + api_base: str, + litellm_params: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + ) -> Dict[str, Any]: + """ + Handle a CountTokens request using httpx with Azure authentication. + + Args: + model: The model identifier (e.g., "claude-3-5-sonnet") + messages: The messages to count tokens for + api_key: The Azure AI API key + api_base: The Azure AI API base URL + litellm_params: Optional LiteLLM parameters + timeout: Optional timeout for the request (defaults to litellm.request_timeout) + + Returns: + Dictionary containing token count response + + Raises: + AnthropicError: If the API request fails + """ + try: + # Validate the request + self.validate_request(model, messages) + + verbose_logger.debug( + f"Processing Azure AI Anthropic CountTokens request for model: {model}" + ) + + # Transform request to Anthropic format + request_body = self.transform_request_to_count_tokens( + model=model, + messages=messages, + ) + + verbose_logger.debug(f"Transformed request: {request_body}") + + # Get endpoint URL + endpoint_url = self.get_count_tokens_endpoint(api_base) + + verbose_logger.debug(f"Making request to: {endpoint_url}") + + # Get required headers with Azure authentication + headers = self.get_required_headers( + api_key=api_key, + litellm_params=litellm_params, + ) + + # Use LiteLLM's async httpx client + async_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders.AZURE_AI + ) + + # Use provided timeout or fall back to litellm.request_timeout + request_timeout = timeout if timeout is not None else litellm.request_timeout + + response = await async_client.post( + endpoint_url, + headers=headers, + json=request_body, + timeout=request_timeout, + ) + + verbose_logger.debug(f"Response status: {response.status_code}") + + if response.status_code != 200: + error_text = response.text + verbose_logger.error(f"Azure AI Anthropic API error: {error_text}") + raise AnthropicError( + status_code=response.status_code, + message=error_text, + ) + + azure_response = response.json() + + verbose_logger.debug(f"Azure AI Anthropic response: {azure_response}") + + # Return Anthropic-compatible response directly - no transformation needed + return azure_response + + except AnthropicError: + # Re-raise Anthropic exceptions as-is + raise + except httpx.HTTPStatusError as e: + # HTTP errors - preserve the actual status code + verbose_logger.error(f"HTTP error in CountTokens handler: {str(e)}") + raise AnthropicError( + status_code=e.response.status_code, + message=e.response.text, + ) + except Exception as e: + verbose_logger.error(f"Error in CountTokens handler: {str(e)}") + raise AnthropicError( + status_code=500, + message=f"CountTokens processing error: {str(e)}", + ) diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py new file mode 100644 index 0000000000..14f9280007 --- /dev/null +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -0,0 +1,119 @@ +""" +Azure AI Anthropic Token Counter implementation using the CountTokens API. +""" + +import os +from typing import Any, Dict, List, Optional + +from litellm._logging import verbose_logger +from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( + AzureAIAnthropicCountTokensHandler, +) +from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.types.utils import LlmProviders, TokenCountResponse + +# Global handler instance - reuse across all token counting requests +azure_ai_anthropic_count_tokens_handler = AzureAIAnthropicCountTokensHandler() + + +class AzureAIAnthropicTokenCounter(BaseTokenCounter): + """Token counter implementation for Azure AI Anthropic provider using the CountTokens API.""" + + def should_use_token_counting_api( + self, + custom_llm_provider: Optional[str] = None, + ) -> bool: + return custom_llm_provider == LlmProviders.AZURE_AI.value + + async def count_tokens( + self, + model_to_use: str, + messages: Optional[List[Dict[str, Any]]], + contents: Optional[List[Dict[str, Any]]], + deployment: Optional[Dict[str, Any]] = None, + request_model: str = "", + ) -> Optional[TokenCountResponse]: + """ + Count tokens using Azure AI Anthropic's CountTokens API. + + Args: + model_to_use: The model identifier + messages: The messages to count tokens for + contents: Alternative content format (not used for Anthropic) + deployment: Deployment configuration containing litellm_params + request_model: The original request model name + + Returns: + TokenCountResponse with token count, or None if counting fails + """ + from litellm.llms.anthropic.common_utils import AnthropicError + + if not messages: + return None + + deployment = deployment or {} + litellm_params = deployment.get("litellm_params", {}) + + # Get Azure AI API key from deployment config or environment + api_key = litellm_params.get("api_key") + if not api_key: + api_key = os.getenv("AZURE_AI_API_KEY") + + # Get API base from deployment config or environment + api_base = litellm_params.get("api_base") + if not api_base: + api_base = os.getenv("AZURE_AI_API_BASE") + + if not api_key: + verbose_logger.warning("No Azure AI API key found for token counting") + return None + + if not api_base: + verbose_logger.warning("No Azure AI API base found for token counting") + return None + + try: + result = await azure_ai_anthropic_count_tokens_handler.handle_count_tokens_request( + model=model_to_use, + messages=messages, + api_key=api_key, + api_base=api_base, + litellm_params=litellm_params, + ) + + if result is not None: + return TokenCountResponse( + total_tokens=result.get("input_tokens", 0), + request_model=request_model, + model_used=model_to_use, + tokenizer_type="azure_ai_anthropic_api", + original_response=result, + ) + except AnthropicError as e: + verbose_logger.warning( + f"Azure AI Anthropic CountTokens API error: status={e.status_code}, message={e.message}" + ) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="azure_ai_anthropic_api", + error=True, + error_message=e.message, + status_code=e.status_code, + ) + except Exception as e: + verbose_logger.warning( + f"Error calling Azure AI Anthropic CountTokens API: {e}" + ) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="azure_ai_anthropic_api", + error=True, + error_message=str(e), + status_code=500, + ) + + return None diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py new file mode 100644 index 0000000000..e284595cc8 --- /dev/null +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -0,0 +1,88 @@ +""" +Azure AI Anthropic CountTokens API transformation logic. + +Extends the base Anthropic CountTokens transformation with Azure authentication. +""" + +from typing import Any, Dict, Optional + +from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION +from litellm.llms.anthropic.count_tokens.transformation import ( + AnthropicCountTokensConfig, +) +from litellm.llms.azure.common_utils import BaseAzureLLM +from litellm.types.router import GenericLiteLLMParams + + +class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): + """ + Configuration and transformation logic for Azure AI Anthropic CountTokens API. + + Extends AnthropicCountTokensConfig with Azure authentication. + Azure AI Anthropic uses the same endpoint format but with Azure auth headers. + """ + + def get_required_headers( + self, + api_key: str, + litellm_params: Optional[Dict[str, Any]] = None, + ) -> Dict[str, str]: + """ + Get the required headers for the Azure AI Anthropic CountTokens API. + + Uses Azure authentication (api-key header) instead of Anthropic's x-api-key. + + Args: + api_key: The Azure AI API key + litellm_params: Optional LiteLLM parameters for additional auth config + + Returns: + Dictionary of required headers with Azure authentication + """ + # Start with base headers + headers = { + "Content-Type": "application/json", + "anthropic-version": "2023-06-01", + "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, + } + + # Use Azure authentication + litellm_params = litellm_params or {} + if "api_key" not in litellm_params: + litellm_params["api_key"] = api_key + + litellm_params_obj = GenericLiteLLMParams(**litellm_params) + + # Get Azure auth headers + azure_headers = BaseAzureLLM._base_validate_azure_environment( + headers={}, litellm_params=litellm_params_obj + ) + + # Merge Azure auth headers + headers.update(azure_headers) + + return headers + + def get_count_tokens_endpoint(self, api_base: str) -> str: + """ + Get the Azure AI Anthropic CountTokens API endpoint. + + Args: + api_base: The Azure AI API base URL + (e.g., https://my-resource.services.ai.azure.com or + https://my-resource.services.ai.azure.com/anthropic) + + Returns: + The endpoint URL for the CountTokens API + """ + # Azure AI Anthropic endpoint format: + # https://.services.ai.azure.com/anthropic/v1/messages/count_tokens + api_base = api_base.rstrip("/") + + # Ensure the URL has /anthropic path + if not api_base.endswith("/anthropic"): + if "/anthropic" not in api_base: + api_base = f"{api_base}/anthropic" + + # Add the count_tokens path + return f"{api_base}/v1/messages/count_tokens" diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 9487c7f83f..01a3f5766c 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -1,17 +1,22 @@ from typing import List, Literal, Optional import litellm -from litellm.llms.base_llm.base_utils import BaseLLMModelInfo +from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues class AzureFoundryModelInfo(BaseLLMModelInfo): + """Model info for Azure AI / Azure Foundry models.""" + + def __init__(self, model: Optional[str] = None): + self._model = model + @staticmethod def get_azure_ai_route(model: str) -> Literal["agents", "default"]: """ Get the Azure AI route for the given model. - + Similar to BedrockModelInfo.get_bedrock_route(). """ if "agents/" in model: @@ -20,34 +25,54 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): @staticmethod def get_api_base(api_base: Optional[str] = None) -> Optional[str]: - return ( - api_base - or litellm.api_base - or get_secret_str("AZURE_AI_API_BASE") - ) - + return api_base or litellm.api_base or get_secret_str("AZURE_AI_API_BASE") + @staticmethod def get_api_key(api_key: Optional[str] = None) -> Optional[str]: return ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("AZURE_AI_API_KEY") - ) - + api_key + or litellm.api_key + or litellm.openai_key + or get_secret_str("AZURE_AI_API_KEY") + ) + @property def api_version(self, api_version: Optional[str] = None) -> Optional[str]: api_version = ( - api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") + api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") ) return api_version - + + def get_token_counter(self) -> Optional[BaseTokenCounter]: + """ + Factory method to create a token counter for Azure AI. + + Returns: + AzureAIAnthropicTokenCounter for Claude models, None otherwise. + """ + # Only return token counter for Claude models + if self._model and "claude" in self._model.lower(): + from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import ( + AzureAIAnthropicTokenCounter, + ) + + return AzureAIAnthropicTokenCounter() + return None + + def get_models( + self, api_key: Optional[str] = None, api_base: Optional[str] = None + ) -> List[str]: + """ + Returns a list of models supported by Azure AI. + + Azure AI doesn't have a standard model listing endpoint, + so this returns an empty list. + """ + return [] + ######################################################### # Not implemented methods ######################################################### - @staticmethod def get_base_model(model: str) -> Optional[str]: @@ -64,4 +89,6 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): api_base: Optional[str] = None, ) -> dict: """Azure Foundry sends api key in query params""" - raise NotImplementedError("Azure Foundry does not support environment validation") + raise NotImplementedError( + "Azure Foundry does not support environment validation" + ) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 293ee1caaf..81ebb5a360 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -271,8 +271,12 @@ class AmazonAnthropicClaudeMessagesConfig( # 4. Remove `ttl` field from cache_control in messages (Bedrock doesn't support it) self._remove_ttl_from_cache_control(anthropic_messages_request) + + # 5. `output_format` is not supported on Bedrock invoke + if "output_format" in anthropic_messages_request: + anthropic_messages_request.pop("output_format", None) - # 5. AUTO-INJECT beta headers based on features used + # 6. AUTO-INJECT beta headers based on features used anthropic_model_info = AnthropicModelInfo() tools = anthropic_messages_optional_request_params.get("tools") messages_typed = cast(List[AllMessageValues], messages) diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 4aefe58394..b4f9cbe42d 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -75,5 +75,16 @@ "gmi": { "base_url": "https://api.gmi-serving.com/v1", "api_key_env": "GMI_API_KEY" + }, + "sarvam": { + "base_url": "https://api.sarvam.ai/v1", + "api_key_env": "SARVAM_API_KEY", + "base_class": "openai_gpt", + "param_mappings": { + "max_completion_tokens": "max_tokens" + }, + "headers": { + "api-subscription-key": "{api_key}" + } } } diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 8f1338db92..96e0963a92 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -72,17 +72,64 @@ def _convert_detail_to_media_resolution_enum( return {"level": "MEDIA_RESOLUTION_MEDIUM"} elif detail == "high": return {"level": "MEDIA_RESOLUTION_HIGH"} + elif detail == "ultra_high": + return {"level": "MEDIA_RESOLUTION_ULTRA_HIGH"} return None -def _process_gemini_image( - image_url: str, +def _apply_gemini_3_metadata( + part: PartType, + model: Optional[str], + media_resolution_enum: Optional[Dict[str, str]], + video_metadata: Optional[Dict[str, Any]], +) -> PartType: + """ + Apply the unique media_resolution and video_metadata parameters of Gemini 3+ + """ + if model is None: + return part + + from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig + + if not VertexGeminiConfig._is_gemini_3_or_newer(model): + return part + + part_dict = dict(part) + + if media_resolution_enum is not None: + part_dict["media_resolution"] = media_resolution_enum + + if video_metadata is not None: + gemini_video_metadata = {} + if "fps" in video_metadata: + gemini_video_metadata["fps"] = video_metadata["fps"] + if "start_offset" in video_metadata: + gemini_video_metadata["startOffset"] = video_metadata["start_offset"] + if "end_offset" in video_metadata: + gemini_video_metadata["endOffset"] = video_metadata["end_offset"] + if gemini_video_metadata: + part_dict["video_metadata"] = gemini_video_metadata + + return cast(PartType, part_dict) + + +def _process_gemini_media( + image_url: str, format: Optional[str] = None, media_resolution_enum: Optional[Dict[str, str]] = None, model: Optional[str] = None, + video_metadata: Optional[Dict[str, Any]] = None, ) -> PartType: """ - Given an image URL, return the appropriate PartType for Gemini + Given a media URL (image, audio, or video), return the appropriate PartType for Gemini + By the way, actually video_metadata can only be used with videos; it cannot be used with images, audio, or files. However, I haven't made any special handling because vertex returns a parameter error. + + Args: + image_url: The URL or base64 string of the media (image, audio, or video) + format: The MIME type of the media + media_resolution_enum: Media resolution level (for Gemini 3+) + model: The model name (to check version compatibility) + video_metadata: Video-specific metadata (fps, start_offset, end_offset) """ try: @@ -104,14 +151,9 @@ def _process_gemini_image( mime_type = format file_data = FileDataType(mime_type=mime_type, file_uri=image_url) part: PartType = {"file_data": file_data} - - if media_resolution_enum is not None and model is not None: - from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig - if VertexGeminiConfig._is_gemini_3_or_newer(model): - part_dict = dict(part) - part_dict["media_resolution"] = media_resolution_enum - return cast(PartType, part_dict) - return part + return _apply_gemini_3_metadata( + part, model, media_resolution_enum, video_metadata + ) elif ( "https://" in image_url and (image_type := format or _get_image_mime_type_from_url(image_url)) @@ -119,27 +161,16 @@ def _process_gemini_image( ): file_data = FileDataType(mime_type=image_type, file_uri=image_url) part = {"file_data": file_data} - - if media_resolution_enum is not None and model is not None: - from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig - if VertexGeminiConfig._is_gemini_3_or_newer(model): - part_dict = dict(part) - part_dict["media_resolution"] = media_resolution_enum - return cast(PartType, part_dict) - return part + return _apply_gemini_3_metadata( + part, model, media_resolution_enum, video_metadata + ) elif "http://" in image_url or "https://" in image_url or "base64" in image_url: image = convert_to_anthropic_image_obj(image_url, format=format) _blob: BlobType = {"data": image["data"], "mime_type": image["media_type"]} - part = {"inline_data": cast(BlobType, _blob)} - - if media_resolution_enum is not None and model is not None: - from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig - if VertexGeminiConfig._is_gemini_3_or_newer(model): - part_dict = dict(part) - part_dict["media_resolution"] = media_resolution_enum - return cast(PartType, part_dict) - return part + return _apply_gemini_3_metadata( + part, model, media_resolution_enum, video_metadata + ) raise Exception("Invalid image received - {}".format(image_url)) except Exception as e: raise e @@ -253,8 +284,8 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 media_resolution_enum = _convert_detail_to_media_resolution_enum(detail) else: image_url = img_element["image_url"] - _part = _process_gemini_image( - image_url=image_url, + _part = _process_gemini_media( + image_url=image_url, format=format, media_resolution_enum=media_resolution_enum, model=model, @@ -279,7 +310,7 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 ) ) ) - _part = _process_gemini_image( + _part = _process_gemini_media( image_url=openai_image_str, format=audio_format_modified, model=model, @@ -290,16 +321,24 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 file_id = file_element["file"].get("file_id") format = file_element["file"].get("format") file_data = file_element["file"].get("file_data") + detail = file_element["file"].get("detail") + video_metadata = file_element["file"].get("video_metadata") passed_file = file_id or file_data if passed_file is None: raise Exception( "Unknown file type. Please pass in a file_id or file_data" ) + + # Convert detail to media_resolution_enum + media_resolution_enum = _convert_detail_to_media_resolution_enum(detail) + try: - _part = _process_gemini_image( - image_url=passed_file, + _part = _process_gemini_media( + image_url=passed_file, format=format, model=model, + media_resolution_enum=media_resolution_enum, + video_metadata=video_metadata, ) _parts.append(_part) except Exception: 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 dd34fbd772..2d2e07e74d 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 @@ -1018,25 +1018,34 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): optional_params["parallel_tool_calls"] = value elif param == "seed": optional_params["seed"] = value - elif param == "reasoning_effort" and isinstance(value, str): - # Validate no conflict with thinking_level - VertexGeminiConfig._validate_thinking_config_conflicts( - optional_params=optional_params, - param_name="reasoning_effort", - param_description="thinking_budget", - ) - if VertexGeminiConfig._is_gemini_3_or_newer(model): - optional_params["thinkingConfig"] = ( - VertexGeminiConfig._map_reasoning_effort_to_thinking_level( - value, model - ) + elif param == "reasoning_effort": + # Extract effort value - handle both string and dict formats + # Dict format comes from OpenAI Agents SDK: {"effort": "high", "summary": "auto"} + effort_value: Optional[str] = None + if isinstance(value, str): + effort_value = value + elif isinstance(value, dict): + effort_value = value.get("effort") + + if effort_value is not None: + # Validate no conflict with thinking_level + VertexGeminiConfig._validate_thinking_config_conflicts( + optional_params=optional_params, + param_name="reasoning_effort", + param_description="thinking_budget", ) - else: - optional_params["thinkingConfig"] = ( - VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( - value, model + if VertexGeminiConfig._is_gemini_3_or_newer(model): + optional_params["thinkingConfig"] = ( + VertexGeminiConfig._map_reasoning_effort_to_thinking_level( + effort_value, model + ) + ) + else: + optional_params["thinkingConfig"] = ( + VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( + effort_value, model + ) ) - ) elif param == "thinking": # Validate no conflict with thinking_level VertexGeminiConfig._validate_thinking_config_conflicts( diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index 0bedef3276..fc75376c0c 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -117,4 +117,9 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert anthropic_messages_request.pop( "model", None ) # do not pass model in request body to vertex ai + + anthropic_messages_request.pop( + "output_format", None + ) # do not pass output_format in request body to vertex ai - vertex ai does not support output_format as yet + return anthropic_messages_request diff --git a/litellm/main.py b/litellm/main.py index ae27b4145b..ea41919e19 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -367,7 +367,7 @@ class AsyncCompletions: @tracer.wrap() @client -async def acompletion( +async def acompletion( # noqa: PLR0915 model: str, # Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create messages: List = [], @@ -599,7 +599,16 @@ async def acompletion( ctx = contextvars.copy_context() func_with_context = partial(ctx.run, func) - init_response = await loop.run_in_executor(None, func_with_context) + # Wrap with timeout if specified + if timeout is not None: + timeout_value = float(timeout) if not isinstance(timeout, (int, float)) else timeout + init_response = await asyncio.wait_for( + loop.run_in_executor(None, func_with_context), + timeout=timeout_value + ) + else: + init_response = await loop.run_in_executor(None, func_with_context) + if isinstance(init_response, dict) or isinstance( init_response, ModelResponse ): ## CACHING SCENARIO @@ -607,7 +616,11 @@ async def acompletion( response = ModelResponse(**init_response) response = init_response elif asyncio.iscoroutine(init_response): - response = await init_response + if timeout is not None: + timeout_value = float(timeout) if not isinstance(timeout, (int, float)) else timeout + response = await asyncio.wait_for(init_response, timeout=timeout_value) + else: + response = await init_response else: response = init_response # type: ignore @@ -624,6 +637,14 @@ async def acompletion( loop=loop ) # sets the logging event loop if the user does sync streaming (e.g. on proxy for sagemaker calls) return response + except asyncio.TimeoutError: + custom_llm_provider = custom_llm_provider or "openai" + from litellm.exceptions import Timeout + raise Timeout( + message=f"Request timed out after {timeout} seconds", + model=model, + llm_provider=custom_llm_provider, + ) except Exception as e: custom_llm_provider = custom_llm_provider or "openai" raise exception_type( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 135b0d46ed..a74b80e737 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -12696,8 +12696,8 @@ "supports_web_search": true }, "gemini-2.5-flash-lite": { - "cache_read_input_token_cost": 2.5e-08, - "input_cost_per_audio_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, @@ -12741,7 +12741,7 @@ "supports_web_search": true }, "gemini-2.5-flash-lite-preview-09-2025": { - "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", @@ -14532,8 +14532,8 @@ "supports_web_search": true }, "gemini/gemini-2.5-flash-lite": { - "cache_read_input_token_cost": 2.5e-08, - "input_cost_per_audio_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", "max_audio_length_hours": 8.4, @@ -14579,7 +14579,7 @@ "tpm": 250000 }, "gemini/gemini-2.5-flash-lite-preview-09-2025": { - "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", @@ -34062,5 +34062,18 @@ "output_cost_per_token": 0, "litellm_provider": "llamagate", "mode": "embedding" + }, + "sarvam/sarvam-m": { + "cache_creation_input_token_cost": 0, + "cache_creation_input_token_cost_above_1hr": 0, + "cache_read_input_token_cost": 0, + "input_cost_per_token": 0, + "litellm_provider": "sarvam", + "max_input_tokens": 8192, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "chat", + "output_cost_per_token": 0, + "supports_reasoning": true } } diff --git a/litellm/passthrough/utils.py b/litellm/passthrough/utils.py index 4bf66d4988..fbbf9cd258 100644 --- a/litellm/passthrough/utils.py +++ b/litellm/passthrough/utils.py @@ -3,6 +3,8 @@ from urllib.parse import parse_qs import httpx +from litellm.constants import PASS_THROUGH_HEADER_PREFIX + class BasePassthroughUtils: @staticmethod @@ -27,7 +29,11 @@ class BasePassthroughUtils: forward_headers: Optional[bool] = False, ): """ - Helper to forward headers from original request + Helper to forward headers from original request. + + Also handles 'x-pass-' prefixed headers which are always forwarded + with the prefix stripped, regardless of forward_headers setting. + e.g., 'x-pass-anthropic-beta: value' becomes 'anthropic-beta: value' """ if forward_headers is True: # Header We Should NOT forward @@ -36,6 +42,14 @@ class BasePassthroughUtils: # Combine request headers with custom headers headers = {**request_headers, **headers} + + # Always process x-pass- prefixed headers (strip prefix and forward) + for header_name, header_value in request_headers.items(): + if header_name.lower().startswith(PASS_THROUGH_HEADER_PREFIX): + # Strip the 'x-pass-' prefix to get the actual header name + actual_header_name = header_name[len(PASS_THROUGH_HEADER_PREFIX) :] + headers[actual_header_name] = header_value + return headers class CommonUtils: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 76a2834485..03652ae155 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -6,6 +6,8 @@ LiteLLM MCP Server Routes import asyncio import contextlib from datetime import datetime +import traceback +import uuid from typing import Any, AsyncIterator, Dict, List, Optional, Tuple, Union, cast from fastapi import FastAPI, HTTPException @@ -13,6 +15,7 @@ from pydantic import AnyUrl, ConfigDict from starlette.types import Receive, Scope, Send from litellm._logging import verbose_logger +from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, @@ -25,8 +28,8 @@ from litellm.proxy._experimental.mcp_server.utils import ( from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer -from litellm.types.utils import StandardLoggingMCPToolCall -from litellm.utils import client +from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall +from litellm.utils import Rules, client, function_setup # Check if MCP is available # "mcp" requires python 3.10 or higher, but several litellm users use python 3.8 @@ -226,6 +229,8 @@ if MCP_AVAILABLE: mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, + log_list_tools_to_spendlogs=True, + list_tools_log_source="mcp_protocol", ) verbose_logger.info( f"MCP list_tools - Successfully returned {len(tools)} tools" @@ -733,13 +738,15 @@ if MCP_AVAILABLE: return server_auth_header, extra_headers - async def _get_tools_from_mcp_servers( + async def _get_tools_from_mcp_servers( # noqa: PLR0915 user_api_key_auth: Optional[UserAPIKeyAuth], mcp_auth_header: Optional[str], mcp_servers: Optional[List[str]], mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, + log_list_tools_to_spendlogs: bool = False, + list_tools_log_source: Optional[str] = None, ) -> List[MCPTool]: """ Helper method to fetch tools from MCP servers based on server filtering criteria. @@ -757,67 +764,188 @@ if MCP_AVAILABLE: if not MCP_AVAILABLE: return [] - allowed_mcp_servers = await _get_allowed_mcp_servers( - user_api_key_auth=user_api_key_auth, - mcp_servers=mcp_servers, - ) + list_tools_start_time = datetime.now() + litellm_logging_obj: Optional[LiteLLMLoggingObj] = None + list_tools_request_data: Dict[str, Any] = {} - # Decide whether to add prefix based on number of allowed servers - add_prefix = not (len(allowed_mcp_servers) == 1) + if log_list_tools_to_spendlogs: + # This is intentionally minimal: only async_success_handler / post_call_failure_hook + rules_obj = Rules() + list_tools_call_id = str(uuid.uuid4()) + spend_logs_metadata: Dict[str, Any] = { + "mcp_operation": "list_tools", + } + if isinstance(list_tools_log_source, str): + spend_logs_metadata["source"] = list_tools_log_source + if isinstance(mcp_servers, list): + spend_logs_metadata["requested_mcp_servers"] = mcp_servers - async def _fetch_and_filter_server_tools(server: MCPServer) -> List[MCPTool]: - """Fetch and filter tools from a single server with error handling.""" - if server is None: - return [] + list_tools_request_data = { + "model": "MCP: list_tools", + "call_type": CallTypes.list_mcp_tools.value, + "litellm_call_id": list_tools_call_id, + "metadata": { + "spend_logs_metadata": spend_logs_metadata, + }, + # Provide a small input payload for standard logging + "input": [ + { + "role": "system", + "content": { + "mcp_operation": "list_tools", + "requested_mcp_servers": mcp_servers, + }, + } + ], + } - server_auth_header, extra_headers = _prepare_mcp_server_headers( - server=server, - mcp_server_auth_headers=mcp_server_auth_headers, - mcp_auth_header=mcp_auth_header, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - ) + # Attach user identifiers when available (matches call_mcp_tool style) + if user_api_key_auth is not None: + user_api_key = getattr(user_api_key_auth, "api_key", None) + if user_api_key: + cast(dict, list_tools_request_data["metadata"])[ + "user_api_key" + ] = user_api_key + + user_identifier = getattr( + user_api_key_auth, "end_user_id", None + ) or getattr(user_api_key_auth, "user_id", None) + if user_identifier: + list_tools_request_data["user"] = user_identifier try: - tools = await global_mcp_server_manager._get_tools_from_server( + litellm_logging_obj, _ = function_setup( + original_function="list_mcp_tools", + rules_obj=rules_obj, + start_time=list_tools_start_time, + **list_tools_request_data, + ) + if litellm_logging_obj: + litellm_logging_obj.call_type = CallTypes.list_mcp_tools.value + litellm_logging_obj.model = "MCP: list_tools" + except Exception as logging_error: + verbose_logger.debug( + "Failed to initialize logging for MCP list_tools: %s", logging_error + ) + litellm_logging_obj = None + + try: + allowed_mcp_servers = await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_servers=mcp_servers, + ) + + # Decide whether to add prefix based on number of allowed servers + add_prefix = not (len(allowed_mcp_servers) == 1) + + async def _fetch_and_filter_server_tools( + server: MCPServer, + ) -> List[MCPTool]: + """Fetch and filter tools from a single server with error handling.""" + if server is None: + return [] + + server_auth_header, extra_headers = _prepare_mcp_server_headers( server=server, - mcp_auth_header=server_auth_header, - extra_headers=extra_headers, - add_prefix=add_prefix, + mcp_server_auth_headers=mcp_server_auth_headers, + mcp_auth_header=mcp_auth_header, + oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) - filtered_tools = filter_tools_by_allowed_tools(tools, server) - filtered_tools = await filter_tools_by_key_team_permissions( - tools=filtered_tools, - server_id=server.server_id, - user_api_key_auth=user_api_key_auth, + try: + tools = await global_mcp_server_manager._get_tools_from_server( + server=server, + mcp_auth_header=server_auth_header, + extra_headers=extra_headers, + add_prefix=add_prefix, + raw_headers=raw_headers, + ) + filtered_tools = filter_tools_by_allowed_tools(tools, server) + + filtered_tools = await filter_tools_by_key_team_permissions( + tools=filtered_tools, + server_id=server.server_id, + user_api_key_auth=user_api_key_auth, + ) + + verbose_logger.debug( + f"Successfully fetched {len(tools)} tools from server {server.name}, {len(filtered_tools)} after filtering" + ) + return filtered_tools + except Exception as e: + verbose_logger.exception( + f"Error getting tools from server {server.name}: {str(e)}" + ) + return [] + + # Fetch tools from all servers in parallel + tasks = [ + _fetch_and_filter_server_tools(server) for server in allowed_mcp_servers + ] + results = await asyncio.gather(*tasks) + + # Flatten results into single list + all_tools: List[MCPTool] = [tool for tools in results for tool in tools] + + # If logging is enabled, enrich spend_logs_metadata with counts + if litellm_logging_obj: + per_server_tool_counts: Dict[str, int] = {} + for server, server_tools in zip(allowed_mcp_servers, results): + if server is None: + continue + server_key = ( + getattr(server, "server_name", None) + or getattr(server, "alias", None) + or getattr(server, "name", None) + or "unknown" + ) + per_server_tool_counts[str(server_key)] = len(server_tools) + + metadata_dict = litellm_logging_obj.model_call_details.get("metadata") + if isinstance(metadata_dict, dict): + spend_meta = metadata_dict.get("spend_logs_metadata") + if not isinstance(spend_meta, dict): + spend_meta = {} + metadata_dict["spend_logs_metadata"] = spend_meta + spend_meta["allowed_server_count"] = len(allowed_mcp_servers) + spend_meta["tool_count_total"] = len(all_tools) + spend_meta["per_server_tool_counts"] = per_server_tool_counts + + end_time = datetime.now() + await litellm_logging_obj.async_success_handler( + result=all_tools, + start_time=list_tools_start_time, + end_time=end_time, ) - verbose_logger.debug( - f"Successfully fetched {len(tools)} tools from server {server.name}, {len(filtered_tools)} after filtering" - ) - return filtered_tools - except Exception as e: - verbose_logger.exception( - f"Error getting tools from server {server.name}: {str(e)}" - ) - return [] + verbose_logger.info( + f"Successfully fetched {len(all_tools)} tools total from all MCP servers" + ) - # Fetch tools from all servers in parallel - tasks = [ - _fetch_and_filter_server_tools(server) for server in allowed_mcp_servers - ] - results = await asyncio.gather(*tasks) + return all_tools + except Exception as e: + # Only fire failure hook if logging was requested for this list-tools execution + if log_list_tools_to_spendlogs and user_api_key_auth is not None: + try: + from litellm.proxy.proxy_server import proxy_logging_obj - # Flatten results into single list - all_tools: List[MCPTool] = [tool for tools in results for tool in tools] - - verbose_logger.info( - f"Successfully fetched {len(all_tools)} tools total from all MCP servers" - ) - - return all_tools + if proxy_logging_obj: + traceback_str = traceback.format_exc( + limit=MAXIMUM_TRACEBACK_LINES_TO_LOG + ) + await proxy_logging_obj.post_call_failure_hook( + request_data=list_tools_request_data or {}, + original_exception=e, + user_api_key_dict=user_api_key_auth, + route="/mcp/list_tools", + traceback_str=traceback_str, + ) + except Exception: + verbose_logger.debug( + "Failed to log MCP list_tools failure via post_call_failure_hook" + ) + raise async def _get_prompts_from_mcp_servers( user_api_key_auth: Optional[UserAPIKeyAuth], @@ -1050,6 +1178,8 @@ if MCP_AVAILABLE: mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, + log_list_tools_to_spendlogs: bool = False, + list_tools_log_source: Optional[str] = None, ) -> List[MCPTool]: """ List all available MCP tools. @@ -1075,6 +1205,8 @@ if MCP_AVAILABLE: mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, + log_list_tools_to_spendlogs=log_list_tools_to_spendlogs, + list_tools_log_source=list_tools_log_source, ) verbose_logger.debug( f"Successfully fetched {len(managed_tools)} tools from managed MCP servers" @@ -1320,33 +1452,6 @@ if MCP_AVAILABLE: content=cast(Any, local_content), isError=False ) - ######################################################### - # Post MCP Tool Call Hook - # Allow modifying the MCP tool call response before it is returned to the user - ######################################################### - if litellm_logging_obj: - litellm_logging_obj.post_call(original_response=response) - end_time = datetime.now() - await litellm_logging_obj.async_post_mcp_tool_call_hook( - kwargs=litellm_logging_obj.model_call_details, - response_obj=response, - start_time=start_time, - end_time=end_time, - ) - # Set call_type to call_mcp_tool so cost calculator recognizes it - from litellm.types.utils import CallTypes - - litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value - # Trigger success logging to build standard_logging_object and call callbacks - # async_success_handler will: - # 1. Call _success_handler_helper_fn which recognizes call_mcp_tool - # 2. Call _process_hidden_params_and_response_cost which: - # - Calculates cost via _response_cost_calculator -> MCPCostCalculator - # - Builds standard_logging_object - # 3. Call async_log_success_event on all callbacks - await litellm_logging_obj.async_success_handler( - result=response, start_time=start_time, end_time=end_time - ) return response @client @@ -1365,49 +1470,82 @@ if MCP_AVAILABLE: Call a specific tool with the provided arguments (handles prefixed tool names). """ start_time = datetime.now() - if arguments is None: - raise HTTPException( - status_code=400, detail="Request arguments are required" + litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get( + "litellm_logging_obj", None + ) + + try: + if arguments is None: + raise HTTPException( + status_code=400, detail="Request arguments are required" + ) + + ## CHECK IF USER IS ALLOWED TO CALL THIS TOOL + allowed_mcp_server_ids = ( + await global_mcp_server_manager.get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, + ) ) - ## CHECK IF USER IS ALLOWED TO CALL THIS TOOL - allowed_mcp_server_ids = ( - await global_mcp_server_manager.get_allowed_mcp_servers( + allowed_mcp_servers: List[MCPServer] = [] + for allowed_mcp_server_id in allowed_mcp_server_ids: + allowed_server = global_mcp_server_manager.get_mcp_server_by_id( + allowed_mcp_server_id + ) + if allowed_server is not None: + allowed_mcp_servers.append(allowed_server) + + allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names( + mcp_servers=mcp_servers, + allowed_mcp_servers=allowed_mcp_servers, + ) + if not allowed_mcp_servers: + raise HTTPException( + status_code=403, + detail="User not allowed to call this tool.", + ) + + # Delegate to execute_mcp_tool for execution + response = await execute_mcp_tool( + name=name, + arguments=arguments, + allowed_mcp_servers=allowed_mcp_servers, + start_time=start_time, user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + **kwargs, ) - ) + except Exception as e: + traceback_str = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) + from litellm.proxy.proxy_server import proxy_logging_obj - allowed_mcp_servers: List[MCPServer] = [] - for allowed_mcp_server_id in allowed_mcp_server_ids: - allowed_server = global_mcp_server_manager.get_mcp_server_by_id( - allowed_mcp_server_id + if proxy_logging_obj and user_api_key_auth: + await proxy_logging_obj.post_call_failure_hook( + request_data=kwargs, + original_exception=e, + user_api_key_dict=user_api_key_auth, + route="/mcp/call_tool", + traceback_str=traceback_str, + ) + raise + + if litellm_logging_obj: + litellm_logging_obj.post_call(original_response=response) + end_time = datetime.now() + await litellm_logging_obj.async_post_mcp_tool_call_hook( + kwargs=litellm_logging_obj.model_call_details, + response_obj=response, + start_time=start_time, + end_time=end_time, ) - if allowed_server is not None: - allowed_mcp_servers.append(allowed_server) - - allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names( - mcp_servers=mcp_servers, - allowed_mcp_servers=allowed_mcp_servers, - ) - if not allowed_mcp_servers: - raise HTTPException( - status_code=403, - detail="User not allowed to call this tool.", + litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value + await litellm_logging_obj.async_success_handler( + result=response, start_time=start_time, end_time=end_time ) - - # Delegate to execute_mcp_tool for execution - return await execute_mcp_tool( - name=name, - arguments=arguments, - allowed_mcp_servers=allowed_mcp_servers, - start_time=start_time, - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - **kwargs, - ) + return response async def mcp_get_prompt( name: str, diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index 3c367eafbc..13bbf2272f 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -26,97 +26,119 @@ class KeyRotationManager: """ Manages automated key rotation based on individual key rotation schedules. """ - + def __init__(self, prisma_client: PrismaClient): self.prisma_client = prisma_client - + async def process_rotations(self): """ Main entry point - find and rotate keys that are due for rotation """ try: verbose_proxy_logger.info("Starting scheduled key rotation check...") - + # Find keys that are due for rotation keys_to_rotate = await self._find_keys_needing_rotation() - + if not keys_to_rotate: verbose_proxy_logger.debug("No keys are due for rotation at this time") return - - verbose_proxy_logger.info(f"Found {len(keys_to_rotate)} keys due for rotation") - + + verbose_proxy_logger.info( + f"Found {len(keys_to_rotate)} keys due for rotation" + ) + # Rotate each key for key in keys_to_rotate: try: await self._rotate_key(key) - key_identifier = key.key_name or (key.token[:8] + "..." if key.token else "unknown") - verbose_proxy_logger.info(f"Successfully rotated key: {key_identifier}") + key_identifier = key.key_name or ( + key.token[:8] + "..." if key.token else "unknown" + ) + verbose_proxy_logger.info( + f"Successfully rotated key: {key_identifier}" + ) except Exception as e: - key_identifier = key.key_name or (key.token[:8] + "..." if key.token else "unknown") - verbose_proxy_logger.error(f"Failed to rotate key {key_identifier}: {e}") - + key_identifier = key.key_name or ( + key.token[:8] + "..." if key.token else "unknown" + ) + verbose_proxy_logger.error( + f"Failed to rotate key {key_identifier}: {e}" + ) + except Exception as e: verbose_proxy_logger.error(f"Key rotation process failed: {e}") - + async def _find_keys_needing_rotation(self) -> List[LiteLLM_VerificationToken]: """ Find keys that are due for rotation based on their key_rotation_at timestamp. - + Logic: - Key has auto_rotate = true - key_rotation_at is null (needs initial setup) OR key_rotation_at <= now """ now = datetime.now(timezone.utc) - - keys_with_rotation = await self.prisma_client.db.litellm_verificationtoken.find_many( - where={ - "auto_rotate": True, # Only keys marked for auto rotation - "OR": [ - {"key_rotation_at": None}, # Keys that need initial rotation time setup - {"key_rotation_at": {"lte": now}} # Keys where rotation time has passed - ] - } + + keys_with_rotation = ( + await self.prisma_client.db.litellm_verificationtoken.find_many( + where={ + "auto_rotate": True, # Only keys marked for auto rotation + "OR": [ + { + "key_rotation_at": None + }, # Keys that need initial rotation time setup + { + "key_rotation_at": {"lte": now} + }, # Keys where rotation time has passed + ], + } + ) ) - + return keys_with_rotation - + def _should_rotate_key(self, key: LiteLLM_VerificationToken, now: datetime) -> bool: """ Determine if a key should be rotated based on key_rotation_at timestamp. """ if not key.rotation_interval: return False - + # If key_rotation_at is not set, rotate immediately (and set it) if key.key_rotation_at is None: return True - + # Check if the rotation time has passed return now >= key.key_rotation_at - + async def _rotate_key(self, key: LiteLLM_VerificationToken): """ Rotate a single key using existing regenerate_key_fn and call the rotation hook """ - # Create regenerate request + # Create regenerate request regenerate_request = RegenerateKeyRequest( - key=key.token or "" + key=key.token or "", + key_alias=key.key_alias, # Pass key alias to ensure correct secret is updated in AWS Secrets Manager ) - + # Create a system user for key rotation from litellm.proxy._types import UserAPIKeyAuth + system_user = UserAPIKeyAuth.get_litellm_internal_jobs_user_api_key_auth() - + # Use existing regenerate key function response = await regenerate_key_fn( data=regenerate_request, user_api_key_dict=system_user, - litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, ) - + # Update the NEW key with rotation info (regenerate_key_fn creates a new token) - if isinstance(response, GenerateKeyResponse) and response.token_id and key.rotation_interval: + if ( + isinstance(response, GenerateKeyResponse) + and response.token_id + and key.rotation_interval + ): # Calculate next rotation time using helper function now = datetime.now(timezone.utc) next_rotation_time = _calculate_key_rotation_time(key.rotation_interval) @@ -125,10 +147,10 @@ class KeyRotationManager: data={ "rotation_count": (key.rotation_count or 0) + 1, "last_rotation_at": now, - "key_rotation_at": next_rotation_time - } + "key_rotation_at": next_rotation_time, + }, ) - + # Call the existing rotation hook for notifications, audit logs, etc. if isinstance(response, GenerateKeyResponse): await KeyManagementEventHooks.async_key_rotated_hook( @@ -136,6 +158,5 @@ class KeyRotationManager: existing_key_row=key, response=response, user_api_key_dict=system_user, - litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, ) - \ No newline at end of file diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 9263bca100..50f8b2a3de 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -152,7 +152,8 @@ class KeyManagementEventHooks: ) await KeyManagementEventHooks._rotate_virtual_key_in_secret_manager( current_secret_name=initial_secret_name, - new_secret_name=data.key_alias + new_secret_name=response.key_alias + or data.key_alias or f"virtual-key-{response.token_id}", new_secret_value=response.key, ) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index c16b4c4b93..8f7dd4f8df 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,5 +1,7 @@ -from typing import Any, Dict, Optional, Union +from typing import Any, Dict, Optional, Union, TYPE_CHECKING +from litellm._logging import verbose_proxy_logger +from litellm.caching import DualCache from litellm.proxy._types import ( KeyRequestBase, LiteLLM_ManagementEndpoint_MetadataFields, @@ -11,6 +13,9 @@ from litellm.proxy._types import ( ) from litellm.proxy.utils import _premium_user_check +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient, ProxyLogging + def _user_has_admin_view(user_api_key_dict: UserAPIKeyAuth) -> bool: return ( @@ -31,6 +36,78 @@ def _is_user_team_admin( return False +async def _user_has_admin_privileges( + user_api_key_dict: UserAPIKeyAuth, + prisma_client: Optional["PrismaClient"] = None, + user_api_key_cache: Optional["DualCache"] = None, + proxy_logging_obj: Optional["ProxyLogging"] = None, +) -> bool: + """ + Check if user has admin privileges (proxy admin, team admin, or org admin). + + Args: + user_api_key_dict: User API key authentication object + prisma_client: Prisma client for database operations + user_api_key_cache: Cache for user API keys + proxy_logging_obj: Proxy logging object + + Returns: + True if user is proxy admin, team admin for any team, or org admin for any organization + """ + # Check if user is proxy admin + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return True + + # If no database connection, can't check team/org admin status + if prisma_client is None or user_api_key_dict.user_id is None: + return False + + # Get user object to check team and org admin status + from litellm.caching import DualCache as DualCacheImport + from litellm.proxy.auth.auth_checks import get_user_object + + try: + user_obj = await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache or DualCacheImport(), + user_id_upsert=False, + proxy_logging_obj=proxy_logging_obj, + ) + + if user_obj is None: + return False + + # Check if user is org admin for any organization + if user_obj.organization_memberships is not None: + for membership in user_obj.organization_memberships: + if membership.user_role == LitellmUserRoles.ORG_ADMIN.value: + return True + + # Check if user is team admin for any team + if user_obj.teams is not None and len(user_obj.teams) > 0: + # Get all teams user is in + teams = await prisma_client.db.litellm_teamtable.find_many( + where={"team_id": {"in": user_obj.teams}} + ) + + for team in teams: + team_obj = LiteLLM_TeamTable(**team.model_dump()) + if _is_user_team_admin( + user_api_key_dict=user_api_key_dict, team_obj=team_obj + ): + return True + + except Exception as e: + # If there's an error checking, default to False for security + verbose_proxy_logger.debug( + f"Error checking admin privileges for user {user_api_key_dict.user_id}: {e}" + ) + return False + + return False + + def _set_object_metadata_field( object_data: Union[ LiteLLM_TeamTable, diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 0622393ec8..6cdadfe216 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -10,7 +10,7 @@ PATCH /config/cost_margin_config - Update cost margin configuration POST /cost/estimate - Estimate cost for a given model and token counts """ -from typing import Dict, Union +from typing import Dict, Optional, Tuple, Union from fastapi import APIRouter, Depends, HTTPException @@ -29,6 +29,52 @@ from litellm.types.utils import LlmProvidersSet router = APIRouter() +def _resolve_model_for_cost_lookup(model: str) -> Tuple[str, Optional[str]]: + """ + Resolve a model name (which may be a router alias/model_group) to the + underlying litellm model name for cost lookup. + + Args: + model: The model name from the request (could be a router alias like 'e-model-router' + or an actual model name like 'azure_ai/gpt-4') + + Returns: + Tuple of (resolved_model_name, custom_llm_provider) + - resolved_model_name: The actual model name to use for cost lookup + - custom_llm_provider: The provider if resolved from router, None otherwise + """ + from litellm.proxy.proxy_server import llm_router + + custom_llm_provider: Optional[str] = None + + # Try to resolve from router if available + if llm_router is not None: + try: + # Get deployments for this model name (handles aliases, wildcards, etc.) + deployments = llm_router.get_model_list(model_name=model) + + if deployments and len(deployments) > 0: + # Get the first deployment's litellm model + first_deployment = deployments[0] + litellm_params = first_deployment.get("litellm_params", {}) + resolved_model = litellm_params.get("model") + + if resolved_model: + verbose_proxy_logger.debug( + f"Resolved model '{model}' to '{resolved_model}' from router" + ) + # Extract custom_llm_provider if present + custom_llm_provider = litellm_params.get("custom_llm_provider") + return resolved_model, custom_llm_provider + except Exception as e: + verbose_proxy_logger.debug( + f"Could not resolve model '{model}' from router: {e}" + ) + + # Return original model if not resolved + return model, custom_llm_provider + + def _calculate_period_costs( num_requests, cost_per_request, input_cost, output_cost, margin_cost ): @@ -413,12 +459,18 @@ async def estimate_cost( ``` """ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.types.utils import Usage - from litellm.utils import ModelResponse + from litellm.types.utils import ModelResponse, Usage + + # Resolve model name (handles router aliases like 'e-model-router' -> 'azure_ai/gpt-4') + resolved_model, resolved_provider = _resolve_model_for_cost_lookup(request.model) + + verbose_proxy_logger.debug( + f"Cost estimate: request.model='{request.model}' resolved to '{resolved_model}'" + ) # Create a mock response with usage for completion_cost mock_response = ModelResponse( - model=request.model, + model=resolved_model, usage=Usage( prompt_tokens=request.input_tokens, completion_tokens=request.output_tokens, @@ -428,7 +480,7 @@ async def estimate_cost( # Create a logging object to capture cost breakdown litellm_logging_obj = LiteLLMLoggingObj( - model=request.model, + model=resolved_model, messages=[], stream=False, call_type="completion", @@ -441,14 +493,14 @@ async def estimate_cost( try: cost_per_request = completion_cost( completion_response=mock_response, - model=request.model, + model=resolved_model, litellm_logging_obj=litellm_logging_obj, ) except Exception as e: raise HTTPException( status_code=404, detail={ - "error": f"Could not calculate cost for model '{request.model}': {str(e)}" + "error": f"Could not calculate cost for model '{request.model}' (resolved to '{resolved_model}'): {str(e)}" }, ) @@ -461,7 +513,7 @@ async def estimate_cost( # Get model info for per-token pricing display try: - model_info = litellm.get_model_info(model=request.model) + model_info = litellm.get_model_info(model=resolved_model) input_cost_per_token = model_info.get("input_cost_per_token") output_cost_per_token = model_info.get("output_cost_per_token") custom_llm_provider = model_info.get("litellm_provider") @@ -470,6 +522,10 @@ async def estimate_cost( output_cost_per_token = None custom_llm_provider = None + # Use provider from router resolution if not found in model_info + if custom_llm_provider is None and resolved_provider is not None: + custom_llm_provider = resolved_provider + # Calculate daily and monthly costs daily_cost, daily_input_cost, daily_output_cost, daily_margin_cost = ( _calculate_period_costs( diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index e48fd22bc8..b079e16151 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -17,7 +17,10 @@ from starlette.websockets import WebSocketState import litellm from litellm._logging import verbose_proxy_logger -from litellm.constants import BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES +from litellm.constants import ( + ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS, + BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES, +) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._types import * from litellm.proxy.auth.route_checks import RouteChecks @@ -1369,6 +1372,27 @@ def get_vertex_base_url(vertex_location: Optional[str]) -> str: return f"https://{vertex_location}-aiplatform.googleapis.com/" +def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict: + """ + Extract only the allowed headers from incoming request for Vertex AI pass-through. + + Uses an allowlist approach for security - only forwards headers we explicitly trust. + This prevents accidentally forwarding sensitive headers like the LiteLLM auth token. + + Args: + request: The FastAPI request object + + Returns: + dict: Headers dictionary with only allowed headers + """ + incoming_headers = dict(request.headers) or {} + headers = {} + for header_name in ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS: + if header_name in incoming_headers: + headers[header_name] = incoming_headers[header_name] + return headers + + def get_vertex_pass_through_handler( call_type: Literal["discovery", "aiplatform"], ) -> BaseVertexAIPassThroughHandler: @@ -1512,9 +1536,10 @@ async def _prepare_vertex_auth_headers( api_base="", ) - headers = { - "Authorization": f"Bearer {auth_header}", - } + # Use allowlist approach - only forward specific safe headers + headers = get_vertex_ai_allowed_incoming_headers(request) + # Add the Authorization header with vendor credentials + headers["Authorization"] = f"Bearer {auth_header}" if base_target_url is not None: base_target_url = get_vertex_pass_through_handler.update_base_target_url_with_credential_location( diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index 0c77b6f851..73e0ece3e2 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -36,13 +36,13 @@ router = APIRouter() def get_base_prompt_id(prompt_id: str) -> str: """ Extract the base prompt ID by stripping the version suffix if present. - + Args: prompt_id: Prompt ID that may include version suffix (e.g., "jack_success.v1" or "jack_success_v1") - + Returns: Base prompt ID without version suffix (e.g., "jack_success") - + Examples: >>> get_base_prompt_id("jack_success.v1") "jack_success" @@ -63,13 +63,13 @@ def get_base_prompt_id(prompt_id: str) -> str: def get_version_number(prompt_id: str) -> int: """ Extract the version number from a versioned prompt ID. - + Args: prompt_id: Prompt ID that may include version suffix (e.g., "jack_success.v2" or "jack_success_v2") - + Returns: Version number (defaults to 1 if no version suffix or invalid format) - + Examples: >>> get_version_number("jack_success.v2") 2 @@ -85,7 +85,7 @@ def get_version_number(prompt_id: str) -> int: return int(version_str) except ValueError: pass - + # Try underscore separator (_v) if "_v" in prompt_id: version_str = prompt_id.split("_v")[1] @@ -93,21 +93,21 @@ def get_version_number(prompt_id: str) -> int: return int(version_str) except ValueError: pass - + return 1 def construct_versioned_prompt_id(prompt_id: str, version: Optional[int] = None) -> str: """ Construct a versioned prompt ID from a base prompt_id and version number. - + Args: prompt_id: Base prompt ID (e.g., "jack_success") version: Version number (if None, returns the base prompt_id unchanged) - + Returns: Versioned prompt ID (e.g., "jack_success.v4") - + Examples: >>> construct_versioned_prompt_id("jack_success", 4) "jack_success.v4" @@ -118,7 +118,7 @@ def construct_versioned_prompt_id(prompt_id: str, version: Optional[int] = None) """ if version is None: return prompt_id - + # Strip any existing version suffix first base_id = get_base_prompt_id(prompt_id) return f"{base_id}.v{version}" @@ -127,14 +127,14 @@ def construct_versioned_prompt_id(prompt_id: str, version: Optional[int] = None) def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: Dict[str, Any]) -> str: """ Find the latest version of a prompt from available prompt IDs. - + Args: prompt_id: Base prompt ID or versioned prompt ID (e.g., "jack_success" or "jack_success.v2") all_prompt_ids: Dictionary of all available prompt IDs (keys are prompt IDs) - + Returns: The prompt ID with the highest version number, or the original prompt_id if no versions exist - + Examples: >>> all_ids = {"jack.v1": {}, "jack.v2": {}, "jack.v3": {}} >>> get_latest_version_prompt_id("jack", all_ids) @@ -146,14 +146,14 @@ def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: Dict[str, Any]) "simple" """ base_id = get_base_prompt_id(prompt_id=prompt_id) - + # Find all versions of this prompt matching_versions = [] for stored_prompt_id in all_prompt_ids.keys(): if get_base_prompt_id(prompt_id=stored_prompt_id) == base_id: version_num = get_version_number(prompt_id=stored_prompt_id) matching_versions.append((version_num, stored_prompt_id)) - + # Use the highest version number if matching_versions: matching_versions.sort(reverse=True) @@ -166,45 +166,47 @@ def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: Dict[str, Any]) def get_latest_prompt_versions(prompts: List[PromptSpec]) -> List[PromptSpec]: """ Filter a list of prompts to return only the latest version of each unique prompt. - + Args: prompts: List of PromptSpec objects - + Returns: List of PromptSpec objects with only the latest version of each prompt """ latest_prompts: Dict[str, PromptSpec] = {} - + for prompt in prompts: base_id = get_base_prompt_id(prompt_id=prompt.prompt_id) version = get_version_number(prompt_id=prompt.prompt_id) - + # Keep the prompt with the highest version number if base_id not in latest_prompts: latest_prompts[base_id] = prompt else: - existing_version = get_version_number(prompt_id=latest_prompts[base_id].prompt_id) + existing_version = get_version_number( + prompt_id=latest_prompts[base_id].prompt_id + ) if version > existing_version: latest_prompts[base_id] = prompt - + return list(latest_prompts.values()) async def get_next_version_for_prompt(prisma_client, prompt_id: str) -> int: """ Get the next version number for a prompt. - + Args: prisma_client: Prisma database client prompt_id: Base prompt ID - + Returns: Next version number (1 if no versions exist, max_version + 1 otherwise) """ existing_prompts = await prisma_client.db.litellm_prompttable.find_many( where={"prompt_id": prompt_id} ) - + if existing_prompts: max_version = max(p.version for p in existing_prompts) return max_version + 1 @@ -215,27 +217,27 @@ async def get_next_version_for_prompt(prisma_client, prompt_id: str) -> int: def create_versioned_prompt_spec(db_prompt) -> PromptSpec: """ Helper function to create a PromptSpec with versioned prompt_id from a DB prompt entry. - + Args: db_prompt: The DB prompt object (from prisma) - + Returns: PromptSpec with versioned prompt_id (e.g., "chat_prompt.v1") """ import json from litellm.types.prompts.init_prompts import PromptLiteLLMParams - + prompt_dict = db_prompt.model_dump() base_prompt_id = prompt_dict["prompt_id"] version = prompt_dict.get("version", 1) - + # Parse litellm_params litellm_params_data = prompt_dict.get("litellm_params") if isinstance(litellm_params_data, str): litellm_params_data = json.loads(litellm_params_data) litellm_params = PromptLiteLLMParams(**litellm_params_data) - + # Parse prompt_info prompt_info_data = prompt_dict.get("prompt_info") if prompt_info_data: @@ -244,10 +246,10 @@ def create_versioned_prompt_spec(db_prompt) -> PromptSpec: prompt_info = PromptInfo(**prompt_info_data) else: prompt_info = PromptInfo(prompt_type="db") - + # Create versioned prompt_id versioned_prompt_id = f"{base_prompt_id}.v{version}" - + return PromptSpec( prompt_id=versioned_prompt_id, litellm_params=litellm_params, @@ -319,10 +321,14 @@ async def list_prompts( prompt_list = [] for prompt_id in prompts: if prompt_id in IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS: - original_prompt = IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS[prompt_id] + original_prompt = IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS[ + prompt_id + ] # Create a copy with base prompt_id (without version suffix) prompt_copy = PromptSpec( - prompt_id=get_base_prompt_id(prompt_id=original_prompt.prompt_id), + prompt_id=get_base_prompt_id( + prompt_id=original_prompt.prompt_id + ), litellm_params=original_prompt.litellm_params, prompt_info=original_prompt.prompt_info, created_at=original_prompt.created_at, @@ -407,32 +413,33 @@ async def get_prompt_versions( raise HTTPException( status_code=403, detail="Only proxy admins can view prompt versions" ) - + # Strip version suffix if provided (e.g., "jack_success.v1" -> "jack_success") base_prompt_id = get_base_prompt_id(prompt_id=prompt_id) - + # Get all prompts and filter by base_prompt_id all_prompts = list(IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS.values()) prompt_versions = [ - prompt for prompt in all_prompts + prompt + for prompt in all_prompts if get_base_prompt_id(prompt_id=prompt.prompt_id) == base_prompt_id ] - + if not prompt_versions: raise HTTPException( status_code=404, detail=f"No versions found for prompt ID {base_prompt_id}" ) - + # Create response with explicit version field for each prompt versioned_prompts = [] for prompt in prompt_versions: # Extract version number from the root prompt_id which has version suffix # (e.g., "jack-sparrow.v3" -> 3) version_number = get_version_number(prompt_id=prompt.prompt_id) - + # Strip version from prompt_id for clean display base_prompt_id = get_base_prompt_id(prompt_id=prompt.prompt_id) - + # Create a copy with explicit version field and clean prompt_id versioned_prompt = PromptSpec( prompt_id=base_prompt_id, # Clean ID without version (e.g., "jack-sparrow") @@ -443,10 +450,10 @@ async def get_prompt_versions( version=version_number, # Explicit version field (e.g., 3) ) versioned_prompts.append(versioned_prompt) - + # Sort by version number (descending - newest first) versioned_prompts.sort(key=lambda p: p.version or 1, reverse=True) - + return ListPromptsResponse(prompts=versioned_prompts) @@ -518,21 +525,21 @@ async def get_prompt_info( # Try to get prompt directly first prompt_spec = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(prompt_id) - + # If not found, try to find the latest version if prompt_spec is None: latest_prompt_id = get_latest_version_prompt_id( prompt_id=prompt_id, - all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS + all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS, ) prompt_spec = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(latest_prompt_id) - + if prompt_spec is None: raise HTTPException(status_code=400, detail=f"Prompt {prompt_id} not found") # Extract version number from the prompt_id version_number = get_version_number(prompt_id=prompt_spec.prompt_id) - + # Create a copy of the prompt spec with the base prompt ID (stripped of version) # and explicit version field for consistency with list_prompts and versions endpoints prompt_spec_response = PromptSpec( @@ -547,7 +554,9 @@ async def get_prompt_info( # Get prompt content from the callback prompt_template: Optional[PromptTemplateBase] = None try: - prompt_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id(prompt_id) + prompt_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id( + prompt_spec.prompt_id + ) if prompt_callback is not None: # Extract content based on integration type integration_name = prompt_callback.integration_name @@ -723,12 +732,12 @@ async def update_prompt( try: # Strip version suffix from prompt_id if present (e.g., "jack_success.v1" -> "jack_success") base_prompt_id = get_base_prompt_id(prompt_id=prompt_id) - + # Check if any version exists existing_prompts = await prisma_client.db.litellm_prompttable.find_many( where={"prompt_id": base_prompt_id} ) - + if not existing_prompts: raise HTTPException( status_code=404, detail=f"Prompt with ID {base_prompt_id} not found" @@ -736,7 +745,10 @@ async def update_prompt( # Check if it's a config prompt existing_in_memory = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(prompt_id) - if existing_in_memory and existing_in_memory.prompt_info.prompt_type == "config": + if ( + existing_in_memory + and existing_in_memory.prompt_info.prompt_type == "config" + ): raise HTTPException( status_code=400, detail="Cannot update config prompts.", @@ -828,17 +840,19 @@ async def delete_prompt( try: # Try to get prompt directly first existing_prompt = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(prompt_id) - + # If not found, try to find the latest version if existing_prompt is None: latest_prompt_id = get_latest_version_prompt_id( prompt_id=prompt_id, - all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS + all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS, + ) + existing_prompt = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id( + latest_prompt_id ) - existing_prompt = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(latest_prompt_id) # Use the resolved prompt_id for deletion prompt_id = latest_prompt_id - + if existing_prompt is None: raise HTTPException( status_code=404, detail=f"Prompt with ID {prompt_id} not found" @@ -850,17 +864,18 @@ async def delete_prompt( detail="Cannot delete config prompts.", ) - # Delete the prompt from the database + # Get the base prompt ID (without version suffix) for database deletion + base_prompt_id = get_base_prompt_id(prompt_id=prompt_id) + + # Delete all versions of the prompt from the database await prisma_client.db.litellm_prompttable.delete_many( - where={"prompt_id": prompt_id} + where={"prompt_id": base_prompt_id} ) - # Remove the prompt from memory - del IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS[prompt_id] - if prompt_id in IN_MEMORY_PROMPT_REGISTRY.prompt_id_to_custom_prompt: - del IN_MEMORY_PROMPT_REGISTRY.prompt_id_to_custom_prompt[prompt_id] + # Remove all versions of the prompt from memory + IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id(base_prompt_id) - return {"message": f"Prompt {prompt_id} deleted successfully"} + return {"message": f"Prompt {base_prompt_id} deleted successfully"} except HTTPException as e: raise e @@ -1036,68 +1051,66 @@ async def test_prompt( user_temperature, version, ) - + try: # Parse the dotprompt content and create PromptTemplate prompt_manager = PromptManager() frontmatter, template_content = prompt_manager._parse_frontmatter( content=request.dotprompt_content ) - + # Create PromptTemplate to leverage existing parameter extraction logic template = PromptTemplate( - content=template_content, - metadata=frontmatter, - template_id="test_prompt" + content=template_content, metadata=frontmatter, template_id="test_prompt" ) - + # Extract model from template if not template.model: raise HTTPException( - status_code=400, - detail="Model is required in dotprompt metadata" + status_code=400, detail="Model is required in dotprompt metadata" ) - + # Always render the template to extract system messages and other metadata variables = request.prompt_variables or {} rendered_content = prompt_manager.jinja_env.from_string( template_content ).render(**variables) - + # Convert rendered content to messages using DotpromptManager's method dotprompt_manager = DotpromptManager() rendered_messages = dotprompt_manager._convert_to_messages( rendered_content=rendered_content ) - + if not rendered_messages: raise HTTPException( - status_code=400, - detail="No messages found in rendered prompt" + status_code=400, detail="No messages found in rendered prompt" ) - + # If conversation history is provided, use it but preserve system messages if request.conversation_history: # Extract system messages from rendered prompt - system_messages = [msg for msg in rendered_messages if msg.get("role") == "system"] + system_messages = [ + msg for msg in rendered_messages if msg.get("role") == "system" + ] # Use conversation history for user/assistant messages messages = system_messages + request.conversation_history else: messages = rendered_messages # type: ignore[assignment] - + # Use PromptTemplate's optional_params which already extracts all parameters optional_params = template.optional_params.copy() - + # Always stream the response optional_params["stream"] = True - + # Build request data for chat completion data = { "model": template.model, "messages": messages, } data.update(optional_params) - + # Use ProxyBaseLLMRequestProcessing to go through all proxy logic base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) result = await base_llm_response_processor.base_process_llm_request( @@ -1118,12 +1131,12 @@ async def test_prompt( user_api_base=user_api_base, version=version, ) - + if isinstance(result, BaseModel): return result.model_dump(exclude_none=True, exclude_unset=True) else: return result - + except HTTPException as e: raise e except Exception as e: @@ -1192,4 +1205,3 @@ async def convert_prompt_file_to_json( temp_file_path.parent.rmdir() except OSError: pass # Directory not empty or other error - diff --git a/litellm/proxy/prompts/prompt_registry.py b/litellm/proxy/prompts/prompt_registry.py index b471768770..58df60a42c 100644 --- a/litellm/proxy/prompts/prompt_registry.py +++ b/litellm/proxy/prompts/prompt_registry.py @@ -97,9 +97,9 @@ class InMemoryPromptRegistry: Prompt id to Prompt object mapping """ - self.prompt_id_to_custom_prompt: Dict[str, Optional[CustomPromptManagement]] = ( - {} - ) + self.prompt_id_to_custom_prompt: Dict[ + str, Optional[CustomPromptManagement] + ] = {} """ Guardrail id to CustomGuardrail object mapping """ @@ -174,5 +174,30 @@ class InMemoryPromptRegistry: """ return self.prompt_id_to_custom_prompt.get(prompt_id) + def delete_prompts_by_base_id(self, base_prompt_id: str) -> list[str]: + """ + Delete all prompts matching the given base prompt ID from memory. -IN_MEMORY_PROMPT_REGISTRY = InMemoryPromptRegistry() \ No newline at end of file + Args: + base_prompt_id: The base prompt ID (without version suffix) + + Returns: + List of prompt IDs that were deleted + """ + from litellm.proxy.prompts.prompt_endpoints import get_base_prompt_id + + prompts_to_delete = [ + pid + for pid in self.IN_MEMORY_PROMPTS.keys() + if get_base_prompt_id(prompt_id=pid) == base_prompt_id + ] + + for pid in prompts_to_delete: + del self.IN_MEMORY_PROMPTS[pid] + if pid in self.prompt_id_to_custom_prompt: + del self.prompt_id_to_custom_prompt[pid] + + return prompts_to_delete + + +IN_MEMORY_PROMPT_REGISTRY = InMemoryPromptRegistry() diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 646a062b72..958ddbf613 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,51 +1,27 @@ model_list: - - model_name: gemini/* + # Anthropic direct + - model_name: anthropic-claude litellm_params: - model: gemini/* - - model_name: -claude-sonnet-4-5-20250929 - litellm_params: - model: bedrock/invoke/us.anthropic.claude-sonnet-4-5-20250929-v1:0 - model_info: - cache_creation_input_token_cost: 3.75e-06 - cache_read_input_token_cost: 3e-07 - input_cost_per_token: 3e-06 - input_cost_per_token_above_200k_tokens: 6e-06 - output_cost_per_token_above_200k_tokens: 2.25e-05 - cache_creation_input_token_cost_above_200k_tokens: 7.5e-06 - cache_read_input_token_cost_above_200k_tokens: 6e-07 - litellm_provider: bedrock_converse - max_input_tokens: 200000 - max_output_tokens: 64000 - max_tokens: 200000 - mode: chat - output_cost_per_token: 1.5e-05 - search_context_cost_per_query: - search_context_size_high: 0.01 - search_context_size_low: 0.01 - search_context_size_medium: 0.01 - supports_assistant_prefill: true - supports_computer_use: true - supports_function_calling: true - supports_pdf_input: true - supports_prompt_caching: true - supports_reasoning: true - supports_response_schema: true - supports_tool_choice: true - supports_vision: true - tool_use_system_prompt_tokens: 346 + model: anthropic/claude-sonnet-4-20250514 + api_key: os.environ/ANTHROPIC_API_KEY - - model_name: us.anthropic.claude-sonnet-4-20250514-v1:0 + # Azure AI Anthropic + - model_name: azure-ai-claude litellm_params: - model: bedrock/converse/us.anthropic.claude-sonnet-4-20250514-v1:0 - model_info: - litellm_provider: bedrock_converse - mode: chat - - model_name: claude-sonnet-4-5-20250929 - litellm_params: - model: azure_ai/claude-opus-4-5 - api_base: https://krish-mh44t553-eastus2.services.ai.azure.com + model: azure_ai/claude-3-5-sonnet + api_base: https://krish-mh44t553-eastus2.services.ai.azure.com/ api_key: os.environ/AZURE_ANTHROPIC_API_KEY + # Azure AI Anthropic (alternate endpoint format) + - model_name: claude-4.5-haiku + litellm_params: + model: anthropic/claude-haiku-4-5 + api_base: https://krish-mh44t553-eastus2.services.ai.azure.com/anthropic/v1/messages + api_version: "2023-06-01" + api_key: os.environ/AZURE_ANTHROPIC_API_KEY + + + # Search Tools Configuration - Define search providers for WebSearch interception # search_tools: # - search_tool_name: "my-perplexity-search" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ff48443efc..ed7c5f8c2f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -548,9 +548,9 @@ except ImportError: server_root_path = os.getenv("SERVER_ROOT_PATH", "") _license_check = LicenseCheck() premium_user: bool = _license_check.is_premium() -premium_user_data: Optional[ - "EnterpriseLicenseData" -] = _license_check.airgapped_license_data +premium_user_data: Optional["EnterpriseLicenseData"] = ( + _license_check.airgapped_license_data +) global_max_parallel_request_retries_env: Optional[str] = os.getenv( "LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES" ) @@ -1208,9 +1208,9 @@ master_key: Optional[str] = None config_agents: Optional[List[AgentConfig]] = None otel_logging = False prisma_client: Optional[PrismaClient] = None -shared_aiohttp_session: Optional[ - "ClientSession" -] = None # Global shared session for connection reuse +shared_aiohttp_session: Optional["ClientSession"] = ( + None # Global shared session for connection reuse +) user_api_key_cache = DualCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) @@ -1218,9 +1218,9 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter( dual_cache=user_api_key_cache ) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) -redis_usage_cache: Optional[ - RedisCache -] = None # redis cache used for tracking spend, tpm/rpm limits +redis_usage_cache: Optional[RedisCache] = ( + None # redis cache used for tracking spend, tpm/rpm limits +) polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None @@ -1559,9 +1559,9 @@ async def update_cache( # noqa: PLR0915 _id = "team_id:{}".format(team_id) try: # Fetch the existing cost for the given user - existing_spend_obj: Optional[ - LiteLLM_TeamTable - ] = await user_api_key_cache.async_get_cache(key=_id) + existing_spend_obj: Optional[LiteLLM_TeamTable] = ( + await user_api_key_cache.async_get_cache(key=_id) + ) if existing_spend_obj is None: # do nothing if team not in api key cache return @@ -3108,17 +3108,19 @@ class ProxyConfig: async def _update_llm_router( self, - new_models: list, + new_models: Optional[Json], proxy_logging_obj: ProxyLogging, ): global llm_router, llm_model_list, master_key, general_settings - + config_data = await proxy_config.get_config() + search_tools = self.parse_search_tools(config_data) try: + models_list: list = new_models if isinstance(new_models, list) else [] if llm_router is None and master_key is not None: - verbose_proxy_logger.debug(f"len new_models: {len(new_models)}") + verbose_proxy_logger.debug(f"len new_models: {len(models_list)}") _model_list: list = self.decrypt_model_list_from_db( - new_models=new_models + new_models=models_list ) if len(_model_list) > 0: verbose_proxy_logger.debug(f"_model_list: {_model_list}") @@ -3127,16 +3129,17 @@ class ProxyConfig: router_general_settings=RouterGeneralSettings( async_only_mode=True # only init async clients ), + search_tools=search_tools, ignore_invalid_deployments=True, ) verbose_proxy_logger.debug(f"updated llm_router: {llm_router}") else: - verbose_proxy_logger.debug(f"len new_models: {len(new_models)}") + verbose_proxy_logger.debug(f"len new_models: {len(models_list)}") ## DELETE MODEL LOGIC - await self._delete_deployment(db_models=new_models) + await self._delete_deployment(db_models=models_list) ## ADD MODEL LOGIC - self._add_deployment(db_models=new_models) + self._add_deployment(db_models=models_list) except Exception as e: verbose_proxy_logger.exception( @@ -3147,7 +3150,6 @@ class ProxyConfig: llm_model_list = llm_router.get_model_list() # check if user set any callbacks in Config Table - config_data = await proxy_config.get_config() self._add_callbacks_from_db_config(config_data) # router settings @@ -3957,10 +3959,10 @@ class ProxyConfig: ) try: - guardrails_in_db: List[ - Guardrail - ] = await GuardrailRegistry.get_all_guardrails_from_db( - prisma_client=prisma_client + guardrails_in_db: List[Guardrail] = ( + await GuardrailRegistry.get_all_guardrails_from_db( + prisma_client=prisma_client + ) ) verbose_proxy_logger.debug( "guardrails from the DB %s", str(guardrails_in_db) @@ -4287,9 +4289,9 @@ async def initialize( # noqa: PLR0915 user_api_base = api_base dynamic_config[user_model]["api_base"] = api_base if api_version: - os.environ[ - "AZURE_API_VERSION" - ] = api_version # set this for azure - litellm can read this from the env + os.environ["AZURE_API_VERSION"] = ( + api_version # set this for azure - litellm can read this from the env + ) if max_tokens: # model-specific param dynamic_config[user_model]["max_tokens"] = max_tokens if temperature: # model-specific param @@ -5081,6 +5083,7 @@ async def model_list( only_model_access_groups: Optional[bool] = False, include_metadata: Optional[bool] = False, fallback_type: Optional[str] = None, + scope: Optional[str] = None, ): """ Use `/model/info` - to get detailed model information, example - pricing, mode, etc. @@ -5091,14 +5094,85 @@ async def model_list( - include_metadata: Include additional metadata in the response with fallback information - fallback_type: Type of fallbacks to include ("general", "context_window", "content_policy") Defaults to "general" when include_metadata=true + - scope: Optional scope parameter. Currently only accepts "expand". + When scope=expand is passed, proxy admins, team admins, and org admins + will receive all proxy models as if they are a proxy admin. """ global llm_model_list, general_settings, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + from litellm.proxy.management_endpoints.common_utils import ( + _user_has_admin_privileges, + ) from litellm.proxy.utils import ( create_model_info_response, get_available_models_for_user, ) + # Validate scope parameter if provided + if scope is not None and scope != "expand": + raise HTTPException( + status_code=400, + detail=f"Invalid scope parameter. Only 'expand' is currently supported. Received: {scope}", + ) + + # Check if scope=expand is requested and user has admin privileges + should_expand_scope = False + if scope == "expand": + should_expand_scope = await _user_has_admin_privileges( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + # If scope=expand and user has admin privileges, return all proxy models + if should_expand_scope: + # Get all proxy models as if user is a proxy admin + if llm_router is None: + proxy_model_list = [] + model_access_groups = {} + else: + proxy_model_list = llm_router.get_model_names() + model_access_groups = llm_router.get_model_access_groups() + + # Include model access groups if requested + if include_model_access_groups: + proxy_model_list = list(set(proxy_model_list + list(model_access_groups.keys()))) + + # Get complete model list including wildcard routes if requested + from litellm.proxy.auth.model_checks import get_complete_model_list + + all_models = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=proxy_model_list, + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=return_wildcard_routes or False, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups or False, + only_model_access_groups=only_model_access_groups or False, + ) + + # Build response data with all proxy models + model_data = [] + for model in all_models: + model_info = create_model_info_response( + model_id=model, + provider="openai", + include_metadata=include_metadata or False, + fallback_type=fallback_type, + llm_router=llm_router, + ) + model_data.append(model_info) + + return dict( + data=model_data, + object="list", + ) + + # Otherwise, use the normal behavior (current implementation) # Get available models for the user all_models = await get_available_models_for_user( user_api_key_dict=user_api_key_dict, @@ -7514,6 +7588,77 @@ async def get_all_team_and_direct_access_models( return all_models +def _enrich_model_info_with_litellm_data( + model: Dict[str, Any], debug: bool = False, llm_router: Optional[Router] = None +) -> Dict[str, Any]: + """ + Enrich a model dictionary with litellm model info (pricing, context window, etc.) + and remove sensitive information. + + Args: + model: Model dictionary to enrich + debug: Whether to include debug information like openai_client + llm_router: Optional router instance for debug info + + Returns: + Enriched model dictionary with sensitive info removed + """ + # provided model_info in config.yaml + model_info = model.get("model_info", {}) + if debug is True: + _openai_client = "None" + if llm_router is not None: + _openai_client = ( + llm_router._get_client( + deployment=model, kwargs={}, client_type="async" + ) + or "None" + ) + else: + _openai_client = "llm_router_is_None" + openai_client = str(_openai_client) + model["openai_client"] = openai_client + + # read litellm model_prices_and_context_window.json to get the following: + # input_cost_per_token, output_cost_per_token, max_tokens + litellm_model_info = get_litellm_model_info(model=model) + + # 2nd pass on the model, try seeing if we can find model in litellm model_cost map + if litellm_model_info == {}: + # use litellm_param model_name to get model_info + litellm_params = model.get("litellm_params", {}) + litellm_model = litellm_params.get("model", None) + try: + litellm_model_info = litellm.get_model_info(model=litellm_model) + except Exception: + litellm_model_info = {} + # 3rd pass on the model, try seeing if we can find model but without the "/" in model cost map + if litellm_model_info == {}: + # use litellm_param model_name to get model_info + litellm_params = model.get("litellm_params", {}) + litellm_model = litellm_params.get("model", None) + if litellm_model: + split_model = litellm_model.split("/") + if len(split_model) > 0: + litellm_model = split_model[-1] + try: + litellm_model_info = litellm.get_model_info( + model=litellm_model, custom_llm_provider=split_model[0] + ) + except Exception: + litellm_model_info = {} + for k, v in litellm_model_info.items(): + if k not in model_info: + model_info[k] = v + model["model_info"] = model_info + # don't return the api key / vertex credentials + # don't return the llm credentials + model = remove_sensitive_info_from_deployment( + model, excluded_keys={"litellm_credential_name"} + ) + return model + + @router.get( "/v2/model/info", description="v2 - returns models available to the user based on their API key permissions. Shows model info from config.yaml (except api key and api base). Filter to just user-added models with ?user_models_only=true", @@ -7533,6 +7678,8 @@ async def model_info_v2( False, description="Return all models across all teams user is in." ), debug: Optional[bool] = False, + page: int = Query(1, description="Page number", ge=1), + size: int = Query(50, description="Page size", ge=1), ): """ BETA ENDPOINT. Might change unexpectedly. Use `/v1/model/info` for now. @@ -7541,7 +7688,13 @@ async def model_info_v2( # Return empty data array when no models are configured (graceful handling for fresh installs) if llm_router is None or not llm_router.model_list: - return {"data": []} + return { + "data": [], + "total_count": 0, + "current_page": page, + "total_pages": 0, + "size": size, + } if prisma_client is None: raise HTTPException( @@ -7576,62 +7729,32 @@ async def model_info_v2( all_models=all_models, ) # fill in model info based on config.yaml and litellm model_prices_and_context_window.json - for _model in all_models: - # provided model_info in config.yaml - model_info = _model.get("model_info", {}) - if debug is True: - _openai_client = "None" - if llm_router is not None: - _openai_client = ( - llm_router._get_client( - deployment=_model, kwargs={}, client_type="async" - ) - or "None" - ) - else: - _openai_client = "llm_router_is_None" - openai_client = str(_openai_client) - _model["openai_client"] = openai_client - - # read litellm model_prices_and_context_window.json to get the following: - # input_cost_per_token, output_cost_per_token, max_tokens - litellm_model_info = get_litellm_model_info(model=_model) - - # 2nd pass on the model, try seeing if we can find model in litellm model_cost map - if litellm_model_info == {}: - # use litellm_param model_name to get model_info - litellm_params = _model.get("litellm_params", {}) - litellm_model = litellm_params.get("model", None) - try: - litellm_model_info = litellm.get_model_info(model=litellm_model) - except Exception: - litellm_model_info = {} - # 3rd pass on the model, try seeing if we can find model but without the "/" in model cost map - if litellm_model_info == {}: - # use litellm_param model_name to get model_info - litellm_params = _model.get("litellm_params", {}) - litellm_model = litellm_params.get("model", None) - split_model = litellm_model.split("/") - if len(split_model) > 0: - litellm_model = split_model[-1] - try: - litellm_model_info = litellm.get_model_info( - model=litellm_model, custom_llm_provider=split_model[0] - ) - except Exception: - litellm_model_info = {} - for k, v in litellm_model_info.items(): - if k not in model_info: - model_info[k] = v - _model["model_info"] = model_info - # don't return the api key / vertex credentials - # don't return the llm credentials - _model = remove_sensitive_info_from_deployment( - _model, excluded_keys={"litellm_credential_name"} + for i, _model in enumerate(all_models): + all_models[i] = _enrich_model_info_with_litellm_data( + model=_model, debug=debug if debug is not None else False, llm_router=llm_router ) verbose_proxy_logger.debug("all_models: %s", all_models) - return {"data": all_models} + + total_count = len(all_models) + + skip = (page - 1) * size + + total_pages = -(-total_count // size) if total_count > 0 else 0 + + paginated_models = all_models[skip : skip + size] + + verbose_proxy_logger.debug( + f"Pagination: skip={skip}, take={size}, total_count={total_count}, total_pages={total_pages}" + ) + + return { + "data": paginated_models, + "total_count": total_count, + "current_page": page, + "total_pages": total_pages, + "size": size, + } @router.get( @@ -9742,9 +9865,9 @@ async def get_config_list( hasattr(sub_field_info, "description") and sub_field_info.description is not None ): - nested_fields[ - idx - ].field_description = sub_field_info.description + nested_fields[idx].field_description = ( + sub_field_info.description + ) idx += 1 _stored_in_db = None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index fcb678ef02..e0855333e2 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -131,7 +131,6 @@ else: unified_guardrail = UnifiedLLMGuardrails() -_anthropic_async_clients = {} def print_verbose(print_statement): """ @@ -961,8 +960,8 @@ class ProxyLogging: Returns: Updated data dictionary if guardrail passes, None if guardrail should be skipped """ - from litellm.types.guardrails import GuardrailEventHooks from litellm.integrations.prometheus import PrometheusLogger + from litellm.types.guardrails import GuardrailEventHooks # Determine the event type based on call type event_type = GuardrailEventHooks.pre_call @@ -4292,74 +4291,6 @@ def construct_database_url_from_env_vars() -> Optional[str]: return None -async def count_tokens_with_anthropic_api( - model_to_use: str, - messages: Optional[List[Dict[str, Any]]], - deployment: Optional[Dict[str, Any]] = None, -) -> Optional[Dict[str, Any]]: - """ - Helper function to count tokens using Anthropic API directly. - - Args: - model_to_use: The model name to use for token counting - messages: The messages to count tokens for - deployment: Optional deployment configuration containing API key - - Returns: - Optional dict with token count and tokenizer info, or None if failed - """ - if not messages: - return None - - try: - import os - - import anthropic - - # Get Anthropic API key from deployment config - anthropic_api_key = None - if deployment is not None: - anthropic_api_key = deployment.get("litellm_params", {}).get("api_key") - - # Fallback to environment variable - if not anthropic_api_key: - anthropic_api_key = os.getenv("ANTHROPIC_API_KEY") - - if anthropic_api_key and messages: - # Call Anthropic API directly for more accurate token counting - - # Use cached client if available to avoid socket exhaustion - if anthropic_api_key not in _anthropic_async_clients: - _anthropic_async_clients[anthropic_api_key] = anthropic.AsyncAnthropic(api_key=anthropic_api_key) - - client = _anthropic_async_clients[anthropic_api_key] - - # Call with explicit parameters to satisfy type checking - # Type ignore for now since messages come from generic dict input - response = await client.beta.messages.count_tokens( - model=model_to_use, - messages=messages, # type: ignore - betas=["token-counting-2024-11-01"], - ) - total_tokens = response.input_tokens - tokenizer_used = "anthropic_api" - - return { - "total_tokens": total_tokens, - "tokenizer_used": tokenizer_used, - } - - except ImportError: - verbose_proxy_logger.warning( - "Anthropic library not available, falling back to LiteLLM tokenizer" - ) - except Exception as e: - verbose_proxy_logger.warning( - f"Error calling Anthropic API: {e}, falling back to LiteLLM tokenizer" - ) - return None - - async def get_available_models_for_user( user_api_key_dict: "UserAPIKeyAuth", llm_router: Optional["Router"], diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index b128452d9a..867c18b6dd 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -16,14 +16,16 @@ from litellm.types.llms.openai import ( ContentPartDoneEvent, ContentPartDonePartOutputText, ContentPartDonePartReasoningText, + FunctionCallArgumentsDeltaEvent, + FunctionCallArgumentsDoneEvent, OutputItemAddedEvent, OutputItemDoneEvent, OutputTextAnnotationAddedEvent, OutputTextDeltaEvent, OutputTextDoneEvent, - FunctionCallArgumentsDeltaEvent, - FunctionCallArgumentsDoneEvent, + ReasoningSummaryPartDoneEvent, ReasoningSummaryTextDeltaEvent, + ReasoningSummaryTextDoneEvent, ResponseCompletedEvent, ResponseCreatedEvent, ResponseInProgressEvent, @@ -88,6 +90,17 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._tool_args_by_call_id: dict[str, str] = {} self._next_tool_output_index: int = 1 # output_index=0 reserved for the message item self._final_tool_events_queued: bool = False + self._sequence_number: int = 0 + self._cached_reasoning_item_id: Optional[str] = None + self._sent_reasoning_summary_text_done_event: bool = False + self._sent_reasoning_summary_part_done_event: bool = False + self._reasoning_summary_text: str = "" + # -- GENERIC RESPONSE-EVENTS PENDING QUEUE as required by fix -- + self._pending_response_events: List[BaseLiteLLMOpenAIResponseObject] = [] + self._reasoning_active = False + self._reasoning_done_emitted = False + self._reasoning_item_id: Optional[str] = None + def _get_or_assign_tool_output_index(self, call_id: str) -> int: existing = self._tool_output_index_by_call_id.get(call_id) @@ -98,13 +111,33 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._tool_output_index_by_call_id[call_id] = idx return idx + + def _is_reasoning_end(self, chunk): + delta = chunk.choices[0].delta + + # if this indicates reasoning content, don't consider reasoning ended + if hasattr(delta, "reasoning_content") and delta.reasoning_content: + return False + if hasattr(delta, "thinking_blocks") and delta.thinking_blocks: + return False + + return ( + delta.content + or delta.function_call + or delta.tool_calls + or chunk.choices[0].finish_reason is not None + ) + def _queue_tool_call_delta_events(self, tool_calls: object) -> None: """ Convert chat-completions streaming `tool_calls` deltas into Responses API streaming events. We emit: - response.output_item.added (function_call) - - response.function_call_arguments.delta + - response.function_call_arguments.delta (split into smaller chunks to match OpenAI behavior) + + Note: Some providers (like Bedrock) send tool call arguments in one large chunk. + We split these into smaller deltas to match OpenAI's token-by-token streaming behavior. """ if not isinstance(tool_calls, list): return @@ -129,33 +162,42 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if call_id not in self._tool_args_by_call_id: self._tool_args_by_call_id[call_id] = "" - self._pending_tool_events.append( - OutputItemAddedEvent( - type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, - output_index=output_index, - item=BaseLiteLLMOpenAIResponseObject( - **{ - "type": "function_call", - "id": call_id, - "call_id": call_id, - "name": fn_name, - "arguments": "", - "status": "in_progress", - } - ), - ) + self._sequence_number += 1 + event = OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=output_index, + item=BaseLiteLLMOpenAIResponseObject( + **{ + "type": "function_call", + "id": call_id, + "call_id": call_id, + "name": fn_name, + "arguments": "", + "status": "in_progress", + } + ), ) + event.__dict__['sequence_number'] = self._sequence_number + self._pending_tool_events.append(event) if fn_args_delta: self._tool_args_by_call_id[call_id] += fn_args_delta - self._pending_tool_events.append( - FunctionCallArgumentsDeltaEvent( + + # Split large argument deltas into smaller chunks to match OpenAI's streaming behavior + # This is especially important for providers like Bedrock that send complete arguments at once + chunk_size = 10 # Match typical OpenAI delta size + for i in range(0, len(fn_args_delta), chunk_size): + delta_chunk = fn_args_delta[i:i + chunk_size] + self._sequence_number += 1 + delta_event: BaseLiteLLMOpenAIResponseObject = FunctionCallArgumentsDeltaEvent( type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, item_id=call_id, output_index=output_index, - delta=fn_args_delta, + delta=delta_chunk, ) - ) + # Add sequence_number as extra field (BaseLiteLLMOpenAIResponseObject allows extra fields) + delta_event.__dict__['sequence_number'] = self._sequence_number + self._pending_tool_events.append(delta_event) def _queue_final_tool_call_done_events(self, litellm_complete_object: ModelResponse) -> None: """ @@ -191,53 +233,79 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): fn_name = str(getattr(fn, "name", "") or "") fn_args = str(getattr(fn, "arguments", "") or "") + # Track if this is a new tool call that wasn't streamed + is_new_tool_call = call_id not in self._tool_args_by_call_id + # If we never sent output_item.added for this call_id, emit it now. - if call_id not in self._tool_args_by_call_id: + if is_new_tool_call: self._tool_args_by_call_id[call_id] = "" - self._pending_tool_events.append( - OutputItemAddedEvent( - type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, - output_index=output_index, - item=BaseLiteLLMOpenAIResponseObject( - **{ - "type": "function_call", - "id": call_id, - "call_id": call_id, - "name": fn_name, - "arguments": "", - "status": "in_progress", - } - ), - ) - ) - - final_args = fn_args or self._tool_args_by_call_id.get(call_id, "") - self._pending_tool_events.append( - FunctionCallArgumentsDoneEvent( - type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE, - item_id=call_id, + self._sequence_number += 1 + event = OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=output_index, - arguments=final_args, - ) - ) - - self._pending_tool_events.append( - OutputItemDoneEvent( - type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, - output_index=output_index, - sequence_number=1, item=BaseLiteLLMOpenAIResponseObject( **{ "type": "function_call", "id": call_id, "call_id": call_id, "name": fn_name, - "arguments": final_args, - "status": "completed", + "arguments": "", + "status": "in_progress", } ), ) + event.__dict__['sequence_number'] = self._sequence_number + self._pending_tool_events.append(event) + + final_args = fn_args or self._tool_args_by_call_id.get(call_id, "") + + # Emit delta events for arguments that weren't streamed yet + # This handles cases where Bedrock sends the complete tool call at the end + already_streamed = self._tool_args_by_call_id.get(call_id, "") + remaining_args = final_args[len(already_streamed):] if final_args else "" + + if remaining_args: + # Split into smaller chunks to match OpenAI's streaming behavior + chunk_size = 10 # Match typical OpenAI delta size + for i in range(0, len(remaining_args), chunk_size): + delta_chunk = remaining_args[i:i + chunk_size] + self._sequence_number += 1 + delta_event = FunctionCallArgumentsDeltaEvent( + type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, + item_id=call_id, + output_index=output_index, + delta=delta_chunk, + ) + delta_event.__dict__['sequence_number'] = self._sequence_number + self._pending_tool_events.append(delta_event) + + self._sequence_number += 1 + done_event = FunctionCallArgumentsDoneEvent( + type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE, + item_id=call_id, + output_index=output_index, + arguments=final_args, ) + done_event.__dict__['sequence_number'] = self._sequence_number + self._pending_tool_events.append(done_event) + + self._sequence_number += 1 + item_done_event = OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=output_index, + sequence_number=self._sequence_number, + item=BaseLiteLLMOpenAIResponseObject( + **{ + "type": "function_call", + "id": call_id, + "call_id": call_id, + "name": fn_name, + "arguments": final_args, + "status": "completed", + } + ), + ) + self._pending_tool_events.append(item_done_event) def _default_response_created_event_data(self) -> dict: response_created_event_data = { @@ -295,24 +363,31 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): """ response_created_event_data = self._default_response_created_event_data() - return ResponseCreatedEvent( + self._sequence_number += 1 + event = ResponseCreatedEvent( type=ResponsesAPIStreamEvents.RESPONSE_CREATED, response=ResponsesAPIResponse(**response_created_event_data), ) + event.__dict__['sequence_number'] = self._sequence_number + return event def create_response_in_progress_event(self) -> ResponseInProgressEvent: response_in_progress_event_data = self._default_response_created_event_data() response_in_progress_event_data["status"] = "in_progress" - return ResponseInProgressEvent( + self._sequence_number += 1 + event = ResponseInProgressEvent( type=ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, response=ResponsesAPIResponse(**response_in_progress_event_data), ) + event.__dict__['sequence_number'] = self._sequence_number + return event def create_output_item_added_event(self) -> OutputItemAddedEvent: if self._cached_item_id is None: self._cached_item_id = f"msg_{str(uuid.uuid4())}" - return OutputItemAddedEvent( + self._sequence_number += 1 + event = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=0, item=BaseLiteLLMOpenAIResponseObject( @@ -325,12 +400,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): } ), ) + event.__dict__['sequence_number'] = self._sequence_number + return event def create_content_part_added_event(self) -> ContentPartAddedEvent: if self._cached_item_id is None: self._cached_item_id = f"msg_{str(uuid.uuid4())}" - return ContentPartAddedEvent( + self._sequence_number += 1 + event = ContentPartAddedEvent( type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED, item_id=self._cached_item_id, output_index=0, @@ -339,6 +417,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): **{"type": "output_text", "text": "", "annotations": []} ), ) + event.__dict__['sequence_number'] = self._sequence_number + return event def create_litellm_model_response( self, @@ -351,6 +431,70 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ), ) + def create_reasoning_summary_text_done_event( + self, + reasoning_item_id: str, + reasoning_content: str, + sequence_number: int, + ) -> ReasoningSummaryTextDoneEvent: + """ + Create response.reasoning_summary_text.done event. + + Example: + { + "type": "response.reasoning_summary_text.done", + "item_id": "rs_0c5dae30e53172980069708ba2f59c8197b71ca9820edad07c", + "output_index": 0, + "sequence_number": 97, + "summary_index": 0, + "text": "**Clarifying the first humans**\n\nThe I'm addressing the user's specific interest." + } + """ + return ReasoningSummaryTextDoneEvent( + type=ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE, + item_id=reasoning_item_id, + output_index=0, + sequence_number=sequence_number, + summary_index=0, + text=reasoning_content, + ) + + def create_reasoning_summary_part_done_event( + self, + reasoning_item_id: str, + reasoning_content: str, + sequence_number: int, + ) -> ReasoningSummaryPartDoneEvent: + """ + Create response.reasoning_summary_part.done event. + + Example: + { + "type": "response.reasoning_summary_part.done", + "item_id": "rs_0c5dae30e53172980069708ba2f59c8197b71ca9820edad07c", + "output_index": 0, + "part": { + "type": "summary_text", + "text": "**Clarifying the first humans**\n\nThe earlier hominins. It feels important to ensure I'm addressing the user's specific interest." + }, + "sequence_number": 98, + "summary_index": 0 + } + """ + return ReasoningSummaryPartDoneEvent( + type=ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE, + item_id=reasoning_item_id, + output_index=0, + sequence_number=sequence_number, + summary_index=0, + part=BaseLiteLLMOpenAIResponseObject( + **{ + "type": "summary_text", + "text": reasoning_content, + } + ), + ) + def create_output_text_done_event( self, litellm_complete_object: ModelResponse ) -> OutputTextDoneEvent: @@ -435,6 +579,50 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ), ) + def create_reasoning_output_item_done_event( + self, + reasoning_item_id: str, + reasoning_content: str, + sequence_number: int, + ) -> OutputItemDoneEvent: + """ + Create response.output_item.done event for reasoning items. + + Example: + { + "type": "response.output_item.done", + "output_index": 0, + "sequence_number": 99, + "item": { + "id": "rs_0c5dae30e53172980069708ba2f59c8197b71ca9820edad07c", + "type": "reasoning", + "summary": [ + { + "type": "summary_text", + "text": "**Clarifying the first humans**..." + } + ] + } + } + """ + return OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=0, + sequence_number=sequence_number, + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": reasoning_item_id, + "type": "reasoning", + "summary": [ + { + "type": "summary_text", + "text": reasoning_content, + } + ], + } + ), + ) + def return_default_done_events( self, litellm_complete_object: ModelResponse ) -> Optional[BaseLiteLLMOpenAIResponseObject]: @@ -458,12 +646,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): elif self.sent_response_in_progress_event is False: self.sent_response_in_progress_event = True return self.create_response_in_progress_event() - elif self.sent_output_item_added_event is False: - self.sent_output_item_added_event = True - return self.create_output_item_added_event() - elif self.sent_content_part_added_event is False: - self.sent_content_part_added_event = True - return self.create_content_part_added_event() return None def is_stream_finished(self) -> bool: @@ -510,6 +692,63 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): else: raise StopAsyncIteration + def _ensure_output_item_for_chunk(self, chunk: ModelResponseStream) -> None: + # Change: Never return a value, just enqueue output item events + if self.sent_output_item_added_event: + return + delta = chunk.choices[0].delta + + self._sequence_number += 1 + self.sent_output_item_added_event = True + + # Reasoning-first + if hasattr(delta, "reasoning_content") and delta.reasoning_content: + self._reasoning_active = True + if self._cached_reasoning_item_id is None: + self._cached_reasoning_item_id = f"rs_{uuid.uuid4()}" + self._reasoning_item_id = self._cached_reasoning_item_id + + event = OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=0, + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": self._cached_reasoning_item_id, + "type": "reasoning", + "status": "in_progress", + "summary": None, + } + ), + ) + event.__dict__['sequence_number'] = self._sequence_number + self._pending_response_events.append(event) + return + + # Tool-first + if hasattr(delta, "tool_calls") and delta.tool_calls: + # Tool calls already handled via _queue_tool_call_delta_events + # DO NOT create message item + return + + # Default: message + self._cached_item_id = self._cached_item_id or f"msg_{uuid.uuid4()}" + event = OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=0, + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": self._cached_item_id, + "type": "message", + "role": "assistant", + "status": "in_progress", + "content": [], + } + ), + ) + event.__dict__['sequence_number'] = self._sequence_number + self._pending_response_events.append(event) + return + async def __anext__( self, ) -> Union[ @@ -525,19 +764,77 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): result = self.return_default_initial_events() if result: return result - # Get the next chunk from the stream + # Emit any pending output_item or other response events before reading a new chunk + if self._pending_response_events: + return self._pending_response_events.pop(0) + # Emit any pending tool events before reading a new chunk + if self._pending_tool_events: + return self._pending_tool_events.pop(0) + try: chunk = await self.litellm_custom_stream_wrapper.__anext__() if chunk is not None: chunk = cast(ModelResponseStream, chunk) + self._ensure_output_item_for_chunk(chunk) + # Proceed to transformation self.collected_chat_completion_chunks.append(chunk) + if self._reasoning_active and not self._reasoning_done_emitted: + # get raw ModelResponse + text_reasoning = self.create_litellm_model_response() + # reasoning_content only + if self._is_reasoning_end(chunk): + reasoning_content = "" + # best effort to obtain reasoning_content from chat model response + if text_reasoning and text_reasoning.choices: + choice = text_reasoning.choices[0] + # Check if it's a Choices object (has message) or StreamingChoices (has delta) + if hasattr(choice, "message"): + reasoning_content = getattr(choice.message, "reasoning_content", "") or "" + + # Ensure we have a valid reasoning_item_id + reasoning_item_id = self._reasoning_item_id or self._cached_reasoning_item_id or f"rs_{uuid.uuid4()}" + + # Create text.done event first with its own sequence number + self._sequence_number += 1 + text_done_event = self.create_reasoning_summary_text_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number + ) + + # Create part.done event second with its own sequence number + self._sequence_number += 1 + part_done_event = self.create_reasoning_summary_part_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number + ) + + self._sequence_number += 1 + reasoning_output_item_done_event = self.create_reasoning_output_item_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number + ) + self._pending_response_events.extend([ + text_done_event, + part_done_event, + reasoning_output_item_done_event, + ]) + self._reasoning_done_emitted = True + self._reasoning_active = False + response_api_chunk = ( self._transform_chat_completion_chunk_to_response_api_chunk( chunk ) ) if response_api_chunk: - return response_api_chunk + self._pending_response_events.append(response_api_chunk) + + if self._pending_response_events: + return self._pending_response_events.pop(0) + except StopAsyncIteration: return self.common_done_event_logic(sync_mode=False) @@ -560,13 +857,21 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): while True: if self.finished is True: raise StopIteration - # Get the next chunk from the stream - result = self.return_default_initial_events() if result: return result + # Emit any pending output_item or other response events before reading a new chunk + if self._pending_response_events: + return self._pending_response_events.pop(0) + # Emit any pending tool events before reading a new chunk + if self._pending_tool_events: + return self._pending_tool_events.pop(0) try: chunk = self.litellm_custom_stream_wrapper.__next__() + self._ensure_output_item_for_chunk(chunk) + # Emit any just-queued output_item event + if self._pending_response_events: + return self._pending_response_events.pop(0) self.collected_chat_completion_chunks.append(chunk) response_api_chunk = ( self._transform_chat_completion_chunk_to_response_api_chunk( @@ -575,6 +880,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) if response_api_chunk: return response_api_chunk + # Otherwise, loop to next chunk except StopIteration: return self.common_done_event_logic(sync_mode=True) except Exception as e: @@ -638,21 +944,26 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Priority 2: Handle text deltas delta_content = self._get_delta_string_from_streaming_choices(chunk.choices) if delta_content: - return OutputTextDeltaEvent( + self._sequence_number += 1 + text_delta_event = OutputTextDeltaEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id=item_id, output_index=0, content_index=0, delta=delta_content, ) + text_delta_event.__dict__['sequence_number'] = self._sequence_number + return text_delta_event # Priority 3: Handle tool call deltas (if any) -> queue events and emit them + # For each tool call delta, we emit events one at a time to match OpenAI's streaming behavior if ( chunk.choices and hasattr(chunk.choices[0].delta, "tool_calls") and chunk.choices[0].delta.tool_calls ): self._queue_tool_call_delta_events(chunk.choices[0].delta.tool_calls) + # Return one pending tool event at a time if self._pending_tool_events: return self._pending_tool_events.pop(0) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 71d94287e8..83c23a5850 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -169,7 +169,9 @@ async def aresponses_api_with_mcp( # Process MCP tools through the complete pipeline (fetch + filter + deduplicate + transform) # Extract user_api_key_auth from litellm_metadata (where it's added by add_user_api_key_auth_to_request_metadata) - user_api_key_auth = kwargs.get("user_api_key_auth") or kwargs.get("litellm_metadata", {}).get("user_api_key_auth") + user_api_key_auth = kwargs.get("user_api_key_auth") or kwargs.get( + "litellm_metadata", {} + ).get("user_api_key_auth") # Get original MCP tools (for events) and OpenAI tools (for LLM) by reusing existing methods ( @@ -280,7 +282,7 @@ async def aresponses_api_with_mcp( user_api_key_auth = kwargs.get("litellm_metadata", {}).get( "user_api_key_auth" ) - + # Extract MCP auth headers from the request to pass to MCP server secret_fields: Optional[Dict[str, Any]] = kwargs.get("secret_fields") ( @@ -292,7 +294,7 @@ async def aresponses_api_with_mcp( secret_fields=secret_fields, tools=tools, ) - + tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( tool_server_map=tool_server_map, tool_calls=tool_calls, @@ -301,6 +303,8 @@ async def aresponses_api_with_mcp( mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers_from_request, + litellm_call_id=kwargs.get("litellm_call_id"), + litellm_trace_id=kwargs.get("litellm_trace_id"), ) if tool_results: @@ -349,6 +353,7 @@ async def aresponses_api_with_mcp( tool_server_map=tool_server_map, base_iterator=final_response, mcp_events=tool_execution_events, + user_api_key_auth=user_api_key_auth, ) # Add custom output elements to the final response (for non-streaming) @@ -587,9 +592,12 @@ def responses( ######################################################### # Update input with provider-specific file IDs if managed files are used ######################################################### - input = cast(Union[str, ResponseInputParam], update_responses_input_with_model_file_ids(input=input)) + input = cast( + Union[str, ResponseInputParam], + update_responses_input_with_model_file_ids(input=input), + ) local_vars["input"] = input - + ######################################################### # Native MCP Responses API ######################################################### @@ -624,11 +632,11 @@ def responses( ) # get provider config - responses_api_provider_config: Optional[BaseResponsesAPIConfig] = ( - ProviderConfigManager.get_provider_responses_api_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), - ) + responses_api_provider_config: Optional[ + BaseResponsesAPIConfig + ] = ProviderConfigManager.get_provider_responses_api_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), ) local_vars.update(kwargs) @@ -823,11 +831,11 @@ def delete_responses( raise ValueError("custom_llm_provider is required but passed as None") # get provider config - responses_api_provider_config: Optional[BaseResponsesAPIConfig] = ( - ProviderConfigManager.get_provider_responses_api_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + responses_api_provider_config: Optional[ + BaseResponsesAPIConfig + ] = ProviderConfigManager.get_provider_responses_api_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if responses_api_provider_config is None: @@ -1003,11 +1011,11 @@ def get_responses( raise ValueError("custom_llm_provider is required but passed as None") # get provider config - responses_api_provider_config: Optional[BaseResponsesAPIConfig] = ( - ProviderConfigManager.get_provider_responses_api_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + responses_api_provider_config: Optional[ + BaseResponsesAPIConfig + ] = ProviderConfigManager.get_provider_responses_api_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if responses_api_provider_config is None: @@ -1160,11 +1168,11 @@ def list_input_items( if custom_llm_provider is None: raise ValueError("custom_llm_provider is required but passed as None") - responses_api_provider_config: Optional[BaseResponsesAPIConfig] = ( - ProviderConfigManager.get_provider_responses_api_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + responses_api_provider_config: Optional[ + BaseResponsesAPIConfig + ] = ProviderConfigManager.get_provider_responses_api_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if responses_api_provider_config is None: @@ -1318,11 +1326,11 @@ def cancel_responses( raise ValueError("custom_llm_provider is required but passed as None") # get provider config - responses_api_provider_config: Optional[BaseResponsesAPIConfig] = ( - ProviderConfigManager.get_provider_responses_api_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + responses_api_provider_config: Optional[ + BaseResponsesAPIConfig + ] = ProviderConfigManager.get_provider_responses_api_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if responses_api_provider_config is None: @@ -1500,11 +1508,11 @@ def compact_responses( raise ValueError("custom_llm_provider is required but passed as None") # get provider config - responses_api_provider_config: Optional[BaseResponsesAPIConfig] = ( - ProviderConfigManager.get_provider_responses_api_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), - ) + responses_api_provider_config: Optional[ + BaseResponsesAPIConfig + ] = ProviderConfigManager.get_provider_responses_api_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), ) if responses_api_provider_config is None: diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 6ce59e3e67..0b0004d154 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -15,6 +15,69 @@ from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper +def _add_mcp_metadata_to_response( + response: Union[ModelResponse, CustomStreamWrapper], + openai_tools: Optional[List], + tool_calls: Optional[List] = None, + tool_results: Optional[List] = None, +) -> None: + """ + Add MCP metadata to response's provider_specific_fields. + + This function adds MCP-related information to the response so that + clients can access which tools were available, which were called, and + what results were returned. + + For ModelResponse: adds to choices[].message.provider_specific_fields + For CustomStreamWrapper: stores in _hidden_params and automatically adds to + final chunk's delta.provider_specific_fields via CustomStreamWrapper._add_mcp_metadata_to_final_chunk() + """ + if isinstance(response, CustomStreamWrapper): + # For streaming, store MCP metadata in _hidden_params + # CustomStreamWrapper._add_mcp_metadata_to_final_chunk() will automatically + # add it to the final chunk's delta.provider_specific_fields + if not hasattr(response, "_hidden_params"): + response._hidden_params = {} + + mcp_metadata = {} + if openai_tools: + mcp_metadata["mcp_list_tools"] = openai_tools + if tool_calls: + mcp_metadata["mcp_tool_calls"] = tool_calls + if tool_results: + mcp_metadata["mcp_call_results"] = tool_results + + if mcp_metadata: + response._hidden_params["mcp_metadata"] = mcp_metadata + return + + if not isinstance(response, ModelResponse): + return + + if not hasattr(response, "choices") or not response.choices: + return + + # Add MCP metadata to all choices' messages + for choice in response.choices: + message = getattr(choice, "message", None) + if message is not None: + # Get existing provider_specific_fields or create new dict + provider_fields = ( + getattr(message, "provider_specific_fields", None) or {} + ) + + # Add MCP metadata + if openai_tools: + provider_fields["mcp_list_tools"] = openai_tools + if tool_calls: + provider_fields["mcp_tool_calls"] = tool_calls + if tool_results: + provider_fields["mcp_call_results"] = tool_results + + # Set the provider_specific_fields + setattr(message, "provider_specific_fields", provider_fields) + + async def acompletion_with_mcp( model: str, messages: List, @@ -103,7 +166,13 @@ async def acompletion_with_mcp( # If not auto-executing, just make the call with transformed tools if not should_auto_execute: - return await litellm_acompletion(**base_call_args) + response = await litellm_acompletion(**base_call_args) + if isinstance(response, (ModelResponse, CustomStreamWrapper)): + _add_mcp_metadata_to_response( + response=response, + openai_tools=openai_tools, + ) + return response # For auto-execute: disable streaming for initial call stream = kwargs.get("stream", False) @@ -130,7 +199,17 @@ async def acompletion_with_mcp( if stream: retry_args = dict(base_call_args) retry_args["stream"] = stream - return await litellm_acompletion(**retry_args) + response = await litellm_acompletion(**retry_args) + if isinstance(response, (ModelResponse, CustomStreamWrapper)): + _add_mcp_metadata_to_response( + response=response, + openai_tools=openai_tools, + ) + return response + _add_mcp_metadata_to_response( + response=initial_response, + openai_tools=openai_tools, + ) return initial_response # Execute tool calls @@ -142,9 +221,16 @@ async def acompletion_with_mcp( mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, + litellm_call_id=kwargs.get("litellm_call_id"), + litellm_trace_id=kwargs.get("litellm_trace_id"), ) if not tool_results: + _add_mcp_metadata_to_response( + response=initial_response, + openai_tools=openai_tools, + tool_calls=tool_calls, + ) return initial_response # Create follow-up messages with tool results @@ -159,4 +245,12 @@ async def acompletion_with_mcp( follow_up_call_args["messages"] = follow_up_messages follow_up_call_args["stream"] = stream - return await litellm_acompletion(**follow_up_call_args) + response = await litellm_acompletion(**follow_up_call_args) + if isinstance(response, (ModelResponse, CustomStreamWrapper)): + _add_mcp_metadata_to_response( + response=response, + openai_tools=openai_tools, + tool_calls=tool_calls, + tool_results=tool_results, + ) + return response diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 9cdcd3894e..4376e076a9 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,3 +1,5 @@ +import traceback +from datetime import datetime from typing import ( TYPE_CHECKING, Any, @@ -11,17 +13,32 @@ from typing import ( ) from litellm._logging import verbose_logger +from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._experimental.mcp_server.utils import split_server_prefix_from_name from litellm.responses.main import aresponses from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator -from litellm.types.llms.openai import ResponsesAPIResponse, ToolParam -from litellm.types.utils import Choices, ModelResponse +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import ( + CallTypes, + Choices, + ModelResponse, + StandardLoggingMCPToolCall, +) +from litellm.utils import Rules, function_setup if TYPE_CHECKING: from mcp.types import Tool as MCPTool + from litellm.proxy.utils import ProxyLogging else: MCPTool = Any +# NOTE: We intentionally keep ToolParam as a broad type here to avoid tight coupling +# to optional OpenAI SDK typing symbols in environments that may not have them available. +# `Any` is used to keep mypy compatible with the broader OpenAI tool union types +# passed around in Responses API while still allowing dict-style access at runtime. +ToolParam = Any + LITELLM_PROXY_MCP_SERVER_URL = "litellm_proxy" LITELLM_PROXY_MCP_SERVER_URL_PREFIX = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" @@ -117,6 +134,8 @@ class LiteLLM_Proxy_MCP_Handler: mcp_auth_header=None, mcp_servers=mcp_servers, mcp_server_auth_headers=None, + log_list_tools_to_spendlogs=True, + list_tools_log_source="responses", ) allowed_mcp_server_ids = ( await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) @@ -462,7 +481,7 @@ class LiteLLM_Proxy_MCP_Handler: return result_text or "Tool executed successfully" @staticmethod - async def _execute_tool_calls( + async def _execute_tool_calls( # noqa: PLR0915 tool_server_map: dict[str, str], tool_calls: List[Any], user_api_key_auth: Any, @@ -470,6 +489,8 @@ class LiteLLM_Proxy_MCP_Handler: mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, + litellm_call_id: Optional[str] = None, + litellm_trace_id: Optional[str] = None, ) -> List[Dict[str, Any]]: """Execute tool calls and return results.""" from fastapi import HTTPException @@ -478,10 +499,16 @@ class LiteLLM_Proxy_MCP_Handler: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + from litellm.proxy.proxy_server import proxy_logging_obj + + from litellm._uuid import uuid tool_results = [] tool_call_id: Optional[str] = None + rules_obj = Rules() for tool_call in tool_calls: + logging_request_data: Dict[str, Any] = {} + tool_name: Optional[str] = None try: ( tool_name, @@ -514,6 +541,103 @@ class LiteLLM_Proxy_MCP_Handler: ): sanitized_tool_name = unprefixed_name + start_time = datetime.now() + logging_input = [ + { + "role": "tool", + "content": { + "tool_name": sanitized_tool_name, + "arguments": parsed_arguments, + }, + } + ] + tool_logging_call_id = litellm_call_id or str(uuid.uuid4()) + logging_request_data = { + "model": f"MCP: {tool_name}", + "metadata": { + "tool_call_id": tool_call_id, + "tool_name": sanitized_tool_name, + "server_name": server_name, + }, + "input": logging_input, + "call_type": CallTypes.call_mcp_tool.value, + "litellm_call_id": tool_logging_call_id, + } + if litellm_trace_id: + logging_request_data["litellm_trace_id"] = litellm_trace_id + user_identifier = None + if user_api_key_auth is not None: + user_api_key = getattr(user_api_key_auth, "api_key", None) + if user_api_key: + logging_request_data["metadata"]["user_api_key"] = user_api_key + + user_identifier = getattr( + user_api_key_auth, "end_user_id", None + ) or getattr(user_api_key_auth, "user_id", None) + if user_identifier: + logging_request_data["user"] = user_identifier + + litellm_logging_obj: Optional[LiteLLMLoggingObj] = None + try: + litellm_logging_obj, _ = function_setup( + original_function="call_mcp_tool", + rules_obj=rules_obj, + start_time=start_time, + **logging_request_data, + ) + except Exception as logging_error: + verbose_logger.debug( + "Failed to initialize logging for MCP tool call %s: %s", + tool_name, + logging_error, + ) + litellm_logging_obj = None + + logging_request_data["litellm_logging_obj"] = litellm_logging_obj + logging_request_data["arguments"] = parsed_arguments + + if litellm_logging_obj: + try: + litellm_logging_obj.pre_call( + input=logging_input, + api_key="", + ) + except Exception: + verbose_logger.exception( + "Failed to run pre_call for MCP tool logging" + ) + + standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = { + "name": sanitized_tool_name, + "arguments": parsed_arguments, + "namespaced_tool_name": tool_name, + } + mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name( + tool_name + ) + if mcp_server: + mcp_info = mcp_server.mcp_info or {} + standard_logging_mcp_tool_call["mcp_server_name"] = ( + mcp_info.get("server_name") + or getattr(mcp_server, "server_name", None) + or server_name + ) + logo_url = mcp_info.get("logo_url") + if logo_url: + standard_logging_mcp_tool_call["mcp_server_logo_url"] = logo_url + cost_info = mcp_info.get("mcp_server_cost_info") + if cost_info: + standard_logging_mcp_tool_call[ + "mcp_server_cost_info" + ] = cost_info + + if litellm_logging_obj: + litellm_logging_obj.model_call_details[ + "mcp_tool_call_metadata" + ] = standard_logging_mcp_tool_call + litellm_logging_obj.model = f"MCP: {tool_name}" + litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value + result = await global_mcp_server_manager.call_tool( server_name=server_name, name=sanitized_tool_name, @@ -526,6 +650,26 @@ class LiteLLM_Proxy_MCP_Handler: proxy_logging_obj=proxy_logging_obj, ) + if litellm_logging_obj: + try: + litellm_logging_obj.post_call(original_response=result) + end_time = datetime.now() + await litellm_logging_obj.async_post_mcp_tool_call_hook( + kwargs=litellm_logging_obj.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + await litellm_logging_obj.async_success_handler( + result=result, + start_time=start_time, + end_time=end_time, + ) + except Exception: + verbose_logger.exception( + "Failed to log MCP tool call success for %s", tool_name + ) + # Format result for inclusion in response result_text = LiteLLM_Proxy_MCP_Handler._parse_mcp_result(result) tool_results.append( @@ -537,6 +681,12 @@ class LiteLLM_Proxy_MCP_Handler: ) except BlockedPiiEntityError as e: + await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( + proxy_logging_obj=proxy_logging_obj, + user_api_key_auth=user_api_key_auth, + request_data=logging_request_data, + error=e, + ) verbose_logger.error( f"BlockedPiiEntityError in MCP tool call: {str(e)}" ) @@ -549,6 +699,12 @@ class LiteLLM_Proxy_MCP_Handler: } ) except GuardrailRaisedException as e: + await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( + proxy_logging_obj=proxy_logging_obj, + user_api_key_auth=user_api_key_auth, + request_data=logging_request_data, + error=e, + ) verbose_logger.error( f"GuardrailRaisedException in MCP tool call: {str(e)}" ) @@ -561,12 +717,28 @@ class LiteLLM_Proxy_MCP_Handler: } ) except HTTPException as e: + await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( + proxy_logging_obj=proxy_logging_obj, + user_api_key_auth=user_api_key_auth, + request_data=logging_request_data, + error=e, + ) verbose_logger.error(f"HTTPException in MCP tool call: {str(e)}") error_message = f"Tool call failed: {str(e.detail) if hasattr(e, 'detail') else str(e)}" tool_results.append( - {"tool_call_id": tool_call_id, "result": error_message} + { + "tool_call_id": tool_call_id, + "result": error_message, + "name": tool_name, + } ) except Exception as e: + await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure( + proxy_logging_obj=proxy_logging_obj, + user_api_key_auth=user_api_key_auth, + request_data=logging_request_data, + error=e, + ) verbose_logger.exception(f"Error executing MCP tool call: {e}") tool_results.append( { @@ -718,6 +890,31 @@ class LiteLLM_Proxy_MCP_Handler: **call_params, ) + @staticmethod + async def _log_mcp_tool_failure( + *, + proxy_logging_obj: Optional["ProxyLogging"], + user_api_key_auth: Any, + request_data: Dict[str, Any], + error: Exception, + ) -> None: + """Log MCP tool failures via proxy logging hooks.""" + + if proxy_logging_obj is None or user_api_key_auth is None: + return + + try: + traceback_str = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) + await proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=error, + user_api_key_dict=user_api_key_auth, + route="/responses/mcp/call_tool", + traceback_str=traceback_str, + ) + except Exception: + verbose_logger.exception("Failed to log MCP tool call failure") + @staticmethod def _create_mcp_streaming_response( input: Union[str, Any], @@ -758,7 +955,8 @@ class LiteLLM_Proxy_MCP_Handler: mcp_events=mcp_discovery_events, # Pre-generated MCP discovery events tool_server_map=tool_server_map, mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, - user_api_key_auth=kwargs.get("user_api_key_auth"), + user_api_key_auth=kwargs.get("user_api_key_auth") + or kwargs.get("litellm_metadata", {}).get("user_api_key_auth"), original_request_params=request_params, ) diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index ac040d3d6e..731aa5c692 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -273,9 +273,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.finished = False # Event queues and generation flags - self.mcp_discovery_events: List[ResponsesAPIStreamingResponse] = ( - mcp_events # Pre-generated MCP discovery events - ) + self.mcp_discovery_events: List[ + ResponsesAPIStreamingResponse + ] = mcp_events # Pre-generated MCP discovery events self.tool_execution_events: List[ResponsesAPIStreamingResponse] = [] self.mcp_discovery_generated = True # Events are already generated self.mcp_events = ( @@ -284,9 +284,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.tool_server_map = tool_server_map # Iterator references - self.base_iterator: Optional[Union[Any, ResponsesAPIResponse]] = ( - base_iterator # Will be created when needed - ) + self.base_iterator: Optional[ + Union[Any, ResponsesAPIResponse] + ] = base_iterator # Will be created when needed self.follow_up_iterator: Optional[Any] = None # Response collection for tool execution @@ -298,12 +298,14 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.custom_llm_provider = self.original_request_params.get( "custom_llm_provider", None ) + self.litellm_call_id = self.original_request_params.get("litellm_call_id") + self.litellm_trace_id = self.original_request_params.get("litellm_trace_id") self._extract_mcp_headers_from_params() # Mark as async iterator self.is_async = True - + def _extract_mcp_headers_from_params(self) -> None: """Extract MCP headers from original request params to pass to tool calls""" from typing import Dict, Optional @@ -311,25 +313,31 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) - + # Extract headers from secret_fields in original_request_params raw_headers_from_request: Optional[Dict[str, str]] = None secret_fields = self.original_request_params.get("secret_fields") if secret_fields and isinstance(secret_fields, dict): raw_headers_from_request = secret_fields.get("raw_headers") - + # Extract MCP-specific headers self.mcp_auth_header: Optional[str] = None self.mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None self.oauth2_headers: Optional[Dict[str, str]] = None self.raw_headers: Optional[Dict[str, str]] = raw_headers_from_request - + if raw_headers_from_request: headers_obj = Headers(raw_headers_from_request) - self.mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(headers_obj) - self.mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers_obj) - self.oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers_obj) - + self.mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers( + headers_obj + ) + self.mcp_server_auth_headers = ( + MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers_obj) + ) + self.oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers( + headers_obj + ) + # Also check if headers are provided in tools array (from request body) tools = self.original_request_params.get("tools") if tools: @@ -339,17 +347,26 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if tool_headers and isinstance(tool_headers, dict): # Merge tool headers into mcp_server_auth_headers headers_obj_from_tool = Headers(tool_headers) - tool_mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers_obj_from_tool) - + tool_mcp_server_auth_headers = ( + MCPRequestHandler._get_mcp_server_auth_headers_from_headers( + headers_obj_from_tool + ) + ) + if tool_mcp_server_auth_headers: if self.mcp_server_auth_headers is None: self.mcp_server_auth_headers = {} # Merge the headers from tool into existing headers - for server_alias, headers_dict in tool_mcp_server_auth_headers.items(): + for ( + server_alias, + headers_dict, + ) in tool_mcp_server_auth_headers.items(): if server_alias not in self.mcp_server_auth_headers: self.mcp_server_auth_headers[server_alias] = {} - self.mcp_server_auth_headers[server_alias].update(headers_dict) - + self.mcp_server_auth_headers[server_alias].update( + headers_dict + ) + # Also merge raw headers if self.raw_headers is None: self.raw_headers = {} @@ -487,9 +504,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Use the pre-fetched all_tools from original_request_params (no re-processing needed) params_for_llm = {} for key, value in params.items(): - params_for_llm[key] = ( - value # Copy all params as-is since tools are already processed - ) + params_for_llm[ + key + ] = value # Copy all params as-is since tools are already processed tools_count = ( len(params_for_llm.get("tools", [])) @@ -543,9 +560,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): return for tool_call in tool_calls: - tool_name, tool_arguments, tool_call_id = ( - LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call) - ) + ( + tool_name, + tool_arguments, + tool_call_id, + ) = LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call) if tool_name and tool_call_id: # Create MCP call events for this tool execution call_events = create_mcp_call_events( @@ -568,6 +587,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): mcp_server_auth_headers=self.mcp_server_auth_headers, oauth2_headers=self.oauth2_headers, raw_headers=self.raw_headers, + litellm_call_id=self.litellm_call_id, + litellm_trace_id=self.litellm_trace_id, ) # Create completion events and output_item.done events for tool execution diff --git a/litellm/router.py b/litellm/router.py index 3e312c154b..62070d2375 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -63,6 +63,7 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker +from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.least_busy import LeastBusyLoggingHandler from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler @@ -5877,7 +5878,8 @@ class Router: ), ) # done reading model["litellm_params"] - if custom_llm_provider not in litellm.provider_list: + # Check if provider is supported: either in enum or JSON-configured + if custom_llm_provider not in litellm.provider_list and not JSONProviderRegistry.exists(custom_llm_provider): raise Exception(f"Unsupported provider - {custom_llm_provider}") #### DEPLOYMENT NAMES INIT ######## diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 779a6950d9..8d18322d37 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -359,6 +359,7 @@ class AnthropicMessagesRequestOptionalParams(TypedDict, total=False): mcp_servers: Optional[List[AnthropicMcpServerTool]] context_management: Optional[Dict[str, Any]] container: Optional[Dict[str, Any]] # Container config with skills for code execution + output_format: Optional[AnthropicOutputSchema] # Structured outputs support class AnthropicMessagesRequest(AnthropicMessagesRequestOptionalParams, total=False): @@ -642,4 +643,8 @@ ANTHROPIC_TOOL_SEARCH_BETA_HEADER = "advanced-tool-use-2025-11-20" # Effort beta header constant ANTHROPIC_EFFORT_BETA_HEADER = "effort-2025-11-24" +# OAuth constants +ANTHROPIC_OAUTH_TOKEN_PREFIX = "sk-ant-oat" +ANTHROPIC_OAUTH_BETA_HEADER = "oauth-2025-04-20" + diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 467c57c33d..3f9842de7d 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -654,6 +654,8 @@ class ChatCompletionFileObjectFile(TypedDict, total=False): file_id: str filename: str format: str + detail: str # For video/image resolution control (low, medium, high, ultra_high) + video_metadata: Dict[str, Any] # For video-specific metadata (fps, start_offset, end_offset) class ChatCompletionFileObject(TypedDict): @@ -1191,7 +1193,7 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): status: Optional[str] = None text: Optional[Union["ResponseText", Dict[str, Any]]] = None truncation: Optional[Literal["auto", "disabled"]] = None - usage: Optional[ResponseAPIUsage] = None + usage: Optional[Any] = None user: Optional[str] = None store: Optional[bool] = None # Define private attributes using PrivateAttr @@ -1248,6 +1250,8 @@ class ResponsesAPIStreamEvents(str, Enum): # Reasoning summary events RESPONSE_PART_ADDED = "response.reasoning_summary_part.added" REASONING_SUMMARY_TEXT_DELTA = "response.reasoning_summary_text.delta" + REASONING_SUMMARY_TEXT_DONE = "response.reasoning_summary_text.done" + REASONING_SUMMARY_PART_DONE = "response.reasoning_summary_part.done" # Output item events OUTPUT_ITEM_ADDED = "response.output_item.added" @@ -1337,6 +1341,24 @@ class ReasoningSummaryTextDeltaEvent(BaseLiteLLMOpenAIResponseObject): delta: str +class ReasoningSummaryTextDoneEvent(BaseLiteLLMOpenAIResponseObject): + type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE] + item_id: str + output_index: int + sequence_number: int + summary_index: int + text: str + + +class ReasoningSummaryPartDoneEvent(BaseLiteLLMOpenAIResponseObject): + type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE] + item_id: str + output_index: int + sequence_number: int + summary_index: int + part: BaseLiteLLMOpenAIResponseObject + + class OutputItemAddedEvent(BaseLiteLLMOpenAIResponseObject): type: Literal[ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED] output_index: int @@ -1591,6 +1613,8 @@ ResponsesAPIStreamingResponse = Annotated[ ResponseIncompleteEvent, ResponsePartAddedEvent, ReasoningSummaryTextDeltaEvent, + ReasoningSummaryTextDoneEvent, + ReasoningSummaryPartDoneEvent, OutputItemAddedEvent, OutputItemDoneEvent, ContentPartAddedEvent, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 53e0a8f185..41a09bd391 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -63,11 +63,11 @@ def _generate_id(): # private helper function return "chatcmpl-" + str(uuid.uuid4()) - class SafeAttributeModel: """ A base model that provides safe attribute access. """ + def __delattr__(self, name): try: super().__delattr__(name) @@ -125,13 +125,14 @@ class SearchContextCostPerQuery(TypedDict, total=False): class AgenticLoopParams(TypedDict, total=False): """ Parameters passed to agentic loop hooks (e.g., WebSearch interception). - + Stored in logging_obj.model_call_details["agentic_loop_params"] to provide agentic hooks with the original request context needed for follow-up calls. """ + model: str """The model string with provider prefix (e.g., 'bedrock/invoke/...')""" - + custom_llm_provider: str """The LLM provider name (e.g., 'bedrock', 'anthropic')""" @@ -384,6 +385,7 @@ class CallTypes(str, Enum): # MCP Call Types ######################################################### call_mcp_tool = "call_mcp_tool" + list_mcp_tools = "list_mcp_tools" ######################################################### # A2A Call Types @@ -448,6 +450,7 @@ CallTypesLiteral = Literal[ "vector_store_file_delete", "avector_store_file_delete", "call_mcp_tool", + "list_mcp_tools", "asend_message", "send_message", "aresponses", @@ -1343,8 +1346,7 @@ class CacheCreationTokenDetails(BaseModel): class PromptTokensDetailsWrapper( - SafeAttributeModel, - PromptTokensDetails + SafeAttributeModel, PromptTokensDetails ): # extends with image generation fields (text_tokens, image_tokens) text_tokens: Optional[int] = None """Text tokens sent to the model.""" diff --git a/litellm/utils.py b/litellm/utils.py index 8fdbc2a3fc..22a2987285 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -771,7 +771,8 @@ def function_setup( # noqa: PLR0915 function_id: Optional[str] = kwargs["id"] if "id" in kwargs else None ## LAZY LOAD COROUTINE CHECKER ## - get_coroutine_checker = getattr(sys.modules[__name__], "get_coroutine_checker") + get_coroutine_checker_fn = getattr(sys.modules[__name__], "get_coroutine_checker") + coroutine_checker = get_coroutine_checker_fn() ## DYNAMIC CALLBACKS ## dynamic_callbacks: Optional[ @@ -825,7 +826,7 @@ def function_setup( # noqa: PLR0915 if len(litellm.input_callback) > 0: removed_async_items = [] for index, callback in enumerate(litellm.input_callback): # type: ignore - if get_coroutine_checker().is_async_callable(callback): + if coroutine_checker.is_async_callable(callback): litellm._async_input_callback.append(callback) removed_async_items.append(index) @@ -835,7 +836,7 @@ def function_setup( # noqa: PLR0915 if len(litellm.success_callback) > 0: removed_async_items = [] for index, callback in enumerate(litellm.success_callback): # type: ignore - if get_coroutine_checker().is_async_callable(callback): + if coroutine_checker.is_async_callable(callback): litellm.logging_callback_manager.add_litellm_async_success_callback( callback ) @@ -860,7 +861,7 @@ def function_setup( # noqa: PLR0915 if len(litellm.failure_callback) > 0: removed_async_items = [] for index, callback in enumerate(litellm.failure_callback): # type: ignore - if get_coroutine_checker().is_async_callable(callback): + if coroutine_checker.is_async_callable(callback): litellm.logging_callback_manager.add_litellm_async_failure_callback( callback ) @@ -893,7 +894,7 @@ def function_setup( # noqa: PLR0915 removed_async_items = [] for index, callback in enumerate(kwargs["success_callback"]): if ( - get_coroutine_checker().is_async_callable(callback) + coroutine_checker.is_async_callable(callback) or callback == "dynamodb" or callback == "s3" ): @@ -2735,6 +2736,7 @@ def register_model(model_cost: Union[str, dict]): # noqa: PLR0915 # Skip get_model_info for these providers during model registration _skip_get_model_info_providers = { LlmProviders.GITHUB_COPILOT.value, + LlmProviders.CHATGPT.value, } for key, value in loaded_model_cost.items(): @@ -8257,6 +8259,10 @@ class ProviderConfigManager: return litellm.ClarifaiConfig() elif LlmProviders.BEDROCK == provider: return litellm.llms.bedrock.common_utils.BedrockModelInfo() + elif LlmProviders.AZURE_AI == provider: + from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + + return AzureFoundryModelInfo(model=model) return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d35f9b80ce..6d87e0b599 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -12696,8 +12696,8 @@ "supports_web_search": true }, "gemini-2.5-flash-lite": { - "cache_read_input_token_cost": 2.5e-08, - "input_cost_per_audio_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", "max_audio_length_hours": 8.4, @@ -12741,7 +12741,7 @@ "supports_web_search": true }, "gemini-2.5-flash-lite-preview-09-2025": { - "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", @@ -14532,8 +14532,8 @@ "supports_web_search": true }, "gemini/gemini-2.5-flash-lite": { - "cache_read_input_token_cost": 2.5e-08, - "input_cost_per_audio_token": 5e-07, + "cache_read_input_token_cost": 1e-08, + "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", "max_audio_length_hours": 8.4, @@ -14579,7 +14579,7 @@ "tpm": 250000 }, "gemini/gemini-2.5-flash-lite-preview-09-2025": { - "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", @@ -17038,14 +17038,14 @@ "supports_vision": true }, "gpt-4o-audio-preview": { - "input_cost_per_audio_token": 0.0001, + "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_audio_token": 0.0002, + "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 1e-05, "supports_audio_input": true, "supports_audio_output": true, @@ -17055,14 +17055,14 @@ "supports_tool_choice": true }, "gpt-4o-audio-preview-2024-10-01": { - "input_cost_per_audio_token": 0.0001, + "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_audio_token": 0.0002, + "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 1e-05, "supports_audio_input": true, "supports_audio_output": true, @@ -17105,6 +17105,186 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "gpt-audio": { + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/realtime", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "gpt-audio-2025-08-28": { + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/realtime", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "gpt-audio-mini": { + "input_cost_per_audio_token": 1e-05, + "input_cost_per_token": 6e-07, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 2e-05, + "output_cost_per_token": 2.4e-06, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/realtime", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "gpt-audio-mini-2025-10-06": { + "input_cost_per_audio_token": 1e-05, + "input_cost_per_token": 6e-07, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 2e-05, + "output_cost_per_token": 2.4e-06, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/realtime", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "gpt-audio-mini-2025-12-15": { + "input_cost_per_audio_token": 1e-05, + "input_cost_per_token": 6e-07, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 2e-05, + "output_cost_per_token": 2.4e-06, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/realtime", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, "gpt-4o-mini": { "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_priority": 1.25e-07, @@ -34237,5 +34417,18 @@ "output_cost_per_token": 0, "litellm_provider": "llamagate", "mode": "embedding" + }, + "sarvam/sarvam-m": { + "cache_creation_input_token_cost": 0, + "cache_creation_input_token_cost_above_1hr": 0, + "cache_read_input_token_cost": 0, + "input_cost_per_token": 0, + "litellm_provider": "sarvam", + "max_input_tokens": 8192, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "chat", + "output_cost_per_token": 0, + "supports_reasoning": true } } diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 343d5bd2c6..a901739c46 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2393,6 +2393,15 @@ "a2a": true, "interactions": true } + }, + "sarvam": { + "display_name": "Sarvam (`sarvam`)", + "url": "https://docs.litellm.ai/docs/providers/sarvam", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true + } } }, "endpoints": { diff --git a/proxy_config.yaml b/proxy_config.yaml new file mode 100644 index 0000000000..57397181cd --- /dev/null +++ b/proxy_config.yaml @@ -0,0 +1,7 @@ +model_list: + - model_name: "*" + litellm_params: + model: "*" + +general_settings: + master_key: sk-1234 diff --git a/pyproject.toml b/pyproject.toml index 0be1692668..bd35bd70f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "1.81.0" +version = "1.81.1" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT" @@ -173,7 +173,7 @@ requires = ["poetry-core", "wheel"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "1.81.0" +version = "1.81.1" version_files = [ "pyproject.toml:^version" ] diff --git a/test_anthropic_messages_structured_outputs_minimal.py b/test_anthropic_messages_structured_outputs_minimal.py new file mode 100644 index 0000000000..3fc7dc9a56 --- /dev/null +++ b/test_anthropic_messages_structured_outputs_minimal.py @@ -0,0 +1,74 @@ +""" +Tests for structured outputs support in Anthropic /v1/messages endpoint. +""" +import pytest +from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, +) + + +def test_output_format_supported_and_transforms_correctly(): + """Test that output_format is supported and properly transformed with beta header.""" + config = AnthropicMessagesConfig() + + # 1. Verify it's in supported parameters + supported_params = config.get_supported_anthropic_messages_params("claude-sonnet-4-5") + assert "output_format" in supported_params + + # 2. Verify transformation preserves output_format and adds beta header + output_format = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"result": {"type": "string"}}} + } + + optional_params = {"max_tokens": 1024, "output_format": output_format} + headers = {} + + # Transform request + result = config.transform_anthropic_messages_request( + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "test"}], + anthropic_messages_optional_request_params=optional_params.copy(), + litellm_params={}, + headers=headers + ) + + # Update headers + headers = config._update_headers_with_anthropic_beta(headers, optional_params) + + # Verify output_format preserved in request body + assert "output_format" in result + assert result["output_format"]["type"] == "json_schema" + + # Verify beta header added + assert "anthropic-beta" in headers + assert "structured-outputs-2025-11-13" in headers["anthropic-beta"] + + +def test_output_format_works_with_bedrock_and_azure(): + """Test that output_format works with Bedrock and Azure Foundry models.""" + config = AnthropicMessagesConfig() + + output_format = {"type": "json_schema", "schema": {"type": "object", "properties": {}}} + optional_params = {"max_tokens": 1024, "output_format": output_format} + messages = [{"role": "user", "content": "test"}] + + # Test Bedrock + bedrock_result = config.transform_anthropic_messages_request( + model="bedrock/anthropic.claude-sonnet-4-5-v2:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params.copy(), + litellm_params={}, + headers={} + ) + assert "output_format" in bedrock_result + + # Test Azure Foundry + azure_result = config.transform_anthropic_messages_request( + model="azure_ai/claude-sonnet-4-5", + messages=messages, + anthropic_messages_optional_request_params=optional_params.copy(), + litellm_params={}, + headers={} + ) + assert "output_format" in azure_result \ No newline at end of file diff --git a/tests/litellm/proxy/management_endpoints/test_cost_estimate_endpoint.py b/tests/litellm/proxy/management_endpoints/test_cost_estimate_endpoint.py index f2d8d87855..bc0f3cf15b 100644 --- a/tests/litellm/proxy/management_endpoints/test_cost_estimate_endpoint.py +++ b/tests/litellm/proxy/management_endpoints/test_cost_estimate_endpoint.py @@ -73,3 +73,67 @@ class TestCostEstimateEndpoint: ) assert exc_info.value.status_code == 404 + + @pytest.mark.asyncio + async def test_estimate_cost_resolves_router_model_alias(self): + """ + Test that estimate_cost resolves router model aliases to underlying models. + + When a user selects a model like 'my-gpt4-alias' from the UI (which is a + router model_name), the endpoint should resolve it to the actual model + (e.g., 'azure/gpt-4') for cost calculation. + + This prevents the bug where custom model names fail cost lookup because + they aren't in model_prices_and_context_window.json. + """ + request = CostEstimateRequest( + model="my-gpt4-alias", # Router alias, not actual model name + input_tokens=1000, + output_tokens=500, + ) + + # Mock the router to return deployment info + mock_router = MagicMock() + mock_router.get_model_list.return_value = [ + { + "model_name": "my-gpt4-alias", + "litellm_params": { + "model": "azure/gpt-4", # Actual model for pricing + "custom_llm_provider": "azure", + }, + } + ] + + with patch( + "litellm.proxy.proxy_server.llm_router", + mock_router, + ): + with patch( + "litellm.proxy.management_endpoints.cost_tracking_settings.completion_cost" + ) as mock_completion_cost: + mock_completion_cost.return_value = 0.05 + + with patch("litellm.get_model_info") as mock_get_model_info: + mock_get_model_info.return_value = { + "input_cost_per_token": 0.00003, + "output_cost_per_token": 0.00006, + "litellm_provider": "azure", + } + + response = await estimate_cost( + request=request, + user_api_key_dict=MagicMock(), + ) + + # Verify router was queried for the alias + mock_router.get_model_list.assert_called_with(model_name="my-gpt4-alias") + + # Verify completion_cost was called with RESOLVED model, not the alias + call_args = mock_completion_cost.call_args + assert call_args.kwargs["model"] == "azure/gpt-4" + + # Verify response contains original model name (for UI display) + assert response.model == "my-gpt4-alias" + assert response.cost_per_request == 0.05 + assert response.provider == "azure" + diff --git a/tests/litellm_utils_tests/base_token_counter_test.py b/tests/litellm_utils_tests/base_token_counter_test.py new file mode 100644 index 0000000000..b5e87021a0 --- /dev/null +++ b/tests/litellm_utils_tests/base_token_counter_test.py @@ -0,0 +1,130 @@ +""" +Base Token Counter Test Suite. + +This module provides an abstract base test class that enforces common tests +across all token counter implementations. Similar to base_llm_unit_tests.py +for LLM chat tests. + +Usage: + Create a test class that inherits from BaseTokenCounterTest and implement + the abstract methods to provide provider-specific configuration. +""" + +import os +import sys +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.types.utils import TokenCountResponse + + +class BaseTokenCounterTest(ABC): + """ + Abstract base test class for token counter implementations. + + Subclasses must implement: + - get_token_counter(): Returns the token counter instance + - get_test_model(): Returns the model name to use for testing + - get_test_messages(): Returns test messages for token counting + - get_deployment_config(): Returns deployment configuration with credentials + - get_custom_llm_provider(): Returns the provider name for should_use_token_counting_api + """ + + @abstractmethod + def get_token_counter(self) -> BaseTokenCounter: + """Must return the token counter instance to test.""" + pass + + @abstractmethod + def get_test_model(self) -> str: + """Must return the model name to use for testing.""" + pass + + @abstractmethod + def get_test_messages(self) -> List[Dict[str, Any]]: + """Must return test messages for token counting.""" + pass + + @abstractmethod + def get_deployment_config(self) -> Dict[str, Any]: + """Must return deployment configuration with credentials.""" + pass + + @abstractmethod + def get_custom_llm_provider(self) -> str: + """Must return the provider name for should_use_token_counting_api check.""" + pass + + @pytest.fixture(autouse=True) + def _handle_missing_credentials(self): + """Fixture to skip tests when credentials are missing.""" + try: + yield + except Exception as e: + error_str = str(e).lower() + if "api key" in error_str or "api_key" in error_str or "unauthorized" in error_str: + pytest.skip(f"Missing or invalid credentials: {e}") + raise + + @pytest.mark.asyncio + async def test_count_tokens_basic(self): + """ + Test basic token counting functionality. + + Verifies that: + - Token counter returns a TokenCountResponse + - total_tokens is greater than 0 + - tokenizer_type is set + - No error occurred + """ + token_counter = self.get_token_counter() + model = self.get_test_model() + messages = self.get_test_messages() + deployment = self.get_deployment_config() + + result = await token_counter.count_tokens( + model_to_use=model, + messages=messages, + contents=None, + deployment=deployment, + request_model=model, + ) + + print(f"Token count result: {result}") + + assert result is not None, "Token counter should return a result" + assert isinstance(result, TokenCountResponse), "Result should be TokenCountResponse" + assert result.total_tokens > 0, f"Token count should be > 0, got {result.total_tokens}" + assert result.tokenizer_type is not None, "tokenizer_type should be set" + assert result.error is not True, f"Token counting should not error: {result.error_message}" + + def test_should_use_token_counting_api(self): + """ + Test that should_use_token_counting_api returns True for the correct provider. + + Verifies that the token counter correctly identifies when it should be used + based on the custom_llm_provider. + """ + token_counter = self.get_token_counter() + provider = self.get_custom_llm_provider() + + result = token_counter.should_use_token_counting_api( + custom_llm_provider=provider + ) + + assert result is True, f"should_use_token_counting_api should return True for {provider}" + + # Also verify it returns False for other providers + other_provider = "some_other_provider_that_doesnt_exist" + result_other = token_counter.should_use_token_counting_api( + custom_llm_provider=other_provider + ) + + assert result_other is False, f"should_use_token_counting_api should return False for {other_provider}" diff --git a/tests/litellm_utils_tests/test_anthropic_token_counter.py b/tests/litellm_utils_tests/test_anthropic_token_counter.py new file mode 100644 index 0000000000..a1fbcecfdd --- /dev/null +++ b/tests/litellm_utils_tests/test_anthropic_token_counter.py @@ -0,0 +1,47 @@ +""" +Anthropic Token Counter Tests. + +Tests for the Anthropic token counter implementation using the base test suite. +""" + +import os +import sys +from typing import Any, Dict, List + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +from litellm.llms.anthropic.count_tokens import AnthropicTokenCounter +from litellm.llms.base_llm.base_utils import BaseTokenCounter +from tests.litellm_utils_tests.base_token_counter_test import BaseTokenCounterTest + + +class TestAnthropicTokenCounter(BaseTokenCounterTest): + """Test suite for Anthropic token counter.""" + + def get_token_counter(self) -> BaseTokenCounter: + return AnthropicTokenCounter() + + def get_test_model(self) -> str: + return "claude-sonnet-4-20250514" + + def get_test_messages(self) -> List[Dict[str, Any]]: + return [ + {"role": "user", "content": "Hello, how are you today?"} + ] + + def get_deployment_config(self) -> Dict[str, Any]: + api_key = os.getenv("ANTHROPIC_API_KEY") + if not api_key: + pytest.skip("ANTHROPIC_API_KEY not set") + return { + "litellm_params": { + "api_key": api_key, + } + } + + def get_custom_llm_provider(self) -> str: + return "anthropic" diff --git a/tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py b/tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py new file mode 100644 index 0000000000..031502cbec --- /dev/null +++ b/tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py @@ -0,0 +1,53 @@ +""" +Azure AI Anthropic Token Counter Tests. + +Tests for the Azure AI Anthropic token counter implementation using the base test suite. +""" + +import os +import sys +from typing import Any, Dict, List + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +from litellm.llms.azure_ai.anthropic.count_tokens import AzureAIAnthropicTokenCounter +from litellm.llms.base_llm.base_utils import BaseTokenCounter +from tests.litellm_utils_tests.base_token_counter_test import BaseTokenCounterTest + + +class TestAzureAIAnthropicTokenCounter(BaseTokenCounterTest): + """Test suite for Azure AI Anthropic token counter.""" + + def get_token_counter(self) -> BaseTokenCounter: + return AzureAIAnthropicTokenCounter() + + def get_test_model(self) -> str: + return "claude-3-5-sonnet" + + def get_test_messages(self) -> List[Dict[str, Any]]: + return [ + {"role": "user", "content": "Hello, how are you today?"} + ] + + def get_deployment_config(self) -> Dict[str, Any]: + api_key = os.getenv("AZURE_AI_API_KEY") + api_base = os.getenv("AZURE_AI_API_BASE") + + if not api_key: + pytest.skip("AZURE_AI_API_KEY not set") + if not api_base: + pytest.skip("AZURE_AI_API_BASE not set") + + return { + "litellm_params": { + "api_key": api_key, + "api_base": api_base, + } + } + + def get_custom_llm_provider(self) -> str: + return "azure_ai" diff --git a/tests/litellm_utils_tests/test_bedrock_token_counter.py b/tests/litellm_utils_tests/test_bedrock_token_counter.py new file mode 100644 index 0000000000..f7c2991882 --- /dev/null +++ b/tests/litellm_utils_tests/test_bedrock_token_counter.py @@ -0,0 +1,101 @@ +""" +Bedrock Token Counter Tests. + +Tests for the Bedrock token counter implementation using the base test suite. + +Note: Not all Bedrock models support token counting. The CountTokens API +is only available for specific models. If the model doesn't support token +counting, the test will be skipped. +""" + +import os +import sys +from typing import Any, Dict, List + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter +from tests.litellm_utils_tests.base_token_counter_test import BaseTokenCounterTest + + +class TestBedrockTokenCounter(BaseTokenCounterTest): + """Test suite for Bedrock token counter. + + Note: Bedrock CountTokens API support varies by model. Some models + (like older Claude versions) may not support token counting. + Use amazon.nova-* models for reliable token counting support. + """ + + def get_token_counter(self) -> BaseTokenCounter: + return BedrockTokenCounter() + + def get_test_model(self) -> str: + # Use Amazon Nova model which supports token counting + # Alternatively, use environment variable to override + return os.getenv("BEDROCK_TEST_MODEL", "amazon.nova-lite-v1:0") + + def get_test_messages(self) -> List[Dict[str, Any]]: + return [ + {"role": "user", "content": "Hello, how are you today?"} + ] + + def get_deployment_config(self) -> Dict[str, Any]: + # Bedrock uses AWS credentials from environment + # Check for AWS credentials + aws_access_key = os.getenv("AWS_ACCESS_KEY_ID") + aws_secret_key = os.getenv("AWS_SECRET_ACCESS_KEY") + aws_region = os.getenv("AWS_REGION_NAME", "us-east-1") + + if not aws_access_key or not aws_secret_key: + pytest.skip("AWS credentials not set (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY)") + + return { + "litellm_params": { + "aws_access_key_id": aws_access_key, + "aws_secret_access_key": aws_secret_key, + "aws_region_name": aws_region, + } + } + + def get_custom_llm_provider(self) -> str: + return "bedrock" + + @pytest.mark.asyncio + async def test_count_tokens_basic(self): + """ + Test basic token counting functionality. + + Override to handle models that don't support token counting. + """ + from litellm.types.utils import TokenCountResponse + + token_counter = self.get_token_counter() + model = self.get_test_model() + messages = self.get_test_messages() + deployment = self.get_deployment_config() + + result = await token_counter.count_tokens( + model_to_use=model, + messages=messages, + contents=None, + deployment=deployment, + request_model=model, + ) + + print(f"Token count result: {result}") + + assert result is not None, "Token counter should return a result" + assert isinstance(result, TokenCountResponse), "Result should be TokenCountResponse" + + # Check if the model doesn't support token counting + if result.error and "doesn't support counting tokens" in str(result.error_message): + pytest.skip(f"Model {model} doesn't support token counting: {result.error_message}") + + assert result.total_tokens > 0, f"Token count should be > 0, got {result.total_tokens}" + assert result.tokenizer_type is not None, "tokenizer_type should be set" + assert result.error is not True, f"Token counting should not error: {result.error_message}" diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index a679fa4e4c..f585aadbfc 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -37,6 +37,7 @@ from litellm.utils import ( trim_messages, validate_environment, ) +from litellm.llms.openai_like.json_loader import JSONProviderRegistry from unittest.mock import AsyncMock, MagicMock, patch @@ -972,11 +973,11 @@ def test_logging_trace_id(langfuse_trace_id, langfuse_existing_trace_id): litellm_logging_obj._get_trace_id(service_name="langfuse") == langfuse_trace_id ) - ## if existing_trace_id exists + ## if no trace_id or existing_trace_id is provided, use litellm_trace_id else: assert ( litellm_logging_obj._get_trace_id(service_name="langfuse") - == litellm_call_id + == litellm_logging_obj.litellm_trace_id ) @@ -1383,7 +1384,7 @@ def test_models_by_provider(): providers.add(v["litellm_provider"]) for provider in providers: - assert provider in models_by_provider.keys() + assert provider in models_by_provider.keys() or JSONProviderRegistry.exists(provider) @pytest.mark.parametrize( diff --git a/tests/local_testing/test_gcs_bucket.py b/tests/local_testing/test_gcs_bucket.py index 2f7d5cd0de..e2f51eb76c 100644 --- a/tests/local_testing/test_gcs_bucket.py +++ b/tests/local_testing/test_gcs_bucket.py @@ -29,6 +29,7 @@ def load_vertex_ai_credentials(): # Define the path to the vertex_key.json file print("loading vertex ai credentials") os.environ["GCS_FLUSH_INTERVAL"] = "1" + os.environ["GCS_USE_BATCHED_LOGGING"] = "false" filepath = os.path.dirname(os.path.abspath(__file__)) vertex_key_path = filepath + "/vertex_key.json" diff --git a/tests/local_testing/test_get_llm_provider.py b/tests/local_testing/test_get_llm_provider.py index 4df910bc73..9b07111ea5 100644 --- a/tests/local_testing/test_get_llm_provider.py +++ b/tests/local_testing/test_get_llm_provider.py @@ -137,6 +137,9 @@ def test_default_api_base(): # Get the API base for the given provider if provider == "github_copilot": continue + # Skip chatgpt as it requires OAuth authentication + if provider == "chatgpt": + continue # Skip ragflow as it requires specific model format: ragflow/chat/{id}/{model} or ragflow/agent/{id}/{model} if provider == "ragflow": continue diff --git a/tests/local_testing/test_timeout.py b/tests/local_testing/test_timeout.py index 7f9837e81d..bb6dbc938b 100644 --- a/tests/local_testing/test_timeout.py +++ b/tests/local_testing/test_timeout.py @@ -285,3 +285,94 @@ async def test_anthropic_timeout(streaming, sync_mode): ) print(type(e)) pass + + +@pytest.mark.asyncio +async def test_timeout_respects_total_time_not_per_retry(): + """ + Test that timeout applies to the TOTAL operation time, not per-retry. + + This test ensures that when a user sets timeout=2, the entire operation + (including all retries) times out at ~2 seconds, not at 2s * num_retries. + + This is a regression test for the issue where timeout was being applied + per-retry attempt, causing the total time to be much longer than expected. + """ + litellm.set_verbose = False + + timeout_value = 2.0 + # Allow for some overhead (network, processing, etc.) + # but ensure we don't wait for multiple retries + max_allowed_time = timeout_value + 1.0 # 3 seconds max + + start_time = time.time() + + try: + # This should timeout because we're asking for a long response + # with a very short timeout + response = await litellm.acompletion( + model="gpt-3.5-turbo", + timeout=timeout_value, + messages=[{"role": "user", "content": "Write a very long detailed essay about the history of computing, at least 5000 words."}], + ) + pytest.fail("Expected timeout error but got a response") + except (openai.APITimeoutError, litellm.exceptions.Timeout) as e: + elapsed_time = time.time() - start_time + + print(f"Timeout occurred after {elapsed_time:.2f} seconds") + print(f"Expected timeout: {timeout_value} seconds") + print(f"Max allowed time: {max_allowed_time} seconds") + + # Verify that the timeout happened within the expected time window + # It should be close to timeout_value, not timeout_value * num_retries + assert elapsed_time < max_allowed_time, ( + f"Timeout took too long! Expected ~{timeout_value}s, " + f"got {elapsed_time:.2f}s. This suggests timeout is being " + f"applied per-retry instead of to the total operation." + ) + + # Also verify it's not TOO fast (sanity check) + assert elapsed_time >= timeout_value * 0.5, ( + f"Timeout happened too quickly: {elapsed_time:.2f}s. " + f"Expected at least {timeout_value * 0.5}s" + ) + + print("✓ Timeout correctly applied to total operation time, not per-retry") + except Exception as e: + pytest.fail( + f"Expected timeout error but got different error: {type(e).__name__}: {e}" + ) + + +@pytest.mark.asyncio +async def test_timeout_with_retries_disabled(): + """ + Test that timeout works correctly when retries are explicitly disabled. + This should timeout even faster since there are no retry attempts. + """ + litellm.set_verbose = False + + timeout_value = 2.0 + max_allowed_time = timeout_value + 0.5 # Even tighter bound with no retries + + start_time = time.time() + + try: + response = await litellm.acompletion( + model="gpt-3.5-turbo", + timeout=timeout_value, + max_retries=0, # Disable retries + messages=[{"role": "user", "content": "Write a very long detailed essay about the history of computing, at least 5000 words."}], + ) + pytest.fail("Expected timeout error but got a response") + except (openai.APITimeoutError, litellm.exceptions.Timeout) as e: + elapsed_time = time.time() - start_time + + print(f"Timeout with no retries occurred after {elapsed_time:.2f} seconds") + + assert elapsed_time < max_allowed_time, ( + f"Timeout took too long even with retries disabled! " + f"Expected ~{timeout_value}s, got {elapsed_time:.2f}s" + ) + + print("✓ Timeout works correctly with retries disabled") diff --git a/tests/logging_callback_tests/test_alerting.py b/tests/logging_callback_tests/test_alerting.py index 8a691e7618..524cc00d5f 100644 --- a/tests/logging_callback_tests/test_alerting.py +++ b/tests/logging_callback_tests/test_alerting.py @@ -866,11 +866,9 @@ async def test_langfuse_trace_id(): assert trace_url is not None - returned_trace_id = int(trace_url.split("/")[-1]) + returned_trace_id = trace_url.split("/")[-1] - assert returned_trace_id == int( - litellm_logging_obj._get_trace_id(service_name="langfuse") - ) + assert returned_trace_id == litellm_logging_obj._get_trace_id(service_name="langfuse") @pytest.mark.asyncio diff --git a/tests/mcp_tests/test_mcp_chat_completions.py b/tests/mcp_tests/test_mcp_chat_completions.py index 973301abfb..8857f016df 100644 --- a/tests/mcp_tests/test_mcp_chat_completions.py +++ b/tests/mcp_tests/test_mcp_chat_completions.py @@ -312,3 +312,209 @@ async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch): # Verify acompletion was called (should be called by acompletion_with_mcp) assert len(acompletion_calls) >= 1, "acompletion should be called" + + +@pytest.mark.asyncio +async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch): + """ + Test that MCP metadata is added to the final streaming chunk's + delta.provider_specific_fields when using MCP tools with streaming. + """ + from types import SimpleNamespace + from unittest.mock import patch + + from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, + ) + from litellm.responses.utils import ResponsesAPIRequestUtils + from litellm.utils import CustomStreamWrapper + from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta + from litellm.litellm_core_utils.litellm_logging import Logging + + dummy_tool = SimpleNamespace( + name="local_search", + description="search", + inputSchema={"type": "object", "properties": {}}, + ) + + async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy): + return [dummy_tool], {"local_search": "local"} + + async def fake_execute(**kwargs): + tool_calls = kwargs.get("tool_calls") or [] + call_entry = tool_calls[0] + call_id = call_entry.get("id") or call_entry.get("call_id") or "call" + return [ + { + "tool_call_id": call_id, + "result": "executed", + "name": call_entry.get("name", "local_search"), + } + ] + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + fake_process, + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + fake_execute, + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(lambda secret_fields, tools: (None, None, None, None)), + ) + + # Create mock streaming chunks + def create_chunk(content, finish_reason=None): + return ModelResponseStream( + id="test-stream", + model="gpt-4o-mini", + created=1234567890, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=content, + role="assistant", + ), + finish_reason=finish_reason, + ) + ], + ) + + chunks = [ + create_chunk("Hello"), + create_chunk(" world"), + create_chunk("!", finish_reason="stop"), # Final chunk + ] + + # Create a proper CustomStreamWrapper with logging_obj + from unittest.mock import MagicMock + logging_obj = MagicMock() + logging_obj.model_call_details = {} + + class MockStreamingResponse(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=logging_obj, + ) + self.chunks = chunks + self._index = 0 + self.sent_last_chunk = False + + def __iter__(self): + return self + + def __next__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + if self._index == len(self.chunks): + self.sent_last_chunk = True + # Call the method that adds MCP metadata to final chunk + chunk = self._add_mcp_metadata_to_final_chunk(chunk) + return chunk + raise StopIteration + + # Track calls to acompletion + acompletion_calls = [] + + async def mock_acompletion(**kwargs): + acompletion_calls.append(kwargs) + # First call (non-streaming for tool extraction) + if not kwargs.get("stream", False): + return ModelResponse( + id="test-1", + model="gpt-4o-mini", + choices=[{ + "message": { + "role": "assistant", + "tool_calls": [{ + "id": "call-1", + "type": "function", + "function": { + "name": "local_search", + "arguments": "{}" + } + }] + }, + "finish_reason": "tool_calls" + }], + created=0, + object="chat.completion", + ) + # Second call (streaming follow-up) + return MockStreamingResponse() + + with patch("litellm.acompletion", side_effect=mock_acompletion): + response = litellm.completion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + tools=[ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp/local", + "server_label": "local", + "require_approval": "never", + } + ], + stream=True, + mock_response="Final answer", + mock_tool_calls=[ + { + "id": "call-1", + "type": "function", + "function": {"name": "local_search", "arguments": "{}"}, + } + ], + ) + + import asyncio + assert asyncio.iscoroutine(response) + result = await response + + assert isinstance(result, CustomStreamWrapper) + + # Verify _hidden_params contains mcp_metadata + assert hasattr(result, "_hidden_params") + assert "mcp_metadata" in result._hidden_params + mcp_metadata = result._hidden_params["mcp_metadata"] + assert "mcp_list_tools" in mcp_metadata + assert "mcp_tool_calls" in mcp_metadata + assert "mcp_call_results" in mcp_metadata + + # Consume the stream and check final chunk + all_chunks = list(result) + assert len(all_chunks) > 0 + + # Find the final chunk (with finish_reason) + final_chunk = None + for chunk in all_chunks: + if hasattr(chunk, "choices") and chunk.choices: + choice = chunk.choices[0] + if hasattr(choice, "finish_reason") and choice.finish_reason: + final_chunk = chunk + break + + # If no chunk with finish_reason, use the last chunk + if final_chunk is None and all_chunks: + final_chunk = all_chunks[-1] + + assert final_chunk is not None, "Should have a final chunk" + + # Verify MCP metadata is in the final chunk's delta.provider_specific_fields + if hasattr(final_chunk, "choices") and final_chunk.choices: + choice = final_chunk.choices[0] + if hasattr(choice, "delta") and choice.delta: + provider_fields = getattr(choice.delta, "provider_specific_fields", None) + assert provider_fields is not None, "Final chunk should have provider_specific_fields" + assert "mcp_list_tools" in provider_fields, "Should have mcp_list_tools" + assert "mcp_tool_calls" in provider_fields, "Should have mcp_tool_calls" + assert "mcp_call_results" in provider_fields, "Should have mcp_call_results" diff --git a/tests/pass_through_unit_tests/messages_api_structured_output/__init__.py b/tests/pass_through_unit_tests/messages_api_structured_output/__init__.py new file mode 100644 index 0000000000..88c85d408a --- /dev/null +++ b/tests/pass_through_unit_tests/messages_api_structured_output/__init__.py @@ -0,0 +1,12 @@ +""" +Anthropic Messages API Structured Outputs Test Suite + +E2E tests for structured outputs functionality across different providers: +- Direct Anthropic API +- Azure AI Foundry Anthropic models +- AWS Bedrock Invoke API +- AWS Bedrock Converse API + +All tests validate that the output_format parameter works correctly +and returns valid JSON instead of Markdown text. +""" \ No newline at end of file diff --git a/tests/pass_through_unit_tests/messages_api_structured_output/base_anthropic_messages_structured_output_test.py b/tests/pass_through_unit_tests/messages_api_structured_output/base_anthropic_messages_structured_output_test.py new file mode 100644 index 0000000000..b0a8cf8b96 --- /dev/null +++ b/tests/pass_through_unit_tests/messages_api_structured_output/base_anthropic_messages_structured_output_test.py @@ -0,0 +1,138 @@ +""" +Base test class for Anthropic Messages API structured outputs E2E tests. + +Tests that structured outputs work correctly via litellm.anthropic.messages interface +by making actual API calls and validating JSON response format. +""" + +import json +import os +import sys +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional + +sys.path.insert(0, os.path.abspath("../../..")) + +import pytest +import litellm + + +class BaseAnthropicMessagesStructuredOutputTest(ABC): + """ + Base test class for structured outputs E2E tests across different providers. + + Subclasses must implement: + - get_model(): Returns the model string to use for tests + + Subclasses may optionally implement: + - get_api_base(): Returns the API base URL (for Azure, etc.) + - get_api_key(): Returns the API key (for Azure, etc.) + """ + + @abstractmethod + def get_model(self) -> str: + """ + Returns the model string to use for tests. + """ + pass + + def get_api_base(self) -> Optional[str]: + """ + Returns the API base URL. Override for providers like Azure. + """ + return None + + def get_api_key(self) -> Optional[str]: + """ + Returns the API key. Override for providers like Azure. + """ + return None + + def get_output_format_schema(self) -> Dict[str, Any]: + """ + Returns a simple JSON schema for testing structured outputs. + """ + return { + "type": "json_schema", + "schema": { + "type": "object", + "properties": { + "sentiment": { + "type": "string", + "enum": ["positive", "negative", "neutral"] + } + }, + "required": ["sentiment"], + "additionalProperties": False + } + } + + def get_test_messages(self) -> List[Dict[str, Any]]: + """ + Returns test messages for structured output testing. + """ + return [ + { + "role": "user", + "content": "What is the sentiment of this text: 'This product is amazing!' Return only the sentiment." + } + ] + + @pytest.mark.asyncio + async def test_structured_output_e2e(self): + """ + E2E test: Make actual API call with structured output and validate JSON response. + """ + litellm._turn_on_debug() + messages = self.get_test_messages() + output_format = self.get_output_format_schema() + + # Build kwargs with optional api_base and api_key + kwargs: Dict[str, Any] = { + "model": self.get_model(), + "messages": messages, + "max_tokens": 100, + "output_format": output_format, + } + + api_base = self.get_api_base() + if api_base: + kwargs["api_base"] = api_base + + api_key = self.get_api_key() + if api_key: + kwargs["api_key"] = api_key + + response = await litellm.anthropic.messages.acreate(**kwargs) + + print(f"Response: {response}") + + # Validate response structure - handle both dict and object responses + if isinstance(response, dict): + assert "content" in response + content_list = response["content"] + else: + assert hasattr(response, "content") + content_list = response.content + + assert len(content_list) > 0 + + content = content_list[0] + + # Handle both dict and object content blocks + if isinstance(content, dict): + assert "text" in content + response_text = content["text"] + else: + assert hasattr(content, "text") + response_text = content.text + + print(f"Response text: {response_text}") + + # The response should be valid JSON + parsed_json = json.loads(response_text) + print(f"Parsed JSON: {parsed_json}") + + # Validate the JSON structure + assert "sentiment" in parsed_json + assert parsed_json["sentiment"] in ["positive", "negative", "neutral"] \ No newline at end of file diff --git a/tests/pass_through_unit_tests/messages_api_structured_output/test_anthropic_api_structured_output.py b/tests/pass_through_unit_tests/messages_api_structured_output/test_anthropic_api_structured_output.py new file mode 100644 index 0000000000..c67c60b49f --- /dev/null +++ b/tests/pass_through_unit_tests/messages_api_structured_output/test_anthropic_api_structured_output.py @@ -0,0 +1,29 @@ +""" +E2E Test suite for Anthropic API structured outputs via litellm.anthropic.messages. + +Tests that structured outputs work correctly with direct Anthropic API calls +by making actual API calls and validating JSON response format. + +Requires ANTHROPIC_API_KEY environment variable. +""" + +import os +import sys + +sys.path.insert(0, os.path.abspath("../../../..")) + +from .base_anthropic_messages_structured_output_test import ( + BaseAnthropicMessagesStructuredOutputTest, +) + + +class TestAnthropicAPIStructuredOutput(BaseAnthropicMessagesStructuredOutputTest): + """ + E2E tests for structured outputs with direct Anthropic API. + + Uses Claude Sonnet 4.5 which supports structured outputs with the + 'anthropic-beta: structured-outputs-2025-11-13' header. + """ + + def get_model(self) -> str: + return "claude-sonnet-4-5-20250929" \ No newline at end of file diff --git a/tests/pass_through_unit_tests/messages_api_structured_output/test_azure_anthropic_structured_output.py b/tests/pass_through_unit_tests/messages_api_structured_output/test_azure_anthropic_structured_output.py new file mode 100644 index 0000000000..da46016b35 --- /dev/null +++ b/tests/pass_through_unit_tests/messages_api_structured_output/test_azure_anthropic_structured_output.py @@ -0,0 +1,36 @@ +""" +E2E Test suite for Azure Anthropic structured outputs via litellm.anthropic.messages. + +Tests that structured outputs work correctly with Azure AI Foundry Anthropic models +by making actual API calls and validating JSON response format. + +Requires Azure AI credentials and model deployment. +""" + +import os +import sys +from typing import Optional + +sys.path.insert(0, os.path.abspath("../../../..")) + +from .base_anthropic_messages_structured_output_test import ( + BaseAnthropicMessagesStructuredOutputTest, +) + + +class TestAzureAnthropicStructuredOutput(BaseAnthropicMessagesStructuredOutputTest): + """ + E2E tests for structured outputs with Azure AI Foundry Anthropic models. + + Uses the azure_ai/ prefix which routes through Azure AI Foundry + while maintaining the Anthropic Messages API format. + """ + + def get_model(self) -> str: + return "azure_ai/claude-opus-4-5" + + def get_api_base(self) -> Optional[str]: + return "https://krish-mh44t553-eastus2.services.ai.azure.com/" + + def get_api_key(self) -> Optional[str]: + return os.environ.get("AZURE_ANTHROPIC_API_KEY") \ No newline at end of file diff --git a/tests/pass_through_unit_tests/messages_api_structured_output/test_bedrock_converse_structured_output.py b/tests/pass_through_unit_tests/messages_api_structured_output/test_bedrock_converse_structured_output.py new file mode 100644 index 0000000000..9229677f32 --- /dev/null +++ b/tests/pass_through_unit_tests/messages_api_structured_output/test_bedrock_converse_structured_output.py @@ -0,0 +1,29 @@ +""" +E2E Test suite for Bedrock Converse API structured outputs via litellm.anthropic.messages. + +Tests that structured outputs work correctly with Bedrock Converse API +by making actual API calls and validating JSON response format. + +Requires AWS credentials and Bedrock model access. +""" + +import os +import sys + +sys.path.insert(0, os.path.abspath("../../../..")) + +from .base_anthropic_messages_structured_output_test import ( + BaseAnthropicMessagesStructuredOutputTest, +) + + +class TestBedrockConverseStructuredOutput(BaseAnthropicMessagesStructuredOutputTest): + """ + E2E tests for structured outputs with Bedrock Converse API. + + Uses the bedrock/converse/ prefix which routes through litellm.completion() + and the AmazonConverseConfig transformation. + """ + + def get_model(self) -> str: + return "bedrock/converse/us.anthropic.claude-3-5-sonnet-20241022-v2:0" \ No newline at end of file diff --git a/tests/pass_through_unit_tests/messages_api_structured_output/test_bedrock_invoke_structured_output.py b/tests/pass_through_unit_tests/messages_api_structured_output/test_bedrock_invoke_structured_output.py new file mode 100644 index 0000000000..d41072c46c --- /dev/null +++ b/tests/pass_through_unit_tests/messages_api_structured_output/test_bedrock_invoke_structured_output.py @@ -0,0 +1,32 @@ +""" +E2E Test suite for Bedrock Invoke API structured outputs via litellm.anthropic.messages. + +Tests that structured outputs work correctly with Bedrock Invoke API (native Anthropic format) +by making actual API calls and validating JSON response format. + +Requires AWS credentials and Bedrock model access. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from .base_anthropic_messages_structured_output_test import ( + BaseAnthropicMessagesStructuredOutputTest, +) + + +@pytest.mark.skip(reason="Skipping Bedrock Invoke structured output tests") +class TestBedrockInvokeStructuredOutput(BaseAnthropicMessagesStructuredOutputTest): + """ + E2E tests for structured outputs with Bedrock Invoke API. + + Uses the bedrock/invoke/ prefix which routes through the native + Anthropic Messages API format on Bedrock. + """ + + def get_model(self) -> str: + return "bedrock/invoke/us.anthropic.claude-3-5-sonnet-20241022-v2:0" \ No newline at end of file diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index 8e1057bb8e..2f161162da 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/tests/proxy_unit_tests/test_proxy_token_counter.py @@ -478,19 +478,20 @@ async def test_anthropic_endpoint_error_handling(): @pytest.mark.asyncio async def test_factory_anthropic_endpoint_calls_anthropic_counter(): """Test that /v1/messages/count_tokens with Anthropic model uses Anthropic counter.""" - from unittest.mock import patch, AsyncMock + from unittest.mock import patch, AsyncMock, MagicMock from fastapi.testclient import TestClient from litellm.proxy.proxy_server import app - # Mock the anthropic token counting function - with patch( - "litellm.proxy.utils.count_tokens_with_anthropic_api" - ) as mock_anthropic_count: - mock_anthropic_count.return_value = { - "total_tokens": 42, - "tokenizer_used": "anthropic", - } + # Mock the global handler instance in token_counter module + mock_handler = MagicMock() + mock_handler.handle_count_tokens_request = AsyncMock( + return_value={"input_tokens": 42} + ) + with patch( + "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", + mock_handler + ): # Mock router to return Anthropic deployment with patch("litellm.proxy.proxy_server.llm_router") as mock_router: mock_router.model_list = [ @@ -510,36 +511,44 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter(): } ) - client = TestClient(app) + # Set ANTHROPIC_API_KEY for the test + with patch.dict("os.environ", {"ANTHROPIC_API_KEY": "test-key"}): + client = TestClient(app) - response = client.post( - "/v1/messages/count_tokens", - json={ - "model": "claude-3-5-sonnet", - "messages": [{"role": "user", "content": "Hello"}], - }, - headers={"Authorization": "Bearer test-key"}, - ) + response = client.post( + "/v1/messages/count_tokens", + json={ + "model": "claude-3-5-sonnet", + "messages": [{"role": "user", "content": "Hello"}], + }, + headers={"Authorization": "Bearer test-key"}, + ) - assert response.status_code == 200 - data = response.json() - assert data["input_tokens"] == 42 + assert response.status_code == 200 + data = response.json() + assert data["input_tokens"] == 42 - # Verify that Anthropic API was called - mock_anthropic_count.assert_called_once() + # Verify that Anthropic handler was called + mock_handler.handle_count_tokens_request.assert_called_once() @pytest.mark.asyncio async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): """Test that /v1/messages/count_tokens with GPT-4 does NOT use Anthropic counter.""" - from unittest.mock import patch, AsyncMock + from unittest.mock import patch, AsyncMock, MagicMock from fastapi.testclient import TestClient from litellm.proxy.proxy_server import app - # Mock the anthropic token counting function + # Mock the global handler instance in token_counter module + mock_handler = MagicMock() + mock_handler.handle_count_tokens_request = AsyncMock( + return_value={"input_tokens": 42} + ) + with patch( - "litellm.proxy.utils.count_tokens_with_anthropic_api" - ) as mock_anthropic_count: + "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", + mock_handler + ): # Mock litellm token counter with patch("litellm.token_counter") as mock_litellm_counter: mock_litellm_counter.return_value = 50 @@ -578,21 +587,27 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): data = response.json() assert data["input_tokens"] == 50 - # Verify that Anthropic API was NOT called - mock_anthropic_count.assert_not_called() + # Verify that Anthropic handler was NOT called + mock_handler.handle_count_tokens_request.assert_not_called() @pytest.mark.asyncio async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): """Test that /utils/token_counter does NOT use Anthropic counter even with Anthropic model.""" - from unittest.mock import patch, AsyncMock + from unittest.mock import patch, AsyncMock, MagicMock from fastapi.testclient import TestClient from litellm.proxy.proxy_server import app - # Mock the anthropic token counting function + # Mock the global handler instance in token_counter module + mock_handler = MagicMock() + mock_handler.handle_count_tokens_request = AsyncMock( + return_value={"input_tokens": 42} + ) + with patch( - "litellm.proxy.utils.count_tokens_with_anthropic_api" - ) as mock_anthropic_count: + "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", + mock_handler + ): # Mock litellm token counter with patch("litellm.token_counter") as mock_litellm_counter: mock_litellm_counter.return_value = 35 @@ -635,8 +650,8 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): data = response.json() assert data["total_tokens"] == 35 - # Verify that Anthropic API was NOT called (since call_endpoint=False) - mock_anthropic_count.assert_not_called() + # Verify that Anthropic handler was NOT called (since call_endpoint=False) + mock_handler.handle_count_tokens_request.assert_not_called() @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/websearch_interception/test_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py similarity index 100% rename from tests/test_litellm/integrations/websearch_interception/test_handler.py rename to tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py new file mode 100644 index 0000000000..0a397d116e --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -0,0 +1,84 @@ +""" +Tests for Anthropic OAuth token handling for Claude Code Max integration. +""" + +import os +import sys + +# Add litellm to path +sys.path.insert(0, os.path.abspath("../../../../..")) + +# Fake OAuth token for testing (not a real secret) +FAKE_OAUTH_TOKEN = "sk-ant-oat01-fake-token-for-testing-123456789abcdef" + + +def test_oauth_detection_in_common_utils(): + """Test 1: OAuth token detection in common_utils""" + from litellm.llms.anthropic.common_utils import optionally_handle_anthropic_oauth + + headers = {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"} + updated_headers, extracted_api_key = optionally_handle_anthropic_oauth(headers, None) + + assert extracted_api_key == FAKE_OAUTH_TOKEN + assert updated_headers["anthropic-beta"] == "oauth-2025-04-20" + assert updated_headers["anthropic-dangerous-direct-browser-access"] == "true" + + +def test_oauth_integration_in_validate_environment(): + """Test 2: OAuth integration in AnthropicConfig validate_environment""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"} + + updated_headers = config.validate_environment( + headers=headers, + model="claude-3-haiku-20240307", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + + assert updated_headers["x-api-key"] == FAKE_OAUTH_TOKEN + assert updated_headers["anthropic-dangerous-direct-browser-access"] == "true" + + +def test_oauth_detection_in_messages_transformation(): + """Test 3: OAuth detection in messages transformation""" + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) + + config = AnthropicMessagesConfig() + headers = {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"} + + updated_headers, _ = config.validate_anthropic_messages_environment( + headers=headers, + model="claude-3-haiku-20240307", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + + assert updated_headers["x-api-key"] == FAKE_OAUTH_TOKEN + assert "oauth-2025-04-20" in updated_headers["anthropic-beta"] + assert updated_headers["anthropic-dangerous-direct-browser-access"] == "true" + + +def test_regular_api_keys_still_work(): + """Test 4: Regular API keys still work (regression test)""" + from litellm.llms.anthropic.common_utils import optionally_handle_anthropic_oauth + + regular_key = "sk-ant-api03-regular-key-123" + headers = {"authorization": f"Bearer {regular_key}"} + + updated_headers, extracted_api_key = optionally_handle_anthropic_oauth(headers, regular_key) + + # Regular key should be unchanged + assert extracted_api_key == regular_key + # OAuth headers should NOT be added + assert "anthropic-dangerous-direct-browser-access" not in updated_headers \ No newline at end of file diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_structured_output.py b/tests/test_litellm/llms/anthropic/test_anthropic_structured_output.py new file mode 100644 index 0000000000..7f6b68881d --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_anthropic_structured_output.py @@ -0,0 +1,169 @@ +""" +Test Anthropic structured output with Pydantic models. + +This test file verifies that Pydantic models with various constraints +are properly converted to Anthropic-compatible JSON schemas. +""" + +import pytest +from pydantic import BaseModel, Field +from typing import List + + +class TestAnthropicStructuredOutput: + """Test Anthropic structured output schema transformations.""" + + def test_max_length_on_list_field_filtered(self): + """ + Test that max_length on List fields is filtered out for Anthropic models. + + Anthropic doesn't support 'maxItems' property for array types in their + output_format.schema, so we need to filter it out. + + Related issue: https://github.com/BerriAI/litellm/issues/19444 + """ + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + # Define a Pydantic model with max_length on a List field + class ResponseModel(BaseModel): + items: List[str] = Field(max_length=5, description="List of items") + name: str = Field(description="Name field") + + config = AnthropicConfig() + + # Get the JSON schema from the Pydantic model + json_schema = config.get_json_schema_from_pydantic_object(ResponseModel) + + # Extract the actual schema + schema = json_schema["json_schema"]["schema"] + + # Verify that maxItems is present in the raw schema (from Pydantic) + assert "maxItems" in schema["properties"]["items"] + + # Now apply the Anthropic output format transformation + response_format = { + "type": "json_schema", + "json_schema": json_schema["json_schema"] + } + + output_format = config.map_response_format_to_anthropic_output_format( + response_format + ) + + # Verify that maxItems is filtered out for Anthropic + assert output_format is not None + assert "schema" in output_format + transformed_schema = output_format["schema"] + + # maxItems should be removed from the items property + assert "maxItems" not in transformed_schema["properties"]["items"] + + # But other properties should remain + assert "type" in transformed_schema["properties"]["items"] + assert transformed_schema["properties"]["items"]["type"] == "array" + assert "description" in transformed_schema["properties"]["items"] + + def test_min_length_on_list_field_filtered(self): + """ + Test that min_length on List fields is filtered out for Anthropic models. + + Anthropic likely doesn't support 'minItems' either. + """ + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + class ResponseModel(BaseModel): + items: List[str] = Field(min_length=2, description="List of items") + + config = AnthropicConfig() + json_schema = config.get_json_schema_from_pydantic_object(ResponseModel) + + response_format = { + "type": "json_schema", + "json_schema": json_schema["json_schema"] + } + + output_format = config.map_response_format_to_anthropic_output_format( + response_format + ) + + assert output_format is not None + transformed_schema = output_format["schema"] + + # minItems should be removed + assert "minItems" not in transformed_schema["properties"]["items"] + + def test_nested_array_constraints_filtered(self): + """ + Test that array constraints are filtered at all nesting levels. + """ + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + class NestedItem(BaseModel): + tags: List[str] = Field(max_length=3) + + class ResponseModel(BaseModel): + items: List[NestedItem] = Field(max_length=5) + + config = AnthropicConfig() + json_schema = config.get_json_schema_from_pydantic_object(ResponseModel) + + response_format = { + "type": "json_schema", + "json_schema": json_schema["json_schema"] + } + + output_format = config.map_response_format_to_anthropic_output_format( + response_format + ) + + assert output_format is not None + transformed_schema = output_format["schema"] + + # Top-level maxItems should be removed + assert "maxItems" not in transformed_schema["properties"]["items"] + + # Nested maxItems should also be removed + if "$defs" in transformed_schema: + nested_item_schema = transformed_schema["$defs"].get("NestedItem", {}) + if "properties" in nested_item_schema and "tags" in nested_item_schema["properties"]: + assert "maxItems" not in nested_item_schema["properties"]["tags"] + + def test_other_constraints_preserved(self): + """ + Test that other valid constraints are preserved. + """ + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + class ResponseModel(BaseModel): + name: str = Field(max_length=100, min_length=1, description="Name") + age: int = Field(ge=0, le=150, description="Age") + items: List[str] = Field(description="Items list") + + config = AnthropicConfig() + json_schema = config.get_json_schema_from_pydantic_object(ResponseModel) + + response_format = { + "type": "json_schema", + "json_schema": json_schema["json_schema"] + } + + output_format = config.map_response_format_to_anthropic_output_format( + response_format + ) + + assert output_format is not None + transformed_schema = output_format["schema"] + + # String constraints should be preserved + name_schema = transformed_schema["properties"]["name"] + assert "maxLength" in name_schema + assert "minLength" in name_schema + assert name_schema["maxLength"] == 100 + assert name_schema["minLength"] == 1 + + # Number constraints should be preserved + age_schema = transformed_schema["properties"]["age"] + assert "minimum" in age_schema + assert "maximum" in age_schema + assert age_schema["minimum"] == 0 + assert age_schema["maximum"] == 150 diff --git a/tests/test_litellm/llms/huggingface/embedding/test_handler.py b/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py similarity index 100% rename from tests/test_litellm/llms/huggingface/embedding/test_handler.py rename to tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py similarity index 100% rename from tests/test_litellm/llms/openai/chat/guardrail_translation/test_handler.py rename to tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 33fd11dfbe..a9c27e3093 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -735,13 +735,13 @@ def test_file_data_field_order(): Related issue: Gemini API returns 400 INVALID_ARGUMENT when fields are in wrong order. """ import json - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_image + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media # Test with HTTPS URL and explicit format (audio file) file_url = "https://generativelanguage.googleapis.com/v1beta/files/test123" format = "audio/mpeg" - result = _process_gemini_image(image_url=file_url, format=format) + result = _process_gemini_media(image_url=file_url, format=format) # Verify the result has file_data assert "file_data" in result @@ -770,12 +770,12 @@ def test_file_data_field_order(): def test_file_data_field_order_gcs_urls(): """Test that GCS URLs also maintain correct field order.""" import json - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_image + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media # Test with GCS URL gcs_url = "gs://bucket/audio.mp3" - result = _process_gemini_image(image_url=gcs_url) + result = _process_gemini_media(image_url=gcs_url) # Verify the result has file_data assert "file_data" in result diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 969199f6ad..5be080b53f 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1980,6 +1980,71 @@ def test_reasoning_effort_maps_to_thinking_level_gemini_3(): assert result["thinkingConfig"]["includeThoughts"] is False +def test_reasoning_effort_dict_format_gemini_3(): + """ + Test that reasoning_effort works when passed as dict format from OpenAI Agents SDK. + + The OpenAI Agents SDK passes reasoning_effort as {"effort": "high", "summary": "auto"} + instead of just a string. This test verifies that we correctly extract the effort value. + + Related issue: https://github.com/BerriAI/litellm/issues/19411 + """ + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + v = VertexGeminiConfig() + model = "gemini-3-pro-preview" + + # Test dict format with effort="high" (OpenAI Agents SDK format) + optional_params = {} + non_default_params = {"reasoning_effort": {"effort": "high", "summary": "auto"}} + result = v.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + assert result["thinkingConfig"]["thinkingLevel"] == "high" + assert result["thinkingConfig"]["includeThoughts"] is True + + # Test dict format with effort="low" + optional_params = {} + non_default_params = {"reasoning_effort": {"effort": "low"}} + result = v.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + assert result["thinkingConfig"]["thinkingLevel"] == "low" + assert result["thinkingConfig"]["includeThoughts"] is True + + # Test dict format with effort="medium" + optional_params = {} + non_default_params = {"reasoning_effort": {"effort": "medium"}} + result = v.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + assert result["thinkingConfig"]["thinkingLevel"] == "high" + assert result["thinkingConfig"]["includeThoughts"] is True + + # Test dict format without effort key - should fall back to Gemini 3 default (low) + optional_params = {} + non_default_params = {"reasoning_effort": {"summary": "auto"}} + result = v.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + # Gemini 3 defaults to thinkingLevel="low" when no explicit effort is set + assert result["thinkingConfig"]["thinkingLevel"] == "low" + + def test_temperature_default_for_gemini_3(): """Test that temperature defaults to 1.0 for Gemini 3+ models when not specified""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -2746,3 +2811,273 @@ def test_gemini_image_gen_usage_metadata_prompt_vs_completion_separation(): # candidatesTokenCount (1290) - image_tokens (1290) = 0 assert result.completion_tokens_details.text_tokens == 0, \ "Completion text tokens should be 0 (image-only response)" + + +def test_file_object_detail_parameter(): + """Test that detail parameter works for type: file objects (Issue #19026)""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What's in this video?"}, + { + "type": "file", + "file": { + "file_id": "https://example.com/video.mp4", + "format": "video/mp4", + "detail": "low" + } + } + ] + } + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + # Verify media_resolution is set for file objects + assert len(contents) == 1 + assert len(contents[0]["parts"]) == 2 # text + file + + # Find the file part + file_part = None + for part in contents[0]["parts"]: + if "file_data" in part: + file_part = part + break + + assert file_part is not None, "File part should exist" + assert "media_resolution" in file_part, "media_resolution should be set for file objects" + assert file_part["media_resolution"] == {"level": "MEDIA_RESOLUTION_LOW"} + + +def test_video_metadata_fps(): + """Test fps parameter in video_metadata (Issue #19026)""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze this video"}, + { + "type": "file", + "file": { + "file_id": "gs://bucket/video.mp4", + "format": "video/mp4", + "video_metadata": {"fps": 5} + } + } + ] + } + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + # Find the file part + file_part = None + for part in contents[0]["parts"]: + if "file_data" in part: + file_part = part + break + + assert file_part is not None + assert "video_metadata" in file_part, "video_metadata should be present" + assert file_part["video_metadata"]["fps"] == 5 + + +def test_video_metadata_complete(): + """Test all video_metadata fields: fps, start_offset, end_offset (Issue #19026)""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze this video clip"}, + { + "type": "file", + "file": { + "file_id": "gs://bucket/video.mp4", + "format": "video/mp4", + "video_metadata": { + "start_offset": "10s", + "end_offset": "60s", + "fps": 5 + } + } + } + ] + } + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + # Find the file part + file_part = None + for part in contents[0]["parts"]: + if "file_data" in part: + file_part = part + break + + assert file_part is not None + assert "video_metadata" in file_part + + # Verify field name conversion: snake_case -> camelCase + vm = file_part["video_metadata"] + assert vm["startOffset"] == "10s", "start_offset should be converted to startOffset" + assert vm["endOffset"] == "60s", "end_offset should be converted to endOffset" + assert vm["fps"] == 5, "fps should remain unchanged" + + +def test_detail_and_video_metadata_combined(): + """Test using both detail and video_metadata together (Issue #19026)""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Analyze video"}, + { + "type": "file", + "file": { + "file_id": "https://example.com/video.mp4", + "format": "video/mp4", + "detail": "high", + "video_metadata": {"fps": 10} + } + } + ] + } + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + # Find the file part + file_part = None + for part in contents[0]["parts"]: + if "file_data" in part: + file_part = part + break + + assert file_part is not None + assert "media_resolution" in file_part + assert file_part["media_resolution"] == {"level": "MEDIA_RESOLUTION_HIGH"} + assert "video_metadata" in file_part + assert file_part["video_metadata"]["fps"] == 10 + + +def test_new_detail_levels(): + """Test new detail levels: medium and ultra_high (Issue #19026)""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _convert_detail_to_media_resolution_enum, + _gemini_convert_messages_with_history, + ) + + # Test mapping function + assert _convert_detail_to_media_resolution_enum("low") == {"level": "MEDIA_RESOLUTION_LOW"} + assert _convert_detail_to_media_resolution_enum("medium") == {"level": "MEDIA_RESOLUTION_MEDIUM"} + assert _convert_detail_to_media_resolution_enum("high") == {"level": "MEDIA_RESOLUTION_HIGH"} + assert _convert_detail_to_media_resolution_enum("ultra_high") == {"level": "MEDIA_RESOLUTION_ULTRA_HIGH"} + + # Test with actual message transformation + messages = [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_id": "https://example.com/video.mp4", + "format": "video/mp4", + "detail": "medium" + } + } + ] + } + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + file_part = None + for part in contents[0]["parts"]: + if "file_data" in part: + file_part = part + break + + assert file_part is not None + assert file_part["media_resolution"] == {"level": "MEDIA_RESOLUTION_MEDIUM"} + + +def test_video_metadata_only_for_gemini_3(): + """Test that video_metadata is only applied for Gemini 3+ models (Issue #19026)""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_id": "https://example.com/video.mp4", + "format": "video/mp4", + "detail": "high", + "video_metadata": {"fps": 5} + } + } + ] + } + ] + + # Test with Gemini 1.5 (should not have video_metadata or media_resolution) + contents_1_5 = _gemini_convert_messages_with_history( + messages=messages, model="gemini-1.5-pro" + ) + + file_part_1_5 = None + for part in contents_1_5[0]["parts"]: + if "file_data" in part: + file_part_1_5 = part + break + + assert file_part_1_5 is not None + assert "media_resolution" not in file_part_1_5, "Gemini 1.5 should not have media_resolution" + assert "video_metadata" not in file_part_1_5, "Gemini 1.5 should not have video_metadata" + + # Test with Gemini 3 (should have both) + contents_3 = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + file_part_3 = None + for part in contents_3[0]["parts"]: + if "file_data" in part: + file_part_3 = part + break + + assert file_part_3 is not None + assert "media_resolution" in file_part_3, "Gemini 3 should have media_resolution" + assert "video_metadata" in file_part_3, "Gemini 3 should have video_metadata" diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/test_litellm/llms/vertex_ai/test_vertex.py index 39ed09f81b..fdba86af4a 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex.py @@ -19,7 +19,7 @@ import pytest import litellm from litellm import get_optional_params -from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_image +from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media from litellm.types.llms.vertex_ai import BlobType @@ -1191,46 +1191,46 @@ def test_logprobs(): assert resp.choices[0].logprobs is not None -def test_process_gemini_image(): - """Test the _process_gemini_image function for different image sources""" - from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_image +def test_process_gemini_media(): + """Test the _process_gemini_media function for different image sources""" + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media from litellm.types.llms.vertex_ai import FileDataType # Test GCS URI - gcs_result = _process_gemini_image("gs://bucket/image.png") + gcs_result = _process_gemini_media("gs://bucket/image.png") assert gcs_result["file_data"] == FileDataType( mime_type="image/png", file_uri="gs://bucket/image.png" ) # Test gs url with format specified - gcs_result = _process_gemini_image("gs://bucket/image", format="image/jpeg") + gcs_result = _process_gemini_media("gs://bucket/image", format="image/jpeg") assert gcs_result["file_data"] == FileDataType( mime_type="image/jpeg", file_uri="gs://bucket/image" ) # Test HTTPS JPG URL - https_result = _process_gemini_image("https://example.com/image.jpg") + https_result = _process_gemini_media("https://example.com/image.jpg") print("https_result JPG", https_result) assert https_result["file_data"] == FileDataType( mime_type="image/jpeg", file_uri="https://example.com/image.jpg" ) # Test HTTPS PNG URL - https_result = _process_gemini_image("https://example.com/image.png") + https_result = _process_gemini_media("https://example.com/image.png") print("https_result PNG", https_result) assert https_result["file_data"] == FileDataType( mime_type="image/png", file_uri="https://example.com/image.png" ) # Test HTTPS VIDEO URL - https_result = _process_gemini_image("https://cloud-samples-data/video/animals.mp4") + https_result = _process_gemini_media("https://cloud-samples-data/video/animals.mp4") print("https_result PNG", https_result) assert https_result["file_data"] == FileDataType( mime_type="video/mp4", file_uri="https://cloud-samples-data/video/animals.mp4" ) # Test HTTPS PDF URL - https_result = _process_gemini_image("https://cloud-samples-data/pdf/animals.pdf") + https_result = _process_gemini_media("https://cloud-samples-data/pdf/animals.pdf") print("https_result PDF", https_result) assert https_result["file_data"] == FileDataType( mime_type="application/pdf", @@ -1239,7 +1239,7 @@ def test_process_gemini_image(): # Test base64 image base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRg..." - base64_result = _process_gemini_image(base64_image) + base64_result = _process_gemini_media(base64_image) print("base64_result", base64_result) assert base64_result["inline_data"]["mime_type"] == "image/jpeg" assert base64_result["inline_data"]["data"] == "/9j/4AAQSkZJRg..." @@ -1368,11 +1368,11 @@ def mock_blob(): "http://subdomain.domain.com/path/to/image.png", ], ) -def test_process_gemini_image_http_url( +def test_process_gemini_media_http_url( http_url: str, mock_convert_url_to_base64: Mock, mock_blob: Mock ) -> None: """ - Test that _process_gemini_image correctly handles HTTP URLs. + Test that _process_gemini_media correctly handles HTTP URLs. Args: http_url: Test HTTP URL @@ -1384,7 +1384,7 @@ def test_process_gemini_image_http_url( expected_image_data = "data:image/jpeg;base64,/9j/4AAQSkZJRg..." mock_convert_url_to_base64.return_value = expected_image_data # Act - result = _process_gemini_image(http_url) + result = _process_gemini_media(http_url) # assert result["file_data"]["file_uri"] == http_url diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 7aa176aaac..8c29cf6a59 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1817,3 +1817,138 @@ class TestMCPServerManagerReload: mock_get_all.assert_awaited_once() mock_build.assert_awaited_once_with(db_row) assert manager.registry["server-1"] is rebuilt_server + + +@pytest.mark.asyncio +async def test_call_mcp_tool_logs_failure_via_post_call_failure_hook(): + """ + Regression test for 6267f168...: + Ensure proxy-side `call_mcp_tool` logs failures via `proxy_logging_obj.post_call_failure_hook`. + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + call_mcp_tool, + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport, UserAPIKeyAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + except ImportError: + pytest.skip("MCP server not available") + + 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"}, + ) + + proxy_logging_mock = MagicMock() + proxy_logging_mock.post_call_failure_hook = AsyncMock() + + user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + + with patch.object( + global_mcp_server_manager, + "get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[mock_server.server_id], + ), patch.object( + global_mcp_server_manager, + "get_mcp_server_by_id", + return_value=mock_server, + ), patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers_from_mcp_server_names", + new_callable=AsyncMock, + return_value=[mock_server], + ), patch( + "litellm.proxy._experimental.mcp_server.server.execute_mcp_tool", + new_callable=AsyncMock, + side_effect=Exception("boom"), + ), patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + proxy_logging_mock, + ): + with pytest.raises(Exception): + await call_mcp_tool( + name="test_server-any_tool", + arguments={"x": 1}, + user_api_key_auth=user_auth, + litellm_call_id="cid", + ) + + proxy_logging_mock.post_call_failure_hook.assert_awaited_once() + assert ( + proxy_logging_mock.post_call_failure_hook.await_args.kwargs.get("route") + == "/mcp/call_tool" + ) + + +@pytest.mark.asyncio +async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enabled(): + """ + Regression test for 872e5b98...: + Ensure list-tools logging path calls `async_success_handler` when enabled. + """ + try: + from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers + from litellm.proxy._types import UserAPIKeyAuth + except ImportError: + pytest.skip("MCP server not available") + + user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + + server_a = MagicMock(name="server_a_obj") + server_a.name = "server_a" + server_a.alias = "server_a" + server_a.server_name = "server_a" + server_a.server_id = "a" + server_a.auth_type = None + server_a.extra_headers = None + + tool_1 = MagicMock() + tool_1.name = "server_a-tool_1" + + dummy_logging_obj = MagicMock() + dummy_logging_obj.model_call_details = {"metadata": {"spend_logs_metadata": {}}} + dummy_logging_obj.async_success_handler = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server_a]), + ), patch( + "litellm.proxy._experimental.mcp_server.server._prepare_mcp_server_headers", + return_value=(None, None), + ), patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + ) as mock_manager, patch( + "litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools", + side_effect=lambda tools, _server: tools, + ), patch( + "litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions", + new=AsyncMock(side_effect=lambda tools, **_: tools), + ), patch( + "litellm.proxy._experimental.mcp_server.server.function_setup", + return_value=(dummy_logging_obj, None), + ): + mock_manager._get_tools_from_server = AsyncMock(return_value=[tool_1]) + + tools = await _get_tools_from_mcp_servers( + user_api_key_auth=user_auth, + mcp_auth_header=None, + mcp_servers=["server_a"], + mcp_server_auth_headers=None, + log_list_tools_to_spendlogs=True, + list_tools_log_source="mcp_protocol", + ) + + assert tools == [tool_1] + dummy_logging_obj.async_success_handler.assert_awaited_once() + assert dummy_logging_obj.async_success_handler.await_args.kwargs["result"] == [tool_1] + + spend_meta = dummy_logging_obj.model_call_details["metadata"]["spend_logs_metadata"] + assert spend_meta["tool_count_total"] == 1 + assert spend_meta["allowed_server_count"] == 1 + assert spend_meta["per_server_tool_counts"]["server_a"] == 1 diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py new file mode 100644 index 0000000000..308c8cdbce --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_integration.py @@ -0,0 +1,144 @@ +""" +Regression test for AWS Secrets Manager Auto-Rotation Bug Fix + +This test verifies that KeyRotationManager correctly passes key_alias +when calling regenerate_key_fn, ensuring the secret is rotated at the +correct location in AWS Secrets Manager. + +Bug Fixed: Key alias was not passed during auto-rotation, causing +secrets to be created at a new location instead of updating in-place. +""" +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import ( + GenerateKeyResponse, + LiteLLM_VerificationToken, + RegenerateKeyRequest, +) +from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager + + +class TestKeyRotationManagerPassesKeyAlias: + """ + Regression tests to ensure KeyRotationManager passes key_alias + to regenerate_key_fn during auto-rotation. + """ + + @pytest.mark.asyncio + async def test_rotate_key_passes_key_alias_to_regenerate_request(self): + """ + Verify that _rotate_key includes key_alias in the RegenerateKeyRequest. + + This is the core fix: previously, key_alias was NOT passed, causing + the secret manager hook to use a generated name instead of the alias. + """ + # Create a mock key with an alias + test_alias = "tenant1/my-important-key" + test_token = "sk-test-token-hash-12345" + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.token = test_token + mock_key.key_alias = test_alias + mock_key.key_name = "sk-...1234" + mock_key.rotation_interval = "30d" + mock_key.rotation_count = 0 + + # Create mock prisma client + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.update = AsyncMock( + return_value=mock_key + ) + + # Create mock response + mock_response = GenerateKeyResponse( + key="sk-new-key-value", + token_id="new-token-hash", + key_alias=test_alias, + ) + + # Capture the RegenerateKeyRequest passed to regenerate_key_fn + captured_request = None + + async def capture_regenerate_key_fn( + data, user_api_key_dict, litellm_changed_by + ): + nonlocal captured_request + captured_request = data + return mock_response + + # Patch regenerate_key_fn to capture the request + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + side_effect=capture_regenerate_key_fn, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + rotation_manager = KeyRotationManager(mock_prisma) + await rotation_manager._rotate_key(mock_key) + + # CRITICAL ASSERTION: key_alias must be passed + assert captured_request is not None, "regenerate_key_fn should have been called" + assert isinstance(captured_request, RegenerateKeyRequest) + assert captured_request.key == test_token, "Token should be passed correctly" + assert captured_request.key_alias == test_alias, ( + f"key_alias should be '{test_alias}' but was '{captured_request.key_alias}'. " + "This is the bug we fixed - key_alias was not being passed!" + ) + + @pytest.mark.asyncio + async def test_rotate_key_passes_none_alias_when_key_has_no_alias(self): + """ + Verify that _rotate_key handles keys without an alias gracefully. + """ + test_token = "sk-test-token-hash-67890" + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.token = test_token + mock_key.key_alias = None # No alias set + mock_key.key_name = "sk-...5678" + mock_key.rotation_interval = "30d" + mock_key.rotation_count = 0 + + mock_prisma = MagicMock() + mock_prisma.db.litellm_verificationtoken.update = AsyncMock( + return_value=mock_key + ) + + mock_response = GenerateKeyResponse( + key="sk-new-key-value", + token_id="new-token-hash", + ) + + captured_request = None + + async def capture_regenerate_key_fn( + data, user_api_key_dict, litellm_changed_by + ): + nonlocal captured_request + captured_request = data + return mock_response + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + side_effect=capture_regenerate_key_fn, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + rotation_manager = KeyRotationManager(mock_prisma) + await rotation_manager._rotate_key(mock_key) + + assert captured_request is not None + assert captured_request.key == test_token + assert ( + captured_request.key_alias is None + ), "key_alias should be None for keys without alias" diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 33f2a75fac..397a6af556 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -12,6 +12,7 @@ sys.path.insert( from litellm.proxy._types import ( LiteLLM_UserTableFiltered, + LitellmUserRoles, NewUserRequest, ProxyException, UpdateUserRequest, @@ -306,6 +307,88 @@ async def test_new_user_license_over_limit(mocker): mock_license_check.is_over_limit.assert_called_once_with(total_users=1000) +@pytest.mark.asyncio +async def test_new_user_non_admin_cannot_create_admin(mocker): + """ + Test that non-admin users cannot create administrative users (PROXY_ADMIN or PROXY_ADMIN_VIEW_ONLY). + This prevents privilege escalation vulnerabilities. + """ + from litellm.proxy.management_endpoints.internal_user_endpoints import new_user + + # Mock the prisma client + mock_prisma_client = mocker.MagicMock() + + # Setup the mock count response (under license limit) + async def mock_count(*args, **kwargs): + return 5 # Low user count, under limit + + mock_prisma_client.db.litellm_usertable.count = mock_count + + # Mock duplicate checks to pass + async def mock_check_duplicate_user_email(*args, **kwargs): + return None # No duplicate found + + async def mock_check_duplicate_user_id(*args, **kwargs): + return None # No duplicate found + + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_email", + mock_check_duplicate_user_email, + ) + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_id", + mock_check_duplicate_user_id, + ) + + # Mock the license check to return False (under limit) + mock_license_check = mocker.MagicMock() + mock_license_check.is_over_limit.return_value = False + + # Patch the imports in the endpoint + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.proxy_server._license_check", mock_license_check) + + # Test Case 1: INTERNAL_USER trying to create PROXY_ADMIN + user_request = NewUserRequest( + user_email="admin@example.com", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + # Mock user_api_key_dict with non-admin role + mock_user_api_key_dict = UserAPIKeyAuth( + user_id="test_internal_user", user_role=LitellmUserRoles.INTERNAL_USER + ) + + # Call new_user function and expect ProxyException + with pytest.raises(ProxyException) as exc_info: + await new_user(data=user_request, user_api_key_dict=mock_user_api_key_dict) + + # Verify the exception details + assert exc_info.value.code == 403 or exc_info.value.code == "403" + assert "Only proxy admins can create administrative users" in str(exc_info.value.message) + assert "proxy_admin" in str(exc_info.value.message) + assert "proxy_admin_viewer" in str(exc_info.value.message) + assert str(LitellmUserRoles.PROXY_ADMIN) in str(exc_info.value.message) + assert str(LitellmUserRoles.INTERNAL_USER) in str(exc_info.value.message) + + # Test Case 2: INTERNAL_USER trying to create PROXY_ADMIN_VIEW_ONLY + user_request_viewer = NewUserRequest( + user_email="admin_viewer@example.com", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + ) + + with pytest.raises(ProxyException) as exc_info2: + await new_user( + data=user_request_viewer, user_api_key_dict=mock_user_api_key_dict + ) + + # Verify the exception details + assert exc_info2.value.code == 403 or exc_info2.value.code == "403" + assert "Only proxy admins can create administrative users" in str( + exc_info2.value.message + ) + assert str(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) in str(exc_info2.value.message) + + @pytest.mark.asyncio async def test_user_info_url_encoding_plus_character(mocker): """ diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py index ceb231eb4c..28b3ba0a17 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py @@ -1,9 +1,14 @@ +from unittest.mock import AsyncMock, MagicMock, patch + import pytest -from unittest.mock import MagicMock, AsyncMock, patch -from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _base_vertex_proxy_route + +from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _base_vertex_proxy_route, +) from litellm.types.router import DeploymentTypedDict + @pytest.mark.asyncio async def test_vertex_passthrough_load_balancing(): """ @@ -220,3 +225,225 @@ async def test_async_get_available_deployment_for_pass_through(): assert deployment is not None assert deployment["litellm_params"]["use_in_pass_through"] is True + +@pytest.mark.asyncio +async def test_vertex_passthrough_forwards_anthropic_beta_header(): + """ + Test that _prepare_vertex_auth_headers forwards the anthropic-beta header + (and other important headers) from the incoming request when credentials are available. + + This test validates the fix for the issue where the 1M context window header + (anthropic-beta: context-1m-2025-08-07) was being dropped when forwarding + requests to Vertex AI. + """ + from starlette.datastructures import Headers + + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _prepare_vertex_auth_headers, + ) + + # Create a mock request with anthropic-beta header + mock_request = MagicMock() + mock_request.headers = Headers({ + "authorization": "Bearer old-token", + "anthropic-beta": "context-1m-2025-08-07", + "content-type": "application/json", + "user-agent": "test-client", + "content-length": "1234", # Should be removed + "host": "localhost:4000", # Should be removed + }) + + # Create mock vertex credentials + mock_vertex_credentials = MagicMock() + mock_vertex_credentials.vertex_project = "test-project" + mock_vertex_credentials.vertex_location = "us-central1" + mock_vertex_credentials.vertex_credentials = "test-credentials" + + # Create mock handler + mock_handler = MagicMock() + mock_handler.update_base_target_url_with_credential_location.return_value = ( + "https://us-central1-aiplatform.googleapis.com" + ) + + with patch.object( + VertexBase, + "_ensure_access_token_async", + new_callable=AsyncMock, + return_value=("test-auth-header", "test-project"), + ) as mock_ensure_token, patch.object( + VertexBase, + "_get_token_and_url", + return_value=("new-access-token", None), + ) as mock_get_token: + + # Call the function + ( + headers, + base_target_url, + headers_passed_through, + vertex_project, + vertex_location, + ) = await _prepare_vertex_auth_headers( + request=mock_request, + vertex_credentials=mock_vertex_credentials, + router_credentials=None, + vertex_project="test-project", + vertex_location="us-central1", + base_target_url="https://us-central1-aiplatform.googleapis.com", + get_vertex_pass_through_handler=mock_handler, + ) + + # Verify that allowlisted headers are preserved + assert "anthropic-beta" in headers + assert headers["anthropic-beta"] == "context-1m-2025-08-07" + assert "content-type" in headers + assert headers["content-type"] == "application/json" + + # Verify that the Authorization header is set with vendor credentials + assert "Authorization" in headers + assert headers["Authorization"] == "Bearer new-access-token" + + # Verify that non-allowlisted headers are NOT forwarded (security) + # Only anthropic-beta, content-type, and Authorization should be present + assert "authorization" not in headers # lowercase auth token not forwarded + assert "user-agent" not in headers # not in allowlist + assert "content-length" not in headers # not in allowlist + assert "host" not in headers # not in allowlist + + # Verify that headers_passed_through is False (since we have credentials) + assert headers_passed_through is False + + +@pytest.mark.asyncio +async def test_vertex_passthrough_does_not_forward_litellm_auth_token(): + """ + Test that the LiteLLM authorization header is NOT forwarded to Vertex AI. + + This test validates the fix for the issue where both the LiteLLM auth token + (lowercase 'authorization') and the Vertex AI token (uppercase 'Authorization') + were being sent, causing 401 errors on the vendor side. + + The incoming request has: + - authorization: Bearer (should NOT be forwarded) + + The outgoing request should only have: + - Authorization: Bearer (vendor credentials) + """ + from starlette.datastructures import Headers + + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _prepare_vertex_auth_headers, + ) + + # Create a mock request with ONLY the litellm auth token (no other headers) + mock_request = MagicMock() + mock_request.headers = Headers({ + "authorization": "Bearer sk-litellm-secret-key", # LiteLLM token - should NOT be forwarded + "Authorization": "Bearer sk-litellm-secret-key-uppercase", # Also try uppercase + }) + + # Create mock vertex credentials + mock_vertex_credentials = MagicMock() + mock_vertex_credentials.vertex_project = "test-project" + mock_vertex_credentials.vertex_location = "us-central1" + mock_vertex_credentials.vertex_credentials = "test-credentials" + + # Create mock handler + mock_handler = MagicMock() + mock_handler.update_base_target_url_with_credential_location.return_value = ( + "https://us-central1-aiplatform.googleapis.com" + ) + + with patch.object( + VertexBase, + "_ensure_access_token_async", + new_callable=AsyncMock, + return_value=("test-auth-header", "test-project"), + ), patch.object( + VertexBase, + "_get_token_and_url", + return_value=("vertex-access-token", None), + ): + + ( + headers, + _base_target_url, + _headers_passed_through, + _vertex_project, + _vertex_location, + ) = await _prepare_vertex_auth_headers( + request=mock_request, + vertex_credentials=mock_vertex_credentials, + router_credentials=None, + vertex_project="test-project", + vertex_location="us-central1", + base_target_url="https://us-central1-aiplatform.googleapis.com", + get_vertex_pass_through_handler=mock_handler, + ) + + # The ONLY Authorization header should be the Vertex token + assert headers["Authorization"] == "Bearer vertex-access-token" + + # The LiteLLM token should NOT be present (neither lowercase nor as a duplicate) + assert "authorization" not in headers + assert headers.get("Authorization") != "Bearer sk-litellm-secret-key" + assert headers.get("Authorization") != "Bearer sk-litellm-secret-key-uppercase" + + # Verify we only have the expected headers (Authorization + any allowlisted ones present) + # Since the request only had auth headers, only Authorization should be in output + assert set(headers.keys()) == {"Authorization"} + + +def test_forward_headers_from_request_x_pass_prefix(): + """ + Test that headers with 'x-pass-' prefix are forwarded with the prefix stripped. + + This allows users to force-forward arbitrary headers to the vendor API: + - 'x-pass-anthropic-beta: value' becomes 'anthropic-beta: value' + - 'x-pass-custom-header: value' becomes 'custom-header: value' + + This is tested on BasePassthroughUtils.forward_headers_from_request which is used + by all pass-through endpoints (not just Vertex AI). + """ + from litellm.passthrough.utils import BasePassthroughUtils + + # Simulate incoming request headers + request_headers = { + "x-pass-anthropic-beta": "context-1m-2025-08-07", + "x-pass-custom-header": "custom-value", + "x-pass-another-header": "another-value", + "authorization": "Bearer sk-litellm-key", + "x-litellm-api-key": "sk-1234", + "content-type": "application/json", + } + + # Start with empty headers dict (simulating custom headers from endpoint config) + headers = {} + + # Call the method with forward_headers=False (default behavior) + # x-pass- headers should still be forwarded + result = BasePassthroughUtils.forward_headers_from_request( + request_headers=request_headers, + headers=headers, + forward_headers=False, + ) + + # Verify x-pass- prefixed headers are forwarded with prefix stripped + assert "anthropic-beta" in result + assert result["anthropic-beta"] == "context-1m-2025-08-07" + assert "custom-header" in result + assert result["custom-header"] == "custom-value" + assert "another-header" in result + assert result["another-header"] == "another-value" + + # Verify other headers are NOT forwarded (since forward_headers=False) + assert "authorization" not in result + assert "x-litellm-api-key" not in result + assert "content-type" not in result + + # Verify original x-pass- prefixed headers are NOT in output (only stripped versions) + assert "x-pass-anthropic-beta" not in result + assert "x-pass-custom-header" not in result + diff --git a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py new file mode 100644 index 0000000000..2c5bc1bf87 --- /dev/null +++ b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py @@ -0,0 +1,189 @@ +import pytest +from unittest.mock import MagicMock, AsyncMock, patch +from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles +from litellm.types.prompts.init_prompts import ( + PromptSpec, + PromptLiteLLMParams, + PromptInfo, +) + + +@pytest.mark.asyncio +async def test_delete_prompt_success(): + """ + Test that delete_prompt correctly identifies the base prompt ID + and deletes all versions from DB and memory. + """ + from litellm.proxy.prompts.prompt_endpoints import delete_prompt + + # Mock user auth + mock_user_auth = UserAPIKeyAuth( + api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + # Mock DB Client + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_prompttable.delete_many = AsyncMock(return_value=None) + + # Mock In-Memory Registry + with patch( + "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" + ) as mock_registry: + # User passes "test_prompt.v2" + # We simulate that get_prompt_by_id returns the prompt spec for v2 + prompt_spec = PromptSpec( + prompt_id="test_prompt.v2", + litellm_params=PromptLiteLLMParams( + prompt_id="test_prompt", prompt_integration="dotprompt" + ), + prompt_info=PromptInfo(prompt_type="db"), + ) + mock_registry.get_prompt_by_id.return_value = prompt_spec + + # Patch the prisma client in the endpoint module + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + response = await delete_prompt( + prompt_id="test_prompt.v2", user_api_key_dict=mock_user_auth + ) + + # Assertions + expected_base_id = "test_prompt" + + # 1. DB deletion should use base ID + mock_prisma_client.db.litellm_prompttable.delete_many.assert_called_once_with( + where={"prompt_id": expected_base_id} + ) + + # 2. Memory deletion should use base ID + mock_registry.delete_prompts_by_base_id.assert_called_once_with( + expected_base_id + ) + + assert response == { + "message": f"Prompt {expected_base_id} deleted successfully" + } + + +@pytest.mark.asyncio +async def test_delete_prompt_by_base_id_success(): + """ + Test that delete_prompt works when passed a base ID directly, + finding the latest version to confirm existence, then deleting. + """ + from litellm.proxy.prompts.prompt_endpoints import delete_prompt + + # Mock user auth + mock_user_auth = UserAPIKeyAuth( + api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + # Mock DB Client + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_prompttable.delete_many = AsyncMock(return_value=None) + + # Mock In-Memory Registry + with patch( + "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" + ) as mock_registry: + # User passes "test_prompt" (base ID) + # 1. get_prompt_by_id("test_prompt") -> None (if it's not registered as base) + # 2. It calls get_latest_version_prompt_id -> returns "test_prompt.v3" + # 3. get_prompt_by_id("test_prompt.v3") -> returns Spec + + # Setup mocks behavior + def get_prompt_side_effect(prompt_id): + if prompt_id == "test_prompt": + return None + if prompt_id == "test_prompt.v3": + return PromptSpec( + prompt_id="test_prompt.v3", + litellm_params=PromptLiteLLMParams( + prompt_id="test_prompt", prompt_integration="dotprompt" + ), + prompt_info=PromptInfo(prompt_type="db"), + ) + return None + + mock_registry.get_prompt_by_id.side_effect = get_prompt_side_effect + mock_registry.IN_MEMORY_PROMPTS = { + "test_prompt.v1": {}, + "test_prompt.v2": {}, + "test_prompt.v3": {}, + } + + # Patch the prisma client in the endpoint module + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + response = await delete_prompt( + prompt_id="test_prompt", user_api_key_dict=mock_user_auth + ) + + # Assertions + expected_base_id = "test_prompt" + + # 1. DB deletion should use base ID + mock_prisma_client.db.litellm_prompttable.delete_many.assert_called_once_with( + where={"prompt_id": expected_base_id} + ) + + # 2. Memory deletion should use base ID + mock_registry.delete_prompts_by_base_id.assert_called_once_with( + expected_base_id + ) + + assert response == { + "message": f"Prompt {expected_base_id} deleted successfully" + } + + +@pytest.mark.asyncio +async def test_get_prompt_info_by_base_id(): + """ + Test that get_prompt_info correctly resolves a base ID to the latest version. + """ + from litellm.proxy.prompts.prompt_endpoints import get_prompt_info + + # Mock user auth + mock_user_auth = UserAPIKeyAuth( + api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + # Mock In-Memory Registry + with patch( + "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" + ) as mock_registry: + # Setup mocks behavior + prompt_spec_v3 = PromptSpec( + prompt_id="test_prompt.v3", + litellm_params=PromptLiteLLMParams( + prompt_id="test_prompt", prompt_integration="dotprompt" + ), + prompt_info=PromptInfo(prompt_type="db"), + ) + + # When get_prompt_by_id is called with "test_prompt", return None (so it searches versions) + # When called with "test_prompt.v3", return the spec + def get_prompt_side_effect(prompt_id): + if prompt_id == "test_prompt": + return None + if prompt_id == "test_prompt.v3": + return prompt_spec_v3 + return None + + mock_registry.get_prompt_by_id.side_effect = get_prompt_side_effect + mock_registry.IN_MEMORY_PROMPTS = { + "test_prompt.v1": {}, + "test_prompt.v2": {}, + "test_prompt.v3": {}, + } + + # We also need to mock get_prompt_callback_by_id to avoid content extraction errors/logic + mock_registry.get_prompt_callback_by_id.return_value = None + + response = await get_prompt_info( + prompt_id="test_prompt", user_api_key_dict=mock_user_auth + ) + + assert ( + response.prompt_spec.prompt_id == "test_prompt" + ) # Should return base ID in spec response + assert response.prompt_spec.version == 3 # Should identify it as version 3 diff --git a/tests/test_litellm/proxy/test_empty_model_list.py b/tests/test_litellm/proxy/test_empty_model_list.py index 6b3e59d319..dd900d3eb5 100644 --- a/tests/test_litellm/proxy/test_empty_model_list.py +++ b/tests/test_litellm/proxy/test_empty_model_list.py @@ -32,7 +32,7 @@ class TestEmptyModelListHandling: self, client, monkeypatch ): """ - Test that /v2/model/info returns {"data": []} instead of 500 + Test that /v2/model/info returns paginated empty response instead of 500 when llm_router is None. """ monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) @@ -56,13 +56,18 @@ class TestEmptyModelListHandling: ) assert response.status_code == 200 - assert response.json() == {"data": []} + data = response.json() + assert data["data"] == [] + assert data["total_count"] == 0 + assert data["current_page"] == 1 + assert data["total_pages"] == 0 + assert data["size"] == 50 # default page size def test_v2_model_info_returns_empty_data_when_model_list_empty( self, client, monkeypatch ): """ - Test that /v2/model/info returns {"data": []} instead of 500 + Test that /v2/model/info returns paginated empty response instead of 500 when llm_router exists but model_list is empty. """ mock_router = MagicMock() @@ -89,7 +94,52 @@ class TestEmptyModelListHandling: ) assert response.status_code == 200 - assert response.json() == {"data": []} + data = response.json() + assert data["data"] == [] + assert data["total_count"] == 0 + assert data["current_page"] == 1 + assert data["total_pages"] == 0 + assert data["size"] == 50 # default page size + + def test_v2_model_info_pagination_with_empty_results( + self, client, monkeypatch + ): + """ + Test that /v2/model/info pagination parameters work correctly + when there are no models (empty results). + """ + mock_router = MagicMock() + mock_router.model_list = [] + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_model_list", []) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + with patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + return_value=MagicMock( + user_id="test-user", + team_id=None, + team_models=[], + models=[], + user_role="proxy_admin", + ), + ): + # Test with custom pagination parameters + response = client.get( + "/v2/model/info", + params={"page": 2, "size": 25}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + data = response.json() + assert data["data"] == [] + assert data["total_count"] == 0 + assert data["current_page"] == 2 # Should respect the page parameter + assert data["total_pages"] == 0 + assert data["size"] == 25 # Should respect the size parameter def test_model_group_info_returns_empty_data_when_model_list_none( self, client, monkeypatch diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 751a903387..cb519e9f50 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3216,3 +3216,851 @@ async def test_get_hierarchical_router_settings(): prisma_client=mock_prisma_client, ) assert result is None + + +@pytest.mark.asyncio +async def test_model_info_v2_pagination_basic(monkeypatch): + """ + Test basic pagination functionality for /v2/model/info endpoint. + Tests multiple pages with different page sizes. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + # Create 75 mock models for testing pagination + mock_models = [ + { + "model_name": f"model-{i}", + "litellm_params": {"model": f"gpt-{i}"}, + "model_info": {"id": f"model-{i}"}, + } + for i in range(1, 76) # 75 models total + ] + + # Mock llm_router + mock_router = MagicMock() + mock_router.model_list = mock_models + + # Mock prisma_client + mock_prisma_client = MagicMock() + + # Mock proxy_config.get_config + mock_get_config = AsyncMock(return_value={}) + + # Mock user authentication + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.user_id = "test-user" + mock_user_api_key_dict.api_key = "test-key" + mock_user_api_key_dict.team_models = [] + mock_user_api_key_dict.models = [] + + # Apply monkeypatches + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + + # Override auth dependency + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_api_key_dict + + client = TestClient(app) + try: + # Test page 1 with size 25 (should return models 1-25) + response = client.get("/v2/model/info", params={"page": 1, "size": 25}) + assert response.status_code == 200 + data = response.json() + assert data["total_count"] == 75 + assert data["current_page"] == 1 + assert data["size"] == 25 + assert data["total_pages"] == 3 # ceil(75/25) = 3 + assert len(data["data"]) == 25 + assert data["data"][0]["model_name"] == "model-1" + assert data["data"][24]["model_name"] == "model-25" + + # Test page 2 with size 25 (should return models 26-50) + response = client.get("/v2/model/info", params={"page": 2, "size": 25}) + assert response.status_code == 200 + data = response.json() + assert data["total_count"] == 75 + assert data["current_page"] == 2 + assert data["size"] == 25 + assert data["total_pages"] == 3 + assert len(data["data"]) == 25 + assert data["data"][0]["model_name"] == "model-26" + assert data["data"][24]["model_name"] == "model-50" + + # Test page 3 with size 25 (should return models 51-75) + response = client.get("/v2/model/info", params={"page": 3, "size": 25}) + assert response.status_code == 200 + data = response.json() + assert data["total_count"] == 75 + assert data["current_page"] == 3 + assert data["size"] == 25 + assert data["total_pages"] == 3 + assert len(data["data"]) == 25 + assert data["data"][0]["model_name"] == "model-51" + assert data["data"][24]["model_name"] == "model-75" + + # Test different page size (size 10) + response = client.get("/v2/model/info", params={"page": 1, "size": 10}) + assert response.status_code == 200 + data = response.json() + assert data["total_count"] == 75 + assert data["current_page"] == 1 + assert data["size"] == 10 + assert data["total_pages"] == 8 # ceil(75/10) = 8 + assert len(data["data"]) == 10 + + finally: + app.dependency_overrides = original_overrides + + +@pytest.mark.asyncio +async def test_model_info_v2_pagination_edge_cases(monkeypatch): + """ + Test edge cases for pagination in /v2/model/info endpoint. + Tests empty results, last page with partial results, and boundary conditions. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + # Mock prisma_client + mock_prisma_client = MagicMock() + + # Mock user authentication + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.user_id = "test-user" + mock_user_api_key_dict.api_key = "test-key" + mock_user_api_key_dict.team_models = [] + mock_user_api_key_dict.models = [] + + # Mock proxy_config.get_config + mock_get_config = AsyncMock(return_value={}) + + # Apply monkeypatches + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + + # Override auth dependency + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_api_key_dict + + client = TestClient(app) + try: + # Test Case 1: Empty model list (no models configured) + mock_router_empty = MagicMock() + mock_router_empty.model_list = [] + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router_empty) + + response = client.get("/v2/model/info", params={"page": 1, "size": 25}) + assert response.status_code == 200 + data = response.json() + assert data["total_count"] == 0 + assert data["current_page"] == 1 + assert data["size"] == 25 + assert data["total_pages"] == 0 + assert len(data["data"]) == 0 + + # Test Case 2: Last page with partial results (23 models, page size 10) + mock_models_partial = [ + { + "model_name": f"model-{i}", + "litellm_params": {"model": f"gpt-{i}"}, + "model_info": {"id": f"model-{i}"}, + } + for i in range(1, 24) # 23 models total + ] + mock_router_partial = MagicMock() + mock_router_partial.model_list = mock_models_partial + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router_partial) + + # Page 1 should have 10 models + response = client.get("/v2/model/info", params={"page": 1, "size": 10}) + assert response.status_code == 200 + data = response.json() + assert data["total_count"] == 23 + assert data["current_page"] == 1 + assert data["total_pages"] == 3 # ceil(23/10) = 3 + assert len(data["data"]) == 10 + + # Page 2 should have 10 models + response = client.get("/v2/model/info", params={"page": 2, "size": 10}) + assert response.status_code == 200 + data = response.json() + assert data["total_count"] == 23 + assert data["current_page"] == 2 + assert data["total_pages"] == 3 + assert len(data["data"]) == 10 + + # Page 3 (last page) should have only 3 models + response = client.get("/v2/model/info", params={"page": 3, "size": 10}) + assert response.status_code == 200 + data = response.json() + assert data["total_count"] == 23 + assert data["current_page"] == 3 + assert data["total_pages"] == 3 + assert len(data["data"]) == 3 + assert data["data"][0]["model_name"] == "model-21" + assert data["data"][2]["model_name"] == "model-23" + + # Test Case 3: Page beyond available pages (should return empty data) + response = client.get("/v2/model/info", params={"page": 4, "size": 10}) + assert response.status_code == 200 + data = response.json() + assert data["total_count"] == 23 + assert data["current_page"] == 4 + assert data["total_pages"] == 3 + assert len(data["data"]) == 0 # No data for page beyond total_pages + + # Test Case 4: Single model with page size 1 + mock_models_single = [ + { + "model_name": "single-model", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "single-model"}, + } + ] + mock_router_single = MagicMock() + mock_router_single.model_list = mock_models_single + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router_single) + + response = client.get("/v2/model/info", params={"page": 1, "size": 1}) + assert response.status_code == 200 + data = response.json() + assert data["total_count"] == 1 + assert data["current_page"] == 1 + assert data["total_pages"] == 1 + assert len(data["data"]) == 1 + assert data["data"][0]["model_name"] == "single-model" + + finally: + app.dependency_overrides = original_overrides + + +def test_enrich_model_info_with_litellm_data(): + """ + Test the _enrich_model_info_with_litellm_data helper function. + Tests model info enrichment, debug mode, and sensitive info removal. + """ + from unittest.mock import MagicMock, patch + + from litellm.proxy.proxy_server import _enrich_model_info_with_litellm_data + + # Test Case 1: Basic model enrichment without debug + model = { + "model_name": "test-model", + "litellm_params": {"model": "gpt-3.5-turbo"}, + "model_info": {"id": "test-model"}, + "api_key": "sk-secret-key", # Should be removed + } + + with patch("litellm.proxy.proxy_server.get_litellm_model_info") as mock_get_info, patch( + "litellm.proxy.proxy_server.remove_sensitive_info_from_deployment" + ) as mock_remove_sensitive: + mock_get_info.return_value = { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "max_tokens": 4096, + } + mock_remove_sensitive.return_value = { + "model_name": "test-model", + "litellm_params": {"model": "gpt-3.5-turbo"}, + "model_info": { + "id": "test-model", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "max_tokens": 4096, + }, + } + + result = _enrich_model_info_with_litellm_data(model=model, debug=False) + + # Verify get_litellm_model_info was called + mock_get_info.assert_called_once_with(model=model) + # Verify remove_sensitive_info_from_deployment was called + mock_remove_sensitive.assert_called_once() + # Verify result doesn't have api_key + assert "api_key" not in result + # Verify model_info was enriched + assert "input_cost_per_token" in result["model_info"] + + # Test Case 2: Model enrichment with debug mode + model_with_debug = { + "model_name": "test-model-debug", + "litellm_params": {"model": "gpt-4"}, + "model_info": {}, + } + + mock_router = MagicMock() + mock_client = MagicMock() + mock_router._get_client.return_value = mock_client + + with patch("litellm.proxy.proxy_server.get_litellm_model_info") as mock_get_info, patch( + "litellm.proxy.proxy_server.remove_sensitive_info_from_deployment" + ) as mock_remove_sensitive: + mock_get_info.return_value = {} + mock_remove_sensitive.return_value = { + "model_name": "test-model-debug", + "litellm_params": {"model": "gpt-4"}, + "model_info": {}, + "openai_client": str(mock_client), + } + + result = _enrich_model_info_with_litellm_data( + model=model_with_debug, debug=True, llm_router=mock_router + ) + + # Verify debug info was added + mock_remove_sensitive.assert_called_once() + call_args = mock_remove_sensitive.call_args[0][0] + assert "openai_client" in call_args + # Verify router._get_client was called for debug + mock_router._get_client.assert_called_once() + + # Test Case 3: Model with fallback to litellm.get_model_info + model_fallback = { + "model_name": "test-model-fallback", + "litellm_params": {"model": "claude-3-opus"}, + "model_info": {}, + } + + with patch("litellm.proxy.proxy_server.get_litellm_model_info") as mock_get_info, patch( + "litellm.get_model_info" + ) as mock_litellm_info, patch( + "litellm.proxy.proxy_server.remove_sensitive_info_from_deployment" + ) as mock_remove_sensitive: + # First call returns empty, triggering fallback + mock_get_info.return_value = {} + mock_litellm_info.return_value = { + "input_cost_per_token": 0.015, + "output_cost_per_token": 0.075, + "max_tokens": 200000, + } + mock_remove_sensitive.return_value = { + "model_name": "test-model-fallback", + "litellm_params": {"model": "claude-3-opus"}, + "model_info": { + "input_cost_per_token": 0.015, + "output_cost_per_token": 0.075, + "max_tokens": 200000, + }, + } + + result = _enrich_model_info_with_litellm_data(model=model_fallback, debug=False) + + # Verify fallback was attempted + mock_litellm_info.assert_called_once_with(model="claude-3-opus") + # Verify model_info was enriched with fallback data + call_args = mock_remove_sensitive.call_args[0][0] + assert call_args["model_info"]["input_cost_per_token"] == 0.015 + + # Test Case 4: Model with split model name fallback + model_split = { + "model_name": "test-model-split", + "litellm_params": {"model": "azure/gpt-4"}, + "model_info": {}, + } + + with patch("litellm.proxy.proxy_server.get_litellm_model_info") as mock_get_info, patch( + "litellm.get_model_info" + ) as mock_litellm_info, patch( + "litellm.proxy.proxy_server.remove_sensitive_info_from_deployment" + ) as mock_remove_sensitive: + # Both first and second pass return empty, triggering third pass + mock_get_info.return_value = {} + # Second pass (no split) + mock_litellm_info.side_effect = [ + {}, # First call returns empty + {"max_tokens": 8192}, # Third pass with split succeeds + ] + mock_remove_sensitive.return_value = { + "model_name": "test-model-split", + "litellm_params": {"model": "azure/gpt-4"}, + "model_info": {"max_tokens": 8192}, + } + + result = _enrich_model_info_with_litellm_data(model=model_split, debug=False) + + # Verify third pass was attempted with split model name + assert mock_litellm_info.call_count == 2 + # Check that second call used split model name + second_call = mock_litellm_info.call_args_list[1] + assert second_call[1]["model"] == "gpt-4" + assert second_call[1]["custom_llm_provider"] == "azure" + + # Test Case 5: Model with existing model_info (should preserve existing keys) + model_existing = { + "model_name": "test-model-existing", + "litellm_params": {"model": "gpt-3.5-turbo"}, + "model_info": {"id": "existing-id", "custom_key": "custom_value"}, + } + + with patch("litellm.proxy.proxy_server.get_litellm_model_info") as mock_get_info, patch( + "litellm.proxy.proxy_server.remove_sensitive_info_from_deployment" + ) as mock_remove_sensitive: + mock_get_info.return_value = { + "input_cost_per_token": 0.001, + "id": "new-id", # Should not override existing "id" + } + mock_remove_sensitive.return_value = { + "model_name": "test-model-existing", + "litellm_params": {"model": "gpt-3.5-turbo"}, + "model_info": { + "id": "existing-id", # Existing key preserved + "custom_key": "custom_value", # Existing key preserved + "input_cost_per_token": 0.001, # New key added + }, + } + + result = _enrich_model_info_with_litellm_data(model=model_existing, debug=False) + + # Verify existing keys are preserved + call_args = mock_remove_sensitive.call_args[0][0] + assert call_args["model_info"]["id"] == "existing-id" + assert call_args["model_info"]["custom_key"] == "custom_value" + assert call_args["model_info"]["input_cost_per_token"] == 0.001 + + +@pytest.mark.asyncio +async def test_model_list_scope_parameter_validation(monkeypatch): + """Test that invalid scope parameter raises HTTPException""" + from fastapi import HTTPException + from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles + from litellm.proxy.proxy_server import model_list + + mock_user_api_key_dict = UserAPIKeyAuth( + user_id="test-user", + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="test-key", + ) + + # Test invalid scope parameter + with pytest.raises(HTTPException) as exc_info: + await model_list( + user_api_key_dict=mock_user_api_key_dict, + scope="invalid_scope", + ) + + assert exc_info.value.status_code == 400 + assert "Invalid scope parameter" in exc_info.value.detail + assert "Only 'expand' is currently supported" in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_model_list_scope_expand_proxy_admin(monkeypatch): + """Test that proxy admin with scope=expand returns all proxy models""" + from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles, LiteLLM_UserTable + from litellm.proxy.proxy_server import model_list + + # Mock user API key dict for proxy admin + mock_user_api_key_dict = UserAPIKeyAuth( + user_id="proxy-admin-user", + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="test-key", + ) + + # Mock llm_router with proxy models + mock_router = MagicMock() + mock_router.get_model_names.return_value = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus"] + mock_router.get_model_access_groups.return_value = {} + + # Mock prisma_client + mock_prisma_client = MagicMock() + + # Mock user_api_key_cache + mock_user_api_key_cache = MagicMock() + + # Mock proxy_logging_obj + mock_proxy_logging_obj = MagicMock() + + # Mock get_complete_model_list + mock_all_models = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus"] + + # Mock create_model_info_response + def mock_create_model_info_response(model_id, provider, include_metadata=False, fallback_type=None, llm_router=None): + return {"id": model_id, "object": "model"} + + # Apply monkeypatches + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr( + "litellm.proxy.auth.model_checks.get_complete_model_list", + lambda **kwargs: mock_all_models, + ) + monkeypatch.setattr( + "litellm.proxy.utils.create_model_info_response", + mock_create_model_info_response, + ) + + # Call model_list with scope=expand + result = await model_list( + user_api_key_dict=mock_user_api_key_dict, + scope="expand", + ) + + # Verify result contains all proxy models + assert result["object"] == "list" + assert len(result["data"]) == 3 + assert all(model["id"] in mock_all_models for model in result["data"]) + + # Verify router methods were called + mock_router.get_model_names.assert_called_once() + mock_router.get_model_access_groups.assert_called_once() + + +@pytest.mark.asyncio +async def test_model_list_scope_expand_org_admin(monkeypatch): + """Test that org admin with scope=expand returns all proxy models""" + from litellm.proxy._types import ( + UserAPIKeyAuth, + LitellmUserRoles, + LiteLLM_UserTable, + ) + from litellm.proxy.proxy_server import model_list + + # Mock user API key dict for org admin + mock_user_api_key_dict = UserAPIKeyAuth( + user_id="org-admin-user", + user_role=LitellmUserRoles.INTERNAL_USER, # Not proxy admin, but org admin + api_key="test-key", + ) + + # Mock user object with org admin membership + from litellm.proxy._types import LiteLLM_OrganizationMembershipTable + from datetime import datetime + + mock_user_obj = LiteLLM_UserTable( + user_id="org-admin-user", + user_email="org-admin@example.com", + organization_memberships=[ + LiteLLM_OrganizationMembershipTable( + user_id="org-admin-user", + organization_id="org-123", + user_role=LitellmUserRoles.ORG_ADMIN.value, + spend=0.0, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + ], + teams=[], + ) + + # Mock llm_router with proxy models + mock_router = MagicMock() + mock_router.get_model_names.return_value = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus"] + mock_router.get_model_access_groups.return_value = {} + + # Mock prisma_client + mock_prisma_client = MagicMock() + + # Mock user_api_key_cache + mock_user_api_key_cache = MagicMock() + + # Mock proxy_logging_obj + mock_proxy_logging_obj = MagicMock() + + # Mock get_user_object to return user with org admin role + async def mock_get_user_object(*args, **kwargs): + return mock_user_obj + + # Mock get_complete_model_list + mock_all_models = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus"] + + # Mock create_model_info_response + def mock_create_model_info_response(model_id, provider, include_metadata=False, fallback_type=None, llm_router=None): + return {"id": model_id, "object": "model"} + + # Apply monkeypatches + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr( + "litellm.proxy.auth.auth_checks.get_user_object", + mock_get_user_object, + ) + monkeypatch.setattr( + "litellm.proxy.auth.model_checks.get_complete_model_list", + lambda **kwargs: mock_all_models, + ) + monkeypatch.setattr( + "litellm.proxy.utils.create_model_info_response", + mock_create_model_info_response, + ) + + # Call model_list with scope=expand + result = await model_list( + user_api_key_dict=mock_user_api_key_dict, + scope="expand", + ) + + # Verify result contains all proxy models + assert result["object"] == "list" + assert len(result["data"]) == 3 + assert all(model["id"] in mock_all_models for model in result["data"]) + + # Verify router methods were called + mock_router.get_model_names.assert_called_once() + mock_router.get_model_access_groups.assert_called_once() + + +@pytest.mark.asyncio +async def test_model_list_scope_expand_team_admin(monkeypatch): + """Test that team admin with scope=expand returns all proxy models""" + from litellm.proxy._types import ( + UserAPIKeyAuth, + LitellmUserRoles, + LiteLLM_UserTable, + LiteLLM_TeamTable, + ) + from litellm.proxy.proxy_server import model_list + + # Mock user API key dict for team admin + mock_user_api_key_dict = UserAPIKeyAuth( + user_id="team-admin-user", + user_role=LitellmUserRoles.INTERNAL_USER, # Not proxy admin, but team admin + api_key="test-key", + ) + + # Mock team with user as admin - use dict structure that matches Prisma return + mock_team = MagicMock() + mock_team.model_dump.return_value = { + "team_id": "team-123", + "members_with_roles": [ + {"user_id": "team-admin-user", "role": "admin"} + ], + } + # Create team object from the dict (validator will convert members_with_roles to Member objects) + mock_team_obj = LiteLLM_TeamTable(**mock_team.model_dump()) + + # Mock user object with team membership + mock_user_obj = LiteLLM_UserTable( + user_id="team-admin-user", + user_email="team-admin@example.com", + organization_memberships=[], + teams=["team-123"], + ) + + # Mock llm_router with proxy models + mock_router = MagicMock() + mock_router.get_model_names.return_value = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus"] + mock_router.get_model_access_groups.return_value = {} + + # Mock prisma_client + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock( + return_value=[mock_team] + ) + + # Mock user_api_key_cache + mock_user_api_key_cache = MagicMock() + + # Mock proxy_logging_obj + mock_proxy_logging_obj = MagicMock() + + # Mock get_user_object to return user with team membership + async def mock_get_user_object(*args, **kwargs): + return mock_user_obj + + # Mock get_complete_model_list + mock_all_models = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus"] + + # Mock create_model_info_response + def mock_create_model_info_response(model_id, provider, include_metadata=False, fallback_type=None, llm_router=None): + return {"id": model_id, "object": "model"} + + # Apply monkeypatches + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr( + "litellm.proxy.auth.auth_checks.get_user_object", + mock_get_user_object, + ) + monkeypatch.setattr( + "litellm.proxy.auth.model_checks.get_complete_model_list", + lambda **kwargs: mock_all_models, + ) + monkeypatch.setattr( + "litellm.proxy.utils.create_model_info_response", + mock_create_model_info_response, + ) + + # Call model_list with scope=expand + result = await model_list( + user_api_key_dict=mock_user_api_key_dict, + scope="expand", + ) + + # Verify result contains all proxy models + assert result["object"] == "list" + assert len(result["data"]) == 3 + assert all(model["id"] in mock_all_models for model in result["data"]) + + # Verify router methods were called + mock_router.get_model_names.assert_called_once() + mock_router.get_model_access_groups.assert_called_once() + + +@pytest.mark.asyncio +async def test_model_list_scope_expand_normal_user(monkeypatch): + """Test that normal internal user with scope=expand returns only their models (not expanded)""" + from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles, LiteLLM_UserTable + from litellm.proxy.proxy_server import model_list + + # Mock user API key dict for normal internal user + mock_user_api_key_dict = UserAPIKeyAuth( + user_id="normal-user", + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="test-key", + models=["gpt-3.5-turbo"], # User only has access to this model + ) + + # Mock user object without admin privileges + mock_user_obj = LiteLLM_UserTable( + user_id="normal-user", + user_email="normal@example.com", + organization_memberships=[], # No org admin + teams=[], # No teams + ) + + # Mock llm_router + mock_router = MagicMock() + mock_router.get_model_names.return_value = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus"] + + # Mock prisma_client + mock_prisma_client = MagicMock() + + # Mock user_api_key_cache + mock_user_api_key_cache = MagicMock() + + # Mock proxy_logging_obj + mock_proxy_logging_obj = MagicMock() + + # Mock get_user_object to return user without admin privileges + async def mock_get_user_object(*args, **kwargs): + return mock_user_obj + + # Mock get_available_models_for_user to return only user's models + async def mock_get_available_models_for_user(*args, **kwargs): + return ["gpt-3.5-turbo"] # Only user's accessible models + + # Mock create_model_info_response + def mock_create_model_info_response(model_id, provider, include_metadata=False, fallback_type=None, llm_router=None): + return {"id": model_id, "object": "model"} + + # Apply monkeypatches + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr( + "litellm.proxy.auth.auth_checks.get_user_object", + mock_get_user_object, + ) + monkeypatch.setattr( + "litellm.proxy.utils.get_available_models_for_user", + mock_get_available_models_for_user, + ) + monkeypatch.setattr( + "litellm.proxy.utils.create_model_info_response", + mock_create_model_info_response, + ) + + # Call model_list with scope=expand + result = await model_list( + user_api_key_dict=mock_user_api_key_dict, + scope="expand", + ) + + # Verify result contains only user's models (not all proxy models) + assert result["object"] == "list" + assert len(result["data"]) == 1 + assert result["data"][0]["id"] == "gpt-3.5-turbo" + + # Verify router methods were NOT called (normal path, not expanded) + mock_router.get_model_names.assert_not_called() + mock_router.get_model_access_groups.assert_not_called() + + +@pytest.mark.asyncio +async def test_model_list_no_scope_parameter(monkeypatch): + """Test that model_list without scope parameter uses normal behavior""" + from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles + from litellm.proxy.proxy_server import model_list + + # Mock user API key dict + mock_user_api_key_dict = UserAPIKeyAuth( + user_id="test-user", + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="test-key", + models=["gpt-3.5-turbo"], + ) + + # Mock llm_router + mock_router = MagicMock() + + # Mock prisma_client + mock_prisma_client = MagicMock() + + # Mock user_api_key_cache + mock_user_api_key_cache = MagicMock() + + # Mock proxy_logging_obj + mock_proxy_logging_obj = MagicMock() + + # Mock get_available_models_for_user + async def mock_get_available_models_for_user(*args, **kwargs): + return ["gpt-3.5-turbo"] + + # Mock create_model_info_response + def mock_create_model_info_response(model_id, provider, include_metadata=False, fallback_type=None, llm_router=None): + return {"id": model_id, "object": "model"} + + # Apply monkeypatches + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr( + "litellm.proxy.utils.get_available_models_for_user", + mock_get_available_models_for_user, + ) + monkeypatch.setattr( + "litellm.proxy.utils.create_model_info_response", + mock_create_model_info_response, + ) + + # Call model_list without scope parameter + result = await model_list( + user_api_key_dict=mock_user_api_key_dict, + scope=None, + ) + + # Verify result uses normal behavior + assert result["object"] == "list" + assert len(result["data"]) == 1 + assert result["data"][0]["id"] == "gpt-3.5-turbo" + + # Verify router methods were NOT called (normal path) + mock_router.get_model_names.assert_not_called() + mock_router.get_model_access_groups.assert_not_called() diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py index 51150383b0..8d324bea61 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_tool_call_streaming_transformation.py @@ -14,7 +14,12 @@ from litellm.responses.litellm_completion_transformation.streaming_iterator impo LiteLLMCompletionStreamingIterator, ) from litellm.types.llms.openai import ResponsesAPIStreamEvents -from litellm.types.utils import Delta, ModelResponse, ModelResponseStream, StreamingChoices +from litellm.types.utils import ( + Delta, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) def test_tool_call_delta_is_emitted_as_responses_events(): @@ -55,12 +60,14 @@ def test_tool_call_delta_is_emitted_as_responses_events(): assert evt1.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED assert evt1.output_index == 1 + # The arguments are now chunked, so we get the first delta chunk evt2 = iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk) assert evt2 is not None assert evt2.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA assert evt2.item_id == "call_1" assert evt2.output_index == 1 - assert evt2.delta == '{"x":1}' + # The delta will be a chunk of the arguments, not the full arguments + assert len(evt2.delta) <= 10 # Chunks are max 10 characters def test_tool_calls_present_only_in_final_response_are_emitted_before_completed(): @@ -104,13 +111,121 @@ def test_tool_calls_present_only_in_final_response_are_emitted_before_completed( assert evt1.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED assert evt1.output_index == 1 - evt2 = iterator.common_done_event_logic(sync_mode=True) - assert evt2.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE - assert evt2.item_id == "call_2" - assert evt2.output_index == 1 - assert evt2.arguments == '{"y":2}' + # Now delta events are emitted (arguments split into chunks) + # Collect all delta events + delta_events = [] + while True: + evt = iterator.common_done_event_logic(sync_mode=True) + if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA: + delta_events.append(evt) + else: + break + + # Verify we got delta events + assert len(delta_events) > 0 + # Verify they reconstruct the original arguments + concatenated_args = ''.join(evt.delta for evt in delta_events) + assert concatenated_args == '{"y":2}' - evt3 = iterator.common_done_event_logic(sync_mode=True) - assert evt3.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE - assert evt3.output_index == 1 + # The last event should be FUNCTION_CALL_ARGUMENTS_DONE + assert evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE + assert evt.item_id == "call_2" + assert evt.output_index == 1 + assert evt.arguments == '{"y":2}' + + evt_final = iterator.common_done_event_logic(sync_mode=True) + assert evt_final.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE + assert evt_final.output_index == 1 + + +def test_tool_call_arguments_are_chunked_to_match_openai_behavior(): + """ + Test that large tool call arguments are split into smaller chunks (size 10) + to replicate OpenAI's native streaming behavior. + + This is especially important for providers like Bedrock that send complete + arguments at once, which need to be split to match OpenAI's token-by-token streaming. + """ + iterator = LiteLLMCompletionStreamingIterator( + model="test-model", + litellm_custom_stream_wrapper=AsyncMock(), + request_input="Test input", + responses_api_request={}, + ) + + # Create a chunk with a large arguments string that should be split + large_arguments = '{"param1": "value1", "param2": "value2", "param3": "value3"}' # 67 chars + chunk = ModelResponseStream( + id="chunk-1", + created=123, + model="test-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + role="assistant", + content="", + tool_calls=[ + { + "id": "call_test", + "type": "function", + "function": {"name": "test_function", "arguments": large_arguments}, + } + ], + ), + ) + ], + ) + + # Process the chunk once - it queues all events internally + evt = iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk) + + # First event should be OUTPUT_ITEM_ADDED + assert evt is not None + assert evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED + assert evt.output_index == 1 + assert hasattr(evt, '__dict__') and 'sequence_number' in evt.__dict__ + + # Collect all remaining delta events from the pending queue by creating empty chunks + delta_events = [] + empty_chunk = ModelResponseStream( + id="chunk-1", + created=123, + model="test-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(role="assistant", content=""), + ) + ], + ) + + # Keep draining pending events (expected: ceil(67 / 10) = 7 delta events) + while iterator._pending_tool_events: + evt = iterator._transform_chat_completion_chunk_to_response_api_chunk(empty_chunk) + if evt and evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA: + delta_events.append(evt) + + # Verify multiple delta events were created (at least 6 chunks for 67 chars) + assert len(delta_events) >= 6 # 67 chars split into chunks of max 10 chars each + + # Verify each delta is at most 10 characters + for evt in delta_events: + assert len(evt.delta) <= 10 + assert evt.item_id == "call_test" + assert evt.output_index == 1 + assert hasattr(evt, '__dict__') and 'sequence_number' in evt.__dict__ + + # Verify all deltas concatenated equal the original arguments + concatenated = ''.join(evt.delta for evt in delta_events) + assert concatenated == large_arguments + + # Verify sequence numbers are increasing + sequence_numbers = [evt.__dict__['sequence_number'] for evt in delta_events] + assert sequence_numbers == sorted(sequence_numbers) + assert len(set(sequence_numbers)) == len(sequence_numbers) # All unique diff --git a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py index 03a749a808..bc1f4fb72b 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -177,3 +177,155 @@ async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch): assert first_call["stream"] is False assert second_call["messages"] == ["follow-up"] assert second_call["stream"] is True + + +@pytest.mark.asyncio +async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch): + """ + Test that acompletion_with_mcp adds MCP metadata to CustomStreamWrapper + and it appears in the final chunk's delta.provider_specific_fields. + """ + from litellm.utils import CustomStreamWrapper + from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta + from litellm.litellm_core_utils.litellm_logging import Logging + + tools = [{"type": "mcp", "server_url": "litellm_proxy/mcp/local"}] + openai_tools = [{"type": "function", "function": {"name": "local_search"}}] + tool_calls = [{"id": "call-1", "type": "function", "function": {"name": "local_search"}}] + tool_results = [{"tool_call_id": "call-1", "result": "executed"}] + + # Create mock streaming chunks + def create_chunk(content, finish_reason=None): + return ModelResponseStream( + id="test-stream", + model="test-model", + created=1234567890, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=content, + role="assistant", + ), + finish_reason=finish_reason, + ) + ], + ) + + chunks = [ + create_chunk("Hello"), + create_chunk(" world", finish_reason="stop"), # Final chunk + ] + + # Create a proper CustomStreamWrapper + from unittest.mock import MagicMock + logging_obj = MagicMock() + logging_obj.model_call_details = {} + + class MockStreamingResponse(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="test-model", + logging_obj=logging_obj, + ) + self.chunks = chunks + self._index = 0 + self.sent_last_chunk = False + + def __iter__(self): + return self + + def __next__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + if self._index == len(self.chunks): + self.sent_last_chunk = True + # Call the method that adds MCP metadata to final chunk + chunk = self._add_mcp_metadata_to_final_chunk(chunk) + return chunk + raise StopIteration + + mock_acompletion = AsyncMock(return_value=MockStreamingResponse()) + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_should_use_litellm_mcp_gateway", + staticmethod(lambda tools: True), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_parse_mcp_tools", + staticmethod(lambda tools: (tools, [])), + ) + async def mock_process(**_): + return (tools, {"local_search": "local"}) + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + mock_process, + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_transform_mcp_tools_to_openai", + staticmethod(lambda *_, **__: openai_tools), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_should_auto_execute_tools", + staticmethod(lambda **_: False), + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(lambda **_: (None, None, None, None)), + ) + + with patch("litellm.acompletion", mock_acompletion): + result = await acompletion_with_mcp( + model="test-model", + messages=[{"role": "user", "content": "hello"}], + tools=tools, + stream=True, + ) + + # Verify result is CustomStreamWrapper + assert isinstance(result, CustomStreamWrapper) + + # Verify _hidden_params contains mcp_metadata + assert hasattr(result, "_hidden_params") + assert "mcp_metadata" in result._hidden_params + mcp_metadata = result._hidden_params["mcp_metadata"] + assert "mcp_list_tools" in mcp_metadata + assert mcp_metadata["mcp_list_tools"] == openai_tools + + # Consume the stream and check final chunk + all_chunks = list(result) + assert len(all_chunks) > 0 + + # Find the final chunk (with finish_reason) + final_chunk = None + for chunk in all_chunks: + if hasattr(chunk, "choices") and chunk.choices: + choice = chunk.choices[0] + if hasattr(choice, "finish_reason") and choice.finish_reason: + final_chunk = chunk + break + + # If no chunk with finish_reason, use the last chunk + if final_chunk is None and all_chunks: + final_chunk = all_chunks[-1] + + assert final_chunk is not None, "Should have a final chunk" + + # Verify MCP metadata is in the final chunk's delta.provider_specific_fields + if hasattr(final_chunk, "choices") and final_chunk.choices: + choice = final_chunk.choices[0] + if hasattr(choice, "delta") and choice.delta: + provider_fields = getattr(choice.delta, "provider_specific_fields", None) + assert provider_fields is not None, "Final chunk should have provider_specific_fields" + assert "mcp_list_tools" in provider_fields, "Should have mcp_list_tools" + assert provider_fields["mcp_list_tools"] == openai_tools diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index b632e72f56..15fdc7bd0c 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -1,12 +1,15 @@ import sys import types -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import HTTPException +import importlib from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) +from typing import Any, cast from litellm.types.utils import ModelResponse from litellm.types.responses.main import OutputFunctionToolCall @@ -22,7 +25,9 @@ def _setup_mcp_call_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) fake_manager = types.SimpleNamespace( - call_tool=AsyncMock(return_value=_DummyMCPResult()) + call_tool=AsyncMock(return_value=_DummyMCPResult()), + # Newer logging path calls this to enrich spend logs metadata + _get_mcp_server_from_tool_name=MagicMock(return_value=None), ) monkeypatch.setattr( "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", @@ -31,6 +36,15 @@ def _setup_mcp_call_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: return fake_manager.call_tool +def _setup_proxy_logging(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + """Patch proxy_logging_obj so failure hook can be asserted.""" + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock() + proxy_module = types.SimpleNamespace(proxy_logging_obj=proxy_logging_obj) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) + return proxy_logging_obj.post_call_failure_hook + + def test_deduplicate_mcp_tools_single_allowed_server(): tools = [{"name": "search"}, {"name": "search"}] # duplicate on purpose @@ -184,7 +198,7 @@ def test_create_follow_up_input_handles_response_function_tool_call(): ) follow_up = LiteLLM_Proxy_MCP_Handler._create_follow_up_input( - response=response, + response=cast(Any, response), tool_results=[], original_input=None, ) @@ -216,6 +230,8 @@ async def test_execute_tool_calls_strips_server_prefix(monkeypatch): user_api_key_auth=None, ) + assert call_tool_mock.await_count == 1 + assert call_tool_mock.await_args is not None assert call_tool_mock.await_args.kwargs["name"] == "read_wiki_structure" @@ -236,6 +252,8 @@ async def test_execute_tool_calls_keeps_tool_name_without_prefix(monkeypatch): user_api_key_auth=None, ) + assert call_tool_mock.await_count == 1 + assert call_tool_mock.await_args is not None assert call_tool_mock.await_args.kwargs["name"] == tool_name @@ -256,4 +274,131 @@ async def test_execute_tool_calls_keeps_tool_name_when_equal_to_server(monkeypat user_api_key_auth=None, ) + assert call_tool_mock.await_count == 1 + assert call_tool_mock.await_args is not None assert call_tool_mock.await_args.kwargs["name"] == tool_name + + +@pytest.mark.asyncio +async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkeypatch): + """ + Regression test for ae4d92ad...: + Ensure responses-side MCP tool execution logs failures via proxy_logging_obj.post_call_failure_hook. + """ + post_call_failure_hook = _setup_proxy_logging(monkeypatch) + + fake_manager = types.SimpleNamespace( + call_tool=AsyncMock( + side_effect=HTTPException(status_code=500, detail="boom") + ) + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + + tool_name = "deepwiki-read_wiki_structure" + tool_calls = [ + {"id": "call-err", "function": {"name": tool_name, "arguments": "{}"}} + ] + + user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user") + + results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=tool_calls, + user_api_key_auth=user_auth, + litellm_call_id="cid", + litellm_trace_id="tid", + ) + + assert len(results) == 1 + assert results[0]["tool_call_id"] == "call-err" + assert results[0]["name"] == tool_name + + post_call_failure_hook.assert_awaited_once() + assert post_call_failure_hook.await_args is not None + assert ( + post_call_failure_hook.await_args.kwargs.get("route") + == "/responses/mcp/call_tool" + ) + + +@pytest.mark.asyncio +async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_function_setup( + monkeypatch, +): + """ + Regression test for ae4d92ad...: + Ensure litellm_call_id / litellm_trace_id are forwarded into function_setup kwargs. + """ + _setup_proxy_logging(monkeypatch) + call_tool_mock = _setup_mcp_call_environment(monkeypatch) + + captured = {} + + def fake_function_setup(*_args, **kwargs): + captured.update(kwargs) + return None, None + + # NOTE: Don't patch via dotted string path here because `litellm.responses` + # is a function attribute on the `litellm` package (shadowing the submodule), + # which breaks monkeypatch's importpath resolution. + handler_module = importlib.import_module( + "litellm.responses.mcp.litellm_proxy_mcp_handler" + ) + monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) + + tool_name = "deepwiki-read_wiki_structure" + tool_calls = [ + {"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}} + ] + + await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=tool_calls, + user_api_key_auth=None, + litellm_call_id="cid", + litellm_trace_id="tid", + ) + + # Ensure the tool call was attempted (sanity) + assert call_tool_mock.await_count == 1 + + assert captured.get("litellm_call_id") == "cid" + assert captured.get("litellm_trace_id") == "tid" + + +@pytest.mark.asyncio +async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch): + """ + Regression test for 872e5b98...: + Ensure responses-side tool discovery enables list-tools SpendLogs logging flags. + """ + mock_get_tools = AsyncMock(return_value=[]) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.server._get_tools_from_mcp_servers", + mock_get_tools, + ) + + # Patch manager methods used by _get_mcp_tools_from_manager to avoid needing full UserAPIKeyAuth fields. + fake_manager = types.SimpleNamespace( + get_allowed_mcp_servers=AsyncMock(return_value=[]), + get_mcp_servers_from_ids=MagicMock(return_value=[]), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + + user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user") + tools, _server_names = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( + user_api_key_auth=user_auth, + mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}], + ) + + assert tools == [] + assert mock_get_tools.await_count == 1 + assert mock_get_tools.await_args is not None + assert mock_get_tools.await_args.kwargs["log_list_tools_to_spendlogs"] is True + assert mock_get_tools.await_args.kwargs["list_tools_log_source"] == "responses" diff --git a/tests/test_litellm/test_utils_custom.py b/tests/test_litellm/test_utils_custom.py deleted file mode 100644 index 3e924e9c71..0000000000 --- a/tests/test_litellm/test_utils_custom.py +++ /dev/null @@ -1,45 +0,0 @@ -import pytest -import sys -from unittest.mock import MagicMock, patch, AsyncMock -from litellm.proxy.utils import count_tokens_with_anthropic_api, _anthropic_async_clients - -@pytest.mark.asyncio -async def test_count_tokens_caching(): - """ - Test that count_tokens_with_anthropic_api caches the client. - """ - # Clear cache - _anthropic_async_clients.clear() - - api_key = "sk-ant-test-key" - messages = [{"role": "user", "content": "hello"}] - model = "claude-3-opus-20240229" - - # Create a mock anthropic module - mock_anthropic = MagicMock() - mock_client = MagicMock() - mock_anthropic.AsyncAnthropic.return_value = mock_client - - # Mock response - mock_response = MagicMock() - mock_response.input_tokens = 10 - - # Setup async return for count_tokens - mock_client.beta.messages.count_tokens = AsyncMock(return_value=mock_response) - - # Patch sys.modules to ensure our mock is used when anthropic is imported - with patch.dict(sys.modules, {"anthropic": mock_anthropic}): - # First call - with patch.dict("os.environ", {"ANTHROPIC_API_KEY": api_key}): - await count_tokens_with_anthropic_api(model, messages) - - assert api_key in _anthropic_async_clients - assert _anthropic_async_clients[api_key] == mock_client - mock_anthropic.AsyncAnthropic.assert_called_once() # Should be called once - - # Second call - with patch.dict("os.environ", {"ANTHROPIC_API_KEY": api_key}): - await count_tokens_with_anthropic_api(model, messages) - - # Should still be called once (cached) - mock_anthropic.AsyncAnthropic.assert_called_once() diff --git a/ui/litellm-dashboard/e2e_tests/constants.ts b/ui/litellm-dashboard/e2e_tests/constants.ts index b07bd68fcf..0de1ab571d 100644 --- a/ui/litellm-dashboard/e2e_tests/constants.ts +++ b/ui/litellm-dashboard/e2e_tests/constants.ts @@ -1 +1,2 @@ export const ADMIN_STORAGE_PATH = "admin.storageState.json"; +export const INTERNAL_USER_VIEWER_STORAGE_PATH = "internalViewer.storageState.json"; \ No newline at end of file diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/users.ts b/ui/litellm-dashboard/e2e_tests/fixtures/users.ts index d1f1eab00e..64d0597ee7 100644 --- a/ui/litellm-dashboard/e2e_tests/fixtures/users.ts +++ b/ui/litellm-dashboard/e2e_tests/fixtures/users.ts @@ -7,4 +7,8 @@ export const users = { email: "admin", password: isCI ? "gm" : "sk-1234", }, + [Role.InternalUserViewer]: { + email: "internalViewer@test.com", + password: "test", + }, }; diff --git a/ui/litellm-dashboard/e2e_tests/globalSetup.ts b/ui/litellm-dashboard/e2e_tests/globalSetup.ts index a725c58f35..e75362a2b3 100644 --- a/ui/litellm-dashboard/e2e_tests/globalSetup.ts +++ b/ui/litellm-dashboard/e2e_tests/globalSetup.ts @@ -1,17 +1,30 @@ import { chromium } from "@playwright/test"; import { users } from "./fixtures/users"; import { Role } from "./fixtures/roles"; +import { ADMIN_STORAGE_PATH, INTERNAL_USER_VIEWER_STORAGE_PATH } from "./constants"; + +async function loginAndSaveState( + browser, + user, + storagePath: string +) { + const context = await browser.newContext(); + const page = await context.newPage(); + + await page.goto("http://localhost:4000/ui/login"); + await page.getByPlaceholder("Enter your username").fill(user.email); + await page.getByPlaceholder("Enter your password").fill(user.password); + await page.getByRole("button", { name: "Login" }).click(); + await page.getByText('AI GATEWAY').waitFor(); + + await context.storageState({ path: storagePath }); + await context.close(); +} async function globalSetup() { const browser = await chromium.launch(); - const page = await browser.newPage(); - await page.goto("http://localhost:4000/ui/login"); - await page.getByPlaceholder("Enter your username").fill(users[Role.ProxyAdmin].email); - await page.getByPlaceholder("Enter your password").fill(users[Role.ProxyAdmin].password); - const loginButton = page.getByRole("button", { name: "Login" }); - await loginButton.click(); - await page.waitForSelector("text=AI Gateway"); - await page.context().storageState({ path: "admin.storageState.json" }); + await loginAndSaveState(browser, users[Role.ProxyAdmin], ADMIN_STORAGE_PATH); + await loginAndSaveState(browser, users[Role.InternalUserViewer], INTERNAL_USER_VIEWER_STORAGE_PATH); await browser.close(); } diff --git a/ui/litellm-dashboard/e2e_tests/tests/keys/createKey.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/keys/createKey.spec.ts new file mode 100644 index 0000000000..4343063b30 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/keys/createKey.spec.ts @@ -0,0 +1,22 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; + +test.describe("Create Key", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Able to create a key with all team models", async ({ page }) => { + await navigateToPage(page, Page.ApiKeys); + await expect(page.getByRole("button", { name: "Next" })).toBeVisible(); + await page.getByRole("button", { name: "+ Create New Key" }).click(); + await page.getByTestId("base-input").click(); + await page.getByTestId("base-input").fill("e2eUITestingCreateKeyAllTeamModels"); + await page.locator(".ant-select-selection-overflow").click(); + await page.getByText("All Team Models").click(); + await page.getByRole("combobox", { name: "* Models info-circle :" }).press("Escape"); + await page.getByRole("button", { name: "Create Key" }).click(); + await page.keyboard.press("Escape"); + await expect(page.getByText("e2eUITestingCreateKeyAllTeamModels")).toBeVisible(); + }); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts index 5ac977ff0c..5e24780dd0 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts @@ -3,7 +3,7 @@ import { users } from "../../fixtures/users"; import { Role } from "../../fixtures/roles"; test("user can log in", async ({ page }) => { - await page.goto("http://localhost:4000/ui/login"); + await page.goto("/ui/login"); await page.getByPlaceholder("Enter your username").fill(users[Role.ProxyAdmin].email); await page.getByPlaceholder("Enter your password").fill(users[Role.ProxyAdmin].password); const loginButton = page.getByRole("button", { name: "Login" }); diff --git a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts index ce07cc2b83..3d78362127 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts @@ -1,6 +1,6 @@ import test, { expect } from "@playwright/test"; import { Role } from "../../fixtures/roles"; -import { ADMIN_STORAGE_PATH } from "../../constants"; +import { ADMIN_STORAGE_PATH, INTERNAL_USER_VIEWER_STORAGE_PATH } from "../../constants"; import { Page } from "../../fixtures/pages"; import { menuLabelToPage } from "../../fixtures/menuMappings"; import { navigateToPage } from "../../helpers/navigation"; @@ -8,17 +8,28 @@ import { navigateToPage } from "../../helpers/navigation"; const sidebarButtons = { [Role.ProxyAdmin]: [ "Virtual Keys", + "MCP Servers", "Playground", "Models", "Usage", + "Logs", "Teams", "Internal Users", "API Reference", "AI Hub", ], + [Role.InternalUserViewer]: [ + "Virtual Keys", + "MCP Servers", + "Usage", + "Teams", + "Logs", + "API Reference", + "AI Hub", + ], }; -const roles = [{ role: Role.ProxyAdmin, storage: ADMIN_STORAGE_PATH }]; +const roles = [{ role: Role.ProxyAdmin, storage: ADMIN_STORAGE_PATH }, { role: Role.InternalUserViewer, storage: INTERNAL_USER_VIEWER_STORAGE_PATH }]; for (const { role, storage } of roles) { test.describe(`${role} sidebar`, () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts index 16ae03044c..1643412d1e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts @@ -2,7 +2,7 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import { renderHook, waitFor } from "@testing-library/react"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import React, { ReactNode } from "react"; -import { useKeys } from "./useKeys"; +import { useKeys, useDeletedKeys } from "./useKeys"; import type { KeyResponse } from "@/components/key_team_helpers/key_list"; // Mock the networking utilities @@ -397,3 +397,293 @@ describe("useKeys", () => { ); }); }); + +describe("useDeletedKeys", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + // Reset fetch mock + mockFetch.mockClear(); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return deleted keys data when query is successful", async () => { + // Mock successful API call + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => mockKeysResponse, + }); + + const { result } = renderHook(() => useDeletedKeys(1, 10), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockKeysResponse); + expect(result.current.error).toBeNull(); + expect(mockFetch).toHaveBeenCalledTimes(1); + expect(mockFetch).toHaveBeenCalledWith( + "/key/list?page=1&size=10&status=deleted&return_full_object=true&include_team_keys=true&include_created_by_keys=true", + { + method: "GET", + headers: { + Authorization: "Bearer test-access-token", + "Content-Type": "application/json", + }, + }, + ); + }); + + it("should pass status=deleted parameter to the API", async () => { + // Mock successful API call + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => mockKeysResponse, + }); + + const { result } = renderHook(() => useDeletedKeys(1, 10), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + // Verify that status=deleted is included in the URL + const callUrl = mockFetch.mock.calls[0][0]; + expect(callUrl).toContain("status=deleted"); + expect(result.current.data).toEqual(mockKeysResponse); + }); + + it("should handle error when deleted keys API call fails", async () => { + const errorMessage = "Failed to fetch deleted keys"; + const errorResponse = { error: errorMessage }; + + // Mock failed API call + mockFetch.mockResolvedValueOnce({ + ok: false, + json: async () => errorResponse, + }); + + const { result } = renderHook(() => useDeletedKeys(1, 10), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toBeDefined(); + expect(result.current.error?.message).toBe(errorMessage); + expect(result.current.data).toBeUndefined(); + expect(mockFetch).toHaveBeenCalledTimes(1); + expect(mockFetch).toHaveBeenCalledWith( + "/key/list?page=1&size=10&status=deleted&return_full_object=true&include_team_keys=true&include_created_by_keys=true", + { + method: "GET", + headers: { + Authorization: "Bearer test-access-token", + "Content-Type": "application/json", + }, + }, + ); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: "Admin", + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useDeletedKeys(1, 10), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(mockFetch).not.toHaveBeenCalled(); + }); + + it("should pass correct page and pageSize parameters to the API", async () => { + // Mock successful API call + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => mockKeysResponse, + }); + + const page = 2; + const pageSize = 20; + + const { result } = renderHook(() => useDeletedKeys(page, pageSize), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(mockFetch).toHaveBeenCalledWith( + `/key/list?page=${page}&size=${pageSize}&status=deleted&return_full_object=true&include_team_keys=true&include_created_by_keys=true`, + { + method: "GET", + headers: { + Authorization: "Bearer test-access-token", + "Content-Type": "application/json", + }, + }, + ); + }); + + it("should return empty deleted keys array when API returns empty data", async () => { + // Mock API returning empty keys array + const emptyResponse = { + keys: [], + total_count: 0, + current_page: 1, + total_pages: 0, + }; + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => emptyResponse, + }); + + const { result } = renderHook(() => useDeletedKeys(1, 10), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(emptyResponse); + expect(mockFetch).toHaveBeenCalledWith( + "/key/list?page=1&size=10&status=deleted&return_full_object=true&include_team_keys=true&include_created_by_keys=true", + { + method: "GET", + headers: { + Authorization: "Bearer test-access-token", + "Content-Type": "application/json", + }, + }, + ); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + mockFetch.mockRejectedValueOnce(timeoutError); + + const { result } = renderHook(() => useDeletedKeys(1, 10), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); + + it("should handle pagination correctly", async () => { + const paginatedResponse = { + keys: [mockKeys[0]], // Only first key + total_count: 15, + current_page: 2, + total_pages: 2, + }; + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => paginatedResponse, + }); + + const { result } = renderHook(() => useDeletedKeys(2, 10), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(result.current.data).toEqual(paginatedResponse); + expect(mockFetch).toHaveBeenCalledWith( + "/key/list?page=2&size=10&status=deleted&return_full_object=true&include_team_keys=true&include_created_by_keys=true", + { + method: "GET", + headers: { + Authorization: "Bearer test-access-token", + "Content-Type": "application/json", + }, + }, + ); + }); + + it("should pass additional options along with status=deleted", async () => { + // Mock successful API call + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => mockKeysResponse, + }); + + const options = { + organizationID: "org-1", + teamID: "team-1", + selectedKeyAlias: "test-alias", + }; + + const { result } = renderHook(() => useDeletedKeys(1, 10, options), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + const callUrl = mockFetch.mock.calls[0][0]; + expect(callUrl).toContain("status=deleted"); + expect(callUrl).toContain("organization_id=org-1"); + expect(callUrl).toContain("team_id=team-1"); + expect(callUrl).toContain("key_alias=test-alias"); + expect(result.current.data).toEqual(mockKeysResponse); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts index 73daf954c1..cf477a2e55 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts @@ -101,12 +101,16 @@ const keyListCall = async ( } }; -export const useKeys = (page: number, pageSize: number): UseQueryResult => { +export const useKeys = ( + page: number, + pageSize: number, + options: KeyListCallOptions = {}, +): UseQueryResult => { const { accessToken } = useAuthorized(); return useQuery({ - queryKey: keyKeys.list({ page, limit: pageSize }), - queryFn: async () => await keyListCall(accessToken!, page, pageSize), + queryKey: keyKeys.list({ page, limit: pageSize, ...options }), + queryFn: async () => await keyListCall(accessToken!, page, pageSize, options), enabled: Boolean(accessToken), staleTime: 30000, // 30 seconds placeholderData: keepPreviousData, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts new file mode 100644 index 0000000000..c6629e5396 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts @@ -0,0 +1,631 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { renderHook, waitFor } from "@testing-library/react"; +import React, { ReactNode } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { + useModelsInfo, + useModelHub, + useAllProxyModels, + useSelectedTeamModels, + type ProxyModel, + type AllProxyModelsResponse, + type PaginatedModelInfoResponse, +} from "./useModels"; + +vi.mock("@/components/networking", () => ({ + modelInfoCall: vi.fn(), + modelHubCall: vi.fn(), + modelAvailableCall: vi.fn(), +})); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +import { modelInfoCall, modelHubCall, modelAvailableCall } from "@/components/networking"; + +const mockProxyModel: ProxyModel = { + id: "model-1", + object: "model", + created: 1234567890, + owned_by: "openai", +}; + +const mockPaginatedModelInfoResponse: PaginatedModelInfoResponse = { + data: [{ id: "model-1", name: "Test Model" }], + total_count: 1, + current_page: 1, + total_pages: 1, + size: 50, +}; + +const mockAllProxyModelsResponse: AllProxyModelsResponse = { + data: [mockProxyModel], +}; + +describe("useModelsInfo", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + vi.clearAllMocks(); + + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should render without crashing", () => { + (modelInfoCall as any).mockResolvedValue(mockPaginatedModelInfoResponse); + + const { result } = renderHook(() => useModelsInfo(), { wrapper }); + + expect(result.current).toBeDefined(); + }); + + it("should return models data when query is successful", async () => { + (modelInfoCall as any).mockResolvedValue(mockPaginatedModelInfoResponse); + + const { result } = renderHook(() => useModelsInfo(), { wrapper }); + + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockPaginatedModelInfoResponse); + expect(result.current.error).toBeNull(); + expect(modelInfoCall).toHaveBeenCalledWith( + "test-access-token", + "test-user-id", + "Admin", + 1, + 50 + ); + expect(modelInfoCall).toHaveBeenCalledTimes(1); + }); + + it("should use custom page and size parameters", async () => { + (modelInfoCall as any).mockResolvedValue(mockPaginatedModelInfoResponse); + + const { result } = renderHook(() => useModelsInfo(2, 25), { wrapper }); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(modelInfoCall).toHaveBeenCalledWith( + "test-access-token", + "test-user-id", + "Admin", + 2, + 25 + ); + }); + + it("should handle error when modelInfoCall fails", async () => { + const errorMessage = "Failed to fetch models"; + const testError = new Error(errorMessage); + + (modelInfoCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useModelsInfo(), { wrapper }); + + expect(result.current.isLoading).toBe(true); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(modelInfoCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: "test-user-id", + userRole: "Admin", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useModelsInfo(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelInfoCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userId is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: null, + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useModelsInfo(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelInfoCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: null, + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useModelsInfo(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelInfoCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when all required auth values are missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: null, + userRole: null, + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useModelsInfo(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelInfoCall).not.toHaveBeenCalled(); + }); +}); + +describe("useModelHub", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + vi.clearAllMocks(); + + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should render without crashing", () => { + (modelHubCall as any).mockResolvedValue({ data: [] }); + + const { result } = renderHook(() => useModelHub(), { wrapper }); + + expect(result.current).toBeDefined(); + }); + + it("should return model hub data when query is successful", async () => { + const mockHubData = { data: [{ id: "hub-1", name: "Test Hub" }] }; + (modelHubCall as any).mockResolvedValue(mockHubData); + + const { result } = renderHook(() => useModelHub(), { wrapper }); + + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockHubData); + expect(result.current.error).toBeNull(); + expect(modelHubCall).toHaveBeenCalledWith("test-access-token"); + expect(modelHubCall).toHaveBeenCalledTimes(1); + }); + + it("should handle error when modelHubCall fails", async () => { + const errorMessage = "Failed to fetch model hub"; + const testError = new Error(errorMessage); + + (modelHubCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useModelHub(), { wrapper }); + + expect(result.current.isLoading).toBe(true); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(modelHubCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: "test-user-id", + userRole: "Admin", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useModelHub(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelHubCall).not.toHaveBeenCalled(); + }); +}); + +describe("useAllProxyModels", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + vi.clearAllMocks(); + + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should render without crashing", () => { + (modelAvailableCall as any).mockResolvedValue(mockAllProxyModelsResponse); + + const { result } = renderHook(() => useAllProxyModels(), { wrapper }); + + expect(result.current).toBeDefined(); + }); + + it("should return all proxy models data when query is successful", async () => { + (modelAvailableCall as any).mockResolvedValue(mockAllProxyModelsResponse); + + const { result } = renderHook(() => useAllProxyModels(), { wrapper }); + + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockAllProxyModelsResponse); + expect(result.current.error).toBeNull(); + expect(modelAvailableCall).toHaveBeenCalledWith( + "test-access-token", + "test-user-id", + "Admin", + true, + null, + true, + false, + "expand" + ); + expect(modelAvailableCall).toHaveBeenCalledTimes(1); + }); + + it("should handle error when modelAvailableCall fails", async () => { + const errorMessage = "Failed to fetch proxy models"; + const testError = new Error(errorMessage); + + (modelAvailableCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useAllProxyModels(), { wrapper }); + + expect(result.current.isLoading).toBe(true); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(modelAvailableCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: "test-user-id", + userRole: "Admin", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAllProxyModels(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelAvailableCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userId is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: null, + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAllProxyModels(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelAvailableCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: null, + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAllProxyModels(), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelAvailableCall).not.toHaveBeenCalled(); + }); +}); + +describe("useSelectedTeamModels", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + vi.clearAllMocks(); + + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should render without crashing", () => { + (modelAvailableCall as any).mockResolvedValue(mockAllProxyModelsResponse); + + const { result } = renderHook(() => useSelectedTeamModels("team-1"), { wrapper }); + + expect(result.current).toBeDefined(); + }); + + it("should return team models data when query is successful", async () => { + (modelAvailableCall as any).mockResolvedValue(mockAllProxyModelsResponse); + + const { result } = renderHook(() => useSelectedTeamModels("team-1"), { wrapper }); + + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockAllProxyModelsResponse); + expect(result.current.error).toBeNull(); + expect(modelAvailableCall).toHaveBeenCalledWith( + "test-access-token", + "test-user-id", + "Admin", + true, + "team-1" + ); + expect(modelAvailableCall).toHaveBeenCalledTimes(1); + }); + + it("should handle error when modelAvailableCall fails", async () => { + const errorMessage = "Failed to fetch team models"; + const testError = new Error(errorMessage); + + (modelAvailableCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useSelectedTeamModels("team-1"), { wrapper }); + + expect(result.current.isLoading).toBe(true); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(modelAvailableCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when teamID is null", () => { + const { result } = renderHook(() => useSelectedTeamModels(null), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelAvailableCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when accessToken is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: "test-user-id", + userRole: "Admin", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useSelectedTeamModels("team-1"), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelAvailableCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userId is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: null, + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useSelectedTeamModels("team-1"), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelAvailableCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: null, + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useSelectedTeamModels("team-1"), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelAvailableCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when teamID is missing and other auth values are present", () => { + const { result } = renderHook(() => useSelectedTeamModels(null), { wrapper }); + + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + expect(modelAvailableCall).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index fa7ab911ec..1dbd79eacf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -14,21 +14,31 @@ export interface AllProxyModelsResponse { data: ProxyModel[]; } +export interface PaginatedModelInfoResponse { + data: any[]; + total_count: number; + current_page: number; + total_pages: number; + size: number; +} + const modelKeys = createQueryKeys("models"); const modelHubKeys = createQueryKeys("modelHub"); const allProxyModelsKeys = createQueryKeys("allProxyModels"); const selectedTeamModelsKeys = createQueryKeys("selectedTeamModels"); -export const useModelsInfo = () => { +export const useModelsInfo = (page: number = 1, size: number = 50) => { const { accessToken, userId, userRole } = useAuthorized(); - return useQuery({ + return useQuery({ queryKey: modelKeys.list({ filters: { ...(userId && { userId }), ...(userRole && { userRole }), + page, + size, }, }), - queryFn: async () => await modelInfoCall(accessToken!, userId!, userRole!), + queryFn: async () => await modelInfoCall(accessToken!, userId!, userRole!, page, size), enabled: Boolean(accessToken && userId && userRole), }); }; @@ -46,7 +56,7 @@ export const useAllProxyModels = () => { const { accessToken, userId, userRole } = useAuthorized(); return useQuery({ queryKey: allProxyModelsKeys.list({}), - queryFn: async () => await modelAvailableCall(accessToken!, userId!, userRole!, true), + queryFn: async () => await modelAvailableCall(accessToken!, userId!, userRole!, true, null, true, false, "expand"), enabled: Boolean(accessToken && userId && userRole), }); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx index 1cce704467..f707ef04ff 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx @@ -94,7 +94,8 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te }, [modelDataResponse?.data]); const allModelsOnProxy = useMemo(() => { - return modelDataResponse?.data?.map((model: any) => model.model_name); + if (!modelDataResponse?.data) return []; + return modelDataResponse.data.map((model: any) => model.model_name); }, [modelDataResponse?.data]); const getProviderFromModel = (model: string) => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx index ae376701a9..8a2298361c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx @@ -4,10 +4,14 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import AllModelsTab from "./AllModelsTab"; // Mock the useModelsInfo hook -const mockUseModelsInfo = vi.fn(() => ({ data: { data: [] } })) as any; +const mockUseModelsInfo = vi.fn(() => ({ + data: { data: [], total_count: 0, current_page: 1, total_pages: 1, size: 50 }, + isLoading: false, + error: null, +})) as any; vi.mock("../../hooks/models/useModels", () => ({ - useModelsInfo: () => mockUseModelsInfo(), + useModelsInfo: (page?: number, size?: number) => mockUseModelsInfo(page, size), })); // Mock the useModelCostMap hook @@ -51,6 +55,21 @@ const createModelCostMapMock = (data: Record) => ({ error: null, }); +// Helper function to create paginated model data mock +const createPaginatedModelData = ( + models: any[], + totalCount: number = models.length, + currentPage: number = 1, + totalPages: number = 1, + size: number = 50, +) => ({ + data: models, + total_count: totalCount, + current_page: currentPage, + total_pages: totalPages, + size: size, +}); + describe("AllModelsTab", () => { const mockSetSelectedModelGroup = vi.fn(); const mockSetSelectedModelId = vi.fn(); @@ -84,7 +103,11 @@ describe("AllModelsTab", () => { }); it("should render with empty data", () => { - mockUseModelsInfo.mockReturnValueOnce({ data: { data: [] } }); + mockUseModelsInfo.mockReturnValueOnce({ + data: createPaginatedModelData([], 0, 1, 1, 50), + isLoading: false, + error: null, + }); mockUseTeams.mockReturnValueOnce({ data: [], @@ -130,28 +153,26 @@ describe("AllModelsTab", () => { }), ); - const modelData = { - data: [ - { - model_name: "gpt-4-accessible", - model_info: { - id: "model-1", - access_via_team_ids: ["team-456"], - access_groups: [], - }, + const modelData = createPaginatedModelData([ + { + model_name: "gpt-4-accessible", + model_info: { + id: "model-1", + access_via_team_ids: ["team-456"], + access_groups: [], }, - { - model_name: "gpt-3.5-turbo-blocked", - model_info: { - id: "model-2", - access_via_team_ids: ["team-789"], - access_groups: [], - }, + }, + { + model_name: "gpt-3.5-turbo-blocked", + model_info: { + id: "model-2", + access_via_team_ids: ["team-789"], + access_groups: [], }, - ], - }; + }, + ], 2, 1, 1, 50); - mockUseModelsInfo.mockReturnValue({ data: modelData }); + mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null }); render(); @@ -191,28 +212,26 @@ describe("AllModelsTab", () => { }), ); - const modelData = { - data: [ - { - model_name: "gpt-4-sales", - model_info: { - id: "model-sales-1", - access_via_team_ids: [], - access_groups: ["sales-model-group"], - }, + const modelData = createPaginatedModelData([ + { + model_name: "gpt-4-sales", + model_info: { + id: "model-sales-1", + access_via_team_ids: [], + access_groups: ["sales-model-group"], }, - { - model_name: "gpt-4-engineering", - model_info: { - id: "model-eng-1", - access_via_team_ids: [], - access_groups: ["engineering-model-group"], - }, + }, + { + model_name: "gpt-4-engineering", + model_info: { + id: "model-eng-1", + access_via_team_ids: [], + access_groups: ["engineering-model-group"], }, - ], - }; + }, + ], 2, 1, 1, 50); - mockUseModelsInfo.mockReturnValue({ data: modelData }); + mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null }); render(); @@ -236,30 +255,28 @@ describe("AllModelsTab", () => { }), ); - const modelData = { - data: [ - { - model_name: "gpt-4-personal", - model_info: { - id: "model-personal-1", - direct_access: true, - access_via_team_ids: [], - access_groups: [], - }, + const modelData = createPaginatedModelData([ + { + model_name: "gpt-4-personal", + model_info: { + id: "model-personal-1", + direct_access: true, + access_via_team_ids: [], + access_groups: [], }, - { - model_name: "gpt-4-team-only", - model_info: { - id: "model-team-1", - direct_access: false, - access_via_team_ids: ["team-123"], - access_groups: [], - }, + }, + { + model_name: "gpt-4-team-only", + model_info: { + id: "model-team-1", + direct_access: false, + access_via_team_ids: ["team-123"], + access_groups: [], }, - ], - }; + }, + ], 2, 1, 1, 50); - mockUseModelsInfo.mockReturnValue({ data: modelData }); + mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null }); render(); @@ -283,42 +300,40 @@ describe("AllModelsTab", () => { }), ); - const modelData = { - data: [ - { - model_name: "gpt-4-config", - litellm_model_name: "gpt-4-config", - provider: "openai", - model_info: { - id: "model-config-1", - db_model: false, - direct_access: true, - access_via_team_ids: [], - access_groups: [], - created_by: "user-123", - created_at: "2024-01-01", - updated_at: "2024-01-01", - }, + const modelData = createPaginatedModelData([ + { + model_name: "gpt-4-config", + litellm_model_name: "gpt-4-config", + provider: "openai", + model_info: { + id: "model-config-1", + db_model: false, + direct_access: true, + access_via_team_ids: [], + access_groups: [], + created_by: "user-123", + created_at: "2024-01-01", + updated_at: "2024-01-01", }, - { - model_name: "gpt-4-db", - litellm_model_name: "gpt-4-db", - provider: "openai", - model_info: { - id: "model-db-1", - db_model: true, - direct_access: true, - access_via_team_ids: [], - access_groups: [], - created_by: "user-123", - created_at: "2024-01-01", - updated_at: "2024-01-01", - }, + }, + { + model_name: "gpt-4-db", + litellm_model_name: "gpt-4-db", + provider: "openai", + model_info: { + id: "model-db-1", + db_model: true, + direct_access: true, + access_via_team_ids: [], + access_groups: [], + created_by: "user-123", + created_at: "2024-01-01", + updated_at: "2024-01-01", }, - ], - }; + }, + ], 2, 1, 1, 50); - mockUseModelsInfo.mockReturnValue({ data: modelData }); + mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null }); render(); @@ -342,27 +357,25 @@ describe("AllModelsTab", () => { }), ); - const modelData = { - data: [ - { - model_name: "gpt-4-config", - litellm_model_name: "gpt-4-config", - provider: "openai", - model_info: { - id: "model-config-1", - db_model: false, - direct_access: true, - access_via_team_ids: [], - access_groups: [], - created_by: "user-123", - created_at: "2024-01-01", - updated_at: "2024-01-01", - }, + const modelData = createPaginatedModelData([ + { + model_name: "gpt-4-config", + litellm_model_name: "gpt-4-config", + provider: "openai", + model_info: { + id: "model-config-1", + db_model: false, + direct_access: true, + access_via_team_ids: [], + access_groups: [], + created_by: "user-123", + created_at: "2024-01-01", + updated_at: "2024-01-01", }, - ], - }; + }, + ], 1, 1, 1, 50); - mockUseModelsInfo.mockReturnValue({ data: modelData }); + mockUseModelsInfo.mockReturnValue({ data: modelData, isLoading: false, error: null }); render(); @@ -370,4 +383,110 @@ describe("AllModelsTab", () => { expect(screen.getByText("Defined in config")).toBeInTheDocument(); }); }); + + it("should handle pagination: Previous button is disabled on first page and Next button works", async () => { + mockUseTeams.mockReturnValue({ + data: [], + isLoading: false, + error: null, + refetch: vi.fn(), + }); + + mockUseModelCostMap.mockReturnValue( + createModelCostMapMock({ + "gpt-4-page1": { litellm_provider: "openai" }, + "gpt-4-page2": { litellm_provider: "openai" }, + }), + ); + + // Mock first page response (page 1 of 2) + const page1Data = createPaginatedModelData( + [ + { + model_name: "gpt-4-page1", + model_info: { + id: "model-page1-1", + direct_access: true, + access_via_team_ids: [], + access_groups: [], + }, + }, + ], + 2, // total_count + 1, // current_page + 2, // total_pages + 50, // size + ); + + // Set up mock to return page1Data for page 1 + mockUseModelsInfo.mockImplementation((page: number = 1) => { + return { data: page1Data, isLoading: false, error: null }; + }); + + render(); + + await waitFor(() => { + expect(screen.getByText("Showing 1 - 1 of 1 results")).toBeInTheDocument(); + }); + + // Check that Previous button is disabled on first page + const previousButton = screen.getByRole("button", { name: /previous/i }); + expect(previousButton).toBeDisabled(); + + // Check that Next button is enabled (since we're on page 1 of 2) + const nextButton = screen.getByRole("button", { name: /next/i }); + expect(nextButton).not.toBeDisabled(); + }); + + it("should handle pagination: Next button is disabled on last page", async () => { + mockUseTeams.mockReturnValue({ + data: [], + isLoading: false, + error: null, + refetch: vi.fn(), + }); + + mockUseModelCostMap.mockReturnValue( + createModelCostMapMock({ + "gpt-4-page2": { litellm_provider: "openai" }, + }), + ); + + // Mock single page response (page 1 of 1 - last page) + const singlePageData = createPaginatedModelData( + [ + { + model_name: "gpt-4-page2", + model_info: { + id: "model-page2-1", + direct_access: true, + access_via_team_ids: [], + access_groups: [], + }, + }, + ], + 1, // total_count + 1, // current_page + 1, // total_pages (only 1 page, so this is the last page) + 50, // size + ); + + mockUseModelsInfo.mockImplementation(() => { + return { data: singlePageData, isLoading: false, error: null }; + }); + + render(); + + await waitFor(() => { + expect(screen.getByText("Showing 1 - 1 of 1 results")).toBeInTheDocument(); + }); + + // When there's only 1 page (last page), Next should be disabled + const nextButton = screen.getByRole("button", { name: /next/i }); + expect(nextButton).toBeDisabled(); + + // Previous should also be disabled on the first (and only) page + const previousButton = screen.getByRole("button", { name: /previous/i }); + expect(previousButton).toBeDisabled(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index a85e651658..04300b7fd1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -31,11 +31,26 @@ const AllModelsTab = ({ setSelectedModelId, setSelectedTeamId, }: AllModelsTabProps) => { - const { data: rawModelData, isLoading: isLoadingModelsInfo } = useModelsInfo(); const { data: modelCostMapData, isLoading: isLoadingModelCostMap } = useModelCostMap(); const { userId, userRole, premiumUser } = useAuthorized(); const { data: teams } = useTeams(); + const [modelNameSearch, setModelNameSearch] = useState(""); + const [modelViewMode, setModelViewMode] = useState("current_team"); + const [currentTeam, setCurrentTeam] = useState("personal"); + const [showFilters, setShowFilters] = useState(false); + const [selectedModelAccessGroupFilter, setSelectedModelAccessGroupFilter] = useState(null); + const [expandedRows, setExpandedRows] = useState>(new Set()); + const [currentPage, setCurrentPage] = useState(1); + const [pageSize] = useState(50); + const [pagination, setPagination] = useState({ + pageIndex: 0, + pageSize: 50, + }); + + const { data: rawModelData, isLoading: isLoadingModelsInfo } = useModelsInfo(currentPage, pageSize); + const isLoading = isLoadingModelsInfo || isLoadingModelCostMap; + const getProviderFromModel = (model: string) => { if (modelCostMapData !== null && modelCostMapData !== undefined) { if (typeof modelCostMapData == "object" && model in modelCostMapData) { @@ -50,18 +65,23 @@ const AllModelsTab = ({ return transformModelData(rawModelData, getProviderFromModel); }, [rawModelData, modelCostMapData]); - const [modelNameSearch, setModelNameSearch] = useState(""); - const [modelViewMode, setModelViewMode] = useState("current_team"); - const [currentTeam, setCurrentTeam] = useState("personal"); - const [showFilters, setShowFilters] = useState(false); - const [selectedModelAccessGroupFilter, setSelectedModelAccessGroupFilter] = useState(null); - const [expandedRows, setExpandedRows] = useState>(new Set()); - const [pagination, setPagination] = useState({ - pageIndex: 0, - pageSize: 50, - }); - - const isLoading = isLoadingModelsInfo || isLoadingModelCostMap; + // Get pagination metadata from the response + const paginationMeta = useMemo(() => { + if (!rawModelData) { + return { + total_count: 0, + current_page: 1, + total_pages: 1, + size: pageSize, + }; + } + return { + total_count: rawModelData.total_count ?? 0, + current_page: rawModelData.current_page ?? 1, + total_pages: rawModelData.total_pages ?? 1, + size: rawModelData.size ?? pageSize, + }; + }, [rawModelData, pageSize]); const filteredData = useMemo(() => { if (!modelData || !modelData.data || modelData.data.length === 0) { @@ -114,6 +134,7 @@ const AllModelsTab = ({ setSelectedModelAccessGroupFilter(null); setCurrentTeam("personal"); setModelViewMode("current_team"); + setCurrentPage(1); setPagination({ pageIndex: 0, pageSize: 50 }); }; @@ -334,10 +355,7 @@ const AllModelsTab = ({ ) : ( {filteredData.length > 0 - ? `Showing ${pagination.pageIndex * pagination.pageSize + 1} - ${Math.min( - (pagination.pageIndex + 1) * pagination.pageSize, - filteredData.length, - )} of ${filteredData.length} results` + ? `Showing 1 - ${filteredData.length} of ${filteredData.length} results` : "Showing 0 results"} )} @@ -347,15 +365,16 @@ const AllModelsTab = ({ ) : ( @@ -365,15 +384,16 @@ const AllModelsTab = ({ ) : ( @@ -391,8 +411,8 @@ const AllModelsTab = ({ setSelectedModelId, setSelectedTeamId, getDisplayModelName, - () => {}, - () => {}, + () => { }, + () => { }, expandedRows, setExpandedRows, )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx index df6d8d3ea8..972c39d49d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx @@ -526,7 +526,7 @@ const CreateTeamModal = ({ valuePropName="checked" help="Bypass global guardrails for this team" > - diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_cost_results.tsx b/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_cost_results.tsx index 0965a153a4..624c297f26 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_cost_results.tsx +++ b/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_cost_results.tsx @@ -113,10 +113,13 @@ const MultiCostResults: React.FC = ({ multiResult, timePe const validEntries = multiResult.entries.filter((e) => e.result !== null); const loadingEntries = multiResult.entries.filter((e) => e.loading); + const errorEntries = multiResult.entries.filter((e) => e.error !== null); const hasAnyResult = validEntries.length > 0; const isAnyLoading = loadingEntries.length > 0; + const hasAnyError = errorEntries.length > 0; - if (!hasAnyResult && !isAnyLoading) { + // Show empty state only if no results, not loading, and no errors + if (!hasAnyResult && !isAnyLoading && !hasAnyError) { return (
@@ -126,7 +129,8 @@ const MultiCostResults: React.FC = ({ multiResult, timePe ); } - if (!hasAnyResult && isAnyLoading) { + // Show loading state only if loading and no results/errors yet + if (!hasAnyResult && isAnyLoading && !hasAnyError) { return (
} /> @@ -135,6 +139,26 @@ const MultiCostResults: React.FC = ({ multiResult, timePe ); } + // Show errors-only view when there are errors but no valid results + if (!hasAnyResult && hasAnyError) { + return ( +
+ +
+ Cost Estimates + {isAnyLoading && } size="small" />} +
+ {/* Error Messages */} + {errorEntries.map((e) => ( +
+ {e.entry.model || "Unknown model"}: + {e.error} +
+ ))} +
+ ); + } + const toggleExpanded = (id: string) => { setExpandedModels((prev) => { const next = new Set(prev); @@ -157,13 +181,28 @@ const MultiCostResults: React.FC = ({ multiResult, timePe title: "Model", dataIndex: "model", key: "model", - render: (text: string, record: { id: string; provider?: string | null }) => ( -
- {text} - {record.provider && ( - - {record.provider} - + render: (text: string, record: { id: string; provider?: string | null; error?: string | null; loading?: boolean; hasZeroCost?: boolean }) => ( +
+
+ {text} + {record.provider && ( + + {record.provider} + + )} + {record.loading && ( + } size="small" /> + )} +
+ {record.error && ( +
+ ⚠️ {record.error} +
+ )} + {record.hasZeroCost && !record.error && ( +
+ ⚠️ No pricing data found for this model. Set base_model in config. +
)}
), @@ -173,17 +212,21 @@ const MultiCostResults: React.FC = ({ multiResult, timePe dataIndex: "cost_per_request", key: "cost_per_request", align: "right" as const, - render: (value: number) => {formatCost(value)}, + render: (value: number | null, record: { error?: string | null }) => ( + record.error ? - : {formatCost(value)} + ), }, { title: "Margin Fee", dataIndex: "margin_cost_per_request", key: "margin_cost_per_request", align: "right" as const, - render: (value: number) => ( - 0 ? "text-amber-600" : "text-gray-400"}`}> - {formatCost(value)} - + render: (value: number | null, record: { error?: string | null }) => ( + record.error ? - : ( + 0 ? "text-amber-600" : "text-gray-400"}`}> + {formatCost(value)} + + ) ), }, { @@ -191,34 +234,43 @@ const MultiCostResults: React.FC = ({ multiResult, timePe dataIndex: periodCostKey, key: "period_cost", align: "right" as const, - render: (value: number | null) => {formatCost(value)}, + render: (value: number | null, record: { error?: string | null }) => ( + record.error ? - : {formatCost(value)} + ), }, { title: "", key: "expand", width: 40, - render: (_: unknown, record: { id: string }) => ( - + render: (_: unknown, record: { id: string; error?: string | null }) => ( + record.error ? null : ( + + ) ), }, ]; - const summaryData = validEntries.map((e) => ({ + // Include both valid results and errors in the table data + const allEntriesWithModels = multiResult.entries.filter((e) => e.entry.model); + const summaryData = allEntriesWithModels.map((e) => ({ key: e.entry.id, id: e.entry.id, - model: e.result!.model, - provider: e.result!.provider, - cost_per_request: e.result!.cost_per_request, - margin_cost_per_request: e.result!.margin_cost_per_request, - daily_cost: e.result!.daily_cost, - monthly_cost: e.result!.monthly_cost, + model: e.result?.model || e.entry.model, + provider: e.result?.provider, + cost_per_request: e.result?.cost_per_request ?? null, + margin_cost_per_request: e.result?.margin_cost_per_request ?? null, + daily_cost: e.result?.daily_cost ?? null, + monthly_cost: e.result?.monthly_cost ?? null, + error: e.error, + loading: e.loading, + hasZeroCost: e.result && e.result.cost_per_request === 0, })); return ( @@ -268,7 +320,7 @@ const MultiCostResults: React.FC = ({ multiResult, timePe {/* Per-Model Table */} - {validEntries.length > 0 && ( + {summaryData.length > 0 && ( = ({ multiResult, timePe }} /> )} - - {/* Error Messages */} - {multiResult.entries - .filter((e) => e.error) - .map((e) => ( -
- {e.entry.model || "Unknown model"}: - {e.error} -
- ))} ); }; diff --git a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx index a3ddeff322..b388135044 100644 --- a/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx +++ b/ui/litellm-dashboard/src/components/ModelSelect/ModelSelect.tsx @@ -53,14 +53,13 @@ const contextFilters: Record { if (selectedOrganization) { - if (selectedOrganization.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value)) { + if (selectedOrganization.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value) || selectedOrganization.models.length === 0) { return allProxyModels; } - // Return organization's models (filtered from allProxyModels) return allProxyModels.filter((model) => selectedOrganization.models.includes(model)); } - return userModels ?? []; + return allProxyModels ?? []; }, organization: ({ allProxyModels, selectedOrganization, options }) => { @@ -102,9 +101,12 @@ export const ModelSelect = (props: ModelSelectProps) => { const isSpecialOption = (value: string) => MODEL_SELECT_SPECIAL_VALUES_ARRAY.some((sv) => sv.value === value); const hasSpecialOptionSelected = value.some(isSpecialOption); const isLoading = isLoadingAllProxyModels || isLoadingTeam || isLoadingOrganization || isCurrentUserLoading; + const organizationHasAllProxyModels = organization?.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value) || organization?.models.length === 0; + console.log("organization:", organization); + console.log("organizationHasAllProxyModels:", organizationHasAllProxyModels); const shouldShowAllProxyModels = showAllProxyModelsOverride || - (organization?.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value) && includeSpecialOptions); + (organizationHasAllProxyModels && includeSpecialOptions); if (isLoading) { return ; @@ -143,51 +145,51 @@ export const ModelSelect = (props: ModelSelectProps) => { options={[ includeSpecialOptions ? { - label: Special Options, - title: "Special Options", - options: [ - ...(shouldShowAllProxyModels - ? [ - { - label: All Proxy Models, - value: MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value, - disabled: - value.length > 0 && - value.some( - (v) => isSpecialOption(v) && v !== MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value, - ), - key: MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value, - }, - ] - : []), - { - label: No Default Models, - value: MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value, - disabled: - value.length > 0 && - value.some((v) => isSpecialOption(v) && v !== MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value), - key: MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value, - }, - ], - } + label: Special Options, + title: "Special Options", + options: [ + ...(shouldShowAllProxyModels + ? [ + { + label: All Proxy Models, + value: MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value, + disabled: + value.length > 0 && + value.some( + (v) => isSpecialOption(v) && v !== MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value, + ), + key: MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value, + }, + ] + : []), + { + label: No Default Models, + value: MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value, + disabled: + value.length > 0 && + value.some((v) => isSpecialOption(v) && v !== MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value), + key: MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value, + }, + ], + } : [], ...(wildcard.length > 0 ? [ - { - label: Wildcard Options, - title: "Wildcard Options", - options: wildcard.map((model) => { - const provider = model.replace("/*", ""); - const capitalizedProvider = provider.charAt(0).toUpperCase() + provider.slice(1); + { + label: Wildcard Options, + title: "Wildcard Options", + options: wildcard.map((model) => { + const provider = model.replace("/*", ""); + const capitalizedProvider = provider.charAt(0).toUpperCase() + provider.slice(1); - return { - label: {`All ${capitalizedProvider} models`}, - value: model, - disabled: hasSpecialOptionSelected, - }; - }), - }, - ] + return { + label: {`All ${capitalizedProvider} models`}, + value: model, + disabled: hasSpecialOptionSelected, + }; + }), + }, + ] : []), { label: Models, diff --git a/ui/litellm-dashboard/src/components/OldTeams.test.tsx b/ui/litellm-dashboard/src/components/OldTeams.test.tsx index 76fc26a884..91b428c8c9 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.test.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.test.tsx @@ -60,6 +60,31 @@ vi.mock("@/components/team/team_info", () => ({ }, })); +vi.mock("./ModelSelect/ModelSelect", () => { + const ModelSelect = React.forwardRef(({ value, onChange, dataTestId, id }: any, ref: any) => { + return ( + { + // Mock onChange - in real usage this would be handled by Ant Design Select + if (onChange) { + onChange(value || []); + } + }} + readOnly + /> + ); + }); + ModelSelect.displayName = "ModelSelect"; + return { + ModelSelect, + }; +}); + vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ useOrganizations: () => mockUseOrganizations(), })); @@ -313,6 +338,7 @@ describe("OldTeams - handleCreate organization handling", () => { created_at: new Date().toISOString(), keys: [], members_with_roles: [], + spend: 0, }, ]} searchParams={{}} @@ -387,6 +413,7 @@ describe("OldTeams - empty state", () => { created_at: new Date().toISOString(), keys: [], members_with_roles: [], + spend: 0, }, ]} searchParams={{}} @@ -555,6 +582,7 @@ describe("OldTeams - premium props", () => { created_at: new Date().toISOString(), keys: [], members_with_roles: [], + spend: 0, }, ]} searchParams={{}} @@ -603,6 +631,7 @@ describe("OldTeams - Default Team Settings tab visibility", () => { created_at: new Date().toISOString(), keys: [], members_with_roles: [], + spend: 0, }, ]} searchParams={{}} @@ -633,6 +662,7 @@ describe("OldTeams - Default Team Settings tab visibility", () => { created_at: new Date().toISOString(), keys: [], members_with_roles: [], + spend: 0, }, ]} searchParams={{}} @@ -663,6 +693,7 @@ describe("OldTeams - Default Team Settings tab visibility", () => { created_at: new Date().toISOString(), keys: [], members_with_roles: [], + spend: 0, }, ]} searchParams={{}} @@ -693,6 +724,7 @@ describe("OldTeams - Default Team Settings tab visibility", () => { created_at: new Date().toISOString(), keys: [], members_with_roles: [], + spend: 0, }, ]} searchParams={{}} @@ -791,6 +823,7 @@ describe("OldTeams - organization alias display", () => { created_at: new Date().toISOString(), keys: [], members_with_roles: [], + spend: 0, }, ]} searchParams={{}} @@ -824,6 +857,7 @@ describe("OldTeams - organization alias display", () => { created_at: new Date().toISOString(), keys: [], members_with_roles: [], + spend: 0, }, ]} searchParams={{}} @@ -856,6 +890,7 @@ describe("OldTeams - organization alias display", () => { created_at: new Date().toISOString(), keys: [], members_with_roles: [], + spend: 0, }, ]} searchParams={{}} diff --git a/ui/litellm-dashboard/src/components/OldTeams.tsx b/ui/litellm-dashboard/src/components/OldTeams.tsx index 5679788d30..1202cc9169 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.tsx @@ -85,6 +85,7 @@ import { updateExistingKeys } from "@/utils/dataUtils"; import DeleteResourceModal from "./common_components/DeleteResourceModal"; import TableIconActionButton from "./common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; import { Member, teamCreateCall, v2TeamListCall } from "./networking"; +import { ModelSelect } from "./ModelSelect/ModelSelect"; interface TeamInfo { members_with_roles: Member[]; @@ -1064,11 +1065,11 @@ const Teams: React.FC = ({ rules={ isOrgAdmin ? [ - { - required: true, - message: "Please select an organization", - }, - ] + { + required: true, + message: "Please select an organization", + }, + ] : [] } help={ @@ -1135,16 +1136,17 @@ const Teams: React.FC = ({ ]} name="models" > - - - No Default Models - - {modelsToPick.map((model) => ( - - {getModelDisplayName(model)} - - ))} - + form.setFieldValue("models", values)} + organizationID={form.getFieldValue("organization_id")} + options={{ + includeSpecialOptions: true, + showAllProxyModelsOverride: !form.getFieldValue("organization_id"), + }} + context="team" + dataTestId="create-team-models-select" + /> diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index c9d11c778f..6d96804792 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -71,12 +71,19 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo pageSize: 50, }); + // Extract sort parameters from sorting state + const sortBy = sorting.length > 0 ? sorting[0].id : null; + const sortOrder = sorting.length > 0 ? (sorting[0].desc ? "desc" : "asc") : null; + const { data: keys, isPending: isLoading, isFetching, refetch, - } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize); + } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, { + sortBy: sortBy || undefined, + sortOrder: sortOrder || undefined, + }); const totalCount = keys?.total_count || 0; const [expandedAccordions, setExpandedAccordions] = useState>({}); @@ -110,6 +117,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo id: "expander", header: () => null, size: 40, + enableSorting: false, cell: ({ row }) => row.getCanExpand() ? (