ai-pipeline-core 0.3.0__py3-none-any.whl → 0.3.4__py3-none-any.whl

ai_pipeline_core/documents/mime_type.py

@@ -24,6 +24,8 @@ EXTENSION_MIME_MAP = {
     "gif": "image/gif",
     "bmp": "image/bmp",
     "webp": "image/webp",
+    "heic": "image/heic",
+    "heif": "image/heif",
     "json": "application/json",
     "yaml": "application/yaml",
     "yml": "application/yaml",
@@ -266,3 +268,29 @@ def is_image_mime_type(mime_type: str) -> bool:
         False
     """
     return mime_type.startswith("image/")
+
+
+LLM_SUPPORTED_IMAGE_MIME_TYPES: frozenset[str] = frozenset({
+    "image/png",
+    "image/jpeg",
+    "image/webp",
+    "image/heic",
+    "image/heif",
+})
+
+
+def is_llm_supported_image(mime_type: str) -> bool:
+    """Check if MIME type is an image format directly supported by LLMs.
+
+    Unsupported image formats (gif, bmp, tiff, svg, etc.) need conversion
+    to PNG before sending to the LLM.
+
+    @public
+
+    Args:
+        mime_type: MIME type string to check.
+
+    Returns:
+        True if the image format is natively supported by LLMs.
+    """
+    return mime_type in LLM_SUPPORTED_IMAGE_MIME_TYPES

ai_pipeline_core/flow/options.py

@@ -41,7 +41,7 @@ class FlowOptions(BaseSettings):
 
         >>> # Or create programmatically:
         >>> options = MyFlowOptions(
-        ...     core_model="gemini-2.5-pro",
+        ...     core_model="gemini-3-pro",
         ...     temperature=0.9
         ... )
 
@@ -61,11 +61,11 @@ class FlowOptions(BaseSettings):
     """
 
     core_model: ModelName = Field(
-        default="gemini-2.5-pro",
+        default="gemini-3-pro",
         description="Primary model for complex analysis and generation tasks.",
     )
     small_model: ModelName = Field(
-        default="grok-4-fast",
+        default="grok-4.1-fast",
         description="Fast, cost-effective model for simple tasks and orchestration.",
     )
 

ai_pipeline_core/images/__init__.py

@@ -0,0 +1,362 @@
+"""Image processing utilities for LLM vision models.
+
+@public
+
+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.
+
+Quick Start:
+    >>> from ai_pipeline_core.images import process_image, ImagePreset
+    >>>
+    >>> result = process_image(screenshot_bytes)
+    >>> for part in result:
+    ...     send_to_llm(part.data, context=part.label)
+    >>>
+    >>> result = process_image(screenshot_bytes, preset=ImagePreset.GEMINI)
+"""
+
+from enum import StrEnum
+
+from pydantic import BaseModel, Field
+
+from ai_pipeline_core.documents import Document, TemporaryDocument
+
+from ._processing import execute_split, load_and_normalize, plan_split
+
+__all__ = [
+    "ImagePreset",
+    "ImageProcessingConfig",
+    "ImagePart",
+    "ProcessedImage",
+    "ImageProcessingError",
+    "process_image",
+    "process_image_to_documents",
+]
+
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+
+class ImagePreset(StrEnum):
+    """Presets for LLM vision model constraints.
+
+    @public
+    """
+
+    GEMINI = "gemini"
+    CLAUDE = "claude"
+    GPT4V = "gpt4v"
+
+
+class ImageProcessingConfig(BaseModel):
+    """Configuration for image processing.
+
+    @public
+
+    Use ``for_preset`` for standard configurations or construct directly for
+    custom constraints.
+
+    Example:
+        >>> config = ImageProcessingConfig.for_preset(ImagePreset.GEMINI)
+        >>> config = ImageProcessingConfig(max_dimension=2000, jpeg_quality=80)
+    """
+
+    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.
+
+        @public
+        """
+        return _PRESETS[preset]
+
+
+_PRESETS: dict[ImagePreset, ImageProcessingConfig] = {
+    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):
+    """A single processed image part.
+
+    @public
+    """
+
+    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.
+
+        @public
+        """
+        if self.total == 1:
+            return "Full image"
+        return f"Part {self.index + 1}/{self.total}"
+
+
+class ProcessedImage(BaseModel):
+    """Result of image processing.
+
+    @public
+
+    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).
+
+        @public
+        """
+        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):
+    """Image processing failed.
+
+    @public
+    """
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def process_image(
+    image: bytes | Document,
+    preset: ImagePreset = ImagePreset.GEMINI,
+    config: ImageProcessingConfig | None = None,
+) -> ProcessedImage:
+    """Process an image for LLM vision models.
+
+    @public
+
+    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.
+
+    Example:
+        >>> result = process_image(screenshot_bytes)
+        >>> for part in result:
+        ...     print(part.label, len(part.data))
+    """
+    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)}")
+
+    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(
+    image: bytes | Document,
+    preset: ImagePreset = ImagePreset.GEMINI,
+    config: ImageProcessingConfig | None = None,
+    name_prefix: str = "image",
+    sources: list[str] | None = None,
+) -> list[TemporaryDocument]:
+    """Process an image and return parts as ``TemporaryDocument`` list.
+
+    @public
+
+    Convenience wrapper around ``process_image`` for direct integration
+    with ``AIMessages``.
+
+    Args:
+        image: Raw image bytes or a Document.
+        preset: Model preset (ignored when *config* is provided).
+        config: Custom configuration.
+        name_prefix: Prefix for generated document names.
+        sources: Optional provenance references attached to each document.
+
+    Returns:
+        List of ``TemporaryDocument`` instances with JPEG image data.
+
+    Example:
+        >>> docs = process_image_to_documents(screenshot_bytes)
+        >>> messages = AIMessages(docs)
+    """
+    result = process_image(image, preset=preset, config=config)
+
+    # Resolve sources
+    doc_sources: list[str] = list(sources or [])
+    if isinstance(image, Document):
+        doc_sources.append(image.sha256)
+
+    documents: list[TemporaryDocument] = []
+    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(
+            TemporaryDocument.create(
+                name=name,
+                content=part.data,
+                description=desc,
+                sources=doc_sources or None,
+            )
+        )
+
+    return documents

ai_pipeline_core/images/_processing.py

@@ -0,0 +1,157 @@
+"""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(
+    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}. "
+            f"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 "
+            f"(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/ai_messages.py

@@ -8,6 +8,7 @@ including text, documents, and model responses.
 
 import base64
 import hashlib
+import io
 import json
 from copy import deepcopy
 from typing import Any, Callable, Iterable, SupportsIndex, Union
@@ -17,9 +18,11 @@ from openai.types.chat import (
     ChatCompletionContentPartParam,
     ChatCompletionMessageParam,
 )
+from PIL import Image
 from prefect.logging import get_logger
 
 from ai_pipeline_core.documents import Document
+from ai_pipeline_core.documents.mime_type import is_llm_supported_image
 
 from .model_response import ModelResponse
 
@@ -53,7 +56,7 @@ class AIMessages(list[AIMessageType]):
     Note: Document conversion is automatic. Text content becomes user text messages.
 
     VISION/PDF MODEL COMPATIBILITY WARNING:
-    Images require vision-capable models (e.g., gpt-4o, gemini-pro-vision, claude-3-haiku).
+    Images require vision-capable models (e.g., gpt-5.1, gemini-3-flash, gemini-3-pro).
     Non-vision models will raise ValueError when encountering image documents.
     PDFs require models with document processing support - check your model's capabilities
     before including PDF documents in messages. Unsupported models may fall back to
@@ -74,7 +77,7 @@ class AIMessages(list[AIMessageType]):
         >>> from ai_pipeline_core import llm
         >>> messages = AIMessages()
         >>> messages.append("What is the capital of France?")
-        >>> response = await llm.generate("gpt-5", messages=messages)
+        >>> response = await llm.generate("gpt-5.1", messages=messages)
         >>> messages.append(response)  # Add the actual response
     """
 
@@ -264,10 +267,31 @@ class AIMessages(list[AIMessageType]):
             elif isinstance(message, Document):
                 messages.append({"role": "user", "content": AIMessages.document_to_prompt(message)})
             elif isinstance(message, ModelResponse):  # type: ignore
-                messages.append({
+                # Build base assistant message
+                assistant_message: ChatCompletionMessageParam = {
                     "role": "assistant",
                     "content": [{"type": "text", "text": message.content}],
-                })
+                }
+
+                # Preserve reasoning_content (Gemini Flash 3+, O1, O3, GPT-5)
+                if reasoning_content := message.reasoning_content:
+                    assistant_message["reasoning_content"] = reasoning_content  # type: ignore[typeddict-item]
+
+                # Preserve thinking_blocks (structured thinking)
+                if hasattr(message.choices[0].message, "thinking_blocks"):
+                    thinking_blocks = getattr(message.choices[0].message, "thinking_blocks", None)
+                    if thinking_blocks:
+                        assistant_message["thinking_blocks"] = thinking_blocks  # type: ignore[typeddict-item]
+
+                # Preserve provider_specific_fields (thought_signatures for Gemini multi-turn)
+                if hasattr(message.choices[0].message, "provider_specific_fields"):
+                    provider_fields = getattr(
+                        message.choices[0].message, "provider_specific_fields", None
+                    )
+                    if provider_fields:
+                        assistant_message["provider_specific_fields"] = provider_fields  # type: ignore[typeddict-item]
+
+                messages.append(assistant_message)
             else:
                 raise ValueError(f"Unsupported message type: {type(message)}")
 
@@ -376,9 +400,19 @@ class AIMessages(list[AIMessageType]):
             "text": f"{header_text}<content>\n",
         })
 
-        # Encode binary content
-        base64_content = base64.b64encode(document.content).decode("utf-8")
-        data_uri = f"data:{document.mime_type};base64,{base64_content}"
+        # Encode binary content, converting unsupported image formats to PNG
+        if document.is_image and not is_llm_supported_image(document.mime_type):
+            img = Image.open(io.BytesIO(document.content))
+            buf = io.BytesIO()
+            img.save(buf, format="PNG")
+            content_bytes = buf.getvalue()
+            mime_type = "image/png"
+        else:
+            content_bytes = document.content
+            mime_type = document.mime_type
+
+        base64_content = base64.b64encode(content_bytes).decode("utf-8")
+        data_uri = f"data:{mime_type};base64,{base64_content}"
 
         # Add appropriate content type
         if document.is_pdf: