diff --git a/docs/my-website/docs/adding_provider/simple_guardrail_tutorial.md b/docs/my-website/docs/adding_provider/simple_guardrail_tutorial.md index 9c654cd1560..884a7397bde 100644 --- a/docs/my-website/docs/adding_provider/simple_guardrail_tutorial.md +++ b/docs/my-website/docs/adding_provider/simple_guardrail_tutorial.md @@ -101,12 +101,11 @@ model_list: - model_name: gpt-4 litellm_params: model: gpt-4 - api_key: os.environ/OPENAI_API_KEY + api_key: os.environ/OPENAI_API_KEY -litellm_settings: - guardrails: +guardrails: - guardrail_name: my_guardrail - litellm_params: + litellm_params: guardrail: my_guardrail mode: during_call api_key: os.environ/MY_GUARDRAIL_API_KEY diff --git a/docs/my-website/docs/providers/github_copilot.md b/docs/my-website/docs/providers/github_copilot.md index 306c9f949ec..e9fd3444f5f 100644 --- a/docs/my-website/docs/providers/github_copilot.md +++ b/docs/my-website/docs/providers/github_copilot.md @@ -35,11 +35,10 @@ from litellm import completion response = completion( model="github_copilot/gpt-4", - messages=[{"role": "user", "content": "Write a Python function to calculate fibonacci numbers"}], - extra_headers={ - "editor-version": "vscode/1.85.1", - "Copilot-Integration-Id": "vscode-chat" - } + messages=[ + {"role": "system", "content": "You are a helpful coding assistant"}, + {"role": "user", "content": "Write a Python function to calculate fibonacci numbers"} + ] ) print(response) ``` @@ -50,11 +49,7 @@ from litellm import completion stream = completion( model="github_copilot/gpt-4", messages=[{"role": "user", "content": "Explain async/await in Python"}], - stream=True, - extra_headers={ - "editor-version": "vscode/1.85.1", - "Copilot-Integration-Id": "vscode-chat" - } + stream=True ) for chunk in stream: @@ -134,11 +129,7 @@ client = OpenAI( # Non-streaming response response = client.chat.completions.create( model="github_copilot/gpt-4", - messages=[{"role": "user", "content": "How do I optimize this SQL query?"}], - extra_headers={ - "editor-version": "vscode/1.85.1", - "Copilot-Integration-Id": "vscode-chat" - } + messages=[{"role": "user", "content": "How do I optimize this SQL query?"}] ) print(response.choices[0].message.content) @@ -156,11 +147,7 @@ response = litellm.completion( model="litellm_proxy/github_copilot/gpt-4", messages=[{"role": "user", "content": "Review this code for bugs"}], api_base="http://localhost:4000", - api_key="your-proxy-api-key", - extra_headers={ - "editor-version": "vscode/1.85.1", - "Copilot-Integration-Id": "vscode-chat" - } + api_key="your-proxy-api-key" ) print(response.choices[0].message.content) @@ -174,8 +161,6 @@ print(response.choices[0].message.content) curl http://localhost:4000/v1/chat/completions \ -H "Content-Type: application/json" \ -H "Authorization: Bearer your-proxy-api-key" \ - -H "editor-version: vscode/1.85.1" \ - -H "Copilot-Integration-Id: vscode-chat" \ -d '{ "model": "github_copilot/gpt-4", "messages": [{"role": "user", "content": "Explain this error message"}] @@ -211,9 +196,11 @@ export GITHUB_COPILOT_API_KEY_FILE="api-key.json" ### Headers -GitHub Copilot supports various editor-specific headers: +LiteLLM automatically injects the required GitHub Copilot headers (simulating VSCode). You don't need to specify them manually. -```python showLineNumbers title="Common Headers" +If you want to override the defaults (e.g., to simulate a different editor), you can use `extra_headers`: + +```python showLineNumbers title="Custom Headers (Optional)" extra_headers = { "editor-version": "vscode/1.85.1", # Editor version "editor-plugin-version": "copilot/1.155.0", # Plugin version diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 385b4b0de32..5cdae51f448 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -94,7 +94,7 @@ litellm_settings: # /chat/completions, /completions, /embeddings, /audio/transcriptions mode: default_off # if default_off, you need to opt in to caching on a per call basis ttl: 600 # ttl for caching - disable_copilot_system_to_assistant: False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior. + disable_copilot_system_to_assistant: False # DEPRECATED - GitHub Copilot API supports system prompts. callback_settings: otel: @@ -197,7 +197,7 @@ router_settings: | disable_add_transform_inline_image_block | boolean | For Fireworks AI models - if true, turns off the auto-add of `#transform=inline` to the url of the image_url, if the model is not a vision model. | | disable_hf_tokenizer_download | boolean | If true, it defaults to using the openai tokenizer for all models (including huggingface models). | | enable_json_schema_validation | boolean | If true, enables json schema validation for all requests. | -| disable_copilot_system_to_assistant | boolean | If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior. Useful for tools (like Claude Code) that send system messages, which Copilot does not support. | +| disable_copilot_system_to_assistant | boolean | **DEPRECATED** - GitHub Copilot API supports system prompts. | ### general_settings - Reference diff --git a/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md b/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md new file mode 100644 index 00000000000..cb246144497 --- /dev/null +++ b/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md @@ -0,0 +1,278 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Custom Code Guardrail + +Write custom guardrail logic using Python-like code that runs in a sandboxed environment. + +## Quick Start + +### 1. Define the guardrail in config + +```yaml +model_list: + - model_name: gpt-4 + litellm_params: + model: gpt-4 + api_key: os.environ/OPENAI_API_KEY + +guardrails: + - guardrail_name: block-ssn + litellm_params: + guardrail: custom_code + mode: pre_call + custom_code: | + def apply_guardrail(inputs, request_data, input_type): + for text in inputs["texts"]: + if regex_match(text, r"\d{3}-\d{2}-\d{4}"): + return block("SSN detected") + return allow() +``` + +### 2. Start proxy + +```bash +litellm --config config.yaml +``` + +### 3. Test + +```bash +curl -X POST http://localhost:4000/chat/completions \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gpt-4", + "messages": [{"role": "user", "content": "My SSN is 123-45-6789"}], + "guardrails": ["block-ssn"] + }' +``` + +## Configuration + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `guardrail` | string | ✅ | Must be `custom_code` | +| `mode` | string | ✅ | When to run: `pre_call`, `post_call`, `during_call` | +| `custom_code` | string | ✅ | Python-like code with `apply_guardrail` function | +| `default_on` | bool | ❌ | Run on all requests (default: `false`) | + +## Writing Custom Code + +### Function Signature + +Your code must define an `apply_guardrail` function: + +```python +def apply_guardrail(inputs, request_data, input_type): + # inputs: see table below + # request_data: {"model": "...", "user_id": "...", "team_id": "...", "metadata": {...}} + # input_type: "request" or "response" + + return allow() # or block() or modify() +``` + +### `inputs` Parameter + +| Field | Type | Description | +|-------|------|-------------| +| `texts` | `List[str]` | Extracted text from the request/response | +| `images` | `List[str]` | Extracted images (for image guardrails) | +| `tools` | `List[dict]` | Tools sent to the LLM | +| `tool_calls` | `List[dict]` | Tool calls returned from the LLM | +| `structured_messages` | `List[dict]` | Full messages with role info (system/user/assistant) | +| `model` | `str` | The model being used | + +### `request_data` Parameter + +| Field | Type | Description | +|-------|------|-------------| +| `model` | `str` | Model name | +| `user_id` | `str` | User ID from API key | +| `team_id` | `str` | Team ID from API key | +| `end_user_id` | `str` | End user ID | +| `metadata` | `dict` | Request metadata | + +### Return Values + +| Function | Description | +|----------|-------------| +| `allow()` | Let request/response through | +| `block(reason)` | Reject with message | +| `modify(texts=[], images=[], tool_calls=[])` | Transform content | + +## Built-in Primitives + +### Regex + +| Function | Description | +|----------|-------------| +| `regex_match(text, pattern)` | Returns `True` if pattern found | +| `regex_replace(text, pattern, replacement)` | Replace all matches | +| `regex_find_all(text, pattern)` | Return list of matches | + +### JSON + +| Function | Description | +|----------|-------------| +| `json_parse(text)` | Parse JSON string, returns `None` on error | +| `json_stringify(obj)` | Convert to JSON string | +| `json_schema_valid(obj, schema)` | Validate against JSON schema | + +### URL + +| Function | Description | +|----------|-------------| +| `extract_urls(text)` | Extract all URLs from text | +| `is_valid_url(url)` | Check if URL is valid | +| `all_urls_valid(text)` | Check all URLs in text are valid | + +### Code Detection + +| Function | Description | +|----------|-------------| +| `detect_code(text)` | Returns `True` if code detected | +| `detect_code_languages(text)` | Returns list of detected languages | +| `contains_code_language(text, ["sql", "python"])` | Check for specific languages | + +### Text Utilities + +| Function | Description | +|----------|-------------| +| `contains(text, substring)` | Check if substring exists | +| `contains_any(text, [substr1, substr2])` | Check if any substring exists | +| `word_count(text)` | Count words | +| `char_count(text)` | Count characters | +| `lower(text)` / `upper(text)` / `trim(text)` | String transforms | + +## Examples + +### Block PII (SSN) + +```python +def apply_guardrail(inputs, request_data, input_type): + for text in inputs["texts"]: + if regex_match(text, r"\d{3}-\d{2}-\d{4}"): + return block("SSN detected") + return allow() +``` + +### Redact Email Addresses + +```python +def apply_guardrail(inputs, request_data, input_type): + pattern = r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}" + modified = [] + for text in inputs["texts"]: + modified.append(regex_replace(text, pattern, "[EMAIL REDACTED]")) + return modify(texts=modified) +``` + +### Block SQL Injection + +```python +def apply_guardrail(inputs, request_data, input_type): + if input_type != "request": + return allow() + for text in inputs["texts"]: + if contains_code_language(text, ["sql"]): + return block("SQL code not allowed") + return allow() +``` + +### Validate JSON Response + +```python +def apply_guardrail(inputs, request_data, input_type): + if input_type != "response": + return allow() + + schema = { + "type": "object", + "required": ["name", "value"] + } + + for text in inputs["texts"]: + obj = json_parse(text) + if obj is None: + return block("Invalid JSON response") + if not json_schema_valid(obj, schema): + return block("Response missing required fields") + return allow() +``` + +### Check URLs in Response + +```python +def apply_guardrail(inputs, request_data, input_type): + if input_type != "response": + return allow() + for text in inputs["texts"]: + if not all_urls_valid(text): + return block("Response contains invalid URLs") + return allow() +``` + +### Combine Multiple Checks + +```python +def apply_guardrail(inputs, request_data, input_type): + modified = [] + + for text in inputs["texts"]: + # Redact SSN + text = regex_replace(text, r"\d{3}-\d{2}-\d{4}", "[SSN]") + # Redact credit cards + text = regex_replace(text, r"\d{16}", "[CARD]") + modified.append(text) + + # Block SQL in requests + if input_type == "request": + for text in inputs["texts"]: + if contains_code_language(text, ["sql"]): + return block("SQL injection blocked") + + return modify(texts=modified) +``` + +## Sandbox Restrictions + +Custom code runs in a restricted environment: + +- ❌ No `import` statements +- ❌ No file I/O +- ❌ No network access +- ❌ No `exec()` or `eval()` +- ✅ Only LiteLLM-provided primitives available + +## Per-Request Usage + +Enable guardrail per request: + +```bash +curl -X POST http://localhost:4000/chat/completions \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "guardrails": ["block-ssn"] + }' +``` + +## Default On + +Run guardrail on all requests: + +```yaml +litellm_settings: + guardrails: + - guardrail_name: block-ssn + litellm_params: + guardrail: custom_code + mode: pre_call + default_on: true + custom_code: | + def apply_guardrail(inputs, request_data, input_type): + ... +``` diff --git a/docs/my-website/docs/proxy/guardrails/grayswan.md b/docs/my-website/docs/proxy/guardrails/grayswan.md index d6efaf15504..6c0ccbc293d 100644 --- a/docs/my-website/docs/proxy/guardrails/grayswan.md +++ b/docs/my-website/docs/proxy/guardrails/grayswan.md @@ -13,20 +13,26 @@ Cygnal returns a `violation` score between `0` and `1` (higher means more likely ### 1. Obtain Credentials -1. Create a Gray Swan account and generate a Cygnal API key. +1. Log in to our Gray Swan platform and generate a Cygnal API key. + + For existing customers, you should already have access to our [platform](https://platform.grayswan.ai). + + For new users, please register at this [page](https://hubs.ly/Q03-sX1J0) and we are more than happy to give you an onboarding! + + 2. Configure environment variables for the LiteLLM proxy host: -```bash -export GRAYSWAN_API_KEY="your-grayswan-key" -export GRAYSWAN_API_BASE="https://api.grayswan.ai" -``` + ```bash + export GRAYSWAN_API_KEY="your-grayswan-key" + export GRAYSWAN_API_BASE="https://api.grayswan.ai" + ``` ### 2. Configure `config.yaml` -Add a guardrail entry that references the Gray Swan integration. Below is a balanced example that monitors both input and output but only blocks once the violation score reaches the configured threshold. +Add a guardrail entry that references the Gray Swan integration. Below is our recommmended settings. ```yaml -model_list: +model_list: # this part is a standard litellm configuration for reference - model_name: openai/gpt-4.1-mini litellm_params: model: openai/gpt-4.1-mini @@ -40,13 +46,14 @@ guardrails: api_key: os.environ/GRAYSWAN_API_KEY api_base: os.environ/GRAYSWAN_API_BASE # optional optional_params: - on_flagged_action: monitor # or "block" + on_flagged_action: passthrough # or "block" or "monitor" violation_threshold: 0.5 # score >= threshold is flagged reasoning_mode: hybrid # off | hybrid | thinking - categories: - safety: "Detect jailbreaks and policy violations" - policy_id: "your-cygnal-policy-id" + policy_id: "your-cygnal-policy-id" # Optional: Your Cygnal policy ID. Defaults to a content safety policy if empty. + streaming_end_of_stream_only: true # For streaming API, only send the assembled message to Cygnal (post_call only). Defaults to false. default_on: true + guardrail_timeout: 30 # Defaults to 30 seconds. Change accordingly. + fail_open: true # Defaults to true; set to false to propagate guardrail errors. general_settings: master_key: "your-litellm-master-key" @@ -65,13 +72,13 @@ litellm --config config.yaml --port 4000 ## Choosing Guardrail Modes -Gray Swan can run during `pre_call`, `during_call`, and `post_call` stages. Combine modes based on your latency and coverage requirements. +Gray Swan can run during `pre_call`, `during_call`, and `post_call` stages. Combine modes based on your latency and coverage requirements. | Mode | When it Runs | Protects | Typical Use Case | |--------------|-------------------|-----------------------|------------------| | `pre_call` | Before LLM call | User input only | Block prompt injection before it reaches the model | | `during_call`| Parallel to call | User input only | Low-latency monitoring without blocking | -| `post_call` | After response | Full conversation | Scan output for policy violations, leaked secrets, or IPI | +| `post_call` | After response | Model Outputs | Scan output for policy violations, leaked secrets, or IPI | When using `during_call` with `on_flagged_action: block` or `on_flagged_action: passthrough`: @@ -81,87 +88,110 @@ When using `during_call` with `on_flagged_action: block` or `on_flagged_action: - The guardrail exception prevents the response from reaching the user, but **does not cancel the running LLM task** - This means you pay full LLM costs while returning an error/passthrough message to the user -**Recommendation:** For cost-sensitive applications, use `pre_call` and `post_call` instead of `during_call` for blocking or passthrough modes. Reserve `during_call` for `monitor` mode where you want low-latency logging without impacting the user experience. +**Recommendation:** Use `pre_call` and `post_call` instead of `during_call` for `passthrough` (or `block`) `on_flagged_action` (see our recommended configuration above). Reserve `during_call` for `monitor` mode ONLY when you want low-latency logging without impacting the user experience. - - +--- -```yaml -guardrails: - - guardrail_name: "cygnal-monitor-only" - litellm_params: - guardrail: grayswan - mode: "during_call" - api_key: os.environ/GRAYSWAN_API_KEY - optional_params: - on_flagged_action: monitor - violation_threshold: 0.6 - default_on: true +## Work with Claude Code + +Follow the official litellm [guide](https://docs.litellm.ai/docs/tutorials/claude_responses_api) on setting up Claude Code with litellm, with the guardrail part mentioned above added to your litellm configuration. Cygnal natively supports coding agent policies defense. Define your own policy or use the provided coding policies on the platform. The example config we show above is also the recommended setup for Claude Code (with the `policy_id` replaced with an appropriate one). + +--- + +## Per-request overrides via `extra_body` + +You can override parts of the Gray Swan guardrail configuration on a per-request basis by passing `litellm_metadata.guardrails[*].grayswan.extra_body`. + +`extra_body` is merged into the Cygnal request body and takes precedence over specific fields from `config.yaml`, which are `policy_id`, `violation_threshold`, and `reasoning_mode`. + +If you include a `metadata` field inside `extra_body`, it is forwarded to the Cygnal API as-is under the request body's `metadata` field. + +Example: + +```bash +curl -X POST "http://0.0.0.0:4000/v1/messages?beta=true" \ + -H "Authorization: Bearer token" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "openrouter/anthropic/claude-sonnet-4.5", + "messages": [{"role": "user", "content": "hello"}], + "litellm_metadata": { + "guardrails": [ + { + "cygnal-monitor": { + "extra_body": { + "policy_id": "specific policy id you want to use", + "metadata": { + "user": "health-check" + } + } + } + } + ] + } + }' ``` -Best for visibility without blocking. Alerts are logged via LiteLLM’s standard logging callbacks. +OpenAI client: - - +```python +from openai import OpenAI -```yaml -guardrails: - - guardrail_name: "cygnal-block-input" - litellm_params: - guardrail: grayswan - mode: "pre_call" - api_key: os.environ/GRAYSWAN_API_KEY - optional_params: - on_flagged_action: block - violation_threshold: 0.4 - categories: - pii: "Detect sensitive data" - default_on: true +client = OpenAI(api_key="anything", base_url="http://0.0.0.0:4000") + +resp = client.responses.create( + model="openrouter/anthropic/claude-sonnet-4.5", + input="hello", + extra_body={ + "litellm_metadata": { + "guardrails": [ + { + "cygnal-monitor": { + "extra_body": { + "policy_id": "69038214e5cdb6befc5e991e", + "metadata": {"trace_id": "trace-123"}, + } + } + } + ] + } + }, +) ``` -Stops malicious or sensitive prompts before any tokens are generated. +Anthropic client: - - +```python +from anthropic import Anthropic -```yaml -guardrails: - - guardrail_name: "cygnal-full-coverage" - litellm_params: - guardrail: grayswan - mode: [pre_call, post_call] - api_key: os.environ/GRAYSWAN_API_KEY - optional_params: - on_flagged_action: block - violation_threshold: 0.5 - reasoning_mode: thinking - policy_id: "policy-id-from-grayswan" - default_on: true +client = Anthropic(api_key="anything", base_url="http://0.0.0.0:4000") + +resp = client.messages.create( + model="openrouter/anthropic/claude-sonnet-4.5", + max_tokens=256, + messages=[{"role": "user", "content": "hello"}], + extra_body={ + "litellm_metadata": { + "guardrails": [ + { + "cygnal-monitor": { + "extra_body": { + "policy_id": "69038214e5cdb6befc5e991e", + "metadata": {"trace_id": "trace-123"}, + } + } + } + ] + } + }, +) ``` -Provides the strongest enforcement by inspecting both prompts and responses. +Notes: - - - -```yaml -guardrails: - - guardrail_name: "cygnal-passthrough" - litellm_params: - guardrail: grayswan - mode: [pre_call, post_call] - api_key: os.environ/GRAYSWAN_API_KEY - optional_params: - on_flagged_action: passthrough - violation_threshold: 0.5 - default_on: true -``` - -Allows requests to proceed without raising a 400 error when content is flagged. Instead of blocking, the model response content is replaced with a detailed violation message including violation score, violated rules, and detection flags (mutation, IPI). **Supported Response Formats:** OpenAI chat/text completions, Anthropic Messages API. Other response types (embeddings, images, etc.) will log a warning and return unchanged. - - - +- The guardrail name (for example, `cygnal-monitor`) must match the `guardrail_name` in `config.yaml`. +- Per-request guardrail overrides may require a premium license, depending on your proxy settings. --- @@ -170,9 +200,14 @@ Allows requests to proceed without raising a 400 error when content is flagged. | Parameter | Type | Description | |---------------------------------------|-----------------|-------------| | `api_key` | string | Gray Swan Cygnal API key. Reads from `GRAYSWAN_API_KEY` if omitted. | +| `api_base` | string | Override for the Gray Swan API base URL. Defaults to `https://api.grayswan.ai` or `GRAYSWAN_API_BASE`. | | `mode` | string or list | Guardrail stages (`pre_call`, `during_call`, `post_call`). | | `optional_params.on_flagged_action` | string | `monitor` (log only), `block` (raise `HTTPException`), or `passthrough` (replace response content with violation message, no 400 error). | -| `.optional_params.violation_threshold`| number (0-1) | Scores at or above this value are considered violations. | +| `optional_params.violation_threshold` | number (0-1) | Scores at or above this value are considered violations. | | `optional_params.reasoning_mode` | string | `off`, `hybrid`, or `thinking`. Enables Cygnal's reasoning capabilities. | | `optional_params.categories` | object | Map of custom category names to descriptions. | | `optional_params.policy_id` | string | Gray Swan policy identifier. | +| `guardrail_timeout` | number | Timeout in seconds for the Cygnal request. Defaults to 30. | +| `fail_open` | boolean | If true, errors contacting Cygnal are logged and the request proceeds; if false, errors propagate. Defaults to treu. | +| `streaming_end_of_stream_only` | boolean | For streaming `post_call`, only send the final assembled response to Cygnal. Defaults to false. | +| `default_on` | boolean | Run the guardrail on every request by default. | diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 033faeda7c9..fda0e3be4e4 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -79,6 +79,7 @@ const sidebars = { "proxy/guardrails/panw_prisma_airs", "proxy/guardrails/secret_detection", "proxy/guardrails/custom_guardrail", + "proxy/guardrails/custom_code_guardrail", "proxy/guardrails/prompt_injection", "proxy/guardrails/tool_permission", "proxy/guardrails/zscaler_ai_guard", diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index b118400b620..dc49036cb15 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -129,6 +129,7 @@ model LiteLLM_TeamTable { team_member_permissions String[] @default([]) policies String[] @default([]) model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases + allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) @@ -160,7 +161,8 @@ model LiteLLM_DeletedTeamTable { team_member_permissions String[] @default([]) policies String[] @default([]) model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases - + allow_team_guardrail_config Boolean @default(false) + // Original timestamps from team creation/updates created_at DateTime? @map("created_at") updated_at DateTime? @map("updated_at") @@ -774,6 +776,7 @@ model LiteLLM_GuardrailsTable { guardrail_name String @unique litellm_params Json guardrail_info Json? + team_id String? created_at DateTime @default(now()) updated_at DateTime @updatedAt } diff --git a/litellm/__init__.py b/litellm/__init__.py index f857e10eed3..8174b9d2655 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -261,6 +261,8 @@ extra_spend_tag_headers: Optional[List[str]] = None in_memory_llm_clients_cache: "LLMClientCache" safe_memory_mode: bool = False enable_azure_ad_token_refresh: Optional[bool] = False +# Proxy Authentication - auto-obtain/refresh OAuth2/JWT tokens for LiteLLM Proxy +proxy_auth: Optional[Any] = None ### DEFAULT AZURE API VERSION ### AZURE_DEFAULT_API_VERSION = "2025-02-01-preview" # this is updated to the latest ### DEFAULT WATSONX API VERSION ### diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 57bd05124aa..753a94295b3 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -329,6 +329,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): else: request_data[key] = value + if headers: + request_data["extra_headers"] = headers + return request_data @staticmethod diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index a5bb530fc56..1652ec2aa0c 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -475,11 +475,18 @@ class CustomGuardrail(CustomLogger): guardrail_config: DynamicGuardrailParams = DynamicGuardrailParams( **guardrail[self.guardrail_name] ) + extra_body = guardrail_config.get("extra_body", {}) if self._validate_premium_user() is not True: + if isinstance(extra_body, dict) and extra_body: + verbose_logger.warning( + "Guardrail %s: ignoring dynamic extra_body keys %s because premium_user is False", + self.guardrail_name, + list(extra_body.keys()), + ) return {} # Return the extra_body if it exists, otherwise empty dict - return guardrail_config.get("extra_body", {}) + return extra_body return {} diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 2c897cb0692..00c38eac188 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1683,6 +1683,108 @@ class PrometheusLogger(CustomLogger): ) pass + def _safe_get(self, obj: Any, key: str, default: Any = None) -> Any: + """Get value from dict or Pydantic model.""" + if obj is None: + return default + if isinstance(obj, dict): + return obj.get(key, default) + return getattr(obj, key, default) + + def _extract_deployment_failure_label_values( + self, request_kwargs: dict + ) -> Dict[str, Optional[str]]: + """ + Extract label values for deployment failure metrics from all available + sources in request_kwargs. Falls back to litellm_params metadata and + user_api_key_auth when standard_logging_payload has None values. + """ + standard_logging_payload = ( + request_kwargs.get("standard_logging_object", {}) or {} + ) + _litellm_params = request_kwargs.get("litellm_params", {}) or {} + _metadata_raw = self._safe_get(standard_logging_payload, "metadata") or {} + if isinstance(_metadata_raw, dict): + _metadata = _metadata_raw + else: + _metadata = { + "user_api_key_alias": getattr( + _metadata_raw, "user_api_key_alias", None + ), + "user_api_key_team_id": getattr( + _metadata_raw, "user_api_key_team_id", None + ), + "user_api_key_team_alias": getattr( + _metadata_raw, "user_api_key_team_alias", None + ), + "user_api_key_hash": getattr(_metadata_raw, "user_api_key_hash", None), + "requester_ip_address": getattr( + _metadata_raw, "requester_ip_address", None + ), + "user_agent": getattr(_metadata_raw, "user_agent", None), + } + _litellm_params_metadata = _litellm_params.get("metadata", {}) or {} + + # Extract user_api_key_auth if present (proxy injects this, skipped in merge) + user_api_key_auth = _litellm_params_metadata.get("user_api_key_auth") + + def _get_api_key_alias() -> Optional[str]: + val = _metadata.get("user_api_key_alias") + if val is not None: + return val + val = _litellm_params_metadata.get("user_api_key_alias") + if val is not None: + return val + if user_api_key_auth is not None: + return getattr(user_api_key_auth, "key_alias", None) + return None + + def _get_team_id() -> Optional[str]: + val = _metadata.get("user_api_key_team_id") + if val is not None: + return val + val = _litellm_params_metadata.get("user_api_key_team_id") + if val is not None: + return val + if user_api_key_auth is not None: + return getattr(user_api_key_auth, "team_id", None) + return None + + def _get_team_alias() -> Optional[str]: + val = _metadata.get("user_api_key_team_alias") + if val is not None: + return val + val = _litellm_params_metadata.get("user_api_key_team_alias") + if val is not None: + return val + if user_api_key_auth is not None: + return getattr(user_api_key_auth, "team_alias", None) + return None + + def _get_hashed_api_key() -> Optional[str]: + val = _metadata.get("user_api_key_hash") + if val is not None: + return val + val = _litellm_params_metadata.get("user_api_key_hash") + if val is not None: + return val + if user_api_key_auth is not None: + return getattr(user_api_key_auth, "api_key", None) or getattr( + user_api_key_auth, "api_key_hash", None + ) + return None + + return { + "api_key_alias": _get_api_key_alias(), + "team": _get_team_id(), + "team_alias": _get_team_alias(), + "hashed_api_key": _get_hashed_api_key(), + "client_ip": _metadata.get("requester_ip_address") + or _litellm_params_metadata.get("requester_ip_address"), + "user_agent": _metadata.get("user_agent") + or _litellm_params_metadata.get("user_agent"), + } + def set_llm_deployment_failure_metrics(self, request_kwargs: dict): """ Sets Failure metrics when an LLM API call fails @@ -1707,6 +1809,21 @@ class PrometheusLogger(CustomLogger): model_id = standard_logging_payload.get("model_id", None) exception = request_kwargs.get("exception", None) + # Fallback: model_id from litellm_metadata.model_info + if model_id is None: + _model_info = ( + (_litellm_params.get("litellm_metadata") or {}).get("model_info") + or (_litellm_params.get("metadata") or {}).get("model_info") + or {} + ) + model_id = _model_info.get("id") + + # Fallback: model_group from litellm_metadata + if model_group is None: + model_group = (_litellm_params.get("litellm_metadata") or {}).get( + "model_group" + ) or (_litellm_params.get("metadata") or {}).get("model_group") + llm_provider = _litellm_params.get("custom_llm_provider", None) if self._should_skip_metrics_for_invalid_key( @@ -1714,9 +1831,37 @@ class PrometheusLogger(CustomLogger): standard_logging_payload=standard_logging_payload, ): return - hashed_api_key = standard_logging_payload.get("metadata", {}).get( + + # Extract context labels from all available sources (fix for None labels) + fallback_values = self._extract_deployment_failure_label_values( + request_kwargs + ) + _metadata = standard_logging_payload.get("metadata", {}) or {} + hashed_api_key = fallback_values.get("hashed_api_key") or _metadata.get( "user_api_key_hash" ) + api_key_alias = fallback_values.get("api_key_alias") or _metadata.get( + "user_api_key_alias" + ) + team = fallback_values.get("team") or _metadata.get("user_api_key_team_id") + team_alias = fallback_values.get("team_alias") or _metadata.get( + "user_api_key_team_alias" + ) + client_ip = fallback_values.get("client_ip") or _metadata.get( + "requester_ip_address" + ) + user_agent = fallback_values.get("user_agent") or _metadata.get( + "user_agent" + ) + + # exception_status: prefer status_code, fallback to exception class for known types + exception_status = None + if exception is not None: + exception_status = str(getattr(exception, "status_code", None)) + if exception_status == "None" or not exception_status: + code = getattr(exception, "code", None) + if code is not None: + exception_status = str(code) # Create enum_values for the label factory (always create for use in different metrics) enum_values = UserAPIKeyLabelValues( @@ -1724,26 +1869,18 @@ class PrometheusLogger(CustomLogger): model_id=model_id, api_base=api_base, api_provider=llm_provider, - exception_status=( - str(getattr(exception, "status_code", None)) if exception else None - ), + exception_status=exception_status, exception_class=( self._get_exception_class_name(exception) if exception else None ), - requested_model=model_group, + requested_model=model_group or litellm_model_name, hashed_api_key=hashed_api_key, - api_key_alias=standard_logging_payload["metadata"][ - "user_api_key_alias" - ], - team=standard_logging_payload["metadata"]["user_api_key_team_id"], - team_alias=standard_logging_payload["metadata"][ - "user_api_key_team_alias" - ], + api_key_alias=api_key_alias, + team=team, + team_alias=team_alias, tags=standard_logging_payload.get("request_tags", []), - client_ip=standard_logging_payload["metadata"].get( - "requester_ip_address" - ), - user_agent=standard_logging_payload["metadata"].get("user_agent"), + client_ip=client_ip, + user_agent=user_agent, ) """ diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 8e1016bd5bd..a14e7d118e8 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -34,6 +34,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.utils import ( ChatCompletionMessageToolCall, + Choices, GenericGuardrailAPIInputs, ModelResponse, ) @@ -76,7 +77,8 @@ class AnthropicMessagesHandler(BaseTranslation): chat_completion_compatible_request, tool_name_mapping = ( LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( - anthropic_message_request=cast(AnthropicMessagesRequest, data) + # Use a shallow copy to avoid mutating request data (pop on litellm_metadata). + anthropic_message_request=cast(AnthropicMessagesRequest, data.copy()) ) ) @@ -84,9 +86,9 @@ class AnthropicMessagesHandler(BaseTranslation): texts_to_check: List[str] = [] images_to_check: List[str] = [] - tools_to_check: List[ChatCompletionToolParam] = ( - chat_completion_compatible_request.get("tools", []) - ) + tools_to_check: List[ + ChatCompletionToolParam + ] = chat_completion_compatible_request.get("tools", []) task_mappings: List[Tuple[int, Optional[int]]] = [] # Track (message_index, content_index) for each text # content_index is None for string content, int for list content @@ -282,7 +284,10 @@ class AnthropicMessagesHandler(BaseTranslation): if hasattr(content_block, "model_dump"): block_dict = content_block.model_dump() else: - block_dict = {"type": block_type, "text": getattr(content_block, "text", None)} + block_dict = { + "type": block_type, + "text": getattr(content_block, "text", None), + } else: continue @@ -358,30 +363,40 @@ class AnthropicMessagesHandler(BaseTranslation): """ has_ended = self._check_streaming_has_ended(responses_so_far) if has_ended: - # build the model response from the responses_so_far - model_response = cast( - ModelResponse, - AnthropicPassthroughLoggingHandler._build_complete_streaming_response( - all_chunks=responses_so_far, - litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj), - model="", - ), + built_response = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=responses_so_far, + litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj), + model="", ) - tool_calls_list = cast(Optional[List[ChatCompletionMessageToolCall]], model_response.choices[0].message.tool_calls) # type: ignore - string_so_far = model_response.choices[0].message.content # type: ignore - guardrail_inputs = GenericGuardrailAPIInputs() - if string_so_far: - guardrail_inputs["texts"] = [string_so_far] - if tool_calls_list: - guardrail_inputs["tool_calls"] = tool_calls_list - _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid - inputs=guardrail_inputs, - request_data={}, - input_type="response", - logging_obj=litellm_logging_obj, - ) + # Check if model_response is valid and has choices before accessing + if ( + built_response is not None + and hasattr(built_response, "choices") + and built_response.choices + ): + model_response = cast(ModelResponse, built_response) + first_choice = cast(Choices, model_response.choices[0]) + tool_calls_list = cast( + Optional[List[ChatCompletionMessageToolCall]], + first_choice.message.tool_calls, + ) + string_so_far = first_choice.message.content + guardrail_inputs = GenericGuardrailAPIInputs() + if string_so_far: + guardrail_inputs["texts"] = [string_so_far] + if tool_calls_list: + guardrail_inputs["tool_calls"] = tool_calls_list + + _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid + inputs=guardrail_inputs, + request_data={}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + else: + verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices") return responses_so_far string_so_far = self.get_streaming_string_so_far(responses_so_far) @@ -648,7 +663,10 @@ class AnthropicMessagesHandler(BaseTranslation): if isinstance(content_block, dict): if content_block.get("type") == "text": cast(Dict[str, Any], content_block)["text"] = guardrail_response - elif hasattr(content_block, "type") and getattr(content_block, "type", None) == "text": + elif ( + hasattr(content_block, "type") + and getattr(content_block, "type", None) == "text" + ): # Update Pydantic object's text attribute if hasattr(content_block, "text"): content_block.text = guardrail_response diff --git a/litellm/llms/github_copilot/chat/transformation.py b/litellm/llms/github_copilot/chat/transformation.py index 50f18cedf9b..945b7ebf7ff 100644 --- a/litellm/llms/github_copilot/chat/transformation.py +++ b/litellm/llms/github_copilot/chat/transformation.py @@ -1,11 +1,15 @@ -from typing import Any, Optional, Tuple, cast, List +from typing import Any, List, Optional, Tuple, cast from litellm.exceptions import AuthenticationError from litellm.llms.openai.openai import OpenAIConfig from litellm.types.llms.openai import AllMessageValues from ..authenticator import Authenticator -from ..common_utils import GetAPIKeyError, GITHUB_COPILOT_API_BASE +from ..common_utils import ( + GITHUB_COPILOT_API_BASE, + GetAPIKeyError, + get_copilot_default_headers, +) class GithubCopilotConfig(OpenAIConfig): @@ -45,14 +49,24 @@ class GithubCopilotConfig(OpenAIConfig): ): import litellm - disable_copilot_system_to_assistant = ( - litellm.disable_copilot_system_to_assistant - ) - if not disable_copilot_system_to_assistant: - for message in messages: - if "role" in message and message["role"] == "system": - cast(Any, message)["role"] = "assistant" - return messages + # Check if system-to-assistant conversion is disabled + if litellm.disable_copilot_system_to_assistant: + # GitHub Copilot API now supports system prompts for all models (Claude, GPT, etc.) + # No conversion needed - just return messages as-is + return messages + + # Default behavior: convert system messages to assistant for compatibility + transformed_messages = [] + for message in messages: + if message.get("role") == "system": + # Convert system message to assistant message + transformed_message = message.copy() + transformed_message["role"] = "assistant" + transformed_messages.append(transformed_message) + else: + transformed_messages.append(message) + + return transformed_messages def validate_environment( self, @@ -69,6 +83,14 @@ class GithubCopilotConfig(OpenAIConfig): headers, model, messages, optional_params, litellm_params, api_key, api_base ) + # Add Copilot-specific headers (editor-version, user-agent, etc.) + try: + copilot_api_key = self.authenticator.get_api_key() + copilot_headers = get_copilot_default_headers(copilot_api_key) + validated_headers = {**copilot_headers, **validated_headers} + except GetAPIKeyError: + pass # Will be handled later in the request flow + # Add X-Initiator header based on message roles initiator = self._determine_initiator(messages) validated_headers["X-Initiator"] = initiator @@ -87,7 +109,7 @@ class GithubCopilotConfig(OpenAIConfig): For other models, returns standard OpenAI parameters (which may include reasoning_effort for o-series models). """ from litellm.utils import supports_reasoning - + # Get base OpenAI parameters base_params = super().get_supported_openai_params(model) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index fb00aa28f45..c406f502b45 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -21,7 +21,13 @@ from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.main import stream_chunk_builder from litellm.types.llms.openai import ChatCompletionToolParam -from litellm.types.utils import Choices, GenericGuardrailAPIInputs, ModelResponse, ModelResponseStream, StreamingChoices +from litellm.types.utils import ( + Choices, + GenericGuardrailAPIInputs, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail @@ -80,9 +86,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if tool_calls_to_check: inputs["tool_calls"] = tool_calls_to_check # type: ignore if messages: - inputs["structured_messages"] = ( - messages # pass the openai /chat/completions messages to the guardrail, as-is - ) + inputs[ + "structured_messages" + ] = messages # pass the openai /chat/completions messages to the guardrail, as-is # Pass tools (function definitions) to the guardrail tools = data.get("tools") if tools: @@ -362,14 +368,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # check if the stream has ended has_stream_ended = False for chunk in responses_so_far: - if chunk.choices[0].finish_reason is not None: + if chunk.choices and chunk.choices[0].finish_reason is not None: has_stream_ended = True break if has_stream_ended: # convert to model response model_response = cast( - ModelResponse, stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj) + ModelResponse, + stream_chunk_builder( + chunks=responses_so_far, logging_obj=litellm_logging_obj + ), ) # run process_output_response await self.process_output_response( diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index d943662f9e4..ad3d4c932d4 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -319,9 +319,7 @@ class OpenAIResponsesHandler(BaseTranslation): return response if not response_output: - verbose_proxy_logger.debug( - "OpenAI Responses API: Empty output in response" - ) + verbose_proxy_logger.debug("OpenAI Responses API: Empty output in response") return response # Step 1: Extract all text content and tool calls from response output @@ -427,27 +425,30 @@ class OpenAIResponsesHandler(BaseTranslation): handle_raw_dict_callback=None, ) - tool_calls = model_response_choices[0].message.tool_calls - text = model_response_choices[0].message.content - guardrail_inputs = GenericGuardrailAPIInputs() - if text: - guardrail_inputs["texts"] = [text] - if tool_calls: - guardrail_inputs["tool_calls"] = cast( - List[ChatCompletionToolCallChunk], tool_calls - ) - # Include model information from the response if available - response_model = final_chunk.get("response", {}).get("model") - if response_model: - guardrail_inputs["model"] = response_model - if tool_calls or text: - _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs=guardrail_inputs, - request_data={}, - input_type="response", - logging_obj=litellm_logging_obj, - ) - return responses_so_far + if model_response_choices: + tool_calls = model_response_choices[0].message.tool_calls + text = model_response_choices[0].message.content + guardrail_inputs = GenericGuardrailAPIInputs() + if text: + guardrail_inputs["texts"] = [text] + if tool_calls: + guardrail_inputs["tool_calls"] = cast( + List[ChatCompletionToolCallChunk], tool_calls + ) + # Include model information from the response if available + response_model = final_chunk.get("response", {}).get("model") + if response_model: + guardrail_inputs["model"] = response_model + if tool_calls or text: + _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=guardrail_inputs, + request_data={}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + else: + verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices") # model_response_stream = OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(final_chunk) # tool_calls = model_response_stream.choices[0].tool_calls # convert openai response to model response @@ -513,11 +514,9 @@ class OpenAIResponsesHandler(BaseTranslation): # Check if it's an OutputText with text if isinstance(content_item, OutputText): if content_item.text: - return True elif isinstance(content_item, dict): if content_item.get("text"): - return True return False diff --git a/litellm/main.py b/litellm/main.py index 6d6bef81c26..bca023e65ec 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1199,6 +1199,13 @@ def completion( # type: ignore # noqa: PLR0915 headers = {} if extra_headers is not None: headers.update(extra_headers) + # Inject proxy auth headers if configured + if litellm.proxy_auth is not None: + try: + proxy_headers = litellm.proxy_auth.get_auth_headers() + headers.update(proxy_headers) + except Exception as e: + verbose_logger.warning(f"Failed to get proxy auth headers: {e}") num_retries = kwargs.get( "num_retries", None ) ## alt. param for 'max_retries'. Use this to pass retries w/ instructor. @@ -2497,6 +2504,20 @@ def completion( # type: ignore # noqa: PLR0915 headers = headers or litellm.headers + # Add GitHub Copilot headers (same as /responses endpoint does) + if custom_llm_provider == "github_copilot": + from litellm.llms.github_copilot.common_utils import ( + get_copilot_default_headers, + ) + from litellm.llms.github_copilot.authenticator import Authenticator + + copilot_auth = Authenticator() + copilot_api_key = copilot_auth.get_api_key() + copilot_headers = get_copilot_default_headers(copilot_api_key) + if extra_headers: + copilot_headers.update(extra_headers) + extra_headers = copilot_headers + if extra_headers is not None: optional_params["extra_headers"] = extra_headers @@ -4597,6 +4618,13 @@ def embedding( # noqa: PLR0915 headers = {} if extra_headers is not None: headers.update(extra_headers) + # Inject proxy auth headers if configured + if litellm.proxy_auth is not None: + try: + proxy_headers = litellm.proxy_auth.get_auth_headers() + headers.update(proxy_headers) + except Exception as e: + verbose_logger.warning(f"Failed to get proxy auth headers: {e}") ### CUSTOM MODEL COST ### input_cost_per_token = kwargs.get("input_cost_per_token", None) output_cost_per_token = kwargs.get("output_cost_per_token", None) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9fa68a51673..b167b92d68b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -27927,7 +27927,9 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 3e-07 + "output_cost_per_token": 3e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/alibaba/qwen3-coder": { "input_cost_per_token": 4e-07, @@ -27936,7 +27938,9 @@ "max_output_tokens": 66536, "max_tokens": 66536, "mode": "chat", - "output_cost_per_token": 1.6e-06 + "output_cost_per_token": 1.6e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/amazon/nova-lite": { "input_cost_per_token": 6e-08, @@ -27945,7 +27949,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.4e-07 + "output_cost_per_token": 2.4e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/amazon/nova-micro": { "input_cost_per_token": 3.5e-08, @@ -27954,7 +27961,9 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.4e-07 + "output_cost_per_token": 1.4e-07, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/amazon/nova-pro": { "input_cost_per_token": 8e-07, @@ -27963,7 +27972,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 3.2e-06 + "output_cost_per_token": 3.2e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/amazon/titan-embed-text-v2": { "input_cost_per_token": 2e-08, @@ -27983,7 +27995,11 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1.25e-06 + "output_cost_per_token": 1.25e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-3-opus": { "cache_creation_input_token_cost": 1.875e-05, @@ -27994,7 +28010,11 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 7.5e-05 + "output_cost_per_token": 7.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-3.5-haiku": { "cache_creation_input_token_cost": 1e-06, @@ -28005,7 +28025,11 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 4e-06 + "output_cost_per_token": 4e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-3.5-sonnet": { "cache_creation_input_token_cost": 3.75e-06, @@ -28016,7 +28040,11 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-3.7-sonnet": { "cache_creation_input_token_cost": 3.75e-06, @@ -28027,7 +28055,11 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-4-opus": { "cache_creation_input_token_cost": 1.875e-05, @@ -28038,7 +28070,11 @@ "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", - "output_cost_per_token": 7.5e-05 + "output_cost_per_token": 7.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-4-sonnet": { "cache_creation_input_token_cost": 3.75e-06, @@ -28049,7 +28085,9 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/cohere/command-a": { "input_cost_per_token": 2.5e-06, @@ -28058,7 +28096,9 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/cohere/command-r": { "input_cost_per_token": 1.5e-07, @@ -28067,7 +28107,9 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 6e-07 + "output_cost_per_token": 6e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/cohere/command-r-plus": { "input_cost_per_token": 2.5e-06, @@ -28076,7 +28118,9 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/cohere/embed-v4.0": { "input_cost_per_token": 1.2e-07, @@ -28094,7 +28138,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.19e-06 + "output_cost_per_token": 2.19e-06, + "supports_tool_choice": true }, "vercel_ai_gateway/deepseek/deepseek-r1-distill-llama-70b": { "input_cost_per_token": 7.5e-07, @@ -28103,7 +28148,10 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 9.9e-07 + "output_cost_per_token": 9.9e-07, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/deepseek/deepseek-v3": { "input_cost_per_token": 9e-07, @@ -28112,7 +28160,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 9e-07 + "output_cost_per_token": 9e-07, + "supports_tool_choice": true }, "vercel_ai_gateway/google/gemini-2.0-flash": { "deprecation_date": "2026-03-31", @@ -28122,7 +28171,11 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 6e-07 + "output_cost_per_token": 6e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/google/gemini-2.0-flash-lite": { "deprecation_date": "2026-03-31", @@ -28132,7 +28185,11 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 3e-07 + "output_cost_per_token": 3e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/google/gemini-2.5-flash": { "input_cost_per_token": 3e-07, @@ -28141,7 +28198,11 @@ "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 2.5e-06 + "output_cost_per_token": 2.5e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/google/gemini-2.5-pro": { "input_cost_per_token": 2.5e-06, @@ -28150,7 +28211,11 @@ "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/google/gemini-embedding-001": { "input_cost_per_token": 1.5e-07, @@ -28168,7 +28233,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2e-07 + "output_cost_per_token": 2e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/google/text-embedding-005": { "input_cost_per_token": 2.5e-08, @@ -28204,7 +28272,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 7.9e-07 + "output_cost_per_token": 7.9e-07, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3-8b": { "input_cost_per_token": 5e-08, @@ -28213,7 +28282,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 8e-08 + "output_cost_per_token": 8e-08, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3.1-70b": { "input_cost_per_token": 7.2e-07, @@ -28222,7 +28292,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 7.2e-07 + "output_cost_per_token": 7.2e-07, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3.1-8b": { "input_cost_per_token": 5e-08, @@ -28231,7 +28302,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 8e-08 + "output_cost_per_token": 8e-08, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/meta/llama-3.2-11b": { "input_cost_per_token": 1.6e-07, @@ -28240,7 +28313,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.6e-07 + "output_cost_per_token": 1.6e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3.2-1b": { "input_cost_per_token": 1e-07, @@ -28258,7 +28334,9 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.5e-07 + "output_cost_per_token": 1.5e-07, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/meta/llama-3.2-90b": { "input_cost_per_token": 7.2e-07, @@ -28267,7 +28345,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 7.2e-07 + "output_cost_per_token": 7.2e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3.3-70b": { "input_cost_per_token": 7.2e-07, @@ -28276,7 +28357,9 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 7.2e-07 + "output_cost_per_token": 7.2e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-4-maverick": { "input_cost_per_token": 2e-07, @@ -28285,7 +28368,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 6e-07 + "output_cost_per_token": 6e-07, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-4-scout": { "input_cost_per_token": 1e-07, @@ -28294,7 +28378,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 3e-07 + "output_cost_per_token": 3e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/codestral": { "input_cost_per_token": 3e-07, @@ -28303,7 +28390,9 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 9e-07 + "output_cost_per_token": 9e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/codestral-embed": { "input_cost_per_token": 1.5e-07, @@ -28321,7 +28410,10 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2.8e-07 + "output_cost_per_token": 2.8e-07, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/mistral/magistral-medium": { "input_cost_per_token": 2e-06, @@ -28330,7 +28422,10 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 5e-06 + "output_cost_per_token": 5e-06, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/mistral/magistral-small": { "input_cost_per_token": 5e-07, @@ -28339,7 +28434,8 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 1.5e-06 + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true }, "vercel_ai_gateway/mistral/ministral-3b": { "input_cost_per_token": 4e-08, @@ -28348,7 +28444,9 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 4e-08 + "output_cost_per_token": 4e-08, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/ministral-8b": { "input_cost_per_token": 1e-07, @@ -28357,7 +28455,10 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 1e-07 + "output_cost_per_token": 1e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/mistral-embed": { "input_cost_per_token": 1e-07, @@ -28375,7 +28476,9 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 6e-06 + "output_cost_per_token": 6e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/mistral-saba-24b": { "input_cost_per_token": 7.9e-07, @@ -28393,7 +28496,10 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 3e-07 + "output_cost_per_token": 3e-07, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/mistral/mixtral-8x22b-instruct": { "input_cost_per_token": 1.2e-06, @@ -28402,7 +28508,8 @@ "max_output_tokens": 2048, "max_tokens": 2048, "mode": "chat", - "output_cost_per_token": 1.2e-06 + "output_cost_per_token": 1.2e-06, + "supports_function_calling": true }, "vercel_ai_gateway/mistral/pixtral-12b": { "input_cost_per_token": 1.5e-07, @@ -28411,7 +28518,11 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 1.5e-07 + "output_cost_per_token": 1.5e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/mistral/pixtral-large": { "input_cost_per_token": 2e-06, @@ -28420,7 +28531,11 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 6e-06 + "output_cost_per_token": 6e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/moonshotai/kimi-k2": { "input_cost_per_token": 5.5e-07, @@ -28429,7 +28544,9 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 2.2e-06 + "output_cost_per_token": 2.2e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/morph/morph-v3-fast": { "input_cost_per_token": 8e-07, @@ -28456,7 +28573,9 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1.5e-06 + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/openai/gpt-3.5-turbo-instruct": { "input_cost_per_token": 1.5e-06, @@ -28474,7 +28593,10 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 3e-05 + "output_cost_per_token": 3e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/openai/gpt-4.1": { "cache_creation_input_token_cost": 0.0, @@ -28485,7 +28607,11 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 8e-06 + "output_cost_per_token": 8e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/gpt-4.1-mini": { "cache_creation_input_token_cost": 0.0, @@ -28496,7 +28622,11 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1.6e-06 + "output_cost_per_token": 1.6e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/gpt-4.1-nano": { "cache_creation_input_token_cost": 0.0, @@ -28507,7 +28637,11 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 4e-07 + "output_cost_per_token": 4e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/gpt-4o": { "cache_creation_input_token_cost": 0.0, @@ -28518,7 +28652,11 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/gpt-4o-mini": { "cache_creation_input_token_cost": 0.0, @@ -28529,7 +28667,11 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 6e-07 + "output_cost_per_token": 6e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/o1": { "cache_creation_input_token_cost": 0.0, @@ -28540,7 +28682,11 @@ "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", - "output_cost_per_token": 6e-05 + "output_cost_per_token": 6e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/o3": { "cache_creation_input_token_cost": 0.0, @@ -28551,7 +28697,11 @@ "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", - "output_cost_per_token": 8e-06 + "output_cost_per_token": 8e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/o3-mini": { "cache_creation_input_token_cost": 0.0, @@ -28562,7 +28712,10 @@ "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", - "output_cost_per_token": 4.4e-06 + "output_cost_per_token": 4.4e-06, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/o4-mini": { "cache_creation_input_token_cost": 0.0, @@ -28573,7 +28726,11 @@ "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", - "output_cost_per_token": 4.4e-06 + "output_cost_per_token": 4.4e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/text-embedding-3-large": { "input_cost_per_token": 1.3e-07, @@ -28645,7 +28802,10 @@ "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/vercel/v0-1.5-md": { "input_cost_per_token": 3e-06, @@ -28654,7 +28814,10 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-2": { "input_cost_per_token": 2e-06, @@ -28663,7 +28826,9 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-2-vision": { "input_cost_per_token": 2e-06, @@ -28672,7 +28837,10 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-3": { "input_cost_per_token": 3e-06, @@ -28681,7 +28849,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-3-fast": { "input_cost_per_token": 5e-06, @@ -28690,7 +28860,8 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.5e-05 + "output_cost_per_token": 2.5e-05, + "supports_function_calling": true }, "vercel_ai_gateway/xai/grok-3-mini": { "input_cost_per_token": 3e-07, @@ -28699,7 +28870,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 5e-07 + "output_cost_per_token": 5e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-3-mini-fast": { "input_cost_per_token": 6e-07, @@ -28708,7 +28881,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 4e-06 + "output_cost_per_token": 4e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-4": { "input_cost_per_token": 3e-06, @@ -28717,7 +28892,9 @@ "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/zai/glm-4.5": { "input_cost_per_token": 6e-07, @@ -28726,7 +28903,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.2e-06 + "output_cost_per_token": 2.2e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/zai/glm-4.5-air": { "input_cost_per_token": 2e-07, @@ -28735,7 +28914,9 @@ "max_output_tokens": 96000, "max_tokens": 96000, "mode": "chat", - "output_cost_per_token": 1.1e-06 + "output_cost_per_token": 1.1e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/zai/glm-4.6": { "litellm_provider": "vercel_ai_gateway", @@ -29912,7 +30093,9 @@ "mode": "chat", "output_cost_per_token": 1e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_regions": ["global"], + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29925,7 +30108,9 @@ "mode": "chat", "output_cost_per_token": 4e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_regions": ["global"], + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29938,7 +30123,9 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_regions": ["global"], + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29951,7 +30138,9 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_regions": ["global"], + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_tool_choice": true }, @@ -34900,4 +35089,4 @@ "output_cost_per_token": 0, "supports_reasoning": true } -} \ No newline at end of file +} diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404/index.html similarity index 100% rename from litellm/proxy/_experimental/out/404.html rename to litellm/proxy/_experimental/out/404/index.html diff --git a/litellm/proxy/_experimental/out/api-reference.html b/litellm/proxy/_experimental/out/api-reference/index.html similarity index 100% rename from litellm/proxy/_experimental/out/api-reference.html rename to litellm/proxy/_experimental/out/api-reference/index.html diff --git a/litellm/proxy/_experimental/out/experimental/api-playground.html b/litellm/proxy/_experimental/out/experimental/api-playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/api-playground.html rename to litellm/proxy/_experimental/out/experimental/api-playground/index.html diff --git a/litellm/proxy/_experimental/out/experimental/budgets.html b/litellm/proxy/_experimental/out/experimental/budgets/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/budgets.html rename to litellm/proxy/_experimental/out/experimental/budgets/index.html diff --git a/litellm/proxy/_experimental/out/experimental/caching.html b/litellm/proxy/_experimental/out/experimental/caching/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/caching.html rename to litellm/proxy/_experimental/out/experimental/caching/index.html diff --git a/litellm/proxy/_experimental/out/experimental/claude-code-plugins.html b/litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/claude-code-plugins.html rename to litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html diff --git a/litellm/proxy/_experimental/out/experimental/old-usage.html b/litellm/proxy/_experimental/out/experimental/old-usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/old-usage.html rename to litellm/proxy/_experimental/out/experimental/old-usage/index.html diff --git a/litellm/proxy/_experimental/out/experimental/prompts.html b/litellm/proxy/_experimental/out/experimental/prompts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/prompts.html rename to litellm/proxy/_experimental/out/experimental/prompts/index.html diff --git a/litellm/proxy/_experimental/out/experimental/tag-management.html b/litellm/proxy/_experimental/out/experimental/tag-management/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/tag-management.html rename to litellm/proxy/_experimental/out/experimental/tag-management/index.html diff --git a/litellm/proxy/_experimental/out/guardrails.html b/litellm/proxy/_experimental/out/guardrails/index.html similarity index 100% rename from litellm/proxy/_experimental/out/guardrails.html rename to litellm/proxy/_experimental/out/guardrails/index.html diff --git a/litellm/proxy/_experimental/out/login.html b/litellm/proxy/_experimental/out/login/index.html similarity index 100% rename from litellm/proxy/_experimental/out/login.html rename to litellm/proxy/_experimental/out/login/index.html diff --git a/litellm/proxy/_experimental/out/logs.html b/litellm/proxy/_experimental/out/logs/index.html similarity index 100% rename from litellm/proxy/_experimental/out/logs.html rename to litellm/proxy/_experimental/out/logs/index.html diff --git a/litellm/proxy/_experimental/out/mcp/oauth/callback.html b/litellm/proxy/_experimental/out/mcp/oauth/callback/index.html similarity index 100% rename from litellm/proxy/_experimental/out/mcp/oauth/callback.html rename to litellm/proxy/_experimental/out/mcp/oauth/callback/index.html diff --git a/litellm/proxy/_experimental/out/model-hub.html b/litellm/proxy/_experimental/out/model-hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model-hub.html rename to litellm/proxy/_experimental/out/model-hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub.html rename to litellm/proxy/_experimental/out/model_hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub_table.html rename to litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/proxy/_experimental/out/models-and-endpoints.html b/litellm/proxy/_experimental/out/models-and-endpoints/index.html similarity index 100% rename from litellm/proxy/_experimental/out/models-and-endpoints.html rename to litellm/proxy/_experimental/out/models-and-endpoints/index.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding/index.html similarity index 100% rename from litellm/proxy/_experimental/out/onboarding.html rename to litellm/proxy/_experimental/out/onboarding/index.html diff --git a/litellm/proxy/_experimental/out/organizations.html b/litellm/proxy/_experimental/out/organizations/index.html similarity index 100% rename from litellm/proxy/_experimental/out/organizations.html rename to litellm/proxy/_experimental/out/organizations/index.html diff --git a/litellm/proxy/_experimental/out/playground.html b/litellm/proxy/_experimental/out/playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/playground.html rename to litellm/proxy/_experimental/out/playground/index.html diff --git a/litellm/proxy/_experimental/out/policies.html b/litellm/proxy/_experimental/out/policies/index.html similarity index 100% rename from litellm/proxy/_experimental/out/policies.html rename to litellm/proxy/_experimental/out/policies/index.html diff --git a/litellm/proxy/_experimental/out/settings/admin-settings.html b/litellm/proxy/_experimental/out/settings/admin-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/admin-settings.html rename to litellm/proxy/_experimental/out/settings/admin-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/logging-and-alerts.html b/litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/logging-and-alerts.html rename to litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html diff --git a/litellm/proxy/_experimental/out/settings/router-settings.html b/litellm/proxy/_experimental/out/settings/router-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/router-settings.html rename to litellm/proxy/_experimental/out/settings/router-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/ui-theme.html b/litellm/proxy/_experimental/out/settings/ui-theme/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/ui-theme.html rename to litellm/proxy/_experimental/out/settings/ui-theme/index.html diff --git a/litellm/proxy/_experimental/out/teams.html b/litellm/proxy/_experimental/out/teams/index.html similarity index 100% rename from litellm/proxy/_experimental/out/teams.html rename to litellm/proxy/_experimental/out/teams/index.html diff --git a/litellm/proxy/_experimental/out/test-key.html b/litellm/proxy/_experimental/out/test-key/index.html similarity index 100% rename from litellm/proxy/_experimental/out/test-key.html rename to litellm/proxy/_experimental/out/test-key/index.html diff --git a/litellm/proxy/_experimental/out/tools/mcp-servers.html b/litellm/proxy/_experimental/out/tools/mcp-servers/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/mcp-servers.html rename to litellm/proxy/_experimental/out/tools/mcp-servers/index.html diff --git a/litellm/proxy/_experimental/out/tools/vector-stores.html b/litellm/proxy/_experimental/out/tools/vector-stores/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/vector-stores.html rename to litellm/proxy/_experimental/out/tools/vector-stores/index.html diff --git a/litellm/proxy/_experimental/out/usage.html b/litellm/proxy/_experimental/out/usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/usage.html rename to litellm/proxy/_experimental/out/usage/index.html diff --git a/litellm/proxy/_experimental/out/users.html b/litellm/proxy/_experimental/out/users/index.html similarity index 100% rename from litellm/proxy/_experimental/out/users.html rename to litellm/proxy/_experimental/out/users/index.html diff --git a/litellm/proxy/_experimental/out/virtual-keys.html b/litellm/proxy/_experimental/out/virtual-keys/index.html similarity index 100% rename from litellm/proxy/_experimental/out/virtual-keys.html rename to litellm/proxy/_experimental/out/virtual-keys/index.html diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 13eeae14485..6f527e268b2 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -14,3 +14,14 @@ model_list: litellm_params: model: openai/gpt-4.1-mini +guardrails: + - guardrail_name: redact-ssn + litellm_params: + guardrail: custom_code + mode: pre_call + custom_code: | + def apply_guardrail(inputs, request_data, input_type): + for text in inputs["texts"]: + if regex_match(text, r"\d{3}-\d{2}-\d{4}"): + return block("SSN detected in message") + return allow() \ No newline at end of file diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 3ce819439cb..a825ce22b25 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -1236,6 +1236,275 @@ async def get_provider_specific_params(): return provider_params +class TestCustomCodeGuardrailRequest(BaseModel): + """Request model for testing custom code guardrails.""" + + custom_code: str + """The Python-like code containing the apply_guardrail function.""" + + test_input: Dict[str, Any] + """The test input to pass to the guardrail. Should contain 'texts', optionally 'images', 'tools', etc.""" + + input_type: str = "request" + """Whether this is a 'request' or 'response' input type.""" + + request_data: Optional[Dict[str, Any]] = None + """Optional mock request_data (model, user_id, team_id, metadata, etc.).""" + + +class TestCustomCodeGuardrailResponse(BaseModel): + """Response model for testing custom code guardrails.""" + + success: bool + """Whether the test executed successfully (no errors).""" + + result: Optional[Dict[str, Any]] = None + """The guardrail result: action (allow/block/modify), reason, modified_texts, etc.""" + + error: Optional[str] = None + """Error message if execution failed.""" + + error_type: Optional[str] = None + """Type of error: 'compilation' or 'execution'.""" + + +@router.post( + "/guardrails/test_custom_code", + tags=["Guardrails"], + dependencies=[Depends(user_api_key_auth)], + response_model=TestCustomCodeGuardrailResponse, +) +async def test_custom_code_guardrail(request: TestCustomCodeGuardrailRequest): + """ + Test custom code guardrail logic without creating a guardrail. + + This endpoint allows admins to experiment with custom code guardrails by: + 1. Compiling the provided code in a sandbox + 2. Executing the apply_guardrail function with test input + 3. Returning the result (allow/block/modify) + + 👉 [Custom Code Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/custom_code_guardrail) + + Example Request: + ```bash + curl -X POST "http://localhost:4000/guardrails/test_custom_code" \\ + -H "Authorization: Bearer " \\ + -H "Content-Type: application/json" \\ + -d '{ + "custom_code": "def apply_guardrail(inputs, request_data, input_type):\\n for text in inputs[\\"texts\\"]:\\n if regex_match(text, r\\"\\\\d{3}-\\\\d{2}-\\\\d{4}\\"):\\n return block(\\"SSN detected\\")\\n return allow()", + "test_input": { + "texts": ["My SSN is 123-45-6789"] + }, + "input_type": "request" + }' + ``` + + Example Success Response (blocked): + ```json + { + "success": true, + "result": { + "action": "block", + "reason": "SSN detected" + }, + "error": null, + "error_type": null + } + ``` + + Example Success Response (allowed): + ```json + { + "success": true, + "result": { + "action": "allow" + }, + "error": null, + "error_type": null + } + ``` + + Example Success Response (modified): + ```json + { + "success": true, + "result": { + "action": "modify", + "texts": ["My SSN is [REDACTED]"] + }, + "error": null, + "error_type": null + } + ``` + + Example Error Response (compilation error): + ```json + { + "success": false, + "result": null, + "error": "Syntax error in custom code: invalid syntax (, line 1)", + "error_type": "compilation" + } + ``` + """ + import concurrent.futures + import re + + from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import ( + get_custom_code_primitives, + ) + + # Security validation patterns + FORBIDDEN_PATTERNS = [ + # Import statements + (r"\bimport\s+", "import statements are not allowed"), + (r"\bfrom\s+\w+\s+import\b", "from...import statements are not allowed"), + (r"__import__\s*\(", "__import__() is not allowed"), + # Dangerous builtins + (r"\bexec\s*\(", "exec() is not allowed"), + (r"\beval\s*\(", "eval() is not allowed"), + (r"\bcompile\s*\(", "compile() is not allowed"), + (r"\bopen\s*\(", "open() is not allowed"), + (r"\bgetattr\s*\(", "getattr() is not allowed"), + (r"\bsetattr\s*\(", "setattr() is not allowed"), + (r"\bdelattr\s*\(", "delattr() is not allowed"), + (r"\bglobals\s*\(", "globals() is not allowed"), + (r"\blocals\s*\(", "locals() is not allowed"), + (r"\bvars\s*\(", "vars() is not allowed"), + (r"\bdir\s*\(", "dir() is not allowed"), + (r"\bbreakpoint\s*\(", "breakpoint() is not allowed"), + (r"\binput\s*\(", "input() is not allowed"), + # Dangerous dunder access + (r"__builtins__", "__builtins__ access is not allowed"), + (r"__globals__", "__globals__ access is not allowed"), + (r"__code__", "__code__ access is not allowed"), + (r"__subclasses__", "__subclasses__ access is not allowed"), + (r"__bases__", "__bases__ access is not allowed"), + (r"__mro__", "__mro__ access is not allowed"), + (r"__class__", "__class__ access is not allowed"), + (r"__dict__", "__dict__ access is not allowed"), + (r"__getattribute__", "__getattribute__ access is not allowed"), + (r"__reduce__", "__reduce__ access is not allowed"), + (r"__reduce_ex__", "__reduce_ex__ access is not allowed"), + # OS/system access + (r"\bos\.", "os module access is not allowed"), + (r"\bsys\.", "sys module access is not allowed"), + (r"\bsubprocess\.", "subprocess module access is not allowed"), + ] + + EXECUTION_TIMEOUT_SECONDS = 5 + + try: + # Step 0: Security validation - check for forbidden patterns + code = request.custom_code + for pattern, error_msg in FORBIDDEN_PATTERNS: + if re.search(pattern, code): + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Security violation: {error_msg}", + error_type="compilation", + ) + + # Step 1: Compile the custom code with restricted environment + exec_globals = get_custom_code_primitives().copy() + + # Remove access to builtins to prevent escape + exec_globals["__builtins__"] = {} + + try: + exec(compile(request.custom_code, "", "exec"), exec_globals) + except SyntaxError as e: + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Syntax error in custom code: {e}", + error_type="compilation", + ) + except Exception as e: + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Failed to compile custom code: {e}", + error_type="compilation", + ) + + # Step 2: Verify apply_guardrail function exists + if "apply_guardrail" not in exec_globals: + return TestCustomCodeGuardrailResponse( + success=False, + error="Custom code must define an 'apply_guardrail' function. " + "Expected signature: apply_guardrail(inputs, request_data, input_type)", + error_type="compilation", + ) + + apply_fn = exec_globals["apply_guardrail"] + if not callable(apply_fn): + return TestCustomCodeGuardrailResponse( + success=False, + error="'apply_guardrail' must be a callable function", + error_type="compilation", + ) + + # Step 3: Prepare test inputs + test_inputs = request.test_input + if "texts" not in test_inputs: + test_inputs["texts"] = [] + + # Prepare mock request_data + mock_request_data = request.request_data or {} + safe_request_data = { + "model": mock_request_data.get("model", "test-model"), + "user_id": mock_request_data.get("user_id"), + "team_id": mock_request_data.get("team_id"), + "end_user_id": mock_request_data.get("end_user_id"), + "metadata": mock_request_data.get("metadata", {}), + } + + # Step 4: Execute the function with timeout protection + + def execute_guardrail(): + return apply_fn(test_inputs, safe_request_data, request.input_type) + + try: + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(execute_guardrail) + try: + result = future.result(timeout=EXECUTION_TIMEOUT_SECONDS) + except concurrent.futures.TimeoutError: + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Execution timeout: code took longer than {EXECUTION_TIMEOUT_SECONDS} seconds", + error_type="execution", + ) + except Exception as e: + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Execution error: {e}", + error_type="execution", + ) + + # Step 5: Validate and return result + if not isinstance(result, dict): + return TestCustomCodeGuardrailResponse( + success=True, + result={ + "action": "allow", + "warning": f"Expected dict result, got {type(result).__name__}. Treating as allow.", + }, + ) + + return TestCustomCodeGuardrailResponse( + success=True, + result=result, + ) + + except Exception as e: + verbose_proxy_logger.exception(f"Error testing custom code guardrail: {e}") + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Unexpected error: {e}", + error_type="execution", + ) + + @router.post("/guardrails/apply_guardrail", response_model=ApplyGuardrailResponse) @router.post("/apply_guardrail", response_model=ApplyGuardrailResponse) async def apply_guardrail( diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py new file mode 100644 index 00000000000..747b188feea --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py @@ -0,0 +1,65 @@ +"""Custom code guardrail integration for LiteLLM. + +This module allows users to write custom guardrail logic using Python-like code +that runs in a sandboxed environment with access to LiteLLM-provided primitives. +""" + +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .custom_code_guardrail import CustomCodeGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail( + litellm_params: "LitellmParams", guardrail: "Guardrail" +) -> CustomCodeGuardrail: + """ + Initialize a custom code guardrail. + + Args: + litellm_params: Configuration parameters including the custom code + guardrail: The guardrail configuration dict + + Returns: + CustomCodeGuardrail instance + """ + import litellm + + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("Custom code guardrail requires a guardrail_name") + + # Get the custom code from litellm_params + custom_code = getattr(litellm_params, "custom_code", None) + if not custom_code: + raise ValueError( + "Custom code guardrail requires 'custom_code' in litellm_params" + ) + + custom_code_guardrail = CustomCodeGuardrail( + guardrail_name=guardrail_name, + custom_code=custom_code, + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + + litellm.logging_callback_manager.add_litellm_callback(custom_code_guardrail) + return custom_code_guardrail + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.CUSTOM_CODE.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.CUSTOM_CODE.value: CustomCodeGuardrail, +} + +__all__ = [ + "CustomCodeGuardrail", + "initialize_guardrail", +] diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py new file mode 100644 index 00000000000..a0ca324411c --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py @@ -0,0 +1,372 @@ +""" +Custom code guardrail for LiteLLM. + +This module provides a guardrail that executes user-defined Python-like code +to implement custom guardrail logic. The code runs in a sandboxed environment +with access to LiteLLM-provided primitives for common guardrail operations. + +Example custom code: + + def apply_guardrail(inputs, request_data, input_type): + '''Block messages containing SSNs''' + for text in inputs["texts"]: + if regex_match(text, r"\\d{3}-\\d{2}-\\d{4}"): + return block("Social Security Number detected") + return allow() +""" + +import threading +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, cast + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +from litellm.types.utils import GenericGuardrailAPIInputs + +from .primitives import get_custom_code_primitives + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + +class CustomCodeGuardrailError(Exception): + """Raised when custom code guardrail execution fails.""" + + def __init__(self, message: str, details: Optional[Dict[str, Any]] = None) -> None: + super().__init__(message) + self.details = details or {} + + +class CustomCodeCompilationError(CustomCodeGuardrailError): + """Raised when custom code fails to compile.""" + + +class CustomCodeExecutionError(CustomCodeGuardrailError): + """Raised when custom code fails during execution.""" + + +class CustomCodeGuardrailConfigModel(GuardrailConfigModel): + """Configuration parameters for the custom code guardrail.""" + + custom_code: str + """The Python-like code containing the apply_guardrail function.""" + + +class CustomCodeGuardrail(CustomGuardrail): + """ + Guardrail that executes user-defined Python-like code. + + The code runs in a sandboxed environment that provides: + - Access to LiteLLM primitives (regex_match, json_parse, etc.) + - No file I/O or network access + - No imports allowed + + Users write an `apply_guardrail(inputs, request_data, input_type)` function + that returns one of: + - allow() - let the request/response through + - block(reason) - reject with a message + - modify(texts=...) - transform the content + + Example: + def apply_guardrail(inputs, request_data, input_type): + for text in inputs["texts"]: + if regex_match(text, r"password"): + return block("Sensitive content detected") + return allow() + """ + + def __init__( + self, + custom_code: str, + guardrail_name: Optional[str] = "custom_code", + **kwargs: Any, + ) -> None: + """ + Initialize the custom code guardrail. + + Args: + custom_code: The source code containing apply_guardrail function + guardrail_name: Name of this guardrail instance + **kwargs: Additional arguments passed to CustomGuardrail + """ + self.custom_code = custom_code + self._compiled_function: Optional[Any] = None + self._compile_lock = threading.Lock() + self._compile_error: Optional[str] = None + + supported_event_hooks = [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.during_call, + GuardrailEventHooks.post_call, + ] + + super().__init__( + guardrail_name=guardrail_name, + supported_event_hooks=supported_event_hooks, + **kwargs, + ) + + # Compile the code on initialization + self._compile_custom_code() + + @staticmethod + def get_config_model() -> Optional[Type[GuardrailConfigModel]]: + """Returns the config model for the UI.""" + return CustomCodeGuardrailConfigModel + + def _compile_custom_code(self) -> None: + """ + Compile the custom code and extract the apply_guardrail function. + + The code runs in a sandboxed environment with only the allowed primitives. + """ + with self._compile_lock: + if self._compiled_function is not None: + return + + try: + # Create a restricted execution environment + # Only include our safe primitives + exec_globals = get_custom_code_primitives().copy() + + # Execute the user code in the restricted environment + exec(compile(self.custom_code, "", "exec"), exec_globals) + + # Extract the apply_guardrail function + if "apply_guardrail" not in exec_globals: + raise CustomCodeCompilationError( + "Custom code must define an 'apply_guardrail' function. " + "Expected signature: apply_guardrail(inputs, request_data, input_type)" + ) + + apply_fn = exec_globals["apply_guardrail"] + if not callable(apply_fn): + raise CustomCodeCompilationError( + "'apply_guardrail' must be a callable function" + ) + + self._compiled_function = apply_fn + verbose_proxy_logger.debug( + f"Custom code guardrail '{self.guardrail_name}' compiled successfully" + ) + + except SyntaxError as e: + self._compile_error = f"Syntax error in custom code: {e}" + raise CustomCodeCompilationError(self._compile_error) from e + except CustomCodeCompilationError: + raise + except Exception as e: + self._compile_error = f"Failed to compile custom code: {e}" + raise CustomCodeCompilationError(self._compile_error) from e + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """ + Apply the custom code guardrail to the inputs. + + This method calls the user-defined apply_guardrail function and + processes its result to determine the appropriate action. + + Args: + inputs: Dictionary containing texts, images, tool_calls + request_data: The original request data with metadata + input_type: "request" for pre-call, "response" for post-call + logging_obj: Optional logging object + + Returns: + GenericGuardrailAPIInputs - possibly modified + + Raises: + HTTPException: If content is blocked + CustomCodeExecutionError: If execution fails + """ + if self._compiled_function is None: + if self._compile_error: + raise CustomCodeExecutionError( + f"Custom code guardrail not compiled: {self._compile_error}" + ) + raise CustomCodeExecutionError("Custom code guardrail not compiled") + + try: + # Prepare inputs dict for the function + + # Prepare request_data with safe subset of information + safe_request_data = self._prepare_safe_request_data(request_data) + + # Execute the custom function + result = self._compiled_function(inputs, safe_request_data, input_type) + + # Process the result + return self._process_result( + result=result, + inputs=inputs, + request_data=request_data, + input_type=input_type, + ) + + except HTTPException: + # Re-raise HTTP exceptions (from block action) + raise + except Exception as e: + verbose_proxy_logger.error( + f"Custom code guardrail '{self.guardrail_name}' execution error: {e}" + ) + raise CustomCodeExecutionError( + f"Custom code guardrail execution failed: {e}", + details={ + "guardrail_name": self.guardrail_name, + "input_type": input_type, + }, + ) from e + + def _prepare_safe_request_data(self, request_data: dict) -> Dict[str, Any]: + """ + Prepare a safe subset of request_data for code execution. + + This filters out sensitive information and provides only what's + needed for guardrail logic. + + Args: + request_data: The full request data + + Returns: + Safe subset of request data + """ + return { + "model": request_data.get("model"), + "user_id": request_data.get("user_api_key_user_id"), + "team_id": request_data.get("user_api_key_team_id"), + "end_user_id": request_data.get("user_api_key_end_user_id"), + "metadata": request_data.get("metadata", {}), + } + + def _process_result( + self, + result: Any, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + ) -> GenericGuardrailAPIInputs: + """ + Process the result from the custom code function. + + Args: + result: The return value from apply_guardrail + inputs: The original inputs + request_data: The request data + input_type: "request" or "response" + + Returns: + GenericGuardrailAPIInputs - possibly modified + + Raises: + HTTPException: If action is "block" + """ + if not isinstance(result, dict): + verbose_proxy_logger.warning( + f"Custom code guardrail '{self.guardrail_name}': " + f"Expected dict result, got {type(result).__name__}. Treating as allow." + ) + return inputs + + action = result.get("action", "allow") + + if action == "allow": + verbose_proxy_logger.debug( + f"Custom code guardrail '{self.guardrail_name}': Allowing {input_type}" + ) + return inputs + + elif action == "block": + reason = result.get("reason", "Blocked by custom code guardrail") + detection_info = result.get("detection_info", {}) + + verbose_proxy_logger.info( + f"Custom code guardrail '{self.guardrail_name}': Blocking {input_type} - {reason}" + ) + + is_output = input_type == "response" + + # For pre-call, raise passthrough exception to return synthetic response + if not is_output: + self.raise_passthrough_exception( + violation_message=reason, + request_data=request_data, + detection_info=detection_info, + ) + + # For post-call, raise HTTP exception + raise HTTPException( + status_code=400, + detail={ + "error": reason, + "guardrail": self.guardrail_name, + "detection_info": detection_info, + }, + ) + + elif action == "modify": + verbose_proxy_logger.debug( + f"Custom code guardrail '{self.guardrail_name}': Modifying {input_type}" + ) + + # Apply modifications + modified_inputs = dict(inputs) + + if "texts" in result and result["texts"] is not None: + modified_inputs["texts"] = result["texts"] + + if "images" in result and result["images"] is not None: + modified_inputs["images"] = result["images"] + + if "tool_calls" in result and result["tool_calls"] is not None: + modified_inputs["tool_calls"] = result["tool_calls"] + + return cast(GenericGuardrailAPIInputs, modified_inputs) + + else: + verbose_proxy_logger.warning( + f"Custom code guardrail '{self.guardrail_name}': " + f"Unknown action '{action}'. Treating as allow." + ) + return inputs + + def update_custom_code(self, new_code: str) -> None: + """ + Update the custom code and recompile. + + This method allows hot-reloading of guardrail logic without + restarting the server. + + Args: + new_code: The new source code + + Raises: + CustomCodeCompilationError: If the new code fails to compile + """ + with self._compile_lock: + # Reset state + old_function = self._compiled_function + old_code = self.custom_code + self._compiled_function = None + self._compile_error = None + + try: + self.custom_code = new_code + self._compile_custom_code() + verbose_proxy_logger.info( + f"Custom code guardrail '{self.guardrail_name}': Code updated successfully" + ) + except CustomCodeCompilationError: + # Rollback on failure + self.custom_code = old_code + self._compiled_function = old_function + raise diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py new file mode 100644 index 00000000000..695e59977c8 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py @@ -0,0 +1,602 @@ +""" +Built-in primitives provided to custom code guardrails. + +These functions are injected into the custom code execution environment +and provide safe, sandboxed functionality for common guardrail operations. +""" + +import json +import re +from typing import Any, Dict, List, Optional, Tuple, Type, Union +from urllib.parse import urlparse + +from litellm._logging import verbose_proxy_logger + +# ============================================================================= +# Result Types - Used by Starlark code to return guardrail decisions +# ============================================================================= + + +def allow() -> Dict[str, Any]: + """ + Allow the request/response to proceed unchanged. + + Returns: + Dict indicating the request should be allowed + """ + return {"action": "allow"} + + +def block( + reason: str, detection_info: Optional[Dict[str, Any]] = None +) -> Dict[str, Any]: + """ + Block the request/response with a reason. + + Args: + reason: Human-readable reason for blocking + detection_info: Optional additional detection metadata + + Returns: + Dict indicating the request should be blocked + """ + result: Dict[str, Any] = {"action": "block", "reason": reason} + if detection_info: + result["detection_info"] = detection_info + return result + + +def modify( + texts: Optional[List[str]] = None, + images: Optional[List[Any]] = None, + tool_calls: Optional[List[Any]] = None, +) -> Dict[str, Any]: + """ + Modify the request/response content. + + Args: + texts: Modified text content (if None, keeps original) + images: Modified image content (if None, keeps original) + tool_calls: Modified tool calls (if None, keeps original) + + Returns: + Dict indicating the content should be modified + """ + result: Dict[str, Any] = {"action": "modify"} + if texts is not None: + result["texts"] = texts + if images is not None: + result["images"] = images + if tool_calls is not None: + result["tool_calls"] = tool_calls + return result + + +# ============================================================================= +# Regex Primitives +# ============================================================================= + + +def regex_match(text: str, pattern: str, flags: int = 0) -> bool: + """ + Check if a regex pattern matches anywhere in the text. + + Args: + text: The text to search in + pattern: The regex pattern to match + flags: Optional regex flags (default: 0) + + Returns: + True if pattern matches, False otherwise + """ + try: + return bool(re.search(pattern, text, flags)) + except re.error as e: + verbose_proxy_logger.warning(f"Starlark regex_match error: {e}") + return False + + +def regex_match_all(text: str, pattern: str, flags: int = 0) -> bool: + """ + Check if a regex pattern matches the entire text. + + Args: + text: The text to match + pattern: The regex pattern + flags: Optional regex flags + + Returns: + True if pattern matches entire text, False otherwise + """ + try: + return bool(re.fullmatch(pattern, text, flags)) + except re.error as e: + verbose_proxy_logger.warning(f"Starlark regex_match_all error: {e}") + return False + + +def regex_replace(text: str, pattern: str, replacement: str, flags: int = 0) -> str: + """ + Replace all occurrences of a pattern in text. + + Args: + text: The text to modify + pattern: The regex pattern to find + replacement: The replacement string + flags: Optional regex flags + + Returns: + The text with replacements applied + """ + try: + return re.sub(pattern, replacement, text, flags=flags) + except re.error as e: + verbose_proxy_logger.warning(f"Starlark regex_replace error: {e}") + return text + + +def regex_find_all(text: str, pattern: str, flags: int = 0) -> List[str]: + """ + Find all occurrences of a pattern in text. + + Args: + text: The text to search + pattern: The regex pattern to find + flags: Optional regex flags + + Returns: + List of all matches + """ + try: + return re.findall(pattern, text, flags) + except re.error as e: + verbose_proxy_logger.warning(f"Starlark regex_find_all error: {e}") + return [] + + +# ============================================================================= +# JSON Primitives +# ============================================================================= + + +def json_parse(text: str) -> Optional[Any]: + """ + Parse a JSON string into a Python object. + + Args: + text: The JSON string to parse + + Returns: + Parsed Python object, or None if parsing fails + """ + try: + return json.loads(text) + except (json.JSONDecodeError, TypeError) as e: + verbose_proxy_logger.debug(f"Starlark json_parse error: {e}") + return None + + +def json_stringify(obj: Any) -> str: + """ + Convert a Python object to a JSON string. + + Args: + obj: The object to serialize + + Returns: + JSON string representation + """ + try: + return json.dumps(obj) + except (TypeError, ValueError) as e: + verbose_proxy_logger.warning(f"Starlark json_stringify error: {e}") + return "" + + +def json_schema_valid(obj: Any, schema: Dict[str, Any]) -> bool: + """ + Validate an object against a JSON schema. + + Args: + obj: The object to validate + schema: The JSON schema to validate against + + Returns: + True if valid, False otherwise + """ + try: + # Try to import jsonschema, fall back to basic validation if not available + try: + import jsonschema + + jsonschema.validate(instance=obj, schema=schema) + return True + except ImportError: + # Basic validation without jsonschema library + return _basic_json_schema_validate(obj, schema) + except Exception as validation_error: + # Catch jsonschema.ValidationError and other validation errors + if "ValidationError" in type(validation_error).__name__: + return False + raise + except Exception as e: + verbose_proxy_logger.warning(f"Custom code json_schema_valid error: {e}") + return False + + +def _basic_json_schema_validate( + obj: Any, schema: Dict[str, Any], max_depth: int = 50 +) -> bool: + """ + Basic JSON schema validation without external library. + Handles: type, required, properties + + Uses an iterative approach with a stack to avoid recursion limits. + max_depth limits nesting to prevent infinite loops from circular schemas. + """ + type_map: Dict[str, Union[Type, Tuple[Type, ...]]] = { + "object": dict, + "array": list, + "string": str, + "number": (int, float), + "integer": int, + "boolean": bool, + "null": type(None), + } + + # Stack of (obj, schema, depth) tuples to process + stack: List[Tuple[Any, Dict[str, Any], int]] = [(obj, schema, 0)] + + while stack: + current_obj, current_schema, depth = stack.pop() + + # Circuit breaker: stop if we've gone too deep + if depth > max_depth: + return False + + # Check type + schema_type = current_schema.get("type") + if schema_type: + expected_type = type_map.get(schema_type) + if expected_type is not None and not isinstance(current_obj, expected_type): + return False + + # Check required fields and properties for dicts + if isinstance(current_obj, dict): + required = current_schema.get("required", []) + for field in required: + if field not in current_obj: + return False + + # Queue property validations + properties = current_schema.get("properties", {}) + for prop_name, prop_schema in properties.items(): + if prop_name in current_obj: + stack.append((current_obj[prop_name], prop_schema, depth + 1)) + + return True + + +# ============================================================================= +# URL Primitives +# ============================================================================= + + +# Common URL pattern for extraction +_URL_PATTERN = re.compile( + r"https?://(?:[-\w.]|(?:%[\da-fA-F]{2}))+[^\s]*", re.IGNORECASE +) + + +def extract_urls(text: str) -> List[str]: + """ + Extract all URLs from text. + + Args: + text: The text to search for URLs + + Returns: + List of URLs found in the text + """ + return _URL_PATTERN.findall(text) + + +def is_valid_url(url: str) -> bool: + """ + Check if a URL is syntactically valid. + + Args: + url: The URL to validate + + Returns: + True if the URL is valid, False otherwise + """ + try: + result = urlparse(url) + return all([result.scheme, result.netloc]) + except Exception: + return False + + +def all_urls_valid(text: str) -> bool: + """ + Check if all URLs in text are valid. + + Args: + text: The text containing URLs + + Returns: + True if all URLs are valid (or no URLs), False otherwise + """ + urls = extract_urls(text) + return all(is_valid_url(url) for url in urls) + + +def get_url_domain(url: str) -> Optional[str]: + """ + Extract the domain from a URL. + + Args: + url: The URL to parse + + Returns: + The domain, or None if invalid + """ + try: + result = urlparse(url) + return result.netloc if result.netloc else None + except Exception: + return None + + +# ============================================================================= +# Code Detection Primitives +# ============================================================================= + + +# Common code patterns for detection +_CODE_PATTERNS = { + "sql": [ + r"\b(SELECT|INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|TRUNCATE)\b.*\b(FROM|INTO|TABLE|SET|WHERE)\b", + r"\b(SELECT)\s+[\w\*,\s]+\s+FROM\s+\w+", + r"\b(INSERT\s+INTO|UPDATE\s+\w+\s+SET|DELETE\s+FROM)\b", + ], + "python": [ + r"^\s*(def|class|import|from|if|for|while|try|except|with)\s+", + r"^\s*@\w+", # decorators + r"\b(print|len|range|str|int|float|list|dict|set)\s*\(", + ], + "javascript": [ + r"\b(function|const|let|var|class|import|export)\s+", + r"=>", # arrow functions + r"\b(console\.(log|error|warn))\s*\(", + ], + "typescript": [ + r":\s*(string|number|boolean|any|void|never)\b", + r"\b(interface|type|enum)\s+\w+", + r"<[A-Z]\w*>", # generics + ], + "java": [ + r"\b(public|private|protected)\s+(static\s+)?(class|void|int|String)\b", + r"\bSystem\.(out|err)\.print", + ], + "go": [ + r"\bfunc\s+\w+\s*\(", + r"\b(package|import)\s+", + r":=", # short variable declaration + ], + "rust": [ + r"\b(fn|let|mut|impl|struct|enum|pub|mod)\s+", + r"->", # return type + r"\b(println!|format!)\s*\(", + ], + "shell": [ + r"^#!.*\b(bash|sh|zsh)\b", + r"\b(echo|grep|sed|awk|cat|ls|cd|mkdir|rm)\s+", + r"\$\{?\w+\}?", # variable expansion + ], + "html": [ + r"<\s*(html|head|body|div|span|p|a|img|script|style)\b[^>]*>", + r"", + ], + "css": [ + r"\{[^}]*:\s*[^}]+;[^}]*\}", + r"@(media|keyframes|import|font-face)\b", + ], +} + + +def detect_code(text: str) -> bool: + """ + Check if text contains code of any language. + + Args: + text: The text to check + + Returns: + True if code is detected, False otherwise + """ + return len(detect_code_languages(text)) > 0 + + +def detect_code_languages(text: str) -> List[str]: + """ + Detect which programming languages are present in text. + + Args: + text: The text to analyze + + Returns: + List of detected language names + """ + detected = [] + for lang, patterns in _CODE_PATTERNS.items(): + for pattern in patterns: + try: + if re.search(pattern, text, re.IGNORECASE | re.MULTILINE): + detected.append(lang) + break # Only add each language once + except re.error: + continue + return detected + + +def contains_code_language(text: str, languages: List[str]) -> bool: + """ + Check if text contains code from specific languages. + + Args: + text: The text to check + languages: List of language names to check for + + Returns: + True if any of the specified languages are detected + """ + detected = detect_code_languages(text) + return any(lang.lower() in [d.lower() for d in detected] for lang in languages) + + +# ============================================================================= +# Text Utility Primitives +# ============================================================================= + + +def contains(text: str, substring: str) -> bool: + """ + Check if text contains a substring. + + Args: + text: The text to search in + substring: The substring to find + + Returns: + True if substring is found, False otherwise + """ + return substring in text + + +def contains_any(text: str, substrings: List[str]) -> bool: + """ + Check if text contains any of the given substrings. + + Args: + text: The text to search in + substrings: List of substrings to find + + Returns: + True if any substring is found, False otherwise + """ + return any(s in text for s in substrings) + + +def contains_all(text: str, substrings: List[str]) -> bool: + """ + Check if text contains all of the given substrings. + + Args: + text: The text to search in + substrings: List of substrings to find + + Returns: + True if all substrings are found, False otherwise + """ + return all(s in text for s in substrings) + + +def word_count(text: str) -> int: + """ + Count the number of words in text. + + Args: + text: The text to count words in + + Returns: + Number of words + """ + return len(text.split()) + + +def char_count(text: str) -> int: + """ + Count the number of characters in text. + + Args: + text: The text to count characters in + + Returns: + Number of characters + """ + return len(text) + + +def lower(text: str) -> str: + """Convert text to lowercase.""" + return text.lower() + + +def upper(text: str) -> str: + """Convert text to uppercase.""" + return text.upper() + + +def trim(text: str) -> str: + """Remove leading and trailing whitespace.""" + return text.strip() + + +# ============================================================================= +# Primitives Registry +# ============================================================================= + + +def get_custom_code_primitives() -> Dict[str, Any]: + """ + Get all primitives to inject into the custom code environment. + + Returns: + Dict of function name to function + """ + return { + # Result types + "allow": allow, + "block": block, + "modify": modify, + # Regex + "regex_match": regex_match, + "regex_match_all": regex_match_all, + "regex_replace": regex_replace, + "regex_find_all": regex_find_all, + # JSON + "json_parse": json_parse, + "json_stringify": json_stringify, + "json_schema_valid": json_schema_valid, + # URL + "extract_urls": extract_urls, + "is_valid_url": is_valid_url, + "all_urls_valid": all_urls_valid, + "get_url_domain": get_url_domain, + # Code detection + "detect_code": detect_code, + "detect_code_languages": detect_code_languages, + "contains_code_language": contains_code_language, + # Text utilities + "contains": contains, + "contains_any": contains_any, + "contains_all": contains_all, + "word_count": word_count, + "char_count": char_count, + "lower": lower, + "upper": upper, + "trim": trim, + # Python builtins (safe subset) + "len": len, + "str": str, + "int": int, + "float": float, + "bool": bool, + "list": list, + "dict": dict, + "True": True, + "False": False, + "None": None, + } diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 2a852cbda08..90f689ed23c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -9,8 +9,10 @@ from fastapi import HTTPException from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, + ModifyResponseException ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -21,6 +23,8 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +GRAYSWAN_BLOCK_ERROR_MSG = "Blocked by Gray Swan Guardrail" + class GraySwanGuardrailMissingSecrets(Exception): """Raised when the Gray Swan API key is missing.""" @@ -205,9 +209,13 @@ class GraySwanGuardrail(CustomGuardrail): # Get dynamic params from request metadata dynamic_body = self.get_guardrail_dynamic_request_body_params(request_data) or {} + if dynamic_body: + verbose_proxy_logger.debug( + "Gray Swan Guardrail: dynamic extra_body=%s", safe_dumps(dynamic_body) + ) # Prepare and send payload - payload = self._prepare_payload(messages, dynamic_body) + payload = self._prepare_payload(messages, dynamic_body, request_data) if payload is None: return inputs @@ -223,6 +231,8 @@ class GraySwanGuardrail(CustomGuardrail): ) return result except Exception as exc: + if self._is_grayswan_exception(exc): + raise end_time = time.time() status_code = getattr(exc, "status_code", None) or getattr( exc, "exception_status_code", None @@ -240,8 +250,20 @@ class GraySwanGuardrail(CustomGuardrail): exc, ) return inputs + if isinstance(exc, GraySwanGuardrailAPIError): + raise exc raise GraySwanGuardrailAPIError(str(exc), status_code=status_code) from exc + def _is_grayswan_exception(self, exc: Exception) -> bool: + # Guardrail decision (passthrough) should always propagate, + # regardless of fail_open. + if isinstance(exc, ModifyResponseException): + return True + detail = getattr(exc, "detail", None) + if isinstance(detail, dict): + return detail.get("error") == GRAYSWAN_BLOCK_ERROR_MSG + return False + # ------------------------------------------------------------------ # Legacy Test Interface (for backward compatibility) # ------------------------------------------------------------------ @@ -324,7 +346,7 @@ class GraySwanGuardrail(CustomGuardrail): raise HTTPException( status_code=400, detail={ - "error": "Blocked by Gray Swan Guardrail", + "error": GRAYSWAN_BLOCK_ERROR_MSG, "violation_location": violation_location, "violation": violation_score, "violated_rules": violated_rules, @@ -445,7 +467,7 @@ class GraySwanGuardrail(CustomGuardrail): raise HTTPException( status_code=400, detail={ - "error": "Blocked by Gray Swan Guardrail", + "error": GRAYSWAN_BLOCK_ERROR_MSG, "violation_location": violation_location, "violation": violation_score, "violated_rules": violated_rules, @@ -494,7 +516,7 @@ class GraySwanGuardrail(CustomGuardrail): } def _prepare_payload( - self, messages: List[Dict[str, str]], dynamic_body: dict + self, messages: List[Dict[str, str]], dynamic_body: dict, request_data: dict ) -> Optional[Dict[str, Any]]: payload: Dict[str, Any] = {"messages": messages} @@ -510,6 +532,18 @@ class GraySwanGuardrail(CustomGuardrail): if reasoning_mode: payload["reasoning_mode"] = reasoning_mode + # Pass through arbitrary metadata when provided via dynamic extra_body. + if "metadata" in dynamic_body: + payload["metadata"] = dynamic_body["metadata"] + + litellm_metadata = request_data.get("litellm_metadata") + if isinstance(litellm_metadata, dict) and litellm_metadata: + cleaned_litellm_metadata = dict(litellm_metadata) + # cleaned_litellm_metadata.pop("user_api_key_auth", None) + sanitized = safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) + if isinstance(sanitized, dict) and sanitized: + payload["litellm_metadata"] = sanitized + return payload def _format_violation_message( diff --git a/litellm/proxy_auth/__init__.py b/litellm/proxy_auth/__init__.py new file mode 100644 index 00000000000..27624a94fb9 --- /dev/null +++ b/litellm/proxy_auth/__init__.py @@ -0,0 +1,30 @@ +""" +Proxy Authentication module for LiteLLM SDK. + +This module provides OAuth2/JWT token management for authenticating +with LiteLLM Proxy or any OAuth2-protected endpoint. + +Usage: + from litellm.proxy_auth import AzureADCredential, ProxyAuthHandler + + litellm.proxy_auth = ProxyAuthHandler( + credential=AzureADCredential(), + scope="api://my-proxy/.default" + ) +""" + +from .credentials import ( + AccessToken, + TokenCredential, + AzureADCredential, + GenericOAuth2Credential, + ProxyAuthHandler, +) + +__all__ = [ + "AccessToken", + "TokenCredential", + "AzureADCredential", + "GenericOAuth2Credential", + "ProxyAuthHandler", +] diff --git a/litellm/proxy_auth/credentials.py b/litellm/proxy_auth/credentials.py new file mode 100644 index 00000000000..cddaf1278f9 --- /dev/null +++ b/litellm/proxy_auth/credentials.py @@ -0,0 +1,240 @@ +""" +Credential providers for proxy authentication. + +This module provides a provider-agnostic interface for obtaining OAuth2/JWT tokens. +It follows the same TokenCredential protocol used by Azure SDK. +""" + +import time +from dataclasses import dataclass +from typing import Any, Optional, Protocol, runtime_checkable + + +@dataclass +class AccessToken: + """ + Represents an OAuth2 access token with expiration. + + This matches the structure used by azure.core.credentials.AccessToken. + + Attributes: + token: The access token string (typically a JWT). + expires_on: Unix timestamp when the token expires. + """ + + token: str + expires_on: int + + +@runtime_checkable +class TokenCredential(Protocol): + """ + Protocol for credential providers. + + This matches the azure.core.credentials.TokenCredential interface, + allowing any Azure SDK credential to be used directly. + + Any class implementing get_token(scope) -> AccessToken can be used. + """ + + def get_token(self, scope: str) -> AccessToken: + """ + Get an access token for the specified scope. + + Args: + scope: The OAuth2 scope to request (e.g., "api://my-app/.default") + + Returns: + AccessToken with the token string and expiration timestamp. + """ + ... + + +class AzureADCredential: + """ + Wrapper for Azure Identity credentials. + + This wraps any azure-identity credential (DefaultAzureCredential, + ClientSecretCredential, ManagedIdentityCredential, etc.) and converts + the token to our AccessToken format. + + If no credential is provided, it will use DefaultAzureCredential + which tries multiple authentication methods automatically. + + Example: + # Use default credential chain (env vars, managed identity, CLI, etc.) + cred = AzureADCredential() + + # Or provide a specific credential + from azure.identity import ClientSecretCredential + azure_cred = ClientSecretCredential(tenant_id, client_id, client_secret) + cred = AzureADCredential(credential=azure_cred) + """ + + def __init__(self, credential: Optional[Any] = None): + """ + Initialize with an optional Azure credential. + + Args: + credential: An azure-identity credential object. If None, + DefaultAzureCredential will be used on first token request. + """ + self._credential = credential + self._initialized = credential is not None + + def get_token(self, scope: str) -> AccessToken: + """ + Get an access token from Azure AD. + + Args: + scope: The OAuth2 scope (e.g., "api://my-app/.default") + + Returns: + AccessToken with the JWT and expiration. + + Raises: + ImportError: If azure-identity is not installed. + """ + if not self._initialized: + try: + from azure.identity import DefaultAzureCredential + + self._credential = DefaultAzureCredential() + self._initialized = True + except ImportError: + raise ImportError( + "azure-identity is required for AzureADCredential. " + "Install it with: pip install azure-identity" + ) + + result = self._credential.get_token(scope) + return AccessToken(token=result.token, expires_on=result.expires_on) + + +class GenericOAuth2Credential: + """ + Generic OAuth2 client credentials flow. + + This works with any OAuth2 provider (Okta, Auth0, Keycloak, etc.) + that supports the client_credentials grant type. + + Example: + cred = GenericOAuth2Credential( + client_id="my-client-id", + client_secret="my-client-secret", + token_url="https://my-idp.com/oauth2/token" + ) + """ + + def __init__(self, client_id: str, client_secret: str, token_url: str): + """ + Initialize OAuth2 client credentials. + + Args: + client_id: OAuth2 client ID + client_secret: OAuth2 client secret + token_url: Token endpoint URL (e.g., "https://idp.com/oauth2/token") + """ + self.client_id = client_id + self.client_secret = client_secret + self.token_url = token_url + self._cached_token: Optional[AccessToken] = None + + def get_token(self, scope: str) -> AccessToken: + """ + Get an access token using OAuth2 client credentials flow. + + Tokens are cached and reused until they expire (with 60s buffer). + + Args: + scope: The OAuth2 scope to request + + Returns: + AccessToken with the token and expiration. + """ + # Return cached token if still valid (with 60s buffer) + if self._cached_token and self._cached_token.expires_on > time.time() + 60: + return self._cached_token + + import httpx + + response = httpx.post( + self.token_url, + data={ + "grant_type": "client_credentials", + "client_id": self.client_id, + "client_secret": self.client_secret, + "scope": scope, + }, + ) + response.raise_for_status() + data = response.json() + + self._cached_token = AccessToken( + token=data["access_token"], + expires_on=int(time.time()) + data.get("expires_in", 3600), + ) + return self._cached_token + + +class ProxyAuthHandler: + """ + Manages OAuth2/JWT token lifecycle for proxy authentication. + + This handler: + - Obtains tokens from the configured credential provider + - Caches tokens to avoid unnecessary requests + - Automatically refreshes tokens before they expire (60s buffer) + - Generates Authorization headers for HTTP requests + + Set this as litellm.proxy_auth to automatically inject auth headers + into all requests to your LiteLLM Proxy. + + Example: + import litellm + from litellm.proxy_auth import AzureADCredential, ProxyAuthHandler + + litellm.proxy_auth = ProxyAuthHandler( + credential=AzureADCredential(), + scope="api://my-litellm-proxy/.default" + ) + litellm.api_base = "https://my-proxy.example.com" + + # Auth headers are now automatically injected + response = litellm.completion(model="gpt-4", messages=[...]) + """ + + def __init__(self, credential: TokenCredential, scope: str): + """ + Initialize the proxy auth handler. + + Args: + credential: A TokenCredential implementation (AzureADCredential, + GenericOAuth2Credential, or any custom implementation) + scope: The OAuth2 scope to request tokens for + """ + self.credential = credential + self.scope = scope + self._cached_token: Optional[AccessToken] = None + + def get_token(self) -> AccessToken: + """ + Get a valid access token, refreshing if necessary. + + Returns: + AccessToken that is valid for at least 60 more seconds. + """ + # Refresh if no token or token expires within 60 seconds + if not self._cached_token or self._cached_token.expires_on <= time.time() + 60: + self._cached_token = self.credential.get_token(self.scope) + return self._cached_token + + def get_auth_headers(self) -> dict: + """ + Get HTTP headers for authentication. + + Returns: + Dict with Authorization header containing Bearer token. + """ + token = self.get_token() + return {"Authorization": f"Bearer {token.token}"} diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index ca22049720e..74ccb34ca6e 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -14,14 +14,14 @@ from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import ( from litellm.types.proxy.guardrails.guardrail_hooks.ibm import ( IBMGuardrailsBaseConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( - ToolPermissionGuardrailConfigModel, +from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( + ContentFilterCategoryConfig, ) from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import ( QualifireGuardrailConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( - ContentFilterCategoryConfig, +from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( + ToolPermissionGuardrailConfigModel, ) """ @@ -68,6 +68,7 @@ class SupportedGuardrailIntegrations(Enum): PROMPT_SECURITY = "prompt_security" GENERIC_GUARDRAIL_API = "generic_guardrail_api" QUALIFIRE = "qualifire" + CUSTOM_CODE = "custom_code" class Role(Enum): @@ -296,13 +297,7 @@ class PresidioConfigModel(PresidioPresidioConfigModelUserInterface): pii_entities_config: Optional[Dict[Union[PiiEntityType, str], PiiAction]] = Field( default=None, description="Configuration for PII entity types and actions" ) - presidio_filter_scope: Literal["input", "output", "both"] = Field( - default="both", - description=( - "Where to apply Presidio checks: 'input' runs on user → model traffic, " - "'output' runs on model → user traffic, and 'both' applies to both." - ), - ) + presidio_score_thresholds: Optional[Dict[Union[PiiEntityType, str], float]] = Field( default=None, description=( @@ -656,6 +651,12 @@ class BaseLitellmParams( description="Additional provider-specific parameters for generic guardrail APIs", ) + # Custom code guardrail params + custom_code: Optional[str] = Field( + default=None, + description="Python-like code containing the apply_guardrail function for custom guardrail logic", + ) + model_config = ConfigDict(extra="allow", protected_namespaces=()) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9fa68a51673..b167b92d68b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -27927,7 +27927,9 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 3e-07 + "output_cost_per_token": 3e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/alibaba/qwen3-coder": { "input_cost_per_token": 4e-07, @@ -27936,7 +27938,9 @@ "max_output_tokens": 66536, "max_tokens": 66536, "mode": "chat", - "output_cost_per_token": 1.6e-06 + "output_cost_per_token": 1.6e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/amazon/nova-lite": { "input_cost_per_token": 6e-08, @@ -27945,7 +27949,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.4e-07 + "output_cost_per_token": 2.4e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/amazon/nova-micro": { "input_cost_per_token": 3.5e-08, @@ -27954,7 +27961,9 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.4e-07 + "output_cost_per_token": 1.4e-07, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/amazon/nova-pro": { "input_cost_per_token": 8e-07, @@ -27963,7 +27972,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 3.2e-06 + "output_cost_per_token": 3.2e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/amazon/titan-embed-text-v2": { "input_cost_per_token": 2e-08, @@ -27983,7 +27995,11 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1.25e-06 + "output_cost_per_token": 1.25e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-3-opus": { "cache_creation_input_token_cost": 1.875e-05, @@ -27994,7 +28010,11 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 7.5e-05 + "output_cost_per_token": 7.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-3.5-haiku": { "cache_creation_input_token_cost": 1e-06, @@ -28005,7 +28025,11 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 4e-06 + "output_cost_per_token": 4e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-3.5-sonnet": { "cache_creation_input_token_cost": 3.75e-06, @@ -28016,7 +28040,11 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-3.7-sonnet": { "cache_creation_input_token_cost": 3.75e-06, @@ -28027,7 +28055,11 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-4-opus": { "cache_creation_input_token_cost": 1.875e-05, @@ -28038,7 +28070,11 @@ "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", - "output_cost_per_token": 7.5e-05 + "output_cost_per_token": 7.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/anthropic/claude-4-sonnet": { "cache_creation_input_token_cost": 3.75e-06, @@ -28049,7 +28085,9 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/cohere/command-a": { "input_cost_per_token": 2.5e-06, @@ -28058,7 +28096,9 @@ "max_output_tokens": 8000, "max_tokens": 8000, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/cohere/command-r": { "input_cost_per_token": 1.5e-07, @@ -28067,7 +28107,9 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 6e-07 + "output_cost_per_token": 6e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/cohere/command-r-plus": { "input_cost_per_token": 2.5e-06, @@ -28076,7 +28118,9 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/cohere/embed-v4.0": { "input_cost_per_token": 1.2e-07, @@ -28094,7 +28138,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.19e-06 + "output_cost_per_token": 2.19e-06, + "supports_tool_choice": true }, "vercel_ai_gateway/deepseek/deepseek-r1-distill-llama-70b": { "input_cost_per_token": 7.5e-07, @@ -28103,7 +28148,10 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 9.9e-07 + "output_cost_per_token": 9.9e-07, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/deepseek/deepseek-v3": { "input_cost_per_token": 9e-07, @@ -28112,7 +28160,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 9e-07 + "output_cost_per_token": 9e-07, + "supports_tool_choice": true }, "vercel_ai_gateway/google/gemini-2.0-flash": { "deprecation_date": "2026-03-31", @@ -28122,7 +28171,11 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 6e-07 + "output_cost_per_token": 6e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/google/gemini-2.0-flash-lite": { "deprecation_date": "2026-03-31", @@ -28132,7 +28185,11 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 3e-07 + "output_cost_per_token": 3e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/google/gemini-2.5-flash": { "input_cost_per_token": 3e-07, @@ -28141,7 +28198,11 @@ "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 2.5e-06 + "output_cost_per_token": 2.5e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/google/gemini-2.5-pro": { "input_cost_per_token": 2.5e-06, @@ -28150,7 +28211,11 @@ "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/google/gemini-embedding-001": { "input_cost_per_token": 1.5e-07, @@ -28168,7 +28233,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2e-07 + "output_cost_per_token": 2e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/google/text-embedding-005": { "input_cost_per_token": 2.5e-08, @@ -28204,7 +28272,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 7.9e-07 + "output_cost_per_token": 7.9e-07, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3-8b": { "input_cost_per_token": 5e-08, @@ -28213,7 +28282,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 8e-08 + "output_cost_per_token": 8e-08, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3.1-70b": { "input_cost_per_token": 7.2e-07, @@ -28222,7 +28292,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 7.2e-07 + "output_cost_per_token": 7.2e-07, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3.1-8b": { "input_cost_per_token": 5e-08, @@ -28231,7 +28302,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 8e-08 + "output_cost_per_token": 8e-08, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/meta/llama-3.2-11b": { "input_cost_per_token": 1.6e-07, @@ -28240,7 +28313,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.6e-07 + "output_cost_per_token": 1.6e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3.2-1b": { "input_cost_per_token": 1e-07, @@ -28258,7 +28334,9 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 1.5e-07 + "output_cost_per_token": 1.5e-07, + "supports_function_calling": true, + "supports_response_schema": true }, "vercel_ai_gateway/meta/llama-3.2-90b": { "input_cost_per_token": 7.2e-07, @@ -28267,7 +28345,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 7.2e-07 + "output_cost_per_token": 7.2e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-3.3-70b": { "input_cost_per_token": 7.2e-07, @@ -28276,7 +28357,9 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 7.2e-07 + "output_cost_per_token": 7.2e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-4-maverick": { "input_cost_per_token": 2e-07, @@ -28285,7 +28368,8 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 6e-07 + "output_cost_per_token": 6e-07, + "supports_tool_choice": true }, "vercel_ai_gateway/meta/llama-4-scout": { "input_cost_per_token": 1e-07, @@ -28294,7 +28378,10 @@ "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 3e-07 + "output_cost_per_token": 3e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/codestral": { "input_cost_per_token": 3e-07, @@ -28303,7 +28390,9 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 9e-07 + "output_cost_per_token": 9e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/codestral-embed": { "input_cost_per_token": 1.5e-07, @@ -28321,7 +28410,10 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2.8e-07 + "output_cost_per_token": 2.8e-07, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/mistral/magistral-medium": { "input_cost_per_token": 2e-06, @@ -28330,7 +28422,10 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 5e-06 + "output_cost_per_token": 5e-06, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/mistral/magistral-small": { "input_cost_per_token": 5e-07, @@ -28339,7 +28434,8 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 1.5e-06 + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true }, "vercel_ai_gateway/mistral/ministral-3b": { "input_cost_per_token": 4e-08, @@ -28348,7 +28444,9 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 4e-08 + "output_cost_per_token": 4e-08, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/ministral-8b": { "input_cost_per_token": 1e-07, @@ -28357,7 +28455,10 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 1e-07 + "output_cost_per_token": 1e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/mistral-embed": { "input_cost_per_token": 1e-07, @@ -28375,7 +28476,9 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 6e-06 + "output_cost_per_token": 6e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/mistral/mistral-saba-24b": { "input_cost_per_token": 7.9e-07, @@ -28393,7 +28496,10 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 3e-07 + "output_cost_per_token": 3e-07, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/mistral/mixtral-8x22b-instruct": { "input_cost_per_token": 1.2e-06, @@ -28402,7 +28508,8 @@ "max_output_tokens": 2048, "max_tokens": 2048, "mode": "chat", - "output_cost_per_token": 1.2e-06 + "output_cost_per_token": 1.2e-06, + "supports_function_calling": true }, "vercel_ai_gateway/mistral/pixtral-12b": { "input_cost_per_token": 1.5e-07, @@ -28411,7 +28518,11 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 1.5e-07 + "output_cost_per_token": 1.5e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/mistral/pixtral-large": { "input_cost_per_token": 2e-06, @@ -28420,7 +28531,11 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 6e-06 + "output_cost_per_token": 6e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/moonshotai/kimi-k2": { "input_cost_per_token": 5.5e-07, @@ -28429,7 +28544,9 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 2.2e-06 + "output_cost_per_token": 2.2e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/morph/morph-v3-fast": { "input_cost_per_token": 8e-07, @@ -28456,7 +28573,9 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 1.5e-06 + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/openai/gpt-3.5-turbo-instruct": { "input_cost_per_token": 1.5e-06, @@ -28474,7 +28593,10 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "chat", - "output_cost_per_token": 3e-05 + "output_cost_per_token": 3e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/openai/gpt-4.1": { "cache_creation_input_token_cost": 0.0, @@ -28485,7 +28607,11 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 8e-06 + "output_cost_per_token": 8e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/gpt-4.1-mini": { "cache_creation_input_token_cost": 0.0, @@ -28496,7 +28622,11 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1.6e-06 + "output_cost_per_token": 1.6e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/gpt-4.1-nano": { "cache_creation_input_token_cost": 0.0, @@ -28507,7 +28637,11 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 4e-07 + "output_cost_per_token": 4e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/gpt-4o": { "cache_creation_input_token_cost": 0.0, @@ -28518,7 +28652,11 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/gpt-4o-mini": { "cache_creation_input_token_cost": 0.0, @@ -28529,7 +28667,11 @@ "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 6e-07 + "output_cost_per_token": 6e-07, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/o1": { "cache_creation_input_token_cost": 0.0, @@ -28540,7 +28682,11 @@ "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", - "output_cost_per_token": 6e-05 + "output_cost_per_token": 6e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/o3": { "cache_creation_input_token_cost": 0.0, @@ -28551,7 +28697,11 @@ "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", - "output_cost_per_token": 8e-06 + "output_cost_per_token": 8e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/o3-mini": { "cache_creation_input_token_cost": 0.0, @@ -28562,7 +28712,10 @@ "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", - "output_cost_per_token": 4.4e-06 + "output_cost_per_token": 4.4e-06, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/o4-mini": { "cache_creation_input_token_cost": 0.0, @@ -28573,7 +28726,11 @@ "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", - "output_cost_per_token": 4.4e-06 + "output_cost_per_token": 4.4e-06, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true }, "vercel_ai_gateway/openai/text-embedding-3-large": { "input_cost_per_token": 1.3e-07, @@ -28645,7 +28802,10 @@ "max_output_tokens": 32000, "max_tokens": 32000, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/vercel/v0-1.5-md": { "input_cost_per_token": 3e-06, @@ -28654,7 +28814,10 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-2": { "input_cost_per_token": 2e-06, @@ -28663,7 +28826,9 @@ "max_output_tokens": 4000, "max_tokens": 4000, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-2-vision": { "input_cost_per_token": 2e-06, @@ -28672,7 +28837,10 @@ "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 1e-05 + "output_cost_per_token": 1e-05, + "supports_vision": true, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-3": { "input_cost_per_token": 3e-06, @@ -28681,7 +28849,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-3-fast": { "input_cost_per_token": 5e-06, @@ -28690,7 +28860,8 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.5e-05 + "output_cost_per_token": 2.5e-05, + "supports_function_calling": true }, "vercel_ai_gateway/xai/grok-3-mini": { "input_cost_per_token": 3e-07, @@ -28699,7 +28870,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 5e-07 + "output_cost_per_token": 5e-07, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-3-mini-fast": { "input_cost_per_token": 6e-07, @@ -28708,7 +28881,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 4e-06 + "output_cost_per_token": 4e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/xai/grok-4": { "input_cost_per_token": 3e-06, @@ -28717,7 +28892,9 @@ "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-05 + "output_cost_per_token": 1.5e-05, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/zai/glm-4.5": { "input_cost_per_token": 6e-07, @@ -28726,7 +28903,9 @@ "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.2e-06 + "output_cost_per_token": 2.2e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/zai/glm-4.5-air": { "input_cost_per_token": 2e-07, @@ -28735,7 +28914,9 @@ "max_output_tokens": 96000, "max_tokens": 96000, "mode": "chat", - "output_cost_per_token": 1.1e-06 + "output_cost_per_token": 1.1e-06, + "supports_function_calling": true, + "supports_tool_choice": true }, "vercel_ai_gateway/zai/glm-4.6": { "litellm_provider": "vercel_ai_gateway", @@ -29912,7 +30093,9 @@ "mode": "chat", "output_cost_per_token": 1e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_regions": ["global"], + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29925,7 +30108,9 @@ "mode": "chat", "output_cost_per_token": 4e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_regions": ["global"], + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29938,7 +30123,9 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_regions": ["global"], + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29951,7 +30138,9 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_regions": ["global"], + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_tool_choice": true }, @@ -34900,4 +35089,4 @@ "output_cost_per_token": 0, "supports_reasoning": true } -} \ No newline at end of file +} diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index ed7595bb023..71e7798b09e 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -40,6 +40,7 @@ IGNORE_FUNCTIONS = [ "filter_exceptions_from_params", # max depth set (default 20) to prevent infinite recursion. "__getattr__", # lazy loading pattern in litellm/__init__.py with proper caching to prevent infinite recursion. "_validate_inheritance_chain", # max depth set (default 100) to prevent infinite recursion in policy inheritance validation. + "_basic_json_schema_validate", # max depth set. "extract_text_from_a2a_message", # max depth set (default 10) to prevent infinite recursion in A2A message parsing. ] diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index 0a57d046c72..c39454728a8 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -1,4 +1,3 @@ -import io import os import sys @@ -10,13 +9,10 @@ from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, call, patch import pytest -from prometheus_client import REGISTRY, CollectorRegistry +from prometheus_client import REGISTRY import litellm -from litellm import completion from litellm._logging import verbose_logger -from litellm._uuid import uuid -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.utils import ( StandardLoggingHiddenParams, StandardLoggingMetadata, @@ -37,7 +33,6 @@ from litellm.proxy._types import UserAPIKeyAuth verbose_logger.setLevel(logging.DEBUG) litellm.set_verbose = True -import time @pytest.fixture @@ -293,7 +288,6 @@ async def test_increment_remaining_budget_metrics(prometheus_logger): ) as mock_get_team, patch( "litellm.proxy.auth.auth_checks.get_key_object" ) as mock_get_key: - mock_get_team.return_value = MagicMock(budget_reset_at=future_reset_time_team) mock_get_key.return_value = MagicMock(budget_reset_at=future_reset_time_key) @@ -648,25 +642,16 @@ async def test_async_log_failure_event(prometheus_logger): ) # litellm_llm_api_failed_requests_metric incremented - """ - Expected metrics - end_user_id, - user_api_key, - user_api_key_alias, - model, - user_api_team, - user_api_team_alias, - user_id, - """ + # Labels: end_user, api_key_hash, api_key_alias, model, team, team_alias, user, model_id prometheus_logger.litellm_llm_api_failed_requests_metric.labels.assert_called_once_with( - None, + None, # end_user_id "test_hash", "test_alias", "gpt-3.5-turbo", "test_team", "test_team_alias", "test_user", - "model-123", + "model-123", # model_id from standard_logging_payload ) prometheus_logger.litellm_llm_api_failed_requests_metric.labels().inc.assert_called_once() @@ -678,38 +663,54 @@ async def test_async_log_failure_event(prometheus_logger): api_provider="openai", ) - # deployment failure responses incremented - prometheus_logger.litellm_deployment_failure_responses.labels.assert_called_once_with( - litellm_model_name="gpt-3.5-turbo", - model_id="model-123", - api_base="https://api.openai.com", - api_provider="openai", - exception_status="None", - exception_class="Exception", - requested_model="openai-gpt", # passed in standard logging payload - hashed_api_key="test_hash", - api_key_alias="test_alias", - team="test_team", - team_alias="test_team_alias", - client_ip="127.0.0.1", # from standard logging payload - user_agent=None, + # deployment failure responses incremented - verify key labels are populated + prometheus_logger.litellm_deployment_failure_responses.labels.assert_called_once() + actual_failure_labels = ( + prometheus_logger.litellm_deployment_failure_responses.labels.call_args.kwargs ) + expected_failure_labels = { + "litellm_model_name": "gpt-3.5-turbo", + "model_id": "model-123", + "api_base": "https://api.openai.com", + "api_provider": "openai", + "exception_class": "Exception", + "requested_model": "openai-gpt", + "hashed_api_key": "test_hash", + "api_key_alias": "test_alias", + "team": "test_team", + "team_alias": "test_team_alias", + } + for key, expected_val in expected_failure_labels.items(): + assert key in actual_failure_labels, f"Missing label {key}" + assert ( + actual_failure_labels[key] == expected_val + ), f"Label {key}: expected {expected_val!r}, got {actual_failure_labels[key]!r}" + assert actual_failure_labels.get("exception_status") in ("None", None) + assert actual_failure_labels.get("client_ip") == "127.0.0.1" prometheus_logger.litellm_deployment_failure_responses.labels().inc.assert_called_once() - # deployment total requests incremented - prometheus_logger.litellm_deployment_total_requests.labels.assert_called_once_with( - litellm_model_name="gpt-3.5-turbo", - model_id="model-123", - api_base="https://api.openai.com", - api_provider="openai", - requested_model="openai-gpt", # passed in standard logging payload - hashed_api_key="test_hash", - api_key_alias="test_alias", - team="test_team", - team_alias="test_team_alias", - client_ip="127.0.0.1", # from standard logging payload - user_agent=None, + # deployment total requests incremented - verify key labels are populated + prometheus_logger.litellm_deployment_total_requests.labels.assert_called_once() + actual_total_labels = ( + prometheus_logger.litellm_deployment_total_requests.labels.call_args.kwargs ) + expected_total_labels = { + "litellm_model_name": "gpt-3.5-turbo", + "model_id": "model-123", + "api_base": "https://api.openai.com", + "api_provider": "openai", + "requested_model": "openai-gpt", + "hashed_api_key": "test_hash", + "api_key_alias": "test_alias", + "team": "test_team", + "team_alias": "test_team_alias", + } + for key, expected_val in expected_total_labels.items(): + assert key in actual_total_labels, f"Missing label {key}" + assert ( + actual_total_labels[key] == expected_val + ), f"Label {key}: expected {expected_val!r}, got {actual_total_labels[key]!r}" + assert actual_total_labels.get("client_ip") == "127.0.0.1" prometheus_logger.litellm_deployment_total_requests.labels().inc.assert_called_once() @@ -1095,7 +1096,7 @@ def test_increment_deployment_cooled_down(prometheus_logger): import inspect method_sig = inspect.signature(prometheus_logger.increment_deployment_cooled_down) - expected_label_count = len([p for p in method_sig.parameters.keys() if p != 'self']) + expected_label_count = len([p for p in method_sig.parameters.keys() if p != "self"]) mock_chain = MagicMock() @@ -1103,11 +1104,15 @@ def test_increment_deployment_cooled_down(prometheus_logger): """Validate label count matches metric definition""" total = len(label_values) + len(label_kwargs) if total != expected_label_count: - raise ValueError(f"Incorrect label count: expected {expected_label_count}, got {total}") + raise ValueError( + f"Incorrect label count: expected {expected_label_count}, got {total}" + ) return mock_chain prometheus_logger.litellm_deployment_cooled_down = MagicMock() - prometheus_logger.litellm_deployment_cooled_down.labels = MagicMock(side_effect=validating_labels) + prometheus_logger.litellm_deployment_cooled_down.labels = MagicMock( + side_effect=validating_labels + ) prometheus_logger.increment_deployment_cooled_down( litellm_model_name="gpt-3.5-turbo", @@ -1179,8 +1184,12 @@ def test_get_custom_labels_from_top_level_metadata(monkeypatch): metadata = { "requester_ip_address": "10.48.203.20", # Top-level field "user_api_key_alias": "TestAlias", # Top-level field - "requester_metadata": {"nested_field": "nested_value"}, # Nested dict (excluded) - "user_api_key_auth_metadata": {"another_nested": "value"}, # Nested dict (excluded) + "requester_metadata": { + "nested_field": "nested_value" + }, # Nested dict (excluded) + "user_api_key_auth_metadata": { + "another_nested": "value" + }, # Nested dict (excluded) } result = get_custom_labels_from_metadata(metadata) assert result == { @@ -1217,7 +1226,9 @@ def test_get_custom_labels_from_top_level_and_nested_metadata(monkeypatch): } -async def test_async_log_success_event_with_top_level_metadata(prometheus_logger, monkeypatch): +async def test_async_log_success_event_with_top_level_metadata( + prometheus_logger, monkeypatch +): """ Test that async_log_success_event correctly extracts custom labels from top-level metadata fields like requester_ip_address, not just from nested dictionaries. @@ -1231,7 +1242,9 @@ async def test_async_log_success_event_with_top_level_metadata(prometheus_logger standard_logging_object = create_standard_logging_payload() standard_logging_object["metadata"]["requester_ip_address"] = "10.48.203.20" standard_logging_object["metadata"]["requester_metadata"] = {} # Empty nested dict - standard_logging_object["metadata"]["user_api_key_auth_metadata"] = {} # Empty nested dict + standard_logging_object["metadata"][ + "user_api_key_auth_metadata" + ] = {} # Empty nested dict kwargs = { "model": "gpt-3.5-turbo", @@ -1273,7 +1286,9 @@ async def test_async_log_success_event_with_top_level_metadata(prometheus_logger prometheus_logger.litellm_remaining_user_budget_metric = create_mock_metric() prometheus_logger.litellm_user_max_budget_metric = create_mock_metric() prometheus_logger.litellm_user_budget_remaining_hours_metric = create_mock_metric() - prometheus_logger.litellm_remaining_api_key_requests_for_model = create_mock_metric() + prometheus_logger.litellm_remaining_api_key_requests_for_model = ( + create_mock_metric() + ) prometheus_logger.litellm_remaining_api_key_tokens_for_model = create_mock_metric() prometheus_logger.litellm_llm_api_time_to_first_token_metric = create_mock_metric() prometheus_logger.litellm_llm_api_latency_metric = create_mock_metric() @@ -1302,7 +1317,7 @@ async def test_async_log_success_event_with_top_level_metadata(prometheus_logger # This confirms that the custom label extraction logic ran without errors assert prometheus_logger.litellm_requests_metric.labels.called assert prometheus_logger.litellm_spend_metric.labels.called - + # Verify that the labels() method was called with some arguments (either positional or keyword) # This ensures the custom label extraction happened and didn't cause a "Incorrect label names" error call_args = prometheus_logger.litellm_requests_metric.labels.call_args @@ -1494,7 +1509,6 @@ async def test_initialize_remaining_budget_metrics(prometheus_logger): with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.management_endpoints.team_endpoints.get_paginated_teams" ) as mock_get_teams: - # Create mock team data with proper datetime objects for budget_reset_at future_reset = datetime.now() + timedelta(hours=24) # Reset 24 hours from now mock_teams = [ @@ -1592,21 +1606,22 @@ async def test_initialize_remaining_budget_metrics_exception_handling( ) as mock_get_teams, patch( "litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper" ) as mock_list_keys: - # Make get_paginated_teams raise an exception mock_get_teams.side_effect = Exception("Database error") mock_list_keys.side_effect = Exception("Key listing error") - + # Mock prisma_client structure to raise an exception for user budget metrics # The code accesses prisma_client.db.litellm_usertable.find_many and count mock_usertable = MagicMock() - mock_usertable.find_many = MagicMock(side_effect=Exception("User database error")) + mock_usertable.find_many = MagicMock( + side_effect=Exception("User database error") + ) mock_usertable.count = MagicMock(side_effect=Exception("User count error")) - + # Mock litellm_teamtable to raise an exception for team count metrics mock_teamtable = MagicMock() mock_teamtable.count = MagicMock(side_effect=Exception("Team count error")) - + mock_db = MagicMock() mock_db.litellm_usertable = mock_usertable mock_db.litellm_teamtable = mock_teamtable @@ -1661,7 +1676,6 @@ async def test_initialize_api_key_budget_metrics(prometheus_logger): with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper" ) as mock_list_keys: - # Create mock key data with proper datetime objects for budget_reset_at future_reset = datetime.now() + timedelta(hours=24) # Reset 24 hours from now key1 = UserAPIKeyAuth( @@ -1916,7 +1930,6 @@ def test_prometheus_label_factory_with_custom_tags(monkeypatch): Test that prometheus_label_factory correctly handles custom tags """ from litellm.integrations.prometheus import ( - get_custom_labels_from_tags, prometheus_label_factory, ) from litellm.types.integrations.prometheus import UserAPIKeyLabelValues @@ -1954,7 +1967,6 @@ def test_prometheus_label_factory_with_no_custom_tags(monkeypatch): Test that prometheus_label_factory works when no custom tags are configured """ from litellm.integrations.prometheus import ( - get_custom_labels_from_tags, prometheus_label_factory, ) from litellm.types.integrations.prometheus import UserAPIKeyLabelValues @@ -2179,9 +2191,7 @@ async def test_prometheus_token_metrics_with_prometheus_config(): All three metrics should be properly incremented when making a successful completion request. """ - from prometheus_client import CollectorRegistry, Counter - import litellm from litellm.types.integrations.prometheus import PrometheusMetricsConfig # Clear registry before test diff --git a/tests/litellm/test_proxy_auth.py b/tests/litellm/test_proxy_auth.py new file mode 100644 index 00000000000..1d73e143e10 --- /dev/null +++ b/tests/litellm/test_proxy_auth.py @@ -0,0 +1,204 @@ +""" +Unit tests for litellm.proxy_auth module. + +Tests the OAuth2/JWT token management for LiteLLM Proxy authentication. +""" + +import time +from unittest.mock import Mock, patch + +import pytest + +from litellm.proxy_auth import ( + AccessToken, + AzureADCredential, + GenericOAuth2Credential, + ProxyAuthHandler, +) + + +class TestAccessToken: + """Tests for AccessToken dataclass.""" + + def test_access_token_creation(self): + """Test AccessToken can be created with required fields.""" + token = AccessToken(token="test-token", expires_on=1234567890) + assert token.token == "test-token" + assert token.expires_on == 1234567890 + + def test_access_token_equality(self): + """Test AccessToken equality comparison.""" + token1 = AccessToken(token="test", expires_on=123) + token2 = AccessToken(token="test", expires_on=123) + assert token1 == token2 + + +class MockCredential: + """Mock credential for testing.""" + + def __init__(self, expires_in_seconds: int = 3600): + self.call_count = 0 + self.expires_in = expires_in_seconds + + def get_token(self, scope: str) -> AccessToken: + self.call_count += 1 + return AccessToken( + token=f"mock-token-{self.call_count}", + expires_on=int(time.time()) + self.expires_in, + ) + + +class TestProxyAuthHandler: + """Tests for ProxyAuthHandler.""" + + def test_get_auth_headers_returns_bearer_token(self): + """Test that get_auth_headers returns correct Authorization header.""" + cred = MockCredential() + handler = ProxyAuthHandler(credential=cred, scope="test-scope") + + headers = handler.get_auth_headers() + + assert "Authorization" in headers + assert headers["Authorization"].startswith("Bearer ") + assert "mock-token-1" in headers["Authorization"] + + def test_token_caching(self): + """Test that tokens are cached and not re-requested.""" + cred = MockCredential(expires_in_seconds=3600) # Long expiry + handler = ProxyAuthHandler(credential=cred, scope="test-scope") + + # Multiple calls should only request token once + handler.get_auth_headers() + handler.get_auth_headers() + handler.get_auth_headers() + + assert cred.call_count == 1 + + def test_token_refresh_when_about_to_expire(self): + """Test that tokens are refreshed when about to expire (within 60s buffer).""" + cred = MockCredential(expires_in_seconds=30) # Expires in 30s (< 60s buffer) + handler = ProxyAuthHandler(credential=cred, scope="test-scope") + + # First call gets token + handler.get_auth_headers() + # Second call should refresh because token expires within 60s buffer + handler.get_auth_headers() + + assert cred.call_count == 2 + + def test_get_token_method(self): + """Test the get_token method returns AccessToken.""" + cred = MockCredential() + handler = ProxyAuthHandler(credential=cred, scope="test-scope") + + token = handler.get_token() + + assert isinstance(token, AccessToken) + assert token.token == "mock-token-1" + + +class TestAzureADCredential: + """Tests for AzureADCredential.""" + + def test_lazy_initialization(self): + """Test that azure-identity is not imported until get_token is called.""" + # This should not raise ImportError even if azure-identity is not installed + cred = AzureADCredential(credential=None) + # _initialized should be False until get_token is called + assert cred._initialized is False + + def test_wraps_azure_credential(self): + """Test that AzureADCredential wraps an azure-identity credential.""" + # Mock Azure credential + mock_azure_cred = Mock() + mock_azure_cred.get_token.return_value = Mock( + token="azure-token", expires_on=9999999999 + ) + + cred = AzureADCredential(credential=mock_azure_cred) + token = cred.get_token("https://graph.microsoft.com/.default") + + assert token.token == "azure-token" + assert token.expires_on == 9999999999 + mock_azure_cred.get_token.assert_called_once_with( + "https://graph.microsoft.com/.default" + ) + + +class TestGenericOAuth2Credential: + """Tests for GenericOAuth2Credential.""" + + def test_token_request(self): + """Test that GenericOAuth2Credential makes correct OAuth2 request.""" + with patch("httpx.post") as mock_post: + mock_response = Mock() + mock_response.json.return_value = { + "access_token": "oauth2-token", + "expires_in": 3600, + } + mock_response.raise_for_status = Mock() + mock_post.return_value = mock_response + + cred = GenericOAuth2Credential( + client_id="test-client", + client_secret="test-secret", + token_url="https://example.com/oauth2/token", + ) + token = cred.get_token("test-scope") + + assert token.token == "oauth2-token" + mock_post.assert_called_once() + call_kwargs = mock_post.call_args + assert call_kwargs[1]["data"]["grant_type"] == "client_credentials" + assert call_kwargs[1]["data"]["client_id"] == "test-client" + assert call_kwargs[1]["data"]["client_secret"] == "test-secret" + assert call_kwargs[1]["data"]["scope"] == "test-scope" + + def test_token_caching(self): + """Test that GenericOAuth2Credential caches tokens.""" + with patch("httpx.post") as mock_post: + mock_response = Mock() + mock_response.json.return_value = { + "access_token": "oauth2-token", + "expires_in": 3600, + } + mock_response.raise_for_status = Mock() + mock_post.return_value = mock_response + + cred = GenericOAuth2Credential( + client_id="test-client", + client_secret="test-secret", + token_url="https://example.com/oauth2/token", + ) + + # Multiple calls should only make one HTTP request + cred.get_token("test-scope") + cred.get_token("test-scope") + cred.get_token("test-scope") + + assert mock_post.call_count == 1 + + +class TestLiteLLMIntegration: + """Tests for integration with litellm module.""" + + def test_proxy_auth_variable_exists(self): + """Test that litellm.proxy_auth variable exists.""" + import litellm + + # Should be None by default + assert hasattr(litellm, "proxy_auth") + + def test_proxy_auth_can_be_set(self): + """Test that litellm.proxy_auth can be set to a ProxyAuthHandler.""" + import litellm + + original_value = litellm.proxy_auth + try: + cred = MockCredential() + handler = ProxyAuthHandler(credential=cred, scope="test") + litellm.proxy_auth = handler + + assert litellm.proxy_auth is handler + finally: + litellm.proxy_auth = original_value diff --git a/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py index adbaf219079..57352eafaf1 100644 --- a/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py @@ -134,3 +134,25 @@ def test_transform_request_with_response_format(): assert result["text"]["format"]["type"] == "json_schema" assert result["text"]["format"]["name"] == "person_schema" assert "schema" in result["text"]["format"] + + +def test_transform_request_includes_extra_headers(): + """Test that transform_request forwards headers as extra_headers for upstream call.""" + handler = LiteLLMResponsesTransformationHandler() + messages = [{"role": "user", "content": "Hello"}] + optional_params = {} + litellm_params = {} + + class MockLoggingObj: + pass + + headers = {"cf-aig-authorization": "secret-token"} + result = handler.transform_request( + model="gpt-5-pro", + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + litellm_logging_obj=MockLoggingObj(), + ) + assert result.get("extra_headers") == headers diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py new file mode 100644 index 00000000000..82517b7af9e --- /dev/null +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -0,0 +1,236 @@ +""" +Unit tests for Anthropic Messages Guardrail Translation Handler + +Tests the handler's ability to process streaming output for Anthropic Messages API +with guardrail transformations, specifically testing edge cases with empty choices. +""" + +import os +import sys +from typing import Any, List, Literal, Optional +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../../../..") +) # Adds the parent directory to the system path + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.anthropic.chat.guardrail_translation.handler import ( + AnthropicMessagesHandler, +) +from litellm.types.utils import GenericGuardrailAPIInputs + + +class MockPassThroughGuardrail(CustomGuardrail): + """Mock guardrail that passes through without blocking - for testing streaming fallback behavior""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + """Simply return inputs unchanged""" + return inputs + + +class MockDynamicGuardrail(CustomGuardrail): + """Mock guardrail that records dynamic params from request metadata.""" + + def __init__(self, guardrail_name: str): + super().__init__(guardrail_name=guardrail_name) + self.dynamic_params: Optional[dict] = None + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.dynamic_params = self.get_guardrail_dynamic_request_body_params( + request_data + ) + return inputs + + +class TestAnthropicMessagesHandlerStreamingOutputProcessing: + """Test streaming output processing functionality""" + + @pytest.mark.asyncio + async def test_process_output_streaming_response_empty_model_response(self): + """Test that streaming response with None model_response doesn't raise error + + This test verifies the fix for the bug where accessing model_response.choices[0] + would raise an error when _build_complete_streaming_response returns None. + """ + handler = AnthropicMessagesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Mock _check_streaming_has_ended to return True (stream ended) + # and _build_complete_streaming_response to return None + with patch.object( + handler, "_check_streaming_has_ended", return_value=True + ), patch( + "litellm.llms.anthropic.chat.guardrail_translation.handler.AnthropicPassthroughLoggingHandler._build_complete_streaming_response", + return_value=None, + ): + responses_so_far = [b"data: some chunk"] + + # This should not raise an error + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=MagicMock(), + ) + + # Should return the responses unchanged + assert result == responses_so_far + + +class TestAnthropicMessagesHandlerInputProcessing: + """Test input processing preserves litellm_metadata for dynamic guardrails.""" + + @pytest.mark.asyncio + async def test_process_input_messages_preserves_litellm_metadata_guardrails(self): + handler = AnthropicMessagesHandler() + guardrail = MockDynamicGuardrail(guardrail_name="cygnal-monitor") + + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [{"role": "user", "content": "hello"}], + "litellm_metadata": { + "guardrails": [ + { + "cygnal-monitor": { + "extra_body": {"policy_id": "policy-123"} + } + } + ] + }, + } + + with patch("litellm.proxy.proxy_server.premium_user", True): + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data.get("litellm_metadata", {}).get("guardrails") + assert guardrail.dynamic_params == {"policy_id": "policy-123"} + + @pytest.mark.asyncio + async def test_process_output_streaming_response_empty_choices(self): + """Test that streaming response with empty choices doesn't raise IndexError + + This test verifies the fix for the bug where accessing model_response.choices[0] + would raise IndexError when the response has an empty choices list. + """ + from litellm.types.utils import ModelResponse + + handler = AnthropicMessagesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Create a mock response with empty choices + mock_response = ModelResponse( + id="msg_123", + created=1234567890, + model="claude-3", + object="chat.completion", + choices=[], # Empty choices + ) + + # Mock _check_streaming_has_ended to return True (stream ended) + # and _build_complete_streaming_response to return the mock response + with patch.object( + handler, "_check_streaming_has_ended", return_value=True + ), patch( + "litellm.llms.anthropic.chat.guardrail_translation.handler.AnthropicPassthroughLoggingHandler._build_complete_streaming_response", + return_value=mock_response, + ): + responses_so_far = [b"data: some chunk"] + + # This should not raise IndexError + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=MagicMock(), + ) + + # Should return the responses unchanged + assert result == responses_so_far + + @pytest.mark.asyncio + async def test_process_output_streaming_response_with_valid_choices(self): + """Test that streaming response with valid choices still works correctly""" + from litellm.types.utils import Choices, Message, ModelResponse + + handler = AnthropicMessagesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Create a mock response with valid choices + mock_response = ModelResponse( + id="msg_123", + created=1234567890, + model="claude-3", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="Hello world", + role="assistant", + ), + ) + ], + ) + + # Mock _check_streaming_has_ended to return True (stream ended) + # and _build_complete_streaming_response to return the mock response + with patch.object( + handler, "_check_streaming_has_ended", return_value=True + ), patch( + "litellm.llms.anthropic.chat.guardrail_translation.handler.AnthropicPassthroughLoggingHandler._build_complete_streaming_response", + return_value=mock_response, + ): + responses_so_far = [b"data: some chunk"] + + # This should process successfully + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=MagicMock(), + ) + + # Should return the responses + assert result == responses_so_far + + @pytest.mark.asyncio + async def test_process_output_streaming_response_stream_not_ended(self): + """Test that streaming response falls back to text processing when stream hasn't ended""" + handler = AnthropicMessagesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Mock _check_streaming_has_ended to return False (stream not ended) + with patch.object( + handler, "_check_streaming_has_ended", return_value=False + ), patch.object( + handler, "get_streaming_string_so_far", return_value="partial text" + ): + responses_so_far = [b"data: some chunk"] + + # This should process successfully using text-based guardrail + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=MagicMock(), + ) + + # Should return the responses + assert result == responses_so_far + + +if __name__ == "__main__": + # Run the tests + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 6c0195d2831..1f5f53d0f0c 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -733,6 +733,154 @@ class TestOpenAIChatCompletionsHandlerToolCallsOutput: assert response.choices[0].finish_reason == "tool_calls" +class MockPassThroughGuardrail(CustomGuardrail): + """Mock guardrail that passes through without blocking - for testing streaming fallback behavior""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + """Simply return inputs unchanged""" + return inputs + + +class TestOpenAIChatCompletionsHandlerStreamingOutput: + """Test streaming output processing functionality""" + + @pytest.mark.asyncio + async def test_process_output_streaming_response_empty_choices(self): + """Test that streaming response with empty choices doesn't raise IndexError + + This test verifies the fix for the bug where accessing chunk.choices[0] + would raise IndexError when a streaming chunk has an empty choices list. + """ + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + handler = OpenAIChatCompletionsHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Create a streaming chunk with empty choices + chunk_with_empty_choices = ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[], # Empty choices - this was causing the IndexError + ) + + responses_so_far = [chunk_with_empty_choices] + + # This should not raise IndexError + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + # Should return the responses unchanged + assert result == responses_so_far + + @pytest.mark.asyncio + async def test_process_output_streaming_response_with_valid_choices(self): + """Test that streaming response with valid choices still works correctly""" + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + handler = OpenAIChatCompletionsHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Create streaming chunks with valid choices + chunk1 = ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Hello"), + finish_reason=None, + ) + ], + ) + + chunk2 = ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=" world"), + finish_reason="stop", + ) + ], + ) + + responses_so_far = [chunk1, chunk2] + + # This should process successfully + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + # Should return the responses + assert result == responses_so_far + + @pytest.mark.asyncio + async def test_process_output_streaming_response_mixed_empty_and_valid_choices_no_finish(self): + """Test streaming response with mix of empty and valid choices chunks (stream not finished) + + This tests the has_stream_ended check when iterating through chunks with mixed choices. + The stream hasn't finished yet (no finish_reason), so it won't trigger stream_chunk_builder. + """ + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + handler = OpenAIChatCompletionsHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Mix of chunks - some with empty choices, some with valid choices + # Stream hasn't finished (no finish_reason) + chunk_empty = ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[], + ) + + chunk_valid = ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Hello"), + finish_reason=None, # Stream not finished + ) + ], + ) + + responses_so_far = [chunk_empty, chunk_valid] + + # This should not raise IndexError when checking has_stream_ended + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + # Should return the responses + assert result == responses_so_far + + if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index a2849ab91a2..ccece8018ff 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -817,3 +817,181 @@ class TestOpenAIResponsesHandlerToolCallExtraction: assert task_mappings[0] == (0, 0) assert task_mappings[1] == (0, 1) assert task_mappings[2] == (0, 2) + + +class MockPassThroughGuardrail(CustomGuardrail): + """Mock guardrail that passes through without blocking - for testing streaming fallback behavior""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + """Simply return inputs unchanged""" + return inputs + + +class TestOpenAIResponsesHandlerStreamingOutputProcessing: + """Test streaming output processing functionality""" + + @pytest.mark.asyncio + async def test_process_output_streaming_response_empty_output(self): + """Test that streaming response with empty output doesn't raise IndexError + + This test verifies the fix for the bug where accessing model_response_choices[0] + would raise IndexError when the response.completed event has an empty output array. + """ + handler = OpenAIResponsesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Simulate a response.completed streaming event with empty output + responses_so_far = [ + { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [], # Empty output - this was causing the IndexError + "status": "completed", + }, + } + ] + + # This should not raise IndexError + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + # Should return the responses unchanged + assert result == responses_so_far + + @pytest.mark.asyncio + async def test_process_output_streaming_response_missing_output_key(self): + """Test that streaming response with missing output key doesn't raise IndexError + + This test verifies the handler gracefully handles when the response dict + doesn't contain an 'output' key at all. + """ + handler = OpenAIResponsesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Simulate a response.completed streaming event with missing output key + responses_so_far = [ + { + "type": "response.completed", + "response": { + "id": "resp_123", + "status": "completed", + # No 'output' key - get() will return [] + }, + } + ] + + # This should not raise IndexError + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + # Should return the responses unchanged + assert result == responses_so_far + + @pytest.mark.asyncio + async def test_process_output_streaming_response_unrecognized_output_type(self): + """Test that streaming response with unrecognized output types doesn't raise IndexError + + This test verifies the handler gracefully handles when output items are of + unrecognized types that _convert_response_output_to_choices skips over. + """ + handler = OpenAIResponsesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Simulate a response.completed streaming event with unrecognized output type + responses_so_far = [ + { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [ + { + "type": "unknown_type", # Unrecognized type + "id": "item_123", + "data": "some data", + } + ], + "status": "completed", + }, + } + ] + + # This should not raise IndexError + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + # Should return the responses unchanged + assert result == responses_so_far + + @pytest.mark.asyncio + async def test_process_output_streaming_response_with_valid_output(self): + """Test that streaming response with valid output still works correctly""" + handler = OpenAIResponsesHandler() + guardrail = MockPassThroughGuardrail(guardrail_name="test") + + # Simulate a response.completed streaming event with valid message output + responses_so_far = [ + { + "type": "response.created", + "response": {"id": "resp_123"}, + }, + { + "type": "response.output_item.added", + "item": {"type": "message", "id": "msg_123"}, + }, + { + "type": "response.content_part.added", + "part": {"type": "output_text", "text": ""}, + }, + { + "type": "response.output_text.delta", + "delta": "Hello", + }, + { + "type": "response.output_text.delta", + "delta": " world", + }, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "output": [ + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "Hello world"}, + ], + } + ], + "status": "completed", + }, + }, + ] + + # This should process successfully + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + # Should return the responses + assert result == responses_so_far diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py index 6dc658827bc..109ad0bfdc8 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py @@ -34,8 +34,9 @@ def test_prepare_payload_uses_dynamic_overrides( "policy_id": "dynamic-policy", "reasoning_mode": "thinking", } + request_data = {} - payload = grayswan_guardrail._prepare_payload(messages, dynamic_body) + payload = grayswan_guardrail._prepare_payload(messages, dynamic_body, request_data) assert payload["messages"] == messages assert payload["categories"] == {"custom": "override"} @@ -47,14 +48,27 @@ def test_prepare_payload_falls_back_to_guardrail_defaults( grayswan_guardrail: GraySwanGuardrail, ) -> None: messages = [{"role": "user", "content": "hello"}] + request_data = {} - payload = grayswan_guardrail._prepare_payload(messages, {}) + payload = grayswan_guardrail._prepare_payload(messages, {}, request_data) assert payload["categories"] == {"safety": "general policy"} assert payload["policy_id"] == "default-policy" assert payload["reasoning_mode"] == "hybrid" +def test_prepare_payload_includes_dynamic_metadata( + grayswan_guardrail: GraySwanGuardrail, +) -> None: + messages = [{"role": "user", "content": "hello"}] + dynamic_body = {"metadata": {"trace_id": "trace-123", "tags": ["a", "b"]}} + request_data = {} + + payload = grayswan_guardrail._prepare_payload(messages, dynamic_body, request_data) + + assert payload["metadata"] == dynamic_body["metadata"] + + def test_process_response_does_not_block_under_threshold( grayswan_guardrail: GraySwanGuardrail, ) -> None: @@ -160,6 +174,119 @@ async def test_run_guardrail_raises_api_error( await grayswan_guardrail.run_grayswan_guardrail(payload) +@pytest.mark.asyncio +async def test_apply_guardrail_passthrough_not_swallowed_by_fail_open( + monkeypatch, +) -> None: + guardrail = GraySwanGuardrail( + guardrail_name="grayswan-passthrough", + api_key="test-key", + on_flagged_action="passthrough", + violation_threshold=0.2, + fail_open=True, + event_hook=GuardrailEventHooks.pre_call, + ) + + async def _fake_call(_payload: dict): + return {"violation": 0.92, "violated_rule_descriptions": []} + + monkeypatch.setattr(guardrail, "_call_grayswan_api", _fake_call) + + with pytest.raises(ModifyResponseException): + await guardrail.apply_guardrail( + inputs={"texts": ["bad"]}, + request_data={"model": "gpt-4"}, + input_type="request", + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_block_not_swallowed_by_fail_open( + monkeypatch, +) -> None: + guardrail = GraySwanGuardrail( + guardrail_name="grayswan-block", + api_key="test-key", + on_flagged_action="block", + violation_threshold=0.2, + fail_open=True, + event_hook=GuardrailEventHooks.pre_call, + ) + + async def _fake_call(_payload: dict): + return {"violation": 0.92, "violated_rule_descriptions": []} + + monkeypatch.setattr(guardrail, "_call_grayswan_api", _fake_call) + + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs={"texts": ["bad"]}, + request_data={"model": "gpt-4"}, + input_type="request", + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_non_grayswan_http_exception_fail_open_true( + monkeypatch, +) -> None: + guardrail = GraySwanGuardrail( + guardrail_name="grayswan-error", + api_key="test-key", + on_flagged_action="monitor", + violation_threshold=0.2, + fail_open=True, + event_hook=GuardrailEventHooks.pre_call, + ) + + async def _fake_call(_payload: dict): + return {"violation": 0.0, "violated_rule_descriptions": []} + + def _fake_process(**_kwargs): + raise HTTPException(status_code=500, detail={"error": "upstream failed"}) + + monkeypatch.setattr(guardrail, "_call_grayswan_api", _fake_call) + monkeypatch.setattr(guardrail, "_process_response_internal", _fake_process) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["ok"]}, + request_data={"model": "gpt-4"}, + input_type="request", + ) + + assert result["texts"] == ["ok"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_non_grayswan_http_exception_fail_open_false( + monkeypatch, +) -> None: + guardrail = GraySwanGuardrail( + guardrail_name="grayswan-error", + api_key="test-key", + on_flagged_action="monitor", + violation_threshold=0.2, + fail_open=False, + event_hook=GuardrailEventHooks.pre_call, + ) + + async def _fake_call(_payload: dict): + return {"violation": 0.0, "violated_rule_descriptions": []} + + def _fake_process(**_kwargs): + raise HTTPException(status_code=500, detail={"error": "upstream failed"}) + + monkeypatch.setattr(guardrail, "_call_grayswan_api", _fake_call) + monkeypatch.setattr(guardrail, "_process_response_internal", _fake_process) + + with pytest.raises(GraySwanGuardrailAPIError): + await guardrail.apply_guardrail( + inputs={"texts": ["ok"]}, + request_data={"model": "gpt-4"}, + input_type="request", + ) + + def test_process_response_passthrough_raises_exception_in_pre_call() -> None: """Test that passthrough mode raises ModifyResponseException in pre_call hook.""" guardrail = GraySwanGuardrail( diff --git a/ui/litellm-dashboard/src/components/guardrails.tsx b/ui/litellm-dashboard/src/components/guardrails.tsx index 26b9acb6a26..d0031b872f7 100644 --- a/ui/litellm-dashboard/src/components/guardrails.tsx +++ b/ui/litellm-dashboard/src/components/guardrails.tsx @@ -1,5 +1,7 @@ import React, { useState, useEffect } from "react"; import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react"; +import { Dropdown } from "antd"; +import { DownOutlined, PlusOutlined, CodeOutlined } from "@ant-design/icons"; import { getGuardrailsList, deleteGuardrailCall } from "./networking"; import AddGuardrailForm from "./guardrails/add_guardrail_form"; import GuardrailTable from "./guardrails/guardrail_table"; @@ -10,6 +12,7 @@ import NotificationsManager from "./molecules/notifications_manager"; import { Guardrail, GuardrailDefinitionLocation } from "./guardrails/types"; import DeleteResourceModal from "./common_components/DeleteResourceModal"; import { getGuardrailLogoAndName } from "./guardrails/guardrail_info_helpers"; +import { CustomCodeModal } from "./guardrails/custom_code"; interface GuardrailsPanelProps { accessToken: string | null; @@ -37,6 +40,7 @@ interface GuardrailsResponse { const GuardrailsPanel: React.FC = ({ accessToken, userRole }) => { const [guardrailsList, setGuardrailsList] = useState([]); const [isAddModalVisible, setIsAddModalVisible] = useState(false); + const [isCustomCodeModalVisible, setIsCustomCodeModalVisible] = useState(false); const [isLoading, setIsLoading] = useState(false); const [isDeleting, setIsDeleting] = useState(false); const [guardrailToDelete, setGuardrailToDelete] = useState(null); @@ -74,10 +78,21 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole setIsAddModalVisible(true); }; + const handleAddCustomCodeGuardrail = () => { + if (selectedGuardrailId) { + setSelectedGuardrailId(null); + } + setIsCustomCodeModalVisible(true); + }; + const handleCloseModal = () => { setIsAddModalVisible(false); }; + const handleCloseCustomCodeModal = () => { + setIsCustomCodeModalVisible(false); + }; + const handleSuccess = () => { fetchGuardrails(); }; @@ -128,9 +143,30 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole
- + , + label: "Add Provider Guardrail", + onClick: handleAddGuardrail, + }, + { + key: "custom_code", + icon: , + label: "Create Custom Code Guardrail", + onClick: handleAddCustomCodeGuardrail, + }, + ], + }} + trigger={["click"]} + disabled={!accessToken} + > + +
{selectedGuardrailId ? ( @@ -159,6 +195,13 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole onSuccess={handleSuccess} /> + + void; + height?: string; + placeholder?: string; + disabled?: boolean; +} + +const CustomCodeEditor: React.FC = ({ + value, + onChange, + height = "350px", + placeholder = `def apply_guardrail(inputs, request_data, input_type): + # inputs: contains texts, images, tools, tool_calls, structured_messages, model + # request_data: contains model, user_id, team_id, end_user_id, metadata + # input_type: "request" or "response" + + for text in inputs["texts"]: + # Example: Block if SSN pattern is detected + if regex_match(text, r"\\d{3}-\\d{2}-\\d{4}"): + return block("SSN detected in message") + + return allow()`, + disabled = false, +}) => { + const textareaRef = useRef(null); + const [activeTab, setActiveTab] = useState("edit"); + const [cursorPosition, setCursorPosition] = useState({ line: 1, column: 1 }); + + // Calculate cursor position + const updateCursorPosition = () => { + if (textareaRef.current) { + const textarea = textareaRef.current; + const textBeforeCursor = value.substring(0, textarea.selectionStart); + const lines = textBeforeCursor.split("\n"); + const line = lines.length; + const column = lines[lines.length - 1].length + 1; + setCursorPosition({ line, column }); + } + }; + + // Handle tab key for indentation + const handleKeyDown = (e: React.KeyboardEvent) => { + if (e.key === "Tab") { + e.preventDefault(); + const textarea = e.currentTarget; + const start = textarea.selectionStart; + const end = textarea.selectionEnd; + + // Insert 4 spaces at cursor position + const newValue = value.substring(0, start) + " " + value.substring(end); + onChange(newValue); + + // Move cursor after the inserted spaces + setTimeout(() => { + textarea.selectionStart = textarea.selectionEnd = start + 4; + }, 0); + } + }; + + const lineCount = value.split("\n").length; + + const tabItems = [ + { + key: "edit", + label: ( + + + Edit + + ), + children: ( +
+ {/* Line numbers */} +
+ {Array.from({ length: Math.max(lineCount, 15) }, (_, i) => ( +
+ {i + 1} +
+ ))} +
+ + {/* Code editor */} +