From 01886ffc7811b281ed0add5903a9a8b97a728e47 Mon Sep 17 00:00:00 2001 From: AlphaRex-pixel Date: Fri, 14 Aug 2026 15:40:34 +0530 Subject: [PATCH] Add TealTiger guardrail: self-contained PII/cost/tool-auth policy engine --- docs/tealtiger.md | 60 +++++++++ .../guardrail_hooks/tealtiger/__init__.py | 32 +++++ .../guardrail_hooks/tealtiger/engine.py | 121 ++++++++++++++++++ .../guardrail_hooks/tealtiger/patterns.py | 104 +++++++++++++++ .../guardrail_hooks/tealtiger/tealtiger.py | 95 ++++++++++++++ litellm/types/guardrails.py | 1 + tests/guardrails_tests/test_tealtiger.py | 95 ++++++++++++++ 7 files changed, 508 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/guardrails_tests/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..bef25066fa9 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py @@ -0,0 +1,32 @@ +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"): + 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 = { + SupportedGuardrailIntegrations.TEALTIGER.value: initialize_guardrail, +} + +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 new file mode 100644 index 00000000000..00fccbf33c6 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py @@ -0,0 +1,121 @@ +""" +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 dataclasses import dataclass, field +from datetime import date +from enum import Enum +from typing import Optional + +from .patterns import default_patterns + + +class PolicyMode(str, Enum): + ENFORCE = "ENFORCE" + MONITOR = "MONITOR" + + +class Action(str, Enum): + ALLOW = "ALLOW" + REDACT = "REDACT" + BLOCK = "BLOCK" + + +@dataclass +class Decision: + action: str + reason_code: str + correlation_id: str + redacted_text: Optional[str] = None + details: dict = field(default_factory=dict) + + def to_receipt(self) -> dict: + return { + "correlation_id": self.correlation_id, + "action": self.action, + "reason_code": self.reason_code, + "details": self.details, + } + + +class TealEngine: + def __init__(self, policies: list[dict], mode: PolicyMode = PolicyMode.ENFORCE): + self.policies = policies + self.mode = mode + self.patterns = default_patterns() + self._lock = threading.Lock() + self._spend_by_day: dict[str, float] = {} + + 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 = str(uuid.uuid4()) + if not text or not self._pii_policy: + return Decision(Action.ALLOW.value, "NO_VIOLATIONS", correlation_id) + + findings = [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 = 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={"types": findings}, + ) + + if self.mode == PolicyMode.MONITOR: + return Decision( + Action.ALLOW.value, + f"PII_DETECTED:{','.join(findings)}", + correlation_id, + details={"types": findings}, + ) + + redacted = text + for name in findings: + redacted = self.patterns[name].sub(f"[REDACTED:{name.upper()}]", redacted) + return Decision( + Action.REDACT.value, + f"PII_DETECTED:{','.join(findings)}", + correlation_id, + redacted_text=redacted, + details={"types": findings}, + ) + + def check_tool(self, tool_name: str) -> bool: + if not self._tool_policy: + return True + allowlist = self._tool_policy.get("allowlist") + blocklist = self._tool_policy.get("blocklist") + if blocklist and tool_name in blocklist: + return False + if allowlist is not None: + 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 + limit = self._cost_policy.get("daily_limit_usd") + if limit is None: + return False, 0.0, 0.0 + today = date.today().isoformat() + with self._lock: + spent = 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 = (tokens / 1000.0) * cost_per_1k_tokens + today = date.today().isoformat() + 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..f8e48b36803 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/patterns.py @@ -0,0 +1,104 @@ +""" +Regex-based PII / secrets detection patterns for TealTiger. + +40+ 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. +""" +import re + +PII_PATTERNS = { + # ---------- 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() -> dict: + """Return a fresh copy of the built-in pattern set.""" + return dict(PII_PATTERNS) + + +PATTERN_CATEGORIES = { + "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..3e1f1d6d69b --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py @@ -0,0 +1,95 @@ +""" +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.96.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 typing import TYPE_CHECKING, Final, List, Literal, Optional + +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +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 + +DEFAULT_POLICIES = [ + {"type": "pii", "action": "REDACT", "patterns": "all"}, + {"type": "cost", "action": "ENFORCE", "daily_limit_usd": 50.0}, + {"type": "tool_auth", "action": "ENFORCE", "allowlist": None}, # None = allow all +] + + +class TealTigerGuardrail(CustomGuardrail): + def __init__( + self, + policies: Optional[list] = None, + policy_mode: str = "ENFORCE", + **kwargs, + ): + self.engine: Final = TealEngine( + policies=policies or DEFAULT_POLICIES, + mode=PolicyMode(policy_mode), + ) + super().__init__(**kwargs) + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = 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". + """ + if input_type == "request": + for tool_call in inputs.get("tool_calls") or []: + tool_name = tool_call.get("name") or tool_call.get("function", {}).get("name") + if tool_name and not self.engine.check_tool(tool_name): + raise ValueError(f"TealTiger: blocked — TOOL_NOT_ALLOWLISTED ({tool_name})") + + over_budget, spent, limit = self.engine.check_budget( + session_id=request_data.get("user") or "default" + ) + if over_budget: + raise ValueError( + f"TealTiger: blocked — DAILY_BUDGET_EXCEEDED (${spent:.2f} / ${limit:.2f})" + ) + + checked_texts: List[str] = [] + for text in inputs.get("texts") or []: + decision = self.engine.evaluate_text(text) + + if decision.action == Action.BLOCK.value: + raise ValueError(f"TealTiger: blocked — {decision.reason_code}") + + checked_texts.append(decision.redacted_text or text) + + inputs["texts"] = checked_texts + return inputs + + # ---- 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 bbb6d758814..f2ecce85ecd 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -134,6 +134,7 @@ class SupportedGuardrailIntegrations(Enum): HEADROOM = "headroom" COMPRESR = "compresr" STRAIKER = "straiker" + TEALTIGER = "tealtiger" class Role(Enum): diff --git a/tests/guardrails_tests/test_tealtiger.py b/tests/guardrails_tests/test_tealtiger.py new file mode 100644 index 00000000000..88916710590 --- /dev/null +++ b/tests/guardrails_tests/test_tealtiger.py @@ -0,0 +1,95 @@ +""" +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", + )