From 05685f8f524c34b59f78b9aeb0cc353936bc8667 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 8 Nov 2024 17:24:47 +0800 Subject: [PATCH] fix issues:Image file not deleted when a doc is removed #9541 Signed-off-by: root --- api/core/tools/utils/web_reader_tool.py | 13 +++++++++++++ api/tasks/clean_dataset_task.py | 8 ++++++++ api/tasks/clean_document_task.py | 8 ++++++++ 3 files changed, 29 insertions(+) diff --git a/api/core/tools/utils/web_reader_tool.py b/api/core/tools/utils/web_reader_tool.py index 5807d61b9409a6..5f04587b7b1b7b 100644 --- a/api/core/tools/utils/web_reader_tool.py +++ b/api/core/tools/utils/web_reader_tool.py @@ -356,3 +356,16 @@ def content_digest(element): digest.update(child.encode("utf-8")) digest = digest.hexdigest() return digest + + +def get_image_upload_file_ids(content): + pattern = r'!\[image\]\((.*?file-preview)\)' + matches = re.findall(pattern, content) + content_pattern = r'files/([^/]+)/file-preview' + image_upload_file_ids = [] + for match in matches: + content_match = re.search(content_pattern, match) + if content_match: + image_upload_file_id = content_match.group(1) + image_upload_file_ids.append(image_upload_file_id) + return image_upload_file_ids diff --git a/api/tasks/clean_dataset_task.py b/api/tasks/clean_dataset_task.py index 36249038011747..85f88a120c2c00 100644 --- a/api/tasks/clean_dataset_task.py +++ b/api/tasks/clean_dataset_task.py @@ -5,6 +5,7 @@ from celery import shared_task from core.rag.index_processor.index_processor_factory import IndexProcessorFactory +from core.tools.utils.web_reader_tool import get_image_upload_file_ids from extensions.ext_database import db from extensions.ext_storage import storage from models.dataset import ( @@ -67,6 +68,13 @@ def clean_dataset_task( db.session.delete(document) for segment in segments: + image_upload_file_ids = get_image_upload_file_ids(segment.content) + for upload_file_id in image_upload_file_ids: + image_file = db.session.query(UploadFile).filter(UploadFile.id == upload_file_id).first() + try: + storage.delete(image_file.key) + except Exception: + logging.exception("Delete image_files failed when storage deleted, image_upload_file_is: {}".format(upload_file_id)) db.session.delete(segment) db.session.query(DatasetProcessRule).filter(DatasetProcessRule.dataset_id == dataset_id).delete() diff --git a/api/tasks/clean_document_task.py b/api/tasks/clean_document_task.py index ae2855aa2ebc4d..e354fab0b6d5f6 100644 --- a/api/tasks/clean_document_task.py +++ b/api/tasks/clean_document_task.py @@ -6,6 +6,7 @@ from celery import shared_task from core.rag.index_processor.index_processor_factory import IndexProcessorFactory +from core.tools.utils.web_reader_tool import get_image_upload_file_ids from extensions.ext_database import db from extensions.ext_storage import storage from models.dataset import Dataset, DocumentSegment @@ -40,6 +41,13 @@ def clean_document_task(document_id: str, dataset_id: str, doc_form: str, file_i index_processor.clean(dataset, index_node_ids) for segment in segments: + image_upload_file_ids = get_image_upload_file_ids(segment.content) + for upload_file_id in image_upload_file_ids: + image_file = db.session.query(UploadFile).filter(UploadFile.id == upload_file_id).first() + try: + storage.delete(image_file.key) + except Exception: + logging.exception("Delete image_files failed when storage deleted, image_upload_file_is: {}".format(upload_file_id)) db.session.delete(segment) db.session.commit()