mildoc/mildoc_wxkf/rag_service.py

"""
RAG服务工具类 (基于LangChain实现)

使用LangChain + Milvus实现RAG服务
支持从Milvus向量数据库检索相关文档并通过大模型生成回答

作者：开发工程师
日期：2025年01月
"""

import logging
from typing import List, Optional, Dict, Any
from pydantic import BaseModel
from langchain_milvus import Milvus
from langchain_openai import ChatOpenAI
from langchain_community.callbacks.manager import get_openai_callback
from langfuse.langchain import CallbackHandler as LangfuseCallbackHandler
from config import Config
from rerank_service import get_rerank_service

# 配置日志
logger = logging.getLogger(__name__)

langfuse_handler = LangfuseCallbackHandler()

class SourceDocument(BaseModel):
    """源文档信息模型"""
    doc_name: str  # 文档名称
    doc_path_name: str  # 文档完整路径
    doc_type: str  # 文档类型
    content_preview: str  # 内容预览（前200字符）
    similarity_score: Optional[float] = None  # 相似度分数


class TokenUsage(BaseModel):
    """Token使用情况模型"""
    prompt_tokens: int  # 输入token数
    completion_tokens: int  # 输出token数
    total_tokens: int  # 总token数


class RAGResponse(BaseModel):
    """RAG服务响应模型"""
    content: str  # 大模型回复给用户的文本内容
    source_documents: List[SourceDocument]  # 检索使用的源文档列表
    token_usage: Optional[TokenUsage] = None  # token使用情况
    success: bool = True  # 查询是否成功
    error_message: Optional[str] = None  # 错误信息
    scene_info: Optional[Dict[str, Any]] = None  # 场景检测信息


