132 lines
4.9 KiB
Python
132 lines
4.9 KiB
Python
# -*- 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()
|