| 
									
										
										
										
											2024-03-28 17:02:35 +08:00
										 |  |  | import concurrent.futures | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | import datetime | 
					
						
							|  |  |  | import json | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  | import logging | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | import re | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  | import threading | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | import time | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  | import uuid | 
					
						
							| 
									
										
										
										
											2024-12-24 18:38:51 +08:00
										 |  |  | from typing import Any, Optional, cast | 
					
						
							| 
									
										
										
										
											2024-02-06 13:21:13 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  | from flask import current_app | 
					
						
							| 
									
										
										
										
											2025-05-06 11:58:49 +08:00
										 |  |  | from flask_login import current_user | 
					
						
							| 
									
										
										
										
											2024-02-06 13:21:13 +08:00
										 |  |  | from sqlalchemy.orm.exc import ObjectDeletedError | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-07-30 11:15:26 +08:00
										 |  |  | from configs import dify_config | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  | from core.entities.knowledge_entities import IndexingEstimate, PreviewDetail, QAPreviewDetail | 
					
						
							| 
									
										
										
										
											2024-01-12 12:34:01 +08:00
										 |  |  | from core.errors.error import ProviderTokenNotInitError | 
					
						
							| 
									
										
										
										
											2024-02-01 18:11:57 +08:00
										 |  |  | from core.model_manager import ModelInstance, ModelManager | 
					
						
							| 
									
										
										
										
											2024-09-04 14:41:47 +08:00
										 |  |  | from core.model_runtime.entities.model_entities import ModelType | 
					
						
							| 
									
										
										
										
											2024-11-01 15:09:22 +08:00
										 |  |  | from core.rag.cleaner.clean_processor import CleanProcessor | 
					
						
							| 
									
										
										
										
											2024-04-09 01:42:58 +08:00
										 |  |  | from core.rag.datasource.keyword.keyword_factory import Keyword | 
					
						
							| 
									
										
										
										
											2024-06-04 20:07:25 +08:00
										 |  |  | from core.rag.docstore.dataset_docstore import DatasetDocumentStore | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  | from core.rag.extractor.entity.extract_setting import ExtractSetting | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  | from core.rag.index_processor.constant.index_type import IndexType | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  | from core.rag.index_processor.index_processor_base import BaseIndexProcessor | 
					
						
							|  |  |  | from core.rag.index_processor.index_processor_factory import IndexProcessorFactory | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  | from core.rag.models.document import ChildDocument, Document | 
					
						
							| 
									
										
										
										
											2024-06-04 20:07:25 +08:00
										 |  |  | from core.rag.splitter.fixed_text_splitter import ( | 
					
						
							|  |  |  |     EnhanceRecursiveCharacterTextSplitter, | 
					
						
							|  |  |  |     FixedRecursiveCharacterTextSplitter, | 
					
						
							|  |  |  | ) | 
					
						
							|  |  |  | from core.rag.splitter.text_splitter import TextSplitter | 
					
						
							| 
									
										
										
										
											2025-02-17 17:05:13 +08:00
										 |  |  | from core.tools.utils.rag_web_reader import get_image_upload_file_ids | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | from extensions.ext_database import db | 
					
						
							|  |  |  | from extensions.ext_redis import redis_client | 
					
						
							|  |  |  | from extensions.ext_storage import storage | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  | from libs import helper | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  | from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegment | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  | from models.dataset import Document as DatasetDocument | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | from models.model import UploadFile | 
					
						
							| 
									
										
										
										
											2024-02-22 17:16:22 +08:00
										 |  |  | from services.feature_service import FeatureService | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | class IndexingRunner: | 
					
						
							| 
									
										
										
										
											2023-08-12 00:57:00 +08:00
										 |  |  |     def __init__(self): | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         self.storage = storage | 
					
						
							| 
									
										
										
										
											2024-01-02 23:42:00 +08:00
										 |  |  |         self.model_manager = ModelManager() | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-02-09 15:21:33 +08:00
										 |  |  |     def run(self, dataset_documents: list[DatasetDocument]): | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         """Run the indexing process.""" | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |         for dataset_document in dataset_documents: | 
					
						
							|  |  |  |             try: | 
					
						
							|  |  |  |                 # get dataset | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |                 dataset = db.session.query(Dataset).filter_by(id=dataset_document.dataset_id).first() | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |                 if not dataset: | 
					
						
							|  |  |  |                     raise ValueError("no dataset found") | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |                 # get the process rule | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 processing_rule = ( | 
					
						
							|  |  |  |                     db.session.query(DatasetProcessRule) | 
					
						
							|  |  |  |                     .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id) | 
					
						
							|  |  |  |                     .first() | 
					
						
							|  |  |  |                 ) | 
					
						
							| 
									
										
										
										
											2024-12-24 18:38:51 +08:00
										 |  |  |                 if not processing_rule: | 
					
						
							|  |  |  |                     raise ValueError("no process rule found") | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |                 index_type = dataset_document.doc_form | 
					
						
							|  |  |  |                 index_processor = IndexProcessorFactory(index_type).init_index_processor() | 
					
						
							|  |  |  |                 # extract | 
					
						
							|  |  |  |                 text_docs = self._extract(index_processor, dataset_document, processing_rule.to_dict()) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |                 # transform | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 documents = self._transform( | 
					
						
							|  |  |  |                     index_processor, dataset, text_docs, dataset_document.doc_language, processing_rule.to_dict() | 
					
						
							|  |  |  |                 ) | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |                 # save segment | 
					
						
							|  |  |  |                 self._load_segments(dataset, dataset_document, documents) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |                 # load | 
					
						
							|  |  |  |                 self._load( | 
					
						
							|  |  |  |                     index_processor=index_processor, | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |                     dataset=dataset, | 
					
						
							|  |  |  |                     dataset_document=dataset_document, | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                     documents=documents, | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |                 ) | 
					
						
							| 
									
										
										
										
											2024-09-11 16:40:52 +08:00
										 |  |  |             except DocumentIsPausedError: | 
					
						
							|  |  |  |                 raise DocumentIsPausedError("Document paused, document id: {}".format(dataset_document.id)) | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             except ProviderTokenNotInitError as e: | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 dataset_document.indexing_status = "error" | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |                 dataset_document.error = str(e.description) | 
					
						
							| 
									
										
										
										
											2024-11-24 13:28:46 +08:00
										 |  |  |                 dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |                 db.session.commit() | 
					
						
							| 
									
										
										
										
											2023-10-12 13:30:44 +08:00
										 |  |  |             except ObjectDeletedError: | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 logging.warning("Document deleted, document id: {}".format(dataset_document.id)) | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             except Exception as e: | 
					
						
							|  |  |  |                 logging.exception("consume document failed") | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 dataset_document.indexing_status = "error" | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |                 dataset_document.error = str(e) | 
					
						
							| 
									
										
										
										
											2024-11-24 13:28:46 +08:00
										 |  |  |                 dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |                 db.session.commit() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |     def run_in_splitting_status(self, dataset_document: DatasetDocument): | 
					
						
							|  |  |  |         """Run the indexing process when the index_status is splitting.""" | 
					
						
							|  |  |  |         try: | 
					
						
							| 
									
										
										
										
											2023-06-16 21:47:51 +08:00
										 |  |  |             # get dataset | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |             dataset = db.session.query(Dataset).filter_by(id=dataset_document.dataset_id).first() | 
					
						
							| 
									
										
										
										
											2023-06-16 21:47:51 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |             if not dataset: | 
					
						
							|  |  |  |                 raise ValueError("no dataset found") | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             # get exist document_segment list and delete | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |             document_segments = ( | 
					
						
							|  |  |  |                 db.session.query(DocumentSegment) | 
					
						
							|  |  |  |                 .filter_by(dataset_id=dataset.id, document_id=dataset_document.id) | 
					
						
							|  |  |  |                 .all() | 
					
						
							|  |  |  |             ) | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-12-07 02:24:39 +01:00
										 |  |  |             for document_segment in document_segments: | 
					
						
							|  |  |  |                 db.session.delete(document_segment) | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |                 if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX: | 
					
						
							|  |  |  |                     # delete child chunks | 
					
						
							|  |  |  |                     db.session.query(ChildChunk).filter(ChildChunk.segment_id == document_segment.id).delete() | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             db.session.commit() | 
					
						
							| 
									
										
										
										
											2023-06-16 21:47:51 +08:00
										 |  |  |             # get the process rule | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             processing_rule = ( | 
					
						
							|  |  |  |                 db.session.query(DatasetProcessRule) | 
					
						
							|  |  |  |                 .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id) | 
					
						
							|  |  |  |                 .first() | 
					
						
							|  |  |  |             ) | 
					
						
							| 
									
										
										
										
											2024-12-24 18:38:51 +08:00
										 |  |  |             if not processing_rule: | 
					
						
							|  |  |  |                 raise ValueError("no process rule found") | 
					
						
							| 
									
										
										
										
											2023-06-16 21:47:51 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |             index_type = dataset_document.doc_form | 
					
						
							|  |  |  |             index_processor = IndexProcessorFactory(index_type).init_index_processor() | 
					
						
							|  |  |  |             # extract | 
					
						
							|  |  |  |             text_docs = self._extract(index_processor, dataset_document, processing_rule.to_dict()) | 
					
						
							| 
									
										
										
										
											2023-06-16 21:47:51 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |             # transform | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             documents = self._transform( | 
					
						
							|  |  |  |                 index_processor, dataset, text_docs, dataset_document.doc_language, processing_rule.to_dict() | 
					
						
							|  |  |  |             ) | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |             # save segment | 
					
						
							|  |  |  |             self._load_segments(dataset, dataset_document, documents) | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |             # load | 
					
						
							|  |  |  |             self._load( | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 index_processor=index_processor, dataset=dataset, dataset_document=dataset_document, documents=documents | 
					
						
							| 
									
										
										
										
											2023-06-16 21:47:51 +08:00
										 |  |  |             ) | 
					
						
							| 
									
										
										
										
											2024-09-11 16:40:52 +08:00
										 |  |  |         except DocumentIsPausedError: | 
					
						
							|  |  |  |             raise DocumentIsPausedError("Document paused, document id: {}".format(dataset_document.id)) | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |         except ProviderTokenNotInitError as e: | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             dataset_document.indexing_status = "error" | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             dataset_document.error = str(e.description) | 
					
						
							| 
									
										
										
										
											2024-11-24 13:28:46 +08:00
										 |  |  |             dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             db.session.commit() | 
					
						
							|  |  |  |         except Exception as e: | 
					
						
							|  |  |  |             logging.exception("consume document failed") | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             dataset_document.indexing_status = "error" | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             dataset_document.error = str(e) | 
					
						
							| 
									
										
										
										
											2024-11-24 13:28:46 +08:00
										 |  |  |             dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             db.session.commit() | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |     def run_in_indexing_status(self, dataset_document: DatasetDocument): | 
					
						
							|  |  |  |         """Run the indexing process when the index_status is indexing.""" | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             # get dataset | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |             dataset = db.session.query(Dataset).filter_by(id=dataset_document.dataset_id).first() | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             if not dataset: | 
					
						
							|  |  |  |                 raise ValueError("no dataset found") | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             # get exist document_segment list and delete | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |             document_segments = ( | 
					
						
							|  |  |  |                 db.session.query(DocumentSegment) | 
					
						
							|  |  |  |                 .filter_by(dataset_id=dataset.id, document_id=dataset_document.id) | 
					
						
							|  |  |  |                 .all() | 
					
						
							|  |  |  |             ) | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |             documents = [] | 
					
						
							|  |  |  |             if document_segments: | 
					
						
							|  |  |  |                 for document_segment in document_segments: | 
					
						
							|  |  |  |                     # transform segment to node | 
					
						
							|  |  |  |                     if document_segment.status != "completed": | 
					
						
							|  |  |  |                         document = Document( | 
					
						
							|  |  |  |                             page_content=document_segment.content, | 
					
						
							|  |  |  |                             metadata={ | 
					
						
							|  |  |  |                                 "doc_id": document_segment.index_node_id, | 
					
						
							|  |  |  |                                 "doc_hash": document_segment.index_node_hash, | 
					
						
							|  |  |  |                                 "document_id": document_segment.document_id, | 
					
						
							|  |  |  |                                 "dataset_id": document_segment.dataset_id, | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                             }, | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |                         ) | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |                         if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX: | 
					
						
							| 
									
										
										
										
											2025-03-25 16:26:14 +08:00
										 |  |  |                             child_chunks = document_segment.get_child_chunks() | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |                             if child_chunks: | 
					
						
							|  |  |  |                                 child_documents = [] | 
					
						
							|  |  |  |                                 for child_chunk in child_chunks: | 
					
						
							|  |  |  |                                     child_document = ChildDocument( | 
					
						
							|  |  |  |                                         page_content=child_chunk.content, | 
					
						
							|  |  |  |                                         metadata={ | 
					
						
							|  |  |  |                                             "doc_id": child_chunk.index_node_id, | 
					
						
							|  |  |  |                                             "doc_hash": child_chunk.index_node_hash, | 
					
						
							|  |  |  |                                             "document_id": document_segment.document_id, | 
					
						
							|  |  |  |                                             "dataset_id": document_segment.dataset_id, | 
					
						
							|  |  |  |                                         }, | 
					
						
							|  |  |  |                                     ) | 
					
						
							|  |  |  |                                     child_documents.append(child_document) | 
					
						
							|  |  |  |                                 document.children = child_documents | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |                         documents.append(document) | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             # build index | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |             # get the process rule | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             processing_rule = ( | 
					
						
							|  |  |  |                 db.session.query(DatasetProcessRule) | 
					
						
							|  |  |  |                 .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id) | 
					
						
							|  |  |  |                 .first() | 
					
						
							|  |  |  |             ) | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |             index_type = dataset_document.doc_form | 
					
						
							| 
									
										
										
										
											2024-03-06 11:33:32 +08:00
										 |  |  |             index_processor = IndexProcessorFactory(index_type).init_index_processor() | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |             self._load( | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 index_processor=index_processor, dataset=dataset, dataset_document=dataset_document, documents=documents | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             ) | 
					
						
							| 
									
										
										
										
											2024-09-11 16:40:52 +08:00
										 |  |  |         except DocumentIsPausedError: | 
					
						
							|  |  |  |             raise DocumentIsPausedError("Document paused, document id: {}".format(dataset_document.id)) | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |         except ProviderTokenNotInitError as e: | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             dataset_document.indexing_status = "error" | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             dataset_document.error = str(e.description) | 
					
						
							| 
									
										
										
										
											2024-11-24 13:28:46 +08:00
										 |  |  |             dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             db.session.commit() | 
					
						
							|  |  |  |         except Exception as e: | 
					
						
							|  |  |  |             logging.exception("consume document failed") | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             dataset_document.indexing_status = "error" | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             dataset_document.error = str(e) | 
					
						
							| 
									
										
										
										
											2024-11-24 13:28:46 +08:00
										 |  |  |             dataset_document.stopped_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             db.session.commit() | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |     def indexing_estimate( | 
					
						
							|  |  |  |         self, | 
					
						
							|  |  |  |         tenant_id: str, | 
					
						
							|  |  |  |         extract_settings: list[ExtractSetting], | 
					
						
							|  |  |  |         tmp_processing_rule: dict, | 
					
						
							| 
									
										
										
										
											2024-10-09 14:36:43 +08:00
										 |  |  |         doc_form: Optional[str] = None, | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |         doc_language: str = "English", | 
					
						
							| 
									
										
										
										
											2024-10-09 14:36:43 +08:00
										 |  |  |         dataset_id: Optional[str] = None, | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |         indexing_technique: str = "economy", | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |     ) -> IndexingEstimate: | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         """
 | 
					
						
							|  |  |  |         Estimate the indexing for the document. | 
					
						
							|  |  |  |         """
 | 
					
						
							| 
									
										
										
										
											2024-02-22 17:16:22 +08:00
										 |  |  |         # check document limit | 
					
						
							|  |  |  |         features = FeatureService.get_features(tenant_id) | 
					
						
							|  |  |  |         if features.billing.enabled: | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |             count = len(extract_settings) | 
					
						
							| 
									
										
										
										
											2024-07-30 11:15:26 +08:00
										 |  |  |             batch_upload_limit = dify_config.BATCH_UPLOAD_LIMIT | 
					
						
							| 
									
										
										
										
											2024-02-22 17:16:22 +08:00
										 |  |  |             if count > batch_upload_limit: | 
					
						
							|  |  |  |                 raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.") | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-01-02 23:42:00 +08:00
										 |  |  |         embedding_model_instance = None | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |         if dataset_id: | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |             dataset = db.session.query(Dataset).filter_by(id=dataset_id).first() | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |             if not dataset: | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 raise ValueError("Dataset not found.") | 
					
						
							|  |  |  |             if dataset.indexing_technique == "high_quality" or indexing_technique == "high_quality": | 
					
						
							| 
									
										
										
										
											2024-01-10 20:48:16 +08:00
										 |  |  |                 if dataset.embedding_model_provider: | 
					
						
							|  |  |  |                     embedding_model_instance = self.model_manager.get_model_instance( | 
					
						
							|  |  |  |                         tenant_id=tenant_id, | 
					
						
							|  |  |  |                         provider=dataset.embedding_model_provider, | 
					
						
							|  |  |  |                         model_type=ModelType.TEXT_EMBEDDING, | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                         model=dataset.embedding_model, | 
					
						
							| 
									
										
										
										
											2024-01-10 20:48:16 +08:00
										 |  |  |                     ) | 
					
						
							|  |  |  |                 else: | 
					
						
							|  |  |  |                     embedding_model_instance = self.model_manager.get_default_model_instance( | 
					
						
							|  |  |  |                         tenant_id=tenant_id, | 
					
						
							|  |  |  |                         model_type=ModelType.TEXT_EMBEDDING, | 
					
						
							|  |  |  |                     ) | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |         else: | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             if indexing_technique == "high_quality": | 
					
						
							| 
									
										
										
										
											2024-01-02 23:42:00 +08:00
										 |  |  |                 embedding_model_instance = self.model_manager.get_default_model_instance( | 
					
						
							|  |  |  |                     tenant_id=tenant_id, | 
					
						
							|  |  |  |                     model_type=ModelType.TEXT_EMBEDDING, | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |                 ) | 
					
						
							| 
									
										
										
										
											2024-12-26 00:16:35 +08:00
										 |  |  |         preview_texts = []  # type: ignore | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-06-16 21:47:51 +08:00
										 |  |  |         total_segments = 0 | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |         index_type = doc_form | 
					
						
							|  |  |  |         index_processor = IndexProcessorFactory(index_type).init_index_processor() | 
					
						
							|  |  |  |         for extract_setting in extract_settings: | 
					
						
							|  |  |  |             # extract | 
					
						
							| 
									
										
										
										
											2023-06-16 21:47:51 +08:00
										 |  |  |             processing_rule = DatasetProcessRule( | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 mode=tmp_processing_rule["mode"], rules=json.dumps(tmp_processing_rule["rules"]) | 
					
						
							| 
									
										
										
										
											2023-06-16 21:47:51 +08:00
										 |  |  |             ) | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |             text_docs = index_processor.extract(extract_setting, process_rule_mode=tmp_processing_rule["mode"]) | 
					
						
							|  |  |  |             documents = index_processor.transform( | 
					
						
							|  |  |  |                 text_docs, | 
					
						
							|  |  |  |                 embedding_model_instance=embedding_model_instance, | 
					
						
							|  |  |  |                 process_rule=processing_rule.to_dict(), | 
					
						
							|  |  |  |                 tenant_id=current_user.current_tenant_id, | 
					
						
							|  |  |  |                 doc_language=doc_language, | 
					
						
							|  |  |  |                 preview=True, | 
					
						
							| 
									
										
										
										
											2023-06-16 21:47:51 +08:00
										 |  |  |             ) | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             total_segments += len(documents) | 
					
						
							|  |  |  |             for document in documents: | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |                 if len(preview_texts) < 10: | 
					
						
							|  |  |  |                     if doc_form and doc_form == "qa_model": | 
					
						
							|  |  |  |                         preview_detail = QAPreviewDetail( | 
					
						
							| 
									
										
										
										
											2024-12-26 00:16:35 +08:00
										 |  |  |                             question=document.page_content, answer=document.metadata.get("answer") or "" | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |                         ) | 
					
						
							|  |  |  |                         preview_texts.append(preview_detail) | 
					
						
							|  |  |  |                     else: | 
					
						
							| 
									
										
										
										
											2024-12-26 00:16:35 +08:00
										 |  |  |                         preview_detail = PreviewDetail(content=document.page_content)  # type: ignore | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |                         if document.children: | 
					
						
							| 
									
										
										
										
											2024-12-26 00:16:35 +08:00
										 |  |  |                             preview_detail.child_chunks = [child.page_content for child in document.children]  # type: ignore | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |                         preview_texts.append(preview_detail) | 
					
						
							| 
									
										
										
										
											2023-08-12 00:57:00 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-11-21 13:03:16 +08:00
										 |  |  |                 # delete image files and related db records | 
					
						
							|  |  |  |                 image_upload_file_ids = get_image_upload_file_ids(document.page_content) | 
					
						
							|  |  |  |                 for upload_file_id in image_upload_file_ids: | 
					
						
							|  |  |  |                     image_file = db.session.query(UploadFile).filter(UploadFile.id == upload_file_id).first() | 
					
						
							|  |  |  |                     try: | 
					
						
							| 
									
										
										
										
											2024-12-24 18:38:51 +08:00
										 |  |  |                         if image_file: | 
					
						
							|  |  |  |                             storage.delete(image_file.key) | 
					
						
							| 
									
										
										
										
											2024-11-21 13:03:16 +08:00
										 |  |  |                     except Exception: | 
					
						
							|  |  |  |                         logging.exception( | 
					
						
							|  |  |  |                             "Delete image_files failed while indexing_estimate, \
 | 
					
						
							|  |  |  |                                           image_upload_file_is: {}".format(upload_file_id) | 
					
						
							|  |  |  |                         ) | 
					
						
							|  |  |  |                     db.session.delete(image_file) | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |         if doc_form and doc_form == "qa_model": | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |             return IndexingEstimate(total_segments=total_segments * 20, qa_preview=preview_texts, preview=[]) | 
					
						
							| 
									
										
										
										
											2024-12-26 00:16:35 +08:00
										 |  |  |         return IndexingEstimate(total_segments=total_segments, preview=preview_texts)  # type: ignore | 
					
						
							| 
									
										
										
										
											2023-06-16 21:47:51 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |     def _extract( | 
					
						
							|  |  |  |         self, index_processor: BaseIndexProcessor, dataset_document: DatasetDocument, process_rule: dict | 
					
						
							|  |  |  |     ) -> list[Document]: | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         # load file | 
					
						
							| 
									
										
										
										
											2024-09-13 22:42:08 +08:00
										 |  |  |         if dataset_document.data_source_type not in {"upload_file", "notion_import", "website_crawl"}: | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |             return [] | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |         data_source_info = dataset_document.data_source_info_dict | 
					
						
							| 
									
										
										
										
											2023-06-16 21:47:51 +08:00
										 |  |  |         text_docs = [] | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |         if dataset_document.data_source_type == "upload_file": | 
					
						
							|  |  |  |             if not data_source_info or "upload_file_id" not in data_source_info: | 
					
						
							| 
									
										
										
										
											2023-06-16 21:47:51 +08:00
										 |  |  |                 raise ValueError("no upload file found") | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             file_detail = ( | 
					
						
							|  |  |  |                 db.session.query(UploadFile).filter(UploadFile.id == data_source_info["upload_file_id"]).one_or_none() | 
					
						
							|  |  |  |             ) | 
					
						
							| 
									
										
										
										
											2023-06-16 21:47:51 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-08-30 11:14:16 +08:00
										 |  |  |             if file_detail: | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |                 extract_setting = ExtractSetting( | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                     datasource_type="upload_file", upload_file=file_detail, document_model=dataset_document.doc_form | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |                 ) | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"]) | 
					
						
							|  |  |  |         elif dataset_document.data_source_type == "notion_import": | 
					
						
							|  |  |  |             if ( | 
					
						
							|  |  |  |                 not data_source_info | 
					
						
							|  |  |  |                 or "notion_workspace_id" not in data_source_info | 
					
						
							|  |  |  |                 or "notion_page_id" not in data_source_info | 
					
						
							|  |  |  |             ): | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |                 raise ValueError("no notion import info found") | 
					
						
							|  |  |  |             extract_setting = ExtractSetting( | 
					
						
							|  |  |  |                 datasource_type="notion_import", | 
					
						
							|  |  |  |                 notion_info={ | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                     "notion_workspace_id": data_source_info["notion_workspace_id"], | 
					
						
							|  |  |  |                     "notion_obj_id": data_source_info["notion_page_id"], | 
					
						
							|  |  |  |                     "notion_page_type": data_source_info["type"], | 
					
						
							| 
									
										
										
										
											2024-02-27 11:39:05 +08:00
										 |  |  |                     "document": dataset_document, | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                     "tenant_id": dataset_document.tenant_id, | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |                 }, | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 document_model=dataset_document.doc_form, | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |             ) | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"]) | 
					
						
							|  |  |  |         elif dataset_document.data_source_type == "website_crawl": | 
					
						
							|  |  |  |             if ( | 
					
						
							|  |  |  |                 not data_source_info | 
					
						
							|  |  |  |                 or "provider" not in data_source_info | 
					
						
							|  |  |  |                 or "url" not in data_source_info | 
					
						
							|  |  |  |                 or "job_id" not in data_source_info | 
					
						
							|  |  |  |             ): | 
					
						
							| 
									
										
										
										
											2024-06-15 02:46:02 +08:00
										 |  |  |                 raise ValueError("no website import info found") | 
					
						
							|  |  |  |             extract_setting = ExtractSetting( | 
					
						
							|  |  |  |                 datasource_type="website_crawl", | 
					
						
							|  |  |  |                 website_info={ | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                     "provider": data_source_info["provider"], | 
					
						
							|  |  |  |                     "job_id": data_source_info["job_id"], | 
					
						
							| 
									
										
										
										
											2024-06-15 02:46:02 +08:00
										 |  |  |                     "tenant_id": dataset_document.tenant_id, | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                     "url": data_source_info["url"], | 
					
						
							|  |  |  |                     "mode": data_source_info["mode"], | 
					
						
							|  |  |  |                     "only_main_content": data_source_info["only_main_content"], | 
					
						
							| 
									
										
										
										
											2024-06-15 02:46:02 +08:00
										 |  |  |                 }, | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 document_model=dataset_document.doc_form, | 
					
						
							| 
									
										
										
										
											2024-06-15 02:46:02 +08:00
										 |  |  |             ) | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"]) | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         # update document status to splitting | 
					
						
							|  |  |  |         self._update_document_index_status( | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             document_id=dataset_document.id, | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |             after_indexing_status="splitting", | 
					
						
							|  |  |  |             extra_update_params={ | 
					
						
							| 
									
										
										
										
											2024-06-27 11:21:31 +08:00
										 |  |  |                 DatasetDocument.word_count: sum(len(text_doc.page_content) for text_doc in text_docs), | 
					
						
							| 
									
										
										
										
											2024-11-24 13:28:46 +08:00
										 |  |  |                 DatasetDocument.parsing_completed_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None), | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             }, | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         ) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |         # replace doc id to document model id | 
					
						
							| 
									
										
										
										
											2024-02-09 15:21:33 +08:00
										 |  |  |         text_docs = cast(list[Document], text_docs) | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         for text_doc in text_docs: | 
					
						
							| 
									
										
										
										
											2024-12-24 18:38:51 +08:00
										 |  |  |             if text_doc.metadata is not None: | 
					
						
							|  |  |  |                 text_doc.metadata["document_id"] = dataset_document.id | 
					
						
							|  |  |  |                 text_doc.metadata["dataset_id"] = dataset_document.dataset_id | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |         return text_docs | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-08-26 20:45:03 +09:00
										 |  |  |     @staticmethod | 
					
						
							|  |  |  |     def filter_string(text): | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |         text = re.sub(r"<\|", "<", text) | 
					
						
							|  |  |  |         text = re.sub(r"\|>", ">", text) | 
					
						
							|  |  |  |         text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F\xEF\xBF\xBE]", "", text) | 
					
						
							| 
									
										
										
										
											2024-01-15 16:52:18 +08:00
										 |  |  |         # Unicode  U+FFFE | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |         text = re.sub("\ufffe", "", text) | 
					
						
							| 
									
										
										
										
											2023-06-28 14:58:40 +08:00
										 |  |  |         return text | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-08-26 20:45:03 +09:00
										 |  |  |     @staticmethod | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |     def _get_splitter( | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         processing_rule_mode: str, | 
					
						
							|  |  |  |         max_tokens: int, | 
					
						
							|  |  |  |         chunk_overlap: int, | 
					
						
							|  |  |  |         separator: str, | 
					
						
							|  |  |  |         embedding_model_instance: Optional[ModelInstance], | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |     ) -> TextSplitter: | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         """
 | 
					
						
							|  |  |  |         Get the NodeParser object according to the processing rule. | 
					
						
							|  |  |  |         """
 | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         if processing_rule_mode in ["custom", "hierarchical"]: | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |             # The user-defined segmentation rule | 
					
						
							| 
									
										
										
										
											2024-07-30 11:15:26 +08:00
										 |  |  |             max_segmentation_tokens_length = dify_config.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |             if max_tokens < 50 or max_tokens > max_segmentation_tokens_length: | 
					
						
							| 
									
										
										
										
											2024-05-20 13:20:27 +08:00
										 |  |  |                 raise ValueError(f"Custom segment length should be between 50 and {max_segmentation_tokens_length}.") | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-05-16 12:57:25 +08:00
										 |  |  |             if separator: | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 separator = separator.replace("\\n", "\n") | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-01-12 18:45:34 +08:00
										 |  |  |             character_splitter = FixedRecursiveCharacterTextSplitter.from_encoder( | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |                 chunk_size=max_tokens, | 
					
						
							| 
									
										
										
										
											2024-03-11 15:36:56 +08:00
										 |  |  |                 chunk_overlap=chunk_overlap, | 
					
						
							| 
									
										
										
										
											2023-05-16 12:57:25 +08:00
										 |  |  |                 fixed_separator=separator, | 
					
						
							| 
									
										
										
										
											2024-05-23 18:05:23 +08:00
										 |  |  |                 separators=["\n\n", "。", ". ", " ", ""], | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 embedding_model_instance=embedding_model_instance, | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |             ) | 
					
						
							|  |  |  |         else: | 
					
						
							|  |  |  |             # Automatic segmentation | 
					
						
							| 
									
										
										
										
											2024-12-24 18:38:51 +08:00
										 |  |  |             automatic_rules: dict[str, Any] = dict(DatasetProcessRule.AUTOMATIC_RULES["segmentation"]) | 
					
						
							| 
									
										
										
										
											2024-01-12 18:45:34 +08:00
										 |  |  |             character_splitter = EnhanceRecursiveCharacterTextSplitter.from_encoder( | 
					
						
							| 
									
										
										
										
											2024-12-24 18:38:51 +08:00
										 |  |  |                 chunk_size=automatic_rules["max_tokens"], | 
					
						
							|  |  |  |                 chunk_overlap=automatic_rules["chunk_overlap"], | 
					
						
							| 
									
										
										
										
											2024-05-23 18:05:23 +08:00
										 |  |  |                 separators=["\n\n", "。", ". ", " ", ""], | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 embedding_model_instance=embedding_model_instance, | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |             ) | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-12-26 00:16:35 +08:00
										 |  |  |         return character_splitter  # type: ignore | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |     def _split_to_documents_for_estimate( | 
					
						
							|  |  |  |         self, text_docs: list[Document], splitter: TextSplitter, processing_rule: DatasetProcessRule | 
					
						
							|  |  |  |     ) -> list[Document]: | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         """
 | 
					
						
							|  |  |  |         Split the text documents into nodes. | 
					
						
							|  |  |  |         """
 | 
					
						
							| 
									
										
										
										
											2024-12-24 18:38:51 +08:00
										 |  |  |         all_documents: list[Document] = [] | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         for text_doc in text_docs: | 
					
						
							|  |  |  |             # document clean | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             document_text = self._document_clean(text_doc.page_content, processing_rule) | 
					
						
							|  |  |  |             text_doc.page_content = document_text | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |             # parse document to nodes | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             documents = splitter.split_documents([text_doc]) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |             split_documents = [] | 
					
						
							|  |  |  |             for document in documents: | 
					
						
							|  |  |  |                 if document.page_content is None or not document.page_content.strip(): | 
					
						
							|  |  |  |                     continue | 
					
						
							| 
									
										
										
										
											2024-12-24 18:38:51 +08:00
										 |  |  |                 if document.metadata is not None: | 
					
						
							|  |  |  |                     doc_id = str(uuid.uuid4()) | 
					
						
							|  |  |  |                     hash = helper.generate_text_hash(document.page_content) | 
					
						
							|  |  |  |                     document.metadata["doc_id"] = doc_id | 
					
						
							|  |  |  |                     document.metadata["doc_hash"] = hash | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |                 split_documents.append(document) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |             all_documents.extend(split_documents) | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |         return all_documents | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-08-26 20:45:03 +09:00
										 |  |  |     @staticmethod | 
					
						
							|  |  |  |     def _document_clean(text: str, processing_rule: DatasetProcessRule) -> str: | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         """
 | 
					
						
							|  |  |  |         Clean the document text according to the processing rules. | 
					
						
							|  |  |  |         """
 | 
					
						
							|  |  |  |         if processing_rule.mode == "automatic": | 
					
						
							|  |  |  |             rules = DatasetProcessRule.AUTOMATIC_RULES | 
					
						
							|  |  |  |         else: | 
					
						
							|  |  |  |             rules = json.loads(processing_rule.rules) if processing_rule.rules else {} | 
					
						
							| 
									
										
										
										
											2024-11-05 14:47:15 +08:00
										 |  |  |         document_text = CleanProcessor.clean(text, {"rules": rules}) | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-11-01 15:09:22 +08:00
										 |  |  |         return document_text | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-08-26 20:45:03 +09:00
										 |  |  |     @staticmethod | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |     def format_split_text(text: str) -> list[QAPreviewDetail]: | 
					
						
							| 
									
										
										
										
											2024-01-12 18:45:34 +08:00
										 |  |  |         regex = r"Q\d+:\s*(.*?)\s*A\d+:\s*([\s\S]*?)(?=Q\d+:|$)" | 
					
						
							| 
									
										
										
										
											2023-12-11 15:53:37 +08:00
										 |  |  |         matches = re.findall(regex, text, re.UNICODE) | 
					
						
							| 
									
										
										
										
											2023-11-13 19:05:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         return [QAPreviewDetail(question=q, answer=re.sub(r"\n\s*", "\n", a.strip())) for q, a in matches if q and a] | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |     def _load( | 
					
						
							|  |  |  |         self, | 
					
						
							|  |  |  |         index_processor: BaseIndexProcessor, | 
					
						
							|  |  |  |         dataset: Dataset, | 
					
						
							|  |  |  |         dataset_document: DatasetDocument, | 
					
						
							|  |  |  |         documents: list[Document], | 
					
						
							|  |  |  |     ) -> None: | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         """
 | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |         insert index and update document/segment status to completed | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         """
 | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-01-02 23:42:00 +08:00
										 |  |  |         embedding_model_instance = None | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |         if dataset.indexing_technique == "high_quality": | 
					
						
							| 
									
										
										
										
											2024-01-02 23:42:00 +08:00
										 |  |  |             embedding_model_instance = self.model_manager.get_model_instance( | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |                 tenant_id=dataset.tenant_id, | 
					
						
							| 
									
										
										
										
											2024-01-02 23:42:00 +08:00
										 |  |  |                 provider=dataset.embedding_model_provider, | 
					
						
							|  |  |  |                 model_type=ModelType.TEXT_EMBEDDING, | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 model=dataset.embedding_model, | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |             ) | 
					
						
							| 
									
										
										
										
											2023-08-12 00:57:00 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         # chunk nodes by chunk size | 
					
						
							|  |  |  |         indexing_start_at = time.perf_counter() | 
					
						
							|  |  |  |         tokens = 0 | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         if dataset_document.doc_form != IndexType.PARENT_CHILD_INDEX: | 
					
						
							|  |  |  |             # create keyword index | 
					
						
							|  |  |  |             create_keyword_thread = threading.Thread( | 
					
						
							|  |  |  |                 target=self._process_keyword_index, | 
					
						
							| 
									
										
										
										
											2024-12-26 00:16:35 +08:00
										 |  |  |                 args=(current_app._get_current_object(), dataset.id, dataset_document.id, documents),  # type: ignore | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |             ) | 
					
						
							|  |  |  |             create_keyword_thread.start() | 
					
						
							| 
									
										
										
										
											2024-01-02 23:42:00 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2025-01-09 15:16:41 +08:00
										 |  |  |         max_workers = 10 | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |         if dataset.indexing_technique == "high_quality": | 
					
						
							| 
									
										
										
										
											2025-01-09 15:16:41 +08:00
										 |  |  |             with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: | 
					
						
							| 
									
										
										
										
											2024-04-09 01:42:58 +08:00
										 |  |  |                 futures = [] | 
					
						
							| 
									
										
										
										
											2025-01-09 15:16:41 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |                 # Distribute documents into multiple groups based on the hash values of page_content | 
					
						
							|  |  |  |                 # This is done to prevent multiple threads from processing the same document, | 
					
						
							|  |  |  |                 # Thereby avoiding potential database insertion deadlocks | 
					
						
							|  |  |  |                 document_groups: list[list[Document]] = [[] for _ in range(max_workers)] | 
					
						
							|  |  |  |                 for document in documents: | 
					
						
							|  |  |  |                     hash = helper.generate_text_hash(document.page_content) | 
					
						
							|  |  |  |                     group_index = int(hash, 16) % max_workers | 
					
						
							|  |  |  |                     document_groups[group_index].append(document) | 
					
						
							|  |  |  |                 for chunk_documents in document_groups: | 
					
						
							|  |  |  |                     if len(chunk_documents) == 0: | 
					
						
							|  |  |  |                         continue | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                     futures.append( | 
					
						
							|  |  |  |                         executor.submit( | 
					
						
							|  |  |  |                             self._process_chunk, | 
					
						
							| 
									
										
										
										
											2024-12-24 18:38:51 +08:00
										 |  |  |                             current_app._get_current_object(),  # type: ignore | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                             index_processor, | 
					
						
							|  |  |  |                             chunk_documents, | 
					
						
							|  |  |  |                             dataset, | 
					
						
							|  |  |  |                             dataset_document, | 
					
						
							|  |  |  |                             embedding_model_instance, | 
					
						
							|  |  |  |                         ) | 
					
						
							|  |  |  |                     ) | 
					
						
							| 
									
										
										
										
											2024-04-09 01:42:58 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |                 for future in futures: | 
					
						
							|  |  |  |                     tokens += future.result() | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         if dataset_document.doc_form != IndexType.PARENT_CHILD_INDEX: | 
					
						
							|  |  |  |             create_keyword_thread.join() | 
					
						
							| 
									
										
										
										
											2024-03-28 17:02:35 +08:00
										 |  |  |         indexing_end_at = time.perf_counter() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |         # update document status to completed | 
					
						
							|  |  |  |         self._update_document_index_status( | 
					
						
							|  |  |  |             document_id=dataset_document.id, | 
					
						
							|  |  |  |             after_indexing_status="completed", | 
					
						
							|  |  |  |             extra_update_params={ | 
					
						
							|  |  |  |                 DatasetDocument.tokens: tokens, | 
					
						
							| 
									
										
										
										
											2024-11-24 13:28:46 +08:00
										 |  |  |                 DatasetDocument.completed_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None), | 
					
						
							| 
									
										
										
										
											2024-03-28 17:02:35 +08:00
										 |  |  |                 DatasetDocument.indexing_latency: indexing_end_at - indexing_start_at, | 
					
						
							| 
									
										
										
										
											2024-08-20 19:30:57 +08:00
										 |  |  |                 DatasetDocument.error: None, | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             }, | 
					
						
							| 
									
										
										
										
											2024-03-28 17:02:35 +08:00
										 |  |  |         ) | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-08-26 20:45:03 +09:00
										 |  |  |     @staticmethod | 
					
						
							|  |  |  |     def _process_keyword_index(flask_app, dataset_id, document_id, documents): | 
					
						
							| 
									
										
										
										
											2024-04-09 01:42:58 +08:00
										 |  |  |         with flask_app.app_context(): | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |             dataset = db.session.query(Dataset).filter_by(id=dataset_id).first() | 
					
						
							| 
									
										
										
										
											2024-04-09 14:34:51 +08:00
										 |  |  |             if not dataset: | 
					
						
							|  |  |  |                 raise ValueError("no dataset found") | 
					
						
							| 
									
										
										
										
											2024-04-09 01:42:58 +08:00
										 |  |  |             keyword = Keyword(dataset) | 
					
						
							|  |  |  |             keyword.create(documents) | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             if dataset.indexing_technique != "high_quality": | 
					
						
							|  |  |  |                 document_ids = [document.metadata["doc_id"] for document in documents] | 
					
						
							| 
									
										
										
										
											2024-04-09 01:42:58 +08:00
										 |  |  |                 db.session.query(DocumentSegment).filter( | 
					
						
							| 
									
										
										
										
											2024-04-09 14:34:51 +08:00
										 |  |  |                     DocumentSegment.document_id == document_id, | 
					
						
							| 
									
										
										
										
											2024-08-27 10:25:24 +08:00
										 |  |  |                     DocumentSegment.dataset_id == dataset_id, | 
					
						
							| 
									
										
										
										
											2024-04-09 01:42:58 +08:00
										 |  |  |                     DocumentSegment.index_node_id.in_(document_ids), | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                     DocumentSegment.status == "indexing", | 
					
						
							|  |  |  |                 ).update( | 
					
						
							|  |  |  |                     { | 
					
						
							|  |  |  |                         DocumentSegment.status: "completed", | 
					
						
							|  |  |  |                         DocumentSegment.enabled: True, | 
					
						
							| 
									
										
										
										
											2024-11-24 13:28:46 +08:00
										 |  |  |                         DocumentSegment.completed_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None), | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                     } | 
					
						
							|  |  |  |                 ) | 
					
						
							| 
									
										
										
										
											2024-04-09 01:42:58 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |                 db.session.commit() | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |     def _process_chunk( | 
					
						
							|  |  |  |         self, flask_app, index_processor, chunk_documents, dataset, dataset_document, embedding_model_instance | 
					
						
							|  |  |  |     ): | 
					
						
							| 
									
										
										
										
											2024-03-28 17:02:35 +08:00
										 |  |  |         with flask_app.app_context(): | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |             # check document is paused | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |             self._check_document_paused_status(dataset_document.id) | 
					
						
							| 
									
										
										
										
											2024-03-28 17:02:35 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |             tokens = 0 | 
					
						
							| 
									
										
										
										
											2024-07-09 23:04:44 +08:00
										 |  |  |             if embedding_model_instance: | 
					
						
							| 
									
										
										
										
											2025-02-17 17:05:13 +08:00
										 |  |  |                 page_content_list = [document.page_content for document in chunk_documents] | 
					
						
							|  |  |  |                 tokens += sum(embedding_model_instance.get_text_embedding_num_tokens(page_content_list)) | 
					
						
							| 
									
										
										
										
											2024-03-28 17:02:35 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |             # load index | 
					
						
							| 
									
										
										
										
											2024-04-09 01:42:58 +08:00
										 |  |  |             index_processor.load(dataset, chunk_documents, with_keywords=False) | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             document_ids = [document.metadata["doc_id"] for document in chunk_documents] | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |             db.session.query(DocumentSegment).filter( | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |                 DocumentSegment.document_id == dataset_document.id, | 
					
						
							| 
									
										
										
										
											2024-08-27 10:25:24 +08:00
										 |  |  |                 DocumentSegment.dataset_id == dataset.id, | 
					
						
							| 
									
										
										
										
											2023-06-25 16:49:14 +08:00
										 |  |  |                 DocumentSegment.index_node_id.in_(document_ids), | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 DocumentSegment.status == "indexing", | 
					
						
							|  |  |  |             ).update( | 
					
						
							|  |  |  |                 { | 
					
						
							|  |  |  |                     DocumentSegment.status: "completed", | 
					
						
							|  |  |  |                     DocumentSegment.enabled: True, | 
					
						
							| 
									
										
										
										
											2024-11-24 13:28:46 +08:00
										 |  |  |                     DocumentSegment.completed_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None), | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                 } | 
					
						
							|  |  |  |             ) | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |             db.session.commit() | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-03-28 17:02:35 +08:00
										 |  |  |             return tokens | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-08-26 20:45:03 +09:00
										 |  |  |     @staticmethod | 
					
						
							|  |  |  |     def _check_document_paused_status(document_id: str): | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |         indexing_cache_key = "document_{}_is_paused".format(document_id) | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         result = redis_client.get(indexing_cache_key) | 
					
						
							|  |  |  |         if result: | 
					
						
							| 
									
										
										
										
											2024-09-11 16:40:52 +08:00
										 |  |  |             raise DocumentIsPausedError() | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-08-26 20:45:03 +09:00
										 |  |  |     @staticmethod | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |     def _update_document_index_status( | 
					
						
							|  |  |  |         document_id: str, after_indexing_status: str, extra_update_params: Optional[dict] = None | 
					
						
							|  |  |  |     ) -> None: | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         """
 | 
					
						
							|  |  |  |         Update the document indexing status. | 
					
						
							|  |  |  |         """
 | 
					
						
							| 
									
										
										
										
											2025-05-13 09:13:12 +08:00
										 |  |  |         count = db.session.query(DatasetDocument).filter_by(id=document_id, is_paused=True).count() | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         if count > 0: | 
					
						
							| 
									
										
										
										
											2024-09-11 16:40:52 +08:00
										 |  |  |             raise DocumentIsPausedError() | 
					
						
							| 
									
										
										
										
											2025-05-13 09:13:12 +08:00
										 |  |  |         document = db.session.query(DatasetDocument).filter_by(id=document_id).first() | 
					
						
							| 
									
										
										
										
											2023-10-12 13:30:44 +08:00
										 |  |  |         if not document: | 
					
						
							| 
									
										
										
										
											2024-09-11 16:40:52 +08:00
										 |  |  |             raise DocumentIsDeletedPausedError() | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |         update_params = {DatasetDocument.indexing_status: after_indexing_status} | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |         if extra_update_params: | 
					
						
							|  |  |  |             update_params.update(extra_update_params) | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2025-05-13 09:13:12 +08:00
										 |  |  |         db.session.query(DatasetDocument).filter_by(id=document_id).update(update_params) | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         db.session.commit() | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-08-26 20:45:03 +09:00
										 |  |  |     @staticmethod | 
					
						
							|  |  |  |     def _update_segments_by_document(dataset_document_id: str, update_params: dict) -> None: | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         """
 | 
					
						
							|  |  |  |         Update the document segment by document id. | 
					
						
							|  |  |  |         """
 | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |         db.session.query(DocumentSegment).filter_by(document_id=dataset_document_id).update(update_params) | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         db.session.commit() | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |     def _transform( | 
					
						
							|  |  |  |         self, | 
					
						
							|  |  |  |         index_processor: BaseIndexProcessor, | 
					
						
							|  |  |  |         dataset: Dataset, | 
					
						
							|  |  |  |         text_docs: list[Document], | 
					
						
							|  |  |  |         doc_language: str, | 
					
						
							|  |  |  |         process_rule: dict, | 
					
						
							|  |  |  |     ) -> list[Document]: | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |         # get embedding model instance | 
					
						
							|  |  |  |         embedding_model_instance = None | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |         if dataset.indexing_technique == "high_quality": | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |             if dataset.embedding_model_provider: | 
					
						
							|  |  |  |                 embedding_model_instance = self.model_manager.get_model_instance( | 
					
						
							|  |  |  |                     tenant_id=dataset.tenant_id, | 
					
						
							|  |  |  |                     provider=dataset.embedding_model_provider, | 
					
						
							|  |  |  |                     model_type=ModelType.TEXT_EMBEDDING, | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |                     model=dataset.embedding_model, | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |                 ) | 
					
						
							|  |  |  |             else: | 
					
						
							|  |  |  |                 embedding_model_instance = self.model_manager.get_default_model_instance( | 
					
						
							|  |  |  |                     tenant_id=dataset.tenant_id, | 
					
						
							|  |  |  |                     model_type=ModelType.TEXT_EMBEDDING, | 
					
						
							|  |  |  |                 ) | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |         documents = index_processor.transform( | 
					
						
							|  |  |  |             text_docs, | 
					
						
							|  |  |  |             embedding_model_instance=embedding_model_instance, | 
					
						
							|  |  |  |             process_rule=process_rule, | 
					
						
							|  |  |  |             tenant_id=dataset.tenant_id, | 
					
						
							|  |  |  |             doc_language=doc_language, | 
					
						
							|  |  |  |         ) | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |         return documents | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |     def _load_segments(self, dataset, dataset_document, documents): | 
					
						
							|  |  |  |         # save node to document segment | 
					
						
							|  |  |  |         doc_store = DatasetDocumentStore( | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             dataset=dataset, user_id=dataset_document.created_by, document_id=dataset_document.id | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |         ) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |         # add document segments | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         doc_store.add_documents(docs=documents, save_child=dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX) | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |         # update document status to indexing | 
					
						
							| 
									
										
										
										
											2024-11-24 13:28:46 +08:00
										 |  |  |         cur_time = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |         self._update_document_index_status( | 
					
						
							|  |  |  |             document_id=dataset_document.id, | 
					
						
							|  |  |  |             after_indexing_status="indexing", | 
					
						
							|  |  |  |             extra_update_params={ | 
					
						
							|  |  |  |                 DatasetDocument.cleaning_completed_at: cur_time, | 
					
						
							|  |  |  |                 DatasetDocument.splitting_completed_at: cur_time, | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             }, | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |         ) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |         # update segment status to indexing | 
					
						
							|  |  |  |         self._update_segments_by_document( | 
					
						
							|  |  |  |             dataset_document_id=dataset_document.id, | 
					
						
							|  |  |  |             update_params={ | 
					
						
							|  |  |  |                 DocumentSegment.status: "indexing", | 
					
						
							| 
									
										
										
										
											2024-11-24 13:28:46 +08:00
										 |  |  |                 DocumentSegment.indexing_at: datetime.datetime.now(datetime.UTC).replace(tzinfo=None), | 
					
						
							| 
									
										
										
										
											2024-09-10 17:00:20 +08:00
										 |  |  |             }, | 
					
						
							| 
									
										
										
										
											2024-02-22 23:31:57 +08:00
										 |  |  |         ) | 
					
						
							|  |  |  |         pass | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-09-11 16:40:52 +08:00
										 |  |  | class DocumentIsPausedError(Exception): | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |     pass | 
					
						
							| 
									
										
										
										
											2023-10-12 13:30:44 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-09-11 16:40:52 +08:00
										 |  |  | class DocumentIsDeletedPausedError(Exception): | 
					
						
							| 
									
										
										
										
											2023-10-12 13:30:44 +08:00
										 |  |  |     pass |