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

ai_pipeline_core/deployment/contract.py

@@ -0,0 +1,84 @@
+"""Unified pipeline run response contract.
+
+@public
+
+Single source of truth for the response shape used by both
+webhook push (ai-pipeline-core) and polling pull (unified-middleware).
+"""
+
+from datetime import datetime
+from typing import Annotated, Literal
+from uuid import UUID
+
+from pydantic import BaseModel, ConfigDict, Discriminator
+
+
+class _RunBase(BaseModel):
+    """Common fields on every run response variant."""
+
+    type: str
+    flow_run_id: UUID
+    project_name: str
+    state: str  # PENDING, RUNNING, COMPLETED, FAILED, CRASHED, CANCELLED
+    timestamp: datetime
+    storage_uri: str = ""
+
+    model_config = ConfigDict(frozen=True)
+
+
+class PendingRun(_RunBase):
+    """Pipeline queued or running but no progress reported yet."""
+
+    type: Literal["pending"] = "pending"  # pyright: ignore[reportIncompatibleVariableOverride]
+
+
+class ProgressRun(_RunBase):
+    """Pipeline running with step-level progress data."""
+
+    type: Literal["progress"] = "progress"  # pyright: ignore[reportIncompatibleVariableOverride]
+    step: int
+    total_steps: int
+    flow_name: str
+    status: str  # "started", "completed", "cached"
+    progress: float  # overall 0.0–1.0
+    step_progress: float  # within step 0.0–1.0
+    message: str
+
+
+class DeploymentResultData(BaseModel):
+    """Typed result payload — always has success + optional error."""
+
+    success: bool
+    error: str | None = None
+
+    model_config = ConfigDict(frozen=True, extra="allow")
+
+
+class CompletedRun(_RunBase):
+    """Pipeline finished (Prefect COMPLETED). Check result.success for business outcome."""
+
+    type: Literal["completed"] = "completed"  # pyright: ignore[reportIncompatibleVariableOverride]
+    result: DeploymentResultData
+
+
+class FailedRun(_RunBase):
+    """Pipeline crashed — execution error, not business logic."""
+
+    type: Literal["failed"] = "failed"  # pyright: ignore[reportIncompatibleVariableOverride]
+    error: str
+    result: DeploymentResultData | None = None
+
+
+RunResponse = Annotated[
+    PendingRun | ProgressRun | CompletedRun | FailedRun,
+    Discriminator("type"),
+]
+
+__all__ = [
+    "CompletedRun",
+    "DeploymentResultData",
+    "FailedRun",
+    "PendingRun",
+    "ProgressRun",
+    "RunResponse",
+]

ai_pipeline_core/deployment/helpers.py

@@ -0,0 +1,98 @@
+"""Helper functions for pipeline deployments."""
+
+import asyncio
+import re
+from typing import Any, Literal, TypedDict
+
+import httpx
+
+from ai_pipeline_core.deployment.contract import CompletedRun, FailedRun, ProgressRun
+from ai_pipeline_core.documents import Document, DocumentList, FlowDocument
+from ai_pipeline_core.logging import get_pipeline_logger
+
+logger = get_pipeline_logger(__name__)
+
+
+class StatusPayload(TypedDict):
+    """Webhook payload for Prefect state transitions (sub-flow level)."""
+
+    type: Literal["status"]
+    flow_run_id: str
+    project_name: str
+    step: int
+    total_steps: int
+    flow_name: str
+    state: str  # RUNNING, COMPLETED, FAILED, CRASHED, CANCELLED
+    state_name: str
+    timestamp: str
+
+
+def class_name_to_deployment_name(class_name: str) -> str:
+    """Convert PascalCase to kebab-case: ResearchPipeline → research-pipeline."""
+    name = re.sub(r"(?<!^)(?=[A-Z])", "-", class_name)
+    return name.lower()
+
+
+def extract_generic_params(cls: type) -> tuple[type | None, type | None]:
+    """Extract TOptions and TResult from PipelineDeployment generic args."""
+    from ai_pipeline_core.deployment.base import PipelineDeployment  # noqa: PLC0415
+
+    for base in getattr(cls, "__orig_bases__", []):
+        origin = getattr(base, "__origin__", None)
+        if origin is PipelineDeployment:
+            args = getattr(base, "__args__", ())
+            if len(args) == 2:
+                return args[0], args[1]
+
+    return None, None
+
+
+async def download_documents(
+    urls: list[str],
+    document_type: type[FlowDocument],
+) -> DocumentList:
+    """Download documents from URLs and return as DocumentList."""
+    documents: list[Document] = []
+    async with httpx.AsyncClient(timeout=60, follow_redirects=True) as client:
+        for url in urls:
+            response = await client.get(url)
+            response.raise_for_status()
+            filename = url.split("/")[-1].split("?")[0] or "document"
+            documents.append(document_type(name=filename, content=response.content))
+    return DocumentList(documents)
+
+
+async def upload_documents(documents: DocumentList, url_mapping: dict[str, str]) -> None:
+    """Upload documents to their mapped URLs."""
+    async with httpx.AsyncClient(timeout=60, follow_redirects=True) as client:
+        for doc in documents:
+            if doc.name in url_mapping:
+                response = await client.put(
+                    url_mapping[doc.name],
+                    content=doc.content,
+                    headers={"Content-Type": doc.mime_type},
+                )
+                response.raise_for_status()
+
+
+async def send_webhook(
+    url: str,
+    payload: ProgressRun | CompletedRun | FailedRun,
+    max_retries: int = 3,
+    retry_delay: float = 10.0,
+) -> None:
+    """Send webhook with retries."""
+    data: dict[str, Any] = payload.model_dump(mode="json")
+    for attempt in range(max_retries):
+        try:
+            async with httpx.AsyncClient(timeout=30) as client:
+                response = await client.post(url, json=data, follow_redirects=True)
+                response.raise_for_status()
+            return
+        except Exception as e:
+            if attempt < max_retries - 1:
+                logger.warning(f"Webhook retry {attempt + 1}/{max_retries}: {e}")
+                await asyncio.sleep(retry_delay)
+            else:
+                logger.error(f"Webhook failed after {max_retries} attempts: {e}")
+                raise

