evaluate_retrieval.py

"""
自适应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
try:
    from langchain_core.documents import Document
except ImportError:
    try:
        from langchain_core.documents import Document
    except ImportError:
        from langchain.schema import Document

# 导入LangChain相关模块
from langchain_community.vectorstores import FAISS, Chroma
from langchain_community.retrievers import BM25Retriever
try:
    from langchain.retrievers import EnsembleRetriever
    from langchain.retrievers import ContextualCompressionRetriever
    from langchain.retrievers.document_compressors import LLMChainExtractor
except ImportError:
    try:
        from langchain_core.retrievers import EnsembleRetriever
        from langchain_core.retrievers import ContextualCompressionRetriever
        from langchain.retrievers.document_compressors import LLMChainExtractor
    except ImportError:
        print("Warning: Could not import advanced retriever components. Some features may be limited.")
        EnsembleRetriever = None
        ContextualCompressionRetriever = None
        LLMChainExtractor = None


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()