ai-pipeline-core 0.2.9__py3-none-any.whl → 0.3.3__py3-none-any.whl

ai_pipeline_core/utils/deploy.py

@@ -18,10 +18,13 @@ Usage:
 
 import argparse
 import asyncio
+import json
 import subprocess
 import sys
+import tempfile
 import tomllib
 import traceback
+from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any, Optional
 
@@ -70,6 +73,8 @@ class Deployer:
         with open(pyproject_path, "rb") as f:
             data = tomllib.load(f)
 
+        self._pyproject_data = data
+
         project = data.get("project", {})
         name = project.get("name")
         version = project.get("version")
@@ -160,6 +165,192 @@ class Deployer:
         self._success(f"Built {tarball_path.name} ({tarball_path.stat().st_size // 1024} KB)")
         return tarball_path
 
+    # -- Agent build/upload support --
+
+    def _load_agent_config(self) -> dict[str, dict[str, Any]]:
+        """Load [tool.deploy.agents] from pyproject.toml.
+
+        Returns:
+            Dict mapping agent name to config (path, extra_vendor).
+            Empty dict if no agents configured.
+        """
+        return self._pyproject_data.get("tool", {}).get("deploy", {}).get("agents", {})
+
+    def _get_cli_agents_source(self) -> str | None:
+        """Get cli_agents_source path from [tool.deploy]."""
+        return self._pyproject_data.get("tool", {}).get("deploy", {}).get("cli_agents_source")
+
+    def _build_wheel_from_source(self, source_dir: Path) -> Path:
+        """Build a wheel from a source directory.
+
+        Args:
+            source_dir: Directory containing pyproject.toml
+
+        Returns:
+            Path to built .whl file in a temp dist directory
+        """
+        if not (source_dir / "pyproject.toml").exists():
+            self._die(f"No pyproject.toml in {source_dir}")
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmp_dist = Path(tmpdir) / "dist"
+            result = subprocess.run(
+                [sys.executable, "-m", "build", "--wheel", "--outdir", str(tmp_dist)],
+                cwd=source_dir,
+                capture_output=True,
+                text=True,
+            )
+            if result.returncode != 0:
+                self._die(f"Wheel build failed for {source_dir.name}:\n{result.stderr}")
+
+            wheels = list(tmp_dist.glob("*.whl"))
+            if not wheels:
+                self._die(f"No wheel produced for {source_dir.name}")
+
+            # Copy to persistent dist/ under source_dir
+            dist_dir = source_dir / "dist"
+            dist_dir.mkdir(exist_ok=True)
+            output = dist_dir / wheels[0].name
+            output.write_bytes(wheels[0].read_bytes())
+            return output
+
+    def _build_agents(self) -> dict[str, dict[str, Any]]:
+        """Build agent wheels and manifests for all configured agents.
+
+        Returns:
+            Dict mapping agent name to build info:
+                {name: {"manifest_json": str, "files": {filename: Path}}}
+            Empty dict if no agents configured.
+        """
+        agent_config = self._load_agent_config()
+        if not agent_config:
+            return {}
+
+        cli_agents_source = self._get_cli_agents_source()
+        if not cli_agents_source:
+            self._die(
+                "Agents configured in [tool.deploy.agents] but "
+                "[tool.deploy].cli_agents_source is not set.\n"
+                "Add to pyproject.toml:\n"
+                '  [tool.deploy]\n  cli_agents_source = "vendor/cli-agents"'
+            )
+
+        self._info(f"Building {len(agent_config)} agent(s): {', '.join(agent_config)}")
+
+        # Build cli-agents wheel once (shared across all agents)
+        cli_agents_dir = Path(cli_agents_source).resolve()
+        if not (cli_agents_dir / "pyproject.toml").exists():
+            self._die(f"cli-agents source not found at {cli_agents_dir}")
+
+        cli_agents_wheel = self._build_wheel_from_source(cli_agents_dir)
+        self._success(f"Built cli-agents wheel: {cli_agents_wheel.name}")
+
+        builds: dict[str, dict[str, Any]] = {}
+
+        for agent_name, config in agent_config.items():
+            agent_path = Path(config["path"]).resolve()
+            if not (agent_path / "pyproject.toml").exists():
+                self._die(
+                    f"Agent '{agent_name}' path not found: {agent_path}\n"
+                    f"Check [tool.deploy.agents.{agent_name}].path in pyproject.toml"
+                )
+
+            # Read module_name from agent's pyproject.toml
+            with open(agent_path / "pyproject.toml", "rb") as f:
+                agent_pyproject = tomllib.load(f)
+
+            module_name = agent_pyproject.get("tool", {}).get("agent", {}).get("module")
+            if not module_name:
+                self._die(
+                    f"Agent '{agent_name}' missing [tool.agent].module in "
+                    f"{agent_path / 'pyproject.toml'}\n"
+                    f'Add:\n  [tool.agent]\n  module = "agent_{agent_name}"'
+                )
+
+            # Build agent wheel
+            agent_wheel = self._build_wheel_from_source(agent_path)
+            self._success(f"Built agent wheel: {agent_wheel.name}")
+
+            # Collect all files for this agent bundle
+            files: dict[str, Path] = {
+                agent_wheel.name: agent_wheel,
+                cli_agents_wheel.name: cli_agents_wheel,
+            }
+
+            # Build extra_vendor packages from repo root
+            vendor_packages: list[str] = []
+            extra_built: set[str] = set()
+            for vendor_name in config.get("extra_vendor", []):
+                extra_source_dir = Path(vendor_name).resolve()
+                if not (extra_source_dir / "pyproject.toml").exists():
+                    self._die(
+                        f"Extra vendor '{vendor_name}' for agent '{agent_name}' "
+                        f"not found at {extra_source_dir}\n"
+                        f"Ensure the directory exists at repo root with pyproject.toml"
+                    )
+                vendor_wheel = self._build_wheel_from_source(extra_source_dir)
+                files[vendor_wheel.name] = vendor_wheel
+                vendor_packages.append(vendor_wheel.name)
+                extra_built.add(extra_source_dir.name.replace("-", "_"))
+                self._success(f"Built vendor wheel: {vendor_wheel.name}")
+
+            # Collect existing vendor/*.whl and vendor/*.tar.gz from agent directory,
+            # skipping packages already built from extra_vendor
+            agent_vendor_dir = agent_path / "vendor"
+            if agent_vendor_dir.exists():
+                for pkg in list(agent_vendor_dir.glob("*.whl")) + list(
+                    agent_vendor_dir.glob("*.tar.gz")
+                ):
+                    pkg_base = pkg.name.split("-")[0].replace("-", "_")
+                    if pkg.name not in files and pkg_base not in extra_built:
+                        files[pkg.name] = pkg
+                        vendor_packages.append(pkg.name)
+
+            # Write manifest (plain JSON dict, compatible with AgentManifest schema)
+            manifest = {
+                "module_name": module_name,
+                "agent_wheel": agent_wheel.name,
+                "cli_agents_wheel": cli_agents_wheel.name,
+                "vendor_packages": vendor_packages,
+                "built_at": datetime.now(timezone.utc).isoformat(),
+            }
+            manifest_json = json.dumps(manifest, indent=2)
+
+            builds[agent_name] = {"manifest_json": manifest_json, "files": files}
+            self._success(f"Agent '{agent_name}' bundle ready ({module_name}, {len(files)} files)")
+
+        return builds
+
+    async def _upload_agents(self, agent_builds: dict[str, dict[str, Any]]):
+        """Upload agent bundles to GCS.
+
+        Args:
+            agent_builds: Output from _build_agents()
+        """
+        if not agent_builds:
+            return
+
+        flow_folder = self.config["folder"].split("/", 1)[1] if "/" in self.config["folder"] else ""
+        base_uri = f"gs://{self.config['bucket']}/flows"
+        base_storage = await Storage.from_uri(base_uri)
+        base_storage = base_storage.with_base(flow_folder)
+
+        for agent_name, build_info in agent_builds.items():
+            agent_storage = base_storage.with_base(f"agents/{agent_name}")
+            self._info(f"Uploading agent '{agent_name}' bundle to {agent_storage.url_for('')}")
+
+            # Upload manifest
+            await agent_storage.write_bytes(
+                "manifest.json",
+                build_info["manifest_json"].encode(),
+            )
+
+            # Upload wheels
+            for filename, filepath in build_info["files"].items():
+                await agent_storage.write_bytes(filename, filepath.read_bytes())
+
+            self._success(f"Agent '{agent_name}' uploaded ({len(build_info['files'])} files)")
+
     async def _upload_package(self, tarball: Path):
         """Upload package tarball to Google Cloud Storage using Storage abstraction.
 
@@ -184,13 +375,17 @@ class Deployer:
 
         self._success(f"Package uploaded to {self.config['folder']}/{tarball.name}")
 
-    async def _deploy_via_api(self):
+    async def _deploy_via_api(self, agent_builds: dict[str, dict[str, Any]] | None = None):
         """Create or update Prefect deployment using RunnerDeployment pattern.
 
         This is the official Prefect approach that:
         1. Automatically creates/updates the flow registration
         2. Handles deployment create vs update logic
         3. Properly formats all parameters for the API
