scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-30 06:43:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 250

fix(guardrails): write a tool visitor's whole payload when it changes more than string values

Browse source
This commit is contained in:
yucheng-berri 2026-10-17 15:42:51 -07:00
parent 8749ad8ade
commit bcf5f2002c
2 changed files with 24 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
litellm/proxy/guardrails/_content_utils.py
View file

@ -116,17 +116,17 @@ def _value_string_spans(arguments: str) -> Iterator[tuple[int, int]]:
def _respliced(arguments: str, parsed: JsonValue, rewritten: JsonValue) -> str:
"""Write the string values that differ between ``parsed`` and ``rewritten`` into ``arguments``, keeping its other text."""
spans: Final = tuple(_value_string_spans(arguments))
before: Final = tuple(iter_tool_json_text(parsed))
after: Final = tuple(iter_tool_json_text(rewritten))
if not len(spans) == len(before) == len(after):
return json.dumps(rewritten, ensure_ascii=False)
changed: Final = tuple((span, new) for span, old, new in zip(spans, before, after) if new != old)
"""Serialize ``rewritten`` by replacing only its changed string values in ``arguments``, when nothing else changed."""
serialized: Final = json.dumps(rewritten, ensure_ascii=False)
strings: Final = iter(tuple(iter_tool_json_text(rewritten)))
if json.dumps(map_tool_json_text(parsed, lambda text: next(strings, text)), ensure_ascii=False) != serialized:
return serialized
values: Final = zip(_value_string_spans(arguments), iter_tool_json_text(parsed), iter_tool_json_text(rewritten))
changed: Final = tuple((span, new) for span, old, new in values if new != old)
starts: Final = (0, *(end for (_, end), _ in changed))
ends: Final = (*(start for (start, _), _ in changed), len(arguments))
replacements: Final = (*(json.dumps(new, ensure_ascii=False) for _, new in changed), "")
return "".join(arguments[start:end] + replacement for start, end, replacement in zip(starts, ends, replacements))
return "".join(arguments[start:end] + text for start, end, text in zip(starts, ends, replacements))
def iter_tool_json_fields(value: object) -> Iterator[tuple[str | None, str]]:

20
tests/unit/proxy/guardrails/test_content_utils.py
View file

@ -1044,10 +1044,22 @@ def test_masked_tool_arguments_change_only_the_masked_strings(arguments: str, ma
assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == masked
def test_tool_visitor_that_changes_the_payload_shape_writes_its_payload_as_arguments() -> None:
data = {"messages": [{"tool_calls": [{"function": {"arguments": '{ "token": "secret", "note": "keep" }'}}]}]}
assert walk_user_text(data, lambda text: text, lambda payload: {"token": "[MASKED]"}) == 1
assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == '{"token": "[MASKED]"}'
@pytest.mark.parametrize(
"rewritten",
(
{"token": "[MASKED]"},
{"token": "[MASKED]", "renamed": "keep", "count": 7},
{"token": "[MASKED]", "note": "keep", "count": 8},
{"token": "[MASKED]", "note": "keep", "count": True},
),
ids=("key-dropped", "key-renamed", "number-changed", "number-retyped"),
)
def test_tool_visitor_that_changes_more_than_strings_writes_its_payload_as_arguments(rewritten: JsonValue) -> None:
data = {
"messages": [{"tool_calls": [{"function": {"arguments": '{ "token": "secret", "note": "keep", "count": 1 }'}}]}]
}
assert walk_user_text(data, lambda text: text, lambda payload: rewritten) == 1
assert data["messages"][0]["tool_calls"][0]["function"]["arguments"] == json.dumps(rewritten)
@pytest.mark.parametrize(