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

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,12 +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 == "grok-4-fast-search":
-        return "x-ai/grok-4-fast: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"):
@@ -186,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.
 
@@ -295,7 +291,7 @@ async def _generate_with_retry(
                 model, span_type="LLM", input=processed_messages
             ) as span:
                 response = await _generate(model, processed_messages, completion_kwargs)
-                span.set_attributes(response.get_laminar_metadata())
+                span.set_attributes(response.get_laminar_metadata())  # pyright: ignore[reportArgumentType]
                 Laminar.set_span_output([
                     r for r in (response.reasoning_content, response.content) if r
                 ])
@@ -341,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.
@@ -369,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
@@ -397,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
@@ -405,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([
@@ -416,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
@@ -511,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..."
         ... )
         >>>
@@ -568,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: ..."
         ... )

ai_pipeline_core/llm/model_response.py

@@ -28,7 +28,7 @@ class ModelResponse(ChatCompletion):
 
     Primary usage is adding to AIMessages for multi-turn conversations:
 
-        >>> response = await llm.generate("gpt-5", messages=messages)
+        >>> response = await llm.generate("gpt-5.1", messages=messages)
         >>> messages.append(response)  # Add assistant response to conversation
         >>> print(response.content)  # Access generated text
 
@@ -43,7 +43,7 @@ class ModelResponse(ChatCompletion):
         >>> from ai_pipeline_core import llm, AIMessages
         >>>
         >>> messages = AIMessages(["Explain quantum computing"])
-        >>> response = await llm.generate("gpt-5", messages=messages)
+        >>> response = await llm.generate("gpt-5.1", messages=messages)
         >>>
         >>> # Primary usage: add to conversation
         >>> messages.append(response)
