""" RAG 核心引擎 提供检索增强生成 (Retrieval-Augmented Generation) 的核心功能： - 模型加载：Embedding 模型、Rerank 模型、LLM 客户端 - 多向量库支持：支持按部门/权限隔离的多向量库模式 - 混合检索：向量检索 + BM25 + RRF 融合 + Rerank 重排 - 自适应检索：根据查询置信度动态调整 TopK - 上下文扩展：自动扩展相邻切片，提升上下文连贯性 - MMR 去重：基于最大边际相关性的结果去重 主要组件： - ONNXReranker: 基于 ONNX Runtime 的 Reranker（可选） - RAGEngine: 核心引擎类（单例模式） 使用方式： from core.engine import get_engine engine = get_engine() result = engine.search_knowledge("你的查询", top_k=10) 依赖配置 (config.py): - API_KEY, BASE_URL, MODEL: LLM 配置 - EMBEDDING_MODEL_PATH, RERANK_MODEL_PATH: 模型路径 - USE_MULTI_KB: 是否启用多向量库模式 - USE_HYBRID_SEARCH: 是否启用混合检索 - USE_RERANK: 是否启用 Rerank """ import os import gc import time import logging import numpy as np from typing import List, Tuple, Optional, Dict, Any logger = logging.getLogger(__name__) from sentence_transformers import SentenceTransformer, CrossEncoder import chromadb from openai import OpenAI from parsers import MINERU_AVAILABLE, PANDAS_AVAILABLE # 缓存支持（延迟导入避免循环依赖） try: from core.cache import get_cache_manager, RAGCacheManager CACHE_AVAILABLE = True except ImportError: CACHE_AVAILABLE = False # 延迟导入，防止循环依赖 _engine_instance = None try: from config import ( API_KEY, BASE_URL, MODEL, MODELS_DIR, EMBEDDING_MODEL_PATH, RERANK_MODEL_PATH, CHROMA_DB_PATH, DOCUMENTS_PATH, BM25_INDEXES_PATH, USE_MULTI_KB, USE_HYBRID_SEARCH, VECTOR_WEIGHT, BM25_WEIGHT, USE_RERANK, RERANK_CANDIDATES, RERANK_TOP_K, RERANK_USE_ONNX, RERANK_BACKEND, RERANK_CLOUD_MODEL, RERANK_CLOUD_API_KEY, RERANK_CLOUD_BASE_URL, RERANK_CLOUD_TIMEOUT, CHUNK_SIZE, CHUNK_OVERLAP, # 设备配置 EMBEDDING_DEVICE, RERANK_DEVICE, # 自适应 TopK 配置 ADAPTIVE_TOPK_ENABLED, ADAPTIVE_LOW_CONFIDENCE, ADAPTIVE_HIGH_CONFIDENCE, ADAPTIVE_EXPAND_RATIO, ADAPTIVE_SHRINK_RATIO, ADAPTIVE_MIN_TOPK, ADAPTIVE_MAX_TOPK, # 切片配置 MIN_CHUNK_SIZE, MAX_CHUNK_SIZE, SECTION_FILTER_ENABLED, # P3 优化配置 MMR_ENABLED, MMR_TOP_K, DYNAMIC_RRF_ENABLED, CONTEXT_EXPANSION_ENABLED, CONTEXT_EXPANSION_BEFORE, CONTEXT_EXPANSION_AFTER, CONTEXT_EXPANSION_MAX_CHUNKS, ENUM_QUERY_DISABLE_TOPK_SHRINK, ENUM_QUERY_MMR_LAMBDA, # Phase 3 扩展精细化 EXPANSION_SCORE_THRESHOLD, MAX_EXPANDED_NEIGHBORS, # 上下文与生成 LLM_TEMPERATURE, LLM_MAX_TOKENS, RECALL_MULTIPLIER, # FAQ 与黑名单 FAQ_RECALL_TOP_K, FAQ_BOOST_AMOUNT, FAQ_DECAY_MONTHS, FAQ_DECAY_RATE, FAQ_DECAY_MAX, BLACKLIST_MIN_DISLIKES, BLACKLIST_CACHE_TTL, # 缓存与融合 CACHE_MIN_SCORE, RRF_K ) except ImportError: # 默认值 ADAPTIVE_TOPK_ENABLED = True ADAPTIVE_LOW_CONFIDENCE = 0.5 ADAPTIVE_HIGH_CONFIDENCE = 0.8 ADAPTIVE_EXPAND_RATIO = 2.0 ADAPTIVE_SHRINK_RATIO = 0.5 ADAPTIVE_MIN_TOPK = 3 ADAPTIVE_MAX_TOPK = 20 MIN_CHUNK_SIZE = 200 MAX_CHUNK_SIZE = 1200 SECTION_FILTER_ENABLED = True # P3 优化默认值 MMR_ENABLED = True MMR_TOP_K = 30 CONTEXT_EXPANSION_ENABLED = True CONTEXT_EXPANSION_BEFORE = 1 CONTEXT_EXPANSION_AFTER = 5 CONTEXT_EXPANSION_MAX_CHUNKS = 24 EXPANSION_SCORE_THRESHOLD = 0.3 MAX_EXPANDED_NEIGHBORS = 4 ENUM_QUERY_DISABLE_TOPK_SHRINK = True ENUM_QUERY_MMR_LAMBDA = 0.85 DYNAMIC_RRF_ENABLED = True EMBEDDING_DEVICE = "auto" RERANK_DEVICE = "auto" RERANK_USE_ONNX = False RERANK_BACKEND = "local" RERANK_CLOUD_MODEL = "qwen3-rerank" RERANK_CLOUD_API_KEY = "" RERANK_CLOUD_BASE_URL = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks" RERANK_CLOUD_TIMEOUT = 15 def _get_device(device_config: str) -> str: """ 解析设备配置 Args: device_config: 设备配置 ("auto", "cuda", "cpu", "cuda:0" 等) Returns: 实际设备字符串 """ if device_config == "auto": try: import torch if torch.cuda.is_available(): device = "cuda" logger.info(f"检测到GPU: {torch.cuda.get_device_name(0)}") return device else: logger.info("未检测到GPU，使用CPU") return "cpu" except ImportError: logger.info("PyTorch未安装，使用CPU") return "cpu" return device_config from core.bm25_index import BM25Index from core.constants import get_empty_result class ONNXReranker: """ 基于 ONNX Runtime 的 Reranker 使用 ONNX Runtime 进行模型推理，在 CPU 上比原生 PyTorch 快 2-3 倍。 自动处理模型转换：首次使用时将 PyTorch 模型导出为 ONNX 格式。 Attributes: tokenizer: HuggingFace tokenizer model: ONNX Runtime 模型 Example: >>> reranker = ONNXReranker("models/bge-reranker-base") >>> scores = reranker.predict([("query", "doc1"), ("query", "doc2")]) """ def __init__(self, model_path: str) -> None: """ 初始化 ONNX Reranker Args: model_path: 模型目录路径，需包含 PyTorch 模型文件 """ from optimum.onnxruntime import ORTModelForSequenceClassification from transformers import AutoTokenizer onnx_path = os.path.join(model_path, "onnx") # 首次使用时导出 ONNX 模型 if not os.path.exists(os.path.join(onnx_path, "model.onnx")): os.makedirs(onnx_path, exist_ok=True) logger.info("正在导出 Rerank 模型为 ONNX 格式...") model = ORTModelForSequenceClassification.from_pretrained( model_path, export=True ) model.save_pretrained(onnx_path) logger.info(f"ONNX 模型已导出: {onnx_path}") else: model = ORTModelForSequenceClassification.from_pretrained(onnx_path) self.tokenizer = AutoTokenizer.from_pretrained(model_path) self.model = model def predict(self, pairs: List[Tuple[str, str]]) -> np.ndarray: """ 计算查询-文档对的相关性分数 兼容 CrossEncoder.predict 接口。 Args: pairs: (query, document) 元组列表 Returns: 相关性分数数组，范围通常为 [0, 1] """ import torch texts = [p[0] for p in pairs] texts_pair = [p[1] for p in pairs] inputs = self.tokenizer( texts, texts_pair, padding=True, truncation=True, max_length=512, return_tensors="pt" ) with torch.no_grad(): outputs = self.model(**inputs) scores = outputs.logits.squeeze(-1) if scores.dim() == 0: scores = scores.unsqueeze(0) return scores.numpy() class CloudReranker: """ 基于 DashScope 云端 Reranker API 的包装类 通过 OpenAI 兼容接口 (/compatible-api/v1/reranks) 调用云端重排序模型， 接口与本地 CrossEncoder/ONNXReranker 的 predict() 保持一致。 优势：不占用本地 CPU/GPU 资源，适合服务器 CPU 推理较慢的场景。 Attributes: model: 云端模型名称（如 qwen3-rerank） api_key: DashScope API Key base_url: API 端点 URL timeout: 请求超时秒数 """ def __init__(self, model: str, api_key: str, base_url: str, timeout: int = 15) -> None: self.model = model self.api_key = api_key self.base_url = base_url self.timeout = timeout self._session = None def _get_session(self): """延迟创建 requests.Session，复用连接池""" if self._session is None: import requests self._session = requests.Session() self._session.headers.update({ "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json" }) return self._session def predict(self, pairs: list) -> np.ndarray: """ 计算查询-文档对的相关性分数 兼容 CrossEncoder.predict / ONNXReranker.predict 接口。 Args: pairs: (query, document) 元组列表 Returns: 相关性分数数组 (np.ndarray) """ if not pairs: return np.array([]) # 从 pairs 中提取 query（取第一个）和 documents query = pairs[0][0] documents = [p[1] for p in pairs] body = { "model": self.model, "query": query, "documents": documents, "top_n": len(documents) } session = self._get_session() resp = session.post(self.base_url, json=body, timeout=self.timeout) if resp.status_code != 200: error_msg = resp.text[:200] logger.error(f"Cloud Reranker API 调用失败: status={resp.status_code}, {error_msg}") # 返回零分，让 pipeline 继续运行（降级处理） return np.zeros(len(pairs)) data = resp.json() results = data.get("results", []) # API 返回的 results 可能不保持原始顺序，需按 index 还原 scores = [0.0] * len(documents) for item in results: idx = item.get("index", 0) score = item.get("relevance_score", 0.0) if 0 <= idx < len(scores): scores[idx] = score return np.array(scores) class RAGEngine: """ RAG 核心引擎 单例模式管理所有共享资源，提供检索增强生成的核心功能。 主要功能： - 模型管理：Embedding 模型、Rerank 模型、LLM 客户端 - 多向量库支持：支持按部门/权限隔离的多向量库模式 - 混合检索：向量检索 + BM25 + RRF 融合 + Rerank 重排 - 自适应检索：根据查询置信度动态调整 TopK - 上下文扩展：自动扩展相邻切片，提升上下文连贯性 Attributes: embedding_model: SentenceTransformer 向量模型 reranker: CrossEncoder 或 ONNXReranker 重排模型 llm_client: OpenAI 客户端 bm25_index: BM25 关键词索引 kb_manager: 多向量库管理器（多向量库模式） kb_router: 知识库路由器（多向量库模式） collection: ChromaDB 集合（单向量库模式） 使用方式： >>> from core.engine import get_engine >>> engine = get_engine() >>> result = engine.search_knowledge("查询", top_k=10) """ @classmethod def get_instance(cls) -> 'RAGEngine': """获取引擎单例实例""" global _engine_instance if _engine_instance is None: _engine_instance = cls() return _engine_instance def __init__(self) -> None: """初始化引擎（延迟加载，需调用 initialize() 完成初始化）""" self.embedding_model = None self.reranker = None self.llm_client = None self.bm25_index = None # 单向量库模式 self.chroma_client = None self.collection = None # 黑名单缓存（避免每次查询都重建 FeedbackService） self._blacklist_cache = None self._blacklist_cache_time = 0 self._blacklist_cache_ttl = BLACKLIST_CACHE_TTL # 多向量库模式 self.kb_manager = None self.kb_router = None # 自适应 TopK 策略 self._adaptive_topk = None self._initialized = False def initialize(self) -> None: """ 显式初始化所有模型和数据库客户端 加载顺序： 1. 向量模型 (SentenceTransformer) 2. 向量数据库 (ChromaDB) 3. BM25 索引 4. LLM 客户端 (OpenAI) 5. Reranker 模型 6. 自适应 TopK 策略 """ global USE_RERANK if self._initialized: return logger.info("初始化 RAG Engine 核心") os.makedirs(MODELS_DIR, exist_ok=True) # 1. 向量模型 if os.path.exists(EMBEDDING_MODEL_PATH): device = _get_device(EMBEDDING_DEVICE) self.embedding_model = SentenceTransformer(EMBEDDING_MODEL_PATH, device=device) logger.info(f"向量模型加载完成: {EMBEDDING_MODEL_PATH} (设备: {device})") else: raise RuntimeError(f"向量模型未找到: {EMBEDDING_MODEL_PATH}") # 2. 向量数据库 if USE_MULTI_KB: from knowledge.manager import KnowledgeBaseManager from knowledge.router import KnowledgeBaseRouter self.kb_manager = KnowledgeBaseManager(CHROMA_DB_PATH) self.kb_router = KnowledgeBaseRouter(use_llm=False) self.collection = self.kb_manager.get_collection('public_kb') logger.info("多向量库模式已启用") else: self.chroma_client = chromadb.PersistentClient(path=CHROMA_DB_PATH) self.collection = self.chroma_client.get_or_create_collection( name="knowledge_base", metadata={"description": "RAG Demo 知识库"} ) logger.info(f"单向量库模式已启用: {CHROMA_DB_PATH}") # 3. BM25 索引 self.bm25_index = BM25Index() if USE_HYBRID_SEARCH and not USE_MULTI_KB: bm25_path = os.path.join(BM25_INDEXES_PATH, "default_bm25.pkl") self.bm25_index.load(bm25_path) logger.info("BM25索引已加载") # 4. LLM 客户端 self.llm_client = OpenAI(api_key=API_KEY, base_url=BASE_URL) logger.info(f"LLM 客户端就绪: {MODEL}") # 5. Reranker 模型（支持 local / cloud / fallback 三种后端） if USE_RERANK: self.reranker = None backend = RERANK_BACKEND # 尝试加载云端 Reranker if backend in ("cloud", "fallback"): if RERANK_CLOUD_API_KEY: try: self.reranker = CloudReranker( model=RERANK_CLOUD_MODEL, api_key=RERANK_CLOUD_API_KEY, base_url=RERANK_CLOUD_BASE_URL, timeout=RERANK_CLOUD_TIMEOUT, ) logger.info(f"Rerank 云端模型就绪: {RERANK_CLOUD_MODEL} ({RERANK_CLOUD_BASE_URL})") except Exception as e: logger.warning(f"云端 Reranker 初始化失败: {e}") if backend == "cloud": USE_RERANK = False # fallback 模式会继续尝试本地模型 else: logger.warning("RERANK_CLOUD_API_KEY 未配置，云端 Reranker 不可用") if backend == "cloud": USE_RERANK = False # 尝试加载本地 Reranker（local 模式，或 fallback 云端失败时） if self.reranker is None and backend in ("local", "fallback"): if os.path.exists(RERANK_MODEL_PATH): try: if RERANK_USE_ONNX: self.reranker = ONNXReranker(RERANK_MODEL_PATH) logger.info(f"Rerank模型加载完成(ONNX加速): {RERANK_MODEL_PATH}") else: device = _get_device(RERANK_DEVICE) self.reranker = CrossEncoder(RERANK_MODEL_PATH, device=device) logger.info(f"Rerank模型加载完成: {RERANK_MODEL_PATH} (设备: {device})") except Exception as e: logger.warning(f"ONNX加载失败，回退到CrossEncoder: {e}") device = _get_device(RERANK_DEVICE) self.reranker = CrossEncoder(RERANK_MODEL_PATH, device=device) logger.info(f"Rerank模型加载完成(回退): {RERANK_MODEL_PATH} (设备: {device})") else: try: os.makedirs(RERANK_MODEL_PATH, exist_ok=True) from transformers import AutoModelForSequenceClassification, AutoTokenizer model = AutoModelForSequenceClassification.from_pretrained("BAAI/bge-reranker-base") tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-reranker-base") model.save_pretrained(RERANK_MODEL_PATH) tokenizer.save_pretrained(RERANK_MODEL_PATH) except Exception as e: logger.warning(f"Rerank模型加载失败: {e}") USE_RERANK = False # 6. 自适应 TopK 策略 if ADAPTIVE_TOPK_ENABLED: from core.adaptive_topk import AdaptiveConfig, AdaptiveTopK config = AdaptiveConfig( enabled=ADAPTIVE_TOPK_ENABLED, low_confidence_threshold=ADAPTIVE_LOW_CONFIDENCE, high_confidence_threshold=ADAPTIVE_HIGH_CONFIDENCE, expand_ratio=ADAPTIVE_EXPAND_RATIO, shrink_ratio=ADAPTIVE_SHRINK_RATIO, min_top_k=ADAPTIVE_MIN_TOPK, max_top_k=ADAPTIVE_MAX_TOPK ) self._adaptive_topk = AdaptiveTopK(config) logger.info(f"自适应TopK已启用 (低={ADAPTIVE_LOW_CONFIDENCE}, 高={ADAPTIVE_HIGH_CONFIDENCE})") self._initialized = True # ---------------- 核心检索逻辑 ---------------- def search_knowledge(self, query, top_k=5, allowed_levels=None, role=None, department=None, collections=None, source_filter=None, _skip_decomposition=False, sub_queries=None): """ 混合检索逻辑封装 Args: query: 查询文本 top_k: 返回结果数量 allowed_levels: 兏许访问的安全级别列表 role: 用户角色（多向量库模式） department: 用户部门（多向量库模式） collections: 指定查询的向量库列表 source_filter: 文件名过滤，精确匹配（如 "2604.09205v1.pdf"） _skip_decomposition: 内部标志，跳过查询拆分（避免递归） sub_queries: 外部传入的子查询列表（由意图分析器生成），如果提供则并行检索 """ if not self._initialized: self.initialize() # ==================== 调试信息收集 + 分阶段计时 ==================== _debug = {'steps': [], 'timing': {}} _overall_start = time.time() # ==================== 查询缓存检查 ==================== cache_enabled = False if CACHE_AVAILABLE: try: from config import QUERY_CACHE_ENABLED cache_enabled = QUERY_CACHE_ENABLED except ImportError: cache_enabled = True # 默认启用 if cache_enabled: cache = get_cache_manager() # 确定缓存键的知识库名称 kb_name = "public_kb" if USE_MULTI_KB and collections: kb_name = collections[0] if len(collections) == 1 else "multi" cached_result = cache.get_query_result(query, kb_name) if cached_result is not None: # 缓存命中，直接返回 _debug['steps'].append({'name': 'cache_hit', 'count': len(cached_result.get('ids', [[]])[0])}) cached_result['_debug'] = _debug return cached_result # ==================== 外部子查询并行检索 ==================== if sub_queries and len(sub_queries) > 1: _debug['steps'].append({'name': 'sub_query_search', 'sub_queries': sub_queries}) result = self._search_with_sub_queries( query, sub_queries, top_k, allowed_levels, role, department, collections, source_filter ) # 缓存结果 if cache_enabled and result.get('ids') and result['ids'][0]: top_dist = result['distances'][0][0] if result.get('distances') and result['distances'][0] else 1.0 top_score = 1.0 - top_dist if top_score >= CACHE_MIN_SCORE: doc_ids = result.get('ids', [[]])[0] if result.get('ids') else [] cache.set_query_result(query, kb_name, result, doc_ids=doc_ids) result['_debug'] = _debug return result # ==================== 查询拆分（复杂查询分解）==================== if not _skip_decomposition: try: from core.query_decomposer import QueryDecomposer decomposer = QueryDecomposer() needs_decompose, decompose_type = decomposer.should_decompose(query) if needs_decompose: _debug['steps'].append({'name': 'query_decomposition', 'type': decompose_type}) result = self._search_with_decomposition( query, decomposer, top_k, allowed_levels, role, department, collections, source_filter ) # 缓存结果 if cache_enabled and result.get('ids') and result['ids'][0]: top_dist = result['distances'][0][0] if result.get('distances') and result['distances'][0] else 1.0 top_score = 1.0 - top_dist if top_score >= CACHE_MIN_SCORE: doc_ids = result.get('ids', [[]])[0] if result.get('ids') else [] cache.set_query_result(query, kb_name, result, doc_ids=doc_ids) result['_debug'] = _debug return result except Exception as e: logger.debug(f"查询拆分检查失败: {e}，使用原始查询") # ==================== Query Expansion（P3-3）==================== # 查询扩展（可选） try: from config import QUERY_EXPANSION_ENABLED, QUERY_EXPANSION_THRESHOLD except ImportError: QUERY_EXPANSION_ENABLED = False QUERY_EXPANSION_THRESHOLD = 0.8 expanded_queries = [query] if QUERY_EXPANSION_ENABLED: try: from core.query_expansion import expand_query_safe expanded_queries = expand_query_safe( query, embedding_model=self.embedding_model, threshold=QUERY_EXPANSION_THRESHOLD, max_expansions=3 ) except Exception as e: pass # 扩展失败时使用原查询 if USE_MULTI_KB and self.kb_manager: _debug['steps'].append({'name': 'multi_kb_search', 'collections': collections}) result = self._search_multi_kb(query, top_k, role, department, collections, source_filter=source_filter, _debug=_debug) # 缓存多知识库结果 if cache_enabled and result.get('ids') and result['ids'][0]: top_dist = result['distances'][0][0] if result.get('distances') and result['distances'][0] else 1.0 top_score = 1.0 - top_dist if top_score >= CACHE_MIN_SCORE: # 传递 doc_ids 实现细粒度缓存失效 doc_ids = result.get('ids', [[]])[0] if result.get('ids') else [] cache.set_query_result(query, kb_name, result, doc_ids=doc_ids) _debug['timing']['total_ms'] = int((time.time() - _overall_start) * 1000) result['_debug'] = _debug return result # 构建 where 条件（支持多个过滤条件组合） conditions = [] if allowed_levels: conditions.append({"security_level": {"$in": allowed_levels}}) if source_filter: conditions.append({"source": source_filter}) where_filter = None if len(conditions) == 1: where_filter = conditions[0] elif len(conditions) > 1: where_filter = {"$and": conditions} query_vector = self.embedding_model.encode(query).tolist() recall_k = RERANK_CANDIDATES if (USE_RERANK or USE_HYBRID_SEARCH) else top_k recall_k = max(recall_k, top_k * RECALL_MULTIPLIER) # ========== P0：图片独立召回通道 ========== # 图片切片独立检索，保证有足够的召回机会 image_recall_k = max(5, top_k // 2) # 图片召回数量独立控制 query_kwargs = { "query_embeddings": [query_vector], "n_results": recall_k } if where_filter: query_kwargs["where"] = where_filter vector_results = self.collection.query(**query_kwargs) _debug['steps'].append({'name': 'vector_search', 'count': len(vector_results['ids'][0]) if vector_results.get('ids') else 0}) # 独立检索图片切片（新增） image_results = self._search_image_chunks(query_vector, image_recall_k, where_filter) if image_results and image_results.get('ids') and image_results['ids'][0]: # 合并图片结果到主结果 vector_results = self._merge_results([vector_results, image_results]) # ========== 独立查询 FAQ 集合 ========== faq_results = self._search_faq_collection(query_vector, top_k=FAQ_RECALL_TOP_K) if faq_results and faq_results.get('ids') and faq_results['ids'][0]: # 合并 FAQ 结果到主结果 vector_results = self._merge_results([vector_results, faq_results]) results_list = [vector_results] weights = [VECTOR_WEIGHT] if USE_HYBRID_SEARCH and self.bm25_index.bm25: bm25_results = self.bm25_index.search(query, top_k=recall_k) _debug['steps'].append({'name': 'bm25_search', 'count': len(bm25_results['ids'][0]) if bm25_results.get('ids') else 0}) if where_filter and bm25_results['metadatas'][0]: allowed_set = set(allowed_levels) # 过滤 BM25 结果 bm25_results = self._filter_results(bm25_results, lambda meta: meta.get('security_level', 'public') in allowed_set) results_list.append(bm25_results) # ========== 动态 RRF 权重（P3-1）========== vector_w, bm25_w = self._get_dynamic_rrf_weights(query) weights = [vector_w, bm25_w] if len(results_list) > 1: fused_results = self.reciprocal_rank_fusion(results_list, weights) _debug['steps'].append({'name': 'rrf_fusion', 'count': len(fused_results['ids'][0]) if fused_results.get('ids') else 0, 'weights': [round(w, 2) for w in weights]}) else: fused_results = results_list[0] # 过滤废止切片 before_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 fused_results = self._filter_deprecated_chunks(fused_results) after_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 if before_count != after_count: _debug['steps'].append({'name': 'deprecated_filter', 'removed': before_count - after_count}) is_enum_query = self._is_enumeration_query(query) fused_results['_enum_query'] = is_enum_query # 章节过滤（如果查询中提到了章节） fused_results = self._filter_by_section(fused_results, query) # 补齐强命中切片周围的连续文本（MMR 前扩展，防止邻居被 MMR 当作冗余去除） before_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 fused_results = self._expand_contiguous_chunks(fused_results, top_k=top_k) after_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 if _debug is not None and after_count != before_count: _debug['steps'].append({'name': 'context_expansion', 'before': before_count, 'after': after_count}) # ========== MMR 去重（P3-2：前置到 rerank 前）========== if MMR_ENABLED: before_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 fused_results = self._apply_mmr(query, fused_results, top_k=MMR_TOP_K, is_enum_query=is_enum_query) after_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 _debug['steps'].append({'name': 'mmr_dedup', 'before': before_count, 'after': after_count}) if USE_RERANK and self.reranker: fused_results = self.rerank_results(query, fused_results, top_k) _debug['steps'].append({ 'name': 'rerank', 'applied': True, 'count': len(fused_results['ids'][0]) if fused_results.get('ids') else 0, 'time_ms': fused_results.get('_rerank_time_ms', 0), 'cached': fused_results.get('_rerank_cached', False) }) else: final_top_k = max(top_k, CONTEXT_EXPANSION_MAX_CHUNKS) if is_enum_query else top_k fused_results = self._truncate_results(fused_results, final_top_k) _debug['steps'].append({'name': 'rerank', 'applied': False}) # FAQ 分数加权（Score Boosting） fused_results = self._boost_faq_chunks(fused_results) # ========== 黑名单过滤（负反馈降权，带缓存）========== blacklist = self._get_blacklist() if blacklist: before_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 fused_results = self.filter_blacklisted_chunks(fused_results, blacklist) after_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 if before_count != after_count: _debug['steps'].append({'name': 'blacklist_filter', 'removed': before_count - after_count}) # 时间衰减（Time Decay） fused_results = self._apply_time_decay(fused_results) # ========== 上下文扩展：补充强命中切片周围的连续文本（rerank 之后，防止被截断）========== # Phase 3：仅对高分种子扩展邻居 before_exp = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 fused_results = self._expand_contiguous_chunks(fused_results, top_k=top_k, min_score=EXPANSION_SCORE_THRESHOLD) after_exp = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 _debug['steps'].append({'name': 'context_expansion', 'before': before_exp, 'after': after_exp}) # 自适应 TopK：根据置信度调整返回数量 if ( self._adaptive_topk and fused_results.get('distances') and fused_results['distances'][0] and not (is_enum_query and ENUM_QUERY_DISABLE_TOPK_SHRINK) and fused_results.get('_score_source') != 'rrf' ): top_score = 1.0 - fused_results['distances'][0][0] # 距离转相似度 adjusted_k, should_retrieve, reason = self._adaptive_topk.adjust(top_score, top_k) if "high_confidence" in reason: # 高置信度时截断结果 fused_results = self._truncate_results(fused_results, adjusted_k) _debug['steps'].append({'name': 'adaptive_topk', 'adjusted_k': adjusted_k, 'reason': reason}) elif _debug is not None and is_enum_query and ENUM_QUERY_DISABLE_TOPK_SHRINK: _debug['steps'].append({'name': 'adaptive_topk', 'skipped': True, 'reason': 'enum_query_preserve_context'}) # ==================== 缓存结果 ==================== if cache_enabled and fused_results.get('ids') and fused_results['ids'][0]: # 只缓存有结果且置信度较高的查询 top_dist = fused_results['distances'][0][0] if fused_results.get('distances') and fused_results['distances'][0] else 1.0 top_score = 1.0 - top_dist # 距离转相似度 if top_score >= CACHE_MIN_SCORE: # 置信度阈值 # 传递 doc_ids 实现细粒度缓存失效 doc_ids = fused_results.get('ids', [[]])[0] if fused_results.get('ids') else [] cache.set_query_result(query, kb_name, fused_results, doc_ids=doc_ids) fused_results['_debug'] = _debug _debug['timing']['total_ms'] = int((time.time() - _overall_start) * 1000) return fused_results def _search_faq_collection(self, query_vector: list, top_k: int = None) -> dict: """ 独立查询 FAQ 集合 FAQ 存储在独立的集合中，与普通文档分离，便于独立管理和清理 Args: query_vector: 查询向量 top_k: 返回结果数量 Returns: FAQ 检索结果 """ if top_k is None: top_k = FAQ_RECALL_TOP_K try: # 获取或创建 FAQ 集合 if self.kb_manager: faq_collection = self.kb_manager.get_collection('faq_kb') else: faq_collection = self.chroma_client.get_or_create_collection( name="faq_collection", metadata={"description": "FAQ 专属向量库"} ) if not faq_collection or faq_collection.count() == 0: return get_empty_result() # 查询 FAQ 集合 results = faq_collection.query( query_embeddings=[query_vector], n_results=top_k ) return results except Exception as e: logger.warning(f"FAQ 集合查询失败: {e}") return get_empty_result() def _search_image_chunks(self, query_vector: list, top_k: int = 5, where_filter: dict = None) -> dict: """ 独立检索图片切片（P0：图片独立召回通道） 图片切片使用独立的召回通道，保证有足够的召回机会， 不被文本切片"淹没"在 MMR/Rerank 阶段。 Args: query_vector: 查询向量 top_k: 返回的图片数量 where_filter: 额外的 where 过滤条件 Returns: 图片切片检索结果 """ try: # 构建 where 条件：只检索图片类型 image_filter = {"chunk_type": {"$in": ["image", "chart", "table"]}} # 合并额外的过滤条件 if where_filter: image_filter = {"$and": [where_filter, image_filter]} # 检索图片切片 results = self.collection.query( query_embeddings=[query_vector], n_results=top_k, where=image_filter ) return results except Exception as e: logger.warning(f"图片切片检索失败: {e}") return get_empty_result() def _search_image_chunks_multi_kb(self, query_vector: list, top_k: int, collection, source_filter: str = None) -> dict: """ 多知识库模式下独立检索图片切片（P0） Args: query_vector: 查询向量 top_k: 返回的图片数量 collection: 向量库集合 source_filter: 文件名过滤 Returns: 图片切片检索结果 """ try: # 构建 where 条件：只检索图片类型 image_filter = {"chunk_type": {"$in": ["image", "chart", "table"]}} # 合并文件来源过滤 if source_filter: image_filter = {"$and": [{"source": source_filter}, image_filter]} # 检索图片切片 results = collection.query( query_embeddings=[query_vector], n_results=top_k, where=image_filter ) return results except Exception as e: logger.warning(f"多知识库图片切片检索失败: {e}") return get_empty_result() def _merge_results(self, results_list: list) -> dict: """ 合并多个检索结果 Args: results_list: 结果列表 Returns: 合并后的结果 """ all_ids = [] all_docs = [] all_metas = [] all_dists = [] for results in results_list: if results and results.get('ids') and results['ids'][0]: all_ids.extend(results['ids'][0]) all_docs.extend(results['documents'][0]) all_metas.extend(results['metadatas'][0]) all_dists.extend(results['distances'][0]) return { 'ids': [all_ids], 'documents': [all_docs], 'metadatas': [all_metas], 'distances': [all_dists] } def _boost_faq_chunks(self, results: dict) -> dict: """ FAQ Chunk 分数加权 FAQ 命中时分数提升 0.1，确保 FAQ 排名靠前 注意：distances 是距离，越小越好，所以减去 0.1 """ if not results.get('metadatas') or not results['metadatas'][0]: return results for i, meta in enumerate(results['metadatas'][0]): if meta.get('chunk_type') == 'faq': # FAQ 加权：距离减小 = 相似度提升 results['distances'][0][i] = max(0, results['distances'][0][i] - FAQ_BOOST_AMOUNT) return results def _apply_time_decay(self, results: dict, decay_months: int = None) -> dict: """ 时间衰减：超过 N 个月的 FAQ 扣分 防止过期 FAQ 成为"钉子户"，确保新内容有机会排在前面 """ from datetime import datetime if decay_months is None: decay_months = FAQ_DECAY_MONTHS if not results.get('metadatas') or not results['metadatas'][0]: return results now = datetime.now() for i, meta in enumerate(results['metadatas'][0]): if meta.get('chunk_type') == 'faq': created_at = meta.get('created_at') if created_at: try: created = datetime.fromisoformat(created_at) age_months = (now - created).days / 30 if age_months > decay_months: # 每超过一个月扣 0.01，最多扣 0.1 decay = min(FAQ_DECAY_MAX, (age_months - decay_months) * FAQ_DECAY_RATE) # 距离增加 = 相似度降低 results['distances'][0][i] += decay except (ValueError, TypeError): pass return results def _get_blacklist(self, min_dislikes: int = None) -> set: """获取黑名单（带缓存，5分钟刷新一次）""" if min_dislikes is None: min_dislikes = BLACKLIST_MIN_DISLIKES import time as _time now = _time.time() if self._blacklist_cache is not None and (now - self._blacklist_cache_time) < self._blacklist_cache_ttl: return self._blacklist_cache try: from services.feedback import FeedbackDB, FeedbackService feedback_service = FeedbackService(FeedbackDB()) self._blacklist_cache = feedback_service.get_chunk_blacklist(min_dislikes=min_dislikes) self._blacklist_cache_time = now except Exception as e: logger.warning(f"获取黑名单缓存失败: {e}") self._blacklist_cache = set() return self._blacklist_cache def filter_blacklisted_chunks(self, results: dict, blacklist: set) -> dict: """ 过滤黑名单 Chunk（负反馈降权机制） Args: results: 检索结果 blacklist: 黑名单 source 集合 Returns: 过滤后的结果 """ if not blacklist or not results.get('metadatas') or not results['metadatas'][0]: return results f_ids, f_docs, f_metas, f_scores = [], [], [], [] filtered_count = 0 for bid, bdoc, bmeta, bscore in zip( results['ids'][0], results['documents'][0], results['metadatas'][0], results['distances'][0] ): source = bmeta.get('source', '') if source not in blacklist: f_ids.append(bid) f_docs.append(bdoc) f_metas.append(bmeta) f_scores.append(bscore) else: filtered_count += 1 if filtered_count > 0: logger.debug(f"过滤了 {filtered_count} 个黑名单 Chunk") filtered = { 'ids': [f_ids], 'documents': [f_docs], 'metadatas': [f_metas], 'distances': [f_scores] } for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'): if key in results: filtered[key] = results[key] return filtered def _filter_results(self, results, condition_func): """通用结果过滤辅助函数""" f_ids, f_docs, f_metas, f_scores = [], [], [], [] for bid, bdoc, bmeta, bscore in zip( results['ids'][0], results['documents'][0], results['metadatas'][0], results['distances'][0] ): if condition_func(bmeta): f_ids.append(bid) f_docs.append(bdoc) f_metas.append(bmeta) f_scores.append(bscore) filtered = { 'ids': [f_ids], 'documents': [f_docs], 'metadatas': [f_metas], 'distances': [f_scores] } for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'): if key in results: filtered[key] = results[key] return filtered def _truncate_results(self, results, top_k): """截断结果到指定的 top_k""" if not results.get('ids') or not results['ids'][0]: return results truncated = { 'ids': [results['ids'][0][:top_k]], 'documents': [results['documents'][0][:top_k]], 'metadatas': [results['metadatas'][0][:top_k]], 'distances': [results['distances'][0][:top_k]] } for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'): if key in results: truncated[key] = results[key] return truncated def _is_enumeration_query(self, query: str) -> bool: """Detect list/clause queries where contiguous chunks are more important than diversity.""" try: from core.query_classifier import is_enumeration_query return is_enumeration_query(query) except Exception as e: logger.debug(f"枚举查询检测失败: {e}") if not query: return False markers = ("哪些", "有哪些", "列出", "严禁", "禁止", "不得", "包括", "要求", "情形", "场景") return any(marker in query for marker in markers) def _to_int(self, value, default=None): try: if value is None or value == "": return default return int(value) except (TypeError, ValueError): return default def _collection_for_meta(self, meta: dict): coll_name = meta.get('_collection') or meta.get('collection') if self.kb_manager and coll_name: try: return self.kb_manager.get_collection(coll_name) except Exception as e: logger.debug(f"获取collection失败: {e}") return None return self.collection def _expand_contiguous_chunks(self, results: dict, top_k: int = None, min_score: float = 0.0) -> dict: """Add same-source same-section neighbor text chunks around strong hits. Args: results: 检索结果 top_k: 最大切片数 min_score: Phase 3 最低分数阈值，仅对 Rerank 分数高于此值的种子扩展 """ if not CONTEXT_EXPANSION_ENABLED: return results if not results.get('ids') or not results['ids'][0]: return results max_chunks = max((top_k or 0) * 2, CONTEXT_EXPANSION_MAX_CHUNKS) base_limit = min(len(results['ids'][0]), max(top_k or 0, 10)) existing_ids = set(results['ids'][0]) items = [] for doc_id, doc, meta, dist in zip( results['ids'][0][:base_limit], results['documents'][0][:base_limit], results['metadatas'][0][:base_limit], (results['distances'][0] if results.get('distances') else [0] * len(results['ids'][0]))[:base_limit] ): items.append((doc_id, doc, meta, dist)) existing_ids = {item[0] for item in items} # 兼容 ID 前缀：同时存储原始 ID 和去前缀版本，确保邻居匹配不遗漏 _raw_id_set = set() for _eid in existing_ids: if '/' in _eid: _raw_id_set.add(_eid.split('/', 1)[1]) else: _raw_id_set.add(_eid) seeds = [ (doc_id, doc, meta, dist) for doc_id, doc, meta, dist in items[:base_limit] if meta.get('chunk_type', 'text') == 'text' and meta.get('source') and self._to_int(meta.get('chunk_index')) is not None ] added = 0 for seed_id, _seed_doc, seed_meta, seed_dist in seeds: if len(items) >= max_chunks: break # Phase 3：跳过分数低于阈值的种子（仅当 min_score > 0 时生效） if min_score > 0 and seed_dist < min_score: continue source = seed_meta.get('source') section = seed_meta.get('section', '') or seed_meta.get('section_path', '') seed_index = self._to_int(seed_meta.get('chunk_index')) if source is None or seed_index is None: continue collection = self._collection_for_meta(seed_meta) if not collection: continue def _get_neighbors(where_filter): try: return collection.get( where=where_filter, include=['documents', 'metadatas'] ) except Exception as e: logger.warning(f"扩展连续切片失败: {e}") return {'ids': [], 'documents': [], 'metadatas': []} where_filter = {"$and": [{"source": source}, {"chunk_type": "text"}]} if section: where_filter["$and"].append({"section": section}) neighbors = _get_neighbors(where_filter) # Continuation chunks often lose section metadata after splitting. # Fall back to source-only and rely on chunk_index window to keep it tight. if not neighbors.get('ids') or len(neighbors.get('ids', [])) <= 1: neighbors = _get_neighbors({"$and": [{"source": source}, {"chunk_type": "text"}]}) neighbor_rows = [] for n_id, n_doc, n_meta in zip( neighbors.get('ids', []), neighbors.get('documents', []), neighbors.get('metadatas', []) ): n_index = self._to_int(n_meta.get('chunk_index')) if n_index is None: continue if seed_index - CONTEXT_EXPANSION_BEFORE <= n_index <= seed_index + CONTEXT_EXPANSION_AFTER: neighbor_rows.append((n_index, n_id, n_doc, n_meta)) seed_neighbors_added = 0 for n_index, n_id, n_doc, n_meta in sorted(neighbor_rows, key=lambda row: row[0]): if len(items) >= max_chunks: break if seed_neighbors_added >= MAX_EXPANDED_NEIGHBORS: break # 兼容前缀 ID：用原始 ID 和去前缀版本双重匹配 if n_id in existing_ids or n_id in _raw_id_set: continue n_meta = dict(n_meta or {}) if seed_meta.get('_collection') and not n_meta.get('_collection'): n_meta['_collection'] = seed_meta.get('_collection') n_meta['_expanded_neighbor'] = True n_meta['_expanded_from_score'] = seed_dist # Phase 3：记录种子分数 distance = seed_dist + 0.0001 * abs(n_index - seed_index) items.append((n_id, n_doc, n_meta, distance)) existing_ids.add(n_id) _raw_id_set.add(n_id.split('/', 1)[1] if '/' in n_id else n_id) added += 1 seed_neighbors_added += 1 if not added: return results expanded = { 'ids': [[item[0] for item in items]], 'documents': [[item[1] for item in items]], 'metadatas': [[item[2] for item in items]], 'distances': [[item[3] for item in items]], '_expanded_context': {'added': added} } for key in ('_debug', '_score_source', '_enum_query'): if key in results: expanded[key] = results[key] return expanded def _search_with_sub_queries( self, original_query, sub_queries, top_k=5, allowed_levels=None, role=None, department=None, collections=None, source_filter=None ): """ 使用外部传入的子查询并行检索（由意图分析器生成） Args: original_query: 原始查询（用于日志和缓存） sub_queries: 子查询列表 top_k: 每个子查询的返回数量 其他参数同 search_knowledge Returns: 合并后的检索结果 """ logger.info(f"子查询并行检索: '{original_query}' → {len(sub_queries)} 个子查询") # 每个子查询取较少的 top_k，避免结果过多 sub_top_k = max(top_k, 5) all_results = [] for sub_q in sub_queries: try: sub_result = self.search_knowledge( sub_q, top_k=sub_top_k, allowed_levels=allowed_levels, role=role, department=department, collections=collections, source_filter=source_filter, _skip_decomposition=True ) if sub_result and sub_result.get('ids') and sub_result['ids'][0]: all_results.append(sub_result) except Exception as e: logger.warning(f"子查询检索失败: '{sub_q}' - {e}") if not all_results: return get_empty_result() # 合并去重 if len(all_results) == 1: return all_results[0] return self._merge_and_deduplicate(all_results, top_k) def _search_with_decomposition( self, query, decomposer, top_k=5, allowed_levels=None, role=None, department=None, collections=None, source_filter=None ): """ 复杂查询拆分检索：将对比/推理类查询拆分为子查询，分别检索后合并 Args: query: 原始查询 decomposer: QueryDecomposer 实例 top_k: 返回结果数量 其他参数同 search_knowledge Returns: 合并后的检索结果 """ decomposed = decomposer.decompose(query) sub_queries = decomposed.sub_queries if len(sub_queries) <= 1: # 无需拆分，走正常流程 return self.search_knowledge( query, top_k, allowed_levels, role, department, collections, source_filter, _skip_decomposition=True ) logger.info(f"查询拆分: '{query}' → {len(sub_queries)} 个子查询 ({decomposed.query_type})") # 并行检索各子查询 all_results = [] for sub_q in sub_queries: try: sub_result = self.search_knowledge( sub_q, top_k=top_k, allowed_levels=allowed_levels, role=role, department=department, collections=collections, source_filter=source_filter, _skip_decomposition=True ) if sub_result and sub_result.get('ids') and sub_result['ids'][0]: all_results.append(sub_result) except Exception as e: logger.warning(f"子查询检索失败: '{sub_q}' - {e}") if not all_results: return get_empty_result() # 合并去重 if len(all_results) == 1: merged = all_results[0] else: merged = self._merge_and_deduplicate(all_results, top_k) return merged def _merge_and_deduplicate(self, results_list, top_k): """ 合并多个检索结果并按文档ID去重 Args: results_list: 多个检索结果列表 top_k: 返回数量 Returns: 去重合并后的结果 """ seen_ids = set() merged_ids = [] merged_docs = [] merged_metas = [] merged_distances = [] # 按距离排序合并（距离越小越好） all_items = [] for result in results_list: ids = result.get('ids', [[]])[0] docs = result.get('documents', [[]])[0] metas = result.get('metadatas', [[]])[0] dists = result.get('distances', [[]])[0] if result.get('distances') else [0] * len(ids) for doc_id, doc, meta, dist in zip(ids, docs, metas, dists): if doc_id not in seen_ids: seen_ids.add(doc_id) all_items.append((doc_id, doc, meta, dist)) # 按距离排序（升序） all_items.sort(key=lambda x: x[3]) # 截取 top_k for doc_id, doc, meta, dist in all_items[:top_k]: merged_ids.append(doc_id) merged_docs.append(doc) merged_metas.append(meta) merged_distances.append(dist) return { 'ids': [merged_ids], 'documents': [merged_docs], 'metadatas': [merged_metas], 'distances': [merged_distances] } def _search_multi_kb(self, query, top_k=5, role=None, department=None, target_collections=None, source_filter=None, _debug=None): """ 多向量库检索 Args: query: 查询文本 top_k: 返回结果数量 role: 用户角色 department: 用户部门 target_collections: 指定查询的向量库列表 source_filter: 文件名过滤，精确匹配 """ if target_collections is None: if role and department: from auth.gateway import get_accessible_collections accessible = get_accessible_collections(role, department, 'read') target_collections = self.kb_router.route(query, role, department, accessible) else: target_collections = ['public_kb'] if not target_collections: return get_empty_result() query_vector = self.embedding_model.encode(query).tolist() # 扩大召回数量，以便过滤废止切片后仍有足够结果 recall_k = RERANK_CANDIDATES if (USE_RERANK or USE_HYBRID_SEARCH) else top_k recall_k = max(recall_k, top_k * RECALL_MULTIPLIER) # ========== P0：图片独立召回数量 ========== image_recall_k = max(5, top_k // 2) # ========== 并行检索各向量库 ========== from concurrent.futures import ThreadPoolExecutor, as_completed def _query_single_collection(coll_name): """查询单个向量库（向量 + BM25）""" coll_results = [] try: coll = self.kb_manager.get_collection(coll_name) if not coll: return coll_results query_kwargs = { "query_embeddings": [query_vector], "n_results": recall_k } if source_filter: query_kwargs["where"] = {"source": source_filter} results = coll.query(**query_kwargs) if results['metadatas'] and results['metadatas'][0]: for meta in results['metadatas'][0]: meta['_collection'] = coll_name coll_results.append(results) if USE_HYBRID_SEARCH: try: bm25 = self.kb_manager.get_bm25_index(coll_name) if bm25.bm25: bm25_res = bm25.search(query, top_k=recall_k) if source_filter and bm25_res['metadatas'] and bm25_res['metadatas'][0]: bm25_res = self._filter_results(bm25_res, lambda meta: meta.get('source') == source_filter) if bm25_res['metadatas'] and bm25_res['metadatas'][0]: for meta in bm25_res['metadatas'][0]: meta['_collection'] = coll_name coll_results.append(bm25_res) except Exception as e: logger.debug(f"向量库 {coll_name} 检索失败: {e}") except Exception as e: logger.debug(f"多向量库检索失败: {e}") return coll_results all_results = [] with ThreadPoolExecutor(max_workers=len(target_collections)) as executor: futures = {executor.submit(_query_single_collection, name): name for name in target_collections} for future in as_completed(futures): all_results.extend(future.result()) # ========== FAQ 检索 ========== faq_results = self._search_faq_collection(query_vector, top_k=FAQ_RECALL_TOP_K) if faq_results and faq_results.get('ids') and faq_results['ids'][0]: for meta in faq_results['metadatas'][0]: meta['_collection'] = 'faq_kb' all_results.append(faq_results) # ========== P0：图片独立检索（多知识库模式）========== for coll_name in target_collections: try: coll = self.kb_manager.get_collection(coll_name) if not coll: continue image_results = self._search_image_chunks_multi_kb( query_vector, image_recall_k, coll, source_filter ) if image_results and image_results.get('ids') and image_results['ids'][0]: for meta in image_results['metadatas'][0]: meta['_collection'] = coll_name all_results.append(image_results) except Exception as e: logger.debug(f"图片切片检索失败: {e}") if not all_results: return get_empty_result() if len(all_results) == 1: fused_results = all_results[0] else: # 动态 RRF 权重（P3-1） vector_w, bm25_w = self._get_dynamic_rrf_weights(query) weights = [vector_w if i % 2 == 0 else bm25_w for i in range(len(all_results))] fused_results = self.reciprocal_rank_fusion(all_results, weights) if _debug is not None: _debug['steps'].append({'name': 'rrf_fusion', 'count': len(fused_results['ids'][0]) if fused_results.get('ids') else 0, 'inputs': len(all_results)}) # 过滤废止切片（status != "active" 或 status == "deprecated"） before_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 fused_results = self._filter_deprecated_chunks(fused_results) after_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 if _debug is not None and before_count != after_count: _debug['steps'].append({'name': 'deprecated_filter', 'removed': before_count - after_count}) is_enum_query = self._is_enumeration_query(query) fused_results['_enum_query'] = is_enum_query # 章节过滤（如果查询中提到了章节） fused_results = self._filter_by_section(fused_results, query) # 补齐强命中切片周围的连续文本（MMR 前扩展，防止邻居被 MMR 当作冗余去除） before_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 fused_results = self._expand_contiguous_chunks(fused_results, top_k=top_k) after_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 if _debug is not None and after_count != before_count: _debug['steps'].append({'name': 'context_expansion', 'before': before_count, 'after': after_count}) # ========== MMR 去重（P3-2：前置到 rerank 前）========== if MMR_ENABLED: before_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 fused_results = self._apply_mmr(query, fused_results, top_k=MMR_TOP_K, is_enum_query=is_enum_query) after_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 if _debug is not None: _debug['steps'].append({'name': 'mmr_dedup', 'before': before_count, 'after': after_count}) if USE_RERANK and self.reranker: fused_results = self.rerank_results(query, fused_results, top_k) if _debug is not None: _debug['steps'].append({'name': 'rerank', 'count': len(fused_results['ids'][0]) if fused_results.get('ids') else 0}) else: final_top_k = max(top_k, CONTEXT_EXPANSION_MAX_CHUNKS) if is_enum_query else top_k fused_results = self._truncate_results(fused_results, final_top_k) # FAQ 分数加权 fused_results = self._boost_faq_chunks(fused_results) # ========== 黑名单过滤（负反馈降权，带缓存）========== blacklist = self._get_blacklist() if blacklist: before_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 fused_results = self.filter_blacklisted_chunks(fused_results, blacklist) after_count = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 if _debug is not None and before_count != after_count: _debug['steps'].append({'name': 'blacklist_filter', 'removed': before_count - after_count}) # 时间衰减 fused_results = self._apply_time_decay(fused_results) # ========== 上下文扩展：补充强命中切片周围的连续文本（rerank 之后，防止被截断）========== # Phase 3：仅对高分种子扩展邻居 before_exp = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 fused_results = self._expand_contiguous_chunks(fused_results, top_k=top_k, min_score=EXPANSION_SCORE_THRESHOLD) after_exp = len(fused_results['ids'][0]) if fused_results.get('ids') else 0 if _debug is not None: _debug['steps'].append({'name': 'context_expansion', 'before': before_exp, 'after': after_exp}) # 自适应 TopK：根据置信度调整返回数量 if ( self._adaptive_topk and fused_results.get('distances') and fused_results['distances'][0] and not (is_enum_query and ENUM_QUERY_DISABLE_TOPK_SHRINK) and fused_results.get('_score_source') != 'rrf' ): top_score = 1.0 - fused_results['distances'][0][0] # 距离转相似度 adjusted_k, should_retrieve, reason = self._adaptive_topk.adjust(top_score, top_k) if "high_confidence" in reason: # 高置信度时截断结果 fused_results = self._truncate_results(fused_results, adjusted_k) if _debug is not None: _debug['steps'].append({'name': 'adaptive_topk', 'adjusted_k': adjusted_k, 'reason': reason}) elif _debug is not None and is_enum_query and ENUM_QUERY_DISABLE_TOPK_SHRINK: _debug['steps'].append({'name': 'adaptive_topk', 'skipped': True, 'reason': 'enum_query_preserve_context'}) return fused_results def _filter_deprecated_chunks(self, results: dict) -> dict: """ 过滤废止切片，只保留 active 状态的切片 Args: results: 检索结果 Returns: 过滤后的结果 """ if not results['metadatas'] or not results['metadatas'][0]: return results filtered_ids = [] filtered_docs = [] filtered_metas = [] filtered_distances = [] for i, (doc_id, doc, meta, dist) in enumerate(zip( results['ids'][0], results['documents'][0], results['metadatas'][0], results['distances'][0] if results['distances'] else [0] * len(results['ids'][0]) )): # 只保留 active 状态的切片（未标记或标记为 active） status = meta.get('status', 'active') if status == 'active': filtered_ids.append(doc_id) filtered_docs.append(doc) filtered_metas.append(meta) filtered_distances.append(dist) filtered = { 'ids': [filtered_ids], 'documents': [filtered_docs], 'metadatas': [filtered_metas], 'distances': [filtered_distances] } for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'): if key in results: filtered[key] = results[key] return filtered def _filter_by_section(self, results: dict, query: str) -> dict: """ 根据查询中的章节信息过滤结果 如果查询中明确提到了章节（如"第一章"、"一、"），则优先返回该章节的切片。 Args: results: 检索结果 query: 用户查询 Returns: 过滤后的结果 """ if not SECTION_FILTER_ENABLED: return results if not results['metadatas'] or not results['metadatas'][0]: return results # 从查询中提取章节关键词 import re section_patterns = [ r'第[一二三四五六七八九十\d]+章', r'第\s*\d+\s*章', r'[一二三四五六七八九十]+、', ] mentioned_sections = [] for pattern in section_patterns: matches = re.findall(pattern, query) mentioned_sections.extend(matches) if not mentioned_sections: return results # 过滤切片 filtered_ids = [] filtered_docs = [] filtered_metas = [] filtered_distances = [] for i, (doc_id, doc, meta, dist) in enumerate(zip( results['ids'][0], results['documents'][0], results['metadatas'][0], results['distances'][0] if results['distances'] else [0] * len(results['ids'][0]) )): section = meta.get('section', meta.get('section_path', '')) # 检查是否匹配任一章节 for section_kw in mentioned_sections: if section_kw in section or section_kw in doc: filtered_ids.append(doc_id) filtered_docs.append(doc) filtered_metas.append(meta) filtered_distances.append(dist) break # 如果过滤后结果为空，返回原始结果 if not filtered_ids: return results filtered = { 'ids': [filtered_ids], 'documents': [filtered_docs], 'metadatas': [filtered_metas], 'distances': [filtered_distances] } for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'): if key in results: filtered[key] = results[key] return filtered # ---------------- 底层融合辅助算法 ---------------- def _get_dynamic_rrf_weights(self, query: str) -> tuple: """ 动态 RRF 权重计算（P3-1） 根据查询长度动态调整权重，不依赖分类： - 短查询（< 15字）：偏向关键词检索（BM25优先） - 中等查询（15-50字）：平衡权重 - 长查询（> 50字）：偏向语义检索（Vector优先） Args: query: 用户查询文本 Returns: (vector_weight, bm25_weight): 权重元组 """ query_len = len(query) # 连续权重函数 # alpha: 0 (短查询) → 1 (长查询) alpha = min(1.0, max(0.0, (query_len - 15) / 35)) # vector_weight: 0.3 (短) → 0.7 (长) vector_weight = 0.3 + 0.4 * alpha # bm25_weight = 1 - vector_weight bm25_weight = 1.0 - vector_weight return (vector_weight, bm25_weight) def _apply_mmr(self, query: str, results: dict, top_k: int = 30, is_enum_query: bool = False) -> dict: """ 应用 MMR 去重（P3-2） 支持两种方案： - 高精度版：基于语义向量，需要重新编码文档（MMR_USE_EMBEDDING=True） - 轻量版：基于文本 Jaccard 相似度，零额外计算（MMR_USE_EMBEDDING=False） Args: query: 用户查询 results: 检索结果 top_k: MMR 保留数量 Returns: 去重后的结果 """ if not results.get('ids') or not results['ids'][0]: return results if len(results['ids'][0]) <= top_k: return results try: # 获取 MMR 方案配置，默认使用高精度版 try: from config import MMR_USE_EMBEDDING except ImportError: MMR_USE_EMBEDDING = True # 默认使用高精度版 if MMR_USE_EMBEDDING: # === 高精度版：基于语义向量 === from core.mmr import mmr_rerank # 获取查询向量 query_emb = np.array(self.embedding_model.encode(query)) # 批量编码所有文档 docs_list = results['documents'][0] all_embeddings = self.embedding_model.encode(docs_list) # 构建候选列表 candidates = [] for i, (doc_id, doc, meta) in enumerate(zip( results['ids'][0], results['documents'][0], results['metadatas'][0] )): doc_emb = all_embeddings[i] candidates.append({ 'id': doc_id, 'content': doc, 'embedding': doc_emb, 'metadata': meta }) # MMR 去重 try: from config import MMR_LAMBDA except ImportError: MMR_LAMBDA = 0.5 lambda_param = ENUM_QUERY_MMR_LAMBDA if is_enum_query else MMR_LAMBDA selected = mmr_rerank( query_emb, candidates, top_k=top_k, lambda_param=lambda_param ) selected_ids = [c['id'] for c in selected] orig_ids = results['ids'][0] orig_dists = results.get('distances', [[0] * len(orig_ids)])[0] id_to_dist = dict(zip(orig_ids, orig_dists)) filtered = { 'ids': [[c['id'] for c in selected]], 'documents': [[c['content'] for c in selected]], 'metadatas': [[c['metadata'] for c in selected]], 'distances': [[id_to_dist.get(doc_id, 0) for doc_id in selected_ids]] } for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'): if key in results: filtered[key] = results[key] return filtered else: # === 轻量版：基于文本相似度 === from core.mmr import mmr_filter_by_content # 构建候选列表（无需 embedding） candidates = [] for doc_id, doc, meta in zip( results['ids'][0], results['documents'][0], results['metadatas'][0] ): candidates.append({ 'id': doc_id, 'content': doc, 'metadata': meta }) # 文本去重（Jaccard 相似度） selected = mmr_filter_by_content( candidates, top_k=top_k, similarity_threshold=0.85 ) # 重建结果格式，保留对应的 distances selected_ids = {c['id'] for c in selected} orig_ids = results['ids'][0] orig_dists = results.get('distances', [[0] * len(orig_ids)])[0] id_to_dist = dict(zip(orig_ids, orig_dists)) filtered = { 'ids': [[c['id'] for c in selected]], 'documents': [[c['content'] for c in selected]], 'metadatas': [[c['metadata'] for c in selected]], 'distances': [[id_to_dist.get(c['id'], 0) for c in selected]] } for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'): if key in results: filtered[key] = results[key] return filtered except Exception as e: # MMR 失败时返回原结果 return results def reciprocal_rank_fusion(self, results_list, weights=None, k=None): if k is None: k = RRF_K if not results_list: return get_empty_result() if weights is None: weights = [1.0] * len(results_list) # 使用 (collection, doc_id) 复合键去重，防止跨库同名文件的结果被吞 doc_scores = {} for results, weight in zip(results_list, weights): if not results['documents'] or not results['documents'][0]: continue for rank, (doc_id, doc, meta) in enumerate(zip( results['ids'][0], results['documents'][0], results['metadatas'][0] )): rrf_score = weight / (k + rank + 1) coll = meta.get('_collection') or meta.get('collection') or '' # 复合键：同 collection 内同名 doc_id 合并分数（向量+BM25），跨 collection 不合并 composite_key = f"{coll}\x00{doc_id}" if coll else doc_id if composite_key not in doc_scores: doc_scores[composite_key] = {'score': 0.0, 'doc': doc, 'meta': meta, 'coll': coll, 'raw_id': doc_id} doc_scores[composite_key]['score'] += rrf_score sorted_items = sorted(doc_scores.items(), key=lambda x: x[1]['score'], reverse=True) # 输出 ID 加 collection 前缀，确保跨库同名切片在全链路中可区分 out_ids = [] for item in sorted_items: coll = item[1]['coll'] raw_id = item[1]['raw_id'] if coll and not raw_id.startswith(f"{coll}/"): out_ids.append(f"{coll}/{raw_id}") else: out_ids.append(raw_id) return { 'ids': [out_ids], 'documents': [[item[1]['doc'] for item in sorted_items]], 'metadatas': [[item[1]['meta'] for item in sorted_items]], 'distances': [[item[1]['score'] for item in sorted_items]], '_score_source': 'rrf' } def rerank_results(self, query, results, top_k=5): """Rerank 重排，支持缓存和计时""" if not self.reranker or not results['documents'][0]: return results t_start = time.time() doc_ids = results['ids'][0] doc_count = len(doc_ids) # 尝试从缓存获取 Rerank 分数（{doc_id: score} 映射，顺序无关） cached_map = None if CACHE_AVAILABLE: try: from config import RERANK_CACHE_ENABLED if RERANK_CACHE_ENABLED: cache = get_cache_manager() cached_map = cache.get_rerank_scores(query, doc_ids) except (ImportError, Exception): pass # 仅当缓存覆盖全部当前文档时才命中，否则重新推理（防止部分缺失） cache_hit = cached_map is not None and all(d in cached_map for d in doc_ids) if cache_hit: # 缓存命中：按当前 doc_ids 顺序查表重建分数，避免位置错位 scores = np.array([cached_map[d] for d in doc_ids]) t_elapsed = time.time() - t_start logger.debug(f"Rerank 缓存命中: {doc_count} 个文档, 耗时 {t_elapsed*1000:.1f}ms") else: # 缓存未命中，执行模型推理 pairs = [(query, doc) for doc in results['documents'][0]] scores = self.reranker.predict(pairs) t_elapsed = time.time() - t_start logger.debug(f"Rerank 推理: {doc_count} 个文档, 耗时 {t_elapsed*1000:.1f}ms") # 写入缓存 if CACHE_AVAILABLE: try: from config import RERANK_CACHE_ENABLED if RERANK_CACHE_ENABLED: cache = get_cache_manager() cache.set_rerank_scores(query, doc_ids, [float(s) for s in scores]) except (ImportError, Exception): pass sorted_indices = np.argsort(scores)[::-1] reranked = { 'ids': [[results['ids'][0][i] for i in sorted_indices[:top_k]]], 'documents': [[results['documents'][0][i] for i in sorted_indices[:top_k]]], 'metadatas': [[results['metadatas'][0][i] for i in sorted_indices[:top_k]]], 'distances': [[float(scores[i]) for i in sorted_indices[:top_k]]], '_reranked': True, '_rerank_time_ms': int(t_elapsed * 1000), '_rerank_cached': cache_hit } # 保留原有标记字段 for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'): if key in results: reranked[key] = results[key] return reranked # ---------------- 安全与工具 ---------------- def check_restricted_documents(self, query, allowed_levels, top_k=3, role=None, department=None): if not self._initialized: self.initialize() if USE_MULTI_KB and self.kb_manager and role and department: from auth.gateway import get_accessible_collections all_colls = [c.name for c in self.kb_manager.list_collections()] accessible = set(get_accessible_collections(role, department, 'read')) restricted = set(all_colls) - accessible if not restricted: return {"has_restricted": False, "restricted_levels": [], "restricted_sources": []} query_vector = self.embedding_model.encode(query).tolist() found_sources = set() top_score = 0.0 for coll_name in restricted: try: coll = self.kb_manager.get_collection(coll_name) if not coll: continue res = coll.query(query_embeddings=[query_vector], n_results=top_k) if res['metadatas'] and res['metadatas'][0]: for meta in res['metadatas'][0]: found_sources.add(meta.get('source', '未知')) for dist in (res.get('distances', [[]])[0] or []): if dist > top_score: top_score = dist except Exception as e: logger.debug(f"权限检查遍历失败: {e}") return { "has_restricted": len(found_sources) > 0, "restricted_levels": [c.replace('dept_', '') for c in restricted if True][:3], "restricted_sources": list(found_sources)[:3], "top_restricted_score": top_score } if not allowed_levels: return {"has_restricted": False, "restricted_levels": [], "restricted_sources": [], "top_restricted_score": 0.0} restricted_levels = {"public", "internal", "confidential", "secret"} - set(allowed_levels) if not restricted_levels: return {"has_restricted": False, "restricted_levels": [], "restricted_sources": [], "top_restricted_score": 0.0} query_vector = self.embedding_model.encode(query).tolist() try: res = self.collection.query( query_embeddings=[query_vector], n_results=top_k, where={"security_level": {"$in": list(restricted_levels)}} ) docs = res.get('documents', [[]])[0] if not docs: return {"has_restricted": False, "restricted_levels": [], "restricted_sources": [], "top_restricted_score": 0.0} metas = res.get('metadatas', [[]])[0] dists = res.get('distances', [[]])[0] found_levels, found_sources, top_score = set(), set(), 0.0 for meta, dist in zip(metas, dists): found_levels.add(meta.get('security_level', 'public')) found_sources.add(meta.get('source', '未知')) if dist > top_score: top_score = dist return { "has_restricted": True, "restricted_levels": list(found_levels), "restricted_sources": list(found_sources)[:3], "top_restricted_score": top_score } except Exception as e: logger.warning(f"受限内容检查失败: {e}") return {"has_restricted": False, "restricted_levels": [], "restricted_sources": [], "top_restricted_score": 0.0} def generate_answer(self, query, context): """底层生成答复能力""" prompt = f"""你是一个严谨的智能助手，请根据以下参考资料回答用户的问题。 ... 参考资料： {context} 用户问题：{query} 请回答：""" try: from core.llm_utils import call_llm result = call_llm( self.llm_client, prompt, MODEL, temperature=LLM_TEMPERATURE, max_tokens=LLM_MAX_TOKENS ) return result or f"调用大模型失败: 返回结果为空" except Exception as e: return f"调用大模型失败: {str(e)}" def generate_answer_stream(self, query, context, history=None): """ 流式生成答复 Args: query: 用户问题 context: 检索到的上下文 history: 对话历史 [{"role": "user/assistant", "content": "..."}] Yields: str: 每个 token """ # 构建消息列表 messages = [] # 添加历史对话 if history: for h in history: messages.append({ "role": h.get("role", "user"), "content": h.get("content", "") }) # Bug 3 修复：添加 system prompt + 强化指令，避免 LLM 编造/忽略参考资料 messages.insert(0, { "role": "system", "content": ( "你是一个严谨的知识库问答助手。" "你必须且只能根据用户提供的【参考资料】回答问题。" "如果参考资料中有答案，必须引用对应内容回答，并在回答末尾标注引用编号（如[1]、[2]）。" "如果参考资料中确实没有相关信息，简短说明即可，不要编造或补充资料外的内容。" "禁止使用参考资料以外的知识进行补充或推测。" ) }) # 添加当前问题（带上下文）- 强化指令 if context: user_message = f"""【参考资料】 {context} 【用户问题】 {query} 请仔细阅读以上全部参考资料后回答。如果参考资料中包含相关内容，必须引用回答并标注编号。如果资料中没有相关信息，请明确说明。""" else: user_message = query messages.append({"role": "user", "content": user_message}) try: from core.llm_utils import call_llm_stream for text in call_llm_stream( self.llm_client, "", MODEL, temperature=LLM_TEMPERATURE, max_tokens=LLM_MAX_TOKENS, messages=messages, error_prefix="[错误]" ): yield text except Exception as e: yield f"[错误] 调用大模型失败: {str(e)}" # 快捷访问单例 def get_engine() -> RAGEngine: """ 获取全局 RAGEngine 单例 首次调用时自动初始化引擎（加载模型、连接数据库）。 Returns: 已初始化的 RAGEngine 实例 Example: >>> engine = get_engine() >>> result = engine.search_knowledge("查询问题", top_k=10) """ engine = RAGEngine.get_instance() if not engine._initialized: engine.initialize() return engine