feat: GenAI SDK client - Support agent engine sandbox http request in genai sdk

PiperOrigin-RevId: 816865842

Author: google-genai-bot

google/genai/_api_client.py

@@ -1121,6 +1121,73 @@ def _request_once(
           response.headers, response if stream else [response.text]
       )
 
+  def _request_once_sandbox(
+      self,
+      http_request: HttpRequest,
+  ) -> HttpResponse:
+    # Only used by Orcas sandbox, which doesn't send request to vertex dataplane
+    # but directly to the sandbox endpoint provisioned in tenant project.
+    # It relays on different auth token and doesn't need to fetch ADC.
+    data: Optional[Union[str, bytes]] = None
+    if http_request.data:
+      if not isinstance(http_request.data, bytes):
+        data = json.dumps(http_request.data) if http_request.data else None
+      else:
+        data = http_request.data
+
+    response = self._httpx_client.request(
+        method=http_request.method,
+        url=http_request.url,
+        headers=http_request.headers,
+        content=data,
+        timeout=http_request.timeout,
+    )
+    errors.APIError.raise_for_response(response)
+    return HttpResponse(
+        response.headers, [response.text]
+    )
+
+  def _request_sandbox(
+      self,
+      http_request: HttpRequest,
+      http_options: Optional[HttpOptionsOrDict] = None,
+  ) -> 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_sandbox, http_request)  # type: ignore[no-any-return]
+
+    return self._retry(self._request_once_sandbox, http_request)  # type: ignore[no-any-return]
+
+  def request_sandbox(
+      self,
+      http_method: str,
+      url: str,
+      request_dict: dict[str, object],
+      http_options: HttpOptions,
+  ) -> SdkHttpResponse:
+    if not http_options.headers:
+      raise ValueError('Request headers must be set.')
+
+    http_request = HttpRequest(
+        method=http_method,
+        url=url,
+        headers=http_options.headers,
+        data=request_dict,
+    )
+    response = self._request_sandbox(http_request, http_options)
+    response_body = (
+        response.response_stream[0] if response.response_stream else ''
+    )
+    return SdkHttpResponse(headers=response.headers, body=response_body)
+
   def _request(
       self,
       http_request: HttpRequest,

google/genai/_replay_api_client.py

@@ -34,7 +34,9 @@
 from ._api_client import HttpRequest
 from ._api_client import HttpResponse
 from ._common import BaseModel
-from .types import HttpOptions, HttpOptionsOrDict
+from .types import HttpOptions
+from .types import HttpOptionsOrDict
+from .types import HttpResponse as SdkHttpResponse
 
 
 def to_snake_case(name: str) -> str:
@@ -679,3 +681,25 @@ async def async_download_file(
       return result
     else:
       return self._build_response_from_replay(request).byte_stream[0]
+
+  def _request_sandbox(
+      self,
+      http_request: HttpRequest,
+      http_options: Optional[HttpOptionsOrDict] = None,
+  ) -> 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_sandbox(
+            http_request, http_options
+        )
+      except errors.APIError as e:
+        self._record_interaction(http_request, e)
+        raise e
+      if result:
+        self._record_interaction(http_request, result)
+      _debug_print('api mode result: %s' % result)
+      return result
+    else:
+      return self._build_response_from_replay(http_request)

google/genai/_test_api_client.py

@@ -17,7 +17,7 @@
 import time
 from unittest.mock import MagicMock, patch
 import pytest
-from .api_client import BaseApiClient
+from .api_client import BaseApiClient, HttpResponse
 
 
 @patch('genai.api_client.BaseApiClient._build_request')
@@ -105,7 +105,7 @@ async def delayed_response(http_request, stream):
 async def test_async_request_streamed_non_blocking(
     mock_async_request, mock_build_request
 ):
-  api_client = ApiClient(api_key='test_api_key')
+  api_client = BaseApiClient(api_key='test_api_key')
   http_method = 'GET'
   path = 'test/path'
   request_dict = {'key': 'value'}
@@ -147,3 +147,28 @@ async def delayed_response(http_request, stream):
   )
   assert chunks == ['{"chunk": 1}', '{"chunk": 2}', '{"chunk": 3}']
   assert end_time - start_time > 0.3
+
+
+@patch('google.genai._api_client.BaseApiClient._request_sandbox')
+def test_request_sandbox(mock_request_sandbox):
+  api_client = ApiClient(api_key='test_api_key')
+  http_method = 'POST'
+  url = 'http://test/url'
+  request_dict = {'key': 'value'}
+
+  mock_request_sandbox.return_value = HttpResponse(
+      headers={'Content-Type': 'application/json'},
+      response_stream=['{"response_key": "response_value"}']
+  )
+
+  response = api_client.request_sandbox(http_method, url, request_dict)
+
+  mock_request_sandbox.assert_called_once()
+  assert response['headers'] == {'Content-Type': 'application/json'}
+  assert response['body'] == '{"response_key": "response_value"}'
+
+  called_with_http_request = mock_request_sandbox.call_args[0][0]
+  assert called_with_http_request.method == http_method
+  assert called_with_http_request.url == url
+  assert called_with_http_request.data == request_dict
+