diff --git a/scripts/kb/build_kb.py b/scripts/kb/build_kb.py new file mode 100644 index 0000000..a14f3fc --- /dev/null +++ b/scripts/kb/build_kb.py @@ -0,0 +1,176 @@ +"""知识库入库脚本（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[a-z_]+)】(?P.+)$") + +# 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()) diff --git a/tests/test_build_kb.py b/tests/test_build_kb.py new file mode 100644 index 0000000..ad93bab --- /dev/null +++ b/tests/test_build_kb.py @@ -0,0 +1,134 @@ +"""T21-3 build_kb 脚本单测：front-matter 解析 / 节切块 / Milvus 行构造（纯函数）。 + +脚本目录非包，sys.path 动态导入（同 test_demo_scripts 模式）。 +不连 Ollama / Milvus / MySQL（真库联调归收尾验证步骤）。 +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +KB_DIR = Path(__file__).resolve().parents[1] / "scripts" / "kb" +sys.path.insert(0, str(KB_DIR)) + +from build_kb import ( # noqa: E402 + DOC_TYPES, + build_chunk_rows, + parse_front_matter, + parse_kb_file, + parse_sections, +) + +PRODUCT = {"product_id": "PROD-X", "product_name": "测试产品", "min_risk_code": "R2"} + +SAMPLE = """--- +product_id: PROD-X +version: 2026.09 +effective_date: 2026-09-01 +--- + +# 测试产品 · 产品手册 + +## 【prospectus】产品概况 + +定位与投资范围说明。 + +## 【fee】费率结构 + +管理费 0.20%/年，托管费 0.10%/年。 + +## 【rule】申赎与交易规则 + +申购 T+1 确认。 + +## 【risk】风险揭示 + +不保本不保收益。 +""" + + +class TestFrontMatter: + def test_parse_meta_and_body_split(self): + meta, body = parse_front_matter(SAMPLE) + assert meta == { + "product_id": "PROD-X", + "version": "2026.09", + "effective_date": "2026-09-01", + } + assert body.startswith("# 测试产品") + + def test_missing_front_matter_raises(self): + with pytest.raises(ValueError, match="front-matter"): + parse_front_matter("# 无 front-matter 正文") + + +class TestSections: + def test_four_sections_with_types(self): + sections = parse_sections(SAMPLE) + assert [s["doc_type"] for s in sections] == list(DOC_TYPES) + assert sections[0]["title"] == "产品概况" + assert "投资范围" in sections[0]["text"] + + def test_section_text_excludes_heading_and_blank_lines(self): + sections = parse_sections(SAMPLE) + # 正文不含标题行、不含空行累计 + assert "【fee】" not in sections[1]["text"] + assert sections[1]["text"].count("\n\n") == 0 + + def test_illegal_doc_type_raises(self): + bad = SAMPLE.replace("【fee】", "【rates】") + with pytest.raises(ValueError, match="doc_type 非法"): + parse_sections(bad) + + def test_empty_body_sections_dropped(self): + body = "## 【fee】费率\n\n## 【risk】风险\n有内容" + sections = parse_sections(body) + assert len(sections) == 1 + assert sections[0]["doc_type"] == "risk" + + +class TestChunkRows: + def test_row_fields_and_chunk_no(self): + meta = {"product_id": "PROD-X", "version": "2026.09", "effective_date": "2026-09-01"} + sections = parse_sections(SAMPLE) + rows = build_chunk_rows(meta, sections, PRODUCT) + assert [r["id"] for r in rows] == ["PROD-X_0", "PROD-X_1", "PROD-X_2", "PROD-X_3"] + first = rows[0] + # 溯源字段必填（03-milvus-collections.md §2.3） + assert first["source_doc_id"] == "KB-PROD-X" + assert first["source_version"] == "2026.09" + assert first["effective_date"] == "2026-09-01" + # product_name / risk_level 取 Core 权威（不入库文件自述值） + assert first["product_name"] == "测试产品" + assert first["risk_level"] == "R2" + assert first["chunk_no"] == 0 + + def test_missing_meta_field_raises(self): + meta = {"product_id": "PROD-X", "version": "2026.09"} # 缺 effective_date + with pytest.raises(ValueError, match="缺少字段"): + build_chunk_rows(meta, [{"doc_type": "fee", "title": "t", "text": "x"}], PRODUCT) + + def test_chunk_text_truncated(self): + from build_kb import CHUNK_TEXT_MAX_CHARS + + meta = {"product_id": "PROD-X", "version": "v", "effective_date": "2026-09-01"} + long_text = "长" * (CHUNK_TEXT_MAX_CHARS + 100) + rows = build_chunk_rows(meta, [{"doc_type": "fee", "title": "t", "text": long_text}], PRODUCT) + assert len(rows[0]["chunk_text"]) == CHUNK_TEXT_MAX_CHARS + + +class TestRealFiles: + """真手册文件抽查：data/kb 全部文件可解析且 4 节齐。""" + + def test_all_kb_files_parse(self): + kb_root = Path(__file__).resolve().parents[1] / "data" / "kb" + files = sorted(kb_root.glob("*.md")) + assert len(files) == 6 + for f in files: + meta, sections = parse_kb_file(f) + assert meta["product_id"] == f.stem + assert [s["doc_type"] for s in sections] == ["prospectus", "fee", "rule", "risk"] + assert all(len(s["text"]) > 20 for s in sections)