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

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,15 +11,16 @@ 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",
-    "ModelResponse",
     "ModelOptions",
+    "ModelResponse",
     "StructuredModelResponse",
     "generate",
     "generate_structured",

ai_pipeline_core/llm/ai_messages.py

@@ -1,33 +1,46 @@
 """AI message handling for LLM interactions.
 
-@public
-
 Provides AIMessages container for managing conversations with mixed content types
 including text, documents, and model responses.
 """
 
 import base64
 import hashlib
+import io
 import json
+from collections.abc import Callable, Iterable
 from copy import deepcopy
-from typing import Any, Callable, Iterable, SupportsIndex, Union
+from typing import Any, SupportsIndex
 
-import tiktoken
 from openai.types.chat import (
     ChatCompletionContentPartParam,
     ChatCompletionMessageParam,
 )
-from prefect.logging import get_logger
+from PIL import Image
 
 from ai_pipeline_core.documents import Document
+from ai_pipeline_core.documents.document import get_tiktoken_encoding
+from ai_pipeline_core.documents.mime_type import is_llm_supported_image
+from ai_pipeline_core.logging import get_pipeline_logger
 
 from .model_response import ModelResponse
 
+logger = get_pipeline_logger(__name__)
+
+
+def _ensure_llm_compatible_image(content: bytes, mime_type: str) -> tuple[bytes, str]:
+    """Convert unsupported image formats to PNG for LLM consumption."""
+    if is_llm_supported_image(mime_type):
+        return content, mime_type
+    img = Image.open(io.BytesIO(content))
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return buf.getvalue(), "image/png"
+
+
 AIMessageType = str | Document | ModelResponse
 """Type for messages in AIMessages container.
 
-@public
-
 Represents the allowed types for conversation messages:
 - str: Plain text messages
 - Document: Structured document content
@@ -35,11 +48,9 @@ Represents the allowed types for conversation messages:
 """
 
 
-class AIMessages(list[AIMessageType]):
+class AIMessages(list[AIMessageType]):  # noqa: PLR0904
     """Container for AI conversation messages supporting mixed types.
 
-    @public
-
     This class extends list to manage conversation messages between user
     and AI, supporting text, Document objects, and ModelResponse instances.
     Messages are converted to OpenAI-compatible format for LLM interactions.
@@ -47,13 +58,14 @@ class AIMessages(list[AIMessageType]):
     Conversion Rules:
         - str: Becomes {"role": "user", "content": text}
         - Document: Becomes {"role": "user", "content": document_content}
-          (automatically handles text, images, PDFs based on MIME type)
+          (automatically handles text, images, PDFs based on MIME type; attachments
+          are rendered as <attachment> XML blocks)
         - ModelResponse: Becomes {"role": "assistant", "content": response.content}
 
     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
@@ -70,12 +82,6 @@ class AIMessages(list[AIMessageType]):
     constructor (`AIMessages("text")`) as this will raise a TypeError to prevent
     accidental character iteration.
 
-    Example:
-        >>> from ai_pipeline_core import llm
-        >>> messages = AIMessages()
-        >>> messages.append("What is the capital of France?")
-        >>> response = await llm.generate("gpt-5", messages=messages)
-        >>> messages.append(response)  # Add the actual response
     """
 
     def __init__(self, iterable: Iterable[AIMessageType] | None = None, *, frozen: bool = False):
@@ -144,8 +150,8 @@ class AIMessages(list[AIMessageType]):
 
     def __setitem__(
         self,
-        index: Union[SupportsIndex, slice],
-        value: Union[AIMessageType, Iterable[AIMessageType]],
+        index: SupportsIndex | slice,
+        value: AIMessageType | Iterable[AIMessageType],
     ) -> None:
         """Set item or slice."""
         self._check_frozen()
@@ -160,7 +166,7 @@ class AIMessages(list[AIMessageType]):
         self._check_frozen()
         return super().__iadd__(other)
 
