diff --git a/litellm/proxy/openai_files_endpoints/batch_guardrails.py b/litellm/proxy/openai_files_endpoints/batch_guardrails.py index 5c886ca0e9b..53d51db2b7f 100644 --- a/litellm/proxy/openai_files_endpoints/batch_guardrails.py +++ b/litellm/proxy/openai_files_endpoints/batch_guardrails.py @@ -260,18 +260,24 @@ def _describe(custom_id: str | None) -> str: return f" (custom_id {safe})" -def _iter_lines(source: BinaryIO) -> Iterator[tuple[int, str]]: - """Yield every non-blank line with its 1-based number, so both passes number records alike.""" +def _iter_lines(source: BinaryIO) -> Iterator[tuple[int, bytes]]: + """ + Yield every non-blank line with its 1-based number, so both passes number records alike. + + Bytes, not text. The upload validation immediately before this parses each line as bytes, + where the json module sniffs the encoding itself and accepts a leading byte order mark or a + lone surrogate. Decoding to `str` first is stricter than that, so a file written by any of + the editors that emit a BOM would pass validation and then fail the scan. + """ for line_number, raw_line in enumerate(source, start=1): - text = raw_line.decode("utf-8") - if text.strip(): - yield line_number, text + if raw_line.strip(): + yield line_number, raw_line def _iter_records(source: BinaryIO) -> Iterator[_ParsedRecord]: """Yield one record per line, relying on the upload validation that already ran.""" - for line_number, text in _iter_lines(source): - yield _ParsedRecord(line_number=line_number, payload=json.loads(text)) + for line_number, raw_line in _iter_lines(source): + yield _ParsedRecord(line_number=line_number, payload=json.loads(raw_line)) def _call_type_from_url(url: str) -> CallTypesLiteral | None: @@ -282,7 +288,13 @@ def _call_type_from_url(url: str) -> CallTypesLiteral | None: ``/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("/") + try: + path: Final = urlsplit(url).path.split("?")[0].rstrip("/") + except ValueError: + # urlsplit rejects a few malformed authorities outright, and the validation that ran + # before this only checks the key is present. An unreadable url is one we do not + # recognize, which is what falling back to the body shape already handles. + return None call_types: Final = get_call_types_for_route(path) if call_types is None: return None @@ -308,8 +320,18 @@ def _scannable_call_type(url: object, body: Mapping[str, object]) -> CallTypesLi def _custom_id_of(payload: Mapping[str, object]) -> str | None: + """ + The record's identifier, rendered as text. + + The batch spec asks for a string, but callers do send numbers, and reporting those as null + would leave the one field a caller reconciles on empty for exactly the records it needs. + """ custom_id: Final = payload.get("custom_id") - return custom_id if isinstance(custom_id, str) else None + if isinstance(custom_id, str): + # A lone surrogate parses out of the file but cannot be encoded back out, and this value + # is echoed in the response, so rendering it would fail the whole upload with a 500. + return custom_id.encode("utf-8", "replace").decode("utf-8") + return str(custom_id) if isinstance(custom_id, (int, float)) and not isinstance(custom_id, bool) else None def _fingerprint(body: Mapping[str, object], keys: frozenset[str]) -> str: @@ -507,9 +529,9 @@ async def scan_batch_input_file( ) -def _read_spooled(redactions: BinaryIO, change: RecordRedacted) -> str: +def _read_spooled(redactions: BinaryIO, change: RecordRedacted) -> bytes: redactions.seek(change.offset) - return redactions.read(change.length).decode("utf-8") + return redactions.read(change.length) def rewrite_batch_input_file(file_source: BinaryIO, result: BatchScanResult) -> BinaryIO: @@ -532,12 +554,12 @@ def rewrite_batch_input_file(file_source: BinaryIO, result: BatchScanResult) -> ) wrote_any = False # rebind-ok: tracks whether a separator is needed try: - for line_number, text in _iter_lines(file_source): + for line_number, raw_line in _iter_lines(file_source): if line_number in dropped: continue change = redacted.get(line_number) - line = text.rstrip("\n") if change is None else _read_spooled(result.redactions, change) - output.write((("\n" if wrote_any else "") + line).encode("utf-8")) + line = raw_line.rstrip(b"\n") if change is None else _read_spooled(result.redactions, change) + output.write(b"\n" + line if wrote_any else line) wrote_any = True except BaseException: output.close() diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 37cfd9d073d..813ce9630a5 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -145,25 +145,37 @@ async def _scan_batch_upload( def get_first_json_object(file_source: bytes | BinaryIO) -> dict | None: + """ + The first record, used to pick a deployment when batch load balancing is on. + + Read the way the upload validation reads it, since a file it accepted must not lose its + routing here: blank lines are not records and are skipped, and the line is parsed as bytes so + the json module sniffs the encoding rather than rejecting a leading byte order mark. Either + difference makes this return None, which silently sends the batch to the default provider. + """ try: if isinstance(file_source, (bytes, bytearray)): - newline: Final = file_source.find(b"\n") - raw: Final = file_source if newline == -1 else file_source[:newline] - first_line = raw.decode("utf-8") + first_record: bytes | None = next((line for line in file_source.splitlines() if line.strip()), None) else: + # lazily, so a batch file that can be gigabytes is not read past its first record file_source.seek(0) - first_line = file_source.readline().decode("utf-8") + first_record = next((line for line in file_source if line.strip()), None) file_source.seek(0) - return json.loads(first_line.strip()) + return None if first_record is None else json.loads(first_record.strip()) except (json.JSONDecodeError, UnicodeDecodeError, OSError, ValueError): return None def get_model_from_json_obj(json_object: dict) -> str | None: - body: Final = json_object.get("body", {}) or {} - model: Final = body.get("model") + """ + The model a record names, or None when it does not name one readably. - return model + The upload validation only checks that `body` is present, not that it is an object, so a + record can carry a string there and reach this. Returning None sends the upload down the + default-provider branch, which is what a record with no resolvable model already did. + """ + body: Final = json_object.get("body") + return body.get("model") if isinstance(body, dict) else None async def _deprecated_loadbalanced_create_file( diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py index 6ce7af1e2ee..a06b1110aa5 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py @@ -363,6 +363,121 @@ async def test_an_absolute_url_resolves_by_path_not_by_body_shape(url, expected_ assert logging_obj.seen[0][0] == expected_call_type +@pytest.mark.asyncio +@pytest.mark.parametrize( + "prefix, label", + [(b"\xef\xbb\xbf", "utf-8 BOM"), (b"", "plain")], + ids=["utf8_bom", "plain"], +) +async def test_a_file_the_upload_validation_accepts_is_a_file_the_scan_can_read(prefix, label): + """The validator parses each line as bytes, which tolerates a BOM; the scan must match it.""" + from litellm.proxy.openai_files_endpoints.batch_file_validation import check_batch_file_upload + + payload = prefix + (json.dumps(_record("a")) + "\n").encode() + assert check_batch_file_upload("in.jsonl", io.BytesIO(payload), None) is None, f"{label} rejected upfront" + + logging_obj = FakeProxyLogging() + assert await _scan(io.BytesIO(payload), logging_obj) is None + assert logging_obj.seen, f"{label} was never scanned" + + +@pytest.mark.asyncio +async def test_a_bom_file_is_rewritten_without_losing_the_untouched_records(): + source = io.BytesIO(b"\xef\xbb\xbf" + ("\n".join( + json.dumps(r) for r in (_record("keep"), _record("dirty", content="my secret is here")) + ) + "\n").encode()) + + result = await _scan_full(source, FakeProxyLogging(_redact_containing("secret"))) + rewritten = rewrite_batch_input_file(source, result).read().decode("utf-8-sig") + + rows = [json.loads(line) for line in rewritten.splitlines()] + assert [row["custom_id"] for row in rows] == ["keep", "dirty"] + assert rows[1]["body"]["messages"][0]["content"] == "my *** is here" + + +@pytest.mark.parametrize( + "prefix", + [b"", b"\xef\xbb\xbf", b"\n", b"\n\xef\xbb\xbf", b" \n"], + ids=["plain", "utf8_bom", "leading_blank", "blank_then_bom", "whitespace_line"], +) +def test_load_balancing_finds_the_routing_record_in_any_file_the_upload_accepts(prefix): + """A file whose routing model cannot be read is silently sent to the default provider.""" + from litellm.proxy.openai_files_endpoints.batch_file_validation import check_batch_file_upload + from litellm.proxy.openai_files_endpoints.files_endpoints import get_first_json_object + + payload = prefix + (json.dumps(_record("a")) + "\n").encode() + assert check_batch_file_upload("in.jsonl", io.BytesIO(payload), None) is None, "rejected upfront" + + assert get_first_json_object(io.BytesIO(payload))["body"]["model"] == "gpt-4o-mini" + assert get_first_json_object(payload)["body"]["model"] == "gpt-4o-mini" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("url", ["http://[", "http://[::1", "https://["], ids=["open_bracket", "unclosed_v6", "https_bracket"]) +async def test_a_malformed_url_does_not_escape_the_scan(url): + """Validation only checks the url key is present, and urlsplit rejects some authorities.""" + record = {**_record("m"), "url": url} + + result = await _scan_full(_jsonl(record), FakeProxyLogging()) + + assert result.changes == () + assert result.scanned_records == 1, "the record should still be scanned by its body shape" + + +@pytest.mark.parametrize( + "custom_id, expected", + [("req-1", "req-1"), ("caf\u00e9-42", "caf\u00e9-42"), ("a\ud800b", "a?b")], + ids=["ascii", "unicode", "lone_surrogate"], +) +def test_a_reported_custom_id_can_always_be_rendered(custom_id, expected): + """The id is echoed in the response; one that cannot be encoded back out would 500 the upload.""" + from litellm.proxy.openai_files_endpoints.batch_guardrails import _custom_id_of + + rendered = _custom_id_of({"custom_id": custom_id}) + + assert rendered == expected + assert json.dumps({"custom_id": rendered}, ensure_ascii=False).encode("utf-8") + + +@pytest.mark.parametrize( + "body", + ["summarize this", ["a"], None, 12345], + ids=["string", "list", "null", "number"], +) +def test_a_record_whose_body_is_not_an_object_does_not_crash_deployment_selection(body): + """Validation only checks that `body` is present, so a record can carry anything there.""" + from litellm.proxy.openai_files_endpoints.batch_file_validation import check_batch_file_upload + from litellm.proxy.openai_files_endpoints.files_endpoints import ( + get_first_json_object, + get_model_from_json_obj, + ) + + record = {"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": body} + payload = b"\xef\xbb\xbf" + (json.dumps(record) + "\n").encode() + assert check_batch_file_upload("in.jsonl", io.BytesIO(payload), None) is None, "rejected upfront" + + found = get_first_json_object(io.BytesIO(payload)) + assert get_model_from_json_obj(json_object=found) is None + + +@pytest.mark.parametrize("payload", [b"", b"\n\n\n"], ids=["empty", "blanks_only"]) +def test_load_balancing_returns_none_when_there_is_no_record(payload): + from litellm.proxy.openai_files_endpoints.files_endpoints import get_first_json_object + + assert get_first_json_object(io.BytesIO(payload)) is None + assert get_first_json_object(payload) is None + + +@pytest.mark.asyncio +async def test_a_numeric_custom_id_is_still_reported(): + """The spec asks for a string, but callers send numbers, and null would break reconciliation.""" + record = {**_record("x", content="tripwire"), "custom_id": 12345} + + result = await _scan_full(_jsonl(record), FakeProxyLogging(_blocking("tripwire"))) + + assert result.changes == (RecordDropped(line_number=1, custom_id="12345", guardrail="block-guard"),) + + @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."""