class RAGService:
    """RAG服务类 (基于LangChain实现)
    
    使用LangChain + Milvus向量数据库实现检索增强生成服务
    支持OpenAI兼容的大模型和嵌入模型
    """
    
    # 场景检测提示词模板
    SCENE_DETECTION_TEMPLATE = """请分析用户问题属于以下哪种客服场景类型，只返回对应的数字：

1. 产品咨询类 - 询问产品功能、规格、价格等基本信息
2. 售后服务类 - 退换货、维修、质量问题等售后相关
3. 账户相关类 - 登录、注册、密码、个人信息等账户问题  
4. 投诉建议类 - 对服务或产品的投诉、意见、建议
5. 技术支持类 - 使用方法、故障排除、技术配置等
6. 其他咨询类 - 不属于以上分类的一般性咨询

用户问题: {question}

请只返回场景类型对应的数字（1-6）："""
    
    # 统一的提示词模板 - 专业客服版本
    PROMPT_TEMPLATE = """你是一位专业的客服人员，请根据提供的知识库内容来回答用户的问题。

知识库内容:
{context}

用户问题: {question}

回答要求：
1. 【角色定位】你是一位专业、耐心、友善的客服代表
2. 【回答原则】严格基于知识库内容回答，不得编造或推测信息
3. 【准确性要求】
   - 如果知识库中有明确答案，请准确完整地回答
   - 如果知识库中信息不完整，说明现有信息并提示用户可联系人工客服获取更详细信息
   - 如果知识库中完全没有相关信息，请礼貌地说明无法找到相关资料，建议用户转接人工客服
4. 【回答格式】
   - 使用纯文本格式，不使用markdown格式
   - 语言简洁明了，适合微信对话环境
   - 使用礼貌、专业的语调
   - 如需列举，使用数字序号或简单的分行
5. 【转人工提示】当遇到以下情况时，主动建议用户转接人工客服：
   - 复杂的售后问题
   - 需要个人账户信息查询的问题
   - 投诉或纠纷相关问题
   - 知识库无法覆盖的专业技术问题

请基于以上要求，为用户提供专业的客服回答："""
    
    def __init__(self):
        """初始化RAG服务"""
        self.vector_store = None
        self.embeddings = None
        self.llm = None
        self.rerank_service = None  # 重排序服务
        self._initialize_components()
    
    def _initialize_components(self):
        """初始化所有组件"""
        try:
            # 初始化嵌入模型
            self._initialize_embeddings()
            
            # 初始化大语言模型
            self._initialize_llm()
            
            # 初始化向量存储
            self._initialize_vector_store()
            
            # 初始化重排序服务
            self._initialize_rerank_service()
            
            logger.info("RAG服务初始化完成")
            
        except Exception as e:
            logger.error(f"RAG服务初始化失败: {e}")
            raise
    
    def _initialize_embeddings(self):
        """初始化嵌入模型"""
        try:
            # 使用自定义嵌入类，兼容OpenAI API
            from openai import OpenAI
            
            class CustomEmbeddings:
                def __init__(self, model_name: str, api_key: str, api_base: str, dimensions: int):
                    self.model_name = model_name
                    self.client = OpenAI(api_key=api_key, base_url=api_base)
                    self.dimensions = dimensions
                
                def embed_query(self, text: str) -> List[float]:
                    """嵌入单个查询"""
                    return self.embed_documents([text])[0]
                
                def embed_documents(self, texts: List[str]) -> List[List[float]]:
                    """嵌入多个文档"""
                    try:
                        response = self.client.embeddings.create(
                            model=self.model_name,
                            input=texts,
                            dimensions=self.dimensions,
                            encoding_format="float"
                        )
                        return [data.embedding for data in response.data]
                    except Exception as e:
                        logger.error(f"嵌入生成失败: {e}")
                        raise
            
            self.embeddings = CustomEmbeddings(
                model_name=Config.LLM_EMBEDDING_MODEL_NAME,
                api_key=Config.LLM_EMBEDDING_API_KEY,
                api_base=Config.LLM_EMBEDDING_BASE_URL,
                dimensions=Config.MILVUS_VECTOR_DIM
            )
            
            # 测试嵌入模型
            test_embedding = self.embeddings.embed_query("测试")
            actual_dim = len(test_embedding)
            
            logger.info(f"嵌入模型初始化成功: {Config.LLM_EMBEDDING_MODEL_NAME}")
            logger.info(f"向量维度: {actual_dim}")
            
            if actual_dim != Config.MILVUS_VECTOR_DIM:
                logger.warning(f"向量维度不匹配! 实际({actual_dim}) != 期望({Config.MILVUS_VECTOR_DIM})")
            
        except Exception as e:
            logger.error(f"嵌入模型初始化失败: {e}")
            raise
    
    def _initialize_llm(self):
        """初始化大语言模型"""
        try:
            self.llm = ChatOpenAI(
                model=Config.LLM_MODEL_NAME,
                openai_api_key=Config.LLM_API_KEY,
                openai_api_base=Config.LLM_BASE_URL,
                temperature=0.1,
                max_tokens=800  # 调整为800，平衡详细度和简洁性
            )
            
            logger.info(f"大语言模型初始化成功: {Config.LLM_MODEL_NAME}")
            
        except Exception as e:
            logger.error(f"大语言模型初始化失败: {e}")
            raise
    
    def _initialize_vector_store(self):
        """初始化向量存储"""
        try:
            # 构建Milvus连接参数
            connection_args = {
                "host": Config.MILVUS_HOST,
                "port": Config.MILVUS_PORT,
                "db_name": Config.MILVUS_DATABASE
            }
            
            # 如果有用户名和密码，添加到连接参数中
            if Config.MILVUS_USER:
                connection_args["user"] = Config.MILVUS_USER
            if Config.MILVUS_PASSWORD:
                connection_args["password"] = Config.MILVUS_PASSWORD
            
            # 配置搜索参数，针对IVF_FLAT索引优化
            search_params = {
                "metric_type": "COSINE",  # 与索引保持一致
                "params": {
                    "nprobe": 64  # 建议设置为nlist的6.25% (64/1024)，平衡性能和召回率
                }
            }
            
            # 初始化Milvus向量存储
            self.vector_store = Milvus(
                embedding_function=self.embeddings,
                collection_name=Config.MILVUS_COLLECTION_NAME,
                connection_args=connection_args,
                text_field="content",  # 文本内容字段
                vector_field="content_vector",  # 向量字段
                auto_id=True,
                search_params=search_params  # 添加搜索参数
            )
            
            logger.info(f"Milvus向量存储初始化成功: {Config.MILVUS_COLLECTION_NAME}")
            logger.info(f"搜索参数: nprobe={search_params['params']['nprobe']}")
            
        except Exception as e:
            logger.error(f"Milvus向量存储初始化失败: {e}")
            raise
    


    def _initialize_rerank_service(self):
        """初始化重排序服务"""
        try:
            self.rerank_service = get_rerank_service()
            if self.rerank_service:
                logger.info("重排序服务初始化成功")
            else:
                logger.info("重排序服务未配置，将跳过rerank步骤")
        except Exception as e:
            logger.warning(f"重排序服务初始化失败: {e}")
            self.rerank_service = None

    def query_service(self, query: str, use_rerank: bool = True) -> RAGResponse:
        """核心查询服务方法
        
        Args:
            query: 用户输入的查询内容
            use_rerank: 是否使用重排序功能
            
        Returns:
            RAGResponse: 包含回答内容、源文档和token使用情况的响应对象
        """
        try:
            logger.info(f"🔍 开始处理查询（rerank={use_rerank}): {query}")
            
            if not query or not query.strip():
                return RAGResponse(
                    content="请输入有效的查询内容",
                    source_documents=[],
                    success=False,
                    error_message="查询内容为空"
                )
            
            # 第0步：场景检测（可选）
            # scene_info = self.detect_user_scene(query)
            scene_info = None # 暂时不使用场景检测
            
            # 第一步：向量检索获取候选文档
            initial_k = 10 if use_rerank and self.rerank_service else 3
            retriever = self.vector_store.as_retriever(
                search_type="similarity",
                search_kwargs={"k": initial_k}
            )
            
            # 获取候选文档
            candidate_docs = retriever.invoke(query)
            logger.info(f"📄 初始检索到 {len(candidate_docs)} 个候选文档")
            
            # 第二步：重排序（如果启用）
            final_docs = candidate_docs
            if use_rerank and self.rerank_service and len(candidate_docs) > 1:
                # 提取文档内容用于重排序
                doc_contents = [doc.page_content for doc in candidate_docs]
                
                # 增加重排序的top_n数量，确保不会过滤掉高相关度文档
                rerank_top_n = min(5, len(candidate_docs))  # 从3增加到5，但不超过候选文档数量
                
                # 执行重排序
                rerank_response = self.rerank_service.rerank_documents(
                    query=query,
                    documents=doc_contents,
                    top_n=rerank_top_n
                )
                
                if rerank_response.success:
                    # 根据重排序结果重新排列文档
                    reranked_docs = []
                    for rerank_doc in rerank_response.documents:
                        if 0 <= rerank_doc.index < len(candidate_docs):
                            original_doc = candidate_docs[rerank_doc.index]
                            # 将相关性分数添加到元数据中
                            if hasattr(original_doc, 'metadata'):
                                original_doc.metadata['rerank_score'] = rerank_doc.relevance_score
                            reranked_docs.append(original_doc)
                    
                    # 安全检查：确保原始最高相似度文档不会被完全过滤掉
                    # 如果原始第1个文档不在重排序结果中，将其添加到结果中
                    if candidate_docs and len(reranked_docs) > 0:
                        first_doc = candidate_docs[0]
                        first_doc_in_rerank = any(
                            hasattr(doc, 'metadata') and 
                            doc.metadata.get('doc_name') == first_doc.metadata.get('doc_name') and
                            doc.page_content == first_doc.page_content
                            for doc in reranked_docs
                        )
                        
                        if not first_doc_in_rerank:
                            # 将原始最高相似度文档添加到重排序结果的开头
                            if hasattr(first_doc, 'metadata'):
                                first_doc.metadata['rerank_score'] = 1.0  # 给予最高分数
                            reranked_docs.insert(0, first_doc)
                            logger.info(f"🔒 安全检查：将原始最高相似度文档添加到重排序结果中")
                    
                    final_docs = reranked_docs[:3]  # 最终仍然只取前3个
                    logger.info(f"🔄 重排序完成，选择了 {len(final_docs)} 个文档")
                else:
                    logger.warning(f"重排序失败，使用原始检索结果: {rerank_response.error_message}")
            
            # 第三步：使用选定的文档生成回答
            with get_openai_callback() as cb:  ## 在上下文中获取 OpenAI 回调处理器，方便地公开令牌和成本信息
                # 构建上下文
                context = "\n\n".join([doc.page_content for doc in final_docs])
                
                # 使用统一的提示模板生成回答
                prompt = self.PROMPT_TEMPLATE.format(context=context, question=query)

                if (Config.LANGFUSE_ENABLE):
                    answer = self.llm.invoke(prompt, config={"callbacks":[langfuse_handler]}).content
                else:
                    answer = self.llm.invoke(prompt).content
            
            # 更新文档引用为最终选择的文档
            source_documents = final_docs
            
            logger.info(f"✅ 查询完成，检索到 {len(source_documents)} 个相关文档")
            logger.info(f"📄 答案长度: {len(answer)} 字符")
            logger.info(f"💰 Token使用: 输入{cb.prompt_tokens}, 输出{cb.completion_tokens}, 总计{cb.total_tokens}")
            
            # 处理源文档信息
            processed_source_docs = []
            for i, doc in enumerate(source_documents):
                try:
                    # 提取文档元数据
                    metadata = doc.metadata if hasattr(doc, 'metadata') else {}
                    doc_name = metadata.get("doc_name", f"文档{i+1}")
                    doc_path_name = metadata.get("doc_path_name", "")
                    doc_type = metadata.get("doc_type", "unknown")
                    rerank_score = metadata.get("rerank_score")
                    
                    # 获取内容预览
                    content_preview = doc.page_content[:200] + "..." if len(doc.page_content) > 200 else doc.page_content
                    
                    source_doc = SourceDocument(
                        doc_name=doc_name,
                        doc_path_name=doc_path_name,
                        doc_type=doc_type,
                        content_preview=content_preview,
                        similarity_score=rerank_score  # 使用rerank分数
                    )
                    processed_source_docs.append(source_doc)
                    
                    score_info = f" (rerank: {rerank_score:.3f})" if rerank_score else ""
                    logger.info(f"📖 文档{i+1}: {doc_name}{score_info} - {content_preview[:50]}...")
                    
                except Exception as e:
                    logger.warning(f"处理源文档{i+1}时出错: {e}")
                    # 添加默认文档信息
                    processed_source_docs.append(SourceDocument(
                        doc_name=f"文档{i+1}",
                        doc_path_name="",
                        doc_type="unknown",
                        content_preview="无法获取文档信息"
                    ))
            
            # 构建token使用情况
            token_usage = TokenUsage(
                prompt_tokens=cb.prompt_tokens,
                completion_tokens=cb.completion_tokens,
                total_tokens=cb.total_tokens
            )
            
            return RAGResponse(
                content=answer if answer else "抱歉，我无法根据现有信息回答您的问题。",
                source_documents=processed_source_docs,
                token_usage=token_usage,
                success=True,
                scene_info=scene_info
            )
            
        except Exception as e:
            logger.error(f"❌ 查询服务失败: {e}")
            return RAGResponse(
                content="",
                source_documents=[],
                success=False,
                error_message=f"查询过程中发生错误：{str(e)}",
                scene_info=None
            )
    
    def get_similar_documents(self, query: str, top_k: int = 3) -> List[Dict[str, Any]]:
        """获取相似文档（用于调试和分析）
        
        Args:
            query: 查询内容
            top_k: 返回文档数量
            
        Returns:
            List[Dict]: 相似文档列表
        """
        try:
            logger.info(f"🔍 搜索相似文档: {query} (top_k={top_k})")
            
            # 使用向量存储进行相似性搜索
            docs = self.vector_store.similarity_search_with_score(query, k=top_k)
            
            results = []
            for doc, score in docs:
                result = {
                    "content": doc.page_content,
                    "metadata": doc.metadata,
                    "similarity_score": float(score),
                    "doc_name": doc.metadata.get("doc_name", "未知文档"),
                    "doc_path_name": doc.metadata.get("doc_path_name", ""),
                    "doc_type": doc.metadata.get("doc_type", "unknown")
                }
                results.append(result)
            
            logger.info(f"✅ 找到 {len(results)} 个相似文档")
            return results
            
        except Exception as e:
            logger.error(f"❌ 获取相似文档失败: {e}")
            return []
    
    def health_check(self) -> Dict[str, Any]:
        """健康检查
        
        Returns:
            Dict: 健康状态信息
        """
        status = {
            "service": "RAGService",
            "status": "healthy",
            "components": {},
            "timestamp": None
        }
        
        try:
            from datetime import datetime
            status["timestamp"] = datetime.now().isoformat()
            
            # 检查嵌入模型
            try:
                test_embedding = self.embeddings.embed_query("健康检查")
                status["components"]["embeddings"] = {
                    "status": "healthy",
                    "model": Config.LLM_EMBEDDING_MODEL_NAME,
                    "dimension": len(test_embedding)
                }
            except Exception as e:
                status["components"]["embeddings"] = {
                    "status": "error",
                    "error": str(e)
                }
                status["status"] = "degraded"
            
            # 检查大语言模型
            try:
                with get_openai_callback() as cb:
                    test_response = self.llm.invoke("你好").content
                status["components"]["llm"] = {
                    "status": "healthy",
                    "model": Config.LLM_MODEL_NAME,
                    "response_length": len(test_response) if test_response else 0,
                    "test_tokens": cb.total_tokens
                }
            except Exception as e:
                status["components"]["llm"] = {
                    "status": "error",
                    "error": str(e)
                }
                status["status"] = "degraded"
            
            # 检查向量存储
            try:
                test_docs = self.vector_store.similarity_search("测试", k=1)
                status["components"]["vector_store"] = {
                    "status": "healthy",
                    "collection": Config.MILVUS_COLLECTION_NAME,
                    "test_results": len(test_docs)
                }
            except Exception as e:
                status["components"]["vector_store"] = {
                    "status": "error",
                    "error": str(e)
                }
                status["status"] = "degraded"
            
            # 检查重排序服务
            try:
                if self.rerank_service:
                    rerank_health = self.rerank_service.health_check()
                    status["components"]["rerank_service"] = rerank_health
                else:
                    status["components"]["rerank_service"] = {
                        "status": "not_configured",
                        "message": "重排序服务未配置"
                    }
            except Exception as e:
                status["components"]["rerank_service"] = {
                    "status": "error",
                    "error": str(e)
                }
                status["status"] = "degraded"
            
        except Exception as e:
            status["status"] = "error"
            status["error"] = str(e)
        
        return status

    def detect_user_scene(self, query: str) -> Dict[str, Any]:
        """检测用户问题的场景类型
        
        Args:
            query: 用户问题
            
        Returns:
            Dict: 包含场景类型和建议的字典
        """
        try:
            # 使用LLM进行场景检测
            prompt = self.SCENE_DETECTION_TEMPLATE.format(question=query)
            response = self.llm.invoke(prompt).content.strip()
            
            # 解析场景类型
            scene_mapping = {
                "1": {"type": "产品咨询类", "priority": "normal", "suggest_human": False},
                "2": {"type": "售后服务类", "priority": "high", "suggest_human": True},
                "3": {"type": "账户相关类", "priority": "high", "suggest_human": True},
                "4": {"type": "投诉建议类", "priority": "urgent", "suggest_human": True},
                "5": {"type": "技术支持类", "priority": "normal", "suggest_human": False},
                "6": {"type": "其他咨询类", "priority": "normal", "suggest_human": False}
            }
            
            scene_info = scene_mapping.get(response, scene_mapping["6"])
            scene_info["detected_number"] = response
            
            logger.info(f"🎯 场景检测结果: {query} -> {scene_info['type']} (优先级: {scene_info['priority']})")
            
            return scene_info
            
        except Exception as e:
            logger.warning(f"场景检测失败: {e}")
            return {"type": "其他咨询类", "priority": "normal", "suggest_human": False, "detected_number": "6"}


