ai-pipeline-core 0.1.12__py3-none-any.whl → 0.4.1__py3-none-any.whl

ai_pipeline_core/images/__init__.py

@@ -0,0 +1,309 @@
+"""Image processing utilities for LLM vision models.
+
+Splits large images, compresses to JPEG, and respects model-specific constraints.
+Designed for website screenshots, document pages, and other visual content
+sent to vision-capable LLMs.
+"""
+
+from enum import StrEnum
+
+from pydantic import BaseModel, Field
+
+from ai_pipeline_core.documents import Document
+
+from ._processing import execute_split, load_and_normalize, plan_split
+
+__all__ = [
+    "ImageDocument",
+    "ImagePart",
+    "ImagePreset",
+    "ImageProcessingConfig",
+    "ImageProcessingError",
+    "ProcessedImage",
+    "process_image",
+    "process_image_to_documents",
+]
+
+
+class ImageDocument(Document):  # noqa: RUF067
+    """Concrete document for processed image parts."""
+
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+
+class ImagePreset(StrEnum):  # noqa: RUF067
+    """Presets for LLM vision model constraints."""
+
+    GEMINI = "gemini"
+    CLAUDE = "claude"
+    GPT4V = "gpt4v"
+
+
+class ImageProcessingConfig(BaseModel):  # noqa: RUF067
+    """Configuration for image processing.
+
+    Use ``for_preset`` for standard configurations or construct directly for
+    custom constraints.
+
+    """
+
+    model_config = {"frozen": True}
+
+    max_dimension: int = Field(
+        default=3000,
+        ge=100,
+        le=8192,
+        description="Maximum width AND height in pixels",
+    )
+    max_pixels: int = Field(
+        default=9_000_000,
+        ge=10_000,
+        description="Maximum total pixels per output image part",
+    )
+    overlap_fraction: float = Field(
+        default=0.20,
+        ge=0.0,
+        le=0.5,
+        description="Overlap between adjacent vertical parts (0.0-0.5)",
+    )
+    max_parts: int = Field(
+        default=20,
+        ge=1,
+        le=100,
+        description="Maximum number of output image parts",
+    )
+    jpeg_quality: int = Field(
+        default=60,
+        ge=10,
+        le=95,
+        description="JPEG compression quality (10-95)",
+    )
+
+    @classmethod
+    def for_preset(cls, preset: ImagePreset) -> "ImageProcessingConfig":
+        """Create configuration from a model preset."""
+        return _PRESETS[preset]
+
+
+_PRESETS: dict[ImagePreset, ImageProcessingConfig] = {  # noqa: RUF067
+    ImagePreset.GEMINI: ImageProcessingConfig(
+        max_dimension=3000,
+        max_pixels=9_000_000,
+        jpeg_quality=75,
+    ),
+    ImagePreset.CLAUDE: ImageProcessingConfig(
+        max_dimension=1568,
+        max_pixels=1_150_000,
+        jpeg_quality=60,
+    ),
+    ImagePreset.GPT4V: ImageProcessingConfig(
+        max_dimension=2048,
+        max_pixels=4_000_000,
+        jpeg_quality=70,
+    ),
+}
+
+
+# ---------------------------------------------------------------------------
+# Result models
+# ---------------------------------------------------------------------------
+
+
+class ImagePart(BaseModel):  # noqa: RUF067
+    """A single processed image part."""
+
+    model_config = {"frozen": True}
+
+    data: bytes = Field(repr=False)
+    width: int
+    height: int
+    index: int = Field(ge=0, description="0-indexed position")
+    total: int = Field(ge=1, description="Total number of parts")
+    source_y: int = Field(ge=0, description="Y offset in original image")
+    source_height: int = Field(ge=1, description="Height of region in original")
+
+    @property
+    def label(self) -> str:
+        """Human-readable label for LLM context, 1-indexed."""
+        if self.total == 1:
+            return "Full image"
+        return f"Part {self.index + 1}/{self.total}"
+
+
+class ProcessedImage(BaseModel):  # noqa: RUF067
+    """Result of image processing.
+
+    Iterable: ``for part in result`` iterates over parts.
+    """
+
+    model_config = {"frozen": True}
+
+    parts: list[ImagePart]
+    original_width: int
+    original_height: int
+    original_bytes: int
+    output_bytes: int
+    was_trimmed: bool = Field(description="True if width was trimmed to fit")
+    warnings: list[str] = Field(default_factory=list)
+
+    @property
+    def compression_ratio(self) -> float:
+        """Output size / input size (lower means more compression)."""
+        if self.original_bytes <= 0:
+            return 1.0
+        return self.output_bytes / self.original_bytes
+
+    def __len__(self) -> int:
+        return len(self.parts)
+
+    def __iter__(self):  # type: ignore[override]
+        return iter(self.parts)
+
+    def __getitem__(self, idx: int) -> ImagePart:
+        return self.parts[idx]
+
+
+# ---------------------------------------------------------------------------
+# Exceptions
+# ---------------------------------------------------------------------------
+
+
+class ImageProcessingError(Exception):  # noqa: RUF067
+    """Image processing failed."""
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def process_image(  # noqa: RUF067
+    image: bytes | Document,
+    preset: ImagePreset = ImagePreset.GEMINI,
+    config: ImageProcessingConfig | None = None,
+) -> ProcessedImage:
+    """Process an image for LLM vision models.
+
+    Splits tall images vertically with overlap, trims width if needed, and
+    compresses to JPEG.  The default preset is **GEMINI** (3 000 px, 9 M pixels).
+
+    Args:
+        image: Raw image bytes or a Document whose content is an image.
+        preset: Model preset (ignored when *config* is provided).
+        config: Custom configuration that overrides the preset.
+
+    Returns:
+        A ``ProcessedImage`` containing one or more ``ImagePart`` objects.
+
+    Raises:
+        ImageProcessingError: If the image cannot be decoded or processed.
+
+    """
+    effective = config if config is not None else ImageProcessingConfig.for_preset(preset)
+
+    # Resolve input bytes
+    raw: bytes
+    if isinstance(image, Document):
+        raw = image.content
+    elif isinstance(image, bytes):  # type: ignore[reportUnnecessaryIsInstance]
+        raw = image
+    else:
+        raise ImageProcessingError(f"Unsupported image input type: {type(image)}")  # pyright: ignore[reportUnreachable]
+
+    if not raw:
+        raise ImageProcessingError("Empty image data")
+
+    original_bytes = len(raw)
+
+    # Load & normalise
+    try:
+        img = load_and_normalize(raw)
+    except Exception as exc:
+        raise ImageProcessingError(f"Failed to decode image: {exc}") from exc
+
+    original_width, original_height = img.size
+
+    # Plan
+    plan = plan_split(
+        width=original_width,
+        height=original_height,
+        max_dimension=effective.max_dimension,
+        max_pixels=effective.max_pixels,
+        overlap_fraction=effective.overlap_fraction,
+        max_parts=effective.max_parts,
+    )
+
+    # Execute
+    raw_parts = execute_split(img, plan, effective.jpeg_quality)
+
+    # Build result
+    parts: list[ImagePart] = []
+    total = len(raw_parts)
+    total_output = 0
+
+    for idx, (data, w, h, sy, sh) in enumerate(raw_parts):
+        total_output += len(data)
+        parts.append(
+            ImagePart(
+                data=data,
+                width=w,
+                height=h,
+                index=idx,
+                total=total,
+                source_y=sy,
+                source_height=sh,
+            )
+        )
+
+    return ProcessedImage(
+        parts=parts,
+        original_width=original_width,
+        original_height=original_height,
+        original_bytes=original_bytes,
+        output_bytes=total_output,
+        was_trimmed=plan.trim_width is not None,
+        warnings=plan.warnings,
+    )
+
+
+def process_image_to_documents(  # noqa: RUF067
+    image: bytes | Document,
+    preset: ImagePreset = ImagePreset.GEMINI,
+    config: ImageProcessingConfig | None = None,
+    name_prefix: str = "image",
+    sources: tuple[str, ...] | None = None,
+) -> list[ImageDocument]:
+    """Process an image and return parts as ImageDocument list.
+
+    Convenience wrapper around ``process_image`` for direct integration
+    with ``AIMessages``.
+    """
+    result = process_image(image, preset=preset, config=config)
+
+    source_list: list[str] = list(sources or ())
+    if isinstance(image, Document):
+        source_list.append(image.sha256)
+    doc_sources = tuple(source_list) if source_list else None
+
+    documents: list[ImageDocument] = []
+    for part in result.parts:
+        if len(result.parts) == 1:
+            name = f"{name_prefix}.jpg"
+            desc = None
+        else:
+            name = f"{name_prefix}_{part.index + 1:02d}_of_{part.total:02d}.jpg"
+            desc = part.label
+
+        documents.append(
+            ImageDocument.create(
+                name=name,
+                content=part.data,
+                description=desc,
+                sources=doc_sources,
+            )
+        )
+
+    return documents