-    def __delitem__(self, index: Union[SupportsIndex, slice]) -> None:
+    def __delitem__(self, index: SupportsIndex | slice) -> None:
         """Delete item or slice from list."""
         self._check_frozen()
         super().__delitem__(index)
@@ -189,9 +195,7 @@ class AIMessages(list[AIMessageType]):
         self._check_frozen()
         super().reverse()
 
-    def sort(
-        self, *, key: Callable[[AIMessageType], Any] | None = None, reverse: bool = False
-    ) -> None:
+    def sort(self, *, key: Callable[[AIMessageType], Any] | None = None, reverse: bool = False) -> None:
         """Sort list in place."""
         self._check_frozen()
         if key is None:
@@ -238,6 +242,8 @@ class AIMessages(list[AIMessageType]):
 
         Transforms the message list into the format expected by OpenAI API.
         Each message type is converted according to its role and content.
+        Documents are rendered as XML with any attachments included as nested
+        <attachment> blocks.
 
         Returns:
             List of ChatCompletionMessageParam dicts (from openai.types.chat)
@@ -247,26 +253,40 @@ class AIMessages(list[AIMessageType]):
         Raises:
             ValueError: If message type is not supported.
 
-        Example:
-            >>> messages = AIMessages(["Hello", response, "Follow up"])
-            >>> prompt = messages.to_prompt()
-            >>> # Result: [
-            >>> #   {"role": "user", "content": "Hello"},
-            >>> #   {"role": "assistant", "content": "..."},
-            >>> #   {"role": "user", "content": "Follow up"}
-            >>> # ]
         """
         messages: list[ChatCompletionMessageParam] = []
 
         for message in self:
             if isinstance(message, str):
-                messages.append({"role": "user", "content": message})
+                messages.append({"role": "user", "content": [{"type": "text", "text": message}]})
             elif isinstance(message, Document):
                 messages.append({"role": "user", "content": AIMessages.document_to_prompt(message)})
             elif isinstance(message, ModelResponse):  # type: ignore
-                messages.append({"role": "assistant", "content": message.content})
+                # 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)}")
+                raise TypeError(f"Unsupported message type: {type(message)}")
 
         return messages
 
@@ -306,8 +326,6 @@ class AIMessages(list[AIMessageType]):
     def approximate_tokens_count(self) -> int:
         """Approximate tokens count for the messages.
 
-        @public
-
         Uses tiktoken with gpt-4 encoding to estimate total token count
         across all messages in the conversation.
 
@@ -317,26 +335,27 @@ class AIMessages(list[AIMessageType]):
         Raises:
             ValueError: If message contains unsupported type.
 
-        Example:
-            >>> messages = AIMessages(["Hello", "World"])
-            >>> messages.approximate_tokens_count  # ~2-3 tokens
         """
         count = 0
+        enc = get_tiktoken_encoding()
         for message in self:
             if isinstance(message, str):
-                count += len(tiktoken.encoding_for_model("gpt-4").encode(message))
+                count += len(enc.encode(message))
             elif isinstance(message, Document):
                 count += message.approximate_tokens_count
             elif isinstance(message, ModelResponse):  # type: ignore
-                count += len(tiktoken.encoding_for_model("gpt-4").encode(message.content))
+                count += len(enc.encode(message.content))
             else:
-                raise ValueError(f"Unsupported message type: {type(message)}")
+                raise TypeError(f"Unsupported message type: {type(message)}")
         return count
 
     @staticmethod
-    def document_to_prompt(document: Document) -> list[ChatCompletionContentPartParam]:
+    def document_to_prompt(document: Document) -> list[ChatCompletionContentPartParam]:  # noqa: PLR0912, PLR0914
         """Convert a document to prompt format for LLM consumption.
 
+        Renders the document as XML with text/image/PDF content, followed by any
+        attachments as separate <attachment> XML blocks with name and description attributes.
+
         Args:
             document: The document to convert.
 
@@ -346,50 +365,80 @@ class AIMessages(list[AIMessageType]):
         prompt: list[ChatCompletionContentPartParam] = []
 
         # Build the text header
