workflow.py

"""Implementation for the 03_excel_visualization Colab tutorial."""

from __future__ import annotations

import csv
import json
import re
import shutil
import subprocess
import sys
from pathlib import Path
from uuid import uuid4

from IPython.display import JSON, FileLink, Image, display
from rath.artifacts import LocalArtifactStore
from rath.context import RunContext
from rath.definition import EffectClass, RetryPolicy, router, step
from rath.flow import Workflow
from rath.runtime import LocalRuntime, RunStatus, SQLiteRunStore
from rath.session import Session

from openrath_colab import (
    checkpoint_rows,
    make_run_directory,
    register_artifact,
    run_to_terminal,
)

FALLBACK_PLAN = {
    "title": "2026 H1 Sales Performance",
    "takeaway": (
        "Revenue increased each month; East leads total revenue while "
        "the overall return rate stays below 5%."
    ),
    "region_chart_title": "Revenue by Region ($)",
    "trend_chart_title": "Monthly Revenue Trend ($)",
}


def validate_plan(payload):
    for key in (
        "title",
        "takeaway",
        "region_chart_title",
        "trend_chart_title",
    ):
        if not isinstance(payload.get(key), str) or not payload[key].strip():
            raise ValueError(f"{key} must be a non-empty string")


class ExcelWorkflow(Workflow):
    def __init__(
        self,
        *,
        llm_client,
        csv_path,
        run_dir,
        output_path,
        artifact_store,
        portable_builder,
    ):
        super().__init__()
        self.llm_client = llm_client
        self.csv_path = csv_path
        self.run_dir = run_dir
        self.output_path = output_path
        self.artifact_store = artifact_store
        self.portable_builder = portable_builder

    @step(
        entry=True,
        successors=("quality_route",),
        effects=EffectClass.READ_ONLY,
    )
    def profile_data(self, state, context):
        with self.csv_path.open(encoding="utf-8", newline="") as handle:
            rows = list(csv.DictReader(handle))
        headers = list(rows[0]) if rows else []
        missing = sum(
            value is None or str(value).strip() == ""
            for row in rows
            for value in row.values()
        )
        identities = [
            (row["Date"], row["Region"], row["Channel"], row["Product"])
            for row in rows
        ]
        duplicates = len(identities) - len(set(identities))
        total_revenue = sum(float(row["Revenue"]) for row in rows)
        total_orders = sum(int(row["Orders"]) for row in rows)
        total_returns = sum(int(row["Returns"]) for row in rows)
        revenue_by_region = {}
        revenue_by_month = {}
        for row in rows:
            revenue = float(row["Revenue"])
            revenue_by_region[row["Region"]] = (
                revenue_by_region.get(row["Region"], 0) + revenue
            )
            revenue_by_month[row["Month"]] = (
                revenue_by_month.get(row["Month"], 0) + revenue
            )
        return {
            **state,
            "profile": {
                "rows": len(rows),
                "columns": headers,
                "missing_cells": missing,
                "duplicate_business_keys": duplicates,
                "regions": sorted({row["Region"] for row in rows}),
                "months": sorted({row["Month"] for row in rows}),
                "total_revenue": total_revenue,
                "total_orders": total_orders,
                "return_rate": total_returns / max(1, total_orders),
                "revenue_by_region": dict(sorted(revenue_by_region.items())),
                "revenue_by_month": dict(sorted(revenue_by_month.items())),
                "quality": (
                    "pass"
                    if rows and missing == 0 and duplicates == 0
                    else "block"
                ),
            },
        }

    @router(successors=("plan_dashboard", "stop_for_quality"))
    def quality_route(self, state):
        return (
            "plan_dashboard"
            if state["profile"]["quality"] == "pass"
            else "stop_for_quality"
        )

    @step(effects=EffectClass.READ_ONLY)
    def stop_for_quality(self, state, context):
        return {**state, "blocked": True}

    @step(
        successors=("build_workbook",),
        effects=EffectClass.NON_IDEMPOTENT,
        timeout_seconds=120,
    )
    def plan_dashboard(self, state, context):
        result = self.llm_client.complete_json(
            system=(
                "你是 Excel dashboard 编辑。只能基于输入 profile 写标题、"
                "一句 takeaway、region_chart_title 和 trend_chart_title。"
                "不得修改数值、发明新指标或输出公式。"
            ),
            user=(
                "请输出 JSON 可视化计划。profile：\n"
                + json.dumps(state["profile"], ensure_ascii=False)
            ),
            fallback=FALLBACK_PLAN,
            validator=validate_plan,
        )
        plan = dict(result.payload)
        profile = state["profile"]
        monthly = profile["revenue_by_month"]
        regional = profile["revenue_by_region"]
        first_month, last_month = list(monthly)[0], list(monthly)[-1]
        leader = max(regional, key=regional.get)
        plan["takeaway"] = (
            f"Revenue rose from ${monthly[first_month]:,.0f} in "
            f"{first_month} to ${monthly[last_month]:,.0f} in {last_month}; "
            f"{leader} leads H1 revenue and the overall return rate is "
            f"{profile['return_rate']:.1%}."
        )
        return {
            **state,
            "dashboard_plan": plan,
            "llm_meta": {
                "mode": result.mode,
                "model": result.model,
                "usage": dict(result.usage),
            },
        }

    @step(
        successors=("qa_and_register",),
        effects=EffectClass.IDEMPOTENT,
        idempotency_key="xlsx-build-v1",
        retry=RetryPolicy(max_attempts=2, base_seconds=0.1),
        timeout_seconds=180,
    )
    def build_workbook(self, state, context):
        build_dir = self.run_dir / "artifact-build"
        build_dir.mkdir(parents=True, exist_ok=True)
        plan_path = build_dir / "dashboard-plan.json"
        plan_path.write_text(
            json.dumps(state["dashboard_plan"], ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        render_dir = self.run_dir / "rendered"
        workspace_xlsx = self.run_dir / "sales-dashboard.xlsx"
        result = subprocess.run(
            [
                sys.executable,
                str(self.portable_builder),
                str(self.csv_path),
                str(plan_path),
                str(render_dir),
                str(workspace_xlsx),
            ],
            check=True,
            capture_output=True,
            text=True,
        )
        return {
            **state,
            "workspace_xlsx_path": str(workspace_xlsx),
            "render_dir": str(render_dir),
            "build_mode": "hosted-colab-public",
            "builder_result": result.stdout.strip().splitlines()[-1],
        }

    @step(
        effects=EffectClass.IDEMPOTENT,
        idempotency_key="xlsx-qa-v1",
    )
    def qa_and_register(self, state, context):
        workspace_xlsx = Path(state["workspace_xlsx_path"])
        render_dir = Path(state["render_dir"])
        qa = json.loads((render_dir / "qa.json").read_text(encoding="utf-8"))
        preview_paths = [Path(item) for item in qa["previews"]]
        if len(preview_paths) != 4 or not all(path.is_file() for path in preview_paths):
            raise RuntimeError("every worksheet must have a visual preview")
        error_pattern = re.compile(r"#REF!|#DIV/0!|#VALUE!|#NAME\?|#N/A")
        error_lines = [
            line
            for line in qa["formula_errors_ndjson"].splitlines()
            if error_pattern.search(line)
            and '"searchTerm"' not in line
        ]
        if error_lines:
            raise RuntimeError(f"formula errors found: {error_lines[:3]}")
        self.output_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(workspace_xlsx, self.output_path)
        artifact = register_artifact(
            self.artifact_store,
            self.output_path,
            tenant_id="local",
            run_id=context.run_id,
            kind="excel-dashboard",
            metadata={"sheets": 4, "rows": state["profile"]["rows"]},
        )
        return {
            **state,
            "xlsx_path": str(self.output_path),
            "qa": {
                "preview_count": len(preview_paths),
                "formula_error_count": len(error_lines),
                "regions": qa["regions"],
                "months": qa["months"],
                "build_mode": state["build_mode"],
            },
            "artifact": {
                "uri": artifact.uri,
                "digest": artifact.digest,
                "size": artifact.size,
            },
        }

    def forward(self, session: Session) -> Session:
        return session

def run_demo(environment):
    llm = environment.llm
    EXAMPLE_DIR = environment.example_dir
    INPUT_DIR = environment.input_dir
    OUTPUT_DIR = environment.output_dir
    PORTABLE_BUILDER = environment.portable_builder

    CSV_PATH = INPUT_DIR / "sales.csv"
    FINAL_OUTPUT = OUTPUT_DIR / "sales-dashboard.xlsx"
    RUN_DIR = make_run_directory(EXAMPLE_DIR, label="excel")
    print(f"run directory: {RUN_DIR}")
    display(
        JSON(
            {
                "input": str(CSV_PATH),
                "output": str(FINAL_OUTPUT),
            }
        )
    )

    store = SQLiteRunStore(RUN_DIR / "runtime.db")
    artifacts = LocalArtifactStore(RUN_DIR / "artifacts")
    workflow = ExcelWorkflow(
        llm_client=llm,
        csv_path=CSV_PATH,
        run_dir=RUN_DIR,
        output_path=FINAL_OUTPUT,
        artifact_store=artifacts,
        portable_builder=PORTABLE_BUILDER,
    )
    revision_id = uuid4()
    runtime = LocalRuntime(store)
    plan = runtime.register(workflow, revision_id=revision_id)
    runtime.submit(
        workflow,
        session_id=uuid4(),
        context=RunContext.local(revision_id=revision_id),
        state={"source": str(CSV_PATH)},
        idempotency_key="excel:synthetic-sales-v1",
    )
    completed = run_to_terminal(runtime, worker_id="excel-worker")
    assert completed.status is RunStatus.SUCCEEDED
    display(
        JSON(
            {
                "profile": dict(completed.state["profile"]),
                "plan_nodes": [
                    {
                        "id": node.id,
                        "kind": node.kind.value,
                        "successors": list(node.successors),
                    }
                    for node in plan.nodes
                ],
                "qa": dict(completed.state["qa"]),
                "artifact": dict(completed.state["artifact"]),
                "checkpoints": checkpoint_rows(store, completed.id),
            }
        )
    )
    display(
        Image(
            filename=str(Path(completed.state["render_dir"]) / "dashboard.png")
        )
    )

    xlsx_path = Path(completed.state["xlsx_path"])
    assert completed.state["profile"]["quality"] == "pass"
    assert completed.state["profile"]["rows"] == 24
    assert completed.state["qa"]["preview_count"] == 4
    assert completed.state["qa"]["formula_error_count"] == 0
    assert completed.state["artifact"]["uri"].startswith("artifact://local/")
    assert xlsx_path.is_file() and xlsx_path.stat().st_size > 10_000
    assert xlsx_path.resolve() == FINAL_OUTPUT.resolve()
    assert len(store.list_checkpoints(completed.id)) == 5
    print("检查通过：质量路由、公式审计、全 Sheet 预览和 artifact 注册均有效。")
    print(f"Excel workbook: {xlsx_path}")
    display(FileLink(str(xlsx_path), result_html_prefix="下载 Excel："))

    return xlsx_path