Skip to content

Commit 2ab2b78

Browse files
stainless-app[bot]megamanics
stainless-app[bot]
authored and
megamanics
committed
feat(client): add retries_taken to raw response class (openai#1601)
1 parent 03dcbaa commit 2ab2b78

File tree

4 files changed

+154
-1
lines changed

4 files changed

+154
-1
lines changed

‎src/openai/_base_client.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1051,6 +1051,7 @@ def _request(
10511051
response=response,
10521052
stream=stream,
10531053
stream_cls=stream_cls,
1054+
retries_taken=options.get_max_retries(self.max_retries) - retries,
10541055
)
10551056

10561057
def _retry_request(
@@ -1092,6 +1093,7 @@ def _process_response(
10921093
response: httpx.Response,
10931094
stream: bool,
10941095
stream_cls: type[Stream[Any]] | type[AsyncStream[Any]] | None,
1096+
retries_taken: int = 0,
10951097
) -> ResponseT:
10961098
if response.request.headers.get(RAW_RESPONSE_HEADER) == "true":
10971099
return cast(
@@ -1103,6 +1105,7 @@ def _process_response(
11031105
stream=stream,
11041106
stream_cls=stream_cls,
11051107
options=options,
1108+
retries_taken=retries_taken,
11061109
),
11071110
)
11081111

@@ -1122,6 +1125,7 @@ def _process_response(
11221125
stream=stream,
11231126
stream_cls=stream_cls,
11241127
options=options,
1128+
retries_taken=retries_taken,
11251129
),
11261130
)
11271131

@@ -1135,6 +1139,7 @@ def _process_response(
11351139
stream=stream,
11361140
stream_cls=stream_cls,
11371141
options=options,
1142+
retries_taken=retries_taken,
11381143
)
11391144
if bool(response.request.headers.get(RAW_RESPONSE_HEADER)):
11401145
return cast(ResponseT, api_response)
@@ -1625,6 +1630,7 @@ async def _request(
16251630
response=response,
16261631
stream=stream,
16271632
stream_cls=stream_cls,
1633+
retries_taken=options.get_max_retries(self.max_retries) - retries,
16281634
)
16291635

16301636
async def _retry_request(
@@ -1664,6 +1670,7 @@ async def _process_response(
16641670
response: httpx.Response,
16651671
stream: bool,
16661672
stream_cls: type[Stream[Any]] | type[AsyncStream[Any]] | None,
1673+
retries_taken: int = 0,
16671674
) -> ResponseT:
16681675
if response.request.headers.get(RAW_RESPONSE_HEADER) == "true":
16691676
return cast(
@@ -1675,6 +1682,7 @@ async def _process_response(
16751682
stream=stream,
16761683
stream_cls=stream_cls,
16771684
options=options,
1685+
retries_taken=retries_taken,
16781686
),
16791687
)
16801688

@@ -1694,6 +1702,7 @@ async def _process_response(
16941702
stream=stream,
16951703
stream_cls=stream_cls,
16961704
options=options,
1705+
retries_taken=retries_taken,
16971706
),
16981707
)
16991708

@@ -1707,6 +1716,7 @@ async def _process_response(
17071716
stream=stream,
17081717
stream_cls=stream_cls,
17091718
options=options,
1719+
retries_taken=retries_taken,
17101720
)
17111721
if bool(response.request.headers.get(RAW_RESPONSE_HEADER)):
17121722
return cast(ResponseT, api_response)

‎src/openai/_legacy_response.py‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,18 @@
55
import logging
66
import datetime
77
import functools
8-
from typing import TYPE_CHECKING, Any, Union, Generic, TypeVar, Callable, Iterator, AsyncIterator, cast, overload
8+
from typing import (
9+
TYPE_CHECKING,
10+
Any,
11+
Union,
12+
Generic,
13+
TypeVar,
14+
Callable,
15+
Iterator,
16+
AsyncIterator,
17+
cast,
18+
overload,
19+
)
920
from typing_extensions import Awaitable, ParamSpec, override, deprecated, get_origin
1021

1122
import anyio
@@ -53,6 +64,9 @@ class LegacyAPIResponse(Generic[R]):
5364

5465
http_response: httpx.Response
5566

67+
retries_taken: int
68+
"""The number of retries made. If no retries happened this will be `0`"""
69+
5670
def __init__(
5771
self,
5872
*,
@@ -62,6 +76,7 @@ def __init__(
6276
stream: bool,
6377
stream_cls: type[Stream[Any]] | type[AsyncStream[Any]] | None,
6478
options: FinalRequestOptions,
79+
retries_taken: int = 0,
6580
) -> None:
6681
self._cast_to = cast_to
6782
self._client = client
@@ -70,6 +85,7 @@ def __init__(
7085
self._stream_cls = stream_cls
7186
self._options = options
7287
self.http_response = raw
88+
self.retries_taken = retries_taken
7389

7490
@property
7591
def request_id(self) -> str | None:

‎src/openai/_response.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,9 @@ class BaseAPIResponse(Generic[R]):
5555

5656
http_response: httpx.Response
5757

58+
retries_taken: int
59+
"""The number of retries made. If no retries happened this will be `0`"""
60+
5861
def __init__(
5962
self,
6063
*,
@@ -64,6 +67,7 @@ def __init__(
6467
stream: bool,
6568
stream_cls: type[Stream[Any]] | type[AsyncStream[Any]] | None,
6669
options: FinalRequestOptions,
70+
retries_taken: int = 0,
6771
) -> None:
6872
self._cast_to = cast_to
6973
self._client = client
@@ -72,6 +76,7 @@ def __init__(
7276
self._stream_cls = stream_cls
7377
self._options = options
7478
self.http_response = raw
79+
self.retries_taken = retries_taken
7580

7681
@property
7782
def headers(self) -> httpx.Headers:

‎tests/test_client.py‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -758,6 +758,65 @@ def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> Non
758758

759759
assert _get_open_connections(self.client) == 0
760760

761+
@pytest.mark.parametrize("failures_before_success", [0, 2, 4])
762+
@mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
763+
@pytest.mark.respx(base_url=base_url)
764+
def test_retries_taken(self, client: OpenAI, failures_before_success: int, respx_mock: MockRouter) -> None:
765+
client = client.with_options(max_retries=4)
766+
767+
nb_retries = 0
768+
769+
def retry_handler(_request: httpx.Request) -> httpx.Response:
770+
nonlocal nb_retries
771+
if nb_retries < failures_before_success:
772+
nb_retries += 1
773+
return httpx.Response(500)
774+
return httpx.Response(200)
775+
776+
respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
777+
778+
response = client.chat.completions.with_raw_response.create(
779+
messages=[
780+
{
781+
"content": "content",
782+
"role": "system",
783+
}
784+
],
785+
model="gpt-4-turbo",
786+
)
787+
788+
assert response.retries_taken == failures_before_success
789+
790+
@pytest.mark.parametrize("failures_before_success", [0, 2, 4])
791+
@mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
792+
@pytest.mark.respx(base_url=base_url)
793+
def test_retries_taken_new_response_class(
794+
self, client: OpenAI, failures_before_success: int, respx_mock: MockRouter
795+
) -> None:
796+
client = client.with_options(max_retries=4)
797+
798+
nb_retries = 0
799+
800+
def retry_handler(_request: httpx.Request) -> httpx.Response:
801+
nonlocal nb_retries
802+
if nb_retries < failures_before_success:
803+
nb_retries += 1
804+
return httpx.Response(500)
805+
return httpx.Response(200)
806+
807+
respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
808+
809+
with client.chat.completions.with_streaming_response.create(
810+
messages=[
811+
{
812+
"content": "content",
813+
"role": "system",
814+
}
815+
],
816+
model="gpt-4-turbo",
817+
) as response:
818+
assert response.retries_taken == failures_before_success
819+
761820

762821
class TestAsyncOpenAI:
763822
client = AsyncOpenAI(base_url=base_url, api_key=api_key, _strict_response_validation=True)
@@ -1488,3 +1547,66 @@ async def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter)
14881547
)
14891548

14901549
assert _get_open_connections(self.client) == 0
1550+
1551+
@pytest.mark.parametrize("failures_before_success", [0, 2, 4])
1552+
@mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1553+
@pytest.mark.respx(base_url=base_url)
1554+
@pytest.mark.asyncio
1555+
async def test_retries_taken(
1556+
self, async_client: AsyncOpenAI, failures_before_success: int, respx_mock: MockRouter
1557+
) -> None:
1558+
client = async_client.with_options(max_retries=4)
1559+
1560+
nb_retries = 0
1561+
1562+
def retry_handler(_request: httpx.Request) -> httpx.Response:
1563+
nonlocal nb_retries
1564+
if nb_retries < failures_before_success:
1565+
nb_retries += 1
1566+
return httpx.Response(500)
1567+
return httpx.Response(200)
1568+
1569+
respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
1570+
1571+
response = await client.chat.completions.with_raw_response.create(
1572+
messages=[
1573+
{
1574+
"content": "content",
1575+
"role": "system",
1576+
}
1577+
],
1578+
model="gpt-4-turbo",
1579+
)
1580+
1581+
assert response.retries_taken == failures_before_success
1582+
1583+
@pytest.mark.parametrize("failures_before_success", [0, 2, 4])
1584+
@mock.patch("openai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1585+
@pytest.mark.respx(base_url=base_url)
1586+
@pytest.mark.asyncio
1587+
async def test_retries_taken_new_response_class(
1588+
self, async_client: AsyncOpenAI, failures_before_success: int, respx_mock: MockRouter
1589+
) -> None:
1590+
client = async_client.with_options(max_retries=4)
1591+
1592+
nb_retries = 0
1593+
1594+
def retry_handler(_request: httpx.Request) -> httpx.Response:
1595+
nonlocal nb_retries
1596+
if nb_retries < failures_before_success:
1597+
nb_retries += 1
1598+
return httpx.Response(500)
1599+
return httpx.Response(200)
1600+
1601+
respx_mock.post("/chat/completions").mock(side_effect=retry_handler)
1602+
1603+
async with client.chat.completions.with_streaming_response.create(
1604+
messages=[
1605+
{
1606+
"content": "content",
1607+
"role": "system",
1608+
}
1609+
],
1610+
model="gpt-4-turbo",
1611+
) as response:
1612+
assert response.retries_taken == failures_before_success

0 commit comments

Comments
 (0)