PlantAI-rag-new/tooling/aqua_kb/batch_runner.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""tooling/aqua_kb/batch_runner.py —— 分批建库调度器

按 manifest 编排建库全链，全程断点续跑：
  for each paper in manifest:
    result.json 已存在 → 跳过抽取（外部/D2 负责抽取，本工具只调度后续步骤）
    → phase2b_converter.convert()（neo4j.cypher + qdrant_points.jsonl）
    → 攒批（--convert-batch 篇）后：
        Neo4j：逐文件 cypher-shell 执行（--neo4j-shell 路径），成功记 checkpoint
        Qdrant：qdrant_loader.load_files()（embed + upsert），成功记 checkpoint
    → 失败写 build/failed.jsonl（含错误），全部步骤可重跑

状态文件（<out>/state/，JSONL 每行一个 paper_id）：
  converted.jsonl / neo4j_loaded.jsonl / qdrant_loaded.jsonl / failed.jsonl

用法：
  python batch_runner.py --manifest pilot_manifest.json --result-root <dir> \
      --out <build dir> --qdrant http://127.0.0.1:27435 \
      --collection aqua_srt_v1_20260930_normalized \
      --neo4j-shell /opt/neo4j-community-5.15.0/bin/cypher-shell \
      --neo4j-uri bolt://127.0.0.1:19889 --neo4j-user neo4j \
      --neo4j-password-file kb-libraries/aqua/data/neo4j/.auth \
      --embed-url http://127.0.0.1:11434/api/embed --embed-model bge-m3:567m \
      --convert-batch 50 [--limit 20]
