diff --git a/api/controllers/console/__init__.py b/api/controllers/console/__init__.py index e25f92399..57dbc8da6 100644 --- a/api/controllers/console/__init__.py +++ b/api/controllers/console/__init__.py @@ -84,6 +84,7 @@ from .datasets import ( external, hit_testing, metadata, + upload_file, website, ) diff --git a/api/controllers/console/datasets/upload_file.py b/api/controllers/console/datasets/upload_file.py new file mode 100644 index 000000000..9b456c771 --- /dev/null +++ b/api/controllers/console/datasets/upload_file.py @@ -0,0 +1,62 @@ +from flask_login import current_user +from flask_restful import Resource +from werkzeug.exceptions import NotFound + +from controllers.console import api +from controllers.console.wraps import ( + account_initialization_required, + setup_required, +) +from core.file import helpers as file_helpers +from extensions.ext_database import db +from models.dataset import Dataset +from models.model import UploadFile +from services.dataset_service import DocumentService + + +class UploadFileApi(Resource): + @setup_required + @account_initialization_required + def get(self, dataset_id, document_id): + """Get upload file.""" + # check dataset + dataset_id = str(dataset_id) + dataset = ( + db.session.query(Dataset) + .filter(Dataset.tenant_id == current_user.current_tenant_id, Dataset.id == dataset_id) + .first() + ) + if not dataset: + raise NotFound("Dataset not found.") + # check document + document_id = str(document_id) + document = DocumentService.get_document(dataset.id, document_id) + if not document: + raise NotFound("Document not found.") + # check upload file + if document.data_source_type != "upload_file": + raise ValueError(f"Document data source type ({document.data_source_type}) is not upload_file.") + data_source_info = document.data_source_info_dict + if data_source_info and "upload_file_id" in data_source_info: + file_id = data_source_info["upload_file_id"] + upload_file = db.session.query(UploadFile).filter(UploadFile.id == file_id).first() + if not upload_file: + raise NotFound("UploadFile not found.") + else: + raise ValueError("Upload file id not found in document data source info.") + + url = file_helpers.get_signed_file_url(upload_file_id=upload_file.id) + return { + "id": upload_file.id, + "name": upload_file.name, + "size": upload_file.size, + "extension": upload_file.extension, + "url": url, + "download_url": f"{url}&as_attachment=true", + "mime_type": upload_file.mime_type, + "created_by": upload_file.created_by, + "created_at": upload_file.created_at.timestamp(), + }, 200 + + +api.add_resource(UploadFileApi, "/datasets//documents//upload-file") diff --git a/web/app/components/datasets/documents/list.tsx b/web/app/components/datasets/documents/list.tsx index 2697580f4..abfa57813 100644 --- a/web/app/components/datasets/documents/list.tsx +++ b/web/app/components/datasets/documents/list.tsx @@ -7,6 +7,7 @@ import { pick, uniq } from 'lodash-es' import { RiArchive2Line, RiDeleteBinLine, + RiDownloadLine, RiEditLine, RiEqualizer2Line, RiLoopLeftLine, @@ -35,6 +36,7 @@ import type { ColorMap, IndicatorProps } from '@/app/components/header/indicator import Indicator from '@/app/components/header/indicator' import { asyncRunSafe } from '@/utils' import { formatNumber } from '@/utils/format' +import { useDocumentDownload } from '@/service/knowledge/use-document' import NotionIcon from '@/app/components/base/notion-icon' import ProgressBar from '@/app/components/base/progress-bar' import { ChunkingMode, DataSourceType, DocumentActionType, type DocumentDisplayStatus, type SimpleDocumentDetail } from '@/models/datasets' @@ -97,6 +99,7 @@ export const StatusItem: FC<{ const { mutateAsync: enableDocument } = useDocumentEnable() const { mutateAsync: disableDocument } = useDocumentDisable() const { mutateAsync: deleteDocument } = useDocumentDelete() + const downloadDocument = useDocumentDownload() const onOperate = async (operationName: OperationName) => { let opApi = deleteDocument @@ -188,6 +191,7 @@ export const OperationAction: FC<{ scene?: 'list' | 'detail' className?: string }> = ({ embeddingAvailable, datasetId, detail, onUpdate, scene = 'list', className = '' }) => { + const downloadDocument = useDocumentDownload() const { id, enabled = false, archived = false, data_source_type, display_status } = detail || {} const [showModal, setShowModal] = useState(false) const [deleting, setDeleting] = useState(false) @@ -296,6 +300,31 @@ export const OperationAction: FC<{ )} {embeddingAvailable && ( <> + + + { }) } +// Download document with authentication (sends Authorization header) +export const useDocumentDownload = () => { + return useMutation({ + mutationFn: async ({ datasetId, documentId }: { datasetId: string; documentId: string }) => { + // The get helper automatically adds the Authorization header from localStorage + return get(`/datasets/${datasetId}/documents/${documentId}/upload-file`) + }, + onError: (error: any) => { + // Show a toast notification if download fails + const message = error?.message || 'Download failed.' + Toast.notify({ type: 'error', message }) + }, + }) +} + export const useSyncWebsite = () => { return useMutation({ mutationFn: ({ datasetId, documentId }: UpdateDocumentBatchParams) => {