@@ -81,7 +81,7 @@ class ModelResponse(ChatCompletion):
             >>> # Usually created internally by generate()
             >>> response = ModelResponse(
             ...     chat_completion=completion,
-            ...     model_options={"temperature": 0.7, "model": "gpt-4"},
+            ...     model_options={"temperature": 0.7, "model": "gpt-5.1"},
             ...     metadata={"time_taken": 1.5, "first_token_time": 0.3}
             ... )
         """
@@ -116,7 +116,7 @@ class ModelResponse(ChatCompletion):
             Generated text from the model, or empty string if none.
 
         Example:
-            >>> response = await generate("gpt-5", messages="Hello")
+            >>> response = await generate("gpt-5.1", messages="Hello")
             >>> text = response.content  # The generated response
             >>>
             >>> # Common pattern: add to messages then use content
@@ -185,7 +185,7 @@ class ModelResponse(ChatCompletion):
 
         Example:
             >>> response = await llm.generate(
-            ...     "gpt-5",
+            ...     "gpt-5.1",
             ...     context=large_doc,
             ...     messages="Summarize this"
             ... )

ai_pipeline_core/llm/model_types.py

@@ -15,18 +15,15 @@ from typing import Literal, TypeAlias
 ModelName: TypeAlias = (
     Literal[
         # Core models
-        "gemini-2.5-pro",
-        "gpt-5",
-        "grok-4",
+        "gemini-3-pro",
+        "gpt-5.1",
         # Small models
-        "gemini-2.5-flash",
+        "gemini-3-flash",
         "gpt-5-mini",
-        "grok-4-fast",
+        "grok-4.1-fast",
         # Search models
-        "gemini-2.5-flash-search",
+        "gemini-3-flash-search",
         "sonar-pro-search",
-        "gpt-4o-search",
-        "grok-4-fast-search",
     ]
     | str
 )
@@ -39,15 +36,15 @@ string for custom models. The type is a union of predefined literals
 and str, giving you the best of both worlds: suggestions for known
 models and flexibility for custom ones.
 
-Note: These are example common model names as of Q3 2025. Actual availability
+Note: These are example common model names as of Q1 2026. Actual availability
 depends on your LiteLLM proxy configuration and provider access.
 
 Model categories:
-    Core models (gemini-2.5-pro, gpt-5, grok-4):
+    Core models (gemini-3-pro, gpt-5.1):
         High-capability models for complex tasks requiring deep reasoning,
         nuanced understanding, or creative generation.
 
-    Small models (gemini-2.5-flash, gpt-5-mini, grok-4-fast):
+    Small models (gemini-3-flash, gpt-5-mini, grok-4.1-fast):
         Efficient models optimized for speed and cost, suitable for
         simpler tasks or high-volume processing.
 
@@ -65,7 +62,7 @@ Example:
     >>> from ai_pipeline_core import llm, ModelName
     >>>
     >>> # Predefined model with IDE autocomplete
-    >>> model: ModelName = "gpt-5"  # IDE suggests common models
+    >>> model: ModelName = "gpt-5.1"  # IDE suggests common models
     >>> response = await llm.generate(model, messages="Hello")
     >>>
     >>> # Custom model works directly
@@ -73,7 +70,7 @@ Example:
     >>> response = await llm.generate(model, messages="Hello")
     >>>
     >>> # Both types work seamlessly
-    >>> models: list[ModelName] = ["gpt-5", "custom-llm", "gemini-2.5-pro"]
+    >>> models: list[ModelName] = ["gpt-5.1", "custom-llm", "gemini-3-pro"]
 
 Note:
     The ModelName type includes both predefined literals and str,

ai_pipeline_core/logging/logging_mixin.py

@@ -117,7 +117,7 @@ class StructuredLoggerMixin(LoggerMixin):
 
         Example:
             self.log_metric("processing_time", 1.23, "seconds",
-                          document_type="pdf", model="gpt-4")
+                          document_type="pdf", model="gpt-5.1")
         """
         self.logger.info(
             f"Metric: {metric_name}",
@@ -140,7 +140,7 @@ class StructuredLoggerMixin(LoggerMixin):
 
         Example:
             self.log_span("llm_generation", 1234.5,
-                         model="gpt-4", tokens=500)
+                         model="gpt-5.1", tokens=500)
         """
         self.logger.info(
             f"Span: {operation}",

ai_pipeline_core/pipeline.py

@@ -605,7 +605,7 @@ def pipeline_flow(
         - pipeline_task: For task-level decoration
         - FlowConfig: Type-safe flow configuration
         - FlowOptions: Base class for flow options
-        - simple_runner.run_pipeline: Execute flows locally
+        - PipelineDeployment: Execute flows locally or remotely
     """
     flow_decorator: Callable[..., Any] = _prefect_flow
 

ai_pipeline_core/progress.py

