diff --git a/litellm/constants.py b/litellm/constants.py index 715d57e594d..764ad1eeb1e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -264,6 +264,9 @@ MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", # Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth LITELLM_ASYNCIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)) TOOL_POLICY_CACHE_TTL_SECONDS = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60)) +GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS = int( + os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60) +) # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger. # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire. MAX_SIZE_IN_MEMORY_QUEUE = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8))) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 856556f7c56..031eeeac9c8 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,3 +1,5 @@ +import hashlib +import json import os import secrets from datetime import datetime @@ -46,7 +48,10 @@ if TYPE_CHECKING: dc = DualCache() -from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY +from litellm.constants import ( + GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, + PRE_CALL_EXECUTED_GUARDRAILS_KEY, +) from litellm.exceptions import ( BlockedPiiEntityError, GuardrailRaisedException, @@ -113,6 +118,7 @@ class CustomGuardrail(CustomLogger): on_sensitive_data: Optional[str] = None, sensitive_data_route_to_model: Optional[str] = None, sticky_session_routing: bool = True, + only_scan_new_messages: bool = False, **kwargs, ): """ @@ -145,6 +151,7 @@ class CustomGuardrail(CustomLogger): self.on_sensitive_data: Optional[str] = on_sensitive_data self.sensitive_data_route_to_model: Optional[str] = sensitive_data_route_to_model self.sticky_session_routing: bool = sticky_session_routing + self.only_scan_new_messages: bool = only_scan_new_messages if supported_event_hooks: ## validate event_hook is in supported_event_hooks @@ -269,6 +276,100 @@ class CustomGuardrail(CustomLogger): """Extract session_id from request data.""" return get_session_id_from_request_data(request_data) + @staticmethod + def _message_content_hash(message: AllMessageValues) -> str: + """Stable content hash for a single message. + + Serializes the whole message (role + content + tool calls) so that an + edited earlier message produces a different hash and gets re-scanned. + """ + serialized = json.dumps(message, sort_keys=True, default=str) + return hashlib.sha256(serialized.encode("utf-8")).hexdigest() + + def _scanned_messages_cache_key(self, session_id: str) -> str: + return f"guardrail_scanned_messages:{self.guardrail_name}:{session_id}" + + async def filter_new_messages_for_session( + self, + messages: list[AllMessageValues] | None, + request_data: dict[str, object], + cache: DualCache, + ) -> list[AllMessageValues] | None: + """Return only the messages not already scanned earlier in this session. + + Returns ``None`` when incremental scanning is inactive (feature off, no + session id, masking enabled, or the cache read failed). ``None`` signals + the caller to fall back to a full scan; a returned list (possibly empty) + signals the caller to scan only that subset and skip masking write-back. + """ + if not self.only_scan_new_messages or not messages: + return None + + if self.mask_request_content or self.mask_response_content: + verbose_logger.warning( + "Guardrail %s: only_scan_new_messages is not supported with masking; scanning full context.", + self.guardrail_name, + ) + return None + + session_id = get_session_id_from_request_data(request_data) + if not session_id: + verbose_logger.debug( + "Guardrail %s: only_scan_new_messages enabled but request has no session id; scanning full context.", + self.guardrail_name, + ) + return None + + try: + cached: object = await cache.async_get_cache(key=self._scanned_messages_cache_key(session_id)) + except Exception as e: # noqa: BLE001 # cache is best-effort; any failure must fall back to a full scan + verbose_logger.warning( + "Guardrail %s: failed to read scanned-message cache (%s); scanning full context.", + self.guardrail_name, + e, + ) + return None + + seen: set[str] = {str(h) for h in cached} if isinstance(cached, list) else set() + return [message for message in messages if self._message_content_hash(message) not in seen] + + async def mark_messages_scanned( + self, + messages: list[AllMessageValues] | None, + request_data: dict[str, object], + cache: DualCache, + ) -> None: + """Record the hashes of all messages present on a successful (non-blocked) scan. + + Called only after the guardrail allows the request, so a blocked message is + never marked scanned and will be re-checked if the client retries. + """ + if not self.only_scan_new_messages or not messages: + return + if self.mask_request_content or self.mask_response_content: + return + session_id = get_session_id_from_request_data(request_data) + if not session_id: + return + + cache_key = self._scanned_messages_cache_key(session_id) + current_hashes = [self._message_content_hash(message) for message in messages] + try: + existing: object = await cache.async_get_cache(key=cache_key) + existing_hashes: list[str] = [str(h) for h in existing] if isinstance(existing, list) else [] + merged: list[str] = list(dict.fromkeys(existing_hashes + current_hashes)) + await cache.async_set_cache( + key=cache_key, + value=merged, + ttl=GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, + ) + except Exception as e: # noqa: BLE001 # cache is best-effort; any failure must not block the request + verbose_logger.warning( + "Guardrail %s: failed to persist scanned-message cache (%s); next call will re-scan.", + self.guardrail_name, + e, + ) + def should_route_on_sensitive_data(self) -> bool: """ Returns True if this guardrail is configured to route requests diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 44ae57f81db..405e284d9ad 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -1172,7 +1172,20 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): verbose_proxy_logger.debug("No messages found for call_type, skipping guardrail") return data - filter_result = self._prepare_guardrail_messages_for_role(messages=new_messages) + scan_messages = await self.filter_new_messages_for_session( + messages=new_messages, + request_data=data, + cache=cache, + ) + incremental_scan = scan_messages is not None + if incremental_scan and not scan_messages: + verbose_proxy_logger.debug("Bedrock Guardrail: no new messages to scan for this session, skipping API call") + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + return data + + messages_to_scan = scan_messages if incremental_scan else new_messages + + filter_result = self._prepare_guardrail_messages_for_role(messages=messages_to_scan) filtered_messages = filter_result.payload_messages if not filtered_messages: @@ -1196,15 +1209,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 2. Update the messages with the guardrail response ########## ######################################################### - updated_subset = self._update_messages_with_updated_bedrock_guardrail_response( - messages=filtered_messages, - bedrock_guardrail_response=bedrock_guardrail_response, - ) - data["messages"] = self._merge_filtered_messages( - original_messages=filter_result.original_messages or new_messages, - updated_target_messages=updated_subset, - target_indices=filter_result.target_indices, - ) + if not incremental_scan: + updated_subset = self._update_messages_with_updated_bedrock_guardrail_response( + messages=filtered_messages, + bedrock_guardrail_response=bedrock_guardrail_response, + ) + data["messages"] = self._merge_filtered_messages( + original_messages=filter_result.original_messages or new_messages, + updated_target_messages=updated_subset, + target_indices=filter_result.target_indices, + ) + + await self.mark_messages_scanned(messages=new_messages, request_data=data, cache=cache) ######################################################### ########## 3. Add the guardrail to the applied guardrails header ########## diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index e8abf66a6f7..8bf2fbb7b78 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -31,6 +31,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): aws_sts_endpoint=litellm_params.aws_sts_endpoint, aws_bedrock_runtime_endpoint=litellm_params.aws_bedrock_runtime_endpoint, experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only, + only_scan_new_messages=litellm_params.only_scan_new_messages or False, ) litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback) return _bedrock_callback diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index c7e080b1363..cfe15d716d1 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -594,6 +594,18 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up description="When True, guardrails only receive the latest message for the relevant role (e.g., newest user input pre-call, newest assistant output post-call)", ) + only_scan_new_messages: Optional[bool] = Field( + default=False, + description=( + "When True, the guardrail only scans messages that have not already been scanned " + "earlier in the same session (identified by litellm_session_id / session_id). " + "Message content is hashed per session and cached; only the diff (new or edited " + "messages) is sent to the guardrail provider on follow-up calls. Falls back to a " + "full scan when the request has no session id or the cache is unavailable. Intended " + "for blocking/detection guardrails; not applied when mask_request_content is set." + ), + ) + skip_system_message_in_guardrail: Optional[bool] = Field( default=None, description=( diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 9289dece83f..7fe56fb777a 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -8,7 +8,6 @@ from litellm.types.utils import GuardrailTracingDetail class TestCustomGuardrailDeploymentHook: - @pytest.mark.asyncio async def test_async_pre_call_deployment_hook_no_guardrails(self): """Test that method returns kwargs unchanged when no guardrails are present""" @@ -21,18 +20,14 @@ class TestCustomGuardrailDeploymentHook: "guardrails": None, } - result = await custom_guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert result == kwargs # Test with guardrails as non-list kwargs["guardrails"] = "not_a_list" - result = await custom_guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert result == kwargs @@ -59,9 +54,7 @@ class TestCustomGuardrailDeploymentHook: "user_api_key_request_route": "test_route", } - result = await custom_guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) # Verify async_pre_call_hook was called with correct parameters custom_guardrail.async_pre_call_hook.assert_called_once() @@ -94,9 +87,7 @@ class TestCustomGuardrailDeploymentHook: super().__init__(guardrail_name="g1", default_on=True) self.pre_call_count = 0 - async def async_pre_call_hook( - self, user_api_key_dict, cache, data, call_type - ): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): self.pre_call_count += 1 return data @@ -109,9 +100,7 @@ class TestCustomGuardrailDeploymentHook: } guardrail.mark_pre_call_hook_ran(kwargs) - await guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert guardrail.pre_call_count == 0 @@ -125,9 +114,7 @@ class TestCustomGuardrailDeploymentHook: super().__init__(guardrail_name="g1", default_on=True) self.pre_call_count = 0 - async def async_pre_call_hook( - self, user_api_key_dict, cache, data, call_type - ): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): self.pre_call_count += 1 return data @@ -139,9 +126,7 @@ class TestCustomGuardrailDeploymentHook: "metadata": {}, } - await guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert guardrail.pre_call_count == 1 @@ -170,9 +155,7 @@ class TestCustomGuardrailDeploymentHook: super().__init__(guardrail_name="g1", default_on=True) self.pre_call_count = 0 - async def async_pre_call_hook( - self, user_api_key_dict, cache, data, call_type - ): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): self.pre_call_count += 1 return data @@ -184,15 +167,12 @@ class TestCustomGuardrailDeploymentHook: "metadata": {PRE_CALL_EXECUTED_GUARDRAILS_KEY: ["g1"]}, } - await guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert guardrail.pre_call_count == 1 class TestCustomGuardrailShouldRunGuardrail: - def test_should_run_guardrail_with_litellm_metadata(self): """Test that should_run_guardrail works with litellm_metadata pattern""" from litellm.types.guardrails import GuardrailEventHooks @@ -209,9 +189,7 @@ class TestCustomGuardrailShouldRunGuardrail: "litellm_metadata": {"guardrails": ["test_guardrail"]}, } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True @@ -231,9 +209,7 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"guardrails": ["test_guardrail"]}, } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True @@ -250,9 +226,7 @@ class TestCustomGuardrailShouldRunGuardrail: # Test with guardrails at root level data = {"model": "gpt-3.5-turbo", "guardrails": ["test_guardrail"]} - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True @@ -272,9 +246,7 @@ class TestCustomGuardrailShouldRunGuardrail: "litellm_metadata": {"guardrails": ["different_guardrail"]}, } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is False @@ -293,9 +265,7 @@ class TestCustomGuardrailShouldRunGuardrail: "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "test"}], } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True, "Global guardrail should run when default_on=True" # Test 2: User-injected disable at root level is IGNORED @@ -307,9 +277,7 @@ class TestCustomGuardrailShouldRunGuardrail: result = custom_guardrail.should_run_guardrail( data=data_with_disable_root, event_type=GuardrailEventHooks.pre_call ) - assert ( - result is True - ), "User-injected disable_global_guardrails should be ignored" + assert result is True, "User-injected disable_global_guardrails should be ignored" # Test 3: User-injected disable in metadata is IGNORED data_with_disable_metadata = { @@ -340,12 +308,8 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}, "litellm_metadata": {"request_tags": ["user-supplied"]}, } - result = custom_guardrail.should_run_guardrail( - data=data_cross_key, event_type=GuardrailEventHooks.pre_call - ) - assert ( - result is False - ), "Admin config in metadata must not be shadowed by user-supplied litellm_metadata" + result = custom_guardrail.should_run_guardrail(data=data_cross_key, event_type=GuardrailEventHooks.pre_call) + assert result is False, "Admin config in metadata must not be shadowed by user-supplied litellm_metadata" # Test 6: After the pre-call strip runs, user-injected # user_api_key_metadata in the non-authoritative metadata key is gone. @@ -356,12 +320,8 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}, "litellm_metadata": {}, # post-strip: attacker payload removed } - result = custom_guardrail.should_run_guardrail( - data=data_post_strip, event_type=GuardrailEventHooks.pre_call - ) - assert ( - result is False - ), "Admin config in metadata must be respected when other metadata key is empty" + result = custom_guardrail.should_run_guardrail(data=data_post_strip, event_type=GuardrailEventHooks.pre_call) + assert result is False, "Admin config in metadata must be respected when other metadata key is empty" def test_should_run_guardrail_key_disable_global_not_overruled_by_team_guardrail_list( self, @@ -427,12 +387,7 @@ class TestCustomGuardrailShouldRunGuardrail: "messages": [{"role": "user", "content": "test"}], "opted_out_global_guardrails": ["global_guardrail"], } - assert ( - custom_guardrail.should_run_guardrail( - data=data_root, event_type=GuardrailEventHooks.pre_call - ) - is True - ) + assert custom_guardrail.should_run_guardrail(data=data_root, event_type=GuardrailEventHooks.pre_call) is True # Test 2: User-injected opt-out in metadata is IGNORED data_metadata = { @@ -441,10 +396,7 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"opted_out_global_guardrails": ["global_guardrail"]}, } assert ( - custom_guardrail.should_run_guardrail( - data=data_metadata, event_type=GuardrailEventHooks.pre_call - ) - is True + custom_guardrail.should_run_guardrail(data=data_metadata, event_type=GuardrailEventHooks.pre_call) is True ) # Test 4: a different guardrail in the opt-out list → still runs @@ -453,12 +405,7 @@ class TestCustomGuardrailShouldRunGuardrail: "messages": [{"role": "user", "content": "test"}], "metadata": {"opted_out_global_guardrails": ["some_other_guardrail"]}, } - assert ( - custom_guardrail.should_run_guardrail( - data=data_other, event_type=GuardrailEventHooks.pre_call - ) - is True - ) + assert custom_guardrail.should_run_guardrail(data=data_other, event_type=GuardrailEventHooks.pre_call) is True # Test 5: empty opt-out list → still runs data_empty = { @@ -466,12 +413,7 @@ class TestCustomGuardrailShouldRunGuardrail: "messages": [{"role": "user", "content": "test"}], "metadata": {"opted_out_global_guardrails": []}, } - assert ( - custom_guardrail.should_run_guardrail( - data=data_empty, event_type=GuardrailEventHooks.pre_call - ) - is True - ) + assert custom_guardrail.should_run_guardrail(data=data_empty, event_type=GuardrailEventHooks.pre_call) is True # Test 6: malformed value (bool instead of list) → safely ignored, guardrail runs data_malformed = { @@ -480,10 +422,7 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"opted_out_global_guardrails": True}, } assert ( - custom_guardrail.should_run_guardrail( - data=data_malformed, event_type=GuardrailEventHooks.pre_call - ) - is True + custom_guardrail.should_run_guardrail(data=data_malformed, event_type=GuardrailEventHooks.pre_call) is True ) def test_should_run_guardrail_opt_out_does_not_affect_non_global(self): @@ -506,12 +445,7 @@ class TestCustomGuardrailShouldRunGuardrail: "guardrails": ["opt_in_guardrail"], }, } - assert ( - non_global.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) - is True - ) + assert non_global.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is True class TestApplyGuardrailCheck: @@ -550,35 +484,33 @@ class TestApplyGuardrailCheck: child_with_override = ChildGuardrailWithOverride() # Test: CustomGuardrail itself has apply_guardrail in its __dict__ - assert ( - "apply_guardrail" in type(CustomGuardrail()).__dict__ - ), "CustomGuardrail should have apply_guardrail in its own __dict__" + assert "apply_guardrail" in type(CustomGuardrail()).__dict__, ( + "CustomGuardrail should have apply_guardrail in its own __dict__" + ) # Test: ParentGuardrail inherits but doesn't override, so it should NOT be in __dict__ - assert ( - "apply_guardrail" not in type(parent_instance).__dict__ - ), "ParentGuardrail should NOT have apply_guardrail in its own __dict__ (only inherited)" + assert "apply_guardrail" not in type(parent_instance).__dict__, ( + "ParentGuardrail should NOT have apply_guardrail in its own __dict__ (only inherited)" + ) # Test: ChildGuardrailWithoutOverride only inherits, should NOT be in __dict__ - assert ( - "apply_guardrail" not in type(child_without_override).__dict__ - ), "ChildGuardrailWithoutOverride should NOT have apply_guardrail in its own __dict__ (only inherited)" + assert "apply_guardrail" not in type(child_without_override).__dict__, ( + "ChildGuardrailWithoutOverride should NOT have apply_guardrail in its own __dict__ (only inherited)" + ) # Test: ChildGuardrailWithOverride overrides the method, SHOULD be in __dict__ - assert ( - "apply_guardrail" in type(child_with_override).__dict__ - ), "ChildGuardrailWithOverride SHOULD have apply_guardrail in its own __dict__ (overridden)" + assert "apply_guardrail" in type(child_with_override).__dict__, ( + "ChildGuardrailWithOverride SHOULD have apply_guardrail in its own __dict__ (overridden)" + ) # Verify that all instances still have the method via inheritance (hasattr) - assert hasattr( - parent_instance, "apply_guardrail" - ), "All instances should have apply_guardrail via inheritance" - assert hasattr( - child_without_override, "apply_guardrail" - ), "All instances should have apply_guardrail via inheritance" - assert hasattr( - child_with_override, "apply_guardrail" - ), "All instances should have apply_guardrail via inheritance" + assert hasattr(parent_instance, "apply_guardrail"), "All instances should have apply_guardrail via inheritance" + assert hasattr(child_without_override, "apply_guardrail"), ( + "All instances should have apply_guardrail via inheritance" + ) + assert hasattr(child_with_override, "apply_guardrail"), ( + "All instances should have apply_guardrail via inheritance" + ) class TestGuardrailLoggingAggregation: @@ -605,11 +537,7 @@ class TestGuardrailLoggingAggregation: def test_appends_to_existing_metadata_list(self): request_data = { - "metadata": { - "standard_logging_guardrail_information": [ - {"guardrail_name": "existing_guardrail"} - ] - } + "metadata": {"standard_logging_guardrail_information": [{"guardrail_name": "existing_guardrail"}]} } self._invoke_add_log(request_data) @@ -621,11 +549,7 @@ class TestGuardrailLoggingAggregation: assert info[1]["guardrail_name"] == "test_guardrail" def test_converts_existing_metadata_dict_to_list(self): - request_data = { - "metadata": { - "standard_logging_guardrail_information": {"guardrail_name": "legacy"} - } - } + request_data = {"metadata": {"standard_logging_guardrail_information": {"guardrail_name": "legacy"}}} self._invoke_add_log(request_data) @@ -637,18 +561,12 @@ class TestGuardrailLoggingAggregation: def test_appends_to_litellm_metadata(self): request_data = { - "litellm_metadata": { - "standard_logging_guardrail_information": [ - {"guardrail_name": "litellm_existing"} - ] - } + "litellm_metadata": {"standard_logging_guardrail_information": [{"guardrail_name": "litellm_existing"}]} } self._invoke_add_log(request_data) - info = request_data["litellm_metadata"][ - "standard_logging_guardrail_information" - ] + info = request_data["litellm_metadata"]["standard_logging_guardrail_information"] assert isinstance(info, list) assert len(info) == 2 assert info[1]["guardrail_name"] == "test_guardrail" @@ -689,9 +607,7 @@ class TestGuardrailOtelSpanEmission: assert len(captured) == 1 emitted = captured[0] - recorded = request_data["metadata"]["standard_logging_guardrail_information"][ - -1 - ] + recorded = request_data["metadata"]["standard_logging_guardrail_information"][-1] assert emitted is recorded assert emitted["guardrail_name"] == "emit_guard" assert emitted["start_time"] == 1.0 @@ -701,9 +617,7 @@ class TestGuardrailOtelSpanEmission: def _boom(_entry): raise RuntimeError("otel exporter down") - monkeypatch.setattr( - "litellm.integrations.otel.logger.emit_guardrail_span", _boom - ) + monkeypatch.setattr("litellm.integrations.otel.logger.emit_guardrail_span", _boom) request_data = {"metadata": {}} self._record(self._make_guardrail(), request_data) @@ -800,9 +714,7 @@ class TestGuardrailSensitiveFieldStripping: duration=1.0, ) - logged_response = request_data["metadata"][ - "standard_logging_guardrail_information" - ][0]["guardrail_response"] + logged_response = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"] assert "secret_fields" not in logged_response assert "sk-live-SHOULD-NOT-APPEAR" not in json.dumps(logged_response) @@ -815,9 +727,7 @@ class TestGuardrailSensitiveFieldStripping: guardrail_json_response=[ { "result": "ok", - "secret_fields": { - "raw_headers": {"authorization": "Bearer sk-secret"} - }, + "secret_fields": {"raw_headers": {"authorization": "Bearer sk-secret"}}, }, {"result": "also_ok"}, ], @@ -871,9 +781,7 @@ class TestGuardrailResponseCredentialMasking: duration=1.0, ) - logged = request_data["metadata"]["standard_logging_guardrail_information"][0][ - "guardrail_response" - ] + logged = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"] masked_key = logged["metadata_snapshot"]["callback_vars"]["langsmith_api_key"] assert masked_key != plaintext_key @@ -882,10 +790,7 @@ class TestGuardrailResponseCredentialMasking: assert logged["model"] == "gpt-4o-mini" assert logged["messages"] == [{"role": "user", "content": "hi"}] - assert ( - logged["metadata_snapshot"]["callback_vars"]["langsmith_project"] - == "proj-name" - ) + assert logged["metadata_snapshot"]["callback_vars"]["langsmith_project"] == "proj-name" def test_nested_user_api_key_auth_metadata_is_masked(self): import json @@ -944,9 +849,7 @@ class TestGuardrailResponseCredentialMasking: request_data: dict = {"metadata": {}} guardrail.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response={ - "filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}] - }, + guardrail_json_response={"filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}]}, request_data=request_data, guardrail_status="success", ) @@ -969,9 +872,7 @@ class TestGuardrailResponseCredentialMasking: guardrail_status="success", ) - logged = request_data["metadata"]["standard_logging_guardrail_information"][0][ - "guardrail_response" - ] + logged = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"] assert logged["flagged"] is True assert logged["score"] == 0.94 assert logged["tokens_used"] == 42 @@ -983,18 +884,14 @@ class TestGuardrailResponseCredentialMasking: plaintext = "lsv2_pt_abcdef1234567890" guardrail.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response={ - "metadata_snapshot": { - "callback_vars": {"langsmith_api_key": plaintext} - } - }, + guardrail_json_response={"metadata_snapshot": {"callback_vars": {"langsmith_api_key": plaintext}}}, request_data=request_data, guardrail_status="success", ) - masked = request_data["metadata"]["standard_logging_guardrail_information"][0][ - "guardrail_response" - ]["metadata_snapshot"]["callback_vars"]["langsmith_api_key"] + masked = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"][ + "metadata_snapshot" + ]["callback_vars"]["langsmith_api_key"] assert masked != plaintext assert masked.startswith(plaintext[:4]) assert masked.endswith(plaintext[-4:]) @@ -1342,9 +1239,7 @@ class TestEventTypeLogging: guardrail = TestGuardrail() request_data = {"metadata": {}} - await guardrail.apply_guardrail( - inputs={"texts": ["x"]}, request_data=request_data - ) + await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data) logged_info = request_data["metadata"]["standard_logging_guardrail_information"] assert len(logged_info) == 1, ( @@ -1386,9 +1281,7 @@ class TestEventTypeLogging: request_data = {"metadata": {}} with pytest.raises(ValueError, match="blocked"): - await guardrail.apply_guardrail( - inputs={"texts": ["x"]}, request_data=request_data - ) + await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data) logged_info = request_data["metadata"]["standard_logging_guardrail_information"] assert len(logged_info) == 1 @@ -1485,9 +1378,7 @@ class TestTracingFieldsPopulation: guardrail_json_response="blocked", request_data=request_data, guardrail_status="guardrail_intervened", - tracing_detail=GuardrailTracingDetail( - policy_template="EU AI Act Article 5" - ), + tracing_detail=GuardrailTracingDetail(policy_template="EU AI Act Article 5"), ) slg_list = request_data["metadata"]["standard_logging_guardrail_information"] @@ -1529,13 +1420,7 @@ class TestCustomGuardrailSpendLogMatchRedaction: cg = CustomGuardrail(guardrail_name="test-rail") raw = { "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "match": "GG", "action": "BLOCKED"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "match": "GG", "action": "BLOCKED"}]}} ] } request_data: dict = {"metadata": {}} @@ -1546,17 +1431,10 @@ class TestCustomGuardrailSpendLogMatchRedaction: ) slg = request_data["metadata"]["standard_logging_guardrail_information"][0] assert ( - slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] + slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" ) - assert ( - raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ - "match" - ] - == "GG" - ) + assert raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "GG" def test_add_standard_logging_redacts_regex_field(self): cg = CustomGuardrail(guardrail_name="test-rail") @@ -1716,3 +1594,130 @@ class TestApplyGuardrailStyleDeploymentDispatch: await guardrail.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) assert guardrail.apply_called is False + + +class TestOnlyScanNewMessages: + """Incremental guardrail scanning: only send messages not already scanned this session.""" + + def _guardrail(self, **overrides): + params = dict(guardrail_name="test-guard", only_scan_new_messages=True) + params.update(overrides) + return CustomGuardrail(**params) + + def _cache(self): + from litellm.caching import DualCache + + return DualCache() + + @pytest.mark.asyncio + async def test_disabled_returns_none(self): + guardrail = self._guardrail(only_scan_new_messages=False) + result = await guardrail.filter_new_messages_for_session( + messages=[{"role": "user", "content": "hi"}], + request_data={"litellm_session_id": "s1"}, + cache=self._cache(), + ) + assert result is None + + @pytest.mark.asyncio + async def test_no_session_id_fails_safe_to_full_scan(self): + guardrail = self._guardrail() + result = await guardrail.filter_new_messages_for_session( + messages=[{"role": "user", "content": "hi"}], + request_data={"metadata": {}}, + cache=self._cache(), + ) + assert result is None + + @pytest.mark.asyncio + async def test_masking_guardrail_not_supported(self): + guardrail = self._guardrail(mask_request_content=True) + result = await guardrail.filter_new_messages_for_session( + messages=[{"role": "user", "content": "hi"}], + request_data={"litellm_session_id": "s1"}, + cache=self._cache(), + ) + assert result is None + + @pytest.mark.asyncio + async def test_cache_read_failure_fails_safe_to_full_scan(self): + from unittest.mock import AsyncMock + + guardrail = self._guardrail() + cache = self._cache() + cache.async_get_cache = AsyncMock(side_effect=RuntimeError("redis down")) + result = await guardrail.filter_new_messages_for_session( + messages=[{"role": "user", "content": "hi"}], + request_data={"litellm_session_id": "s1"}, + cache=cache, + ) + assert result is None + + @pytest.mark.asyncio + async def test_dedupes_previously_scanned_messages(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-dedupe"} + turn1 = [ + {"role": "system", "content": "you are helpful"}, + {"role": "user", "content": "first question"}, + ] + + first = await guardrail.filter_new_messages_for_session(messages=turn1, request_data=request, cache=cache) + assert first == turn1 + await guardrail.mark_messages_scanned(messages=turn1, request_data=request, cache=cache) + + turn2 = turn1 + [ + {"role": "assistant", "content": "an answer"}, + {"role": "user", "content": "second question"}, + ] + second = await guardrail.filter_new_messages_for_session(messages=turn2, request_data=request, cache=cache) + assert second == [ + {"role": "assistant", "content": "an answer"}, + {"role": "user", "content": "second question"}, + ] + + @pytest.mark.asyncio + async def test_no_new_messages_returns_empty(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-empty"} + messages = [{"role": "user", "content": "only message"}] + + await guardrail.filter_new_messages_for_session(messages=messages, request_data=request, cache=cache) + await guardrail.mark_messages_scanned(messages=messages, request_data=request, cache=cache) + + again = await guardrail.filter_new_messages_for_session(messages=messages, request_data=request, cache=cache) + assert again == [] + + @pytest.mark.asyncio + async def test_modified_earlier_message_is_rescanned(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-edit"} + original = [{"role": "user", "content": "original"}] + + await guardrail.filter_new_messages_for_session(messages=original, request_data=request, cache=cache) + await guardrail.mark_messages_scanned(messages=original, request_data=request, cache=cache) + + edited = [{"role": "user", "content": "original EDITED"}] + result = await guardrail.filter_new_messages_for_session(messages=edited, request_data=request, cache=cache) + assert result == edited + + @pytest.mark.asyncio + async def test_tool_calls_are_hashed_and_deduped(self): + guardrail = self._guardrail() + cache = self._cache() + request = {"litellm_session_id": "sess-tools"} + tool_message = { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}], + } + messages = [{"role": "user", "content": "call the tool"}, tool_message] + + await guardrail.filter_new_messages_for_session(messages=messages, request_data=request, cache=cache) + await guardrail.mark_messages_scanned(messages=messages, request_data=request, cache=cache) + + result = await guardrail.filter_new_messages_for_session(messages=messages, request_data=request, cache=cache) + assert result == [] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 15827b80bcf..75b1cf7a98a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -53,9 +53,7 @@ async def test__redact_pii_matches_function(): redacted_response = _redact_pii_matches(response_with_pii) # Verify that PII matches are redacted - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] assert pii_entities[0]["match"] == "[REDACTED]", "Name should be redacted" assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted" @@ -173,12 +171,8 @@ async def test__redact_pii_matches_multiple_assessments(): redacted_response = _redact_pii_matches(response_multiple_assessments) # Verify all PII in all assessments are redacted - assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] - assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"][ - "piiEntities" - ] + assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"]["piiEntities"] assert assessment1_pii[0]["match"] == "[REDACTED]", "Email should be redacted" assert assessment2_pii[0]["match"] == "[REDACTED]", "Credit card should be redacted" @@ -199,9 +193,7 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the Bedrock API response with PII mock_bedrock_response = MagicMock() @@ -239,20 +231,11 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): # Mock AWS-related methods to ensure test runs without external dependencies with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch( - "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug" - ) as mock_debug, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ) as mock_load_creds, - patch.object( - guardrail, "_prepare_request", return_value=MagicMock() - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch("litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug") as mock_debug, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")) as mock_load_creds, + patch.object(guardrail, "_prepare_request", return_value=MagicMock()) as mock_prepare_request, ): - mock_post.return_value = mock_bedrock_response # Call the method that should log the redacted response @@ -275,37 +258,23 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): bedrock_response_log_call = call break - assert ( - bedrock_response_log_call is not None - ), "Should have logged Bedrock AI response" + assert bedrock_response_log_call is not None, "Should have logged Bedrock AI response" # Extract the logged response data - logged_response = bedrock_response_log_call[0][ - 1 - ] # Second argument to debug call + logged_response = bedrock_response_log_call[0][1] # Second argument to debug call # Verify that the logged response has redacted PII assert ( - logged_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] - == "[REDACTED]" + logged_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" ) # Verify other fields are preserved assert logged_response["action"] == "GUARDRAIL_INTERVENED" - assert ( - logged_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["type"] - == "PHONE" - ) + assert logged_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["type"] == "PHONE" slg_list = request_data["metadata"]["standard_logging_guardrail_information"] assert ( - slg_list[0]["guardrail_response"]["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"][0]["match"] + slg_list[0]["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" ) @@ -319,9 +288,7 @@ async def test_bedrock_guardrail_original_response_not_modified(): # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the Bedrock API response with PII original_response_data = { @@ -361,17 +328,10 @@ async def test_bedrock_guardrail_original_response_not_modified(): # Mock AWS-related methods to ensure test runs without external dependencies with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ) as mock_load_creds, - patch.object( - guardrail, "_prepare_request", return_value=MagicMock() - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")) as mock_load_creds, + patch.object(guardrail, "_prepare_request", return_value=MagicMock()) as mock_prepare_request, ): - mock_post.return_value = mock_bedrock_response # Call the method @@ -385,19 +345,12 @@ async def test_bedrock_guardrail_original_response_not_modified(): # (The json() method should return the original data) original_data = mock_bedrock_response.json() assert ( - original_data["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] + original_data["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "+1 412 555 1212" ) # Verify that the returned BedrockGuardrailResponse contains original data - assert ( - result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ - "match" - ] - == "+1 412 555 1212" - ) + assert result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "+1 412 555 1212" print("Original response not modified test passed") @@ -454,18 +407,14 @@ async def test__redact_pii_matches_preserves_non_pii_entities(): redacted_response = _redact_pii_matches(response_with_mixed_data) # Verify that PII entity matches are redacted - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] assert pii_entities[0]["match"] == "[REDACTED]", "PII match should be redacted" assert pii_entities[0]["type"] == "EMAIL", "PII type should be preserved" assert pii_entities[0]["action"] == "ANONYMIZED", "PII action should be preserved" assert pii_entities[0]["confidence"] == "HIGH", "PII confidence should be preserved" # Verify that regex matches are also redacted (updated behavior) - regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "regexes" - ] + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted" assert regexes[0]["name"] == "custom_pattern", "Regex name should be preserved" assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved" @@ -496,9 +445,7 @@ async def test_pii_redaction_matches_debug_output_format(): "assessments": [ { "invocationMetrics": { - "guardrailCoverage": { - "textCharacters": {"guarded": 84, "total": 84} - }, + "guardrailCoverage": {"textCharacters": {"guarded": 84, "total": 84}}, "guardrailProcessingLatency": 322, "usage": { "contentPolicyImageUnits": 0, @@ -553,9 +500,7 @@ async def test_pii_redaction_matches_debug_output_format(): redacted_response = _redact_pii_matches(original_response) # Verify the redacted response matches your expected debug output - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] # All PII matches should be redacted assert pii_entities[0]["match"] == "[REDACTED]", "NAME should be redacted" @@ -570,34 +515,19 @@ async def test_pii_redaction_matches_debug_output_format(): assert pii_entities[0]["detected"] == True # Verify that the original response is unchanged - original_pii_entities = original_response["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"] - assert ( - original_pii_entities[0]["match"] == "John Smith" - ), "Original should be unchanged" - assert ( - original_pii_entities[1]["match"] == "324-12-3212" - ), "Original should be unchanged" - assert ( - original_pii_entities[2]["match"] == "607-456-7890" - ), "Original should be unchanged" + original_pii_entities = original_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assert original_pii_entities[0]["match"] == "John Smith", "Original should be unchanged" + assert original_pii_entities[1]["match"] == "324-12-3212", "Original should be unchanged" + assert original_pii_entities[2]["match"] == "607-456-7890", "Original should be unchanged" # Verify all other metadata is preserved in redacted response assert redacted_response["action"] == "GUARDRAIL_INTERVENED" assert redacted_response["actionReason"] == "Guardrail blocked." assert redacted_response["blockedResponse"] == "Input blocked by PII policy" - assert ( - redacted_response["assessments"][0]["invocationMetrics"][ - "guardrailProcessingLatency" - ] - == 322 - ) + assert redacted_response["assessments"][0]["invocationMetrics"]["guardrailProcessingLatency"] == 322 print("PII redaction matches debug output format test passed") - print( - f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}" - ) + print(f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}") print(f"Redacted PII values: {[e['match'] for e in pii_entities]}") @@ -632,14 +562,10 @@ async def test__redact_pii_matches_with_regex_matches(): redacted_response = _redact_pii_matches(response_with_regex) # Verify that regex matches are redacted - regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "regexes" - ] + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert regexes[0]["match"] == "[REDACTED]", "SSN regex match should be redacted" - assert ( - regexes[1]["match"] == "[REDACTED]" - ), "Credit card regex match should be redacted" + assert regexes[1]["match"] == "[REDACTED]", "Credit card regex match should be redacted" # Verify other fields are preserved assert regexes[0]["name"] == "SSN_PATTERN", "Regex name should be preserved" @@ -648,13 +574,9 @@ async def test__redact_pii_matches_with_regex_matches(): assert regexes[1]["action"] == "ANONYMIZED", "Regex action should be preserved" # Verify original response is unchanged - original_regexes = response_with_regex["assessments"][0][ - "sensitiveInformationPolicy" - ]["regexes"] + original_regexes = response_with_regex["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert original_regexes[0]["match"] == "123-45-6789", "Original should be unchanged" - assert ( - original_regexes[1]["match"] == "4111-1111-1111-1111" - ), "Original should be unchanged" + assert original_regexes[1]["match"] == "4111-1111-1111-1111", "Original should be unchanged" print("Regex matches redaction test passed") @@ -690,31 +612,17 @@ async def test__redact_pii_matches_with_custom_words(): # Verify that custom word matches are redacted custom_words = redacted_response["assessments"][0]["wordPolicy"]["customWords"] - assert ( - custom_words[0]["match"] == "[REDACTED]" - ), "First custom word match should be redacted" - assert ( - custom_words[1]["match"] == "[REDACTED]" - ), "Second custom word match should be redacted" + assert custom_words[0]["match"] == "[REDACTED]", "First custom word match should be redacted" + assert custom_words[1]["match"] == "[REDACTED]", "Second custom word match should be redacted" # Verify other fields are preserved - assert ( - custom_words[0]["action"] == "BLOCKED" - ), "Custom word action should be preserved" - assert ( - custom_words[1]["action"] == "ANONYMIZED" - ), "Custom word action should be preserved" + assert custom_words[0]["action"] == "BLOCKED", "Custom word action should be preserved" + assert custom_words[1]["action"] == "ANONYMIZED", "Custom word action should be preserved" # Verify original response is unchanged - original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"][ - "customWords" - ] - assert ( - original_custom_words[0]["match"] == "confidential_data" - ), "Original should be unchanged" - assert ( - original_custom_words[1]["match"] == "secret_information" - ), "Original should be unchanged" + original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"]["customWords"] + assert original_custom_words[0]["match"] == "confidential_data", "Original should be unchanged" + assert original_custom_words[1]["match"] == "secret_information", "Original should be unchanged" print("Custom words redaction test passed") @@ -750,41 +658,21 @@ async def test__redact_pii_matches_with_managed_words(): redacted_response = _redact_pii_matches(response_with_managed_words) # Verify that managed word matches are redacted - managed_words = redacted_response["assessments"][0]["wordPolicy"][ - "managedWordLists" - ] + managed_words = redacted_response["assessments"][0]["wordPolicy"]["managedWordLists"] - assert ( - managed_words[0]["match"] == "[REDACTED]" - ), "First managed word match should be redacted" - assert ( - managed_words[1]["match"] == "[REDACTED]" - ), "Second managed word match should be redacted" + assert managed_words[0]["match"] == "[REDACTED]", "First managed word match should be redacted" + assert managed_words[1]["match"] == "[REDACTED]", "Second managed word match should be redacted" # Verify other fields are preserved - assert ( - managed_words[0]["action"] == "BLOCKED" - ), "Managed word action should be preserved" - assert ( - managed_words[0]["type"] == "PROFANITY" - ), "Managed word type should be preserved" - assert ( - managed_words[1]["action"] == "ANONYMIZED" - ), "Managed word action should be preserved" - assert ( - managed_words[1]["type"] == "HATE_SPEECH" - ), "Managed word type should be preserved" + assert managed_words[0]["action"] == "BLOCKED", "Managed word action should be preserved" + assert managed_words[0]["type"] == "PROFANITY", "Managed word type should be preserved" + assert managed_words[1]["action"] == "ANONYMIZED", "Managed word action should be preserved" + assert managed_words[1]["type"] == "HATE_SPEECH", "Managed word type should be preserved" # Verify original response is unchanged - original_managed_words = response_with_managed_words["assessments"][0][ - "wordPolicy" - ]["managedWordLists"] - assert ( - original_managed_words[0]["match"] == "inappropriate_word" - ), "Original should be unchanged" - assert ( - original_managed_words[1]["match"] == "offensive_term" - ), "Original should be unchanged" + original_managed_words = response_with_managed_words["assessments"][0]["wordPolicy"]["managedWordLists"] + assert original_managed_words[0]["match"] == "inappropriate_word", "Original should be unchanged" + assert original_managed_words[1]["match"] == "offensive_term", "Original should be unchanged" print("Managed words redaction test passed") @@ -841,9 +729,7 @@ async def test__redact_pii_matches_comprehensive_coverage(): # PII entities pii_entities = assessment["sensitiveInformationPolicy"]["piiEntities"] - assert ( - pii_entities[0]["match"] == "[REDACTED]" - ), "PII entity match should be redacted" + assert pii_entities[0]["match"] == "[REDACTED]", "PII entity match should be redacted" # Regex matches regexes = assessment["sensitiveInformationPolicy"]["regexes"] @@ -851,15 +737,11 @@ async def test__redact_pii_matches_comprehensive_coverage(): # Custom words custom_words = assessment["wordPolicy"]["customWords"] - assert ( - custom_words[0]["match"] == "[REDACTED]" - ), "Custom word match should be redacted" + assert custom_words[0]["match"] == "[REDACTED]", "Custom word match should be redacted" # Managed words managed_words = assessment["wordPolicy"]["managedWordLists"] - assert ( - managed_words[0]["match"] == "[REDACTED]" - ), "Managed word match should be redacted" + assert managed_words[0]["match"] == "[REDACTED]", "Managed word match should be redacted" # Verify all other fields are preserved assert pii_entities[0]["type"] == "EMAIL" @@ -868,21 +750,10 @@ async def test__redact_pii_matches_comprehensive_coverage(): # Verify original response is unchanged original_assessment = comprehensive_response["assessments"][0] - assert ( - original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"] - == "user@example.com" - ) - assert ( - original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"] - == "555-123-4567" - ) - assert ( - original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential" - ) - assert ( - original_assessment["wordPolicy"]["managedWordLists"][0]["match"] - == "inappropriate" - ) + assert original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "user@example.com" + assert original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"] == "555-123-4567" + assert original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential" + assert original_assessment["wordPolicy"]["managedWordLists"][0]["match"] == "inappropriate" print("Comprehensive coverage redaction test passed") @@ -914,9 +785,7 @@ async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): aws_region_name = "us-east-1" # Mock the _load_credentials method to avoid actual AWS credential loading - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -926,10 +795,12 @@ async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): ) # Verify that the custom endpoint is used in the URL - expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected URL to contain custom endpoint. Got: {prepped_request.url}" + expected_url = ( + f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, ( + f"Expected URL to contain custom endpoint. Got: {prepped_request.url}" + ) print(f"Custom runtime endpoint test passed. URL: {prepped_request.url}") @@ -944,9 +815,7 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", custom_endpoint) # Create guardrail without explicit aws_bedrock_runtime_endpoint - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock credentials mock_credentials = MagicMock() @@ -960,9 +829,7 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): aws_region_name = "us-east-1" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -972,10 +839,10 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): ) # Verify that the custom endpoint from environment is used in the URL - expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected URL to contain env endpoint. Got: {prepped_request.url}" + expected_url = ( + f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, f"Expected URL to contain env endpoint. Got: {prepped_request.url}" print(f"Environment runtime endpoint test passed. URL: {prepped_request.url}") @@ -988,9 +855,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) # Create guardrail without any custom endpoint - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock credentials mock_credentials = MagicMock() @@ -1004,9 +869,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey aws_region_name = "us-west-2" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -1017,9 +880,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey # Verify that the default endpoint is used expected_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected default URL. Got: {prepped_request.url}" + assert prepped_request.url == expected_url, f"Expected default URL. Got: {prepped_request.url}" print(f"Default endpoint test passed. URL: {prepped_request.url}") @@ -1057,9 +918,7 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch aws_region_name = "us-east-1" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -1069,10 +928,12 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch ) # Verify that the parameter takes precedence over environment variable - expected_url = f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" + expected_url = ( + f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, ( + f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" + ) print(f"Parameter precedence test passed. URL: {prepped_request.url}") @@ -1081,14 +942,10 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): """Test that apply_guardrail handles response with tool_calls (no text content) without calling Bedrock API""" # Create a BedrockGuardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the make_bedrock_api_request method - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api_request: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api_request: # Test the apply_guardrail method with tool_calls in response inputs = { "texts": [], @@ -1115,14 +972,9 @@ async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): assert guardrailed_inputs is not None assert "tool_calls" in guardrailed_inputs assert len(guardrailed_inputs["tool_calls"]) == 1 - assert ( - guardrailed_inputs["tool_calls"][0]["id"] == "call_eFSCWFsyL7MclHYnzKrcQnMK" - ) + assert guardrailed_inputs["tool_calls"][0]["id"] == "call_eFSCWFsyL7MclHYnzKrcQnMK" assert guardrailed_inputs["tool_calls"][0]["function"]["name"] == "get_weather" - assert ( - guardrailed_inputs["tool_calls"][0]["function"]["arguments"] - == '{"location":"São Paulo"}' - ) + assert guardrailed_inputs["tool_calls"][0]["function"]["arguments"] == '{"location":"São Paulo"}' # Verify that the Bedrock API was NOT called since there's no text to process mock_api_request.assert_not_called() print("✅ apply_guardrail with tool_calls test passed - no API call made") @@ -1136,14 +988,10 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source(): policies (e.g. PII on model output) then returned action=NONE for non-streaming completions that go through unified_guardrail -> process_output_response. """ - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") bedrock_none = {"action": "NONE", "output": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = bedrock_none await guardrail.apply_guardrail( @@ -1168,14 +1016,10 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source(): @pytest.mark.asyncio async def test_bedrock_apply_guardrail_request_uses_INPUT_source(): """input_type='request' must call Bedrock with source=INPUT and user messages.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") bedrock_none = {"action": "NONE", "output": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = bedrock_none await guardrail.apply_guardrail( @@ -1258,12 +1102,8 @@ async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): # Mock AWS-related methods with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = mock_bedrock_response @@ -1431,9 +1271,7 @@ class TestShouldRaiseGuardrailBlockedExceptionNullSafety: """Tests for _should_raise_guardrail_blocked_exception handling of null list fields.""" def _create_guardrail(self) -> BedrockGuardrail: - return BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + return BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") @pytest.mark.asyncio async def test_should_handle_all_null_policy_sub_lists(self): @@ -1554,9 +1392,7 @@ class TestShouldRaiseGuardrailBlockedExceptionNullSafety: { "sensitiveInformationPolicy": { "piiEntities": None, - "regexes": [ - {"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"} - ], + "regexes": [{"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"}], }, } ], @@ -1611,18 +1447,14 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_should_handle_none_texts_in_inputs(self): """inputs[\"texts\"] is explicitly None — should not crash.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") inputs = {"texts": None} # Explicit None mock_credentials = MagicMock() with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, patch.object( guardrail, "_load_credentials", @@ -1645,18 +1477,14 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_should_handle_missing_texts_key(self): """inputs has no \"texts\" key at all — should not crash.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") inputs = {} # No "texts" key mock_credentials = MagicMock() with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, patch.object( guardrail, "_load_credentials", @@ -1677,9 +1505,7 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): """Test that BLOCKED actions raise exceptions but ANONYMIZED actions do not""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Test 1: ANONYMIZED action should NOT raise exception anonymized_response = { @@ -1700,9 +1526,7 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): ], } - should_raise = guardrail._should_raise_guardrail_blocked_exception( - anonymized_response - ) + should_raise = guardrail._should_raise_guardrail_blocked_exception(anonymized_response) assert should_raise is False, "ANONYMIZED actions should not raise exceptions" # Test 2: BLOCKED action should raise exception @@ -1710,13 +1534,7 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "I can't provide that information."}], "assessments": [ - { - "topicPolicy": { - "topics": [ - {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} - ] - } - } + {"topicPolicy": {"topics": [{"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"}]}} ], } @@ -1738,19 +1556,13 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): } ] }, - "topicPolicy": { - "topics": [ - {"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"} - ] - }, + "topicPolicy": {"topics": [{"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"}]}, } ], } should_raise = guardrail._should_raise_guardrail_blocked_exception(mixed_response) - assert ( - should_raise is True - ), "Mixed actions with any BLOCKED should raise exceptions" + assert should_raise is True, "Mixed actions with any BLOCKED should raise exceptions" # Test 4: NONE action should not raise exception none_response = { @@ -1782,9 +1594,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): When logging_event_type is set, it must be forwarded to standard guardrail logging. When omitted, INPUT maps to pre_call (legacy). """ - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "test-access-key" mock_credentials.secret_key = "test-secret-key" @@ -1795,13 +1605,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): mock_bedrock_response.json.return_value = { "action": "NONE", "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "match": "GG", "action": "BLOCKED"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "match": "GG", "action": "BLOCKED"}]}} ], } @@ -1811,12 +1615,8 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): } with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -1831,15 +1631,13 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): request_data=request_data, logging_event_type=GuardrailEventHooks.during_call, ) - assert ( - mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.during_call - ) + assert mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.during_call # Raw Bedrock JSON is forwarded; redaction runs once in # CustomGuardrail.add_standard_logging_guardrail_information_to_request_data. assert ( - mock_log.call_args.kwargs["guardrail_json_response"]["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"][0]["match"] + mock_log.call_args.kwargs["guardrail_json_response"]["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["match"] == "GG" ) @@ -1855,9 +1653,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): @pytest.mark.asyncio async def test_make_bedrock_api_request_filters_dynamic_evaluation_overrides(): - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "test-access-key" mock_credentials.secret_key = "test-secret-key" @@ -1873,15 +1669,9 @@ async def test_make_bedrock_api_request_filters_dynamic_evaluation_overrides(): prepared_request.headers = {} with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), - patch.object( - guardrail, "_prepare_request", return_value=prepared_request - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=prepared_request) as mock_prepare_request, patch.object( guardrail, "get_guardrail_dynamic_request_body_params", @@ -1933,9 +1723,7 @@ async def test_during_call_hook_invokes_bedrock_async_moderation_hook(): "model": "gpt-4", "messages": [{"role": "user", "content": "test"}], }, - user_api_key_dict=UserAPIKeyAuth( - api_key="test_key", user_id="test_user" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), call_type="completion", ) finally: @@ -1991,11 +1779,7 @@ def test_extract_blocked_assessments_multiple_policies(): "action": "GUARDRAIL_INTERVENED", "assessments": [ { - "topicPolicy": { - "topics": [ - {"name": "Investment", "type": "DENY", "action": "BLOCKED"} - ] - }, + "topicPolicy": {"topics": [{"name": "Investment", "type": "DENY", "action": "BLOCKED"}]}, "contentPolicy": { "filters": [ { @@ -2006,9 +1790,7 @@ def test_extract_blocked_assessments_multiple_policies(): } ] }, - "wordPolicy": { - "customWords": [{"match": "forbidden", "action": "BLOCKED"}] - }, + "wordPolicy": {"customWords": [{"match": "forbidden", "action": "BLOCKED"}]}, } ], } @@ -2023,13 +1805,7 @@ def test_extract_blocked_assessments_only_anonymized_returns_empty(): response = { "action": "GUARDRAIL_INTERVENED", "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "ANONYMIZED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED", "match": "Jack"}]}} ], } assert g._extract_blocked_assessments(response) == [] @@ -2049,23 +1825,14 @@ def test_get_http_exception_includes_assessments_and_identifier(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "Sorry, the model cannot answer this question."}], "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "BLOCKED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "BLOCKED", "match": "Jack"}]}} ], } exc = g._get_http_exception_for_blocked_guardrail(response) assert isinstance(exc, HTTPException) assert exc.status_code == 400 assert exc.detail["error"] == "Violated guardrail policy" - assert ( - exc.detail["bedrock_guardrail_response"] - == "Sorry, the model cannot answer this question." - ) + assert exc.detail["bedrock_guardrail_response"] == "Sorry, the model cannot answer this question." assert exc.detail["guardrailIdentifier"] == "amgllac6xf3r" assert exc.detail["guardrailVersion"] == "1" assert exc.detail["assessments"][0]["policy"] == "sensitiveInformationPolicy" @@ -2088,15 +1855,11 @@ def test_extract_violation_category_names_mixed_policies(): {"name": "Tax Advice", "action": "BLOCKED"}, ] }, - "contentPolicy": { - "filters": [{"type": "VIOLENCE", "action": "BLOCKED"}] - }, + "contentPolicy": {"filters": [{"type": "VIOLENCE", "action": "BLOCKED"}]}, "wordPolicy": { "managedWordLists": [{"type": "PROFANITY", "action": "BLOCKED"}], }, - "sensitiveInformationPolicy": { - "piiEntities": [{"type": "EMAIL", "action": "BLOCKED"}] - }, + "sensitiveInformationPolicy": {"piiEntities": [{"type": "EMAIL", "action": "BLOCKED"}]}, } ], } @@ -2120,13 +1883,9 @@ def test_extract_violation_category_names_does_not_leak_user_input(): "assessments": [ { "wordPolicy": { - "customWords": [ - {"match": "secret-codeword-abc-123", "action": "BLOCKED"} - ], - }, - "sensitiveInformationPolicy": { - "regexes": [{"match": "4111-1111-1111-1111", "action": "BLOCKED"}] + "customWords": [{"match": "secret-codeword-abc-123", "action": "BLOCKED"}], }, + "sensitiveInformationPolicy": {"regexes": [{"match": "4111-1111-1111-1111", "action": "BLOCKED"}]}, } ], } @@ -2166,13 +1925,7 @@ def test_extract_violation_category_names_skips_anonymized(): g = _make_guardrail() response = { "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [{"type": "NAME", "action": "ANONYMIZED"}] - } - } - ], + "assessments": [{"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED"}]}}], } assert g._extract_violation_category_names(response) == [] @@ -2190,9 +1943,7 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): ``tracing_detail`` so downstream loggers (OTEL, ...) can surface the raw provider verdict as a queryable attribute without re-parsing the redacted guardrail_response blob.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "k" mock_credentials.secret_key = "s" @@ -2202,13 +1953,7 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "topicPolicy": { - "topics": [{"name": "Fiduciary Advice", "action": "BLOCKED"}] - } - } - ], + "assessments": [{"topicPolicy": {"topics": [{"name": "Fiduciary Advice", "action": "BLOCKED"}]}}], } request_data = { @@ -2217,12 +1962,8 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): } with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -2253,9 +1994,7 @@ async def test_make_bedrock_api_request_omits_guardrail_action_when_missing(): """If the Bedrock response omits ``action`` (older / partial payloads), the field must be left off ``tracing_detail`` rather than written as ``None`` — downstream code expects strings or absence, not nulls.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "k" mock_credentials.secret_key = "s" @@ -2266,12 +2005,8 @@ async def test_make_bedrock_api_request_omits_guardrail_action_when_missing(): mock_bedrock_response.json.return_value = {"assessments": []} with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -2300,13 +2035,7 @@ def test_get_http_exception_no_blocked_assessments_omits_field(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "blocked"}], "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "ANONYMIZED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED", "match": "Jack"}]}} ], } exc = g._get_http_exception_for_blocked_guardrail(response) @@ -2370,9 +2099,7 @@ async def test_streaming_post_call_only_runs_output_scan(): yield c minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: out = [] async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth(), @@ -2382,18 +2109,11 @@ async def test_streaming_post_call_only_runs_output_scan(): out.append(chunk) assert len(out) >= 1 - output_calls = [ - c for c in mock_make.call_args_list if c.kwargs.get("source") == "OUTPUT" - ] + output_calls = [c for c in mock_make.call_args_list if c.kwargs.get("source") == "OUTPUT"] assert len(output_calls) == 1 assert output_calls[0].kwargs.get("request_data") is request_data - assert ( - output_calls[0].kwargs.get("logging_event_type") - == GuardrailEventHooks.post_call - ) - input_calls = [ - c for c in mock_make.call_args_list if c.kwargs.get("source") == "INPUT" - ] + assert output_calls[0].kwargs.get("logging_event_type") == GuardrailEventHooks.post_call + input_calls = [c for c in mock_make.call_args_list if c.kwargs.get("source") == "INPUT"] assert len(input_calls) == 0 @@ -2432,9 +2152,7 @@ async def test_streaming_post_call_output_only_path_passes_request_data_to_make_ yield c minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: async for _ in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth(), response=mock_stream(), @@ -2488,9 +2206,7 @@ async def test_post_call_success_hook_only_runs_output_scan(): ) minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: await guardrail.async_post_call_success_hook( data=request_data, user_api_key_dict=UserAPIKeyAuth(), @@ -2499,10 +2215,7 @@ async def test_post_call_success_hook_only_runs_output_scan(): sources = [c.kwargs.get("source") for c in mock_make.call_args_list] assert sources == ["OUTPUT"] - assert ( - mock_make.call_args.kwargs.get("logging_event_type") - == GuardrailEventHooks.post_call - ) + assert mock_make.call_args.kwargs.get("logging_event_type") == GuardrailEventHooks.post_call # --------------------------------------------------------------------------- @@ -2522,9 +2235,7 @@ _GROUNDING_RESPONSE_TEXT = "The capital of Japan is Tokyo." def _grounding_guardrail() -> BedrockGuardrail: - return BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + return BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") def _grounding_messages() -> list: @@ -2556,27 +2267,19 @@ def _model_response(content: str) -> ModelResponse: # Expected OUTPUT content blocks, keyed by their grounding qualifier, so the # per-test assertions read as the block sequence they expect. -_GROUNDING_SOURCE_BLOCK = { - "text": {"text": _GROUNDING_SOURCE_TEXT, "qualifiers": ["grounding_source"]} -} +_GROUNDING_SOURCE_BLOCK = {"text": {"text": _GROUNDING_SOURCE_TEXT, "qualifiers": ["grounding_source"]}} _QUERY_BLOCK = {"text": {"text": _GROUNDING_QUERY_TEXT, "qualifiers": ["query"]}} -_GUARD_BLOCK = { - "text": {"text": _GROUNDING_RESPONSE_TEXT, "qualifiers": ["guard_content"]} -} +_GUARD_BLOCK = {"text": {"text": _GROUNDING_RESPONSE_TEXT, "qualifiers": ["guard_content"]}} def _input_request(messages: list) -> dict: """Arrange a guardrail and act: build the Bedrock INPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format( - source="INPUT", messages=messages - ) + return _grounding_guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) def _output_request(messages: list, response=None) -> dict: """Arrange a guardrail and act: build the Bedrock OUTPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format( - source="OUTPUT", response=response, messages=messages - ) + return _grounding_guardrail().convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages) def test_grounding_input_strips_grounding_and_query_qualifiers(): @@ -2600,9 +2303,7 @@ def test_grounding_input_leaves_existing_guarded_text_unqualified(): """An existing guarded_text input block keeps its legacy unqualified payload.""" expected_request = {"source": "INPUT", "content": [{"text": {"text": "policy"}}]} - actual_request = _input_request( - [{"role": "user", "content": [{"type": "guarded_text", "text": "policy"}]}] - ) + actual_request = _input_request([{"role": "user", "content": [{"type": "guarded_text", "text": "policy"}]}]) assert actual_request == expected_request @@ -2615,9 +2316,7 @@ def test_grounding_output_assembles_source_query_and_response(): "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK], } - actual_request = _output_request( - _grounding_messages(), _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(_grounding_messages(), _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == expected_request @@ -2629,9 +2328,7 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): "content": [{"text": {"text": "Hi there."}}], } - actual_request = _output_request( - [{"role": "user", "content": "hello"}], _model_response("Hi there.") - ) + actual_request = _output_request([{"role": "user", "content": "hello"}], _model_response("Hi there.")) assert actual_request == expected_request @@ -2639,9 +2336,7 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): def test_grounding_output_combines_multiple_sources(): """Every grounding_source block is emitted; Bedrock combines them into one corpus.""" uk_source_text = "London is the capital of UK." - uk_source_block = { - "text": {"text": uk_source_text, "qualifiers": ["grounding_source"]} - } + uk_source_block = {"text": {"text": uk_source_text, "qualifiers": ["grounding_source"]}} messages = [ { "role": "system", @@ -2662,9 +2357,7 @@ def test_grounding_output_combines_multiple_sources(): ], } - actual_request = _output_request( - messages, _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == expected_request @@ -2709,9 +2402,7 @@ def test_grounding_source_trusted_only_from_app_roles(role, is_trusted): if is_trusted: expected_content = [_GROUNDING_SOURCE_BLOCK, *expected_content] - actual_request = _output_request( - messages, _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == {"source": "OUTPUT", "content": expected_content} @@ -2748,12 +2439,8 @@ async def test_grounding_output_blocked_raises_400(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = mock_bedrock_response @@ -2791,13 +2478,7 @@ def _blocked_bedrock_httpx_response() -> MagicMock: response.json.return_value = { "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "Sorry, the model cannot answer this question."}], - "assessments": [ - { - "topicPolicy": { - "topics": [{"name": "Denied", "type": "DENY", "action": "BLOCKED"}] - } - } - ], + "assessments": [{"topicPolicy": {"topics": [{"name": "Denied", "type": "DENY", "action": "BLOCKED"}]}}], } return response @@ -2820,12 +2501,8 @@ async def test_make_bedrock_api_request_block_raises_modify_response_when_flag_s mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2857,12 +2534,8 @@ async def test_make_bedrock_api_request_block_raises_http_400_when_flag_unset(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2903,12 +2576,8 @@ async def test_async_pre_call_hook_propagates_modify_response_on_block(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2953,12 +2622,8 @@ async def test_async_moderation_hook_propagates_modify_response_on_block(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2998,12 +2663,8 @@ async def test_async_post_call_success_hook_attaches_original_response_on_block( mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -3032,9 +2693,7 @@ async def test_apply_guardrail_propagates_modify_response_on_block(): disable_exception_on_block=True, ) - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.side_effect = ModifyResponseException( message="Sorry, the model cannot answer this question.", model="bedrock-nova-micro", @@ -3274,3 +2933,109 @@ async def test_chat_completion_modify_response_exception_streaming_logging_obj_n # CustomStreamWrapper would raise AttributeError inside __init__ and this # call would never reach here. assert response is not None + + +class TestBedrockOnlyScanNewMessages: + """Bedrock pre_call hook honors only_scan_new_messages: scans only the per-session diff.""" + + def _guardrail(self): + return BedrockGuardrail( + guardrail_name="bedrock-incremental", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + only_scan_new_messages=True, + ) + + @pytest.mark.asyncio + async def test_second_turn_scans_only_new_messages(self): + guardrail = self._guardrail() + cache = DualCache() + user_api_key_dict = UserAPIKeyAuth() + bedrock_none = {"action": "NONE", "output": [], "outputs": []} + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = bedrock_none + + turn1 = { + "model": "anthropic/claude-3-5-sonnet", + "messages": [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "first question"}, + ], + "litellm_session_id": "sess-bedrock-1", + } + await guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=cache, data=turn1, call_type="completion" + ) + assert mock_api.call_count == 1 + first_scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in first_scanned] == ["be helpful", "first question"] + + mock_api.reset_mock() + + turn2 = { + "model": "anthropic/claude-3-5-sonnet", + "messages": [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "first question"}, + {"role": "assistant", "content": "first answer"}, + {"role": "user", "content": "second question"}, + ], + "litellm_session_id": "sess-bedrock-1", + } + await guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=cache, data=turn2, call_type="completion" + ) + assert mock_api.call_count == 1 + second_scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in second_scanned] == ["first answer", "second question"] + assert len(turn2["messages"]) == 4 + + @pytest.mark.asyncio + async def test_identical_resend_skips_api_call(self): + guardrail = self._guardrail() + cache = DualCache() + user_api_key_dict = UserAPIKeyAuth() + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + data = { + "model": "anthropic/claude-3-5-sonnet", + "messages": [{"role": "user", "content": "only question"}], + "litellm_session_id": "sess-bedrock-2", + } + await guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=cache, data=dict(data), call_type="completion" + ) + assert mock_api.call_count == 1 + + mock_api.reset_mock() + await guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=cache, data=dict(data), call_type="completion" + ) + mock_api.assert_not_called() + + @pytest.mark.asyncio + async def test_no_session_id_scans_full_context(self): + guardrail = self._guardrail() + cache = DualCache() + user_api_key_dict = UserAPIKeyAuth() + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + data = { + "model": "anthropic/claude-3-5-sonnet", + "messages": [ + {"role": "user", "content": "q1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "q2"}, + ], + "metadata": {}, + } + await guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=cache, data=data, call_type="completion" + ) + assert mock_api.call_count == 1 + scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in scanned] == ["q1", "a1", "q2"]