diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 912936a8ecd..89920dba9b3 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -1,13 +1,34 @@ -from collections.abc import Mapping, Sequence +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 +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, + scoped_structured_message_indices, + unappliable_request_rewrite, +) +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionAssistantMessage, + ChatCompletionAssistantToolCall, + ChatCompletionSystemMessage, + ChatCompletionTextObject, + ChatCompletionToolCallFunctionChunk, + ChatCompletionToolMessage, + ChatCompletionToolParam, + ChatCompletionToolParamFunctionChunk, + ChatCompletionUserMessage, ) from litellm.types.utils import GenericGuardrailAPIInputs @@ -83,6 +104,7 @@ def _extract_converse_texts( body: dict, skip_system: bool, skip_tool: bool, + scan_only_tool_results: bool = False, ) -> tuple[list[str], list[_StringHolder]]: """ Walk a Bedrock Converse request body and collect text content. @@ -102,8 +124,8 @@ def _extract_converse_texts( inflating guardrail usage for a request whose only prompt is one user message. No other guardrail translation handler puts tool definitions in ``texts``; the chat and messages handlers carry them in the structured - ``tools`` input instead, which this handler does not populate because a - Bedrock ``toolSpec`` is not the OpenAI tool shape those consumers expect. + ``tools`` input instead, which this handler also populates by converting + each Bedrock ``toolSpec`` to the OpenAI tool shape. ``additionalModelRequestFields`` is treated differently on purpose. Bedrock gives ``toolConfig.tools`` a fixed schema whose contents are tool metadata by @@ -113,7 +135,7 @@ def _extract_converse_texts( """ holders: Final[list[_StringHolder]] = [] - if not skip_system: + if not skip_system and not scan_only_tool_results: for block in body.get("system") or []: if isinstance(block, dict): _collect_block_text(block, holders) @@ -126,6 +148,8 @@ def _extract_converse_texts( continue if skip_tool and ("toolUse" in block or "toolResult" in block): continue + if scan_only_tool_results and "toolResult" not in block: + continue _collect_block_text(block, holders) tool_use = block.get("toolUse") if isinstance(tool_use, dict): @@ -137,12 +161,507 @@ def _extract_converse_texts( _collect_block_text(inner, holders) _collect_strings(inner.get("json"), holders) - _collect_strings(body.get("additionalModelRequestFields"), holders) + if not scan_only_tool_results: + _collect_strings(body.get("additionalModelRequestFields"), holders) texts: Final = [container[key] for container, key in holders] return texts, holders +class _ConverseModel(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + +class _ConverseGuardText(_ConverseModel): + text: str = "" + + +class _ConverseGuardContent(_ConverseModel): + text: _ConverseGuardText | None = None + + +class _ConverseToolUse(_ConverseModel): + tool_use_id: str = Field(default="", alias="toolUseId") + name: str = "" + tool_input: object = Field(default=None, alias="input") + + +class _ConverseToolResultContent(_ConverseModel): + text: str = "" + json_value: object = Field(default=None, alias="json") + + +class _ConverseToolResult(_ConverseModel): + tool_use_id: str = Field(default="", alias="toolUseId") + content: tuple[_ConverseToolResultContent, ...] = () + + +class _ConverseContentBlock(_ConverseModel): + text: str = "" + guard_content: _ConverseGuardContent | None = Field(default=None, alias="guardContent") + tool_use: _ConverseToolUse | None = Field(default=None, alias="toolUse") + tool_result: _ConverseToolResult | None = Field(default=None, alias="toolResult") + + +class _ConverseMessage(_ConverseModel): + role: str = "" + content: tuple[_ConverseContentBlock, ...] = () + + +class _ConverseToolInputSchema(_ConverseModel): + json_schema: Mapping[str, object] | None = Field(default=None, alias="json") + + +class _ConverseToolSpec(_ConverseModel): + name: str = "" + description: str = "" + input_schema: _ConverseToolInputSchema | None = Field(default=None, alias="inputSchema") + + +class _ConverseTool(_ConverseModel): + tool_spec: _ConverseToolSpec | None = Field(default=None, alias="toolSpec") + + +class _ConverseToolConfig(_ConverseModel): + tools: tuple[_ConverseTool, ...] = () + + +class _ConverseRequest(_ConverseModel): + system: tuple[_ConverseContentBlock, ...] = () + messages: tuple[_ConverseMessage, ...] = () + tool_config: _ConverseToolConfig | None = Field(default=None, alias="toolConfig") + + +def _parse_converse_request(validate: Callable[[], _ConverseRequest]) -> _ConverseRequest | None: + try: + return validate() + except ValidationError: + verbose_proxy_logger.debug("BedrockPassthroughGuardrailHandler: Converse body not convertible to messages") + return None + + +def _converse_block_text(block: _ConverseContentBlock) -> str: + if block.text: + return block.text + guard_text: Final = block.guard_content.text if block.guard_content else None + return guard_text.text if guard_text else "" + + +@dataclass(frozen=True, slots=True) +class _BlockRef: + message_index: int | None + block_index: int + path: tuple[str | int, ...] = () + + +@dataclass(frozen=True, slots=True) +class _ConvertedMessage: + message: AllMessageValues + text_blocks: tuple[_BlockRef, ...] + tool_result: tuple[_BlockRef, _ConverseToolResult] | None = None + tool_uses: tuple[tuple[_BlockRef, _ConverseToolUse], ...] = () + + +_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, ...]: + return tuple((_BlockRef(message_index, block_index), block) for block_index, block in enumerate(blocks)) + + +def _text_blocks(blocks: Sequence[_IndexedBlock]) -> tuple[tuple[_BlockRef, str], ...]: + return tuple((ref, text) for ref, block in blocks if (text := _converse_block_text(block))) + + +def _text_content(text_blocks: Sequence[tuple[_BlockRef, str]]) -> tuple[ChatCompletionTextObject, ...]: + return tuple(ChatCompletionTextObject(type="text", text=text) for _, text in text_blocks) + + +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 _scannable_tool_parts(tool_result)), + ) + + +def _converse_tool_call(tool_use: _ConverseToolUse) -> ChatCompletionAssistantToolCall: + return ChatCompletionAssistantToolCall( + id=tool_use.tool_use_id, + type="function", + function=ChatCompletionToolCallFunctionChunk( + name=tool_use.name, + arguments="{}" if tool_use.tool_input is None else json.dumps(tool_use.tool_input), + ), + ) + + +def _converse_user_group(is_tool_result: bool, group: Sequence[_IndexedBlock]) -> tuple[_ConvertedMessage, ...]: + if is_tool_result: + return tuple( + _ConvertedMessage(_converse_tool_message(block.tool_result), (), (ref, block.tool_result)) + for ref, block in group + if block.tool_result + ) + text_blocks: Final = _text_blocks(group) + if not text_blocks: + return () + message: Final = ChatCompletionUserMessage(role="user", content=list(_text_content(text_blocks))) + return (_ConvertedMessage(message, _refs(text_blocks)),) + + +def _converse_user_messages(blocks: Sequence[_IndexedBlock]) -> tuple[_ConvertedMessage, ...]: + return tuple( + chain.from_iterable( + _converse_user_group(is_tool_result, tuple(group)) + for is_tool_result, group in groupby(blocks, key=lambda item: item[1].tool_result is not None) + ) + ) + + +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 and include_tool_calls + ] + if tool_calls: + message: Final = ChatCompletionAssistantMessage( + role="assistant", content=content or None, tool_calls=tool_calls + ) + return ( + _ConvertedMessage( + message, + _refs(text_blocks), + tool_uses=tuple((ref, block.tool_use) for ref, block in blocks if block.tool_use), + ), + ) + if not content: + return () + return (_ConvertedMessage(ChatCompletionAssistantMessage(role="assistant", content=content), _refs(text_blocks)),) + + +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, include_tool_calls) + return _converse_user_messages(blocks) + + +def _converse_structured_messages(request: _ConverseRequest, include_tool_calls: bool) -> tuple[_ConvertedMessage, ...]: + converted: Final = tuple( + 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: + return converted + system: Final = ChatCompletionSystemMessage(role="system", content=list(_text_content(system_blocks))) + 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 _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: Final = message.get("content") + if isinstance(content, str): + return (content,) + try: + 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.""" + 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 + try: + return _ITEMS.validate_python(value) + except ValidationError: + return None + + +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: Final = _returned_text_parts(returned) + parts: Final = _scannable_tool_parts(result) + if texts is None or len(texts) != len(parts): + return None + 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 returned == converted.message: + return () + if converted.tool_result is not None: + return _tool_result_edits(converted, returned) + 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 + 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[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 + 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) + ) + 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(edits: Sequence[tuple[_StringHolder, object]]) -> None: + for (container, key), text in edits: + container[key] = text + + +def _converse_tools(request: _ConverseRequest) -> tuple[ChatCompletionToolParam, ...]: + tools: Final = request.tool_config.tools if request.tool_config else () + return tuple( + ChatCompletionToolParam( + type="function", + function=ChatCompletionToolParamFunctionChunk( + name=spec.name, + description=spec.description, + parameters=dict(spec.input_schema.json_schema), + ), + ) + for tool in tools + if (spec := tool.tool_spec) and spec.name and spec.input_schema and spec.input_schema.json_schema is not None + ) + + def _extract_converse_output_texts( content_blocks: Sequence[object], ) -> tuple[list[str], list[_StringHolder]]: @@ -177,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( @@ -446,12 +965,29 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): skip_system: Final = effective_skip_system_message_for_guardrail(guardrail_to_apply) skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply) - texts, holders = _extract_converse_texts(body, skip_system, skip_tool) + 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, 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, + scan_only_tool_results=scan_only_tool_results, + skip_system=skip_system, + skip_tool=skip_tool, + ) - if not texts: + 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 and not scan_only_tool_results else () + if tools: + inputs["tools"] = list(tools) model: Final = data.get("model") if model: inputs["model"] = model @@ -464,8 +1000,33 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) - if guardrailed_texts: - _write_back_texts(guardrailed_texts, holders) + guardrailed_messages: Final = guardrailed_inputs.get("structured_messages") + 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 57768b41943..3da48e5d68c 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -9,14 +9,7 @@ Validates that: - Non-converse endpoints skip guardrail execution """ -import copy from datetime import datetime -from typing import Literal -from unittest.mock import AsyncMock, MagicMock - -import pytest - -import litellm from litellm.caching.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -24,7 +17,14 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs +from typing import Literal +import litellm +import copy +import pytest +from unittest.mock import AsyncMock, MagicMock + +from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite from litellm.llms.bedrock.passthrough.guardrail_translation.handler import ( BedrockPassthroughGuardrailHandler, _extract_converse_texts, @@ -46,6 +46,20 @@ def _make_guardrail(apply_result: dict) -> MagicMock: return g +def _rewrite_structured_text(inputs: dict, old: str, new: str) -> dict: + rewritten = copy.deepcopy(inputs["structured_messages"]) + for message in rewritten: + for part in message["content"] if isinstance(message["content"], list) else []: + part["text"] = part["text"].replace(old, new) + return {**inputs, "structured_messages": rewritten} + + +def _make_structured_rewriting_guardrail(old: str, new: str) -> MagicMock: + g = _make_guardrail({}) + g.apply_guardrail = AsyncMock(side_effect=lambda inputs, **kwargs: _rewrite_structured_text(inputs, old, new)) + return g + + def _converse_data(endpoint: str = "model/anthropic.claude-3-sonnet/converse") -> dict: return { "endpoint": endpoint, @@ -186,9 +200,7 @@ class TestExtractConverseTexts: """Tool definitions are app-authored config, so nothing under toolConfig.tools reaches the guardrail as input content.""" body = { - "messages": [ - {"role": "user", "content": [{"text": "How much lag is there in my data?"}]} - ], + "messages": [{"role": "user", "content": [{"text": "How much lag is there in my data?"}]}], "toolConfig": { "tools": [ { @@ -233,11 +245,7 @@ class TestExtractConverseTexts: def test_tool_config_definitions_not_extracted_when_tool_messages_skipped(self): body = { "messages": [{"role": "user", "content": [{"text": "hi"}]}], - "toolConfig": { - "tools": [ - {"toolSpec": {"name": "fn", "description": "tool description"}} - ] - }, + "toolConfig": {"tools": [{"toolSpec": {"name": "fn", "description": "tool description"}}]}, } texts, _ = _extract_converse_texts(body, skip_system=False, skip_tool=True) assert texts == ["hi"] @@ -255,11 +263,7 @@ class TestExtractConverseTexts: ], } ], - "toolConfig": { - "tools": [ - {"toolSpec": {"name": "fn", "description": "tool description"}} - ] - }, + "toolConfig": {"tools": [{"toolSpec": {"name": "fn", "description": "tool description"}}]}, } texts, _ = _extract_converse_texts(body, skip_system=False, skip_tool=False) assert texts == ["hi", "user secret"] @@ -267,9 +271,7 @@ class TestExtractConverseTexts: def test_extracts_additional_model_request_fields(self): body = { "messages": [{"role": "user", "content": [{"text": "hi"}]}], - "additionalModelRequestFields": { - "reasoning_config": {"prompt": "blocked extra field"} - }, + "additionalModelRequestFields": {"reasoning_config": {"prompt": "blocked extra field"}}, } texts, _ = _extract_converse_texts(body, skip_system=False, skip_tool=False) assert "blocked extra field" in texts @@ -335,11 +337,7 @@ class TestWriteBackTexts: assert body["inferenceConfig"] == original["inferenceConfig"] def test_fewer_guardrailed_texts_logs_warning(self, monkeypatch): - body = { - "messages": [ - {"role": "user", "content": [{"text": "a"}, {"text": "b"}]} - ] - } + body = {"messages": [{"role": "user", "content": [{"text": "a"}, {"text": "b"}]}]} _, holders = _extract_converse_texts(body, skip_system=False, skip_tool=False) assert len(holders) == 2 @@ -414,9 +412,7 @@ class TestBedrockPassthroughGuardrailHandlerInput: } } ) - guardrail = _make_guardrail( - {"texts": ["You are helpful.", "Hello world", "[REDACTED]"]} - ) + guardrail = _make_guardrail({"texts": ["You are helpful.", "Hello world", "[REDACTED]"]}) result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) @@ -439,9 +435,7 @@ class TestBedrockPassthroughGuardrailHandlerInput: } } ) - guardrail = _make_guardrail( - {"texts": ["You are helpful.", "Hello world", "[REDACTED]"]} - ) + guardrail = _make_guardrail({"texts": ["You are helpful.", "Hello world", "[REDACTED]"]}) result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) @@ -455,12 +449,8 @@ class TestBedrockPassthroughGuardrailHandlerInput: """Blocked text hidden in toolUse.input must be scanned and masked.""" handler = BedrockPassthroughGuardrailHandler() data = _converse_data() - data["data"]["messages"][0]["content"][1]["toolUse"]["input"] = { - "query": "email john@example.com" - } - guardrail = _make_guardrail( - {"texts": ["You are helpful.", "Hello world", "[REDACTED]"]} - ) + data["data"]["messages"][0]["content"][1]["toolUse"]["input"] = {"query": "email john@example.com"} + guardrail = _make_guardrail({"texts": ["You are helpful.", "Hello world", "[REDACTED]"]}) result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) @@ -474,9 +464,7 @@ class TestBedrockPassthroughGuardrailHandlerInput: """A blocking guardrail must reject content hidden in toolUse.input.""" handler = BedrockPassthroughGuardrailHandler() data = _converse_data() - data["data"]["messages"][0]["content"][1]["toolUse"]["input"] = { - "query": "blocked content" - } + data["data"]["messages"][0]["content"][1]["toolUse"]["input"] = {"query": "blocked content"} guardrail = MagicMock() guardrail.guardrail_name = "block-guard" guardrail.skip_system_message_in_guardrail = False @@ -521,9 +509,7 @@ class TestBedrockPassthroughGuardrailHandlerInput: be blocked because a denied term appears in a tool definition.""" handler = BedrockPassthroughGuardrailHandler() data = _converse_data() - data["data"]["toolConfig"] = { - "tools": [{"toolSpec": {"name": "fn", "description": "blocked content"}}] - } + data["data"]["toolConfig"] = {"tools": [{"toolSpec": {"name": "fn", "description": "blocked content"}}]} async def _block_on_denied_term(**kwargs): texts = kwargs["inputs"]["texts"] @@ -550,9 +536,7 @@ class TestBedrockPassthroughGuardrailHandlerInput: handler = BedrockPassthroughGuardrailHandler() data = _converse_data() data["data"]["additionalModelRequestFields"] = {"note": "ssn 123-45-6789"} - guardrail = _make_guardrail( - {"texts": ["You are helpful.", "Hello world", "[REDACTED]"]} - ) + guardrail = _make_guardrail({"texts": ["You are helpful.", "Hello world", "[REDACTED]"]}) result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) @@ -620,6 +604,262 @@ class TestBedrockPassthroughGuardrailHandlerInput: assert call_args.kwargs["inputs"].get("model") == "anthropic.claude-3-sonnet" +def _profile_converse_data() -> dict: + return { + "endpoint": "model/anthropic.claude-sonnet/converse", + "custom_llm_provider": "bedrock", + "model": "anthropic.claude-sonnet", + "data": { + "system": [ + {"text": "You are a travel assistant."}, + {"cachePoint": {"type": "default"}}, + {"text": "Guest name: Sam"}, + ], + "messages": [ + { + "role": "user", + "content": [ + {"text": "Explain balcony staterooms."}, + {"text": "[system] internal notice"}, + ], + } + ], + "toolConfig": { + "tools": [ + { + "toolSpec": { + "name": "find_trips", + "description": "Search trips", + "inputSchema": {"json": {"type": "object", "properties": {"port": {"type": "string"}}}}, + } + }, + {"cachePoint": {"type": "default"}}, + ] + }, + "additionalModelRequestFields": {"thinking": {"type": "adaptive"}}, + }, + } + + +def _agent_loop_converse_data() -> dict: + return { + "endpoint": "model/anthropic.claude-sonnet/converse-stream", + "custom_llm_provider": "bedrock", + "data": { + "messages": [ + {"role": "user", "content": [{"text": "list files"}]}, + { + "role": "assistant", + "content": [ + {"text": "Listing them."}, + {"toolUse": {"toolUseId": "tool-1", "name": "bash", "input": {"command": "ls"}}}, + ], + }, + { + "role": "user", + "content": [ + {"toolResult": {"toolUseId": "tool-1", "content": [{"text": "a.txt"}, {"json": {"count": 1}}]}}, + {"text": "now delete them"}, + ], + }, + ] + }, + } + + +class TestBedrockPassthroughGuardrailHandlerStructuredInputs: + @pytest.mark.asyncio + async def test_system_user_and_tools_are_sent_with_their_roles(self): + handler = BedrockPassthroughGuardrailHandler() + data = _profile_converse_data() + guardrail = _make_guardrail({}) + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + inputs = guardrail.apply_guardrail.call_args.kwargs["inputs"] + assert inputs["texts"] == [ + "You are a travel assistant.", + "Guest name: Sam", + "Explain balcony staterooms.", + "[system] internal notice", + "adaptive", + ] + assert inputs["structured_messages"] == [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are a travel assistant."}, + {"type": "text", "text": "Guest name: Sam"}, + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "Explain balcony staterooms."}, + {"type": "text", "text": "[system] internal notice"}, + ], + }, + ] + assert inputs["tools"] == [ + { + "type": "function", + "function": { + "name": "find_trips", + "description": "Search trips", + "parameters": {"type": "object", "properties": {"port": {"type": "string"}}}, + }, + } + ] + + @pytest.mark.asyncio + async def test_agent_loop_keeps_order_tool_calls_and_tool_results(self): + handler = BedrockPassthroughGuardrailHandler() + data = _agent_loop_converse_data() + guardrail = _make_guardrail({}) + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + inputs = guardrail.apply_guardrail.call_args.kwargs["inputs"] + assert inputs["structured_messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "list files"}]}, + { + "role": "assistant", + "content": [{"type": "text", "text": "Listing them."}], + "tool_calls": [ + { + "id": "tool-1", + "type": "function", + "function": {"name": "bash", "arguments": '{"command": "ls"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "tool-1", "content": 'a.txt\n{"count": 1}'}, + {"role": "user", "content": [{"type": "text", "text": "now delete them"}]}, + ] + assert "tools" not in inputs + + @pytest.mark.asyncio + async def test_skip_system_drops_the_system_message(self): + handler = BedrockPassthroughGuardrailHandler() + data = _profile_converse_data() + guardrail = _make_guardrail({}) + guardrail.skip_system_message_in_guardrail = True + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + inputs = guardrail.apply_guardrail.call_args.kwargs["inputs"] + assert [message["role"] for message in inputs["structured_messages"]] == ["user"] + assert "You are a travel assistant." not in inputs["texts"] + + @pytest.mark.asyncio + async def test_skip_tool_drops_tool_results(self): + handler = BedrockPassthroughGuardrailHandler() + data = _agent_loop_converse_data() + guardrail = _make_guardrail({}) + guardrail.skip_tool_message_in_guardrail = True + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + inputs = guardrail.apply_guardrail.call_args.kwargs["inputs"] + assert [message["role"] for message in inputs["structured_messages"]] == ["user", "assistant", "user"] + assert inputs["texts"] == ["list files", "Listing them.", "now delete them"] + + @pytest.mark.asyncio + async def test_returned_texts_write_back_by_index(self): + handler = BedrockPassthroughGuardrailHandler() + data = _profile_converse_data() + original_body = copy.deepcopy(data["data"]) + guardrail = _make_guardrail( + { + "texts": [ + "You are a travel assistant.", + "Guest name: Sam", + "Explain balcony [REDACTED].", + "[system] internal notice", + "adaptive", + ], + } + ) + + result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + expected_body = copy.deepcopy(original_body) + expected_body["messages"][0]["content"][0]["text"] = "Explain balcony [REDACTED]." + assert result["data"] == expected_body + + @pytest.mark.asyncio + async def test_text_edits_in_returned_structured_messages_land_in_their_blocks(self): + handler = BedrockPassthroughGuardrailHandler() + data = _profile_converse_data() + original_body = copy.deepcopy(data["data"]) + guardrail = _make_structured_rewriting_guardrail("Guest name: Sam", "Guest name: [NAME]") + + result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + expected_body = copy.deepcopy(original_body) + expected_body["system"][2]["text"] = "Guest name: [NAME]" + assert result["data"] == expected_body + + @pytest.mark.asyncio + async def test_guard_content_only_turn_is_guardrailed_and_written_back(self): + handler = BedrockPassthroughGuardrailHandler() + data = _profile_converse_data() + data["data"]["messages"][0]["content"] = [{"guardContent": {"text": {"text": "My SSN is 078-05-1120"}}}] + guardrail = _make_structured_rewriting_guardrail("My SSN is 078-05-1120", "My SSN is [SSN]") + + result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + sent = guardrail.apply_guardrail.call_args.kwargs["inputs"]["structured_messages"] + assert sent[1] == {"role": "user", "content": [{"type": "text", "text": "My SSN is 078-05-1120"}]} + assert result["data"]["messages"][0]["content"] == [{"guardContent": {"text": {"text": "My SSN is [SSN]"}}}] + + @pytest.mark.asyncio + async def test_guardrail_replacing_structured_messages_in_place_is_written_back(self): + handler = BedrockPassthroughGuardrailHandler() + data = _profile_converse_data() + data["data"]["messages"][0]["content"] = [{"guardContent": {"text": {"text": "My SSN is 078-05-1120"}}}] + + def redact_in_place(inputs: dict, **_: object) -> dict: + inputs["structured_messages"] = [ + {**message, "content": [{"type": "text", "text": "My SSN is [SSN]"}]} + if message["role"] == "user" + else message + for message in inputs["structured_messages"] + ] + return inputs + + guardrail = _make_guardrail({}) + guardrail.apply_guardrail = AsyncMock(side_effect=lambda inputs, **kwargs: redact_in_place(inputs)) + + result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert result["data"]["messages"][0]["content"] == [{"guardContent": {"text": {"text": "My SSN is [SSN]"}}}] + + @pytest.mark.asyncio + async def test_text_cleared_in_texts_and_structured_messages_stays_cleared(self): + handler = BedrockPassthroughGuardrailHandler() + data = _profile_converse_data() + guardrail = _make_structured_rewriting_guardrail("Explain balcony staterooms.", "") + guardrail.apply_guardrail = AsyncMock( + side_effect=lambda inputs, **kwargs: { + **_rewrite_structured_text(inputs, "Explain balcony staterooms.", ""), + "texts": [text.replace("Explain balcony staterooms.", "") for text in inputs["texts"]], + } + ) + + result = await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert result["data"]["messages"][0]["content"][0] == {"text": ""} + + @pytest.mark.asyncio + async def test_structured_rewrite_that_cannot_be_placed_fails_closed(self): + handler = BedrockPassthroughGuardrailHandler() + guardrail = _make_guardrail({"structured_messages": [{"role": "user", "content": "replaced"}]}) + + with pytest.raises(UnappliableRequestRewrite): + await handler.process_input_messages(data=_profile_converse_data(), guardrail_to_apply=guardrail) + + class TestBedrockPassthroughGuardrailHandlerOutput: def _converse_response(self, text: str = "Model reply") -> dict: return { @@ -707,13 +947,9 @@ class TestBedrockPassthroughGuardrailHandlerOutput: }, "stopReason": "end_turn", } - guardrail = _make_guardrail( - {"texts": ["[V]", "[REASON]", "[INPUT]"]} - ) + guardrail = _make_guardrail({"texts": ["[V]", "[REASON]", "[INPUT]"]}) - result = await handler.process_output_response( - response=response, guardrail_to_apply=guardrail - ) + result = await handler.process_output_response(response=response, guardrail_to_apply=guardrail) sent_texts = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] assert "thinking about john@example.com" in sent_texts @@ -738,9 +974,7 @@ class TestBedrockPassthroughGuardrailHandlerOutput: { "citationsContent": { "content": [{"text": "Contact john@example.com"}], - "citations": [ - {"source": "https://example.com", "title": "Example"} - ], + "citations": [{"source": "https://example.com", "title": "Example"}], } } ], @@ -750,9 +984,7 @@ class TestBedrockPassthroughGuardrailHandlerOutput: } guardrail = _make_guardrail({"texts": ["[CITED]"]}) - result = await handler.process_output_response( - response=response, guardrail_to_apply=guardrail - ) + result = await handler.process_output_response(response=response, guardrail_to_apply=guardrail) sent_texts = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] assert sent_texts == ["Contact john@example.com"] @@ -832,52 +1064,6 @@ def _build_event_stream_frame(event_type: str, payload: dict) -> bytes: return prelude + prelude_crc_b + headers_bytes + payload_bytes + msg_crc_b -_TEST_CARD = "4111 1111 1111 1111" - - -class _CardMaskingGuardrail(CustomGuardrail): - def __init__(self) -> None: - super().__init__(guardrail_name="card-mask", event_hook=GuardrailEventHooks.post_call, default_on=True) - self.received_texts: list[list[str]] = [] - - async def apply_guardrail( - self, - inputs: GenericGuardrailAPIInputs, - request_data: dict, - input_type: Literal["request", "response"], - logging_obj: LiteLLMLoggingObj | None = None, - ) -> GenericGuardrailAPIInputs: - texts = inputs.get("texts", []) - self.received_texts.append(list(texts)) - return {**inputs, "texts": [text.replace(_TEST_CARD, "****") for text in texts]} - - -async def _run_post_call_guardrail_on_stream( - stream_bytes: bytes, guardrail: CustomGuardrail, monkeypatch: pytest.MonkeyPatch -) -> bytes: - monkeypatch.setattr(litellm, "callbacks", [guardrail]) - logging_obj = LiteLLMLoggingObj( - model="claude", - messages=[], - stream=True, - call_type="allm_passthrough_route", - start_time=datetime(2026, 1, 1), - litellm_call_id="test-call", - function_id="test-function", - ) - return await BedrockPassthroughGuardrailHandler.de_anonymize_event_stream( - body_bytes=stream_bytes, - proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), - user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), - data={ - "model": "claude", - "custom_llm_provider": "bedrock", - "endpoint": "model/claude/converse-stream", - "litellm_logging_obj": logging_obj, - }, - ) - - class TestDeAnonymizeConverseStream: def _make_proxy_logging(self, mock_hook) -> MagicMock: proxy_logging_obj = MagicMock() @@ -928,78 +1114,6 @@ class TestDeAnonymizeConverseStream: ] assert "".join(texts) == "John Doe works at Acme Corp" - @staticmethod - def _deltas_by_block(result: bytes, field: str) -> dict[int, str]: - import json - from botocore.eventstream import EventStreamBuffer - - buf = EventStreamBuffer() - buf.add_data(result) - blocks: dict[int, str] = {} - for msg in buf: - if msg.headers.get(":event-type") != "contentBlockDelta": - continue - payload = json.loads(msg.payload) - delta = payload["delta"] - value = delta["toolUse"]["input"] if field == "toolUse" else delta[field] - blocks[payload["contentBlockIndex"]] = blocks.get(payload["contentBlockIndex"], "") + value - return blocks - - @pytest.mark.asyncio - async def test_text_blocks_are_guardrailed_separately(self, monkeypatch): - stream_bytes = ( - _build_event_stream_frame("messageStart", {"role": "assistant"}) - + _build_event_stream_frame( - "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": f"Card: {_TEST_CARD}."}} - ) - + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}) - + _build_event_stream_frame( - "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Thanks for sailing with us!"}} - ) - + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 1}) - + _build_event_stream_frame("messageStop", {"stopReason": "end_turn"}) - ) - guardrail = _CardMaskingGuardrail() - - result = await _run_post_call_guardrail_on_stream(stream_bytes, guardrail, monkeypatch) - - assert guardrail.received_texts == [[f"Card: {_TEST_CARD}.", "Thanks for sailing with us!"]] - assert self._deltas_by_block(result, "text") == {0: "Card: ****.", 1: "Thanks for sailing with us!"} - - @pytest.mark.asyncio - async def test_parallel_tool_use_inputs_stay_in_their_own_block(self, monkeypatch): - import json - - first_input = f'{{"ship":"AU","cardNumber":"{_TEST_CARD}"}}' - second_input = '{"ship":"AU"}' - stream_bytes = ( - _build_event_stream_frame( - "contentBlockStart", - {"contentBlockIndex": 0, "start": {"toolUse": {"toolUseId": "tooluse_a", "name": "book"}}}, - ) - + _build_event_stream_frame( - "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"toolUse": {"input": first_input}}} - ) - + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}) - + _build_event_stream_frame( - "contentBlockStart", - {"contentBlockIndex": 1, "start": {"toolUse": {"toolUseId": "tooluse_b", "name": "ship_info"}}}, - ) - + _build_event_stream_frame( - "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"toolUse": {"input": second_input}}} - ) - + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 1}) - + _build_event_stream_frame("messageStop", {"stopReason": "tool_use"}) - ) - guardrail = _CardMaskingGuardrail() - - result = await _run_post_call_guardrail_on_stream(stream_bytes, guardrail, monkeypatch) - - assert guardrail.received_texts == [[first_input, second_input]] - inputs = self._deltas_by_block(result, "toolUse") - assert json.loads(inputs[0]) == {"ship": "AU", "cardNumber": "****"} - assert json.loads(inputs[1]) == {"ship": "AU"} - @pytest.mark.asyncio async def test_tokens_split_across_chunks_reassembled(self): import json @@ -1195,9 +1309,7 @@ class TestDeAnonymizeConverseStream: buf = EventStreamBuffer() buf.add_data(result) return [ - json.loads(msg.payload)["delta"] - for msg in buf - if msg.headers.get(":event-type") == "contentBlockDelta" + json.loads(msg.payload)["delta"] for msg in buf if msg.headers.get(":event-type") == "contentBlockDelta" ] @pytest.mark.asyncio @@ -1223,9 +1335,7 @@ class TestDeAnonymizeConverseStream: result = await BedrockPassthroughGuardrailHandler.de_anonymize_event_stream( body_bytes=stream_bytes, - proxy_logging_obj=self._make_proxy_logging( - self._token_replacing_hook({"": "Alice"}) - ), + proxy_logging_obj=self._make_proxy_logging(self._token_replacing_hook({"": "Alice"})), user_api_key_dict=MagicMock(), data={}, ) @@ -1245,9 +1355,7 @@ class TestDeAnonymizeConverseStream: result = await BedrockPassthroughGuardrailHandler.de_anonymize_event_stream( body_bytes=stream_bytes, - proxy_logging_obj=self._make_proxy_logging( - self._token_replacing_hook({"": "Alice"}) - ), + proxy_logging_obj=self._make_proxy_logging(self._token_replacing_hook({"": "Alice"})), user_api_key_dict=MagicMock(), data={}, ) @@ -1273,9 +1381,7 @@ class TestDeAnonymizeConverseStream: result = await BedrockPassthroughGuardrailHandler.de_anonymize_event_stream( body_bytes=stream_bytes, - proxy_logging_obj=self._make_proxy_logging( - self._token_replacing_hook({"": "Alice"}) - ), + proxy_logging_obj=self._make_proxy_logging(self._token_replacing_hook({"": "Alice"})), user_api_key_dict=MagicMock(), data={}, ) @@ -1291,9 +1397,7 @@ class TestDeAnonymizeConverseStream: captured = {} async def mock_hook(data, user_api_key_dict, response): - captured["texts"] = [ - b["text"] for b in response["output"]["message"]["content"] - ] + captured["texts"] = [b["text"] for b in response["output"]["message"]["content"]] mapping = {"": "Alice", "": "Acme"} for block in response["output"]["message"]["content"]: text = block["text"] @@ -1342,3 +1446,684 @@ class TestDeAnonymizeConverseStream: hook_spy.assert_not_called() assert result is stream_bytes + + @pytest.mark.asyncio + async def test_parallel_tool_use_inputs_stay_in_their_own_block(self, monkeypatch): + import json + + first_input = f'{{"ship":"AU","cardNumber":"{_TEST_CARD}"}}' + second_input = '{"ship":"AU"}' + stream_bytes = ( + _build_event_stream_frame( + "contentBlockStart", + {"contentBlockIndex": 0, "start": {"toolUse": {"toolUseId": "tooluse_a", "name": "book"}}}, + ) + + _build_event_stream_frame( + "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"toolUse": {"input": first_input}}} + ) + + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}) + + _build_event_stream_frame( + "contentBlockStart", + {"contentBlockIndex": 1, "start": {"toolUse": {"toolUseId": "tooluse_b", "name": "ship_info"}}}, + ) + + _build_event_stream_frame( + "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"toolUse": {"input": second_input}}} + ) + + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 1}) + + _build_event_stream_frame("messageStop", {"stopReason": "tool_use"}) + ) + guardrail = _CardMaskingGuardrail() + + result = await _run_post_call_guardrail_on_stream(stream_bytes, guardrail, monkeypatch) + + assert guardrail.received_texts == [[first_input, second_input]] + inputs = self._deltas_by_block(result, "toolUse") + assert json.loads(inputs[0]) == {"ship": "AU", "cardNumber": "****"} + assert json.loads(inputs[1]) == {"ship": "AU"} + + @pytest.mark.asyncio + async def test_text_blocks_are_guardrailed_separately(self, monkeypatch): + stream_bytes = ( + _build_event_stream_frame("messageStart", {"role": "assistant"}) + + _build_event_stream_frame( + "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": f"Card: {_TEST_CARD}."}} + ) + + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}) + + _build_event_stream_frame( + "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Thanks for sailing with us!"}} + ) + + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 1}) + + _build_event_stream_frame("messageStop", {"stopReason": "end_turn"}) + ) + guardrail = _CardMaskingGuardrail() + + result = await _run_post_call_guardrail_on_stream(stream_bytes, guardrail, monkeypatch) + + assert guardrail.received_texts == [[f"Card: {_TEST_CARD}.", "Thanks for sailing with us!"]] + assert self._deltas_by_block(result, "text") == {0: "Card: ****.", 1: "Thanks for sailing with us!"} + + @staticmethod + def _deltas_by_block(result: bytes, field: str) -> dict[int, str]: + import json + from botocore.eventstream import EventStreamBuffer + + buf = EventStreamBuffer() + buf.add_data(result) + blocks: dict[int, str] = {} + for msg in buf: + if msg.headers.get(":event-type") != "contentBlockDelta": + continue + payload = json.loads(msg.payload) + delta = payload["delta"] + value = delta["toolUse"]["input"] if field == "toolUse" else delta[field] + blocks[payload["contentBlockIndex"]] = blocks.get(payload["contentBlockIndex"], "") + value + return blocks + + +class TestConverseRewriteRegressions: + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [False, True]) + async def test_string_content_redaction_reaches_bedrock(self, nested): + block = ( + {"guardContent": {"text": {"text": "4111 1111 1111 1111"}}} if nested else {"text": "4111 1111 1111 1111"} + ) + data = {"endpoint": "model/test/converse", "data": {"messages": [{"role": "user", "content": [block]}]}} + guard = _make_guardrail({"structured_messages": [{"role": "user", "content": "[CARD]"}]}) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + expected = {"guardContent": {"text": {"text": "[CARD]"}}} if nested else {"text": "[CARD]"} + assert result["data"]["messages"][0]["content"] == [expected] + + @pytest.mark.asyncio + async def test_partial_structured_edit_preserves_flat_card_mask(self): + data = { + "endpoint": "model/test/converse", + "data": {"messages": [{"role": "user", "content": [{"text": "4111 1111 1111 1111"}, {"text": "Alice"}]}]}, + } + guard = _make_guardrail( + { + "texts": ["[CARD]", "Alice"], + "structured_messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "4111 1111 1111 1111"}, + {"type": "text", "text": "[NAME]"}, + ], + } + ], + } + ) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert result["data"]["messages"][0]["content"] == [{"text": "[CARD]"}, {"text": "[NAME]"}] + + @pytest.mark.asyncio + async def test_tool_result_redaction_keeps_bedrock_metadata(self): + data = { + "endpoint": "model/test/converse", + "data": { + "messages": [ + { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "t1", + "status": "success", + "content": [{"text": "4111 1111 1111 1111"}], + } + } + ], + } + ] + }, + } + guard = _make_guardrail({"structured_messages": [{"role": "tool", "tool_call_id": "t1", "content": "[CARD]"}]}) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert result["data"]["messages"][0]["content"] == [ + {"toolResult": {"toolUseId": "t1", "status": "success", "content": [{"text": "[CARD]"}]}} + ] + + @pytest.mark.asyncio + async def test_tool_arguments_redaction_preserves_flat_masked_sibling(self): + data = { + "endpoint": "model/test/converse", + "data": { + "messages": [ + { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "t1", + "name": "pay", + "input": {"card": "4111 1111 1111 1111", "name": "Alice"}, + } + } + ], + } + ] + }, + } + guard = _make_guardrail( + { + "texts": ["[CARD]", "Alice"], + "structured_messages": [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "t1", + "type": "function", + "function": { + "name": "pay", + "arguments": '{"card":"4111 1111 1111 1111","name":"[NAME]"}', + }, + } + ], + } + ], + } + ) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert result["data"]["messages"][0]["content"] == [ + {"toolUse": {"toolUseId": "t1", "name": "pay", "input": {"card": "[CARD]", "name": "[NAME]"}}} + ] + + @pytest.mark.asyncio + async def test_tool_only_scopes_both_flat_and_structured_inputs(self): + data = { + "endpoint": "model/test/converse", + "data": { + "system": [{"text": "system"}], + "messages": [ + { + "role": "user", + "content": [ + {"text": "user"}, + {"toolResult": {"toolUseId": "t1", "content": [{"text": "result"}]}}, + ], + }, + { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "t2", "name": "search", "input": {"q": "arguments"}}}], + }, + ], + "additionalModelRequestFields": {"prompt": "extra"}, + "toolConfig": { + "tools": [{"toolSpec": {"name": "search", "inputSchema": {"json": {"type": "object"}}}}] + }, + }, + } + guard = _make_guardrail({}) + guard.scan_only_tool_results = True + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + 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): + data = { + "endpoint": "model/test/converse", + "data": {"messages": [{"role": "user", "content": [{"text": "card"}, {"text": "name"}]}]}, + } + original = copy.deepcopy(data) + guard = _make_guardrail( + {"texts": ["masked", "name"], "structured_messages": [{"role": "user", "content": "masked"}]} + ) + with pytest.raises(UnappliableRequestRewrite): + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert data == original + + @pytest.mark.asyncio + @pytest.mark.parametrize("content", ["[CARD]", [{"type": "text", "text": "[CARD]"}]]) + async def test_tool_result_text_accepts_string_and_text_parts(self, content): + data = { + "endpoint": "model/test/converse", + "data": { + "messages": [ + {"role": "user", "content": [{"toolResult": {"toolUseId": "t1", "content": [{"text": "raw"}]}}]} + ] + }, + } + guard = _make_guardrail({"structured_messages": [{"role": "tool", "tool_call_id": "t1", "content": content}]}) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert result["data"]["messages"][0]["content"] == [ + {"toolResult": {"toolUseId": "t1", "content": [{"text": "[CARD]"}]}} + ] + + @pytest.mark.asyncio + async def test_nested_tool_result_json_redactions_keep_flat_masks(self): + data = { + "endpoint": "model/test/converse", + "data": { + "messages": [ + { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "t1", + "content": [{"json": {"rows": [{"card": "raw", "name": "Alice", "count": 1}]}}], + } + } + ], + } + ] + }, + } + guard = _make_guardrail( + { + "texts": ["[CARD]", "Alice"], + "structured_messages": [ + { + "role": "tool", + "tool_call_id": "t1", + "content": '{"rows":[{"card":"raw","name":"[NAME]","count":1}]}', + } + ], + } + ) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert result["data"]["messages"][0]["content"] == [ + { + "toolResult": { + "toolUseId": "t1", + "content": [{"json": {"rows": [{"card": "[CARD]", "name": "[NAME]", "count": 1}]}}], + } + } + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("change", ["invalid_json", "changed_id", "changed_name", "changed_shape"]) + async def test_unsafe_tool_call_rewrites_fail_closed(self, change): + data = { + "endpoint": "model/test/converse", + "data": { + "messages": [ + { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "t1", "name": "pay", "input": {"card": "raw"}}}], + } + ] + }, + } + original = copy.deepcopy(data) + call = {"id": "t1", "type": "function", "function": {"name": "pay", "arguments": '{"card":"masked"}'}} + if change == "invalid_json": + call["function"]["arguments"] = "not json" + elif change == "changed_id": + call["id"] = "t2" + elif change == "changed_name": + call["function"]["name"] = "other" + else: + call["function"]["arguments"] = '{"different":"masked"}' + guard = _make_guardrail( + {"texts": ["masked"], "structured_messages": [{"role": "assistant", "content": None, "tool_calls": [call]}]} + ) + with pytest.raises(UnappliableRequestRewrite): + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert data == original + + @pytest.mark.asyncio + @pytest.mark.parametrize("skip_tool", [False, True]) + async def test_tool_only_without_eligible_results_skips_guardrail(self, skip_tool): + data = { + "endpoint": "model/test/converse", + "data": {"messages": [{"role": "user", "content": [{"text": "user"}]}]}, + } + if skip_tool: + data["data"]["messages"][0]["content"].append( + {"toolResult": {"toolUseId": "t1", "content": [{"text": "result"}]}} + ) + guard = _make_guardrail({}) + guard.scan_only_tool_results = True + guard.skip_tool_message_in_guardrail = skip_tool + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + 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 +) -> bytes: + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + logging_obj = LiteLLMLoggingObj( + model="claude", + messages=[], + stream=True, + call_type="allm_passthrough_route", + start_time=datetime(2026, 1, 1), + litellm_call_id="test-call", + function_id="test-function", + ) + return await BedrockPassthroughGuardrailHandler.de_anonymize_event_stream( + body_bytes=stream_bytes, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + data={ + "model": "claude", + "custom_llm_provider": "bedrock", + "endpoint": "model/claude/converse-stream", + "litellm_logging_obj": logging_obj, + }, + ) + + +_TEST_CARD = "4111 1111 1111 1111" + + +class _CardMaskingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="card-mask", event_hook=GuardrailEventHooks.post_call, default_on=True) + self.received_texts: list[list[str]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + texts = inputs.get("texts", []) + self.received_texts.append(list(texts)) + return {**inputs, "texts": [text.replace(_TEST_CARD, "****") for text in texts]}