feat: 完善重新分块功能,优化代码逻辑,保存请求参数

This commit is contained in:
supreme0597 2025-11-12 23:14:56 +08:00 committed by Wenjie Zhang
parent b7ba646fae
commit 1700628624
6 changed files with 101 additions and 79 deletions

View File

@ -293,6 +293,7 @@ class KnowledgeBase(ABC):
"type": file_info.get("file_type", ""), "type": file_info.get("file_type", ""),
"status": file_info.get("status", "done"), "status": file_info.get("status", "done"),
"created_at": created_at, "created_at": created_at,
"processing_params": file_info.get("processing_params", None),
} }
# 按创建时间倒序排序文件列表 # 按创建时间倒序排序文件列表

View File

@ -47,10 +47,6 @@ class ChromaKB(KnowledgeBase):
# 存储集合映射 {db_id: collection} # 存储集合映射 {db_id: collection}
self.collections: dict[str, Any] = {} self.collections: dict[str, Any] = {}
# 元数据锁
self._metadata_lock = asyncio.Lock()
logger.info("ChromaKB initialized") logger.info("ChromaKB initialized")
@property @property
@ -290,6 +286,7 @@ class ChromaKB(KnowledgeBase):
try: try:
# 更新状态为处理中 # 更新状态为处理中
self.files_meta[file_id]["processing_params"] = params.copy()
self.files_meta[file_id]["status"] = "processing" self.files_meta[file_id]["status"] = "processing"
self._save_metadata() self._save_metadata()
@ -361,24 +358,6 @@ class ChromaKB(KnowledgeBase):
return processed_items_info return processed_items_info
async def delete_file_chunks_only(self, db_id: str, file_id: str) -> None:
"""仅删除文件的chunks数据保留元数据用于更新操作"""
collection = await self._get_chroma_collection(db_id)
if collection:
try:
# 查找所有相关的chunks
results = collection.get(where={"full_doc_id": file_id}, include=["metadatas"])
# 删除所有相关chunks
if results and results.get("ids"):
collection.delete(ids=results["ids"])
logger.info(f"Deleted {len(results['ids'])} chunks for file {file_id}")
except Exception as e:
logger.error(f"Error deleting file {file_id} from ChromaDB: {e}")
# 注意:这里不删除 files_meta[file_id],保留元数据用于后续操作
async def aquery(self, query_text: str, db_id: str, **kwargs) -> list[dict]: async def aquery(self, query_text: str, db_id: str, **kwargs) -> list[dict]:
"""异步查询知识库""" """异步查询知识库"""
collection = await self._get_chroma_collection(db_id) collection = await self._get_chroma_collection(db_id)
@ -473,16 +452,32 @@ class ChromaKB(KnowledgeBase):
logger.error(f"ChromaDB query error: {e}, {traceback.format_exc()}") logger.error(f"ChromaDB query error: {e}, {traceback.format_exc()}")
return [] return []
async def delete_file_chunks_only(self, db_id: str, file_id: str) -> None:
"""仅删除文件的chunks数据保留元数据用于更新操作"""
collection = await self._get_chroma_collection(db_id)
if collection:
try:
# 查找所有相关的chunks
results = collection.get(where={"full_doc_id": file_id}, include=["metadatas"])
# 删除所有相关chunks
if results and results.get("ids"):
collection.delete(ids=results["ids"])
logger.info(f"Deleted {len(results['ids'])} chunks for file {file_id}")
except Exception as e:
logger.error(f"Error deleting file {file_id} from ChromaDB: {e}")
# 注意:这里不删除 files_meta[file_id],保留元数据用于后续操作
async def delete_file(self, db_id: str, file_id: str) -> None: async def delete_file(self, db_id: str, file_id: str) -> None:
"""删除文件(包括元数据)""" """删除文件(包括元数据)"""
# 先删除 ChromaDB 中的 chunks 数据 # 先删除 ChromaDB 中的 chunks 数据
await self.delete_file_chunks_only(db_id, file_id) await self.delete_file_chunks_only(db_id, file_id)
# 使用锁确保元数据操作的原子性 # 删除文件记录
async with self._metadata_lock: if file_id in self.files_meta:
if file_id in self.files_meta: del self.files_meta[file_id]
del self.files_meta[file_id] self._save_metadata()
self._save_metadata()
async def get_file_basic_info(self, db_id: str, file_id: str) -> dict: async def get_file_basic_info(self, db_id: str, file_id: str) -> dict:
"""获取文件基本信息(仅元数据)""" """获取文件基本信息(仅元数据)"""

View File

