Feat/support multimodal embedding (#29115)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/tasks/add_document_to_index_task.py

@@ -4,9 +4,10 @@ import time
 import click
 from celery import shared_task
 
-from core.rag.index_processor.constant.index_type import IndexType
+from core.rag.index_processor.constant.doc_type import DocType
+from core.rag.index_processor.constant.index_type import IndexStructureType
 from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
-from core.rag.models.document import ChildDocument, Document
+from core.rag.models.document import AttachmentDocument, ChildDocument, Document
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from libs.datetime_utils import naive_utc_now
@@ -55,6 +56,7 @@ def add_document_to_index_task(dataset_document_id: str):
         )
 
         documents = []
+        multimodal_documents = []
         for segment in segments:
             document = Document(
                 page_content=segment.content,
@@ -65,7 +67,7 @@ def add_document_to_index_task(dataset_document_id: str):
                     "dataset_id": segment.dataset_id,
                 },
             )
-            if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
+            if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
                 child_chunks = segment.get_child_chunks()
                 if child_chunks:
                     child_documents = []
@@ -81,11 +83,25 @@ def add_document_to_index_task(dataset_document_id: str):
                         )
                         child_documents.append(child_document)
                     document.children = child_documents
+            if dataset.is_multimodal:
+                for attachment in segment.attachments:
+                    multimodal_documents.append(
+                        AttachmentDocument(
+                            page_content=attachment["name"],
+                            metadata={
+                                "doc_id": attachment["id"],
+                                "doc_hash": "",
+                                "document_id": segment.document_id,
+                                "dataset_id": segment.dataset_id,
+                                "doc_type": DocType.IMAGE,
+                            },
+                        )
+                    )
             documents.append(document)
 
         index_type = dataset.doc_form
         index_processor = IndexProcessorFactory(index_type).init_index_processor()
-        index_processor.load(dataset, documents)
+        index_processor.load(dataset, documents, multimodal_documents=multimodal_documents)
 
         # delete auto disable log
         db.session.query(DatasetAutoDisableLog).where(DatasetAutoDisableLog.document_id == dataset_document.id).delete()

api/tasks/clean_dataset_task.py

@@ -18,6 +18,7 @@ from models.dataset import (
     DatasetQuery,
     Document,
     DocumentSegment,
+    SegmentAttachmentBinding,
 )
 from models.model import UploadFile
 
@@ -58,14 +59,20 @@ def clean_dataset_task(
         )
         documents = db.session.scalars(select(Document).where(Document.dataset_id == dataset_id)).all()
         segments = db.session.scalars(select(DocumentSegment).where(DocumentSegment.dataset_id == dataset_id)).all()
+        # Use JOIN to fetch attachments with bindings in a single query
+        attachments_with_bindings = db.session.execute(
+            select(SegmentAttachmentBinding, UploadFile)
+            .join(UploadFile, UploadFile.id == SegmentAttachmentBinding.attachment_id)
+            .where(SegmentAttachmentBinding.tenant_id == tenant_id, SegmentAttachmentBinding.dataset_id == dataset_id)
+        ).all()
 
         # Enhanced validation: Check if doc_form is None, empty string, or contains only whitespace
         # This ensures all invalid doc_form values are properly handled
         if doc_form is None or (isinstance(doc_form, str) and not doc_form.strip()):
             # Use default paragraph index type for empty/invalid datasets to enable vector database cleanup
-            from core.rag.index_processor.constant.index_type import IndexType
+            from core.rag.index_processor.constant.index_type import IndexStructureType
 
-            doc_form = IndexType.PARAGRAPH_INDEX
+            doc_form = IndexStructureType.PARAGRAPH_INDEX
             logger.info(
                 click.style(f"Invalid doc_form detected, using default index type for cleanup: {doc_form}", fg="yellow")
             )
