feat: 完善重新分块功能,优化代码逻辑,保存请求参数
This commit is contained in:
parent
b7ba646fae
commit
1700628624
@ -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),
|
||||||
}
|
}
|
||||||
|
|
||||||
# 按创建时间倒序排序文件列表
|
# 按创建时间倒序排序文件列表
|
||||||
|
|||||||
@ -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:
|
||||||
"""获取文件基本信息(仅元数据)"""
|
"""获取文件基本信息(仅元数据)"""
|
||||||
|
|||||||
@ -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:
|
||||||
"""获取文件基本信息(仅元数据)"""
|
"""获取文件基本信息(仅元数据)"""
|
||||||
|
|||||||
@ -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()
|
||||||
|
|
||||||
|
|||||||
@ -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;
|
||||||
// 重置参数为默认值
|
// 重置参数为默认值
|
||||||
|
|||||||
@ -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,
|
||||||
|
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user