scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-20 05:04:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 155

feat(guardrails): add only_scan_new_messages for per-session incremental scanning

Browse source 
Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-07-14 21:34:14 +00:00
parent edd3bce0ec
commit e32b2324d4
7 changed files with 607 additions and 704 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/constants.py
View file

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

103
litellm/integrations/custom_guardrail.py
View file

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

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

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

1
litellm/proxy/guardrails/guardrail_initializers.py
View file

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

12
litellm/types/guardrails.py
View file

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

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

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

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

File diff suppressed because it is too large Load diff