90b915232a
第一批（快速修复）： - H6: main.py --debug 默认值 True→False，防止 Werkzeug RCE - M2+M3: /search 增加 validate_query + top_k 范围限制(1-50) - M4: context_count 范围限制(0-10) + 异常捕获 - L3: assert → raise RuntimeError（生产环境 API Key 检查） - H1: SSE 错误事件移除 traceback 字段 第二批（安全加固）： - H2+H3: 文档接口路径遍历 realpath 校验 + 文件类型/大小限制 - H4+H5: 批量上传文件大小检查 - M6: LIKE 查询通配符转义 - M1: 37 处 str(e) 异常信息统一脱敏（6 文件） - M5: CORS 生产环境限制来源 - M7: SESSION_MANAGER None 保护（503） - M11: subprocess 参数注入防护（白名单 + -- 分隔符） 第三批（架构改进）： - M8+M9: 提取 JSON 解析共享工具（extract_json_object/list） - M10: Prompt 注入检测防御（prompt_guard.py） - M12: 解析器文件大小限制（Excel 50MB/TXT 20MB/PDF 100MB） - M13: 全局单例竞态条件双重检查锁定（engine/bm25/intent_analyzer）
145 lines
3.8 KiB
Python
145 lines
3.8 KiB
Python
"""
|
||
BM25 关键词检索索引
|
||
|
||
使用 rank_bm25 + jieba 分词实现中文关键词检索。
|
||
支持索引的序列化/反序列化。
|
||
|
||
使用方式：
|
||
from core.bm25_index import BM25Index
|
||
|
||
index = BM25Index()
|
||
index.add_documents(ids, documents, metadatas)
|
||
results = index.search("查询内容", top_k=5)
|
||
"""
|
||
|
||
import os
|
||
import pickle
|
||
import threading
|
||
import numpy as np
|
||
from rank_bm25 import BM25Okapi
|
||
import jieba
|
||
import logging
|
||
from core.constants import get_empty_result
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class BM25Index:
|
||
"""BM25索引管理器，用于关键词检索"""
|
||
|
||
def __init__(self):
|
||
self.bm25 = None
|
||
self.documents = [] # 原始文档
|
||
self.metadatas = [] # 元数据
|
||
self.ids = [] # 文档ID
|
||
|
||
def tokenize(self, text):
|
||
"""中文分词"""
|
||
return list(jieba.cut(text))
|
||
|
||
def add_documents(self, ids, documents, metadatas):
|
||
"""添加文档到索引"""
|
||
self.ids = ids
|
||
self.documents = documents
|
||
self.metadatas = metadatas
|
||
|
||
# 分词并构建BM25索引
|
||
tokenized_docs = [self.tokenize(doc) for doc in documents]
|
||
self.bm25 = BM25Okapi(tokenized_docs)
|
||
|
||
def search(self, query, top_k=10):
|
||
"""BM25检索"""
|
||
if not self.bm25:
|
||
return get_empty_result()
|
||
|
||
tokenized_query = self.tokenize(query)
|
||
scores = self.bm25.get_scores(tokenized_query)
|
||
|
||
# 获取top_k个结果
|
||
top_indices = np.argsort(scores)[::-1][:top_k]
|
||
|
||
return {
|
||
'ids': [[self.ids[i] for i in top_indices]],
|
||
'documents': [[self.documents[i] for i in top_indices]],
|
||
'metadatas': [[self.metadatas[i] for i in top_indices]],
|
||
'distances': [[float(scores[i]) for i in top_indices]]
|
||
}
|
||
|
||
def save(self, path):
|
||
"""保存索引到文件"""
|
||
data = {
|
||
'ids': self.ids,
|
||
'documents': self.documents,
|
||
'metadatas': self.metadatas
|
||
}
|
||
with open(path, 'wb') as f:
|
||
pickle.dump(data, f)
|
||
logger.info(f"BM25索引已保存: {path}")
|
||
|
||
def load(self, path):
|
||
"""从文件加载索引"""
|
||
if not os.path.exists(path):
|
||
return False
|
||
|
||
with open(path, 'rb') as f:
|
||
data = pickle.load(f)
|
||
|
||
self.ids = data['ids']
|
||
self.documents = data['documents']
|
||
self.metadatas = data['metadatas']
|
||
|
||
# 重建BM25索引
|
||
tokenized_docs = [self.tokenize(doc) for doc in self.documents]
|
||
self.bm25 = BM25Okapi(tokenized_docs)
|
||
|
||
logger.info(f"BM25索引已加载: {len(self.documents)} 个文档")
|
||
return True
|
||
|
||
def clear(self):
|
||
"""清空索引"""
|
||
self.bm25 = None
|
||
self.documents = []
|
||
self.metadatas = []
|
||
self.ids = []
|
||
|
||
|
||
# ==================== 全局 BM25 索引管理器 ====================
|
||
|
||
_bm25_indexer: BM25Index = None
|
||
_bm25_lock = threading.Lock() # 单例创建锁，防止多线程竞态条件
|
||
|
||
|
||
def get_bm25_indexer() -> BM25Index:
|
||
"""
|
||
获取全局 BM25 索引器实例（双重检查锁定，线程安全）
|
||
|
||
Returns:
|
||
BM25Index 实例
|
||
"""
|
||
global _bm25_indexer
|
||
if _bm25_indexer is None:
|
||
with _bm25_lock:
|
||
if _bm25_indexer is None:
|
||
_bm25_indexer = BM25Index()
|
||
return _bm25_indexer
|
||
|
||
|
||
def init_bm25_indexer(ids=None, documents=None, metadatas=None) -> BM25Index:
|
||
"""
|
||
初始化 BM25 索引器并添加文档（加锁，线程安全）
|
||
|
||
Args:
|
||
ids: 文档 ID 列表
|
||
documents: 文档内容列表
|
||
metadatas: 元数据列表
|
||
|
||
Returns:
|
||
初始化后的 BM25Index 实例
|
||
"""
|
||
global _bm25_indexer
|
||
with _bm25_lock:
|
||
_bm25_indexer = BM25Index()
|
||
if ids and documents:
|
||
_bm25_indexer.add_documents(ids, documents, metadatas or [])
|
||
return _bm25_indexer
|