diff --git a/docs/my-website/docs/proxy/guardrails/tool_permission.md b/docs/my-website/docs/proxy/guardrails/tool_permission.md index 9ed05ed46a8..22ecdd2251e 100644 --- a/docs/my-website/docs/proxy/guardrails/tool_permission.md +++ b/docs/my-website/docs/proxy/guardrails/tool_permission.md @@ -46,6 +46,43 @@ guardrails: - `pre_call` Run **before** LLM call, on **input** - `post_call` Run **after** LLM call, on **input & output** +### `on_disallowed_action` behavior + +| Value | What happens | +| --- | --- | +| `block` | The request is immediately rejected. Pre-call checks raise a `400` HTTP error. Post-call checks raise `GuardrailRaisedException`, so the proxy responds with an error instead of the model output. Use when invoking the forbidden tool must halt the workflow. | +| `rewrite` | LiteLLM silently strips disallowed tools from the payload before it reaches the model (pre-call) or rewrites the model response/tool calls after the fact. The guardrail inserts error text into `message.content`/`tool_result` entries so the client learns the tool was blocked while the rest of the completion continues. Use when you want graceful degradation instead of hard failures. | + +### Custom denial message + +Set `violation_message_template` when you want the guardrail to return a branded error (e.g., “this violates our org policy…”). LiteLLM replaces placeholders from the denied tool: + +- `{tool_name}` – the tool/function name (e.g., `Read`) +- `{rule_id}` – the matching rule ID (or `None` when the default action kicks in) +- `{default_message}` – the original LiteLLM message if you need to append it + +Example: + +```yaml +guardrails: + - guardrail_name: "tool-permission-guardrail" + litellm_params: + guardrail: tool_permission + mode: "post_call" + violation_message_template: "this violates our org policy, we don't support executing {tool_name} commands" + rules: + - id: "allow_bash" + tool_name: "Bash" + decision: "allow" + - id: "deny_read" + tool_name: "Read" + decision: "deny" + default_action: "deny" + on_disallowed_action: "block" +``` + +If a request tries to invoke `Read`, the proxy now returns “this violates our org policy, we don't support executing Read commands” instead of the stock error text. Omit the field to keep the default messaging. + ### 2. Start the Proxy ```shell @@ -57,7 +94,7 @@ litellm --config config.yaml --port 4000 -**Block requset** +**Block request (`on_disallowed_action: block`)** ```bash # Test @@ -96,7 +133,7 @@ curl -X POST "http://localhost:4000/v1/chat/completions" \ -**Rewrite requset** +**Rewrite request (`on_disallowed_action: rewrite`)** ```bash # Test @@ -118,7 +155,7 @@ curl -X POST "http://localhost:4000/v1/chat/completions" \ }' ``` -**Expected response:** +**Expected response (tool removed, completion continues):** ```json { diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index b50d05ed2ec..b52f1b3095e 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -36,6 +36,7 @@ class CustomGuardrail(CustomLogger): default_on: bool = False, mask_request_content: bool = False, mask_response_content: bool = False, + violation_message_template: Optional[str] = None, **kwargs, ): """ @@ -57,12 +58,34 @@ class CustomGuardrail(CustomLogger): self.default_on: bool = default_on self.mask_request_content: bool = mask_request_content self.mask_response_content: bool = mask_response_content + self.violation_message_template: Optional[str] = violation_message_template if supported_event_hooks: ## validate event_hook is in supported_event_hooks self._validate_event_hook(event_hook, supported_event_hooks) super().__init__(**kwargs) + def render_violation_message( + self, default: str, context: Optional[Dict[str, Any]] = None + ) -> str: + """Return a custom violation message if template is configured.""" + + if not self.violation_message_template: + return default + + format_context: Dict[str, Any] = {"default_message": default} + if context: + format_context.update(context) + try: + return self.violation_message_template.format(**format_context) + except Exception as e: + verbose_logger.warning( + "Failed to format violation message template for guardrail %s: %s", + self.guardrail_name, + e, + ) + return default + @staticmethod def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: """ @@ -279,7 +302,7 @@ class CustomGuardrail(CustomLogger): data, self.event_hook ) if result is not None: - return result + return result return True def _event_hook_is_event_type(self, event_type: GuardrailEventHooks) -> bool: diff --git a/litellm/proxy/example_config_yaml/tool_permission_example.yaml b/litellm/proxy/example_config_yaml/tool_permission_example.yaml index e18425ba383..735b4bb7ed2 100644 --- a/litellm/proxy/example_config_yaml/tool_permission_example.yaml +++ b/litellm/proxy/example_config_yaml/tool_permission_example.yaml @@ -10,6 +10,7 @@ guardrails: guardrail: tool_permission mode: "post_call" default_on: true # Apply to all requests by default + violation_message_template: "this violates our org policy, we don't support executing {tool_name} commands" rules: - id: "allow_bash" tool_name: "Bash" @@ -33,4 +34,4 @@ general_settings: # Optional: Add logging configuration litellm_settings: success_callback: ["langfuse"] - failure_callback: ["langfuse"] \ No newline at end of file + failure_callback: ["langfuse"] diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 97f8dd76bd4..19060fa9d6d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -120,13 +120,27 @@ class ToolPermissionGuardrail(CustomGuardrail): for rule in self.rules: if self._matches_pattern(tool_name, rule.tool_name): is_allowed = rule.decision == "allow" - message = f"Tool '{tool_name}' {'allowed' if is_allowed else 'denied'} by rule '{rule.id}'" + default_message = f"Tool '{tool_name}' {'allowed' if is_allowed else 'denied'} by rule '{rule.id}'" + message = self.render_violation_message( + default=default_message, + context={ + "tool_name": tool_name, + "rule_id": rule.id, + }, + ) verbose_proxy_logger.debug(message) return is_allowed, rule.id, message # No rule matched, use default action is_allowed = self.default_action == "allow" - message = f"Tool '{tool_name}' {'allowed' if is_allowed else 'denied'} by default action" + default_message = f"Tool '{tool_name}' {'allowed' if is_allowed else 'denied'} by default action" + message = self.render_violation_message( + default=default_message, + context={ + "tool_name": tool_name, + "rule_id": None, + }, + ) verbose_proxy_logger.debug(message) return is_allowed, None, message @@ -449,7 +463,9 @@ class ToolPermissionGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Tool Permission Guardrail: Checking response") # Extract tool_calls from the response - tool_calls = self._extract_tool_calls_from_response(assembled_model_response) + tool_calls = self._extract_tool_calls_from_response( + assembled_model_response + ) if not tool_calls: verbose_proxy_logger.debug( diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 6a5ba22419b..f2083e9c67e 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -135,6 +135,7 @@ def initialize_tool_permission(litellm_params: LitellmParams, guardrail: Guardra default_action=getattr(litellm_params, "default_action", "deny"), on_disallowed_action=getattr(litellm_params, "on_disallowed_action", "block"), default_on=litellm_params.default_on, + violation_message_template=litellm_params.violation_message_template, ) litellm.logging_callback_manager.add_litellm_callback(_tool_permission_callback) return _tool_permission_callback @@ -172,9 +173,12 @@ def initialize_panw_prisma_airs(litellm_params, guardrail): raise ValueError("PANW Prisma AIRS: profile_name is required") _panw_callback = PanwPrismaAirsHandler( - guardrail_name=guardrail.get("guardrail_name", "panw_prisma_airs"), # Use .get() with default + guardrail_name=guardrail.get( + "guardrail_name", "panw_prisma_airs" + ), # Use .get() with default api_key=litellm_params.api_key, - api_base=litellm_params.api_base or "https://service.api.aisecurity.paloaltonetworks.com/v1/scan/sync/request", + api_base=litellm_params.api_base + or "https://service.api.aisecurity.paloaltonetworks.com/v1/scan/sync/request", profile_name=litellm_params.profile_name, default_on=litellm_params.default_on, ) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index cae9623b44b..f2b9d71cca6 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -16,7 +16,6 @@ from litellm.types.proxy.guardrails.guardrail_hooks.ibm import ( ) - """ Pydantic object defining how to set guardrails on litellm proxy @@ -51,7 +50,7 @@ class SupportedGuardrailIntegrations(Enum): OPENAI_MODERATION = "openai_moderation" NOMA = "noma" TOOL_PERMISSION = "tool_permission" - ZSCALER_AI_GUARD = "zscaler_ai_guard" + ZSCALER_AI_GUARD = "zscaler_ai_guard" JAVELIN = "javelin" ENKRYPTAI = "enkryptai" IBM_GUARDRAILS = "ibm_guardrails" @@ -432,7 +431,7 @@ class ZscalerAIGuardConfigModel(BaseModel): policy_id: Optional[int] = Field( default=None, - description="Policy ID for Zscaler AI Guard. Can also be set via ZSCALER_AI_GUARD_POLICY_ID environment variable" + description="Policy ID for Zscaler AI Guard. Can also be set via ZSCALER_AI_GUARD_POLICY_ID environment variable", ) send_user_api_key_alias: Optional[bool] = Field( default=False, description="Whether to send user_API_key_alias in headers" @@ -444,6 +443,7 @@ class ZscalerAIGuardConfigModel(BaseModel): default=False, description="Whether to send user_API_key_team_id in headers" ) + class JavelinGuardrailConfigModel(BaseModel): """Configuration parameters for the Javelin guardrail""" @@ -479,7 +479,8 @@ class BlockedWord(BaseModel): description="Action to take when keyword is detected (BLOCK or MASK)" ) description: Optional[str] = Field( - default=None, description="Optional description explaining why this keyword is sensitive" + default=None, + description="Optional description explaining why this keyword is sensitive", ) @@ -491,15 +492,15 @@ class ContentFilterPattern(BaseModel): ) pattern_name: Optional[str] = Field( default=None, - description="Name of prebuilt pattern (e.g., 'us_ssn', 'credit_card'). Required if pattern_type is 'prebuilt'" + description="Name of prebuilt pattern (e.g., 'us_ssn', 'credit_card'). Required if pattern_type is 'prebuilt'", ) pattern: Optional[str] = Field( default=None, - description="Custom regex pattern. Required if pattern_type is 'regex'" + description="Custom regex pattern. Required if pattern_type is 'regex'", ) name: Optional[str] = Field( default=None, - description="Name for this pattern (used in logging and error messages)" + description="Name for this pattern (used in logging and error messages)", ) action: ContentFilterAction = Field( description="Action to take when pattern matches (BLOCK or MASK)" @@ -511,15 +512,13 @@ class ContentFilterConfigModel(BaseModel): patterns: Optional[List[ContentFilterPattern]] = Field( default=None, - description="List of patterns (prebuilt or custom regex) to detect" + description="List of patterns (prebuilt or custom regex) to detect", ) blocked_words: Optional[List[BlockedWord]] = Field( - default=None, - description="List of blocked words with individual actions" + default=None, description="List of blocked words with individual actions" ) blocked_words_file: Optional[str] = Field( - default=None, - description="Path to YAML file containing blocked_words list" + default=None, description="Path to YAML file containing blocked_words list" ) @@ -575,6 +574,11 @@ class BaseLitellmParams(BaseModel): # works for new and patch update guardrails description="Optional field if guardrail requires a 'model' parameter", ) + violation_message_template: Optional[str] = Field( + default=None, + description="Custom message when a guardrail blocks an action. Supports placeholders like {tool_name}, {rule_id}, and {default_message}.", + ) + # Model Armor params template_id: Optional[str] = Field( default=None, description="The ID of your Model Armor template" @@ -613,7 +617,7 @@ class LitellmParams( GraySwanGuardrailConfigModel, NomaGuardrailConfigModel, ToolPermissionGuardrailConfigModel, - ZscalerAIGuardConfigModel, + ZscalerAIGuardConfigModel, JavelinGuardrailConfigModel, ContentFilterConfigModel, BaseLitellmParams, @@ -671,10 +675,12 @@ class GuardrailEventHooks(str, Enum): class DynamicGuardrailParams(TypedDict): extra_body: Dict[str, Any] + class GUARDRAIL_DEFINITION_LOCATION(str, Enum): DB = "db" CONFIG = "config" + class GuardrailInfoResponse(BaseModel): guardrail_id: Optional[str] = None guardrail_name: str @@ -682,7 +688,9 @@ class GuardrailInfoResponse(BaseModel): guardrail_info: Optional[Dict] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None - guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = GUARDRAIL_DEFINITION_LOCATION.CONFIG + guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = ( + GUARDRAIL_DEFINITION_LOCATION.CONFIG + ) def __init__(self, **kwargs): super().__init__(**kwargs) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index a9d87398217..8c88b22f60e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -129,6 +129,24 @@ class TestToolPermissionGuardrail: assert rule_id is None assert "default" in (msg or "") + def test_check_tool_permission_custom_template(self): + guardrail = ToolPermissionGuardrail( + guardrail_name="custom-template", + rules=self.test_rules, + default_action="deny", + violation_message_template="custom {tool_name} {rule_id} :: {default_message}", + ) + + _, rule_id, message = guardrail._check_tool_permission("Read") + assert rule_id == "deny_read" + assert message.startswith("custom Read deny_read") + assert "Tool 'Read' denied" in message + + _, rule_id, message = guardrail._check_tool_permission("UnknownTool") + assert rule_id is None + assert message.startswith("custom UnknownTool None") + assert "Tool 'UnknownTool' denied by default action" in message + def test_extract_tool_calls_openai_format(self): tool_call = { "id": "call_123", @@ -224,6 +242,39 @@ class TestToolPermissionGuardrail: ) assert excinfo.value.status_code == 400 + @pytest.mark.asyncio + async def test_async_pre_call_hook_uses_custom_template(self): + guardrail = ToolPermissionGuardrail( + guardrail_name="custom-template", + rules=self.test_rules, + default_action="deny", + on_disallowed_action="block", + violation_message_template="blocked {tool_name} by policy", + ) + + data = { + "tools": [ + {"type": "function", "function": {"name": "Read"}}, + ] + } + user_api_key_dict = UserAPIKeyAuth() + cache = DualCache(default_in_memory_ttl=1) + + with patch.object(guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(HTTPException) as excinfo: + await guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="completion", + ) + + assert excinfo.value.status_code == 400 + assert ( + excinfo.value.detail.get("detection_message") + == "blocked Read by policy" + ) + @pytest.mark.asyncio async def test_async_pre_call_hook_rewrite_mode(self): guardrail = ToolPermissionGuardrail(