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

google/genai/models.py

@@ -16,7 +16,7 @@
 # Code generated by the Google Gen AI SDK generator DO NOT EDIT.
 
 import logging
-from typing import AsyncIterator, Awaitable, Iterator, Optional, Union
+from typing import Any, AsyncIterator, Awaitable, Iterator, Optional, Union
 from urllib.parse import urlencode
 from . import _api_module
 from . import _common
@@ -904,6 +904,9 @@ def _GenerateContentConfig_to_mldev(
         ),
     )
 
+  if getv(from_object, ['labels']) is not None:
+    raise ValueError('labels parameter is not supported in Gemini API.')
+
   if getv(from_object, ['cached_content']) is not None:
     setv(
         parent_object,
@@ -1066,6 +1069,9 @@ def _GenerateContentConfig_to_vertex(
         ),
     )
 
+  if getv(from_object, ['labels']) is not None:
+    setv(parent_object, ['labels'], getv(from_object, ['labels']))
+
   if getv(from_object, ['cached_content']) is not None:
     setv(
         parent_object,
@@ -1430,11 +1436,7 @@ def _GenerateImagesConfig_to_mldev(
     )
 
   if getv(from_object, ['enhance_prompt']) is not None:
-    setv(
-        parent_object,
-        ['parameters', 'enhancePrompt'],
-        getv(from_object, ['enhance_prompt']),
-    )
+    raise ValueError('enhance_prompt parameter is not supported in Gemini API.')
 
   return to_object
 
@@ -2785,7 +2787,7 @@ def _ComputeTokensParameters_to_vertex(
   return to_object
 
 
-def _MediaResolution_to_mldev_enum_validate(enum_value: any):
+def _MediaResolution_to_mldev_enum_validate(enum_value: Any):
   if enum_value in set([
       'MEDIA_RESOLUTION_UNSPECIFIED',
       'MEDIA_RESOLUTION_LOW',
@@ -2795,17 +2797,17 @@ def _MediaResolution_to_mldev_enum_validate(enum_value: any):
     raise ValueError(f'{enum_value} enum value is not supported in Gemini API.')
 
 
-def _SafetyFilterLevel_to_mldev_enum_validate(enum_value: any):
+def _SafetyFilterLevel_to_mldev_enum_validate(enum_value: Any):
   if enum_value in set(['BLOCK_NONE']):
     raise ValueError(f'{enum_value} enum value is not supported in Gemini API.')
 
 
-def _PersonGeneration_to_mldev_enum_validate(enum_value: any):
+def _PersonGeneration_to_mldev_enum_validate(enum_value: Any):
   if enum_value in set(['ALLOW_ALL']):
     raise ValueError(f'{enum_value} enum value is not supported in Gemini API.')
 
 
-def _MaskReferenceMode_to_mldev_enum_validate(enum_value: any):
+def _MaskReferenceMode_to_mldev_enum_validate(enum_value: Any):
   if enum_value in set([
       'MASK_MODE_DEFAULT',
       'MASK_MODE_USER_PROVIDED',
@@ -2816,7 +2818,7 @@ def _MaskReferenceMode_to_mldev_enum_validate(enum_value: any):
     raise ValueError(f'{enum_value} enum value is not supported in Gemini API.')
 
 
-def _ControlReferenceType_to_mldev_enum_validate(enum_value: any):
+def _ControlReferenceType_to_mldev_enum_validate(enum_value: Any):
   if enum_value in set([
       'CONTROL_TYPE_DEFAULT',
       'CONTROL_TYPE_CANNY',
@@ -2826,7 +2828,7 @@ def _ControlReferenceType_to_mldev_enum_validate(enum_value: any):
     raise ValueError(f'{enum_value} enum value is not supported in Gemini API.')
 
 
-def _SubjectReferenceType_to_mldev_enum_validate(enum_value: any):
+def _SubjectReferenceType_to_mldev_enum_validate(enum_value: Any):
   if enum_value in set([
       'SUBJECT_TYPE_DEFAULT',
       'SUBJECT_TYPE_PERSON',
@@ -2836,7 +2838,7 @@ def _SubjectReferenceType_to_mldev_enum_validate(enum_value: any):
     raise ValueError(f'{enum_value} enum value is not supported in Gemini API.')
 
 
-def _EditMode_to_mldev_enum_validate(enum_value: any):
+def _EditMode_to_mldev_enum_validate(enum_value: Any):
   if enum_value in set([
       'EDIT_MODE_DEFAULT',
       'EDIT_MODE_INPAINT_REMOVAL',
@@ -3389,6 +3391,9 @@ def _GeneratedImage_from_vertex(
         getv(from_object, ['raiFilteredReason']),
     )
 
+  if getv(from_object, ['prompt']) is not None:
+    setv(to_object, ['enhanced_prompt'], getv(from_object, ['prompt']))
+
   return to_object
 
 
@@ -3850,7 +3855,7 @@ class Models(_api_module.BaseModule):
     self._api_client._verify_response(return_value)
     return return_value
 
-  def generate_content_stream(
+  def _generate_content_stream(
       self,
       *,
       model: str,
@@ -4668,10 +4673,14 @@ class Models(_api_module.BaseModule):
         f'AFC is enabled with max remote calls: {remaining_remote_calls_afc}.'
     )
     automatic_function_calling_history = []
+    response = None
+    i = 0
     while remaining_remote_calls_afc > 0:
+      i += 1
       response = self._generate_content(
           model=model, contents=contents, config=config
       )
+      logging.info(f'AFC remote call {i} is done.')
       remaining_remote_calls_afc -= 1
       if remaining_remote_calls_afc == 0:
         logging.info('Reached max remote calls for automatic function calling.')
@@ -4679,6 +4688,8 @@ class Models(_api_module.BaseModule):
       function_map = _extra_utils.get_function_map(config)
       if not function_map:
         break
+      if not response:
+        break
       if (
           not response.candidates
           or not response.candidates[0].content
@@ -4690,21 +4701,165 @@ class Models(_api_module.BaseModule):
       )
       if not func_response_parts:
         break
-      contents = t.t_contents(self._api_client, contents)
-      contents.append(response.candidates[0].content)
-      contents.append(
-          types.Content(
-              role='user',
-              parts=func_response_parts,
-          )
+      func_call_content = response.candidates[0].content
+      func_response_content = types.Content(
+          role='user',
+          parts=func_response_parts,
       )
-      automatic_function_calling_history.extend(contents)
+      contents = t.t_contents(self._api_client, contents)
+      if not automatic_function_calling_history:
+        automatic_function_calling_history.extend(contents)
+      contents.append(func_call_content)
+      contents.append(func_response_content)
+      automatic_function_calling_history.append(func_call_content)
+      automatic_function_calling_history.append(func_response_content)
     if _extra_utils.should_append_afc_history(config):
       response.automatic_function_calling_history = (
           automatic_function_calling_history
       )
     return response
 
+  def generate_content_stream(
+      self,
+      *,
+      model: str,
+      contents: Union[types.ContentListUnion, types.ContentListUnionDict],
+      config: Optional[types.GenerateContentConfigOrDict] = None,
+  ) -> Iterator[types.GenerateContentResponse]:
+    """Makes an API request to generate content using a model and yields the model's response in chunks.
+
+    For the `model` parameter, supported format for Vertex AI API includes:
+    - the Gemini model ID, for example: 'gemini-1.5-flash-002'
+    - the full resource name starts with 'projects/', for example:
+      'projects/my-project-id/locations/us-central1/publishers/google/models/gemini-1.5-flash-002'
+    - the partial resource name with 'publishers/', for example:
+      'publishers/google/models/gemini-1.5-flash-002' or
+      'publishers/meta/models/llama-3.1-405b-instruct-maas'
+    - `/` separated publisher and model name, for example:
+      'google/gemini-1.5-flash-002' or 'meta/llama-3.1-405b-instruct-maas'
+
+    For the `model` parameter, supported format for Gemini API includes:
+    - the Gemini model ID, for example: 'gemini-1.5-flash-002'
+    - the model name starts with 'models/', for example:
+      'models/gemini-1.5-flash-002'
+    - if you would like to use a tuned model, the model name starts with
+      'tunedModels/', for example:
+      'tunedModels/1234567890123456789'
+
+    Some models support multimodal input and output.
+
+    Usage:
+
+    .. code-block:: python
+
+      from google.genai import types
+      from google import genai
+
+      client = genai.Client(
+          vertexai=True, project='my-project-id', location='us-central1'
+      )
+
+      for chunk in client.models.generate_content_stream(
+        model='gemini-1.5-flash-002',
+        contents='''What is a good name for a flower shop that specializes in
+          selling bouquets of dried flowers?'''
+      ):
+        print(chunk.text)
+      # **Elegant & Classic:**
+      # * The Dried Bloom
+      # * Everlasting Florals
+      # * Timeless Petals
+
+      for chunk in client.models.generate_content_stream(
+        model='gemini-1.5-flash-002',
+        contents=[
+          types.Part.from_text('What is shown in this image?'),
+          types.Part.from_uri('gs://generativeai-downloads/images/scones.jpg',
+          'image/jpeg')
+        ]
+      ):
+        print(chunk.text)
+      # The image shows a flat lay arrangement of freshly baked blueberry
+      # scones.
+    """
+
+    if _extra_utils.should_disable_afc(config):
+      return self._generate_content_stream(
+          model=model, contents=contents, config=config
+      )
+    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
+    i = 0
+    while remaining_remote_calls_afc > 0:
+      i += 1
+      response = self._generate_content_stream(
+          model=model, contents=contents, config=config
+      )
+      logging.info(f'AFC remote call {i} is done.')
+      remaining_remote_calls_afc -= 1
+      if remaining_remote_calls_afc == 0:
+        logging.info('Reached max remote calls for automatic function calling.')
+
+      function_map = _extra_utils.get_function_map(config)
+
+      if i == 1:
+        # First request gets a function call.
+        # Then get function response parts.
+        # Yield chunks only if there's no function response parts.
+        for chunk in response:
+          if not function_map:
+            yield chunk
+          else:
+            func_response_parts = _extra_utils.get_function_response_parts(
+                chunk, function_map
+            )
+            if not func_response_parts:
+              yield chunk
+
+      else:
+        #  Second request and beyond, yield chunks.
+        for chunk in response:
+          if _extra_utils.should_append_afc_history(config):
+            chunk.automatic_function_calling_history = (
+                automatic_function_calling_history
+            )
+          yield chunk
+      if not chunk:
+        break
+      if (
+          not chunk
+          or not chunk.candidates
+          or not chunk.candidates[0].content
+          or not chunk.candidates[0].content.parts
+      ):
+        break
+
+      if not function_map:
+        break
+      func_response_parts = _extra_utils.get_function_response_parts(
+          chunk, function_map
+      )
+      if not func_response_parts:
+        break
+
+      # Append function response parts to contents for the next request.
+      func_call_content = chunk.candidates[0].content
+      func_response_content = types.Content(
+          role='user',
+          parts=func_response_parts,
+      )
+      contents = t.t_contents(self._api_client, contents)
+      if not automatic_function_calling_history:
+        automatic_function_calling_history.extend(contents)
+      contents.append(func_call_content)
+      contents.append(func_response_content)
+      automatic_function_calling_history.append(func_call_content)
+      automatic_function_calling_history.append(func_response_content)
+
   def upscale_image(
       self,
       *,
@@ -4768,8 +4923,8 @@ class Models(_api_module.BaseModule):
   ) -> Pager[types.Model]:
     """Makes an API request to list the available models.
 
-    If `query_base` is set to True in the config, the API will return all
-    available base models. If set to False or not set (default), it will return
+    If `query_base` is set to True in the config or not set (default), the
+    API will return all available base models. If set to False, it will return
     all tuned models.
 
     Args:
@@ -4792,6 +4947,8 @@ class Models(_api_module.BaseModule):
         types._ListModelsParameters(config=config).config
         or types.ListModelsConfig()
     )
+    if config.query_base is None:
+      config.query_base = True
     if self._api_client.vertexai:
       config = config.copy()
       if not config.query_base:
@@ -4802,8 +4959,6 @@ class Models(_api_module.BaseModule):
             if filter_value
             else 'labels.tune-type:*'
         )
-    if not config.query_base:
-      config.query_base = False
     return Pager(
         'models',
         self._list,
@@ -4871,7 +5026,7 @@ class AsyncModels(_api_module.BaseModule):
     self._api_client._verify_response(return_value)
     return return_value
 
-  async def generate_content_stream(
+  async def _generate_content_stream(
       self,
       *,
       model: str,
@@ -5666,6 +5821,7 @@ class AsyncModels(_api_module.BaseModule):
         f'AFC is enabled with max remote calls: {remaining_remote_calls_afc}.'
     )
     automatic_function_calling_history = []
+    response = None
     while remaining_remote_calls_afc > 0:
       response = await self._generate_content(
           model=model, contents=contents, config=config
@@ -5677,6 +5833,8 @@ class AsyncModels(_api_module.BaseModule):
       function_map = _extra_utils.get_function_map(config)
       if not function_map:
         break
+      if not response:
+        break
       if (
           not response.candidates
           or not response.candidates[0].content
@@ -5688,15 +5846,18 @@ class AsyncModels(_api_module.BaseModule):
       )
       if not func_response_parts:
         break
-      contents = t.t_contents(self._api_client, contents)
-      contents.append(response.candidates[0].content)
-      contents.append(
-          types.Content(
-              role='user',
-              parts=func_response_parts,
-          )
+      func_call_content = response.candidates[0].content
+      func_response_content = types.Content(
+          role='user',
+          parts=func_response_parts,
       )
-      automatic_function_calling_history.extend(contents)
+      contents = t.t_contents(self._api_client, contents)
+      if not automatic_function_calling_history:
+        automatic_function_calling_history.extend(contents)
+      contents.append(func_call_content)
+      contents.append(func_response_content)
+      automatic_function_calling_history.append(func_call_content)
+      automatic_function_calling_history.append(func_response_content)
 
     if _extra_utils.should_append_afc_history(config):
       response.automatic_function_calling_history = (
@@ -5704,6 +5865,154 @@ class AsyncModels(_api_module.BaseModule):
       )
     return response
 
+  async def generate_content_stream(
+      self,
+      *,
+      model: str,
+      contents: Union[types.ContentListUnion, types.ContentListUnionDict],
+      config: Optional[types.GenerateContentConfigOrDict] = None,
+  ) -> Awaitable[AsyncIterator[types.GenerateContentResponse]]:
+    """Makes an API request to generate content using a model and yields the model's response in chunks.
+
+    For the `model` parameter, supported format for Vertex AI API includes:
+    - the Gemini model ID, for example: 'gemini-1.5-flash-002'
+    - the full resource name starts with 'projects/', for example:
+      'projects/my-project-id/locations/us-central1/publishers/google/models/gemini-1.5-flash-002'
+    - the partial resource name with 'publishers/', for example:
+      'publishers/google/models/gemini-1.5-flash-002' or
+      'publishers/meta/models/llama-3.1-405b-instruct-maas'
+    - `/` separated publisher and model name, for example:
+      'google/gemini-1.5-flash-002' or 'meta/llama-3.1-405b-instruct-maas'
+
+    For the `model` parameter, supported format for Gemini API includes:
+    - the Gemini model ID, for example: 'gemini-1.5-flash-002'
+    - the model name starts with 'models/', for example:
+      'models/gemini-1.5-flash-002'
+    - if you would like to use a tuned model, the model name starts with
+      'tunedModels/', for example:
+      'tunedModels/1234567890123456789'
+
+    Some models support multimodal input and output.
+
+    Usage:
+
+    .. code-block:: python
+
+      from google.genai import types
+      from google import genai
+
+      client = genai.Client(
+          vertexai=True, project='my-project-id', location='us-central1'
+      )
+
+      async for chunk in await client.aio.models.generate_content_stream(
+        model='gemini-1.5-flash-002',
+        contents='''What is a good name for a flower shop that specializes in
+          selling bouquets of dried flowers?'''
+      ):
+        print(chunk.text)
+      # **Elegant & Classic:**
+      # * The Dried Bloom
+      # * Everlasting Florals
+      # * Timeless Petals
+
+      async for chunk in awiat client.aio.models.generate_content_stream(
+        model='gemini-1.5-flash-002',
+        contents=[
+          types.Part.from_text('What is shown in this image?'),
+          types.Part.from_uri('gs://generativeai-downloads/images/scones.jpg',
+          'image/jpeg')
+        ]
+      ):
+        print(chunk.text)
+      # The image shows a flat lay arrangement of freshly baked blueberry
+      # scones.
+    """
+
+    if _extra_utils.should_disable_afc(config):
+      return self._generate_content_stream(
+          model=model, contents=contents, config=config
+      )
+
+    async def async_generator(model, contents, config):
+      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
+      i = 0
+      while remaining_remote_calls_afc > 0:
+        i += 1
+        response = await self._generate_content_stream(
+            model=model, contents=contents, config=config
+        )
+        logging.info(f'AFC remote call {i} is done.')
+        remaining_remote_calls_afc -= 1
+        if remaining_remote_calls_afc == 0:
+          logging.info(
+              'Reached max remote calls for automatic function calling.'
+          )
+
+        function_map = _extra_utils.get_function_map(config)
+
+        if i == 1:
+          # First request gets a function call.
+          # Then get function response parts.
+          # Yield chunks only if there's no function response parts.
+          async for chunk in response:
+            if not function_map:
+              yield chunk
+            else:
+              func_response_parts = _extra_utils.get_function_response_parts(
+                  chunk, function_map
+              )
+              if not func_response_parts:
+                yield chunk
+
+        else:
+          #  Second request and beyond, yield chunks.
+          async for chunk in response:
+
+            if _extra_utils.should_append_afc_history(config):
+              chunk.automatic_function_calling_history = (
+                  automatic_function_calling_history
+              )
+            yield chunk
+
+        if not chunk:
+          break
+        if (
+            not chunk
+            or not chunk.candidates
+            or not chunk.candidates[0].content
+            or not chunk.candidates[0].content.parts
+        ):
+          break
+        if not function_map:
+          break
+        func_response_parts = _extra_utils.get_function_response_parts(
+            chunk, function_map
+        )
+        if not func_response_parts:
+          break
+
+        # Append function response parts to contents for the next request.
+        func_call_content = chunk.candidates[0].content
+        func_response_content = types.Content(
+            role='user',
+            parts=func_response_parts,
+        )
+        contents = t.t_contents(self._api_client, contents)
+        if not automatic_function_calling_history:
+          automatic_function_calling_history.extend(contents)
+        contents.append(func_call_content)
+        contents.append(func_response_content)
+        automatic_function_calling_history.append(func_call_content)
+        automatic_function_calling_history.append(func_response_content)
+
+    return async_generator(model, contents, config)
+
   async def list(
       self,
       *,