diff --git a/litellm/llms/bedrock/chat/converse_to_openai.py b/litellm/llms/bedrock/chat/converse_to_openai.py index 02b94453e54..ce80a2dc8bb 100644 --- a/litellm/llms/bedrock/chat/converse_to_openai.py +++ b/litellm/llms/bedrock/chat/converse_to_openai.py @@ -295,26 +295,37 @@ def _user_row(slots: Sequence[Slot]) -> ConverseRow: return ConverseRow(ChatCompletionUserMessage(role="user", content=parts), tuple(slots)) -def _with_request_fields(rows: tuple[ConverseRow, ...], value: object) -> tuple[ConverseRow, ...]: +def _with_request_fields(rows: tuple[ConverseRow, ...], value: object, user_turn: int) -> tuple[ConverseRow, ...]: """ ``rows`` with the ``additionalModelRequestFields`` strings after the text of the latest user row, so a guardrail - that judges only the latest user message reads both. They get a user row of their own when there is none. + that judges only the latest user message reads both. When no user row has text they get a user row of their own + at ``user_turn``, where the rows of the latest user turn that is not only tool results start. """ slots: Final = _string_leaves(("additionalModelRequestFields",), value) if not slots: return rows latest: Final = next((index for index in reversed(range(len(rows))) if rows[index].message["role"] == "user"), None) if latest is None: - return (*rows, _user_row(slots)) + return (*rows[:user_turn], _user_row(slots), *rows[user_turn:]) return (*rows[:latest], _user_row((*rows[latest].content, *slots)), *rows[latest + 1 :]) +def _is_user_turn(message: _Message) -> bool: + return message.role != "assistant" and any(block.tool_result is None for block in message.content) + + def converse_rows(request: ConverseRequest, include_tool_calls: bool) -> tuple[ConverseRow, ...]: - messages: Final = ( + system: Final = _system_rows(request.system) + messages: Final = tuple( _message_rows(index, message, include_tool_calls) for index, message in enumerate(request.messages) ) - rows: Final = (*_system_rows(request.system), *chain.from_iterable(messages)) - return _with_request_fields(rows, request.additional_model_request_fields) + starts: Final = tuple(accumulate(map(len, messages), initial=len(system))) + rows: Final = (*system, *chain.from_iterable(messages)) + user_turn: Final = next( + (starts[index] for index in reversed(range(len(messages))) if _is_user_turn(request.messages[index])), + len(rows), + ) + return _with_request_fields(rows, request.additional_model_request_fields, user_turn) @dataclass(frozen=True, slots=True) diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 340de027bfa..a90b4829634 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -105,12 +105,7 @@ class _Remove: holder: _StringHolder -@dataclass(frozen=True, slots=True) -class _DropToolConfig: - pass - - -_Write: TypeAlias = _Assign | _DropToolConfig +_Write: TypeAlias = _Assign _ContentHolder: TypeAlias = tuple[object, str] _BlankDrop: TypeAlias = tuple[_ContentHolder, frozenset[int]] _BODY_BY_REFERENCE: Final = TypeAdapter(InstanceOf[dict[str, object]]) @@ -332,8 +327,8 @@ def _replaced_tools_write(body: Mapping[str, object], edit: ReplacedTools) -> _W """ Rebuild ``toolConfig.tools`` in the order the guardrail returned. Entries it was never sent (cache points, system tools) stay right after the tool they followed, and a cache point goes away with its tool. Refuses when - ``toolChoice`` names a tool that is no longer there, drops ``toolConfig`` when nothing is left, and creates it - when the request had none. + ``toolChoice`` names a tool that is no longer there, and creates ``toolConfig`` when the request had none. A + ``toolConfig`` left with no tool stays until the request is sent, so later guardrails still see the tool history. """ tool_config: Final = _child(body, "toolConfig") tools: Final = () if tool_config is None else as_json_array(_child(tool_config, "tools")) @@ -357,7 +352,7 @@ def _replaced_tools_write(body: Mapping[str, object], edit: ReplacedTools) -> _W return None if tool_config is None: return None if all(map(_is_cache_point, rebuilt)) else _Assign((body, "toolConfig"), {"tools": rebuilt}) - return _DropToolConfig() if all(map(_is_cache_point, rebuilt)) else _Assign((tool_config, "tools"), rebuilt) + return _Assign((tool_config, "tools"), rebuilt) def _resolved(body: Mapping[str, object], edit: Edit) -> _Write | None: @@ -372,11 +367,7 @@ def _resolved(body: Mapping[str, object], edit: Edit) -> _Write | None: def _kept_write(write: _Write, flat_changes: frozenset[_LeafIdentity]) -> bool: - match write: - case _Assign(): - return _leaf_identity(write.holder) not in flat_changes - case _DropToolConfig(): - return True + return _leaf_identity(write.holder) not in flat_changes def _holder_edits( @@ -397,52 +388,6 @@ def _apply(writes: Sequence[_Assign | _Remove]) -> None: del container[key] -def _tool_result_part(item: object) -> str | None: - text: Final = _child(item, "text") - if isinstance(text, str): - return text - fields: Final = as_json_object(item) - return json.dumps(fields["json"]) if fields is not None and "json" in fields else None - - -def _tool_result_text(tool_result: object) -> str: - parts: Final = as_json_array(_child(tool_result, "content")) or () - rendered: Final = "\n".join(part for item in parts if (part := _tool_result_part(item)) is not None).strip() - return rendered or "" - - -def _neutralized(block: object) -> object: - tool_use: Final = _child(block, "toolUse") - if as_json_object(tool_use) is not None: - arguments: Final = json.dumps(_child(tool_use, "input")) - return {"text": f"[tool call {_child(tool_use, 'toolUseId')}: {_child(tool_use, 'name')}({arguments})]"} - tool_result: Final = _child(block, "toolResult") - if as_json_object(tool_result) is None: - return block - return {"text": f"[tool result for {_child(tool_result, 'toolUseId')}: {_tool_result_text(tool_result)}]"} - - -def _has_tool_block(content: Sequence[object]) -> bool: - return any(_child(block, "toolUse") is not None or _child(block, "toolResult") is not None for block in content) - - -def _tool_config_drop(body: Mapping[str, object]) -> tuple[_Assign | _Remove, ...]: - """Remove ``toolConfig`` and turn every ``toolUse`` and ``toolResult`` block into text.""" - turns: Final = as_json_array(body.get("messages")) or () - neutralized: Final = tuple( - _Assign((message, "content"), [_neutralized(block) for block in content]) - for message in turns - if (content := as_json_array(_child(message, "content"))) is not None and _has_tool_block(content) - ) - return (*neutralized, _Remove((body, "toolConfig"))) - - -def _write_back_structured_edits(body: Mapping[str, object], writes: Sequence[_Write]) -> None: - _apply(tuple(write for write in writes if isinstance(write, _Assign))) - if any(isinstance(write, _DropToolConfig) for write in writes): - _apply(_tool_config_drop(body)) - - def _block_text_holder(block: object) -> _ContentHolder | None: if isinstance(_child(block, "text"), str): return block, "text" @@ -454,7 +399,7 @@ def _final_texts( holders: Sequence[_StringHolder], applied_texts: Sequence[str], writes: Sequence[_Write] ) -> Mapping[_LeafIdentity, object]: flat: Final = {_leaf_identity(holder): text for holder, text in zip(holders, applied_texts)} - structured: Final = {_leaf_identity(write.holder): write.value for write in writes if isinstance(write, _Assign)} + structured: Final = {_leaf_identity(write.holder): write.value for write in writes} return MappingProxyType({**flat, **structured}) @@ -478,16 +423,15 @@ def _empties_turn(drop: _BlankDrop) -> bool: return bool(dropped) and all(id(block) in dropped or _is_cache_point(block) for block in blocks) -def _blank_drops(body: Mapping[str, object], final: Mapping[_LeafIdentity, object]) -> tuple[_BlankDrop, ...] | None: - """The system and message text blocks the guardrail left blank, or None when dropping them empties a turn.""" - turns: Final = tuple( - ((message, "content"), _blank_blocks(final, (message, "content"))) - for message in as_json_array(body.get("messages")) or () - ) - if any(_empties_turn(turn) for turn in turns): - return None +def _turn_drop(final: Mapping[_LeafIdentity, object], message: object) -> _BlankDrop: + return (message, "content"), _blank_blocks(final, (message, "content")) + + +def _blank_drops(body: Mapping[str, object], final: Mapping[_LeafIdentity, object]) -> tuple[_BlankDrop, ...]: + """The system and message text blocks the guardrail left blank, except in a turn they would leave empty.""" + turns: Final = (_turn_drop(final, message) for message in as_json_array(body.get("messages")) or ()) system: Final = ((body, "system"), _blank_blocks(final, (body, "system"))) - return tuple(drop for drop in (system, *turns) if drop[1]) + return tuple(drop for drop in (system, *(turn for turn in turns if not _empties_turn(turn))) if drop[1]) def _blank_drop_write(drop: _BlankDrop) -> _Assign | _Remove: @@ -519,7 +463,7 @@ def _edits_drop_a_text_change( writes: Final = _holder_edits(rewritten, edits, frozenset()) if writes is None: return True - _write_back_structured_edits(rewritten, writes) + _apply(writes) edited: Final = [container[key] for container, key in holders] return any( returned != before and after == before for before, returned, after in zip(texts, guardrailed_texts, edited) @@ -1375,7 +1319,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): sent_messages: Final = [rows[index].message for index in scoped_indices] inputs: Final = GenericGuardrailAPIInputs(texts=list(texts)) if sent_messages: - inputs["structured_messages"] = copy.deepcopy(sent_messages) + inputs["structured_messages"] = json.loads(json.dumps(sent_messages)) tools: Final = converse_tools(request) if request and not scope.scan_only_tool_results else () if tools: inputs["tools"] = copy.deepcopy([tool.param for tool in tools]) @@ -1437,17 +1381,14 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): request_body, (*row_edits, *tool_changes), _changed_leaves(holders, texts, applied_texts) ) ) - blank_drops: Final = ( - None - if holder_edits is None - else _blank_drops(request_body, _final_texts(holders, applied_texts, holder_edits)) - ) - if holder_edits is None or blank_drops is None: + if holder_edits is None: raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) - if applied_texts: - _write_back_texts(applied_texts, holders) - _write_back_structured_edits(request_body, holder_edits) - _apply(tuple(_blank_drop_write(drop) for drop in blank_drops)) + if holder_edits or (applied_texts and list(applied_texts) != texts): + blank_drops: Final = _blank_drops(request_body, _final_texts(holders, applied_texts, holder_edits)) + if applied_texts: + _write_back_texts(applied_texts, holders) + _apply(holder_edits) + _apply(tuple(_blank_drop_write(drop) for drop in blank_drops)) return data diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index fb0ecd5b825..d7d6341fe7a 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -12,6 +12,7 @@ from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector +from litellm.llms.bedrock.chat.converse_to_openai import as_json_array, as_json_object from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage from litellm.types.utils import ModelResponseStream @@ -34,6 +35,78 @@ def is_bedrock_streaming_endpoint(endpoint: str) -> bool: return endpoint.partition("?")[0].rstrip("/").rsplit("/", 1)[-1] in BEDROCK_STREAMING_ACTIONS +_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) + + +def _is_converse_url(url: str) -> bool: + return httpx.URL(url).path.rstrip("/").rsplit("/", 1)[-1] in _CONVERSE_ACTIONS + + +def _is_cache_point(entry: object) -> bool: + fields: Final = as_json_object(entry) + return fields is not None and "cachePoint" in fields + + +def _tool_result_part(item: object) -> str | None: + fields: Final = as_json_object(item) + if fields is None: + return None + text: Final = fields.get("text") + if isinstance(text, str): + return text + return json.dumps(fields["json"]) if "json" in fields else None + + +def _tool_result_text(tool_result: Mapping[str, object]) -> str: + parts: Final = as_json_array(tool_result.get("content")) or () + rendered: Final = "\n".join(part for item in parts if (part := _tool_result_part(item)) is not None).strip() + return rendered or "" + + +def _block_field(block: object, key: str) -> Mapping[str, object] | None: + fields: Final = as_json_object(block) + return None if fields is None else as_json_object(fields.get(key)) + + +def _as_text_block(block: object) -> object: + tool_use: Final = _block_field(block, "toolUse") + if tool_use is not None: + arguments: Final = json.dumps(tool_use.get("input")) + return {"text": f"[tool call {tool_use.get('toolUseId')}: {tool_use.get('name')}({arguments})]"} + tool_result: Final = _block_field(block, "toolResult") + if tool_result is None: + return block + return {"text": f"[tool result for {tool_result.get('toolUseId')}: {_tool_result_text(tool_result)}]"} + + +def _has_tool_block(content: Sequence[object]) -> bool: + return any( + _block_field(block, "toolUse") is not None or _block_field(block, "toolResult") is not None for block in content + ) + + +def _as_text_message(message: object) -> object: + fields: Final = as_json_object(message) + content: Final = None if fields is None else as_json_array(fields.get("content")) + if fields is None or content is None or not _has_tool_block(content): + return message + return {**fields, "content": [_as_text_block(block) for block in content]} + + +def _converse_body_to_send(body: Mapping[str, object]) -> Mapping[str, object]: + """ + ``body`` without a Converse ``toolConfig`` that holds no tool, with every ``toolUse`` and ``toolResult`` turned into + text, since Bedrock rejects tool blocks without a tool. Any other ``body`` comes back as it is. + """ + tool_config: Final = as_json_object(body.get("toolConfig")) + tools: Final = None if tool_config is None else as_json_array(tool_config.get("tools")) + if tools is None or not all(map(_is_cache_point, tools)): + return body + messages: Final = as_json_array(body.get("messages")) + rest: Final = {key: value for key, value in body.items() if key != "toolConfig"} + return rest if messages is None else {**rest, "messages": [_as_text_message(message) for message in messages]} + + _TEXT_ONLY_DELTA_FIELDS: Final = frozenset({"content", "role"}) @@ -286,16 +359,17 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD self, headers: dict, litellm_params: dict, - request_data: dict | None, + request_data: dict[str, object] | None, api_base: str, model: str | None = None, ) -> tuple[dict, bytes | None]: + body: Final[Mapping[str, object]] = request_data or {} optional_params: Final = litellm_params.copy() return self._sign_request( service_name="bedrock", headers=headers, optional_params=optional_params, - request_data=request_data or {}, + request_data=dict(_converse_body_to_send(body) if _is_converse_url(api_base) else body), api_base=api_base, model=model, api_key=self.get_bedrock_bearer_token(optional_params), diff --git a/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py b/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py index 23a2951d1fb..08303236c09 100644 --- a/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py +++ b/tests/integration/observability/test_bedrock_converse_guardrail_rewrites.py @@ -261,6 +261,7 @@ def _guardrails(vendor_url: str) -> tuple[dict[str, JsonValue], ...]: _guardrail(_generic(vendor), {**generic, "api_base": f"{vendor_url}/gga/{vendor.value}"}) for vendor in _Vendor ), + _guardrail("mcp-end-user", {"guardrail": "mcp_end_user_permission"}), _guardrail( "gga-block-skip-tool", {**generic, "api_base": f"{vendor_url}/gga/{_Vendor.BLOCKS.value}", "skip_tool_message_in_guardrail": True}, @@ -274,7 +275,6 @@ def _guardrails(vendor_url: str) -> tuple[dict[str, JsonValue], ...]: }, ), _guardrail("tool-policy", {"guardrail": "tool_policy"}), - _guardrail("mcp-end-user", {"guardrail": "mcp_end_user_permission"}), ) @@ -329,7 +329,14 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: owned_proxy(gateway, directory, {}, config=_config(directory, guardrails, bedrock), workers=2) as owned, owned.scenario() as scenario, ): - yield _Rig(owned, vendor, bedrock, {name: scenario.key(guardrails=[name]) for name in names}) + keys: Final = {name: scenario.key(guardrails=[name]) for name in names} + tool_remover_then_masker: Final = ["mcp-end-user", "gga-mask-tool-only"] + yield _Rig( + owned, + vendor, + bedrock, + {**keys, "+".join(tool_remover_then_masker): scenario.key(guardrails=tool_remover_then_masker)}, + ) @pytest.fixture(autouse=True) @@ -579,9 +586,10 @@ def test_text_block_the_guardrail_cleared_is_dropped_before_bedrock(rig: _Rig) - assert sent["messages"] == [_user({"text": "Order details below."})] -def test_turn_the_guardrail_cleared_entirely_is_refused_before_bedrock(rig: _Rig) -> None: +def test_turn_the_guardrail_cleared_entirely_is_sent_as_written(rig: _Rig) -> None: body: Final = {"messages": [_user({"text": f"Card {_CARD} please."})]} - assert _UNAPPLIABLE in _refused(rig, _converse(rig, _generic(_Vendor.CLEARS_CARD_TEXTS), body)) + sent: Final = _forwarded(rig, _converse(rig, _generic(_Vendor.CLEARS_CARD_TEXTS), body)) + assert sent["messages"] == [_user({"text": ""})] def test_deeply_nested_tool_result_json_reaches_bedrock_with_the_message_mask(rig: _Rig) -> None: @@ -644,18 +652,12 @@ def test_mcp_end_user_permission_drops_tools_of_servers_the_end_user_cannot_use( @pytest.mark.parametrize( - ("guardrail_config", "result_block"), - [ - ({}, {"text": "[tool result for tooluse_post: posted to #ops]"}), - ( - {"guardrailConfig": {"guardrailIdentifier": "gr-scripted", "guardrailVersion": "1"}}, - {"guardContent": {"text": {"text": "[tool result for tooluse_post: posted to #ops]"}}}, - ), - ], + "guardrail_config", + [{}, {"guardrailConfig": {"guardrailIdentifier": "gr-scripted", "guardrailVersion": "1"}}], ids=["plain", "with-guardrail-config"], ) def test_mcp_end_user_permission_removing_every_tool_drops_the_tool_config_and_turns_tool_history_into_text( - rig: _Rig, guardrail_config: Mapping[str, JsonValue], result_block: Mapping[str, JsonValue] + rig: _Rig, guardrail_config: Mapping[str, JsonValue] ) -> None: tool_use: Final = {"toolUseId": "tooluse_post", "name": "slack-post_message", "input": {"channel": "#ops"}} body: Final = { @@ -682,6 +684,35 @@ def test_mcp_end_user_permission_removing_every_tool_drops_the_tool_config_and_t "messages": [ _user({"text": "Post the update."}), _assistant({"text": '[tool call tooluse_post: slack-post_message({"channel": "#ops"})]'}), - _user(dict(result_block), {"text": "Again."}), + _user({"text": "[tool result for tooluse_post: posted to #ops]"}, {"text": "Again."}), ], } + + +def test_tool_result_of_a_tool_mcp_end_user_permission_removed_is_masked_by_a_later_tool_results_only_guardrail( + rig: _Rig, +) -> None: + tool_use: Final = {"toolUseId": "tooluse_post", "name": "slack-post_message", "input": {"channel": "#ops"}} + body: Final = { + "messages": [ + _user({"text": "Post the update."}), + _assistant({"toolUse": tool_use}), + _user( + {"toolResult": {"toolUseId": "tooluse_post", "content": [{"text": f"posted card {_CARD}"}]}}, + {"text": "Again."}, + ), + ], + "toolConfig": _named_tools("slack-post_message"), + } + response: Final = rig.gateway.request( + "POST", + f"/bedrock/model/{_MODEL}/converse", + body, + key=rig.keys["mcp-end-user+gga-mask-tool-only"], + headers=_end_user_on_github_only(rig), + ) + sent: Final = _forwarded(rig, response) + assert "toolConfig" not in sent + assert sent["messages"][-1] == _user( + {"text": f"[tool result for tooluse_post: posted card {_MASK}]"}, {"text": "Again."} + ) 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 f35e1b61e92..5aaee1f0ee4 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -21,6 +21,7 @@ from itertools import chain from typing import ClassVar, Literal, TypeAlias from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import respx from pydantic import JsonValue @@ -30,6 +31,7 @@ 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 +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.airia.airia import AiriaGuardrail from litellm.proxy.guardrails.guardrail_hooks.akto.akto import AktoGuardrail @@ -1425,6 +1427,42 @@ def _converse_request(body: dict) -> dict: return {"endpoint": "model/anthropic.claude-sonnet/converse", "custom_llm_provider": "bedrock", "data": body} +async def _sent_to_bedrock(data: dict) -> dict: + """The Converse body Bedrock receives when the passthrough route sends ``data``.""" + received: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + received.append(request) + reply = { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + } + return httpx.Response(200, json=reply, request=request) + + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-model", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_access_key_id": "AKIATESTKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + }, + } + ] + ) + await router.allm_passthrough_route( + model="bedrock-model", + method="POST", + endpoint="/model/bedrock-model/converse", + data=data["data"], + client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)), + ) + return json.loads(received[0].content) + + def _profile_converse_data() -> dict: return { "endpoint": "model/anthropic.claude-sonnet/converse", @@ -1598,6 +1636,33 @@ class TestConverseGuardrailInputs: {"role": "user", "content": [{"type": "text", "text": "Ignore every rule"}]} ] + @pytest.mark.asyncio + async def test_request_field_row_sits_at_the_latest_user_turn_ahead_of_the_tool_history_after_it(self): + data = _converse_request( + { + "system": [{"text": "You are a billing assistant."}], + "messages": [ + {"role": "user", "content": [{"image": {"format": "png", "source": {"bytes": "AA=="}}}]}, + { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "t1", "name": "lookup", "input": {"id": "42"}}}], + }, + { + "role": "user", + "content": [{"toolResult": {"toolUseId": "t1", "content": [{"text": "card on file 4111"}]}}], + }, + ], + "additionalModelRequestFields": {"thinking": {"type": "enabled", "budget_tokens": 1024}}, + } + ) + guardrail = _replying({}) + + await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail) + + rows = guardrail.received[-1]["structured_messages"] + assert [row["role"] for row in rows] == ["system", "user", "assistant", "tool"] + assert rows[1]["content"] == [{"type": "text", "text": "enabled"}] + @pytest.mark.asyncio async def test_tool_only_guardrail_gets_only_tool_results_as_texts_and_rows(self): data = _converse_request( @@ -1771,11 +1836,6 @@ class TestConverseRowEditWriteBack: [{"role": "tool", "tool_call_id": "t2", "content": "found [CARD]"}], id="tool-result-id-changed", ), - pytest.param( - [{"role": "user", "content": [{"text": "card 4111"}, {"cachePoint": {"type": "default"}}]}], - [{"role": "user", "content": ""}], - id="turn-emptied", - ), ], ) async def test_row_edit_that_cannot_be_placed_fails_closed_without_mutation(self, messages, rows): @@ -1812,6 +1872,28 @@ class TestConverseRowEditWriteBack: {"role": "user", "content": [{"text": "card [CARD]"}, cache_point, masked_last_block]} ] + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("content", "blanked", "forwarded"), + [ + ([{"text": " 4111 1111 1111 1111 "}], " ", [{"text": " "}]), + ( + [{"text": "card 4111"}, {"cachePoint": {"type": "default"}}], + "", + [{"text": ""}, {"cachePoint": {"type": "default"}}], + ), + ], + ids=["whitespace-left", "nothing-left-before-a-cache-point"], + ) + async def test_turn_the_guardrail_blanks_entirely_is_forwarded_as_written(self, content, blanked, forwarded): + data = _converse_request({"messages": [{"role": "user", "content": content}]}) + + result = await BedrockPassthroughGuardrailHandler().process_input_messages( + data, _replying({"structured_messages": [{"role": "user", "content": blanked}]}) + ) + + assert result["data"]["messages"] == [{"role": "user", "content": forwarded}] + @pytest.mark.asyncio async def test_single_block_row_keeps_a_newline_the_guardrail_wrote(self): data = _converse_request({"messages": [{"role": "user", "content": [{"text": "card 4111"}]}]}) @@ -2026,7 +2108,7 @@ class TestConverseToolListRewrites: assert result["data"]["toolConfig"] == {"tools": [_REFUND, system_tool], "toolChoice": {"auto": {}}} @pytest.mark.asyncio - async def test_every_tool_removed_drops_the_tool_config_and_turns_tool_history_into_text(self): + async def test_every_tool_removed_reaches_bedrock_without_tool_config_and_with_tool_history_as_text(self): data = _tool_config_data() data["data"]["messages"] = [ { @@ -2055,7 +2137,8 @@ class TestConverseToolListRewrites: result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail) - assert result["data"] == { + assert result["data"]["toolConfig"] == {"tools": [], "toolChoice": {"auto": {}}} + assert await _sent_to_bedrock(result) == { "messages": [ { "role": "assistant", @@ -2371,9 +2454,46 @@ class TestConverseGuardContent: result = await handler.process_input_messages( without_tools, masker_class(lambda inputs: {"texts": _masked_texts(inputs)}) ) + sent = await _sent_to_bedrock(result) - assert "toolConfig" not in result["data"] - assert "4111" not in json.dumps(result["data"]) + assert "toolConfig" not in sent + assert "4111" not in json.dumps(sent) + + @pytest.mark.asyncio + async def test_tool_result_of_a_removed_tool_is_masked_by_a_later_tool_results_only_guardrail(self): + data = _converse_request( + { + "messages": [ + {"role": "user", "content": [{"text": "look up my order"}]}, + { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "t1", "name": "lookup", "input": {"q": "7"}}}], + }, + { + "role": "user", + "content": [ + {"toolResult": {"toolUseId": "t1", "content": [{"text": "card on file 4111"}]}}, + {"text": "summarize that"}, + ], + }, + ], + "toolConfig": { + "tools": [{"toolSpec": {"name": "lookup", "inputSchema": {"json": {"type": "object"}}}}] + }, + } + ) + handler = BedrockPassthroughGuardrailHandler() + tool_remover = _ScriptedGuardrail(lambda inputs: {"texts": inputs["texts"], "tools": []}) + results_masker = _ScriptedGuardrail(lambda inputs: {"texts": _masked_texts(inputs)}, scan_only_tool_results=True) + + without_tools = await handler.process_input_messages(data, tool_remover) + result = await handler.process_input_messages(without_tools, results_masker) + + assert results_masker.received[-1]["texts"] == ["card on file 4111"] + assert (await _sent_to_bedrock(result))["messages"][-1] == { + "role": "user", + "content": [{"text": "[tool result for t1: card on file [CARD]]"}, {"text": "summarize that"}], + } @pytest.mark.asyncio async def test_tool_history_written_as_text_adds_no_guard_content_block(self): @@ -2382,8 +2502,9 @@ class TestConverseGuardContent: tool_remover = _ScriptedGuardrail(lambda inputs: {"texts": inputs["texts"], "tools": []}) result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, tool_remover) + sent = await _sent_to_bedrock(result) - blocks = tuple(chain.from_iterable(message["content"] for message in result["data"]["messages"])) + blocks = tuple(chain.from_iterable(message["content"] for message in sent["messages"])) assert all(block.keys() == {"text"} for block in blocks), ( "a guardContent block makes Bedrock's own guardrail skip every plain text block", blocks,