mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
feat(guardrails): add prebuilt credential patterns for OpenAI, Anthropic, and Gemini (#50268)
* feat(guardrails): add prebuilt credential patterns for JWT, and OpenAI, Anthropic, Gemini API keys * fix(guardrails): prevent quadratic backtracking in jwt_token pattern * test(guardrails): move credential pattern tests into test_patterns.py * feat(guardrails): remove jwt_token to keep credential patterns strictly bounded * fix(guardrails): mask whole provider keys and keep openai_api_key to OpenAI's key shapes End each pattern so a key ending in "-" is masked in full, require OpenAI's T3BlbkFJ marker on bare sk- keys so other vendors' sk- keys and LiteLLM virtual keys are left alone, add sk-None- user keys, and replace the registration-only tests with whole-match, length-boundary, and negative cases * fix(guardrails): require the T3BlbkFJ marker on every OpenAI key form and match keys after escaped separators openai_api_key now needs the marker on the sk-proj-, sk-svcacct-, sk-admin- and sk-None- forms too, as gitleaks and TruffleHog do. anthropic_api_key and gemini_api_key also start a match right after a literal \n or \t or a percent-encoded byte, so keys in pasted JSON strings and URL-encoded logs are masked. * fix(guardrails): bound the OpenAI key regex, match Gemini prefixes case-sensitively, and hold back long keys in streamed replies * fix(guardrails): size the streaming holdback per pattern, stop OpenAI key tails before the next key, and allow non-ASCII text after a key * fix(guardrails): mask content filter pattern matches found in the original text so overlapping patterns cannot split a key * fix(guardrails): let a block pattern see the text after earlier masking patterns, as before Masking patterns still match the original text and merge overlapping matches into one tag, but a BLOCK pattern checks the text with the masking patterns listed before it applied, so a request main masked and passed is not blocked now. Detections name only the patterns whose tag was written. * fix(guardrails): keep content filter masks in order and hold back the longest OpenAI key Masking patterns run in configured order on the masked text as on main, so a contextual pattern still matches next to an earlier tag, and each also matches the original text so an earlier tag cannot split a key. Detections follow the configured pattern, not its name, and the OpenAI streaming holdback covers the pattern's 419-character longest match. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Yogeshwaran Ravichandran
GitHub
parent
5bc25ff8ea
commit
8d9ddb349e
5 changed files with 468 additions and 36 deletions
|
|
@ -6,15 +6,17 @@ to detect and block/mask sensitive content.
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import bisect
|
||||
import contextvars
|
||||
import itertools
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, AsyncIterable, Callable, Coroutine, Mapping, Sequence
|
||||
from collections.abc import AsyncGenerator, AsyncIterable, Callable, Coroutine, Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass, replace
|
||||
from datetime import datetime
|
||||
from functools import partial
|
||||
from re import Pattern
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, cast
|
||||
|
||||
|
|
@ -79,7 +81,7 @@ from .competitor_intent import (
|
|||
AirlineCompetitorIntentChecker,
|
||||
BaseCompetitorIntentChecker,
|
||||
)
|
||||
from .patterns import PATTERN_EXTRA_CONFIG, get_compiled_pattern
|
||||
from .patterns import PATTERN_EXTRA_CONFIG, PATTERN_STREAMING_HOLDBACK_CHARS, get_compiled_pattern
|
||||
|
||||
MAX_KEYWORD_VALUE_GAP_WORDS: Final = 1
|
||||
_DECODED_ARGUMENTS: Final = TypeAdapter(object)
|
||||
|
|
@ -223,10 +225,70 @@ async def _stream_then_end_of_stream_chunk(
|
|||
@dataclass(frozen=True, slots=True)
|
||||
class _StreamedScanPlan:
|
||||
context_chars: int
|
||||
holdback_chars: int
|
||||
exception_phrases: tuple[str, ...]
|
||||
conditional_words: tuple[str, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RedactionSpan:
|
||||
start: int
|
||||
end: int
|
||||
pattern_index: int
|
||||
pattern_name: str
|
||||
|
||||
|
||||
def _merge_redaction_spans(spans: Iterable[_RedactionSpan]) -> tuple[_RedactionSpan, ...]:
|
||||
ordered: Final = sorted(spans, key=lambda span: (span.start, -span.end))
|
||||
reach: Final = tuple(itertools.accumulate((span.end for span in ordered), max))
|
||||
first_indices: Final = tuple(
|
||||
index for index, span in enumerate(ordered) if index == 0 or span.start >= reach[index - 1]
|
||||
)
|
||||
last_indices: Final = (*(index - 1 for index in first_indices[1:]), len(ordered) - 1)
|
||||
return tuple(replace(ordered[first], end=reach[last]) for first, last in zip(first_indices, last_indices))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RedactedText:
|
||||
spans: tuple[_RedactionSpan, ...]
|
||||
tags: tuple[str, ...]
|
||||
tag_starts: tuple[int, ...]
|
||||
|
||||
def render(self, text: str) -> str:
|
||||
previous_ends: Final = (0, *(span.end for span in self.spans))
|
||||
kept_parts: Final = (
|
||||
text[previous_end : span.start] + tag
|
||||
for previous_end, span, tag in zip(previous_ends, self.spans, self.tags)
|
||||
)
|
||||
return "".join((*kept_parts, text[previous_ends[-1] :]))
|
||||
|
||||
def original_start(self, position: int) -> int:
|
||||
index: Final = bisect.bisect_right(self.tag_starts, position) - 1
|
||||
if index < 0:
|
||||
return position
|
||||
tag_end: Final = self.tag_starts[index] + len(self.tags[index])
|
||||
if position < tag_end:
|
||||
return self.spans[index].start
|
||||
return position - tag_end + self.spans[index].end
|
||||
|
||||
def original_end(self, position: int) -> int:
|
||||
index: Final = bisect.bisect_left(self.tag_starts, position) - 1
|
||||
if index < 0:
|
||||
return position
|
||||
tag_end: Final = self.tag_starts[index] + len(self.tags[index])
|
||||
if position <= tag_end:
|
||||
return self.spans[index].end
|
||||
return position - tag_end + self.spans[index].end
|
||||
|
||||
|
||||
def _redacted_text(spans: tuple[_RedactionSpan, ...], redaction_format: str) -> _RedactedText:
|
||||
tags: Final = tuple(redaction_format.format(pattern_name=span.pattern_name.upper()) for span in spans)
|
||||
shifts: Final = itertools.accumulate(
|
||||
(len(tag) - (span.end - span.start) for span, tag in zip(spans, tags)), initial=0
|
||||
)
|
||||
return _RedactedText(spans, tags, tuple(span.start + shift for span, shift in zip(spans, shifts)))
|
||||
|
||||
|
||||
class CategoryFileData(TypedDict, total=False):
|
||||
category_name: str
|
||||
description: str
|
||||
|
|
@ -987,21 +1049,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."""
|
||||
|
||||
|
|
@ -1392,10 +1439,8 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
self,
|
||||
pattern_name: str,
|
||||
action: ContentFilterAction,
|
||||
text: str,
|
||||
spans: list[tuple[int, int]],
|
||||
detections: list[ContentFilterDetection] | None,
|
||||
) -> str:
|
||||
) -> None:
|
||||
"""Handle regex pattern match detection and action."""
|
||||
if detections is not None:
|
||||
pattern_detection: Final[PatternDetection] = {
|
||||
|
|
@ -1413,11 +1458,49 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
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 _mask_step(
|
||||
self,
|
||||
text: str,
|
||||
spans: tuple[_RedactionSpan, ...],
|
||||
indexed_entry: tuple[int, CompiledPatternEntry],
|
||||
) -> tuple[_RedactionSpan, ...]:
|
||||
index, pattern_entry = indexed_entry
|
||||
if pattern_entry["action"] != ContentFilterAction.MASK:
|
||||
return spans
|
||||
original_matches: Final = self._find_pattern_spans(text, pattern_entry)
|
||||
redacted_matches: Final = self._redacted_text_matches(text, spans, pattern_entry) if spans else ()
|
||||
return _merge_redaction_spans(
|
||||
(
|
||||
*spans,
|
||||
*(
|
||||
_RedactionSpan(start, end, index, pattern_entry["pattern_name"])
|
||||
for start, end in (*original_matches, *redacted_matches)
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
def _redacted_text_matches(
|
||||
self, text: str, spans: tuple[_RedactionSpan, ...], pattern_entry: CompiledPatternEntry
|
||||
) -> tuple[tuple[int, int], ...]:
|
||||
redacted: Final = _redacted_text(spans, self.pattern_redaction_format)
|
||||
return tuple(
|
||||
(redacted.original_start(start), redacted.original_end(end))
|
||||
for start, end in self._find_pattern_spans(redacted.render(text), pattern_entry)
|
||||
)
|
||||
|
||||
def _pattern_detected(
|
||||
self,
|
||||
text: str,
|
||||
index: int,
|
||||
pattern_entry: CompiledPatternEntry,
|
||||
span_states: Sequence[tuple[_RedactionSpan, ...]],
|
||||
) -> bool:
|
||||
if pattern_entry["action"] == ContentFilterAction.MASK:
|
||||
return any(span.pattern_index == index for span in span_states[-1])
|
||||
text_after_earlier_masks: Final = _redacted_text(span_states[index], self.pattern_redaction_format).render(text)
|
||||
return bool(self._find_pattern_spans(text_after_earlier_masks, pattern_entry))
|
||||
|
||||
def _handle_blocked_word_match(
|
||||
self,
|
||||
|
|
@ -1506,13 +1589,14 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
keyword, category_name, severity, action = category_keyword_match
|
||||
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)
|
||||
no_spans: Final[tuple[_RedactionSpan, ...]] = ()
|
||||
span_states: Final = tuple(
|
||||
itertools.accumulate(enumerate(self.compiled_patterns), partial(self._mask_step, text), initial=no_spans)
|
||||
)
|
||||
for index, pattern_entry in enumerate(self.compiled_patterns):
|
||||
if self._pattern_detected(text, index, pattern_entry, span_states):
|
||||
self._handle_pattern_match(pattern_entry["pattern_name"], pattern_entry["action"], detections)
|
||||
text = _redacted_text(span_states[-1], self.pattern_redaction_format).render(text)
|
||||
|
||||
# Check blocked words - iterate through ALL blocked words
|
||||
text_lower = text.lower()
|
||||
|
|
@ -2071,7 +2155,9 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
def _streamed_scan_plan(self) -> _StreamedScanPlan:
|
||||
"""
|
||||
Per-stream inputs for buffer trimming: the retained tail length (the default
|
||||
context, widened to the longest configured keyword), the category exception
|
||||
context, widened to the longest configured keyword), the unreleased tail length
|
||||
(widened to the longest holdback a configured prebuilt pattern declares, since
|
||||
provider keys run longer than the default), the category exception
|
||||
phrases, which suppress matches anywhere in the scanned text, and the conditional
|
||||
category words, which only match when paired inside one sentence.
|
||||
"""
|
||||
|
|
@ -2081,6 +2167,13 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
)
|
||||
return _StreamedScanPlan(
|
||||
context_chars=max(CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS, longest_keyword),
|
||||
holdback_chars=max(
|
||||
(
|
||||
PATTERN_STREAMING_HOLDBACK_CHARS.get(entry["pattern_name"], CONTENT_FILTER_STREAMING_HOLDBACK_CHARS)
|
||||
for entry in self.compiled_patterns
|
||||
),
|
||||
default=CONTENT_FILTER_STREAMING_HOLDBACK_CHARS,
|
||||
),
|
||||
exception_phrases=tuple(
|
||||
phrase for category in self.loaded_categories.values() for phrase in category.exceptions
|
||||
),
|
||||
|
|
@ -2221,7 +2314,7 @@ class ContentFilterGuardrail(CustomGuardrail):
|
|||
|
||||
safe_to_yield_len = max(
|
||||
previous_state.yielded_masked_text_len,
|
||||
len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS),
|
||||
len(masked_text) - (0 if is_final else plan.holdback_chars),
|
||||
)
|
||||
released_text = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len]
|
||||
if isinstance(choice, TextChoices):
|
||||
|
|
|
|||
|
|
@ -98,6 +98,28 @@
|
|||
"category": "Credential Patterns",
|
||||
"description": "Detects generic API key patterns"
|
||||
},
|
||||
{
|
||||
"name": "openai_api_key",
|
||||
"display_name": "OpenAI API Key",
|
||||
"pattern": "sk-(?:(?:proj|svcacct|admin|None)-[A-Za-z0-9_\\-]{20,200}T3BlbkFJ(?:(?!sk-(?:proj|svcacct|admin|None)-)[A-Za-z0-9_\\-]){20,200}|[A-Za-z0-9]{20}T3BlbkFJ[A-Za-z0-9]{20}(?![A-Za-z0-9_]))",
|
||||
"streaming_holdback_chars": 419,
|
||||
"category": "Credential Patterns",
|
||||
"description": "Detects OpenAI API keys (sk-proj-..., sk-svcacct-..., sk-admin-..., sk-None-... and legacy sk-... keys, all carrying the T3BlbkFJ marker)"
|
||||
},
|
||||
{
|
||||
"name": "anthropic_api_key",
|
||||
"display_name": "Anthropic API Key",
|
||||
"pattern": "(?:(?<![A-Za-z0-9])|(?<=\\\\[nrt])|(?<=\\\\u[0-9A-F]{4})|(?<=%[0-9A-F]{2})|(?<=%25[0-9A-F]{2}))sk-ant-[A-Za-z0-9_\\-]{20,}",
|
||||
"category": "Credential Patterns",
|
||||
"description": "Detects Anthropic API keys (sk-ant-... format)"
|
||||
},
|
||||
{
|
||||
"name": "gemini_api_key",
|
||||
"display_name": "Gemini / Google API Key",
|
||||
"pattern": "(?:(?<![A-Za-z0-9])|(?<=\\\\[nrt])|(?<=\\\\u[0-9A-F]{4})|(?<=%[0-9A-F]{2})|(?<=%25[0-9A-F]{2}))(?:(?-i:AIza)[A-Za-z0-9_\\-]{33,36}|(?-i:AQ\\.)[A-Za-z0-9_\\-]{45,55})(?![A-Za-z0-9_])",
|
||||
"category": "Credential Patterns",
|
||||
"description": "Detects Google / Gemini API keys (legacy AIza... and new AQ.... authorization key formats)"
|
||||
},
|
||||
{
|
||||
"name": "ipv4",
|
||||
"display_name": "IP Address (IPv4)",
|
||||
|
|
@ -626,4 +648,4 @@
|
|||
"allow_word_numbers": false
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,10 +11,11 @@ import re
|
|||
from collections.abc import Iterator, Sequence
|
||||
from enum import Enum
|
||||
from re import Pattern
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter, with_config
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.proxy.guardrails.content_filter_data import DATA_ROOTS, category_dirs
|
||||
|
||||
|
|
@ -26,6 +27,7 @@ class _PatternEntry(TypedDict):
|
|||
pattern: ReadOnly[str]
|
||||
category: ReadOnly[str]
|
||||
description: ReadOnly[str]
|
||||
streaming_holdback_chars: ReadOnly[NotRequired[int]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
|
|
@ -84,6 +86,7 @@ KNOWN_PATTERN_KEYS: Final = {
|
|||
"category",
|
||||
"action",
|
||||
"description",
|
||||
"streaming_holdback_chars",
|
||||
}
|
||||
|
||||
PATTERN_EXTRA_CONFIG: Final[dict[str, _PatternExtraConfig]] = {}
|
||||
|
|
@ -132,6 +135,14 @@ for pattern_data in _PATTERNS_DATA["patterns"]:
|
|||
PATTERN_CATEGORIES[category] = []
|
||||
PATTERN_CATEGORIES[category].append(pattern_data["name"])
|
||||
|
||||
PATTERN_STREAMING_HOLDBACK_CHARS: Final = MappingProxyType(
|
||||
{
|
||||
pattern_data["name"]: pattern_data["streaming_holdback_chars"]
|
||||
for pattern_data in _PATTERNS_DATA["patterns"]
|
||||
if "streaming_holdback_chars" in pattern_data
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# Build display names mapping from JSON
|
||||
PATTERN_DISPLAY_NAMES: Final[dict[str, str]] = {
|
||||
|
|
|
|||
|
|
@ -247,6 +247,88 @@ class TestContentFilterGuardrail:
|
|||
assert "[EMAIL_REDACTED]" in result[0]
|
||||
assert "test@example.com" not in result[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"pattern_names",
|
||||
[("aws_secret_key", "openai_api_key"), ("openai_api_key", "aws_secret_key")],
|
||||
)
|
||||
async def test_apply_guardrail_masks_whole_key_when_another_pattern_matches_inside_it(
|
||||
self, pattern_names: tuple[str, str]
|
||||
):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-overlapping-patterns",
|
||||
patterns=[
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name=name, action=ContentFilterAction.MASK)
|
||||
for name in pattern_names
|
||||
],
|
||||
)
|
||||
openai_key: Final = "sk-proj-" + "Ab3" * 13 + "x-" + "Cd4" * 11 + "T3BlbkFJ" + "Ef5" * 24 + "gh"
|
||||
aws_secret_key: Final = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
|
||||
|
||||
guardrailed_inputs: Final = await guardrail.apply_guardrail(
|
||||
inputs={"texts": [f"key {openai_key} and secret {aws_secret_key}"]},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert guardrailed_inputs.get("texts") == [
|
||||
"key [OPENAI_API_KEY_REDACTED] and secret [AWS_SECRET_KEY_REDACTED]"
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_does_not_block_a_match_an_earlier_mask_pattern_already_covers(self):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-mask-then-block",
|
||||
patterns=[
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="generic_api_key", action=ContentFilterAction.MASK
|
||||
),
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt", pattern_name="aws_secret_key", action=ContentFilterAction.BLOCK
|
||||
),
|
||||
],
|
||||
)
|
||||
guardrailed_inputs: Final = await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["config api_key=" + "A1b2C3d4E5" * 4]},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
assert guardrailed_inputs.get("texts") == ["config [GENERIC_API_KEY_REDACTED]"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_masks_a_contextual_match_an_earlier_mask_moved_next_to_its_keyword(self):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-contextual-after-mask",
|
||||
patterns=[
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name="email", action=ContentFilterAction.MASK),
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name="au_tfn", action=ContentFilterAction.MASK),
|
||||
],
|
||||
)
|
||||
guardrailed_inputs: Final = await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["TFN jane.doe@example.com 123456782"]},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
assert guardrailed_inputs.get("texts") == ["TFN [EMAIL_REDACTED] [AU_TFN_REDACTED]"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_counts_only_the_unnamed_regex_that_matched(self):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-unnamed-regex-count",
|
||||
patterns=[
|
||||
ContentFilterPattern(pattern_type="regex", pattern=r"Key\d+", action=ContentFilterAction.MASK),
|
||||
ContentFilterPattern(pattern_type="regex", pattern=r"Token\d+", action=ContentFilterAction.MASK),
|
||||
],
|
||||
)
|
||||
request_data: Final = {"metadata": {}}
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["use Key1 here"]},
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
)
|
||||
entry: Final = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
||||
assert entry["masked_entity_count"] == {"custom_regex": 1}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_blocked_word_mask(self):
|
||||
"""
|
||||
|
|
@ -919,12 +1001,21 @@ class TestContentFilterGuardrail:
|
|||
# masked_entity_count for email is the real count, not N×.
|
||||
assert entry["masked_entity_count"].get("email") == 1
|
||||
|
||||
@staticmethod
|
||||
@classmethod
|
||||
async def _collect_streamed_text(
|
||||
cls,
|
||||
guardrail: ContentFilterGuardrail,
|
||||
chunks: list[str],
|
||||
metadata: dict[str, list[StandardLoggingGuardrailInformation]],
|
||||
) -> str:
|
||||
return "".join(await cls._collect_streamed_deltas(guardrail, chunks, metadata))
|
||||
|
||||
@staticmethod
|
||||
async def _collect_streamed_deltas(
|
||||
guardrail: ContentFilterGuardrail,
|
||||
chunks: list[str],
|
||||
metadata: dict[str, list[StandardLoggingGuardrailInformation]],
|
||||
) -> list[str]:
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
async def mock_stream():
|
||||
|
|
@ -951,7 +1042,7 @@ class TestContentFilterGuardrail:
|
|||
request_data={"messages": [], "model": "gpt-4o", "metadata": metadata},
|
||||
):
|
||||
yielded.append(chunk.choices[0].delta.content or "")
|
||||
return "".join(yielded)
|
||||
return yielded
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_scans_bounded_window_per_chunk(self):
|
||||
|
|
@ -1068,6 +1159,53 @@ class TestContentFilterGuardrail:
|
|||
assert entry["guardrail_status"] == "guardrail_intervened"
|
||||
assert [d["keyword"] for d in entry["guardrail_response"]] == [phrase]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("pattern_name", "key"),
|
||||
[
|
||||
("openai_api_key", "sk-proj-" + "a" * 74 + "T3BlbkFJ" + "b" * 74),
|
||||
("openai_api_key", "sk-svcacct-" + "a" * 200 + "T3BlbkFJ" + "b" * 200),
|
||||
("anthropic_api_key", "sk-ant-api03-" + "a" * 93 + "AA"),
|
||||
("gemini_api_key", "AQ." + "a" * 55),
|
||||
],
|
||||
)
|
||||
async def test_streaming_hook_masks_whole_provider_key_split_across_chunks(self, pattern_name: str, key: str):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-streaming-provider-key",
|
||||
patterns=[
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt",
|
||||
pattern_name=pattern_name,
|
||||
action=ContentFilterAction.MASK,
|
||||
)
|
||||
],
|
||||
event_hook=GuardrailEventHooks.post_call,
|
||||
)
|
||||
text: Final = f"Use {key} to call the API."
|
||||
chunks: Final = [text[i : i + 4] for i in range(0, len(text), 4)]
|
||||
|
||||
streamed: Final = await self._collect_streamed_text(guardrail, chunks, {})
|
||||
|
||||
assert streamed == f"Use [{pattern_name.upper()}_REDACTED] to call the API."
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_keeps_default_holdback_without_provider_key_patterns(self):
|
||||
guardrail: Final = ContentFilterGuardrail(
|
||||
guardrail_name="test-streaming-older-credential-patterns",
|
||||
patterns=[
|
||||
ContentFilterPattern(pattern_type="prebuilt", pattern_name=name, action=ContentFilterAction.MASK)
|
||||
for name in ("aws_access_key", "aws_secret_key", "github_token", "slack_token", "generic_api_key")
|
||||
],
|
||||
event_hook=GuardrailEventHooks.post_call,
|
||||
)
|
||||
text: Final = "plain words with nothing secret in them. " * 12
|
||||
chunks: Final = [text[i : i + 10] for i in range(0, len(text), 10)]
|
||||
|
||||
deltas: Final = await self._collect_streamed_deltas(guardrail, chunks, {})
|
||||
|
||||
assert "".join(deltas) == text
|
||||
assert len(next(delta for delta in reversed(deltas) if delta)) <= CONTENT_FILTER_STREAMING_HOLDBACK_CHARS
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_blocks_keyword_longer_than_scan_context(self):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ Tests for content filter pattern loading from JSON
|
|||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -183,3 +185,169 @@ def test_eu_patterns_have_category():
|
|||
assert (
|
||||
pattern_name in eu_category_patterns
|
||||
), f"Pattern {pattern_name} not in EU PII Patterns category"
|
||||
|
||||
|
||||
def test_provider_key_patterns_are_offered_in_the_ui_as_credential_patterns() -> None:
|
||||
ui_categories: Final = {pattern["name"]: pattern["category"] for pattern in get_pattern_metadata()}
|
||||
assert [ui_categories.get(name) for name in ("openai_api_key", "anthropic_api_key", "gemini_api_key")] == [
|
||||
"Credential Patterns"
|
||||
] * 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("pattern_name", "key"),
|
||||
[
|
||||
pytest.param("openai_api_key", "sk-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai-legacy"),
|
||||
pytest.param("openai_api_key", "sk-proj-" + "a" * 74 + "T3BlbkFJ" + "b" * 74, id="openai-project"),
|
||||
pytest.param("openai_api_key", "sk-proj-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai-project-shortest"),
|
||||
pytest.param("openai_api_key", "sk-svcacct-" + "a_b-" * 15 + "T3BlbkFJ" + "c-d_" * 15, id="openai-service-account"),
|
||||
pytest.param("openai_api_key", "sk-admin-" + "a" * 58 + "T3BlbkFJ" + "b" * 58, id="openai-admin"),
|
||||
pytest.param("openai_api_key", "sk-None-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai-user-key"),
|
||||
pytest.param("openai_api_key", "sk-proj-" + "a" * 20 + "T3BlbkFJ" + "b" * 19 + "-", id="openai-trailing-hyphen"),
|
||||
pytest.param(
|
||||
"anthropic_api_key", "sk-ant-api03-abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOP", id="anthropic-versioned"
|
||||
),
|
||||
pytest.param(
|
||||
"anthropic_api_key", "sk-ant-abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTU", id="anthropic-unversioned"
|
||||
),
|
||||
pytest.param("anthropic_api_key", "sk-ant-" + "a" * 20, id="anthropic-shortest"),
|
||||
pytest.param("anthropic_api_key", "sk-ant-api03-" + "a_b-" * 23 + "AA", id="anthropic-urlsafe"),
|
||||
pytest.param(
|
||||
"anthropic_api_key", "sk-ant-api03-abcdefghijklmnopqrstuvwxyz0123456789-", id="anthropic-trailing-hyphen"
|
||||
),
|
||||
pytest.param("gemini_api_key", "AIzaSyD-9tSrke72I6_AQp2h8HI4-uE6QNKQV0", id="gemini-legacy"),
|
||||
pytest.param("gemini_api_key", "AIza" + "a" * 33, id="gemini-legacy-shortest"),
|
||||
pytest.param("gemini_api_key", "AIza" + "a" * 36, id="gemini-legacy-longest"),
|
||||
pytest.param("gemini_api_key", "AIza" + "a" * 34 + "-", id="gemini-legacy-trailing-hyphen"),
|
||||
pytest.param("gemini_api_key", "AQ.Ab8RN6I-7xK9pL2mQ0vW5yZ1uX3tS4rE8qD7cB6aN5mP4oI3uY2", id="gemini-auth-key"),
|
||||
pytest.param("gemini_api_key", "AQ." + "a" * 45, id="gemini-auth-key-shortest"),
|
||||
pytest.param("gemini_api_key", "AQ." + "a" * 55, id="gemini-auth-key-longest"),
|
||||
pytest.param("gemini_api_key", "AQ." + "a" * 49 + "-", id="gemini-auth-key-trailing-hyphen"),
|
||||
],
|
||||
)
|
||||
def test_credential_pattern_matches_whole_key(pattern_name: str, key: str) -> None:
|
||||
match: Final = get_compiled_pattern(pattern_name).search(f"I keep getting 401 with {key}, is it expired?")
|
||||
assert match is not None
|
||||
assert match.group(0) == key
|
||||
|
||||
|
||||
@pytest.mark.parametrize("separator", ["\\n", "\\t", "\\u003c", "Bearer%20", "?key%3D", "?key%253D", "GOOGLE_KEY_"])
|
||||
@pytest.mark.parametrize(
|
||||
("pattern_name", "key"),
|
||||
[
|
||||
pytest.param("openai_api_key", "sk-proj-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai"),
|
||||
pytest.param("anthropic_api_key", "sk-ant-api03-" + "a" * 20, id="anthropic"),
|
||||
pytest.param("gemini_api_key", "AIza" + "a" * 35, id="gemini-legacy"),
|
||||
pytest.param("gemini_api_key", "AQ." + "a" * 50, id="gemini-auth-key"),
|
||||
],
|
||||
)
|
||||
def test_credential_pattern_matches_key_after_escaped_separator(pattern_name: str, key: str, separator: str) -> None:
|
||||
match: Final = get_compiled_pattern(pattern_name).search(f'{{"log": "retrying with{separator}{key} failed"}}')
|
||||
assert match is not None
|
||||
assert match.group(0) == key
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("pattern_name", "text"),
|
||||
[
|
||||
pytest.param("openai_api_key", "sk-short", id="openai-too-short"),
|
||||
pytest.param("openai_api_key", "sk-proj-" + "a" * 19, id="openai-project-too-short"),
|
||||
pytest.param("openai_api_key", "sk-" + "a" * 20 + "T3BlbkFJ" + "b" * 21, id="openai-legacy-too-long"),
|
||||
pytest.param("openai_api_key", "sk-proj-" + "a" * 48, id="openai-project-without-marker"),
|
||||
pytest.param("openai_api_key", "sk-proj-" + "a" * 19 + "T3BlbkFJ" + "b" * 20, id="openai-project-marker-too-early"),
|
||||
pytest.param("openai_api_key", "sk-proj-" + "a" * 20 + "T3BlbkFJ" + "b" * 19, id="openai-project-marker-too-late"),
|
||||
pytest.param("openai_api_key", "SK-ADMIN-dashboard-redesign-phase-two", id="openai-prefixed-slug"),
|
||||
pytest.param("openai_api_key", "desk-proj-management-dashboard-v2", id="openai-inside-longer-word"),
|
||||
pytest.param(
|
||||
"openai_api_key", "sk-abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRST", id="openai-legacy-without-marker"
|
||||
),
|
||||
pytest.param("openai_api_key", "sk-0123456789abcdef0123456789abcdef", id="openai-other-vendor-sk-key"),
|
||||
pytest.param("openai_api_key", "sk-Xq3vR9tLm2Pa8KdW5yZgQw", id="openai-litellm-virtual-key"),
|
||||
pytest.param(
|
||||
"openai_api_key",
|
||||
"sk-AbCdEfGhIjKlMnOpQrStUvWxYz0123456789AbCd-_x",
|
||||
id="openai-litellm-virtual-key-raised-length",
|
||||
),
|
||||
pytest.param("openai_api_key", "not-a-key", id="openai-no-prefix"),
|
||||
pytest.param(
|
||||
"openai_api_key", "sk-ant-api03-thisisanthropicnotopenaiabc123", id="openai-ignores-anthropic-key"
|
||||
),
|
||||
pytest.param("openai_api_key", "ordinary sentence without any api key in it", id="openai-plain-text"),
|
||||
pytest.param("anthropic_api_key", "sk-ant-" + "a" * 19, id="anthropic-too-short"),
|
||||
pytest.param("anthropic_api_key", "flask-ant-design-system-components", id="anthropic-inside-longer-word"),
|
||||
pytest.param(
|
||||
"anthropic_api_key",
|
||||
"sk-proj-abcdefghijklmnopqrstuvwxyzABCDEF1234567890",
|
||||
id="anthropic-ignores-openai-project-key",
|
||||
),
|
||||
pytest.param(
|
||||
"anthropic_api_key",
|
||||
"sk-admin-abcdefghijklmnopqrstuvwxyzABCDEF123456",
|
||||
id="anthropic-ignores-openai-admin-key",
|
||||
),
|
||||
pytest.param("anthropic_api_key", "sk-ant", id="anthropic-prefix-only"),
|
||||
pytest.param("anthropic_api_key", "plain text with nothing sensitive", id="anthropic-plain-text"),
|
||||
pytest.param("gemini_api_key", "AIza", id="gemini-prefix-only"),
|
||||
pytest.param("gemini_api_key", "AIza" + "a" * 32, id="gemini-legacy-too-short"),
|
||||
pytest.param("gemini_api_key", "AIza" + "a" * 37, id="gemini-legacy-too-long"),
|
||||
pytest.param("gemini_api_key", "BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB", id="gemini-wrong-prefix"),
|
||||
pytest.param("gemini_api_key", "AQ.", id="gemini-auth-prefix-only"),
|
||||
pytest.param("gemini_api_key", "AQ." + "a" * 44, id="gemini-auth-key-too-short"),
|
||||
pytest.param("gemini_api_key", "AQ." + "a" * 56, id="gemini-auth-key-too-long"),
|
||||
pytest.param("gemini_api_key", "AQ-" + "a" * 50, id="gemini-auth-key-wrong-separator"),
|
||||
pytest.param("gemini_api_key", "x" + "AIza" + "a" * 35, id="gemini-inside-longer-token"),
|
||||
pytest.param("gemini_api_key", "The AQ. was reported at 42 ppm today.", id="gemini-aq-acronym"),
|
||||
pytest.param("gemini_api_key", "Is the question AQ. Maybe it is BQ.", id="gemini-aq-sentence-end"),
|
||||
pytest.param("gemini_api_key", "plain text with no google key", id="gemini-plain-text"),
|
||||
pytest.param(
|
||||
"gemini_api_key",
|
||||
"https://example.com/blog/aizawa-takeshi-interview-2024-tokyo-01",
|
||||
id="gemini-lowercase-legacy-prefix",
|
||||
),
|
||||
pytest.param(
|
||||
"gemini_api_key",
|
||||
"ssh aq.internal-observability-metrics-collector-eu-west-01",
|
||||
id="gemini-lowercase-auth-prefix",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_credential_pattern_ignores_non_keys(pattern_name: str, text: str) -> None:
|
||||
assert get_compiled_pattern(pattern_name).search(text) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("suffix", ["请检查", "é", "ü der Schlüssel"])
|
||||
@pytest.mark.parametrize(
|
||||
("pattern_name", "key"),
|
||||
[
|
||||
pytest.param("openai_api_key", "sk-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai-legacy"),
|
||||
pytest.param("gemini_api_key", "AIza" + "a" * 35, id="gemini-legacy"),
|
||||
pytest.param("gemini_api_key", "AQ." + "a" * 50, id="gemini-auth-key"),
|
||||
],
|
||||
)
|
||||
def test_credential_pattern_matches_key_followed_by_non_ascii_letters(pattern_name: str, key: str, suffix: str) -> None:
|
||||
match: Final = get_compiled_pattern(pattern_name).search(f"我的密钥是{key}{suffix}")
|
||||
assert match is not None
|
||||
assert match.group(0) == key
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"leading",
|
||||
[
|
||||
pytest.param("sk-proj-" + "a" * 74 + "T3BlbkFJ" + "b" * 74, id="another-key"),
|
||||
pytest.param("sk-proj-" + "x" * 20 + "T3BlbkFJ" + "y" * 192, id="key-shaped-filler"),
|
||||
pytest.param("sk-proj-" + "x" * 100 + "T3BlbkFJ" + "y" * 100, id="long-key-shaped-filler"),
|
||||
],
|
||||
)
|
||||
def test_openai_key_pattern_masks_a_key_glued_to_preceding_key_shaped_text(leading: str) -> None:
|
||||
key: Final = "sk-proj-" + "c" * 74 + "T3BlbkFJ" + "d" * 74
|
||||
assert get_compiled_pattern("openai_api_key").sub("[KEY]", leading + key) == "[KEY][KEY]"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", ["sk-proj-", "sk-svcacct-", "sk-admin-", "sk-None-"])
|
||||
def test_openai_key_pattern_scans_repeated_prefixes_without_stalling(prefix: str) -> None:
|
||||
prompt: Final = prefix * (128_000 // len(prefix))
|
||||
started: Final = time.perf_counter()
|
||||
match: Final = get_compiled_pattern("openai_api_key").search(prompt)
|
||||
elapsed: Final = time.perf_counter() - started
|
||||
assert match is None
|
||||
assert elapsed < 1.0
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue