scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-19 04:54:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 70

fix(guardrails): run only_scan_new_messages in the unified apply_guardrail path

Browse source 
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>
This commit is contained in:
Devin AI 2026-07-14 23:03:33 +00:00
parent c54b54bd5e
commit 29ca44115e
4 changed files with 183 additions and 147 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
litellm/integrations/custom_guardrail.py
View file

@ -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 []

87
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -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(

98
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -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

100
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -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"]