diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 86ca6625629..76060623564 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -364,6 +364,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): ) guardrailed_texts = guardrailed_inputs.get("texts", []) + guardrail_returned_tool_calls = "tool_calls" in guardrailed_inputs + guardrailed_tool_calls = guardrailed_inputs.get("tool_calls", []) # Step 3: Map guardrail responses back to original response structure if guardrailed_texts and texts_to_check: @@ -372,12 +374,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation): responses=guardrailed_texts, task_mappings=text_task_mappings, ) + elif guardrailed_texts: + self._replace_output_with_guardrail_text( + response=response, + guardrailed_texts=guardrailed_texts, + ) # Step 4: Apply guardrailed tool calls back to response - if tool_calls_to_check: + if guardrail_returned_tool_calls and tool_calls_to_check: await self._apply_guardrail_responses_to_output_tool_calls( response=response, - tool_calls=tool_calls_to_check, + tool_calls=guardrailed_tool_calls, # type: ignore[arg-type] task_mappings=tool_call_task_mappings, ) @@ -743,6 +750,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Override this method to customize how tool call responses are applied. """ + if not tool_calls: + for choice_idx, _ in task_mappings: + choice = cast(Choices, response.choices[choice_idx]) + choice.message.tool_calls = None + if choice.finish_reason == "tool_calls" and choice.message.content: + choice.finish_reason = "stop" + return + 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] @@ -765,6 +780,25 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if "name" in func_dict: existing_tool_call.function.name = func_dict["name"] + def _replace_output_with_guardrail_text( + self, + response: "ModelResponse", + guardrailed_texts: List[str], + ) -> None: + """ + Apply guardrail text when the original response had no text slots. + + Tool-call-only responses have no text mappings, but response guardrails can + still intentionally replace the output with a block/mask message. + """ + if not response.choices or not guardrailed_texts: + return + + choice = cast(Choices, response.choices[0]) + choice.message.content = guardrailed_texts[0] + if choice.finish_reason == "tool_calls": + choice.finish_reason = "stop" + async def _apply_guardrail_responses_to_output_streaming( self, responses: List["ModelResponseStream"], diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index a2c37002942..3b41463dc11 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -8,8 +8,7 @@ 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 +from typing import Any, Literal, Optional import pytest @@ -509,6 +508,60 @@ class TestOpenAIChatCompletionsHandlerToolCallsInput: class TestOpenAIChatCompletionsHandlerToolCallsOutput: """Test output processing with tool calls""" + @pytest.mark.asyncio + async def test_response_guardrail_can_replace_tool_call_only_output(self): + """Test a response guardrail can replace tool-call-only output with text.""" + handler = OpenAIChatCompletionsHandler() + + class BlockingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + assert input_type == "response" + assert len(inputs.get("tool_calls", [])) == 1 + return { + "texts": ["guardrail blocked the request"], + "tool_calls": [], + } + + response = ModelResponse( + id="chatcmpl-tool-call-blocked", + 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="get_weather", + arguments=json.dumps({"location": "Paris"}), + ), + ) + ], + ), + ) + ], + ) + + await handler.process_output_response(response, BlockingGuardrail()) + + choice = response.choices[0] + assert choice.message.content == "guardrail blocked the request" + assert choice.message.tool_calls is None + assert choice.finish_reason == "stop" + @pytest.mark.asyncio async def test_extract_tool_calls_from_output_response(self): """Test that tool calls are extracted from output responses""" @@ -765,7 +818,7 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: This test verifies the fix for the bug where accessing chunk.choices[0] would raise IndexError when a streaming chunk has an empty choices list. """ - from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + from litellm.types.utils import ModelResponseStream handler = OpenAIChatCompletionsHandler() guardrail = MockPassThroughGuardrail(guardrail_name="test")