100d1a06eb
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
126 lines
3.7 KiB
Python
126 lines
3.7 KiB
Python
"""
|
||
知识库管理器 - BM25索引管理 Mixin
|
||
|
||
提供 BM25 关键词检索索引的管理功能，包括索引的加载、保存和重建。
|
||
|
||
BM25 是一种基于概率检索模型的排序函数，适合中文关键词检索。
|
||
每个向量库对应一个独立的 BM25 索引，使用 jieba 进行中文分词。
|
||
|
||
主要方法：
|
||
- get_bm25_index: 获取或加载 BM25 索引
|
||
- save_bm25_index: 保存 BM25 索引到磁盘
|
||
- rebuild_bm25_index: 从向量库重建 BM25 索引
|
||
"""
|
||
|
||
import os
|
||
import logging
|
||
from typing import Optional
|
||
|
||
from .base import BM25Index
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class IndexMixin:
|
||
"""
|
||
BM25索引管理 Mixin
|
||
|
||
提供 BM25 索引的生命周期管理，支持：
|
||
- 索引的懒加载（首次访问时从磁盘加载）
|
||
- 索引的持久化（保存到 .pkl 文件）
|
||
- 索引的重建（从向量库同步）
|
||
|
||
依赖属性（需由主类提供）：
|
||
- self.bm25_base_path: BM25 索引存储路径
|
||
- self._bm25_indexes: BM25 索引缓存字典
|
||
"""
|
||
|
||
def get_bm25_index(self, kb_name: str) -> BM25Index:
|
||
"""
|
||
获取或加载 BM25 索引
|
||
|
||
如果索引已在内存中，直接返回。
|
||
否则尝试从磁盘加载，加载失败则创建空索引。
|
||
|
||
Args:
|
||
kb_name: 向量库名称
|
||
|
||
Returns:
|
||
BM25Index 实例
|
||
|
||
Note:
|
||
索引文件路径: {bm25_base_path}/{kb_name}.pkl
|
||
"""
|
||
if kb_name not in self._bm25_indexes:
|
||
self._bm25_indexes[kb_name] = BM25Index()
|
||
bm25_path = os.path.join(self.bm25_base_path, f"{kb_name}.pkl")
|
||
if os.path.exists(bm25_path):
|
||
self._bm25_indexes[kb_name].load(bm25_path)
|
||
return self._bm25_indexes[kb_name]
|
||
|
||
def save_bm25_index(self, kb_name: str) -> None:
|
||
"""
|
||
保存 BM25 索引到磁盘
|
||
|
||
将指定向量库的 BM25 索引序列化保存到 .pkl 文件。
|
||
|
||
Args:
|
||
kb_name: 向量库名称
|
||
"""
|
||
if kb_name in self._bm25_indexes:
|
||
bm25_path = os.path.join(self.bm25_base_path, f"{kb_name}.pkl")
|
||
self._bm25_indexes[kb_name].save(bm25_path)
|
||
logger.info(f"保存 BM25 索引: {kb_name}")
|
||
|
||
def rebuild_bm25_index(self, kb_name: str) -> bool:
|
||
"""
|
||
重建 BM25 索引
|
||
|
||
从 ChromaDB 向量库中读取所有文档，重新构建 BM25 索引。
|
||
用于文档变更后同步索引状态。
|
||
|
||
Args:
|
||
kb_name: 向量库名称
|
||
|
||
Returns:
|
||
重建成功返回 True，失败返回 False
|
||
|
||
Warning:
|
||
大量文档时重建可能耗时较长。
|
||
"""
|
||
try:
|
||
collection = self.get_collection(kb_name)
|
||
if not collection:
|
||
return False
|
||
|
||
result = collection.get()
|
||
|
||
bm25_index = BM25Index()
|
||
if result['ids']:
|
||
bm25_index.add_documents(
|
||
ids=result['ids'],
|
||
documents=result['documents'],
|
||
metadatas=result['metadatas']
|
||
)
|
||
|
||
self._bm25_indexes[kb_name] = bm25_index
|
||
self.save_bm25_index(kb_name)
|
||
|
||
logger.info(f"重建 BM25 索引: {kb_name}, 文档数: {len(result['ids'])}")
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"重建 BM25 索引失败: {e}")
|
||
return False
|
||
|
||
def _rebuild_bm25_index(self, kb_name: str) -> None:
|
||
"""
|
||
重建指定向量库的 BM25 索引（内部方法）
|
||
|
||
兼容旧代码的内部方法，直接调用 rebuild_bm25_index。
|
||
|
||
Args:
|
||
kb_name: 向量库名称
|
||
"""
|
||
self.rebuild_bm25_index(kb_name)
|