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 33KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755
  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
  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. # keep separate, avoid union-list ambiguity
  252. preview_texts: list[PreviewDetail] = []
  253. qa_preview_texts: list[QAPreviewDetail] = []
  254. total_segments = 0
  255. index_type = doc_form
  256. index_processor = IndexProcessorFactory(index_type).init_index_processor()
  257. for extract_setting in extract_settings:
  258. # extract
  259. processing_rule = DatasetProcessRule(
  260. mode=tmp_processing_rule["mode"], rules=json.dumps(tmp_processing_rule["rules"])
  261. )
  262. text_docs = index_processor.extract(extract_setting, process_rule_mode=tmp_processing_rule["mode"])
  263. documents = index_processor.transform(
  264. text_docs,
  265. embedding_model_instance=embedding_model_instance,
  266. process_rule=processing_rule.to_dict(),
  267. tenant_id=tenant_id,
  268. doc_language=doc_language,
  269. preview=True,
  270. )
  271. total_segments += len(documents)
  272. for document in documents:
  273. if len(preview_texts) < 10:
  274. if doc_form and doc_form == "qa_model":
  275. qa_detail = QAPreviewDetail(
  276. question=document.page_content, answer=document.metadata.get("answer") or ""
  277. )
  278. qa_preview_texts.append(qa_detail)
  279. else:
  280. preview_detail = PreviewDetail(content=document.page_content)
  281. if document.children:
  282. preview_detail.child_chunks = [child.page_content for child in document.children]
  283. preview_texts.append(preview_detail)
  284. # delete image files and related db records
  285. image_upload_file_ids = get_image_upload_file_ids(document.page_content)
  286. for upload_file_id in image_upload_file_ids:
  287. stmt = select(UploadFile).where(UploadFile.id == upload_file_id)
  288. image_file = db.session.scalar(stmt)
  289. if image_file is None:
  290. continue
  291. try:
  292. storage.delete(image_file.key)
  293. except Exception:
  294. logger.exception(
  295. "Delete image_files failed while indexing_estimate, \
  296. image_upload_file_is: %s",
  297. upload_file_id,
  298. )
  299. db.session.delete(image_file)
  300. if doc_form and doc_form == "qa_model":
  301. return IndexingEstimate(total_segments=total_segments * 20, qa_preview=qa_preview_texts, preview=[])
  302. return IndexingEstimate(total_segments=total_segments, preview=preview_texts)
  303. def _extract(
  304. self, index_processor: BaseIndexProcessor, dataset_document: DatasetDocument, process_rule: dict
  305. ) -> list[Document]:
  306. # load file
  307. if dataset_document.data_source_type not in {"upload_file", "notion_import", "website_crawl"}:
  308. return []
  309. data_source_info = dataset_document.data_source_info_dict
  310. text_docs = []
  311. if dataset_document.data_source_type == "upload_file":
  312. if not data_source_info or "upload_file_id" not in data_source_info:
  313. raise ValueError("no upload file found")
  314. stmt = select(UploadFile).where(UploadFile.id == data_source_info["upload_file_id"])
  315. file_detail = db.session.scalars(stmt).one_or_none()
  316. if file_detail:
  317. extract_setting = ExtractSetting(
  318. datasource_type=DatasourceType.FILE.value,
  319. upload_file=file_detail,
  320. document_model=dataset_document.doc_form,
  321. )
  322. text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"])
  323. elif dataset_document.data_source_type == "notion_import":
  324. if (
  325. not data_source_info
  326. or "notion_workspace_id" not in data_source_info
  327. or "notion_page_id" not in data_source_info
  328. ):
  329. raise ValueError("no notion import info found")
  330. extract_setting = ExtractSetting(
  331. datasource_type=DatasourceType.NOTION.value,
  332. notion_info={
  333. "notion_workspace_id": data_source_info["notion_workspace_id"],
  334. "notion_obj_id": data_source_info["notion_page_id"],
  335. "notion_page_type": data_source_info["type"],
  336. "document": dataset_document,
  337. "tenant_id": dataset_document.tenant_id,
  338. },
  339. document_model=dataset_document.doc_form,
  340. )
  341. text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"])
  342. elif dataset_document.data_source_type == "website_crawl":
  343. if (
  344. not data_source_info
  345. or "provider" not in data_source_info
  346. or "url" not in data_source_info
  347. or "job_id" not in data_source_info
  348. ):
  349. raise ValueError("no website import info found")
  350. extract_setting = ExtractSetting(
  351. datasource_type=DatasourceType.WEBSITE.value,
  352. website_info={
  353. "provider": data_source_info["provider"],
  354. "job_id": data_source_info["job_id"],
  355. "tenant_id": dataset_document.tenant_id,
  356. "url": data_source_info["url"],
  357. "mode": data_source_info["mode"],
  358. "only_main_content": data_source_info["only_main_content"],
  359. },
  360. document_model=dataset_document.doc_form,
  361. )
  362. text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule["mode"])
  363. # update document status to splitting
  364. self._update_document_index_status(
  365. document_id=dataset_document.id,
  366. after_indexing_status="splitting",
  367. extra_update_params={
  368. DatasetDocument.word_count: sum(len(text_doc.page_content) for text_doc in text_docs),
  369. DatasetDocument.parsing_completed_at: naive_utc_now(),
  370. },
  371. )
  372. # replace doc id to document model id
  373. for text_doc in text_docs:
  374. if text_doc.metadata is not None:
  375. text_doc.metadata["document_id"] = dataset_document.id
  376. text_doc.metadata["dataset_id"] = dataset_document.dataset_id
  377. return text_docs
  378. @staticmethod
  379. def filter_string(text):
  380. text = re.sub(r"<\|", "<", text)
  381. text = re.sub(r"\|>", ">", text)
  382. text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F\xEF\xBF\xBE]", "", text)
  383. # Unicode U+FFFE
  384. text = re.sub("\ufffe", "", text)
  385. return text
  386. @staticmethod
  387. def _get_splitter(
  388. processing_rule_mode: str,
  389. max_tokens: int,
  390. chunk_overlap: int,
  391. separator: str,
  392. embedding_model_instance: Optional[ModelInstance],
  393. ) -> TextSplitter:
  394. """
  395. Get the NodeParser object according to the processing rule.
  396. """
  397. character_splitter: TextSplitter
  398. if processing_rule_mode in ["custom", "hierarchical"]:
  399. # The user-defined segmentation rule
  400. max_segmentation_tokens_length = dify_config.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH
  401. if max_tokens < 50 or max_tokens > max_segmentation_tokens_length:
  402. raise ValueError(f"Custom segment length should be between 50 and {max_segmentation_tokens_length}.")
  403. if separator:
  404. separator = separator.replace("\\n", "\n")
  405. character_splitter = FixedRecursiveCharacterTextSplitter.from_encoder(
  406. chunk_size=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 segmentation
  414. automatic_rules: dict[str, Any] = dict(DatasetProcessRule.AUTOMATIC_RULES["segmentation"])
  415. character_splitter = EnhanceRecursiveCharacterTextSplitter.from_encoder(
  416. chunk_size=automatic_rules["max_tokens"],
  417. chunk_overlap=automatic_rules["chunk_overlap"],
  418. separators=["\n\n", "。", ". ", " ", ""],
  419. embedding_model_instance=embedding_model_instance,
  420. )
  421. return character_splitter
  422. def _split_to_documents_for_estimate(
  423. self, text_docs: list[Document], splitter: TextSplitter, processing_rule: DatasetProcessRule
  424. ) -> list[Document]:
  425. """
  426. Split the text documents into nodes.
  427. """
  428. all_documents: list[Document] = []
  429. for text_doc in text_docs:
  430. # document clean
  431. document_text = self._document_clean(text_doc.page_content, processing_rule)
  432. text_doc.page_content = document_text
  433. # parse document to nodes
  434. documents = splitter.split_documents([text_doc])
  435. split_documents = []
  436. for document in documents:
  437. if document.page_content is None or not document.page_content.strip():
  438. continue
  439. if document.metadata is not None:
  440. doc_id = str(uuid.uuid4())
  441. hash = helper.generate_text_hash(document.page_content)
  442. document.metadata["doc_id"] = doc_id
  443. document.metadata["doc_hash"] = hash
  444. split_documents.append(document)
  445. all_documents.extend(split_documents)
  446. return all_documents
  447. @staticmethod
  448. def _document_clean(text: str, processing_rule: DatasetProcessRule) -> str:
  449. """
  450. Clean the document text according to the processing rules.
  451. """
  452. if processing_rule.mode == "automatic":
  453. rules = DatasetProcessRule.AUTOMATIC_RULES
  454. else:
  455. rules = json.loads(processing_rule.rules) if processing_rule.rules else {}
  456. document_text = CleanProcessor.clean(text, {"rules": rules})
  457. return document_text
  458. @staticmethod
  459. def format_split_text(text: str) -> list[QAPreviewDetail]:
  460. regex = r"Q\d+:\s*(.*?)\s*A\d+:\s*([\s\S]*?)(?=Q\d+:|$)"
  461. matches = re.findall(regex, text, re.UNICODE)
  462. return [QAPreviewDetail(question=q, answer=re.sub(r"\n\s*", "\n", a.strip())) for q, a in matches if q and a]
  463. def _load(
  464. self,
  465. index_processor: BaseIndexProcessor,
  466. dataset: Dataset,
  467. dataset_document: DatasetDocument,
  468. documents: list[Document],
  469. ):
  470. """
  471. insert index and update document/segment status to completed
  472. """
  473. embedding_model_instance = None
  474. if dataset.indexing_technique == "high_quality":
  475. embedding_model_instance = self.model_manager.get_model_instance(
  476. tenant_id=dataset.tenant_id,
  477. provider=dataset.embedding_model_provider,
  478. model_type=ModelType.TEXT_EMBEDDING,
  479. model=dataset.embedding_model,
  480. )
  481. # chunk nodes by chunk size
  482. indexing_start_at = time.perf_counter()
  483. tokens = 0
  484. create_keyword_thread = None
  485. if dataset_document.doc_form != IndexType.PARENT_CHILD_INDEX and dataset.indexing_technique == "economy":
  486. # create keyword index
  487. create_keyword_thread = threading.Thread(
  488. target=self._process_keyword_index,
  489. args=(current_app._get_current_object(), dataset.id, dataset_document.id, documents), # type: ignore
  490. )
  491. create_keyword_thread.start()
  492. max_workers = 10
  493. if dataset.indexing_technique == "high_quality":
  494. with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
  495. futures = []
  496. # Distribute documents into multiple groups based on the hash values of page_content
  497. # This is done to prevent multiple threads from processing the same document,
  498. # Thereby avoiding potential database insertion deadlocks
  499. document_groups: list[list[Document]] = [[] for _ in range(max_workers)]
  500. for document in documents:
  501. hash = helper.generate_text_hash(document.page_content)
  502. group_index = int(hash, 16) % max_workers
  503. document_groups[group_index].append(document)
  504. for chunk_documents in document_groups:
  505. if len(chunk_documents) == 0:
  506. continue
  507. futures.append(
  508. executor.submit(
  509. self._process_chunk,
  510. current_app._get_current_object(), # type: ignore
  511. index_processor,
  512. chunk_documents,
  513. dataset,
  514. dataset_document,
  515. embedding_model_instance,
  516. )
  517. )
  518. for future in futures:
  519. tokens += future.result()
  520. if (
  521. dataset_document.doc_form != IndexType.PARENT_CHILD_INDEX
  522. and dataset.indexing_technique == "economy"
  523. and create_keyword_thread is not None
  524. ):
  525. create_keyword_thread.join()
  526. indexing_end_at = time.perf_counter()
  527. # update document status to completed
  528. self._update_document_index_status(
  529. document_id=dataset_document.id,
  530. after_indexing_status="completed",
  531. extra_update_params={
  532. DatasetDocument.tokens: tokens,
  533. DatasetDocument.completed_at: naive_utc_now(),
  534. DatasetDocument.indexing_latency: indexing_end_at - indexing_start_at,
  535. DatasetDocument.error: None,
  536. },
  537. )
  538. @staticmethod
  539. def _process_keyword_index(flask_app, dataset_id, document_id, documents):
  540. with flask_app.app_context():
  541. dataset = db.session.query(Dataset).filter_by(id=dataset_id).first()
  542. if not dataset:
  543. raise ValueError("no dataset found")
  544. keyword = Keyword(dataset)
  545. keyword.create(documents)
  546. if dataset.indexing_technique != "high_quality":
  547. document_ids = [document.metadata["doc_id"] for document in documents]
  548. db.session.query(DocumentSegment).where(
  549. DocumentSegment.document_id == document_id,
  550. DocumentSegment.dataset_id == dataset_id,
  551. DocumentSegment.index_node_id.in_(document_ids),
  552. DocumentSegment.status == "indexing",
  553. ).update(
  554. {
  555. DocumentSegment.status: "completed",
  556. DocumentSegment.enabled: True,
  557. DocumentSegment.completed_at: naive_utc_now(),
  558. }
  559. )
  560. db.session.commit()
  561. def _process_chunk(
  562. self, flask_app, index_processor, chunk_documents, dataset, dataset_document, embedding_model_instance
  563. ):
  564. with flask_app.app_context():
  565. # check document is paused
  566. self._check_document_paused_status(dataset_document.id)
  567. tokens = 0
  568. if embedding_model_instance:
  569. page_content_list = [document.page_content for document in chunk_documents]
  570. tokens += sum(embedding_model_instance.get_text_embedding_num_tokens(page_content_list))
  571. # load index
  572. index_processor.load(dataset, chunk_documents, with_keywords=False)
  573. document_ids = [document.metadata["doc_id"] for document in chunk_documents]
  574. db.session.query(DocumentSegment).where(
  575. DocumentSegment.document_id == dataset_document.id,
  576. DocumentSegment.dataset_id == dataset.id,
  577. DocumentSegment.index_node_id.in_(document_ids),
  578. DocumentSegment.status == "indexing",
  579. ).update(
  580. {
  581. DocumentSegment.status: "completed",
  582. DocumentSegment.enabled: True,
  583. DocumentSegment.completed_at: naive_utc_now(),
  584. }
  585. )
  586. db.session.commit()
  587. return tokens
  588. @staticmethod
  589. def _check_document_paused_status(document_id: str):
  590. indexing_cache_key = f"document_{document_id}_is_paused"
  591. result = redis_client.get(indexing_cache_key)
  592. if result:
  593. raise DocumentIsPausedError()
  594. @staticmethod
  595. def _update_document_index_status(
  596. document_id: str, after_indexing_status: str, extra_update_params: Optional[dict] = None
  597. ):
  598. """
  599. Update the document indexing status.
  600. """
  601. count = db.session.query(DatasetDocument).filter_by(id=document_id, is_paused=True).count()
  602. if count > 0:
  603. raise DocumentIsPausedError()
  604. document = db.session.query(DatasetDocument).filter_by(id=document_id).first()
  605. if not document:
  606. raise DocumentIsDeletedPausedError()
  607. update_params = {DatasetDocument.indexing_status: after_indexing_status}
  608. if extra_update_params:
  609. update_params.update(extra_update_params)
  610. db.session.query(DatasetDocument).filter_by(id=document_id).update(update_params) # type: ignore
  611. db.session.commit()
  612. @staticmethod
  613. def _update_segments_by_document(dataset_document_id: str, update_params: dict):
  614. """
  615. Update the document segment by document id.
  616. """
  617. db.session.query(DocumentSegment).filter_by(document_id=dataset_document_id).update(update_params)
  618. db.session.commit()
  619. def _transform(
  620. self,
  621. index_processor: BaseIndexProcessor,
  622. dataset: Dataset,
  623. text_docs: list[Document],
  624. doc_language: str,
  625. process_rule: dict,
  626. ) -> list[Document]:
  627. # get embedding model instance
  628. embedding_model_instance = None
  629. if dataset.indexing_technique == "high_quality":
  630. if dataset.embedding_model_provider:
  631. embedding_model_instance = self.model_manager.get_model_instance(
  632. tenant_id=dataset.tenant_id,
  633. provider=dataset.embedding_model_provider,
  634. model_type=ModelType.TEXT_EMBEDDING,
  635. model=dataset.embedding_model,
  636. )
  637. else:
  638. embedding_model_instance = self.model_manager.get_default_model_instance(
  639. tenant_id=dataset.tenant_id,
  640. model_type=ModelType.TEXT_EMBEDDING,
  641. )
  642. documents = index_processor.transform(
  643. text_docs,
  644. embedding_model_instance=embedding_model_instance,
  645. process_rule=process_rule,
  646. tenant_id=dataset.tenant_id,
  647. doc_language=doc_language,
  648. )
  649. return documents
  650. def _load_segments(self, dataset, dataset_document, documents):
  651. # save node to document segment
  652. doc_store = DatasetDocumentStore(
  653. dataset=dataset, user_id=dataset_document.created_by, document_id=dataset_document.id
  654. )
  655. # add document segments
  656. doc_store.add_documents(docs=documents, save_child=dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX)
  657. # update document status to indexing
  658. cur_time = naive_utc_now()
  659. self._update_document_index_status(
  660. document_id=dataset_document.id,
  661. after_indexing_status="indexing",
  662. extra_update_params={
  663. DatasetDocument.cleaning_completed_at: cur_time,
  664. DatasetDocument.splitting_completed_at: cur_time,
  665. },
  666. )
  667. # update segment status to indexing
  668. self._update_segments_by_document(
  669. dataset_document_id=dataset_document.id,
  670. update_params={
  671. DocumentSegment.status: "indexing",
  672. DocumentSegment.indexing_at: naive_utc_now(),
  673. },
  674. )
  675. pass
  676. class DocumentIsPausedError(Exception):
  677. pass
  678. class DocumentIsDeletedPausedError(Exception):
  679. pass