ai-pipeline-core 0.2.6__py3-none-any.whl → 0.4.1__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",
@@ -65,19 +67,8 @@ def detect_mime_type(content: bytes, name: str) -> str:
         Only the first 1024 bytes are analyzed for content detection.
         Extension-based detection is O(1) lookup.
 
-    Note:
-        Extension-based detection is preferred for text formats as
-        content analysis can sometimes misidentify structured text.
-
-    Example:
-        >>> detect_mime_type(b'{"key": "value"}', "data.json")
-        'application/json'
-        >>> detect_mime_type(b'Hello World', "text.txt")
-        'text/plain'
-        >>> detect_mime_type(b'', "empty.txt")
-        'text/plain'
-        >>> detect_mime_type(b'\\x89PNG', "image.xyz")
-        'image/png'  # Magic detects PNG despite wrong extension
+    Extension-based detection is preferred for text formats as
+    content analysis can sometimes misidentify structured text.
     """
     # Check for empty content
     if len(content) == 0:
@@ -97,40 +88,13 @@ def detect_mime_type(content: bytes, name: str) -> str:
             return mime
     except (AttributeError, OSError, magic.MagicException) as e:
         logger.warning(f"MIME detection failed for {name}: {e}")
-    except Exception as e:
-        logger.error(f"Unexpected error in MIME detection for {name}: {e}")
+    except Exception:
+        logger.exception(f"Unexpected error in MIME detection for {name}")
 
     # Final fallback based on extension or default
     return EXTENSION_MIME_MAP.get(ext, "application/octet-stream")
 
 
-def mime_type_from_extension(name: str) -> str:
-    """Get MIME type based solely on file extension.
-
-    Simple extension-based MIME type detection without content analysis.
-    This is a legacy function maintained for backward compatibility.
-
-    Args:
-        name: Filename with extension.
-
-    Returns:
-        MIME type based on extension, or 'application/octet-stream'
-        if extension is unknown.
-
-    Note:
-        Prefer detect_mime_type() for more accurate detection.
-        This function only checks the file extension.
-
-    Example:
-        >>> mime_type_from_extension("document.pdf")
-        'application/pdf'
-        >>> mime_type_from_extension("unknown.xyz")
-        'application/octet-stream'
-    """
-    ext = name.lower().split(".")[-1] if "." in name else ""
-    return EXTENSION_MIME_MAP.get(ext, "application/octet-stream")
-
-
 def is_text_mime_type(mime_type: str) -> bool:
     """Check if MIME type represents text-based content.
 
@@ -151,13 +115,6 @@ def is_text_mime_type(mime_type: str) -> bool:
         - application/yaml
         - application/x-yaml
 
-    Example:
-        >>> is_text_mime_type('text/plain')
-        True
-        >>> is_text_mime_type('application/json')
-        True
-        >>> is_text_mime_type('image/png')
-        False
     """
     text_types = [
         "text/",
@@ -179,15 +136,8 @@ def is_json_mime_type(mime_type: str) -> bool:
     Returns:
         True if MIME type is 'application/json', False otherwise.
 
-    Note:
-        Only matches exact 'application/json', not variants like
-        'application/ld+json' or 'application/vnd.api+json'.
-
-    Example:
-        >>> is_json_mime_type('application/json')
-        True
-        >>> is_json_mime_type('text/json')  # Not standard JSON MIME
-        False
+    Only matches exact 'application/json', not variants like
+    'application/ld+json' or 'application/vnd.api+json'.
     """
     return mime_type == "application/json"
 
@@ -207,13 +157,8 @@ def is_yaml_mime_type(mime_type: str) -> bool:
         - application/yaml (standard)
         - application/x-yaml (legacy)
 
-    Example:
-        >>> is_yaml_mime_type('application/yaml')
-        True
-        >>> is_yaml_mime_type('application/x-yaml')
-        True
     """
-    return mime_type == "application/yaml" or mime_type == "application/x-yaml"
+    return mime_type in {"application/yaml", "application/x-yaml"}
 
 
 def is_pdf_mime_type(mime_type: str) -> bool:
@@ -225,15 +170,8 @@ def is_pdf_mime_type(mime_type: str) -> bool:
     Returns:
         True if MIME type is 'application/pdf', False otherwise.
 
-    Note:
-        PDF documents require special handling in the LLM module
-        and are supported by certain vision-capable models.
-
-    Example:
-        >>> is_pdf_mime_type('application/pdf')
-        True
-        >>> is_pdf_mime_type('text/plain')
-        False
+    PDF documents require special handling in the LLM module
+    and are supported by certain vision-capable models.
     """
     return mime_type == "application/pdf"
 
