From 45b6ece18c15dfd9b9ec98783c40c002ccb3cb88 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 24 Jun 2026 13:19:57 -0700 Subject: [PATCH 1/6] 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 | 268 ++++++- litellm/llms/vertex_ai/files/handler.py | 82 +-- .../llms/vertex_ai/files/transformation.py | 522 ++++++------- litellm/proxy/hooks/batch_rate_limiter.py | 67 +- .../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 | 451 +++++++++--- .../test_files_endpoint.py | 77 ++ tests/test_litellm/test_router.py | 30 + 22 files changed, 2265 insertions(+), 698 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 8486e37384e..af4870bb1a5 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 fc3c25e0d95..c88f8b77dc2 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 fe34731759f..bf9ce3b0acb 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 8ac5b47c6e7..d1daa4eb80f 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 @@ -103,6 +105,7 @@ from litellm.types.llms.openai import ( ResponseInputParam, ResponsesAPIResponse, ) +from litellm.types.realtime import RealtimeQueryParams from litellm.types.rerank import RerankResponse from litellm.types.responses.main import DeleteResponseResult from litellm.types.router import GenericLiteLLMParams @@ -128,7 +131,6 @@ from litellm.types.vector_stores import ( VectorStoreSearchOptionalRequestParams, VectorStoreSearchResponse, ) -from litellm.types.realtime import RealtimeQueryParams from litellm.types.videos.main import VideoObject from litellm.utils import ( CustomStreamWrapper, @@ -3238,6 +3240,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 ): @@ -3319,7 +3338,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, }, @@ -3396,6 +3423,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 ): @@ -3439,6 +3483,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 5b691beccbf..91c604e6204 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -21,6 +21,7 @@ from typing import ( TYPE_CHECKING, Any, Dict, + Iterable, List, Literal, NoReturn, @@ -32,13 +33,15 @@ from typing import ( 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 ( + _count_entry_tokens, + _estimate_batch_entry_tokens, _extract_file_access_credentials, - _get_batch_job_input_file_usage, - _get_file_content_as_dictionary, - _get_models_from_batch_input_file_content, + _iter_batch_input_lines, ) from litellm.exceptions import RateLimitErrorCategory from litellm.integrations.custom_logger import CustomLogger @@ -537,6 +540,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 @@ -568,7 +573,38 @@ class _PROXY_BatchRateLimiter(CustomLogger): f"Expected bytes content from file retrieval for {file_id}, " f"got {type(file_content_bytes)}" ) - file_content_as_dict = _get_file_content_as_dictionary(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 @@ -578,17 +614,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, ) @@ -614,14 +645,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 standard auth helpers so the same model access rules the proxy enforces on `/chat/completions` apply here. @@ -640,10 +672,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 team_object = None if ( diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 944423632ef..7c19804dde3 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 ( @@ -97,16 +97,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 @@ -327,9 +329,15 @@ async def create_file( 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) @@ -454,13 +462,13 @@ async def create_file( ) # 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 607bfd584fd..b5285f11f8b 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -180,6 +180,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 bccb5eaaacb..8a2d5f33805 100644 --- a/tests/batches_tests/test_openai_batches_and_files.py +++ b/tests/batches_tests/test_openai_batches_and_files.py @@ -27,7 +27,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import StandardLoggingPayload import socket import httpx -from unittest.mock import patch, MagicMock +from unittest.mock import patch, MagicMock, AsyncMock def _can_resolve_openai(): @@ -513,10 +513,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..8c5305ee67b 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 a6f6e651487..1d4d39ec140 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -14,155 +14,131 @@ from fastapi import HTTPException from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + +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") + ] + + # --------------------------------------------------------------------------- # Token counter — covers all three batch payload shapes # --------------------------------------------------------------------------- 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"] + assert tokens == 6 # --------------------------------------------------------------------------- @@ -211,7 +187,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 @@ -250,7 +226,7 @@ async def test_pre_call_allows_all_team_models_key_when_model_in_team_allowlist( with patch("litellm.proxy.proxy_server.llm_router", None): await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), ) @@ -297,7 +273,7 @@ async def test_pre_call_uses_current_team_allowlist_for_all_team_models_key(): ): 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_info.value.status_code == 403 @@ -358,7 +334,7 @@ async def test_pre_call_allows_all_team_models_key_via_current_team_object(): ): await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, - file_content_as_dict=file_dict, + models=_models(file_dict), ) mock_get_team_object.assert_awaited_once() @@ -421,7 +397,7 @@ async def test_pre_call_denies_all_team_models_key_via_member_scope(): ): 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_info.value.status_code == 403 @@ -479,7 +455,7 @@ async def test_pre_call_fails_closed_when_current_team_fetch_fails_for_all_team_ ): 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_info.value.status_code == expected_status @@ -524,7 +500,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), ) @@ -744,7 +720,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], ) @@ -837,7 +813,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], ) @@ -863,11 +839,11 @@ 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": {}}]), ) @@ -1390,3 +1366,272 @@ async def test_count_input_file_usage_raises_on_non_bytes_content(): user_api_key_dict=UserAPIKeyAuth(api_key="sk", models=["*"]), data={}, ) + + +# Streaming input counting — peak memory must not scale with a full dict list +# --------------------------------------------------------------------------- + + +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 f42639cee8a..46ecb31e1c8 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 @@ -398,6 +398,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 c2aa4a095d4..f80acc40ccf 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3062,6 +3062,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 From 6a1d2b0c668360c3b19203c0e096ee3f1133cc13 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 25 Jun 2026 17:35:15 -0700 Subject: [PATCH 2/6] fix(proxy/client): redact api key from key/info client error messages (#31342) * fix(proxy/client): redact api key from key/info client error messages The keys management client builds GET /key/info?key= and lets the requests HTTPError propagate. str(HTTPError) renders the failing request URL verbatim ("... for url: .../key/info?key=sk-..."), so any caller that logs the exception leaks the full key; the 401 branch leaked the same way through UnauthorizedError(str(orig_exception)) Redact both branches with the existing redact_secrets helper so the secret-bearing query param is scrubbed to ?REDACTED while the status code, reason, and response object are preserved. Server-side responses already mask the key, so this closes the remaining client-side surface * fix: preserve key info unauthorized response --------- Co-authored-by: Cursor Agent (cherry picked from commit 71ee1a852a02787097c025d71f24202fdfc6b8c0) --- litellm/proxy/client/exceptions.py | 20 +++- litellm/proxy/client/keys.py | 9 +- tests/test_litellm/proxy/client/test_keys.py | 96 +++++++++++++++++++- 3 files changed, 118 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/client/exceptions.py b/litellm/proxy/client/exceptions.py index fffd1b78b7e..c4089381e30 100644 --- a/litellm/proxy/client/exceptions.py +++ b/litellm/proxy/client/exceptions.py @@ -2,18 +2,30 @@ from typing import Union import requests +from litellm.litellm_core_utils.secret_redaction import redact_string + + +def _redact_orig_exception( + orig_exception: Union[requests.exceptions.HTTPError, str], +) -> Union[requests.exceptions.HTTPError, str]: + if isinstance(orig_exception, requests.exceptions.HTTPError): + return requests.exceptions.HTTPError( + redact_string(str(orig_exception)), response=orig_exception.response + ) + return redact_string(str(orig_exception)) + class UnauthorizedError(Exception): """Exception raised when the API returns a 401 Unauthorized response.""" def __init__(self, orig_exception: Union[requests.exceptions.HTTPError, str]): - self.orig_exception = orig_exception - super().__init__(str(orig_exception)) + self.orig_exception = _redact_orig_exception(orig_exception) + super().__init__(str(self.orig_exception)) class NotFoundError(Exception): """Exception raised when the API returns a 404 Not Found response or indicates a resource was not found.""" def __init__(self, orig_exception: Union[requests.exceptions.HTTPError, str]): - self.orig_exception = orig_exception - super().__init__(str(orig_exception)) + self.orig_exception = _redact_orig_exception(orig_exception) + super().__init__(str(self.orig_exception)) diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index d8687cbad16..845b49d1581 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -2,6 +2,8 @@ from typing import Any, Dict, List, Optional, Union import requests +from litellm.litellm_core_utils.secret_redaction import redact_string + from .exceptions import UnauthorizedError @@ -314,6 +316,9 @@ class KeysManagementClient: response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: + redacted_message = redact_string(str(e)) if e.response.status_code == 401: - raise UnauthorizedError(e) - raise + raise UnauthorizedError(e) from None + raise requests.exceptions.HTTPError( + redacted_message, response=e.response + ) from None diff --git a/tests/test_litellm/proxy/client/test_keys.py b/tests/test_litellm/proxy/client/test_keys.py index 136408e01ac..620daefb39e 100644 --- a/tests/test_litellm/proxy/client/test_keys.py +++ b/tests/test_litellm/proxy/client/test_keys.py @@ -1,5 +1,6 @@ import os import sys +import traceback import pytest import requests @@ -11,7 +12,7 @@ sys.path.insert( import responses -from litellm.proxy.client.exceptions import UnauthorizedError +from litellm.proxy.client.exceptions import NotFoundError, UnauthorizedError from litellm.proxy.client.keys import KeysManagementClient @@ -420,3 +421,96 @@ def test_info_server_error(client): ) with pytest.raises(requests.exceptions.HTTPError): client.info(key="test-key") + + +LEAKY_KEY = "sk-1234567890abcdefghijklmnop" + + +def _render_full_traceback(exc: BaseException) -> str: + return "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)) + + +@responses.activate +def test_info_not_found_redacts_key_everywhere(client): + """A 404 must not echo the raw key embedded in the request URL. + + Covers str(exc) and the rendered traceback, since the chain through + __cause__ / __context__ is what logging.exception() and the default + excepthook print. + """ + responses.add( + responses.GET, + f"{client._base_url}/key/info?key={LEAKY_KEY}", + status=404, + json={"error": {"message": "Key not found", "code": "404"}}, + ) + with pytest.raises(requests.exceptions.HTTPError) as excinfo: + client.info(key=LEAKY_KEY) + + exc = excinfo.value + assert LEAKY_KEY not in str(exc) + assert "REDACTED" in str(exc) + assert LEAKY_KEY not in _render_full_traceback(exc) + assert exc.__cause__ is None and exc.__suppress_context__ + assert exc.response is not None + assert exc.response.status_code == 404 + assert exc.request is not None + # Known residual: the live request URL still carries the key, since the + # response is preserved so callers keep status_code / text. str(exc) and the + # traceback are scrubbed; the URL-borne key is the root issue tracked in + # LIT-4013 (move the lookup key out of the query string server-side). + assert LEAKY_KEY in exc.response.request.url + + +@responses.activate +def test_info_unauthorized_redacts_key_everywhere(client): + """A 401 surfaced as UnauthorizedError must not echo the raw key in the + message, the retained original, or the rendered traceback chain.""" + responses.add( + responses.GET, + f"{client._base_url}/key/info?key={LEAKY_KEY}", + status=401, + json={"error": "Unauthorized"}, + ) + with pytest.raises(UnauthorizedError) as excinfo: + client.info(key=LEAKY_KEY) + + exc = excinfo.value + assert LEAKY_KEY not in str(exc) + assert "REDACTED" in str(exc) + assert LEAKY_KEY not in str(exc.orig_exception) + assert LEAKY_KEY not in _render_full_traceback(exc) + assert exc.__cause__ is None and exc.__suppress_context__ + assert isinstance(exc.orig_exception, requests.exceptions.HTTPError) + assert exc.orig_exception.response is not None + assert exc.orig_exception.response.status_code == 401 + + +def _http_error_with_key(prefix: str, status: int) -> requests.exceptions.HTTPError: + resp = requests.Response() + resp.status_code = status + return requests.exceptions.HTTPError( + f"{prefix} for url: http://x/key/info?key={LEAKY_KEY}", response=resp + ) + + +def test_unauthorized_error_redacts_wrapped_key(): + """UnauthorizedError scrubs the key in str(exc) and in the retained + orig_exception, while preserving the response for structured access.""" + wrapped = UnauthorizedError( + _http_error_with_key("401 Client Error: Unauthorized", 401) + ) + assert LEAKY_KEY not in str(wrapped) + assert "REDACTED" in str(wrapped) + assert LEAKY_KEY not in str(wrapped.orig_exception) + assert wrapped.orig_exception.response.status_code == 401 + + +def test_not_found_error_redacts_wrapped_key(): + """NotFoundError scrubs the key in str(exc) and in the retained + orig_exception, while preserving the response for structured access.""" + wrapped = NotFoundError(_http_error_with_key("404 Client Error: Not Found", 404)) + assert LEAKY_KEY not in str(wrapped) + assert "REDACTED" in str(wrapped) + assert LEAKY_KEY not in str(wrapped.orig_exception) + assert wrapped.orig_exception.response.status_code == 404 From ccd7bb67542eaea6e466cf8795c51fb05e822882 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 29 Jun 2026 17:31:32 -0700 Subject: [PATCH 3/6] fix(vertex_ai/files): single media upload for batch files to fix 499s on large uploads (#31653) * fix(vertex_ai/files): upload batch files in a single media request to fix 499s on large uploads PR #31036 switched the vertex batch file upload from a single GCS media upload to a chunked resumable session. The resumable path sends the body as many sequential PUTs, each waiting a full round-trip to GCS before the next, so a multi-GB upload accumulates hundreds of round-trips and overruns the client/load-balancer request timeout, surfacing as 499s (client closed connection) on files as small as 500MB. This was a regression from the last-known-good commit, where the upload completed as one continuous request. Revert the batch upload to a single uploadType=media request, but stage the transformed payload to a temp file first so peak memory stays bounded (the goal of the resumable rewrite) without the per-chunk round-trips. The temp file is closed deterministically (TemporaryFile unlinks on close), not left to the GC. The now-unused resumable chunked-upload plumbing is removed. Also swap the per-row transform's stdlib json for orjson (parse + serialize), which is ~4x faster on this hot path; the streaming body now emits compact orjson bytes. The request stays synchronous, so the returned file object is real and POST /v1/batches keeps working immediately against the uploaded object. Tests: single media request carries the whole payload with a real Content-Length (no chunked transfer-encoding); failed upload raises; the staged temp file is closed deterministically; byte-for-byte transform parity. * test(vertex_ai/files): mock single media upload POST instead of removed resumable method test_avertex_batch_prediction patched BaseLLMHTTPHandler._aresumable_chunked_upload, which was removed when the batch jsonl upload moved from a chunked resumable GCS session to a single uploadType=media request. Patch the raw httpx.AsyncClient.post that _astage_and_upload_media issues so the real staging, upload and response transform run while the GCS object response is mocked, and assert the media URL and Content-Type. * fix(vertex_ai/files): forward request timeout to media upload, drop orjson, sort imports Forward the per-request timeout through _stage_and_upload_media / _astage_and_upload_media to the GCS POST. Every other upload branch forwards it; the new media path was dropping it, so a caller-provided timeout was silently ignored (the files path passes 600s by default, but a custom request_timeout would not have reached this upload). Regression test asserts the resolved timeout reaches the request (mutation-verified). Revert the orjson swap in the batch transform: importing orjson at module load in this core-path file broke `import litellm` on environments without orjson (the Windows import test). Back to stdlib json; the upload leg dominates large uploads anyway, so the transform-side win was marginal. Fix import ordering in llm_http_handler.py (I001) introduced by the new imports. * fix(vertex_ai/files): stream batch upload to GCS instead of staging to a temp file Addresses a disk-exhaustion concern: staging the full transformed batch body to a local temp file before the GCS request meant an authenticated user could fill the proxy's temp volume with large concurrent uploads (on top of Starlette's input spool). GCS's simple/media upload accepts chunked transfer-encoding, so stream the transform straight to the single media request instead. Each block is produced on a worker thread (the transform never runs on the event loop) and sent chunked, so the body is neither buffered in memory nor written to disk, and the upload is still one continuous request (no per-chunk round-trips, no 499). Drops the temp-file staging, the tempfile/IO imports, and Content-Length computation. Regression test asserts the upload streams (chunked transfer-encoding, no Content-Length) and creates no temp file; mutation-verified that reintroducing staging fails it. (cherry picked from commit 85840aef513ce93e11d7c203996426c14d4e4582) --- litellm/llms/custom_httpx/llm_http_handler.py | 285 +++++------------ .../llms/vertex_ai/files/transformation.py | 25 +- litellm/types/files.py | 17 +- .../test_openai_batches_and_files.py | 38 ++- .../test_vertex_ai_binary_file_upload.py | 43 +-- .../files/test_vertex_ai_files_streaming.py | 294 ++++++------------ 6 files changed, 217 insertions(+), 485 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d1daa4eb80f..70fb4cd95cf 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -44,7 +44,10 @@ from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.base_llm.containers.transformation import BaseContainerConfig from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.base_llm.evals.transformation import BaseEvalsAPIConfig -from litellm.llms.base_llm.files.transformation import BaseFilesConfig +from litellm.llms.base_llm.files.transformation import ( + BaseFilesConfig, + BaseFileUploadStream, +) from litellm.llms.base_llm.google_genai.transformation import ( BaseGoogleGenAIGenerateContentConfig, ) @@ -83,7 +86,7 @@ from litellm.types.containers.main import ( ContainerObject, DeleteContainerResult, ) -from litellm.types.files import TwoStepFileUploadConfig +from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig from litellm.types.integrations.custom_logger import ( AgenticLoopPlan, AgenticLoopRequestPatch, @@ -3240,18 +3243,15 @@ class BaseLLMHTTPHandler: data=presigned_request["data"], timeout=timeout, ) - elif ( - isinstance(transformed_request, dict) - and "resumable_chunked_upload" in transformed_request - ): + elif isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request: + media_cfg = cast(StreamingMediaUploadConfig, transformed_request["streaming_media_upload"]) try: - upload_response = self._resumable_chunked_upload( + upload_response = self._upload_media( client=sync_httpx_client, - initiate_url=api_base, + url=api_base, base_headers=headers, - config=cast(Dict[str, Any], transformed_request)[ - "resumable_chunked_upload" - ], + body_stream=cast(BaseFileUploadStream, media_cfg["body_stream"]), + content_type=media_cfg.get("content_type") or "application/octet-stream", timeout=timeout, ) except Exception as e: @@ -3338,13 +3338,12 @@ class BaseLLMHTTPHandler: input="", api_key="", additional_args={ - # A resumable upload config holds a reference to the (potentially + # A streaming 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 + "" + if isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request else transformed_request ), "api_base": api_base, @@ -3423,18 +3422,15 @@ class BaseLLMHTTPHandler: data=presigned_request["data"], timeout=timeout, ) - elif ( - isinstance(transformed_request, dict) - and "resumable_chunked_upload" in transformed_request - ): + elif isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request: + media_cfg = cast(StreamingMediaUploadConfig, transformed_request["streaming_media_upload"]) try: - upload_response = await self._aresumable_chunked_upload( + upload_response = await self._aupload_media( client=async_httpx_client, - initiate_url=api_base, + url=api_base, base_headers=headers, - config=cast(Dict[str, Any], transformed_request)[ - "resumable_chunked_upload" - ], + body_stream=cast(BaseFileUploadStream, media_cfg["body_stream"]), + content_type=media_cfg.get("content_type") or "application/octet-stream", timeout=timeout, ) except Exception as e: @@ -3483,222 +3479,81 @@ 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 + # The fine-grained transform stream (one piece per JSONL row) is regrouped + # into blocks of this size before upload, so the request yields a manageable + # number of chunks; never more than one block is buffered. + _MEDIA_UPLOAD_BLOCK_SIZE = 4 * 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. - """ + def _iter_in_blocks(byte_iter: Iterator[bytes], block_size: int) -> Iterator[bytes]: buf = bytearray() for piece in byte_iter: buf.extend(piece) - while len(buf) >= chunk_size: - yield bytes(buf[:chunk_size]) - del buf[:chunk_size] + while len(buf) >= block_size: + yield bytes(buf[:block_size]) + del buf[:block_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}" + def _check_media_upload_response(self, resp: httpx.Response) -> None: + if resp.status_code not in (200, 201): + resp.raise_for_status() + raise ValueError(f"media upload: unexpected status {resp.status_code}") - @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( + def _upload_media( 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, + base_headers: Dict[str, str], + body_stream: BaseFileUploadStream, + content_type: str, timeout: Optional[Union[float, httpx.Timeout]], ) -> httpx.Response: - headers = { - **base_headers, - "Content-Range": self._resumable_content_range(offset, len(data), is_final), + headers = {**base_headers, "Content-Type": content_type} + kwargs: Dict[str, Any] = { + "headers": headers, + "content": self._iter_in_blocks(body_stream.iter_bytes(), self._MEDIA_UPLOAD_BLOCK_SIZE), } - 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}") + if timeout is not None: + kwargs["timeout"] = timeout + resp = client.client.post(url, **kwargs) + self._check_media_upload_response(resp) return resp - async def _aresumable_chunked_upload( + async def _aupload_media( 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, + base_headers: Dict[str, str], + body_stream: BaseFileUploadStream, + content_type: str, 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) + """Stream the transformed body straight to a single media upload. Each + block is produced on a worker thread (the transform never runs on the + event loop) and sent with chunked transfer-encoding, so the body is + neither buffered in memory nor staged to disk, and the upload is one + continuous request rather than the many sequential round-trips of the + resumable path that overran client/LB timeouts.""" + headers = {**base_headers, "Content-Type": content_type} + block_iter = iter(self._iter_in_blocks(body_stream.iter_bytes(), self._MEDIA_UPLOAD_BLOCK_SIZE)) + done = object() + + async def _abody() -> AsyncIterator[bytes]: + while True: + block = await asyncio.to_thread(next, block_iter, done) + if block is done: + break + yield cast(bytes, block) + + kwargs: Dict[str, Any] = {"headers": headers, "content": _abody()} + if timeout is not None: + kwargs["timeout"] = timeout + resp = await client.client.post(url, **kwargs) 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}") + self._check_media_upload_response(resp) return resp def create_batch( diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index d5164d8c1c2..5abaedae5a8 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -60,7 +60,7 @@ from litellm.types.llms.openai import ( OpenAIFileObject, PathLike, ) -from litellm.types.files import ResumableChunkedUploadConfig +from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.vertex_ai import GcsBucketResponse from litellm.types.utils import LlmProviders, ModelResponse @@ -389,21 +389,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): raise ValueError("file is required") if purpose is None: raise ValueError("purpose is required") - _, 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) - # 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}" + endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_object_name}" api_base = api_base or "https://storage.googleapis.com" if not api_base: raise ValueError("api_base is required") @@ -455,8 +445,9 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ 2 Cases: 1. Handle basic file upload - 2. Handle batch file upload (.jsonl), streamed to a GCS resumable - session so large uploads stay memory-bounded. + 2. Handle batch file upload (.jsonl), staged to a temp file and uploaded + in a single media request so large uploads stay memory-bounded without + the per-chunk round-trips of a resumable session. """ file_data = create_file_data.get("file") if file_data is None: @@ -468,14 +459,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): content_type=content_type, ): return { - "resumable_chunked_upload": ResumableChunkedUploadConfig( + "streaming_media_upload": StreamingMediaUploadConfig( body_stream=_OpenAIToVertexBatchUploadStream( file_data, self._map_openai_to_vertex_params, ), - initiate_headers={ - "X-Upload-Content-Type": "application/json", - }, + content_type="application/json", ) } diff --git a/litellm/types/files.py b/litellm/types/files.py index 1b2d7e30f1f..97c4dbb3abc 100644 --- a/litellm/types/files.py +++ b/litellm/types/files.py @@ -323,19 +323,18 @@ class TwoStepFileUploadConfig(TypedDict, total=False): upload_url_key: str -class ResumableChunkedUploadConfig(TypedDict, total=False): - """Drives a memory-bounded resumable upload (GCS JSON API). +class StreamingMediaUploadConfig(TypedDict, total=False): + """Drives a memory-bounded single-request upload (GCS simple/media upload). - 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. + The handler stages ``body_stream`` to a temp file off the event loop (so peak + memory stays bounded), then PUTs/POSTs it in one request with a known + Content-Length. Unlike a resumable chunked upload this incurs no per-chunk + round-trips, so a multi-GB upload finishes in one continuous transfer instead + of hundreds of sequential PUTs that overrun client/LB timeouts. ``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] + content_type: str diff --git a/tests/batches_tests/test_openai_batches_and_files.py b/tests/batches_tests/test_openai_batches_and_files.py index 8a2d5f33805..0a49b3d77d1 100644 --- a/tests/batches_tests/test_openai_batches_and_files.py +++ b/tests/batches_tests/test_openai_batches_and_files.py @@ -513,25 +513,26 @@ async def test_avertex_batch_prediction(monkeypatch): mock_response.status_code = 200 return mock_response - # 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"), - ) + # Batch jsonl creation now stages the body to a temp file and issues a single + # uploadType=media POST against the raw httpx.AsyncClient (client.client) inside + # _astage_and_upload_media, not AsyncHTTPHandler.post. Patch that raw POST so the + # real staging/upload + response transform run while the GCS object response is + # mocked; AsyncHTTPHandler.post still handles the batch-prediction call. 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, ), + patch.object( + httpx.AsyncClient, + "post", + new_callable=AsyncMock, + return_value=httpx.Response( + 200, + json=mock_file_response, + request=httpx.Request("POST", "https://storage.googleapis.com/upload"), + ), + ) as mock_gcs_upload, ): litellm.set_verbose = True litellm._turn_on_debug() @@ -552,6 +553,15 @@ async def test_avertex_batch_prediction(monkeypatch): == "gs://litellm-local/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/5f7b99ad-9203-4430-98bf-3b45451af4cb" ) + mock_gcs_upload.assert_awaited_once() + upload_url = str(mock_gcs_upload.call_args.args[0]) + assert "uploadType=media" in upload_url + assert "/b/litellm-local/o" in upload_url + assert ( + mock_gcs_upload.call_args.kwargs["headers"]["Content-Type"] + == "application/json" + ) + # Create batch create_batch_response = await litellm.acreate_batch( completion_window="24h", 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 122518d4acb..d2ee9d7d659 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 @@ -37,9 +37,7 @@ class TestVertexAIBinaryFileUpload: # Create mock PDF binary data (with non-UTF-8 bytes) # PDF files start with %PDF- and contain binary data mock_pdf_content = b"%PDF-1.4\n%\xc4\xe5\xf2\xe5\xeb\xa7\xf3\xa0\xd0\xc4\xc6\n" - mock_pdf_content += ( - b"\x00\x01\x02\x03\xff\xfe\xfd" * 100 - ) # Add more binary data + mock_pdf_content += b"\x00\x01\x02\x03\xff\xfe\xfd" * 100 # Add more binary data # Create file object file_obj = io.BytesIO(mock_pdf_content) @@ -60,14 +58,12 @@ class TestVertexAIBinaryFileUpload: ) # Verify the transformation returns bytes (not string) - assert isinstance( - transformed_request, bytes - ), f"Expected bytes for binary file, got {type(transformed_request)}" + assert isinstance(transformed_request, bytes), ( + f"Expected bytes for binary file, got {type(transformed_request)}" + ) # Verify the bytes match the original content - assert ( - transformed_request == mock_pdf_content - ), "Transformed request should preserve binary content exactly" + assert transformed_request == mock_pdf_content, "Transformed request should preserve binary content exactly" # Verify that the bytes contain non-UTF-8 characters # This should raise UnicodeDecodeError if we try to decode @@ -132,16 +128,14 @@ class TestVertexAIBinaryFileUpload: pytest.fail(f"httpx should accept bytes in data parameter: {e}") # Document the expected behavior - assert isinstance( - mock_binary_data, bytes - ), "Binary file data should remain as bytes" + assert isinstance(mock_binary_data, bytes), "Binary file data should remain as bytes" @pytest.mark.asyncio - async def test_jsonl_file_upload_returns_resumable_stream(self): + async def test_jsonl_file_upload_returns_streaming_body(self): """ - Test that JSONL batch files are transformed into a resumable-upload config + Test that JSONL batch files are transformed into a streaming-media config carrying a streaming body (not a buffered bytes payload), so the handler - can stream the upload to GCS in bounded chunks. + can stage the upload to a temp file and send it in one media request. """ # Create mock JSONL content mock_jsonl_content = ( @@ -164,16 +158,13 @@ class TestVertexAIBinaryFileUpload: litellm_params={}, ) - assert ( - isinstance(transformed_request, dict) - and "resumable_chunked_upload" in transformed_request - ), f"Expected a resumable upload config for JSONL, got {type(transformed_request)}" + assert isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request, ( + f"Expected a streaming media upload config for JSONL, got {type(transformed_request)}" + ) - stream = transformed_request["resumable_chunked_upload"]["body_stream"] + stream = transformed_request["streaming_media_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': ...}" + assert "request" in decoded, "JSONL transform must wrap each row in {'request': ...}" @pytest.mark.asyncio async def test_mixed_file_types_in_sequence(self): @@ -214,7 +205,7 @@ class TestVertexAIBinaryFileUpload: optional_params={}, litellm_params={}, ) - assert isinstance(result2, dict) and "resumable_chunked_upload" in result2 + assert isinstance(result2, dict) and "streaming_media_upload" in result2 # Test 3: Upload another binary file binary_content2 = b"\xc4\xe5\xf2\xe5\xeb" @@ -264,7 +255,5 @@ class TestVertexAIBinaryFileUpload: }, } - assert ( - expected_behavior["binary_files"]["encoding"] == "none - preserve raw bytes" - ) + assert expected_behavior["binary_files"]["encoding"] == "none - preserve raw bytes" assert expected_behavior["text_files"]["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 index cd556c48b6b..2e3280c0ed1 100644 --- 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 @@ -20,6 +20,7 @@ replaced by a list-based pipeline: import gc import io import json +import tempfile import time import tracemalloc @@ -41,15 +42,15 @@ from litellm.llms.vertex_ai.files.transformation import ( 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 _upload_stream(transformed) -> BaseFileUploadStream: + """Pull the streaming body out of the upload transform result.""" + return transformed["streaming_media_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, dict) and "streaming_media_upload" in transformed: + return b"".join(_upload_stream(transformed).iter_bytes()) if isinstance(transformed, BaseFileUploadStream): return b"".join(transformed.iter_bytes()) if isinstance(transformed, str): @@ -83,17 +84,13 @@ def _reference_vertex_jsonl_string(cfg: VertexAIFilesConfig, content: str) -> st 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 - ) - ) + 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): + def test_transform_create_file_request_returns_streaming_body_parity(self): cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(300) request: CreateFileRequest = { @@ -105,14 +102,12 @@ class TestStreamingOutputParity: 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") - ) + # A batch upload must be a streaming-media config carrying a streaming + # body, so the handler can stream it to GCS; a buffered bytes/str return + # would defeat the OOM fix. + assert isinstance(out, dict) and "streaming_media_upload" in out + assert isinstance(_upload_stream(out), BaseFileUploadStream) + assert _join_upload_body(out).decode("utf-8") == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) class TestFileLikeInputNotPartiallyConsumed: @@ -215,9 +210,7 @@ class TestStreamingLineIterator: def seek(self, *args): raise io.UnsupportedOperation("not seekable") - handle = _NonSeekable( - b'{"custom_id": "request-0"}\n{"custom_id": "request-1"}\n' - ) + 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)) @@ -235,13 +228,8 @@ class TestGetObjectNameLazyParse: 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" - ) + raw = b'{"custom_id": "r-0", "body": {"model": "gemini-2.5-flash"}}\ngarbage 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 @@ -278,15 +266,11 @@ class TestStreamingPeakMemory: 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(): + 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) - ) + 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 @@ -305,9 +289,7 @@ class TestStreamingPeakMemory: # 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" + assert peak / len(raw) < 2.0, "get_object_name should not copy the whole payload" class TestPathSourcedStreaming: @@ -341,12 +323,10 @@ class TestPathSourcedStreaming: litellm_params={"gcs_bucket_name": "test-bucket"}, data=data, ) - assert "uploadType=resumable" in url + assert "uploadType=media" 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 + out = cfg.transform_create_file_request(model="", create_file_data=data, optional_params={}, litellm_params={}) + assert isinstance(out, dict) and "streaming_media_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() @@ -371,7 +351,7 @@ class TestPathSourcedStreaming: out = cfg.transform_create_file_request( model="", create_file_data=data, optional_params={}, litellm_params={} ) - for _ in _resumable_stream(out).iter_bytes(): + for _ in _upload_stream(out).iter_bytes(): pass # drain without accumulating gc.collect() @@ -384,20 +364,15 @@ class TestPathSourcedStreaming: # 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})" - ) + assert peak < len(raw) * 0.3, f"peak {peak} not bounded vs payload {len(raw)} (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) + out = cfg.transform_create_file_request(model="", create_file_data=data, optional_params={}, litellm_params={}) + stream = _upload_stream(out) first = b"".join(stream.iter_bytes()) second = b"".join(stream.iter_bytes()) assert first == second and len(first) > 0 @@ -436,25 +411,20 @@ def _logging_obj() -> Logging: ) -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": []} +def _gcs_media_mock(status: int = 200): + """A fake GCS simple-media endpoint: one request carries the whole object; + capture the body and headers and return the object resource.""" + state = {"received": bytearray(), "methods": [], "urls": [], "headers": [], "timeouts": []} 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) + state["headers"].append(dict(request.headers)) + # httpx records the resolved per-request timeout here, so the test can + # assert the caller's timeout was forwarded rather than the client default. + state["timeouts"].append(request.extensions.get("timeout")) + state["received"].extend(await request.aread()) + return httpx.Response(status, json=_GCS_OBJECT_JSON) return handler, state @@ -465,8 +435,8 @@ def _async_handler_with(mock) -> AsyncHTTPHandler: return handler -class TestResumableUploadUrl: - def test_batch_jsonl_uses_resumable_upload_type(self): +class TestUploadUrl: + def test_batch_jsonl_uses_media_upload_type(self): cfg = VertexAIFilesConfig() request: CreateFileRequest = { "file": ("batch.jsonl", _make_openai_jsonl_bytes(3), "application/jsonl"), @@ -480,29 +450,12 @@ class TestResumableUploadUrl: litellm_params={"gcs_bucket_name": "test-bucket"}, data=request, ) - assert "uploadType=resumable" in url - assert "uploadType=media" not in url + # A single media upload is one continuous transfer (no per-chunk + # round-trips), which is what keeps large uploads under client/LB timeouts. + assert "uploadType=media" in url + assert "uploadType=resumable" 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): + def test_binary_upload_uses_media_upload_type(self): cfg = VertexAIFilesConfig() request: CreateFileRequest = { "file": ("doc.pdf", b"%PDF-1.4 binary", "application/pdf"), @@ -520,14 +473,12 @@ class TestResumableUploadUrl: assert "uploadType=resumable" not in url -class TestResumableStreamBody: +class TestUploadStreamBody: 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")) + 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; @@ -545,59 +496,19 @@ class TestResumableStreamBody: # an empty body silently. cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(40) - stream = _OpenAIToVertexBatchUploadStream( - io.BytesIO(raw), cfg._map_openai_to_vertex_params - ) + 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.""" +class TestStreamingMediaUpload: + """End-to-end against a faked GCS media endpoint. These fail if the handler + buffers the payload in memory, drops bytes, omits Content-Length (which would + flip httpx to chunked transfer-encoding), or makes more than one request.""" - async def _run(self, raw: bytes, chunk_size: int, final_status: int = 200): + async def _run(self, raw: bytes, status: int = 200, timeout=None): cfg = VertexAIFilesConfig() request: CreateFileRequest = { "file": ("batch.jsonl", raw, "application/jsonl"), @@ -614,11 +525,8 @@ class TestResumableUploadProtocol: 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) + mock, state = _gcs_media_mock(status=status) response = await BaseLLMHTTPHandler().async_create_file( transformed_request=transformed, litellm_params={}, @@ -627,70 +535,52 @@ class TestResumableUploadProtocol: api_base=api_base, logging_obj=_logging_obj(), client=_async_handler_with(mock), - timeout=None, + timeout=timeout, ) - return expected, state, response, session_url, api_base + return expected, state, response - async def test_streams_in_chunks_and_reassembles(self): + async def test_single_request_carries_whole_payload(self): raw = _make_openai_jsonl_bytes(300) - chunk_size = 4096 - expected, state, response, session_url, api_base = await self._run( - raw, chunk_size - ) + expected, state, response = await self._run(raw) - # 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" + # Exactly one request (the single media upload), and it lands on the + # media endpoint, not a resumable session. + assert state["methods"] == ["POST"] + assert "uploadType=media" in state["urls"][0] - # 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. + # The body is streamed with chunked transfer-encoding and no + # Content-Length, which is what proves it is neither buffered in memory + # nor staged to a temp file (the disk-exhaustion guard) before sending. + headers = state["headers"][0] + assert headers.get("transfer-encoding") == "chunked" + assert "content-length" not in headers + # httpx reassembles the chunked body; GCS receives exactly the transform. 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): + async def test_failed_upload_raises(self): raw = _make_openai_jsonl_bytes(80) with pytest.raises(Exception): - await self._run(raw, chunk_size=4096, final_status=403) + await self._run(raw, status=403) + + async def test_request_timeout_is_forwarded(self): + # The caller's per-request timeout must reach the GCS upload; every other + # upload branch forwards it. httpx records the resolved timeout in + # request.extensions["timeout"]; a dropped timeout would show the client + # default instead of the value passed here. + raw = _make_openai_jsonl_bytes(20) + _, state, _ = await self._run(raw, timeout=httpx.Timeout(137.0)) + forwarded = state["timeouts"][0] + assert forwarded is not None + assert forwarded.get("read") == 137.0 and forwarded.get("write") == 137.0 + + async def test_upload_does_not_stage_to_disk(self, monkeypatch): + # Disk-exhaustion guard: the transformed body must stream to GCS, never be + # written to a temp file first. If any tempfile is created during the + # upload, an attacker could fill the proxy's temp volume with large + # concurrent uploads. + created = [] + real_tempfile = tempfile.TemporaryFile + monkeypatch.setattr(tempfile, "TemporaryFile", lambda *a, **k: (created.append(1), real_tempfile(*a, **k))[1]) + await self._run(_make_openai_jsonl_bytes(50)) + assert created == [] From 3d06452b362523f2994d85f3cd0f5e1f1351188e Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 29 Jun 2026 18:15:00 -0700 Subject: [PATCH 4/6] chore(release): bump litellm-enterprise 0.1.43 -> 0.1.43.post1 for stable/1.90.x --- enterprise/pyproject.toml | 4 ++-- pyproject.toml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index b032942427c..d0538a2b258 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.43" +version = "0.1.43.post1" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.43" +version = "0.1.43.post1" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/pyproject.toml b/pyproject.toml index 91de8683968..f677f00e0e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,7 +63,7 @@ proxy = [ "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.26.0,<2.0", "litellm-proxy-extras==0.4.74", - "litellm-enterprise==0.1.43", + "litellm-enterprise==0.1.43.post1", "RestrictedPython>=8.1,<9.0", "rich>=13.9.4,<14.0", "polars>=1.38.1,<2.0", From 22421112ea0a0a872fa9f8d3cced41e2a51a3ed3 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 29 Jun 2026 18:15:57 -0700 Subject: [PATCH 5/6] =?UTF-8?q?bump:=20version=201.90.0=20=E2=86=92=201.90?= =?UTF-8?q?.1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f677f00e0e9..8c8cc6aa290 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.90.0" +version = "1.90.1" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -272,7 +272,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.90.0" +version = "1.90.1" version_files = [ "pyproject.toml:^version", ] From 92da671c689f8491225ea11501081f9f12a6a2c5 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 29 Jun 2026 18:16:13 -0700 Subject: [PATCH 6/6] chore: refresh uv.lock for 1.90.1 and litellm-enterprise 0.1.43.post1 --- uv.lock | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/uv.lock b/uv.lock index f3ca7b4e626..594b68b0bd2 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-06-17T22:13:32.966924Z" +exclude-newer = "2026-06-27T01:16:05.524641Z" exclude-newer-span = "P3D" [manifest] @@ -3245,7 +3245,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.90.0" +version = "1.90.1" source = { editable = "." } dependencies = [ { name = "aiohttp" }, @@ -3610,7 +3610,7 @@ proxy-dev = [ [[package]] name = "litellm-enterprise" -version = "0.1.43" +version = "0.1.43.post1" source = { editable = "enterprise" } [[package]]