|
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 |
@@ -1376,6 +1377,81 @@ class Part(_common.BaseModel): |
1376 | 1377 | description="""Optional. Video metadata. The metadata should only be specified while the video data is presented in inline_data or file_data.""", |
1377 | 1378 | ) |
1378 | 1379 |
|
| 1380 | + def __init__( |
| 1381 | + self, |
| 1382 | + value: Optional['PartUnionDict'] = None, |
| 1383 | + /, |
| 1384 | + *, |
| 1385 | + video_metadata: Optional[VideoMetadata] = None, |
| 1386 | + thought: Optional[bool] = None, |
| 1387 | + inline_data: Optional[Blob] = None, |
| 1388 | + file_data: Optional[FileData] = None, |
| 1389 | + thought_signature: Optional[bytes] = None, |
| 1390 | + function_call: Optional[FunctionCall] = None, |
| 1391 | + code_execution_result: Optional[CodeExecutionResult] = None, |
| 1392 | + executable_code: Optional[ExecutableCode] = None, |
| 1393 | + function_response: Optional[FunctionResponse] = None, |
| 1394 | + text: Optional[str] = None, |
| 1395 | + # Pydantic allows CamelCase in addition to snake_case attribute |
| 1396 | + # names. kwargs here catch these aliases. |
| 1397 | + **kwargs: Any, |
| 1398 | + ): |
| 1399 | + part_dict = dict( |
| 1400 | + video_metadata=video_metadata, |
| 1401 | + thought=thought, |
| 1402 | + inline_data=inline_data, |
| 1403 | + file_data=file_data, |
| 1404 | + thought_signature=thought_signature, |
| 1405 | + function_call=function_call, |
| 1406 | + code_execution_result=code_execution_result, |
| 1407 | + executable_code=executable_code, |
| 1408 | + function_response=function_response, |
| 1409 | + text=text, |
| 1410 | + ) |
| 1411 | + part_dict = {k: v for k, v in part_dict.items() if v is not None} |
| 1412 | + |
| 1413 | + if part_dict and value is not None: |
| 1414 | + raise ValueError( |
| 1415 | + 'Positional and keyword arguments can not be combined when ' |
| 1416 | + 'initializing a Part.' |
| 1417 | + ) |
| 1418 | + |
| 1419 | + if value is None: |
| 1420 | + pass |
| 1421 | + elif isinstance(value, str): |
| 1422 | + part_dict['text'] = value |
| 1423 | + elif isinstance(value, File): |
| 1424 | + if not value.uri or not value.mime_type: |
| 1425 | + raise ValueError('file uri and mime_type are required.') |
| 1426 | + part_dict['file_data'] = FileData( |
| 1427 | + file_uri=value.uri, |
| 1428 | + mime_type=value.mime_type, |
| 1429 | + display_name=value.display_name, |
| 1430 | + ) |
| 1431 | + elif isinstance(value, dict): |
| 1432 | + try: |
| 1433 | + Part.model_validate(value) |
| 1434 | + part_dict.update(value) # type: ignore[arg-type] |
| 1435 | + except pydantic.ValidationError: |
| 1436 | + part_dict['file_data'] = FileData.model_validate(value) |
| 1437 | + elif isinstance(value, Part): |
| 1438 | + part_dict.update(value.dict()) |
| 1439 | + elif 'image' in value.__class__.__name__.lower(): |
| 1440 | + # PIL.Image case. |
| 1441 | + |
| 1442 | + suffix = value.format.lower() if value.format else 'jpeg' |
| 1443 | + mimetype = f'image/{suffix}' |
| 1444 | + bytes_io = io.BytesIO() |
| 1445 | + value.save(bytes_io, suffix.upper()) |
| 1446 | + |
| 1447 | + part_dict['inline_data'] = Blob( |
| 1448 | + data=bytes_io.getvalue(), mime_type=mimetype |
| 1449 | + ) |
| 1450 | + else: |
| 1451 | + raise ValueError(f'Unsupported content part type: {type(value)}') |
| 1452 | + |
| 1453 | + super().__init__(**part_dict, **kwargs) |
| 1454 | + |
1379 | 1455 | def as_image(self) -> Optional['Image']: |
1380 | 1456 | """Returns the part as a PIL Image, or None if the part is not an image.""" |
1381 | 1457 | if not self.inline_data: |
|
0 commit comments