| 1 | from __future__ import annotations |
| 2 | |
| 3 | import time |
| 4 | import uuid |
| 5 | from dataclasses import dataclass |
| 6 | from pathlib import Path |
| 7 | from typing import Literal |
| 8 | |
| 9 | from helpers import files, media_artifacts |
| 10 | |
| 11 | |
| 12 | DEFAULT_MAX_IMAGE_BYTES = media_artifacts.DEFAULT_MAX_ARTIFACT_SIZE_BYTES |
| 13 | ImageCategory = Literal["images", "screenshots"] |
| 14 | |
| 15 | |
| 16 | @dataclass(frozen=True) |
| 17 | class ChatImage: |
| 18 | path: str |
| 19 | a0_path: str |
| 20 | mime: str |
| 21 | size: int |
| 22 | |
| 23 | |
| 24 | def screenshot_dir(context_id: str, source: str) -> Path: |
| 25 | return artifact_dir(context_id, category="screenshots", source=source) |
| 26 | |
| 27 | |
| 28 | def artifact_dir( |
| 29 | context_id: str, |
| 30 | *, |
| 31 | category: ImageCategory = "images", |
| 32 | source: str = "vision-load", |
| 33 | ) -> Path: |
| 34 | context_segment = files.safe_file_name(str(context_id or "default")).strip("._") or "default" |
| 35 | safe_category = files.safe_file_name(category).strip("._") or "images" |
| 36 | safe_source = files.safe_file_name(source).strip("._") or "vision-load" |
| 37 | |
| 38 | return Path(files.get_abs_path("usr/chats", context_segment)) / safe_category / safe_source |
| 39 | |
| 40 | |
| 41 | def save_image_bytes( |
| 42 | *, |
| 43 | context_id: str, |
| 44 | payload: bytes, |
| 45 | mime_type: str = "image/png", |
| 46 | category: ImageCategory = "images", |
| 47 | source: str = "vision-load", |
| 48 | preferred_name: str = "", |
| 49 | max_bytes: int | None = DEFAULT_MAX_IMAGE_BYTES, |
| 50 | ) -> ChatImage: |
| 51 | data = bytes(payload or b"") |
| 52 | if not data: |
| 53 | raise media_artifacts.EmptyBase64Data("image payload is empty") |
| 54 | if max_bytes is not None and len(data) > max_bytes: |
| 55 | raise media_artifacts.ArtifactTooLarge(len(data), max_bytes) |
| 56 | |
| 57 | safe_mime = media_artifacts.normalize_mime( |
| 58 | mime_type, |
| 59 | default="image/png", |
| 60 | required_prefix="image/", |
| 61 | ) |
| 62 | default_extension = media_artifacts.guess_extension(safe_mime, ".png") |
| 63 | default_filename = f"{source or 'image'}{default_extension}" |
| 64 | filename = media_artifacts.safe_filename( |
| 65 | preferred_name, |
| 66 | default=default_filename, |
| 67 | default_extension=default_extension, |
| 68 | ) |
| 69 | filename_path = Path(filename) |
| 70 | stem = filename_path.stem or Path(default_filename).stem or "image" |
| 71 | suffix = filename_path.suffix or default_extension |
| 72 | timestamp = time.strftime("%Y%m%d-%H%M%S") |
| 73 | path = artifact_dir(context_id, category=category, source=source) / ( |
| 74 | f"{stem}-{timestamp}-{uuid.uuid4().hex[:8]}{suffix}" |
| 75 | ) |
| 76 | path.parent.mkdir(parents=True, exist_ok=True) |
| 77 | path.write_bytes(data) |
| 78 | return ChatImage( |
| 79 | path=str(path), |
| 80 | a0_path=files.normalize_a0_path(str(path)), |
| 81 | mime=safe_mime, |
| 82 | size=len(data), |
| 83 | ) |
| 84 | |
| 85 | |
| 86 | def save_image_base64( |
| 87 | *, |
| 88 | context_id: str, |
| 89 | data: str, |
| 90 | mime_type: str = "image/png", |
| 91 | category: ImageCategory = "images", |
| 92 | source: str = "vision-load", |
| 93 | preferred_name: str = "", |
| 94 | max_bytes: int | None = DEFAULT_MAX_IMAGE_BYTES, |
| 95 | ) -> ChatImage: |
| 96 | payload = media_artifacts.decode_base64_payload(data, max_bytes=max_bytes) |
| 97 | return save_image_bytes( |
| 98 | context_id=context_id, |
| 99 | payload=payload.payload, |
| 100 | mime_type=mime_type, |
| 101 | category=category, |
| 102 | source=source, |
| 103 | preferred_name=preferred_name, |
| 104 | max_bytes=max_bytes, |
| 105 | ) |
| 106 | |
| 107 | |
| 108 | def save_image_file( |
| 109 | *, |
| 110 | context_id: str, |
| 111 | path: str | Path, |
| 112 | category: ImageCategory = "images", |
| 113 | source: str = "vision-load", |
| 114 | preferred_name: str = "", |
| 115 | max_bytes: int | None = DEFAULT_MAX_IMAGE_BYTES, |
| 116 | ) -> ChatImage: |
| 117 | image_path = Path(path) |
| 118 | payload = image_path.read_bytes() |
| 119 | mime = media_artifacts.normalize_mime( |
| 120 | _guess_image_mime(image_path), |
| 121 | default="image/png", |
| 122 | required_prefix="image/", |
| 123 | ) |
| 124 | return save_image_bytes( |
| 125 | context_id=context_id, |
| 126 | payload=payload, |
| 127 | mime_type=mime, |
| 128 | category=category, |
| 129 | source=source, |
| 130 | preferred_name=preferred_name or image_path.name, |
| 131 | max_bytes=max_bytes, |
| 132 | ) |
| 133 | |
| 134 | |
| 135 | def save_image_data_url( |
| 136 | *, |
| 137 | context_id: str, |
| 138 | data_url: str, |
| 139 | category: ImageCategory = "images", |
| 140 | source: str = "vision-load", |
| 141 | preferred_name: str = "", |
| 142 | max_bytes: int | None = DEFAULT_MAX_IMAGE_BYTES, |
| 143 | ) -> ChatImage: |
| 144 | header, encoded = _split_image_data_url(data_url) |
| 145 | mime = header.removeprefix("data:").split(";", 1)[0] or "image/png" |
| 146 | return save_image_base64( |
| 147 | context_id=context_id, |
| 148 | data=encoded, |
| 149 | mime_type=mime, |
| 150 | category=category, |
| 151 | source=source, |
| 152 | preferred_name=preferred_name, |
| 153 | max_bytes=max_bytes, |
| 154 | ) |
| 155 | |
| 156 | |
| 157 | def materialize_image_ref( |
| 158 | *, |
| 159 | context_id: str, |
| 160 | url: str, |
| 161 | source: str = "", |
| 162 | preferred_name: str = "", |
| 163 | max_bytes: int | None = DEFAULT_MAX_IMAGE_BYTES, |
| 164 | ) -> str: |
| 165 | value = str(url or "").strip() |
| 166 | if not value or not str(context_id or "").strip(): |
| 167 | return value |
| 168 | |
| 169 | resolved_source = source or infer_source(value, preferred_name) |
| 170 | category = category_for_source(resolved_source) |
| 171 | if _is_data_image_url(value): |
| 172 | saved = save_image_data_url( |
| 173 | context_id=context_id, |
| 174 | data_url=value, |
| 175 | category=category, |
| 176 | source=resolved_source, |
| 177 | preferred_name=preferred_name, |
| 178 | max_bytes=max_bytes, |
| 179 | ) |
| 180 | return saved.a0_path |
| 181 | |
| 182 | from helpers import images |
| 183 | |
| 184 | source_path = images.resolve_ref(value) |
| 185 | if is_chat_scoped_path(context_id=context_id, path=source_path): |
| 186 | return files.normalize_a0_path(str(source_path)) |
| 187 | saved = save_image_file( |
| 188 | context_id=context_id, |
| 189 | path=source_path, |
| 190 | category=category, |
| 191 | source=resolved_source, |
| 192 | preferred_name=preferred_name or source_path.name, |
| 193 | max_bytes=max_bytes, |
| 194 | ) |
| 195 | return saved.a0_path |
| 196 | |
| 197 | |
| 198 | def is_chat_scoped_path(*, context_id: str, path: str | Path) -> bool: |
| 199 | if not str(context_id or "").strip(): |
| 200 | return False |
| 201 | try: |
| 202 | target = Path(path).resolve(strict=False) |
| 203 | root = artifact_dir(context_id, category="images", source="vision-load").parents[1].resolve(strict=False) |
| 204 | return target == root or root in target.parents |
| 205 | except OSError: |
| 206 | return False |
| 207 | |
| 208 | |
| 209 | def infer_source(value: str = "", preferred_name: str = "") -> str: |
| 210 | raw = f"{value or ''} {preferred_name or ''}".lower() |
| 211 | if "computer-use" in raw or "computer_use" in raw or "_a0_connector/computer_use" in raw: |
| 212 | return "computer-use" |
| 213 | if "/desktop/screenshots/" in raw or "\\desktop\\screenshots\\" in raw or "desktop-" in raw: |
| 214 | return "desktop" |
| 215 | if ( |
| 216 | "/browser/screenshots/" in raw |
| 217 | or "\\browser\\screenshots\\" in raw |
| 218 | or "host-browser" in raw |
| 219 | or "browser-" in raw |
| 220 | ): |
| 221 | return "browser" |
| 222 | return "vision-load" |
| 223 | |
| 224 | |
| 225 | def category_for_source(source: str) -> ImageCategory: |
| 226 | return "screenshots" if source in {"desktop", "browser", "computer-use"} else "images" |
| 227 | |
| 228 | |
| 229 | def _guess_image_mime(path: Path) -> str: |
| 230 | import mimetypes |
| 231 | |
| 232 | return mimetypes.guess_type(path.name)[0] or "image/png" |
| 233 | |
| 234 | |
| 235 | def _is_data_image_url(value: str) -> bool: |
| 236 | normalized = str(value or "").strip().lower() |
| 237 | return normalized.startswith("data:image/") and ";base64," in normalized |
| 238 | |
| 239 | |
| 240 | def _split_image_data_url(data_url: str) -> tuple[str, str]: |
| 241 | value = str(data_url or "").strip() |
| 242 | if not _is_data_image_url(value) or "," not in value: |
| 243 | raise ValueError("image data URL must be data:image/*;base64,...") |
| 244 | return value.split(",", 1) |