ai-pipeline-core 0.4.4__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.4"
+__version__ = "0.4.6"
 
 __all__ = [
     "AIMessageType",

ai_pipeline_core/deployment/deploy.py

@@ -312,6 +312,28 @@ class Deployer:
 
         return builds
 
+    def _build_vendor_packages(self) -> list[Path]:
+        """Build vendor wheels from [tool.deploy].vendor_packages paths.
+
+        Returns:
+            List of built wheel paths, deduplicated by filename.
+        """
+        vendor_paths: list[str] = self._pyproject_data.get("tool", {}).get("deploy", {}).get("vendor_packages", [])
+        if not vendor_paths:
+            return []
+
+        self._info(f"Building {len(vendor_paths)} vendor package(s)")
+        wheels: list[Path] = []
+        seen: set[str] = set()
+        for vendor_path_str in vendor_paths:
+            vendor_dir = Path(vendor_path_str).resolve()
+            wheel = self._build_wheel_from_source(vendor_dir)
+            if wheel.name not in seen:
+                wheels.append(wheel)
+                seen.add(wheel.name)
+                self._success(f"Built vendor wheel: {wheel.name}")
+        return wheels
+
     def _create_gcs_bucket(self, bucket_folder: str) -> Any:
         """Create a GcsBucket instance for uploading files.
 
@@ -361,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}")
@@ -457,6 +479,14 @@ class Deployer:
             paused=False,
         )
 
+        # Populate parameter schema from flow function signature
+        deployment._set_defaults_from_flow(flow)  # pyright: ignore[reportPossiblyUnboundVariable]
+
+        # Inject result type schema so consumers can discover the response shape
+        return_type = getattr(flow.fn, "__annotations__", {}).get("return")  # pyright: ignore[reportPossiblyUnboundVariable]
+        if return_type is not None and hasattr(return_type, "model_json_schema"):
+            deployment._parameter_openapi_schema.definitions["_ResultSchema"] = return_type.model_json_schema()
+
         # Verify work pool exists before deploying
         async with get_client() as client:
             try:
@@ -494,21 +524,27 @@ class Deployer:
         # Phase 2: Build agent bundles (if configured)
         agent_builds = self._build_agents()
 
-        # Phase 3: Upload flow package (include private dependency wheels from agent builds)
-        vendor_wheels: list[Path] = []
-        if agent_builds:
-            seen: 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 and "cli_agents" in filename:
-                        vendor_wheels.append(filepath)
-                        seen.add(filename)
+        # Phase 3: Build vendor packages from [tool.deploy].vendor_packages
+        vendor_wheels = self._build_vendor_packages()
+
+        # 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)
 
-        # Phase 4: Upload agent bundles
+        # Phase 5: Upload agent bundles
         await self._upload_agents(agent_builds)
 
-        # Phase 5: Create/update Prefect deployment
+        # Phase 6: Create/update Prefect deployment
         await self._deploy_via_api(agent_builds)
 
         print()

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.
+
+    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
 
-    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
+    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
-
-            full_name = f"{deployment_class.name}/{deployment_name or deployment_class.name}"
-
-            result = await run_remote_deployment(full_name, parameters)
+        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, 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.4.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.4
+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.4.dist-info → ai_pipeline_core-0.4.6.dist-info}/RECORD

@@ -1,4 +1,4 @@
-ai_pipeline_core/__init__.py,sha256=VcU67cNpAB6EJf2V_FxUx4ZFosbpb6IzFFwpJauHJzk,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=TCF4fH5f-K1ADODZHEyf-x7PJzDbv4qtWxlpoCe_mTs,22909
+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=tOexisKEeeBoHLGYZWqcjr2H-nqqYc6kvoDL72AW78w,4661
+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.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,,
+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,,