From 63cd43143ede04bf17722a2f740ffd78e20e7b3d Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:19:10 -0700 Subject: [PATCH] fix(bedrock): apply Converse guardrail rewrites texts-first and keep main's write-backs Map tool results onto their text and json parts only, accept a structured row the flat texts already reproduce, let a flat-text mask win over a structured leaf, ignore texts that do not line up with the request when structured rows come back, keep tool call arguments out of the rows under skip_tool_message_in_guardrail and send no tool definitions to tool-only guardrails --- .../guardrail_translation/handler.py | 400 +++++++++++----- ...est_bedrock_converse_guardrail_rewrites.py | 451 ++++++++++++++++++ .../guardrail_translation/test_handler.py | 301 ++++++++++++ 3 files changed, 1031 insertions(+), 121 deletions(-) create mode 100644 tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 70ed49ebbc5..89920dba9b3 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -2,6 +2,7 @@ import copy import json from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from functools import partial, reduce from itertools import chain, groupby from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, TypeAlias @@ -37,8 +38,8 @@ if TYPE_CHECKING: from litellm.llms.pass_through.guardrail_translation.handler import ( PassThroughEndpointHandler, ) - from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging + from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth _CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) _EVENT_STREAM_CONTENT_TYPE: Final = "vnd.amazon.eventstream" @@ -168,7 +169,7 @@ def _extract_converse_texts( class _ConverseModel(BaseModel): - model_config = ConfigDict(extra="ignore", frozen=True, populate_by_name=True) + model_config = ConfigDict(extra="ignore", frozen=True) class _ConverseGuardText(_ConverseModel): @@ -263,7 +264,11 @@ class _ConvertedMessage: _IndexedBlock: TypeAlias = tuple[_BlockRef, _ConverseContentBlock] _TEXT_PARTS: Final = TypeAdapter(tuple[ChatCompletionTextObject, ...]) +_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_ITEMS: Final = TypeAdapter(Sequence[object]) +_PARSED_JSON: Final = TypeAdapter(object) _TextEdit: TypeAlias = tuple[_BlockRef, object] +_JsonPair: TypeAlias = tuple[tuple[str | int, ...], object, object] def _indexed_blocks(message_index: int | None, blocks: Sequence[_ConverseContentBlock]) -> tuple[_IndexedBlock, ...]: @@ -282,15 +287,17 @@ def _refs(text_blocks: Sequence[tuple[_BlockRef, str]]) -> tuple[_BlockRef, ...] return tuple(ref for ref, _ in text_blocks) +def _scannable_tool_parts(tool_result: _ConverseToolResult) -> tuple[tuple[int, _ConverseToolResultContent], ...]: + return tuple( + (index, part) for index, part in enumerate(tool_result.content) if part.text or part.json_value is not None + ) + + def _converse_tool_message(tool_result: _ConverseToolResult) -> ChatCompletionToolMessage: return ChatCompletionToolMessage( role="tool", tool_call_id=tool_result.tool_use_id, - content="\n".join( - part.text or json.dumps(part.json_value) - for part in tool_result.content - if part.text or part.json_value is not None - ), + content="\n".join(part.text or json.dumps(part.json_value) for _, part in _scannable_tool_parts(tool_result)), ) @@ -328,10 +335,14 @@ def _converse_user_messages(blocks: Sequence[_IndexedBlock]) -> tuple[_Converted ) -def _converse_assistant_messages(blocks: Sequence[_IndexedBlock]) -> tuple[_ConvertedMessage, ...]: +def _converse_assistant_messages( + blocks: Sequence[_IndexedBlock], include_tool_calls: bool +) -> tuple[_ConvertedMessage, ...]: text_blocks: Final = _text_blocks(blocks) content: Final = list(_text_content(text_blocks)) - tool_calls: Final = [_converse_tool_call(block.tool_use) for _, block in blocks if block.tool_use] + tool_calls: Final = [ + _converse_tool_call(block.tool_use) for _, block in blocks if block.tool_use and include_tool_calls + ] if tool_calls: message: Final = ChatCompletionAssistantMessage( role="assistant", content=content or None, tool_calls=tool_calls @@ -348,16 +359,20 @@ def _converse_assistant_messages(blocks: Sequence[_IndexedBlock]) -> tuple[_Conv return (_ConvertedMessage(ChatCompletionAssistantMessage(role="assistant", content=content), _refs(text_blocks)),) -def _converse_message(message_index: int, message: _ConverseMessage) -> tuple[_ConvertedMessage, ...]: +def _converse_message( + message_index: int, message: _ConverseMessage, include_tool_calls: bool +) -> tuple[_ConvertedMessage, ...]: blocks: Final = _indexed_blocks(message_index, message.content) if message.role == "assistant": - return _converse_assistant_messages(blocks) + return _converse_assistant_messages(blocks, include_tool_calls) return _converse_user_messages(blocks) -def _converse_structured_messages(request: _ConverseRequest) -> tuple[_ConvertedMessage, ...]: +def _converse_structured_messages(request: _ConverseRequest, include_tool_calls: bool) -> tuple[_ConvertedMessage, ...]: converted: Final = tuple( - chain.from_iterable(_converse_message(index, message) for index, message in enumerate(request.messages)) + chain.from_iterable( + _converse_message(index, message, include_tool_calls) for index, message in enumerate(request.messages) + ) ) system_blocks: Final = _text_blocks(_indexed_blocks(None, request.system)) if not system_blocks: @@ -366,143 +381,269 @@ def _converse_structured_messages(request: _ConverseRequest) -> tuple[_Converted return (_ConvertedMessage(system, _refs(system_blocks)), *converted) +def _without(mapping: Mapping[str, object], key: str) -> tuple[tuple[str, object], ...]: + return tuple(sorted(((name, value) for name, value in mapping.items() if name != key), key=lambda item: item[0])) + + def _without_content(message: Mapping[str, object]) -> tuple[tuple[str, object], ...]: - return tuple(sorted(((key, value) for key, value in message.items() if key != "content"), key=lambda item: item[0])) + return _without(message, "content") + + +def _fields(message: AllMessageValues) -> Mapping[str, object]: + return message + + +def _as_object(value: object) -> Mapping[str, object] | None: + try: + return _OBJECT.validate_python(value) + except ValidationError: + return None def _returned_text_parts(message: Mapping[str, object]) -> tuple[str, ...] | None: - content = message.get("content") + content: Final = message.get("content") if isinstance(content, str): return (content,) try: - parts: Final = _TEXT_PARTS.validate_python(message.get("content")) + parts: Final = _TEXT_PARTS.validate_python(content) except ValidationError: return None return tuple(part["text"] for part in parts) +def _merged_edits(children: Sequence[tuple[_TextEdit, ...] | None]) -> tuple[_TextEdit, ...] | None: + if any(child is None for child in children): + return None + return tuple(chain.from_iterable(child for child in children if child is not None)) + + +def _json_children(path: tuple[str | int, ...], original: object, returned: object) -> tuple[_JsonPair, ...] | None: + original_object: Final = _as_object(original) + returned_object: Final = _as_object(returned) + if original_object is not None and returned_object is not None: + if original_object.keys() != returned_object.keys(): + return None + return tuple(((*path, key), value, returned_object[key]) for key, value in original_object.items()) + original_items: Final = _as_items(original) + returned_items: Final = _as_items(returned) + if original_items is None or returned_items is None or len(original_items) != len(returned_items): + return None + return tuple(((*path, index), value, returned_items[index]) for index, value in enumerate(original_items)) + + def _json_edits( ref: _BlockRef, path: tuple[str | int, ...], original: object, returned: object ) -> tuple[_TextEdit, ...] | None: """Only write changed leaves, preserving independent flat-text redactions.""" - if original == returned: - return () - if isinstance(original, dict) and isinstance(returned, dict) and original.keys() == returned.keys(): - children = [_json_edits(ref, (*path, key), value, returned[key]) for key, value in original.items()] - elif isinstance(original, list) and isinstance(returned, list) and len(original) == len(returned): - children = [_json_edits(ref, (*path, index), value, returned[index]) for index, value in enumerate(original)] - elif isinstance(original, str) and isinstance(returned, str): - return ((_BlockRef(ref.message_index, ref.block_index, path), returned),) - else: + pending: Final[list[_JsonPair]] = [(path, original, returned)] # mutable-ok: work stack instead of recursion + edits: Final[list[_TextEdit]] = [] # mutable-ok: collects the leaves the work stack visits + while pending: + leaf_path, before, after = pending.pop() + if before == after: + continue + if isinstance(before, str) and isinstance(after, str): + edits.append((_BlockRef(ref.message_index, ref.block_index, leaf_path), after)) + continue + if (children := _json_children(leaf_path, before, after)) is None: + return None + pending.extend(children) + return tuple(edits) + + +def _as_items(value: object) -> Sequence[object] | None: + if not isinstance(value, list): return None - if any(child is None for child in children): + try: + return _ITEMS.validate_python(value) + except ValidationError: return None - return tuple(edit for child in children if child is not None for edit in child) + + +def _tool_part_edits( + ref: _BlockRef, index: int, part: _ConverseToolResultContent, text: str +) -> tuple[_TextEdit, ...] | None: + if part.text: + return ((_BlockRef(ref.message_index, ref.block_index, ("toolResult", "content", index, "text")), text),) + try: + value: Final = _PARSED_JSON.validate_json(text) + except ValidationError: + return None + return _json_edits(ref, ("toolResult", "content", index, "json"), part.json_value, value) def _tool_result_edits(converted: _ConvertedMessage, returned: Mapping[str, object]) -> tuple[_TextEdit, ...] | None: if converted.tool_result is None or _without_content(returned) != _without_content(converted.message): return None ref, result = converted.tool_result - texts = _returned_text_parts(returned) - # A joined string cannot be split safely across multiple source blocks. - if texts is None or len(texts) != len(result.content): + texts: Final = _returned_text_parts(returned) + parts: Final = _scannable_tool_parts(result) + if texts is None or len(texts) != len(parts): return None - edits: list[_TextEdit] = [] - for index, (part, text) in enumerate(zip(result.content, texts)): - if part.text: - if text != part.text: - edits.append( - (_BlockRef(ref.message_index, ref.block_index, ("toolResult", "content", index, "text")), text) - ) - elif part.json_value is not None: - try: - value = json.loads(text) - except ValueError: - return None - changes = _json_edits(ref, ("toolResult", "content", index, "json"), part.json_value, value) - if changes is None: - return None - edits.extend(changes) - else: - return None - return tuple(edits) + return _merged_edits(tuple(_tool_part_edits(ref, index, part, text) for (index, part), text in zip(parts, texts))) + + +def _call_function(call: object) -> Mapping[str, object] | None: + call_object: Final = _as_object(call) + return _as_object(call_object.get("function")) if call_object is not None else None + + +def _tool_call_edits( + tool_use: tuple[_BlockRef, _ConverseToolUse], before: object, after: object +) -> tuple[_TextEdit, ...] | None: + if before == after: + return () + before_call: Final = _as_object(before) + after_call: Final = _as_object(after) + before_function: Final = _call_function(before) + after_function: Final = _call_function(after) + if before_call is None or after_call is None or before_function is None or after_function is None: + return None + arguments: Final = after_function.get("arguments") + if ( + not isinstance(arguments, str) + or _without(after_call, "function") != _without(before_call, "function") + or _without(after_function, "arguments") != _without(before_function, "arguments") + ): + return None + try: + value: Final = _PARSED_JSON.validate_json(arguments) + except ValidationError: + return None + ref, tool = tool_use + return _json_edits(ref, ("toolUse", "input"), tool.tool_input, value) + + +def _tool_calls_edits( + tool_uses: Sequence[tuple[_BlockRef, _ConverseToolUse]], before: object, after: object +) -> tuple[_TextEdit, ...] | None: + before_calls: Final = _as_items(before) + after_calls: Final = _as_items(after) + if before_calls is None or after_calls is None or len(before_calls) != len(after_calls): + return None + return _merged_edits( + tuple( + _tool_call_edits(tool_use, original, returned) + for tool_use, original, returned in zip(tool_uses, before_calls, after_calls) + ) + ) + + +def _content_edits(converted: _ConvertedMessage, returned: Mapping[str, object]) -> tuple[_TextEdit, ...] | None: + original: Final = _fields(converted.message) + if returned.get("content") == original.get("content"): + return () + texts: Final = _returned_text_parts(returned) + before_texts: Final = _returned_text_parts(original) + if texts is None or before_texts is None or len(texts) != len(converted.text_blocks): + return None + return tuple((ref, text) for ref, before, text in zip(converted.text_blocks, before_texts, texts) if text != before) def _message_text_edits(converted: _ConvertedMessage, returned: Mapping[str, object]) -> tuple[_TextEdit, ...] | None: - if not isinstance(returned, Mapping): - return None if returned == converted.message: return () if converted.tool_result is not None: return _tool_result_edits(converted, returned) - edits: list[_TextEdit] = [] - original = converted.message - comparable = dict(returned) - if converted.tool_uses: - calls = comparable.get("tool_calls") - original_calls = original.get("tool_calls") - if not isinstance(calls, list) or not isinstance(original_calls, list) or len(calls) != len(original_calls): - return None - for (ref, tool), before, after in zip(converted.tool_uses, original_calls, calls): - if before == after: - continue - if not isinstance(after, dict) or not isinstance(after.get("function"), dict): - return None - expected = copy.deepcopy(before) - expected["function"]["arguments"] = after["function"].get("arguments") - if expected != after or not isinstance(after["function"].get("arguments"), str): - return None - try: - arguments = json.loads(after["function"]["arguments"]) - except ValueError: - return None - changes = _json_edits(ref, ("toolUse", "input"), tool.tool_input, arguments) - if changes is None: - return None - edits.extend(changes) - comparable["tool_calls"] = original_calls - if _without_content(comparable) != _without_content(original): + original_calls: Final = _fields(converted.message).get("tool_calls") + call_edits: Final = ( + _tool_calls_edits(converted.tool_uses, original_calls, returned.get("tool_calls")) + if converted.tool_uses + else () + ) + comparable: Final = {**returned, "tool_calls": original_calls} if converted.tool_uses else returned + if call_edits is None or _without_content(comparable) != _without_content(converted.message): return None - if comparable.get("content") != original.get("content"): - texts = _returned_text_parts(comparable) - before_texts = _returned_text_parts(original) - if texts is None or before_texts is None or len(texts) != len(converted.text_blocks): - return None - edits.extend( - (ref, text) for ref, before, text in zip(converted.text_blocks, before_texts, texts) if text != before - ) - return tuple(edits) + return _merged_edits((call_edits, _content_edits(converted, comparable))) + + +def _leaf(body: Mapping[str, Any], ref: _BlockRef) -> _StringHolder: + blocks: Final = body["system"] if ref.message_index is None else body["messages"][ref.message_index]["content"] + block: Final = blocks[ref.block_index] + if not ref.path: + return (block, "text") if block.get("text") else (block["guardContent"]["text"], "text") + return reduce(lambda node, key: node[key], ref.path[:-1], block), ref.path[-1] + + +def _leaf_identity(holder: _StringHolder) -> tuple[int, str | int]: + container, key = holder + return id(container), key + + +def _changed_leaves( + holders: Sequence[_StringHolder], texts: Sequence[str], guardrailed_texts: Sequence[str] +) -> frozenset[tuple[int, str | int]]: + return frozenset( + _leaf_identity(holder) for holder, before, after in zip(holders, texts, guardrailed_texts) if after != before + ) + + +def _texts_applied_view( + body: Mapping[str, Any], texts: Sequence[str], skip_system: bool, skip_tool: bool, scan_only_tool_results: bool +) -> tuple[_ConvertedMessage, ...]: + rewritten: Final = copy.deepcopy(dict(body)) + _, holders = _extract_converse_texts(rewritten, skip_system, skip_tool, scan_only_tool_results) + _write_back_texts(texts, holders) + return _converse_structured_messages(_ConverseRequest.model_validate(rewritten), not skip_tool) + + +def _same_content(expected: Mapping[str, object], returned: Mapping[str, object]) -> bool: + if expected.get("content") == returned.get("content"): + return True + expected_parts: Final = _returned_text_parts(expected) + returned_parts: Final = _returned_text_parts(returned) + return ( + expected_parts is not None + and returned_parts is not None + and "\n".join(expected_parts) == "\n".join(returned_parts) + ) + + +def _reproduces(rewritten: _ConvertedMessage, returned: Mapping[str, object] | None) -> bool: + """Whether the flat-text rewrite already yields the row the guardrail returned.""" + return ( + returned is not None + and _without_content(returned) == _without_content(rewritten.message) + and _same_content(rewritten.message, returned) + ) def _structured_text_edits( + leaf: Callable[[_BlockRef], _StringHolder], converted: Sequence[_ConvertedMessage], scoped_indices: Sequence[int], - returned: Sequence[Mapping[str, object]], -) -> tuple[_TextEdit, ...] | None: + returned: Sequence[object], + flat_changes: frozenset[tuple[int, str | int]], + flat_view: Callable[[], tuple[_ConvertedMessage, ...]], +) -> tuple[tuple[_StringHolder, object], ...] | None: + """ + Structured edits for the leaves the flat texts left unchanged. + + A row whose rewrite cannot be mapped back onto its Converse blocks is accepted only when + the flat-text rewrite already produces it. + """ if len(returned) != len(scoped_indices): return None - per_message: Final = tuple( - _message_text_edits(converted[index], message) for index, message in zip(scoped_indices, returned) + rows: Final = tuple(_as_object(message) for message in returned) + proposed: Final = tuple( + _message_text_edits(converted[index], row) if row is not None else None + for index, row in zip(scoped_indices, rows) ) - if any(edits is None for edits in per_message): - return None - return tuple(chain.from_iterable(edits for edits in per_message if edits is not None)) + unmapped: Final = tuple((index, row) for index, row, edits in zip(scoped_indices, rows, proposed) if edits is None) + if unmapped: + view: Final = flat_view() + if len(view) != len(converted): + return None + if not all(_reproduces(view[index], row) for index, row in unmapped): + return None + resolved: Final = tuple( + (leaf(ref), text) for ref, text in chain.from_iterable(edits for edits in proposed if edits is not None) + ) + return tuple((holder, text) for holder, text in resolved if _leaf_identity(holder) not in flat_changes) -def _write_back_structured_texts(body: Mapping[str, Any], edits: Sequence[_TextEdit]) -> None: - for ref, text in edits: - blocks = body["system"] if ref.message_index is None else body["messages"][ref.message_index]["content"] - block = blocks[ref.block_index] - if ref.path: - target = block - for key in ref.path[:-1]: - target = target[key] - target[ref.path[-1]] = text - elif "text" in block: - block["text"] = text - else: - block["guardContent"]["text"]["text"] = text +def _write_back_structured_texts(edits: Sequence[tuple[_StringHolder, object]]) -> None: + for (container, key), text in edits: + container[key] = text def _converse_tools(request: _ConverseRequest) -> tuple[ChatCompletionToolParam, ...]: @@ -555,8 +696,8 @@ def _extract_converse_output_texts( def _write_back_texts( - guardrailed_texts: list[str], - holders: list[_StringHolder], + guardrailed_texts: Sequence[str], + holders: Sequence[_StringHolder], ) -> None: if len(guardrailed_texts) < len(holders): verbose_proxy_logger.warning( @@ -827,7 +968,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply) texts, holders = _extract_converse_texts(body, skip_system, skip_tool, scan_only_tool_results) converse_request: Final = _parse_converse_request(lambda: _ConverseRequest.model_validate(body)) - converted: Final = _converse_structured_messages(converse_request) if converse_request else () + converted: Final = _converse_structured_messages(converse_request, not skip_tool) if converse_request else () structured_messages: Final = tuple(item.message for item in converted) scoped_message_indices: Final = scoped_structured_message_indices( structured_messages, @@ -839,12 +980,12 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): if not texts and not scoped_message_indices: return data - inputs: Final = GenericGuardrailAPIInputs(texts=texts) + inputs: Final = GenericGuardrailAPIInputs(texts=list(texts)) if scoped_message_indices: inputs["structured_messages"] = copy.deepcopy( [structured_messages[index] for index in scoped_message_indices] ) - tools: Final = _converse_tools(converse_request) if converse_request else () + tools: Final = _converse_tools(converse_request) if converse_request and not scan_only_tool_results else () if tools: inputs["tools"] = list(tools) model: Final = data.get("model") @@ -859,16 +1000,33 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) - edits: tuple[_TextEdit, ...] = () guardrailed_messages: Final = guardrailed_inputs.get("structured_messages") - if guardrailed_messages is not None: - proposed_edits = _structured_text_edits(converted, scoped_message_indices, guardrailed_messages) - if proposed_edits is None: - raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) - edits = proposed_edits - if guardrailed_texts: - _write_back_texts(guardrailed_texts, holders) - _write_back_structured_texts(body, edits) + sent_messages: Final = [structured_messages[index] for index in scoped_message_indices] + rewritten_rows: Final = ( + guardrailed_messages + if sent_messages and isinstance(guardrailed_messages, list) and guardrailed_messages != sent_messages + else None + ) + applied_texts: Final = ( + texts if rewritten_rows is not None and len(guardrailed_texts) != len(texts) else guardrailed_texts + ) + edits: Final = ( + () + if rewritten_rows is None + else _structured_text_edits( + partial(_leaf, body), + converted, + scoped_message_indices, + rewritten_rows, + _changed_leaves(holders, texts, applied_texts), + partial(_texts_applied_view, body, applied_texts, skip_system, skip_tool, scan_only_tool_results), + ) + ) + if edits is None: + raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) + if applied_texts: + _write_back_texts(applied_texts, holders) + _write_back_structured_texts(edits) return data diff --git a/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py b/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py new file mode 100644 index 00000000000..46f87cad7c9 --- /dev/null +++ b/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py @@ -0,0 +1,451 @@ +from __future__ import annotations + +import json +import re +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from pydantic import JsonValue, TypeAdapter + +from tests.integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment, list_value, object_value +from tests.integration._support.process import owned_proxy +from tests.integration._support.wire import Reply, Request, Wire, wire_server + +_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" +_MODEL: Final = "converse-guarded" +_CARD: Final = "4111 1111 1111 1111" +_CARD_PATTERN: Final = re.compile(r"\b(?:\d{4} ){3}\d{4}\b") +_MASK: Final = "****" +_FLAGGED: Final = "flagged-by-moderation" +_GUARD_PATH: Final = "/beta/litellm_basic_guardrail_api" +_PNG: Final = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=" +_JSON: Final = TypeAdapter(JsonValue) +_UNAPPLIABLE: Final = "rewrote the request in a way this endpoint cannot apply" +_GENERIC_MODES: Final = ( + "masksm", + "masksm_flat", + "smonly", + "sm_partial", + "echo_sm", + "mask", + "block", + "rowcount", +) + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + vendor: Wire + bedrock: Wire + keys: Mapping[str, str] + + +def _masked(value: JsonValue) -> JsonValue: + return _JSON.validate_json(_CARD_PATTERN.sub(_MASK, json.dumps(value))) + + +def _flattened(row: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + content: Final = row.get("content") + if not isinstance(content, list): + return dict(row) + parts: Final = tuple(object_value(part).get("text") for part in content) + return {**row, "content": "\n".join(part for part in parts if isinstance(part, str))} + + +def _normalised_first_part(row: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + content: Final = row.get("content") + if not isinstance(content, list) or not content: + return dict(row) + first: Final = object_value(content[0]) + text: Final = first.get("text") + if not isinstance(text, str): + return dict(row) + return {**row, "content": [{**first, "text": text.strip() + " "}, *content[1:]]} + + +def _generic_verdict(mode: str, body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + texts: Final = list_value(body.get("texts") or []) + rows: Final = tuple(object_value(row) for row in list_value(body.get("structured_messages") or [])) + masked_texts: Final = _masked(list(texts)) + masked_rows: Final = list_value(_masked([dict(row) for row in rows])) + if _CARD_PATTERN.search(json.dumps(body)) is None: + return {"action": "NONE"} + verdicts: Final[Mapping[str, Callable[[], dict[str, JsonValue]]]] = { + "masksm": lambda: {"action": "GUARDRAIL_INTERVENED", "texts": masked_texts, "structured_messages": masked_rows}, + "masksm_flat": lambda: { + "action": "GUARDRAIL_INTERVENED", + "texts": masked_texts, + "structured_messages": [_flattened(object_value(row)) for row in masked_rows], + }, + "smonly": lambda: {"action": "GUARDRAIL_INTERVENED", "structured_messages": masked_rows}, + "sm_partial": lambda: { + "action": "GUARDRAIL_INTERVENED", + "texts": masked_texts, + "structured_messages": [_normalised_first_part(row) for row in rows], + }, + "echo_sm": lambda: { + "action": "GUARDRAIL_INTERVENED", + "texts": masked_texts, + "structured_messages": [dict(row) for row in rows], + }, + "mask": lambda: {"action": "GUARDRAIL_INTERVENED", "texts": masked_texts}, + "block": lambda: {"action": "BLOCKED", "blocked_reason": "card number"}, + "rowcount": lambda: { + "action": "GUARDRAIL_INTERVENED", + "texts": masked_texts, + "structured_messages": masked_rows[:-1], + }, + } + return verdicts[mode]() + + +def _vendor(request: Request) -> Reply: + body: Final = JSON_OBJECT.validate_json(request.body) + if request.target == "/moderation/moderations": + flagged: Final = _FLAGGED in str(body.get("input")) + result: Final = { + "flagged": flagged, + "categories": {"violence": flagged}, + "category_scores": {"violence": 0.99 if flagged else 0.0}, + "category_applied_input_types": {"violence": ["text"]}, + } + return Reply( + body=json.dumps({"id": "modr-scripted", "model": "omni-moderation-latest", "results": [result]}).encode() + ) + if request.target == "/promptguard/api/v1/guard": + messages: Final = body.get("messages") + if _CARD_PATTERN.search(json.dumps(messages)) is None: + return Reply(body=json.dumps({"decision": "allow"}).encode()) + return Reply(body=json.dumps({"decision": "redact", "redacted_messages": _masked(messages)}).encode()) + mode: Final = request.target.removeprefix("/gga/").removesuffix(_GUARD_PATH) + return Reply(body=json.dumps(_generic_verdict(mode, body)).encode()) + + +def _bedrock(request: Request) -> Reply: + assert request.target == f"/model/{_MODEL_ID}/converse", request.target + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "guarded reply"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 7, "outputTokens": 2, "totalTokens": 9}, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + + +def _guardrail(name: str, params: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {"guardrail_name": name, "litellm_params": {"mode": "pre_call", "default_on": False, **params}} + + +def _config(directory: Path, vendor: Wire, bedrock: Wire) -> Path: + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["model_list"] = [ + *configuration["model_list"], + { + "model_name": _MODEL, + "litellm_params": { + "model": f"bedrock/{_MODEL_ID}", + "api_base": bedrock.url, + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + }, + }, + ] + generic: Final = {"guardrail": "generic_guardrail_api", "api_key": "synthetic-guardrail-key"} + configuration["guardrails"] = [ + *(_guardrail(f"gga-{mode}", {**generic, "api_base": f"{vendor.url}/gga/{mode}"}) for mode in _GENERIC_MODES), + _guardrail( + "gga-block-skip-tool", + {**generic, "api_base": f"{vendor.url}/gga/block", "skip_tool_message_in_guardrail": True}, + ), + _guardrail( + "gga-mask-tool-only", + {**generic, "api_base": f"{vendor.url}/gga/mask", "scan_only_tool_results": True}, + ), + _guardrail( + "promptguard", + {"guardrail": "promptguard", "api_base": f"{vendor.url}/promptguard", "api_key": "synthetic-pg-key"}, + ), + _guardrail( + "moderation", + {"guardrail": "openai_moderation", "api_base": f"{vendor.url}/moderation", "api_key": "synthetic-key"}, + ), + ] + path: Final = directory / "converse-guardrails.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("converse-guardrail-rewrites") + with ( + gateway_from_environment() as gateway, + wire_server(_vendor) as vendor, + wire_server(_bedrock) as bedrock, + owned_proxy(gateway, directory, {}, config=_config(directory, vendor, bedrock), workers=2) as owned, + owned.scenario() as scenario, + ): + names: Final = ( + *(f"gga-{mode}" for mode in _GENERIC_MODES), + "gga-block-skip-tool", + "gga-mask-tool-only", + "promptguard", + "moderation", + ) + yield _Rig(owned, vendor, bedrock, {name: scenario.key(guardrails=[name]) for name in names}) + + +@pytest.fixture(autouse=True) +def _drained(rig: _Rig) -> Iterator[None]: + rig.vendor.drain() + rig.bedrock.drain() + yield + rig.vendor.drain() + rig.bedrock.drain() + + +def _converse(rig: _Rig, guardrail: str, body: Mapping[str, JsonValue]) -> httpx.Response: + return rig.gateway.request("POST", f"/bedrock/model/{_MODEL}/converse", body, key=rig.keys[guardrail]) + + +def _forwarded(rig: _Rig, response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text + sent: Final = rig.bedrock.drain() + assert len(sent) == 1, response.text + return JSON_OBJECT.validate_json(sent[0].body) + + +def _refused(rig: _Rig, response: httpx.Response) -> str: + assert response.status_code >= 400, response.text + assert rig.bedrock.drain() == (), response.text + return response.text + + +def _vendor_bodies(rig: _Rig) -> tuple[dict[str, JsonValue], ...]: + return tuple(JSON_OBJECT.validate_json(request.body) for request in rig.vendor.drain()) + + +def _user(*blocks: JsonValue) -> dict[str, JsonValue]: + return {"role": "user", "content": list(blocks)} + + +def _assistant(*blocks: JsonValue) -> dict[str, JsonValue]: + return {"role": "assistant", "content": list(blocks)} + + +def _tool_use(card: str) -> dict[str, JsonValue]: + return {"toolUse": {"toolUseId": "tooluse_lookup", "name": "lookup_card", "input": {"card": card, "limit": 2}}} + + +def _tool_result(*content: JsonValue) -> dict[str, JsonValue]: + return {"toolResult": {"toolUseId": "tooluse_lookup", "content": list(content), "status": "success"}} + + +def test_multi_block_user_message_with_flattened_vendor_reply_reaches_bedrock_masked(rig: _Rig) -> None: + body: Final = {"messages": [_user({"text": "Order details below."}, {"text": f"Card {_CARD} please."})]} + sent: Final = _forwarded(rig, _converse(rig, "gga-masksm_flat", body)) + assert sent["messages"] == [_user({"text": "Order details below."}, {"text": f"Card {_MASK} please."})] + + +def test_multi_block_system_prompt_with_flattened_vendor_reply_reaches_bedrock_masked(rig: _Rig) -> None: + body: Final = { + "system": [{"text": "You are a billing assistant."}, {"text": f"Card on file {_CARD}."}], + "messages": [_user({"text": "What card is on file?"})], + } + sent: Final = _forwarded(rig, _converse(rig, "gga-masksm_flat", body)) + assert sent["system"] == [{"text": "You are a billing assistant."}, {"text": f"Card on file {_MASK}."}] + assert sent["messages"] == body["messages"] + + +def test_tool_result_with_text_and_json_parts_reaches_bedrock_masked(rig: _Rig) -> None: + body: Final = { + "messages": [ + _user({"text": "Look up my card."}), + _assistant(_tool_use("last four 1111")), + _user(_tool_result({"text": f"found {_CARD}"}, {"json": {"card": _CARD, "status": "active"}})), + ] + } + sent: Final = _forwarded(rig, _converse(rig, "gga-masksm", body)) + assert sent["messages"] == [ + _user({"text": "Look up my card."}), + _assistant(_tool_use("last four 1111")), + _user(_tool_result({"text": f"found {_MASK}"}, {"json": {"card": _MASK, "status": "active"}})), + ] + + +def test_tool_result_with_text_and_image_parts_reaches_bedrock_masked_with_image_intact(rig: _Rig) -> None: + image: Final = {"image": {"format": "png", "source": {"bytes": _PNG}}} + body: Final = { + "messages": [ + _user({"text": "Read the screenshot."}), + _assistant(_tool_use("screenshot")), + _user(_tool_result({"text": f"screenshot shows {_CARD}"}, image)), + ] + } + sent: Final = _forwarded(rig, _converse(rig, "gga-masksm", body)) + assert sent["messages"] == [ + _user({"text": "Read the screenshot."}), + _assistant(_tool_use("screenshot")), + _user(_tool_result({"text": f"screenshot shows {_MASK}"}, image)), + ] + + +def test_flat_text_mask_wins_over_an_unmasked_structured_leaf(rig: _Rig) -> None: + body: Final = {"messages": [_user({"text": f"Card {_CARD} on file."}, {"text": "Thanks."})]} + sent: Final = _forwarded(rig, _converse(rig, "gga-sm_partial", body)) + assert sent["messages"] == [_user({"text": f"Card {_MASK} on file."}, {"text": "Thanks."})] + + +def test_promptguard_redaction_keeps_every_turn_in_place(rig: _Rig) -> None: + body: Final = { + "system": [{"text": "You are a support bot."}], + "messages": [ + _user({"text": "My name is Ann."}), + _assistant({"text": "Hi Ann, how can I help?"}), + _user({"text": f"Repeat exactly: card {_CARD} ok"}), + ], + } + sent: Final = _forwarded(rig, _converse(rig, "promptguard", body)) + assert sent["system"] == [{"text": "You are a support bot."}] + assert sent["messages"] == [ + _user({"text": "My name is Ann."}), + _assistant({"text": "Hi Ann, how can I help?"}), + _user({"text": f"Repeat exactly: card {_MASK} ok"}), + ] + + +def test_moderation_blocks_a_flagged_latest_user_turn_before_bedrock(rig: _Rig) -> None: + body: Final = { + "system": [{"text": "You are a support bot."}], + "messages": [ + _user({"text": "Say hi."}), + _assistant({"text": "Hi."}), + _user({"text": f"please {_FLAGGED} now"}), + ], + } + refused: Final = _refused(rig, _converse(rig, "moderation", body)) + assert "Violated OpenAI moderation policy" in refused + + +def test_skip_tool_messages_keeps_tool_call_arguments_away_from_the_guardrail(rig: _Rig) -> None: + body: Final = { + "messages": [ + _user({"text": "Check my card."}), + _assistant(_tool_use(_CARD)), + _user(_tool_result({"text": "card is active"})), + ] + } + sent: Final = _forwarded(rig, _converse(rig, "gga-block-skip-tool", body)) + assert sent["messages"] == body["messages"] + assert all(_CARD not in json.dumps(scanned) for scanned in _vendor_bodies(rig)) + + +def test_tool_only_guardrail_receives_no_tool_definitions(rig: _Rig) -> None: + body: Final = { + "messages": [ + _user({"text": f"My card is {_CARD}."}), + _assistant(_tool_use("lookup")), + _user(_tool_result({"text": f"found {_CARD}"})), + ], + "toolConfig": { + "tools": [ + { + "toolSpec": { + "name": "lookup_card", + "description": "Look up a card", + "inputSchema": {"json": {"type": "object", "properties": {"card": {"type": "string"}}}}, + } + } + ] + }, + } + sent: Final = _forwarded(rig, _converse(rig, "gga-mask-tool-only", body)) + scanned: Final = _vendor_bodies(rig) + assert len(scanned) == 1 + assert scanned[0].get("tools") is None, scanned[0] + assert scanned[0]["texts"] == [f"found {_CARD}"] + assert sent["messages"] == [ + _user({"text": f"My card is {_CARD}."}), + _assistant(_tool_use("lookup")), + _user(_tool_result({"text": f"found {_MASK}"})), + ] + + +def test_structured_only_vendor_mask_reaches_bedrock(rig: _Rig) -> None: + body: Final = {"messages": [_user({"text": f"Card {_CARD} please."})]} + sent: Final = _forwarded(rig, _converse(rig, "gga-smonly", body)) + assert sent["messages"] == [_user({"text": f"Card {_MASK} please."})] + + +def test_guard_content_turn_reaches_bedrock_masked(rig: _Rig) -> None: + body: Final = {"messages": [_user({"text": "Check this."}, {"guardContent": {"text": {"text": f"Card {_CARD}."}}})]} + sent: Final = _forwarded(rig, _converse(rig, "gga-masksm", body)) + assert sent["messages"] == [_user({"text": "Check this."}, {"guardContent": {"text": {"text": f"Card {_MASK}."}}})] + + +def test_tool_call_arguments_in_history_reach_bedrock_masked(rig: _Rig) -> None: + body: Final = { + "messages": [ + _user({"text": "Check my card."}), + _assistant(_tool_use(_CARD)), + _user(_tool_result({"text": "card is active"})), + ] + } + sent: Final = _forwarded(rig, _converse(rig, "gga-masksm", body)) + assert sent["messages"] == [ + _user({"text": "Check my card."}), + _assistant(_tool_use(_MASK)), + _user(_tool_result({"text": "card is active"})), + ] + + +def test_echoed_structured_rows_keep_the_text_mask(rig: _Rig) -> None: + body: Final = {"messages": [_user({"text": f"Card {_CARD} please."})]} + sent: Final = _forwarded(rig, _converse(rig, "gga-echo_sm", body)) + assert sent["messages"] == [_user({"text": f"Card {_MASK} please."})] + + +def test_single_block_string_reply_reaches_bedrock_masked(rig: _Rig) -> None: + body: Final = {"messages": [_user({"text": f"Card {_CARD} please."})]} + sent: Final = _forwarded(rig, _converse(rig, "gga-masksm_flat", body)) + assert sent["messages"] == [_user({"text": f"Card {_MASK} please."})] + + +def test_missing_structured_row_is_refused_before_bedrock(rig: _Rig) -> None: + body: Final = {"messages": [_user({"text": "Hello."}), _assistant({"text": "Hi."}), _user({"text": f"{_CARD}"})]} + assert _UNAPPLIABLE in _refused(rig, _converse(rig, "gga-rowcount", body)) + + +def test_blocked_verdict_never_reaches_bedrock(rig: _Rig) -> None: + body: Final = {"messages": [_user({"text": f"Card {_CARD} please."})]} + _refused(rig, _converse(rig, "gga-block", body)) + + +def test_concurrent_requests_keep_their_own_rewrites(rig: _Rig) -> None: + cards: Final = tuple(f"4111 1111 1111 {index:04d}" for index in range(24)) + bodies: Final = tuple( + {"messages": [_user({"text": f"request {index}"}, {"text": f"card {card}"})]} + for index, card in enumerate(cards) + ) + with ThreadPoolExecutor(max_workers=12) as pool: + responses: Final = tuple(pool.map(lambda body: _converse(rig, "gga-masksm_flat", body), bodies)) + assert [response.status_code for response in responses] == [200] * len(bodies) + sent: Final = sorted( + (JSON_OBJECT.validate_json(request.body)["messages"] for request in rig.bedrock.drain()), key=json.dumps + ) + expected: Final = sorted( + ([_user({"text": f"request {index}"}, {"text": f"card {_MASK}"})] for index in range(len(cards))), + key=json.dumps, + ) + assert sent == expected diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index 8dc0ee2c9a5..3da48e5d68c 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -1650,6 +1650,9 @@ class TestConverseRewriteRegressions: }, ], "additionalModelRequestFields": {"prompt": "extra"}, + "toolConfig": { + "tools": [{"toolSpec": {"name": "search", "inputSchema": {"json": {"type": "object"}}}}] + }, }, } guard = _make_guardrail({}) @@ -1658,6 +1661,7 @@ class TestConverseRewriteRegressions: inputs = guard.apply_guardrail.call_args.kwargs["inputs"] assert inputs["texts"] == ["result"] assert inputs["structured_messages"] == [{"role": "tool", "tool_call_id": "t1", "content": "result"}] + assert "tools" not in inputs @pytest.mark.asyncio async def test_ambiguous_string_over_multiple_blocks_fails_closed_without_mutation(self): @@ -1781,6 +1785,303 @@ class TestConverseRewriteRegressions: assert result == data guard.apply_guardrail.assert_not_called() + @pytest.mark.asyncio + async def test_tool_result_image_part_is_skipped_when_mapping_the_text_part(self): + image = {"image": {"format": "png", "source": {"bytes": "iVBORw0KGgo="}}} + data = { + "endpoint": "model/test/converse", + "data": { + "messages": [ + { + "role": "user", + "content": [{"toolResult": {"toolUseId": "t1", "content": [{"text": "card 4111"}, image]}}], + } + ] + }, + } + guard = _make_guardrail( + {"structured_messages": [{"role": "tool", "tool_call_id": "t1", "content": "card [CARD]"}]} + ) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert result["data"]["messages"][0]["content"] == [ + {"toolResult": {"toolUseId": "t1", "content": [{"text": "card [CARD]"}, image]}} + ] + + @pytest.mark.asyncio + async def test_joined_tool_result_row_is_accepted_when_the_flat_texts_reproduce_it(self): + data = { + "endpoint": "model/test/converse", + "data": { + "messages": [ + { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "t1", + "content": [{"text": "found 4111"}, {"json": {"card": "4111"}}], + } + } + ], + } + ] + }, + } + guard = _make_guardrail( + { + "texts": ["found [CARD]", "[CARD]"], + "structured_messages": [ + {"role": "tool", "tool_call_id": "t1", "content": 'found [CARD]\n{"card": "[CARD]"}'} + ], + } + ) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert result["data"]["messages"][0]["content"] == [ + {"toolResult": {"toolUseId": "t1", "content": [{"text": "found [CARD]"}, {"json": {"card": "[CARD]"}}]}} + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("role", ["system", "user"]) + async def test_flattened_multi_block_row_is_accepted_when_the_flat_texts_reproduce_it(self, role): + blocks = [{"text": "Order details."}, {"text": "card 4111"}] + body = ( + {"system": blocks, "messages": [{"role": "user", "content": [{"text": "hi"}]}]} + if role == "system" + else {"messages": [{"role": "user", "content": blocks}]} + ) + data = {"endpoint": "model/test/converse", "data": body} + rows = [{"role": role, "content": "Order details.\ncard [CARD]"}] + texts = ["Order details.", "card [CARD]"] + if role == "system": + rows.append({"role": "user", "content": [{"type": "text", "text": "hi"}]}) + texts.append("hi") + guard = _make_guardrail({"texts": texts, "structured_messages": rows}) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + written = result["data"]["system"] if role == "system" else result["data"]["messages"][0]["content"] + assert written == [{"text": "Order details."}, {"text": "card [CARD]"}] + + @pytest.mark.asyncio + async def test_flat_text_mask_wins_over_a_changed_but_unmasked_structured_leaf(self): + data = { + "endpoint": "model/test/converse", + "data": {"messages": [{"role": "user", "content": [{"text": "card 4111 on file."}, {"text": "Alice"}]}]}, + } + guard = _make_guardrail( + { + "texts": ["card [CARD] on file.", "Alice"], + "structured_messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "card 4111 on file. "}, + {"type": "text", "text": "[NAME]"}, + ], + } + ], + } + ) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert result["data"]["messages"][0]["content"] == [{"text": "card [CARD] on file."}, {"text": "[NAME]"}] + + @pytest.mark.asyncio + async def test_texts_that_do_not_line_up_with_the_request_are_ignored_next_to_structured_rows(self): + data = { + "endpoint": "model/test/converse", + "data": { + "system": [{"text": "You are a support bot."}], + "messages": [ + {"role": "user", "content": [{"text": "My name is Ann."}]}, + {"role": "assistant", "content": [{"text": "Hi Ann."}]}, + {"role": "user", "content": [{"text": "card 4111 ok"}]}, + ], + }, + } + guard = _make_guardrail( + { + "texts": ["My name is Ann.", "card [CARD] ok"], + "structured_messages": [ + {"role": "system", "content": [{"type": "text", "text": "You are a support bot."}]}, + {"role": "user", "content": [{"type": "text", "text": "My name is Ann."}]}, + {"role": "assistant", "content": [{"type": "text", "text": "Hi Ann."}]}, + {"role": "user", "content": [{"type": "text", "text": "card [CARD] ok"}]}, + ], + } + ) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert result["data"]["system"] == [{"text": "You are a support bot."}] + assert result["data"]["messages"] == [ + {"role": "user", "content": [{"text": "My name is Ann."}]}, + {"role": "assistant", "content": [{"text": "Hi Ann."}]}, + {"role": "user", "content": [{"text": "card [CARD] ok"}]}, + ] + + @pytest.mark.asyncio + async def test_structured_messages_that_are_not_rows_fall_back_to_texts(self): + data = { + "endpoint": "model/test/converse", + "data": {"messages": [{"role": "user", "content": [{"text": "4111"}]}]}, + } + guard = _make_guardrail({"texts": ["[CARD]"], "structured_messages": {"evaluation": {"action": "Redact"}}}) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert result["data"]["messages"][0]["content"] == [{"text": "[CARD]"}] + + @pytest.mark.asyncio + async def test_skip_tool_keeps_tool_call_arguments_out_of_structured_rows(self): + data = { + "endpoint": "model/test/converse", + "data": { + "messages": [ + { + "role": "assistant", + "content": [ + {"text": "Checking."}, + {"toolUse": {"toolUseId": "t1", "name": "pay", "input": {"card": "4111"}}}, + ], + }, + {"role": "user", "content": [{"text": "thanks"}]}, + ] + }, + } + guard = _make_guardrail({}) + guard.skip_tool_message_in_guardrail = True + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert guard.apply_guardrail.call_args.kwargs["inputs"]["structured_messages"] == [ + {"role": "assistant", "content": [{"type": "text", "text": "Checking."}]}, + {"role": "user", "content": [{"type": "text", "text": "thanks"}]}, + ] + + @pytest.mark.asyncio + async def test_unchanged_structured_rows_keep_the_flat_text_write_back(self): + data = { + "endpoint": "model/test/converse", + "data": {"messages": [{"role": "user", "content": [{"text": "4111"}, {"text": "Alice"}]}]}, + } + guard = _make_guardrail( + { + "texts": ["[CARD]"], + "structured_messages": [ + {"role": "user", "content": [{"type": "text", "text": "4111"}, {"type": "text", "text": "Alice"}]} + ], + } + ) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert result["data"]["messages"][0]["content"] == [{"text": "[CARD]"}, {"text": "Alice"}] + + @pytest.mark.asyncio + async def test_rows_returned_when_none_were_sent_are_ignored(self): + data = { + "endpoint": "model/test/converse", + "data": {"messages": [{"role": "user", "content": [{"text": "4111"}, "not a block"]}]}, + } + guard = _make_guardrail({"texts": ["[CARD]"], "structured_messages": [{"role": "user", "content": "[CARD]"}]}) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert "structured_messages" not in guard.apply_guardrail.call_args.kwargs["inputs"] + assert result["data"]["messages"][0]["content"] == [{"text": "[CARD]"}, "not a block"] + + @pytest.mark.asyncio + async def test_missing_structured_row_fails_closed_without_mutation(self): + data = { + "endpoint": "model/test/converse", + "data": { + "messages": [ + {"role": "user", "content": [{"text": "hello"}]}, + {"role": "assistant", "content": [{"text": "hi"}]}, + {"role": "user", "content": [{"text": "4111"}]}, + ] + }, + } + original = copy.deepcopy(data) + guard = _make_guardrail( + { + "texts": ["hello", "hi", "[CARD]"], + "structured_messages": [ + {"role": "user", "content": [{"type": "text", "text": "hello"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "hi"}]}, + ], + } + ) + with pytest.raises(UnappliableRequestRewrite): + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert data == original + + @pytest.mark.asyncio + async def test_flattened_row_after_a_turn_the_texts_emptied_fails_closed(self): + data = { + "endpoint": "model/test/converse", + "data": { + "messages": [ + {"role": "user", "content": [{"text": "hello"}]}, + {"role": "assistant", "content": [{"text": "hi"}]}, + {"role": "user", "content": [{"text": "card"}, {"text": "4111"}]}, + ] + }, + } + guard = _make_guardrail( + { + "texts": ["hello", "", "card", "[CARD]"], + "structured_messages": [ + {"role": "user", "content": [{"type": "text", "text": "hello"}]}, + {"role": "assistant", "content": [{"type": "text", "text": ""}]}, + {"role": "user", "content": "card\n[CARD]"}, + ], + } + ) + with pytest.raises(UnappliableRequestRewrite): + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + + @pytest.mark.asyncio + async def test_joined_tool_result_row_the_flat_texts_do_not_reproduce_fails_closed(self): + data = { + "endpoint": "model/test/converse", + "data": { + "messages": [ + { + "role": "user", + "content": [ + {"toolResult": {"toolUseId": "t1", "content": [{"text": "found 4111"}, {"text": "ok"}]}} + ], + } + ] + }, + } + original = copy.deepcopy(data) + guard = _make_guardrail( + {"structured_messages": [{"role": "tool", "tool_call_id": "t1", "content": "found [CARD]\nok"}]} + ) + with pytest.raises(UnappliableRequestRewrite): + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert data == original + + @pytest.mark.asyncio + async def test_flat_mask_written_into_the_sent_texts_list_still_wins(self): + data = { + "endpoint": "model/test/converse", + "data": {"messages": [{"role": "user", "content": [{"text": "ssn 123 email a@b.c"}]}]}, + } + + def mask_in_place(inputs: dict, **_: object) -> dict: + inputs["texts"][0] = "ssn [SSN] email a@b.c" + rows = copy.deepcopy(inputs["structured_messages"]) + rows[0]["content"][0]["text"] = "ssn 123 email [EMAIL]" + return {**inputs, "structured_messages": rows} + + guard = _make_guardrail({}) + guard.apply_guardrail = AsyncMock(side_effect=lambda inputs, **kwargs: mask_in_place(inputs)) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert result["data"]["messages"][0]["content"] == [{"text": "ssn [SSN] email a@b.c"}] + + @pytest.mark.asyncio + async def test_snake_case_block_keys_are_not_read_as_converse_fields(self): + data = { + "endpoint": "model/test/converse", + "data": {"messages": [{"role": "user", "content": [{"guard_content": {"text": {"text": "4111"}}}]}]}, + } + original = copy.deepcopy(data) + guard = _make_guardrail({"structured_messages": [{"role": "user", "content": "[CARD]"}]}) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert result == original + guard.apply_guardrail.assert_not_called() + async def _run_post_call_guardrail_on_stream( stream_bytes: bytes, guardrail: CustomGuardrail, monkeypatch: pytest.MonkeyPatch