From 55d5879622c9d2d23ee7788eb8b17c7e0fd99216 Mon Sep 17 00:00:00 2001 From: Wang Qi Date: Tue, 14 Jul 2026 20:12:31 +0800 Subject: [PATCH] Fix file link-to-datasets to suport both replace and add (#817) (#16908) --- api/apps/restful_apis/file2document_api.py | 43 ++++++++++------ .../restful_api/test_file_routes_unit.py | 44 +++++++++++++++++ .../test_file2document_routes_unit.py | 9 ++++ web/src/hooks/use-file-request.ts | 49 +++++++++++++++++-- web/src/interfaces/request/file-manager.ts | 2 + web/src/pages/files/hooks.ts | 28 ++++++++++- 6 files changed, 157 insertions(+), 18 deletions(-) diff --git a/api/apps/restful_apis/file2document_api.py b/api/apps/restful_apis/file2document_api.py index b5524bcadd..212234eb7e 100644 --- a/api/apps/restful_apis/file2document_api.py +++ b/api/apps/restful_apis/file2document_api.py @@ -18,6 +18,8 @@ import asyncio import logging from pathlib import Path +from quart import request + from api.common.check_team_permission import check_file_team_permission, check_kb_team_permission from api.db.services.file2document_service import File2DocumentService from api.db.services.file_service import FileService @@ -25,6 +27,7 @@ from api.db.services.file_service import FileService from api.apps import login_required, current_user from api.db.services.knowledgebase_service import KnowledgebaseService from api.utils.api_utils import get_data_error_result, get_json_result, get_request_json, server_error_response, validate_request +from common.constants import RetCode from common.misc_utils import get_uuid from api.db import FileType from api.db.services.document_service import DocumentService @@ -32,25 +35,34 @@ from api.db.services.document_service import DocumentService logger = logging.getLogger(__name__) -def _convert_files(file_ids, kb_ids, user_id): - """Synchronous worker: add new docs for the given file/kb pairs while preserving existing links. - - Previously this function replaced all existing links with the new ones, which caused - multi-select "link to knowledge base" to overwrite previous links. Now it only creates - documents for knowledge bases that are not already linked to the file, and leaves - existing links untouched. - """ +def _convert_files(file_ids, kb_ids, user_id, mode): + """Synchronous worker: add missing links or replace existing links.""" + replace_existing = mode == "replace" for id in file_ids: e, file = FileService.get_by_id(id) if not e: continue - existing_links = {inform.document_id for inform in File2DocumentService.get_by_file_id(id)} + existing_links = File2DocumentService.get_by_file_id(id) existing_kb_ids = set() - for doc_id in existing_links: - e, doc = DocumentService.get_by_id(doc_id) - if e and doc: - existing_kb_ids.add(doc.kb_id) + if replace_existing: + for inform in existing_links: + doc_id = inform.document_id + e, doc = DocumentService.get_by_id(doc_id) + if e and doc: + tenant_id = DocumentService.get_tenant_id(doc_id) + if not tenant_id: + raise RuntimeError("Tenant not found!") + if not DocumentService.remove_document(doc, tenant_id): + raise RuntimeError("Database error (Document removal)!") + File2DocumentService.delete_by_document_id(doc_id) + if existing_links: + File2DocumentService.delete_by_file_id(id) + else: + for inform in existing_links: + e, doc = DocumentService.get_by_id(inform.document_id) + if e and doc: + existing_kb_ids.add(doc.kb_id) for kb_id in kb_ids: if kb_id in existing_kb_ids: @@ -89,6 +101,9 @@ async def convert(): req = await get_request_json() kb_ids = req["kb_ids"] file_ids = req["file_ids"] + mode = (request.args.get("mode", "replace") or "replace").lower() + if mode not in {"replace", "add"}: + return get_json_result(code=RetCode.ARGUMENT_ERROR, message="mode must be 'add' or 'replace'") try: files = FileService.get_by_ids(file_ids) @@ -167,7 +182,7 @@ async def convert(): # For large folders this prevents 504 Gateway Timeout by returning as # soon as the background task is scheduled. loop = asyncio.get_running_loop() - future = loop.run_in_executor(None, _convert_files, all_file_ids, kb_ids, user_id) + future = loop.run_in_executor(None, _convert_files, all_file_ids, kb_ids, user_id, mode) future.add_done_callback(lambda f: logging.error("_convert_files failed: %s", f.exception()) if f.exception() else None) logger.info( "user_id=%s resource_type=file_to_dataset_link resource_id=batch action=schedule_convert result=scheduled file_ids=%s kb_ids=%s", diff --git a/test/testcases/restful_api/test_file_routes_unit.py b/test/testcases/restful_api/test_file_routes_unit.py index 6d055936cf..b4d16bdf27 100644 --- a/test/testcases/restful_api/test_file_routes_unit.py +++ b/test/testcases/restful_api/test_file_routes_unit.py @@ -402,6 +402,10 @@ def _set_request_json(monkeypatch, module, payload_state): monkeypatch.setattr(module, "get_request_json", _req_json) +def _set_request_args(monkeypatch, module, args): + monkeypatch.setattr(module, "request", SimpleNamespace(args=args)) + + @pytest.fixture(scope="session") def auth(): return "unit-auth" @@ -462,6 +466,10 @@ def _load_file2document_module(monkeypatch): def delete_by_file_id(*_args, **_kwargs): return None + @staticmethod + def delete_by_document_id(*_args, **_kwargs): + return None + @staticmethod def insert(_payload): return SimpleNamespace(to_json=lambda: {}) @@ -585,6 +593,7 @@ def test_convert_branch_matrix_unit(monkeypatch): module = _load_file2document_module(monkeypatch) req_state = {"kb_ids": ["kb-1"], "file_ids": ["f1"]} _set_request_json(monkeypatch, module, req_state) + _set_request_args(monkeypatch, module, {}) # Falsy file returns "File not found!" during synchronous validation. monkeypatch.setattr(module.FileService, "get_by_ids", lambda _ids: [_FalsyFile("f1", module.FileType.DOC.value)]) @@ -640,6 +649,41 @@ def test_convert_branch_matrix_unit(monkeypatch): assert "convert boom" in res["message"] +@pytest.mark.p2 +def test_convert_files_mode_add_and_replace_unit(monkeypatch): + module = _load_file2document_module(monkeypatch) + inserted = [] + removed = [] + deleted_doc_links = [] + deleted_file_links = [] + file = _DummyFile("f1", module.FileType.DOC.value, name="a.txt") + kb = SimpleNamespace(id="kb-new", parser_id="naive", pipeline_id="p1", parser_config={}) + + monkeypatch.setattr(module.FileService, "get_by_id", lambda _file_id: (True, file)) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb)) + monkeypatch.setattr(module.File2DocumentService, "get_by_file_id", lambda file_id: [SimpleNamespace(document_id=f"doc-{file_id}")]) + monkeypatch.setattr(module.File2DocumentService, "delete_by_document_id", lambda doc_id: deleted_doc_links.append(doc_id)) + monkeypatch.setattr(module.File2DocumentService, "delete_by_file_id", lambda file_id: deleted_file_links.append(file_id)) + monkeypatch.setattr(module.File2DocumentService, "insert", lambda payload: inserted.append(payload) or SimpleNamespace(to_json=lambda: {})) + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda doc_id: (True, SimpleNamespace(id=doc_id, kb_id="kb-old"))) + monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "tenant-1") + monkeypatch.setattr(module.DocumentService, "remove_document", lambda doc, tenant_id: removed.append((doc.id, tenant_id)) or True) + monkeypatch.setattr(module.DocumentService, "insert", lambda payload: SimpleNamespace(id=f"new-{payload['kb_id']}")) + + module._convert_files(["f1", "f2"], ["kb-old", "kb-new"], "user-1", "add") + assert len(inserted) == 2 + assert removed == [] + assert deleted_doc_links == [] + assert deleted_file_links == [] + + inserted.clear() + module._convert_files(["f1", "f2"], ["kb-new"], "user-1", "replace") + assert len(inserted) == 2 + assert removed == [("doc-f1", "tenant-1"), ("doc-f2", "tenant-1")] + assert deleted_doc_links == ["doc-f1", "doc-f2"] + assert deleted_file_links == ["f1", "f2"] + + def _load_file_api_service(monkeypatch): LOGGER.debug("_load_file_api_service: entry") repo_root = Path(__file__).resolve().parents[3] diff --git a/test/testcases/test_web_api/test_file_app/test_file2document_routes_unit.py b/test/testcases/test_web_api/test_file_app/test_file2document_routes_unit.py index e4850e3643..5279e36626 100644 --- a/test/testcases/test_web_api/test_file_app/test_file2document_routes_unit.py +++ b/test/testcases/test_web_api/test_file_app/test_file2document_routes_unit.py @@ -59,6 +59,10 @@ def _set_request_json(monkeypatch, module, payload_state): monkeypatch.setattr(module, "get_request_json", _req_json) +def _set_request_args(monkeypatch, module, args): + monkeypatch.setattr(module, "request", SimpleNamespace(args=args)) + + @pytest.fixture(scope="session") def auth(): return "unit-auth" @@ -119,6 +123,10 @@ def _load_file2document_module(monkeypatch): def delete_by_file_id(*_args, **_kwargs): return None + @staticmethod + def delete_by_document_id(*_args, **_kwargs): + return None + @staticmethod def insert(_payload): return SimpleNamespace(to_json=lambda: {}) @@ -242,6 +250,7 @@ def test_convert_branch_matrix_unit(monkeypatch): module = _load_file2document_module(monkeypatch) req_state = {"kb_ids": ["kb-1"], "file_ids": ["f1"]} _set_request_json(monkeypatch, module, req_state) + _set_request_args(monkeypatch, module, {}) # Falsy file returns "File not found!" during synchronous validation. monkeypatch.setattr(module.FileService, "get_by_ids", lambda _ids: [_FalsyFile("f1", module.FileType.DOC.value)]) diff --git a/web/src/hooks/use-file-request.ts b/web/src/hooks/use-file-request.ts index 8c975c5b8e..5d3fa58ac8 100644 --- a/web/src/hooks/use-file-request.ts +++ b/web/src/hooks/use-file-request.ts @@ -2,11 +2,17 @@ import message from '@/components/ui/message'; import { PaginationProps } from '@/interfaces/antd-compat'; import { IFetchFileListResult, + IFile, IFolder, } from '@/interfaces/database/file-manager'; -import { IConnectRequestBody } from '@/interfaces/request/file-manager'; +import { + ConnectFileToKnowledgeMode, + IConnectRequestBody, +} from '@/interfaces/request/file-manager'; import fileManagerService from '@/services/file-manager-service'; +import api from '@/utils/api'; import { downloadFileFromBlob } from '@/utils/file-util'; +import request from '@/utils/request'; import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; import { useDebounce } from 'ahooks'; import { useCallback } from 'react'; @@ -294,12 +300,49 @@ export const useConnectToKnowledge = () => { mutateAsync, } = useMutation({ mutationKey: [FileApiAction.ConnectFileToKnowledge], - mutationFn: async (params: IConnectRequestBody) => { - const { data } = await fileManagerService.connectFileToKnowledge(params); + mutationFn: async ( + params: IConnectRequestBody & { + mode: ConnectFileToKnowledgeMode; + kbsInfo: IFile['kbs_info']; + }, + ) => { + const { data } = await request.post(api.connectFileToKnowledge, { + data: { fileIds: params.fileIds, kbIds: params.kbIds }, + params: { mode: params.mode }, + }); if (data.code === 0) { message.success(t('message.operated')); + const fileIdSet = new Set(params.fileIds); + queryClient.setQueriesData( + { + queryKey: [FileApiAction.FetchFileList], + }, + (oldData) => { + if (!oldData?.files) return oldData; + return { + ...oldData, + files: oldData.files.map((file) => { + if (!fileIdSet.has(file.id)) return file; + const kbsInfo = + params.mode === 'replace' + ? params.kbsInfo + : [ + ...(file.kbs_info ?? []), + ...params.kbsInfo.filter( + (kb) => + !(file.kbs_info ?? []).some( + (item) => item.kb_id === kb.kb_id, + ), + ), + ]; + return { ...file, kbs_info: kbsInfo }; + }), + }; + }, + ); queryClient.invalidateQueries({ queryKey: [FileApiAction.FetchFileList], + refetchType: 'none', }); } return data.code; diff --git a/web/src/interfaces/request/file-manager.ts b/web/src/interfaces/request/file-manager.ts index 7c79e81c3f..3458b915a6 100644 --- a/web/src/interfaces/request/file-manager.ts +++ b/web/src/interfaces/request/file-manager.ts @@ -8,3 +8,5 @@ export interface IConnectRequestBody { fileIds: string[]; kbIds: string[]; } + +export type ConnectFileToKnowledgeMode = 'add' | 'replace'; diff --git a/web/src/pages/files/hooks.ts b/web/src/pages/files/hooks.ts index 4fea3879da..86f268947b 100644 --- a/web/src/pages/files/hooks.ts +++ b/web/src/pages/files/hooks.ts @@ -1,7 +1,9 @@ import { useSetModalState } from '@/hooks/common-hooks'; import { useConnectToKnowledge, useRenameFile } from '@/hooks/use-file-request'; +import { useSelectKnowledgeOptions } from '@/hooks/use-knowledge-request'; import { TableRowSelection } from '@/interfaces/antd-compat'; import { IFile } from '@/interfaces/database/file-manager'; +import { ConnectFileToKnowledgeMode } from '@/interfaces/request/file-manager'; import { useCallback, useMemo, useState } from 'react'; import { useNavigate, useSearchParams } from 'react-router'; @@ -83,6 +85,8 @@ export const useHandleConnectToKnowledge = () => { useConnectToKnowledge(); const [record, setRecord] = useState({} as IFile); const [documentIds, setDocumentIds] = useState([]); + const [mode, setMode] = useState('replace'); + const knowledgeOptions = useSelectKnowledgeOptions(); const initialValue = useMemo(() => { return Array.isArray(record?.kbs_info) @@ -90,11 +94,25 @@ export const useHandleConnectToKnowledge = () => { : []; }, [record?.kbs_info]); + const knowledgeNameMap = useMemo(() => { + return new Map( + knowledgeOptions?.map((option) => [ + option.value, + typeof option.label === 'string' ? option.label : String(option.label), + ]) ?? [], + ); + }, [knowledgeOptions]); + const onConnectToKnowledgeOk = useCallback( async (knowledgeIds: string[]) => { const ret = await connectToKnowledge({ fileIds: documentIds, kbIds: knowledgeIds, + mode, + kbsInfo: knowledgeIds.map((id) => ({ + kb_id: id, + kb_name: knowledgeNameMap.get(id) ?? id, + })), }); if (ret === 0) { @@ -102,7 +120,13 @@ export const useHandleConnectToKnowledge = () => { } return ret; }, - [connectToKnowledge, hideConnectToKnowledgeModal, documentIds], + [ + connectToKnowledge, + hideConnectToKnowledgeModal, + documentIds, + mode, + knowledgeNameMap, + ], ); const handleShowConnectToKnowledgeModal = useCallback( @@ -110,9 +134,11 @@ export const useHandleConnectToKnowledge = () => { if (Array.isArray(documents)) { setDocumentIds(documents); setRecord({} as IFile); + setMode('add'); } else { setRecord(documents); setDocumentIds([documents.id]); + setMode('replace'); } showConnectToKnowledgeModal();