diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 14c3b7736a3..1bcf1b571bd 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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": ( - "" - if isinstance(transformed_request, dict) - and "resumable_chunked_upload" in transformed_request + "" + if isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request else transformed_request ), "api_base": api_base, @@ -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( diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index d5164d8c1c2..5abaedae5a8 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -60,7 +60,7 @@ from litellm.types.llms.openai import ( OpenAIFileObject, PathLike, ) -from litellm.types.files import ResumableChunkedUploadConfig +from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.vertex_ai import GcsBucketResponse from litellm.types.utils import LlmProviders, ModelResponse @@ -389,21 +389,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): raise ValueError("file is required") if purpose is None: raise ValueError("purpose is required") - _, content_type = extract_file_metadata(file_data) object_name = self.get_object_name(file_data, purpose) if object_prefix: object_name = f"{object_prefix}/{object_name}" encoded_object_name = encode_gcs_object_name_for_url(object_name) - # Batch jsonl is streamed via a resumable session (bounded memory on - # large uploads); everything else is a single simple-media upload. - upload_type = ( - "resumable" - if FilesAPIUtils.is_batch_jsonl_request( - create_file_data=data, content_type=content_type - ) - else "media" - ) - endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType={upload_type}&name={encoded_object_name}" + endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_object_name}" api_base = api_base or "https://storage.googleapis.com" if not api_base: raise ValueError("api_base is required") @@ -455,8 +445,9 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ 2 Cases: 1. Handle basic file upload - 2. Handle batch file upload (.jsonl), streamed to a GCS resumable - session so large uploads stay memory-bounded. + 2. Handle batch file upload (.jsonl), staged to a temp file and uploaded + in a single media request so large uploads stay memory-bounded without + the per-chunk round-trips of a resumable session. """ file_data = create_file_data.get("file") if file_data is None: @@ -468,14 +459,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): content_type=content_type, ): return { - "resumable_chunked_upload": ResumableChunkedUploadConfig( + "streaming_media_upload": StreamingMediaUploadConfig( body_stream=_OpenAIToVertexBatchUploadStream( file_data, self._map_openai_to_vertex_params, ), - initiate_headers={ - "X-Upload-Content-Type": "application/json", - }, + content_type="application/json", ) } diff --git a/litellm/proxy/client/exceptions.py b/litellm/proxy/client/exceptions.py index fffd1b78b7e..c4089381e30 100644 --- a/litellm/proxy/client/exceptions.py +++ b/litellm/proxy/client/exceptions.py @@ -2,18 +2,30 @@ from typing import Union import requests +from litellm.litellm_core_utils.secret_redaction import redact_string + + +def _redact_orig_exception( + orig_exception: Union[requests.exceptions.HTTPError, str], +) -> Union[requests.exceptions.HTTPError, str]: + if isinstance(orig_exception, requests.exceptions.HTTPError): + return requests.exceptions.HTTPError( + redact_string(str(orig_exception)), response=orig_exception.response + ) + return redact_string(str(orig_exception)) + class UnauthorizedError(Exception): """Exception raised when the API returns a 401 Unauthorized response.""" def __init__(self, orig_exception: Union[requests.exceptions.HTTPError, str]): - self.orig_exception = orig_exception - super().__init__(str(orig_exception)) + self.orig_exception = _redact_orig_exception(orig_exception) + super().__init__(str(self.orig_exception)) class NotFoundError(Exception): """Exception raised when the API returns a 404 Not Found response or indicates a resource was not found.""" def __init__(self, orig_exception: Union[requests.exceptions.HTTPError, str]): - self.orig_exception = orig_exception - super().__init__(str(orig_exception)) + self.orig_exception = _redact_orig_exception(orig_exception) + super().__init__(str(self.orig_exception)) diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index d8687cbad16..845b49d1581 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -2,6 +2,8 @@ from typing import Any, Dict, List, Optional, Union import requests +from litellm.litellm_core_utils.secret_redaction import redact_string + from .exceptions import UnauthorizedError @@ -314,6 +316,9 @@ class KeysManagementClient: response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: + redacted_message = redact_string(str(e)) if e.response.status_code == 401: - raise UnauthorizedError(e) - raise + raise UnauthorizedError(e) from None + raise requests.exceptions.HTTPError( + redacted_message, response=e.response + ) from None diff --git a/litellm/types/files.py b/litellm/types/files.py index 1b2d7e30f1f..97c4dbb3abc 100644 --- a/litellm/types/files.py +++ b/litellm/types/files.py @@ -323,19 +323,18 @@ class TwoStepFileUploadConfig(TypedDict, total=False): upload_url_key: str -class ResumableChunkedUploadConfig(TypedDict, total=False): - """Drives a memory-bounded resumable upload (GCS JSON API). +class StreamingMediaUploadConfig(TypedDict, total=False): + """Drives a memory-bounded single-request upload (GCS simple/media upload). - The handler POSTs to the upload URL to open a session, reads the session URI - from ``session_url_header``, then PUTs ``body_stream`` to that URI in - ``chunk_size``-byte chunks (a 256 KiB multiple) using Content-Range, so the - payload is never buffered in full and the transfer is resumable. + The handler stages ``body_stream`` to a temp file off the event loop (so peak + memory stays bounded), then PUTs/POSTs it in one request with a known + Content-Length. Unlike a resumable chunked upload this incurs no per-chunk + round-trips, so a multi-GB upload finishes in one continuous transfer instead + of hundreds of sequential PUTs that overrun client/LB timeouts. ``body_stream`` is a ``BaseFileUploadStream``; it is typed ``Any`` here to avoid importing the llms layer into types. """ body_stream: Required[Any] - chunk_size: int - session_url_header: str - initiate_headers: Dict[str, str] + content_type: str diff --git a/pyproject.toml b/pyproject.toml index 9eb07d8f108..1a9845bd636 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", ] diff --git a/tests/batches_tests/test_openai_batches_and_files.py b/tests/batches_tests/test_openai_batches_and_files.py index 8a2d5f33805..0a49b3d77d1 100644 --- a/tests/batches_tests/test_openai_batches_and_files.py +++ b/tests/batches_tests/test_openai_batches_and_files.py @@ -513,25 +513,26 @@ async def test_avertex_batch_prediction(monkeypatch): mock_response.status_code = 200 return mock_response - # Batch jsonl file creation now streams to a GCS resumable session via - # _aresumable_chunked_upload (httpx send), not AsyncHTTPHandler.post, so mock - # that entry point to return the GCS object response. The resumable protocol - # itself is covered in test_vertex_ai_files_streaming.py. - mock_upload_response = httpx.Response( - 200, - json=mock_file_response, - request=httpx.Request("PUT", "https://storage.googleapis.com/upload"), - ) + # Batch jsonl creation now stages the body to a temp file and issues a single + # uploadType=media POST against the raw httpx.AsyncClient (client.client) inside + # _astage_and_upload_media, not AsyncHTTPHandler.post. Patch that raw POST so the + # real staging/upload + response transform run while the GCS object response is + # mocked; AsyncHTTPHandler.post still handles the batch-prediction call. with ( patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", side_effect=mock_side_effect, - ) as mock_global_post, - patch( - "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler._aresumable_chunked_upload", - new_callable=AsyncMock, - return_value=mock_upload_response, ), + patch.object( + httpx.AsyncClient, + "post", + new_callable=AsyncMock, + return_value=httpx.Response( + 200, + json=mock_file_response, + request=httpx.Request("POST", "https://storage.googleapis.com/upload"), + ), + ) as mock_gcs_upload, ): litellm.set_verbose = True litellm._turn_on_debug() @@ -552,6 +553,15 @@ async def test_avertex_batch_prediction(monkeypatch): == "gs://litellm-local/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/5f7b99ad-9203-4430-98bf-3b45451af4cb" ) + mock_gcs_upload.assert_awaited_once() + upload_url = str(mock_gcs_upload.call_args.args[0]) + assert "uploadType=media" in upload_url + assert "/b/litellm-local/o" in upload_url + assert ( + mock_gcs_upload.call_args.kwargs["headers"]["Content-Type"] + == "application/json" + ) + # Create batch create_batch_response = await litellm.acreate_batch( completion_window="24h", diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py index 122518d4acb..d2ee9d7d659 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py @@ -37,9 +37,7 @@ class TestVertexAIBinaryFileUpload: # Create mock PDF binary data (with non-UTF-8 bytes) # PDF files start with %PDF- and contain binary data mock_pdf_content = b"%PDF-1.4\n%\xc4\xe5\xf2\xe5\xeb\xa7\xf3\xa0\xd0\xc4\xc6\n" - mock_pdf_content += ( - b"\x00\x01\x02\x03\xff\xfe\xfd" * 100 - ) # Add more binary data + mock_pdf_content += b"\x00\x01\x02\x03\xff\xfe\xfd" * 100 # Add more binary data # Create file object file_obj = io.BytesIO(mock_pdf_content) @@ -60,14 +58,12 @@ class TestVertexAIBinaryFileUpload: ) # Verify the transformation returns bytes (not string) - assert isinstance( - transformed_request, bytes - ), f"Expected bytes for binary file, got {type(transformed_request)}" + assert isinstance(transformed_request, bytes), ( + f"Expected bytes for binary file, got {type(transformed_request)}" + ) # Verify the bytes match the original content - assert ( - transformed_request == mock_pdf_content - ), "Transformed request should preserve binary content exactly" + assert transformed_request == mock_pdf_content, "Transformed request should preserve binary content exactly" # Verify that the bytes contain non-UTF-8 characters # This should raise UnicodeDecodeError if we try to decode @@ -132,16 +128,14 @@ class TestVertexAIBinaryFileUpload: pytest.fail(f"httpx should accept bytes in data parameter: {e}") # Document the expected behavior - assert isinstance( - mock_binary_data, bytes - ), "Binary file data should remain as bytes" + assert isinstance(mock_binary_data, bytes), "Binary file data should remain as bytes" @pytest.mark.asyncio - async def test_jsonl_file_upload_returns_resumable_stream(self): + async def test_jsonl_file_upload_returns_streaming_body(self): """ - Test that JSONL batch files are transformed into a resumable-upload config + Test that JSONL batch files are transformed into a streaming-media config carrying a streaming body (not a buffered bytes payload), so the handler - can stream the upload to GCS in bounded chunks. + can stage the upload to a temp file and send it in one media request. """ # Create mock JSONL content mock_jsonl_content = ( @@ -164,16 +158,13 @@ class TestVertexAIBinaryFileUpload: litellm_params={}, ) - assert ( - isinstance(transformed_request, dict) - and "resumable_chunked_upload" in transformed_request - ), f"Expected a resumable upload config for JSONL, got {type(transformed_request)}" + assert isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request, ( + f"Expected a streaming media upload config for JSONL, got {type(transformed_request)}" + ) - stream = transformed_request["resumable_chunked_upload"]["body_stream"] + stream = transformed_request["streaming_media_upload"]["body_stream"] decoded = json.loads(b"".join(stream.iter_bytes()).decode("utf-8")) - assert ( - "request" in decoded - ), "JSONL transform must wrap each row in {'request': ...}" + assert "request" in decoded, "JSONL transform must wrap each row in {'request': ...}" @pytest.mark.asyncio async def test_mixed_file_types_in_sequence(self): @@ -214,7 +205,7 @@ class TestVertexAIBinaryFileUpload: optional_params={}, litellm_params={}, ) - assert isinstance(result2, dict) and "resumable_chunked_upload" in result2 + assert isinstance(result2, dict) and "streaming_media_upload" in result2 # Test 3: Upload another binary file binary_content2 = b"\xc4\xe5\xf2\xe5\xeb" @@ -264,7 +255,5 @@ class TestVertexAIBinaryFileUpload: }, } - assert ( - expected_behavior["binary_files"]["encoding"] == "none - preserve raw bytes" - ) + assert expected_behavior["binary_files"]["encoding"] == "none - preserve raw bytes" assert expected_behavior["text_files"]["encoding"] == "UTF-8" diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index cd556c48b6b..2e3280c0ed1 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -20,6 +20,7 @@ replaced by a list-based pipeline: import gc import io import json +import tempfile import time import tracemalloc @@ -41,15 +42,15 @@ from litellm.llms.vertex_ai.files.transformation import ( from litellm.types.llms.openai import CreateFileRequest -def _resumable_stream(transformed) -> BaseFileUploadStream: - """Pull the streaming body out of a resumable-upload transform result.""" - return transformed["resumable_chunked_upload"]["body_stream"] +def _upload_stream(transformed) -> BaseFileUploadStream: + """Pull the streaming body out of the upload transform result.""" + return transformed["streaming_media_upload"]["body_stream"] def _join_upload_body(transformed) -> bytes: """Materialize a transform result's upload body for byte-level assertions.""" - if isinstance(transformed, dict) and "resumable_chunked_upload" in transformed: - return b"".join(_resumable_stream(transformed).iter_bytes()) + if isinstance(transformed, dict) and "streaming_media_upload" in transformed: + return b"".join(_upload_stream(transformed).iter_bytes()) if isinstance(transformed, BaseFileUploadStream): return b"".join(transformed.iter_bytes()) if isinstance(transformed, str): @@ -83,17 +84,13 @@ def _reference_vertex_jsonl_string(cfg: VertexAIFilesConfig, content: str) -> st transform, so the streaming path can be checked against it for parity.""" entries = [json.loads(line) for line in content.splitlines() if line.strip()] return "\n".join( - json.dumps( - _openai_batch_jsonl_entry_to_vertex_wrapped_request( - entry, cfg._map_openai_to_vertex_params - ) - ) + json.dumps(_openai_batch_jsonl_entry_to_vertex_wrapped_request(entry, cfg._map_openai_to_vertex_params)) for entry in entries ) class TestStreamingOutputParity: - def test_transform_create_file_request_returns_resumable_stream_parity(self): + def test_transform_create_file_request_returns_streaming_body_parity(self): cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(300) request: CreateFileRequest = { @@ -105,14 +102,12 @@ class TestStreamingOutputParity: model="", create_file_data=request, optional_params={}, litellm_params={} ) - # A batch upload must be a resumable-upload config carrying a streaming - # body, so the handler can chunk it; a buffered bytes/str return would - # defeat the OOM fix. - assert isinstance(out, dict) and "resumable_chunked_upload" in out - assert isinstance(_resumable_stream(out), BaseFileUploadStream) - assert _join_upload_body(out).decode("utf-8") == _reference_vertex_jsonl_string( - cfg, raw.decode("utf-8") - ) + # A batch upload must be a streaming-media config carrying a streaming + # body, so the handler can stream it to GCS; a buffered bytes/str return + # would defeat the OOM fix. + assert isinstance(out, dict) and "streaming_media_upload" in out + assert isinstance(_upload_stream(out), BaseFileUploadStream) + assert _join_upload_body(out).decode("utf-8") == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) class TestFileLikeInputNotPartiallyConsumed: @@ -215,9 +210,7 @@ class TestStreamingLineIterator: def seek(self, *args): raise io.UnsupportedOperation("not seekable") - handle = _NonSeekable( - b'{"custom_id": "request-0"}\n{"custom_id": "request-1"}\n' - ) + handle = _NonSeekable(b'{"custom_id": "request-0"}\n{"custom_id": "request-1"}\n') with pytest.raises(ValueError, match="seekable"): list(_iter_openai_jsonl_lines(handle)) @@ -235,13 +228,8 @@ class TestGetObjectNameLazyParse: cfg = VertexAIFilesConfig() # Tail rows are deliberately not valid JSON. Parsing the whole payload # would raise here; a first-row-only parse must not. - raw = ( - b'{"custom_id": "r-0", "body": {"model": "gemini-2.5-flash"}}\n' - b"garbage line that is not json\n" - ) - object_name = cfg.get_object_name( - ("batch.jsonl", raw, "application/jsonl"), purpose="batch" - ) + raw = b'{"custom_id": "r-0", "body": {"model": "gemini-2.5-flash"}}\ngarbage line that is not json\n' + object_name = cfg.get_object_name(("batch.jsonl", raw, "application/jsonl"), purpose="batch") assert "gemini-2.5-flash" in object_name @@ -278,15 +266,11 @@ class TestStreamingPeakMemory: def drain_stream(): # Consume the upload body one row at a time, as the chunked uploader # does, without accumulating it. - for _ in _OpenAIToVertexBatchUploadStream( - raw, cfg._map_openai_to_vertex_params - ).iter_bytes(): + for _ in _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params).iter_bytes(): pass streaming_peak = self._measure(drain_stream) - list_peak = self._measure( - lambda: _reference_vertex_jsonl_string(cfg, content_str) - ) + list_peak = self._measure(lambda: _reference_vertex_jsonl_string(cfg, content_str)) # Core guard: the lazily consumed streaming body peaks well under a list # pipeline that materializes every transformed row. Building full @@ -305,9 +289,7 @@ class TestStreamingPeakMemory: # first-row parse should allocate only a small fraction of the payload; # parsing every row would blow past this bound. peak = self._measure(lambda: cfg.get_object_name(file_data, purpose="batch")) - assert ( - peak / len(raw) < 2.0 - ), "get_object_name should not copy the whole payload" + assert peak / len(raw) < 2.0, "get_object_name should not copy the whole payload" class TestPathSourcedStreaming: @@ -341,12 +323,10 @@ class TestPathSourcedStreaming: litellm_params={"gcs_bucket_name": "test-bucket"}, data=data, ) - assert "uploadType=resumable" in url + assert "uploadType=media" in url - out = cfg.transform_create_file_request( - model="", create_file_data=data, optional_params={}, litellm_params={} - ) - assert isinstance(out, dict) and "resumable_chunked_upload" in out + out = cfg.transform_create_file_request(model="", create_file_data=data, optional_params={}, litellm_params={}) + assert isinstance(out, dict) and "streaming_media_upload" in out body = _join_upload_body(out).decode("utf-8") assert body == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) lines = body.splitlines() @@ -371,7 +351,7 @@ class TestPathSourcedStreaming: out = cfg.transform_create_file_request( model="", create_file_data=data, optional_params={}, litellm_params={} ) - for _ in _resumable_stream(out).iter_bytes(): + for _ in _upload_stream(out).iter_bytes(): pass # drain without accumulating gc.collect() @@ -384,20 +364,15 @@ class TestPathSourcedStreaming: # Streaming from disk must not materialize the payload. Reading the whole # file into bytes (the pre-fix path) would push peak past the file size. - assert peak < len(raw) * 0.3, ( - f"peak {peak} not bounded vs payload {len(raw)} " - f"(ratio {peak / len(raw):.2f})" - ) + assert peak < len(raw) * 0.3, f"peak {peak} not bounded vs payload {len(raw)} (ratio {peak / len(raw):.2f})" def test_path_source_stream_is_reiterable(self, tmp_path): cfg = VertexAIFilesConfig() path, _ = self._write_jsonl(tmp_path, 50) data = self._batch_request(path) - out = cfg.transform_create_file_request( - model="", create_file_data=data, optional_params={}, litellm_params={} - ) - stream = _resumable_stream(out) + out = cfg.transform_create_file_request(model="", create_file_data=data, optional_params={}, litellm_params={}) + stream = _upload_stream(out) first = b"".join(stream.iter_bytes()) second = b"".join(stream.iter_bytes()) assert first == second and len(first) > 0 @@ -436,25 +411,20 @@ def _logging_obj() -> Logging: ) -def _gcs_resumable_mock(session_url: str, final_status: int = 200): - """A fake GCS resumable endpoint: POST opens a session (URI in Location), - each PUT appends and returns 308 until the final chunk returns 200/201.""" - state = {"received": bytearray(), "ranges": [], "methods": [], "urls": []} +def _gcs_media_mock(status: int = 200): + """A fake GCS simple-media endpoint: one request carries the whole object; + capture the body and headers and return the object resource.""" + state = {"received": bytearray(), "methods": [], "urls": [], "headers": [], "timeouts": []} async def handler(request: httpx.Request) -> httpx.Response: state["methods"].append(request.method) state["urls"].append(str(request.url)) - if request.method == "POST": - return httpx.Response(200, headers={"location": session_url}) - body = await request.aread() - content_range = request.headers["content-range"] - state["ranges"].append(content_range) - state["received"].extend(body) - if content_range.rsplit("/", 1)[-1] == "*": - return httpx.Response( - 308, headers={"range": f"bytes=0-{len(state['received']) - 1}"} - ) - return httpx.Response(final_status, json=_GCS_OBJECT_JSON) + state["headers"].append(dict(request.headers)) + # httpx records the resolved per-request timeout here, so the test can + # assert the caller's timeout was forwarded rather than the client default. + state["timeouts"].append(request.extensions.get("timeout")) + state["received"].extend(await request.aread()) + return httpx.Response(status, json=_GCS_OBJECT_JSON) return handler, state @@ -465,8 +435,8 @@ def _async_handler_with(mock) -> AsyncHTTPHandler: return handler -class TestResumableUploadUrl: - def test_batch_jsonl_uses_resumable_upload_type(self): +class TestUploadUrl: + def test_batch_jsonl_uses_media_upload_type(self): cfg = VertexAIFilesConfig() request: CreateFileRequest = { "file": ("batch.jsonl", _make_openai_jsonl_bytes(3), "application/jsonl"), @@ -480,29 +450,12 @@ class TestResumableUploadUrl: litellm_params={"gcs_bucket_name": "test-bucket"}, data=request, ) - assert "uploadType=resumable" in url - assert "uploadType=media" not in url + # A single media upload is one continuous transfer (no per-chunk + # round-trips), which is what keeps large uploads under client/LB timeouts. + assert "uploadType=media" in url + assert "uploadType=resumable" not in url - def test_batch_text_plain_uses_resumable_upload_type(self): - # Clients often label a .jsonl batch upload as text/plain; it must still - # take the streaming/resumable path, not the buffered media path. - cfg = VertexAIFilesConfig() - request: CreateFileRequest = { - "file": ("batch.jsonl", _make_openai_jsonl_bytes(3), "text/plain"), - "purpose": "batch", - } - url = cfg.get_complete_file_url( - api_base=None, - api_key=None, - model="", - optional_params={}, - litellm_params={"gcs_bucket_name": "test-bucket"}, - data=request, - ) - assert "uploadType=resumable" in url - assert "uploadType=media" not in url - - def test_binary_upload_stays_simple_media(self): + def test_binary_upload_uses_media_upload_type(self): cfg = VertexAIFilesConfig() request: CreateFileRequest = { "file": ("doc.pdf", b"%PDF-1.4 binary", "application/pdf"), @@ -520,14 +473,12 @@ class TestResumableUploadUrl: assert "uploadType=resumable" not in url -class TestResumableStreamBody: +class TestUploadStreamBody: def test_stream_matches_legacy_pipeline(self): cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(120) stream = _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params) - assert b"".join(stream.iter_bytes()).decode( - "utf-8" - ) == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) + assert b"".join(stream.iter_bytes()).decode("utf-8") == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) def test_stream_is_reiterable_for_retries(self): # A one-shot generator would make a transport retry upload an empty body; @@ -545,59 +496,19 @@ class TestResumableStreamBody: # an empty body silently. cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(40) - stream = _OpenAIToVertexBatchUploadStream( - io.BytesIO(raw), cfg._map_openai_to_vertex_params - ) + stream = _OpenAIToVertexBatchUploadStream(io.BytesIO(raw), cfg._map_openai_to_vertex_params) first = b"".join(stream.iter_bytes()) second = b"".join(stream.iter_bytes()) assert first == second and len(first) > 0 -class TestResumableChunking: - def test_intermediate_chunks_are_exactly_chunk_size(self): - pieces = list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([b"x" * 10]), 4)) - assert pieces == [b"xxxx", b"xxxx", b"xx"] - - def test_exact_multiple_yields_no_trailing_empty(self): - # An exactly chunk-aligned stream yields only full chunks; the upload - # finalizes on the last data chunk instead of an extra empty request. - pieces = list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([b"x" * 8]), 4)) - assert pieces == [b"xxxx", b"xxxx"] - - def test_empty_stream_yields_nothing(self): - # A 0-byte stream yields no chunks; the caller finalizes with one empty - # request (bytes */0). - assert list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([]), 4)) == [] - - def test_default_chunk_size_is_256kib_multiple(self): - assert BaseLLMHTTPHandler._RESUMABLE_CHUNK_SIZE % (256 * 1024) == 0 - - def test_content_range_intermediate_uses_star_total(self): - assert ( - BaseLLMHTTPHandler._resumable_content_range(0, 4096, is_final=False) - == "bytes 0-4095/*" - ) - - def test_content_range_final_uses_real_total(self): - assert ( - BaseLLMHTTPHandler._resumable_content_range(8192, 100, is_final=True) - == "bytes 8192-8291/8292" - ) - - def test_content_range_empty_finalize(self): - assert ( - BaseLLMHTTPHandler._resumable_content_range(8192, 0, is_final=True) - == "bytes */8192" - ) - - @pytest.mark.asyncio -class TestResumableUploadProtocol: - """End-to-end against a faked GCS resumable endpoint. These are the tests - that fail if the handler buffers the whole body, drops bytes, mislabels a - Content-Range, follows the 308 instead of continuing, or skips finalize.""" +class TestStreamingMediaUpload: + """End-to-end against a faked GCS media endpoint. These fail if the handler + buffers the payload in memory, drops bytes, omits Content-Length (which would + flip httpx to chunked transfer-encoding), or makes more than one request.""" - async def _run(self, raw: bytes, chunk_size: int, final_status: int = 200): + async def _run(self, raw: bytes, status: int = 200, timeout=None): cfg = VertexAIFilesConfig() request: CreateFileRequest = { "file": ("batch.jsonl", raw, "application/jsonl"), @@ -614,11 +525,8 @@ class TestResumableUploadProtocol: transformed = cfg.transform_create_file_request( model="", create_file_data=request, optional_params={}, litellm_params={} ) - transformed["resumable_chunked_upload"]["chunk_size"] = chunk_size expected = _join_upload_body(transformed) - - session_url = "https://storage.googleapis.com/upload/sess?upload_id=SID" - mock, state = _gcs_resumable_mock(session_url, final_status=final_status) + mock, state = _gcs_media_mock(status=status) response = await BaseLLMHTTPHandler().async_create_file( transformed_request=transformed, litellm_params={}, @@ -627,70 +535,52 @@ class TestResumableUploadProtocol: api_base=api_base, logging_obj=_logging_obj(), client=_async_handler_with(mock), - timeout=None, + timeout=timeout, ) - return expected, state, response, session_url, api_base + return expected, state, response - async def test_streams_in_chunks_and_reassembles(self): + async def test_single_request_carries_whole_payload(self): raw = _make_openai_jsonl_bytes(300) - chunk_size = 4096 - expected, state, response, session_url, api_base = await self._run( - raw, chunk_size - ) + expected, state, response = await self._run(raw) - # One session-open POST, then a sequence of chunk PUTs. - assert state["methods"][0] == "POST" - assert set(state["methods"][1:]) == {"PUT"} - assert state["methods"].count("PUT") >= 2, "payload must span multiple chunks" + # Exactly one request (the single media upload), and it lands on the + # media endpoint, not a resumable session. + assert state["methods"] == ["POST"] + assert "uploadType=media" in state["urls"][0] - # POST opens a resumable session; every chunk goes to the session URI. - assert "uploadType=resumable" in state["urls"][0] - assert all(u == session_url for u in state["urls"][1:]) - - # Every non-final chunk is exactly chunk_size with an unknown-total range; - # the final chunk carries the real total. - intermediate = state["ranges"][:-1] - for index, content_range in enumerate(intermediate): - assert ( - content_range - == f"bytes {index * chunk_size}-{(index + 1) * chunk_size - 1}/*" - ) - total = len(expected) - last_offset = len(intermediate) * chunk_size - if last_offset == total: # payload landed on a chunk boundary - assert state["ranges"][-1] == f"bytes */{total}" - else: - assert state["ranges"][-1] == f"bytes {last_offset}-{total - 1}/{total}" - - # The bytes GCS received are exactly the transformed batch payload. + # The body is streamed with chunked transfer-encoding and no + # Content-Length, which is what proves it is neither buffered in memory + # nor staged to a temp file (the disk-exhaustion guard) before sending. + headers = state["headers"][0] + assert headers.get("transfer-encoding") == "chunked" + assert "content-length" not in headers + # httpx reassembles the chunked body; GCS receives exactly the transform. assert bytes(state["received"]) == expected assert response.object == "file" - async def test_exact_multiple_finalizes_on_last_data_chunk(self): - # A body that is an exact multiple of the chunk size finalizes on its - # last data chunk (bytes (TOTAL-chunk)-(TOTAL-1)/TOTAL), with no extra - # empty finalize request. - chunk_size = 256 - total = chunk_size * 3 - stream = _FixedBytesStream(b"a" * total) - config = {"body_stream": stream, "chunk_size": chunk_size} - session_url = "https://storage.googleapis.com/upload/sess?upload_id=SID" - mock, state = _gcs_resumable_mock(session_url) - - response = await BaseLLMHTTPHandler()._aresumable_chunked_upload( - client=_async_handler_with(mock), - initiate_url="https://storage.googleapis.com/upload?uploadType=resumable", - base_headers={"Authorization": "Bearer x"}, - config=config, - timeout=None, - ) - - assert state["ranges"][-1] == f"bytes {total - chunk_size}-{total - 1}/{total}" - assert "*" not in state["ranges"][-1] - assert bytes(state["received"]) == b"a" * total - assert response.status_code == 200 - - async def test_failed_chunk_raises(self): + async def test_failed_upload_raises(self): raw = _make_openai_jsonl_bytes(80) with pytest.raises(Exception): - await self._run(raw, chunk_size=4096, final_status=403) + await self._run(raw, status=403) + + async def test_request_timeout_is_forwarded(self): + # The caller's per-request timeout must reach the GCS upload; every other + # upload branch forwards it. httpx records the resolved timeout in + # request.extensions["timeout"]; a dropped timeout would show the client + # default instead of the value passed here. + raw = _make_openai_jsonl_bytes(20) + _, state, _ = await self._run(raw, timeout=httpx.Timeout(137.0)) + forwarded = state["timeouts"][0] + assert forwarded is not None + assert forwarded.get("read") == 137.0 and forwarded.get("write") == 137.0 + + async def test_upload_does_not_stage_to_disk(self, monkeypatch): + # Disk-exhaustion guard: the transformed body must stream to GCS, never be + # written to a temp file first. If any tempfile is created during the + # upload, an attacker could fill the proxy's temp volume with large + # concurrent uploads. + created = [] + real_tempfile = tempfile.TemporaryFile + monkeypatch.setattr(tempfile, "TemporaryFile", lambda *a, **k: (created.append(1), real_tempfile(*a, **k))[1]) + await self._run(_make_openai_jsonl_bytes(50)) + assert created == [] diff --git a/tests/test_litellm/proxy/client/test_keys.py b/tests/test_litellm/proxy/client/test_keys.py index 136408e01ac..620daefb39e 100644 --- a/tests/test_litellm/proxy/client/test_keys.py +++ b/tests/test_litellm/proxy/client/test_keys.py @@ -1,5 +1,6 @@ import os import sys +import traceback import pytest import requests @@ -11,7 +12,7 @@ sys.path.insert( import responses -from litellm.proxy.client.exceptions import UnauthorizedError +from litellm.proxy.client.exceptions import NotFoundError, UnauthorizedError from litellm.proxy.client.keys import KeysManagementClient @@ -420,3 +421,96 @@ def test_info_server_error(client): ) with pytest.raises(requests.exceptions.HTTPError): client.info(key="test-key") + + +LEAKY_KEY = "sk-1234567890abcdefghijklmnop" + + +def _render_full_traceback(exc: BaseException) -> str: + return "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)) + + +@responses.activate +def test_info_not_found_redacts_key_everywhere(client): + """A 404 must not echo the raw key embedded in the request URL. + + Covers str(exc) and the rendered traceback, since the chain through + __cause__ / __context__ is what logging.exception() and the default + excepthook print. + """ + responses.add( + responses.GET, + f"{client._base_url}/key/info?key={LEAKY_KEY}", + status=404, + json={"error": {"message": "Key not found", "code": "404"}}, + ) + with pytest.raises(requests.exceptions.HTTPError) as excinfo: + client.info(key=LEAKY_KEY) + + exc = excinfo.value + assert LEAKY_KEY not in str(exc) + assert "REDACTED" in str(exc) + assert LEAKY_KEY not in _render_full_traceback(exc) + assert exc.__cause__ is None and exc.__suppress_context__ + assert exc.response is not None + assert exc.response.status_code == 404 + assert exc.request is not None + # Known residual: the live request URL still carries the key, since the + # response is preserved so callers keep status_code / text. str(exc) and the + # traceback are scrubbed; the URL-borne key is the root issue tracked in + # LIT-4013 (move the lookup key out of the query string server-side). + assert LEAKY_KEY in exc.response.request.url + + +@responses.activate +def test_info_unauthorized_redacts_key_everywhere(client): + """A 401 surfaced as UnauthorizedError must not echo the raw key in the + message, the retained original, or the rendered traceback chain.""" + responses.add( + responses.GET, + f"{client._base_url}/key/info?key={LEAKY_KEY}", + status=401, + json={"error": "Unauthorized"}, + ) + with pytest.raises(UnauthorizedError) as excinfo: + client.info(key=LEAKY_KEY) + + exc = excinfo.value + assert LEAKY_KEY not in str(exc) + assert "REDACTED" in str(exc) + assert LEAKY_KEY not in str(exc.orig_exception) + assert LEAKY_KEY not in _render_full_traceback(exc) + assert exc.__cause__ is None and exc.__suppress_context__ + assert isinstance(exc.orig_exception, requests.exceptions.HTTPError) + assert exc.orig_exception.response is not None + assert exc.orig_exception.response.status_code == 401 + + +def _http_error_with_key(prefix: str, status: int) -> requests.exceptions.HTTPError: + resp = requests.Response() + resp.status_code = status + return requests.exceptions.HTTPError( + f"{prefix} for url: http://x/key/info?key={LEAKY_KEY}", response=resp + ) + + +def test_unauthorized_error_redacts_wrapped_key(): + """UnauthorizedError scrubs the key in str(exc) and in the retained + orig_exception, while preserving the response for structured access.""" + wrapped = UnauthorizedError( + _http_error_with_key("401 Client Error: Unauthorized", 401) + ) + assert LEAKY_KEY not in str(wrapped) + assert "REDACTED" in str(wrapped) + assert LEAKY_KEY not in str(wrapped.orig_exception) + assert wrapped.orig_exception.response.status_code == 401 + + +def test_not_found_error_redacts_wrapped_key(): + """NotFoundError scrubs the key in str(exc) and in the retained + orig_exception, while preserving the response for structured access.""" + wrapped = NotFoundError(_http_error_with_key("404 Client Error: Not Found", 404)) + assert LEAKY_KEY not in str(wrapped) + assert "REDACTED" in str(wrapped) + assert LEAKY_KEY not in str(wrapped.orig_exception) + assert wrapped.orig_exception.response.status_code == 404 diff --git a/uv.lock b/uv.lock index a8d563931e2..faaa5e5e741 100644 --- a/uv.lock +++ b/uv.lock @@ -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" },