@ -32,9 +32,6 @@ class LightRagKB(KnowledgeBase):
# 存储 LightRAG 实例映射 {db_id: LightRAG} # 存储 LightRAG 实例映射 {db_id: LightRAG}
self.instances: dict[str, LightRAG] = {} self.instances: dict[str, LightRAG] = {}
# 元数据锁
self._metadata_lock = asyncio.Lock()
# 设置 LightRAG 日志 # 设置 LightRAG 日志
log_dir = os.path.join(work_dir, "logs", "lightrag") log_dir = os.path.join(work_dir, "logs", "lightrag")
os.makedirs(log_dir, exist_ok=True) os.makedirs(log_dir, exist_ok=True)
@ -306,6 +303,7 @@ class LightRagKB(KnowledgeBase):
try: try:
# 更新状态为处理中 # 更新状态为处理中
self.files_meta[file_id]["processing_params"] = params.copy()
self.files_meta[file_id]["status"] = "processing" self.files_meta[file_id]["status"] = "processing"
self._save_metadata() self._save_metadata()
@ -357,19 +355,6 @@ class LightRagKB(KnowledgeBase):
return processed_items_info return processed_items_info
async def delete_file_chunks_only(self, db_id: str, file_id: str) -> None:
"""仅删除文件的chunks数据保留元数据用于更新操作"""
rag = await self._get_lightrag_instance(db_id)
if rag:
try:
# 使用 LightRAG 删除文档
await rag.adelete_by_doc_id(file_id)
logger.info(f"Deleted chunks for file {file_id} from LightRAG")
except Exception as e:
logger.error(f"Error deleting file {file_id} from LightRAG: {e}")
# 注意:这里不删除 files_meta[file_id],保留元数据用于后续操作
async def aquery(self, query_text: str, db_id: str, **kwargs) -> str: async def aquery(self, query_text: str, db_id: str, **kwargs) -> str:
"""异步查询知识库""" """异步查询知识库"""
rag = await self._get_lightrag_instance(db_id) rag = await self._get_lightrag_instance(db_id)
@ -395,16 +380,27 @@ class LightRagKB(KnowledgeBase):
logger.error(f"Query error: {e}, {traceback.format_exc()}") logger.error(f"Query error: {e}, {traceback.format_exc()}")
return "" return ""
async def delete_file_chunks_only(self, db_id: str, file_id: str) -> None:
"""仅删除文件的chunks数据保留元数据用于更新操作"""
rag = await self._get_lightrag_instance(db_id)
if rag:
try:
# 使用 LightRAG 删除文档
await rag.adelete_by_doc_id(file_id)
logger.info(f"Deleted chunks for file {file_id} from LightRAG")
except Exception as e:
logger.error(f"Error deleting file {file_id} from LightRAG: {e}")
# 注意:这里不删除 files_meta[file_id],保留元数据用于后续操作
async def delete_file(self, db_id: str, file_id: str) -> None: async def delete_file(self, db_id: str, file_id: str) -> None:
"""删除文件(包括元数据)""" """删除文件(包括元数据)"""
# 先删除 LightRAG 中的 chunks 数据 # 先删除 LightRAG 中的 chunks 数据
await self.delete_file_chunks_only(db_id, file_id) await self.delete_file_chunks_only(db_id, file_id)
# 使用锁确保元数据操作的原子性 # 删除文件记录
async with self._metadata_lock: if file_id in self.files_meta:
if file_id in self.files_meta: del self.files_meta[file_id]
del self.files_meta[file_id] self._save_metadata()
self._save_metadata()
async def get_file_basic_info(self, db_id: str, file_id: str) -> dict: async def get_file_basic_info(self, db_id: str, file_id: str) -> dict:
"""获取文件基本信息(仅元数据)""" """获取文件基本信息(仅元数据)"""

View File

@ -332,6 +332,7 @@ class MilvusKB(KnowledgeBase):
try: try:
# 更新状态为处理中 # 更新状态为处理中
async with self._metadata_lock: async with self._metadata_lock:
self.files_meta[file_id]["processing_params"] = params.copy()
self.files_meta[file_id]["status"] = "processing" self.files_meta[file_id]["status"] = "processing"
self._save_metadata() self._save_metadata()

View File

