scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-31 06:56:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 245

Merge 31774a7b2e into a392a39627

Browse source
This commit is contained in:
SilkForgeAi 2026-10-31 14:33:28 +08:00 • committed by GitHub
commit 45772e19ba
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 173 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

77
litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py
View file

@ -91,12 +91,13 @@ def _get_request_route_from_data(request_data: dict) -> str | None:
def _resolve_tool_names_from_messages(messages: list[dict]) -> dict[str, str]:
"""
Build a map of tool_call_id -> tool_name from assistant messages' tool_calls.
Build a map of tool_call_id -> tool_name from assistant messages' tool_calls
(OpenAI Chat Completions) and tool_use content blocks (Anthropic /v1/messages).
Used to resolve which tool produced each tool result in the conversation.
"""
mapping: Final[dict[str, str]] = {}
for msg in messages:
if msg.get("role") != "assistant":
if not isinstance(msg, dict) or msg.get("role") != "assistant":
continue
tool_calls = msg.get("tool_calls") or []
for tc in tool_calls:
@ -109,9 +110,63 @@ def _resolve_tool_names_from_messages(messages: list[dict]) -> dict[str, str]:
fn = getattr(fn_obj, "name", None) if fn_obj else None
if tc_id and fn:
mapping[tc_id] = fn
content = msg.get("content")
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_use":
tc_id = block.get("id")
fn = block.get("name")
if isinstance(tc_id, str) and tc_id and isinstance(fn, str) and fn:
mapping[tc_id] = fn
return mapping
def _extract_executed_tool_names(request_data: dict) -> list[str]:
"""
Return the ordered list of tool names whose outputs appear in the conversation
across OpenAI Chat Completions (`role == "tool"`), Anthropic `/v1/messages`
(`type == "tool_result"` blocks), and `/v1/responses` (`type == "function_call_output"`).
"""
executed: Final[list[str]] = []
messages: Final = request_data.get("messages") or []
if isinstance(messages, list) and messages:
tc_id_to_name: Final = _resolve_tool_names_from_messages(messages)
for msg in messages:
if not isinstance(msg, dict):
continue
if msg.get("role") == "tool":
tool_call_id = msg.get("tool_call_id")
source_tool = tc_id_to_name.get(tool_call_id, "") if isinstance(tool_call_id, str) else ""
if source_tool:
executed.append(source_tool)
content = msg.get("content")
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_result":
tool_use_id = block.get("tool_use_id")
source_tool = tc_id_to_name.get(tool_use_id, "") if isinstance(tool_use_id, str) else ""
if source_tool:
executed.append(source_tool)
input_items: Final = request_data.get("input")
if isinstance(input_items, list) and input_items:
call_id_to_name: Final[dict[str, str]] = {}
for item in input_items:
if isinstance(item, dict) and item.get("type") == "function_call":
call_id = item.get("call_id") or item.get("id")
fn = item.get("name")
if isinstance(call_id, str) and call_id and isinstance(fn, str) and fn:
call_id_to_name[call_id] = fn
for item in input_items:
if isinstance(item, dict) and item.get("type") == "function_call_output":
call_id = item.get("call_id")
source_tool = call_id_to_name.get(call_id, "") if isinstance(call_id, str) else ""
if source_tool:
executed.append(source_tool)
return executed
class ToolPolicyGuardrail(CustomGuardrail):
"""
Guardrail that enforces per-tool input/output policies from the in-memory
@ -146,11 +201,10 @@ class ToolPolicyGuardrail(CustomGuardrail):
if isinstance(t, dict) and isinstance(t.get("function"), dict) and t["function"].get("name")
]
route: Final = _get_request_route_from_data(request_data)
route_tool_names: Final = extract_request_tool_names(route, request_data) if route else []
replayed_tool_names: Final = replayed_tool_search_tool_names(_REQUEST_DATA.validate_python(request_data))
tool_names = (
extract_request_tool_names(route, request_data)
if route and not function_tool_names
else [*function_tool_names, *replayed_tool_names]
tool_names = list(
dict.fromkeys([*function_tool_names, *route_tool_names, *replayed_tool_names])
)
else:
tool_calls: Final = inputs.get("tool_calls") or []
@ -204,17 +258,8 @@ class ToolPolicyGuardrail(CustomGuardrail):
if input_type == "response":
trusted_input_tools: Final = [name for name in tool_names if policy_map.get(name) == "trusted"]
if trusted_input_tools:
messages: Final = request_data.get("messages") or []
tc_id_to_name: Final = _resolve_tool_names_from_messages(messages)
untrusted_sources: Final[list[str]] = []
for msg in messages:
if msg.get("role") != "tool":
continue
tool_call_id = msg.get("tool_call_id")
source_tool = tc_id_to_name.get(tool_call_id, "") if tool_call_id else ""
if not source_tool:
continue
for source_tool in _extract_executed_tool_names(request_data):
if registry.get_output_policy(source_tool) == "untrusted":
if source_tool not in untrusted_sources:
untrusted_sources.append(source_tool)

112
tests/unit/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py
View file

@ -222,3 +222,115 @@ async def test_blocked_tool_loaded_by_a_replayed_tool_search_raises_next_to_allo
)
assert exc_info.value.status_code == 400
assert exc_info.value.detail["blocked_tools"] == ["run_shell"]
def _registry_mock_with_output(input_policies: dict, output_policies: dict):
reg = _registry_mock(input_policies)
reg.get_output_policy.side_effect = lambda name: output_policies.get(name, "untrusted")
return reg
@pytest.mark.asyncio
async def test_anthropic_messages_tool_result_blocks_trusted_tool_call(guardrail):
request_data = {
"user_api_key_request_route": "/v1/messages",
"messages": [
{
"role": "assistant",
"content": [
{
"type": "tool_use",
"id": "toolu_01",
"name": "read_external_note",
"input": {},
}
],
},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "toolu_01",
"content": "Note contents",
}
],
},
],
}
with patch(
"litellm.proxy.db.tool_registry_writer.get_tool_policy_registry",
return_value=_registry_mock_with_output(
{"read_external_note": "untrusted", "send_summary": "trusted"},
{"read_external_note": "untrusted", "send_summary": "trusted"},
),
):
with pytest.raises(HTTPException) as exc_info:
await guardrail.apply_guardrail(
inputs=_tool_response_inputs(["send_summary"]),
request_data=request_data,
input_type="response",
)
assert exc_info.value.status_code == 400
assert exc_info.value.detail["blocked_tools"] == ["send_summary"]
assert exc_info.value.detail["untrusted_sources"] == ["read_external_note"]
@pytest.mark.asyncio
async def test_responses_api_function_call_output_blocks_trusted_tool_call(guardrail):
request_data = {
"user_api_key_request_route": "/v1/responses",
"input": [
{
"type": "function_call",
"call_id": "call_1",
"name": "read_external_note",
"arguments": "{}",
},
{
"type": "function_call_output",
"call_id": "call_1",
"output": "Note contents",
},
],
}
with patch(
"litellm.proxy.db.tool_registry_writer.get_tool_policy_registry",
return_value=_registry_mock_with_output(
{"read_external_note": "untrusted", "send_summary": "trusted"},
{"read_external_note": "untrusted", "send_summary": "trusted"},
),
):
with pytest.raises(HTTPException) as exc_info:
await guardrail.apply_guardrail(
inputs=_tool_response_inputs(["send_summary"]),
request_data=request_data,
input_type="response",
)
assert exc_info.value.status_code == 400
assert exc_info.value.detail["blocked_tools"] == ["send_summary"]
assert exc_info.value.detail["untrusted_sources"] == ["read_external_note"]
@pytest.mark.asyncio
async def test_request_stage_blocks_route_tool_when_function_tool_names_non_empty(guardrail):
request_data = {
"user_api_key_request_route": "/v1/messages",
"tools": [
{"name": "ping", "input_schema": {"type": "object"}},
{"type": "bash_20250124", "name": "bash"},
],
}
with patch(
"litellm.proxy.db.tool_registry_writer.get_tool_policy_registry",
return_value=_registry_mock({"ping": "untrusted", "bash": "blocked"}),
):
with pytest.raises(HTTPException) as exc_info:
await guardrail.apply_guardrail(
inputs=_tool_request_inputs(["ping"]),
request_data=request_data,
input_type="request",
)
assert exc_info.value.status_code == 400
assert exc_info.value.detail["blocked_tools"] == ["bash"]