google-genai 0.4.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

google/genai/_api_client.py

@@ -20,7 +20,9 @@ import asyncio
 import copy
 from dataclasses import dataclass
 import datetime
+import io
 import json
+import logging
 import os
 import sys
 from typing import Any, Optional, Tuple, TypedDict, Union
@@ -60,6 +62,10 @@ class HttpOptions(BaseModel):
       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):
@@ -69,13 +75,14 @@ class HttpOptionsDict(TypedDict):
   """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, Union[str, list[str]]]] = None
+  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]
 
@@ -133,7 +140,7 @@ def _join_url_path(base_url: str, path: str) -> str:
 
 @dataclass
 class HttpRequest:
-  headers: dict[str, Union[str, list[str]]]
+  headers: dict[str, str]
   url: str
   method: str
   data: Union[dict[str, object], bytes]
@@ -142,10 +149,16 @@ class HttpRequest:
 
 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
 
   @property
   def text(self) -> str:
@@ -158,6 +171,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():
@@ -168,6 +183,17 @@ class HttpResponse:
             chunk = chunk[len(b'data: ') :]
           yield json.loads(str(chunk, 'utf-8'))
 
+  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]):
     for attribute in dir(self):
       response_payload[attribute] = copy.deepcopy(getattr(self, attribute))
@@ -193,11 +219,17 @@ 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(
           'Project/location and API key are mutually exclusive in the client initializer.'
       )
+    elif credentials and api_key:
+      # API cannot consume both credentials and api_key.
+      raise ValueError(
+          'Credentials and API key are mutually exclusive in the client initializer.'
+      )
 
     # Validate http_options if a dict is provided.
     if isinstance(http_options, dict):
@@ -208,26 +240,66 @@ class ApiClient:
     elif(isinstance(http_options, HttpOptions)):
       http_options = http_options.model_dump()
 
-    self.api_key: Optional[str] = None
-    self.project = project or os.environ.get('GOOGLE_CLOUD_PROJECT', None)
-    self.location = location or os.environ.get('GOOGLE_CLOUD_LOCATION', None)
+    # Retrieve implicitly set values from the environment.
+    env_project = os.environ.get('GOOGLE_CLOUD_PROJECT', None)
+    env_location = os.environ.get('GOOGLE_CLOUD_LOCATION', None)
+    env_api_key = os.environ.get('GOOGLE_API_KEY', None)
+    self.project = project or env_project
+    self.location = location or env_location
+    self.api_key = api_key or env_api_key
+
     self._credentials = credentials
     self._http_options = HttpOptionsDict()
 
+    # Handle when to use Vertex AI in express mode (api key).
+    # Explicit initializer arguments are already validated above.
     if self.vertexai:
-      if not self.project:
+      if credentials:
+        # Explicit credentials take precedence over implicit api_key.
+        logging.info(
+            'The user provided Google Cloud credentials will take precedence'
+            + ' over the API key from the environment variable.'
+        )
+        self.api_key = None
+      elif (env_location or env_project) and api_key:
+        # Explicit api_key takes precedence over implicit project/location.
+        logging.info(
+            'The user provided Vertex AI API key will take precedence over the'
+            + ' project/location from the environment variables.'
+        )
+        self.project = None
+        self.location = None
+      elif (project or location) and env_api_key:
+        # Explicit project/location takes precedence over implicit api_key.
+        logging.info(
+            'The user provided project/location will take precedence over the'
+            + ' Vertex AI API key from the environment variable.'
+        )
+        self.api_key = None
+      elif (env_location or env_project) and env_api_key:
+        # Implicit project/location takes precedence over implicit api_key.
+        logging.info(
+            'The project/location from the environment variables will take'
+            + ' precedence over the API key from the environment variables.'
+        )
+        self.api_key = None
+      if not self.project and not self.api_key:
         self.project = google.auth.default()[1]
-      # Will change this to support EasyGCP in the future.
-      if not self.project or not self.location:
+      if not ((self.project and self.location) or self.api_key):
         raise ValueError(
-            'Project and location 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'] = (
+            f'https://aiplatform.googleapis.com/'
+        )
+      else:
+        self._http_options['base_url'] = (
+            f'https://{self.location}-aiplatform.googleapis.com/'
         )
-      self._http_options['base_url'] = (
-          f'https://{self.location}-aiplatform.googleapis.com/'
-      )
       self._http_options['api_version'] = 'v1beta1'
     else:  # ML Dev API
-      self.api_key = api_key or os.environ.get('GOOGLE_API_KEY', None)
       if not self.api_key:
         raise ValueError('API key must be set when using the Google AI API.')
       self._http_options['base_url'] = (
@@ -236,7 +308,7 @@ class ApiClient:
       self._http_options['api_version'] = 'v1beta'
     # Default options for both clients.
     self._http_options['headers'] = {'Content-Type': 'application/json'}
-    if self.api_key:
+    if self.api_key and not self.vertexai:
       self._http_options['headers']['x-goog-api-key'] = self.api_key
     # Update the http options with the user provided http options.
     if http_options:
@@ -266,8 +338,18 @@ class ApiClient:
       )
     else:
       patched_http_options = self._http_options
-    if self.vertexai and not path.startswith('projects/'):
+    skip_project_and_location_in_path_val = patched_http_options.get(
+        'skip_project_and_location_in_path', False
+    )
+    if (
+        self.vertexai
+        and not path.startswith('projects/')
+        and not skip_project_and_location_in_path_val
+        and not self.api_key
+    ):
       path = f'projects/{self.project}/locations/{self.location}/' + path
+    elif self.vertexai and self.api_key:
+      path = f'{path}?key={self.api_key}'
     url = _join_url_path(
         patched_http_options['base_url'],
         patched_http_options['api_version'] + '/' + path,
@@ -285,7 +367,7 @@ class ApiClient:
       http_request: HttpRequest,
       stream: bool = False,
   ) -> HttpResponse:
-    if self.vertexai:
+    if self.vertexai and not self.api_key:
       if not self._credentials:
         self._credentials, _ = google.auth.default(
             scopes=["https://www.googleapis.com/auth/cloud-platform"],
@@ -296,7 +378,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,
@@ -316,7 +398,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,12 +509,35 @@ class ApiClient:
     if http_options and 'response_payload' in http_options:
       response.copy_to_dict(http_options['response_payload'])
 
-  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.
@@ -442,37 +547,36 @@ 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.text}'
+        )
 
     if response.headers['X-Goog-Upload-Status'] != 'final':
       raise ValueError(
@@ -481,12 +585,52 @@ class ApiClient:
       )
     return response.text
 
+  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:
@@ -506,23 +650,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
-
-
-class RequestJsonEncoder(json.JSONEncoder):
-  """Encode bytes as strings without modify its content."""
-
-  def default(self, o):
-    if isinstance(o, bytes):
-      return o.decode()
-    elif 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/_common.py

@@ -17,7 +17,7 @@
 
 import base64
 import datetime
-import json
+import enum
 import typing
 from typing import Union
 import uuid
@@ -116,7 +116,7 @@ def get_value_by_path(data: object, keys: list[str]):
 class BaseModule:
 
   def __init__(self, api_client_: _api_client.ApiClient):
-    self.api_client = api_client_
+    self._api_client = api_client_
 
 
 def convert_to_dict(obj: dict[str, object]) -> dict[str, object]:
@@ -145,7 +145,7 @@ def _remove_extra_fields(
 ) -> None:
   """Removes extra fields from the response that are not in the model.
 
