diff --git a/src/openai/_base_client.py b/src/openai/_base_client.py index 216b36aabd..3c31f9c757 100644 --- a/src/openai/_base_client.py +++ b/src/openai/_base_client.py @@ -41,7 +41,7 @@ from . import _exceptions from ._qs import Querystring -from ._files import to_httpx_files, async_to_httpx_files +from ._files import ReplayableFiles, to_httpx_files, async_to_httpx_files from ._types import ( Body, Omit, @@ -453,6 +453,14 @@ def _custom_auth( ) -> httpx.Auth | None: return None + def _prepare_request_for_retry(self, request: httpx.Request) -> None: + replayable_files = request.extensions.get("openai_replayable_files") + if isinstance(replayable_files, ReplayableFiles): + replayable_files.rewind() + + def _should_prepare_replayable_files(self, options: FinalRequestOptions) -> bool: + return options.files is not None and options.get_max_retries(self.max_retries) > 0 + def _build_headers(self, options: FinalRequestOptions, *, retries_taken: int = 0) -> httpx.Headers: custom_headers = options.headers or {} headers_dict = _merge_mappings({**self._auth_headers(options.security), **self.default_headers}, custom_headers) @@ -1010,16 +1018,46 @@ def request( # ensure the idempotency key is reused between requests input_options.idempotency_key = self._idempotency_key() + replayable_files: ReplayableFiles | None = None + if self._should_prepare_replayable_files(input_options): + replayable_files = ReplayableFiles(input_options.files) + input_options.files = replayable_files.files + try: + return self._request_with_retries( + cast_to=cast_to, + input_options=input_options, + replayable_files=replayable_files, + stream=stream, + stream_cls=stream_cls, + ) + finally: + if replayable_files is not None: + replayable_files.close() + + def _request_with_retries( + self, + *, + cast_to: type[ResponseT], + input_options: FinalRequestOptions, + replayable_files: ReplayableFiles | None, + stream: bool, + stream_cls: type[_StreamT] | None, + ) -> ResponseT | _StreamT: response: httpx.Response | None = None max_retries = input_options.get_max_retries(self.max_retries) retries_taken = 0 + options: FinalRequestOptions | None = None for retries_taken in range(max_retries + 1): + if replayable_files is not None: + replayable_files.rewind() options = model_copy(input_options) options = self._prepare_options(options) remaining_retries = max_retries - retries_taken request = self._build_request(options, retries_taken=retries_taken) + if replayable_files is not None: + request.extensions["openai_replayable_files"] = replayable_files self._prepare_request(request) kwargs: HttpxSendArgs = {} @@ -1107,6 +1145,7 @@ def request( break assert response is not None, "could not resolve response (should never happen)" + assert options is not None, "could not resolve options (should never happen)" return self._process_response( cast_to=cast_to, options=options, @@ -1622,16 +1661,46 @@ async def request( # ensure the idempotency key is reused between requests input_options.idempotency_key = self._idempotency_key() + replayable_files: ReplayableFiles | None = None + if self._should_prepare_replayable_files(input_options): + replayable_files = ReplayableFiles(input_options.files) + input_options.files = replayable_files.files + try: + return await self._request_with_retries( + cast_to=cast_to, + input_options=input_options, + replayable_files=replayable_files, + stream=stream, + stream_cls=stream_cls, + ) + finally: + if replayable_files is not None: + replayable_files.close() + + async def _request_with_retries( + self, + *, + cast_to: type[ResponseT], + input_options: FinalRequestOptions, + replayable_files: ReplayableFiles | None, + stream: bool, + stream_cls: type[_AsyncStreamT] | None, + ) -> ResponseT | _AsyncStreamT: response: httpx.Response | None = None max_retries = input_options.get_max_retries(self.max_retries) retries_taken = 0 + options: FinalRequestOptions | None = None for retries_taken in range(max_retries + 1): + if replayable_files is not None: + replayable_files.rewind() options = model_copy(input_options) options = await self._prepare_options(options) remaining_retries = max_retries - retries_taken request = self._build_request(options, retries_taken=retries_taken) + if replayable_files is not None: + request.extensions["openai_replayable_files"] = replayable_files await self._prepare_request(request) kwargs: HttpxSendArgs = {} @@ -1718,6 +1787,7 @@ async def request( break assert response is not None, "could not resolve response (should never happen)" + assert options is not None, "could not resolve options (should never happen)" return await self._process_response( cast_to=cast_to, options=options, diff --git a/src/openai/_client.py b/src/openai/_client.py index 66d03b23dd..b08e6298da 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -463,10 +463,19 @@ def _send_with_auth_retry( response.close() self._workload_identity_auth.invalidate_token() request.headers["Authorization"] = f"Bearer {self._workload_identity_auth.get_token()}" + self._prepare_request_for_retry(request) return self._send_with_auth_retry(request, stream=stream, retried=True, **kwargs) return response + @override + def _should_prepare_replayable_files(self, options: FinalRequestOptions) -> bool: + return super()._should_prepare_replayable_files(options) or ( + options.files is not None + and self._workload_identity_auth is not None + and options.security.get("bearer_auth", False) + ) + @override def _send_request( self, @@ -1059,10 +1068,19 @@ async def _send_with_auth_retry( await response.aclose() self._workload_identity_auth.invalidate_token() request.headers["Authorization"] = f"Bearer {await self._workload_identity_auth.get_token_async()}" + self._prepare_request_for_retry(request) return await self._send_with_auth_retry(request, stream=stream, retried=True, **kwargs) return response + @override + def _should_prepare_replayable_files(self, options: FinalRequestOptions) -> bool: + return super()._should_prepare_replayable_files(options) or ( + options.files is not None + and self._workload_identity_auth is not None + and options.security.get("bearer_auth", False) + ) + @override async def _send_request( self, diff --git a/src/openai/_files.py b/src/openai/_files.py index 1a2cc77478..83faacf892 100644 --- a/src/openai/_files.py +++ b/src/openai/_files.py @@ -2,8 +2,10 @@ import io import os +import shutil import pathlib -from typing import Sequence, cast, overload +import tempfile +from typing import IO, Sequence, cast, overload from typing_extensions import TypeVar, TypeGuard import anyio @@ -21,6 +23,69 @@ _T = TypeVar("_T") +_SPOOLED_FILE_MAX_SIZE = 8 * 1024 * 1024 + + +class ReplayableFiles: + def __init__(self, files: HttpxRequestFiles | None) -> None: + self._positions: list[tuple[IO[bytes], int]] = [] + self._temporary_files: list[IO[bytes]] = [] + try: + self.files = self._transform_files(files) + except Exception: + self.close() + raise + + def _transform_files(self, files: HttpxRequestFiles | None) -> HttpxRequestFiles | None: + if files is None: + return None + + if is_mapping_t(files): + return {key: self._transform_file(file) for key, file in files.items()} + + sequence = cast(Sequence[tuple[str, HttpxFileTypes]], files) + return [(key, self._transform_file(file)) for key, file in sequence] + + def _transform_file(self, file: HttpxFileTypes) -> HttpxFileTypes: + if isinstance(file, tuple): + return cast(HttpxFileTypes, (file[0], self._transform_file_content(file[1]), *file[2:])) + + return self._transform_file_content(file) + + def _transform_file_content(self, file: HttpxFileContent) -> HttpxFileContent: + if not isinstance(file, io.IOBase): + return file + + file = cast(IO[bytes], file) + try: + if file.seekable(): + position = file.tell() + file.seek(position) + self._positions.append((file, position)) + return file + except (OSError, ValueError): + pass + + buffered = cast(IO[bytes], tempfile.SpooledTemporaryFile(max_size=_SPOOLED_FILE_MAX_SIZE)) + try: + shutil.copyfileobj(file, buffered) + except Exception: + buffered.close() + raise + + buffered.seek(0) + self._positions.append((buffered, 0)) + self._temporary_files.append(buffered) + return buffered + + def rewind(self) -> None: + for file, position in self._positions: + file.seek(position) + + def close(self) -> None: + for file in self._temporary_files: + file.close() + def is_base64_file_input(obj: object) -> TypeGuard[Base64FileInput]: return isinstance(obj, io.IOBase) or isinstance(obj, os.PathLike) diff --git a/tests/test_client.py b/tests/test_client.py index 2d8955a58e..21ae647331 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -3,6 +3,7 @@ from __future__ import annotations import gc +import io import os import sys import json @@ -68,6 +69,43 @@ def mirror_request_content(request: httpx.Request) -> httpx.Response: return httpx.Response(200, content=request.content) +class NonSeekableBytesIO(io.BytesIO): + def __init__(self, initial_bytes: bytes) -> None: + super().__init__(initial_bytes) + self.read_count = 0 + + @override + def read(self, size: int | None = -1) -> bytes: + self.read_count += 1 + return super().read(size) + + @override + def seekable(self) -> bool: + return False + + @override + def seek(self, offset: int, whence: int = 0) -> int: + raise io.UnsupportedOperation("seek") + + @override + def tell(self) -> int: + raise io.UnsupportedOperation("tell") + + +class MockWorkloadIdentityAuth: + def __init__(self) -> None: + self.token_number = 1 + + def get_token(self) -> str: + return f"openai-access-token-{self.token_number}" + + async def get_token_async(self) -> str: + return self.get_token() + + def invalidate_token(self) -> None: + self.token_number += 1 + + # note: we can't use the httpx.MockTransport class as it consumes the request # body itself, which means we can't test that the body is read lazily class MockTransport(httpx.BaseTransport, httpx.AsyncBaseTransport): @@ -1140,6 +1178,88 @@ def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, client ).__enter__() assert _get_open_connections(client) == 0 + @pytest.mark.parametrize("failure_mode", ["status", "connection", "timeout"]) + @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) + def test_multipart_retry_replays_non_seekable_file(self, failure_mode: str) -> None: + file_content = b"non-seekable multipart content" + request_bodies: list[bytes] = [] + + def handler(request: httpx.Request) -> httpx.Response: + request_bodies.append(b"".join(cast(httpx.SyncByteStream, request.stream))) + if len(request_bodies) == 1: + if failure_mode == "status": + return httpx.Response(500) + if failure_mode == "timeout": + raise httpx.ReadTimeout("timed out", request=request) + raise httpx.ConnectError("connection failed", request=request) + return httpx.Response(200) + + with OpenAI( + base_url=base_url, + api_key=api_key, + max_retries=1, + http_client=httpx.Client(transport=MockTransport(handler=handler)), + ) as client: + response = client.post( + "/upload", + cast_to=httpx.Response, + files={"file": ("input.txt", NonSeekableBytesIO(file_content))}, + options={"headers": {"Content-Type": "multipart/form-data"}}, + ) + + assert response.status_code == 200 + assert [body.count(file_content) for body in request_bodies] == [1, 1] + + def test_multipart_without_retries_streams_non_seekable_file(self) -> None: + file_content = b"streamed non-seekable multipart content" + file = NonSeekableBytesIO(file_content) + + def handler(request: httpx.Request) -> httpx.Response: + assert file.read_count == 0 + body = b"".join(cast(httpx.SyncByteStream, request.stream)) + return httpx.Response(200, content=body) + + with OpenAI( + base_url=base_url, + api_key=api_key, + max_retries=0, + http_client=httpx.Client(transport=MockTransport(handler=handler)), + ) as client: + response = client.post( + "/upload", + cast_to=httpx.Response, + files={"file": ("input.txt", file)}, + options={"headers": {"Content-Type": "multipart/form-data"}}, + ) + + assert response.content.count(file_content) == 1 + assert file.read_count > 0 + + @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) + def test_multipart_retry_rewinds_seekable_file(self) -> None: + file_content = b"seekable multipart content" + file = io.BytesIO(file_content) + request_bodies: list[bytes] = [] + + def handler(request: httpx.Request) -> httpx.Response: + request_bodies.append(b"".join(cast(httpx.SyncByteStream, request.stream))) + return httpx.Response(500 if len(request_bodies) == 1 else 200) + + with OpenAI( + base_url=base_url, + api_key=api_key, + max_retries=1, + http_client=httpx.Client(transport=MockTransport(handler=handler)), + ) as client: + client.post( + "/upload", + cast_to=httpx.Response, + files={"file": ("input.txt", file)}, + options={"headers": {"Content-Type": "multipart/form-data"}}, + ) + + assert [body.count(file_content) for body in request_bodies] == [1, 1] + @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) @pytest.mark.respx(base_url=base_url) @@ -2396,6 +2516,65 @@ async def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, ).__aenter__() assert _get_open_connections(async_client) == 0 + @pytest.mark.parametrize("failure_mode", ["status", "connection", "timeout"]) + @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) + async def test_multipart_retry_replays_non_seekable_file(self, failure_mode: str) -> None: + file_content = b"non-seekable multipart content" + request_bodies: list[bytes] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + stream = cast(httpx.AsyncByteStream, request.stream) + request_bodies.append(b"".join([chunk async for chunk in stream])) + if len(request_bodies) == 1: + if failure_mode == "status": + return httpx.Response(500) + if failure_mode == "timeout": + raise httpx.ReadTimeout("timed out", request=request) + raise httpx.ConnectError("connection failed", request=request) + return httpx.Response(200) + + async with AsyncOpenAI( + base_url=base_url, + api_key=api_key, + max_retries=1, + http_client=httpx.AsyncClient(transport=MockTransport(handler=handler)), + ) as client: + response = await client.post( + "/upload", + cast_to=httpx.Response, + files={"file": ("input.txt", NonSeekableBytesIO(file_content))}, + options={"headers": {"Content-Type": "multipart/form-data"}}, + ) + + assert response.status_code == 200 + assert [body.count(file_content) for body in request_bodies] == [1, 1] + + async def test_multipart_without_retries_streams_non_seekable_file(self) -> None: + file_content = b"streamed non-seekable multipart content" + file = NonSeekableBytesIO(file_content) + + async def handler(request: httpx.Request) -> httpx.Response: + assert file.read_count == 0 + stream = cast(httpx.AsyncByteStream, request.stream) + body = b"".join([chunk async for chunk in stream]) + return httpx.Response(200, content=body) + + async with AsyncOpenAI( + base_url=base_url, + api_key=api_key, + max_retries=0, + http_client=httpx.AsyncClient(transport=MockTransport(handler=handler)), + ) as client: + response = await client.post( + "/upload", + cast_to=httpx.Response, + files={"file": ("input.txt", file)}, + options={"headers": {"Content-Type": "multipart/form-data"}}, + ) + + assert response.content.count(file_content) == 1 + assert file.read_count > 0 + @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) @mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) @pytest.mark.respx(base_url=base_url) @@ -2656,6 +2835,31 @@ async def token_provider_2() -> str: class TestWorkloadIdentity401Retry: + def test_workload_identity_401_retry_replays_non_seekable_file(self) -> None: + file_content = b"non-seekable workload identity content" + request_bodies: list[bytes] = [] + + def handler(request: httpx.Request) -> httpx.Response: + request_bodies.append(b"".join(cast(httpx.SyncByteStream, request.stream))) + return httpx.Response(401 if len(request_bodies) == 1 else 200) + + with OpenAI( + base_url=base_url, + workload_identity=workload_identity, + max_retries=0, + http_client=httpx.Client(transport=MockTransport(handler=handler)), + ) as client: + client._workload_identity_auth = cast(Any, MockWorkloadIdentityAuth()) + response = client.post( + "/upload", + cast_to=httpx.Response, + files={"file": ("input.txt", NonSeekableBytesIO(file_content))}, + options={"headers": {"Content-Type": "multipart/form-data"}}, + ) + + assert response.status_code == 200 + assert [body.count(file_content) for body in request_bodies] == [1, 1] + @pytest.mark.respx() def test_workload_identity_401_retry(self, respx_mock: MockRouter) -> None: provider_call_count = 0 @@ -2804,6 +3008,32 @@ def provider() -> str: class TestAsyncWorkloadIdentity401Retry: + async def test_workload_identity_401_retry_replays_non_seekable_file(self) -> None: + file_content = b"non-seekable workload identity content" + request_bodies: list[bytes] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + stream = cast(httpx.AsyncByteStream, request.stream) + request_bodies.append(b"".join([chunk async for chunk in stream])) + return httpx.Response(401 if len(request_bodies) == 1 else 200) + + async with AsyncOpenAI( + base_url=base_url, + workload_identity=workload_identity, + max_retries=0, + http_client=httpx.AsyncClient(transport=MockTransport(handler=handler)), + ) as client: + client._workload_identity_auth = cast(Any, MockWorkloadIdentityAuth()) + response = await client.post( + "/upload", + cast_to=httpx.Response, + files={"file": ("input.txt", NonSeekableBytesIO(file_content))}, + options={"headers": {"Content-Type": "multipart/form-data"}}, + ) + + assert response.status_code == 200 + assert [body.count(file_content) for body in request_bodies] == [1, 1] + @pytest.mark.respx() async def test_workload_identity_401_retry(self, respx_mock: MockRouter) -> None: provider_call_count = 0