Self-RAG实现指南：让模型自己决定何时检索

[Retrieve] - 是否需要检索？
[IsRel]    - 检索结果是否相关？
[IsSup]    - 生成是否被检索结果支持？
[IsUse]    - 生成是否有用？

用户问题
    ↓
模型判断 [Retrieve]? → No → 直接生成
                   → Yes → 执行检索
                            ↓
                    评估每个文档 [IsRel]? → 过滤不相关
                            ↓
                    基于相关文档生成
                            ↓
                    评估生成 [IsSup]? → No → 重新生成
                                   → Yes → 评估 [IsUse]?
                                              ↓
                                         输出最终答案

class SelfRAG:
    def __init__(self, llm, retriever):
        self.llm = llm
        self.retriever = retriever
    
    async def generate(self, question):
        # 1. 判断是否需要检索
        need_retrieve = await self.predict_retrieve_token(question)
        
        if not need_retrieve:
            # 不需要检索，直接生成
            answer = await self.llm.generate(question)
            return {"answer": answer, "retrieved": False}
        
        # 2. 执行检索
        docs = await self.retriever.search(question, top_k=5)
        
        # 3. 评估每个文档的相关性
        relevant_docs = []
        for doc in docs:
            is_rel = await self.predict_isrel_token(question, doc)
            if is_rel:
                relevant_docs.append(doc)
        
        if not relevant_docs:
            # 没有相关文档，直接生成
            answer = await self.llm.generate(question)
            return {"answer": answer, "retrieved": True, "relevant": 0}
        
        # 4. 基于相关文档生成多个候选
        candidates = []
        for doc in relevant_docs:
            candidate = await self.llm.generate(
                f"基于以下信息回答问题:\n{doc.text}\n\n问题: {question}"
            )
            
            # 5. 评估每个候选是否被文档支持
            is_sup = await self.predict_issup_token(candidate, doc)
            
            # 6. 评估是否有用
            is_use = await self.predict_isuse_token(question, candidate)
            
            candidates.append({
                "answer": candidate,
                "support_score": is_sup,
                "usefulness_score": is_use,
                "source_doc": doc
            })
        
        # 7. 选择最佳候选
        best = max(candidates, key=lambda x: x["support_score"] * x["usefulness_score"])
        
        return {
            "answer": best["answer"],
            "retrieved": True,
            "relevant_docs": len(relevant_docs),
            "source": best["source_doc"]
        }

class ReflectionTokenPredictor:
    def __init__(self, llm):
        self.llm = llm
    
    async def predict_retrieve(self, question):
        """预测是否需要检索"""
        prompt = f"""
        问题: {question}
        
        这个问题是否需要查阅外部资料才能准确回答？
        考虑:
        - 问题是否涉及事实性知识？
        - 你的训练数据中是否有足够的知识？
        - 问题是否涉及最新信息？
        
        输出: "RETRIEVE" 或 "NO_RETRIEVE"
        """
        result = await self.llm.call(prompt)
        return "RETRIEVE" in result
    
    async def predict_isrel(self, question, document):
        """预测文档是否与问题相关"""
        prompt = f"""
        问题: {question}
        文档: {document.text[:500]}
        
        这个文档是否包含回答问题所需的信息？
        
        输出: "RELEVANT" 或 "IRRELEVANT"
        """
        result = await self.llm.call(prompt)
        return "RELEVANT" in result
    
    async def predict_issup(self, answer, document):
        """预测答案是否被文档支持"""
        prompt = f"""
        文档: {document.text[:500]}
        答案: {answer}
        
        答案中的所有陈述是否都能从文档中找到支持？
        
        输出: "FULLY_SUPPORTED" / "PARTIALLY_SUPPORTED" / "NOT_SUPPORTED"
        """
        result = await self.llm.call(prompt)
        if "FULLY" in result:
            return 1.0
        elif "PARTIALLY" in result:
            return 0.5
        else:
            return 0.0
    
    async def predict_isuse(self, question, answer):
        """预测答案是否有用"""
        prompt = f"""
        问题: {question}
        答案: {answer}
        
        这个答案对提问者有多大帮助？（1-5分）
        """
        result = await self.llm.call(prompt)
        return int(result.strip()) / 5.0

class SelfRAGTrainingDataBuilder:
    async def build_training_data(self, base_dataset):
        """构建Self-RAG训练数据"""
        training_samples = []
        
        for sample in base_dataset:
            question = sample["question"]
            answer = sample["answer"]
            relevant_docs = sample.get("relevant_docs", [])
            
            # 1. 生成Retrieve标签
            retrieve_label = "RETRIEVE" if relevant_docs else "NO_RETRIEVE"
            
            # 2. 生成IsRel标签
            for doc in relevant_docs:
                isrel_label = "RELEVANT"
                training_samples.append({
                    "input": f"问题: {question}\n文档: {doc}",
                    "label": isrel_label,
                    "token_type": "IsRel"
                })
            
            # 3. 生成IsSup标签
            for doc in relevant_docs:
                issup_label = "FULLY_SUPPORTED"  # 假设标准答案被文档支持
                training_samples.append({
                    "input": f"文档: {doc}\n答案: {answer}",
                    "label": issup_label,
                    "token_type": "IsSup"
                })
            
            # 4. 生成IsUse标签
            isuse_label = "5"  # 标准答案评为5分
            training_samples.append({
                "input": f"问题: {question}\n答案: {answer}",
                "label": isuse_label,
                "token_type": "IsUse"
            })
        
        return training_samples

