rag/knowledge/router.py

"""
知识库路由器 - 智能选择查询目标

功能：
1. 查询意图分析 - 判断查询是否涉及特定部门
2. 知识库路由 - 根据意图和权限选择目标向量库
3. 单库优化 - 如果只需查询单库，避免不必要的并行检索

使用方式：
    from knowledge.router import KnowledgeBaseRouter

    router = KnowledgeBaseRouter()

    # 获取目标向量库
    target_kbs = router.route(
        query="财务部的报销流程是什么",
        role="user",
        department="tech"
    )
    # 返回: ["public_kb", "dept_finance"]  # 如果有权限
"""

import os
import re
import json
import logging
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass
from openai import OpenAI

# 导入配置
from config import API_KEY, BASE_URL, MODEL

# 导入 LLM 工具函数
from core.llm_utils import call_llm, parse_json_from_response

# 导入权限管理
from auth.gateway import get_accessible_collections, normalize_department_name

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


# ==================== 部门关键词配置 ====================

# 部门关键词映射（可根据实际情况扩展）
DEPARTMENT_KEYWORDS = {
    "finance": [
        "财务", "报销", "发票", "预算", "支出", "收入", "成本",
        "账目", "会计", "审计", "税务", "工资", "奖金", "补贴",
        "费用", "付款", "收款", "借款", "报销单", "财务部"
    ],
    "hr": [
        "人事", "招聘", "入职", "离职", "考勤", "请假", "休假",
        "员工", "培训", "绩效", "晋升", "调岗", "合同", "档案",
        "社保", "公积金", "福利", "加班", "年假", "人事部", "人力资源"
    ],
    "tech": [
        "技术", "开发", "代码", "系统", "服务器", "数据库", "API",
        "接口", "部署", "测试", "Bug", "需求", "架构", "运维",
        "网络安全", "服务器", "云服务", "技术部", "研发", "IT"
    ],
    "operation": [
        "运营", "推广", "营销", "活动", "用户", "增长", "数据",
        "分析", "客服", "售后", "投诉", "反馈", "运营部", "运营中心"
    ],
    "marketing": [
        "市场", "品牌", "宣传", "广告", "公关", "媒体", "推广",
        "展会", "活动策划", "市场部", "营销部"
    ],
    "legal": [
        "法务", "合同", "法律", "诉讼", "合规", "风险", "版权",
        "知识产权", "协议", "法务部"
    ],
    "admin": [
        "行政", "办公室", "会议室", "采购", "固定资产", "办公用品",
        "印章", "档案", "行政部", "总务"
    ]
}

# 通用关键词（查询 public_kb）
GENERAL_KEYWORDS = [
    "公司", "企业", "组织", "介绍", "简介", "文化", "价值观",
    "制度", "规定", "流程", "政策", "手册", "指南", "帮助",
    "联系方式", "地址", "电话", "邮箱"
]


# ==================== 数据结构 ====================

@dataclass
class QueryIntent:
    """查询意图"""
    is_general: bool              # 是否为通用问题
    department: Optional[str]      # 涉及的部门（如果有）
    confidence: float             # 置信度
    keywords: List[str]           # 匹配到的关键词
    reason: str                   # 判断理由


# ==================== 知识库路由器 ====================

