mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-20 05:04:24 +00:00
fix(guardrails): scan tool-result-only messages under experimental_use_latest_role_message_only
_select_messages_for_apply_guardrail used _count_message_texts, which only recognizes a top-level "text" key, to decide whether the latest user message had scannable content. A message whose only content is a tool_result block passed _count_message_texts(...) == 0 and the whole INPUT scan was skipped via skip_scan=True, bypassing Bedrock's ApplyGuardrail API entirely for that message under experimental_use_latest_role_message_only. Use get_content_items_for_message (the tool-result-aware extractor already used elsewhere in this fix) instead, so a tool-result-only message is correctly recognized as having scannable content. Addresses veria-ai review comment on #33092.
This commit is contained in:
parent
5817d7cdc9
commit
598eac84bb
2 changed files with 59 additions and 1 deletions
|
|
@ -532,7 +532,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
return ApplyGuardrailMessageSelection(None, None, True, skip_scan=True)
|
||||
|
||||
selected_message = structured_messages[latest_user_index]
|
||||
if self._count_message_texts(selected_message) == 0:
|
||||
# _count_message_texts only recognizes a top-level "text" key and does not know
|
||||
# about tool_result-nested text, so a message whose only content is a tool_result
|
||||
# block would otherwise look empty here and skip_scan, bypassing the guardrail
|
||||
# entirely. get_content_items_for_message is the tool-result-aware extractor.
|
||||
if not self.get_content_items_for_message(message=selected_message):
|
||||
verbose_proxy_logger.debug(
|
||||
"Bedrock Guardrail: latest user message has no text content, skipping INPUT scan"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3684,3 +3684,57 @@ def test_locate_message_texts_slice_returns_none_on_tool_result_count_mismatch()
|
|||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_select_messages_for_apply_guardrail_does_not_skip_tool_result_only_message():
|
||||
"""veria-ai review (PR #33092): _select_messages_for_apply_guardrail used
|
||||
_count_message_texts -- which only recognizes a top-level "text" key -- to
|
||||
decide whether the latest user message has scannable content. A message
|
||||
whose *only* content is a tool_result block passed
|
||||
_count_message_texts(...) == 0, so the whole INPUT scan was skipped via
|
||||
skip_scan=True: a complete guardrail bypass under
|
||||
experimental_use_latest_role_message_only, since Bedrock's ApplyGuardrail
|
||||
API is never even called (worse than a masking-misalignment bug -- BLOCKED
|
||||
detection doesn't happen either). Must use the tool-result-aware
|
||||
extractor (get_content_items_for_message) instead."""
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
experimental_use_latest_role_message_only=True,
|
||||
)
|
||||
tool_result_only_message = _string_tool_result_message()
|
||||
|
||||
selection = guardrail._select_messages_for_apply_guardrail(
|
||||
texts=[],
|
||||
inputs={"texts": [], "structured_messages": [tool_result_only_message]},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert selection.skip_scan is False
|
||||
assert selection.filtered_messages == [tool_result_only_message]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_scans_tool_result_only_message_under_latest_role_only():
|
||||
"""End-to-end companion to the unit test above: apply_guardrail must
|
||||
actually call Bedrock's ApplyGuardrail API with the tool_result-only
|
||||
message's content, not silently skip_scan and let it through unscanned."""
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
experimental_use_latest_role_message_only=True,
|
||||
)
|
||||
tool_result_only_message = _string_tool_result_message()
|
||||
|
||||
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": [], "structured_messages": [tool_result_only_message]},
|
||||
request_data={"messages": [tool_result_only_message]},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
mock_api.assert_called_once()
|
||||
assert mock_api.call_args.kwargs["messages"] == [tool_result_only_message]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue