google-genai 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

google/genai/_api_client.py

@@ -20,6 +20,7 @@ import asyncio
 import copy
 from dataclasses import dataclass
 import datetime
+import io
 import json
 import logging
 import os
@@ -35,55 +36,7 @@ import requests
 
 from . import errors
 from . import version
-
-
-class HttpOptions(BaseModel):
-  """HTTP options for the api client."""
-  model_config = ConfigDict(extra='forbid')
-
-  base_url: Optional[str] = Field(
-      default=None,
-      description="""The base URL for the AI platform service endpoint.""",
-  )
-  api_version: Optional[str] = Field(
-      default=None,
-      description="""Specifies the version of the API to use.""",
-  )
-  headers: Optional[dict[str, str]] = Field(
-      default=None,
-      description="""Additional HTTP headers to be sent with the request.""",
-  )
-  response_payload: Optional[dict] = Field(
-      default=None,
-      description="""If set, the response payload will be returned int the supplied dict.""",
-  )
-  timeout: Optional[Union[float, Tuple[float, float]]] = Field(
-      default=None,
-      description="""Timeout for the request in seconds.""",
-  )
-  skip_project_and_location_in_path: bool = Field(
-      default=False,
-      description="""If set to True, the project and location will not be appended to the path.""",
-  )
-
-
-class HttpOptionsDict(TypedDict):
-  """HTTP options for the api client."""
-
-  base_url: Optional[str] = None
-  """The base URL for the AI platform service endpoint."""
-  api_version: Optional[str] = None
-  """Specifies the version of the API to use."""
-  headers: Optional[dict[str, str]] = None
-  """Additional HTTP headers to be sent with the request."""
-  response_payload: Optional[dict] = None
-  """If set, the response payload will be returned int the supplied dict."""
-  timeout: Optional[Union[float, Tuple[float, float]]] = None
-  """Timeout for the request in seconds."""
-  skip_project_and_location_in_path: bool = False
-  """If set to True, the project and location will not be appended to the path."""
-
-HttpOptionsOrDict = Union[HttpOptions, HttpOptionsDict]
+from .types import HttpOptions, HttpOptionsDict, HttpOptionsOrDict
 
 
 def _append_library_version_headers(headers: dict[str, str]) -> None:
@@ -143,18 +96,35 @@ class HttpRequest:
   url: str
   method: str
   data: Union[dict[str, object], bytes]
-  timeout: Optional[Union[float, Tuple[float, float]]] = None
+  timeout: Optional[float] = None
 
 
 class HttpResponse:
 
-  def __init__(self, headers: dict[str, str], response_stream: Union[Any, str]):
+  def __init__(
+      self,
+      headers: dict[str, str],
+      response_stream: Union[Any, str] = None,
+      byte_stream: Union[Any, bytes] = None,
+  ):
     self.status_code = 200
     self.headers = headers
     self.response_stream = response_stream
+    self.byte_stream = byte_stream
+    self.segment_iterator = self.segments()
+
+  # Async iterator for async streaming.
+  def __aiter__(self):
+    return self
+
+  async def __anext__(self):
+    try:
+      return next(self.segment_iterator)
+    except StopIteration:
+      raise StopAsyncIteration
 
   @property
-  def text(self) -> str:
+  def json(self) -> Any:
     if not self.response_stream[0]:  # Empty response
       return ''
     return json.loads(self.response_stream[0])
@@ -164,6 +134,8 @@ class HttpResponse:
       # list of objects retrieved from replay or from non-streaming API.
       for chunk in self.response_stream:
         yield json.loads(chunk) if chunk else {}
+    elif self.response_stream is None:
+      yield from []
     else:
       # Iterator of objects retrieved from the API.
       for chunk in self.response_stream.iter_lines():
@@ -174,7 +146,20 @@ class HttpResponse:
             chunk = chunk[len(b'data: ') :]
           yield json.loads(str(chunk, 'utf-8'))
 
