ai-pipeline-core 0.3.3__tar.gz → 0.3.4__tar.gz

{ai_pipeline_core-0.3.3 → ai_pipeline_core-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-pipeline-core
-Version: 0.3.3
+Version: 0.3.4
 Summary: Core utilities for AI-powered processing pipelines using prefect
 Project-URL: Homepage, https://github.com/bbarwik/ai-pipeline-core
 Project-URL: Repository, https://github.com/bbarwik/ai-pipeline-core

{ai_pipeline_core-0.3.3 → ai_pipeline_core-0.3.4}/ai_pipeline_core/__init__.py

@@ -82,6 +82,26 @@ Optional Environment Variables:
     - LMNR_DEBUG: Set to "true" to enable debug-level traces
 """
 
+import os
+import sys
+
+# Disable Prefect's built-in OpenTelemetry spans to prevent duplicates.
+# All tracing is handled by our @trace decorator and Laminar SDK.
+# Must be set before Prefect is imported by submodules below.
+os.environ.setdefault("PREFECT_CLOUD_ENABLE_ORCHESTRATION_TELEMETRY", "false")
+
+# If Prefect was already imported (user imported it before us), refresh its cached settings.
+if "prefect" in sys.modules:
+    try:
+        from prefect.settings import get_current_settings  # noqa: PLC0415
+
+        if get_current_settings().cloud.enable_orchestration_telemetry:
+            from prefect.context import refresh_global_settings_context  # noqa: PLC0415
+
+            refresh_global_settings_context()
+    except (ImportError, AttributeError):
+        pass
+
 from . import llm, progress
 from .deployment import DeploymentContext, DeploymentResult, PipelineDeployment
 from .documents import (
@@ -130,7 +150,7 @@ from .settings import Settings
 from .tracing import TraceInfo, TraceLevel, set_trace_cost, trace
 from .utils.remote_deployment import remote_deployment
 
-__version__ = "0.3.3"
+__version__ = "0.3.4"
 
 __all__ = [
     # Config/Settings

{ai_pipeline_core-0.3.3 → ai_pipeline_core-0.3.4}/ai_pipeline_core/debug/content.py

@@ -85,6 +85,7 @@ class ArtifactStore:
             "image/png": ".png",
             "image/jpeg": ".jpg",
             "image/gif": ".gif",
+            "image/webp": ".webp",
             "application/pdf": ".pdf",
         }
         ext = ext_map.get(mime_type, ".bin")

{ai_pipeline_core-0.3.3 → ai_pipeline_core-0.3.4}/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-0.3.3 → ai_pipeline_core-0.3.4}/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
 
@@ -397,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:

{ai_pipeline_core-0.3.3 → ai_pipeline_core-0.3.4}/ai_pipeline_core/llm/client.py

@@ -13,6 +13,7 @@ Key functions:
 
 import asyncio
 import time
+from io import BytesIO
 from typing import Any, TypeVar
 
 from lmnr import Laminar
@@ -21,19 +22,77 @@ from openai.lib.streaming.chat import ChunkEvent, ContentDeltaEvent, ContentDone
 from openai.types.chat import (
     ChatCompletionMessageParam,
 )
+from PIL import Image
 from prefect.logging import get_logger
 from pydantic import BaseModel, ValidationError
 
+from ai_pipeline_core.documents import Document
 from ai_pipeline_core.exceptions import LLMError
+from ai_pipeline_core.images import ImageProcessingConfig, process_image_to_documents
 from ai_pipeline_core.settings import settings
 
-from .ai_messages import AIMessages
+from .ai_messages import AIMessages, AIMessageType
 from .model_options import ModelOptions
 from .model_response import ModelResponse, StructuredModelResponse
 from .model_types import ModelName
 
 logger = get_logger()
 
+# Image splitting configs for automatic large-image handling at the LLM boundary.
+# Gemini supports up to 3000x3000; all other models use a conservative 1000x1000 default.
+_GEMINI_IMAGE_CONFIG = ImageProcessingConfig(
+    max_dimension=3000, max_pixels=9_000_000, jpeg_quality=75
+)
+_DEFAULT_IMAGE_CONFIG = ImageProcessingConfig(
+    max_dimension=1000, max_pixels=1_000_000, jpeg_quality=75
+)
+
+
+def _get_image_config(model: str) -> ImageProcessingConfig:
+    """Return the image splitting config for a model."""
+    if "gemini" in model.lower():
+        return _GEMINI_IMAGE_CONFIG
+    return _DEFAULT_IMAGE_CONFIG
+
+
+def _prepare_images_for_model(messages: AIMessages, model: str) -> AIMessages:
+    """Split image documents that exceed model constraints.
+
+    Returns a new AIMessages with oversized images replaced by tiles.
+    Returns the original instance unchanged if no splitting is needed.
+    """
+    if not any(isinstance(m, Document) and m.is_image for m in messages):
+        return messages
+
+    config = _get_image_config(model)
+    result: list[AIMessageType] = []
+    changed = False
+
+    for msg in messages:
+        if not (isinstance(msg, Document) and msg.is_image):
+            result.append(msg)
+            continue
+
+        try:
+            with Image.open(BytesIO(msg.content)) as img:
+                w, h = img.size
+        except Exception:
+            result.append(msg)
+            continue
+
+        if w <= config.max_dimension and h <= config.max_dimension and w * h <= config.max_pixels:
+            result.append(msg)
+            continue
+
+        name_prefix = msg.name.rsplit(".", 1)[0] if "." in msg.name else msg.name
+        tiles = process_image_to_documents(msg, config=config, name_prefix=name_prefix)
+        result.extend(tiles)
+        changed = True
+
+    if not changed:
+        return messages
+    return AIMessages(result)
+
 
 def _process_messages(
     context: AIMessages,
@@ -271,6 +330,10 @@ async def _generate_with_retry(
     if not context and not messages:
         raise ValueError("Either context or messages must be provided")
 
+    # Auto-split large images based on model-specific constraints
+    context = _prepare_images_for_model(context, model)
+    messages = _prepare_images_for_model(messages, model)
+
     if "gemini" in model.lower() and context.approximate_tokens_count < 10000:
         # Bug fix for minimum explicit context size for Gemini models
         options.cache_ttl = None

{ai_pipeline_core-0.3.3 → ai_pipeline_core-0.3.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ai-pipeline-core"
-version = "0.3.3"
+version = "0.3.4"
 description = "Core utilities for AI-powered processing pipelines using prefect"
 readme = "README.md"
 license = {text = "MIT"}
@@ -177,9 +177,11 @@ reportIncompatibleMethodOverride = "error"
 reportIncompatibleVariableOverride = "error"
 # Report missing parameter types as warnings
 reportMissingParameterType = "warning"
+# Allow tests to import private internals for unit testing
+reportPrivateUsage = "warning"
 
 [tool.bumpversion]
-current_version = "0.3.2"
+current_version = "0.3.4"
 commit = true
 tag = true
 tag_name = "v{new_version}"