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

google/genai/_automatic_function_calling_util.py

@@ -62,8 +62,8 @@ def _raise_for_default_if_mldev(schema: types.Schema):
     )
 
 
-def _raise_if_schema_unsupported(client, schema: types.Schema):
-  if not client.vertexai:
+def _raise_if_schema_unsupported(api_option: Literal['VERTEX_AI', 'GEMINI_API'], schema: types.Schema):
+  if api_option == 'GEMINI_API':
     _raise_for_any_of_if_mldev(schema)
     _raise_for_default_if_mldev(schema)
 
@@ -114,7 +114,7 @@ def _is_default_value_compatible(
 
 
 def _parse_schema_from_parameter(
-    client,
+    api_option: Literal['VERTEX_AI', 'GEMINI_API'],
     param: inspect.Parameter,
     func_name: str,
 ) -> types.Schema:
@@ -134,7 +134,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(api_option, schema)
     return schema
   if (
       isinstance(param.annotation, VersionedUnionType)
@@ -153,7 +153,7 @@ def _parse_schema_from_parameter(
         schema.nullable = True
         continue
       schema_in_any_of = _parse_schema_from_parameter(
-          client,
+          api_option,
           inspect.Parameter(
               'item', inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=arg
           ),
@@ -175,7 +175,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(api_option, schema)
     return schema
   if isinstance(param.annotation, _GenericAlias) or isinstance(
       param.annotation, builtin_types.GenericAlias
@@ -188,7 +188,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(api_option, schema)
       return schema
     if origin is Literal:
       if not all(isinstance(arg, str) for arg in args):
@@ -201,12 +201,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(api_option, schema)
       return schema
     if origin is list:
       schema.type = 'ARRAY'
       schema.items = _parse_schema_from_parameter(
-          client,
+          api_option,
           inspect.Parameter(
               'item',
               inspect.Parameter.POSITIONAL_OR_KEYWORD,
@@ -218,7 +218,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(api_option, schema)
       return schema
     if origin is Union:
       schema.any_of = []
@@ -233,7 +233,7 @@ def _parse_schema_from_parameter(
           schema.nullable = True
           continue
         schema_in_any_of = _parse_schema_from_parameter(
-            client,
+            api_option,
             inspect.Parameter(
                 'item',
                 inspect.Parameter.POSITIONAL_OR_KEYWORD,
@@ -268,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(client, schema)
+      _raise_if_schema_unsupported(api_option, schema)
       return schema
       # all other generic alias will be invoked in raise branch
   if (
@@ -284,7 +284,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,
+          api_option,
           inspect.Parameter(
               field_name,
               inspect.Parameter.POSITIONAL_OR_KEYWORD,
@@ -292,9 +292,9 @@ def _parse_schema_from_parameter(
           ),
           func_name,
       )
-    if client.vertexai:
+    if api_option == 'VERTEX_AI':
       schema.required = _get_required_fields(schema)
-    _raise_if_schema_unsupported(client, schema)
+    _raise_if_schema_unsupported(api_option, schema)
     return schema
   raise ValueError(
       f'Failed to parse the parameter {param} of function {func_name} for'

google/genai/_common.py

@@ -220,7 +220,8 @@ class CaseInSensitiveEnum(str, enum.Enum):
         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
+          # We need to use super() to avoid infinite recursion.
+          unknown_enum_val = super().__new__(cls, value)
           unknown_enum_val._name_ = str(value)  # pylint: disable=protected-access
           unknown_enum_val._value_ = value  # pylint: disable=protected-access
           return unknown_enum_val

google/genai/_extra_utils.py

@@ -271,8 +271,8 @@ def should_disable_afc(
       and int(config_model.automatic_function_calling.maximum_remote_calls) > 0
   ):
     logging.warning(
-        '`automatic_function_calling.disable` is set to `True`. But'
-        ' `automatic_function_calling.maximum_remote_calls` is set to be a'
+        '`automatic_function_calling.disable` is set to `True`. And'
+        ' `automatic_function_calling.maximum_remote_calls` is a'
         ' positive number'
         f' {config_model.automatic_function_calling.maximum_remote_calls}.'
         ' Disabling automatic function calling. If you want to enable'

google/genai/_transformers.py

@@ -374,8 +374,8 @@ def handle_null_fields(schema: dict[str, Any]):
         schema['anyOf'].remove({'type': 'null'})
         if len(schema['anyOf']) == 1:
           # If there is only one type left after removing null, remove the anyOf field.
-          field_type = schema['anyOf'][0]['type']
-          schema['type'] = field_type
+          for key,val in schema['anyOf'][0].items():
+            schema[key] = val
           del schema['anyOf']
 
 
@@ -446,9 +446,17 @@ def process_schema(
 
     if schema.get('default') is not None:
       raise ValueError(
-          'Default value is not supported in the response schema for the Gemmini API.'
+          'Default value is not supported in the response schema for the Gemini API.'
       )
 
+  if schema.get('title') == 'PlaceholderLiteralEnum':
+    schema.pop('title', None)
+
+  # If a dict is provided directly to response_schema, it may use `any_of`
+  # instead of `anyOf`. Otherwise model_json_schema() uses `anyOf`
+  if schema.get('any_of', None) is not None:
+    schema['anyOf'] = schema.pop('any_of')
+
   if defs is None:
     defs = schema.pop('$defs', {})
     for _, sub_schema in defs.items():
@@ -456,6 +464,15 @@ def process_schema(
 
   handle_null_fields(schema)
 
+  # After removing null fields, Optional fields with only one possible type
+  # will have a $ref key that needs to be flattened
+  # For example: {'default': None, 'description': 'Name of the person', 'nullable': True, '$ref': '#/$defs/TestPerson'}
+  if schema.get('$ref', None):
+    ref = defs[schema.get('$ref').split('defs/')[-1]]
+    for schema_key in list(ref.keys()):
+      schema[schema_key] = ref[schema_key]
+    del schema['$ref']
+
   any_of = schema.get('anyOf', None)
   if any_of is not None:
     if not client.vertexai:
@@ -478,6 +495,16 @@ def process_schema(
     schema_type = schema_type.value
   schema_type = schema_type.upper()
 
+  # model_json_schema() returns a schema with a 'const' field when a Literal with one value is provided as a pydantic field
+  # For example `genre: Literal['action']` becomes: {'const': 'action', 'title': 'Genre', 'type': 'string'}
+  const = schema.get('const', None)
+  if const is not None:
+    if schema_type == 'STRING':
+      schema['enum'] = [const]
+      del schema['const']
+    else:
+      raise ValueError('Literal values must be strings.')
+
   if schema_type == 'OBJECT':
     properties = schema.get('properties', None)
     if properties is None:
@@ -502,6 +529,7 @@ def process_schema(
       process_schema(ref, client, defs)
       schema['items'] = ref
 
+
 def _process_enum(
     enum: EnumMeta, client: Optional[_api_client.ApiClient] = None
 ) -> types.Schema:

google/genai/chats.py

@@ -57,12 +57,16 @@ class Chat(_BaseChat):
   """Chat session."""
 
   def send_message(
-      self, message: Union[list[PartUnionDict], PartUnionDict]
+      self,
+      message: Union[list[PartUnionDict], PartUnionDict],
+      config: Optional[GenerateContentConfigOrDict] = None,
   ) -> GenerateContentResponse:
     """Sends the conversation history with the additional message and returns the model's response.
 
     Args:
       message: The message to send to the model.
+      config:  Optional config to override the default Chat config for this
+        request.
 
     Returns:
       The model's response.
@@ -79,7 +83,7 @@ class Chat(_BaseChat):
     response = self._modules.generate_content(
         model=self._model,
         contents=self._curated_history + [input_content],
-        config=self._config,
+        config=config if config else self._config,
     )
     if _validate_response(response):
       if response.automatic_function_calling_history:
@@ -92,12 +96,16 @@ class Chat(_BaseChat):
     return response
 
   def send_message_stream(
-      self, message: Union[list[PartUnionDict], PartUnionDict]
+      self,
+      message: Union[list[PartUnionDict], PartUnionDict],
+      config: Optional[GenerateContentConfigOrDict] = None,
   ):
     """Sends the conversation history with the additional message and yields the model's response in chunks.
 
     Args:
       message: The message to send to the model.
+      config: Optional config to override the default Chat config for this
+        request.
 
     Yields:
       The model's response in chunks.
@@ -117,7 +125,7 @@ class Chat(_BaseChat):
     for chunk in self._modules.generate_content_stream(
         model=self._model,
         contents=self._curated_history + [input_content],
-        config=self._config,
+        config=config if config else self._config,
     ):
       if _validate_response(chunk):
         output_contents.append(chunk.candidates[0].content)
@@ -164,12 +172,16 @@ class AsyncChat(_BaseChat):
   """Async chat session."""
 
   async def send_message(
-      self, message: Union[list[PartUnionDict], PartUnionDict]
+      self,
+      message: Union[list[PartUnionDict], PartUnionDict],
+      config: Optional[GenerateContentConfigOrDict] = None,
   ) -> GenerateContentResponse:
     """Sends the conversation history with the additional message and returns model's response.
 
     Args:
       message: The message to send to the model.
+      config: Optional config to override the default Chat config for this
+        request.
 
     Returns:
       The model's response.
@@ -186,7 +198,7 @@ class AsyncChat(_BaseChat):
     response = await self._modules.generate_content(
         model=self._model,
         contents=self._curated_history + [input_content],
-        config=self._config,
+        config=config if config else self._config,
     )
     if _validate_response(response):
       if response.automatic_function_calling_history:
@@ -199,12 +211,16 @@ class AsyncChat(_BaseChat):
     return response
 
   async def send_message_stream(
-      self, message: Union[list[PartUnionDict], PartUnionDict]
+      self,
+      message: Union[list[PartUnionDict], PartUnionDict],
+      config: Optional[GenerateContentConfigOrDict] = None,
   ) -> Awaitable[AsyncIterator[GenerateContentResponse]]:
     """Sends the conversation history with the additional message and yields the model's response in chunks.
 
     Args:
       message: The message to send to the model.
+      config: Optional config to override the default Chat config for this
+        request.
 
     Yields:
       The model's response in chunks.
@@ -225,7 +241,7 @@ class AsyncChat(_BaseChat):
       async for chunk in await self._modules.generate_content_stream(
           model=self._model,
           contents=self._curated_history + [input_content],
-          config=self._config,
+          config=config if config else self._config,
       ):
         if _validate_response(chunk):
           output_contents.append(chunk.candidates[0].content)

google/genai/models.py

@@ -4784,15 +4784,18 @@ class Models(_api_module.BaseModule):
     """
 
     if _extra_utils.should_disable_afc(config):
-      return self._generate_content_stream(
+      yield from self._generate_content_stream(
           model=model, contents=contents, config=config
       )
+      return
+
     remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc(config)
     logging.info(
         f'AFC is enabled with max remote calls: {remaining_remote_calls_afc}.'
     )
     automatic_function_calling_history = []
     chunk = None
+    func_response_parts = None
     i = 0
     while remaining_remote_calls_afc > 0:
       i += 1
@@ -4828,6 +4831,10 @@ class Models(_api_module.BaseModule):
                 automatic_function_calling_history
             )
           yield chunk
+        func_response_parts = _extra_utils.get_function_response_parts(
+            chunk, function_map
+        )
+
       if not chunk:
         break
       if (
@@ -4840,9 +4847,6 @@ class Models(_api_module.BaseModule):
 
       if not function_map:
         break
-      func_response_parts = _extra_utils.get_function_response_parts(
-          chunk, function_map
-      )
       if not func_response_parts:
         break
 
@@ -5930,10 +5934,16 @@ class AsyncModels(_api_module.BaseModule):
     """
 
     if _extra_utils.should_disable_afc(config):
-      return self._generate_content_stream(
+      response = await self._generate_content_stream(
           model=model, contents=contents, config=config
       )
 
+      async def base_async_generator(model, contents, config):
+        async for chunk in response:
+          yield chunk
+
+      return base_async_generator(model, contents, config)
+
     async def async_generator(model, contents, config):
       remaining_remote_calls_afc = _extra_utils.get_max_remote_calls_afc(config)
       logging.info(
@@ -5979,7 +5989,9 @@ class AsyncModels(_api_module.BaseModule):
                   automatic_function_calling_history
               )
             yield chunk
-
+          func_response_parts = _extra_utils.get_function_response_parts(
+              chunk, function_map
+          )
         if not chunk:
           break
         if (
@@ -5991,9 +6003,7 @@ class AsyncModels(_api_module.BaseModule):
           break
         if not function_map:
           break
-        func_response_parts = _extra_utils.get_function_response_parts(
-            chunk, function_map
-        )
+
         if not func_response_parts:
           break
 

google/genai/tunings.py

@@ -15,8 +15,7 @@
 
 # Code generated by the Google Gen AI SDK generator DO NOT EDIT.
 
-import time
-from typing import Any, Optional, Union
+from typing import Optional, Union
 from urllib.parse import urlencode
 from . import _api_module
 from . import _common
@@ -998,11 +997,14 @@ class Tunings(_api_module.BaseModule):
           config=config,
       )
       operation_dict = operation.to_json_dict()
-      tuned_model_dict = _resolve_operation(self._api_client, operation_dict)
-      tuning_job_dict = _TuningJob_from_mldev(
-          self._api_client, tuned_model_dict
+      try:
+        tuned_model_name = operation_dict['metadata']['tunedModel']
+      except KeyError:
+        tuned_model_name = operation_dict['name'].partition('/operations/')[0]
+      tuning_job = types.TuningJob(
+          name=tuned_model_name,
+          state=types.JobState.JOB_STATE_QUEUED,
       )
-      tuning_job = types.TuningJob._from_response(tuning_job_dict, None)
     if tuning_job.name and self._api_client.vertexai:
       _IpythonUtils.display_model_tuning_button(
           tuning_job_resource=tuning_job.name
@@ -1010,72 +1012,6 @@ class Tunings(_api_module.BaseModule):
     return tuning_job
 
 
-_LRO_POLLING_INITIAL_DELAY_SECONDS = 1.0
-_LRO_POLLING_MAXIMUM_DELAY_SECONDS = 20.0
-_LRO_POLLING_TIMEOUT_SECONDS = 900.0
-_LRO_POLLING_MULTIPLIER = 1.5
-
-
-def _resolve_operation(api_client: ApiClient, struct: dict[str, Any]):
-  if (name := struct.get('name')) and '/operations/' in name:
-    operation: dict[str, Any] = struct
-    total_seconds = 0.0
-    delay_seconds = _LRO_POLLING_INITIAL_DELAY_SECONDS
-    while not operation.get('done'):
-      if total_seconds > _LRO_POLLING_TIMEOUT_SECONDS:
-        raise RuntimeError(f'Operation {name} timed out.\n{operation}')
-      # TODO(b/374433890): Replace with LRO module once it's available.
-      operation: dict[str, Any] = api_client.request(
-          http_method='GET', path=name, request_dict={}
-      )
-      if 'ReplayApiClient' not in type(api_client).__name__:
-        time.sleep(delay_seconds)
-      total_seconds += total_seconds
-      # Exponential backoff
-      delay_seconds = min(
-          delay_seconds * _LRO_POLLING_MULTIPLIER,
-          _LRO_POLLING_MAXIMUM_DELAY_SECONDS,
-      )
-    if error := operation.get('error'):
-      raise RuntimeError(
-          f'Operation {name} failed with error: {error}.\n{operation}'
-      )
-    return operation.get('response')
-  else:
-    return struct
-
-
-async def _resolve_operation_async(
-    api_client: ApiClient, struct: dict[str, Any]
-):
-  if (name := struct.get('name')) and '/operations/' in name:
-    operation: dict[str, Any] = struct
-    total_seconds = 0.0
-    delay_seconds = _LRO_POLLING_INITIAL_DELAY_SECONDS
-    while not operation.get('done'):
-      if total_seconds > _LRO_POLLING_TIMEOUT_SECONDS:
-        raise RuntimeError(f'Operation {name} timed out.\n{operation}')
-      # TODO(b/374433890): Replace with LRO module once it's available.
-      operation: dict[str, Any] = await api_client.async_request(
-          http_method='GET', path=name, request_dict={}
-      )
-      if 'ReplayApiClient' not in type(api_client).__name__:
-        time.sleep(delay_seconds)
-      total_seconds += total_seconds
-      # Exponential backoff
-      delay_seconds = min(
-          delay_seconds * _LRO_POLLING_MULTIPLIER,
-          _LRO_POLLING_MAXIMUM_DELAY_SECONDS,
-      )
-    if error := operation.get('error'):
-      raise RuntimeError(
-          f'Operation {name} failed with error: {error}.\n{operation}'
-      )
-    return operation.get('response')
-  else:
-    return struct
-
-
 class AsyncTunings(_api_module.BaseModule):
 
   async def _get(
@@ -1370,13 +1306,14 @@ class AsyncTunings(_api_module.BaseModule):
           config=config,
       )
       operation_dict = operation.to_json_dict()
-      tuned_model_dict = await _resolve_operation_async(
-          self._api_client, operation_dict
-      )
-      tuning_job_dict = _TuningJob_from_mldev(
-          self._api_client, tuned_model_dict
+      try:
+        tuned_model_name = operation_dict['metadata']['tunedModel']
+      except KeyError:
+        tuned_model_name = operation_dict['name'].partition('/operations/')[0]
+      tuning_job = types.TuningJob(
+          name=tuned_model_name,
+          state=types.JobState.JOB_STATE_QUEUED,
       )
-      tuning_job = types.TuningJob._from_response(tuning_job_dict, None)
     if tuning_job.name and self._api_client.vertexai:
       _IpythonUtils.display_model_tuning_button(
           tuning_job_resource=tuning_job.name

google/genai/types.py

@@ -971,13 +971,24 @@ class FunctionDeclaration(_common.BaseModel):
   )
 
   @classmethod
-  def from_callable(
+  def from_callable_with_api_option(
       cls,
       *,
-      client,
       callable: Callable,
+      api_option: Literal['VERTEX_AI', 'GEMINI_API'] = 'GEMINI_API',
   ) -> 'FunctionDeclaration':
-    """Converts a Callable to a FunctionDeclaration based on the client."""
+    """Converts a Callable to a FunctionDeclaration based on the API option.
+
+    Supported API option is 'VERTEX_AI' or 'GEMINI_API'. If api_option is unset,
+    it will default to 'GEMINI_API'. If unsupported api_option is provided, it
+    will raise ValueError.
+    """
+    supported_api_options = ['VERTEX_AI', 'GEMINI_API']
+    if api_option not in supported_api_options:
+      raise ValueError(
+          f'Unsupported api_option value: {api_option}. Supported api_option'
+          f' value is one of: {supported_api_options}.'
+      )
     from . import _automatic_function_calling_util
 
     parameters_properties = {}
@@ -988,7 +999,7 @@ class FunctionDeclaration(_common.BaseModel):
           inspect.Parameter.POSITIONAL_ONLY,
       ):
         schema = _automatic_function_calling_util._parse_schema_from_parameter(
-            client, param, callable.__name__
+            api_option, param, callable.__name__
         )
         parameters_properties[name] = schema
     declaration = FunctionDeclaration(
@@ -1000,13 +1011,13 @@ class FunctionDeclaration(_common.BaseModel):
           type='OBJECT',
           properties=parameters_properties,
       )
-      if client.vertexai:
+      if api_option == 'VERTEX_AI':
         declaration.parameters.required = (
             _automatic_function_calling_util._get_required_fields(
                 declaration.parameters
             )
         )
-    if not client.vertexai:
+    if api_option == 'GEMINI_API':
       return declaration
 
     return_annotation = inspect.signature(callable).return_annotation
@@ -1015,7 +1026,7 @@ class FunctionDeclaration(_common.BaseModel):
 
     declaration.response = (
         _automatic_function_calling_util._parse_schema_from_parameter(
-            client,
+            api_option,
             inspect.Parameter(
                 'return_value',
                 inspect.Parameter.POSITIONAL_OR_KEYWORD,
@@ -1026,6 +1037,23 @@ class FunctionDeclaration(_common.BaseModel):
     )
     return declaration
 
+  @classmethod
+  def from_callable(
+      cls,
+      *,
+      client,
+      callable: Callable,
+  ) -> 'FunctionDeclaration':
+    """Converts a Callable to a FunctionDeclaration based on the client."""
+    if client.vertexai:
+      return cls.from_callable_with_api_option(
+          callable=callable, api_option='VERTEX_AI'
+      )
+    else:
+      return cls.from_callable_with_api_option(
+          callable=callable, api_option='GEMINI_API'
+      )
+
 
 class FunctionDeclarationDict(TypedDict, total=False):
   """Defines a function that the model can generate JSON inputs for.
@@ -1909,6 +1937,20 @@ class GenerateContentConfig(_common.BaseModel):
       """,
   )
 
+  @pydantic.field_validator('response_schema', mode='before')
+  @classmethod
+  def _convert_literal_to_enum(cls, value):
+    if typing.get_origin(value) is typing.Literal:
+      enum_vals = typing.get_args(value)
+      if not all(isinstance(arg, str) for arg in enum_vals):
+        # This doesn't stop execution, it tells pydantic to raise a ValidationError
+        # when the class is instantiated with an unsupported Literal
+        raise ValueError(f'Literal type {value} must be a list of strings.')
+      # The title 'PlaceholderLiteralEnum' is removed from the generated Schema
+      # before sending the request
+      return Enum('PlaceholderLiteralEnum', {s: s for s in enum_vals})
+    return value
+
 
 class GenerateContentConfigDict(TypedDict, total=False):
   """Optional model configuration parameters.
@@ -2815,7 +2857,7 @@ class GenerateContentResponse(_common.BaseModel):
     text = ''
     any_text_part_text = False
     for part in self.candidates[0].content.parts:
-      for field_name, field_value in part.dict(
+      for field_name, field_value in part.model_dump(
           exclude={'text', 'thought'}
       ).items():
         if field_value is not None:
@@ -2883,6 +2925,11 @@ class GenerateContentResponse(_common.BaseModel):
       enum_value = result.text.replace('"', '')
       try:
         result.parsed = response_schema(enum_value)
+        if (
+            hasattr(response_schema, '__name__')
+            and response_schema.__name__ == 'PlaceholderLiteralEnum'
+        ):
+          result.parsed = str(response_schema(enum_value).name)
       except ValueError:
         pass
     elif isinstance(response_schema, GenericAlias) or isinstance(

google/genai/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 #
 
-__version__ = '1.0.0' # x-release-please-version
+__version__ = '1.1.0' # x-release-please-version