diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 44ae57f81db..a61f516cf57 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -19,6 +19,7 @@ from typing import ( Dict, List, Literal, + Mapping, NamedTuple, Optional, Tuple, @@ -103,6 +104,28 @@ class GuardrailMessageFilterResult(NamedTuple): target_indices: Optional[List[int]] +def _tool_result_texts(tool_result_content: object) -> list[str]: + """Ordered scannable text carried by an Anthropic ``tool_result`` block. + + A ``tool_result`` stores its text one level deeper than an ordinary content + block: under ``content``, which is either a string or a list of blocks whose + ``text`` field holds the string (images/documents carry no text). Returning + the strings in document order lets the scan payload and the masking write-back + agree on how many text units the block contributes and in what sequence. + """ + if isinstance(tool_result_content, str): + return [tool_result_content] + if isinstance(tool_result_content, list): + return [ + text + for block in tool_result_content + if isinstance(block, dict) + for text in (block.get("text"),) + if isinstance(text, str) + ] + return [] + + class ApplyGuardrailMessageSelection(NamedTuple): """Messages selected for an apply_guardrail scan + write-back metadata.""" @@ -294,11 +317,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): blocks.append(QualifiedTextBlock(text=content, qualifier=None)) elif isinstance(content, list): for item in content: - if isinstance(item, dict) and "text" in item: + if isinstance(item, str): + blocks.append(QualifiedTextBlock(text=item, qualifier=None)) + elif isinstance(item, dict) and "text" in item: qualifier = _CONTENT_TYPE_TO_QUALIFIER.get(item.get("type", "")) blocks.append(QualifiedTextBlock(text=item["text"], qualifier=qualifier)) - elif isinstance(item, str): - blocks.append(QualifiedTextBlock(text=item, qualifier=None)) + elif isinstance(item, dict): + tool_result = cast(Mapping[str, object], item) # cast-ok: dynamic Anthropic tool_result JSON + if tool_result.get("type") == "tool_result": + blocks.extend( + QualifiedTextBlock(text=text, qualifier=None) + for text in _tool_result_texts(tool_result.get("content")) + ) return blocks def _build_content_item(self, block: QualifiedTextBlock) -> BedrockContentItem: @@ -1566,6 +1596,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): new_item["text"] = masked_texts[masking_index] masking_index += 1 new_content.append(new_item) + elif isinstance(item, dict) and item.get("type") == "tool_result": + masked_item, masking_index = self._mask_tool_result_item( + tool_result=item, masked_texts=masked_texts, masking_index=masking_index + ) + new_content.append(masked_item) elif isinstance(item, str): if masking_index < len(masked_texts): item = masked_texts[masking_index] @@ -1575,6 +1610,38 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return new_content, masking_index + def _mask_tool_result_item( + self, tool_result: dict, masked_texts: list[str], masking_index: int + ) -> tuple[dict, int]: + """Write masked text back into an Anthropic ``tool_result`` block. + + Mirrors :func:`_tool_result_texts`: the block's scannable text lives under + ``content`` as either a string or a list of ``text`` blocks, so the same + units are consumed here in the same order the scan payload sent them. This + keeps ``masking_index`` aligned with Bedrock's per-block masked outputs and + anonymizes the tool output itself. + """ + new_item = tool_result.copy() + tool_content = new_item.get("content") + if isinstance(tool_content, str): + if masking_index < len(masked_texts): + new_item["content"] = masked_texts[masking_index] + masking_index += 1 + return new_item, masking_index + if isinstance(tool_content, list): + new_blocks: list[Any] = [] + for block in tool_content: + if isinstance(block, dict) and isinstance(block.get("text"), str): + new_block = block.copy() + if masking_index < len(masked_texts): + new_block["text"] = masked_texts[masking_index] + masking_index += 1 + new_blocks.append(new_block) + else: + new_blocks.append(block) + new_item["content"] = new_blocks + return new_item, masking_index + def _apply_masking_to_response( self, response: Union[ModelResponse, Any], diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 15827b80bcf..1740addbda0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2769,6 +2769,174 @@ async def test_grounding_output_blocked_raises_400(): assert exc_info.value.status_code == 400 +############################################################################### +# #33086: Anthropic tool_result content blocks reach the native pre_call / +# during_call hooks with their text nested under `content` (a string or a list +# of `text` blocks), not a top-level `text` key. Before the fix the INPUT scan +# dropped tool_result text entirely, so PII returned by a tool call bypassed +# Bedrock scanning; the ANONYMIZE write-back must also mask that text and stay +# index-aligned with Bedrock's per-block masked outputs. +############################################################################### + + +def test_input_scans_tool_result_string_content(): + """A tool_result whose `content` is a plain string must be sent to Bedrock.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": "SSN 324-12-3212", + } + ], + } + ] + + assert _input_request(messages) == { + "source": "INPUT", + "content": [{"text": {"text": "SSN 324-12-3212"}}], + } + + +def test_input_scans_tool_result_list_content_skipping_non_text(): + """text blocks nested in a tool_result are scanned; images/documents are skipped.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [ + {"type": "text", "text": "email a@b.com"}, + {"type": "image", "source": {"type": "base64", "data": "xx"}}, + {"type": "text", "text": "phone 607-456-7890"}, + ], + } + ], + } + ] + + assert _input_request(messages) == { + "source": "INPUT", + "content": [ + {"text": {"text": "email a@b.com"}}, + {"text": {"text": "phone 607-456-7890"}}, + ], + } + + +def test_input_scans_tool_result_alongside_plain_text_in_order(): + """A plain text block and a tool_result in the same message keep document order.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "please review"}, + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": "secret data", + }, + ], + } + ] + + assert _input_request(messages) == { + "source": "INPUT", + "content": [ + {"text": {"text": "please review"}}, + {"text": {"text": "secret data"}}, + ], + } + + +def test_masking_writes_back_into_tool_result_and_keeps_alignment(): + """Masked outputs return 1:1 in scan order. Because the tool_result text was + scanned it owns its masked output, so the write-back masks the tool_result and + the following text block stays aligned with its own masked output (before the + fix the tool_result was skipped, shifting every later block's mask).""" + guardrail = _grounding_guardrail() + messages = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": "SSN 324-12-3212", + }, + {"type": "text", "text": "and my name is John Smith"}, + ], + } + ] + + masked = guardrail._apply_masking_to_messages( + messages=messages, + masked_texts=["SSN {US_SSN}", "and my name is {NAME}"], + ) + + assert masked == [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": "SSN {US_SSN}", + }, + {"type": "text", "text": "and my name is {NAME}"}, + ], + } + ] + + +def test_masking_writes_back_into_tool_result_list_blocks(): + """Each text block nested in a tool_result is masked in order; non-text blocks + are left untouched and do not consume a masked output.""" + guardrail = _grounding_guardrail() + messages = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [ + {"type": "text", "text": "email a@b.com"}, + {"type": "image", "source": {"data": "xx"}}, + {"type": "text", "text": "phone 607-456-7890"}, + ], + } + ], + } + ] + + masked = guardrail._apply_masking_to_messages( + messages=messages, + masked_texts=["email {EMAIL}", "phone {PHONE}"], + ) + + assert masked == [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": [ + {"type": "text", "text": "email {EMAIL}"}, + {"type": "image", "source": {"data": "xx"}}, + {"type": "text", "text": "phone {PHONE}"}, + ], + } + ], + } + ] + + ############################################################################### # LIT-4186: disable_exception_on_block regression tests #