rag_chain.py

from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

class RAGChain:
    """
    RAG链类，用于处理检索增强生成
    """
    
    def __init__(self, vector_db, llm_model="gpt-3.5-turbo"):
        """
        初始化RAG链
        
        Args:
            vector_db: VectorDB实例
            llm_model: LLM模型名称
        """
        self.vector_db = vector_db
        self.llm_model = llm_model
        self.llm = self._get_llm()
        self.chain = self._build_chain()
    
    def _get_llm(self):
        """
        获取LLM模型
        
        Returns:
            LLM实例
        """
        try:
            if self.llm_model.startswith("gpt-"):
                # 使用OpenAI模型
                return ChatOpenAI(model=self.llm_model, temperature=0.3)
            elif self.llm_model.startswith("claude-"):
                # 使用Anthropic模型
                return ChatAnthropic(model=self.llm_model, temperature=0.3)
            else:
                # 默认使用gpt-3.5-turbo
                return ChatOpenAI(model="gpt-3.5-turbo", temperature=0.3)
        except Exception as e:
            print(f"初始化LLM模型失败: {e}")
            # 可以添加fallback逻辑
            raise
    
    def _build_chain(self):
        """
        构建RAG链
        
        Returns:
            RAG链实例
        """
        # 定义提示模板
        prompt = ChatPromptTemplate.from_template("""
        请根据以下上下文回答用户的问题。
        确保你的回答完全基于提供的上下文，不要添加任何外部信息。
        如果上下文没有相关信息，请直接说"根据提供的文档，我无法回答这个问题"。
        
        上下文：
        {context}
        
        问题：
        {question}
        
        回答：
        """)
        
        # 构建检索器
        retriever = self.vector_db.db.as_retriever(search_kwargs={"k": 4})
        
        # 构建链
        chain = (
            {"context": retriever, "question": RunnablePassthrough()}
            | prompt
            | self.llm
            | StrOutputParser()
        )
        
        return chain
    
    def invoke(self, question):
        """
        调用RAG链生成回答
        
        Args:
            question: 用户问题
            
        Returns:
            生成的回答
        """
        try:
            return self.chain.invoke(question)
        except Exception as e:
            print(f"调用RAG链失败: {e}")
            raise
    
    def update_llm_model(self, llm_model):
        """
        更新LLM模型
        
        Args:
            llm_model: 新的LLM模型名称
        """
        self.llm_model = llm_model
        self.llm = self._get_llm()
        self.chain = self._build_chain()