@@ -0,0 +1,127 @@
+"""@public Intra-flow progress tracking with order-preserving webhook delivery."""
+
+import asyncio
+from collections.abc import Generator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from uuid import UUID
+
+from ai_pipeline_core.deployment.contract import ProgressRun
+from ai_pipeline_core.logging import get_pipeline_logger
+
+logger = get_pipeline_logger(__name__)
+
+
+@dataclass(frozen=True, slots=True)
+class ProgressContext:
+    """Internal context holding state for progress calculation and webhook delivery."""
+
+    webhook_url: str
+    project_name: str
+    run_id: str
+    flow_run_id: str
+    flow_name: str
+    step: int
+    total_steps: int
+    weights: tuple[float, ...]
+    completed_weight: float
+    current_flow_weight: float
+    queue: asyncio.Queue[ProgressRun | None]
+
+
+_context: ContextVar[ProgressContext | None] = ContextVar("progress_context", default=None)
+
+
+async def update(fraction: float, message: str = "") -> None:
+    """@public Report intra-flow progress (0.0-1.0). No-op without context."""
+    ctx = _context.get()
+    if ctx is None or not ctx.webhook_url:
+        return
+
+    fraction = max(0.0, min(1.0, fraction))
+
+    total_weight = sum(ctx.weights)
+    if total_weight > 0:
+        overall = (ctx.completed_weight + ctx.current_flow_weight * fraction) / total_weight
+    else:
+        overall = fraction
+    overall = round(max(0.0, min(1.0, overall)), 4)
+
+    payload = ProgressRun(
+        flow_run_id=UUID(ctx.flow_run_id) if ctx.flow_run_id else UUID(int=0),
+        project_name=ctx.project_name,
+        state="RUNNING",
+        timestamp=datetime.now(timezone.utc),
+        step=ctx.step,
+        total_steps=ctx.total_steps,
+        flow_name=ctx.flow_name,
+        status="progress",
+        progress=overall,
+        step_progress=round(fraction, 4),
+        message=message,
+    )
+
+    ctx.queue.put_nowait(payload)
+
+
+async def webhook_worker(
+    queue: asyncio.Queue[ProgressRun | None],
+    webhook_url: str,
+    max_retries: int = 3,
+    retry_delay: float = 10.0,
+) -> None:
+    """Process webhooks sequentially with retries, preserving order."""
+    from ai_pipeline_core.deployment.helpers import send_webhook  # noqa: PLC0415
+
+    while True:
+        payload = await queue.get()
+        if payload is None:
+            queue.task_done()
+            break
+
+        try:
+            await send_webhook(webhook_url, payload, max_retries, retry_delay)
+        except Exception:
+            pass  # Already logged in send_webhook
+
+        queue.task_done()
+
+
+@contextmanager
+def flow_context(
+    webhook_url: str,
+    project_name: str,
+    run_id: str,
+    flow_run_id: str,
+    flow_name: str,
+    step: int,
+    total_steps: int,
+    weights: tuple[float, ...],
+    completed_weight: float,
+    queue: asyncio.Queue[ProgressRun | None],
+) -> Generator[None, None, None]:
+    """Set up progress context for a flow. Framework internal use."""
+    current_flow_weight = weights[step - 1] if step <= len(weights) else 1.0
+    ctx = ProgressContext(
+        webhook_url=webhook_url,
+        project_name=project_name,
+        run_id=run_id,
+        flow_run_id=flow_run_id,
+        flow_name=flow_name,
+        step=step,
+        total_steps=total_steps,
+        weights=weights,
+        completed_weight=completed_weight,
+        current_flow_weight=current_flow_weight,
+        queue=queue,
+    )
+    token = _context.set(ctx)
+    try:
+        yield
+    finally:
+        _context.reset(token)
+
+
+__all__ = ["update", "webhook_worker", "flow_context", "ProgressContext"]

ai_pipeline_core/prompt_builder/__init__.py

@@ -0,0 +1,5 @@
+"""@public Prompt builder for document-aware LLM interactions with caching."""
+
+from .prompt_builder import EnvironmentVariable, PromptBuilder
+
+__all__ = ["EnvironmentVariable", "PromptBuilder"]

ai_pipeline_core/prompt_builder/documents_prompt.jinja2

@@ -0,0 +1,23 @@
+You were provided with the following documents:
+- **core documents** - these are already a reviewed documents which are part of official project documentation.
+- **source documents** (called also **sources**) - these are not part of official project documentation and they will be deleted after your task is completed.
+
+{% if core_documents %}
+There are the following **core documents** available during this session:
+{% for document in core_documents %}
+- {{ document.id }} - {{ document.name }}
+{% endfor %}
+{% else %}
+There are no **core documents** available during this session.
+{% endif %}
+
+{% if new_documents %}
+There are the following **source documents** (called also **sources**) available during this session:
+{% for document in new_documents %}
+- {{ document.id }} - {{ document.name }}
+{% endfor %}
+{% else %}
+There are no **source documents** (called also **sources**) available during this session.
+{% endif %}
+
+There won't be more **core documents** and **source documents** provided during this conversation, however **new core documents** may be provided.