fix(guardrails): scan Responses API input in Azure Prompt Shield

litellm/proxy/guardrails/guardrail_hooks/azure/base.py

@@ -1,5 +1,6 @@
 import re
-from typing import TYPE_CHECKING, Any, Final
+from collections.abc import Mapping
+from typing import Any, Final, cast
 
 from litellm._logging import verbose_proxy_logger
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
@@ -9,9 +10,9 @@
     get_async_httpx_client,
     httpxSpecialProvider,
 )
-
-if TYPE_CHECKING:
-    from litellm.types.llms.openai import AllMessageValues
+from litellm.responses.utils import ResponsesAPIRequestUtils
+from litellm.types.llms.openai import AllMessageValues, ResponseInputParam
+from litellm.types.utils import CallTypes, CallTypesLiteral
 
 # Azure Content Safety APIs have a 10,000 character limit per request.
 AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH: Final = 10000
@@ -23,6 +24,8 @@
 AZURE_CONTENT_SAFETY_DEFAULT_API_VERSION: Final = "2024-09-01"
 JAVELIN_API_VERSION_STORED_BY_OLDER_RELEASES: Final = "v1"
 
+_RESPONSES_API_CALL_TYPES: Final = frozenset({CallTypes.responses, CallTypes.aresponses})
+
 
 def resolve_content_safety_api_version(configured: str | None) -> str:
     if not configured or configured == JAVELIN_API_VERSION_STORED_BY_OLDER_RELEASES:
@@ -131,7 +134,20 @@ def split_text_by_words(text: str, max_length: int) -> list[str]:
 
         return chunks
 
-    def get_user_prompt(self, messages: list["AllMessageValues"]) -> str | None:
+    def get_user_prompt_from_request(self, data: Mapping[str, object], call_type: CallTypesLiteral) -> str | None:
+        if call_type in _RESPONSES_API_CALL_TYPES:
+            responses_input: Final = data.get("input")
+            if not isinstance(responses_input, (str, list)):
+                return None
+            validated_input: Final = cast(ResponseInputParam, responses_input)  # cast-ok: narrowed to str | list
+            return get_last_user_message(ResponsesAPIRequestUtils.responses_input_to_chat_messages(validated_input))
+
+        messages: Final = data.get("messages")
+        if not isinstance(messages, list):
+            return None
+        return get_last_user_message(cast(list[AllMessageValues], messages))  # cast-ok: narrowed to list
+
+    def get_user_prompt(self, messages: list[AllMessageValues]) -> str | None:
         """
         Get the last consecutive block of messages from the user.
 

litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py

@@ -33,7 +33,6 @@
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
     from litellm.proxy._types import UserAPIKeyAuth
     from litellm.types.guardrails import LitellmParams
-    from litellm.types.llms.openai import AllMessageValues
     from litellm.types.proxy.guardrails.guardrail_hooks.azure.azure_prompt_shield import (
         AzurePromptShieldGuardrailResponse,
     )
@@ -250,11 +249,7 @@ async def async_pre_call_hook(
             "Azure Prompt Shield: Running pre-call prompt scan, on call_type: %s",
             call_type,
         )
-        new_messages: Final[list[AllMessageValues] | None] = data.get("messages")
-        if new_messages is None:
-            verbose_proxy_logger.warning("Azure Prompt Shield: not running guardrail. No messages in data")
-            return data
-        user_prompt: Final = self.get_user_prompt(new_messages)
+        user_prompt: Final = self.get_user_prompt_from_request(data, call_type)
 
         if user_prompt:
             verbose_proxy_logger.debug("Azure Prompt Shield: User prompt: %s", user_prompt)