Rollback making ProcessorPart to inherit from genai.Part

kibergus · kibergus · commit ebcf642b3372 · 2025-11-06T13:21:54.000Z
The initial transition has been made in a faith that frameworks can unify on a shared content representation and avoid unnecessary conversions. However it requires that representation to be flexible enough to accomodate needs of its users. Unfortunately this attempt has failed and adding anything not directly needed for Gemini models to genai.Part is nearly impossible.

PiperOrigin-RevId: 828318675
diff --git a/genai_processors/content_api.py b/genai_processors/content_api.py
@@ -26,10 +26,9 @@
 from genai_processors import mime_types
 from google.genai import types as genai_types
 import PIL.Image
-import pydantic
 
 
-class ProcessorPart(genai_types.Part):
+class ProcessorPart:
   """A wrapper around `Part` with additional metadata.
 
   Represents a single piece of content that can be processed by an agentic
@@ -39,11 +38,6 @@ class ProcessorPart(genai_types.Part):
   belongs to, the MIME type of the content, and arbitrary metadata.
   """
 
-  _metadata: dict[str, Any] = pydantic.PrivateAttr(default_factory=dict)
-  _role: str = pydantic.PrivateAttr(default='')
-  _substream_name: str = pydantic.PrivateAttr(default='')
-  _mimetype: str = pydantic.PrivateAttr(default='')
-
   def __init__(
       self,
       value: 'ProcessorPartTypes',
@@ -56,43 +50,48 @@ def __init__(
     """Constructs a ProcessorPart.
 
     Args:
-      value: The content to use to construct the ProcessorPart. Any keyword
-        arguments after this one overrides any properties in value.
+      value: The content to use to construct the ProcessorPart.
       role: Optional. The producer of the content. In Genai models, must be
         either 'user' or 'model', but the user can set their own semantics.
         Useful to set for multi-turn conversations, otherwise can be empty.
       substream_name: (Optional) ProcessorPart stream can be split into multiple
         independent streams. They may have specific semantics, e.g. a song and
         its lyrics, or can be just alternative responses. Prefer using a default
-        substream with an empty name.
+        substream with an empty name. If the `ProcessorPart` is created using
+        another `ProcessorPart`, this ProcessorPart inherits the existing
+        substream_name, unless it is overridden in this argument.
       mimetype: Mime type of the data.
       metadata: (Optional) Auxiliary information about the part. If the
         `ProcessorPart` is created using another `ProcessorPart` or a
         `content_pb2.Part`, this ProcessorPart inherits the existing metadata,
         unless it is overridden in this argument.
     """
+    super().__init__()
+    self._metadata = {}
+
     match value:
+      case genai_types.Part():
+        self._part = value
       case ProcessorPart():
-        super().__init__(**value.model_dump(exclude_unset=True))
+        self._part = value.part
         role = role or value.role
         substream_name = substream_name or value.substream_name
         mimetype = mimetype or value.mimetype
         self._metadata = value.metadata
-      case genai_types.Part():
-        super().__init__(**value.model_dump(exclude_unset=True))
+        self._metadata.update(metadata or {})
       case str():
-        super().__init__(text=value)
+        self._part = genai_types.Part(text=value)
       case bytes():
         if not mimetype:
           raise ValueError(
               'MIME type must be specified when constructing a ProcessorPart'
               ' from bytes.'
           )
         if is_text(mimetype):
-          super().__init__(text=value.decode('utf-8'))
+          self._part = genai_types.Part(text=value.decode('utf-8'))
         else:
-          super().__init__(
-              inline_data=genai_types.Blob(data=value, mime_type=mimetype)
+          self._part = genai_types.Part.from_bytes(
+              data=value, mime_type=mimetype
           )
       case PIL.Image.Image():
         if mimetype:
@@ -114,10 +113,8 @@ def __init__(
           mimetype = f'image/{suffix}'
         bytes_io = io.BytesIO()
         value.save(bytes_io, suffix.upper())
-        super().__init__(
-            inline_data=genai_types.Blob(
-                data=bytes_io.getvalue(), mime_type=mimetype
-            )
+        self._part = genai_types.Part.from_bytes(
+            data=bytes_io.getvalue(), mime_type=mimetype
         )
       case _:
         raise ValueError(f"Can't construct ProcessorPart from {type(value)}.")
@@ -130,10 +127,10 @@ def __init__(
     if mimetype:
       self._mimetype = mimetype
     # Otherwise, if MIME type is specified using inline data, use that.
-    elif self.inline_data and self.inline_data.mime_type:
-      self._mimetype = self.inline_data.mime_type
+    elif self._part.inline_data and self._part.inline_data.mime_type:
+      self._mimetype = self._part.inline_data.mime_type
     # Otherwise, if text is not empty, assume 'text/plain' MIME type.
-    elif self.text:
+    elif self._part.text:
       self._mimetype = 'text/plain'
     else:
       self._mimetype = ''
@@ -147,23 +144,24 @@ def __repr__(self) -> str:
     if self.role:
       optional_args += f', role={self.role!r}'
     return (
-        f'ProcessorPart({self.to_json_dict()!r},'
+        f'ProcessorPart({self.part.to_json_dict()!r},'
         f' mimetype={self.mimetype!r}{optional_args})'
     )
 
   def __eq__(self, other: Any) -> bool:
     if not isinstance(other, ProcessorPart):
       return False
-    return self.__dict__ == other.__dict__
+    return (
+        self._part == other._part
+        and self._role.lower() == other._role.lower()
+        and self._substream_name.lower() == other._substream_name.lower()
+        and self._metadata == other._metadata
+    )
 
   @property
   def part(self) -> genai_types.Part:
-    """Returns the underlying Genai Part.
-
-    DEPRECATED: Use the ProcessorPart itself, it now inherits from genai.Part.
-    This property is provided for backward compatibility reasons.
-    """
-    return self
+    """Returns the underlying Genai Part."""
+    return self._part
 
   @property
   def role(self) -> str:
@@ -189,10 +187,10 @@ def bytes(self) -> bytes | None:
       Text encoded into bytes or bytes from inline data if the underlying part
       is a Blob.
     """
-    if self.text:
+    if self.part.text:
       return self.text.encode()
-    if isinstance(self.inline_data, genai_types.Blob):
-      return self.inline_data.data
+    if isinstance(self.part.inline_data, genai_types.Blob):
+      return self.part.inline_data.data
     return None
 
   @property
@@ -215,6 +213,25 @@ def mimetype(self) -> str:
     """
     return self._mimetype or 'text/plain'
 
+  @property
+  def text(self) -> str:
+    """Returns part text as string.
+
+    Returns:
+      The text of the part.
+
+    Raises:
+      ValueError if part has no text.
+    """
+    if not mime_types.is_text(self.mimetype):
+      raise ValueError('Part is not text.')
+    return self.part.text or ''
+
+  @text.setter
+  def text(self, value: str) -> None:
+    """Sets part to a text part."""
+    self._part = genai_types.Part(text=value)
+
   @property
   def metadata(self) -> dict[str, Any]:
     """Returns metadata."""
@@ -229,6 +246,16 @@ def get_metadata(self, key: str, default=None) -> Any:
     """Returns metadata for a given key."""
     return self._metadata.get(key, default)
 
+  @property
+  def function_call(self) -> genai_types.FunctionCall | None:
+    """Returns function call."""
+    return self.part.function_call
+
+  @property
+  def function_response(self) -> genai_types.FunctionResponse | None:
+    """Returns function response."""
+    return self.part.function_response
+
   @property
   def tool_cancellation(self) -> str | None:
     """Returns an id of a function call to be cancelled.
@@ -239,13 +266,13 @@ def tool_cancellation(self) -> str | None:
       The id of the function call to be cancelled or None if this part is not a
       tool cancellation from the model.
     """
-    if not self.function_response:
+    if not self.part.function_response:
       return None
-    if self.function_response.name != 'tool_cancellation':
+    if self.part.function_response.name != 'tool_cancellation':
       return None
-    if not self.function_response.response:
+    if not self.part.function_response.response:
       return None
-    return self.function_response.response.get('function_call_id', None)
+    return self.part.function_response.response.get('function_call_id', None)
 
   T = TypeVar('T')
 
@@ -274,8 +301,8 @@ def pil_image(self) -> PIL.Image.Image:
     if not mime_types.is_image(self.mimetype):
       raise ValueError(f'Part is not an image. Mime type is {self.mimetype}.')
     bytes_io = io.BytesIO()
-    if self.inline_data is not None:
-      bytes_io.write(self.inline_data.data)
+    if self.part.inline_data is not None:
+      bytes_io.write(self.part.inline_data.data)
     bytes_io.seek(0)
     return PIL.Image.open(bytes_io)
 
@@ -452,7 +479,7 @@ def to_dict(self) -> dict[str, Any]:
     ```
     """
     return {
-        'part': self.model_dump(mode='json', exclude_none=True),
+        'part': self.part.model_dump(mode='json', exclude_none=True),
         'role': self.role,
         'substream_name': self.substream_name,
         'mimetype': self.mimetype,
@@ -826,9 +853,10 @@ def to_genai_contents(
   """
   processor_content = ProcessorContent(content)
   contents = []
-  for role, content_parts in itertools.groupby(
+  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
diff --git a/genai_processors/tests/content_api_test.py b/genai_processors/tests/content_api_test.py
@@ -72,14 +72,15 @@ def test_from_bytes_with_text_mimetype(self):
     mimetype = 'text/plain'
     part = content_api.ProcessorPart(bytes_data, mimetype=mimetype)
     self.assertEqual(part.text, 'hello')
+    self.assertEqual(part.part.text, 'hello')
 
   def test_from_bytes_with_non_text_mimetype(self):
     bytes_data = b'hello'
     mimetype = 'application/octet-stream'
     part = content_api.ProcessorPart(bytes_data, mimetype=mimetype)
     self.assertEqual(part.bytes, bytes_data)
-    self.assertEqual(part.inline_data.data, bytes_data)
-    self.assertEqual(part.inline_data.mime_type, mimetype)
+    self.assertEqual(part.part.inline_data.data, bytes_data)
+    self.assertEqual(part.part.inline_data.mime_type, mimetype)
 
   def test_eq_part_and_non_part(self):
     part = content_api.ProcessorPart('foo')
@@ -416,17 +417,17 @@ def test_to_genai_contents(self):
     expected_genai_contents = [
         genai_types.Content(
             parts=[
-                content_api.ProcessorPart('part1'),
-                content_api.ProcessorPart('part2'),
+                genai_types.Part(text='part1'),
+                genai_types.Part(text='part2'),
             ],
             role='user',
         ),
         genai_types.Content(
-            parts=[content_api.ProcessorPart('part3')],
+            parts=[genai_types.Part(text='part3')],
             role='model',
         ),
         genai_types.Content(
-            parts=[content_api.ProcessorPart('part4')],
+            parts=[genai_types.Part(text='part4')],
             role='user',
         ),
     ]
@@ -442,7 +443,7 @@ def test_to_genai_contents_single_part(self):
     genai_contents = content_api.to_genai_contents(parts)
     expected_genai_contents = [
         genai_types.Content(
-            parts=[content_api.ProcessorPart('part1')],
+            parts=[genai_types.Part(text='part1')],
             role='user',
         ),
     ]
diff --git a/genai_processors/tests/function_calling_test.py b/genai_processors/tests/function_calling_test.py
@@ -82,7 +82,9 @@ async def test_one_function_call(self):
         output,
         [
             content_api.ProcessorPart.from_function_call(
-                name='get_weather', args={'location': 'London'}
+                name='get_weather',
+                args={'location': 'London'},
+                substream_name=function_calling.FUNCTION_CALL_SUBTREAM_NAME,
             ),
             content_api.ProcessorPart.from_function_response(
                 name='get_weather',
@@ -179,7 +181,9 @@ async def test_max_function_calls(self):
         output,
         [
             content_api.ProcessorPart.from_function_call(
-                name='get_time', args={}
+                name='get_time',
+                args={},
+                substream_name=function_calling.FUNCTION_CALL_SUBTREAM_NAME,
             ),
             content_api.ProcessorPart.from_function_response(
                 name='get_time',
@@ -226,7 +230,9 @@ async def test_failing_function(self):
         output,
         [
             content_api.ProcessorPart.from_function_call(
-                name='failing_function', args={}
+                name='failing_function',
+                args={},
+                substream_name=function_calling.FUNCTION_CALL_SUBTREAM_NAME,
             ),
             content_api.ProcessorPart.from_function_response(
                 name='failing_function',
