scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

Merge pull request #31667 from BerriAI/litellm_backport_1_90_x_31036_31342_31653

Browse source 
chore(release): backport #31036, #31342, #31653 to stable/1.90.x and cut 1.90.1 (litellm-enterprise 0.1.43.post1)
This commit is contained in:
yuneng-jiang 2026-06-29 18:40:47 -07:00 • committed by GitHub
commit ca6149ab99
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
28 changed files with 2139 additions and 729 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -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")

4
enterprise/pyproject.toml
View file

@ -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==",

130
litellm/batches/batch_utils.py
View file

@ -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:

35
litellm/files/utils.py
View file

@ -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

1
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -14,6 +14,7 @@ OPTIONAL_KWARGS_KEYS = frozenset(
"azure_password",
"azure_scope",
"timeout",
"gcs_bucket_name",
"bucket_name",
"vertex_credentials",
"vertex_project",

40
litellm/litellm_core_utils/prompt_templates/common_utils.py
View file

@ -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.

18
litellm/llms/base_llm/files/transformation.py
View file

@ -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

127
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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
@ -42,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,
)
@ -81,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,
@ -103,6 +108,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 +134,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 +3243,20 @@ class BaseLLMHTTPHandler:
data=presigned_request["data"],
timeout=timeout,
)
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._upload_media(
client=sync_httpx_client,
url=api_base,
base_headers=headers,
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:
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,14 @@ class BaseLLMHTTPHandler:
input="",
api_key="",
additional_args={
"complete_input_dict": transformed_request,
# 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": (
"<streaming media upload>"
if isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request
else transformed_request
),
"api_base": api_base,
"headers": headers,
},
@ -3396,6 +3422,20 @@ class BaseLLMHTTPHandler:
data=presigned_request["data"],
timeout=timeout,
)
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._aupload_media(
client=async_httpx_client,
url=api_base,
base_headers=headers,
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:
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 +3479,83 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params,
)
# 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_in_blocks(byte_iter: Iterator[bytes], block_size: int) -> Iterator[bytes]:
buf = bytearray()
for piece in byte_iter:
buf.extend(piece)
while len(buf) >= block_size:
yield bytes(buf[:block_size])
del buf[:block_size]
if buf:
yield bytes(buf)
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}")
def _upload_media(
self,
*,
client: HTTPHandler,
url: str,
base_headers: Dict[str, str],
body_stream: BaseFileUploadStream,
content_type: str,
timeout: Optional[Union[float, httpx.Timeout]],
) -> httpx.Response:
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),
}
if timeout is not None:
kwargs["timeout"] = timeout
resp = client.client.post(url, **kwargs)
self._check_media_upload_response(resp)
return resp
async def _aupload_media(
self,
*,
client: AsyncHTTPHandler,
url: str,
base_headers: Dict[str, str],
body_stream: BaseFileUploadStream,
content_type: str,
timeout: Optional[Union[float, httpx.Timeout]],
) -> httpx.Response:
"""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()
self._check_media_upload_response(resp)
return resp
def create_batch(
self,
create_batch_data: "CreateBatchRequest",

82
litellm/llms/vertex_ai/files/handler.py
View file

@ -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,

509
litellm/llms/vertex_ai/files/transformation.py
View file

@ -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 StreamingMediaUploadConfig
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": <request_body>}
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,8 +389,7 @@ 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)
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)
@ -366,14 +435,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 +445,33 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
"""
2 Cases:
1. Handle basic file upload
2. Handle batch file upload (.jsonl)
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:
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 {
"streaming_media_upload": StreamingMediaUploadConfig(
body_stream=_OpenAIToVertexBatchUploadStream(
file_data,
self._map_openai_to_vertex_params,
),
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 +696,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 +744,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 +846,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",
)

20
litellm/proxy/client/exceptions.py
View file

@ -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))

9
litellm/proxy/client/keys.py
View file

@ -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

67
litellm/proxy/hooks/batch_rate_limiter.py
View file

@ -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 (

38
litellm/proxy/openai_files_endpoints/files_endpoints.py
View file

@ -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(

100
litellm/router_utils/batch_utils.py
View file

@ -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

17
litellm/types/files.py
View file

@ -321,3 +321,20 @@ class TwoStepFileUploadConfig(TypedDict, total=False):
upload_request: Required[TwoStepFileUploadRequest]
upload_url_location: Required[Literal["headers", "body"]]
upload_url_key: str
class StreamingMediaUploadConfig(TypedDict, total=False):
"""Drives a memory-bounded single-request upload (GCS simple/media upload).
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]
content_type: str

3
litellm/types/router.py
View file

@ -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

6
pyproject.toml
View file

@ -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"
@ -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",
@ -272,7 +272,7 @@ source-exclude = [
profile = "black"
[tool.commitizen]
version = "1.90.0"
version = "1.90.1"
version_files = [
"pyproject.toml:^version",
]

36
tests/batches_tests/test_openai_batches_and_files.py
View file

@ -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,27 @@ 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 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,
),
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()
file_name = "vertex_batch_completions.jsonl"
@ -536,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",

93
tests/router_unit_tests/test_router_batch_utils.py
View file

@ -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

47
tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py
View file

@ -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
@ -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_string(self):
async def test_jsonl_file_upload_returns_streaming_body(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 streaming-media config
carrying a streaming body (not a buffered bytes payload), so the handler
can stage the upload to a temp file and send it in one media request.
"""
# Create mock JSONL content
mock_jsonl_content = (
@ -164,10 +158,13 @@ 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 "streaming_media_upload" in transformed_request, (
f"Expected a streaming media upload config for JSONL, got {type(transformed_request)}"
)
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': ...}"
@pytest.mark.asyncio
async def test_mixed_file_types_in_sequence(self):
@ -208,7 +205,7 @@ class TestVertexAIBinaryFileUpload:
optional_params={},
litellm_params={},
)
assert isinstance(result2, str)
assert isinstance(result2, dict) and "streaming_media_upload" in result2
# Test 3: Upload another binary file
binary_content2 = b"\xc4\xe5\xf2\xe5\xeb"
@ -251,14 +248,12 @@ 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",
},
}
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"

586
tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py Normal file
View file

@ -0,0 +1,586 @@
"""
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 tempfile
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 _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 "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):
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_streaming_body_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 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:
"""
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"}}\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
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=media" in url
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()
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 _upload_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)} (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 = _upload_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_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))
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
def _async_handler_with(mock) -> AsyncHTTPHandler:
handler = AsyncHTTPHandler()
handler.client = httpx.AsyncClient(transport=httpx.MockTransport(mock))
return handler
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"),
"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,
)
# 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_binary_upload_uses_media_upload_type(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 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"))
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
@pytest.mark.asyncio
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, status: int = 200, timeout=None):
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={}
)
expected = _join_upload_body(transformed)
mock, state = _gcs_media_mock(status=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=timeout,
)
return expected, state, response
async def test_single_request_carries_whole_payload(self):
raw = _make_openai_jsonl_bytes(300)
expected, state, response = await self._run(raw)
# 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]
# 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_failed_upload_raises(self):
raw = _make_openai_jsonl_bytes(80)
with pytest.raises(Exception):
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 == []

232
tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py
View file

@ -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"

96
tests/test_litellm/proxy/client/test_keys.py
View file

@ -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

451
tests/test_litellm/proxy/hooks/test_batch_file_validation.py
View file

@ -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

77
tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py
View file

@ -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

30
tests/test_litellm/test_router.py
View file

@ -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

6
uv.lock generated
View file

@ -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]]