mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
Merge pull request #31662 from BerriAI/litellm_backport_1_89_x_bp-31342-31653-0629
chore(release): backport #31342, #31653 to stable/1.89.x and cut 1.89.5
This commit is contained in:
yuneng-jiang
GitHub
commit
b216c4d80d
11 changed files with 339 additions and 496 deletions
|
|
@ -41,7 +41,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,
|
||||
)
|
||||
|
|
@ -80,7 +83,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,
|
||||
|
|
@ -3232,18 +3235,15 @@ class BaseLLMHTTPHandler:
|
|||
data=presigned_request["data"],
|
||||
timeout=timeout,
|
||||
)
|
||||
elif (
|
||||
isinstance(transformed_request, dict)
|
||||
and "resumable_chunked_upload" in transformed_request
|
||||
):
|
||||
elif isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request:
|
||||
media_cfg = cast(StreamingMediaUploadConfig, transformed_request["streaming_media_upload"])
|
||||
try:
|
||||
upload_response = self._resumable_chunked_upload(
|
||||
upload_response = self._upload_media(
|
||||
client=sync_httpx_client,
|
||||
initiate_url=api_base,
|
||||
url=api_base,
|
||||
base_headers=headers,
|
||||
config=cast(Dict[str, Any], transformed_request)[
|
||||
"resumable_chunked_upload"
|
||||
],
|
||||
body_stream=cast(BaseFileUploadStream, media_cfg["body_stream"]),
|
||||
content_type=media_cfg.get("content_type") or "application/octet-stream",
|
||||
timeout=timeout,
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -3330,13 +3330,12 @@ class BaseLLMHTTPHandler:
|
|||
input="",
|
||||
api_key="",
|
||||
additional_args={
|
||||
# A resumable upload config holds a reference to the (potentially
|
||||
# A streaming upload config holds a reference to the (potentially
|
||||
# huge) upload payload; logging deep-copies additional_args, so log
|
||||
# a placeholder instead of re-materializing the payload.
|
||||
"complete_input_dict": (
|
||||
"<resumable chunked upload>"
|
||||
if isinstance(transformed_request, dict)
|
||||
and "resumable_chunked_upload" in transformed_request
|
||||
"<streaming media upload>"
|
||||
if isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request
|
||||
else transformed_request
|
||||
),
|
||||
"api_base": api_base,
|
||||
|
|
@ -3415,18 +3414,15 @@ class BaseLLMHTTPHandler:
|
|||
data=presigned_request["data"],
|
||||
timeout=timeout,
|
||||
)
|
||||
elif (
|
||||
isinstance(transformed_request, dict)
|
||||
and "resumable_chunked_upload" in transformed_request
|
||||
):
|
||||
elif isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request:
|
||||
media_cfg = cast(StreamingMediaUploadConfig, transformed_request["streaming_media_upload"])
|
||||
try:
|
||||
upload_response = await self._aresumable_chunked_upload(
|
||||
upload_response = await self._aupload_media(
|
||||
client=async_httpx_client,
|
||||
initiate_url=api_base,
|
||||
url=api_base,
|
||||
base_headers=headers,
|
||||
config=cast(Dict[str, Any], transformed_request)[
|
||||
"resumable_chunked_upload"
|
||||
],
|
||||
body_stream=cast(BaseFileUploadStream, media_cfg["body_stream"]),
|
||||
content_type=media_cfg.get("content_type") or "application/octet-stream",
|
||||
timeout=timeout,
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -3475,222 +3471,81 @@ class BaseLLMHTTPHandler:
|
|||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
# 8 MiB; a 256 KiB multiple, which GCS requires for every non-final chunk.
|
||||
_RESUMABLE_CHUNK_SIZE = 8 * 1024 * 1024
|
||||
# The fine-grained transform stream (one piece per JSONL row) is regrouped
|
||||
# into blocks of this size before upload, so the request yields a manageable
|
||||
# number of chunks; never more than one block is buffered.
|
||||
_MEDIA_UPLOAD_BLOCK_SIZE = 4 * 1024 * 1024
|
||||
|
||||
@staticmethod
|
||||
def _iter_resumable_chunks(
|
||||
byte_iter: Iterator[bytes], chunk_size: int
|
||||
) -> Iterator[bytes]:
|
||||
"""Regroup a byte stream into ``chunk_size`` pieces, yielding a final
|
||||
partial piece only when it is non-empty. Every full piece is exactly
|
||||
``chunk_size`` bytes (kept a 256 KiB multiple for GCS) and never more than
|
||||
one chunk is buffered. An exactly chunk-aligned stream yields only full
|
||||
chunks, so the upload finalizes on its last data chunk instead of making
|
||||
an extra empty request; a 0-byte stream yields nothing and the caller
|
||||
finalizes with a single empty request.
|
||||
"""
|
||||
def _iter_in_blocks(byte_iter: Iterator[bytes], block_size: int) -> Iterator[bytes]:
|
||||
buf = bytearray()
|
||||
for piece in byte_iter:
|
||||
buf.extend(piece)
|
||||
while len(buf) >= chunk_size:
|
||||
yield bytes(buf[:chunk_size])
|
||||
del buf[:chunk_size]
|
||||
while len(buf) >= block_size:
|
||||
yield bytes(buf[:block_size])
|
||||
del buf[:block_size]
|
||||
if buf:
|
||||
yield bytes(buf)
|
||||
|
||||
@staticmethod
|
||||
def _resumable_content_range(offset: int, data_len: int, is_final: bool) -> str:
|
||||
if not is_final:
|
||||
return f"bytes {offset}-{offset + data_len - 1}/*"
|
||||
total = offset + data_len
|
||||
if data_len == 0:
|
||||
return f"bytes */{total}"
|
||||
return f"bytes {offset}-{total - 1}/{total}"
|
||||
def _check_media_upload_response(self, resp: httpx.Response) -> None:
|
||||
if resp.status_code not in (200, 201):
|
||||
resp.raise_for_status()
|
||||
raise ValueError(f"media upload: unexpected status {resp.status_code}")
|
||||
|
||||
@staticmethod
|
||||
def _resumable_request_kwargs(
|
||||
headers: dict,
|
||||
content: bytes,
|
||||
timeout: Optional[Union[float, httpx.Timeout]],
|
||||
) -> dict:
|
||||
kwargs: Dict[str, Any] = {"headers": headers, "content": content}
|
||||
if timeout is not None:
|
||||
kwargs["timeout"] = timeout
|
||||
return kwargs
|
||||
|
||||
def _resumable_chunked_upload(
|
||||
def _upload_media(
|
||||
self,
|
||||
*,
|
||||
client: HTTPHandler,
|
||||
initiate_url: str,
|
||||
base_headers: dict,
|
||||
config: dict,
|
||||
timeout: Optional[Union[float, httpx.Timeout]],
|
||||
) -> httpx.Response:
|
||||
"""Open a GCS resumable session, then PUT the body in bounded chunks so a
|
||||
large upload is never held in memory in full."""
|
||||
stream = config["body_stream"]
|
||||
chunk_size = config.get("chunk_size", self._RESUMABLE_CHUNK_SIZE)
|
||||
session_url_header = config.get("session_url_header", "location")
|
||||
httpx_client = client.client
|
||||
|
||||
init_headers = {**base_headers, **config.get("initiate_headers", {})}
|
||||
init_req = httpx_client.build_request(
|
||||
"POST",
|
||||
initiate_url,
|
||||
**self._resumable_request_kwargs(init_headers, b"", timeout),
|
||||
)
|
||||
init_resp = httpx_client.send(init_req, follow_redirects=False)
|
||||
init_resp.read()
|
||||
if init_resp.status_code not in (200, 201):
|
||||
init_resp.raise_for_status()
|
||||
session_url = init_resp.headers.get(session_url_header)
|
||||
if not session_url:
|
||||
raise ValueError(
|
||||
f"resumable upload: no session URL in '{session_url_header}' header"
|
||||
)
|
||||
|
||||
offset = 0
|
||||
pending: Optional[bytes] = None
|
||||
for chunk in self._iter_resumable_chunks(stream.iter_bytes(), chunk_size):
|
||||
if pending is not None:
|
||||
self._send_resumable_chunk(
|
||||
httpx_client,
|
||||
session_url,
|
||||
base_headers,
|
||||
pending,
|
||||
offset,
|
||||
is_final=False,
|
||||
timeout=timeout,
|
||||
)
|
||||
offset += len(pending)
|
||||
pending = chunk
|
||||
return self._send_resumable_chunk(
|
||||
httpx_client,
|
||||
session_url,
|
||||
base_headers,
|
||||
pending or b"",
|
||||
offset,
|
||||
is_final=True,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
def _send_resumable_chunk(
|
||||
self,
|
||||
httpx_client: httpx.Client,
|
||||
url: str,
|
||||
base_headers: dict,
|
||||
data: bytes,
|
||||
offset: int,
|
||||
*,
|
||||
is_final: bool,
|
||||
base_headers: Dict[str, str],
|
||||
body_stream: BaseFileUploadStream,
|
||||
content_type: str,
|
||||
timeout: Optional[Union[float, httpx.Timeout]],
|
||||
) -> httpx.Response:
|
||||
headers = {
|
||||
**base_headers,
|
||||
"Content-Range": self._resumable_content_range(offset, len(data), is_final),
|
||||
headers = {**base_headers, "Content-Type": content_type}
|
||||
kwargs: Dict[str, Any] = {
|
||||
"headers": headers,
|
||||
"content": self._iter_in_blocks(body_stream.iter_bytes(), self._MEDIA_UPLOAD_BLOCK_SIZE),
|
||||
}
|
||||
req = httpx_client.build_request(
|
||||
"PUT", url, **self._resumable_request_kwargs(headers, data, timeout)
|
||||
)
|
||||
resp = httpx_client.send(req, follow_redirects=False)
|
||||
resp.read()
|
||||
if resp.status_code not in ((200, 201) if is_final else (308,)):
|
||||
# 4xx/5xx raise here; the ValueError catches an unexpected success
|
||||
# status (e.g. a 200 where the protocol expects a 308 between chunks).
|
||||
resp.raise_for_status()
|
||||
raise ValueError(f"resumable upload: unexpected status {resp.status_code}")
|
||||
if timeout is not None:
|
||||
kwargs["timeout"] = timeout
|
||||
resp = client.client.post(url, **kwargs)
|
||||
self._check_media_upload_response(resp)
|
||||
return resp
|
||||
|
||||
async def _aresumable_chunked_upload(
|
||||
async def _aupload_media(
|
||||
self,
|
||||
*,
|
||||
client: AsyncHTTPHandler,
|
||||
initiate_url: str,
|
||||
base_headers: dict,
|
||||
config: dict,
|
||||
timeout: Optional[Union[float, httpx.Timeout]],
|
||||
) -> httpx.Response:
|
||||
stream = config["body_stream"]
|
||||
chunk_size = config.get("chunk_size", self._RESUMABLE_CHUNK_SIZE)
|
||||
session_url_header = config.get("session_url_header", "location")
|
||||
httpx_client = client.client
|
||||
|
||||
init_headers = {**base_headers, **config.get("initiate_headers", {})}
|
||||
init_req = httpx_client.build_request(
|
||||
"POST",
|
||||
initiate_url,
|
||||
**self._resumable_request_kwargs(init_headers, b"", timeout),
|
||||
)
|
||||
init_resp = await httpx_client.send(init_req, follow_redirects=False)
|
||||
await init_resp.aread()
|
||||
if init_resp.status_code not in (200, 201):
|
||||
init_resp.raise_for_status()
|
||||
session_url = init_resp.headers.get(session_url_header)
|
||||
if not session_url:
|
||||
raise ValueError(
|
||||
f"resumable upload: no session URL in '{session_url_header}' header"
|
||||
)
|
||||
|
||||
offset = 0
|
||||
pending: Optional[bytes] = None
|
||||
# Producing each chunk runs the synchronous per-row transform for that
|
||||
# chunk's worth of rows. Pull it off the event loop thread so a large
|
||||
# upload does not block other concurrent requests between PUTs.
|
||||
chunk_iter = self._iter_resumable_chunks(stream.iter_bytes(), chunk_size)
|
||||
done = object()
|
||||
while True:
|
||||
chunk = await asyncio.to_thread(next, chunk_iter, done)
|
||||
if chunk is done:
|
||||
break
|
||||
if pending is not None:
|
||||
await self._asend_resumable_chunk(
|
||||
httpx_client,
|
||||
session_url,
|
||||
base_headers,
|
||||
pending,
|
||||
offset,
|
||||
is_final=False,
|
||||
timeout=timeout,
|
||||
)
|
||||
offset += len(pending)
|
||||
pending = chunk
|
||||
return await self._asend_resumable_chunk(
|
||||
httpx_client,
|
||||
session_url,
|
||||
base_headers,
|
||||
pending or b"",
|
||||
offset,
|
||||
is_final=True,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
async def _asend_resumable_chunk(
|
||||
self,
|
||||
httpx_client: httpx.AsyncClient,
|
||||
url: str,
|
||||
base_headers: dict,
|
||||
data: bytes,
|
||||
offset: int,
|
||||
*,
|
||||
is_final: bool,
|
||||
base_headers: Dict[str, str],
|
||||
body_stream: BaseFileUploadStream,
|
||||
content_type: str,
|
||||
timeout: Optional[Union[float, httpx.Timeout]],
|
||||
) -> httpx.Response:
|
||||
headers = {
|
||||
**base_headers,
|
||||
"Content-Range": self._resumable_content_range(offset, len(data), is_final),
|
||||
}
|
||||
req = httpx_client.build_request(
|
||||
"PUT", url, **self._resumable_request_kwargs(headers, data, timeout)
|
||||
)
|
||||
resp = await httpx_client.send(req, follow_redirects=False)
|
||||
"""Stream the transformed body straight to a single media upload. Each
|
||||
block is produced on a worker thread (the transform never runs on the
|
||||
event loop) and sent with chunked transfer-encoding, so the body is
|
||||
neither buffered in memory nor staged to disk, and the upload is one
|
||||
continuous request rather than the many sequential round-trips of the
|
||||
resumable path that overran client/LB timeouts."""
|
||||
headers = {**base_headers, "Content-Type": content_type}
|
||||
block_iter = iter(self._iter_in_blocks(body_stream.iter_bytes(), self._MEDIA_UPLOAD_BLOCK_SIZE))
|
||||
done = object()
|
||||
|
||||
async def _abody() -> AsyncIterator[bytes]:
|
||||
while True:
|
||||
block = await asyncio.to_thread(next, block_iter, done)
|
||||
if block is done:
|
||||
break
|
||||
yield cast(bytes, block)
|
||||
|
||||
kwargs: Dict[str, Any] = {"headers": headers, "content": _abody()}
|
||||
if timeout is not None:
|
||||
kwargs["timeout"] = timeout
|
||||
resp = await client.client.post(url, **kwargs)
|
||||
await resp.aread()
|
||||
if resp.status_code not in ((200, 201) if is_final else (308,)):
|
||||
# 4xx/5xx raise here; the ValueError catches an unexpected success
|
||||
# status (e.g. a 200 where the protocol expects a 308 between chunks).
|
||||
resp.raise_for_status()
|
||||
raise ValueError(f"resumable upload: unexpected status {resp.status_code}")
|
||||
self._check_media_upload_response(resp)
|
||||
return resp
|
||||
|
||||
def create_batch(
|
||||
|
|
|
|||
|
|
@ -60,7 +60,7 @@ from litellm.types.llms.openai import (
|
|||
OpenAIFileObject,
|
||||
PathLike,
|
||||
)
|
||||
from litellm.types.files import ResumableChunkedUploadConfig
|
||||
from litellm.types.files import StreamingMediaUploadConfig
|
||||
from litellm.types.llms.vertex_ai import GcsBucketResponse
|
||||
from litellm.types.utils import LlmProviders, ModelResponse
|
||||
|
||||
|
|
@ -389,21 +389,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
raise ValueError("file is required")
|
||||
if purpose is None:
|
||||
raise ValueError("purpose is required")
|
||||
_, content_type = extract_file_metadata(file_data)
|
||||
object_name = self.get_object_name(file_data, purpose)
|
||||
if object_prefix:
|
||||
object_name = f"{object_prefix}/{object_name}"
|
||||
encoded_object_name = encode_gcs_object_name_for_url(object_name)
|
||||
# Batch jsonl is streamed via a resumable session (bounded memory on
|
||||
# large uploads); everything else is a single simple-media upload.
|
||||
upload_type = (
|
||||
"resumable"
|
||||
if FilesAPIUtils.is_batch_jsonl_request(
|
||||
create_file_data=data, content_type=content_type
|
||||
)
|
||||
else "media"
|
||||
)
|
||||
endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType={upload_type}&name={encoded_object_name}"
|
||||
endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_object_name}"
|
||||
api_base = api_base or "https://storage.googleapis.com"
|
||||
if not api_base:
|
||||
raise ValueError("api_base is required")
|
||||
|
|
@ -455,8 +445,9 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
"""
|
||||
2 Cases:
|
||||
1. Handle basic file upload
|
||||
2. Handle batch file upload (.jsonl), streamed to a GCS resumable
|
||||
session so large uploads stay memory-bounded.
|
||||
2. Handle batch file upload (.jsonl), staged to a temp file and uploaded
|
||||
in a single media request so large uploads stay memory-bounded without
|
||||
the per-chunk round-trips of a resumable session.
|
||||
"""
|
||||
file_data = create_file_data.get("file")
|
||||
if file_data is None:
|
||||
|
|
@ -468,14 +459,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
content_type=content_type,
|
||||
):
|
||||
return {
|
||||
"resumable_chunked_upload": ResumableChunkedUploadConfig(
|
||||
"streaming_media_upload": StreamingMediaUploadConfig(
|
||||
body_stream=_OpenAIToVertexBatchUploadStream(
|
||||
file_data,
|
||||
self._map_openai_to_vertex_params,
|
||||
),
|
||||
initiate_headers={
|
||||
"X-Upload-Content-Type": "application/json",
|
||||
},
|
||||
content_type="application/json",
|
||||
)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -323,19 +323,18 @@ class TwoStepFileUploadConfig(TypedDict, total=False):
|
|||
upload_url_key: str
|
||||
|
||||
|
||||
class ResumableChunkedUploadConfig(TypedDict, total=False):
|
||||
"""Drives a memory-bounded resumable upload (GCS JSON API).
|
||||
class StreamingMediaUploadConfig(TypedDict, total=False):
|
||||
"""Drives a memory-bounded single-request upload (GCS simple/media upload).
|
||||
|
||||
The handler POSTs to the upload URL to open a session, reads the session URI
|
||||
from ``session_url_header``, then PUTs ``body_stream`` to that URI in
|
||||
``chunk_size``-byte chunks (a 256 KiB multiple) using Content-Range, so the
|
||||
payload is never buffered in full and the transfer is resumable.
|
||||
The handler stages ``body_stream`` to a temp file off the event loop (so peak
|
||||
memory stays bounded), then PUTs/POSTs it in one request with a known
|
||||
Content-Length. Unlike a resumable chunked upload this incurs no per-chunk
|
||||
round-trips, so a multi-GB upload finishes in one continuous transfer instead
|
||||
of hundreds of sequential PUTs that overrun client/LB timeouts.
|
||||
|
||||
``body_stream`` is a ``BaseFileUploadStream``; it is typed ``Any`` here to
|
||||
avoid importing the llms layer into types.
|
||||
"""
|
||||
|
||||
body_stream: Required[Any]
|
||||
chunk_size: int
|
||||
session_url_header: str
|
||||
initiate_headers: Dict[str, str]
|
||||
content_type: str
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm"
|
||||
version = "1.89.4"
|
||||
version = "1.89.5"
|
||||
description = "Library to easily interface with LLM API providers"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10, <3.14"
|
||||
|
|
@ -261,7 +261,7 @@ source-exclude = [
|
|||
profile = "black"
|
||||
|
||||
[tool.commitizen]
|
||||
version = "1.89.4"
|
||||
version = "1.89.5"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -513,25 +513,26 @@ async def test_avertex_batch_prediction(monkeypatch):
|
|||
mock_response.status_code = 200
|
||||
return mock_response
|
||||
|
||||
# Batch jsonl file creation now streams to a GCS resumable session via
|
||||
# _aresumable_chunked_upload (httpx send), not AsyncHTTPHandler.post, so mock
|
||||
# that entry point to return the GCS object response. The resumable protocol
|
||||
# itself is covered in test_vertex_ai_files_streaming.py.
|
||||
mock_upload_response = httpx.Response(
|
||||
200,
|
||||
json=mock_file_response,
|
||||
request=httpx.Request("PUT", "https://storage.googleapis.com/upload"),
|
||||
)
|
||||
# Batch jsonl creation now stages the body to a temp file and issues a single
|
||||
# uploadType=media POST against the raw httpx.AsyncClient (client.client) inside
|
||||
# _astage_and_upload_media, not AsyncHTTPHandler.post. Patch that raw POST so the
|
||||
# real staging/upload + response transform run while the GCS object response is
|
||||
# mocked; AsyncHTTPHandler.post still handles the batch-prediction call.
|
||||
with (
|
||||
patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
side_effect=mock_side_effect,
|
||||
) as mock_global_post,
|
||||
patch(
|
||||
"litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler._aresumable_chunked_upload",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_upload_response,
|
||||
),
|
||||
patch.object(
|
||||
httpx.AsyncClient,
|
||||
"post",
|
||||
new_callable=AsyncMock,
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json=mock_file_response,
|
||||
request=httpx.Request("POST", "https://storage.googleapis.com/upload"),
|
||||
),
|
||||
) as mock_gcs_upload,
|
||||
):
|
||||
litellm.set_verbose = True
|
||||
litellm._turn_on_debug()
|
||||
|
|
@ -552,6 +553,15 @@ async def test_avertex_batch_prediction(monkeypatch):
|
|||
== "gs://litellm-local/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/5f7b99ad-9203-4430-98bf-3b45451af4cb"
|
||||
)
|
||||
|
||||
mock_gcs_upload.assert_awaited_once()
|
||||
upload_url = str(mock_gcs_upload.call_args.args[0])
|
||||
assert "uploadType=media" in upload_url
|
||||
assert "/b/litellm-local/o" in upload_url
|
||||
assert (
|
||||
mock_gcs_upload.call_args.kwargs["headers"]["Content-Type"]
|
||||
== "application/json"
|
||||
)
|
||||
|
||||
# Create batch
|
||||
create_batch_response = await litellm.acreate_batch(
|
||||
completion_window="24h",
|
||||
|
|
|
|||
|
|
@ -37,9 +37,7 @@ class TestVertexAIBinaryFileUpload:
|
|||
# Create mock PDF binary data (with non-UTF-8 bytes)
|
||||
# PDF files start with %PDF- and contain binary data
|
||||
mock_pdf_content = b"%PDF-1.4\n%\xc4\xe5\xf2\xe5\xeb\xa7\xf3\xa0\xd0\xc4\xc6\n"
|
||||
mock_pdf_content += (
|
||||
b"\x00\x01\x02\x03\xff\xfe\xfd" * 100
|
||||
) # Add more binary data
|
||||
mock_pdf_content += b"\x00\x01\x02\x03\xff\xfe\xfd" * 100 # Add more binary data
|
||||
|
||||
# Create file object
|
||||
file_obj = io.BytesIO(mock_pdf_content)
|
||||
|
|
@ -60,14 +58,12 @@ class TestVertexAIBinaryFileUpload:
|
|||
)
|
||||
|
||||
# Verify the transformation returns bytes (not string)
|
||||
assert isinstance(
|
||||
transformed_request, bytes
|
||||
), f"Expected bytes for binary file, got {type(transformed_request)}"
|
||||
assert isinstance(transformed_request, bytes), (
|
||||
f"Expected bytes for binary file, got {type(transformed_request)}"
|
||||
)
|
||||
|
||||
# Verify the bytes match the original content
|
||||
assert (
|
||||
transformed_request == mock_pdf_content
|
||||
), "Transformed request should preserve binary content exactly"
|
||||
assert transformed_request == mock_pdf_content, "Transformed request should preserve binary content exactly"
|
||||
|
||||
# Verify that the bytes contain non-UTF-8 characters
|
||||
# This should raise UnicodeDecodeError if we try to decode
|
||||
|
|
@ -132,16 +128,14 @@ class TestVertexAIBinaryFileUpload:
|
|||
pytest.fail(f"httpx should accept bytes in data parameter: {e}")
|
||||
|
||||
# Document the expected behavior
|
||||
assert isinstance(
|
||||
mock_binary_data, bytes
|
||||
), "Binary file data should remain as bytes"
|
||||
assert isinstance(mock_binary_data, bytes), "Binary file data should remain as bytes"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_jsonl_file_upload_returns_resumable_stream(self):
|
||||
async def test_jsonl_file_upload_returns_streaming_body(self):
|
||||
"""
|
||||
Test that JSONL batch files are transformed into a resumable-upload config
|
||||
Test that JSONL batch files are transformed into a streaming-media config
|
||||
carrying a streaming body (not a buffered bytes payload), so the handler
|
||||
can stream the upload to GCS in bounded chunks.
|
||||
can stage the upload to a temp file and send it in one media request.
|
||||
"""
|
||||
# Create mock JSONL content
|
||||
mock_jsonl_content = (
|
||||
|
|
@ -164,16 +158,13 @@ class TestVertexAIBinaryFileUpload:
|
|||
litellm_params={},
|
||||
)
|
||||
|
||||
assert (
|
||||
isinstance(transformed_request, dict)
|
||||
and "resumable_chunked_upload" in transformed_request
|
||||
), f"Expected a resumable upload config for JSONL, got {type(transformed_request)}"
|
||||
assert isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request, (
|
||||
f"Expected a streaming media upload config for JSONL, got {type(transformed_request)}"
|
||||
)
|
||||
|
||||
stream = transformed_request["resumable_chunked_upload"]["body_stream"]
|
||||
stream = transformed_request["streaming_media_upload"]["body_stream"]
|
||||
decoded = json.loads(b"".join(stream.iter_bytes()).decode("utf-8"))
|
||||
assert (
|
||||
"request" in decoded
|
||||
), "JSONL transform must wrap each row in {'request': ...}"
|
||||
assert "request" in decoded, "JSONL transform must wrap each row in {'request': ...}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mixed_file_types_in_sequence(self):
|
||||
|
|
@ -214,7 +205,7 @@ class TestVertexAIBinaryFileUpload:
|
|||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert isinstance(result2, dict) and "resumable_chunked_upload" in result2
|
||||
assert isinstance(result2, dict) and "streaming_media_upload" in result2
|
||||
|
||||
# Test 3: Upload another binary file
|
||||
binary_content2 = b"\xc4\xe5\xf2\xe5\xeb"
|
||||
|
|
@ -264,7 +255,5 @@ class TestVertexAIBinaryFileUpload:
|
|||
},
|
||||
}
|
||||
|
||||
assert (
|
||||
expected_behavior["binary_files"]["encoding"] == "none - preserve raw bytes"
|
||||
)
|
||||
assert expected_behavior["binary_files"]["encoding"] == "none - preserve raw bytes"
|
||||
assert expected_behavior["text_files"]["encoding"] == "UTF-8"
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ replaced by a list-based pipeline:
|
|||
import gc
|
||||
import io
|
||||
import json
|
||||
import tempfile
|
||||
import time
|
||||
import tracemalloc
|
||||
|
||||
|
|
@ -41,15 +42,15 @@ from litellm.llms.vertex_ai.files.transformation import (
|
|||
from litellm.types.llms.openai import CreateFileRequest
|
||||
|
||||
|
||||
def _resumable_stream(transformed) -> BaseFileUploadStream:
|
||||
"""Pull the streaming body out of a resumable-upload transform result."""
|
||||
return transformed["resumable_chunked_upload"]["body_stream"]
|
||||
def _upload_stream(transformed) -> BaseFileUploadStream:
|
||||
"""Pull the streaming body out of the upload transform result."""
|
||||
return transformed["streaming_media_upload"]["body_stream"]
|
||||
|
||||
|
||||
def _join_upload_body(transformed) -> bytes:
|
||||
"""Materialize a transform result's upload body for byte-level assertions."""
|
||||
if isinstance(transformed, dict) and "resumable_chunked_upload" in transformed:
|
||||
return b"".join(_resumable_stream(transformed).iter_bytes())
|
||||
if isinstance(transformed, dict) and "streaming_media_upload" in transformed:
|
||||
return b"".join(_upload_stream(transformed).iter_bytes())
|
||||
if isinstance(transformed, BaseFileUploadStream):
|
||||
return b"".join(transformed.iter_bytes())
|
||||
if isinstance(transformed, str):
|
||||
|
|
@ -83,17 +84,13 @@ def _reference_vertex_jsonl_string(cfg: VertexAIFilesConfig, content: str) -> st
|
|||
transform, so the streaming path can be checked against it for parity."""
|
||||
entries = [json.loads(line) for line in content.splitlines() if line.strip()]
|
||||
return "\n".join(
|
||||
json.dumps(
|
||||
_openai_batch_jsonl_entry_to_vertex_wrapped_request(
|
||||
entry, cfg._map_openai_to_vertex_params
|
||||
)
|
||||
)
|
||||
json.dumps(_openai_batch_jsonl_entry_to_vertex_wrapped_request(entry, cfg._map_openai_to_vertex_params))
|
||||
for entry in entries
|
||||
)
|
||||
|
||||
|
||||
class TestStreamingOutputParity:
|
||||
def test_transform_create_file_request_returns_resumable_stream_parity(self):
|
||||
def test_transform_create_file_request_returns_streaming_body_parity(self):
|
||||
cfg = VertexAIFilesConfig()
|
||||
raw = _make_openai_jsonl_bytes(300)
|
||||
request: CreateFileRequest = {
|
||||
|
|
@ -105,14 +102,12 @@ class TestStreamingOutputParity:
|
|||
model="", create_file_data=request, optional_params={}, litellm_params={}
|
||||
)
|
||||
|
||||
# A batch upload must be a resumable-upload config carrying a streaming
|
||||
# body, so the handler can chunk it; a buffered bytes/str return would
|
||||
# defeat the OOM fix.
|
||||
assert isinstance(out, dict) and "resumable_chunked_upload" in out
|
||||
assert isinstance(_resumable_stream(out), BaseFileUploadStream)
|
||||
assert _join_upload_body(out).decode("utf-8") == _reference_vertex_jsonl_string(
|
||||
cfg, raw.decode("utf-8")
|
||||
)
|
||||
# A batch upload must be a streaming-media config carrying a streaming
|
||||
# body, so the handler can stream it to GCS; a buffered bytes/str return
|
||||
# would defeat the OOM fix.
|
||||
assert isinstance(out, dict) and "streaming_media_upload" in out
|
||||
assert isinstance(_upload_stream(out), BaseFileUploadStream)
|
||||
assert _join_upload_body(out).decode("utf-8") == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8"))
|
||||
|
||||
|
||||
class TestFileLikeInputNotPartiallyConsumed:
|
||||
|
|
@ -215,9 +210,7 @@ class TestStreamingLineIterator:
|
|||
def seek(self, *args):
|
||||
raise io.UnsupportedOperation("not seekable")
|
||||
|
||||
handle = _NonSeekable(
|
||||
b'{"custom_id": "request-0"}\n{"custom_id": "request-1"}\n'
|
||||
)
|
||||
handle = _NonSeekable(b'{"custom_id": "request-0"}\n{"custom_id": "request-1"}\n')
|
||||
with pytest.raises(ValueError, match="seekable"):
|
||||
list(_iter_openai_jsonl_lines(handle))
|
||||
|
||||
|
|
@ -235,13 +228,8 @@ class TestGetObjectNameLazyParse:
|
|||
cfg = VertexAIFilesConfig()
|
||||
# Tail rows are deliberately not valid JSON. Parsing the whole payload
|
||||
# would raise here; a first-row-only parse must not.
|
||||
raw = (
|
||||
b'{"custom_id": "r-0", "body": {"model": "gemini-2.5-flash"}}\n'
|
||||
b"garbage line that is not json\n"
|
||||
)
|
||||
object_name = cfg.get_object_name(
|
||||
("batch.jsonl", raw, "application/jsonl"), purpose="batch"
|
||||
)
|
||||
raw = b'{"custom_id": "r-0", "body": {"model": "gemini-2.5-flash"}}\ngarbage line that is not json\n'
|
||||
object_name = cfg.get_object_name(("batch.jsonl", raw, "application/jsonl"), purpose="batch")
|
||||
assert "gemini-2.5-flash" in object_name
|
||||
|
||||
|
||||
|
|
@ -278,15 +266,11 @@ class TestStreamingPeakMemory:
|
|||
def drain_stream():
|
||||
# Consume the upload body one row at a time, as the chunked uploader
|
||||
# does, without accumulating it.
|
||||
for _ in _OpenAIToVertexBatchUploadStream(
|
||||
raw, cfg._map_openai_to_vertex_params
|
||||
).iter_bytes():
|
||||
for _ in _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params).iter_bytes():
|
||||
pass
|
||||
|
||||
streaming_peak = self._measure(drain_stream)
|
||||
list_peak = self._measure(
|
||||
lambda: _reference_vertex_jsonl_string(cfg, content_str)
|
||||
)
|
||||
list_peak = self._measure(lambda: _reference_vertex_jsonl_string(cfg, content_str))
|
||||
|
||||
# Core guard: the lazily consumed streaming body peaks well under a list
|
||||
# pipeline that materializes every transformed row. Building full
|
||||
|
|
@ -305,9 +289,7 @@ class TestStreamingPeakMemory:
|
|||
# first-row parse should allocate only a small fraction of the payload;
|
||||
# parsing every row would blow past this bound.
|
||||
peak = self._measure(lambda: cfg.get_object_name(file_data, purpose="batch"))
|
||||
assert (
|
||||
peak / len(raw) < 2.0
|
||||
), "get_object_name should not copy the whole payload"
|
||||
assert peak / len(raw) < 2.0, "get_object_name should not copy the whole payload"
|
||||
|
||||
|
||||
class TestPathSourcedStreaming:
|
||||
|
|
@ -341,12 +323,10 @@ class TestPathSourcedStreaming:
|
|||
litellm_params={"gcs_bucket_name": "test-bucket"},
|
||||
data=data,
|
||||
)
|
||||
assert "uploadType=resumable" in url
|
||||
assert "uploadType=media" in url
|
||||
|
||||
out = cfg.transform_create_file_request(
|
||||
model="", create_file_data=data, optional_params={}, litellm_params={}
|
||||
)
|
||||
assert isinstance(out, dict) and "resumable_chunked_upload" in out
|
||||
out = cfg.transform_create_file_request(model="", create_file_data=data, optional_params={}, litellm_params={})
|
||||
assert isinstance(out, dict) and "streaming_media_upload" in out
|
||||
body = _join_upload_body(out).decode("utf-8")
|
||||
assert body == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8"))
|
||||
lines = body.splitlines()
|
||||
|
|
@ -371,7 +351,7 @@ class TestPathSourcedStreaming:
|
|||
out = cfg.transform_create_file_request(
|
||||
model="", create_file_data=data, optional_params={}, litellm_params={}
|
||||
)
|
||||
for _ in _resumable_stream(out).iter_bytes():
|
||||
for _ in _upload_stream(out).iter_bytes():
|
||||
pass # drain without accumulating
|
||||
|
||||
gc.collect()
|
||||
|
|
@ -384,20 +364,15 @@ class TestPathSourcedStreaming:
|
|||
|
||||
# Streaming from disk must not materialize the payload. Reading the whole
|
||||
# file into bytes (the pre-fix path) would push peak past the file size.
|
||||
assert peak < len(raw) * 0.3, (
|
||||
f"peak {peak} not bounded vs payload {len(raw)} "
|
||||
f"(ratio {peak / len(raw):.2f})"
|
||||
)
|
||||
assert peak < len(raw) * 0.3, f"peak {peak} not bounded vs payload {len(raw)} (ratio {peak / len(raw):.2f})"
|
||||
|
||||
def test_path_source_stream_is_reiterable(self, tmp_path):
|
||||
cfg = VertexAIFilesConfig()
|
||||
path, _ = self._write_jsonl(tmp_path, 50)
|
||||
data = self._batch_request(path)
|
||||
|
||||
out = cfg.transform_create_file_request(
|
||||
model="", create_file_data=data, optional_params={}, litellm_params={}
|
||||
)
|
||||
stream = _resumable_stream(out)
|
||||
out = cfg.transform_create_file_request(model="", create_file_data=data, optional_params={}, litellm_params={})
|
||||
stream = _upload_stream(out)
|
||||
first = b"".join(stream.iter_bytes())
|
||||
second = b"".join(stream.iter_bytes())
|
||||
assert first == second and len(first) > 0
|
||||
|
|
@ -436,25 +411,20 @@ def _logging_obj() -> Logging:
|
|||
)
|
||||
|
||||
|
||||
def _gcs_resumable_mock(session_url: str, final_status: int = 200):
|
||||
"""A fake GCS resumable endpoint: POST opens a session (URI in Location),
|
||||
each PUT appends and returns 308 until the final chunk returns 200/201."""
|
||||
state = {"received": bytearray(), "ranges": [], "methods": [], "urls": []}
|
||||
def _gcs_media_mock(status: int = 200):
|
||||
"""A fake GCS simple-media endpoint: one request carries the whole object;
|
||||
capture the body and headers and return the object resource."""
|
||||
state = {"received": bytearray(), "methods": [], "urls": [], "headers": [], "timeouts": []}
|
||||
|
||||
async def handler(request: httpx.Request) -> httpx.Response:
|
||||
state["methods"].append(request.method)
|
||||
state["urls"].append(str(request.url))
|
||||
if request.method == "POST":
|
||||
return httpx.Response(200, headers={"location": session_url})
|
||||
body = await request.aread()
|
||||
content_range = request.headers["content-range"]
|
||||
state["ranges"].append(content_range)
|
||||
state["received"].extend(body)
|
||||
if content_range.rsplit("/", 1)[-1] == "*":
|
||||
return httpx.Response(
|
||||
308, headers={"range": f"bytes=0-{len(state['received']) - 1}"}
|
||||
)
|
||||
return httpx.Response(final_status, json=_GCS_OBJECT_JSON)
|
||||
state["headers"].append(dict(request.headers))
|
||||
# httpx records the resolved per-request timeout here, so the test can
|
||||
# assert the caller's timeout was forwarded rather than the client default.
|
||||
state["timeouts"].append(request.extensions.get("timeout"))
|
||||
state["received"].extend(await request.aread())
|
||||
return httpx.Response(status, json=_GCS_OBJECT_JSON)
|
||||
|
||||
return handler, state
|
||||
|
||||
|
|
@ -465,8 +435,8 @@ def _async_handler_with(mock) -> AsyncHTTPHandler:
|
|||
return handler
|
||||
|
||||
|
||||
class TestResumableUploadUrl:
|
||||
def test_batch_jsonl_uses_resumable_upload_type(self):
|
||||
class TestUploadUrl:
|
||||
def test_batch_jsonl_uses_media_upload_type(self):
|
||||
cfg = VertexAIFilesConfig()
|
||||
request: CreateFileRequest = {
|
||||
"file": ("batch.jsonl", _make_openai_jsonl_bytes(3), "application/jsonl"),
|
||||
|
|
@ -480,29 +450,12 @@ class TestResumableUploadUrl:
|
|||
litellm_params={"gcs_bucket_name": "test-bucket"},
|
||||
data=request,
|
||||
)
|
||||
assert "uploadType=resumable" in url
|
||||
assert "uploadType=media" not in url
|
||||
# A single media upload is one continuous transfer (no per-chunk
|
||||
# round-trips), which is what keeps large uploads under client/LB timeouts.
|
||||
assert "uploadType=media" in url
|
||||
assert "uploadType=resumable" not in url
|
||||
|
||||
def test_batch_text_plain_uses_resumable_upload_type(self):
|
||||
# Clients often label a .jsonl batch upload as text/plain; it must still
|
||||
# take the streaming/resumable path, not the buffered media path.
|
||||
cfg = VertexAIFilesConfig()
|
||||
request: CreateFileRequest = {
|
||||
"file": ("batch.jsonl", _make_openai_jsonl_bytes(3), "text/plain"),
|
||||
"purpose": "batch",
|
||||
}
|
||||
url = cfg.get_complete_file_url(
|
||||
api_base=None,
|
||||
api_key=None,
|
||||
model="",
|
||||
optional_params={},
|
||||
litellm_params={"gcs_bucket_name": "test-bucket"},
|
||||
data=request,
|
||||
)
|
||||
assert "uploadType=resumable" in url
|
||||
assert "uploadType=media" not in url
|
||||
|
||||
def test_binary_upload_stays_simple_media(self):
|
||||
def test_binary_upload_uses_media_upload_type(self):
|
||||
cfg = VertexAIFilesConfig()
|
||||
request: CreateFileRequest = {
|
||||
"file": ("doc.pdf", b"%PDF-1.4 binary", "application/pdf"),
|
||||
|
|
@ -520,14 +473,12 @@ class TestResumableUploadUrl:
|
|||
assert "uploadType=resumable" not in url
|
||||
|
||||
|
||||
class TestResumableStreamBody:
|
||||
class TestUploadStreamBody:
|
||||
def test_stream_matches_legacy_pipeline(self):
|
||||
cfg = VertexAIFilesConfig()
|
||||
raw = _make_openai_jsonl_bytes(120)
|
||||
stream = _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params)
|
||||
assert b"".join(stream.iter_bytes()).decode(
|
||||
"utf-8"
|
||||
) == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8"))
|
||||
assert b"".join(stream.iter_bytes()).decode("utf-8") == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8"))
|
||||
|
||||
def test_stream_is_reiterable_for_retries(self):
|
||||
# A one-shot generator would make a transport retry upload an empty body;
|
||||
|
|
@ -545,59 +496,19 @@ class TestResumableStreamBody:
|
|||
# an empty body silently.
|
||||
cfg = VertexAIFilesConfig()
|
||||
raw = _make_openai_jsonl_bytes(40)
|
||||
stream = _OpenAIToVertexBatchUploadStream(
|
||||
io.BytesIO(raw), cfg._map_openai_to_vertex_params
|
||||
)
|
||||
stream = _OpenAIToVertexBatchUploadStream(io.BytesIO(raw), cfg._map_openai_to_vertex_params)
|
||||
first = b"".join(stream.iter_bytes())
|
||||
second = b"".join(stream.iter_bytes())
|
||||
assert first == second and len(first) > 0
|
||||
|
||||
|
||||
class TestResumableChunking:
|
||||
def test_intermediate_chunks_are_exactly_chunk_size(self):
|
||||
pieces = list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([b"x" * 10]), 4))
|
||||
assert pieces == [b"xxxx", b"xxxx", b"xx"]
|
||||
|
||||
def test_exact_multiple_yields_no_trailing_empty(self):
|
||||
# An exactly chunk-aligned stream yields only full chunks; the upload
|
||||
# finalizes on the last data chunk instead of an extra empty request.
|
||||
pieces = list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([b"x" * 8]), 4))
|
||||
assert pieces == [b"xxxx", b"xxxx"]
|
||||
|
||||
def test_empty_stream_yields_nothing(self):
|
||||
# A 0-byte stream yields no chunks; the caller finalizes with one empty
|
||||
# request (bytes */0).
|
||||
assert list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([]), 4)) == []
|
||||
|
||||
def test_default_chunk_size_is_256kib_multiple(self):
|
||||
assert BaseLLMHTTPHandler._RESUMABLE_CHUNK_SIZE % (256 * 1024) == 0
|
||||
|
||||
def test_content_range_intermediate_uses_star_total(self):
|
||||
assert (
|
||||
BaseLLMHTTPHandler._resumable_content_range(0, 4096, is_final=False)
|
||||
== "bytes 0-4095/*"
|
||||
)
|
||||
|
||||
def test_content_range_final_uses_real_total(self):
|
||||
assert (
|
||||
BaseLLMHTTPHandler._resumable_content_range(8192, 100, is_final=True)
|
||||
== "bytes 8192-8291/8292"
|
||||
)
|
||||
|
||||
def test_content_range_empty_finalize(self):
|
||||
assert (
|
||||
BaseLLMHTTPHandler._resumable_content_range(8192, 0, is_final=True)
|
||||
== "bytes */8192"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestResumableUploadProtocol:
|
||||
"""End-to-end against a faked GCS resumable endpoint. These are the tests
|
||||
that fail if the handler buffers the whole body, drops bytes, mislabels a
|
||||
Content-Range, follows the 308 instead of continuing, or skips finalize."""
|
||||
class TestStreamingMediaUpload:
|
||||
"""End-to-end against a faked GCS media endpoint. These fail if the handler
|
||||
buffers the payload in memory, drops bytes, omits Content-Length (which would
|
||||
flip httpx to chunked transfer-encoding), or makes more than one request."""
|
||||
|
||||
async def _run(self, raw: bytes, chunk_size: int, final_status: int = 200):
|
||||
async def _run(self, raw: bytes, status: int = 200, timeout=None):
|
||||
cfg = VertexAIFilesConfig()
|
||||
request: CreateFileRequest = {
|
||||
"file": ("batch.jsonl", raw, "application/jsonl"),
|
||||
|
|
@ -614,11 +525,8 @@ class TestResumableUploadProtocol:
|
|||
transformed = cfg.transform_create_file_request(
|
||||
model="", create_file_data=request, optional_params={}, litellm_params={}
|
||||
)
|
||||
transformed["resumable_chunked_upload"]["chunk_size"] = chunk_size
|
||||
expected = _join_upload_body(transformed)
|
||||
|
||||
session_url = "https://storage.googleapis.com/upload/sess?upload_id=SID"
|
||||
mock, state = _gcs_resumable_mock(session_url, final_status=final_status)
|
||||
mock, state = _gcs_media_mock(status=status)
|
||||
response = await BaseLLMHTTPHandler().async_create_file(
|
||||
transformed_request=transformed,
|
||||
litellm_params={},
|
||||
|
|
@ -627,70 +535,52 @@ class TestResumableUploadProtocol:
|
|||
api_base=api_base,
|
||||
logging_obj=_logging_obj(),
|
||||
client=_async_handler_with(mock),
|
||||
timeout=None,
|
||||
timeout=timeout,
|
||||
)
|
||||
return expected, state, response, session_url, api_base
|
||||
return expected, state, response
|
||||
|
||||
async def test_streams_in_chunks_and_reassembles(self):
|
||||
async def test_single_request_carries_whole_payload(self):
|
||||
raw = _make_openai_jsonl_bytes(300)
|
||||
chunk_size = 4096
|
||||
expected, state, response, session_url, api_base = await self._run(
|
||||
raw, chunk_size
|
||||
)
|
||||
expected, state, response = await self._run(raw)
|
||||
|
||||
# One session-open POST, then a sequence of chunk PUTs.
|
||||
assert state["methods"][0] == "POST"
|
||||
assert set(state["methods"][1:]) == {"PUT"}
|
||||
assert state["methods"].count("PUT") >= 2, "payload must span multiple chunks"
|
||||
# Exactly one request (the single media upload), and it lands on the
|
||||
# media endpoint, not a resumable session.
|
||||
assert state["methods"] == ["POST"]
|
||||
assert "uploadType=media" in state["urls"][0]
|
||||
|
||||
# POST opens a resumable session; every chunk goes to the session URI.
|
||||
assert "uploadType=resumable" in state["urls"][0]
|
||||
assert all(u == session_url for u in state["urls"][1:])
|
||||
|
||||
# Every non-final chunk is exactly chunk_size with an unknown-total range;
|
||||
# the final chunk carries the real total.
|
||||
intermediate = state["ranges"][:-1]
|
||||
for index, content_range in enumerate(intermediate):
|
||||
assert (
|
||||
content_range
|
||||
== f"bytes {index * chunk_size}-{(index + 1) * chunk_size - 1}/*"
|
||||
)
|
||||
total = len(expected)
|
||||
last_offset = len(intermediate) * chunk_size
|
||||
if last_offset == total: # payload landed on a chunk boundary
|
||||
assert state["ranges"][-1] == f"bytes */{total}"
|
||||
else:
|
||||
assert state["ranges"][-1] == f"bytes {last_offset}-{total - 1}/{total}"
|
||||
|
||||
# The bytes GCS received are exactly the transformed batch payload.
|
||||
# The body is streamed with chunked transfer-encoding and no
|
||||
# Content-Length, which is what proves it is neither buffered in memory
|
||||
# nor staged to a temp file (the disk-exhaustion guard) before sending.
|
||||
headers = state["headers"][0]
|
||||
assert headers.get("transfer-encoding") == "chunked"
|
||||
assert "content-length" not in headers
|
||||
# httpx reassembles the chunked body; GCS receives exactly the transform.
|
||||
assert bytes(state["received"]) == expected
|
||||
assert response.object == "file"
|
||||
|
||||
async def test_exact_multiple_finalizes_on_last_data_chunk(self):
|
||||
# A body that is an exact multiple of the chunk size finalizes on its
|
||||
# last data chunk (bytes (TOTAL-chunk)-(TOTAL-1)/TOTAL), with no extra
|
||||
# empty finalize request.
|
||||
chunk_size = 256
|
||||
total = chunk_size * 3
|
||||
stream = _FixedBytesStream(b"a" * total)
|
||||
config = {"body_stream": stream, "chunk_size": chunk_size}
|
||||
session_url = "https://storage.googleapis.com/upload/sess?upload_id=SID"
|
||||
mock, state = _gcs_resumable_mock(session_url)
|
||||
|
||||
response = await BaseLLMHTTPHandler()._aresumable_chunked_upload(
|
||||
client=_async_handler_with(mock),
|
||||
initiate_url="https://storage.googleapis.com/upload?uploadType=resumable",
|
||||
base_headers={"Authorization": "Bearer x"},
|
||||
config=config,
|
||||
timeout=None,
|
||||
)
|
||||
|
||||
assert state["ranges"][-1] == f"bytes {total - chunk_size}-{total - 1}/{total}"
|
||||
assert "*" not in state["ranges"][-1]
|
||||
assert bytes(state["received"]) == b"a" * total
|
||||
assert response.status_code == 200
|
||||
|
||||
async def test_failed_chunk_raises(self):
|
||||
async def test_failed_upload_raises(self):
|
||||
raw = _make_openai_jsonl_bytes(80)
|
||||
with pytest.raises(Exception):
|
||||
await self._run(raw, chunk_size=4096, final_status=403)
|
||||
await self._run(raw, status=403)
|
||||
|
||||
async def test_request_timeout_is_forwarded(self):
|
||||
# The caller's per-request timeout must reach the GCS upload; every other
|
||||
# upload branch forwards it. httpx records the resolved timeout in
|
||||
# request.extensions["timeout"]; a dropped timeout would show the client
|
||||
# default instead of the value passed here.
|
||||
raw = _make_openai_jsonl_bytes(20)
|
||||
_, state, _ = await self._run(raw, timeout=httpx.Timeout(137.0))
|
||||
forwarded = state["timeouts"][0]
|
||||
assert forwarded is not None
|
||||
assert forwarded.get("read") == 137.0 and forwarded.get("write") == 137.0
|
||||
|
||||
async def test_upload_does_not_stage_to_disk(self, monkeypatch):
|
||||
# Disk-exhaustion guard: the transformed body must stream to GCS, never be
|
||||
# written to a temp file first. If any tempfile is created during the
|
||||
# upload, an attacker could fill the proxy's temp volume with large
|
||||
# concurrent uploads.
|
||||
created = []
|
||||
real_tempfile = tempfile.TemporaryFile
|
||||
monkeypatch.setattr(tempfile, "TemporaryFile", lambda *a, **k: (created.append(1), real_tempfile(*a, **k))[1])
|
||||
await self._run(_make_openai_jsonl_bytes(50))
|
||||
assert created == []
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
4
uv.lock
generated
4
uv.lock
generated
|
|
@ -9,7 +9,7 @@ resolution-markers = [
|
|||
]
|
||||
|
||||
[options]
|
||||
exclude-newer = "2026-06-22T01:17:31.93796Z"
|
||||
exclude-newer = "2026-06-27T00:53:21.064066Z"
|
||||
exclude-newer-span = "P3D"
|
||||
|
||||
[manifest]
|
||||
|
|
@ -3294,7 +3294,7 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "litellm"
|
||||
version = "1.89.4"
|
||||
version = "1.89.5"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue