Files
hkuds--vimax/tools/image_response.py
T
2026-07-13 12:23:54 +08:00

41 lines
1.1 KiB
Python

from __future__ import annotations
import base64
from io import BytesIO
from typing import Any
from PIL import Image
def image_from_response_part(part: Any) -> Image.Image | None:
inline_data = getattr(part, "inline_data", None)
if inline_data is None and isinstance(part, dict):
inline_data = part.get("inline_data")
if inline_data is None:
return None
as_image = getattr(part, "as_image", None)
if callable(as_image):
image = as_image()
if isinstance(image, Image.Image):
return image
data = _value(inline_data, "data")
if data is None:
return None
if isinstance(data, str):
if data.startswith("data:") and "," in data:
data = data.split(",", 1)[1]
data = base64.b64decode(data)
if isinstance(data, bytearray):
data = bytes(data)
if not isinstance(data, bytes):
return None
return Image.open(BytesIO(data)).convert("RGB")
def _value(obj: Any, key: str) -> Any:
if isinstance(obj, dict):
return obj.get(key)
return getattr(obj, key, None)