group_fqcd_jr/tools/load_knowledge_milvus.py

"""把 knowledge/_chunks.jsonl 灌入 Milvus 四个知识集合（一次性灌库脚本，可反复重跑）。

设计要点：

1. **集合 schema 不再由本脚本定义** —— 从 `tools/setup_milvus_knowledge_collections.py`
   import `build_schema` / `build_index_params` / `NUM_PARTITIONS` / `FIELD_LIMITS`。
   2026-09-18 之前这里自带第二份字段表（本脚本 `doc_id`/`content` 一套 vs 建表脚本
   `knowledge_id`/`snippet` 一套），**同名集合两套定义**正是 `H-05` ④ 要求收敛的缺陷；
   收敛方向取本脚本这一套，理由见建表脚本的 docstring。
2. **`visibility` 是分区键**，写入时必须**显式**给出，且只接受 `public` / `registered`。
   本脚本在向量化**之前**先逐条校验（fail-closed），不依赖集合默认值兜底 ——
   档位缺省不是"写空串"，而是"不计入任何档位"，属档位越权风险。
3. **向量输入用「标题 + 正文」**而不是只喂正文：标题里带条款号与章节名（如
   「第五章 … 第十三条 … 第一类：资金流转异常」），是比正文更干净的检索信号。
4. **幂等**：用 upsert，同一 doc_id 重复灌库不会产生重复行，可以反复重跑。
5. 灌完立刻做检索自检（拿几个真实客户问题去查），不看自检结果不算灌成功。
"""

import asyncio
import json
import os
import sys
from collections import Counter, defaultdict
from pathlib import Path

import httpx
from dotenv import load_dotenv
from pymilvus import MilvusClient

load_dotenv(override=False)

# 建表脚本与本脚本要能被同一进程 import（`tools/` 不是包，按路径挂进去）。
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
sys.path.insert(0, str(Path(__file__).resolve().parent))

import setup_milvus_knowledge_collections as collections_setup  # noqa: E402
from app.core.knowledge_contracts import ALLOWED_COLLECTIONS, VECTOR_DIM  # noqa: E402
from setup_milvus_knowledge_collections import FIELD_LIMITS  # noqa: E402

MILVUS_URI = os.environ.get("MILVUS_URI", "http://127.0.0.1:19530")
MILVUS_TOKEN = os.environ.get("MILVUS_TOKEN") or None
EMBED_BASE = "https://dashscope.aliyuncs.com/compatible-mode/v1"
#: 与 `model_endpoint_config.id=1`（`knowledge-embedding-qwen-v3`）及 `VECTOR_DIM` 契约一致。
#: 2026-09-18 实测：`text-embedding-v3` 与 `qwen3.7-text-embedding-flash` 都可用、都输出 1024 维，
#: 故此处不一致**不致命**，但必须收敛为一种口径 —— 取现库与文档都在用的 `text-embedding-v3`。
EMBED_MODEL = "text-embedding-v3"
#: 密钥名必须与端点的 `secret_ref`（`env:QWEN_API_KEY`）**逐字一致**，否则
#: `EnvironmentSecretResolver` 报"模型密钥未配置"，表现为客服答不出 → 强制转人工。
EMBED_KEY = os.environ.get("QWEN_API_KEY", "")
DIM = VECTOR_DIM
BATCH = 10

#: 档位取值域。写入侧**只接受**这两个值；空值 / 未知值一律中止灌库（fail-closed）。
VISIBILITY_VALUES = frozenset({"public", "registered"})

#: 与检索白名单同源，避免"灌进了检索查不到的集合"。
COLLECTIONS = sorted(ALLOWED_COLLECTIONS)

# 检索自检用例：(自然语言问题, 期望命中的 doc_id 前缀)。
#
# ⚠️ 期望值是**语料版本的函数**，改语料就要跟着改：2026-09-18 重灌时（FAQ 由 44 条
# V1.x 换成 `D6.1.3` 的 64 条 V2.0、政策与产品手册同步 V2.0）旧期望值全部失效，
# 表现为「自检只有 3/7 命中」——但逐条看会发现**排序其实是对的**，只是编号变了
# （例如「基金赎回到账」从 `FAQ-0016` 变成 `FAQ-0026`）。这组值即重灌当日的实测基线。
CHECKS = [
    ("基金赎回到账需要多长时间", "FAQ-0026"),
    # 分等级的「能买什么」：`POL-AST-012` 的匹配矩阵与 FAQ 是同一份内容，
    # 但这一问的措辞更近条款语言，故政策集合排前（实测 0.792）。
    ("C1 保守型客户可以买哪些风险等级的产品", "POL-AST-012"),
    ("C1 客户能买什么", "FAQ-0019"),
    ("南方季季盈90天的起投金额是多少", "PROD-006"),
    ("开户需要准备哪些材料", "FAQ-0036"),
    ("业绩比较基准是什么意思", "FAQ-0022"),
    ("高净值客户能享受什么费率优惠", "HNW-005"),
    # `乙-7` 第 4 集合（金融行业基础信息）：这三条**只有**新集合能答，
    # 用来证明第 4 集合确实灌进去了、也确实进了检索面。
    ("基金定投是什么", "BAS-TRD-012"),
    ("场内基金和场外基金有什么区别", "BAS-CON-006"),
    # 期望值是 FAQ-0065（`W6` 补的专条「什么是T日、T+1？」）而不是第 4 集合 ——
    # 实测 FAQ-0065 得 0.703 胜出，这是**更好**的结果：专条比通用常识更贴题。
    # 第 4 集合是「FAQ 答不了时的兜底」，不是"抢答 FAQ 能答的题"。
    ("T+1 是什么意思", "FAQ-0065"),
]


async def embed(texts: list[str]) -> list[list[float]]:
    async with httpx.AsyncClient(timeout=90) as client:
        response = await client.post(
            f"{EMBED_BASE}/embeddings",
            headers={"Authorization": f"Bearer {EMBED_KEY}"},
            json={"model": EMBED_MODEL, "input": texts},
        )
        response.raise_for_status()
        payload = response.json()
    items = sorted(payload["data"], key=lambda item: item["index"])
    return [item["embedding"] for item in items]


async def main() -> None:
    if not EMBED_KEY:
        print("缺少 QWEN_API_KEY（必须与 model_endpoint_config.secret_ref 逐字一致）")
        return

    records = [
        json.loads(line)
        for line in (Path("knowledge") / "_chunks.jsonl").read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]
    print(f"待入库块数：{len(records)}")

    # 档位先校验后向量化：非法/缺失时**一条向量都不调**，避免花了钱才发现白灌。
    invalid = [r for r in records if str(r.get("visibility", "")) not in VISIBILITY_VALUES]
    if invalid:
        print(f"档位非法或缺失：{len(invalid)} 条（只接受 {sorted(VISIBILITY_VALUES)}），已中止：")
        for record in invalid[:5]:
            print(f"   {record.get('doc_id')} visibility={record.get('visibility')!r}")
        return
    print("档位分布：" + str(dict(Counter(str(r['visibility']) for r in records))))

    print("\n== 建集合（幂等；结构冲突即中止，绝不覆盖） ==")
    _created, _existed, conflicting = await collections_setup.ensure_collections(
        MILVUS_URI, MILVUS_TOKEN or ""
    )
    if conflicting:
        print("集合结构冲突，已中止（未做任何覆盖）：")
        for line in conflicting:
            print(f"  - {line}")
        return

    client = MilvusClient(uri=MILVUS_URI, token=MILVUS_TOKEN)

    grouped: dict[str, list[dict[str, object]]] = defaultdict(list)
    for record in records:
        grouped[str(record["collection"])].append(record)

    print("\n== 生成向量并写入 ==")
    truncated: dict[str, int] = {}
    for name, group in grouped.items():
        rows: list[dict[str, object]] = []
        for start in range(0, len(group), BATCH):
            batch = group[start:start + BATCH]
            vectors = await embed(
                [f"{record['title']}\n{record['content']}" for record in batch]
            )
            for record, vector in zip(batch, vectors, strict=True):
                # 截断长度**只从集合定义取**（`FIELD_LIMITS` 由建表模块导出）：
                # 这里若写死数字，就会出现"脚本截到 8000、集合只给 4096"这类
                # 只在写入那一刻才暴露的错配，而且改一处必须记得改另一处。
                row: dict[str, object] = {}
                for field, limit in FIELD_LIMITS.items():
                    value = str(record[field])
                    if len(value) > limit:
                        truncated[field] = truncated.get(field, 0) + 1
                        value = value[:limit]
                    row[field] = value
                row["embedding"] = vector
                rows.append(row)
            print(f"  {name}: 已向量化 {min(start + BATCH, len(group))}/{len(group)}")
        client.upsert(collection_name=name, data=rows)
        client.flush(collection_name=name)

    if truncated:
        print("⚠️ 有字段被截断到集合上限（计数）：", truncated)
    else:
        print("所有字段均在集合定义的长度上限内，无截断。")

    print("\n== 各集合条目数 ==")
    for name in COLLECTIONS:
        stats = client.get_collection_stats(collection_name=name)
        print(f"  {name}: {stats.get('row_count')} 条")

    print("\n== 检索自检 ==")
    queries = [question for question, _ in CHECKS]
    query_vectors = await embed(queries)
    hits = 0
    for (question, expected), vector in zip(CHECKS, query_vectors, strict=True):
        results = client.search(
            collection_name="fin_policy_collection", data=[vector], limit=3,
            output_fields=["doc_id", "title", "visibility"],
        )
        top = client.search(
            collection_name="fin_faq_collection", data=[vector], limit=3,
            output_fields=["doc_id", "title"],
        )
        product = client.search(
            collection_name="fin_product_collection", data=[vector], limit=3,
            output_fields=["doc_id", "title"],
        )
        basic = client.search(
            collection_name="fin_basic_collection", data=[vector], limit=3,
            output_fields=["doc_id", "title"],
        )
        merged = [item for group in (basic, top, product, results) for item in group[0]]
        merged.sort(key=lambda item: item["distance"], reverse=True)
        best = merged[0] if merged else None
        # 注意：COSINE 下 pymilvus 返回的 distance 越大越相似
        found = ""
        if best is not None:
            found = f"{best['entity']['doc_id']} (score={best['distance']:.3f}) {best['entity']['title'][:44]}"
        mark = "OK " if best is not None and str(best["entity"]["doc_id"]).startswith(expected) else "检查"
        hits += 1 if mark == "OK " else 0
        print(f"  [{mark}] {question}\n        → {found}")
    print(f"\n自检命中 {hits}/{len(CHECKS)}")


asyncio.run(main())