diff --git a/litellm/proxy/openai_files_endpoints/batch_guardrails.py b/litellm/proxy/openai_files_endpoints/batch_guardrails.py new file mode 100644 index 000000000000..349039a050a3 --- /dev/null +++ b/litellm/proxy/openai_files_endpoints/batch_guardrails.py @@ -0,0 +1,304 @@ +""" +Run the configured pre-call guardrails over every record of a batch input file. + +Runs after ``batch_file_validation.check_batch_file_upload``, so every line here is already known +to parse as a JSON object carrying ``custom_id``, ``method``, ``url`` and ``body``. +""" + +from __future__ import annotations + +import asyncio +import copy +import json +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, BinaryIO, Final, NoReturn, TypeAlias +from urllib.parse import urlsplit + +from fastapi import HTTPException +from typing_extensions import assert_never + +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.utils import CallTypes, CallTypesLiteral + +if TYPE_CHECKING: + from litellm.proxy.utils import ProxyLogging + +EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) + +_SCAN_WINDOW: Final = 32 + +_SCAN_METADATA_KEY: Final = "litellm_metadata" + +# `metadata` is dropped rather than diffed: guardrail dispatch writes its bookkeeping into it +# whenever the payload has one, and a record's own metadata is not scanned content on the +# online path either. +_INJECTED_KEYS: Final = frozenset({_SCAN_METADATA_KEY, "metadata"}) + +# Only what guardrail dispatch reads. The parent OTel span is deliberately left out: parenting one +# guardrail span per record would put tens of thousands of spans on a single upload's trace. +_SCAN_METADATA_KEYS: Final = frozenset( + { + "guardrails", + "_guardrail_pipelines", + "_pipeline_managed_guardrails", + "user_api_key_metadata", + "user_api_key_team_metadata", + "tags", + "headers", + } +) + +_SCANNABLE_CALL_TYPES: Final = frozenset( + { + CallTypes.acompletion, + CallTypes.atext_completion, + CallTypes.aembedding, + CallTypes.aresponses, + CallTypes.anthropic_messages, + } +) + +# Mirrors the record classifier in litellm/llms/bedrock/files/transformation.py, so a record +# litellm already accepts without a url keeps working. +_BODY_SHAPE_CALL_TYPES: Final = ( + ("messages", CallTypes.acompletion), + ("prompt", CallTypes.atext_completion), + ("input", CallTypes.aembedding), +) + + +@dataclass(frozen=True, slots=True) +class UnparseableRecord: + line_number: int + + +@dataclass(frozen=True, slots=True) +class UnscannableRecord: + line_number: int + custom_id: str | None + url: str | None + + +@dataclass(frozen=True, slots=True) +class RedactionRequired: + line_number: int + custom_id: str | None + + +BatchScanFailure: TypeAlias = UnparseableRecord | UnscannableRecord | RedactionRequired + + +@dataclass(frozen=True, slots=True) +class _ParsedRecord: + line_number: int + payload: Mapping[str, object] + + +def _rejected(message: str) -> HTTPException: + return HTTPException(status_code=400, detail={"error": message}) # mutable-ok: FastAPI detail shape + + +def raise_public(failure: BatchScanFailure) -> NoReturn: + """Map a scan failure onto the 400 contract the files endpoint already returns.""" + match failure: + case UnparseableRecord(line_number=line_number): + raise _rejected( + f"The 'body' of batch input line {line_number} is not an object, so guardrails cannot be applied to it" + ) + case UnscannableRecord(line_number=line_number, custom_id=custom_id, url=url): + raise _rejected( + f"Batch input line {line_number}{_describe(custom_id)} targets {url or 'no url'} " + "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 _describe(custom_id: str | None) -> str: + return f" (custom_id {custom_id})" if custom_id else "" + + +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)) + + +def _call_type_from_url(url: str) -> CallTypesLiteral | None: + """ + Resolve the route a record names, tolerating how callers actually write it. + + An absolute url has to reduce to its path or nothing matches, and a record naming + ``/v1/responses`` in full would fall through to its body, where ``input`` reads as an + embedding and the record gets scanned as the wrong call type rather than the right one. + """ + path: Final = urlsplit(url).path.split("?")[0].rstrip("/") + call_types: Final = get_call_types_for_route(path) + if call_types is None: + return None + scannable: Final = next((c for c in call_types if c in _SCANNABLE_CALL_TYPES), None) + return None if scannable is None else scannable.value + + +def _call_type_from_body(body: Mapping[str, object]) -> CallTypesLiteral | None: + shape: Final = next((call_type for field, call_type in _BODY_SHAPE_CALL_TYPES if field in body), None) + return None if shape is None else shape.value + + +def _scannable_call_type(url: object, body: Mapping[str, object]) -> CallTypesLiteral | None: + """ + Resolve how to scan a record: its url when we recognize one, otherwise its body shape. + + An unrecognized url falls through to the body rather than rejecting, because a record we can + still read is a record we can still scan, and the provider transformers treat an unknown url + as chat rather than as an error. + """ + from_url: Final = _call_type_from_url(url) if isinstance(url, str) and url else None + return from_url if from_url is not None else _call_type_from_body(body) + + +def _custom_id_of(payload: Mapping[str, object]) -> str | None: + custom_id: Final = payload.get("custom_id") + return custom_id if isinstance(custom_id, str) else None + + +def _fingerprint(body: Mapping[str, object], keys: frozenset[str]) -> str: + """ + Order-insensitive projection, so a guardrail re-serializing a dict does not read as a change. + + An absent key projects to ``null`` while a key holding ``None`` projects to the string + ``"null"``, so adding or dropping a null-valued key still reads as a change. + """ + return json.dumps( + tuple( + (key, json.dumps(body[key], sort_keys=True, default=str) if key in body else None) for key in sorted(keys) + ) + ) + + +def build_scan_metadata(request_metadata: Mapping[str, object]) -> Mapping[str, object]: + """ + Narrow the request metadata to the keys guardrail dispatch reads. + + Passing the whole thing through would carry values that cannot be copied, such as the parent + OTel span, and would hand every record proxy state it has no business seeing. + """ + return MappingProxyType( + {key: value for key, value in request_metadata.items() if key in _SCAN_METADATA_KEYS} + ) # mutable-ok: MappingProxyType freezes the comprehension + + +async def _scan_record( + record: _ParsedRecord, + scan_metadata: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, +) -> BatchScanFailure | None: + body: Final = record.payload.get("body") + if not isinstance(body, dict): + return UnparseableRecord(line_number=record.line_number) + + custom_id: Final = _custom_id_of(record.payload) + url: Final = record.payload.get("url") + call_type: Final = _scannable_call_type(url, body) + if call_type is None: + return UnscannableRecord( + line_number=record.line_number, + custom_id=custom_id, + url=url if isinstance(url, str) else None, + ) + + scan_input: Final[dict[str, object]] = copy.deepcopy(body) # mutable-ok: pre_call_hook mutates the dict it is given + 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 + # place would otherwise leak across records and back into the request. The narrowing above + # 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, + ) + + 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 + + +async def _scan_window( + window: tuple[_ParsedRecord, ...], + scan_metadata: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, +) -> tuple[tuple[int, BatchScanFailure | 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), + return_exceptions=True, + ) + return tuple((record.line_number, outcome) for record, outcome in zip(window, outcomes) if outcome is not None) + + +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)) + return min(raised or problems, key=lambda problem: problem[0])[1] + + +async def scan_batch_input_file( + *, + file_source: BinaryIO, + request_metadata: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, +) -> BatchScanFailure | None: + """ + 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. + """ + scan_metadata: Final = build_scan_metadata(request_metadata) + problems: Final[list[tuple[int, BatchScanFailure | BaseException]]] = [] # mutable-ok: spans windows + window: Final[list[_ParsedRecord]] = [] # mutable-ok: bounded read-ahead buffer + + async def drain() -> None: + if window: + problems.extend(await _scan_window(tuple(window), scan_metadata, user_api_key_dict, proxy_logging_obj)) + window.clear() + + try: + for item in _iter_records(file_source): + window.append(item) + if len(window) >= _SCAN_WINDOW: + await drain() + if problems: + break + if not problems: + await drain() + finally: + file_source.seek(0) + + if not problems: + return None + worst: Final = _worst(tuple(problems)) + if isinstance(worst, BaseException): + raise worst + return worst diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index b7200de8fb69..5794d490bb35 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -46,6 +46,11 @@ check_batch_file_upload, raise_batch_file_validation_failure, ) +from litellm.proxy.openai_files_endpoints.batch_guardrails import ( + EMPTY_MAPPING, + raise_public, + scan_batch_input_file, +) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, add_internal_model_credentials, @@ -471,6 +476,22 @@ async def create_file( proxy_config=proxy_config, ) + # /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, + ) + 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) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1d042e2521b7..58e6f5a94e2d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1522,6 +1522,23 @@ def _handle_pipeline_result( return data + def has_pre_call_guardrails(self, request_metadata: Mapping[str, object]) -> bool: + """ + Whether any guardrail or guardrail pipeline would inspect a request carrying this metadata. + + Evaluated with the same predicate the pre-call loop uses, so a proxy configured only with + post-call guardrails answers False. Callers that must pay a real cost to build the hook's + input, such as streaming a batch input file off disk, use this to skip that work. + """ + if request_metadata.get("_guardrail_pipelines"): + return True + probe: Final = {"metadata": dict(request_metadata)} # mutable-ok: should_run_guardrail takes a dict + return any( + isinstance(callback, CustomGuardrail) + and callback.should_run_guardrail(data=probe, event_type=GuardrailEventHooks.pre_call) + for callback in ProxyLogging._callback_capabilities().resolved_callbacks + ) + # The actual implementation of the function @overload async def pre_call_hook( @@ -1529,6 +1546,7 @@ async def pre_call_hook( user_api_key_dict: UserAPIKeyAuth, data: None, call_type: CallTypesLiteral, + guardrails_only: bool = False, ) -> None: pass @@ -1538,6 +1556,7 @@ async def pre_call_hook( user_api_key_dict: UserAPIKeyAuth, data: dict, call_type: CallTypesLiteral, + guardrails_only: bool = False, ) -> dict: pass @@ -1546,6 +1565,7 @@ async def pre_call_hook( user_api_key_dict: UserAPIKeyAuth, data: dict | None, call_type: CallTypesLiteral, + guardrails_only: bool = False, ) -> dict | None: """ Allows users to modify/reject the incoming request to the proxy, without having to deal with parsing Request body. @@ -1554,10 +1574,15 @@ async def pre_call_hook( 1. /chat/completions 2. /embeddings 3. /image/generation + + With ``guardrails_only`` the walk is limited to guardrails and guardrail pipelines: rate + limiting, budget accounting, prompt templates and hanging-request alerting are skipped. + Use it to scan a payload that is not itself a request, such as one record of a batch file. """ verbose_proxy_logger.debug("Inside Proxy Logging Pre-call hook!") - self._init_response_taking_too_long_task(data=data) + if not guardrails_only: + self._init_response_taking_too_long_task(data=data) if data is None: return None @@ -1569,7 +1594,8 @@ async def pre_call_hook( ## PROMPT TEMPLATE CHECK ## if ( - litellm_logging_obj is not None + not guardrails_only + and litellm_logging_obj is not None and prompt_id is not None and (call_type == "completion" or call_type == "acompletion") ): @@ -1600,7 +1626,7 @@ async def pre_call_hook( # CustomGuardrail is configured. Saves the loop overhead + # ``time.time()`` x2 per registered callback for the common # "callbacks=[]" case on small / dev deployments. - if not caps.has_guardrail and not caps.has_pre_call_override: + if not caps.has_guardrail and (guardrails_only or not caps.has_pre_call_override): if data is not None: self._process_guardrail_metadata(data) return data @@ -1637,7 +1663,8 @@ async def pre_call_hook( data = result elif ( - _callback is not None + not guardrails_only + and _callback is not None and isinstance(_callback, CustomLogger) and "async_pre_call_hook" in vars(_callback.__class__) and _callback.__class__.async_pre_call_hook != CustomLogger.async_pre_call_hook 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 new file mode 100644 index 000000000000..c54debc24189 --- /dev/null +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py @@ -0,0 +1,532 @@ +import io +import json + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.openai_files_endpoints.batch_guardrails import ( + RedactionRequired, + UnparseableRecord, + UnscannableRecord, + raise_public, + scan_batch_input_file, +) + + +def _record(custom_id, content="hello", url="/v1/chat/completions"): + return { + "custom_id": custom_id, + "method": "POST", + "url": url, + "body": { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": content}], + }, + } + + +def _jsonl(*records): + return io.BytesIO("\n".join(json.dumps(r) for r in records).encode()) + + +class FakeProxyLogging: + """Stands in for ProxyLogging so the scan can be driven without a live proxy.""" + + def __init__(self, on_record=None): + self.on_record = on_record or (lambda data: None) + self.seen = [] + + async def pre_call_hook(self, user_api_key_dict, data, call_type, guardrails_only=False): + self.seen.append((call_type, json.dumps(data.get("messages"), sort_keys=True))) + self.on_record(data) + return data + + def has_pre_call_guardrails(self, request_metadata): + return True + + +def _redact_containing(needle): + def _hook(data): + for message in data.get("messages") or []: + if isinstance(message.get("content"), str) and needle in message["content"]: + message["content"] = message["content"].replace(needle, "***") + + return _hook + + +def _raise_on(needle, exc): + def _hook(data): + for message in data.get("messages") or []: + if isinstance(message.get("content"), str) and needle in message["content"]: + raise exc + + return _hook + + +async def _scan(source, logging_obj, metadata=None): + return await scan_batch_input_file( + file_source=source, + request_metadata=metadata if metadata is not None else {}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=logging_obj, + ) + + +@pytest.mark.asyncio +async def test_clean_file_passes_and_rewinds_the_handle(): + source = _jsonl(_record("a"), _record("b"), _record("c")) + logging_obj = FakeProxyLogging() + + assert await _scan(source, logging_obj) is None + assert len(logging_obj.seen) == 3 + assert source.tell() == 0, "handle must be rewound so the upload still sees the whole file" + + +@pytest.mark.asyncio +async def test_every_record_is_scanned_not_just_the_first(): + records = [_record(f"r{i}") for i in range(70)] + logging_obj = FakeProxyLogging() + + assert await _scan(_jsonl(*records), logging_obj) is None + assert len(logging_obj.seen) == 70, "records past the first scan window must still be scanned" + + +@pytest.mark.asyncio +async def test_redaction_is_reported_with_line_and_custom_id(): + source = _jsonl(_record("keep-1"), _record("dirty", content="my secret is here"), _record("keep-2")) + + failure = await _scan(source, FakeProxyLogging(_redact_containing("secret"))) + + assert failure == RedactionRequired(line_number=2, custom_id="dirty") + + +@pytest.mark.asyncio +async def test_body_carrying_its_own_metadata_is_not_reported_as_redacted(): + record = _record("has-meta") + record["body"]["metadata"] = {"team": "finance"} + + failure = await _scan(_jsonl(record), FakeProxyLogging(), metadata={"guardrails": ["x"]}) + + assert failure is None, "the metadata the proxy injects must not be diffed as record content" + + +@pytest.mark.asyncio +async def test_guardrail_writing_bookkeeping_into_metadata_is_not_a_redaction(): + def _touch_metadata(data): + data["litellm_metadata"]["applied_guardrails"] = ["some-guard"] + + assert await _scan(_jsonl(_record("a")), FakeProxyLogging(_touch_metadata)) is None + + +@pytest.mark.asyncio +async def test_records_own_metadata_is_left_out_of_the_scan_and_the_diff(): + """Guardrail dispatch writes bookkeeping into `metadata`; diffing it would reject every such record.""" + record = _record("has-meta") + record["body"]["metadata"] = {"team": "finance"} + seen = [] + + def _write_bookkeeping(data): + seen.append("metadata" in data) + data.setdefault("metadata", {})["applied_guardrails"] = ["g"] + + assert await _scan(_jsonl(record), FakeProxyLogging(_write_bookkeeping)) is None + assert seen == [False], "the record's own metadata must not be handed to guardrail dispatch" + assert record["body"]["metadata"] == {"team": "finance"} + + +@pytest.mark.asyncio +async def test_request_metadata_is_narrowed_to_what_guardrails_read(): + """An OTel-enabled proxy puts a lock-bearing span here; a per-record copy of it is a crash.""" + import threading + + seen = [] + metadata = { + "guardrails": ["g"], + "tags": ["t"], + "headers": {"x-noma-application-id": "app-1"}, + "litellm_parent_otel_span": threading.RLock(), + "user_api_key": "sk-secret", + } + + failure = await _scan( + _jsonl(_record("a")), + FakeProxyLogging(lambda d: seen.append(dict(d["litellm_metadata"]))), + metadata=metadata, + ) + + assert failure is None + assert seen == [{"guardrails": ["g"], "tags": ["t"], "headers": {"x-noma-application-id": "app-1"}}] + + +@pytest.mark.asyncio +async def test_one_record_cannot_leak_a_metadata_write_into_the_next_one(): + """`headers` and `tags` are nested and shared; an in-place write must not cross records.""" + seen = [] + + def _tamper(data): + bag = data["litellm_metadata"] + seen.append((dict(bag["headers"]), list(bag["tags"]))) + bag["headers"]["x-injected"] = "from-record-1" + bag["tags"].append("from-record-1") + + metadata = {"guardrails": ["g"], "headers": {"x-real": "yes"}, "tags": ["real"]} + await _scan(_jsonl(_record("a"), _record("b")), FakeProxyLogging(_tamper), metadata=metadata) + + assert seen == [({"x-real": "yes"}, ["real"]), ({"x-real": "yes"}, ["real"])] + assert metadata == {"guardrails": ["g"], "headers": {"x-real": "yes"}, "tags": ["real"]} + + +@pytest.mark.asyncio +async def test_records_are_scanned_under_the_headers_the_upload_carried(): + """Guardrails such as noma pick their application from a header, so dropping it changes the policy.""" + seen = [] + + await _scan( + _jsonl(_record("a")), + FakeProxyLogging(lambda d: seen.append(d["litellm_metadata"].get("headers"))), + metadata={"guardrails": ["g"], "headers": {"x-noma-application-id": "app-1"}}, + ) + + assert seen == [{"x-noma-application-id": "app-1"}] + + +@pytest.mark.asyncio +async def test_guardrail_that_adds_a_key_is_detected(): + def _add_key(data): + data["mock_response"] = "intercepted" + + failure = await _scan(_jsonl(_record("a")), FakeProxyLogging(_add_key)) + + assert failure == RedactionRequired(line_number=1, custom_id="a") + + +@pytest.mark.asyncio +async def test_guardrail_that_adds_a_null_valued_key_is_detected(): + """A null value must not read the same as a missing key, or dropping one hides a change.""" + + def _add_null_key(data): + data["response_format"] = None + + failure = await _scan(_jsonl(_record("a")), FakeProxyLogging(_add_null_key)) + + assert failure == RedactionRequired(line_number=1, custom_id="a") + + +@pytest.mark.asyncio +async def test_guardrail_that_drops_a_null_valued_key_is_detected(): + def _drop_null_key(data): + data.pop("response_format") + + record = _record("a") + record["body"]["response_format"] = None + + failure = await _scan(_jsonl(record), FakeProxyLogging(_drop_null_key)) + + assert failure == RedactionRequired(line_number=1, custom_id="a") + + +@pytest.mark.asyncio +async def test_guardrail_that_only_reorders_a_nested_dict_is_not_a_redaction(): + def _reorder(data): + message = data["messages"][0] + data["messages"][0] = {key: message[key] for key in reversed(list(message))} + + assert await _scan(_jsonl(_record("a")), FakeProxyLogging(_reorder)) is None + + +@pytest.mark.asyncio +async def test_record_without_a_url_falls_back_to_its_body_shape(): + logging_obj = FakeProxyLogging() + record = _record("no-url") + del record["url"] + + assert await _scan(_jsonl(record), logging_obj) is None + assert logging_obj.seen[0][0] == "acompletion" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body, expected_call_type", + [ + ({"messages": [{"role": "user", "content": "x"}]}, "acompletion"), + ({"prompt": "x"}, "atext_completion"), + ({"input": "x"}, "aembedding"), + ], +) +async def test_empty_url_falls_back_to_its_body_shape(body, expected_call_type): + logging_obj = FakeProxyLogging() + record = {"custom_id": "c", "url": "", "body": {"model": "m", **body}} + + assert await _scan(_jsonl(record), logging_obj) is None + 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")) + + await _scan(source, FakeProxyLogging(_redact_containing("secret"))) + + assert source.tell() == 0 + + +@pytest.mark.asyncio +async def test_record_without_a_body_object_is_rejected(): + source = io.BytesIO(b'{"custom_id": "no-body", "url": "/v1/chat/completions"}\n') + + assert await _scan(source, FakeProxyLogging()) == UnparseableRecord(line_number=1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "url, expected_call_type", + [ + ("/v1/chat/completions", "acompletion"), + ("/v1/completions", "atext_completion"), + ("/v1/embeddings", "aembedding"), + ("/v1/responses", "aresponses"), + ("/v1/messages", "anthropic_messages"), + ], +) +async def test_supported_urls_scan_under_the_matching_call_type(url, expected_call_type): + logging_obj = FakeProxyLogging() + + assert await _scan(_jsonl(_record("a", url=url)), logging_obj) is None + assert logging_obj.seen[0][0] == expected_call_type + + +@pytest.mark.asyncio +async def test_unrecognized_url_falls_back_to_the_body_shape(): + """A record we can still read is a record we can still scan, so the url alone must not reject it.""" + logging_obj = FakeProxyLogging() + + assert await _scan(_jsonl(_record("img", url="/v1/images/generations")), logging_obj) is None + assert logging_obj.seen[0][0] == "acompletion" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "url", + ["/chat/completions", "/v1/chat/completions/", "https://api.openai.com/v1/chat/completions"], +) +async def test_url_variants_callers_actually_write_are_accepted(url): + logging_obj = FakeProxyLogging() + + assert await _scan(_jsonl(_record("v", url=url)), logging_obj) is None + assert logging_obj.seen[0][0] == "acompletion" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "url, expected_call_type", + [ + ("https://api.openai.com/v1/responses", "aresponses"), + ("https://api.openai.com/v1/embeddings", "aembedding"), + ("https://api.openai.com/v1/messages", "anthropic_messages"), + ("https://api.openai.com/v1/responses?api-version=1", "aresponses"), + ], +) +async def test_an_absolute_url_resolves_by_path_not_by_body_shape(url, expected_call_type): + """A Responses body carries `input`, which reads as an embedding if the host is not stripped first.""" + logging_obj = FakeProxyLogging() + record = {"custom_id": "abs", "method": "POST", "url": url, "body": {"model": "m", "input": "x"}} + + assert await _scan(_jsonl(record), logging_obj) is None + assert logging_obj.seen[0][0] == expected_call_type + + +@pytest.mark.asyncio +async def test_query_string_on_a_known_url_does_not_change_the_call_type(): + """The body carries `messages`, so only stripping the query string can yield aembedding.""" + logging_obj = FakeProxyLogging() + record = { + "custom_id": "q", + "url": "/v1/embeddings?api-version=1", + "body": {"model": "m", "input": "x", "messages": [{"role": "user", "content": "y"}]}, + } + + assert await _scan(_jsonl(record), logging_obj) is None + assert logging_obj.seen[0][0] == "aembedding" + + +@pytest.mark.asyncio +async def test_record_whose_body_cannot_be_read_is_rejected(): + source = _jsonl({"custom_id": "opaque", "url": "/v1/rerank", "body": {"model": "m", "documents": ["a"]}}) + + failure = await _scan(source, FakeProxyLogging()) + + assert failure == UnscannableRecord(line_number=1, custom_id="opaque", url="/v1/rerank") + + +@pytest.mark.asyncio +async def test_url_less_record_whose_body_shape_is_unknown_is_rejected(): + record = {"custom_id": "opaque", "body": {"model": "m", "something_else": 1}} + + assert await _scan(_jsonl(record), FakeProxyLogging()) == UnscannableRecord( + line_number=1, custom_id="opaque", url=None + ) + + +@pytest.mark.asyncio +async def test_blocking_guardrail_exception_propagates_unwrapped(): + blocked = HTTPException(status_code=403, detail={"error": "Violated guardrail policy"}) + 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 + + +@pytest.mark.asyncio +async def test_records_are_not_mutated_by_the_scan(): + record = _record("a", content="my secret is here") + payload = json.dumps(record) + source = io.BytesIO(payload.encode()) + + await _scan(source, FakeProxyLogging(_redact_containing("secret"))) + + assert source.getvalue().decode() == payload, "the scan must never rewrite the uploaded bytes" + + +@pytest.mark.parametrize( + "failure, fragment", + [ + (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): + with pytest.raises(HTTPException) as raised: + raise_public(failure) + + assert raised.value.status_code == 400 + assert fragment in raised.value.detail["error"] + + +@pytest.mark.asyncio +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")) + + failure = await _scan_record( + record, + {}, + UserAPIKeyAuth(api_key="sk-test"), + FakeProxyLogging(_redact_containing("secret")), + ) + + assert failure == RedactionRequired(line_number=1, custom_id="a") + assert record.payload["body"]["messages"][0]["content"] == "my secret is here", ( + "the guardrail redacted the record itself instead of a copy" + ) + + +@pytest.mark.asyncio +async def test_scan_is_bounded_so_a_huge_file_cannot_fan_out_without_limit(): + import asyncio + + from litellm.proxy.openai_files_endpoints.batch_guardrails import _SCAN_WINDOW + + in_flight = {"now": 0, "peak": 0} + + class CountingLogging(FakeProxyLogging): + async def pre_call_hook(self, user_api_key_dict, data, call_type, guardrails_only=False): + in_flight["now"] += 1 + in_flight["peak"] = max(in_flight["peak"], in_flight["now"]) + await asyncio.sleep(0) + in_flight["now"] -= 1 + return data + + records = [_record(f"r{i}") for i in range(_SCAN_WINDOW * 3)] + + assert await _scan(_jsonl(*records), CountingLogging()) is None + assert in_flight["peak"] <= _SCAN_WINDOW, ( + f"peak {in_flight['peak']} exceeded the scan window; a gigabyte file would fan out unbounded" + ) + + +@pytest.mark.asyncio +async def test_scan_runs_guardrails_only(): + """Rate limiters, budget hooks and the hanging-request alert must not fire once per record.""" + flags = [] + + class FlagCapturingLogging(FakeProxyLogging): + async def pre_call_hook(self, user_api_key_dict, data, call_type, guardrails_only=False): + flags.append(guardrails_only) + return data + + await _scan(_jsonl(_record("a"), _record("b")), FlagCapturingLogging()) + + assert flags == [True, True] + + +@pytest.mark.asyncio +async def test_guardrail_that_returns_a_replacement_dict_is_detected(): + """async_pre_call_hook may return a NEW dict instead of mutating; that result is the real input.""" + + class ReplacingLogging(FakeProxyLogging): + async def pre_call_hook(self, user_api_key_dict, data, call_type, guardrails_only=False): + replacement = json.loads(json.dumps(data)) + replacement["messages"][0]["content"] = "***" + return replacement + + failure = await _scan(_jsonl(_record("a", content="my secret is here")), ReplacingLogging()) + + assert failure == RedactionRequired(line_number=1, custom_id="a") 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 99fb19f0d60c..8f96fcd796a6 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 @@ -3516,3 +3516,89 @@ def test_create_file_non_batch_purpose_skips_batch_validation(monkeypatch, llm_r assert response.status_code == 200, response.text assert len(forwarded_calls) == 1 + + +def _batch_upload(client_, content: bytes, purpose: str = "batch"): + return client_.post( + "/v1/files", + files={"file": ("batch.jsonl", content, "application/jsonl")}, + data={"purpose": purpose}, + headers={"Authorization": "Bearer test-key"}, + ) + + +@pytest.mark.parametrize( + "content, purpose, expected_status, expected_fragment", + [ + ( + b'{"custom_id":"r-0","method":"POST","url":"/v1/chat/completions",' + b'"body":{"model":"gpt-3.5-turbo","messages":[{"role":"user","content":"hi"}]}}\n', + "batch", + 200, + None, + ), + ( + 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", + ), + ( + 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', + "assistants", + 200, + None, + ), + (b"{ not json\n", "batch", 400, "line 1"), + ], +) +def test_batch_upload_runs_guardrails_on_each_record( + monkeypatch, llm_router: Router, content, purpose, expected_status, expected_fragment +): + """POST /v1/files with purpose=batch must reach the guardrail chain; other purposes must not.""" + 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"] = "***" + 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() + + async def fake_route_create_file(**kwargs): + 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" + ) + try: + resp = _batch_upload(client, content, purpose) + assert resp.status_code == expected_status, resp.text + if expected_fragment is not None: + assert expected_fragment in resp.text + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + ProxyLogging._callback_capabilities_cache.clear() diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index 05005dae7970..12fc9310d487 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -166,3 +166,123 @@ def fake_process(d): ) assert out is data assert invoked["data"] is data + + +@pytest.mark.asyncio +async def test_guardrails_only_skips_non_guardrail_pre_call_callbacks( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """Rate limiters and budget hooks ride this same loop; a content scan must not trip them.""" + calls: list[str] = [] + + class _RateLimiterLike(CustomLogger): + async def async_pre_call_hook(self, **kwargs): # type: ignore[override] + calls.append("ran") + return None + + monkeypatch.setattr(litellm, "callbacks", [_RateLimiterLike()]) + proxy_logging.slack_alerting_instance = MagicMock(alerting=None) + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={"messages": [{"x": "input"}], "model": "m"}, + call_type="completion", + guardrails_only=True, + ) + assert calls == [] + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={"messages": [{"x": "input"}], "model": "m"}, + call_type="completion", + ) + assert calls == ["ran"], "the default path must still run non-guardrail pre-call callbacks" + + +@pytest.mark.asyncio +async def test_guardrails_only_skips_the_hanging_request_alert(proxy_logging, make_user_api_key_auth, monkeypatch): + monkeypatch.setattr(litellm, "callbacks", []) + alerting = MagicMock(alerting=True) + proxy_logging.slack_alerting_instance = alerting + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={"messages": [{"x": "input"}], "model": "m"}, + call_type="completion", + guardrails_only=True, + ) + alerting.response_taking_too_long.assert_not_called() + + +@pytest.mark.asyncio +async def test_guardrails_only_skips_prompt_template_rewriting(proxy_logging, make_user_api_key_auth, monkeypatch): + """A prompt template would rewrite messages, which a per-record content diff would misread.""" + monkeypatch.setattr(litellm, "callbacks", []) + proxy_logging.slack_alerting_instance = MagicMock(alerting=None) + process = AsyncMock() + monkeypatch.setattr(proxy_logging, "_process_prompt_template", process) + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={"messages": [{"x": "input"}], "model": "m", "prompt_id": "p1", "litellm_logging_obj": MagicMock()}, + call_type="acompletion", + guardrails_only=True, + ) + process.assert_not_awaited() + + +@pytest.mark.parametrize( + "event_hook, expected", + [("pre_call", True), ("post_call", False), ("during_call", False)], +) +def test_has_pre_call_guardrails_follows_the_guardrail_event_hook(proxy_logging, monkeypatch, event_hook, expected): + """A post-call-only guardrail must not make callers pay for pre-call work.""" + from litellm.integrations.custom_guardrail import CustomGuardrail + + guardrail = CustomGuardrail(guardrail_name="g", event_hook=event_hook, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + assert proxy_logging.has_pre_call_guardrails({}) is expected + + +def test_has_pre_call_guardrails_is_false_without_callbacks(proxy_logging, monkeypatch): + monkeypatch.setattr(litellm, "callbacks", []) + + assert proxy_logging.has_pre_call_guardrails({}) is False + + +def test_has_pre_call_guardrails_is_true_for_a_configured_pipeline(proxy_logging, monkeypatch): + monkeypatch.setattr(litellm, "callbacks", []) + + assert proxy_logging.has_pre_call_guardrails({"_guardrail_pipelines": ["p1"]}) is True + + +@pytest.mark.asyncio +async def test_default_path_still_arms_the_hanging_request_alert(proxy_logging, make_user_api_key_auth, monkeypatch): + """Pins the other side of the gate: without the flag, the alert must still fire.""" + monkeypatch.setattr(litellm, "callbacks", []) + alerting = MagicMock(alerting=True) + alerting.response_taking_too_long = AsyncMock() + proxy_logging.slack_alerting_instance = alerting + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={"messages": [{"x": "input"}], "model": "m"}, + call_type="completion", + ) + alerting.response_taking_too_long.assert_called_once() + + +@pytest.mark.asyncio +async def test_default_path_still_applies_prompt_templates(proxy_logging, make_user_api_key_auth, monkeypatch): + monkeypatch.setattr(litellm, "callbacks", []) + proxy_logging.slack_alerting_instance = MagicMock(alerting=None) + process = AsyncMock() + monkeypatch.setattr(proxy_logging, "_process_prompt_template", process) + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={"messages": [{"x": "input"}], "model": "m", "prompt_id": "p1", "litellm_logging_obj": MagicMock()}, + call_type="acompletion", + ) + process.assert_awaited_once()