""" Tool Policy Guardrail Reads input_policy / output_policy from LiteLLM_ToolTable and enforces them. Input policy values: "untrusted" - allow through (default for newly discovered tools) "trusted" - only allow if conversation contains no untrusted tool output "blocked" - raise HTTPException, preventing the tool call Output policy values: "untrusted" - output may be tainted (default) "trusted" - output is verified safe Configuration in proxy config YAML: guardrails: - guardrail_name: "tool_policy" litellm_params: guardrail: tool_policy mode: post_call """ from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple from fastapi import HTTPException from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) from litellm.proxy.guardrails.tool_name_extraction import extract_request_tool_names from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj GUARDRAIL_NAME = "tool_policy" def _get_request_object_permission_ids( request_data: dict, ) -> Tuple[Optional[str], Optional[str]]: """Extract object_permission_id and team_object_permission_id from request_data.""" if not request_data: return None, None for key in ("litellm_metadata", "metadata"): meta = request_data.get(key) if not isinstance(meta, dict): continue auth = meta.get("user_api_key_auth") if auth is not None and hasattr(auth, "object_permission_id"): key_op = getattr(auth, "object_permission_id", None) team_op = getattr(auth, "team_object_permission_id", None) if key_op is not None or team_op is not None: return ( str(key_op).strip() if key_op else None, str(team_op).strip() if team_op else None, ) key_op = meta.get("user_api_key_object_permission_id") team_op = meta.get("user_api_key_team_object_permission_id") if key_op is not None or team_op is not None: return ( str(key_op).strip() if key_op else None, str(team_op).strip() if team_op else None, ) return None, None def _get_request_route_from_data(request_data: dict) -> Optional[str]: """Get request route from request_data (metadata or top-level).""" route = request_data.get("user_api_key_request_route") if route: return route meta = request_data.get("metadata") or request_data.get("litellm_metadata") or {} return meta.get("user_api_key_request_route") 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. Used to resolve which tool produced each tool result in the conversation. """ mapping: Dict[str, str] = {} for msg in messages: if msg.get("role") != "assistant": continue tool_calls = msg.get("tool_calls") or [] for tc in tool_calls: if isinstance(tc, dict): tc_id = tc.get("id") fn = (tc.get("function") or {}).get("name") else: tc_id = getattr(tc, "id", None) fn_obj = getattr(tc, "function", None) fn = getattr(fn_obj, "name", None) if fn_obj else None if tc_id and fn: mapping[tc_id] = fn return mapping class ToolPolicyGuardrail(CustomGuardrail): """ Guardrail that enforces per-tool input/output policies from the in-memory ToolPolicyRegistry (synced from DB). """ def __init__(self, **kwargs: Any) -> None: if "supported_event_hooks" not in kwargs: kwargs["supported_event_hooks"] = [ GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call, GuardrailEventHooks.during_call, ] super().__init__(**kwargs) @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: """ Enforce input_policy and output_policy trust chain on request tools / response tool_calls. """ if input_type == "request": tools = inputs.get("tools") or [] tool_names = [ t["function"]["name"] for t in tools if isinstance(t, dict) and isinstance(t.get("function"), dict) and t["function"].get("name") ] if not tool_names: route = _get_request_route_from_data(request_data) if route: tool_names = extract_request_tool_names(route, request_data) else: tool_calls = inputs.get("tool_calls") or [] tool_names = [] for tc in tool_calls: fn = None if isinstance(tc, dict): fn = (tc.get("function") or {}).get("name") elif hasattr(tc, "function"): fn = getattr(tc.function, "name", None) if fn: tool_names.append(fn) if not tool_names: return inputs ( object_permission_id, team_object_permission_id, ) = _get_request_object_permission_ids(request_data) from litellm.proxy.db.tool_registry_writer import get_tool_policy_registry registry = get_tool_policy_registry() if not registry.is_initialized(): return inputs # Stage 1: Check for blocked tools (input_policy=blocked or per-key/team override) policy_map = registry.get_effective_policies( tool_names, object_permission_id=object_permission_id, team_object_permission_id=team_object_permission_id, ) blocked = [name for name in tool_names if policy_map.get(name) == "blocked"] if blocked: verbose_proxy_logger.warning( "ToolPolicyGuardrail: blocking tool(s) %s (input_policy=blocked)", blocked, ) raise HTTPException( status_code=400, detail={ "error": "Violated tool policy", "blocked_tools": blocked, "message": f"Tool(s) {blocked} are blocked by policy.", }, ) # Stage 2: Trust chain enforcement (response path only) # For each tool with input_policy=trusted, check if conversation # contains output from tools with output_policy=untrusted if input_type == "response": trusted_input_tools = [ name for name in tool_names if policy_map.get(name) == "trusted" ] if trusted_input_tools: messages = request_data.get("messages") or [] tc_id_to_name = _resolve_tool_names_from_messages(messages) untrusted_sources: 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 if registry.get_output_policy(source_tool) == "untrusted": if source_tool not in untrusted_sources: untrusted_sources.append(source_tool) if untrusted_sources: verbose_proxy_logger.warning( "ToolPolicyGuardrail: trust chain violation — %s require trusted input " "but conversation has untrusted output from %s", trusted_input_tools, untrusted_sources, ) raise HTTPException( status_code=400, detail={ "error": "Violated tool policy", "blocked_tools": trusted_input_tools, "untrusted_sources": untrusted_sources, "message": ( f"{', '.join(trusted_input_tools)} requires trusted input but " f"conversation contains untrusted output from {', '.join(untrusted_sources)}." ), }, ) return inputs