- import concurrent.futures
 - import json
 - import logging
 - import re
 - import threading
 - import time
 - import uuid
 - from typing import Any, Optional
 - 
 - from flask import current_app
 - from sqlalchemy import select
 - from sqlalchemy.orm.exc import ObjectDeletedError
 - 
 - from configs import dify_config
 - from core.entities.knowledge_entities import IndexingEstimate, PreviewDetail, QAPreviewDetail
 - from core.errors.error import ProviderTokenNotInitError
 - from core.model_manager import ModelInstance, ModelManager
 - from core.model_runtime.entities.model_entities import ModelType
 - from core.rag.cleaner.clean_processor import CleanProcessor
 - from core.rag.datasource.keyword.keyword_factory import Keyword
 - from core.rag.docstore.dataset_docstore import DatasetDocumentStore
 - from core.rag.extractor.entity.datasource_type import DatasourceType
 - from core.rag.extractor.entity.extract_setting import ExtractSetting
 - from core.rag.index_processor.constant.index_type import IndexType
 - from core.rag.index_processor.index_processor_base import BaseIndexProcessor
 - from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
 - from core.rag.models.document import ChildDocument, Document
 - from core.rag.splitter.fixed_text_splitter import (
 -     EnhanceRecursiveCharacterTextSplitter,
 -     FixedRecursiveCharacterTextSplitter,
 - )
 - from core.rag.splitter.text_splitter import TextSplitter
 - from core.tools.utils.web_reader_tool import get_image_upload_file_ids
 - from extensions.ext_database import db
 - from extensions.ext_redis import redis_client
 - from extensions.ext_storage import storage
 - from libs import helper
 - from libs.datetime_utils import naive_utc_now
 - from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegment
 - from models.dataset import Document as DatasetDocument
 - from models.model import UploadFile
 - from services.feature_service import FeatureService
 - 
 - logger = logging.getLogger(__name__)
 - 
 - 
 - class IndexingRunner:
 -     def __init__(self):
 -         self.storage = storage
 -         self.model_manager = ModelManager()
 - 
 -     def run(self, dataset_documents: list[DatasetDocument]):
 -         """Run the indexing process."""
 -         for dataset_document in dataset_documents:
 -             try:
 -                 # get dataset
 -                 dataset = db.session.query(Dataset).filter_by(id=dataset_document.dataset_id).first()
 - 
 -                 if not dataset:
 -                     raise ValueError("no dataset found")
 -                 # get the process rule
 -                 stmt = select(DatasetProcessRule).where(
 -                     DatasetProcessRule.id == dataset_document.dataset_process_rule_id
 -                 )
 -                 processing_rule = db.session.scalar(stmt)
 -                 if not processing_rule:
 -                     raise ValueError("no process rule found")
 -                 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
 -                 documents = self._transform(
 -                     index_processor, dataset, text_docs, dataset_document.doc_language, processing_rule.to_dict()
 -                 )
 -                 # save segment
 -                 self._load_segments(dataset, dataset_document, documents)
 - 
 -                 # load
 -                 self._load(
 -                     index_processor=index_processor,
 -                     dataset=dataset,
 -                     dataset_document=dataset_document,
 -                     documents=documents,
 -                 )
 -             except DocumentIsPausedError:
 -                 raise DocumentIsPausedError(f"Document paused, document id: {dataset_document.id}")
 -             except ProviderTokenNotInitError as e:
 -                 dataset_document.indexing_status = "error"
 -                 dataset_document.error = str(e.description)
 -                 dataset_document.stopped_at = naive_utc_now()
 -                 db.session.commit()
 -             except ObjectDeletedError:
 -                 logger.warning("Document deleted, document id: %s", dataset_document.id)
 -             except Exception as e:
 -                 logger.exception("consume document failed")
 -                 dataset_document.indexing_status = "error"
 -                 dataset_document.error = str(e)
 -                 dataset_document.stopped_at = naive_utc_now()
 -                 db.session.commit()
 - 
 -     def run_in_splitting_status(self, dataset_document: DatasetDocument):
 -         """Run the indexing process when the index_status is splitting."""
 -         try:
 -             # get dataset
 -             dataset = db.session.query(Dataset).filter_by(id=dataset_document.dataset_id).first()
 - 
 -             if not dataset:
 -                 raise ValueError("no dataset found")
 - 
 -             # get exist document_segment list and delete
 -             document_segments = (
 -                 db.session.query(DocumentSegment)
 -                 .filter_by(dataset_id=dataset.id, document_id=dataset_document.id)
 -                 .all()
 -             )
 - 
 -             for document_segment in document_segments:
 -                 db.session.delete(document_segment)
 -                 if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
 -                     # delete child chunks
 -                     db.session.query(ChildChunk).where(ChildChunk.segment_id == document_segment.id).delete()
 -             db.session.commit()
 -             # get the process rule
 -             stmt = select(DatasetProcessRule).where(DatasetProcessRule.id == dataset_document.dataset_process_rule_id)
 -             processing_rule = db.session.scalar(stmt)
 -             if not processing_rule:
 -                 raise ValueError("no process rule found")
 - 
 -             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
 -             documents = self._transform(
 -                 index_processor, dataset, text_docs, dataset_document.doc_language, processing_rule.to_dict()
 -             )
 -             # save segment
 -             self._load_segments(dataset, dataset_document, documents)
 - 
 -             # load
 -             self._load(
 -                 index_processor=index_processor, dataset=dataset, dataset_document=dataset_document, documents=documents
 -             )
 -         except DocumentIsPausedError:
 -             raise DocumentIsPausedError(f"Document paused, document id: {dataset_document.id}")
 -         except ProviderTokenNotInitError as e:
 -             dataset_document.indexing_status = "error"
 -             dataset_document.error = str(e.description)
 -             dataset_document.stopped_at = naive_utc_now()
 -             db.session.commit()
 -         except Exception as e:
 -             logger.exception("consume document failed")
 -             dataset_document.indexing_status = "error"
 -             dataset_document.error = str(e)
 -             dataset_document.stopped_at = naive_utc_now()
 -             db.session.commit()
 - 
 -     def run_in_indexing_status(self, dataset_document: DatasetDocument):
 -         """Run the indexing process when the index_status is indexing."""
 -         try:
 -             # get dataset
 -             dataset = db.session.query(Dataset).filter_by(id=dataset_document.dataset_id).first()
 - 
 -             if not dataset:
 -                 raise ValueError("no dataset found")
 - 
 -             # get exist document_segment list and delete
 -             document_segments = (
 -                 db.session.query(DocumentSegment)
 -                 .filter_by(dataset_id=dataset.id, document_id=dataset_document.id)
 -                 .all()
 -             )
 - 
 -             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,
 -                             },
 -                         )
 -                         if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
 -                             child_chunks = document_segment.get_child_chunks()
 -                             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
 -                         documents.append(document)
 -             # build index
 -             index_type = dataset_document.doc_form
 -             index_processor = IndexProcessorFactory(index_type).init_index_processor()
 -             self._load(
 -                 index_processor=index_processor, dataset=dataset, dataset_document=dataset_document, documents=documents
 -             )
 -         except DocumentIsPausedError:
 -             raise DocumentIsPausedError(f"Document paused, document id: {dataset_document.id}")
 -         except ProviderTokenNotInitError as e:
 -             dataset_document.indexing_status = "error"
 -             dataset_document.error = str(e.description)
 -             dataset_document.stopped_at = naive_utc_now()
 -             db.session.commit()
 -         except Exception as e:
 -             logger.exception("consume document failed")
 -             dataset_document.indexing_status = "error"
 -             dataset_document.error = str(e)
 -             dataset_document.stopped_at = naive_utc_now()
 -             db.session.commit()
 - 
 -     def indexing_estimate(
 -         self,
 -         tenant_id: str,
 -         extract_settings: list[ExtractSetting],
 -         tmp_processing_rule: dict,
 -         doc_form: Optional[str] = None,
 -         doc_language: str = "English",
 -         dataset_id: Optional[str] = None,
 -         indexing_technique: str = "economy",
 -     ) -> IndexingEstimate:
 -         """
 -         Estimate the indexing for the document.
 -         """
 -         # check document limit
 -         features = FeatureService.get_features(tenant_id)
 -         if features.billing.enabled:
 -             count = len(extract_settings)
 -             batch_upload_limit = dify_config.BATCH_UPLOAD_LIMIT
 -             if count > batch_upload_limit:
 -                 raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")
 - 
 -         embedding_model_instance = None
 -         if dataset_id:
 -             dataset = db.session.query(Dataset).filter_by(id=dataset_id).first()
 -             if not dataset:
 -                 raise ValueError("Dataset not found.")
 -             if dataset.indexing_technique == "high_quality" or indexing_technique == "high_quality":
 -                 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,
 -                         model=dataset.embedding_model,
 -                     )
 -                 else:
 -                     embedding_model_instance = self.model_manager.get_default_model_instance(
 -                         tenant_id=tenant_id,
 -                         model_type=ModelType.TEXT_EMBEDDING,
 -                     )
 -         else:
 -             if indexing_technique == "high_quality":
 -                 embedding_model_instance = self.model_manager.get_default_model_instance(
 -                     tenant_id=tenant_id,
 -                     model_type=ModelType.TEXT_EMBEDDING,
 -                 )
 -         # keep separate, avoid union-list ambiguity
 -         preview_texts: list[PreviewDetail] = []
 -         qa_preview_texts: list[QAPreviewDetail] = []
 - 
 -         total_segments = 0
 -         index_type = doc_form
 -         index_processor = IndexProcessorFactory(index_type).init_index_processor()
 -         for extract_setting in extract_settings:
 -             # extract
 -             processing_rule = DatasetProcessRule(
 -                 mode=tmp_processing_rule["mode"], rules=json.dumps(tmp_processing_rule["rules"])
 -             )
 -             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=tenant_id,
 -                 doc_language=doc_language,
 -                 preview=True,
 -             )
 -             total_segments += len(documents)
 -             for document in documents:
 -                 if len(preview_texts) < 10:
 -                     if doc_form and doc_form == "qa_model":
 -                         qa_detail = QAPreviewDetail(
 -                             question=document.page_content, answer=document.metadata.get("answer") or ""
 -                         )
 -                         qa_preview_texts.append(qa_detail)
 -                     else:
 -                         preview_detail = PreviewDetail(content=document.page_content)
 -                         if document.children:
 -                             preview_detail.child_chunks = [child.page_content for child in document.children]
 -                         preview_texts.append(preview_detail)
 - 
 -                 # 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:
 -                     stmt = select(UploadFile).where(UploadFile.id == upload_file_id)
 -                     image_file = db.session.scalar(stmt)
 -                     if image_file is None:
 -                         continue
 -                     try:
 -                         storage.delete(image_file.key)
 -                     except Exception:
 -                         logger.exception(
 -                             "Delete image_files failed while indexing_estimate, \
 -                                           image_upload_file_is: %s",
 -                             upload_file_id,
 -                         )
 -                     db.session.delete(image_file)
 - 
 -         if doc_form and doc_form == "qa_model":
 -             return IndexingEstimate(total_segments=total_segments * 20, qa_preview=qa_preview_texts, preview=[])
 -         return IndexingEstimate(total_segments=total_segments, preview=preview_texts)
 - 
 -     def _extract(
 -         self, index_processor: BaseIndexProcessor, dataset_document: DatasetDocument, process_rule: dict
 -     ) -> list[Document]:
 -         # load file
 -         if dataset_document.data_source_type not in {"upload_file", "notion_import", "website_crawl"}:
 -             return []
 - 
 -         data_source_info = dataset_document.data_source_info_dict
 -         text_docs = []
 -         if dataset_document.data_source_type == "upload_file":
 -             if not data_source_info or "upload_file_id" not in data_source_info:
 -                 raise ValueError("no upload file found")
 -             stmt = select(UploadFile).where(UploadFile.id == data_source_info["upload_file_id"])
 -             file_detail = db.session.scalars(stmt).one_or_none()
 - 
 -             if file_detail:
 -                 extract_setting = ExtractSetting(
 -                     datasource_type=DatasourceType.FILE.value,
 -                     upload_file=file_detail,
 -                     document_model=dataset_document.doc_form,
 -                 )
 -                 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
 -             ):
 -                 raise ValueError("no notion import info found")
 -             extract_setting = ExtractSetting(
 -                 datasource_type=DatasourceType.NOTION.value,
 -                 notion_info={
 -                     "credential_id": data_source_info["credential_id"],
 -                     "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"],
 -                     "document": dataset_document,
 -                     "tenant_id": dataset_document.tenant_id,
 -                 },
 -                 document_model=dataset_document.doc_form,
 -             )
 -             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
 -             ):
 -                 raise ValueError("no website import info found")
 -             extract_setting = ExtractSetting(
 -                 datasource_type=DatasourceType.WEBSITE.value,
 -                 website_info={
 -                     "provider": data_source_info["provider"],
 -                     "job_id": data_source_info["job_id"],
 -                     "tenant_id": dataset_document.tenant_id,
 -                     "url": data_source_info["url"],
 -                     "mode": data_source_info["mode"],
 -                     "only_main_content": data_source_info["only_main_content"],
 -                 },
 -                 document_model=dataset_document.doc_form,
 -             )
 -             text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"])
 -         # update document status to splitting
 -         self._update_document_index_status(
 -             document_id=dataset_document.id,
 -             after_indexing_status="splitting",
 -             extra_update_params={
 -                 DatasetDocument.word_count: sum(len(text_doc.page_content) for text_doc in text_docs),
 -                 DatasetDocument.parsing_completed_at: naive_utc_now(),
 -             },
 -         )
 - 
 -         # replace doc id to document model id
 -         for text_doc in text_docs:
 -             if text_doc.metadata is not None:
 -                 text_doc.metadata["document_id"] = dataset_document.id
 -                 text_doc.metadata["dataset_id"] = dataset_document.dataset_id
 - 
 -         return text_docs
 - 
 -     @staticmethod
 -     def filter_string(text):
 -         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)
 -         # Unicode  U+FFFE
 -         text = re.sub("\ufffe", "", text)
 -         return text
 - 
 -     @staticmethod
 -     def _get_splitter(
 -         processing_rule_mode: str,
 -         max_tokens: int,
 -         chunk_overlap: int,
 -         separator: str,
 -         embedding_model_instance: Optional[ModelInstance],
 -     ) -> TextSplitter:
 -         """
 -         Get the NodeParser object according to the processing rule.
 -         """
 -         character_splitter: TextSplitter
 -         if processing_rule_mode in ["custom", "hierarchical"]:
 -             # The user-defined segmentation rule
 -             max_segmentation_tokens_length = dify_config.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH
 -             if max_tokens < 50 or max_tokens > max_segmentation_tokens_length:
 -                 raise ValueError(f"Custom segment length should be between 50 and {max_segmentation_tokens_length}.")
 - 
 -             if separator:
 -                 separator = separator.replace("\\n", "\n")
 - 
 -             character_splitter = FixedRecursiveCharacterTextSplitter.from_encoder(
 -                 chunk_size=max_tokens,
 -                 chunk_overlap=chunk_overlap,
 -                 fixed_separator=separator,
 -                 separators=["\n\n", "。", ". ", " ", ""],
 -                 embedding_model_instance=embedding_model_instance,
 -             )
 -         else:
 -             # Automatic segmentation
 -             automatic_rules: dict[str, Any] = dict(DatasetProcessRule.AUTOMATIC_RULES["segmentation"])
 -             character_splitter = EnhanceRecursiveCharacterTextSplitter.from_encoder(
 -                 chunk_size=automatic_rules["max_tokens"],
 -                 chunk_overlap=automatic_rules["chunk_overlap"],
 -                 separators=["\n\n", "。", ". ", " ", ""],
 -                 embedding_model_instance=embedding_model_instance,
 -             )
 - 
 -         return character_splitter
 - 
 -     def _split_to_documents_for_estimate(
 -         self, text_docs: list[Document], splitter: TextSplitter, processing_rule: DatasetProcessRule
 -     ) -> list[Document]:
 -         """
 -         Split the text documents into nodes.
 -         """
 -         all_documents: list[Document] = []
 -         for text_doc in text_docs:
 -             # document clean
 -             document_text = self._document_clean(text_doc.page_content, processing_rule)
 -             text_doc.page_content = document_text
 - 
 -             # parse document to nodes
 -             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
 -                 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
 - 
 -                 split_documents.append(document)
 - 
 -             all_documents.extend(split_documents)
 - 
 -         return all_documents
 - 
 -     @staticmethod
 -     def _document_clean(text: str, processing_rule: DatasetProcessRule) -> str:
 -         """
 -         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 {}
 -         document_text = CleanProcessor.clean(text, {"rules": rules})
 - 
 -         return document_text
 - 
 -     @staticmethod
 -     def format_split_text(text: str) -> list[QAPreviewDetail]:
 -         regex = r"Q\d+:\s*(.*?)\s*A\d+:\s*([\s\S]*?)(?=Q\d+:|$)"
 -         matches = re.findall(regex, text, re.UNICODE)
 - 
 -         return [QAPreviewDetail(question=q, answer=re.sub(r"\n\s*", "\n", a.strip())) for q, a in matches if q and a]
 - 
 -     def _load(
 -         self,
 -         index_processor: BaseIndexProcessor,
 -         dataset: Dataset,
 -         dataset_document: DatasetDocument,
 -         documents: list[Document],
 -     ):
 -         """
 -         insert index and update document/segment status to completed
 -         """
 - 
 -         embedding_model_instance = None
 -         if dataset.indexing_technique == "high_quality":
 -             embedding_model_instance = self.model_manager.get_model_instance(
 -                 tenant_id=dataset.tenant_id,
 -                 provider=dataset.embedding_model_provider,
 -                 model_type=ModelType.TEXT_EMBEDDING,
 -                 model=dataset.embedding_model,
 -             )
 - 
 -         # chunk nodes by chunk size
 -         indexing_start_at = time.perf_counter()
 -         tokens = 0
 -         if dataset_document.doc_form != IndexType.PARENT_CHILD_INDEX and dataset.indexing_technique == "economy":
 -             # create keyword index
 -             create_keyword_thread = threading.Thread(
 -                 target=self._process_keyword_index,
 -                 args=(current_app._get_current_object(), dataset.id, dataset_document.id, documents),  # type: ignore
 -             )
 -             create_keyword_thread.start()
 - 
 -         max_workers = 10
 -         if dataset.indexing_technique == "high_quality":
 -             with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
 -                 futures = []
 - 
 -                 # 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
 -                     futures.append(
 -                         executor.submit(
 -                             self._process_chunk,
 -                             current_app._get_current_object(),  # type: ignore
 -                             index_processor,
 -                             chunk_documents,
 -                             dataset,
 -                             dataset_document,
 -                             embedding_model_instance,
 -                         )
 -                     )
 - 
 -                 for future in futures:
 -                     tokens += future.result()
 -         if dataset_document.doc_form != IndexType.PARENT_CHILD_INDEX and dataset.indexing_technique == "economy":
 -             create_keyword_thread.join()
 -         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,
 -                 DatasetDocument.completed_at: naive_utc_now(),
 -                 DatasetDocument.indexing_latency: indexing_end_at - indexing_start_at,
 -                 DatasetDocument.error: None,
 -             },
 -         )
 - 
 -     @staticmethod
 -     def _process_keyword_index(flask_app, dataset_id, document_id, documents):
 -         with flask_app.app_context():
 -             dataset = db.session.query(Dataset).filter_by(id=dataset_id).first()
 -             if not dataset:
 -                 raise ValueError("no dataset found")
 -             keyword = Keyword(dataset)
 -             keyword.create(documents)
 -             if dataset.indexing_technique != "high_quality":
 -                 document_ids = [document.metadata["doc_id"] for document in documents]
 -                 db.session.query(DocumentSegment).where(
 -                     DocumentSegment.document_id == document_id,
 -                     DocumentSegment.dataset_id == dataset_id,
 -                     DocumentSegment.index_node_id.in_(document_ids),
 -                     DocumentSegment.status == "indexing",
 -                 ).update(
 -                     {
 -                         DocumentSegment.status: "completed",
 -                         DocumentSegment.enabled: True,
 -                         DocumentSegment.completed_at: naive_utc_now(),
 -                     }
 -                 )
 - 
 -                 db.session.commit()
 - 
 -     def _process_chunk(
 -         self, flask_app, index_processor, chunk_documents, dataset, dataset_document, embedding_model_instance
 -     ):
 -         with flask_app.app_context():
 -             # check document is paused
 -             self._check_document_paused_status(dataset_document.id)
 - 
 -             tokens = 0
 -             if embedding_model_instance:
 -                 page_content_list = [document.page_content for document in chunk_documents]
 -                 tokens += sum(embedding_model_instance.get_text_embedding_num_tokens(page_content_list))
 - 
 -             # load index
 -             index_processor.load(dataset, chunk_documents, with_keywords=False)
 - 
 -             document_ids = [document.metadata["doc_id"] for document in chunk_documents]
 -             db.session.query(DocumentSegment).where(
 -                 DocumentSegment.document_id == dataset_document.id,
 -                 DocumentSegment.dataset_id == dataset.id,
 -                 DocumentSegment.index_node_id.in_(document_ids),
 -                 DocumentSegment.status == "indexing",
 -             ).update(
 -                 {
 -                     DocumentSegment.status: "completed",
 -                     DocumentSegment.enabled: True,
 -                     DocumentSegment.completed_at: naive_utc_now(),
 -                 }
 -             )
 - 
 -             db.session.commit()
 - 
 -             return tokens
 - 
 -     @staticmethod
 -     def _check_document_paused_status(document_id: str):
 -         indexing_cache_key = f"document_{document_id}_is_paused"
 -         result = redis_client.get(indexing_cache_key)
 -         if result:
 -             raise DocumentIsPausedError()
 - 
 -     @staticmethod
 -     def _update_document_index_status(
 -         document_id: str, after_indexing_status: str, extra_update_params: Optional[dict] = None
 -     ):
 -         """
 -         Update the document indexing status.
 -         """
 -         count = db.session.query(DatasetDocument).filter_by(id=document_id, is_paused=True).count()
 -         if count > 0:
 -             raise DocumentIsPausedError()
 -         document = db.session.query(DatasetDocument).filter_by(id=document_id).first()
 -         if not document:
 -             raise DocumentIsDeletedPausedError()
 - 
 -         update_params = {DatasetDocument.indexing_status: after_indexing_status}
 - 
 -         if extra_update_params:
 -             update_params.update(extra_update_params)
 -         db.session.query(DatasetDocument).filter_by(id=document_id).update(update_params)  # type: ignore
 -         db.session.commit()
 - 
 -     @staticmethod
 -     def _update_segments_by_document(dataset_document_id: str, update_params: dict):
 -         """
 -         Update the document segment by document id.
 -         """
 -         db.session.query(DocumentSegment).filter_by(document_id=dataset_document_id).update(update_params)
 -         db.session.commit()
 - 
 -     def _transform(
 -         self,
 -         index_processor: BaseIndexProcessor,
 -         dataset: Dataset,
 -         text_docs: list[Document],
 -         doc_language: str,
 -         process_rule: dict,
 -     ) -> list[Document]:
 -         # get embedding model instance
 -         embedding_model_instance = None
 -         if dataset.indexing_technique == "high_quality":
 -             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,
 -                     model=dataset.embedding_model,
 -                 )
 -             else:
 -                 embedding_model_instance = self.model_manager.get_default_model_instance(
 -                     tenant_id=dataset.tenant_id,
 -                     model_type=ModelType.TEXT_EMBEDDING,
 -                 )
 - 
 -         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,
 -         )
 - 
 -         return documents
 - 
 -     def _load_segments(self, dataset, dataset_document, documents):
 -         # save node to document segment
 -         doc_store = DatasetDocumentStore(
 -             dataset=dataset, user_id=dataset_document.created_by, document_id=dataset_document.id
 -         )
 - 
 -         # add document segments
 -         doc_store.add_documents(docs=documents, save_child=dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX)
 - 
 -         # update document status to indexing
 -         cur_time = naive_utc_now()
 -         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,
 -             },
 -         )
 - 
 -         # update segment status to indexing
 -         self._update_segments_by_document(
 -             dataset_document_id=dataset_document.id,
 -             update_params={
 -                 DocumentSegment.status: "indexing",
 -                 DocumentSegment.indexing_at: naive_utc_now(),
 -             },
 -         )
 -         pass
 - 
 - 
 - class DocumentIsPausedError(Exception):
 -     pass
 - 
 - 
 - class DocumentIsDeletedPausedError(Exception):
 -     pass
 
 
  |