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

google/genai/_api_client.py

@@ -21,6 +21,7 @@ import copy
 from dataclasses import dataclass
 import datetime
 import json
+import logging
 import os
 import sys
 from typing import Any, Optional, Tuple, TypedDict, Union
@@ -60,6 +61,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 +74,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 +139,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]
@@ -195,9 +201,15 @@ class ApiClient:
 
     # Validate explicitly set intializer 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 +220,65 @@ 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 or self.location) and not self.api_key:
         raise ValueError(
-            'Project and location must be set when using the Vertex AI API.'
+            'Project/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 +287,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 +317,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 +346,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"],
@@ -513,13 +574,12 @@ class ApiClient:
     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, bytes):
-      return o.decode()
-    elif isinstance(o, datetime.datetime):
+    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')

google/genai/_common.py

@@ -17,7 +17,6 @@
 
 import base64
 import datetime
-import json
 import typing
 from typing import Union
 import uuid
@@ -116,7 +115,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]:
@@ -190,6 +189,8 @@ class BaseModel(pydantic.BaseModel):
       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 +202,10 @@ 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')
 
 
 def timestamped_unique_name() -> str:
@@ -217,40 +221,21 @@ def timestamped_unique_name() -> str:
 
 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 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]:
   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, dict):
-      processed_data[key] = process_bytes_fields(value, encode)
+      processed_data[key] = apply_base64_encoding(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
+        ]
       else:
-        processed_data[key] = [process_bytes_fields(v, encode) for v in value]
+        processed_data[key] = [apply_base64_encoding(v) for v in value]
     else:
       processed_data[key] = value
   return processed_data
-

google/genai/_replay_api_client.py

@@ -25,7 +25,6 @@ import datetime
 from typing import Any, Literal, Optional, Union
 
 import google.auth
-from pydantic import BaseModel
 from requests.exceptions import HTTPError
 
 from . import errors
@@ -34,6 +33,7 @@ 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."""
@@ -72,6 +72,11 @@ def _redact_request_url(url: str) -> str:
       '{VERTEX_URL_PREFIX}/',
       url,
   )
+  result = re.sub(
+      r'.*-aiplatform.googleapis.com/[^/]+/',
+      '{VERTEX_URL_PREFIX}/',
+      result,
+  )
   result = re.sub(
       r'https://generativelanguage.googleapis.com/[^/]+',
       '{MLDEV_URL_PREFIX}',
@@ -259,18 +264,9 @@ 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:
-      replay_session_dict = self.replay_session.model_dump()
-      # Use for non-utf-8 bytes in image/video... output.
-      for interaction in replay_session_dict['interactions']:
-        segments = []
-        for response in interaction['response']['sdk_response_segments']:
-          segments.append(json.loads(json.dumps(
-              response, cls=ResponseJsonEncoder
-          )))
-        interaction['response']['sdk_response_segments'] = segments
       f.write(
           json.dumps(
-              replay_session_dict, indent=2, cls=RequestJsonEncoder
+              self.replay_session.model_dump(mode='json'), indent=2, cls=ResponseJsonEncoder
           )
       )
     self.replay_session = None
@@ -371,15 +367,8 @@ class ReplayApiClient(ApiClient):
     if isinstance(response_model, list):
       response_model = response_model[0]
     print('response_model: ', response_model.model_dump(exclude_none=True))
-    actual = json.dumps(
-        response_model.model_dump(exclude_none=True),
-        cls=ResponseJsonEncoder,
-        sort_keys=True,
-    )
-    expected = json.dumps(
-        interaction.response.sdk_response_segments[self._sdk_response_index],
-        sort_keys=True,
-    )
+    actual = response_model.model_dump(exclude_none=True, mode='json')
+    expected = interaction.response.sdk_response_segments[self._sdk_response_index]
     assert (
         actual == expected
     ), f'SDK response mismatch:\nActual: {actual}\nExpected: {expected}'
@@ -432,36 +421,12 @@ class ReplayApiClient(ApiClient):
       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.
-
-  We need RequestJsonEncoder and ResponseJsonEncoder because:
-    1. In production, we only need RequestJsonEncoder to help json module
-    to convert non-stringable and stringable types to json string. Especially
-    for bytes type, the value of bytes field is encoded to base64 string so it
-    is always stringable and the RequestJsonEncoder doesn't have to deal with
-    utf-8 JSON broken issue.
-    2. In replay test, we also need ResponseJsonEncoder to help json module
-    convert non-stringable and stringable types to json string. But response
-    object returned from SDK method is different from the request api_client
-    sent to server. For the bytes type, there is no base64 string in response
-    anymore, because SDK handles it internally. So bytes type in Response is
-    non-stringable. The ResponseJsonEncoder uses different encoding
-    strategy than the RequestJsonEncoder to deal with utf-8 JSON broken issue.
   """
   def default(self, o):
-    if isinstance(o, bytes):
-      # Use base64.b64encode() to encode bytes to string so that the media bytes
-      # fields are serializable.
-      # o.decode(encoding='utf-8', errors='replace') doesn't work because it
-      # uses a fixed error string `\ufffd` for all non-utf-8 characters,
-      # which cannot be converted back to original bytes. And other languages
-      # only have the original bytes to compare with.
-      # Since we use base64.b64encoding() in replay test, a change that breaks
-      # native bytes can be captured by
-      # test_compute_tokens.py::test_token_bytes_deserialization.
-      return base64.b64encode(o).decode(encoding='utf-8')
-    elif isinstance(o, datetime.datetime):
+    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'):

google/genai/_transformers.py

@@ -142,6 +142,30 @@ def t_model(client: _api_client.ApiClient, model: str):
     else:
       return f'models/{model}'
 
+def t_models_url(api_client: _api_client.ApiClient, base_models: bool) -> str:
+  if api_client.vertexai:
+    if base_models:
+      return 'publishers/google/models'
+    else:
+      return 'models'
+  else:
+    if base_models:
+      return 'models'
+    else:
+      return 'tunedModels'
+
+
+def t_extract_models(api_client: _api_client.ApiClient, response: dict) -> list[types.Model]:
+  if response.get('models') is not None:
+    return response.get('models')
+  elif response.get('tunedModels') is not None:
+    return response.get('tunedModels')
+  elif response.get('publisherModels') is not None:
+    return response.get('publisherModels')
+  else:
+    raise ValueError('Cannot determine the models type.')
+
+
 def t_caches_model(api_client: _api_client.ApiClient, model: str):
   model = t_model(api_client, model)
   if not model:
@@ -452,3 +476,17 @@ def t_tuning_job_status(
     return 'JOB_STATE_FAILED'
   else:
     return status
+
+
+# Some fields don't accept url safe base64 encoding.
+# We shouldn't use this transformer if the backend adhere to Cloud Type
+# format https://cloud.google.com/docs/discovery/type-format.
+# TODO(b/389133914): Remove the hack after Vertex backend fix the issue.
+def t_bytes(api_client: _api_client.ApiClient, data: bytes) -> str:
+  if not isinstance(data, bytes):
+    return data
+  if api_client.vertexai:
+    return base64.b64encode(data).decode('ascii')
+  else:
+    return base64.urlsafe_encode(data).decode('ascii')
+