diff --git a/docs/my-website/docs/proxy/guardrails/tool_permission.md b/docs/my-website/docs/proxy/guardrails/tool_permission.md index 2e0b72a8a8a..19b674c9e55 100644 --- a/docs/my-website/docs/proxy/guardrails/tool_permission.md +++ b/docs/my-website/docs/proxy/guardrails/tool_permission.md @@ -2,9 +2,9 @@ import Image from '@theme/IdealImage'; import Tabs from '@theme/Tabs'; import TabItem from '@theme/TabItem'; -# Tool Permission Guardrail +# LiteLLM Tool Permission Guardrail -LiteLLM provides a Tool Permission Guardrail that lets you control which **tool calls** a model is allowed to invoke, using configurable allow/deny rules. This offers fine-grained, provider-agnostic control over tool execution (e.g., OpenAI Chat Completions `tool_calls`, Anthropic Messages `tool_use`, MCP tools). +LiteLLM provides the LiteLLM Tool Permission Guardrail that lets you control which **tool calls** a model is allowed to invoke, using configurable allow/deny rules. This offers fine-grained, provider-agnostic control over tool execution (e.g., OpenAI Chat Completions `tool_calls`, Anthropic Messages `tool_use`, MCP tools). ## Quick Start ### 1. Define Guardrails on your LiteLLM config.yaml diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index e64fbe9084e..a1cfead9bb2 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -31,6 +31,7 @@ from litellm.types.guardrails import ( PiiEntityType, PresidioPresidioConfigModelUserInterface, SupportedGuardrailIntegrations, + ToolPermissionGuardrailConfigModel, ) #### GUARDRAILS ENDPOINTS #### @@ -635,7 +636,9 @@ async def get_guardrail_info(guardrail_id: str): raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = GUARDRAIL_DEFINITION_LOCATION.DB + guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = ( + GUARDRAIL_DEFINITION_LOCATION.DB + ) result = await GUARDRAIL_REGISTRY.get_guardrail_by_id_from_db( guardrail_id=guardrail_id, prisma_client=prisma_client ) @@ -702,10 +705,12 @@ async def get_guardrail_ui_settings(): # Convert the PII_ENTITY_CATEGORIES_MAP to the format expected by the UI category_maps = [] for category, entities in PII_ENTITY_CATEGORIES_MAP.items(): - category_maps.append({ - "category": category.value, - "entities": [entity.value for entity in entities] - }) + category_maps.append( + { + "category": category.value, + "entities": [entity.value for entity in entities], + } + ) return GuardrailUIAddGuardrailSettings( supported_entities=[entity.value for entity in PiiEntityType], @@ -728,20 +733,20 @@ async def get_guardrail_ui_settings(): async def validate_blocked_words_file(request: Dict[str, str]): """ Validate a blocked_words YAML file content. - + Args: request: Dictionary with 'file_content' key containing the YAML string - + Returns: Dictionary with 'valid' boolean and either 'message'/'errors' depending on result - + Example Request: ```json { "file_content": "blocked_words:\\n - keyword: \\"test\\"\\n action: \\"BLOCK\\"" } ``` - + Example Success Response: ```json { @@ -749,7 +754,7 @@ async def validate_blocked_words_file(request: Dict[str, str]): "message": "Valid YAML file with 2 blocked words" } ``` - + Example Error Response: ```json { @@ -759,56 +764,54 @@ async def validate_blocked_words_file(request: Dict[str, str]): ``` """ import yaml - + try: file_content = request.get("file_content", "") if not file_content: - return { - "valid": False, - "error": "No file content provided" - } - + return {"valid": False, "error": "No file content provided"} + data = yaml.safe_load(file_content) - + if not isinstance(data, dict) or "blocked_words" not in data: return { "valid": False, - "error": "Invalid format: file must contain 'blocked_words' key with a list" + "error": "Invalid format: file must contain 'blocked_words' key with a list", } - + blocked_words_list = data["blocked_words"] if not isinstance(blocked_words_list, list): - return { - "valid": False, - "error": "'blocked_words' must be a list" - } - + return {"valid": False, "error": "'blocked_words' must be a list"} + # Validate each entry errors = [] for idx, word_data in enumerate(blocked_words_list): if not isinstance(word_data, dict): errors.append(f"Entry {idx}: must be an object") continue - + if "keyword" not in word_data: errors.append(f"Entry {idx}: missing 'keyword' field") elif not isinstance(word_data["keyword"], str): errors.append(f"Entry {idx}: 'keyword' must be a string") - + if "action" not in word_data: errors.append(f"Entry {idx}: missing 'action' field") elif word_data["action"] not in ["BLOCK", "MASK"]: - errors.append(f"Entry {idx}: action must be 'BLOCK' or 'MASK', got '{word_data['action']}'") - - if "description" in word_data and not isinstance(word_data["description"], str): + errors.append( + f"Entry {idx}: action must be 'BLOCK' or 'MASK', got '{word_data['action']}'" + ) + + if "description" in word_data and not isinstance( + word_data["description"], str + ): errors.append(f"Entry {idx}: 'description' must be a string") - + if errors: return {"valid": False, "errors": errors} - + return { "valid": True, - "message": f"Valid YAML file with {len(blocked_words_list)} blocked word(s)" + "message": f"Valid YAML file with {len(blocked_words_list)} blocked word(s)", } except yaml.YAMLError as e: return {"valid": False, "error": f"Invalid YAML syntax: {str(e)}"} @@ -931,30 +934,32 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool """Check if optional_params field should be skipped (not meaningfully overridden).""" if field_name != "optional_params": return False - + if field_annotation is None: return True - + # Check if the annotation is still a generic TypeVar (not specialized) if isinstance(field_annotation, TypeVar) or ( hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is TypeVar ): return True - + # Also skip if it's a generic type that wasn't specialized if hasattr(field_annotation, "__name__") and field_annotation.__name__ in ( "T", "TypeVar", ): return True - + # Handle Optional[T] where T is still a TypeVar if hasattr(field_annotation, "__args__"): - non_none_args = [arg for arg in field_annotation.__args__ if arg is not type(None)] + non_none_args = [ + arg for arg in field_annotation.__args__ if arg is not type(None) + ] if non_none_args and isinstance(non_none_args[0], TypeVar): return True - + return False @@ -1041,9 +1046,11 @@ def _extract_fields_recursive( for field_name, field in model.model_fields.items(): field_annotation = field.annotation - + # Skip optional_params if it's not meaningfully overridden - if _should_skip_optional_params(field_name=field_name, field_annotation=field_annotation): + if _should_skip_optional_params( + field_name=field_name, field_annotation=field_annotation + ): continue # Handle Optional types and get the actual type @@ -1153,12 +1160,18 @@ async def get_provider_specific_params(): bedrock_fields = _get_fields_from_model(BedrockGuardrailConfigModel) presidio_fields = _get_fields_from_model(PresidioPresidioConfigModelUserInterface) lakera_v2_fields = _get_fields_from_model(LakeraV2GuardrailConfigModel) + tool_permission_fields = _get_fields_from_model(ToolPermissionGuardrailConfigModel) + + tool_permission_fields[ + "ui_friendly_name" + ] = ToolPermissionGuardrailConfigModel.ui_friendly_name() # Return the provider-specific parameters provider_params = { SupportedGuardrailIntegrations.BEDROCK.value: bedrock_fields, SupportedGuardrailIntegrations.PRESIDIO.value: presidio_fields, SupportedGuardrailIntegrations.LAKERA_V2.value: lakera_v2_fields, + SupportedGuardrailIntegrations.TOOL_PERMISSION.value: tool_permission_fields, } ### get the config model for the guardrail - go through the registry and get the config model for the guardrail @@ -1175,6 +1188,7 @@ async def get_provider_specific_params(): return provider_params + @router.post("/guardrails/apply_guardrail", response_model=ApplyGuardrailResponse) @router.post("/apply_guardrail", response_model=ApplyGuardrailResponse) async def apply_guardrail( @@ -1183,11 +1197,11 @@ async def apply_guardrail( ): """ Apply a guardrail to text input and return the processed result. - + This endpoint allows testing guardrails by applying them to custom text inputs. """ from litellm.proxy.utils import handle_exception_on_proxy - + try: active_guardrail: Optional[ CustomGuardrail @@ -1207,4 +1221,3 @@ async def apply_guardrail( return ApplyGuardrailResponse(response_text=response_text) except Exception as e: raise handle_exception_on_proxy(e) - diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 95d8f894dc2..eef8043b237 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -62,8 +62,11 @@ class ToolPermissionGuardrail(CustomGuardrail): self.rules: List[ToolPermissionRule] = [] self._compiled_rule_patterns: Dict[str, Dict[str, re.Pattern]] = {} if rules: - for rule_dict in rules: - rule = ToolPermissionRule(**rule_dict) + for rule_item in rules: + if isinstance(rule_item, ToolPermissionRule): + rule = rule_item + else: + rule = ToolPermissionRule(**rule_item) self.rules.append(rule) if rule.allowed_param_patterns: @@ -88,6 +91,14 @@ class ToolPermissionGuardrail(CustomGuardrail): self.default_action, ) + @staticmethod + def get_config_model(): + from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( + ToolPermissionGuardrailConfigModel, + ) + + return ToolPermissionGuardrailConfigModel + def _matches_pattern(self, tool_name: str, pattern: str) -> bool: """ Check if a tool name matches a pattern diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index f2083e9c67e..9bb965ef14e 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -1,4 +1,6 @@ # litellm/proxy/guardrails/guardrail_initializers.py +from typing import Any, Dict, List, Optional + import litellm from litellm.proxy._types import CommonProxyErrors from litellm.types.guardrails import * @@ -128,10 +130,19 @@ def initialize_tool_permission(litellm_params: LitellmParams, guardrail: Guardra ToolPermissionGuardrail, ) + rules: Optional[List[Dict[str, Any]]] = None + if litellm_params.rules: + rules = [] + for rule in litellm_params.rules: + if hasattr(rule, "model_dump"): + rules.append(rule.model_dump()) + else: + rules.append(dict(rule)) + _tool_permission_callback = ToolPermissionGuardrail( guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, - rules=litellm_params.rules, + rules=rules, default_action=getattr(litellm_params, "default_action", "deny"), on_disallowed_action=getattr(litellm_params, "on_disallowed_action", "block"), default_on=litellm_params.default_on, diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 5f6295e151f..24a235def59 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -14,6 +14,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import ( from litellm.types.proxy.guardrails.guardrail_hooks.ibm import ( IBMGuardrailsBaseConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( + ToolPermissionGuardrailConfigModel, +) """ @@ -415,18 +418,6 @@ class NomaGuardrailConfigModel(BaseModel): ) -class ToolPermissionGuardrailConfigModel(BaseModel): - """Configuration parameters for the Tool Permission guardrail""" - - rules: Optional[List[Dict]] = Field( - default=None, description="List of permission rules for tool usage" - ) - default_action: Optional[str] = Field( - default="Deny", - description="Default action when no rule matches (Allow or Deny)", - ) - - class ZscalerAIGuardConfigModel(BaseModel): """Configuration parameters for the Zscaler AI Guard guardrail""" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py index b2248c51930..e78cfad8bdb 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -1,8 +1,10 @@ # Tool Permission Guardrail Type Definitions -from typing import Dict, Literal, Optional +from typing import Dict, List, Literal, Optional from pydantic import BaseModel, Field +from .base import GuardrailConfigModel + class ToolPermissionRule(BaseModel): """ @@ -43,3 +45,23 @@ class PermissionError(BaseModel): tool_name: str = Field(description="Name of the denied tool") rule_id: Optional[str] = Field(description="ID of the rule that caused denial") message: str = Field(description="Error message") + + +class ToolPermissionGuardrailConfigModel(GuardrailConfigModel): + """Configuration parameters exposed to the UI for the Tool Permission guardrail.""" + + rules: Optional[List[ToolPermissionRule]] = Field( + default=None, + description="Ordered allow/deny rules. Patterns support * wildcards and optional regex constraints on tool arguments.", + ) + default_action: Literal["allow", "deny"] = Field( + default="deny", description="Fallback decision when no rule matches" + ) + on_disallowed_action: Literal["block", "rewrite"] = Field( + default="block", + description="Choose whether disallowed tools block the request or get rewritten out of the payload", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "LiteLLM Tool Permission Guardrail" diff --git a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx index 55fe800a431..20ca36f6d16 100644 --- a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx @@ -1,4 +1,4 @@ -import React, { useState, useEffect } from "react"; +import React, { useState, useEffect, useMemo } from "react"; import { Form, Typography, Select, Modal, Tag, Steps } from "antd"; import { Button, TextInput } from "@tremor/react"; import { @@ -16,6 +16,9 @@ import GuardrailProviderFields from "./guardrail_provider_fields"; import GuardrailOptionalParams from "./guardrail_optional_params"; import NotificationsManager from "../molecules/notifications_manager"; import ContentFilterConfiguration from "./content_filter/ContentFilterConfiguration"; +import ToolPermissionRulesEditor, { + ToolPermissionConfig, +} from "./tool_permission/ToolPermissionRulesEditor"; const { Title, Text, Link } = Typography; const { Option } = Select; @@ -100,6 +103,20 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a // Content Filter state const [selectedPatterns, setSelectedPatterns] = useState([]); const [blockedWords, setBlockedWords] = useState([]); + const [toolPermissionConfig, setToolPermissionConfig] = useState({ + rules: [], + default_action: "deny", + on_disallowed_action: "block", + violation_message_template: "", + }); + + const isToolPermissionProvider = useMemo(() => { + if (!selectedProvider) { + return false; + } + const providerValue = guardrail_provider_map[selectedProvider]; + return (providerValue || "").toLowerCase() === "tool_permission"; + }, [selectedProvider]); // Fetch guardrail UI settings + provider params on mount / accessToken change useEffect(() => { @@ -145,6 +162,13 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a setSelectedCategories([]); setGlobalSeverityThreshold(2); setCategorySpecificThresholds({}); + + setToolPermissionConfig({ + rules: [], + default_action: "deny", + on_disallowed_action: "block", + violation_message_template: "", + }); }; const handleEntitySelect = (entity: string) => { @@ -225,6 +249,14 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a setSelectedCategories([]); setGlobalSeverityThreshold(2); setCategorySpecificThresholds({}); + setSelectedPatterns([]); + setBlockedWords([]); + setToolPermissionConfig({ + rules: [], + default_action: "deny", + on_disallowed_action: "block", + violation_message_template: "", + }); setCurrentStep(0); }; @@ -315,6 +347,20 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a } } + if (guardrailProvider === "tool_permission") { + if (toolPermissionConfig.rules.length === 0) { + NotificationsManager.fromBackend("Add at least one tool permission rule"); + setLoading(false); + return; + } + guardrailData.litellm_params.rules = toolPermissionConfig.rules; + guardrailData.litellm_params.default_action = toolPermissionConfig.default_action; + guardrailData.litellm_params.on_disallowed_action = toolPermissionConfig.on_disallowed_action; + if (toolPermissionConfig.violation_message_template) { + guardrailData.litellm_params.violation_message_template = toolPermissionConfig.violation_message_template; + } + } + /****************************** * Add provider-specific params * ---------------------------------- @@ -535,11 +581,13 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a {/* Use the GuardrailProviderFields component to render provider-specific fields */} - + {!isToolPermissionProvider && ( + + )} ); }; @@ -593,7 +641,20 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a }; const renderOptionalParams = () => { - if (!selectedProvider || !providerParams) return null; + if (!selectedProvider) return null; + + if (isToolPermissionProvider) { + return ( + + ); + } + + if (!providerParams) { + return null; + } console.log("guardrail_provider_map: ", guardrail_provider_map); console.log("selectedProvider: ", selectedProvider); diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx index fa34cf55636..0d1a205f66f 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx @@ -26,6 +26,9 @@ import PiiConfiguration from "./pii_configuration"; import GuardrailProviderFields from "./guardrail_provider_fields"; import GuardrailOptionalParams from "./guardrail_optional_params"; import ContentFilterManager, { formatContentFilterDataForAPI } from "./content_filter/ContentFilterManager"; +import ToolPermissionRulesEditor, { + ToolPermissionConfig, +} from "./tool_permission/ToolPermissionRulesEditor"; import { ArrowLeftIcon } from "@heroicons/react/outline"; import { copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils"; import { CheckIcon, CopyIcon } from "lucide-react"; @@ -83,6 +86,14 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, } | null>(null); const [copiedStates, setCopiedStates] = useState>({}); const [hasUnsavedContentFilterChanges, setHasUnsavedContentFilterChanges] = useState(false); + const emptyToolPermissionConfig: ToolPermissionConfig = { + rules: [], + default_action: "deny", + on_disallowed_action: "block", + violation_message_template: "", + }; + const [toolPermissionConfig, setToolPermissionConfig] = useState(emptyToolPermissionConfig); + const [toolPermissionDirty, setToolPermissionDirty] = useState(false); // Content Filter data ref (managed by ContentFilterManager) const contentFilterDataRef = React.useRef<{ patterns: any[]; blockedWords: any[] }>({ @@ -180,6 +191,29 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, } }, [guardrailData, guardrailProviderSpecificParams, form]); + const resetToolPermissionEditor = useCallback(() => { + if (guardrailData?.litellm_params?.guardrail === "tool_permission") { + setToolPermissionConfig({ + rules: (guardrailData.litellm_params?.rules as ToolPermissionConfig["rules"]) || [], + default_action: ((guardrailData.litellm_params?.default_action || "deny") as ToolPermissionConfig["default_action"]).toLowerCase() as ToolPermissionConfig["default_action"], + on_disallowed_action: ((guardrailData.litellm_params?.on_disallowed_action || "block") as ToolPermissionConfig["on_disallowed_action"]).toLowerCase() as ToolPermissionConfig["on_disallowed_action"], + violation_message_template: guardrailData.litellm_params?.violation_message_template || "", + }); + } else { + setToolPermissionConfig(emptyToolPermissionConfig); + } + setToolPermissionDirty(false); + }, [guardrailData]); + + useEffect(() => { + resetToolPermissionEditor(); + }, [resetToolPermissionEditor]); + + const handleToolPermissionConfigChange = (config: ToolPermissionConfig) => { + setToolPermissionConfig(config); + setToolPermissionDirty(true); + }; + const handlePiiEntitySelect = (entity: string) => { setSelectedPiiEntities((prev) => { if (prev.includes(entity)) { @@ -255,6 +289,31 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, } } + if (guardrailData.litellm_params?.guardrail === "tool_permission") { + const originalRules = guardrailData.litellm_params?.rules || []; + const currentRules = toolPermissionConfig.rules || []; + const rulesChanged = JSON.stringify(originalRules) !== JSON.stringify(currentRules); + + const originalDefault = (guardrailData.litellm_params?.default_action || "deny").toLowerCase(); + const currentDefault = (toolPermissionConfig.default_action || "deny").toLowerCase(); + const defaultChanged = originalDefault !== currentDefault; + + const originalOnDisallowed = (guardrailData.litellm_params?.on_disallowed_action || "block").toLowerCase(); + const currentOnDisallowed = (toolPermissionConfig.on_disallowed_action || "block").toLowerCase(); + const onDisallowedChanged = originalOnDisallowed !== currentOnDisallowed; + + const originalMessage = guardrailData.litellm_params?.violation_message_template || ""; + const currentMessage = toolPermissionConfig.violation_message_template || ""; + const messageChanged = originalMessage !== currentMessage; + + if (toolPermissionDirty || rulesChanged || defaultChanged || onDisallowedChanged || messageChanged) { + updateData.litellm_params.rules = currentRules; + updateData.litellm_params.default_action = currentDefault; + updateData.litellm_params.on_disallowed_action = currentOnDisallowed; + updateData.litellm_params.violation_message_template = currentMessage || null; + } + } + /****************************** * Add provider-specific params (reusing logic from add_guardrail_form.tsx) * ---------------------------------- @@ -273,7 +332,8 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, console.log("currentProvider: ", currentProvider); // Use pre-fetched provider params to copy recognised params - if (guardrailProviderSpecificParams && currentProvider) { + const isToolPermissionGuardrail = guardrailData.litellm_params?.guardrail === "tool_permission"; + if (guardrailProviderSpecificParams && currentProvider && !isToolPermissionGuardrail) { const providerKey = guardrail_provider_map[currentProvider]?.toLowerCase(); const providerSpecificParams = guardrailProviderSpecificParams[providerKey] || {}; @@ -488,6 +548,12 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, )} + {guardrailData.litellm_params?.guardrail === "tool_permission" && ( + + + + )} + {/* Content Filter Configuration Display */} = ({ guardrailId, onClose, Provider Settings - {/* Provider-specific fields */} - guardrail_provider_map[key] === guardrailData.litellm_params?.guardrail, - ) || null - } - accessToken={accessToken} - providerParams={guardrailProviderSpecificParams} - value={guardrailData.litellm_params} - /> + {guardrailData.litellm_params?.guardrail === "tool_permission" ? ( + + ) : ( + <> + {/* Provider-specific fields */} + guardrail_provider_map[key] === guardrailData.litellm_params?.guardrail, + ) || null + } + accessToken={accessToken} + providerParams={guardrailProviderSpecificParams} + value={guardrailData.litellm_params} + /> - {/* Optional parameters */} - {guardrailProviderSpecificParams && - (() => { - const currentProvider = Object.keys(guardrail_provider_map).find( - (key) => guardrail_provider_map[key] === guardrailData.litellm_params?.guardrail, - ); - if (!currentProvider) return null; + {/* Optional parameters */} + {guardrailProviderSpecificParams && + (() => { + const currentProvider = Object.keys(guardrail_provider_map).find( + (key) => guardrail_provider_map[key] === guardrailData.litellm_params?.guardrail, + ); + if (!currentProvider) return null; - const providerKey = guardrail_provider_map[currentProvider]?.toLowerCase(); - const providerFields = guardrailProviderSpecificParams[providerKey]; + const providerKey = guardrail_provider_map[currentProvider]?.toLowerCase(); + const providerFields = guardrailProviderSpecificParams[providerKey]; - if (!providerFields || !providerFields.optional_params) return null; + if (!providerFields || !providerFields.optional_params) return null; - return ( - - ); - })()} + return ( + + ); + })()} + + )} Advanced Settings @@ -619,6 +694,7 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, onClick={() => { setIsEditing(false); setHasUnsavedContentFilterChanges(false); + resetToolPermissionEditor(); }} > Cancel @@ -672,6 +748,10 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, Last Updated
{formatDate(guardrailData.updated_at)}
+ + {guardrailData.litellm_params?.guardrail === "tool_permission" && ( + + )} )} diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx index fa665d4911c..c6314c95bef 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx @@ -46,6 +46,7 @@ export const guardrail_provider_map: Record = { Bedrock: "bedrock", Lakera: "lakera_v2", LitellmContentFilter: "litellm_content_filter", + ToolPermission: "tool_permission", }; // Function to populate provider map from API response - updates the original map diff --git a/ui/litellm-dashboard/src/components/guardrails/tool_permission/ToolPermissionRulesEditor.test.tsx b/ui/litellm-dashboard/src/components/guardrails/tool_permission/ToolPermissionRulesEditor.test.tsx new file mode 100644 index 00000000000..59736694888 --- /dev/null +++ b/ui/litellm-dashboard/src/components/guardrails/tool_permission/ToolPermissionRulesEditor.test.tsx @@ -0,0 +1,65 @@ +import React from "react"; +import { describe, it, expect, vi } from "vitest"; +import { render, screen, waitFor, fireEvent } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import ToolPermissionRulesEditor, { + ToolPermissionConfig, +} from "./ToolPermissionRulesEditor"; + +describe("ToolPermissionRulesEditor", () => { + it("renders empty state and lets users add a new rule", async () => { + const onChange = vi.fn(); + render(); + + expect(screen.getByText(/No tool rules added yet/i)).toBeInTheDocument(); + + await userEvent.click(screen.getByRole("button", { name: /add rule/i })); + + expect(onChange).toHaveBeenCalled(); + const payload = onChange.mock.calls[0][0] as ToolPermissionConfig; + expect(payload.rules).toHaveLength(1); + expect(payload.rules[0].decision).toBe("allow"); + }); + + it("captures violation message and argument constraints", async () => { + let latestConfig: ToolPermissionConfig | null = null; + const initialConfig: ToolPermissionConfig = { + rules: [ + { + id: "allow_bash", + tool_name: "Bash", + decision: "allow", + }, + ], + default_action: "deny", + on_disallowed_action: "block", + violation_message_template: "", + }; + + const Wrapper = () => { + const [state, setState] = React.useState(initialConfig); + const handleChange = (next: ToolPermissionConfig) => { + latestConfig = next; + setState(next); + }; + return ; + }; + + render(); + + await userEvent.click(screen.getByRole("button", { name: /restrict tool arguments/i })); + const initialInput = await screen.findByPlaceholderText(/messages\[0\].content/i); + await userEvent.clear(initialInput); + fireEvent.change(initialInput, { target: { value: "input.location" } }); + + const violationArea = await screen.findByPlaceholderText(/violates our org policy/i); + await userEvent.clear(violationArea); + fireEvent.change(violationArea, { target: { value: "Do not run bash" } }); + + await waitFor(() => { + expect(latestConfig).not.toBeNull(); + expect(latestConfig?.rules[0].allowed_param_patterns).toEqual({ "input.location": "" }); + expect(latestConfig?.violation_message_template).toBe("Do not run bash"); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/guardrails/tool_permission/ToolPermissionRulesEditor.tsx b/ui/litellm-dashboard/src/components/guardrails/tool_permission/ToolPermissionRulesEditor.tsx new file mode 100644 index 00000000000..790876ed3f0 --- /dev/null +++ b/ui/litellm-dashboard/src/components/guardrails/tool_permission/ToolPermissionRulesEditor.tsx @@ -0,0 +1,322 @@ +import React from "react"; +import { Card, Text } from "@tremor/react"; +import { Button, Divider, Empty, Input, Select, Space, Tooltip } from "antd"; +import { InfoCircleOutlined, PlusOutlined, DeleteOutlined } from "@ant-design/icons"; + +export type ToolPermissionDecision = "allow" | "deny"; +export type ToolPermissionDefaultAction = "allow" | "deny"; +export type ToolPermissionOnDisallowedAction = "block" | "rewrite"; + +export interface ToolPermissionRuleConfig { + id: string; + tool_name: string; + decision: ToolPermissionDecision; + allowed_param_patterns?: Record; +} + +export interface ToolPermissionConfig { + rules: ToolPermissionRuleConfig[]; + default_action: ToolPermissionDefaultAction; + on_disallowed_action: ToolPermissionOnDisallowedAction; + violation_message_template?: string; +} + +interface ToolPermissionRulesEditorProps { + value?: ToolPermissionConfig; + onChange?: (config: ToolPermissionConfig) => void; + disabled?: boolean; +} + +const DEFAULT_CONFIG: ToolPermissionConfig = { + rules: [], + default_action: "deny", + on_disallowed_action: "block", + violation_message_template: "", +}; + +const ensureConfig = (config?: ToolPermissionConfig): ToolPermissionConfig => ({ + ...DEFAULT_CONFIG, + ...(config || {}), + rules: config?.rules ? [...config.rules] : [], +}); + +const ToolPermissionRulesEditor: React.FC = ({ + value, + onChange, + disabled = false, +}) => { + const config = ensureConfig(value); + + const updateConfig = (partial: Partial) => { + const nextConfig: ToolPermissionConfig = { + ...config, + ...partial, + }; + onChange?.(nextConfig); + }; + + const updateRule = (ruleIndex: number, updates: Partial) => { + const nextRules = config.rules.map((rule, index) => + index === ruleIndex ? { ...rule, ...updates } : rule, + ); + updateConfig({ rules: nextRules }); + }; + + const addRule = () => { + const nextRules = [ + ...config.rules, + { + id: `rule_${Math.random().toString(36).slice(2, 8)}`, + tool_name: "", + decision: "allow" as ToolPermissionDecision, + allowed_param_patterns: undefined, + }, + ]; + updateConfig({ rules: nextRules }); + }; + + const removeRule = (ruleIndex: number) => { + const nextRules = config.rules.filter((_, index) => index !== ruleIndex); + updateConfig({ rules: nextRules }); + }; + + const updateAllowedParamEntries = ( + ruleIndex: number, + mutate: (entries: [string, string][]) => void, + ) => { + const targetRule = config.rules[ruleIndex]; + if (!targetRule) { + return; + } + const entries = Object.entries(targetRule.allowed_param_patterns || {}); + mutate(entries); + const updatedObject: Record = {}; + entries.forEach(([key, value]) => { + updatedObject[key] = value; + }); + updateRule(ruleIndex, { + allowed_param_patterns: + Object.keys(updatedObject).length > 0 ? updatedObject : undefined, + }); + }; + + const updateAllowedParamPath = ( + ruleIndex: number, + entryIndex: number, + nextPath: string, + ) => { + updateAllowedParamEntries(ruleIndex, (entries) => { + if (!entries[entryIndex]) { + return; + } + const [, value] = entries[entryIndex]; + entries[entryIndex] = [nextPath, value]; + }); + }; + + const updateAllowedParamPattern = ( + ruleIndex: number, + entryIndex: number, + pattern: string, + ) => { + updateAllowedParamEntries(ruleIndex, (entries) => { + if (!entries[entryIndex]) { + return; + } + const [path] = entries[entryIndex]; + entries[entryIndex] = [path, pattern]; + }); + }; + + const renderAllowedParamPatterns = (rule: ToolPermissionRuleConfig, index: number) => { + const entries = Object.entries(rule.allowed_param_patterns || {}); + if (entries.length === 0) { + return ( + + ); + } + + return ( +
+ Argument constraints (dot or array paths) + {entries.map(([path, pattern], patternIndex) => ( + + updateAllowedParamPath(index, patternIndex, e.target.value)} + /> + updateAllowedParamPattern(index, patternIndex, e.target.value)} + /> + +
+ ); + }; + + return ( + +
+
+ LiteLLM Tool Permission Guardrail + + Use wildcards (e.g., mcp__github_*) to scope which tools can run and optionally constrain + payload fields. + +
+ {!disabled && ( + + )} +
+ + + + {config.rules.length === 0 ? ( + + ) : ( +
+ {config.rules.map((rule, index) => ( + +
+ Rule {index + 1} + +
+
+
+ Rule ID + updateRule(index, { id: e.target.value })} + /> +
+
+ Tool Name / Pattern + updateRule(index, { tool_name: e.target.value })} + /> +
+
+ +
+ Decision + +
+ +
{renderAllowedParamPatterns(rule, index)}
+
+ ))} +
+ )} + + + +
+
+ Default action + +
+
+ + On disallowed action + + + + + +
+
+ +
+ Violation message (optional) + updateConfig({ violation_message_template: e.target.value })} + /> +
+
+ ); +}; + +export default ToolPermissionRulesEditor;