chore: process proxy and base url settings for file uploads

PiperOrigin-RevId: 882260641

Author: speedstorm1

google/genai/_api_client.py

@@ -1538,6 +1538,21 @@ def _upload_fd(
           The HttpResponse object from the finalize request.
     """
     offset = 0
+    http_options = http_options if http_options else self._http_options
+    base_url = (
+        http_options.get('base_url')
+        if isinstance(http_options, dict)
+        else getattr(http_options, 'base_url', None)
+    )
+    if base_url:
+      parsed_base = urlparse(base_url)
+      parsed_upload = urlparse(upload_url)
+      upload_url = urlunparse(
+          parsed_upload._replace(
+              scheme=parsed_base.scheme, netloc=parsed_base.netloc
+          )
+      )
+
     # Upload the file in chunks
     while True:
       file_chunk = file.read(CHUNK_SIZE)
@@ -1548,7 +1563,6 @@ def _upload_fd(
       # If last chunk, finalize the upload.
       if chunk_size + offset >= upload_size:
         upload_command += ', finalize'
-      http_options = http_options if http_options else self._http_options
       timeout = (
           http_options.get('timeout')
           if isinstance(http_options, dict)
@@ -1562,11 +1576,17 @@ def _upload_fd(
             else self._http_options.timeout
         )
       timeout_in_seconds = get_timeout_in_seconds(timeout)
-      upload_headers = {
+      user_headers = (
+          http_options.get('headers', {})
+          if isinstance(http_options, dict)
+          else (getattr(http_options, 'headers', {}) or {})
+      )
+      upload_headers = dict(user_headers) if user_headers else {}
+      upload_headers.update({
           'X-Goog-Upload-Command': upload_command,
           'X-Goog-Upload-Offset': str(offset),
           'Content-Length': str(chunk_size),
-      }
+      })
       populate_server_timeout_header(upload_headers, timeout_in_seconds)
       retry_count = 0
       while retry_count < MAX_RETRY_COUNT:
@@ -1689,6 +1709,21 @@ async def _async_upload_fd(
           The HttpResponse object from the finalized request.
     """
     offset = 0
+    http_options = http_options if http_options else self._http_options
+    base_url = (
+        http_options.get('base_url')
+        if isinstance(http_options, dict)
+        else getattr(http_options, 'base_url', None)
+    )
+    if base_url:
+      parsed_base = urlparse(base_url)
+      parsed_upload = urlparse(upload_url)
+      upload_url = urlunparse(
+          parsed_upload._replace(
+              scheme=parsed_base.scheme, netloc=parsed_base.netloc
+          )
+      )
+
     # Upload the file in chunks
     if self._use_aiohttp():  # pylint: disable=g-import-not-at-top
       self._aiohttp_session = await self._get_aiohttp_session()
@@ -1704,7 +1739,6 @@ async def _async_upload_fd(
         # If last chunk, finalize the upload.
         if chunk_size + offset >= upload_size:
           upload_command += ', finalize'
-        http_options = http_options if http_options else self._http_options
         timeout = (
             http_options.get('timeout')
              if isinstance(http_options, dict)
@@ -1718,11 +1752,17 @@ async def _async_upload_fd(
               else self._http_options.timeout
           )
         timeout_in_seconds = get_timeout_in_seconds(timeout)
-        upload_headers = {
+        user_headers = (
+            http_options.get('headers', {})
+            if isinstance(http_options, dict)
+            else (getattr(http_options, 'headers', {}) or {})
+        )
+        upload_headers = dict(user_headers) if user_headers else {}
+        upload_headers.update({
             'X-Goog-Upload-Command': upload_command,
-             'X-Goog-Upload-Offset': str(offset),
+            'X-Goog-Upload-Offset': str(offset),
             'Content-Length': str(chunk_size),
-        }
+        })
         populate_server_timeout_header(upload_headers, timeout_in_seconds)
 
         retry_count = 0
@@ -1780,7 +1820,6 @@ async def _async_upload_fd(
         # If last chunk, finalize the upload.
         if chunk_size + offset >= upload_size:
           upload_command += ', finalize'
-        http_options = http_options if http_options else self._http_options
         timeout = (
             http_options.get('timeout')
             if isinstance(http_options, dict)
@@ -1794,11 +1833,17 @@ async def _async_upload_fd(
               else self._http_options.timeout
           )
         timeout_in_seconds = get_timeout_in_seconds(timeout)
-        upload_headers = {
+        user_headers = (
+            http_options.get('headers', {})
+            if isinstance(http_options, dict)
+            else (getattr(http_options, 'headers', {}) or {})
+        )
+        upload_headers = dict(user_headers) if user_headers else {}
+        upload_headers.update({
             'X-Goog-Upload-Command': upload_command,
             'X-Goog-Upload-Offset': str(offset),
             'Content-Length': str(chunk_size),
-        }
+        })
         populate_server_timeout_header(upload_headers, timeout_in_seconds)
 
         retry_count = 0

google/genai/tests/client/test_upload_errors.py

@@ -45,6 +45,34 @@ def _httpx_response(code: int, headers: dict = None, content: bytes = b''):
 def client():
   return api_client.BaseApiClient(vertexai=False, api_key='test_api_key')
 
+
+def test_upload_url_rewrite(client: api_client.BaseApiClient):
+  mock_httpx_client = mock.MagicMock(spec=httpx.Client)
+  mock_httpx_client.request.side_effect = [
+      _httpx_response(
+          200, headers={"X-Goog-Upload-Status": "final"}
+      ),  # Upload request succeeding
+  ]
+  client._httpx_client = mock_httpx_client
+
+  http_options = types.HttpOptions(base_url="https://my-proxy.company.com")
+
+  with io.BytesIO(b"test") as f:
+    client._upload_fd(
+        f,
+        "https://generativelanguage.googleapis.com/upload/v1beta/files?uploadType=resumable",
+        4,
+        http_options=http_options,
+    )
+
+  assert mock_httpx_client.request.call_count == 1
+  call_args = mock_httpx_client.request.call_args[1]
+  assert (
+      call_args["url"]
+      == "https://my-proxy.company.com/upload/v1beta/files?uploadType=resumable"
+  )
+
+
 def test_upload_fd_error(client: api_client.BaseApiClient):
   error_content = json.dumps({
       "error": {
@@ -71,6 +99,39 @@ def test_upload_fd_error(client: api_client.BaseApiClient):
 
   assert mock_httpx_client.request.call_count == 2
 
+
+@pytest.mark.asyncio
+async def test_async_upload_url_rewrite_httpx(client: api_client.BaseApiClient):
+  mock_async_httpx_client = mock.MagicMock(spec=httpx.AsyncClient)
+  mock_async_httpx_client.request = mock.AsyncMock(
+      side_effect=[
+          _httpx_response(
+              200, headers={"X-Goog-Upload-Status": "final"}
+          ),  # Upload request
+      ]
+  )
+  client._async_httpx_client = mock_async_httpx_client
+
+  http_options = types.HttpOptions(base_url="https://my-proxy.company.com")
+
+  with mock.patch.object(
+      client, "_use_aiohttp", return_value=False
+  ), io.BytesIO(b"test") as f:
+    await client._async_upload_fd(
+        f,
+        "https://generativelanguage.googleapis.com/upload/v1beta/files?uploadType=resumable",
+        4,
+        http_options=http_options,
+    )
+
+  assert mock_async_httpx_client.request.call_count == 1
+  call_args = mock_async_httpx_client.request.call_args[1]
+  assert (
+      call_args["url"]
+      == "https://my-proxy.company.com/upload/v1beta/files?uploadType=resumable"
+  )
+
+
 @pytest.mark.asyncio
 async def test_async_upload_fd_error_httpx(client: api_client.BaseApiClient):
   error_content = json.dumps({