delete vectara

evaluate_retrieval.py

@@ -0,0 +1,346 @@
+"""
+自适应RAG系统检索效果评估脚本
+评估不同检索策略和配置的效果
+"""
+
+import os
+import sys
+import time
+import json
+import argparse
+from typing import List, Dict, Any, Optional
+from dotenv import load_dotenv
+
+# 加载环境变量
+load_dotenv()
+
+# 导入项目模块
+from main import AdaptiveRAGSystem
+from document_processor import DocumentProcessor
+from retrieval_evaluation import RetrievalEvaluator, RetrievalResult, RetrievalTestSet
+from langchain.schema import Document
+
+# 导入LangChain相关模块
+from langchain_community.vectorstores import FAISS, Chroma
+from langchain_community.retrievers import BM25Retriever
+from langchain.retrievers import EnsembleRetriever
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import LLMChainExtractor
+
+
+class AdaptiveRAGRetriever:
+    """自适应RAG系统检索器包装器"""
+    
+    def __init__(self, system_config: Dict[str, Any], retriever_type: str = "default"):
+        """
+        初始化检索器
+        
+        Args:
+            system_config: 系统配置
+            retriever_type: 检索器类型
+        """
+        self.system_config = system_config
+        self.retriever_type = retriever_type
+        self.system = None
+        self._initialize_system()
+    
+    def _initialize_system(self):
+        """初始化RAG系统"""
+        try:
+            # 根据检索器类型调整配置
+            config = self.system_config.copy()
+            
+            if self.retriever_type == "vector_only":
+                config["retrieval_strategy"] = "vector"
+            elif self.retriever_type == "bm25_only":
+                config["retrieval_strategy"] = "bm25"
+            elif self.retriever_type == "hybrid":
+                config["retrieval_strategy"] = "hybrid"
+            elif self.retriever_type == "graph":
+                config["retrieval_strategy"] = "graph"
+            elif self.retriever_type == "compression":
+                config["use_compression"] = True
+            elif self.retriever_type == "rerank":
+                config["use_reranking"] = True
+            elif self.retriever_type == "query_expansion":
+                config["use_query_expansion"] = True
+            
+            # 创建系统实例
+            self.system = AdaptiveRAGSystem(config)
+            
+            # 初始化文档处理器（如果需要）
+            if not hasattr(self.system, 'document_processor') or self.system.document_processor is None:
+                self.system.document_processor = DocumentProcessor(config)
+            
+        except Exception as e:
+            print(f"初始化RAG系统失败: {e}")
+            raise
+    
+    def retrieve(self, query: str, top_k: int = 10) -> List[Document]:
+        """
+        检索文档
+        
+        Args:
+            query: 查询文本
+            top_k: 返回的文档数量
+            
+        Returns:
+            检索到的文档列表
+        """
+        try:
+            # 使用系统的检索方法
+            if hasattr(self.system, 'retrieve'):
+                docs = self.system.retrieve(query, top_k)
+            else:
+                # 如果没有直接的retrieve方法，尝试通过文档处理器检索
+                if self.system.document_processor:
+                    docs = self.system.document_processor.retrieve(query, top_k)
+                else:
+                    raise ValueError("无法找到检索方法")
+            
+            return docs[:top_k]
+        except Exception as e:
+            print(f"检索失败: {e}")
+            return []
+
+
+def create_evaluation_dataset(data_dir: str = "data", num_queries: int = 20) -> RetrievalTestSet:
+    """
+    从项目数据创建评估数据集
+    
+    Args:
+        data_dir: 数据目录
+        num_queries: 查询数量
+        
+    Returns:
+        检索测试集
+    """
+    # 检查数据目录
+    if not os.path.exists(data_dir):
+        print(f"数据目录 {data_dir} 不存在，创建示例数据集")
+        from retrieval_evaluation import create_sample_test_set
+        return create_sample_test_set()
+    
+    # 尝试从现有数据创建测试集
+    try:
+        # 加载文档
+        documents = []
+        doc_files = []
+        
+        # 查找所有文本文件
+        for root, dirs, files in os.walk(data_dir):
+            for file in files:
+                if file.endswith('.txt') or file.endswith('.md'):
+                    doc_files.append(os.path.join(root, file))
+        
+        # 如果没有找到文档文件，创建示例数据集
+        if not doc_files:
+            print(f"在 {data_dir} 中未找到文档文件，创建示例数据集")
+            from retrieval_evaluation import create_sample_test_set
+            return create_sample_test_set()
+        
+        # 读取文档内容
+        for i, file_path in enumerate(doc_files):
+            with open(file_path, 'r', encoding='utf-8') as f:
+                content = f.read().strip()
+                if content:
+                    documents.append(Document(page_content=content, metadata={"source": file_path, "doc_id": str(i)}))
+        
+        # 生成查询（这里简化处理，实际应用中应该使用真实查询）
+        queries = []
+        qrels = {}
+        
+        # 从文档中提取关键句子作为查询
+        for i in range(min(num_queries, len(documents))):
+            doc = documents[i]
+            sentences = doc.page_content.split('.')
+            if sentences:
+                # 取第一个非空句子作为查询
+                for sentence in sentences:
+                    sentence = sentence.strip()
+                    if sentence and len(sentence) > 10:  # 确保查询有足够长度
+                        queries.append(sentence)
+                        # 假设查询与当前文档相关
+                        qrels[str(i)] = {str(i): 2}  # 高度相关
+                        # 可能与其他文档也相关
+                        for j in range(min(3, len(documents))):
+                            if j != i:
+                                qrels[str(i)][str(j)] = 1  # 部分相关
+                        break
+        
+        # 保存查询文件
+        with open("eval_queries.txt", "w", encoding="utf-8") as f:
+            for query in queries:
+                f.write(query + "\n")
+        
+        # 保存文档文件
+        with open("eval_documents.txt", "w", encoding="utf-8") as f:
+            for doc in documents:
+                f.write(doc.page_content + "\n")
+        
+        # 保存相关性标注文件
+        with open("eval_qrels.csv", "w", encoding="utf-8") as f:
+            for query_id, doc_relevance in qrels.items():
+                for doc_id, relevance in doc_relevance.items():
+                    f.write(f"{query_id},{doc_id},{relevance}\n")
+        
+        print(f"评估数据集已创建:")
+        print(f"- 查询数量: {len(queries)}")
+        print(f"- 文档数量: {len(documents)}")
+        print(f"- eval_queries.txt: 查询文件")
+        print(f"- eval_documents.txt: 文档文件")
+        print(f"- eval_qrels.csv: 相关性标注文件")
+        
+        return RetrievalTestSet("eval_queries.txt", "eval_documents.txt", "eval_qrels.csv")
+    
+    except Exception as e:
+        print(f"创建评估数据集失败: {e}")
+        print("创建示例数据集")
+        from retrieval_evaluation import create_sample_test_set
+        return create_sample_test_set()
+
+
+def evaluate_retrievers(system_config: Dict[str, Any], 
+                       retriever_types: List[str],
+                       test_set: RetrievalTestSet,
+                       output_dir: str = "evaluation_results") -> Dict[str, Any]:
+    """
+    评估多个检索器
+    
+    Args:
+        system_config: 系统配置
+        retriever_types: 检索器类型列表
+        test_set: 测试集
+        output_dir: 输出目录
+        
+    Returns:
+        评估结果
+    """
+    # 创建输出目录
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # 初始化评估器
+    evaluator = RetrievalEvaluator()
+    
+    # 存储所有检索结果
+    all_results = {}
+    
+    # 评估每个检索器
+    for retriever_type in retriever_types:
+        print(f"\n评估检索器: {retriever_type}")
+        print("=" * 50)
+        
+        try:
+            # 创建检索器
+            retriever = AdaptiveRAGRetriever(system_config, retriever_type)
+            
+            # 获取检索结果
+            results = test_set.get_retrieval_results(retriever)
+            all_results[retriever_type] = results
+            
+            print(f"完成 {len(results)} 个查询的检索")
+            
+        except Exception as e:
+            print(f"评估检索器 {retriever_type} 失败: {e}")
+            continue
+    
+    # 比较检索器
+    if len(all_results) > 1:
+        print("\n比较检索器性能")
+        print("=" * 50)
+        metrics = evaluator.compare_retrievers(all_results)
+        
+        # 生成报告
+        report = evaluator.generate_report(
+            metrics, 
+            os.path.join(output_dir, "retrieval_evaluation_report.md")
+        )
+        
+        # 绘制比较图
+        evaluator.plot_metrics_comparison(
+            metrics, 
+            os.path.join(output_dir, "retrieval_evaluation_comparison.png")
+        )
+        
+        # 保存详细指标
+        metrics_data = {}
+        for name, metric in metrics.items():
+            metrics_data[name] = {
+                "precision_at_k": metric.precision_at_k,
+                "recall_at_k": metric.recall_at_k,
+                "f1_at_k": metric.f1_at_k,
+                "map_score": metric.map_score,
+                "mrr": metric.mrr,
+                "ndcg_at_k": metric.ndcg_at_k,
+                "coverage": metric.coverage,
+                "diversity": metric.diversity,
+                "novelty": metric.novelty,
+                "latency": metric.latency
+            }
+        
+        with open(os.path.join(output_dir, "metrics.json"), "w", encoding="utf-8") as f:
+            json.dump(metrics_data, f, indent=2, ensure_ascii=False)
+        
+        return {
+            "metrics": metrics,
+            "metrics_data": metrics_data,
+            "report": report,
+            "results": all_results
+        }
+    else:
+        print("只有一个检索器成功评估，跳过比较")
+        return {"results": all_results}
+
+
+def main():
+    """主函数"""
+    parser = argparse.ArgumentParser(description="评估自适应RAG系统的检索效果")
+    parser.add_argument("--config", type=str, default="config.py", help="配置文件路径")
+    parser.add_argument("--data_dir", type=str, default="data", help="数据目录")
+    parser.add_argument("--output_dir", type=str, default="evaluation_results", help="输出目录")
+    parser.add_argument("--num_queries", type=int, default=20, help="查询数量")
+    parser.add_argument("--retrievers", nargs="+", 
+                       default=["default", "vector_only", "bm25_only", "hybrid"],
+                       help="要评估的检索器类型")
+    
+    args = parser.parse_args()
+    
+    # 加载配置
+    try:
+        if args.config.endswith('.py'):
+            # 动态导入Python配置文件
+            import importlib.util
+            spec = importlib.util.spec_from_file_location("config", args.config)
+            config_module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(config_module)
+            system_config = config_module.config
+        else:
+            # 加载JSON配置文件
+            with open(args.config, 'r', encoding='utf-8') as f:
+                system_config = json.load(f)
+    except Exception as e:
+        print(f"加载配置文件失败: {e}")
+        print("使用默认配置")
+        system_config = {
+            "model_name": "gpt-3.5-turbo",
+            "vector_store": "faiss",
+            "retrieval_strategy": "hybrid",
+            "use_reranking": False,
+            "use_compression": False,
+            "use_query_expansion": False
+        }
+    
+    # 创建评估数据集
+    print("创建评估数据集")
+    test_set = create_evaluation_dataset(args.data_dir, args.num_queries)
+    
+    # 评估检索器
+    print("\n开始评估检索器")
+    results = evaluate_retrievers(system_config, args.retrievers, test_set, args.output_dir)
+    
+    print("\n评估完成!")
+    print(f"结果保存在: {args.output_dir}")
+
+
+if __name__ == "__main__":
+    main()

