From a8adc5ec198a3a2115d054aef702f6e844d4802e Mon Sep 17 00:00:00 2001 From: AlphaRex-pixel Date: Sun, 4 Oct 2026 21:46:35 +0530 Subject: [PATCH 1/7] feat: add TealTiger as a custom guardrail plugin --- docs/tealtiger.md | 60 ++++ .../guardrail_hooks/tealtiger/__init__.py | 37 +++ .../guardrail_hooks/tealtiger/engine.py | 137 ++++++++ .../guardrail_hooks/tealtiger/patterns.py | 144 ++++++++ .../guardrail_hooks/tealtiger/tealtiger.py | 175 ++++++++++ litellm/types/guardrails.py | 1 + .../guardrail_hooks/test_tealtiger.py | 311 ++++++++++++++++++ 7 files changed, 865 insertions(+) create mode 100644 docs/tealtiger.md create mode 100644 litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/tealtiger/patterns.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py create mode 100644 tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py diff --git a/docs/tealtiger.md b/docs/tealtiger.md new file mode 100644 index 00000000000..b45316e4362 --- /dev/null +++ b/docs/tealtiger.md @@ -0,0 +1,60 @@ +# TealTiger + +Self-contained PII detection, cost governance, and tool-authorization +guardrail. No API keys, no network calls in the governance path. + +## Quick Start + +Add to your proxy `config.yaml`: + +```yaml +model_list: + - model_name: gpt-4 + litellm_params: + model: gpt-4 + api_key: os.environ/OPENAI_API_KEY + +guardrails: + - guardrail_name: "tealtiger" + litellm_params: + guardrail: tealtiger + mode: "pre_call" # required for redaction to actually modify the outbound request; during_call/post_call do not rewrite the request before it's sent + policy_mode: "ENFORCE" # or "MONITOR" to dry-run without blocking +``` + +## Test it + +```bash +curl -i http://localhost:4000/chat/completions \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gpt-4", + "messages": [{"role": "user", "content": "My SSN is 123-45-6789"}], + "guardrails": ["tealtiger"] + }' +``` + +The request goes through with the SSN redacted rather than sent to the +model. Set `policy_mode: ENFORCE` + a `pii` policy `action: BLOCK` to reject +the call outright instead. + +## Policy types + +| type | action options | key params | +|-------------|-------------------|--------------------------------------| +| `pii` | `REDACT`, `BLOCK` | `patterns`: `"all"` or list of names | +| `cost` | `ENFORCE` | `daily_limit_usd` | +| `tool_auth` | `ENFORCE` | `allowlist`, `blocklist` | + +## Known limitations (tracked for follow-up PRs) + +- Cost tracking currently needs a hook into token-usage data that + `apply_guardrail` doesn't receive directly — see the `NOTE for reviewers` + comment in `tealtiger.py`. +- PII pattern set ships with 47 built-in patterns across 5 categories + (government ID, financial, contact, network/device, credentials) — see + `PATTERN_CATEGORIES` in `patterns.py`. Several are intentionally broad + (`credit_card_generic`, `us_bank_account`) and will overlap on shared + input; scope `patterns: [...]` explicitly rather than using `"all"` if + that matters for your use case. 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..5584b713c30 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py @@ -0,0 +1,37 @@ +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +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 + + _tealtiger_callback: Final = TealTigerGuardrail( + policies=getattr(litellm_params, "policies", None), + policy_mode=getattr(litellm_params, "policy_mode", "ENFORCE"), + guardrail_name=guardrail.get("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 + + +guardrail_initializer_registry: Final = MappingProxyType( + { + SupportedGuardrailIntegrations.TEALTIGER.value: initialize_guardrail, + } +) + +guardrail_class_registry: Final = MappingProxyType( + { + SupportedGuardrailIntegrations.TEALTIGER.value: TealTigerGuardrail, + } +) 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..f0b6f03eeb9 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py @@ -0,0 +1,137 @@ +""" +TealEngine: deterministic, dependency-free PII / cost / tool-auth policy +evaluation. No network calls in the governance path. +""" + +from __future__ import annotations + +import threading +import uuid +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from datetime import date +from enum import Enum +from types import MappingProxyType +from typing import Final + +from .patterns import default_patterns + + +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, + } + ) + + +class TealEngine: + def __init__( + self, + policies: Sequence[Mapping[str, object]], + mode: PolicyMode = PolicyMode.ENFORCE, + ) -> 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) + + 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.patterns.items() if pat.search(text)) + if not findings: + return Decision(Action.ALLOW.value, "NO_VIOLATIONS", correlation_id) + + 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 _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 + + 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 () + 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 () + return tool_name in allowlist + return True + + 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 = 0.002) -> float: + cost: Final = (tokens / 1000.0) * cost_per_1k_tokens + 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 + return cost 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..b9c90856acb --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/patterns.py @@ -0,0 +1,144 @@ +""" +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 +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"\b(?:\+?1[-.\s]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b"), + "phone_intl": re.compile(r"\+\d{1,3}[-.\s]?\(?\d{1,4}\)?(?:[-.\s]?\d{2,4}){2,4}\b"), + "us_zip_plus4": re.compile(r"\b\d{5}-\d{4}\b"), + "po_box": re.compile(r"\bP\.?O\.?\s?Box\s?\d+\b", re.IGNORECASE), + "lat_long_coordinates": re.compile(r"\b-?\d{1,3}\.\d{3,},\s?-?\d{1,3}\.\d{3,}\b"), + # ---------- network / device identifiers ---------- + "ipv4": re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"), + "ipv6": re.compile(r"\b(?:[A-Fa-f0-9]{1,4}:){7}[A-Fa-f0-9]{1,4}\b"), + "mac_address": re.compile(r"\b(?:[0-9A-Fa-f]{2}:){5}[0-9A-Fa-f]{2}\b"), + "imei": re.compile(r"\b\d{15}\b"), + "url_with_credentials": re.compile(r"\bhttps?://[^:\s]+:[^@\s]+@[^\s]+\b"), + # ---------- credentials / API keys / secrets ---------- + "aws_access_key": re.compile(r"\b(AKIA|ASIA)[0-9A-Z]{16}\b"), + "aws_secret_key": re.compile(r"\b(?=.*[A-Za-z])(?=.*\d)(?=.*[+/])[A-Za-z0-9+/]{40}\b"), + "gcp_api_key": re.compile(r"\bAIza[0-9A-Za-z_-]{35}\b"), + "azure_client_secret": re.compile(r"\b[A-Za-z0-9_~.-]{3}\dQ~[A-Za-z0-9_~.-]{31,34}\b"), + "github_token": re.compile(r"\bgh[pousr]_[A-Za-z0-9]{36,255}\b"), + "gitlab_token": re.compile(r"\bglpat-[A-Za-z0-9_-]{20}\b"), + "slack_token": re.compile(r"\bxox[baprs]-[A-Za-z0-9-]{10,72}\b"), + "stripe_live_key": re.compile(r"\b(sk|pk)_live_[A-Za-z0-9]{24,}\b"), + "openai_api_key": re.compile(r"\bsk-[A-Za-z0-9]{20,}\b"), + "jwt": re.compile(r"\beyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\b"), + "private_key_pem": re.compile(r"-----BEGIN (RSA |EC |OPENSSH |)PRIVATE KEY-----"), + "generic_bearer_token": re.compile(r"\bBearer\s+[A-Za-z0-9_\-.]{20,}\b"), + "basic_auth_header": re.compile(r"\bBasic\s+[A-Za-z0-9+/]{20,}={0,2}\b"), + } +) + + +def default_patterns() -> Mapping[str, re.Pattern[str]]: + """Return the built-in, immutable pattern mapping.""" + return PII_PATTERNS + + +PATTERN_CATEGORIES: Final = MappingProxyType( + { + "government_id": ( + "ssn", + "ssn_no_dash", + "itin", + "ein", + "us_passport", + "passport_generic_intl", + "drivers_license_generic", + "medicare_id_us", + "npi_number", + "uk_nino", + "date_of_birth", + "vin", + ), + "financial": ( + "credit_card_visa", + "credit_card_mastercard", + "credit_card_amex", + "credit_card_discover", + "credit_card_generic", + "iban", + "swift_bic", + "us_bank_routing", + "us_bank_account", + "bitcoin_address", + "ethereum_address", + ), + "contact": ( + "email", + "phone_us", + "phone_intl", + "us_zip_plus4", + "po_box", + "lat_long_coordinates", + ), + "network_device": ( + "ipv4", + "ipv6", + "mac_address", + "imei", + "url_with_credentials", + ), + "credentials": ( + "aws_access_key", + "aws_secret_key", + "gcp_api_key", + "azure_client_secret", + "github_token", + "gitlab_token", + "slack_token", + "stripe_live_key", + "openai_api_key", + "jwt", + "private_key_pem", + "generic_bearer_token", + "basic_auth_header", + ), + } +) 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..1d81a7b4abe --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py @@ -0,0 +1,175 @@ +""" +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. + +No network calls, no LLM calls, no external API keys required. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal + +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.types.guardrails import GuardrailEventHooks, Mode +from litellm.types.utils import GenericGuardrailAPIInputs + +from .engine import Action, PolicyMode, TealEngine + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_EMPTY_MAPPING: Final = MappingProxyType({}) + +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 +) + + +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__ actually accepts, converting bare strings to the real enum.""" + 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 _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") + nested_name: Final = function.get("name") if isinstance(function, Mapping) else None + return nested_name if isinstance(nested_name, str) else None + + +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) + return _tool_call_name_from_object(tool_call) + + +class TealTigerGuardrail(CustomGuardrail): + def __init__( + self, + policies: Sequence[Mapping[str, object]] | None = None, + policy_mode: str = "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. + 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 or DEFAULT_POLICIES, + mode=PolicyMode(policy_mode), + ) + super().__init__( + guardrail_name=guardrail_name, + event_hook=_to_event_hook(event_hook), + default_on=default_on or False, + ) + + @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 both request (pre-call) and response (post-call) content, + since TealTiger scans for PII leakage in both directions. + + Tool-call authorization and budget checks only make sense on the + request path, so they're gated on input_type == "request". + + Returns a new mapping rather than mutating `inputs` in place, since + the caller's object should not be rewritten at a distance. + """ + if input_type == "request": + self._check_tool_calls(inputs) + self._check_budget(request_data) + + checked_texts: Final = tuple(self._check_text_or_raise(text) for text in inputs.get("texts") or ()) + return {**inputs, "texts": list(checked_texts)} + + def _check_tool_calls(self, inputs: GenericGuardrailAPIInputs) -> None: + for tool_call in inputs.get("tool_calls") or (): + tool_name = _tool_call_name(tool_call) + if tool_name and not self.engine.check_tool(tool_name): + raise ValueError(f"TealTiger: blocked — 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: + raise ValueError(f"TealTiger: blocked — 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 + + # ---- cost tracking on successful response ---- + # NOTE for reviewers: apply_guardrail is only handed extracted `texts` / + # `images` / `tool_calls`, not token-usage data — confirmed by reading + # litellm/integrations/custom_guardrail.py's apply_guardrail docstring + # directly, not just documentation. Wiring TealEngine.track_cost() into + # something more precise than a fixed per-request estimate needs either + # (a) an async_post_call_success_hook override that reads response.usage + # (loses the endpoint-agnostic behavior apply_guardrail gives us), or + # (b) whatever usage data logging_obj exposes at call time — worth + # asking maintainers directly rather than guessing further. diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 46026c12d24..18ae8943473 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -146,6 +146,7 @@ class SupportedGuardrailIntegrations(Enum): ALICE = "alice" AGENT_365 = "agent_365" CONDUCT = "conduct" + TEALTIGER = "tealtiger" class Role(Enum): 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..88c2762e8d2 --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py @@ -0,0 +1,311 @@ +""" +Unit tests for the TealTiger guardrail. + +Run with: pytest tests/guardrails_tests/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 pytest + +from litellm.proxy.guardrails.guardrail_hooks.tealtiger.engine import ( + Action, + PolicyMode, + TealEngine, +) +from litellm.proxy.guardrails.guardrail_hooks.tealtiger.tealtiger import ( + DEFAULT_POLICIES, + TealTigerGuardrail, +) + + +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"]}, + ] + 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 + + +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_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 "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"}], + } + + 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 result["texts"] == ["hello"] + + +# ---------- coverage for _to_event_hook and _tool_call_name helpers ---------- + +from litellm.proxy.guardrails.guardrail_hooks.tealtiger.tealtiger import ( + _to_event_hook, + _tool_call_name, +) +from litellm.types.guardrails import GuardrailEventHooks + + +def test_to_event_hook_passes_through_none(): + assert _to_event_hook(None) is None + + +def test_to_event_hook_passes_through_real_enum(): + result = _to_event_hook(GuardrailEventHooks.pre_call) + assert result == GuardrailEventHooks.pre_call + + +def test_to_event_hook_converts_bare_string(): + result = _to_event_hook("pre_call") + assert result == GuardrailEventHooks.pre_call + + +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_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_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_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 "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 From f11dc48032e4350b923ec904c74f4cc9e41bd01b Mon Sep 17 00:00:00 2001 From: AlphaRex-pixel Date: Sun, 11 Oct 2026 17:43:07 +0530 Subject: [PATCH 2/7] fix(guardrails): address tealtiger review findings --- docs/tealtiger.md | 60 ------- .../guardrail_hooks/tealtiger/__init__.py | 18 +-- .../guardrail_hooks/tealtiger/engine.py | 51 +++++- .../guardrail_hooks/tealtiger/tealtiger.py | 43 ++--- .../guardrail_hooks/test_tealtiger.py | 147 +++++++++++++++++- 5 files changed, 222 insertions(+), 97 deletions(-) delete mode 100644 docs/tealtiger.md diff --git a/docs/tealtiger.md b/docs/tealtiger.md deleted file mode 100644 index b45316e4362..00000000000 --- a/docs/tealtiger.md +++ /dev/null @@ -1,60 +0,0 @@ -# TealTiger - -Self-contained PII detection, cost governance, and tool-authorization -guardrail. No API keys, no network calls in the governance path. - -## Quick Start - -Add to your proxy `config.yaml`: - -```yaml -model_list: - - model_name: gpt-4 - litellm_params: - model: gpt-4 - api_key: os.environ/OPENAI_API_KEY - -guardrails: - - guardrail_name: "tealtiger" - litellm_params: - guardrail: tealtiger - mode: "pre_call" # required for redaction to actually modify the outbound request; during_call/post_call do not rewrite the request before it's sent - policy_mode: "ENFORCE" # or "MONITOR" to dry-run without blocking -``` - -## Test it - -```bash -curl -i http://localhost:4000/chat/completions \ - -H "Authorization: Bearer sk-1234" \ - -H "Content-Type: application/json" \ - -d '{ - "model": "gpt-4", - "messages": [{"role": "user", "content": "My SSN is 123-45-6789"}], - "guardrails": ["tealtiger"] - }' -``` - -The request goes through with the SSN redacted rather than sent to the -model. Set `policy_mode: ENFORCE` + a `pii` policy `action: BLOCK` to reject -the call outright instead. - -## Policy types - -| type | action options | key params | -|-------------|-------------------|--------------------------------------| -| `pii` | `REDACT`, `BLOCK` | `patterns`: `"all"` or list of names | -| `cost` | `ENFORCE` | `daily_limit_usd` | -| `tool_auth` | `ENFORCE` | `allowlist`, `blocklist` | - -## Known limitations (tracked for follow-up PRs) - -- Cost tracking currently needs a hook into token-usage data that - `apply_guardrail` doesn't receive directly — see the `NOTE for reviewers` - comment in `tealtiger.py`. -- PII pattern set ships with 47 built-in patterns across 5 categories - (government ID, financial, contact, network/device, credentials) — see - `PATTERN_CATEGORIES` in `patterns.py`. Several are intentionally broad - (`credit_card_generic`, `us_bank_account`) and will overlap on shared - input; scope `patterns: [...]` explicitly rather than using `"all"` if - that matters for your use case. diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py index 5584b713c30..110a714a191 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py @@ -1,4 +1,3 @@ -from types import MappingProxyType from typing import TYPE_CHECKING, Final from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -24,14 +23,11 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" return _tealtiger_callback -guardrail_initializer_registry: Final = MappingProxyType( - { - SupportedGuardrailIntegrations.TEALTIGER.value: initialize_guardrail, - } -) +# 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 = MappingProxyType( - { - SupportedGuardrailIntegrations.TEALTIGER.value: TealTigerGuardrail, - } -) +guardrail_class_registry: Final = { + SupportedGuardrailIntegrations.TEALTIGER.value: TealTigerGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py index f0b6f03eeb9..acb82d159ec 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py @@ -5,16 +5,27 @@ evaluation. No network calls in the governance path. 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 types import MappingProxyType from typing import Final -from .patterns import default_patterns +from .patterns import PATTERN_CATEGORIES, default_patterns + +DEFAULT_COST_PER_1K_TOKENS: Final = 0.002 +_CHARS_PER_TOKEN: Final = 4 + + +def estimate_tokens(texts: Sequence[str]) -> int: + """Rough, dependency-free token estimate (about 4 characters per token).""" + return math.ceil(sum(len(text) for text in texts) / _CHARS_PER_TOKEN) class PolicyMode(str, Enum): @@ -62,13 +73,39 @@ class TealEngine: 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 + ) + + 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 + entries: Final = (scope,) if isinstance(scope, str) else scope + if not isinstance(entries, (list, tuple)): + raise TypeError("TealTiger: pii policy `patterns` must be 'all' or a list of names") + 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.patterns.items() if pat.search(text)) + findings: Final = tuple(name for name, pat in self._active_patterns.items() if pat.search(text)) if not findings: return Decision(Action.ALLOW.value, "NO_VIOLATIONS", correlation_id) @@ -129,9 +166,13 @@ class TealEngine: 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 = 0.002) -> float: - cost: Final = (tokens / 1000.0) * cost_per_1k_tokens + 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 - return 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)) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py index 1d81a7b4abe..ead4955fd17 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py @@ -16,6 +16,7 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal +from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, @@ -131,45 +132,49 @@ class TealTigerGuardrail(CustomGuardrail): Runs on both request (pre-call) and response (post-call) content, since TealTiger scans for PII leakage in both directions. - Tool-call authorization and budget checks only make sense on the - request path, so they're gated on input_type == "request". + 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. + + 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. """ + self._check_tools(inputs, input_type) if input_type == "request": - self._check_tool_calls(inputs) self._check_budget(request_data) 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)} - def _check_tool_calls(self, inputs: GenericGuardrailAPIInputs) -> None: - for tool_call in inputs.get("tool_calls") or (): - tool_name = _tool_call_name(tool_call) + 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) + + def _check_tools(self, inputs: GenericGuardrailAPIInputs, input_type: Literal["request", "response"]) -> None: + 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): - raise ValueError(f"TealTiger: blocked — TOOL_NOT_ALLOWLISTED ({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: - raise ValueError(f"TealTiger: blocked — DAILY_BUDGET_EXCEEDED (${spent:.2f} / ${limit:.2f})") + 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 - - # ---- cost tracking on successful response ---- - # NOTE for reviewers: apply_guardrail is only handed extracted `texts` / - # `images` / `tool_calls`, not token-usage data — confirmed by reading - # litellm/integrations/custom_guardrail.py's apply_guardrail docstring - # directly, not just documentation. Wiring TealEngine.track_cost() into - # something more precise than a fixed per-request estimate needs either - # (a) an async_post_call_success_hook override that reads response.usage - # (loses the endpoint-agnostic behavior apply_guardrail gives us), or - # (b) whatever usage data logging_obj exposes at call time — worth - # asking maintainers directly rather than guessing further. diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py index 88c2762e8d2..947ae4dd3a2 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py @@ -1,18 +1,20 @@ """ Unit tests for the TealTiger guardrail. -Run with: pytest tests/guardrails_tests/test_tealtiger.py -v +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 pytest from litellm.proxy.guardrails.guardrail_hooks.tealtiger.engine import ( Action, PolicyMode, TealEngine, + estimate_tokens, ) from litellm.proxy.guardrails.guardrail_hooks.tealtiger.tealtiger import ( DEFAULT_POLICIES, @@ -145,11 +147,12 @@ async def test_apply_guardrail_session_id_falls_back_when_user_not_a_string(): # ---------- coverage for _to_event_hook and _tool_call_name helpers ---------- +from litellm.types.guardrails import GuardrailEventHooks + from litellm.proxy.guardrails.guardrail_hooks.tealtiger.tealtiger import ( _to_event_hook, _tool_call_name, ) -from litellm.types.guardrails import GuardrailEventHooks def test_to_event_hook_passes_through_none(): @@ -309,3 +312,143 @@ 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 + + +_TOOL_POLICIES = [ + {"type": "tool_auth", "action": "ENFORCE", "allowlist": ["search"]}, +] + + +@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"}}]} + + with pytest.raises(ValueError, match="TOOL_NOT_ALLOWLISTED"): + 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") + + +@pytest.mark.asyncio +async def test_apply_guardrail_allows_allowlisted_tool_definition(): + guardrail = TealTigerGuardrail(policies=_TOOL_POLICIES, policy_mode="ENFORCE") + 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"] + + +@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"}]} + + result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + assert result["texts"] == ["hi"] + + +@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) + + result = await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={}, input_type="request") + assert result["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 + + 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 + + with pytest.raises(ValueError, match="DAILY_BUDGET_EXCEEDED"): + await guardrail.apply_guardrail(inputs={"texts": ["next"]}, request_data={}, input_type="request") + + +@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") + + await guardrail.apply_guardrail(inputs={"texts": ["y" * 4000]}, request_data={}, input_type="response") + _, spent, _ = guardrail.engine.check_budget() + assert spent == pytest.approx(1.0) + + +def test_estimate_tokens(): + assert estimate_tokens([]) == 0 + assert estimate_tokens(["abcd"]) == 1 + assert estimate_tokens(["abcde", "fg"]) == 2 + + +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 "") + + +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_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 + + +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() From e0e2e009ed778a27067fa3c368c716c40a278f2b Mon Sep 17 00:00:00 2001 From: AlphaRex-pixel Date: Sun, 11 Oct 2026 21:57:05 +0530 Subject: [PATCH 3/7] fix(guardrails): restore main's guardrails types and re-add tealtiger enum member --- litellm/types/guardrails.py | 247 +++++++++++++++++++++++++++++------- 1 file changed, 200 insertions(+), 47 deletions(-) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 2725653d9ac..cbc54015aca 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -2,12 +2,14 @@ from collections.abc import Mapping from datetime import datetime from enum import Enum from types import MappingProxyType -from typing import Final, Literal +from typing import Final, Literal, cast -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import ConfigDict, Field, field_validator, model_validator from typing_extensions import ReadOnly, Required, TypedDict +from litellm._logging import verbose_logger from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS +from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.proxy.guardrails.guardrail_hooks.agent_365 import ( Agent365GuardrailConfigModel, ) @@ -23,6 +25,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.cisco_ai_defense import ( from litellm.types.proxy.guardrails.guardrail_hooks.compresr import ( CompresrGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.decision_model import ( + DecisionModelCheck, +) from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import ( EnkryptAIGuardrailConfigs, ) @@ -145,9 +150,10 @@ class SupportedGuardrailIntegrations(Enum): STRAIKER = "straiker" ALICE = "alice" AGENT_365 = "agent_365" + LLM_SHIELD_PROXY = "llm_shield_proxy" CONDUCT = "conduct" - TEALTIGER = "tealtiger" DECISION_MODEL = "decision_model" + TEALTIGER = 'tealtiger' class Role(Enum): @@ -167,7 +173,7 @@ class GuardrailItemSpec(TypedDict, total=False): callback_args: dict[str, dict] -class GuardrailItem(BaseModel): +class GuardrailItem(LiteLLMBaseModel): callbacks: list[str] default_on: bool logging_only: bool | None @@ -437,7 +443,7 @@ class GuardrailParamUITypes(str, Enum): PERCENTAGE = "percentage" -class PresidioPresidioConfigModelUserInterface(BaseModel): +class PresidioPresidioConfigModelUserInterface(LiteLLMBaseModel): """Configuration parameters for the Presidio PII masking guardrail on LiteLLM UI""" presidio_analyzer_api_base: str | None = Field( @@ -546,31 +552,31 @@ BedrockChecksSensitiveInformationEntity = Literal[ ] -class BedrockChecksContentFilterCategoryItem(BaseModel): +class BedrockChecksContentFilterCategoryItem(LiteLLMBaseModel): category: BedrockChecksContentFilterCategory -class BedrockChecksContentFilterModel(BaseModel): +class BedrockChecksContentFilterModel(LiteLLMBaseModel): categories: list[BedrockChecksContentFilterCategoryItem] -class BedrockChecksPromptAttackCategoryItem(BaseModel): +class BedrockChecksPromptAttackCategoryItem(LiteLLMBaseModel): category: BedrockChecksPromptAttackCategory -class BedrockChecksPromptAttackModel(BaseModel): +class BedrockChecksPromptAttackModel(LiteLLMBaseModel): categories: list[BedrockChecksPromptAttackCategoryItem] -class BedrockChecksSensitiveInformationEntityItem(BaseModel): +class BedrockChecksSensitiveInformationEntityItem(LiteLLMBaseModel): type: BedrockChecksSensitiveInformationEntity -class BedrockChecksSensitiveInformationModel(BaseModel): +class BedrockChecksSensitiveInformationModel(LiteLLMBaseModel): entities: list[BedrockChecksSensitiveInformationEntityItem] -class BedrockChecksConfigModel(BaseModel): +class BedrockChecksConfigModel(LiteLLMBaseModel): """Inline `checks` config for the resource-less Bedrock InvokeGuardrailChecks API. Include only the checks you want to run; at least one must be set. @@ -589,7 +595,7 @@ class BedrockChecksConfigModel(BaseModel): return self -class BedrockGuardrailConfigModel(BaseModel): +class BedrockGuardrailConfigModel(LiteLLMBaseModel): """Configuration parameters for the AWS Bedrock guardrail""" guardrailIdentifier: str | None = Field(default=None, description="The ID of your guardrail on Bedrock") @@ -664,7 +670,7 @@ class BedrockGuardrailConfigModel(BaseModel): ) -class BedrockGuardrailStreamingParams(BaseModel): +class BedrockGuardrailStreamingParams(LiteLLMBaseModel): streaming_buffer_until_moderated: bool = Field( default=True, description="If True (default), withhold every streamed chunk until the end-of-stream " @@ -704,7 +710,7 @@ class BedrockGuardrailStreamingParams(BaseModel): ) -class LakeraV2GuardrailConfigModel(BaseModel): +class LakeraV2GuardrailConfigModel(LiteLLMBaseModel): """Configuration parameters for the Lakera AI v2 guardrail""" api_key: str | None = Field(default=None, description="API key for the Lakera AI service") @@ -729,7 +735,7 @@ class LakeraV2GuardrailConfigModel(BaseModel): ) -class LassoGuardrailConfigModel(BaseModel): +class LassoGuardrailConfigModel(LiteLLMBaseModel): """Configuration parameters for the Lasso guardrail""" lasso_user_id: str | None = Field(default=None, description="User ID for the Lasso guardrail") @@ -737,7 +743,7 @@ class LassoGuardrailConfigModel(BaseModel): mask: bool | None = Field(default=False, description="Enable content masking using Lasso classifix API") -class DeepKeepGuardrailConfigModel(BaseModel): +class DeepKeepGuardrailConfigModel(LiteLLMBaseModel): """Configuration parameters for the DeepKeep AI Firewall guardrail""" deepkeep_firewall_id: str | None = Field( @@ -749,7 +755,7 @@ class DeepKeepGuardrailConfigModel(BaseModel): ) -class PillarGuardrailConfigModel(BaseModel): +class PillarGuardrailConfigModel(LiteLLMBaseModel): """Configuration parameters for the Pillar Security guardrail""" on_flagged_action: str | None = Field( @@ -774,7 +780,7 @@ class PillarGuardrailConfigModel(BaseModel): ) -class NomaGuardrailConfigModel(BaseModel): +class NomaGuardrailConfigModel(LiteLLMBaseModel): """Configuration parameters for the Noma Security guardrail""" use_v2: bool | None = Field( @@ -803,7 +809,7 @@ class NomaGuardrailConfigModel(BaseModel): ) -class ZscalerAIGuardConfigModel(BaseModel): +class ZscalerAIGuardConfigModel(LiteLLMBaseModel): """Configuration parameters for the Zscaler AI Guard guardrail""" policy_id: int | None = Field( @@ -821,7 +827,7 @@ class ZscalerAIGuardConfigModel(BaseModel): ) -class JavelinGuardrailConfigModel(BaseModel): +class JavelinGuardrailConfigModel(LiteLLMBaseModel): """Configuration parameters for the Javelin guardrail""" guard_name: str | None = Field(default=None, description="Name of the Javelin guard to use") @@ -838,7 +844,7 @@ class ContentFilterAction(str, Enum): MASK = "MASK" -class BlockedWord(BaseModel): +class BlockedWord(LiteLLMBaseModel): """Represents a blocked word with its action and optional description""" keyword: str = Field(description="The keyword to block or mask") @@ -849,7 +855,7 @@ class BlockedWord(BaseModel): ) -class ContentFilterPattern(BaseModel): +class ContentFilterPattern(LiteLLMBaseModel): """Represents a content filter pattern (prebuilt or custom regex)""" pattern_type: Literal["prebuilt", "regex"] = Field( @@ -870,7 +876,7 @@ class ContentFilterPattern(BaseModel): action: ContentFilterAction = Field(description="Action to take when pattern matches (BLOCK or MASK)") -class ContentFilterConfigModel(BaseModel): +class ContentFilterConfigModel(LiteLLMBaseModel): """Configuration parameters for the content filter guardrail""" patterns: list[ContentFilterPattern] | None = Field( @@ -901,6 +907,95 @@ class ContentFilterConfigModel(BaseModel): MCP_SECURITY_ON_VIOLATION: Final = frozenset({"block", "alert"}) +GuardrailStreamScope = Literal["streaming", "non_streaming", "both"] +DEFAULT_GUARDRAIL_STREAM_SCOPE: Final[GuardrailStreamScope] = "both" + + +class GuardrailEventHooks(str, Enum): + pre_call = "pre_call" + post_call = "post_call" + during_call = "during_call" + logging_only = "logging_only" + pre_mcp_call = "pre_mcp_call" + during_mcp_call = "during_mcp_call" + post_mcp_call = "post_mcp_call" + realtime_input_transcription = "realtime_input_transcription" + + +GUARDRAIL_EVENT_HOOK_VALUES: Final = frozenset(member.value for member in GuardrailEventHooks) + +_GUARDRAIL_STREAM_SCOPES: Final[Mapping[str, GuardrailStreamScope]] = MappingProxyType( + { + "streaming": "streaming", + "non_streaming": "non_streaming", + "both": "both", + } +) + + +def _as_guardrail_stream_scope(value: object) -> GuardrailStreamScope: + if not isinstance(value, str): + raise ValueError(f"stream_scope values must be strings, got {type(value).__name__}") + scope: Final = _GUARDRAIL_STREAM_SCOPES.get(value.lower()) + if scope is None: + raise ValueError(f"stream_scope must be one of both, streaming, non_streaming, got {value!r}") + return scope + + +def _validated_stream_scope_hook(key: object) -> str: + if not isinstance(key, str): + raise ValueError(f"stream_scope keys must be strings, got {type(key).__name__}") + hook: Final = key.lower() + if hook not in GUARDRAIL_EVENT_HOOK_VALUES: + raise ValueError( + f"stream_scope keys must be guardrail modes ({sorted(GUARDRAIL_EVENT_HOOK_VALUES)}), got {key!r}" + ) + return hook + + +def coerce_stream_scope(value: object) -> GuardrailStreamScope | dict[str, GuardrailStreamScope] | None: + if value is None: + return None + if isinstance(value, str): + return _as_guardrail_stream_scope(value) + if isinstance(value, Mapping): + scope_map: Final[Mapping[str, object]] = cast(Mapping[str, object], value) # cast-ok: keys validated below + return { + _validated_stream_scope_hook(key): _as_guardrail_stream_scope(scope) for key, scope in scope_map.items() + } + raise ValueError(f"stream_scope must be a string or mapping, got {type(value).__name__}") + + +def stored_stream_scope(value: object) -> GuardrailStreamScope | dict[str, GuardrailStreamScope] | None: + try: + return coerce_stream_scope(value) + except ValueError: + verbose_logger.warning("Ignoring invalid stored stream_scope value of type %s", type(value).__name__) + return None + + +def with_tolerated_stream_scope(params: Mapping[str, object]) -> dict[str, object]: + if "stream_scope" not in params: + return dict(params) + return { + **params, + "stream_scope": stored_stream_scope(params["stream_scope"]), + } + + +def runtime_stream_scope( + stream_scope: object, +) -> tuple[GuardrailStreamScope, MappingProxyType[str, GuardrailStreamScope]]: + coerced: Final = coerce_stream_scope(stream_scope) + if coerced is None: + return DEFAULT_GUARDRAIL_STREAM_SCOPE, MappingProxyType({}) + if isinstance(coerced, str): + return coerced, MappingProxyType({}) + return DEFAULT_GUARDRAIL_STREAM_SCOPE, MappingProxyType(coerced) + + +LoggingOnlyScope = Literal["input", "output"] + class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch update guardrails api_key: str | None = Field(default=None, description="API key for the guardrail service") @@ -1043,9 +1138,11 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up skip_unscannable_attachments: bool | None = Field( default=False, description=( - "Implemented by guardrail='model_armor'. When True, attachment references that carry no " - "inline bytes (file_id, gs://, or http(s) URLs) pass through unscanned instead of blocking, " - "while fail_on_error still governs real Model Armor API errors. Default False blocks them." + "Implemented by guardrail='model_armor' and guardrail='bedrock'. When True, attachments the " + "guardrail cannot scan pass through unscanned instead of blocking. For Model Armor these are " + "references with no inline bytes (file_id, gs://, or http(s) URLs), and fail_on_error still " + "governs real Model Armor API errors. For Bedrock these are documents, files, audio, video, " + "and images that are not inline PNG or JPEG up to 4 MB. Default False blocks them." ), ) sanitize_error_detail: bool | None = Field( @@ -1066,10 +1163,23 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up description=( "Behavior when a guardrail endpoint is unreachable due to network errors. " "Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', and 'typesafe'. " - "'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed." + "'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed. " + "Also implemented by guardrail='decision_model'." ), ) + max_input_chars: int | None = Field( + default=None, + gt=0, + description="Character budget for each text sent to the guardrail's model. Implemented by guardrail='decision_model'.", + ) + + max_concurrent_decision_calls: int | None = Field( + default=None, + gt=0, + description="Maximum decisions calls in flight at once on the guardrail in each proxy worker. Implemented by guardrail='decision_model'.", + ) + extra_headers: list[str] | None = Field( default=None, description=( @@ -1142,6 +1252,39 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up ), ) + stream_scope: GuardrailStreamScope | dict[str, GuardrailStreamScope] | None = Field( + default=None, + description=( + "Whether this guardrail runs on streaming requests, non-streaming requests, or both. " + "A string applies to every configured mode. A map overrides named modes " + "(pre_call, during_call, post_call, ...); omitted keys default to both. " + "Unset means both, matching historical behavior." + ), + ) + + @field_validator("stream_scope", mode="before") + @classmethod + def normalize_stream_scope(cls, v: object) -> GuardrailStreamScope | dict[str, GuardrailStreamScope] | None: + return coerce_stream_scope(v) + + logging_only_scope: LoggingOnlyScope | None = Field( + default=None, + description=( + "which direction a logging_only scan observes: 'input' (request) or 'output' (response); " + "unset scans both directions. Only applies to mode logging_only; pre_call/post_call on the " + "same guardrail keep blocking." + ), + ) + + logging_only_continue_on_input_failure: bool | None = Field( + default=None, + description=( + "when True, a flagged or raising logging_only request scan is logged and the response is " + "still scanned, so both verdicts land. Only applies to mode logging_only and is ignored " + "when logging_only_scope is 'input' or 'output'." + ), + ) + @field_validator( "mode", "default_action", @@ -1169,7 +1312,7 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up model_config = ConfigDict(extra="allow", protected_namespaces=()) -class Mode(BaseModel): +class Mode(LiteLLMBaseModel): tags: dict[str, str | list[str]] = Field(description="Tags for the guardrail mode") default: str | list[str] | None = Field(default=None, description="Default mode when no tags match") @@ -1207,6 +1350,17 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o Agent365GuardrailConfigModel, ): guardrail: str = Field(description="The type of guardrail integration to use") + decision_model: str | None = Field( + default=None, + description="For guardrail='decision_model': the Decisions-API model that scores each check", + ) + checks: BedrockChecksConfigModel | tuple[DecisionModelCheck, ...] | None = Field( # pyright: ignore[reportIncompatibleVariableOverride] # widened to also accept decision-model checks + default=None, + description=( + "Inline Bedrock InvokeGuardrailChecks config for guardrail='bedrock', or the predicate " + "checks a decision_model guardrail scores for guardrail='decision_model'" + ), + ) mode: str | list[str] | Mode = Field( description="When to apply the guardrail (pre_call, post_call, during_call, logging_only)" ) @@ -1232,6 +1386,14 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o raise ValueError(f"on_violation={self.on_violation!r} is only supported by guardrail='mcp_security'") return self + @model_validator(mode="after") + def validate_checks_list_for_guardrail(self) -> "LitellmParams": + if isinstance(self.checks, tuple) and self.guardrail != SupportedGuardrailIntegrations.DECISION_MODEL.value: + raise ValueError( + f"checks as a list is only supported by guardrail='decision_model', got guardrail={self.guardrail!r}" + ) + return self + def __init__(self, **kwargs) -> None: default_on: Final = kwargs.pop("default_on", None) if default_on is not None: @@ -1268,17 +1430,6 @@ class guardrailConfig(TypedDict): guardrails: list[Guardrail] -class GuardrailEventHooks(str, Enum): - pre_call = "pre_call" - post_call = "post_call" - during_call = "during_call" - logging_only = "logging_only" - pre_mcp_call = "pre_mcp_call" - during_mcp_call = "during_mcp_call" - post_mcp_call = "post_mcp_call" - realtime_input_transcription = "realtime_input_transcription" - - class DynamicGuardrailParams(TypedDict): extra_body: ReadOnly[dict[str, object]] @@ -1288,7 +1439,7 @@ class GUARDRAIL_DEFINITION_LOCATION(str, Enum): CONFIG = "config" -class GuardrailInfoResponse(BaseModel): +class GuardrailInfoResponse(LiteLLMBaseModel): guardrail_id: str | None = None guardrail_name: str litellm_params: BaseLitellmParams | None = None @@ -1301,20 +1452,22 @@ class GuardrailInfoResponse(BaseModel): super().__init__(**kwargs) -class ListGuardrailsResponse(BaseModel): +class ListGuardrailsResponse(LiteLLMBaseModel): guardrails: list[GuardrailInfoResponse] -class GuardrailUIAddGuardrailSettings(BaseModel): +class GuardrailUIAddGuardrailSettings(LiteLLMBaseModel): supported_entities: list[str] supported_actions: list[str] supported_modes: list[str] supported_modes_by_provider: dict[str, list[str]] + providers_without_directional_logging_only_scope: tuple[str, ...] pii_entity_categories: list[PiiEntityCategoryMap] content_filter_settings: dict[str, object] | None = None + decision_model_providers: tuple[str, ...] = Field(default_factory=tuple) -class PresidioPerRequestConfig(BaseModel): +class PresidioPerRequestConfig(LiteLLMBaseModel): """ presdio params that can be controlled per request, api key """ @@ -1323,7 +1476,7 @@ class PresidioPerRequestConfig(BaseModel): entities: list[PiiEntityType] | None = None -class ApplyGuardrailRequest(BaseModel): +class ApplyGuardrailRequest(LiteLLMBaseModel): guardrail_name: str text: str language: str | None = None @@ -1333,11 +1486,11 @@ class ApplyGuardrailRequest(BaseModel): metadata: dict[str, object] | None = None -class ApplyGuardrailResponse(BaseModel): +class ApplyGuardrailResponse(LiteLLMBaseModel): response_text: str -class PatchGuardrailRequest(BaseModel): +class PatchGuardrailRequest(LiteLLMBaseModel): guardrail_name: str | None = None litellm_params: BaseLitellmParams | None = None guardrail_info: dict[str, object] | None = None From a8cb6f98401c7de4a236146769a40544f76b35ac Mon Sep 17 00:00:00 2001 From: AlphaRex-pixel Date: Sun, 11 Oct 2026 22:28:55 +0530 Subject: [PATCH 4/7] style(guardrails): fix line endings for ruff format --- litellm/types/guardrails.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index cbc54015aca..e6ae818e462 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -153,7 +153,7 @@ class SupportedGuardrailIntegrations(Enum): LLM_SHIELD_PROXY = "llm_shield_proxy" CONDUCT = "conduct" DECISION_MODEL = "decision_model" - TEALTIGER = 'tealtiger' + TEALTIGER = "tealtiger" class Role(Enum): From f299dd191e75d3693669f453c510d0bbe38bb5eb Mon Sep 17 00:00:00 2001 From: AlphaRex-pixel Date: Sun, 11 Oct 2026 22:47:01 +0530 Subject: [PATCH 5/7] fix(guardrails): remove new basedpyright unknown-type errors --- .../guardrail_hooks/tealtiger/__init__.py | 2 +- .../guardrails/guardrail_hooks/tealtiger/engine.py | 14 +++++++++----- .../guardrail_hooks/tealtiger/tealtiger.py | 12 ++++++------ 3 files changed, 16 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py index 110a714a191..84d1cc54546 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py @@ -14,7 +14,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" _tealtiger_callback: Final = TealTigerGuardrail( policies=getattr(litellm_params, "policies", None), policy_mode=getattr(litellm_params, "policy_mode", "ENFORCE"), - guardrail_name=guardrail.get("guardrail_name", ""), + 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/engine.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py index acb82d159ec..a84a49b9470 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py @@ -94,9 +94,13 @@ class TealEngine: scope: Final = self._pii_policy.get("patterns", "all") if scope == "all": return self.patterns - entries: Final = (scope,) if isinstance(scope, str) else scope - if not isinstance(entries, (list, tuple)): - raise TypeError("TealTiger: pii policy `patterns` must be 'all' or a list of names") + 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}) @@ -145,12 +149,12 @@ class TealEngine: 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 () + 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 () + 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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py index ead4955fd17..768d7c0909a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py @@ -19,7 +19,7 @@ from typing import TYPE_CHECKING, Final, Literal from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, - log_guardrail_information, + log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # upstream decorator is untyped ) from litellm.types.guardrails import GuardrailEventHooks, Mode from litellm.types.utils import GenericGuardrailAPIInputs @@ -29,8 +29,6 @@ from .engine import Action, PolicyMode, TealEngine if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -_EMPTY_MAPPING: Final = MappingProxyType({}) - DEFAULT_POLICIES: Final = ( MappingProxyType({"type": "pii", "action": "REDACT", "patterns": "all"}), MappingProxyType({"type": "cost", "action": "ENFORCE", "daily_limit_usd": 50.0}), @@ -64,7 +62,9 @@ def _tool_call_name_from_mapping(tool_call: Mapping[str, object]) -> str | None: if isinstance(direct_name, str): return direct_name function: Final = tool_call.get("function") - nested_name: Final = function.get("name") if isinstance(function, Mapping) else None + 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 @@ -86,7 +86,7 @@ def _tool_call_name(tool_call: object) -> str | None: (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) + return _tool_call_name_from_mapping(tool_call) # pyright: ignore[reportUnknownArgumentType] # Mapping[Unknown, Unknown] from isinstance return _tool_call_name_from_object(tool_call) @@ -114,7 +114,7 @@ class TealTigerGuardrail(CustomGuardrail): policies=policies or DEFAULT_POLICIES, mode=PolicyMode(policy_mode), ) - super().__init__( + 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, From c63a6d86c176d18b222cce3207d47864aa3f3ab7 Mon Sep 17 00:00:00 2001 From: AlphaRex-pixel Date: Mon, 12 Oct 2026 22:54:54 +0530 Subject: [PATCH 6/7] Rewrite TealTiger guardrail: per-user budgets, precise PII defaults, typed policies, decision receipts --- .../guardrail_hooks/tealtiger/__init__.py | 8 +- .../guardrail_hooks/tealtiger/budget.py | 38 + .../guardrail_hooks/tealtiger/config.py | 34 + .../guardrail_hooks/tealtiger/decision.py | 47 + .../guardrail_hooks/tealtiger/engine.py | 253 +++--- .../guardrail_hooks/tealtiger/patterns.py | 64 +- .../guardrail_hooks/tealtiger/request.py | 61 ++ .../guardrail_hooks/tealtiger/tealtiger.py | 162 ++-- .../guardrails/guardrail_hooks/tealtiger.py | 65 ++ .../guardrail_hooks/test_tealtiger.py | 804 ++++++++++-------- 10 files changed, 916 insertions(+), 620 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/tealtiger/budget.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/tealtiger/config.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/tealtiger/decision.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/tealtiger/request.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/tealtiger.py 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 From 717d1dd14c67dced150e35977037886ab974833d Mon Sep 17 00:00:00 2001 From: AlphaRex-pixel Date: Mon, 12 Oct 2026 23:05:50 +0530 Subject: [PATCH 7/7] Fix import ordering in TealTiger guardrail --- .../guardrails/guardrail_hooks/tealtiger/tealtiger.py | 7 ++++--- .../proxy/guardrails/guardrail_hooks/test_tealtiger.py | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py index b7a4851f7a1..45f45b1f912 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py @@ -14,24 +14,25 @@ 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.utils import GenericGuardrailAPIInputs - 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 - from pydantic import BaseModel def _to_event_hook( diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py index e991b8fa2a1..aaf564b3c8c 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tealtiger.py @@ -6,7 +6,6 @@ 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 import budget as budget_module @@ -21,6 +20,7 @@ from litellm.proxy.guardrails.guardrail_hooks.tealtiger.request import ( 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,