-  Muates the response in place.
+  Mutates the response in place.
   """
 
   key_values = list(response.items())
@@ -186,10 +186,12 @@ class BaseModel(pydantic.BaseModel):
       alias_generator=alias_generators.to_camel,
       populate_by_name=True,
       from_attributes=True,
-      protected_namespaces={},
+      protected_namespaces=(),
       extra='forbid',
       # This allows us to use arbitrary types in the model. E.g. PIL.Image.
       arbitrary_types_allowed=True,
+      ser_json_bytes='base64',
+      val_json_bytes='base64',
   )
 
   @classmethod
@@ -201,7 +203,24 @@ class BaseModel(pydantic.BaseModel):
     # We will provide another mechanism to allow users to access these fields.
     _remove_extra_fields(cls, response)
     validated_response = cls.model_validate(response)
-    return apply_base64_decoding_for_model(validated_response)
+    return validated_response
+
+  def to_json_dict(self) -> dict[str, object]:
+    return self.model_dump(exclude_none=True, mode='json')
+
+
+class CaseInSensitiveEnum(str, enum.Enum):
+  """Case insensitive enum."""
+
+  @classmethod
+  def _missing_(cls, value):
+    try:
+      return cls[value.upper()]  # Try to access directly with uppercase
+    except KeyError:
+      try:
+        return cls[value.lower()]  # Try to access directly with lowercase
+      except KeyError as e:
+        raise ValueError(f"{value} is not a valid {cls.__name__}") from e
 
 
 def timestamped_unique_name() -> str:
@@ -215,42 +234,39 @@ def timestamped_unique_name() -> str:
   return f'{timestamp}_{unique_id}'
 
 
-def apply_base64_encoding(data: dict[str, object]) -> dict[str, object]:
-  """Applies base64 encoding to bytes values in the given data."""
-  return process_bytes_fields(data, encode=True)
-
-
-def apply_base64_decoding(data: dict[str, object]) -> dict[str, object]:
-  """Applies base64 decoding to bytes values in the given data."""
-  return process_bytes_fields(data, encode=False)
+def encode_unserializable_types(data: dict[str, object]) -> dict[str, object]:
+  """Converts unserializable types in dict to json.dumps() compatible types.
 
