Skip to main content

Custom Guardrail

Use this is you want to write code to run a custom guardrail

Quick Start

1. Write a CustomGuardrail Class

A CustomGuardrail has 3 methods to enforce guardrails

  • async_pre_call_hook - (Optional) modify input or reject request before making LLM API call
  • async_moderation_hook - (Optional) reject request, runs while making LLM API call (help to lower latency)
  • async_post_call_success_hook- (Optional) apply guardrail on input/output, runs after making LLM API call

See detailed spec of methods here

Example CustomGuardrail Class

Create a new file called custom_guardrail.py and add this code to it

from typing import Any, Dict, List, Literal, Optional, Union

import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata
from litellm.types.guardrails import GuardrailEventHooks


class myCustomGuardrail(CustomGuardrail):
def __init__(
self,
**kwargs,
):
# store kwargs as optional_params
self.optional_params = kwargs

super().__init__(**kwargs)

async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank"
],
) -> Optional[Union[Exception, str, dict]]:
"""
Runs before the LLM API call
Runs on only Input
Use this if you want to MODIFY the input
"""

# In this guardrail, if a user inputs `litellm` we will mask it and then send it to the LLM
_messages = data.get("messages")
if _messages:
for message in _messages:
_content = message.get("content")
if isinstance(_content, str):
if "litellm" in _content.lower():
_content = _content.replace("litellm", "********")
message["content"] = _content

verbose_proxy_logger.debug(
"async_pre_call_hook: Message after masking %s", _messages
)

return data

async def async_moderation_hook(
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: Literal["completion", "embeddings", "image_generation", "moderation", "audio_transcription"],
):
"""
Runs in parallel to LLM API call
Runs on only Input

This can NOT modify the input, only used to reject or accept a call before going to LLM API
"""

# this works the same as async_pre_call_hook, but just runs in parallel as the LLM API Call
# In this guardrail, if a user inputs `litellm` we will mask it.
_messages = data.get("messages")
if _messages:
for message in _messages:
_content = message.get("content")
if isinstance(_content, str):
if "litellm" in _content.lower():
raise ValueError("Guardrail failed words - `litellm` detected")

async def async_post_call_success_hook(
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
response,
):
"""
Runs on response from LLM API call

It can be used to reject a response

If a response contains the word "coffee" -> we will raise an exception
"""
verbose_proxy_logger.debug("async_pre_call_hook response: %s", response)
if isinstance(response, litellm.ModelResponse):
for choice in response.choices:
if isinstance(choice, litellm.Choices):
verbose_proxy_logger.debug("async_pre_call_hook choice: %s", choice)
if (
choice.message.content
and isinstance(choice.message.content, str)
and "coffee" in choice.message.content
):
raise ValueError("Guardrail failed Coffee Detected")


2. Pass your custom guardrail class in LiteLLM config.yaml

In the config below, we point the guardrail to our custom guardrail by setting guardrail: custom_guardrail.myCustomGuardrail

  • Python Filename: custom_guardrail.py
  • Guardrail class name : myCustomGuardrail. This is defined in Step 1

guardrail: custom_guardrail.myCustomGuardrail

model_list:
- model_name: gpt-4
litellm_params:
model: openai/gpt-4o
api_key: os.environ/OPENAI_API_KEY

guardrails:
- guardrail_name: "custom-pre-guard"
litellm_params:
guardrail: custom_guardrail.myCustomGuardrail # 👈 Key change
mode: "pre_call" # runs async_pre_call_hook
- guardrail_name: "custom-during-guard"
litellm_params:
guardrail: custom_guardrail.myCustomGuardrail
mode: "during_call" # runs async_moderation_hook
- guardrail_name: "custom-post-guard"
litellm_params:
guardrail: custom_guardrail.myCustomGuardrail
mode: "post_call" # runs async_post_call_success_hook

3. Start LiteLLM Gateway

Mount your custom_guardrail.py on the LiteLLM Docker container

This mounts your custom_guardrail.py file from your local directory to the /app directory in the Docker container, making it accessible to the LiteLLM Gateway.

docker run -d \
-p 4000:4000 \
-e OPENAI_API_KEY=$OPENAI_API_KEY \
--name my-app \
-v $(pwd)/my_config.yaml:/app/config.yaml \
-v $(pwd)/custom_guardrail.py:/app/custom_guardrail.py \
my-app:latest \
--config /app/config.yaml \
--port 4000 \
--detailed_debug \

4. Test it

Test "custom-pre-guard"

Langchain, OpenAI SDK Usage Examples

Expect this to mask the word litellm before sending the request to the LLM API. This runs the async_pre_call_hook

curl -i  -X POST http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-1234" \
-d '{
"model": "gpt-4",
"messages": [
{
"role": "user",
"content": "say the word - `litellm`"
}
],
"guardrails": ["custom-pre-guard"]
}'

Expected response after pre-guard

{
"id": "chatcmpl-9zREDkBIG20RJB4pMlyutmi1hXQWc",
"choices": [
{
"finish_reason": "stop",
"index": 0,
"message": {
"content": "It looks like you've chosen a string of asterisks. This could be a way to censor or hide certain text. However, without more context, I can't provide a specific word or phrase. If there's something specific you'd like me to say or if you need help with a topic, feel free to let me know!",
"role": "assistant",
"tool_calls": null,
"function_call": null
}
}
],
"created": 1724429701,
"model": "gpt-4o-2024-05-13",
"object": "chat.completion",
"system_fingerprint": "fp_3aa7262c27",
"usage": {
"completion_tokens": 65,
"prompt_tokens": 14,
"total_tokens": 79
},
"service_tier": null
}

Test "custom-during-guard"

Langchain, OpenAI SDK Usage Examples

Expect this to fail since since litellm is in the message content. This runs the async_moderation_hook

curl -i  -X POST http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-1234" \
-d '{
"model": "gpt-4",
"messages": [
{
"role": "user",
"content": "say the word - `litellm`"
}
],
"guardrails": ["custom-during-guard"]
}'

Expected response after running during-guard

{
"error": {
"message": "Guardrail failed words - `litellm` detected",
"type": "None",
"param": "None",
"code": "500"
}
}

Test "custom-post-guard"

Langchain, OpenAI SDK Usage Examples

Expect this to fail since since coffee will be in the response content. This runs the async_post_call_success_hook

curl -i  -X POST http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-1234" \
-d '{
"model": "gpt-4",
"messages": [
{
"role": "user",
"content": "what is coffee"
}
],
"guardrails": ["custom-post-guard"]
}'

Expected response after running during-guard

{
"error": {
"message": "Guardrail failed Coffee Detected",
"type": "None",
"param": "None",
"code": "500"
}
}

CustomGuardrail methods

ComponentDescriptionOptionalChecked DataCan Modify InputCan Modify OutputCan Fail Call
async_pre_call_hookA hook that runs before the LLM API callINPUT
async_moderation_hookA hook that runs during the LLM API callINPUT
async_post_call_success_hookA hook that runs after a successful LLM API callINPUT, OUTPUT