class KnowledgeBaseRouter:
    """
    知识库路由器

    根据查询内容和用户权限，智能选择需要查询的向量库。
    支持规则匹配和 LLM 意图分析两种方式。
    """

    def __init__(self, use_llm: bool = True):
        """
        初始化

        Args:
            use_llm: 是否使用 LLM 进行意图分析（更准确但更慢）
        """
        self.use_llm = use_llm
        self.llm_client = None

        if use_llm:
            try:
                self.llm_client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
                logger.info("LLM 客户端初始化成功，将使用 LLM 进行意图分析")
            except Exception as e:
                logger.warning(f"LLM 客户端初始化失败: {e}，将使用规则匹配")
                self.use_llm = False

    def route(
        self,
        query: str,
        role: str,
        department: str,
        accessible_collections: List[str] = None
    ) -> List[str]:
        """
        根据查询意图和用户权限，决定查询哪些向量库

        Args:
            query: 用户查询
            role: 用户角色
            department: 用户部门
            accessible_collections: 可访问的向量库列表（可选）

        Returns:
            需要查询的向量库名称列表
        """
        # 1. 获取可访问的向量库
        if accessible_collections is None:
            accessible_collections = get_accessible_collections(role, department)

        if not accessible_collections:
            logger.warning(f"用户无可访问的向量库: role={role}, dept={department}")
            return []

        # 2. 分析查询意图
        intent = self.analyze_intent(query)

        # 3. 根据意图选择目标库
        target_kbs = self._select_knowledge_bases(
            intent, accessible_collections, role, department
        )

        logger.info(
            f"路由决策: query='{query[:30]}...', "
            f"intent={intent.department or 'general'}, "
            f"targets={target_kbs}"
        )

        return target_kbs

    def analyze_intent(self, query: str) -> QueryIntent:
        """
        分析查询意图

        Args:
            query: 用户查询

        Returns:
            QueryIntent 对象
        """
        # 先尝试规则匹配（快速）
        rule_intent = self._analyze_by_rules(query)

        # 如果规则匹配置信度高，直接返回
        if rule_intent.confidence > 0.8:
            return rule_intent

        # 否则使用 LLM 分析（更准确）
        if self.use_llm and self.llm_client:
            llm_intent = self._analyze_by_llm(query)
            if llm_intent:
                # 取两者中置信度高的
                return llm_intent if llm_intent.confidence > rule_intent.confidence else rule_intent

        return rule_intent

    def _analyze_by_rules(self, query: str) -> QueryIntent:
        """基于规则的意图分析"""
        query_lower = query.lower()
        matched_departments = {}
        matched_general = []

        # 检查部门关键词
        for dept, keywords in DEPARTMENT_KEYWORDS.items():
            for keyword in keywords:
                if keyword in query_lower:
                    if dept not in matched_departments:
                        matched_departments[dept] = []
                    matched_departments[dept].append(keyword)

        # 检查通用关键词
        for keyword in GENERAL_KEYWORDS:
            if keyword in query_lower:
                matched_general.append(keyword)

        # 判断结果
        if matched_departments:
            # 找到匹配最多的部门
            best_dept = max(
                matched_departments.keys(),
                key=lambda d: len(matched_departments[d])
            )
            keywords = matched_departments[best_dept]
            confidence = min(0.9, 0.5 + len(keywords) * 0.1)

            return QueryIntent(
                is_general=False,
                department=best_dept,
                confidence=confidence,
                keywords=keywords,
                reason=f"匹配到部门关键词: {', '.join(keywords)}"
            )

        elif matched_general:
            return QueryIntent(
                is_general=True,
                department=None,
                confidence=0.7,
                keywords=matched_general,
                reason=f"匹配到通用关键词: {', '.join(matched_general)}"
            )

        else:
            return QueryIntent(
                is_general=False,
                department=None,
                confidence=0.3,
                keywords=[],
                reason="未匹配到关键词，需要查询所有可访问的库"
            )

    def _analyze_by_llm(self, query: str) -> Optional[QueryIntent]:
        """使用 LLM 进行意图分析"""
        prompt = f"""分析以下问题的意图，判断：

1. 是否为通用问题（涉及公司整体、产品、文化等，不特指某部门）
2. 是否涉及特定部门（财务、人事、技术等）

问题：{query}

请直接返回 JSON 格式（不要包含其他内容）：
{{"is_general": true/false, "department": "部门英文名或null", "confidence": 0.0-1.0}}

部门英文名对照：
- finance: 财务
- hr: 人事
- tech: 技术
- operation: 运营
- marketing: 市场
- legal: 法务
- admin: 行政

注意：
- 如果问题涉及多个部门，返回 null
- 如果问题明显指向某个部门，返回对应英文名
- confidence 表示判断置信度，0-1之间"""

        content = call_llm(
            self.llm_client, prompt, MODEL,
            temperature=0.1,
            max_tokens=512
        )

        if content is None:
            logger.warning("LLM 意图分析失败: 调用返回空")
            return None

        # 使用 parse_json_from_response 解析 JSON
        result = parse_json_from_response(content)

        if result is None:
            logger.warning(f"LLM 意图分析失败: JSON 解析失败，原始内容: {content[:100]}")
            return None

        return QueryIntent(
            is_general=result.get("is_general", False),
            department=result.get("department"),
            confidence=result.get("confidence", 0.5),
            keywords=[],
            reason="LLM 意图分析"
        )

    def _select_knowledge_bases(
        self,
        intent: QueryIntent,
        accessible_collections: List[str],
        role: str,
        department: str
    ) -> List[str]:
        """
        选择要查询的知识库

        Args:
            intent: 查询意图
            accessible_collections: 可访问的向量库
            role: 用户角色
            department: 用户部门

        Returns:
            目标向量库列表
        """
        result = []
        public_kb = "public_kb"

        # 通用问题：优先查 public_kb
        if intent.is_general:
            if public_kb in accessible_collections:
                result.append(public_kb)
            # 但也可能需要查其他库（取决于置信度）
            if intent.confidence < 0.7:
                result.extend([kb for kb in accessible_collections if kb not in result])

        # 涉及特定部门
        elif intent.department:
            dept_kb = f"dept_{intent.department}"

            # 检查是否有权限访问该部门
            if dept_kb in accessible_collections:
                result.append(dept_kb)
                # 也查 public_kb（可能有相关政策）
                if public_kb in accessible_collections and public_kb not in result:
                    result.append(public_kb)
            else:
                # 没有权限访问目标部门，查 public_kb
                if public_kb in accessible_collections:
                    result.append(public_kb)
                logger.info(
                    f"用户无权访问部门 {intent.department} 的知识库，"
                    f"只查 public_kb"
                )

        # 未识别意图
        else:
            # admin 查所有
            if role == "admin":
                result = accessible_collections
            # 其他用户查 public 和本部门
            else:
                if public_kb in accessible_collections:
                    result.append(public_kb)
                # 使用标准化的部门名称
                normalized_dept = normalize_department_name(department)
                if normalized_dept:
                    user_dept_kb = f"dept_{normalized_dept}"
                    if user_dept_kb in accessible_collections and user_dept_kb not in result:
                        result.append(user_dept_kb)

        # 去重并保持顺序
        seen = set()
        unique_result = []
        for kb in result:
            if kb not in seen:
                seen.add(kb)
                unique_result.append(kb)

        return unique_result

    def get_routing_stats(self) -> Dict:
        """获取路由统计信息（用于监控）"""
        return {
            "use_llm": self.use_llm,
            "department_keywords": {
                dept: len(keywords)
                for dept, keywords in DEPARTMENT_KEYWORDS.items()
            },
            "general_keywords_count": len(GENERAL_KEYWORDS)
        }

    # ==================== 版本感知检索 ====================

    def route_with_version_awareness(
        self,
        query: str,
        role: str,
        department: str,
        accessible_collections: List[str] = None,
        include_deprecated: bool = False,
        top_k: int = 5
    ) -> Dict:
        """
        版本感知的路由

        在普通路由基础上，额外查询已废止的相关文档，
        为用户提供版本提示。

        Args:
            query: 用户查询
            role: 用户角色
            department: 用户部门
            accessible_collections: 可访问的向量库列表
            include_deprecated: 是否包含废止版本在结果中
            top_k: 返回数量

        Returns:
            {
                "target_collections": ["public_kb", "dept_finance"],
                "version_hints": [
                    {
                        "document": "报销制度.pdf",
                        "status": "deprecated",
                        "message": "该文档已于2026-03-01废止"
                    }
                ]
            }
        """
        # 1. 获取目标向量库（复用现有逻辑）
        target_kbs = self.route(query, role, department, accessible_collections)

        if not target_kbs:
            return {
                "target_collections": [],
                "version_hints": []
            }

        # 2. 查询是否有相关的废止版本
        version_hints = []
        if not include_deprecated:
            version_hints = self._find_deprecated_versions(query, target_kbs, top_k=3)

        logger.info(
            f"版本感知路由: query='{query[:30]}...', "
            f"targets={target_kbs}, hints={len(version_hints)}"
        )

        return {
            "target_collections": target_kbs,
            "version_hints": version_hints
        }

    def _find_deprecated_versions(
        self,
        query: str,
        collections: List[str],
        top_k: int = 3
    ) -> List[Dict]:
        """
        查找与查询相关的已废止版本

        Args:
            query: 用户查询
            collections: 目标向量库列表
            top_k: 每个库返回数量

        Returns:
            已废止版本提示列表
        """
        try:
            from knowledge.manager import get_kb_manager

            kb_manager = get_kb_manager()

            # 获取查询向量
            query_vector = self._get_query_vector(query)
            if query_vector is None:
                return []

            # 使用知识库管理器查找废止版本
            hints = kb_manager.find_deprecated_versions(
                kb_names=collections,
                query_vector=query_vector,
                top_k=top_k
            )

            # 去重（同一文档只提示一次）
            seen_docs = set()
            unique_hints = []
            for hint in hints:
                doc_key = f"{hint['collection']}/{hint['document']}"
                if doc_key not in seen_docs:
                    seen_docs.add(doc_key)
                    unique_hints.append(hint)

            return unique_hints

        except Exception as e:
            logger.warning(f"查找废止版本失败: {e}")
            return []

    def _get_query_vector(self, query: str) -> Optional[List[float]]:
        """
        获取查询向量

        Args:
            query: 查询文本

        Returns:
            查询向量，失败返回None
        """
        try:
            # 尝试使用 RAGEngine 的 embedding_model
            from core.engine import get_engine
            embedding_model = get_engine().embedding_model
            return embedding_model.encode(query).tolist()
        except Exception as e:
            logger.debug(f"无法从 RAGEngine 获取向量模型: {e}")

        try:
            # 尝试使用 sentence-transformers
            from sentence_transformers import SentenceTransformer
            model = SentenceTransformer('BAAI/bge-base-zh-v1.5')
            return model.encode(query).tolist()
        except Exception as e:
            logger.debug(f"Embedding 编码失败: {e}")

        logger.warning("无法加载向量模型，跳过废止版本检测")
        return None

    def search_with_version_context(
        self,
        query: str,
        role: str,
        department: str,
        top_k: int = 5
    ) -> Dict:
        """
        带版本上下文的搜索

        执行完整搜索流程：
        1. 版本感知路由
        2. 执行检索（只返回生效版本）
        3. 返回结果 + 废止版本提示

        Args:
            query: 用户查询
            role: 用户角色
            department: 用户部门
            top_k: 返回数量

        Returns:
            {
                "results": [...],  # 生效版本的检索结果
                "version_hints": [...],  # 废止版本提示
                "target_collections": [...]
            }
        """
        from knowledge.manager import get_kb_manager

        kb_manager = get_kb_manager()

        # 1. 版本感知路由
        route_result = self.route_with_version_awareness(
            query, role, department, include_deprecated=False
        )

        target_kbs = route_result["target_collections"]
        version_hints = route_result["version_hints"]

        if not target_kbs:
            return {
                "results": [],
                "version_hints": version_hints,
                "target_collections": []
            }

        # 2. 执行检索（只返回生效版本）
        query_vector = self._get_query_vector(query)
        if query_vector is None:
            return {
                "results": [],
                "version_hints": version_hints,
                "target_collections": target_kbs
            }

        # 多库检索，只返回active状态的chunks
        search_result = kb_manager.search_multiple(
            kb_names=target_kbs,
            query_vector=query_vector,
            query_text=query,
            top_k=top_k,
            use_bm25=True
        )

        # 过滤只返回active状态的chunks
        active_results = []
        if search_result.ids:
            for i, (doc_id, doc, meta, score) in enumerate(zip(
                search_result.ids,
                search_result.documents,
                search_result.metadatas,
                search_result.distances
            )):
                if meta.get("status", "active") == "active":
                    active_results.append({
                        "id": doc_id,
                        "document": doc,
                        "metadata": meta,
                        "score": score
                    })

        return {
            "results": active_results[:top_k],
            "version_hints": version_hints,
            "target_collections": target_kbs
        }


