2138 lines
87 KiB
Python
2138 lines
87 KiB
Python
"""
|
||
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 re
|
||
import time
|
||
import logging
|
||
import threading
|
||
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
|
||
_engine_lock = threading.Lock() # 单例创建锁，防止多线程竞态条件
|
||
|
||
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 = "cloud"
|
||
RERANK_CLOUD_MODEL = "xop3qwen8breranker"
|
||
RERANK_CLOUD_API_KEY = ""
|
||
RERANK_CLOUD_BASE_URL = "https://maas-api.cn-huabei-1.xf-yun.com/v2/rerank"
|
||
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
|
||
}
|
||
|
||
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:
|
||
with _engine_lock:
|
||
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._encode_cached(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 _encode_cached(self, text):
|
||
"""
|
||
带缓存的 embedding 编码
|
||
|
||
优先从 Embedding Cache（LRU）读取，未命中再调用模型编码并写入缓存。
|
||
支持单文本和批量文本输入。
|
||
|
||
Args:
|
||
text: 单个文本字符串 或 文本列表
|
||
|
||
Returns:
|
||
numpy 数组（单文本为一维，批量为二维）
|
||
"""
|
||
import numpy as _np
|
||
|
||
# 检查 embedding 缓存是否启用（缓存配置查询结果，避免每次重复导入）
|
||
if not hasattr(self, '_emb_cache_enabled'):
|
||
self._emb_cache_enabled = True # 默认启用
|
||
if CACHE_AVAILABLE:
|
||
try:
|
||
from config import EMBEDDING_CACHE_ENABLED
|
||
self._emb_cache_enabled = EMBEDDING_CACHE_ENABLED
|
||
except ImportError:
|
||
pass
|
||
|
||
if not self._emb_cache_enabled:
|
||
return self.embedding_model.encode(text)
|
||
|
||
try:
|
||
_cache = get_cache_manager()
|
||
except Exception:
|
||
return self.embedding_model.encode(text)
|
||
|
||
# 批量输入
|
||
if isinstance(text, list):
|
||
try:
|
||
cached_embs, missed_indices = _cache.get_embeddings_batch(text)
|
||
if missed_indices:
|
||
missed_texts = [text[i] for i in missed_indices]
|
||
# encode(list) 始终返回 2D ndarray，直接按行索引即可
|
||
new_embs = self.embedding_model.encode(missed_texts)
|
||
if len(missed_indices) == 1:
|
||
# 单条时 encode 可能返回 1D，需统一处理
|
||
if new_embs.ndim == 1:
|
||
new_embs = new_embs.reshape(1, -1)
|
||
for idx, mi in enumerate(missed_indices):
|
||
emb_list = new_embs[idx].tolist()
|
||
cached_embs[mi] = emb_list
|
||
try:
|
||
_cache.set_embedding(text[mi], emb_list)
|
||
except Exception:
|
||
pass
|
||
return _np.array(cached_embs)
|
||
except Exception:
|
||
# 缓存故障时优雅降级为直接编码
|
||
return self.embedding_model.encode(text)
|
||
|
||
# 单文本输入
|
||
cached = _cache.get_embedding(text)
|
||
if cached is not None:
|
||
return _np.array(cached)
|
||
|
||
embedding = self.embedding_model.encode(text)
|
||
try:
|
||
emb_list = embedding.tolist() if hasattr(embedding, 'tolist') else list(embedding)
|
||
_cache.set_embedding(text, emb_list)
|
||
except Exception:
|
||
pass
|
||
return embedding
|
||
|
||
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': []}
|
||
|
||
# 扩展同 section 的 text 邻居
|
||
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"}]})
|
||
|
||
# 同时扩展同 section 的 table 邻居（table 切片的 rerank 分数往往偏低，
|
||
# 但与同 section 的 text 切片属于同一语义单元，不应割裂）
|
||
table_where = {"$and": [{"source": source}, {"chunk_type": "table"}]}
|
||
if section:
|
||
table_where["$and"].append({"section": section})
|
||
table_neighbors = _get_neighbors(table_where)
|
||
|
||
# 当 section 为空时，table 查询只有 source 条件，可能拉入大量无关表格，
|
||
# 缩小 chunk_index 窗口至 ±1 以降低噪音；有 section 时使用正常窗口
|
||
if section:
|
||
_t_before, _t_after = CONTEXT_EXPANSION_BEFORE, CONTEXT_EXPANSION_AFTER
|
||
else:
|
||
_t_before, _t_after = 1, 1
|
||
|
||
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))
|
||
|
||
# 同 section 的 table 邻居也加入扩展范围
|
||
for n_id, n_doc, n_meta in zip(
|
||
table_neighbors.get('ids', []),
|
||
table_neighbors.get('documents', []),
|
||
table_neighbors.get('metadatas', [])
|
||
):
|
||
n_index = self._to_int(n_meta.get('chunk_index'))
|
||
if n_index is None:
|
||
continue
|
||
if seed_index - _t_before <= n_index <= seed_index + _t_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._encode_cached(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
|
||
|
||
# 获取查询向量（使用 embedding 缓存）
|
||
query_emb = np.array(self._encode_cached(query))
|
||
|
||
# 批量编码所有文档（使用 embedding 缓存）
|
||
docs_list = results['documents'][0]
|
||
all_embeddings = self._encode_cached(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 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]）。"
|
||
"如果参考资料中确实没有相关信息，简短说明即可，不要编造或补充资料外的内容。"
|
||
"禁止使用参考资料以外的知识进行补充或推测。"
|
||
"【重要-表格处理规则】当用户询问表格、要求展示表格内容时，你必须将参考资料中的 Markdown 表格原样输出（保留 | 分隔符和表格结构），"
|
||
"不要仅用文字描述表格存在或仅列出章节名称。如果参考资料中多个章节都有表格，"
|
||
"优先展示与用户问题最相关的表格完整内容。"
|
||
)
|
||
})
|
||
|
||
# 添加当前问题（带上下文）- 强化指令
|
||
if context:
|
||
# 检测用户问题是否涉及表格，加入针对性指令
|
||
_table_hint = ""
|
||
# 检测上下文中是否包含 Markdown 表格（数据驱动，无需硬编码关键词）
|
||
_has_table_in_context = bool(re.search(r'\|.+\|', context)) if context else False
|
||
if _has_table_in_context:
|
||
_table_hint = "\n注意：参考资料中包含 Markdown 格式的表格数据，请务必将相关表格以原始 Markdown 表格格式完整展示在回答中，不要仅用文字描述。"
|
||
|
||
user_message = f"""【参考资料】
|
||
{context}
|
||
|
||
【用户问题】
|
||
{query}
|
||
|
||
请仔细阅读以上全部参考资料后回答。注意参考资料中标有章节路径，请根据章节路径准确定位相关内容。如果参考资料中包含相关内容，必须引用回答并标注编号。如果资料中没有相关信息，请明确说明。{_table_hint}"""
|
||
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
|