From 1d424104cfda5f37bb7b0229d805c779ba0fdbf8 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 24 Jun 2026 13:19:57 -0700 Subject: [PATCH] fix(vertex/files): stream OpenAI->Vertex batch JSONL uploads (#31036) * fix(vertex/files): stream OpenAI->Vertex batch JSONL uploads to fix OOM on large files Large (1GB+) batch JSONL uploads to Vertex AI / GCS caused OOM or killed the worker because the request body was buffered and multiplied 2-3x in size. The create-file path is now streaming end-to-end: transform_create_file_request returns a ResumableChunkedUploadConfig carrying a lazy _OpenAIToVertexBatchUploadStream, and the HTTP handler opens a GCS resumable session and PUTs the body in bounded 8 MiB chunks (Content-Range, 308 between chunks) so the transformed payload is never held in full. The proxy /v1/files endpoint streams from Starlette's spooled upload handle instead of reading the whole body, and batch rate limiting counts tokens and models in a single streaming pass. Only gcs_bucket_name is supported for the GCS target; the legacy bucket_name key is intentionally not read. Also removes the unreachable VertexAIFilesHandler create path and everything only it kept alive (VertexAIJsonlFilesTransformation, _stream_openai_jsonl_to_vertex, the legacy transform helpers), plus the orphaned batch_utils helpers the streaming rewrite replaced. * fix(batches): return original JSONL on unparseable row to avoid silent batch truncation The streaming rewrite of replace_model_in_jsonl accumulated physical lines and skipped a row on JSONDecodeError to support multi-line objects, but a genuinely malformed or truncated row never completes: it poisons the buffer, swallows every following row, and the function still returned the partial rewrite (the rows before the bad one, already model-rewritten) as if the batch were complete. That turned the pre-rewrite behavior of returning the original file unchanged (so the provider rejects the bad batch loudly) into a silent partial submission. Restore the original-content fallback: when an unparseable remainder is left after the loop, return the original file_content (rewinding a consumed seekable source) instead of the truncated output. The multi-line happy path is unchanged. * test(batches): mock resumable GCS upload in vertex batch prediction test The vertex batch file-create path now streams to a GCS resumable session via _aresumable_chunked_upload (httpx send) instead of AsyncHTTPHandler.post, so the existing test's post mock no longer intercepted the upload and a real request hit GCS (401). Mock _aresumable_chunked_upload to return the GCS object response; the resumable protocol itself is covered in test_vertex_ai_files_streaming.py. * fix(batches): resilient per-row token accounting; no hard-block on count failure The batch input-file pass iterated a generator whose json.loads raised on a malformed line; the outer except caught it and stopped the loop, so any body.model on rows after a bad line was never collected and the model allowlist check ran against a partial set. It also hard-blocked the batch with a 400 whenever token counting raised, a backwards-incompatible change from the prior swallow-and-proceed behavior that breaks legitimate rows the token counter cannot measure (e.g. some multimodal content). Iterate the JSONL line-by-line and account each row independently. A malformed line is skipped (its request cannot run upstream anyway) and a row the counter cannot measure falls back to a conservative size-based estimate. The loop never aborts, so the allowlist check always sees every parseable model, and the token total is never zeroed, so a crafted uncountable row still cannot evade the TPM limit, without hard-rejecting a legitimate batch. * perf(vertex/files): unblock async upload; drop empty finalize; widen batch MIME types Three review follow-ups on the resumable batch upload: - _aresumable_chunked_upload pulled chunks from a synchronous generator that runs the per-row transform inline on the event loop thread, blocking other requests between PUTs on large uploads. Each chunk is now produced via asyncio.to_thread. - _iter_resumable_chunks no longer yields a trailing empty chunk, so an exactly chunk-aligned upload finalizes on its last data chunk instead of an extra zero-byte PUT; a 0-byte stream still finalizes via the caller's empty request. - valid_content_type now accepts the MIME types clients label .jsonl batch uploads with (text/plain, application/json, ndjson, ...), so such a batch file no longer silently bypasses the streaming path into the buffered media upload. * fix(vertex/files): keep legacy bucket_name as GCS bucket fallback The rename to gcs_bucket_name dropped the legacy bucket_name key entirely, so an SDK caller passing bucket_name to a Vertex AI file create/retrieve/content call with GCS_BUCKET_NAME unset got ValueError("GCS bucket_name is required") where it previously resolved the bucket. _get_configured_bucket_name now reads gcs_bucket_name, then bucket_name, then the env var, and bucket_name is restored to OPTIONAL_KWARGS_KEYS so it survives get_litellm_params on the retrieve and content paths. gcs_bucket_name keeps precedence when both are present * style: sort imports in llm_http_handler to satisfy I001 budget --------- Co-authored-by: Yuneng Jiang (cherry picked from commit 56825926af7f23969e47e2979e71431861a8701e) --- .../proxy/hooks/managed_files.py | 8 +- litellm/batches/batch_utils.py | 130 ++-- litellm/files/utils.py | 35 +- .../litellm_core_utils/get_litellm_params.py | 1 + .../prompt_templates/common_utils.py | 40 + litellm/llms/base_llm/files/transformation.py | 18 +- litellm/llms/custom_httpx/llm_http_handler.py | 266 ++++++- litellm/llms/vertex_ai/files/handler.py | 82 +-- .../llms/vertex_ai/files/transformation.py | 522 ++++++------- litellm/proxy/hooks/batch_rate_limiter.py | 74 +- .../openai_files_endpoints/files_endpoints.py | 38 +- litellm/router_utils/batch_utils.py | 100 ++- litellm/types/files.py | 18 + litellm/types/router.py | 3 + .../test_openai_batches_and_files.py | 26 +- .../test_router_batch_utils.py | 93 ++- .../test_vertex_ai_binary_file_upload.py | 28 +- .../files/test_vertex_ai_files_streaming.py | 696 ++++++++++++++++++ .../test_vertex_ai_files_transformation.py | 232 ++++-- .../proxy/hooks/test_batch_file_validation.py | 448 ++++++++--- .../test_files_endpoint.py | 77 ++ tests/test_litellm/test_router.py | 30 + 22 files changed, 2264 insertions(+), 701 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index ae5905f9cdf..8722ba4b429 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -13,7 +13,9 @@ from litellm import Router, verbose_logger from litellm._uuid import uuid from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_metadata, +) from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.llms.base_llm.managed_resources.isolation import ( build_list_page, @@ -981,9 +983,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): target_model_names_list: List[str], ) -> OpenAIFileObject: ## GET THE FILE TYPE FROM THE CREATE FILE REQUEST - file_data = extract_file_data(create_file_request["file"]) - - file_type = file_data["content_type"] + _, file_type = extract_file_metadata(create_file_request["file"]) output_file_id = file_objects[0].id model_id = file_objects[0]._hidden_params.get("model_id") diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 74e753b09ea..aeec58f1dfc 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -1,5 +1,5 @@ import json -from typing import Any, List, Literal, Optional, Tuple +from typing import Any, Iterator, List, Literal, Optional, Tuple import litellm from litellm._logging import verbose_logger @@ -314,6 +314,70 @@ def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]: raise e +def _iter_batch_input_lines(file_content: bytes) -> Iterator[bytes]: + """ + Yield non-empty JSONL lines (unparsed) one at a time, so a caller can parse + each row in its own try/except and a single malformed line cannot abort the + whole pass. Peak memory stays bounded for large batch files. + """ + start, length, newline = 0, len(file_content), ord("\n") + while start < length: + idx = file_content.find(newline, start) + if idx == -1: + chunk, start = file_content[start:], length + else: + chunk, start = file_content[start:idx], idx + 1 + line = chunk.strip() + if line: + yield line + + +def _iter_batch_input_entries(file_content: bytes) -> Iterator[dict]: + """ + Yield parsed batch input JSONL entries one at a time without materializing the + whole file as a list, so peak memory stays bounded. Raises on a malformed line; + callers that must survive bad rows should iterate ``_iter_batch_input_lines`` + and parse per-row instead. + """ + for line in _iter_batch_input_lines(file_content): + yield json.loads(line) + + +# A batch request's input tokens scale roughly with its serialized size, so this +# is a conservative per-row fallback when the token counter cannot measure a row. +_BATCH_TOKEN_ESTIMATE_BYTES_PER_TOKEN = 4 + + +def _estimate_batch_entry_tokens(raw_line: bytes) -> int: + """Conservative token estimate for a batch row the token counter cannot measure + (or that cannot be parsed). Keeps the batch token total non-zero so a crafted + row cannot evade the TPM limit, without hard-rejecting a legitimate batch.""" + return max(1, len(raw_line) // _BATCH_TOKEN_ESTIMATE_BYTES_PER_TOKEN) + + +def _count_entry_tokens( + entry: dict, + model_name: Optional[str] = None, +) -> int: + """Token-count a single batch input entry's body (chat / text / embedding).""" + body = entry.get("body", {}) or {} + model = body.get("model", model_name or "") + + messages = body.get("messages") + if messages: + return token_counter(model=model, messages=messages) + + prompt = body.get("prompt") + if prompt: + return _count_prompt_or_input_tokens(model=model, value=prompt) + + input_data = body.get("input") + if input_data: + return _count_prompt_or_input_tokens(model=model, value=input_data) + + return 0 + + def _get_batch_job_cost_from_file_content( file_content_dictionary: List[dict], custom_llm_provider: Literal[ @@ -396,70 +460,6 @@ def _get_batch_job_total_usage_from_file_content( ) -def _get_models_from_batch_input_file_content( - file_content_dictionary: List[dict], -) -> List[str]: - """Extract the distinct ``body.model`` values from a batch *input* file. - - Used by the proxy's batch pre-call hook to enforce that the caller is - authorized for every model named inside the JSONL — not just the one - on the outer request — so the proxy's per-key model allowlist isn't - bypassed by smuggling expensive models into the batch file. - """ - models: List[str] = [] - seen: set = set() - for _item in file_content_dictionary: - body = _item.get("body") or {} - model = body.get("model") - if model and model not in seen: - seen.add(model) - models.append(model) - return models - - -def _get_batch_job_input_file_usage( - file_content_dictionary: List[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", - model_name: Optional[str] = None, -) -> Usage: - """ - Count the number of tokens in the input file - - Used for batch rate limiting to count the number of tokens in the input file - """ - prompt_tokens: int = 0 - completion_tokens: int = 0 - - for _item in file_content_dictionary: - body = _item.get("body", {}) - model = body.get("model", model_name or "") - - # Chat completion payloads. - messages = body.get("messages") - if messages: - prompt_tokens += token_counter(model=model, messages=messages) - continue - - # Text completion payloads (`prompt`). - prompt = body.get("prompt") - if prompt: - prompt_tokens += _count_prompt_or_input_tokens(model=model, value=prompt) - continue - - # Embedding payloads (`input`). - input_data = body.get("input") - if input_data: - prompt_tokens += _count_prompt_or_input_tokens( - model=model, value=input_data - ) - - return Usage( - total_tokens=prompt_tokens + completion_tokens, - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - ) - - def _count_prompt_or_input_tokens(model: str, value: Any) -> int: """Token-count a ``prompt`` / ``input`` field that the OpenAI batch schema allows in four shapes: diff --git a/litellm/files/utils.py b/litellm/files/utils.py index a2b9a42c154..a0df7a89b0f 100644 --- a/litellm/files/utils.py +++ b/litellm/files/utils.py @@ -3,6 +3,22 @@ from typing import Optional from litellm.types.llms.openai import CreateFileRequest from litellm.types.utils import ExtractedFileData +# MIME types a .jsonl batch upload is plausibly labeled with. Clients are +# inconsistent (text/plain, application/json, octet-stream, ndjson, ...), so a +# batch file must not silently bypass the streaming path just because of its +# declared type. ``purpose == "batch"`` is the authoritative signal; non-JSONL +# content still fails loudly when the rows are parsed. +_BATCH_JSONL_CONTENT_TYPES = frozenset( + { + "application/jsonl", + "application/json", + "application/octet-stream", + "application/x-ndjson", + "application/x-jsonlines", + "text/plain", + } +) + class FilesAPIUtils: """ @@ -24,9 +40,24 @@ class FilesAPIUtils: and extracted_file_data.get("content") is not None ) + @staticmethod + def is_batch_jsonl_request( + create_file_data: CreateFileRequest, content_type: Optional[str] + ) -> bool: + """ + Batch-jsonl check from metadata only, so the body can stay a streamable + Path/handle instead of being read into memory. + """ + return ( + create_file_data.get("purpose") == "batch" + and FilesAPIUtils.valid_content_type(content_type) + and create_file_data.get("file") is not None + ) + @staticmethod def valid_content_type(content_type: Optional[str]) -> bool: """ - Check if the content type is valid + Whether the upload's MIME type is one a batch JSONL file is plausibly + sent as (see ``_BATCH_JSONL_CONTENT_TYPES``). """ - return content_type in set(["application/jsonl", "application/octet-stream"]) + return content_type in _BATCH_JSONL_CONTENT_TYPES diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index b32803b5dfc..68ca78a70d8 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -14,6 +14,7 @@ _OPTIONAL_KWARGS_KEYS = frozenset( "azure_password", "azure_scope", "timeout", + "gcs_bucket_name", "bucket_name", "vertex_credentials", "vertex_project", diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index b44d21368f8..e50d621bb18 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -757,6 +757,46 @@ def update_responses_tools_with_model_file_ids( return updated_tools +def extract_file_metadata(file_data: FileTypes) -> Tuple[Optional[str], Optional[str]]: + """ + Resolve (filename, content_type) without reading the file body. + + Mirrors extract_file_data's metadata resolution but never calls .read(), so + it stays O(1) on large uploads. Use this when only metadata is needed (batch + detection, GCS object naming) and the body must remain a streamable Path/handle. + """ + filename: Optional[str] = None + content_type: Optional[str] = None + file_content: Any = None + + if isinstance(file_data, tuple): + if len(file_data) == 2: + filename, file_content = file_data + elif len(file_data) == 3: + filename, file_content, content_type = file_data + elif len(file_data) == 4: + filename, file_content, content_type, _ = file_data + elif isinstance(file_data, InMemoryFile): + filename = file_data.name + content_type = file_data.content_type + else: + file_content = file_data + + if filename is None: + if isinstance(file_content, PathLike): + filename = Path(file_content).name + elif isinstance(file_content, io.IOBase): + name_attr = getattr(file_content, "name", None) + if isinstance(name_attr, str): + filename = Path(name_attr).name + + if not content_type: + guessed = mimetypes.guess_type(filename)[0] if filename else None + content_type = guessed or "application/octet-stream" + + return filename, content_type + + def extract_file_data(file_data: FileTypes) -> ExtractedFileData: """ Extracts and processes file data from various input formats. diff --git a/litellm/llms/base_llm/files/transformation.py b/litellm/llms/base_llm/files/transformation.py index c3abfafc552..85016c7a5c4 100644 --- a/litellm/llms/base_llm/files/transformation.py +++ b/litellm/llms/base_llm/files/transformation.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union import httpx from openai.types.file_deleted import FileDeleted @@ -32,6 +32,22 @@ else: Router = Any +class BaseFileUploadStream(ABC): + """Re-iterable request body that yields an upload's bytes lazily. + + A provider returns one of these (inside the upload config from + ``transform_create_file_request``) when the upload body can be produced + incrementally; the HTTP handler then sends it in bounded chunks instead of + buffering the whole payload, which is what exhausts memory on large uploads. + + ``iter_bytes`` must return a fresh iterator each call so the body can be + replayed if the upload is retried. + """ + + @abstractmethod + def iter_bytes(self) -> Iterator[bytes]: ... + + class BaseFilesConfig(BaseConfig): @property @abstractmethod diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index cabf6cce7cc..2e5de8ea599 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,12 +1,13 @@ +import asyncio import json import ssl -from urllib.parse import parse_qs, urlencode, urlparse, urlunparse from typing import ( TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, + Iterator, List, Literal, Optional, @@ -14,6 +15,7 @@ from typing import ( Union, cast, ) +from urllib.parse import parse_qs, urlencode, urlparse, urlunparse import httpx # type: ignore from openai.types.file_deleted import FileDeleted @@ -3223,6 +3225,23 @@ class BaseLLMHTTPHandler: data=presigned_request["data"], timeout=timeout, ) + elif ( + isinstance(transformed_request, dict) + and "resumable_chunked_upload" in transformed_request + ): + try: + upload_response = self._resumable_chunked_upload( + client=sync_httpx_client, + initiate_url=api_base, + base_headers=headers, + config=cast(Dict[str, Any], transformed_request)[ + "resumable_chunked_upload" + ], + timeout=timeout, + ) + except Exception as e: + verbose_logger.exception(f"Error creating file: {e}") + raise self._handle_error(e=e, provider_config=provider_config) elif isinstance(transformed_request, str) or isinstance( transformed_request, bytes ): @@ -3304,7 +3323,15 @@ class BaseLLMHTTPHandler: input="", api_key="", additional_args={ - "complete_input_dict": transformed_request, + # A resumable upload config holds a reference to the (potentially + # huge) upload payload; logging deep-copies additional_args, so log + # a placeholder instead of re-materializing the payload. + "complete_input_dict": ( + "" + if isinstance(transformed_request, dict) + and "resumable_chunked_upload" in transformed_request + else transformed_request + ), "api_base": api_base, "headers": headers, }, @@ -3381,6 +3408,23 @@ class BaseLLMHTTPHandler: data=presigned_request["data"], timeout=timeout, ) + elif ( + isinstance(transformed_request, dict) + and "resumable_chunked_upload" in transformed_request + ): + try: + upload_response = await self._aresumable_chunked_upload( + client=async_httpx_client, + initiate_url=api_base, + base_headers=headers, + config=cast(Dict[str, Any], transformed_request)[ + "resumable_chunked_upload" + ], + timeout=timeout, + ) + except Exception as e: + verbose_logger.exception(f"Error creating file: {e}") + raise self._handle_error(e=e, provider_config=provider_config) elif isinstance(transformed_request, str) or isinstance( transformed_request, bytes ): @@ -3424,6 +3468,224 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, ) + # 8 MiB; a 256 KiB multiple, which GCS requires for every non-final chunk. + _RESUMABLE_CHUNK_SIZE = 8 * 1024 * 1024 + + @staticmethod + def _iter_resumable_chunks( + byte_iter: Iterator[bytes], chunk_size: int + ) -> Iterator[bytes]: + """Regroup a byte stream into ``chunk_size`` pieces, yielding a final + partial piece only when it is non-empty. Every full piece is exactly + ``chunk_size`` bytes (kept a 256 KiB multiple for GCS) and never more than + one chunk is buffered. An exactly chunk-aligned stream yields only full + chunks, so the upload finalizes on its last data chunk instead of making + an extra empty request; a 0-byte stream yields nothing and the caller + finalizes with a single empty request. + """ + buf = bytearray() + for piece in byte_iter: + buf.extend(piece) + while len(buf) >= chunk_size: + yield bytes(buf[:chunk_size]) + del buf[:chunk_size] + if buf: + yield bytes(buf) + + @staticmethod + def _resumable_content_range(offset: int, data_len: int, is_final: bool) -> str: + if not is_final: + return f"bytes {offset}-{offset + data_len - 1}/*" + total = offset + data_len + if data_len == 0: + return f"bytes */{total}" + return f"bytes {offset}-{total - 1}/{total}" + + @staticmethod + def _resumable_request_kwargs( + headers: dict, + content: bytes, + timeout: Optional[Union[float, httpx.Timeout]], + ) -> dict: + kwargs: Dict[str, Any] = {"headers": headers, "content": content} + if timeout is not None: + kwargs["timeout"] = timeout + return kwargs + + def _resumable_chunked_upload( + self, + *, + client: HTTPHandler, + initiate_url: str, + base_headers: dict, + config: dict, + timeout: Optional[Union[float, httpx.Timeout]], + ) -> httpx.Response: + """Open a GCS resumable session, then PUT the body in bounded chunks so a + large upload is never held in memory in full.""" + stream = config["body_stream"] + chunk_size = config.get("chunk_size", self._RESUMABLE_CHUNK_SIZE) + session_url_header = config.get("session_url_header", "location") + httpx_client = client.client + + init_headers = {**base_headers, **config.get("initiate_headers", {})} + init_req = httpx_client.build_request( + "POST", + initiate_url, + **self._resumable_request_kwargs(init_headers, b"", timeout), + ) + init_resp = httpx_client.send(init_req, follow_redirects=False) + init_resp.read() + if init_resp.status_code not in (200, 201): + init_resp.raise_for_status() + session_url = init_resp.headers.get(session_url_header) + if not session_url: + raise ValueError( + f"resumable upload: no session URL in '{session_url_header}' header" + ) + + offset = 0 + pending: Optional[bytes] = None + for chunk in self._iter_resumable_chunks(stream.iter_bytes(), chunk_size): + if pending is not None: + self._send_resumable_chunk( + httpx_client, + session_url, + base_headers, + pending, + offset, + is_final=False, + timeout=timeout, + ) + offset += len(pending) + pending = chunk + return self._send_resumable_chunk( + httpx_client, + session_url, + base_headers, + pending or b"", + offset, + is_final=True, + timeout=timeout, + ) + + def _send_resumable_chunk( + self, + httpx_client: httpx.Client, + url: str, + base_headers: dict, + data: bytes, + offset: int, + *, + is_final: bool, + timeout: Optional[Union[float, httpx.Timeout]], + ) -> httpx.Response: + headers = { + **base_headers, + "Content-Range": self._resumable_content_range(offset, len(data), is_final), + } + req = httpx_client.build_request( + "PUT", url, **self._resumable_request_kwargs(headers, data, timeout) + ) + resp = httpx_client.send(req, follow_redirects=False) + resp.read() + if resp.status_code not in ((200, 201) if is_final else (308,)): + # 4xx/5xx raise here; the ValueError catches an unexpected success + # status (e.g. a 200 where the protocol expects a 308 between chunks). + resp.raise_for_status() + raise ValueError(f"resumable upload: unexpected status {resp.status_code}") + return resp + + async def _aresumable_chunked_upload( + self, + *, + client: AsyncHTTPHandler, + initiate_url: str, + base_headers: dict, + config: dict, + timeout: Optional[Union[float, httpx.Timeout]], + ) -> httpx.Response: + stream = config["body_stream"] + chunk_size = config.get("chunk_size", self._RESUMABLE_CHUNK_SIZE) + session_url_header = config.get("session_url_header", "location") + httpx_client = client.client + + init_headers = {**base_headers, **config.get("initiate_headers", {})} + init_req = httpx_client.build_request( + "POST", + initiate_url, + **self._resumable_request_kwargs(init_headers, b"", timeout), + ) + init_resp = await httpx_client.send(init_req, follow_redirects=False) + await init_resp.aread() + if init_resp.status_code not in (200, 201): + init_resp.raise_for_status() + session_url = init_resp.headers.get(session_url_header) + if not session_url: + raise ValueError( + f"resumable upload: no session URL in '{session_url_header}' header" + ) + + offset = 0 + pending: Optional[bytes] = None + # Producing each chunk runs the synchronous per-row transform for that + # chunk's worth of rows. Pull it off the event loop thread so a large + # upload does not block other concurrent requests between PUTs. + chunk_iter = self._iter_resumable_chunks(stream.iter_bytes(), chunk_size) + done = object() + while True: + chunk = await asyncio.to_thread(next, chunk_iter, done) + if chunk is done: + break + if pending is not None: + await self._asend_resumable_chunk( + httpx_client, + session_url, + base_headers, + pending, + offset, + is_final=False, + timeout=timeout, + ) + offset += len(pending) + pending = chunk + return await self._asend_resumable_chunk( + httpx_client, + session_url, + base_headers, + pending or b"", + offset, + is_final=True, + timeout=timeout, + ) + + async def _asend_resumable_chunk( + self, + httpx_client: httpx.AsyncClient, + url: str, + base_headers: dict, + data: bytes, + offset: int, + *, + is_final: bool, + timeout: Optional[Union[float, httpx.Timeout]], + ) -> httpx.Response: + headers = { + **base_headers, + "Content-Range": self._resumable_content_range(offset, len(data), is_final), + } + req = httpx_client.build_request( + "PUT", url, **self._resumable_request_kwargs(headers, data, timeout) + ) + resp = await httpx_client.send(req, follow_redirects=False) + await resp.aread() + if resp.status_code not in ((200, 201) if is_final else (308,)): + # 4xx/5xx raise here; the ValueError catches an unexpected success + # status (e.g. a 200 where the protocol expects a 308 between chunks). + resp.raise_for_status() + raise ValueError(f"resumable upload: unexpected status {resp.status_code}") + return resp + def create_batch( self, create_batch_data: "CreateBatchRequest", diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index c31bfde69e7..176cfe98411 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -17,17 +17,13 @@ from litellm.litellm_core_utils.cloud_storage_security import ( ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.openai import ( - CreateFileRequest, FileContentRequest, HttpxBinaryResponseContent, - OpenAIFileObject, ) from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES -from .transformation import VertexAIFilesConfig, VertexAIJsonlFilesTransformation - -vertex_ai_files_transformation = VertexAIJsonlFilesTransformation() +from .transformation import VertexAIFilesConfig class VertexAIFilesHandler(GCSBucketBase): @@ -43,82 +39,6 @@ class VertexAIFilesHandler(GCSBucketBase): llm_provider=LlmProviders.VERTEX_AI, ) - async def async_create_file( - self, - create_file_data: CreateFileRequest, - api_base: Optional[str], - vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES], - vertex_project: Optional[str], - vertex_location: Optional[str], - timeout: Union[float, httpx.Timeout], - max_retries: Optional[int], - ) -> OpenAIFileObject: - gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config( - kwargs={} - ) - headers = await self.construct_request_headers( - vertex_instance=gcs_logging_config["vertex_instance"], - service_account_json=gcs_logging_config["path_service_account"], - ) - bucket_name = gcs_logging_config["bucket_name"] - ( - logging_payload, - object_name, - ) = vertex_ai_files_transformation.transform_openai_file_content_to_vertex_ai_file_content( - openai_file_content=create_file_data.get("file") - ) - gcs_upload_response = await self._log_json_data_on_gcs( - headers=headers, - bucket_name=bucket_name, - object_name=object_name, - logging_payload=logging_payload, - ) - - return vertex_ai_files_transformation.transform_gcs_bucket_response_to_openai_file_object( - create_file_data=create_file_data, - gcs_upload_response=gcs_upload_response, - ) - - def create_file( - self, - _is_async: bool, - create_file_data: CreateFileRequest, - api_base: Optional[str], - vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES], - vertex_project: Optional[str], - vertex_location: Optional[str], - timeout: Union[float, httpx.Timeout], - max_retries: Optional[int], - ) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]: - """ - Creates a file on VertexAI GCS Bucket - - Only supported for Async litellm.acreate_file - """ - - if _is_async: - return self.async_create_file( - create_file_data=create_file_data, - api_base=api_base, - vertex_credentials=vertex_credentials, - vertex_project=vertex_project, - vertex_location=vertex_location, - timeout=timeout, - max_retries=max_retries, - ) - else: - return asyncio.run( - self.async_create_file( - create_file_data=create_file_data, - api_base=api_base, - vertex_credentials=vertex_credentials, - vertex_project=vertex_project, - vertex_location=vertex_location, - timeout=timeout, - max_retries=max_retries, - ) - ) - def _extract_bucket_and_object_from_file_id( self, file_id: str, diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index f30518bc7ca..d5164d8c1c2 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -1,9 +1,21 @@ import base64 +import io +import itertools import json import os import re import time -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import ( + Any, + Callable, + Dict, + Iterable, + Iterator, + List, + Optional, + Tuple, + Union, +) import httpx from httpx import Headers, Response @@ -22,9 +34,13 @@ from litellm.litellm_core_utils.cloud_storage_security import ( validate_managed_cloud_file_id, ) from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + extract_file_metadata, +) from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( + BaseFileUploadStream, BaseFilesConfig, LiteLLMLoggingObj, ) @@ -44,8 +60,9 @@ from litellm.types.llms.openai import ( OpenAIFileObject, PathLike, ) +from litellm.types.files import ResumableChunkedUploadConfig from litellm.types.llms.vertex_ai import GcsBucketResponse -from litellm.types.utils import ExtractedFileData, LlmProviders, ModelResponse +from litellm.types.utils import LlmProviders, ModelResponse from ..common_utils import VertexAIError from ..vertex_llm_base import VertexBase @@ -137,42 +154,140 @@ def _get_litellm_batch_custom_id_from_labels(labels: Dict[str, Any]) -> str: return str(labels.get("litellm_custom_id", "unknown")) -def _openai_batch_jsonl_entries_to_vertex_wrapped_requests( - openai_jsonl_content: List[Dict[str, Any]], +def _openai_batch_jsonl_entry_to_vertex_wrapped_request( + openai_entry: Dict[str, Any], map_openai_to_vertex_params: Callable[[Dict[str, Any]], Dict[str, Any]], -) -> List[Dict[str, Any]]: +) -> Dict[str, Any]: """ - Transforms OpenAI JSONL batch entries to Vertex AI JSONL lines. + Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request. jsonl body for vertex is {"request": } Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} - {"request":{"contents": [{"role": "user", "parts": [{"text": "Describe what is happening in this video."}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/another_video.mov", "mimeType": "video/mov"}}]}]}} + """ + openai_request_body = openai_entry.get("body") or {} + vertex_request_body = _transform_request_body( + messages=openai_request_body.get("messages", []), + model=openai_request_body.get("model", ""), + optional_params=map_openai_to_vertex_params(openai_request_body), + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=None, + ) + + custom_id = openai_entry.get("custom_id") + if custom_id is not None: + if "labels" not in vertex_request_body: + vertex_request_body["labels"] = {} + _set_litellm_batch_custom_id_labels(vertex_request_body["labels"], custom_id) + + return {"request": vertex_request_body} + + +def _iter_stripped_lines(raw_lines: Iterable[Union[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() + if line: + yield line + + +def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[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. + """ + content: Any = openai_file_content + if isinstance(content, tuple): + content = content[1] + + if isinstance(content, (bytes, bytearray)): + # Scan for newlines in place so a large in-memory payload is not copied + # into a BytesIO just to iterate it line by line. + newline = ord("\n") + start, length = 0, len(content) + 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() + if line: + yield line + return + + if isinstance(content, str): + yield from _iter_stripped_lines(io.StringIO(content)) + return + + if isinstance(content, PathLike): + with open(str(content), "rb") as handle: + yield from _iter_stripped_lines(handle) + return + + if hasattr(content, "read"): + # The handle is read twice per upload (first-row probe for the GCS + # object name, then the body stream), so it must rewind to 0. A + # non-seekable handle would silently resume mid-stream and drop the + # already-consumed first row, so reject it loudly instead. + seek = getattr(content, "seek", None) + if seek is None: + raise ValueError( + "Batch upload file handle must be seekable; got a non-seekable " + "stream. Pass bytes, a path, or a seekable handle." + ) + try: + seek(0) + except (OSError, ValueError) as e: + raise ValueError( + "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) + return + + raise ValueError("Unsupported file content type") + + +def _iter_openai_jsonl_entries( + openai_file_content: FileTypes, +) -> Iterator[Dict[str, Any]]: + for line in _iter_openai_jsonl_lines(openai_file_content): + yield json.loads(line) + + +class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): + """Streams an OpenAI batch JSONL upload as Vertex-wrapped JSONL one row at a + time, so the transformed payload is never held in full. + + The transform runs lazily as the HTTP client pulls each chunk, which keeps + peak memory at one row regardless of how large the batch file is. """ - vertex_jsonl_content = [] - for _openai_jsonl_content in openai_jsonl_content: - openai_request_body = _openai_jsonl_content.get("body") or {} - vertex_request_body = _transform_request_body( - messages=openai_request_body.get("messages", []), - model=openai_request_body.get("model", ""), - optional_params=map_openai_to_vertex_params(openai_request_body), - custom_llm_provider="vertex_ai", - litellm_params={}, - cached_content=None, - ) + def __init__( + self, + openai_file_content: FileTypes, + map_openai_to_vertex_params: Callable[[Dict[str, Any]], Dict[str, Any]], + ) -> None: + self._openai_file_content = openai_file_content + self._map_openai_to_vertex_params = map_openai_to_vertex_params - # Add custom_id as a label for correlation in batch outputs - custom_id = _openai_jsonl_content.get("custom_id") - if custom_id is not None: - if "labels" not in vertex_request_body: - vertex_request_body["labels"] = {} - _set_litellm_batch_custom_id_labels( - vertex_request_body["labels"], custom_id + def _iter_vertex_jsonl_chunks(self) -> Iterator[bytes]: + first = True + for entry in _iter_openai_jsonl_entries(self._openai_file_content): + wrapped = _openai_batch_jsonl_entry_to_vertex_wrapped_request( + entry, self._map_openai_to_vertex_params ) + prefix = b"" if first else b"\n" + first = False + yield prefix + json.dumps(wrapped).encode("utf-8") - vertex_jsonl_content.append({"request": vertex_request_body}) - return vertex_jsonl_content + def iter_bytes(self) -> Iterator[bytes]: + return self._iter_vertex_jsonl_chunks() class VertexAIFilesConfig(VertexBase, BaseFilesConfig): @@ -181,7 +296,6 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ def __init__(self): - self.jsonl_transformation = VertexAIJsonlFilesTransformation() super().__init__() @property @@ -208,43 +322,6 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): headers["Authorization"] = f"Bearer {api_key}" return headers - def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str: - """ - Helper to extract content from various OpenAI file types and return as string. - - Handles: - - Direct content (str, bytes, IO[bytes]) - - Tuple formats: (filename, content, [content_type], [headers]) - - PathLike objects - """ - content: Union[str, bytes] = b"" - # Extract file content from tuple if necessary - if isinstance(openai_file_content, tuple): - # Take the second element which is always the file content - file_content = openai_file_content[1] - else: - file_content = openai_file_content - - # Handle different file content types - if isinstance(file_content, str): - # String content can be used directly - content = file_content - elif isinstance(file_content, bytes): - # Bytes content can be decoded - content = file_content - elif isinstance(file_content, PathLike): # PathLike - with open(str(file_content), "rb") as f: - content = f.read() - elif hasattr(file_content, "read"): # IO[bytes] - # File-like objects need to be read - content = file_content.read() - - # Ensure content is string - if isinstance(content, bytes): - content = content.decode("utf-8") - - return content - def _get_gcs_object_name_from_batch_jsonl( self, openai_jsonl_content: List[Dict[str, Any]], @@ -261,32 +338,21 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): object_name = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" return object_name - def get_object_name( - self, extracted_file_data: ExtractedFileData, purpose: str - ) -> str: + def get_object_name(self, file_data: FileTypes, purpose: str) -> str: """ - Get the object name for the request + Get the object name for the request. + + Reads only the first JSONL entry (streamed) for batch files, so a large + upload is never materialized just to derive the GCS object name. """ - extracted_file_data_content = extracted_file_data.get("content") - - if extracted_file_data_content is None: - raise ValueError("file content is required") - if purpose == "batch": - ## 1. If jsonl, check if there's a model name - file_content = self._get_content_from_openai_file( - extracted_file_data_content - ) - - # Split into lines and parse each line as JSON - openai_jsonl_content = [ - json.loads(line) for line in file_content.splitlines() if line.strip() - ] - if len(openai_jsonl_content) > 0: - return self._get_gcs_object_name_from_batch_jsonl(openai_jsonl_content) + ## 1. If jsonl, derive the object name from the first entry's model + first_entry = 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]) ## 2. If not jsonl, store under a server-generated managed object name - filename = extracted_file_data.get("filename") + filename, _ = extract_file_metadata(file_data) return build_managed_cloud_object_name( prefix=f"{VERTEX_AI_MANAGED_GCS_PREFIX}uploads/", filename=filename, @@ -294,7 +360,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): ) def _get_configured_bucket_name(self, litellm_params: Dict) -> str: - bucket_name = litellm_params.get("bucket_name") or os.getenv("GCS_BUCKET_NAME") + bucket_name = ( + litellm_params.get("gcs_bucket_name") + or litellm_params.get("bucket_name") + or os.getenv("GCS_BUCKET_NAME") + ) if not bucket_name: raise ValueError("GCS bucket_name is required") return bucket_name @@ -319,12 +389,21 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): raise ValueError("file is required") if purpose is None: raise ValueError("purpose is required") - extracted_file_data = extract_file_data(file_data) - object_name = self.get_object_name(extracted_file_data, purpose) + _, content_type = extract_file_metadata(file_data) + object_name = self.get_object_name(file_data, purpose) if object_prefix: object_name = f"{object_prefix}/{object_name}" encoded_object_name = encode_gcs_object_name_for_url(object_name) - endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_object_name}" + # Batch jsonl is streamed via a resumable session (bounded memory on + # large uploads); everything else is a single simple-media upload. + upload_type = ( + "resumable" + if FilesAPIUtils.is_batch_jsonl_request( + create_file_data=data, content_type=content_type + ) + else "media" + ) + endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType={upload_type}&name={encoded_object_name}" api_base = api_base or "https://storage.googleapis.com" if not api_base: raise ValueError("api_base is required") @@ -366,14 +445,6 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): ) return vertex_params - def _transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - self, openai_jsonl_content: List[Dict[str, Any]] - ) -> List[Dict[str, Any]]: - return _openai_batch_jsonl_entries_to_vertex_wrapped_requests( - openai_jsonl_content=openai_jsonl_content, - map_openai_to_vertex_params=self._map_openai_to_vertex_params, - ) - def transform_create_file_request( self, model: str, @@ -384,40 +455,34 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ 2 Cases: 1. Handle basic file upload - 2. Handle batch file upload (.jsonl) + 2. Handle batch file upload (.jsonl), streamed to a GCS resumable + session so large uploads stay memory-bounded. """ file_data = create_file_data.get("file") if file_data is None: raise ValueError("file is required") - extracted_file_data = extract_file_data(file_data) - extracted_file_data_content = extracted_file_data.get("content") - if extracted_file_data_content is None: - raise ValueError("file content is required") - - if FilesAPIUtils.is_batch_jsonl_file( + _, content_type = extract_file_metadata(file_data) + if FilesAPIUtils.is_batch_jsonl_request( create_file_data=create_file_data, - extracted_file_data=extracted_file_data, + content_type=content_type, ): - ## 1. If jsonl, check if there's a model name - file_content = self._get_content_from_openai_file( - extracted_file_data_content - ) - - # Split into lines and parse each line as JSON - openai_jsonl_content = [ - json.loads(line) for line in file_content.splitlines() if line.strip() - ] - vertex_jsonl_content = ( - self._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_jsonl_content + return { + "resumable_chunked_upload": ResumableChunkedUploadConfig( + body_stream=_OpenAIToVertexBatchUploadStream( + file_data, + self._map_openai_to_vertex_params, + ), + initiate_headers={ + "X-Upload-Content-Type": "application/json", + }, ) - ) - return "\n".join(json.dumps(item) for item in vertex_jsonl_content) - elif isinstance(extracted_file_data_content, bytes): + } + + extracted_file_data_content = extract_file_data(file_data).get("content") + if isinstance(extracted_file_data_content, bytes): return extracted_file_data_content - else: - raise ValueError("Unsupported file content type") + raise ValueError("Unsupported file content type") def transform_create_file_response( self, @@ -642,39 +707,38 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): } """ try: - # Decode content - content_str = content.decode("utf-8") - - # Check if it's JSONL (multiple lines) - lines = content_str.strip().split("\n") - if not lines: + # Read the result file one row at a time. Batch output files can be + # 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 = _iter_openai_jsonl_lines(content) + try: + first_line = next(lines) + except StopIteration: return content - # Try to parse the first line to see if it's Vertex AI batch output - first_line = json.loads(lines[0]) - - # Check if it has Vertex AI batch output structure with discriminating fields - # Must have request, response, and processed_time - # Plus either candidates (success) or status (error) - has_base_structure = ( - "response" in first_line - and "request" in first_line - and "processed_time" in first_line + # Identify a Vertex AI batch output from the first row's + # discriminating fields. Anything else (e.g. a binary file whose + # first line is not valid UTF-8/JSON) raises and falls through to the + # passthrough below, leaving the content untouched. + first_row = json.loads(first_line) + is_vertex_batch_output = ( + "request" in first_row + and "response" in first_row + and "processed_time" in first_row + and ( + "candidates" in first_row.get("response", {}) + or "promptFeedback" in first_row.get("response", {}) + or bool(first_row.get("status")) + ) ) - has_success_or_error = ( - "candidates" in first_line.get("response", {}) - or "promptFeedback" in first_line.get("response", {}) - or bool(first_line.get("status")) - ) - - if not (has_base_structure and has_success_or_error): - # Not a Vertex AI batch output, return as-is + if not is_vertex_batch_output: return content vertex_gemini_config = VertexGeminiConfig() - # Always use a fresh local Logging object for the per-line transformation - # so we never mutate the caller's logging_obj (which already went through - # pre_call and has its own model/start_time/optional_params set). + # Use a fresh Logging object for the per-row transform so we never + # mutate the caller's (which already ran pre_call with its own + # model/start_time/optional_params). batch_transform_logging_obj = Logging( model="", messages=[], @@ -691,29 +755,27 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): request=httpx.Request(method="POST", url="https://example.com"), ) - # Transform all lines - transformed_lines = [] - for line in lines: - if not line.strip(): - continue - + # Transform each row straight into the output buffer, so peak memory + # stays at ~one row plus the output. If any row fails, return the + # original content unchanged. + output = bytearray() + for line in itertools.chain([first_line], lines): try: - vertex_output = json.loads(line) openai_output = ( self._transform_single_vertex_batch_output_to_openai( - vertex_output=vertex_output, + vertex_output=json.loads(line), vertex_gemini_config=vertex_gemini_config, logging_obj=batch_transform_logging_obj, mock_httpx_response=mock_httpx_response, ) ) - transformed_lines.append(json.dumps(openai_output)) except Exception: - # If any line fails, return original content return content + if output: + output += b"\n" + output += json.dumps(openai_output).encode("utf-8") - # Return transformed content - return "\n".join(transformed_lines).encode("utf-8") + return bytes(output) except Exception: # If anything fails, return original content @@ -795,137 +857,3 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): "message": f"Failed to transform response: {str(e)}", }, } - - -class VertexAIJsonlFilesTransformation(VertexGeminiConfig): - """ - Transforms OpenAI /v1/files/* requests to VertexAI /v1/files/* requests - """ - - def transform_openai_file_content_to_vertex_ai_file_content( - self, openai_file_content: Optional[FileTypes] = None - ) -> Tuple[str, str]: - """ - Transforms OpenAI FileContentRequest to VertexAI FileContentRequest - """ - - if openai_file_content is None: - raise ValueError("contents of file are None") - # Read the content of the file - file_content = self._get_content_from_openai_file(openai_file_content) - - # Split into lines and parse each line as JSON - openai_jsonl_content = [ - json.loads(line) for line in file_content.splitlines() if line.strip() - ] - vertex_jsonl_content = ( - self._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_jsonl_content - ) - ) - vertex_jsonl_string = "\n".join( - json.dumps(item) for item in vertex_jsonl_content - ) - object_name = self._get_gcs_object_name( - openai_jsonl_content=openai_jsonl_content - ) - return vertex_jsonl_string, object_name - - def _transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - self, openai_jsonl_content: List[Dict[str, Any]] - ) -> List[Dict[str, Any]]: - return _openai_batch_jsonl_entries_to_vertex_wrapped_requests( - openai_jsonl_content=openai_jsonl_content, - map_openai_to_vertex_params=self._map_openai_to_vertex_params, - ) - - def _get_gcs_object_name( - self, - openai_jsonl_content: List[Dict[str, Any]], - ) -> str: - """ - Gets a unique GCS object name for the VertexAI batch prediction job - - named as: litellm-vertex-{model}-{uuid} - """ - _model = openai_jsonl_content[0].get("body", {}).get("model", "") - if "publishers/google/models" not in _model: - _model = f"publishers/google/models/{_model}" - safe_model_path = sanitize_cloud_object_path(_model, fallback="model") - object_name = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" - return object_name - - def _map_openai_to_vertex_params( - self, - openai_request_body: Dict[str, Any], - ) -> Dict[str, Any]: - """ - wrapper to call VertexGeminiConfig.map_openai_params - """ - _model = openai_request_body.get("model", "") - vertex_params = self.map_openai_params( - model=_model, - non_default_params=openai_request_body, - optional_params={}, - drop_params=False, - ) - return vertex_params - - def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str: - """ - Helper to extract content from various OpenAI file types and return as string. - - Handles: - - Direct content (str, bytes, IO[bytes]) - - Tuple formats: (filename, content, [content_type], [headers]) - - PathLike objects - """ - content: Union[str, bytes] = b"" - # Extract file content from tuple if necessary - if isinstance(openai_file_content, tuple): - # Take the second element which is always the file content - file_content = openai_file_content[1] - else: - file_content = openai_file_content - - # Handle different file content types - if isinstance(file_content, str): - # String content can be used directly - content = file_content - elif isinstance(file_content, bytes): - # Bytes content can be decoded - content = file_content - elif isinstance(file_content, PathLike): # PathLike - with open(str(file_content), "rb") as f: - content = f.read() - elif hasattr(file_content, "read"): # IO[bytes] - # File-like objects need to be read - content = file_content.read() - - # Ensure content is string - if isinstance(content, bytes): - content = content.decode("utf-8") - - return content - - def transform_gcs_bucket_response_to_openai_file_object( - self, create_file_data: CreateFileRequest, gcs_upload_response: Dict[str, Any] - ) -> OpenAIFileObject: - """ - Transforms GCS Bucket upload file response to OpenAI FileObject - """ - gcs_id = gcs_upload_response.get("id", "") - # Remove the last numeric ID from the path - gcs_id = "/".join(gcs_id.split("/")[:-1]) if gcs_id else "" - - return OpenAIFileObject( - purpose=create_file_data.get("purpose", "batch"), - id=f"gs://{gcs_id}", - filename=gcs_upload_response.get("name", ""), - created_at=_convert_vertex_datetime_to_openai_datetime( - vertex_datetime=gcs_upload_response.get("timeCreated", "") - ), - status="uploaded", - bytes=gcs_upload_response.get("size", 0), - object="file", - ) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 561b6ec8a77..0f4dcbdbfd1 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -17,17 +17,19 @@ Quick summary: - async_log_success_event() fires on GET /v1/batches/{id} (batch completion) """ -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Literal, Optional, Union from fastapi import HTTPException from pydantic import BaseModel +import json + import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import ( - _get_batch_job_input_file_usage, - _get_file_content_as_dictionary, - _get_models_from_batch_input_file_content, + _count_entry_tokens, + _estimate_batch_entry_tokens, + _iter_batch_input_lines, ) from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth @@ -233,6 +235,8 @@ class _PROXY_BatchRateLimiter(CustomLogger): # Managed files require bypassing the HTTP endpoint (which runs access-check hooks) # and calling the managed files hook directly with the user's credentials. is_managed_file = _is_base64_encoded_unified_file_id(file_id) + # For managed files the unified file id encodes the proxy model + # alias(es) the file was uploaded for; auth validates against those. target_model_names = ( get_models_from_unified_file_id(is_managed_file) if is_managed_file @@ -251,7 +255,44 @@ class _PROXY_BatchRateLimiter(CustomLogger): user_api_key_dict=user_api_key_dict, ) - file_content_as_dict = _get_file_content_as_dictionary(file_content.content) + file_content_bytes = getattr(file_content, "content", None) + if not isinstance(file_content_bytes, bytes): + raise ValueError( + f"Expected bytes content from file retrieval for {file_id}, " + f"got {type(file_content_bytes)}" + ) + + # Single streaming pass over the JSONL lines, accounting each row + # independently. One bad row can never abort the pass: a malformed + # line is skipped (its request can't run upstream anyway) and a row + # the token counter can't measure falls back to a conservative + # size-based estimate. This guarantees two things a restricted caller + # must not be able to break by crafting a row that raises: + # 1. The allowlist check below always sees every parseable + # ``body.model`` (the loop never stops early), so models can't be + # smuggled in after a bad row. + # 2. The token total is never silently zeroed, so the TPM limit + # can't be evaded by sending uncountable rows. + # Counting stays best-effort, so a legitimate (e.g. multimodal) row + # the counter can't measure is estimated, not hard-rejected. + models: set = set() + total_tokens = 0 + request_count = 0 + for raw_line in _iter_batch_input_lines(file_content_bytes): + request_count += 1 + try: + entry = json.loads(raw_line) + except Exception: + total_tokens += _estimate_batch_entry_tokens(raw_line) + continue + if isinstance(entry, dict): + model = (entry.get("body") or {}).get("model") + if model: + models.add(model) + try: + total_tokens += _count_entry_tokens(entry) + except Exception: + total_tokens += _estimate_batch_entry_tokens(raw_line) # Validate every model named in the batch JSONL against the # caller's per-key model allowlist. Without this, a caller @@ -261,17 +302,12 @@ class _PROXY_BatchRateLimiter(CustomLogger): if user_api_key_dict is not None: await self._enforce_batch_file_model_access( user_api_key_dict=user_api_key_dict, - file_content_as_dict=file_content_as_dict, + models=models, target_model_names=target_model_names or None, ) - input_file_usage = _get_batch_job_input_file_usage( - file_content_dictionary=file_content_as_dict, - custom_llm_provider=custom_llm_provider, - ) - request_count = len(file_content_as_dict) return BatchFileUsage( - total_tokens=input_file_usage.total_tokens, + total_tokens=total_tokens, request_count=request_count, ) @@ -297,14 +333,15 @@ class _PROXY_BatchRateLimiter(CustomLogger): async def _enforce_batch_file_model_access( self, user_api_key_dict: UserAPIKeyAuth, - file_content_as_dict: List[dict], + models: Optional[Iterable[str]] = None, target_model_names: Optional[List[str]] = None, ) -> None: """Reject the batch if the caller is not authorized for the upload target. For managed files, ``target_model_names`` (from the unified file id) is - the proxy alias the file was uploaded for and is used directly for auth. - For legacy/non-managed files, falls back to ``body.model`` values in the JSONL. + the proxy alias the file was uploaded for and is checked directly. + Otherwise the ``body.model`` values collected from the JSONL (``models``) + are checked. Reuses ``can_key_call_model`` so the same allowlist semantics (wildcards, access groups, ``all-proxy-models``, team aliases) @@ -315,10 +352,9 @@ class _PROXY_BatchRateLimiter(CustomLogger): if target_model_names: models = target_model_names - else: - models = _get_models_from_batch_input_file_content(file_content_as_dict) - if not models: - return + + if not models: + return llm_model_list = llm_router.model_list if llm_router is not None else None for model in models: diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 378cbbda89c..1d2cc726154 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -7,7 +7,7 @@ import asyncio import traceback -from typing import Any, Optional, cast, get_args +from typing import Any, BinaryIO, Optional, Union, cast, get_args import httpx from fastapi import ( @@ -91,16 +91,18 @@ def get_files_provider_config( return None -def get_first_json_object(file_content_bytes: bytes) -> Optional[dict]: +def get_first_json_object(file_source: Union[bytes, BinaryIO]) -> Optional[dict]: try: - # Decode the bytes to a string and split into lines - file_content = file_content_bytes.decode("utf-8") - first_line = file_content.splitlines()[0].strip() - - # Parse the JSON object from the first line - json_object = json.loads(first_line) - return json_object - except (json.JSONDecodeError, UnicodeDecodeError): + if isinstance(file_source, (bytes, bytearray)): + newline = file_source.find(b"\n") + raw = file_source if newline == -1 else file_source[:newline] + first_line = raw.decode("utf-8") + else: + file_source.seek(0) + first_line = file_source.readline().decode("utf-8") + file_source.seek(0) + return json.loads(first_line.strip()) + except (json.JSONDecodeError, UnicodeDecodeError, OSError, ValueError): return None @@ -321,9 +323,15 @@ async def create_file( # noqa: PLR0915 data: Dict = {} try: - # Use orjson to parse JSON data, orjson speeds up requests significantly - # Read the file content - file_content = await file.read() + # Batch uploads can be gigabytes. Starlette has already spooled the upload + # to disk, so stream from that handle instead of reading it into memory. + # Other uploads are small and stay in-memory bytes. + file_source: Union[bytes, BinaryIO] + if purpose == "batch": + await file.seek(0) + file_source = file.file + else: + file_source = await file.read() custom_llm_provider = ( provider or get_custom_llm_provider_from_request_headers(request=request) @@ -443,13 +451,13 @@ async def create_file( # noqa: PLR0915 ) # Prepare the file data according to FileTypes - file_data = (file.filename, file_content, file.content_type) + file_data = (file.filename, file_source, file.content_type) ## check if model is a loadbalanced model router_model: Optional[str] = None is_router_model = False if litellm.enable_loadbalancing_on_batch_endpoints is True: - json_obj = get_first_json_object(file_content_bytes=file_content) + json_obj = get_first_json_object(file_source) if json_obj: router_model = get_model_from_json_obj(json_object=json_obj) is_router_model = is_known_model( diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index 5e58479825b..ddec753d362 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -82,43 +82,83 @@ def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> File if isinstance(file_content, PathLike): return file_content - # Decode the bytes to a string and split into lines - # If file_content is a file-like object, read the bytes - if hasattr(file_content, "read"): - file_content_bytes = file_content.read() # type: ignore - elif isinstance(file_content, tuple): - file_content_bytes = file_content[1] - else: - file_content_bytes = file_content - - # Decode the bytes to a string and split into lines - if isinstance(file_content_bytes, bytes): - file_content_str = file_content_bytes.decode("utf-8") - elif isinstance(file_content_bytes, str): - file_content_str = file_content_bytes + # Iterate the source line-by-line WITHOUT reading it all into memory. A + # 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 = file_content[1] if isinstance(file_content, tuple) else file_content + if hasattr(source, "read"): + if hasattr(source, "seek"): + try: + source.seek(0) # type: ignore[attr-defined] + 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) else: return file_content - # Parse JSONL properly, handling potential multiline JSON objects - json_objects = parse_jsonl_with_embedded_newlines(file_content_str) + # Rewrite one row at a time, writing straight into the output buffer + # instead of holding every parsed row in a list. Peak memory stays at + # ~one row plus the output rather than several full copies of the file, + # which the managed-files path depends on (it re-runs this rewrite once + # per target model). Lines are accumulated so JSON objects that span + # 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). + output = InMemoryFile( + b"", name="modified_file.jsonl", content_type="application/jsonl" + ) + wrote_any = False + buffer = "" + for raw_line in line_iter: # type: ignore[attr-defined] + buffer += ( + raw_line.decode("utf-8") + if isinstance(raw_line, (bytes, bytearray)) + else raw_line + ) + stripped = buffer.strip() + if not stripped: + buffer = "" + continue + try: + json_object = json.loads(stripped) + except json.JSONDecodeError: + 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 + buffer = "" + + if buffer.strip(): + # A row never parsed (truncated/malformed, or it swallowed the rows + # that followed it). Returning the partial `output` would silently + # drop those rows; return the unchanged original so the provider + # rejects the batch loudly instead of accepting a truncated one. + verbose_logger.error( + f"error parsing trailing batch content: {buffer[:100]}..." + ) + if hasattr(source, "seek"): + try: + source.seek(0) # type: ignore[attr-defined] + except (OSError, ValueError): + pass + return file_content # If no valid JSON objects were found, return the original content - if len(json_objects) == 0: + if not wrote_any: return file_content - modified_lines = [] - for json_object in json_objects: - # Replace the model name if it exists - if "body" in json_object: - json_object["body"]["model"] = new_model_name - - # Convert the modified JSON object back to a string - modified_lines.append(json.dumps(json_object)) - - # Reassemble the modified lines and return as bytes - modified_file_content = "\n".join(modified_lines).encode("utf-8") - - return InMemoryFile(modified_file_content, name="modified_file.jsonl", content_type="application/jsonl") # type: ignore + output.seek(0) + return output # type: ignore except (json.JSONDecodeError, UnicodeDecodeError, TypeError): # return the original file content if there is an error replacing the model name diff --git a/litellm/types/files.py b/litellm/types/files.py index bf56894329c..1b2d7e30f1f 100644 --- a/litellm/types/files.py +++ b/litellm/types/files.py @@ -321,3 +321,21 @@ class TwoStepFileUploadConfig(TypedDict, total=False): upload_request: Required[TwoStepFileUploadRequest] upload_url_location: Required[Literal["headers", "body"]] upload_url_key: str + + +class ResumableChunkedUploadConfig(TypedDict, total=False): + """Drives a memory-bounded resumable upload (GCS JSON API). + + The handler POSTs to the upload URL to open a session, reads the session URI + from ``session_url_header``, then PUTs ``body_stream`` to that URI in + ``chunk_size``-byte chunks (a 256 KiB multiple) using Content-Range, so the + payload is never buffered in full and the transfer is resumable. + + ``body_stream`` is a ``BaseFileUploadStream``; it is typed ``Any`` here to + avoid importing the llms layer into types. + """ + + body_stream: Required[Any] + chunk_size: int + session_url_header: str + initiate_headers: Dict[str, str] diff --git a/litellm/types/router.py b/litellm/types/router.py index ef7eb05d087..c36f064a91e 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -173,6 +173,9 @@ class CredentialLiteLLMParams(BaseModel): ## UNIFIED PROJECT/REGION ## region_name: Optional[str] = None + ## OBJECT STORAGE (files / batches) ## + gcs_bucket_name: Optional[str] = None + ## AWS BEDROCK / SAGEMAKER ## aws_access_key_id: Optional[str] = None aws_secret_access_key: Optional[str] = None diff --git a/tests/batches_tests/test_openai_batches_and_files.py b/tests/batches_tests/test_openai_batches_and_files.py index 2e89381cba9..6fd6b777952 100644 --- a/tests/batches_tests/test_openai_batches_and_files.py +++ b/tests/batches_tests/test_openai_batches_and_files.py @@ -31,7 +31,7 @@ from litellm.types.utils import StandardLoggingPayload import random import socket import httpx -from unittest.mock import patch, MagicMock +from unittest.mock import patch, MagicMock, AsyncMock def _can_resolve_openai(): @@ -511,10 +511,26 @@ async def test_avertex_batch_prediction(monkeypatch): mock_response.status_code = 200 return mock_response - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - side_effect=mock_side_effect, - ) as mock_global_post: + # Batch jsonl file creation now streams to a GCS resumable session via + # _aresumable_chunked_upload (httpx send), not AsyncHTTPHandler.post, so mock + # that entry point to return the GCS object response. The resumable protocol + # itself is covered in test_vertex_ai_files_streaming.py. + mock_upload_response = httpx.Response( + 200, + json=mock_file_response, + request=httpx.Request("PUT", "https://storage.googleapis.com/upload"), + ) + with ( + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + side_effect=mock_side_effect, + ) as mock_global_post, + patch( + "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler._aresumable_chunked_upload", + new_callable=AsyncMock, + return_value=mock_upload_response, + ), + ): litellm.set_verbose = True litellm._turn_on_debug() file_name = "vertex_batch_completions.jsonl" diff --git a/tests/router_unit_tests/test_router_batch_utils.py b/tests/router_unit_tests/test_router_batch_utils.py index 1b8f713a437..b8760906645 100644 --- a/tests/router_unit_tests/test_router_batch_utils.py +++ b/tests/router_unit_tests/test_router_batch_utils.py @@ -1,17 +1,10 @@ import sys import os -import traceback -from dotenv import load_dotenv -from fastapi import Request -from datetime import datetime sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -from litellm import Router import pytest -import litellm -from unittest.mock import patch, MagicMock, AsyncMock import json from io import BytesIO @@ -76,6 +69,29 @@ def test_tuple_input(sample_jsonl_bytes): assert result.content_type == "application/jsonl" +def test_tuple_with_file_handle_rewrites_model(sample_jsonl_bytes): + """Security regression: when the tuple's content element is a file handle + (batch uploads stream from the spooled upload handle), the model must still + be rewritten. Otherwise a restricted body.model survives unmodified and + bypasses the batch model allowlist, which only checks the upload target.""" + new_model = "approved-target-model" + handle = BytesIO(sample_jsonl_bytes) + test_tuple = ("test.jsonl", handle, "application/json") + + result = replace_model_in_jsonl(test_tuple, new_model) + + assert isinstance(result, InMemoryFile) + rows = [ + json.loads(line) + for line in result.getvalue().decode("utf-8").splitlines() + if line.strip() + ] + assert rows, "rewrite must produce rows" + # every row now carries the rewritten target, not the original (restricted) model + assert all(row["body"]["model"] == new_model for row in rows) + assert all(row["body"]["model"] != "gpt-5.5" for row in rows) + + def test_file_like_object(sample_file_like): """Test with file-like object input""" new_model = "claude-3" @@ -129,9 +145,9 @@ def test_should_replace_model_in_jsonl(): """Test that should_replace_model_in_jsonl returns the correct value""" from litellm.router_utils.batch_utils import should_replace_model_in_jsonl - assert should_replace_model_in_jsonl(purpose="batch") == True - assert should_replace_model_in_jsonl(purpose="test") == False - assert should_replace_model_in_jsonl(purpose="user_data") == False + assert should_replace_model_in_jsonl(purpose="batch") is True + assert should_replace_model_in_jsonl(purpose="test") is False + assert should_replace_model_in_jsonl(purpose="user_data") is False def test_parse_jsonl_with_embedded_newlines_simple(): @@ -217,6 +233,63 @@ def test_parse_jsonl_with_embedded_newlines_whitespace_only(): assert len(result) == 0 +def test_replace_model_in_jsonl_malformed_middle_row_returns_original(): + """Regression: a malformed/truncated middle row must not silently drop the + rows that follow it. The streaming rewrite accumulates physical lines into a + buffer; a row that never parses poisons the buffer so every later valid row + is concatenated into it and dropped. Returning that partial rewrite would + ship a truncated batch with no error to the caller. Instead the original + content is returned unchanged so the provider rejects the bad batch loudly.""" + content = ( + b'{"custom_id":"a","body":{"model":"x"}}\n' + b'{"custom_id":"b","body":{"model":\n' # truncated, never completes + b'{"custom_id":"c","body":{"model":"x"}}\n' + ) + + result = replace_model_in_jsonl(content, "new-model") + + assert ( + result == content + ), "must return the original unchanged, not a partial rewrite" + + +def test_replace_model_in_jsonl_malformed_row_seekable_handle_rewound(): + """When the source is a seekable handle that gets consumed during the failed + rewrite, it must be rewound to 0 so the caller can re-read the full original.""" + content = ( + b'{"custom_id":"a","body":{"model":"x"}}\n' + b'{"custom_id":"b","body":{"model":\n' + b'{"custom_id":"c","body":{"model":"x"}}\n' + ) + handle = BytesIO(content) + + result = replace_model_in_jsonl(handle, "new-model") + + assert result is handle + assert handle.read() == content, "handle must be rewound for the caller to re-read" + + +def test_replace_model_in_jsonl_multi_row_rewrites_every_model(): + """Happy path: a well-formed multi-row file gets every row's model rewritten + and no row is dropped.""" + content = ( + b'{"custom_id":"a","body":{"model":"old1"}}\n' + b'{"custom_id":"b","body":{"model":"old2"}}\n' + b'{"custom_id":"c","body":{"model":"old3"}}\n' + ) + + result = replace_model_in_jsonl(content, "new-model") + + assert isinstance(result, InMemoryFile) + rows = [ + json.loads(line) + for line in result.getvalue().decode("utf-8").splitlines() + if line.strip() + ] + assert [row["custom_id"] for row in rows] == ["a", "b", "c"] + assert all(row["body"]["model"] == "new-model" for row in rows) + + 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/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py index d4586134b13..122518d4acb 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py @@ -8,8 +8,8 @@ Regression test for: UTF-8 codec error when uploading binary files """ import io +import json import pytest -from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -137,11 +137,11 @@ class TestVertexAIBinaryFileUpload: ), "Binary file data should remain as bytes" @pytest.mark.asyncio - async def test_jsonl_file_upload_returns_string(self): + async def test_jsonl_file_upload_returns_resumable_stream(self): """ - Test that JSONL files (text) are correctly transformed to strings. - - This ensures we handle both binary and text files correctly. + Test that JSONL batch files are transformed into a resumable-upload config + carrying a streaming body (not a buffered bytes payload), so the handler + can stream the upload to GCS in bounded chunks. """ # Create mock JSONL content mock_jsonl_content = ( @@ -164,10 +164,16 @@ class TestVertexAIBinaryFileUpload: litellm_params={}, ) - # JSONL files should be transformed to string - assert isinstance( - transformed_request, str - ), f"Expected string for JSONL file, got {type(transformed_request)}" + assert ( + isinstance(transformed_request, dict) + and "resumable_chunked_upload" in transformed_request + ), f"Expected a resumable upload config for JSONL, got {type(transformed_request)}" + + stream = transformed_request["resumable_chunked_upload"]["body_stream"] + decoded = json.loads(b"".join(stream.iter_bytes()).decode("utf-8")) + assert ( + "request" in decoded + ), "JSONL transform must wrap each row in {'request': ...}" @pytest.mark.asyncio async def test_mixed_file_types_in_sequence(self): @@ -208,7 +214,7 @@ class TestVertexAIBinaryFileUpload: optional_params={}, litellm_params={}, ) - assert isinstance(result2, str) + assert isinstance(result2, dict) and "resumable_chunked_upload" in result2 # Test 3: Upload another binary file binary_content2 = b"\xc4\xe5\xf2\xe5\xeb" @@ -251,7 +257,7 @@ class TestVertexAIBinaryFileUpload: }, "text_files": { "input_type": "str or bytes", - "output_type": "str", + "output_type": "bytes", "examples": ["JSONL", "CSV", "TXT"], "http_method": "POST", "encoding": "UTF-8", diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py new file mode 100644 index 00000000000..cd556c48b6b --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -0,0 +1,696 @@ +""" +Tests for the streaming OpenAI -> Vertex JSONL batch transform. + +The transform converts batch uploads entry-by-entry rather than materializing +the payload in full intermediate lists (decoded str, parsed dicts, transformed +dicts, joined output), which keeps peak memory bounded on large uploads. + +These tests lock in the behaviour that would regress if the streaming path were +replaced by a list-based pipeline: + 1. Byte-for-byte output parity with a list pipeline (wire format). + 2. The streaming transform peaks at a clear fraction of a list pipeline on the + same input (relative differential, robust to GC noise). + 3. ``get_object_name`` only parses the first JSONL row, so a payload whose + later rows are not valid JSON does not raise. + 4. A tuple-wrapped file handle uploaded through the real create_file ordering + keeps every row, including entry 0 (no partial upload from a consumed + cursor). +""" + +import gc +import io +import json +import time +import tracemalloc + +import httpx +import pytest + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.base_llm.files.transformation import BaseFileUploadStream +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.vertex_ai.files.transformation import ( + VertexAIFilesConfig, + _OpenAIToVertexBatchUploadStream, + _get_litellm_batch_custom_id_from_labels, + _iter_openai_jsonl_entries, + _iter_openai_jsonl_lines, + _openai_batch_jsonl_entry_to_vertex_wrapped_request, +) +from litellm.types.llms.openai import CreateFileRequest + + +def _resumable_stream(transformed) -> BaseFileUploadStream: + """Pull the streaming body out of a resumable-upload transform result.""" + return transformed["resumable_chunked_upload"]["body_stream"] + + +def _join_upload_body(transformed) -> bytes: + """Materialize a transform result's upload body for byte-level assertions.""" + if isinstance(transformed, dict) and "resumable_chunked_upload" in transformed: + return b"".join(_resumable_stream(transformed).iter_bytes()) + if isinstance(transformed, BaseFileUploadStream): + return b"".join(transformed.iter_bytes()) + if isinstance(transformed, str): + return transformed.encode("utf-8") + return transformed + + +def _make_openai_jsonl_bytes(n_rows: int, padding: int = 400) -> bytes: + pad = "x" * padding + rows = [] + for i in range(n_rows): + rows.append( + json.dumps( + { + "custom_id": f"request-{i}", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.5-flash", + "messages": [{"role": "user", "content": f"{pad} {i}"}], + "max_tokens": 4, + }, + } + ) + ) + return ("\n".join(rows)).encode("utf-8") + + +def _reference_vertex_jsonl_string(cfg: VertexAIFilesConfig, content: str) -> str: + """Row-by-row reference output built eagerly from the live single-entry + transform, so the streaming path can be checked against it for parity.""" + entries = [json.loads(line) for line in content.splitlines() if line.strip()] + return "\n".join( + json.dumps( + _openai_batch_jsonl_entry_to_vertex_wrapped_request( + entry, cfg._map_openai_to_vertex_params + ) + ) + for entry in entries + ) + + +class TestStreamingOutputParity: + def test_transform_create_file_request_returns_resumable_stream_parity(self): + cfg = VertexAIFilesConfig() + raw = _make_openai_jsonl_bytes(300) + request: CreateFileRequest = { + "file": ("batch.jsonl", raw, "application/jsonl"), + "purpose": "batch", + } + + out = cfg.transform_create_file_request( + model="", create_file_data=request, optional_params={}, litellm_params={} + ) + + # A batch upload must be a resumable-upload config carrying a streaming + # body, so the handler can chunk it; a buffered bytes/str return would + # defeat the OOM fix. + assert isinstance(out, dict) and "resumable_chunked_upload" in out + assert isinstance(_resumable_stream(out), BaseFileUploadStream) + assert _join_upload_body(out).decode("utf-8") == _reference_vertex_jsonl_string( + cfg, raw.decode("utf-8") + ) + + +class TestFileLikeInputNotPartiallyConsumed: + """ + In ``llm_http_handler.create_file`` the object-name step + (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`` + 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 + locks the full-payload invariant in. + """ + + def test_filehandle_create_file_keeps_first_entry(self): + cfg = VertexAIFilesConfig() + n_rows = 25 + raw = _make_openai_jsonl_bytes(n_rows) + create_file_data: CreateFileRequest = { + "file": ("batch.jsonl", io.BytesIO(raw), "application/jsonl"), + "purpose": "batch", + } + + # Object-name step first (as the handler does), then the transform, both + # reading the same live BytesIO handle. + cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=create_file_data, + ) + out = cfg.transform_create_file_request( + model="", + create_file_data=create_file_data, + optional_params={}, + litellm_params={}, + ) + + lines = _join_upload_body(out).decode("utf-8").splitlines() + assert len(lines) == n_rows, "no batch row may be dropped from the upload" + first_labels = json.loads(lines[0])["request"]["labels"] + assert _get_litellm_batch_custom_id_from_labels(first_labels) == "request-0" + + +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}'] + + 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)] == [ + {"a": 1}, + {"b": 2}, + ] + + def test_accepts_str_bytes_tuple_and_filelike(self): + expected = [{"a": 1}, {"b": 2}] + text = '{"a": 1}\n{"b": 2}\n' + for source in ( + text, + text.encode("utf-8"), + ("name.jsonl", text.encode("utf-8"), "application/jsonl"), + io.BytesIO(text.encode("utf-8")), + ): + 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}', + ] + + 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}] + + 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] + + 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 + # handle can't rewind, so it must fail loudly rather than silently resume + # mid-stream and omit the opening batch request. + class _NonSeekable: + def __init__(self, raw: bytes): + self._buf = io.BytesIO(raw) + + def read(self, *args): + return self._buf.read(*args) + + def __iter__(self): + return iter(self._buf) + + def seek(self, *args): + raise io.UnsupportedOperation("not seekable") + + 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)) + + 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): + next(gen) + + +class TestGetObjectNameLazyParse: + def test_only_parses_first_row_for_model(self): + cfg = VertexAIFilesConfig() + # Tail rows are deliberately not valid JSON. Parsing the whole payload + # would raise here; a first-row-only parse must not. + raw = ( + b'{"custom_id": "r-0", "body": {"model": "gemini-2.5-flash"}}\n' + b"garbage line that is not json\n" + ) + object_name = cfg.get_object_name( + ("batch.jsonl", raw, "application/jsonl"), purpose="batch" + ) + assert "gemini-2.5-flash" in object_name + + +class TestStreamingPeakMemory: + """ + Differential guard: the streaming transform must stay well under the peak + that a list pipeline incurs on the same input. If the hot path builds full + intermediate lists, the streaming assertion fails. + + The assertion that matters is the *relative* one: ``streaming_peak`` must be + a clear fraction of ``list_peak`` on the identical input. Absolute + ``tracemalloc`` ratios drift with GC timing and the live set carried in from + earlier tests, so they make poor CI gates; the relative comparison cancels + that shared noise and is exactly what regresses (toward 1.0) when the hot + path builds full intermediate lists. ``gc.collect()`` before each + measurement removes any garbage the previous run left behind. + """ + + def _measure(self, fn): + gc.collect() + tracemalloc.start() + try: + fn() + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + return peak + + def test_streaming_peak_well_below_list_pipeline(self): + cfg = VertexAIFilesConfig() + raw = _make_openai_jsonl_bytes(8000) + content_str = raw.decode("utf-8") + + def drain_stream(): + # Consume the upload body one row at a time, as the chunked uploader + # does, without accumulating it. + for _ in _OpenAIToVertexBatchUploadStream( + raw, cfg._map_openai_to_vertex_params + ).iter_bytes(): + pass + + streaming_peak = self._measure(drain_stream) + list_peak = self._measure( + lambda: _reference_vertex_jsonl_string(cfg, content_str) + ) + + # Core guard: the lazily consumed streaming body peaks well under a list + # pipeline that materializes every transformed row. Building full + # intermediate lists in the hot path pushes this ratio back toward 1.0. + assert streaming_peak < list_peak * 0.6, ( + f"streaming peak {streaming_peak} not a clear win over list pipeline " + f"{list_peak} (ratio {streaming_peak / list_peak:.2f})" + ) + + def test_get_object_name_does_not_scale_with_payload(self): + cfg = VertexAIFilesConfig() + raw = _make_openai_jsonl_bytes(8000) + file_data = ("batch.jsonl", raw, "application/jsonl") + + # The payload bytes already exist before measurement starts, so a lazy + # first-row parse should allocate only a small fraction of the payload; + # parsing every row would blow past this bound. + peak = self._measure(lambda: cfg.get_object_name(file_data, purpose="batch")) + assert ( + peak / len(raw) < 2.0 + ), "get_object_name should not copy the whole payload" + + +class TestPathSourcedStreaming: + """ + The proxy spools large batch uploads to a temp file and passes a pathlib.Path + as the file content instead of pre-reading bytes, so the transform streams + from disk. These lock in that a Path source yields identical output, keeps + every row, stays memory-bounded, and is re-iterable (multi-model uploads). + """ + + def _write_jsonl(self, tmp_path, n_rows, padding=400): + raw = _make_openai_jsonl_bytes(n_rows, padding=padding) + path = tmp_path / "batch.jsonl" + path.write_bytes(raw) + return path, raw + + def _batch_request(self, path) -> CreateFileRequest: + return {"file": ("batch.jsonl", path, "application/jsonl"), "purpose": "batch"} + + def test_transform_from_path_matches_legacy_and_keeps_all_rows(self, tmp_path): + cfg = VertexAIFilesConfig() + n_rows = 200 + path, raw = self._write_jsonl(tmp_path, n_rows) + data = self._batch_request(path) + + url = cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=data, + ) + assert "uploadType=resumable" in url + + out = cfg.transform_create_file_request( + model="", create_file_data=data, optional_params={}, litellm_params={} + ) + assert isinstance(out, dict) and "resumable_chunked_upload" in out + body = _join_upload_body(out).decode("utf-8") + assert body == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) + lines = body.splitlines() + assert len(lines) == n_rows, "no batch row may be dropped from a Path source" + first_labels = json.loads(lines[0])["request"]["labels"] + assert _get_litellm_batch_custom_id_from_labels(first_labels) == "request-0" + + def test_path_source_peak_stays_below_payload(self, tmp_path): + cfg = VertexAIFilesConfig() + path, raw = self._write_jsonl(tmp_path, 8000) + data = self._batch_request(path) + + def run(): + cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=data, + ) + out = cfg.transform_create_file_request( + model="", create_file_data=data, optional_params={}, litellm_params={} + ) + for _ in _resumable_stream(out).iter_bytes(): + pass # drain without accumulating + + gc.collect() + tracemalloc.start() + try: + run() + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + + # Streaming from disk must not materialize the payload. Reading the whole + # file into bytes (the pre-fix path) would push peak past the file size. + assert peak < len(raw) * 0.3, ( + f"peak {peak} not bounded vs payload {len(raw)} " + f"(ratio {peak / len(raw):.2f})" + ) + + def test_path_source_stream_is_reiterable(self, tmp_path): + cfg = VertexAIFilesConfig() + path, _ = self._write_jsonl(tmp_path, 50) + data = self._batch_request(path) + + out = cfg.transform_create_file_request( + model="", create_file_data=data, optional_params={}, litellm_params={} + ) + stream = _resumable_stream(out) + first = b"".join(stream.iter_bytes()) + second = b"".join(stream.iter_bytes()) + assert first == second and len(first) > 0 + + +_GCS_OBJECT_JSON = { + "id": "test-bucket/litellm-vertex-files/x/123", + "name": "litellm-vertex-files/x", + "size": "0", + "timeCreated": "2026-01-01T00:00:00.000000Z", + "purpose": "batch", +} + + +class _FixedBytesStream(BaseFileUploadStream): + """Streaming body of exact, controllable bytes for protocol-edge tests.""" + + def __init__(self, data: bytes, piece: int = 64): + self._data = data + self._piece = piece + + def iter_bytes(self): + for i in range(0, len(self._data), self._piece): + yield self._data[i : i + self._piece] + + +def _logging_obj() -> Logging: + return Logging( + model="", + messages=[], + stream=False, + call_type="acreate_file", + start_time=time.time(), + litellm_call_id="test", + function_id="", + ) + + +def _gcs_resumable_mock(session_url: str, final_status: int = 200): + """A fake GCS resumable endpoint: POST opens a session (URI in Location), + each PUT appends and returns 308 until the final chunk returns 200/201.""" + state = {"received": bytearray(), "ranges": [], "methods": [], "urls": []} + + async def handler(request: httpx.Request) -> httpx.Response: + state["methods"].append(request.method) + state["urls"].append(str(request.url)) + if request.method == "POST": + return httpx.Response(200, headers={"location": session_url}) + body = await request.aread() + content_range = request.headers["content-range"] + state["ranges"].append(content_range) + state["received"].extend(body) + if content_range.rsplit("/", 1)[-1] == "*": + return httpx.Response( + 308, headers={"range": f"bytes=0-{len(state['received']) - 1}"} + ) + return httpx.Response(final_status, json=_GCS_OBJECT_JSON) + + return handler, state + + +def _async_handler_with(mock) -> AsyncHTTPHandler: + handler = AsyncHTTPHandler() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(mock)) + return handler + + +class TestResumableUploadUrl: + def test_batch_jsonl_uses_resumable_upload_type(self): + cfg = VertexAIFilesConfig() + request: CreateFileRequest = { + "file": ("batch.jsonl", _make_openai_jsonl_bytes(3), "application/jsonl"), + "purpose": "batch", + } + url = cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=request, + ) + assert "uploadType=resumable" in url + assert "uploadType=media" not in url + + def test_batch_text_plain_uses_resumable_upload_type(self): + # Clients often label a .jsonl batch upload as text/plain; it must still + # take the streaming/resumable path, not the buffered media path. + cfg = VertexAIFilesConfig() + request: CreateFileRequest = { + "file": ("batch.jsonl", _make_openai_jsonl_bytes(3), "text/plain"), + "purpose": "batch", + } + url = cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=request, + ) + assert "uploadType=resumable" in url + assert "uploadType=media" not in url + + def test_binary_upload_stays_simple_media(self): + cfg = VertexAIFilesConfig() + request: CreateFileRequest = { + "file": ("doc.pdf", b"%PDF-1.4 binary", "application/pdf"), + "purpose": "user_data", + } + url = cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=request, + ) + assert "uploadType=media" in url + assert "uploadType=resumable" not in url + + +class TestResumableStreamBody: + def test_stream_matches_legacy_pipeline(self): + cfg = VertexAIFilesConfig() + raw = _make_openai_jsonl_bytes(120) + stream = _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params) + assert b"".join(stream.iter_bytes()).decode( + "utf-8" + ) == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) + + def test_stream_is_reiterable_for_retries(self): + # A one-shot generator would make a transport retry upload an empty body; + # iter_bytes() must yield the full payload every call. + cfg = VertexAIFilesConfig() + raw = _make_openai_jsonl_bytes(40) + stream = _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params) + first = b"".join(stream.iter_bytes()) + second = b"".join(stream.iter_bytes()) + assert first == second and len(first) > 0 + + def test_stream_is_reiterable_for_seekable_file_like_input(self): + # A seekable handle (BytesIO, temp file) must be rewound between calls; + # otherwise the first iter_bytes() exhausts it and a retry would upload + # an empty body silently. + cfg = VertexAIFilesConfig() + raw = _make_openai_jsonl_bytes(40) + stream = _OpenAIToVertexBatchUploadStream( + io.BytesIO(raw), cfg._map_openai_to_vertex_params + ) + first = b"".join(stream.iter_bytes()) + second = b"".join(stream.iter_bytes()) + assert first == second and len(first) > 0 + + +class TestResumableChunking: + def test_intermediate_chunks_are_exactly_chunk_size(self): + pieces = list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([b"x" * 10]), 4)) + assert pieces == [b"xxxx", b"xxxx", b"xx"] + + def test_exact_multiple_yields_no_trailing_empty(self): + # An exactly chunk-aligned stream yields only full chunks; the upload + # finalizes on the last data chunk instead of an extra empty request. + pieces = list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([b"x" * 8]), 4)) + assert pieces == [b"xxxx", b"xxxx"] + + def test_empty_stream_yields_nothing(self): + # A 0-byte stream yields no chunks; the caller finalizes with one empty + # request (bytes */0). + assert list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([]), 4)) == [] + + def test_default_chunk_size_is_256kib_multiple(self): + assert BaseLLMHTTPHandler._RESUMABLE_CHUNK_SIZE % (256 * 1024) == 0 + + def test_content_range_intermediate_uses_star_total(self): + assert ( + BaseLLMHTTPHandler._resumable_content_range(0, 4096, is_final=False) + == "bytes 0-4095/*" + ) + + def test_content_range_final_uses_real_total(self): + assert ( + BaseLLMHTTPHandler._resumable_content_range(8192, 100, is_final=True) + == "bytes 8192-8291/8292" + ) + + def test_content_range_empty_finalize(self): + assert ( + BaseLLMHTTPHandler._resumable_content_range(8192, 0, is_final=True) + == "bytes */8192" + ) + + +@pytest.mark.asyncio +class TestResumableUploadProtocol: + """End-to-end against a faked GCS resumable endpoint. These are the tests + that fail if the handler buffers the whole body, drops bytes, mislabels a + Content-Range, follows the 308 instead of continuing, or skips finalize.""" + + async def _run(self, raw: bytes, chunk_size: int, final_status: int = 200): + cfg = VertexAIFilesConfig() + request: CreateFileRequest = { + "file": ("batch.jsonl", raw, "application/jsonl"), + "purpose": "batch", + } + api_base = cfg.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "test-bucket"}, + data=request, + ) + transformed = cfg.transform_create_file_request( + model="", create_file_data=request, optional_params={}, litellm_params={} + ) + transformed["resumable_chunked_upload"]["chunk_size"] = chunk_size + expected = _join_upload_body(transformed) + + session_url = "https://storage.googleapis.com/upload/sess?upload_id=SID" + mock, state = _gcs_resumable_mock(session_url, final_status=final_status) + response = await BaseLLMHTTPHandler().async_create_file( + transformed_request=transformed, + litellm_params={}, + provider_config=cfg, + headers={"Authorization": "Bearer x"}, + api_base=api_base, + logging_obj=_logging_obj(), + client=_async_handler_with(mock), + timeout=None, + ) + return expected, state, response, session_url, api_base + + async def test_streams_in_chunks_and_reassembles(self): + raw = _make_openai_jsonl_bytes(300) + chunk_size = 4096 + expected, state, response, session_url, api_base = await self._run( + raw, chunk_size + ) + + # One session-open POST, then a sequence of chunk PUTs. + assert state["methods"][0] == "POST" + assert set(state["methods"][1:]) == {"PUT"} + assert state["methods"].count("PUT") >= 2, "payload must span multiple chunks" + + # POST opens a resumable session; every chunk goes to the session URI. + assert "uploadType=resumable" in state["urls"][0] + assert all(u == session_url for u in state["urls"][1:]) + + # Every non-final chunk is exactly chunk_size with an unknown-total range; + # the final chunk carries the real total. + intermediate = state["ranges"][:-1] + for index, content_range in enumerate(intermediate): + assert ( + content_range + == f"bytes {index * chunk_size}-{(index + 1) * chunk_size - 1}/*" + ) + total = len(expected) + last_offset = len(intermediate) * chunk_size + if last_offset == total: # payload landed on a chunk boundary + assert state["ranges"][-1] == f"bytes */{total}" + else: + assert state["ranges"][-1] == f"bytes {last_offset}-{total - 1}/{total}" + + # The bytes GCS received are exactly the transformed batch payload. + assert bytes(state["received"]) == expected + assert response.object == "file" + + async def test_exact_multiple_finalizes_on_last_data_chunk(self): + # A body that is an exact multiple of the chunk size finalizes on its + # last data chunk (bytes (TOTAL-chunk)-(TOTAL-1)/TOTAL), with no extra + # empty finalize request. + chunk_size = 256 + total = chunk_size * 3 + stream = _FixedBytesStream(b"a" * total) + config = {"body_stream": stream, "chunk_size": chunk_size} + session_url = "https://storage.googleapis.com/upload/sess?upload_id=SID" + mock, state = _gcs_resumable_mock(session_url) + + response = await BaseLLMHTTPHandler()._aresumable_chunked_upload( + client=_async_handler_with(mock), + initiate_url="https://storage.googleapis.com/upload?uploadType=resumable", + base_headers={"Authorization": "Bearer x"}, + config=config, + timeout=None, + ) + + assert state["ranges"][-1] == f"bytes {total - chunk_size}-{total - 1}/{total}" + assert "*" not in state["ranges"][-1] + assert bytes(state["received"]) == b"a" * total + assert response.status_code == 200 + + async def test_failed_chunk_raises(self): + raw = _make_openai_jsonl_bytes(80) + with pytest.raises(Exception): + await self._run(raw, chunk_size=4096, final_status=403) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 7c063c72607..7ac5a50525b 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -14,8 +14,8 @@ from unittest.mock import MagicMock from litellm.llms.vertex_ai.files.transformation import ( VertexAIFilesConfig, - VertexAIJsonlFilesTransformation, _get_litellm_batch_custom_id_from_labels, + _openai_batch_jsonl_entry_to_vertex_wrapped_request, _sanitize_gcp_label_value, ) from litellm.types.llms.openai import OpenAIFileObject, HttpxBinaryResponseContent @@ -33,7 +33,7 @@ class TestParseGcsUri: def test_should_parse_standard_gs_uri(self, config): file_id = "gs://my-bucket/litellm-vertex-files/path/to/object.jsonl" bucket, encoded = config._parse_gcs_uri( - file_id, litellm_params={"bucket_name": "my-bucket"} + file_id, litellm_params={"gcs_bucket_name": "my-bucket"} ) assert bucket == "my-bucket" assert encoded == urllib.parse.quote( @@ -43,7 +43,7 @@ class TestParseGcsUri: def test_should_parse_uri_with_nested_publisher_path(self, config): uri = "gs://litellm-local/litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123" bucket, encoded = config._parse_gcs_uri( - uri, litellm_params={"bucket_name": "litellm-local"} + uri, litellm_params={"gcs_bucket_name": "litellm-local"} ) assert bucket == "litellm-local" expected_path = ( @@ -56,7 +56,7 @@ class TestParseGcsUri: "gs://my-bucket/litellm-vertex-files/some/path", safe="" ) bucket, encoded = config._parse_gcs_uri( - encoded_uri, litellm_params={"bucket_name": "my-bucket"} + encoded_uri, litellm_params={"gcs_bucket_name": "my-bucket"} ) assert bucket == "my-bucket" assert encoded == urllib.parse.quote("litellm-vertex-files/some/path", safe="") @@ -64,21 +64,21 @@ class TestParseGcsUri: def test_should_reject_bucket_only(self, config): with pytest.raises(ValueError, match="object name"): config._parse_gcs_uri( - "gs://my-bucket", litellm_params={"bucket_name": "my-bucket"} + "gs://my-bucket", litellm_params={"gcs_bucket_name": "my-bucket"} ) def test_should_reject_no_gs_prefix(self, config): with pytest.raises(ValueError, match="gs://"): config._parse_gcs_uri( "my-bucket/litellm-vertex-files/object.txt", - litellm_params={"bucket_name": "my-bucket"}, + litellm_params={"gcs_bucket_name": "my-bucket"}, ) def test_should_reject_unmanaged_object_path(self, config): with pytest.raises(ValueError, match="LiteLLM-managed"): config._parse_gcs_uri( "gs://my-bucket/private/object.txt", - litellm_params={"bucket_name": "my-bucket"}, + litellm_params={"gcs_bucket_name": "my-bucket"}, ) def test_should_reject_request_supplied_legacy_flag(self, config): @@ -86,7 +86,7 @@ class TestParseGcsUri: config._parse_gcs_uri( "gs://my-bucket/private/object.txt", litellm_params={ - "bucket_name": "my-bucket", + "gcs_bucket_name": "my-bucket", "allow_legacy_cloud_file_ids": True, }, ) @@ -96,7 +96,7 @@ class TestParseGcsUri: bucket, encoded = config._parse_gcs_uri( "gs://my-bucket/private/object.txt", litellm_params={ - "bucket_name": "my-bucket", + "gcs_bucket_name": "my-bucket", "_litellm_internal_model_credentials": trusted_credentials, }, ) @@ -109,7 +109,7 @@ class TestParseGcsUri: config._parse_gcs_uri( "gs://my-bucket/private/object.txt", litellm_params={ - "bucket_name": "my-bucket", + "gcs_bucket_name": "my-bucket", "_litellm_internal_model_credentials": { "allow_legacy_cloud_file_ids": True }, @@ -121,7 +121,7 @@ class TestParseGcsUri: bucket, encoded = config._parse_gcs_uri( "gs://my-bucket/team-a/private/object.txt", litellm_params={ - "bucket_name": "my-bucket/team-a", + "gcs_bucket_name": "my-bucket/team-a", "_litellm_internal_model_credentials": trusted_credentials, }, ) @@ -135,7 +135,7 @@ class TestParseGcsUri: config._parse_gcs_uri( "gs://my-bucket/team-b/private/object.txt", litellm_params={ - "bucket_name": "my-bucket/team-a", + "gcs_bucket_name": "my-bucket/team-a", "_litellm_internal_model_credentials": trusted_credentials, }, ) @@ -144,7 +144,7 @@ class TestParseGcsUri: with pytest.raises(ValueError, match="configured storage bucket"): config._parse_gcs_uri( "gs://other-bucket/litellm-vertex-files/object.txt", - litellm_params={"bucket_name": "my-bucket"}, + litellm_params={"gcs_bucket_name": "my-bucket"}, ) @@ -156,7 +156,7 @@ class TestCreateFileUrl: model="", optional_params={}, litellm_params={ - "bucket_name": "safe-bucket", + "gcs_bucket_name": "safe-bucket", "litellm_metadata": {"gcs_bucket_name": "attacker-bucket"}, }, data={ @@ -182,7 +182,7 @@ class TestTransformRetrieveFile: url, params = config.transform_retrieve_file_request( file_id=file_id, optional_params={}, - litellm_params={"bucket_name": "my-bucket"}, + litellm_params={"gcs_bucket_name": "my-bucket"}, ) expected_encoded = urllib.parse.quote( "litellm-vertex-files/path/to/file.jsonl", safe="" @@ -243,7 +243,7 @@ class TestTransformFileContent: url, params = config.transform_file_content_request( file_content_request={"file_id": file_id}, optional_params={}, - litellm_params={"bucket_name": "my-bucket"}, + litellm_params={"gcs_bucket_name": "my-bucket"}, ) encoded = urllib.parse.quote("litellm-vertex-files/path/to/file.jsonl", safe="") assert ( @@ -378,7 +378,7 @@ class TestTransformDeleteFile: url, params = config.transform_delete_file_request( file_id=file_id, optional_params={}, - litellm_params={"bucket_name": "my-bucket"}, + litellm_params={"gcs_bucket_name": "my-bucket"}, ) encoded = urllib.parse.quote("litellm-vertex-files/path/to/file.jsonl", safe="") assert ( @@ -854,6 +854,106 @@ class TestVertexBatchOutputTransformation: ) assert transformed_content == invalid_content + def test_binary_content_passthrough(self, config): + """A binary file (PDF/video) whose first bytes are not valid UTF-8 must be + returned unchanged. The row-by-row transform only engages for a JSONL + batch output and must never line-parse or corrupt binary content.""" + binary = b"%PDF-1.4\n%\xc4\xe5\xf2\xe5\xeb\xa7\n" + b"\x00\x01\x02\xff\xfe" * 64 + assert config._try_transform_vertex_batch_output_to_openai(binary) == binary + + def test_streaming_transform_peaks_below_list_pipeline(self, config): + """The output transform must stream row-by-row, not build a list of every + parsed row and a second list of transformed rows. This guards against a + regression to the list pipeline, which peaks at several full copies and + OOMs on large result files. The relative comparison cancels shared noise + (per-row transform cost, GC timing) and only the list overhead differs. + """ + import gc + import tracemalloc + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + def vertex_row(index: int) -> dict: + return { + "status": "", + "processed_time": "2024-11-01T18:13:16.826+00:00", + "request": { + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "labels": {"litellm_custom_id": f"r-{index}"}, + }, + "response": { + "candidates": [ + { + "content": { + "parts": [{"text": "hello " * 20}], + "role": "model", + }, + "finishReason": "STOP", + } + ], + "modelVersion": "gemini-2.0-flash-001", + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 20, + "totalTokenCount": 30, + }, + }, + } + + content = ("\n".join(json.dumps(vertex_row(i)) for i in range(4000))).encode( + "utf-8" + ) + + def list_pipeline() -> bytes: + gemini_config = VertexGeminiConfig() + logging_obj = Logging( + model="", + messages=[], + stream=False, + call_type="batch_transform", + start_time=0.1, + litellm_call_id="", + function_id="", + ) + logging_obj.optional_params = {} + mock_response = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + request=httpx.Request("POST", "https://example.com"), + ) + rows = content.decode("utf-8").strip().split("\n") + transformed = [ + json.dumps( + config._transform_single_vertex_batch_output_to_openai( + json.loads(row), gemini_config, logging_obj, mock_response + ) + ) + for row in rows + ] + return "\n".join(transformed).encode("utf-8") + + def peak_of(fn) -> int: + gc.collect() + tracemalloc.start() + try: + fn() + return tracemalloc.get_traced_memory()[1] + finally: + tracemalloc.stop() + + streaming_peak = peak_of( + lambda: config._try_transform_vertex_batch_output_to_openai(content) + ) + list_peak = peak_of(list_pipeline) + + assert streaming_peak < list_peak * 0.75, ( + f"streaming peak {streaming_peak} is not a clear win over the list " + f"pipeline {list_peak} (ratio {streaming_peak / list_peak:.2f})" + ) + class TestTryTransformDoesNotMutateCallerLoggingObj: """Regression tests: _try_transform_vertex_batch_output_to_openai must not mutate @@ -953,12 +1053,23 @@ class TestTryTransformDoesNotMutateCallerLoggingObj: assert transformed["response"]["status_code"] == 200 +def _wrap_entries(openai_jsonl_content): + """Vertex-wrapped requests for a list of OpenAI batch entries, built via the + live single-entry transform that the streaming upload path uses.""" + cfg = VertexAIFilesConfig() + return [ + _openai_batch_jsonl_entry_to_vertex_wrapped_request( + entry, cfg._map_openai_to_vertex_params + ) + for entry in openai_jsonl_content + ] + + class TestVertexBatchCustomIdLabels: """Test custom_id handling in batch transformations""" def test_custom_id_added_to_labels_in_vertex_request(self): """Test that custom_id from OpenAI format is added as a label in Vertex AI format""" - transformation = VertexAIJsonlFilesTransformation() openai_jsonl_content = [ { @@ -973,11 +1084,7 @@ class TestVertexBatchCustomIdLabels: } ] - vertex_jsonl_content = ( - transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_jsonl_content - ) - ) + vertex_jsonl_content = _wrap_entries(openai_jsonl_content) assert len(vertex_jsonl_content) == 1 vertex_request = vertex_jsonl_content[0] @@ -992,7 +1099,6 @@ class TestVertexBatchCustomIdLabels: def test_long_custom_id_round_trips_across_raw_label_chunks(self): """Test that long custom_ids are not truncated in raw labels.""" - transformation = VertexAIJsonlFilesTransformation() custom_id_a = "shared-prefix-that-is-longer-than-thirty-six-bytes-A" custom_id_b = "shared-prefix-that-is-longer-than-thirty-six-bytes-B" @@ -1009,11 +1115,7 @@ class TestVertexBatchCustomIdLabels: for custom_id in (custom_id_a, custom_id_b) ] - vertex_jsonl_content = ( - transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_jsonl_content - ) - ) + vertex_jsonl_content = _wrap_entries(openai_jsonl_content) labels_a = vertex_jsonl_content[0]["request"]["labels"] labels_b = vertex_jsonl_content[1]["request"]["labels"] @@ -1028,7 +1130,6 @@ class TestVertexBatchCustomIdLabels: def test_multiple_requests_each_get_their_own_label(self): """Test that multiple requests each get their own custom_id label""" - transformation = VertexAIJsonlFilesTransformation() openai_jsonl_content = [ { @@ -1043,11 +1144,7 @@ class TestVertexBatchCustomIdLabels: for i in range(3) ] - vertex_jsonl_content = ( - transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_jsonl_content - ) - ) + vertex_jsonl_content = _wrap_entries(openai_jsonl_content) assert len(vertex_jsonl_content) == 3 @@ -1063,7 +1160,6 @@ class TestVertexBatchCustomIdLabels: def test_request_without_custom_id_has_no_label(self): """Test that requests without custom_id don't get a label""" - transformation = VertexAIJsonlFilesTransformation() openai_jsonl_content = [ { @@ -1076,11 +1172,7 @@ class TestVertexBatchCustomIdLabels: } ] - vertex_jsonl_content = ( - transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_jsonl_content - ) - ) + vertex_jsonl_content = _wrap_entries(openai_jsonl_content) # Should not have labels if no custom_id was provided assert "labels" not in vertex_jsonl_content[0]["request"] @@ -1090,7 +1182,6 @@ class TestVertexBatchCustomIdLabels: Test the full round trip: OpenAI format -> Vertex AI format -> Vertex AI output -> OpenAI output Verify that custom_id is preserved through the entire flow. """ - transformation = VertexAIJsonlFilesTransformation() config = VertexAIFilesConfig() # Step 1: Transform OpenAI input to Vertex AI format (mixed case exercises raw label) @@ -1106,11 +1197,7 @@ class TestVertexBatchCustomIdLabels: } ] - vertex_input = ( - transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_input - ) - ) + vertex_input = _wrap_entries(openai_input) # Verify both labels are GCP-safe and encoded raw preserves round-trip. assert ( @@ -1154,7 +1241,6 @@ class TestVertexBatchCustomIdLabels: def test_custom_id_label_sanitization(self): """Test that custom_id values are sanitized to meet GCP label constraints""" - transformation = VertexAIJsonlFilesTransformation() # Test sanitization function assert _sanitize_gcp_label_value("MyRequest-1") == "myrequest-1" @@ -1179,11 +1265,7 @@ class TestVertexBatchCustomIdLabels: } ] - vertex_input = ( - transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_input - ) - ) + vertex_input = _wrap_entries(openai_input) # Verify both labels are safe for GCP labels. assert ( @@ -1192,3 +1274,47 @@ class TestVertexBatchCustomIdLabels: raw_label = vertex_input[0]["request"]["labels"]["litellm_custom_id_raw"] assert raw_label != "MyRequest-1" assert _sanitize_gcp_label_value(raw_label) == raw_label + + +class TestConfiguredBucketNameResolution: + def test_should_resolve_new_gcs_bucket_name_key(self, config, monkeypatch): + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + assert ( + config._get_configured_bucket_name({"gcs_bucket_name": "my-new-bucket"}) + == "my-new-bucket" + ) + + def test_should_resolve_legacy_bucket_name_key(self, config, monkeypatch): + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + assert ( + config._get_configured_bucket_name({"bucket_name": "my-legacy-bucket"}) + == "my-legacy-bucket" + ) + + def test_should_prefer_new_key_over_legacy(self, config, monkeypatch): + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + assert ( + config._get_configured_bucket_name( + {"gcs_bucket_name": "new", "bucket_name": "legacy"} + ) + == "new" + ) + + def test_should_fall_back_to_env(self, config, monkeypatch): + monkeypatch.setenv("GCS_BUCKET_NAME", "env-bucket") + assert config._get_configured_bucket_name({}) == "env-bucket" + + def test_should_raise_when_no_bucket_anywhere(self, config, monkeypatch): + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + with pytest.raises(ValueError, match="GCS bucket_name is required"): + config._get_configured_bucket_name({}) + + def test_legacy_kwarg_survives_get_litellm_params(self): + from litellm.litellm_core_utils.get_litellm_params import ( + _OPTIONAL_KWARGS_KEYS, + get_litellm_params, + ) + + assert "bucket_name" in _OPTIONAL_KWARGS_KEYS + params = get_litellm_params(bucket_name="my-legacy-bucket") + assert params.get("bucket_name") == "my-legacy-bucket" diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index a076a09706b..1ca22ebaa75 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -7,167 +7,138 @@ VERIA-39 regression tests: models the caller is not authorized to use. """ + from unittest.mock import AsyncMock, MagicMock, patch + import pytest + + from fastapi import HTTPException + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth -# --------------------------------------------------------------------------- -# Token counter — covers all three batch payload shapes -# --------------------------------------------------------------------------- + +def _models(file_content_as_dict): + """Distinct body.model values, mirroring how the rate limiter collects the + models from a streamed batch file before the access check.""" + return [ + entry["body"]["model"] + for entry in file_content_as_dict + if (entry.get("body") or {}).get("model") + ] def test_token_counter_counts_chat_messages(): - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - { - "body": { - "model": "gpt-4o-mini", - "messages": [{"role": "user", "content": "hello"}], - } + tokens = _count_entry_tokens( + { + "body": { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], } - ] + } ) - assert usage.prompt_tokens > 0 + assert tokens > 0 def test_token_counter_counts_text_completion_prompt(): - """Pre-fix this returned 0 tokens (the function only inspected + """Pre-fix this returned 0 tokens (the counter only inspected `messages`), letting `prompt`-style batches slip past TPM limits.""" - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - {"body": {"model": "gpt-3.5-turbo-instruct", "prompt": "hello world"}} - ] + tokens = _count_entry_tokens( + {"body": {"model": "gpt-3.5-turbo-instruct", "prompt": "hello world"}} ) - assert usage.prompt_tokens > 0 + assert tokens > 0 def test_token_counter_counts_embedding_input_string(): - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - {"body": {"model": "text-embedding-3-small", "input": "hello world"}} - ] + tokens = _count_entry_tokens( + {"body": {"model": "text-embedding-3-small", "input": "hello world"}} ) - assert usage.prompt_tokens > 0 + assert tokens > 0 def test_token_counter_counts_embedding_input_list(): - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - { - "body": { - "model": "text-embedding-3-small", - "input": ["hello", "world"], - } + tokens = _count_entry_tokens( + { + "body": { + "model": "text-embedding-3-small", + "input": ["hello", "world"], } - ] + } ) - assert usage.prompt_tokens > 0 + assert tokens > 0 def test_token_counter_counts_text_completion_prompt_list(): - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - { - "body": { - "model": "gpt-3.5-turbo-instruct", - "prompt": ["alpha", "beta"], - } + tokens = _count_entry_tokens( + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": ["alpha", "beta"], } - ] + } ) - assert usage.prompt_tokens > 0 + assert tokens > 0 def test_token_counter_counts_pre_tokenized_prompt_int_list(): """OpenAI's text-completion API accepts a single pre-tokenized prompt as a list of ints. Each int is one token; pre-fix this shape was silently counted as zero, leaving a TPM bypass.""" - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - { - "body": { - "model": "gpt-3.5-turbo-instruct", - "prompt": [1, 2, 3, 4, 5], - } + tokens = _count_entry_tokens( + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": [1, 2, 3, 4, 5], } - ] + } ) - assert usage.prompt_tokens == 5 + assert tokens == 5 def test_token_counter_counts_pre_tokenized_prompt_list_of_int_lists(): """Multiple pre-tokenized prompts (`list[list[int]]`) — the most important bypass shape. A 1000-token batch must report 1000 tokens, not zero.""" - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - { - "body": { - "model": "gpt-3.5-turbo-instruct", - "prompt": [[1] * 250, [2] * 250, [3] * 500], - } + tokens = _count_entry_tokens( + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": [[1] * 250, [2] * 250, [3] * 500], } - ] + } ) - assert usage.prompt_tokens == 1000 + assert tokens == 1000 def test_token_counter_counts_pre_tokenized_input_for_embeddings(): """Same shape applies to embeddings (`input`).""" - from litellm.batches.batch_utils import _get_batch_job_input_file_usage + from litellm.batches.batch_utils import _count_entry_tokens - usage = _get_batch_job_input_file_usage( - file_content_dictionary=[ - { - "body": { - "model": "text-embedding-3-small", - "input": [[1, 2, 3], [4, 5, 6]], - } + tokens = _count_entry_tokens( + { + "body": { + "model": "text-embedding-3-small", + "input": [[1, 2, 3], [4, 5, 6]], } - ] + } ) - assert usage.prompt_tokens == 6 - - -# --------------------------------------------------------------------------- -# Model extractor -# --------------------------------------------------------------------------- - - -def test_model_extractor_returns_distinct_models(): - from litellm.batches.batch_utils import _get_models_from_batch_input_file_content - - models = _get_models_from_batch_input_file_content( - [ - {"body": {"model": "gpt-4o", "messages": []}}, - {"body": {"model": "gpt-4o", "messages": []}}, # duplicate - {"body": {"model": "gpt-4o-mini", "messages": []}}, - {"body": {}}, # missing model - ] - ) - assert models == ["gpt-4o", "gpt-4o-mini"] - - -# --------------------------------------------------------------------------- -# Pre-call hook model validation -# --------------------------------------------------------------------------- + assert tokens == 6 @pytest.mark.asyncio @@ -211,7 +182,7 @@ async def test_pre_call_rejects_unauthorized_model_in_batch_file(): with pytest.raises(HTTPException) as exc: await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), ) assert exc.value.status_code == 403 @@ -255,7 +226,7 @@ async def test_pre_call_allows_authorized_model_in_batch_file(): # Should not raise await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), ) @@ -293,7 +264,7 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias( ): await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), target_model_names=[proxy_alias], ) @@ -386,7 +357,7 @@ async def test_pre_call_uses_target_model_names_not_stripped_reverse_lookup( ): await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), target_model_names=[batch_alias], ) @@ -412,9 +383,274 @@ async def test_pre_call_skips_check_when_no_models_present(): # entirely. await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=[], + models=_models([]), ) await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=[{"body": {}}], + models=_models([{"body": {}}]), ) + + +def _make_batch_input_bytes(n_rows: int, padding: int = 200) -> bytes: + import json as _json + + pad = "x" * padding + rows = [] + for i in range(n_rows): + rows.append( + _json.dumps( + { + "custom_id": f"request-{i}", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gpt-4o" if i % 2 else "gpt-3.5-turbo", + "messages": [{"role": "user", "content": f"{pad} {i}"}], + }, + } + ) + ) + return ("\n".join(rows)).encode("utf-8") + + +def test_iter_batch_input_entries_matches_dict_list(): + from litellm.batches.batch_utils import ( + _get_file_content_as_dictionary, + _iter_batch_input_entries, + ) + + raw = _make_batch_input_bytes(50) + streamed = list(_iter_batch_input_entries(raw)) + assert streamed == _get_file_content_as_dictionary(raw) + assert streamed[0]["custom_id"] == "request-0" + # tolerant of blank lines and a missing trailing newline + assert list(_iter_batch_input_entries(raw + b"\n\n")) == streamed + + +def test_streaming_count_peak_below_dict_list(): + import gc + import tracemalloc + + from litellm.batches.batch_utils import ( + _get_file_content_as_dictionary, + _iter_batch_input_entries, + ) + + raw = _make_batch_input_bytes(8000) + + def _measure(fn): + gc.collect() + tracemalloc.start() + try: + fn() + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + return peak + + def _stream(): + count = 0 + models: set = set() + for entry in _iter_batch_input_entries(raw): + count += 1 + model = (entry.get("body") or {}).get("model") + if model: + models.add(model) + return count + + def _build_list(): + return len(_get_file_content_as_dictionary(raw)) + + stream_peak = _measure(_stream) + list_peak = _measure(_build_list) + assert stream_peak < list_peak * 0.5, ( + f"streaming count peak {stream_peak} is not a clear win over the dict " + f"list {list_peak} (ratio {stream_peak / list_peak:.2f})" + ) + + +@pytest.mark.asyncio +async def test_count_input_file_usage_streams_without_building_list(): + """count_input_file_usage must count requests/tokens in one streaming pass. + Mocks the download; asserts the count is correct and that the dict-list + helper is never called (a revert to the list approach would call it).""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + raw = _make_batch_input_bytes(10) + fake_content = MagicMock() + fake_content.content = raw + + with ( + patch("litellm.afile_content", new=AsyncMock(return_value=fake_content)), + patch( + "litellm.batches.batch_utils._get_file_content_as_dictionary" + ) as mock_dict_list, + ): + usage = await rate_limiter.count_input_file_usage( + file_id="file-not-managed", + custom_llm_provider="openai", + user_api_key_dict=None, + ) + + assert usage.request_count == 10 + assert usage.total_tokens > 0 + mock_dict_list.assert_not_called() + + +def _one_row_batch_bytes(model: str) -> bytes: + import json as _json + + return ( + _json.dumps( + { + "custom_id": "r0", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model, + "messages": [{"role": "user", "content": "x"}], + }, + } + ) + + "\n" + ).encode("utf-8") + + +@pytest.mark.asyncio +async def test_count_input_file_usage_enforces_models_when_token_counting_fails(): + """Security regression: a row whose content makes token counting raise must + NOT skip the model allowlist check. async_pre_call_hook swallows non-HTTP + exceptions and submits the batch, so a raised counting error would otherwise + fail open. The access check must still run and deny the restricted model.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + fake_content = MagicMock() + fake_content.content = _one_row_batch_bytes("restricted-model") + user = UserAPIKeyAuth( + api_key="sk-x", + user_id="bob", + models=["only-allowed"], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + def _boom(*args, **kwargs): + raise ValueError("unsupported content part: input_audio") + + deny = AsyncMock(side_effect=Exception("model not in allowlist")) + + with ( + patch("litellm.afile_content", new=AsyncMock(return_value=fake_content)), + patch("litellm.proxy.hooks.batch_rate_limiter._count_entry_tokens", new=_boom), + patch("litellm.proxy.auth.auth_checks.can_key_call_model", new=deny), + patch("litellm.proxy.proxy_server.llm_router", MagicMock(model_list=[])), + ): + with pytest.raises(HTTPException) as exc: + await rate_limiter.count_input_file_usage( + file_id="file-not-managed", + custom_llm_provider="openai", + user_api_key_dict=user, + ) + + # The access check ran despite token counting failing, and denied the model. + deny.assert_awaited() + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_count_input_file_usage_estimates_tokens_when_counting_fails_for_allowed_model(): + """A token-counting failure for an allowed model must not hard-block the batch + (the pre-streaming behavior let such batches through), but it also must not + zero the token total, which would let a caller evade the TPM limit by sending + rows the counter cannot measure. The row falls back to a conservative + size-based estimate so the batch proceeds with a non-zero count.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + fake_content = MagicMock() + fake_content.content = _one_row_batch_bytes("allowed-model") + user = UserAPIKeyAuth( + api_key="sk-x", + user_id="bob", + models=["allowed-model"], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + def _boom(*args, **kwargs): + raise ValueError("unsupported content part: file") + + allow = AsyncMock(return_value=True) + + with ( + patch("litellm.afile_content", new=AsyncMock(return_value=fake_content)), + patch("litellm.proxy.hooks.batch_rate_limiter._count_entry_tokens", new=_boom), + patch("litellm.proxy.auth.auth_checks.can_key_call_model", new=allow), + patch("litellm.proxy.proxy_server.llm_router", MagicMock(model_list=[])), + ): + usage = await rate_limiter.count_input_file_usage( + file_id="file-not-managed", + custom_llm_provider="openai", + user_api_key_dict=user, + ) + + allow.assert_awaited() + assert usage.request_count == 1 + # Estimated, not zeroed: a crafted uncountable row can't evade the TPM limit. + assert usage.total_tokens > 0 + + +@pytest.mark.asyncio +async def test_count_input_file_usage_collects_models_after_malformed_line(): + """A malformed JSONL line must not abort model collection. A restricted model + named on a row AFTER a malformed line must still be collected and denied by the + allowlist check, otherwise a caller could hide a restricted model behind a bad + row.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + fake_content = MagicMock() + fake_content.content = ( + _one_row_batch_bytes("only-allowed") + + b"{ this is not valid json\n" + + _one_row_batch_bytes("restricted-model") + ) + user = UserAPIKeyAuth( + api_key="sk-x", + user_id="bob", + models=["only-allowed"], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + async def _deny_restricted(model, **kwargs): + if model == "restricted-model": + raise Exception("model not in allowlist") + return True + + deny = AsyncMock(side_effect=_deny_restricted) + + with ( + patch("litellm.afile_content", new=AsyncMock(return_value=fake_content)), + patch("litellm.proxy.auth.auth_checks.can_key_call_model", new=deny), + patch("litellm.proxy.proxy_server.llm_router", MagicMock(model_list=[])), + ): + with pytest.raises(HTTPException) as exc: + await rate_limiter.count_input_file_usage( + file_id="file-not-managed", + custom_llm_provider="openai", + user_api_key_dict=user, + ) + + assert exc.value.status_code == 403 diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 5fc36b71f2b..7564844d54a 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -378,6 +378,83 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: app.dependency_overrides.pop(ps.user_api_key_auth, None) +def test_create_file_batch_streams_from_upload_spool(monkeypatch, llm_router: Router): + """ + Batch uploads must be passed downstream as the upload's streamable file handle + (Starlette's already-spooled file), not read into an in-memory bytes object, so + the proxy never buffers the whole payload. Non-batch uploads keep the bytes path. + """ + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.openai_files_endpoints import files_endpoints as fe + from litellm.types.llms.openai import OpenAIFileObject + + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + captured: dict = {} + + async def fake_route_create_file(*, _create_file_request, **kwargs): + file_elem = _create_file_request["file"][1] + captured["file_elem"] = file_elem + if hasattr(file_elem, "read") and hasattr(file_elem, "seek"): + file_elem.seek(0) + captured["streamed_content"] = file_elem.read() + return OpenAIFileObject( + id="dummy-id", + object="file", + bytes=0, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(fe, "route_create_file", fake_route_create_file) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" + ) + + content = ( + b'{"custom_id":"r-0","method":"POST","url":"/v1/chat/completions",' + b'"body":{"model":"gpt-3.5-turbo","messages":[{"role":"user","content":"hi"}]}}\n' + ) + try: + resp = client.post( + "/v1/files", + files={"file": ("batch.jsonl", content, "application/jsonl")}, + data={"purpose": "batch"}, + headers={"Authorization": "Bearer test-key"}, + ) + assert resp.status_code == 200, resp.text + file_elem = captured["file_elem"] + assert not isinstance( + file_elem, (bytes, bytearray) + ), "batch upload must be a streamable handle, not in-memory bytes" + assert hasattr(file_elem, "read") and hasattr( + file_elem, "seek" + ), "batch upload must be a seekable file handle" + assert ( + captured["streamed_content"] == content + ), "the handle must stream the uploaded bytes" + + captured.clear() + resp = client.post( + "/v1/files", + files={"file": ("data.jsonl", content, "application/jsonl")}, + data={"purpose": "user_data"}, + headers={"Authorization": "Bearer test-key"}, + ) + assert resp.status_code == 200, resp.text + assert isinstance( + captured["file_elem"], (bytes, bytearray) + ), "non-batch upload must stay in-memory bytes" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.flaky(retries=3, delay=2) def test_target_storage_invokes_storage_backend( mocker: MockerFixture, monkeypatch, llm_router: Router diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 7a28eacca8d..e75720f9548 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2655,6 +2655,36 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() assert credentials["custom_llm_provider"] == "bedrock" +def test_get_deployment_credentials_with_provider_includes_bucket_name(): + """ + Regression: bucket_name must survive the CredentialLiteLLMParams filter so + managed-files batch retrieval can resolve the GCS/S3 bucket. Previously it was + dropped, causing "GCS bucket_name is required" when fetching batch output files. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "vertex-gemini", + "litellm_params": { + "model": "vertex_ai/gemini-3.5-flash", + "vertex_project": "my-project", + "vertex_location": "global", + "gcs_bucket_name": "my-batch-bucket", + }, + } + ], + ) + + credentials = router.get_deployment_credentials_with_provider( + model_id="vertex-gemini" + ) + + assert credentials is not None + assert credentials["gcs_bucket_name"] == "my-batch-bucket" + assert credentials["vertex_project"] == "my-project" + assert credentials["custom_llm_provider"] == "vertex_ai" + + def test_get_deployment_credentials_with_provider_resolves_credential_name(): """ Test that get_deployment_credentials_with_provider correctly resolves