diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index cfafca4d174..06a37141b47 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -506,6 +506,19 @@ def _collect_ws_project_quota_callbacks() -> tuple[ProjectQuotaCallback, ...]: ) +def _unwrap_provider_error(error: Exception) -> Exception: + if not isinstance(error, httpx.TransportError): + return error + current: BaseException | None = error.__cause__ # rebind-ok: walks the bounded exception cause chain + for _ in range(10): + if current is None: + return error + if isinstance(current, BaseLLMException): + return current + current = current.__cause__ + return error + + class BaseLLMHTTPHandler: async def _make_common_async_call( self, @@ -6319,16 +6332,19 @@ class BaseLLMHTTPHandler: BaseRealtimeHTTPConfig, ], ): + error: Final = _unwrap_provider_error(e) received_status_code: Final = ( - e.response.status_code if isinstance(e, httpx.HTTPStatusError) else getattr(e, "status_code", None) + error.response.status_code + if isinstance(error, httpx.HTTPStatusError) + else getattr(error, "status_code", None) ) status_code = received_status_code if isinstance(received_status_code, int) else 500 - error_headers = getattr(e, "headers", None) - if isinstance(e, httpx.HTTPStatusError): - error_text = e.response.text + error_headers = getattr(error, "headers", None) + if isinstance(error, httpx.HTTPStatusError): + error_text = error.response.text else: - error_text = getattr(e, "text", str(e)) - error_response: Final = getattr(e, "response", None) + error_text = getattr(error, "text", str(error)) + error_response: Final = getattr(error, "response", None) if error_headers is None and error_response: error_headers = getattr(error_response, "headers", None) if error_response and hasattr(error_response, "text"): diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 7f3ef479ab5..794208cc2af 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -684,7 +684,9 @@ def _openai_batch_jsonl_entry_to_vertex_rows( if _is_embeddings_batch_entry(openai_entry): return _openai_batch_jsonl_entry_to_vertex_embeddings_rows(openai_entry) - openai_request_body: Final = openai_entry.get("body") or {} + if not isinstance(openai_entry.get("body"), dict): + raise TypeError("`body` on batch requests must be a JSON object, but was missing or not an object") + openai_request_body: Final = openai_entry["body"] chat_request_body: Final = ( responses_batch_body_to_chat_body(openai_request_body, custom_llm_provider="vertex_ai") if _is_responses_batch_entry(openai_entry) @@ -708,21 +710,40 @@ def _openai_batch_jsonl_entry_to_vertex_rows( return ({"request": vertex_request_body},) -def _iter_stripped_lines(raw_lines: Iterable[str | bytes]) -> Iterator[str]: - """Decode (when needed), strip, and drop blank lines from an iterable of lines.""" - for raw in raw_lines: - line = raw.decode("utf-8") if isinstance(raw, (bytes, bytearray)) else raw - line = line.strip() +def _decode_batch_input_line(lineno: int, raw: str | bytes | bytearray) -> str: + if isinstance(raw, str): + return raw.removeprefix("\ufeff") + try: + return raw.decode("utf-8-sig") + except UnicodeDecodeError as e: + raise VertexAIError( + status_code=400, + message=f"Batch input file line {lineno} is not valid UTF-8: {e.reason}", + ) from e + + +def _describe_batch_row_error(error: Exception) -> str: + match error: + case KeyError(): + return f"missing field {error}" + case _: + return str(error) + + +def _iter_numbered_stripped_lines(raw_lines: Iterable[str | bytes]) -> Iterator[tuple[int, str]]: + """Decode (when needed), strip trailing whitespace, and drop blank lines from an iterable of lines.""" + for lineno, raw in enumerate(raw_lines, start=1): + line = _decode_batch_input_line(lineno, raw).rstrip() if line: - yield line + yield lineno, line -def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: +def _iter_numbered_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[tuple[int, str]]: """ Yield non-empty JSONL lines one at a time without materializing the whole payload, so peak memory stays bounded regardless of payload size. Mirrors - ``str.splitlines()`` + ``line.strip()`` for ``\\n`` / ``\\r\\n`` delimited - JSONL. + ``str.splitlines()`` + ``line.rstrip()`` for ``\\n`` / ``\\r\\n`` delimited + JSONL, keeping leading whitespace so JSON error columns match the file. """ content: FileTypes | str = openai_file_content if isinstance(content, tuple): @@ -733,24 +754,26 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: # into a BytesIO just to iterate it line by line. newline: Final = ord("\n") start, length = 0, len(content) + lineno: int = 0 # rebind-ok: line counter advances for each streamed chunk while start < length: idx = content.find(newline, start) if idx == -1: chunk, start = content[start:], length else: chunk, start = content[start:idx], idx + 1 - line = chunk.decode("utf-8").strip() + lineno += 1 + line = _decode_batch_input_line(lineno, chunk).rstrip() if line: - yield line + yield lineno, line return if isinstance(content, str): - yield from _iter_stripped_lines(io.StringIO(content)) + yield from _iter_numbered_stripped_lines(io.StringIO(content)) return if isinstance(content, PathLike): with open(str(content), "rb") as handle: - yield from _iter_stripped_lines(handle) + yield from _iter_numbered_stripped_lines(handle) return if hasattr(content, "read"): @@ -771,7 +794,7 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: "Batch upload file handle must be seekable so it can be re-read " "for the GCS object name and the upload body." ) from e - yield from _iter_stripped_lines(content) + yield from _iter_numbered_stripped_lines(content) return raise ValueError("Unsupported file content type") @@ -779,9 +802,30 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: def _iter_openai_jsonl_entries( openai_file_content: FileTypes, -) -> Iterator[dict[str, object]]: - for line in _iter_openai_jsonl_lines(openai_file_content): - yield json.loads(line) +) -> Iterator[tuple[int, dict[str, object]]]: + for lineno, line in _iter_numbered_openai_jsonl_lines(openai_file_content): + entry = _load_openai_jsonl_row(lineno, line) + if not isinstance(entry, dict): + raise VertexAIError( + status_code=400, + message=f"Batch input file line {lineno} is not a JSON object", + ) + yield lineno, entry + + +def _load_openai_jsonl_row(lineno: int, line: str) -> object: + try: + return json.loads(line) + except json.JSONDecodeError as e: + raise VertexAIError( + status_code=400, + message=f"Batch input file line {lineno} is not valid JSON: {e.msg} at column {e.colno}", + ) from e + except (ValueError, RecursionError) as e: + raise VertexAIError( + status_code=400, + message=f"Batch input file line {lineno} could not be parsed: {e}", + ) from e def _parse_vertex_batch_output_row(line: str) -> _VertexBatchRow: @@ -807,8 +851,16 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): def _iter_vertex_jsonl_chunks(self) -> Iterator[bytes]: first = True - for entry in _iter_openai_jsonl_entries(self._openai_file_content): - for wrapped in _openai_batch_jsonl_entry_to_vertex_rows(entry, self._map_openai_to_vertex_params): + for lineno, entry in _iter_openai_jsonl_entries(self._openai_file_content): + try: + vertex_rows = _openai_batch_jsonl_entry_to_vertex_rows(entry, self._map_openai_to_vertex_params) + except Exception as e: + raise VertexAIError( + status_code=400, + message=f"Batch input file line {lineno} could not be converted to a Vertex AI request: " + f"{_describe_batch_row_error(e)}", + ) from e + for wrapped in vertex_rows: prefix = b"" if first else b"\n" first = False yield prefix + json.dumps(wrapped).encode("utf-8") @@ -921,11 +973,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): `projects/../locations/../endpoints/` batch job model instead of a nonexistent publisher model. """ - raw_model: Final = ( - deployment_model.removeprefix("vertex_ai/") - if deployment_model - else openai_jsonl_content[0].get("body", {}).get("model", "") - ) + first_entry: Final = openai_jsonl_content[0] + body_model: Final = first_entry["body"].get("model") if isinstance(first_entry.get("body"), dict) else None + jsonl_model: Final = body_model if isinstance(body_model, str) else "" + raw_model: Final = deployment_model.removeprefix("vertex_ai/") if deployment_model else jsonl_model return _managed_batch_object_name(raw_model, passthrough=False) def _get_passthrough_gcs_object_name(self, deployment_model: str | None) -> str: @@ -956,9 +1007,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): return self._get_passthrough_gcs_object_name(deployment_model) if purpose == "batch": ## 1. If jsonl, derive the object name from the deployment model (or the first entry's) - first_entry: Final = next(_iter_openai_jsonl_entries(file_data), None) - if first_entry is not None: - return self._get_gcs_object_name_from_batch_jsonl([first_entry], deployment_model=deployment_model) + first_numbered_entry: Final = next(_iter_openai_jsonl_entries(file_data), None) + if first_numbered_entry is not None: + return self._get_gcs_object_name_from_batch_jsonl( + [first_numbered_entry[1]], deployment_model=deployment_model + ) ## 2. If not jsonl, store under a server-generated managed object name filename, _ = extract_file_metadata(file_data) @@ -1406,7 +1459,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # as large as the (multi-GB) input, so splitting into a list of rows # and building a second list of transformed rows peaks at several full # copies and OOMs on retrieval. - lines: Final = _iter_openai_jsonl_lines(content) + lines: Final = (line for _, line in _iter_numbered_openai_jsonl_lines(content)) try: first_line: Final = next(lines) except StopIteration: diff --git a/litellm/proxy/openai_files_endpoints/batch_file_validation.py b/litellm/proxy/openai_files_endpoints/batch_file_validation.py index d49aaaeba65..1aa24eed4cd 100644 --- a/litellm/proxy/openai_files_endpoints/batch_file_validation.py +++ b/litellm/proxy/openai_files_endpoints/batch_file_validation.py @@ -96,7 +96,7 @@ def _iter_lines(file_source: bytes | BinaryIO) -> Iterator[bytes]: def _check_line(line_number: int, raw_line: bytes, line_shape: BatchLineShape) -> BatchFileValidationFailure | None: try: parsed: Final = json.loads(raw_line) - except (json.JSONDecodeError, UnicodeDecodeError): + except (ValueError, RecursionError): return BatchFileInvalidJsonLine(line_number=line_number) if not isinstance(parsed, dict): return BatchFileLineNotObject(line_number=line_number) diff --git a/litellm/router.py b/litellm/router.py index fb8dbe40eb8..99d5e515bdc 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6531,6 +6531,7 @@ class Router: file = replace_model_in_jsonl( file_content=file, new_model_name=stripped_model, + keep_line_numbers=custom_llm_provider == LlmProviders.VERTEX_AI.value, ) kwargs_copy["file"] = file diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index d0725addb30..5753d0d8db0 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -1,5 +1,6 @@ import io import json +from collections.abc import Iterable from os import PathLike from typing import Final @@ -73,7 +74,7 @@ def should_replace_model_in_jsonl( return purpose == "batch" and not passthrough -def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> FileTypes: +def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str, keep_line_numbers: bool = False) -> FileTypes: try: ## if pathlike, return the original file content if isinstance(file_content, PathLike): @@ -83,17 +84,17 @@ def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> File # spooled upload handle (managed batches stream from it) is read straight # off its backing; bytes/str are wrapped so they iterate line-by-line. source: Final = file_content[1] if isinstance(file_content, tuple) else file_content - if hasattr(source, "read"): + if isinstance(source, (bytes, bytearray)): + line_iter: Iterable[bytes] | Iterable[str] = io.BytesIO(bytes(source)) + elif isinstance(source, str): + line_iter = io.StringIO(source) + elif hasattr(source, "read") and not isinstance(source, PathLike): if hasattr(source, "seek"): try: source.seek(0) except (OSError, ValueError): pass - line_iter: object = source - elif isinstance(source, (bytes, bytearray)): - line_iter = io.BytesIO(bytes(source)) - elif isinstance(source, str): - line_iter = io.StringIO(source) + line_iter = source else: return file_content @@ -105,12 +106,20 @@ def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> File # multiple physical lines still parse. Streaming the handle also means # the model rewrite is actually applied to tuple-wrapped upload handles; # otherwise a restricted body.model would survive and bypass the batch - # model allowlist (which validates the upload target alias). + # model allowlist (which validates the upload target alias). With + # keep_line_numbers each row is written on the line it ends on, blank + # lines included, so a row error naming a line still points into the + # caller's file; otherwise rows go back to back, since OpenAI rejects a + # batch file whose first line is empty. output: Final = InMemoryFile(b"", name="modified_file.jsonl", content_type="application/jsonl") - wrote_any = False + last_written_line = 0 buffer = "" - for raw_line in line_iter: - buffer += raw_line.decode("utf-8") if isinstance(raw_line, (bytes, bytearray)) else raw_line + for line_number, raw_line in enumerate(line_iter, start=1): + buffer += ( + raw_line.decode("utf-8-sig") + if isinstance(raw_line, (bytes, bytearray)) + else raw_line.removeprefix("\ufeff") + ) stripped = buffer.strip() if not stripped: buffer = "" @@ -121,8 +130,13 @@ def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> File continue # object not complete yet; keep accumulating if isinstance(json_object, dict) and isinstance(json_object.get("body"), dict): json_object["body"]["model"] = new_model_name - output.write((("\n" if wrote_any else "") + json.dumps(json_object)).encode("utf-8")) - wrote_any = True + output.write( + ( + "\n" * (line_number - max(last_written_line, 1) if keep_line_numbers else min(last_written_line, 1)) + + json.dumps(json_object) + ).encode("utf-8") + ) + last_written_line = line_number buffer = "" if buffer.strip(): @@ -139,7 +153,7 @@ def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> File return file_content # If no valid JSON objects were found, return the original content - if not wrote_any: + if not last_written_line: return file_content output.seek(0) diff --git a/tests/integration/providers/_vertex_batch_rows.py b/tests/integration/providers/_vertex_batch_rows.py new file mode 100644 index 00000000000..9d4f9c57e75 --- /dev/null +++ b/tests/integration/providers/_vertex_batch_rows.py @@ -0,0 +1,174 @@ +import json +import socket +import threading +from collections.abc import Callable +from email.message import Message +from email.parser import BytesParser +from email.policy import HTTP +from typing import Final, Literal +from urllib.parse import unquote, urlsplit + +import httpx +from integration._support.client import Gateway, Scenario +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request +from pydantic import JsonValue, TypeAdapter + +PROJECT: Final = "scripted-project" +LOCATION: Final = "us-central1" +BUCKET: Final = "scripted-bucket" +UPLOAD_PREFIX: Final = f"/upload/storage/v1/b/{BUCKET}/o?uploadType=media&name=" +CHAT_URL: Final = "/v1/chat/completions" +BOM: Final = b"\xef\xbb\xbf" +BODY_NOT_OBJECT: Final = "`body` on batch requests must be a JSON object, but was missing or not an object" +Route = Literal["model", "target_model_names"] +ROUTES: Final[tuple[Route, ...]] = ("model", "target_model_names") +_ERROR_BODY: Final = TypeAdapter(dict[str, dict[str, JsonValue]]) + + +def stored_object(request: Request) -> Reply: + name: Final = unquote(request.target.removeprefix(UPLOAD_PREFIX)) + return Reply( + body=json.dumps( + { + "kind": "storage#object", + "id": f"{BUCKET}/{name}/1759950000000000", + "name": name, + "bucket": BUCKET, + "size": str(len(request.body)), + "timeCreated": "2026-10-15T06:00:00.000Z", + "contentType": "application/octet-stream", + } + ).encode() + ) + + +def is_gcs_upload(request: Request) -> bool: + return request.method == "POST" and request.target.startswith(UPLOAD_PREFIX) + + +def gcs_peer(request: Request) -> Reply: + assert is_gcs_upload(request), (request.method, request.target) + assert request.headers["authorization"] == "Bearer scripted-token" + return stored_object(request) + + +def held_gcs_peer(release: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert release.wait(timeout=120), "the held bucket upload was never released" + return gcs_peer(request) + + return respond + + +def vertex_params(gateway: Gateway, url: str, backend: str = "gemini-2.5-flash") -> dict[str, JsonValue]: + return { + "model": f"vertex_ai/{backend}", + "api_key": None, + "api_base": url, + "vertex_project": PROJECT, + "vertex_location": LOCATION, + "vertex_credentials": service_account_json(PROJECT, gateway.upstream_url), + "gcs_bucket_name": BUCKET, + } + + +def vertex_model(gateway: Gateway, scenario: Scenario, url: str, backend: str = "gemini-2.5-flash") -> str: + return scenario.model( + model=f"vertex_ai/{backend}", + api_key=None, + api_base=url, + vertex_project=PROJECT, + vertex_location=LOCATION, + vertex_credentials=service_account_json(PROJECT, gateway.upstream_url), + gcs_bucket_name=BUCKET, + ) + + +def chat_row(custom_id: str, model: str, text: str) -> dict[str, JsonValue]: + return { + "custom_id": custom_id, + "method": "POST", + "url": CHAT_URL, + "body": {"model": model, "messages": [{"role": "user", "content": text}], "max_tokens": 8}, + } + + +def roleless_row(custom_id: str, model: str) -> dict[str, JsonValue]: + return {"custom_id": custom_id, "method": "POST", "url": CHAT_URL, "body": {"model": model, "messages": [{}]}} + + +def encoded(row: dict[str, JsonValue]) -> bytes: + return json.dumps(row).encode() + + +def jsonl(*lines: bytes) -> bytes: + return b"\n".join(lines) + b"\n" + + +def upload( + gateway: Gateway, route: Route | None, model: str, content: bytes, *, client: httpx.Client | None = None, **extra: str +) -> httpx.Response: + fields: Final = {"purpose": "batch", **({route: model} if route is not None else {}), **extra} + return (client or gateway.client).post( + "/v1/files", + data=fields, + files={"file": ("rows.jsonl", content, "application/jsonl")}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) + + +def error_of(response: httpx.Response) -> dict[str, JsonValue]: + return _ERROR_BODY.validate_json(response.content)["error"] + + +def unconvertible_line_message(line: int) -> str: + return f"Batch input file line {line} could not be converted to a Vertex AI request: " + + +def stored_custom_ids(request: Request) -> tuple[str, ...]: + return tuple( + json.loads(line)["request"]["labels"]["litellm_custom_id"] for line in request.body.decode().split("\n") if line + ) + + +def closed_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.bind(("127.0.0.1", 0)) + return int(probe.getsockname()[1]) + + +def multipart_file(request: Request) -> bytes: + envelope: Final = f"content-type: {request.headers['content-type']}\r\n\r\n".encode() + request.body + parts: Final[tuple[Message, ...]] = tuple(BytesParser(policy=HTTP).parsebytes(envelope).iter_parts()) + files: Final = tuple(part.get_payload(decode=True) for part in parts if part.get_filename() is not None) + assert len(files) == 1, [part.get("content-disposition") for part in parts] + payload: Final = files[0] + assert isinstance(payload, bytes) + return payload + + +def is_openai_upload(request: Request) -> bool: + return request.method == "POST" and urlsplit(request.target).path == "/v1/files" + + +def openai_files_peer(request: Request) -> Reply: + if not is_openai_upload(request): + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + return Reply( + body=json.dumps( + { + "id": "file-scripted", + "object": "file", + "bytes": len(multipart_file(request)), + "created_at": 1792040000, + "filename": "modified_file.jsonl", + "purpose": "batch", + "status": "uploaded", + } + ).encode() + ) + + +def openai_uploaded_lines(request: Request) -> list[str]: + return multipart_file(request).decode().split("\n") diff --git a/tests/integration/providers/test_vertex_batch_row_line_errors_chaos.py b/tests/integration/providers/test_vertex_batch_row_line_errors_chaos.py new file mode 100644 index 00000000000..12e58ec6e41 --- /dev/null +++ b/tests/integration/providers/test_vertex_batch_row_line_errors_chaos.py @@ -0,0 +1,166 @@ +import threading +import uuid +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass +from itertools import chain +from typing import Final + +import httpx +from integration._support.client import Gateway, eventually +from integration._support.wire import Request, wire_server + +from integration.providers._vertex_batch_rows import ( + ROUTES, + Route, + chat_row, + closed_port, + encoded, + error_of, + gcs_peer, + held_gcs_peer, + jsonl, + roleless_row, + stored_custom_ids, + unconvertible_line_message, + upload, + vertex_model, +) + + +@dataclass(frozen=True, slots=True) +class _Upload: + route: Route + bad_line: int | None + identity: str + + +@dataclass(frozen=True, slots=True) +class _Answer: + sent: _Upload + response: httpx.Response + + +def _uploads(count: int, *, bad: bool) -> tuple[_Upload, ...]: + return tuple( + _Upload(ROUTES[index % 2], 2 + index if bad else None, f"{'bad' if bad else 'good'}-{uuid.uuid4().hex}") + for index in range(count) + ) + + +def _content(item: _Upload, model: str) -> bytes: + if item.bad_line is None: + return jsonl(encoded(chat_row(item.identity, model, "kept"))) + leading: Final = tuple( + encoded(chat_row(f"{item.identity}-{index}", model, "fine")) for index in range(1, item.bad_line) + ) + return jsonl(*leading, encoded(roleless_row(item.identity, model))) + + +def _burst_client(gateway: Gateway) -> httpx.Client: + return httpx.Client( + base_url=str(gateway.client.base_url), + timeout=180, + trust_env=False, + limits=httpx.Limits(max_connections=64, max_keepalive_connections=0), + ) + + +def _send(gateway: Gateway, client: httpx.Client, model: str, item: _Upload) -> _Answer: + return _Answer(item, upload(gateway, item.route, model, _content(item, model), client=client)) + + +def _submit( + pool: ThreadPoolExecutor, gateway: Gateway, client: httpx.Client, model: str, items: tuple[_Upload, ...] +) -> tuple[Future[_Answer], ...]: + return tuple(pool.submit(_send, gateway, client, model, item) for item in items) + + +def _burst(gateway: Gateway, client: httpx.Client, model: str, items: tuple[_Upload, ...]) -> tuple[_Answer, ...]: + with ThreadPoolExecutor(max_workers=len(items)) as pool: + return tuple(future.result(timeout=180) for future in _submit(pool, gateway, client, model, items)) + + +def _landed(requests: tuple[Request, ...]) -> tuple[str, ...]: + return tuple(sorted(chain.from_iterable(stored_custom_ids(request) for request in requests))) + + +def _assert_names_its_own_line(answer: _Answer) -> None: + assert answer.sent.bad_line is not None + assert answer.response.status_code == 400, (answer.sent, answer.response.text) + message: Final = str(error_of(answer.response)["message"]) + assert message.startswith(unconvertible_line_message(answer.sent.bad_line) + "missing field 'role'"), ( + answer.sent, + message, + ) + + +def _assert_each_bad_file_names_its_own_line(answers: tuple[_Answer, ...]) -> None: + for answer in answers: + _assert_names_its_own_line(answer) + + +def _assert_each_good_file_landed_once(answers: tuple[_Answer, ...], requests: tuple[Request, ...]) -> None: + assert all(answer.response.status_code == 200 for answer in answers), [a.response.text for a in answers] + assert _landed(requests) == tuple(sorted(answer.sent.identity for answer in answers)) + + +def _mixed(gateway: Gateway, client: httpx.Client, model: str, bad: int, good: int) -> tuple[_Answer, ...]: + return _burst(gateway, client, model, _uploads(bad, bad=True) + _uploads(good, bad=False)) + + +def _split(answers: tuple[_Answer, ...]) -> tuple[tuple[_Answer, ...], tuple[_Answer, ...]]: + return ( + tuple(answer for answer in answers if answer.sent.bad_line is not None), + tuple(answer for answer in answers if answer.sent.bad_line is None), + ) + + +def test_concurrent_mixed_burst_names_each_bad_files_own_line_and_lands_each_good_file_once(gateway: Gateway) -> None: + with wire_server(gcs_peer) as wire, gateway.scenario() as scenario, _burst_client(gateway) as client: + model: Final = vertex_model(gateway, scenario, wire.url) + bad, good = _split(_mixed(gateway, client, model, 12, 12)) + _assert_each_bad_file_names_its_own_line(bad) + _assert_each_good_file_landed_once(good, wire.drain()) + + +def test_bucket_outage_answers_500_without_blaming_a_row_and_recovers_on_restart(gateway: Gateway) -> None: + port: Final = closed_port() + with gateway.scenario() as scenario, _burst_client(gateway) as client: + model: Final = vertex_model(gateway, scenario, f"http://127.0.0.1:{port}") + with wire_server(gcs_peer, port=port) as before: + bad_before, good_before = _split(_mixed(gateway, client, model, 6, 6)) + _assert_each_bad_file_names_its_own_line(bad_before) + _assert_each_good_file_landed_once(good_before, before.drain()) + during: Final = _mixed(gateway, client, model, 6, 6) + assert all(answer.response.status_code == 500 for answer in during), [a.response.text for a in during] + assert not any("Batch input file line" in answer.response.text for answer in during) + assert gateway.client.get("/health/liveliness").status_code == 200 + with wire_server(gcs_peer, port=port) as after: + bad_after, good_after = _split(_mixed(gateway, client, model, 6, 6)) + _assert_each_bad_file_names_its_own_line(bad_after) + _assert_each_good_file_landed_once(good_after, after.drain()) + + +def test_slow_bucket_holds_good_uploads_without_delaying_bad_rows_400s(gateway: Gateway) -> None: + release: Final = threading.Event() + good_items: Final = _uploads(6, bad=False) + with ( + wire_server(held_gcs_peer(release)) as wire, + gateway.scenario() as scenario, + _burst_client(gateway) as client, + ThreadPoolExecutor(max_workers=12) as pool, + ): + try: + model: Final = vertex_model(gateway, scenario, wire.url) + held: Final = _submit(pool, gateway, client, model, good_items) + eventually(wire.received.qsize, lambda size: size == len(good_items), seconds=60) + bad: Final = tuple( + future.result(timeout=60) + for future in _submit(pool, gateway, client, model, _uploads(6, bad=True)) + ) + _assert_each_bad_file_names_its_own_line(bad) + assert not any(future.done() for future in held) + release.set() + _assert_each_good_file_landed_once(tuple(future.result(timeout=120) for future in held), wire.drain()) + finally: + release.set() diff --git a/tests/integration/providers/test_vertex_batch_row_line_errors_owned_proxy.py b/tests/integration/providers/test_vertex_batch_row_line_errors_owned_proxy.py new file mode 100644 index 00000000000..f5df41c16b0 --- /dev/null +++ b/tests/integration/providers/test_vertex_batch_row_line_errors_owned_proxy.py @@ -0,0 +1,288 @@ +import json +import threading +from itertools import chain +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from functools import partial +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Request, wire_server +from pydantic import JsonValue, TypeAdapter + +from integration.providers._vertex_batch_rows import ( + Route, + chat_row, + closed_port, + encoded, + error_of, + gcs_peer, + held_gcs_peer, + is_openai_upload, + jsonl, + openai_files_peer, + openai_uploaded_lines, + roleless_row, + stored_custom_ids, + unconvertible_line_message, + upload, + vertex_params, +) + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +VERTEX: Final = "lb-vertex" +OPENAI: Final = "lb-openai" +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +UploadRoute = Route | Literal["load_balanced"] +UPLOAD_ROUTES: Final[tuple[UploadRoute, ...]] = ("load_balanced", "model", "target_model_names") + + +@dataclass(frozen=True, slots=True) +class _Rig: + proxy: OwnedProxy + vertex_port: int + openai_port: int + + @property + def gateway(self) -> Gateway: + return self.proxy.gateway + + +def _config(gateway: Gateway, directory: Path, vertex_port: int, openai_port: int) -> Path: + base: Final = _JSON.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final[dict[str, JsonValue]] = { + **base, + "model_list": [ + {"model_name": VERTEX, "litellm_params": vertex_params(gateway, f"http://127.0.0.1:{vertex_port}")}, + { + "model_name": OPENAI, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "integration-provider-key", + "api_base": f"http://127.0.0.1:{openai_port}/v1", + }, + }, + ], + "litellm_settings": { + **_JSON.validate_python(base["litellm_settings"]), + "enable_loadbalancing_on_batch_endpoints": True, + }, + "guardrails": [ + { + "guardrail_name": "tripwire-filter", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "default_on": True, + "blocked_words": [{"keyword": "tripwire", "action": "BLOCK"}], + }, + } + ], + } + path: Final = directory / "vertex-batch-row-lines.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("vertex-batch-row-lines") + vertex_port: Final = closed_port() + openai_port: Final = closed_port() + with gateway_from_environment() as shared: + config: Final = _config(shared, directory, vertex_port, openai_port) + with owned_proxy_process(shared, directory, {}, config=config, workers=2) as owned: + yield _Rig(owned, vertex_port, openai_port) + + +def _upload(gateway: Gateway, route: UploadRoute, model: str, content: bytes) -> httpx.Response: + return upload(gateway, None if route == "load_balanced" else route, model, content) + + +def _tripwire_row(custom_id: str, model: str) -> bytes: + return encoded(chat_row(custom_id, model, "this row trips the tripwire filter")) + + +LAYOUTS: Final[tuple[tuple[str, int], ...]] = (("dropped_row_first", 1), ("dropped_row_second", 2)) + + +def _layout(name: str, model: str) -> bytes: + roleless: Final = encoded(roleless_row("bad", model)) + match name: + case "dropped_row_first": + return jsonl(_tripwire_row("dropped", model), b"", roleless) + case "dropped_row_second": + return jsonl(encoded(chat_row("first", model, "one")), _tripwire_row("dropped", model), b"", roleless) + case _: + raise AssertionError(name) + + +@pytest.mark.parametrize("route", UPLOAD_ROUTES) +def test_unconvertible_row_after_a_blank_line_names_its_file_line(rig: _Rig, route: UploadRoute) -> None: + content: Final = jsonl( + encoded(chat_row("first", VERTEX, "one")), + b"", + encoded(chat_row("third", VERTEX, "three")), + encoded(roleless_row("bad", VERTEX)), + ) + with wire_server(gcs_peer, port=rig.vertex_port) as wire: + response: Final = _upload(rig.gateway, route, VERTEX, content) + assert response.status_code == 400, response.text + error: Final = error_of(response) + assert str(error["message"]).startswith(unconvertible_line_message(4) + "missing field 'role'"), error + assert error["type"] == "invalid_request_error", error + assert wire.drain() == () + + +@pytest.mark.parametrize(("layout", "bad_line"), LAYOUTS, ids=tuple(name for name, _ in LAYOUTS)) +@pytest.mark.parametrize("route", UPLOAD_ROUTES) +def test_row_after_a_guardrail_dropped_row_names_its_line_in_the_uploaded_file( + rig: _Rig, route: UploadRoute, layout: str, bad_line: int +) -> None: + with wire_server(gcs_peer, port=rig.vertex_port) as wire: + response: Final = _upload(rig.gateway, route, VERTEX, _layout(layout, VERTEX)) + assert response.status_code == 400, response.text + message: Final = str(error_of(response)["message"]) + assert message.startswith(unconvertible_line_message(bad_line) + "missing field 'role'"), message + assert wire.drain() == () + + +def _kept_and_dropped(model: str) -> bytes: + return jsonl( + encoded(chat_row("first", model, "one")), + _tripwire_row("dropped", model), + b"", + encoded(chat_row("fourth", model, "four")), + ) + + +def _dropped_first(model: str) -> bytes: + return jsonl( + _tripwire_row("dropped", model), + b"", + encoded(chat_row("first", model, "one")), + encoded(chat_row("fourth", model, "four")), + ) + + +@pytest.mark.parametrize("content", (_kept_and_dropped, _dropped_first), ids=("dropped_row_second", "dropped_row_first")) +@pytest.mark.parametrize("route", UPLOAD_ROUTES) +def test_guardrail_dropped_and_blank_rows_never_reach_openai( + rig: _Rig, route: UploadRoute, content: Callable[[str], bytes] +) -> None: + with wire_server(openai_files_peer, port=rig.openai_port) as wire: + response: Final = _upload(rig.gateway, route, OPENAI, content(OPENAI)) + assert response.status_code == 200, response.text + uploads: Final = tuple(request for request in wire.drain() if is_openai_upload(request)) + assert len(uploads) == 1, uploads + sent_model: Final = OPENAI if route == "model" else "gpt-4o-mini" + rows: Final = tuple(json.loads(line) for line in openai_uploaded_lines(uploads[0])) + assert [(row["custom_id"], row["body"]["model"]) for row in rows] == [ + ("first", sent_model), + ("fourth", sent_model), + ] + + +@pytest.mark.parametrize("route", UPLOAD_ROUTES) +def test_guardrail_dropped_row_never_reaches_the_bucket(rig: _Rig, route: UploadRoute) -> None: + with wire_server(gcs_peer, port=rig.vertex_port) as wire: + response: Final = _upload(rig.gateway, route, VERTEX, _kept_and_dropped(VERTEX)) + assert response.status_code == 200, response.text + uploads: Final = wire.drain() + assert len(uploads) == 1, uploads + assert stored_custom_ids(uploads[0]) == ("first", "fourth") + + +def _bad_file(line: int) -> bytes: + good: Final = tuple(encoded(chat_row(f"row-{line}-{index}", VERTEX, "fine")) for index in range(1, line)) + return jsonl(*good, encoded(roleless_row(f"bad-{line}", VERTEX))) + + +def _guarded_upload(client: httpx.Client, gateway: Gateway, content: bytes) -> httpx.Response | None: + try: + return upload(gateway, "model", VERTEX, content, client=client) + except httpx.TransportError: + return None + + +def _landed_ids(requests: tuple[Request, ...]) -> tuple[str, ...]: + return tuple(chain.from_iterable(stored_custom_ids(request) for request in requests)) + + +def _held_connections(pid: int, url: str) -> int: + port: Final = urlsplit(url).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +def _upload_bad_file(client: httpx.Client, gateway: Gateway, line: int) -> httpx.Response | None: + return _guarded_upload(client, gateway, _bad_file(line)) + + +def _assert_names_its_line(response: httpx.Response | None, line: int) -> None: + assert response is not None + assert response.status_code == 400, response.text + assert str(error_of(response)["message"]).startswith(unconvertible_line_message(line)), response.text + + +@pytest.mark.timeout(4 * graceful_stop_seconds() + 240) +def test_worker_killed_mid_upload_leaves_the_sibling_naming_lines_and_landing_each_file_once(rig: _Rig) -> None: + bad_lines: Final = tuple(range(2, 10)) + held_ids: Final = tuple(f"held-{index}" for index in range(8)) + release: Final = threading.Event() + with ( + httpx.Client( + base_url=str(rig.gateway.client.base_url), + timeout=180, + trust_env=False, + limits=httpx.Limits(max_keepalive_connections=0), + ) as client, + ThreadPoolExecutor(max_workers=16) as pool, + ): + with wire_server(gcs_peer, port=rig.vertex_port) as wire: + answers: Final = tuple( + pool.map(partial(_upload_bad_file, client, rig.gateway), bad_lines) + ) + for line, answer in zip(bad_lines, answers, strict=True): + _assert_names_its_line(answer, line) + assert wire.drain() == () + with wire_server(held_gcs_peer(release), port=rig.vertex_port) as held: + try: + workers: Final = eventually(rig.proxy.worker_pids, lambda pids: len(pids) == 2, seconds=30) + pending: Final = tuple( + pool.submit(_guarded_upload, client, rig.gateway, jsonl(encoded(chat_row(identity, VERTEX, "x")))) + for identity in held_ids + ) + eventually(held.received.qsize, lambda size: size == len(held_ids), seconds=60) + held_by: Final = MappingProxyType({pid: _held_connections(pid, held.url) for pid in workers}) + assert sum(held_by.values()) == len(held_ids), held_by + survivor_pid, victim_pid = sorted(workers, key=held_by.__getitem__) + with rig.proxy.replacing_worker(victim_pid) as kill: + kill.freeze() + kill() + release.set() + outcomes: Final = tuple(future.result(timeout=120) for future in pending) + served: Final = tuple(outcome for outcome in outcomes if outcome is not None) + assert len(served) == held_by[survivor_pid], (held_by, [outcome and outcome.text for outcome in outcomes]) + assert all(outcome.status_code == 200 for outcome in served), [outcome.text for outcome in served] + landed: Final = _landed_ids(held.drain()) + assert sorted(landed) == sorted(held_ids) + _assert_names_its_line(_guarded_upload(client, rig.gateway, _bad_file(5)), 5) + after: Final = _guarded_upload(client, rig.gateway, jsonl(encoded(chat_row("after", VERTEX, "x")))) + assert after is not None and after.status_code == 200, after and after.text + assert _landed_ids(held.drain()) == ("after",) + finally: + release.set() diff --git a/tests/integration/providers/test_vertex_batch_row_line_errors_wire.py b/tests/integration/providers/test_vertex_batch_row_line_errors_wire.py new file mode 100644 index 00000000000..e47dd42413d --- /dev/null +++ b/tests/integration/providers/test_vertex_batch_row_line_errors_wire.py @@ -0,0 +1,405 @@ +import json +from collections.abc import Callable, Iterable +from typing import Final + +import anthropic +import openai +import pytest +from integration._support.client import Gateway, Scenario +from integration._support.wire import Reply, wire_server +from pydantic import JsonValue + +from integration.providers._vertex_batch_rows import ( + BOM, + BODY_NOT_OBJECT, + CHAT_URL, + ROUTES, + Route, + chat_row, + closed_port, + encoded, + error_of, + gcs_peer, + is_openai_upload, + jsonl, + openai_files_peer, + openai_uploaded_lines, + roleless_row, + stored_custom_ids, + unconvertible_line_message, + upload, + vertex_model, +) + + +def _with_body(body: JsonValue) -> bytes: + return encoded({"custom_id": "bad", "method": "POST", "url": CHAT_URL, "body": body}) + + +def _row_case(case: str, model: str) -> tuple[bytes, str]: + valid_body: Final = json.dumps(chat_row("bad", model, "fine")["body"]) + match case: + case "body_null": + return _with_body(None), BODY_NOT_OBJECT + case "body_int": + return _with_body(7), BODY_NOT_OBJECT + case "body_list": + return _with_body([{"model": model}]), BODY_NOT_OBJECT + case "body_empty_string": + return _with_body(""), BODY_NOT_OBJECT + case "body_5kb_string": + return _with_body("x" * 5120), BODY_NOT_OBJECT + case "body_twice_null_last": + twice: Final = f'{{"custom_id": "bad", "method": "POST", "url": "{CHAT_URL}", "body": {valid_body}, "body": null}}' + return twice.encode(), BODY_NOT_OBJECT + case "message_without_role": + return encoded(roleless_row("bad", model)), "missing field 'role'" + case "responses_body_null": + return encoded({"custom_id": "bad", "method": "POST", "url": "/v1/responses", "body": None}), BODY_NOT_OBJECT + case _: + raise AssertionError(case) + + +ROW_CASES: Final = ( + "body_null", + "body_int", + "body_list", + "body_empty_string", + "body_5kb_string", + "body_twice_null_last", + "message_without_role", + "responses_body_null", +) + + +@pytest.mark.parametrize("case", ROW_CASES) +@pytest.mark.parametrize("route", ROUTES) +def test_row_vertex_cannot_convert_answers_400_naming_its_file_line_every_time( + gateway: Gateway, route: Route, case: str +) -> None: + with wire_server(gcs_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url) + bad_line, detail = _row_case(case, model) + content: Final = jsonl( + encoded(chat_row("first", model, "one")), b"", bad_line, encoded(chat_row("last", model, "two")) + ) + responses: Final = tuple(upload(gateway, route, model, content) for _ in range(2)) + assert tuple(response.status_code for response in responses) == (400, 400), [r.text for r in responses] + errors: Final = tuple(error_of(response) for response in responses) + assert tuple(error["type"] for error in errors) == ("invalid_request_error",) * 2, errors + assert str(errors[0]["message"]).startswith(unconvertible_line_message(3) + detail), errors[0] + assert errors[0]["message"] == errors[1]["message"] + assert wire.drain() == (), "a batch file with an unconvertible row reached the bucket" + + +def test_embeddings_row_vertex_cannot_convert_names_its_file_line(gateway: Gateway) -> None: + block: Final[dict[str, JsonValue]] = {"type": "file", "file": {"file_id": "gs://scripted-bucket/clip.mp4"}} + with wire_server(gcs_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url, backend="gemini-embedding-2-preview") + fine: Final = encoded( + {"custom_id": "fine", "method": "POST", "url": "/v1/embeddings", "body": {"model": model, "input": "bus"}} + ) + bare: Final = encoded( + {"custom_id": "bare", "method": "POST", "url": "/v1/embeddings", "body": {"model": model, "input": block}} + ) + response: Final = upload(gateway, "target_model_names", model, jsonl(fine, b" ", b"", bare)) + assert response.status_code == 400, response.text + assert str(error_of(response)["message"]).startswith(unconvertible_line_message(4)), response.text + assert wire.drain() == () + + +def _clean_and_messy(model: str) -> tuple[bytes, bytes]: + first: Final = encoded(chat_row("first", model, "one")) + second: Final = encoded(chat_row("second", model, "two")) + third: Final = encoded(chat_row("third", model, "three")) + messy: Final = b"".join( + (BOM + first + b"\r\n", b"\r\n", b" \t\r\n", BOM + second + b" \r\n", b" " + third + b"\n", b"\n") + ) + return jsonl(first, second, third), messy + + +@pytest.mark.parametrize("route", ROUTES) +def test_byte_order_marks_blank_lines_and_crlf_upload_the_same_rows_as_a_clean_file( + gateway: Gateway, route: Route +) -> None: + with wire_server(gcs_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url) + clean, messy = _clean_and_messy(model) + clean_response: Final = upload(gateway, route, model, clean) + assert clean_response.status_code == 200, clean_response.text + clean_uploads: Final = wire.drain() + messy_response: Final = upload(gateway, route, model, messy) + assert messy_response.status_code == 200, messy_response.text + messy_uploads: Final = wire.drain() + assert (len(clean_uploads), len(messy_uploads)) == (1, 1), (clean_uploads, messy_uploads) + assert messy_uploads[0].body == clean_uploads[0].body + assert stored_custom_ids(messy_uploads[0]) == ("first", "second", "third") + stored_texts: Final = tuple( + json.loads(line)["request"]["contents"][0]["parts"][0]["text"] + for line in messy_uploads[0].body.decode().split("\n") + ) + assert stored_texts == ("one", "two", "three") + + +def _deeply_nested(model: str) -> bytes: + depth: Final = 100_000 + return ( + b'{"custom_id": "deep", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "' + + model.encode() + + b'", "messages": ' + + b"[" * depth + + b"]" * depth + + b"}}" + ) + + +def _long_integer(model: str) -> bytes: + body: Final = json.dumps(chat_row("long", model, "x")["body"]).encode() + return b'{"custom_id": "long", "method": "POST", "url": "/v1/chat/completions", "body": ' + ( + body[:-1] + b', "seed": ' + b"7" * 5000 + b"}}" + ) + + +@pytest.mark.parametrize("row", (_deeply_nested, _long_integer), ids=("nested_100k_deep", "integer_5000_digits")) +def test_row_too_deep_or_long_to_parse_answers_400_naming_its_line( + gateway: Gateway, row: Callable[[str], bytes] +) -> None: + with wire_server(gcs_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url) + response: Final = upload( + gateway, "model", model, jsonl(encoded(chat_row("first", model, "one")), b"", row(model)) + ) + assert response.status_code == 400, response.text + assert error_of(response) == { + "message": "Batch input file line 3 is not valid JSON. The file was not forwarded to the provider.", + "type": "invalid_request_error", + "param": "file", + "code": "400", + }, response.text + assert wire.drain() == () + assert gateway.client.get("/health/liveliness").status_code == 200 + follow_up: Final = upload(gateway, "model", model, jsonl(encoded(chat_row("after", model, "ok")))) + assert follow_up.status_code == 200, follow_up.text + assert len(wire.drain()) == 1 + + +VALIDATION_CASES: Final[tuple[tuple[str, bytes, str, str], ...]] = ( + ("truncated_json", b'{"custom_id": "bad", ', "Batch input file line 3 is not valid JSON.", "file"), + ("invalid_utf8", b'{"custom_id": "\xff\xfe"}', "Batch input file line 3 is not valid JSON.", "file"), + ("json_array", b'[{"custom_id": "bad"}]', "Batch input file line 3 must be a JSON object.", "file"), + ( + "body_missing", + b'{"custom_id": "bad", "method": "POST", "url": "/v1/chat/completions"}', + "Missing required parameter: 'body' (batch input file line 3).", + "body", + ), +) + + +@pytest.mark.parametrize( + ("bad_line", "message", "param"), + tuple(case[1:] for case in VALIDATION_CASES), + ids=tuple(case[0] for case in VALIDATION_CASES), +) +@pytest.mark.parametrize("route", ROUTES) +def test_rows_the_upload_validation_rejects_keep_their_400( + gateway: Gateway, route: Route, bad_line: bytes, message: str, param: str +) -> None: + with wire_server(gcs_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url) + response: Final = upload(gateway, route, model, jsonl(encoded(chat_row("first", model, "one")), b"", bad_line)) + assert response.status_code == 400, response.text + error: Final = error_of(response) + assert str(error["message"]).startswith(message), response.text + assert str(error["message"]).endswith("The file was not forwarded to the provider."), response.text + assert (error["type"], error["param"]) == ("invalid_request_error", param), response.text + assert wire.drain() == () + + +@pytest.mark.parametrize("route", ROUTES) +def test_empty_object_body_still_uploads(gateway: Gateway, route: Route) -> None: + with wire_server(gcs_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url) + empty_body: Final = encoded({"custom_id": "empty", "method": "POST", "url": CHAT_URL, "body": {}}) + response: Final = upload(gateway, route, model, jsonl(encoded(chat_row("first", model, "one")), empty_body)) + assert response.status_code == 200, response.text + uploads: Final = wire.drain() + assert len(uploads) == 1 + assert stored_custom_ids(uploads[0]) == ("first", "empty") + + +def test_passthrough_upload_stores_the_callers_bytes_unchanged(gateway: Gateway) -> None: + content: Final = b"".join( + ( + BOM + b'{"request": {"contents": [{"role": "user", "parts": [{"text": "one"}]}]}}\r\n', + b"\r\n", + b' {"request": {"contents": [{"role": "user", "parts": [{"text": "two"}]}]}} \n', + ) + ) + with wire_server(gcs_peer) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url) + response: Final = upload(gateway, "target_model_names", model, content, passthrough="true") + assert response.status_code == 200, response.text + uploads: Final = wire.drain() + assert len(uploads) == 1 + assert uploads[0].body == content + + +GCS_FAILURES: Final[tuple[tuple[str, Reply | None, int], ...]] = ( + ( + "bucket_answers_403", + Reply(status=403, body=json.dumps({"error": {"code": 403, "message": "scripted denial"}}).encode()), + 403, + ), + ("bucket_drops_the_connection", Reply(drop_connection=True), 500), + ("bucket_port_closed", None, 500), +) + + +@pytest.mark.parametrize( + ("reply", "status"), tuple(case[1:] for case in GCS_FAILURES), ids=tuple(case[0] for case in GCS_FAILURES) +) +@pytest.mark.parametrize("route", ROUTES) +def test_bucket_failures_keep_their_status_and_never_blame_a_row( + gateway: Gateway, route: Route, reply: Reply | None, status: int +) -> None: + with wire_server(lambda _request: reply or Reply()) as wire, gateway.scenario() as scenario: + url: Final = wire.url if reply is not None else f"http://127.0.0.1:{closed_port()}" + model: Final = vertex_model(gateway, scenario, url) + response: Final = upload(gateway, route, model, jsonl(encoded(chat_row("first", model, "one")))) + assert response.status_code == status, response.text + assert "Batch input file line" not in response.text, response.text + + +def _bom_then_blank_lines(first: bytes, fourth: bytes) -> bytes: + return b"".join((BOM + first + b"\n", b"\n", b" \n", fourth + b"\n")) + + +def _leading_blank_line(first: bytes, fourth: bytes) -> bytes: + return b"".join((b"\n", first + b"\n", fourth + b"\n")) + + +@pytest.mark.parametrize("layout", (_bom_then_blank_lines, _leading_blank_line), ids=("bom_then_blank", "leading_blank")) +def test_router_rewrite_sends_openai_only_the_rows(gateway: Gateway, layout: Callable[[bytes, bytes], bytes]) -> None: + with wire_server(openai_files_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4o-mini", api_key="integration-provider-key", api_base=f"{wire.url}/v1" + ) + content: Final = layout(encoded(chat_row("first", model, "one")), encoded(chat_row("fourth", model, "four"))) + response: Final = upload(gateway, "target_model_names", model, content) + assert response.status_code == 200, response.text + uploads: Final = tuple(request for request in wire.drain() if is_openai_upload(request)) + assert len(uploads) == 1, [request.target for request in uploads] + rows: Final = tuple(json.loads(line) for line in openai_uploaded_lines(uploads[0])) + assert [(row["custom_id"], row["body"]["model"]) for row in rows] == [ + ("first", "gpt-4o-mini"), + ("fourth", "gpt-4o-mini"), + ] + + +def test_one_upload_to_openai_then_vertex_sends_openai_only_the_rows_and_names_the_vertex_file_line( + gateway: Gateway, +) -> None: + with ( + wire_server(openai_files_peer) as openai_wire, + wire_server(gcs_peer) as gcs_wire, + gateway.scenario() as scenario, + ): + openai_target: Final = scenario.model( + model="openai/gpt-4o-mini", api_key="integration-provider-key", api_base=f"{openai_wire.url}/v1" + ) + vertex_target: Final = vertex_model(gateway, scenario, gcs_wire.url) + content: Final = jsonl( + b"", encoded(chat_row("first", openai_target, "one")), b"", encoded(roleless_row("bad", openai_target)) + ) + response: Final = upload(gateway, "target_model_names", f"{openai_target},{vertex_target}", content) + assert response.status_code == 400, response.text + message: Final = str(error_of(response)["message"]) + assert message.startswith(unconvertible_line_message(4) + "missing field 'role'"), message + uploads: Final = tuple(request for request in openai_wire.drain() if is_openai_upload(request)) + assert len(uploads) == 1, [request.target for request in uploads] + rows: Final = tuple(json.loads(line) for line in openai_uploaded_lines(uploads[0])) + assert [(row["custom_id"], row["body"]["model"]) for row in rows] == [ + ("first", "gpt-4o-mini"), + ("bad", "gpt-4o-mini"), + ] + assert gcs_wire.drain() == () + + +def _unreachable_anthropic_model(scenario: Scenario) -> str: + return scenario.model( + model="anthropic/claude-haiku-4-5", + api_key="synthetic-anthropic-key", + api_base=f"http://127.0.0.1:{closed_port()}", + num_retries=0, + ) + + +def _proxy_url(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def _consume(reply: object, stream: bool) -> None: + if stream and isinstance(reply, Iterable): + tuple(reply) + + +def _chat(client: openai.OpenAI, model: str, stream: bool) -> None: + _consume( + client.chat.completions.create( + model=model, messages=[{"role": "user", "content": "hello"}], stream=stream, max_tokens=8 + ), + stream, + ) + + +def _messages(client: anthropic.Anthropic, model: str, stream: bool) -> None: + _consume( + client.messages.create( + model=model, messages=[{"role": "user", "content": "hello"}], max_tokens=8, stream=stream + ), + stream, + ) + + +def _responses(client: openai.OpenAI, model: str, stream: bool) -> None: + _consume(client.responses.create(model=model, input="hello", stream=stream, max_output_tokens=16), stream) + + +@pytest.mark.parametrize("stream", (False, True), ids=("whole", "stream")) +def test_chat_transport_failure_keeps_its_500(gateway: Gateway, stream: bool) -> None: + with gateway.scenario() as scenario: + model: Final = _unreachable_anthropic_model(scenario) + with openai.OpenAI(base_url=f"{_proxy_url(gateway)}/v1", api_key=gateway.key, max_retries=0) as client: + with pytest.raises(openai.InternalServerError) as raised: + _chat(client, model, stream) + assert raised.value.status_code == 500 + assert "Batch input file" not in str(raised.value) + raw: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "hi"}], "stream": stream}, + ) + assert raw.status_code == 500, raw.text + + +@pytest.mark.parametrize("stream", (False, True), ids=("whole", "stream")) +def test_messages_transport_failure_keeps_its_500(gateway: Gateway, stream: bool) -> None: + with gateway.scenario() as scenario: + model: Final = _unreachable_anthropic_model(scenario) + with anthropic.Anthropic(base_url=_proxy_url(gateway), api_key=gateway.key, max_retries=0) as client: + with pytest.raises(anthropic.InternalServerError) as raised: + _messages(client, model, stream) + assert raised.value.status_code == 500 + assert "Batch input file" not in str(raised.value) + + +@pytest.mark.parametrize("stream", (False, True), ids=("whole", "stream")) +def test_responses_transport_failure_keeps_its_500(gateway: Gateway, stream: bool) -> None: + with gateway.scenario() as scenario: + model: Final = _unreachable_anthropic_model(scenario) + with openai.OpenAI(base_url=f"{_proxy_url(gateway)}/v1", api_key=gateway.key, max_retries=0) as client: + with pytest.raises(openai.InternalServerError) as raised: + _responses(client, model, stream) + assert raised.value.status_code == 500 + assert "Batch input file" not in str(raised.value) diff --git a/tests/integration/sdk/test_vertex_batch_row_line_errors_sdk.py b/tests/integration/sdk/test_vertex_batch_row_line_errors_sdk.py new file mode 100644 index 00000000000..1f973367ad7 --- /dev/null +++ b/tests/integration/sdk/test_vertex_batch_row_line_errors_sdk.py @@ -0,0 +1,140 @@ +import asyncio +import json +from collections.abc import Callable +from typing import Final, Literal + +import httpx +import pytest +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from typing_extensions import ReadOnly, TypedDict + +import litellm +from integration.providers._vertex_batch_rows import ( + BOM, + BUCKET, + LOCATION, + PROJECT, + chat_row, + encoded, + is_gcs_upload, + jsonl, + roleless_row, + stored_object, + unconvertible_line_message, +) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.vertex_ai.common_utils import VertexAIError + +Client = Literal["sync", "async_aiohttp", "async_httpx"] +CLIENTS: Final[tuple[Client, ...]] = ("sync", "async_aiohttp", "async_httpx") +_TOKEN_REPLY: Final = Reply( + body=json.dumps({"access_token": "scripted-token", "token_type": "Bearer", "expires_in": 3600}).encode() +) + + +def _peer(request: Request) -> Reply: + if request.target == "/_oauth/token": + return _TOKEN_REPLY + assert is_gcs_upload(request), (request.method, request.target) + return stored_object(request) + + +class _VertexKwargs(TypedDict): + custom_llm_provider: ReadOnly[Literal["vertex_ai"]] + api_base: ReadOnly[str] + vertex_project: ReadOnly[str] + vertex_location: ReadOnly[str] + vertex_credentials: ReadOnly[str] + gcs_bucket_name: ReadOnly[str] + + +def _vertex_kwargs(wire: Wire) -> _VertexKwargs: + return { + "custom_llm_provider": "vertex_ai", + "api_base": wire.url, + "vertex_project": PROJECT, + "vertex_location": LOCATION, + "vertex_credentials": service_account_json(PROJECT, wire.url), + "gcs_bucket_name": BUCKET, + } + + +async def _create_file(client: Client, wire: Wire, content: bytes) -> object: + file: Final = ("rows.jsonl", content, "application/jsonl") + match client: + case "sync": + return await asyncio.to_thread(litellm.create_file, file=file, purpose="batch", **_vertex_kwargs(wire)) + case "async_aiohttp": + return await litellm.acreate_file(file=file, purpose="batch", **_vertex_kwargs(wire)) + case "async_httpx": + handler: Final = AsyncHTTPHandler(transport=httpx.AsyncHTTPTransport(trust_env=False)) + return await litellm.acreate_file(file=file, purpose="batch", client=handler, **_vertex_kwargs(wire)) + + +def _good(custom_id: str) -> bytes: + return encoded(chat_row(custom_id, "gemini-2.5-flash", "fine")) + + +ROWS: Final[tuple[tuple[str, Callable[[], bytes], str], ...]] = ( + ( + "malformed_json_line_3", + lambda: jsonl(_good("first"), b"", b'{"custom_id": "bad", '), + "Batch input file line 3 is not valid JSON: ", + ), + ( + "invalid_utf8_line_3", + lambda: jsonl(_good("first"), b"", b'{"custom_id": "\xff\xfe"}'), + "Batch input file line 3 is not valid UTF-8: ", + ), + ( + "malformed_json_line_1", + lambda: jsonl(b'{"custom_id": "bad", ', _good("second")), + "Batch input file line 1 is not valid JSON: ", + ), + ( + "roleless_row_line_4", + lambda: jsonl(BOM + _good("first"), b"", b" ", encoded(roleless_row("bad", "gemini-2.5-flash"))), + unconvertible_line_message(4) + "missing field 'role'", + ), +) + + +@pytest.mark.parametrize( + ("content", "message"), tuple(row[1:] for row in ROWS), ids=tuple(row[0] for row in ROWS) +) +@pytest.mark.parametrize("client", CLIENTS) +async def test_sdk_create_file_names_the_file_line_vertex_cannot_read( + client: Client, content: Callable[[], bytes], message: str +) -> None: + with wire_server(_peer) as wire: + with pytest.raises(VertexAIError) as raised: + await _create_file(client, wire, content()) + assert raised.value.status_code == 400 + assert message in str(raised.value), str(raised.value) + assert not any(is_gcs_upload(request) for request in wire.drain()) + + +async def test_router_create_file_names_the_line_after_rewriting_each_rows_model() -> None: + with wire_server(_peer) as wire: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "vertex-batch", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", **_vertex_kwargs(wire)}, + } + ] + ) + content: Final = jsonl( + BOM + encoded(chat_row("first", "vertex-batch", "one")), + b"", + b" ", + encoded(roleless_row("bad", "vertex-batch")), + ) + with pytest.raises(VertexAIError) as raised: + await router.acreate_file( + model="vertex-batch", file=("rows.jsonl", content, "application/jsonl"), purpose="batch" + ) + assert raised.value.status_code == 400 + assert unconvertible_line_message(4) + "missing field 'role'" in str(raised.value), str(raised.value) + assert not any(is_gcs_upload(request) for request in wire.drain()) diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 200768bc49f..5337f607719 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -2668,6 +2668,62 @@ def test_sync_retrieve_file_content_raises_on_http_error(): assert exc_info.value.status_code == 404 +@pytest.mark.asyncio +async def test_handle_error_surfaces_provider_error_raised_while_aiohttp_streams_the_body( + monkeypatch: pytest.MonkeyPatch, +): + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig + + monkeypatch.setenv("NO_PROXY", "127.0.0.1") + monkeypatch.setenv("no_proxy", "127.0.0.1") + message: Final = "Batch input file line 3 is not valid JSON: Expecting value at column 1" + + async def read_until_client_hangs_up(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + await reader.read() + writer.close() + + async def body_with_bad_third_row(): + yield b'{"custom_id": "r1"}\n' + raise VertexAIError(status_code=400, message=message) + + server: Final = await asyncio.start_server(read_until_client_hangs_up, "127.0.0.1", 0) + port: Final = server.sockets[0].getsockname()[1] + handler: Final = AsyncHTTPHandler() + try: + async with server: + with pytest.raises(httpx.HTTPError) as sent: + await handler.post(f"http://127.0.0.1:{port}/upload", content=body_with_bad_third_row()) + finally: + await handler.close() + + with pytest.raises(VertexAIError) as surfaced: + BaseLLMHTTPHandler()._handle_error(e=sent.value, provider_config=VertexAIFilesConfig()) + + assert surfaced.value.status_code == 400 + assert surfaced.value.message == message + + +def _runtime_error_caused_by(cause: Exception) -> RuntimeError: + try: + raise RuntimeError("upload bookkeeping failed") from cause + except RuntimeError as outer: + return outer + + +def test_handle_error_reports_its_own_failure_when_it_is_not_a_transport_error(): + from litellm.llms.vertex_ai.common_utils import VertexAIError + from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig + + outer: Final = _runtime_error_caused_by(VertexAIError(status_code=400, message="inner provider error")) + + with pytest.raises(VertexAIError) as surfaced: + BaseLLMHTTPHandler()._handle_error(e=outer, provider_config=VertexAIFilesConfig()) + + assert surfaced.value.status_code == 500 + assert "upload bookkeeping failed" in surfaced.value.message + + _FILE_CONTENT_WIF_ENV = { "ANTHROPIC_FEDERATION_RULE_ID": "fdrl_llm_http_handler_seam", "ANTHROPIC_ORGANIZATION_ID": "org-llm-http-handler-seam", diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index cc15422afa8..f59c584eb90 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -26,9 +26,11 @@ import gc import gzip import io import json +import sys import tempfile import time import tracemalloc +from collections.abc import Callable from pathlib import Path from typing import Final @@ -46,7 +48,7 @@ from litellm.llms.vertex_ai.files.transformation import ( VertexAIFilesConfig, _get_litellm_batch_custom_id_from_labels, _iter_openai_jsonl_entries, - _iter_openai_jsonl_lines, + _iter_numbered_openai_jsonl_lines, _openai_batch_jsonl_entry_to_vertex_rows, _OpenAIToVertexBatchUploadStream, ) @@ -139,7 +141,7 @@ class TestFileLikeInputNotPartiallyConsumed: (get_complete_file_url -> get_object_name) runs before transform_create_file_request, and both read the same create_file_data source. When the file is a tuple-wrapped open handle, the streaming reader - must still emit every row including entry 0: ``_iter_openai_jsonl_lines`` + must still emit every row including entry 0: ``_iter_numbered_openai_jsonl_lines`` rewinds a seekable source (seek(0)) before each pass, so the object-name step's partial read of the cursor does not consume the upload. A partial upload missing the first request would be silent and hard to catch, so this @@ -181,17 +183,17 @@ class TestFileLikeInputNotPartiallyConsumed: class TestStreamingLineIterator: def test_skips_blank_and_whitespace_lines(self): content = b'{"a": 1}\n\n \n{"b": 2}\n' - assert list(_iter_openai_jsonl_lines(content)) == ['{"a": 1}', '{"b": 2}'] + assert list(_iter_numbered_openai_jsonl_lines(content)) == [(1, '{"a": 1}'), (4, '{"b": 2}')] def test_handles_crlf_and_missing_trailing_newline(self): content = b'{"a": 1}\r\n{"b": 2}' - assert [json.loads(line) for line in _iter_openai_jsonl_lines(content)] == [ + assert [json.loads(line) for _, line in _iter_numbered_openai_jsonl_lines(content)] == [ {"a": 1}, {"b": 2}, ] def test_accepts_str_bytes_tuple_and_filelike(self): - expected = [{"a": 1}, {"b": 2}] + expected = [(1, {"a": 1}), (2, {"b": 2})] text = '{"a": 1}\n{"b": 2}\n' for source in ( text, @@ -202,19 +204,19 @@ class TestStreamingLineIterator: assert list(_iter_openai_jsonl_entries(source)) == expected def test_str_input_without_trailing_newline(self): - assert list(_iter_openai_jsonl_lines('{"a": 1}\n{"b": 2}')) == [ - '{"a": 1}', - '{"b": 2}', + assert list(_iter_numbered_openai_jsonl_lines('{"a": 1}\n{"b": 2}')) == [ + (1, '{"a": 1}'), + (2, '{"b": 2}'), ] def test_pathlike_input_is_read_line_by_line(self, tmp_path): path = tmp_path / "batch.jsonl" path.write_bytes(b'{"a": 1}\n{"b": 2}\n') - assert list(_iter_openai_jsonl_entries(path)) == [{"a": 1}, {"b": 2}] + assert list(_iter_openai_jsonl_entries(path)) == [(1, {"a": 1}), (2, {"b": 2})] def test_unsupported_content_type_raises(self): with pytest.raises(ValueError, match="Unsupported file content type"): - list(_iter_openai_jsonl_lines(12345)) # type: ignore[arg-type] + list(_iter_numbered_openai_jsonl_lines(12345)) # type: ignore[arg-type] def test_non_seekable_handle_raises_instead_of_dropping_first_row(self): # The handle is read twice (object-name probe, then body). A non-seekable @@ -235,15 +237,17 @@ class TestStreamingLineIterator: handle = _NonSeekable(b'{"custom_id": "request-0"}\n{"custom_id": "request-1"}\n') with pytest.raises(ValueError, match="seekable"): - list(_iter_openai_jsonl_lines(handle)) + list(_iter_numbered_openai_jsonl_lines(handle)) def test_is_lazy_does_not_parse_past_first_entry(self): # Second row is invalid JSON; pulling only the first entry must not raise. content = b'{"custom_id": "first"}\nnot-json-at-all\n' gen = _iter_openai_jsonl_entries(content) - assert next(gen)["custom_id"] == "first" - with pytest.raises(json.JSONDecodeError): + assert next(gen) == (1, {"custom_id": "first"}) + with pytest.raises(VertexAIError) as exc_info: next(gen) + assert exc_info.value.status_code == 400 + assert "line 2" in str(exc_info.value) class TestGetObjectNameLazyParse: @@ -255,6 +259,24 @@ class TestGetObjectNameLazyParse: object_name = cfg.get_object_name(("batch.jsonl", raw, "application/jsonl"), purpose="batch") assert "gemini-2.5-flash" in object_name + @pytest.mark.parametrize( + "first_row", + [ + b'{"custom_id": "r-0", "body": null}', + b'{"custom_id": "r-0", "body": "hi"}', + b'{"custom_id": "r-0", "body": {"model": null, "messages": [{"role": "user", "content": "hi"}]}}', + b'{"custom_id": "r-0", "body": {"model": 5, "messages": [{"role": "user", "content": "hi"}]}}', + ], + ) + def test_bad_body_on_line_1_is_reported_as_a_400_naming_the_line(self, first_row: bytes): + cfg: Final = VertexAIFilesConfig() + file_data: Final = ("batch.jsonl", first_row + b"\n", "application/jsonl") + cfg.get_object_name(file_data, purpose="batch") + with pytest.raises(VertexAIError) as exc_info: + list(_OpenAIToVertexBatchUploadStream(file_data, cfg._map_openai_to_vertex_params).iter_bytes()) + assert exc_info.value.status_code == 400 + assert exc_info.value.message.startswith("Batch input file line 1 could not be converted") + class TestStreamingLaziness: def test_upload_stream_transforms_rows_only_as_they_are_pulled(self) -> None: @@ -268,8 +290,10 @@ class TestStreamingLaziness: ) assert custom_ids == ("request-0", "request-1", "request-2") - with pytest.raises(json.JSONDecodeError): + with pytest.raises(VertexAIError) as exc_info: next(chunks) + assert exc_info.value.status_code == 400 + assert "line 4" in exc_info.value.message def test_get_object_name_does_not_scale_with_payload(self): cfg = VertexAIFilesConfig() @@ -368,6 +392,17 @@ class TestPathSourcedStreaming: assert first == second and len(first) > 0 +_LINE_2_NOT_CONVERTED: Final = "Batch input file line 2 could not be converted to a Vertex AI request" + + +def _as_utf8_bytes(text: str) -> bytes: + return text.encode("utf-8") + + +def _as_str(text: str) -> str: + return text + + _GCS_OBJECT_JSON = { "id": "test-bucket/litellm-vertex-files/x/123", "name": "litellm-vertex-files/x", @@ -464,6 +499,144 @@ class TestUploadUrl: class TestUploadStreamBody: + def test_stream_rejects_malformed_single_row(self): + stream = _OpenAIToVertexBatchUploadStream( + openai_file_content=b"this is not json\n", + map_openai_to_vertex_params=lambda body: body, + ) + with pytest.raises(VertexAIError) as exc_info: + list(stream.iter_bytes()) + assert exc_info.value.status_code == 400 + assert "line 1" in str(exc_info.value) + + def test_stream_reports_file_line_number_for_malformed_file_like_input(self): + valid_row = json.dumps( + { + "custom_id": "r1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.5-flash", + "messages": [{"role": "user", "content": "hi"}], + }, + } + ).encode("utf-8") + raw = valid_row + b"\n\nthis is not json\n" + stream = _OpenAIToVertexBatchUploadStream( + openai_file_content=io.BytesIO(raw), + map_openai_to_vertex_params=lambda body: body, + ) + with pytest.raises(VertexAIError) as exc_info: + list(stream.iter_bytes()) + assert "line 3" in str(exc_info.value) + + def test_stream_reports_file_line_number_for_malformed_bytes_input(self): + valid_row = json.dumps( + { + "custom_id": "r1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.5-flash", + "messages": [{"role": "user", "content": "hi"}], + }, + } + ).encode("utf-8") + raw = valid_row + b"\n\nthis is not json\n" + stream = _OpenAIToVertexBatchUploadStream( + openai_file_content=raw, + map_openai_to_vertex_params=lambda body: body, + ) + with pytest.raises(VertexAIError) as exc_info: + list(stream.iter_bytes()) + assert "line 3" in str(exc_info.value) + + @pytest.mark.parametrize( + ("bad_row", "expected_message"), + [ + (b"[1, 2]", "Batch input file line 2 is not a JSON object"), + (b'"text"', "Batch input file line 2 is not a JSON object"), + (b"42", "Batch input file line 2 is not a JSON object"), + ( + b'{"custom_id": "r2", "body": "hi"}', + f"{_LINE_2_NOT_CONVERTED}: `body` on batch requests must be a JSON object", + ), + ( + b'{"custom_id": "r2", "body": [1]}', + f"{_LINE_2_NOT_CONVERTED}: `body` on batch requests must be a JSON object", + ), + ( + b'{"custom_id": "r2", "body": null}', + f"{_LINE_2_NOT_CONVERTED}: `body` on batch requests must be a JSON object", + ), + ( + b'{"custom_id": "r2", "body": ""}', + f"{_LINE_2_NOT_CONVERTED}: `body` on batch requests must be a JSON object", + ), + (b'{"custom_id": "r2"}', f"{_LINE_2_NOT_CONVERTED}: `body` on batch requests must be a JSON object"), + (b"[" * 1_000_000, "Batch input file line 2 could not be parsed"), + ( + b'{"n": ' + b"9" * (sys.get_int_max_str_digits() + 1) + b"}", + "Batch input file line 2 could not be parsed", + ), + (b'{"custom_id": "r2", "body": {"a": "\xff"}}', "Batch input file line 2 is not valid UTF-8"), + ( + b'{"custom_id": "r2", "url": "/v1/embeddings", "body": {"model": "gemini-embedding-001"}}', + f"{_LINE_2_NOT_CONVERTED}: `input` is required", + ), + ( + b'{"custom_id": "r2", "url": "/v1/embeddings", "body": {"model": "gemini-embedding-001", "input": 5}}', + _LINE_2_NOT_CONVERTED, + ), + ( + b'{"custom_id": "r2", "body": {"model": "gemini-2.5-flash", "messages": [{"content": "hi"}]}}', + f"{_LINE_2_NOT_CONVERTED}: missing field 'role'", + ), + ( + b'{"custom_id": "r2", "body": {"model": "gemini-2.5-flash", "messages": [{"role": 5, "content": "hi"}]}}', + _LINE_2_NOT_CONVERTED, + ), + (b'{"custom_id": "r2", "body": {"a": 1,}}', "Batch input file line 2 is not valid JSON"), + ], + ) + def test_stream_rejects_malformed_row_naming_its_line(self, bad_row: bytes, expected_message: str): + valid_row = json.dumps({"custom_id": "r1", "body": {"model": "gemini-2.5-flash"}}).encode("utf-8") + stream = _OpenAIToVertexBatchUploadStream( + openai_file_content=valid_row + b"\n" + bad_row + b"\n", + map_openai_to_vertex_params=lambda body: body, + ) + with pytest.raises(VertexAIError) as exc_info: + list(stream.iter_bytes()) + assert exc_info.value.status_code == 400 + assert expected_message in exc_info.value.message + + @pytest.mark.parametrize("encode", [_as_utf8_bytes, _as_str]) + @pytest.mark.parametrize("bad_row", ['{"custom_id": "r1",}', " {bad", '\t\t{"a": tru}']) + def test_invalid_json_error_names_the_column(self, bad_row: str, encode: Callable[[str], bytes | str]): + stream = _OpenAIToVertexBatchUploadStream( + openai_file_content=encode(bad_row + "\n"), + map_openai_to_vertex_params=lambda body: body, + ) + with pytest.raises(VertexAIError) as exc_info: + list(stream.iter_bytes()) + with pytest.raises(json.JSONDecodeError) as parsed: + json.loads(bad_row) + assert exc_info.value.message.endswith(f"{parsed.value.msg} at column {parsed.value.colno}") + + @pytest.mark.parametrize("encode", [_as_utf8_bytes, _as_str]) + def test_stream_uploads_a_file_that_starts_with_a_byte_order_mark(self, encode: Callable[[str], bytes | str]): + row: Final = json.dumps( + {"custom_id": "r1", "body": {"model": "gemini-2.5-flash", "messages": [{"role": "user", "content": "hi"}]}} + ) + with_bom: Final = _OpenAIToVertexBatchUploadStream(encode("\ufeff" + row + "\n"), lambda body: body) + without_bom: Final = _OpenAIToVertexBatchUploadStream((row + "\n").encode("utf-8"), lambda body: body) + assert b"".join(with_bom.iter_bytes()) == b"".join(without_bom.iter_bytes()) + + def test_get_object_name_rejects_malformed_batch_jsonl(self): + with pytest.raises(VertexAIError) as exc_info: + VertexAIFilesConfig().get_object_name(file_data=b"not json\n", purpose="batch") + assert exc_info.value.status_code == 400 + def test_stream_matches_legacy_pipeline(self): cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(120) diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_batch_file_validation.py b/tests/unit/proxy/openai_files_endpoint/test_files_batch_file_validation.py index b55ae1500be..5cdc8667118 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_batch_file_validation.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_batch_file_validation.py @@ -1,4 +1,5 @@ import io +import sys import pytest @@ -92,6 +93,15 @@ def test_non_utf8_line_rejected_as_invalid_json(): ) +@pytest.mark.parametrize( + "unparseable_line", + [b"[" * 1_000_000, b'{"n": ' + b"9" * (sys.get_int_max_str_digits() + 1) + b"}"], +) +def test_line_too_deep_or_too_long_to_parse_rejected_as_invalid_json(unparseable_line: bytes): + content = VALID_LINE + b"\n" + unparseable_line + b"\n" + assert check_batch_file_upload("batch.jsonl", content, None) == BatchFileInvalidJsonLine(line_number=2) + + def test_non_object_line_rejected(): content = VALID_LINE + b"\n" + b'["custom_id", "method"]\n' assert check_batch_file_upload("batch.jsonl", content, None) == BatchFileLineNotObject(line_number=2) diff --git a/tests/unit/router_utils/test_batch_utils.py b/tests/unit/router_utils/test_batch_utils.py index b1cfe1ac911..4aa503ee528 100644 --- a/tests/unit/router_utils/test_batch_utils.py +++ b/tests/unit/router_utils/test_batch_utils.py @@ -276,6 +276,40 @@ def test_replace_model_in_jsonl_multi_row_rewrites_every_model(): assert all(row["body"]["model"] == "new-model" for row in rows) +def test_replace_model_in_jsonl_keeps_every_row_on_its_original_line(): + content = ( + b'{"custom_id":"a","body":{"model":"old"}}\n' + b"\n" + b'{"custom_id":"c","body":{"model":"old"}}\n' + b'{"custom_id":"d",\n' + b'"body":{"model":"old"}}\n' + b'{"custom_id":"f","body":{"model":"old"}}\n' + ) + + result = replace_model_in_jsonl(content, "new-model", keep_line_numbers=True) + + assert isinstance(result, InMemoryFile) + lines = result.getvalue().decode("utf-8").split("\n") + assert [json.loads(line)["custom_id"] if line else None for line in lines] == ["a", None, "c", None, "d", "f"] + assert all(json.loads(line)["body"]["model"] == "new-model" for line in lines if line) + + +@pytest.mark.parametrize( + "content", + [ + b'\xef\xbb\xbf{"custom_id":"a","body":{"model":"old"}}\n\xef\xbb\xbf{"custom_id":"b","body":{"model":"old"}}\n', + '\ufeff{"custom_id":"a","body":{"model":"old"}}\n\ufeff{"custom_id":"b","body":{"model":"old"}}\n', + ], + ids=["bytes", "str"], +) +def test_replace_model_in_jsonl_rewrites_rows_that_start_with_a_bom(content): + result = replace_model_in_jsonl(content, "new-model") + + assert isinstance(result, InMemoryFile) + rows = [json.loads(line) for line in result.getvalue().decode("utf-8").split("\n")] + assert [(row["custom_id"], row["body"]["model"]) for row in rows] == [("a", "new-model"), ("b", "new-model")] + + def test_replace_model_in_jsonl_with_embedded_newlines(): """Test that replace_model_in_jsonl works correctly with embedded newlines in content""" # Create a JSONL with embedded newlines in the message content diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index a7c11cc55f3..3c6ba5ed7b4 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -578,6 +578,30 @@ async def test_async_router_acreate_file_with_jsonl(): assert first_call_content == non_jsonl_content +@pytest.mark.asyncio +@pytest.mark.parametrize( + "litellm_params, uploaded_rows", + [ + ({"model": "gpt-5.5"}, ["a", "c"]), + ({"model": "vertex_ai/gemini-3.8-flash", "vertex_project": "p"}, [None, "a", None, "c"]), + ], + ids=["openai", "vertex_ai"], +) +async def test_async_router_acreate_file_keeps_blank_lines_only_for_vertex(litellm_params, uploaded_rows): + import json + from io import BytesIO + from unittest.mock import MagicMock, patch + + content = b'\n{"custom_id": "a", "body": {"model": "batch-model"}}\n\n{"custom_id": "c", "body": {"model": "batch-model"}}\n' + router = litellm.Router(model_list=[{"model_name": "batch-model", "litellm_params": litellm_params}]) + + with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file: + await router.acreate_file(model="batch-model", purpose="batch", file=BytesIO(content)) + + lines = mock_acreate_file.call_args.kwargs["file"].read().decode("utf-8").split("\n") + assert [json.loads(line)["custom_id"] if line else None for line in lines] == uploaded_rows + + @pytest.mark.asyncio async def test_async_router_acreate_file_passthrough_keeps_the_file_and_forwards_the_flag(): """A passthrough batch upload must reach the provider byte for byte: the router