group_xinghuo_jinrong/app/service/milvus_service.py

"""Milvus 接入层（T21-2 · FLOW §2「milvus_tool」/ §3 知识库入库）。

Collection 设计（冻结）：docs/项目框架设计/表设计/03-milvus-collections.md §2
`kb_product_rules`——产品手册/交易规则的语义检索，客户 + 代理人共用，
检索必带溯源字段（source_doc_id + source_version），且只返回已生效文档
（effective_date <= today，§2.3 合规约束）。

环境坑（2026-09-07 实测，含评审 P1 修正）：pymilvus 3.x 在 **import 阶段**
load_dotenv() 并急切解析环境变量 MILVUS_URI，文件 URI（如 .env 的
./data/milvus.db）会让 ``import pymilvus`` 直接抛 Illegal uri。本模块的
防御顺序：先初始化 settings（干净读 .env），再把 os.environ 的 MILVUS_URI
归位为合法 http 占位（仅替换文件 URI），最后 import pymilvus——任何其他
顺序都会以 import 炸 / settings 遮蔽（真连接事故）两种方式踩坑。真实连接
恒走 MilvusClient(uri=...) 显式传参，不经过全局单例。

测试：单测用真 Milvus Lite（tempfile 临时 uri，module 级 fixture 共享），
本机实测可用（0904 交接 + 0907 复测）；不做 mock 双轨。
"""

from __future__ import annotations

import os

from app.config.settings import settings  # noqa: E402

# ---- 环境防御（顺序敏感，评审 P1 修正）：必须在 settings 初始化之后、
# from pymilvus import 之前。原因：pymilvus 3.x import 时 load_dotenv() +
# 急切解析 MILVUS_URI，.env 的文件 URI 会炸 import；但若先改 os.environ 再
# 初始化 settings，pydantic 会读到被污染的环境变量、把 .env 的 MILVUS_URI
# 也遮蔽掉（import 顺序敏感的真连接事故）。故口径：
#   ① settings 先初始化（此时环境干净，milvus_uri 拿到 .env 正确值）；
#   ② 再把 os.environ 的 MILVUS_URI 归位为合法 http 占位（仅文件 URI 被换，
#      合法 http 值保留），pymilvus 全局单例不再炸 import；
#   ③ 真实连接恒走 MilvusClient(uri=settings.milvus_uri) 显式传参。
_milvus_env_uri = os.environ.get("MILVUS_URI")
if _milvus_env_uri is None or not _milvus_env_uri.startswith(("http://", "https://")):
    os.environ["MILVUS_URI"] = "http://localhost:19530"

from pymilvus import DataType, MilvusClient  # noqa: E402

from app.config.settings import settings

COLLECTION_NAME = "kb_product_rules"

# effective_date 用 VARCHAR(10) 存 ISO "YYYY-MM-DD"：Milvus 无原生 DATE 标量
# 类型；ISO 字符串序与日期序一致，filter 比较语义等价（口径登记，评审对照）
EFFECTIVE_DATE_LEN = 10

# 向量检索默认返回条数（TopK；对话 Tool 一期取 3，脚本/服务可覆盖）
DEFAULT_TOP_K = 3


def milvus_client() -> MilvusClient:
    """连接工厂（settings.milvus_uri；测试/脚本注入点）。

    Milvus Lite：uri 为本地文件路径（ MilvusClient 生成的是**目录**，
    交接文档 §八 既有坑）；Standalone：http://host:19530。
    """
    return MilvusClient(uri=settings.milvus_uri)


def _kb_schema(dim: int):
    """kb_product_rules Collection schema（字段与 03-milvus-collections.md §2.1 一致）。"""
    schema = MilvusClient.create_schema(auto_id=False, enable_dynamic_field=False)
    schema.add_field("id", DataType.VARCHAR, is_primary=True, max_length=64)  # {product_id}_{chunk_no}
    schema.add_field("embedding", DataType.FLOAT_VECTOR, dim=dim)
    schema.add_field("product_id", DataType.VARCHAR, max_length=64)
    schema.add_field("product_name", DataType.VARCHAR, max_length=256)
    schema.add_field("doc_type", DataType.VARCHAR, max_length=32)  # prospectus/fee/rule/risk
    schema.add_field("risk_level", DataType.VARCHAR, max_length=2)  # R1~R5
    schema.add_field("source_doc_id", DataType.VARCHAR, max_length=64)  # 溯源必填
    schema.add_field("source_version", DataType.VARCHAR, max_length=32)  # 溯源必填
    schema.add_field("effective_date", DataType.VARCHAR, max_length=EFFECTIVE_DATE_LEN)
    schema.add_field("chunk_text", DataType.VARCHAR, max_length=8192)
    schema.add_field("chunk_no", DataType.INT64)
    return schema


