scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-14 04:09:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 97

fix(generic_guardrail_api): return tool_calls on the guardrail response

Browse source 
`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.
This commit is contained in:
Ninad Phalak 2026-09-13 13:40:18 -05:00
parent c2c2a623c0
commit 0b19b55b3c
No known key found for this signature in database
GPG key ID: 59119ED515433744
3 changed files with 158 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py
View file

@ -24,14 +24,14 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.llms.openai import ChatCompletionToolCallChunk, ChatCompletionToolParam
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIMetadata,
GenericGuardrailAPIRequest,
GenericGuardrailAPIResponse,
GuardrailToolParam,
)
from litellm.types.utils import GenericGuardrailAPIInputs
from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -323,6 +323,7 @@ class GenericGuardrailAPI(CustomGuardrail):
images: list[str] | None,
tools: list[ChatCompletionToolParam] | None,
guardrail_response: GenericGuardrailAPIResponse,
tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] | None = None,
) -> GenericGuardrailAPIInputs:
# Action is NONE or no modifications needed
return_inputs: Final = GenericGuardrailAPIInputs(texts=texts)
@ -336,6 +337,14 @@ class GenericGuardrailAPI(CustomGuardrail):
return_inputs["tools"] = guardrail_response.tools
elif tools:
return_inputs["tools"] = tools
# Mirrors ``tools``: the guardrail's replacements win, otherwise the calls
# are echoed back unchanged. Handlers compare the returned list against the
# one they sent and fall back to their own copy on a length mismatch, so a
# guardrail that cannot reorder them still cannot change their count.
if guardrail_response.tool_calls:
return_inputs["tool_calls"] = guardrail_response.tool_calls
elif tool_calls:
return_inputs["tool_calls"] = tool_calls
if guardrail_response.stream_holdback_chars is not None:
return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars
return return_inputs
@ -473,6 +482,7 @@ class GenericGuardrailAPI(CustomGuardrail):
texts=texts,
images=images,
tools=tools,
tool_calls=tool_calls,
guardrail_response=guardrail_response,
)

8
litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py
View file

@ -167,6 +167,7 @@ class GenericGuardrailAPIResponse:
action: str
blocked_reason: str | None
stream_holdback_chars: list[int] | None
tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] | None
def __init__(
self,
@ -176,6 +177,7 @@ class GenericGuardrailAPIResponse:
images: list[str] | None = None,
tools: list[GuardrailToolParam] | None = None,
stream_holdback_chars: list[int] | None = None,
tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] | None = None,
) -> None:
self.action = action
self.blocked_reason = blocked_reason
@ -186,6 +188,11 @@ class GenericGuardrailAPIResponse:
# framework must withhold from streaming emission until the next
# processing round (word-boundary safety for text transformations).
self.stream_holdback_chars = stream_holdback_chars
# Replacement tool calls. ``tool_calls`` is sent on the request, so the
# response mirrors it back the way ``tools`` and ``images`` are: without
# this field a guardrail that rewrites tool arguments has no channel to
# return them on, and the caller keeps the values it sent in.
self.tool_calls = tool_calls
@classmethod
def from_dict(cls, data: dict) -> "GenericGuardrailAPIResponse":
@ -200,4 +207,5 @@ class GenericGuardrailAPIResponse:
images=data.get("images"),
tools=data.get("tools"),
stream_holdback_chars=stream_holdback_chars,
tool_calls=data.get("tool_calls"),
)

138
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py
View file

@ -1380,6 +1380,35 @@ class TestGenericGuardrailAPIResponseParsing:
assert response.stream_holdback_chars == [3, 0, 0, 0]
def test_from_dict_parses_tool_calls(self):
"""``tool_calls`` goes out on the request, so the response has to carry it back."""
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIResponse,
)
tool_calls = [
{
"id": "call_1",
"type": "function",
"function": {"name": "send_email", "arguments": '{"to": "jane@example.com"}'},
}
]
response = GenericGuardrailAPIResponse.from_dict(
{"action": "GUARDRAIL_INTERVENED", "texts": ["Done"], "tool_calls": tool_calls}
)
assert response.tool_calls == tool_calls
def test_from_dict_tool_calls_absent_is_none(self):
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIResponse,
)
response = GenericGuardrailAPIResponse.from_dict({"action": "NONE", "texts": ["hi"]})
assert response.tool_calls is None
@pytest.mark.asyncio
async def test_apply_guardrail_flows_holdback_back_to_inputs(self, generic_guardrail):
"""A GUARDRAIL_INTERVENED response with stream_holdback_chars is surfaced on
@ -1876,6 +1905,115 @@ class TestToolSupport:
assert forwarded_tools == tools
@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.
"""
guardrailed_tool_calls = [
{
"id": "call_1",
"type": "function",
"function": {"name": "send_email", "arguments": '{"to": "jane@example.com"}'},
}
]
mock_response = MagicMock()
mock_response.json.return_value = {
"action": "GUARDRAIL_INTERVENED",
"texts": ["Done"],
"tool_calls": guardrailed_tool_calls,
}
mock_response.raise_for_status = MagicMock()
with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response):
result = await generic_guardrail.apply_guardrail(
inputs={"texts": ["Done"]},
request_data={},
input_type="response",
)
assert result["tool_calls"] == guardrailed_tool_calls
def test_guardrail_tool_calls_take_precedence_over_the_sent_ones(self, generic_guardrail):
"""The returned calls win over the ones that went out.
Built through ``_build_guardrail_return_inputs`` rather than ``apply_guardrail`` so
the assertion is about this mapping, not about request-model coercion of the calls
the caller supplied.
"""
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIResponse,
)
sent = [{"id": "call_1", "type": "function", "function": {"name": "send_email", "arguments": "{}"}}]
returned = [
{
"id": "call_1",
"type": "function",
"function": {"name": "send_email", "arguments": '{"to": "jane@example.com"}'},
}
]
result = generic_guardrail._build_guardrail_return_inputs(
texts=["Done"],
images=None,
tools=None,
tool_calls=sent,
guardrail_response=GenericGuardrailAPIResponse.from_dict(
{"action": "GUARDRAIL_INTERVENED", "texts": ["Done"], "tool_calls": returned}
),
)
assert result["tool_calls"] == returned
def test_untouched_tool_calls_are_echoed_not_dropped(self, generic_guardrail):
"""A guardrail that ignores tool calls must not blank them out.
The downstream handlers compare what comes back against what they sent and fall back
to their own copy when the lengths differ, so echoing an unchanged list leaves them
behaving exactly as they did before this field existed.
"""
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIResponse,
)
sent = [
{
"id": "call_1",
"type": "function",
"function": {"name": "send_email", "arguments": '{"to": "jane@example.com"}'},
}
]
result = generic_guardrail._build_guardrail_return_inputs(
texts=["Done"],
images=None,
tools=None,
tool_calls=sent,
guardrail_response=GenericGuardrailAPIResponse.from_dict({"action": "NONE", "texts": ["Done"]}),
)
assert result["tool_calls"] == sent
def test_no_tool_calls_on_either_side_stay_absent(self, generic_guardrail):
"""An empty list is not the same as no list: the key must stay out of the result."""
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIResponse,
)
result = generic_guardrail._build_guardrail_return_inputs(
texts=["Done"],
images=None,
tools=None,
guardrail_response=GenericGuardrailAPIResponse.from_dict({"action": "NONE", "texts": ["Done"]}),
)
assert "tool_calls" not in result
class TestFailOnError:
"""Test fail_on_error: complete fail-open on any guardrail error"""