ai-pipeline-core 0.4.2__tar.gz → 0.4.4__tar.gz

{ai_pipeline_core-0.4.2 → ai_pipeline_core-0.4.4}/.gitignore

@@ -160,7 +160,6 @@ ehthumbs.db
 .prefect/
 
 # Test artifacts
-tests/test_data/
 test_output/
 
 # temporary specification

{ai_pipeline_core-0.4.2 → ai_pipeline_core-0.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-pipeline-core
-Version: 0.4.2
+Version: 0.4.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.4.2 → ai_pipeline_core-0.4.4}/ai_pipeline_core/__init__.py

@@ -64,7 +64,7 @@ from .prompt_manager import PromptManager
 from .settings import Settings
 from .testing import disable_run_logger, prefect_test_harness
 
-__version__ = "0.4.1"
+__version__ = "0.4.4"
 
 __all__ = [
     "AIMessageType",

{ai_pipeline_core-0.4.2 → ai_pipeline_core-0.4.4}/ai_pipeline_core/deployment/base.py

@@ -51,6 +51,7 @@ from .helpers import (
     send_webhook,
     upload_documents,
 )
+from .progress import flow_context, webhook_worker
 
 logger = get_pipeline_logger(__name__)
 
@@ -518,16 +519,39 @@ class PipelineDeployment(Generic[TOptions, TResult]):
                 else:
                     current_docs = input_docs
 
-                try:
-                    await active_flow(project_name, current_docs, options.model_dump())
-                except Exception as e:
-                    # Upload partial results on failure
-                    if context.output_documents_urls and store:
-                        all_docs = await store.load(run_scope, self._all_document_types())
-                        await upload_documents(all_docs, context.output_documents_urls)
-                    await self._send_completion(context, flow_run_id, project_name, result=None, error=str(e))
-                    completion_sent = True
-                    raise
+                # Set up intra-flow progress context so progress_update() works inside flows
+                flow_minutes = tuple(getattr(f, "estimated_minutes", 1) for f in self.flows)
+                completed_mins = sum(flow_minutes[: max(step - 1, 0)])
+                progress_queue: asyncio.Queue[ProgressRun | None] = asyncio.Queue()
+                wh_url = context.progress_webhook_url or ""
+                worker = asyncio.create_task(webhook_worker(progress_queue, wh_url)) if wh_url else None
+
+                with flow_context(
+                    webhook_url=wh_url,
+                    project_name=project_name,
+                    run_id=flow_run_id,
+                    flow_run_id=flow_run_id,
+                    flow_name=flow_name,
+                    step=step,
+                    total_steps=total_steps,
+                    flow_minutes=flow_minutes,
+                    completed_minutes=completed_mins,
+                    queue=progress_queue,
+                ):
+                    try:
+                        await active_flow(project_name, current_docs, options.model_dump())
+                    except Exception as e:
+                        # Upload partial results on failure
+                        if context.output_documents_urls and store:
+                            all_docs = await store.load(run_scope, self._all_document_types())
+                            await upload_documents(all_docs, context.output_documents_urls)
+                        await self._send_completion(context, flow_run_id, project_name, result=None, error=str(e))
+                        completion_sent = True
+                        raise
+                    finally:
+                        progress_queue.put_nowait(None)
+                        if worker:
+                            await worker
 
                 # Per-flow upload (load from store since @pipeline_flow saves there)
                 if context.output_documents_urls and store and output_types:
@@ -797,7 +821,31 @@ class PipelineDeployment(Generic[TOptions, TResult]):
                 else:
                     current_docs = initial_documents or []
 
-                await flow_fn(project_name, current_docs, options)
+                # Set up intra-flow progress context so progress_update() works inside flows
+                flow_minutes = tuple(getattr(f, "estimated_minutes", 1) for f in self.flows)
+                completed_mins = sum(flow_minutes[: max(step - 1, 0)])
+                progress_queue: asyncio.Queue[ProgressRun | None] = asyncio.Queue()
+                wh_url = context.progress_webhook_url or ""
+                worker = asyncio.create_task(webhook_worker(progress_queue, wh_url)) if wh_url else None
+
+                with flow_context(
+                    webhook_url=wh_url,
+                    project_name=project_name,
+                    run_id=str(run_uuid) if run_uuid else "",
+                    flow_run_id=str(run_uuid) if run_uuid else "",
+                    flow_name=flow_name,
+                    step=step,
+                    total_steps=total_steps,
+                    flow_minutes=flow_minutes,
+                    completed_minutes=completed_mins,
+                    queue=progress_queue,
+                ):
+                    try:
+                        await flow_fn(project_name, current_docs, options)
+                    finally:
+                        progress_queue.put_nowait(None)
+                        if worker:
+                            await worker
 
             # Build result from all documents in store
             if store:
@@ -828,12 +876,6 @@ class PipelineDeployment(Generic[TOptions, TResult]):
         """
         deployment = self
 
-        @flow(
-            name=self.name,
-            flow_run_name=f"{self.name}-{{project_name}}",
-            persist_result=True,
-            result_serializer="json",
-        )
         async def _deployment_flow(
             project_name: str,
             documents: list[Document],
@@ -851,7 +893,16 @@ class PipelineDeployment(Generic[TOptions, TResult]):
                 store.shutdown()
                 set_document_store(None)
 
-        return _deployment_flow
+        # Patch annotations so Prefect generates the parameter schema from the concrete types
+        _deployment_flow.__annotations__["options"] = self.options_type
+        _deployment_flow.__annotations__["return"] = self.result_type
+
+        return flow(
+            name=self.name,
+            flow_run_name=f"{self.name}-{{project_name}}",
+            persist_result=True,
+            result_serializer="json",
+        )(_deployment_flow)
 
 
 __all__ = [

{ai_pipeline_core-0.4.2 → ai_pipeline_core-0.4.4}/ai_pipeline_core/deployment/progress.py

@@ -9,6 +9,8 @@ from dataclasses import dataclass
 from datetime import UTC, datetime
 from uuid import UUID
 
+from prefect import get_client
+
 from ai_pipeline_core.logging import get_pipeline_logger
 
 from .contract import ProgressRun
@@ -38,9 +40,14 @@ _context: ContextVar[ProgressContext | None] = ContextVar("progress_context", de
 
 
 async def update(fraction: float, message: str = "") -> None:
-    """Report intra-flow progress (0.0-1.0). No-op without context."""
+    """Report intra-flow progress (0.0-1.0). No-op without context.
+
+    Sends webhook payload (if webhook_url configured) AND updates Prefect
+    flow run labels (if flow_run_id available) so both push and poll consumers
+    see progress, and staleness detection stays current.
+    """
     ctx = _context.get()
-    if ctx is None or not ctx.webhook_url:
+    if ctx is None:
         return
 
     fraction = max(0.0, min(1.0, fraction))
@@ -50,22 +57,43 @@ async def update(fraction: float, message: str = "") -> None:
     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(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)
+    step_progress = round(fraction, 4)
+
+    # Enqueue webhook payload for async delivery
+    if ctx.webhook_url:
+        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(UTC),
+            step=ctx.step,
+            total_steps=ctx.total_steps,
+            flow_name=ctx.flow_name,
+            status="progress",
+            progress=overall,
+            step_progress=step_progress,
+            message=message,
+        )
+        ctx.queue.put_nowait(payload)
+
+    # Update Prefect labels so polling consumers and staleness detection stay current
+    if ctx.flow_run_id:
+        try:
+            async with get_client() as client:
+                await client.update_flow_run_labels(
+                    flow_run_id=UUID(ctx.flow_run_id),
+                    labels={
+                        "progress.step": ctx.step,
+                        "progress.total_steps": ctx.total_steps,
+                        "progress.flow_name": ctx.flow_name,
+                        "progress.status": "progress",
+                        "progress.progress": overall,
+                        "progress.step_progress": step_progress,
+                        "progress.message": message,
+                    },
+                )
+        except Exception as e:
+            logger.warning(f"Progress label update failed: {e}")
 
 
 async def webhook_worker(

{ai_pipeline_core-0.4.2 → ai_pipeline_core-0.4.4}/ai_pipeline_core/llm/client.py

@@ -258,10 +258,10 @@ def _model_name_to_openrouter_model(model: ModelName) -> str:
     Returns:
         OpenRouter model name.
     """
-    if model == "gemini-3-flash-search":
-        return "google/gemini-3-flash:online"
     if model == "sonar-pro-search":
         return "perplexity/sonar-pro-search"
+    if model.endswith("-search"):
+        model = model.replace("-search", ":online")
     if model.startswith("gemini"):
         return f"google/{model}"
     elif model.startswith("gpt"):

{ai_pipeline_core-0.4.2 → ai_pipeline_core-0.4.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ai-pipeline-core"
-version = "0.4.2"
+version = "0.4.4"
 description = "Core utilities for AI-powered processing pipelines using prefect"
 readme = "README.md"
 license = {text = "MIT"}
@@ -61,7 +61,7 @@ dev = [
 minversion = "8.0"
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"
-addopts = "-q -n 8 -m 'not integration and not clickhouse'"
+addopts = "-q -m 'not integration and not clickhouse'"
 testpaths = ["tests"]
 markers = [
     "integration: marks tests as integration tests (deselect with '-m \"not integration\"')",