diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py new file mode 100644 index 00000000000..a470b3f8bc1 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py @@ -0,0 +1,37 @@ +from typing import TYPE_CHECKING, Final + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .config import parse_config +from .tealtiger import TealTigerGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +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=config.policies, + policy_mode=config.policy_mode, + guardrail_name=guardrail["guardrail_name"], + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + + litellm.logging_callback_manager.add_litellm_callback(_tealtiger_callback) + return _tealtiger_callback + + +# Plain dicts on purpose: guardrail_registry.py discovers hooks with isinstance(registry, dict) +guardrail_initializer_registry: Final = { + SupportedGuardrailIntegrations.TEALTIGER.value: initialize_guardrail, +} + +guardrail_class_registry: Final = { + SupportedGuardrailIntegrations.TEALTIGER.value: TealTigerGuardrail, +} 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 new file mode 100644 index 00000000000..09d9e49345a --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py @@ -0,0 +1,139 @@ +""" +TealEngine: deterministic, dependency-free PII / cost / tool-auth policy evaluation. +No network calls in the governance path. Every check returns a Decision value. +""" + +import math +import re +import uuid +from collections.abc import Mapping, Sequence +from dataclasses import replace +from functools import reduce +from types import MappingProxyType +from typing import Final + +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 + +_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 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) + + +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[Policy], + mode: PolicyMode = PolicyMode.ENFORCE, + ledger: BudgetLedger | None = None, + ) -> None: + 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 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 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) + + 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: + return reduce( + lambda redacted, name: self._patterns[name].sub(f"[REDACTED:{name.upper()}]", redacted), findings, text + ) + + 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 _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 new file mode 100644 index 00000000000..766e45e9915 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/patterns.py @@ -0,0 +1,198 @@ +""" +Regex-based PII / secrets detection patterns for TealTiger. + +47 patterns across five categories: government IDs, financial data, +contact info, network/device identifiers, and credentials/API keys. + +These are deterministic regex matches — fast (no ML, no network calls) but +inherently heuristic. Expect some false positives/negatives, especially on +the loosest patterns (e.g. generic credit card, generic driver's license). +Callers can narrow scope via TealEngine(policies=[{"patterns": [...]}]) +instead of "all", or override this module's PII_PATTERNS entirely. +""" + +from __future__ import annotations + +import re +from collections.abc import Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final + +PII_PATTERNS: Final = MappingProxyType( + { + # ---------- government / national IDs ---------- + "ssn": re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), + "ssn_no_dash": re.compile(r"\b(?!000|666|9\d{2})\d{3}(?!00)\d{2}(?!0000)\d{4}\b"), + "itin": re.compile(r"\b9\d{2}-(7\d|8[0-8])-\d{4}\b"), + "ein": re.compile(r"\b\d{2}-\d{7}\b"), + "us_passport": re.compile(r"\b[A-Z]{1,2}\d{7,9}\b"), + "passport_generic_intl": re.compile(r"\b[A-Z]{2}\d{6,9}\b"), + "drivers_license_generic": re.compile(r"\b[A-Z]{1,2}\d{6,8}\b"), + "medicare_id_us": re.compile(r"\b\d[A-Z]\d{2}-[A-Z]\d{2}-[A-Z]\d{2}\b"), + "npi_number": re.compile(r"\b\d{10}\b"), + "uk_nino": re.compile(r"\b[A-CEGHJ-PR-TW-Z]{2}\d{6}[A-D]\b"), + "date_of_birth": re.compile(r"\b(0[1-9]|1[0-2])[/-](0[1-9]|[12]\d|3[01])[/-](19|20)\d{2}\b"), + "vin": re.compile(r"\b[A-HJ-NPR-Z0-9]{17}\b"), + # ---------- financial ---------- + "credit_card_visa": re.compile(r"\b4\d{3}[ -]?\d{4}[ -]?\d{4}[ -]?\d{4}\b"), + "credit_card_mastercard": re.compile(r"\b5[1-5]\d{2}[ -]?\d{4}[ -]?\d{4}[ -]?\d{4}\b"), + "credit_card_amex": re.compile(r"\b3[47]\d{2}[ -]?\d{6}[ -]?\d{5}\b"), + "credit_card_discover": re.compile(r"\b6(?:011|5\d{2})[ -]?\d{4}[ -]?\d{4}[ -]?\d{4}\b"), + "credit_card_generic": re.compile(r"\b(?:\d[ -]*?){13,16}\b"), + "iban": re.compile(r"\b[A-Z]{2}\d{2}[A-Z0-9]{10,30}\b"), + "swift_bic": re.compile(r"\b[A-Z]{4}[A-Z]{2}[A-Z0-9]{2}(?:[A-Z0-9]{3})?\b"), + "us_bank_routing": re.compile(r"\b\d{9}\b"), + "us_bank_account": re.compile(r"\b\d{8,17}\b"), + "bitcoin_address": re.compile(r"\b(bc1|[13])[a-zA-HJ-NP-Z0-9]{25,39}\b"), + "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"(? 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 new file mode 100644 index 00000000000..45f45b1f912 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py @@ -0,0 +1,141 @@ +""" +TealTigerGuardrail: self-contained PII / cost / tool-authorization guardrail. + +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. + +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 typing import TYPE_CHECKING, Final, Literal + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # upstream decorator is untyped +) +from litellm.types.guardrails import GuardrailEventHooks, Mode +from litellm.types.proxy.guardrails.guardrail_hooks.tealtiger import DEFAULT_POLICIES, Policy, PolicyMode +from litellm.types.utils import GenericGuardrailAPIInputs + +from .decision import NO_VIOLATIONS, Action, Decision +from .engine import TealEngine +from .request import identity_of, tool_name_of + +if TYPE_CHECKING: + from pydantic import BaseModel + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + + +def _to_event_hook( + mode: GuardrailEventHooks # mutable-ok: matches LitellmParams.mode's real type + | list[GuardrailEventHooks] + | Mode + | str + | list[str] + | None, +) -> ( + 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__ accepts.""" + if mode is None or isinstance(mode, (GuardrailEventHooks, Mode)): + return mode + if isinstance(mode, str): + return GuardrailEventHooks(mode) + return [item if isinstance(item, GuardrailEventHooks) else GuardrailEventHooks(item) for item in mode] + + +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 _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[Policy] = DEFAULT_POLICIES, + policy_mode: PolicyMode = PolicyMode.ENFORCE, + guardrail_name: str | None = None, + # 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 + | str + | list[str] + | None = None, + default_on: bool | None = False, + ) -> None: + 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, + inputs: GenericGuardrailAPIInputs, + request_data: Mapping[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + """ + 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 as receipts and never raised. + + 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. + """ + 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) + + decisions: Final = (*tool_decisions, *budget_decisions, *text_decisions) + for decision in decisions: + _log_receipt(decision) + _raise_if_blocked(decisions) + + 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)} + + @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 () + return tuple(name for name in map(tool_name_of, (*definitions, *calls)) if name is not None) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 718a651bdee..02da743af64 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -163,6 +163,7 @@ class SupportedGuardrailIntegrations(Enum): AIRIA = "airia" WINGBACK = "wingback" DECISION_MODEL = "decision_model" + TEALTIGER = "tealtiger" class Role(Enum): 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 new file mode 100644 index 00000000000..aaf564b3c8c --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py @@ -0,0 +1,530 @@ +"""Unit tests for the TealTiger guardrail. Nothing on the other side of the call: no network, no API keys.""" + +import json +import re +from unittest.mock import patch + +import pytest +from fastapi import HTTPException +from pydantic import ValidationError + +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 TealTigerGuardrail, _to_event_hook +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.tealtiger import ( + DEFAULT_POLICIES, + 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 _engine(*policies, mode: PolicyMode = PolicyMode.ENFORCE) -> TealEngine: + return TealEngine(policies=policies, mode=mode) + + +# ---------- PII ---------- + + +def test_default_policies_redact_an_ssn_and_mark_it_critical(): + decision = _engine(*DEFAULT_POLICIES).check_text("My SSN is 123-45-6789") + + assert (decision.action, decision.reason_code, decision.risk_level, decision.redacted_text) == ( + Action.REDACT, + "PII_DETECTED:ssn", + RiskLevel.CRITICAL, + "My SSN is [REDACTED:SSN]", + ) + + +def test_email_is_redacted_at_medium_risk(): + decision = _engine(*DEFAULT_POLICIES).check_text("mail jane@example.com") + + assert (decision.reason_code, decision.risk_level, decision.redacted_text) == ( + "PII_DETECTED:email", + RiskLevel.MEDIUM, + "mail [REDACTED:EMAIL]", + ) + + +def test_clean_text_is_allowed_without_a_redaction(): + decision = _engine(*DEFAULT_POLICIES).check_text("What's the weather like today?") + + assert (decision.action, decision.reason_code, decision.risk_level, decision.redacted_text) == ( + Action.ALLOW, + "NO_VIOLATIONS", + RiskLevel.NONE, + None, + ) + + +@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}" + + +@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_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_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_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_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_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") + + assert (decision.action, decision.reason_code, decision.redacted_text, dict(decision.details)) == ( + Action.ALLOW, + "PII_DETECTED:ssn", + None, + {"types": ("ssn",), "monitor_would": "REDACT"}, + ) + + +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"]}, + } + + +# ---------- 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_request_text_is_redacted_and_other_inputs_are_preserved(): + inputs = {"texts": ["Contact me at jane@example.com"], "images": ["img"], "model": "gpt-4"} + + 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") + + assert (blocked.value.status_code, bool(_BLOCKED_TOOL.fullmatch(blocked.value.detail))) == (400, True) + + +@pytest.mark.asyncio +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"], "tools": [{"type": "function", "function": {"name": "search"}}]} + + +@pytest.mark.asyncio +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 == inputs + + +@pytest.mark.asyncio +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 + + 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_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": ["y" * 4000]}, request_data={"user": "alice"}, input_type="response" + ) + + 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_every_non_clean_decision_is_logged_as_a_receipt(): + guardrail = _guardrail(PiiPolicy()) + + 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" + ) + + +@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", + ) + + +# ---------- wiring ---------- + + +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 (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