feat: Support retry configuration at request level

PiperOrigin-RevId: 787236850

Author: yinghsienwu

google/genai/_api_client.py

@@ -965,8 +965,21 @@ def _request_once(
   def _request(
       self,
       http_request: HttpRequest,
+      http_options: Optional[HttpOptionsOrDict] = None,
       stream: bool = False,
   ) -> HttpResponse:
+    if http_options:
+      parameter_model = (
+          HttpOptions(**http_options)
+          if isinstance(http_options, dict)
+          else http_options
+      )
+      # Support per request retry options.
+      if parameter_model.retry_options:
+        retry_kwargs = _retry_args(parameter_model.retry_options)
+        retry = tenacity.Retrying(**retry_kwargs)
+        return retry(self._request_once, http_request, stream)  # type: ignore[no-any-return]
+
     return self._retry(self._request_once, http_request, stream)  # type: ignore[no-any-return]
 
   async def _async_request_once(
@@ -1111,8 +1124,20 @@ async def _async_request_once(
   async def _async_request(
       self,
       http_request: HttpRequest,
+      http_options: Optional[HttpOptionsOrDict] = None,
       stream: bool = False,
   ) -> HttpResponse:
+    if http_options:
+      parameter_model = (
+          HttpOptions(**http_options)
+          if isinstance(http_options, dict)
+          else http_options
+      )
+      # Support per request retry options.
+      if parameter_model.retry_options:
+        retry_kwargs = _retry_args(parameter_model.retry_options)
+        retry = tenacity.AsyncRetrying(**retry_kwargs)
+        return await retry(self._async_request_once, http_request, stream)  # type: ignore[no-any-return]
     return await self._async_retry(  # type: ignore[no-any-return]
         self._async_request_once, http_request, stream
     )
@@ -1134,7 +1159,7 @@ def request(
     http_request = self._build_request(
         http_method, path, request_dict, http_options
     )
-    response = self._request(http_request, stream=False)
+    response = self._request(http_request, http_options, stream=False)
     response_body = (
         response.response_stream[0] if response.response_stream else ''
     )
@@ -1151,7 +1176,7 @@ def request_streamed(
         http_method, path, request_dict, http_options
     )
 
-    session_response = self._request(http_request, stream=True)
+    session_response = self._request(http_request, http_options, stream=True)
     for chunk in session_response.segments():
       yield SdkHttpResponse(
           headers=session_response.headers, body=json.dumps(chunk)
@@ -1168,7 +1193,9 @@ async def async_request(
         http_method, path, request_dict, http_options
     )
 
-    result = await self._async_request(http_request=http_request, stream=False)
+    result = await self._async_request(
+        http_request=http_request, http_options=http_options, stream=False
+    )
     response_body = result.response_stream[0] if result.response_stream else ''
     return SdkHttpResponse(headers=result.headers, body=response_body)
 

google/genai/_replay_api_client.py

@@ -479,13 +479,14 @@ def _verify_response(self, response_model: BaseModel) -> None:
   def _request(
       self,
       http_request: HttpRequest,
+      http_options: Optional[HttpOptionsOrDict] = None,
       stream: bool = False,
   ) -> HttpResponse:
     self._initialize_replay_session_if_not_loaded()
     if self._should_call_api():
       _debug_print('api mode request: %s' % http_request)
       try:
-        result = super()._request(http_request, stream)
+        result = super()._request(http_request, http_options, stream)
       except errors.APIError as e:
         self._record_interaction(http_request, e)
         raise e
@@ -507,13 +508,16 @@ def _request(
   async def _async_request(
       self,
       http_request: HttpRequest,
+      http_options: Optional[HttpOptionsOrDict] = None,
       stream: bool = False,
   ) -> HttpResponse:
     self._initialize_replay_session_if_not_loaded()
     if self._should_call_api():
       _debug_print('api mode request: %s' % http_request)
       try:
-        result = await super()._async_request(http_request, stream)
+        result = await super()._async_request(
+            http_request, http_options, stream
+        )
       except errors.APIError as e:
         self._record_interaction(http_request, e)
         raise e

google/genai/tests/client/test_retries.py

@@ -264,6 +264,35 @@ def test_retries_failed_request_retries_successfully():
     assert response.headers['status-code'] == '200'
 
 
+def test_retries_failed_request_retries_successfully_at_request_level():
+  mock_transport = mock.Mock(spec=httpx.BaseTransport)
+  mock_transport.handle_request.side_effect = (
+      _httpx_response(429),
+      _httpx_response(200),
+  )
+
+  client = api_client.BaseApiClient(
+      vertexai=True,
+      project='test_project',
+      location='global',
+      http_options=_transport_options(
+          transport=mock_transport,
+      ),
+  )
+
+  with _patch_auth_default():
+    response = client.request(
+        http_method='GET',
+        path='path',
+        request_dict={},
+        http_options=types.HttpOptions(
+            retry_options=_RETRY_OPTIONS
+        ),  # At request level.
+    )
+    mock_transport.handle_request.assert_called()
+    assert response.headers['status-code'] == '200'
+
+
 def test_retries_failed_request_retries_unsuccessfully():
   mock_transport = mock.Mock(spec=httpx.BaseTransport)
   mock_transport.handle_request.side_effect = (
@@ -290,6 +319,36 @@ def test_retries_failed_request_retries_unsuccessfully():
     mock_transport.handle_request.assert_called()
 
 
+def test_retries_failed_request_retries_unsuccessfully_at_request_level():
+  mock_transport = mock.Mock(spec=httpx.BaseTransport)
+  mock_transport.handle_request.side_effect = (
+      _httpx_response(429),
+      _httpx_response(504),
+  )
+
+  client = api_client.BaseApiClient(
+      vertexai=True,
+      project='test_project',
+      location='global',
+      http_options=_transport_options(
+          transport=mock_transport,
+      ),
+  )
+
+  with _patch_auth_default():
+    try:
+      client.request(
+          http_method='GET',
+          path='path',
+          request_dict={},
+          http_options={'retry_options': _RETRY_OPTIONS},  # At request level.
+      )
+      assert False, 'Expected APIError to be raised.'
+    except errors.APIError as e:
+      assert e.code == 504
+    mock_transport.handle_request.assert_called()
+
+
 # Async httpx
 
 
@@ -401,6 +460,40 @@ async def run():
   asyncio.run(run())
 
 
+def test_async_retries_failed_request_retries_successfully_at_request_level():
+  api_client.has_aiohttp = False
+
+  async def run():
+    mock_transport = mock.Mock(spec=httpx.AsyncBaseTransport)
+    mock_transport.handle_async_request.side_effect = (
+        _httpx_response(429),
+        _httpx_response(200),
+    )
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+        http_options=_transport_options(
+            async_transport=mock_transport,
+        ),
+    )
+
+    with _patch_auth_default():
+      response = await client.async_request(
+          http_method='GET',
+          path='path',
+          request_dict={},
+          http_options=types.HttpOptions(
+              retry_options=_RETRY_OPTIONS
+          ),  # At request level.
+      )
+      mock_transport.handle_async_request.assert_called()
+      assert response.headers['status-code'] == '200'
+
+  asyncio.run(run())
+
+
 def test_async_retries_failed_request_retries_unsuccessfully():
   api_client.has_aiohttp = False
 
@@ -434,6 +527,41 @@ async def run():
   asyncio.run(run())
 
 
+def test_async_retries_failed_request_retries_unsuccessfully_at_request_level():
+  api_client.has_aiohttp = False
+
+  async def run():
+    mock_transport = mock.Mock(spec=httpx.AsyncBaseTransport)
+    mock_transport.handle_async_request.side_effect = (
+        _httpx_response(429),
+        _httpx_response(504),
+    )
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+        http_options=_transport_options(
+            async_transport=mock_transport,
+        ),
+    )
+
+    with _patch_auth_default():
+      try:
+        await client.async_request(
+            http_method='GET',
+            path='path',
+            request_dict={},
+            http_options=types.HttpOptions(retry_options=_RETRY_OPTIONS),
+        )
+        assert False, 'Expected APIError to be raised.'
+      except errors.APIError as e:
+        assert e.code == 504
+      mock_transport.handle_async_request.assert_called()
+
+  asyncio.run(run())
+
+
 # Async aiohttp
 
 
@@ -553,6 +681,39 @@ async def run():
   asyncio.run(run())
 
 
+@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
+def test_aiohttp_retries_failed_request_retries_successfully_at_request_level(
+    mock_request,
+):
+  api_client.has_aiohttp = True
+
+  async def run():
+    mock_request.side_effect = (
+        _aiohttp_async_response(429),
+        _aiohttp_async_response(200),
+    )
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+    )
+
+    with _patch_auth_default():
+      response = await client.async_request(
+          http_method='GET',
+          path='path',
+          request_dict={},
+          http_options=types.HttpOptions(
+              retry_options=_RETRY_OPTIONS
+          ),  # At request level.
+      )
+      mock_request.assert_called()
+      assert response.headers['status-code'] == '200'
+
+  asyncio.run(run())
+
+
 @mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
 def test_aiohttp_retries_failed_request_retries_unsuccessfully(mock_request):
   api_client.has_aiohttp = True
@@ -583,3 +744,37 @@ async def run():
       mock_request.assert_called()
 
   asyncio.run(run())
+
+
+@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
+def test_aiohttp_retries_failed_request_retries_unsuccessfully_at_request_level(
+    mock_request,
+):
+  api_client.has_aiohttp = True
+
+  async def run():
+    mock_request.side_effect = (
+        _aiohttp_async_response(429),
+        _aiohttp_async_response(504),
+    )
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+    )
+
+    with _patch_auth_default():
+      try:
+        await client.async_request(
+            http_method='GET',
+            path='path',
+            request_dict={},
+            http_options={'retry_options': _RETRY_OPTIONS},  # At request level.
+        )
+        assert False, 'Expected APIError to be raised.'
+      except errors.APIError as e:
+        assert e.code == 504
+      mock_request.assert_called()
+
+  asyncio.run(run())