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

google/genai/_api_client.py

@@ -99,6 +99,19 @@ class HttpRequest:
   timeout: Optional[float] = None
 
 
+# TODO(b/394358912): Update this class to use a SDKResponse class that can be
+# generated and used for all languages.
+@dataclass
+class BaseResponse:
+  http_headers: dict[str, str]
+
+  @property
+  def dict(self) -> dict[str, Any]:
+    if isinstance(self, dict):
+      return self
+    return {'httpHeaders': self.http_headers}
+
+
 class HttpResponse:
 
   def __init__(
@@ -255,7 +268,7 @@ class ApiClient:
             'Project and location or API key must be set when using the Vertex '
             'AI API.'
         )
-      if self.api_key:
+      if self.api_key or self.location == 'global':
         self._http_options['base_url'] = (
             f'https://aiplatform.googleapis.com/'
         )
@@ -273,7 +286,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 and not self.vertexai:
+    if self.api_key:
       self._http_options['headers']['x-goog-api-key'] = self.api_key
     # Update the http options with the user provided http options.
     if http_options:
@@ -323,8 +336,6 @@ class ApiClient:
         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,
@@ -436,18 +447,12 @@ class ApiClient:
         http_method, path, request_dict, http_options
     )
     response = self._request(http_request, stream=False)
-    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
+    json_response = response.json
+    if not json_response:
+      base_response = BaseResponse(response.headers).dict
+      return base_response
+
+    return json_response
 
   def request_streamed(
       self,
@@ -461,10 +466,6 @@ class ApiClient:
     )
 
     session_response = self._request(http_request, stream=True)
-    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
 
@@ -480,9 +481,11 @@ class ApiClient:
     )
 
     result = await self._async_request(http_request=http_request, stream=False)