-  def copy_to_dict(self, response_payload: dict[str, object]):
+  def byte_segments(self):
+    if isinstance(self.byte_stream, list):
+      # list of objects retrieved from replay or from non-streaming API.
+      yield from self.byte_stream
+    elif self.byte_stream is None:
+      yield from []
+    else:
+      raise ValueError(
+          'Byte segments are not supported for streaming responses.'
+      )
+
+  def _copy_to_dict(self, response_payload: dict[str, object]):
+    # Cannot pickle 'generator' object.
+    delattr(self, 'segment_iterator')
     for attribute in dir(self):
       response_payload[attribute] = copy.deepcopy(getattr(self, attribute))
 
@@ -199,7 +184,7 @@ class ApiClient:
       ]:
         self.vertexai = True
 
-    # Validate explicitly set intializer values.
+    # Validate explicitly set initializer values.
     if (project or location) and api_key:
       # API cannot consume both project/location and api_key.
       raise ValueError(
@@ -265,9 +250,10 @@ class ApiClient:
         self.api_key = None
       if not self.project and not self.api_key:
         self.project = google.auth.default()[1]
-      if not (self.project or self.location) and not self.api_key:
+      if not ((self.project and self.location) or self.api_key):
         raise ValueError(
-            'Project/location or API key must be set when using the Vertex AI API.'
+            'Project and location or API key must be set when using the Vertex '
+            'AI API.'
         )
       if self.api_key:
         self._http_options['base_url'] = (
@@ -304,7 +290,7 @@ class ApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: HttpOptionsDict = None,
+      http_options: HttpOptionsOrDict = None,
   ) -> HttpRequest:
     # Remove all special dict keys such as _url and _query.
     keys_to_delete = [key for key in request_dict.keys() if key.startswith('_')]
@@ -312,18 +298,28 @@ class ApiClient:
       del request_dict[key]
     # patch the http options with the user provided settings.
     if http_options:
-      patched_http_options = _patch_http_options(
-          self._http_options, http_options
-      )
+      if isinstance(http_options, HttpOptions):
+        patched_http_options = _patch_http_options(
+            self._http_options, http_options.model_dump()
+        )
+      else:
+        patched_http_options = _patch_http_options(
+            self._http_options, http_options
+        )
     else:
       patched_http_options = self._http_options
-    skip_project_and_location_in_path_val = patched_http_options.get(
-        'skip_project_and_location_in_path', False
-    )
+    # Skip adding project and locations when getting Vertex AI base models.
+    query_vertex_base_models = False
+    if (
+        self.vertexai
+        and http_method == 'get'
+        and path.startswith('publishers/google/models')
+    ):
+      query_vertex_base_models = True
     if (
         self.vertexai
         and not path.startswith('projects/')
-        and not skip_project_and_location_in_path_val
+        and not query_vertex_base_models
         and not self.api_key
     ):
       path = f'projects/{self.project}/locations/{self.location}/' + path
@@ -333,12 +329,19 @@ class ApiClient:
         patched_http_options['base_url'],
         patched_http_options['api_version'] + '/' + path,
     )
+
+    timeout_in_seconds = patched_http_options.get('timeout', None)
+    if timeout_in_seconds:
+      timeout_in_seconds = timeout_in_seconds / 1000.0
+    else:
+      timeout_in_seconds = None
+
     return HttpRequest(
         method=http_method,
         url=url,
         headers=patched_http_options['headers'],
         data=request_dict,
-        timeout=patched_http_options.get('timeout', None),
+        timeout=timeout_in_seconds,
     )
 
   def _request(
@@ -357,7 +360,7 @@ class ApiClient:
           http_request.method.upper(),
           http_request.url,
           headers=http_request.headers,
-          data=json.dumps(http_request.data, cls=RequestJsonEncoder)
+          data=json.dumps(http_request.data)
           if http_request.data
           else None,
           timeout=http_request.timeout,
@@ -377,7 +380,7 @@ class ApiClient:
     data = None
     if http_request.data:
       if not isinstance(http_request.data, bytes):
-        data = json.dumps(http_request.data, cls=RequestJsonEncoder)
+        data = json.dumps(http_request.data)
       else:
         data = http_request.data
 
@@ -427,15 +430,24 @@ class ApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: HttpOptionsDict = None,
+      http_options: HttpOptionsOrDict = None,
   ):
     http_request = self._build_request(
         http_method, path, request_dict, http_options
     )
     response = self._request(http_request, stream=False)
