PlantAI-rag-new/tooling/scripts/docgen_import.py

#!/usr/bin/env python3
"""docgen_import.py —— 批量导入工具（P2，不进编译链，直连 vstore HTTP）。

用途：把本地资料目录（含可选 metadata.csv）批量 ingest 进 8094 vstore，
导入完成后删除源明文目录（--keep-source 可保留）。as_template 默认 true。

用法：
    python tooling/scripts/docgen_import.py --dir /path/to/明文资料 \
        --vstore http://127.0.0.1:8094 [--dry-run] [--keep-source] \
        [--no-rebuild] [--workers 8]

安全纪律（docgen_team_context.md 第 9 条）：
- 明文文件只在内存中读取、经 HTTP 发往 vstore，本机不落任何明文副本
- 日志不打印文件内容；--dry-run 不读文件内容、不删目录
- 删除源明文目录前必须 ingest 成功（单个文件失败默认中止且不删源目录）
"""
from __future__ import annotations

import argparse
import csv
import json
import shutil
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import httpx

SUPPORTED_EXTS = {
    ".md", ".markdown", ".txt", ".rst", ".docx", ".doc",
    ".pdf", ".csv", ".html", ".htm", ".json",
}
CSV_META = "metadata.csv"
ALLOWED_CSV_HEADERS = {"name", "category", "tags", "description", "as_template"}


def fail(msg: str) -> "SystemExit":
    print(f"[导入工具] 错误：{msg}", file=sys.stderr)
    raise SystemExit(2)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="批量导入资料到 docgen vstore")
    parser.add_argument("--dir", required=True, type=Path, help="资料明文目录")
    parser.add_argument("--vstore", default="http://127.0.0.1:8094", help="vstore 基地址")
    parser.add_argument("--dry-run", action="store_true", help="只列文件与元数据，不发起请求、不删除任何文件")
    parser.add_argument("--keep-source", action="store_true", help="导入成功后保留源明文目录")
    parser.add_argument("--workers", type=int, default=8, help="并发 ingest 数（默认 8）")
    parser.add_argument("--no-rebuild", action="store_true", help="跳过导入后的 rebuild-index")
    parser.add_argument("--timeout", type=float, default=300.0, help="单文件 ingest 读超时（秒）")
    args = parser.parse_args()
    if not args.dir.is_dir():
        fail(f"目录不存在：{args.dir}")
    if args.workers < 1 or args.workers > 32:
        fail("--workers 需在 1..32 之间")
    return args


def load_metadata(source_dir: Path) -> dict[str, dict]:
    """读取 metadata.csv（可选）：按 file 列映射（无 file 列则跳过）。"""
    csv_path = source_dir / CSV_META
    if not csv_path.is_file():
        return {}
    meta: dict[str, dict] = {}
    try:
        with csv_path.open(encoding="utf-8-sig", newline="") as fh:
            reader = csv.DictReader(fh)
            if reader.fieldnames is None:
                return {}
            headers = {h.strip() for h in reader.fieldnames}
            if "file" not in headers:
                print(f"[导入工具] 提示：{CSV_META} 无 file 列，忽略")
                return {}
            unknown = headers - {"file"} - ALLOWED_CSV_HEADERS
            if unknown:
                print(f"[导入工具] 提示：{CSV_META} 存在未识别列（忽略）：{sorted(unknown)}")
            for row in reader:
                key = (row.get("file") or "").strip()
                if not key:
                    continue
                item: dict = {}
                if row.get("name"):
                    item["name"] = row["name"].strip()
                if row.get("category"):
                    item["category"] = row["category"].strip()
                if row.get("tags"):
                    item["tags"] = row["tags"].strip()
                if row.get("description"):
                    item["description"] = row["description"].strip()
                if str(row.get("as_template") or "").strip().lower() in ("0", "false", "no"):
                    item["as_template"] = False
                meta[key] = item
    except (OSError, csv.Error) as exc:
        fail(f"{CSV_META} 读取失败：{exc}")
    return meta


def collect_files(source_dir: Path) -> list[Path]:
    """收集支持扩展名的资料文件（跳过 metadata.csv 自身）。"""
    files: list[Path] = []
    for p in sorted(source_dir.rglob("*")):
        if not p.is_file() or p.name == CSV_META:
            continue
        if p.suffix.lower() in SUPPORTED_EXTS:
            files.append(p)
    return files


