161 lines
6.7 KiB
Python
161 lines
6.7 KiB
Python
"""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 实测）：pymilvus 3.x 在 **import 阶段** load_dotenv()
|
||
并急切解析环境变量 MILVUS_URI，文件 URI（如本项目 .env 的
|
||
./data/milvus.db）会让 ``import pymilvus`` 直接抛 Illegal uri。因此本模块
|
||
在 import pymilvus **之前**用 setdefault 预置一个合法 URI（dotenv 不覆盖
|
||
已存在变量）；真实连接由 MilvusClient(uri=...) 显式传入 settings.milvus_uri，
|
||
不经过全局单例。若未来接 Standalone（Docker），.env 设合法 http URI 即可。
|
||
|
||
测试：单测用真 Milvus Lite（tempfile 临时 uri，module 级 fixture 共享），
|
||
本机实测可用（0904 交接 + 0907 复测）；不做 mock 双轨。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
|
||
# ---- 必须先于 from pymilvus import：防 .env 的文件 URI 炸掉 import（见模块注释）----
|
||
os.environ.setdefault("MILVUS_URI", "http://localhost:19530")
|
||
|
||
from pymilvus import DataType, MilvusClient
|
||
|
||
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()
|