google-genai 0.5.0__py3-none-any.whl → 0.6.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
@@ -148,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:
@@ -164,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():
@@ -174,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))
@@ -199,7 +219,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 +285,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'] = (
@@ -357,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,
@@ -377,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
 
@@ -488,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.
@@ -503,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,
+      )
+
+      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 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(
@@ -542,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:
@@ -567,22 +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
-
-
-# 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/_common.py

@@ -17,6 +17,7 @@
 
 import base64
 import datetime
+import enum
 import typing
 from typing import Union
 import uuid
@@ -144,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())
@@ -185,7 +186,7 @@ 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,
@@ -208,6 +209,20 @@ class BaseModel(pydantic.BaseModel):
     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:
   """Composes a timestamped unique name.
 
@@ -219,23 +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."""
+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.
+
+  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):
       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] = apply_base64_encoding(value)
+      processed_data[key] = encode_unserializable_types(value)
     elif isinstance(value, list):
       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] = [apply_base64_encoding(v) 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}.'

google/genai/_replay_api_client.py

@@ -17,11 +17,12 @@
 
 import base64
 import copy
+import datetime
 import inspect
+import io
 import json
 import os
 import re
-import datetime
 from typing import Any, Literal, Optional, Union
 
 import google.auth
@@ -32,9 +33,9 @@ from ._api_client import ApiClient
 from ._api_client import HttpOptions
 from ._api_client import HttpRequest
 from ._api_client import HttpResponse
-from ._api_client import RequestJsonEncoder
 from ._common import BaseModel
 
+
 def _redact_version_numbers(version_string: str) -> str:
   """Redacts version numbers in the form x.y.z from a string."""
   return re.sub(r'\d+\.\d+\.\d+', '{VERSION_NUMBER}', version_string)
@@ -145,6 +146,7 @@ class ReplayResponse(BaseModel):
   status_code: int = 200
   headers: dict[str, str]
   body_segments: list[dict[str, object]]
+  byte_segments: Optional[list[bytes]] = None
   sdk_response_segments: list[dict[str, object]]
 
   def model_post_init(self, __context: Any) -> None:
@@ -264,17 +266,13 @@ class ReplayApiClient(ApiClient):
     replay_file_path = self._get_replay_file_path()
     os.makedirs(os.path.dirname(replay_file_path), exist_ok=True)
     with open(replay_file_path, 'w') as f:
-      f.write(
-          json.dumps(
-              self.replay_session.model_dump(mode='json'), indent=2, cls=ResponseJsonEncoder
-          )
-      )
+      f.write(self.replay_session.model_dump_json(exclude_unset=True, indent=2))
     self.replay_session = None
 
   def _record_interaction(
       self,
       http_request: HttpRequest,
-      http_response: Union[HttpResponse, errors.APIError],
+      http_response: Union[HttpResponse, errors.APIError, bytes],
   ):
     if not self._should_update_replay():
       return
@@ -289,6 +287,9 @@ class ReplayApiClient(ApiClient):
       response = ReplayResponse(
           headers=dict(http_response.headers),
           body_segments=list(http_response.segments()),
+          byte_segments=[
+              seg[:100] + b'...' for seg in http_response.byte_segments()
+          ],
           status_code=http_response.status_code,
           sdk_response_segments=[],
       )
@@ -322,11 +323,7 @@ class ReplayApiClient(ApiClient):
     # so that the comparison is fair.
     _redact_request_body(request_data_copy)
 
-    # Need to call dumps() and loads() to convert dict bytes values to strings.
-    # Because the expected_request_body dict never contains bytes values.
-    actual_request_body = [
-        json.loads(json.dumps(request_data_copy, cls=RequestJsonEncoder))
-    ]
+    actual_request_body = [request_data_copy]
     expected_request_body = interaction.request.body_segments
     assert actual_request_body == expected_request_body, (
         'Request body mismatch:\n'
@@ -349,6 +346,7 @@ class ReplayApiClient(ApiClient):
             json.dumps(segment)
             for segment in interaction.response.body_segments
         ],
+        byte_stream=interaction.response.byte_segments,
     )
 
   def _verify_response(self, response_model: BaseModel):
@@ -368,7 +366,9 @@ class ReplayApiClient(ApiClient):
       response_model = response_model[0]
     print('response_model: ', response_model.model_dump(exclude_none=True))
     actual = response_model.model_dump(exclude_none=True, mode='json')
-    expected = interaction.response.sdk_response_segments[self._sdk_response_index]
+    expected = interaction.response.sdk_response_segments[
+        self._sdk_response_index
+    ]
     assert (
         actual == expected
     ), f'SDK response mismatch:\nActual: {actual}\nExpected: {expected}'
@@ -402,10 +402,21 @@ class ReplayApiClient(ApiClient):
     else:
       return self._build_response_from_replay(http_request)
 
-  def upload_file(self, file_path: str, upload_url: str, upload_size: int):
-    request = HttpRequest(
-        method='POST', url='', data={'file_path': file_path}, headers={}
-    )
+  def upload_file(self, file_path: Union[str, io.IOBase], upload_url: str, upload_size: int):
+    if isinstance(file_path, io.IOBase):
+      offset = file_path.tell()
+      content = file_path.read()
+      file_path.seek(offset, os.SEEK_SET)
+      request = HttpRequest(
+          method='POST',
+          url='',
+          data={'bytes': base64.b64encode(content).decode('utf-8')},
+          headers={}
+      )
+    else:
+      request = HttpRequest(
+          method='POST', url='', data={'file_path': file_path}, headers={}
+      )
     if self._should_call_api():
       try:
         result = super().upload_file(file_path, upload_url, upload_size)
@@ -420,18 +431,19 @@ class ReplayApiClient(ApiClient):
     else:
       return self._build_response_from_replay(request).text
 
-
-# TODO(b/389693448): Cleanup datetime hacks.
-class ResponseJsonEncoder(json.JSONEncoder):
-  """The replay test json encoder for response.
-  """
-  def default(self, o):
-    if isinstance(o, datetime.datetime):
-      # dt.isoformat() prints "2024-11-15T23:27:45.624657+00:00"
-      # but replay files want "2024-11-15T23:27:45.624657Z"
-      if o.isoformat().endswith('+00:00'):
-        return o.isoformat().replace('+00:00', 'Z')
-      else:
-        return o.isoformat()
+  def _download_file_request(self, request):
+    self._initialize_replay_session_if_not_loaded()
+    if self._should_call_api():
+      try:
+        result = super()._download_file_request(request)
+      except HTTPError as e:
+        result = HttpResponse(
+            e.response.headers, [json.dumps({'reason': e.response.reason})]
+        )
+        result.status_code = e.response.status_code
+        raise e
+      self._record_interaction(request, result)
+      return result
     else:
-      return super().default(o)
+      return self._build_response_from_replay(request)
+