scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Guardrail API - support tool call checks on OpenAI /chat/completions, OpenAI /responses, Anthropic /v1/messages (#17459)

Browse source 
* fix(unified_guardrail.py): correctly map a v1/messages call to the anthropic unified guardrail

* fix: add more rigorous call type checks

* fix(anthropic_endpoints/endpoints.py): initialize logging object at the beginning of endpoint

ensures call id + trace id are emitted to guardrail api

* feat(anthropic/chat/guardrail_translation): support streaming guardrails

sample on every 5 chunks

* fix(openai/chat/guardrail_translation): support openai streaming guardrails

* fix: initial commit fixing output guardrails for responses api

* feat(openai/responses/guardrail_translation): handler.py - fix output checks on responses api

* fix(openai/responses/guardrail_translation/handler.py): ensure responses api guardrails work on streaming

* test: update tests

* test: update tests

* fix: support multiple kinds of input to the guardrail api

* feat(guardrail_translation/handler.py): support extracting tool calls from openai chat completions for guardrail api's

* feat(generic_guardrail_api.py): support extracting + returning modified tool calls on generic_guardrails_api

allows guardrail api to analyze tool call being sent to provider - to run any analysis on it

* fix(guardrails.py): support anthropic /v1/messages tool calls

* feat(responses_api/): extract tool calls for guardrail processing

* docs(generic_guardrail_api.md): document tools param support

* docs: generic_guardrail_api.md

improve documentation
This commit is contained in:
Krish Dholakia 2025-12-03 21:20:39 -08:00 • committed by GitHub
parent be0530a6b3
commit 32013f63a0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
29 changed files with 1234 additions and 275 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
cookbook/mock_guardrail_server/mock_bedrock_guardrail_server.py
View file

@ -426,6 +426,7 @@ This is a beta API. Please help us improve it.
class LitellmBasicGuardrailRequest(BaseModel):
texts: List[str]
images: Optional[List[str]] = None
tools: Optional[List[dict]] = None
request_data: Dict[str, Any] = Field(default_factory=dict)
additional_provider_specific_params: Dict[str, Any] = Field(default_factory=dict)
input_type: Literal["request", "response"]

87
docs/my-website/docs/adding_provider/generic_guardrail_api.md
View file

@ -21,6 +21,20 @@ The **Generic Guardrail API** lets you integrate with LiteLLM **instantly** by i
5. **Custom Parameters** - Pass provider-specific params via config
6. **Full Control** - You own and maintain your guardrail API
## Supported Endpoints
The Generic Guardrail API works with the following LiteLLM endpoints:
- `/v1/chat/completions` - OpenAI Chat Completions
- `/v1/completions` - OpenAI Text Completions
- `/v1/responses` - OpenAI Responses API
- `/v1/images/generations` - OpenAI Image Generation
- `/v1/audio/transcriptions` - OpenAI Audio Transcriptions
- `/v1/audio/speech` - OpenAI Text-to-Speech
- `/v1/messages` - Anthropic Messages
- `/v1/rerank` - Cohere Rerank
- Pass-through endpoints
## How It Works
1. LiteLLM extracts text and images from any request (chat messages, embeddings, image prompts, etc.)
@ -40,6 +54,21 @@ Implement `POST /beta/litellm_basic_guardrail_api`
{
"texts": ["extracted text from the request"], // array of text strings
"images": ["base64_encoded_image_data"], // optional array of images
"tools": [ // optional array of tools (OpenAI ChatCompletionToolParam format)
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"}
}
}
}
}
],
"request_data": {
"user_api_key_hash": "hash of the litellm virtual key used",
"user_api_key_alias": "alias of the litellm virtual key used",
@ -75,6 +104,49 @@ Implement `POST /beta/litellm_basic_guardrail_api`
- `NONE` - Request proceeds unchanged
- `GUARDRAIL_INTERVENED` - Request proceeds with modified texts/images (provide `texts` and/or `images` fields)
## Parameters
### `tools` Parameter
The `tools` parameter provides information about available function/tool definitions in the request.
**Format:** OpenAI `ChatCompletionToolParam` format (see [OpenAI API reference](https://platform.openai.com/docs/api-reference/chat/create#chat-create-tools))
**Example:**
```json
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather in a location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "City and state, e.g. San Francisco, CA"
},
"unit": {
"type": "string",
"enum": ["celsius", "fahrenheit"]
}
},
"required": ["location"]
}
}
}
```
**Limitations:**
- **Input only:** Tools are only passed for `input_type="request"` (pre-call guardrails). Output/response guardrails do not currently receive tool information.
- **Supported endpoints:** The `tools` parameter is supported on: `/v1/chat/completions`, `/v1/responses`, and `/v1/messages`. Other endpoints do not have tool support.
**Use cases:**
- Enforce tool permission policies (e.g., only allow certain users/teams to access specific tools)
- Validate tool schemas before sending to LLM
- Log tool usage for audit purposes
- Block sensitive tools based on user context
## LiteLLM Configuration
Add to `config.yaml`:
@ -138,6 +210,7 @@ app = FastAPI()
class GuardrailRequest(BaseModel):
texts: List[str]
images: Optional[List[str]] = None
tools: Optional[List[Dict[str, Any]]] = None # OpenAI ChatCompletionToolParam format
request_data: Dict[str, Any]
input_type: str # "request" or "response"
litellm_call_id: Optional[str] = None
@ -153,6 +226,8 @@ class GuardrailResponse(BaseModel):
@app.post("/beta/litellm_basic_guardrail_api")
async def apply_guardrail(request: GuardrailRequest):
# Your guardrail logic here
# Example: Check text content
for text in request.texts:
if "badword" in text.lower():
return GuardrailResponse(
@ -160,6 +235,18 @@ async def apply_guardrail(request: GuardrailRequest):
blocked_reason="Content contains prohibited terms"
)
# Example: Check tools (if present in request)
if request.tools:
for tool in request.tools:
if tool.get("type") == "function":
function_name = tool.get("function", {}).get("name", "")
# Block sensitive tools
if function_name in ["delete_data", "access_admin_panel"]:
return GuardrailResponse(
action="BLOCKED",
blocked_reason=f"Tool '{function_name}' is not allowed"
)
return GuardrailResponse(action="NONE")
```

17
litellm/integrations/custom_guardrail.py
View file

@ -17,6 +17,7 @@ from litellm.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.guardrails import (
DynamicGuardrailParams,
GenericGuardrailAPIInputs,
GuardrailEventHooks,
LitellmParams,
Mode,
@ -449,20 +450,22 @@ class CustomGuardrail(CustomLogger):
async def apply_guardrail(
self,
texts: List[str],
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
images: Optional[List[str]] = None,
) -> Tuple[List[str], Optional[List[str]]]:
) -> GenericGuardrailAPIInputs:
"""
Apply your guardrail logic to the given text
Apply your guardrail logic to the given inputs
Args:
texts: The texts to apply the guardrail to
images: The images to apply the guardrail to
inputs: Dictionary containing:
- texts: List of texts to apply the guardrail to
- images: Optional list of images to apply the guardrail to
- tool_calls: Optional list of tool calls to apply the guardrail to
request_data: The request data dictionary - containing user api key metadata (e.g. user_id, team_id, etc.)
input_type: The type of input to apply the guardrail to - "request" or "response"
logging_obj: Optional logging object for tracking the guardrail execution
Any of the custom guardrails can override this method to provide custom guardrail logic
@ -473,7 +476,7 @@ class CustomGuardrail(CustomLogger):
- If the guardrail raises an exception
"""
return texts, images
return inputs
def _process_response(
self,

93
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -16,7 +16,13 @@ import json
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast
from litellm._logging import verbose_proxy_logger
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import (
LiteLLMAnthropicMessagesAdapter,
)
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.guardrails import GenericGuardrailAPIInputs
from litellm.types.llms.anthropic import AllAnthropicToolsValues
from litellm.types.llms.openai import ChatCompletionToolParam
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
@ -37,6 +43,10 @@ class AnthropicMessagesHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
def __init__(self):
super().__init__()
self.adapter = LiteLLMAnthropicMessagesAdapter()
async def process_input_messages(
self,
data: dict,
@ -47,11 +57,13 @@ class AnthropicMessagesHandler(BaseTranslation):
Process input messages by applying guardrails to text content.
"""
messages = data.get("messages")
tools = data.get("tools", None)
if messages is None:
return data
texts_to_check: List[str] = []
images_to_check: List[str] = []
tools_to_check: List[ChatCompletionToolParam] = []
task_mappings: List[Tuple[int, Optional[int]]] = []
# Track (message_index, content_index) for each text
# content_index is None for string content, int for list content
@ -66,18 +78,28 @@ class AnthropicMessagesHandler(BaseTranslation):
task_mappings=task_mappings,
)
if tools is not None:
self._extract_input_tools(
tools=tools,
tools_to_check=tools_to_check,
)
# Step 2: Apply guardrail to all texts in batch
if texts_to_check:
guardrailed_texts, guardrailed_images = (
await guardrail_to_apply.apply_guardrail(
texts=texts_to_check,
request_data=data,
input_type="request",
images=images_to_check if images_to_check else None,
logging_obj=litellm_logging_obj,
)
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
inputs["images"] = images_to_check
if tools_to_check:
inputs["tools"] = tools_to_check
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
# Step 3: Map guardrail responses back to original message structure
await self._apply_guardrail_responses_to_input(
messages=messages,
@ -105,15 +127,17 @@ class AnthropicMessagesHandler(BaseTranslation):
Override this method to customize text/image extraction logic.
"""
content = message.get("content", None)
if content is None:
tools = message.get("tools", None)
if content is None and tools is None:
return
if isinstance(content, str):
## CHECK FOR TEXT + IMAGES
if content is not None and isinstance(content, str):
# Simple string content
texts_to_check.append(content)
task_mappings.append((msg_idx, None))
elif isinstance(content, list):
elif content is not None and isinstance(content, list):
# List content (e.g., multimodal with text and images)
for content_idx, content_item in enumerate(content):
# Extract text
@ -131,6 +155,22 @@ class AnthropicMessagesHandler(BaseTranslation):
if data:
images_to_check.append(data)
def _extract_input_tools(
self,
tools: List[Dict[str, Any]],
tools_to_check: List[ChatCompletionToolParam],
) -> None:
"""
Extract tools from a message.
"""
## CHECK FOR TOOLS
if tools is not None and isinstance(tools, list):
# TRANSFORM ANTHROPIC TOOLS TO OPENAI TOOLS
openai_tools = self.adapter.translate_anthropic_tools_to_openai(
tools=cast(List[AllAnthropicToolsValues], tools)
)
tools_to_check.extend(openai_tools)
async def _apply_guardrail_responses_to_input(
self,
messages: List[Dict[str, Any]],
@ -224,16 +264,18 @@ class AnthropicMessagesHandler(BaseTranslation):
if user_metadata:
request_data["litellm_metadata"] = user_metadata
guardrailed_texts, guardrailed_images = (
await guardrail_to_apply.apply_guardrail(
texts=texts_to_check,
request_data=request_data,
input_type="response",
images=images_to_check if images_to_check else None,
logging_obj=litellm_logging_obj,
)
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
inputs["images"] = images_to_check
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
# Step 3: Map guardrail responses back to original response structure
await self._apply_guardrail_responses_to_output(
response=response,
@ -260,14 +302,11 @@ class AnthropicMessagesHandler(BaseTranslation):
Get the string so far, check the apply guardrail to the string so far, and return the list of responses so far.
"""
string_so_far = self.get_streaming_string_so_far(responses_so_far)
_, _ = (
await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid
texts=[string_so_far],
request_data={},
input_type="response",
logging_obj=litellm_logging_obj,
images=None,
)
guardrailed_inputs = await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid
inputs={"texts": [string_so_far]},
request_data={},
input_type="response",
logging_obj=litellm_logging_obj,
)
return responses_so_far

5
litellm/llms/cohere/rerank/guardrail_translation/handler.py
View file

@ -49,12 +49,13 @@ class CohereRerankHandler(BaseTranslation):
# Process query only
query = data.get("query")
if query is not None and isinstance(query, str):
guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail(
texts=[query],
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [query]},
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
data["query"] = guardrailed_texts[0] if guardrailed_texts else query
verbose_proxy_logger.debug(

357
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -19,6 +19,8 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.guardrails import GenericGuardrailAPIInputs
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.utils import Choices, StreamingChoices
if TYPE_CHECKING:
@ -52,38 +54,60 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
texts_to_check: List[str] = []
images_to_check: List[str] = []
task_mappings: List[Tuple[int, Optional[int]]] = []
# Track (message_index, content_index) for each text
tool_calls_to_check: List[ChatCompletionToolParam] = []
text_task_mappings: List[Tuple[int, Optional[int]]] = []
tool_call_task_mappings: List[Tuple[int, int]] = []
# text_task_mappings: Track (message_index, content_index) for each text
# content_index is None for string content, int for list content
# tool_call_task_mappings: Track (message_index, tool_call_index) for each tool call
# Step 1: Extract all text content and images
# Step 1: Extract all text content, images, and tool calls
for msg_idx, message in enumerate(messages):
self._extract_input_text_and_images(
self._extract_inputs(
message=message,
msg_idx=msg_idx,
texts_to_check=texts_to_check,
images_to_check=images_to_check,
task_mappings=task_mappings,
tool_calls_to_check=tool_calls_to_check,
text_task_mappings=text_task_mappings,
tool_call_task_mappings=tool_call_task_mappings,
)
# Step 2: Apply guardrail to all texts in batch
if texts_to_check:
guardrailed_texts, guardrailed_images = (
await guardrail_to_apply.apply_guardrail(
texts=texts_to_check,
request_data=data,
input_type="request",
images=images_to_check if images_to_check else None,
logging_obj=litellm_logging_obj,
)
# Step 2: Apply guardrail to all texts and tool calls in batch
if texts_to_check or tool_calls_to_check:
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
inputs["images"] = images_to_check
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check # type: ignore
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
guardrailed_tool_calls = guardrailed_inputs.get("tools", [])
# Step 3: Map guardrail responses back to original message structure
await self._apply_guardrail_responses_to_input(
messages=messages,
responses=guardrailed_texts,
task_mappings=task_mappings,
)
if guardrailed_texts and texts_to_check:
await self._apply_guardrail_responses_to_input_texts(
messages=messages,
responses=guardrailed_texts,
task_mappings=text_task_mappings,
)
# Step 4: Apply guardrailed tool calls back to messages
if guardrailed_tool_calls:
# Note: The guardrail may modify tool_calls_to_check in place
# or we may need to handle returned tool calls differently
await self._apply_guardrail_responses_to_input_tool_calls(
messages=messages,
tool_calls=guardrailed_tool_calls, # type: ignore
task_mappings=tool_call_task_mappings,
)
verbose_proxy_logger.debug(
"OpenAI Chat Completions: Processed input messages: %s", messages
@ -91,61 +115,71 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
return data
def _extract_input_text_and_images(
def _extract_inputs(
self,
message: Dict[str, Any],
msg_idx: int,
texts_to_check: List[str],
images_to_check: List[str],
task_mappings: List[Tuple[int, Optional[int]]],
tool_calls_to_check: List[ChatCompletionToolParam],
text_task_mappings: List[Tuple[int, Optional[int]]],
tool_call_task_mappings: List[Tuple[int, int]],
) -> None:
"""
Extract text content and images from a message.
Extract text content, images, and tool calls from a message.
Override this method to customize text/image extraction logic.
Override this method to customize text/image/tool call extraction logic.
"""
content = message.get("content", None)
if content is None:
return
if content is not None:
if isinstance(content, str):
# Simple string content
texts_to_check.append(content)
text_task_mappings.append((msg_idx, None))
if isinstance(content, str):
# Simple string content
texts_to_check.append(content)
task_mappings.append((msg_idx, None))
elif isinstance(content, list):
# List content (e.g., multimodal with text and images)
for content_idx, content_item in enumerate(content):
# Extract text
text_str = content_item.get("text", None)
if text_str is not None:
texts_to_check.append(text_str)
text_task_mappings.append((msg_idx, int(content_idx)))
elif isinstance(content, list):
# List content (e.g., multimodal with text and images)
for content_idx, content_item in enumerate(content):
# Extract text
text_str = content_item.get("text", None)
if text_str is not None:
texts_to_check.append(text_str)
task_mappings.append((msg_idx, int(content_idx)))
# Extract images (image_url)
if content_item.get("type") == "image_url":
image_url = content_item.get("image_url", {})
if isinstance(image_url, dict):
url = image_url.get("url")
if url:
images_to_check.append(url)
# Extract images (image_url)
if content_item.get("type") == "image_url":
image_url = content_item.get("image_url", {})
if isinstance(image_url, dict):
url = image_url.get("url")
if url:
images_to_check.append(url)
# Extract tool calls (typically in assistant messages)
tool_calls = message.get("tools", None)
if tool_calls is not None and isinstance(tool_calls, list):
for tool_call_idx, tool_call in enumerate(tool_calls):
if isinstance(tool_call, dict):
# Add the full tool call object to the list
tool_calls_to_check.append(ChatCompletionToolParam(**tool_call))
tool_call_task_mappings.append((msg_idx, int(tool_call_idx)))
async def _apply_guardrail_responses_to_input(
async def _apply_guardrail_responses_to_input_texts(
self,
messages: List[Dict[str, Any]],
responses: List[str],
task_mappings: List[Tuple[int, Optional[int]]],
) -> None:
"""
Apply guardrail responses back to input messages.
Apply guardrail responses back to input message text content.
Override this method to customize how responses are applied.
Override this method to customize how text responses are applied.
"""
for task_idx, guardrail_response in enumerate(responses):
mapping = task_mappings[task_idx]
msg_idx = cast(int, mapping[0])
content_idx_optional = cast(Optional[int], mapping[1])
# Handle content
content = messages[msg_idx].get("content", None)
if content is None:
continue
@ -160,6 +194,31 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
"text"
] = guardrail_response
async def _apply_guardrail_responses_to_input_tool_calls(
self,
messages: List[Dict[str, Any]],
tool_calls: List[Dict[str, Any]],
task_mappings: List[Tuple[int, int]],
) -> None:
"""
Apply guardrailed tool calls back to input messages.
The guardrail may have modified the tool_calls list in place,
so we apply the modified tool calls back to the original messages.
Override this method to customize how tool call responses are applied.
"""
for task_idx, (msg_idx, tool_call_idx) in enumerate(task_mappings):
if task_idx < len(tool_calls):
guardrailed_tool_call = tool_calls[task_idx]
message_tool_calls = messages[msg_idx].get("tool_calls", None)
if message_tool_calls is not None and isinstance(
message_tool_calls, list
):
if tool_call_idx < len(message_tool_calls):
# Replace the tool call with the guardrailed version
message_tool_calls[tool_call_idx] = guardrailed_tool_call
async def process_output_response(
self,
response: "ModelResponse",
@ -193,21 +252,27 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
texts_to_check: List[str] = []
images_to_check: List[str] = []
task_mappings: List[Tuple[int, Optional[int]]] = []
# Track (choice_index, content_index) for each text
tool_calls_to_check: List[Dict[str, Any]] = []
text_task_mappings: List[Tuple[int, Optional[int]]] = []
tool_call_task_mappings: List[Tuple[int, int]] = []
# text_task_mappings: Track (choice_index, content_index) for each text
# content_index is None for string content, int for list content
# tool_call_task_mappings: Track (choice_index, tool_call_index) for each tool call
# Step 1: Extract all text content and images from response choices
# Step 1: Extract all text content, images, and tool calls from response choices
for choice_idx, choice in enumerate(response.choices):
self._extract_output_text_and_images(
choice=choice,
choice_idx=choice_idx,
texts_to_check=texts_to_check,
images_to_check=images_to_check,
task_mappings=task_mappings,
tool_calls_to_check=tool_calls_to_check,
text_task_mappings=text_task_mappings,
tool_call_task_mappings=tool_call_task_mappings,
)
# Step 2: Apply guardrail to all texts in batch
if texts_to_check:
# Step 2: Apply guardrail to all texts and tool calls in batch
if texts_to_check or tool_calls_to_check:
# Create a request_data dict with response info and user API key metadata
request_data: dict = {"response": response}
@ -218,22 +283,36 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if user_metadata:
request_data["litellm_metadata"] = user_metadata
guardrailed_texts, guardrailed_images = (
await guardrail_to_apply.apply_guardrail(
texts=texts_to_check,
request_data=request_data,
input_type="response",
images=images_to_check if images_to_check else None,
logging_obj=litellm_logging_obj,
)
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
inputs["images"] = images_to_check
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check # type: ignore
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
# Step 3: Map guardrail responses back to original response structure
await self._apply_guardrail_responses_to_output(
response=response,
responses=guardrailed_texts,
task_mappings=task_mappings,
)
if guardrailed_texts and texts_to_check:
await self._apply_guardrail_responses_to_output_texts(
response=response,
responses=guardrailed_texts,
task_mappings=text_task_mappings,
)
# Step 4: Apply guardrailed tool calls back to response
if tool_calls_to_check:
await self._apply_guardrail_responses_to_output_tool_calls(
response=response,
tool_calls=tool_calls_to_check,
task_mappings=tool_call_task_mappings,
)
verbose_proxy_logger.debug(
"OpenAI Chat Completions: Processed output response: %s", response
@ -334,16 +413,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if user_metadata:
request_data["litellm_metadata"] = user_metadata
guardrailed_texts, guardrailed_images = (
await guardrail_to_apply.apply_guardrail(
texts=texts_to_check,
request_data=request_data,
input_type="response",
images=images_to_check if images_to_check else None,
logging_obj=litellm_logging_obj,
)
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
inputs["images"] = images_to_check
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
# Step 4: Apply guardrailed text back to all streaming chunks
# For each choice, replace the combined text across all chunks
await self._apply_guardrail_responses_to_output_streaming(
@ -363,7 +444,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
self, response: Union["ModelResponse", "ModelResponseStream"]
) -> bool:
"""
Check if response has any text content to process.
Check if response has any text content or tool calls to process.
Override this method to customize text content detection.
"""
@ -372,15 +453,29 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if isinstance(response, ModelResponse):
for choice in response.choices:
if isinstance(choice, litellm.Choices):
# Check for text content
if choice.message.content and isinstance(
choice.message.content, str
):
return True
# Check for tool calls
if choice.message.tool_calls and isinstance(
choice.message.tool_calls, list
):
if len(choice.message.tool_calls) > 0:
return True
elif isinstance(response, ModelResponseStream):
for choice in response.choices:
if isinstance(choice, litellm.StreamingChoices):
# Check for text content
if choice.delta.content and isinstance(choice.delta.content, str):
return True
# Check for tool calls
if choice.delta.tool_calls and isinstance(
choice.delta.tool_calls, list
):
if len(choice.delta.tool_calls) > 0:
return True
return False
def _extract_output_text_and_images(
@ -389,23 +484,28 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
choice_idx: int,
texts_to_check: List[str],
images_to_check: List[str],
task_mappings: List[Tuple[int, Optional[int]]],
tool_calls_to_check: List[Dict[str, Any]],
text_task_mappings: List[Tuple[int, Optional[int]]],
tool_call_task_mappings: List[Tuple[int, int]],
) -> None:
"""
Extract text content and images from a response choice.
Extract text content, images, and tool calls from a response choice.
Override this method to customize text/image extraction logic.
Override this method to customize text/image/tool call extraction logic.
"""
verbose_proxy_logger.debug(
"OpenAI Chat Completions: Processing choice: %s", choice
)
# Determine content source based on choice type
# Determine content source and tool calls based on choice type
content = None
tool_calls = None
if isinstance(choice, litellm.Choices):
content = choice.message.content
tool_calls = choice.message.tool_calls
elif isinstance(choice, litellm.StreamingChoices):
content = choice.delta.content
tool_calls = choice.delta.tool_calls
else:
# Unknown choice type, skip processing
return
@ -414,7 +514,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if content and isinstance(content, str):
# Simple string content
texts_to_check.append(content)
task_mappings.append((choice_idx, None))
text_task_mappings.append((choice_idx, None))
elif content and isinstance(content, list):
# List content (e.g., multimodal response)
@ -423,7 +523,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
content_text = content_item.get("text")
if content_text:
texts_to_check.append(content_text)
task_mappings.append((choice_idx, int(content_idx)))
text_task_mappings.append((choice_idx, int(content_idx)))
# Extract images
if content_item.get("type") == "image_url":
@ -433,39 +533,108 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if url:
images_to_check.append(url)
async def _apply_guardrail_responses_to_output(
# Process tool calls if they exist
if tool_calls is not None and isinstance(tool_calls, list):
for tool_call_idx, tool_call in enumerate(tool_calls):
# Convert tool call to dict format for guardrail processing
tool_call_dict = self._convert_tool_call_to_dict(tool_call)
if tool_call_dict:
tool_calls_to_check.append(tool_call_dict)
tool_call_task_mappings.append((choice_idx, int(tool_call_idx)))
def _convert_tool_call_to_dict(
self, tool_call: Union[Dict[str, Any], Any]
) -> Optional[Dict[str, Any]]:
"""
Convert a tool call object to dictionary format.
Tool calls can be either dict or object depending on the type.
"""
if isinstance(tool_call, dict):
return tool_call
elif hasattr(tool_call, "id") and hasattr(tool_call, "function"):
# Convert object to dict
function = tool_call.function
function_dict = {}
if hasattr(function, "name"):
function_dict["name"] = function.name
if hasattr(function, "arguments"):
function_dict["arguments"] = function.arguments
tool_call_dict = {
"id": tool_call.id if hasattr(tool_call, "id") else None,
"type": tool_call.type if hasattr(tool_call, "type") else "function",
"function": function_dict,
}
return tool_call_dict
return None
async def _apply_guardrail_responses_to_output_texts(
self,
response: "ModelResponse",
responses: List[str],
task_mappings: List[Tuple[int, Optional[int]]],
) -> None:
"""
Apply guardrail responses back to output response.
Apply guardrail text responses back to output response.
Override this method to customize how responses are applied.
Override this method to customize how text responses are applied.
"""
for task_idx, guardrail_response in enumerate(responses):
mapping = task_mappings[task_idx]
choice_idx = cast(int, mapping[0])
content_idx_optional = cast(Optional[int], mapping[1])
content = cast(Choices, response.choices[choice_idx]).message.content
choice = cast(Choices, response.choices[choice_idx])
# Handle content
content = choice.message.content
if content is None:
continue
if isinstance(content, str) and content_idx_optional is None:
# Replace string content with guardrail response
cast(Choices, response.choices[choice_idx]).message.content = (
guardrail_response
)
choice.message.content = guardrail_response
elif isinstance(content, list) and content_idx_optional is not None:
# Replace specific text item in list content
cast(Choices, response.choices[choice_idx]).message.content[ # type: ignore
content_idx_optional
][
"text"
] = guardrail_response
choice.message.content[content_idx_optional]["text"] = guardrail_response # type: ignore
async def _apply_guardrail_responses_to_output_tool_calls(
self,
response: "ModelResponse",
tool_calls: List[Dict[str, Any]],
task_mappings: List[Tuple[int, int]],
) -> None:
"""
Apply guardrailed tool calls back to output response.
The guardrail may have modified the tool_calls list in place,
so we apply the modified tool calls back to the original response.
Override this method to customize how tool call responses are applied.
"""
for task_idx, (choice_idx, tool_call_idx) in enumerate(task_mappings):
if task_idx < len(tool_calls):
guardrailed_tool_call = tool_calls[task_idx]
choice = cast(Choices, response.choices[choice_idx])
choice_tool_calls = choice.message.tool_calls
if choice_tool_calls is not None and isinstance(
choice_tool_calls, list
):
if tool_call_idx < len(choice_tool_calls):
# Update the tool call with guardrailed version
existing_tool_call = choice_tool_calls[tool_call_idx]
# Update object attributes (output responses always have typed objects)
if "function" in guardrailed_tool_call:
func_dict = guardrailed_tool_call["function"]
if "arguments" in func_dict:
existing_tool_call.function.arguments = func_dict[
"arguments"
]
if "name" in func_dict:
existing_tool_call.function.name = func_dict["name"]
async def _apply_guardrail_responses_to_output_streaming(
self,

15
litellm/llms/openai/completion/guardrail_translation/handler.py
View file

@ -53,12 +53,13 @@ class OpenAITextCompletionHandler(BaseTranslation):
if isinstance(prompt, str):
# Single string prompt
guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail(
texts=[prompt],
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [prompt]},
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
data["prompt"] = guardrailed_texts[0] if guardrailed_texts else prompt
verbose_proxy_logger.debug(
@ -79,12 +80,13 @@ class OpenAITextCompletionHandler(BaseTranslation):
text_indices.append(idx)
if texts_to_check:
guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail(
texts=texts_to_check,
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": texts_to_check},
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
# Replace guardrailed texts back
for guardrail_idx, prompt_idx in enumerate(text_indices):
@ -152,12 +154,13 @@ class OpenAITextCompletionHandler(BaseTranslation):
if user_metadata:
request_data["litellm_metadata"] = user_metadata
guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail(
texts=texts_to_check,
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": texts_to_check},
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
# Apply guardrailed texts back to choices
for guardrail_idx, choice_idx in enumerate(choice_indices):

5
litellm/llms/openai/image_generation/guardrail_translation/handler.py
View file

@ -52,12 +52,13 @@ class OpenAIImageGenerationHandler(BaseTranslation):
# Apply guardrail to the prompt
if isinstance(prompt, str):
guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail(
texts=[prompt],
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [prompt]},
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
data["prompt"] = guardrailed_texts[0] if guardrailed_texts else prompt
verbose_proxy_logger.debug(

96
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -28,12 +28,19 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
- text: str
"""
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union, cast
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
from openai import BaseModel
from openai import BaseModel
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
from litellm.types.guardrails import GenericGuardrailAPIInputs
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.responses.main import GenericResponseOutputItem, OutputText
if TYPE_CHECKING:
@ -65,17 +72,28 @@ class OpenAIResponsesHandler(BaseTranslation):
Handles both string input and list of message objects.
"""
input_data: Optional[Union[str, "ResponseInputParam"]] = data.get("input")
tools_to_check: List[ChatCompletionToolParam] = []
if input_data is None:
return data
# Handle simple string input
if isinstance(input_data, str):
guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail(
texts=[input_data],
inputs = GenericGuardrailAPIInputs(texts=[input_data])
# Extract and transform tools if present
if "tools" in data and data["tools"]:
self._extract_and_transform_tools(data["tools"], tools_to_check)
if tools_to_check:
inputs["tools"] = tools_to_check
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data
verbose_proxy_logger.debug("OpenAI Responses API: Processed string input")
return data
@ -90,7 +108,7 @@ class OpenAIResponsesHandler(BaseTranslation):
# Track (message_index, content_index) for each text
# content_index is None for string content, int for list content
# Step 1: Extract all text content and images
# Step 1: Extract all text content, images, and tools
for msg_idx, message in enumerate(input_data):
self._extract_input_text_and_images(
message=message,
@ -100,18 +118,26 @@ class OpenAIResponsesHandler(BaseTranslation):
task_mappings=task_mappings,
)
# Extract and transform tools if present
if "tools" in data and data["tools"]:
self._extract_and_transform_tools(data["tools"], tools_to_check)
# Step 2: Apply guardrail to all texts in batch
if texts_to_check:
guardrailed_texts, guardrailed_images = (
await guardrail_to_apply.apply_guardrail(
texts=texts_to_check,
request_data=data,
input_type="request",
images=images_to_check if images_to_check else None,
logging_obj=litellm_logging_obj,
)
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
inputs["images"] = images_to_check
if tools_to_check:
inputs["tools"] = tools_to_check
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
# Step 3: Map guardrail responses back to original input structure
await self._apply_guardrail_responses_to_input(
messages=input_data,
@ -125,6 +151,29 @@ class OpenAIResponsesHandler(BaseTranslation):
return data
def _extract_and_transform_tools(
self,
tools: List[Dict[str, Any]],
tools_to_check: List[ChatCompletionToolParam],
) -> None:
"""
Extract and transform tools from Responses API format to Chat Completion format.
Uses the LiteLLM transformation function to convert Responses API tools
to Chat Completion tools that can be passed to guardrails.
"""
if tools is not None and isinstance(tools, list):
# Transform Responses API tools to Chat Completion tools
(
transformed_tools,
_,
) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(
tools # type: ignore
)
tools_to_check.extend(
cast(List[ChatCompletionToolParam], transformed_tools)
)
def _extract_input_text_and_images(
self,
message: Any, # Can be Dict[str, Any] or ResponseInputParam
@ -254,16 +303,18 @@ class OpenAIResponsesHandler(BaseTranslation):
if user_metadata:
request_data["litellm_metadata"] = user_metadata
guardrailed_texts, guardrailed_images = (
await guardrail_to_apply.apply_guardrail(
texts=texts_to_check,
request_data=request_data,
input_type="response",
images=images_to_check if images_to_check else None,
logging_obj=litellm_logging_obj,
)
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
inputs["images"] = images_to_check
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
# Step 3: Map guardrail responses back to original response structure
await self._apply_guardrail_responses_to_output(
response=response,
@ -288,12 +339,11 @@ class OpenAIResponsesHandler(BaseTranslation):
Process output streaming response by applying guardrails to text content.
"""
string_so_far = self.get_streaming_string_so_far(responses_so_far)
guardrailed_text, _ = await guardrail_to_apply.apply_guardrail(
texts=[string_so_far],
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [string_so_far]},
request_data={},
input_type="response",
logging_obj=litellm_logging_obj,
images=None,
)
return responses_so_far

5
litellm/llms/openai/speech/guardrail_translation/handler.py
View file

@ -50,12 +50,13 @@ class OpenAITextToSpeechHandler(BaseTranslation):
return data
if isinstance(input_text, str):
guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail(
texts=[input_text],
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [input_text]},
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
data["input"] = guardrailed_texts[0] if guardrailed_texts else input_text
verbose_proxy_logger.debug(

5
litellm/llms/openai/transcriptions/guardrail_translation/handler.py
View file

@ -88,12 +88,13 @@ class OpenAIAudioTranscriptionHandler(BaseTranslation):
if user_metadata:
request_data["litellm_metadata"] = user_metadata
guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail(
texts=[original_text],
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [original_text]},
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
)
guardrailed_texts = guardrailed_inputs.get("texts", [])
response.text = guardrailed_texts[0] if guardrailed_texts else original_text
verbose_proxy_logger.debug(

8
litellm/llms/pass_through/guardrail_translation/handler.py
View file

@ -118,8 +118,8 @@ class PassThroughEndpointHandler(BaseTranslation):
return data
# Apply guardrail (pass-through doesn't modify the text, just checks it)
await guardrail_to_apply.apply_guardrail(
texts=[text_to_check],
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [text_to_check]},
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
@ -178,8 +178,8 @@ class PassThroughEndpointHandler(BaseTranslation):
request_data["litellm_metadata"] = user_metadata
# Apply guardrail (pass-through doesn't modify the text, just checks it)
await guardrail_to_apply.apply_guardrail(
texts=[text_to_check],
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [text_to_check]},
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,

10
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -1214,13 +1214,15 @@ async def apply_guardrail(
detail=f"Guardrail '{request.guardrail_name}' not found. Please ensure the guardrail is configured in your LiteLLM proxy.",
)
response_text = await active_guardrail.apply_guardrail(
texts=[request.text],
guardrailed_inputs = await active_guardrail.apply_guardrail(
inputs={"texts": [request.text]},
request_data={},
input_type="request",
images=None,
)
response_text = guardrailed_inputs.get("texts", [])
return ApplyGuardrailResponse(response_text=response_text[0][0])
return ApplyGuardrailResponse(
response_text=response_text[0] if response_text else request.text
)
except Exception as e:
raise handle_exception_on_proxy(e)

20
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -42,7 +42,7 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.secret_managers.main import get_secret_str
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.guardrails import GenericGuardrailAPIInputs, GuardrailEventHooks
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
BedrockContentItem,
@ -54,6 +54,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.utils import (
CallTypes,
CallTypesLiteral,
@ -1245,12 +1246,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
async def apply_guardrail(
self,
texts: List[str],
inputs: "GenericGuardrailAPIInputs",
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
images: Optional[List[str]] = None,
) -> Tuple[List[str], Optional[List[str]]]:
) -> "GenericGuardrailAPIInputs":
"""
Apply Bedrock guardrail to a batch of texts for testing purposes.
@ -1258,17 +1258,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
It creates mock messages to test the guardrail functionality.
Args:
texts: List of texts to analyze
inputs: Dictionary containing texts and optional images
request_data: Request data dictionary for logging metadata
input_type: Whether this is a "request" or "response"
images: Optional list of images (not processed separately)
logging_obj: Optional logging object
Returns:
Tuple of (processed_texts, images) - texts may be masked, images unchanged
GenericGuardrailAPIInputs - processed_texts may be masked, images unchanged
Raises:
Exception: If content is blocked by Bedrock guardrail
"""
texts = inputs.get("texts", [])
try:
verbose_proxy_logger.debug(
f"Bedrock Guardrail: Applying guardrail to {len(texts)} text(s)"
@ -1279,6 +1280,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
mock_messages: List[AllMessageValues] = [
ChatCompletionUserMessage(role="user", content=text) for text in texts
]
request_messages = mock_messages
filter_result = self._prepare_guardrail_messages_for_role(
messages=request_messages
@ -1297,6 +1299,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
# Apply any masking that was applied by the guardrail
output_list = bedrock_response.get("output")
if output_list:
# If the guardrail returned modified content, use that
@ -1319,7 +1322,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
"Bedrock Guardrail: Successfully applied guardrail"
)
return masked_texts or texts, images
inputs["texts"] = masked_texts
return inputs
except Exception as e:
verbose_proxy_logger.error(

18
litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py
View file

@ -15,7 +15,6 @@ from typing import (
List,
Literal,
Optional,
Tuple,
Union,
)
@ -30,7 +29,7 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.guardrails import GenericGuardrailAPIInputs, GuardrailEventHooks
from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import (
EnkryptAIProcessedResult,
EnkryptAIResponse,
@ -481,27 +480,28 @@ class EnkryptAIGuardrails(CustomGuardrail):
async def apply_guardrail(
self,
texts: List[str],
inputs: "GenericGuardrailAPIInputs",
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
images: Optional[List[str]] = None,
) -> Tuple[List[str], Optional[List[str]]]:
) -> "GenericGuardrailAPIInputs":
"""
Apply EnkryptAI guardrail to a batch of texts.
Args:
texts: List of texts to check for attacks
inputs: Dictionary containing texts and optional images
request_data: Request data dictionary containing metadata
input_type: Whether this is a "request" or "response"
images: Optional list of images (not used by EnkryptAI)
logging_obj: Optional logging object
Returns:
Tuple of (texts, images) - texts unchanged if passed, images unchanged
GenericGuardrailAPIInputs - texts unchanged if passed, images unchanged
Raises:
ValueError: If any attacks are detected
"""
texts = inputs.get("texts", [])
# Check each text for attacks
for text in texts:
result = await self._call_enkryptai_guardrails(
@ -517,7 +517,7 @@ class EnkryptAIGuardrails(CustomGuardrail):
error_message = self._create_error_message(processed_result)
raise ValueError(error_message)
return texts, images
return inputs
async def async_post_call_streaming_iterator_hook(
self,

44
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py
View file

@ -14,7 +14,7 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.guardrails import GenericGuardrailAPIInputs, GuardrailEventHooks
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIMetadata,
GenericGuardrailAPIRequest,
@ -144,22 +144,24 @@ class GenericGuardrailAPI(CustomGuardrail):
async def apply_guardrail(
self,
texts: List[str],
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
images: Optional[List[str]] = None,
) -> Tuple[List[str], Optional[List[str]]]:
) -> GenericGuardrailAPIInputs:
"""
Apply the Generic Guardrail API to the given text.
Apply the Generic Guardrail API to the given inputs.
This is the main method that gets called by the framework.
Args:
texts: List of texts to check
inputs: Dictionary containing:
- texts: List of texts to check
- images: Optional list of images to check
- tool_calls: Optional list of tool calls to check
request_data: Request data dictionary containing user_api_key_dict and other metadata
input_type: Whether this is a "request" or "response" guardrail
images: Optional list of images to check
logging_obj: Optional logging object for tracking the guardrail execution
Returns:
Tuple of (processed texts, processed images)
@ -169,6 +171,11 @@ class GenericGuardrailAPI(CustomGuardrail):
"""
verbose_proxy_logger.debug("Generic Guardrail API: Applying guardrail to text")
# Extract texts and images from inputs
texts = inputs.get("texts", [])
images = inputs.get("images")
tools = inputs.get("tools")
# Use provided request_data or create an empty dict
if request_data is None:
request_data = {}
@ -193,6 +200,7 @@ class GenericGuardrailAPI(CustomGuardrail):
texts=texts,
request_data=user_metadata,
images=images,
tools=tools,
additional_provider_specific_params=additional_params,
input_type=input_type,
)
@ -230,17 +238,19 @@ class GenericGuardrailAPI(CustomGuardrail):
)
raise Exception(f"Content blocked by guardrail: {error_message}")
elif guardrail_response.action == "GUARDRAIL_INTERVENED":
# Content was modified by the guardrail
if guardrail_response.texts:
verbose_proxy_logger.debug("Generic Guardrail API modified text")
return guardrail_response.texts, guardrail_response.images
# Action is NONE or no modifications needed
return (
guardrail_response.texts or texts,
guardrail_response.images or images,
)
return_inputs = GenericGuardrailAPIInputs(texts=texts)
if guardrail_response.texts:
return_inputs["texts"] = guardrail_response.texts
if guardrail_response.images:
return_inputs["images"] = guardrail_response.images
elif images:
return_inputs["images"] = images
if guardrail_response.tools:
return_inputs["tools"] = guardrail_response.tools
elif tools:
return_inputs["tools"] = tools
return return_inputs
except Exception as e:
# Check if it's already an exception we raised

28
litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py
View file

@ -27,6 +27,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.guardrails import GenericGuardrailAPIInputs
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import (
BlockedWord,
@ -304,12 +305,11 @@ class ContentFilterGuardrail(CustomGuardrail):
async def apply_guardrail(
self,
texts: List[str],
inputs: "GenericGuardrailAPIInputs",
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
images: Optional[List[str]] = None,
) -> Tuple[List[str], Optional[List[str]]]:
) -> "GenericGuardrailAPIInputs":
"""
Apply content filtering guardrail to a batch of texts.
@ -317,17 +317,19 @@ class ContentFilterGuardrail(CustomGuardrail):
either blocking the request or masking the sensitive content.
Args:
texts: List of texts to apply the guardrail to
inputs: Dictionary containing texts and optional images
request_data: Request data dictionary for logging metadata
input_type: Whether this is a "request" or "response"
images: Optional list of images (not processed)
logging_obj: Optional logging object
Returns:
Tuple of (processed_texts, images) - texts may be masked, images unchanged
GenericGuardrailAPIInputs - processed_texts may be masked, images unchanged
Raises:
HTTPException: If sensitive content is detected and action is BLOCK
"""
texts = inputs.get("texts", [])
verbose_proxy_logger.debug(
f"ContentFilterGuardrail: Applying guardrail to {len(texts)} text(s)"
)
@ -386,7 +388,8 @@ class ContentFilterGuardrail(CustomGuardrail):
verbose_proxy_logger.debug(
"ContentFilterGuardrail: Guardrail applied successfully"
)
return processed_texts, images
inputs["texts"] = processed_texts
return inputs
async def async_post_call_streaming_iterator_hook(
self,
@ -423,12 +426,17 @@ class ContentFilterGuardrail(CustomGuardrail):
if isinstance(choice.delta.content, str):
# Check the chunk content using apply_guardrail
try:
processed_content = await self.apply_guardrail(
texts=[choice.delta.content],
guardrailed_inputs = await self.apply_guardrail(
inputs={"texts": [choice.delta.content]},
input_type="response",
images=None,
request_data=request_data,
)
processed_texts = guardrailed_inputs.get("texts", [])
processed_content = (
processed_texts[0]
if processed_texts
else choice.delta.content
)
if processed_content != choice.delta.content:
choice.delta.content = processed_content
verbose_proxy_logger.debug(

12
litellm/proxy/guardrails/guardrail_hooks/presidio.py
View file

@ -29,9 +29,11 @@ import aiohttp
import litellm # noqa: E401
from litellm import get_secret
from litellm._logging import verbose_proxy_logger
from litellm.types.guardrails import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm._uuid import uuid
from litellm.caching.caching import DualCache
from litellm.exceptions import BlockedPiiEntityError
@ -713,17 +715,18 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
async def apply_guardrail(
self,
texts: List[str],
inputs: "GenericGuardrailAPIInputs",
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
images: Optional[List[str]] = None,
) -> Tuple[List[str], Optional[List[str]]]:
) -> "GenericGuardrailAPIInputs":
"""
UI will call this function to check:
1. If the connection to the guardrail is working
2. When Testing the guardrail with some text, this function will be called with the input text and returns a text after applying the guardrail
"""
texts = inputs.get("texts", [])
new_texts = []
for text in texts:
modified_text = await self.check_pii(
@ -733,7 +736,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
request_data=request_data or {},
)
new_texts.append(modified_text)
return new_texts, images
inputs["texts"] = new_texts
return inputs
def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
"""

1
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -193,6 +193,7 @@ class UnifiedLLMGuardrails(CustomLogger):
"guardrail_to_apply", None
)
# Get sampling rate from guardrail config or optional_params, default to 5
sampling_rate = 5
if guardrail_to_apply is not None:

17
litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py
View file

@ -4,7 +4,7 @@
#
# +-------------------------------------------------------------+
import os
from typing import TYPE_CHECKING, List, Literal, Optional, Tuple
from typing import TYPE_CHECKING, Literal, Optional
from fastapi import HTTPException
@ -14,6 +14,7 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.guardrails import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -71,27 +72,27 @@ class ZscalerAIGuard(CustomGuardrail):
async def apply_guardrail(
self,
texts: List[str],
inputs: "GenericGuardrailAPIInputs",
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
images: Optional[List[str]] = None,
) -> Tuple[List[str], Optional[List[str]]]:
) -> "GenericGuardrailAPIInputs":
"""
Apply Zscaler AI Guard guardrail to batch of texts.
Args:
texts: List of texts to check
inputs: Dictionary containing texts and optional images
request_data: Request data dictionary containing metadata
input_type: Whether this is a "request" or "response"
images: Optional list of images (not used by Zscaler)
logging_obj: Optional logging object
Returns:
Tuple of (processed_texts, images) - texts unchanged if passed, images unchanged
GenericGuardrailAPIInputs - texts unchanged if passed, images unchanged
Raises:
Exception: If content is blocked by Zscaler AI Guard
"""
texts = inputs.get("texts", [])
try:
verbose_proxy_logger.debug(f"ZscalerAIGuard: Checking {len(texts)} text(s)")
@ -143,7 +144,7 @@ class ZscalerAIGuard(CustomGuardrail):
raise e
verbose_proxy_logger.debug("ZscalerAIGuard: Successfully applied guardrail.")
return texts, images
return inputs
def extract_blocking_info(self, response):
"""

7
litellm/types/guardrails.py
View file

@ -5,6 +5,7 @@ from typing import Any, Dict, List, Literal, Optional, Union
from pydantic import BaseModel, ConfigDict, Field
from typing_extensions import Required, TypedDict
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import (
EnkryptAIGuardrailConfigs,
)
@ -738,3 +739,9 @@ class PatchGuardrailRequest(BaseModel):
guardrail_name: Optional[str] = None
litellm_params: Optional[BaseLitellmParams] = None
guardrail_info: Optional[Dict[str, Any]] = None
class GenericGuardrailAPIInputs(TypedDict, total=False):
texts: List[str]
images: List[str]
tools: List[ChatCompletionToolParam]

7
litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py
View file

@ -3,6 +3,7 @@ from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel, Field
from typing_extensions import TypedDict
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
@ -59,6 +60,7 @@ class GenericGuardrailAPIRequest:
litellm_trace_id: Optional[str],
additional_provider_specific_params: Optional[Dict[str, Any]] = None,
images: Optional[List[str]] = None,
tools: Optional[List[ChatCompletionToolParam]] = None,
):
self.texts = texts
self.request_data = request_data
@ -69,12 +71,14 @@ class GenericGuardrailAPIRequest:
self.input_type = input_type
self.litellm_call_id = litellm_call_id
self.litellm_trace_id = litellm_trace_id
self.tools = tools
def to_dict(self) -> dict:
return {
"texts": self.texts,
"request_data": self.request_data,
"images": self.images,
"tools": self.tools,
"additional_provider_specific_params": self.additional_provider_specific_params,
"input_type": self.input_type,
"litellm_call_id": self.litellm_call_id,
@ -87,6 +91,7 @@ class GenericGuardrailAPIResponse:
texts: Optional[List[str]]
images: Optional[List[str]]
tools: Optional[List[ChatCompletionToolParam]]
action: str
blocked_reason: Optional[str]
@ -96,11 +101,13 @@ class GenericGuardrailAPIResponse:
texts: Optional[List[str]] = None,
blocked_reason: Optional[str] = None,
images: Optional[List[str]] = None,
tools: Optional[List[ChatCompletionToolParam]] = None,
):
self.action = action
self.blocked_reason = blocked_reason
self.texts = texts
self.images = images
self.tools = tools
@classmethod
def from_dict(cls, data: dict) -> "GenericGuardrailAPIResponse":

25
tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py
View file

@ -39,11 +39,12 @@ async def test_bedrock_apply_guardrail_success():
mock_api_request.return_value = mock_response
# Test the apply_guardrail method with new signature
result, _ = await guardrail.apply_guardrail(
texts=["This is a test message with some content"],
guardrailed_inputs = await guardrail.apply_guardrail(
inputs={"texts": ["This is a test message with some content"]},
request_data={},
input_type="request",
)
result = guardrailed_inputs.get("texts", [])
# Verify the result
assert result == ["This is a test message with some content"]
@ -71,7 +72,9 @@ async def test_bedrock_apply_guardrail_blocked():
# Test the apply_guardrail method should raise an exception
with pytest.raises(Exception) as exc_info:
await guardrail.apply_guardrail(
texts=["This is blocked content"], request_data={}, input_type="request"
inputs={"texts": ["This is blocked content"]},
request_data={},
input_type="request",
)
assert "Content blocked by Bedrock guardrail" in str(exc_info.value)
@ -100,11 +103,12 @@ async def test_bedrock_apply_guardrail_with_masking():
mock_api_request.return_value = mock_response
# Test the apply_guardrail method with new signature
result, _ = await guardrail.apply_guardrail(
texts=["This is a test message with sensitive content"],
guardrailed_inputs = await guardrail.apply_guardrail(
inputs={"texts": ["This is a test message with sensitive content"]},
request_data={},
input_type="request",
)
result = guardrailed_inputs.get("texts", [])
# Verify the result contains the masked content
assert result == ["This is a test message with [REDACTED] content"]
@ -130,7 +134,9 @@ async def test_bedrock_apply_guardrail_api_failure():
# Test the apply_guardrail method should raise an exception
with pytest.raises(Exception) as exc_info:
await guardrail.apply_guardrail(
texts=["This is a test message"], request_data={}, input_type="request"
inputs={"texts": ["This is a test message"]},
request_data={},
input_type="request",
)
# The error message should contain the original exception
@ -215,11 +221,12 @@ async def test_bedrock_apply_guardrail_filters_request_messages_when_flag_enable
) as mock_api:
mock_api.return_value = {"action": "ALLOWED"}
result, _ = await guardrail.apply_guardrail(
texts=["latest question"],
guardrailed_inputs = await guardrail.apply_guardrail(
inputs={"texts": ["latest question"]},
request_data=request_data,
input_type="request",
)
result = guardrailed_inputs.get("texts", [])
assert mock_api.called
_, kwargs = mock_api.call_args
@ -250,7 +257,7 @@ async def test_bedrock_apply_guardrail_filters_request_messages_when_flag_enable
with pytest.raises(Exception, match="policy") as exc_info:
await guardrail.apply_guardrail(
texts=["blocked"],
inputs={"texts": ["blocked"]},
request_data=request_data,
input_type="request",
)

0
tests/test_litellm/llms/openai/chat/__init__.py Normal file
View file

0
tests/test_litellm/llms/openai/chat/guardrail_translation/__init__.py Normal file
View file

518
tests/test_litellm/llms/openai/chat/guardrail_translation/test_handler.py Normal file
View file

@ -0,0 +1,518 @@
"""
Unit tests for OpenAI Chat Completions Guardrail Translation Handler
Tests the handler's ability to process input/output for Chat Completions API
with guardrail transformations, including tool calls.
"""
import json
import os
import sys
from typing import Any, List, Literal, Optional, Tuple
from unittest.mock import AsyncMock, MagicMock
import pytest
sys.path.insert(
0, os.path.abspath("../../../../../../..")
) # Adds the parent directory to the system path
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.openai.chat.guardrail_translation.handler import (
OpenAIChatCompletionsHandler,
)
from litellm.types.guardrails import GenericGuardrailAPIInputs
from litellm.types.utils import (
ChatCompletionMessageToolCall,
Choices,
Function,
Message,
ModelResponse,
)
class MockGuardrail(CustomGuardrail):
"""Mock guardrail for testing that transforms text and tool calls"""
def __init__(self, guardrail_name: str = "test"):
super().__init__(guardrail_name=guardrail_name)
self.last_inputs = None
self.last_request_data = None
self.tool_calls_modified = False
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> Tuple[List[str], Optional[List[str]]]:
"""Mock apply_guardrail that uppercases text and modifies tool calls"""
self.last_inputs = inputs
self.last_request_data = request_data
# Return modified texts (uppercase for testing)
texts = inputs.get("texts", [])
modified_texts = [text.upper() for text in texts]
# Modify tool calls in place if present
tool_calls = inputs.get("tool_calls", [])
if tool_calls:
self.tool_calls_modified = True
for tool_call in tool_calls:
if isinstance(tool_call, dict) and "function" in tool_call:
function = tool_call["function"]
if "arguments" in function:
# Modify arguments to uppercase JSON string
try:
args_dict = json.loads(function["arguments"])
# Uppercase all string values
for key, value in args_dict.items():
if isinstance(value, str):
args_dict[key] = value.upper()
function["arguments"] = json.dumps(args_dict)
except json.JSONDecodeError:
# If not JSON, just uppercase the string
function["arguments"] = function["arguments"].upper()
return modified_texts, []
class TestOpenAIChatCompletionsHandlerToolCallsInput:
"""Test input processing with tool calls"""
@pytest.mark.asyncio
async def test_extract_tool_calls_from_input_messages(self):
"""Test that tool calls are extracted from input messages"""
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail()
# Create input data with tool calls (assistant message)
data = {
"messages": [
{"role": "user", "content": "What's the weather in SF?"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_123",
"type": "function",
"function": {
"name": "get_weather",
"arguments": json.dumps(
{"location": "San Francisco", "unit": "celsius"}
),
},
}
],
},
]
}
# Process the input
await handler.process_input_messages(data, guardrail)
# Verify tool calls were extracted and passed to guardrail
assert guardrail.last_inputs is not None
assert "tool_calls" in guardrail.last_inputs
assert len(guardrail.last_inputs["tool_calls"]) == 1
tool_call = guardrail.last_inputs["tool_calls"][0]
assert tool_call["id"] == "call_123"
assert tool_call["function"]["name"] == "get_weather"
# Note: tool call arguments may already be modified by guardrail
# Check that it contains location parameter
assert "location" in tool_call["function"]["arguments"]
# Verify tool call was modified by guardrail
assert guardrail.tool_calls_modified is True
# Verify the message was updated with modified tool call
modified_tool_call = data["messages"][1]["tool_calls"][0]
args = json.loads(modified_tool_call["function"]["arguments"])
assert args["location"] == "SAN FRANCISCO" # Should be uppercased
assert args["unit"] == "CELSIUS" # Should be uppercased
@pytest.mark.asyncio
async def test_extract_tool_calls_and_text_from_input_messages(self):
"""Test that both tool calls and text content are extracted"""
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail()
# Create input data with both text and tool calls
data = {
"messages": [
{"role": "user", "content": "What's the weather?"},
{
"role": "assistant",
"content": "Let me check that for you.",
"tool_calls": [
{
"id": "call_456",
"type": "function",
"function": {
"name": "get_current_weather",
"arguments": json.dumps({"location": "Boston"}),
},
}
],
},
]
}
# Process the input
await handler.process_input_messages(data, guardrail)
# Verify both texts and tool calls were extracted
assert guardrail.last_inputs is not None
assert "texts" in guardrail.last_inputs
assert "tool_calls" in guardrail.last_inputs
# Should have 2 texts (user message + assistant message)
assert len(guardrail.last_inputs["texts"]) == 2
assert "What's the weather?" in guardrail.last_inputs["texts"]
assert "Let me check that for you." in guardrail.last_inputs["texts"]
# Should have 1 tool call
assert len(guardrail.last_inputs["tool_calls"]) == 1
assert (
guardrail.last_inputs["tool_calls"][0]["function"]["name"]
== "get_current_weather"
)
# Verify text content was modified
assert data["messages"][0]["content"] == "WHAT'S THE WEATHER?"
assert data["messages"][1]["content"] == "LET ME CHECK THAT FOR YOU."
# Verify tool call was modified
modified_tool_call = data["messages"][1]["tool_calls"][0]
args = json.loads(modified_tool_call["function"]["arguments"])
assert args["location"] == "BOSTON"
@pytest.mark.asyncio
async def test_extract_multiple_tool_calls_from_input(self):
"""Test extraction of multiple tool calls"""
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail()
# Create input data with multiple tool calls
data = {
"messages": [
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {
"name": "get_weather",
"arguments": json.dumps({"location": "NYC"}),
},
},
{
"id": "call_2",
"type": "function",
"function": {
"name": "get_time",
"arguments": json.dumps({"timezone": "EST"}),
},
},
],
}
]
}
# Process the input
await handler.process_input_messages(data, guardrail)
# Verify multiple tool calls were extracted
assert guardrail.last_inputs is not None
assert "tool_calls" in guardrail.last_inputs
assert len(guardrail.last_inputs["tool_calls"]) == 2
# Verify both tool calls
tool_calls = guardrail.last_inputs["tool_calls"]
assert tool_calls[0]["function"]["name"] == "get_weather"
assert tool_calls[1]["function"]["name"] == "get_time"
# Verify both were modified
modified_tool_calls = data["messages"][0]["tool_calls"]
args1 = json.loads(modified_tool_calls[0]["function"]["arguments"])
args2 = json.loads(modified_tool_calls[1]["function"]["arguments"])
assert args1["location"] == "NYC"
assert args2["timezone"] == "EST"
@pytest.mark.asyncio
async def test_tool_calls_separate_from_texts(self):
"""Test that tool calls are passed as a separate parameter, not mixed with texts"""
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail()
data = {
"messages": [
{"role": "user", "content": "Get weather for LA"},
{
"role": "assistant",
"content": "Sure!",
"tool_calls": [
{
"id": "call_xyz",
"type": "function",
"function": {
"name": "get_weather",
"arguments": json.dumps({"city": "Los Angeles"}),
},
}
],
},
]
}
# Process the input
await handler.process_input_messages(data, guardrail)
# Verify tool calls and texts are separate
assert guardrail.last_inputs is not None
texts = guardrail.last_inputs.get("texts", [])
tool_calls = guardrail.last_inputs.get("tool_calls", [])
# Texts should only contain the content strings
assert len(texts) == 2
assert "Get weather for LA" in texts
assert "Sure!" in texts
# Tool call arguments should NOT be in texts
assert not any("Los Angeles" in text for text in texts)
# Tool calls should be separate
assert len(tool_calls) == 1
assert tool_calls[0]["function"]["name"] == "get_weather"
# Check that it contains city parameter (may be modified by guardrail)
assert "city" in tool_calls[0]["function"]["arguments"]
@pytest.mark.asyncio
async def test_no_tool_calls_in_input(self):
"""Test that messages without tool calls work correctly"""
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail()
# Create input data without tool calls
data = {
"messages": [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there!"},
]
}
# Process the input
await handler.process_input_messages(data, guardrail)
# Verify no tool calls were passed to guardrail
assert guardrail.last_inputs is not None
tool_calls = guardrail.last_inputs.get("tool_calls", [])
assert len(tool_calls) == 0
# Verify text was still processed
assert len(guardrail.last_inputs["texts"]) == 2
assert data["messages"][0]["content"] == "HELLO"
assert data["messages"][1]["content"] == "HI THERE!"
@pytest.mark.asyncio
async def test_empty_tool_calls_list(self):
"""Test that empty tool_calls list is handled correctly"""
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail()
data = {
"messages": [
{"role": "assistant", "content": "Hello", "tool_calls": []},
]
}
# Process the input
await handler.process_input_messages(data, guardrail)
# Verify empty tool_calls doesn't cause issues
assert guardrail.last_inputs is not None
tool_calls = guardrail.last_inputs.get("tool_calls", [])
assert len(tool_calls) == 0
class TestOpenAIChatCompletionsHandlerToolCallsOutput:
"""Test output processing with tool calls"""
@pytest.mark.asyncio
async def test_extract_tool_calls_from_output_response(self):
"""Test that tool calls are extracted from output responses"""
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail()
# Create a mock response with tool calls
response = ModelResponse(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion",
choices=[
Choices(
finish_reason="tool_calls",
index=0,
message=Message(
content=None,
role="assistant",
tool_calls=[
ChatCompletionMessageToolCall(
id="call_789",
type="function",
function=Function(
name="search_database",
arguments=json.dumps({"query": "python tutorials"}),
),
)
],
),
)
],
)
# Process the output
await handler.process_output_response(response, guardrail)
# Verify tool calls were extracted and passed to guardrail
assert guardrail.last_inputs is not None
assert "tool_calls" in guardrail.last_inputs
assert len(guardrail.last_inputs["tool_calls"]) == 1
tool_call = guardrail.last_inputs["tool_calls"][0]
assert tool_call["function"]["name"] == "search_database"
# Check that it contains query parameter (may be modified by guardrail)
assert "query" in tool_call["function"]["arguments"]
# Verify tool call was modified in response
response_tool_call = response.choices[0].message.tool_calls[0]
args = json.loads(response_tool_call.function.arguments)
assert args["query"] == "PYTHON TUTORIALS" # Should be uppercased
@pytest.mark.asyncio
async def test_extract_tool_calls_and_content_from_output(self):
"""Test extraction of both content and tool calls from output"""
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail()
response = ModelResponse(
id="chatcmpl-456",
created=1234567890,
model="gpt-4",
object="chat.completion",
choices=[
Choices(
finish_reason="tool_calls",
index=0,
message=Message(
content="I'll search for that information.",
role="assistant",
tool_calls=[
ChatCompletionMessageToolCall(
id="call_999",
type="function",
function=Function(
name="web_search",
arguments=json.dumps(
{"keywords": "litellm documentation"}
),
),
)
],
),
)
],
)
# Process the output
await handler.process_output_response(response, guardrail)
# Verify both texts and tool calls were extracted
assert guardrail.last_inputs is not None
assert "texts" in guardrail.last_inputs
assert "tool_calls" in guardrail.last_inputs
assert len(guardrail.last_inputs["texts"]) == 1
assert "I'll search for that information." in guardrail.last_inputs["texts"]
assert len(guardrail.last_inputs["tool_calls"]) == 1
# Verify both were modified
assert (
response.choices[0].message.content == "I'LL SEARCH FOR THAT INFORMATION."
)
response_tool_call = response.choices[0].message.tool_calls[0]
args = json.loads(response_tool_call.function.arguments)
assert args["keywords"] == "LITELLM DOCUMENTATION"
@pytest.mark.asyncio
async def test_extract_multiple_tool_calls_from_output(self):
"""Test extraction of multiple tool calls from output"""
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail()
response = ModelResponse(
id="chatcmpl-789",
created=1234567890,
model="gpt-4",
object="chat.completion",
choices=[
Choices(
finish_reason="tool_calls",
index=0,
message=Message(
content=None,
role="assistant",
tool_calls=[
ChatCompletionMessageToolCall(
id="call_1",
type="function",
function=Function(
name="get_weather",
arguments=json.dumps({"location": "Tokyo"}),
),
),
ChatCompletionMessageToolCall(
id="call_2",
type="function",
function=Function(
name="get_news",
arguments=json.dumps({"topic": "technology"}),
),
),
],
),
)
],
)
# Process the output
await handler.process_output_response(response, guardrail)
# Verify multiple tool calls were extracted
assert guardrail.last_inputs is not None
assert "tool_calls" in guardrail.last_inputs
assert len(guardrail.last_inputs["tool_calls"]) == 2
# Verify both tool calls
tool_calls = guardrail.last_inputs["tool_calls"]
assert tool_calls[0]["function"]["name"] == "get_weather"
assert tool_calls[1]["function"]["name"] == "get_news"
# Verify both were modified
response_tool_calls = response.choices[0].message.tool_calls
args1 = json.loads(response_tool_calls[0].function.arguments)
args2 = json.loads(response_tool_calls[1].function.arguments)
assert args1["location"] == "TOKYO"
assert args2["topic"] == "TECHNOLOGY"
if __name__ == "__main__":
# Run the tests
pytest.main([__file__, "-v"])

19
tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py
View file

@ -191,7 +191,9 @@ class TestContentFilterGuardrail:
with pytest.raises(HTTPException) as exc_info:
await guardrail.apply_guardrail(
texts=["My SSN is 123-45-6789"], request_data={}, input_type="request"
inputs={"texts": ["My SSN is 123-45-6789"]},
request_data={},
input_type="request",
)
assert exc_info.value.status_code == 400
@ -215,11 +217,12 @@ class TestContentFilterGuardrail:
patterns=patterns,
)
result, _ = await guardrail.apply_guardrail(
texts=["Contact me at test@example.com"],
guardrailed_inputs = await guardrail.apply_guardrail(
inputs={"texts": ["Contact me at test@example.com"]},
request_data={},
input_type="request",
)
result = guardrailed_inputs.get("texts", [])
assert result is not None
assert len(result) == 1
@ -243,11 +246,12 @@ class TestContentFilterGuardrail:
blocked_words=blocked_words,
)
result, _ = await guardrail.apply_guardrail(
texts=["This is PROPRIETARY information"],
guardrailed_inputs = await guardrail.apply_guardrail(
inputs={"texts": ["This is PROPRIETARY information"]},
request_data={},
input_type="request",
)
result = guardrailed_inputs.get("texts", [])
assert result is not None
assert len(result) == 1
@ -277,11 +281,12 @@ class TestContentFilterGuardrail:
patterns=patterns,
)
result, _ = await guardrail.apply_guardrail(
texts=["Contact user@test.com or SSN: 123-45-6789"],
guardrailed_inputs = await guardrail.apply_guardrail(
inputs={"texts": ["Contact user@test.com or SSN: 123-45-6789"]},
request_data={},
input_type="request",
)
result = guardrailed_inputs.get("texts", [])
assert result is not None
assert len(result) == 1

85
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py
View file

@ -4,6 +4,7 @@ Tests for Generic Guardrail API integration
This test file tests the Generic Guardrail API implementation,
specifically focusing on metadata extraction and passing.
"""
import os
from unittest.mock import AsyncMock, MagicMock, patch
@ -108,9 +109,14 @@ class TestGenericGuardrailAPIConfiguration:
headers={"Authorization": "Bearer test-key"},
additional_provider_specific_params={"custom_param": "value"},
)
assert guardrail.api_base == "https://api.test.guardrail.com/beta/litellm_basic_guardrail_api"
assert (
guardrail.api_base
== "https://api.test.guardrail.com/beta/litellm_basic_guardrail_api"
)
assert guardrail.headers == {"Authorization": "Bearer test-key"}
assert guardrail.additional_provider_specific_params == {"custom_param": "value"}
assert guardrail.additional_provider_specific_params == {
"custom_param": "value"
}
def test_init_with_env_vars(self):
"""Test initialization with environment variables"""
@ -121,7 +127,10 @@ class TestGenericGuardrailAPIConfiguration:
},
):
guardrail = GenericGuardrailAPI()
assert guardrail.api_base == "https://env.api.guardrail.com/beta/litellm_basic_guardrail_api"
assert (
guardrail.api_base
== "https://env.api.guardrail.com/beta/litellm_basic_guardrail_api"
)
def test_init_without_api_base_raises_error(self):
"""Test that initialization without API base raises ValueError"""
@ -134,14 +143,20 @@ class TestGenericGuardrailAPIConfiguration:
guardrail = GenericGuardrailAPI(
api_base="https://api.test.guardrail.com/v1",
)
assert guardrail.api_base == "https://api.test.guardrail.com/v1/beta/litellm_basic_guardrail_api"
assert (
guardrail.api_base
== "https://api.test.guardrail.com/v1/beta/litellm_basic_guardrail_api"
)
def test_api_base_not_duplicated(self):
"""Test that endpoint path is not duplicated if already present"""
guardrail = GenericGuardrailAPI(
api_base="https://api.test.guardrail.com/beta/litellm_basic_guardrail_api",
)
assert guardrail.api_base == "https://api.test.guardrail.com/beta/litellm_basic_guardrail_api"
assert (
guardrail.api_base
== "https://api.test.guardrail.com/beta/litellm_basic_guardrail_api"
)
class TestMetadataExtraction:
@ -164,7 +179,7 @@ class TestMetadataExtraction:
generic_guardrail.async_handler, "post", return_value=mock_response
) as mock_post:
await generic_guardrail.apply_guardrail(
texts=["Who is Ishaan?"],
inputs=GenericGuardrailAPIInputs(texts=["Who is Ishaan?"]),
request_data=mock_request_data_input,
input_type="request",
)
@ -180,7 +195,10 @@ class TestMetadataExtraction:
request_metadata = json_payload["request_data"]
# Verify metadata was extracted from request_data["metadata"]
assert request_metadata["user_api_key_hash"] == "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"
assert (
request_metadata["user_api_key_hash"]
== "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"
)
assert request_metadata["user_api_key_user_id"] == "default_user_id"
assert request_metadata["user_api_key_user_email"] == "test@example.com"
assert request_metadata["user_api_key_team_id"] == "test-team"
@ -192,7 +210,7 @@ class TestMetadataExtraction:
"""Test extracting metadata from output response (litellm_metadata field)"""
# Create request_data as it would be created by the handler
user_dict = mock_user_api_key_dict.model_dump()
# Transform to prefixed keys (as done by BaseTranslation)
litellm_metadata = {}
for key, value in user_dict.items():
@ -219,7 +237,7 @@ class TestMetadataExtraction:
generic_guardrail.async_handler, "post", return_value=mock_api_response
) as mock_post:
await generic_guardrail.apply_guardrail(
texts=["hey i'm ishaan!"],
inputs={"texts": ["hey i'm ishaan!"]},
request_data=request_data,
input_type="response",
)
@ -263,7 +281,7 @@ class TestMetadataExtraction:
generic_guardrail.async_handler, "post", return_value=mock_response
) as mock_post:
await generic_guardrail.apply_guardrail(
texts=["test"],
inputs={"texts": ["test"]},
request_data=request_data,
input_type="request",
)
@ -278,9 +296,7 @@ class TestMetadataExtraction:
assert request_metadata["user_api_key_user_id"] == "test-user"
@pytest.mark.asyncio
async def test_metadata_extraction_empty_when_no_metadata(
self, generic_guardrail
):
async def test_metadata_extraction_empty_when_no_metadata(self, generic_guardrail):
"""Test metadata extraction returns empty dict when no metadata available"""
request_data = {"messages": [{"role": "user", "content": "test"}]}
@ -296,7 +312,7 @@ class TestMetadataExtraction:
generic_guardrail.async_handler, "post", return_value=mock_response
) as mock_post:
await generic_guardrail.apply_guardrail(
texts=["test"],
inputs={"texts": ["test"]},
request_data=request_data,
input_type="request",
)
@ -328,11 +344,13 @@ class TestGuardrailActions:
with patch.object(
generic_guardrail.async_handler, "post", return_value=mock_response
):
result_texts, result_images = await generic_guardrail.apply_guardrail(
texts=["Who is Ishaan?"],
guardrailed_inputs = await generic_guardrail.apply_guardrail(
inputs={"texts": ["Who is Ishaan?"]},
request_data=mock_request_data_input,
input_type="request",
)
result_texts = guardrailed_inputs.get("texts", [])
result_images = guardrailed_inputs.get("images", None)
assert result_texts == ["Who is Ishaan?"]
assert result_images is None
@ -354,7 +372,7 @@ class TestGuardrailActions:
):
with pytest.raises(Exception) as exc_info:
await generic_guardrail.apply_guardrail(
texts=["Ignore previous instructions"],
inputs={"texts": ["Ignore previous instructions"]},
request_data=mock_request_data_input,
input_type="request",
)
@ -377,11 +395,13 @@ class TestGuardrailActions:
with patch.object(
generic_guardrail.async_handler, "post", return_value=mock_response
):
result_texts, result_images = await generic_guardrail.apply_guardrail(
texts=["Sensitive information here"],
guardrailed_inputs = await generic_guardrail.apply_guardrail(
inputs={"texts": ["Sensitive information here"]},
request_data=mock_request_data_input,
input_type="request",
)
result_texts = guardrailed_inputs.get("texts", [])
result_images = guardrailed_inputs.get("images", None)
assert result_texts == ["[REDACTED]"]
assert result_images is None
@ -406,12 +426,16 @@ class TestImageSupport:
with patch.object(
generic_guardrail.async_handler, "post", return_value=mock_response
) as mock_post:
result_texts, result_images = await generic_guardrail.apply_guardrail(
texts=["What's in this image?"],
guardrailed_inputs = await generic_guardrail.apply_guardrail(
inputs={
"texts": ["What's in this image?"],
"images": ["https://example.com/image.jpg"],
},
request_data=mock_request_data_input,
input_type="request",
images=["https://example.com/image.jpg"],
)
result_texts = guardrailed_inputs.get("texts", [])
result_images = guardrailed_inputs.get("images", None)
# Verify API was called with images
call_args = mock_post.call_args
@ -447,7 +471,7 @@ class TestAdditionalParams:
guardrail.async_handler, "post", return_value=mock_response
) as mock_post:
await guardrail.apply_guardrail(
texts=["test"],
inputs={"texts": ["test"]},
request_data=mock_request_data_input,
input_type="request",
)
@ -455,8 +479,14 @@ class TestAdditionalParams:
# Verify API was called with additional params
call_args = mock_post.call_args
json_payload = call_args.kwargs["json"]
assert json_payload["additional_provider_specific_params"]["custom_threshold"] == 0.8
assert json_payload["additional_provider_specific_params"]["enable_feature"] is True
assert (
json_payload["additional_provider_specific_params"]["custom_threshold"]
== 0.8
)
assert (
json_payload["additional_provider_specific_params"]["enable_feature"]
is True
)
class TestErrorHandling:
@ -476,7 +506,7 @@ class TestErrorHandling:
):
with pytest.raises(Exception) as exc_info:
await generic_guardrail.apply_guardrail(
texts=["test"],
inputs={"texts": ["test"]},
request_data=mock_request_data_input,
input_type="request",
)
@ -495,10 +525,9 @@ class TestErrorHandling:
):
with pytest.raises(Exception) as exc_info:
await generic_guardrail.apply_guardrail(
texts=["test"],
inputs={"texts": ["test"]},
request_data=mock_request_data_input,
input_type="request",
)
assert "Generic Guardrail API failed" in str(exc_info.value)

4
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py
View file

@ -571,7 +571,7 @@ async def test_presidio_sets_guardrail_information_in_request_data():
with patch.object(presidio, "check_pii", mock_check_pii):
await presidio.apply_guardrail(
texts=["Test message"],
inputs={"texts": ["Test message"]},
request_data=request_data,
input_type="request",
)
@ -623,7 +623,7 @@ async def test_request_data_flows_to_apply_guardrail():
with patch.object(presidio, "check_pii", mock_check_pii):
result = await presidio.apply_guardrail(
texts=["Test message"],
inputs={"texts": ["Test message"]},
request_data=request_data,
input_type="request",
)