# 全局RAG服务实例
_rag_service_instance = None

def get_rag_service() -> Optional[RAGService]:
    """获取RAG服务实例（单例模式）"""
    global _rag_service_instance
    
    if _rag_service_instance is None:
        try:
            _rag_service_instance = RAGService()
            logger.info("RAG服务实例创建成功")
        except Exception as e:
            logger.error(f"RAG服务实例创建失败: {e}")
            return None
    
    return _rag_service_instance


# 便捷函数
def query_question(question: str) -> RAGResponse:
    """查询问题的便捷函数
    
    Args:
        question: 用户问题
        
    Returns:
        RAGResponse: 查询响应
    """
    rag_service = get_rag_service()
    if rag_service is None:
        return RAGResponse(
            content="",
            source_documents=[],
            success=False,
            error_message="RAG服务初始化失败"
        )
    
    return rag_service.query_service(question)


if __name__ == "__main__":
    # 测试代码 - 专业客服RAG系统
    rag = get_rag_service()


    # 检测健康状态
    health = rag.health_check()
    print(f"健康状态: {health}")


    # 测试不同场景的客服问题
    test_cases = [
        "帮我介绍一下盗窃罪",  # 其他咨询类
        "我要退货，怎么办理？",    # 售后服务类
        "忘记密码了，如何重置？",  # 账户相关类
        "你们的服务太差了！",      # 投诉建议类
        "产品无法连接WiFi",       # 技术支持类
    ]

    print(100 * "=")
    
    for i, question in enumerate(test_cases, 1):
        print(100 * "*")
        print(f"\n=== 测试案例 {i}: {question} ===")
        
        # 测试场景检测
        scene_info = rag.detect_user_scene(question)
        print(f"🎯 场景类型: {scene_info['type']}")
        print(f"🚨 优先级: {scene_info['priority']}")
        print(f"👤 建议转人工: {'是' if scene_info['suggest_human'] else '否'}")
        
        # 测试完整查询
        response = rag.query_service(question, use_rerank=True)
        print(f"💬 客服回答: {response.content}")
        if response.scene_info:
            print(f"📊 场景信息: {response.scene_info['type']}")
        if response.token_usage:
            print(f"💰 Token使用: {response.token_usage.total_tokens}")
        print(f"📚 参考文档数: {len(response.source_documents)}")
        
    print(100 * "=")
    
    print("\n=== 测试经典产品咨询 ===")
    response = rag.query_service("介绍一下老人与海这本书", use_rerank=True)
    print(f"回答: {response.content}")
    print(f"场景信息: {response.scene_info}")
    print(f"源文档数量: {len(response.source_documents)}")
    if response.token_usage:
        print(f"Token使用: {response.token_usage.total_tokens}")

    for doc in response.source_documents:
        print("--------------------------------")
        print(f"源文档: {doc.doc_name}")
        print(f"相似度分数: {doc.similarity_score}")
        print(f"内容预览: {doc.content_preview[:50]}...")