ai-pipeline-core 0.4.5__py3-none-any.whl → 0.4.6__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.5"
+__version__ = "0.4.6"
 
 __all__ = [
     "AIMessageType",

ai_pipeline_core/deployment/deploy.py

@@ -383,7 +383,7 @@ class Deployer:
         dest_uri = f"gs://{self.config['bucket']}/{flow_folder}/{tarball.name}"
         self._info(f"Uploading to {dest_uri}")
 
-        tarball_bytes = tarball.read_bytes()  # noqa: ASYNC240
+        tarball_bytes = tarball.read_bytes()
         await bucket.write_path(tarball.name, tarball_bytes)
 
         self._success(f"Package uploaded to {flow_folder}/{tarball.name}")
@@ -527,15 +527,16 @@ class Deployer:
         # Phase 3: Build vendor packages from [tool.deploy].vendor_packages
         vendor_wheels = self._build_vendor_packages()
 
-        # Also include cli_agents wheels from agent builds
-        if agent_builds:
-            seen_agent: set[str] = set()
-            for build_info in agent_builds.values():
-                for filename, filepath in build_info["files"].items():
-                    if filename.endswith(".whl") and filename not in seen_agent and "cli_agents" in filename:
-                        if filename not in {w.name for w in vendor_wheels}:
-                            vendor_wheels.append(filepath)
-                        seen_agent.add(filename)
+        # Build cli-agents wheel if source is configured — it's a private package
+        # not on PyPI, so the worker needs the wheel even when no agents are deployed
+        cli_agents_source = self._get_cli_agents_source()
+        if cli_agents_source:
+            cli_dir = Path(cli_agents_source).resolve()
+            if (cli_dir / "pyproject.toml").exists():
+                cli_wheel = self._build_wheel_from_source(cli_dir)
+                if cli_wheel.name not in {w.name for w in vendor_wheels}:
+                    vendor_wheels.append(cli_wheel)
+                    self._success(f"Built cli-agents vendor wheel: {cli_wheel.name}")
 
         # Phase 4: Upload flow package + vendor wheels
         await self._upload_package(tarball, vendor_wheels)

ai_pipeline_core/deployment/remote.py

@@ -1,9 +1,10 @@
 """Remote deployment utilities for calling PipelineDeployment flows via Prefect."""
 
-import inspect
-from collections.abc import Callable
+import asyncio
+from collections.abc import Awaitable, Callable, Coroutine
 from functools import wraps
-from typing import Any, ParamSpec, TypeVar, cast
+from typing import Any, TypeVar, cast
+from uuid import UUID
 
 from prefect import get_client
 from prefect.client.orchestration import PrefectClient
@@ -13,14 +14,20 @@ from prefect.deployments.flow_runs import run_deployment
 from prefect.exceptions import ObjectNotFound
 
 from ai_pipeline_core.deployment import DeploymentContext, DeploymentResult, PipelineDeployment
+from ai_pipeline_core.documents import Document
+from ai_pipeline_core.logging import get_pipeline_logger
 from ai_pipeline_core.observability.tracing import TraceLevel, set_trace_cost, trace
 from ai_pipeline_core.pipeline.options import FlowOptions
 from ai_pipeline_core.settings import settings
 
-P = ParamSpec("P")
+logger = get_pipeline_logger(__name__)
+
 TOptions = TypeVar("TOptions", bound=FlowOptions)
 TResult = TypeVar("TResult", bound=DeploymentResult)
 
+ProgressCallback = Callable[[float, str], Awaitable[None]]
+"""Signature for remote deployment progress callbacks: (fraction, message) -> None."""
+
 
 def _is_already_traced(func: Callable[..., Any]) -> bool:
     """Check if function or its __wrapped__ has __is_traced__ attribute."""
@@ -30,17 +37,80 @@ def _is_already_traced(func: Callable[..., Any]) -> bool:
     return getattr(wrapped, "__is_traced__", False) if wrapped else False
 
 
-async def run_remote_deployment(deployment_name: str, parameters: dict[str, Any]) -> Any:
-    """Run a remote Prefect deployment, trying local client first then remote."""
+_POLL_INTERVAL = 5.0
+
+
+async def _poll_remote_flow_run(
+    client: PrefectClient,
+    flow_run_id: UUID,
+    deployment_name: str,
+    poll_interval: float = _POLL_INTERVAL,
+    on_progress: ProgressCallback | None = None,
+) -> Any:
+    """Poll a remote flow run until final, invoking on_progress callback with progress.
 
-    async def _run(client: PrefectClient, as_subflow: bool) -> Any:
-        fr: FlowRun = await run_deployment(client=client, name=deployment_name, parameters=parameters, as_subflow=as_subflow)  # type: ignore
-        return await fr.state.result()  # type: ignore
+    Reads the remote flow run's progress labels on each poll cycle and calls
+    on_progress(fraction, message) if provided. Without a callback, no progress
+    is reported. Only sends 1.0 on successful completion (not failure).
+    """
+    last_fraction = 0.0
+
+    while True:
+        try:
+            flow_run = await client.read_flow_run(flow_run_id)
+        except Exception:
+            logger.warning("Failed to poll remote flow run %s", flow_run_id, exc_info=True)
+            await asyncio.sleep(poll_interval)
+            continue
+
+        state = flow_run.state
+        if state and state.is_final():
+            if on_progress and state.is_completed():
+                await on_progress(1.0, f"[{deployment_name}] Completed")
+            return await state.result()  # type: ignore[union-attr]
+
+        if on_progress:
+            labels: dict[str, Any] = flow_run.labels or {}
+            progress_val = labels.get("progress.progress")
+
+            if progress_val is not None:
+                fraction = max(float(progress_val), last_fraction)
+                last_fraction = fraction
+                flow_name = str(labels.get("progress.flow_name", ""))
+                message = str(labels.get("progress.message", ""))
+                display = f"[{deployment_name}] {flow_name}: {message}" if flow_name else f"[{deployment_name}] Running"
+                await on_progress(fraction, display)
+            else:
+                await on_progress(last_fraction, f"[{deployment_name}] Waiting to start")
+
+        await asyncio.sleep(poll_interval)
+
+
+async def run_remote_deployment(
+    deployment_name: str,
+    parameters: dict[str, Any],
+    on_progress: ProgressCallback | None = None,
+) -> Any:
+    """Run a remote Prefect deployment with optional progress callback.
+
+    Creates the remote flow run immediately (timeout=0) then polls its state,
+    invoking on_progress(fraction, message) on each poll cycle if provided.
+    """
+
+    async def _create_and_poll(client: PrefectClient, as_subflow: bool) -> Any:
+        fr: FlowRun = await run_deployment(
+            client=client,
+            name=deployment_name,
+            parameters=parameters,
+            as_subflow=as_subflow,
+            timeout=0,
+        )  # type: ignore
+        return await _poll_remote_flow_run(client, fr.id, deployment_name, on_progress=on_progress)
 
     async with get_client() as client:
         try:
             await client.read_deployment_by_name(name=deployment_name)
-            return await _run(client, True)  # noqa: FBT003
+            return await _create_and_poll(client, True)  # noqa: FBT003
         except ObjectNotFound:
             pass
 
@@ -56,7 +126,7 @@ async def run_remote_deployment(deployment_name: str, parameters: dict[str, Any]
             await client.read_deployment_by_name(name=deployment_name)
             ctx = AsyncClientContext.model_construct(client=client, _httpx_settings=None, _context_stack=0)
             with ctx:
-                return await _run(client, False)  # noqa: FBT003
+                return await _create_and_poll(client, False)  # noqa: FBT003
         except ObjectNotFound:
             pass
 
@@ -70,32 +140,38 @@ def remote_deployment(
     name: str | None = None,
     trace_level: TraceLevel = "always",
     trace_cost: float | None = None,
-) -> Callable[[Callable[P, TResult]], Callable[P, TResult]]:
-    """Decorator to call PipelineDeployment flows remotely with automatic serialization."""
+) -> Callable[[Callable[..., Any]], Callable[..., Coroutine[Any, Any, TResult]]]:
+    """Decorator to call PipelineDeployment flows remotely with automatic serialization.
+
+    The decorated function's body is never executed — it serves as a typed stub.
+    The wrapper enforces the deployment contract: (project_name, documents, options, context).
+    """
 
-    def decorator(func: Callable[P, TResult]) -> Callable[P, TResult]:
+    def decorator(func: Callable[..., Any]) -> Callable[..., Coroutine[Any, Any, TResult]]:
         fname = getattr(func, "__name__", deployment_class.name)
 
         if _is_already_traced(func):
             raise TypeError(f"@remote_deployment target '{fname}' already has @trace")
 
         @wraps(func)
-        async def _wrapper(*args: P.args, **kwargs: P.kwargs) -> TResult:
-            sig = inspect.signature(func)
-            bound = sig.bind(*args, **kwargs)
-            bound.apply_defaults()
-
-            # Pass parameters with proper types - Prefect handles Pydantic serialization
-            parameters: dict[str, Any] = {}
-            for pname, value in bound.arguments.items():
-                if value is None and pname == "context":
-                    parameters[pname] = DeploymentContext()
-                else:
-                    parameters[pname] = value
+        async def _wrapper(
+            project_name: str,
+            documents: list[Document],
+            options: TOptions,
+            context: DeploymentContext | None = None,
+            *,
+            on_progress: ProgressCallback | None = None,
+        ) -> TResult:
+            parameters: dict[str, Any] = {
+                "project_name": project_name,
+                "documents": documents,
+                "options": options,
+                "context": context if context is not None else DeploymentContext(),
+            }
 
             full_name = f"{deployment_class.name}/{deployment_name or deployment_class.name.replace('-', '_')}"
 
-            result = await run_remote_deployment(full_name, parameters)
+            result = await run_remote_deployment(full_name, parameters, on_progress=on_progress)
 
             if trace_cost is not None and trace_cost > 0:
                 set_trace_cost(trace_cost)
@@ -111,6 +187,6 @@ def remote_deployment(
             name=name or deployment_class.name,
         )(_wrapper)
 
-        return traced_wrapper  # type: ignore[return-value]
+        return traced_wrapper
 
     return decorator

{ai_pipeline_core-0.4.5.dist-info → ai_pipeline_core-0.4.6.dist-info}/METADATA

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

@@ -1,4 +1,4 @@
-ai_pipeline_core/__init__.py,sha256=ydSl96FjnoGhcGRC_Pnr1n7R0ub3vZWnKxyLWILJ54o,3270
+ai_pipeline_core/__init__.py,sha256=QAQSyrKafsNov4dy9vNqTVarh2nDdrthMfYh7X-3Mcg,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
@@ -7,10 +7,10 @@ ai_pipeline_core/testing.py,sha256=jIRrLxNvTwdamucfJoHET2qMeRhhMZV9uEJXO5vAfis,2
 ai_pipeline_core/deployment/__init__.py,sha256=wTkVK6gcEQvqBajFMTAuodRONpN25yHbR1jtcumf0WQ,900
 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=3Pojw_HonW46RbjWf6AGNtv-1F46Ed4y71pgQL0teAk,24576
+ai_pipeline_core/deployment/deploy.py,sha256=y5FxKMm7nGwkjzA74pTffO7A82MaDuajx6LHGTem8bI,24662
 ai_pipeline_core/deployment/helpers.py,sha256=yVtGFUs4AFXkpLkiQ_ale0nXXt5btfWSb5PAbikQHNs,3312
 ai_pipeline_core/deployment/progress.py,sha256=rO2g8VIh7EpzxzGGAroXEpveWoWZkk66jkDW22BY4j8,4827
-ai_pipeline_core/deployment/remote.py,sha256=pOle4Dw_U4j9mLbk8NQdmKUys4uaGJyOKMZXhRvnw3E,4679
+ai_pipeline_core/deployment/remote.py,sha256=tOBbICtPXbJHN8QA9juCqkeP9PqEc16mPyBrwaOwEt4,7434
 ai_pipeline_core/docs_generator/__init__.py,sha256=JbWbk-Lw5GgWrCMRuw8zvKNTZY2jXv7XqoMiBYudvRI,1255
 ai_pipeline_core/docs_generator/__main__.py,sha256=CH4agiM2suFJ63MhTg5m0GuXdc40z-6o4ojR72JQWVA,145
 ai_pipeline_core/docs_generator/cli.py,sha256=8OjdMtzQraPxWN3uPapSNJnKyPLPtnygKL0rF5JL2GY,7172
@@ -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.5.dist-info/METADATA,sha256=U-sw_nFzaA55Z_YzxVvvh9Mb8N3L50nU7ME6YkQ9tO4,29947
-ai_pipeline_core-0.4.5.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-ai_pipeline_core-0.4.5.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
-ai_pipeline_core-0.4.5.dist-info/RECORD,,
+ai_pipeline_core-0.4.6.dist-info/METADATA,sha256=hyy3vHyR5xZ5GRg4Nrp8BxKojp3pa-RDg1KWPPm5_O8,29947
+ai_pipeline_core-0.4.6.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+ai_pipeline_core-0.4.6.dist-info/licenses/LICENSE,sha256=kKj8mfbdWwkyG3U6n7ztB3bAZlEwShTkAsvaY657i3I,1074
+ai_pipeline_core-0.4.6.dist-info/RECORD,,