Repository navigation
fix(generic_guardrail_api): return tool_calls on the guardrail response - #41013
ninadphalak wants to merge 5 commits into
Conversation
`GenericGuardrailAPIRequest` carries `tool_calls` to the guardrail, but
`GenericGuardrailAPIResponse` has no matching field and
`_build_guardrail_return_inputs` never set the key, so a guardrail that rewrites
tool arguments cannot hand them back. The read side already exists:
- `litellm/llms/openai/chat/guardrail_translation/handler.py` reads
`guardrailed_inputs.get("tool_calls", [])`, and the output path reads
`returned_tool_calls` with a length guard
- `litellm/llms/anthropic/chat/guardrail_translation/handler.py` reads
`returned_tool_calls`, also length-guarded
- `litellm/llms/openai/responses/guardrail_translation/handler.py` passes it to
`_post_guardrail_tool_call_shapes`, which falls back and logs "guardrail
returned no tool calls for the N scanned" whenever the key is absent
Only the generic hook's write side was missing, so all of those lookups missed and
a guardrail's returned tool calls were dropped.
Add the field and thread it through, mirroring `tools`/`images`: the guardrail's
replacements win, otherwise the calls that were sent are echoed back unchanged.
Every consumer compares the returned list against the one it sent and falls back to
its own copy on a length mismatch, so an unchanged echo leaves existing guardrails
behaving exactly as before.
 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
…tra comments A same-length list of tool calls with no function block was forwarded as-is, and the chat handlers index function.name and function.arguments on whatever comes back, so an unusable list raised there instead of returning a response. Returned calls are now checked for those two fields, and the calls already sent are kept when one is missing them, which is the fallback the responses handler already takes. The new field is annotated with the request's own field type rather than spelling the list union out twice, which keeps it exact and keeps the type-discipline budget flat, and the comment blocks that restated the code are gone.
| @pytest.mark.asyncio | ||
| async def test_guardrail_returned_tool_calls_flow_back_to_inputs(self, generic_guardrail): | ||
| """A guardrail that rewrites tool arguments has to be able to return them. | ||
|
|
||
| ``tool_calls`` is sent on the request, so the matching response field is what closes | ||
| the loop. Without it the calls were dropped on the floor and the caller silently | ||
| kept the values it sent in -- for a rewriting guardrail (reversible redaction, for | ||
| instance) the rewritten arguments never reached the client. | ||
| """ |
There was a problem hiding this comment.
These docstrings restate straightforward behavior, including similar blocks near lines 1937 and 1999. Repository policy requires removing them before merge
Context Used: CLAUDE.md (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
There was a problem hiding this comment.
Removed the docstrings from the new tests in e5cab21, and from the new source helpers too, per the comment rule in AGENTS.md
… never rebind The type alias that named this field read the request's field through a TypedDict subscript, which ruff reads as a forward reference, and it failed the basedpyright budget. The annotations spell the union out again, which is exact and checks against GenericGuardrailAPIInputs without an alias and without a suppression comment. The unusable-reply path no longer rebinds a Final local, so it needs no suppression comment either: the calls that will be applied are picked in one expression.
Resolve conflicts with structured_messages, which landed in the same two files; keep both fields. The response now parses tool_calls in from_dict (tool_calls_from_response): every row needs a function name and string arguments, otherwise the list is dropped with a warning and the calls that were sent are kept. Fields are read-only Sequences, and the sent calls are echoed by _echo_sent_tool_calls in apply_guardrail, so the change adds no LIT001 or basedpyright findings.
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
| if rows_to_write_back is not None: | ||
| return_inputs["structured_messages"] = list(rows_to_write_back) | ||
| if guardrail_response.tool_calls: | ||
| return_inputs["tool_calls"] = list(guardrail_response.tool_calls) |
There was a problem hiding this comment.
Medium: Multi-choice streams bypass text redaction
In a post-call policy pipeline using this guardrail, a caller can request stream: true, n: 2 and receive original, unredacted assistant text when the guardrail also rewrites a tool call: _write_ended_stream_tool_call_rewrites rejects multi-choice streams, and PipelineExecutor._run_streaming_step catches that rejection and restores all original chunks. The PR changes this integration from omitting tool-call rewrites to forwarding them, introducing a rollback path that also discards text redactions.
Before forwarding these rewrites on this path, implement per-choice tool-call delivery or make unsupported rewrites fail closed rather than releasing the original stream.
There was a problem hiding this comment.
Leaving this: outside pipelines it fails closed, and inside them the executor rolls back every guardrail's undeliverable rewrite. Happy to change the executor separately
PR overviewThis PR updates the generic guardrail API integration to return tool-call rewrites in guardrail responses and forward them through response processing. One issue remains open, with none addressed so far. A caller requesting a multi-choice stream can receive original, unredacted assistant text when the guardrail also rewrites a tool call: the unsupported rewrite triggers a fallback that restores the original chunks, discarding text redactions. This creates a conditional bypass of post-call redaction policies for affected streaming responses. Open issues (1)
Fixed/addressed: 0 · PR risk: 6/10 |
|
Merged main and fixed the lint failures. CI is green except osv-scan, which fails on every open PR. @yucheng-berri could you review this? |
|
@greptileai please re-review |
|
Added a live Docker run with Gemini to the description: main-latest returns the sent postal code, this branch returns the guardrail's |
This branch has not been deployed
1 participant
TLDR
Problem this solves:
How it solves it:
tool_calls, parsed infrom_dictbytool_calls_from_responseUser Flow
Before: an operator rewrites tool arguments in a guardrail and the rewritten values never reach the application
"tools"present and gets backfinish_reason: "tool_calls"After: the same request comes back with the guardrail's values in the tool call
argumentsPre-Submission checklist
tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py, 80 passedosv-scan, which fails on every open PR right now (asharpadvisory inui/litellm-dashboard/package-lock.json)Screenshots / Proof of Fix
A real Gemini call through the proxy, with a generic guardrail that rewrites postal code 94107 to 10001 in returned tool calls. Before is
ghcr.io/berriai/litellm:main-latest(sha256:67831417c141), After is the same image with this branch'slitellm/mounted over the installed packageBefore, the guardrail returns 10001 and the client still gets 94107
After (
LITELLM_SRC=/path/to/this/branch docker compose -f compose.yaml -f pr.yaml up -d), the client gets the guardrail's valuecompose.yaml, pr.yaml, config.yaml, request.json, rewrite_guardrail.py
{"model": "gemini-3.7-flash", "messages": [{"role": "user", "content": "Ship order 1182 to postal code 94107."}], "tools": [{"type": "function", "function": {"name": "ship_order", "description": "Ship an order", "parameters": {"type": "object", "properties": {"order_id": {"type": "string"}, "postal_code": {"type": "string"}}, "required": ["order_id", "postal_code"]}}}], "tool_choice": {"type": "function", "function": {"name": "ship_order"}}}Mutation check on the 8 new unit tests, each mutation reverted afterwards:
tool_callsback in_build_guardrail_return_inputs: 2 fail_echo_sent_tool_calls: 2 failtool_calls_from_response: 2 failType of Pull Request
🐛 Bug Fix
Caveats (if any)
Low
incremental_diffis documented for stringdelta.contentonlyPipelineExecutor, the same rule every guardrail there follows. Outside a pipeline the request fails instead