From 7897f5c497368356a06f7f5c439e1f440482a751 Mon Sep 17 00:00:00 2001 From: Oliver Fei Date: Tue, 20 Oct 2026 14:44:41 -0400 Subject: [PATCH] fix(guardrails): align TrendAI with current lint rules --- .../guardrail_hooks/trendai/__init__.py | 14 +++----- .../guardrail_hooks/trendai/_models.py | 22 ++++++------- .../guardrail_hooks/trendai/_text.py | 2 +- .../guardrail_hooks/trendai/trendai.py | 32 ++++++++++--------- 4 files changed, 34 insertions(+), 36 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py index 518089f0d02..14faca620ae 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py @@ -11,12 +11,12 @@ if TYPE_CHECKING: def _normalize_event_hook( - mode: str | list[str] | Mode, -) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode: + mode: str | list[str] | Mode, # mutable-ok: guardrail configuration accepts lists of modes +) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode: # mutable-ok: guardrail event hooks require a list if isinstance(mode, str): return GuardrailEventHooks(mode) if isinstance(mode, list): - return [GuardrailEventHooks(item) for item in mode] # mutable-ok: guardrail event hooks require a list + return [GuardrailEventHooks(item) for item in mode] return mode @@ -41,11 +41,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" return callback -guardrail_initializer_registry: Final = { # mutable-ok: registry discovery requires a dict - GUARDRAIL_NAME: initialize_guardrail -} -guardrail_class_registry: Final = { # mutable-ok: registry discovery requires a dict - GUARDRAIL_NAME: TrendAIGuardrail -} +guardrail_initializer_registry: Final = {GUARDRAIL_NAME: initialize_guardrail} +guardrail_class_registry: Final = {GUARDRAIL_NAME: TrendAIGuardrail} __all__ = ("TrendAIGuardrail",) diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py index 81759370f56..9ba9ed2e10a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py @@ -8,7 +8,7 @@ from typing import Literal, TypeAlias from pydantic import BaseModel, ConfigDict, Field -class TrendAISettings(BaseModel): +class TrendAISettings(BaseModel, frozen=True): app_name: str | None = None fallback_on_error: Literal["block", "allow"] = "block" timeout: float = 5.0 @@ -16,18 +16,18 @@ class TrendAISettings(BaseModel): response_content_chunk_size_bytes: int = 49_500 -class TrendAIChatMessage(BaseModel): +class TrendAIChatMessage(BaseModel, frozen=True): role: Literal["assistant"] = "assistant" content: str -class TrendAIChatChoice(BaseModel): +class TrendAIChatChoice(BaseModel, frozen=True): index: int = 0 message: TrendAIChatMessage finish_reason: Literal["stop"] = "stop" -class TrendAIChatCompletionPayload(BaseModel): +class TrendAIChatCompletionPayload(BaseModel, frozen=True): """The OpenAI chat-completion shape Trend AI Guard scans model output as.""" id: str = "chatcmpl-stream" @@ -41,44 +41,44 @@ class TrendAIChatCompletionPayload(BaseModel): return cls(model=model, choices=(TrendAIChatChoice(message=TrendAIChatMessage(content=content)),)) -class TrendAIRedactedMessage(BaseModel): +class TrendAIRedactedMessage(BaseModel, frozen=True): model_config = ConfigDict(extra="ignore") content: str | None = None -class TrendAIRedactedChoice(BaseModel): +class TrendAIRedactedChoice(BaseModel, frozen=True): model_config = ConfigDict(extra="ignore") message: TrendAIRedactedMessage | None = None -class TrendAIRedactedResponse(BaseModel): +class TrendAIRedactedResponse(BaseModel, frozen=True): model_config = ConfigDict(extra="ignore") choices: tuple[TrendAIRedactedChoice, ...] = () -class TrendAIRedactedPrompt(BaseModel): +class TrendAIRedactedPrompt(BaseModel, frozen=True): model_config = ConfigDict(extra="ignore") prompt: str | None = None -class TrendAISensitiveRule(BaseModel): +class TrendAISensitiveRule(BaseModel, frozen=True): model_config = ConfigDict(extra="ignore") id: str = "" -class TrendAISensitiveInformation(BaseModel): +class TrendAISensitiveInformation(BaseModel, frozen=True): model_config = ConfigDict(extra="ignore") has_policy_violation: bool = Field(default=False, alias="hasPolicyViolation") rules: tuple[TrendAISensitiveRule, ...] = () -class TrendAIResponse(BaseModel): +class TrendAIResponse(BaseModel, frozen=True): model_config = ConfigDict(extra="ignore") action: str diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py index 522c081dc48..19e655b7302 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py @@ -62,7 +62,7 @@ def apply_window_redaction(content: str, window: TrendAITextWindow, redacted: st def _last_user_text_parts(structured_messages: Sequence[AllMessageValues]) -> tuple[str, ...] | None: last_user_message: Final = next( - (message for message in reversed(structured_messages) if message.get("role") == "user"), None + (message for message in reversed(structured_messages) if message["role"] == "user"), None ) return message_slot_texts(last_user_message) if last_user_message is not None else None diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py index eb84eda3550..d2d70cabaf6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py @@ -5,6 +5,7 @@ import asyncio import os import time from collections.abc import AsyncIterator, Mapping, Sequence +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar, Final, Literal, NoReturn, Protocol from urllib.parse import SplitResult, urlsplit, urlunsplit @@ -61,8 +62,8 @@ class _AsyncHTTPClient(Protocol): self, url: str, *, - json: dict[str, object] | None = None, - headers: dict[str, str] | None = None, + json: dict[str, object] | None = None, # mutable-ok: injected HTTP handler requires dict payloads + headers: dict[str, str] | None = None, # mutable-ok: injected HTTP handler requires dict headers timeout: float | httpx.Timeout | None = None, ) -> httpx.Response: ... @@ -81,7 +82,9 @@ class TrendAIGuardrail(CustomGuardrail): response_content_chunk_size_bytes: int = RESPONSE_CONTENT_CHUNK_SIZE_BYTES, async_handler: _AsyncHTTPClient | None = None, guardrail_name: str | None = None, - event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None, + event_hook: ( # mutable-ok: CustomGuardrail accepts lists of hooks + GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None + ) = None, default_on: bool = False, ) -> None: resolved_api_key: Final = api_key or os.environ.get("TMV1_API_KEY") @@ -109,7 +112,7 @@ class TrendAIGuardrail(CustomGuardrail): self.api_url: str = _build_apply_guardrails_url(resolved_api_base) self.app_name: str = app_name or os.environ.get("TMV1_APPLICATION_NAME", "litellm") self.fallback_on_error: Literal["block", "allow"] = fallback_on_error - self.timeout: float = timeout + self.timeout = timeout self.stream_overlap_size: int = stream_overlap_size self.response_content_chunk_size_bytes: int = response_content_chunk_size_bytes self.async_handler: _AsyncHTTPClient = async_handler or get_async_httpx_client( @@ -125,7 +128,7 @@ class TrendAIGuardrail(CustomGuardrail): @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: # mutable-ok: CustomGuardrail contract - return [ # mutable-ok: CustomGuardrail contract + return [ GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.post_call, @@ -167,7 +170,7 @@ class TrendAIGuardrail(CustomGuardrail): return inputs redacted_inputs: Final[GenericGuardrailAPIInputs] = { **inputs, - "texts": list(redact_request_texts(texts, prompt, redacted)), # mutable-ok: GenericGuardrailAPIInputs field + "texts": list(redact_request_texts(texts, prompt, redacted)), } return redacted_inputs @@ -183,7 +186,7 @@ class TrendAIGuardrail(CustomGuardrail): model: Final = inputs.get("model") or _RESPONSE_MODEL started_at: Final = time.time() scans: Final = tuple([await self._scan_response_windows(text, model) for text in texts]) - all_scans: Final = tuple(scan for text_scans in scans for scan in text_scans) + all_scans: Final = tuple(chain.from_iterable(scans)) verdict: Final = _window_verdict(all_scans) self._record_scan( verdict, @@ -201,7 +204,7 @@ class TrendAIGuardrail(CustomGuardrail): self._raise_redaction_failure(request_data, started_at, event_type=None) redacted_inputs: Final[GenericGuardrailAPIInputs] = { **inputs, - "texts": list(merged_texts), # mutable-ok: GenericGuardrailAPIInputs field + "texts": list(merged_texts), } return redacted_inputs @@ -339,14 +342,14 @@ class TrendAIGuardrail(CustomGuardrail): ) -> None: ended_at: Final = time.time() self.add_standard_logging_guardrail_information_to_request_data( # pyright: ignore[reportUnknownMemberType] # base signature takes an untyped dict - guardrail_json_response=dict(guardrail_json_response), # mutable-ok: base signature takes a plain dict + guardrail_json_response=dict(guardrail_json_response), request_data=request_data, guardrail_status=guardrail_status, start_time=started_at, end_time=ended_at, duration=ended_at - started_at, event_type=event_type, - masked_entity_count=dict(masked_entity_count) or None, # mutable-ok: base signature takes a plain dict + masked_entity_count=dict(masked_entity_count) or None, ) def _build_request_headers(self, request_type: str | None = None) -> Mapping[str, str]: @@ -370,8 +373,8 @@ class TrendAIGuardrail(CustomGuardrail): try: response: Final = await self.async_handler.post( self.api_url, - json=dict(payload), # mutable-ok: httpx takes a plain dict - headers=dict(self._build_request_headers(request_type)), # mutable-ok: httpx takes a plain dict + json=dict(payload), + headers=dict(self._build_request_headers(request_type)), timeout=self.timeout, ) response.raise_for_status() @@ -436,9 +439,8 @@ def _window_verdict(scans: Sequence[TrendAIWindowScan]) -> TrendAIScanResult: terminal: Final = next((scan.result for scan in scans if not isinstance(scan.result, TrendAIAllow)), None) if terminal is not None: return terminal - counts: Final = tuple( - pair for scan in scans if isinstance(scan.result, TrendAIAllow) for pair in scan.result.masked_entity_count - ) + allowed_counts: Final = (scan.result.masked_entity_count for scan in scans if isinstance(scan.result, TrendAIAllow)) + counts: Final = tuple(chain.from_iterable(allowed_counts)) return TrendAIAllow(masked_entity_count=tuple(sorted(_merge_entity_counts(_NO_ENTITIES, counts).items())))