copper-pilot-cli 0.1.1__py3-none-any.whl

copper_pilot_cli/media_utils.py

@@ -0,0 +1,626 @@
+"""Utilities for handling image and video media from clipboard and files."""
+
+import base64
+import io
+import logging
+import os
+import pathlib
+import re
+import shutil
+
+# S404: subprocess needed for clipboard access via pngpaste/osascript
+import subprocess  # noqa: S404
+import sys
+import tempfile
+from collections import Counter
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from langchain_core.messages.content import VideoContentBlock
+
+logger = logging.getLogger(__name__)
+
+IMAGE_EXTENSIONS: frozenset[str] = frozenset(
+    {
+        ".png",
+        ".jpg",
+        ".jpeg",
+        ".gif",
+        ".bmp",
+        ".tiff",
+        ".tif",
+        ".webp",
+        ".ico",
+    }
+)
+"""Common image file extensions supported by PIL."""
+
+VIDEO_EXTENSIONS: frozenset[str] = frozenset(
+    {
+        ".mp4",
+        ".mov",
+        ".avi",
+        ".webm",
+        ".m4v",
+        ".wmv",
+    }
+)
+"""Video file extensions with validated magic-byte support."""
+
+MAX_MEDIA_BYTES: int = 20 * 1024 * 1024
+"""Maximum media file size (20 MB). Keeps base64 payload under ~27 MB."""
+
+
+def strip_media_placeholders(
+    text: str,
+    placeholders: Iterable[str],
+    *,
+    placeholder_spans: Iterable[tuple[int, int]] | None = None,
+) -> str:
+    """Remove display-only media placeholders from user text.
+
+    Placeholders like `[image 1]` are inserted into the terminal input purely for
+    display; the actual media travels as structured content blocks. They must not
+    leak into the canonical model-facing message or LangSmith trace as if the user
+    typed them.
+
+    When available, tracked placeholder spans identify the exact display tokens to
+    strip so user-authored literal duplicates with the same token are preserved.
+    The token fallback removes one matching occurrence per tracked media item for
+    callers that only have placeholder text.
+
+    Args:
+        text: Raw user text that may contain media placeholders.
+        placeholders: Exact placeholder tokens for the media actually attached to
+            this message (e.g. ``["[image 1]", "[video 1]"]``).
+        placeholder_spans: Exact `(start, end)` spans for tracked display tokens
+            in `text`, when known.
+
+    Returns:
+        Text with the given media placeholders removed and surrounding whitespace
+        tidied. Newlines are preserved so multi-line prompts keep their structure.
+        Returns an empty string when only whitespace remains after removal, so
+        callers can treat a placeholder-only message as having no text block.
+    """
+    tokens = [p for p in placeholders if p]
+    if not tokens:
+        return text
+
+    valid_spans = _valid_placeholder_spans(text, tokens, placeholder_spans)
+    spans = [(start, end) for start, end, _token in valid_spans]
+    counts = Counter(tokens)
+    for _start, _end, token in valid_spans:
+        counts[token] -= 1
+    for token, count in counts.items():
+        if count <= 0:
+            continue
+        pattern = re.compile(r"[ \t]*" + re.escape(token))
+        matches = list(pattern.finditer(text))
+        if len(matches) > count:
+            # No (or too few) tracked spans, yet the token appears more times
+            # than we intend to strip: we can't tell the display token from a
+            # user-typed literal and fall back to removing the leading
+            # occurrence(s). That guess can strip a literal and leave the real
+            # display token behind, so leave a breadcrumb (never the text).
+            logger.debug(
+                "Ambiguous media placeholder strip: removing %d of %d "
+                "occurrences of %r with no tracked span to disambiguate",
+                count,
+                len(matches),
+                token,
+            )
+        for index, match in enumerate(matches):
+            if index >= count:
+                break
+            spans.append(match.span())
+
+    # Only strip spaces/tabs (not newlines) so code indentation on lines after a
+    # removed placeholder is preserved. A full .strip() would collapse
+    # "[image 1]\n    def foo():" to "def foo():", losing the leading indent.
+    cleaned = text
+    for start, end in sorted(spans, reverse=True):
+        cleaned = cleaned[:start] + cleaned[end:]
+    cleaned = cleaned.strip(" \t")
+    return cleaned if cleaned.strip() else ""
+
+
+def _valid_placeholder_spans(
+    text: str,
+    tokens: list[str],
+    spans: Iterable[tuple[int, int]] | None,
+) -> list[tuple[int, int, str]]:
+    """Return valid display placeholder spans expanded over adjacent padding.
+
+    A span is kept only when it is in range and its slice equals one of the
+    bound tokens, so a stale span (e.g. left by an offset shift) is dropped
+    rather than used to delete arbitrary text; the caller then falls back to
+    token matching for that item.
+
+    Args:
+        text: Text the spans index into.
+        tokens: Bound placeholder tokens for the attached media.
+        spans: Candidate `(start, end)` display-token spans, or `None`.
+
+    Returns:
+        `(start, end, token)` triples with `start` expanded left over any
+        adjacent spaces/tabs, for spans that validate against `text`.
+    """
+    if spans is None:
+        return []
+
+    token_set = set(tokens)
+    valid: list[tuple[int, int, str]] = []
+    for start, end in spans:
+        if not (0 <= start < end <= len(text)):
+            continue
+        token = text[start:end]
+        if token not in token_set:
+            continue
+        expanded_start = start
+        while expanded_start > 0 and text[expanded_start - 1] in " \t":
+            expanded_start -= 1
+        valid.append((expanded_start, end, token))
+    return valid
+
+
+def _get_executable(name: str) -> str | None:
+    """Get full path to an executable using shutil.which().
+
+    Args:
+        name: Name of the executable to find
+
+    Returns:
+        Full path to executable, or None if not found.
+    """
+    return shutil.which(name)
+
+
+@dataclass
+class ImageData:
+    """Represents a pasted image with its base64 encoding."""
+
+    base64_data: str
+    format: str  # "png", "jpeg", etc.
+    placeholder: str  # Display text like "[image 1]"
+    placeholder_span: tuple[int, int] | None = None
+
+    def to_message_content(self) -> dict:
+        """Convert to LangChain message content format.
+
+        Returns:
+            Dict with type and image_url for multimodal messages.
+        """
+        return {
+            "type": "image_url",
+            "image_url": {"url": f"data:image/{self.format};base64,{self.base64_data}"},
+        }
+
+
+@dataclass
+class VideoData:
+    """Represents a pasted video with its base64 encoding."""
+
+    base64_data: str
+    format: str  # "mp4", "quicktime", etc.
+    placeholder: str  # Display text like "[video 1]"
+    placeholder_span: tuple[int, int] | None = None
+
+    def to_message_content(self) -> "VideoContentBlock":
+        """Convert to LangChain `VideoContentBlock` format.
+
+        Returns:
+            `VideoContentBlock` with base64 data and mime_type.
+        """
+        from langchain_core.messages.content import create_video_block
+
+        return create_video_block(
+            base64=self.base64_data,
+            mime_type=f"video/{self.format}",
+        )
+
+
+def get_clipboard_image() -> ImageData | None:
+    """Attempt to read an image from the system clipboard.
+
+    Supports macOS via `pngpaste` or `osascript`.
+
+    Returns:
+        ImageData if an image is found, None otherwise.
+    """
+    if sys.platform == "darwin":
+        return _get_macos_clipboard_image()
+    logger.warning(
+        "Clipboard image paste is not supported on %s. "
+        "Only macOS is currently supported. "
+        "You can still attach images by dragging and dropping file paths.",
+        sys.platform,
+    )
+    return None
+
+
+def get_image_from_path(path: pathlib.Path) -> ImageData | None:
+    """Read and encode an image file from disk.
+
+    Args:
+        path: Path to the image file.
+
+    Returns:
+        `ImageData` when the file is a valid image, otherwise `None`.
+    """
+    from PIL import Image, UnidentifiedImageError
+
+    try:
+        file_size = path.stat().st_size
+        if file_size == 0:
+            logger.debug("Image file is empty: %s", path)
+            return None
+        if file_size > MAX_MEDIA_BYTES:
+            logger.warning(
+                "Image file %s is too large (%d MB, max %d MB)",
+                path,
+                file_size // (1024 * 1024),
+                MAX_MEDIA_BYTES // (1024 * 1024),
+            )
+            return None
+
+        image_bytes = path.read_bytes()
+        if not image_bytes:
+            return None
+
+        with Image.open(io.BytesIO(image_bytes)) as image:
+            image_format = (image.format or "").lower()
+
+        if image_format == "jpg":
+            image_format = "jpeg"
+        if not image_format:
+            suffix = path.suffix.lower().removeprefix(".")
+            image_format = "jpeg" if suffix == "jpg" else suffix
+        if not image_format:
+            image_format = "png"
+
+        return ImageData(
+            base64_data=encode_to_base64(image_bytes),
+            format=image_format,
+            placeholder="[image]",
+        )
+    except (UnidentifiedImageError, OSError) as e:
+        logger.debug("Failed to load image from %s: %s", path, e, exc_info=True)
+        return None
+
+
+def _detect_video_format(data: bytes) -> str | None:
+    """Detect video MIME subtype from magic bytes.
+
+    Args:
+        data: Raw file bytes (at least 12 bytes for reliable detection).
+
+    Returns:
+        MIME subtype (e.g. "mp4", "webm") or `None` if unrecognized.
+    """
+    min_avi_len = 12
+    if data[4:8] == b"ftyp":
+        # ftyp box: major brand at bytes 8-12 distinguishes MOV vs MP4
+        brand = data[8:12]
+        if brand == b"qt  ":
+            return "quicktime"
+        return "mp4"
+    if data[:4] == b"RIFF" and len(data) >= min_avi_len and data[8:12] == b"AVI ":
+        return "avi"
+    if data[:4] == b"\x30\x26\xb2\x75":  # ASF/WMV
+        return "x-ms-wmv"
+    if data[:4] == b"\x1a\x45\xdf\xa3":  # WebM/Matroska (EBML header)
+        return "webm"
+    return None
+
+
+def get_video_from_path(path: pathlib.Path) -> VideoData | None:
+    """Read and encode a video file from disk.
+
+    Args:
+        path: Path to the video file.
+
+    Returns:
+        `VideoData` when the file is a valid video, otherwise `None`.
+    """
+    suffix = path.suffix.lower()
+    if suffix not in VIDEO_EXTENSIONS:
+        return None
+
+    try:
+        file_size = path.stat().st_size
+        if file_size == 0:
+            logger.debug("Video file is empty: %s", path)
+            return None
+        if file_size > MAX_MEDIA_BYTES:
+            logger.warning(
+                "Video file %s is too large (%d MB, max %d MB)",
+                path,
+                file_size // (1024 * 1024),
+                MAX_MEDIA_BYTES // (1024 * 1024),
+            )
+            return None
+
+        video_bytes = path.read_bytes()
+
+        # Validate it's a real video file by checking magic bytes
+        # MP4 starts with ftyp, MOV also uses ftyp, AVI starts with RIFF
+        min_video_len = 8
+        if len(video_bytes) < min_video_len:
+            logger.debug("Video file too small (%d bytes): %s", len(video_bytes), path)
+            return None
+
+        # Detect format from magic bytes (not extension) so renamed files
+        # get the correct MIME type.
+        detected_format = _detect_video_format(video_bytes)
+        if detected_format is None:
+            logger.warning(
+                "Video file %s has unrecognized signature for extension '%s'; "
+                "skipping. If this is a valid video, the format may not be "
+                "supported yet.",
+                path,
+                suffix,
+            )
+            return None
+
+        return VideoData(
+            base64_data=encode_to_base64(video_bytes),
+            format=detected_format,
+            placeholder="[video]",
+        )
+    except OSError as e:
+        logger.warning("Failed to load video from %s: %s", path, e, exc_info=True)
+        return None
+
+
+def is_media_path(path: pathlib.Path) -> bool:
+    """Return whether a path's extension is a known image or video extension.
+
+    This is a cheap, extension-only check (no file read or decode). Use it to
+    classify a dropped file path as media without loading it, e.g. to reject a
+    dragged image in a text-only input.
+
+    It is a heuristic, not a support check, and disagrees with the loaders in
+    both directions: `get_image_from_path` accepts anything Pillow can decode
+    regardless of extension, while `get_video_from_path` additionally requires a
+    recognized magic-byte signature. Call `get_media_from_path` when the answer
+    must match what can actually be attached.
+
+    Args:
+        path: Path whose suffix is inspected.
+
+    Returns:
+        `True` when the suffix is a known image or video extension.
+    """
+    suffix = path.suffix.lower()
+    return suffix in IMAGE_EXTENSIONS or suffix in VIDEO_EXTENSIONS
+
+
+def get_media_from_path(path: pathlib.Path) -> ImageData | VideoData | None:
+    """Try to load a file as an image first, then as a video.
+
+    Args:
+        path: Path to the media file.
+
+    Returns:
+        `ImageData` or `VideoData` if the file is valid media, otherwise `None`.
+    """
+    result: ImageData | VideoData | None = get_image_from_path(path)
+    if result is not None:
+        return result
+    return get_video_from_path(path)
+
+
+def _get_macos_clipboard_image() -> ImageData | None:
+    """Get clipboard image on macOS using pngpaste or osascript.
+
+    First tries pngpaste (faster if installed), then falls back to osascript.
+
+    Returns:
+        ImageData if an image is found, None otherwise.
+    """
+    from PIL import Image, UnidentifiedImageError
+
+    # Try pngpaste first (fast if installed)
+    pngpaste_path = _get_executable("pngpaste")
+    if pngpaste_path:
+        try:
+            # S603: pngpaste_path is validated via shutil.which(), args are hardcoded
+            result = subprocess.run(  # noqa: S603
+                [pngpaste_path, "-"],
+                capture_output=True,
+                check=False,
+                timeout=2,
+            )
+            if result.returncode == 0 and result.stdout:
+                # Successfully got PNG data - validate it's a real image
+                try:
+                    Image.open(io.BytesIO(result.stdout))
+                    base64_data = base64.b64encode(result.stdout).decode("utf-8")
+                    return ImageData(
+                        base64_data=base64_data,
+                        format="png",  # 'pngpaste -' always outputs PNG
+                        placeholder="[image]",
+                    )
+                except (
+                    # UnidentifiedImageError: corrupted or non-image data
+                    UnidentifiedImageError,
+                    OSError,  # OSError: I/O errors during image processing
+                ) as e:
+                    logger.debug("Invalid image data from pngpaste: %s", e, exc_info=True)
+        except FileNotFoundError:
+            # pngpaste not installed - expected on systems without it
+            logger.debug("pngpaste not found, falling back to osascript")
+        except subprocess.TimeoutExpired:
+            logger.debug("pngpaste timed out after 2 seconds")
+
+    # Fallback to osascript with temp file (built-in but slower)
+    return _get_clipboard_via_osascript()
+
+
+def _get_clipboard_via_osascript() -> ImageData | None:
+    """Get clipboard image via osascript using a temp file.
+
+    osascript outputs data in a special format that can't be captured as raw binary,
+    so we write to a temp file instead.
+
+    Returns:
+        ImageData if an image is found, None otherwise.
+    """
+    from PIL import Image, UnidentifiedImageError
+
+    # Get osascript path - it's a macOS builtin so should always exist
+    osascript_path = _get_executable("osascript")
+    if not osascript_path:
+        return None
+
+    # Create a temp file for the image
+    fd, temp_path = tempfile.mkstemp(suffix=".png")
+    os.close(fd)
+
+    try:
+        # First check if clipboard has PNG data
+        # S603: osascript_path is validated via shutil.which(), args are hardcoded
+        check_result = subprocess.run(  # noqa: S603
+            [osascript_path, "-e", "clipboard info"],
+            capture_output=True,
+            check=False,
+            timeout=2,
+            text=True,
+        )
+
+        if check_result.returncode != 0:
+            return None
+
+        # Check for PNG or TIFF in clipboard info
+        clipboard_info = check_result.stdout.lower()
+        if "pngf" not in clipboard_info and "tiff" not in clipboard_info:
+            return None
+
+        # Try to get PNG first, fall back to TIFF
+        if "pngf" in clipboard_info:
+            get_script = f"""
+            set pngData to the clipboard as «class PNGf»
+            set theFile to open for access POSIX file "{temp_path}" with write permission
+            write pngData to theFile
+            close access theFile
+            return "success"
+            """  # noqa: E501
+        else:
+            get_script = f"""
+            set tiffData to the clipboard as TIFF picture
+            set theFile to open for access POSIX file "{temp_path}" with write permission
+            write tiffData to theFile
+            close access theFile
+            return "success"
+            """  # noqa: E501
+
+        # S603: osascript_path validated via shutil.which(), script is internal
+        result = subprocess.run(  # noqa: S603
+            [osascript_path, "-e", get_script],
+            capture_output=True,
+            check=False,
+            timeout=3,
+            text=True,
+        )
+
+        if result.returncode != 0 or "success" not in result.stdout:
+            return None
+
+        # Check if file was created and has content
+        if not pathlib.Path(temp_path).exists() or pathlib.Path(temp_path).stat().st_size == 0:
+            return None
+
+        # Read and validate the image
+        image_data = pathlib.Path(temp_path).read_bytes()
+
+        try:
+            image = Image.open(io.BytesIO(image_data))
+            # Convert to PNG if it's not already (e.g., if we got TIFF)
+            buffer = io.BytesIO()
+            image.save(buffer, format="PNG")
+            buffer.seek(0)
+            base64_data = base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+            return ImageData(
+                base64_data=base64_data,
+                format="png",
+                placeholder="[image]",
+            )
+        except (
+            # UnidentifiedImageError: corrupted or non-image data
+            UnidentifiedImageError,
+            OSError,  # OSError: I/O errors during image processing
+        ) as e:
+            logger.debug("Failed to process clipboard image via osascript: %s", e, exc_info=True)
+            return None
+
+    except subprocess.TimeoutExpired:
+        logger.debug("osascript timed out while accessing clipboard")
+        return None
+    except OSError as e:
+        logger.debug("OSError accessing clipboard via osascript: %s", e)
+        return None
+    finally:
+        # Clean up temp file
+        try:
+            pathlib.Path(temp_path).unlink()
+        except OSError as e:
+            logger.debug("Failed to clean up temp file %s: %s", temp_path, e)
+
+
+def encode_to_base64(data: bytes) -> str:
+    """Encode raw bytes to a base64 string.
+
+    Args:
+        data: Raw bytes to encode.
+
+    Returns:
+        Base64-encoded string.
+    """
+    return base64.b64encode(data).decode("utf-8")
+
+
+def create_multimodal_content(
+    text: str, images: list[ImageData], videos: list[VideoData] | None = None
+) -> list[Any]:
+    """Create multimodal message content with text, images, and videos.
+
+    Args:
+        text: Text content of the message
+        images: List of ImageData objects
+        videos: Optional list of VideoData objects
+
+    Returns:
+        List of content blocks in LangChain message format.
+    """
+    content_blocks = []
+
+    # Add text block. Strip only the display-only placeholders bound to the media
+    # actually attached here (e.g. "[image 1]") so the canonical/model-facing text
+    # never contains fake user-authored placeholder text. When a span is known,
+    # text that merely resembles the schema is preserved exactly; without a span
+    # `strip_media_placeholders` falls back to removing one occurrence per item,
+    # which can catch a look-alike literal. The media itself is carried by the
+    # structured blocks below.
+    #
+    # `placeholders` and `spans` are passed as parallel-but-unzipped lists on
+    # purpose: `strip_media_placeholders` recovers each span's token from the
+    # text slice, and each tracked media item has a unique token, so it never
+    # needs index alignment between the two.
+    media = [*images, *(videos or [])]
+    placeholders = [item.placeholder for item in media]
+    spans = [item.placeholder_span for item in media if item.placeholder_span is not None]
+    clean_text = strip_media_placeholders(text, placeholders, placeholder_spans=spans)
+    if clean_text:
+        content_blocks.append({"type": "text", "text": clean_text})
+
+    # Add image blocks
+    content_blocks.extend(image.to_message_content() for image in images)
+
+    # Add video blocks
+    if videos:
+        content_blocks.extend(video.to_message_content() for video in videos)
+
+    return content_blocks