Add a to_genai_content function to content_api to convert ProcessorContent to genai types.

PiperOrigin-RevId: 818528069

Author: aelissee

genai_processors/content_api.py

@@ -18,6 +18,7 @@
 import dataclasses
 import functools
 import io
+import itertools
 import json
 from typing import Any, TypeVar
 
@@ -413,15 +414,13 @@ def from_dict(cls, *, data: dict[str, Any]) -> 'ProcessorPart':
 
     Args:
       data: A JSON-compatible dictionary containing the serialized data for the
-        ProcessorPart.
-
-        It is expected to have the following keys:
-          * 'part' (dict): A dictionary representing the underlying
-            `google.genai.types.Part` object.
-          * 'role' (str): The role of the part (e.g., 'user', 'model').
-          * 'substream_name' (str): The substream name.
-          * 'mimetype' (str): The MIME type of the part.
-          * 'metadata' (dict[str, Any]): Auxiliary metadata.
+        ProcessorPart.  It is expected to have the following keys:
+        * 'part' (dict): A dictionary representing the underlying
+          `google.genai.types.Part` object.
+        * 'role' (str): The role of the part (e.g., 'user', 'model').
+        * 'substream_name' (str): The substream name.
+        * 'mimetype' (str): The MIME type of the part.
+        * 'metadata' (dict[str, Any]): Auxiliary metadata.
 
     Returns:
       A new ProcessorPart instance.
@@ -438,7 +437,7 @@ def from_dict(cls, *, data: dict[str, Any]) -> 'ProcessorPart':
     reconstructed = ProcessorPart.from_dict(data=part_as_dict)
     print(reconstructed)
     ```
-    """
+    """  # fmt: skip
     return cls(
         genai_types.Part.model_validate(data['part']),
         role=data.get('role', ''),
@@ -830,3 +829,29 @@ def to_genai_part(
     raise ValueError(
         f'Unsupported type for to_genai_part: {type(part_content)}'
     )
+
+
+def to_genai_contents(
+    content: ProcessorContentTypes,
+) -> list[genai_types.Content]:
+  """Converts a list of ProcessorParts into a list of Genai Content objects.
+
+  Consecutive parts with the same role are grouped together into a single
+  `genai_types.Content` object.
+
+  Args:
+    content: Processor content, e.g. a list of `ProcessorPartTypes`.
+
+  Returns:
+    A list of `genai_types.Content` objects, where each object represents
+    content from a single role.
+  """
+  processor_content = ProcessorContent(content)
+  contents = []
+  for role, content_part in itertools.groupby(
+      processor_content, lambda p: p.role
+  ):
+    content_parts = [p.part for p in content_part]
+    contents.append(genai_types.Content(parts=content_parts, role=role))
+
+  return contents

genai_processors/core/genai_model.py

@@ -56,6 +56,7 @@
 from typing import Any
 from genai_processors import content_api
 from genai_processors import processor
+from genai_processors import streams
 from genai_processors.core import constrained_decoding
 from google.genai import client
 from google.genai import types as genai_types
@@ -164,26 +165,13 @@ async def _generate_from_api(
       self, content: AsyncIterable[content_api.ProcessorPartTypes]
   ) -> AsyncIterable[content_api.ProcessorPartTypes]:
     """Internal method to call the GenAI API and stream results."""
-    turn = genai_types.Content(parts=[])
-    contents = []
-    async for content_part in content:
-      content_part = content_api.ProcessorPart(content_part)
-      if turn.role and content_part.role != turn.role:
-        contents.append(turn)
-        turn = genai_types.Content(parts=[])
-
-      turn.role = content_part.role or 'user'
-      turn.parts.append(content_api.to_genai_part(content_part))  # pylint: disable=attribute-error
-
-    if turn.role:
-      contents.append(turn)
-
+    contents = await streams.gather_stream(content)
     if not contents:
       return
 
     async for res in await self._client.aio.models.generate_content_stream(
         model=self._model_name,
-        contents=contents,
+        contents=content_api.to_genai_contents(contents),
         config=self._generate_content_config,
     ):
       res: genai_types.GenerateContentResponse = res

genai_processors/tests/content_api_test.py

@@ -407,5 +407,48 @@ def test_eq_content_and_non_content(self):
     content = content_api.ProcessorContent('foo')
     self.assertNotEqual(content, object())
 
+  def test_to_genai_contents(self):
+    parts = [
+        content_api.ProcessorPart('part1', role='user'),
+        content_api.ProcessorPart('part2', role='user'),
+        content_api.ProcessorPart('part3', role='model'),
+        content_api.ProcessorPart('part4', role='user'),
+    ]
+    expected_genai_contents = [
+        genai_types.Content(
+            parts=[
+                genai_types.Part(text='part1'),
+                genai_types.Part(text='part2'),
+            ],
+            role='user',
+        ),
+        genai_types.Content(
+            parts=[genai_types.Part(text='part3')],
+            role='model',
+        ),
+        genai_types.Content(
+            parts=[genai_types.Part(text='part4')],
+            role='user',
+        ),
+    ]
+    genai_contents = content_api.to_genai_contents(parts)
+    self.assertEqual(genai_contents, expected_genai_contents)
+
+  def test_to_genai_contents_empty(self):
+    genai_contents = content_api.to_genai_contents([])
+    self.assertEmpty(genai_contents)
+
+  def test_to_genai_contents_single_part(self):
+    parts = [content_api.ProcessorPart('part1', role='user')]
+    genai_contents = content_api.to_genai_contents(parts)
+    expected_genai_contents = [
+        genai_types.Content(
+            parts=[genai_types.Part(text='part1')],
+            role='user',
+        ),
+    ]
+    self.assertEqual(genai_contents, expected_genai_contents)
+
+
 if __name__ == '__main__':
   absltest.main()