main.py

@@ -110,13 +110,14 @@ class AdaptiveRAGSystem:
             verbose (bool): 是否显示详细输出
             
         Returns:
-            str: 最终答案
+            dict: 包含最终答案和评估指标的字典
         """
         print(f"\n🔍 处理问题: {question}")
         print("=" * 50)
         
         inputs = {"question": question, "retry_count": 0}  # 初始化重试计数器
         final_generation = None
+        retrieval_metrics = None
         
         # 设置配置，增加递归限制
         config = {"recursion_limit": 50}  # 增加到 50，默认是 25
@@ -128,6 +129,9 @@ class AdaptiveRAGSystem:
                     # 可选：在每个节点打印完整状态
                     # pprint(value, indent=2, width=80, depth=None)
                 final_generation = value.get("generation", final_generation)
+                # 保存检索评估指标
+                if "retrieval_metrics" in value:
+                    retrieval_metrics = value["retrieval_metrics"]
             if verbose:
                 pprint("\n---\n")
         
@@ -136,7 +140,11 @@ class AdaptiveRAGSystem:
         print(final_generation)
         print("=" * 50)
         
-        return final_generation
+        # 返回包含答案和评估指标的字典
+        return {
+            "answer": final_generation,
+            "retrieval_metrics": retrieval_metrics
+        }
     
     def interactive_mode(self):
         """交互模式，允许用户持续提问"""
@@ -156,7 +164,17 @@ class AdaptiveRAGSystem:
                     print("⚠️  请输入一个有效的问题")
                     continue
                 
-                self.query(question)
+                result = self.query(question)
+                
+                # 显示检索评估摘要
+                if result.get("retrieval_metrics"):
+                    metrics = result["retrieval_metrics"]
+                    print("\n📊 检索评估摘要:")
+                    print(f"   - 检索耗时: {metrics.get('latency', 0):.4f}秒")
+                    print(f"   - 检索文档数: {metrics.get('retrieved_docs_count', 0)}")
+                    print(f"   - Precision@3: {metrics.get('precision_at_3', 0):.4f}")
+                    print(f"   - Recall@3: {metrics.get('recall_at_3', 0):.4f}")
+                    print(f"   - MAP: {metrics.get('map_score', 0):.4f}")
                 
             except KeyboardInterrupt:
                 print("\n👋 感谢使用，再见!")
@@ -175,7 +193,17 @@ def main():
         # 测试查询
         test_question = "AlphaCodium论文讲的是什么？"
         # test_question = "解释embedding嵌入的原理，最好列举实现过程的具体步骤"
-        rag_system.query(test_question)
+        result = rag_system.query(test_question)
+        
+        # 显示测试查询的检索评估摘要
+        if result.get("retrieval_metrics"):
+            metrics = result["retrieval_metrics"]
+            print("\n📊 测试查询检索评估摘要:")
+            print(f"   - 检索耗时: {metrics.get('latency', 0):.4f}秒")
+            print(f"   - 检索文档数: {metrics.get('retrieved_docs_count', 0)}")
+            print(f"   - Precision@3: {metrics.get('precision_at_3', 0):.4f}")
+            print(f"   - Recall@3: {metrics.get('recall_at_3', 0):.4f}")
+            print(f"   - MAP: {metrics.get('map_score', 0):.4f}")
         
         # 启动交互模式
         rag_system.interactive_mode()

retrieval_evaluation.py

@@ -0,0 +1,674 @@
+"""
+检索效果评估模块
+提供多种评估指标和方法，用于评估RAG系统中检索结果的质量
+"""
+
+import time
+import json
+import numpy as np
+from typing import List, Dict, Tuple, Any, Optional, Union
+from dataclasses import dataclass, asdict
+from langchain.schema import Document
+from sklearn.metrics import ndcg_score, precision_score, recall_score, f1_score
+from sentence_transformers import SentenceTransformer, util
+import matplotlib.pyplot as plt
+import seaborn as sns
+import pandas as pd
+import torch
+
+
+@dataclass
+class RetrievalResult:
+    """检索结果数据类"""
+    query: str
+    retrieved_docs: List[Document]
+    relevant_docs: List[Document]  # 真实相关的文档
+    retrieval_time: float
+    scores: Optional[List[float]] = None  # 检索分数
+
+
+@dataclass
+class EvaluationMetrics:
+    """评估指标数据类"""
+    precision_at_k: Dict[int, float]
+    recall_at_k: Dict[int, float]
+    f1_at_k: Dict[int, float]
+    map_score: float  # 平均精度均值
+    mrr: float  # 平均倒数排名
+    ndcg_at_k: Dict[int, float]
+    coverage: float  # 覆盖率
+    diversity: float  # 多样性
+    novelty: float  # 新颖性
+    latency: float  # 平均检索延迟
+
+
+class RetrievalEvaluator:
+    """检索效果评估器"""
+    
+    def __init__(self, embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2"):
+        """
+        初始化评估器
+        
+        Args:
+            embedding_model: 用于计算语义相似度的嵌入模型
+        """
+        self.embedding_model = SentenceTransformer(embedding_model)
+        
+    def evaluate_retrieval(self, results: List[RetrievalResult], k_values: List[int] = [1, 3, 5, 10]) -> EvaluationMetrics:
+        """
+        评估检索结果
+        
+        Args:
+            results: 检索结果列表
+            k_values: 要计算的k值列表
+            
+        Returns:
+            评估指标
+        """
+        precision_at_k = {}
+        recall_at_k = {}
+        f1_at_k = {}
+        ndcg_at_k = {}
+        
+        total_precision = {k: 0 for k in k_values}
+        total_recall = {k: 0 for k in k_values}
+        total_f1 = {k: 0 for k in k_values}
+        total_ndcg = {k: 0 for k in k_values}
+        
+        all_precisions = []
+        all_reciprocal_ranks = []
+        all_latencies = []
+        
+        for result in results:
+            query = result.query
+            retrieved_docs = result.retrieved_docs
+            relevant_docs = result.relevant_docs
+            retrieval_time = result.retrieval_time
+            
+            all_latencies.append(retrieval_time)
+            
+            # 获取相关文档的ID或内容
+            relevant_ids = set()
+            for doc in relevant_docs:
+                # 使用文档内容作为ID，实际应用中可以使用文档ID
+                doc_id = doc.page_content[:50]  # 使用前50个字符作为ID
+                relevant_ids.add(doc_id)
+            
+            # 计算每个k值的指标
+            for k in k_values:
+                retrieved_k = retrieved_docs[:k]
+                retrieved_k_ids = set()
+                
+                for doc in retrieved_k:
+                    doc_id = doc.page_content[:50]
+                    retrieved_k_ids.add(doc_id)
+                
+                # 计算交集
+                intersection = len(relevant_ids.intersection(retrieved_k_ids))
+                
+                # 计算Precision@K
+                precision_k = intersection / k if k > 0 else 0
+                total_precision[k] += precision_k
+                
+                # 计算Recall@K
+                recall_k = intersection / len(relevant_ids) if len(relevant_ids) > 0 else 0
+                total_recall[k] += recall_k
+                
+                # 计算F1@K
+                if precision_k + recall_k > 0:
+                    f1_k = 2 * (precision_k * recall_k) / (precision_k + recall_k)
+                else:
+                    f1_k = 0
+                total_f1[k] += f1_k
+                
+                # 计算NDCG@K
+                if result.scores:
+                    # 创建相关性分数 (1表示相关，0表示不相关)
+                    relevance_scores = []
+                    for doc in retrieved_k:
+                        doc_id = doc.page_content[:50]
+                        relevance = 1 if doc_id in relevant_ids else 0
+                        relevance_scores.append(relevance)
+                    
+                    # 理想排序 (所有相关文档排在前面)
+                    ideal_relevance = sorted(relevance_scores, reverse=True)
+                    
+                    # 计算NDCG
+                    if len(relevance_scores) > 1 and sum(ideal_relevance) > 0:
+                        try:
+                            ndcg_k = ndcg_score([ideal_relevance], [relevance_scores], k=k)
+                            total_ndcg[k] += ndcg_k
+                        except:
+                            # 如果计算失败，使用简化的NDCG计算
+                            dcg = 0
+                            idcg = 0
+                            for i, rel in enumerate(relevance_scores):
+                                dcg += rel / np.log2(i + 2) if rel > 0 else 0
+                            for i, rel in enumerate(ideal_relevance):
+                                idcg += rel / np.log2(i + 2) if rel > 0 else 0
+                            ndcg_k = dcg / idcg if idcg > 0 else 0
+                            total_ndcg[k] += ndcg_k
+                    else:
+                        total_ndcg[k] += 1.0  # 如果没有相关文档或只有一个文档，NDCG为1
+            
+            # 计算平均精度 (AP)
+            precisions = []
+            for i, doc in enumerate(retrieved_docs):
+                doc_id = doc.page_content[:50]
+                if doc_id in relevant_ids:
+                    precision_at_i = len(relevant_ids.intersection(set(
+                        d.page_content[:50] for d in retrieved_docs[:i+1]
+                    ))) / (i + 1)
+                    precisions.append(precision_at_i)
+            
+            ap = sum(precisions) / len(relevant_ids) if precisions else 0
+            all_precisions.append(ap)
+            
+            # 计算倒数排名 (RR)
+            for i, doc in enumerate(retrieved_docs):
+                doc_id = doc.page_content[:50]
+                if doc_id in relevant_ids:
+                    rr = 1 / (i + 1)
+                    all_reciprocal_ranks.append(rr)
+                    break
+            else:
+                all_reciprocal_ranks.append(0)
+        
+        # 计算平均指标
+        num_results = len(results)
+        for k in k_values:
+            precision_at_k[k] = total_precision[k] / num_results
+            recall_at_k[k] = total_recall[k] / num_results
+            f1_at_k[k] = total_f1[k] / num_results
+            ndcg_at_k[k] = total_ndcg[k] / num_results
+        
+        map_score = sum(all_precisions) / num_results if all_precisions else 0
+        mrr = sum(all_reciprocal_ranks) / num_results if all_reciprocal_ranks else 0
+        latency = sum(all_latencies) / num_results if all_latencies else 0
+        
+        # 计算覆盖率、多样性和新颖性
+        coverage = self._calculate_coverage(results)
+        diversity = self._calculate_diversity(results)
+        novelty = self._calculate_novelty(results)
+        
+        return EvaluationMetrics(
+            precision_at_k=precision_at_k,
+            recall_at_k=recall_at_k,
+            f1_at_k=f1_at_k,
+            map_score=map_score,
+            mrr=mrr,
+            ndcg_at_k=ndcg_at_k,
+            coverage=coverage,
+            diversity=diversity,
+            novelty=novelty,
+            latency=latency
+        )
+    
+    def _calculate_coverage(self, results: List[RetrievalResult]) -> float:
+        """计算覆盖率 - 检索到的唯一文档数与总文档数的比例"""
+        all_retrieved = set()
+        all_relevant = set()
+        
+        for result in results:
+            for doc in result.retrieved_docs:
+                doc_id = doc.page_content[:50]
+                all_retrieved.add(doc_id)
+            
+            for doc in result.relevant_docs:
+                doc_id = doc.page_content[:50]
+                all_relevant.add(doc_id)
+        
+        coverage = len(all_retrieved) / len(all_relevant) if all_relevant else 0
+        return coverage
+    
+    def _calculate_diversity(self, results: List[RetrievalResult]) -> float:
+        """计算多样性 - 检索结果之间的平均语义差异"""
+        all_similarities = []
+        
+        for result in results:
+            if len(result.retrieved_docs) < 2:
+                continue
+                
+            # 获取文档嵌入
+            doc_texts = [doc.page_content for doc in result.retrieved_docs]
+            embeddings = self.embedding_model.encode(doc_texts, convert_to_tensor=True)
+            
+            # 计算文档之间的余弦相似度
+            cos_sim = util.pytorch_cos_sim(embeddings, embeddings)
+            
+            # 获取上三角矩阵（排除对角线）
+            upper_triangle_indices = torch.triu_indices(len(cos_sim), len(cos_sim), offset=1)
+            similarities = cos_sim[upper_triangle_indices[0], upper_triangle_indices[1]]
+            
+            # 多样性 = 1 - 平均相似度
+            diversity = 1 - similarities.mean().item()
+            all_similarities.append(diversity)
+        
+        return sum(all_similarities) / len(all_similarities) if all_similarities else 0
+    
+    def _calculate_novelty(self, results: List[RetrievalResult]) -> float:
+        """计算新颖性 - 检索结果中不重复内容的比例"""
+        total_docs = 0
+        unique_docs = set()
+        
+        for result in results:
+            for doc in result.retrieved_docs:
+                total_docs += 1
+                doc_id = doc.page_content[:50]
+                unique_docs.add(doc_id)
+        
+        novelty = len(unique_docs) / total_docs if total_docs > 0 else 0
+        return novelty
+    
+    def compare_retrievers(self, retriever_results: Dict[str, List[RetrievalResult]], 
+                          k_values: List[int] = [1, 3, 5, 10]) -> Dict[str, EvaluationMetrics]:
+        """
+        比较多个检索器的性能
+        
+        Args:
+            retriever_results: 检索器名称到检索结果的映射
+            k_values: 要计算的k值列表
+            
+        Returns:
+            检索器名称到评估指标的映射
+        """
+        metrics = {}
+        
+        for name, results in retriever_results.items():
+            print(f"评估检索器: {name}")
+            metrics[name] = self.evaluate_retrieval(results, k_values)
+        
+        return metrics
+    
+    def generate_report(self, metrics: Dict[str, EvaluationMetrics], 
+                        save_path: Optional[str] = None) -> str:
+        """
+        生成评估报告
+        
+        Args:
+            metrics: 检索器名称到评估指标的映射
+            save_path: 报告保存路径
+            
+        Returns:
+            报告文本
+        """
+        report = []
+        report.append("# 检索效果评估报告\n")
+        
+        # 创建比较表
+        df_data = []
+        for name, metric in metrics.items():
+            row = {"检索器": name}
+            row.update({
+                f"Precision@{k}": f"{metric.precision_at_k[k]:.4f}" 
+                for k in sorted(metric.precision_at_k.keys())
+            })
+            row.update({
+                f"Recall@{k}": f"{metric.recall_at_k[k]:.4f}" 
+                for k in sorted(metric.recall_at_k.keys())
+            })
+            row.update({
+                f"F1@{k}": f"{metric.f1_at_k[k]:.4f}" 
+                for k in sorted(metric.f1_at_k.keys())
+            })
+            row.update({
+                f"NDCG@{k}": f"{metric.ndcg_at_k[k]:.4f}" 
+                for k in sorted(metric.ndcg_at_k.keys())
+            })
+            row.update({
+                "MAP": f"{metric.map_score:.4f}",
+                "MRR": f"{metric.mrr:.4f}",
+                "覆盖率": f"{metric.coverage:.4f}",
+                "多样性": f"{metric.diversity:.4f}",
+                "新颖性": f"{metric.novelty:.4f}",
+                "延迟(ms)": f"{metric.latency*1000:.2f}"
+            })
+            df_data.append(row)
+        
+        df = pd.DataFrame(df_data)
+        report.append("## 指标比较表\n")
+        report.append(df.to_string(index=False))
+        report.append("\n\n")
+        
+        # 添加指标解释
+        report.append("## 指标解释\n")
+        report.append("- **Precision@K**: 前K个结果中相关文档的比例\n")
+        report.append("- **Recall@K**: 前K个结果中相关文档占所有相关文档的比例\n")
+        report.append("- **F1@K**: Precision和Recall的调和平均数\n")
+        report.append("- **NDCG@K**: 归一化折扣累积增益，考虑排序位置\n")
+        report.append("- **MAP**: 平均精度均值，所有查询的平均精度\n")
+        report.append("- **MRR**: 平均倒数排名，第一个相关文档排名的倒数平均值\n")
+        report.append("- **覆盖率**: 检索到的唯一文档数与总文档数的比例\n")
+        report.append("- **多样性**: 检索结果之间的平均语义差异\n")
+        report.append("- **新颖性**: 检索结果中不重复内容的比例\n")
+        report.append("- **延迟**: 平均检索时间\n")
+        
+        # 添加最佳检索器
+        report.append("## 最佳检索器\n")
+        
+        # 找出每个指标的最佳检索器
+        best_metrics = {}
+        for metric_name in ["precision_at_5", "recall_at_5", "f1_at_5", "ndcg_at_5", "map_score", "mrr"]:
+            best_name = max(metrics.keys(), key=lambda x: getattr(metrics[x], metric_name))
+            best_metrics[metric_name] = best_name
+            report.append(f"- **{metric_name}**: {best_name}\n")
+        
+        report_text = "".join(report)
+        
+        # 保存报告
+        if save_path:
+            with open(save_path, "w", encoding="utf-8") as f:
+                f.write(report_text)
+            print(f"报告已保存到: {save_path}")
+        
+        return report_text
+    
+    def plot_metrics_comparison(self, metrics: Dict[str, EvaluationMetrics], 
+                              save_path: Optional[str] = None):
+        """
+        绘制指标比较图
+        
+        Args:
+            metrics: 检索器名称到评估指标的映射
+            save_path: 图表保存路径
+        """
+        # 准备数据
+        retriever_names = list(metrics.keys())
+        
+        # 创建子图
+        fig, axes = plt.subplots(2, 3, figsize=(18, 12))
+        fig.suptitle("检索器性能比较", fontsize=16)
+        
+        # Precision@K
+        ax = axes[0, 0]
+        k_values = sorted(list(metrics[retriever_names[0]].precision_at_k.keys()))
+        for name in retriever_names:
+            precision_values = [metrics[name].precision_at_k[k] for k in k_values]
+            ax.plot(k_values, precision_values, marker='o', label=name)
+        ax.set_title("Precision@K")
+        ax.set_xlabel("K")
+        ax.set_ylabel("Precision")
+        ax.legend()
+        ax.grid(True)
+        
+        # Recall@K
+        ax = axes[0, 1]
+        for name in retriever_names:
+            recall_values = [metrics[name].recall_at_k[k] for k in k_values]
+            ax.plot(k_values, recall_values, marker='o', label=name)
+        ax.set_title("Recall@K")
+        ax.set_xlabel("K")
+        ax.set_ylabel("Recall")
+        ax.legend()
+        ax.grid(True)
+        
+        # F1@K
+        ax = axes[0, 2]
+        for name in retriever_names:
+            f1_values = [metrics[name].f1_at_k[k] for k in k_values]
+            ax.plot(k_values, f1_values, marker='o', label=name)
+        ax.set_title("F1@K")
+        ax.set_xlabel("K")
+        ax.set_ylabel("F1")
+        ax.legend()
+        ax.grid(True)
+        
+        # NDCG@K
+        ax = axes[1, 0]
+        for name in retriever_names:
+            ndcg_values = [metrics[name].ndcg_at_k[k] for k in k_values]
+            ax.plot(k_values, ndcg_values, marker='o', label=name)
+        ax.set_title("NDCG@K")
+        ax.set_xlabel("K")
+        ax.set_ylabel("NDCG")
+        ax.legend()
+        ax.grid(True)
+        
+        # MAP和MRR
+        ax = axes[1, 1]
+        map_values = [metrics[name].map_score for name in retriever_names]
+        mrr_values = [metrics[name].mrr for name in retriever_names]
+        x = np.arange(len(retriever_names))
+        width = 0.35
+        ax.bar(x - width/2, map_values, width, label='MAP')
+        ax.bar(x + width/2, mrr_values, width, label='MRR')
+        ax.set_title("MAP和MRR")
+        ax.set_xticks(x)
+        ax.set_xticklabels(retriever_names)
+        ax.legend()
+        ax.grid(True)
+        
+        # 其他指标
+        ax = axes[1, 2]
+        other_metrics = ['coverage', 'diversity', 'novelty']
+        metric_values = {metric: [] for metric in other_metrics}
+        for name in retriever_names:
+            for metric in other_metrics:
+                metric_values[metric].append(getattr(metrics[name], metric))
+        
+        x = np.arange(len(retriever_names))
+        width = 0.25
+        for i, metric in enumerate(other_metrics):
+            ax.bar(x + i*width, metric_values[metric], width, label=metric)
+        ax.set_title("其他指标")
+        ax.set_xticks(x + width)
+        ax.set_xticklabels(retriever_names)
+        ax.legend()
+        ax.grid(True)
+        
+        plt.tight_layout()
+        
+        # 保存图表
+        if save_path:
+            plt.savefig(save_path, dpi=300, bbox_inches='tight')
+            print(f"图表已保存到: {save_path}")
+        
+        plt.show()
+
+
+class RetrievalTestSet:
+    """检索测试集"""
+    
+    def __init__(self, queries_file: str, documents_file: str, qrels_file: str):
+        """
+        初始化测试集
+        
+        Args:
+            queries_file: 查询文件路径，每行一个查询
+            documents_file: 文档文件路径，每行一个文档
+            qrels_file: 相关性标注文件路径，格式为: query_id,doc_id,relevance
+        """
+        self.queries = self._load_queries(queries_file)
+        self.documents = self._load_documents(documents_file)
+        self.qrels = self._load_qrels(qrels_file)
+    
+    def _load_queries(self, file_path: str) -> Dict[str, str]:
+        """加载查询"""
+        queries = {}
+        with open(file_path, 'r', encoding='utf-8') as f:
+            for i, line in enumerate(f):
+                queries[str(i)] = line.strip()
+        return queries
+    
+    def _load_documents(self, file_path: str) -> Dict[str, Document]:
+        """加载文档"""
+        documents = {}
+        with open(file_path, 'r', encoding='utf-8') as f:
+            for i, line in enumerate(f):
+                doc = Document(page_content=line.strip(), metadata={"doc_id": str(i)})
+                documents[str(i)] = doc
+        return documents
+    
+    def _load_qrels(self, file_path: str) -> Dict[str, Dict[str, int]]:
+        """加载相关性标注"""
+        qrels = {}
+        with open(file_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                parts = line.strip().split(',')
+                if len(parts) >= 3:
+                    query_id, doc_id, relevance = parts[0], parts[1], int(parts[2])
+                    if query_id not in qrels:
+                        qrels[query_id] = {}
+                    qrels[query_id][doc_id] = relevance
+        return qrels
+    
+    def get_retrieval_results(self, retriever, top_k: int = 10) -> List[RetrievalResult]:
+        """
+        使用指定检索器获取检索结果
+        
+        Args:
+            retriever: 检索器，需要有一个retrieve(query, top_k)方法
+            top_k: 返��的文档数量
+            
+        Returns:
+            检索结果列表
+        """
+        results = []
+        
+        for query_id, query_text in self.queries.items():
+            start_time = time.time()
+            retrieved_docs = retriever.retrieve(query_text, top_k)
+            retrieval_time = time.time() - start_time
+            
+            # 获取相关文档
+            relevant_docs = []
+            if query_id in self.qrels:
+                for doc_id, relevance in self.qrels[query_id].items():
+                    if relevance > 0 and doc_id in self.documents:
+                        relevant_docs.append(self.documents[doc_id])
+            
+            result = RetrievalResult(
+                query=query_text,
+                retrieved_docs=retrieved_docs,
+                relevant_docs=relevant_docs,
+                retrieval_time=retrieval_time
+            )
+            results.append(result)
+        
+        return results
+
+
+def create_sample_test_set():
+    """创建示例测试集"""
+    # 创建示例查询
+    queries = [
+        "什么是机器学习？",
+        "深度学习和机器学习的区别是什么？",
+        "如何评估机器学习模型的性能？",
+        "自然语言处理有哪些应用？",
+        "计算机视觉的基本任务是什么？"
+    ]
+    
+    # 创建示例文档
+    documents = [
+        "机器学习是人工智能的一个分支，它使计算机能够在没有明确编程的情况下学习和改进。",
+        "深度学习是机器学习的一个子集，它使用多层神经网络来模拟人脑的工作方式。",
+        "评估机器学习模型的常用指标包括准确率、精确率、召回率和F1分数。",
+        "自然语言处理是计算机科学和人工智能的一个分支，专注于计算机与人类语言之间的交互。",
+        "计算机视觉是人工智能的一个领域，训练计算机解释和理解视觉世界。",
+        "强化学习是机器学习的一个类型，它关注软件代理应该如何在环境中采取行动以最大化累积奖励。",
+        "数据预处理是机器学习流程中的重要步骤，包括数据清洗、特征选择和特征工程。",
+        "过拟合是机器学习中的一个常见问题，指模型在训练数据上表现良好但在新数据上表现不佳。",
+        "卷积神经网络（CNN）是一类深度神经网络，最常用于分析视觉图像。",
+        "循环神经网络（RNN）是一类人工神经网络，其中节点之间的连接形成有向图沿时间序列。"
+    ]
+    
+    # 创建相关性标注
+    qrels = {
+        "0": {"0": 2, "1": 1, "6": 1, "7": 1},  # 什么是机器学习？
+        "1": {"0": 1, "1": 2, "8": 1, "9": 1},  # 深度学习和机器学习的区别
+        "2": {"2": 2, "7": 1},  # 如何评估机器学习模型的性能
+        "3": {"3": 2, "9": 1},  # 自然语言处理的应用
+        "4": {"4": 2, "8": 1}   # 计算机视觉的基本任务
+    }
+    
+    # 保存文件
+    with open("sample_queries.txt", "w", encoding="utf-8") as f:
+        for query in queries:
+            f.write(query + "\n")
+    
+    with open("sample_documents.txt", "w", encoding="utf-8") as f:
+        for doc in documents:
+            f.write(doc + "\n")
+    
+    with open("sample_qrels.csv", "w", encoding="utf-8") as f:
+        for query_id, doc_relevance in qrels.items():
+            for doc_id, relevance in doc_relevance.items():
+                f.write(f"{query_id},{doc_id},{relevance}\n")
+    
+    print("示例测试集已创建:")
+    print("- sample_queries.txt: 查询文件")
+    print("- sample_documents.txt: 文档文件")
+    print("- sample_qrels.csv: 相关性标注文件")
+    
+    return RetrievalTestSet("sample_queries.txt", "sample_documents.txt", "sample_qrels.csv")
+
+
+if __name__ == "__main__":
+    # 创建示例测试集
+    test_set = create_sample_test_set()
+    
+    # 创建评估器
+    evaluator = RetrievalEvaluator()
+    
+    # 这里应该使用您的实际检索器
+    # 以下是一个模拟的检索器，用于演示
+    class MockRetriever:
+        def __init__(self, name):
+            self.name = name
+        
+        def retrieve(self, query, top_k=10):
+            # 模拟检索结果
+            import random
+            all_docs = list(test_set.documents.values())
+            # 模拟不同质量的检索器
+            if self.name == "good":
+                # 好的检索器：有更高概率返回相关文档
+                relevant_docs = [doc for doc in all_docs if any(keyword in doc.page_content.lower() 
+                                for keyword in query.lower().split()[:2])]
+                if relevant_docs:
+                    results = relevant_docs[:min(top_k//2, len(relevant_docs))]
+                    results += random.sample(all_docs, min(top_k-len(results), len(all_docs)))
+                else:
+                    results = random.sample(all_docs, min(top_k, len(all_docs)))
+            elif self.name == "medium":
+                # 中等检索器
+                relevant_docs = [doc for doc in all_docs if any(keyword in doc.page_content.lower() 
+                                for keyword in [query.lower().split()[0]])]
+                if relevant_docs:
+                    results = relevant_docs[:min(top_k//3, len(relevant_docs))]
+                    results += random.sample(all_docs, min(top_k-len(results), len(all_docs)))
+                else:
+                    results = random.sample(all_docs, min(top_k, len(all_docs)))
+            else:
+                # 差的检索器：随机返回
+                results = random.sample(all_docs, min(top_k, len(all_docs)))
+            
+            return results
+    
+    # 创建不同质量的检索器
+    good_retriever = MockRetriever("good")
+    medium_retriever = MockRetriever("medium")
+    poor_retriever = MockRetriever("poor")
+    
+    # 获取检索结果
+    good_results = test_set.get_retrieval_results(good_retriever)
+    medium_results = test_set.get_retrieval_results(medium_retriever)
+    poor_results = test_set.get_retrieval_results(poor_retriever)
+    
+    # 比较检索器
+    retriever_results = {
+        "好的检索器": good_results,
+        "中等检索器": medium_results,
+        "差的检索器": poor_results
+    }
+    
+    # 评估检索器
+    metrics = evaluator.compare_retrievers(retriever_results)
+    
+    # 生成报告
+    report = evaluator.generate_report(metrics, "retrieval_evaluation_report.md")
+    print(report)
+    
+    # 绘制比较图
+    evaluator.plot_metrics_comparison(metrics, "retrieval_evaluation_comparison.png")

workflow_nodes.py

@@ -3,6 +3,7 @@
 包含所有工作流节点函数和状态管理
 """
 
