feat: T21-4 RAG 检索编排——rag_service.search_knowledge(query→bge-m3 向量→Milvus 检索→chunks+source_refs 溯源去重), 异常原样上抛不吞(防 LLM 编造回答), 空白 query 短路, ensure_collection 幂等防御; 单测 8 例注入 mock, 412 绿

2026-09-07 12:01:45 +08:00
parent 974fff9507
commit e2a444453e
2 changed files with 202 additions and 1 deletions
@@ -1 +1,77 @@
-"""RAG 编排：检索 kb_product_rules / kb_business_ops + 溯源。"""
+"""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)}
@@ -0,0 +1,125 @@
+"""T21-4 rag_service 单测：检索编排 + source_refs 溯源（monkeypatch 注入，不连外部）。
+
+embedding 与 Milvus client 均经模块级 _embed/_client 注入点 mock；
+Milvus 检索行为用假 client（FakeClient）模拟，search_kb 参数透传断言。
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from app.service import rag_service as rs
+from app.service.embedding import EmbeddingError
+
+
+class FakeClient:
+    """假 Milvus client：记录 ensure/close 调用，search 透传给 stub。"""
+
+    def __init__(self, search_stub):
+        self.search_stub = search_stub
+        self.closed = False
+        self.last_kwargs: dict = {}
+
+    def close(self):
+        self.closed = True
+
+
+def _client_with(monkeypatch, stub) -> FakeClient:
+    """构造假 client 并注入 rs._client（embed 走假 8 维向量）。"""
+    client = FakeClient(stub)
+    monkeypatch.setattr(rs, "_client", lambda: client)
+    monkeypatch.setattr(rs, "_embed", lambda q: [0.1] * 8)
+    monkeypatch.setattr(
+        rs.milvus_service, "ensure_collection", lambda c, dim=None: None
+    )
+    monkeypatch.setattr(
+        rs.milvus_service,
+        "search_kb",
+        lambda c, vector, **kwargs: (client.last_kwargs.update(kwargs), stub(vector, **kwargs))[1],
+    )
+    return client
+
+
+def _hit(idx: str, doc: str = "KB-P1", ver: str = "v1", pid: str = "PROD-1") -> dict:
+    return {
+        "id": idx,
+        "score": 0.9,
+        "product_id": pid,
+        "product_name": "产品一",
+        "doc_type": "fee",
+        "risk_level": "R2",
+        "source_doc_id": doc,
+        "source_version": ver,
+        "effective_date": "2026-09-01",
+        "chunk_text": "申购费率 1.5%",
+        "chunk_no": 0,
+    }
+
+
+class TestSearch:
+    def test_basic_flow_and_source_refs(self, monkeypatch):
+        _client_with(monkeypatch, lambda v, **k: [_hit("P1_0"), _hit("P1_1", ver="v2")])
+        out = rs.search_knowledge("基金申购费率")
+        assert len(out["results"]) == 2
+        # 同文档不同版本 → 溯源按 (doc, ver, pid) 去重，应为 2 条
+        assert len(out["source_refs"]) == 2
+        assert out["source_refs"][0]["source_doc_id"] == "KB-P1"
+
+    def test_blank_query_short_circuits(self, monkeypatch):
+        # 空白 query：不发 embedding、不连 Milvus
+        def bad_client():
+            raise AssertionError("空白 query 不应触发 Milvus 连接")
+
+        monkeypatch.setattr(rs, "_client", bad_client)
+        monkeypatch.setattr(rs, "_embed", lambda q: [] if False else (_ for _ in ()).throw(AssertionError("不应向量化")))
+        out = rs.search_knowledge("   ")
+        assert out == {"query": "", "results": [], "source_refs": []}
+
+    def test_filter_params_passthrough(self, monkeypatch):
+        client = _client_with(monkeypatch, lambda v, **k: [])
+        rs.search_knowledge("定投规则", product_id="PROD-9", doc_type="rule", top_k=5)
+        assert client.last_kwargs["product_id"] == "PROD-9"
+        assert client.last_kwargs["doc_type"] == "rule"
+        assert client.last_kwargs["top_k"] == 5
+
+    def test_query_whitespace_stripped(self, monkeypatch):
+        _client_with(monkeypatch, lambda v, **k: [])
+        out = rs.search_knowledge("  赎回到账时间  ")
+        assert out["query"] == "赎回到账时间"
+
+
+class TestSourceRefs:
+    def test_dedup_same_doc_version(self, monkeypatch):
+        _client_with(monkeypatch, lambda v, **k: [_hit("P1_0"), _hit("P1_1")])
+        out = rs.search_knowledge("费率")
+        assert len(out["source_refs"]) == 1
+
+    def test_multi_product_refs(self, monkeypatch):
+        _client_with(
+            monkeypatch,
+            lambda v, **k: [_hit("P1_0"), _hit("P2_0", doc="KB-P2", pid="PROD-2")],
+        )
+        out = rs.search_knowledge("费率")
+        assert len(out["source_refs"]) == 2
+        pids = {r["product_id"] for r in out["source_refs"]}
+        assert pids == {"PROD-1", "PROD-2"}
+
+
+class TestFailurePropagation:
+    """失败口径：异常原样上抛，不吞不降级（防止 LLM 编造回答）。"""
+
+    def test_embed_error_propagates(self, monkeypatch):
+        def boom(q):
+            raise EmbeddingError("Ollama 连接失败")
+
+        monkeypatch.setattr(rs, "_embed", boom)
+        with pytest.raises(EmbeddingError):
+            rs.search_knowledge("任意")
+
+    def test_milvus_error_propagates(self, monkeypatch):
+        def stub(vector, **kwargs):
+            raise RuntimeError("milvus down")
+
+        _client_with(monkeypatch, stub)
+        with pytest.raises(RuntimeError):
+            rs.search_knowledge("任意")