demand-agent-new/examples/demand/run_existing_execution_local.py

"""Local JSON entrypoint for an existing PG Pattern V2 execution.

This entrypoint intentionally requires an existing PG pattern_mining_execution.id.
It redirects result/.trace/output into one local run directory and never calls
prepare/run_mining.
"""
import argparse
import asyncio
import json
import os
import sys
from datetime import datetime
from pathlib import Path
from zoneinfo import ZoneInfo


sys.path.insert(0, str(Path(__file__).parent.parent.parent))


def _load_project_env() -> None:
    try:
        from dotenv import load_dotenv
    except ImportError:
        return
    load_dotenv(Path(__file__).parent.parent.parent / ".env")


def _default_output_root(execution_id: int) -> Path:
    ts = datetime.now(ZoneInfo("Asia/Shanghai")).strftime("%Y%m%d_%H%M%S")
    return Path(__file__).parent / "test_output_data" / f"execution_{execution_id}_{ts}"


def _test_output_base_dir() -> Path:
    return Path(__file__).parent / "test_output_data"


def _resolve_output_root(execution_id: int, run_id: str | None, output_root: Path | None) -> Path:
    base_dir = _test_output_base_dir().resolve()
    if run_id:
        return base_dir / run_id
    if output_root is None:
        return _default_output_root(execution_id)

    resolved = output_root.expanduser().resolve()
    if resolved != base_dir and base_dir not in resolved.parents:
        raise ValueError(
            f"local MVP 输出必须位于 {base_dir} 下；请改用 --run-id 或传入该目录下的 --output-root"
        )
    return resolved


def _configure_local_env(execution_id: int, output_root: Path) -> None:
    os.environ["DEMAND_OUTPUT_MODE"] = "local_json"
    os.environ["DEMAND_LOCAL_ENTRYPOINT"] = "run_existing_execution_local"
    os.environ["DEMAND_LOCAL_OUTPUT_ROOT"] = str(output_root)
    os.environ["DEMAND_RESULT_BASE_DIR"] = str(output_root / "intermediate" / "result")
    os.environ["DEMAND_WEIGHT_DATA_DIR"] = str(output_root / "intermediate" / "data")
    os.environ["DEMAND_TRACE_STORE_PATH"] = str(output_root / ".trace")
    os.environ["DEMAND_OUTPUT_BASE_DIR"] = str(output_root / "output")
    os.environ.setdefault("DEMAND_LOCAL_TASK_ID", "1")


def _validate_execution_success(execution_id: int) -> None:
    from examples.demand.db_manager import query_execution_for_evidence

    execution = query_execution_for_evidence(execution_id)
    if not execution:
        raise ValueError(f"execution_id={execution_id} 不存在，local MVP 只允许跑已有 execution")
    if str(execution.get("status") or "").strip().lower() != "success":
        raise ValueError(
            f"execution_id={execution_id} status={execution.get('status')}，不是 success，拒绝生成本地输出"
        )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run DemandAgent against an existing PG Pattern V2 execution_id and write local JSON outputs."
    )
    parser.add_argument("--execution-id", type=int, required=True, help="Existing PG pattern_mining_execution.id")
    parser.add_argument("--merge-level2", required=True, help="Current merge_leve2 / cluster name")
    parser.add_argument("--platform-type", default="piaoquan", help="Platform label for manifest only")
    parser.add_argument("--count", type=int, default=5, help="Approximate demand count requested from the prompt")
    parser.add_argument("--run-id", default=None, help="Run directory name under examples/demand/test_output_data")
    parser.add_argument(
        "--output-root",
        type=Path,
        default=None,
        help="Optional full run directory; must be under examples/demand/test_output_data",
    )
    return parser.parse_args()


async def async_main() -> dict:
    _load_project_env()
    args = parse_args()
    output_root = _resolve_output_root(args.execution_id, args.run_id, args.output_root)
    _configure_local_env(args.execution_id, output_root)
    _validate_execution_success(args.execution_id)

    from examples.demand.run import main

    return await main(
        cluster_name=args.merge_level2,
        platform_type=args.platform_type,
        count=args.count,
        execution_id=args.execution_id,
        task_id=int(os.environ.get("DEMAND_LOCAL_TASK_ID", "1")),
    )


def main() -> None:
    result = asyncio.run(async_main())
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()