mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
Merge 038a1bda1e into a392a39627
This commit is contained in:
Vivek Kalyanarangan
GitHub
commit
5170c3ae3f
6 changed files with 1044 additions and 54 deletions
|
|
@ -7,12 +7,14 @@ to detect and block/mask sensitive content.
|
|||
|
||||
import asyncio
|
||||
import contextvars
|
||||
import functools
|
||||
import itertools
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, AsyncIterable, Callable, Coroutine, Mapping, Sequence
|
||||
from bisect import bisect_left, bisect_right
|
||||
from collections.abc import AsyncGenerator, AsyncIterable, Callable, Coroutine, Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass, replace
|
||||
from datetime import datetime
|
||||
from re import Pattern
|
||||
|
|
@ -21,6 +23,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, Type
|
|||
import yaml
|
||||
from fastapi import HTTPException
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
from litellm import Router
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -79,6 +82,7 @@ from .competitor_intent import (
|
|||
BaseCompetitorIntentChecker,
|
||||
)
|
||||
from .patterns import PATTERN_EXTRA_CONFIG, get_compiled_pattern
|
||||
from .validators import PATTERN_VALIDATORS
|
||||
|
||||
MAX_KEYWORD_VALUE_GAP_WORDS: Final = 1
|
||||
_DECODED_ARGUMENTS: Final = TypeAdapter(object)
|
||||
|
|
@ -109,6 +113,12 @@ WORD_NUMBER_SEQUENCE_PATTERN: Final = re.compile(
|
|||
)
|
||||
WORD_NUMBER_TOKEN_FINDER: Final = re.compile(rf"(?:{WORD_NUMBER_TOKEN_REGEX})", re.IGNORECASE)
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _checksum_window_regex(regex: Pattern[str]) -> Pattern[str]:
|
||||
return re.compile("(?=(" + regex.pattern.replace(r"\b", "") + "))", regex.flags)
|
||||
|
||||
|
||||
_logging_only_scan_guardrail: Final[contextvars.ContextVar["ContentFilterGuardrail | None"]] = contextvars.ContextVar(
|
||||
"litellm_content_filter_logging_only_scan", default=None
|
||||
)
|
||||
|
|
@ -130,11 +140,13 @@ class CompiledPatternEntry(TypedDict):
|
|||
action: ContentFilterAction
|
||||
keyword_regex: Pattern[str] | None
|
||||
allow_word_numbers: bool
|
||||
validator: ReadOnly[Callable[[str], bool] | None]
|
||||
|
||||
|
||||
class _PatternExtraLookup(TypedDict):
|
||||
keyword_pattern: str | None
|
||||
allow_word_numbers: bool
|
||||
validator: ReadOnly[Callable[[str], bool] | None]
|
||||
|
||||
|
||||
class _CategoryConfigView(TypedDict):
|
||||
|
|
@ -219,6 +231,114 @@ async def _stream_then_end_of_stream_chunk(
|
|||
yield held_item, False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _SortedSpans:
|
||||
starts: tuple[int, ...]
|
||||
furthest_ends: tuple[int, ...]
|
||||
|
||||
@classmethod
|
||||
def of(cls, spans: Iterable[tuple[int, int]]) -> "_SortedSpans":
|
||||
ordered: Final = sorted(spans)
|
||||
return cls(
|
||||
starts=tuple(start for start, _ in ordered),
|
||||
furthest_ends=tuple(itertools.accumulate((end for _, end in ordered), max)),
|
||||
)
|
||||
|
||||
def overlap(self, span: tuple[int, int]) -> bool:
|
||||
index: Final = bisect_left(self.starts, span[1])
|
||||
return index > 0 and self.furthest_ends[index - 1] > span[0]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Redaction:
|
||||
start: int
|
||||
end: int
|
||||
original_start: int
|
||||
original_end: int
|
||||
failed_checksum: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _MaskedText:
|
||||
text: str
|
||||
redactions: tuple[_Redaction, ...] = ()
|
||||
tracks_redactions: bool = True
|
||||
|
||||
def original_span(self, span: tuple[int, int]) -> tuple[int, int]:
|
||||
return self._original_position(span[0], is_end=False), self._original_position(span[1], is_end=True)
|
||||
|
||||
def _original_position(self, position: int, *, is_end: bool) -> int:
|
||||
find: Final = bisect_left if is_end else bisect_right
|
||||
index: Final = find(self.redactions, position, key=lambda redaction: redaction.start) - 1
|
||||
if index < 0:
|
||||
return position
|
||||
redaction: Final = self.redactions[index]
|
||||
if position >= redaction.end:
|
||||
return redaction.original_end + position - redaction.end
|
||||
return redaction.original_end if is_end else redaction.original_start
|
||||
|
||||
def _redactions_overlapping(self, span: tuple[int, int]) -> Sequence[_Redaction]:
|
||||
first: Final = bisect_right(self.redactions, span[0], key=lambda redaction: redaction.end)
|
||||
return self.redactions[
|
||||
first : bisect_left(self.redactions, span[1], lo=first, key=lambda redaction: redaction.start)
|
||||
]
|
||||
|
||||
def masked(self, spans: Sequence[tuple[int, int]], tag: str, failed_checksum: Sequence[bool]) -> "_MaskedText":
|
||||
span_ends: Final = tuple(end for _, end in spans)
|
||||
text_pieces: Final = itertools.chain.from_iterable(
|
||||
(self.text[previous_end:start], tag) for previous_end, (start, _) in zip((0, *span_ends), spans)
|
||||
)
|
||||
text: Final = "".join((*text_pieces, self.text[span_ends[-1] :]))
|
||||
if not self.tracks_redactions:
|
||||
return _MaskedText(text=text, tracks_redactions=False)
|
||||
growth: Final = tuple(itertools.accumulate((len(tag) - end + start for start, end in spans), initial=0))
|
||||
absorbed: Final = tuple(self._redactions_overlapping(span) for span in spans)
|
||||
dropped: Final = frozenset(itertools.chain.from_iterable(absorbed))
|
||||
shifted: Final = (
|
||||
replace(redaction, start=redaction.start + shift, end=redaction.end + shift)
|
||||
for redaction, shift in zip(
|
||||
self.redactions, (growth[bisect_right(span_ends, redaction.start)] for redaction in self.redactions)
|
||||
)
|
||||
if redaction not in dropped
|
||||
)
|
||||
added: Final = (
|
||||
_Redaction(
|
||||
start=start + growth[index],
|
||||
end=start + growth[index] + len(tag),
|
||||
original_start=self._original_position(start, is_end=False),
|
||||
original_end=self._original_position(end, is_end=True),
|
||||
failed_checksum=failed_checksum[index]
|
||||
and all(redaction.failed_checksum for redaction in absorbed[index]),
|
||||
)
|
||||
for index, (start, end) in enumerate(spans)
|
||||
)
|
||||
return _MaskedText(
|
||||
text=text,
|
||||
redactions=tuple(sorted((*shifted, *added), key=lambda redaction: redaction.start)),
|
||||
)
|
||||
|
||||
def restored(self, original: str) -> tuple[str, tuple[tuple[int, int], ...]]:
|
||||
failed: Final = tuple(redaction for redaction in self.redactions if redaction.failed_checksum)
|
||||
if not failed:
|
||||
return self.text, ()
|
||||
gaps: Final = tuple(
|
||||
self.text[previous_end : redaction.start]
|
||||
for previous_end, redaction in zip((0, *(redaction.end for redaction in failed)), failed)
|
||||
)
|
||||
numbers: Final = tuple(original[redaction.original_start : redaction.original_end] for redaction in failed)
|
||||
ends: Final = tuple(itertools.accumulate(len(gap) + len(number) for gap, number in zip(gaps, numbers)))
|
||||
return (
|
||||
"".join((*itertools.chain.from_iterable(zip(gaps, numbers)), self.text[failed[-1].end :])),
|
||||
tuple((end - len(number), end) for end, number in zip(ends, numbers)),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _FirstPass:
|
||||
masked: _MaskedText
|
||||
detected_rules: frozenset[int] = frozenset()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _StreamedScanPlan:
|
||||
context_chars: int
|
||||
|
|
@ -820,7 +940,11 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
pattern_config: ContentFilterPattern configuration
|
||||
"""
|
||||
try:
|
||||
extra_config: _PatternExtraLookup = {"keyword_pattern": None, "allow_word_numbers": False}
|
||||
extra_config: _PatternExtraLookup = {
|
||||
"keyword_pattern": None,
|
||||
"allow_word_numbers": False,
|
||||
"validator": None,
|
||||
}
|
||||
if pattern_config.pattern_type == "prebuilt":
|
||||
if not pattern_config.pattern_name:
|
||||
raise ValueError("pattern_name is required for prebuilt patterns")
|
||||
|
|
@ -845,6 +969,7 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
"action": pattern_config.action,
|
||||
"keyword_regex": keyword_regex,
|
||||
"allow_word_numbers": extra_config["allow_word_numbers"],
|
||||
"validator": extra_config["validator"],
|
||||
}
|
||||
)
|
||||
verbose_proxy_logger.debug("Added pattern: %s with action %s", pattern_name, pattern_config.action)
|
||||
|
|
@ -855,9 +980,11 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
@staticmethod
|
||||
def _lookup_pattern_extra(pattern_name: str) -> _PatternExtraLookup:
|
||||
extra: Final = PATTERN_EXTRA_CONFIG.get(pattern_name)
|
||||
validator_name: Final = extra.get("validator") if extra is not None else None
|
||||
return {
|
||||
"keyword_pattern": extra.get("keyword_pattern") if extra is not None else None,
|
||||
"allow_word_numbers": bool(extra.get("allow_word_numbers")) if extra is not None else False,
|
||||
"validator": PATTERN_VALIDATORS[validator_name] if validator_name is not None else None,
|
||||
}
|
||||
|
||||
def _load_blocked_words_file(self, file_path: str) -> None:
|
||||
|
|
@ -934,6 +1061,31 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
|
||||
return self._merge_spans(match_spans)
|
||||
|
||||
def _checksum_windows(self, text: str) -> _SortedSpans:
|
||||
return _SortedSpans.of(
|
||||
itertools.chain.from_iterable(self._valid_checksum_windows(text, entry) for entry in self.compiled_patterns)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _valid_checksum_windows(text: str, pattern_entry: CompiledPatternEntry) -> tuple[tuple[int, int], ...]:
|
||||
validator: Final = pattern_entry["validator"]
|
||||
if validator is None:
|
||||
return ()
|
||||
return tuple(
|
||||
candidate.span(1)
|
||||
for candidate in _checksum_window_regex(pattern_entry["regex"]).finditer(text)
|
||||
if validator(candidate.group(1))
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _fails_checksum(
|
||||
masked: _MaskedText,
|
||||
span: tuple[int, int],
|
||||
pattern_entry: CompiledPatternEntry,
|
||||
checksum_windows: Callable[[], _SortedSpans],
|
||||
) -> bool:
|
||||
return pattern_entry["validator"] is not None and not checksum_windows().overlap(masked.original_span(span))
|
||||
|
||||
def _match_near_keyword(
|
||||
self,
|
||||
value_start: int,
|
||||
|
|
@ -986,21 +1138,6 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
merged.append((start, end))
|
||||
return merged
|
||||
|
||||
def _mask_spans(self, text: str, spans: list[tuple[int, int]], redaction: str) -> str:
|
||||
"""Apply masking for the provided spans using the given redaction tag."""
|
||||
|
||||
if not spans:
|
||||
return text
|
||||
|
||||
result_parts: Final[list[str]] = []
|
||||
previous_end = 0
|
||||
for start, end in spans:
|
||||
result_parts.append(text[previous_end:start])
|
||||
result_parts.append(redaction)
|
||||
previous_end = end
|
||||
result_parts.append(text[previous_end:])
|
||||
return "".join(result_parts)
|
||||
|
||||
def _convert_word_number_sequence(self, sequence: str) -> str | None:
|
||||
"""Convert a spelled-out digit sequence (e.g., 'One-Two') into digits."""
|
||||
|
||||
|
|
@ -1027,8 +1164,14 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
Returns:
|
||||
Tuple of (matched_text, pattern_name, action) if match found, None otherwise
|
||||
"""
|
||||
checksum_windows: Final = functools.cache(functools.partial(self._checksum_windows, text))
|
||||
unmasked: Final = _MaskedText(text=text)
|
||||
for pattern_entry in self.compiled_patterns:
|
||||
spans = self._find_pattern_spans(text, pattern_entry)
|
||||
spans = [
|
||||
span
|
||||
for span in self._find_pattern_spans(text, pattern_entry)
|
||||
if not self._fails_checksum(unmasked, span, pattern_entry, checksum_windows)
|
||||
]
|
||||
if spans:
|
||||
start, end = spans[0]
|
||||
matched_text = text[start:end]
|
||||
|
|
@ -1039,7 +1182,7 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
return None
|
||||
|
||||
def _check_conditional_categories(
|
||||
self, text: str, exceptions: list[str]
|
||||
self, text: str, exceptions: Sequence[str]
|
||||
) -> tuple[str, str, str, ContentFilterAction] | None:
|
||||
"""
|
||||
Check text for conditional category matches (identifier + block word in same sentence).
|
||||
|
|
@ -1127,7 +1270,7 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
return None
|
||||
|
||||
def _check_phrase_patterns(
|
||||
self, text: str, exceptions: list[str]
|
||||
self, text: str, exceptions: Sequence[str]
|
||||
) -> tuple[str, str, str, ContentFilterAction] | None:
|
||||
"""
|
||||
Check text against phrase patterns from loaded categories.
|
||||
|
|
@ -1176,7 +1319,7 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
return None
|
||||
|
||||
def _check_category_keywords(
|
||||
self, text: str, exceptions: list[str]
|
||||
self, text: str, exceptions: Sequence[str]
|
||||
) -> tuple[str, str, str, ContentFilterAction] | None:
|
||||
"""
|
||||
Check text for category keywords.
|
||||
|
|
@ -1391,11 +1534,9 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
self,
|
||||
pattern_name: str,
|
||||
action: ContentFilterAction,
|
||||
text: str,
|
||||
spans: list[tuple[int, int]],
|
||||
detections: list[ContentFilterDetection] | None,
|
||||
) -> str:
|
||||
"""Handle regex pattern match detection and action."""
|
||||
) -> None:
|
||||
"""Record a regex pattern detection and raise when its action blocks."""
|
||||
if detections is not None:
|
||||
pattern_detection: Final[PatternDetection] = {
|
||||
"type": "pattern",
|
||||
|
|
@ -1411,12 +1552,104 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
status_code=400,
|
||||
detail={"error": error_msg, "pattern": pattern_name},
|
||||
)
|
||||
elif action == ContentFilterAction.MASK:
|
||||
redaction_tag: Final = self.pattern_redaction_format.format(pattern_name=pattern_name.upper())
|
||||
text = self._mask_spans(text, spans, redaction_tag)
|
||||
verbose_proxy_logger.info("Masked all %s matches in content", pattern_name)
|
||||
|
||||
return text
|
||||
def _apply_patterns(
|
||||
self,
|
||||
text: str,
|
||||
detections: list[ContentFilterDetection] | None,
|
||||
validates_checksums: bool,
|
||||
) -> str:
|
||||
checks_checksums: Final = validates_checksums and any(
|
||||
pattern_entry["validator"] is not None for pattern_entry in self.compiled_patterns
|
||||
)
|
||||
checksum_windows: Final = (
|
||||
functools.cache(functools.partial(self._checksum_windows, text)) if checks_checksums else None
|
||||
)
|
||||
first_pass: Final = functools.reduce(
|
||||
lambda so_far, indexed_entry: self._act_on_first_pass_matches(
|
||||
so_far, *indexed_entry, checksum_windows, detections
|
||||
),
|
||||
enumerate(self.compiled_patterns),
|
||||
_FirstPass(masked=_MaskedText(text=text, tracks_redactions=checks_checksums)),
|
||||
)
|
||||
put_back_text, put_back = first_pass.masked.restored(text)
|
||||
if not put_back:
|
||||
return put_back_text
|
||||
put_back_spans: Final = _SortedSpans.of(put_back)
|
||||
return functools.reduce(
|
||||
lambda masked_so_far, indexed_entry: self._act_on_spans(
|
||||
masked_so_far,
|
||||
indexed_entry[1],
|
||||
*self._matches_on_put_back_numbers(masked_so_far, indexed_entry[1], put_back_spans),
|
||||
None if indexed_entry[0] in first_pass.detected_rules else detections,
|
||||
),
|
||||
(
|
||||
(index, pattern_entry)
|
||||
for index, pattern_entry in enumerate(self.compiled_patterns)
|
||||
if pattern_entry["validator"] is None
|
||||
),
|
||||
_MaskedText(text=put_back_text),
|
||||
).text
|
||||
|
||||
def _act_on_first_pass_matches(
|
||||
self,
|
||||
so_far: _FirstPass,
|
||||
index: int,
|
||||
pattern_entry: CompiledPatternEntry,
|
||||
checksum_windows: Callable[[], _SortedSpans] | None,
|
||||
detections: list[ContentFilterDetection] | None,
|
||||
) -> _FirstPass:
|
||||
spans, failed_checksum = self._checked_matches(so_far.masked, pattern_entry, checksum_windows)
|
||||
return _FirstPass(
|
||||
masked=self._act_on_spans(so_far.masked, pattern_entry, spans, failed_checksum, detections),
|
||||
detected_rules=(
|
||||
so_far.detected_rules if all(failed_checksum) else so_far.detected_rules | frozenset((index,))
|
||||
),
|
||||
)
|
||||
|
||||
def _checked_matches(
|
||||
self,
|
||||
masked: _MaskedText,
|
||||
pattern_entry: CompiledPatternEntry,
|
||||
checksum_windows: Callable[[], _SortedSpans] | None,
|
||||
) -> tuple[Sequence[tuple[int, int]], tuple[bool, ...]]:
|
||||
spans: Final = self._find_pattern_spans(masked.text, pattern_entry)
|
||||
if checksum_windows is None:
|
||||
return spans, (False,) * len(spans)
|
||||
return spans, tuple(self._fails_checksum(masked, span, pattern_entry, checksum_windows) for span in spans)
|
||||
|
||||
def _matches_on_put_back_numbers(
|
||||
self,
|
||||
masked: _MaskedText,
|
||||
pattern_entry: CompiledPatternEntry,
|
||||
put_back: _SortedSpans,
|
||||
) -> tuple[Sequence[tuple[int, int]], tuple[bool, ...]]:
|
||||
spans: Final = tuple(
|
||||
span
|
||||
for span in self._find_pattern_spans(masked.text, pattern_entry)
|
||||
if put_back.overlap(masked.original_span(span))
|
||||
)
|
||||
return spans, (False,) * len(spans)
|
||||
|
||||
def _act_on_spans(
|
||||
self,
|
||||
masked: _MaskedText,
|
||||
pattern_entry: CompiledPatternEntry,
|
||||
spans: Sequence[tuple[int, int]],
|
||||
failed_checksum: Sequence[bool],
|
||||
detections: list[ContentFilterDetection] | None,
|
||||
) -> _MaskedText:
|
||||
if not spans:
|
||||
return masked
|
||||
pattern_name: Final = pattern_entry["pattern_name"]
|
||||
if not all(failed_checksum):
|
||||
self._handle_pattern_match(pattern_name, pattern_entry["action"], detections)
|
||||
if pattern_entry["action"] != ContentFilterAction.MASK:
|
||||
return masked
|
||||
verbose_proxy_logger.info("Masked all %s matches in content", pattern_name)
|
||||
return masked.masked(
|
||||
spans, self.pattern_redaction_format.format(pattern_name=pattern_name.upper()), failed_checksum
|
||||
)
|
||||
|
||||
def _handle_blocked_word_match(
|
||||
self,
|
||||
|
|
@ -1464,6 +1697,17 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
return text
|
||||
|
||||
def _filter_single_text(self, text: str, detections: list[ContentFilterDetection] | None = None) -> str:
|
||||
return self._filter_text(text, detections, validates_checksums=True)
|
||||
|
||||
def _filter_streamed_text(self, text: str, detections: list[ContentFilterDetection] | None = None) -> str:
|
||||
return self._filter_text(text, detections, validates_checksums=False)
|
||||
|
||||
def _filter_text(
|
||||
self,
|
||||
text: str,
|
||||
detections: list[ContentFilterDetection] | None,
|
||||
validates_checksums: bool,
|
||||
) -> str:
|
||||
"""
|
||||
Apply all content filtering checks to a single text.
|
||||
|
||||
|
|
@ -1475,6 +1719,8 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
Args:
|
||||
text: Text to filter
|
||||
detections: Optional list to append detection information
|
||||
validates_checksums: Whether a match that fails its pattern's checksum is put back. Streams pass False,
|
||||
since later text can still change what a rule matches, so they mask every match as before
|
||||
|
||||
Returns:
|
||||
Filtered text (with masking applied if action is MASK)
|
||||
|
|
@ -1483,9 +1729,9 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
HTTPException: If sensitive content is detected and action is BLOCK
|
||||
"""
|
||||
# Collect all exceptions from loaded categories
|
||||
all_exceptions: Final[list[str]] = []
|
||||
for category in self.loaded_categories.values():
|
||||
all_exceptions.extend(category.exceptions)
|
||||
all_exceptions: Final = tuple(
|
||||
itertools.chain.from_iterable(category.exceptions for category in self.loaded_categories.values())
|
||||
)
|
||||
|
||||
# Check conditional categories first
|
||||
conditional_match: Final = self._check_conditional_categories(text, all_exceptions)
|
||||
|
|
@ -1506,12 +1752,7 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
text = self._handle_category_keyword_match(keyword, category_name, severity, action, text, detections)
|
||||
|
||||
# Check regex patterns - process ALL patterns, not just first match
|
||||
for pattern_entry in self.compiled_patterns:
|
||||
spans = self._find_pattern_spans(text, pattern_entry)
|
||||
if spans:
|
||||
pattern_name = pattern_entry["pattern_name"]
|
||||
action = pattern_entry["action"]
|
||||
text = self._handle_pattern_match(pattern_name, action, text, spans, detections)
|
||||
text = self._apply_patterns(text, detections, validates_checksums)
|
||||
|
||||
# Check blocked words - iterate through ALL blocked words
|
||||
text_lower = text.lower()
|
||||
|
|
@ -2124,8 +2365,8 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
return deferred
|
||||
head_detections: Final[list[ContentFilterDetection]] = [] # mutable-ok: filled by _filter_single_text
|
||||
try:
|
||||
masked_head: Final = self._filter_single_text(head, detections=head_detections)
|
||||
masked_tail: Final = self._filter_single_text(tail)
|
||||
masked_head: Final = self._filter_streamed_text(head, detections=head_detections)
|
||||
masked_tail: Final = self._filter_streamed_text(tail)
|
||||
except Exception:
|
||||
return deferred
|
||||
if masked_head + masked_tail != masked_text or len(masked_head) > state.yielded_masked_text_len:
|
||||
|
|
@ -2202,7 +2443,7 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
scan_started = time.perf_counter()
|
||||
|
||||
try:
|
||||
masked_text = self._filter_single_text(text_to_scan, detections=choice_detections)
|
||||
masked_text = self._filter_streamed_text(text_to_scan, detections=choice_detections)
|
||||
if is_final and masked_text.endswith(" "):
|
||||
masked_text = masked_text[:-1]
|
||||
latest_detections = tuple(choice_detections)
|
||||
|
|
|
|||
|
|
@ -32,36 +32,41 @@
|
|||
"name": "visa",
|
||||
"display_name": "Credit Card (Visa)",
|
||||
"pattern": "\\b4\\d{3}[\\s\\-]?\\d{4}[\\s\\-]?\\d{4}[\\s\\-]?\\d{4}\\b",
|
||||
"validator": "luhn",
|
||||
"category": "Payment Card Patterns",
|
||||
"description": "Detects Visa credit card numbers"
|
||||
"description": "Detects Visa credit card numbers that pass the Luhn checksum"
|
||||
},
|
||||
{
|
||||
"name": "mastercard",
|
||||
"display_name": "Credit Card (Mastercard)",
|
||||
"pattern": "\\b5[1-5]\\d{2}[\\s\\-]?\\d{4}[\\s\\-]?\\d{4}[\\s\\-]?\\d{4}\\b",
|
||||
"validator": "luhn",
|
||||
"category": "Payment Card Patterns",
|
||||
"description": "Detects Mastercard credit card numbers"
|
||||
"description": "Detects Mastercard credit card numbers that pass the Luhn checksum"
|
||||
},
|
||||
{
|
||||
"name": "amex",
|
||||
"display_name": "Credit Card (American Express)",
|
||||
"pattern": "\\b3[47]\\d{2}[\\s\\-]?\\d{6}[\\s\\-]?\\d{5}\\b",
|
||||
"validator": "luhn",
|
||||
"category": "Payment Card Patterns",
|
||||
"description": "Detects American Express credit card numbers"
|
||||
"description": "Detects American Express credit card numbers that pass the Luhn checksum"
|
||||
},
|
||||
{
|
||||
"name": "discover",
|
||||
"display_name": "Credit Card (Discover)",
|
||||
"pattern": "\\b6(?:011|5\\d{2})[\\s\\-]?\\d{4}[\\s\\-]?\\d{4}[\\s\\-]?\\d{4}\\b",
|
||||
"validator": "luhn",
|
||||
"category": "Payment Card Patterns",
|
||||
"description": "Detects Discover credit card numbers"
|
||||
"description": "Detects Discover credit card numbers that pass the Luhn checksum"
|
||||
},
|
||||
{
|
||||
"name": "credit_card",
|
||||
"display_name": "Credit Card (Any)",
|
||||
"pattern": "(?:\\b4\\d{3}[\\s\\-]?\\d{4}[\\s\\-]?\\d{4}[\\s\\-]?\\d{4}\\b|\\b5[1-5]\\d{2}[\\s\\-]?\\d{4}[\\s\\-]?\\d{4}[\\s\\-]?\\d{4}\\b|\\b3[47]\\d{2}[\\s\\-]?\\d{6}[\\s\\-]?\\d{5}\\b|\\b6(?:011|5\\d{2})[\\s\\-]?\\d{4}[\\s\\-]?\\d{4}[\\s\\-]?\\d{4}\\b)",
|
||||
"validator": "luhn",
|
||||
"category": "Payment Card Patterns",
|
||||
"description": "Detects any major credit card number"
|
||||
"description": "Detects any major credit card number that passes the Luhn checksum"
|
||||
},
|
||||
{
|
||||
"name": "aws_access_key",
|
||||
|
|
@ -395,8 +400,9 @@
|
|||
"name": "iban",
|
||||
"display_name": "IBAN (International Bank Account Number)",
|
||||
"pattern": "\\b[A-Z]{2}\\d{2}[A-Z0-9]{11,30}\\b",
|
||||
"validator": "iban_mod97",
|
||||
"category": "Payment Card Patterns",
|
||||
"description": "Detects IBANs (2 letter country code + 2 check digits + 4 char bank code + 7 digit base + optional 0-16 alphanumeric)"
|
||||
"description": "Detects IBANs (2 letter country code + 2 check digits + 4 char bank code + 7 digit base + optional 0-16 alphanumeric) that pass the mod-97 check"
|
||||
},
|
||||
{
|
||||
"name": "fr_nir",
|
||||
|
|
@ -409,8 +415,9 @@
|
|||
"name": "eu_iban_enhanced",
|
||||
"display_name": "IBAN (Enhanced EU Format)",
|
||||
"pattern": "\\b[A-Z]{2}[0-9]{2}[A-Z0-9]{4}[0-9]{7}[A-Z0-9]{0,16}\\b",
|
||||
"validator": "iban_mod97",
|
||||
"category": "EU PII Patterns",
|
||||
"description": "Enhanced IBAN detection with more specific format validation (2 letter country + 2 check digits + 4 char bank code + 7 digit account base + optional 0-16 alphanumeric)"
|
||||
"description": "Enhanced IBAN detection with more specific format validation (2 letter country + 2 check digits + 4 char bank code + 7 digit account base + optional 0-16 alphanumeric) that pass the mod-97 check"
|
||||
},
|
||||
{
|
||||
"name": "fr_phone",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,155 @@
|
|||
import re
|
||||
import string
|
||||
from collections.abc import Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
_SEPARATORS: Final = re.compile(r"[\s\-]")
|
||||
_IBAN_CHARACTERS: Final = re.compile(r"[0-9A-Z]{5,34}")
|
||||
_IGNORECASE_ASCII_LETTER_ALIASES: Final = str.maketrans({"\u0130": "I", "\u212a": "K"})
|
||||
_ISO_13616_IBAN_LENGTHS: Final[Mapping[str, int]] = MappingProxyType(
|
||||
{
|
||||
"AD": 24,
|
||||
"AE": 23,
|
||||
"AL": 28,
|
||||
"AT": 20,
|
||||
"AX": 18,
|
||||
"AZ": 28,
|
||||
"BA": 20,
|
||||
"BE": 16,
|
||||
"BG": 22,
|
||||
"BH": 22,
|
||||
"BI": 27,
|
||||
"BL": 27,
|
||||
"BR": 29,
|
||||
"BY": 28,
|
||||
"CH": 21,
|
||||
"CR": 22,
|
||||
"CY": 28,
|
||||
"CZ": 24,
|
||||
"DE": 22,
|
||||
"DJ": 27,
|
||||
"DK": 18,
|
||||
"DO": 28,
|
||||
"EE": 20,
|
||||
"EG": 29,
|
||||
"ES": 24,
|
||||
"FI": 18,
|
||||
"FK": 18,
|
||||
"FO": 18,
|
||||
"FR": 27,
|
||||
"GB": 22,
|
||||
"GE": 22,
|
||||
"GF": 27,
|
||||
"GG": 22,
|
||||
"GI": 23,
|
||||
"GL": 18,
|
||||
"GP": 27,
|
||||
"GR": 27,
|
||||
"GT": 28,
|
||||
"HN": 28,
|
||||
"HR": 21,
|
||||
"HU": 28,
|
||||
"IE": 22,
|
||||
"IL": 23,
|
||||
"IM": 22,
|
||||
"IQ": 23,
|
||||
"IS": 26,
|
||||
"IT": 27,
|
||||
"JE": 22,
|
||||
"JO": 30,
|
||||
"KW": 30,
|
||||
"KZ": 20,
|
||||
"LB": 28,
|
||||
"LC": 32,
|
||||
"LI": 21,
|
||||
"LT": 20,
|
||||
"LU": 20,
|
||||
"LV": 21,
|
||||
"LY": 25,
|
||||
"MC": 27,
|
||||
"MD": 24,
|
||||
"ME": 22,
|
||||
"MF": 27,
|
||||
"MK": 19,
|
||||
"MN": 20,
|
||||
"MQ": 27,
|
||||
"MR": 27,
|
||||
"MT": 31,
|
||||
"MU": 30,
|
||||
"NC": 27,
|
||||
"NI": 28,
|
||||
"NL": 18,
|
||||
"NO": 15,
|
||||
"OM": 23,
|
||||
"PF": 27,
|
||||
"PK": 24,
|
||||
"PL": 28,
|
||||
"PM": 27,
|
||||
"PS": 29,
|
||||
"PT": 25,
|
||||
"QA": 29,
|
||||
"RE": 27,
|
||||
"RO": 24,
|
||||
"RS": 22,
|
||||
"RU": 33,
|
||||
"SA": 24,
|
||||
"SC": 31,
|
||||
"SD": 18,
|
||||
"SE": 24,
|
||||
"SI": 19,
|
||||
"SK": 24,
|
||||
"SM": 27,
|
||||
"SO": 23,
|
||||
"ST": 25,
|
||||
"SV": 28,
|
||||
"TF": 27,
|
||||
"TL": 23,
|
||||
"TN": 24,
|
||||
"TR": 26,
|
||||
"UA": 29,
|
||||
"VA": 22,
|
||||
"VG": 24,
|
||||
"WF": 27,
|
||||
"XK": 20,
|
||||
"YE": 30,
|
||||
"YT": 27,
|
||||
}
|
||||
)
|
||||
_IBAN_LETTER_VALUES: Final = str.maketrans(
|
||||
{letter: str(value) for value, letter in enumerate(string.ascii_uppercase, 10)}
|
||||
)
|
||||
_DIGIT_VALUES: Final = bytes.maketrans(b"0123456789", bytes(range(10)))
|
||||
_DOUBLED_DIGIT_SUMS: Final = bytes.maketrans(b"0123456789", bytes((0, 2, 4, 6, 8, 1, 3, 5, 7, 9)))
|
||||
|
||||
|
||||
def _to_ascii(value: str) -> str:
|
||||
if value.isascii():
|
||||
return value
|
||||
folded: Final = value.translate(_IGNORECASE_ASCII_LETTER_ALIASES)
|
||||
return "".join(str(int(character)) if character.isdecimal() else character for character in folded)
|
||||
|
||||
|
||||
def passes_luhn(value: str) -> bool:
|
||||
digits: Final = _SEPARATORS.sub("", value)
|
||||
if not digits.isdecimal():
|
||||
return False
|
||||
reversed_digits: Final = _to_ascii(digits)[::-1].encode()
|
||||
checksum: Final = sum(reversed_digits[::2].translate(_DIGIT_VALUES)) + sum(
|
||||
reversed_digits[1::2].translate(_DOUBLED_DIGIT_SUMS)
|
||||
)
|
||||
return checksum % 10 == 0
|
||||
|
||||
|
||||
def passes_iban_mod97(value: str) -> bool:
|
||||
compact: Final = _to_ascii(_SEPARATORS.sub("", value)).upper()
|
||||
iban: Final = compact[: _ISO_13616_IBAN_LENGTHS.get(compact[:2], len(compact))]
|
||||
if _IBAN_CHARACTERS.fullmatch(iban) is None:
|
||||
return False
|
||||
rearranged: Final = iban[4:] + iban[:4]
|
||||
return int(rearranged.translate(_IBAN_LETTER_VALUES)) % 97 == 1
|
||||
|
||||
|
||||
PATTERN_VALIDATORS: Final[Mapping[str, Callable[[str], bool]]] = MappingProxyType(
|
||||
{"luhn": passes_luhn, "iban_mod97": passes_iban_mod97}
|
||||
)
|
||||
|
|
@ -44,6 +44,10 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
|
||||
def _in_chunks(text: str, size: int) -> tuple[str, ...]:
|
||||
return tuple(text[start : start + size] for start in range(0, len(text), size))
|
||||
|
||||
|
||||
class TestContentFilterGuardrail:
|
||||
"""Test the ContentFilterGuardrail class"""
|
||||
|
||||
|
|
@ -378,10 +382,521 @@ class TestContentFilterGuardrail:
|
|||
)
|
||||
|
||||
# Test Visa card
|
||||
result = guardrail._check_patterns("My card is 4532-1234-5678-9010")
|
||||
result = guardrail._check_patterns("My card is 4532-0151-1283-0366")
|
||||
assert result is not None
|
||||
assert result[1] == "visa"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("pattern_name", "checksum_valid", "checksum_invalid"),
|
||||
[
|
||||
("credit_card", "4532 0151 1283 0366", "4532 0151 1283 0367"),
|
||||
("visa", "4532-0151-1283-0366", "4532-1234-5678-9010"),
|
||||
("mastercard", "5555555555554444", "5555555555554445"),
|
||||
("amex", "3782 822463 10005", "3782 822463 10006"),
|
||||
("discover", "6011111111111117", "6011111111111118"),
|
||||
("iban", "GB82WEST12345698765432", "GB83WEST12345698765432"),
|
||||
("eu_iban_enhanced", "DE89370400440532013000", "DE88370400440532013000"),
|
||||
],
|
||||
)
|
||||
async def test_checksum_pattern_masks_only_values_that_pass_the_checksum(
|
||||
self, pattern_name: str, checksum_valid: str, checksum_invalid: str
|
||||
):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-checksum",
|
||||
patterns=[
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name=pattern_name, action=ContentFilterAction.MASK)
|
||||
],
|
||||
)
|
||||
|
||||
guardrailed_inputs: Final = await guardrail.apply_guardrail(
|
||||
inputs={"texts": [f"real {checksum_valid} and look-alike {checksum_invalid}"]},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert guardrailed_inputs.get("texts") == [
|
||||
f"real [{pattern_name.upper()}_REDACTED] and look-alike {checksum_invalid}"
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("PIN 4821 4532 0151 1283 0366 thanks", "PIN [CREDIT_CARD_REDACTED] 0366 thanks"),
|
||||
("order 4111 4532 0151 1283 0366", "order [CREDIT_CARD_REDACTED] 0366"),
|
||||
("codes 4532 4532 4532 4532 4532 0151 1283 0366 thanks", "codes 4532 4532 4532 4532 [CREDIT_CARD_REDACTED] thanks"),
|
||||
("order 4532 1234 5678 9010 paid by 4532 0151 1283 0366", "order 4532 1234 5678 9010 paid by [CREDIT_CARD_REDACTED]"),
|
||||
("card 4532 4151 4283 4000 4777 4888 4999 ok", "card [CREDIT_CARD_REDACTED] 4777 4888 4999 ok"),
|
||||
("card 4532 4111 1111 1116 2025 ok", "card [CREDIT_CARD_REDACTED] 2025 ok"),
|
||||
("VISA4000 4000 0000 0004 1111 ok", "VISA4000 [CREDIT_CARD_REDACTED] ok"),
|
||||
],
|
||||
)
|
||||
async def test_checksum_pattern_keeps_the_mask_on_a_match_that_overlaps_a_valid_card(self, text: str, expected: str):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-checksum-overlap",
|
||||
patterns=[
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK)
|
||||
],
|
||||
)
|
||||
|
||||
guardrailed_inputs: Final = await guardrail.apply_guardrail(
|
||||
inputs={"texts": [text]},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert guardrailed_inputs.get("texts") == [expected]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("ids 4853 4091 4909 4475 5508 5016 6232 6258 end", "ids [VISA_REDACTED] [MASTERCARD_REDACTED] end"),
|
||||
("ids 4336 4262 4755 4637 3753 786027 13221 end", "ids [VISA_REDACTED] [AMEX_REDACTED] end"),
|
||||
("4000 4000 0000 0000\n5555 5555 5555 4444", "[VISA_REDACTED]\n[MASTERCARD_REDACTED]"),
|
||||
("cards 4378 4937 4618 4485 5555 5555 5555 4444 ok", "cards [VISA_REDACTED] [MASTERCARD_REDACTED] ok"),
|
||||
("pay 5472 3038 5280 1841 1485 4534 3853 6336 1330 ok", "pay [MASTERCARD_REDACTED] 1485 [VISA_REDACTED] ok"),
|
||||
("cards 4532 0151 1283 0366 5555 5555 5555 4444 ok", "cards [VISA_REDACTED] [MASTERCARD_REDACTED] ok"),
|
||||
(
|
||||
"5521-5568-8763-1997 4966 7452 4225 5188 4264 ref 580355",
|
||||
"[MASTERCARD_REDACTED] [VISA_REDACTED] 4264 ref 580355",
|
||||
),
|
||||
("ref XX00DE89370400440532013000 ok", "ref [IBAN_REDACTED] ok"),
|
||||
],
|
||||
)
|
||||
async def test_checksum_patterns_mask_each_valid_card_where_two_brands_overlap(self, text: str, expected: str):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-checksum-cross-pattern",
|
||||
patterns=[
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name=pattern_name, action=ContentFilterAction.MASK)
|
||||
for pattern_name in ("visa", "mastercard", "amex", "discover", "credit_card", "iban")
|
||||
],
|
||||
)
|
||||
|
||||
guardrailed_inputs: Final = await guardrail.apply_guardrail(
|
||||
inputs={"texts": [text]},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert guardrailed_inputs.get("texts") == [expected]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("masking_pattern", "blocking_pattern", "text"),
|
||||
[
|
||||
("visa", "mastercard", "4004 5555 5555 5557 5555 5555 5555 4444"),
|
||||
("visa", "mastercard", "4000 4000 0000 0000 5555 5555 5555 4444"),
|
||||
("visa", "mastercard", "ids 4853 4091 4909 4475 5508 5016 6232 6258 end"),
|
||||
("discover", "credit_card", "4629 4611-6260-8050-5592-2700 65405563 93269332"),
|
||||
("visa", "mastercard", "pay 5472 3038 5280 1841 1485 4534 3853 6336 1330 ok"),
|
||||
("credit_card", "nl_bsn_contextual", "BSN 4111 1111 1111 1112 123456789"),
|
||||
],
|
||||
)
|
||||
async def test_checksum_mask_leaves_a_later_blocking_pattern_its_match(
|
||||
self, masking_pattern: str, blocking_pattern: str, text: str
|
||||
):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-checksum-mask-then-block",
|
||||
patterns=[
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name=masking_pattern, action=ContentFilterAction.MASK
|
||||
),
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name=blocking_pattern, action=ContentFilterAction.BLOCK
|
||||
),
|
||||
],
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.apply_guardrail(inputs={"texts": [text]}, request_data={}, input_type="request")
|
||||
assert exc_info.value.detail["pattern"] == blocking_pattern
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checksum_mask_before_a_blocking_pattern_still_masks_the_card_both_match(self):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-checksum-mask-before-block",
|
||||
patterns=[
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK
|
||||
),
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name="visa", action=ContentFilterAction.BLOCK),
|
||||
],
|
||||
)
|
||||
guardrailed_inputs: Final = await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["card 4111 1111 1111 1111 ok"]}, request_data={}, input_type="request"
|
||||
)
|
||||
assert guardrailed_inputs.get("texts") == ["card [CREDIT_CARD_REDACTED] ok"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("later_pattern", "text", "expected"),
|
||||
[
|
||||
(
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="nl_bsn_contextual", action=ContentFilterAction.MASK
|
||||
),
|
||||
"BSN 4111 1111 1111 1112 123456789",
|
||||
"BSN 4111 1111 1111 1112 [NL_BSN_CONTEXTUAL_REDACTED]",
|
||||
),
|
||||
(
|
||||
ContentFilterPattern(
|
||||
pattern_type="regex", pattern=r"\b\d{4}[\s-]?\d{4}\b", name="acct8", action=ContentFilterAction.MASK
|
||||
),
|
||||
"4381 4099 4560 4729 4064 4577 4061",
|
||||
"[CREDIT_CARD_REDACTED] [ACCT8_REDACTED] 4061",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_later_patterns_run_on_the_text_with_every_checksum_match_masked(
|
||||
self, later_pattern: ContentFilterPattern, text: str, expected: str
|
||||
):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-checksum-mask-then-later-pattern",
|
||||
patterns=[
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK),
|
||||
later_pattern,
|
||||
],
|
||||
)
|
||||
guardrailed_inputs: Final = await guardrail.apply_guardrail(
|
||||
inputs={"texts": [text]}, request_data={}, input_type="request"
|
||||
)
|
||||
assert guardrailed_inputs.get("texts") == [expected]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("later_pattern", "text", "expected"),
|
||||
[
|
||||
(
|
||||
ContentFilterPattern(
|
||||
pattern_type="regex", pattern=r"\b\d{4}\b", name="group4", action=ContentFilterAction.MASK
|
||||
),
|
||||
"card 4111 1111 1111 1111 and 4111 1111 1111 1112 end",
|
||||
"card [CREDIT_CARD_REDACTED] and [GROUP4_REDACTED] [GROUP4_REDACTED] [GROUP4_REDACTED] [GROUP4_REDACTED] end",
|
||||
),
|
||||
(
|
||||
ContentFilterPattern(
|
||||
pattern_type="regex", pattern=r"\b\d{4}[\s-]?\d{4}\b", name="acct8", action=ContentFilterAction.MASK
|
||||
),
|
||||
"order 4398 2296 8107 4767 4479 3291",
|
||||
"order [ACCT8_REDACTED] [ACCT8_REDACTED] [ACCT8_REDACTED]",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_later_patterns_mask_a_number_put_back_after_failing_the_checksum(
|
||||
self, later_pattern: ContentFilterPattern, text: str, expected: str
|
||||
):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-checksum-mask-then-later-pattern",
|
||||
patterns=[
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK),
|
||||
later_pattern,
|
||||
],
|
||||
)
|
||||
guardrailed_inputs: Final = await guardrail.apply_guardrail(
|
||||
inputs={"texts": [text]}, request_data={}, input_type="request"
|
||||
)
|
||||
assert guardrailed_inputs.get("texts") == [expected]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_later_blocking_pattern_blocks_a_number_put_back_after_failing_the_checksum(self):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-checksum-mask-then-later-block",
|
||||
patterns=[
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK),
|
||||
ContentFilterPattern(
|
||||
pattern_type="regex", pattern=r"\b\d{4}\b", name="group4", action=ContentFilterAction.BLOCK
|
||||
),
|
||||
],
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["order 4111 1111 1111 1112"]}, request_data={}, input_type="request"
|
||||
)
|
||||
assert exc_info.value.detail["pattern"] == "group4"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_later_pattern_counts_once_when_it_also_masks_a_number_put_back(self):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-checksum-mask-then-later-pattern-count",
|
||||
patterns=[
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK),
|
||||
ContentFilterPattern(
|
||||
pattern_type="regex", pattern=r"\b\d{4}\b", name="group4", action=ContentFilterAction.MASK
|
||||
),
|
||||
],
|
||||
)
|
||||
request_data: Final[dict[str, dict[str, list[StandardLoggingGuardrailInformation]]]] = {"metadata": {}}
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["pin 1234 and card 4111 1111 1111 1112 end"]},
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
entry: Final = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
||||
assert entry["masked_entity_count"] == {"group4": 1}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_later_pattern_sharing_a_name_counts_its_matches_on_a_number_put_back(self):
|
||||
later_patterns: Final = (
|
||||
ContentFilterPattern(pattern_type="regex", pattern=r"\bpin\b", action=ContentFilterAction.MASK),
|
||||
ContentFilterPattern(pattern_type="regex", pattern=r"\b\d{4}\b", action=ContentFilterAction.MASK),
|
||||
)
|
||||
card_rule: Final = ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK
|
||||
)
|
||||
|
||||
async def masked_entity_count(patterns: tuple[ContentFilterPattern, ...]) -> dict[str, int]:
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-checksum-shared-name-count", patterns=list(patterns)
|
||||
)
|
||||
request_data: Final[dict[str, dict[str, list[StandardLoggingGuardrailInformation]]]] = {"metadata": {}}
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["pin and card 4111 1111 1111 1112 end"]},
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
return request_data["metadata"]["standard_logging_guardrail_information"][0]["masked_entity_count"]
|
||||
|
||||
assert await masked_entity_count((card_rule, *later_patterns)) == await masked_entity_count(later_patterns)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("pattern_name", "checksum_valid", "checksum_invalid"),
|
||||
[
|
||||
("visa", "4111 1111 1111 1111", "4111 1111 1111 1112"),
|
||||
("mastercard", "5555555555554444", "5555555555554445"),
|
||||
("iban", "GB82WEST12345698765432", "GB83WEST12345698765432"),
|
||||
],
|
||||
)
|
||||
async def test_checksum_pattern_blocks_only_values_that_pass_the_checksum(
|
||||
self, pattern_name: str, checksum_valid: str, checksum_invalid: str
|
||||
):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-checksum-block",
|
||||
patterns=[
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name=pattern_name, action=ContentFilterAction.BLOCK)
|
||||
],
|
||||
)
|
||||
|
||||
passed: Final = await guardrail.apply_guardrail(
|
||||
inputs={"texts": [f"look-alike {checksum_invalid}"]}, request_data={}, input_type="request"
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": [f"real {checksum_valid}"]}, request_data={}, input_type="request"
|
||||
)
|
||||
|
||||
assert passed.get("texts") == [f"look-alike {checksum_invalid}"]
|
||||
assert exc_info.value.detail["pattern"] == pattern_name
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("patterns", "blocked_words", "chunks", "hidden"),
|
||||
[
|
||||
(
|
||||
(
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK
|
||||
),
|
||||
),
|
||||
(),
|
||||
("4532 " * 20, "4532 0151 1283 0366 thanks"),
|
||||
"0366",
|
||||
),
|
||||
(
|
||||
(
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name="iban", action=ContentFilterAction.MASK),
|
||||
ContentFilterPattern(
|
||||
pattern_type="regex", pattern="904580", name="ref", action=ContentFilterAction.BLOCK
|
||||
),
|
||||
),
|
||||
(),
|
||||
("send it to DE07841266521904580", "435 thanks"),
|
||||
"904580",
|
||||
),
|
||||
(
|
||||
(
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK
|
||||
),
|
||||
ContentFilterPattern(
|
||||
pattern_type="regex", pattern=r"\b\d{4}\b", name="group4", action=ContentFilterAction.MASK
|
||||
),
|
||||
),
|
||||
(),
|
||||
("card 4111 1111 11", "11 1112", " end"),
|
||||
"1112",
|
||||
),
|
||||
(
|
||||
(
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK
|
||||
),
|
||||
ContentFilterPattern(
|
||||
pattern_type="regex", pattern=r"\b\d{4}\b", name="group4", action=ContentFilterAction.MASK
|
||||
),
|
||||
),
|
||||
(BlockedWord(keyword="ok", action=ContentFilterAction.MASK),),
|
||||
_in_chunks("card 4111 1111 1111 1112 ok ok ok ok ok ok ok ok ok ok ok ok and more text after it end", 8),
|
||||
"1112",
|
||||
),
|
||||
(
|
||||
(
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK
|
||||
),
|
||||
ContentFilterPattern(
|
||||
pattern_type="regex", pattern=r"\b\d{4}[\s-]?\d{4}\b", name="acct8", action=ContentFilterAction.MASK
|
||||
),
|
||||
),
|
||||
(BlockedWord(keyword="ok", action=ContentFilterAction.MASK),),
|
||||
_in_chunks("order 4398 2296 8107 4767 4479 3291 ok ok ok ok ok ok ok ok ok ok ok ok and more text after it end", 8),
|
||||
"4767",
|
||||
),
|
||||
(
|
||||
(
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name="amex", action=ContentFilterAction.MASK),
|
||||
ContentFilterPattern(pattern_type="regex", pattern=r"\d{6}", name="d6", action=ContentFilterAction.MASK),
|
||||
ContentFilterPattern(
|
||||
pattern_type="regex", pattern=r"\d{4} \d{4}", name="pair", action=ContentFilterAction.MASK
|
||||
),
|
||||
),
|
||||
(),
|
||||
_in_chunks("FR27 5892 9262 5322 3897 2114 796 it 3797 621746 14951 3722 314364 73911 textok 4336 2547 3487 7245 safe and more words", 8),
|
||||
"621746",
|
||||
),
|
||||
(
|
||||
(
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK
|
||||
),
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="uoft_tcard", action=ContentFilterAction.MASK
|
||||
),
|
||||
),
|
||||
(),
|
||||
_in_chunks("id 4111111111111112" + "\n" * 60 + "tcard holder ok", 8),
|
||||
"4111111111111112",
|
||||
),
|
||||
(
|
||||
(
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK
|
||||
),
|
||||
ContentFilterPattern(
|
||||
pattern_type="regex",
|
||||
pattern=r"\b\d{4}(?: \d{4}){3}\b(?=.{0,90}\bcvv\b)",
|
||||
name="card_with_cvv",
|
||||
action=ContentFilterAction.BLOCK,
|
||||
),
|
||||
),
|
||||
(),
|
||||
_in_chunks(
|
||||
"the number on file is 4111 1111 1111 1112, and the three digits printed on the back are the cvv 123",
|
||||
4,
|
||||
),
|
||||
"1112",
|
||||
),
|
||||
(
|
||||
(
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK
|
||||
),
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="au_medicare", action=ContentFilterAction.BLOCK
|
||||
),
|
||||
),
|
||||
(),
|
||||
_in_chunks(
|
||||
"ref 8834 6011 1484 9325 7256 904580\nmore text here to push it along, please keep going for a while now",
|
||||
3,
|
||||
),
|
||||
"6011",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_checksum_pattern_stream_sends_no_digit_of_a_number_it_matched(
|
||||
self,
|
||||
patterns: tuple[ContentFilterPattern, ...],
|
||||
blocked_words: tuple[BlockedWord, ...],
|
||||
chunks: tuple[str, ...],
|
||||
hidden: str,
|
||||
):
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-checksum-streaming", patterns=list(patterns), blocked_words=list(blocked_words) or None
|
||||
)
|
||||
|
||||
async def mock_stream():
|
||||
for index, content in enumerate(chunks):
|
||||
yield ModelResponseStream(
|
||||
id=f"chunk{index}",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
delta=Delta(content=content),
|
||||
index=0,
|
||||
finish_reason="stop" if index == len(chunks) - 1 else None,
|
||||
)
|
||||
],
|
||||
model="gpt-4",
|
||||
)
|
||||
|
||||
streamed: Final = "".join(
|
||||
[
|
||||
chunk.choices[0].delta.content or ""
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=MagicMock(), response=mock_stream(), request_data={}
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert hidden not in streamed
|
||||
assert streamed.endswith(chunks[-1][-2:])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"iban",
|
||||
[
|
||||
"GB66\u212aEST12345698765432",
|
||||
"GB73W\u0130ST12345698765432",
|
||||
"GB\u0668\u0662WEST12345698765432",
|
||||
],
|
||||
)
|
||||
async def test_checksum_pattern_masks_a_valid_iban_with_characters_the_regex_case_folds(self, iban: str):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-checksum-folding",
|
||||
patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="iban", action=ContentFilterAction.MASK)],
|
||||
)
|
||||
|
||||
guardrailed_inputs: Final = await guardrail.apply_guardrail(
|
||||
inputs={"texts": [f"send it to {iban}"]},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert guardrailed_inputs.get("texts") == ["send it to [IBAN_REDACTED]"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("pattern_name", ["iban", "eu_iban_enhanced"])
|
||||
async def test_checksum_pattern_masks_a_valid_iban_with_characters_attached_to_its_end(self, pattern_name: str):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-checksum-attached",
|
||||
patterns=[
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name=pattern_name, action=ContentFilterAction.MASK)
|
||||
],
|
||||
)
|
||||
|
||||
guardrailed_inputs: Final = await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["pay DE89370400440532013000EUR or DE88370400440532013000EUR"]},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert guardrailed_inputs.get("texts") == [f"pay [{pattern_name.upper()}_REDACTED] or DE88370400440532013000EUR"]
|
||||
|
||||
def test_api_key_patterns(self):
|
||||
"""
|
||||
Test API key pattern detection
|
||||
|
|
@ -964,13 +1479,13 @@ class TestContentFilterGuardrail:
|
|||
scanned_lengths: Final[list[int]] = []
|
||||
|
||||
class RecordingGuardrail(ContentFilterGuardrail):
|
||||
def _filter_single_text(
|
||||
def _filter_streamed_text(
|
||||
self,
|
||||
text: str,
|
||||
detections: list[ContentFilterDetection] | None = None,
|
||||
) -> str:
|
||||
scanned_lengths.append(len(text))
|
||||
return super()._filter_single_text(text, detections=detections)
|
||||
return super()._filter_streamed_text(text, detections=detections)
|
||||
|
||||
guardrail: Final = RecordingGuardrail(
|
||||
guardrail_name="test-streaming-bounded-scan",
|
||||
|
|
@ -1006,13 +1521,13 @@ class TestContentFilterGuardrail:
|
|||
scanned_lengths: Final[list[int]] = []
|
||||
|
||||
class RecordingGuardrail(ContentFilterGuardrail):
|
||||
def _filter_single_text(
|
||||
def _filter_streamed_text(
|
||||
self,
|
||||
text: str,
|
||||
detections: list[ContentFilterDetection] | None = None,
|
||||
) -> str:
|
||||
scanned_lengths.append(len(text))
|
||||
return super()._filter_single_text(text, detections=detections)
|
||||
return super()._filter_streamed_text(text, detections=detections)
|
||||
|
||||
guardrail: Final = RecordingGuardrail(
|
||||
guardrail_name="test-streaming-refused-cut-backoff",
|
||||
|
|
|
|||
|
|
@ -1938,6 +1938,25 @@ async def test_logging_only_block_records_every_match_as_flagged(
|
|||
assert response_entry["guardrail_status"] == "success"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logging_only_flags_an_unnamed_block_rule_on_a_number_put_back_after_failing_the_checksum() -> None:
|
||||
entries: Final = await _log_call(
|
||||
_observing_filter(
|
||||
patterns=(
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="credit_card", action=ContentFilterAction.MASK
|
||||
),
|
||||
ContentFilterPattern(pattern_type="regex", pattern="secret", action=ContentFilterAction.MASK),
|
||||
ContentFilterPattern(pattern_type="regex", pattern=r"\b\d{4}\b", action=ContentFilterAction.BLOCK),
|
||||
)
|
||||
),
|
||||
"secret 4111 1111 1111 1112",
|
||||
Message(role="assistant", content="here you go"),
|
||||
)
|
||||
|
||||
assert entries[0]["guardrail_status"] == "guardrail_flagged"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logging_only_block_scans_the_response_after_a_request_match() -> None:
|
||||
entries: Final = await _log_call(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,53 @@
|
|||
from itertools import product
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.validators import (
|
||||
passes_iban_mod97,
|
||||
passes_luhn,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("value", "expected"),
|
||||
[
|
||||
("4111 1111 1111 1111", True),
|
||||
("4111-1111-1111-1112", False),
|
||||
("٤" + "١" * 15, True),
|
||||
("4111 1111 1111 111x", False),
|
||||
("", False),
|
||||
],
|
||||
)
|
||||
def test_passes_luhn(value: str, expected: bool):
|
||||
assert passes_luhn(value) is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"valid_number",
|
||||
["4111111111111111", "4012888888881881", "5555555555554444", "378282246310005", "6011000990139424"],
|
||||
)
|
||||
def test_passes_luhn_rejects_every_single_digit_change(valid_number: str):
|
||||
single_digit_changes: Final = tuple(
|
||||
valid_number[:position] + replacement + valid_number[position + 1 :]
|
||||
for position, replacement in product(range(len(valid_number)), "0123456789")
|
||||
)
|
||||
|
||||
assert passes_luhn(valid_number)
|
||||
assert not any(passes_luhn(changed) for changed in single_digit_changes if changed != valid_number)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("value", "expected"),
|
||||
[
|
||||
("GB82 WEST 1234 5698 7654 32", True),
|
||||
("gb82west12345698765432", True),
|
||||
("GB82WEST12345698765433", False),
|
||||
("GB82", False),
|
||||
("GB82WEST1234569876543é", False),
|
||||
("DE89370400440532013000EUR", True),
|
||||
("DE88370400440532013000EUR", False),
|
||||
],
|
||||
)
|
||||
def test_passes_iban_mod97(value: str, expected: bool):
|
||||
assert passes_iban_mod97(value) is expected
|
||||
Loading…
Add table
Reference in a new issue