From 32013f63a0e03f6fc198cba40dea7cf9ee48b117 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Wed, 3 Dec 2025 21:20:39 -0800 Subject: [PATCH] Guardrail API - support tool call checks on OpenAI `/chat/completions`, OpenAI `/responses`, Anthropic `/v1/messages` (#17459) * 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 --- .../mock_bedrock_guardrail_server.py | 1 + .../adding_provider/generic_guardrail_api.md | 87 +++ litellm/integrations/custom_guardrail.py | 17 +- .../chat/guardrail_translation/handler.py | 93 +++- .../rerank/guardrail_translation/handler.py | 5 +- .../chat/guardrail_translation/handler.py | 357 ++++++++---- .../guardrail_translation/handler.py | 15 +- .../guardrail_translation/handler.py | 5 +- .../guardrail_translation/handler.py | 96 +++- .../speech/guardrail_translation/handler.py | 5 +- .../guardrail_translation/handler.py | 5 +- .../guardrail_translation/handler.py | 8 +- .../proxy/guardrails/guardrail_endpoints.py | 10 +- .../guardrail_hooks/bedrock_guardrails.py | 20 +- .../guardrail_hooks/enkryptai/enkryptai.py | 18 +- .../generic_guardrail_api.py | 44 +- .../litellm_content_filter/content_filter.py | 28 +- .../guardrails/guardrail_hooks/presidio.py | 12 +- .../unified_guardrail/unified_guardrail.py | 1 + .../zscaler_ai_guard/zscaler_ai_guard.py | 17 +- litellm/types/guardrails.py | 7 + .../guardrail_hooks/generic_guardrail_api.py | 7 + .../test_bedrock_apply_guardrail.py | 25 +- .../test_litellm/llms/openai/chat/__init__.py | 0 .../chat/guardrail_translation/__init__.py | 0 .../guardrail_translation/test_handler.py | 518 ++++++++++++++++++ .../content_filter/test_content_filter.py | 19 +- .../test_generic_guardrail_api.py | 85 ++- .../guardrail_hooks/test_presidio.py | 4 +- 29 files changed, 1234 insertions(+), 275 deletions(-) create mode 100644 tests/test_litellm/llms/openai/chat/__init__.py create mode 100644 tests/test_litellm/llms/openai/chat/guardrail_translation/__init__.py create mode 100644 tests/test_litellm/llms/openai/chat/guardrail_translation/test_handler.py diff --git a/cookbook/mock_guardrail_server/mock_bedrock_guardrail_server.py b/cookbook/mock_guardrail_server/mock_bedrock_guardrail_server.py index c9f7f13f6aa..fd53ece6604 100644 --- a/cookbook/mock_guardrail_server/mock_bedrock_guardrail_server.py +++ b/cookbook/mock_guardrail_server/mock_bedrock_guardrail_server.py @@ -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"] diff --git a/docs/my-website/docs/adding_provider/generic_guardrail_api.md b/docs/my-website/docs/adding_provider/generic_guardrail_api.md index 9392deeb5db..f8c07b25f9d 100644 --- a/docs/my-website/docs/adding_provider/generic_guardrail_api.md +++ b/docs/my-website/docs/adding_provider/generic_guardrail_api.md @@ -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") ``` diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 507a754a7e0..782f9460044 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -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, diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index e248e409de4..9b6511f151c 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -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 diff --git a/litellm/llms/cohere/rerank/guardrail_translation/handler.py b/litellm/llms/cohere/rerank/guardrail_translation/handler.py index e042295ab0b..6893a5991c3 100644 --- a/litellm/llms/cohere/rerank/guardrail_translation/handler.py +++ b/litellm/llms/cohere/rerank/guardrail_translation/handler.py @@ -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( diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 2b381e28bc0..76aa6f730c2 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -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, diff --git a/litellm/llms/openai/completion/guardrail_translation/handler.py b/litellm/llms/openai/completion/guardrail_translation/handler.py index f8b733567d3..73d08cfead4 100644 --- a/litellm/llms/openai/completion/guardrail_translation/handler.py +++ b/litellm/llms/openai/completion/guardrail_translation/handler.py @@ -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): diff --git a/litellm/llms/openai/image_generation/guardrail_translation/handler.py b/litellm/llms/openai/image_generation/guardrail_translation/handler.py index 53ee994c48f..842a64b1878 100644 --- a/litellm/llms/openai/image_generation/guardrail_translation/handler.py +++ b/litellm/llms/openai/image_generation/guardrail_translation/handler.py @@ -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( diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 7ef6576159e..9377eb4e193 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -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 diff --git a/litellm/llms/openai/speech/guardrail_translation/handler.py b/litellm/llms/openai/speech/guardrail_translation/handler.py index 71edd4f2801..4c2f71477be 100644 --- a/litellm/llms/openai/speech/guardrail_translation/handler.py +++ b/litellm/llms/openai/speech/guardrail_translation/handler.py @@ -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( diff --git a/litellm/llms/openai/transcriptions/guardrail_translation/handler.py b/litellm/llms/openai/transcriptions/guardrail_translation/handler.py index 18678a9878b..ac416f42c81 100644 --- a/litellm/llms/openai/transcriptions/guardrail_translation/handler.py +++ b/litellm/llms/openai/transcriptions/guardrail_translation/handler.py @@ -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( diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index ae96335cec3..0702086a578 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -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, diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 395fdb249d3..bb78383ce44 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 3932e1220b4..e0fb3192401 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -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( diff --git a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py index 49401ca7921..493d432eebb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py @@ -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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index f806e8b4802..33912b8fcd9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 776b9d36d5e..6756b188847 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -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( diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 6b19685f07e..d183b688edd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -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: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index d234dffc511..c16cb89b785 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -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: diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py index 562f0d40272..bebb87f8d21 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py @@ -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): """ diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 2c0b10ec67d..ea3a6cc4ede 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -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] diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index d2acc3e3403..66823270e87 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -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": diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py index 203bd05c57c..8d98f56cd3a 100644 --- a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py +++ b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py @@ -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", ) diff --git a/tests/test_litellm/llms/openai/chat/__init__.py b/tests/test_litellm/llms/openai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/__init__.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_handler.py new file mode 100644 index 00000000000..09f3f8dcf0f --- /dev/null +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_handler.py @@ -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"]) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index 2b192dad5d5..b5c093fdf06 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index e6ebada4d2e..f3578fbcefa 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -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) - diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 3e19437fe83..a6b2ae5b3a0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -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", )