diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py index d64c786d9c6..baf3884238b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py @@ -136,12 +136,19 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai ) new_messages: Optional[List[AllMessageValues]] = data.get("messages") if new_messages is None: - # Responses API uses "input" instead of "messages" + # Responses API uses "input" instead of "messages". + # Input can be a plain string or a list that may contain + # non-message items (e.g. tool-call outputs without a "role" + # key). Filter to only message-shaped dicts so get_user_prompt() + # doesn't break its reverse iteration on a non-message item. input_data = data.get("input") if isinstance(input_data, str): new_messages = [{"role": "user", "content": input_data}] elif isinstance(input_data, list): - new_messages = input_data + new_messages = [ + item for item in input_data + if isinstance(item, dict) and "role" in item + ] if new_messages is None: verbose_proxy_logger.warning( "Azure Prompt Shield: not running guardrail. No messages in data" diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py index bafec44d214..e3c8a23369b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py @@ -229,12 +229,19 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr ) new_messages: Optional[List[AllMessageValues]] = data.get("messages") if new_messages is None: - # Responses API uses "input" instead of "messages" + # Responses API uses "input" instead of "messages". + # Input can be a plain string or a list that may contain + # non-message items (e.g. tool-call outputs without a "role" + # key). Filter to only message-shaped dicts so get_user_prompt() + # doesn't break its reverse iteration on a non-message item. input_data = data.get("input") if isinstance(input_data, str): new_messages = [{"role": "user", "content": input_data}] elif isinstance(input_data, list): - new_messages = input_data + new_messages = [ + item for item in input_data + if isinstance(item, dict) and "role" in item + ] if new_messages is None: verbose_proxy_logger.warning( "Azure Text Moderation: not running guardrail. No messages in data" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py index 69535789b12..e31d517ba7a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py @@ -265,3 +265,123 @@ def test_split_preserves_whitespace(): original = ("line one\n" + "line two\t\tcol\n" + " indented\n") * 200 chunks = guardrail.split_text_by_words(original, 500) assert "".join(chunks) == original + + +# ---------- Responses API (data["input"]) tests ---------- + + +@pytest.mark.asyncio +async def test_azure_prompt_shield_responses_api_string_input(): + """Test guardrail works when Responses API sends input as a plain string.""" + guardrail = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure_prompt_shield", + api_key="test_key", + api_base="test_base", + ) + with patch.object(guardrail, "async_make_request") as mock_request: + mock_request.return_value = { + "userPromptAnalysis": {"attackDetected": False}, + "documentsAnalysis": [], + } + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + cache=None, + data={"input": "What is the capital of France?"}, + call_type="completion", + ) + + mock_request.assert_called_once() + assert ( + mock_request.call_args.kwargs["user_prompt"] + == "What is the capital of France?" + ) + + +@pytest.mark.asyncio +async def test_azure_prompt_shield_responses_api_list_input(): + """Test guardrail works when Responses API sends input as a message list.""" + guardrail = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure_prompt_shield", + api_key="test_key", + api_base="test_base", + ) + with patch.object(guardrail, "async_make_request") as mock_request: + mock_request.return_value = { + "userPromptAnalysis": {"attackDetected": False}, + "documentsAnalysis": [], + } + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + cache=None, + data={ + "input": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there"}, + {"role": "user", "content": "Tell me about Paris"}, + ] + }, + call_type="completion", + ) + + mock_request.assert_called_once() + assert ( + mock_request.call_args.kwargs["user_prompt"] + == "Tell me about Paris" + ) + + +@pytest.mark.asyncio +async def test_azure_prompt_shield_responses_api_multi_turn_with_tool_output(): + """Test guardrail filters out non-message items (tool outputs) from input list. + + Responses API lists can contain tool-call output items that lack a "role" + key. These must be filtered out so get_user_prompt() can find user messages + that appear before them in the list. + """ + guardrail = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure_prompt_shield", + api_key="test_key", + api_base="test_base", + ) + with patch.object(guardrail, "async_make_request") as mock_request: + mock_request.return_value = { + "userPromptAnalysis": {"attackDetected": False}, + "documentsAnalysis": [], + } + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + cache=None, + data={ + "input": [ + {"role": "user", "content": "Search for weather in Tokyo"}, + {"role": "assistant", "content": "Let me look that up."}, + {"role": "user", "content": "Thanks, also check Paris"}, + # Tool output — no "role" key, should be filtered out + {"type": "tool_result", "tool_use_id": "abc", "content": "72°F"}, + ] + }, + call_type="completion", + ) + + mock_request.assert_called_once() + assert "Paris" in mock_request.call_args.kwargs["user_prompt"] + + +@pytest.mark.asyncio +async def test_azure_prompt_shield_responses_api_no_input(): + """Test guardrail skips gracefully when neither messages nor input is present.""" + guardrail = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure_prompt_shield", + api_key="test_key", + api_base="test_base", + ) + with patch.object(guardrail, "async_make_request") as mock_request: + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + cache=None, + data={"model": "gpt-4o"}, + call_type="completion", + ) + + mock_request.assert_not_called() + assert result == {"model": "gpt-4o"}