-    if http_options and 'deprecated_response_payload' in http_options:
-      result._copy_to_dict(http_options['deprecated_response_payload'])
-    return result.json
+    json_response = result.json
+    if not json_response:
+      base_response = BaseResponse(result.headers).dict
+      return base_response
+    return json_response
 
   async def async_request_streamed(
       self,
@@ -497,8 +500,6 @@ class ApiClient:
 
     response = await self._async_request(http_request=http_request, stream=True)
 
-    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

google/genai/_automatic_function_calling_util.py

@@ -17,14 +17,18 @@ import inspect
 import sys
 import types as builtin_types
 import typing
-from typing import Any, Callable, Literal, Union, _GenericAlias, get_args, get_origin
+from typing import _GenericAlias, Any, Callable, get_args, get_origin, Literal, Union
+
 import pydantic
+
+from . import _extra_utils
 from . import types
 
+
 if sys.version_info >= (3, 10):
-  UnionType = builtin_types.UnionType
+  VersionedUnionType = builtin_types.UnionType
 else:
-  UnionType = typing._UnionGenericAlias
+  VersionedUnionType = typing._UnionGenericAlias
 
 _py_builtin_type_to_schema_type = {
     str: 'STRING',
@@ -45,7 +49,8 @@ def _is_builtin_primitive_or_compound(
 def _raise_for_any_of_if_mldev(schema: types.Schema):
   if schema.any_of:
     raise ValueError(
-        'AnyOf is not supported in function declaration schema for Google AI.'
+        'AnyOf is not supported in function declaration schema for'
+        ' the Gemini API.'
     )
 
 
@@ -53,15 +58,7 @@ def _raise_for_default_if_mldev(schema: types.Schema):
   if schema.default is not None:
     raise ValueError(
         'Default value is not supported in function declaration schema for'
-        ' Google AI.'
-    )
-
-
-def _raise_for_nullable_if_mldev(schema: types.Schema):
-  if schema.nullable:
-    raise ValueError(
-        'Nullable is not supported in function declaration schema for'
-        ' Google AI.'
+        ' the Gemini API.'
     )
 
 
@@ -69,7 +66,6 @@ def _raise_if_schema_unsupported(client, schema: types.Schema):
   if not client.vertexai:
     _raise_for_any_of_if_mldev(schema)
     _raise_for_default_if_mldev(schema)
-    _raise_for_nullable_if_mldev(schema)
 
 
 def _is_default_value_compatible(
@@ -82,10 +78,10 @@ def _is_default_value_compatible(
   if (
       isinstance(annotation, _GenericAlias)
       or isinstance(annotation, builtin_types.GenericAlias)
-      or isinstance(annotation, UnionType)
+      or isinstance(annotation, VersionedUnionType)
   ):
     origin = get_origin(annotation)
-    if origin in (Union, UnionType):
+    if origin in (Union, VersionedUnionType):
       return any(
           _is_default_value_compatible(default_value, arg)
           for arg in get_args(annotation)
@@ -141,7 +137,7 @@ def _parse_schema_from_parameter(
     _raise_if_schema_unsupported(client, schema)
     return schema
   if (
-      isinstance(param.annotation, UnionType)
+      isinstance(param.annotation, VersionedUnionType)
       # only parse simple UnionType, example int | str | float | bool
       # complex UnionType will be invoked in raise branch
       and all(
@@ -229,7 +225,11 @@ def _parse_schema_from_parameter(
       schema.type = 'OBJECT'
       unique_types = set()
       for arg in args:
-        if arg.__name__ == 'NoneType':  # Optional type
+        # The first check is for NoneType in Python 3.9, since the __name__
+        # attribute is not available in Python 3.9
+        if type(arg) is type(None) or (
+            hasattr(arg, '__name__') and arg.__name__ == 'NoneType'
+        ):  # Optional type
           schema.nullable = True
           continue
         schema_in_any_of = _parse_schema_from_parameter(
@@ -272,9 +272,8 @@ def _parse_schema_from_parameter(
       return schema
       # all other generic alias will be invoked in raise branch
   if (
-      inspect.isclass(param.annotation)
       # for user defined class, we only support pydantic model
-      and issubclass(param.annotation, pydantic.BaseModel)
+      _extra_utils.is_annotation_pydantic_model(param.annotation)
   ):
     if (
         param.default is not inspect.Parameter.empty

google/genai/_common.py

@@ -18,14 +18,17 @@
 import base64
 import datetime
 import enum
+import functools
 import typing
 from typing import Union
 import uuid
+import warnings
 
 import pydantic
 from pydantic import alias_generators
 
 from . import _api_client
+from . import errors
 
 
 def set_value_by_path(data, keys, value):
@@ -213,8 +216,16 @@ class CaseInSensitiveEnum(str, enum.Enum):
     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
+      except KeyError:
+        warnings.warn(f"{value} is not a valid {cls.__name__}")
+        try:
+          # Creating a enum instance based on the value
+          unknown_enum_val = cls._new_member_(cls)  # pylint: disable=protected-access,attribute-error
+          unknown_enum_val._name_ = str(value)  # pylint: disable=protected-access
+          unknown_enum_val._value_ = value  # pylint: disable=protected-access
+          return unknown_enum_val
+        except:
+          return None
 
 
 def timestamped_unique_name() -> str:
@@ -264,3 +275,23 @@ def encode_unserializable_types(data: dict[str, object]) -> dict[str, object]:
     else:
       processed_data[key] = value
   return processed_data
+
+
+def experimental_warning(message: str):
+  """Experimental warning, only warns once."""
+  def decorator(func):
+    warning_done = False
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+      nonlocal warning_done
+      if not warning_done:
+        warning_done = True
+        warnings.warn(
+            message=message,
+            category=errors.ExperimentalWarning,
+            stacklevel=2,
+        )
+      return func(*args, **kwargs)
+    return wrapper
+  return decorator
+

google/genai/_extra_utils.py

@@ -108,16 +108,22 @@ def convert_number_values_for_function_call_args(
   return args
 
 
-def _is_annotation_pydantic_model(annotation: Any) -> bool:
-  return inspect.isclass(annotation) and issubclass(
-      annotation, pydantic.BaseModel
-  )
+def is_annotation_pydantic_model(annotation: Any) -> bool:
+  try:
+    return inspect.isclass(annotation) and issubclass(
+        annotation, pydantic.BaseModel
+    )
+  # for python 3.10 and below, inspect.isclass(annotation) has inconsistent
+  # results with versions above. for example, inspect.isclass(dict[str, int]) is
+  # True in 3.10 and below but False in 3.11 and above.
+  except TypeError:
+    return False
 
 
 def convert_if_exist_pydantic_model(
     value: Any, annotation: Any, param_name: str, func_name: str
 ) -> Any:
-  if isinstance(value, dict) and _is_annotation_pydantic_model(annotation):
+  if isinstance(value, dict) and is_annotation_pydantic_model(annotation):
     try:
       return annotation(**value)
     except pydantic.ValidationError as e:
@@ -146,7 +152,7 @@ def convert_if_exist_pydantic_model(
       if (
           (get_args(arg) and get_origin(arg) is list)
           or isinstance(value, arg)
-          or (isinstance(value, dict) and _is_annotation_pydantic_model(arg))
+          or (isinstance(value, dict) and is_annotation_pydantic_model(arg))
       ):
         try:
           return convert_if_exist_pydantic_model(

google/genai/_operations.py

@@ -0,0 +1,365 @@
+# 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.
+#
+
+# Code generated by the Google Gen AI SDK generator DO NOT EDIT.
+
+from typing import Optional, Union
+from urllib.parse import urlencode
+from . import _api_module
+from . import _common
+from . import types
+from ._api_client import ApiClient
+from ._common import get_value_by_path as getv
+from ._common import set_value_by_path as setv
+
+
+def _GetOperationParameters_to_mldev(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['operation_name']) is not None:
+    setv(
+        to_object,
+        ['_url', 'operationName'],
+        getv(from_object, ['operation_name']),
+    )
+
+  if getv(from_object, ['config']) is not None:
+    setv(to_object, ['config'], getv(from_object, ['config']))
+
+  return to_object
+
+
+def _GetOperationParameters_to_vertex(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['operation_name']) is not None:
+    setv(
+        to_object,
+        ['_url', 'operationName'],
+        getv(from_object, ['operation_name']),
+    )
+
+  if getv(from_object, ['config']) is not None:
+    setv(to_object, ['config'], getv(from_object, ['config']))
+
+  return to_object
+
+
+def _FetchPredictOperationParameters_to_mldev(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['operation_name']) is not None:
+    raise ValueError('operation_name parameter is not supported in Gemini API.')
+
+  if getv(from_object, ['resource_name']) is not None:
+    raise ValueError('resource_name parameter is not supported in Gemini API.')
+
+  if getv(from_object, ['config']) is not None:
+    raise ValueError('config parameter is not supported in Gemini API.')
+
+  return to_object
+
+
+def _FetchPredictOperationParameters_to_vertex(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['operation_name']) is not None:
+    setv(to_object, ['operationName'], getv(from_object, ['operation_name']))
+
+  if getv(from_object, ['resource_name']) is not None:
+    setv(
+        to_object,
+        ['_url', 'resourceName'],
+        getv(from_object, ['resource_name']),
+    )
+
+  if getv(from_object, ['config']) is not None:
+    setv(to_object, ['config'], getv(from_object, ['config']))
+
+  return to_object
+
+
+def _Operation_from_mldev(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['name']) is not None:
+    setv(to_object, ['name'], getv(from_object, ['name']))
+
+  if getv(from_object, ['metadata']) is not None:
+    setv(to_object, ['metadata'], getv(from_object, ['metadata']))
+
+  if getv(from_object, ['done']) is not None:
+    setv(to_object, ['done'], getv(from_object, ['done']))
+
+  if getv(from_object, ['error']) is not None:
+    setv(to_object, ['error'], getv(from_object, ['error']))
+
+  if getv(from_object, ['response']) is not None:
+    setv(to_object, ['response'], getv(from_object, ['response']))
+
+  return to_object
+
+
+def _Operation_from_vertex(
+    api_client: ApiClient,
+    from_object: Union[dict, object],
+    parent_object: dict = None,
+) -> dict:
+  to_object = {}
+  if getv(from_object, ['name']) is not None:
+    setv(to_object, ['name'], getv(from_object, ['name']))
+
+  if getv(from_object, ['metadata']) is not None:
+    setv(to_object, ['metadata'], getv(from_object, ['metadata']))
+
+  if getv(from_object, ['done']) is not None:
+    setv(to_object, ['done'], getv(from_object, ['done']))
+
+  if getv(from_object, ['error']) is not None:
+    setv(to_object, ['error'], getv(from_object, ['error']))
+
+  if getv(from_object, ['response']) is not None:
+    setv(to_object, ['response'], getv(from_object, ['response']))
+
+  return to_object
+
+
+class _operations(_api_module.BaseModule):
+
+  def _get_operation(
+      self,
+      *,
+      operation_name: str,
+      config: Optional[types.GetOperationConfigOrDict] = None,
+  ) -> types.Operation:
+    parameter_model = types._GetOperationParameters(
+        operation_name=operation_name,
+        config=config,
+    )
+
+    if self._api_client.vertexai:
+      request_dict = _GetOperationParameters_to_vertex(
+          self._api_client, parameter_model
+      )
+      path = '{operationName}'.format_map(request_dict.get('_url'))
+    else:
+      request_dict = _GetOperationParameters_to_mldev(
+          self._api_client, parameter_model
+      )
+      path = '{operationName}'.format_map(request_dict.get('_url'))
+    query_params = request_dict.get('_query')
+    if query_params:
+      path = f'{path}?{urlencode(query_params)}'
+    # TODO: remove the hack that pops config.
+    request_dict.pop('config', None)
+
+    http_options = None
+    if isinstance(config, dict):
+      http_options = config.get('http_options', None)
+    elif hasattr(config, 'http_options'):
+      http_options = config.http_options
+
+    request_dict = _common.convert_to_dict(request_dict)
+    request_dict = _common.encode_unserializable_types(request_dict)
+
+    response_dict = self._api_client.request(
+        'get', path, request_dict, http_options
+    )
+
+    if self._api_client.vertexai:
+      response_dict = _Operation_from_vertex(self._api_client, response_dict)
+    else:
+      response_dict = _Operation_from_mldev(self._api_client, response_dict)
+
+    return_value = types.Operation._from_response(
+        response=response_dict, kwargs=parameter_model
+    )
+    self._api_client._verify_response(return_value)
+    return return_value
+
+  def _fetch_predict_operation(
+      self,
+      *,
+      operation_name: str,
+      resource_name: str,
+      config: Optional[types.FetchPredictOperationConfigOrDict] = None,
+  ) -> types.Operation:
+    parameter_model = types._FetchPredictOperationParameters(
+        operation_name=operation_name,
+        resource_name=resource_name,
+        config=config,
+    )
+
+    if not self._api_client.vertexai:
+      raise ValueError('This method is only supported in the Vertex AI client.')
+    else:
+      request_dict = _FetchPredictOperationParameters_to_vertex(
+          self._api_client, parameter_model
+      )
+      path = '{resourceName}:fetchPredictOperation'.format_map(
+          request_dict.get('_url')
+      )
+
+    query_params = request_dict.get('_query')
+    if query_params:
+      path = f'{path}?{urlencode(query_params)}'
+    # TODO: remove the hack that pops config.
+    request_dict.pop('config', None)
+
+    http_options = None
+    if isinstance(config, dict):
+      http_options = config.get('http_options', None)
+    elif hasattr(config, 'http_options'):
+      http_options = config.http_options
+
+    request_dict = _common.convert_to_dict(request_dict)
+    request_dict = _common.encode_unserializable_types(request_dict)
+
+    response_dict = self._api_client.request(
+        'post', path, request_dict, http_options
+    )
+
+    if self._api_client.vertexai:
+      response_dict = _Operation_from_vertex(self._api_client, response_dict)
+    else:
+      response_dict = _Operation_from_mldev(self._api_client, response_dict)
+
+    return_value = types.Operation._from_response(
+        response=response_dict, kwargs=parameter_model
+    )
+    self._api_client._verify_response(return_value)
+    return return_value
+
+
+class Async_operations(_api_module.BaseModule):
+
+  async def _get_operation(
+      self,
+      *,
+      operation_name: str,
+      config: Optional[types.GetOperationConfigOrDict] = None,
+  ) -> types.Operation:
+    parameter_model = types._GetOperationParameters(
+        operation_name=operation_name,
+        config=config,
+    )
+
+    if self._api_client.vertexai:
+      request_dict = _GetOperationParameters_to_vertex(
+          self._api_client, parameter_model
+      )
+      path = '{operationName}'.format_map(request_dict.get('_url'))
+    else:
+      request_dict = _GetOperationParameters_to_mldev(
+          self._api_client, parameter_model
+      )
+      path = '{operationName}'.format_map(request_dict.get('_url'))
+    query_params = request_dict.get('_query')
+    if query_params:
+      path = f'{path}?{urlencode(query_params)}'
+    # TODO: remove the hack that pops config.
+    request_dict.pop('config', None)
+
+    http_options = None
+    if isinstance(config, dict):
+      http_options = config.get('http_options', None)
+    elif hasattr(config, 'http_options'):
+      http_options = config.http_options
+
+    request_dict = _common.convert_to_dict(request_dict)
+    request_dict = _common.encode_unserializable_types(request_dict)
+
+    response_dict = await self._api_client.async_request(
+        'get', path, request_dict, http_options
+    )
+
+    if self._api_client.vertexai:
+      response_dict = _Operation_from_vertex(self._api_client, response_dict)
+    else:
+      response_dict = _Operation_from_mldev(self._api_client, response_dict)
+
+    return_value = types.Operation._from_response(
+        response=response_dict, kwargs=parameter_model
+    )
+    self._api_client._verify_response(return_value)
+    return return_value
+
+  async def _fetch_predict_operation(
+      self,
+      *,
+      operation_name: str,
+      resource_name: str,
+      config: Optional[types.FetchPredictOperationConfigOrDict] = None,
+  ) -> types.Operation:
+    parameter_model = types._FetchPredictOperationParameters(
+        operation_name=operation_name,
+        resource_name=resource_name,
+        config=config,
+    )
+
+    if not self._api_client.vertexai:
+      raise ValueError('This method is only supported in the Vertex AI client.')
+    else:
+      request_dict = _FetchPredictOperationParameters_to_vertex(
+          self._api_client, parameter_model
+      )
+      path = '{resourceName}:fetchPredictOperation'.format_map(
+          request_dict.get('_url')
+      )
+
+    query_params = request_dict.get('_query')
+    if query_params:
+      path = f'{path}?{urlencode(query_params)}'
+    # TODO: remove the hack that pops config.
+    request_dict.pop('config', None)
+
+    http_options = None
+    if isinstance(config, dict):
+      http_options = config.get('http_options', None)
+    elif hasattr(config, 'http_options'):
+      http_options = config.http_options
+
+    request_dict = _common.convert_to_dict(request_dict)
+    request_dict = _common.encode_unserializable_types(request_dict)
+
+    response_dict = await self._api_client.async_request(
+        'post', path, request_dict, http_options
+    )
+
+    if self._api_client.vertexai:
+      response_dict = _Operation_from_vertex(self._api_client, response_dict)
+    else:
+      response_dict = _Operation_from_mldev(self._api_client, response_dict)
+
+    return_value = types.Operation._from_response(
+        response=response_dict, kwargs=parameter_model
+    )
+    self._api_client._verify_response(return_value)
+    return return_value

google/genai/_replay_api_client.py

@@ -78,6 +78,11 @@ def _redact_request_url(url: str) -> str:
       '{VERTEX_URL_PREFIX}/',
       result,
   )
+  result = re.sub(
+      r'.*aiplatform.googleapis.com/[^/]+/',
+      '{VERTEX_URL_PREFIX}/',
+      result,
+  )
   result = re.sub(
       r'https://generativelanguage.googleapis.com/[^/]+',
       '{MLDEV_URL_PREFIX}',
@@ -357,6 +362,8 @@ class ReplayApiClient(ApiClient):
     if self._should_update_replay():
       if isinstance(response_model, list):
         response_model = response_model[0]
+      if response_model and 'http_headers' in response_model.model_fields:
+        response_model.http_headers.pop('Date', None)
       interaction.response.sdk_response_segments.append(
           response_model.model_dump(exclude_none=True)
       )