mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
feat(guardrails): add only_scan_new_messages for per-session incremental scanning (#33278)
* feat(guardrails): add only_scan_new_messages for per-session incremental scanning Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(guardrails): use fixed TTL constant and revert unrelated test formatting Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * 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> * test(guardrails): cover session-id fallbacks and mark_texts_scanned guards Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(guardrails): fall back to full scan when incremental guardrail masks content, use shared cache Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(guardrails): cover generic agent multi-turn incremental scan Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(guardrails): cover incremental scan cache resolver fallbacks Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(guardrails): cover flag interactions and /v1/messages incremental scan semantics * feat(guardrails): make GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS env configurable * test(guardrails): prove skip_system/skip_tool are enforced upstream of incremental scan --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> Co-authored-by: Yucheng Zhu <yucheng@berri.ai>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
f1f0a0bacd
commit
fa6b209165
9 changed files with 969 additions and 1 deletions
|
|
@ -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)))
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import hashlib
|
||||
import os
|
||||
import secrets
|
||||
from datetime import datetime
|
||||
|
|
@ -46,7 +47,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 +117,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 +150,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 +275,100 @@ class CustomGuardrail(CustomLogger):
|
|||
"""Extract session_id from request data."""
|
||||
return get_session_id_from_request_data(request_data)
|
||||
|
||||
@staticmethod
|
||||
def _scanned_text_hash(text: str) -> str:
|
||||
"""Stable content hash for a single scannable text segment.
|
||||
|
||||
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.
|
||||
"""
|
||||
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
||||
|
||||
def _scanned_texts_cache_key(self, session_id: str) -> str:
|
||||
return f"guardrail_scanned_texts:{self.guardrail_name}:{session_id}"
|
||||
|
||||
async def filter_new_texts_for_session(
|
||||
self,
|
||||
texts: list[str] | None,
|
||||
request_data: dict[str, object],
|
||||
cache: DualCache,
|
||||
) -> 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 texts:
|
||||
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_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.",
|
||||
self.guardrail_name,
|
||||
e,
|
||||
)
|
||||
return None
|
||||
|
||||
seen: set[str] = {str(h) for h in cached} if isinstance(cached, list) else set()
|
||||
return [text for text in texts if self._scanned_text_hash(text) not in seen]
|
||||
|
||||
async def mark_texts_scanned(
|
||||
self,
|
||||
texts: list[str] | None,
|
||||
request_data: dict[str, object],
|
||||
cache: DualCache,
|
||||
) -> None:
|
||||
"""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 segment is
|
||||
never marked scanned and will be re-checked if the client retries.
|
||||
"""
|
||||
if not self.only_scan_new_messages or not texts:
|
||||
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_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 []
|
||||
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
|
||||
|
|
|
|||
|
|
@ -2046,6 +2046,90 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
masking_index += 1
|
||||
verbose_proxy_logger.debug("Applied masking to choice text content")
|
||||
|
||||
@staticmethod
|
||||
def _incremental_scan_cache() -> DualCache:
|
||||
"""Resolve the cache used to remember which segments a session already scanned.
|
||||
|
||||
Prefers the proxy's shared cache (``internal_usage_cache.dual_cache``), which is
|
||||
backed by Redis when the deployment configures it, so incremental state is shared
|
||||
across proxy instances. Falls back to a process-local ``DualCache`` singleton when
|
||||
the proxy is not running (e.g. unit tests), where sharing does not apply.
|
||||
"""
|
||||
from litellm.integrations.custom_guardrail import dc as fallback_cache
|
||||
|
||||
try:
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging
|
||||
except Exception: # noqa: BLE001 # proxy not importable outside the server; use local fallback
|
||||
return fallback_cache
|
||||
if _proxy_logging is not None:
|
||||
return _proxy_logging.internal_usage_cache.dual_cache
|
||||
return fallback_cache
|
||||
|
||||
def _bedrock_response_has_masked_output(self, response: BedrockGuardrailResponse) -> bool:
|
||||
"""Return True if the guardrail rewrote (masked/anonymized) any scanned text.
|
||||
|
||||
Bedrock returns non-empty ``output``/``outputs`` text only when it changed the
|
||||
content; an ``action == "NONE"`` response leaves both empty.
|
||||
"""
|
||||
for field in ("output", "outputs"):
|
||||
items = response.get(field) or []
|
||||
if any(isinstance(item, dict) and item.get("text") for item in items):
|
||||
return True
|
||||
return False
|
||||
|
||||
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) or when the guardrail
|
||||
turns out to mask content, telling the caller to run the normal full scan.
|
||||
Otherwise scans only the new segments and skips the Bedrock call entirely
|
||||
when nothing is new. Incremental mode is for blocking/detection guardrails
|
||||
only: if the guardrail returns masked output it cannot be applied to the
|
||||
skipped context, so the scan falls back to the full path and no session
|
||||
state is recorded.
|
||||
"""
|
||||
cache = self._incremental_scan_cache()
|
||||
|
||||
new_texts = await self.filter_new_texts_for_session(
|
||||
texts=texts,
|
||||
request_data=request_data,
|
||||
cache=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
|
||||
|
||||
bedrock_response = 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,
|
||||
)
|
||||
|
||||
if self._bedrock_response_has_masked_output(bedrock_response):
|
||||
verbose_proxy_logger.warning(
|
||||
"Bedrock Guardrail %s: guardrail returned masked/anonymized content; "
|
||||
"only_scan_new_messages cannot apply masking to skipped context, falling back to a full-context scan",
|
||||
self.guardrail_name,
|
||||
)
|
||||
return None
|
||||
|
||||
await self.mark_texts_scanned(
|
||||
texts=texts,
|
||||
request_data=request_data,
|
||||
cache=cache,
|
||||
)
|
||||
return inputs
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: "GenericGuardrailAPIInputs",
|
||||
|
|
@ -2077,6 +2161,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(
|
||||
|
|
|
|||
|
|
@ -35,6 +35,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
|
||||
|
|
|
|||
|
|
@ -725,6 +725,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=(
|
||||
|
|
|
|||
|
|
@ -1716,3 +1716,201 @@ 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 text segments 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_texts_for_session(
|
||||
texts=["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_texts_for_session(
|
||||
texts=["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_texts_for_session(
|
||||
texts=["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_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_texts(self):
|
||||
guardrail = self._guardrail()
|
||||
cache = self._cache()
|
||||
request = {"litellm_session_id": "sess-dedupe"}
|
||||
turn1 = ["you are helpful", "first question"]
|
||||
|
||||
first = await guardrail.filter_new_texts_for_session(texts=turn1, request_data=request, cache=cache)
|
||||
assert first == turn1
|
||||
await guardrail.mark_texts_scanned(texts=turn1, request_data=request, cache=cache)
|
||||
|
||||
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_texts_returns_empty(self):
|
||||
guardrail = self._guardrail()
|
||||
cache = self._cache()
|
||||
request = {"litellm_session_id": "sess-empty"}
|
||||
texts = ["only message"]
|
||||
|
||||
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_texts_for_session(texts=texts, request_data=request, cache=cache)
|
||||
assert again == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_modified_earlier_text_is_rescanned(self):
|
||||
guardrail = self._guardrail()
|
||||
cache = self._cache()
|
||||
request = {"litellm_session_id": "sess-edit"}
|
||||
original = ["original"]
|
||||
|
||||
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 = ["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_blocked_scan_does_not_persist_hashes(self):
|
||||
guardrail = self._guardrail()
|
||||
cache = self._cache()
|
||||
request = {"litellm_session_id": "sess-blocked"}
|
||||
texts = ["please block me"]
|
||||
|
||||
filtered = await guardrail.filter_new_texts_for_session(texts=texts, request_data=request, cache=cache)
|
||||
assert filtered == texts
|
||||
|
||||
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
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_id_from_metadata_is_used_for_dedupe(self):
|
||||
guardrail = self._guardrail()
|
||||
cache = self._cache()
|
||||
request = {"metadata": {"session_id": "sess-meta"}}
|
||||
texts = ["shared message"]
|
||||
|
||||
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_texts_for_session(texts=texts, request_data=request, cache=cache)
|
||||
assert again == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_id_from_litellm_metadata_is_used_for_dedupe(self):
|
||||
guardrail = self._guardrail()
|
||||
cache = self._cache()
|
||||
request = {"litellm_metadata": {"session_id": "sess-lmeta"}}
|
||||
texts = ["shared message"]
|
||||
|
||||
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_texts_for_session(texts=texts, request_data=request, cache=cache)
|
||||
assert again == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mark_texts_scanned_disabled_does_not_persist(self):
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
guardrail = self._guardrail(only_scan_new_messages=False)
|
||||
cache = self._cache()
|
||||
cache.async_set_cache = AsyncMock()
|
||||
|
||||
await guardrail.mark_texts_scanned(texts=["a"], request_data={"litellm_session_id": "s1"}, cache=cache)
|
||||
cache.async_set_cache.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mark_texts_scanned_masking_does_not_persist(self):
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
guardrail = self._guardrail(mask_request_content=True)
|
||||
cache = self._cache()
|
||||
cache.async_set_cache = AsyncMock()
|
||||
|
||||
await guardrail.mark_texts_scanned(texts=["a"], request_data={"litellm_session_id": "s1"}, cache=cache)
|
||||
cache.async_set_cache.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mark_texts_scanned_without_session_does_not_persist(self):
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
guardrail = self._guardrail()
|
||||
cache = self._cache()
|
||||
cache.async_set_cache = AsyncMock()
|
||||
|
||||
await guardrail.mark_texts_scanned(texts=["a"], request_data={"metadata": {}}, cache=cache)
|
||||
cache.async_set_cache.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mark_texts_scanned_survives_cache_write_failure(self):
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
guardrail = self._guardrail()
|
||||
cache = self._cache()
|
||||
cache.async_set_cache = AsyncMock(side_effect=RuntimeError("redis down"))
|
||||
|
||||
await guardrail.mark_texts_scanned(texts=["a"], request_data={"litellm_session_id": "s1"}, cache=cache)
|
||||
|
|
|
|||
|
|
@ -373,3 +373,132 @@ class TestAnthropicMessagesHandlerToolInjection:
|
|||
if __name__ == "__main__":
|
||||
# Run the tests
|
||||
pytest.main([__file__, "-v"])
|
||||
|
||||
|
||||
class TestAnthropicMessagesIncrementalScan:
|
||||
"""PR #33278: only_scan_new_messages through the real /v1/messages translation
|
||||
handler (the path Claude Code uses). Encodes the wire payloads observed in the
|
||||
live validation against a real Bedrock guardrail.
|
||||
"""
|
||||
|
||||
def _bedrock_guardrail(self):
|
||||
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail
|
||||
|
||||
return BedrockGuardrail(
|
||||
guardrail_name="bedrock-incremental-anthropic",
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
default_on=True,
|
||||
only_scan_new_messages=True,
|
||||
)
|
||||
|
||||
def _data(self, messages, session_id):
|
||||
return {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": messages,
|
||||
"system": "You are a helpful geography assistant.",
|
||||
"litellm_session_id": session_id,
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_first_turn_scans_all_eligible_then_second_turn_scans_only_diff(self):
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
handler = AnthropicMessagesHandler()
|
||||
guardrail = self._bedrock_guardrail()
|
||||
sid = "anth-sess-diff"
|
||||
turn1 = [{"role": "user", "content": "What is the capital of France?"}]
|
||||
turn2 = turn1 + [
|
||||
{"role": "assistant", "content": "Paris."},
|
||||
{"role": "user", "content": "What is the capital of Germany?"},
|
||||
]
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
|
||||
await handler.process_input_messages(
|
||||
data=self._data(turn1, sid), guardrail_to_apply=guardrail
|
||||
)
|
||||
assert mock_api.call_count == 1
|
||||
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [
|
||||
"What is the capital of France?"
|
||||
]
|
||||
mock_api.reset_mock()
|
||||
await handler.process_input_messages(
|
||||
data=self._data(turn2, sid), guardrail_to_apply=guardrail
|
||||
)
|
||||
assert mock_api.call_count == 1
|
||||
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [
|
||||
"Paris.",
|
||||
"What is the capital of Germany?",
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_identical_resend_makes_no_guardrail_call(self):
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
handler = AnthropicMessagesHandler()
|
||||
guardrail = self._bedrock_guardrail()
|
||||
sid = "anth-sess-resend"
|
||||
msgs = [
|
||||
{"role": "user", "content": "What is the capital of France?"},
|
||||
{"role": "assistant", "content": "Paris."},
|
||||
{"role": "user", "content": "What is the capital of Germany?"},
|
||||
]
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
|
||||
await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail)
|
||||
assert mock_api.call_count == 1
|
||||
mock_api.reset_mock()
|
||||
await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail)
|
||||
mock_api.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edited_history_message_is_rescanned(self):
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
handler = AnthropicMessagesHandler()
|
||||
guardrail = self._bedrock_guardrail()
|
||||
sid = "anth-sess-edit"
|
||||
msgs = [{"role": "user", "content": "What is the capital of France?"}]
|
||||
edited = [{"role": "user", "content": "What is the capital and population of France?"}]
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
|
||||
await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail)
|
||||
mock_api.reset_mock()
|
||||
await handler.process_input_messages(data=self._data(edited, sid), guardrail_to_apply=guardrail)
|
||||
assert mock_api.call_count == 1
|
||||
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [
|
||||
"What is the capital and population of France?"
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mixed_text_and_tool_use_keeps_text_segments(self):
|
||||
"""A message carrying both text and a tool_use block must not lose its text.
|
||||
(tool_use inputs and tool_result content are dropped from texts on the
|
||||
anthropic input path today; that is pre-existing baseline behavior.)"""
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
handler = AnthropicMessagesHandler()
|
||||
guardrail = self._bedrock_guardrail()
|
||||
sid = "anth-sess-tools"
|
||||
msgs = [
|
||||
{"role": "user", "content": "Search for the weather in Paris"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Let me look that up for you."},
|
||||
{"type": "tool_use", "id": "toolu_1", "name": "search", "input": {"query": "canary-args"}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "canary-result"}],
|
||||
},
|
||||
{"role": "user", "content": "Thanks, summarize the result."},
|
||||
]
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
|
||||
await handler.process_input_messages(data=self._data(msgs, sid), guardrail_to_apply=guardrail)
|
||||
scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]]
|
||||
assert "Let me look that up for you." in scanned, "text beside a tool_use must be scanned"
|
||||
assert "Search for the weather in Paris" in scanned
|
||||
assert "Thanks, summarize the result." in scanned
|
||||
|
|
|
|||
|
|
@ -1137,3 +1137,95 @@ class TestGetStructuredMessages:
|
|||
if __name__ == "__main__":
|
||||
# Run the tests
|
||||
pytest.main([__file__, "-v"])
|
||||
|
||||
|
||||
class TestIncrementalScanRespectsSkipFlags:
|
||||
"""PR #33278: skip_system_message_in_guardrail and skip_tool_message_in_guardrail
|
||||
are enforced while this handler builds inputs["texts"] (_extract_inputs early
|
||||
returns for system/tool roles), upstream of BedrockGuardrail's incremental path.
|
||||
Bypassing _select_messages_for_apply_guardrail therefore cannot resurrect skipped
|
||||
content on any turn, including a session's first turn where every segment is new.
|
||||
Verified live against a real Bedrock ApplyGuardrail before being encoded here.
|
||||
The flags are set as instance attributes, mirroring how guardrail_registry
|
||||
applies litellm_params to the callback (they are not constructor kwargs).
|
||||
"""
|
||||
|
||||
def _bedrock_guardrail(self):
|
||||
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail
|
||||
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="bedrock-incremental-skip-flags",
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
default_on=True,
|
||||
only_scan_new_messages=True,
|
||||
)
|
||||
guardrail.skip_system_message_in_guardrail = True
|
||||
guardrail.skip_tool_message_in_guardrail = True
|
||||
return guardrail
|
||||
|
||||
def _messages(self, followup=None):
|
||||
base = [
|
||||
{"role": "system", "content": "SYSTEM-PROMPT-must-not-be-scanned"},
|
||||
{"role": "user", "content": "Search for the weather in Paris"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Let me look that up.",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "search", "arguments": '{"query": "weather"}'},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT-must-not-be-scanned"},
|
||||
{"role": "user", "content": "Thanks, summarize."},
|
||||
]
|
||||
return base + (followup or [])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_first_turn_scans_no_system_or_tool_content(self):
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
guardrail = self._bedrock_guardrail()
|
||||
data = {"messages": self._messages(), "litellm_session_id": "skip-flags-turn1"}
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
assert mock_api.call_count == 1
|
||||
scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]]
|
||||
assert scanned == [
|
||||
"Search for the weather in Paris",
|
||||
"Let me look that up.",
|
||||
"Thanks, summarize.",
|
||||
]
|
||||
assert not any("SYSTEM-PROMPT" in text for text in scanned)
|
||||
assert not any("TOOL-RESULT" in text for text in scanned)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_second_turn_scans_only_new_eligible_content(self):
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
handler = OpenAIChatCompletionsHandler()
|
||||
guardrail = self._bedrock_guardrail()
|
||||
session = "skip-flags-turn2"
|
||||
followup = [
|
||||
{"role": "assistant", "content": "It is sunny in Paris."},
|
||||
{"role": "user", "content": "And tomorrow?"},
|
||||
]
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
|
||||
await handler.process_input_messages(
|
||||
data={"messages": self._messages(), "litellm_session_id": session},
|
||||
guardrail_to_apply=guardrail,
|
||||
)
|
||||
mock_api.reset_mock()
|
||||
await handler.process_input_messages(
|
||||
data={"messages": self._messages(followup), "litellm_session_id": session},
|
||||
guardrail_to_apply=guardrail,
|
||||
)
|
||||
assert mock_api.call_count == 1
|
||||
scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]]
|
||||
assert scanned == ["It is sunny in Paris.", "And tomorrow?"]
|
||||
|
|
|
|||
|
|
@ -3274,3 +3274,343 @@ 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 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(
|
||||
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()
|
||||
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
|
||||
|
||||
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"]
|
||||
assert [m["content"] for m in first_scanned] == ["be helpful", "first question"]
|
||||
|
||||
mock_api.reset_mock()
|
||||
|
||||
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"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_identical_resend_skips_api_call(self):
|
||||
guardrail = self._guardrail()
|
||||
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": []}
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["only question"]}, request_data=session, input_type="request"
|
||||
)
|
||||
assert mock_api.call_count == 1
|
||||
|
||||
mock_api.reset_mock()
|
||||
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()
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
|
||||
|
||||
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_masking_guardrail_falls_back_and_does_not_persist(self):
|
||||
"""A guardrail that anonymizes content must not be short-circuited.
|
||||
|
||||
Regression: the incremental fast path used to ignore the guardrail response,
|
||||
so masked/anonymized output was dropped, the raw text reached the model, and
|
||||
the segment was marked scanned so it was never re-checked. Detecting masked
|
||||
output must force a full-context scan (which applies the masking) and must not
|
||||
persist session state, so an identical resend is scanned again.
|
||||
"""
|
||||
guardrail = self._guardrail()
|
||||
session = {"litellm_session_id": "sess-bedrock-mask"}
|
||||
masked = {
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
"output": [],
|
||||
"outputs": [{"text": "my ssn is [REDACTED]"}],
|
||||
}
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = masked
|
||||
|
||||
result = await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["my ssn is 123-45-6789"]},
|
||||
request_data=session,
|
||||
input_type="request",
|
||||
)
|
||||
assert mock_api.call_count == 2
|
||||
assert result["texts"] == ["my ssn is [REDACTED]"]
|
||||
|
||||
mock_api.reset_mock()
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["my ssn is 123-45-6789"]},
|
||||
request_data=session,
|
||||
input_type="request",
|
||||
)
|
||||
assert mock_api.call_count >= 1
|
||||
first_scanned = mock_api.call_args_list[0].kwargs.get("messages")
|
||||
assert first_scanned is not None
|
||||
assert [m["content"] for m in first_scanned] == ["my ssn is 123-45-6789"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_agent_multi_turn_scans_only_new_each_turn(self):
|
||||
"""A generic agent (not Claude Code) opts in by propagating a session id.
|
||||
|
||||
Agent frameworks on the OpenAI SDK carry the session through the request
|
||||
body (metadata.session_id here), not the x-claude-code-session-id header.
|
||||
Across a growing multi-turn conversation every turn after the first must
|
||||
send Bedrock only the newly appended segments, never the whole context.
|
||||
"""
|
||||
guardrail = self._guardrail()
|
||||
session = {"metadata": {"session_id": "agent-multi-turn"}}
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["system prompt", "turn 1 question"]},
|
||||
request_data=session,
|
||||
input_type="request",
|
||||
)
|
||||
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [
|
||||
"system prompt",
|
||||
"turn 1 question",
|
||||
]
|
||||
|
||||
mock_api.reset_mock()
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["system prompt", "turn 1 question", "turn 1 answer", "turn 2 question"]},
|
||||
request_data=session,
|
||||
input_type="request",
|
||||
)
|
||||
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [
|
||||
"turn 1 answer",
|
||||
"turn 2 question",
|
||||
]
|
||||
|
||||
mock_api.reset_mock()
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={
|
||||
"texts": [
|
||||
"system prompt",
|
||||
"turn 1 question",
|
||||
"turn 1 answer",
|
||||
"turn 2 question",
|
||||
"turn 2 answer",
|
||||
"turn 3 question",
|
||||
]
|
||||
},
|
||||
request_data=session,
|
||||
input_type="request",
|
||||
)
|
||||
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [
|
||||
"turn 2 answer",
|
||||
"turn 3 question",
|
||||
]
|
||||
|
||||
def test_incremental_scan_cache_prefers_proxy_shared_cache(self):
|
||||
guardrail = self._guardrail()
|
||||
shared = DualCache()
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.internal_usage_cache.dual_cache = shared
|
||||
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging):
|
||||
assert guardrail._incremental_scan_cache() is shared
|
||||
|
||||
def test_incremental_scan_cache_falls_back_when_proxy_logging_missing(self):
|
||||
from litellm.integrations.custom_guardrail import dc as fallback_cache
|
||||
|
||||
guardrail = self._guardrail()
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj", None):
|
||||
assert guardrail._incremental_scan_cache() is fallback_cache
|
||||
|
||||
def test_incremental_scan_cache_falls_back_when_proxy_not_importable(self):
|
||||
from litellm.integrations.custom_guardrail import dc as fallback_cache
|
||||
|
||||
guardrail = self._guardrail()
|
||||
with patch.dict(sys.modules, {"litellm.proxy.proxy_server": None}):
|
||||
assert guardrail._incremental_scan_cache() is fallback_cache
|
||||
|
||||
@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"]
|
||||
|
||||
|
||||
class TestBedrockIncrementalFlagInteractions:
|
||||
"""Regression coverage for only_scan_new_messages combined with the other
|
||||
Bedrock guardrail flags, from the PR #33278 live validation. Live evidence:
|
||||
each of these was reproduced against a real Bedrock ApplyGuardrail first;
|
||||
the mocks here encode the wire payloads observed there.
|
||||
"""
|
||||
|
||||
def _guardrail(self, **overrides):
|
||||
params = dict(
|
||||
guardrail_name="bedrock-incremental-flags",
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
default_on=True,
|
||||
only_scan_new_messages=True,
|
||||
)
|
||||
params.update(overrides)
|
||||
return BedrockGuardrail(**params)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_edited_history_segment_rescans_only_that_segment(self):
|
||||
guardrail = self._guardrail()
|
||||
session = {"litellm_session_id": "sess-flags-edit"}
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["q1", "a1", "q2"]}, request_data=session, input_type="request"
|
||||
)
|
||||
mock_api.reset_mock()
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["q1 EDITED", "a1", "q2"]}, request_data=session, input_type="request"
|
||||
)
|
||||
assert mock_api.call_count == 1
|
||||
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == ["q1 EDITED"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_content_different_session_rescans_everything(self):
|
||||
guardrail = self._guardrail()
|
||||
texts = ["shared question", "shared answer"]
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": list(texts)}, request_data={"litellm_session_id": "sess-x1"}, input_type="request"
|
||||
)
|
||||
mock_api.reset_mock()
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": list(texts)}, request_data={"litellm_session_id": "sess-x2"}, input_type="request"
|
||||
)
|
||||
assert mock_api.call_count == 1
|
||||
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == texts
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_masking_flag_disables_incremental_single_full_scan(self):
|
||||
"""mask_request_content must fall back to exactly ONE full scan per turn
|
||||
and never persist hashes (verified live: 1 call/turn, no cache writes)."""
|
||||
guardrail = self._guardrail(mask_request_content=True)
|
||||
session = {"litellm_session_id": "sess-flags-mask"}
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["q1"]}, request_data=session, input_type="request"
|
||||
)
|
||||
assert mock_api.call_count == 1
|
||||
mock_api.reset_mock()
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["q1"]}, request_data=session, input_type="request"
|
||||
)
|
||||
assert mock_api.call_count == 1, "masking mode must re-scan every turn, exactly once"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_server_side_anonymize_falls_back_full_scan_and_never_persists(self):
|
||||
"""A guardrail that rewrites content (Bedrock-side ANONYMIZE) must fall back
|
||||
to the full scan so masking applies, and record no session state. Live
|
||||
validation showed this costs 2 provider calls per turn; the count is
|
||||
asserted here as documentation of that intended-tradeoff behavior."""
|
||||
guardrail = self._guardrail()
|
||||
session = {"litellm_session_id": "sess-flags-anon"}
|
||||
masked = {"action": "NONE", "output": [{"text": "MASKED q1"}], "outputs": [{"text": "MASKED q1"}]}
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = masked
|
||||
result = await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["q1"]}, request_data=session, input_type="request"
|
||||
)
|
||||
assert mock_api.call_count == 2, "incremental attempt + full-scan fallback"
|
||||
assert result["texts"] == ["MASKED q1"], "masked content must be applied"
|
||||
mock_api.reset_mock()
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["q1"]}, request_data=session, input_type="request"
|
||||
)
|
||||
assert mock_api.call_count == 2, "no hashes persisted, so the double scan repeats"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.xfail(
|
||||
reason="PR #33278 known gap: incremental path bypasses _select_messages_for_apply_guardrail, "
|
||||
"so experimental_use_latest_role_message_only is silently ignored. Intended semantics "
|
||||
"(pending DRI decision): incremental mode defers to the latest-role selection.",
|
||||
strict=False,
|
||||
)
|
||||
async def test_latest_role_only_is_respected_with_incremental(self):
|
||||
guardrail = self._guardrail(experimental_use_latest_role_message_only=True)
|
||||
session = {"litellm_session_id": "sess-flags-latestrole"}
|
||||
structured = [
|
||||
{"role": "system", "content": "sys"},
|
||||
{"role": "user", "content": "q1"},
|
||||
]
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {"action": "NONE", "output": [], "outputs": []}
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["sys", "q1"], "structured_messages": structured},
|
||||
request_data=session,
|
||||
input_type="request",
|
||||
)
|
||||
scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]]
|
||||
assert scanned == ["q1"], "latest-role selection must exclude the system prompt"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue