From e36a96403785ac617022f65fb0c2f40a6aaba19b Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Thu, 20 Aug 2026 00:54:41 -0700 Subject: [PATCH] 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. --- litellm/exceptions.py | 13 + litellm/integrations/custom_guardrail.py | 75 +-- .../guardrail_hooks/deepkeep/deepkeep.py | 1 + .../generic_guardrail_api.py | 1 + .../promptguard/promptguard.py | 1 + .../guardrail_hooks/singulr/singulr.py | 1 + .../guardrail_hooks/straiker/straiker.py | 1 + .../guardrail_hooks/tool_permission.py | 4 +- .../vigil_guard/vigil_guard.py | 2 + .../batch_guardrails.py | 268 +++++++++-- .../openai_files_endpoints/files_endpoints.py | 74 ++- litellm/types/llms/openai.py | 40 ++ .../test_batch_guardrails.py | 452 +++++++++++++++--- .../test_files_endpoint.py | 82 +++- 14 files changed, 864 insertions(+), 151 deletions(-) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 2eb4232fef9..286f7528896 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -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) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 0172c789d1e..f2e390625f5 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py index e1c0653ebd3..c0f72af7576 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py +++ b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py @@ -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( diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 16768a4b08f..e3cf645ceaf 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -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( diff --git a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py index 4775a8b3caa..c25f704567e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py @@ -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": diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py index cd9da8a58b7..3865ba4ed0e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py index eee66f93b7a..bb866db4009 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py +++ b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py @@ -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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 0514d2ab6f7..3c5625bc272 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -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( ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py b/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py index ee1aade8ea6..6b8148645aa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py @@ -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": diff --git a/litellm/proxy/openai_files_endpoints/batch_guardrails.py b/litellm/proxy/openai_files_endpoints/batch_guardrails.py index 349039a050a..237d3f6577a 100644 --- a/litellm/proxy/openai_files_endpoints/batch_guardrails.py +++ b/litellm/proxy/openai_files_endpoints/batch_guardrails.py @@ -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 diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 5794d490bb3..9dbd62fb7b5 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -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 "" diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 6457b285cb5..1588c650177 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -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: diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py index c54debc2418..5715f2558b3 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py @@ -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),) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 8f96fcd796a..ca073f9bb10 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -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()