From 47ac65fd67c4c4a5c793c8a3aaddb4408a60a1b9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 23:49:33 -0700 Subject: [PATCH] fix(guardrails): scan compact input in vendor guardrails and warn on skipped input guardrails --- litellm/integrations/custom_guardrail.py | 3 +- .../guardrails/guardrail_hooks/akto/akto.py | 5 +- .../guardrails/guardrail_hooks/azure/base.py | 4 +- .../guardrail_hooks/bedrock_guardrails.py | 6 +- .../microsoft_purview/purview_dlp.py | 5 +- .../unified_guardrail/unified_guardrail.py | 39 ++- litellm/types/utils.py | 9 + .../integrations/test_custom_guardrail.py | 15 + .../guardrail_hooks/akto/test_akto.py | 5 +- .../azure/test_azure_prompt_shield.py | 7 +- .../azure/test_azure_text_moderation.py | 5 +- .../guardrail_hooks/test_microsoft_purview.py | 321 +++++------------- .../test_unified_guardrail.py | 26 +- tests/unit/test_utils.py | 29 ++ 14 files changed, 208 insertions(+), 271 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9e2dc902d55..107d1ac7996 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -33,6 +33,7 @@ from litellm.types.guardrails import ( from litellm.types.llms.openai import AllMessageValues from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel from litellm.types.utils import ( + RESPONSES_INPUT_CALL_TYPES, CallTypes, GenericGuardrailAPIInputs, GuardrailStatus, @@ -1632,7 +1633,7 @@ class CustomGuardrail(CustomLogger): # /responses # User/System messages are stored in the "input" key, use litellm transformation to get the messages ######################################################### - if call_type == CallTypes.responses.value or call_type == CallTypes.aresponses.value: + if call_type in RESPONSES_INPUT_CALL_TYPES: from typing import cast from litellm.responses.litellm_completion_transformation.transformation import ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py b/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py index da131931aec..1cd7e922b13 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py +++ b/litellm/proxy/guardrails/guardrail_hooks/akto/akto.py @@ -43,7 +43,7 @@ from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers from litellm.proxy.mcp_server.utils import JSONLeafPath, json_string_leaves from litellm.types.guardrails import GuardrailEventHooks, LitellmParams, Mode from litellm.types.proxy.guardrails.guardrail_hooks.akto import AktoGuardrailConfigModelOptionalParams -from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs +from litellm.types.utils import RESPONSES_INPUT_CALL_TYPES, CallTypes, GenericGuardrailAPIInputs from .akto_attachments import request_attachments, without_attachment_content @@ -81,7 +81,6 @@ DEFAULT_REQUEST_PATH: Final = "/v1/chat/completions" MCP_PATH: Final = "/mcp" MCP_TOOL_PREFIX: Final = "mcp" DEFAULT_BLOCK_REASON: Final = "Blocked by Akto Guardrails" -RESPONSES_API_CALL_TYPES: Final = frozenset((CallTypes.responses.value, CallTypes.aresponses.value)) MESSAGES_API_CALL_TYPES: Final = frozenset((CallTypes.anthropic_messages.value, CallTypes.aanthropic_messages.value)) UNMASKABLE_REASON: Final = "Content masked by Akto guardrail policy could not be applied" MALFORMED_ATTACHMENT_REASON: Final = "Attachment could not be read for the Akto guardrail check" @@ -431,7 +430,7 @@ class AktoGuardrail(CustomGuardrail): api_messages: Final = self.messages_api_messages(request_data) # The Responses API sends "input", so a "messages" key there is a decoy raw_messages: Final = ( - None if call_type_of(request_data) in RESPONSES_API_CALL_TYPES else request_data.get("messages") + None if call_type_of(request_data) in RESPONSES_INPUT_CALL_TYPES else request_data.get("messages") ) messages: Final = ( api_messages diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py index 5f57969c310..37e93ccfad2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py @@ -14,7 +14,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import AllMessageValues, ResponseInputParam -from litellm.types.utils import CallTypes, CallTypesLiteral +from litellm.types.utils import RESPONSES_INPUT_CALL_TYPES, CallTypesLiteral # Azure Content Safety APIs have a 10,000 character limit per request. AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH: Final = 10000 @@ -26,7 +26,7 @@ AZURE_CONTENT_SAFETY_TEXT_RECORD_LENGTH: Final = 1000 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}) +RESPONSES_API_CALL_TYPES: Final = RESPONSES_INPUT_CALL_TYPES _RESPONSES_API_CALL_TYPES: Final = RESPONSES_API_CALL_TYPES diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index adda465ad95..95c963a94d2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -92,6 +92,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.utils import ( + RESPONSES_INPUT_CALL_TYPES, CallTypes, CallTypesLiteral, Choices, @@ -226,9 +227,6 @@ def _redact_assessment_match_fields(assessments: list[dict]) -> list[dict]: return redacted if isinstance(redacted, list) else assessments -_RESPONSES_API_CALL_TYPES: Final = frozenset({CallTypes.responses, CallTypes.aresponses}) - - def _without_image_bytes(content: Sequence[BedrockContentItem]) -> tuple[BedrockContentItem, ...]: return tuple( BedrockContentItem( @@ -255,7 +253,7 @@ def _is_responses_api_route(request_route: str | None) -> bool: if request_route is None: return False call_types: Final = get_call_types_for_route(request_route) - return call_types is not None and any(call_type in _RESPONSES_API_CALL_TYPES for call_type in call_types) + return call_types is not None and any(call_type in RESPONSES_INPUT_CALL_TYPES for call_type in call_types) class _JsonResponse(Protocol): diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py index 72474ba27c3..9cb57dcf547 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py @@ -24,6 +24,7 @@ from litellm.integrations.custom_guardrail import ( ) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( + RESPONSES_INPUT_CALL_TYPES, Choices, GuardrailStatus, ModelResponse, @@ -402,7 +403,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): user_id: Final = self._resolve_user_id_for_blocking(data, user_api_key_dict) prompt_text: str | None = None - if call_type in ("responses", "aresponses"): + if call_type in RESPONSES_INPUT_CALL_TYPES: # Route Responses API calls to the responses-specific extractor # before the generic ``messages`` branch. This mirrors # ``async_logging_hook`` and ensures ``instructions`` (system @@ -663,7 +664,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail): # Log prompt (uploadText) try: prompt_text: str | None = None - if call_type in ("responses", "aresponses"): + if call_type in RESPONSES_INPUT_CALL_TYPES: # Responses API: route to the responses-specific extractor # before the generic ``messages`` branch. litellm's logging # pipeline stores the raw responses ``input`` (a string or a 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 be5a2ccaa03..d0c70f8423a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -110,6 +110,19 @@ def _translated_call_type(call_type: object, mappings: Mapping[CallTypes, object return known_call_type if known_call_type in mappings else None +def _warn_untranslated_call_type( + guardrail_name: str | None, request_route: str | None, call_type: object, stage: str +) -> None: + verbose_proxy_logger.warning( + "Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; " + "skipping %s scanning.", + guardrail_name, + request_route, + call_type, + stage, + ) + + def resolve_endpoint_translation( user_api_key_dict: UserAPIKeyAuth, first_response_item: object | None ) -> "tuple[str, BaseTranslation] | None": @@ -269,7 +282,9 @@ class UnifiedLLMGuardrails(CustomLogger): verbose_proxy_logger.debug("Running UnifiedLLMGuardrails pre-call hook") - guardrail_to_apply: Final[CustomGuardrail] = data.pop("guardrail_to_apply", None) + guardrail_to_apply: Final = cast( # cast-ok: the proxy stores the selected CustomGuardrail under this key + "CustomGuardrail | None", data.pop("guardrail_to_apply", None) + ) if guardrail_to_apply is None: return data @@ -287,6 +302,9 @@ class UnifiedLLMGuardrails(CustomLogger): mappings: Final = load_guardrail_translation_mappings() input_call_type: Final = _translated_call_type(call_type, mappings) if input_call_type is None: + _warn_untranslated_call_type( + guardrail_to_apply.guardrail_name, user_api_key_dict.request_route, call_type, "pre-call" + ) return data endpoint_translation: Final = _as_endpoint_translation(mappings[input_call_type]()) @@ -318,7 +336,9 @@ class UnifiedLLMGuardrails(CustomLogger): """ verbose_proxy_logger.debug("Running UnifiedLLMGuardrails moderation hook") - guardrail_to_apply: Final[CustomGuardrail] = data.pop("guardrail_to_apply", None) + guardrail_to_apply: Final = cast( # cast-ok: the proxy stores the selected CustomGuardrail under this key + "CustomGuardrail | None", data.pop("guardrail_to_apply", None) + ) if guardrail_to_apply is None: return data @@ -336,6 +356,9 @@ class UnifiedLLMGuardrails(CustomLogger): mappings: Final = load_guardrail_translation_mappings() input_call_type: Final = _translated_call_type(call_type, mappings) if input_call_type is None: + _warn_untranslated_call_type( + guardrail_to_apply.guardrail_name, user_api_key_dict.request_route, call_type, "during-call" + ) return data endpoint_translation: Final = _as_endpoint_translation(mappings[input_call_type]()) @@ -369,7 +392,9 @@ class UnifiedLLMGuardrails(CustomLogger): ) from litellm.types.guardrails import GuardrailEventHooks - guardrail_to_apply: Final[CustomGuardrail] = data.pop("guardrail_to_apply", None) + guardrail_to_apply: Final = cast( # cast-ok: the proxy stores the selected CustomGuardrail under this key + "CustomGuardrail | None", data.pop("guardrail_to_apply", None) + ) if guardrail_to_apply is None: return @@ -412,12 +437,8 @@ class UnifiedLLMGuardrails(CustomLogger): mappings: Final = load_guardrail_translation_mappings() if CallTypes(call_type) not in mappings: - verbose_proxy_logger.warning( - "Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; " - "skipping post-call scanning.", - guardrail_to_apply.guardrail_name, - user_api_key_dict.request_route, - call_type, + _warn_untranslated_call_type( + guardrail_to_apply.guardrail_name, user_api_key_dict.request_route, call_type, "post-call" ) return response diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 260b5049e1c..0e2a40431e2 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -797,6 +797,15 @@ MCP_GUARDRAIL_CALL_TYPES: Final[frozenset[str]] = frozenset( {CallTypes.call_mcp_tool.value, CallTypes.list_mcp_tools.value} ) +RESPONSES_INPUT_CALL_TYPES: Final[frozenset[str]] = frozenset( + { + CallTypes.responses.value, + CallTypes.aresponses.value, + CallTypes.compact_responses.value, + CallTypes.acompact_responses.value, + } +) + # Mapping of API routes to their corresponding call types API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = { # Chat Completions diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 614ff97553a..ed48a2488bb 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -3786,3 +3786,18 @@ def test_default_on_guardrail() -> None: ) is True ) + + +@pytest.mark.parametrize( + "call_type", + [CallTypes.responses, CallTypes.aresponses, CallTypes.compact_responses, CallTypes.acompact_responses], +) +def test_guardrail_messages_for_responses_and_compact_come_from_input(call_type: CallTypes) -> None: + guardrail: Final = CustomGuardrail(guardrail_name="test-guardrail") + + messages: Final = guardrail.get_guardrails_messages_for_call_type( + call_type=call_type, + data={"input": "Remember forbiddenword", "messages": [{"role": "user", "content": "decoy"}]}, + ) + + assert messages == [{"role": "user", "content": "Remember forbiddenword"}] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/akto/test_akto.py b/tests/unit/proxy/guardrails/guardrail_hooks/akto/test_akto.py index e2f10225568..638e94ed93c 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/akto/test_akto.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/akto/test_akto.py @@ -1652,11 +1652,12 @@ async def test_a_client_sent_response_key_cannot_skip_recording_or_tool_call_che assert calls["/v1/chat/completions"].get("ingest_data") == "true", "the reply is still recorded" -def test_a_decoy_messages_key_cannot_replace_the_responses_api_input(akto_post_call): +@pytest.mark.parametrize("call_type", ["aresponses", "acompact_responses"]) +def test_a_decoy_messages_key_cannot_replace_the_responses_api_input(akto_post_call, call_type): request_data = { "input": [{"role": "user", "content": "the real prompt"}], "messages": [{"role": "user", "content": "hello"}], - "litellm_logging_obj": SimpleNamespace(call_type="aresponses", model_call_details={}), + "litellm_logging_obj": SimpleNamespace(call_type=call_type, model_call_details={}), } payload = akto_post_call.build_akto_payload( diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py b/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py index 3784ddb4694..a59a20a235e 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py @@ -530,8 +530,11 @@ async def test_prompt_shield_messages_less_embeddings_return_data_and_log_allow( ), ], ) +@pytest.mark.parametrize("call_type", ["aresponses", "acompact_responses"]) @pytest.mark.asyncio -async def test_responses_input_is_scanned_and_billing_is_logged(responses_input: object, expected_prompt: str) -> None: +async def test_responses_input_is_scanned_and_billing_is_logged( + responses_input: object, expected_prompt: str, call_type: str +) -> None: guardrail: Final = _priced_shield_guardrail(cost_tier="paid", price_per_1000_text_records=0.38) data: Final[dict[str, object]] = {"input": responses_input} @@ -540,7 +543,7 @@ async def test_responses_input_is_scanned_and_billing_is_logged(responses_input: user_api_key_dict=UserAPIKeyAuth(api_key="k"), cache=None, data=data, - call_type="aresponses", + call_type=call_type, ) mock_post.assert_called_once() diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py b/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py index c57c54afc73..8ee04008bc9 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py @@ -54,8 +54,9 @@ async def test_azure_text_moderation_guardrail_pre_call_hook(): assert mock_async_make_request.call_args.kwargs["text"] == "Hello, how are you?" +@pytest.mark.parametrize("call_type", ["aresponses", "acompact_responses"]) @pytest.mark.asyncio -async def test_azure_text_moderation_scans_responses_input() -> None: +async def test_azure_text_moderation_scans_responses_input(call_type: str) -> None: guardrail: Final = AzureContentSafetyTextModerationGuardrail( guardrail_name="azure_text_moderation", api_key="azure_text_moderation_api_key", @@ -78,7 +79,7 @@ async def test_azure_text_moderation_scans_responses_input() -> None: user_api_key_dict=UserAPIKeyAuth(api_key="azure_text_moderation_api_key"), cache=None, data={"input": "Review this response input"}, - call_type="aresponses", + call_type=call_type, ) assert exc_info.value.status_code == 400 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py index 25d2556bf26..55aa3ccac3f 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py @@ -17,6 +17,7 @@ from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview.base import ( from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview.purview_dlp import ( MicrosoftPurviewDLPGuardrail, ) +from litellm.types.utils import CallTypesLiteral def _make_guardrail(**kwargs) -> MicrosoftPurviewDLPGuardrail: @@ -201,9 +202,7 @@ class TestResolveUserId: def test_end_user_id_after_key_user_id(self): """When both key user_id and end_user_id exist, key user_id is used first.""" guardrail = _make_guardrail() - auth = UserAPIKeyAuth( - api_key="test", user_id="key-owner", end_user_id="end-user-101" - ) + auth = UserAPIKeyAuth(api_key="test", user_id="key-owner", end_user_id="end-user-101") assert guardrail._resolve_user_id({}, auth) == "key-owner" def test_none_when_missing(self): @@ -222,9 +221,7 @@ class TestPreCallHook: async def test_pre_call_allow(self): guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_pre_call_hook( @@ -247,9 +244,7 @@ class TestPreCallHook: "litellm_call_id": "call-abc", } - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} out = await guardrail.async_pre_call_hook( @@ -265,9 +260,7 @@ class TestPreCallHook: async def test_pre_call_block(self): guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.side_effect = HTTPException( status_code=400, detail={"error": "Microsoft Purview DLP: Content blocked by policy"}, @@ -275,9 +268,7 @@ class TestPreCallHook: with pytest.raises(HTTPException) as exc_info: await guardrail.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth( - api_key="test", user_id="user-123" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), cache=None, data={ "messages": [ @@ -296,9 +287,7 @@ class TestPreCallHook: async def test_pre_call_no_user_id_raises(self): guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: with pytest.raises(HTTPException) as exc_info: await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test"), @@ -314,9 +303,7 @@ class TestPreCallHook: async def test_pre_call_no_messages_skips(self): guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), cache=None, @@ -333,9 +320,7 @@ class TestPreCallFullTranscript: """DLP text must include system / prior turns, not only the last user block.""" guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_pre_call_hook( @@ -371,16 +356,10 @@ class TestPostCallHook: guardrail = _make_guardrail() response = ModelResponse( - choices=[ - Choices( - index=0, message=Message(content="Safe response", role="assistant") - ) - ], + choices=[Choices(index=0, message=Message(content="Safe response", role="assistant"))], ) - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} result = await guardrail.async_post_call_success_hook( @@ -410,9 +389,7 @@ class TestPostCallHook: ], ) - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.side_effect = HTTPException( status_code=400, detail={"error": "Microsoft Purview DLP: Content blocked by policy"}, @@ -421,9 +398,7 @@ class TestPostCallHook: with pytest.raises(HTTPException) as exc_info: await guardrail.async_post_call_success_hook( data={}, - user_api_key_dict=UserAPIKeyAuth( - api_key="test", user_id="user-123" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), response=response, ) @@ -435,14 +410,10 @@ class TestPostCallHook: guardrail = _make_guardrail() response = ModelResponse( - choices=[ - Choices(index=0, message=Message(content="Response", role="assistant")) - ], + choices=[Choices(index=0, message=Message(content="Response", role="assistant"))], ) - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: with pytest.raises(HTTPException) as exc_info: await guardrail.async_post_call_success_hook( data={}, @@ -471,9 +442,7 @@ class TestPostCallHook: ], ) - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_post_call_success_hook( @@ -493,9 +462,7 @@ class TestTextCompletionHooks: async def test_pre_call_text_completion_uses_prompt(self): guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_pre_call_hook( @@ -522,9 +489,7 @@ class TestTextCompletionHooks: ], ) - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_post_call_success_hook( @@ -605,9 +570,7 @@ class TestResponsesAPIHooks: """Pre-call hook must scan plain-string ``input`` on responses call type.""" guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_pre_call_hook( @@ -621,21 +584,20 @@ class TestResponsesAPIHooks: assert mock_check.call_args.kwargs["activity"] == "uploadText" assert "SSN: 123-45-6789" in mock_check.call_args.kwargs["text"] + @pytest.mark.parametrize("call_type", ["aresponses", "acompact_responses"]) @pytest.mark.asyncio - async def test_pre_call_aresponses_string_input(self): + async def test_pre_call_aresponses_string_input(self, call_type: CallTypesLiteral): """Pre-call hook must scan ``input`` on ``aresponses`` call type too.""" guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), cache=None, data={"input": "sensitive content"}, - call_type="aresponses", + call_type=call_type, ) mock_check.assert_called_once() @@ -646,17 +608,13 @@ class TestResponsesAPIHooks: """Pre-call hook must extract text from structured list ``input``.""" guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), cache=None, - data={ - "input": [{"role": "user", "content": "Secret phrase: alpha bravo"}] - }, + data={"input": [{"role": "user", "content": "Secret phrase: alpha bravo"}]}, call_type="responses", ) @@ -668,9 +626,7 @@ class TestResponsesAPIHooks: """Pre-call hook must not call _check_content when ``input`` is absent.""" guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), cache=None, @@ -685,9 +641,7 @@ class TestResponsesAPIHooks: """Benign string ``input`` must still scan ``instructions`` (system message).""" guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_pre_call_hook( @@ -710,9 +664,7 @@ class TestResponsesAPIHooks: """Requests with only ``instructions`` (no ``input``) must still be scanned.""" guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_pre_call_hook( @@ -723,10 +675,7 @@ class TestResponsesAPIHooks: ) mock_check.assert_called_once() - assert ( - "policy text in instructions only" - in mock_check.call_args.kwargs["text"] - ) + assert "policy text in instructions only" in mock_check.call_args.kwargs["text"] @pytest.mark.asyncio async def test_post_call_responses_api_output_text(self): @@ -743,16 +692,12 @@ class TestResponsesAPIHooks: "id": "msg-1", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "card 4111-1111-1111-1111"} - ], + "content": [{"type": "output_text", "text": "card 4111-1111-1111-1111"}], } ], ) - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} result = await guardrail.async_post_call_success_hook( @@ -778,9 +723,7 @@ class TestResponsesAPIHooks: output=[], ) - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: await guardrail.async_post_call_success_hook( data={}, user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), @@ -789,8 +732,9 @@ class TestResponsesAPIHooks: mock_check.assert_not_called() + @pytest.mark.parametrize("call_type", ["responses", "compact_responses"]) @pytest.mark.asyncio - async def test_logging_hook_responses_api_input_and_output(self): + async def test_logging_hook_responses_api_input_and_output(self, call_type: str): """Logging hook must scan both ``input`` and ``ResponsesAPIResponse.output``.""" from litellm.types.llms.openai import ResponsesAPIResponse @@ -809,20 +753,16 @@ class TestResponsesAPIHooks: ], ) - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_logging_hook( kwargs={ "input": "prompt body", - "litellm_params": { - "metadata": {"user_api_key_user_id": "user-123"} - }, + "litellm_params": {"metadata": {"user_api_key_user_id": "user-123"}}, }, result=result_response, - call_type="responses", + call_type=call_type, ) assert mock_check.call_count == 2 @@ -860,9 +800,7 @@ class TestResponsesAPIHooks: ], ) - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_logging_hook( @@ -872,9 +810,7 @@ class TestResponsesAPIHooks: # Simulate litellm's logging path which mirrors the raw # responses input under "messages". "messages": "prompt body", - "litellm_params": { - "metadata": {"user_api_key_user_id": "user-123"} - }, + "litellm_params": {"metadata": {"user_api_key_user_id": "user-123"}}, }, result=result_response, call_type="aresponses", @@ -883,11 +819,7 @@ class TestResponsesAPIHooks: assert mock_check.call_count == 2 activities = {c.kwargs["activity"] for c in mock_check.call_args_list} assert activities == {"uploadText", "downloadText"} - upload_calls = [ - c - for c in mock_check.call_args_list - if c.kwargs["activity"] == "uploadText" - ] + upload_calls = [c for c in mock_check.call_args_list if c.kwargs["activity"] == "uploadText"] assert len(upload_calls) == 1 upload_text = upload_calls[0].kwargs["text"] assert "prompt body" in upload_text @@ -910,10 +842,7 @@ class TestLoggingResolveUserId: } } } - assert ( - guardrail._resolve_user_id_from_logging_kwargs(kwargs) - == "trusted-from-proxy" - ) + assert guardrail._resolve_user_id_from_logging_kwargs(kwargs) == "trusted-from-proxy" def test_logging_ignores_caller_supplied_user_id_field(self): """Caller-controlled ``metadata[user_id_field]`` must not drive Purview audit attribution.""" @@ -928,9 +857,7 @@ class TestLoggingResolveUserId: "user_api_key_user_id": "from-top-level", "litellm_params": {"metadata": {}}, } - assert ( - guardrail._resolve_user_id_from_logging_kwargs(kwargs) == "from-top-level" - ) + assert guardrail._resolve_user_id_from_logging_kwargs(kwargs) == "from-top-level" def test_logging_returns_none_when_no_trusted_identity(self): guardrail = _make_guardrail() @@ -1064,9 +991,7 @@ class TestTokenCaching: async def test_token_cached(self): guardrail = _make_guardrail() - with patch.object( - guardrail.async_handler, "post", return_value=_mock_token_response() - ) as mock_post: + with patch.object(guardrail.async_handler, "post", return_value=_mock_token_response()) as mock_post: token1 = await guardrail._get_access_token() token2 = await guardrail._get_access_token() @@ -1079,9 +1004,7 @@ class TestTokenCaching: async def test_token_refreshed_on_expiry(self): guardrail = _make_guardrail() - with patch.object( - guardrail.async_handler, "post", return_value=_mock_token_response() - ) as mock_post: + with patch.object(guardrail.async_handler, "post", return_value=_mock_token_response()) as mock_post: await guardrail._get_access_token() # Expire the token @@ -1130,9 +1053,7 @@ class TestGraphPostHttpError: guardrail = _make_guardrail() error_resp = Mock() - error_resp.json.return_value = { - "error": {"code": "Forbidden", "message": "no access"} - } + error_resp.json.return_value = {"error": {"code": "Forbidden", "message": "no access"}} error_resp.headers = {} error_resp.raise_for_status = Mock( side_effect=httpx.HTTPStatusError( @@ -1143,9 +1064,7 @@ class TestGraphPostHttpError: ) with ( - patch.object( - guardrail, "_get_access_token", new_callable=AsyncMock - ) as mock_token, + patch.object(guardrail, "_get_access_token", new_callable=AsyncMock) as mock_token, patch.object(guardrail.async_handler, "post", return_value=error_resp), ): mock_token.return_value = "mock-token" @@ -1161,9 +1080,7 @@ class TestGraphPostHttpError: """A Graph error on protectionScopes/compute must not be cached as success.""" guardrail = _make_guardrail() - with patch.object( - guardrail, "_graph_post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail, "_graph_post", new_callable=AsyncMock) as mock_post: mock_post.side_effect = httpx.HTTPStatusError( "429 Too Many Requests", request=httpx.Request("POST", "https://graph.microsoft.com/"), @@ -1187,15 +1104,9 @@ class TestScopeCaching: async def test_scope_cached(self): guardrail = _make_guardrail() - with patch.object( - guardrail, "_graph_post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail, "_graph_post", new_callable=AsyncMock) as mock_post: mock_post.return_value = ( - { - "value": [ - {"activities": "uploadText", "executionMode": "evaluateInline"} - ] - }, + {"value": [{"activities": "uploadText", "executionMode": "evaluateInline"}]}, {"ETag": "scope-etag"}, ) @@ -1213,17 +1124,11 @@ class TestScopeCaching: guardrail._scope_cache_maxsize = 3 scope_payload = ( - { - "value": [ - {"activities": "uploadText", "executionMode": "evaluateInline"} - ] - }, + {"value": [{"activities": "uploadText", "executionMode": "evaluateInline"}]}, {"ETag": "scope-etag"}, ) - with patch.object( - guardrail, "_graph_post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail, "_graph_post", new_callable=AsyncMock) as mock_post: mock_post.return_value = scope_payload await guardrail._compute_protection_scopes("user-a") @@ -1258,9 +1163,7 @@ class TestScopeCaching: {"ETag": "scope-etag"}, ) - with patch.object( - guardrail, "_graph_post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail, "_graph_post", new_callable=AsyncMock) as mock_post: mock_post.return_value = scope_payload # Populate cache: user-a (older), user-b (newer) @@ -1281,18 +1184,14 @@ class TestScopeCaching: assert mock_post.call_count == 4 assert "user-a" in guardrail._scope_cache, "user-a was wrongly evicted" - assert ( - "user-b" not in guardrail._scope_cache - ), "user-b should have been evicted" + assert "user-b" not in guardrail._scope_cache, "user-b should have been evicted" assert "user-c" in guardrail._scope_cache @pytest.mark.asyncio async def test_scope_invalidated_on_modified(self): guardrail = _make_guardrail() - with patch.object( - guardrail, "_graph_post", new_callable=AsyncMock - ) as mock_post: + with patch.object(guardrail, "_graph_post", new_callable=AsyncMock) as mock_post: # First call: compute scopes mock_post.return_value = ( {"value": []}, @@ -1320,9 +1219,7 @@ class TestGetPromptTextForDlp: def test_single_message_no_extra_separator(self): """A single message is returned as-is (no leading/trailing separator).""" guardrail = _make_guardrail() - result = guardrail.get_prompt_text_for_dlp( - [{"role": "user", "content": "Hello"}] - ) + result = guardrail.get_prompt_text_for_dlp([{"role": "user", "content": "Hello"}]) assert result == "Hello" def test_messages_separated_by_double_newline(self): @@ -1542,9 +1439,7 @@ class TestCheckContentApiErrorHandling: assert exc_info.value.status_code == 400 assert isinstance(exc_info.value.detail, dict) - assert "upstream policy evaluation failed" in exc_info.value.detail.get( - "error", "" - ) + assert "upstream policy evaluation failed" in exc_info.value.detail.get("error", "") assert "network failure" in exc_info.value.detail.get("exception", "") assert isinstance(exc_info.value.__cause__, RuntimeError) @@ -1633,9 +1528,7 @@ class TestCheckContentApiErrorHandling: headers={"Retry-After": "30"}, request=httpx.Request("POST", "https://graph.microsoft.com/v1.0/x"), ) - upstream_err = httpx.HTTPStatusError( - "rate limited", request=upstream_resp.request, response=upstream_resp - ) + upstream_err = httpx.HTTPStatusError("rate limited", request=upstream_resp.request, response=upstream_resp) with patch.object( guardrail, @@ -1666,9 +1559,7 @@ class TestCheckContentApiErrorHandling: status_code=401, request=httpx.Request("POST", "https://graph.microsoft.com/v1.0/x"), ) - upstream_err = httpx.HTTPStatusError( - "unauthorized", request=upstream_resp.request, response=upstream_resp - ) + upstream_err = httpx.HTTPStatusError("unauthorized", request=upstream_resp.request, response=upstream_resp) with patch.object( guardrail, @@ -1723,9 +1614,7 @@ class TestAsyncLoggingHookIndependence: await guardrail.async_logging_hook( kwargs={ "messages": [{"role": "user", "content": "prompt"}], - "litellm_params": { - "metadata": {"user_api_key_user_id": "user-123"} - }, + "litellm_params": {"metadata": {"user_api_key_user_id": "user-123"}}, }, result=response, call_type="completion", @@ -1762,9 +1651,7 @@ class TestAsyncLoggingHookIndependence: await guardrail.async_logging_hook( kwargs={ "messages": [{"role": "user", "content": "prompt"}], - "litellm_params": { - "metadata": {"user_api_key_user_id": "user-123"} - }, + "litellm_params": {"metadata": {"user_api_key_user_id": "user-123"}}, }, result=response, call_type="completion", @@ -1924,15 +1811,13 @@ class TestGetPromptTextToolCalls: """async_pre_call_hook must include tool_call arguments in the text sent to Purview.""" guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_pre_call_hook( - user_api_key_dict=__import__( - "litellm.proxy._types", fromlist=["UserAPIKeyAuth"] - ).UserAPIKeyAuth(api_key="test", user_id="user-123"), + user_api_key_dict=__import__("litellm.proxy._types", fromlist=["UserAPIKeyAuth"]).UserAPIKeyAuth( + api_key="test", user_id="user-123" + ), cache=None, data={ "messages": [ @@ -2055,16 +1940,14 @@ class TestCompletionResponseTextPartsToolCalls: ], ) - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_post_call_success_hook( data={}, - user_api_key_dict=__import__( - "litellm.proxy._types", fromlist=["UserAPIKeyAuth"] - ).UserAPIKeyAuth(api_key="test", user_id="user-123"), + user_api_key_dict=__import__("litellm.proxy._types", fromlist=["UserAPIKeyAuth"]).UserAPIKeyAuth( + api_key="test", user_id="user-123" + ), response=response, ) @@ -2152,9 +2035,7 @@ class TestGraphUserIdEncoding: guardrail.async_handler.post = AsyncMock(side_effect=_capture_post) - with patch.object( - guardrail, "_get_access_token", new_callable=AsyncMock - ) as mock_token: + with patch.object(guardrail, "_get_access_token", new_callable=AsyncMock) as mock_token: mock_token.return_value = "tok" await guardrail._compute_protection_scopes("user/with%special") @@ -2194,9 +2075,7 @@ class TestResolveTrustedUserId: def test_trusted_prefers_key_user_id_over_end_user_id(self): guardrail = _make_guardrail() - auth = UserAPIKeyAuth( - api_key="test", user_id="key-owner", end_user_id="end-user" - ) + auth = UserAPIKeyAuth(api_key="test", user_id="key-owner", end_user_id="end-user") assert guardrail._resolve_trusted_user_id({}, auth) == "key-owner" @@ -2285,9 +2164,7 @@ class TestTokenIdPromptHandling: """Pure token-id prompts must be rejected in blocking pre_call mode.""" guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: with pytest.raises(HTTPException) as exc_info: await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"), @@ -2307,9 +2184,7 @@ class TestTokenIdPromptHandling: guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: with caplog.at_level(logging.WARNING): await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"), @@ -2326,9 +2201,7 @@ class TestTokenIdPromptHandling: """Normal string prompts must still be sent to Purview.""" guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: mock_check.return_value = {"policyActions": []} await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"), @@ -2346,9 +2219,7 @@ class TestTokenIdPromptHandling: """Empty/whitespace-only string prompts must not be flagged as token-id prompts.""" guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: data = {"prompt": empty_prompt} result = await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"), @@ -2374,9 +2245,7 @@ class TestTokenIdPromptHandling: """Nested/mixed token-id prompts must also be rejected in blocking pre_call mode.""" guardrail = _make_guardrail() - with patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check: + with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check: with pytest.raises(HTTPException) as exc_info: await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"), @@ -2435,12 +2304,8 @@ class TestStreamingIteratorHook: with ( patch("litellm.main.stream_chunk_builder", return_value=assembled_response), - patch( - "litellm.llms.base_llm.base_model_iterator.MockResponseIterator" - ) as mock_iterator_cls, - patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check, + patch("litellm.llms.base_llm.base_model_iterator.MockResponseIterator") as mock_iterator_cls, + patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check, ): mock_check.return_value = {"policyActions": []} @@ -2491,18 +2356,15 @@ class TestStreamingIteratorHook: new_callable=AsyncMock, side_effect=HTTPException( status_code=400, - detail={ - "error": "Microsoft Purview DLP: Content blocked by policy" - }, + detail={"error": "Microsoft Purview DLP: Content blocked by policy"}, ), ), ): chunks = [] + async def _drain(): async for chunk in guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=UserAPIKeyAuth( - api_key="test", user_id="user-123" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), response=fake_response_stream(), request_data={"metadata": {"user_id": "user-123"}}, ): @@ -2533,10 +2395,9 @@ class TestStreamingIteratorHook: async def fake_response_stream(): yield assembled_response - with patch( - "litellm.main.stream_chunk_builder", return_value=assembled_response - ): + with patch("litellm.main.stream_chunk_builder", return_value=assembled_response): chunks = [] + async def _drain(): async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test"), # no user_id @@ -2568,9 +2429,7 @@ class TestStreamingIteratorHook: with ( patch("litellm.main.stream_chunk_builder", return_value=assembled_response), - patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check, + patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check, ): mock_check.return_value = {"policyActions": []} @@ -2626,9 +2485,7 @@ class TestStreamingIteratorHook: with ( patch("litellm.main.stream_chunk_builder") as mock_stream_builder, - patch.object( - guardrail, "_check_content", new_callable=AsyncMock - ) as mock_check, + patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check, ): mock_check.return_value = {"policyActions": []} @@ -2688,11 +2545,10 @@ class TestStreamingIteratorHook: ), ): chunks = [] + async def _drain(): async for chunk in guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=UserAPIKeyAuth( - api_key="test", user_id="user-123" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), response=fake_response_stream(), request_data={}, ): @@ -2719,7 +2575,4 @@ class TestRegistration: assert "microsoft_purview" in guardrail_initializer_registry assert "microsoft_purview" in guardrail_class_registry - assert ( - guardrail_class_registry["microsoft_purview"] - is MicrosoftPurviewDLPGuardrail - ) + assert guardrail_class_registry["microsoft_purview"] is MicrosoftPurviewDLPGuardrail diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 279c10153d3..e2dbe4eff18 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -3128,22 +3128,28 @@ async def test_compact_input_is_guardrailed_like_create_while_stored_response_lo @pytest.mark.parametrize("route_call_type", ["avector_store_list", "not_a_call_type"]) @pytest.mark.asyncio async def test_a_route_without_a_call_type_passes_input_guardrails_unscanned( - monkeypatch: pytest.MonkeyPatch, hook: str, route_call_type: str + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, hook: str, route_call_type: str ) -> None: _patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings()) guardrail: Final = RecordingGuardrail() data: Final = {"model": "gpt-4o", "input": "Remember forbiddenword", "guardrail_to_apply": guardrail} - user_api_key_dict: Final = UserAPIKeyAuth(token="hashed-key") + user_api_key_dict: Final = UserAPIKeyAuth(token="hashed-key", request_route="/v1/vector_stores") - result: Final = ( - await UnifiedLLMGuardrails().async_pre_call_hook( - user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type=route_call_type + with caplog.at_level(logging.WARNING): + result: Final = ( + await UnifiedLLMGuardrails().async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type=route_call_type + ) + if hook == "async_pre_call_hook" + else await UnifiedLLMGuardrails().async_moderation_hook( + data=data, user_api_key_dict=user_api_key_dict, call_type=route_call_type + ) ) - if hook == "async_pre_call_hook" - else await UnifiedLLMGuardrails().async_moderation_hook( - data=data, user_api_key_dict=user_api_key_dict, call_type=route_call_type - ) - ) assert guardrail.apply_calls == [] assert result == {"model": "gpt-4o", "input": "Remember forbiddenword"} + stage: Final = "pre-call" if hook == "async_pre_call_hook" else "during-call" + assert ( + f"route '/v1/vector_stores' but call type '{route_call_type}' has no guardrail translation handler; " + f"skipping {stage} scanning" in caplog.text + ) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 2adf9f45f0b..4781ade8f54 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -55,6 +55,7 @@ from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.proxy.utils import is_valid_api_key from litellm.responses.main import aresponses, responses from litellm.types.caching import CachingSupportedCallTypes +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY, STREAM_OPTIONS_STASH_KEYS from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import CredentialLiteLLMParams, GenericLiteLLMParams, LiteLLM_Params @@ -10920,6 +10921,34 @@ async def test_client_passes_stored_response_call_types_to_both_deployment_hooks assert callback.hook_calls == (("pre", CallTypes(call_type)), ("post", CallTypes(call_type))) +@pytest.mark.parametrize("call_type", ["aget_responses", "acompact_responses"]) +@pytest.mark.asyncio +async def test_a_deployment_post_call_guardrail_blocks_a_stored_response_reply( + call_type: str, monkeypatch: pytest.MonkeyPatch +) -> None: + class NoPineappleGuardrail(CustomGuardrail): + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: object, response: object + ) -> object: + if "pineapple" in str(response): + raise ValueError("pineapple is not allowed") + return response + + monkeypatch.setattr( + litellm, + "callbacks", + [NoPineappleGuardrail(guardrail_name="no-pineapple", event_hook=GuardrailEventHooks.post_call)], + ) + + async def stored_response_call(**kwargs: object) -> dict[str, str]: + return {"id": "resp_abc", "output_text": "I like pineapple"} + + stored_response_call.__name__ = call_type + + with pytest.raises(ValueError, match="pineapple is not allowed"): + await client(stored_response_call)(model="gpt-4o", response_id="resp_abc", guardrails=["no-pineapple"]) + + @pytest.mark.parametrize( "call_type", [