mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-01 16:21:12 +00:00
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
This commit is contained in:
Krish Dholakia
@@ -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
|
||||
|
||||
@@ -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):
|
||||
...
|
||||
```
|
||||
@@ -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",
|
||||
|
||||
@@ -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()
|
||||
@@ -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 <your_api_key>" \\
|
||||
-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 (<guardrail>, 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, "<guardrail>", "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(
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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, "<guardrail>", "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
|
||||
@@ -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"</\s*(html|head|body|div|span|p|a|script|style)\s*>",
|
||||
],
|
||||
"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,
|
||||
}
|
||||
+12
-11
@@ -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=())
|
||||
|
||||
|
||||
|
||||
@@ -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.
|
||||
]
|
||||
|
||||
|
||||
@@ -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<GuardrailsPanelProps> = ({ accessToken, userRole }) => {
|
||||
const [guardrailsList, setGuardrailsList] = useState<Guardrail[]>([]);
|
||||
const [isAddModalVisible, setIsAddModalVisible] = useState(false);
|
||||
const [isCustomCodeModalVisible, setIsCustomCodeModalVisible] = useState(false);
|
||||
const [isLoading, setIsLoading] = useState(false);
|
||||
const [isDeleting, setIsDeleting] = useState(false);
|
||||
const [guardrailToDelete, setGuardrailToDelete] = useState<Guardrail | null>(null);
|
||||
@@ -74,10 +78,21 @@ const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ 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<GuardrailsPanelProps> = ({ accessToken, userRole
|
||||
<TabPanels>
|
||||
<TabPanel>
|
||||
<div className="flex justify-between items-center mb-4">
|
||||
<Button onClick={handleAddGuardrail} disabled={!accessToken}>
|
||||
+ Add New Guardrail
|
||||
</Button>
|
||||
<Dropdown
|
||||
menu={{
|
||||
items: [
|
||||
{
|
||||
key: "provider",
|
||||
icon: <PlusOutlined />,
|
||||
label: "Add Provider Guardrail",
|
||||
onClick: handleAddGuardrail,
|
||||
},
|
||||
{
|
||||
key: "custom_code",
|
||||
icon: <CodeOutlined />,
|
||||
label: "Create Custom Code Guardrail",
|
||||
onClick: handleAddCustomCodeGuardrail,
|
||||
},
|
||||
],
|
||||
}}
|
||||
trigger={["click"]}
|
||||
disabled={!accessToken}
|
||||
>
|
||||
<Button disabled={!accessToken}>
|
||||
+ Add New Guardrail <DownOutlined className="ml-2" />
|
||||
</Button>
|
||||
</Dropdown>
|
||||
</div>
|
||||
|
||||
{selectedGuardrailId ? (
|
||||
@@ -159,6 +195,13 @@ const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ accessToken, userRole
|
||||
onSuccess={handleSuccess}
|
||||
/>
|
||||
|
||||
<CustomCodeModal
|
||||
visible={isCustomCodeModalVisible}
|
||||
onClose={handleCloseCustomCodeModal}
|
||||
accessToken={accessToken}
|
||||
onSuccess={handleSuccess}
|
||||
/>
|
||||
|
||||
<DeleteResourceModal
|
||||
isOpen={isDeleteModalOpen}
|
||||
title="Delete Guardrail"
|
||||
|
||||
@@ -0,0 +1,188 @@
|
||||
import React, { useRef, useEffect, useState } from "react";
|
||||
import { Input, Tabs, Typography } from "antd";
|
||||
import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
|
||||
import { vscDarkPlus } from "react-syntax-highlighter/dist/esm/styles/prism";
|
||||
import { CodeOutlined, EyeOutlined } from "@ant-design/icons";
|
||||
|
||||
const { TextArea } = Input;
|
||||
const { Text } = Typography;
|
||||
|
||||
interface CustomCodeEditorProps {
|
||||
value: string;
|
||||
onChange: (value: string) => void;
|
||||
height?: string;
|
||||
placeholder?: string;
|
||||
disabled?: boolean;
|
||||
}
|
||||
|
||||
const CustomCodeEditor: React.FC<CustomCodeEditorProps> = ({
|
||||
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<HTMLTextAreaElement>(null);
|
||||
const [activeTab, setActiveTab] = useState<string>("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<HTMLTextAreaElement>) => {
|
||||
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: (
|
||||
<span className="flex items-center gap-1.5">
|
||||
<CodeOutlined />
|
||||
Edit
|
||||
</span>
|
||||
),
|
||||
children: (
|
||||
<div className="relative" style={{ height }}>
|
||||
{/* Line numbers */}
|
||||
<div
|
||||
className="absolute left-0 top-0 bottom-0 w-12 bg-[#1e1e1e] border-r border-[#3c3c3c] text-right pr-2 pt-3 overflow-hidden select-none"
|
||||
style={{ fontFamily: "monospace", fontSize: "13px", lineHeight: "1.5" }}
|
||||
>
|
||||
{Array.from({ length: Math.max(lineCount, 15) }, (_, i) => (
|
||||
<div key={i + 1} className="text-gray-500 h-[19.5px]">
|
||||
{i + 1}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
|
||||
{/* Code editor */}
|
||||
<textarea
|
||||
ref={textareaRef as any}
|
||||
value={value}
|
||||
onChange={(e) => {
|
||||
onChange(e.target.value);
|
||||
updateCursorPosition();
|
||||
}}
|
||||
onKeyDown={handleKeyDown}
|
||||
onClick={updateCursorPosition}
|
||||
onKeyUp={updateCursorPosition}
|
||||
placeholder={placeholder}
|
||||
disabled={disabled}
|
||||
spellCheck={false}
|
||||
className="w-full h-full pl-14 pr-4 pt-3 pb-3 font-mono text-sm resize-none focus:outline-none focus:ring-2 focus:ring-blue-500"
|
||||
style={{
|
||||
backgroundColor: "#1e1e1e",
|
||||
color: "#d4d4d4",
|
||||
border: "1px solid #3c3c3c",
|
||||
borderRadius: "8px",
|
||||
lineHeight: "1.5",
|
||||
tabSize: 4,
|
||||
}}
|
||||
/>
|
||||
|
||||
{/* Status bar */}
|
||||
<div className="absolute bottom-0 left-0 right-0 h-6 bg-[#252526] border-t border-[#3c3c3c] flex items-center justify-between px-3 text-xs text-gray-400 rounded-b-lg">
|
||||
<span>Python-like (Sandboxed)</span>
|
||||
<span>Ln {cursorPosition.line}, Col {cursorPosition.column}</span>
|
||||
</div>
|
||||
</div>
|
||||
),
|
||||
},
|
||||
{
|
||||
key: "preview",
|
||||
label: (
|
||||
<span className="flex items-center gap-1.5">
|
||||
<EyeOutlined />
|
||||
Preview
|
||||
</span>
|
||||
),
|
||||
children: (
|
||||
<div style={{ height }} className="overflow-auto rounded-lg border border-gray-200">
|
||||
<SyntaxHighlighter
|
||||
language="python"
|
||||
style={vscDarkPlus}
|
||||
showLineNumbers
|
||||
wrapLines
|
||||
customStyle={{
|
||||
margin: 0,
|
||||
borderRadius: "8px",
|
||||
fontSize: "13px",
|
||||
minHeight: height,
|
||||
}}
|
||||
lineNumberStyle={{
|
||||
minWidth: "3em",
|
||||
paddingRight: "1em",
|
||||
color: "#6e7681",
|
||||
borderRight: "1px solid #3c3c3c",
|
||||
marginRight: "1em",
|
||||
}}
|
||||
>
|
||||
{value || placeholder}
|
||||
</SyntaxHighlighter>
|
||||
</div>
|
||||
),
|
||||
},
|
||||
];
|
||||
|
||||
return (
|
||||
<div className="custom-code-editor">
|
||||
<Tabs
|
||||
activeKey={activeTab}
|
||||
onChange={setActiveTab}
|
||||
items={tabItems}
|
||||
className="custom-code-tabs"
|
||||
size="small"
|
||||
/>
|
||||
<style>{`
|
||||
.custom-code-tabs .ant-tabs-nav {
|
||||
margin-bottom: 8px;
|
||||
}
|
||||
.custom-code-tabs .ant-tabs-tab {
|
||||
padding: 4px 12px;
|
||||
}
|
||||
.custom-code-editor textarea::placeholder {
|
||||
color: #6e7681;
|
||||
}
|
||||
`}</style>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default CustomCodeEditor;
|
||||
@@ -0,0 +1,540 @@
|
||||
import React, { useState, useRef, useEffect } from "react";
|
||||
import { Modal, Select, Switch, Collapse, Input, Spin } from "antd";
|
||||
import { Button, TextInput } from "@tremor/react";
|
||||
import {
|
||||
CodeOutlined,
|
||||
PlayCircleOutlined,
|
||||
CheckCircleOutlined,
|
||||
CloseCircleOutlined,
|
||||
CaretRightOutlined,
|
||||
SaveOutlined,
|
||||
} from "@ant-design/icons";
|
||||
import { createGuardrailCall, testCustomCodeGuardrail } from "../../networking";
|
||||
import NotificationsManager from "../../molecules/notifications_manager";
|
||||
|
||||
const { Panel } = Collapse;
|
||||
const { TextArea } = Input;
|
||||
|
||||
// Code templates
|
||||
const CODE_TEMPLATES = {
|
||||
empty: {
|
||||
name: "Empty Template",
|
||||
code: `def apply_guardrail(inputs, request_data, input_type):
|
||||
# inputs: {texts, images, tools, tool_calls, structured_messages, model}
|
||||
# request_data: {model, user_id, team_id, end_user_id, metadata}
|
||||
# input_type: "request" or "response"
|
||||
return allow()`,
|
||||
},
|
||||
blockSSN: {
|
||||
name: "Block SSN",
|
||||
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()`,
|
||||
},
|
||||
redactEmail: {
|
||||
name: "Redact Emails",
|
||||
code: `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)`,
|
||||
},
|
||||
blockSQL: {
|
||||
name: "Block SQL Injection",
|
||||
code: `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()`,
|
||||
},
|
||||
validateJSON: {
|
||||
name: "Validate JSON",
|
||||
code: `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()`,
|
||||
},
|
||||
};
|
||||
|
||||
// Available primitives organized by category
|
||||
const PRIMITIVES = {
|
||||
"Return Values": [
|
||||
{ name: "allow()", desc: "Let request/response through" },
|
||||
{ name: "block(reason)", desc: "Reject with message" },
|
||||
{ name: "modify(texts=[], images=[], tool_calls=[])", desc: "Transform content" },
|
||||
],
|
||||
"Regex Functions": [
|
||||
{ name: "regex_match(text, pattern)", desc: "Returns True if pattern found" },
|
||||
{ name: "regex_replace(text, pattern, replacement)", desc: "Replace all matches" },
|
||||
{ name: "regex_find_all(text, pattern)", desc: "Return list of matches" },
|
||||
],
|
||||
"JSON Functions": [
|
||||
{ name: "json_parse(text)", desc: "Parse JSON string, returns None on error" },
|
||||
{ name: "json_stringify(obj)", desc: "Convert to JSON string" },
|
||||
{ name: "json_schema_valid(obj, schema)", desc: "Validate against JSON schema" },
|
||||
],
|
||||
"URL Functions": [
|
||||
{ name: "extract_urls(text)", desc: "Extract all URLs from text" },
|
||||
{ name: "is_valid_url(url)", desc: "Check if URL is valid" },
|
||||
{ name: "all_urls_valid(text)", desc: "Check all URLs in text are valid" },
|
||||
],
|
||||
"Code Detection": [
|
||||
{ name: "detect_code(text)", desc: "Returns True if code detected" },
|
||||
{ name: "detect_code_languages(text)", desc: "Returns list of detected languages" },
|
||||
{ name: 'contains_code_language(text, ["sql"])', desc: "Check for specific languages" },
|
||||
],
|
||||
"Text Utilities": [
|
||||
{ name: "contains(text, substring)", desc: "Check if substring exists" },
|
||||
{ name: "contains_any(text, [substr1, substr2])", desc: "Check if any substring exists" },
|
||||
{ name: "word_count(text)", desc: "Count words" },
|
||||
{ name: "char_count(text)", desc: "Count characters" },
|
||||
{ name: "lower(text) / upper(text) / trim(text)", desc: "String transforms" },
|
||||
],
|
||||
};
|
||||
|
||||
const MODE_OPTIONS = [
|
||||
{ value: "pre_call", label: "pre_call (Request)" },
|
||||
{ value: "post_call", label: "post_call (Response)" },
|
||||
{ value: "during_call", label: "during_call (Parallel)" },
|
||||
{ value: "logging_only", label: "logging_only" },
|
||||
];
|
||||
|
||||
interface CustomCodeModalProps {
|
||||
visible: boolean;
|
||||
onClose: () => void;
|
||||
onSuccess: () => void;
|
||||
accessToken: string | null;
|
||||
}
|
||||
|
||||
const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
|
||||
visible,
|
||||
onClose,
|
||||
onSuccess,
|
||||
accessToken,
|
||||
}) => {
|
||||
const [guardrailName, setGuardrailName] = useState("");
|
||||
const [mode, setMode] = useState<string>("pre_call");
|
||||
const [defaultOn, setDefaultOn] = useState(false);
|
||||
const [selectedTemplate, setSelectedTemplate] = useState<string>("empty");
|
||||
const [code, setCode] = useState(CODE_TEMPLATES.empty.code);
|
||||
const [isSaving, setIsSaving] = useState(false);
|
||||
const [isTesting, setIsTesting] = useState(false);
|
||||
const [testExpanded, setTestExpanded] = useState(false);
|
||||
const [testInput, setTestInput] = useState('{"texts": ["Hello, my SSN is 123-45-6789"], "images": [], "tools": [], "tool_calls": [], "structured_messages": [], "model": "gpt-4"}');
|
||||
const [testResult, setTestResult] = useState<any>(null);
|
||||
const [copiedPrimitive, setCopiedPrimitive] = useState<string | null>(null);
|
||||
const textareaRef = useRef<HTMLTextAreaElement>(null);
|
||||
|
||||
// Handle template change
|
||||
const handleTemplateChange = (templateKey: string) => {
|
||||
setSelectedTemplate(templateKey);
|
||||
setCode(CODE_TEMPLATES[templateKey as keyof typeof CODE_TEMPLATES].code);
|
||||
};
|
||||
|
||||
// Reset form when modal opens
|
||||
useEffect(() => {
|
||||
if (visible) {
|
||||
setGuardrailName("");
|
||||
setMode("pre_call");
|
||||
setDefaultOn(false);
|
||||
setSelectedTemplate("empty");
|
||||
setCode(CODE_TEMPLATES.empty.code);
|
||||
setTestResult(null);
|
||||
setTestExpanded(false);
|
||||
}
|
||||
}, [visible]);
|
||||
|
||||
// Copy primitive to clipboard
|
||||
const copyPrimitive = async (primitive: string) => {
|
||||
try {
|
||||
await navigator.clipboard.writeText(primitive);
|
||||
setCopiedPrimitive(primitive);
|
||||
setTimeout(() => setCopiedPrimitive(null), 2000);
|
||||
} catch (err) {
|
||||
console.error("Failed to copy:", err);
|
||||
}
|
||||
};
|
||||
|
||||
// Handle tab key in textarea
|
||||
const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
|
||||
if (e.key === "Tab") {
|
||||
e.preventDefault();
|
||||
const textarea = e.currentTarget;
|
||||
const start = textarea.selectionStart;
|
||||
const end = textarea.selectionEnd;
|
||||
const newValue = code.substring(0, start) + " " + code.substring(end);
|
||||
setCode(newValue);
|
||||
setTimeout(() => {
|
||||
textarea.selectionStart = textarea.selectionEnd = start + 4;
|
||||
}, 0);
|
||||
}
|
||||
};
|
||||
|
||||
// Save guardrail
|
||||
const handleSave = async () => {
|
||||
if (!guardrailName.trim()) {
|
||||
NotificationsManager.fromBackend("Please enter a guardrail name");
|
||||
return;
|
||||
}
|
||||
if (!code.trim()) {
|
||||
NotificationsManager.fromBackend("Please enter custom code");
|
||||
return;
|
||||
}
|
||||
if (!accessToken) {
|
||||
NotificationsManager.fromBackend("No access token available");
|
||||
return;
|
||||
}
|
||||
|
||||
setIsSaving(true);
|
||||
try {
|
||||
const guardrailData = {
|
||||
guardrail_name: guardrailName,
|
||||
litellm_params: {
|
||||
guardrail: "custom_code",
|
||||
mode: mode,
|
||||
default_on: defaultOn,
|
||||
custom_code: code,
|
||||
},
|
||||
guardrail_info: {},
|
||||
};
|
||||
|
||||
await createGuardrailCall(accessToken, guardrailData);
|
||||
NotificationsManager.success("Custom code guardrail created successfully");
|
||||
onSuccess();
|
||||
onClose();
|
||||
} catch (error) {
|
||||
console.error("Failed to create guardrail:", error);
|
||||
NotificationsManager.fromBackend(
|
||||
"Failed to create guardrail: " + (error instanceof Error ? error.message : String(error))
|
||||
);
|
||||
} finally {
|
||||
setIsSaving(false);
|
||||
}
|
||||
};
|
||||
|
||||
// Test guardrail using backend endpoint
|
||||
const handleTest = async () => {
|
||||
if (!accessToken) {
|
||||
setTestResult({ error: "No access token available" });
|
||||
return;
|
||||
}
|
||||
|
||||
setIsTesting(true);
|
||||
setTestResult(null);
|
||||
|
||||
try {
|
||||
// Parse test input JSON
|
||||
let parsedInput;
|
||||
try {
|
||||
parsedInput = JSON.parse(testInput);
|
||||
} catch (e) {
|
||||
setTestResult({ error: "Invalid test input JSON" });
|
||||
setIsTesting(false);
|
||||
return;
|
||||
}
|
||||
|
||||
// Ensure texts array exists
|
||||
if (!parsedInput.texts) {
|
||||
parsedInput.texts = [];
|
||||
}
|
||||
|
||||
const response = await testCustomCodeGuardrail(accessToken, {
|
||||
custom_code: code,
|
||||
test_input: parsedInput,
|
||||
input_type: mode as "request" | "response",
|
||||
request_data: {
|
||||
model: "test-model",
|
||||
metadata: {},
|
||||
},
|
||||
});
|
||||
|
||||
if (response.success && response.result) {
|
||||
setTestResult(response.result);
|
||||
} else if (response.error) {
|
||||
setTestResult({
|
||||
error: response.error,
|
||||
error_type: response.error_type,
|
||||
});
|
||||
} else {
|
||||
setTestResult({ error: "Unknown error occurred" });
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("Failed to test custom code:", error);
|
||||
setTestResult({
|
||||
error: error instanceof Error ? error.message : "Failed to test custom code",
|
||||
});
|
||||
} finally {
|
||||
setIsTesting(false);
|
||||
}
|
||||
};
|
||||
|
||||
const lineCount = code.split("\n").length;
|
||||
|
||||
return (
|
||||
<Modal
|
||||
open={visible}
|
||||
onCancel={onClose}
|
||||
footer={null}
|
||||
width={1200}
|
||||
className="custom-code-modal"
|
||||
closable={true}
|
||||
destroyOnClose
|
||||
>
|
||||
<div className="flex flex-col h-[80vh]">
|
||||
{/* Header */}
|
||||
<div className="pb-4 border-b border-gray-200">
|
||||
<h2 className="text-xl font-semibold text-gray-900">Create Custom Guardrail</h2>
|
||||
<p className="text-sm text-gray-500 mt-1">Define custom logic using Python-like syntax</p>
|
||||
</div>
|
||||
|
||||
{/* Top Controls */}
|
||||
<div className="flex items-center gap-4 py-4 border-b border-gray-100">
|
||||
<div className="flex-1 max-w-[200px]">
|
||||
<label className="block text-xs font-medium text-gray-600 mb-1">Guardrail Name</label>
|
||||
<TextInput
|
||||
value={guardrailName}
|
||||
onValueChange={setGuardrailName}
|
||||
placeholder="e.g., block-pii-custom"
|
||||
/>
|
||||
</div>
|
||||
<div className="w-[180px]">
|
||||
<label className="block text-xs font-medium text-gray-600 mb-1">Mode</label>
|
||||
<Select
|
||||
value={mode}
|
||||
onChange={setMode}
|
||||
options={MODE_OPTIONS}
|
||||
className="w-full"
|
||||
size="middle"
|
||||
/>
|
||||
</div>
|
||||
<div className="w-[180px]">
|
||||
<label className="block text-xs font-medium text-gray-600 mb-1">Template</label>
|
||||
<Select
|
||||
value={selectedTemplate}
|
||||
onChange={handleTemplateChange}
|
||||
className="w-full"
|
||||
size="middle"
|
||||
>
|
||||
{Object.entries(CODE_TEMPLATES).map(([key, template]) => (
|
||||
<Select.Option key={key} value={key}>
|
||||
{template.name}
|
||||
</Select.Option>
|
||||
))}
|
||||
</Select>
|
||||
</div>
|
||||
<div className="flex items-center gap-2 pt-5">
|
||||
<span className="text-sm text-gray-600">Default On</span>
|
||||
<Switch checked={defaultOn} onChange={setDefaultOn} />
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Main Content */}
|
||||
<div className="flex flex-1 overflow-hidden mt-4 gap-4">
|
||||
{/* Code Editor */}
|
||||
<div className="flex-1 flex flex-col min-w-0">
|
||||
<div className="flex items-center justify-between mb-2">
|
||||
<span className="text-xs font-semibold text-gray-500 uppercase tracking-wide">Python Logic</span>
|
||||
<span className="text-xs text-gray-400">Restricted environment (no imports)</span>
|
||||
</div>
|
||||
<div className="flex-1 relative rounded-lg overflow-hidden border border-gray-700 bg-[#1e1e1e]">
|
||||
{/* Line numbers */}
|
||||
<div
|
||||
className="absolute left-0 top-0 bottom-0 w-10 bg-[#1e1e1e] border-r border-gray-700 text-right pr-2 pt-3 select-none overflow-hidden"
|
||||
style={{ fontFamily: "monospace", fontSize: "13px", lineHeight: "1.5" }}
|
||||
>
|
||||
{Array.from({ length: Math.max(lineCount, 20) }, (_, i) => (
|
||||
<div key={i + 1} className="text-gray-500 h-[19.5px]">{i + 1}</div>
|
||||
))}
|
||||
</div>
|
||||
{/* Code textarea */}
|
||||
<textarea
|
||||
ref={textareaRef}
|
||||
value={code}
|
||||
onChange={(e) => setCode(e.target.value)}
|
||||
onKeyDown={handleKeyDown}
|
||||
spellCheck={false}
|
||||
className="w-full h-full pl-12 pr-4 pt-3 pb-3 font-mono text-sm resize-none focus:outline-none bg-transparent text-gray-200"
|
||||
style={{ lineHeight: "1.5", tabSize: 4 }}
|
||||
/>
|
||||
</div>
|
||||
|
||||
{/* Test Section */}
|
||||
<Collapse
|
||||
activeKey={testExpanded ? ["test"] : []}
|
||||
onChange={(keys) => setTestExpanded(keys.includes("test"))}
|
||||
className="mt-3 bg-white border border-gray-200 rounded-lg"
|
||||
expandIcon={({ isActive }) => <CaretRightOutlined rotate={isActive ? 90 : 0} />}
|
||||
>
|
||||
<Panel
|
||||
header={
|
||||
<span className="flex items-center gap-2 text-sm font-medium">
|
||||
<PlayCircleOutlined className="text-blue-500" />
|
||||
Test Your Guardrail
|
||||
</span>
|
||||
}
|
||||
key="test"
|
||||
>
|
||||
<div className="space-y-3">
|
||||
<div>
|
||||
<label className="block text-xs font-medium text-gray-600 mb-1">Test Input (JSON)</label>
|
||||
<TextArea
|
||||
value={testInput}
|
||||
onChange={(e) => setTestInput(e.target.value)}
|
||||
rows={4}
|
||||
className="font-mono text-xs"
|
||||
placeholder='{"texts": ["test message"], ...}'
|
||||
/>
|
||||
</div>
|
||||
<div className="flex items-center gap-3">
|
||||
<Button
|
||||
size="xs"
|
||||
onClick={handleTest}
|
||||
disabled={isTesting}
|
||||
icon={PlayCircleOutlined}
|
||||
>
|
||||
{isTesting ? "Running..." : "Run Test"}
|
||||
</Button>
|
||||
{testResult && (
|
||||
<div className={`flex items-center gap-2 text-sm ${
|
||||
testResult.error ? "text-red-600" :
|
||||
testResult.action === "allow" ? "text-green-600" :
|
||||
testResult.action === "block" ? "text-orange-600" :
|
||||
"text-blue-600"
|
||||
}`}>
|
||||
{testResult.error ? (
|
||||
<>
|
||||
<CloseCircleOutlined />
|
||||
<span>
|
||||
{testResult.error_type && <span className="font-medium">[{testResult.error_type}] </span>}
|
||||
{testResult.error}
|
||||
</span>
|
||||
</>
|
||||
) : testResult.action === "allow" ? (
|
||||
<><CheckCircleOutlined /> Allowed</>
|
||||
) : testResult.action === "block" ? (
|
||||
<><CloseCircleOutlined /> Blocked: {testResult.reason}</>
|
||||
) : testResult.action === "modify" ? (
|
||||
<>
|
||||
<CheckCircleOutlined /> Modified
|
||||
{testResult.texts && testResult.texts.length > 0 && (
|
||||
<span className="text-xs text-gray-500 ml-1">
|
||||
→ {testResult.texts[0].substring(0, 50)}{testResult.texts[0].length > 50 ? "..." : ""}
|
||||
</span>
|
||||
)}
|
||||
</>
|
||||
) : (
|
||||
<><CheckCircleOutlined /> {testResult.action || "Unknown"}</>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</Panel>
|
||||
</Collapse>
|
||||
</div>
|
||||
|
||||
{/* Primitives Panel */}
|
||||
<div className="w-[280px] flex-shrink-0 overflow-auto">
|
||||
<div className="flex items-center gap-2 mb-3">
|
||||
<CodeOutlined className="text-blue-500" />
|
||||
<span className="font-semibold text-gray-700">Available Primitives</span>
|
||||
</div>
|
||||
<p className="text-xs text-gray-500 mb-3">Click to copy functions to clipboard</p>
|
||||
|
||||
<Collapse
|
||||
defaultActiveKey={["Return Values"]}
|
||||
className="primitives-collapse bg-transparent border-0"
|
||||
expandIconPosition="end"
|
||||
>
|
||||
{Object.entries(PRIMITIVES).map(([category, primitives]) => (
|
||||
<Panel
|
||||
header={<span className="text-sm font-medium text-gray-700">{category}</span>}
|
||||
key={category}
|
||||
className="bg-white mb-2 rounded-lg border border-gray-200"
|
||||
>
|
||||
<div className="space-y-2">
|
||||
{primitives.map((p) => (
|
||||
<button
|
||||
key={p.name}
|
||||
onClick={() => copyPrimitive(p.name)}
|
||||
className={`w-full text-left px-2 py-2 rounded transition-colors ${
|
||||
copiedPrimitive === p.name
|
||||
? "bg-green-100"
|
||||
: "bg-gray-50 hover:bg-blue-50"
|
||||
}`}
|
||||
>
|
||||
{copiedPrimitive === p.name ? (
|
||||
<span className="flex items-center gap-1 text-xs font-mono text-green-700">
|
||||
<CheckCircleOutlined /> Copied!
|
||||
</span>
|
||||
) : (
|
||||
<>
|
||||
<div className="text-xs font-mono text-gray-800">{p.name}</div>
|
||||
<div className="text-[10px] text-gray-500 mt-0.5">{p.desc}</div>
|
||||
</>
|
||||
)}
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
</Panel>
|
||||
))}
|
||||
</Collapse>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Footer */}
|
||||
<div className="flex items-center justify-between pt-4 mt-4 border-t border-gray-200">
|
||||
<span className="text-xs text-gray-400">Changes are auto-saved to local draft</span>
|
||||
<div className="flex items-center gap-3">
|
||||
<Button variant="secondary" onClick={onClose}>
|
||||
Cancel
|
||||
</Button>
|
||||
<Button
|
||||
onClick={handleSave}
|
||||
loading={isSaving}
|
||||
disabled={isSaving || !guardrailName.trim()}
|
||||
icon={SaveOutlined}
|
||||
>
|
||||
Save Guardrail
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<style>{`
|
||||
.custom-code-modal .ant-modal-content {
|
||||
padding: 24px;
|
||||
}
|
||||
.custom-code-modal .ant-modal-close {
|
||||
top: 20px;
|
||||
right: 20px;
|
||||
}
|
||||
.primitives-collapse .ant-collapse-item {
|
||||
border: none !important;
|
||||
}
|
||||
.primitives-collapse .ant-collapse-header {
|
||||
padding: 8px 12px !important;
|
||||
}
|
||||
.primitives-collapse .ant-collapse-content-box {
|
||||
padding: 8px 12px !important;
|
||||
}
|
||||
`}</style>
|
||||
</Modal>
|
||||
);
|
||||
};
|
||||
|
||||
export default CustomCodeModal;
|
||||
@@ -0,0 +1,588 @@
|
||||
import React, { useState, useCallback } from "react";
|
||||
import { Card, Title, Text, Button } from "@tremor/react";
|
||||
import { Collapse, Typography, Tooltip, Spin, Alert, Tabs } from "antd";
|
||||
import {
|
||||
PlayCircleOutlined,
|
||||
InfoCircleOutlined,
|
||||
CodeOutlined,
|
||||
CheckCircleOutlined,
|
||||
CloseCircleOutlined,
|
||||
EditOutlined,
|
||||
BookOutlined,
|
||||
ExperimentOutlined,
|
||||
} from "@ant-design/icons";
|
||||
import NotificationsManager from "../../molecules/notifications_manager";
|
||||
import { testCustomCodeGuardrail } from "../../networking";
|
||||
import CustomCodeEditor from "./CustomCodeEditor";
|
||||
import { CUSTOM_CODE_PRIMITIVES, CUSTOM_CODE_EXAMPLES, DEFAULT_CUSTOM_CODE } from "./custom_code_constants";
|
||||
|
||||
const { Panel } = Collapse;
|
||||
const { Paragraph } = Typography;
|
||||
|
||||
interface CustomCodePlaygroundProps {
|
||||
accessToken: string | null;
|
||||
initialCode?: string;
|
||||
onCodeChange?: (code: string) => void;
|
||||
showTestingPanel?: boolean;
|
||||
}
|
||||
|
||||
interface TestResult {
|
||||
action: "allow" | "block" | "modify";
|
||||
reason?: string;
|
||||
modified_texts?: string[];
|
||||
modified_images?: string[];
|
||||
modified_tool_calls?: any[];
|
||||
execution_time_ms?: number;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
const CustomCodePlayground: React.FC<CustomCodePlaygroundProps> = ({
|
||||
accessToken,
|
||||
initialCode = DEFAULT_CUSTOM_CODE,
|
||||
onCodeChange,
|
||||
showTestingPanel = true,
|
||||
}) => {
|
||||
const [customCode, setCustomCode] = useState(initialCode);
|
||||
const [testInput, setTestInput] = useState(
|
||||
JSON.stringify(
|
||||
{
|
||||
texts: ["Hello, my SSN is 123-45-6789"],
|
||||
images: [],
|
||||
tools: [],
|
||||
tool_calls: [],
|
||||
structured_messages: [
|
||||
{ role: "user", content: "Hello, my SSN is 123-45-6789" },
|
||||
],
|
||||
model: "gpt-4",
|
||||
},
|
||||
null,
|
||||
2
|
||||
)
|
||||
);
|
||||
const [requestData, setRequestData] = useState(
|
||||
JSON.stringify(
|
||||
{
|
||||
model: "gpt-4",
|
||||
user_id: "test-user",
|
||||
team_id: "test-team",
|
||||
end_user_id: "end-user-123",
|
||||
metadata: {},
|
||||
},
|
||||
null,
|
||||
2
|
||||
)
|
||||
);
|
||||
const [inputType, setInputType] = useState<"request" | "response">("request");
|
||||
const [testResult, setTestResult] = useState<TestResult | null>(null);
|
||||
const [isTesting, setIsTesting] = useState(false);
|
||||
const [activeTab, setActiveTab] = useState<string>("editor");
|
||||
|
||||
const handleCodeChange = useCallback(
|
||||
(code: string) => {
|
||||
setCustomCode(code);
|
||||
onCodeChange?.(code);
|
||||
},
|
||||
[onCodeChange]
|
||||
);
|
||||
|
||||
const handleRunTest = async () => {
|
||||
if (!accessToken) {
|
||||
NotificationsManager.fromBackend("No access token available");
|
||||
return;
|
||||
}
|
||||
|
||||
setIsTesting(true);
|
||||
setTestResult(null);
|
||||
|
||||
try {
|
||||
let parsedInputs: any;
|
||||
let parsedRequestData: any;
|
||||
|
||||
try {
|
||||
parsedInputs = JSON.parse(testInput);
|
||||
} catch (e) {
|
||||
throw new Error("Invalid JSON in test input");
|
||||
}
|
||||
|
||||
try {
|
||||
parsedRequestData = JSON.parse(requestData);
|
||||
} catch (e) {
|
||||
throw new Error("Invalid JSON in request data");
|
||||
}
|
||||
|
||||
const response = await testCustomCodeGuardrail(accessToken, {
|
||||
custom_code: customCode,
|
||||
test_input: parsedInputs,
|
||||
input_type: inputType,
|
||||
request_data: parsedRequestData,
|
||||
});
|
||||
|
||||
if (response.success && response.result) {
|
||||
setTestResult(response.result);
|
||||
|
||||
if (response.result.action === "allow") {
|
||||
NotificationsManager.success("Guardrail allowed the request");
|
||||
} else if (response.result.action === "block") {
|
||||
NotificationsManager.fromBackend(`Guardrail blocked: ${response.result.reason || "No reason provided"}`);
|
||||
} else if (response.result.action === "modify") {
|
||||
NotificationsManager.success("Guardrail modified the content");
|
||||
}
|
||||
} else if (response.error) {
|
||||
setTestResult({
|
||||
action: "block",
|
||||
error: response.error,
|
||||
});
|
||||
NotificationsManager.fromBackend(`Test failed: ${response.error}`);
|
||||
}
|
||||
} catch (error) {
|
||||
const errorMessage = error instanceof Error ? error.message : "Unknown error";
|
||||
setTestResult({
|
||||
action: "block",
|
||||
error: errorMessage,
|
||||
});
|
||||
NotificationsManager.fromBackend(`Test failed: ${errorMessage}`);
|
||||
} finally {
|
||||
setIsTesting(false);
|
||||
}
|
||||
};
|
||||
|
||||
const loadExample = (exampleKey: keyof typeof CUSTOM_CODE_EXAMPLES) => {
|
||||
setCustomCode(CUSTOM_CODE_EXAMPLES[exampleKey]);
|
||||
onCodeChange?.(CUSTOM_CODE_EXAMPLES[exampleKey]);
|
||||
setActiveTab("editor");
|
||||
};
|
||||
|
||||
const renderTestResult = () => {
|
||||
if (!testResult) return null;
|
||||
|
||||
const isError = !!testResult.error;
|
||||
const isAllow = testResult.action === "allow";
|
||||
const isBlock = testResult.action === "block";
|
||||
const isModify = testResult.action === "modify";
|
||||
|
||||
return (
|
||||
<div className="mt-4">
|
||||
<Text className="font-medium text-gray-700 block mb-2">Test Result</Text>
|
||||
<div
|
||||
className={`rounded-lg p-4 border ${
|
||||
isError
|
||||
? "bg-red-50 border-red-200"
|
||||
: isAllow
|
||||
? "bg-green-50 border-green-200"
|
||||
: isBlock
|
||||
? "bg-orange-50 border-orange-200"
|
||||
: "bg-blue-50 border-blue-200"
|
||||
}`}
|
||||
>
|
||||
<div className="flex items-center gap-2 mb-2">
|
||||
{isError ? (
|
||||
<CloseCircleOutlined className="text-red-500 text-lg" />
|
||||
) : isAllow ? (
|
||||
<CheckCircleOutlined className="text-green-500 text-lg" />
|
||||
) : isBlock ? (
|
||||
<CloseCircleOutlined className="text-orange-500 text-lg" />
|
||||
) : (
|
||||
<EditOutlined className="text-blue-500 text-lg" />
|
||||
)}
|
||||
<span
|
||||
className={`font-semibold ${
|
||||
isError
|
||||
? "text-red-700"
|
||||
: isAllow
|
||||
? "text-green-700"
|
||||
: isBlock
|
||||
? "text-orange-700"
|
||||
: "text-blue-700"
|
||||
}`}
|
||||
>
|
||||
{isError ? "Error" : testResult.action.toUpperCase()}
|
||||
</span>
|
||||
{testResult.execution_time_ms && (
|
||||
<span className="text-xs text-gray-500 ml-auto">
|
||||
{testResult.execution_time_ms.toFixed(2)}ms
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{isError && (
|
||||
<Alert type="error" message={testResult.error} className="mt-2" />
|
||||
)}
|
||||
|
||||
{isBlock && testResult.reason && (
|
||||
<Paragraph className="text-orange-700 mb-0 mt-2">
|
||||
<strong>Reason:</strong> {testResult.reason}
|
||||
</Paragraph>
|
||||
)}
|
||||
|
||||
{isModify && testResult.modified_texts && testResult.modified_texts.length > 0 && (
|
||||
<div className="mt-2">
|
||||
<Text className="font-medium text-blue-700 block mb-1">Modified Texts:</Text>
|
||||
<pre className="bg-white rounded p-2 text-xs overflow-auto max-h-32 border border-blue-100">
|
||||
{JSON.stringify(testResult.modified_texts, null, 2)}
|
||||
</pre>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
const renderPrimitivesReference = () => (
|
||||
<div className="space-y-4">
|
||||
{Object.entries(CUSTOM_CODE_PRIMITIVES).map(([category, primitives]) => (
|
||||
<div key={category}>
|
||||
<Text className="font-semibold text-gray-700 block mb-2">{category}</Text>
|
||||
<div className="bg-gray-50 rounded-lg p-3">
|
||||
<table className="w-full text-sm">
|
||||
<thead>
|
||||
<tr className="border-b border-gray-200">
|
||||
<th className="text-left py-1 pr-4 font-medium text-gray-600">Function</th>
|
||||
<th className="text-left py-1 font-medium text-gray-600">Description</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
{primitives.map((primitive) => (
|
||||
<tr key={primitive.name} className="border-b border-gray-100 last:border-0">
|
||||
<td className="py-1.5 pr-4">
|
||||
<code className="text-xs bg-blue-50 text-blue-700 px-1.5 py-0.5 rounded font-mono">
|
||||
{primitive.signature}
|
||||
</code>
|
||||
</td>
|
||||
<td className="py-1.5 text-gray-600">{primitive.description}</td>
|
||||
</tr>
|
||||
))}
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
</div>
|
||||
))}
|
||||
|
||||
<div>
|
||||
<Text className="font-semibold text-gray-700 block mb-2">Return Values</Text>
|
||||
<div className="bg-gray-50 rounded-lg p-3">
|
||||
<table className="w-full text-sm">
|
||||
<thead>
|
||||
<tr className="border-b border-gray-200">
|
||||
<th className="text-left py-1 pr-4 font-medium text-gray-600">Function</th>
|
||||
<th className="text-left py-1 font-medium text-gray-600">Description</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr className="border-b border-gray-100">
|
||||
<td className="py-1.5 pr-4"><code className="text-xs bg-green-100 text-green-700 px-1.5 py-0.5 rounded font-mono">allow()</code></td>
|
||||
<td className="py-1.5 text-gray-600">Let request/response through</td>
|
||||
</tr>
|
||||
<tr className="border-b border-gray-100">
|
||||
<td className="py-1.5 pr-4"><code className="text-xs bg-red-100 text-red-700 px-1.5 py-0.5 rounded font-mono">block(reason)</code></td>
|
||||
<td className="py-1.5 text-gray-600">Reject with message</td>
|
||||
</tr>
|
||||
<tr className="border-b border-gray-100">
|
||||
<td className="py-1.5 pr-4"><code className="text-xs bg-yellow-100 text-yellow-700 px-1.5 py-0.5 rounded font-mono">modify(texts=[], images=[], tool_calls=[])</code></td>
|
||||
<td className="py-1.5 text-gray-600">Transform content</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
|
||||
const renderInputParamsReference = () => (
|
||||
<div className="space-y-4">
|
||||
<div>
|
||||
<Text className="font-semibold text-gray-700 block mb-2">`inputs` Parameter</Text>
|
||||
<div className="bg-gray-50 rounded-lg p-3">
|
||||
<table className="w-full text-sm">
|
||||
<thead>
|
||||
<tr className="border-b border-gray-200">
|
||||
<th className="text-left py-1 pr-4 font-medium text-gray-600">Field</th>
|
||||
<th className="text-left py-1 pr-4 font-medium text-gray-600">Type</th>
|
||||
<th className="text-left py-1 font-medium text-gray-600">Description</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr className="border-b border-gray-100">
|
||||
<td className="py-1.5 pr-4"><code className="text-xs bg-gray-200 px-1 rounded font-mono">texts</code></td>
|
||||
<td className="py-1.5 pr-4 text-gray-500">List[str]</td>
|
||||
<td className="py-1.5 text-gray-600">Extracted text from the request/response</td>
|
||||
</tr>
|
||||
<tr className="border-b border-gray-100">
|
||||
<td className="py-1.5 pr-4"><code className="text-xs bg-gray-200 px-1 rounded font-mono">images</code></td>
|
||||
<td className="py-1.5 pr-4 text-gray-500">List[str]</td>
|
||||
<td className="py-1.5 text-gray-600">Extracted images (for image guardrails)</td>
|
||||
</tr>
|
||||
<tr className="border-b border-gray-100">
|
||||
<td className="py-1.5 pr-4"><code className="text-xs bg-gray-200 px-1 rounded font-mono">tools</code></td>
|
||||
<td className="py-1.5 pr-4 text-gray-500">List[dict]</td>
|
||||
<td className="py-1.5 text-gray-600">Tools sent to the LLM</td>
|
||||
</tr>
|
||||
<tr className="border-b border-gray-100">
|
||||
<td className="py-1.5 pr-4"><code className="text-xs bg-gray-200 px-1 rounded font-mono">tool_calls</code></td>
|
||||
<td className="py-1.5 pr-4 text-gray-500">List[dict]</td>
|
||||
<td className="py-1.5 text-gray-600">Tool calls returned from the LLM</td>
|
||||
</tr>
|
||||
<tr className="border-b border-gray-100">
|
||||
<td className="py-1.5 pr-4"><code className="text-xs bg-gray-200 px-1 rounded font-mono">structured_messages</code></td>
|
||||
<td className="py-1.5 pr-4 text-gray-500">List[dict]</td>
|
||||
<td className="py-1.5 text-gray-600">Full messages with role info</td>
|
||||
</tr>
|
||||
<tr className="border-b border-gray-100">
|
||||
<td className="py-1.5 pr-4"><code className="text-xs bg-gray-200 px-1 rounded font-mono">model</code></td>
|
||||
<td className="py-1.5 pr-4 text-gray-500">str</td>
|
||||
<td className="py-1.5 text-gray-600">The model being used</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div>
|
||||
<Text className="font-semibold text-gray-700 block mb-2">`request_data` Parameter</Text>
|
||||
<div className="bg-gray-50 rounded-lg p-3">
|
||||
<table className="w-full text-sm">
|
||||
<thead>
|
||||
<tr className="border-b border-gray-200">
|
||||
<th className="text-left py-1 pr-4 font-medium text-gray-600">Field</th>
|
||||
<th className="text-left py-1 pr-4 font-medium text-gray-600">Type</th>
|
||||
<th className="text-left py-1 font-medium text-gray-600">Description</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr className="border-b border-gray-100">
|
||||
<td className="py-1.5 pr-4"><code className="text-xs bg-gray-200 px-1 rounded font-mono">model</code></td>
|
||||
<td className="py-1.5 pr-4 text-gray-500">str</td>
|
||||
<td className="py-1.5 text-gray-600">Model name</td>
|
||||
</tr>
|
||||
<tr className="border-b border-gray-100">
|
||||
<td className="py-1.5 pr-4"><code className="text-xs bg-gray-200 px-1 rounded font-mono">user_id</code></td>
|
||||
<td className="py-1.5 pr-4 text-gray-500">str</td>
|
||||
<td className="py-1.5 text-gray-600">User ID from API key</td>
|
||||
</tr>
|
||||
<tr className="border-b border-gray-100">
|
||||
<td className="py-1.5 pr-4"><code className="text-xs bg-gray-200 px-1 rounded font-mono">team_id</code></td>
|
||||
<td className="py-1.5 pr-4 text-gray-500">str</td>
|
||||
<td className="py-1.5 text-gray-600">Team ID from API key</td>
|
||||
</tr>
|
||||
<tr className="border-b border-gray-100">
|
||||
<td className="py-1.5 pr-4"><code className="text-xs bg-gray-200 px-1 rounded font-mono">end_user_id</code></td>
|
||||
<td className="py-1.5 pr-4 text-gray-500">str</td>
|
||||
<td className="py-1.5 text-gray-600">End user ID</td>
|
||||
</tr>
|
||||
<tr className="border-b border-gray-100">
|
||||
<td className="py-1.5 pr-4"><code className="text-xs bg-gray-200 px-1 rounded font-mono">metadata</code></td>
|
||||
<td className="py-1.5 pr-4 text-gray-500">dict</td>
|
||||
<td className="py-1.5 text-gray-600">Request metadata</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
|
||||
const renderExamplesTab = () => (
|
||||
<div className="grid grid-cols-1 md:grid-cols-2 gap-3">
|
||||
<button
|
||||
onClick={() => loadExample("blockSSN")}
|
||||
className="text-left p-4 border border-gray-200 rounded-lg hover:border-blue-400 hover:bg-blue-50 transition-colors"
|
||||
>
|
||||
<Text className="font-medium text-gray-800 block mb-1">🔒 Block PII (SSN)</Text>
|
||||
<Text className="text-xs text-gray-500">Detect and block Social Security Numbers</Text>
|
||||
</button>
|
||||
<button
|
||||
onClick={() => loadExample("redactEmail")}
|
||||
className="text-left p-4 border border-gray-200 rounded-lg hover:border-blue-400 hover:bg-blue-50 transition-colors"
|
||||
>
|
||||
<Text className="font-medium text-gray-800 block mb-1">📧 Redact Emails</Text>
|
||||
<Text className="text-xs text-gray-500">Replace email addresses with [EMAIL REDACTED]</Text>
|
||||
</button>
|
||||
<button
|
||||
onClick={() => loadExample("blockSQL")}
|
||||
className="text-left p-4 border border-gray-200 rounded-lg hover:border-blue-400 hover:bg-blue-50 transition-colors"
|
||||
>
|
||||
<Text className="font-medium text-gray-800 block mb-1">🛡️ Block SQL Injection</Text>
|
||||
<Text className="text-xs text-gray-500">Prevent SQL code in requests</Text>
|
||||
</button>
|
||||
<button
|
||||
onClick={() => loadExample("validateJSON")}
|
||||
className="text-left p-4 border border-gray-200 rounded-lg hover:border-blue-400 hover:bg-blue-50 transition-colors"
|
||||
>
|
||||
<Text className="font-medium text-gray-800 block mb-1">✅ Validate JSON Response</Text>
|
||||
<Text className="text-xs text-gray-500">Ensure responses have required fields</Text>
|
||||
</button>
|
||||
<button
|
||||
onClick={() => loadExample("checkURLs")}
|
||||
className="text-left p-4 border border-gray-200 rounded-lg hover:border-blue-400 hover:bg-blue-50 transition-colors"
|
||||
>
|
||||
<Text className="font-medium text-gray-800 block mb-1">🔗 Check URLs</Text>
|
||||
<Text className="text-xs text-gray-500">Validate all URLs in responses</Text>
|
||||
</button>
|
||||
<button
|
||||
onClick={() => loadExample("combined")}
|
||||
className="text-left p-4 border border-gray-200 rounded-lg hover:border-blue-400 hover:bg-blue-50 transition-colors"
|
||||
>
|
||||
<Text className="font-medium text-gray-800 block mb-1">🔄 Combined Checks</Text>
|
||||
<Text className="text-xs text-gray-500">Multiple checks with redaction and blocking</Text>
|
||||
</button>
|
||||
</div>
|
||||
);
|
||||
|
||||
const tabItems = [
|
||||
{
|
||||
key: "editor",
|
||||
label: (
|
||||
<span className="flex items-center gap-1.5">
|
||||
<CodeOutlined />
|
||||
Code Editor
|
||||
</span>
|
||||
),
|
||||
children: (
|
||||
<div className="space-y-4">
|
||||
<CustomCodeEditor value={customCode} onChange={handleCodeChange} height="350px" />
|
||||
|
||||
<div className="p-3 bg-yellow-50 border border-yellow-200 rounded-lg text-sm text-yellow-800">
|
||||
<strong>⚠️ Sandbox Restrictions:</strong> No imports, no file I/O, no network access, no exec() or eval().
|
||||
Only LiteLLM-provided primitives are available.
|
||||
</div>
|
||||
</div>
|
||||
),
|
||||
},
|
||||
{
|
||||
key: "examples",
|
||||
label: (
|
||||
<span className="flex items-center gap-1.5">
|
||||
<BookOutlined />
|
||||
Examples
|
||||
</span>
|
||||
),
|
||||
children: renderExamplesTab(),
|
||||
},
|
||||
{
|
||||
key: "primitives",
|
||||
label: (
|
||||
<span className="flex items-center gap-1.5">
|
||||
<InfoCircleOutlined />
|
||||
Primitives Reference
|
||||
</span>
|
||||
),
|
||||
children: renderPrimitivesReference(),
|
||||
},
|
||||
{
|
||||
key: "params",
|
||||
label: (
|
||||
<span className="flex items-center gap-1.5">
|
||||
<InfoCircleOutlined />
|
||||
Input Parameters
|
||||
</span>
|
||||
),
|
||||
children: renderInputParamsReference(),
|
||||
},
|
||||
];
|
||||
|
||||
if (showTestingPanel) {
|
||||
tabItems.push({
|
||||
key: "test",
|
||||
label: (
|
||||
<span className="flex items-center gap-1.5">
|
||||
<ExperimentOutlined />
|
||||
Test
|
||||
</span>
|
||||
),
|
||||
children: (
|
||||
<div className="space-y-4">
|
||||
<div className="grid grid-cols-1 lg:grid-cols-2 gap-4">
|
||||
<div>
|
||||
<div className="flex items-center justify-between mb-2">
|
||||
<Text className="font-medium text-gray-700">Test Input (inputs parameter)</Text>
|
||||
<Tooltip title="This represents the 'inputs' parameter passed to your apply_guardrail function">
|
||||
<InfoCircleOutlined className="text-gray-400" />
|
||||
</Tooltip>
|
||||
</div>
|
||||
<textarea
|
||||
value={testInput}
|
||||
onChange={(e) => setTestInput(e.target.value)}
|
||||
className="w-full h-48 p-3 font-mono text-sm border border-gray-200 rounded-lg focus:ring-2 focus:ring-blue-500 focus:border-blue-500"
|
||||
placeholder="Enter test input JSON..."
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div>
|
||||
<div className="flex items-center justify-between mb-2">
|
||||
<Text className="font-medium text-gray-700">Request Data (request_data parameter)</Text>
|
||||
<Tooltip title="This represents the 'request_data' parameter passed to your apply_guardrail function">
|
||||
<InfoCircleOutlined className="text-gray-400" />
|
||||
</Tooltip>
|
||||
</div>
|
||||
<textarea
|
||||
value={requestData}
|
||||
onChange={(e) => setRequestData(e.target.value)}
|
||||
className="w-full h-48 p-3 font-mono text-sm border border-gray-200 rounded-lg focus:ring-2 focus:ring-blue-500 focus:border-blue-500"
|
||||
placeholder="Enter request data JSON..."
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="flex items-center gap-4">
|
||||
<div className="flex items-center gap-2">
|
||||
<Text className="text-sm text-gray-600">Input Type:</Text>
|
||||
<select
|
||||
value={inputType}
|
||||
onChange={(e) => setInputType(e.target.value as "request" | "response")}
|
||||
className="border border-gray-200 rounded px-3 py-1.5 text-sm focus:ring-2 focus:ring-blue-500 focus:border-blue-500"
|
||||
>
|
||||
<option value="request">request</option>
|
||||
<option value="response">response</option>
|
||||
</select>
|
||||
</div>
|
||||
|
||||
<Button
|
||||
onClick={handleRunTest}
|
||||
disabled={!accessToken || isTesting}
|
||||
icon={isTesting ? undefined : PlayCircleOutlined}
|
||||
className="ml-auto"
|
||||
>
|
||||
{isTesting ? (
|
||||
<span className="flex items-center gap-2">
|
||||
<Spin size="small" /> Running Test...
|
||||
</span>
|
||||
) : (
|
||||
"Run Test"
|
||||
)}
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
{renderTestResult()}
|
||||
</div>
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
return (
|
||||
<Card className="p-0 overflow-hidden">
|
||||
<Tabs
|
||||
activeKey={activeTab}
|
||||
onChange={setActiveTab}
|
||||
items={tabItems}
|
||||
className="custom-code-playground-tabs"
|
||||
tabBarStyle={{ padding: "0 16px", marginBottom: 0 }}
|
||||
/>
|
||||
<div className="p-4">
|
||||
{tabItems.find(tab => tab.key === activeTab)?.children}
|
||||
</div>
|
||||
<style>{`
|
||||
.custom-code-playground-tabs .ant-tabs-nav {
|
||||
background: #f9fafb;
|
||||
border-bottom: 1px solid #e5e7eb;
|
||||
}
|
||||
.custom-code-playground-tabs .ant-tabs-tab {
|
||||
padding: 12px 16px;
|
||||
}
|
||||
.custom-code-playground-tabs .ant-tabs-tab-active {
|
||||
background: white;
|
||||
}
|
||||
`}</style>
|
||||
</Card>
|
||||
);
|
||||
};
|
||||
|
||||
export default CustomCodePlayground;
|
||||
@@ -0,0 +1,194 @@
|
||||
// Custom Code Guardrail Constants
|
||||
|
||||
export const DEFAULT_CUSTOM_CODE = `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()
|
||||
`;
|
||||
|
||||
export const CUSTOM_CODE_PRIMITIVES = {
|
||||
"Regex Functions": [
|
||||
{
|
||||
name: "regex_match",
|
||||
signature: "regex_match(text, pattern)",
|
||||
description: "Returns True if pattern found in text",
|
||||
},
|
||||
{
|
||||
name: "regex_replace",
|
||||
signature: "regex_replace(text, pattern, replacement)",
|
||||
description: "Replace all matches of pattern with replacement",
|
||||
},
|
||||
{
|
||||
name: "regex_find_all",
|
||||
signature: "regex_find_all(text, pattern)",
|
||||
description: "Return list of all matches",
|
||||
},
|
||||
],
|
||||
"JSON Functions": [
|
||||
{
|
||||
name: "json_parse",
|
||||
signature: "json_parse(text)",
|
||||
description: "Parse JSON string, returns None on error",
|
||||
},
|
||||
{
|
||||
name: "json_stringify",
|
||||
signature: "json_stringify(obj)",
|
||||
description: "Convert object to JSON string",
|
||||
},
|
||||
{
|
||||
name: "json_schema_valid",
|
||||
signature: "json_schema_valid(obj, schema)",
|
||||
description: "Validate object against JSON schema",
|
||||
},
|
||||
],
|
||||
"URL Functions": [
|
||||
{
|
||||
name: "extract_urls",
|
||||
signature: "extract_urls(text)",
|
||||
description: "Extract all URLs from text",
|
||||
},
|
||||
{
|
||||
name: "is_valid_url",
|
||||
signature: "is_valid_url(url)",
|
||||
description: "Check if URL is valid",
|
||||
},
|
||||
{
|
||||
name: "all_urls_valid",
|
||||
signature: "all_urls_valid(text)",
|
||||
description: "Check all URLs in text are valid",
|
||||
},
|
||||
],
|
||||
"Code Detection": [
|
||||
{
|
||||
name: "detect_code",
|
||||
signature: "detect_code(text)",
|
||||
description: "Returns True if code detected",
|
||||
},
|
||||
{
|
||||
name: "detect_code_languages",
|
||||
signature: "detect_code_languages(text)",
|
||||
description: "Returns list of detected languages",
|
||||
},
|
||||
{
|
||||
name: "contains_code_language",
|
||||
signature: 'contains_code_language(text, ["sql", "python"])',
|
||||
description: "Check for specific languages",
|
||||
},
|
||||
],
|
||||
"Text Utilities": [
|
||||
{
|
||||
name: "contains",
|
||||
signature: "contains(text, substring)",
|
||||
description: "Check if substring exists in text",
|
||||
},
|
||||
{
|
||||
name: "contains_any",
|
||||
signature: "contains_any(text, [substr1, substr2])",
|
||||
description: "Check if any substring exists",
|
||||
},
|
||||
{
|
||||
name: "word_count",
|
||||
signature: "word_count(text)",
|
||||
description: "Count words in text",
|
||||
},
|
||||
{
|
||||
name: "char_count",
|
||||
signature: "char_count(text)",
|
||||
description: "Count characters in text",
|
||||
},
|
||||
{
|
||||
name: "lower",
|
||||
signature: "lower(text)",
|
||||
description: "Convert text to lowercase",
|
||||
},
|
||||
{
|
||||
name: "upper",
|
||||
signature: "upper(text)",
|
||||
description: "Convert text to uppercase",
|
||||
},
|
||||
{
|
||||
name: "trim",
|
||||
signature: "trim(text)",
|
||||
description: "Remove leading/trailing whitespace",
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
export const CUSTOM_CODE_EXAMPLES = {
|
||||
blockSSN: `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()
|
||||
`,
|
||||
|
||||
redactEmail: `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)
|
||||
`,
|
||||
|
||||
blockSQL: `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()
|
||||
`,
|
||||
|
||||
validateJSON: `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()
|
||||
`,
|
||||
|
||||
checkURLs: `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()
|
||||
`,
|
||||
|
||||
combined: `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)
|
||||
`,
|
||||
};
|
||||
@@ -0,0 +1 @@
|
||||
export { default as CustomCodeModal } from "./CustomCodeModal";
|
||||
@@ -7607,6 +7607,89 @@ export const applyGuardrail = async (
|
||||
}
|
||||
};
|
||||
|
||||
export interface TestCustomCodeGuardrailRequest {
|
||||
custom_code: string;
|
||||
test_input: {
|
||||
texts: string[];
|
||||
images?: string[];
|
||||
tools?: Record<string, any>[];
|
||||
tool_calls?: Record<string, any>[];
|
||||
structured_messages?: Record<string, any>[];
|
||||
model?: string;
|
||||
};
|
||||
input_type?: "request" | "response";
|
||||
request_data?: {
|
||||
model?: string;
|
||||
user_id?: string;
|
||||
team_id?: string;
|
||||
end_user_id?: string;
|
||||
metadata?: Record<string, any>;
|
||||
};
|
||||
}
|
||||
|
||||
export interface TestCustomCodeGuardrailResponse {
|
||||
success: boolean;
|
||||
result?: {
|
||||
action: "allow" | "block" | "modify";
|
||||
reason?: string;
|
||||
texts?: string[];
|
||||
images?: string[];
|
||||
tool_calls?: Record<string, any>[];
|
||||
detection_info?: Record<string, any>;
|
||||
warning?: string;
|
||||
};
|
||||
error?: string;
|
||||
error_type?: "compilation" | "execution";
|
||||
}
|
||||
|
||||
export const testCustomCodeGuardrail = async (
|
||||
accessToken: string,
|
||||
request: TestCustomCodeGuardrailRequest
|
||||
): Promise<TestCustomCodeGuardrailResponse> => {
|
||||
try {
|
||||
const url = proxyBaseUrl
|
||||
? `${proxyBaseUrl}/guardrails/test_custom_code`
|
||||
: `/guardrails/test_custom_code`;
|
||||
|
||||
const response = await fetch(url, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify(request),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const errorData = await response.text();
|
||||
let errorMessage = "Failed to test custom code guardrail";
|
||||
|
||||
try {
|
||||
const errorJson = JSON.parse(errorData);
|
||||
if (errorJson.error?.message) {
|
||||
errorMessage = errorJson.error.message;
|
||||
} else if (errorJson.detail) {
|
||||
errorMessage = errorJson.detail;
|
||||
} else if (errorJson.message) {
|
||||
errorMessage = errorJson.message;
|
||||
}
|
||||
} catch (e) {
|
||||
errorMessage = errorData || errorMessage;
|
||||
}
|
||||
|
||||
handleError(errorData);
|
||||
throw new Error(errorMessage);
|
||||
}
|
||||
|
||||
const data = await response.json();
|
||||
console.log("Test custom code guardrail response:", data);
|
||||
return data;
|
||||
} catch (error) {
|
||||
console.error("Failed to test custom code guardrail:", error);
|
||||
throw error;
|
||||
}
|
||||
};
|
||||
|
||||
export const validateBlockedWordsFile = async (accessToken: string, fileContent: string) => {
|
||||
try {
|
||||
const url = proxyBaseUrl
|
||||
|
||||
Reference in New Issue
Block a user