ai-pipeline-core 0.3.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

ai_pipeline_core/deployment/contract.py

@@ -1,7 +1,5 @@
 """Unified pipeline run response contract.
 
-@public
-
 Single source of truth for the response shape used by both
 webhook push (ai-pipeline-core) and polling pull (unified-middleware).
 """
@@ -16,12 +14,10 @@ from pydantic import BaseModel, ConfigDict, Discriminator
 class _RunBase(BaseModel):
     """Common fields on every run response variant."""
 
-    type: str
     flow_run_id: UUID
     project_name: str
     state: str  # PENDING, RUNNING, COMPLETED, FAILED, CRASHED, CANCELLED
     timestamp: datetime
-    storage_uri: str = ""
 
     model_config = ConfigDict(frozen=True)
 
@@ -29,19 +25,19 @@ class _RunBase(BaseModel):
 class PendingRun(_RunBase):
     """Pipeline queued or running but no progress reported yet."""
 
-    type: Literal["pending"] = "pending"  # pyright: ignore[reportIncompatibleVariableOverride]
+    type: Literal["pending"] = "pending"
 
 
 class ProgressRun(_RunBase):
     """Pipeline running with step-level progress data."""
 
-    type: Literal["progress"] = "progress"  # pyright: ignore[reportIncompatibleVariableOverride]
+    type: Literal["progress"] = "progress"
     step: int
     total_steps: int
     flow_name: str
     status: str  # "started", "completed", "cached"
-    progress: float  # overall 0.0–1.0
-    step_progress: float  # within step 0.0–1.0
+    progress: float  # overall 0.0-1.0
+    step_progress: float  # within step 0.0-1.0
     message: str
 
 
@@ -57,14 +53,14 @@ class DeploymentResultData(BaseModel):
 class CompletedRun(_RunBase):
     """Pipeline finished (Prefect COMPLETED). Check result.success for business outcome."""
 
-    type: Literal["completed"] = "completed"  # pyright: ignore[reportIncompatibleVariableOverride]
+    type: Literal["completed"] = "completed"
     result: DeploymentResultData
 
 
 class FailedRun(_RunBase):
     """Pipeline crashed — execution error, not business logic."""
 
-    type: Literal["failed"] = "failed"  # pyright: ignore[reportIncompatibleVariableOverride]
+    type: Literal["failed"] = "failed"
     error: str
     result: DeploymentResultData | None = None
 

ai_pipeline_core/{utils → deployment}/deploy.py

@@ -13,7 +13,7 @@ Requirements:
 - Local package installed for flow metadata extraction
 
 Usage:
-    python -m ai_pipeline_core.utils.deploy
+    python -m ai_pipeline_core.deployment.deploy
 """
 
 import argparse
@@ -24,17 +24,17 @@ import sys
 import tempfile
 import tomllib
 import traceback
-from datetime import datetime, timezone
+from datetime import UTC, datetime
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any
 
 from prefect.cli.deploy._storage import _PullStepStorage  # type: ignore
 from prefect.client.orchestration import get_client
 from prefect.deployments.runner import RunnerDeployment
 from prefect.flows import load_flow_from_entrypoint
+from prefect_gcp.cloud_storage import GcpCredentials, GcsBucket  # pyright: ignore[reportMissingTypeStubs]
 
 from ai_pipeline_core.settings import settings
-from ai_pipeline_core.storage import Storage
 
 # ============================================================================
 # Deployer Class
@@ -60,11 +60,7 @@ class Deployer:
             Configuration dictionary with project metadata and deployment settings.
         """
         if not settings.prefect_gcs_bucket:
-            self._die(
-                "PREFECT_GCS_BUCKET not configured in settings.\n"
-                "Configure via environment variable or .env file:\n"
-                "  PREFECT_GCS_BUCKET=your-bucket-name"
-            )
+            self._die("PREFECT_GCS_BUCKET not configured in settings.\nConfigure via environment variable or .env file:\n  PREFECT_GCS_BUCKET=your-bucket-name")
 
         pyproject_path = Path("pyproject.toml")
         if not pyproject_path.exists():
