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

google/genai/_api_client.py

@@ -36,55 +36,7 @@ import requests
 
 from . import errors
 from . import version
-
-
-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: Optional[str] = None
-  """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, 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]
+from .types import HttpOptions, HttpOptionsDict, HttpOptionsOrDict
 
 
 def _append_library_version_headers(headers: dict[str, str]) -> None:
@@ -144,7 +96,7 @@ class HttpRequest:
   url: str
   method: str
   data: Union[dict[str, object], bytes]
-  timeout: Optional[Union[float, Tuple[float, float]]] = None
+  timeout: Optional[float] = None
 
 
 class HttpResponse:
@@ -159,9 +111,20 @@ class HttpResponse:
     self.headers = headers
     self.response_stream = response_stream
     self.byte_stream = byte_stream
+    self.segment_iterator = self.segments()
+
+  # Async iterator for async streaming.
+  def __aiter__(self):
+    return self
+
+  async def __anext__(self):
+    try:
+      return next(self.segment_iterator)
+    except StopIteration:
+      raise StopAsyncIteration
 
   @property
-  def text(self) -> str:
+  def json(self) -> Any:
     if not self.response_stream[0]:  # Empty response
       return ''
     return json.loads(self.response_stream[0])
@@ -194,7 +157,9 @@ class HttpResponse:
           'Byte segments are not supported for streaming responses.'
       )
 
-  def copy_to_dict(self, response_payload: dict[str, object]):
+  def _copy_to_dict(self, response_payload: dict[str, object]):
+    # Cannot pickle 'generator' object.
+    delattr(self, 'segment_iterator')
     for attribute in dir(self):
       response_payload[attribute] = copy.deepcopy(getattr(self, attribute))
 
@@ -290,7 +255,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/'
         )
@@ -308,7 +273,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:
@@ -325,7 +290,7 @@ class ApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: HttpOptionsDict = None,
+      http_options: HttpOptionsOrDict = 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('_')]
@@ -333,33 +298,48 @@ class ApiClient:
       del request_dict[key]
     # patch the http options with the user provided settings.
     if http_options:
-      patched_http_options = _patch_http_options(
-          self._http_options, http_options
-      )
+      if isinstance(http_options, HttpOptions):
+        patched_http_options = _patch_http_options(
+            self._http_options, http_options.model_dump()
+        )
+      else:
+        patched_http_options = _patch_http_options(
+            self._http_options, http_options
+        )
     else:
       patched_http_options = self._http_options
-    skip_project_and_location_in_path_val = patched_http_options.get(
-        'skip_project_and_location_in_path', False
-    )
+    # Skip adding project and locations when getting Vertex AI base models.
+    query_vertex_base_models = False
+    if (
+        self.vertexai
+        and http_method == 'get'
+        and path.startswith('publishers/google/models')
+    ):
+      query_vertex_base_models = True
     if (
         self.vertexai
         and not path.startswith('projects/')
-        and not skip_project_and_location_in_path_val
+        and not query_vertex_base_models
         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,
     )
+
+    timeout_in_seconds = patched_http_options.get('timeout', None)
+    if timeout_in_seconds:
+      timeout_in_seconds = timeout_in_seconds / 1000.0
+    else:
+      timeout_in_seconds = None
+
     return HttpRequest(
         method=http_method,
         url=url,
         headers=patched_http_options['headers'],
         data=request_dict,
-        timeout=patched_http_options.get('timeout', None),
+        timeout=timeout_in_seconds,
     )
 
   def _request(
@@ -448,15 +428,24 @@ class ApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: HttpOptionsDict = None,
+      http_options: HttpOptionsOrDict = None,
   ):
     http_request = self._build_request(
         http_method, path, request_dict, http_options
     )
     response = self._request(http_request, stream=False)
-    if http_options and 'response_payload' in http_options:
-      response.copy_to_dict(http_options['response_payload'])
-    return response.text
+    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
 
   def request_streamed(
       self,
@@ -470,8 +459,10 @@ class ApiClient:
     )
 
     session_response = self._request(http_request, stream=True)
-    if http_options and 'response_payload' in http_options:
-      session_response.copy_to_dict(http_options['response_payload'])
+    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
 
@@ -487,9 +478,9 @@ class ApiClient:
     )
 
     result = await self._async_request(http_request=http_request, stream=False)
-    if http_options and 'response_payload' in http_options:
-      result.copy_to_dict(http_options['response_payload'])
-    return result.text
+    if http_options and 'deprecated_response_payload' in http_options:
+      result._copy_to_dict(http_options['deprecated_response_payload'])
+    return result.json
 
   async def async_request_streamed(
       self,
@@ -504,10 +495,12 @@ class ApiClient:
 
     response = await self._async_request(http_request=http_request, stream=True)
 
-    for chunk in response.segments():
-      yield chunk
-    if http_options and 'response_payload' in http_options:
-      response.copy_to_dict(http_options['response_payload'])
+    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
+    return async_generator()
 
   def upload_file(
       self, file_path: Union[str, io.IOBase], upload_url: str, upload_size: int
@@ -575,15 +568,15 @@ class ApiClient:
       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}'
+            f' finalized. {response.headers}, body: {response.json}'
         )
 
     if response.headers['X-Goog-Upload-Status'] != 'final':
       raise ValueError(
           'Failed to upload file: Upload status is not finalized. headers:'
-          f' {response.headers}, body: {response.text}'
+          f' {response.headers}, body: {response.json}'
       )
-    return response.text
+    return response.json
 
   def download_file(self, path: str, http_options):
     """Downloads the file data.
