mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-30 06:43:27 +00:00
fix(guardrails): cover Responses tool call items and instructions
Responses carries the same tool call arguments and system prompt as chat and messages under function_call.arguments, custom_tool_call.input and instructions, so secrets there reached the provider unscanned
This commit is contained in:
parent
e734719f50
commit
e488c1a169
3 changed files with 105 additions and 12 deletions
|
|
@ -201,6 +201,9 @@ def _iter_native_content_text(content: object) -> Iterator[str]:
|
|||
# built-in shape; ``custom_tool_call_output`` is the custom-tool
|
||||
# counterpart (see ``ChatCompletionCustomToolCallOutput``).
|
||||
_OUTPUT_ITEM_TYPES: Final[frozenset[str]] = frozenset({"function_call_output", "custom_tool_call_output"})
|
||||
# Responses input items carrying the model's tool call: ``function_call`` holds
|
||||
# JSON ``arguments`` and ``custom_tool_call`` holds a freeform ``input`` string.
|
||||
_CALL_ITEM_FIELDS: Final[Mapping[str, str]] = {"function_call": "arguments", "custom_tool_call": "input"}
|
||||
|
||||
|
||||
def _part_text(part: Mapping[str, object]) -> str | None:
|
||||
|
|
@ -278,12 +281,13 @@ def _iter_inspection_messages(data: Mapping[str, object]) -> Iterator[object]:
|
|||
|
||||
|
||||
def iter_message_text(data: Mapping[str, object]) -> Iterator[str]:
|
||||
"""Yield every text fragment from ``messages`` AND ``input``.
|
||||
"""Yield every text fragment from ``system``, ``instructions``, ``messages`` AND ``input``.
|
||||
|
||||
Walks every role (user, assistant, system, …) — guardrails inspect
|
||||
the entire conversation, not just user turns.
|
||||
"""
|
||||
yield from _iter_native_content_text(data.get("system"))
|
||||
yield from _iter_native_content_text(data.get("instructions"))
|
||||
for message in _iter_inspection_messages(data):
|
||||
fields = _text_object(message)
|
||||
if fields is None:
|
||||
|
|
@ -295,16 +299,31 @@ def iter_message_text(data: Mapping[str, object]) -> Iterator[str]:
|
|||
call_fields = _text_object(call)
|
||||
function = _text_object(call_fields.get("function")) if call_fields is not None else None
|
||||
if function is not None:
|
||||
arguments = function.get("arguments")
|
||||
if isinstance(arguments, str) and arguments:
|
||||
try:
|
||||
parsed = _tool_arguments(arguments)
|
||||
except (ValueError, RecursionError):
|
||||
yield arguments
|
||||
else:
|
||||
yield from _tool_payload_text(parsed)
|
||||
elif isinstance(arguments, dict):
|
||||
yield from _tool_payload_text(arguments)
|
||||
yield from _iter_arguments_text(function.get("arguments"))
|
||||
items: Final = _object_list(data.get("input")) if isinstance(data.get("input"), list) else None
|
||||
for item in items or ():
|
||||
call = _call_item(item)
|
||||
if call is not None:
|
||||
yield from _iter_arguments_text(call[1])
|
||||
|
||||
|
||||
def _call_item(item: object) -> tuple[str, object] | None:
|
||||
"""Return ``(field, value)`` for a Responses tool-call input item, else None."""
|
||||
fields: Final = _text_object(item)
|
||||
field: Final = _CALL_ITEM_FIELDS.get(str(fields.get("type"))) if fields is not None else None
|
||||
return (field, fields.get(field)) if fields is not None and field is not None and field in fields else None
|
||||
|
||||
|
||||
def _iter_arguments_text(arguments: object) -> Iterator[str]:
|
||||
if isinstance(arguments, str) and arguments:
|
||||
try:
|
||||
parsed: Final = _tool_arguments(arguments)
|
||||
except (ValueError, RecursionError):
|
||||
yield arguments
|
||||
else:
|
||||
yield from _tool_payload_text(parsed)
|
||||
elif isinstance(arguments, dict):
|
||||
yield from _tool_payload_text(arguments)
|
||||
|
||||
|
||||
def walk_user_text(
|
||||
|
|
@ -314,7 +333,8 @@ def walk_user_text(
|
|||
) -> int:
|
||||
"""Rewrite every text fragment in place via ``visit``.
|
||||
|
||||
Mutates ``data["messages"]`` (including tool payloads), ``data["system"]`` and ``data["input"]``. Returns the number
|
||||
Mutates ``data["messages"]`` and ``data["input"]`` (including tool payloads), ``data["system"]`` and
|
||||
``data["instructions"]``. Returns the number
|
||||
of fragments visited so callers can short-circuit when nothing was
|
||||
inspected.
|
||||
"""
|
||||
|
|
@ -393,6 +413,8 @@ def walk_user_text(
|
|||
|
||||
if "system" in data:
|
||||
data["system"] = _rewrite_content(data["system"]) # rebind-ok: preserve the native hook's in-place contract
|
||||
if "instructions" in data:
|
||||
data["instructions"] = _rewrite_content(data["instructions"]) # rebind-ok: Responses system prompt, in place
|
||||
|
||||
input_value: Final = data.get("input")
|
||||
if isinstance(input_value, str):
|
||||
|
|
@ -419,6 +441,8 @@ def walk_user_text(
|
|||
item["content"] = _rewrite_content(item["content"])
|
||||
elif item.get("type") in _OUTPUT_ITEM_TYPES and "output" in item:
|
||||
item["output"] = _rewrite_content(item["output"])
|
||||
elif (call := _call_item(item)) is not None:
|
||||
item[call[0]] = _rewrite_arguments(call[1])
|
||||
return visited
|
||||
|
||||
return visited
|
||||
|
|
|
|||
|
|
@ -353,6 +353,45 @@ def test_existing_text_shapes_still_redact(caller: Caller, path: str) -> None:
|
|||
_stored(caller, request_id)
|
||||
|
||||
|
||||
_RESPONSES_ITEMS: Final[dict[str, tuple[JsonValue, ...]]] = {
|
||||
"instructions": (),
|
||||
"function-call": (
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": "call_1",
|
||||
"name": "write_file",
|
||||
"arguments": json.dumps({"path": ".env", "content": TEXT}),
|
||||
},
|
||||
{"type": "function_call_output", "call_id": "call_1", "output": "written"},
|
||||
),
|
||||
"custom-tool-call": (
|
||||
{"type": "custom_tool_call", "call_id": "call_1", "name": "shell", "input": f"export {TEXT}"},
|
||||
{"type": "custom_tool_call_output", "call_id": "call_1", "output": "done"},
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", tuple(_RESPONSES_ITEMS))
|
||||
def test_responses_tool_calls_and_instructions_are_redacted(caller: Caller, shape: str) -> None:
|
||||
marker: Final = _identity()
|
||||
body: Final = {
|
||||
"model": "gpt",
|
||||
"cache": {"no-cache": True},
|
||||
"input": [{"role": "user", "content": f"{marker} write the env file"}, *_RESPONSES_ITEMS[shape]],
|
||||
**({"instructions": TEXT} if shape == "instructions" else {}),
|
||||
}
|
||||
response: Final = caller.rig.proxy.request(
|
||||
"POST", "/v1/responses", body, key=caller.key, headers={"x-litellm-call-id": marker}
|
||||
)
|
||||
assert response.status_code == 200 and "accepted" in response.text, response.text
|
||||
sent: Final = _wire(caller, marker)
|
||||
if shape == "function-call":
|
||||
assert json.loads(sent["input"][1]["arguments"]) == {"path": ".env", "content": "AWS_ACCESS_KEY_ID=[REDACTED]"}
|
||||
if shape == "custom-tool-call":
|
||||
assert sent["input"][1]["input"] == "export AWS_ACCESS_KEY_ID=[REDACTED]"
|
||||
_stored(caller, marker)
|
||||
|
||||
|
||||
def test_cache_hit_keeps_masked_guardrail_metadata(caller: Caller) -> None:
|
||||
marker: Final = _identity()
|
||||
body: Final = {**_body("openai", marker), "cache": {"ttl": 60}}
|
||||
|
|
|
|||
|
|
@ -707,6 +707,36 @@ def test_walk_user_text_redacts_custom_tool_call_output():
|
|||
assert data["input"][0]["output"] == "[MASKED]"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("item", "field", "masked"),
|
||||
(
|
||||
(
|
||||
{"type": "function_call", "name": "secret_tool_name", "arguments": '{"token": "secret", "count": 7}'},
|
||||
"arguments",
|
||||
'{"token": "[MASKED]", "count": 7}',
|
||||
),
|
||||
({"type": "function_call", "arguments": '{"token":"secret",broken'}, "arguments", '{"token":"[MASKED]",broken'),
|
||||
({"type": "custom_tool_call", "name": "secret_tool_name", "input": "echo secret"}, "input", "echo [MASKED]"),
|
||||
),
|
||||
ids=("function-call", "function-call-malformed", "custom-tool-call"),
|
||||
)
|
||||
def test_responses_tool_call_items_are_read_and_masked(item, field, masked):
|
||||
data = {"input": [{"role": "user", "content": "go"}, item]}
|
||||
assert any("secret" in text and "secret_tool_name" not in text for text in iter_message_text(data))
|
||||
assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 2
|
||||
assert data["input"][1][field] == masked
|
||||
assert data["input"][1].get("name", "secret_tool_name") == "secret_tool_name"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("instructions", ("secret rules", [{"type": "input_text", "text": "secret rules"}]))
|
||||
def test_responses_instructions_are_read_and_masked(instructions):
|
||||
data = {"instructions": instructions, "input": "hi"}
|
||||
assert "secret rules" in tuple(iter_message_text(data))
|
||||
assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]")) == 2
|
||||
masked = data["instructions"] if isinstance(instructions, str) else data["instructions"][0]["text"]
|
||||
assert masked == "[MASKED] rules"
|
||||
|
||||
|
||||
def test_build_inspection_messages_custom_tool_call_output():
|
||||
"""build_inspection_messages should include custom_tool_call_output text."""
|
||||
data = {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue