mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
fix(vertex_ai): return 400 with file line number for malformed batch JSONL row (#40097)
* fix(vertex_ai): return 400 with file line number for malformed batch JSONL row Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(vertex_ai): keep yield outside the JSON decode guard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(http_handler): surface BaseLLMException raised while streaming a request body Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(http_handler): only follow __cause__ when unwrapping streamed body errors Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(http_handler): avoid recursive streamed error unwrapping Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(vertex_ai): return a 400 naming the line for wrong-shape and non-UTF-8 batch rows Keep main's error path for a directly raised provider error and only unwrap one found behind a transport error. Add the parser column to invalid JSON errors, and cover the real aiohttp chain in the _handle_error test * fix(vertex_ai): report every row conversion error as a 400 naming its line A message with no role, a non-string role, a bad temperature, or a litellm BadRequestError raised while converting a row still returned a 500, and the async upload lost the 400 for litellm's own exceptions. The stream now turns any error from one row's conversion into a 400 naming that line. _handle_error only unwraps an httpx.TransportError, which is what the aiohttp transport raises around a failing request body, so other callers keep their own status. A file that starts with a UTF-8 byte order mark now uploads, and provider messages start with 'Batch input file line N' like the proxy check. * fix(vertex_ai): reject null or missing batch bodies and unparseable rows with a 400 naming the line A chat row whose body is missing or not an object now raises a 400 naming its line instead of uploading a blank prompt, and a non-object body on line 1 no longer crashes the GCS object name lookup. Rows too deeply nested or too long to parse return a 400 from both the SDK and the proxy's own batch file check. JSON error columns now count leading whitespace. * fix(vertex_ai): keep batch rows on their file line through router and guardrail rewrites The Router model rewrite and batch guardrails now keep blank lines and blank out dropped records, so a provider error naming a line still points into the caller's file. The rewrite strips a leading BOM per line, and a non-string body.model on line 1 is reported as a 400 naming the line * test(integration): cover Vertex batch row line errors across routes, rewrites, SDK and worker chaos * fix(router): keep blank lines only for Vertex batch uploads and leave guardrail rewrites compact * test(integration): score the whole OpenAI upload after a guardrail drop and cover a mixed OpenAI and Vertex batch upload --------- Co-authored-by: jesus <jesus@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
f9d45f8a43
commit
c00706218b
15 changed files with 1618 additions and 64 deletions
|
|
@ -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"):
|
||||
|
|
|
|||
|
|
@ -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/<id>` 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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
174
tests/integration/providers/_vertex_batch_rows.py
Normal file
174
tests/integration/providers/_vertex_batch_rows.py
Normal file
|
|
@ -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")
|
||||
|
|
@ -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()
|
||||
|
|
@ -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()
|
||||
|
|
@ -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)
|
||||
140
tests/integration/sdk/test_vertex_batch_row_line_errors_sdk.py
Normal file
140
tests/integration/sdk/test_vertex_batch_row_line_errors_sdk.py
Normal file
|
|
@ -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())
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue