feat: Make genai.Part constructible from PartUnionDict.

PiperOrigin-RevId: 813307548

Author: kibergus

google/genai/tests/types/test_types.py

@@ -18,6 +18,7 @@
 import sys
 import typing
 from typing import Optional, assert_never
+import PIL.Image
 import pydantic
 import pytest
 from ... import types
@@ -262,6 +263,58 @@ def test_factory_method_from_mcp_call_tool_function_response_embedded_resource()
   assert isinstance(my_function_response, types.FunctionResponse)
 
 
+def test_part_constructor_with_string_value():
+  part = types.Part('hello')
+  assert part.text == 'hello'
+  assert part.file_data is None
+  assert part.inline_data is None
+
+
+def test_part_constructor_with_part_value():
+  other_part = types.Part(text='hello from other part')
+  part = types.Part(other_part)
+  assert part.text == 'hello from other part'
+
+
+def test_part_constructor_with_part_dict_value():
+  part = types.Part({'text': 'hello from dict'})
+  assert part.text == 'hello from dict'
+
+
+def test_part_constructor_with_file_data_dict_value():
+  part = types.Part(
+      {'file_uri': 'gs://my-bucket/file-data', 'mime_type': 'text/plain'}
+  )
+  assert part.file_data.file_uri == 'gs://my-bucket/file-data'
+  assert part.file_data.mime_type == 'text/plain'
+
+
+def test_part_constructor_with_kwargs_and_value_fails():
+  with pytest.raises(
+      ValueError, match='Positional and keyword arguments can not be combined'
+  ):
+    types.Part('hello', text='world')
+
+
+def test_part_constructor_with_file_value():
+  f = types.File(
+      uri='gs://my-bucket/my-file',
+      mime_type='text/plain',
+      display_name='test file',
+  )
+  part = types.Part(f)
+  assert part.file_data.file_uri == 'gs://my-bucket/my-file'
+  assert part.file_data.mime_type == 'text/plain'
+  assert part.file_data.display_name == 'test file'
+
+
+def test_part_constructor_with_pil_image():
+  img = PIL.Image.new('RGB', (1, 1), color='red')
+  part = types.Part(img)
+  assert part.inline_data.mime_type == 'image/jpeg'
+  assert isinstance(part.inline_data.data, bytes)
+
+
 class FakeClient:
 
   def __init__(self, vertexai=False) -> None:

google/genai/types.py

@@ -19,12 +19,13 @@
 import datetime
 from enum import Enum, EnumMeta
 import inspect
+import io
 import json
 import logging
 import sys
 import types as builtin_types
 import typing
-from typing import Any, Callable, Literal, Optional, Sequence, Union, _UnionGenericAlias  # type: ignore
+from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Union, _UnionGenericAlias  # type: ignore
 import pydantic
 from pydantic import ConfigDict, Field, PrivateAttr, model_validator
 from typing_extensions import Self, TypedDict
@@ -1260,6 +1261,81 @@ class Part(_common.BaseModel):
       default=None, description="""Optional. Text part (can be code)."""
   )
 
+  def __init__(
+      self,
+      value: Optional['PartUnionDict'] = None,
+      /,
+      *,
+      video_metadata: Optional[VideoMetadata] = None,
+      thought: Optional[bool] = None,
+      inline_data: Optional[Blob] = None,
+      file_data: Optional[FileData] = None,
+      thought_signature: Optional[bytes] = None,
+      function_call: Optional[FunctionCall] = None,
+      code_execution_result: Optional[CodeExecutionResult] = None,
+      executable_code: Optional[ExecutableCode] = None,
+      function_response: Optional[FunctionResponse] = None,
+      text: Optional[str] = None,
+      # Pydantic allows CamelCase in addition to snake_case attribute
+      # names. kwargs here catch these aliases.
+      **kwargs: Any,
+  ):
+    part_dict = dict(
+        video_metadata=video_metadata,
+        thought=thought,
+        inline_data=inline_data,
+        file_data=file_data,
+        thought_signature=thought_signature,
+        function_call=function_call,
+        code_execution_result=code_execution_result,
+        executable_code=executable_code,
+        function_response=function_response,
+        text=text,
+    )
+    part_dict = {k: v for k, v in part_dict.items() if v is not None}
+
+    if part_dict and value is not None:
+      raise ValueError(
+          'Positional and keyword arguments can not be combined when '
+          'initializing a Part.'
+      )
+
+    if value is None:
+      pass
+    elif isinstance(value, str):
+      part_dict['text'] = value
+    elif isinstance(value, File):
+      if not value.uri or not value.mime_type:
+        raise ValueError('file uri and mime_type are required.')
+      part_dict['file_data'] = FileData(
+          file_uri=value.uri,
+          mime_type=value.mime_type,
+          display_name=value.display_name,
+      )
+    elif isinstance(value, dict):
+      try:
+        Part.model_validate(value)
+        part_dict.update(value)  # type: ignore[arg-type]
+      except pydantic.ValidationError:
+        part_dict['file_data'] = FileData.model_validate(value)
+    elif isinstance(value, Part):
+      part_dict.update(value.dict())
+    elif 'image' in value.__class__.__name__.lower():
+      # PIL.Image case.
+
+      suffix = value.format.lower() if value.format else 'jpeg'
+      mimetype = f'image/{suffix}'
+      bytes_io = io.BytesIO()
+      value.save(bytes_io, suffix.upper())
+
+      part_dict['inline_data'] = Blob(
+          data=bytes_io.getvalue(), mime_type=mimetype
+      )
+    else:
+      raise ValueError(f'Unsupported content part type: {type(value)}')
+
+    super().__init__(**part_dict, **kwargs)
+
   def as_image(self) -> Optional['Image']:
     """Returns the part as a PIL Image, or None if the part is not an image."""
     if not self.inline_data: