feat: sync core logic from main (cache fix, intent analyzer, engine improvements)
- core/cache.py: fix GET/SET key mismatch, use coarse-grained kb_version keys - core/intent_analyzer.py: add cache_type tagging, enhance history parsing - core/engine.py: adaptive topk score source fix, BM25 top-3 preservation, section cluster boost
This commit is contained in:
lacerate551
@@ -224,39 +224,36 @@ class RAGCacheManager:
|
||||
"""
|
||||
获取查询缓存结果
|
||||
|
||||
始终使用粗粒度 key（基于 kb_version），确保 GET/SET key 一致。
|
||||
|
||||
Args:
|
||||
query: 查询文本
|
||||
kb_name: 知识库名称
|
||||
doc_ids: 相关文档 ID 列表（用于细粒度缓存 key）
|
||||
doc_ids: 保留参数以兼容调用方签名（当前未使用）
|
||||
"""
|
||||
kb_version = self.get_kb_version(kb_name)
|
||||
|
||||
# 计算文档哈希（如果提供了 doc_ids）
|
||||
doc_hash = ""
|
||||
if doc_ids:
|
||||
doc_hash = self._compute_doc_hash(kb_name, doc_ids)
|
||||
|
||||
key = self._make_query_cache_key(query, kb_name, kb_version, doc_hash)
|
||||
# 使用粗粒度 key，与 SET 保持一致
|
||||
key = self._make_query_cache_key(query, kb_name, kb_version)
|
||||
return self.query_cache.get(key)
|
||||
|
||||
def set_query_result(self, query: str, kb_name: str, result: Dict, doc_ids: List[str] = None) -> None:
|
||||
"""
|
||||
设置查询缓存结果
|
||||
|
||||
始终使用粗粒度 key（基于 kb_version），确保 GET/SET key 一致。
|
||||
kb_version 在文档变更时自增，触发整个知识库的缓存失效。
|
||||
|
||||
Args:
|
||||
query: 查询文本
|
||||
kb_name: 知识库名称
|
||||
result: 缓存结果
|
||||
doc_ids: 相关文档 ID 列表（用于细粒度失效）
|
||||
doc_ids: 保留参数以兼容调用方签名（当前未使用）
|
||||
"""
|
||||
kb_version = self.get_kb_version(kb_name)
|
||||
|
||||
# 计算相关文档的版本哈希（细粒度失效）
|
||||
doc_hash = ""
|
||||
if doc_ids:
|
||||
doc_hash = self._compute_doc_hash(kb_name, doc_ids)
|
||||
|
||||
key = self._make_query_cache_key(query, kb_name, kb_version, doc_hash)
|
||||
# 使用与 GET 相同的粗粒度 key，确保缓存可命中
|
||||
key = self._make_query_cache_key(query, kb_name, kb_version)
|
||||
self.query_cache.set(key, result, kb_version=kb_version)
|
||||
|
||||
def _compute_doc_hash(self, kb_name: str, doc_ids: List[str]) -> str:
|
||||
|
||||
570
core/engine.py
570
core/engine.py
@@ -124,10 +124,10 @@ except ImportError:
|
||||
EMBEDDING_DEVICE = "auto"
|
||||
RERANK_DEVICE = "auto"
|
||||
RERANK_USE_ONNX = False
|
||||
RERANK_BACKEND = "local"
|
||||
RERANK_CLOUD_MODEL = "qwen3-rerank"
|
||||
RERANK_BACKEND = "cloud"
|
||||
RERANK_CLOUD_MODEL = "xop3qwen8breranker"
|
||||
RERANK_CLOUD_API_KEY = ""
|
||||
RERANK_CLOUD_BASE_URL = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks"
|
||||
RERANK_CLOUD_BASE_URL = "https://maas-api.cn-huabei-1.xf-yun.com/v2/rerank"
|
||||
RERANK_CLOUD_TIMEOUT = 15
|
||||
|
||||
|
||||
@@ -285,8 +285,7 @@ class CloudReranker:
|
||||
body = {
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"documents": documents,
|
||||
"top_n": len(documents)
|
||||
"documents": documents
|
||||
}
|
||||
|
||||
session = self._get_session()
|
||||
@@ -670,6 +669,7 @@ class RAGEngine:
|
||||
|
||||
results_list = [vector_results]
|
||||
weights = [VECTOR_WEIGHT]
|
||||
bm25_results = None # 初始化，防止 NameError
|
||||
|
||||
if USE_HYBRID_SEARCH and self.bm25_index.bm25:
|
||||
bm25_results = self.bm25_index.search(query, top_k=recall_k)
|
||||
@@ -684,6 +684,22 @@ class RAGEngine:
|
||||
vector_w, bm25_w = self._get_dynamic_rrf_weights(query)
|
||||
weights = [vector_w, bm25_w]
|
||||
|
||||
# ========== 保留 BM25 原始 top-3 完整信息，用于下游分歧检测救援 ==========
|
||||
_bm25_raw_top3 = []
|
||||
if USE_HYBRID_SEARCH and bm25_results and bm25_results.get('ids') and bm25_results['ids'][0]:
|
||||
_bm25_ids = bm25_results['ids'][0][:3]
|
||||
_bm25_docs = bm25_results['documents'][0][:3]
|
||||
_bm25_metas = bm25_results['metadatas'][0][:3]
|
||||
_bm25_dists = (bm25_results.get('distances', [[]])[0] or [0]*3)[:3]
|
||||
for i in range(len(_bm25_ids)):
|
||||
_bm25_raw_top3.append({
|
||||
'id': _bm25_ids[i],
|
||||
'doc': _bm25_docs[i],
|
||||
'meta': _bm25_metas[i],
|
||||
'bm25_score': _bm25_dists[i],
|
||||
'rank': i + 1
|
||||
})
|
||||
|
||||
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]})
|
||||
@@ -699,6 +715,8 @@ class RAGEngine:
|
||||
|
||||
is_enum_query = self._is_enumeration_query(query)
|
||||
fused_results['_enum_query'] = is_enum_query
|
||||
# 传递 BM25 原始 top-3 到路由层，用于分歧检测救援
|
||||
fused_results['_bm25_top3'] = _bm25_raw_top3
|
||||
|
||||
# 章节过滤（如果查询中提到了章节）
|
||||
fused_results = self._filter_by_section(fused_results, query)
|
||||
@@ -769,7 +787,15 @@ class RAGEngine:
|
||||
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] # 距离转相似度
|
||||
# 根据分数来源计算相似度分数（越高越好）
|
||||
score_source = fused_results.get('_score_source')
|
||||
top_dist = fused_results['distances'][0][0]
|
||||
if score_source == 'rerank':
|
||||
# Rerank 后 distances 是相关性分数，越大越好，直接使用
|
||||
top_score = top_dist
|
||||
else:
|
||||
# 向量距离，越小越好，转为相似度
|
||||
top_score = 1.0 - top_dist
|
||||
adjusted_k, should_retrieve, reason = self._adaptive_topk.adjust(top_score, top_k)
|
||||
if "high_confidence" in reason:
|
||||
# 高置信度时截断结果
|
||||
@@ -1108,7 +1134,7 @@ class RAGEngine:
|
||||
'metadatas': [f_metas],
|
||||
'distances': [f_scores]
|
||||
}
|
||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'):
|
||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||
if key in results:
|
||||
filtered[key] = results[key]
|
||||
return filtered
|
||||
@@ -1131,7 +1157,7 @@ class RAGEngine:
|
||||
'metadatas': [f_metas],
|
||||
'distances': [f_scores]
|
||||
}
|
||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'):
|
||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||
if key in results:
|
||||
filtered[key] = results[key]
|
||||
return filtered
|
||||
@@ -1146,7 +1172,7 @@ class RAGEngine:
|
||||
'metadatas': [results['metadatas'][0][:top_k]],
|
||||
'distances': [results['distances'][0][:top_k]]
|
||||
}
|
||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'):
|
||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||
if key in results:
|
||||
truncated[key] = results[key]
|
||||
return truncated
|
||||
@@ -1181,166 +1207,6 @@ class RAGEngine:
|
||||
return None
|
||||
return self.collection
|
||||
|
||||
def _expand_contiguous_chunks(self, results: dict, top_k: int = None,
|
||||
min_score: float = 0.0, query: str = '') -> dict:
|
||||
"""Add same-source same-section neighbor text chunks around strong hits.
|
||||
|
||||
Args:
|
||||
results: 检索结果
|
||||
top_k: 最大切片数
|
||||
min_score: Phase 3 最低分数阈值，仅对 Rerank 分数高于此值的种子扩展
|
||||
query: 查询文本，用于词法匹配辅助种子资格判定
|
||||
"""
|
||||
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' or meta.get('_cluster_boosted'))
|
||||
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:
|
||||
# 词法匹配豁免：CrossEncoder 低分但关键词重叠度高时仍允许作为种子
|
||||
if query and self._chunk_lexical_score(_seed_doc, query) > 0.3:
|
||||
pass # 词法匹配度高，允许作为种子
|
||||
else:
|
||||
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"}]})
|
||||
|
||||
# 同时扩展同 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
|
||||
|
||||
@staticmethod
|
||||
def _normalize_section_path(section_path: str, levels: int = None) -> str:
|
||||
"""归一化 section_path：取前 N 级路径，容忍 MinerU 标题检测误差。
|
||||
@@ -1465,6 +1331,167 @@ class RAGEngine:
|
||||
matched = sum(1 for w in bigrams if w in chunk_text)
|
||||
return matched / len(bigrams)
|
||||
|
||||
def _expand_contiguous_chunks(self, results: dict, top_k: int = None,
|
||||
min_score: float = 0.0, query: str = '') -> dict:
|
||||
"""Add same-source same-section neighbor text chunks around strong hits.
|
||||
|
||||
Args:
|
||||
results: 检索结果
|
||||
top_k: 最大切片数
|
||||
min_score: Phase 3 最低分数阈值，仅对 Rerank 分数高于此值的种子扩展
|
||||
query: 查询文本，用于词法匹配辅助种子资格判定
|
||||
"""
|
||||
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' or meta.get('_cluster_boosted'))
|
||||
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 时生效）
|
||||
# 词法匹配豁免：CrossEncoder 低分但关键词重叠度高时仍允许作为种子
|
||||
if min_score > 0 and seed_dist < min_score:
|
||||
if query and self._chunk_lexical_score(_seed_doc, query) > 0.3:
|
||||
pass # 词法匹配度高，允许作为种子
|
||||
else:
|
||||
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', '_bm25_top3'):
|
||||
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
|
||||
@@ -1487,6 +1514,7 @@ class RAGEngine:
|
||||
sub_top_k = max(top_k, 5)
|
||||
|
||||
all_results = []
|
||||
_all_bm25_top3 = [] # 收集各子查询的 BM25 top3
|
||||
for sub_q in sub_queries:
|
||||
try:
|
||||
sub_result = self.search_knowledge(
|
||||
@@ -1497,6 +1525,9 @@ class RAGEngine:
|
||||
)
|
||||
if sub_result and sub_result.get('ids') and sub_result['ids'][0]:
|
||||
all_results.append(sub_result)
|
||||
# 收集子查询的 BM25 top3
|
||||
if sub_result.get('_bm25_top3'):
|
||||
_all_bm25_top3.extend(sub_result['_bm25_top3'])
|
||||
except Exception as e:
|
||||
logger.warning(f"子查询检索失败: '{sub_q}' - {e}")
|
||||
|
||||
@@ -1505,9 +1536,19 @@ class RAGEngine:
|
||||
|
||||
# 合并去重
|
||||
if len(all_results) == 1:
|
||||
return all_results[0]
|
||||
merged = all_results[0]
|
||||
else:
|
||||
merged = self._merge_and_deduplicate(all_results, top_k)
|
||||
|
||||
return self._merge_and_deduplicate(all_results, top_k)
|
||||
# 将收集的 BM25 top3 传递到合并结果中
|
||||
if _all_bm25_top3:
|
||||
_all_bm25_top3.sort(key=lambda x: x.get('bm25_score', 0), reverse=True)
|
||||
_all_bm25_top3 = _all_bm25_top3[:3]
|
||||
for rank, item in enumerate(_all_bm25_top3):
|
||||
item['rank'] = rank + 1
|
||||
merged['_bm25_top3'] = _all_bm25_top3
|
||||
|
||||
return merged
|
||||
|
||||
def _search_with_decomposition(
|
||||
self, query, decomposer, top_k=5, allowed_levels=None,
|
||||
@@ -1539,6 +1580,7 @@ class RAGEngine:
|
||||
|
||||
# 并行检索各子查询
|
||||
all_results = []
|
||||
_all_bm25_top3 = [] # 收集各子查询的 BM25 top3
|
||||
for sub_q in sub_queries:
|
||||
try:
|
||||
sub_result = self.search_knowledge(
|
||||
@@ -1549,6 +1591,9 @@ class RAGEngine:
|
||||
)
|
||||
if sub_result and sub_result.get('ids') and sub_result['ids'][0]:
|
||||
all_results.append(sub_result)
|
||||
# 收集子查询的 BM25 top3
|
||||
if sub_result.get('_bm25_top3'):
|
||||
_all_bm25_top3.extend(sub_result['_bm25_top3'])
|
||||
except Exception as e:
|
||||
logger.warning(f"子查询检索失败: '{sub_q}' - {e}")
|
||||
|
||||
@@ -1561,6 +1606,14 @@ class RAGEngine:
|
||||
else:
|
||||
merged = self._merge_and_deduplicate(all_results, top_k)
|
||||
|
||||
# 将收集的 BM25 top3 传递到合并结果中
|
||||
if _all_bm25_top3:
|
||||
_all_bm25_top3.sort(key=lambda x: x.get('bm25_score', 0), reverse=True)
|
||||
_all_bm25_top3 = _all_bm25_top3[:3]
|
||||
for rank, item in enumerate(_all_bm25_top3):
|
||||
item['rank'] = rank + 1
|
||||
merged['_bm25_top3'] = _all_bm25_top3
|
||||
|
||||
return merged
|
||||
|
||||
def _merge_and_deduplicate(self, results_list, top_k):
|
||||
@@ -1645,12 +1698,13 @@ class RAGEngine:
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
def _query_single_collection(coll_name):
|
||||
"""查询单个向量库（向量 + BM25）"""
|
||||
"""查询单个向量库（向量 + BM25），返回 (coll_results, bm25_raw_items)"""
|
||||
coll_results = []
|
||||
bm25_raw_items = [] # 该 collection 的 BM25 原始结果
|
||||
try:
|
||||
coll = self.kb_manager.get_collection(coll_name)
|
||||
if not coll:
|
||||
return coll_results
|
||||
return coll_results, bm25_raw_items
|
||||
|
||||
query_kwargs = {
|
||||
"query_embeddings": [query_vector],
|
||||
@@ -1668,25 +1722,61 @@ class RAGEngine:
|
||||
if USE_HYBRID_SEARCH:
|
||||
try:
|
||||
bm25 = self.kb_manager.get_bm25_index(coll_name)
|
||||
if bm25.bm25:
|
||||
if bm25 and bm25.bm25:
|
||||
bm25_res = bm25.search(query, top_k=recall_k)
|
||||
if source_filter and bm25_res['metadatas'] and bm25_res['metadatas'][0]:
|
||||
# 兼容两种 BM25Index：core.bm25_index 返回 dict，knowledge.base 返回 tuple
|
||||
if isinstance(bm25_res, tuple):
|
||||
_ids, _docs, _metas, _dists = bm25_res
|
||||
bm25_res = {
|
||||
'ids': [_ids],
|
||||
'documents': [_docs],
|
||||
'metadatas': [_metas],
|
||||
'distances': [_dists]
|
||||
}
|
||||
if source_filter 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)
|
||||
# 提取 BM25 原始 top-3（在此处直接捕获，避免与向量结果混淆）
|
||||
_bm25_ids = bm25_res['ids'][0][:3]
|
||||
_bm25_docs = bm25_res['documents'][0][:3]
|
||||
_bm25_metas = bm25_res['metadatas'][0][:3]
|
||||
_bm25_dists = (bm25_res.get('distances', [[]])[0] or [0]*3)[:3]
|
||||
for i in range(len(_bm25_ids)):
|
||||
# 确保 meta 包含 _collection（用于路由层注入时下游处理）
|
||||
bm25_meta = _bm25_metas[i]
|
||||
if '_collection' not in bm25_meta:
|
||||
bm25_meta = {**bm25_meta, '_collection': coll_name}
|
||||
bm25_raw_items.append({
|
||||
'id': _bm25_ids[i],
|
||||
'doc': _bm25_docs[i],
|
||||
'meta': bm25_meta,
|
||||
'bm25_score': _bm25_dists[i],
|
||||
})
|
||||
logger.debug(f"[BM25] {coll_name}: captured {len(bm25_raw_items)} raw items")
|
||||
except Exception as e:
|
||||
logger.debug(f"向量库 {coll_name} 检索失败: {e}")
|
||||
logger.debug(f"向量库 {coll_name} BM25检索失败: {e}")
|
||||
except Exception as e:
|
||||
logger.debug(f"多向量库检索失败: {e}")
|
||||
return coll_results
|
||||
return coll_results, bm25_raw_items
|
||||
|
||||
all_results = []
|
||||
_bm25_raw_top3 = []
|
||||
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())
|
||||
coll_results, bm25_raw_items = future.result()
|
||||
all_results.extend(coll_results)
|
||||
_bm25_raw_top3.extend(bm25_raw_items)
|
||||
|
||||
# 按 bm25_score 降序取全局 top-3
|
||||
if _bm25_raw_top3:
|
||||
_bm25_raw_top3.sort(key=lambda x: x['bm25_score'], reverse=True)
|
||||
_bm25_raw_top3 = _bm25_raw_top3[:3]
|
||||
for rank, item in enumerate(_bm25_raw_top3):
|
||||
item['rank'] = rank + 1
|
||||
|
||||
# ========== FAQ 检索 ==========
|
||||
faq_results = self._search_faq_collection(query_vector, top_k=FAQ_RECALL_TOP_K)
|
||||
@@ -1734,6 +1824,8 @@ class RAGEngine:
|
||||
|
||||
is_enum_query = self._is_enumeration_query(query)
|
||||
fused_results['_enum_query'] = is_enum_query
|
||||
# 传递 BM25 原始 top-3 到路由层，用于分歧检测救援
|
||||
fused_results['_bm25_top3'] = _bm25_raw_top3
|
||||
|
||||
# 章节过滤（如果查询中提到了章节）
|
||||
fused_results = self._filter_by_section(fused_results, query)
|
||||
@@ -1801,7 +1893,15 @@ class RAGEngine:
|
||||
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] # 距离转相似度
|
||||
# 根据分数来源计算相似度分数（越高越好）
|
||||
score_source = fused_results.get('_score_source')
|
||||
top_dist = fused_results['distances'][0][0]
|
||||
if score_source == 'rerank':
|
||||
# Rerank 后 distances 是相关性分数，越大越好，直接使用
|
||||
top_score = top_dist
|
||||
else:
|
||||
# 向量距离，越小越好，转为相似度
|
||||
top_score = 1.0 - top_dist
|
||||
adjusted_k, should_retrieve, reason = self._adaptive_topk.adjust(top_score, top_k)
|
||||
if "high_confidence" in reason:
|
||||
# 高置信度时截断结果
|
||||
@@ -1851,7 +1951,7 @@ class RAGEngine:
|
||||
'metadatas': [filtered_metas],
|
||||
'distances': [filtered_distances]
|
||||
}
|
||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'):
|
||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||
if key in results:
|
||||
filtered[key] = results[key]
|
||||
return filtered
|
||||
@@ -1924,7 +2024,7 @@ class RAGEngine:
|
||||
'metadatas': [filtered_metas],
|
||||
'distances': [filtered_distances]
|
||||
}
|
||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'):
|
||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||
if key in results:
|
||||
filtered[key] = results[key]
|
||||
return filtered
|
||||
@@ -2040,7 +2140,7 @@ class RAGEngine:
|
||||
'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'):
|
||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||
if key in results:
|
||||
filtered[key] = results[key]
|
||||
return filtered
|
||||
@@ -2080,7 +2180,7 @@ class RAGEngine:
|
||||
'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'):
|
||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||
if key in results:
|
||||
filtered[key] = results[key]
|
||||
return filtered
|
||||
@@ -2189,109 +2289,15 @@ class RAGEngine:
|
||||
'_rerank_cached': cache_hit
|
||||
}
|
||||
# 保留原有标记字段
|
||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'):
|
||||
for key in ('_debug', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||
if key in results:
|
||||
reranked[key] = results[key]
|
||||
# Rerank 后 distances 语义变为 CrossEncoder 分数，更新 _score_source
|
||||
# 使自适应 TopK 能正确应用（之前 _score_source='rrf' 会导致自适应 TopK 被跳过）
|
||||
reranked['_score_source'] = 'rerank'
|
||||
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):
|
||||
"""
|
||||
|
||||
@@ -316,7 +316,8 @@ class IntentAnalyzer:
|
||||
|
||||
if cache_emb is not None:
|
||||
cached = cache.get(cache_emb)
|
||||
if cached:
|
||||
# 确保缓存条目是意图分析结果（非 RAG 回答缓存）
|
||||
if cached and cached.get("cache_type") != "rag_answer":
|
||||
logger.info(f"意图分析缓存命中: {cached.get('reason', '')[:50]}")
|
||||
return IntentAnalysis.from_dict(cached)
|
||||
else:
|
||||
@@ -386,9 +387,11 @@ class IntentAnalyzer:
|
||||
intent=intent_type
|
||||
)
|
||||
|
||||
# 存入语义缓存
|
||||
# 存入语义缓存（标记类型，避免与 RAG 回答缓存混淆）
|
||||
if cache and cache_emb is not None:
|
||||
cache.set(cache_emb, analysis.to_dict())
|
||||
cache_data = analysis.to_dict()
|
||||
cache_data["cache_type"] = "intent_analysis"
|
||||
cache.set(cache_emb, cache_data)
|
||||
|
||||
# 存入精确匹配缓存
|
||||
if len(self._exact_cache) < self._exact_cache_max:
|
||||
@@ -446,6 +449,7 @@ class IntentAnalyzer:
|
||||
if not history:
|
||||
return "（无历史对话）"
|
||||
|
||||
import re
|
||||
parts = []
|
||||
|
||||
# 提取最近 3 轮对话
|
||||
@@ -454,6 +458,7 @@ class IntentAnalyzer:
|
||||
for msg in recent_history:
|
||||
role = "用户" if msg.get("role") == "user" else "助手"
|
||||
content = msg.get("content", "")
|
||||
original_content = content # 保留原始内容用于结构化提取
|
||||
|
||||
# 截断过长的内容
|
||||
if len(content) > 500:
|
||||
@@ -461,9 +466,10 @@ class IntentAnalyzer:
|
||||
|
||||
parts.append(f"【{role}】{content}")
|
||||
|
||||
# 提取图片信息
|
||||
# 提取结构化信息（从 assistant 消息中提取章节、表格、来源等）
|
||||
metadata = msg.get("metadata", {})
|
||||
if isinstance(metadata, dict):
|
||||
# 已有的图片提取
|
||||
images = metadata.get("images", [])
|
||||
if images:
|
||||
for img in images[:3]:
|
||||
@@ -472,6 +478,43 @@ class IntentAnalyzer:
|
||||
img_type = img.get("type", "图片")
|
||||
parts.append(f" └─ {img_type}: {desc}")
|
||||
|
||||
# 来源文件提取
|
||||
sources = metadata.get("sources", [])
|
||||
if sources:
|
||||
source_names = []
|
||||
for s in sources[:3]:
|
||||
if isinstance(s, dict):
|
||||
name = s.get("source", "") or s.get("name", "")
|
||||
if name:
|
||||
source_names.append(name)
|
||||
elif isinstance(s, str):
|
||||
source_names.append(s)
|
||||
if source_names:
|
||||
parts.append(f" └─ 来源文件: {', '.join(source_names)}")
|
||||
|
||||
# collections（检索知识库）提取
|
||||
colls = metadata.get("collections", [])
|
||||
if colls:
|
||||
parts.append(f" └─ 检索知识库: {', '.join(colls)}")
|
||||
|
||||
# 从 assistant 原始内容中提取章节路径和表格结构
|
||||
if role == "助手" and original_content:
|
||||
# 提取章节路径：━ xxx ━ 格式
|
||||
sections = re.findall(r'━\s*(.+?)\s*━', original_content)
|
||||
if sections:
|
||||
unique_sections = list(dict.fromkeys(sections)) # 去重保序
|
||||
parts.append(f" └─ 涉及章节: {'; '.join(unique_sections[:3])}")
|
||||
|
||||
# 提取表格列名：| A | B | C | 格式的表头行
|
||||
table_headers = re.findall(r'^\|\s*(.+?)\s*\|', original_content, re.MULTILINE)
|
||||
if table_headers:
|
||||
# 取第一个表格的列名
|
||||
first_header = table_headers[0]
|
||||
cols = [c.strip() for c in first_header.split('|') if c.strip()]
|
||||
# 排除分隔符行（--- 格式）
|
||||
if cols and not all(re.match(r'^[-:]+$', c) for c in cols):
|
||||
parts.append(f" └─ 含表格，列名: {', '.join(cols[:6])}")
|
||||
|
||||
# 添加图片上下文
|
||||
if context_images:
|
||||
parts.append("\n【上下文中的图片】")
|
||||
|
||||
Reference in New Issue
Block a user