-        description = (
-            f"<description>{document.description}</description>\n" if document.description else ""
-        )
-        header_text = (
-            f"<document>\n<id>{document.id}</id>\n<name>{document.name}</name>\n{description}"
-        )
+        description = f"<description>{document.description}</description>\n" if document.description else ""
+        header_text = f"<document>\n<id>{document.id}</id>\n<name>{document.name}</name>\n{description}"
 
         # Handle text documents
         if document.is_text:
             text_content = document.content.decode("utf-8")
-            content_text = f"{header_text}<content>\n{text_content}\n</content>\n</document>\n"
+            content_text = f"{header_text}<content>\n{text_content}\n</content>\n"
             prompt.append({"type": "text", "text": content_text})
-            return prompt
 
-        # Handle non-text documents
-        if not document.is_image and not document.is_pdf:
-            get_logger(__name__).error(
-                f"Document is not a text, image or PDF: {document.name} - {document.mime_type}"
-            )
+        # Handle binary documents (image/PDF)
+        elif document.is_image or document.is_pdf:
+            prompt.append({"type": "text", "text": f"{header_text}<content>\n"})
+
+            if document.is_image:
+                content_bytes, mime_type = _ensure_llm_compatible_image(document.content, document.mime_type)
+            else:
+                content_bytes, mime_type = document.content, document.mime_type
+            base64_content = base64.b64encode(content_bytes).decode("utf-8")
+            data_uri = f"data:{mime_type};base64,{base64_content}"
+
+            if document.is_pdf:
+                prompt.append({
+                    "type": "file",
+                    "file": {"file_data": data_uri},
+                })
+            else:
+                prompt.append({
+                    "type": "image_url",
+                    "image_url": {"url": data_uri, "detail": "high"},
+                })
+
+            prompt.append({"type": "text", "text": "</content>\n"})
+
+        else:
+            logger.error(f"Document is not a text, image or PDF: {document.name} - {document.mime_type}")
             return []
 
-        # Add header for binary content
-        prompt.append({
-            "type": "text",
-            "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}"
-
-        # Add appropriate content type
-        if document.is_pdf:
-            prompt.append({
-                "type": "file",
-                "file": {"file_data": data_uri},
-            })
-        else:  # is_image
-            prompt.append({
-                "type": "image_url",
-                "image_url": {"url": data_uri, "detail": "high"},
-            })
-
-        # Close the document tag
-        prompt.append({"type": "text", "text": "</content>\n</document>\n"})
+        # Render attachments
+        for att in document.attachments:
+            desc_attr = f' description="{att.description}"' if att.description else ""
+            att_open = f'<attachment name="{att.name}"{desc_attr}>\n'
+
+            if att.is_text:
+                prompt.append({"type": "text", "text": f"{att_open}{att.text}\n</attachment>\n"})
+            elif att.is_image or att.is_pdf:
+                prompt.append({"type": "text", "text": att_open})
+
+                if att.is_image:
+                    att_bytes, att_mime = _ensure_llm_compatible_image(att.content, att.mime_type)
+                else:
+                    att_bytes, att_mime = att.content, att.mime_type
+                att_b64 = base64.b64encode(att_bytes).decode("utf-8")
+                att_uri = f"data:{att_mime};base64,{att_b64}"
+
+                if att.is_pdf:
+                    prompt.append({
+                        "type": "file",
+                        "file": {"file_data": att_uri},
+                    })
+                else:
+                    prompt.append({
+                        "type": "image_url",
+                        "image_url": {"url": att_uri, "detail": "high"},
+                    })
+
+                prompt.append({"type": "text", "text": "</attachment>\n"})
+            else:
+                logger.warning(f"Skipping unsupported attachment type: {att.name} - {att.mime_type}")
+
+        # Close document — merge into last text part to preserve JSON structure (and cache key)
+        last = prompt[-1]
+        if last["type"] == "text":
+            prompt[-1] = {"type": "text", "text": last["text"] + "</document>\n"}
+        else:
+            prompt.append({"type": "text", "text": "</document>\n"})
 
         return prompt