""" 重排业务逻辑模块 本模块包含 RAG 相关的重排业务逻辑(文档处理、排序、top_n) 使用 model_services/rerank_services.py 提供的纯服务层 """ import logging from typing import List from langchain_core.documents import Document from ..model_services import get_rerank_service logger = logging.getLogger(__name__) class DocumentReranker: """ 文档重排器 - 业务逻辑层 负责: - 从 Document 提取内容 - 调用 rerank service 获取得分 - 根据得分排序 - 返回 top_n 文档 """ def __init__(self, rerank_service=None): """ 初始化文档重排器 Args: rerank_service: 重排服务(可选,默认通过 get_rerank_service() 获取) """ self._rerank_service = rerank_service or get_rerank_service() def compress_documents( self, documents: List[Document], query: str, top_n: int = 5 ) -> List[Document]: """ 对文档进行重排 - 业务逻辑 Args: documents: 待排序的文档列表 query: 查询字符串 top_n: 返回前 N 个结果 Returns: List[Document]: 排序后的文档列表,每个文档的 metadata 中包含 rerank_score """ if not documents: return [] try: # 1. 从 Document 提取内容 doc_contents = [doc.page_content for doc in documents] logger.info(f"[Rerank] 收到 {len(documents)} 个文档待重排") # 2. 调用重排服务计算得分 scores = self._rerank_service.compute_scores(query, doc_contents) logger.info(f"[Rerank] 获取到 {len(scores)} 个得分") # 3. 构建 (文档, 分数) 对并排序 doc_score_pairs = list(zip(documents, scores)) doc_score_pairs_sorted = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True) # 4. 取 top_n,并添加 rerank_score 到 metadata top_docs = [] for doc, score in doc_score_pairs_sorted[:top_n]: # 创建新文档,添加 rerank_score new_doc = Document( page_content=doc.page_content, metadata={**doc.metadata, "rerank_score": score} ) top_docs.append(new_doc) return top_docs except Exception as e: logger.warning(f"[Rerank] 重排失败,返回原始结果: {e}") return documents[:top_n] def create_document_reranker(rerank_service=None) -> DocumentReranker: """ 创建文档重排器的工厂函数 Args: rerank_service: 重排服务(可选) Returns: DocumentReranker: 文档重排器实例 """ return DocumentReranker(rerank_service)