ai_pipeline_core/documents/flow_document.py

@@ -18,7 +18,7 @@ class FlowDocument(Document):
 
     FlowDocument is used for data that needs to be saved between pipeline
     steps and across multiple flow executions. These documents are typically
-    written to the file system using the simple_runner utilities.
+    written to the file system using the deployment utilities.
 
     Key characteristics:
     - Persisted to file system between pipeline steps

ai_pipeline_core/documents/task_document.py

@@ -40,7 +40,7 @@ class TaskDocument(Document):
 
     Note:
         - Cannot instantiate TaskDocument directly - must subclass
-        - Not saved by simple_runner utilities
+        - Not saved by deployment utilities
         - Reduces I/O overhead for temporary data
         - No additional abstract methods to implement
     """

ai_pipeline_core/documents/temporary_document.py

@@ -23,7 +23,7 @@ class TemporaryDocument(Document):
     - Can be instantiated directly (not abstract)
     - Cannot be subclassed (annotated with Python's @final decorator in code)
     - Useful for transient data like API responses or intermediate calculations
-    - Ignored by simple_runner save operations
+    - Ignored by deployment save operations
     - Useful for tests and debugging
 
     Creating TemporaryDocuments:

ai_pipeline_core/flow/config.py

@@ -39,11 +39,13 @@ class FlowConfig(ABC):
     Class Variables:
         INPUT_DOCUMENT_TYPES: List of FlowDocument types this flow accepts
         OUTPUT_DOCUMENT_TYPE: Single FlowDocument type this flow produces
+        WEIGHT: Weight for progress calculation (default 1.0, based on avg duration)
 
     Validation Rules:
         - INPUT_DOCUMENT_TYPES and OUTPUT_DOCUMENT_TYPE must be defined
         - OUTPUT_DOCUMENT_TYPE cannot be in INPUT_DOCUMENT_TYPES (prevents cycles)
         - Field names must be exact (common typos are detected)
+        - WEIGHT must be a positive number
 
     Why this matters:
         Flows connect in pipelines where one flow's output becomes another's input.
@@ -54,6 +56,7 @@ class FlowConfig(ABC):
         >>> class ProcessingFlowConfig(FlowConfig):
         ...     INPUT_DOCUMENT_TYPES = [RawDataDocument]
         ...     OUTPUT_DOCUMENT_TYPE = ProcessedDocument  # Different type!
+        ...     WEIGHT = 45.0  # Average ~45 minutes
         >>>
         >>> # Use in @pipeline_flow - RECOMMENDED PATTERN
         >>> @pipeline_flow(config=ProcessingFlowConfig, name="processing")
@@ -72,11 +75,12 @@ class FlowConfig(ABC):
     Note:
         - Validation happens at class definition time
         - Helps catch configuration errors early
-        - Used by simple_runner to manage document flow
+        - Used by PipelineDeployment to manage document flow
     """
 
     INPUT_DOCUMENT_TYPES: ClassVar[list[type[FlowDocument]]]
     OUTPUT_DOCUMENT_TYPE: ClassVar[type[FlowDocument]]