@@ -110,7 +106,7 @@ class Deployer:
                 "  PREFECT_API_URL=https://api.prefect.cloud/api/accounts/.../workspaces/..."
             )
 
-    def _run(self, cmd: str, check: bool = True) -> Optional[str]:
+    def _run(self, cmd: str, *, check: bool = True) -> str | None:
         """Execute shell command and return output.
 
         Args:
@@ -120,22 +116,25 @@ class Deployer:
         Returns:
             Command stdout if successful, None if failed and check=False
         """
-        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
+        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, check=False)
 
         if check and result.returncode != 0:
             self._die(f"Command failed: {cmd}\n{result.stderr}")
 
         return result.stdout.strip() if result.returncode == 0 else None
 
-    def _info(self, msg: str):
+    @staticmethod
+    def _info(msg: str):
         """Print info message."""
         print(f"→ {msg}")
 
-    def _success(self, msg: str):
+    @staticmethod
+    def _success(msg: str):
         """Print success message."""
         print(f"✓ {msg}")
 
-    def _die(self, msg: str):
+    @staticmethod
+    def _die(msg: str):
         """Print error and exit."""
         print(f"✗ {msg}", file=sys.stderr)
         sys.exit(1)
@@ -156,11 +155,7 @@ class Deployer:
         # Verify tarball was created
         tarball_path = Path("dist") / self.config["tarball"]
         if not tarball_path.exists():
-            self._die(
-                f"Build artifact not found: {tarball_path}\n"
-                f"Expected tarball name: {self.config['tarball']}\n"
-                f"Check that pyproject.toml version matches."
-            )
+            self._die(f"Build artifact not found: {tarball_path}\nExpected tarball name: {self.config['tarball']}\nCheck that pyproject.toml version matches.")
 
         self._success(f"Built {tarball_path.name} ({tarball_path.stat().st_size // 1024} KB)")
         return tarball_path
@@ -199,6 +194,7 @@ class Deployer:
                 cwd=source_dir,
                 capture_output=True,
                 text=True,
+                check=False,
             )
             if result.returncode != 0:
                 self._die(f"Wheel build failed for {source_dir.name}:\n{result.stderr}")
@@ -214,7 +210,7 @@ class Deployer:
             output.write_bytes(wheels[0].read_bytes())
             return output
 
-    def _build_agents(self) -> dict[str, dict[str, Any]]:
+    def _build_agents(self) -> dict[str, dict[str, Any]]:  # noqa: PLR0914
         """Build agent wheels and manifests for all configured agents.
 
         Returns:
@@ -238,7 +234,7 @@ class Deployer:
         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()
+        cli_agents_dir = Path(cli_agents_source).resolve()  # pyright: ignore[reportArgumentType]
         if not (cli_agents_dir / "pyproject.toml").exists():
             self._die(f"cli-agents source not found at {cli_agents_dir}")
 
@@ -250,10 +246,7 @@ class Deployer:
         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"
-                )
+                self._die(f"Agent '{agent_name}' path not found: {agent_path}\nCheck [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:
@@ -298,9 +291,7 @@ class Deployer:
             # 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")
-                ):
+                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
@@ -312,7 +303,7 @@ class Deployer:
                 "agent_wheel": agent_wheel.name,
                 "cli_agents_wheel": cli_agents_wheel.name,
                 "vendor_packages": vendor_packages,
-                "built_at": datetime.now(timezone.utc).isoformat(),
+                "built_at": datetime.now(UTC).isoformat(),
             }
             manifest_json = json.dumps(manifest, indent=2)
 
@@ -321,6 +312,17 @@ class Deployer:
 
         return builds
 
+    def _create_gcs_bucket(self, bucket_folder: str) -> Any:
+        """Create a GcsBucket instance for uploading files.
+
+        Args:
+            bucket_folder: Folder path within the bucket.
+        """
+        creds = GcpCredentials()
+        if hasattr(settings, "gcs_service_account_file") and settings.gcs_service_account_file:
+            creds = GcpCredentials(service_account_file=Path(settings.gcs_service_account_file))
+        return GcsBucket(bucket=self.config["bucket"], bucket_folder=bucket_folder, gcp_credentials=creds)
+
     async def _upload_agents(self, agent_builds: dict[str, dict[str, Any]]):
         """Upload agent bundles to GCS.
 
