ai-pipeline-core 0.4.3__py3-none-any.whl → 0.4.4__py3-none-any.whl

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.3"
+__version__ = "0.4.4"
 
 __all__ = [
     "AIMessageType",

ai_pipeline_core/deployment/base.py

@@ -821,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:
@@ -852,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],
@@ -875,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/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/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.3.dist-info → ai_pipeline_core-0.4.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-pipeline-core
-Version: 0.4.3
+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.3.dist-info → ai_pipeline_core-0.4.4.dist-info}/RECORD

@@ -1,15 +1,15 @@
-ai_pipeline_core/__init__.py,sha256=YxuYpXNT3kliHsZSnbbwCZxtDAdTgT5TkKp-ONSl-q4,3270
+ai_pipeline_core/__init__.py,sha256=VcU67cNpAB6EJf2V_FxUx4ZFosbpb6IzFFwpJauHJzk,3270
 ai_pipeline_core/exceptions.py,sha256=csAl7vq6xjSFBF8-UM9WZODCbhsOdOG5zH6IbA8iteM,1280
 ai_pipeline_core/prompt_manager.py,sha256=3wFkL5rrjtUT1cLInkgyhS8hKnO4MeD1cdXAEuLhgoE,9459
 ai_pipeline_core/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_pipeline_core/settings.py,sha256=BUz8JEFfJQrdE4rNOhQWwxnTrfekLjWkoy-3wDZQ7PY,5142
 ai_pipeline_core/testing.py,sha256=jIRrLxNvTwdamucfJoHET2qMeRhhMZV9uEJXO5vAfis,279
 ai_pipeline_core/deployment/__init__.py,sha256=wTkVK6gcEQvqBajFMTAuodRONpN25yHbR1jtcumf0WQ,900
-ai_pipeline_core/deployment/base.py,sha256=4fY9JXLW5PMtn7QFSt8OVRlpKSFTRQ2YnmvB_cZr8H0,36025
+ai_pipeline_core/deployment/base.py,sha256=bGSnDdrw6cLM_TItAiwptnwApbw5wkoIGY9pnwDvOTQ,37485
 ai_pipeline_core/deployment/contract.py,sha256=a1qbHhneTGB27oSOUy79CUIhOIzOoq37M63XoIMzA4Y,1952
 ai_pipeline_core/deployment/deploy.py,sha256=TCF4fH5f-K1ADODZHEyf-x7PJzDbv4qtWxlpoCe_mTs,22909
 ai_pipeline_core/deployment/helpers.py,sha256=yVtGFUs4AFXkpLkiQ_ale0nXXt5btfWSb5PAbikQHNs,3312
-ai_pipeline_core/deployment/progress.py,sha256=5tVD9nW0N-b8Z2BxazcWCWHFpLu6pJ-eqPmRyj68X6Y,3591
+ai_pipeline_core/deployment/progress.py,sha256=rO2g8VIh7EpzxzGGAroXEpveWoWZkk66jkDW22BY4j8,4827
 ai_pipeline_core/deployment/remote.py,sha256=tOexisKEeeBoHLGYZWqcjr2H-nqqYc6kvoDL72AW78w,4661
 ai_pipeline_core/docs_generator/__init__.py,sha256=JbWbk-Lw5GgWrCMRuw8zvKNTZY2jXv7XqoMiBYudvRI,1255
 ai_pipeline_core/docs_generator/__main__.py,sha256=CH4agiM2suFJ63MhTg5m0GuXdc40z-6o4ojR72JQWVA,145
@@ -38,7 +38,7 @@ ai_pipeline_core/images/__init__.py,sha256=Hc2QKR27Q2Q-h5nH-EbzfxdE3dHArBm-st5_x
 ai_pipeline_core/images/_processing.py,sha256=MrCuPGsyyEl9UlXYIPhZs0wN8CPTMZmejV2Lo2wyCZk,4362
 ai_pipeline_core/llm/__init__.py,sha256=oyRvYD5DLDl7JIRTBUaiVz6jUC5dLLujkMNFpfRp2zc,795
 ai_pipeline_core/llm/ai_messages.py,sha256=Ycmntk5d6NUFqVVsnNR_IDwJUFuHYEH7CPvmmDfYaJI,17424
-ai_pipeline_core/llm/client.py,sha256=CjxOiniuy5CEsA_Xz0KPLCBthbnUfC43fTpuDcqkIUM,30276
+ai_pipeline_core/llm/client.py,sha256=N8eH9bY2rF28U5kGK0HQ3ibKvphcipSMLVVxtxtut8Y,30275
 ai_pipeline_core/llm/model_options.py,sha256=hg8xR0RJdJKp8QJNA4EbLnfFsnkE4HnxD85aYxc--hM,9164
 ai_pipeline_core/llm/model_response.py,sha256=Ml9wcssSssqibReJxCc9EQu488pz69Cmq_XNBs_xmak,12219
 ai_pipeline_core/llm/model_types.py,sha256=qHoUPPEkHu9B4kJ5xcIC09fk72v667ZxvzigxtgLpVo,2174
@@ -70,7 +70,7 @@ ai_pipeline_core/observability/_tracking/_writer.py,sha256=xZjwYyIxDzzzPxqkKjYAY
 ai_pipeline_core/pipeline/__init__.py,sha256=uMv1jwSyq8Ym8Hbn5097twBJLdwN1iMeqnVM4EWyrhA,282
 ai_pipeline_core/pipeline/decorators.py,sha256=CDJAeOjGLt5Ewc0Jc9zEuwLZwKyutOv89LSRS9dcXmI,37456
 ai_pipeline_core/pipeline/options.py,sha256=KF4FcT085-IwX8r649v0a9ua5xnApM0qG2wJHWbq39A,438
-ai_pipeline_core-0.4.3.dist-info/METADATA,sha256=j3lgsgzMN1D_4twQ3WqCzzaY1yfQThSKOxwMfZILk50,29947
-ai_pipeline_core-0.4.3.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-ai_pipeline_core-0.4.3.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
-ai_pipeline_core-0.4.3.dist-info/RECORD,,
+ai_pipeline_core-0.4.4.dist-info/METADATA,sha256=DWL9spC2_pILAtrr6BARsTIjqWk3-VJO6WNyTqdEqyQ,29947
+ai_pipeline_core-0.4.4.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+ai_pipeline_core-0.4.4.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
+ai_pipeline_core-0.4.4.dist-info/RECORD,,