| 1 | from __future__ import annotations |
| 2 | |
| 3 | import base64 |
| 4 | import os |
| 5 | from io import BytesIO |
| 6 | from pathlib import Path |
| 7 | from typing import Any |
| 8 | |
| 9 | from PIL import Image, ImageOps |
| 10 | |
| 11 | from .models import ToolResult |
| 12 | |
| 13 | |
| 14 | SUPPORTED_IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp", ".gif"} |
| 15 | |
| 16 | |
| 17 | class ViewImageHandler: |
| 18 | def __init__(self, workspace_root: str | Path, session_index: Any) -> None: |
| 19 | self.workspace_root = Path(workspace_root).resolve() |
| 20 | self.session_index = session_index |
| 21 | |
| 22 | def __call__(self, args: dict[str, Any]) -> ToolResult: |
| 23 | try: |
| 24 | session = self.session_index.get_or_create_active() |
| 25 | session_root = self.session_index.working_dir(session["session_id"]).resolve() |
| 26 | path = _resolve_session_path(self.workspace_root, session_root, args["path"]) |
| 27 | image, original_size = _load_image(path) |
| 28 | try: |
| 29 | data_url, display = _encode_for_model( |
| 30 | image, |
| 31 | max_dimension=_env_int("VIMAX_IMAGE_VIEW_MAX_DIMENSION", 1568, minimum=256), |
| 32 | max_bytes=_env_int("VIMAX_IMAGE_VIEW_MAX_BYTES", 5_000_000, minimum=100_000), |
| 33 | ) |
| 34 | finally: |
| 35 | image.close() |
| 36 | except (OSError, ValueError) as exc: |
| 37 | return ToolResult("view_image", False, str(exc), {"error_type": "invalid_input"}) |
| 38 | |
| 39 | relative = path.relative_to(session_root).as_posix() |
| 40 | workspace_path = path.relative_to(self.workspace_root).as_posix() |
| 41 | metadata = { |
| 42 | "path": relative, |
| 43 | "workspace_path": workspace_path, |
| 44 | "session_id": session["session_id"], |
| 45 | "mime_type": display["mime_type"], |
| 46 | "original_bytes": path.stat().st_size, |
| 47 | "original_width": original_size[0], |
| 48 | "original_height": original_size[1], |
| 49 | "display_width": display["width"], |
| 50 | "display_height": display["height"], |
| 51 | "display_bytes": display["bytes"], |
| 52 | "camera_metadata": _read_camera_metadata(path), |
| 53 | } |
| 54 | return ToolResult( |
| 55 | "view_image", |
| 56 | True, |
| 57 | f"Image loaded for visual inspection: {relative} ({original_size[0]}x{original_size[1]}).", |
| 58 | metadata, |
| 59 | model_content=[{"type": "image_url", "image_url": {"url": data_url, "detail": "high"}}], |
| 60 | ) |
| 61 | |
| 62 | |
| 63 | def _resolve_session_path(workspace_root: Path, session_root: Path, raw: Any) -> Path: |
| 64 | text = str(raw).strip() |
| 65 | if not text: |
| 66 | raise ValueError("view_image path is required") |
| 67 | supplied = Path(text) |
| 68 | if supplied.is_absolute(): |
| 69 | path = supplied.resolve() |
| 70 | elif supplied.parts and supplied.parts[0] == ".working_dir": |
| 71 | path = (workspace_root / supplied).resolve() |
| 72 | else: |
| 73 | path = (session_root / supplied).resolve() |
| 74 | if path != session_root and session_root not in path.parents: |
| 75 | raise ValueError(f"Image path escapes active session workspace: {text}") |
| 76 | if not path.exists(): |
| 77 | raise ValueError(f"Image not found in active session: {text}") |
| 78 | if not path.is_file(): |
| 79 | raise ValueError(f"Image path is not a file: {text}") |
| 80 | return path |
| 81 | |
| 82 | |
| 83 | def _load_image(path: Path) -> tuple[Image.Image, tuple[int, int]]: |
| 84 | if path.suffix.lower() not in SUPPORTED_IMAGE_SUFFIXES: |
| 85 | raise ValueError(f"Unsupported image type: {path.suffix or '<none>'}") |
| 86 | try: |
| 87 | with Image.open(path) as source: |
| 88 | source.seek(0) |
| 89 | original_size = source.size |
| 90 | image = ImageOps.exif_transpose(source).convert("RGB") |
| 91 | image.load() |
| 92 | return image, original_size |
| 93 | except Exception as exc: |
| 94 | raise ValueError(f"Cannot decode image {path.name}: {exc}") from exc |
| 95 | |
| 96 | |
| 97 | def _encode_for_model(image: Image.Image, *, max_dimension: int, max_bytes: int) -> tuple[str, dict[str, Any]]: |
| 98 | rendered = image.copy() |
| 99 | try: |
| 100 | rendered.thumbnail((max_dimension, max_dimension), Image.Resampling.LANCZOS) |
| 101 | quality = 90 |
| 102 | payload = b"" |
| 103 | while quality >= 35: |
| 104 | buffer = BytesIO() |
| 105 | rendered.save(buffer, format="JPEG", quality=quality, optimize=True) |
| 106 | payload = buffer.getvalue() |
| 107 | if len(payload) <= max_bytes: |
| 108 | break |
| 109 | quality -= 10 |
| 110 | while len(payload) > max_bytes and min(rendered.size) > 320: |
| 111 | rendered.thumbnail( |
| 112 | (max(320, int(rendered.width * 0.8)), max(320, int(rendered.height * 0.8))), |
| 113 | Image.Resampling.LANCZOS, |
| 114 | ) |
| 115 | buffer = BytesIO() |
| 116 | rendered.save(buffer, format="JPEG", quality=55, optimize=True) |
| 117 | payload = buffer.getvalue() |
| 118 | if len(payload) > max_bytes: |
| 119 | raise ValueError(f"Image cannot be reduced below configured {max_bytes} byte limit") |
| 120 | encoded = base64.b64encode(payload).decode("ascii") |
| 121 | return f"data:image/jpeg;base64,{encoded}", { |
| 122 | "mime_type": "image/jpeg", |
| 123 | "width": rendered.width, |
| 124 | "height": rendered.height, |
| 125 | "bytes": len(payload), |
| 126 | } |
| 127 | finally: |
| 128 | rendered.close() |
| 129 | |
| 130 | |
| 131 | def _read_camera_metadata(path: Path) -> dict[str, Any]: |
| 132 | try: |
| 133 | with Image.open(path) as source: |
| 134 | exif = source.getexif() |
| 135 | except Exception: |
| 136 | return {} |
| 137 | if not exif: |
| 138 | return {} |
| 139 | metadata: dict[str, Any] = {} |
| 140 | for tag, name in {271: "make", 272: "model", 42036: "lens_model"}.items(): |
| 141 | value = exif.get(tag) |
| 142 | if value: |
| 143 | metadata[name] = str(value).strip() |
| 144 | focal_length = _numeric_exif_value(exif.get(37386)) |
| 145 | if focal_length is not None: |
| 146 | metadata["focal_length_mm"] = round(focal_length, 2) |
| 147 | equivalent = _numeric_exif_value(exif.get(41989)) |
| 148 | if equivalent is not None: |
| 149 | metadata["focal_length_35mm_equivalent"] = round(equivalent, 2) |
| 150 | return metadata |
| 151 | |
| 152 | |
| 153 | def _numeric_exif_value(value: Any) -> float | None: |
| 154 | if value is None: |
| 155 | return None |
| 156 | try: |
| 157 | if isinstance(value, tuple) and len(value) == 2: |
| 158 | denominator = float(value[1]) |
| 159 | return float(value[0]) / denominator if denominator else None |
| 160 | return float(value) |
| 161 | except (TypeError, ValueError, ZeroDivisionError): |
| 162 | return None |
| 163 | |
| 164 | |
| 165 | def _env_int(name: str, default: int, *, minimum: int) -> int: |
| 166 | try: |
| 167 | return max(minimum, int(os.environ.get(name, str(default)))) |
| 168 | except ValueError: |
| 169 | return default |
| 170 |