220 lines
9.0 KiB
Python
220 lines
9.0 KiB
Python
|
#!/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())
|