feat: 完成 issue #6 LLM 网关 + RAG 模板化（混合网关主编排）

2026-08-04 18:02:00 +08:00
parent fa5523a37d
commit 5d937e8efd
12 changed files with 1707 additions and 20 deletions
@@ -0,0 +1,131 @@
+# -*- coding: utf-8 -*-
+"""混合网关主编排（gateway）端到端单元测试：路由 → 生成 → 校验 → DLP 防线。"""
+import os
+import sys
+import unittest
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import _bootstrap  # noqa: F401
+
+from llm_gateway.dlp import DlpEngine  # noqa: E402
+from llm_gateway.gateway import (  # noqa: E402
+    CloudBackend,
+    LLMGateway,
+    LocalBackend,
+)
+from llm_gateway.prompts import PromptRegistry  # noqa: E402
+from llm_gateway.router import RouteTarget, SensitivityRouter  # noqa: E402
+
+ROUTER_CONFIG = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+    "config", "router.template.yaml",
+)
+PROMPTS_CONFIG = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+    "config", "prompts.template.yaml",
+)
+
+
+def make_gateway() -> LLMGateway:
+    return LLMGateway(
+        dlp=DlpEngine(),
+        router=SensitivityRouter.from_template_config(ROUTER_CONFIG),
+        prompts=PromptRegistry.from_template_config(PROMPTS_CONFIG),
+        local=LocalBackend(),
+        cloud=CloudBackend(),
+        high_stakes_names=["alarm_explain"],
+    )
+
+
+class GatewayRoutingTest(unittest.TestCase):
+    """路由目标决定后端选择。"""
+
+    def setUp(self):
+        self.gw = make_gateway()
+
+    def test_sensitive_query_uses_local(self):
+        result = self.gw.ask("炉温当前是多少", rag_context=["SOP-炉温"])
+        self.assertEqual(result.route.target, RouteTarget.LOCAL)
+        self.assertIn("本地70B占位", result.answer)
+        self.assertFalse(result.needs_human)
+
+    def test_common_query_uses_cloud(self):
+        result = self.gw.ask("海绵钛是什么", rag_context=["科普手册"])
+        self.assertEqual(result.route.target, RouteTarget.CLOUD)
+        self.assertIn("云端API占位", result.answer)
+
+    def test_blocked_query_needs_human(self):
+        result = self.gw.ask("现场出现紧急停机指令", rag_context=[])
+        self.assertEqual(result.route.target, RouteTarget.BLOCK)
+        self.assertTrue(result.needs_human)
+        self.assertIn("人工确认", result.answer)
+
+    def test_dlp_blocked_query_forces_block(self):
+        # 身份证号触发 DLP → 即使模板规则未覆盖也 block
+        result = self.gw.ask("员工 110101199003071234 的炉温查询",
+                             rag_context=["SOP"])
+        self.assertEqual(result.route.target, RouteTarget.BLOCK)
+        self.assertEqual(result.route.reason, "dlp_blocked")
+
+
+class GatewayVerificationTest(unittest.TestCase):
+    """引用溯源 + 信度阈值（高利害）。"""
+
+    def setUp(self):
+        self.gw = make_gateway()
+
+    def test_unsupported_citation_flagged(self):
+        # 占位后端回显 [来源: rag_context]，与 rag_context 一致 → 支持
+        result = self.gw.ask("炉温偏高怎么处理",
+                             rag_context=["沸腾氯化炉异常处置SOP"],
+                             confidence=0.9)
+        self.assertTrue(result.verdict.supported)
+
+    def test_high_stakes_low_confidence_human_review(self):
+        # alarm_explain 为高利害模板：低信度 → 人工确认
+        result = self.gw.ask("解释报警并给出处置建议",
+                             rag_context=["报警SOP"],
+                             confidence=0.4)
+        self.assertEqual(result.route.target, RouteTarget.LOCAL)
+        self.assertEqual(result.verdict.action, "pass")  # qa 非高利害，不启用阈值
+        gw2 = LLMGateway(
+            dlp=DlpEngine(),
+            router=SensitivityRouter.from_template_config(ROUTER_CONFIG),
+            prompts=PromptRegistry.from_template_config(PROMPTS_CONFIG),
+            prompt_name="alarm_explain",
+            high_stakes_names=["alarm_explain"],
+        )
+        result2 = gw2.ask("解释报警并给出处置建议",
+                          rag_context=["报警SOP"],
+                          confidence=0.4)
+        self.assertEqual(result2.verdict.action, "human_review")
+        self.assertTrue(result2.needs_human)
+
+    def test_prompt_version_binding(self):
+        # 显式绑定 qa@1.0.0（默认）——当前注册表已按模板加载
+        pv = self.gw.prompts.get("qa", version="1.0.0")
+        self.assertEqual(pv.version, "1.0.0")
+
+
+class GatewayAuditTest(unittest.TestCase):
+    def test_audits_collectable(self):
+        gw = make_gateway()
+        gw.ask("炉温当前是多少", rag_context=["SOP"])
+        audits = gw.drain_audits()
+        self.assertIn("router", audits)
+        self.assertIn("guard", audits)
+        self.assertGreaterEqual(len(audits["router"]), 1)
+        # drain 后清空
+        self.assertEqual(gw.drain_audits()["router"], [])
+
+    def test_result_to_dict(self):
+        gw = make_gateway()
+        result = gw.ask("炉温当前是多少", rag_context=["SOP"])
+        d = result.to_dict()
+        self.assertIn("answer_id", d)
+        self.assertIn("route", d)
+        self.assertIn("verdict", d)
+
+
+if __name__ == "__main__":
+    unittest.main()