mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-24 05:45:11 +00:00
fix(guardrails): split a newline-joined Converse message back into its parts
A vendor that returns a multi-part message as one newline-joined string now has its masks written back to each part, guardContent included, when no part sent already holds a newline.
This commit is contained in:
parent
20c7fb359d
commit
a499c28029
2 changed files with 54 additions and 8 deletions
|
|
@ -166,12 +166,25 @@ def _carries_no_text(content: object) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _content_texts(row: ConverseRow, returned: Mapping[str, object]) -> tuple[str, ...] | None:
|
||||
content: Final = returned.get("content")
|
||||
sent: Final = _text_parts(row.message)
|
||||
if (
|
||||
isinstance(content, str)
|
||||
and len(row.content) > 1
|
||||
and sent is not None
|
||||
and all("\n" not in text for text in sent)
|
||||
):
|
||||
return tuple(content.split("\n"))
|
||||
return _text_parts(returned)
|
||||
|
||||
|
||||
def _content_edits(row: ConverseRow, original: object, returned: Mapping[str, object]) -> tuple[Edit, ...] | None:
|
||||
if returned.get("content") == original:
|
||||
return ()
|
||||
if not row.content and _carries_no_text(returned.get("content")):
|
||||
return ()
|
||||
texts: Final = _text_parts(returned)
|
||||
texts: Final = _content_texts(row, returned)
|
||||
if texts is None or len(texts) != len(row.content):
|
||||
return None
|
||||
return _merged(tuple(_slot_edits(slot, text) for slot, text in zip(row.content, texts)))
|
||||
|
|
|
|||
|
|
@ -1741,14 +1741,14 @@ class TestConverseRowEditWriteBack:
|
|||
id="role-changed",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"text": "card 4111"}, {"cachePoint": {"type": "default"}}, {"text": "thanks"}],
|
||||
}
|
||||
],
|
||||
[{"role": "user", "content": [{"text": "card\n4111"}, {"text": "thanks"}]}],
|
||||
[{"role": "user", "content": "card\n[CARD] thanks"}],
|
||||
id="blocks-that-hold-newlines-joined-with-a-space",
|
||||
),
|
||||
pytest.param(
|
||||
[{"role": "user", "content": [{"text": "card 4111"}, {"text": "thanks"}, {"text": "bye"}]}],
|
||||
[{"role": "user", "content": "card [CARD]\nthanks"}],
|
||||
id="blocks-merged-around-a-cache-point",
|
||||
id="newline-joined-blocks-with-one-missing",
|
||||
),
|
||||
pytest.param([_pay_turn(_CARD_TOOL_CALL)], [_pay_row(content="added text")], id="text-added-to-a-call"),
|
||||
pytest.param([_pay_turn(_CARD_TOOL_CALL)], [_pay_row(call_id="t2")], id="tool-call-id-changed"),
|
||||
|
|
@ -1789,6 +1789,39 @@ class TestConverseRowEditWriteBack:
|
|||
|
||||
assert data["data"] == original
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("last_block", "masked_last_block"),
|
||||
[
|
||||
({"text": "thanks 4111"}, {"text": "thanks [CARD]"}),
|
||||
({"guardContent": {"text": {"text": "thanks 4111"}}}, {"guardContent": {"text": {"text": "thanks [CARD]"}}}),
|
||||
],
|
||||
ids=["text", "guard-content"],
|
||||
)
|
||||
async def test_row_joined_with_newlines_is_split_back_into_its_blocks(self, last_block, masked_last_block):
|
||||
cache_point = {"cachePoint": {"type": "default"}}
|
||||
data = _converse_request(
|
||||
{"messages": [{"role": "user", "content": [{"text": "card 4111"}, cache_point, last_block]}]}
|
||||
)
|
||||
|
||||
result = await BedrockPassthroughGuardrailHandler().process_input_messages(
|
||||
data, _replying({"structured_messages": [{"role": "user", "content": "card [CARD]\nthanks [CARD]"}]})
|
||||
)
|
||||
|
||||
assert result["data"]["messages"] == [
|
||||
{"role": "user", "content": [{"text": "card [CARD]"}, cache_point, masked_last_block]}
|
||||
]
|
||||
|
||||
@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"}]}]})
|
||||
|
||||
result = await BedrockPassthroughGuardrailHandler().process_input_messages(
|
||||
data, _replying({"structured_messages": [{"role": "user", "content": "card\n[CARD]"}]})
|
||||
)
|
||||
|
||||
assert result["data"]["messages"] == [{"role": "user", "content": [{"text": "card\n[CARD]"}]}]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("content", [None, "", []], ids=["none", "empty-string", "empty-list"])
|
||||
async def test_tool_call_argument_mask_in_a_row_is_written_back(self, content):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue