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

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

@@ -53,7 +53,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 +74,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 +264,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)}")
 

ai_pipeline_core/llm/client.py

@@ -150,10 +150,8 @@ def _model_name_to_openrouter_model(model: ModelName) -> str:
     Returns:
         OpenRouter model name.
     """
-    if model == "gpt-4o-search":
-        return "openai/gpt-4o-search-preview"
-    if model == "gemini-2.5-flash-search":
-        return "google/gemini-2.5-flash:online"
+    if model == "gemini-3-flash-search":
+        return "google/gemini-3-flash:online"
     if model == "sonar-pro-search":
         return "perplexity/sonar-pro-search"
     if model.startswith("gemini"):
@@ -184,7 +182,7 @@ async def _generate(
     Handles both regular and structured output generation.
 
     Args:
-        model: Model identifier (e.g., "gpt-5", "gemini-2.5-pro").
+        model: Model identifier (e.g., "gpt-5.1", "gemini-3-pro").
         messages: Formatted messages for the API.
         completion_kwargs: Additional parameters for the completion API.
 
@@ -339,7 +337,7 @@ async def generate(
         4. CONFIGURATION: Configure model behavior via LiteLLM proxy or environment variables
 
     Args:
-        model: Model to use (e.g., "gpt-5", "gemini-2.5-pro", "grok-4").
+        model: Model to use (e.g., "gpt-5.1", "gemini-3-pro", "grok-4.1-fast").
                Accepts predefined models or any string for custom models.
         context: Static context to cache (documents, examples, instructions).
                 Defaults to None (empty context). Cached for 5 minutes by default.
@@ -367,17 +365,17 @@ async def generate(
         Wrap Documents in AIMessages - DO NOT pass directly or convert to .text:
 
         # CORRECT - wrap Document in AIMessages
-        response = await llm.generate("gpt-5", messages=AIMessages([my_document]))
+        response = await llm.generate("gpt-5.1", messages=AIMessages([my_document]))
 
         # WRONG - don't pass Document directly
-        response = await llm.generate("gpt-5", messages=my_document)  # NO!
+        response = await llm.generate("gpt-5.1", messages=my_document)  # NO!
 
         # WRONG - don't convert to string yourself
-        response = await llm.generate("gpt-5", messages=my_document.text)  # NO!
+        response = await llm.generate("gpt-5.1", messages=my_document.text)  # NO!
 
     VISION/PDF MODEL COMPATIBILITY:
         When using Documents containing images or PDFs, ensure your model supports these formats:
-        - Images require vision-capable models (gpt-4o, gemini-pro-vision, claude-3-sonnet)
+        - Images require vision-capable models (gpt-5.1, gemini-3-flash, gemini-3-pro)
         - PDFs require document processing support (varies by provider)
         - Non-compatible models will raise ValueError or fall back to text extraction
         - Check model capabilities before including visual/PDF content
@@ -395,7 +393,7 @@ async def generate(
 
     Example:
         >>> # CORRECT - No options parameter (this is the recommended pattern)
-        >>> response = await llm.generate("gpt-5", messages="Explain quantum computing")
+        >>> response = await llm.generate("gpt-5.1", messages="Explain quantum computing")
         >>> print(response.content)  # In production, use get_pipeline_logger instead of print
 
         >>> # With context caching for efficiency
@@ -403,10 +401,10 @@ async def generate(
         >>> static_doc = AIMessages([large_document, "few-shot example: ..."])
         >>>
         >>> # First call: caches context
-        >>> r1 = await llm.generate("gpt-5", context=static_doc, messages="Summarize")
+        >>> r1 = await llm.generate("gpt-5.1", context=static_doc, messages="Summarize")
         >>>
         >>> # Second call: reuses cache, saves tokens!
-        >>> r2 = await llm.generate("gpt-5", context=static_doc, messages="Key points?")
+        >>> r2 = await llm.generate("gpt-5.1", context=static_doc, messages="Key points?")
 
         >>> # Multi-turn conversation
         >>> messages = AIMessages([
@@ -414,7 +412,7 @@ async def generate(
         ...     previous_response,
         ...     "Can you give an example?"
         ... ])
-        >>> response = await llm.generate("gpt-5", messages=messages)
+        >>> response = await llm.generate("gpt-5.1", messages=messages)
 
     Performance:
         - Context caching saves ~50-90% tokens on repeated calls
@@ -509,7 +507,7 @@ async def generate_structured(
 
         >>> # Step 1: Research/analysis with generate() - no options parameter
         >>> research = await llm.generate(
-        ...     "gpt-5",
+        ...     "gpt-5.1",
         ...     messages="Research and analyze this complex topic..."
         ... )
         >>>
@@ -566,7 +564,7 @@ async def generate_structured(
         >>>
         >>> # CORRECT - No options parameter
         >>> response = await llm.generate_structured(
-        ...     "gpt-5",
+        ...     "gpt-5.1",
         ...     response_format=Analysis,
         ...     messages="Analyze this product review: ..."
         ... )