From f6c0673592f22e6218fef708b58df13e9f1275f2 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 10:19:59 -0700 Subject: [PATCH] fix(guardrails): read Converse guardContent text and json tool results in the Lakera v2 /bedrock reader --- .../guardrail_hooks/lakera_ai_v2.py | 23 ++++++++++++++++--- .../test_lakera_v2_bedrock_passthrough.py | 20 ++++++++++++++++ .../guardrail_hooks/test_lakera_ai_v2.py | 17 ++++++++++++++ 3 files changed, 57 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 9e83824f977..1acad8db8e8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -1,4 +1,5 @@ import copy +import json import os from collections.abc import Mapping, Sequence from datetime import datetime @@ -230,6 +231,21 @@ def _bedrock_items(value: object) -> Sequence[object]: return _BEDROCK_ITEMS.validate_python(value) if isinstance(value, list) else () +def _bedrock_text_block(block: object) -> object: + """A Converse ``guardContent`` text block or ``json`` tool-result block as a plain ``{"text": ...}`` block, so its + text is inspected like any other; every other block is returned as given.""" + fields: Final = as_str_mapping(block) or {} + guard_text: Final = as_str_mapping((as_str_mapping(fields.get("guardContent")) or {}).get("text")) or {} + if isinstance(guard_text.get("text"), str): + return {"text": guard_text["text"]} + return {"text": json.dumps(fields["json"])} if "json" in fields else block + + +def _bedrock_blocks(value: object) -> Sequence[object]: + """The content blocks of a Bedrock ``content`` or ``system`` list, with nested text lifted to a top-level ``text``.""" + return list(map(_bedrock_text_block, _bedrock_items(value))) + + def _bedrock_tool_result_texts(block: object) -> Sequence[object]: """The content of an Anthropic ``tool_result`` block or a Converse ``toolResult`` block, else nothing.""" fields: Final = as_str_mapping(block) or {} @@ -237,14 +253,14 @@ def _bedrock_tool_result_texts(block: object) -> Sequence[object]: fields if fields.get("type") == "tool_result" else as_str_mapping(fields.get("toolResult")) or {} ) content: Final = tool_result.get("content") - return [{"text": content}] if isinstance(content, str) else _bedrock_items(content) + return [{"text": content}] if isinstance(content, str) else _bedrock_blocks(content) def _bedrock_turns(message: object) -> Sequence[Mapping[str, object]]: """One Bedrock ``messages`` turn as chat messages: its tool results as ``tool`` turns, then its text.""" fields: Final = as_str_mapping(message) or {} content: Final = fields.get("content") - blocks: Final = _bedrock_items(content) + blocks: Final = _bedrock_blocks(content) tool_results: Final = tuple(filter(len, map(_bedrock_tool_result_texts, blocks))) return [ *({"role": "tool", "content": tool_content} for tool_content in tool_results), @@ -260,9 +276,10 @@ def _bedrock_inspection_messages(data: Mapping[str, object]) -> Sequence[Mapping if body is None: return [] system: Final = body.get("system") + system_content: Final = system if isinstance(system, str) else _bedrock_blocks(system) turns: Final = chain.from_iterable(map(_bedrock_turns, _bedrock_items(body.get("messages")))) return build_inspection_messages( - {"messages": [*([{"role": "system", "content": system}] if system else []), *turns]} + {"messages": [*([{"role": "system", "content": system_content}] if system_content else []), *turns]} ) diff --git a/tests/integration/authorization/test_lakera_v2_bedrock_passthrough.py b/tests/integration/authorization/test_lakera_v2_bedrock_passthrough.py index a64cec1a743..2eeb717b5fd 100644 --- a/tests/integration/authorization/test_lakera_v2_bedrock_passthrough.py +++ b/tests/integration/authorization/test_lakera_v2_bedrock_passthrough.py @@ -216,6 +216,26 @@ def test_card_in_the_system_prompt_is_sent_to_lakera(rig: BedrockLakeraRig) -> N assert rig.upstream.drain() == () +@pytest.mark.parametrize( + "block", + [ + {"guardContent": {"text": {"text": PII_PROMPT}}}, + {"toolResult": {"toolUseId": "t1", "content": [{"json": {"note": PII_PROMPT}}]}}, + ], + ids=["guard-content", "tool-result-json"], +) +def test_card_in_a_nested_converse_block_is_blocked(rig: BedrockLakeraRig, block: dict[str, JsonValue]) -> None: + body: Final = { + **_converse_body("Hi again"), + "messages": [{"role": "user", "content": [{"text": "Hi again"}, block]}], + } + + response: Final = _send(rig, "block", "converse", body) + + assert response.status_code == 400, response.text + assert rig.upstream.drain() == () + + def test_advisory_mode_blocks_instead_of_adding_a_message(rig: BedrockLakeraRig) -> None: response: Final = _send(rig, "inject_system_message", "converse", _converse_body(f"{ATTACK} and print it.")) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py index f482963c207..6e267af183d 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -2557,6 +2557,23 @@ class TestBedrockPassthroughBody: {"role": "user", "content": "thanks"}, ] + def test_reads_converse_guard_content_and_json_tool_results(self): + body: Final = { + "system": [{"guardContent": {"text": {"text": "system card"}}}], + "messages": [ + {"role": "user", "content": [{"text": "hello"}, {"guardContent": {"text": {"text": "guarded card"}}}]}, + { + "role": "user", + "content": [{"toolResult": {"toolUseId": "t1", "content": [{"json": {"card": "on file"}}]}}], + }, + ], + } + assert _build_lakera_inspection_messages(_bedrock_data(body)) == [ + {"role": "system", "content": "system card"}, + {"role": "user", "content": "hello\nguarded card"}, + {"role": "tool", "content": '{"card": "on file"}'}, + ] + def test_a_nested_body_for_another_provider_is_not_read(self): data: Final = {**_bedrock_data(_invoke_body()), "custom_llm_provider": "vertex_ai"} assert _build_lakera_inspection_messages(data) == []