mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-16 12:47:19 +08:00
This commit is contained in:
@@ -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",
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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)])
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -8,3 +8,5 @@ export interface IConnectRequestBody {
|
||||
fileIds: string[];
|
||||
kbIds: string[];
|
||||
}
|
||||
|
||||
export type ConnectFileToKnowledgeMode = 'add' | 'replace';
|
||||
|
||||
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user