diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 54fce61c713..a7ccad92804 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -25,7 +25,11 @@ from litellm.llms.bedrock.chat.converse_to_openai import ( converse_tools, parse_converse_request, ) -from litellm.llms.bedrock.passthrough.guardrail_translation.structured_edits import Edit, structured_edits +from litellm.llms.bedrock.passthrough.guardrail_translation.structured_edits import ( + Edit, + ToolResultText, + structured_edits, +) from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -61,6 +65,7 @@ def _generic_passthrough_handler() -> "PassThroughEndpointHandler": _StringHolder = tuple[Any, str | int] +_Write: TypeAlias = tuple[_StringHolder, object] _BODY_BY_REFERENCE: Final = TypeAdapter(InstanceOf[dict[str, object]]) @@ -208,20 +213,40 @@ def _changed_leaves( ) +def _tool_result_write(body: Mapping[str, object], edit: ToolResultText) -> _Write | None: + container: Final = reduce(_child, edit.content[:-1], body) + blocks: Final = as_json_array(_child(container, edit.content[-1])) + if blocks is None: + return None + dropped: Final = frozenset(edit.indices[1:]) + return (container, edit.content[-1]), [ + {"text": edit.text} if index == edit.indices[0] else block + for index, block in enumerate(blocks) + if index not in dropped + ] + + +def _resolved(body: Mapping[str, object], edit: Edit) -> _Write | None: + match edit: + case ToolResultText(): + return _tool_result_write(body, edit) + case (path, text): + holder: Final = _holder(body, path) + return None if holder is None else (holder, text) + + def _holder_edits( body: Mapping[str, object], edits: Sequence[Edit], flat_changes: frozenset[tuple[int, str | int]] -) -> tuple[tuple[_StringHolder, str], ...] | None: - resolved: Final = tuple((_holder(body, path), text) for path, text in edits) - if any(holder is None for holder, _ in resolved): +) -> tuple[_Write, ...] | None: + resolved: Final = tuple(_resolved(body, edit) for edit in edits) + if any(write is None for write in resolved): return None - return tuple( - (holder, text) for holder, text in resolved if holder is not None and _leaf_identity(holder) not in flat_changes - ) + return tuple(write for write in resolved if write is not None and _leaf_identity(write[0]) not in flat_changes) -def _write_back_structured_texts(edits: Sequence[tuple[_StringHolder, str]]) -> None: - for (container, key), text in edits: - container[key] = text +def _write_back_structured_edits(edits: Sequence[_Write]) -> None: + for (container, key), value in edits: + container[key] = value def _extract_converse_output_texts( @@ -589,7 +614,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) if applied_texts: _write_back_texts(applied_texts, holders) - _write_back_structured_texts(holder_edits) + _write_back_structured_edits(holder_edits) return data diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/structured_edits.py b/litellm/llms/bedrock/passthrough/guardrail_translation/structured_edits.py index 8c83786d42d..fd21a40db2f 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/structured_edits.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/structured_edits.py @@ -1,4 +1,5 @@ from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from itertools import chain from typing import Final, TypeAlias @@ -16,7 +17,15 @@ from litellm.llms.bedrock.chat.converse_to_openai import ( ) from litellm.types.llms.openai import ChatCompletionTextObject -Edit: TypeAlias = tuple[JsonPath, str] + +@dataclass(frozen=True, slots=True) +class ToolResultText: + content: JsonPath + indices: tuple[int, ...] + text: str + + +Edit: TypeAlias = tuple[JsonPath, str] | ToolResultText _JsonPair: TypeAlias = tuple[JsonPath, object, object] _TEXT_PARTS: Final = TypeAdapter(tuple[ChatCompletionTextObject, ...]) _PARSED_JSON: Final = TypeAdapter(object) @@ -184,6 +193,28 @@ def _reproduces(expected: ConverseRow, returned: Mapping[str, object] | None) -> ) +def _tool_result_text(row: ConverseRow, returned: Mapping[str, object]) -> ToolResultText | None: + texts: Final = _text_parts(returned) + indices: Final = tuple(index for slot in row.content if isinstance(index := slot.path[-2], int)) + if ( + row.message["role"] != "tool" + or not indices + or texts is None + or _without_content(returned) != _without_content(row.message) + ): + return None + return ToolResultText(row.content[0].path[:-2], indices, "\n".join(texts)) + + +def _unmapped_edits( + row: ConverseRow, view: ConverseRow, returned: Mapping[str, object] | None +) -> tuple[Edit, ...] | None: + if _reproduces(view, returned): + return () + rewrite: Final = _tool_result_text(row, returned) if returned is not None and view == row else None + return None if rewrite is None else (rewrite,) + + def structured_edits( rows: Sequence[ConverseRow], scoped_indices: Sequence[int], @@ -199,8 +230,8 @@ def structured_edits( unmapped: Final = tuple( (index, row) for index, row, edits in zip(scoped_indices, returned_rows, proposed) if edits is None ) - if unmapped: - view: Final = texts_applied_rows() - if len(view) != len(rows) or not all(_reproduces(view[index], row) for index, row in unmapped): - return None - return tuple(chain.from_iterable(edits for edits in proposed if edits is not None)) + view: Final = texts_applied_rows() if unmapped else () + if unmapped and len(view) != len(rows): + return None + fallbacks: Final = tuple(_unmapped_edits(rows[index], view[index], row) for index, row in unmapped) + return _merged((*(edits for edits in proposed if edits is not None), *fallbacks)) diff --git a/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py b/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py index eeb6082931e..95d99a3c56f 100644 --- a/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py +++ b/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py @@ -41,6 +41,8 @@ class _Vendor(str, Enum): MASKS_TEXTS_ONLY = "masks-texts-only" BLOCKS = "blocks" DROPS_A_ROW = "drops-a-row" + COMPACTS_TOOL_RESULTS = "compacts-tool-results" + COMPACTS_RAW_TOOL_RESULTS = "compacts-raw-tool-results" @dataclass(frozen=True, slots=True) @@ -78,6 +80,10 @@ def _blank_content(row: Mapping[str, JsonValue]) -> dict[str, JsonValue]: return {**row, "content": ""} if row.get("content") is None else dict(row) +def _compacted(row: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {**row, "content": f"compacted: {row.get('content')}"} if row.get("role") == "tool" else dict(row) + + def _generic_verdict(vendor: _Vendor, body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: if _CARD_PATTERN.search(json.dumps(body)) is None: return {"action": "NONE"} @@ -119,6 +125,14 @@ def _generic_verdict(vendor: _Vendor, body: Mapping[str, JsonValue]) -> dict[str return {"action": "BLOCKED", "blocked_reason": "card number"} case _Vendor.DROPS_A_ROW: return {"action": "GUARDRAIL_INTERVENED", "texts": texts, "structured_messages": list(map(dict, rows[:-1]))} + case _Vendor.COMPACTS_TOOL_RESULTS: + return {"action": "GUARDRAIL_INTERVENED", "structured_messages": list(map(_compacted, rows))} + case _Vendor.COMPACTS_RAW_TOOL_RESULTS: + return { + "action": "GUARDRAIL_INTERVENED", + "texts": texts, + "structured_messages": list(map(_compacted, sent_rows)), + } def _akto_verdict(body: bytes) -> Reply: @@ -339,6 +353,35 @@ def test_tool_result_with_text_and_image_parts_reaches_bedrock_masked_with_image ] +def test_compacted_tool_result_reaches_bedrock_as_one_text_block_with_the_image_in_place(rig: _Rig) -> None: + image: Final = {"image": {"format": "png", "source": {"bytes": _PNG}}} + body: Final = { + "messages": [ + _user({"text": "Look up my card."}), + _assistant(_tool_use("last four 1111")), + _user(_tool_result({"text": f"found {_CARD}"}, image, {"json": {"card": _CARD, "status": "active"}})), + ] + } + sent: Final = _forwarded(rig, _converse(rig, _generic(_Vendor.COMPACTS_TOOL_RESULTS), body)) + compacted: Final = f"compacted: found {_MASK}\n{json.dumps({'card': _MASK, 'status': 'active'})}" + assert sent["messages"] == [ + _user({"text": "Look up my card."}), + _assistant(_tool_use("last four 1111")), + _user(_tool_result({"text": compacted}, image)), + ] + + +def test_compacted_tool_result_that_drops_the_text_mask_is_refused_before_bedrock(rig: _Rig) -> None: + body: Final = { + "messages": [ + _user({"text": "Look up my card."}), + _assistant(_tool_use("last four 1111")), + _user(_tool_result({"json": {"card": _CARD, "status": "active"}})), + ] + } + assert _UNAPPLIABLE in _refused(rig, _converse(rig, _generic(_Vendor.COMPACTS_RAW_TOOL_RESULTS), body)) + + 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, _generic(_Vendor.RESPACES_UNMASKED_ROWS), body)) 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 b3ca059fc1e..094d6ce2494 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -2388,7 +2388,48 @@ class TestConverseRewriteRegressions: 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): + @pytest.mark.parametrize( + ("content", "text"), + [ + ('{"rows":"[1]{id}\\n1"}', '{"rows":"[1]{id}\\n1"}'), + ([{"type": "text", "text": "rows"}, {"type": "text", "text": "[1]{id}"}], "rows\n[1]{id}"), + ], + ) + async def test_reshaped_tool_result_row_is_sent_as_one_text_block_where_its_parts_were(self, content, text): + image = {"image": {"format": "png", "source": {"bytes": "iVBORw0KGgo="}}} + data = { + "endpoint": "model/test/converse", + "data": { + "messages": [ + { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "t1", + "content": [image, {"text": "found 4111"}, {"json": {"rows": [{"id": 1}]}}], + "status": "success", + } + } + ], + } + ] + }, + } + guard = _replying({"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": [image, {"text": text}], + "status": "success", + } + } + ] + + @pytest.mark.asyncio + async def test_joined_tool_result_row_is_sent_as_one_text_block(self): data = { "endpoint": "model/test/converse", "data": { @@ -2402,10 +2443,66 @@ class TestConverseRewriteRegressions: ] }, } - original = copy.deepcopy(data["data"]) guard = _replying( {"structured_messages": [{"role": "tool", "tool_call_id": "t1", "content": "found [CARD]\nok"}]} ) + result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert result["data"]["messages"][0]["content"] == [ + {"toolResult": {"toolUseId": "t1", "content": [{"text": "found [CARD]\nok"}]}} + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("content", "texts", "row"), + [ + pytest.param( + [{"json": {"card": "4111", "rows": [1, 2]}}], + ["[CARD]"], + {"role": "tool", "tool_call_id": "t1", "content": '{"card": "4111", "rows": "1,2"}'}, + id="flat-mask-on-the-tool-result", + ), + pytest.param( + [{"json": {"rows": [1, 2]}}], + None, + {"role": "tool", "tool_call_id": "t2", "content": '{"rows": "1,2"}'}, + id="changed-tool-call-id", + ), + pytest.param( + [{"image": {"format": "png", "source": {"bytes": "iVBORw0KGgo="}}}], + None, + {"role": "tool", "tool_call_id": "t1", "content": "an image"}, + id="no-text-parts-to-replace", + ), + ], + ) + async def test_tool_result_reshape_that_would_drop_a_mask_or_misplace_text_fails_closed(self, content, texts, row): + data = { + "endpoint": "model/test/converse", + "data": { + "messages": [{"role": "user", "content": [{"toolResult": {"toolUseId": "t1", "content": content}}]}] + }, + } + original = copy.deepcopy(data["data"]) + guard = _replying({"structured_messages": [row], **({"texts": texts} if texts is not None else {})}) + with pytest.raises(UnappliableRequestRewrite): + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) + assert data["data"] == original + + @pytest.mark.asyncio + async def test_joined_user_row_the_flat_texts_do_not_reproduce_fails_closed(self): + data = { + "endpoint": "model/test/converse", + "data": { + "messages": [ + { + "role": "user", + "content": [{"text": "card 4111"}, {"cachePoint": {"type": "default"}}, {"text": "thanks"}], + } + ] + }, + } + original = copy.deepcopy(data["data"]) + guard = _replying({"structured_messages": [{"role": "user", "content": "card [CARD]\nthanks"}]}) with pytest.raises(UnappliableRequestRewrite): await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard) assert data["data"] == original