Skip to main content

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

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

litellm --config config.yaml

3. Test

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

ParameterTypeRequiredDescription
guardrailstringMust be custom_code
modestringWhen to run: pre_call, post_call, during_call
custom_codestringPython-like code with apply_guardrail function
default_onboolRun on all requests (default: false)

Writing Custom Code

Function Signature

Your code must define an apply_guardrail function:

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

FieldTypeDescription
textsList[str]Extracted text from the request/response
imagesList[str]Extracted images (for image guardrails)
toolsList[dict]Tools sent to the LLM
tool_callsList[dict]Tool calls returned from the LLM
structured_messagesList[dict]Full messages with role info (system/user/assistant)
modelstrThe model being used

request_data Parameter

FieldTypeDescription
modelstrModel name
user_idstrUser ID from API key
team_idstrTeam ID from API key
end_user_idstrEnd user ID
metadatadictRequest metadata

Return Values

FunctionDescription
allow()Let request/response through
block(reason)Reject with message
modify(texts=[], images=[], tool_calls=[])Transform content

Built-in Primitives

Regex

FunctionDescription
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

FunctionDescription
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

FunctionDescription
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

FunctionDescription
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

FunctionDescription
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)

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

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

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

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

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

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:

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:

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):
...