64 lines
2.5 KiB
Python
64 lines
2.5 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
"""知识库入库脚本：分块 → LLM 生成问答 → 双路向量索引 → 落盘 kb_index.json / kb_qa.json。
|
||
|
||
知识源 = backend/knowledge/ 下全部 *.md（park.md / opc.md / park_data.md …）。
|
||
- 幂等：缓存有效时直接加载（不重新构建）；--force 强制删除缓存重建。
|
||
- 首次运行会调用 DashScope（text-embedding-v4 向量化 + qwen-plus 生成问答），
|
||
耗时约数分钟，仅构建一次，之后命中缓存秒开。
|
||
|
||
用法（在 backend/ 目录下执行）：
|
||
python scripts/ingest_kb.py # 有缓存则加载，无缓存则构建
|
||
python scripts/ingest_kb.py --force # 强制重建（任一知识文件变更后）
|
||
"""
|
||
import argparse
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
BACKEND_DIR = Path(__file__).resolve().parent.parent
|
||
if str(BACKEND_DIR) not in sys.path:
|
||
sys.path.insert(0, str(BACKEND_DIR))
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(description="知识库入库（分块 + 问答 + 双路向量索引）")
|
||
parser.add_argument("--force", action="store_true", help="删除缓存并强制重建")
|
||
parser.add_argument("--dry", action="store_true", help="只做分块预览，不调 API、不落盘")
|
||
args = parser.parse_args()
|
||
|
||
from app import rag
|
||
|
||
if args.dry:
|
||
chunks = rag._chunk_md(rag._kb_text())
|
||
print(f"[dry] 知识源 {len(list(rag.KB_DIR.glob(rag.KB_GLOB)))} 个文件，全库 "
|
||
f"{len(rag._kb_text())} 字 → {len(chunks)} 块（最长 {max(map(len, chunks))} 字）")
|
||
return
|
||
|
||
if args.force:
|
||
for p in (rag.META_FILE, rag.VECTOR_FILE, rag.QA_FILE, rag.LEGACY_INDEX_FILE):
|
||
if p.exists():
|
||
p.unlink()
|
||
print(f"[ingest] 已删除缓存 {p.name}")
|
||
|
||
print("[ingest] 知识源：")
|
||
for f in sorted(rag.KB_DIR.glob(rag.KB_GLOB)):
|
||
print(f" - {f.name} ({f.stat().st_size} B)")
|
||
|
||
idx = rag._load_index()
|
||
n_chunk = len(idx["chunks"])
|
||
n_vec = len(idx["embed_texts"])
|
||
dim = idx["embeddings"].shape[1]
|
||
n_q = sum(1 for qs in idx["questions"] for _ in qs)
|
||
print("[ingest] 入库完成：")
|
||
print(f" - 分块数: {n_chunk}")
|
||
print(f" - 生成问句数: {n_q}")
|
||
print(f" - 向量数: {n_vec}（dim={dim}，含原句+问句双路）")
|
||
print(f" - 向量矩阵: {rag.VECTOR_FILE}")
|
||
print(f" - 元数据: {rag.META_FILE}")
|
||
print(f" - 问答缓存: {rag.QA_FILE}")
|
||
print("[ingest] 成功。")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|