# ==================== 全局实例 ====================

_kb_router: Optional[KnowledgeBaseRouter] = None


def get_kb_router() -> KnowledgeBaseRouter:
    """获取全局知识库路由器实例"""
    global _kb_router
    if _kb_router is None:
        _kb_router = KnowledgeBaseRouter()
    return _kb_router


# ==================== 便捷函数 ====================

def route_query(
    query: str,
    role: str,
    department: str,
    accessible_collections: List[str] = None
) -> List[str]:
    """
    路由查询到目标知识库（便捷函数）

    Args:
        query: 用户查询
        role: 用户角色
        department: 用户部门
        accessible_collections: 可访问的向量库列表

    Returns:
        目标向量库列表
    """
    router = get_kb_router()
    return router.route(query, role, department, accessible_collections)


def route_query_with_version(
    query: str,
    role: str,
    department: str,
    accessible_collections: List[str] = None,
    include_deprecated: bool = False
) -> Dict:
    """
    版本感知的路由（便捷函数）

    Args:
        query: 用户查询
        role: 用户角色
        department: 用户部门
        accessible_collections: 可访问的向量库列表
        include_deprecated: 是否包含废止版本

    Returns:
        {
            "target_collections": [...],
            "version_hints": [...]
        }
    """
    router = get_kb_router()
    return router.route_with_version_awareness(
        query, role, department, accessible_collections, include_deprecated
    )


def search_with_version_context(
    query: str,
    role: str,
    department: str,
    top_k: int = 5
) -> Dict:
    """
    带版本上下文的搜索（便捷函数）

    Args:
        query: 用户查询
        role: 用户角色
        department: 用户部门
        top_k: 返回数量

    Returns:
        {
            "results": [...],
            "version_hints": [...],
            "target_collections": [...]
        }
    """
    router = get_kb_router()
    return router.search_with_version_context(query, role, department, top_k)