mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
fix(proxy): run pre-call guardrails on batch input file uploads (#37519)
* fix(proxy): run pre-call guardrails on batch input file uploads POST /v1/files with purpose=batch was the only route in files_endpoints that never reached pre_call_hook, so guardrails did not see batch content at all and records reached the provider unscanned. Stream the uploaded JSONL a record at a time and run each record's body through the existing pre_call_hook dispatch under the call type its url maps to, so guardrail resolution, key and team config, and the per-endpoint translations are reused rather than reimplemented. The hook gains a guardrails_only mode for this, since the same callback loop also drives rate limiters, budget hooks, prompt templates and hanging-request alerting, none of which should fire once per record. A guardrail that blocks raises its own exception, which propagates untouched so its status code survives. A record a guardrail would rewrite, a record that cannot be parsed, and a record whose url cannot be scanned all reject the upload, since silently skipping any of them is the bypass this is meant to close. Per-record redaction lands separately. The scan only runs when a guardrail that actually runs pre_call, or a guardrail pipeline, is configured, so deployments without one are byte for byte unchanged. * fix(proxy): compare the dict a batch guardrail returns, not the one it was given async_pre_call_hook may return a replacement dict instead of mutating its input, and process_pre_call_hook_response then makes that replacement the request. The scan only inspected the dict it passed in, so a guardrail that redacts by returning a copy was treated as a no-op and its record uploaded unchanged. * fix(proxy): treat a missing batch body key as different from a null one Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(proxy): document the guardrails_only mode on pre_call_hook * fix(proxy): resolve a batch record's scan type from its body when the url is unfamiliar The scanner only accepted five exact urls, but callers write that field by hand and the provider transformers are far more permissive: bedrock treats any non-empty url as chat and vertex strips query strings and trailing slashes. Uploads that work today would have started failing the moment a pre-call guardrail was configured. Normalize the url before lookup and fall back to the body shape when it is unfamiliar, so a record we can still read is a record we still scan. Only a body with no messages, prompt or input is now refused, and the error says so instead of listing urls that were never the whole set. Also pins the default side of the guardrails_only gate: the hanging-request alert and prompt templating are asserted to still fire when the flag is absent. * refactor(proxy): drop batch guardrail checks the upload validation already makes check_batch_file_upload now runs first and rejects a line that does not parse, a line that is not an object, and a line missing custom_id, method, url or body, so the guardrail scan can rely on all four. Its own parse handling was unreachable through the endpoint and is gone, along with the tests for it. What is left is the case that validation does not cover, a body whose value is not an object, since it only checks that the key is present. * fix(proxy): resolve a batch record's call type from the url path, not the whole url A record naming its route in full, which is how callers actually write batch files, matched no known route, so it fell through to the body shape. A Responses record carries `input`, and that reads as an embedding, so the record was scanned as the wrong call type and any guardrail scoped to chat or Responses skipped it while the upload was accepted. Chat records survived only because their body shape happens to map back to the same call type. The url is now reduced to its path before matching. Guardrails that pick their policy from a request header, such as noma choosing an application id, saw no headers at all during the scan and fell back to a default, so a batch record could be evaluated under a different policy than the same content sent online. The sanitized headers the proxy already stores in request metadata now travel with the scan. Also drops the bare `dict` annotation, the unreachable non-dict branch on the guardrail chain's own return, and the type alias that was missing its `TypeAlias`, which together were failing the lint gate. * fix(proxy): give each batch record its own copy of the scan metadata The narrowed metadata was handed to every record as a shallow copy, so `headers` and `tags` stayed shared with the upload request and with the other records in the same window. A guardrail that writes into one of those in place, which several do to record their own bookkeeping, would have its write show up in every record scanned after it and in the request itself. The narrowing already removed the values that cannot be copied, so each record now gets a deep copy. --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng-berri
GitHub
parent
122675c309
commit
3a31331435
6 changed files with 1094 additions and 4 deletions
304
litellm/proxy/openai_files_endpoints/batch_guardrails.py
Normal file
304
litellm/proxy/openai_files_endpoints/batch_guardrails.py
Normal file
|
|
@ -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
|
||||
|
|
@ -46,6 +46,11 @@ from litellm.proxy.openai_files_endpoints.batch_file_validation import (
|
|||
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)
|
||||
|
||||
|
|
|
|||
|
|
@ -1522,6 +1522,23 @@ class ProxyLogging:
|
|||
|
||||
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 @@ class ProxyLogging:
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
data: None,
|
||||
call_type: CallTypesLiteral,
|
||||
guardrails_only: bool = False,
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
|
|
@ -1538,6 +1556,7 @@ class ProxyLogging:
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
data: dict,
|
||||
call_type: CallTypesLiteral,
|
||||
guardrails_only: bool = False,
|
||||
) -> dict:
|
||||
pass
|
||||
|
||||
|
|
@ -1546,6 +1565,7 @@ class ProxyLogging:
|
|||
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 @@ class ProxyLogging:
|
|||
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 @@ class ProxyLogging:
|
|||
## 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 @@ class ProxyLogging:
|
|||
# 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 @@ class ProxyLogging:
|
|||
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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -166,3 +166,123 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox
|
|||
)
|
||||
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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue