341 lines
11 KiB
Python
341 lines
11 KiB
Python
"""
|
||
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) |