our-claude-skills/custom-skills/90-reference-curator/07-pipeline-orchestrator/code/scripts/pipeline.py

#!/usr/bin/env python3
"""Pipeline Orchestrator CLI — manage pipeline runs and state.

This script provides state management for the 6-stage pipeline.
Claude orchestrates the actual stages via slash commands.

Usage:
    python pipeline.py init --input "prompt engineering" --type topic [--options '{"max_sources": 10}']
    python pipeline.py advance --run-id 1 --stage crawling [--stats '{"pages_crawled": 45}']
    python pipeline.py pause --run-id 1 --error "Crawl timeout" --stage crawling
    python pipeline.py resume --run-id 1
    python pipeline.py complete --run-id 1 [--export-path ~/reference-library/exports/]
    python pipeline.py status [--run-id 1]
    python pipeline.py track-iteration --run-id 1 --doc-id 42 --action refactor
"""

import json
import sys
from datetime import datetime
from pathlib import Path

import click
from rich.console import Console
from rich.table import Table

from refcurator.db import db_session
from refcurator.config import get_pipeline_config

console = Console()

# QA loop limits
MAX_REFACTOR = 3
MAX_DEEP_RESEARCH = 2
MAX_TOTAL = 5


@click.group()
def cli():
    """Pipeline Orchestrator — manage pipeline run state."""
    pass


@cli.command()
@click.option("--input", "input_value", required=True, help="Topic, URL(s), or manifest path")
@click.option("--type", "run_type", required=True,
              type=click.Choice(["topic", "urls", "manifest"]))
@click.option("--options", type=str, help="Pipeline options as JSON string")
def init(input_value, run_type, options):
    """Initialize a new pipeline run."""
    opts = json.loads(options) if options else {}

    # Merge with config defaults
    try:
        cfg = get_pipeline_config()
        defaults = cfg.get("pipeline", {})
        for key in ("max_sources", "max_pages", "auto_approve", "approval_threshold"):
            if key not in opts and key in defaults:
                opts[key] = defaults[key]
    except FileNotFoundError:
        pass

    stats = {
        "sources_discovered": 0,
        "pages_crawled": 0,
        "documents_stored": 0,
        "documents_distilled": 0,
        "approved": 0,
        "refactored": 0,
        "deep_researched": 0,
        "rejected": 0,
        "needs_manual_review": 0,
    }

    # Determine starting stage
    start_stage = "discovery" if run_type == "topic" else "crawling"

    with db_session() as db:
        run_id = db.insert_returning_id(
            """INSERT INTO pipeline_runs
               (run_type, input_value, status, current_stage, options, stats, started_at)
               VALUES (%s, %s, %s, %s, %s, %s, %s)""",
            (run_type, input_value, "running", start_stage,
             json.dumps(opts), json.dumps(stats), datetime.now().isoformat()),
        )

    console.print(f"[green]Pipeline initialized:[/green] run_id={run_id}")
    console.print(f"  Type: {run_type}")
    console.print(f"  Input: {input_value}")
    console.print(f"  Starting stage: {start_stage}")
    click.echo(json.dumps({"run_id": run_id, "status": "running", "stage": start_stage}))


@cli.command()
@click.option("--run-id", required=True, type=int, help="Pipeline run ID")
@click.option("--stage", required=True,
              type=click.Choice(["discovery", "crawling", "storing", "evaluating",
                                 "distilling", "exporting"]))
@click.option("--stats", type=str, help="Stats update as JSON string (merged with existing)")
def advance(run_id, stage, stats):
    """Advance pipeline to the next stage."""
    stats_update = json.loads(stats) if stats else {}

    with db_session() as db:
        run = db.fetch_one(
            "SELECT * FROM pipeline_runs WHERE run_id = %s", (run_id,)
        )
        if not run:
            console.print(f"[red]Error:[/red] run_id={run_id} not found")
            sys.exit(1)

        if run["status"] != "running":
            console.print(f"[red]Error:[/red] Pipeline is {run['status']}, not running")
            sys.exit(1)

        # Merge stats
        current_stats = json.loads(run["stats"]) if isinstance(run["stats"], str) else (run["stats"] or {})
        current_stats.update(stats_update)

        db.execute(
            """UPDATE pipeline_runs
               SET current_stage = %s, stats = %s
               WHERE run_id = %s""",
            (stage, json.dumps(current_stats), run_id),
        )

    stage_num = ["discovery", "crawling", "storing", "evaluating", "distilling", "exporting"].index(stage) + 1
    console.print(f"[green]Pipeline advanced:[/green] Stage {stage_num}/6 — {stage}")
    click.echo(json.dumps({"run_id": run_id, "stage": stage, "stats": current_stats}))


