diff --git a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py index bef25066fa9..d87316d5d57 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/__init__.py @@ -1,3 +1,4 @@ +from types import MappingProxyType from typing import TYPE_CHECKING, Final from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -23,10 +24,14 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" return _tealtiger_callback -guardrail_initializer_registry: Final = { - SupportedGuardrailIntegrations.TEALTIGER.value: initialize_guardrail, -} +guardrail_initializer_registry: Final = MappingProxyType( + { + SupportedGuardrailIntegrations.TEALTIGER.value: initialize_guardrail, + } +) -guardrail_class_registry: Final = { - SupportedGuardrailIntegrations.TEALTIGER.value: TealTigerGuardrail, -} +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 index 3d0f075f197..c22b5e658d8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/engine.py @@ -7,9 +7,12 @@ from __future__ import annotations import threading import uuid +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field -from datetime import datetime +from datetime import date from enum import Enum +from types import MappingProxyType +from typing import Final from .patterns import default_patterns @@ -25,51 +28,61 @@ class Action(str, Enum): BLOCK = "BLOCK" -@dataclass +@dataclass(frozen=True) class Decision: action: str reason_code: str correlation_id: str redacted_text: str | None = None - details: dict = field(default_factory=dict) + details: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) - def to_receipt(self) -> dict: - return { - "correlation_id": self.correlation_id, - "action": self.action, - "reason_code": self.reason_code, - "details": self.details, - } + 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: list[dict], mode: PolicyMode = PolicyMode.ENFORCE) -> None: + 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] = {} + 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._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()) + correlation_id: Final = 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)] + 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 = self._pii_policy.get("action", "REDACT") + 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={"types": findings}, + details=MappingProxyType({"types": findings}), ) if self.mode == PolicyMode.MONITOR: @@ -77,25 +90,29 @@ class TealEngine: Action.ALLOW.value, f"PII_DETECTED:{','.join(findings)}", correlation_id, - details={"types": findings}, + details=MappingProxyType({"types": findings}), ) - redacted = text - for name in findings: - redacted = self.patterns[name].sub(f"[REDACTED:{name.upper()}]", redacted) + redacted: Final = self._redact(text, findings) return Decision( Action.REDACT.value, f"PII_DETECTED:{','.join(findings)}", correlation_id, redacted_text=redacted, - details={"types": findings}, + 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 - allowlist = self._tool_policy.get("allowlist") - blocklist = self._tool_policy.get("blocklist") + allowlist: Final = self._tool_policy.get("allowlist") + blocklist: Final = self._tool_policy.get("blocklist") if blocklist and tool_name in blocklist: return False if allowlist is not None: @@ -105,17 +122,17 @@ class TealEngine: 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") + limit: Final = self._cost_policy.get("daily_limit_usd") if limit is None: return False, 0.0, 0.0 - today = datetime.now().astimezone().date().isoformat() + today: Final = date.today().isoformat() # noqa: DTZ011 # daily budget boundary is intentionally local-server-time, not UTC with self._lock: - spent = self._spend_by_day.get(today, 0.0) + 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 = (tokens / 1000.0) * cost_per_1k_tokens - today = datetime.now().astimezone().date().isoformat() + 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 index c386e06bb1d..98dc8ecf519 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/patterns.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/patterns.py @@ -1,7 +1,7 @@ """ Regex-based PII / secrets detection patterns for TealTiger. -40+ patterns across five categories: government IDs, financial data, +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 @@ -11,125 +11,154 @@ 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 = { - # ---------- 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"), -} +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() -> dict: - """Return a fresh copy of the built-in pattern set.""" - return dict(PII_PATTERNS) +def default_patterns() -> Mapping[str, re.Pattern[str]]: + """Return the built-in, immutable pattern mapping.""" + return 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", - ], -} +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 index 38e7925cd4b..a0e5ef03336 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tealtiger/tealtiger.py @@ -3,14 +3,18 @@ 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 +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 typing import TYPE_CHECKING, Final, Literal, Optional +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, @@ -23,19 +27,23 @@ 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 -] +_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 +) class TealTigerGuardrail(CustomGuardrail): def __init__( self, - policies: list | None = None, + policies: Sequence[Mapping[str, object]] | None = None, policy_mode: str = "ENFORCE", - **kwargs: object, + **kwargs, # kwargs-ok: forwards unknown/future CustomGuardrail.__init__ params (guardrail_name, event_hook, default_on, ...) to the base class ) -> None: self.engine: Final = TealEngine( policies=policies or DEFAULT_POLICIES, @@ -47,9 +55,9 @@ class TealTigerGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: Mapping[str, object], input_type: Literal["request", "response"], - logging_obj: Optional["LiteLLMLoggingObj"] = None, + logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: """ Runs on both request (pre-call) and response (post-call) content, @@ -57,28 +65,43 @@ class TealTigerGuardrail(CustomGuardrail): 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": - 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})") + self._check_tool_calls(inputs) + self._check_budget(request_data) - 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: Final = tuple( + self._check_text_or_raise(text) for text in inputs.get("texts") or () + ) + return {**inputs, "texts": list(checked_texts)} # mutable-ok: dict interop - checked_texts: list[str] = [] - for text in inputs.get("texts") or []: - decision = self.engine.evaluate_text(text) + def _check_tool_calls(self, inputs: GenericGuardrailAPIInputs) -> None: + for tool_call in inputs.get("tool_calls") or (): + tool_name = tool_call.get("name") or tool_call.get( + "function", _EMPTY_MAPPING + ).get("name") + if tool_name and not self.engine.check_tool(tool_name): + raise ValueError( + f"TealTiger: blocked — TOOL_NOT_ALLOWLISTED ({tool_name})" + ) - if decision.action == Action.BLOCK.value: - raise ValueError(f"TealTiger: blocked — {decision.reason_code}") + def _check_budget(self, request_data: Mapping[str, object]) -> None: + 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.append(decision.redacted_text or text) - - inputs["texts"] = checked_texts - return inputs + 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` /