mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
feat: add claude_code guardrails + provider-agnostic block_hosted_tools
Adds three new guardrail hooks for Claude Code proxy deployments:
**block_hosted_tools** (provider-agnostic, new top-level hook)
- Blocks platform-executed hosted tools from Anthropic, OpenAI, and Gemini
- Per-provider YAML configs (anthropic.yaml, openai.yaml, gemini.yaml) define
tool type prefixes (Anthropic versioned types), exact type strings (OpenAI),
and native top-level keys (Gemini googleSearch, codeExecution, etc.)
- Registered as guardrail type: "block_hosted_tools"
**claude_code_prompt_cache**
- Auto-injects cache_control: {type: ephemeral} into system messages
- Only applies to Anthropic API models (detected via get_llm_provider)
- Reuses AnthropicCacheControlHook._safe_insert_cache_control_in_message
**claude_code_block_expensive_flags**
- Blocks speed=fast (~6x pricing), inference_geo, thinking.type=enabled
- Config-driven via expensive_api_flags.yaml; inherits Anthropic hosted tool
prefixes from block_hosted_tools/anthropic.yaml via inherit_from
Policy templates updated:
- claude-code-pii-protection: uses block_hosted_tools + litellm_content_filter
- claude-code-cost-optimization: uses prompt_cache + block_expensive_flags
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Krrish Dholakia
parent
cf6b0e45a7
commit
44ee768307
12 changed files with 917 additions and 0 deletions
|
|
@ -674,5 +674,124 @@
|
|||
],
|
||||
"guardrails_remove": []
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "claude-code-pii-protection",
|
||||
"title": "Claude Code PII Protection",
|
||||
"description": "Prevents data leakage in Claude Code sessions by blocking all Anthropic-hosted tools (bash, web_search, web_fetch, code_execution, memory, text_editor) and masking PII and credentials found in request messages.",
|
||||
"icon": "ShieldCheckIcon",
|
||||
"iconColor": "text-indigo-500",
|
||||
"iconBg": "bg-indigo-50",
|
||||
"guardrails": [
|
||||
"claude-code-block-hosted-tools",
|
||||
"claude-code-pii-credentials",
|
||||
"claude-code-pii-personal-data"
|
||||
],
|
||||
"complexity": "Medium",
|
||||
"guardrailDefinitions": [
|
||||
{
|
||||
"guardrail_name": "claude-code-block-hosted-tools",
|
||||
"litellm_params": {
|
||||
"guardrail": "block_hosted_tools",
|
||||
"mode": "pre_call"
|
||||
},
|
||||
"guardrail_info": {
|
||||
"description": "Blocks platform-hosted tools from Anthropic (bash, web_search, web_fetch, code_execution, memory, text_editor, tool_search), OpenAI (code_interpreter, file_search, web_search_preview), and Gemini (googleSearch, code_execution, url_context). Only custom tool schemas are permitted."
|
||||
}
|
||||
},
|
||||
{
|
||||
"guardrail_name": "claude-code-pii-credentials",
|
||||
"litellm_params": {
|
||||
"guardrail": "litellm_content_filter",
|
||||
"mode": "pre_call",
|
||||
"patterns": [
|
||||
{"pattern_type": "prebuilt", "pattern_name": "aws_access_key", "action": "BLOCK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "aws_secret_key", "action": "BLOCK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "github_token", "action": "BLOCK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "slack_token", "action": "BLOCK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "generic_api_key", "action": "BLOCK"}
|
||||
],
|
||||
"pattern_redaction_format": "[{pattern_name}_REDACTED]"
|
||||
},
|
||||
"guardrail_info": {
|
||||
"description": "Blocks requests containing credentials or API keys (AWS, GitHub, Slack) that could be exfiltrated through the LLM context."
|
||||
}
|
||||
},
|
||||
{
|
||||
"guardrail_name": "claude-code-pii-personal-data",
|
||||
"litellm_params": {
|
||||
"guardrail": "litellm_content_filter",
|
||||
"mode": "pre_call",
|
||||
"patterns": [
|
||||
{"pattern_type": "prebuilt", "pattern_name": "us_ssn", "action": "MASK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "email", "action": "MASK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "us_phone", "action": "MASK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "credit_card", "action": "MASK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "visa", "action": "MASK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "mastercard", "action": "MASK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "iban", "action": "MASK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "ipv4", "action": "MASK"}
|
||||
],
|
||||
"pattern_redaction_format": "[{pattern_name}_REDACTED]"
|
||||
},
|
||||
"guardrail_info": {
|
||||
"description": "Masks SSNs, email addresses, phone numbers, credit card numbers, IBANs, and IP addresses before they enter the LLM context."
|
||||
}
|
||||
}
|
||||
],
|
||||
"templateData": {
|
||||
"policy_name": "claude-code-pii-protection",
|
||||
"description": "Prevents data leakage in Claude Code sessions. Blocks Anthropic-hosted tools and masks PII and credentials in request messages.",
|
||||
"guardrails_add": [
|
||||
"claude-code-block-hosted-tools",
|
||||
"claude-code-pii-credentials",
|
||||
"claude-code-pii-personal-data"
|
||||
],
|
||||
"guardrails_remove": []
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "claude-code-cost-optimization",
|
||||
"title": "Claude Code Cost Optimization",
|
||||
"description": "Reduces Claude Code API spend by blocking expensive inference modes (fast/turbo, inference_geo, and extended thinking) and automatically injecting prompt caching headers into system messages to maximize cache hit rates.",
|
||||
"icon": "CurrencyDollarIcon",
|
||||
"iconColor": "text-green-500",
|
||||
"iconBg": "bg-green-50",
|
||||
"guardrails": [
|
||||
"claude-code-inject-prompt-cache",
|
||||
"claude-code-block-expensive-flags"
|
||||
],
|
||||
"complexity": "Low",
|
||||
"guardrailDefinitions": [
|
||||
{
|
||||
"guardrail_name": "claude-code-inject-prompt-cache",
|
||||
"litellm_params": {
|
||||
"guardrail": "claude_code_prompt_cache",
|
||||
"mode": "pre_call"
|
||||
},
|
||||
"guardrail_info": {
|
||||
"description": "Automatically adds cache_control: {type: ephemeral} to system messages so Anthropic can cache the system prompt prefix. Only applies to Anthropic API models. Reduces cost on repeated calls that share the same system prompt."
|
||||
}
|
||||
},
|
||||
{
|
||||
"guardrail_name": "claude-code-block-expensive-flags",
|
||||
"litellm_params": {
|
||||
"guardrail": "claude_code_block_expensive_flags",
|
||||
"mode": "pre_call"
|
||||
},
|
||||
"guardrail_info": {
|
||||
"description": "Blocks expensive Anthropic API flags: fast/turbo inference (speed=fast, ~6x pricing), inference geo-routing (inference_geo, 1.1x pricing), and extended thinking (thinking.type=enabled). Also blocks Anthropic-hosted tools inherited from the hosted tools policy."
|
||||
}
|
||||
}
|
||||
],
|
||||
"templateData": {
|
||||
"policy_name": "claude-code-cost-optimization",
|
||||
"description": "Reduces Claude Code API spend by blocking fast mode, inference_geo, extended thinking, and Anthropic-hosted tools, while auto-injecting prompt caching into system messages.",
|
||||
"guardrails_add": [
|
||||
"claude-code-inject-prompt-cache",
|
||||
"claude-code-block-expensive-flags"
|
||||
],
|
||||
"guardrails_remove": []
|
||||
}
|
||||
}
|
||||
]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,43 @@
|
|||
"""
|
||||
Block Hosted Tools guardrail for LiteLLM.
|
||||
|
||||
Blocks platform-executed server-side tools from Anthropic, OpenAI, and Gemini.
|
||||
Provider tool lists are maintained in per-provider YAML files (anthropic.yaml,
|
||||
openai.yaml, gemini.yaml) alongside this module.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import litellm
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
||||
from .guardrail import BlockHostedToolsGuardrail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
def initialize_guardrail(
|
||||
litellm_params: "LitellmParams", guardrail: "Guardrail"
|
||||
) -> BlockHostedToolsGuardrail:
|
||||
guardrail_name = guardrail.get("guardrail_name")
|
||||
if not guardrail_name:
|
||||
raise ValueError("BlockHostedToolsGuardrail requires a guardrail_name")
|
||||
instance = BlockHostedToolsGuardrail(
|
||||
guardrail_name=guardrail_name,
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback(instance)
|
||||
return instance
|
||||
|
||||
|
||||
guardrail_initializer_registry = {
|
||||
SupportedGuardrailIntegrations.BLOCK_HOSTED_TOOLS.value: initialize_guardrail,
|
||||
}
|
||||
|
||||
guardrail_class_registry = {
|
||||
SupportedGuardrailIntegrations.BLOCK_HOSTED_TOOLS.value: BlockHostedToolsGuardrail,
|
||||
}
|
||||
|
||||
__all__ = ["BlockHostedToolsGuardrail", "initialize_guardrail"]
|
||||
|
|
@ -0,0 +1,24 @@
|
|||
# Anthropic server-side hosted tools.
|
||||
# Tools are matched by prefix against the tool 'type' field, because Anthropic
|
||||
# appends a version date suffix (e.g. bash_20250124, web_search_20250305).
|
||||
#
|
||||
# Reference: https://docs.anthropic.com/en/docs/agents-and-tools/computer-use
|
||||
|
||||
provider: "anthropic"
|
||||
description: "Anthropic server-side hosted tools executed on Anthropic infrastructure"
|
||||
|
||||
tool_type_prefixes:
|
||||
- prefix: "bash_"
|
||||
description: "Bash shell execution tool"
|
||||
- prefix: "web_search_"
|
||||
description: "Web search tool"
|
||||
- prefix: "web_fetch_"
|
||||
description: "Web fetch / URL retrieval tool"
|
||||
- prefix: "code_execution_"
|
||||
description: "Code execution sandbox tool"
|
||||
- prefix: "memory_"
|
||||
description: "Memory read/write tool"
|
||||
- prefix: "text_editor_"
|
||||
description: "Text editor tool"
|
||||
- prefix: "tool_search_"
|
||||
description: "Tool search / discovery tool"
|
||||
|
|
@ -0,0 +1,35 @@
|
|||
# Google Gemini / Vertex AI platform-hosted tools executed on Google infrastructure.
|
||||
# Gemini tools appear in two formats in the API:
|
||||
# - OpenAI-compat: tool dict has 'type' set to the value below
|
||||
# - Gemini native: tool dict has a top-level key matching the tool name
|
||||
# Both formats are checked.
|
||||
#
|
||||
# Reference: https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/function-calling
|
||||
|
||||
provider: "gemini"
|
||||
description: "Google Gemini / Vertex AI platform-hosted tools"
|
||||
|
||||
# Matched against the 'type' field (OpenAI-compat format)
|
||||
tool_type_exact:
|
||||
- type: "code_execution"
|
||||
description: "Code execution sandbox (OpenAI-compat format)"
|
||||
- type: "url_context"
|
||||
description: "URL fetching and context grounding (OpenAI-compat format)"
|
||||
|
||||
# Matched as top-level dict keys (Gemini native format)
|
||||
# e.g. {"googleSearch": {}}, {"codeExecution": {}}
|
||||
tool_top_level_keys:
|
||||
- key: "googleSearch"
|
||||
description: "Google Search grounding"
|
||||
- key: "googleSearchRetrieval"
|
||||
description: "Google Search retrieval with data store integration"
|
||||
- key: "enterpriseWebSearch"
|
||||
description: "Enterprise web search"
|
||||
- key: "url_context"
|
||||
description: "URL context grounding"
|
||||
- key: "code_execution"
|
||||
description: "Code execution sandbox"
|
||||
- key: "googleMaps"
|
||||
description: "Google Maps integration"
|
||||
- key: "computerUse"
|
||||
description: "Computer use / automated screen interaction"
|
||||
|
|
@ -0,0 +1,177 @@
|
|||
"""
|
||||
Block Hosted Tools Guardrail
|
||||
|
||||
Provider-agnostic guardrail that blocks platform-executed hosted tools from
|
||||
Anthropic, OpenAI, and Google Gemini. Tool definitions are loaded from
|
||||
per-provider YAML files in this directory so the list can be extended without
|
||||
touching Python code.
|
||||
|
||||
Detection strategy per provider:
|
||||
Anthropic — prefix match on tool 'type' field (versioned: bash_20250124, etc.)
|
||||
OpenAI — exact match on tool 'type' field (code_interpreter, file_search, …)
|
||||
Gemini — exact match on tool 'type' field OR presence of a top-level key
|
||||
(googleSearch, codeExecution, …) in the tool dict
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Dict, FrozenSet, List, Literal, Optional, Tuple
|
||||
|
||||
import yaml
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
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
|
||||
|
||||
_DIR = os.path.dirname(__file__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config loading — runs once at import time
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _load_provider_configs() -> Dict[str, dict]:
|
||||
"""Load all *.yaml files in this directory, keyed by provider name."""
|
||||
configs: Dict[str, dict] = {}
|
||||
for fname in os.listdir(_DIR):
|
||||
if not fname.endswith(".yaml"):
|
||||
continue
|
||||
with open(os.path.join(_DIR, fname)) as f:
|
||||
cfg = yaml.safe_load(f)
|
||||
provider = cfg.get("provider")
|
||||
if provider:
|
||||
configs[provider] = cfg
|
||||
return configs
|
||||
|
||||
|
||||
def _build_match_sets(
|
||||
configs: Dict[str, dict],
|
||||
) -> Tuple[tuple, FrozenSet[str], FrozenSet[str]]:
|
||||
"""
|
||||
Derive three matching structures from all provider configs:
|
||||
|
||||
Returns:
|
||||
type_prefixes — tuple of prefix strings (Anthropic-style versioned types)
|
||||
exact_types — frozenset of exact type strings (OpenAI / Gemini-compat)
|
||||
top_level_keys — frozenset of Gemini native top-level dict keys
|
||||
"""
|
||||
prefixes: List[str] = []
|
||||
exact: List[str] = []
|
||||
keys: List[str] = []
|
||||
|
||||
for cfg in configs.values():
|
||||
for entry in cfg.get("tool_type_prefixes", []):
|
||||
if isinstance(entry, dict) and "prefix" in entry:
|
||||
prefixes.append(entry["prefix"])
|
||||
for entry in cfg.get("tool_type_exact", []):
|
||||
if isinstance(entry, dict) and "type" in entry:
|
||||
exact.append(entry["type"])
|
||||
for entry in cfg.get("tool_top_level_keys", []):
|
||||
if isinstance(entry, dict) and "key" in entry:
|
||||
keys.append(entry["key"])
|
||||
|
||||
return tuple(prefixes), frozenset(exact), frozenset(keys)
|
||||
|
||||
|
||||
_PROVIDER_CONFIGS = _load_provider_configs()
|
||||
_TYPE_PREFIXES, _EXACT_TYPES, _TOP_LEVEL_KEYS = _build_match_sets(_PROVIDER_CONFIGS)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tool-matching helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _match_tool(tool: dict) -> Optional[str]:
|
||||
"""
|
||||
Return a human-readable description if the tool is a known hosted tool,
|
||||
or None if it is a user-defined tool (pass through).
|
||||
"""
|
||||
if not isinstance(tool, dict):
|
||||
return None
|
||||
|
||||
raw_type: Optional[str] = tool.get("type")
|
||||
|
||||
# Skip generic OpenAI function wrapper — never a hosted tool
|
||||
if raw_type == "function":
|
||||
return None
|
||||
|
||||
if raw_type:
|
||||
# Prefix match (Anthropic versioned types)
|
||||
if raw_type.startswith(_TYPE_PREFIXES):
|
||||
name = tool.get("name") or raw_type
|
||||
return f"{name} ({raw_type})"
|
||||
# Exact match (OpenAI / Gemini-compat)
|
||||
if raw_type in _EXACT_TYPES:
|
||||
name = tool.get("name") or raw_type
|
||||
return f"{name} ({raw_type})"
|
||||
|
||||
# Gemini native top-level keys ({"googleSearch": {}, ...})
|
||||
for key in _TOP_LEVEL_KEYS:
|
||||
if key in tool:
|
||||
return f"{key}"
|
||||
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Guardrail class
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class BlockHostedToolsGuardrail(CustomGuardrail):
|
||||
"""
|
||||
Guardrail that blocks platform-hosted tools from Anthropic, OpenAI, and Gemini.
|
||||
|
||||
Raises HTTP 403 if any tool in the request matches a known hosted tool
|
||||
from any supported provider. Provider tool lists are defined in YAML
|
||||
files alongside this module and loaded at import time.
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
if "supported_event_hooks" not in kwargs:
|
||||
kwargs["supported_event_hooks"] = [GuardrailEventHooks.pre_call]
|
||||
super().__init__(**kwargs)
|
||||
verbose_proxy_logger.debug(
|
||||
f"BlockHostedToolsGuardrail initialized "
|
||||
f"(providers: {list(_PROVIDER_CONFIGS.keys())})"
|
||||
)
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
if input_type != "request":
|
||||
return inputs
|
||||
|
||||
tools: List[dict] = list(inputs.get("tools") or []) # type: ignore[assignment]
|
||||
blocked: List[str] = []
|
||||
|
||||
for tool in tools:
|
||||
desc = _match_tool(tool)
|
||||
if desc:
|
||||
blocked.append(desc)
|
||||
|
||||
if blocked:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": (
|
||||
"the following platform-hosted tools are disabled by "
|
||||
"your organization's policy: " + ", ".join(blocked)
|
||||
),
|
||||
"guardrail": self.guardrail_name,
|
||||
"blocked_tools": blocked,
|
||||
},
|
||||
)
|
||||
|
||||
return inputs
|
||||
|
|
@ -0,0 +1,19 @@
|
|||
# OpenAI platform-hosted tools executed on OpenAI's infrastructure.
|
||||
# These are matched by exact 'type' field value (no version suffix).
|
||||
#
|
||||
# Reference: https://platform.openai.com/docs/guides/tools
|
||||
|
||||
provider: "openai"
|
||||
description: "OpenAI platform-hosted tools executed on OpenAI infrastructure"
|
||||
|
||||
tool_type_exact:
|
||||
- type: "code_interpreter"
|
||||
description: "Python code execution and file output sandbox"
|
||||
- type: "file_search"
|
||||
description: "Semantic file search across vector stores"
|
||||
- type: "web_search_preview"
|
||||
description: "Bing-powered live web search"
|
||||
- type: "computer_use_preview"
|
||||
description: "Computer use / automated screen interaction"
|
||||
- type: "shell"
|
||||
description: "Shell command execution in containers"
|
||||
|
|
@ -0,0 +1,77 @@
|
|||
"""
|
||||
Claude Code guardrail integrations for LiteLLM.
|
||||
|
||||
Two focused policy-enforcement guardrails for Claude Code deployments:
|
||||
|
||||
1. claude_code_prompt_cache — auto-inject Anthropic prompt-caching headers
|
||||
2. claude_code_block_expensive_flags — block expensive API flags (fast mode, etc.)
|
||||
|
||||
Hosted tool blocking is handled by the provider-agnostic 'block_hosted_tools'
|
||||
guardrail (see guardrail_hooks/block_hosted_tools/).
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import litellm
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
||||
from .block_expensive_flags import ClaudeCodeBlockExpensiveFlagsGuardrail
|
||||
from .prompt_cache import ClaudeCodePromptCacheGuardrail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Per-guardrail initializer functions #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
|
||||
def _init_prompt_cache(
|
||||
litellm_params: "LitellmParams", guardrail: "Guardrail"
|
||||
) -> ClaudeCodePromptCacheGuardrail:
|
||||
guardrail_name = guardrail.get("guardrail_name")
|
||||
if not guardrail_name:
|
||||
raise ValueError("ClaudeCodePromptCacheGuardrail requires a guardrail_name")
|
||||
instance = ClaudeCodePromptCacheGuardrail(
|
||||
guardrail_name=guardrail_name,
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback(instance)
|
||||
return instance
|
||||
|
||||
|
||||
def _init_block_expensive_flags(
|
||||
litellm_params: "LitellmParams", guardrail: "Guardrail"
|
||||
) -> ClaudeCodeBlockExpensiveFlagsGuardrail:
|
||||
guardrail_name = guardrail.get("guardrail_name")
|
||||
if not guardrail_name:
|
||||
raise ValueError("ClaudeCodeBlockExpensiveFlagsGuardrail requires a guardrail_name")
|
||||
instance = ClaudeCodeBlockExpensiveFlagsGuardrail(
|
||||
guardrail_name=guardrail_name,
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback(instance)
|
||||
return instance
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Registries consumed by the guardrail loader #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
guardrail_initializer_registry = {
|
||||
SupportedGuardrailIntegrations.CLAUDE_CODE_PROMPT_CACHE.value: _init_prompt_cache,
|
||||
SupportedGuardrailIntegrations.CLAUDE_CODE_BLOCK_EXPENSIVE_FLAGS.value: _init_block_expensive_flags,
|
||||
}
|
||||
|
||||
guardrail_class_registry = {
|
||||
SupportedGuardrailIntegrations.CLAUDE_CODE_PROMPT_CACHE.value: ClaudeCodePromptCacheGuardrail,
|
||||
SupportedGuardrailIntegrations.CLAUDE_CODE_BLOCK_EXPENSIVE_FLAGS.value: ClaudeCodeBlockExpensiveFlagsGuardrail,
|
||||
}
|
||||
|
||||
__all__ = [
|
||||
"ClaudeCodePromptCacheGuardrail",
|
||||
"ClaudeCodeBlockExpensiveFlagsGuardrail",
|
||||
]
|
||||
|
|
@ -0,0 +1,175 @@
|
|||
"""
|
||||
Claude Code - Block Expensive API Flags Guardrail
|
||||
|
||||
Blocks Anthropic API parameters that trigger feature-specific pricing surcharges
|
||||
(fast mode, inference_geo, extended thinking). Also inherits the hosted tool
|
||||
type prefixes from hosted_tool_types.yaml so hosted tools are blocked here too.
|
||||
|
||||
Blocked params are driven by expensive_api_flags.yaml which references
|
||||
hosted_tool_types.yaml via `inherit_from`, following the same pattern as
|
||||
harmful_child_safety.yaml inherits from harm_toxic_abuse.json.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any, List, Literal, Optional
|
||||
|
||||
import yaml
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
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
|
||||
|
||||
_DIR = os.path.dirname(__file__)
|
||||
_FLAGS_YAML = os.path.join(_DIR, "expensive_api_flags.yaml")
|
||||
_TOOLS_YAML = os.path.join(_DIR, "hosted_tool_types.yaml")
|
||||
|
||||
|
||||
def _load_config() -> dict:
|
||||
"""Load expensive_api_flags.yaml, merging any inherited config."""
|
||||
with open(_FLAGS_YAML) as f:
|
||||
config: dict = yaml.safe_load(f)
|
||||
|
||||
inherited_prefixes: List[str] = []
|
||||
inherit_from = config.get("inherit_from")
|
||||
if inherit_from:
|
||||
inherit_path = os.path.join(_DIR, inherit_from)
|
||||
with open(inherit_path) as f:
|
||||
inherited: dict = yaml.safe_load(f)
|
||||
inherited_prefixes = [
|
||||
entry["prefix"]
|
||||
for entry in inherited.get("tool_type_prefixes", [])
|
||||
if isinstance(entry, dict) and "prefix" in entry
|
||||
]
|
||||
|
||||
config["_inherited_tool_type_prefixes"] = inherited_prefixes
|
||||
return config
|
||||
|
||||
|
||||
_CONFIG: dict = _load_config()
|
||||
_BLOCKED_PARAMS: List[dict] = _CONFIG.get("blocked_params", [])
|
||||
_INHERITED_TOOL_TYPE_PREFIXES: tuple = tuple(_CONFIG.get("_inherited_tool_type_prefixes", []))
|
||||
|
||||
|
||||
def _tool_type(tool: dict) -> Optional[str]:
|
||||
t = tool.get("type")
|
||||
if t and t != "function":
|
||||
return t
|
||||
return None
|
||||
|
||||
|
||||
def _is_hosted_tool(tool: dict) -> bool:
|
||||
if not _INHERITED_TOOL_TYPE_PREFIXES:
|
||||
return False
|
||||
t = _tool_type(tool)
|
||||
if not t:
|
||||
return False
|
||||
return t.startswith(_INHERITED_TOOL_TYPE_PREFIXES)
|
||||
|
||||
|
||||
def _check_param(
|
||||
request_data: dict, param_cfg: dict
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Return an error message if the param in request_data matches a blocked value.
|
||||
Returns None when the param is not blocked.
|
||||
"""
|
||||
param = param_cfg.get("param")
|
||||
if not param:
|
||||
return None
|
||||
|
||||
value: Any = request_data.get(param)
|
||||
if value is None:
|
||||
return None
|
||||
|
||||
nested_key = param_cfg.get("nested_key")
|
||||
if nested_key:
|
||||
# e.g. thinking.type — value must be a dict
|
||||
if not isinstance(value, dict):
|
||||
return None
|
||||
value = value.get(nested_key)
|
||||
if value is None:
|
||||
return None
|
||||
|
||||
blocked_values: List[str] = param_cfg.get("blocked_values", [])
|
||||
if "*" in blocked_values or str(value) in blocked_values:
|
||||
return param_cfg.get("error_message", f"{param} is disabled by your organization's policy")
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class ClaudeCodeBlockExpensiveFlagsGuardrail(CustomGuardrail):
|
||||
"""
|
||||
Guardrail that blocks expensive Anthropic API flags.
|
||||
|
||||
Checks request_data for feature-specific pricing flags (fast mode,
|
||||
inference_geo, extended thinking) and Anthropic-hosted tools inherited
|
||||
from hosted_tool_types.yaml. Raises HTTP 403 on the first violation.
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
if "supported_event_hooks" not in kwargs:
|
||||
kwargs["supported_event_hooks"] = [GuardrailEventHooks.pre_call]
|
||||
super().__init__(**kwargs)
|
||||
verbose_proxy_logger.debug("ClaudeCodeBlockExpensiveFlagsGuardrail initialized")
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
if input_type != "request":
|
||||
return inputs
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 1. Check request-level blocked params (speed, inference_geo, etc.) #
|
||||
# ------------------------------------------------------------------ #
|
||||
for param_cfg in _BLOCKED_PARAMS:
|
||||
error_msg = _check_param(request_data, param_cfg)
|
||||
if error_msg:
|
||||
param = param_cfg.get("param", "unknown")
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": error_msg,
|
||||
"guardrail": self.guardrail_name,
|
||||
"blocked_param": param,
|
||||
},
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 2. Check for inherited hosted tools (from hosted_tool_types.yaml) #
|
||||
# ------------------------------------------------------------------ #
|
||||
if _INHERITED_TOOL_TYPE_PREFIXES:
|
||||
tools: List[dict] = list(inputs.get("tools") or []) # type: ignore[assignment]
|
||||
blocked_tools: List[str] = []
|
||||
for tool in tools:
|
||||
if _is_hosted_tool(tool):
|
||||
tool_type = _tool_type(tool) or "unknown"
|
||||
tool_name = tool.get("name") or tool_type
|
||||
blocked_tools.append(f"{tool_name} ({tool_type})")
|
||||
|
||||
if blocked_tools:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": (
|
||||
"the following Anthropic-hosted tools are disabled by "
|
||||
"your organization's policy: " + ", ".join(blocked_tools)
|
||||
),
|
||||
"guardrail": self.guardrail_name,
|
||||
"blocked_tools": blocked_tools,
|
||||
},
|
||||
)
|
||||
|
||||
return inputs
|
||||
|
|
@ -0,0 +1,35 @@
|
|||
# Expensive Anthropic API flags.
|
||||
# Blocks request parameters that trigger feature-specific pricing surcharges.
|
||||
# Also inherits hosted tool type prefixes so hosted tools are blocked as well.
|
||||
#
|
||||
# Pricing reference: https://platform.claude.com/docs/en/about-claude/pricing#feature-specific-pricing
|
||||
|
||||
category_name: "claude_code_expensive_flags"
|
||||
description: "Blocks expensive Anthropic API features that incur pricing surcharges"
|
||||
default_action: "BLOCK"
|
||||
|
||||
# Inherit hosted tool type prefixes — hosted tools are also expensive.
|
||||
inherit_from: "../block_hosted_tools/anthropic.yaml"
|
||||
|
||||
# Request-level parameters to block.
|
||||
# Each entry maps a top-level request_data key to one or more blocked values.
|
||||
# Use "*" as the value to block the flag entirely regardless of its value.
|
||||
blocked_params:
|
||||
- param: "speed"
|
||||
blocked_values:
|
||||
- "fast"
|
||||
description: "Fast/turbo inference mode (~6x pricing surcharge)"
|
||||
error_message: "fast inference mode (speed=fast) is disabled by your organization's policy"
|
||||
|
||||
- param: "inference_geo"
|
||||
blocked_values:
|
||||
- "*"
|
||||
description: "Inference geo-routing (1.1x pricing surcharge)"
|
||||
error_message: "inference geo-routing (inference_geo) is disabled by your organization's policy"
|
||||
|
||||
- param: "thinking"
|
||||
blocked_values:
|
||||
- "enabled"
|
||||
nested_key: "type"
|
||||
description: "Extended thinking / budget-token thinking mode"
|
||||
error_message: "extended thinking mode (thinking.type=enabled) is disabled by your organization's policy"
|
||||
|
|
@ -0,0 +1,91 @@
|
|||
"""
|
||||
Claude Code - Prompt Cache Injection Guardrail
|
||||
|
||||
Automatically injects cache_control: {type: ephemeral} into system messages
|
||||
so Anthropic can cache the prompt prefix, reducing costs on repeated calls.
|
||||
|
||||
Only runs when the request targets an Anthropic API model. Uses the existing
|
||||
AnthropicCacheControlHook._safe_insert_cache_control_in_message utility so
|
||||
the injection logic is not duplicated.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Literal, Optional
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.llms.openai import ChatCompletionCachedContent
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
from litellm.utils import get_llm_provider
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
|
||||
def _is_anthropic_model(model: Optional[str]) -> bool:
|
||||
"""Return True when the model resolves to the Anthropic provider."""
|
||||
if not model:
|
||||
return False
|
||||
try:
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model=model)
|
||||
return custom_llm_provider == "anthropic"
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
class ClaudeCodePromptCacheGuardrail(CustomGuardrail):
|
||||
"""
|
||||
Guardrail that injects Anthropic prompt-caching headers into system messages.
|
||||
|
||||
Targets only Anthropic API models — the provider is detected from the
|
||||
model string so this guardrail is safe to apply globally without
|
||||
restricting other providers.
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
if "supported_event_hooks" not in kwargs:
|
||||
kwargs["supported_event_hooks"] = [GuardrailEventHooks.pre_call]
|
||||
super().__init__(**kwargs)
|
||||
verbose_proxy_logger.debug("ClaudeCodePromptCacheGuardrail initialized")
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
if input_type != "request":
|
||||
return inputs
|
||||
|
||||
model: Optional[str] = request_data.get("model") or inputs.get("model") # type: ignore[assignment]
|
||||
if not _is_anthropic_model(model):
|
||||
verbose_proxy_logger.debug(
|
||||
f"ClaudeCodePromptCacheGuardrail: skipping non-Anthropic model '{model}'"
|
||||
)
|
||||
return inputs
|
||||
|
||||
messages = request_data.get("messages") or []
|
||||
if not messages:
|
||||
return inputs
|
||||
|
||||
control = ChatCompletionCachedContent(type="ephemeral")
|
||||
modified = []
|
||||
for msg in messages:
|
||||
if msg.get("role") == "system":
|
||||
msg = AnthropicCacheControlHook._safe_insert_cache_control_in_message(
|
||||
message=msg, # type: ignore[arg-type]
|
||||
control=control,
|
||||
)
|
||||
modified.append(msg)
|
||||
|
||||
request_data["messages"] = modified
|
||||
verbose_proxy_logger.debug(
|
||||
"ClaudeCodePromptCacheGuardrail: injected cache_control into system messages"
|
||||
)
|
||||
return inputs
|
||||
|
|
@ -69,6 +69,9 @@ class SupportedGuardrailIntegrations(Enum):
|
|||
GENERIC_GUARDRAIL_API = "generic_guardrail_api"
|
||||
QUALIFIRE = "qualifire"
|
||||
CUSTOM_CODE = "custom_code"
|
||||
BLOCK_HOSTED_TOOLS = "block_hosted_tools"
|
||||
CLAUDE_CODE_PROMPT_CACHE = "claude_code_prompt_cache"
|
||||
CLAUDE_CODE_BLOCK_EXPENSIVE_FLAGS = "claude_code_block_expensive_flags"
|
||||
|
||||
|
||||
class Role(Enum):
|
||||
|
|
|
|||
|
|
@ -674,5 +674,124 @@
|
|||
],
|
||||
"guardrails_remove": []
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "claude-code-pii-protection",
|
||||
"title": "Claude Code PII Protection",
|
||||
"description": "Prevents data leakage in Claude Code sessions by blocking all Anthropic-hosted tools (bash, web_search, web_fetch, code_execution, memory, text_editor) and masking PII and credentials found in request messages.",
|
||||
"icon": "ShieldCheckIcon",
|
||||
"iconColor": "text-indigo-500",
|
||||
"iconBg": "bg-indigo-50",
|
||||
"guardrails": [
|
||||
"claude-code-block-hosted-tools",
|
||||
"claude-code-pii-credentials",
|
||||
"claude-code-pii-personal-data"
|
||||
],
|
||||
"complexity": "Medium",
|
||||
"guardrailDefinitions": [
|
||||
{
|
||||
"guardrail_name": "claude-code-block-hosted-tools",
|
||||
"litellm_params": {
|
||||
"guardrail": "block_hosted_tools",
|
||||
"mode": "pre_call"
|
||||
},
|
||||
"guardrail_info": {
|
||||
"description": "Blocks platform-hosted tools from Anthropic (bash, web_search, web_fetch, code_execution, memory, text_editor, tool_search), OpenAI (code_interpreter, file_search, web_search_preview), and Gemini (googleSearch, code_execution, url_context). Only custom tool schemas are permitted."
|
||||
}
|
||||
},
|
||||
{
|
||||
"guardrail_name": "claude-code-pii-credentials",
|
||||
"litellm_params": {
|
||||
"guardrail": "litellm_content_filter",
|
||||
"mode": "pre_call",
|
||||
"patterns": [
|
||||
{"pattern_type": "prebuilt", "pattern_name": "aws_access_key", "action": "BLOCK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "aws_secret_key", "action": "BLOCK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "github_token", "action": "BLOCK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "slack_token", "action": "BLOCK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "generic_api_key", "action": "BLOCK"}
|
||||
],
|
||||
"pattern_redaction_format": "[{pattern_name}_REDACTED]"
|
||||
},
|
||||
"guardrail_info": {
|
||||
"description": "Blocks requests containing credentials or API keys (AWS, GitHub, Slack) that could be exfiltrated through the LLM context."
|
||||
}
|
||||
},
|
||||
{
|
||||
"guardrail_name": "claude-code-pii-personal-data",
|
||||
"litellm_params": {
|
||||
"guardrail": "litellm_content_filter",
|
||||
"mode": "pre_call",
|
||||
"patterns": [
|
||||
{"pattern_type": "prebuilt", "pattern_name": "us_ssn", "action": "MASK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "email", "action": "MASK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "us_phone", "action": "MASK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "credit_card", "action": "MASK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "visa", "action": "MASK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "mastercard", "action": "MASK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "iban", "action": "MASK"},
|
||||
{"pattern_type": "prebuilt", "pattern_name": "ipv4", "action": "MASK"}
|
||||
],
|
||||
"pattern_redaction_format": "[{pattern_name}_REDACTED]"
|
||||
},
|
||||
"guardrail_info": {
|
||||
"description": "Masks SSNs, email addresses, phone numbers, credit card numbers, IBANs, and IP addresses before they enter the LLM context."
|
||||
}
|
||||
}
|
||||
],
|
||||
"templateData": {
|
||||
"policy_name": "claude-code-pii-protection",
|
||||
"description": "Prevents data leakage in Claude Code sessions. Blocks Anthropic-hosted tools and masks PII and credentials in request messages.",
|
||||
"guardrails_add": [
|
||||
"claude-code-block-hosted-tools",
|
||||
"claude-code-pii-credentials",
|
||||
"claude-code-pii-personal-data"
|
||||
],
|
||||
"guardrails_remove": []
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "claude-code-cost-optimization",
|
||||
"title": "Claude Code Cost Optimization",
|
||||
"description": "Reduces Claude Code API spend by blocking expensive inference modes (fast/turbo, inference_geo, and extended thinking) and automatically injecting prompt caching headers into system messages to maximize cache hit rates.",
|
||||
"icon": "CurrencyDollarIcon",
|
||||
"iconColor": "text-green-500",
|
||||
"iconBg": "bg-green-50",
|
||||
"guardrails": [
|
||||
"claude-code-inject-prompt-cache",
|
||||
"claude-code-block-expensive-flags"
|
||||
],
|
||||
"complexity": "Low",
|
||||
"guardrailDefinitions": [
|
||||
{
|
||||
"guardrail_name": "claude-code-inject-prompt-cache",
|
||||
"litellm_params": {
|
||||
"guardrail": "claude_code_prompt_cache",
|
||||
"mode": "pre_call"
|
||||
},
|
||||
"guardrail_info": {
|
||||
"description": "Automatically adds cache_control: {type: ephemeral} to system messages so Anthropic can cache the system prompt prefix. Only applies to Anthropic API models. Reduces cost on repeated calls that share the same system prompt."
|
||||
}
|
||||
},
|
||||
{
|
||||
"guardrail_name": "claude-code-block-expensive-flags",
|
||||
"litellm_params": {
|
||||
"guardrail": "claude_code_block_expensive_flags",
|
||||
"mode": "pre_call"
|
||||
},
|
||||
"guardrail_info": {
|
||||
"description": "Blocks expensive Anthropic API flags: fast/turbo inference (speed=fast, ~6x pricing), inference geo-routing (inference_geo, 1.1x pricing), and extended thinking (thinking.type=enabled). Also blocks Anthropic-hosted tools inherited from the hosted tools policy."
|
||||
}
|
||||
}
|
||||
],
|
||||
"templateData": {
|
||||
"policy_name": "claude-code-cost-optimization",
|
||||
"description": "Reduces Claude Code API spend by blocking fast mode, inference_geo, extended thinking, and Anthropic-hosted tools, while auto-injecting prompt caching into system messages.",
|
||||
"guardrails_add": [
|
||||
"claude-code-inject-prompt-cache",
|
||||
"claude-code-block-expensive-flags"
|
||||
],
|
||||
"guardrails_remove": []
|
||||
}
|
||||
}
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue