mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
test(guardrails): type the new tool_permission pre_call test parameters
This commit is contained in:
parent
9d23a8ad03
commit
c25cee4e92
1 changed files with 10 additions and 4 deletions
|
|
@ -5,6 +5,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics)
|
|||
import json
|
||||
import logging
|
||||
import re
|
||||
from typing import Literal
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -24,6 +25,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
|
|||
PermissionError,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
CallTypesLiteral,
|
||||
ChatCompletionMessageToolCall,
|
||||
Choices,
|
||||
ModelResponse,
|
||||
|
|
@ -1177,7 +1179,7 @@ class TestToolPermissionGuardrailAnthropicMessages:
|
|||
def _tool_use(self, name, tool_id="tu_1"):
|
||||
return {"type": "tool_use", "id": tool_id, "name": name, "input": {"command": "ls"}}
|
||||
|
||||
def _always_on_pre_call(self, on_disallowed_action):
|
||||
def _always_on_pre_call(self, on_disallowed_action: Literal["block", "rewrite"]) -> ToolPermissionGuardrail:
|
||||
return ToolPermissionGuardrail(
|
||||
guardrail_name=f"anthropic-pre-call-{on_disallowed_action}",
|
||||
rules=self.rules,
|
||||
|
|
@ -1197,7 +1199,7 @@ class TestToolPermissionGuardrailAnthropicMessages:
|
|||
],
|
||||
ids=["anthropic", "anthropic_custom_type", "responses_api_flat_function"],
|
||||
)
|
||||
async def test_pre_call_blocks_denied_request_tool_in_flat_format(self, denied_tool):
|
||||
async def test_pre_call_blocks_denied_request_tool_in_flat_format(self, denied_tool: dict[str, object]) -> None:
|
||||
data = {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}], "tools": [denied_tool]}
|
||||
|
||||
with pytest.raises(HTTPException) as excinfo:
|
||||
|
|
@ -1231,8 +1233,12 @@ class TestToolPermissionGuardrailAnthropicMessages:
|
|||
ids=["anthropic", "responses_api"],
|
||||
)
|
||||
async def test_pre_call_rewrite_strips_denied_flat_tool_and_forced_choice(
|
||||
self, tool_shape, tool_choice, call_type, expected_tool_choice
|
||||
):
|
||||
self,
|
||||
tool_shape: dict[str, object],
|
||||
tool_choice: dict[str, str],
|
||||
call_type: CallTypesLiteral,
|
||||
expected_tool_choice: dict[str, str] | str,
|
||||
) -> None:
|
||||
data = {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue