diff --git a/app/service/rag_service.py b/app/service/rag_service.py index 18ca77b..a575946 100644 --- a/app/service/rag_service.py +++ b/app/service/rag_service.py @@ -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)} diff --git a/tests/test_rag_service.py b/tests/test_rag_service.py new file mode 100644 index 0000000..a2ca86c --- /dev/null +++ b/tests/test_rag_service.py @@ -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("任意")