From 7056d9984ea5f7146fb346d77e4e21ee222c2be1 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Tue, 3 Feb 2026 19:57:24 -0800 Subject: [PATCH] Custom Code Guardrails UI Playground (#20377) * feat(guardrails/): allow custom code execution for guardrails first step in allowing teams to submit custom code for guardrails * feat: custom_code_guardrail.md support passing custom code for guardrails * feat: initial commit adding ui for custom code guardrails allows users to write guardrails based on custom code * feat: expose new test custom code guardrail endpoint allows ui testing playground to sanity check if guardrail is working as expected * fix: fix linting errors * fix: fix max recursion check * fix: fix linting error --- .../simple_guardrail_tutorial.md | 7 +- .../proxy/guardrails/custom_code_guardrail.md | 278 ++++++++ docs/my-website/sidebars.js | 1 + .../out/{404.html => 404/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{budgets.html => budgets/index.html} | 0 .../{caching.html => caching/index.html} | 0 .../index.html} | 0 .../{old-usage.html => old-usage/index.html} | 0 .../{prompts.html => prompts/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../out/{login.html => login/index.html} | 0 .../out/{logs.html => logs/index.html} | 0 .../{callback.html => callback/index.html} | 0 .../{model-hub.html => model-hub/index.html} | 0 .../{model_hub.html => model_hub/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{policies.html => policies/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{ui-theme.html => ui-theme/index.html} | 0 .../out/{teams.html => teams/index.html} | 0 .../{test-key.html => test-key/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../out/{usage.html => usage/index.html} | 0 .../out/{users.html => users/index.html} | 0 .../index.html} | 0 litellm/proxy/_new_secret_config.yaml | 11 + .../proxy/guardrails/guardrail_endpoints.py | 269 ++++++++ .../guardrail_hooks/custom_code/__init__.py | 65 ++ .../custom_code/custom_code_guardrail.py | 372 +++++++++++ .../guardrail_hooks/custom_code/primitives.py | 602 ++++++++++++++++++ litellm/types/guardrails.py | 23 +- .../code_coverage_tests/recursive_detector.py | 1 + .../src/components/guardrails.tsx | 49 +- .../custom_code/CustomCodeEditor.tsx | 188 ++++++ .../custom_code/CustomCodeModal.tsx | 540 ++++++++++++++++ .../custom_code/CustomCodePlayground.tsx | 588 +++++++++++++++++ .../custom_code/custom_code_constants.ts | 194 ++++++ .../guardrails/custom_code/index.ts | 1 + .../src/components/networking.tsx | 83 +++ 49 files changed, 3254 insertions(+), 18 deletions(-) create mode 100644 docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md rename litellm/proxy/_experimental/out/{404.html => 404/index.html} (100%) rename litellm/proxy/_experimental/out/{api-reference.html => api-reference/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{api-playground.html => api-playground/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{budgets.html => budgets/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{caching.html => caching/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{claude-code-plugins.html => claude-code-plugins/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{old-usage.html => old-usage/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{prompts.html => prompts/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{tag-management.html => tag-management/index.html} (100%) rename litellm/proxy/_experimental/out/{guardrails.html => guardrails/index.html} (100%) rename litellm/proxy/_experimental/out/{login.html => login/index.html} (100%) rename litellm/proxy/_experimental/out/{logs.html => logs/index.html} (100%) rename litellm/proxy/_experimental/out/mcp/oauth/{callback.html => callback/index.html} (100%) rename litellm/proxy/_experimental/out/{model-hub.html => model-hub/index.html} (100%) rename litellm/proxy/_experimental/out/{model_hub.html => model_hub/index.html} (100%) rename litellm/proxy/_experimental/out/{model_hub_table.html => model_hub_table/index.html} (100%) rename litellm/proxy/_experimental/out/{models-and-endpoints.html => models-and-endpoints/index.html} (100%) rename litellm/proxy/_experimental/out/{onboarding.html => onboarding/index.html} (100%) rename litellm/proxy/_experimental/out/{organizations.html => organizations/index.html} (100%) rename litellm/proxy/_experimental/out/{playground.html => playground/index.html} (100%) rename litellm/proxy/_experimental/out/{policies.html => policies/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{admin-settings.html => admin-settings/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{logging-and-alerts.html => logging-and-alerts/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{router-settings.html => router-settings/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{ui-theme.html => ui-theme/index.html} (100%) rename litellm/proxy/_experimental/out/{teams.html => teams/index.html} (100%) rename litellm/proxy/_experimental/out/{test-key.html => test-key/index.html} (100%) rename litellm/proxy/_experimental/out/tools/{mcp-servers.html => mcp-servers/index.html} (100%) rename litellm/proxy/_experimental/out/tools/{vector-stores.html => vector-stores/index.html} (100%) rename litellm/proxy/_experimental/out/{usage.html => usage/index.html} (100%) rename litellm/proxy/_experimental/out/{users.html => users/index.html} (100%) rename litellm/proxy/_experimental/out/{virtual-keys.html => virtual-keys/index.html} (100%) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py create mode 100644 ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeEditor.tsx create mode 100644 ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx create mode 100644 ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodePlayground.tsx create mode 100644 ui/litellm-dashboard/src/components/guardrails/custom_code/custom_code_constants.ts create mode 100644 ui/litellm-dashboard/src/components/guardrails/custom_code/index.ts 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 9c654cd156..884a7397bd 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/proxy/guardrails/custom_code_guardrail.md b/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md new file mode 100644 index 0000000000..cb24614449 --- /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/sidebars.js b/docs/my-website/sidebars.js index 98c8ee6eaa..20e66fb532 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/_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 13eeae1448..6f527e268b 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 b80b02fdc6..aec5dc1ede 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -1344,6 +1344,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 0000000000..747b188fee --- /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 0000000000..a0ca324411 --- /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 0000000000..695e59977c --- /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/types/guardrails.py b/litellm/types/guardrails.py index ed46349104..3f7de10dde 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/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index ed7595bb02..71e7798b09 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/ui/litellm-dashboard/src/components/guardrails.tsx b/ui/litellm-dashboard/src/components/guardrails.tsx index 26b9acb6a2..d0031b872f 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 */} +