Agentic_Rag/app/decision_engine.py

from __future__ import annotations

import json
import re
from typing import Protocol

from app.config import Settings
from app.schemas import (
    Evidence,
    PlanStep,
    QualityGrade,
    RetrievalPlan,
    RouteDecision,
    RouteName,
)


USER_ID_PATTERN = re.compile(r"\bU\d{4}\b", re.IGNORECASE)
PRODUCT_NAMES = ("XPhone 15 Pro", "AirBuds Pro", "USB-C Charger", "SmartWatch S")
POLICY_KEYWORDS = {
    "return_policy": ("退货", "退款", "七日无理由", "换货", "售后"),
    "shipping_policy": ("运费", "包邮", "配送", "物流", "送达"),
    "price_protection": ("价保", "保价", "差价", "降价"),
}


def extract_user_id(query: str) -> str:
    match = USER_ID_PATTERN.search(query.upper())
    return match.group(0) if match else ""


def extract_product_name(query: str) -> str:
    lowered = query.lower()
    return next((name for name in PRODUCT_NAMES if name.lower() in lowered), "")


def extract_policy_type(query: str) -> str:
    for policy_type, keywords in POLICY_KEYWORDS.items():
        if any(keyword in query for keyword in keywords):
            return policy_type
    return ""


class DecisionEngine(Protocol):
    def route(self, query: str, max_rounds: int) -> RouteDecision: ...

    def plan(self, query: str, decision: RouteDecision) -> RetrievalPlan: ...

    def grade(
        self,
        query: str,
        decision: RouteDecision,
        evidence: list[Evidence],
        current_round: int,
        max_rounds: int,
        min_score: float,
    ) -> QualityGrade: ...

    def rewrite(self, query: str, grade: QualityGrade) -> str: ...

    def answer(self, query: str, evidence: list[Evidence], partial: bool = False) -> str: ...