@@ -624,7 +617,6 @@ class ApiClient:
     errors.APIError.raise_for_response(response)
     return HttpResponse(response.headers, byte_stream=[response.content])
 
-
   async def async_upload_file(
       self,
       file_path: Union[str, io.IOBase],

google/genai/_api_module.py

@@ -0,0 +1,24 @@
+# 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.
+#
+
+"""Utilities for the API Modules of the Google Gen AI SDK."""
+
+from . import _api_client
+
+
+class BaseModule:
+
+  def __init__(self, api_client_: _api_client.ApiClient):
+    self._api_client = api_client_

google/genai/_automatic_function_calling_util.py

@@ -14,11 +14,18 @@
 #
 
 import inspect
-import types as typing_types
+import sys
+import types as builtin_types
+import typing
 from typing import Any, Callable, Literal, Union, _GenericAlias, get_args, get_origin
 import pydantic
 from . import types
 
+if sys.version_info >= (3, 10):
+  UnionType = builtin_types.UnionType
+else:
+  UnionType = typing._UnionGenericAlias
+
 _py_builtin_type_to_schema_type = {
     str: 'STRING',
     int: 'INTEGER',
@@ -58,8 +65,8 @@ def _raise_for_nullable_if_mldev(schema: types.Schema):
     )
 
 
-def _raise_if_schema_unsupported(variant: str, schema: types.Schema):
-  if not variant == 'VERTEX_AI':
+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)
@@ -74,11 +81,11 @@ def _is_default_value_compatible(
 
   if (
       isinstance(annotation, _GenericAlias)
-      or isinstance(annotation, typing_types.GenericAlias)
-      or isinstance(annotation, typing_types.UnionType)
+      or isinstance(annotation, builtin_types.GenericAlias)
+      or isinstance(annotation, UnionType)
   ):
     origin = get_origin(annotation)
-    if origin in (Union, typing_types.UnionType):
+    if origin in (Union, UnionType):
       return any(
           _is_default_value_compatible(default_value, arg)
           for arg in get_args(annotation)
@@ -107,12 +114,13 @@ def _is_default_value_compatible(
       return default_value in get_args(annotation)
 
   # return False for any other unrecognized annotation
-  # let caller handle the raise
   return False
 
 
 def _parse_schema_from_parameter(
-    variant: str, param: inspect.Parameter, func_name: str
+    client,
+    param: inspect.Parameter,
+    func_name: str,
 ) -> types.Schema:
   """parse schema from parameter.
 
@@ -130,12 +138,12 @@ 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(variant, schema)
+    _raise_if_schema_unsupported(client, schema)
     return schema
   if (
-      isinstance(param.annotation, typing_types.UnionType)
+      isinstance(param.annotation, UnionType)
       # only parse simple UnionType, example int | str | float | bool
-      # complex types.UnionType will be invoked in raise branch
+      # complex UnionType will be invoked in raise branch
       and all(
           (_is_builtin_primitive_or_compound(arg) or arg is type(None))
           for arg in get_args(param.annotation)
@@ -149,7 +157,7 @@ def _parse_schema_from_parameter(
         schema.nullable = True
         continue
       schema_in_any_of = _parse_schema_from_parameter(
-          variant,
+          client,
           inspect.Parameter(
               'item', inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=arg
           ),
@@ -171,10 +179,10 @@ 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(variant, schema)
+    _raise_if_schema_unsupported(client, schema)
     return schema
   if isinstance(param.annotation, _GenericAlias) or isinstance(
-      param.annotation, typing_types.GenericAlias
+      param.annotation, builtin_types.GenericAlias
   ):
     origin = get_origin(param.annotation)
     args = get_args(param.annotation)
@@ -184,7 +192,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(variant, schema)
+      _raise_if_schema_unsupported(client, schema)
       return schema
     if origin is Literal:
       if not all(isinstance(arg, str) for arg in args):
@@ -197,12 +205,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(variant, schema)
+      _raise_if_schema_unsupported(client, schema)
       return schema
     if origin is list:
       schema.type = 'ARRAY'
       schema.items = _parse_schema_from_parameter(
-          variant,
+          client,
           inspect.Parameter(
               'item',
               inspect.Parameter.POSITIONAL_OR_KEYWORD,
@@ -214,7 +222,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(variant, schema)
+      _raise_if_schema_unsupported(client, schema)
       return schema
     if origin is Union:
       schema.any_of = []
@@ -225,7 +233,7 @@ def _parse_schema_from_parameter(
           schema.nullable = True
           continue
         schema_in_any_of = _parse_schema_from_parameter(
-            variant,
+            client,
             inspect.Parameter(
                 'item',
                 inspect.Parameter.POSITIONAL_OR_KEYWORD,
@@ -233,6 +241,17 @@ def _parse_schema_from_parameter(
             ),
             func_name,
         )
+        if (
+            len(param.annotation.__args__) == 2
+            and type(None) in param.annotation.__args__
+        ):  # Optional type
+          for optional_arg in param.annotation.__args__:
+            if (
+                hasattr(optional_arg, '__origin__')
+                and optional_arg.__origin__ is list
+            ):
+              # Optional type with list, for example Optional[list[str]]
+              schema.items = schema_in_any_of.items
         if (
             schema_in_any_of.model_dump_json(exclude_none=True)
             not in unique_types
@@ -249,7 +268,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(variant, schema)
+      _raise_if_schema_unsupported(client, schema)
       return schema
       # all other generic alias will be invoked in raise branch
   if (
@@ -266,7 +285,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(
-          variant,
+          client,
           inspect.Parameter(
               field_name,
               inspect.Parameter.POSITIONAL_OR_KEYWORD,
@@ -274,7 +293,9 @@ def _parse_schema_from_parameter(
           ),
           func_name,
       )
-    _raise_if_schema_unsupported(variant, schema)
+    if client.vertexai:
+      schema.required = _get_required_fields(schema)
+    _raise_if_schema_unsupported(client, schema)
     return schema
   raise ValueError(
       f'Failed to parse the parameter {param} of function {func_name} for'

google/genai/_common.py

@@ -21,6 +21,7 @@ import enum
 import typing
 from typing import Union
 import uuid
+import warnings
 
 import pydantic
 from pydantic import alias_generators
@@ -113,12 +114,6 @@ def get_value_by_path(data: object, keys: list[str]):
   return data
 
 
-class BaseModule:
-
-  def __init__(self, api_client_: _api_client.ApiClient):
-    self._api_client = api_client_
-
-
 def convert_to_dict(obj: dict[str, object]) -> dict[str, object]:
   """Recursively converts a given object to a dictionary.
 
@@ -219,8 +214,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:

google/genai/_extra_utils.py

@@ -13,12 +13,13 @@
 # limitations under the License.
 #
 
-"""Extra utils depending on types that are shared between sync and async modules.
-"""
+"""Extra utils depending on types that are shared between sync and async modules."""
 
 import inspect
 import logging
-from typing import Any, Callable, Dict, get_args, get_origin, Optional, types as typing_types, Union
+import typing
+from typing import Any, Callable, Dict, Optional, Union, get_args, get_origin
+import sys
 
 import pydantic
 
@@ -26,6 +27,10 @@ from . import _common
 from . import errors
 from . import types
 
+if sys.version_info >= (3, 10):
+  from types import UnionType
+else:
+  UnionType = typing._UnionGenericAlias
 
 _DEFAULT_MAX_REMOTE_CALLS_AFC = 10
 
@@ -78,8 +83,8 @@ def get_function_map(
         if inspect.iscoroutinefunction(tool):
           raise errors.UnsupportedFunctionError(
               f'Function {tool.__name__} is a coroutine function, which is not'
-              ' supported for automatic function calling. Please manually invoke'
-              f' {tool.__name__} to get the function response.'
+              ' supported for automatic function calling. Please manually'
+              f' invoke {tool.__name__} to get the function response.'
           )
         function_map[tool.__name__] = tool
   return function_map
@@ -135,11 +140,13 @@ def convert_if_exist_pydantic_model(
         for k, v in value.items()
     }
   # example 1: typing.Union[int, float]
-  # example 2: int | float equivalent to typing.types.UnionType[int, float]
-  if get_origin(annotation) in (Union, typing_types.UnionType):
+  # example 2: int | float equivalent to UnionType[int, float]
+  if get_origin(annotation) in (Union, UnionType):
     for arg in get_args(annotation):
-      if isinstance(value, arg) or (
-          isinstance(value, dict) and _is_annotation_pydantic_model(arg)
+      if (
+          (get_args(arg) and get_origin(arg) is list)
+          or isinstance(value, arg)
+          or (isinstance(value, dict) and _is_annotation_pydantic_model(arg))
       ):
         try:
           return convert_if_exist_pydantic_model(
@@ -209,7 +216,9 @@ def get_function_response_parts(
       response = {'result': invoke_function_from_dict_args(args, func)}
     except Exception as e:  # pylint: disable=broad-except
       response = {'error': str(e)}
-    func_response = types.Part.from_function_response(func_name, response)
+    func_response = types.Part.from_function_response(
+        name=func_name, response=response
+    )
 
     func_response_parts.append(func_response)
   return func_response_parts
@@ -231,8 +240,7 @@ def should_disable_afc(
       and config_model.automatic_function_calling
       and config_model.automatic_function_calling.maximum_remote_calls
       is not None
-      and int(config_model.automatic_function_calling.maximum_remote_calls)
-      <= 0
+      and int(config_model.automatic_function_calling.maximum_remote_calls) <= 0
   ):
     logging.warning(
         'max_remote_calls in automatic_function_calling_config'
@@ -294,6 +302,7 @@ def get_max_remote_calls_afc(
     return _DEFAULT_MAX_REMOTE_CALLS_AFC
   return int(config_model.automatic_function_calling.maximum_remote_calls)
 
+
 def should_append_afc_history(
     config: Optional[types.GenerateContentConfigOrDict] = None,
 ) -> bool:
@@ -302,9 +311,6 @@ def should_append_afc_history(
       if config and isinstance(config, dict)
       else config
   )
-  if (
-      not config_model
-      or not config_model.automatic_function_calling
-  ):
+  if not config_model or not config_model.automatic_function_calling:
     return True
   return not config_model.automatic_function_calling.ignore_call_history