-    if http_options and 'response_payload' in http_options:
-      response.copy_to_dict(http_options['response_payload'])
-    return response.text
+    if http_options:
+      if (
+          isinstance(http_options, HttpOptions)
+          and http_options.deprecated_response_payload is not None
+      ):
+        response._copy_to_dict(http_options.deprecated_response_payload)
+      elif (
+          isinstance(http_options, dict)
+          and 'deprecated_response_payload' in http_options
+      ):
+        response._copy_to_dict(http_options['deprecated_response_payload'])
+    return response.json
 
   def request_streamed(
       self,
@@ -449,8 +461,10 @@ class ApiClient:
     )
 
     session_response = self._request(http_request, stream=True)
-    if http_options and 'response_payload' in http_options:
-      session_response.copy_to_dict(http_options['response_payload'])
+    if http_options and 'deprecated_response_payload' in http_options:
+      session_response._copy_to_dict(
+          http_options['deprecated_response_payload']
+      )
     for chunk in session_response.segments():
       yield chunk
 
@@ -466,9 +480,9 @@ class ApiClient:
     )
 
     result = await self._async_request(http_request=http_request, stream=False)
-    if http_options and 'response_payload' in http_options:
-      result.copy_to_dict(http_options['response_payload'])
-    return result.text
+    if http_options and 'deprecated_response_payload' in http_options:
+      result._copy_to_dict(http_options['deprecated_response_payload'])
+    return result.json
 
   async def async_request_streamed(
       self,
@@ -483,17 +497,42 @@ class ApiClient:
 
     response = await self._async_request(http_request=http_request, stream=True)
 
-    for chunk in response.segments():
-      yield chunk
-    if http_options and 'response_payload' in http_options:
-      response.copy_to_dict(http_options['response_payload'])
+    if http_options and 'deprecated_response_payload' in http_options:
+      response._copy_to_dict(http_options['deprecated_response_payload'])
+    async def async_generator():
+      async for chunk in response:
+        yield chunk
+    return async_generator()
 
-  def upload_file(self, file_path: str, upload_url: str, upload_size: int):
+  def upload_file(
+      self, file_path: Union[str, io.IOBase], upload_url: str, upload_size: int
+  ) -> str:
     """Transfers a file to the given URL.
 
     Args:
-      file_path: The full path to the file. If the local file path is not found,
-        an error will be raised.
+      file_path: The full path to the file or a file like object inherited from
+        io.BytesIO. If the local file path is not found, an error will be
+        raised.
+      upload_url: The URL to upload the file to.
+      upload_size: The size of file content to be uploaded, this will have to
+        match the size requested in the resumable upload request.
+
+    returns:
+          The response json object from the finalize request.
+    """
+    if isinstance(file_path, io.IOBase):
+      return self._upload_fd(file_path, upload_url, upload_size)
+    else:
+      with open(file_path, 'rb') as file:
+        return self._upload_fd(file, upload_url, upload_size)
+
+  def _upload_fd(
+      self, file: io.IOBase, upload_url: str, upload_size: int
+  ) -> str:
+    """Transfers a file to the given URL.
+
+    Args:
+      file: A file like object inherited from io.BytesIO.
       upload_url: The URL to upload the file to.
       upload_size: The size of file content to be uploaded, this will have to
         match the size requested in the resumable upload request.
@@ -503,51 +542,89 @@ class ApiClient:
     """
     offset = 0
     # Upload the file in chunks
-    with open(file_path, 'rb') as file:
-      while True:
-        file_chunk = file.read(1024 * 1024 * 8)  # 8 MB 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'
-
-        request = HttpRequest(
-            method='POST',
-            url=upload_url,
-            headers={
-                'X-Goog-Upload-Command': upload_command,
-                'X-Goog-Upload-Offset': str(offset),
-                'Content-Length': str(chunk_size),
-            },
-            data=file_chunk,
-        )
-        response = self._request_unauthorized(request, stream=False)
-        offset += chunk_size
-        if response.headers['X-Goog-Upload-Status'] != 'active':
-          break  # upload is complete or it has been interrupted.
+    while True:
+      file_chunk = file.read(1024 * 1024 * 8)  # 8 MB 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'
+      request = HttpRequest(
+          method='POST',
+          url=upload_url,
+          headers={
+              'X-Goog-Upload-Command': upload_command,
+              'X-Goog-Upload-Offset': str(offset),
+              'Content-Length': str(chunk_size),
+          },
+          data=file_chunk,
+      )
 
-        if upload_size <= offset:  # Status is not finalized.
-          raise ValueError(
-              'All content has been uploaded, but the upload status is not'
-              f' finalized. {response.headers}, body: {response.text}'
-          )
+      response = self._request_unauthorized(request, stream=False)
+      offset += chunk_size
+      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' finalized. {response.headers}, body: {response.json}'
+        )
 
     if response.headers['X-Goog-Upload-Status'] != 'final':
       raise ValueError(
           'Failed to upload file: Upload status is not finalized. headers:'
-          f' {response.headers}, body: {response.text}'
+          f' {response.headers}, body: {response.json}'
       )