@@ -330,50 +332,38 @@ class Deployer:
         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)
+        flow_folder = self.config["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('')}")
+            agent_folder = f"{flow_folder}/agents/{agent_name}"
+            bucket = self._create_gcs_bucket(agent_folder)
+            self._info(f"Uploading agent '{agent_name}' bundle to gs://{self.config['bucket']}/{agent_folder}")
 
             # Upload manifest
-            await agent_storage.write_bytes(
-                "manifest.json",
-                build_info["manifest_json"].encode(),
-            )
+            await bucket.write_path("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())
+                await bucket.write_path(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.
+        """Upload package tarball to Google Cloud Storage.
 
         Args:
             tarball: Path to the tarball to upload
         """
-        # Extract flow_folder from the config folder path
-        # e.g., "flows/ai-document-writer" -> "ai-document-writer"
-        flow_folder = self.config["folder"].split("/", 1)[1] if "/" in self.config["folder"] else ""
-
-        # Initialize storage with gs://bucket-name/flows and set subfolder to flow_folder
-        base_uri = f"gs://{self.config['bucket']}/flows"
-        storage = await Storage.from_uri(base_uri)
-        storage = storage.with_base(flow_folder)
+        flow_folder = self.config["folder"]
+        bucket = self._create_gcs_bucket(flow_folder)
 
-        dest_uri = storage.url_for(tarball.name)
+        dest_uri = f"gs://{self.config['bucket']}/{flow_folder}/{tarball.name}"
         self._info(f"Uploading to {dest_uri}")
 
-        # Read and upload the tarball
-        tarball_bytes = tarball.read_bytes()
-        await storage.write_bytes(tarball.name, tarball_bytes)
+        tarball_bytes = tarball.read_bytes()  # noqa: ASYNC240
+        await bucket.write_path(tarball.name, tarball_bytes)
 
-        self._success(f"Package uploaded to {self.config['folder']}/{tarball.name}")
+        self._success(f"Package uploaded to {flow_folder}/{tarball.name}")
 
     async def _deploy_via_api(self, agent_builds: dict[str, dict[str, Any]] | None = None):
         """Create or update Prefect deployment using RunnerDeployment pattern.
@@ -437,7 +427,7 @@ class Deployer:
 
         # Create RunnerDeployment
         # This is the official Prefect pattern that handles all the complexity
-        self._info(f"Creating deployment for flow '{flow.name}'")
+        self._info(f"Creating deployment for flow '{flow.name}'")  # pyright: ignore[reportPossiblyUnboundVariable]
 
         # Set AGENT_BUNDLES_URI env var if agents were built
         job_variables: dict[str, Any] = {}
@@ -448,14 +438,13 @@ class Deployer:
 
         deployment = RunnerDeployment(
             name=self.config["package"],
-            flow_name=flow.name,
+            flow_name=flow.name,  # pyright: ignore[reportPossiblyUnboundVariable]
             entrypoint=entrypoint,
             work_pool_name=self.config["work_pool"],
             work_queue_name=self.config["work_queue"],
             tags=[self.config["name"]],
             version=self.config["version"],
-            description=flow.description
-            or f"Deployment for {self.config['package']} v{self.config['version']}",
+            description=flow.description or f"Deployment for {self.config['package']} v{self.config['version']}",  # pyright: ignore[reportPossiblyUnboundVariable]
             storage=_PullStepStorage(pull_steps),
             parameters={},
             job_variables=job_variables,
@@ -466,14 +455,9 @@ class Deployer:
         async with get_client() as client:
             try:
                 work_pool = await client.read_work_pool(self.config["work_pool"])
-                self._success(
-                    f"Work pool '{self.config['work_pool']}' verified (type: {work_pool.type})"
-                )
+                self._success(f"Work pool '{self.config['work_pool']}' verified (type: {work_pool.type})")
             except Exception as e:
-                self._die(
-                    f"Work pool '{self.config['work_pool']}' not accessible: {e}\n"
-                    "Create it in the Prefect UI or with: prefect work-pool create"
-                )
+                self._die(f"Work pool '{self.config['work_pool']}' not accessible: {e}\nCreate it in the Prefect UI or with: prefect work-pool create")
 
         # Apply deployment
         # This automatically handles create vs update based on whether deployment exists
@@ -486,7 +470,7 @@ class Deployer:
             if self.api_url:
                 ui_url = self.api_url.replace("/api/", "/")
                 print(f"\n🌐 View deployment: {ui_url}/deployments/deployment/{deployment_id}")
-                print(f"🚀 Run now: prefect deployment run '{flow.name}/{self.config['package']}'")
+                print(f"🚀 Run now: prefect deployment run '{flow.name}/{self.config['package']}'")  # pyright: ignore[reportPossiblyUnboundVariable]
         except Exception as e:
             self._die(f"Failed to apply deployment: {e}")
 
@@ -530,9 +514,6 @@ def main():
         description="Deploy Prefect flows to GCP using the official RunnerDeployment pattern",
         formatter_class=argparse.RawDescriptionHelpFormatter,
         epilog="""
-Example:
-  python -m ai_pipeline_core.utils.deploy
-
 Prerequisites:
   - Settings configured with PREFECT_API_URL (and optionally PREFECT_API_KEY)
   - Settings configured with PREFECT_GCS_BUCKET

ai_pipeline_core/deployment/helpers.py

@@ -7,12 +7,16 @@ from typing import Any, Literal, TypedDict
 import httpx
 
 from ai_pipeline_core.deployment.contract import CompletedRun, FailedRun, ProgressRun
-from ai_pipeline_core.documents import Document, DocumentList, FlowDocument
+from ai_pipeline_core.documents import Document
 from ai_pipeline_core.logging import get_pipeline_logger
 
 logger = get_pipeline_logger(__name__)
 
 
+class DownloadedDocument(Document):
+    """Concrete document for downloaded content."""
+
+
 class StatusPayload(TypedDict):
     """Webhook payload for Prefect state transitions (sub-flow level)."""
 
@@ -22,24 +26,22 @@ class StatusPayload(TypedDict):
     step: int
     total_steps: int
     flow_name: str
-    state: str  # RUNNING, COMPLETED, FAILED, CRASHED, CANCELLED
+    state: str
     state_name: str
     timestamp: str
 
 
 def class_name_to_deployment_name(class_name: str) -> str:
-    """Convert PascalCase to kebab-case: ResearchPipeline → research-pipeline."""
+    """Convert PascalCase to kebab-case: ResearchPipeline -> research-pipeline."""
     name = re.sub(r"(?<!^)(?=[A-Z])", "-", class_name)
     return name.lower()
 
 
-def extract_generic_params(cls: type) -> tuple[type | None, type | None]:
-    """Extract TOptions and TResult from PipelineDeployment generic args."""
-    from ai_pipeline_core.deployment.base import PipelineDeployment  # noqa: PLC0415
-
+def extract_generic_params(cls: type, base_class: type) -> tuple[type | None, type | None]:
+    """Extract TOptions and TResult from a generic base class's args."""
     for base in getattr(cls, "__orig_bases__", []):
         origin = getattr(base, "__origin__", None)
-        if origin is PipelineDeployment:
+        if origin is base_class:
             args = getattr(base, "__args__", ())
             if len(args) == 2:
                 return args[0], args[1]
@@ -47,22 +49,19 @@ def extract_generic_params(cls: type) -> tuple[type | None, type | None]:
     return None, None
 
 
-async def download_documents(
-    urls: list[str],
-    document_type: type[FlowDocument],
-) -> DocumentList:
-    """Download documents from URLs and return as DocumentList."""
+async def download_documents(urls: list[str]) -> list[Document]:
+    """Download documents from URLs."""
     documents: list[Document] = []
     async with httpx.AsyncClient(timeout=60, follow_redirects=True) as client:
         for url in urls:
             response = await client.get(url)
             response.raise_for_status()
             filename = url.split("/")[-1].split("?")[0] or "document"
-            documents.append(document_type(name=filename, content=response.content))
-    return DocumentList(documents)
+            documents.append(DownloadedDocument(name=filename, content=response.content))
+    return documents
 
 
-async def upload_documents(documents: DocumentList, url_mapping: dict[str, str]) -> None:
+async def upload_documents(documents: list[Document], url_mapping: dict[str, str]) -> None:
     """Upload documents to their mapped URLs."""
     async with httpx.AsyncClient(timeout=60, follow_redirects=True) as client:
         for doc in documents:
@@ -94,5 +93,5 @@ async def send_webhook(
                 logger.warning(f"Webhook retry {attempt + 1}/{max_retries}: {e}")
                 await asyncio.sleep(retry_delay)
             else:
-                logger.error(f"Webhook failed after {max_retries} attempts: {e}")
+                logger.exception(f"Webhook failed after {max_retries} attempts")
                 raise

ai_pipeline_core/{progress.py → deployment/progress.py}

@@ -1,16 +1,19 @@
-"""@public Intra-flow progress tracking with order-preserving webhook delivery."""
+"""Intra-flow progress tracking with order-preserving webhook delivery."""
 
 import asyncio
+import contextlib
 from collections.abc import Generator
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
-from datetime import datetime, timezone
+from datetime import UTC, datetime
 from uuid import UUID
 
-from ai_pipeline_core.deployment.contract import ProgressRun
 from ai_pipeline_core.logging import get_pipeline_logger
 
+from .contract import ProgressRun
+from .helpers import send_webhook
+
 logger = get_pipeline_logger(__name__)
 
 
@@ -25,9 +28,9 @@ class ProgressContext:
     flow_name: str
     step: int
     total_steps: int
-    weights: tuple[float, ...]
-    completed_weight: float
-    current_flow_weight: float
+    total_minutes: float
+    completed_minutes: float
+    current_flow_minutes: float
     queue: asyncio.Queue[ProgressRun | None]
 
 
@@ -35,16 +38,15 @@ _context: ContextVar[ProgressContext | None] = ContextVar("progress_context", de
 
 
 async def update(fraction: float, message: str = "") -> None:
-    """@public Report intra-flow progress (0.0-1.0). No-op without context."""
+    """Report intra-flow progress (0.0-1.0). No-op without context."""
     ctx = _context.get()
     if ctx is None or not ctx.webhook_url:
         return
 
     fraction = max(0.0, min(1.0, fraction))
 
-    total_weight = sum(ctx.weights)
-    if total_weight > 0:
-        overall = (ctx.completed_weight + ctx.current_flow_weight * fraction) / total_weight
+    if ctx.total_minutes > 0:
+        overall = (ctx.completed_minutes + ctx.current_flow_minutes * fraction) / ctx.total_minutes
     else:
         overall = fraction
     overall = round(max(0.0, min(1.0, overall)), 4)
@@ -53,7 +55,7 @@ async def update(fraction: float, message: str = "") -> None:
         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(timezone.utc),
+        timestamp=datetime.now(UTC),
         step=ctx.step,
         total_steps=ctx.total_steps,
         flow_name=ctx.flow_name,
@@ -73,24 +75,20 @@ async def webhook_worker(
     retry_delay: float = 10.0,
 ) -> None:
     """Process webhooks sequentially with retries, preserving order."""
-    from ai_pipeline_core.deployment.helpers import send_webhook  # noqa: PLC0415
-
     while True:
         payload = await queue.get()
         if payload is None:
             queue.task_done()
             break
 
-        try:
+        with contextlib.suppress(Exception):
             await send_webhook(webhook_url, payload, max_retries, retry_delay)
-        except Exception:
-            pass  # Already logged in send_webhook
 
         queue.task_done()
 
 
 @contextmanager
-def flow_context(
+def flow_context(  # noqa: PLR0917
     webhook_url: str,
     project_name: str,
     run_id: str,
@@ -98,12 +96,13 @@ def flow_context(
     flow_name: str,
     step: int,
     total_steps: int,
-    weights: tuple[float, ...],
-    completed_weight: float,
+    flow_minutes: tuple[float, ...],
+    completed_minutes: float,
     queue: asyncio.Queue[ProgressRun | None],
 ) -> Generator[None, None, None]:
     """Set up progress context for a flow. Framework internal use."""
-    current_flow_weight = weights[step - 1] if step <= len(weights) else 1.0
+    current_flow_minutes = flow_minutes[step - 1] if step <= len(flow_minutes) else 1.0
+    total_minutes = sum(flow_minutes) if flow_minutes else current_flow_minutes
     ctx = ProgressContext(
         webhook_url=webhook_url,
         project_name=project_name,
@@ -112,9 +111,9 @@ def flow_context(
         flow_name=flow_name,
         step=step,
         total_steps=total_steps,
-        weights=weights,
-        completed_weight=completed_weight,
-        current_flow_weight=current_flow_weight,
+        total_minutes=total_minutes,
+        completed_minutes=completed_minutes,
+        current_flow_minutes=current_flow_minutes,
         queue=queue,
     )
     token = _context.set(ctx)
@@ -124,4 +123,4 @@ def flow_context(
         _context.reset(token)
 
 
-__all__ = ["update", "webhook_worker", "flow_context", "ProgressContext"]
+__all__ = ["ProgressContext", "flow_context", "update", "webhook_worker"]

ai_pipeline_core/{utils/remote_deployment.py → deployment/remote.py}

@@ -1,8 +1,9 @@
-"""@public Remote deployment utilities for calling PipelineDeployment flows via Prefect."""
+"""Remote deployment utilities for calling PipelineDeployment flows via Prefect."""
 
 import inspect
+from collections.abc import Callable
 from functools import wraps
-from typing import Any, Callable, ParamSpec, TypeVar, cast
+from typing import Any, ParamSpec, TypeVar, cast
 
 from prefect import get_client
 from prefect.client.orchestration import PrefectClient
@@ -12,9 +13,9 @@ 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.flow.options import FlowOptions
+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
-from ai_pipeline_core.tracing import TraceLevel, set_trace_cost, trace
 
 P = ParamSpec("P")
 TOptions = TypeVar("TOptions", bound=FlowOptions)
@@ -33,15 +34,13 @@ async def run_remote_deployment(deployment_name: str, parameters: dict[str, Any]
     """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(
-            client=client, name=deployment_name, parameters=parameters, as_subflow=as_subflow
-        )  # type: ignore
+        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
 
     async with get_client() as client:
         try:
             await client.read_deployment_by_name(name=deployment_name)
-            return await _run(client, True)
+            return await _run(client, True)  # noqa: FBT003
         except ObjectNotFound:
             pass
 
@@ -55,11 +54,9 @@ async def run_remote_deployment(deployment_name: str, parameters: dict[str, Any]
     ) as client:
         try:
             await client.read_deployment_by_name(name=deployment_name)
-            ctx = AsyncClientContext.model_construct(
-                client=client, _httpx_settings=None, _context_stack=0
-            )
+            ctx = AsyncClientContext.model_construct(client=client, _httpx_settings=None, _context_stack=0)
             with ctx:
-                return await _run(client, False)
+                return await _run(client, False)  # noqa: FBT003
         except ObjectNotFound:
             pass
 
@@ -74,7 +71,7 @@ def remote_deployment(
     trace_level: TraceLevel = "always",
     trace_cost: float | None = None,
 ) -> Callable[[Callable[P, TResult]], Callable[P, TResult]]:
-    """@public Decorator to call PipelineDeployment flows remotely with automatic serialization."""
+    """Decorator to call PipelineDeployment flows remotely with automatic serialization."""
 
     def decorator(func: Callable[P, TResult]) -> Callable[P, TResult]:
         fname = getattr(func, "__name__", deployment_class.name)
@@ -106,7 +103,7 @@ def remote_deployment(
             if isinstance(result, DeploymentResult):
                 return cast(TResult, result)
             if isinstance(result, dict):
-                return cast(TResult, deployment_class.result_type(**result))
+                return cast(TResult, deployment_class.result_type(**cast(dict[str, Any], result)))
             raise TypeError(f"Expected DeploymentResult, got {type(result).__name__}")
 
         traced_wrapper = trace(