| 1 | import base64 |
| 2 | import io |
| 3 | import math |
| 4 | import mimetypes |
| 5 | from pathlib import Path |
| 6 | from typing import Any |
| 7 | from urllib.parse import unquote, urlparse |
| 8 | |
| 9 | from PIL import Image |
| 10 | |
| 11 | |
| 12 | def prepare_content(content: Any) -> Any: |
| 13 | if isinstance(content, list): |
| 14 | return [prepare_content(item) for item in content] |
| 15 | if not isinstance(content, dict): |
| 16 | return content |
| 17 | |
| 18 | if content.get("type") == "image_url": |
| 19 | image_url = content.get("image_url") |
| 20 | if isinstance(image_url, dict): |
| 21 | url = str(image_url.get("url", "") or "").strip() |
| 22 | if is_local_ref(url): |
| 23 | return {**content, "image_url": {**image_url, "url": to_data_url(url)}} |
| 24 | elif isinstance(image_url, str): |
| 25 | url = image_url.strip() |
| 26 | if is_local_ref(url): |
| 27 | return {**content, "image_url": {"url": to_data_url(url)}} |
| 28 | |
| 29 | return {key: prepare_content(value) for key, value in content.items()} |
| 30 | |
| 31 | |
| 32 | def is_local_ref(url: str) -> bool: |
| 33 | if not url: |
| 34 | return False |
| 35 | lowered = url.lower() |
| 36 | if lowered.startswith(("http://", "https://", "data:")): |
| 37 | return False |
| 38 | return lowered.startswith("file://") or url.startswith(("/", "./", "../", "~")) |
| 39 | |
| 40 | |
| 41 | def to_data_url(url: str) -> str: |
| 42 | path = resolve_ref(url) |
| 43 | mime_type = mimetypes.guess_type(path.name)[0] |
| 44 | if not mime_type or not mime_type.startswith("image/"): |
| 45 | raise ValueError(f"Image attachment must have an image MIME type: {path}") |
| 46 | encoded = base64.b64encode(path.read_bytes()).decode("utf-8") |
| 47 | return f"data:{mime_type};base64,{encoded}" |
| 48 | |
| 49 | |
| 50 | def resolve_ref(url: str) -> Path: |
| 51 | raw_path = unquote(urlparse(url).path) if url.lower().startswith("file://") else url |
| 52 | path = Path(raw_path).expanduser() |
| 53 | candidates = [path] |
| 54 | if raw_path.startswith("/a0/"): |
| 55 | from helpers import files |
| 56 | |
| 57 | candidates.append(Path(files.fix_dev_path(raw_path))) |
| 58 | elif not path.is_absolute(): |
| 59 | from helpers import files |
| 60 | |
| 61 | candidates.append(Path(files.get_abs_path(raw_path))) |
| 62 | |
| 63 | seen: set[str] = set() |
| 64 | for candidate in candidates: |
| 65 | key = str(candidate) |
| 66 | if key in seen: |
| 67 | continue |
| 68 | seen.add(key) |
| 69 | if candidate.exists() and candidate.is_file(): |
| 70 | return candidate |
| 71 | |
| 72 | raise FileNotFoundError(f"Image attachment path does not exist: {raw_path}") |
| 73 | |
| 74 | |
| 75 | def compress_image(image_data: bytes, *, max_pixels: int = 256_000, quality: int = 50) -> bytes: |
| 76 | """Compress an image by scaling it down and converting to JPEG with quality settings. |
| 77 | |
| 78 | Args: |
| 79 | image_data: Raw image bytes |
| 80 | max_pixels: Maximum number of pixels in the output image (width * height) |
| 81 | quality: JPEG quality setting (1-100) |
| 82 | |
| 83 | Returns: |
| 84 | Compressed image as bytes |
| 85 | """ |
| 86 | # load image from bytes |
| 87 | img = Image.open(io.BytesIO(image_data)) |
| 88 | |
| 89 | # calculate scaling factor to get to max_pixels |
| 90 | current_pixels = img.width * img.height |
| 91 | if current_pixels > max_pixels: |
| 92 | scale = math.sqrt(max_pixels / current_pixels) |
| 93 | new_width = int(img.width * scale) |
| 94 | new_height = int(img.height * scale) |
| 95 | img = img.resize((new_width, new_height), Image.Resampling.LANCZOS) |
| 96 | |
| 97 | # convert to RGB if needed (for JPEG) |
| 98 | if img.mode in ('RGBA', 'P'): |
| 99 | img = img.convert('RGB') |
| 100 | |
| 101 | # save as JPEG with compression |
| 102 | output = io.BytesIO() |
| 103 | img.save(output, format='JPEG', quality=quality, optimize=True) |
| 104 | return output.getvalue() |