class DemoDecisionEngine:
    """确定性演示引擎：便于无 API Key 启动、调试和自动化测试。"""

    SQL_TERMS = ("订单", "几单", "多少单", "实付", "总额", "金额", "统计", "最近")
    DOC_TERMS = ("规则", "政策", "退货", "退款", "运费", "包邮", "价保", "售后")
    WEB_TERMS = ("网页", "官网", "公开资料", "最新公告", "新闻", "互联网")
    DIRECT_TERMS = ("改写", "润色", "翻译", "缩短", "转成")

    def route(self, query: str, max_rounds: int) -> RouteDecision:
        stripped = query.strip()
        user_id = extract_user_id(stripped)
        product_name = extract_product_name(stripped)
        policy_type = extract_policy_type(stripped)

        # 安全门控和缺失槽位检查必须先于数据源选择，避免先访问再拒绝。
        if any(term in stripped for term in ("所有买家手机号", "密码", "支付密钥")):
            return RouteDecision(
                needs_retrieval=False,
                intent="restricted_request",
                routes=[RouteName.REFUSE],
                confidence="high",
                reason_code="SECURITY_POLICY",
                max_rounds=max_rounds,
            )

        if any(term in stripped for term in self.DIRECT_TERMS):
            return RouteDecision(
                needs_retrieval=False,
                intent="text_transformation",
                routes=[RouteName.DIRECT_ANSWER],
                confidence="high",
                reason_code="NO_EXTERNAL_KNOWLEDGE_REQUIRED",
                max_rounds=max_rounds,
            )

        needs_sql = any(term in stripped for term in self.SQL_TERMS)
        needs_docs = any(term in stripped for term in self.DOC_TERMS)
        needs_web = any(term in stripped for term in self.WEB_TERMS)
        if needs_sql and not user_id:
            return RouteDecision(
                needs_retrieval=False,
                intent="missing_authenticated_user",
                routes=[RouteName.CLARIFY],
                confidence="high",
                reason_code="MISSING_USER_ID",
                max_rounds=max_rounds,
            )
        if "这件商品" in stripped and not product_name:
            return RouteDecision(
                needs_retrieval=False,
                intent="missing_product",
                routes=[RouteName.CLARIFY],
                confidence="high",
                reason_code="MISSING_PRODUCT_NAME",
                max_rounds=max_rounds,
            )

        filters = {
            key: value
            for key, value in {
                "user_id": user_id,
                "product_name": product_name,
                "policy_type": policy_type,
            }.items()
            if value
        }

        # routes 是可并行的数据源集合，不使用单标签覆盖复合意图。
        routes: list[RouteName] = []
        if needs_docs:
            routes.append(RouteName.MILVUS_SEARCH)
        if needs_sql:
            routes.append(RouteName.SQL_QUERY)
        if needs_web:
            routes.append(RouteName.WEB_SEARCH)

        if len(routes) > 1:
            return RouteDecision(
                needs_retrieval=True,
                intent="multi_source_research",
                routes=routes,
                requires_decomposition=True,
                confidence="high" if not needs_sql or user_id else "medium",
                reason_code="MULTIPLE_DATA_SOURCES_REQUIRED",
                filters=filters,
                max_rounds=max_rounds,
            )
        if needs_sql:
            return RouteDecision(
                needs_retrieval=True,
                intent="order_statistics",
                routes=[RouteName.SQL_QUERY],
                confidence="high",
                reason_code="EXACT_AGGREGATION_REQUIRED",
                filters=filters,
                max_rounds=max_rounds,
            )
        if needs_web:
            return RouteDecision(
                needs_retrieval=True,
                intent="public_web_research",
                routes=[RouteName.WEB_SEARCH],
                confidence="high",
                reason_code="CURRENT_PUBLIC_INFORMATION_REQUIRED",
                filters=filters,
                max_rounds=max_rounds,
            )

        return RouteDecision(
            needs_retrieval=True,
            intent="commerce_policy_qa",
            routes=[RouteName.MILVUS_SEARCH],
            confidence="high" if policy_type else "medium",
            reason_code="INTERNAL_DOCUMENT_REQUIRED",
            filters=filters,
            max_rounds=max_rounds,
        )

    def plan(self, query: str, decision: RouteDecision) -> RetrievalPlan:
        user_id = str(decision.filters.get("user_id", ""))
        product_name = str(decision.filters.get("product_name", ""))
        policy_type = str(decision.filters.get("policy_type", ""))
        steps: list[PlanStep] = []
        # 当前三类读取任务互不依赖，因此不填写 depends_on，由 Executor 并行执行。
        if RouteName.MILVUS_SEARCH in decision.routes:
            steps.append(
                PlanStep(
                    id="search_policy",
                    tool=RouteName.MILVUS_SEARCH,
                    query=query,
                    arguments={"policy_type": policy_type, "top_k": 4},
                )
            )
        if RouteName.SQL_QUERY in decision.routes:
            steps.append(
                PlanStep(
                    id="query_statistics",
                    tool=RouteName.SQL_QUERY,
                    query=(
                        f"统计用户 {user_id} 最近 30 天"
                        f"{product_name or '全部商品'}的有效订单"
                    ),
                    arguments={
                        "user_id": user_id,
                        "days": 30,
                        "product_keyword": product_name,
                    },
                )
            )
        if RouteName.WEB_SEARCH in decision.routes:
            steps.append(
                PlanStep(
                    id="search_web",
                    tool=RouteName.WEB_SEARCH,
                    query=query,
                    arguments={"max_results": 5},
                )
            )
        return RetrievalPlan(goal=query, steps=steps)

    def grade(
        self,
        query: str,
        decision: RouteDecision,
        evidence: list[Evidence],
        current_round: int,
        max_rounds: int,
        min_score: float,
    ) -> QualityGrade:
        # 先做确定性的来源覆盖检查，再决定是否需要模型参与语义评分。
        required_sources = set()
        if RouteName.MILVUS_SEARCH in decision.routes:
            required_sources.add("milvus")
        if RouteName.SQL_QUERY in decision.routes:
            required_sources.add("sql")
        if RouteName.WEB_SEARCH in decision.routes:
            required_sources.add("web")

        available_sources = {item.source_type for item in evidence}
        missing = sorted(required_sources - available_sources)
        vector_evidence = [item for item in evidence if item.source_type == "milvus"]
        vector_relevant = any(
            item.score is None or item.score >= min_score for item in vector_evidence
        )
        if "milvus" not in required_sources:
            vector_relevant = True

        if not missing and vector_relevant:
            return QualityGrade(
                relevant=True,
                sufficient=True,
                recommended_action="accept",
                reason="所需数据源均返回证据，向量结果达到最低质量门槛。",
            )

        # 只有预算尚未耗尽时才能 Rewrite；否则必须有界终止。
        if current_round < max_rounds:
            return QualityGrade(
                relevant=bool(evidence),
                sufficient=False,
                missing_aspects=missing or ["高相关技术文档"],
                recommended_action="rewrite_query",
                reason="证据不完整或向量结果相关性不足，允许改写后重试。",
            )

        return QualityGrade(
            relevant=bool(evidence),
            sufficient=False,
            missing_aspects=missing or ["高相关技术文档"],
            recommended_action="stop",
            reason="达到最大检索轮数，停止继续调用工具。",
        )

    def rewrite(self, query: str, grade: QualityGrade) -> str:
        missing = "、".join(grade.missing_aspects)
        suffix = f"电商平台内部政策 订单范围 官网公告 {missing}".strip()
        return f"{query}；补充检索条件：{suffix}"

    def answer(self, query: str, evidence: list[Evidence], partial: bool = False) -> str:
        if not evidence:
            if any(term in query for term in self.DIRECT_TERMS):
                return "该请求不需要外部检索，应由文本处理模型直接完成。"
            return "未获得能够支持回答的证据。"

        prefix = "当前仅能给出部分结论。" if partial else "基于已检索证据，结论如下。"
        lines = [prefix]
        for index, item in enumerate(evidence, start=1):
            lines.append(f"- {item.content} [{index}]")
        return "\n".join(lines)


