diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py index 84d1cc54546..a470b3f8bc1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py @@ -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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/budget.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/budget.py new file mode 100644 index 00000000000..fd6bfd55278 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/budget.py @@ -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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/config.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/config.py new file mode 100644 index 00000000000..1a533c724f4 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/config.py @@ -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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/decision.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/decision.py new file mode 100644 index 00000000000..56ee8ce7796 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/decision.py @@ -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" diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py index a84a49b9470..09d9e49345a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py @@ -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}), + ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/patterns.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/patterns.py index b9c90856acb..766e45e9915 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/patterns.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/patterns.py @@ -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"(? 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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/request.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/request.py new file mode 100644 index 00000000000..9f90df6f030 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/request.py @@ -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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py index 768d7c0909a..b7a4851f7a1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py @@ -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) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/tealtiger.py b/litellm/types/proxy/guardrails/guardrail_hooks/tealtiger.py new file mode 100644 index 00000000000..b851f38e8e5 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/tealtiger.py @@ -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" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py index 947ae4dd3a2..e991b8fa2a1 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py @@ -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