ai_pipeline_core/images/_processing.py

@@ -0,0 +1,151 @@
+"""Internal image processing logic: planning, splitting, encoding."""
+
+from dataclasses import dataclass
+from io import BytesIO
+from math import ceil
+
+from PIL import Image, ImageOps
+
+PIL_MAX_PIXELS = 100_000_000  # 100MP security limit
+
+
+@dataclass(frozen=True)
+class SplitPlan:
+    """Describes how to split an image into parts."""
+
+    tile_width: int
+    tile_height: int
+    step_y: int
+    num_parts: int
+    trim_width: int | None  # None = no trim needed
+    warnings: list[str]
+
+
+def plan_split(  # noqa: PLR0917
+    width: int,
+    height: int,
+    max_dimension: int,
+    max_pixels: int,
+    overlap_fraction: float,
+    max_parts: int,
+) -> SplitPlan:
+    """Calculate how to split an image. Pure function, no side effects.
+
+    Returns a SplitPlan describing tile size, step, and number of parts.
+    """
+    warnings: list[str] = []
+
+    # Effective tile size respecting both max_dimension and max_pixels
+    tile_size = max_dimension
+    while tile_size * tile_size > max_pixels and tile_size > 100:
+        tile_size -= 10
+
+    # Width: trim if needed (left-aligned, web content is left-aligned)
+    trim_width = tile_size if width > tile_size else None
+
+    effective_width = min(width, tile_size)
+
+    # If single-tile pixel budget is still exceeded by width * tile_height, reduce tile_height
+    tile_h = tile_size
+    while effective_width * tile_h > max_pixels and tile_h > 100:
+        tile_h -= 10
+
+    # No vertical split needed
+    if height <= tile_h:
+        return SplitPlan(
+            tile_width=effective_width,
+            tile_height=height,
+            step_y=0,
+            num_parts=1,
+            trim_width=trim_width,
+            warnings=warnings,
+        )
+
+    # Vertical split with overlap
+    overlap_px = int(tile_h * overlap_fraction)
+    step = tile_h - overlap_px
+    if step <= 0:
+        step = 1
+
+    num_parts = 1 + ceil((height - tile_h) / step)
+
+    # Auto-reduce if exceeds max_parts
+    if num_parts > max_parts:
+        warnings.append(f"Image requires {num_parts} parts but max is {max_parts}. Reducing to {max_parts} parts with larger step.")
+        num_parts = max_parts
+        if num_parts > 1:
+            step = (height - tile_h) // (num_parts - 1)
+        else:
+            step = 0
+
+    return SplitPlan(
+        tile_width=effective_width,
+        tile_height=tile_h,
+        step_y=step,
+        num_parts=num_parts,
+        trim_width=trim_width,
+        warnings=warnings,
+    )
+
+
+def load_and_normalize(data: bytes) -> Image.Image:
+    """Load image from bytes, apply EXIF orientation, validate size."""
+    img = Image.open(BytesIO(data))
+    img.load()
+
+    if img.width * img.height > PIL_MAX_PIXELS:
+        raise ValueError(f"Image too large: {img.width}x{img.height} = {img.width * img.height:,} pixels (limit: {PIL_MAX_PIXELS:,})")
+
+    # Fix EXIF orientation (important for mobile photos)
+    img = ImageOps.exif_transpose(img)
+    return img
+
+
+def encode_jpeg(img: Image.Image, quality: int) -> bytes:
+    """Encode PIL Image as JPEG bytes."""
+    # Convert to RGB if needed (JPEG doesn't support alpha)
+    if img.mode not in {"RGB", "L"}:
+        img = img.convert("RGB")
+
+    buf = BytesIO()
+    img.save(buf, format="JPEG", quality=quality, optimize=True)
+    return buf.getvalue()
+
+
+def execute_split(
+    img: Image.Image,
+    plan: SplitPlan,
+    jpeg_quality: int,
+) -> list[tuple[bytes, int, int, int, int]]:
+    """Execute a split plan on an image.
+
+    Returns list of (data, width, height, source_y, source_height) tuples.
+    """
+    width, height = img.size
+
+    # Trim width if needed (left-aligned crop)
+    if plan.trim_width is not None and width > plan.trim_width:
+        img = img.crop((0, 0, plan.trim_width, height))
+        width = plan.trim_width
+
+    # Convert to RGB once for JPEG
+    if img.mode not in {"RGB", "L"}:
+        img = img.convert("RGB")
+
+    parts: list[tuple[bytes, int, int, int, int]] = []
+
+    for i in range(plan.num_parts):
+        if plan.num_parts == 1:
+            y = 0
+        else:
+            y = i * plan.step_y
+            # Clamp so last tile aligns to bottom
+            y = min(y, max(0, height - plan.tile_height))
+
+        h = min(plan.tile_height, height - y)
+        tile = img.crop((0, y, width, y + h))
+
+        data = encode_jpeg(tile, jpeg_quality)
+        parts.append((data, width, h, y, h))
+
+    return parts

ai_pipeline_core/llm/__init__.py

@@ -1,7 +1,8 @@
 """Large Language Model integration via LiteLLM proxy.
 
 This package provides OpenAI API-compatible LLM interactions with built-in retry logic,
-LMNR tracing, and structured output generation using Pydantic models.
+LMNR tracing, and structured output generation using Pydantic models. Supports per-call
+observability via purpose and expected_cost parameters for span naming and cost tracking.
 """
 
 from .ai_messages import AIMessages, AIMessageType
@@ -10,12 +11,13 @@ from .client import (
     generate_structured,
 )
 from .model_options import ModelOptions
-from .model_response import ModelResponse, StructuredModelResponse
+from .model_response import Citation, ModelResponse, StructuredModelResponse
 from .model_types import ModelName
 
 __all__ = [
-    "AIMessages",
     "AIMessageType",
+    "AIMessages",
+    "Citation",
     "ModelName",
     "ModelOptions",
     "ModelResponse",