@@ -255,14 +193,31 @@ def is_image_mime_type(mime_type: str) -> bool:
         - image/webp
         - image/svg+xml
 
-    Note:
-        Image documents are automatically encoded for vision-capable
-        LLM models in the AIMessages.document_to_prompt() method.
-
-    Example:
-        >>> is_image_mime_type('image/png')
-        True
-        >>> is_image_mime_type('application/pdf')
-        False
+    Image documents are automatically encoded for vision-capable
+    LLM models in the AIMessages.document_to_prompt() method.
     """
     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.
+
+    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/documents/utils.py

@@ -5,15 +5,14 @@ canonical key generation, and hash validation used throughout the document syste
 """
 
 import re
-from typing import Any, Iterable, Type
+from collections.abc import Iterable
+from typing import Any
 from urllib.parse import urlparse
 
 
 def sanitize_url(url: str) -> str:
     """Sanitize URL or query string for use in filenames.
 
-    @public
-
     Removes or replaces characters that are invalid in filenames.
 
     Args:
@@ -63,15 +62,13 @@ def camel_to_snake(name: str) -> str:
 
 
 def canonical_name_key(
-    obj_or_name: Type[Any] | str,
+    obj_or_name: type[Any] | str,
     *,
     max_parent_suffixes: int = 3,
     extra_suffixes: Iterable[str] = (),
 ) -> str:
     """Produce a canonical snake_case key from a class or name.
 
-    @public
-
     Process:
       1) Starting with the class name (or given string),
       2) Stripping any trailing parent class names (up to `max_parent_suffixes` from the MRO),
@@ -120,8 +117,6 @@ def canonical_name_key(
 def is_document_sha256(value: str) -> bool:
     """Check if a string is a valid base32-encoded SHA256 hash with proper entropy.
 
-    @public
-
     This function validates that a string is not just formatted like a SHA256 hash,
     but actually has the entropy characteristics of a real hash. It checks:
     1. Correct length (52 characters without padding)
@@ -174,7 +169,4 @@ def is_document_sha256(value: str) -> bool:
     # Require at least 8 unique characters (out of 32 possible in base32)
     # This prevents patterns like "AAAAAAA..." from being identified as real hashes
     unique_chars = len(set(value))
-    if unique_chars < 8:
-        return False
-
-    return True
+    return unique_chars >= 8

ai_pipeline_core/exceptions.py

@@ -1,97 +1,45 @@
 """Exception hierarchy for AI Pipeline Core.
 
-@public
-
 This module defines the exception hierarchy used throughout the AI Pipeline Core library.
 All exceptions inherit from PipelineCoreError, providing a consistent error handling interface.
 """
 
 
 class PipelineCoreError(Exception):
-    """Base exception for all AI Pipeline Core errors.
-
-    @public
-    """
-
-    pass
+    """Base exception for all AI Pipeline Core errors."""
 
 
 class DocumentError(PipelineCoreError):
-    """Base exception for document-related errors.
-
-    @public
-    """
-
-    pass
+    """Base exception for document-related errors."""
 
 
 class DocumentValidationError(DocumentError):
-    """Raised when document validation fails.
-
-    @public
-    """
-
-    pass
+    """Raised when document validation fails."""
 
 
 class DocumentSizeError(DocumentValidationError):
-    """Raised when document content exceeds MAX_CONTENT_SIZE limit.
-
-    @public
-    """
-
-    pass
+    """Raised when document content exceeds MAX_CONTENT_SIZE limit."""
 
 
 class DocumentNameError(DocumentValidationError):
-    """Raised when document name contains invalid characters or patterns.
-
-    @public
-    """
-
-    pass
+    """Raised when document name contains invalid characters or patterns."""
 
 
 class LLMError(PipelineCoreError):
-    """Raised when LLM generation fails after all retries.
-
-    @public
-    """
-
-    pass
+    """Raised when LLM generation fails after all retries."""
 
 
 class PromptError(PipelineCoreError):
-    """Base exception for prompt template errors.
-
-    @public
-    """
-
-    pass
+    """Base exception for prompt template errors."""
 
 
 class PromptRenderError(PromptError):
-    """Raised when Jinja2 template rendering fails.
-
-    @public
-    """
-
-    pass
+    """Raised when Jinja2 template rendering fails."""
 
 
 class PromptNotFoundError(PromptError):
-    """Raised when prompt template file is not found in search paths.
-
-    @public
-    """
-
-    pass
+    """Raised when prompt template file is not found in search paths."""
 
 
 class MimeTypeError(DocumentError):
-    """Raised when MIME type detection or validation fails.
-
-    @public
-    """
-
-    pass
+    """Raised when MIME type detection or validation fails."""

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