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

import asyncio
import json
from typing import List, Dict, Any, Optional
from foundation.ai.models.rerank_model import rerank_model
from foundation.observability.monitoring.time_statistics import track_execution_time
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)

        # 重排序模型配置
        self.rerank_model_type = config_handler.get('retrieval', 'RERANK_MODEL_TYPE', 'bge')  # 'bge' 或 'qwen3'
        self.logger.info(f"初始化重排序模型类型: {self.rerank_model_type}")

    def set_rerank_model(self, model_type: str):
        """
        设置重排序模型类型

        Args:
            model_type: 模型类型 ('bge' 或 'qwen3')
        """
        if model_type not in ['bge', 'qwen3']:
            raise ValueError("model_type 必须是 'bge' 或 'qwen3'")

        self.rerank_model_type = model_type
        self.logger.info(f"重排序模型类型已设置为: {model_type}")

    def _clean_document(self, doc: str) -> str:
        """
        清理文档文本，移除HTML标签和特殊字符

        Args:
            doc: 原始文档文本

        Returns:
            str: 清理后的文档文本
        """
        if not isinstance(doc, str):
            self.logger.debug(f"文档类型转换: {type(doc)} -> str")
            return str(doc)

        original_length = len(doc)

        # 移除HTML标签
        import re
        doc = re.sub(r'<[^>]+>', '', doc)

        # 移除多余的空白字符
        doc = re.sub(r'\s+', ' ', doc)

        # 更宽松的字符过滤 - 保留更多字符
        doc = re.sub(r'[^\u4e00-\u9fff\w\s.,;:!?()（）。，；：！？\-\+\=\*/%&@#￥$【】「」""''""\n\r]', '', doc)

        # 截断过长的文本
        if len(doc) > 8000:  # 设置最大长度限制
            doc = doc[:8000] + "..."

        cleaned_doc = doc.strip()
        self.logger.debug(f"文档清理: {original_length} -> {len(cleaned_doc)} 字符")

        return cleaned_doc

    def _get_rerank_results(self, query_text: str, documents: List[str], top_k: int = None) -> List[Dict[str, Any]]:
        """
        根据配置选择重排序模型并执行重排序

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

        Returns:
            List[Dict]: 重排序后的结果列表
        """
        try:
            # 清理和验证文档列表
            cleaned_documents = []
            valid_original_docs = []

            for doc in documents:
                if doc and isinstance(doc, str) and doc.strip():
                    cleaned_doc = self._clean_document(doc)
                    if cleaned_doc and len(cleaned_doc) > 3:
                        cleaned_documents.append(cleaned_doc)
                        valid_original_docs.append(doc)

            if not cleaned_documents:
                return []

            if self.rerank_model_type == 'qwen3':
                self.logger.info("使用 Qwen3-Reranker-8B 进行重排序")
                rerank_results = rerank_model.qwen3_rerank(query_text, cleaned_documents, top_k)

                # 将清理后的文本映射回原始文本
                for result in rerank_results:
                    cleaned_text = result.get('text', '')
                    # 查找原始文本
                    for i, cleaned in enumerate(cleaned_documents):
                        if cleaned == cleaned_text:
                            result['text'] = valid_original_docs[i]
                            break

                return rerank_results
            else:
                self.logger.info("使用 BGE Reranker 进行重排序")
                rerank_results = rerank_model.bge_rerank(query_text, cleaned_documents, top_k)

                # 将清理后的文本映射回原始文本
                for result in rerank_results:
                    cleaned_text = result.get('text', '')
                    # 查找原始文本
                    for i, cleaned in enumerate(cleaned_documents):
                        if cleaned == cleaned_text:
                            result['text'] = valid_original_docs[i]
                            break

                return rerank_results

        except Exception as e:
            self.logger.error(f"重排序失败，模型类型: {self.rerank_model_type}, 错误: {str(e)}")
            # 返回原始顺序作为fallback
            return [{"text": doc, "score": 0.0} for i, doc in enumerate(documents[:top_k])]

    @track_execution_time
    async def entity_recall(self, main_entity: str,assisted_search_entity: list,
                   top_k: int = 5) -> List[Dict[str, Any]]:
        """
        执行实体召回
        :param main_entity: 查询实体
        :param assisted_search_entity: 辅助搜索实体
        :param top_k: 返回结果数量
        """
        collection_name = "first_bfp_collection_entity"
        # 主实体搜索 - 使用异步方法
        entity_result = await self.async_multi_stage_recall(
            collection_name=collection_name,
            query_text=main_entity,
            hybrid_top_k=20,  # 从默认50降到20
            top_k=top_k
        )

        assist_tasks = [
            self.async_multi_stage_recall(
                collection_name=collection_name,
                query_text=assisted_search_entity,
                hybrid_top_k=20,  # 从默认50降到20
                top_k=top_k
            ) for assisted_search_entity in assisted_search_entity
        ]
        # 辅助搜索,异步并发
        assist_results_list = await asyncio.gather(*assist_tasks,return_exceptions=True)
        assist_results = []
        for res in assist_results_list:
            if isinstance(res, Exception):
                self.logger.error(f"辅助实体召回失败: {str(res)}")
            else:
                assist_results.extend(res)

        all_results = entity_result + assist_results

        entity_list = list(set([item['text_content'] for item in all_results]))
        self.logger.info(f"entity_list:{entity_list}")

        return entity_list
    
    @track_execution_time
    async def async_bfp_recall(self, entity_list: List[str],background: str ,
                           top_k: int = 3,) -> List[Dict[str, Any]]:
        """
        混合搜索召回 - 向量+BM25召回

        Args:
            entity_list: 实体列表
            background: 背景/上下文信息，用于二次重排
            top_k: 返回结果数量
        """
        import time
        start_time = time.time()

        # 异步并发召回编制依据
        collection_name = "first_bfp_collection_test"

        gather_start = time.time()
        # 优化：降低hybrid_top_k参数从50到20，减少混合搜索时间
        bfp_tasks = [
            self.async_multi_stage_recall(
                collection_name=collection_name,
                query_text=entity,
                hybrid_top_k=20,  # 从50降到20，减少60%的混合搜索时间
                top_k=top_k
            ) for entity in entity_list
        ]

        bfp_tasks_list = await asyncio.gather(*bfp_tasks,return_exceptions=True)
        gather_end = time.time()
        gather_time = gather_end - gather_start

        bfp_results = []
        for res in bfp_tasks_list:
            if isinstance(res, Exception):
                self.logger.error(f"辅助实体召回失败: {str(res)}")
            else:
                bfp_results.extend(res)

        # BFP召回结果已经通过multi_stage_recall进行了重排序，保持原有顺序
        # 只对第一次重排序得分大于0.8的文档进行二次重排序
        high_score_results = [item for item in bfp_results if item.get('rerank_score', 0) > 0.8]
        low_score_results = [item for item in bfp_results if item.get('rerank_score', 0) <= 0.8]

        self.logger.info(f"筛选结果：高分文档(>0.8) {len(high_score_results)} 个，低分文档(≤0.8) {len(low_score_results)} 个")

        # 如果没有高分文档，直接返回原始结果
        if not high_score_results:
            self.logger.info("没有得分大于0.8的文档，跳过二次重排序，直接返回原始结果")
            return bfp_results

        # 提取高分文档的文本内容用于二次重排
        high_score_text_content = list(set([item['text_content'] for item in high_score_results]))
        self.logger.info(f"提取高分文档文本内容，共 {len(high_score_text_content)} 个，准备二次重排")

        # 二次重排 - 使用配置的重排序模型
        rerank_start = time.time()
        bfp_rerank_result = self._get_rerank_results(background, high_score_text_content, 5)
        rerank_end = time.time()
        self.logger.info(f"二次重排序耗时: {rerank_end - rerank_start:.3f}秒")

        # 根据重排结果重新组织数据
        reorganize_start = time.time()
        final_results = []
        text_to_metadata = {item['text_content']: item for item in high_score_results}

        # 处理二次重排序的高分文档
        for rerank_item in bfp_rerank_result:
            text = rerank_item.get('text', '')
            score = rerank_item.get('score', 0.0)

            if text in text_to_metadata:
                original_item = text_to_metadata[text].copy()
                original_item['bfp_rerank_score'] = score
                final_results.append(original_item)

        reorganize_end = time.time()
        total_time = reorganize_end - start_time

        self.logger.info(f"结果重组耗时: {reorganize_end - reorganize_start:.3f}秒")
        self.logger.info(f"二次重排完成，返回 {len(final_results)} 个高分文档，丢弃 {len(low_score_results)} 个低分文档")
        self.logger.info(f"[async_bfp_recall] 总耗时: {total_time:.3f}秒 (召回: {gather_end-gather_start:.3f}s + 重排: {rerank_end-rerank_start:.3f}s + 其他: {total_time-(gather_end-gather_start)-(rerank_end-rerank_start):.3f}s)")

        return final_results


    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"开始混合检索")

            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)} 个结果")
            # for i, result in enumerate(results):
            #     text_content = result.get('text_content', '')
            #     metadata = result.get('metadata', {})
            #     title = metadata.get('title', 'N/A')
            #     file = metadata.get('file', 'N/A')
            #     self.logger.info(f"混合搜索结果 {i+1}: 标题='{title}', 文件='{file}', 内容长度={len(text_content)}")
            #     # self.logger.info(f"  完整元数据: {metadata}")
            #     # self.logger.info(f"  文本内容: '{text_content}'")

            return results

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

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

        Args:
            candidates_with_metadata: 候选文档列表，包含文本内容和元数据
            query_text: 查询文本
            top_k: 返回结果数量

        Returns:
            List[Dict]: 重排序后的结果列表，包含原始索引信息
        """
        try:
            # 第一步：基于文本内容+元数据的组合去重
            unique_candidates = []
            original_indices_map = []  # 记录每个去重后的候选文档对应的原始索引列表
            unique_combinations = set()  # 记录已见过的文本+元数据组合

            for original_index, candidate in enumerate(candidates_with_metadata):
                text_content = candidate.get('text_content', '')
                metadata = candidate.get('metadata', {})

                # 处理嵌套的metadata字符串
                title = ''
                file = ''
                if 'metadata' in metadata and isinstance(metadata['metadata'], str):
                    import json
                    try:
                        # 解析JSON格式的metadata
                        inner_metadata = json.loads(metadata['metadata'])
                        title = inner_metadata.get('title', '')
                        file = inner_metadata.get('file', '')
                    except (json.JSONDecodeError, TypeError):
                        pass
                else:
                    title = metadata.get('title', '')
                    file = metadata.get('file', '')

                # 创建组合键：文本内容 + 关键元数据
                combination_key = (text_content, title, file)

                if combination_key not in unique_combinations:
                    # 新的唯一组合
                    unique_candidates.append(candidate)
                    original_indices_map.append([original_index])
                    unique_combinations.add(combination_key)
                else:
                    # 找到对应的唯一候选并添加索引
                    for unique_idx, unique_candidate in enumerate(unique_candidates):
                        if unique_candidate.get('text_content', '') == text_content:
                            # 解析唯一候选的元数据
                            unique_metadata = unique_candidate.get('metadata', {})
                            unique_title = ''
                            unique_file = ''
                            if 'metadata' in unique_metadata and isinstance(unique_metadata['metadata'], str):
                                import json
                                try:
                                    inner_metadata = json.loads(unique_metadata['metadata'])
                                    unique_title = inner_metadata.get('title', '')
                                    unique_file = inner_metadata.get('file', '')
                                except (json.JSONDecodeError, TypeError):
                                    pass
                            else:
                                unique_title = unique_metadata.get('title', '')
                                unique_file = unique_metadata.get('file', '')

                            if unique_title == title and unique_file == file:
                                original_indices_map[unique_idx].append(original_index)
                                break

            # 提取唯一候选文档的文本内容用于重排序
            unique_texts = [candidate.get('text_content', '') for candidate in unique_candidates]

            # 使用配置的重排序模型进行重排序
            rerank_results = self._get_rerank_results(query_text, unique_texts, top_k)

            # 转换结果格式，使用索引映射来处理原始索引
            scored_docs = []
            for i, api_result in enumerate(rerank_results):
                rerank_text = api_result.get('text', '')
                rerank_score = float(api_result.get('score', '0.0'))

                # 使用去重时的索引映射
                original_index = original_indices_map[i][0]  # 取第一个原始索引
                original_candidate = unique_candidates[i]  # 获取原始候选文档（包含元数据）

                # 获取原始混合搜索的评分信息
                hybrid_distance = original_candidate.get('distance', 0.0)
                hybrid_similarity = original_candidate.get('similarity', 0.0)

                # 解析元数据获取标题用于日志
                metadata = original_candidate.get('metadata', {})
                title = 'N/A'
                if 'metadata' in metadata and isinstance(metadata['metadata'], str):
                    try:
                        import json
                        inner_metadata = json.loads(metadata['metadata'])
                        title = inner_metadata.get('title', 'N/A')
                    except:
                        pass

                scored_docs.append({
                    'text_content': rerank_text,
                    'metadata': original_candidate.get('metadata', {}),  # 保留原始元数据
                    'rerank_score': rerank_score,
                    'original_index': original_index,
                    'rerank_rank': i,
                    'duplicate_count': len(original_indices_map[i]),  # 记录重复数量
                    'hybrid_distance': hybrid_distance,  # 保留原始混合搜索评分
                    'hybrid_similarity': hybrid_similarity
                })

                # 输出双重评分信息
                # self.logger.info(f"重排序评分 #{i+1}: 标题='{title}' | 混合搜索相似度={hybrid_similarity:.4f} | BGE重排序评分={rerank_score:.6f}")

            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"执行多路召回")

            # 第一阶段：混合搜索召回（向量+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 []

            # 第二阶段：重排序召回，传递完整的混合搜索结果（包含元数据）
            rerank_results = self.rerank_recall(
                candidates_with_metadata=hybrid_results,
                query_text=query_text,
                top_k=top_k
            )

            # 优化重排序结果的元数据结构
            final_results = []
            for rerank_result in rerank_results:
                metadata = rerank_result.get('metadata', {}).copy()
                duplicate_count = rerank_result.get('duplicate_count', 1)

                # 如果内层有metadata字段，将其提取到外层
                if 'metadata' in metadata and isinstance(metadata['metadata'], str):
                    import json
                    try:
                        # 解析JSON格式的metadata
                        inner_metadata = json.loads(metadata['metadata'])
                        metadata.update(inner_metadata)
                        # 移除内层的metadata字符串，避免重复
                        del metadata['metadata']
                    except (json.JSONDecodeError, TypeError):
                        # 如果解析失败，保持原样
                        pass

                # 移除重复的content字段
                if 'content' in metadata:
                    del metadata['content']

                # 添加重复计数信息到元数据中
                if duplicate_count > 1:
                    metadata['duplicate_count'] = duplicate_count

                # 输出优化后的结果，包含双重评分
                final_result = {
                    'text_content': rerank_result['text_content'],
                    'metadata': metadata,
                    'hybrid_similarity': rerank_result.get('hybrid_similarity', 0.0),  # 混合搜索相似度
                    'rerank_score': rerank_result.get('rerank_score', 0.0)  # BGE重排序评分
                }
                final_results.append(final_result)

                self.logger.debug(f"元数据优化完成: 重排序排名{rerank_result.get('rerank_rank')}, 重复数量={duplicate_count}")

            return final_results

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


    async def async_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]: 重排序后的结果列表，只包含重排序分数
        """
        import time
        try:
            start_time = time.time()

            # 第一阶段：混合搜索召回（向量+BM25）
            hybrid_start = time.time()
            hybrid_results = await asyncio.to_thread(
                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:
                return []

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

            # 优化重排序结果的元数据结构
            final_results = []
            for rerank_result in rerank_results:
                metadata = rerank_result.get('metadata', {}).copy()
                duplicate_count = rerank_result.get('duplicate_count', 1)

                # 如果内层有metadata字段，将其提取到外层
                if 'metadata' in metadata and isinstance(metadata['metadata'], str):
                    import json
                    try:
                        # 解析JSON格式的metadata
                        inner_metadata = json.loads(metadata['metadata'])
                        metadata.update(inner_metadata)
                        # 移除内层的metadata字符串，避免重复
                        del metadata['metadata']
                    except (json.JSONDecodeError, TypeError):
                        # 如果解析失败，保持原样
                        pass

                # 移除重复的content字段
                if 'content' in metadata:
                    del metadata['content']

                # 添加重复计数信息到元数据中
                if duplicate_count > 1:
                    metadata['duplicate_count'] = duplicate_count

                # 输出优化后的结果，包含双重评分
                final_result = {
                    'text_content': rerank_result['text_content'],
                    'metadata': metadata,
                    'hybrid_similarity': rerank_result.get('hybrid_similarity', 0.0),  # 混合搜索相似度
                    'rerank_score': rerank_result.get('rerank_score', 0.0)  # BGE重排序评分
                }
                final_results.append(final_result)

                self.logger.debug(f"元数据优化完成: 重排序排名{rerank_result.get('rerank_rank')}, 重复数量={duplicate_count}")

            return final_results

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

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