feat: T21-3 知识库首批数据+入库脚本——data/kb 6 只种子产品手册(每只 prospectus/fee/rule/risk 四节 24 块, 费率与 core_product seed 对齐)+build_kb.py(front-matter+节标题切块/doc_type 白名单/Core 权威元数据缺省即中止/一次批量 embedding/upsert 幂等/--dry-run); 脚本单测 10 例, 404 绿

2026-09-07 11:55:24 +08:00
parent 1056a95343
commit 7edf522aff
2 changed files with 310 additions and 0 deletions
@@ -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<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())
@@ -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)