feat: Save prompt safety attributes in dedicated field for generate_images

PiperOrigin-RevId: 739252571

Author: matthew29tang

google/genai/models.py

@@ -3585,6 +3585,9 @@ def _SafetyAttributes_from_mldev(
         to_object, ['scores'], getv(from_object, ['safetyAttributes', 'scores'])
     )
 
+  if getv(from_object, ['contentType']) is not None:
+    setv(to_object, ['content_type'], getv(from_object, ['contentType']))
+
   return to_object
 
 
@@ -3606,6 +3609,9 @@ def _SafetyAttributes_from_vertex(
         to_object, ['scores'], getv(from_object, ['safetyAttributes', 'scores'])
     )
 
+  if getv(from_object, ['contentType']) is not None:
+    setv(to_object, ['content_type'], getv(from_object, ['contentType']))
+
   return to_object
 
 
@@ -3692,6 +3698,17 @@ def _GenerateImagesResponse_from_mldev(
         ],
     )
 
+  if getv(from_object, ['positivePromptSafetyAttributes']) is not None:
+    setv(
+        to_object,
+        ['positive_prompt_safety_attributes'],
+        _SafetyAttributes_from_mldev(
+            api_client,
+            getv(from_object, ['positivePromptSafetyAttributes']),
+            to_object,
+        ),
+    )
+
   return to_object
 
 
@@ -3711,6 +3728,17 @@ def _GenerateImagesResponse_from_vertex(
         ],
     )
 
+  if getv(from_object, ['positivePromptSafetyAttributes']) is not None:
+    setv(
+        to_object,
+        ['positive_prompt_safety_attributes'],
+        _SafetyAttributes_from_vertex(
+            api_client,
+            getv(from_object, ['positivePromptSafetyAttributes']),
+            to_object,
+        ),
+    )
+
   return to_object
 
 
