Merge af38256db8 into 72228db5c6

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)

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

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"]