@ -203,7 +203,8 @@ const selectedRowKeys = computed({
// //
const rechunkModalVisible = ref(false); const rechunkModalVisible = ref(false);
const rechunkModalLoading = ref(false); const rechunkModalLoading = computed(() => store.state.chunkLoading);
const rechunkParams = ref({ const rechunkParams = ref({
chunk_size: 1000, chunk_size: 1000,
chunk_overlap: 200, chunk_overlap: 200,
@ -448,18 +449,17 @@ const handleDownloadFile = async (record) => {
}; };
const handleRechunkFile = async (record) => { const handleRechunkFile = async (record) => {
const dbId = store.databaseId;
if (!dbId) {
console.error('无法获取数据库ID数据库ID:', store.databaseId, '记录:', record);
message.error('无法获取数据库ID请刷新页面后重试');
return;
}
try { try {
// ID // ID
currentRechunkFileIds.value = [record.file_id]; currentRechunkFileIds.value = [record.file_id];
isBatchRechunk.value = false; isBatchRechunk.value = false;
if (record?.processing_params) {
rechunkParams.value = {
...record?.processing_params
};
}
// //
rechunkModalVisible.value = true; rechunkModalVisible.value = true;
} catch (error) { } catch (error) {
@ -470,34 +470,25 @@ const handleRechunkFile = async (record) => {
// //
const handleRechunkConfirm = async () => { const handleRechunkConfirm = async () => {
const dbId = store.databaseId;
if (!dbId) {
console.error('无法获取数据库ID数据库ID:', store.databaseId);
message.error('无法获取数据库ID请刷新页面后重试');
return;
}
if (currentRechunkFileIds.value.length === 0) {
message.warning('请选择要重新分块的文件');
return;
}
rechunkModalLoading.value = true;
try { try {
// rechunks // rechunks
const result = await documentApi.rechunksDocuments(dbId, currentRechunkFileIds.value, rechunkParams.value); const result = await store.rechunksFiles({fileIds: currentRechunkFileIds.value, params: rechunkParams.value});
if (result) {
if (result.status === 'queued') { currentRechunkFileIds.value = [];
message.success('重新分块任务已提交,请在任务中心查看进度');
//
store.getDatabaseInfo(undefined, true);
// //
if (isBatchRechunk.value) { if (isBatchRechunk.value) {
selectedRowKeys.value = []; selectedRowKeys.value = [];
} }
// //
rechunkModalVisible.value = false; rechunkModalVisible.value = false;
//
rechunkParams.value = {
chunk_size: 1000,
chunk_overlap: 200,
use_qa_split: false,
qa_separator: '\n\n\n'
};
} else { } else {
message.error(`重新分块失败: ${result.message}`); message.error(`重新分块失败: ${result.message}`);
} }
@ -505,15 +496,13 @@ const handleRechunkConfirm = async () => {
console.error('重新分块失败:', error); console.error('重新分块失败:', error);
const errorMessage = error.message || '重新分块失败,请稍后重试'; const errorMessage = error.message || '重新分块失败,请稍后重试';
message.error(errorMessage); message.error(errorMessage);
} finally {
rechunkModalLoading.value = false;
} }
}; };
// //
const handleRechunkCancel = () => { const handleRechunkCancel = () => {
rechunkModalVisible.value = false; rechunkModalVisible.value = false;
rechunkModalLoading.value = false; // rechunkModalLoading.value = false;
currentRechunkFileIds.value = []; currentRechunkFileIds.value = [];
isBatchRechunk.value = false; isBatchRechunk.value = false;
// //

View File

@ -256,6 +256,45 @@ export const useDatabaseStore = defineStore('database', () => {
} }
} }
async function rechunksFiles({ fileIds, params }) {
if (fileIds.length === 0) {
message.error('请选择要重新分块的文件!');
return;
}
state.chunkLoading = true;
try {
const data = await documentApi.rechunksDocuments(databaseId.value, fileIds, { ...params });
if (data.status === 'success' || data.status === 'queued') {
enableAutoRefresh('auto');
message.success(data.message || `文档已提交处理,请在任务中心查看进度`);
if (data.task_id) {
taskerStore.registerQueuedTask({
task_id: data.task_id,
name: `文档重新分块 (${databaseId.value || ''})`,
task_type: 'knowledge_rechunks',
message: data.message,
payload: {
db_id: databaseId.value,
count: fileIds.length,
}
});
}
await getDatabaseInfo(undefined, true); // Skip query params when adding files
return true; // Indicate success
} else {
message.error(data.message || '处理失败');
return false;
}
} catch (error) {
console.error(error);
message.error(error.message || '处理请求失败');
return false;
} finally {
state.chunkLoading = false;
}
}
async function openFileDetail(record) { async function openFileDetail(record) {
if (record.status !== 'done') { if (record.status !== 'done') {
message.error('文件未处理完成,请稍后再试'); message.error('文件未处理完成,请稍后再试');
@ -379,6 +418,7 @@ export const useDatabaseStore = defineStore('database', () => {
handleDeleteFile, handleDeleteFile,
handleBatchDelete, handleBatchDelete,
addFiles, addFiles,
rechunksFiles,
openFileDetail, openFileDetail,
loadQueryParams, loadQueryParams,