Skip to content

Commit c0c89b7

Browse files
feat(client): support passing chunk size for binary responses (#747)
1 parent 1afd138 commit c0c89b7

File tree

2 files changed

+27
-6
lines changed

2 files changed

+27
-6
lines changed

src/openai/_base_client.py

+14-4
Original file line numberDiff line numberDiff line change
@@ -1727,9 +1727,14 @@ def iter_raw(self, chunk_size: Optional[int] = None) -> Iterator[bytes]:
17271727
return self.response.iter_raw(chunk_size)
17281728

17291729
@override
1730-
def stream_to_file(self, file: str | os.PathLike[str]) -> None:
1730+
def stream_to_file(
1731+
self,
1732+
file: str | os.PathLike[str],
1733+
*,
1734+
chunk_size: int | None = None,
1735+
) -> None:
17311736
with open(file, mode="wb") as f:
1732-
for data in self.response.iter_bytes():
1737+
for data in self.response.iter_bytes(chunk_size):
17331738
f.write(data)
17341739

17351740
@override
@@ -1757,10 +1762,15 @@ async def aiter_raw(self, chunk_size: Optional[int] = None) -> AsyncIterator[byt
17571762
return self.response.aiter_raw(chunk_size)
17581763

17591764
@override
1760-
async def astream_to_file(self, file: str | os.PathLike[str]) -> None:
1765+
async def astream_to_file(
1766+
self,
1767+
file: str | os.PathLike[str],
1768+
*,
1769+
chunk_size: int | None = None,
1770+
) -> None:
17611771
path = anyio.Path(file)
17621772
async with await path.open(mode="wb") as f:
1763-
async for data in self.response.aiter_bytes():
1773+
async for data in self.response.aiter_bytes(chunk_size):
17641774
await f.write(data)
17651775

17661776
@override

src/openai/_types.py

+13-2
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,12 @@ def iter_raw(self, chunk_size: Optional[int] = None) -> Iterator[bytes]:
123123
pass
124124

125125
@abstractmethod
126-
def stream_to_file(self, file: str | PathLike[str]) -> None:
126+
def stream_to_file(
127+
self,
128+
file: str | PathLike[str],
129+
*,
130+
chunk_size: int | None = None,
131+
) -> None:
127132
"""
128133
Stream the output to the given file.
129134
"""
@@ -172,7 +177,13 @@ async def aiter_raw(self, chunk_size: Optional[int] = None) -> AsyncIterator[byt
172177
"""
173178
pass
174179

175-
async def astream_to_file(self, file: str | PathLike[str]) -> None:
180+
@abstractmethod
181+
async def astream_to_file(
182+
self,
183+
file: str | PathLike[str],
184+
*,
185+
chunk_size: int | None = None,
186+
) -> None:
176187
"""
177188
Stream the output to the given file.
178189
"""

0 commit comments

Comments
 (0)