78 lines
3.0 KiB
Python
78 lines
3.0 KiB
Python
"""RAG 检索编排（T21-4 · FLOW §2「milvus_tool：产品规则 RAG + source_refs」）。
|
||
|
||
职责：query 文本 → embedding（Ollama bge-m3）→ Milvus kb_product_rules
|
||
向量检索 → 返回 chunks + **source_refs 溯源清单**（source_doc_id +
|
||
source_version 必带，03-milvus-collections.md §2.3 合规口径：回答必须可溯源）。
|
||
|
||
kb_business_ops（代理人内部制度）一期不接——按 05-底座清单「代理人组开发
|
||
时再建」；本模块仅 kb_product_rules（客户 + 代理人共用，用户拍板 2026-09-07：
|
||
Tool 仅 customer/advisor 开放）。
|
||
|
||
失败口径：EmbeddingError / Milvus 异常**原样上抛**（不吞不降级）——检索
|
||
失败由调用方决定呈现（对话 Tool 层 run_tool 统一转 TOOL_ERROR 留痕），
|
||
此处假装「检索到空结果」会让 LLM 编造回答，比报错更危险。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import Any
|
||
|
||
from app.service import embedding, milvus_service
|
||
|
||
# 对外默认 TopK（对话 Tool 引用；脚本/服务可显式覆盖）
|
||
DEFAULT_TOP_K = 3
|
||
|
||
|
||
def _embed(query: str) -> list[float]:
|
||
"""query 向量化（测试注入点）。"""
|
||
return embedding.embed_text(query)
|
||
|
||
|
||
def _client() -> "milvus_service.MilvusClient":
|
||
"""Milvus 连接（测试注入点；与 milvus_service.milvus_client 同源）。"""
|
||
return milvus_service.milvus_client()
|
||
|
||
|
||
def _source_refs(results: list[dict[str, Any]]) -> list[dict[str, str]]:
|
||
"""结果 → 去重溯源清单（来源文档 × 版本 × 产品维度）。"""
|
||
refs: dict[tuple[str, str, str], dict[str, str]] = {}
|
||
for r in results:
|
||
key = (r.get("source_doc_id", ""), r.get("source_version", ""), r.get("product_id", ""))
|
||
refs.setdefault(
|
||
key,
|
||
{
|
||
"source_doc_id": r.get("source_doc_id", ""),
|
||
"source_version": r.get("source_version", ""),
|
||
"product_id": r.get("product_id", ""),
|
||
"product_name": r.get("product_name", ""),
|
||
},
|
||
)
|
||
return list(refs.values())
|
||
|
||
|
||
def search_knowledge(
|
||
query: str,
|
||
*,
|
||
product_id: str | None = None,
|
||
doc_type: str | None = None,
|
||
top_k: int = DEFAULT_TOP_K,
|
||
) -> dict[str, Any]:
|
||
"""知识检索主入口：query → chunks（含溯源字段）+ source_refs。
|
||
|
||
product_id / doc_type 为可选标量过滤；effective_date 合规过滤
|
||
（只返回已生效文档）内建在 milvus_service.search_kb。
|
||
"""
|
||
if not query or not query.strip():
|
||
return {"query": query.strip(), "results": [], "source_refs": []}
|
||
vector = _embed(query.strip())
|
||
client = _client()
|
||
try:
|
||
# ensure_collection 幂等防御：空库/首访时明确空结果而非报错
|
||
milvus_service.ensure_collection(client)
|
||
results = milvus_service.search_kb(
|
||
client, vector, top_k=top_k, product_id=product_id, doc_type=doc_type
|
||
)
|
||
finally:
|
||
client.close()
|
||
return {"query": query.strip(), "results": results, "source_refs": _source_refs(results)}
|