def ensure_collection(client: MilvusClient, dim: int | None = None) -> None:
    """幂等建集合 + 确保加载（exists → load；不存在 → 建 schema + 索引 + load）。

    索引 AUTOINDEX + COSINE：Milvus Lite 支持 AUTOINDEX（自动选型），
    COSINE 语义直观且 bge-m3 未归一化时也正确。dim 缺省取 settings.embed_dim，
    与 embedding 服务同源——两边配置天然一致，不必各配一份。

    load 口径（2026-09-07 实测坑）：Milvus Lite 文件库**跨进程重开**后
    collection 处于 released 状态，直接 search 报 MilvusException(code=101)；
    单测同进程覆盖不到，须每次 ensure 时显式 load_collection。
    """
    dim = dim or settings.embed_dim
    if client.has_collection(COLLECTION_NAME):
        client.load_collection(COLLECTION_NAME)
        return
    index = client.prepare_index_params()
    index.add_index(field_name="embedding", index_type="AUTOINDEX", metric_type="COSINE")
    client.create_collection(
        COLLECTION_NAME, schema=_kb_schema(dim), index_params=index
    )
    client.load_collection(COLLECTION_NAME)


def insert_chunks(client: MilvusClient, rows: list[dict]) -> int:
    """批量写入 chunk（id={product_id}_{chunk_no} 主键 → 重复导入自动覆盖）。

    用 upsert 而非 insert：入库脚本重跑（文档改版重建向量）时按主键覆盖
    旧块，不产生重复 id 报错。返回写入条数（rows 长度，Milvus 侧不二次校验）。
    """
    if not rows:
        return 0
    client.upsert(COLLECTION_NAME, data=rows)
    return len(rows)


def _escape(value: str) -> str:
    """filter 表达式字符串转义（单引号；防 product_id 等值破坏表达式）。"""
    return value.replace("'", "\\'")


def search_kb(
    client: MilvusClient,
    query_vector: list[float],
    *,
    top_k: int = DEFAULT_TOP_K,
    product_id: str | None = None,
    doc_type: str | None = None,
    effective_on: str | None = None,
) -> list[dict]:
    """向量检索 kb_product_rules（§2.3 合规约束内建在本函数）。

    合规口径：只返回已生效文档——effective_date <= effective_on（缺省当天，
    ISO 字符串比较）。product_id/doc_type 为可选标量过滤（检索某产品的
    手册块 / 某类文档）。返回按相似度降序的扁平 list，每项含溯源字段
    （source_doc_id/source_version）+ 原文 chunk_text + score（COSINE 距离，
    越大越相似）。
    """
    conditions = [f"effective_date <= '{_escape(effective_on or _today_iso())}'"]
    if product_id:
        conditions.append(f"product_id == '{_escape(product_id)}'")
    if doc_type:
        conditions.append(f"doc_type == '{_escape(doc_type)}'")
    results = client.search(
        COLLECTION_NAME,
        data=[query_vector],
        limit=top_k,
        filter=" and ".join(conditions),
        output_fields=[
            "product_id",
            "product_name",
            "doc_type",
            "risk_level",
            "source_doc_id",
            "source_version",
            "effective_date",
            "chunk_text",
            "chunk_no",
        ],
    )
    hits = results[0] if results else []
    return [
        {
            "id": hit["id"],
            "score": float(hit["distance"]),
            **hit.get("entity", {}),
        }
        for hit in hits
    ]


def _today_iso() -> str:
    """当天 ISO 日期（集中改点；时区收敛统一走 B8 挂账，与 core_tools._now_naive 同口径）。"""
    import datetime as _dt

    return _dt.date.today().isoformat()