diff --git a/litellm/litellm_core_utils/llm_response_utils/tool_calls.py b/litellm/litellm_core_utils/llm_response_utils/tool_calls.py new file mode 100644 index 00000000000..d59ff9019c4 --- /dev/null +++ b/litellm/litellm_core_utils/llm_response_utils/tool_calls.py @@ -0,0 +1,125 @@ +import json +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from typing import Final, Literal + +from pydantic import InstanceOf, TypeAdapter, ValidationError + +from litellm.litellm_core_utils.core_helpers import as_object_list +from litellm.types.utils import ChatCompletionMessageToolCall, Function + +_FIELDS: Final = TypeAdapter(InstanceOf[dict[str, object]]) + + +def _fields( + value: object, +) -> dict[str, object] | None: # mutable-ok: rewrites must update the original provider container + try: + return _FIELDS.validate_python(value) + except ValidationError: + return None + + +def _items(value: object) -> tuple[object, ...]: + return tuple(as_object_list(value) or ()) + + +@dataclass(frozen=True, slots=True) +class RawToolCallGroup: + format: Literal["openai", "legacy", "gemini", "bedrock"] + message: dict[str, object] # mutable-ok: rewrite only denied calls in the original message + finish: dict[str, object] # mutable-ok: synchronize the original finish reason after removing calls + blocks: tuple[object, ...] + calls: tuple[ChatCompletionMessageToolCall | None, ...] + + def rewrite(self, errors: Mapping[str, str]) -> None: + messages: Final = tuple(errors[call.id] for call in self.calls if call is not None and call.id in errors) + if not messages: + return + kept: Final = [block for block, call in zip(self.blocks, self.calls) if call is None or call.id not in errors] + text: Final = "\n".join(messages) + match self.format: + case "openai" | "legacy": + key: Final = "tool_calls" if self.format == "openai" else "function_call" + self.message[key] = (kept or None) if self.format == "openai" else None + content: Final = self.message.get("content") + parts: Final = as_object_list(content) + self.message["content"] = ( + [*parts, {"type": "text", "text": text}] + if parts is not None + else f"{content}\n\n{text}" + if content + else text + ) + if not self.message.get("tool_calls") and not self.message.get("function_call"): + self.finish["finish_reason"] = "stop" + case "gemini": + self.message["parts"] = [*kept, {"text": text}] + case "bedrock": + self.message["content"] = [*kept, {"text": text}] + if not any((fields := _fields(block)) is not None and "toolUse" in fields for block in kept): + self.finish["stopReason"] = "end_turn" + + +def _call(block: object, format: str, identifier: str) -> ChatCompletionMessageToolCall | None: + fields: Final = _fields(block) + if fields is None: + return None + function_key: Final = {"openai": "function", "gemini": "functionCall", "bedrock": "toolUse"}.get(format) + function: Final = fields if function_key is None else _fields(fields.get(function_key)) + if function is None: + return None + name: Final = function.get("name") + if not isinstance(name, str) or not name: + return None + arguments_key: Final = "args" if format == "gemini" else "input" if format == "bedrock" else "arguments" + arguments: Final = function.get(arguments_key) + serialized: Final = ( + arguments if isinstance(arguments, str) else json.dumps(arguments) if arguments is not None else "" + ) + return ChatCompletionMessageToolCall( + id=identifier, type="function", function=Function(name=name, arguments=serialized) + ) + + +def _group( + format: Literal["openai", "legacy", "gemini", "bedrock"], + message: dict[str, object], # mutable-ok: the group retains the provider message for in-place rewrites + finish: dict[str, object], # mutable-ok: the group retains the provider finish reason for in-place rewrites + blocks: tuple[object, ...], + index: int, +) -> RawToolCallGroup: + # Positions keep missing or duplicate provider IDs from removing an allowed call. + calls: Final = tuple( + _call(block, format, f"raw_{format}_{index}_{position}") for position, block in enumerate(blocks) + ) + return RawToolCallGroup(format, message, finish, blocks, calls) + + +def _choice_groups(index: int, choice: object) -> Iterator[RawToolCallGroup]: + choice_fields: Final = _fields(choice) + message: Final = _fields(choice_fields.get("message")) if choice_fields is not None else None + if message is not None and choice_fields is not None: + yield _group("openai", message, choice_fields, _items(message.get("tool_calls")), index) + yield _group("legacy", message, choice_fields, (message.get("function_call"),), index) + + +def _candidate_groups(index: int, candidate: object) -> Iterator[RawToolCallGroup]: + candidate_fields: Final = _fields(candidate) + content: Final = _fields(candidate_fields.get("content")) if candidate_fields is not None else None + if content is not None and candidate_fields is not None: + yield _group("gemini", content, candidate_fields, _items(content.get("parts")), index) + + +def raw_tool_call_groups(response: object) -> Iterator[RawToolCallGroup]: + fields: Final = _fields(response) + if fields is None: + return + for index, choice in enumerate(_items(fields.get("choices"))): + yield from _choice_groups(index, choice) + for index, candidate in enumerate(_items(fields.get("candidates"))): + yield from _candidate_groups(index, candidate) + output: Final = _fields(fields.get("output")) + bedrock_message: Final = _fields(output.get("message")) if output is not None else None + if bedrock_message is not None: + yield _group("bedrock", bedrock_message, fields, _items(bedrock_message.get("content")), 0) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index f2697c0aa11..8c5bb58f1f3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -23,6 +23,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none +from litellm.litellm_core_utils.llm_response_utils.tool_calls import raw_tool_call_groups from litellm.litellm_core_utils.tool_fields import declared_function_name, qualified_member_name from litellm.llms.anthropic.common_utils import ( anthropic_server_tool_call_params, @@ -1612,7 +1613,7 @@ class ToolPermissionGuardrail(CustomGuardrail): anthropic_content: Final = ( None if isinstance(response, ModelResponse) else self._get_anthropic_content_blocks(response) ) - if not isinstance(response, ModelResponse) and anthropic_content is None: + if not isinstance(response, (ModelResponse, dict)): return response verbose_proxy_logger.debug("Tool Permission Guardrail Post-Call Hook: Checking response") @@ -1621,11 +1622,14 @@ class ToolPermissionGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Tool Permission Guardrail: Skipping check (not enabled)") return response - # Extract tool_calls from the response + raw_groups: Final = tuple(raw_tool_call_groups(response)) if anthropic_content is None else () + raw_calls: Final = chain.from_iterable(group.calls for group in raw_groups) tool_calls: Final = ( self._extract_tool_calls_from_response(response) if isinstance(response, ModelResponse) else self._extract_tool_calls_from_anthropic_content(anthropic_content or ()) + if anthropic_content is not None + else tuple(call for call in raw_calls if call is not None) ) if not tool_calls: @@ -1640,8 +1644,14 @@ class ToolPermissionGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Tool Permission Guardrail Post-Call Hook: All tools allowed") elif isinstance(response, ModelResponse): self._modify_response_with_permission_errors(response, denied_tools) - else: + elif anthropic_content is not None: self._modify_anthropic_content_with_permission_errors(response, anthropic_content or (), denied_tools) + else: + errors: Final = { + call.id: self._create_permission_error_result(call, error).content for call, error in denied_tools + } + for group in raw_groups: + group.rewrite(errors) add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return response diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission_raw_response.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission_raw_response.py new file mode 100644 index 00000000000..15dd2449fbc --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission_raw_response.py @@ -0,0 +1,261 @@ +import copy +import json +from typing import Final, Literal + +import pytest + +from litellm.exceptions import GuardrailRaisedException +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.tool_permission import ToolPermissionGuardrail +from litellm.types.guardrails import GuardrailEventHooks + +_Format = Literal["openai", "legacy", "gemini", "bedrock"] + + +def _response(format: _Format, names: tuple[str, ...], with_ids: bool = True) -> dict[str, object]: + arguments: Final = {"path": "/private/secret"} + match format: + case "openai": + calls: Final = [ + { + **({"id": "duplicate-id"} if with_ids else {}), + "type": "function", + "function": {"name": name, "arguments": json.dumps(arguments)}, + "provider_extra": name, + } + for name in names + ] + return { + "choices": [ + { + "message": {"role": "assistant", "content": "working", "tool_calls": calls}, + "finish_reason": "tool_calls", + "logprobs": {"retained": True}, + } + ], + "usage": {"total_tokens": 17}, + "provider_extra": {"request_id": "keep"}, + } + case "legacy": + return { + "choices": [ + { + "message": { + "role": "assistant", + "content": "working", + "function_call": {"name": names[0], "arguments": json.dumps(arguments)}, + }, + "finish_reason": "function_call", + "logprobs": {"retained": True}, + } + ], + "usage": {"total_tokens": 17}, + "provider_extra": {"request_id": "keep"}, + } + case "gemini": + parts: Final = [ + {"functionCall": {"name": name, "args": arguments}, "thoughtSignature": name} for name in names + ] + return { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "working"}, *parts]}, + "finishReason": "STOP", + "safetyRatings": [{"blocked": False}], + } + ], + "usageMetadata": {"totalTokenCount": 17}, + "provider_extra": {"request_id": "keep"}, + } + case "bedrock": + blocks: Final = [ + {"toolUse": {**({"toolUseId": "duplicate-id"} if with_ids else {}), "name": name, "input": arguments}} + for name in names + ] + return { + "output": {"message": {"role": "assistant", "content": [{"text": "working"}, *blocks]}}, + "stopReason": "tool_use", + "usage": {"totalTokens": 17}, + "provider_extra": {"request_id": "keep"}, + } + + +def _guardrail(action: Literal["block", "rewrite"]) -> ToolPermissionGuardrail: + return ToolPermissionGuardrail( + guardrail_name="raw-tools", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + default_action="allow", + on_disallowed_action=action, + rules=[{"id": "deny-read", "tool_name": "^Read$", "decision": "deny"}], + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("format", ["openai", "legacy", "gemini", "bedrock"]) +async def test_raw_response_blocks_denied_tool_without_mutating_response(format: _Format) -> None: + response: Final = _response(format, ("Read",)) + before: Final = copy.deepcopy(response) + + with pytest.raises(GuardrailRaisedException, match="Read"): + await _guardrail("block").async_post_call_success_hook( + data={}, user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + assert response == before + + +@pytest.mark.asyncio +@pytest.mark.parametrize("format", ["openai", "legacy", "gemini", "bedrock"]) +async def test_raw_response_allowed_tool_is_identity_passthrough(format: _Format) -> None: + response: Final = _response(format, ("Bash",)) + before: Final = copy.deepcopy(response) + + result: Final = await _guardrail("rewrite").async_post_call_success_hook( + data={}, user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + assert result is response + assert response == before + + +@pytest.mark.asyncio +@pytest.mark.parametrize("format", ["openai", "legacy", "gemini", "bedrock"]) +@pytest.mark.parametrize("with_ids", [True, False]) +async def test_raw_response_rewrite_removes_denied_call_and_preserves_metadata(format: _Format, with_ids: bool) -> None: + response: Final = _response(format, ("Read", "Bash"), with_ids) + before: Final = copy.deepcopy(response) + + result: Final = await _guardrail("rewrite").async_post_call_success_hook( + data={}, user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + assert result is response + assert response["provider_extra"] == before["provider_extra"] + assert "Permission denied" in json.dumps(response) + match response: + case {"choices": [{"message": message, "finish_reason": finish, "logprobs": logprobs}]}: + assert logprobs == {"retained": True} + assert response["usage"] == before["usage"] + assert message["content"].startswith("working\n\nPermission denied") + if format == "legacy": + assert message["function_call"] is None + assert finish == "stop" + else: + assert [call["function"]["name"] for call in message["tool_calls"]] == ["Bash"] + assert message["tool_calls"][0]["provider_extra"] == "Bash" + assert finish == "tool_calls" + case {"candidates": [{"content": {"parts": parts}, "safetyRatings": ratings, "finishReason": finish}]}: + assert parts[0] == {"text": "working"} + assert parts[1] == { + "functionCall": {"name": "Bash", "args": {"path": "/private/secret"}}, + "thoughtSignature": "Bash", + } + assert len(parts) == 3 + assert ratings == [{"blocked": False}] + assert finish == "STOP" + assert response["usageMetadata"] == before["usageMetadata"] + case {"output": {"message": {"content": blocks}}, "stopReason": stop}: + assert blocks[0] == {"text": "working"} + assert blocks[1]["toolUse"]["name"] == "Bash" + assert len(blocks) == 3 + assert stop == "tool_use" + assert response["usage"] == before["usage"] + case _: + pytest.fail("response shape changed") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("format", ["openai", "legacy", "gemini", "bedrock"]) +async def test_raw_response_rewrite_ends_turn_when_all_tools_denied(format: _Format) -> None: + response: Final = _response(format, ("Read",)) + + await _guardrail("rewrite").async_post_call_success_hook( + data={}, user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + assert "/private/secret" not in json.dumps(response) + match response: + case {"choices": [{"message": message, "finish_reason": finish}]}: + assert not message.get("tool_calls") + assert not message.get("function_call") + assert finish == "stop" + case {"candidates": [{"content": {"parts": parts}}]}: + assert all("functionCall" not in part for part in parts) + case {"output": {"message": {"content": blocks}}, "stopReason": stop}: + assert all("toolUse" not in block for block in blocks) + assert stop == "end_turn" + case _: + pytest.fail("response shape changed") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("format", ["openai", "legacy", "gemini", "bedrock"]) +async def test_raw_response_checks_argument_rules(format: _Format) -> None: + guardrail: Final = ToolPermissionGuardrail( + guardrail_name="raw-paths", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + default_action="deny", + on_disallowed_action="block", + rules=[ + { + "id": "public-read", + "tool_name": "^Read$", + "decision": "allow", + "allowed_param_patterns": {"path": "^/public/.*$"}, + } + ], + ) + + with pytest.raises(GuardrailRaisedException, match="path"): + await guardrail.async_post_call_success_hook( + data={}, user_api_key_dict=UserAPIKeyAuth(), response=_response(format, ("Read",)) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("format", ["openai", "gemini"]) +async def test_raw_response_rewrite_isolates_candidates(format: _Format) -> None: + key: Final = "choices" if format == "openai" else "candidates" + allowed: Final = _response(format, ("Bash",))[key] + denied: Final = _response(format, ("Read",))[key] + assert isinstance(allowed, list) + assert isinstance(denied, list) + response: Final = {key: [*allowed, *denied]} + before: Final = copy.deepcopy(allowed) + + result: Final = await _guardrail("rewrite").async_post_call_success_hook( + data={}, user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + assert result[key][0] == before[0] + assert "Permission denied" in json.dumps(result[key][1]) + assert "/private/secret" not in json.dumps(result[key][1]) + + +@pytest.mark.asyncio +async def test_raw_openai_rewrite_retains_content_parts() -> None: + content: Final = [{"type": "text", "text": "working", "provider_extra": 17}] + response: Final = { + "choices": [ + { + "message": { + "role": "assistant", + "content": content, + "tool_calls": [{"type": "function", "function": {"name": "Read", "arguments": "{}"}}], + }, + "finish_reason": "tool_calls", + } + ] + } + + result: Final = await _guardrail("rewrite").async_post_call_success_hook( + data={}, user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + parts: Final = result["choices"][0]["message"]["content"] + assert parts[0] is content[0] + assert len(parts) == 2 + assert parts[1]["type"] == "text" + assert "Permission denied" in parts[1]["text"]