def ingest_one(vstore: str, source_dir: Path, rel: Path, meta: dict, timeout: float) -> dict:
    """单文件 ingest（在 worker 线程执行，同步 httpx）。"""
    data = (source_dir / rel).read_bytes()  # 绝对路径读取（worker 线程 cwd 不依赖源目录）
    payload = {
        "name": meta.get("name", rel.stem),
        "category": meta.get("category", ""),
        "tags": meta.get("tags", ""),
        "description": meta.get("description", ""),
        "as_template": meta.get("as_template", True),
    }
    with httpx.Client(timeout=httpx.Timeout(connect=10.0, read=timeout, write=60.0, pool=10.0)) as client:
        resp = client.post(
            f"{vstore}/ingest",
            data=payload,
            files={"file": (rel.name, data)},
        )
        if resp.status_code >= 400:
            raise RuntimeError(f"HTTP {resp.status_code}: {resp.text[:200]}")
        try:
            body = resp.json()
        except ValueError:
            body = {"document_id": ""}
        document_id = body.get("document_id") or body.get("id") or ""
        return {"rel": rel, "name": rel.name, "size": len(data), "document_id": document_id, "ok": True}


def rebuild_index_sync(vstore: str) -> dict:
    with httpx.Client(timeout=httpx.Timeout(connect=10.0, read=120.0, write=30.0, pool=10.0)) as client:
        resp = client.post(f"{vstore}/rebuild-index")
        if resp.status_code >= 400:
            raise RuntimeError(f"HTTP {resp.status_code}: {resp.text[:200]}")
        try:
            return resp.json()
        except ValueError:
            return {}


def main() -> int:
    args = parse_args()
    source_dir: Path = args.dir.resolve()
    files = collect_files(source_dir)
    if not files:
        print("[导入工具] 未发现可导入文件（支持扩展名：%s）" % ", ".join(sorted(SUPPORTED_EXTS)))
        return 0

    meta = load_metadata(source_dir)

    print(f"[导入工具] 目录：{source_dir}")
    print(f"[导入工具] vstore：{args.vstore}")
    print(f"[导入工具] 待导入：{len(files)} 个文件，并发 {args.workers}")
    if args.dry_run:
        print("[导入工具] DRY-RUN：仅列出文件与元数据，不发起请求")
        for p in files:
            rel = p.relative_to(source_dir)
            m = meta.get(rel.as_posix(), {})
            print(f"  - {rel.as_posix()}  name={m.get('name', rel.stem)} "
                  f"as_template={m.get('as_template', True)} category={m.get('category', '')}")
        print(f"[导入工具] DRY-RUN 完成（{len(files)} 个文件待导入；源目录未改动）")
        return 0

    started = time.time()
    ok: list[dict] = []
    failed: list[tuple[str, str]] = []
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = {}
        for p in files:
            rel = p.relative_to(source_dir)
            m = meta.get(rel.as_posix(), {})
            futures[pool.submit(ingest_one, args.vstore, source_dir, rel, m, args.timeout)] = rel
        for fut in as_completed(futures):
            rel = futures[fut]
            try:
                ok.append(fut.result())
                print(f"  [ok] {rel.as_posix()}")
            except Exception as exc:
                failed.append((rel.as_posix(), str(exc)))
                print(f"  [失败] {rel.as_posix()}: {exc}", file=sys.stderr)

    elapsed = time.time() - started
    print(f"[导入工具] ingest 完成：成功 {len(ok)}，失败 {len(failed)}，耗时 {elapsed:.1f}s")
    for name, err in failed:
        print(f"  - {name}: {err}", file=sys.stderr)

    if failed:
        fail(f"有 {len(failed)} 个文件 ingest 失败——未重建索引、未删除源目录（可重跑，幂等）")

    if not args.no_rebuild:
        print("[导入工具] rebuild-index（原子切换）…")
        try:
            result = rebuild_index_sync(args.vstore)
            print(f"[导入工具] rebuild-index 完成：{json.dumps(result, ensure_ascii=False)[:200]}")
        except Exception as exc:
            fail(f"rebuild-index 失败：{exc}（未删除源目录，修复后重跑）")

    if args.keep_source:
        print(f"[导入工具] --keep-source：保留源目录 {source_dir}")
    else:
        # 全部成功才删除源明文目录（纪律：不落未加密明文）
        try:
            shutil.rmtree(source_dir)
            print(f"[导入工具] 已删除源明文目录：{source_dir}")
        except OSError as exc:
            print(f"[导入工具] 警告：源目录删除失败（请手动清理）：{exc}", file=sys.stderr)

    print("[导入工具] 完成")
    return 0


if __name__ == "__main__":
    sys.exit(main())