google-genai 1.6.0__py3-none-any.whl → 1.8.0__py3-none-any.whl

google/genai/_api_client.py

@@ -68,26 +68,30 @@ def _append_library_version_headers(headers: dict[str, str]) -> None:
 
 
 def _patch_http_options(
-    options: HttpOptionsDict, patch_options: dict[str, Any]
-) -> HttpOptionsDict:
-  # use shallow copy so we don't override the original objects.
-  copy_option = HttpOptionsDict()
-  copy_option.update(options)
-  for patch_key, patch_value in patch_options.items():
-    # if both are dicts, update the copy.
-    # This is to handle cases like merging headers.
-    if isinstance(patch_value, dict) and isinstance(
-        copy_option.get(patch_key, None), dict
-    ):
-      copy_option[patch_key] = {}
-      copy_option[patch_key].update(
-          options[patch_key]
-      )  # shallow copy from original options.
-      copy_option[patch_key].update(patch_value)
-    elif patch_value is not None:  # Accept empty values.
-      copy_option[patch_key] = patch_value
-  if copy_option['headers']:
-    _append_library_version_headers(copy_option['headers'])
+    options: HttpOptions, patch_options: HttpOptions
+) -> HttpOptions:
+  copy_option = options.model_copy()
+
+  options_headers = copy_option.headers or {}
+  patch_options_headers = patch_options.headers or {}
+  copy_option.headers = {
+      **options_headers,
+      **patch_options_headers,
+  }
+
+  http_options_keys = HttpOptions.model_fields.keys()
+
+  for key in http_options_keys:
+    if key == 'headers':
+      continue
+    patch_value = getattr(patch_options, key, None)
+    if patch_value is not None:
+      setattr(copy_option, key, patch_value)
+    else:
+      setattr(copy_option, key, getattr(options, key))
+
+  if copy_option.headers is not None:
+    _append_library_version_headers(copy_option.headers)
   return copy_option
 
 
@@ -188,9 +192,11 @@ class HttpResponse:
         if chunk:
           # In streaming mode, the chunk of JSON is prefixed with "data:" which
           # we must strip before parsing.
-          if chunk.startswith(b'data: '):
-            chunk = chunk[len(b'data: ') :]
-          yield json.loads(str(chunk, 'utf-8'))
+          if not isinstance(chunk, str):
+            chunk = chunk.decode('utf-8')
+          if chunk.startswith('data: '):
+            chunk = chunk[len('data: ') :]
+          yield json.loads(chunk)
 
   async def async_segments(self) -> AsyncIterator[Any]:
     if isinstance(self.response_stream, list):
@@ -198,7 +204,7 @@ class HttpResponse:
       for chunk in self.response_stream:
         yield json.loads(chunk) if chunk else {}
     elif self.response_stream is None:
-      async for c in []:
+      async for c in []:  # type: ignore[attr-defined]
         yield c
     else:
       # Iterator of objects retrieved from the API.
@@ -206,15 +212,15 @@ class HttpResponse:
         async for chunk in self.response_stream.aiter_lines():
           # This is httpx.Response.
           if chunk:
-            # In async streaming mode, the chunk of JSON is prefixed with "data:"
-            # which we must strip before parsing.
+            # In async streaming mode, the chunk of JSON is prefixed with
+            # "data:" which we must strip before parsing.
+            if not isinstance(chunk, str):
+              chunk = chunk.decode('utf-8')
             if chunk.startswith('data: '):
               chunk = chunk[len('data: ') :]
             yield json.loads(chunk)
       else:
-        raise ValueError(
-            'Error parsing streaming response.'
-        )
+        raise ValueError('Error parsing streaming response.')
 
   def byte_segments(self):
     if isinstance(self.byte_stream, list):
@@ -234,6 +240,41 @@ class HttpResponse:
       response_payload[attribute] = copy.deepcopy(getattr(self, attribute))
 
 