@cli.command()
@click.option("--run-id", required=True, type=int, help="Pipeline run ID")
@click.option("--error", required=True, help="Error message")
@click.option("--stage", required=True, help="Stage where error occurred")
def pause(run_id, error, stage):
    """Pause pipeline due to error."""
    with db_session() as db:
        db.execute(
            """UPDATE pipeline_runs
               SET status = %s, error_message = %s, error_stage = %s
               WHERE run_id = %s""",
            ("paused", error, stage, run_id),
        )

    console.print(f"[yellow]Pipeline paused:[/yellow] run_id={run_id} at {stage}")
    console.print(f"  Error: {error}")
    console.print(f"  Resume with: python pipeline.py resume --run-id {run_id}")


@cli.command()
@click.option("--run-id", required=True, type=int, help="Pipeline run ID to resume")
def resume(run_id):
    """Resume a paused pipeline."""
    with db_session() as db:
        run = db.fetch_one(
            "SELECT * FROM pipeline_runs WHERE run_id = %s", (run_id,)
        )
        if not run:
            console.print(f"[red]Error:[/red] run_id={run_id} not found")
            sys.exit(1)

        if run["status"] != "paused":
            console.print(f"[red]Error:[/red] Pipeline is {run['status']}, not paused")
            sys.exit(1)

        db.execute(
            """UPDATE pipeline_runs
               SET status = %s, error_message = %s
               WHERE run_id = %s""",
            ("running", None, run_id),
        )

    stage = run["current_stage"]
    console.print(f"[green]Pipeline resumed:[/green] run_id={run_id}, continuing from {stage}")
    click.echo(json.dumps({"run_id": run_id, "status": "running", "resume_stage": stage}))


@cli.command()
@click.option("--run-id", required=True, type=int, help="Pipeline run ID")
@click.option("--export-path", type=click.Path(), help="Export output path")
@click.option("--export-count", type=int, help="Number of documents exported")
def complete(run_id, export_path, export_count):
    """Mark pipeline as completed."""
    with db_session() as db:
        db.execute(
            """UPDATE pipeline_runs
               SET status = %s, completed_at = %s, export_path = %s, export_document_count = %s
               WHERE run_id = %s""",
            ("completed", datetime.now().isoformat(), export_path, export_count, run_id),
        )

        run = db.fetch_one("SELECT * FROM pipeline_runs WHERE run_id = %s", (run_id,))

    stats = json.loads(run["stats"]) if isinstance(run["stats"], str) else (run["stats"] or {})

    console.print(f"\n[bold green]Pipeline Complete![/bold green] run_id={run_id}")
    console.print(f"  Sources: {stats.get('sources_discovered', 0)}")
    console.print(f"  Crawled: {stats.get('pages_crawled', 0)}")
    console.print(f"  Stored: {stats.get('documents_stored', 0)}")
    console.print(f"  Approved: {stats.get('approved', 0)}")
    console.print(f"  Rejected: {stats.get('rejected', 0)}")
    if export_path:
        console.print(f"  Exports: {export_path}")

    click.echo(json.dumps({"run_id": run_id, "status": "completed", "stats": stats}, default=str))


@cli.command()
@click.option("--run-id", type=int, help="Show specific run (or latest if omitted)")
@click.option("--all", "show_all", is_flag=True, help="Show all runs")
def status(run_id, show_all):
    """Show pipeline run status."""
    with db_session() as db:
        if run_id:
            rows = db.fetch_all(
                "SELECT * FROM pipeline_runs WHERE run_id = %s", (run_id,)
            )
        elif show_all:
            rows = db.fetch_all(
                "SELECT * FROM pipeline_runs ORDER BY started_at DESC LIMIT %s", (20,)
            )
        else:
            rows = db.fetch_all(
                "SELECT * FROM pipeline_runs ORDER BY started_at DESC LIMIT %s", (1,)
            )

    if not rows:
        console.print("[dim]No pipeline runs found.[/dim]")
        return

    for run in rows:
        stats = json.loads(run["stats"]) if isinstance(run["stats"], str) else (run["stats"] or {})
        status_color = {
            "running": "green", "completed": "blue",
            "failed": "red", "paused": "yellow",
        }.get(run["status"], "white")

        console.print(f"\n[bold]Pipeline Run #{run['run_id']}[/bold]")
        console.print(f"  Status: [{status_color}]{run['status']}[/{status_color}]")
        console.print(f"  Type: {run['run_type']}")
        console.print(f"  Input: {run['input_value']}")
        console.print(f"  Stage: {run['current_stage']}")
        console.print(f"  Started: {run['started_at']}")

        if run.get("error_message"):
            console.print(f"  [red]Error:[/red] {run['error_message']} (at {run.get('error_stage')})")

        if stats:
            table = Table(title="Pipeline Stats", show_header=False)
            table.add_column("Metric")
            table.add_column("Value", justify="right")
            for k, v in stats.items():
                if v:
                    table.add_row(k.replace("_", " ").title(), str(v))
            console.print(table)


@cli.command("track-iteration")
@click.option("--run-id", required=True, type=int, help="Pipeline run ID")
@click.option("--doc-id", required=True, type=int, help="Document ID")
@click.option("--action", required=True, type=click.Choice(["refactor", "deep_research"]),
              help="QA loop action")
def track_iteration(run_id, doc_id, action):
    """Track QA loop iteration for a document.

    Returns the routing decision:
    - 're_distill': proceed with refactor
    - 're_crawl_and_distill': proceed with deep research
    - 'needs_manual_review': max iterations exceeded
    """
    with db_session() as db:
        tracker = db.fetch_one(
            """SELECT * FROM pipeline_iteration_tracker
               WHERE run_id = %s AND doc_id = %s""",
            (run_id, doc_id),
        )

        if tracker:
            refactor_count = tracker.get("refactor_count", 0)
            deep_research_count = tracker.get("deep_research_count", 0)
        else:
            refactor_count = 0
            deep_research_count = 0

        total = refactor_count + deep_research_count

        # Check limits
        if total >= MAX_TOTAL:
            result = "needs_manual_review"
        elif action == "refactor" and refactor_count >= MAX_REFACTOR:
            result = "needs_manual_review"
        elif action == "deep_research" and deep_research_count >= MAX_DEEP_RESEARCH:
            result = "needs_manual_review"
        elif action == "refactor":
            refactor_count += 1
            result = "re_distill"
        else:
            deep_research_count += 1
            result = "re_crawl_and_distill"

        # Upsert tracker
        if tracker:
            db.execute(
                """UPDATE pipeline_iteration_tracker
                   SET refactor_count = %s, deep_research_count = %s,
                       final_decision = %s
                   WHERE run_id = %s AND doc_id = %s""",
                (refactor_count, deep_research_count,
                 "needs_manual_review" if result == "needs_manual_review" else None,
                 run_id, doc_id),
            )
        else:
            db.insert_returning_id(
                """INSERT INTO pipeline_iteration_tracker
                   (run_id, doc_id, refactor_count, deep_research_count, final_decision)
                   VALUES (%s, %s, %s, %s, %s)""",
                (run_id, doc_id, refactor_count, deep_research_count,
                 "needs_manual_review" if result == "needs_manual_review" else None),
            )

    console.print(
        f"[{'yellow' if result == 'needs_manual_review' else 'green'}]"
        f"Iteration tracked:[/] doc_id={doc_id}, {action} "
        f"(refactor: {refactor_count}/{MAX_REFACTOR}, "
        f"research: {deep_research_count}/{MAX_DEEP_RESEARCH}) → {result}"
    )
    click.echo(json.dumps({
        "run_id": run_id,
        "doc_id": doc_id,
        "action": action,
        "result": result,
        "refactor_count": refactor_count,
        "deep_research_count": deep_research_count,
        "total_iterations": refactor_count + deep_research_count,
    }))


if __name__ == "__main__":
    cli()