scripts/kb/build_kb.py

"""知识库入库脚本（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

# 项目根加入 path（直跑脚本时 sys.path[0] 为脚本目录，须自行引导；同 rebuild_alerts）
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))

# 允许 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())