Support multimodal content in ProcessorPart.from_function_response.

PiperOrigin-RevId: 828443519

Author: kibergus

genai_processors/content_api.py

@@ -20,7 +20,7 @@
 import io
 import itertools
 import json
-from typing import Any, TypeVar
+from typing import Any, TypeVar, Union
 
 from absl import logging
 from genai_processors import mime_types
@@ -328,23 +328,66 @@ def from_function_response(
       cls,
       *,
       name: str,
-      response: dict[str, Any],
+      response: Union[dict[str, Any], 'ProcessorContentTypes'],
       function_call_id: str | None = None,
       will_continue: bool | None = None,
       scheduling: genai_types.FunctionResponseScheduling | None = None,
       **kwargs,
   ) -> 'ProcessorPart':
-    """Constructs a ProcessorPart as a function response."""
-    part = genai_types.Part(
-        function_response=genai_types.FunctionResponse(
-            id=function_call_id,
-            name=name,
-            response=response,
-            will_continue=will_continue,
-            scheduling=scheduling,
-        )
+    """Constructs a ProcessorPart as a function response.
+
+    Args:
+      name: The name of the invoked function.
+      response: The value returned by the function. It can be a
+        JSON-serializable object or `ProcessorContentTypes`. If response is a
+        JSON-serializable or is text-only it will be stored in the '.result'
+        field of the function response. Otherwise it will be stored in `.parts`.
+      function_call_id: The ID of the function call this is a response to if
+        known.
+      will_continue: Whether this is the last response from a generator
+        function.
+      scheduling: The scheduling policy for the function response. Controls
+        whether the response generation will be triggered immediately or the
+        function response is just added to the context.
+      **kwargs: Additional keyword arguments to pass to the `ProcessorPart`
+        constructor.
+
+    Returns:
+      A ProcessorPart representing the function response.
+    """
+    function_response_args = dict(
+        id=function_call_id,
+        name=name,
+        will_continue=will_continue,
+        scheduling=scheduling,
     )
-    return cls(part, **kwargs)
+    # For maximal compatibility we first try to interpret response as JSON
+    # serializable object. This is what tools in Gemini API return historically.
+    try:
+      function_response = genai_types.FunctionResponse(
+          response={'result': response}, **function_response_args
+      )
+      function_response.json()
+    except ValueError:
+      # Response is not JSON serializable. Then try to construct content.
+      response_content = ProcessorContent(response)
+      try:
+        function_response = genai_types.FunctionResponse(
+            response={'result': as_text(response_content, strict=True)},
+            **function_response_args,
+        )
+      except ValueError:
+        parts = [
+            genai_types.FunctionResponsePart.from_bytes(
+                data=part.bytes, mime_type=part.mimetype
+            )
+            for part in response_content
+        ]
+        function_response = genai_types.FunctionResponse(
+            parts=parts, **function_response_args
+        )
+
+    return cls(genai_types.Part(function_response=function_response), **kwargs)
 
   @classmethod
   def from_executable_code(

genai_processors/tests/content_api_test.py

@@ -82,6 +82,67 @@ def test_from_bytes_with_non_text_mimetype(self):
     self.assertEqual(part.part.inline_data.data, bytes_data)
     self.assertEqual(part.part.inline_data.mime_type, mimetype)
 
+  def test_from_function_response_with_text_content(self):
+    part = content_api.ProcessorPart.from_function_response(
+        name='foo',
+        response=content_api.ProcessorContent('bar'),
+    )
+    self.assertEqual(part.function_response.name, 'foo')
+    self.assertEqual(part.function_response.response, {'result': 'bar'})
+    self.assertIsNone(part.function_response.parts)
+
+  def test_from_function_response_with_dataclass(self):
+    part = content_api.ProcessorPart.from_function_response(
+        name='foo',
+        response=Dataclass(foo='foo', bar=1),
+    )
+    self.assertEqual(part.function_response.name, 'foo')
+    self.assertEqual(
+        part.function_response.response, {'result': Dataclass(foo='foo', bar=1)}
+    )
+    self.assertIsNone(part.function_response.parts)
+
+  def test_from_function_response_with_image_content(self):
+    image_bytes = _png_image_bytes()
+    part = content_api.ProcessorPart.from_function_response(
+        name='foo',
+        response=content_api.ProcessorPart(image_bytes, mimetype='image/png'),
+    )
+    self.assertEqual(part.function_response.name, 'foo')
+    self.assertIsNone(part.function_response.response)
+    self.assertLen(part.function_response.parts, 1)
+    self.assertEqual(
+        part.function_response.parts[0].inline_data.mime_type, 'image/png'
+    )
+    self.assertEqual(
+        part.function_response.parts[0].inline_data.data, image_bytes
+    )
+
+  def test_from_function_response_with_mixed_content(self):
+    image_bytes = _png_image_bytes()
+    part = content_api.ProcessorPart.from_function_response(
+        name='foo',
+        response=[
+            'Here is a black cat in a black room: ',
+            content_api.ProcessorPart(image_bytes, mimetype='image/png'),
+        ],
+    )
+    self.assertEqual(
+        part.function_response,
+        genai_types.FunctionResponse(
+            name='foo',
+            parts=[
+                genai_types.FunctionResponsePart.from_bytes(
+                    data=b'Here is a black cat in a black room: ',
+                    mime_type='text/plain',
+                ),
+                genai_types.FunctionResponsePart.from_bytes(
+                    data=image_bytes, mime_type='image/png'
+                ),
+            ],
+        ),
+    )
+
   def test_eq_part_and_non_part(self):
     part = content_api.ProcessorPart('foo')
     self.assertNotEqual(part, object())