PlantAI-rag-new/vstore/scripts/rebuild_chunks.py

"""rebuild_chunks.py —— 存量资料 chunk 索引重建 CLI（FR-PRIV-007）。

不进 HTTP（避免超时），由 TL 在部署根后台执行：
    cd {部署根}/vstore && ../venv/bin/python scripts/rebuild_chunks.py

路径推导（零硬编码，优先级从高到低）：
- CLI 参数 --root / --vault / --index / --chunks（供测试/特殊部署）
- 环境变量 VSTORE_ROOT / VSTORE_VAULT_PATH / VSTORE_INDEX_PATH / VSTORE_CHUNK_PATH
- 默认：root = 脚本所在目录的父目录（部署根下的 vstore/）

密钥加载：cwd=部署根时加载 vstore/.env（VSTORE_KEY 必须可得），否则用环境变量。

进度打印 + 完成汇总；重建期间 vstore HTTP 服务（若运行中）不阻塞——
本脚本是独立子进程，只经文件锁写 chunks 目录。
"""
from __future__ import annotations

import argparse
import os
import sys
import time
from pathlib import Path


def _load_env_file(path: Path) -> None:
    """vstore/.env → os.environ（已存在的键不覆盖）。"""
    if not path.exists():
        return
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, value = line.partition("=")
        key = key.strip()
        value = value.strip().strip('"').strip("'")
        if key and key not in os.environ:
            os.environ[key] = value


def main() -> int:
    parser = argparse.ArgumentParser(description="重建私有资料 chunk 索引（FR-PRIV-007）")
    parser.add_argument("--root", help="vstore 数据根（默认 VSTORE_ROOT 或脚本父目录）")
    parser.add_argument("--vault", help="密文目录（默认 {root}/vault）")
    parser.add_argument("--index", help="doc 级索引目录（默认 {root}/index）")
    parser.add_argument("--chunks", help="chunk 索引目录（默认 {root}/chunks）")
    parser.add_argument("--batch-size", type=int, default=16, help="批量嵌入每批条数")
    parser.add_argument("--workers", type=int, default=4, help="批量嵌入并发线程数")
    args = parser.parse_args()

    script_dir = Path(__file__).resolve().parent
    vstore_dir = script_dir.parent          # 部署根下的 vstore/
    deploy_root = vstore_dir.parent         # 部署根
    _load_env_file(vstore_dir / ".env")     # cwd=部署根时密钥在此

    root = args.root or os.environ.get("VSTORE_ROOT") or str(vstore_dir)
    vault_path = args.vault or os.environ.get("VSTORE_VAULT_PATH") or str(Path(root) / "vault")
    index_path = args.index or os.environ.get("VSTORE_INDEX_PATH") or str(Path(root) / "index")
    chunk_path = args.chunks or os.environ.get("VSTORE_CHUNK_PATH") or str(Path(root) / "chunks")

    # 密钥：cwd=部署根时先找 {deploy_root}/vstore/.env，再环境变量
    _load_env_file(deploy_root / "vstore" / ".env")

    sys.path.insert(0, str(vstore_dir.parent))  # 保证 `import vstore` 可用
    from vstore import settings as _settings

    _settings.set_paths(
        _settings.VstorePaths(
            root=root,
            vault_path=vault_path,
            index_path=index_path,
            chunk_path=chunk_path,
        )
    )
    try:
        _settings.get_fernet_key()
    except RuntimeError as exc:
        print(f"[rebuild-chunks] 密钥缺失：{exc}", file=sys.stderr)
        return 2

    from vstore import vault as _vault
    from vstore import embed as _embed

    vault = _vault.Vault(vault_path, index_path)
    vault.index.load()
    vault.chunks.load()

    enc_files = sorted(Path(vault_path).glob("*.enc"))
    print(f"[rebuild-chunks] 密文目录：{vault_path}")
    print(f"[rebuild-chunks] 待处理密文：{len(enc_files)} 条")

    started = time.perf_counter()
    docs: dict[str, list[dict[str, str | int | bool | list[float]]]] = {}
    total_texts = 0
    for n, enc in enumerate(enc_files, 1):
        doc_id = enc.stem
        try:
            doc = _vault._deserialize_doc(_vault.load_encrypted(enc))
        except Exception as exc:
            print(f"  [{n}/{len(enc_files)}] {doc_id} 解密失败，跳过：{exc}")
            continue
        meta = doc.get("meta") or {}
        extracted = doc.get("extracted") or {}
        name = str(meta.get("name") or doc_id)
        built = _vault.build_chunks(
            extracted.get("title") or name,
            extracted.get("outline") or [],
            [str(p or "") for p in extracted.get("paragraphs") or []],
            extracted.get("tables") or [],
        )
        if not built:
            print(f"  [{n}/{len(enc_files)}] {doc_id} 无正文，0 块")
            continue
        vecs = _embed.embed_batch_chunked(
            [c.text for c in built],
            batch_size=args.batch_size,
            workers=args.workers,
        )
        docs[doc_id] = [
            {
                "seq": c.seq,
                "text": c.text,
                "vector": list(vec),
                "doc_name": name,
                "doc_category": str(meta.get("category") or ""),
                "doc_tags": str(meta.get("tags") or ""),
                "as_template": bool(meta.get("as_template")),
                "created_at": str(meta.get("created_at") or ""),
            }
            for c, vec in zip(built, vecs)
        ]
        total_texts += len(built)
        print(f"  [{n}/{len(enc_files)}] {doc_id}（{name[:30]}）→ {len(built)} 块")

    count = vault.chunks.rebuild_all(docs)
    duration_ms = int((time.perf_counter() - started) * 1000)
    print(f"[rebuild-chunks] 完成：docs={len(docs)}, chunks={count}, "
          f"嵌入文本={total_texts}, duration_ms={duration_ms}")
    return 0


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