+    WEIGHT: ClassVar[float] = 1.0
 
     def __init_subclass__(cls, **kwargs: Any):
         """Validate flow configuration at subclass definition time.
@@ -106,7 +110,7 @@ class FlowConfig(ABC):
             return
 
         # Check for invalid field names (common mistakes)
-        allowed_fields = {"INPUT_DOCUMENT_TYPES", "OUTPUT_DOCUMENT_TYPE"}
+        allowed_fields = {"INPUT_DOCUMENT_TYPES", "OUTPUT_DOCUMENT_TYPE", "WEIGHT"}
         class_attrs = {name for name in dir(cls) if not name.startswith("_") and name.isupper()}
 
         # Find fields that look like they might be mistakes
@@ -145,6 +149,13 @@ class FlowConfig(ABC):
                 f"({cls.OUTPUT_DOCUMENT_TYPE.__name__}) cannot be in INPUT_DOCUMENT_TYPES"
             )
 
+        # Validate WEIGHT
+        weight = getattr(cls, "WEIGHT", 1.0)
+        if not isinstance(weight, (int, float)) or weight <= 0:
+            raise TypeError(
+                f"FlowConfig {cls.__name__}: WEIGHT must be a positive number, got {weight}"
+            )
+
     @classmethod
     def get_input_document_types(cls) -> list[type[FlowDocument]]:
         """Get the list of input document types this flow accepts.

ai_pipeline_core/flow/options.py

@@ -53,7 +53,7 @@ class FlowOptions(BaseSettings):
         - Frozen (immutable) after creation
         - Extra fields ignored (not strict)
         - Can be populated from environment variables
-        - Used by simple_runner.cli for command-line parsing
+        - Used by PipelineDeployment.run_cli for command-line parsing
 
     Note:
         The base class provides model selection. Subclasses should

ai_pipeline_core/llm/client.py

@@ -45,31 +45,30 @@ def _process_messages(
 
     Internal function that combines context and messages into a single
     list of API-compatible messages. Applies caching directives to
-    context messages for efficiency.
+    system prompt and context messages for efficiency.
 
     Args:
         context: Messages to be cached (typically expensive/static content).
         messages: Regular messages without caching (dynamic queries).
         system_prompt: Optional system instructions for the model.
-        cache_ttl: Cache TTL for context messages (e.g. "120s", "300s", "1h").
+        cache_ttl: Cache TTL for system and context messages (e.g. "120s", "300s", "1h").
                    Set to None or empty string to disable caching.
 
     Returns:
         List of formatted messages ready for API calls, with:
-        - System prompt at the beginning (if provided)
-        - Context messages with cache_control on the last one (if cache_ttl)
+        - System prompt at the beginning with cache_control (if provided and cache_ttl set)
+        - Context messages with cache_control on all messages (if cache_ttl set)
         - Regular messages without caching
 
     System Prompt Location:
         The system prompt parameter is always injected as the FIRST message
-        with role="system". It is NOT cached with context, allowing dynamic
-        system prompts without breaking cache efficiency.
+        with role="system". It is cached along with context when cache_ttl is set.
 
     Cache behavior:
-        The last context message gets ephemeral caching with specified TTL
+        All system and context messages get ephemeral caching with specified TTL
         to reduce token usage on repeated calls with same context.
         If cache_ttl is None or empty string (falsy), no caching is applied.
-        Only the last context message receives cache_control to maximize efficiency.
+        All system and context messages receive cache_control to maximize cache efficiency.
 
     Note:
         This is an internal function used by _generate_with_retry().
@@ -79,26 +78,28 @@ def _process_messages(
 
     # Add system prompt if provided
     if system_prompt:
-        processed_messages.append({"role": "system", "content": system_prompt})
+        processed_messages.append({
+            "role": "system",
+            "content": [{"type": "text", "text": system_prompt}],
+        })
 
     # Process context messages with caching if provided
     if context:
         # Use AIMessages.to_prompt() for context
         context_messages = context.to_prompt()
+        processed_messages.extend(context_messages)
 
-        # Apply caching to last context message and last content part if cache_ttl is set
-        if cache_ttl:
-            context_messages[-1]["cache_control"] = {  # type: ignore
-                "type": "ephemeral",
-                "ttl": cache_ttl,
-            }
-            assert isinstance(context_messages[-1]["content"], list)  # type: ignore
-            context_messages[-1]["content"][-1]["cache_control"] = {  # type: ignore
+    if cache_ttl:
+        for message in processed_messages:
+            message["cache_control"] = {  # type: ignore
                 "type": "ephemeral",
                 "ttl": cache_ttl,
             }
-
-        processed_messages.extend(context_messages)
+            if isinstance(message["content"], list):  # type: ignore
+                message["content"][-1]["cache_control"] = {  # type: ignore
+                    "type": "ephemeral",
+                    "ttl": cache_ttl,
+                }
 
     # Process regular messages without caching
     if messages:
@@ -153,10 +154,8 @@ def _model_name_to_openrouter_model(model: ModelName) -> str:
         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 == "sonar-pro-search":
-        return "perplexity/sonar-reasoning-pro"
+        return "perplexity/sonar-pro-search"
     if model.startswith("gemini"):
         return f"google/{model}"
     elif model.startswith("gpt"):
@@ -294,7 +293,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
                 ])

ai_pipeline_core/llm/model_response.py

@@ -88,10 +88,13 @@ class ModelResponse(ChatCompletion):
         data = chat_completion.model_dump()
 
         # fixes issue where the role is "assistantassistant" instead of "assistant"
+        valid_finish_reasons = {"stop", "length", "tool_calls", "content_filter", "function_call"}
         for i in range(len(data["choices"])):
-            if role := data["choices"][i]["message"].get("role"):
-                if role.startswith("assistant") and role != "assistant":
-                    data["choices"][i]["message"]["role"] = "assistant"
+            data["choices"][i]["message"]["role"] = "assistant"
+            # Only update finish_reason if it's not already a valid value
+            current_finish_reason = data["choices"][i].get("finish_reason")
+            if current_finish_reason not in valid_finish_reasons:
+                data["choices"][i]["finish_reason"] = "stop"
 
         super().__init__(**data)
 

ai_pipeline_core/llm/model_types.py

@@ -26,7 +26,6 @@ ModelName: TypeAlias = (
         "gemini-2.5-flash-search",
         "sonar-pro-search",
         "gpt-4o-search",
-        "grok-4-fast-search",
     ]
     | str
 )

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.

