177 lines
7.4 KiB
Python
177 lines
7.4 KiB
Python
"""知识库入库脚本（T21-3 · FLOW §3 主链路「md 手册 → 切块 → embedding → Milvus」）。
|
||
|
||
用法（项目根目录）：
|
||
python scripts/kb/build_kb.py --dry-run # 只解析打印，不连任何外部服务
|
||
python scripts/kb/build_kb.py # 解析 + embedding(Ollama) + 写入 Milvus
|
||
python scripts/kb/build_kb.py --kb-dir data/kb
|
||
|
||
约定：
|
||
- 手册文件 ``data/kb/{product_id}.md``；front-matter 提供
|
||
version / effective_date（溯源与生效控制）；正文节标题格式
|
||
``## 【doc_type】标题``（doc_type ∈ prospectus/fee/rule/risk），
|
||
一节 = 一个 chunk。
|
||
- product_name / risk_level 以 **jinrong_core.core_product 为权威**
|
||
（L0 口径：画像/知识不得覆盖 Core 正式数据），Core 查不到直接报错退出，
|
||
不允许静默用缺省值入库。
|
||
- 幂等：chunk 主键 {product_id}_{chunk_no} + upsert，重跑覆盖不重复。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import re
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
# 允许 doc_type 白名单（03-milvus-collections.md §2.1 冻结枚举）
|
||
DOC_TYPES = ("prospectus", "fee", "rule", "risk")
|
||
|
||
# 节标题解析：## 【type】标题
|
||
_SECTION_RE = re.compile(r"^##\s*【(?P<type>[a-z_]+)】(?P<title>.+)$")
|
||
|
||
# chunk_text 写入上限（字符）。Milvus VARCHAR(max_length=8192) 按字符计；
|
||
# 演示手册单节数百字，2500 为保守防御（防手滑写爆 + 中英文混排余量）
|
||
CHUNK_TEXT_MAX_CHARS = 2500
|
||
|
||
FRONT_MATTER_RE = re.compile(r"^---\s*\n(?P<body>.*?)\n---\s*\n", re.DOTALL)
|
||
|
||
|
||
def parse_front_matter(text: str) -> tuple[dict[str, str], str]:
|
||
"""提取 --- 包围的 front-matter（key: value 逐行），返回 (meta, 正文)。"""
|
||
match = FRONT_MATTER_RE.match(text)
|
||
if not match:
|
||
raise ValueError("缺少 front-matter（--- 包围的 product_id/version/effective_date）")
|
||
meta: dict[str, str] = {}
|
||
for line in match.group("body").splitlines():
|
||
line = line.strip()
|
||
if not line or ":" not in line:
|
||
continue
|
||
key, _, value = line.partition(":")
|
||
meta[key.strip()] = value.strip()
|
||
return meta, text[match.end():]
|
||
|
||
|
||
def parse_sections(body: str) -> list[dict[str, str]]:
|
||
"""正文按 ``## 【type】标题`` 切节；一节一个 chunk（含标题行下的全部文本）。
|
||
|
||
非法 doc_type（不在白名单）直接报错——比静默跳过更早暴露手册笔误。
|
||
"""
|
||
sections: list[dict[str, str]] = []
|
||
current: dict[str, str] | None = None
|
||
for line in body.splitlines():
|
||
matched = _SECTION_RE.match(line.strip())
|
||
if matched:
|
||
doc_type = matched.group("type")
|
||
if doc_type not in DOC_TYPES:
|
||
raise ValueError(f"节标题 doc_type 非法：{doc_type}（白名单 {DOC_TYPES}）")
|
||
current = {"doc_type": doc_type, "title": matched.group("title").strip(), "text": ""}
|
||
sections.append(current)
|
||
elif current is not None and line.strip():
|
||
# 节内正文（空行不累计；一级标题/结尾注释行在首个节标题前会被忽略）
|
||
current["text"] = (current["text"] + "\n" + line.strip()).strip()
|
||
return [s for s in sections if s["text"]]
|
||
|
||
|
||
def build_chunk_rows(
|
||
meta: dict[str, str], sections: list[dict[str, str]], product: dict
|
||
) -> list[dict]:
|
||
"""front-matter + 节列表 + Core 产品行 → Milvus 行（不含 embedding 字段）。
|
||
|
||
溯源：source_doc_id 固定 KB-{product_id}（对应 data/kb 源文件），
|
||
source_version / effective_date 取 front-matter；chunk_no 全文档递增。
|
||
"""
|
||
required = ("product_id", "version", "effective_date")
|
||
missing = [k for k in required if not meta.get(k)]
|
||
if missing:
|
||
raise ValueError(f"front-matter 缺少字段：{missing}")
|
||
rows: list[dict] = []
|
||
for no, sec in enumerate(sections):
|
||
rows.append(
|
||
{
|
||
"id": f"{meta['product_id']}_{no}",
|
||
"product_id": meta["product_id"],
|
||
"product_name": product["product_name"],
|
||
"doc_type": sec["doc_type"],
|
||
"risk_level": product["min_risk_code"],
|
||
"source_doc_id": f"KB-{meta['product_id']}",
|
||
"source_version": meta["version"],
|
||
"effective_date": meta["effective_date"],
|
||
"chunk_text": sec["text"][:CHUNK_TEXT_MAX_CHARS],
|
||
"chunk_no": no,
|
||
}
|
||
)
|
||
return rows
|
||
|
||
|
||
def parse_kb_file(path: Path) -> tuple[dict[str, str], list[dict[str, str]]]:
|
||
"""单文件解析（front-matter + 节）；供 main 与单测共用。"""
|
||
meta, body = parse_front_matter(path.read_text(encoding="utf-8"))
|
||
if not meta.get("product_id"):
|
||
raise ValueError(f"{path.name}: front-matter 缺 product_id")
|
||
return meta, parse_sections(body)
|
||
|
||
|
||
def main(argv: list[str] | None = None) -> int:
|
||
parser = argparse.ArgumentParser(description="知识库产品手册入库（T21-3）")
|
||
parser.add_argument("--kb-dir", default="data/kb", help="手册目录（默认 data/kb）")
|
||
parser.add_argument("--dry-run", action="store_true", help="只解析打印，不连外部服务")
|
||
args = parser.parse_args(argv)
|
||
|
||
kb_dir = Path(args.kb_dir)
|
||
files = sorted(kb_dir.glob("*.md"))
|
||
if not files:
|
||
print(f"[kb] {kb_dir} 下没有 .md 手册")
|
||
return 1
|
||
|
||
# ---------- 解析全部文件（不依赖外部服务） ----------
|
||
parsed: list[tuple[dict, list[dict]]] = []
|
||
for path in files:
|
||
try:
|
||
meta, sections = parse_kb_file(path)
|
||
except ValueError as exc:
|
||
print(f"[kb] 解析失败 {path.name}: {exc}")
|
||
return 1
|
||
parsed.append((meta, sections))
|
||
print(f"[kb] {path.name}: product={meta['product_id']} 节数={len(sections)} "
|
||
f"types={[s['doc_type'] for s in sections]}")
|
||
|
||
if args.dry_run:
|
||
print("[kb] dry-run 完成，未连接 Ollama / Milvus / MySQL")
|
||
return 0
|
||
|
||
# ---------- 延迟导入：dry-run 不触发 pymilvus / Ollama / MySQL 依赖 ----------
|
||
from app.repository.core_ro import CoreReadOnlyRepository
|
||
from app.service import embedding, milvus_service
|
||
|
||
core_ro = CoreReadOnlyRepository()
|
||
|
||
# 全部 rows 聚合后一次批量 embedding（1 次 HTTP，24 块级规模足够）
|
||
all_rows: list[dict] = []
|
||
for meta, sections in parsed:
|
||
product = core_ro.get_product(meta["product_id"])
|
||
if product is None:
|
||
# 权威元数据缺失直接失败：禁止用缺省名/缺省风险等级入库（L0 口径）
|
||
print(f"[kb] core_product 无 {meta['product_id']}，中止（请先灌 Core 库）")
|
||
return 1
|
||
all_rows.extend(build_chunk_rows(meta, sections, product))
|
||
|
||
if not all_rows:
|
||
print("[kb] 没有可入库的 chunk（手册节全为空？）")
|
||
return 1
|
||
|
||
print(f"[kb] 批量向量化 {len(all_rows)} 块（模型 {embedding.settings.embed_model}）…")
|
||
vectors = embedding.embed_texts([r["chunk_text"] for r in all_rows])
|
||
for row, vec in zip(all_rows, vectors):
|
||
row["embedding"] = vec
|
||
|
||
client = milvus_service.milvus_client()
|
||
milvus_service.ensure_collection(client)
|
||
written = milvus_service.insert_chunks(client, all_rows)
|
||
client.close()
|
||
print(f"[kb] 写入 {milvus_service.COLLECTION_NAME} 成功：{written} 块（upsert 幂等，可重跑）")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|