mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
fix(generic_guardrail_api): validate returned tool calls, drop the extra 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.
This commit is contained in:
parent
0b19b55b3c
commit
f0f06aec07
3 changed files with 54 additions and 17 deletions
|
|
@ -24,14 +24,15 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.llms.openai import ChatCompletionToolCallChunk, ChatCompletionToolParam
|
||||
from litellm.types.llms.openai import ChatCompletionToolParam
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
GenericGuardrailAPIMetadata,
|
||||
GenericGuardrailAPIRequest,
|
||||
GenericGuardrailAPIResponse,
|
||||
GenericGuardrailAPIToolCalls,
|
||||
GuardrailToolParam,
|
||||
)
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
|
@ -104,6 +105,16 @@ def _sanitize_inbound_headers(
|
|||
return sanitized or None
|
||||
|
||||
|
||||
def _returned_tool_call_is_usable(tool_call: object) -> bool:
|
||||
"""Whether a guardrail's returned tool call carries what the chat handlers index."""
|
||||
function: Final = tool_call.get("function") if isinstance(tool_call, dict) else getattr(tool_call, "function", None)
|
||||
if function is None:
|
||||
return False
|
||||
name: Final = function.get("name") if isinstance(function, dict) else getattr(function, "name", None)
|
||||
arguments: Final = function.get("arguments") if isinstance(function, dict) else getattr(function, "arguments", None)
|
||||
return isinstance(name, str) and isinstance(arguments, str)
|
||||
|
||||
|
||||
def _extract_inbound_headers(
|
||||
request_data: dict,
|
||||
logging_obj: Optional["LiteLLMLoggingObj"],
|
||||
|
|
@ -323,7 +334,7 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
images: list[str] | None,
|
||||
tools: list[ChatCompletionToolParam] | None,
|
||||
guardrail_response: GenericGuardrailAPIResponse,
|
||||
tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] | None = None,
|
||||
tool_calls: GenericGuardrailAPIToolCalls | None = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
# Action is NONE or no modifications needed
|
||||
return_inputs: Final = GenericGuardrailAPIInputs(texts=texts)
|
||||
|
|
@ -337,12 +348,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
|
||||
returned_tool_calls: Final = guardrail_response.tool_calls
|
||||
if returned_tool_calls and not all(_returned_tool_call_is_usable(call) for call in returned_tool_calls):
|
||||
verbose_proxy_logger.warning(
|
||||
"Generic Guardrail API returned tool calls missing a function name or arguments; keeping the ones already sent",
|
||||
)
|
||||
returned_tool_calls = None # rebind-ok: an unusable reply is discarded
|
||||
if returned_tool_calls:
|
||||
return_inputs["tool_calls"] = returned_tool_calls
|
||||
elif tool_calls:
|
||||
return_inputs["tool_calls"] = tool_calls
|
||||
if guardrail_response.stream_holdback_chars is not None:
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from typing import Any, Final, Literal
|
||||
from typing import Any, Final, Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from typing_extensions import TypedDict
|
||||
|
|
@ -8,7 +8,9 @@ from litellm.types.llms.openai import (
|
|||
ChatCompletionToolCallChunk,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs
|
||||
|
||||
GenericGuardrailAPIToolCalls: TypeAlias = GenericGuardrailAPIInputs["tool_calls"] # noqa: F821 # TypedDict field reference, not a name
|
||||
|
||||
|
||||
class GuardrailToolParam(BaseModel):
|
||||
|
|
@ -167,7 +169,7 @@ class GenericGuardrailAPIResponse:
|
|||
action: str
|
||||
blocked_reason: str | None
|
||||
stream_holdback_chars: list[int] | None
|
||||
tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] | None
|
||||
tool_calls: GenericGuardrailAPIToolCalls | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -177,7 +179,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,
|
||||
tool_calls: GenericGuardrailAPIToolCalls | None = None,
|
||||
) -> None:
|
||||
self.action = action
|
||||
self.blocked_reason = blocked_reason
|
||||
|
|
@ -188,10 +190,6 @@ 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
|
||||
|
|
|
|||
|
|
@ -1999,6 +1999,32 @@ class TestToolSupport:
|
|||
|
||||
assert result["tool_calls"] == sent
|
||||
|
||||
def test_malformed_returned_tool_calls_are_discarded(self, generic_guardrail):
|
||||
"""A same-length list with no function block must not reach chat handling.
|
||||
|
||||
The chat handlers index ``function.name`` and ``function.arguments`` on whatever the
|
||||
guardrail returns, so forwarding an unusable list raises there instead of returning a
|
||||
response. The calls that were sent are kept instead.
|
||||
"""
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
GenericGuardrailAPIResponse,
|
||||
)
|
||||
|
||||
sent = [{"id": "call_1", "type": "function", "function": {"name": "send_email", "arguments": "{}"}}]
|
||||
malformed = [{"id": "call_1", "type": "function"}]
|
||||
|
||||
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": malformed}
|
||||
),
|
||||
)
|
||||
|
||||
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 (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue