Underground-Digital/Workflow-Engine
0
1import concurrent.futures2import datetime3import json4import logging5import re6import threading7import time8import uuid9from typing import Optional, cast10 11from flask import Flask, current_app12from flask_login import current_user13from sqlalchemy.orm.exc import ObjectDeletedError14 15from configs import dify_config16from core.errors.error import ProviderTokenNotInitError17from core.llm_generator.llm_generator import LLMGenerator18from core.model_manager import ModelInstance, ModelManager19from core.model_runtime.entities.model_entities import ModelType20from core.rag.cleaner.clean_processor import CleanProcessor21from core.rag.datasource.keyword.keyword_factory import Keyword22from core.rag.docstore.dataset_docstore import DatasetDocumentStore23from core.rag.extractor.entity.extract_setting import ExtractSetting24from core.rag.index_processor.index_processor_base import BaseIndexProcessor25from core.rag.index_processor.index_processor_factory import IndexProcessorFactory26from core.rag.models.document import Document27from core.rag.splitter.fixed_text_splitter import (28 EnhanceRecursiveCharacterTextSplitter,29 FixedRecursiveCharacterTextSplitter,30)31from core.rag.splitter.text_splitter import TextSplitter32from extensions.ext_database import db33from extensions.ext_redis import redis_client34from extensions.ext_storage import storage35from libs import helper36from models.dataset import Dataset, DatasetProcessRule, DocumentSegment37from models.dataset import Document as DatasetDocument38from models.model import UploadFile39from services.feature_service import FeatureService40 41 42class IndexingRunner:43 def __init__(self):44 self.storage = storage45 self.model_manager = ModelManager()46 47 def run(self, dataset_documents: list[DatasetDocument]):48 """Run the indexing process."""49 for dataset_document in dataset_documents:50 try:51 # get dataset52 dataset = Dataset.query.filter_by(id=dataset_document.dataset_id).first()53 54 if not dataset:55 raise ValueError("no dataset found")56 57 # get the process rule58 processing_rule = (59 db.session.query(DatasetProcessRule)60 .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id)61 .first()62 )63 index_type = dataset_document.doc_form64 index_processor = IndexProcessorFactory(index_type).init_index_processor()65 # extract66 text_docs = self._extract(index_processor, dataset_document, processing_rule.to_dict())67 68 # transform69 documents = self._transform(70 index_processor, dataset, text_docs, dataset_document.doc_language, processing_rule.to_dict()71 )72 # save segment73 self._load_segments(dataset, dataset_document, documents)74 75 # load76 self._load(77 index_processor=index_processor,78 dataset=dataset,79 dataset_document=dataset_document,80 documents=documents,81 )82 except DocumentIsPausedError:83 raise DocumentIsPausedError("Document paused, document id: {}".format(dataset_document.id))84 except ProviderTokenNotInitError as e:85 dataset_document.indexing_status = "error"86 dataset_document.error = str(e.description)87 dataset_document.stopped_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)88 db.session.commit()89 except ObjectDeletedError:90 logging.warning("Document deleted, document id: {}".format(dataset_document.id))91 except Exception as e:92 logging.exception("consume document failed")93 dataset_document.indexing_status = "error"94 dataset_document.error = str(e)95 dataset_document.stopped_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)96 db.session.commit()97 98 def run_in_splitting_status(self, dataset_document: DatasetDocument):99 """Run the indexing process when the index_status is splitting."""100 try:101 # get dataset102 dataset = Dataset.query.filter_by(id=dataset_document.dataset_id).first()103 104 if not dataset:105 raise ValueError("no dataset found")106 107 # get exist document_segment list and delete108 document_segments = DocumentSegment.query.filter_by(109 dataset_id=dataset.id, document_id=dataset_document.id110 ).all()111 112 for document_segment in document_segments:113 db.session.delete(document_segment)114 db.session.commit()115 # get the process rule116 processing_rule = (117 db.session.query(DatasetProcessRule)118 .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id)119 .first()120 )121 122 index_type = dataset_document.doc_form123 index_processor = IndexProcessorFactory(index_type).init_index_processor()124 # extract125 text_docs = self._extract(index_processor, dataset_document, processing_rule.to_dict())126 127 # transform128 documents = self._transform(129 index_processor, dataset, text_docs, dataset_document.doc_language, processing_rule.to_dict()130 )131 # save segment132 self._load_segments(dataset, dataset_document, documents)133 134 # load135 self._load(136 index_processor=index_processor, dataset=dataset, dataset_document=dataset_document, documents=documents137 )138 except DocumentIsPausedError:139 raise DocumentIsPausedError("Document paused, document id: {}".format(dataset_document.id))140 except ProviderTokenNotInitError as e:141 dataset_document.indexing_status = "error"142 dataset_document.error = str(e.description)143 dataset_document.stopped_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)144 db.session.commit()145 except Exception as e:146 logging.exception("consume document failed")147 dataset_document.indexing_status = "error"148 dataset_document.error = str(e)149 dataset_document.stopped_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)150 db.session.commit()151 152 def run_in_indexing_status(self, dataset_document: DatasetDocument):153 """Run the indexing process when the index_status is indexing."""154 try:155 # get dataset156 dataset = Dataset.query.filter_by(id=dataset_document.dataset_id).first()157 158 if not dataset:159 raise ValueError("no dataset found")160 161 # get exist document_segment list and delete162 document_segments = DocumentSegment.query.filter_by(163 dataset_id=dataset.id, document_id=dataset_document.id164 ).all()165 166 documents = []167 if document_segments:168 for document_segment in document_segments:169 # transform segment to node170 if document_segment.status != "completed":171 document = Document(172 page_content=document_segment.content,173 metadata={174 "doc_id": document_segment.index_node_id,175 "doc_hash": document_segment.index_node_hash,176 "document_id": document_segment.document_id,177 "dataset_id": document_segment.dataset_id,178 },179 )180 181 documents.append(document)182 183 # build index184 # get the process rule185 processing_rule = (186 db.session.query(DatasetProcessRule)187 .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id)188 .first()189 )190 191 index_type = dataset_document.doc_form192 index_processor = IndexProcessorFactory(index_type).init_index_processor()193 self._load(194 index_processor=index_processor, dataset=dataset, dataset_document=dataset_document, documents=documents195 )196 except DocumentIsPausedError:197 raise DocumentIsPausedError("Document paused, document id: {}".format(dataset_document.id))198 except ProviderTokenNotInitError as e:199 dataset_document.indexing_status = "error"200 dataset_document.error = str(e.description)201 dataset_document.stopped_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)202 db.session.commit()203 except Exception as e:204 logging.exception("consume document failed")205 dataset_document.indexing_status = "error"206 dataset_document.error = str(e)207 dataset_document.stopped_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)208 db.session.commit()209 210 def indexing_estimate(211 self,212 tenant_id: str,213 extract_settings: list[ExtractSetting],214 tmp_processing_rule: dict,215 doc_form: Optional[str] = None,216 doc_language: str = "English",217 dataset_id: Optional[str] = None,218 indexing_technique: str = "economy",219 ) -> dict:220 """221 Estimate the indexing for the document.222 """223 # check document limit224 features = FeatureService.get_features(tenant_id)225 if features.billing.enabled:226 count = len(extract_settings)227 batch_upload_limit = dify_config.BATCH_UPLOAD_LIMIT228 if count > batch_upload_limit:229 raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")230 231 embedding_model_instance = None232 if dataset_id:233 dataset = Dataset.query.filter_by(id=dataset_id).first()234 if not dataset:235 raise ValueError("Dataset not found.")236 if dataset.indexing_technique == "high_quality" or indexing_technique == "high_quality":237 if dataset.embedding_model_provider:238 embedding_model_instance = self.model_manager.get_model_instance(239 tenant_id=tenant_id,240 provider=dataset.embedding_model_provider,241 model_type=ModelType.TEXT_EMBEDDING,242 model=dataset.embedding_model,243 )244 else:245 embedding_model_instance = self.model_manager.get_default_model_instance(246 tenant_id=tenant_id,247 model_type=ModelType.TEXT_EMBEDDING,248 )249 else:250 if indexing_technique == "high_quality":251 embedding_model_instance = self.model_manager.get_default_model_instance(252 tenant_id=tenant_id,253 model_type=ModelType.TEXT_EMBEDDING,254 )255 preview_texts = []256 total_segments = 0257 index_type = doc_form258 index_processor = IndexProcessorFactory(index_type).init_index_processor()259 all_text_docs = []260 for extract_setting in extract_settings:261 # extract262 text_docs = index_processor.extract(extract_setting, process_rule_mode=tmp_processing_rule["mode"])263 all_text_docs.extend(text_docs)264 processing_rule = DatasetProcessRule(265 mode=tmp_processing_rule["mode"], rules=json.dumps(tmp_processing_rule["rules"])266 )267 268 # get splitter269 splitter = self._get_splitter(processing_rule, embedding_model_instance)270 271 # split to documents272 documents = self._split_to_documents_for_estimate(273 text_docs=text_docs, splitter=splitter, processing_rule=processing_rule274 )275 276 total_segments += len(documents)277 for document in documents:278 if len(preview_texts) < 5:279 preview_texts.append(document.page_content)280 281 if doc_form and doc_form == "qa_model":282 if len(preview_texts) > 0:283 # qa model document284 response = LLMGenerator.generate_qa_document(285 current_user.current_tenant_id, preview_texts[0], doc_language286 )287 document_qa_list = self.format_split_text(response)288 289 return {"total_segments": total_segments * 20, "qa_preview": document_qa_list, "preview": preview_texts}290 return {"total_segments": total_segments, "preview": preview_texts}291 292 def _extract(293 self, index_processor: BaseIndexProcessor, dataset_document: DatasetDocument, process_rule: dict294 ) -> list[Document]:295 # load file296 if dataset_document.data_source_type not in {"upload_file", "notion_import", "website_crawl"}:297 return []298 299 data_source_info = dataset_document.data_source_info_dict300 text_docs = []301 if dataset_document.data_source_type == "upload_file":302 if not data_source_info or "upload_file_id" not in data_source_info:303 raise ValueError("no upload file found")304 305 file_detail = (306 db.session.query(UploadFile).filter(UploadFile.id == data_source_info["upload_file_id"]).one_or_none()307 )308 309 if file_detail:310 extract_setting = ExtractSetting(311 datasource_type="upload_file", upload_file=file_detail, document_model=dataset_document.doc_form312 )313 text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"])314 elif dataset_document.data_source_type == "notion_import":315 if (316 not data_source_info317 or "notion_workspace_id" not in data_source_info318 or "notion_page_id" not in data_source_info319 ):320 raise ValueError("no notion import info found")321 extract_setting = ExtractSetting(322 datasource_type="notion_import",323 notion_info={324 "notion_workspace_id": data_source_info["notion_workspace_id"],325 "notion_obj_id": data_source_info["notion_page_id"],326 "notion_page_type": data_source_info["type"],327 "document": dataset_document,328 "tenant_id": dataset_document.tenant_id,329 },330 document_model=dataset_document.doc_form,331 )332 text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"])333 elif dataset_document.data_source_type == "website_crawl":334 if (335 not data_source_info336 or "provider" not in data_source_info337 or "url" not in data_source_info338 or "job_id" not in data_source_info339 ):340 raise ValueError("no website import info found")341 extract_setting = ExtractSetting(342 datasource_type="website_crawl",343 website_info={344 "provider": data_source_info["provider"],345 "job_id": data_source_info["job_id"],346 "tenant_id": dataset_document.tenant_id,347 "url": data_source_info["url"],348 "mode": data_source_info["mode"],349 "only_main_content": data_source_info["only_main_content"],350 },351 document_model=dataset_document.doc_form,352 )353 text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"])354 # update document status to splitting355 self._update_document_index_status(356 document_id=dataset_document.id,357 after_indexing_status="splitting",358 extra_update_params={359 DatasetDocument.word_count: sum(len(text_doc.page_content) for text_doc in text_docs),360 DatasetDocument.parsing_completed_at: datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),361 },362 )363 364 # replace doc id to document model id365 text_docs = cast(list[Document], text_docs)366 for text_doc in text_docs:367 text_doc.metadata["document_id"] = dataset_document.id368 text_doc.metadata["dataset_id"] = dataset_document.dataset_id369 370 return text_docs371 372 @staticmethod373 def filter_string(text):374 text = re.sub(r"<\|", "<", text)375 text = re.sub(r"\|>", ">", text)376 text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F\xEF\xBF\xBE]", "", text)377 # Unicode U+FFFE378 text = re.sub("\ufffe", "", text)379 return text380 381 @staticmethod382 def _get_splitter(383 processing_rule: DatasetProcessRule, embedding_model_instance: Optional[ModelInstance]384 ) -> TextSplitter:385 """386 Get the NodeParser object according to the processing rule.387 """388 if processing_rule.mode == "custom":389 # The user-defined segmentation rule390 rules = json.loads(processing_rule.rules)391 segmentation = rules["segmentation"]392 max_segmentation_tokens_length = dify_config.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH393 if segmentation["max_tokens"] < 50 or segmentation["max_tokens"] > max_segmentation_tokens_length:394 raise ValueError(f"Custom segment length should be between 50 and {max_segmentation_tokens_length}.")395 396 separator = segmentation["separator"]397 if separator:398 separator = separator.replace("\\n", "\n")399 400 if segmentation.get("chunk_overlap"):401 chunk_overlap = segmentation["chunk_overlap"]402 else:403 chunk_overlap = 0404 405 character_splitter = FixedRecursiveCharacterTextSplitter.from_encoder(406 chunk_size=segmentation["max_tokens"],407 chunk_overlap=chunk_overlap,408 fixed_separator=separator,409 separators=["\n\n", "ใ", ". ", " ", ""],410 embedding_model_instance=embedding_model_instance,411 )412 else:413 # Automatic segmentation414 character_splitter = EnhanceRecursiveCharacterTextSplitter.from_encoder(415 chunk_size=DatasetProcessRule.AUTOMATIC_RULES["segmentation"]["max_tokens"],416 chunk_overlap=DatasetProcessRule.AUTOMATIC_RULES["segmentation"]["chunk_overlap"],417 separators=["\n\n", "ใ", ". ", " ", ""],418 embedding_model_instance=embedding_model_instance,419 )420 421 return character_splitter422 423 def _step_split(424 self,425 text_docs: list[Document],426 splitter: TextSplitter,427 dataset: Dataset,428 dataset_document: DatasetDocument,429 processing_rule: DatasetProcessRule,430 ) -> list[Document]:431 """432 Split the text documents into documents and save them to the document segment.433 """434 documents = self._split_to_documents(435 text_docs=text_docs,436 splitter=splitter,437 processing_rule=processing_rule,438 tenant_id=dataset.tenant_id,439 document_form=dataset_document.doc_form,440 document_language=dataset_document.doc_language,441 )442 443 # save node to document segment444 doc_store = DatasetDocumentStore(445 dataset=dataset, user_id=dataset_document.created_by, document_id=dataset_document.id446 )447 448 # add document segments449 doc_store.add_documents(documents)450 451 # update document status to indexing452 cur_time = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)453 self._update_document_index_status(454 document_id=dataset_document.id,455 after_indexing_status="indexing",456 extra_update_params={457 DatasetDocument.cleaning_completed_at: cur_time,458 DatasetDocument.splitting_completed_at: cur_time,459 },460 )461 462 # update segment status to indexing463 self._update_segments_by_document(464 dataset_document_id=dataset_document.id,465 update_params={466 DocumentSegment.status: "indexing",467 DocumentSegment.indexing_at: datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),468 },469 )470 471 return documents472 473 def _split_to_documents(474 self,475 text_docs: list[Document],476 splitter: TextSplitter,477 processing_rule: DatasetProcessRule,478 tenant_id: str,479 document_form: str,480 document_language: str,481 ) -> list[Document]:482 """483 Split the text documents into nodes.484 """485 all_documents = []486 all_qa_documents = []487 for text_doc in text_docs:488 # document clean489 document_text = self._document_clean(text_doc.page_content, processing_rule)490 text_doc.page_content = document_text491 492 # parse document to nodes493 documents = splitter.split_documents([text_doc])494 split_documents = []495 for document_node in documents:496 if document_node.page_content.strip():497 doc_id = str(uuid.uuid4())498 hash = helper.generate_text_hash(document_node.page_content)499 document_node.metadata["doc_id"] = doc_id500 document_node.metadata["doc_hash"] = hash501 # delete Splitter character502 page_content = document_node.page_content503 if page_content.startswith(".") or page_content.startswith("ใ"):504 page_content = page_content[1:]505 else:506 page_content = page_content507 document_node.page_content = page_content508 509 if document_node.page_content:510 split_documents.append(document_node)511 all_documents.extend(split_documents)512 # processing qa document513 if document_form == "qa_model":514 for i in range(0, len(all_documents), 10):515 threads = []516 sub_documents = all_documents[i : i + 10]517 for doc in sub_documents:518 document_format_thread = threading.Thread(519 target=self.format_qa_document,520 kwargs={521 "flask_app": current_app._get_current_object(),522 "tenant_id": tenant_id,523 "document_node": doc,524 "all_qa_documents": all_qa_documents,525 "document_language": document_language,526 },527 )528 threads.append(document_format_thread)529 document_format_thread.start()530 for thread in threads:531 thread.join()532 return all_qa_documents533 return all_documents534 535 def format_qa_document(self, flask_app: Flask, tenant_id: str, document_node, all_qa_documents, document_language):536 format_documents = []537 if document_node.page_content is None or not document_node.page_content.strip():538 return539 with flask_app.app_context():540 try:541 # qa model document542 response = LLMGenerator.generate_qa_document(tenant_id, document_node.page_content, document_language)543 document_qa_list = self.format_split_text(response)544 qa_documents = []545 for result in document_qa_list:546 qa_document = Document(547 page_content=result["question"], metadata=document_node.metadata.model_copy()548 )549 doc_id = str(uuid.uuid4())550 hash = helper.generate_text_hash(result["question"])551 qa_document.metadata["answer"] = result["answer"]552 qa_document.metadata["doc_id"] = doc_id553 qa_document.metadata["doc_hash"] = hash554 qa_documents.append(qa_document)555 format_documents.extend(qa_documents)556 except Exception as e:557 logging.exception(e)558 559 all_qa_documents.extend(format_documents)560 561 def _split_to_documents_for_estimate(562 self, text_docs: list[Document], splitter: TextSplitter, processing_rule: DatasetProcessRule563 ) -> list[Document]:564 """565 Split the text documents into nodes.566 """567 all_documents = []568 for text_doc in text_docs:569 # document clean570 document_text = self._document_clean(text_doc.page_content, processing_rule)571 text_doc.page_content = document_text572 573 # parse document to nodes574 documents = splitter.split_documents([text_doc])575 576 split_documents = []577 for document in documents:578 if document.page_content is None or not document.page_content.strip():579 continue580 doc_id = str(uuid.uuid4())581 hash = helper.generate_text_hash(document.page_content)582 583 document.metadata["doc_id"] = doc_id584 document.metadata["doc_hash"] = hash585 586 split_documents.append(document)587 588 all_documents.extend(split_documents)589 590 return all_documents591 592 @staticmethod593 def _document_clean(text: str, processing_rule: DatasetProcessRule) -> str:594 """595 Clean the document text according to the processing rules.596 """597 if processing_rule.mode == "automatic":598 rules = DatasetProcessRule.AUTOMATIC_RULES599 else:600 rules = json.loads(processing_rule.rules) if processing_rule.rules else {}601 document_text = CleanProcessor.clean(text, {"rules": rules})602 603 return document_text604 605 @staticmethod606 def format_split_text(text):607 regex = r"Q\d+:\s*(.*?)\s*A\d+:\s*([\s\S]*?)(?=Q\d+:|$)"608 matches = re.findall(regex, text, re.UNICODE)609 610 return [{"question": q, "answer": re.sub(r"\n\s*", "\n", a.strip())} for q, a in matches if q and a]611 612 def _load(613 self,614 index_processor: BaseIndexProcessor,615 dataset: Dataset,616 dataset_document: DatasetDocument,617 documents: list[Document],618 ) -> None:619 """620 insert index and update document/segment status to completed621 """622 623 embedding_model_instance = None624 if dataset.indexing_technique == "high_quality":625 embedding_model_instance = self.model_manager.get_model_instance(626 tenant_id=dataset.tenant_id,627 provider=dataset.embedding_model_provider,628 model_type=ModelType.TEXT_EMBEDDING,629 model=dataset.embedding_model,630 )631 632 # chunk nodes by chunk size633 indexing_start_at = time.perf_counter()634 tokens = 0635 chunk_size = 10636 637 # create keyword index638 create_keyword_thread = threading.Thread(639 target=self._process_keyword_index,640 args=(current_app._get_current_object(), dataset.id, dataset_document.id, documents),641 )642 create_keyword_thread.start()643 if dataset.indexing_technique == "high_quality":644 with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:645 futures = []646 for i in range(0, len(documents), chunk_size):647 chunk_documents = documents[i : i + chunk_size]648 futures.append(649 executor.submit(650 self._process_chunk,651 current_app._get_current_object(),652 index_processor,653 chunk_documents,654 dataset,655 dataset_document,656 embedding_model_instance,657 )658 )659 660 for future in futures:661 tokens += future.result()662 663 create_keyword_thread.join()664 indexing_end_at = time.perf_counter()665 666 # update document status to completed667 self._update_document_index_status(668 document_id=dataset_document.id,669 after_indexing_status="completed",670 extra_update_params={671 DatasetDocument.tokens: tokens,672 DatasetDocument.completed_at: datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),673 DatasetDocument.indexing_latency: indexing_end_at - indexing_start_at,674 DatasetDocument.error: None,675 },676 )677 678 @staticmethod679 def _process_keyword_index(flask_app, dataset_id, document_id, documents):680 with flask_app.app_context():681 dataset = Dataset.query.filter_by(id=dataset_id).first()682 if not dataset:683 raise ValueError("no dataset found")684 keyword = Keyword(dataset)685 keyword.create(documents)686 if dataset.indexing_technique != "high_quality":687 document_ids = [document.metadata["doc_id"] for document in documents]688 db.session.query(DocumentSegment).filter(689 DocumentSegment.document_id == document_id,690 DocumentSegment.dataset_id == dataset_id,691 DocumentSegment.index_node_id.in_(document_ids),692 DocumentSegment.status == "indexing",693 ).update(694 {695 DocumentSegment.status: "completed",696 DocumentSegment.enabled: True,697 DocumentSegment.completed_at: datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),698 }699 )700 701 db.session.commit()702 703 def _process_chunk(704 self, flask_app, index_processor, chunk_documents, dataset, dataset_document, embedding_model_instance705 ):706 with flask_app.app_context():707 # check document is paused708 self._check_document_paused_status(dataset_document.id)709 710 tokens = 0711 if embedding_model_instance:712 tokens += sum(713 embedding_model_instance.get_text_embedding_num_tokens([document.page_content])714 for document in chunk_documents715 )716 717 # load index718 index_processor.load(dataset, chunk_documents, with_keywords=False)719 720 document_ids = [document.metadata["doc_id"] for document in chunk_documents]721 db.session.query(DocumentSegment).filter(722 DocumentSegment.document_id == dataset_document.id,723 DocumentSegment.dataset_id == dataset.id,724 DocumentSegment.index_node_id.in_(document_ids),725 DocumentSegment.status == "indexing",726 ).update(727 {728 DocumentSegment.status: "completed",729 DocumentSegment.enabled: True,730 DocumentSegment.completed_at: datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),731 }732 )733 734 db.session.commit()735 736 return tokens737 738 @staticmethod739 def _check_document_paused_status(document_id: str):740 indexing_cache_key = "document_{}_is_paused".format(document_id)741 result = redis_client.get(indexing_cache_key)742 if result:743 raise DocumentIsPausedError()744 745 @staticmethod746 def _update_document_index_status(747 document_id: str, after_indexing_status: str, extra_update_params: Optional[dict] = None748 ) -> None:749 """750 Update the document indexing status.751 """752 count = DatasetDocument.query.filter_by(id=document_id, is_paused=True).count()753 if count > 0:754 raise DocumentIsPausedError()755 document = DatasetDocument.query.filter_by(id=document_id).first()756 if not document:757 raise DocumentIsDeletedPausedError()758 759 update_params = {DatasetDocument.indexing_status: after_indexing_status}760 761 if extra_update_params:762 update_params.update(extra_update_params)763 764 DatasetDocument.query.filter_by(id=document_id).update(update_params)765 db.session.commit()766 767 @staticmethod768 def _update_segments_by_document(dataset_document_id: str, update_params: dict) -> None:769 """770 Update the document segment by document id.771 """772 DocumentSegment.query.filter_by(document_id=dataset_document_id).update(update_params)773 db.session.commit()774 775 @staticmethod776 def batch_add_segments(segments: list[DocumentSegment], dataset: Dataset):777 """778 Batch add segments index processing779 """780 documents = []781 for segment in segments:782 document = Document(783 page_content=segment.content,784 metadata={785 "doc_id": segment.index_node_id,786 "doc_hash": segment.index_node_hash,787 "document_id": segment.document_id,788 "dataset_id": segment.dataset_id,789 },790 )791 documents.append(document)792 # save vector index793 index_type = dataset.doc_form794 index_processor = IndexProcessorFactory(index_type).init_index_processor()795 index_processor.load(dataset, documents)796 797 def _transform(798 self,799 index_processor: BaseIndexProcessor,800 dataset: Dataset,801 text_docs: list[Document],802 doc_language: str,803 process_rule: dict,804 ) -> list[Document]:805 # get embedding model instance806 embedding_model_instance = None807 if dataset.indexing_technique == "high_quality":808 if dataset.embedding_model_provider:809 embedding_model_instance = self.model_manager.get_model_instance(810 tenant_id=dataset.tenant_id,811 provider=dataset.embedding_model_provider,812 model_type=ModelType.TEXT_EMBEDDING,813 model=dataset.embedding_model,814 )815 else:816 embedding_model_instance = self.model_manager.get_default_model_instance(817 tenant_id=dataset.tenant_id,818 model_type=ModelType.TEXT_EMBEDDING,819 )820 821 documents = index_processor.transform(822 text_docs,823 embedding_model_instance=embedding_model_instance,824 process_rule=process_rule,825 tenant_id=dataset.tenant_id,826 doc_language=doc_language,827 )828 829 return documents830 831 def _load_segments(self, dataset, dataset_document, documents):832 # save node to document segment833 doc_store = DatasetDocumentStore(834 dataset=dataset, user_id=dataset_document.created_by, document_id=dataset_document.id835 )836 837 # add document segments838 doc_store.add_documents(documents)839 840 # update document status to indexing841 cur_time = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)842 self._update_document_index_status(843 document_id=dataset_document.id,844 after_indexing_status="indexing",845 extra_update_params={846 DatasetDocument.cleaning_completed_at: cur_time,847 DatasetDocument.splitting_completed_at: cur_time,848 },849 )850 851 # update segment status to indexing852 self._update_segments_by_document(853 dataset_document_id=dataset_document.id,854 update_params={855 DocumentSegment.status: "indexing",856 DocumentSegment.indexing_at: datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),857 },858 )859 pass860 861 862class DocumentIsPausedError(Exception):863 pass864 865 866class DocumentIsDeletedPausedError(Exception):867 pass868