ingest.py

"""Photo ingestion: scan folder, caption new/changed images via Modal."""

import logging
import os
from pathlib import Path
from typing import Generator

from caption_store import get_entry, mark_error, upsert_entry

logger = logging.getLogger(__name__)

IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".tiff"}


def scan_images(folder: str) -> list[str]:
    folder_path = Path(folder)
    if not folder_path.exists() or not folder_path.is_dir():
        raise ValueError(f"Folder not found or not readable: {folder}")
    return [
        str(p)
        for p in folder_path.rglob("*")
        if p.suffix.lower() in IMAGE_EXTENSIONS and p.is_file()
    ]


def _needs_captioning(image_path: str) -> bool:
    entry = get_entry(image_path)
    if entry is None:
        return True
    return os.path.getmtime(image_path) != entry.get("mtime", -1)


def _get_captioner():
    """Return a handle to the deployed Modal Captioner."""
    import modal
    Captioner = modal.Cls.from_name("photographers-archive", "Captioner")
    return Captioner()


def ingest_folder(folder: str, collection: str = "General") -> Generator[tuple[int, int, str], None, None]:
    """
    Yields (processed_count, total_count, status_message) during ingestion.
    Raises ValueError for invalid folder paths.
    """
    images = scan_images(folder)
    total = len(images)

    if total == 0:
        yield (0, 0, "No supported images found in the selected folder.")
        return

    captioner = _get_captioner()
    processed = 0

    for image_path in images:
        if not _needs_captioning(image_path):
            processed += 1
            yield (processed, total, f"Skipped (cached): {os.path.basename(image_path)}")
            continue

        try:
            with open(image_path, "rb") as f:
                image_bytes = f.read()
            caption = captioner.caption.remote(image_bytes, os.path.basename(image_path))
            upsert_entry(image_path, caption, os.path.getmtime(image_path), collection=collection)
            processed += 1
            yield (processed, total, f"Captioned: {os.path.basename(image_path)}")
        except Exception as e:
            logger.error("Failed to caption %s: %s", image_path, e)
            mark_error(image_path, str(e), collection=collection)
            processed += 1
            yield (processed, total, f"Error: {os.path.basename(image_path)}")

    yield (total, total, f"Done. {total} images processed.")