scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-01 07:03:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

fix(guardrails): align TrendAI with current lint rules

Browse source
This commit is contained in:
Oliver Fei 2026-10-20 14:44:41 -04:00
parent a022e9e8aa
commit 7897f5c497
4 changed files with 34 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py
View file

@ -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",)

22
litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py
View file

@ -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

2
litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py
View file

@ -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

32
litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py
View file

@ -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())))