mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-30 06:43:27 +00:00
Merge 15a5371e06 into bc18e10f6b
This commit is contained in:
Oliver Fei
GitHub
commit
c1a0741143
6 changed files with 1503 additions and 0 deletions
47
litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py
Normal file
47
litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import litellm
|
||||
from litellm.types.guardrails import GuardrailEventHooks, Mode
|
||||
|
||||
from ._models import TrendAISettings
|
||||
from .trendai import GUARDRAIL_NAME, TrendAIGuardrail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
def _normalize_event_hook(
|
||||
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]
|
||||
return mode
|
||||
|
||||
|
||||
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> TrendAIGuardrail:
|
||||
settings: Final = TrendAISettings.model_validate(litellm_params.model_dump(mode="python", exclude_none=True))
|
||||
guardrail_name: Final = guardrail["guardrail_name"]
|
||||
callback: Final = TrendAIGuardrail(
|
||||
api_key=litellm_params.api_key,
|
||||
api_base=litellm_params.api_base,
|
||||
app_name=settings.app_name,
|
||||
fallback_on_error=settings.fallback_on_error,
|
||||
timeout=settings.timeout,
|
||||
stream_overlap_size=settings.stream_overlap_size,
|
||||
response_content_chunk_size_bytes=settings.response_content_chunk_size_bytes,
|
||||
guardrail_name=guardrail_name,
|
||||
event_hook=_normalize_event_hook(litellm_params.mode),
|
||||
default_on=litellm_params.default_on is True,
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] # callback union is partially untyped
|
||||
callback
|
||||
)
|
||||
return callback
|
||||
|
||||
|
||||
guardrail_initializer_registry: Final = {GUARDRAIL_NAME: initialize_guardrail}
|
||||
guardrail_class_registry: Final = {GUARDRAIL_NAME: TrendAIGuardrail}
|
||||
|
||||
__all__ = ("TrendAIGuardrail",)
|
||||
130
litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py
Normal file
130
litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py
Normal file
|
|
@ -0,0 +1,130 @@
|
|||
# Derived from tm-v1-ai-guard-litellm-plugin revision 6cafc143f62962a98d4eb5abe9f608c61ff194d4.
|
||||
# This file has been modified for integration into LiteLLM.
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class TrendAISettings(BaseModel, frozen=True):
|
||||
app_name: str | None = None
|
||||
fallback_on_error: Literal["block", "allow"] = "block"
|
||||
timeout: float = 5.0
|
||||
stream_overlap_size: int = 256
|
||||
response_content_chunk_size_bytes: int = 49_500
|
||||
|
||||
|
||||
class TrendAIChatMessage(BaseModel, frozen=True):
|
||||
role: Literal["assistant"] = "assistant"
|
||||
content: str
|
||||
|
||||
|
||||
class TrendAIChatChoice(BaseModel, frozen=True):
|
||||
index: int = 0
|
||||
message: TrendAIChatMessage
|
||||
finish_reason: Literal["stop"] = "stop"
|
||||
|
||||
|
||||
class TrendAIChatCompletionPayload(BaseModel, frozen=True):
|
||||
"""The OpenAI chat-completion shape Trend AI Guard scans model output as."""
|
||||
|
||||
id: str = "chatcmpl-stream"
|
||||
object: Literal["chat.completion"] = "chat.completion"
|
||||
created: int = 0
|
||||
model: str
|
||||
choices: tuple[TrendAIChatChoice, ...]
|
||||
|
||||
@classmethod
|
||||
def for_content(cls, content: str, model: str) -> "TrendAIChatCompletionPayload":
|
||||
return cls(model=model, choices=(TrendAIChatChoice(message=TrendAIChatMessage(content=content)),))
|
||||
|
||||
|
||||
class TrendAIRedactedMessage(BaseModel, frozen=True):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
content: str | None = None
|
||||
|
||||
|
||||
class TrendAIRedactedChoice(BaseModel, frozen=True):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
message: TrendAIRedactedMessage | None = None
|
||||
|
||||
|
||||
class TrendAIRedactedResponse(BaseModel, frozen=True):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
choices: tuple[TrendAIRedactedChoice, ...] = ()
|
||||
|
||||
|
||||
class TrendAIRedactedPrompt(BaseModel, frozen=True):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
prompt: str | None = None
|
||||
|
||||
|
||||
class TrendAISensitiveRule(BaseModel, frozen=True):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
id: str = ""
|
||||
|
||||
|
||||
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, frozen=True):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
action: str
|
||||
reasons: tuple[str, ...] = ()
|
||||
reason: str = ""
|
||||
redacted_request: Mapping[str, object] | None = Field(default=None, alias="redactedRequest")
|
||||
sensitive_information: TrendAISensitiveInformation | None = Field(default=None, alias="sensitiveInformation")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TrendAIAllow:
|
||||
kind: Literal["allow"] = "allow"
|
||||
redacted_content: str | None = None
|
||||
masked_entity_count: tuple[tuple[str, int], ...] = ()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TrendAIBlock:
|
||||
reason: str
|
||||
status_code: int = 400
|
||||
kind: Literal["block"] = "block"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TrendAIProviderFailure:
|
||||
reason: str
|
||||
status_code: int | None = None
|
||||
kind: Literal["provider_failure"] = "provider_failure"
|
||||
|
||||
|
||||
TrendAIScanResult: TypeAlias = TrendAIAllow | TrendAIBlock | TrendAIProviderFailure
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TrendAITextWindow:
|
||||
text: str
|
||||
start: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TrendAIWindowScan:
|
||||
window: TrendAITextWindow
|
||||
result: TrendAIScanResult
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TrendAIRequestPrompt:
|
||||
prompt: str
|
||||
text_range: tuple[int, int] | None
|
||||
110
litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py
Normal file
110
litellm/proxy/guardrails/guardrail_hooks/trendai/_text.py
Normal file
|
|
@ -0,0 +1,110 @@
|
|||
# Derived from tm-v1-ai-guard-litellm-plugin revision 6cafc143f62962a98d4eb5abe9f608c61ff194d4.
|
||||
# This file has been modified for integration into LiteLLM.
|
||||
|
||||
from collections.abc import Iterator, Sequence
|
||||
from itertools import chain
|
||||
from typing import Final
|
||||
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts, message_text_slot_count
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
||||
from ._models import TrendAIRequestPrompt, TrendAITextWindow
|
||||
|
||||
|
||||
def utf8_windows(content: str, *, chunk_size_bytes: int, overlap_chars: int) -> tuple[TrendAITextWindow, ...]:
|
||||
"""Split ``content`` into windows of at most ``chunk_size_bytes`` UTF-8 bytes.
|
||||
|
||||
Consecutive windows share their last ``overlap_chars`` characters so a finding that straddles a
|
||||
window boundary is still seen whole by one scan. The overlap is capped below the window length so
|
||||
every window makes forward progress.
|
||||
"""
|
||||
return tuple(_iter_utf8_windows(content, chunk_size_bytes, overlap_chars))
|
||||
|
||||
|
||||
def _iter_utf8_windows(content: str, chunk_size_bytes: int, overlap_chars: int) -> Iterator[TrendAITextWindow]:
|
||||
encoded: Final = content.encode("utf-8")
|
||||
byte_offset = 0 # rebind-ok: window cursor advances across the loop
|
||||
char_offset = 0 # rebind-ok: window cursor advances across the loop
|
||||
while byte_offset < len(encoded):
|
||||
text = (
|
||||
encoded[byte_offset : byte_offset + chunk_size_bytes].decode("utf-8", errors="ignore")
|
||||
or content[char_offset]
|
||||
)
|
||||
yield TrendAITextWindow(text=text, start=char_offset)
|
||||
end_byte_offset = byte_offset + len(text.encode("utf-8"))
|
||||
if end_byte_offset >= len(encoded):
|
||||
return
|
||||
overlap = min(max(overlap_chars, 0), len(text) - 1)
|
||||
byte_offset = end_byte_offset - len(text[len(text) - overlap :].encode("utf-8"))
|
||||
char_offset += len(text) - overlap
|
||||
|
||||
|
||||
def merge_redaction(original: str, current: str, redacted: str) -> str | None:
|
||||
"""Apply the masks ``redacted`` adds over ``original`` on top of the masks already in ``current``.
|
||||
|
||||
Returns None when the three texts do not line up character for character, since a positional
|
||||
merge would then scramble the output.
|
||||
"""
|
||||
if len(original) != len(current) or len(redacted) != len(original):
|
||||
return None
|
||||
return "".join(
|
||||
redacted_char if redacted_char != original_char else current_char
|
||||
for original_char, current_char, redacted_char in zip(original, current, redacted, strict=True)
|
||||
)
|
||||
|
||||
|
||||
def apply_window_redaction(content: str, window: TrendAITextWindow, redacted: str) -> str | None:
|
||||
end: Final = window.start + len(window.text)
|
||||
merged: Final = merge_redaction(window.text, content[window.start : end], redacted)
|
||||
if merged is None:
|
||||
return None
|
||||
return f"{content[: window.start]}{merged}{content[end:]}"
|
||||
|
||||
|
||||
def locate_request_prompt(
|
||||
texts: Sequence[str],
|
||||
structured_messages: Sequence[AllMessageValues] | None,
|
||||
) -> TrendAIRequestPrompt | None:
|
||||
"""Locate the last user turn by message order, never by matching its text.
|
||||
|
||||
A different flat-text layout can still be scanned, but cannot be safely rewritten.
|
||||
"""
|
||||
messages: Final = tuple(structured_messages or ())
|
||||
user_index: Final = next(
|
||||
(index for index in reversed(range(len(messages))) if messages[index]["role"] == "user"), None
|
||||
)
|
||||
if messages and user_index is None:
|
||||
return None
|
||||
parts: Final = message_slot_texts(messages[user_index]) if user_index is not None else tuple(texts)
|
||||
prompt: Final = "".join(parts).strip()
|
||||
if not prompt:
|
||||
return None
|
||||
start: Final = sum(map(message_text_slot_count, messages[:user_index])) if user_index is not None else 0
|
||||
normalized_texts: Final = (
|
||||
tuple(chain.from_iterable(map(message_slot_texts, messages))) if messages else tuple(texts)
|
||||
)
|
||||
text_range: Final = (start, start + len(parts)) if normalized_texts == tuple(texts) else None
|
||||
return TrendAIRequestPrompt(prompt=prompt, text_range=text_range)
|
||||
|
||||
|
||||
def redact_request_texts(texts: Sequence[str], prompt: TrendAIRequestPrompt, redacted: str) -> tuple[str, ...] | None:
|
||||
"""Write a redacted prompt back over the ``texts`` entries the prompt was assembled from.
|
||||
|
||||
A single-part prompt is replaced outright. A multi-part prompt whose redaction kept its length
|
||||
is split back into the original parts; otherwise the first part carries the whole redaction and
|
||||
the rest are blanked, so no original part can leak.
|
||||
"""
|
||||
if prompt.text_range is None:
|
||||
return None
|
||||
start, end = prompt.text_range
|
||||
parts: Final = tuple(texts[start:end])
|
||||
if len(parts) == 1:
|
||||
return (*texts[:start], redacted, *texts[end:])
|
||||
joined: Final = "".join(parts)
|
||||
leading: Final = len(joined) - len(joined.lstrip())
|
||||
aligned: Final = f"{joined[:leading]}{redacted}{joined[len(joined.rstrip()) :]}"
|
||||
if len(aligned) != len(joined):
|
||||
return (*texts[:start], redacted, *("" for _ in parts[1:]), *texts[end:])
|
||||
offsets: Final = tuple(sum(len(part) for part in parts[:index]) for index in range(len(parts) + 1))
|
||||
split: Final = tuple(aligned[offsets[index] : offsets[index + 1]] for index in range(len(parts)))
|
||||
return (*texts[:start], *split, *texts[end:])
|
||||
509
litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py
Normal file
509
litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py
Normal file
|
|
@ -0,0 +1,509 @@
|
|||
# Derived from tm-v1-ai-guard-litellm-plugin revision 6cafc143f62962a98d4eb5abe9f608c61ff194d4.
|
||||
# This file has been modified for integration into LiteLLM.
|
||||
|
||||
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
|
||||
|
||||
import httpx
|
||||
from pydantic import ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._version import version as litellm_version
|
||||
from litellm.exceptions import GuardrailRaisedException, Timeout
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # legacy decorator has an untyped signature
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # legacy client factory has an untyped params map
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks, Mode
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus
|
||||
|
||||
from ._models import (
|
||||
TrendAIAllow,
|
||||
TrendAIBlock,
|
||||
TrendAIChatCompletionPayload,
|
||||
TrendAIProviderFailure,
|
||||
TrendAIRedactedPrompt,
|
||||
TrendAIRedactedResponse,
|
||||
TrendAIResponse,
|
||||
TrendAIScanResult,
|
||||
TrendAIWindowScan,
|
||||
)
|
||||
from ._text import apply_window_redaction, locate_request_prompt, redact_request_texts, utf8_windows
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
GUARDRAIL_NAME: Final = "trendai"
|
||||
OPENAI_CHAT_COMPLETION_RESPONSE_V1: Final = "OpenAIChatCompletionResponseV1"
|
||||
RESPONSE_CONTENT_CHUNK_SIZE_BYTES: Final = 49_500
|
||||
TMV1_CLIENT_NAME: Final = "litellm"
|
||||
PLUGIN_VERSION: Final = "v0.1.3"
|
||||
PROVIDER_UNAVAILABLE_STATUS: Final = 503
|
||||
_APPLY_GUARDRAILS_PATH: Final = "/applyGuardrails"
|
||||
_TREND_AI_SECURITY_PATH: Final = "/v3.0/aiSecurity"
|
||||
_RESPONSE_MODEL: Final = "guardrailed-response"
|
||||
_REDACTION_FAILED_MESSAGE: Final = "Trend AI Guard could not apply the requested redaction"
|
||||
_NO_ENTITIES: Final[Mapping[str, int]] = MappingProxyType({})
|
||||
|
||||
|
||||
class _AsyncHTTPClient(Protocol):
|
||||
async def post(
|
||||
self,
|
||||
url: str,
|
||||
*,
|
||||
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: ...
|
||||
|
||||
|
||||
class TrendAIGuardrail(CustomGuardrail):
|
||||
records_own_guardrail_information: ClassVar[bool] = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
app_name: str | None = None,
|
||||
fallback_on_error: Literal["block", "allow"] = "block",
|
||||
timeout: float = 5.0,
|
||||
stream_overlap_size: int = 256,
|
||||
response_content_chunk_size_bytes: int = RESPONSE_CONTENT_CHUNK_SIZE_BYTES,
|
||||
async_handler: _AsyncHTTPClient | None = None,
|
||||
guardrail_name: str | 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")
|
||||
if not resolved_api_key:
|
||||
raise ValueError(
|
||||
"Trend AI Guard requires an API key. Pass api_key or set the TMV1_API_KEY environment variable."
|
||||
)
|
||||
|
||||
resolved_api_base: Final = api_base or os.environ.get("TRENDAI_AI_GUARD_BASE_URL")
|
||||
if not resolved_api_base:
|
||||
raise ValueError(
|
||||
"Trend AI Guard requires an API base URL. Pass api_base or set the "
|
||||
"TRENDAI_AI_GUARD_BASE_URL environment variable."
|
||||
)
|
||||
if fallback_on_error not in ("block", "allow"):
|
||||
raise ValueError("fallback_on_error must be 'block' or 'allow'")
|
||||
if timeout <= 0:
|
||||
raise ValueError("timeout must be greater than zero")
|
||||
if stream_overlap_size < 0:
|
||||
raise ValueError("stream_overlap_size must be non-negative")
|
||||
if response_content_chunk_size_bytes < 1:
|
||||
raise ValueError("response_content_chunk_size_bytes must be greater than zero")
|
||||
|
||||
self.api_key: str = resolved_api_key
|
||||
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 = 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(
|
||||
llm_provider=httpxSpecialProvider.GuardrailCallback
|
||||
)
|
||||
|
||||
super().__init__( # pyright: ignore[reportUnknownMemberType] # base constructor retains untyped extension kwargs
|
||||
guardrail_name=guardrail_name,
|
||||
event_hook=event_hook,
|
||||
default_on=default_on,
|
||||
supported_event_hooks=self.get_supported_event_hooks(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: # mutable-ok: CustomGuardrail contract
|
||||
return [
|
||||
GuardrailEventHooks.pre_call,
|
||||
GuardrailEventHooks.during_call,
|
||||
GuardrailEventHooks.post_call,
|
||||
GuardrailEventHooks.logging_only,
|
||||
]
|
||||
|
||||
def supports_scan_only_tool_results(self) -> bool:
|
||||
return False
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object], # mutable-ok: CustomGuardrail hook contract is a plain dict
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: "LiteLLMLoggingObj | None" = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
texts: Final = tuple(inputs.get("texts") or ())
|
||||
match input_type:
|
||||
case "request":
|
||||
return await self._apply_request_guardrail(inputs, texts, request_data)
|
||||
case "response":
|
||||
return await self._apply_response_guardrail(inputs, texts, request_data)
|
||||
case _:
|
||||
assert_never(input_type)
|
||||
|
||||
async def _apply_request_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
texts: tuple[str, ...],
|
||||
request_data: dict[str, object], # mutable-ok: CustomGuardrail hook contract is a plain dict
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
prompt: Final = locate_request_prompt(texts, inputs.get("structured_messages"))
|
||||
if prompt is None:
|
||||
verbose_proxy_logger.debug("Trend AI Guard: no user prompt to scan in request inputs")
|
||||
return inputs
|
||||
started_at: Final = time.time()
|
||||
result: Final = await self._scan_payload(MappingProxyType({"prompt": prompt.prompt}))
|
||||
self._record_scan(result, request_data, started_at, event_type=None)
|
||||
redacted: Final = self._enforce(result)
|
||||
if redacted is None:
|
||||
return inputs
|
||||
redacted_texts: Final = redact_request_texts(texts, prompt, redacted)
|
||||
if redacted_texts is None:
|
||||
self._raise_redaction_failure(request_data, started_at, event_type=None)
|
||||
redacted_inputs: Final[GenericGuardrailAPIInputs] = {
|
||||
**inputs,
|
||||
"texts": list(redacted_texts),
|
||||
}
|
||||
return redacted_inputs
|
||||
|
||||
async def _apply_response_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
texts: tuple[str, ...],
|
||||
request_data: dict[str, object], # mutable-ok: CustomGuardrail hook contract is a plain dict
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
if not texts:
|
||||
verbose_proxy_logger.debug("Trend AI Guard: no response text to scan")
|
||||
return inputs
|
||||
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(chain.from_iterable(scans))
|
||||
verdict: Final = _window_verdict(all_scans)
|
||||
self._record_scan(
|
||||
verdict,
|
||||
request_data,
|
||||
started_at,
|
||||
event_type=None,
|
||||
redacted=any(_is_redacting(scan.result) for scan in all_scans),
|
||||
)
|
||||
self._enforce(verdict)
|
||||
redacted_texts: Final = tuple(
|
||||
_merge_window_redactions(text, text_scans) for text, text_scans in zip(texts, scans, strict=True)
|
||||
)
|
||||
merged_texts: Final = tuple(text for text in redacted_texts if text is not None)
|
||||
if len(merged_texts) != len(texts):
|
||||
self._raise_redaction_failure(request_data, started_at, event_type=None)
|
||||
redacted_inputs: Final[GenericGuardrailAPIInputs] = {
|
||||
**inputs,
|
||||
"texts": list(merged_texts),
|
||||
}
|
||||
return redacted_inputs
|
||||
|
||||
async def _scan_response_windows(self, content: str, model: str) -> tuple[TrendAIWindowScan, ...]:
|
||||
return tuple([scan async for scan in self._iter_response_window_scans(content, model)])
|
||||
|
||||
async def _iter_response_window_scans(self, content: str, model: str) -> AsyncIterator[TrendAIWindowScan]:
|
||||
"""Scan ``content`` window by window, stopping at the first block or provider failure."""
|
||||
for window in utf8_windows(
|
||||
content,
|
||||
chunk_size_bytes=self.response_content_chunk_size_bytes,
|
||||
overlap_chars=self.stream_overlap_size,
|
||||
):
|
||||
result = await self._scan_payload(
|
||||
_chat_completion_payload(window.text, model),
|
||||
request_type=OPENAI_CHAT_COMPLETION_RESPONSE_V1,
|
||||
)
|
||||
yield TrendAIWindowScan(window=window, result=result)
|
||||
if not isinstance(result, TrendAIAllow):
|
||||
return
|
||||
|
||||
def _raise_redaction_failure(
|
||||
self,
|
||||
request_data: dict[str, object], # mutable-ok: CustomGuardrail hook contract is a plain dict
|
||||
started_at: float,
|
||||
*,
|
||||
event_type: GuardrailEventHooks | None,
|
||||
) -> NoReturn:
|
||||
self._record_failure(_REDACTION_FAILED_MESSAGE, request_data, started_at, event_type=event_type)
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=_REDACTION_FAILED_MESSAGE,
|
||||
should_wrap_with_default_message=False,
|
||||
status_code=PROVIDER_UNAVAILABLE_STATUS,
|
||||
)
|
||||
|
||||
def _enforce(self, result: TrendAIScanResult) -> str | None:
|
||||
"""Raise for a block or a fail-closed provider failure; otherwise return any redacted content."""
|
||||
match result:
|
||||
case TrendAIAllow(redacted_content=redacted_content):
|
||||
return redacted_content
|
||||
case TrendAIBlock(reason=reason, status_code=status_code):
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=f"Blocked by Trend AI Guard. Security violation: {reason}",
|
||||
should_wrap_with_default_message=False,
|
||||
status_code=status_code,
|
||||
blocked_content=True,
|
||||
)
|
||||
case TrendAIProviderFailure(reason=reason, status_code=status_code):
|
||||
if self.fallback_on_error == "allow":
|
||||
verbose_proxy_logger.warning(
|
||||
"Trend AI Guard: %s (status=%s); allowing traffic (fallback_on_error=allow)",
|
||||
reason,
|
||||
status_code,
|
||||
)
|
||||
return None
|
||||
verbose_proxy_logger.error(
|
||||
"Trend AI Guard: %s (status=%s); blocking traffic (fallback_on_error=block)", reason, status_code
|
||||
)
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=f"Security Guard Error: {reason}",
|
||||
should_wrap_with_default_message=False,
|
||||
status_code=PROVIDER_UNAVAILABLE_STATUS,
|
||||
)
|
||||
case _:
|
||||
assert_never(result)
|
||||
|
||||
def _record_scan(
|
||||
self,
|
||||
result: TrendAIScanResult,
|
||||
request_data: dict[str, object], # mutable-ok: CustomGuardrail hook contract is a plain dict
|
||||
started_at: float,
|
||||
*,
|
||||
event_type: GuardrailEventHooks | None,
|
||||
redacted: bool | None = None,
|
||||
) -> None:
|
||||
match result:
|
||||
case TrendAIAllow(redacted_content=redacted_content, masked_entity_count=masked_entity_count):
|
||||
self._record(
|
||||
_allow_record(redacted_content is not None if redacted is None else redacted),
|
||||
"success",
|
||||
request_data,
|
||||
started_at,
|
||||
event_type=event_type,
|
||||
masked_entity_count=_merge_entity_counts(_NO_ENTITIES, masked_entity_count),
|
||||
)
|
||||
case TrendAIBlock(reason=reason):
|
||||
self._record(
|
||||
MappingProxyType({"action": "block", "reason": reason}),
|
||||
"guardrail_intervened",
|
||||
request_data,
|
||||
started_at,
|
||||
event_type=event_type,
|
||||
)
|
||||
case TrendAIProviderFailure(reason=reason, status_code=status_code):
|
||||
self._record(
|
||||
MappingProxyType(
|
||||
{"description": reason, "status_code": status_code, "fallback_on_error": self.fallback_on_error}
|
||||
),
|
||||
"guardrail_failed_to_respond",
|
||||
request_data,
|
||||
started_at,
|
||||
event_type=event_type,
|
||||
)
|
||||
case _:
|
||||
assert_never(result)
|
||||
|
||||
def _record_failure(
|
||||
self,
|
||||
description: str,
|
||||
request_data: dict[str, object], # mutable-ok: CustomGuardrail hook contract is a plain dict
|
||||
started_at: float,
|
||||
*,
|
||||
event_type: GuardrailEventHooks | None,
|
||||
) -> None:
|
||||
self._record(
|
||||
MappingProxyType({"description": description}),
|
||||
"guardrail_failed_to_respond",
|
||||
request_data,
|
||||
started_at,
|
||||
event_type=event_type,
|
||||
)
|
||||
|
||||
def _record(
|
||||
self,
|
||||
guardrail_json_response: Mapping[str, object],
|
||||
guardrail_status: GuardrailStatus,
|
||||
request_data: dict[str, object], # mutable-ok: CustomGuardrail hook contract is a plain dict
|
||||
started_at: float,
|
||||
*,
|
||||
event_type: GuardrailEventHooks | None,
|
||||
masked_entity_count: Mapping[str, int] = _NO_ENTITIES,
|
||||
) -> 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),
|
||||
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,
|
||||
)
|
||||
|
||||
def _build_request_headers(self, request_type: str | None = None) -> Mapping[str, str]:
|
||||
headers: Final = (
|
||||
("TMV1-Application-Name", self.app_name),
|
||||
("Authorization", f"Bearer {self.api_key}"),
|
||||
("Content-Type", "application/json"),
|
||||
("TMV1-Client-Name", TMV1_CLIENT_NAME),
|
||||
("TMV1-Client-Version", litellm_version),
|
||||
("TMV1-Plugin-Version", PLUGIN_VERSION),
|
||||
("prefer", "redact-pii,return=representation"),
|
||||
*_optional_header("TMV1-Request-Type", request_type),
|
||||
)
|
||||
return MappingProxyType(dict(headers))
|
||||
|
||||
async def _scan_payload(
|
||||
self,
|
||||
payload: Mapping[str, object],
|
||||
request_type: str | None = None,
|
||||
) -> TrendAIScanResult:
|
||||
try:
|
||||
response: Final = await self.async_handler.post(
|
||||
self.api_url,
|
||||
json=dict(payload),
|
||||
headers=dict(self._build_request_headers(request_type)),
|
||||
timeout=self.timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
parsed: Final = TrendAIResponse.model_validate_json(response.content)
|
||||
except httpx.HTTPStatusError as error:
|
||||
return TrendAIProviderFailure(
|
||||
reason="Trend AI Guard returned an HTTP error",
|
||||
status_code=error.response.status_code,
|
||||
)
|
||||
except (httpx.RequestError, Timeout, asyncio.TimeoutError) as error:
|
||||
return TrendAIProviderFailure(reason=f"Trend AI Guard request failed: {type(error).__name__}")
|
||||
except ValidationError:
|
||||
return TrendAIProviderFailure(reason="Trend AI Guard returned an invalid response")
|
||||
|
||||
normalized_action: Final = parsed.action.strip().lower()
|
||||
if normalized_action == "block":
|
||||
reason: Final = ", ".join(parsed.reasons) or parsed.reason or "Content policy violation"
|
||||
return TrendAIBlock(reason=reason)
|
||||
if normalized_action != "allow":
|
||||
return TrendAIProviderFailure(reason=f"Trend AI Guard returned an unsupported action: {parsed.action!r}")
|
||||
|
||||
redacted_content: Final = _extract_redacted_content(request_type, parsed)
|
||||
if parsed.redacted_request is not None and redacted_content is None:
|
||||
return TrendAIProviderFailure(reason="Trend AI Guard returned an invalid redacted payload")
|
||||
|
||||
entity_counts: Final = (
|
||||
tuple(
|
||||
sorted(
|
||||
(
|
||||
rule_id,
|
||||
sum(1 for candidate in parsed.sensitive_information.rules if candidate.id.strip() == rule_id),
|
||||
)
|
||||
for rule_id in frozenset(
|
||||
rule.id.strip() for rule in parsed.sensitive_information.rules if rule.id.strip()
|
||||
)
|
||||
)
|
||||
)
|
||||
if parsed.sensitive_information is not None and redacted_content is not None
|
||||
else ()
|
||||
)
|
||||
return TrendAIAllow(redacted_content=redacted_content, masked_entity_count=entity_counts)
|
||||
|
||||
|
||||
def _chat_completion_payload(content: str, model: str) -> Mapping[str, object]:
|
||||
return MappingProxyType(TrendAIChatCompletionPayload.for_content(content, model).model_dump())
|
||||
|
||||
|
||||
def _optional_header(name: str, value: str | None) -> tuple[tuple[str, str], ...]:
|
||||
return () if value is None else ((name, value),)
|
||||
|
||||
|
||||
def _allow_record(redacted: bool) -> Mapping[str, object]:
|
||||
return MappingProxyType({"action": "allow", "redacted": redacted})
|
||||
|
||||
|
||||
def _is_redacting(result: TrendAIScanResult) -> bool:
|
||||
return isinstance(result, TrendAIAllow) and result.redacted_content is not None
|
||||
|
||||
|
||||
def _window_verdict(scans: Sequence[TrendAIWindowScan]) -> TrendAIScanResult:
|
||||
"""Collapse window scans into one verdict: the first non-allow result, else an allow carrying every entity count."""
|
||||
terminal: Final = next((scan.result for scan in scans if not isinstance(scan.result, TrendAIAllow)), None)
|
||||
if terminal is not None:
|
||||
return terminal
|
||||
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())))
|
||||
|
||||
|
||||
def _merge_entity_counts(counts: Mapping[str, int], additions: Sequence[tuple[str, int]]) -> Mapping[str, int]:
|
||||
"""Per-entity maxima: overlapping scans may report the same entity more than once."""
|
||||
entities: Final = frozenset(counts) | frozenset(entity for entity, _ in additions)
|
||||
return MappingProxyType(
|
||||
{
|
||||
entity: max((counts.get(entity, 0), *(count for candidate, count in additions if candidate == entity)))
|
||||
for entity in entities
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _merge_window_redactions(content: str, scans: Sequence[TrendAIWindowScan]) -> str | None:
|
||||
"""Fold every window's redaction into ``content``; None if any window cannot be merged positionally."""
|
||||
merged = content # rebind-ok: folded across the windows
|
||||
for scan in scans:
|
||||
if not isinstance(scan.result, TrendAIAllow) or scan.result.redacted_content is None:
|
||||
continue
|
||||
merged = apply_window_redaction(merged, scan.window, scan.result.redacted_content)
|
||||
if merged is None:
|
||||
return None
|
||||
return merged
|
||||
|
||||
|
||||
def _build_apply_guardrails_url(api_base: str) -> str:
|
||||
parsed: Final = urlsplit(api_base.strip())
|
||||
normalized_path: Final = parsed.path.rstrip("/")
|
||||
path_with_product: Final = (
|
||||
f"{normalized_path}{_TREND_AI_SECURITY_PATH}"
|
||||
if parsed.hostname is not None
|
||||
and parsed.hostname.lower().endswith(".trendmicro.com")
|
||||
and "/aiSecurity" not in normalized_path
|
||||
else normalized_path
|
||||
)
|
||||
final_path: Final = (
|
||||
path_with_product
|
||||
if path_with_product.endswith(_APPLY_GUARDRAILS_PATH)
|
||||
else f"{path_with_product}{_APPLY_GUARDRAILS_PATH}"
|
||||
)
|
||||
return urlunsplit(SplitResult(parsed.scheme, parsed.netloc, final_path, parsed.query, parsed.fragment))
|
||||
|
||||
|
||||
def _extract_redacted_content(request_type: str | None, response: TrendAIResponse) -> str | None:
|
||||
payload: Final = response.redacted_request
|
||||
if payload is None:
|
||||
return None
|
||||
try:
|
||||
if request_type == OPENAI_CHAT_COMPLETION_RESPONSE_V1:
|
||||
parsed_response: Final = TrendAIRedactedResponse.model_validate(payload)
|
||||
content: Final = " ".join(
|
||||
choice.message.content
|
||||
for choice in parsed_response.choices
|
||||
if choice.message is not None and choice.message.content
|
||||
)
|
||||
return content or None
|
||||
return TrendAIRedactedPrompt.model_validate(payload).prompt or None
|
||||
except ValidationError:
|
||||
return None
|
||||
|
|
@ -0,0 +1,707 @@
|
|||
import json
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from copy import deepcopy
|
||||
from functools import reduce
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.llms.bedrock.passthrough.guardrail_translation.handler import BedrockPassthroughGuardrailHandler
|
||||
from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler
|
||||
from litellm.proxy.guardrails.guardrail_hooks.trendai import TrendAIGuardrail, initialize_guardrail
|
||||
from litellm.proxy.guardrails.guardrail_hooks.trendai._models import (
|
||||
TrendAIAllow,
|
||||
TrendAIBlock,
|
||||
TrendAIProviderFailure,
|
||||
TrendAIScanResult,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_hooks.trendai._text import utf8_windows
|
||||
from litellm.proxy.guardrails.guardrail_hooks.trendai.trendai import (
|
||||
OPENAI_CHAT_COMPLETION_RESPONSE_V1,
|
||||
PLUGIN_VERSION,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler
|
||||
from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
Choices,
|
||||
GenericGuardrailAPIInputs,
|
||||
Message,
|
||||
ModelResponse,
|
||||
)
|
||||
|
||||
|
||||
def _guardrail(
|
||||
*,
|
||||
app_name: str | None = None,
|
||||
fallback_on_error: Literal["block", "allow"] = "block",
|
||||
timeout: float = 5.0,
|
||||
stream_overlap_size: int = 256,
|
||||
response_content_chunk_size_bytes: int = 49_500,
|
||||
async_handler: httpx.AsyncClient | None = None,
|
||||
api_base: str = "https://guard.example.com/v3.0/aiSecurity",
|
||||
event_hook: GuardrailEventHooks = GuardrailEventHooks.pre_call,
|
||||
) -> TrendAIGuardrail:
|
||||
return TrendAIGuardrail(
|
||||
api_key="test-key",
|
||||
api_base=api_base,
|
||||
app_name=app_name,
|
||||
fallback_on_error=fallback_on_error,
|
||||
timeout=timeout,
|
||||
stream_overlap_size=stream_overlap_size,
|
||||
response_content_chunk_size_bytes=response_content_chunk_size_bytes,
|
||||
async_handler=async_handler,
|
||||
guardrail_name="trendai",
|
||||
event_hook=event_hook,
|
||||
)
|
||||
|
||||
|
||||
async def _scan(
|
||||
guardrail: TrendAIGuardrail,
|
||||
payload: Mapping[str, object],
|
||||
request_type: str | None = None,
|
||||
) -> TrendAIScanResult:
|
||||
return await guardrail._scan_payload( # pyright: ignore[reportPrivateUsage] # verifies the section 2 transport seam
|
||||
payload,
|
||||
request_type=request_type,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("api_base", "expected"),
|
||||
[
|
||||
(
|
||||
"https://api.xdr.trendmicro.com",
|
||||
"https://api.xdr.trendmicro.com/v3.0/aiSecurity/applyGuardrails",
|
||||
),
|
||||
(
|
||||
"https://api.xdr.trendmicro.com/v3.0/aiSecurity/",
|
||||
"https://api.xdr.trendmicro.com/v3.0/aiSecurity/applyGuardrails",
|
||||
),
|
||||
(
|
||||
"https://self-hosted.example.com/guard/",
|
||||
"https://self-hosted.example.com/guard/applyGuardrails",
|
||||
),
|
||||
(
|
||||
"https://self-hosted.example.com/guard/applyGuardrails",
|
||||
"https://self-hosted.example.com/guard/applyGuardrails",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_build_apply_guardrails_url(api_base: str, expected: str) -> None:
|
||||
assert _guardrail(api_base=api_base).api_url == expected
|
||||
|
||||
|
||||
def test_environment_fallbacks(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("TMV1_API_KEY", "env-key")
|
||||
monkeypatch.setenv("TRENDAI_AI_GUARD_BASE_URL", "https://guard.example.com")
|
||||
monkeypatch.setenv("TMV1_APPLICATION_NAME", "env-app")
|
||||
|
||||
guardrail = TrendAIGuardrail(
|
||||
guardrail_name="trendai",
|
||||
event_hook=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
|
||||
assert guardrail.api_key == "env-key"
|
||||
assert guardrail.api_url == "https://guard.example.com/applyGuardrails"
|
||||
assert guardrail.app_name == "env-app"
|
||||
|
||||
|
||||
def test_explicit_configuration_takes_precedence(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("TMV1_API_KEY", "env-key")
|
||||
monkeypatch.setenv("TRENDAI_AI_GUARD_BASE_URL", "https://env.example.com")
|
||||
monkeypatch.setenv("TMV1_APPLICATION_NAME", "env-app")
|
||||
|
||||
guardrail = _guardrail(app_name="configured-app")
|
||||
|
||||
assert guardrail.api_key == "test-key"
|
||||
assert guardrail.api_url == "https://guard.example.com/v3.0/aiSecurity/applyGuardrails"
|
||||
assert guardrail.app_name == "configured-app"
|
||||
|
||||
|
||||
def test_negative_stream_overlap_is_rejected() -> None:
|
||||
with pytest.raises(ValueError, match="stream_overlap_size"):
|
||||
_guardrail(stream_overlap_size=-1)
|
||||
|
||||
|
||||
def test_invalid_timeout_is_rejected() -> None:
|
||||
with pytest.raises(ValueError, match="timeout"):
|
||||
_guardrail(timeout=0)
|
||||
|
||||
|
||||
def test_pii_masking_is_not_configurable_on_the_guardrail() -> None:
|
||||
with pytest.raises(TypeError, match="mask_pii"):
|
||||
TrendAIGuardrail(
|
||||
api_key="test-key",
|
||||
api_base="https://guard.example.com/v3.0/aiSecurity",
|
||||
mask_pii=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scan_uses_injected_client_and_required_headers() -> None:
|
||||
captured_request: httpx.Request | None = None
|
||||
|
||||
async def respond(request: httpx.Request) -> httpx.Response:
|
||||
nonlocal captured_request
|
||||
captured_request = request
|
||||
return httpx.Response(200, json={"action": "allow"})
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
result = await _scan(
|
||||
_guardrail(async_handler=client, app_name="test-app"),
|
||||
{"prompt": "hello"},
|
||||
request_type=OPENAI_CHAT_COMPLETION_RESPONSE_V1,
|
||||
)
|
||||
|
||||
assert isinstance(result, TrendAIAllow)
|
||||
assert captured_request is not None
|
||||
assert str(captured_request.url) == "https://guard.example.com/v3.0/aiSecurity/applyGuardrails"
|
||||
assert captured_request.headers["TMV1-Application-Name"] == "test-app"
|
||||
assert captured_request.headers["Authorization"] == "Bearer test-key"
|
||||
assert captured_request.headers["Content-Type"] == "application/json"
|
||||
assert captured_request.headers["TMV1-Client-Name"] == "litellm"
|
||||
assert captured_request.headers["TMV1-Plugin-Version"] == PLUGIN_VERSION
|
||||
assert captured_request.headers["TMV1-Request-Type"] == OPENAI_CHAT_COMPLETION_RESPONSE_V1
|
||||
assert captured_request.headers["prefer"] == "redact-pii,return=representation"
|
||||
assert json.loads(captured_request.content) == {"prompt": "hello"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scan_parses_response_redaction_and_entities() -> None:
|
||||
async def respond(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"action": "allow",
|
||||
"redactedRequest": {"choices": [{"message": {"content": "email: *****"}}]},
|
||||
"sensitiveInformation": {"rules": [{"id": "EMAIL"}, {"id": "EMAIL"}]},
|
||||
},
|
||||
)
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
result = await _scan(
|
||||
_guardrail(async_handler=client),
|
||||
{"choices": []},
|
||||
request_type=OPENAI_CHAT_COMPLETION_RESPONSE_V1,
|
||||
)
|
||||
|
||||
assert result == TrendAIAllow(redacted_content="email: *****", masked_entity_count=(("EMAIL", 2),))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scan_models_block_and_provider_failure_separately() -> None:
|
||||
responses = iter(
|
||||
(
|
||||
httpx.Response(200, json={"action": "block", "reasons": ["malware", "credential theft"]}),
|
||||
httpx.Response(200, json={"unexpected": True}),
|
||||
httpx.Response(503, text="unavailable"),
|
||||
)
|
||||
)
|
||||
|
||||
async def respond(request: httpx.Request) -> httpx.Response:
|
||||
return next(responses)
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
guardrail = _guardrail(async_handler=client)
|
||||
blocked = await _scan(guardrail, {"prompt": "blocked"})
|
||||
malformed = await _scan(guardrail, {"prompt": "malformed"})
|
||||
unavailable = await _scan(guardrail, {"prompt": "unavailable"})
|
||||
|
||||
assert blocked == TrendAIBlock(reason="malware, credential theft")
|
||||
assert isinstance(malformed, TrendAIProviderFailure)
|
||||
assert isinstance(unavailable, TrendAIProviderFailure)
|
||||
assert unavailable.status_code == 503
|
||||
|
||||
|
||||
def test_guardrail_is_discovered_by_global_registries() -> None:
|
||||
from litellm.proxy.guardrails.guardrail_registry import (
|
||||
guardrail_class_registry,
|
||||
guardrail_initializer_registry,
|
||||
)
|
||||
|
||||
assert guardrail_class_registry["trendai"] is TrendAIGuardrail
|
||||
assert "trendai" in guardrail_initializer_registry
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("timeout_config", "expected_timeout"),
|
||||
[({}, 5.0), ({"timeout": None}, 5.0), ({"timeout": 2.5}, 2.5)],
|
||||
ids=["omitted", "null", "configured"],
|
||||
)
|
||||
def test_initializer_defaults_null_timeout_and_preserves_configured_timeout(
|
||||
timeout_config: Mapping[str, float | None], expected_timeout: float
|
||||
) -> None:
|
||||
params: Final = LitellmParams.model_validate(
|
||||
{
|
||||
"guardrail": "trendai",
|
||||
"mode": "pre_call",
|
||||
"api_key": "test-key",
|
||||
"api_base": "https://guard.example.com",
|
||||
**timeout_config,
|
||||
}
|
||||
)
|
||||
callback: Final = initialize_guardrail(params, {"guardrail_name": "trendai-timeout"})
|
||||
try:
|
||||
assert callback.timeout == expected_timeout
|
||||
finally:
|
||||
litellm.logging_callback_manager.remove_callback_from_all_lists(callback)
|
||||
|
||||
|
||||
def test_registry_rejects_tool_result_only_scanning() -> None:
|
||||
with pytest.raises(ValueError, match="never scans tool results"):
|
||||
InMemoryGuardrailHandler().initialize_guardrail(
|
||||
{
|
||||
"guardrail_name": "trendai-tool-results",
|
||||
"litellm_params": {
|
||||
"guardrail": "trendai",
|
||||
"mode": "pre_call",
|
||||
"api_key": "test-key",
|
||||
"api_base": "https://guard.example.com",
|
||||
"timeout": 5.0,
|
||||
"scan_only_tool_results": True,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
Responder = Callable[[httpx.Request], httpx.Response]
|
||||
|
||||
|
||||
def _engine(
|
||||
*,
|
||||
block_on: str | None = None,
|
||||
redact: Mapping[str, str] | None = None,
|
||||
entity: str = "PII",
|
||||
fail_on: str | None = None,
|
||||
) -> tuple[Responder, list[str]]:
|
||||
scanned: list[str] = []
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
body = json.loads(request.content)
|
||||
text = body["prompt"] if "prompt" in body else body["choices"][0]["message"]["content"]
|
||||
scanned.append(text)
|
||||
if fail_on is not None and fail_on in text:
|
||||
return httpx.Response(503, text="down")
|
||||
if block_on is not None and block_on in text:
|
||||
return httpx.Response(200, json={"action": "block", "reasons": ["policy"]})
|
||||
hits = {needle: mask for needle, mask in (redact or {}).items() if needle in text}
|
||||
if not hits:
|
||||
return httpx.Response(200, json={"action": "allow"})
|
||||
redacted = reduce(lambda value, pair: value.replace(*pair), hits.items(), text)
|
||||
payload = {"prompt": redacted} if "prompt" in body else {"choices": [{"message": {"content": redacted}}]}
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"action": "allow",
|
||||
"redactedRequest": payload,
|
||||
"sensitiveInformation": {"rules": [{"id": entity}]},
|
||||
},
|
||||
)
|
||||
|
||||
return respond, scanned
|
||||
|
||||
|
||||
def _guardrail_records(request_data: Mapping[str, object]) -> list[dict[str, object]]:
|
||||
metadata = request_data["metadata"]
|
||||
assert isinstance(metadata, dict)
|
||||
return list(metadata.get("standard_logging_guardrail_information") or [])
|
||||
|
||||
|
||||
async def _apply(
|
||||
guardrail: TrendAIGuardrail,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
input_type: Literal["request", "response"],
|
||||
) -> tuple[GenericGuardrailAPIInputs, dict[str, object]]:
|
||||
request_data: dict[str, object] = {"metadata": {}}
|
||||
result = await guardrail.apply_guardrail(inputs=inputs, request_data=request_data, input_type=input_type)
|
||||
return result, request_data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_scans_only_the_last_user_turn_and_writes_redaction_back() -> None:
|
||||
respond, scanned = _engine(redact={"a@b.com": "[EMAIL]"}, entity="EMAIL")
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
result, request_data = await _apply(
|
||||
_guardrail(async_handler=client),
|
||||
{
|
||||
"texts": ["system prompt", "first turn a@b.com", "reply", "mail a@b.com now"],
|
||||
"structured_messages": [
|
||||
{"role": "system", "content": "system prompt"},
|
||||
{"role": "user", "content": "first turn a@b.com"},
|
||||
{"role": "assistant", "content": "reply"},
|
||||
{"role": "user", "content": "mail a@b.com now"},
|
||||
],
|
||||
},
|
||||
"request",
|
||||
)
|
||||
|
||||
assert scanned == ["mail a@b.com now"]
|
||||
assert result["texts"] == ["system prompt", "first turn a@b.com", "reply", "mail [EMAIL] now"]
|
||||
(record,) = _guardrail_records(request_data)
|
||||
assert record["guardrail_status"] == "success"
|
||||
assert record["masked_entity_count"] == {"EMAIL": 1}
|
||||
assert record["guardrail_response"] == {"action": "allow", "redacted": True}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("repeated_role", ["assistant", "tool"])
|
||||
@pytest.mark.parametrize("multipart", [False, True], ids=["string", "multipart"])
|
||||
async def test_request_redaction_targets_the_last_user_even_when_later_messages_repeat_it(
|
||||
repeated_role: Literal["assistant", "tool"], multipart: bool
|
||||
) -> None:
|
||||
respond, scanned = _engine(redact={"a@b.com": "[EMAIL]"})
|
||||
user_content: Final = (
|
||||
[
|
||||
{"type": "text", "text": "mail "},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}},
|
||||
{"type": "text", "text": "a@b.com"},
|
||||
]
|
||||
if multipart
|
||||
else "mail a@b.com"
|
||||
)
|
||||
repeated_message: Final = {
|
||||
"role": repeated_role,
|
||||
"content": (
|
||||
[{"type": "text", "text": "mail "}, {"type": "text", "text": "a@b.com"}] if multipart else "mail a@b.com"
|
||||
),
|
||||
**({"tool_call_id": "call-1"} if repeated_role == "tool" else {}),
|
||||
}
|
||||
expected_repeated_message: Final = deepcopy(repeated_message)
|
||||
expected_content: Final = (
|
||||
[
|
||||
{"type": "text", "text": "mail "},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}},
|
||||
{"type": "text", "text": "[EMAIL]"},
|
||||
]
|
||||
if multipart
|
||||
else "mail [EMAIL]"
|
||||
)
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
out: Final = await OpenAIChatCompletionsHandler().process_input_messages(
|
||||
data={
|
||||
"messages": [
|
||||
{"role": "user", "content": "mail a@b.com"},
|
||||
{"role": "assistant", "content": "previous reply"},
|
||||
{"role": "user", "content": user_content},
|
||||
repeated_message,
|
||||
],
|
||||
"metadata": {},
|
||||
},
|
||||
guardrail_to_apply=_guardrail(async_handler=client),
|
||||
)
|
||||
|
||||
assert scanned == ["mail a@b.com"]
|
||||
assert out["messages"] == [
|
||||
{"role": "user", "content": "mail a@b.com"},
|
||||
{"role": "assistant", "content": "previous reply"},
|
||||
{"role": "user", "content": expected_content},
|
||||
expected_repeated_message,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_converse_redaction_with_unaligned_text_slots_fails_closed() -> None:
|
||||
respond, scanned = _engine(redact={"a@b.com": "[EMAIL]"})
|
||||
body: Final = {
|
||||
"messages": [
|
||||
{"role": "user", "content": [{"text": "lookup"}]},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"toolUse": {"toolUseId": "call-1", "name": "lookup", "input": {}}}],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"text": "mail a@b.com"},
|
||||
{"toolResult": {"toolUseId": "call-1", "content": [{"text": "found"}]}},
|
||||
],
|
||||
},
|
||||
],
|
||||
"additionalModelRequestFields": {"system": "extra"},
|
||||
}
|
||||
original_body: Final = deepcopy(body)
|
||||
request_data: Final[dict[str, object]] = {
|
||||
"endpoint": "/model/test-model/converse",
|
||||
"data": body,
|
||||
"metadata": {},
|
||||
}
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
with pytest.raises(GuardrailRaisedException, match="could not apply the requested redaction"):
|
||||
await BedrockPassthroughGuardrailHandler().process_input_messages(
|
||||
request_data, _guardrail(async_handler=client, fallback_on_error="allow")
|
||||
)
|
||||
|
||||
assert scanned == ["mail a@b.comextra"]
|
||||
assert request_data["data"] == original_body
|
||||
assert _guardrail_records(request_data)[-1]["guardrail_status"] == "guardrail_failed_to_respond"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_redaction_is_split_back_across_multipart_user_content() -> None:
|
||||
respond, scanned = _engine(redact={"4111": "####"})
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
result, _ = await _apply(
|
||||
_guardrail(async_handler=client),
|
||||
{
|
||||
"texts": ["card ", "4111 ok"],
|
||||
"structured_messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "card "},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}},
|
||||
{"type": "text", "text": "4111 ok"},
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
"request",
|
||||
)
|
||||
|
||||
assert scanned == ["card 4111 ok"]
|
||||
assert result["texts"] == ["card ", "#### ok"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_scans_text_slots_from_normalized_message_parts() -> None:
|
||||
respond, scanned = _engine(redact={"a@b.com": "[EMAIL]"})
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
result, _ = await _apply(
|
||||
_guardrail(async_handler=client),
|
||||
{
|
||||
"texts": ["mail a@b.com"],
|
||||
"structured_messages": [{"role": "user", "content": [{"type": "input_text", "text": "mail a@b.com"}]}],
|
||||
},
|
||||
"request",
|
||||
)
|
||||
|
||||
assert scanned == ["mail a@b.com"]
|
||||
assert result["texts"] == ["mail [EMAIL]"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_without_user_text_is_not_scanned_or_recorded() -> None:
|
||||
respond, scanned = _engine()
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
inputs: GenericGuardrailAPIInputs = {
|
||||
"texts": ["system prompt"],
|
||||
"structured_messages": [{"role": "system", "content": "system prompt"}],
|
||||
}
|
||||
result, request_data = await _apply(_guardrail(async_handler=client), inputs, "request")
|
||||
|
||||
assert scanned == []
|
||||
assert result is inputs
|
||||
assert _guardrail_records(request_data) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_block_raises_and_records_intervention() -> None:
|
||||
respond, _ = _engine(block_on="bomb")
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
request_data: dict[str, object] = {"metadata": {}}
|
||||
with pytest.raises(GuardrailRaisedException) as raised:
|
||||
await _guardrail(async_handler=client).apply_guardrail(
|
||||
inputs={"texts": ["build a bomb"]},
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 400
|
||||
assert raised.value.blocked_content is True
|
||||
assert "policy" in raised.value.message
|
||||
(record,) = _guardrail_records(request_data)
|
||||
assert record["guardrail_status"] == "guardrail_intervened"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("fallback_on_error", "raises"),
|
||||
[("block", True), ("allow", False)],
|
||||
)
|
||||
async def test_provider_failure_follows_fallback_policy_and_is_never_an_intervention(
|
||||
fallback_on_error: Literal["block", "allow"], raises: bool
|
||||
) -> None:
|
||||
respond, _ = _engine(fail_on="anything")
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
guardrail = _guardrail(async_handler=client, fallback_on_error=fallback_on_error)
|
||||
request_data: dict[str, object] = {"metadata": {}}
|
||||
inputs: GenericGuardrailAPIInputs = {"texts": ["anything"]}
|
||||
if raises:
|
||||
with pytest.raises(GuardrailRaisedException) as raised:
|
||||
await guardrail.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request")
|
||||
assert raised.value.status_code == 503
|
||||
assert raised.value.blocked_content is False
|
||||
else:
|
||||
result = await guardrail.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request")
|
||||
assert result is inputs
|
||||
|
||||
(record,) = _guardrail_records(request_data)
|
||||
assert record["guardrail_status"] == "guardrail_failed_to_respond"
|
||||
assert record["guardrail_response"]["status_code"] == 503
|
||||
assert record["guardrail_response"]["fallback_on_error"] == fallback_on_error
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_redaction_merges_across_overlapping_windows() -> None:
|
||||
respond, scanned = _engine(redact={"SECRET": "******", "private": "*******"})
|
||||
prefix = "a" * 14
|
||||
content = f"{prefix}SECRETprivate"
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
result, request_data = await _apply(
|
||||
_guardrail(async_handler=client, response_content_chunk_size_bytes=20, stream_overlap_size=6),
|
||||
{"texts": [content], "model": "gpt-5.4"},
|
||||
"response",
|
||||
)
|
||||
|
||||
assert scanned == [f"{prefix}SECRET", "SECRETprivate"]
|
||||
assert result["texts"] == [f"{prefix}*************"]
|
||||
(record,) = _guardrail_records(request_data)
|
||||
assert record["guardrail_status"] == "success"
|
||||
assert record["masked_entity_count"] == {"PII": 1}
|
||||
|
||||
|
||||
def _redacted_choice(content: str) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200, json={"action": "allow", "redactedRequest": {"choices": [{"message": {"content": content}}]}}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_later_overlap_scan_cannot_undo_an_earlier_redaction() -> None:
|
||||
"""Window two sees ``SECRET`` in its overlap unmasked, flags only ``tail``, and must not restore ``SECRET``."""
|
||||
responses = iter((_redacted_choice("aaaaa******"), _redacted_choice("SECRET####")))
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: next(responses))) as client:
|
||||
result, _ = await _apply(
|
||||
_guardrail(async_handler=client, response_content_chunk_size_bytes=11, stream_overlap_size=6),
|
||||
{"texts": ["aaaaaSECRETtail"]},
|
||||
"response",
|
||||
)
|
||||
|
||||
assert result["texts"] == ["aaaaa******####"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_block_in_a_later_window_stops_scanning() -> None:
|
||||
respond, scanned = _engine(block_on="zzz")
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
request_data: dict[str, object] = {"metadata": {}}
|
||||
with pytest.raises(GuardrailRaisedException, match="policy"):
|
||||
await _guardrail(
|
||||
async_handler=client, response_content_chunk_size_bytes=5, stream_overlap_size=0
|
||||
).apply_guardrail(
|
||||
inputs={"texts": ["aaaaabbbbbzzzzzccccc"]},
|
||||
request_data=request_data,
|
||||
input_type="response",
|
||||
)
|
||||
|
||||
assert scanned == ["aaaaa", "bbbbb", "zzzzz"]
|
||||
assert _guardrail_records(request_data)[0]["guardrail_status"] == "guardrail_intervened"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_scans_every_choice() -> None:
|
||||
respond, scanned = _engine(redact={"SECRET": "******"})
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
result, _ = await _apply(
|
||||
_guardrail(async_handler=client),
|
||||
{"texts": ["first SECRET", "second SECRET"]},
|
||||
"response",
|
||||
)
|
||||
|
||||
assert scanned == ["first SECRET", "second SECRET"]
|
||||
assert result["texts"] == ["first ******", "second ******"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unmergeable_response_redaction_blocks_instead_of_leaking() -> None:
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={"action": "allow", "redactedRequest": {"choices": [{"message": {"content": "short"}}]}},
|
||||
)
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
request_data: dict[str, object] = {"metadata": {}}
|
||||
with pytest.raises(GuardrailRaisedException, match="redaction"):
|
||||
await _guardrail(async_handler=client, fallback_on_error="allow").apply_guardrail(
|
||||
inputs={"texts": ["a much longer sensitive response"]},
|
||||
request_data=request_data,
|
||||
input_type="response",
|
||||
)
|
||||
|
||||
assert _guardrail_records(request_data)[-1]["guardrail_status"] == "guardrail_failed_to_respond"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_redaction_flows_through_the_chat_completions_handler() -> None:
|
||||
respond, _ = _engine(redact={"a@b.com": "[EMAIL]"})
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
data: dict[str, object] = {
|
||||
"model": "gpt-5.4",
|
||||
"messages": [
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": [{"type": "text", "text": "mail a@b.com"}]},
|
||||
],
|
||||
"metadata": {},
|
||||
}
|
||||
out = await OpenAIChatCompletionsHandler().process_input_messages(
|
||||
data=data, guardrail_to_apply=_guardrail(async_handler=client)
|
||||
)
|
||||
|
||||
assert out["messages"][0]["content"] == "be terse"
|
||||
assert out["messages"][1]["content"][0]["text"] == "mail [EMAIL]"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("content", "chunk_size_bytes", "overlap", "expected"),
|
||||
[
|
||||
("abcdéX", 5, 0, ["abcd", "éX"]),
|
||||
("abcéX", 5, 0, ["abcé", "X"]),
|
||||
("abcdéXYZ", 5, 2, ["abcd", "cdéX", "éXYZ"]),
|
||||
("abcdef", 5, 100, ["abcde", "bcdef"]),
|
||||
("", 5, 2, []),
|
||||
("é", 1, 0, ["é"]),
|
||||
],
|
||||
)
|
||||
def test_utf8_windows_respect_byte_limits_and_character_overlap(
|
||||
content: str, chunk_size_bytes: int, overlap: int, expected: Sequence[str]
|
||||
) -> None:
|
||||
windows = utf8_windows(content, chunk_size_bytes=chunk_size_bytes, overlap_chars=overlap)
|
||||
|
||||
assert [window.text for window in windows] == list(expected)
|
||||
assert all(content[window.start : window.start + len(window.text)] == window.text for window in windows)
|
||||
|
||||
|
||||
def _logged_call(user_text: str, assistant_text: str) -> tuple[dict[str, object], ModelResponse]:
|
||||
response = ModelResponse(choices=[Choices(message=Message(role="assistant", content=assistant_text))])
|
||||
kwargs: dict[str, object] = {
|
||||
"model": "gpt-5.4",
|
||||
"messages": [{"role": "user", "content": user_text}],
|
||||
"litellm_call_id": "call-1",
|
||||
"litellm_params": {"metadata": {}},
|
||||
"optional_params": {},
|
||||
"standard_logging_object": {"guardrail_information": None},
|
||||
}
|
||||
return kwargs, response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logging_only_scans_both_sides_without_modifying_them() -> None:
|
||||
respond, scanned = _engine(redact={"a@b.com": "[EMAIL]", "SECRET": "******"})
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
|
||||
guardrail = _guardrail(async_handler=client, event_hook=GuardrailEventHooks.logging_only)
|
||||
kwargs, response = _logged_call("user a@b.com", "assistant SECRET")
|
||||
out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
|
||||
|
||||
assert scanned == ["user a@b.com", "assistant SECRET"]
|
||||
assert out_kwargs["messages"] == [{"role": "user", "content": "user a@b.com"}]
|
||||
assert out_response is response
|
||||
assert response.choices[0].message.content == "assistant SECRET"
|
||||
entries = out_kwargs["standard_logging_object"]["guardrail_information"]
|
||||
assert [entry["guardrail_status"] for entry in entries] == ["success", "success"]
|
||||
assert all(entry["guardrail_mode"] == "logging_only" for entry in entries)
|
||||
Loading…
Add table
Reference in a new issue