+class SyncHttpxClient(httpx.Client):
+  """Sync httpx client."""
+
+  def __init__(self, **kwargs: Any) -> None:
+    """Initializes the httpx client."""
+    kwargs.setdefault('follow_redirects', True)
+    super().__init__(**kwargs)
+
+  def __del__(self) -> None:
+    """Closes the httpx client."""
+    if self.is_closed:
+      return
+    try:
+      self.close()
+    except Exception:
+      pass
+
+
+class AsyncHttpxClient(httpx.AsyncClient):
+  """Async httpx client."""
+
+  def __init__(self, **kwargs: Any) -> None:
+    """Initializes the httpx client."""
+    kwargs.setdefault('follow_redirects', True)
+    super().__init__(**kwargs)
+
+  def __del__(self) -> None:
+    if self.is_closed:
+      return
+    try:
+      asyncio.get_running_loop().create_task(self.aclose())
+    except Exception:
+      pass
+
+
 class BaseApiClient:
   """Client for calling HTTP APIs sending and receiving JSON."""
 
@@ -269,16 +310,14 @@ class BaseApiClient:
       )
 
     # Validate http_options if it is provided.
-    validated_http_options: dict[str, Any]
+    validated_http_options = HttpOptions()
     if isinstance(http_options, dict):
       try:
-        validated_http_options = HttpOptions.model_validate(
-            http_options
-        ).model_dump()
+        validated_http_options = HttpOptions.model_validate(http_options)
       except ValidationError as e:
         raise ValueError(f'Invalid http_options: {e}')
     elif isinstance(http_options, HttpOptions):
-      validated_http_options = http_options.model_dump()
+      validated_http_options = http_options
 
     # Retrieve implicitly set values from the environment.
     env_project = os.environ.get('GOOGLE_CLOUD_PROJECT', None)
@@ -289,11 +328,15 @@ class BaseApiClient:
     self.api_key = api_key or env_api_key
 
     self._credentials = credentials
-    self._http_options = HttpOptionsDict()
+    self._http_options = HttpOptions()
     # Initialize the lock. This lock will be used to protect access to the
     # credentials. This is crucial for thread safety when multiple coroutines
     # might be accessing the credentials at the same time.
-    self._auth_lock = asyncio.Lock()
+    try:
+      self._auth_lock = asyncio.Lock()
+    except RuntimeError:
+      asyncio.set_event_loop(asyncio.new_event_loop())
+      self._auth_lock = asyncio.Lock()
 
     # Handle when to use Vertex AI in express mode (api key).
     # Explicit initializer arguments are already validated above.
@@ -337,12 +380,12 @@ class BaseApiClient:
             'AI API.'
         )
       if self.api_key or self.location == 'global':
-        self._http_options['base_url'] = f'https://aiplatform.googleapis.com/'
+        self._http_options.base_url = f'https://aiplatform.googleapis.com/'
       else:
-        self._http_options['base_url'] = (
+        self._http_options.base_url = (
             f'https://{self.location}-aiplatform.googleapis.com/'
         )
-      self._http_options['api_version'] = 'v1beta1'
+      self._http_options.api_version = 'v1beta1'
     else:  # Implicit initialization or missing arguments.
       if not self.api_key:
         raise ValueError(
@@ -350,24 +393,27 @@ class BaseApiClient:
             'provide (`api_key`) arguments. To use the Google Cloud API,'
             ' provide (`vertexai`, `project` & `location`) arguments.'
         )
-      self._http_options['base_url'] = (
-          'https://generativelanguage.googleapis.com/'
-      )
-      self._http_options['api_version'] = 'v1beta'
+      self._http_options.base_url = 'https://generativelanguage.googleapis.com/'
+      self._http_options.api_version = 'v1beta'
     # Default options for both clients.
-    self._http_options['headers'] = {'Content-Type': 'application/json'}
+    self._http_options.headers = {'Content-Type': 'application/json'}
     if self.api_key:
-      self._http_options['headers']['x-goog-api-key'] = self.api_key
+      if self._http_options.headers is not None:
+        self._http_options.headers['x-goog-api-key'] = self.api_key
     # Update the http options with the user provided http options.
     if http_options:
       self._http_options = _patch_http_options(
           self._http_options, validated_http_options
       )
     else:
-      _append_library_version_headers(self._http_options['headers'])
+      if self._http_options.headers is not None:
+        _append_library_version_headers(self._http_options.headers)
+    # Initialize the httpx client.
+    self._httpx_client = SyncHttpxClient()
+    self._async_httpx_client = AsyncHttpxClient()
 
   def _websocket_base_url(self):
-    url_parts = urlparse(self._http_options['base_url'])
+    url_parts = urlparse(self._http_options.base_url)
     return url_parts._replace(scheme='wss').geturl()
 
   def _access_token(self) -> str:
@@ -378,9 +424,7 @@ class BaseApiClient:
         self.project = project
 
     if self._credentials:
-      if (
-          self._credentials.expired or not self._credentials.token
-      ):
+      if self._credentials.expired or not self._credentials.token:
         # Only refresh when it needs to. Default expiration is 3600 seconds.
         _refresh_auth(self._credentials)
       if not self._credentials.token:
@@ -404,9 +448,7 @@ class BaseApiClient:
             self.project = project
 
     if self._credentials:
-      if (
-          self._credentials.expired or not self._credentials.token
-      ):
+      if self._credentials.expired or not self._credentials.token:
         # Only refresh when it needs to. Default expiration is 3600 seconds.
         async with self._auth_lock:
           if self._credentials.expired or not self._credentials.token:
@@ -435,11 +477,12 @@ class BaseApiClient:
     if http_options:
       if isinstance(http_options, HttpOptions):
         patched_http_options = _patch_http_options(
-            self._http_options, http_options.model_dump()
+            self._http_options,
+            http_options,
         )
       else:
         patched_http_options = _patch_http_options(
-            self._http_options, http_options
+            self._http_options, HttpOptions.model_validate(http_options)
         )
     else:
       patched_http_options = self._http_options
@@ -458,13 +501,27 @@ class BaseApiClient:
         and not self.api_key
     ):
       path = f'projects/{self.project}/locations/{self.location}/' + path
+
+    if patched_http_options.api_version is None:
+      versioned_path = f'/{path}'
+    else:
+      versioned_path = f'{patched_http_options.api_version}/{path}'
+
+    if (
+        patched_http_options.base_url is None
+        or not patched_http_options.base_url
+    ):
+      raise ValueError('Base URL must be set.')
+    else:
+      base_url = patched_http_options.base_url
+
     url = _join_url_path(
-        patched_http_options.get('base_url', ''),
-        patched_http_options.get('api_version', '') + '/' + path,
+        base_url,
+        versioned_path,
     )
 
-    timeout_in_seconds: Optional[Union[float, int]] = patched_http_options.get(
-        'timeout', None
+    timeout_in_seconds: Optional[Union[float, int]] = (
+        patched_http_options.timeout
     )
     if timeout_in_seconds:
       # HttpOptions.timeout is in milliseconds. But httpx.Client.request()
@@ -473,10 +530,12 @@ class BaseApiClient:
     else:
       timeout_in_seconds = None
 
+    if patched_http_options.headers is None:
+      raise ValueError('Request headers must be set.')
     return HttpRequest(
         method=http_method,
         url=url,
-        headers=patched_http_options['headers'],
+        headers=patched_http_options.headers,
         data=request_dict,
         timeout=timeout_in_seconds,
     )
@@ -488,48 +547,44 @@ class BaseApiClient:
   ) -> HttpResponse:
     data: Optional[Union[str, bytes]] = None
     if self.vertexai and not self.api_key:
