init: RAG 知识库服务初始提交
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
This commit is contained in:
lacerate551
385
knowledge/collection.py
Normal file
385
knowledge/collection.py
Normal file
@@ -0,0 +1,385 @@
|
||||
"""
|
||||
知识库管理器 - 集合管理 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
|
||||
|
||||
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
|
||||
"""
|
||||
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
|
||||
|
||||
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}")
|
||||
|
||||
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: 描述
|
||||
"""
|
||||
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()
|
||||
|
||||
for name, info in self._metadata.get("collections", {}).items():
|
||||
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", "")
|
||||
))
|
||||
|
||||
return result
|
||||
|
||||
def collection_exists(self, kb_name: str) -> bool:
|
||||
"""
|
||||
检查向量库是否存在
|
||||
|
||||
Args:
|
||||
kb_name: 向量库名称
|
||||
|
||||
Returns:
|
||||
存在返回 True，不存在返回 False
|
||||
"""
|
||||
return kb_name in self._metadata.get("collections", {})
|
||||
Reference in New Issue
Block a user