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

google/genai/__init__.py

@@ -16,7 +16,8 @@
 """Google Gen AI SDK"""
 
 from .client import Client
+from . import  version
 
-__version__ = '0.3.0'
+__version__ = version.__version__
 
 __all__ = ['Client']

google/genai/_api_client.py

@@ -21,37 +21,74 @@ import copy
 from dataclasses import dataclass
 import datetime
 import json
+import logging
 import os
 import sys
-from typing import Any, Optional, TypedDict, Union
+from typing import Any, Optional, Tuple, TypedDict, Union
 from urllib.parse import urlparse, urlunparse
 
 import google.auth
 import google.auth.credentials
 from google.auth.transport.requests import AuthorizedSession
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
 import requests
 
 from . import errors
+from . import version
 
 
-class HttpOptions(TypedDict):
+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: str = None
+  base_url: Optional[str] = None
   """The base URL for the AI platform service endpoint."""
-  api_version: str = None
+  api_version: Optional[str] = None
   """Specifies the version of the API to use."""
-  headers: dict[str, dict] = None
+  headers: Optional[dict[str, str]] = None
   """Additional HTTP headers to be sent with the request."""
-  response_payload: dict = None
+  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]
 
 
 def _append_library_version_headers(headers: dict[str, str]) -> None:
   """Appends the telemetry header to the headers dict."""
-  # TODO: Automate revisions to the SDK library version.
-  library_label = f'google-genai-sdk/0.3.0'
+  library_label = f'google-genai-sdk/{version.__version__}'
   language_label = 'gl-python/' + sys.version.split()[0]
   version_header_value = f'{library_label} {language_label}'
   if (
@@ -71,20 +108,24 @@ def _append_library_version_headers(headers: dict[str, str]) -> None:
 
 
 def _patch_http_options(
-    options: HttpOptions, patch_options: HttpOptions
-) -> HttpOptions:
+    options: HttpOptionsDict, patch_options: HttpOptionsDict
+) -> HttpOptionsDict:
   # use shallow copy so we don't override the original objects.
-  copy_option = HttpOptions()
+  copy_option = HttpOptionsDict()
   copy_option.update(options)
-  for k, v in patch_options.items():
+  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(v, dict) and isinstance(copy_option.get(k, None), dict):
-      copy_option[k] = {}
-      copy_option[k].update(options[k])  # shallow copy from original options.
-      copy_option[k].update(v)
-    elif v is not None:  # Accept empty values.
-      copy_option[k] = v
+    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
   _append_library_version_headers(copy_option['headers'])
   return copy_option
 
@@ -102,6 +143,7 @@ class HttpRequest:
   url: str
   method: str
   data: Union[dict[str, object], bytes]
+  timeout: Optional[Union[float, Tuple[float, float]]] = None
 
 
 class HttpResponse:
@@ -147,7 +189,7 @@ class ApiClient:
       credentials: google.auth.credentials.Credentials = None,
       project: Union[str, None] = None,
       location: Union[str, None] = None,
-      http_options: HttpOptions = None,
+      http_options: HttpOptionsOrDict = None,
   ):
     self.vertexai = vertexai
     if self.vertexai is None:
@@ -159,30 +201,84 @@ 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):
+      try:
+        HttpOptions.model_validate(http_options)
+      except ValidationError as e:
+        raise ValueError(f'Invalid http_options: {e}')
+    elif(isinstance(http_options, HttpOptions)):
+      http_options = http_options.model_dump()
+
+    # 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.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)
     self._credentials = credentials
-    self._http_options = HttpOptions()
+    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'] = (
@@ -191,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:
@@ -208,7 +304,7 @@ class ApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: HttpOptions = None,
+      http_options: HttpOptionsDict = 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('_')]
@@ -221,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,
@@ -232,6 +338,7 @@ class ApiClient:
         url=url,
         headers=patched_http_options['headers'],
         data=request_dict,
+        timeout=patched_http_options.get('timeout', None),
     )
 
   def _request(
@@ -239,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"],
@@ -250,10 +357,10 @@ class ApiClient:
           http_request.method.upper(),
           http_request.url,
           headers=http_request.headers,
-          data=json.dumps(http_request.data, cls=RequestJsonEncoder) if http_request.data else None,
-          # TODO: support timeout in RequestOptions so it can be configured
-          # per methods.
-          timeout=None,
+          data=json.dumps(http_request.data, cls=RequestJsonEncoder)
+          if http_request.data
+          else None,
+          timeout=http_request.timeout,
       )
       errors.APIError.raise_for_response(response)
       return HttpResponse(
@@ -275,13 +382,14 @@ class ApiClient:
         data = http_request.data
 
     http_session = requests.Session()
-    request = requests.Request(
+    response = http_session.request(
         method=http_request.method,
         url=http_request.url,
         headers=http_request.headers,
         data=data,
-    ).prepare()
-    response = http_session.send(request, stream=stream)
+        timeout=http_request.timeout,
+        stream=stream,
+    )
     errors.APIError.raise_for_response(response)
     return HttpResponse(
         response.headers, response if stream else [response.text]
@@ -307,8 +415,10 @@ class ApiClient:
           stream=stream,
       )
 
-  def get_read_only_http_options(self) -> HttpOptions:
-    copied = HttpOptions()
+  def get_read_only_http_options(self) -> HttpOptionsDict:
+    copied = HttpOptionsDict()
+    if isinstance(self._http_options, BaseModel):
+      self._http_options = self._http_options.model_dump()
     copied.update(self._http_options)
     return copied
 
@@ -317,7 +427,7 @@ class ApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: HttpOptions = None,
+      http_options: HttpOptionsDict = None,
   ):
     http_request = self._build_request(
         http_method, path, request_dict, http_options
@@ -332,7 +442,7 @@ class ApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: HttpOptions = None,
+      http_options: HttpOptionsDict = None,
   ):
     http_request = self._build_request(
         http_method, path, request_dict, http_options
@@ -349,7 +459,7 @@ class ApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: HttpOptions = None,
+      http_options: HttpOptionsDict = None,
   ) -> dict[str, object]:
     http_request = self._build_request(
         http_method, path, request_dict, http_options
@@ -365,7 +475,7 @@ class ApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: HttpOptions = None,
+      http_options: HttpOptionsDict = None,
   ):
     http_request = self._build_request(
         http_method, path, request_dict, http_options
@@ -464,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/_automatic_function_calling_util.py

@@ -58,8 +58,8 @@ def _raise_for_nullable_if_mldev(schema: types.Schema):
     )
 
 