-    return response.text
+    return response.json
+
+  def download_file(self, path: str, http_options):
+    """Downloads the file data.
+
+    Args:
+      path: The request path with query params.
+      http_options: The http options to use for the request.
+
+    returns:
+          The file bytes
+    """
+    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 = None
+    if http_request.data:
+      if not isinstance(http_request.data, bytes):
+        data = json.dumps(http_request.data, cls=RequestJsonEncoder)
+      else:
+        data = http_request.data
+
+    http_session = requests.Session()
+    response = http_session.request(
+        method=http_request.method,
+        url=http_request.url,
+        headers=http_request.headers,
+        data=data,
+        timeout=http_request.timeout,
+        stream=False,
+    )
+
+    errors.APIError.raise_for_response(response)
+    return HttpResponse(response.headers, byte_stream=[response.content])
 
   async def async_upload_file(
       self,
-      file_path: str,
+      file_path: Union[str, io.IOBase],
       upload_url: str,
       upload_size: int,
-  ):
+  ) -> str:
     """Transfers a file asynchronously to the given URL.
 
     Args:
@@ -567,22 +644,48 @@ class ApiClient:
         upload_size,
     )
 
+  async def _async_upload_fd(
+      self,
+      file: io.IOBase,
+      upload_url: str,
+      upload_size: int,
+  ) -> str:
+    """Transfers a file asynchronously to the given URL.
+
+    Args:
+      file: A file like object inherited from io.BytesIO.
+      upload_url: The URL to upload the file to.
+      upload_size: The size of file content to be uploaded, this will have to
+        match the size requested in the resumable upload request.
+
+    returns:
+          The response json object from the finalize request.
+    """
+    return await asyncio.to_thread(
+        self._upload_fd,
+        file,
+        upload_url,
+        upload_size,
+    )
+
+  async def async_download_file(self, path: str, http_options):
+    """Downloads the file data.
+
+    Args:
+      path: The request path with query params.
+      http_options: The http options to use for the request.
+
+    returns:
+          The file bytes
+    """
+    return await asyncio.to_thread(
+        self.download_file,
+        path,
+        http_options,
+    )
+
   # 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
   # recorded response.
   def _verify_response(self, response_model: BaseModel):
     pass
-
-
-# TODO(b/389693448): Cleanup datetime hacks.
-class RequestJsonEncoder(json.JSONEncoder):
-  """Encode bytes as strings without modify its content."""
-
-  def default(self, o):
-    if isinstance(o, datetime.datetime):
-      # This Zulu time format is used by the Vertex AI API and the test recorder
-      # Using strftime works well, but we want to align with the replay encoder.
-      # o.astimezone(datetime.timezone.utc).strftime('%Y-%m-%dT%H:%M:%S.%fZ')
-      return o.isoformat().replace('+00:00', 'Z')
-    else:
-      return super().default(o)

google/genai/_api_module.py

@@ -0,0 +1,24 @@
+# Copyright 2024 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.
+#
+
+"""Utilities for the API Modules of the Google Gen AI SDK."""
+
+from . import _api_client
+
+
+class BaseModule:
+
+  def __init__(self, api_client_: _api_client.ApiClient):
+    self._api_client = api_client_