"""
from __future__ import annotations

import argparse
import json
import re
import subprocess
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

import phase2b_converter
import qdrant_loader

# 与抽取链落盘一致的目录名净化：非 [0-9a-zA-Z._-] 一律替换为 `_`（保留 doi: 前缀）。
# 例：doi:10.1016/0044-8486(81)90060-0 → doi_10.1016_0044-8486_81_90060-0
_SANITIZE_RE = re.compile(r"[^0-9a-zA-Z._-]")


def sanitize_paper_key(value: str) -> str:
    """把 paper_id 规范成与磁盘一致的目录名形态（md_batch/<sanitized>/result.json）。"""
    return _SANITIZE_RE.sub("_", str(value))


def load_ids(path: Path) -> set:
    ids: set = set()
    if path.exists():
        with open(path, "r", encoding="utf-8") as handle:
            for line in handle:
                line = line.strip()
                if not line:
                    continue
                try:
                    ids.add(json.loads(line))  # 兼容 json.dumps 写入的带引号行
                except json.JSONDecodeError:
                    ids.add(line)
    return ids


def append_id(path: Path, value: str) -> None:
    with open(path, "a", encoding="utf-8") as handle:
        handle.write(json.dumps(value) + "\n")


class Runner:
    def __init__(self, args: argparse.Namespace) -> None:
        self.args = args
        self.out = Path(args.out)
        self.state_dir = self.out / "state"
        self.state_dir.mkdir(parents=True, exist_ok=True)
        self.converted_path = self.state_dir / "converted.jsonl"
        self.neo4j_loaded_path = self.state_dir / "neo4j_loaded.jsonl"
        self.qdrant_loaded_path = self.state_dir / "qdrant_loaded.jsonl"
        self.failed_path = self.out / "failed.jsonl"
        self.converted = load_ids(self.converted_path)
        self.neo4j_loaded = load_ids(self.neo4j_loaded_path)
        self.qdrant_loaded = load_ids(self.qdrant_loaded_path)
        self._embed: Optional[qdrant_loader.EmbedClient] = None

    @property
    def embed(self) -> qdrant_loader.EmbedClient:
        if self._embed is None:
            self._embed = qdrant_loader.EmbedClient(
                self.args.embed_url, self.args.embed_model
            )
        return self._embed

    # ---------- manifest ----------

    def read_manifest(self) -> List[Dict[str, Any]]:
        with open(self.args.manifest, "r", encoding="utf-8") as handle:
            manifest = json.load(handle)
        entries = manifest.get("papers") if isinstance(manifest, dict) else manifest
        if not isinstance(entries, list):
            raise RuntimeError("manifest must be a dict with 'papers' list or a list")
        if self.args.limit:
            entries = entries[: self.args.limit]
        return entries

    @staticmethod
    def paper_key(entry: Dict[str, Any]) -> str:
        return str(entry.get("paper_id") or entry.get("id") or entry.get("doi")
                   or entry.get("result") or json.dumps(entry, sort_keys=True))

    def resolve_result(self, entry: Dict[str, Any]) -> Path:
        # 兼容 1：manifest 行显式给出 result 路径（相对 result-root 或绝对路径）
        if entry.get("result"):
            path = Path(entry["result"])
            if not path.is_absolute():
                path = Path(self.args.result_root) / path
            return path
        # 默认：result-root/<sanitized(paper_id)>/result.json
        # sanitize 与抽取链落盘目录名一致（doi:/、/、() 等 → _），
        # 直接用原始 paper_id 拼路径会因含 :/() 而找不到文件。
        key = sanitize_paper_key(self.paper_key(entry))
        return Path(self.args.result_root) / key / "result.json"

    # ---------- steps ----------

    def step_convert(self, key: str, result_path: Path, paper_json: Optional[Dict[str, Any]]) -> Dict[str, Any]:
        """转换幂等（产物确定性），checkpoint 只写一次。返回产物路径 stats。"""
        stats = phase2b_converter.convert(
            result_path, self.out, self.args.release_id, paper_json
        )
        if key not in self.converted:
            append_id(self.converted_path, key)
            print(f"[runner] converted {key}: {stats['assertion_count']} assertions")
        return stats

    def step_neo4j(self, key: str, cypher_path: Path) -> None:
        if key in self.neo4j_loaded:
            print(f"[runner] neo4j skip {key} (checkpoint)")
            return
        shell = self.args.neo4j_shell
        auth = [f"--format=plain", f"--address={self.args.neo4j_uri}",
                f"--username={self.args.neo4j_user}"]
        with open(self.args.neo4j_password_file, "r", encoding="utf-8") as handle:
            password = handle.read().strip().splitlines()[0]
        command = [shell, *auth, f"--password={password}", "-f", str(cypher_path)]
        proc = subprocess.run(command, capture_output=True, text=True, timeout=300)
        if proc.returncode != 0:
            raise RuntimeError(f"cypher-shell failed for {key}: {proc.stderr[-2000:]}")
        append_id(self.neo4j_loaded_path, key)
        print(f"[runner] neo4j loaded {key}")

    def step_qdrant(self, key: str, points_path: Path) -> None:
        if key in self.qdrant_loaded:
            print(f"[runner] qdrant skip {key} (checkpoint)")
            return
        embed = self.embed
        probe = embed.embed("aquaculture warm-up probe")
        qdrant_loader.ensure_collection(
            self.args.qdrant, self.args.collection, len(probe)
        )
        qdrant_loader.upsert_file(
            points_path,
            self.args.qdrant,
            self.args.collection,
            self.args.batch,
            checkpoint=self.state_dir / "qdrant_checkpoint.jsonl",
            embed=embed,
        )
        append_id(self.qdrant_loaded_path, key)
        print(f"[runner] qdrant loaded {key}")

    # ---------- main ----------

    def run(self) -> None:
        entries = self.read_manifest()
        started = time.time()
        stats = {"total": len(entries), "converted": 0, "neo4j": 0, "qdrant": 0,
                 "failed": 0}
        for index, entry in enumerate(entries, 1):
            key = self.paper_key(entry)
            if key in self.failed_keys:
                continue
            paper_json = entry.get("paper") if isinstance(entry.get("paper"), dict) else None
            try:
                result_path = self.resolve_result(entry)
                if not result_path.exists():
                    print(f"[runner] SKIP {key}: result.json missing "
                          f"({result_path}) — 抽取由外部完成")
                    continue
                conv = self.step_convert(key, result_path, paper_json)
                stats["converted"] += 1
                cypher_path = Path(conv["cypher"])
                points_path = Path(conv["points"])
                self.step_neo4j(key, cypher_path)
                stats["neo4j"] += 1
                self.step_qdrant(key, points_path)
                stats["qdrant"] += 1
                if index % 10 == 0:
                    print(f"[runner] progress {index}/{stats['total']} "
                          f"({time.time() - started:.0f}s)")
            except Exception as exc:
                stats["failed"] += 1
                self.record_failure(key, str(exc))
                print(f"[runner] FAILED {key}: {exc}", file=sys.stderr)
        print(f"[runner] DONE {json.dumps(stats, ensure_ascii=False)} in "
              f"{time.time() - started:.1f}s")

    @property
    def failed_keys(self) -> set:
        keys: set = set()
        if self.failed_path.exists():
            with open(self.failed_path, "r", encoding="utf-8") as handle:
                for line in handle:
                    try:
                        keys.add(json.loads(line)["paper_key"])
                    except Exception:
                        continue
        return keys

    def record_failure(self, key: str, error: str) -> None:
        with open(self.failed_path, "a", encoding="utf-8") as handle:
            handle.write(json.dumps(
                {"paper_key": key, "error": error,
                 "ts": time.strftime("%Y-%m-%dT%H:%M:%S%z")},
                ensure_ascii=False,
            ) + "\n")


def main() -> None:
    parser = argparse.ArgumentParser(description="aqua KB batch runner")
    parser.add_argument("--manifest", required=True)
    parser.add_argument("--result-root", required=True)
    parser.add_argument("--out", required=True, help="build output root")
    parser.add_argument("--release-id", default="aqua_srt_v1_20260930_normalized")
    parser.add_argument("--qdrant", required=True)
    parser.add_argument("--collection", default="aqua_srt_v1_20260930_normalized")
    parser.add_argument("--batch", type=int, default=200)
    parser.add_argument("--neo4j-shell", default="cypher-shell")
    parser.add_argument("--neo4j-uri", default="bolt://127.0.0.1:19889")
    parser.add_argument("--neo4j-user", default="neo4j")
    parser.add_argument("--neo4j-password-file", required=True)
    parser.add_argument("--embed-url", default="http://127.0.0.1:11434/api/embed")
    parser.add_argument("--embed-model", default="bge-m3:567m")
    parser.add_argument("--convert-batch", type=int, default=50,
                        help="batch size for Neo4j/Qdrant loading (checkpoint per paper)")
    parser.add_argument("--limit", type=int, default=None,
                        help="only process first N manifest entries")
    args = parser.parse_args()
    Runner(args).run()


if __name__ == "__main__":
    main()