Skip to content

Commit 0a93ede

Browse files
kiberguscopybara-github
authored andcommitted
feat: Make genai.Part constructible from PartUnionDict.
PiperOrigin-RevId: 813307548
1 parent 7e0beba commit 0a93ede

File tree

3 files changed

+131
-26
lines changed

3 files changed

+131
-26
lines changed

google/genai/_transformers.py

Lines changed: 1 addition & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -381,31 +381,7 @@ def t_audio_blob(blob: types.BlobOrDict) -> types.Blob:
381381
def t_part(part: Optional[types.PartUnionDict]) -> types.Part:
382382
if part is None:
383383
raise ValueError('content part is required.')
384-
if isinstance(part, str):
385-
return types.Part(text=part)
386-
if _is_duck_type_of(part, types.File):
387-
if not part.uri or not part.mime_type: # type: ignore[union-attr]
388-
raise ValueError('file uri and mime_type are required.')
389-
return types.Part.from_uri(file_uri=part.uri, mime_type=part.mime_type) # type: ignore[union-attr]
390-
if isinstance(part, dict):
391-
try:
392-
return types.Part.model_validate(part)
393-
except pydantic.ValidationError:
394-
return types.Part(file_data=types.FileData.model_validate(part))
395-
if _is_duck_type_of(part, types.Part):
396-
return part # type: ignore[return-value]
397-
398-
if 'image' in part.__class__.__name__.lower():
399-
try:
400-
import PIL.Image
401-
402-
PIL_Image = PIL.Image.Image
403-
except ImportError:
404-
PIL_Image = None
405-
406-
if PIL_Image is not None and isinstance(part, PIL_Image):
407-
return types.Part(inline_data=pil_to_blob(part))
408-
raise ValueError(f'Unsupported content part type: {type(part)}')
384+
return types.Part(part)
409385

410386

411387
def t_parts(

google/genai/tests/types/test_types.py

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import sys
1919
import typing
2020
from typing import Optional, assert_never
21+
import PIL.Image
2122
import pydantic
2223
import pytest
2324
from ... import types
@@ -301,6 +302,58 @@ def test_factory_method_from_mcp_call_tool_function_response_embedded_resource()
301302
assert isinstance(my_function_response, types.FunctionResponse)
302303

303304

305+
def test_part_constructor_with_string_value():
306+
part = types.Part('hello')
307+
assert part.text == 'hello'
308+
assert part.file_data is None
309+
assert part.inline_data is None
310+
311+
312+
def test_part_constructor_with_part_value():
313+
other_part = types.Part(text='hello from other part')
314+
part = types.Part(other_part)
315+
assert part.text == 'hello from other part'
316+
317+
318+
def test_part_constructor_with_part_dict_value():
319+
part = types.Part({'text': 'hello from dict'})
320+
assert part.text == 'hello from dict'
321+
322+
323+
def test_part_constructor_with_file_data_dict_value():
324+
part = types.Part(
325+
{'file_uri': 'gs://my-bucket/file-data', 'mime_type': 'text/plain'}
326+
)
327+
assert part.file_data.file_uri == 'gs://my-bucket/file-data'
328+
assert part.file_data.mime_type == 'text/plain'
329+
330+
331+
def test_part_constructor_with_kwargs_and_value_fails():
332+
with pytest.raises(
333+
ValueError, match='Positional and keyword arguments can not be combined'
334+
):
335+
types.Part('hello', text='world')
336+
337+
338+
def test_part_constructor_with_file_value():
339+
f = types.File(
340+
uri='gs://my-bucket/my-file',
341+
mime_type='text/plain',
342+
display_name='test file',
343+
)
344+
part = types.Part(f)
345+
assert part.file_data.file_uri == 'gs://my-bucket/my-file'
346+
assert part.file_data.mime_type == 'text/plain'
347+
assert part.file_data.display_name == 'test file'
348+
349+
350+
def test_part_constructor_with_pil_image():
351+
img = PIL.Image.new('RGB', (1, 1), color='red')
352+
part = types.Part(img)
353+
assert part.inline_data.mime_type == 'image/jpeg'
354+
assert isinstance(part.inline_data.data, bytes)
355+
356+
304357
class FakeClient:
305358

306359
def __init__(self, vertexai=False) -> None:

google/genai/types.py

Lines changed: 77 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,13 @@
1919
import datetime
2020
from enum import Enum, EnumMeta
2121
import inspect
22+
import io
2223
import json
2324
import logging
2425
import sys
2526
import types as builtin_types
2627
import typing
27-
from typing import Any, Callable, Literal, Optional, Sequence, Union, _UnionGenericAlias # type: ignore
28+
from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Union, _UnionGenericAlias # type: ignore
2829
import pydantic
2930
from pydantic import ConfigDict, Field, PrivateAttr, model_validator
3031
from typing_extensions import Self, TypedDict
@@ -1294,6 +1295,81 @@ class Part(_common.BaseModel):
12941295
default=None, description="""Optional. Text part (can be code)."""
12951296
)
12961297

1298+
def __init__(
1299+
self,
1300+
value: Optional['PartUnionDict'] = None,
1301+
/,
1302+
*,
1303+
video_metadata: Optional[VideoMetadata] = None,
1304+
thought: Optional[bool] = None,
1305+
inline_data: Optional[Blob] = None,
1306+
file_data: Optional[FileData] = None,
1307+
thought_signature: Optional[bytes] = None,
1308+
function_call: Optional[FunctionCall] = None,
1309+
code_execution_result: Optional[CodeExecutionResult] = None,
1310+
executable_code: Optional[ExecutableCode] = None,
1311+
function_response: Optional[FunctionResponse] = None,
1312+
text: Optional[str] = None,
1313+
# Pydantic allows CamelCase in addition to snake_case attribute
1314+
# names. kwargs here catch these aliases.
1315+
**kwargs: Any,
1316+
):
1317+
part_dict = dict(
1318+
video_metadata=video_metadata,
1319+
thought=thought,
1320+
inline_data=inline_data,
1321+
file_data=file_data,
1322+
thought_signature=thought_signature,
1323+
function_call=function_call,
1324+
code_execution_result=code_execution_result,
1325+
executable_code=executable_code,
1326+
function_response=function_response,
1327+
text=text,
1328+
)
1329+
part_dict = {k: v for k, v in part_dict.items() if v is not None}
1330+
1331+
if part_dict and value is not None:
1332+
raise ValueError(
1333+
'Positional and keyword arguments can not be combined when '
1334+
'initializing a Part.'
1335+
)
1336+
1337+
if value is None:
1338+
pass
1339+
elif isinstance(value, str):
1340+
part_dict['text'] = value
1341+
elif isinstance(value, File):
1342+
if not value.uri or not value.mime_type:
1343+
raise ValueError('file uri and mime_type are required.')
1344+
part_dict['file_data'] = FileData(
1345+
file_uri=value.uri,
1346+
mime_type=value.mime_type,
1347+
display_name=value.display_name,
1348+
)
1349+
elif isinstance(value, dict):
1350+
try:
1351+
Part.model_validate(value)
1352+
part_dict.update(value) # type: ignore[arg-type]
1353+
except pydantic.ValidationError:
1354+
part_dict['file_data'] = FileData.model_validate(value)
1355+
elif isinstance(value, Part):
1356+
part_dict.update(value.dict())
1357+
elif 'image' in value.__class__.__name__.lower():
1358+
# PIL.Image case.
1359+
1360+
suffix = value.format.lower() if value.format else 'jpeg'
1361+
mimetype = f'image/{suffix}'
1362+
bytes_io = io.BytesIO()
1363+
value.save(bytes_io, suffix.upper())
1364+
1365+
part_dict['inline_data'] = Blob(
1366+
data=bytes_io.getvalue(), mime_type=mimetype
1367+
)
1368+
else:
1369+
raise ValueError(f'Unsupported content part type: {type(value)}')
1370+
1371+
super().__init__(**part_dict, **kwargs)
1372+
12971373
def as_image(self) -> Optional['Image']:
12981374
"""Returns the part as a PIL Image, or None if the part is not an image."""
12991375
if not self.inline_data:

0 commit comments

Comments
 (0)