video-image-production-build/run_global_data.py

"""从 production_final.json 运行完整 Global Data 闭环。"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from time import perf_counter
from typing import Any
from uuid import uuid4

from langchain_core.language_models import BaseChatModel

from production_build_agents.global_data.graph import create_global_data_graph
from production_build_agents.global_data_stage_finalization import (
    validate_completed_run,
)
from production_build_agents.observability import GlobalDataObservation
from production_build_agents.run.langgraph_checkpointer import (
    create_sqlite_checkpointer,
)
from production_build_agents.run.artifacts import file_sha256
from production_build_agents.run.lock import acquire_run_lock
from production_build_agents.run.layout import run_directory
from production_build_agents.run.metrics import record_run_invocation
from production_build_agents.run.registry import (
    register_thread,
    validate_resume_identity,
)


TERMINAL_STATUSES = {
    "FAILED",
    "COMPLETED",
}
PROTOCOL_VERSION = "0.3"


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="运行完整 Global Data DAG，直到阶段完成或明确失败"
    )
    parser.add_argument("input", type=Path, help="原始 production_final.json")
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("demo_output"),
        help="所有 Run 的根目录；每个 thread_id 使用独立子目录",
    )
    parser.add_argument("--thread-id", default=None)
    return parser


def _initial_state(
    *,
    thread_id: str,
    input_path: Path,
    input_sha256: str,
    run_dir: Path,
) -> dict[str, Any]:
    return {
        "run_id": thread_id,
        "protocol_version": PROTOCOL_VERSION,
        "input_path": str(input_path),
        "input_sha256": input_sha256,
        "output_dir": str(run_dir),
        "status": "PENDING",
        "phase": "PREPROCESS",
        "replan_scope": None,
        "failure_code": None,
        "error": None,
        "global_data_plan_path": None,
        "plan_history": {},
        "task_records": {},
        "current_task_path": None,
        "current_executor_delivery_path": None,
        "current_validation_report_path": None,
        "current_stage_candidate_path": None,
        "current_stage_validation_report_path": None,
        "global_data_delivery_path": None,
        "planner_calls": 0,
        "executor_calls": 0,
        "validator_calls": 0,
        "stage_validator_calls": 0,
        "replan_count": 0,
        "event_log": [],
    }


def run_global_data(
    input_path: Path,
    *,
    output_root: Path,
    thread_id: str,
    planner_model: BaseChatModel | None = None,
    executor_model: BaseChatModel | None = None,
    validator_model: BaseChatModel | None = None,
    registry_dir: Path | None = None,
) -> dict[str, Any]:
    """创建或恢复一个持久化 Global Data Run；终态重启不再执行。"""

    source = input_path.resolve()
    if not source.is_file():
        raise FileNotFoundError(f"输入文件不存在：{source}")
    input_sha256 = file_sha256(source)
    run_dir = run_directory(output_root.resolve(), thread_id).resolve()
    register_thread(
        registry_dir=(
            registry_dir
            or Path(__file__).resolve().parent / ".run_registry"
        ),
        thread_id=thread_id,
        protocol_version=PROTOCOL_VERSION,
        input_path=source,
        input_sha256=input_sha256,
        run_dir=run_dir,
    )
    config = {
        "configurable": {"thread_id": thread_id},
        "recursion_limit": 256,
    }
    started = perf_counter()
    graph_invoked = False

    with acquire_run_lock(run_dir):
        with create_sqlite_checkpointer(run_dir) as checkpointer:
            graph = create_global_data_graph(
                checkpointer=checkpointer,
                planner_model=planner_model,
                executor_model=executor_model,
                validator_model=validator_model,
            )
            saved = dict(graph.get_state(config).values)
            execution_mode = (
                (
                    "completed_noop"
                    if saved.get("status") == "COMPLETED"
                    else "failed_noop"
                )
                if saved and saved.get("status") in TERMINAL_STATUSES
                else ("resume" if saved else "fresh")
            )
            with GlobalDataObservation(
                project_root=Path(__file__).resolve().parent,
                thread_id=thread_id,
                input_sha256=input_sha256,
                protocol_version=PROTOCOL_VERSION,
                execution_mode=execution_mode,
                graph=graph,
            ) as observation:
                if saved:
                    validate_resume_identity(
                        saved,
                        thread_id=thread_id,
                        protocol_version=PROTOCOL_VERSION,
                        input_path=source,
                        input_sha256=input_sha256,
                        run_dir=run_dir,
                    )
                    if saved.get("status") in TERMINAL_STATUSES:
                        if saved.get("status") == "COMPLETED":
                            validate_completed_run(saved)
                        observation.finish(saved, graph_invoked=False)
                        return saved
                    graph_invoked = True
                    result = dict(graph.invoke(None, config))
                else:
                    graph_invoked = True
                    result = dict(
                        graph.invoke(
                            _initial_state(
                                thread_id=thread_id,
                                input_path=source,
                                input_sha256=input_sha256,
                                run_dir=run_dir,
                            ),
                            config,
                        )
                    )
                observation.finish(result, graph_invoked=graph_invoked)
        if graph_invoked:
            try:
                record_run_invocation(
                    run_dir,
                    run_id=thread_id,
                    duration_ms=max(
                        0,
                        round((perf_counter() - started) * 1000),
                    ),
                    status=str(result.get("status") or "UNKNOWN"),
                )
            except (KeyError, OSError, TypeError, ValueError):
                pass
    return result

def main() -> None:
    args = _build_parser().parse_args()
    thread_id = args.thread_id or f"demo-{uuid4()}"
    result = run_global_data(
        args.input,
        output_root=args.output_dir,
        thread_id=thread_id,
    )
    print(json.dumps(result, ensure_ascii=False, indent=2))
    if result.get("status") != "COMPLETED":
        raise SystemExit(1)


if __name__ == "__main__":
    main()