You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

indexing_runner.py 32KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748
  1. import concurrent.futures
  2. import json
  3. import logging
  4. import re
  5. import threading
  6. import time
  7. import uuid
  8. from typing import Any, Optional, cast
  9. from flask import current_app
  10. from sqlalchemy import select
  11. from sqlalchemy.orm.exc import ObjectDeletedError
  12. from configs import dify_config
  13. from core.entities.knowledge_entities import IndexingEstimate, PreviewDetail, QAPreviewDetail
  14. from core.errors.error import ProviderTokenNotInitError
  15. from core.model_manager import ModelInstance, ModelManager
  16. from core.model_runtime.entities.model_entities import ModelType
  17. from core.rag.cleaner.clean_processor import CleanProcessor
  18. from core.rag.datasource.keyword.keyword_factory import Keyword
  19. from core.rag.docstore.dataset_docstore import DatasetDocumentStore
  20. from core.rag.extractor.entity.datasource_type import DatasourceType
  21. from core.rag.extractor.entity.extract_setting import ExtractSetting
  22. from core.rag.index_processor.constant.index_type import IndexType
  23. from core.rag.index_processor.index_processor_base import BaseIndexProcessor
  24. from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
  25. from core.rag.models.document import ChildDocument, Document
  26. from core.rag.splitter.fixed_text_splitter import (
  27. EnhanceRecursiveCharacterTextSplitter,
  28. FixedRecursiveCharacterTextSplitter,
  29. )
  30. from core.rag.splitter.text_splitter import TextSplitter
  31. from core.tools.utils.web_reader_tool import get_image_upload_file_ids
  32. from extensions.ext_database import db
  33. from extensions.ext_redis import redis_client
  34. from extensions.ext_storage import storage
  35. from libs import helper
  36. from libs.datetime_utils import naive_utc_now
  37. from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegment
  38. from models.dataset import Document as DatasetDocument
  39. from models.model import UploadFile
  40. from services.feature_service import FeatureService
  41. logger = logging.getLogger(__name__)
  42. class IndexingRunner:
  43. def __init__(self):
  44. self.storage = storage
  45. self.model_manager = ModelManager()
  46. def run(self, dataset_documents: list[DatasetDocument]):
  47. """Run the indexing process."""
  48. for dataset_document in dataset_documents:
  49. try:
  50. # get dataset
  51. dataset = db.session.query(Dataset).filter_by(id=dataset_document.dataset_id).first()
  52. if not dataset:
  53. raise ValueError("no dataset found")
  54. # get the process rule
  55. stmt = select(DatasetProcessRule).where(
  56. DatasetProcessRule.id == dataset_document.dataset_process_rule_id
  57. )
  58. processing_rule = db.session.scalar(stmt)
  59. if not processing_rule:
  60. raise ValueError("no process rule found")
  61. index_type = dataset_document.doc_form
  62. index_processor = IndexProcessorFactory(index_type).init_index_processor()
  63. # extract
  64. text_docs = self._extract(index_processor, dataset_document, processing_rule.to_dict())
  65. # transform
  66. documents = self._transform(
  67. index_processor, dataset, text_docs, dataset_document.doc_language, processing_rule.to_dict()
  68. )
  69. # save segment
  70. self._load_segments(dataset, dataset_document, documents)
  71. # load
  72. self._load(
  73. index_processor=index_processor,
  74. dataset=dataset,
  75. dataset_document=dataset_document,
  76. documents=documents,
  77. )
  78. except DocumentIsPausedError:
  79. raise DocumentIsPausedError(f"Document paused, document id: {dataset_document.id}")
  80. except ProviderTokenNotInitError as e:
  81. dataset_document.indexing_status = "error"
  82. dataset_document.error = str(e.description)
  83. dataset_document.stopped_at = naive_utc_now()
  84. db.session.commit()
  85. except ObjectDeletedError:
  86. logger.warning("Document deleted, document id: %s", dataset_document.id)
  87. except Exception as e:
  88. logger.exception("consume document failed")
  89. dataset_document.indexing_status = "error"
  90. dataset_document.error = str(e)
  91. dataset_document.stopped_at = naive_utc_now()
  92. db.session.commit()
  93. def run_in_splitting_status(self, dataset_document: DatasetDocument):
  94. """Run the indexing process when the index_status is splitting."""
  95. try:
  96. # get dataset
  97. dataset = db.session.query(Dataset).filter_by(id=dataset_document.dataset_id).first()
  98. if not dataset:
  99. raise ValueError("no dataset found")
  100. # get exist document_segment list and delete
  101. document_segments = (
  102. db.session.query(DocumentSegment)
  103. .filter_by(dataset_id=dataset.id, document_id=dataset_document.id)
  104. .all()
  105. )
  106. for document_segment in document_segments:
  107. db.session.delete(document_segment)
  108. if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
  109. # delete child chunks
  110. db.session.query(ChildChunk).where(ChildChunk.segment_id == document_segment.id).delete()
  111. db.session.commit()
  112. # get the process rule
  113. stmt = select(DatasetProcessRule).where(DatasetProcessRule.id == dataset_document.dataset_process_rule_id)
  114. processing_rule = db.session.scalar(stmt)
  115. if not processing_rule:
  116. raise ValueError("no process rule found")
  117. index_type = dataset_document.doc_form
  118. index_processor = IndexProcessorFactory(index_type).init_index_processor()
  119. # extract
  120. text_docs = self._extract(index_processor, dataset_document, processing_rule.to_dict())
  121. # transform
  122. documents = self._transform(
  123. index_processor, dataset, text_docs, dataset_document.doc_language, processing_rule.to_dict()
  124. )
  125. # save segment
  126. self._load_segments(dataset, dataset_document, documents)
  127. # load
  128. self._load(
  129. index_processor=index_processor, dataset=dataset, dataset_document=dataset_document, documents=documents
  130. )
  131. except DocumentIsPausedError:
  132. raise DocumentIsPausedError(f"Document paused, document id: {dataset_document.id}")
  133. except ProviderTokenNotInitError as e:
  134. dataset_document.indexing_status = "error"
  135. dataset_document.error = str(e.description)
  136. dataset_document.stopped_at = naive_utc_now()
  137. db.session.commit()
  138. except Exception as e:
  139. logger.exception("consume document failed")
  140. dataset_document.indexing_status = "error"
  141. dataset_document.error = str(e)
  142. dataset_document.stopped_at = naive_utc_now()
  143. db.session.commit()
  144. def run_in_indexing_status(self, dataset_document: DatasetDocument):
  145. """Run the indexing process when the index_status is indexing."""
  146. try:
  147. # get dataset
  148. dataset = db.session.query(Dataset).filter_by(id=dataset_document.dataset_id).first()
  149. if not dataset:
  150. raise ValueError("no dataset found")
  151. # get exist document_segment list and delete
  152. document_segments = (
  153. db.session.query(DocumentSegment)
  154. .filter_by(dataset_id=dataset.id, document_id=dataset_document.id)
  155. .all()
  156. )
  157. documents = []
  158. if document_segments:
  159. for document_segment in document_segments:
  160. # transform segment to node
  161. if document_segment.status != "completed":
  162. document = Document(
  163. page_content=document_segment.content,
  164. metadata={
  165. "doc_id": document_segment.index_node_id,
  166. "doc_hash": document_segment.index_node_hash,
  167. "document_id": document_segment.document_id,
  168. "dataset_id": document_segment.dataset_id,
  169. },
  170. )
  171. if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
  172. child_chunks = document_segment.get_child_chunks()
  173. if child_chunks:
  174. child_documents = []
  175. for child_chunk in child_chunks:
  176. child_document = ChildDocument(
  177. page_content=child_chunk.content,
  178. metadata={
  179. "doc_id": child_chunk.index_node_id,
  180. "doc_hash": child_chunk.index_node_hash,
  181. "document_id": document_segment.document_id,
  182. "dataset_id": document_segment.dataset_id,
  183. },
  184. )
  185. child_documents.append(child_document)
  186. document.children = child_documents
  187. documents.append(document)
  188. # build index
  189. index_type = dataset_document.doc_form
  190. index_processor = IndexProcessorFactory(index_type).init_index_processor()
  191. self._load(
  192. index_processor=index_processor, dataset=dataset, dataset_document=dataset_document, documents=documents
  193. )
  194. except DocumentIsPausedError:
  195. raise DocumentIsPausedError(f"Document paused, document id: {dataset_document.id}")
  196. except ProviderTokenNotInitError as e:
  197. dataset_document.indexing_status = "error"
  198. dataset_document.error = str(e.description)
  199. dataset_document.stopped_at = naive_utc_now()
  200. db.session.commit()
  201. except Exception as e:
  202. logger.exception("consume document failed")
  203. dataset_document.indexing_status = "error"
  204. dataset_document.error = str(e)
  205. dataset_document.stopped_at = naive_utc_now()
  206. db.session.commit()
  207. def indexing_estimate(
  208. self,
  209. tenant_id: str,
  210. extract_settings: list[ExtractSetting],
  211. tmp_processing_rule: dict,
  212. doc_form: Optional[str] = None,
  213. doc_language: str = "English",
  214. dataset_id: Optional[str] = None,
  215. indexing_technique: str = "economy",
  216. ) -> IndexingEstimate:
  217. """
  218. Estimate the indexing for the document.
  219. """
  220. # check document limit
  221. features = FeatureService.get_features(tenant_id)
  222. if features.billing.enabled:
  223. count = len(extract_settings)
  224. batch_upload_limit = dify_config.BATCH_UPLOAD_LIMIT
  225. if count > batch_upload_limit:
  226. raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")
  227. embedding_model_instance = None
  228. if dataset_id:
  229. dataset = db.session.query(Dataset).filter_by(id=dataset_id).first()
  230. if not dataset:
  231. raise ValueError("Dataset not found.")
  232. if dataset.indexing_technique == "high_quality" or indexing_technique == "high_quality":
  233. if dataset.embedding_model_provider:
  234. embedding_model_instance = self.model_manager.get_model_instance(
  235. tenant_id=tenant_id,
  236. provider=dataset.embedding_model_provider,
  237. model_type=ModelType.TEXT_EMBEDDING,
  238. model=dataset.embedding_model,
  239. )
  240. else:
  241. embedding_model_instance = self.model_manager.get_default_model_instance(
  242. tenant_id=tenant_id,
  243. model_type=ModelType.TEXT_EMBEDDING,
  244. )
  245. else:
  246. if indexing_technique == "high_quality":
  247. embedding_model_instance = self.model_manager.get_default_model_instance(
  248. tenant_id=tenant_id,
  249. model_type=ModelType.TEXT_EMBEDDING,
  250. )
  251. preview_texts = [] # type: ignore
  252. total_segments = 0
  253. index_type = doc_form
  254. index_processor = IndexProcessorFactory(index_type).init_index_processor()
  255. for extract_setting in extract_settings:
  256. # extract
  257. processing_rule = DatasetProcessRule(
  258. mode=tmp_processing_rule["mode"], rules=json.dumps(tmp_processing_rule["rules"])
  259. )
  260. text_docs = index_processor.extract(extract_setting, process_rule_mode=tmp_processing_rule["mode"])
  261. documents = index_processor.transform(
  262. text_docs,
  263. embedding_model_instance=embedding_model_instance,
  264. process_rule=processing_rule.to_dict(),
  265. tenant_id=tenant_id,
  266. doc_language=doc_language,
  267. preview=True,
  268. )
  269. total_segments += len(documents)
  270. for document in documents:
  271. if len(preview_texts) < 10:
  272. if doc_form and doc_form == "qa_model":
  273. preview_detail = QAPreviewDetail(
  274. question=document.page_content, answer=document.metadata.get("answer") or ""
  275. )
  276. preview_texts.append(preview_detail)
  277. else:
  278. preview_detail = PreviewDetail(content=document.page_content) # type: ignore
  279. if document.children:
  280. preview_detail.child_chunks = [child.page_content for child in document.children] # type: ignore
  281. preview_texts.append(preview_detail)
  282. # delete image files and related db records
  283. image_upload_file_ids = get_image_upload_file_ids(document.page_content)
  284. for upload_file_id in image_upload_file_ids:
  285. stmt = select(UploadFile).where(UploadFile.id == upload_file_id)
  286. image_file = db.session.scalar(stmt)
  287. if image_file is None:
  288. continue
  289. try:
  290. storage.delete(image_file.key)
  291. except Exception:
  292. logger.exception(
  293. "Delete image_files failed while indexing_estimate, \
  294. image_upload_file_is: %s",
  295. upload_file_id,
  296. )
  297. db.session.delete(image_file)
  298. if doc_form and doc_form == "qa_model":
  299. return IndexingEstimate(total_segments=total_segments * 20, qa_preview=preview_texts, preview=[])
  300. return IndexingEstimate(total_segments=total_segments, preview=preview_texts) # type: ignore
  301. def _extract(
  302. self, index_processor: BaseIndexProcessor, dataset_document: DatasetDocument, process_rule: dict
  303. ) -> list[Document]:
  304. # load file
  305. if dataset_document.data_source_type not in {"upload_file", "notion_import", "website_crawl"}:
  306. return []
  307. data_source_info = dataset_document.data_source_info_dict
  308. text_docs = []
  309. if dataset_document.data_source_type == "upload_file":
  310. if not data_source_info or "upload_file_id" not in data_source_info:
  311. raise ValueError("no upload file found")
  312. stmt = select(UploadFile).where(UploadFile.id == data_source_info["upload_file_id"])
  313. file_detail = db.session.scalars(stmt).one_or_none()
  314. if file_detail:
  315. extract_setting = ExtractSetting(
  316. datasource_type=DatasourceType.FILE.value,
  317. upload_file=file_detail,
  318. document_model=dataset_document.doc_form,
  319. )
  320. text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"])
  321. elif dataset_document.data_source_type == "notion_import":
  322. if (
  323. not data_source_info
  324. or "notion_workspace_id" not in data_source_info
  325. or "notion_page_id" not in data_source_info
  326. ):
  327. raise ValueError("no notion import info found")
  328. extract_setting = ExtractSetting(
  329. datasource_type=DatasourceType.NOTION.value,
  330. notion_info={
  331. "notion_workspace_id": data_source_info["notion_workspace_id"],
  332. "notion_obj_id": data_source_info["notion_page_id"],
  333. "notion_page_type": data_source_info["type"],
  334. "document": dataset_document,
  335. "tenant_id": dataset_document.tenant_id,
  336. },
  337. document_model=dataset_document.doc_form,
  338. )
  339. text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"])
  340. elif dataset_document.data_source_type == "website_crawl":
  341. if (
  342. not data_source_info
  343. or "provider" not in data_source_info
  344. or "url" not in data_source_info
  345. or "job_id" not in data_source_info
  346. ):
  347. raise ValueError("no website import info found")
  348. extract_setting = ExtractSetting(
  349. datasource_type=DatasourceType.WEBSITE.value,
  350. website_info={
  351. "provider": data_source_info["provider"],
  352. "job_id": data_source_info["job_id"],
  353. "tenant_id": dataset_document.tenant_id,
  354. "url": data_source_info["url"],
  355. "mode": data_source_info["mode"],
  356. "only_main_content": data_source_info["only_main_content"],
  357. },
  358. document_model=dataset_document.doc_form,
  359. )
  360. text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"])
  361. # update document status to splitting
  362. self._update_document_index_status(
  363. document_id=dataset_document.id,
  364. after_indexing_status="splitting",
  365. extra_update_params={
  366. DatasetDocument.word_count: sum(len(text_doc.page_content) for text_doc in text_docs),
  367. DatasetDocument.parsing_completed_at: naive_utc_now(),
  368. },
  369. )
  370. # replace doc id to document model id
  371. text_docs = cast(list[Document], text_docs)
  372. for text_doc in text_docs:
  373. if text_doc.metadata is not None:
  374. text_doc.metadata["document_id"] = dataset_document.id
  375. text_doc.metadata["dataset_id"] = dataset_document.dataset_id
  376. return text_docs
  377. @staticmethod
  378. def filter_string(text):
  379. text = re.sub(r"<\|", "<", text)
  380. text = re.sub(r"\|>", ">", text)
  381. text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F\xEF\xBF\xBE]", "", text)
  382. # Unicode U+FFFE
  383. text = re.sub("\ufffe", "", text)
  384. return text
  385. @staticmethod
  386. def _get_splitter(
  387. processing_rule_mode: str,
  388. max_tokens: int,
  389. chunk_overlap: int,
  390. separator: str,
  391. embedding_model_instance: Optional[ModelInstance],
  392. ) -> TextSplitter:
  393. """
  394. Get the NodeParser object according to the processing rule.
  395. """
  396. if processing_rule_mode in ["custom", "hierarchical"]:
  397. # The user-defined segmentation rule
  398. max_segmentation_tokens_length = dify_config.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH
  399. if max_tokens < 50 or max_tokens > max_segmentation_tokens_length:
  400. raise ValueError(f"Custom segment length should be between 50 and {max_segmentation_tokens_length}.")
  401. if separator:
  402. separator = separator.replace("\\n", "\n")
  403. character_splitter = FixedRecursiveCharacterTextSplitter.from_encoder(
  404. chunk_size=max_tokens,
  405. chunk_overlap=chunk_overlap,
  406. fixed_separator=separator,
  407. separators=["\n\n", "。", ". ", " ", ""],
  408. embedding_model_instance=embedding_model_instance,
  409. )
  410. else:
  411. # Automatic segmentation
  412. automatic_rules: dict[str, Any] = dict(DatasetProcessRule.AUTOMATIC_RULES["segmentation"])
  413. character_splitter = EnhanceRecursiveCharacterTextSplitter.from_encoder(
  414. chunk_size=automatic_rules["max_tokens"],
  415. chunk_overlap=automatic_rules["chunk_overlap"],
  416. separators=["\n\n", "。", ". ", " ", ""],
  417. embedding_model_instance=embedding_model_instance,
  418. )
  419. return character_splitter # type: ignore
  420. def _split_to_documents_for_estimate(
  421. self, text_docs: list[Document], splitter: TextSplitter, processing_rule: DatasetProcessRule
  422. ) -> list[Document]:
  423. """
  424. Split the text documents into nodes.
  425. """
  426. all_documents: list[Document] = []
  427. for text_doc in text_docs:
  428. # document clean
  429. document_text = self._document_clean(text_doc.page_content, processing_rule)
  430. text_doc.page_content = document_text
  431. # parse document to nodes
  432. documents = splitter.split_documents([text_doc])
  433. split_documents = []
  434. for document in documents:
  435. if document.page_content is None or not document.page_content.strip():
  436. continue
  437. if document.metadata is not None:
  438. doc_id = str(uuid.uuid4())
  439. hash = helper.generate_text_hash(document.page_content)
  440. document.metadata["doc_id"] = doc_id
  441. document.metadata["doc_hash"] = hash
  442. split_documents.append(document)
  443. all_documents.extend(split_documents)
  444. return all_documents
  445. @staticmethod
  446. def _document_clean(text: str, processing_rule: DatasetProcessRule) -> str:
  447. """
  448. Clean the document text according to the processing rules.
  449. """
  450. if processing_rule.mode == "automatic":
  451. rules = DatasetProcessRule.AUTOMATIC_RULES
  452. else:
  453. rules = json.loads(processing_rule.rules) if processing_rule.rules else {}
  454. document_text = CleanProcessor.clean(text, {"rules": rules})
  455. return document_text
  456. @staticmethod
  457. def format_split_text(text: str) -> list[QAPreviewDetail]:
  458. regex = r"Q\d+:\s*(.*?)\s*A\d+:\s*([\s\S]*?)(?=Q\d+:|$)"
  459. matches = re.findall(regex, text, re.UNICODE)
  460. return [QAPreviewDetail(question=q, answer=re.sub(r"\n\s*", "\n", a.strip())) for q, a in matches if q and a]
  461. def _load(
  462. self,
  463. index_processor: BaseIndexProcessor,
  464. dataset: Dataset,
  465. dataset_document: DatasetDocument,
  466. documents: list[Document],
  467. ) -> None:
  468. """
  469. insert index and update document/segment status to completed
  470. """
  471. embedding_model_instance = None
  472. if dataset.indexing_technique == "high_quality":
  473. embedding_model_instance = self.model_manager.get_model_instance(
  474. tenant_id=dataset.tenant_id,
  475. provider=dataset.embedding_model_provider,
  476. model_type=ModelType.TEXT_EMBEDDING,
  477. model=dataset.embedding_model,
  478. )
  479. # chunk nodes by chunk size
  480. indexing_start_at = time.perf_counter()
  481. tokens = 0
  482. if dataset_document.doc_form != IndexType.PARENT_CHILD_INDEX and dataset.indexing_technique == "economy":
  483. # create keyword index
  484. create_keyword_thread = threading.Thread(
  485. target=self._process_keyword_index,
  486. args=(current_app._get_current_object(), dataset.id, dataset_document.id, documents), # type: ignore
  487. )
  488. create_keyword_thread.start()
  489. max_workers = 10
  490. if dataset.indexing_technique == "high_quality":
  491. with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
  492. futures = []
  493. # Distribute documents into multiple groups based on the hash values of page_content
  494. # This is done to prevent multiple threads from processing the same document,
  495. # Thereby avoiding potential database insertion deadlocks
  496. document_groups: list[list[Document]] = [[] for _ in range(max_workers)]
  497. for document in documents:
  498. hash = helper.generate_text_hash(document.page_content)
  499. group_index = int(hash, 16) % max_workers
  500. document_groups[group_index].append(document)
  501. for chunk_documents in document_groups:
  502. if len(chunk_documents) == 0:
  503. continue
  504. futures.append(
  505. executor.submit(
  506. self._process_chunk,
  507. current_app._get_current_object(), # type: ignore
  508. index_processor,
  509. chunk_documents,
  510. dataset,
  511. dataset_document,
  512. embedding_model_instance,
  513. )
  514. )
  515. for future in futures:
  516. tokens += future.result()
  517. if dataset_document.doc_form != IndexType.PARENT_CHILD_INDEX and dataset.indexing_technique == "economy":
  518. create_keyword_thread.join()
  519. indexing_end_at = time.perf_counter()
  520. # update document status to completed
  521. self._update_document_index_status(
  522. document_id=dataset_document.id,
  523. after_indexing_status="completed",
  524. extra_update_params={
  525. DatasetDocument.tokens: tokens,
  526. DatasetDocument.completed_at: naive_utc_now(),
  527. DatasetDocument.indexing_latency: indexing_end_at - indexing_start_at,
  528. DatasetDocument.error: None,
  529. },
  530. )
  531. @staticmethod
  532. def _process_keyword_index(flask_app, dataset_id, document_id, documents):
  533. with flask_app.app_context():
  534. dataset = db.session.query(Dataset).filter_by(id=dataset_id).first()
  535. if not dataset:
  536. raise ValueError("no dataset found")
  537. keyword = Keyword(dataset)
  538. keyword.create(documents)
  539. if dataset.indexing_technique != "high_quality":
  540. document_ids = [document.metadata["doc_id"] for document in documents]
  541. db.session.query(DocumentSegment).where(
  542. DocumentSegment.document_id == document_id,
  543. DocumentSegment.dataset_id == dataset_id,
  544. DocumentSegment.index_node_id.in_(document_ids),
  545. DocumentSegment.status == "indexing",
  546. ).update(
  547. {
  548. DocumentSegment.status: "completed",
  549. DocumentSegment.enabled: True,
  550. DocumentSegment.completed_at: naive_utc_now(),
  551. }
  552. )
  553. db.session.commit()
  554. def _process_chunk(
  555. self, flask_app, index_processor, chunk_documents, dataset, dataset_document, embedding_model_instance
  556. ):
  557. with flask_app.app_context():
  558. # check document is paused
  559. self._check_document_paused_status(dataset_document.id)
  560. tokens = 0
  561. if embedding_model_instance:
  562. page_content_list = [document.page_content for document in chunk_documents]
  563. tokens += sum(embedding_model_instance.get_text_embedding_num_tokens(page_content_list))
  564. # load index
  565. index_processor.load(dataset, chunk_documents, with_keywords=False)
  566. document_ids = [document.metadata["doc_id"] for document in chunk_documents]
  567. db.session.query(DocumentSegment).where(
  568. DocumentSegment.document_id == dataset_document.id,
  569. DocumentSegment.dataset_id == dataset.id,
  570. DocumentSegment.index_node_id.in_(document_ids),
  571. DocumentSegment.status == "indexing",
  572. ).update(
  573. {
  574. DocumentSegment.status: "completed",
  575. DocumentSegment.enabled: True,
  576. DocumentSegment.completed_at: naive_utc_now(),
  577. }
  578. )
  579. db.session.commit()
  580. return tokens
  581. @staticmethod
  582. def _check_document_paused_status(document_id: str):
  583. indexing_cache_key = f"document_{document_id}_is_paused"
  584. result = redis_client.get(indexing_cache_key)
  585. if result:
  586. raise DocumentIsPausedError()
  587. @staticmethod
  588. def _update_document_index_status(
  589. document_id: str, after_indexing_status: str, extra_update_params: Optional[dict] = None
  590. ) -> None:
  591. """
  592. Update the document indexing status.
  593. """
  594. count = db.session.query(DatasetDocument).filter_by(id=document_id, is_paused=True).count()
  595. if count > 0:
  596. raise DocumentIsPausedError()
  597. document = db.session.query(DatasetDocument).filter_by(id=document_id).first()
  598. if not document:
  599. raise DocumentIsDeletedPausedError()
  600. update_params = {DatasetDocument.indexing_status: after_indexing_status}
  601. if extra_update_params:
  602. update_params.update(extra_update_params)
  603. db.session.query(DatasetDocument).filter_by(id=document_id).update(update_params) # type: ignore
  604. db.session.commit()
  605. @staticmethod
  606. def _update_segments_by_document(dataset_document_id: str, update_params: dict) -> None:
  607. """
  608. Update the document segment by document id.
  609. """
  610. db.session.query(DocumentSegment).filter_by(document_id=dataset_document_id).update(update_params)
  611. db.session.commit()
  612. def _transform(
  613. self,
  614. index_processor: BaseIndexProcessor,
  615. dataset: Dataset,
  616. text_docs: list[Document],
  617. doc_language: str,
  618. process_rule: dict,
  619. ) -> list[Document]:
  620. # get embedding model instance
  621. embedding_model_instance = None
  622. if dataset.indexing_technique == "high_quality":
  623. if dataset.embedding_model_provider:
  624. embedding_model_instance = self.model_manager.get_model_instance(
  625. tenant_id=dataset.tenant_id,
  626. provider=dataset.embedding_model_provider,
  627. model_type=ModelType.TEXT_EMBEDDING,
  628. model=dataset.embedding_model,
  629. )
  630. else:
  631. embedding_model_instance = self.model_manager.get_default_model_instance(
  632. tenant_id=dataset.tenant_id,
  633. model_type=ModelType.TEXT_EMBEDDING,
  634. )
  635. documents = index_processor.transform(
  636. text_docs,
  637. embedding_model_instance=embedding_model_instance,
  638. process_rule=process_rule,
  639. tenant_id=dataset.tenant_id,
  640. doc_language=doc_language,
  641. )
  642. return documents
  643. def _load_segments(self, dataset, dataset_document, documents):
  644. # save node to document segment
  645. doc_store = DatasetDocumentStore(
  646. dataset=dataset, user_id=dataset_document.created_by, document_id=dataset_document.id
  647. )
  648. # add document segments
  649. doc_store.add_documents(docs=documents, save_child=dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX)
  650. # update document status to indexing
  651. cur_time = naive_utc_now()
  652. self._update_document_index_status(
  653. document_id=dataset_document.id,
  654. after_indexing_status="indexing",
  655. extra_update_params={
  656. DatasetDocument.cleaning_completed_at: cur_time,
  657. DatasetDocument.splitting_completed_at: cur_time,
  658. },
  659. )
  660. # update segment status to indexing
  661. self._update_segments_by_document(
  662. dataset_document_id=dataset_document.id,
  663. update_params={
  664. DocumentSegment.status: "indexing",
  665. DocumentSegment.indexing_at: naive_utc_now(),
  666. },
  667. )
  668. pass
  669. class DocumentIsPausedError(Exception):
  670. pass
  671. class DocumentIsDeletedPausedError(Exception):
  672. pass