From 1700628624984a4ad3470bece4e92159476ce232 Mon Sep 17 00:00:00 2001 From: supreme0597 Date: Wed, 12 Nov 2025 23:14:56 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=AE=8C=E5=96=84=E9=87=8D=E6=96=B0?= =?UTF-8?q?=E5=88=86=E5=9D=97=E5=8A=9F=E8=83=BD=EF=BC=8C=E4=BC=98=E5=8C=96?= =?UTF-8?q?=E4=BB=A3=E7=A0=81=E9=80=BB=E8=BE=91=EF=BC=8C=E4=BF=9D=E5=AD=98?= =?UTF-8?q?=E8=AF=B7=E6=B1=82=E5=8F=82=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/knowledge/base.py | 1 + src/knowledge/implementations/chroma.py | 49 ++++++++++------------ src/knowledge/implementations/lightrag.py | 38 ++++++++--------- src/knowledge/implementations/milvus.py | 1 + web/src/components/FileTable.vue | 51 +++++++++-------------- web/src/stores/database.js | 40 ++++++++++++++++++ 6 files changed, 101 insertions(+), 79 deletions(-) diff --git a/src/knowledge/base.py b/src/knowledge/base.py index efcefe9a..b98e0ea1 100644 --- a/src/knowledge/base.py +++ b/src/knowledge/base.py @@ -293,6 +293,7 @@ class KnowledgeBase(ABC): "type": file_info.get("file_type", ""), "status": file_info.get("status", "done"), "created_at": created_at, + "processing_params": file_info.get("processing_params", None), } # 按创建时间倒序排序文件列表 diff --git a/src/knowledge/implementations/chroma.py b/src/knowledge/implementations/chroma.py index 4f661b28..7350510e 100644 --- a/src/knowledge/implementations/chroma.py +++ b/src/knowledge/implementations/chroma.py @@ -47,10 +47,6 @@ class ChromaKB(KnowledgeBase): # 存储集合映射 {db_id: collection} self.collections: dict[str, Any] = {} - - # 元数据锁 - self._metadata_lock = asyncio.Lock() - logger.info("ChromaKB initialized") @property @@ -290,6 +286,7 @@ class ChromaKB(KnowledgeBase): try: # 更新状态为处理中 + self.files_meta[file_id]["processing_params"] = params.copy() self.files_meta[file_id]["status"] = "processing" self._save_metadata() @@ -361,24 +358,6 @@ class ChromaKB(KnowledgeBase): 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]: """异步查询知识库""" 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()}") 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: """删除文件（包括元数据）""" # 先删除 ChromaDB 中的 chunks 数据 await self.delete_file_chunks_only(db_id, file_id) - # 使用锁确保元数据操作的原子性 - async with self._metadata_lock: - if file_id in self.files_meta: - del self.files_meta[file_id] - self._save_metadata() + # 删除文件记录 + if file_id in self.files_meta: + del self.files_meta[file_id] + self._save_metadata() async def get_file_basic_info(self, db_id: str, file_id: str) -> dict: """获取文件基本信息（仅元数据）""" diff --git a/src/knowledge/implementations/lightrag.py b/src/knowledge/implementations/lightrag.py index d7465f5d..313a2a4d 100644 --- a/src/knowledge/implementations/lightrag.py +++ b/src/knowledge/implementations/lightrag.py @@ -32,9 +32,6 @@ class LightRagKB(KnowledgeBase): # 存储 LightRAG 实例映射 {db_id: LightRAG} self.instances: dict[str, LightRAG] = {} - # 元数据锁 - self._metadata_lock = asyncio.Lock() - # 设置 LightRAG 日志 log_dir = os.path.join(work_dir, "logs", "lightrag") os.makedirs(log_dir, exist_ok=True) @@ -306,6 +303,7 @@ class LightRagKB(KnowledgeBase): try: # 更新状态为处理中 + self.files_meta[file_id]["processing_params"] = params.copy() self.files_meta[file_id]["status"] = "processing" self._save_metadata() @@ -357,19 +355,6 @@ class LightRagKB(KnowledgeBase): 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: """异步查询知识库""" rag = await self._get_lightrag_instance(db_id) @@ -395,16 +380,27 @@ class LightRagKB(KnowledgeBase): logger.error(f"Query error: {e}, {traceback.format_exc()}") 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: """删除文件（包括元数据）""" # 先删除 LightRAG 中的 chunks 数据 await self.delete_file_chunks_only(db_id, file_id) - # 使用锁确保元数据操作的原子性 - async with self._metadata_lock: - if file_id in self.files_meta: - del self.files_meta[file_id] - self._save_metadata() + # 删除文件记录 + if file_id in self.files_meta: + del self.files_meta[file_id] + self._save_metadata() async def get_file_basic_info(self, db_id: str, file_id: str) -> dict: """获取文件基本信息（仅元数据）""" diff --git a/src/knowledge/implementations/milvus.py b/src/knowledge/implementations/milvus.py index 3d5fc669..cae06236 100644 --- a/src/knowledge/implementations/milvus.py +++ b/src/knowledge/implementations/milvus.py @@ -332,6 +332,7 @@ class MilvusKB(KnowledgeBase): try: # 更新状态为处理中 async with self._metadata_lock: + self.files_meta[file_id]["processing_params"] = params.copy() self.files_meta[file_id]["status"] = "processing" self._save_metadata() diff --git a/web/src/components/FileTable.vue b/web/src/components/FileTable.vue index 86efba6d..72d0c2fb 100644 --- a/web/src/components/FileTable.vue +++ b/web/src/components/FileTable.vue @@ -203,7 +203,8 @@ const selectedRowKeys = computed({ // 重新分块参数配置相关 const rechunkModalVisible = ref(false); -const rechunkModalLoading = ref(false); +const rechunkModalLoading = computed(() => store.state.chunkLoading); + const rechunkParams = ref({ chunk_size: 1000, chunk_overlap: 200, @@ -448,18 +449,17 @@ const handleDownloadFile = async (record) => { }; const handleRechunkFile = async (record) => { - const dbId = store.databaseId; - if (!dbId) { - console.error('无法获取数据库ID，数据库ID:', store.databaseId, '记录:', record); - message.error('无法获取数据库ID，请刷新页面后重试'); - return; - } - try { // 设置当前重新分块的文件ID currentRechunkFileIds.value = [record.file_id]; isBatchRechunk.value = false; + if (record?.processing_params) { + rechunkParams.value = { + ...record?.processing_params + }; + } + // 显示参数配置模态框 rechunkModalVisible.value = true; } catch (error) { @@ -470,34 +470,25 @@ const handleRechunkFile = async (record) => { // 重新分块确认 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 { // 调用 rechunks 接口 - const result = await documentApi.rechunksDocuments(dbId, currentRechunkFileIds.value, rechunkParams.value); - - if (result.status === 'queued') { - message.success('重新分块任务已提交，请在任务中心查看进度'); - // 刷新文件列表 - store.getDatabaseInfo(undefined, true); + const result = await store.rechunksFiles({fileIds: currentRechunkFileIds.value, params: rechunkParams.value}); + if (result) { + currentRechunkFileIds.value = []; // 清空选择 if (isBatchRechunk.value) { selectedRowKeys.value = []; } // 关闭模态框 rechunkModalVisible.value = false; + + // 重置参数为默认值 + rechunkParams.value = { + chunk_size: 1000, + chunk_overlap: 200, + use_qa_split: false, + qa_separator: '\n\n\n' + }; } else { message.error(`重新分块失败: ${result.message}`); } @@ -505,15 +496,13 @@ const handleRechunkConfirm = async () => { console.error('重新分块失败:', error); const errorMessage = error.message || '重新分块失败，请稍后重试'; message.error(errorMessage); - } finally { - rechunkModalLoading.value = false; } }; // 重新分块取消 const handleRechunkCancel = () => { rechunkModalVisible.value = false; - rechunkModalLoading.value = false; + // rechunkModalLoading.value = false; currentRechunkFileIds.value = []; isBatchRechunk.value = false; // 重置参数为默认值 diff --git a/web/src/stores/database.js b/web/src/stores/database.js index ae320d91..4164196a 100644 --- a/web/src/stores/database.js +++ b/web/src/stores/database.js @@ -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) { if (record.status !== 'done') { message.error('文件未处理完成，请稍后再试'); @@ -379,6 +418,7 @@ export const useDatabaseStore = defineStore('database', () => { handleDeleteFile, handleBatchDelete, addFiles, + rechunksFiles, openFileDetail, loadQueryParams,