Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 71 additions & 1 deletion src/openai/_base_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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 = {}
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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 = {}
Expand Down Expand Up @@ -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,
Expand Down
18 changes: 18 additions & 0 deletions src/openai/_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
67 changes: 66 additions & 1 deletion src/openai/_files.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Avoid spooling when retries are disabled

Because ReplayableFiles is constructed unconditionally before _request_with_retries reads max_retries, this copy now drains every non-seekable multipart upload into a temporary file even when the caller set max_retries=0 and is using normal API-key auth, where there is no retry path to replay. Large pipe/socket/generated uploads that previously streamed directly can now block until EOF or fill temporary storage; please only spool non-seekable streams when a retry or auth-retry can actually occur.

Useful? React with 👍 / 👎.

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)
Expand Down
Loading