class SelfRAGTrainer:
    def __init__(self, base_model):
        self.base_model = base_model
    
    async def train(self, training_data):
        """微调模型学习反思令牌预测"""
        # 1. 格式化训练数据
        formatted = self.format_training_data(training_data)
        
        # 2. 微调
        # 方式1: 在基础模型上微调（需要大量数据）
        # 方式2: 使用few-shot prompting（更实用）
        
        # 使用方式2
        self.model = self.create_few_shot_model(formatted)
    
    def create_few_shot_model(self, training_data):
        """使用few-shot examples创建Self-RAG模型"""
        # 为每种反思令牌选择最好的示例
        retrieve_examples = self.select_examples(training_data, "Retrieve", n=3)
        isrel_examples = self.select_examples(training_data, "IsRel", n=3)
        issup_examples = self.select_examples(training_data, "IsSup", n=3)
        isuse_examples = self.select_examples(training_data, "IsUse", n=3)
        
        system_prompt = f"""
        你是一个Self-RAG助手。在回答问题时，你需要:
        
        1. 首先判断是否需要检索外部资料:
        {self.format_examples(retrieve_examples)}
        
        2. 如果检索，评估每个文档的相关性:
        {self.format_examples(isrel_examples)}
        
        3. 评估答案是否被文档支持:
        {self.format_examples(issup_examples)}
        
        4. 评估答案的有用程度:
        {self.format_examples(isuse_examples)}
        """
        
        return FewShotModel(self.base_model, system_prompt)

class HybridSelfRAG:
    """Self-RAG与规则结合的混合模式"""
    
    async def generate(self, question):
        # 规则前置：某些类型的问题不需要检索
        if self.is_simple_factual(question):
            return await self.llm.generate(question)
        
        if self.is_current_event(question):
            need_retrieve = True
        else:
            need_retrieve = await self.predict_retrieve(question)
        
        if need_retrieve:
            return await self.rag_generate(question)
        else:
            return await self.llm.generate(question)
    
    def is_simple_factual(self, question):
        """简单事实型问题"""
        simple_patterns = [
            r"^\d+[\+\-\*/]\d+",  # 数学计算
            r"什么是.*定义",       # 定义类
        ]
        return any(re.match(p, question) for p in simple_patterns)
    
    def is_current_event(self, question):
        """时事类问题"""
        event_keywords = ["最新", "今天", "最近", "现在", "current", "latest"]
        return any(kw in question.lower() for kw in event_keywords)

class CostAwareSelfRAG:
    """成本感知的Self-RAG"""
    
    def __init__(self):
        self.token_budget = 100000  # 每次会话的Token预算
        self.token_used = 0
    
    async def generate(self, question):
        # 估算检索+生成的Token消耗
        retrieve_cost = 2000  # 检索+反思令牌的估算成本
        direct_cost = 500     # 直接生成的估算成本
        
        remaining_budget = self.token_budget - self.token_used
        
        if remaining_budget < direct_cost:
            # 预算不足，使用最省钱的模式
            return await self.llm.generate(question)
        
        # 正常Self-RAG流程
        need_retrieve = await self.predict_retrieve(question)
        
        if need_retrieve and remaining_budget > retrieve_cost:
            result = await self.rag_generate(question)
            self.token_used += retrieve_cost
        else:
            result = await self.llm.generate(question)
            self.token_used += direct_cost
        
        return result

class SelfRAGEvaluator:
    async def evaluate(self, self_rag, test_cases):
        metrics = {
            "retrieve_accuracy": [],    # 检索决策准确率
            "relevance_accuracy": [],   # 相关性判断准确率
            "support_accuracy": [],     # 支持度评估准确率
            "answer_quality": [],       # 答案质量
            "efficiency": []            # 效率（Token消耗）
        }
        
        for case in test_cases:
            result = await self_rag.generate(case.question)
            
            # 检索决策是否正确
            should_retrieve = bool(case.relevant_docs)
            actual_retrieve = result["retrieved"]
            metrics["retrieve_accuracy"].append(should_retrieve == actual_retrieve)
            
            # 答案质量
            quality = await self.evaluate_answer(result["answer"], case.expected_answer)
            metrics["answer_quality"].append(quality)
            
            # 效率
            metrics["efficiency"].append(result.get("tokens_used", 0))
        
        return {k: np.mean(v) for k, v in metrics.items()}