#!/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())