ai-pipeline-core 0.4.1__tar.gz → 0.4.3__tar.gz

{ai_pipeline_core-0.4.1 → ai_pipeline_core-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-pipeline-core
-Version: 0.4.1
+Version: 0.4.3
 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.1 → ai_pipeline_core-0.4.3}/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.3"
 
 __all__ = [
     "AIMessageType",

{ai_pipeline_core-0.4.1 → ai_pipeline_core-0.4.3}/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__)
 
@@ -132,7 +133,7 @@ class FlowCallable(Protocol):
     output_document_types: list[type[Document]]
     estimated_minutes: int
 
-    def __call__(self, project_name: str, documents: list[Document], flow_options: FlowOptions) -> Any:  # type: ignore[type-arg]
+    def __call__(self, project_name: str, documents: list[Document], flow_options: FlowOptions | dict[str, Any]) -> Any:  # type: ignore[type-arg]
         """Execute the flow with standard pipeline signature."""
         ...
 
@@ -518,16 +519,39 @@ class PipelineDeployment(Generic[TOptions, TResult]):
                 else:
                     current_docs = input_docs
 
-                try:
-                    await active_flow(project_name, current_docs, options)
-                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:

{ai_pipeline_core-0.4.1 → ai_pipeline_core-0.4.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ai-pipeline-core"
-version = "0.4.1"
+version = "0.4.3"
 description = "Core utilities for AI-powered processing pipelines using prefect"
 readme = "README.md"
 license = {text = "MIT"}