+
+        Args:
+            agent_builds: Output from _build_agents(). If non-empty, sets
+                AGENT_BUNDLES_URI env var on the deployment.
         """
         # Define entrypoint (assumes flow function has same name as package)
         entrypoint = f"{self.config['package']}:{self.config['package']}"
@@ -244,6 +439,13 @@ class Deployer:
         # This is the official Prefect pattern that handles all the complexity
         self._info(f"Creating deployment for flow '{flow.name}'")
 
+        # Set AGENT_BUNDLES_URI env var if agents were built
+        job_variables: dict[str, Any] = {}
+        if agent_builds:
+            bundles_uri = f"gs://{self.config['bucket']}/{self.config['folder']}/agents"
+            job_variables["env"] = {"AGENT_BUNDLES_URI": bundles_uri}
+            self._info(f"Setting AGENT_BUNDLES_URI={bundles_uri}")
+
         deployment = RunnerDeployment(
             name=self.config["package"],
             flow_name=flow.name,
@@ -256,7 +458,7 @@ class Deployer:
             or f"Deployment for {self.config['package']} v{self.config['version']}",
             storage=_PullStepStorage(pull_steps),
             parameters={},
-            job_variables={},
+            job_variables=job_variables,
             paused=False,
         )
 
@@ -296,14 +498,20 @@ class Deployer:
         print("=" * 70)
         print()
 
-        # Phase 1: Build
+        # Phase 1: Build flow package
         tarball = self._build_package()
 
-        # Phase 2: Upload
+        # Phase 2: Build agent bundles (if configured)
+        agent_builds = self._build_agents()
+
+        # Phase 3: Upload flow package
         await self._upload_package(tarball)
 
-        # Phase 3: Deploy
-        await self._deploy_via_api()
+        # Phase 4: Upload agent bundles
+        await self._upload_agents(agent_builds)
+
+        # Phase 5: Create/update Prefect deployment
+        await self._deploy_via_api(agent_builds)
 
         print()
         print("=" * 70)

ai_pipeline_core/utils/remote_deployment.py

@@ -1,12 +1,8 @@
-"""Experimental remote deployment utilities.
-
-EXPERIMENTAL: This module provides utilities for calling remotely deployed Prefect flows.
-Subject to change in future versions.
-"""
+"""@public Remote deployment utilities for calling PipelineDeployment flows via Prefect."""
 
 import inspect
 from functools import wraps
-from typing import Any, Callable, ParamSpec, Type, TypeVar
+from typing import Any, Callable, ParamSpec, TypeVar, cast
 
 from prefect import get_client
 from prefect.client.orchestration import PrefectClient
@@ -15,85 +11,26 @@ from prefect.context import AsyncClientContext
 from prefect.deployments.flow_runs import run_deployment
 from prefect.exceptions import ObjectNotFound
 
-from ai_pipeline_core import DocumentList, FlowDocument
+from ai_pipeline_core.deployment import DeploymentContext, DeploymentResult, PipelineDeployment
+from ai_pipeline_core.flow.options import FlowOptions
 from ai_pipeline_core.settings import settings
 from ai_pipeline_core.tracing import TraceLevel, set_trace_cost, trace
 
-# --------------------------------------------------------------------------- #
-# Utility functions (copied from pipeline.py for consistency)
-# --------------------------------------------------------------------------- #
-
-
-def _callable_name(obj: Any, fallback: str) -> str:
-    """Safely extract callable's name for error messages.
-
-    Args:
-        obj: Any object that might have a __name__ attribute.
-        fallback: Default name if extraction fails.
-
-    Returns:
-        The callable's __name__ if available, fallback otherwise.
-
-    Note:
-        Internal helper that never raises exceptions.
-    """
-    try:
-        n = getattr(obj, "__name__", None)
-        return n if isinstance(n, str) else fallback
-    except Exception:
-        return fallback
+P = ParamSpec("P")
+TOptions = TypeVar("TOptions", bound=FlowOptions)
+TResult = TypeVar("TResult", bound=DeploymentResult)
 
 
 def _is_already_traced(func: Callable[..., Any]) -> bool:
-    """Check if a function has already been wrapped by the trace decorator.
-
-    This checks both for the explicit __is_traced__ marker and walks
-    the __wrapped__ chain to detect nested trace decorations.
-
-    Args:
-        func: Function to check for existing trace decoration.
-
-    Returns:
-        True if the function is already traced, False otherwise.
-    """
-    # Check for explicit marker
-    if hasattr(func, "__is_traced__") and func.__is_traced__:  # type: ignore[attr-defined]
+    """Check if function or its __wrapped__ has __is_traced__ attribute."""
+    if getattr(func, "__is_traced__", False):
         return True
-
-    # Walk the __wrapped__ chain to detect nested traces
-    current = func
-    depth = 0
-    max_depth = 10  # Prevent infinite loops
-
-    while hasattr(current, "__wrapped__") and depth < max_depth:
-        wrapped = current.__wrapped__  # type: ignore[attr-defined]
-        # Check if the wrapped function has the trace marker
-        if hasattr(wrapped, "__is_traced__") and wrapped.__is_traced__:  # type: ignore[attr-defined]
-            return True
-        current = wrapped
-        depth += 1
-
-    return False
-
-
-# --------------------------------------------------------------------------- #
-# Remote deployment execution
-# --------------------------------------------------------------------------- #
+    wrapped = getattr(func, "__wrapped__", None)
+    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.
-
-    Args:
-        deployment_name: Name of the deployment to run.
-        parameters: Parameters to pass to the deployment.
-
-    Returns:
-        Result from the deployment execution.
-
-    Raises:
-        ValueError: If deployment is not found in local or remote Prefect API.
-    """
+    """Run a remote Prefect deployment, trying local client first then remote."""
 
     async def _run(client: PrefectClient, as_subflow: bool) -> Any:
         fr: FlowRun = await run_deployment(
@@ -109,7 +46,7 @@ async def run_remote_deployment(deployment_name: str, parameters: dict[str, Any]
             pass
 
     if not settings.prefect_api_url:
-        raise ValueError(f"{deployment_name} deployment not found, PREFECT_API_URL is not set")
+        raise ValueError(f"{deployment_name} not found, PREFECT_API_URL not set")
 
     async with PrefectClient(
         api=settings.prefect_api_url,
@@ -118,9 +55,10 @@ async def run_remote_deployment(deployment_name: str, parameters: dict[str, Any]
     ) as client:
         try:
             await client.read_deployment_by_name(name=deployment_name)
-            with AsyncClientContext.model_construct(
+            ctx = AsyncClientContext.model_construct(
                 client=client, _httpx_settings=None, _context_stack=0
-            ):
+            )
+            with ctx:
                 return await _run(client, False)
         except ObjectNotFound:
             pass
@@ -128,142 +66,54 @@ async def run_remote_deployment(deployment_name: str, parameters: dict[str, Any]
     raise ValueError(f"{deployment_name} deployment not found")
 
 
-P = ParamSpec("P")
-T = TypeVar("T")
-
-
 def remote_deployment(
-    output_document_type: Type[FlowDocument],
+    deployment_class: type[PipelineDeployment[TOptions, TResult]],
     *,
-    # tracing
+    deployment_name: str | None = None,
     name: str | None = None,
     trace_level: TraceLevel = "always",
-    trace_ignore_input: bool = False,
-    trace_ignore_output: bool = False,
-    trace_ignore_inputs: list[str] | None = None,
-    trace_input_formatter: Callable[..., str] | None = None,
-    trace_output_formatter: Callable[..., str] | None = None,
     trace_cost: float | None = None,
-    trace_trim_documents: bool = True,
-) -> Callable[[Callable[P, T]], Callable[P, T]]:
-    """Decorator for calling remote Prefect deployments with automatic tracing.
-
-    EXPERIMENTAL: Decorator for calling remote Prefect deployments with automatic
-    parameter serialization, result deserialization, and LMNR tracing.
-
-    IMPORTANT: Never combine with @trace decorator - this includes tracing automatically.
-    The framework will raise TypeError if you try to use both decorators together.
-
-    Best Practice - Use Defaults:
-        For most use cases, only specify output_document_type. The defaults provide
-        automatic tracing with optimal settings.
-
-    Args:
-        output_document_type: The FlowDocument type to deserialize results into.
-        name: Custom trace name (defaults to function name).
-        trace_level: When to trace ("always", "debug", "off").
-                    - "always": Always trace (default)
-                    - "debug": Only trace when LMNR_DEBUG="true"
-                    - "off": Disable tracing
-        trace_ignore_input: Don't trace input arguments.
-        trace_ignore_output: Don't trace return value.
-        trace_ignore_inputs: List of parameter names to exclude from tracing.
-        trace_input_formatter: Custom formatter for input tracing.
-        trace_output_formatter: Custom formatter for output tracing.
-        trace_cost: Optional cost value to track in metadata. When provided and > 0,
-             sets gen_ai.usage.output_cost, gen_ai.usage.cost, and cost metadata.
-        trace_trim_documents: Trim document content in traces to first 100 chars (default True).
-                             Reduces trace size with large documents.
+) -> Callable[[Callable[P, TResult]], Callable[P, TResult]]:
+    """@public Decorator to call PipelineDeployment flows remotely with automatic serialization."""
 
-    Returns:
-        Decorator function that wraps the target function.
+    def decorator(func: Callable[P, TResult]) -> Callable[P, TResult]:
+        fname = getattr(func, "__name__", deployment_class.name)
 
-    Example:
-        >>> # RECOMMENDED - Minimal usage
-        >>> @remote_deployment(output_document_type=OutputDoc)
-        >>> async def process_remotely(
-        ...     project_name: str,
-        ...     documents: DocumentList,
-        ...     flow_options: FlowOptions
-        >>> ) -> DocumentList:
-        ...     pass  # This stub is replaced by remote call
-        >>>
-        >>> # With custom tracing
-        >>> @remote_deployment(
-        ...     output_document_type=OutputDoc,
-        ...     trace_cost=0.05,  # Track cost of remote execution
-        ...     trace_level="debug"  # Only trace in debug mode
-        >>> )
-        >>> async def debug_remote_flow(...) -> DocumentList:
-        ...     pass
-
-    Note:
-        - Remote calls are automatically traced with LMNR
-        - The decorated function's body is never executed - it serves as a signature template
-        - Deployment name is auto-derived from function name
-        - DocumentList parameters are automatically serialized/deserialized
-
-    Raises:
-        TypeError: If function is already decorated with @trace.
-        ValueError: If deployment is not found.
-    """
-
-    def decorator(func: Callable[P, T]) -> Callable[P, T]:
-        fname = _callable_name(func, "remote_deployment")
-
-        # Check if function is already traced
         if _is_already_traced(func):
-            raise TypeError(
-                f"@remote_deployment target '{fname}' is already decorated "
-                f"with @trace. Remove the @trace decorator - @remote_deployment includes "
-                f"tracing automatically."
-            )
+            raise TypeError(f"@remote_deployment target '{fname}' already has @trace")
 
         @wraps(func)
-        async def _wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        async def _wrapper(*args: P.args, **kwargs: P.kwargs) -> TResult:
             sig = inspect.signature(func)
             bound = sig.bind(*args, **kwargs)
             bound.apply_defaults()
 
-            # Serialize parameters, converting DocumentList to list[dict]
-            parameters = {}
+            # Pass parameters with proper types - Prefect handles Pydantic serialization
+            parameters: dict[str, Any] = {}
             for pname, value in bound.arguments.items():
-                if isinstance(value, DocumentList):
-                    parameters[pname] = [doc for doc in value]
+                if value is None and pname == "context":
+                    parameters[pname] = DeploymentContext()
                 else:
                     parameters[pname] = value
 
-            # Auto-derive deployment name
-            deployment_name = f"{func.__name__.replace('_', '-')}/{func.__name__}"
+            full_name = f"{deployment_class.name}/{deployment_name or deployment_class.name}"
 
-            result = await run_remote_deployment(
-                deployment_name=deployment_name, parameters=parameters
-            )
+            result = await run_remote_deployment(full_name, parameters)
 
-            # Set trace cost if provided
             if trace_cost is not None and trace_cost > 0:
                 set_trace_cost(trace_cost)
 
-            assert isinstance(result, list), "Result must be a list"
-
-            # Auto-handle return type conversion from list[dict] to DocumentList
-            return_type = sig.return_annotation
-
-            assert return_type is DocumentList, "Return type must be a DocumentList"
-            return DocumentList([output_document_type(**item) for item in result])  # type: ignore
+            if isinstance(result, DeploymentResult):
+                return cast(TResult, result)
+            if isinstance(result, dict):
+                return cast(TResult, deployment_class.result_type(**result))
+            raise TypeError(f"Expected DeploymentResult, got {type(result).__name__}")
 
-        # Apply trace decorator
         traced_wrapper = trace(
             level=trace_level,
-            name=name or fname,
-            ignore_input=trace_ignore_input,
-            ignore_output=trace_ignore_output,
-            ignore_inputs=trace_ignore_inputs,
-            input_formatter=trace_input_formatter,
-            output_formatter=trace_output_formatter,
-            trim_documents=trace_trim_documents,
+            name=name or deployment_class.name,
         )(_wrapper)
 
-        return traced_wrapper  # type: ignore
+        return traced_wrapper  # type: ignore[return-value]
 
     return decorator