+  This function is called in models.py after calling convert_to_dict(). The
+  convert_to_dict() can convert pydantic object to dict. However, the input to
+  convert_to_dict() is dict mixed of pydantic object and nested dict(the output
+  of converters). So they may be bytes in the dict and they are out of
+  `ser_json_bytes` control in model_dump(mode='json') called in
+  `convert_to_dict`, as well as datetime deserialization in Pydantic json mode.
 
-def apply_base64_decoding_for_model(data: BaseModel) -> BaseModel:
-  d = data.model_dump(exclude_none=True)
-  d = apply_base64_decoding(d)
-  return data.model_validate(d)
-
-
-def process_bytes_fields(data: dict[str, object], encode=True) -> dict[str, object]:
+  Returns:
+    A dictionary with json.dumps() incompatible type (e.g. bytes datetime)
+    to compatible type (e.g. base64 encoded string, isoformat date string).
+  """
   processed_data = {}
   if not isinstance(data, dict):
     return data
   for key, value in data.items():
     if isinstance(value, bytes):
-      if encode:
-        processed_data[key] = base64.b64encode(value)
-      else:
-        processed_data[key] = base64.b64decode(value)
+      processed_data[key] = base64.urlsafe_b64encode(value).decode('ascii')
+    elif isinstance(value, datetime.datetime):
+      processed_data[key] = value.isoformat()
     elif isinstance(value, dict):
-      processed_data[key] = process_bytes_fields(value, encode)
+      processed_data[key] = encode_unserializable_types(value)
     elif isinstance(value, list):
-      if encode and all(isinstance(v, bytes) for v in value):
-        processed_data[key] = [base64.b64encode(v) for v in value]
-      elif all(isinstance(v, bytes) for v in value):
-        processed_data[key] = [base64.b64decode(v) for v in value]
+      if all(isinstance(v, bytes) for v in value):
+        processed_data[key] = [
+            base64.urlsafe_b64encode(v).decode('ascii') for v in value
+        ]
+      if all(isinstance(v, datetime.datetime) for v in value):
+        processed_data[key] = [v.isoformat() for v in value]
       else:
-        processed_data[key] = [process_bytes_fields(v, encode) for v in value]
+        processed_data[key] = [encode_unserializable_types(v) for v in value]
     else:
       processed_data[key] = value
   return processed_data
-

google/genai/_extra_utils.py

@@ -116,7 +116,7 @@ def convert_if_exist_pydantic_model(
     try:
       return annotation(**value)
     except pydantic.ValidationError as e:
-      raise errors.UnkownFunctionCallArgumentError(
+      raise errors.UnknownFunctionCallArgumentError(
           f'Failed to parse parameter {param_name} for function'
           f' {func_name} from function call part because function call argument'
           f' value {value} is not compatible with parameter annotation'
@@ -150,7 +150,7 @@ def convert_if_exist_pydantic_model(
         except pydantic.ValidationError:
           continue
     # if none of the union type is matched, raise error
-    raise errors.UnkownFunctionCallArgumentError(
+    raise errors.UnknownFunctionCallArgumentError(
         f'Failed to parse parameter {param_name} for function'
         f' {func_name} from function call part because function call argument'
         f' value {value} cannot be converted to parameter annotation'
@@ -161,7 +161,7 @@ def convert_if_exist_pydantic_model(
   if isinstance(value, int) and annotation is float:
     return value
   if not isinstance(value, annotation):
-    raise errors.UnkownFunctionCallArgumentError(
+    raise errors.UnknownFunctionCallArgumentError(
         f'Failed to parse parameter {param_name} for function {func_name} from'
         f' function call part because function call argument value {value} is'
         f' not compatible with parameter annotation {annotation}.'