mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
Rewrite TealTiger guardrail: per-user budgets, precise PII defaults, typed policies, decision receipts
This commit is contained in:
AlphaRex-pixel
parent
f299dd191e
commit
c63a6d86c1
10 changed files with 916 additions and 620 deletions
|
|
@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Final
|
|||
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
||||
from .config import parse_config
|
||||
from .tealtiger import TealTigerGuardrail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -11,9 +12,12 @@ if TYPE_CHECKING:
|
|||
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> TealTigerGuardrail:
|
||||
import litellm
|
||||
|
||||
config: Final = parse_config(
|
||||
getattr(litellm_params, "policies", None), getattr(litellm_params, "policy_mode", None)
|
||||
)
|
||||
_tealtiger_callback: Final = TealTigerGuardrail(
|
||||
policies=getattr(litellm_params, "policies", None),
|
||||
policy_mode=getattr(litellm_params, "policy_mode", "ENFORCE"),
|
||||
policies=config.policies,
|
||||
policy_mode=config.policy_mode,
|
||||
guardrail_name=guardrail["guardrail_name"],
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
|
|
|
|||
38
litellm/proxy/guardrails/guardrail_hooks/tealtiger/budget.py
Normal file
38
litellm/proxy/guardrails/guardrail_hooks/tealtiger/budget.py
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
import threading
|
||||
from datetime import date
|
||||
from typing import Final
|
||||
|
||||
|
||||
def _today() -> str:
|
||||
return date.today().isoformat() # noqa: DTZ011 # daily boundary is intentionally local-server-time, not UTC
|
||||
|
||||
|
||||
class BudgetLedger:
|
||||
"""In-memory spend per user (daily) and per session. Per-process, and reset when the local day rolls over."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._lock: Final = threading.Lock()
|
||||
self._day = _today()
|
||||
self._by_user: dict[str, float] = {} # mutable-ok: live counters guarded by the lock
|
||||
self._by_session: dict[str, float] = {} # mutable-ok: live counters guarded by the lock
|
||||
|
||||
def _roll_day_locked(self) -> None:
|
||||
today: Final = _today()
|
||||
if today == self._day:
|
||||
return
|
||||
self._day = today
|
||||
self._by_user = {}
|
||||
self._by_session = {}
|
||||
|
||||
def spent(self, user: str, session: str | None) -> tuple[float, float]:
|
||||
"""Return (spent today by this user, spent by this session)."""
|
||||
with self._lock:
|
||||
self._roll_day_locked()
|
||||
return self._by_user.get(user, 0.0), self._by_session.get(session, 0.0) if session else 0.0
|
||||
|
||||
def record(self, user: str, session: str | None, cost_usd: float) -> None:
|
||||
with self._lock:
|
||||
self._roll_day_locked()
|
||||
self._by_user[user] = self._by_user.get(user, 0.0) + cost_usd
|
||||
if session:
|
||||
self._by_session[session] = self._by_session.get(session, 0.0) + cost_usd
|
||||
34
litellm/proxy/guardrails/guardrail_hooks/tealtiger/config.py
Normal file
34
litellm/proxy/guardrails/guardrail_hooks/tealtiger/config.py
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
from itertools import chain
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.tealtiger import (
|
||||
DEFAULT_POLICIES,
|
||||
PiiPolicy,
|
||||
Policy,
|
||||
PolicyMode,
|
||||
)
|
||||
|
||||
from .patterns import unknown_pattern_names
|
||||
|
||||
|
||||
class TealTigerConfig(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
policies: tuple[Policy, ...] = DEFAULT_POLICIES
|
||||
policy_mode: PolicyMode = PolicyMode.ENFORCE
|
||||
|
||||
|
||||
def parse_config(policies: object, policy_mode: object) -> TealTigerConfig:
|
||||
"""Validate the raw `policies` / `policy_mode` litellm_params into typed config, failing at startup."""
|
||||
raw: Final = {
|
||||
key: value for key, value in (("policies", policies), ("policy_mode", policy_mode)) if value is not None
|
||||
}
|
||||
config: Final = TealTigerConfig.model_validate(raw)
|
||||
unknown: Final = tuple(
|
||||
chain.from_iterable(unknown_pattern_names(p.patterns) for p in config.policies if isinstance(p, PiiPolicy))
|
||||
)
|
||||
if unknown:
|
||||
raise ValueError(f"TealTiger: unknown pii pattern or category name(s): {', '.join(unknown)}")
|
||||
return config
|
||||
|
|
@ -0,0 +1,47 @@
|
|||
import json
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
|
||||
class Action(str, Enum):
|
||||
ALLOW = "ALLOW"
|
||||
REDACT = "REDACT"
|
||||
BLOCK = "BLOCK"
|
||||
|
||||
|
||||
class RiskLevel(str, Enum):
|
||||
NONE = "NONE"
|
||||
MEDIUM = "MEDIUM"
|
||||
HIGH = "HIGH"
|
||||
CRITICAL = "CRITICAL"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Decision:
|
||||
"""The outcome of one governance check. Every check returns one, and each is logged as a receipt."""
|
||||
|
||||
action: Action
|
||||
reason_code: str
|
||||
risk_level: RiskLevel
|
||||
correlation_id: str
|
||||
redacted_text: str | None = None
|
||||
details: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({}))
|
||||
|
||||
def receipt_json(self) -> str:
|
||||
"""Typed, JSON-serializable decision receipt for audit trails."""
|
||||
return json.dumps(
|
||||
{
|
||||
"correlation_id": self.correlation_id,
|
||||
"action": self.action.value,
|
||||
"reason_code": self.reason_code,
|
||||
"risk_level": self.risk_level.value,
|
||||
"details": dict(self.details),
|
||||
},
|
||||
sort_keys=True,
|
||||
)
|
||||
|
||||
|
||||
NO_VIOLATIONS: Final = "NO_VIOLATIONS"
|
||||
|
|
@ -1,182 +1,139 @@
|
|||
"""
|
||||
TealEngine: deterministic, dependency-free PII / cost / tool-auth policy
|
||||
evaluation. No network calls in the governance path.
|
||||
TealEngine: deterministic, dependency-free PII / cost / tool-auth policy evaluation.
|
||||
No network calls in the governance path. Every check returns a Decision value.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import re
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import date
|
||||
from enum import Enum
|
||||
from itertools import chain
|
||||
from dataclasses import replace
|
||||
from functools import reduce
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from .patterns import PATTERN_CATEGORIES, default_patterns
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.tealtiger import (
|
||||
CostPolicy,
|
||||
PiiPolicy,
|
||||
Policy,
|
||||
PolicyMode,
|
||||
ToolAuthPolicy,
|
||||
)
|
||||
|
||||
from .budget import BudgetLedger
|
||||
from .decision import NO_VIOLATIONS, Action, Decision, RiskLevel
|
||||
from .patterns import CRITICAL_PATTERNS, resolve_scope
|
||||
from .request import Identity
|
||||
|
||||
DEFAULT_COST_PER_1K_TOKENS: Final = 0.002
|
||||
_CHARS_PER_TOKEN: Final = 4
|
||||
_NO_DETAILS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_NO_PATTERNS: Final[Mapping[str, re.Pattern[str]]] = MappingProxyType({})
|
||||
|
||||
|
||||
def estimate_tokens(texts: Sequence[str]) -> int:
|
||||
"""Rough, dependency-free token estimate (about 4 characters per token)."""
|
||||
"""Rough token estimate (about 4 characters per token). apply_guardrail is not handed usage data."""
|
||||
return math.ceil(sum(len(text) for text in texts) / _CHARS_PER_TOKEN)
|
||||
|
||||
|
||||
class PolicyMode(str, Enum):
|
||||
ENFORCE = "ENFORCE"
|
||||
MONITOR = "MONITOR"
|
||||
|
||||
|
||||
class Action(str, Enum):
|
||||
ALLOW = "ALLOW"
|
||||
REDACT = "REDACT"
|
||||
BLOCK = "BLOCK"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Decision:
|
||||
action: str
|
||||
reason_code: str
|
||||
correlation_id: str
|
||||
redacted_text: str | None = None
|
||||
details: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({}))
|
||||
|
||||
def to_receipt(self) -> Mapping[str, object]:
|
||||
return MappingProxyType(
|
||||
{
|
||||
"correlation_id": self.correlation_id,
|
||||
"action": self.action,
|
||||
"reason_code": self.reason_code,
|
||||
"details": self.details,
|
||||
}
|
||||
)
|
||||
def _estimate_cost_usd(policy: CostPolicy, texts: Sequence[str]) -> float:
|
||||
return estimate_tokens(texts) / 1000.0 * policy.cost_per_1k_tokens
|
||||
|
||||
|
||||
class TealEngine:
|
||||
def __init__(
|
||||
self,
|
||||
policies: Sequence[Mapping[str, object]],
|
||||
policies: Sequence[Policy],
|
||||
mode: PolicyMode = PolicyMode.ENFORCE,
|
||||
ledger: BudgetLedger | None = None,
|
||||
) -> None:
|
||||
self.policies = policies
|
||||
self.mode = mode
|
||||
self.patterns = default_patterns()
|
||||
self._lock = threading.Lock()
|
||||
self._spend_by_day: dict[str, float] = {} # mutable-ok: live counter under lock
|
||||
|
||||
self._pii_policy = next((p for p in policies if p["type"] == "pii"), None)
|
||||
self._cost_policy = next((p for p in policies if p["type"] == "cost"), None)
|
||||
self._tool_policy = next((p for p in policies if p["type"] == "tool_auth"), None)
|
||||
self._active_patterns: Final = self._resolve_active_patterns()
|
||||
raw_rate: Final = self._cost_policy.get("cost_per_1k_tokens") if self._cost_policy else None
|
||||
self.cost_per_1k_tokens: Final[float] = (
|
||||
float(raw_rate) if isinstance(raw_rate, (int, float)) else DEFAULT_COST_PER_1K_TOKENS
|
||||
self.mode: Final = mode
|
||||
self._ledger: Final = ledger or BudgetLedger()
|
||||
self._pii: Final = next((p for p in policies if isinstance(p, PiiPolicy)), None)
|
||||
self._cost: Final = next((p for p in policies if isinstance(p, CostPolicy)), None)
|
||||
self._tool: Final = next((p for p in policies if isinstance(p, ToolAuthPolicy)), None)
|
||||
self._patterns: Final[Mapping[str, re.Pattern[str]]] = (
|
||||
resolve_scope(self._pii.patterns) if self._pii else _NO_PATTERNS
|
||||
)
|
||||
|
||||
def _expand_scope_entry(self, entry: object) -> tuple[str, ...]:
|
||||
if isinstance(entry, str) and entry in PATTERN_CATEGORIES:
|
||||
return tuple(PATTERN_CATEGORIES[entry])
|
||||
if isinstance(entry, str) and entry in self.patterns:
|
||||
return (entry,)
|
||||
raise ValueError(f"TealTiger: unknown pii pattern or category {entry!r}")
|
||||
|
||||
def _resolve_active_patterns(self) -> Mapping[str, re.Pattern[str]]:
|
||||
"""Honor the pii policy's `patterns` scope: "all", or a list of pattern
|
||||
names and/or category names. Unknown names fail at startup, not silently."""
|
||||
if not self._pii_policy:
|
||||
return MappingProxyType({})
|
||||
scope: Final = self._pii_policy.get("patterns", "all")
|
||||
if scope == "all":
|
||||
return self.patterns
|
||||
if isinstance(scope, str):
|
||||
return self._select_patterns((scope,))
|
||||
if isinstance(scope, (list, tuple)):
|
||||
return self._select_patterns(scope) # pyright: ignore[reportUnknownArgumentType] # isinstance on object yields list[Unknown]
|
||||
raise TypeError("TealTiger: pii policy `patterns` must be 'all' or a list of names")
|
||||
|
||||
def _select_patterns(self, entries: Sequence[object]) -> Mapping[str, re.Pattern[str]]:
|
||||
names: Final = tuple(chain.from_iterable(map(self._expand_scope_entry, entries)))
|
||||
return MappingProxyType({name: self.patterns[name] for name in names})
|
||||
|
||||
def evaluate_text(self, text: str) -> Decision:
|
||||
correlation_id: Final = str(uuid.uuid4())
|
||||
if not text or not self._pii_policy:
|
||||
return Decision(Action.ALLOW.value, "NO_VIOLATIONS", correlation_id)
|
||||
|
||||
findings: Final = tuple(name for name, pat in self._active_patterns.items() if pat.search(text))
|
||||
def check_text(self, text: str) -> Decision:
|
||||
if not text or self._pii is None:
|
||||
return self._decide(Action.ALLOW, NO_VIOLATIONS, RiskLevel.NONE)
|
||||
findings: Final = tuple(name for name, pattern in self._patterns.items() if pattern.search(text))
|
||||
if not findings:
|
||||
return Decision(Action.ALLOW.value, "NO_VIOLATIONS", correlation_id)
|
||||
return self._decide(Action.ALLOW, NO_VIOLATIONS, RiskLevel.NONE)
|
||||
risk: Final = RiskLevel.CRITICAL if any(f in CRITICAL_PATTERNS for f in findings) else RiskLevel.MEDIUM
|
||||
reason: Final = f"PII_DETECTED:{','.join(findings)}"
|
||||
details: Final[Mapping[str, object]] = {"types": findings}
|
||||
if self._pii.action == "BLOCK":
|
||||
return self._decide(Action.BLOCK, reason, risk, details=details)
|
||||
return self._decide(Action.REDACT, reason, risk, text=self._redact(text, findings), details=details)
|
||||
|
||||
pii_action: Final = self._pii_policy.get("action", "REDACT")
|
||||
if pii_action == "BLOCK" and self.mode == PolicyMode.ENFORCE:
|
||||
return Decision(
|
||||
Action.BLOCK.value,
|
||||
f"PII_DETECTED:{','.join(findings)}",
|
||||
correlation_id,
|
||||
details=MappingProxyType({"types": findings}),
|
||||
)
|
||||
|
||||
if self.mode == PolicyMode.MONITOR:
|
||||
return Decision(
|
||||
Action.ALLOW.value,
|
||||
f"PII_DETECTED:{','.join(findings)}",
|
||||
correlation_id,
|
||||
details=MappingProxyType({"types": findings}),
|
||||
)
|
||||
|
||||
redacted: Final = self._redact(text, findings)
|
||||
return Decision(
|
||||
Action.REDACT.value,
|
||||
f"PII_DETECTED:{','.join(findings)}",
|
||||
correlation_id,
|
||||
redacted_text=redacted,
|
||||
details=MappingProxyType({"types": findings}),
|
||||
def check_tool(self, tool_name: str) -> Decision:
|
||||
policy: Final = self._tool
|
||||
if policy is None:
|
||||
return self._decide(Action.ALLOW, NO_VIOLATIONS, RiskLevel.NONE)
|
||||
allowed: Final = tool_name not in policy.blocklist and (
|
||||
policy.allowlist is None or tool_name in policy.allowlist
|
||||
)
|
||||
if allowed:
|
||||
return self._decide(Action.ALLOW, NO_VIOLATIONS, RiskLevel.NONE)
|
||||
return self._decide(Action.BLOCK, "TOOL_NOT_ALLOWLISTED", RiskLevel.HIGH, details={"tool_name": tool_name})
|
||||
|
||||
def check_budget(self, identity: Identity, texts: Sequence[str]) -> Decision:
|
||||
policy: Final = self._cost
|
||||
if policy is None:
|
||||
return self._decide(Action.ALLOW, NO_VIOLATIONS, RiskLevel.NONE)
|
||||
daily, session = self._ledger.spent(identity.user, identity.session)
|
||||
candidates: Final = (
|
||||
("REQUEST_BUDGET_EXCEEDED", _estimate_cost_usd(policy, texts), policy.request_limit_usd),
|
||||
("SESSION_BUDGET_EXCEEDED", session, policy.session_limit_usd),
|
||||
("DAILY_BUDGET_EXCEEDED", daily, policy.daily_limit_usd),
|
||||
)
|
||||
exceeded: Final = next(
|
||||
((reason, spent, limit) for reason, spent, limit in candidates if limit is not None and spent >= limit),
|
||||
None,
|
||||
)
|
||||
if exceeded is None:
|
||||
return self._decide(Action.ALLOW, NO_VIOLATIONS, RiskLevel.NONE)
|
||||
reason, spent, limit = exceeded
|
||||
return self._decide(
|
||||
Action.BLOCK,
|
||||
reason,
|
||||
RiskLevel.HIGH,
|
||||
details={"spent_usd": spent, "limit_usd": limit, "user": identity.user},
|
||||
)
|
||||
|
||||
def record_usage(self, identity: Identity, texts: Sequence[str]) -> None:
|
||||
"""Accrue the estimated spend of text that passed through the guardrail."""
|
||||
if self._cost is not None:
|
||||
self._ledger.record(identity.user, identity.session, _estimate_cost_usd(self._cost, texts))
|
||||
|
||||
def _redact(self, text: str, findings: Sequence[str]) -> str:
|
||||
result = text # rebind-ok: progressively transformed across findings, one sub() pass per match
|
||||
for name in findings:
|
||||
result = self.patterns[name].sub(f"[REDACTED:{name.upper()}]", result)
|
||||
return result
|
||||
return reduce(
|
||||
lambda redacted, name: self._patterns[name].sub(f"[REDACTED:{name.upper()}]", redacted), findings, text
|
||||
)
|
||||
|
||||
def check_tool(self, tool_name: str) -> bool:
|
||||
if not self._tool_policy:
|
||||
return True
|
||||
raw_blocklist: Final = self._tool_policy.get("blocklist")
|
||||
blocklist: Final = raw_blocklist if isinstance(raw_blocklist, (list, tuple)) else () # pyright: ignore[reportUnknownVariableType] # isinstance on object yields list[Unknown]
|
||||
if tool_name in blocklist:
|
||||
return False
|
||||
raw_allowlist: Final = self._tool_policy.get("allowlist")
|
||||
if raw_allowlist is not None:
|
||||
allowlist: Final = raw_allowlist if isinstance(raw_allowlist, (list, tuple)) else () # pyright: ignore[reportUnknownVariableType] # isinstance on object yields list[Unknown]
|
||||
return tool_name in allowlist
|
||||
return True
|
||||
def _decide(
|
||||
self,
|
||||
action: Action,
|
||||
reason_code: str,
|
||||
risk_level: RiskLevel,
|
||||
text: str | None = None,
|
||||
details: Mapping[str, object] = _NO_DETAILS,
|
||||
) -> Decision:
|
||||
decision: Final = Decision(
|
||||
action, reason_code, risk_level, str(uuid.uuid4()), text, MappingProxyType(dict(details))
|
||||
)
|
||||
return self._apply_mode(decision)
|
||||
|
||||
def check_budget(self, session_id: str = "default") -> tuple[bool, float, float]:
|
||||
if not self._cost_policy:
|
||||
return False, 0.0, 0.0
|
||||
raw_limit: Final = self._cost_policy.get("daily_limit_usd")
|
||||
if not isinstance(raw_limit, (int, float)):
|
||||
return False, 0.0, 0.0
|
||||
limit: Final[float] = float(raw_limit)
|
||||
today: Final = date.today().isoformat() # noqa: DTZ011 # daily budget boundary is intentionally local-server-time, not UTC
|
||||
with self._lock:
|
||||
spent: Final = self._spend_by_day.get(today, 0.0)
|
||||
return spent >= limit, spent, limit
|
||||
|
||||
def track_cost(self, tokens: int, cost_per_1k_tokens: float | None = None) -> None:
|
||||
rate: Final = self.cost_per_1k_tokens if cost_per_1k_tokens is None else cost_per_1k_tokens
|
||||
cost: Final = (tokens / 1000.0) * rate
|
||||
today: Final = date.today().isoformat() # noqa: DTZ011 # daily budget boundary is intentionally local-server-time, not UTC
|
||||
with self._lock:
|
||||
self._spend_by_day[today] = self._spend_by_day.get(today, 0.0) + cost
|
||||
|
||||
def record_usage(self, texts: Sequence[str]) -> None:
|
||||
"""Accrue estimated spend for text that passed through the guardrail."""
|
||||
self.track_cost(tokens=estimate_tokens(texts))
|
||||
def _apply_mode(self, decision: Decision) -> Decision:
|
||||
"""MONITOR mode never blocks or rewrites: the decision is downgraded to ALLOW and the receipt records what
|
||||
ENFORCE mode would have done."""
|
||||
if self.mode is PolicyMode.ENFORCE or decision.action is Action.ALLOW:
|
||||
return decision
|
||||
return replace(
|
||||
decision,
|
||||
action=Action.ALLOW,
|
||||
redacted_text=None,
|
||||
details=MappingProxyType({**decision.details, "monitor_would": decision.action.value}),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -14,7 +14,8 @@ instead of "all", or override this module's PII_PATTERNS entirely.
|
|||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -47,7 +48,7 @@ PII_PATTERNS: Final = MappingProxyType(
|
|||
"ethereum_address": re.compile(r"\b0x[a-fA-F0-9]{40}\b"),
|
||||
# ---------- contact info ----------
|
||||
"email": re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b"),
|
||||
"phone_us": re.compile(r"\b(?:\+?1[-.\s]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b"),
|
||||
"phone_us": re.compile(r"(?<!\w)(?:\+?1[-.\s]?)?(?:\(\d{3}\)\s?|\d{3}[-.\s])\d{3}[-.\s]\d{4}\b"),
|
||||
"phone_intl": re.compile(r"\+\d{1,3}[-.\s]?\(?\d{1,4}\)?(?:[-.\s]?\d{2,4}){2,4}\b"),
|
||||
"us_zip_plus4": re.compile(r"\b\d{5}-\d{4}\b"),
|
||||
"po_box": re.compile(r"\bP\.?O\.?\s?Box\s?\d+\b", re.IGNORECASE),
|
||||
|
|
@ -76,9 +77,34 @@ PII_PATTERNS: Final = MappingProxyType(
|
|||
)
|
||||
|
||||
|
||||
def default_patterns() -> Mapping[str, re.Pattern[str]]:
|
||||
"""Return the built-in, immutable pattern mapping."""
|
||||
return PII_PATTERNS
|
||||
# Patterns that match ordinary text too often to scan by default (any 8-letter uppercase word looks like a
|
||||
# SWIFT code, any 10-digit number looks like an NPI or phone, and so on). They stay available when named
|
||||
# explicitly in a pii policy's `patterns` list.
|
||||
LOOSE_PATTERNS: Final = frozenset(
|
||||
{
|
||||
"ssn_no_dash",
|
||||
"ein",
|
||||
"us_passport",
|
||||
"passport_generic_intl",
|
||||
"drivers_license_generic",
|
||||
"npi_number",
|
||||
"date_of_birth",
|
||||
"vin",
|
||||
"credit_card_generic",
|
||||
"swift_bic",
|
||||
"us_bank_routing",
|
||||
"us_bank_account",
|
||||
"bitcoin_address",
|
||||
"us_zip_plus4",
|
||||
"lat_long_coordinates",
|
||||
"ipv4",
|
||||
"imei",
|
||||
}
|
||||
)
|
||||
|
||||
DEFAULT_PATTERNS: Final[Mapping[str, re.Pattern[str]]] = MappingProxyType(
|
||||
{name: pattern for name, pattern in PII_PATTERNS.items() if name not in LOOSE_PATTERNS}
|
||||
)
|
||||
|
||||
|
||||
PATTERN_CATEGORIES: Final = MappingProxyType(
|
||||
|
|
@ -142,3 +168,31 @@ PATTERN_CATEGORIES: Final = MappingProxyType(
|
|||
),
|
||||
}
|
||||
)
|
||||
|
||||
# Matches that expose a secret or a government/financial identifier get the highest risk level.
|
||||
CRITICAL_PATTERNS: Final = frozenset(
|
||||
{
|
||||
"ssn",
|
||||
"itin",
|
||||
"credit_card_visa",
|
||||
"credit_card_mastercard",
|
||||
"credit_card_amex",
|
||||
"credit_card_discover",
|
||||
*PATTERN_CATEGORIES["credentials"],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def resolve_scope(scope: str | Sequence[str]) -> Mapping[str, re.Pattern[str]]:
|
||||
"""Turn a pii policy's `patterns` value into the patterns to scan: "all" is the high-precision default set,
|
||||
anything else is exactly the pattern and category names listed (unknown names are rejected by parse_config)."""
|
||||
if isinstance(scope, str):
|
||||
return DEFAULT_PATTERNS
|
||||
names: Final = tuple(chain.from_iterable(PATTERN_CATEGORIES.get(entry, (entry,)) for entry in scope))
|
||||
return MappingProxyType({name: PII_PATTERNS[name] for name in names if name in PII_PATTERNS})
|
||||
|
||||
|
||||
def unknown_pattern_names(scope: str | Sequence[str]) -> tuple[str, ...]:
|
||||
if isinstance(scope, str):
|
||||
return ()
|
||||
return tuple(entry for entry in scope if entry not in PATTERN_CATEGORIES and entry not in PII_PATTERNS)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,61 @@
|
|||
"""Typed views over the untyped request shapes LiteLLM hands to a guardrail."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Identity:
|
||||
user: str
|
||||
session: str | None
|
||||
|
||||
|
||||
ANONYMOUS: Final = Identity(user="anonymous", session=None)
|
||||
|
||||
|
||||
class _FunctionName(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
name: str
|
||||
|
||||
|
||||
class _NamedTool(BaseModel):
|
||||
"""Both tool shapes in play: {"name": ...} / {"function": {"name": ...}} dicts and pydantic tool-call objects."""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
name: str | None = None
|
||||
function: _FunctionName | None = None
|
||||
|
||||
|
||||
class _Metadata(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
user_api_key_user_id: str | None = None
|
||||
|
||||
|
||||
class _RequestIdentity(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
user: str | None = None
|
||||
litellm_session_id: str | None = None
|
||||
metadata: _Metadata | None = None
|
||||
|
||||
|
||||
def tool_name_of(tool: object) -> str | None:
|
||||
try:
|
||||
parsed: Final = _NamedTool.model_validate(tool, from_attributes=True)
|
||||
except ValidationError:
|
||||
return None
|
||||
return parsed.name or (parsed.function.name if parsed.function else None)
|
||||
|
||||
|
||||
def identity_of(request_data: object) -> Identity:
|
||||
try:
|
||||
parsed: Final = _RequestIdentity.model_validate(request_data, from_attributes=True)
|
||||
except ValidationError:
|
||||
return ANONYMOUS
|
||||
user: Final = parsed.user or (parsed.metadata.user_api_key_user_id if parsed.metadata else None)
|
||||
return Identity(user=user or ANONYMOUS.user, session=parsed.litellm_session_id)
|
||||
|
|
@ -1,21 +1,19 @@
|
|||
"""
|
||||
TealTigerGuardrail — self-contained PII / cost / tool-authorization guardrail.
|
||||
TealTigerGuardrail: self-contained PII / cost / tool-authorization guardrail.
|
||||
|
||||
Implements apply_guardrail(), the current recommended pattern for guardrails
|
||||
that need to run uniformly across chat completions, /v1/messages, responses
|
||||
API, embeddings, etc. Verified against a live litellm proxy (v1.97.0) and
|
||||
against litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py as a reference
|
||||
implementation of the same interface.
|
||||
Implements apply_guardrail(), the recommended interface for guardrails that run uniformly across chat completions,
|
||||
/v1/messages, the responses API, embeddings and so on. No network calls, no LLM calls, no external API keys.
|
||||
|
||||
No network calls, no LLM calls, no external API keys required.
|
||||
Every governance check returns a Decision value. LiteLLM's guardrail contract blocks a call by raising, so the
|
||||
single raise lives at the edge of apply_guardrail.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
from fastapi import HTTPException
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
|
|
@ -24,16 +22,16 @@ from litellm.integrations.custom_guardrail import (
|
|||
from litellm.types.guardrails import GuardrailEventHooks, Mode
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
from .engine import Action, PolicyMode, TealEngine
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.tealtiger import DEFAULT_POLICIES, Policy, PolicyMode
|
||||
|
||||
from .decision import NO_VIOLATIONS, Action, Decision
|
||||
from .engine import TealEngine
|
||||
from .request import identity_of, tool_name_of
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
DEFAULT_POLICIES: Final = (
|
||||
MappingProxyType({"type": "pii", "action": "REDACT", "patterns": "all"}),
|
||||
MappingProxyType({"type": "cost", "action": "ENFORCE", "daily_limit_usd": 50.0}),
|
||||
MappingProxyType({"type": "tool_auth", "action": "ENFORCE", "allowlist": None}), # None = allow all
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
def _to_event_hook(
|
||||
|
|
@ -46,8 +44,7 @@ def _to_event_hook(
|
|||
) -> (
|
||||
GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None # mutable-ok: matches CustomGuardrail's event_hook
|
||||
):
|
||||
"""Narrow LitellmParams.mode's wider type down to what CustomGuardrail's
|
||||
own __init__ actually accepts, converting bare strings to the real enum."""
|
||||
"""Narrow LitellmParams.mode's wider type down to what CustomGuardrail's own __init__ accepts."""
|
||||
if mode is None or isinstance(mode, (GuardrailEventHooks, Mode)):
|
||||
return mode
|
||||
if isinstance(mode, str):
|
||||
|
|
@ -55,53 +52,30 @@ def _to_event_hook(
|
|||
return [item if isinstance(item, GuardrailEventHooks) else GuardrailEventHooks(item) for item in mode]
|
||||
|
||||
|
||||
def _tool_call_name_from_mapping(tool_call: Mapping[str, object]) -> str | None:
|
||||
"""Extract a tool call's function name from the TypedDict-shaped variant
|
||||
(ChatCompletionToolCallChunk)."""
|
||||
direct_name: Final = tool_call.get("name")
|
||||
if isinstance(direct_name, str):
|
||||
return direct_name
|
||||
function: Final = tool_call.get("function")
|
||||
if not isinstance(function, Mapping):
|
||||
return None
|
||||
nested_name: Final = function.get("name") # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # isinstance on object yields Mapping[Unknown, Unknown]
|
||||
return nested_name if isinstance(nested_name, str) else None
|
||||
def _log_receipt(decision: Decision) -> None:
|
||||
if decision.reason_code == NO_VIOLATIONS:
|
||||
verbose_proxy_logger.debug("tealtiger.decision_receipt=%s", decision.receipt_json())
|
||||
return
|
||||
verbose_proxy_logger.info("tealtiger.decision_receipt=%s", decision.receipt_json())
|
||||
|
||||
|
||||
def _tool_call_name_from_object(tool_call: object) -> str | None:
|
||||
"""Extract a tool call's function name from the pydantic-object variant
|
||||
(ChatCompletionMessageToolCall)."""
|
||||
direct_name: Final = getattr(tool_call, "name", None)
|
||||
if isinstance(direct_name, str):
|
||||
return direct_name
|
||||
function_obj: Final = getattr(tool_call, "function", None)
|
||||
nested_name: Final = getattr(function_obj, "name", None)
|
||||
return nested_name if isinstance(nested_name, str) else None
|
||||
|
||||
|
||||
def _tool_call_name(tool_call: object) -> str | None:
|
||||
"""Extract a tool call's function name across both real shapes this can
|
||||
take (see GenericGuardrailAPIInputs.tool_calls in litellm/types/utils.py):
|
||||
a plain dict/TypedDict (ChatCompletionToolCallChunk) or a pydantic object
|
||||
(ChatCompletionMessageToolCall). Avoids chained .get() calls, since one
|
||||
branch of that union doesn't reliably support dict-style access."""
|
||||
if isinstance(tool_call, Mapping):
|
||||
return _tool_call_name_from_mapping(tool_call) # pyright: ignore[reportUnknownArgumentType] # Mapping[Unknown, Unknown] from isinstance
|
||||
return _tool_call_name_from_object(tool_call)
|
||||
def _raise_if_blocked(decisions: Sequence[Decision]) -> None:
|
||||
blocked: Final = next((d for d in decisions if d.action is Action.BLOCK), None)
|
||||
if blocked is None:
|
||||
return
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"TealTiger: blocked: {blocked.reason_code} (correlation_id={blocked.correlation_id})",
|
||||
)
|
||||
|
||||
|
||||
class TealTigerGuardrail(CustomGuardrail):
|
||||
def __init__(
|
||||
self,
|
||||
policies: Sequence[Mapping[str, object]] | None = None,
|
||||
policy_mode: str = "ENFORCE",
|
||||
policies: Sequence[Policy] = DEFAULT_POLICIES,
|
||||
policy_mode: PolicyMode = PolicyMode.ENFORCE,
|
||||
guardrail_name: str | None = None,
|
||||
# Widened beyond CustomGuardrail's own event_hook type: this repo's
|
||||
# LitellmParams.mode field is str | list[str] | Mode (see
|
||||
# litellm/types/guardrails.py), so callers like this repo's own
|
||||
# __init__.py hand us that wider type directly. We narrow it to what
|
||||
# the base class actually accepts in _to_event_hook() below, rather
|
||||
# than requiring every caller to pre-convert it themselves.
|
||||
# LitellmParams.mode is str | list[str] | Mode, wider than CustomGuardrail's event_hook; narrowed below.
|
||||
event_hook: GuardrailEventHooks # mutable-ok: matches LitellmParams.mode / CustomGuardrail.event_hook
|
||||
| list[GuardrailEventHooks]
|
||||
| Mode
|
||||
|
|
@ -110,16 +84,19 @@ class TealTigerGuardrail(CustomGuardrail):
|
|||
| None = None,
|
||||
default_on: bool | None = False,
|
||||
) -> None:
|
||||
self.engine: Final = TealEngine(
|
||||
policies=policies or DEFAULT_POLICIES,
|
||||
mode=PolicyMode(policy_mode),
|
||||
)
|
||||
self.engine: Final = TealEngine(policies=policies, mode=policy_mode)
|
||||
super().__init__( # pyright: ignore[reportUnknownMemberType] # upstream CustomGuardrail.__init__ takes untyped **kwargs
|
||||
guardrail_name=guardrail_name,
|
||||
event_hook=_to_event_hook(event_hook),
|
||||
default_on=default_on or False,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> type[GuardrailConfigModel[BaseModel]] | None:
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.tealtiger import TealTigerGuardrailConfigModel
|
||||
|
||||
return TealTigerGuardrailConfigModel
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
|
|
@ -129,52 +106,35 @@ class TealTigerGuardrail(CustomGuardrail):
|
|||
logging_obj: LiteLLMLoggingObj | None = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
"""
|
||||
Runs on both request (pre-call) and response (post-call) content,
|
||||
since TealTiger scans for PII leakage in both directions.
|
||||
Runs on request and response content, since TealTiger scans for PII leakage in both directions.
|
||||
|
||||
Tool authorization covers tool definitions sent with the request and
|
||||
tool calls in both directions (a model response can invoke a tool the
|
||||
request never declared). The budget check only runs on the request path.
|
||||
In MONITOR mode violations are logged, never raised.
|
||||
Tool authorization covers tool definitions sent with the request and tool calls in both directions (a model
|
||||
response can invoke a tool the request never declared). The budget check only runs on the request path.
|
||||
In MONITOR mode violations are logged as receipts and never raised.
|
||||
|
||||
Spend is accrued per call from an estimate (about 4 characters per
|
||||
token) of the text that passes through, since apply_guardrail is not
|
||||
handed usage data.
|
||||
|
||||
Returns a new mapping rather than mutating `inputs` in place, since
|
||||
the caller's object should not be rewritten at a distance.
|
||||
Spend accrues per call from an estimate (about 4 characters per token) of the text passing through, since
|
||||
apply_guardrail is not handed usage data.
|
||||
"""
|
||||
self._check_tools(inputs, input_type)
|
||||
if input_type == "request":
|
||||
self._check_budget(request_data)
|
||||
identity: Final = identity_of(request_data)
|
||||
texts: Final = tuple(inputs.get("texts") or ())
|
||||
tool_decisions: Final = tuple(self.engine.check_tool(name) for name in self._tool_names(inputs, input_type))
|
||||
budget_decisions: Final = (self.engine.check_budget(identity, texts),) if input_type == "request" else ()
|
||||
text_decisions: Final = tuple(self.engine.check_text(text) for text in texts)
|
||||
|
||||
checked_texts: Final = tuple(self._check_text_or_raise(text) for text in inputs.get("texts") or ())
|
||||
self.engine.record_usage(checked_texts)
|
||||
return {**inputs, "texts": list(checked_texts)}
|
||||
decisions: Final = (*tool_decisions, *budget_decisions, *text_decisions)
|
||||
for decision in decisions:
|
||||
_log_receipt(decision)
|
||||
_raise_if_blocked(decisions)
|
||||
|
||||
def _violation(self, message: str) -> None:
|
||||
"""Raise in ENFORCE mode; only log in MONITOR mode."""
|
||||
if self.engine.mode == PolicyMode.ENFORCE:
|
||||
raise ValueError(f"TealTiger: blocked — {message}")
|
||||
verbose_proxy_logger.warning("TealTiger (MONITOR, not blocking): %s", message)
|
||||
self.engine.record_usage(identity, texts)
|
||||
redacted: Final = tuple(
|
||||
text if decision.redacted_text is None else decision.redacted_text
|
||||
for text, decision in zip(texts, text_decisions, strict=True)
|
||||
)
|
||||
return {**inputs, "texts": list(redacted)}
|
||||
|
||||
def _check_tools(self, inputs: GenericGuardrailAPIInputs, input_type: Literal["request", "response"]) -> None:
|
||||
@staticmethod
|
||||
def _tool_names(inputs: GenericGuardrailAPIInputs, input_type: Literal["request", "response"]) -> tuple[str, ...]:
|
||||
definitions: Final[Sequence[object]] = (inputs.get("tools") or ()) if input_type == "request" else ()
|
||||
calls: Final[Sequence[object]] = inputs.get("tool_calls") or ()
|
||||
for candidate in (*definitions, *calls):
|
||||
tool_name = _tool_call_name(candidate)
|
||||
if tool_name and not self.engine.check_tool(tool_name):
|
||||
self._violation(f"TOOL_NOT_ALLOWLISTED ({tool_name})")
|
||||
|
||||
def _check_budget(self, request_data: Mapping[str, object]) -> None:
|
||||
raw_user: Final = request_data.get("user")
|
||||
session_id: Final[str] = raw_user if isinstance(raw_user, str) else "default"
|
||||
over_budget, spent, limit = self.engine.check_budget(session_id=session_id)
|
||||
if over_budget:
|
||||
self._violation(f"DAILY_BUDGET_EXCEEDED (${spent:.2f} / ${limit:.2f})")
|
||||
|
||||
def _check_text_or_raise(self, text: str) -> str:
|
||||
decision: Final = self.engine.evaluate_text(text)
|
||||
if decision.action == Action.BLOCK.value:
|
||||
raise ValueError(f"TealTiger: blocked — {decision.reason_code}")
|
||||
return decision.redacted_text or text
|
||||
return tuple(name for name in map(tool_name_of, (*definitions, *calls)) if name is not None)
|
||||
|
|
|
|||
65
litellm/types/proxy/guardrails/guardrail_hooks/tealtiger.py
Normal file
65
litellm/types/proxy/guardrails/guardrail_hooks/tealtiger.py
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
from enum import Enum
|
||||
from typing import Annotated, Final, Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from .base import GuardrailConfigModel
|
||||
|
||||
|
||||
class PolicyMode(str, Enum):
|
||||
ENFORCE = "ENFORCE"
|
||||
MONITOR = "MONITOR"
|
||||
|
||||
|
||||
class PiiPolicy(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
type: Literal["pii"] = "pii"
|
||||
action: Literal["REDACT", "BLOCK"] = Field(default="REDACT", description="Redact matches, or block the request.")
|
||||
patterns: Literal["all"] | tuple[str, ...] = Field(
|
||||
default="all",
|
||||
description=(
|
||||
"'all' scans the high-precision default set. A list of pattern names and/or category names "
|
||||
"scans exactly those, including loose patterns such as ssn_no_dash or swift_bic."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class CostPolicy(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
type: Literal["cost"] = "cost"
|
||||
daily_limit_usd: float | None = Field(default=None, gt=0, description="Per-user daily spend limit.")
|
||||
session_limit_usd: float | None = Field(default=None, gt=0, description="Per-session spend limit.")
|
||||
request_limit_usd: float | None = Field(default=None, gt=0, description="Estimated cost limit for one request.")
|
||||
cost_per_1k_tokens: float = Field(default=0.002, ge=0, description="Rate used to turn estimated tokens into USD.")
|
||||
|
||||
|
||||
class ToolAuthPolicy(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
type: Literal["tool_auth"] = "tool_auth"
|
||||
allowlist: tuple[str, ...] | None = Field(default=None, description="Only these tools may be used.")
|
||||
blocklist: tuple[str, ...] = Field(default=(), description="These tools are always blocked.")
|
||||
|
||||
|
||||
Policy: TypeAlias = Annotated[PiiPolicy | CostPolicy | ToolAuthPolicy, Field(discriminator="type")]
|
||||
|
||||
DEFAULT_POLICIES: Final[tuple[Policy, ...]] = (PiiPolicy(),)
|
||||
|
||||
|
||||
class TealTigerGuardrailConfigModel(GuardrailConfigModel[BaseModel]):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
policies: tuple[Policy, ...] | None = Field(
|
||||
default=None,
|
||||
description="Policies to enforce: pii, cost and tool_auth. Defaults to PII redaction only.",
|
||||
)
|
||||
policy_mode: PolicyMode | None = Field(
|
||||
default=None,
|
||||
description="ENFORCE blocks or redacts. MONITOR only logs decision receipts. Defaults to ENFORCE.",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "TealTiger"
|
||||
|
|
@ -1,454 +1,530 @@
|
|||
"""
|
||||
Unit tests for the TealTiger guardrail.
|
||||
"""Unit tests for the TealTiger guardrail. Nothing on the other side of the call: no network, no API keys."""
|
||||
|
||||
Run with: pytest tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py -v
|
||||
|
||||
These import from the vendored litellm/proxy/guardrails/guardrail_hooks/
|
||||
tealtiger/ package. Placed inside an actual litellm checkout, no path
|
||||
changes should be needed since the tests import via the standard package path.
|
||||
"""
|
||||
import json
|
||||
import re
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.proxy.guardrails.guardrail_hooks.tealtiger.engine import (
|
||||
Action,
|
||||
PolicyMode,
|
||||
TealEngine,
|
||||
estimate_tokens,
|
||||
from litellm.proxy.guardrails.guardrail_hooks.tealtiger import budget as budget_module
|
||||
from litellm.proxy.guardrails.guardrail_hooks.tealtiger.budget import BudgetLedger
|
||||
from litellm.proxy.guardrails.guardrail_hooks.tealtiger.config import parse_config
|
||||
from litellm.proxy.guardrails.guardrail_hooks.tealtiger.decision import Action, RiskLevel
|
||||
from litellm.proxy.guardrails.guardrail_hooks.tealtiger.engine import TealEngine, estimate_tokens
|
||||
from litellm.proxy.guardrails.guardrail_hooks.tealtiger.request import (
|
||||
ANONYMOUS,
|
||||
Identity,
|
||||
identity_of,
|
||||
tool_name_of,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_hooks.tealtiger.tealtiger import (
|
||||
from litellm.proxy.guardrails.guardrail_hooks.tealtiger.tealtiger import TealTigerGuardrail, _to_event_hook
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.tealtiger import (
|
||||
DEFAULT_POLICIES,
|
||||
TealTigerGuardrail,
|
||||
CostPolicy,
|
||||
PiiPolicy,
|
||||
PolicyMode,
|
||||
TealTigerGuardrailConfigModel,
|
||||
ToolAuthPolicy,
|
||||
)
|
||||
|
||||
_BLOCKED_TOOL = re.compile(r"TealTiger: blocked: TOOL_NOT_ALLOWLISTED \(correlation_id=[0-9a-f-]{36}\)")
|
||||
_ALICE = Identity(user="alice", session="s1")
|
||||
|
||||
def make_engine(mode=PolicyMode.ENFORCE):
|
||||
policies = [
|
||||
{"type": "pii", "action": "REDACT", "patterns": "all"},
|
||||
{"type": "cost", "action": "ENFORCE", "daily_limit_usd": 1.0},
|
||||
{"type": "tool_auth", "action": "ENFORCE", "allowlist": ["search", "calculator"]},
|
||||
]
|
||||
|
||||
def _engine(*policies, mode: PolicyMode = PolicyMode.ENFORCE) -> TealEngine:
|
||||
return TealEngine(policies=policies, mode=mode)
|
||||
|
||||
|
||||
def test_engine_allows_clean_text():
|
||||
engine = make_engine()
|
||||
decision = engine.evaluate_text("What's the weather like today?")
|
||||
assert decision.action == Action.ALLOW.value
|
||||
# ---------- PII ----------
|
||||
|
||||
|
||||
def test_engine_redacts_ssn():
|
||||
engine = make_engine()
|
||||
decision = engine.evaluate_text("My SSN is 123-45-6789")
|
||||
assert decision.action == Action.REDACT.value
|
||||
assert "123-45-6789" not in decision.redacted_text
|
||||
def test_default_policies_redact_an_ssn_and_mark_it_critical():
|
||||
decision = _engine(*DEFAULT_POLICIES).check_text("My SSN is 123-45-6789")
|
||||
|
||||
|
||||
def test_engine_monitor_mode_never_redacts():
|
||||
engine = make_engine(mode=PolicyMode.MONITOR)
|
||||
decision = engine.evaluate_text("My SSN is 123-45-6789")
|
||||
assert decision.action == Action.ALLOW.value
|
||||
|
||||
|
||||
def test_engine_tool_allowlist():
|
||||
engine = make_engine()
|
||||
assert engine.check_tool("calculator") is True
|
||||
assert engine.check_tool("shell_exec") is False
|
||||
|
||||
|
||||
def test_engine_budget_enforcement():
|
||||
engine = make_engine()
|
||||
engine.track_cost(tokens=1_000_000) # pushes spend over $1 limit
|
||||
over_budget, spent, limit = engine.check_budget()
|
||||
assert over_budget is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_redacts_request_text():
|
||||
guardrail = TealTigerGuardrail(policies=DEFAULT_POLICIES, policy_mode="ENFORCE")
|
||||
inputs = {"texts": ["Contact me at jane@example.com"], "images": [], "tool_calls": []}
|
||||
|
||||
result = await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data={"model": "gpt-4", "user": "test-user"},
|
||||
input_type="request",
|
||||
assert (decision.action, decision.reason_code, decision.risk_level, decision.redacted_text) == (
|
||||
Action.REDACT,
|
||||
"PII_DETECTED:ssn",
|
||||
RiskLevel.CRITICAL,
|
||||
"My SSN is [REDACTED:SSN]",
|
||||
)
|
||||
assert "jane@example.com" not in result["texts"][0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_blocks_disallowed_tool():
|
||||
policies = [
|
||||
{"type": "pii", "action": "REDACT", "patterns": "all"},
|
||||
{"type": "tool_auth", "action": "ENFORCE", "allowlist": ["search"]},
|
||||
]
|
||||
guardrail = TealTigerGuardrail(policies=policies, policy_mode="ENFORCE")
|
||||
inputs = {
|
||||
"texts": ["run it"],
|
||||
"images": [],
|
||||
"tool_calls": [{"name": "shell_exec"}],
|
||||
}
|
||||
def test_email_is_redacted_at_medium_risk():
|
||||
decision = _engine(*DEFAULT_POLICIES).check_text("mail jane@example.com")
|
||||
|
||||
with pytest.raises(ValueError, match="TOOL_NOT_ALLOWLISTED"):
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data={"model": "gpt-4"},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
|
||||
def test_engine_malformed_blocklist_fails_safe():
|
||||
"""A non-list/tuple blocklist value shouldn't crash — treated as empty."""
|
||||
policies = [
|
||||
{"type": "tool_auth", "action": "ENFORCE", "blocklist": "not-a-list"},
|
||||
]
|
||||
engine = TealEngine(policies=policies, mode=PolicyMode.ENFORCE)
|
||||
# malformed blocklist -> treated as empty -> nothing gets blocked via it
|
||||
assert engine.check_tool("anything") is True
|
||||
|
||||
|
||||
def test_engine_malformed_allowlist_denies_by_default():
|
||||
"""A non-list/tuple allowlist value shouldn't crash — fails closed (deny)."""
|
||||
policies = [
|
||||
{"type": "tool_auth", "action": "ENFORCE", "allowlist": "not-a-list"},
|
||||
]
|
||||
engine = TealEngine(policies=policies, mode=PolicyMode.ENFORCE)
|
||||
# malformed allowlist -> treated as empty -> nothing is in it -> denied
|
||||
assert engine.check_tool("anything") is False
|
||||
|
||||
|
||||
def test_engine_malformed_budget_limit_treated_as_no_limit():
|
||||
"""A non-numeric daily_limit_usd shouldn't crash — treated as no budget cap."""
|
||||
policies = [
|
||||
{"type": "cost", "action": "ENFORCE", "daily_limit_usd": "fifty dollars"},
|
||||
]
|
||||
engine = TealEngine(policies=policies, mode=PolicyMode.ENFORCE)
|
||||
engine.track_cost(tokens=999_999_999) # would blow any real limit
|
||||
over_budget, spent, limit = engine.check_budget()
|
||||
assert over_budget is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_session_id_falls_back_when_user_not_a_string():
|
||||
"""A non-string 'user' in request_data shouldn't crash session_id handling."""
|
||||
policies = [
|
||||
{"type": "cost", "action": "ENFORCE", "daily_limit_usd": 1.0},
|
||||
]
|
||||
guardrail = TealTigerGuardrail(policies=policies, policy_mode="ENFORCE")
|
||||
inputs = {"texts": ["hello"], "images": [], "tool_calls": []}
|
||||
# user is an int here, not a str -- should fall back to "default" cleanly
|
||||
result = await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data={"model": "gpt-4", "user": 12345},
|
||||
input_type="request",
|
||||
assert (decision.reason_code, decision.risk_level, decision.redacted_text) == (
|
||||
"PII_DETECTED:email",
|
||||
RiskLevel.MEDIUM,
|
||||
"mail [REDACTED:EMAIL]",
|
||||
)
|
||||
assert result["texts"] == ["hello"]
|
||||
|
||||
|
||||
# ---------- coverage for _to_event_hook and _tool_call_name helpers ----------
|
||||
def test_clean_text_is_allowed_without_a_redaction():
|
||||
decision = _engine(*DEFAULT_POLICIES).check_text("What's the weather like today?")
|
||||
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
assert (decision.action, decision.reason_code, decision.risk_level, decision.redacted_text) == (
|
||||
Action.ALLOW,
|
||||
"NO_VIOLATIONS",
|
||||
RiskLevel.NONE,
|
||||
None,
|
||||
)
|
||||
|
||||
from litellm.proxy.guardrails.guardrail_hooks.tealtiger.tealtiger import (
|
||||
_to_event_hook,
|
||||
_tool_call_name,
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text",
|
||||
["Enter your PASSWORD", "Order 1234567890 shipped", "upgrade to 1.2.3.4", "due 12/25/2024", "INTERNATIONAL"],
|
||||
)
|
||||
def test_ordinary_text_is_not_flagged_by_the_default_scope(text: str):
|
||||
decision = _engine(*DEFAULT_POLICIES).check_text(text)
|
||||
|
||||
assert (decision.action, decision.redacted_text) == (Action.ALLOW, None), f"flagged {text!r}: {decision}"
|
||||
|
||||
|
||||
def test_to_event_hook_passes_through_none():
|
||||
assert _to_event_hook(None) is None
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("call (555) 123-4567 now", "call [REDACTED:PHONE_US] now"),
|
||||
("call 555-123-4567 now", "call [REDACTED:PHONE_US] now"),
|
||||
("card 4111 1111 1111 1111", "card [REDACTED:CREDIT_CARD_VISA]"),
|
||||
("key AKIAIOSFODNN7EXAMPLE", "key [REDACTED:AWS_ACCESS_KEY]"),
|
||||
],
|
||||
)
|
||||
def test_real_pii_is_still_redacted_by_the_default_scope(text: str, expected: str):
|
||||
assert _engine(*DEFAULT_POLICIES).check_text(text).redacted_text == expected
|
||||
|
||||
|
||||
def test_to_event_hook_passes_through_real_enum():
|
||||
result = _to_event_hook(GuardrailEventHooks.pre_call)
|
||||
assert result == GuardrailEventHooks.pre_call
|
||||
def test_block_action_blocks_instead_of_redacting():
|
||||
decision = _engine(PiiPolicy(action="BLOCK")).check_text("My SSN is 123-45-6789")
|
||||
|
||||
assert (decision.action, decision.reason_code, decision.redacted_text) == (Action.BLOCK, "PII_DETECTED:ssn", None)
|
||||
|
||||
|
||||
def test_to_event_hook_converts_bare_string():
|
||||
result = _to_event_hook("pre_call")
|
||||
assert result == GuardrailEventHooks.pre_call
|
||||
def test_pattern_list_scans_only_the_listed_patterns():
|
||||
engine = _engine(PiiPolicy(patterns=("email",)))
|
||||
|
||||
assert engine.check_text("SSN 123-45-6789").action == Action.ALLOW
|
||||
assert engine.check_text("mail jane@example.com").redacted_text == "mail [REDACTED:EMAIL]"
|
||||
|
||||
|
||||
def test_to_event_hook_converts_list_of_strings():
|
||||
result = _to_event_hook(["pre_call", "post_call"])
|
||||
assert result == [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call]
|
||||
def test_category_name_scans_every_pattern_in_that_category():
|
||||
engine = _engine(PiiPolicy(patterns=("credentials",)))
|
||||
|
||||
assert engine.check_text("mail jane@example.com").action == Action.ALLOW
|
||||
assert engine.check_text("key AKIAIOSFODNN7EXAMPLE").redacted_text == "key [REDACTED:AWS_ACCESS_KEY]"
|
||||
|
||||
|
||||
def test_to_event_hook_converts_mixed_list():
|
||||
result = _to_event_hook([GuardrailEventHooks.pre_call, "post_call"])
|
||||
assert result == [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call]
|
||||
def test_loose_pattern_is_available_when_named_explicitly():
|
||||
decision = _engine(PiiPolicy(patterns=("swift_bic",))).check_text("Enter your PASSWORD")
|
||||
|
||||
assert decision.redacted_text == "Enter your [REDACTED:SWIFT_BIC]"
|
||||
|
||||
|
||||
def test_tool_call_name_from_dict_shape():
|
||||
"""ChatCompletionToolCallChunk-style: a plain dict."""
|
||||
tool_call = {"id": "1", "type": "function", "function": {"name": "shell_exec"}, "index": 0}
|
||||
assert _tool_call_name(tool_call) == "shell_exec"
|
||||
def test_monitor_mode_keeps_the_text_and_records_what_enforce_would_have_done():
|
||||
decision = _engine(*DEFAULT_POLICIES, mode=PolicyMode.MONITOR).check_text("My SSN is 123-45-6789")
|
||||
|
||||
|
||||
def test_tool_call_name_from_dict_shape_missing_function():
|
||||
tool_call = {"id": "1", "type": "function", "index": 0}
|
||||
assert _tool_call_name(tool_call) is None
|
||||
|
||||
|
||||
class _FakeFunction:
|
||||
def __init__(self, name):
|
||||
self.name = name
|
||||
|
||||
|
||||
class _FakeToolCallObject:
|
||||
"""Mimics ChatCompletionMessageToolCall's real shape: attribute access."""
|
||||
|
||||
def __init__(self, name):
|
||||
self.function = _FakeFunction(name)
|
||||
self.id = "1"
|
||||
self.type = "function"
|
||||
|
||||
|
||||
def test_tool_call_name_from_object_shape():
|
||||
"""ChatCompletionMessageToolCall-style: a pydantic-like object."""
|
||||
assert _tool_call_name(_FakeToolCallObject("calculator")) == "calculator"
|
||||
|
||||
|
||||
def test_tool_call_name_from_object_shape_missing_function():
|
||||
class _Empty:
|
||||
pass
|
||||
|
||||
assert _tool_call_name(_Empty()) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_blocks_object_shaped_tool_call():
|
||||
"""The object-shaped (pydantic) tool_call variant should be checked too,
|
||||
not just the dict-shaped one."""
|
||||
policies = [
|
||||
{"type": "pii", "action": "REDACT", "patterns": "all"},
|
||||
{"type": "tool_auth", "action": "ENFORCE", "allowlist": ["search"]},
|
||||
]
|
||||
guardrail = TealTigerGuardrail(policies=policies, policy_mode="ENFORCE")
|
||||
inputs = {
|
||||
"texts": ["run it"],
|
||||
"images": [],
|
||||
"tool_calls": [_FakeToolCallObject("shell_exec")],
|
||||
}
|
||||
with pytest.raises(ValueError, match="TOOL_NOT_ALLOWLISTED"):
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data={"model": "gpt-4"},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
|
||||
# ---------- additional edge-case coverage ----------
|
||||
|
||||
|
||||
def test_engine_no_pii_match_leaves_redacted_text_none():
|
||||
engine = make_engine()
|
||||
decision = engine.evaluate_text("Contact me at clean_user_handle")
|
||||
assert decision.action == Action.ALLOW.value
|
||||
assert decision.redacted_text is None
|
||||
|
||||
|
||||
def test_engine_tool_auth_blocklist_only():
|
||||
policies = [
|
||||
{"type": "tool_auth", "action": "ENFORCE", "blocklist": ["dangerous_exec", "terminal"]},
|
||||
]
|
||||
engine = TealEngine(policies=policies, mode=PolicyMode.ENFORCE)
|
||||
assert engine.check_tool("calculator") is True
|
||||
assert engine.check_tool("dangerous_exec") is False
|
||||
|
||||
|
||||
def test_engine_tool_auth_no_allowlist_or_blocklist_allows_everything():
|
||||
policies = [
|
||||
{"type": "tool_auth", "action": "ENFORCE"},
|
||||
]
|
||||
engine = TealEngine(policies=policies, mode=PolicyMode.ENFORCE)
|
||||
assert engine.check_tool("any_tool") is True
|
||||
|
||||
|
||||
def test_engine_budget_tracking_stays_under_limit():
|
||||
policies = [
|
||||
{"type": "cost", "action": "ENFORCE", "daily_limit_usd": 10.0},
|
||||
]
|
||||
engine = TealEngine(policies=policies, mode=PolicyMode.ENFORCE)
|
||||
engine.track_cost(tokens=10)
|
||||
over_budget, spent, limit = engine.check_budget()
|
||||
assert over_budget is False
|
||||
assert spent > 0
|
||||
assert limit == 10.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_budget_exceeded_raises_error():
|
||||
policies = [
|
||||
{"type": "cost", "action": "ENFORCE", "daily_limit_usd": 0.00001},
|
||||
]
|
||||
guardrail = TealTigerGuardrail(policies=policies, policy_mode="ENFORCE")
|
||||
guardrail.engine.track_cost(tokens=1_000_000)
|
||||
|
||||
inputs = {"texts": ["hello"], "images": [], "tool_calls": []}
|
||||
|
||||
with pytest.raises(ValueError, match="BUDGET_EXCEEDED"):
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data={"model": "gpt-4", "user": "test-user"},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_redacts_response_text():
|
||||
guardrail = TealTigerGuardrail(policies=DEFAULT_POLICIES, policy_mode="ENFORCE")
|
||||
inputs = {"texts": ["Found secret token: 123-45-6789"], "images": [], "tool_calls": []}
|
||||
|
||||
result = await guardrail.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data={"model": "gpt-4"},
|
||||
input_type="response",
|
||||
assert (decision.action, decision.reason_code, decision.redacted_text, dict(decision.details)) == (
|
||||
Action.ALLOW,
|
||||
"PII_DETECTED:ssn",
|
||||
None,
|
||||
{"types": ("ssn",), "monitor_would": "REDACT"},
|
||||
)
|
||||
assert "123-45-6789" not in result["texts"][0]
|
||||
|
||||
|
||||
def test_tool_call_name_handles_invalid_types():
|
||||
assert _tool_call_name(None) is None
|
||||
assert _tool_call_name("invalid_str") is None
|
||||
assert _tool_call_name(12345) is None
|
||||
def test_receipt_is_typed_json_with_risk_level_reason_and_correlation_id():
|
||||
decision = _engine(*DEFAULT_POLICIES).check_text("My SSN is 123-45-6789")
|
||||
|
||||
assert json.loads(decision.receipt_json()) == {
|
||||
"correlation_id": decision.correlation_id,
|
||||
"action": "REDACT",
|
||||
"reason_code": "PII_DETECTED:ssn",
|
||||
"risk_level": "CRITICAL",
|
||||
"details": {"types": ["ssn"]},
|
||||
}
|
||||
|
||||
|
||||
_TOOL_POLICIES = [
|
||||
{"type": "tool_auth", "action": "ENFORCE", "allowlist": ["search"]},
|
||||
]
|
||||
# ---------- tools ----------
|
||||
|
||||
|
||||
def test_tool_allowlist_allows_listed_and_blocks_unlisted_tools():
|
||||
engine = _engine(ToolAuthPolicy(allowlist=("search",)))
|
||||
|
||||
assert engine.check_tool("search").action == Action.ALLOW
|
||||
blocked = engine.check_tool("shell_exec")
|
||||
assert (blocked.action, blocked.reason_code, blocked.risk_level, dict(blocked.details)) == (
|
||||
Action.BLOCK,
|
||||
"TOOL_NOT_ALLOWLISTED",
|
||||
RiskLevel.HIGH,
|
||||
{"tool_name": "shell_exec"},
|
||||
)
|
||||
|
||||
|
||||
def test_tool_blocklist_wins_over_a_missing_allowlist():
|
||||
engine = _engine(ToolAuthPolicy(blocklist=("shell_exec",)))
|
||||
|
||||
assert [engine.check_tool(name).action for name in ("search", "shell_exec")] == [Action.ALLOW, Action.BLOCK]
|
||||
|
||||
|
||||
def test_no_tool_policy_allows_every_tool():
|
||||
assert _engine(*DEFAULT_POLICIES).check_tool("anything").action == Action.ALLOW
|
||||
|
||||
|
||||
# ---------- budgets ----------
|
||||
|
||||
|
||||
def _cost_engine(**limits: float) -> TealEngine:
|
||||
return _engine(CostPolicy(cost_per_1k_tokens=1.0, **limits))
|
||||
|
||||
|
||||
def test_daily_budget_is_tracked_per_user():
|
||||
engine = _cost_engine(daily_limit_usd=1.0)
|
||||
engine.record_usage(_ALICE, ["x" * 4000]) # 1000 estimated tokens at $1 per 1k
|
||||
|
||||
alice = engine.check_budget(_ALICE, ["hi"])
|
||||
bob = engine.check_budget(Identity(user="bob", session=None), ["hi"])
|
||||
|
||||
assert (alice.action, alice.reason_code, bob.action) == (Action.BLOCK, "DAILY_BUDGET_EXCEEDED", Action.ALLOW)
|
||||
assert dict(alice.details) == {"spent_usd": 1.0, "limit_usd": 1.0, "user": "alice"}
|
||||
|
||||
|
||||
def test_session_budget_is_tracked_per_session():
|
||||
engine = _cost_engine(session_limit_usd=1.0)
|
||||
engine.record_usage(_ALICE, ["x" * 4000])
|
||||
|
||||
same_session = engine.check_budget(Identity(user="bob", session="s1"), ["hi"])
|
||||
other_session = engine.check_budget(Identity(user="alice", session="s2"), ["hi"])
|
||||
|
||||
assert (same_session.reason_code, other_session.action) == ("SESSION_BUDGET_EXCEEDED", Action.ALLOW)
|
||||
|
||||
|
||||
def test_request_budget_blocks_one_oversized_request():
|
||||
engine = _cost_engine(request_limit_usd=1.0)
|
||||
|
||||
small = engine.check_budget(_ALICE, ["x" * 3996]) # 999 tokens: $0.999
|
||||
large = engine.check_budget(_ALICE, ["x" * 4000]) # 1000 tokens: $1.00
|
||||
|
||||
assert (small.action, large.reason_code) == (Action.ALLOW, "REQUEST_BUDGET_EXCEEDED")
|
||||
|
||||
|
||||
def test_no_cost_policy_means_no_budget():
|
||||
engine = _engine(*DEFAULT_POLICIES)
|
||||
engine.record_usage(_ALICE, ["x" * 4_000_000])
|
||||
|
||||
assert engine.check_budget(_ALICE, ["hi"]).action == Action.ALLOW
|
||||
|
||||
|
||||
def test_monitor_mode_allows_an_exhausted_budget_and_records_it():
|
||||
engine = _engine(CostPolicy(daily_limit_usd=1.0, cost_per_1k_tokens=1.0), mode=PolicyMode.MONITOR)
|
||||
engine.record_usage(_ALICE, ["x" * 4000])
|
||||
|
||||
decision = engine.check_budget(_ALICE, ["hi"])
|
||||
|
||||
assert (decision.action, decision.reason_code, decision.details["monitor_would"]) == (
|
||||
Action.ALLOW,
|
||||
"DAILY_BUDGET_EXCEEDED",
|
||||
"BLOCK",
|
||||
)
|
||||
|
||||
|
||||
def test_ledger_resets_when_the_local_day_rolls_over():
|
||||
ledger = BudgetLedger()
|
||||
ledger.record("alice", "s1", 2.0)
|
||||
|
||||
with patch.object(budget_module, "_today", return_value="2999-01-01"):
|
||||
assert ledger.spent("alice", "s1") == (0.0, 0.0)
|
||||
|
||||
|
||||
def test_ledger_accumulates_within_a_day():
|
||||
ledger = BudgetLedger()
|
||||
ledger.record("alice", "s1", 0.25)
|
||||
ledger.record("alice", "s1", 0.5)
|
||||
|
||||
assert ledger.spent("alice", "s1") == (0.75, 0.75)
|
||||
|
||||
|
||||
def test_estimate_tokens_is_about_four_characters_per_token():
|
||||
assert [estimate_tokens(texts) for texts in ([], ["abcd"], ["abcde", "fg"])] == [0, 1, 2]
|
||||
|
||||
|
||||
# ---------- config ----------
|
||||
|
||||
|
||||
def test_config_defaults_to_pii_redaction_in_enforce_mode_with_no_budget():
|
||||
config = parse_config(None, None)
|
||||
|
||||
assert (config.policies, config.policy_mode) == ((PiiPolicy(),), PolicyMode.ENFORCE)
|
||||
|
||||
|
||||
def test_config_parses_a_yaml_style_policy_list_into_typed_policies():
|
||||
config = parse_config(
|
||||
[
|
||||
{"type": "pii", "action": "BLOCK", "patterns": ["email"]},
|
||||
{"type": "cost", "daily_limit_usd": 5},
|
||||
{"type": "tool_auth", "allowlist": ["search"]},
|
||||
],
|
||||
"MONITOR",
|
||||
)
|
||||
|
||||
assert config.policies == (
|
||||
PiiPolicy(action="BLOCK", patterns=("email",)),
|
||||
CostPolicy(daily_limit_usd=5.0),
|
||||
ToolAuthPolicy(allowlist=("search",)),
|
||||
)
|
||||
assert config.policy_mode == PolicyMode.MONITOR
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"policies",
|
||||
[
|
||||
[{"type": "nope"}],
|
||||
[{"type": "cost", "daily_limit": 5}],
|
||||
[{"type": "cost", "daily_limit_usd": -1}],
|
||||
[{"type": "pii", "patterns": "email"}],
|
||||
],
|
||||
)
|
||||
def test_config_rejects_invalid_policies_at_startup(policies: list[dict[str, object]]):
|
||||
with pytest.raises(ValidationError):
|
||||
parse_config(policies, None)
|
||||
|
||||
|
||||
def test_config_rejects_an_unknown_pattern_name_at_startup():
|
||||
with pytest.raises(ValueError, match="unknown pii pattern or category name"):
|
||||
parse_config([{"type": "pii", "patterns": ["email", "not_a_pattern"]}], None)
|
||||
|
||||
|
||||
def test_config_rejects_an_unknown_policy_mode():
|
||||
with pytest.raises(ValidationError):
|
||||
parse_config(None, "WATCH")
|
||||
|
||||
|
||||
# ---------- request shapes ----------
|
||||
|
||||
|
||||
class _Function:
|
||||
name = "search"
|
||||
|
||||
|
||||
class _ToolCallObject:
|
||||
function = _Function()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tool", "expected"),
|
||||
[
|
||||
({"name": "search"}, "search"),
|
||||
({"type": "function", "function": {"name": "search"}}, "search"),
|
||||
(_ToolCallObject(), "search"),
|
||||
({"function": "not-a-mapping"}, None),
|
||||
({}, None),
|
||||
(None, None),
|
||||
("invalid", None),
|
||||
(12345, None),
|
||||
],
|
||||
)
|
||||
def test_tool_name_of_reads_every_tool_shape(tool: object, expected: str | None):
|
||||
assert tool_name_of(tool) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("request_data", "expected"),
|
||||
[
|
||||
({"user": "alice", "litellm_session_id": "s1"}, Identity(user="alice", session="s1")),
|
||||
({"metadata": {"user_api_key_user_id": "key-user"}}, Identity(user="key-user", session=None)),
|
||||
({"user": "alice", "metadata": {"user_api_key_user_id": "key-user"}}, Identity(user="alice", session=None)),
|
||||
({}, ANONYMOUS),
|
||||
({"user": 42}, ANONYMOUS),
|
||||
],
|
||||
)
|
||||
def test_identity_of_prefers_the_request_user_and_falls_back_to_anonymous(request_data: dict, expected: Identity):
|
||||
assert identity_of(request_data) == expected
|
||||
|
||||
|
||||
# ---------- guardrail ----------
|
||||
|
||||
|
||||
def _guardrail(*policies, mode: PolicyMode = PolicyMode.ENFORCE) -> TealTigerGuardrail:
|
||||
return TealTigerGuardrail(policies=policies or DEFAULT_POLICIES, policy_mode=mode)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_blocks_disallowed_tool_definition_openai_shape():
|
||||
guardrail = TealTigerGuardrail(policies=_TOOL_POLICIES, policy_mode="ENFORCE")
|
||||
inputs = {"texts": ["hi"], "tools": [{"type": "function", "function": {"name": "shell_exec"}}]}
|
||||
async def test_request_text_is_redacted_and_other_inputs_are_preserved():
|
||||
inputs = {"texts": ["Contact me at jane@example.com"], "images": ["img"], "model": "gpt-4"}
|
||||
|
||||
with pytest.raises(ValueError, match="TOOL_NOT_ALLOWLISTED"):
|
||||
result = await _guardrail().apply_guardrail(inputs=inputs, request_data={}, input_type="request")
|
||||
|
||||
assert result == {"texts": ["Contact me at [REDACTED:EMAIL]"], "images": ["img"], "model": "gpt-4"}
|
||||
assert inputs["texts"] == ["Contact me at jane@example.com"], "the caller's inputs must not be mutated"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_text_is_redacted_too():
|
||||
result = await _guardrail().apply_guardrail(
|
||||
inputs={"texts": ["token: 123-45-6789"]}, request_data={}, input_type="response"
|
||||
)
|
||||
|
||||
assert result == {"texts": ["token: [REDACTED:SSN]"]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_block_policy_raises_a_400_with_the_reason_and_correlation_id():
|
||||
guardrail = _guardrail(PiiPolicy(action="BLOCK"))
|
||||
|
||||
with pytest.raises(HTTPException) as blocked:
|
||||
await guardrail.apply_guardrail(inputs={"texts": ["SSN 123-45-6789"]}, request_data={}, input_type="request")
|
||||
|
||||
assert blocked.value.status_code == 400
|
||||
assert re.fullmatch(r"TealTiger: blocked: PII_DETECTED:ssn \(correlation_id=[0-9a-f-]{36}\)", blocked.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"inputs",
|
||||
[
|
||||
{"texts": ["hi"], "tools": [{"type": "function", "function": {"name": "shell_exec"}}]},
|
||||
{"texts": ["hi"], "tools": [{"name": "shell_exec"}]},
|
||||
{"texts": ["hi"], "tool_calls": [{"name": "shell_exec"}]},
|
||||
{"texts": ["hi"], "tool_calls": [_ToolCallObject()]},
|
||||
],
|
||||
)
|
||||
async def test_unlisted_tool_in_a_request_is_blocked(inputs: dict):
|
||||
guardrail = _guardrail(ToolAuthPolicy(allowlist=("other",)))
|
||||
|
||||
with pytest.raises(HTTPException) as blocked:
|
||||
await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_blocks_disallowed_tool_definition_anthropic_shape():
|
||||
guardrail = TealTigerGuardrail(policies=_TOOL_POLICIES, policy_mode="ENFORCE")
|
||||
inputs = {"texts": ["hi"], "tools": [{"name": "shell_exec"}]}
|
||||
|
||||
with pytest.raises(ValueError, match="TOOL_NOT_ALLOWLISTED"):
|
||||
await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request")
|
||||
assert (blocked.value.status_code, bool(_BLOCKED_TOOL.fullmatch(blocked.value.detail))) == (400, True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_allows_allowlisted_tool_definition():
|
||||
guardrail = TealTigerGuardrail(policies=_TOOL_POLICIES, policy_mode="ENFORCE")
|
||||
async def test_unlisted_tool_call_in_a_response_is_blocked():
|
||||
guardrail = _guardrail(ToolAuthPolicy(allowlist=("search",)))
|
||||
|
||||
with pytest.raises(HTTPException) as blocked:
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": [], "tool_calls": [{"name": "shell_exec"}]}, request_data={}, input_type="response"
|
||||
)
|
||||
|
||||
assert _BLOCKED_TOOL.fullmatch(blocked.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_allowlisted_tool_definition_passes():
|
||||
guardrail = _guardrail(ToolAuthPolicy(allowlist=("search",)))
|
||||
inputs = {"texts": ["hi"], "tools": [{"type": "function", "function": {"name": "search"}}]}
|
||||
|
||||
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request")
|
||||
assert result["texts"] == ["hi"]
|
||||
|
||||
assert result == {"texts": ["hi"], "tools": [{"type": "function", "function": {"name": "search"}}]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_blocks_disallowed_tool_call_in_response():
|
||||
guardrail = TealTigerGuardrail(policies=_TOOL_POLICIES, policy_mode="ENFORCE")
|
||||
inputs = {"texts": [], "tool_calls": [{"name": "shell_exec"}]}
|
||||
|
||||
with pytest.raises(ValueError, match="TOOL_NOT_ALLOWLISTED"):
|
||||
await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="response")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_monitor_mode_does_not_block_disallowed_tool():
|
||||
guardrail = TealTigerGuardrail(policies=_TOOL_POLICIES, policy_mode="MONITOR")
|
||||
inputs = {"texts": ["hi"], "tools": [{"name": "shell_exec"}], "tool_calls": [{"name": "shell_exec"}]}
|
||||
async def test_monitor_mode_never_raises_and_leaves_text_untouched():
|
||||
guardrail = _guardrail(PiiPolicy(action="BLOCK"), ToolAuthPolicy(allowlist=("other",)), mode=PolicyMode.MONITOR)
|
||||
inputs = {"texts": ["SSN 123-45-6789"], "tools": [{"name": "shell_exec"}], "tool_calls": [{"name": "shell_exec"}]}
|
||||
|
||||
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request")
|
||||
assert result["texts"] == ["hi"]
|
||||
|
||||
assert result == inputs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_monitor_mode_does_not_block_exhausted_budget():
|
||||
policies = [{"type": "cost", "action": "ENFORCE", "daily_limit_usd": 0.00001}]
|
||||
guardrail = TealTigerGuardrail(policies=policies, policy_mode="MONITOR")
|
||||
guardrail.engine.track_cost(tokens=1_000_000)
|
||||
async def test_traffic_accrues_spend_per_user_until_the_daily_budget_blocks_that_user_only():
|
||||
guardrail = _guardrail(CostPolicy(daily_limit_usd=0.01, cost_per_1k_tokens=1.0))
|
||||
big_request = {"texts": ["x" * 40_000]} # about 10k tokens: $10 against a $0.01 limit
|
||||
|
||||
result = await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={}, input_type="request")
|
||||
assert result["texts"] == ["hi"]
|
||||
await guardrail.apply_guardrail(inputs=big_request, request_data={"user": "alice"}, input_type="request")
|
||||
|
||||
with pytest.raises(HTTPException) as blocked:
|
||||
await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={"user": "alice"}, input_type="request")
|
||||
bob = await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={"user": "bob"}, input_type="request")
|
||||
|
||||
assert blocked.value.detail.startswith("TealTiger: blocked: DAILY_BUDGET_EXCEEDED")
|
||||
assert bob == {"texts": ["hi"]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_accrues_from_traffic_and_eventually_blocks():
|
||||
policies = [{"type": "cost", "action": "ENFORCE", "daily_limit_usd": 0.01, "cost_per_1k_tokens": 1.0}]
|
||||
guardrail = TealTigerGuardrail(policies=policies, policy_mode="ENFORCE")
|
||||
big_text = "x" * 40_000 # about 10k tokens, about $10 at $1 per 1k tokens
|
||||
async def test_response_text_also_accrues_spend():
|
||||
guardrail = _guardrail(CostPolicy(daily_limit_usd=1.0, cost_per_1k_tokens=1.0))
|
||||
|
||||
await guardrail.apply_guardrail(inputs={"texts": [big_text]}, request_data={}, input_type="request")
|
||||
over_budget, spent, _ = guardrail.engine.check_budget()
|
||||
assert over_budget is True
|
||||
assert spent > 0
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["y" * 4000]}, request_data={"user": "alice"}, input_type="response"
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="DAILY_BUDGET_EXCEEDED"):
|
||||
await guardrail.apply_guardrail(inputs={"texts": ["next"]}, request_data={}, input_type="request")
|
||||
with pytest.raises(HTTPException) as blocked:
|
||||
await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={"user": "alice"}, input_type="request")
|
||||
assert blocked.value.detail.startswith("TealTiger: blocked: DAILY_BUDGET_EXCEEDED")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_accrues_on_response_path():
|
||||
policies = [{"type": "cost", "action": "ENFORCE", "daily_limit_usd": 100.0, "cost_per_1k_tokens": 1.0}]
|
||||
guardrail = TealTigerGuardrail(policies=policies, policy_mode="ENFORCE")
|
||||
async def test_every_non_clean_decision_is_logged_as_a_receipt():
|
||||
guardrail = _guardrail(PiiPolicy())
|
||||
|
||||
await guardrail.apply_guardrail(inputs={"texts": ["y" * 4000]}, request_data={}, input_type="response")
|
||||
_, spent, _ = guardrail.engine.check_budget()
|
||||
assert spent == pytest.approx(1.0)
|
||||
with patch("litellm.proxy.guardrails.guardrail_hooks.tealtiger.tealtiger.verbose_proxy_logger") as logger:
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["SSN 123-45-6789", "hello"]}, request_data={}, input_type="request"
|
||||
)
|
||||
|
||||
(info_call,) = logger.info.call_args_list
|
||||
template, receipt_json = info_call.args
|
||||
receipt = json.loads(receipt_json)
|
||||
assert (template, receipt["action"], receipt["reason_code"], receipt["risk_level"]) == (
|
||||
"tealtiger.decision_receipt=%s",
|
||||
"REDACT",
|
||||
"PII_DETECTED:ssn",
|
||||
"CRITICAL",
|
||||
)
|
||||
clean_reasons = [json.loads(call.args[1])["reason_code"] for call in logger.debug.call_args_list]
|
||||
assert clean_reasons == ["NO_VIOLATIONS", "NO_VIOLATIONS"], (
|
||||
"the budget check and the clean text each get a debug receipt"
|
||||
)
|
||||
|
||||
|
||||
def test_estimate_tokens():
|
||||
assert estimate_tokens([]) == 0
|
||||
assert estimate_tokens(["abcd"]) == 1
|
||||
assert estimate_tokens(["abcde", "fg"]) == 2
|
||||
@pytest.mark.asyncio
|
||||
async def test_monitor_mode_logs_the_pii_finding_it_does_not_enforce():
|
||||
guardrail = _guardrail(PiiPolicy(), mode=PolicyMode.MONITOR)
|
||||
|
||||
with patch("litellm.proxy.guardrails.guardrail_hooks.tealtiger.tealtiger.verbose_proxy_logger") as logger:
|
||||
await guardrail.apply_guardrail(inputs={"texts": ["SSN 123-45-6789"]}, request_data={}, input_type="request")
|
||||
|
||||
receipt = json.loads(logger.info.call_args.args[1])
|
||||
assert (receipt["action"], receipt["reason_code"], receipt["details"]["monitor_would"]) == (
|
||||
"ALLOW",
|
||||
"PII_DETECTED:ssn",
|
||||
"REDACT",
|
||||
)
|
||||
|
||||
|
||||
def test_pattern_scope_list_only_scans_listed_patterns():
|
||||
engine = TealEngine(policies=[{"type": "pii", "action": "REDACT", "patterns": ["email"]}])
|
||||
|
||||
ssn = engine.evaluate_text("SSN 123-45-6789")
|
||||
assert ssn.action == Action.ALLOW.value
|
||||
|
||||
email = engine.evaluate_text("mail jane@example.com")
|
||||
assert email.action == Action.REDACT.value
|
||||
assert "jane@example.com" not in (email.redacted_text or "")
|
||||
# ---------- wiring ----------
|
||||
|
||||
|
||||
def test_pattern_scope_accepts_category_names():
|
||||
engine = TealEngine(policies=[{"type": "pii", "action": "REDACT", "patterns": ["credentials"]}])
|
||||
|
||||
assert engine.evaluate_text("mail jane@example.com").action == Action.ALLOW.value
|
||||
assert engine.evaluate_text("key AKIAIOSFODNN7EXAMPLE").action == Action.REDACT.value
|
||||
|
||||
|
||||
def test_pattern_scope_unknown_name_fails_at_startup():
|
||||
with pytest.raises(ValueError, match="unknown pii pattern"):
|
||||
TealEngine(policies=[{"type": "pii", "action": "REDACT", "patterns": ["not_a_pattern"]}])
|
||||
|
||||
|
||||
def test_pattern_scope_invalid_type_fails_at_startup():
|
||||
with pytest.raises(TypeError, match="must be 'all' or a list"):
|
||||
TealEngine(policies=[{"type": "pii", "action": "REDACT", "patterns": 5}])
|
||||
|
||||
|
||||
def test_cost_rate_read_from_policy_and_defaults():
|
||||
custom = TealEngine(policies=[{"type": "cost", "daily_limit_usd": 1.0, "cost_per_1k_tokens": 0.5}])
|
||||
assert custom.cost_per_1k_tokens == 0.5
|
||||
default = TealEngine(policies=[{"type": "cost", "daily_limit_usd": 1.0}])
|
||||
assert default.cost_per_1k_tokens == 0.002
|
||||
def test_get_config_model_returns_the_tealtiger_config_model():
|
||||
assert TealTigerGuardrail.get_config_model() is TealTigerGuardrailConfigModel
|
||||
assert TealTigerGuardrailConfigModel.ui_friendly_name() == "TealTiger"
|
||||
|
||||
|
||||
def test_registries_are_plain_dicts_so_litellm_discovers_them():
|
||||
from litellm.proxy.guardrails.guardrail_hooks import tealtiger as hook
|
||||
|
||||
assert isinstance(hook.guardrail_initializer_registry, dict)
|
||||
assert isinstance(hook.guardrail_class_registry, dict)
|
||||
assert "tealtiger" in hook.guardrail_initializer_registry
|
||||
assert (type(hook.guardrail_initializer_registry), type(hook.guardrail_class_registry)) == (dict, dict)
|
||||
assert list(hook.guardrail_initializer_registry) == ["tealtiger"]
|
||||
|
||||
|
||||
def test_litellm_registry_discovers_tealtiger():
|
||||
from litellm.proxy.guardrails.guardrail_registry import get_guardrail_initializer_from_hooks
|
||||
|
||||
assert "tealtiger" in get_guardrail_initializer_from_hooks()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mode", "expected"),
|
||||
[
|
||||
(None, None),
|
||||
(GuardrailEventHooks.pre_call, GuardrailEventHooks.pre_call),
|
||||
("pre_call", GuardrailEventHooks.pre_call),
|
||||
(["pre_call", GuardrailEventHooks.post_call], [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call]),
|
||||
],
|
||||
)
|
||||
def test_to_event_hook_converts_litellm_params_mode(mode, expected):
|
||||
assert _to_event_hook(mode) == expected
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue