|
19 | 19 | import datetime |
20 | 20 | from enum import Enum, EnumMeta |
21 | 21 | import inspect |
| 22 | +import io |
22 | 23 | import json |
23 | 24 | import logging |
24 | 25 | import sys |
25 | 26 | import types as builtin_types |
26 | 27 | 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 |
28 | 29 | import pydantic |
29 | 30 | from pydantic import ConfigDict, Field, PrivateAttr, model_validator |
30 | 31 | from typing_extensions import Self, TypedDict |
@@ -1294,6 +1295,81 @@ class Part(_common.BaseModel): |
1294 | 1295 | default=None, description="""Optional. Text part (can be code).""" |
1295 | 1296 | ) |
1296 | 1297 |
|
| 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 | + |
1297 | 1373 | def as_image(self) -> Optional['Image']: |
1298 | 1374 | """Returns the part as a PIL Image, or None if the part is not an image.""" |
1299 | 1375 | if not self.inline_data: |
|
0 commit comments