From 29ca44115e13876ed0403e52824ef49bc267ff1f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 14 Jul 2026 23:03:33 +0000 Subject: [PATCH] fix(guardrails): run only_scan_new_messages in the unified apply_guardrail path The initial wiring lived in BedrockGuardrail.async_pre_call_hook, but the proxy routes Bedrock through the unified apply_guardrail interface, so the flag had no effect live. Move incremental selection into apply_guardrail: filter the flat texts list against per-session scanned hashes, skip the Bedrock call when nothing is new, and mark hashes only after a successful (non-blocked) scan. Full-context fallback is preserved when there is no session id, the cache is unavailable, or a masking guardrail is configured. Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 45 ++++---- .../guardrail_hooks/bedrock_guardrails.py | 87 ++++++++++----- .../integrations/test_custom_guardrail.py | 98 ++++++++--------- .../test_bedrock_guardrails.py | 100 +++++++++--------- 4 files changed, 183 insertions(+), 147 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 031eeeac9c8..cf9dafcb222 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,5 +1,4 @@ import hashlib -import json import os import secrets from datetime import datetime @@ -277,32 +276,32 @@ class CustomGuardrail(CustomLogger): return get_session_id_from_request_data(request_data) @staticmethod - def _message_content_hash(message: AllMessageValues) -> str: - """Stable content hash for a single message. + def _scanned_text_hash(text: str) -> str: + """Stable content hash for a single scannable text segment. - Serializes the whole message (role + content + tool calls) so that an - edited earlier message produces a different hash and gets re-scanned. + Hashing the exact text the provider would receive means an edited earlier + segment produces a different hash and gets re-scanned, while an unchanged + segment repeated on a later turn is skipped. """ - serialized = json.dumps(message, sort_keys=True, default=str) - return hashlib.sha256(serialized.encode("utf-8")).hexdigest() + return hashlib.sha256(text.encode("utf-8")).hexdigest() - def _scanned_messages_cache_key(self, session_id: str) -> str: - return f"guardrail_scanned_messages:{self.guardrail_name}:{session_id}" + def _scanned_texts_cache_key(self, session_id: str) -> str: + return f"guardrail_scanned_texts:{self.guardrail_name}:{session_id}" - async def filter_new_messages_for_session( + async def filter_new_texts_for_session( self, - messages: list[AllMessageValues] | None, + texts: list[str] | None, request_data: dict[str, object], cache: DualCache, - ) -> list[AllMessageValues] | None: - """Return only the messages not already scanned earlier in this session. + ) -> list[str] | None: + """Return only the text segments 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: + if not self.only_scan_new_messages or not texts: return None if self.mask_request_content or self.mask_response_content: @@ -321,7 +320,7 @@ class CustomGuardrail(CustomLogger): return None try: - cached: object = await cache.async_get_cache(key=self._scanned_messages_cache_key(session_id)) + cached: object = await cache.async_get_cache(key=self._scanned_texts_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.", @@ -331,20 +330,20 @@ class CustomGuardrail(CustomLogger): 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] + return [text for text in texts if self._scanned_text_hash(text) not in seen] - async def mark_messages_scanned( + async def mark_texts_scanned( self, - messages: list[AllMessageValues] | None, + texts: list[str] | None, request_data: dict[str, object], cache: DualCache, ) -> None: - """Record the hashes of all messages present on a successful (non-blocked) scan. + """Record the hashes of all text segments present on a successful (non-blocked) scan. - Called only after the guardrail allows the request, so a blocked message is + Called only after the guardrail allows the request, so a blocked segment is never marked scanned and will be re-checked if the client retries. """ - if not self.only_scan_new_messages or not messages: + if not self.only_scan_new_messages or not texts: return if self.mask_request_content or self.mask_response_content: return @@ -352,8 +351,8 @@ class CustomGuardrail(CustomLogger): 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] + cache_key = self._scanned_texts_cache_key(session_id) + current_hashes = [self._scanned_text_hash(text) for text in texts] 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 [] diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 405e284d9ad..1c40ba5365b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -1172,20 +1172,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): verbose_proxy_logger.debug("No messages found for call_type, skipping guardrail") return data - 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) + filter_result = self._prepare_guardrail_messages_for_role(messages=new_messages) filtered_messages = filter_result.payload_messages if not filtered_messages: @@ -1209,18 +1196,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 2. Update the messages with the guardrail response ########## ######################################################### - 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) + 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, + ) ######################################################### ########## 3. Add the guardrail to the applied guardrails header ########## @@ -1651,6 +1635,48 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): masking_index += 1 verbose_proxy_logger.debug("Applied masking to choice text content") + async def _apply_incremental_request_scan( + self, + texts: list[str], + inputs: "GenericGuardrailAPIInputs", + request_data: dict, + ) -> Optional["GenericGuardrailAPIInputs"]: + """Scan only the text segments not already seen earlier in this session. + + Returns ``None`` when incremental scanning is inactive (feature off, no + session id, masking enabled, or cache unavailable), telling the caller to + run the normal full scan. Otherwise scans only the new segments, skips the + Bedrock call entirely when nothing is new, and never writes masked content + back (incremental mode is for blocking/detection guardrails only). + """ + from litellm.integrations.custom_guardrail import dc as guardrail_dual_cache + + new_texts = await self.filter_new_texts_for_session( + texts=texts, + request_data=request_data, + cache=guardrail_dual_cache, + ) + if new_texts is None: + return None + + if not new_texts: + verbose_proxy_logger.debug("Bedrock Guardrail: no new messages to scan for this session, skipping API call") + return inputs + + await self.make_bedrock_api_request( + source="INPUT", + messages=[ChatCompletionUserMessage(role="user", content=text) for text in new_texts], + request_data=request_data, + logging_event_type=GuardrailEventHooks.pre_call, + ) + + await self.mark_texts_scanned( + texts=texts, + request_data=request_data, + cache=guardrail_dual_cache, + ) + return inputs + async def apply_guardrail( self, inputs: "GenericGuardrailAPIInputs", @@ -1682,6 +1708,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): try: verbose_proxy_logger.debug(f"Bedrock Guardrail: Applying guardrail to {len(texts)} text(s)") + if input_type == "request": + incremental_result = await self._apply_incremental_request_scan( + texts=texts, + inputs=inputs, + request_data=request_data, + ) + if incremental_result is not None: + return incremental_result + masked_texts = [] selection = self._select_messages_for_apply_guardrail( diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 1e2155e795f..4f0e30c356e 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1719,7 +1719,7 @@ class TestApplyGuardrailStyleDeploymentDispatch: class TestOnlyScanNewMessages: - """Incremental guardrail scanning: only send messages not already scanned this session.""" + """Incremental guardrail scanning: only send text segments not already scanned this session.""" def _guardrail(self, **overrides): params = dict(guardrail_name="test-guard", only_scan_new_messages=True) @@ -1734,8 +1734,8 @@ class TestOnlyScanNewMessages: @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"}], + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], request_data={"litellm_session_id": "s1"}, cache=self._cache(), ) @@ -1744,8 +1744,8 @@ class TestOnlyScanNewMessages: @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"}], + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], request_data={"metadata": {}}, cache=self._cache(), ) @@ -1754,8 +1754,8 @@ class TestOnlyScanNewMessages: @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"}], + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], request_data={"litellm_session_id": "s1"}, cache=self._cache(), ) @@ -1768,78 +1768,80 @@ class TestOnlyScanNewMessages: 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"}], + result = await guardrail.filter_new_texts_for_session( + texts=["hi"], request_data={"litellm_session_id": "s1"}, cache=cache, ) assert result is None @pytest.mark.asyncio - async def test_dedupes_previously_scanned_messages(self): + async def test_dedupes_previously_scanned_texts(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"}, - ] + turn1 = ["you are helpful", "first question"] - first = await guardrail.filter_new_messages_for_session(messages=turn1, request_data=request, cache=cache) + first = await guardrail.filter_new_texts_for_session(texts=turn1, request_data=request, cache=cache) assert first == turn1 - await guardrail.mark_messages_scanned(messages=turn1, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=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"}, - ] + turn2 = turn1 + ["an answer", "second question"] + second = await guardrail.filter_new_texts_for_session(texts=turn2, request_data=request, cache=cache) + assert second == ["an answer", "second question"] @pytest.mark.asyncio - async def test_no_new_messages_returns_empty(self): + async def test_no_new_texts_returns_empty(self): guardrail = self._guardrail() cache = self._cache() request = {"litellm_session_id": "sess-empty"} - messages = [{"role": "user", "content": "only message"}] + texts = ["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) + await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=texts, request_data=request, cache=cache) - again = await guardrail.filter_new_messages_for_session(messages=messages, request_data=request, cache=cache) + again = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) assert again == [] @pytest.mark.asyncio - async def test_modified_earlier_message_is_rescanned(self): + async def test_modified_earlier_text_is_rescanned(self): guardrail = self._guardrail() cache = self._cache() request = {"litellm_session_id": "sess-edit"} - original = [{"role": "user", "content": "original"}] + original = ["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) + await guardrail.filter_new_texts_for_session(texts=original, request_data=request, cache=cache) + await guardrail.mark_texts_scanned(texts=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) + edited = ["original EDITED"] + result = await guardrail.filter_new_texts_for_session(texts=edited, request_data=request, cache=cache) assert result == edited @pytest.mark.asyncio - async def test_tool_calls_are_hashed_and_deduped(self): + async def test_blocked_scan_does_not_persist_hashes(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] + request = {"litellm_session_id": "sess-blocked"} + texts = ["please block me"] - 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) + filtered = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert filtered == texts - result = await guardrail.filter_new_messages_for_session(messages=messages, request_data=request, cache=cache) - assert result == [] + again = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache) + assert again == texts + + @pytest.mark.asyncio + async def test_scanned_hashes_written_with_fixed_ttl(self): + from unittest.mock import AsyncMock + + from litellm.constants import GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS + + guardrail = self._guardrail() + cache = self._cache() + cache.async_set_cache = AsyncMock() + request = {"litellm_session_id": "sess-ttl"} + + await guardrail.mark_texts_scanned(texts=["a", "b"], request_data=request, cache=cache) + + cache.async_set_cache.assert_awaited_once() + assert cache.async_set_cache.await_args.kwargs["ttl"] == GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS 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 cf1473a0b11..30caa861f76 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 @@ -3277,7 +3277,13 @@ async def test_chat_completion_modify_response_exception_streaming_logging_obj_n class TestBedrockOnlyScanNewMessages: - """Bedrock pre_call hook honors only_scan_new_messages: scans only the per-session diff.""" + """Bedrock apply_guardrail honors only_scan_new_messages: scans only the per-session diff. + + apply_guardrail is the path the proxy actually runs for Bedrock (via the unified + guardrail interface), so these tests exercise it directly rather than the legacy + async_pre_call_hook. Each test uses a unique session id to isolate the process-wide + incremental cache. + """ def _guardrail(self): return BedrockGuardrail( @@ -3291,23 +3297,16 @@ class TestBedrockOnlyScanNewMessages: @pytest.mark.asyncio async def test_second_turn_scans_only_new_messages(self): guardrail = self._guardrail() - cache = DualCache() - user_api_key_dict = UserAPIKeyAuth() + session = {"litellm_session_id": "sess-bedrock-diff"} 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" + await guardrail.apply_guardrail( + inputs={"texts": ["be helpful", "first question"]}, + request_data=session, + input_type="request", ) assert mock_api.call_count == 1 first_scanned = mock_api.call_args.kwargs["messages"] @@ -3315,68 +3314,69 @@ class TestBedrockOnlyScanNewMessages: 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" + await guardrail.apply_guardrail( + inputs={"texts": ["be helpful", "first question", "first answer", "second question"]}, + request_data=session, + input_type="request", ) 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() + session = {"litellm_session_id": "sess-bedrock-resend"} 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" + + await guardrail.apply_guardrail( + inputs={"texts": ["only question"]}, request_data=session, input_type="request" ) 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" + result = await guardrail.apply_guardrail( + inputs={"texts": ["only question"]}, request_data=session, input_type="request" ) mock_api.assert_not_called() + assert result["texts"] == ["only question"] @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" + + await guardrail.apply_guardrail( + inputs={"texts": ["q1", "a1", "q2"]}, + request_data={"metadata": {}}, + input_type="request", ) assert mock_api.call_count == 1 scanned = mock_api.call_args.kwargs["messages"] assert [m["content"] for m in scanned] == ["q1", "a1", "q2"] + + @pytest.mark.asyncio + async def test_blocked_turn_is_rescanned_on_retry(self): + guardrail = self._guardrail() + session = {"litellm_session_id": "sess-bedrock-blocked"} + + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.side_effect = HTTPException(status_code=400, detail="blocked") + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs={"texts": ["blocked prompt"]}, request_data=session, input_type="request" + ) + + mock_api.reset_mock() + mock_api.side_effect = None + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await guardrail.apply_guardrail( + inputs={"texts": ["blocked prompt"]}, request_data=session, input_type="request" + ) + assert mock_api.call_count == 1 + scanned = mock_api.call_args.kwargs["messages"] + assert [m["content"] for m in scanned] == ["blocked prompt"]