class DeepSeekDecisionEngine(DemoDecisionEngine):
    """使用 DeepSeek 的 OpenAI 兼容接口完成路由、评分和答案生成。"""

    def __init__(self, settings: Settings) -> None:
        # 凭证只从 Settings 读取，并在创建 SDK 客户端时短暂解密。
        api_key = settings.deepseek_api_key.get_secret_value()
        if not api_key:
            raise RuntimeError("LLM_PROVIDER=deepseek 时必须配置 DEEPSEEK_API_KEY")
        try:
            from langchain_openai import ChatOpenAI
        except ImportError as exc:
            raise RuntimeError(
                "缺少 DeepSeek 依赖，请执行 uv sync --extra test --extra deepseek"
            ) from exc

        common_kwargs = {
            "model": settings.deepseek_model_name,
            "api_key": api_key,
            "base_url": settings.deepseek_base_url,
            "max_retries": 2,
        }

        # 最终答案可以使用 Thinking；开关由 .env 控制。
        answer_thinking = "enabled" if settings.deepseek_answer_thinking else "disabled"
        self.llm = ChatOpenAI(
            **common_kwargs,
            extra_body={"thinking": {"type": answer_thinking}},
        )

        # with_structured_output(function_calling) 会发送强制 tool_choice。
        # DeepSeek Thinking 模式不接受该组合，因此路由与评分必须关闭 Thinking。
        self.structured_llm = ChatOpenAI(
            **common_kwargs,
            temperature=0,
            extra_body={"thinking": {"type": "disabled"}},
        )
        self.router = self.structured_llm.with_structured_output(
            RouteDecision,
            method="function_calling",
        )
        self.grader = self.structured_llm.with_structured_output(
            QualityGrade,
            method="function_calling",
        )

    def route(self, query: str, max_rounds: int) -> RouteDecision:
        prompt = f"""
你是电商售后系统的路由器。只输出符合 Schema 的结果。
可选路径：direct_answer、milvus_search、sql_query、web_search、clarify、refuse。
内部退货、运费和价保政策使用 milvus_search；订单数量和金额使用 sql_query；
官网、新闻、最新公告和公开资料使用 web_search；
多类需求同时存在时返回对应的多条路径，并设置 requires_decomposition=true。
订单统计缺少受控用户编号时选择 clarify，敏感越权请求选择 refuse。
filters 中只填写已从原问题明确提取的 user_id、product_name、policy_type。
最大检索轮数：{max_rounds}
用户问题：{query}
""".strip()
        result = self.router.invoke(prompt)
        # 部分兼容接口可能返回普通文本而非 Function Call，LangChain 此时返回 None。
        # 路由属于关键控制决策，结构化输出缺失时回退到确定性规则，不允许继续空值。
        if result is None:
            return super().route(query, max_rounds)
        # 关键业务槽位使用确定性解析覆盖模型结果，避免模型改写用户身份。
        deterministic_filters = {
            key: value
            for key, value in {
                "user_id": extract_user_id(query),
                "product_name": extract_product_name(query),
                "policy_type": extract_policy_type(query),
            }.items()
            if value
        }
        result.filters = {**result.filters, **deterministic_filters}
        if RouteName.SQL_QUERY in result.routes and "user_id" not in result.filters:
            return RouteDecision(
                needs_retrieval=False,
                intent="missing_authenticated_user",
                routes=[RouteName.CLARIFY],
                confidence="high",
                reason_code="MISSING_USER_ID",
                max_rounds=max_rounds,
            )
        result.max_rounds = max_rounds
        return result

    def grade(
        self,
        query: str,
        decision: RouteDecision,
        evidence: list[Evidence],
        current_round: int,
        max_rounds: int,
        min_score: float,
    ) -> QualityGrade:
        prompt = f"""
判断证据是否足以回答问题。recommended_action 只能是 accept、rewrite_query 或 stop。
当前轮次：{current_round}/{max_rounds}
问题：{query}
路由：{decision.model_dump_json()}
证据：{json.dumps([item.model_dump() for item in evidence], ensure_ascii=False)}
""".strip()
        result = self.grader.invoke(prompt)
        # Grader 输出为空时使用确定性来源覆盖、相似度和轮次判断兜底。
        if result is None:
            return super().grade(
                query=query,
                decision=decision,
                evidence=evidence,
                current_round=current_round,
                max_rounds=max_rounds,
                min_score=min_score,
            )
        return result

    def answer(self, query: str, evidence: list[Evidence], partial: bool = False) -> str:
        prompt = f"""
严格依据证据回答，不得补充证据之外的事实。每个关键结论使用 [序号] 引用。
问题：{query}
是否为部分证据：{partial}
证据：{json.dumps([item.model_dump() for item in evidence], ensure_ascii=False)}
""".strip()
        return str(self.llm.invoke(prompt).content)


def create_decision_engine(settings: Settings) -> DecisionEngine:
    if settings.llm_provider == "demo":
        return DemoDecisionEngine()
    if settings.llm_provider == "deepseek":
        return DeepSeekDecisionEngine(settings)
    raise ValueError(f"不支持的 LLM_PROVIDER: {settings.llm_provider}")