-def _raise_if_schema_unsupported(client, schema: types.Schema):
-  if not client.vertexai:
+def _raise_if_schema_unsupported(variant: str, schema: types.Schema):
+  if not variant == 'VERTEX_AI':
     _raise_for_any_of_if_mldev(schema)
     _raise_for_default_if_mldev(schema)
     _raise_for_nullable_if_mldev(schema)
@@ -112,7 +112,7 @@ def _is_default_value_compatible(
 
 
 def _parse_schema_from_parameter(
-    client, param: inspect.Parameter, func_name: str
+    variant: str, param: inspect.Parameter, func_name: str
 ) -> types.Schema:
   """parse schema from parameter.
 
@@ -130,7 +130,7 @@ def _parse_schema_from_parameter(
         raise ValueError(default_value_error_msg)
       schema.default = param.default
     schema.type = _py_builtin_type_to_schema_type[param.annotation]
-    _raise_if_schema_unsupported(client, schema)
+    _raise_if_schema_unsupported(variant, schema)
     return schema
   if (
       isinstance(param.annotation, typing_types.UnionType)
@@ -149,7 +149,7 @@ def _parse_schema_from_parameter(
         schema.nullable = True
         continue
       schema_in_any_of = _parse_schema_from_parameter(
-          client,
+          variant,
           inspect.Parameter(
               'item', inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=arg
           ),
@@ -171,7 +171,7 @@ def _parse_schema_from_parameter(
       if not _is_default_value_compatible(param.default, param.annotation):
         raise ValueError(default_value_error_msg)
       schema.default = param.default
-    _raise_if_schema_unsupported(client, schema)
+    _raise_if_schema_unsupported(variant, schema)
     return schema
   if isinstance(param.annotation, _GenericAlias) or isinstance(
       param.annotation, typing_types.GenericAlias
@@ -184,7 +184,7 @@ def _parse_schema_from_parameter(
         if not _is_default_value_compatible(param.default, param.annotation):
           raise ValueError(default_value_error_msg)
         schema.default = param.default
-      _raise_if_schema_unsupported(client, schema)
+      _raise_if_schema_unsupported(variant, schema)
       return schema
     if origin is Literal:
       if not all(isinstance(arg, str) for arg in args):
@@ -197,12 +197,12 @@ def _parse_schema_from_parameter(
         if not _is_default_value_compatible(param.default, param.annotation):
           raise ValueError(default_value_error_msg)
         schema.default = param.default
-      _raise_if_schema_unsupported(client, schema)
+      _raise_if_schema_unsupported(variant, schema)
       return schema
     if origin is list:
       schema.type = 'ARRAY'
       schema.items = _parse_schema_from_parameter(
-          client,
+          variant,
           inspect.Parameter(
               'item',
               inspect.Parameter.POSITIONAL_OR_KEYWORD,
@@ -214,7 +214,7 @@ def _parse_schema_from_parameter(
         if not _is_default_value_compatible(param.default, param.annotation):
           raise ValueError(default_value_error_msg)
         schema.default = param.default
-      _raise_if_schema_unsupported(client, schema)
+      _raise_if_schema_unsupported(variant, schema)
       return schema
     if origin is Union:
       schema.any_of = []
@@ -225,7 +225,7 @@ def _parse_schema_from_parameter(
           schema.nullable = True
           continue
         schema_in_any_of = _parse_schema_from_parameter(
-            client,
+            variant,
             inspect.Parameter(
                 'item',
                 inspect.Parameter.POSITIONAL_OR_KEYWORD,
@@ -249,7 +249,7 @@ def _parse_schema_from_parameter(
         if not _is_default_value_compatible(param.default, param.annotation):
           raise ValueError(default_value_error_msg)
         schema.default = param.default
-      _raise_if_schema_unsupported(client, schema)
+      _raise_if_schema_unsupported(variant, schema)
       return schema
       # all other generic alias will be invoked in raise branch
   if (
@@ -266,7 +266,7 @@ def _parse_schema_from_parameter(
     schema.properties = {}
     for field_name, field_info in param.annotation.model_fields.items():
       schema.properties[field_name] = _parse_schema_from_parameter(
-          client,
+          variant,
           inspect.Parameter(
               field_name,
               inspect.Parameter.POSITIONAL_OR_KEYWORD,
@@ -274,7 +274,7 @@ def _parse_schema_from_parameter(
           ),
           func_name,
       )
-    _raise_if_schema_unsupported(client, schema)
+    _raise_if_schema_unsupported(variant, schema)
     return schema
   raise ValueError(
       f'Failed to parse the parameter {param} of function {func_name} for'

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
-