114 lines
3.7 KiB
Python
114 lines
3.7 KiB
Python
|
"""知识库入库脚本：解析桌面知识库文件 → 向量化 → 存入 Milvus。
|
||
|
|
|
||
|
|
用法：
|
||
|
|
python scripts/kb/build_collections.py # 仅创建不存在的 collection
|
||
|
|
python scripts/kb/build_collections.py --rebuild # 先删除再重建
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
# 把项目根目录加入 sys.path
|
||
|
|
_ROOT = Path(__file__).resolve().parent.parent.parent
|
||
|
|
sys.path.insert(0, str(_ROOT))
|
||
|
|
|
||
|
|
from app.config.settings import settings
|
||
|
|
from app.tool.document_parser import parse_collection_dir
|
||
|
|
from app.tool.embedding_tool import get_embedder
|
||
|
|
from app.tool.milvus_tool import get_milvus_client
|
||
|
|
|
||
|
|
|
||
|
|
# Collection → 子目录映射
|
||
|
|
KB_LAYOUT: dict[str, str] = {
|
||
|
|
"fin_faq": "fin_faq_collection",
|
||
|
|
"fin_product": "fin_product_collection",
|
||
|
|
"fin_policy": "fin_policy_collection",
|
||
|
|
}
|
||
|
|
|
||
|
|
# 每个 collection 的标量字段（从 chunk.metadata 中取值）
|
||
|
|
COLLECTION_FIELDS: dict[str, list[str]] = {
|
||
|
|
"fin_faq": ["question", "answer", "category", "source_doc", "chunk_no", "chunk_text"],
|
||
|
|
"fin_product": ["product_name", "risk_level", "doc_type", "source_doc", "chunk_no", "chunk_text"],
|
||
|
|
"fin_policy": ["policy_name", "chapter", "source_doc", "chunk_no", "chunk_text"],
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def chunk_to_record(chunk, collection: str, vector: list[float]) -> dict:
|
||
|
|
"""把 Chunk + 向量转为 Milvus 插入记录。"""
|
||
|
|
record = {
|
||
|
|
"id": chunk.chunk_id,
|
||
|
|
"embedding": vector,
|
||
|
|
}
|
||
|
|
for field in COLLECTION_FIELDS[collection]:
|
||
|
|
record[field] = chunk.metadata.get(field, "")
|
||
|
|
# 类型修正：chunk_no 必须是 int
|
||
|
|
record["chunk_no"] = int(chunk.metadata.get("chunk_no", 0))
|
||
|
|
return record
|
||
|
|
|
||
|
|
|
||
|
|
def main(rebuild: bool = False) -> None:
|
||
|
|
embedder = get_embedder()
|
||
|
|
milvus = get_milvus_client()
|
||
|
|
kb_root = Path(settings.kb_root_dir)
|
||
|
|
|
||
|
|
if not kb_root.exists():
|
||
|
|
print(f"[ERROR] 知识库目录不存在: {kb_root}")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
total_inserted = 0
|
||
|
|
|
||
|
|
for coll_name, sub_dir in KB_LAYOUT.items():
|
||
|
|
dir_path = kb_root / sub_dir
|
||
|
|
if not dir_path.exists():
|
||
|
|
print(f"[WARN] 跳过 {coll_name}: 目录不存在 {dir_path}")
|
||
|
|
continue
|
||
|
|
|
||
|
|
# 1. 解析文档
|
||
|
|
print(f"\n{'=' * 60}")
|
||
|
|
print(f"[{coll_name}] 开始解析: {dir_path}")
|
||
|
|
chunks, _ = parse_collection_dir(dir_path)
|
||
|
|
print(f"[{coll_name}] 解析完成: {len(chunks)} 块")
|
||
|
|
|
||
|
|
if not chunks:
|
||
|
|
print(f"[WARN] {coll_name} 无有效内容，跳过")
|
||
|
|
continue
|
||
|
|
|
||
|
|
# 2. 重建或确保 collection
|
||
|
|
if rebuild:
|
||
|
|
milvus.drop_collection(coll_name)
|
||
|
|
milvus.ensure_collection(coll_name)
|
||
|
|
|
||
|
|
# 3. 向量化
|
||
|
|
print(f"[{coll_name}] 开始向量化 ({len(chunks)} 条)...")
|
||
|
|
texts = [c.text for c in chunks]
|
||
|
|
vectors = embedder.embed_batch(texts)
|
||
|
|
|
||
|
|
# 4. 组装记录并插入
|
||
|
|
records = []
|
||
|
|
skipped = 0
|
||
|
|
for chunk, vec in zip(chunks, vectors):
|
||
|
|
if not vec:
|
||
|
|
skipped += 1
|
||
|
|
continue
|
||
|
|
records.append(chunk_to_record(chunk, coll_name, vec))
|
||
|
|
|
||
|
|
if skipped:
|
||
|
|
print(f"[{coll_name}] 跳过 {skipped} 条（向量化失败）")
|
||
|
|
|
||
|
|
insert_count = milvus.insert(coll_name, records)
|
||
|
|
print(f"[{coll_name}] 插入 {insert_count} 条记录")
|
||
|
|
total_inserted += insert_count
|
||
|
|
|
||
|
|
print(f"\n{'=' * 60}")
|
||
|
|
print(f"入库完成：共插入 {total_inserted} 条记录到 {len(KB_LAYOUT)} 个 collection")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
parser = argparse.ArgumentParser(description="知识库入库")
|
||
|
|
parser.add_argument("--rebuild", action="store_true", help="先删除再重建 collection")
|
||
|
|
args = parser.parse_args()
|
||
|
|
main(rebuild=args.rebuild)
|