Fix file link-to-datasets to suport both replace and add (#817) (#16908)

This commit is contained in:
Wang Qi
2026-07-14 20:12:31 +08:00
committed by GitHub
parent 376a9057a8
commit 55d5879622
6 changed files with 157 additions and 18 deletions

View File

@@ -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",

View File

@@ -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]

View File

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

View File

@@ -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<IFetchFileListResult>(
{
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;

View File

@@ -8,3 +8,5 @@ export interface IConnectRequestBody {
fileIds: string[];
kbIds: string[];
}
export type ConnectFileToKnowledgeMode = 'add' | 'replace';

View File

@@ -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<IFile>({} as IFile);
const [documentIds, setDocumentIds] = useState<string[]>([]);
const [mode, setMode] = useState<ConnectFileToKnowledgeMode>('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();