kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
src/knowledge/base.py
View File

@ -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),
}
# 按创建时间倒序排序文件列表

49
src/knowledge/implementations/chroma.py
View File

@ -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:
"""获取文件基本信息（仅元数据）"""

38
src/knowledge/implementations/lightrag.py
View File

@ -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:
"""获取文件基本信息（仅元数据）"""

1
src/knowledge/implementations/milvus.py
View File

@ -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()

51
web/src/components/FileTable.vue
View File

@ -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;
// 重置参数为默认值

40
web/src/stores/database.js
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) {
if (record.status !== 'done') {
message.error('文件未处理完成，请稍后再试');
@ -379,6 +418,7 @@ export const useDatabaseStore = defineStore('database', () => {
handleDeleteFile,
handleBatchDelete,
addFiles,
rechunksFiles,
openFileDetail,
loadQueryParams,