google-genai 1.12.1__py3-none-any.whl → 1.14.0__py3-none-any.whl

google/genai/_api_client.py

@@ -58,7 +58,9 @@ from .types import HttpOptionsOrDict
 
 logger = logging.getLogger('google_genai._api_client')
 CHUNK_SIZE = 8 * 1024 * 1024  # 8 MB chunk size
-
+MAX_RETRY_COUNT = 3
+INITIAL_RETRY_DELAY = 1  # second
+DELAY_MULTIPLIER = 2
 
 def _append_library_version_headers(headers: dict[str, str]) -> None:
   """Appends the telemetry header to the headers dict."""
@@ -283,6 +285,18 @@ class SyncHttpxClient(httpx.Client):
     kwargs.setdefault('follow_redirects', True)
     super().__init__(**kwargs)
 
+  def __del__(self) -> None:
+    """Closes the httpx client."""
+    try:
+      if self.is_closed:
+        return
+    except Exception:
+      pass
+    try:
+      self.close()
+    except Exception:
+      pass
+
 
 class AsyncHttpxClient(httpx.AsyncClient):
   """Async httpx client."""
@@ -292,6 +306,17 @@ class AsyncHttpxClient(httpx.AsyncClient):
     kwargs.setdefault('follow_redirects', True)
     super().__init__(**kwargs)
 
+  def __del__(self) -> None:
+    try:
+      if self.is_closed:
+        return
+    except Exception:
+      pass
+    try:
+      asyncio.get_running_loop().create_task(self.aclose())
+    except Exception:
+      pass
+
 
 class BaseApiClient:
   """Client for calling HTTP APIs sending and receiving JSON."""
@@ -865,15 +890,23 @@ class BaseApiClient:
           'Content-Length': str(chunk_size),
       }
       _populate_server_timeout_header(upload_headers, timeout_in_seconds)
-      response = self._httpx_client.request(
-          method='POST',
-          url=upload_url,
-          headers=upload_headers,
-          content=file_chunk,
-          timeout=timeout_in_seconds,
-      )
+      retry_count = 0
+      while retry_count < MAX_RETRY_COUNT:
+        response = self._httpx_client.request(
+            method='POST',
+            url=upload_url,
+            headers=upload_headers,
+            content=file_chunk,
+            timeout=timeout_in_seconds,
+        )
+        if response.headers.get('x-goog-upload-status'):
+          break
+        delay_seconds = INITIAL_RETRY_DELAY * (DELAY_MULTIPLIER**retry_count)
+        retry_count += 1
+        time.sleep(delay_seconds)
+
       offset += chunk_size
-      if response.headers['x-goog-upload-status'] != 'active':
+      if response.headers.get('x-goog-upload-status') != 'active':
         break  # upload is complete or it has been interrupted.
       if upload_size <= offset:  # Status is not finalized.
         raise ValueError(
@@ -881,7 +914,7 @@ class BaseApiClient:
             f' finalized.'
         )
 
-    if response.headers['x-goog-upload-status'] != 'final':
+    if response.headers.get('x-goog-upload-status') != 'final':
       raise ValueError(
           'Failed to upload file: Upload status is not finalized.'
       )
@@ -1013,13 +1046,22 @@ class BaseApiClient:
           'Content-Length': str(chunk_size),
       }
       _populate_server_timeout_header(upload_headers, timeout_in_seconds)
-      response = await self._async_httpx_client.request(
-          method='POST',
-          url=upload_url,
-          content=file_chunk,
-          headers=upload_headers,
-          timeout=timeout_in_seconds,
-      )
+
+      retry_count = 0
+      while retry_count < MAX_RETRY_COUNT:
+        response = await self._async_httpx_client.request(
+            method='POST',
+            url=upload_url,
+            content=file_chunk,
+            headers=upload_headers,
+            timeout=timeout_in_seconds,
+        )
+        if response.headers.get('x-goog-upload-status'):
+          break
+        delay_seconds = INITIAL_RETRY_DELAY * (DELAY_MULTIPLIER**retry_count)
+        retry_count += 1
+        time.sleep(delay_seconds)
+
       offset += chunk_size
       if response.headers.get('x-goog-upload-status') != 'active':
         break  # upload is complete or it has been interrupted.

google/genai/_base_url.py

@@ -0,0 +1,50 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+from typing import Optional
+
+from .types import HttpOptions
+
+_default_base_gemini_url = None
+_default_base_vertex_url = None
+
+
+def set_default_base_urls(
+    gemini_url: Optional[str], vertex_url: Optional[str]
+) -> None:
+  """Overrides the base URLs for the Gemini API and Vertex AI API."""
+  global _default_base_gemini_url, _default_base_vertex_url
+  _default_base_gemini_url = gemini_url
+  _default_base_vertex_url = vertex_url
+
+
+def get_base_url(
+    vertexai: bool,
+    http_options: Optional[HttpOptions] = None,
+) -> Optional[str]:
+  """Returns the default base URL based on the following priority.
+
+  1. Base URLs set via HttpOptions.
+  2. Base URLs set via the latest call to setDefaultBaseUrls.
+  3. Base URLs set via environment variables.
+  """
+  if http_options and http_options.base_url:
+    return http_options.base_url
+
+  if vertexai:
+    return _default_base_vertex_url or os.getenv('GOOGLE_VERTEX_BASE_URL')
+  else:
+    return _default_base_gemini_url or os.getenv('GOOGLE_GEMINI_BASE_URL')