ai-pipeline-core 0.4.3__tar.gz → 0.4.5__tar.gz

{ai_pipeline_core-0.4.3 → ai_pipeline_core-0.4.5}/.gitignore

@@ -160,7 +160,6 @@ ehthumbs.db
 .prefect/
 
 # Test artifacts
-tests/test_data/
 test_output/
 
 # temporary specification

{ai_pipeline_core-0.4.3 → ai_pipeline_core-0.4.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-pipeline-core
-Version: 0.4.3
+Version: 0.4.5
 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 → ai_pipeline_core-0.4.5}/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.5"
 
 __all__ = [
     "AIMessageType",

{ai_pipeline_core-0.4.3 → ai_pipeline_core-0.4.5}/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-0.4.3 → ai_pipeline_core-0.4.5}/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.
 
@@ -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,26 @@ 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] = []
+        # 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: set[str] = set()
+            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 and "cli_agents" in filename:
-                        vendor_wheels.append(filepath)
-                        seen.add(filename)
+                    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)
+
+        # 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-0.4.3 → ai_pipeline_core-0.4.5}/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-0.4.3 → ai_pipeline_core-0.4.5}/ai_pipeline_core/deployment/remote.py

@@ -93,7 +93,7 @@ def remote_deployment(
                 else:
                     parameters[pname] = value
 
-            full_name = f"{deployment_class.name}/{deployment_name or deployment_class.name}"
+            full_name = f"{deployment_class.name}/{deployment_name or deployment_class.name.replace('-', '_')}"
 
             result = await run_remote_deployment(full_name, parameters)
 

{ai_pipeline_core-0.4.3 → ai_pipeline_core-0.4.5}/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 → ai_pipeline_core-0.4.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ai-pipeline-core"
-version = "0.4.3"
+version = "0.4.5"
 description = "Core utilities for AI-powered processing pipelines using prefect"
 readme = "README.md"
 license = {text = "MIT"}
@@ -61,7 +61,7 @@ dev = [
 minversion = "8.0"
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"
-addopts = "-q -n 8 -m 'not integration and not clickhouse'"
+addopts = "-q -m 'not integration and not clickhouse'"
 testpaths = ["tests"]
 markers = [
     "integration: marks tests as integration tests (deselect with '-m \"not integration\"')",