180 lines
7.4 KiB
Python
180 lines
7.4 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""端到端联调用例（Issue #86）—— 问答流链路。
|
|||
|
|
|
|||
|
|
验证 PRD 5.4（LLM 网关 + RAG）的端到端协作：
|
|||
|
|
rag-kb 模板化知识库检索（命中文档片段 + 来源）→ llm-gateway 混合网关
|
|||
|
|
（路由 → 生成 → 溯源校验 → DLP 出站防线）。
|
|||
|
|
|
|||
|
|
不依赖真实 LLM / 向量库：全部走内存桩，可在 CI 直接执行。
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
# 引导加载四个内核模块（必须在测试导入前执行）
|
|||
|
|
import tests.e2e._bootstrap # noqa: F401
|
|||
|
|
|
|||
|
|
import unittest
|
|||
|
|
|
|||
|
|
from rag_kb import RagKnowledgeBase, build_document
|
|||
|
|
from llm_gateway import (
|
|||
|
|
DLP_DEFAULT_RULES,
|
|||
|
|
CloudBackend,
|
|||
|
|
DlpEngine,
|
|||
|
|
GatewayResult,
|
|||
|
|
HallucinationGuard,
|
|||
|
|
LLMGateway,
|
|||
|
|
LocalBackend,
|
|||
|
|
PromptRegistry,
|
|||
|
|
RouteTarget,
|
|||
|
|
SensitivityRouter,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _build_kb() -> RagKnowledgeBase:
|
|||
|
|
"""构造一个模板化知识库：工艺规范 / SOP / 国标 三类知识源。"""
|
|||
|
|
kb = RagKnowledgeBase()
|
|||
|
|
docs = [
|
|||
|
|
build_document(
|
|||
|
|
title="氯化车间操作规程",
|
|||
|
|
text="氯化车间 1#炉正常运行温度区间 850~920℃，超过 950℃ 属于超温，"
|
|||
|
|
"应立即减少通氯量并检查冷却系统。停机检修须挂牌上锁。",
|
|||
|
|
category="process",
|
|||
|
|
),
|
|||
|
|
build_document(
|
|||
|
|
title="海绵钛氯化工序 SOP",
|
|||
|
|
text="氯化工序标准作业指导书：开机前确认氯气流量计归零，"
|
|||
|
|
"升温阶段按 50℃/h 速率升温至 850℃。异常停机时按紧急停机程序处置。",
|
|||
|
|
category="sop",
|
|||
|
|
),
|
|||
|
|
build_document(
|
|||
|
|
title="工业氯化工艺国家标准",
|
|||
|
|
text="GB/T XXXX 工业氯化工艺安全规范：氯化炉设计压力不低于 0.6MPa，"
|
|||
|
|
"操作人员持证上岗，关键参数实时记录留存不少于 3 年。",
|
|||
|
|
category="standard",
|
|||
|
|
),
|
|||
|
|
]
|
|||
|
|
kb.add_documents(docs)
|
|||
|
|
return kb
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _build_gateway() -> LLMGateway:
|
|||
|
|
"""构造一个可运行的混合网关：注册 qa 提示词 + 默认 DLP/路由/校验。"""
|
|||
|
|
prompts = PromptRegistry()
|
|||
|
|
prompts.update(
|
|||
|
|
name="qa",
|
|||
|
|
version="1.0.0",
|
|||
|
|
text="你是工业 AI 优化助手。基于知识库回答：{query}",
|
|||
|
|
description="问答主提示词 v1.0.0",
|
|||
|
|
)
|
|||
|
|
return LLMGateway(
|
|||
|
|
dlp=DlpEngine(),
|
|||
|
|
router=SensitivityRouter(),
|
|||
|
|
prompts=prompts,
|
|||
|
|
guard=HallucinationGuard(),
|
|||
|
|
local=LocalBackend(echo_context=True),
|
|||
|
|
cloud=CloudBackend(echo_context=True),
|
|||
|
|
prompt_name="qa",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class QueryPipelineE2ETest(unittest.TestCase):
|
|||
|
|
"""RAG 检索 → LLM 网关编排 全链路联调。"""
|
|||
|
|
|
|||
|
|
@classmethod
|
|||
|
|
def setUpClass(cls) -> None:
|
|||
|
|
cls.kb = _build_kb()
|
|||
|
|
cls.gateway = _build_gateway()
|
|||
|
|
|
|||
|
|
# -- RAG 检索 ----------------------------------------------------------
|
|||
|
|
def test_kb_retrieval_returns_relevant_chunks_with_source(self) -> None:
|
|||
|
|
"""检索命中文档片段并带来源（引用溯源基础）。"""
|
|||
|
|
hits = self.kb.search("氯化炉温度", top_k=3)
|
|||
|
|
self.assertGreater(len(hits), 0)
|
|||
|
|
# 命中片段必须携带来源信息（文档标题 / 类别）
|
|||
|
|
for hit in hits:
|
|||
|
|
self.assertIsNotNone(hit.chunk.title)
|
|||
|
|
self.assertIn(hit.chunk.category, ("process", "sop", "standard"))
|
|||
|
|
# 相关片段应命中"温度"相关内容
|
|||
|
|
joined = " ".join(h.chunk.text for h in hits)
|
|||
|
|
self.assertIn("温度", joined)
|
|||
|
|
|
|||
|
|
def test_kb_category_filter(self) -> None:
|
|||
|
|
"""类别过滤：仅检索 SOP 知识源。"""
|
|||
|
|
from rag_kb import KnowledgeSourceKind
|
|||
|
|
hits = self.kb.search("升温", top_k=5, categories=[KnowledgeSourceKind.SOP])
|
|||
|
|
self.assertGreater(len(hits), 0)
|
|||
|
|
for hit in hits:
|
|||
|
|
self.assertEqual(hit.chunk.category, "sop")
|
|||
|
|
|
|||
|
|
# -- 问答闭环（普通问题 → 本地后端）------------------------------------
|
|||
|
|
def test_normal_question_routes_local_with_rag_context(self) -> None:
|
|||
|
|
"""普通工艺问题：路由到本地后端，回答含 RAG 来源引用。"""
|
|||
|
|
query = "氯化车间 1#炉的正常运行温度是多少？"
|
|||
|
|
hits = self.kb.search(query, top_k=3)
|
|||
|
|
sources = [h.chunk.title for h in hits]
|
|||
|
|
|
|||
|
|
result = self.gateway.ask(query, rag_context=sources, confidence=0.95)
|
|||
|
|
|
|||
|
|
self.assertIsInstance(result, GatewayResult)
|
|||
|
|
self.assertEqual(result.route.target, RouteTarget.LOCAL)
|
|||
|
|
# 本地后端 echo_context 时输出含来源标记
|
|||
|
|
self.assertGreater(len(result.answer), 0)
|
|||
|
|
if sources:
|
|||
|
|
self.assertIn(sources[0], result.answer)
|
|||
|
|
# 溯源校验通过（有来源支撑）
|
|||
|
|
self.assertTrue(result.verdict.supported)
|
|||
|
|
self.assertFalse(result.needs_human)
|
|||
|
|
|
|||
|
|
# -- 敏感数据 DLP 拦截（fail-closed）-----------------------------------
|
|||
|
|
def test_sensitive_query_blocked_by_dlp(self) -> None:
|
|||
|
|
"""含敏感数据的问题被 DLP 拦截 → 路由 block → 转人工。"""
|
|||
|
|
# 身份证号（DLP 默认规则命中）
|
|||
|
|
query = "请查询员工 110101199003078834 的工资"
|
|||
|
|
result = self.gateway.ask(query, confidence=1.0)
|
|||
|
|
|
|||
|
|
self.assertEqual(result.route.target, RouteTarget.BLOCK)
|
|||
|
|
self.assertTrue(result.needs_human)
|
|||
|
|
# block 时不调用后端，给出人工确认占位
|
|||
|
|
self.assertIn("人工", result.answer)
|
|||
|
|
|
|||
|
|
# -- 脱敏/通用问题路由到云端 ------------------------------------------
|
|||
|
|
def test_generic_question_routes_cloud(self) -> None:
|
|||
|
|
"""通用（非敏感）问题经 DLP 放行后可路由云端（这里默认路由 local，
|
|||
|
|
需显式配置 cloud 规则才走云端）。验证默认 local 闭环正常。"""
|
|||
|
|
query = "今天的天气如何？"
|
|||
|
|
result = self.gateway.ask(query, confidence=0.9)
|
|||
|
|
# 默认无规则命中 → local
|
|||
|
|
self.assertEqual(result.route.target, RouteTarget.LOCAL)
|
|||
|
|
self.assertFalse(result.needs_human)
|
|||
|
|
|
|||
|
|
# -- 高利害低信度转人工 ------------------------------------------------
|
|||
|
|
def test_high_stakes_low_confidence_to_human(self) -> None:
|
|||
|
|
"""高利害提示词 + 低信度 → 转人工复核（幻觉防线）。"""
|
|||
|
|
prompts = PromptRegistry()
|
|||
|
|
prompts.update(name="alarm_explain", version="1.0.0",
|
|||
|
|
text="解释报警：{query}", description="报警解释（高利害）")
|
|||
|
|
gateway = LLMGateway(
|
|||
|
|
prompts=prompts,
|
|||
|
|
prompt_name="alarm_explain",
|
|||
|
|
high_stakes_names=["alarm_explain"],
|
|||
|
|
)
|
|||
|
|
result = gateway.ask("1#炉超温报警", rag_context=["氯化车间操作规程"],
|
|||
|
|
confidence=0.3) # 低信度
|
|||
|
|
self.assertTrue(result.needs_human)
|
|||
|
|
self.assertEqual(result.verdict.action, "human_review")
|
|||
|
|
|
|||
|
|
# -- 审计可追溯 --------------------------------------------------------
|
|||
|
|
def test_audit_drain_after_query(self) -> None:
|
|||
|
|
"""每次 ask() 后各组件审计记录可统一导出（DLP/路由/Prompt/幻觉）。"""
|
|||
|
|
self.gateway.ask("氯化炉温度区间", confidence=0.9)
|
|||
|
|
audits = self.gateway.drain_audits()
|
|||
|
|
self.assertIn("dlp", audits)
|
|||
|
|
self.assertIn("router", audits)
|
|||
|
|
self.assertIn("prompts", audits)
|
|||
|
|
self.assertIn("guard", audits)
|
|||
|
|
# 路由审计应有记录
|
|||
|
|
self.assertGreater(len(audits["router"]), 0)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
unittest.main()
|