+import time
 from typing import List
 from typing_extensions import TypedDict
 try:
@@ -19,6 +20,7 @@ except ImportError:
 
 from config import LOCAL_LLM, WEB_SEARCH_RESULTS_COUNT, ENABLE_HYBRID_SEARCH, ENABLE_QUERY_EXPANSION, ENABLE_MULTIMODAL
 from document_processor import DocumentProcessor
+from retrieval_evaluation import RetrievalEvaluator, RetrievalResult
 from pprint import pprint
 
 
@@ -31,11 +33,13 @@ class GraphState(TypedDict):
         generation: LLM生成
         documents: 文档列表
         retry_count: 重试计数器，防止无限循环
+        retrieval_metrics: 检索评估指标
     """
     question: str
     generation: str
     documents: List[str]
     retry_count: int
+    retrieval_metrics: dict  # 添加检索评估指标
 
 
 class WorkflowNodes:
@@ -46,6 +50,9 @@ class WorkflowNodes:
         self.retriever = retriever if retriever is not None else getattr(doc_processor, 'retriever', None)
         self.graders = graders
         
+        # 初始化检索评估器
+        self.retrieval_evaluator = RetrievalEvaluator()
+        
         # 设置RAG链 - 使用本地提示模板
         rag_prompt_template = PromptTemplate(
             template="""你是一个问答助手。使用以下检索到的上下文来回答问题。
