mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
feat(proxy): redact or drop individual batch records instead of rejecting the file
A single record tripping a guardrail rejected the whole upload, which is unusable for a file holding thousands of rows. A record a guardrail rewrites is now submitted in its rewritten form, a record it blocks is left out, and the create response reports every changed record by both custom_id and line so a caller can reconcile against the file it sent. The same outcome is written to the proxy log and to request metadata, so it is not visible only to the caller. A rewritten record goes straight to a spool and only its offset is carried, so a masking guardrail touching most rows of a large upload does not build a second copy of the file on the heap, and the rewrite runs off the event loop the way the sibling full-file validation does. Both proxy-injected metadata keys are captured from the record and restored exactly, including an explicit null, so a masked row keeps the tags that decide how it is attributed. A record is dropped only when a guardrail judged its content. `GuardrailRaisedException` now carries `blocked_content` for that, because half its raise sites in the repo signal an unreachable or unparseable backend under a fail-closed policy, and treating those as blocks would turn "refuse this request" into "drop this record and submit the rest". The default is off, so a raise that does not say what it means aborts the upload instead of silently shrinking the file.
This commit is contained in:
Yucheng Zhu
parent
3a31331435
commit
e36a964037
14 changed files with 864 additions and 151 deletions
|
|
@ -1039,16 +1039,29 @@ class LiteLLMUnknownProvider(BadRequestError):
|
|||
|
||||
|
||||
class GuardrailRaisedException(Exception):
|
||||
"""
|
||||
Raised both when a guardrail judged content and when it could not judge it at all, since a
|
||||
guardrail that fails closed refuses the request the same way a policy violation does.
|
||||
|
||||
``blocked_content`` separates the two. Set it only where the guardrail actually reached a
|
||||
verdict on the payload; leave it alone for an unreachable backend, a timeout, or a response
|
||||
the integration could not parse. Callers that treat a block as something other than a plain
|
||||
failure, such as the batch path dropping one record and submitting the rest, must gate on it,
|
||||
because dropping a record no guardrail ever inspected is a silent loss of enforcement.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
guardrail_name: str | None = None,
|
||||
message: str = "",
|
||||
should_wrap_with_default_message: bool = True,
|
||||
status_code: int = 400,
|
||||
blocked_content: bool = False,
|
||||
):
|
||||
default_message: Final = f"Guardrail raised an exception, Guardrail: {guardrail_name}, Message: {message}"
|
||||
self.guardrail_name = guardrail_name
|
||||
self.status_code = status_code
|
||||
self.blocked_content = blocked_content
|
||||
self.message = default_message if should_wrap_with_default_message else message
|
||||
super().__init__(self.message)
|
||||
|
||||
|
|
|
|||
|
|
@ -65,6 +65,41 @@ _guardrail_self_recorded: Final[contextvars.ContextVar[bool]] = contextvars.Cont
|
|||
)
|
||||
|
||||
|
||||
def is_guardrail_intervention(e: Exception) -> bool:
|
||||
"""
|
||||
Returns True if the exception represents an intentional guardrail block
|
||||
(this was logged previously as an API failure - guardrail_failed_to_respond).
|
||||
|
||||
Guardrails signal intentional blocks by raising:
|
||||
- GuardrailRaisedException (generic guardrail API, tool permission)
|
||||
- BlockedPiiEntityError (Presidio PII detection)
|
||||
- SensitiveDataRouteException (sensitive-data reroute to on-premise model)
|
||||
- HTTPException with a block-signalling status (400, 403, 422)
|
||||
- ModifyResponseException (passthrough mode violation)
|
||||
|
||||
Only the statuses guardrails use in-tree to signal a deliberate rejection
|
||||
count as an intervention: 400 (content policy), 403 (e.g. akto) and 422
|
||||
(e.g. llm_as_a_judge). Other 4xx codes are commonly propagated from an
|
||||
upstream guardrail provider response (401 bad key, 408 timeout, 429 rate
|
||||
limit, or a raw upstream status), which are technical failures, not
|
||||
blocks, so they stay guardrail_failed_to_respond.
|
||||
"""
|
||||
if isinstance(e, ModifyResponseException):
|
||||
return True
|
||||
if isinstance(
|
||||
e,
|
||||
(
|
||||
GuardrailRaisedException,
|
||||
BlockedPiiEntityError,
|
||||
SensitiveDataRouteException,
|
||||
),
|
||||
):
|
||||
return True
|
||||
if HTTPException is not None and isinstance(e, HTTPException) and e.status_code in _GUARDRAIL_BLOCK_STATUS_CODES:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _strict_guardrail_modes_enabled() -> bool:
|
||||
"""Whether guardrail-mode validation raises (default) or logs a warning.
|
||||
|
||||
|
|
@ -429,11 +464,13 @@ class CustomGuardrail(CustomLogger):
|
|||
f"Sensitive data detected by {self.guardrail_name} (routing skipped: request has no session_id)"
|
||||
),
|
||||
guardrail_name=self.guardrail_name,
|
||||
blocked_content=True,
|
||||
)
|
||||
else:
|
||||
raise GuardrailRaisedException(
|
||||
message=f"Sensitive data detected by {self.guardrail_name}",
|
||||
guardrail_name=self.guardrail_name,
|
||||
blocked_content=True,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -1068,42 +1105,8 @@ class CustomGuardrail(CustomLogger):
|
|||
|
||||
@staticmethod
|
||||
def _is_guardrail_intervention(e: Exception) -> bool:
|
||||
"""
|
||||
Returns True if the exception represents an intentional guardrail block
|
||||
(this was logged previously as an API failure - guardrail_failed_to_respond).
|
||||
|
||||
Guardrails signal intentional blocks by raising:
|
||||
- GuardrailRaisedException (generic guardrail API, tool permission)
|
||||
- BlockedPiiEntityError (Presidio PII detection)
|
||||
- SensitiveDataRouteException (sensitive-data reroute to on-premise model)
|
||||
- HTTPException with a block-signalling status (400, 403, 422)
|
||||
- ModifyResponseException (passthrough mode violation)
|
||||
|
||||
Only the statuses guardrails use in-tree to signal a deliberate rejection
|
||||
count as an intervention: 400 (content policy), 403 (e.g. akto) and 422
|
||||
(e.g. llm_as_a_judge). Other 4xx codes are commonly propagated from an
|
||||
upstream guardrail provider response (401 bad key, 408 timeout, 429 rate
|
||||
limit, or a raw upstream status), which are technical failures, not
|
||||
blocks, so they stay guardrail_failed_to_respond.
|
||||
"""
|
||||
if isinstance(e, ModifyResponseException):
|
||||
return True
|
||||
if isinstance(
|
||||
e,
|
||||
(
|
||||
GuardrailRaisedException,
|
||||
BlockedPiiEntityError,
|
||||
SensitiveDataRouteException,
|
||||
),
|
||||
):
|
||||
return True
|
||||
if (
|
||||
HTTPException is not None
|
||||
and isinstance(e, HTTPException)
|
||||
and e.status_code in _GUARDRAIL_BLOCK_STATUS_CODES
|
||||
):
|
||||
return True
|
||||
return False
|
||||
"""Retained spelling for existing callers; prefer ``is_guardrail_intervention``."""
|
||||
return is_guardrail_intervention(e)
|
||||
|
||||
def _process_error(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -356,6 +356,7 @@ class DeepKeepGuardrail(CustomGuardrail):
|
|||
guardrail_name=GUARDRAIL_NAME,
|
||||
message=error_message,
|
||||
should_wrap_with_default_message=False,
|
||||
blocked_content=True,
|
||||
)
|
||||
|
||||
return self._build_return_inputs(
|
||||
|
|
|
|||
|
|
@ -464,6 +464,7 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
guardrail_name=GUARDRAIL_NAME,
|
||||
message=error_message,
|
||||
should_wrap_with_default_message=False,
|
||||
blocked_content=True,
|
||||
)
|
||||
|
||||
return self._build_guardrail_return_inputs(
|
||||
|
|
|
|||
|
|
@ -169,6 +169,7 @@ class PromptGuardGuardrail(CustomGuardrail):
|
|||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=(f"Blocked by PromptGuard: {threat_type} (confidence={confidence}, event_id={event_id})"),
|
||||
blocked_content=True,
|
||||
)
|
||||
|
||||
if decision == "redact":
|
||||
|
|
|
|||
|
|
@ -211,6 +211,7 @@ class SingulrGuardrail(CustomGuardrail):
|
|||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=f"Blocked by Singulr: {result.blocking_due_to or 'unknown'}",
|
||||
blocked_content=True,
|
||||
)
|
||||
|
||||
return inputs
|
||||
|
|
|
|||
|
|
@ -542,6 +542,7 @@ class StraikerGuardrail(CustomGuardrail):
|
|||
guardrail_name=self.guardrail_name or GUARDRAIL_NAME,
|
||||
message=message,
|
||||
should_wrap_with_default_message=False,
|
||||
blocked_content=True,
|
||||
)
|
||||
raise ModifyResponseException(
|
||||
message=message,
|
||||
|
|
|
|||
|
|
@ -527,7 +527,9 @@ class ToolPermissionGuardrail(CustomGuardrail):
|
|||
if not is_allowed and message is not None:
|
||||
verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message)
|
||||
if self.on_disallowed_action == "block":
|
||||
raise GuardrailRaisedException(guardrail_name=self.guardrail_name, message=message)
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name, message=message, blocked_content=True
|
||||
)
|
||||
|
||||
return tuple(
|
||||
(
|
||||
|
|
|
|||
|
|
@ -205,6 +205,7 @@ class VigilGuardGuardrail(CustomGuardrail):
|
|||
guardrail_name=self.guardrail_name,
|
||||
message=self._build_block_reason(analysis),
|
||||
should_wrap_with_default_message=False,
|
||||
blocked_content=True,
|
||||
)
|
||||
|
||||
if decision == "SANITIZED":
|
||||
|
|
@ -245,6 +246,7 @@ class VigilGuardGuardrail(CustomGuardrail):
|
|||
guardrail_name=self.guardrail_name,
|
||||
message=self._build_block_reason(analysis),
|
||||
should_wrap_with_default_message=False,
|
||||
blocked_content=True,
|
||||
)
|
||||
|
||||
if decision == "SANITIZED":
|
||||
|
|
|
|||
|
|
@ -10,6 +10,8 @@ from __future__ import annotations
|
|||
import asyncio
|
||||
import copy
|
||||
import json
|
||||
import re
|
||||
import tempfile
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
|
|
@ -19,8 +21,11 @@ from urllib.parse import urlsplit
|
|||
from fastapi import HTTPException
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.integrations.custom_guardrail import is_guardrail_intervention
|
||||
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.llms.openai import BatchGuardrailRecord, BatchGuardrailReport
|
||||
from litellm.types.utils import CallTypes, CallTypesLiteral
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -30,6 +35,16 @@ EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({})
|
|||
|
||||
_SCAN_WINDOW: Final = 32
|
||||
|
||||
# Past this the rewrite rolls to disk, keeping the router's per-deployment deepcopy of the handle
|
||||
# as cheap as it is for the spooled upload this replaces.
|
||||
_REWRITE_SPOOL_BYTES: Final = 1024 * 1024
|
||||
|
||||
# custom_id is caller-supplied and reaches a log line, so it is stripped of control characters
|
||||
# and capped rather than rendered as given.
|
||||
_CONTROL_CHARACTERS: Final = re.compile(r"[\x00-\x1f\x7f]")
|
||||
_CUSTOM_ID_LOG_LIMIT: Final = 128
|
||||
_SUMMARY_LIMIT: Final = 50
|
||||
|
||||
_SCAN_METADATA_KEY: Final = "litellm_metadata"
|
||||
|
||||
# `metadata` is dropped rather than diffed: guardrail dispatch writes its bookkeeping into it
|
||||
|
|
@ -82,13 +97,74 @@ class UnscannableRecord:
|
|||
url: str | None
|
||||
|
||||
|
||||
BatchScanFailure: TypeAlias = UnparseableRecord | UnscannableRecord
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RedactionRequired:
|
||||
class _Redaction:
|
||||
"""A rewritten record on its way to the scan spool, held only for the window it was scanned in."""
|
||||
|
||||
line_number: int
|
||||
custom_id: str | None
|
||||
text: str
|
||||
|
||||
|
||||
BatchScanFailure: TypeAlias = UnparseableRecord | UnscannableRecord | RedactionRequired
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RecordRedacted:
|
||||
line_number: int
|
||||
custom_id: str | None
|
||||
offset: int
|
||||
length: int
|
||||
"""Where the re-serialized record sits in the scan spool, so a large file's rewrites stay off the heap."""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RecordDropped:
|
||||
line_number: int
|
||||
custom_id: str | None
|
||||
guardrail: str | None = None
|
||||
|
||||
|
||||
_RecordChange: TypeAlias = RecordRedacted | RecordDropped
|
||||
_ScanOutcome: TypeAlias = BatchScanFailure | _Redaction | RecordDropped
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BatchScanResult:
|
||||
"""What the scan decided, per record. Empty changes means the upload proceeds untouched."""
|
||||
|
||||
changes: tuple[_RecordChange, ...]
|
||||
scanned_records: int
|
||||
redactions: BinaryIO
|
||||
"""Spool holding every rewritten record, keyed by the offsets on each ``RecordRedacted``."""
|
||||
|
||||
@property
|
||||
def submitted_records(self) -> int:
|
||||
return self.scanned_records - sum(1 for change in self.changes if isinstance(change, RecordDropped))
|
||||
|
||||
def summary(self) -> str:
|
||||
"""Compact per-record outcome for the server-side log line, capped so one upload cannot flood it."""
|
||||
shown: Final = ", ".join(
|
||||
f"line {change.line_number}{_describe(change.custom_id)} "
|
||||
f"{'redacted' if isinstance(change, RecordRedacted) else 'dropped'}"
|
||||
for change in self.changes[:_SUMMARY_LIMIT]
|
||||
)
|
||||
remaining: Final = len(self.changes) - _SUMMARY_LIMIT
|
||||
return shown if remaining <= 0 else f"{shown}, and {remaining} more"
|
||||
|
||||
def report(self) -> BatchGuardrailReport:
|
||||
return BatchGuardrailReport(
|
||||
submitted_records=self.submitted_records,
|
||||
modified_records=tuple(
|
||||
BatchGuardrailRecord(
|
||||
line=change.line_number,
|
||||
custom_id=change.custom_id,
|
||||
action="redacted" if isinstance(change, RecordRedacted) else "dropped",
|
||||
guardrail=change.guardrail if isinstance(change, RecordDropped) else None,
|
||||
)
|
||||
for change in self.changes
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -114,25 +190,62 @@ def raise_public(failure: BatchScanFailure) -> NoReturn:
|
|||
"and its body has no messages, prompt or input, so guardrails cannot read it. "
|
||||
"Give the record a chat, completion, embedding, responses or messages body"
|
||||
)
|
||||
case RedactionRequired(line_number=line_number, custom_id=custom_id):
|
||||
raise _rejected(
|
||||
f"A guardrail changed batch input line {line_number}{_describe(custom_id)}. "
|
||||
"Per-record redaction is not enabled, so the file was rejected rather than modified"
|
||||
)
|
||||
case _:
|
||||
assert_never(failure)
|
||||
|
||||
|
||||
def raise_nothing_to_submit() -> NoReturn:
|
||||
"""Every record was blocked, so there is no batch left to create."""
|
||||
raise _rejected(
|
||||
"Every record in the batch input file was blocked by a guardrail, so there is nothing left to submit"
|
||||
)
|
||||
|
||||
|
||||
def _is_content_block(exc: BaseException) -> bool:
|
||||
"""
|
||||
Whether the guardrail judged the record, as opposed to failing to judge it.
|
||||
|
||||
Stricter than ``is_guardrail_intervention``, which answers a different question and counts
|
||||
every ``GuardrailRaisedException`` as a block. Several integrations raise that same exception
|
||||
for an unreachable backend or an unparseable response, and only when the operator configured
|
||||
the guardrail to fail closed, so treating it as a block would turn "refuse this request" into
|
||||
"drop this record and submit the rest", which is the silent loss of enforcement this whole
|
||||
path exists to prevent. A guardrail that does not say it blocked content aborts the upload.
|
||||
"""
|
||||
if isinstance(exc, GuardrailRaisedException):
|
||||
return exc.blocked_content
|
||||
return is_guardrail_intervention(exc)
|
||||
|
||||
|
||||
def _naming_guardrail(exc: BaseException) -> str | None:
|
||||
"""The guardrail that raised, from whichever place it recorded its own name."""
|
||||
named: Final = getattr(exc, "guardrail_name", None)
|
||||
if isinstance(named, str):
|
||||
return named
|
||||
detail: Final = getattr(exc, "detail", None)
|
||||
enriched: Final = detail.get("guardrail_name") if isinstance(detail, dict) else None
|
||||
return enriched if isinstance(enriched, str) else None
|
||||
|
||||
|
||||
def _describe(custom_id: str | None) -> str:
|
||||
return f" (custom_id {custom_id})" if custom_id else ""
|
||||
if not custom_id:
|
||||
return ""
|
||||
safe: Final = _CONTROL_CHARACTERS.sub(" ", custom_id)[:_CUSTOM_ID_LOG_LIMIT]
|
||||
return f" (custom_id {safe})"
|
||||
|
||||
|
||||
def _iter_lines(source: BinaryIO) -> Iterator[tuple[int, str]]:
|
||||
"""Yield every non-blank line with its 1-based number, so both passes number records alike."""
|
||||
for line_number, raw_line in enumerate(source, start=1):
|
||||
text = raw_line.decode("utf-8")
|
||||
if text.strip():
|
||||
yield line_number, text
|
||||
|
||||
|
||||
def _iter_records(source: BinaryIO) -> Iterator[_ParsedRecord]:
|
||||
"""Yield one record per line, relying on the upload validation that already ran."""
|
||||
for line_number, raw_line in enumerate(source, start=1):
|
||||
text = raw_line.decode("utf-8")
|
||||
if text.strip():
|
||||
yield _ParsedRecord(line_number=line_number, payload=json.loads(text))
|
||||
for line_number, text in _iter_lines(source):
|
||||
yield _ParsedRecord(line_number=line_number, payload=json.loads(text))
|
||||
|
||||
|
||||
def _call_type_from_url(url: str) -> CallTypesLiteral | None:
|
||||
|
|
@ -204,7 +317,7 @@ async def _scan_record(
|
|||
scan_metadata: Mapping[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> BatchScanFailure | None:
|
||||
) -> _ScanOutcome | None:
|
||||
body: Final = record.payload.get("body")
|
||||
if not isinstance(body, dict):
|
||||
return UnparseableRecord(line_number=record.line_number)
|
||||
|
|
@ -220,6 +333,7 @@ async def _scan_record(
|
|||
)
|
||||
|
||||
scan_input: Final[dict[str, object]] = copy.deepcopy(body) # mutable-ok: pre_call_hook mutates the dict it is given
|
||||
own_injected: Final = MappingProxyType({key: body[key] for key in _INJECTED_KEYS if key in body})
|
||||
scan_input.pop("metadata", None)
|
||||
# Deep, and per record: `headers` and `tags` are nested containers shared with the upload
|
||||
# request and with every other record in the window, and a guardrail that writes into one in
|
||||
|
|
@ -227,19 +341,31 @@ async def _scan_record(
|
|||
# already removed the values that cannot be copied.
|
||||
scan_input[_SCAN_METADATA_KEY] = copy.deepcopy(dict(scan_metadata)) # mutable-ok: guardrails write here
|
||||
|
||||
# The chain hands back the body it produced, which may be a replacement for the dict it was
|
||||
# given rather than that same dict mutated, so this is what gets compared.
|
||||
scanned: Final[dict] = await proxy_logging_obj.pre_call_hook( # mutable-ok: the guardrails' own dict
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=scan_input,
|
||||
call_type=call_type,
|
||||
guardrails_only=True,
|
||||
)
|
||||
try:
|
||||
# The chain hands back the body it produced, which may be a replacement for the dict it was
|
||||
# given rather than that same dict mutated, so this is what gets compared.
|
||||
scanned: Final[dict] = await proxy_logging_obj.pre_call_hook( # mutable-ok: the guardrails' own dict
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=scan_input,
|
||||
call_type=call_type,
|
||||
guardrails_only=True,
|
||||
)
|
||||
except Exception as exc:
|
||||
if _is_content_block(exc):
|
||||
return RecordDropped(line_number=record.line_number, custom_id=custom_id, guardrail=_naming_guardrail(exc))
|
||||
raise
|
||||
|
||||
compared: Final = (frozenset(body) | frozenset(scanned)) - _INJECTED_KEYS
|
||||
if _fingerprint(scanned, compared) != _fingerprint(body, compared):
|
||||
return RedactionRequired(line_number=record.line_number, custom_id=custom_id)
|
||||
return None
|
||||
if _fingerprint(scanned, compared) == _fingerprint(body, compared):
|
||||
return None
|
||||
for injected in _INJECTED_KEYS:
|
||||
scanned.pop(injected, None)
|
||||
scanned.update(own_injected)
|
||||
return _Redaction(
|
||||
line_number=record.line_number,
|
||||
custom_id=custom_id,
|
||||
text=json.dumps({**record.payload, "body": scanned}), # mutable-ok: json.dumps needs a plain dict
|
||||
)
|
||||
|
||||
|
||||
async def _scan_window(
|
||||
|
|
@ -247,7 +373,7 @@ async def _scan_window(
|
|||
scan_metadata: Mapping[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> tuple[tuple[int, BatchScanFailure | BaseException], ...]:
|
||||
) -> tuple[tuple[int, _ScanOutcome | BaseException], ...]:
|
||||
"""``return_exceptions=True`` so one record raising never leaves its siblings unobserved."""
|
||||
outcomes: Final = await asyncio.gather(
|
||||
*(_scan_record(record, scan_metadata, user_api_key_dict, proxy_logging_obj) for record in window),
|
||||
|
|
@ -256,6 +382,20 @@ async def _scan_window(
|
|||
return tuple((record.line_number, outcome) for record, outcome in zip(window, outcomes) if outcome is not None)
|
||||
|
||||
|
||||
def _spool(redactions: BinaryIO, redaction: _Redaction) -> RecordRedacted:
|
||||
"""Park the rewritten record on disk so only its location is carried for the rest of the scan."""
|
||||
encoded: Final = redaction.text.encode("utf-8")
|
||||
redactions.seek(0, 2)
|
||||
offset: Final = redactions.tell()
|
||||
redactions.write(encoded)
|
||||
return RecordRedacted(
|
||||
line_number=redaction.line_number,
|
||||
custom_id=redaction.custom_id,
|
||||
offset=offset,
|
||||
length=len(encoded),
|
||||
)
|
||||
|
||||
|
||||
def _worst(problems: tuple[tuple[int, BatchScanFailure | BaseException], ...]) -> BatchScanFailure | BaseException:
|
||||
"""A guardrail that blocked outranks a record we merely refused; then earliest line wins."""
|
||||
raised: Final = tuple(problem for problem in problems if isinstance(problem[1], BaseException))
|
||||
|
|
@ -268,20 +408,36 @@ async def scan_batch_input_file(
|
|||
request_metadata: Mapping[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> BatchScanFailure | None:
|
||||
) -> BatchScanFailure | BatchScanResult:
|
||||
"""
|
||||
Stream a batch input file and run the pre-call guardrail chain against every record.
|
||||
|
||||
Returns the record to reject, or None when every record passed. A guardrail that blocks raises
|
||||
its own exception, which is re-raised untouched so its status code survives.
|
||||
A record a guardrail rewrites is kept in its rewritten form and a record it blocks is dropped,
|
||||
which is what the online path does per request. Both are returned for reporting. A guardrail
|
||||
exception that is not a block is re-raised untouched so its status code survives, since dropping
|
||||
a record that was never inspected is worse than refusing the file.
|
||||
"""
|
||||
scan_metadata: Final = build_scan_metadata(request_metadata)
|
||||
problems: Final[list[tuple[int, BatchScanFailure | BaseException]]] = [] # mutable-ok: spans windows
|
||||
changes: Final[list[_RecordChange]] = [] # mutable-ok: accumulates across windows
|
||||
window: Final[list[_ParsedRecord]] = [] # mutable-ok: bounded read-ahead buffer
|
||||
scanned: Final[list[int]] = [] # mutable-ok: counts records the scan actually reached
|
||||
redactions: Final = tempfile.SpooledTemporaryFile( # noqa: SIM115 # the rewrite reads this back
|
||||
max_size=_REWRITE_SPOOL_BYTES
|
||||
)
|
||||
|
||||
async def drain() -> None:
|
||||
if window:
|
||||
problems.extend(await _scan_window(tuple(window), scan_metadata, user_api_key_dict, proxy_logging_obj))
|
||||
scanned.append(len(window))
|
||||
for line_number, outcome in await _scan_window(
|
||||
tuple(window), scan_metadata, user_api_key_dict, proxy_logging_obj
|
||||
):
|
||||
if isinstance(outcome, _Redaction):
|
||||
changes.append(_spool(redactions, outcome))
|
||||
elif isinstance(outcome, RecordDropped):
|
||||
changes.append(outcome)
|
||||
else:
|
||||
problems.append((line_number, outcome))
|
||||
window.clear()
|
||||
|
||||
try:
|
||||
|
|
@ -296,9 +452,51 @@ async def scan_batch_input_file(
|
|||
finally:
|
||||
file_source.seek(0)
|
||||
|
||||
if not problems:
|
||||
return None
|
||||
worst: Final = _worst(tuple(problems))
|
||||
if isinstance(worst, BaseException):
|
||||
raise worst
|
||||
return worst
|
||||
if problems:
|
||||
worst: Final = _worst(tuple(problems))
|
||||
if isinstance(worst, BaseException):
|
||||
raise worst
|
||||
return worst
|
||||
return BatchScanResult(
|
||||
changes=tuple(sorted(changes, key=lambda change: change.line_number)),
|
||||
scanned_records=sum(scanned),
|
||||
redactions=redactions,
|
||||
)
|
||||
|
||||
|
||||
def _read_spooled(redactions: BinaryIO, change: RecordRedacted) -> str:
|
||||
redactions.seek(change.offset)
|
||||
return redactions.read(change.length).decode("utf-8")
|
||||
|
||||
|
||||
def rewrite_batch_input_file(file_source: BinaryIO, result: BatchScanResult) -> BinaryIO:
|
||||
"""
|
||||
Re-emit the file with redacted records rewritten and dropped records left out.
|
||||
|
||||
Untouched records are copied through as written rather than re-serialized, so enabling the
|
||||
feature does not reformat records no guardrail objected to. Blank lines between records are
|
||||
not carried over, since they are not records. Rewritten records are read back from the scan's
|
||||
spool rather than from memory, so a file whose records are mostly rewritten does not put a
|
||||
second copy of itself on the heap.
|
||||
"""
|
||||
redacted: Final = MappingProxyType(
|
||||
{change.line_number: change for change in result.changes if isinstance(change, RecordRedacted)}
|
||||
) # mutable-ok: MappingProxyType freezes the lookup table
|
||||
dropped: Final = frozenset(change.line_number for change in result.changes if isinstance(change, RecordDropped))
|
||||
|
||||
output: Final = tempfile.SpooledTemporaryFile( # noqa: SIM115 # the caller uploads this handle
|
||||
max_size=_REWRITE_SPOOL_BYTES
|
||||
)
|
||||
wrote_any = False # rebind-ok: tracks whether a separator is needed
|
||||
try:
|
||||
for line_number, text in _iter_lines(file_source):
|
||||
if line_number in dropped:
|
||||
continue
|
||||
change = redacted.get(line_number)
|
||||
line = text.rstrip("\n") if change is None else _read_spooled(result.redactions, change)
|
||||
output.write((("\n" if wrote_any else "") + line).encode("utf-8"))
|
||||
wrote_any = True
|
||||
finally:
|
||||
file_source.seek(0)
|
||||
output.seek(0)
|
||||
return output
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@
|
|||
|
||||
import asyncio
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, BinaryIO, Final, cast, get_args
|
||||
|
||||
import httpx
|
||||
|
|
@ -29,6 +30,7 @@ from litellm._logging import verbose_proxy_logger
|
|||
from litellm.litellm_core_utils.cloud_storage_security import (
|
||||
is_managed_cloud_storage_uri,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
|
||||
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
|
@ -48,7 +50,10 @@ from litellm.proxy.openai_files_endpoints.batch_file_validation import (
|
|||
)
|
||||
from litellm.proxy.openai_files_endpoints.batch_guardrails import (
|
||||
EMPTY_MAPPING,
|
||||
BatchScanResult,
|
||||
raise_nothing_to_submit,
|
||||
raise_public,
|
||||
rewrite_batch_input_file,
|
||||
scan_batch_input_file,
|
||||
)
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
|
|
@ -111,6 +116,34 @@ def get_files_provider_config(
|
|||
return None
|
||||
|
||||
|
||||
async def _scan_batch_upload(
|
||||
*,
|
||||
file_source: bytes | BinaryIO,
|
||||
purpose: str,
|
||||
request_metadata: Mapping[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> BatchScanResult | None:
|
||||
"""Guardrail the records of a batch input file, or None when this upload has nothing to scan."""
|
||||
if (
|
||||
purpose != "batch"
|
||||
or isinstance(file_source, bytes)
|
||||
or not proxy_logging_obj.has_pre_call_guardrails(request_metadata)
|
||||
):
|
||||
return None
|
||||
outcome: Final = await scan_batch_input_file(
|
||||
file_source=file_source,
|
||||
request_metadata=request_metadata,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if not isinstance(outcome, BatchScanResult):
|
||||
raise_public(outcome)
|
||||
if outcome.changes and outcome.submitted_records == 0:
|
||||
raise_nothing_to_submit()
|
||||
return outcome
|
||||
|
||||
|
||||
def get_first_json_object(file_source: bytes | BinaryIO) -> dict | None:
|
||||
try:
|
||||
if isinstance(file_source, (bytes, bytearray)):
|
||||
|
|
@ -478,28 +511,38 @@ async def create_file(
|
|||
|
||||
# /v1/files stores its proxy metadata under litellm_metadata, not metadata
|
||||
request_metadata: Final = data.get("metadata") or data.get("litellm_metadata") or EMPTY_MAPPING
|
||||
if (
|
||||
purpose == "batch"
|
||||
and not isinstance(file_source, bytes)
|
||||
and proxy_logging_obj.has_pre_call_guardrails(request_metadata)
|
||||
):
|
||||
scan_failure: Final = await scan_batch_input_file(
|
||||
file_source=file_source,
|
||||
request_metadata=request_metadata,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
scan_result: Final = await _scan_batch_upload(
|
||||
file_source=file_source,
|
||||
purpose=purpose,
|
||||
request_metadata=request_metadata,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if scan_result is not None and scan_result.changes:
|
||||
# The caller sees this in the response; a proxy admin needs it server side too,
|
||||
# and it has to land before the post-call hook for logging callbacks to pick it up.
|
||||
get_or_create_metadata_bucket(data)[1]["batch_guardrail"] = scan_result.report().model_dump()
|
||||
verbose_proxy_logger.warning(
|
||||
"batch guardrails changed %s of %s records in %s: %s",
|
||||
len(scan_result.changes),
|
||||
scan_result.scanned_records,
|
||||
file.filename,
|
||||
scan_result.summary(),
|
||||
)
|
||||
if scan_failure is not None:
|
||||
raise_public(scan_failure)
|
||||
|
||||
# Prepare the file data according to FileTypes
|
||||
file_data: Final = (file.filename, file_source, file.content_type)
|
||||
upload_source: Final = (
|
||||
await asyncio.to_thread(rewrite_batch_input_file, file_source, scan_result)
|
||||
if scan_result is not None and scan_result.changes
|
||||
else file_source
|
||||
)
|
||||
file_data: Final = (file.filename, upload_source, file.content_type)
|
||||
|
||||
## check if model is a loadbalanced model
|
||||
router_model: str | None = None
|
||||
is_router_model = False
|
||||
if litellm.enable_loadbalancing_on_batch_endpoints is True:
|
||||
json_obj: Final = get_first_json_object(file_source)
|
||||
json_obj: Final = get_first_json_object(upload_source)
|
||||
if json_obj:
|
||||
router_model = get_model_from_json_obj(json_object=json_obj)
|
||||
is_router_model = is_known_model(model=router_model, llm_router=llm_router)
|
||||
|
|
@ -567,6 +610,9 @@ async def create_file(
|
|||
if _response is not None and isinstance(_response, OpenAIFileObject):
|
||||
response = _response
|
||||
|
||||
if scan_result is not None and scan_result.changes:
|
||||
response.litellm_batch_guardrail = scan_result.report()
|
||||
|
||||
### RESPONSE HEADERS ###
|
||||
hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
|
||||
model_id: Final = hidden_params.get("model_id", None) or ""
|
||||
|
|
|
|||
|
|
@ -279,6 +279,40 @@ OpenAIFilesPurpose = Literal[
|
|||
]
|
||||
|
||||
|
||||
class BatchGuardrailRecord(BaseModel):
|
||||
"""One batch input record a guardrail acted on."""
|
||||
|
||||
line: int
|
||||
"""The 1-based line of the uploaded file the record started on."""
|
||||
|
||||
custom_id: str | None = None
|
||||
"""The record's own `custom_id`, when it carried one."""
|
||||
|
||||
action: Literal["redacted", "dropped"]
|
||||
"""`redacted` means the record was submitted with the guardrail's rewrite applied.
|
||||
|
||||
`dropped` means the guardrail blocked it and it was left out of the submitted file.
|
||||
"""
|
||||
|
||||
guardrail: str | None = None
|
||||
"""Which guardrail dropped the record, when it named itself.
|
||||
|
||||
Set for dropped records only. A guardrail refusing content and a guardrail that is
|
||||
unreachable under a fail-closed setting raise the same way, so this names the guardrail
|
||||
to check rather than claiming a reason it cannot distinguish.
|
||||
"""
|
||||
|
||||
|
||||
class BatchGuardrailReport(BaseModel):
|
||||
"""What guardrails did to a batch input file, per record."""
|
||||
|
||||
submitted_records: int
|
||||
"""How many records reached the provider."""
|
||||
|
||||
modified_records: tuple[BatchGuardrailRecord, ...]
|
||||
"""Every record that was redacted or dropped, in file order."""
|
||||
|
||||
|
||||
class OpenAIFileObject(BaseModel):
|
||||
id: str
|
||||
"""The file identifier, which can be referenced in the API endpoints."""
|
||||
|
|
@ -319,6 +353,12 @@ class OpenAIFileObject(BaseModel):
|
|||
`error` field on `fine_tuning.job`.
|
||||
"""
|
||||
|
||||
litellm_batch_guardrail: BatchGuardrailReport | None = None
|
||||
"""Set by the proxy when guardrails acted on a `purpose=batch` upload.
|
||||
|
||||
Absent on every other upload, so OpenAI-shaped clients see an unchanged response.
|
||||
"""
|
||||
|
||||
_hidden_params: dict = {"response_cost": 0.0} # no cost for writing a file
|
||||
|
||||
def __contains__(self, key) -> bool:
|
||||
|
|
|
|||
|
|
@ -4,9 +4,14 @@ import json
|
|||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.openai_files_endpoints.batch_guardrails import (
|
||||
RedactionRequired,
|
||||
BatchScanResult,
|
||||
RecordDropped,
|
||||
RecordRedacted,
|
||||
rewrite_batch_input_file,
|
||||
UnparseableRecord,
|
||||
UnscannableRecord,
|
||||
raise_public,
|
||||
|
|
@ -64,7 +69,7 @@ def _raise_on(needle, exc):
|
|||
return _hook
|
||||
|
||||
|
||||
async def _scan(source, logging_obj, metadata=None):
|
||||
async def _scan_full(source, logging_obj, metadata=None):
|
||||
return await scan_batch_input_file(
|
||||
file_source=source,
|
||||
request_metadata=metadata if metadata is not None else {},
|
||||
|
|
@ -73,6 +78,14 @@ async def _scan(source, logging_obj, metadata=None):
|
|||
)
|
||||
|
||||
|
||||
async def _scan(source, logging_obj, metadata=None):
|
||||
"""Collapses "the scan found nothing to do" to None so the reject-mode cases read plainly."""
|
||||
result = await _scan_full(source, logging_obj, metadata)
|
||||
if isinstance(result, BatchScanResult):
|
||||
return None if not result.changes else result
|
||||
return result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clean_file_passes_and_rewinds_the_handle():
|
||||
source = _jsonl(_record("a"), _record("b"), _record("c"))
|
||||
|
|
@ -98,7 +111,7 @@ async def test_redaction_is_reported_with_line_and_custom_id():
|
|||
|
||||
failure = await _scan(source, FakeProxyLogging(_redact_containing("secret")))
|
||||
|
||||
assert failure == RedactionRequired(line_number=2, custom_id="dirty")
|
||||
assert [(c.line_number, c.custom_id) for c in failure.changes] == [(2, "dirty")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -198,7 +211,7 @@ async def test_guardrail_that_adds_a_key_is_detected():
|
|||
|
||||
failure = await _scan(_jsonl(_record("a")), FakeProxyLogging(_add_key))
|
||||
|
||||
assert failure == RedactionRequired(line_number=1, custom_id="a")
|
||||
assert [(c.line_number, c.custom_id) for c in failure.changes] == [(1, "a")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -210,7 +223,7 @@ async def test_guardrail_that_adds_a_null_valued_key_is_detected():
|
|||
|
||||
failure = await _scan(_jsonl(_record("a")), FakeProxyLogging(_add_null_key))
|
||||
|
||||
assert failure == RedactionRequired(line_number=1, custom_id="a")
|
||||
assert [(c.line_number, c.custom_id) for c in failure.changes] == [(1, "a")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -223,7 +236,7 @@ async def test_guardrail_that_drops_a_null_valued_key_is_detected():
|
|||
|
||||
failure = await _scan(_jsonl(record), FakeProxyLogging(_drop_null_key))
|
||||
|
||||
assert failure == RedactionRequired(line_number=1, custom_id="a")
|
||||
assert [(c.line_number, c.custom_id) for c in failure.changes] == [(1, "a")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -262,26 +275,6 @@ async def test_empty_url_falls_back_to_its_body_shape(body, expected_call_type):
|
|||
assert logging_obj.seen[0][0] == expected_call_type
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_blocking_guardrail_outranks_an_earlier_refused_record():
|
||||
"""PR 2 turns RedactionRequired into a non-failure; a block must not be lost behind it."""
|
||||
blocked = HTTPException(status_code=403, detail={"error": "Violated guardrail policy"})
|
||||
|
||||
def _hook(data):
|
||||
content = data["messages"][0]["content"]
|
||||
if content == "raiser":
|
||||
raise blocked
|
||||
if content == "redact":
|
||||
data["messages"][0]["content"] = "***"
|
||||
|
||||
source = _jsonl(_record("a", content="redact"), _record("b", content="raiser"))
|
||||
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await _scan(source, FakeProxyLogging(_hook))
|
||||
|
||||
assert raised.value is blocked
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_is_rewound_even_when_a_record_is_refused():
|
||||
source = _jsonl(_record("a", content="secret"))
|
||||
|
|
@ -390,45 +383,14 @@ async def test_url_less_record_whose_body_shape_is_unknown_is_rejected():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_blocking_guardrail_exception_propagates_unwrapped():
|
||||
blocked = HTTPException(status_code=403, detail={"error": "Violated guardrail policy"})
|
||||
blocked = HTTPException(status_code=503, detail={"error": "guardrail service unavailable"})
|
||||
source = _jsonl(_record("a"), _record("b", content="tripwire"))
|
||||
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await _scan(source, FakeProxyLogging(_raise_on("tripwire", blocked)))
|
||||
|
||||
assert raised.value is blocked, "the guardrail's own exception must survive so its status code does"
|
||||
assert raised.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_earliest_refused_record_is_the_one_reported():
|
||||
source = _jsonl(_record("a"), _record("b", content="secret"), _record("c", content="secret"))
|
||||
|
||||
failure = await _scan(source, FakeProxyLogging(_redact_containing("secret")))
|
||||
|
||||
assert failure == RedactionRequired(line_number=2, custom_id="b")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_earliest_failing_record_wins_when_the_raise_comes_first():
|
||||
blocked = HTTPException(status_code=400, detail="blocked")
|
||||
|
||||
def _hook(data):
|
||||
content = data["messages"][0]["content"]
|
||||
if content == "raiser":
|
||||
raise blocked
|
||||
if content == "redact":
|
||||
data["messages"][0]["content"] = "***"
|
||||
|
||||
source = _jsonl(
|
||||
_record("a", content="raiser"),
|
||||
_record("b", content="redact"),
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await _scan(source, FakeProxyLogging(_hook))
|
||||
|
||||
assert raised.value is blocked
|
||||
assert raised.value.status_code == 503
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -447,7 +409,6 @@ async def test_records_are_not_mutated_by_the_scan():
|
|||
[
|
||||
(UnparseableRecord(line_number=7), "line 7"),
|
||||
(UnscannableRecord(line_number=3, custom_id="x", url="/v1/audio/speech"), "custom_id x"),
|
||||
(RedactionRequired(line_number=2, custom_id=None), "line 2"),
|
||||
],
|
||||
)
|
||||
def test_every_failure_maps_to_a_400_naming_the_record(failure, fragment):
|
||||
|
|
@ -463,7 +424,8 @@ async def test_scan_does_not_mutate_the_parsed_record():
|
|||
"""The guardrail must redact a copy. Mutating the record would corrupt what PR 2 writes out."""
|
||||
from litellm.proxy.openai_files_endpoints.batch_guardrails import _ParsedRecord, _scan_record
|
||||
|
||||
record = _ParsedRecord(line_number=1, payload=_record("a", content="my secret is here"))
|
||||
payload = _record("a", content="my secret is here")
|
||||
record = _ParsedRecord(line_number=1, payload=payload)
|
||||
|
||||
failure = await _scan_record(
|
||||
record,
|
||||
|
|
@ -472,7 +434,7 @@ async def test_scan_does_not_mutate_the_parsed_record():
|
|||
FakeProxyLogging(_redact_containing("secret")),
|
||||
)
|
||||
|
||||
assert failure == RedactionRequired(line_number=1, custom_id="a")
|
||||
assert (failure.line_number, failure.custom_id) == (1, "a")
|
||||
assert record.payload["body"]["messages"][0]["content"] == "my secret is here", (
|
||||
"the guardrail redacted the record itself instead of a copy"
|
||||
)
|
||||
|
|
@ -529,4 +491,368 @@ async def test_guardrail_that_returns_a_replacement_dict_is_detected():
|
|||
|
||||
failure = await _scan(_jsonl(_record("a", content="my secret is here")), ReplacingLogging())
|
||||
|
||||
assert failure == RedactionRequired(line_number=1, custom_id="a")
|
||||
assert [(c.line_number, c.custom_id) for c in failure.changes] == [(1, "a")]
|
||||
|
||||
|
||||
def _blocking(needle, status_code=400, guardrail_name="block-guard"):
|
||||
def _hook(data):
|
||||
for message in data.get("messages") or []:
|
||||
if isinstance(message.get("content"), str) and needle in message["content"]:
|
||||
raise HTTPException(
|
||||
status_code=status_code,
|
||||
detail={"error": "Violated guardrail policy", "guardrail_name": guardrail_name},
|
||||
)
|
||||
|
||||
return _hook
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redact_mode_keeps_a_masked_record_instead_of_rejecting():
|
||||
source = _jsonl(_record("a"), _record("b", content="my secret is here"), _record("c"))
|
||||
|
||||
result = await _scan_full(source, FakeProxyLogging(_redact_containing("secret")))
|
||||
|
||||
assert [(c.line_number, c.custom_id) for c in result.changes] == [(2, "b")]
|
||||
rewritten = json.loads(rewrite_batch_input_file(source, result).read().decode().splitlines()[1])
|
||||
assert rewritten["body"]["messages"][0]["content"] == "my *** is here"
|
||||
assert "litellm_metadata" not in rewritten["body"], "proxy metadata must not reach the uploaded file"
|
||||
assert result.submitted_records == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("status_code", [400, 403, 422], ids=["content_policy", "akto", "llm_as_a_judge"])
|
||||
async def test_every_status_litellm_calls_a_block_drops_the_record(status_code):
|
||||
"""Follows CustomGuardrail._is_guardrail_intervention, so drop matches what litellm logs as a block."""
|
||||
source = _jsonl(_record("a"), _record("b", content="tripwire"))
|
||||
|
||||
result = await _scan_full(source, FakeProxyLogging(_blocking("tripwire", status_code)))
|
||||
|
||||
assert result.changes == (RecordDropped(line_number=2, custom_id="b", guardrail="block-guard"),)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redact_mode_drops_a_blocked_record_and_submits_the_rest():
|
||||
source = _jsonl(_record("a"), _record("b", content="tripwire"), _record("c"))
|
||||
|
||||
result = await _scan_full(source, FakeProxyLogging(_blocking("tripwire")))
|
||||
|
||||
assert result.changes == (RecordDropped(line_number=2, custom_id="b", guardrail="block-guard"),)
|
||||
assert result.submitted_records == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("status_code", [500, 502, 408, 429, 401])
|
||||
async def test_redact_mode_does_not_drop_a_record_on_an_infrastructure_failure(status_code):
|
||||
"""A guardrail service that is down must abort the upload, never silently cost the caller records."""
|
||||
source = _jsonl(_record("a"), _record("b", content="tripwire"))
|
||||
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await _scan_full(source, FakeProxyLogging(_blocking("tripwire", status_code)))
|
||||
|
||||
assert raised.value.status_code == status_code
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_every_record_blocked_leaves_nothing_to_submit():
|
||||
source = _jsonl(_record("a", content="tripwire"), _record("b", content="tripwire"))
|
||||
|
||||
result = await _scan_full(source, FakeProxyLogging(_blocking("tripwire")))
|
||||
|
||||
assert result.submitted_records == 0
|
||||
assert [change.line_number for change in result.changes] == [1, 2]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rewrite_drops_blocked_records_and_masks_redacted_ones():
|
||||
records = [_record("a"), _record("b", content="my secret is here"), _record("c", content="tripwire"), _record("d")]
|
||||
source = _jsonl(*records)
|
||||
|
||||
def _hook(data):
|
||||
_redact_containing("secret")(data)
|
||||
_blocking("tripwire")(data)
|
||||
|
||||
result = await _scan_full(source, FakeProxyLogging(_hook))
|
||||
rewritten = rewrite_batch_input_file(source, result)
|
||||
|
||||
lines = [json.loads(line) for line in (rewritten.seek(0), rewritten.read().decode())[1].splitlines()]
|
||||
assert [line["custom_id"] for line in lines] == ["a", "b", "d"]
|
||||
assert lines[1]["body"]["messages"][0]["content"] == "my *** is here"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rewrite_copies_untouched_records_byte_for_byte():
|
||||
"""Enabling the feature must not reformat records no guardrail objected to."""
|
||||
untouched = '{"custom_id":"keep","url":"/v1/chat/completions","body":{"messages":[{"role":"user","content":"hi"}],"model":"m"}}'
|
||||
dirty = json.dumps(_record("dirty", content="my secret is here"))
|
||||
source = io.BytesIO((untouched + "\n" + dirty).encode())
|
||||
|
||||
result = await _scan_full(source, FakeProxyLogging(_redact_containing("secret")))
|
||||
rewritten = rewrite_batch_input_file(source, result)
|
||||
|
||||
assert (rewritten.seek(0), rewritten.read().decode())[1].splitlines()[0] == untouched
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_report_names_every_changed_record_in_file_order():
|
||||
records = [_record("a"), _record("b", content="tripwire"), _record("c", content="my secret is here")]
|
||||
|
||||
def _hook(data):
|
||||
_redact_containing("secret")(data)
|
||||
_blocking("tripwire")(data)
|
||||
|
||||
result = await _scan_full(_jsonl(*records), FakeProxyLogging(_hook))
|
||||
report = result.report()
|
||||
|
||||
assert report.submitted_records == 2
|
||||
assert [(r.line, r.custom_id, r.action, r.guardrail) for r in report.modified_records] == [
|
||||
(2, "b", "dropped", "block-guard"),
|
||||
(3, "c", "redacted", None),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clean_file_needs_no_rewrite():
|
||||
"""A file nothing objected to keeps streaming off disk rather than being buffered in memory."""
|
||||
result = await _scan_full(_jsonl(_record("a"), _record("b")), FakeProxyLogging())
|
||||
|
||||
assert result.changes == ()
|
||||
assert result.submitted_records == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"exc",
|
||||
[
|
||||
GuardrailRaisedException(guardrail_name="g", message="blocked", blocked_content=True),
|
||||
BlockedPiiEntityError(entity_type="US_SSN", guardrail_name="presidio"),
|
||||
],
|
||||
ids=["guardrail_raised", "blocked_pii_entity"],
|
||||
)
|
||||
async def test_litellm_native_block_exceptions_drop_the_record(exc):
|
||||
"""Presidio and friends raise these rather than an HTTPException; they are still policy blocks."""
|
||||
|
||||
def _hook(data):
|
||||
if "tripwire" in data["messages"][0]["content"]:
|
||||
raise exc
|
||||
|
||||
source = _jsonl(_record("a"), _record("b", content="tripwire"))
|
||||
|
||||
result = await _scan_full(source, FakeProxyLogging(_hook))
|
||||
|
||||
assert result.changes == (RecordDropped(line_number=2, custom_id="b", guardrail=exc.guardrail_name),)
|
||||
assert result.submitted_records == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raising_a_native_block_exception_drops_whatever_status_it_carries():
|
||||
"""Raising this type IS the block signal in litellm, so the drop set matches what it logs as a block."""
|
||||
|
||||
def _hook(data):
|
||||
if "tripwire" in data["messages"][0]["content"]:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name="g", message="refused", status_code=503, blocked_content=True
|
||||
)
|
||||
|
||||
result = await _scan_full(_jsonl(_record("b", content="tripwire")), FakeProxyLogging(_hook))
|
||||
|
||||
assert result.changes == (RecordDropped(line_number=1, custom_id="b", guardrail="g"),)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_unreachable_guardrail_aborts_instead_of_quietly_dropping_the_record():
|
||||
"""Several integrations raise this same exception when their backend is down and they fail closed."""
|
||||
|
||||
def _hook(data):
|
||||
if "tripwire" in data["messages"][0]["content"]:
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name="g", message="Singulr API unreachable (block_on_error=True): timed out"
|
||||
)
|
||||
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
await _scan_full(_jsonl(_record("a"), _record("b", content="tripwire")), FakeProxyLogging(_hook))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_400_that_is_not_a_guardrail_decision_still_aborts():
|
||||
"""A guardrail's own HTTP client can raise a 400 because OUR payload was rejected, not the content."""
|
||||
from litellm.exceptions import BadRequestError
|
||||
|
||||
def _hook(data):
|
||||
raise BadRequestError(message="guardrail service rejected the payload", model="m", llm_provider="p")
|
||||
|
||||
with pytest.raises(BadRequestError):
|
||||
await _scan_full(_jsonl(_record("a")), FakeProxyLogging(_hook))
|
||||
|
||||
|
||||
async def _rewritten_body(record, hook):
|
||||
"""Scan one record and hand back the body as it lands in the uploaded file."""
|
||||
source = _jsonl(record)
|
||||
result = await _scan_full(source, FakeProxyLogging(hook))
|
||||
rewritten = rewrite_batch_input_file(source, result)
|
||||
return json.loads(rewritten.read().decode())["body"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_redacted_record_keeps_its_own_body_metadata():
|
||||
"""`metadata` is a real chat-completions parameter; redaction must not silently drop it."""
|
||||
record = _record("m", content="my secret is here")
|
||||
record["body"]["metadata"] = {"team": "finance"}
|
||||
|
||||
body = await _rewritten_body(record, _redact_containing("secret"))
|
||||
|
||||
assert body["metadata"] == {"team": "finance"}
|
||||
assert body["messages"][0]["content"] == "my *** is here"
|
||||
assert "litellm_metadata" not in body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_redacted_record_keeps_its_own_litellm_metadata():
|
||||
"""Tags ride in litellm_metadata; a guardrail firing must not change how the record is attributed."""
|
||||
record = _record("m", content="my secret is here")
|
||||
record["body"]["litellm_metadata"] = {"tags": ["cost-center-42"]}
|
||||
|
||||
body = await _rewritten_body(record, _redact_containing("secret"))
|
||||
|
||||
assert body["litellm_metadata"] == {"tags": ["cost-center-42"]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_redacted_record_keeps_an_explicitly_null_metadata():
|
||||
"""An absent key and a null one are different records, so redaction must not collapse them."""
|
||||
record = _record("m", content="my secret is here")
|
||||
record["body"]["metadata"] = None
|
||||
|
||||
body = await _rewritten_body(record, _redact_containing("secret"))
|
||||
|
||||
assert "metadata" in body and body["metadata"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_log_summary_cannot_be_used_to_forge_log_lines():
|
||||
"""custom_id is caller-supplied and lands in a log line, so control characters must not survive."""
|
||||
forged = "a\nWARNING: proxy shutting down"
|
||||
result = await _scan_full(_jsonl(_record(forged, content="tripwire")), FakeProxyLogging(_blocking("tripwire")))
|
||||
|
||||
summary = result.summary()
|
||||
|
||||
assert "\n" not in summary
|
||||
assert "a WARNING: proxy shutting down" in summary
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_log_summary_is_capped_so_one_upload_cannot_flood_it():
|
||||
records = [_record(f"row-{index}", content="tripwire") for index in range(60)]
|
||||
result = await _scan_full(_jsonl(*records), FakeProxyLogging(_blocking("tripwire")))
|
||||
|
||||
summary = result.summary()
|
||||
|
||||
assert summary.endswith("and 10 more")
|
||||
assert "row-49" in summary and "row-50" not in summary
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_scan_keeps_rewritten_records_off_the_heap():
|
||||
"""A file whose records are mostly rewritten must not build a second copy of itself in memory."""
|
||||
import dataclasses
|
||||
|
||||
bulky = "my secret is here" + ("x" * 50_000)
|
||||
result = await _scan_full(
|
||||
_jsonl(*(_record(str(index), content=bulky) for index in range(4))),
|
||||
FakeProxyLogging(_redact_containing("secret")),
|
||||
)
|
||||
|
||||
retained = sum(
|
||||
len(value)
|
||||
for change in result.changes
|
||||
for value in (getattr(change, field.name) for field in dataclasses.fields(change))
|
||||
if isinstance(value, str)
|
||||
)
|
||||
assert len(result.changes) == 4
|
||||
assert retained < 100, f"{retained} bytes of record text retained per scan"
|
||||
assert result.redactions.tell() > 200_000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_uploaded_file_is_what_the_loadbalancing_model_sniff_reads():
|
||||
"""If line 1 is dropped, the router must not pick its model from a record nobody submitted."""
|
||||
dropped_first = {
|
||||
"custom_id": "gone",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": "model-a", "messages": [{"role": "user", "content": "tripwire"}]},
|
||||
}
|
||||
kept = {
|
||||
"custom_id": "kept",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": "model-b", "messages": [{"role": "user", "content": "fine"}]},
|
||||
}
|
||||
source = _jsonl(dropped_first, kept)
|
||||
|
||||
result = await _scan_full(source, FakeProxyLogging(_blocking("tripwire")))
|
||||
rewritten = rewrite_batch_input_file(source, result)
|
||||
|
||||
first_line = json.loads((rewritten.seek(0), rewritten.read().decode())[1].splitlines()[0])
|
||||
assert first_line["custom_id"] == "kept"
|
||||
assert first_line["body"]["model"] == "model-b"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_infrastructure_failure_outranks_a_redaction_and_aborts():
|
||||
"""A record we could not inspect must abort the upload even when an earlier record was rewritten."""
|
||||
down = HTTPException(status_code=503, detail={"error": "guardrail service unavailable"})
|
||||
|
||||
def _hook(data):
|
||||
content = data["messages"][0]["content"]
|
||||
if content == "raiser":
|
||||
raise down
|
||||
if content == "redact":
|
||||
data["messages"][0]["content"] = "***"
|
||||
|
||||
source = _jsonl(_record("a", content="redact"), _record("b", content="raiser"))
|
||||
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await _scan_full(source, FakeProxyLogging(_hook))
|
||||
|
||||
assert raised.value is down
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_earliest_unscannable_record_is_the_one_reported():
|
||||
source = _jsonl(
|
||||
_record("a"),
|
||||
{"custom_id": "bad-1", "url": "/v1/rerank", "body": {"model": "m"}},
|
||||
{"custom_id": "bad-2", "url": "/v1/rerank", "body": {"model": "m"}},
|
||||
)
|
||||
|
||||
failure = await _scan_full(source, FakeProxyLogging())
|
||||
|
||||
assert failure == UnscannableRecord(line_number=2, custom_id="bad-1", url="/v1/rerank")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_dropped_record_names_the_guardrail_from_an_enriched_http_detail():
|
||||
"""litellm stamps guardrail_name into a block's detail dict; the report should carry it through."""
|
||||
blocked = HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": "Violated guardrail policy", "guardrail_name": "zscaler"},
|
||||
)
|
||||
|
||||
def _hook(data):
|
||||
if "tripwire" in data["messages"][0]["content"]:
|
||||
raise blocked
|
||||
|
||||
result = await _scan_full(_jsonl(_record("b", content="tripwire")), FakeProxyLogging(_hook))
|
||||
|
||||
assert result.changes == (RecordDropped(line_number=1, custom_id="b", guardrail="zscaler"),)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_dropped_record_without_a_named_guardrail_reports_none():
|
||||
"""An unnamed block still drops; the report just cannot say which guardrail did it."""
|
||||
|
||||
def _hook(data):
|
||||
if "tripwire" in data["messages"][0]["content"]:
|
||||
raise HTTPException(status_code=400, detail="blocked")
|
||||
|
||||
result = await _scan_full(_jsonl(_record("b", content="tripwire")), FakeProxyLogging(_hook))
|
||||
|
||||
assert result.changes == (RecordDropped(line_number=1, custom_id="b", guardrail=None),)
|
||||
|
|
|
|||
|
|
@ -3541,8 +3541,8 @@ def _batch_upload(client_, content: bytes, purpose: str = "batch"):
|
|||
b'{"custom_id":"r-0","method":"POST","url":"/v1/chat/completions",'
|
||||
b'"body":{"model":"gpt-3.5-turbo","messages":[{"role":"user","content":"leak me"}]}}\n',
|
||||
"batch",
|
||||
400,
|
||||
"A guardrail changed batch input line 1",
|
||||
200,
|
||||
None,
|
||||
),
|
||||
(
|
||||
b'{"custom_id":"r-0","method":"POST","url":"/v1/chat/completions",'
|
||||
|
|
@ -3602,3 +3602,81 @@ def test_batch_upload_runs_guardrails_on_each_record(
|
|||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
||||
|
||||
def test_batch_upload_redacts_per_record(monkeypatch, llm_router: Router):
|
||||
"""An offending record is submitted masked, matching what the online path does per request."""
|
||||
expected_custom_ids = ["keep-1", "dirty", "keep-2"]
|
||||
import json as _json
|
||||
|
||||
import litellm
|
||||
import litellm.proxy.openai_files_endpoints.files_endpoints as fe
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
class _Redactor(CustomGuardrail):
|
||||
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
||||
for message in data.get("messages") or []:
|
||||
if isinstance(message.get("content"), str) and "leak" in message["content"]:
|
||||
message["content"] = message["content"].replace("leak", "***")
|
||||
return data
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router)
|
||||
setup_proxy_logging_object(monkeypatch, llm_router)
|
||||
monkeypatch.setattr(litellm, "callbacks", [_Redactor(guardrail_name="g", default_on=True)])
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
||||
uploaded = {}
|
||||
|
||||
async def fake_route_create_file(**kwargs):
|
||||
handle = kwargs["_create_file_request"]["file"][1]
|
||||
uploaded["body"] = handle.read() if hasattr(handle, "read") else handle
|
||||
return OpenAIFileObject(
|
||||
id="dummy-id",
|
||||
object="file",
|
||||
bytes=0,
|
||||
created_at=1234567890,
|
||||
filename="batch.jsonl",
|
||||
purpose="batch",
|
||||
status="uploaded",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(fe, "route_create_file", fake_route_create_file)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user"
|
||||
)
|
||||
|
||||
def _row(custom_id, content):
|
||||
return _json.dumps(
|
||||
{
|
||||
"custom_id": custom_id,
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": content}]},
|
||||
}
|
||||
)
|
||||
|
||||
content = ("\n".join([_row("keep-1", "fine"), _row("dirty", "please leak this"), _row("keep-2", "fine")])).encode()
|
||||
try:
|
||||
resp = client.post(
|
||||
"/v1/files",
|
||||
files={"file": ("batch.jsonl", content, "application/jsonl")},
|
||||
data={"purpose": "batch"},
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
rows = [_json.loads(line) for line in uploaded["body"].decode().splitlines()]
|
||||
assert [row["custom_id"] for row in rows] == expected_custom_ids
|
||||
assert rows[1]["body"]["messages"][0]["content"] == "please *** this"
|
||||
report = resp.json()["litellm_batch_guardrail"]
|
||||
assert report["submitted_records"] == 3
|
||||
assert report["modified_records"] == [
|
||||
{"line": 2, "custom_id": "dirty", "action": "redacted", "guardrail": None}
|
||||
]
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue