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:
@@ -224,39 +224,36 @@ class RAGCacheManager:
|
|||||||
"""
|
"""
|
||||||
获取查询缓存结果
|
获取查询缓存结果
|
||||||
|
|
||||||
|
始终使用粗粒度 key(基于 kb_version),确保 GET/SET key 一致。
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
query: 查询文本
|
query: 查询文本
|
||||||
kb_name: 知识库名称
|
kb_name: 知识库名称
|
||||||
doc_ids: 相关文档 ID 列表(用于细粒度缓存 key)
|
doc_ids: 保留参数以兼容调用方签名(当前未使用)
|
||||||
"""
|
"""
|
||||||
kb_version = self.get_kb_version(kb_name)
|
kb_version = self.get_kb_version(kb_name)
|
||||||
|
|
||||||
# 计算文档哈希(如果提供了 doc_ids)
|
# 使用粗粒度 key,与 SET 保持一致
|
||||||
doc_hash = ""
|
key = self._make_query_cache_key(query, kb_name, kb_version)
|
||||||
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)
|
|
||||||
return self.query_cache.get(key)
|
return self.query_cache.get(key)
|
||||||
|
|
||||||
def set_query_result(self, query: str, kb_name: str, result: Dict, doc_ids: List[str] = None) -> None:
|
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:
|
Args:
|
||||||
query: 查询文本
|
query: 查询文本
|
||||||
kb_name: 知识库名称
|
kb_name: 知识库名称
|
||||||
result: 缓存结果
|
result: 缓存结果
|
||||||
doc_ids: 相关文档 ID 列表(用于细粒度失效)
|
doc_ids: 保留参数以兼容调用方签名(当前未使用)
|
||||||
"""
|
"""
|
||||||
kb_version = self.get_kb_version(kb_name)
|
kb_version = self.get_kb_version(kb_name)
|
||||||
|
|
||||||
# 计算相关文档的版本哈希(细粒度失效)
|
# 使用与 GET 相同的粗粒度 key,确保缓存可命中
|
||||||
doc_hash = ""
|
key = self._make_query_cache_key(query, kb_name, kb_version)
|
||||||
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)
|
|
||||||
self.query_cache.set(key, result, kb_version=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:
|
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"
|
EMBEDDING_DEVICE = "auto"
|
||||||
RERANK_DEVICE = "auto"
|
RERANK_DEVICE = "auto"
|
||||||
RERANK_USE_ONNX = False
|
RERANK_USE_ONNX = False
|
||||||
RERANK_BACKEND = "local"
|
RERANK_BACKEND = "cloud"
|
||||||
RERANK_CLOUD_MODEL = "qwen3-rerank"
|
RERANK_CLOUD_MODEL = "xop3qwen8breranker"
|
||||||
RERANK_CLOUD_API_KEY = ""
|
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
|
RERANK_CLOUD_TIMEOUT = 15
|
||||||
|
|
||||||
|
|
||||||
@@ -285,8 +285,7 @@ class CloudReranker:
|
|||||||
body = {
|
body = {
|
||||||
"model": self.model,
|
"model": self.model,
|
||||||
"query": query,
|
"query": query,
|
||||||
"documents": documents,
|
"documents": documents
|
||||||
"top_n": len(documents)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
session = self._get_session()
|
session = self._get_session()
|
||||||
@@ -670,6 +669,7 @@ class RAGEngine:
|
|||||||
|
|
||||||
results_list = [vector_results]
|
results_list = [vector_results]
|
||||||
weights = [VECTOR_WEIGHT]
|
weights = [VECTOR_WEIGHT]
|
||||||
|
bm25_results = None # 初始化,防止 NameError
|
||||||
|
|
||||||
if USE_HYBRID_SEARCH and self.bm25_index.bm25:
|
if USE_HYBRID_SEARCH and self.bm25_index.bm25:
|
||||||
bm25_results = self.bm25_index.search(query, top_k=recall_k)
|
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)
|
vector_w, bm25_w = self._get_dynamic_rrf_weights(query)
|
||||||
weights = [vector_w, bm25_w]
|
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:
|
if len(results_list) > 1:
|
||||||
fused_results = self.reciprocal_rank_fusion(results_list, weights)
|
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]})
|
_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)
|
is_enum_query = self._is_enumeration_query(query)
|
||||||
fused_results['_enum_query'] = is_enum_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)
|
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 not (is_enum_query and ENUM_QUERY_DISABLE_TOPK_SHRINK)
|
||||||
and fused_results.get('_score_source') != 'rrf'
|
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)
|
adjusted_k, should_retrieve, reason = self._adaptive_topk.adjust(top_score, top_k)
|
||||||
if "high_confidence" in reason:
|
if "high_confidence" in reason:
|
||||||
# 高置信度时截断结果
|
# 高置信度时截断结果
|
||||||
@@ -1108,7 +1134,7 @@ class RAGEngine:
|
|||||||
'metadatas': [f_metas],
|
'metadatas': [f_metas],
|
||||||
'distances': [f_scores]
|
'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:
|
if key in results:
|
||||||
filtered[key] = results[key]
|
filtered[key] = results[key]
|
||||||
return filtered
|
return filtered
|
||||||
@@ -1131,7 +1157,7 @@ class RAGEngine:
|
|||||||
'metadatas': [f_metas],
|
'metadatas': [f_metas],
|
||||||
'distances': [f_scores]
|
'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:
|
if key in results:
|
||||||
filtered[key] = results[key]
|
filtered[key] = results[key]
|
||||||
return filtered
|
return filtered
|
||||||
@@ -1146,7 +1172,7 @@ class RAGEngine:
|
|||||||
'metadatas': [results['metadatas'][0][:top_k]],
|
'metadatas': [results['metadatas'][0][:top_k]],
|
||||||
'distances': [results['distances'][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:
|
if key in results:
|
||||||
truncated[key] = results[key]
|
truncated[key] = results[key]
|
||||||
return truncated
|
return truncated
|
||||||
@@ -1181,166 +1207,6 @@ class RAGEngine:
|
|||||||
return None
|
return None
|
||||||
return self.collection
|
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
|
@staticmethod
|
||||||
def _normalize_section_path(section_path: str, levels: int = None) -> str:
|
def _normalize_section_path(section_path: str, levels: int = None) -> str:
|
||||||
"""归一化 section_path:取前 N 级路径,容忍 MinerU 标题检测误差。
|
"""归一化 section_path:取前 N 级路径,容忍 MinerU 标题检测误差。
|
||||||
@@ -1465,6 +1331,167 @@ class RAGEngine:
|
|||||||
matched = sum(1 for w in bigrams if w in chunk_text)
|
matched = sum(1 for w in bigrams if w in chunk_text)
|
||||||
return matched / len(bigrams)
|
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(
|
def _search_with_sub_queries(
|
||||||
self, original_query, sub_queries, top_k=5, allowed_levels=None,
|
self, original_query, sub_queries, top_k=5, allowed_levels=None,
|
||||||
role=None, department=None, collections=None, source_filter=None
|
role=None, department=None, collections=None, source_filter=None
|
||||||
@@ -1487,6 +1514,7 @@ class RAGEngine:
|
|||||||
sub_top_k = max(top_k, 5)
|
sub_top_k = max(top_k, 5)
|
||||||
|
|
||||||
all_results = []
|
all_results = []
|
||||||
|
_all_bm25_top3 = [] # 收集各子查询的 BM25 top3
|
||||||
for sub_q in sub_queries:
|
for sub_q in sub_queries:
|
||||||
try:
|
try:
|
||||||
sub_result = self.search_knowledge(
|
sub_result = self.search_knowledge(
|
||||||
@@ -1497,6 +1525,9 @@ class RAGEngine:
|
|||||||
)
|
)
|
||||||
if sub_result and sub_result.get('ids') and sub_result['ids'][0]:
|
if sub_result and sub_result.get('ids') and sub_result['ids'][0]:
|
||||||
all_results.append(sub_result)
|
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:
|
except Exception as e:
|
||||||
logger.warning(f"子查询检索失败: '{sub_q}' - {e}")
|
logger.warning(f"子查询检索失败: '{sub_q}' - {e}")
|
||||||
|
|
||||||
@@ -1505,9 +1536,19 @@ class RAGEngine:
|
|||||||
|
|
||||||
# 合并去重
|
# 合并去重
|
||||||
if len(all_results) == 1:
|
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(
|
def _search_with_decomposition(
|
||||||
self, query, decomposer, top_k=5, allowed_levels=None,
|
self, query, decomposer, top_k=5, allowed_levels=None,
|
||||||
@@ -1539,6 +1580,7 @@ class RAGEngine:
|
|||||||
|
|
||||||
# 并行检索各子查询
|
# 并行检索各子查询
|
||||||
all_results = []
|
all_results = []
|
||||||
|
_all_bm25_top3 = [] # 收集各子查询的 BM25 top3
|
||||||
for sub_q in sub_queries:
|
for sub_q in sub_queries:
|
||||||
try:
|
try:
|
||||||
sub_result = self.search_knowledge(
|
sub_result = self.search_knowledge(
|
||||||
@@ -1549,6 +1591,9 @@ class RAGEngine:
|
|||||||
)
|
)
|
||||||
if sub_result and sub_result.get('ids') and sub_result['ids'][0]:
|
if sub_result and sub_result.get('ids') and sub_result['ids'][0]:
|
||||||
all_results.append(sub_result)
|
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:
|
except Exception as e:
|
||||||
logger.warning(f"子查询检索失败: '{sub_q}' - {e}")
|
logger.warning(f"子查询检索失败: '{sub_q}' - {e}")
|
||||||
|
|
||||||
@@ -1561,6 +1606,14 @@ class RAGEngine:
|
|||||||
else:
|
else:
|
||||||
merged = self._merge_and_deduplicate(all_results, top_k)
|
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
|
return merged
|
||||||
|
|
||||||
def _merge_and_deduplicate(self, results_list, top_k):
|
def _merge_and_deduplicate(self, results_list, top_k):
|
||||||
@@ -1645,12 +1698,13 @@ class RAGEngine:
|
|||||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||||
|
|
||||||
def _query_single_collection(coll_name):
|
def _query_single_collection(coll_name):
|
||||||
"""查询单个向量库(向量 + BM25)"""
|
"""查询单个向量库(向量 + BM25),返回 (coll_results, bm25_raw_items)"""
|
||||||
coll_results = []
|
coll_results = []
|
||||||
|
bm25_raw_items = [] # 该 collection 的 BM25 原始结果
|
||||||
try:
|
try:
|
||||||
coll = self.kb_manager.get_collection(coll_name)
|
coll = self.kb_manager.get_collection(coll_name)
|
||||||
if not coll:
|
if not coll:
|
||||||
return coll_results
|
return coll_results, bm25_raw_items
|
||||||
|
|
||||||
query_kwargs = {
|
query_kwargs = {
|
||||||
"query_embeddings": [query_vector],
|
"query_embeddings": [query_vector],
|
||||||
@@ -1668,25 +1722,61 @@ class RAGEngine:
|
|||||||
if USE_HYBRID_SEARCH:
|
if USE_HYBRID_SEARCH:
|
||||||
try:
|
try:
|
||||||
bm25 = self.kb_manager.get_bm25_index(coll_name)
|
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)
|
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)
|
bm25_res = self._filter_results(bm25_res, lambda meta: meta.get('source') == source_filter)
|
||||||
if bm25_res['metadatas'] and bm25_res['metadatas'][0]:
|
if bm25_res['metadatas'] and bm25_res['metadatas'][0]:
|
||||||
for meta in bm25_res['metadatas'][0]:
|
for meta in bm25_res['metadatas'][0]:
|
||||||
meta['_collection'] = coll_name
|
meta['_collection'] = coll_name
|
||||||
coll_results.append(bm25_res)
|
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:
|
except Exception as e:
|
||||||
logger.debug(f"向量库 {coll_name} 检索失败: {e}")
|
logger.debug(f"向量库 {coll_name} BM25检索失败: {e}")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.debug(f"多向量库检索失败: {e}")
|
logger.debug(f"多向量库检索失败: {e}")
|
||||||
return coll_results
|
return coll_results, bm25_raw_items
|
||||||
|
|
||||||
all_results = []
|
all_results = []
|
||||||
|
_bm25_raw_top3 = []
|
||||||
with ThreadPoolExecutor(max_workers=len(target_collections)) as executor:
|
with ThreadPoolExecutor(max_workers=len(target_collections)) as executor:
|
||||||
futures = {executor.submit(_query_single_collection, name): name for name in target_collections}
|
futures = {executor.submit(_query_single_collection, name): name for name in target_collections}
|
||||||
for future in as_completed(futures):
|
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 检索 ==========
|
||||||
faq_results = self._search_faq_collection(query_vector, top_k=FAQ_RECALL_TOP_K)
|
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)
|
is_enum_query = self._is_enumeration_query(query)
|
||||||
fused_results['_enum_query'] = is_enum_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)
|
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 not (is_enum_query and ENUM_QUERY_DISABLE_TOPK_SHRINK)
|
||||||
and fused_results.get('_score_source') != 'rrf'
|
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)
|
adjusted_k, should_retrieve, reason = self._adaptive_topk.adjust(top_score, top_k)
|
||||||
if "high_confidence" in reason:
|
if "high_confidence" in reason:
|
||||||
# 高置信度时截断结果
|
# 高置信度时截断结果
|
||||||
@@ -1851,7 +1951,7 @@ class RAGEngine:
|
|||||||
'metadatas': [filtered_metas],
|
'metadatas': [filtered_metas],
|
||||||
'distances': [filtered_distances]
|
'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:
|
if key in results:
|
||||||
filtered[key] = results[key]
|
filtered[key] = results[key]
|
||||||
return filtered
|
return filtered
|
||||||
@@ -1924,7 +2024,7 @@ class RAGEngine:
|
|||||||
'metadatas': [filtered_metas],
|
'metadatas': [filtered_metas],
|
||||||
'distances': [filtered_distances]
|
'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:
|
if key in results:
|
||||||
filtered[key] = results[key]
|
filtered[key] = results[key]
|
||||||
return filtered
|
return filtered
|
||||||
@@ -2040,7 +2140,7 @@ class RAGEngine:
|
|||||||
'metadatas': [[c['metadata'] 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]]
|
'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:
|
if key in results:
|
||||||
filtered[key] = results[key]
|
filtered[key] = results[key]
|
||||||
return filtered
|
return filtered
|
||||||
@@ -2080,7 +2180,7 @@ class RAGEngine:
|
|||||||
'metadatas': [[c['metadata'] for c in selected]],
|
'metadatas': [[c['metadata'] for c in selected]],
|
||||||
'distances': [[id_to_dist.get(c['id'], 0) 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:
|
if key in results:
|
||||||
filtered[key] = results[key]
|
filtered[key] = results[key]
|
||||||
return filtered
|
return filtered
|
||||||
@@ -2189,109 +2289,15 @@ class RAGEngine:
|
|||||||
'_rerank_cached': cache_hit
|
'_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:
|
if key in results:
|
||||||
reranked[key] = results[key]
|
reranked[key] = results[key]
|
||||||
|
# Rerank 后 distances 语义变为 CrossEncoder 分数,更新 _score_source
|
||||||
|
# 使自适应 TopK 能正确应用(之前 _score_source='rrf' 会导致自适应 TopK 被跳过)
|
||||||
|
reranked['_score_source'] = 'rerank'
|
||||||
return reranked
|
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):
|
def generate_answer_stream(self, query, context, history=None):
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -316,7 +316,8 @@ class IntentAnalyzer:
|
|||||||
|
|
||||||
if cache_emb is not None:
|
if cache_emb is not None:
|
||||||
cached = cache.get(cache_emb)
|
cached = cache.get(cache_emb)
|
||||||
if cached:
|
# 确保缓存条目是意图分析结果(非 RAG 回答缓存)
|
||||||
|
if cached and cached.get("cache_type") != "rag_answer":
|
||||||
logger.info(f"意图分析缓存命中: {cached.get('reason', '')[:50]}")
|
logger.info(f"意图分析缓存命中: {cached.get('reason', '')[:50]}")
|
||||||
return IntentAnalysis.from_dict(cached)
|
return IntentAnalysis.from_dict(cached)
|
||||||
else:
|
else:
|
||||||
@@ -386,9 +387,11 @@ class IntentAnalyzer:
|
|||||||
intent=intent_type
|
intent=intent_type
|
||||||
)
|
)
|
||||||
|
|
||||||
# 存入语义缓存
|
# 存入语义缓存(标记类型,避免与 RAG 回答缓存混淆)
|
||||||
if cache and cache_emb is not None:
|
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:
|
if len(self._exact_cache) < self._exact_cache_max:
|
||||||
@@ -446,6 +449,7 @@ class IntentAnalyzer:
|
|||||||
if not history:
|
if not history:
|
||||||
return "(无历史对话)"
|
return "(无历史对话)"
|
||||||
|
|
||||||
|
import re
|
||||||
parts = []
|
parts = []
|
||||||
|
|
||||||
# 提取最近 3 轮对话
|
# 提取最近 3 轮对话
|
||||||
@@ -454,6 +458,7 @@ class IntentAnalyzer:
|
|||||||
for msg in recent_history:
|
for msg in recent_history:
|
||||||
role = "用户" if msg.get("role") == "user" else "助手"
|
role = "用户" if msg.get("role") == "user" else "助手"
|
||||||
content = msg.get("content", "")
|
content = msg.get("content", "")
|
||||||
|
original_content = content # 保留原始内容用于结构化提取
|
||||||
|
|
||||||
# 截断过长的内容
|
# 截断过长的内容
|
||||||
if len(content) > 500:
|
if len(content) > 500:
|
||||||
@@ -461,9 +466,10 @@ class IntentAnalyzer:
|
|||||||
|
|
||||||
parts.append(f"【{role}】{content}")
|
parts.append(f"【{role}】{content}")
|
||||||
|
|
||||||
# 提取图片信息
|
# 提取结构化信息(从 assistant 消息中提取章节、表格、来源等)
|
||||||
metadata = msg.get("metadata", {})
|
metadata = msg.get("metadata", {})
|
||||||
if isinstance(metadata, dict):
|
if isinstance(metadata, dict):
|
||||||
|
# 已有的图片提取
|
||||||
images = metadata.get("images", [])
|
images = metadata.get("images", [])
|
||||||
if images:
|
if images:
|
||||||
for img in images[:3]:
|
for img in images[:3]:
|
||||||
@@ -472,6 +478,43 @@ class IntentAnalyzer:
|
|||||||
img_type = img.get("type", "图片")
|
img_type = img.get("type", "图片")
|
||||||
parts.append(f" └─ {img_type}: {desc}")
|
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:
|
if context_images:
|
||||||
parts.append("\n【上下文中的图片】")
|
parts.append("\n【上下文中的图片】")
|
||||||
|
|||||||
Reference in New Issue
Block a user