LQAgentPlatform/foundation/ai/rag/retrieval/retrieval.py

from typing import List, Dict, Any, Optional
from foundation.ai.models.rerank_model import rerank_model
from foundation.infrastructure.config.config import config_handler
from foundation.observability.logger.loggering import server_logger
from foundation.database.base.vector.milvus_vector import MilvusVectorManager

class RetrievalManager:
    """
    召回管理器，实现多路召回功能
    """

    def __init__(self):
        """
        初始化召回管理器
        """
        self.vector_manager = MilvusVectorManager()
        self.logger = server_logger
        self.dense_weight = config_handler.get('hybrid_search', 'DENSE_WEIGHT', 0.7)
        self.sparse_weight = config_handler.get('hybrid_search', 'SPARSE_WEIGHT', 0.3)

    def hybrid_search_recall(self, collection_name: str, query_text: str,
                           top_k: int = 10, ranker_type: str = "weighted",
                           dense_weight: float = 0.7, sparse_weight: float = 0.3) -> List[Dict[str, Any]]:
        """
        混合搜索召回 - 向量+BM25召回

        Args:
            collection_name: 集合名称
            query_text: 查询文本
            top_k: 返回结果数量
            ranker_type: 重排序类型 "weighted" 或 "rrf"
            dense_weight: 密集向量权重
            sparse_weight: 稀疏向量权重

        Returns:
            List[Dict]: 搜索结果列表
        """
        try:
            self.logger.info(f"开始混合搜索召回，查询: {query_text}")

            param = {'collection_name': collection_name}
            results = self.vector_manager.hybrid_search(
                param=param,
                query_text=query_text,
                top_k=top_k,
                ranker_type=ranker_type,
                dense_weight=dense_weight,
                sparse_weight=sparse_weight
            )

            self.logger.info(f"混合搜索召回返回 {len(results)} 个结果")
            return results

        except Exception as e:
            self.logger.error(f"混合搜索召回失败: {str(e)}")
            return []

    def rerank_recall(self, candidates: List[str], query_text: str,
                  top_k: int = None  ) -> List[Dict[str, Any]]:
        """
        重排序召回 - 使用BGE重排序模型对候选文档重新排序

        Args:
            candidates: 候选文档列表
            query_text: 查询文本
            top_k: 返回结果数量

        Returns:
            List[Dict]: 重排序后的结果列表
        """
        try:
            self.logger.info(f"开始重排序召回，候选文档数量: {len(candidates)}")

            # 调用重排序执行器
            rerank_results = rerank_model.bge_rerank(query_text, candidates, top_k)

            # 转换结果格式
            scored_docs = []
            for i, result in enumerate(rerank_results):
                scored_docs.append({
                    'text_content': result.get('text', ''),
                    'rerank_score': float(result.get('score', '0.0')),
                    'index': i
                })

            self.logger.info(f"重排序召回返回 {len(scored_docs)} 个结果")
            return scored_docs

        except Exception as e:
            self.logger.error(f"重排序召回失败: {str(e)}")
            return []

    def multi_stage_recall(self, collection_name: str, query_text: str,
                          hybrid_top_k: int = 50, top_k: int = 10,
                          ranker_type: str = "weighted") -> List[Dict[str, Any]]:
        """
        多路召回 - 先混合搜索召回，再重排序，只返回重排序结果

        Args:
            collection_name: 集合名称
            query_text: 查询文本
            hybrid_top_k: 混合搜索召回的文档数量
            top_k: 最终返回的文档数量
            ranker_type: 混合搜索的重排序类型

        Returns:
            List[Dict]: 重排序后的结果列表，只包含重排序分数
        """
        try:
            self.logger.info(f"开始多路召回，查询: {query_text}")

            # 第一阶段：混合搜索召回（向量+BM25）
            hybrid_results = self.hybrid_search_recall(
                collection_name=collection_name,
                query_text=query_text,
                top_k=hybrid_top_k,
                ranker_type=ranker_type
            )

            if not hybrid_results:
                self.logger.warning("混合搜索召回无结果，返回空列表")
                return []

            # 提取候选文档文本
            candidates = [result['text_content'] for result in hybrid_results]

            # 第二阶段：重排序召回
            rerank_results = self.rerank_recall(
                candidates=candidates,
                query_text=query_text,
                top_k=top_k
            )

            # 为重排序结果添加混合搜索的原始元数据
            for rerank_result in rerank_results:
                original_index = rerank_result.get('index', 0)
                if original_index < len(hybrid_results):
                    original_metadata = hybrid_results[original_index].get('metadata', {})
                    rerank_result['metadata'] = original_metadata

            self.logger.info(f"多路召回完成，返回 {len(rerank_results)} 个重排序结果")
            return rerank_results

        except Exception as e:
            self.logger.error(f"多路召回失败: {str(e)}")
            return []

    # 创建全局召回管理器实例
retrieval_manager = RetrievalManager()