Skip to content

Commit 0437e72

Browse files
stainless-app[bot]stainless-bot
authored andcommitted
feat(client): add retry_count to raw response class (#197)
1 parent 91dd974 commit 0437e72

File tree

3 files changed

+58
-0
lines changed

3 files changed

+58
-0
lines changed

src/prompt_foundry_python_sdk/_base_client.py

+8
Original file line numberDiff line numberDiff line change
@@ -1049,6 +1049,7 @@ def _request(
10491049
response=response,
10501050
stream=stream,
10511051
stream_cls=stream_cls,
1052+
retries_taken=options.get_max_retries(self.max_retries) - retries,
10521053
)
10531054

10541055
def _retry_request(
@@ -1090,6 +1091,7 @@ def _process_response(
10901091
response: httpx.Response,
10911092
stream: bool,
10921093
stream_cls: type[Stream[Any]] | type[AsyncStream[Any]] | None,
1094+
retries_taken: int = 0,
10931095
) -> ResponseT:
10941096
origin = get_origin(cast_to) or cast_to
10951097

@@ -1107,6 +1109,7 @@ def _process_response(
11071109
stream=stream,
11081110
stream_cls=stream_cls,
11091111
options=options,
1112+
retries_taken=retries_taken,
11101113
),
11111114
)
11121115

@@ -1120,6 +1123,7 @@ def _process_response(
11201123
stream=stream,
11211124
stream_cls=stream_cls,
11221125
options=options,
1126+
retries_taken=retries_taken,
11231127
)
11241128
if bool(response.request.headers.get(RAW_RESPONSE_HEADER)):
11251129
return cast(ResponseT, api_response)
@@ -1610,6 +1614,7 @@ async def _request(
16101614
response=response,
16111615
stream=stream,
16121616
stream_cls=stream_cls,
1617+
retries_taken=options.get_max_retries(self.max_retries) - retries,
16131618
)
16141619

16151620
async def _retry_request(
@@ -1649,6 +1654,7 @@ async def _process_response(
16491654
response: httpx.Response,
16501655
stream: bool,
16511656
stream_cls: type[Stream[Any]] | type[AsyncStream[Any]] | None,
1657+
retries_taken: int = 0,
16521658
) -> ResponseT:
16531659
origin = get_origin(cast_to) or cast_to
16541660

@@ -1666,6 +1672,7 @@ async def _process_response(
16661672
stream=stream,
16671673
stream_cls=stream_cls,
16681674
options=options,
1675+
retries_taken=retries_taken,
16691676
),
16701677
)
16711678

@@ -1679,6 +1686,7 @@ async def _process_response(
16791686
stream=stream,
16801687
stream_cls=stream_cls,
16811688
options=options,
1689+
retries_taken=retries_taken,
16821690
)
16831691
if bool(response.request.headers.get(RAW_RESPONSE_HEADER)):
16841692
return cast(ResponseT, api_response)

src/prompt_foundry_python_sdk/_response.py

+5
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

+45
Original file line numberDiff line numberDiff line change
@@ -746,6 +746,27 @@ def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> Non
746746

747747
assert _get_open_connections(self.client) == 0
748748

749+
@pytest.mark.parametrize("failures_before_success", [0, 2, 4])
750+
@mock.patch("prompt_foundry_python_sdk._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
751+
@pytest.mark.respx(base_url=base_url)
752+
def test_retries_taken(self, client: PromptFoundry, failures_before_success: int, respx_mock: MockRouter) -> None:
753+
client = client.with_options(max_retries=4)
754+
755+
nb_retries = 0
756+
757+
def retry_handler(_request: httpx.Request) -> httpx.Response:
758+
nonlocal nb_retries
759+
if nb_retries < failures_before_success:
760+
nb_retries += 1
761+
return httpx.Response(500)
762+
return httpx.Response(200)
763+
764+
respx_mock.post("/sdk/v1/prompts/1212121").mock(side_effect=retry_handler)
765+
766+
response = client.prompts.with_raw_response.get_parameters(id="1212121")
767+
768+
assert response.retries_taken == failures_before_success
769+
749770

750771
class TestAsyncPromptFoundry:
751772
client = AsyncPromptFoundry(base_url=base_url, api_key=api_key, _strict_response_validation=True)
@@ -1442,3 +1463,27 @@ async def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter)
14421463
)
14431464

14441465
assert _get_open_connections(self.client) == 0
1466+
1467+
@pytest.mark.parametrize("failures_before_success", [0, 2, 4])
1468+
@mock.patch("prompt_foundry_python_sdk._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
1469+
@pytest.mark.respx(base_url=base_url)
1470+
@pytest.mark.asyncio
1471+
async def test_retries_taken(
1472+
self, async_client: AsyncPromptFoundry, failures_before_success: int, respx_mock: MockRouter
1473+
) -> None:
1474+
client = async_client.with_options(max_retries=4)
1475+
1476+
nb_retries = 0
1477+
1478+
def retry_handler(_request: httpx.Request) -> httpx.Response:
1479+
nonlocal nb_retries
1480+
if nb_retries < failures_before_success:
1481+
nb_retries += 1
1482+
return httpx.Response(500)
1483+
return httpx.Response(200)
1484+
1485+
respx_mock.post("/sdk/v1/prompts/1212121").mock(side_effect=retry_handler)
1486+
1487+
response = await client.prompts.with_raw_response.get_parameters(id="1212121")
1488+
1489+
assert response.retries_taken == failures_before_success

0 commit comments

Comments
 (0)