-      http_request.headers['Authorization'] = (
-          f'Bearer {self._access_token()}'
-      )
+      http_request.headers['Authorization'] = f'Bearer {self._access_token()}'
       if self._credentials and self._credentials.quota_project_id:
         http_request.headers['x-goog-user-project'] = (
             self._credentials.quota_project_id
         )
-      data = json.dumps(http_request.data)
+      data = json.dumps(http_request.data) if http_request.data else None
     else:
       if http_request.data:
         if not isinstance(http_request.data, bytes):
-          data = json.dumps(http_request.data)
+          data = json.dumps(http_request.data) if http_request.data else None
         else:
           data = http_request.data
 
     if stream:
-      client = httpx.Client()
-      httpx_request = client.build_request(
+      httpx_request = self._httpx_client.build_request(
           method=http_request.method,
           url=http_request.url,
           content=data,
           headers=http_request.headers,
           timeout=http_request.timeout,
       )
-      response = client.send(httpx_request, stream=stream)
+      response = self._httpx_client.send(httpx_request, stream=stream)
       errors.APIError.raise_for_response(response)
       return HttpResponse(
           response.headers, response if stream else [response.text]
       )
     else:
-      with httpx.Client() as client:
-        response = 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 if stream else [response.text]
-        )
+      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 if stream else [response.text]
+      )
 
   async def _async_request(
       self, http_request: HttpRequest, stream: bool = False
@@ -543,50 +598,48 @@ class BaseApiClient:
         http_request.headers['x-goog-user-project'] = (
             self._credentials.quota_project_id
         )
-      data = json.dumps(http_request.data)
+      data = json.dumps(http_request.data) if http_request.data else None
     else:
       if http_request.data:
         if not isinstance(http_request.data, bytes):
-          data = json.dumps(http_request.data)
+          data = json.dumps(http_request.data) if http_request.data else None
         else:
           data = http_request.data
 
     if stream:
-      aclient = httpx.AsyncClient()
-      httpx_request = aclient.build_request(
+      httpx_request = self._async_httpx_client.build_request(
           method=http_request.method,
           url=http_request.url,
           content=data,
           headers=http_request.headers,
           timeout=http_request.timeout,
       )
-      response = await aclient.send(
+      response = await self._async_httpx_client.send(
           httpx_request,
           stream=stream,
       )
-      errors.APIError.raise_for_response(response)
+      await errors.APIError.raise_for_async_response(response)
       return HttpResponse(
           response.headers, response if stream else [response.text]
       )
     else:
-      async with httpx.AsyncClient() as aclient:
-        response = await aclient.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 if stream else [response.text]
-        )
+      response = await self._async_httpx_client.request(
+          method=http_request.method,
+          url=http_request.url,
+          headers=http_request.headers,
+          content=data,
+          timeout=http_request.timeout,
+      )
+      await errors.APIError.raise_for_async_response(response)
+      return HttpResponse(
+          response.headers, response if stream else [response.text]
+      )
 
-  def get_read_only_http_options(self) -> HttpOptionsDict:
-    copied = HttpOptionsDict()
+  def get_read_only_http_options(self) -> dict[str, Any]:
     if isinstance(self._http_options, BaseModel):
-      self._http_options = self._http_options.model_dump()
-    copied.update(self._http_options)
+      copied = self._http_options.model_dump()
+    else:
+      copied = self._http_options
     return copied
 
   def request(
@@ -612,7 +665,7 @@ class BaseApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: Optional[HttpOptionsDict] = None,
+      http_options: Optional[HttpOptionsOrDict] = None,
   ):
     http_request = self._build_request(
         http_method, path, request_dict, http_options
@@ -644,7 +697,7 @@ class BaseApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: Optional[HttpOptionsDict] = None,
+      http_options: Optional[HttpOptionsOrDict] = None,
   ):
     http_request = self._build_request(
         http_method, path, request_dict, http_options
@@ -660,7 +713,7 @@ class BaseApiClient:
 
   def upload_file(
       self, file_path: Union[str, io.IOBase], upload_url: str, upload_size: int
-  ) -> dict[str, str]:
+  ) -> HttpResponse:
     """Transfers a file to the given URL.
 
     Args:
@@ -672,7 +725,7 @@ class BaseApiClient:
         match the size requested in the resumable upload request.
 
     returns:
-          The response json object from the finalize request.
+          The HttpResponse object from the finalize request.
     """
     if isinstance(file_path, io.IOBase):
       return self._upload_fd(file_path, upload_url, upload_size)
@@ -682,7 +735,7 @@ class BaseApiClient:
 
   def _upload_fd(
       self, file: io.IOBase, upload_url: str, upload_size: int
-  ) -> dict[str, str]:
+  ) -> HttpResponse:
     """Transfers a file to the given URL.
 
     Args:
@@ -692,7 +745,7 @@ class BaseApiClient:
         match the size requested in the resumable upload request.
 
     returns:
-          The response json object from the finalize request.
+          The HttpResponse object from the finalize request.
     """
     offset = 0
     # Upload the file in chunks
@@ -705,7 +758,7 @@ class BaseApiClient:
       # If last chunk, finalize the upload.
       if chunk_size + offset >= upload_size:
         upload_command += ', finalize'
-      request = HttpRequest(
+      response = self._httpx_client.request(
           method='POST',
           url=upload_url,
           headers={
@@ -713,25 +766,22 @@ class BaseApiClient:
               'X-Goog-Upload-Offset': str(offset),
               'Content-Length': str(chunk_size),
           },
-          data=file_chunk,
+          content=file_chunk,
       )
-
-      response = self._request(request, stream=False)
       offset += chunk_size
-      if response.headers['X-Goog-Upload-Status'] != 'active':
+      if response.headers['x-goog-upload-status'] != 'active':
         break  # upload is complete or it has been interrupted.
-
       if upload_size <= offset:  # Status is not finalized.
         raise ValueError(
-            'All content has been uploaded, but the upload status is not'
+            f'All content has been uploaded, but the upload status is not'
             f' finalized.'
         )
 
-    if response.headers['X-Goog-Upload-Status'] != 'final':
+    if response.headers['x-goog-upload-status'] != 'final':
       raise ValueError(
           'Failed to upload file: Upload status is not finalized.'
       )
-    return response.json
+    return HttpResponse(response.headers, response_stream=[response.text])
 
   def download_file(self, path: str, http_options):
     """Downloads the file data.
@@ -746,12 +796,7 @@ class BaseApiClient:
     http_request = self._build_request(
         'get', path=path, request_dict={}, http_options=http_options
     )
-    return self._download_file_request(http_request).byte_stream[0]
 
-  def _download_file_request(
-      self,
-      http_request: HttpRequest,
-  ) -> HttpResponse:
     data: Optional[Union[str, bytes]] = None
     if http_request.data:
       if not isinstance(http_request.data, bytes):
@@ -759,24 +804,25 @@ class BaseApiClient:
       else:
         data = http_request.data
 
-    with httpx.Client(follow_redirects=True) as client:
-      response = client.request(
-          method=http_request.method,
-          url=http_request.url,
-          headers=http_request.headers,
-          content=data,
-          timeout=http_request.timeout,
-      )
+    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, byte_stream=[response.read()])
+    errors.APIError.raise_for_response(response)
+    return HttpResponse(
+        response.headers, byte_stream=[response.read()]
+    ).byte_stream[0]
 
   async def async_upload_file(
       self,
       file_path: Union[str, io.IOBase],
       upload_url: str,
       upload_size: int,
-  ) -> dict[str, str]:
+  ) -> HttpResponse:
     """Transfers a file asynchronously to the given URL.
 
     Args:
@@ -787,7 +833,7 @@ class BaseApiClient:
         match the size requested in the resumable upload request.
 
     returns:
-          The response json object from the finalize request.
+          The HttpResponse object from the finalize request.
     """
     if isinstance(file_path, io.IOBase):
       return await self._async_upload_fd(file_path, upload_url, upload_size)
@@ -802,7 +848,7 @@ class BaseApiClient:
       file: Union[io.IOBase, anyio.AsyncFile],
       upload_url: str,
       upload_size: int,
-  ) -> dict[str, str]:
+  ) -> HttpResponse:
     """Transfers a file asynchronously to the given URL.
 
     Args:
@@ -812,47 +858,46 @@ class BaseApiClient:
         match the size requested in the resumable upload request.
 
     returns:
-          The response json object from the finalize request.
+          The HttpResponse object from the finalized request.
     """
-    async with httpx.AsyncClient() as aclient:
-      offset = 0
-      # Upload the file in chunks
-      while True:
-        if isinstance(file, io.IOBase):
-          file_chunk = file.read(CHUNK_SIZE)
-        else:
-          file_chunk = await file.read(CHUNK_SIZE)
-        chunk_size = 0
-        if file_chunk:
-          chunk_size = len(file_chunk)
-        upload_command = 'upload'
-        # If last chunk, finalize the upload.
-        if chunk_size + offset >= upload_size:
-          upload_command += ', finalize'
-        response = await aclient.request(
-            method='POST',
-            url=upload_url,
-            content=file_chunk,
-            headers={
-                'X-Goog-Upload-Command': upload_command,
-                'X-Goog-Upload-Offset': str(offset),
-                'Content-Length': str(chunk_size),
-            },
-        )
-        offset += chunk_size
-        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(
-              'All content has been uploaded, but the upload status is not'
-              f' finalized.'
-          )
-      if response.headers.get('x-goog-upload-status') != 'final':
+    offset = 0
+    # Upload the file in chunks
+    while True:
+      if isinstance(file, io.IOBase):
+        file_chunk = file.read(CHUNK_SIZE)
+      else:
+        file_chunk = await file.read(CHUNK_SIZE)
+      chunk_size = 0
+      if file_chunk:
+        chunk_size = len(file_chunk)
+      upload_command = 'upload'
+      # If last chunk, finalize the upload.
+      if chunk_size + offset >= upload_size:
+        upload_command += ', finalize'
+      response = await self._async_httpx_client.request(
+          method='POST',
+          url=upload_url,
+          content=file_chunk,
+          headers={
+              'X-Goog-Upload-Command': upload_command,
+              'X-Goog-Upload-Offset': str(offset),
+              'Content-Length': str(chunk_size),
+          },
+      )
+      offset += chunk_size
+      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(
-            'Failed to upload file: Upload status is not finalized.'
+            'All content has been uploaded, but the upload status is not'
+            f' finalized.'
         )
-      return response.json()
+    if response.headers.get('x-goog-upload-status') != 'final':
+      raise ValueError(
+          'Failed to upload file: Upload status is not finalized.'
+      )
+    return HttpResponse(response.headers, response_stream=[response.text])
 
   async def async_download_file(self, path: str, http_options):
     """Downloads the file data.
@@ -875,19 +920,18 @@ class BaseApiClient:
       else:
         data = http_request.data
 
-    async with httpx.AsyncClient(follow_redirects=True) as aclient:
-      response = await aclient.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)
+    response = await self._async_httpx_client.request(
+        method=http_request.method,
+        url=http_request.url,
+        headers=http_request.headers,
+        content=data,
+        timeout=http_request.timeout,
+    )
+    await errors.APIError.raise_for_async_response(response)
 
-      return HttpResponse(
-          response.headers, byte_stream=[response.read()]
-      ).byte_stream[0]
+    return HttpResponse(
+        response.headers, byte_stream=[response.read()]
+    ).byte_stream[0]
 
   # This method does nothing in the real api client. It is used in the
   # replay_api_client to verify the response from the SDK method matches the