iAOP/tests/e2e/test_query_pipeline.py

# -*- 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()