@@ -4515,7 +4543,7 @@ def embed_content(
     self._api_client._verify_response(return_value)
     return return_value
 
-  def generate_images(
+  def _generate_images(
       self,
       *,
       model: str,
@@ -4528,21 +4556,6 @@ def generate_images(
       model (str): The model to use.
       prompt (str): A text description of the images to generate.
       config (GenerateImagesConfig): Configuration for generation.
-
-    Usage:
-
-    .. code-block:: python
-
-      response = client.models.generate_images(
-        model='imagen-3.0-generate-002',
-        prompt='Man with a dog',
-        config=types.GenerateImagesConfig(
-            number_of_images= 1,
-            include_rai_reason= True,
-        )
-      )
-      response.generated_images[0].image.show()
-      # Shows a man with a dog.
     """
 
     parameter_model = types._GenerateImagesParameters(
@@ -5558,6 +5571,61 @@ def generate_content_stream(
       automatic_function_calling_history.append(func_call_content)
       automatic_function_calling_history.append(func_response_content)
 
+  def generate_images(
+      self,
+      *,
+      model: str,
+      prompt: str,
+      config: Optional[types.GenerateImagesConfigOrDict] = None,
+  ) -> types.GenerateImagesResponse:
+    """Generates images based on a text description and configuration.
+
+    Args:
+      model (str): The model to use.
+      prompt (str): A text description of the images to generate.
+      config (GenerateImagesConfig): Configuration for generation.
+
+    Usage:
+
+    .. code-block:: python
+
+      response = client.models.generate_images(
+        model='imagen-3.0-generate-002',
+        prompt='Man with a dog',
+        config=types.GenerateImagesConfig(
+            number_of_images= 1,
+            include_rai_reason= True,
+        )
+      )
+      response.generated_images[0].image.show()
+      # Shows a man with a dog.
+    """
+    api_response = self._generate_images(
+        model=model,
+        prompt=prompt,
+        config=config,
+    )
+    positive_prompt_safety_attributes = None
+    generated_images = []
+    if not api_response or not api_response.generated_images:
+      return api_response
+
+    for generated_image in api_response.generated_images:
+      if (
+          generated_image.safety_attributes
+          and generated_image.safety_attributes.content_type
+          == 'Positive Prompt'
+      ):
+        positive_prompt_safety_attributes = generated_image.safety_attributes
+      else:
+        generated_images.append(generated_image)
+
+    response = types.GenerateImagesResponse(
+        generated_images=generated_images,
+        positive_prompt_safety_attributes=positive_prompt_safety_attributes,
+    )
+    return response
+
   def edit_image(
       self,
       *,
@@ -5956,7 +6024,7 @@ async def embed_content(
     self._api_client._verify_response(return_value)
     return return_value
 
-  async def generate_images(
+  async def _generate_images(
       self,
       *,
       model: str,
@@ -5969,21 +6037,6 @@ async def generate_images(
       model (str): The model to use.
       prompt (str): A text description of the images to generate.
       config (GenerateImagesConfig): Configuration for generation.
-
-    Usage:
-
-    .. code-block:: python
-
-      response = await client.aio.models.generate_images(
-        model='imagen-3.0-generate-002',
-        prompt='Man with a dog',
-        config=types.GenerateImagesConfig(
-            number_of_images= 1,
-            include_rai_reason= True,
-        )
-      )
-      response.generated_images[0].image.show()
-      # Shows a man with a dog.
     """
 
     parameter_model = types._GenerateImagesParameters(
@@ -7087,6 +7140,61 @@ async def list(
         config,
     )
 
+  async def generate_images(
+      self,
+      *,
+      model: str,
+      prompt: str,
+      config: Optional[types.GenerateImagesConfigOrDict] = None,
+  ) -> types.GenerateImagesResponse:
+    """Generates images based on a text description and configuration.
+
+    Args:
+      model (str): The model to use.
+      prompt (str): A text description of the images to generate.
+      config (GenerateImagesConfig): Configuration for generation.
+
+    Usage:
+
+    .. code-block:: python
+
+      response = await client.aio.models.generate_images(
+        model='imagen-3.0-generate-002',
+        prompt='Man with a dog',
+        config=types.GenerateImagesConfig(
+            number_of_images= 1,
+            include_rai_reason= True,
+        )
+      )
+      response.generated_images[0].image.show()
+      # Shows a man with a dog.
+    """
+    api_response = await self._generate_images(
+        model=model,
+        prompt=prompt,
+        config=config,
+    )
+    positive_prompt_safety_attributes = None
+    generated_images = []
+    if not api_response or not api_response.generated_images:
+      return api_response
+
+    for generated_image in api_response.generated_images:
+      if (
+          generated_image.safety_attributes
+          and generated_image.safety_attributes.content_type
+          == 'Positive Prompt'
+      ):
+        positive_prompt_safety_attributes = generated_image.safety_attributes
+      else:
+        generated_images.append(generated_image)
+
+    response = types.GenerateImagesResponse(
+        generated_images=generated_images,
+        positive_prompt_safety_attributes=positive_prompt_safety_attributes,
+    )
+    return response
+
   async def upscale_image(
       self,
       *,

google/genai/tests/models/test_generate_images.py

@@ -182,4 +182,6 @@ async def test_simple_prompt_async(client):
   )
 
   assert response.generated_images[0].image.image_bytes
-  assert len(response.generated_images) == 2
+  # TODO(tangmatthew): Re-enable this check once the bug is fixed.
+  assert len(response.generated_images) == 1
+  assert response.positive_prompt_safety_attributes.content_type == 'Positive Prompt'

google/genai/types.py

@@ -3512,7 +3512,8 @@ class GenerateImagesConfig(_common.BaseModel):
   )
   include_safety_attributes: Optional[bool] = Field(
       default=None,
-      description="""Whether to report the safety scores of each image in the response.
+      description="""Whether to report the safety scores of each generated image and
+      the positive prompt in the response.
       """,
   )
   include_rai_reason: Optional[bool] = Field(
@@ -3591,7 +3592,8 @@ class GenerateImagesConfigDict(TypedDict, total=False):
       """
 
   include_safety_attributes: Optional[bool]
-  """Whether to report the safety scores of each image in the response.
+  """Whether to report the safety scores of each generated image and
+      the positive prompt in the response.
       """
 
   include_rai_reason: Optional[bool]
@@ -3837,6 +3839,11 @@ class SafetyAttributes(_common.BaseModel):
       description="""List of scores of each categories.
       """,
   )
+  content_type: Optional[str] = Field(
+      default=None,
+      description="""Internal use only.
+      """,
+  )
 
 
 class SafetyAttributesDict(TypedDict, total=False):
@@ -3850,6 +3857,10 @@ class SafetyAttributesDict(TypedDict, total=False):
   """List of scores of each categories.
       """
 
+  content_type: Optional[str]
+  """Internal use only.
+      """
+
 
 SafetyAttributesOrDict = Union[SafetyAttributes, SafetyAttributesDict]
 
@@ -3916,6 +3927,12 @@ class GenerateImagesResponse(_common.BaseModel):
       description="""List of generated images.
       """,
   )
+  positive_prompt_safety_attributes: Optional[SafetyAttributes] = Field(
+      default=None,
+      description="""Safety attributes of the positive prompt. Only populated if
+      ``include_safety_attributes`` is set to True.
+      """,
+  )
 
 
 class GenerateImagesResponseDict(TypedDict, total=False):
@@ -3925,6 +3942,11 @@ class GenerateImagesResponseDict(TypedDict, total=False):
   """List of generated images.
       """
 
+  positive_prompt_safety_attributes: Optional[SafetyAttributesDict]
+  """Safety attributes of the positive prompt. Only populated if
+      ``include_safety_attributes`` is set to True.
+      """
+
 
 GenerateImagesResponseOrDict = Union[
     GenerateImagesResponse, GenerateImagesResponseDict
@@ -4161,7 +4183,8 @@ class EditImageConfig(_common.BaseModel):
   )
   include_safety_attributes: Optional[bool] = Field(
       default=None,
-      description="""Whether to report the safety scores of each image in the response.
+      description="""Whether to report the safety scores of each generated image and
+      the positive prompt in the response.
       """,
   )
   include_rai_reason: Optional[bool] = Field(
@@ -4239,7 +4262,8 @@ class EditImageConfigDict(TypedDict, total=False):
       """
 
   include_safety_attributes: Optional[bool]
-  """Whether to report the safety scores of each image in the response.
+  """Whether to report the safety scores of each generated image and
+      the positive prompt in the response.
       """
 
   include_rai_reason: Optional[bool]