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

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""tooling/aqua_kb/registry_builder.py —— paper metadata registry 生成器

聚合各 result.json 的 paper 元数据为 registry JSONL，字段对齐植物库
paper_metadata_registry_v1_20260717/paper_metadata_registry_full.jsonl 首行结构：
  paper_id/doi/pmid/pmcid/title/journal/publication_year/metadata_sources/
  lookup_status/resolution_source/conflict_status/conflicts/identifier_candidates/checked_at

幂等：结果按 paper_id 排序，确定性输出。缺字段留空、conflicts 空数组，
lookup_status 按数据来源标记 resolved / doi_candidate / fallback。

用法：
  python registry_builder.py --results <dir of result.json...> --out <registry.jsonl>
"""
from __future__ import annotations

import argparse
import json
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List

# doi 字面量
_DOI_LITERAL = re.compile(r"10\.\d{4,9}/[^\s\"']+", re.IGNORECASE)


def _text(value: Any) -> str:
    return "" if value is None else str(value).strip()


def find_doi(text: str) -> str:
    if not text:
        return ""
    match = _DOI_LITERAL.search(text.replace("_", "/"))
    if match:
        return match.group(0).rstrip(".,;)]")
    return ""


def extract_paper_meta(result: Dict[str, Any], result_path: Path) -> Dict[str, Any]:
    paper: Dict[str, Any] = {}
    cand = result.get("paper")
    if isinstance(cand, dict):
        paper = dict(cand)
    meta = result.get("metadata") or {}
    if not paper and isinstance(meta.get("paper"), dict):
        paper = dict(meta["paper"])

    def get(*keys: str) -> str:
        for key in keys:
            value = paper.get(key)
            if value not in (None, ""):
                return _text(value)
        return ""

    doi = get("doi", "DOI") or find_doi(
        get("title", "") or json.dumps(paper, ensure_ascii=False)
    ) or find_doi(result_path.stem)
    paper_id = get("paper_id", "id") or (f"doi:{doi}" if doi else result_path.stem)
    if doi and paper_id == result_path.stem:
        paper_id = f"doi:{doi}"
    title = get("title")
    if not title:
        title = get("paper_title")
    pmid = get("pmid", "pubmed_id")
    pmcid = get("pmcid", "PMCID")
    journal = get("journal", "journal_name", "venue")
    year = get("publication_year", "year")

    lookup_status = "resolved" if (doi or pmid or pmcid) else "fallback"
    resolution_source = "doi_candidate" if doi else "paper_id"
    metadata_sources = []
    if paper:
        metadata_sources.append("result_paper")
    if doi:
        metadata_sources.append("doi_candidate")
    if not metadata_sources:
        metadata_sources = ["paper_id"]

    return {
        "paper_id": paper_id,
        "title": title,
        "doi": doi,
        "pmid": pmid,
        "pmcid": pmcid,
        "journal": journal,
        "publication_year": year,
        "metadata_sources": metadata_sources,
        "lookup_status": lookup_status,
        "resolution_source": resolution_source,
        "conflict_status": "clean",
        "conflicts": [],
        "identifier_candidates": [doi] if doi else [],
        "checked_at": datetime.now(timezone.utc).isoformat(),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="aqua paper metadata registry builder")
    parser.add_argument("--results", required=True, help="dir containing result.json files (recursive)")
    parser.add_argument("--out", required=True, help="output registry JSONL path")
    args = parser.parse_args()

    root = Path(args.results)
    files = sorted(root.rglob("result.json")) if root.is_dir() else [root]
    rows: Dict[str, Dict[str, Any]] = {}
    for result_path in files:
        try:
            with open(result_path, "r", encoding="utf-8") as handle:
                result = json.load(handle)
        except Exception as exc:
            print(f"[registry] WARN skip {result_path}: {exc}")
            continue
        meta = extract_paper_meta(result, result_path)
        paper_id = meta["paper_id"]
        if paper_id in rows:
            existing = rows[paper_id]
            # 合并：非空字段优先保留已解析信息；metadata_sources 去重合并
            for key, value in meta.items():
                if key in ("metadata_sources", "identifier_candidates", "conflicts"):
                    continue
                if not existing.get(key) and value:
                    existing[key] = value
            merged_sources: List[str] = list(existing["metadata_sources"])
            for source in meta["metadata_sources"]:
                if source not in merged_sources:
                    merged_sources.append(source)
            existing["metadata_sources"] = merged_sources
            cands: List[str] = list(existing["identifier_candidates"])
            for cand in meta["identifier_candidates"]:
                if cand and cand not in cands:
                    cands.append(cand)
            existing["identifier_candidates"] = cands
            if existing["lookup_status"] == "fallback" and meta["lookup_status"] == "resolved":
                existing["lookup_status"] = "resolved"
                existing["resolution_source"] = meta["resolution_source"]
        else:
            rows[paper_id] = meta

    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with open(out_path, "w", encoding="utf-8") as handle:
        for paper_id in sorted(rows):
            handle.write(json.dumps(rows[paper_id], ensure_ascii=False) + "\n")
    print(f"[registry] wrote {len(rows)} rows -> {out_path}")


if __name__ == "__main__":
    main()