From bc8eeab1681cb12dfc0229941174075a5c11dd30 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 26 Dec 2025 09:59:12 +0900 Subject: [PATCH] feat: add run_guardrails_on_stream_item_complete param and tool override in guardrail --- .../transformation.py | 2 +- litellm/integrations/custom_guardrail.py | 4 + .../guardrail_translation/base_translation.py | 14 + .../guardrail_translation/handler.py | 306 +++++++++++++++++- .../unified_guardrail/unified_guardrail.py | 23 +- .../guardrails/guardrail_initializers.py | 8 + litellm/types/guardrails.py | 4 + ...test_openai_responses_guardrail_handler.py | 26 ++ 8 files changed, 369 insertions(+), 18 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 55a8e665bbd..96b24c08688 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -934,7 +934,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): function_chunk = ChatCompletionToolCallFunctionChunk( name=output_item.get("name", None), - arguments="", # responses API sends everything again, we don't + arguments=output_item.get("arguments", None), ) # Add provider_specific_fields to function if present diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index fe0ce208ee6..964d1de836b 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -86,6 +86,7 @@ class CustomGuardrail(CustomLogger): mask_request_content: bool = False, mask_response_content: bool = False, violation_message_template: Optional[str] = None, + run_guardrails_on_stream_item_complete: bool = False, **kwargs, ): """ @@ -98,6 +99,8 @@ class CustomGuardrail(CustomLogger): default_on: If True, the guardrail will be run by default on all requests mask_request_content: If True, the guardrail will mask the request content mask_response_content: If True, the guardrail will mask the response content + violation_message_template: If provided, this template will override the default violation message returned when the guardrail is triggered + run_guardrails_on_stream_item_complete: If True, the guardrail will be executed when a streamed output item is completed """ self.guardrail_name = guardrail_name self.supported_event_hooks = supported_event_hooks @@ -108,6 +111,7 @@ class CustomGuardrail(CustomLogger): self.mask_request_content: bool = mask_request_content self.mask_response_content: bool = mask_response_content self.violation_message_template: Optional[str] = violation_message_template + self.run_guardrails_on_stream_item_complete: bool = run_guardrails_on_stream_item_complete if supported_event_hooks: ## validate event_hook is in supported_event_hooks diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 7106c207bd6..d69d53143a6 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -98,3 +98,17 @@ class BaseTranslation(ABC): Optional to override in subclasses. """ return responses_so_far + + def is_stream_item_complete(self, responses_so_far: List[Any]) -> bool: + """ + Check if a streaming item has completed. + + This is used to determine when to run guardrails on a complete item + + Args: + responses_so_far: List of streaming responses received so far + + Returns: + bool: True if a stream item has completed, False otherwise + """ + return False diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 9b8f15c7623..02db45ea2f8 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -45,13 +45,14 @@ from litellm.responses.litellm_completion_transformation.transformation import ( from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, ChatCompletionToolParam, + ResponsesAPIStreamEvents, ) from litellm.types.responses.main import ( GenericResponseOutputItem, OutputFunctionToolCall, OutputText, ) -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail @@ -113,7 +114,12 @@ class OpenAIResponsesHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) guardrailed_texts = guardrailed_inputs.get("texts", []) + guardrailed_tools = guardrailed_inputs.get("tools") data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data + if guardrailed_tools is not None: + data["tools"] = self._convert_chat_completion_tools_to_responses_tools( + cast(List[ChatCompletionToolParam], guardrailed_tools) + ) verbose_proxy_logger.debug("OpenAI Responses API: Processed string input") return data @@ -158,6 +164,7 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts = guardrailed_inputs.get("texts", []) + guardrailed_tools = guardrailed_inputs.get("tools") # Step 3: Map guardrail responses back to original input structure await self._apply_guardrail_responses_to_input( @@ -166,6 +173,11 @@ class OpenAIResponsesHandler(BaseTranslation): task_mappings=task_mappings, ) + if guardrailed_tools is not None: + data["tools"] = self._convert_chat_completion_tools_to_responses_tools( + cast(List[ChatCompletionToolParam], guardrailed_tools) + ) + verbose_proxy_logger.debug( "OpenAI Responses API: Processed input messages: %s", input_data ) @@ -187,13 +199,156 @@ class OpenAIResponsesHandler(BaseTranslation): # Transform Responses API tools to Chat Completion tools ( transformed_tools, - _, + web_search_options, ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( tools # type: ignore ) tools_to_check.extend( cast(List[ChatCompletionToolParam], transformed_tools) ) + if web_search_options is not None: + # For guardrail checks we surface web_search as an explicit tool entry. + tools_to_check.append( + cast(ChatCompletionToolParam, {"type": "web_search"}) + ) + + + def _convert_chat_completion_tools_to_responses_tools( + self, chat_completion_tools: List[ChatCompletionToolParam] + ) -> List[Dict[str, Any]]: + """Convert Chat Completion-style tools back to Responses API definitions.""" + + transformation_handler = LiteLLMResponsesTransformationHandler() + responses_tools = transformation_handler._convert_tools_to_responses_format( + cast(List[Dict[str, Any]], chat_completion_tools) + ) + return cast(List[Dict[str, Any]], responses_tools) + + def _convert_chat_completion_tool_calls_to_responses_tool_calls( + self, tool_calls: List[ChatCompletionToolCallChunk] + ) -> List[ResponseFunctionToolCall]: + """Convert Chat Completion tool calls to Responses API function call objects.""" + + if not tool_calls: + return [] + + normalized_tool_calls: List[Dict[str, Any]] = [] + for tool_call in tool_calls: + if isinstance(tool_call, dict): + normalized_tool_calls.append(cast(Dict[str, Any], tool_call)) + continue + + function_payload: Dict[str, Any] = {} + function_object = getattr(tool_call, "function", {}) + if isinstance(function_object, dict): + function_payload = function_object + elif hasattr(function_object, "model_dump"): + function_payload = cast(Dict[str, Any], function_object.model_dump()) + else: + function_payload = cast(Dict[str, Any], getattr(function_object, "__dict__", {})) + + normalized_tool_calls.append( + { + "id": getattr(tool_call, "id", None), + "type": getattr(tool_call, "type", "function"), + "function": function_payload, + } + ) + + chat_completion_response = ModelResponse( + choices=[ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": "", + "tool_calls": normalized_tool_calls, + }, + } + ] + ) + + return LiteLLMCompletionResponsesConfig.transform_chat_completion_tools_to_responses_tools( + chat_completion_response=chat_completion_response + ) + + def _apply_guardrail_tool_calls_to_response_output( + self, + response: Any, + guardrail_tool_calls: List[ChatCompletionToolCallChunk], + ) -> None: + """Replace response output tool-call items using guardrail-modified tool calls.""" + + converted_tool_calls = self._convert_chat_completion_tool_calls_to_responses_tool_calls( + guardrail_tool_calls + ) + output_items = self._get_response_output(response) + if output_items is None: + return + + filtered_output_items = [] + for output_item in output_items: + if not self._is_tool_call_output_item(output_item): + filtered_output_items.append(output_item) + + filtered_output_items.extend(converted_tool_calls) + if isinstance(response, dict): + response["output"] = filtered_output_items + else: + response.output = filtered_output_items + + def _apply_guardrail_tool_calls_to_stream_chunk( + self, + chunk: Any, + guardrail_tool_calls: List[ChatCompletionToolCallChunk], + ) -> None: + """Update streaming chunk tool call payload with guardrail output.""" + + converted_tool_calls = self._convert_chat_completion_tool_calls_to_responses_tool_calls( + guardrail_tool_calls + ) + + serialized_tool_call = ( + self._serialize_response_tool_call(converted_tool_calls[0], as_dict=True) + if converted_tool_calls + else {} + ) + + # Handle both dict and Pydantic model chunks + if isinstance(chunk, dict): + chunk["item"] = serialized_tool_call + elif hasattr(chunk, "item"): + setattr(chunk, "item", serialized_tool_call) + else: + verbose_proxy_logger.warning( + "Cannot set item on chunk of type %s", type(chunk) + ) + + def _get_response_output(self, response: Any) -> Optional[List[Any]]: + if isinstance(response, dict): + return response.get("output") + if hasattr(response, "output"): + return getattr(response, "output") + return None + + def _serialize_response_tool_call( + self, tool_call: ResponseFunctionToolCall, as_dict: bool + ) -> Any: + if as_dict: + if hasattr(tool_call, "model_dump"): + return tool_call.model_dump(exclude_none=True) + return cast(Dict[str, Any], dict(tool_call)) + return tool_call + + def _is_tool_call_output_item(self, output_item: Any) -> bool: + if isinstance(output_item, (ResponseFunctionToolCall, OutputFunctionToolCall)): + return True + if isinstance(output_item, BaseModel): + return getattr(output_item, "type", None) == "function_call" + if isinstance(output_item, dict): + return output_item.get("type") == "function_call" + return False def _extract_input_text_and_images( self, @@ -353,6 +508,7 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts = guardrailed_inputs.get("texts", []) + guardrailed_tool_calls = guardrailed_inputs.get("tool_calls") # Step 3: Map guardrail responses back to original response structure await self._apply_guardrail_responses_to_output( @@ -361,12 +517,47 @@ class OpenAIResponsesHandler(BaseTranslation): task_mappings=task_mappings, ) + if guardrailed_tool_calls is not None: + self._apply_guardrail_tool_calls_to_response_output( + response=response, + guardrail_tool_calls=cast( + List[ChatCompletionToolCallChunk], guardrailed_tool_calls + ), + ) + verbose_proxy_logger.debug( "OpenAI Responses API: Processed output response: %s", response ) return response + def _convert_streaming_tool_calls_to_dicts( + self, tool_calls: List[Any] + ) -> List[ChatCompletionToolCallChunk]: + """ + Convert streaming tool calls (ChatCompletionDeltaToolCall) to dict format. + + ChatCompletionDeltaToolCall objects from streaming responses need to be + converted to dicts so they can be processed by guardrails. + + Args: + tool_calls: List of ChatCompletionDeltaToolCall objects + + Returns: + List of tool calls as ChatCompletionToolCallChunk dicts + """ + converted: List[ChatCompletionToolCallChunk] = [] + for tc in tool_calls: + if hasattr(tc, "model_dump"): + converted.append(cast(ChatCompletionToolCallChunk, tc.model_dump())) + elif hasattr(tc, "__dict__"): + converted.append(cast(ChatCompletionToolCallChunk, dict(tc.__dict__))) + else: + verbose_proxy_logger.warning( + "Unexpected tool call type in streaming response: %s", type(tc) + ) + return converted + async def process_output_streaming_response( self, responses_so_far: List[Any], @@ -388,16 +579,24 @@ class OpenAIResponsesHandler(BaseTranslation): tool_calls = model_response_stream.choices[0].delta.tool_calls if tool_calls: - _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={ - "tool_calls": cast( - List[ChatCompletionToolCallChunk], tool_calls - ) - }, + tool_calls_as_dicts = self._convert_streaming_tool_calls_to_dicts( + tool_calls + ) + + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs={"tool_calls": tool_calls_as_dicts}, request_data={}, input_type="response", logging_obj=litellm_logging_obj, ) + guardrailed_tool_calls = guardrailed_inputs.get("tool_calls") + if guardrailed_tool_calls is not None: + self._apply_guardrail_tool_calls_to_stream_chunk( + chunk=final_chunk, + guardrail_tool_calls=cast( + List[ChatCompletionToolCallChunk], guardrailed_tool_calls + ), + ) return responses_so_far elif final_chunk.get("type") == "response.completed": # convert openai response to model response @@ -408,22 +607,37 @@ class OpenAIResponsesHandler(BaseTranslation): handle_raw_dict_callback=None, ) + if not model_response_choices: + return responses_so_far + tool_calls = model_response_choices[0].message.tool_calls text = model_response_choices[0].message.content guardrail_inputs = GenericGuardrailAPIInputs() if text: guardrail_inputs["texts"] = [text] if tool_calls: - guardrail_inputs["tool_calls"] = cast( - List[ChatCompletionToolCallChunk], tool_calls + tool_calls_as_dicts = self._convert_streaming_tool_calls_to_dicts( + tool_calls ) + guardrail_inputs["tool_calls"] = tool_calls_as_dicts if tool_calls: - _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=guardrail_inputs, request_data={}, input_type="response", logging_obj=litellm_logging_obj, ) + guardrailed_tool_calls = guardrailed_inputs.get("tool_calls") + if guardrailed_tool_calls is not None: + response_payload = final_chunk.get("response") + if response_payload is not None: + self._apply_guardrail_tool_calls_to_response_output( + response=response_payload, + guardrail_tool_calls=cast( + List[ChatCompletionToolCallChunk], + guardrailed_tool_calls, + ), + ) return responses_so_far # model_response_stream = OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(final_chunk) # tool_calls = model_response_stream.choices[0].tool_calls @@ -452,6 +666,27 @@ class OpenAIResponsesHandler(BaseTranslation): """ return "".join([response.get("text", "") for response in responses_so_far]) + def is_stream_item_complete(self, responses_so_far: List[Any]) -> bool: + """ + Check if a streaming item has completed. + + For OpenAI Responses API, an item is complete when we receive + a chunk with type "response.output_item.done" or "response.completed". + + Args: + responses_so_far: List of streaming responses received so far + + Returns: + bool: True if the last chunk indicates an item has completed, False otherwise + """ + if not responses_so_far: + return False + + final_chunk = responses_so_far[-1] + + chunk_type = final_chunk.get("type") + return chunk_type in (ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ResponsesAPIStreamEvents.RESPONSE_COMPLETED) + def _has_text_content(self, response: "ResponsesAPIResponse") -> bool: """ Check if response has any text content to process. @@ -589,9 +824,9 @@ class OpenAIResponsesHandler(BaseTranslation): else: continue + task_mappings.append((output_idx, int(content_idx))) if text_content: texts_to_check.append(text_content) - task_mappings.append((output_idx, int(content_idx))) async def _apply_guardrail_responses_to_output( self, @@ -606,13 +841,56 @@ class OpenAIResponsesHandler(BaseTranslation): """ # Handle both dict and Pydantic object responses if isinstance(response, dict): - response_output = response.get("output", []) + response_output = cast(List[Any], response.get("output", [])) elif hasattr(response, "output"): - response_output = response.output or [] + response_output = cast(List[Any], response.output or []) else: return + # If task_mappings is empty but we have responses, create new output items + # This happens when response.output only contains tool calls but guardrail returns text + if not task_mappings and responses: + verbose_proxy_logger.debug( + "OpenAI Responses API: task_mappings is empty but responses exist. Creating new output items." + ) + if isinstance(response, dict): + for idx, guardrail_response in enumerate(responses): + new_output_item = { + "type": "message", + "id": f"msg_guardrail_{idx + 1}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": guardrail_response}], + } + response_output.append(new_output_item) + response["output"] = response_output + else: + for idx, guardrail_response in enumerate(responses): + new_output_item = GenericResponseOutputItem( + type="message", + id=f"msg_guardrail_{idx + 1}", + status="completed", + role="assistant", + content=[ + OutputText( + type="output_text", + text=guardrail_response, + annotations=None, + ) + ], + ) + response_output.append(new_output_item) + response.output = response_output + return + for task_idx, guardrail_response in enumerate(responses): + if task_idx >= len(task_mappings): + verbose_proxy_logger.warning( + "OpenAI Responses API: task_idx %d exceeds task_mappings length %d. Skipping remaining responses.", + task_idx, + len(task_mappings), + ) + break mapping = task_mappings[task_idx] output_idx = cast(int, mapping[0]) content_idx = cast(int, mapping[1]) 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 a1bbf36ac0c..f19bc9f2381 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -274,6 +274,8 @@ class UnifiedLLMGuardrails(CustomLogger): yield item return + run_guardrails_on_stream_item_complete = guardrail_to_apply.run_guardrails_on_stream_item_complete + event_type: GuardrailEventHooks = GuardrailEventHooks.post_call if ( guardrail_to_apply.should_run_guardrail( @@ -328,19 +330,34 @@ class UnifiedLLMGuardrails(CustomLogger): yield item continue - # Process chunk based on sampling rate - if chunk_counter % sampling_rate == 0: + # Check if we should process this chunk + should_process = False + endpoint_translation = None + # Process based on sampling rate + if chunk_counter % sampling_rate == 0: + should_process = True verbose_proxy_logger.debug( "Processing streaming chunk %s (sampling_rate=%s) with guardrail %s", chunk_counter, sampling_rate, guardrail_to_apply.guardrail_name, ) - + # Or if run_guardrails_on_stream_item_complete is enabled and item is complete + elif run_guardrails_on_stream_item_complete: endpoint_translation = endpoint_guardrail_translation_mappings[ CallTypes(call_type) ]() + if endpoint_translation.is_stream_item_complete(responses_so_far): + should_process = True + + # Process chunk if conditions are met + if should_process: + # Reuse endpoint_translation if already created for completion check + if endpoint_translation is None: + endpoint_translation = endpoint_guardrail_translation_mappings[ + CallTypes(call_type) + ]() processed_items = ( await endpoint_translation.process_output_streaming_response( diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 66b41005c4e..580bee18c6a 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -18,6 +18,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): guardrailVersion=litellm_params.guardrailVersion, default_on=litellm_params.default_on, disable_exception_on_block=litellm_params.disable_exception_on_block, + run_guardrails_on_stream_item_complete=litellm_params.run_guardrails_on_stream_item_complete, mask_request_content=litellm_params.mask_request_content, mask_response_content=litellm_params.mask_response_content, aws_region_name=litellm_params.aws_region_name, @@ -46,6 +47,7 @@ def initialize_lakera(litellm_params: LitellmParams, guardrail: Guardrail): event_hook=litellm_params.mode, category_thresholds=litellm_params.category_thresholds, default_on=litellm_params.default_on, + run_guardrails_on_stream_item_complete=litellm_params.run_guardrails_on_stream_item_complete, ) litellm.logging_callback_manager.add_litellm_callback(_lakera_callback) return _lakera_callback @@ -60,6 +62,7 @@ def initialize_lakera_v2(litellm_params: LitellmParams, guardrail: Guardrail): guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, default_on=litellm_params.default_on, + run_guardrails_on_stream_item_complete=litellm_params.run_guardrails_on_stream_item_complete, project_id=litellm_params.project_id, payload=litellm_params.payload, breakdown=litellm_params.breakdown, @@ -88,6 +91,7 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): presidio_ad_hoc_recognizers=litellm_params.presidio_ad_hoc_recognizers, mock_redacted_text=litellm_params.mock_redacted_text, default_on=litellm_params.default_on, + run_guardrails_on_stream_item_complete=litellm_params.run_guardrails_on_stream_item_complete, pii_entities_config=litellm_params.pii_entities_config, presidio_score_thresholds=litellm_params.presidio_score_thresholds, presidio_analyzer_api_base=litellm_params.presidio_analyzer_api_base, @@ -139,6 +143,7 @@ def initialize_hide_secrets(litellm_params: LitellmParams, guardrail: Guardrail) event_hook=litellm_params.mode, guardrail_name=guardrail.get("guardrail_name", ""), default_on=litellm_params.default_on, + run_guardrails_on_stream_item_complete=litellm_params.run_guardrails_on_stream_item_complete, ) litellm.logging_callback_manager.add_litellm_callback(_secret_detection_object) return _secret_detection_object @@ -166,6 +171,7 @@ def initialize_tool_permission(litellm_params: LitellmParams, guardrail: Guardra on_disallowed_action=getattr(litellm_params, "on_disallowed_action", "block"), default_on=litellm_params.default_on, violation_message_template=litellm_params.violation_message_template, + run_guardrails_on_stream_item_complete=litellm_params.run_guardrails_on_stream_item_complete, ) litellm.logging_callback_manager.add_litellm_callback(_tool_permission_callback) return _tool_permission_callback @@ -186,6 +192,7 @@ def initialize_lasso( mask=litellm_params.mask, event_hook=litellm_params.mode, default_on=litellm_params.default_on, + run_guardrails_on_stream_item_complete=litellm_params.run_guardrails_on_stream_item_complete, ) litellm.logging_callback_manager.add_litellm_callback(_lasso_callback) @@ -211,6 +218,7 @@ def initialize_panw_prisma_airs(litellm_params, guardrail): or "https://service.api.aisecurity.paloaltonetworks.com/v1/scan/sync/request", profile_name=litellm_params.profile_name, default_on=litellm_params.default_on, + run_guardrails_on_stream_item_complete=litellm_params.run_guardrails_on_stream_item_complete, mask_on_block=getattr(litellm_params, "mask_on_block", False), mask_request_content=getattr(litellm_params, "mask_request_content", False), mask_response_content=getattr(litellm_params, "mask_response_content", False), diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 7a1388ed8ba..e8c12d224e2 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -615,6 +615,10 @@ class BaseLitellmParams(BaseModel): # works for new and patch update guardrails description="Custom message when a guardrail blocks an action. Supports placeholders like {tool_name}, {rule_id}, and {default_message}.", ) + run_guardrails_on_stream_item_complete: Optional[bool] = Field( + default=None, description="Will be executed when a streamed output item is completed" + ) + # Model Armor params template_id: Optional[str] = Field( default=None, description="The ID of your Model Armor template" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index a2849ab91a2..9bdb71ece41 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -182,6 +182,32 @@ class TestOpenAIResponsesHandlerInputProcessing: # Empty string should be processed assert result["input"][1]["content"] == " [GUARDRAILED]" + @pytest.mark.asyncio + async def test_process_input_preserves_web_search_tool(self): + """Ensure web_search tools remain after guardrail processing""" + handler = OpenAIResponsesHandler() + guardrail = MockGuardrail(guardrail_name="test") + + data = { + "input": "Search the news", + "model": "gpt-4", + "tools": [ + { + "type": "web_search_preview", + "search_context_size": "medium", + "user_location": {"country": "US"}, + } + ], + } + + result = await handler.process_input_messages(data, guardrail) + + assert "tools" in result + assert any( + tool.get("type") in {"web_search", "web_search_preview"} + for tool in result["tools"] + ) + class TestOpenAIResponsesHandlerOutputProcessing: """Test output processing functionality"""