437 lines
16 KiB
Python
437 lines
16 KiB
Python
"""
|
||
知识库管理器 - 集合管理 Mixin
|
||
|
||
提供向量库（ChromaDB Collection）的创建、删除、查询等管理功能。
|
||
每个向量库对应一个独立的 ChromaDB 集合，支持按部门隔离。
|
||
|
||
主要方法：
|
||
- get_collection: 获取或创建向量库集合
|
||
- create_collection: 创建新向量库
|
||
- delete_collection: 删除向量库
|
||
- list_collections: 列出所有向量库
|
||
- collection_exists: 检查向量库是否存在
|
||
"""
|
||
|
||
import os
|
||
import gc
|
||
import time
|
||
import logging
|
||
from typing import List, Optional, Tuple
|
||
from datetime import datetime
|
||
|
||
import chromadb
|
||
from chromadb import Collection
|
||
|
||
from .base import (
|
||
CollectionInfo, PUBLIC_KB_NAME
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class CollectionMixin:
|
||
"""
|
||
集合管理 Mixin
|
||
|
||
提供 ChromaDB 向量库集合的管理功能，包括：
|
||
- 集合的创建与删除
|
||
- 集合元数据管理
|
||
- 集合列表查询
|
||
|
||
依赖属性（需由主类提供）：
|
||
- self.base_path: 向量库存储根路径
|
||
- self.bm25_base_path: BM25 索引存储路径
|
||
- self._collections: 集合缓存字典
|
||
- self._clients: 客户端缓存字典
|
||
- self._bm25_indexes: BM25 索引缓存字典
|
||
- self._metadata: 元数据字典
|
||
- self._lock: 线程锁
|
||
"""
|
||
|
||
def _get_client(self, kb_name: str) -> chromadb.PersistentClient:
|
||
"""
|
||
获取或创建 ChromaDB 客户端
|
||
|
||
每个向量库使用独立的客户端实例，客户端会被缓存以避免重复创建。
|
||
|
||
Args:
|
||
kb_name: 向量库名称
|
||
|
||
Returns:
|
||
ChromaDB PersistentClient 实例
|
||
"""
|
||
if kb_name not in self._clients:
|
||
db_path = os.path.join(self.base_path, kb_name)
|
||
os.makedirs(db_path, exist_ok=True)
|
||
self._clients[kb_name] = chromadb.PersistentClient(path=db_path)
|
||
return self._clients[kb_name]
|
||
|
||
def get_collection(self, kb_name: str) -> Optional[Collection]:
|
||
"""
|
||
获取或创建向量库集合
|
||
|
||
如果集合已存在于缓存中，直接返回缓存实例。
|
||
否则从 ChromaDB 获取或创建新集合。
|
||
|
||
Args:
|
||
kb_name: 向量库名称（只能包含字母、数字和下划线）
|
||
|
||
Returns:
|
||
ChromaDB Collection 对象；失败时返回 None
|
||
|
||
Note:
|
||
使用 cosine 相似度作为向量距离度量。
|
||
sync_threshold 设置为 100000 以支持大批量写入。
|
||
"""
|
||
with self._lock:
|
||
if kb_name in self._collections:
|
||
return self._collections[kb_name]
|
||
try:
|
||
client = self._get_client(kb_name)
|
||
collection = client.get_or_create_collection(
|
||
name=kb_name,
|
||
metadata={
|
||
"hnsw:space": "cosine",
|
||
"hnsw:sync_threshold": 100000
|
||
}
|
||
)
|
||
self._collections[kb_name] = collection
|
||
logger.info(f"获取向量库: {kb_name}, 文档数: {collection.count()}")
|
||
return collection
|
||
except Exception as e:
|
||
logger.error(f"获取向量库失败: {kb_name}, 错误: {e}")
|
||
return None
|
||
|
||
def create_collection(
|
||
self,
|
||
kb_name: str,
|
||
display_name: str = "",
|
||
department: str = "",
|
||
description: str = ""
|
||
) -> Tuple[bool, str]:
|
||
"""
|
||
创建新向量库
|
||
|
||
创建一个新的 ChromaDB 集合，同时初始化对应的 BM25 索引。
|
||
|
||
Args:
|
||
kb_name: 向量库名称（只能包含字母、数字和下划线）
|
||
display_name: 显示名称（用于 UI 展示）
|
||
department: 所属部门（用于权限控制）
|
||
description: 向量库描述
|
||
|
||
Returns:
|
||
元组 (success, message)：
|
||
- success: 创建是否成功
|
||
- message: 成功或失败的描述信息
|
||
|
||
Example:
|
||
>>> success, msg = kb_manager.create_collection(
|
||
... "dept_finance",
|
||
... display_name="财务部知识库",
|
||
... department="finance"
|
||
... )
|
||
"""
|
||
from .base import BM25Index
|
||
|
||
# 从磁盘重新加载元数据，确保多 worker 进程间状态一致
|
||
self._metadata = self._load_metadata()
|
||
|
||
if not kb_name or not kb_name.replace('_', '').isalnum():
|
||
return False, "向量库名称只能包含字母、数字和下划线"
|
||
|
||
if kb_name in self._metadata.get("collections", {}):
|
||
return False, f"向量库 '{kb_name}' 已存在"
|
||
|
||
chroma_dir = os.path.join(self.base_path, kb_name)
|
||
if os.path.exists(chroma_dir):
|
||
import shutil
|
||
shutil.rmtree(chroma_dir)
|
||
logger.warning(f"清理残留向量库文件夹: {chroma_dir}")
|
||
|
||
try:
|
||
collection = self.get_collection(kb_name)
|
||
if not collection:
|
||
return False, "创建向量库失败"
|
||
|
||
self._bm25_indexes[kb_name] = BM25Index()
|
||
|
||
if "collections" not in self._metadata:
|
||
self._metadata["collections"] = {}
|
||
|
||
self._metadata["collections"][kb_name] = {
|
||
"display_name": display_name or kb_name,
|
||
"department": department,
|
||
"description": description,
|
||
"created_at": datetime.now().isoformat()
|
||
}
|
||
self._save_metadata()
|
||
|
||
logger.info(f"创建向量库: {kb_name}")
|
||
return True, f"向量库 '{kb_name}' 创建成功"
|
||
|
||
except Exception as e:
|
||
logger.error(f"创建向量库失败: {e}")
|
||
return False, f"创建失败: {str(e)}"
|
||
|
||
def update_collection_metadata(
|
||
self,
|
||
kb_name: str,
|
||
display_name: str = None,
|
||
description: str = None
|
||
) -> bool:
|
||
"""
|
||
更新向量库元数据
|
||
|
||
仅更新 display_name 和 description，不影响集合中的数据。
|
||
|
||
Args:
|
||
kb_name: 向量库名称
|
||
display_name: 新的显示名称（None 表示不更新）
|
||
description: 新的描述（None 表示不更新）
|
||
|
||
Returns:
|
||
更新成功返回 True，向量库不存在返回 False
|
||
"""
|
||
# 从磁盘重新加载元数据，确保多 worker 进程间状态一致
|
||
self._metadata = self._load_metadata()
|
||
|
||
collections = self._metadata.get("collections", {})
|
||
if kb_name not in collections:
|
||
return False
|
||
|
||
if display_name is not None:
|
||
collections[kb_name]["display_name"] = display_name
|
||
if description is not None:
|
||
collections[kb_name]["description"] = description
|
||
|
||
self._save_metadata()
|
||
logger.info(f"更新向量库元数据: {kb_name}")
|
||
return True
|
||
|
||
def delete_collection(self, kb_name: str, delete_documents: bool = False) -> Tuple[bool, str]:
|
||
"""
|
||
删除向量库
|
||
|
||
删除向量库及其所有数据，包括：
|
||
- ChromaDB 集合
|
||
- BM25 索引文件
|
||
- 向量库文件夹
|
||
- 可选：原始文档文件
|
||
|
||
Args:
|
||
kb_name: 向量库名称
|
||
delete_documents: 是否同时删除原始文档文件
|
||
|
||
Returns:
|
||
元组 (success, message)：
|
||
- success: 删除是否成功
|
||
- message: 成功或失败的描述信息
|
||
|
||
Warning:
|
||
公开知识库 (public_kb) 不能被删除。
|
||
删除操作不可逆，请谨慎使用。
|
||
"""
|
||
import shutil
|
||
|
||
# 从磁盘重新加载元数据，确保多 worker 进程间状态一致
|
||
self._metadata = self._load_metadata()
|
||
|
||
if kb_name == PUBLIC_KB_NAME:
|
||
return False, "公开知识库不能删除"
|
||
|
||
if kb_name not in self._metadata.get("collections", {}):
|
||
return False, f"向量库 '{kb_name}' 不存在"
|
||
|
||
try:
|
||
client = self._clients.get(kb_name)
|
||
if client:
|
||
try:
|
||
client.delete_collection(kb_name)
|
||
except Exception as e:
|
||
logger.debug(f"删除集合失败: {e}")
|
||
try:
|
||
client.close()
|
||
except Exception as e:
|
||
logger.debug(f"关闭客户端失败: {e}")
|
||
|
||
if kb_name in self._collections:
|
||
del self._collections[kb_name]
|
||
if kb_name in self._bm25_indexes:
|
||
del self._bm25_indexes[kb_name]
|
||
if kb_name in self._clients:
|
||
del self._clients[kb_name]
|
||
|
||
gc.collect()
|
||
time.sleep(1)
|
||
|
||
bm25_path = os.path.join(self.bm25_base_path, f"{kb_name}.pkl")
|
||
if os.path.exists(bm25_path):
|
||
os.remove(bm25_path)
|
||
|
||
chroma_dir = os.path.join(self.base_path, kb_name)
|
||
if os.path.exists(chroma_dir):
|
||
for attempt in range(3):
|
||
try:
|
||
shutil.rmtree(chroma_dir)
|
||
logger.info(f"删除向量库文件夹: {chroma_dir}")
|
||
break
|
||
except PermissionError:
|
||
if attempt < 2:
|
||
logger.warning(f"删除文件夹失败，重试 {attempt + 1}/3")
|
||
gc.collect()
|
||
time.sleep(1)
|
||
else:
|
||
logger.error(f"删除文件夹失败，文件被占用: {chroma_dir}")
|
||
raise
|
||
|
||
if delete_documents:
|
||
from config import DOCUMENTS_PATH
|
||
docs_dir = os.path.join(DOCUMENTS_PATH, kb_name)
|
||
if os.path.exists(docs_dir):
|
||
shutil.rmtree(docs_dir)
|
||
logger.info(f"删除文档文件夹: {docs_dir}")
|
||
|
||
# 删除该向量库所有文档的哈希记录
|
||
try:
|
||
from knowledge.sync import SyncDatabase
|
||
sync_db = SyncDatabase()
|
||
# 获取所有以 kb_name/ 开头的哈希记录
|
||
all_hashes = sync_db.get_all_document_hashes()
|
||
deleted_count = 0
|
||
for doc_id in list(all_hashes.keys()):
|
||
if doc_id.startswith(f"{kb_name}/"):
|
||
sync_db.delete_document_hash(doc_id)
|
||
deleted_count += 1
|
||
if deleted_count > 0:
|
||
logger.info(f"清理向量库哈希记录: {kb_name}, 共 {deleted_count} 条")
|
||
except Exception as e:
|
||
logger.warning(f"清理哈希记录失败: {e}")
|
||
|
||
# 清理该向量库的 SQLite 版本记录和变更日志
|
||
try:
|
||
from data.db import get_connection
|
||
with get_connection("knowledge") as conn:
|
||
cursor1 = conn.execute(
|
||
"DELETE FROM document_versions WHERE collection = ?",
|
||
(kb_name,)
|
||
)
|
||
cursor2 = conn.execute(
|
||
"DELETE FROM version_change_logs WHERE collection = ?",
|
||
(kb_name,)
|
||
)
|
||
conn.commit()
|
||
ver_cleaned = cursor1.rowcount
|
||
log_cleaned = cursor2.rowcount
|
||
if ver_cleaned > 0 or log_cleaned > 0:
|
||
logger.info(
|
||
f"清理向量库版本记录: {kb_name}, "
|
||
f"版本 {ver_cleaned} 条, 日志 {log_cleaned} 条"
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"清理版本记录失败: {e}")
|
||
|
||
if kb_name in self._metadata.get("collections", {}):
|
||
del self._metadata["collections"][kb_name]
|
||
self._save_metadata()
|
||
|
||
logger.info(f"删除向量库: {kb_name}")
|
||
return True, f"向量库 '{kb_name}' 已删除"
|
||
|
||
except Exception as e:
|
||
logger.error(f"删除向量库失败: {e}")
|
||
return False, f"删除失败: {str(e)}"
|
||
|
||
def list_collections(self) -> List[CollectionInfo]:
|
||
"""
|
||
列出所有向量库
|
||
|
||
返回所有已创建的向量库信息，包括名称、文档数量、创建时间等。
|
||
|
||
Returns:
|
||
CollectionInfo 对象列表，每个对象包含：
|
||
- name: 向量库名称
|
||
- display_name: 显示名称
|
||
- document_count: 文档数量
|
||
- created_at: 创建时间
|
||
- department: 所属部门
|
||
- description: 描述
|
||
"""
|
||
# 从磁盘重新加载元数据，确保多 worker 进程间状态一致
|
||
self._metadata = self._load_metadata()
|
||
result = []
|
||
|
||
# 扫描 base_path 下的所有子目录作为向量库
|
||
# 每个向量库使用独立目录：base_path/my_ky, base_path/public_kb 等
|
||
actual_collections = []
|
||
try:
|
||
if os.path.exists(self.base_path):
|
||
for item in os.listdir(self.base_path):
|
||
item_path = os.path.join(self.base_path, item)
|
||
if os.path.isdir(item_path) and not item.startswith('.'):
|
||
# 检查是否包含 chroma.sqlite3（有效的向量库目录）
|
||
if os.path.exists(os.path.join(item_path, 'chroma.sqlite3')):
|
||
actual_collections.append(item)
|
||
except Exception as e:
|
||
logger.warning(f"扫描向量库目录失败: {e}")
|
||
|
||
# 如果扫描失败，回退到元数据中的集合列表
|
||
if not actual_collections:
|
||
actual_collections = list(self._metadata.get("collections", {}).keys())
|
||
|
||
for name in actual_collections:
|
||
if name not in self._metadata.get("collections", {}):
|
||
if "collections" not in self._metadata:
|
||
self._metadata["collections"] = {}
|
||
self._metadata["collections"][name] = {
|
||
"display_name": name,
|
||
"department": "",
|
||
"description": "",
|
||
"created_at": datetime.now().isoformat()
|
||
}
|
||
logger.info(f"自动补充向量库元数据: {name}")
|
||
|
||
self._save_metadata()
|
||
|
||
stale_collections = []
|
||
|
||
for name, info in self._metadata.get("collections", {}).items():
|
||
try:
|
||
collection = self.get_collection(name)
|
||
result.append(CollectionInfo(
|
||
name=name,
|
||
display_name=info.get("display_name", name),
|
||
document_count=collection.count() if collection else 0,
|
||
created_at=info.get("created_at", ""),
|
||
department=info.get("department", ""),
|
||
description=info.get("description", "")
|
||
))
|
||
except Exception as e:
|
||
logger.warning(
|
||
f"跳过异常向量库 '{name}': {e}，可能是 ChromaDB 数据目录已丢失"
|
||
)
|
||
stale_collections.append(name)
|
||
|
||
# 清理元数据中指向已失效集合的条目
|
||
if stale_collections:
|
||
for name in stale_collections:
|
||
self._metadata.get("collections", {}).pop(name, None)
|
||
logger.info(f"清理失效向量库元数据: {name}")
|
||
self._save_metadata()
|
||
|
||
return result
|
||
|
||
def collection_exists(self, kb_name: str) -> bool:
|
||
"""
|
||
检查向量库是否存在
|
||
|
||
Args:
|
||
kb_name: 向量库名称
|
||
|
||
Returns:
|
||
存在返回 True，不存在返回 False
|
||
"""
|
||
# 从磁盘重新加载元数据，确保多 worker 进程间状态一致
|
||
self._metadata = self._load_metadata()
|
||
return kb_name in self._metadata.get("collections", {})
|