ai-pipeline-core 0.2.6__py3-none-any.whl → 0.4.1__py3-none-any.whl

ai_pipeline_core/deployment/helpers.py

@@ -0,0 +1,97 @@
+"""Helper functions for pipeline deployments."""
+
+import asyncio
+import re
+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
+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)."""
+
+    type: Literal["status"]
+    flow_run_id: str
+    project_name: str
+    step: int
+    total_steps: int
+    flow_name: str
+    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."""
+    name = re.sub(r"(?<!^)(?=[A-Z])", "-", class_name)
+    return name.lower()
+
+
+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 base_class:
+            args = getattr(base, "__args__", ())
+            if len(args) == 2:
+                return args[0], args[1]
+
+    return None, None
+
+
+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(DownloadedDocument(name=filename, content=response.content))
+    return documents
+
+
+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:
+            if doc.name in url_mapping:
+                response = await client.put(
+                    url_mapping[doc.name],
+                    content=doc.content,
+                    headers={"Content-Type": doc.mime_type},
+                )
+                response.raise_for_status()
+
+
+async def send_webhook(
+    url: str,
+    payload: ProgressRun | CompletedRun | FailedRun,
+    max_retries: int = 3,
+    retry_delay: float = 10.0,
+) -> None:
+    """Send webhook with retries."""
+    data: dict[str, Any] = payload.model_dump(mode="json")
+    for attempt in range(max_retries):
+        try:
+            async with httpx.AsyncClient(timeout=30) as client:
+                response = await client.post(url, json=data, follow_redirects=True)
+                response.raise_for_status()
+            return
+        except Exception as e:
+            if attempt < max_retries - 1:
+                logger.warning(f"Webhook retry {attempt + 1}/{max_retries}: {e}")
+                await asyncio.sleep(retry_delay)
+            else:
+                logger.exception(f"Webhook failed after {max_retries} attempts")
+                raise

ai_pipeline_core/deployment/progress.py

@@ -0,0 +1,126 @@
+"""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 UTC, datetime
+from uuid import UUID
+
+from ai_pipeline_core.logging import get_pipeline_logger
+
+from .contract import ProgressRun
+from .helpers import send_webhook
+
+logger = get_pipeline_logger(__name__)
+
+
+@dataclass(frozen=True, slots=True)
+class ProgressContext:
+    """Internal context holding state for progress calculation and webhook delivery."""
+
+    webhook_url: str
+    project_name: str
+    run_id: str
+    flow_run_id: str
+    flow_name: str
+    step: int
+    total_steps: int
+    total_minutes: float
+    completed_minutes: float
+    current_flow_minutes: float
+    queue: asyncio.Queue[ProgressRun | None]
+
+
+_context: ContextVar[ProgressContext | None] = ContextVar("progress_context", default=None)
+
+
+async def update(fraction: float, message: str = "") -> None:
+    """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))
+
+    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)
+
+    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)
+
+
+async def webhook_worker(
+    queue: asyncio.Queue[ProgressRun | None],
+    webhook_url: str,
+    max_retries: int = 3,
+    retry_delay: float = 10.0,
+) -> None:
+    """Process webhooks sequentially with retries, preserving order."""
+    while True:
+        payload = await queue.get()
+        if payload is None:
+            queue.task_done()
+            break
+
+        with contextlib.suppress(Exception):
+            await send_webhook(webhook_url, payload, max_retries, retry_delay)
+
+        queue.task_done()
+
+
+@contextmanager
+def flow_context(  # noqa: PLR0917
+    webhook_url: str,
+    project_name: str,
+    run_id: str,
+    flow_run_id: str,
+    flow_name: str,
+    step: int,
+    total_steps: int,
+    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_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,
+        run_id=run_id,
+        flow_run_id=flow_run_id,
+        flow_name=flow_name,
+        step=step,
+        total_steps=total_steps,
+        total_minutes=total_minutes,
+        completed_minutes=completed_minutes,
+        current_flow_minutes=current_flow_minutes,
+        queue=queue,
+    )
+    token = _context.set(ctx)
+    try:
+        yield
+    finally:
+        _context.reset(token)
+
+
+__all__ = ["ProgressContext", "flow_context", "update", "webhook_worker"]

ai_pipeline_core/deployment/remote.py

@@ -0,0 +1,116 @@
+"""Remote deployment utilities for calling PipelineDeployment flows via Prefect."""
+
+import inspect
+from collections.abc import Callable
+from functools import wraps
+from typing import Any, ParamSpec, TypeVar, cast
+
+from prefect import get_client
+from prefect.client.orchestration import PrefectClient
+from prefect.client.schemas import FlowRun
+from prefect.context import AsyncClientContext
+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.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")
+TOptions = TypeVar("TOptions", bound=FlowOptions)
+TResult = TypeVar("TResult", bound=DeploymentResult)
+
+
+def _is_already_traced(func: Callable[..., Any]) -> bool:
+    """Check if function or its __wrapped__ has __is_traced__ attribute."""
+    if getattr(func, "__is_traced__", False):
+        return True
+    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, 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
+        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)  # noqa: FBT003
+        except ObjectNotFound:
+            pass
+
+    if not settings.prefect_api_url:
+        raise ValueError(f"{deployment_name} not found, PREFECT_API_URL not set")
+
+    async with PrefectClient(
+        api=settings.prefect_api_url,
+        api_key=settings.prefect_api_key,
+        auth_string=settings.prefect_api_auth_string,
+    ) as client:
+        try:
+            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
+        except ObjectNotFound:
+            pass
+
+    raise ValueError(f"{deployment_name} deployment not found")
+
+
+def remote_deployment(
+    deployment_class: type[PipelineDeployment[TOptions, TResult]],
+    *,
+    deployment_name: str | None = None,
+    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."""
+
+    def decorator(func: Callable[P, TResult]) -> Callable[P, 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)
+
+            if trace_cost is not None and trace_cost > 0:
+                set_trace_cost(trace_cost)
+
+            if isinstance(result, DeploymentResult):
+                return cast(TResult, result)
+            if isinstance(result, dict):
+                return cast(TResult, deployment_class.result_type(**cast(dict[str, Any], result)))
+            raise TypeError(f"Expected DeploymentResult, got {type(result).__name__}")
+
+        traced_wrapper = trace(
+            level=trace_level,
+            name=name or deployment_class.name,
+        )(_wrapper)
+
+        return traced_wrapper  # type: ignore[return-value]
+
+    return decorator

ai_pipeline_core/docs_generator/__init__.py

@@ -0,0 +1,54 @@
+"""AI-focused documentation generator.
+
+Generates dense, self-contained guides from source code and test suite
+for AI coding agents. Uses AST parsing, dependency resolution, and
+size management for guides with a 50KB warning threshold.
+"""
+
+from ai_pipeline_core.docs_generator.extractor import (
+    ClassInfo,
+    FunctionInfo,
+    MethodInfo,
+    ModuleInfo,
+    SymbolTable,
+    is_public_name,
+    parse_module,
+)
+from ai_pipeline_core.docs_generator.guide_builder import (
+    GuideData,
+    TestExample,
+    build_guide,
+    discover_tests,
+    select_examples,
+)
+from ai_pipeline_core.docs_generator.trimmer import manage_guide_size
+from ai_pipeline_core.docs_generator.validator import (
+    ValidationResult,
+    compute_source_hash,
+    validate_all,
+    validate_completeness,
+    validate_freshness,
+    validate_size,
+)
+
+__all__ = [
+    "ClassInfo",
+    "FunctionInfo",
+    "GuideData",
+    "MethodInfo",
+    "ModuleInfo",
+    "SymbolTable",
+    "TestExample",
+    "ValidationResult",
+    "build_guide",
+    "compute_source_hash",
+    "discover_tests",
+    "is_public_name",
+    "manage_guide_size",
+    "parse_module",
+    "select_examples",
+    "validate_all",
+    "validate_completeness",
+    "validate_freshness",
+    "validate_size",
+]

ai_pipeline_core/docs_generator/__main__.py

@@ -0,0 +1,5 @@
+"""Entry point for python -m ai_pipeline_core.docs_generator."""
+
+from ai_pipeline_core.docs_generator.cli import main
+
+raise SystemExit(main())

ai_pipeline_core/docs_generator/cli.py

@@ -0,0 +1,196 @@
+"""CLI for AI documentation generation and validation."""
+
+import argparse
+import sys
+from pathlib import Path
+
+from ai_pipeline_core.docs_generator.extractor import build_symbol_table
+from ai_pipeline_core.docs_generator.guide_builder import build_guide, render_guide
+from ai_pipeline_core.docs_generator.trimmer import manage_guide_size
+from ai_pipeline_core.docs_generator.validator import (
+    HASH_FILE,
+    compute_source_hash,
+    validate_all,
+)
+
+EXCLUDED_MODULES: frozenset[str] = frozenset({"docs_generator"})
+
+
+def _normalize_whitespace(content: str) -> str:
+    """Strip trailing whitespace from each line and ensure final newline."""
+    lines = [line.rstrip() for line in content.splitlines()]
+    return "\n".join(lines) + "\n"
+
+
+TEST_DIR_OVERRIDES: dict[str, str] = {}  # nosemgrep: no-mutable-module-globals
+
+
+def _discover_modules(source_dir: Path) -> list[str]:
+    """Discover all public module groupings from package structure."""
+    modules: set[str] = set()
+    for py_file in sorted(source_dir.rglob("*.py")):
+        if py_file.name.startswith("_") and py_file.name != "__init__.py":
+            continue
+        relative = py_file.relative_to(source_dir)
+        if len(relative.parts) > 1:
+            modules.add(relative.parts[0])
+        else:
+            modules.add(relative.stem)
+    return sorted(modules - EXCLUDED_MODULES)
+
+
+def main(argv: list[str] | None = None) -> int:
+    """Entry point for AI docs CLI with generate/check subcommands."""
+    parser = argparse.ArgumentParser(description="AI documentation generator")
+    parser.add_argument("--source-dir", type=Path, help="Source package directory")
+    parser.add_argument("--tests-dir", type=Path, help="Tests directory")
+    parser.add_argument("--output-dir", type=Path, help="Output .ai-docs directory")
+    subparsers = parser.add_subparsers(dest="command")
+    subparsers.add_parser("generate", help="Generate .ai-docs/ documentation")
+    subparsers.add_parser("check", help="Validate .ai-docs/ is up-to-date")
+
+    args = parser.parse_args(argv)
+    if not args.command:
+        parser.print_help()
+        return 1
+
+    source_dir, tests_dir, output_dir, repo_root = _resolve_paths(args)
+
+    if args.command == "generate":
+        return _run_generate(source_dir, tests_dir, output_dir, repo_root)
+    return _run_check(source_dir, tests_dir, output_dir)
+
+
+def _resolve_paths(args: argparse.Namespace) -> tuple[Path, Path, Path, Path]:
+    """Resolve source, tests, output directories and repo root from args or auto-detect."""
+    cli_file = Path(__file__).resolve()
+    repo_root = cli_file.parent.parent.parent
+    source_dir = args.source_dir or (repo_root / "ai_pipeline_core")
+    tests_dir = args.tests_dir or (repo_root / "tests")
+    output_dir = args.output_dir or (repo_root / ".ai-docs")
+    return source_dir, tests_dir, output_dir, repo_root
+
+
+def _run_generate(source_dir: Path, tests_dir: Path, output_dir: Path, repo_root: Path) -> int:
+    """Generate all module guides, INDEX.md, and .hash file."""
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Clean stale files
+    for existing in output_dir.glob("*.md"):
+        existing.unlink()
+    hash_file = output_dir / HASH_FILE
+    if hash_file.exists():
+        hash_file.unlink()
+
+    table = build_symbol_table(source_dir)
+    generated: list[tuple[str, int]] = []
+
+    for module_name in _discover_modules(source_dir):
+        data = build_guide(module_name, source_dir, tests_dir, table, TEST_DIR_OVERRIDES, repo_root)
+        if not data.classes and not data.functions:
+            print(f"  skip {module_name} (no public symbols)")
+            continue
+
+        content = render_guide(data)
+        content = manage_guide_size(data, content)
+        content = _normalize_whitespace(content)
+
+        guide_path = output_dir / f"{module_name}.md"
+        guide_path.write_text(content)
+        size = len(content.encode("utf-8"))
+        generated.append((module_name, size))
+        print(f"  wrote {module_name}.md ({size:,} bytes)")
+
+    # INDEX.md
+    index_content = _normalize_whitespace(_render_index(generated))
+    (output_dir / "INDEX.md").write_text(index_content)
+    print(f"  wrote INDEX.md ({len(index_content):,} bytes)")
+
+    # .hash
+    source_hash = compute_source_hash(source_dir, tests_dir)
+    (output_dir / HASH_FILE).write_text(source_hash + "\n")
+    print(f"  wrote {HASH_FILE}")
+
+    total = sum(size for _, size in generated)
+    print(f"\nGenerated {len(generated)} guides ({total:,} bytes total)")
+    return 0
+
+
+def _run_check(source_dir: Path, tests_dir: Path, output_dir: Path) -> int:
+    """Validate .ai-docs/ freshness, completeness, and size."""
+    if not output_dir.is_dir():
+        print("FAIL: .ai-docs/ directory does not exist. Run 'generate' first.", file=sys.stderr)
+        return 1
+
+    result = validate_all(output_dir, source_dir, tests_dir, excluded_modules=EXCLUDED_MODULES)
+
+    if not result.is_fresh:
+        print("FAIL: .ai-docs/ is stale (source hash mismatch)")
+    if result.missing_symbols:
+        print(f"FAIL: {len(result.missing_symbols)} public symbols missing from guides:")
+        for sym in result.missing_symbols:
+            print(f"  - {sym}")
+    if result.size_violations:
+        print(f"WARNING: {len(result.size_violations)} guides exceed size limit:")
+        for name, size in result.size_violations:
+            print(f"  - {name}: {size:,} bytes")
+
+    if result.is_valid:
+        print("OK: .ai-docs/ is up-to-date")
+        return 0
+    return 1
+
+
+def _render_index(generated: list[tuple[str, int]]) -> str:
+    """Render INDEX.md with reading order, task lookup, imports, and size table."""
+    lines: list[str] = [
+        "# AI Documentation Index",
+        "",
+        "Auto-generated guide index. Do not edit manually.",
+        "",
+        "## Reading Order",
+        "",
+    ]
+    for i, (name, _) in enumerate(generated, 1):
+        lines.append(f"{i}. [{name}]({name}.md)")
+
+    lines.extend([
+        "",
+        "## Task-Based Lookup",
+        "",
+        "| Task | Guide |",
+        "| ---- | ----- |",
+    ])
+    task_map = {
+        "Create/read documents": "documents",
+        "Store/retrieve documents": "document_store",
+        "Call LLMs": "llm",
+        "Deploy pipelines": "deployment",
+        "Load templates": "prompt_manager",
+        "Process images": "images",
+        "Define flows/tasks": "pipeline",
+        "Configure settings": "settings",
+        "Handle errors": "exceptions",
+        "Log messages": "logging",
+        "Debug & observe traces": "observability",
+        "Test pipelines": "testing",
+    }
+    guide_set = {name for name, _ in generated}
+    for task, guide in task_map.items():
+        if guide in guide_set:
+            lines.append(f"| {task} | [{guide}]({guide}.md) |")
+
+    lines.extend([
+        "",
+        "## Module Sizes",
+        "",
+        "| Module | Size |",
+        "| ------ | ---- |",
+    ])
+    for name, size in generated:
+        lines.append(f"| {name} | {size:,} bytes |")
+    total = sum(size for _, size in generated)
+    lines.append(f"| **Total** | **{total:,} bytes** |")
+    lines.append("")
+
+    return "\n".join(lines)