""" RAG 检索流水线 组合检索器、重排序器、查询改写器等组件，提供完整的 RAG 检索功能。 """ import time from typing import List, Dict, Any, Optional, Union from dataclasses import dataclass, field from enum import Enum from langchain_core.documents import Document from langchain_core.embeddings import Embeddings from langchain_core.language_models import BaseLanguageModel from .retriever import ( create_base_retriever, create_hybrid_retriever, create_ensemble_retriever, create_qdrant_client, ) from .reranker import CrossEncoderReranker from .query_transform import MultiQueryTransformer, create_rag_fusion_pipeline class RAGLevel(Enum): """RAG 功能级别""" BASIC = 1 # 基础向量搜索 HYBRID = 2 # 混合检索 + 重排序 FUSION = 3 # RAG-Fusion AGENTIC = 4 # Agentic RAG @dataclass class RAGConfig: """RAG 配置""" # Qdrant 配置 collection_name: str = "documents" qdrant_url: Optional[str] = None qdrant_api_key: Optional[str] = None # 检索配置 rag_level: RAGLevel = RAGLevel.FUSION dense_k: int = 10 # 向量检索数量 sparse_k: int = 10 # BM25 检索数量 total_k: int = 20 # 总检索数量 rerank_top_n: int = 5 # 重排序返回数量 # 查询改写配置 num_queries: int = 3 # RAG-Fusion 查询数量 # 模型配置 reranker_model: str = "BAAI/bge-reranker-base" device: Optional[str] = None # 性能配置 enable_cache: bool = True verbose: bool = True @dataclass class RetrievalResult: """检索结果""" documents: List[Document] query_time: float level: RAGLevel metadata: Dict[str, Any] = field(default_factory=dict) class RAGPipeline: """ RAG 检索流水线 支持从 Level 1 到 Level 4 的所有功能。 """ def __init__( self, embeddings: Embeddings, llm: Optional[BaseLanguageModel] = None, config: Optional[RAGConfig] = None, ): """ 初始化 RAG 流水线 Args: embeddings: 嵌入模型 llm: 语言模型（用于查询改写，Level 3+ 需要） config: 配置 """ self.embeddings = embeddings self.llm = llm self.config = config or RAGConfig() # 初始化组件 self._client = None self._reranker = None self._query_transformer = None self._retriever = None # 缓存 self._cache = {} def _get_client(self): """获取 Qdrant 客户端""" if self._client is None: self._client = create_qdrant_client( url=self.config.qdrant_url, api_key=self.config.qdrant_api_key, ) return self._client def _get_reranker(self): """获取重排序器""" if self._reranker is None: self._reranker = CrossEncoderReranker( model_name=self.config.reranker_model, top_n=self.config.rerank_top_n, device=self.config.device, ) return self._reranker def _get_query_transformer(self): """获取查询改写器""" if self._query_transformer is None and self.llm is not None: self._query_transformer = MultiQueryTransformer( llm=self.llm, num_queries=self.config.num_queries, ) return self._query_transformer def _create_basic_retriever(self): """创建基础检索器（Level 1）""" return create_base_retriever( collection_name=self.config.collection_name, embeddings=self.embeddings, search_kwargs={"k": self.config.total_k}, client=self._get_client(), ) def _create_hybrid_retriever(self): """创建混合检索器（Level 2）""" base_retriever = create_hybrid_retriever( collection_name=self.config.collection_name, embeddings=self.embeddings, dense_k=self.config.dense_k, sparse_k=self.config.sparse_k, client=self._get_client(), ) # 应用重排序 reranker = self._get_reranker() return reranker.create_contextual_compression_retriever(base_retriever) def _create_fusion_retriever(self): """创建 RAG-Fusion 检索器（Level 3）""" if self.llm is None: raise ValueError("Level 3 (RAG-Fusion) 需要语言模型进行查询改写") # 创建基础混合检索器 base_retriever = create_hybrid_retriever( collection_name=self.config.collection_name, embeddings=self.embeddings, dense_k=self.config.dense_k, sparse_k=self.config.sparse_k, client=self._get_client(), ) # 创建 RAG-Fusion 流水线 reranker = self._get_reranker() return create_rag_fusion_pipeline( base_retriever=base_retriever, llm=self.llm, reranker=reranker, num_queries=self.config.num_queries, ) def _get_retriever(self): """根据配置级别获取检索器""" if self._retriever is None: if self.config.rag_level == RAGLevel.BASIC: self._retriever = self._create_basic_retriever() elif self.config.rag_level == RAGLevel.HYBRID: self._retriever = self._create_hybrid_retriever() elif self.config.rag_level == RAGLevel.FUSION: self._retriever = self._create_fusion_retriever() elif self.config.rag_level == RAGLevel.AGENTIC: # Agentic RAG 使用 Fusion 作为基础，在 tools.py 中包装 self._retriever = self._create_fusion_retriever() else: raise ValueError(f"不支持的 RAG 级别: {self.config.rag_level}") return self._retriever def retrieve( self, query: str, use_cache: Optional[bool] = None, **kwargs, ) -> RetrievalResult: """ 执行检索 Args: query: 查询文本 use_cache: 是否使用缓存 **kwargs: 额外参数 Returns: 检索结果 """ start_time = time.time() # 检查缓存 if use_cache is None: use_cache = self.config.enable_cache cache_key = f"{query}:{self.config.rag_level.value}" if use_cache and cache_key in self._cache: if self.config.verbose: print(f"使用缓存结果: {query}") return self._cache[cache_key] # 获取检索器并执行检索 retriever = self._get_retriever() documents = retriever.invoke(query, **kwargs) # 计算查询时间 query_time = time.time() - start_time # 创建结果 result = RetrievalResult( documents=documents, query_time=query_time, level=self.config.rag_level, metadata={ "query": query, "collection": self.config.collection_name, "doc_count": len(documents), }, ) # 缓存结果 if use_cache: self._cache[cache_key] = result if self.config.verbose: print(f"检索完成: {len(documents)} 文档, 耗时: {query_time:.2f}s") return result def format_context( self, documents: List[Document], max_length: Optional[int] = None, ) -> str: """ 格式化检索到的文档为上下文文本 Args: documents: 文档列表 max_length: 最大长度（字符数） Returns: 格式化后的上下文文本 """ context_parts = [] total_length = 0 for i, doc in enumerate(documents): # 提取内容和元数据 content = doc.page_content.strip() metadata = doc.metadata # 格式化文档 doc_text = f"[文档 {i+1}]\n" if metadata.get("source"): doc_text += f"来源: {metadata['source']}\n" if metadata.get("page"): doc_text += f"页码: {metadata['page']}\n" doc_text += f"内容: {content}\n\n" # 检查长度限制 if max_length is not None: if total_length + len(doc_text) > max_length: # 如果添加这个文档会超限，则截断并添加说明 remaining = max_length - total_length if remaining > 100: # 至少保留100字符 doc_text = doc_text[:remaining] + "...\n\n[内容已截断]" context_parts.append(doc_text) break else: break context_parts.append(doc_text) total_length += len(doc_text) return "".join(context_parts).strip() def clear_cache(self): """清空缓存""" self._cache.clear() @classmethod def create_from_config( cls, embeddings: Embeddings, llm: Optional[BaseLanguageModel] = None, config_dict: Optional[Dict[str, Any]] = None, ) -> "RAGPipeline": """ 从配置字典创建流水线 Args: embeddings: 嵌入模型 llm: 语言模型 config_dict: 配置字典 Returns: RAGPipeline 实例 """ config_dict = config_dict or {} # 创建配置对象 config = RAGConfig( collection_name=config_dict.get("collection_name", "documents"), qdrant_url=config_dict.get("qdrant_url"), qdrant_api_key=config_dict.get("qdrant_api_key"), rag_level=RAGLevel(config_dict.get("rag_level", RAGLevel.FUSION.value)), dense_k=config_dict.get("dense_k", 10), sparse_k=config_dict.get("sparse_k", 10), total_k=config_dict.get("total_k", 20), rerank_top_n=config_dict.get("rerank_top_n", 5), num_queries=config_dict.get("num_queries", 3), reranker_model=config_dict.get("reranker_model", "BAAI/bge-reranker-base"), device=config_dict.get("device"), enable_cache=config_dict.get("enable_cache", True), verbose=config_dict.get("verbose", True), ) return cls(embeddings=embeddings, llm=llm, config=config)