@@ -90,6 +97,7 @@ def clean_dataset_task(
 
             for document in documents:
                 db.session.delete(document)
+                # delete document file
 
             for segment in segments:
                 image_upload_file_ids = get_image_upload_file_ids(segment.content)
@@ -107,6 +115,19 @@ def clean_dataset_task(
                         )
                     db.session.delete(image_file)
                 db.session.delete(segment)
+        # delete segment attachments
+        if attachments_with_bindings:
+            for binding, attachment_file in attachments_with_bindings:
+                try:
+                    storage.delete(attachment_file.key)
+                except Exception:
+                    logger.exception(
+                        "Delete attachment_file failed when storage deleted, \
+                                        attachment_file_id: %s",
+                        binding.attachment_id,
+                    )
+                db.session.delete(attachment_file)
+                db.session.delete(binding)
 
         db.session.query(DatasetProcessRule).where(DatasetProcessRule.dataset_id == dataset_id).delete()
         db.session.query(DatasetQuery).where(DatasetQuery.dataset_id == dataset_id).delete()

api/tasks/clean_document_task.py

@@ -9,7 +9,7 @@ from core.rag.index_processor.index_processor_factory import IndexProcessorFacto
 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, DatasetMetadataBinding, DocumentSegment
+from models.dataset import Dataset, DatasetMetadataBinding, DocumentSegment, SegmentAttachmentBinding
 from models.model import UploadFile
 
 logger = logging.getLogger(__name__)
@@ -36,6 +36,16 @@ def clean_document_task(document_id: str, dataset_id: str, doc_form: str, file_i
             raise Exception("Document has no dataset")
 
         segments = db.session.scalars(select(DocumentSegment).where(DocumentSegment.document_id == document_id)).all()
+        # Use JOIN to fetch attachments with bindings in a single query
+        attachments_with_bindings = db.session.execute(
+            select(SegmentAttachmentBinding, UploadFile)
+            .join(UploadFile, UploadFile.id == SegmentAttachmentBinding.attachment_id)
+            .where(
+                SegmentAttachmentBinding.tenant_id == dataset.tenant_id,
+                SegmentAttachmentBinding.dataset_id == dataset_id,
+                SegmentAttachmentBinding.document_id == document_id,
+            )
+        ).all()
         # check segment is exist
         if segments:
             index_node_ids = [segment.index_node_id for segment in segments]
@@ -69,6 +79,19 @@ def clean_document_task(document_id: str, dataset_id: str, doc_form: str, file_i
                     logger.exception("Delete file failed when document deleted, file_id: %s", file_id)
                 db.session.delete(file)
                 db.session.commit()
+        # delete segment attachments
+        if attachments_with_bindings:
+            for binding, attachment_file in attachments_with_bindings:
+                try:
+                    storage.delete(attachment_file.key)
+                except Exception:
+                    logger.exception(
+                        "Delete attachment_file failed when storage deleted, \
+                                        attachment_file_id: %s",
+                        binding.attachment_id,
+                    )
+                db.session.delete(attachment_file)
+                db.session.delete(binding)
 
         # delete dataset metadata binding
         db.session.query(DatasetMetadataBinding).where(

api/tasks/deal_dataset_index_update_task.py

@@ -4,9 +4,10 @@ import time
 import click
 from celery import shared_task  # type: ignore
 
-from core.rag.index_processor.constant.index_type import IndexType
+from core.rag.index_processor.constant.doc_type import DocType
+from core.rag.index_processor.constant.index_type import IndexStructureType
 from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
-from core.rag.models.document import ChildDocument, Document
+from core.rag.models.document import AttachmentDocument, ChildDocument, Document
 from extensions.ext_database import db
 from models.dataset import Dataset, DocumentSegment
 from models.dataset import Document as DatasetDocument
@@ -28,7 +29,7 @@ def deal_dataset_index_update_task(dataset_id: str, action: str):
 
         if not dataset:
             raise Exception("Dataset not found")
-        index_type = dataset.doc_form or IndexType.PARAGRAPH_INDEX
+        index_type = dataset.doc_form or IndexStructureType.PARAGRAPH_INDEX
         index_processor = IndexProcessorFactory(index_type).init_index_processor()
         if action == "upgrade":
             dataset_documents = (
@@ -119,6 +120,7 @@ def deal_dataset_index_update_task(dataset_id: str, action: str):
                         )
                         if segments:
                             documents = []
+                            multimodal_documents = []
                             for segment in segments:
                                 document = Document(
                                     page_content=segment.content,
@@ -129,7 +131,7 @@ def deal_dataset_index_update_task(dataset_id: str, action: str):
                                         "dataset_id": segment.dataset_id,
                                     },
                                 )
-                                if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
+                                if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
                                     child_chunks = segment.get_child_chunks()
                                     if child_chunks:
                                         child_documents = []
@@ -145,9 +147,25 @@ def deal_dataset_index_update_task(dataset_id: str, action: str):
                                             )
                                             child_documents.append(child_document)
                                         document.children = child_documents
+                                if dataset.is_multimodal:
+                                    for attachment in segment.attachments:
+                                        multimodal_documents.append(
+                                            AttachmentDocument(
+                                                page_content=attachment["name"],
+                                                metadata={
+                                                    "doc_id": attachment["id"],
+                                                    "doc_hash": "",
+                                                    "document_id": segment.document_id,
+                                                    "dataset_id": segment.dataset_id,
+                                                    "doc_type": DocType.IMAGE,
+                                                },
+                                            )
+                                        )
                                 documents.append(document)
                             # save vector index
-                            index_processor.load(dataset, documents, with_keywords=False)
+                            index_processor.load(
+                                dataset, documents, multimodal_documents=multimodal_documents, with_keywords=False
+                            )
                         db.session.query(DatasetDocument).where(DatasetDocument.id == dataset_document.id).update(
                             {"indexing_status": "completed"}, synchronize_session=False
                         )

api/tasks/deal_dataset_vector_index_task.py

@@ -1,14 +1,14 @@
 import logging
 import time
-from typing import Literal
 
 import click
 from celery import shared_task
 from sqlalchemy import select
 
-from core.rag.index_processor.constant.index_type import IndexType
+from core.rag.index_processor.constant.doc_type import DocType
+from core.rag.index_processor.constant.index_type import IndexStructureType
 from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
-from core.rag.models.document import ChildDocument, Document
+from core.rag.models.document import AttachmentDocument, ChildDocument, Document
 from extensions.ext_database import db
 from models.dataset import Dataset, DocumentSegment
 from models.dataset import Document as DatasetDocument
@@ -17,7 +17,7 @@ logger = logging.getLogger(__name__)
 
 
 @shared_task(queue="dataset")
-def deal_dataset_vector_index_task(dataset_id: str, action: Literal["remove", "add", "update"]):
+def deal_dataset_vector_index_task(dataset_id: str, action: str):
     """
     Async deal dataset from index
     :param dataset_id: dataset_id
@@ -32,7 +32,7 @@ def deal_dataset_vector_index_task(dataset_id: str, action: Literal["remove", "a
 
         if not dataset:
             raise Exception("Dataset not found")
-        index_type = dataset.doc_form or IndexType.PARAGRAPH_INDEX
+        index_type = dataset.doc_form or IndexStructureType.PARAGRAPH_INDEX
         index_processor = IndexProcessorFactory(index_type).init_index_processor()
         if action == "remove":
             index_processor.clean(dataset, None, with_keywords=False)
@@ -119,6 +119,7 @@ def deal_dataset_vector_index_task(dataset_id: str, action: Literal["remove", "a
                         )
                         if segments:
                             documents = []
+                            multimodal_documents = []
                             for segment in segments:
                                 document = Document(
                                     page_content=segment.content,
@@ -129,7 +130,7 @@ def deal_dataset_vector_index_task(dataset_id: str, action: Literal["remove", "a
                                         "dataset_id": segment.dataset_id,
                                     },
                                 )
-                                if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
+                                if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
                                     child_chunks = segment.get_child_chunks()
                                     if child_chunks:
                                         child_documents = []
@@ -145,9 +146,25 @@ def deal_dataset_vector_index_task(dataset_id: str, action: Literal["remove", "a
                                             )
                                             child_documents.append(child_document)
                                         document.children = child_documents
+                                if dataset.is_multimodal:
+                                    for attachment in segment.attachments:
+                                        multimodal_documents.append(
+                                            AttachmentDocument(
+                                                page_content=attachment["name"],
+                                                metadata={
+                                                    "doc_id": attachment["id"],
+                                                    "doc_hash": "",
+                                                    "document_id": segment.document_id,
+                                                    "dataset_id": segment.dataset_id,
+                                                    "doc_type": DocType.IMAGE,
+                                                },
+                                            )
+                                        )
                                 documents.append(document)
                             # save vector index
-                            index_processor.load(dataset, documents, with_keywords=False)
+                            index_processor.load(
+                                dataset, documents, multimodal_documents=multimodal_documents, with_keywords=False
+                            )
                         db.session.query(DatasetDocument).where(DatasetDocument.id == dataset_document.id).update(
                             {"indexing_status": "completed"}, synchronize_session=False
                         )

api/tasks/delete_segment_from_index_task.py

@@ -6,14 +6,15 @@ from celery import shared_task
 
 from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
 from extensions.ext_database import db
-from models.dataset import Dataset, Document
+from models.dataset import Dataset, Document, SegmentAttachmentBinding
+from models.model import UploadFile
 
 logger = logging.getLogger(__name__)
 
 
 @shared_task(queue="dataset")
 def delete_segment_from_index_task(
-    index_node_ids: list, dataset_id: str, document_id: str, child_node_ids: list | None = None
+    index_node_ids: list, dataset_id: str, document_id: str, segment_ids: list, child_node_ids: list | None = None
 ):
     """
     Async Remove segment from index
@@ -49,6 +50,21 @@ def delete_segment_from_index_task(
             delete_child_chunks=True,
             precomputed_child_node_ids=child_node_ids,
         )
+        if dataset.is_multimodal:
+            # delete segment attachment binding
+            segment_attachment_bindings = (
+                db.session.query(SegmentAttachmentBinding)
+                .where(SegmentAttachmentBinding.segment_id.in_(segment_ids))
+                .all()
+            )
+            if segment_attachment_bindings:
+                attachment_ids = [binding.attachment_id for binding in segment_attachment_bindings]
+                index_processor.clean(dataset=dataset, node_ids=attachment_ids, with_keywords=False)
+                for binding in segment_attachment_bindings:
+                    db.session.delete(binding)
+                # delete upload file
+                db.session.query(UploadFile).where(UploadFile.id.in_(attachment_ids)).delete(synchronize_session=False)
+                db.session.commit()
 
         end_at = time.perf_counter()
         logger.info(click.style(f"Segment deleted from index latency: {end_at - start_at}", fg="green"))

api/tasks/disable_segments_from_index_task.py

@@ -8,7 +8,7 @@ from sqlalchemy import select
 from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
-from models.dataset import Dataset, DocumentSegment
+from models.dataset import Dataset, DocumentSegment, SegmentAttachmentBinding
 from models.dataset import Document as DatasetDocument
 
 logger = logging.getLogger(__name__)
@@ -59,6 +59,16 @@ def disable_segments_from_index_task(segment_ids: list, dataset_id: str, documen
 
     try:
         index_node_ids = [segment.index_node_id for segment in segments]
+        if dataset.is_multimodal:
+            segment_ids = [segment.id for segment in segments]
+            segment_attachment_bindings = (
+                db.session.query(SegmentAttachmentBinding)
+                .where(SegmentAttachmentBinding.segment_id.in_(segment_ids))
+                .all()
+            )
+            if segment_attachment_bindings:
+                attachment_ids = [binding.attachment_id for binding in segment_attachment_bindings]
+                index_node_ids.extend(attachment_ids)
         index_processor.clean(dataset, index_node_ids, with_keywords=True, delete_child_chunks=False)
 
         end_at = time.perf_counter()

api/tasks/enable_segment_to_index_task.py

@@ -4,9 +4,10 @@ import time
 import click
 from celery import shared_task
 
-from core.rag.index_processor.constant.index_type import IndexType
+from core.rag.index_processor.constant.doc_type import DocType
+from core.rag.index_processor.constant.index_type import IndexStructureType
 from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
-from core.rag.models.document import ChildDocument, Document
+from core.rag.models.document import AttachmentDocument, ChildDocument, Document
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from libs.datetime_utils import naive_utc_now
@@ -67,7 +68,7 @@ def enable_segment_to_index_task(segment_id: str):
             return
 
         index_processor = IndexProcessorFactory(dataset_document.doc_form).init_index_processor()
-        if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
+        if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
             child_chunks = segment.get_child_chunks()
             if child_chunks:
                 child_documents = []
@@ -83,8 +84,24 @@ def enable_segment_to_index_task(segment_id: str):
                     )
                     child_documents.append(child_document)
                 document.children = child_documents
+        multimodel_documents = []
+        if dataset.is_multimodal:
+            for attachment in segment.attachments:
+                multimodel_documents.append(
+                    AttachmentDocument(
+                        page_content=attachment["name"],
+                        metadata={
+                            "doc_id": attachment["id"],
+                            "doc_hash": "",
+                            "document_id": segment.document_id,
+                            "dataset_id": segment.dataset_id,
+                            "doc_type": DocType.IMAGE,
+                        },
+                    )
+                )
+
         # save vector index
-        index_processor.load(dataset, [document])
+        index_processor.load(dataset, [document], multimodal_documents=multimodel_documents)
 
         end_at = time.perf_counter()
         logger.info(click.style(f"Segment enabled to index: {segment.id} latency: {end_at - start_at}", fg="green"))

api/tasks/enable_segments_to_index_task.py

@@ -5,9 +5,10 @@ import click
 from celery import shared_task
 from sqlalchemy import select
 
-from core.rag.index_processor.constant.index_type import IndexType
+from core.rag.index_processor.constant.doc_type import DocType
+from core.rag.index_processor.constant.index_type import IndexStructureType
 from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
-from core.rag.models.document import ChildDocument, Document
+from core.rag.models.document import AttachmentDocument, ChildDocument, Document
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from libs.datetime_utils import naive_utc_now
@@ -60,6 +61,7 @@ def enable_segments_to_index_task(segment_ids: list, dataset_id: str, document_i
 
     try:
         documents = []
+        multimodal_documents = []
         for segment in segments:
             document = Document(
                 page_content=segment.content,
@@ -71,7 +73,7 @@ def enable_segments_to_index_task(segment_ids: list, dataset_id: str, document_i
                 },
             )
 
-            if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
+            if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
                 child_chunks = segment.get_child_chunks()
                 if child_chunks:
                     child_documents = []
@@ -87,9 +89,24 @@ def enable_segments_to_index_task(segment_ids: list, dataset_id: str, document_i
                         )
                         child_documents.append(child_document)
                     document.children = child_documents
+
+            if dataset.is_multimodal:
+                for attachment in segment.attachments:
+                    multimodal_documents.append(
+                        AttachmentDocument(
+                            page_content=attachment["name"],
+                            metadata={
+                                "doc_id": attachment["id"],
+                                "doc_hash": "",
+                                "document_id": segment.document_id,
+                                "dataset_id": segment.dataset_id,
+                                "doc_type": DocType.IMAGE,
+                            },
+                        )
+                    )
             documents.append(document)
         # save vector index
-        index_processor.load(dataset, documents)
+        index_processor.load(dataset, documents, multimodal_documents=multimodal_documents)
 
         end_at = time.perf_counter()
         logger.info(click.style(f"Segments enabled to index latency: {end_at - start_at}", fg="green"))