update dataset embedding model, update document status to be indexing (#7145)

api/tasks/deal_dataset_vector_index_task.py

@@ -42,13 +42,20 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
             ).all()
 
             if dataset_documents:
-                documents = []
+                dataset_documents_ids = [doc.id for doc in dataset_documents]
+                db.session.query(DatasetDocument).filter(DatasetDocument.id.in_(dataset_documents_ids)) \
+                    .update({"indexing_status": "indexing"}, synchronize_session=False)
+                db.session.commit()
+
                 for dataset_document in dataset_documents:
-                    # delete from vector index
+                    try:
+                        # add from vector index
                         segments = db.session.query(DocumentSegment).filter(
                             DocumentSegment.document_id == dataset_document.id,
                             DocumentSegment.enabled == True
                         ) .order_by(DocumentSegment.position.asc()).all()
+                        if segments:
+                            documents = []
                             for segment in segments:
                                 document = Document(
                                     page_content=segment.content,
@@ -61,12 +68,16 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
                                 )
 
                                 documents.append(document)
-
                             # save vector index
                             index_processor.load(dataset, documents, with_keywords=False)
+                        db.session.query(DatasetDocument).filter(DatasetDocument.id == dataset_document.id) \
+                            .update({"indexing_status": "completed"}, synchronize_session=False)
+                        db.session.commit()
+                    except Exception as e:
+                        db.session.query(DatasetDocument).filter(DatasetDocument.id == dataset_document.id) \
+                            .update({"indexing_status": "error", "error": str(e)}, synchronize_session=False)
+                        db.session.commit()
         elif action == 'update':
-            # clean index
-            index_processor.clean(dataset, None, with_keywords=False)
             dataset_documents = db.session.query(DatasetDocument).filter(
                 DatasetDocument.dataset_id == dataset_id,
                 DatasetDocument.indexing_status == 'completed',
@@ -75,13 +86,24 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
             ).all()
             # add new index
             if dataset_documents:
-                documents = []
+                # update document status
+                dataset_documents_ids = [doc.id for doc in dataset_documents]
+                db.session.query(DatasetDocument).filter(DatasetDocument.id.in_(dataset_documents_ids)) \
+                    .update({"indexing_status": "indexing"}, synchronize_session=False)
+                db.session.commit()
+
+                # clean index
+                index_processor.clean(dataset, None, with_keywords=False)
+
                 for dataset_document in dataset_documents:
-                    # delete from vector index
+                    # update from vector index
+                    try:
                         segments = db.session.query(DocumentSegment).filter(
                             DocumentSegment.document_id == dataset_document.id,
                             DocumentSegment.enabled == True
                         ).order_by(DocumentSegment.position.asc()).all()
+                        if segments:
+                            documents = []
                             for segment in segments:
                                 document = Document(
                                     page_content=segment.content,
@@ -94,9 +116,16 @@ def deal_dataset_vector_index_task(dataset_id: str, action: str):
                                 )
 
                                 documents.append(document)
-
                             # save vector index
                             index_processor.load(dataset, documents, with_keywords=False)
+                        db.session.query(DatasetDocument).filter(DatasetDocument.id == dataset_document.id) \
+                            .update({"indexing_status": "completed"}, synchronize_session=False)
+                        db.session.commit()
+                    except Exception as e:
+                        db.session.query(DatasetDocument).filter(DatasetDocument.id == dataset_document.id) \
+                            .update({"indexing_status": "error", "error": str(e)}, synchronize_session=False)
+                        db.session.commit()
+
 
         end_at = time.perf_counter()
         logging.info(