ai_pipeline_core/prompt_builder/global_cache.py

@@ -0,0 +1,78 @@
+"""Prompt cache coordination for concurrent LLM calls."""
+
+import asyncio
+import time
+from asyncio import Lock
+
+from ai_pipeline_core.documents import Document
+from ai_pipeline_core.llm import AIMessages, ModelName
+from ai_pipeline_core.llm.model_response import ModelResponse
+
+CACHED_PROMPTS: dict[str, Lock | int] = {}
+
+_cache_lock = Lock()
+CACHE_TTL = 600
+MIN_SIZE_FOR_CACHE = 32 * 1024
+
+
+class GlobalCacheLock:
+    """Serialize first prompt per cache key so subsequent calls get cache hits.
+
+    Waits for the first caller to complete before allowing others to execute,
+    ensuring the prompt cache is populated.
+    """
+
+    wait_time: float = 0
+    use_cache: bool = False
+
+    def _context_size(self, context: AIMessages) -> int:
+        length = 0
+        for msg in context:
+            if isinstance(msg, Document):
+                if msg.is_text:
+                    length += msg.size
+                else:
+                    length += 1024
+            elif isinstance(msg, str):
+                length += len(msg)
+            elif isinstance(msg, ModelResponse):  # type: ignore[arg-type]
+                length += len(msg.content)
+        return length
+
+    def __init__(self, model: ModelName, context: AIMessages, cache_lock: bool):  # noqa: D107
+        self.use_cache = cache_lock and self._context_size(context) > MIN_SIZE_FOR_CACHE
+        self.cache_key = f"{model}-{context.get_prompt_cache_key()}"
+        self.new_cache = False
+
+    async def __aenter__(self) -> "GlobalCacheLock":
+        wait_start = time.time()
+        if not self.use_cache:
+            return self
+
+        async with _cache_lock:
+            cache = CACHED_PROMPTS.get(self.cache_key)
+            if isinstance(cache, int):
+                if time.time() > cache + CACHE_TTL:
+                    cache = None
+                else:
+                    CACHED_PROMPTS[self.cache_key] = int(time.time())
+                    self.wait_time = time.time() - wait_start
+                    return self
+            if not cache:
+                self.new_cache = True
+                CACHED_PROMPTS[self.cache_key] = Lock()
+                await CACHED_PROMPTS[self.cache_key].acquire()  # type: ignore[union-attr]
+
+        if not self.new_cache and isinstance(cache, Lock):
+            async with cache:
+                pass  # waiting for lock to be released
+
+        self.wait_time = time.time() - wait_start
+        return self
+
+    async def __aexit__(self, exc_type: type | None, exc: BaseException | None, tb: object) -> None:
+        if self.new_cache:
+            await asyncio.sleep(1)  # give time for cache to be prepared
+            async with _cache_lock:
+                CACHED_PROMPTS[self.cache_key].release()  # type: ignore[union-attr]
+                CACHED_PROMPTS[self.cache_key] = int(time.time())

ai_pipeline_core/prompt_builder/new_core_documents_prompt.jinja2

@@ -0,0 +1,6 @@
+During this session the followiing **new core documents** were created:
+{% for document in new_core_documents %}
+- {{ document.id }} - {{ document.name }}
+{% endfor %}
+
+There won't be more documents provided during this session.