@@ -77,6 +84,7 @@ class WorkflowNodes:
         print("---检索---")
         question = state["question"]
         retry_count = state.get("retry_count", 0)
+        retrieval_start_time = time.time()
         
         # 使用增强检索方法，支持混合检索、查询扩展和多模态
         try:
@@ -118,8 +126,19 @@ class WorkflowNodes:
             except Exception as fallback_e:
                 print(f"❌ 回退检索也失败: {fallback_e}")
                 documents = []
-            
-        return {"documents": documents, "question": question, "retry_count": retry_count}
+        
+        # 计算检索时间
+        retrieval_time = time.time() - retrieval_start_time
+        
+        # 评估检索结果
+        retrieval_metrics = self._evaluate_retrieval_results(question, documents, retrieval_time)
+        
+        return {
+            "documents": documents, 
+            "question": question, 
+            "retry_count": retry_count,
+            "retrieval_metrics": retrieval_metrics
+        }
     
     def generate(self, state):
         """
@@ -295,6 +314,73 @@ class WorkflowNodes:
             return "not supported"
 
 
+def _evaluate_retrieval_results(self, question, documents, retrieval_time):
+        """
+        评估检索结果的质量
+        
+        Args:
+            question: 查询问题
+            documents: 检索到的文档
+            retrieval_time: 检索耗时
+            
+        Returns:
+            dict: 评估指标
+        """
+        try:
+            # 创建模拟的相关文档（在实际应用中，这些应该是真实的相关文档）
+            # 这里我们假设前几个文档是相关的，用于演示评估功能
+            relevant_docs = documents[:min(2, len(documents))] if documents else []
+            
+            # 创建检索结果对象
+            retrieval_result = RetrievalResult(
+                query=question,
+                retrieved_docs=documents,
+                relevant_docs=relevant_docs,
+                retrieval_time=retrieval_time
+            )
+            
+            # 评估检索结果
+            metrics = self.retrieval_evaluator.evaluate_retrieval([retrieval_result], k_values=[1, 3, 5])
+            
+            # 提取关键指标
+            result_metrics = {
+                "precision_at_1": metrics.precision_at_k.get(1, 0),
+                "precision_at_3": metrics.precision_at_k.get(3, 0),
+                "precision_at_5": metrics.precision_at_k.get(5, 0),
+                "recall_at_1": metrics.recall_at_k.get(1, 0),
+                "recall_at_3": metrics.recall_at_k.get(3, 0),
+                "recall_at_5": metrics.recall_at_k.get(5, 0),
+                "map_score": metrics.map_score,
+                "mrr": metrics.mrr,
+                "latency": metrics.latency,
+                "retrieved_docs_count": len(documents)
+            }
+            
+            # 打印评估结果
+            print("\n---检索评估结果---")
+            print(f"检索耗时: {result_metrics['latency']:.4f}秒")
+            print(f"检索文档数: {result_metrics['retrieved_docs_count']}")
+            print(f"Precision@1: {result_metrics['precision_at_1']:.4f}")
+            print(f"Precision@3: {result_metrics['precision_at_3']:.4f}")
+            print(f"Precision@5: {result_metrics['precision_at_5']:.4f}")
+            print(f"Recall@1: {result_metrics['recall_at_1']:.4f}")
+            print(f"Recall@3: {result_metrics['recall_at_3']:.4f}")
+            print(f"Recall@5: {result_metrics['recall_at_5']:.4f}")
+            print(f"MAP: {result_metrics['map_score']:.4f}")
+            print(f"MRR: {result_metrics['mrr']:.4f}")
+            print("--------------------\n")
+            
+            return result_metrics
+            
+        except Exception as e:
+            print(f"⚠️ 检索评估失败: {e}")
+            return {
+                "error": str(e),
+                "latency": retrieval_time,
+                "retrieved_docs_count": len(documents)
+            }
+
+
 def format_docs(docs):
     """格式化文档用于显示"""
     return "\n\n".join(doc.page_content for doc in docs)