mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
fix(bedrock): send a reshaped Converse tool result as one text block
A vendor that compresses or summarises a tool result, such as Headroom collapsing a JSON list into a table string, returned a 500 because the row could not be split back into the original text and JSON parts. The tool result now carries the vendor's string as one text block where its parts were, keeping images and documents in place. It still fails closed when the flat texts masked that tool result or the row changed its tool call id.
This commit is contained in:
yucheng-berri
parent
e0c91941e8
commit
7eb0c68d2b
4 changed files with 215 additions and 19 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue