indexing_runner.py 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773
  1. import datetime
  2. import json
  3. import logging
  4. import re
  5. import threading
  6. import time
  7. import uuid
  8. from typing import Optional, List, cast
  9. from flask import current_app, Flask
  10. from flask_login import current_user
  11. from langchain.schema import Document
  12. from langchain.text_splitter import RecursiveCharacterTextSplitter, TextSplitter
  13. from sqlalchemy.orm.exc import ObjectDeletedError
  14. from core.data_loader.file_extractor import FileExtractor
  15. from core.data_loader.loader.notion import NotionLoader
  16. from core.docstore.dataset_docstore import DatesetDocumentStore
  17. from core.generator.llm_generator import LLMGenerator
  18. from core.index.index import IndexBuilder
  19. from core.model_providers.error import ProviderTokenNotInitError
  20. from core.model_providers.model_factory import ModelFactory
  21. from core.model_providers.models.entity.message import MessageType
  22. from core.spiltter.fixed_text_splitter import FixedRecursiveCharacterTextSplitter
  23. from extensions.ext_database import db
  24. from extensions.ext_redis import redis_client
  25. from extensions.ext_storage import storage
  26. from libs import helper
  27. from models.dataset import Document as DatasetDocument
  28. from models.dataset import Dataset, DocumentSegment, DatasetProcessRule
  29. from models.model import UploadFile
  30. from models.source import DataSourceBinding
  31. class IndexingRunner:
  32. def __init__(self):
  33. self.storage = storage
  34. def run(self, dataset_documents: List[DatasetDocument]):
  35. """Run the indexing process."""
  36. for dataset_document in dataset_documents:
  37. try:
  38. # get dataset
  39. dataset = Dataset.query.filter_by(
  40. id=dataset_document.dataset_id
  41. ).first()
  42. if not dataset:
  43. raise ValueError("no dataset found")
  44. # get the process rule
  45. processing_rule = db.session.query(DatasetProcessRule). \
  46. filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id). \
  47. first()
  48. # load file
  49. text_docs = self._load_data(dataset_document)
  50. # get splitter
  51. splitter = self._get_splitter(processing_rule)
  52. # split to documents
  53. documents = self._step_split(
  54. text_docs=text_docs,
  55. splitter=splitter,
  56. dataset=dataset,
  57. dataset_document=dataset_document,
  58. processing_rule=processing_rule
  59. )
  60. self._build_index(
  61. dataset=dataset,
  62. dataset_document=dataset_document,
  63. documents=documents
  64. )
  65. except DocumentIsPausedException:
  66. raise DocumentIsPausedException('Document paused, document id: {}'.format(dataset_document.id))
  67. except ProviderTokenNotInitError as e:
  68. dataset_document.indexing_status = 'error'
  69. dataset_document.error = str(e.description)
  70. dataset_document.stopped_at = datetime.datetime.utcnow()
  71. db.session.commit()
  72. except ObjectDeletedError:
  73. logging.warning('Document deleted, document id: {}'.format(dataset_document.id))
  74. except Exception as e:
  75. logging.exception("consume document failed")
  76. dataset_document.indexing_status = 'error'
  77. dataset_document.error = str(e)
  78. dataset_document.stopped_at = datetime.datetime.utcnow()
  79. db.session.commit()
  80. def run_in_splitting_status(self, dataset_document: DatasetDocument):
  81. """Run the indexing process when the index_status is splitting."""
  82. try:
  83. # get dataset
  84. dataset = Dataset.query.filter_by(
  85. id=dataset_document.dataset_id
  86. ).first()
  87. if not dataset:
  88. raise ValueError("no dataset found")
  89. # get exist document_segment list and delete
  90. document_segments = DocumentSegment.query.filter_by(
  91. dataset_id=dataset.id,
  92. document_id=dataset_document.id
  93. ).all()
  94. db.session.delete(document_segments)
  95. db.session.commit()
  96. # load file
  97. text_docs = self._load_data(dataset_document)
  98. # get the process rule
  99. processing_rule = db.session.query(DatasetProcessRule). \
  100. filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id). \
  101. first()
  102. # get splitter
  103. splitter = self._get_splitter(processing_rule)
  104. # split to documents
  105. documents = self._step_split(
  106. text_docs=text_docs,
  107. splitter=splitter,
  108. dataset=dataset,
  109. dataset_document=dataset_document,
  110. processing_rule=processing_rule
  111. )
  112. # build index
  113. self._build_index(
  114. dataset=dataset,
  115. dataset_document=dataset_document,
  116. documents=documents
  117. )
  118. except DocumentIsPausedException:
  119. raise DocumentIsPausedException('Document paused, document id: {}'.format(dataset_document.id))
  120. except ProviderTokenNotInitError as e:
  121. dataset_document.indexing_status = 'error'
  122. dataset_document.error = str(e.description)
  123. dataset_document.stopped_at = datetime.datetime.utcnow()
  124. db.session.commit()
  125. except Exception as e:
  126. logging.exception("consume document failed")
  127. dataset_document.indexing_status = 'error'
  128. dataset_document.error = str(e)
  129. dataset_document.stopped_at = datetime.datetime.utcnow()
  130. db.session.commit()
  131. def run_in_indexing_status(self, dataset_document: DatasetDocument):
  132. """Run the indexing process when the index_status is indexing."""
  133. try:
  134. # get dataset
  135. dataset = Dataset.query.filter_by(
  136. id=dataset_document.dataset_id
  137. ).first()
  138. if not dataset:
  139. raise ValueError("no dataset found")
  140. # get exist document_segment list and delete
  141. document_segments = DocumentSegment.query.filter_by(
  142. dataset_id=dataset.id,
  143. document_id=dataset_document.id
  144. ).all()
  145. documents = []
  146. if document_segments:
  147. for document_segment in document_segments:
  148. # transform segment to node
  149. if document_segment.status != "completed":
  150. document = Document(
  151. page_content=document_segment.content,
  152. metadata={
  153. "doc_id": document_segment.index_node_id,
  154. "doc_hash": document_segment.index_node_hash,
  155. "document_id": document_segment.document_id,
  156. "dataset_id": document_segment.dataset_id,
  157. }
  158. )
  159. documents.append(document)
  160. # build index
  161. self._build_index(
  162. dataset=dataset,
  163. dataset_document=dataset_document,
  164. documents=documents
  165. )
  166. except DocumentIsPausedException:
  167. raise DocumentIsPausedException('Document paused, document id: {}'.format(dataset_document.id))
  168. except ProviderTokenNotInitError as e:
  169. dataset_document.indexing_status = 'error'
  170. dataset_document.error = str(e.description)
  171. dataset_document.stopped_at = datetime.datetime.utcnow()
  172. db.session.commit()
  173. except Exception as e:
  174. logging.exception("consume document failed")
  175. dataset_document.indexing_status = 'error'
  176. dataset_document.error = str(e)
  177. dataset_document.stopped_at = datetime.datetime.utcnow()
  178. db.session.commit()
  179. def file_indexing_estimate(self, tenant_id: str, file_details: List[UploadFile], tmp_processing_rule: dict,
  180. doc_form: str = None, doc_language: str = 'English', dataset_id: str = None,
  181. indexing_technique: str = 'economy') -> dict:
  182. """
  183. Estimate the indexing for the document.
  184. """
  185. embedding_model = None
  186. if dataset_id:
  187. dataset = Dataset.query.filter_by(
  188. id=dataset_id
  189. ).first()
  190. if not dataset:
  191. raise ValueError('Dataset not found.')
  192. if dataset.indexing_technique == 'high_quality' or indexing_technique == 'high_quality':
  193. embedding_model = ModelFactory.get_embedding_model(
  194. tenant_id=dataset.tenant_id,
  195. model_provider_name=dataset.embedding_model_provider,
  196. model_name=dataset.embedding_model
  197. )
  198. else:
  199. if indexing_technique == 'high_quality':
  200. embedding_model = ModelFactory.get_embedding_model(
  201. tenant_id=tenant_id
  202. )
  203. tokens = 0
  204. preview_texts = []
  205. total_segments = 0
  206. for file_detail in file_details:
  207. # load data from file
  208. text_docs = FileExtractor.load(file_detail)
  209. processing_rule = DatasetProcessRule(
  210. mode=tmp_processing_rule["mode"],
  211. rules=json.dumps(tmp_processing_rule["rules"])
  212. )
  213. # get splitter
  214. splitter = self._get_splitter(processing_rule)
  215. # split to documents
  216. documents = self._split_to_documents_for_estimate(
  217. text_docs=text_docs,
  218. splitter=splitter,
  219. processing_rule=processing_rule
  220. )
  221. total_segments += len(documents)
  222. for document in documents:
  223. if len(preview_texts) < 5:
  224. preview_texts.append(document.page_content)
  225. if indexing_technique == 'high_quality' or embedding_model:
  226. tokens += embedding_model.get_num_tokens(self.filter_string(document.page_content))
  227. if doc_form and doc_form == 'qa_model':
  228. text_generation_model = ModelFactory.get_text_generation_model(
  229. tenant_id=tenant_id
  230. )
  231. if len(preview_texts) > 0:
  232. # qa model document
  233. response = LLMGenerator.generate_qa_document(current_user.current_tenant_id, preview_texts[0],
  234. doc_language)
  235. document_qa_list = self.format_split_text(response)
  236. return {
  237. "total_segments": total_segments * 20,
  238. "tokens": total_segments * 2000,
  239. "total_price": '{:f}'.format(
  240. text_generation_model.calc_tokens_price(total_segments * 2000, MessageType.USER)),
  241. "currency": embedding_model.get_currency(),
  242. "qa_preview": document_qa_list,
  243. "preview": preview_texts
  244. }
  245. return {
  246. "total_segments": total_segments,
  247. "tokens": tokens,
  248. "total_price": '{:f}'.format(embedding_model.calc_tokens_price(tokens)) if embedding_model else 0,
  249. "currency": embedding_model.get_currency() if embedding_model else 'USD',
  250. "preview": preview_texts
  251. }
  252. def notion_indexing_estimate(self, tenant_id: str, notion_info_list: list, tmp_processing_rule: dict,
  253. doc_form: str = None, doc_language: str = 'English', dataset_id: str = None,
  254. indexing_technique: str = 'economy') -> dict:
  255. """
  256. Estimate the indexing for the document.
  257. """
  258. embedding_model = None
  259. if dataset_id:
  260. dataset = Dataset.query.filter_by(
  261. id=dataset_id
  262. ).first()
  263. if not dataset:
  264. raise ValueError('Dataset not found.')
  265. if dataset.indexing_technique == 'high_quality' or indexing_technique == 'high_quality':
  266. embedding_model = ModelFactory.get_embedding_model(
  267. tenant_id=dataset.tenant_id,
  268. model_provider_name=dataset.embedding_model_provider,
  269. model_name=dataset.embedding_model
  270. )
  271. else:
  272. if indexing_technique == 'high_quality':
  273. embedding_model = ModelFactory.get_embedding_model(
  274. tenant_id=tenant_id
  275. )
  276. # load data from notion
  277. tokens = 0
  278. preview_texts = []
  279. total_segments = 0
  280. for notion_info in notion_info_list:
  281. workspace_id = notion_info['workspace_id']
  282. data_source_binding = DataSourceBinding.query.filter(
  283. db.and_(
  284. DataSourceBinding.tenant_id == current_user.current_tenant_id,
  285. DataSourceBinding.provider == 'notion',
  286. DataSourceBinding.disabled == False,
  287. DataSourceBinding.source_info['workspace_id'] == f'"{workspace_id}"'
  288. )
  289. ).first()
  290. if not data_source_binding:
  291. raise ValueError('Data source binding not found.')
  292. for page in notion_info['pages']:
  293. loader = NotionLoader(
  294. notion_access_token=data_source_binding.access_token,
  295. notion_workspace_id=workspace_id,
  296. notion_obj_id=page['page_id'],
  297. notion_page_type=page['type']
  298. )
  299. documents = loader.load()
  300. processing_rule = DatasetProcessRule(
  301. mode=tmp_processing_rule["mode"],
  302. rules=json.dumps(tmp_processing_rule["rules"])
  303. )
  304. # get splitter
  305. splitter = self._get_splitter(processing_rule)
  306. # split to documents
  307. documents = self._split_to_documents_for_estimate(
  308. text_docs=documents,
  309. splitter=splitter,
  310. processing_rule=processing_rule
  311. )
  312. total_segments += len(documents)
  313. for document in documents:
  314. if len(preview_texts) < 5:
  315. preview_texts.append(document.page_content)
  316. if indexing_technique == 'high_quality' or embedding_model:
  317. tokens += embedding_model.get_num_tokens(document.page_content)
  318. if doc_form and doc_form == 'qa_model':
  319. text_generation_model = ModelFactory.get_text_generation_model(
  320. tenant_id=tenant_id
  321. )
  322. if len(preview_texts) > 0:
  323. # qa model document
  324. response = LLMGenerator.generate_qa_document(current_user.current_tenant_id, preview_texts[0],
  325. doc_language)
  326. document_qa_list = self.format_split_text(response)
  327. return {
  328. "total_segments": total_segments * 20,
  329. "tokens": total_segments * 2000,
  330. "total_price": '{:f}'.format(
  331. text_generation_model.calc_tokens_price(total_segments * 2000, MessageType.USER)),
  332. "currency": embedding_model.get_currency(),
  333. "qa_preview": document_qa_list,
  334. "preview": preview_texts
  335. }
  336. return {
  337. "total_segments": total_segments,
  338. "tokens": tokens,
  339. "total_price": '{:f}'.format(embedding_model.calc_tokens_price(tokens)) if embedding_model else 0,
  340. "currency": embedding_model.get_currency() if embedding_model else 'USD',
  341. "preview": preview_texts
  342. }
  343. def _load_data(self, dataset_document: DatasetDocument, automatic: bool = False) -> List[Document]:
  344. # load file
  345. if dataset_document.data_source_type not in ["upload_file", "notion_import"]:
  346. return []
  347. data_source_info = dataset_document.data_source_info_dict
  348. text_docs = []
  349. if dataset_document.data_source_type == 'upload_file':
  350. if not data_source_info or 'upload_file_id' not in data_source_info:
  351. raise ValueError("no upload file found")
  352. file_detail = db.session.query(UploadFile). \
  353. filter(UploadFile.id == data_source_info['upload_file_id']). \
  354. one_or_none()
  355. if file_detail:
  356. text_docs = FileExtractor.load(file_detail, is_automatic=False)
  357. elif dataset_document.data_source_type == 'notion_import':
  358. loader = NotionLoader.from_document(dataset_document)
  359. text_docs = loader.load()
  360. # update document status to splitting
  361. self._update_document_index_status(
  362. document_id=dataset_document.id,
  363. after_indexing_status="splitting",
  364. extra_update_params={
  365. DatasetDocument.word_count: sum([len(text_doc.page_content) for text_doc in text_docs]),
  366. DatasetDocument.parsing_completed_at: datetime.datetime.utcnow()
  367. }
  368. )
  369. # replace doc id to document model id
  370. text_docs = cast(List[Document], text_docs)
  371. for text_doc in text_docs:
  372. # remove invalid symbol
  373. text_doc.page_content = self.filter_string(text_doc.page_content)
  374. text_doc.metadata['document_id'] = dataset_document.id
  375. text_doc.metadata['dataset_id'] = dataset_document.dataset_id
  376. return text_docs
  377. def filter_string(self, text):
  378. text = re.sub(r'<\|', '<', text)
  379. text = re.sub(r'\|>', '>', text)
  380. text = re.sub(r'[\x00-\x08\x0B\x0C\x0E-\x1F\x7F\x80-\xFF]', '', text)
  381. return text
  382. def _get_splitter(self, processing_rule: DatasetProcessRule) -> TextSplitter:
  383. """
  384. Get the NodeParser object according to the processing rule.
  385. """
  386. if processing_rule.mode == "custom":
  387. # The user-defined segmentation rule
  388. rules = json.loads(processing_rule.rules)
  389. segmentation = rules["segmentation"]
  390. if segmentation["max_tokens"] < 50 or segmentation["max_tokens"] > 1000:
  391. raise ValueError("Custom segment length should be between 50 and 1000.")
  392. separator = segmentation["separator"]
  393. if separator:
  394. separator = separator.replace('\\n', '\n')
  395. character_splitter = FixedRecursiveCharacterTextSplitter.from_tiktoken_encoder(
  396. chunk_size=segmentation["max_tokens"],
  397. chunk_overlap=0,
  398. fixed_separator=separator,
  399. separators=["\n\n", "。", ".", " ", ""]
  400. )
  401. else:
  402. # Automatic segmentation
  403. character_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
  404. chunk_size=DatasetProcessRule.AUTOMATIC_RULES['segmentation']['max_tokens'],
  405. chunk_overlap=0,
  406. separators=["\n\n", "。", ".", " ", ""]
  407. )
  408. return character_splitter
  409. def _step_split(self, text_docs: List[Document], splitter: TextSplitter,
  410. dataset: Dataset, dataset_document: DatasetDocument, processing_rule: DatasetProcessRule) \
  411. -> List[Document]:
  412. """
  413. Split the text documents into documents and save them to the document segment.
  414. """
  415. documents = self._split_to_documents(
  416. text_docs=text_docs,
  417. splitter=splitter,
  418. processing_rule=processing_rule,
  419. tenant_id=dataset.tenant_id,
  420. document_form=dataset_document.doc_form,
  421. document_language=dataset_document.doc_language
  422. )
  423. # save node to document segment
  424. doc_store = DatesetDocumentStore(
  425. dataset=dataset,
  426. user_id=dataset_document.created_by,
  427. document_id=dataset_document.id
  428. )
  429. # add document segments
  430. doc_store.add_documents(documents)
  431. # update document status to indexing
  432. cur_time = datetime.datetime.utcnow()
  433. self._update_document_index_status(
  434. document_id=dataset_document.id,
  435. after_indexing_status="indexing",
  436. extra_update_params={
  437. DatasetDocument.cleaning_completed_at: cur_time,
  438. DatasetDocument.splitting_completed_at: cur_time,
  439. }
  440. )
  441. # update segment status to indexing
  442. self._update_segments_by_document(
  443. dataset_document_id=dataset_document.id,
  444. update_params={
  445. DocumentSegment.status: "indexing",
  446. DocumentSegment.indexing_at: datetime.datetime.utcnow()
  447. }
  448. )
  449. return documents
  450. def _split_to_documents(self, text_docs: List[Document], splitter: TextSplitter,
  451. processing_rule: DatasetProcessRule, tenant_id: str,
  452. document_form: str, document_language: str) -> List[Document]:
  453. """
  454. Split the text documents into nodes.
  455. """
  456. all_documents = []
  457. all_qa_documents = []
  458. for text_doc in text_docs:
  459. # document clean
  460. document_text = self._document_clean(text_doc.page_content, processing_rule)
  461. text_doc.page_content = document_text
  462. # parse document to nodes
  463. documents = splitter.split_documents([text_doc])
  464. split_documents = []
  465. for document_node in documents:
  466. if document_node.page_content.strip():
  467. doc_id = str(uuid.uuid4())
  468. hash = helper.generate_text_hash(document_node.page_content)
  469. document_node.metadata['doc_id'] = doc_id
  470. document_node.metadata['doc_hash'] = hash
  471. split_documents.append(document_node)
  472. all_documents.extend(split_documents)
  473. # processing qa document
  474. if document_form == 'qa_model':
  475. for i in range(0, len(all_documents), 10):
  476. threads = []
  477. sub_documents = all_documents[i:i + 10]
  478. for doc in sub_documents:
  479. document_format_thread = threading.Thread(target=self.format_qa_document, kwargs={
  480. 'flask_app': current_app._get_current_object(),
  481. 'tenant_id': tenant_id, 'document_node': doc, 'all_qa_documents': all_qa_documents,
  482. 'document_language': document_language})
  483. threads.append(document_format_thread)
  484. document_format_thread.start()
  485. for thread in threads:
  486. thread.join()
  487. return all_qa_documents
  488. return all_documents
  489. def format_qa_document(self, flask_app: Flask, tenant_id: str, document_node, all_qa_documents, document_language):
  490. format_documents = []
  491. if document_node.page_content is None or not document_node.page_content.strip():
  492. return
  493. with flask_app.app_context():
  494. try:
  495. # qa model document
  496. response = LLMGenerator.generate_qa_document(tenant_id, document_node.page_content, document_language)
  497. document_qa_list = self.format_split_text(response)
  498. qa_documents = []
  499. for result in document_qa_list:
  500. qa_document = Document(page_content=result['question'], metadata=document_node.metadata.copy())
  501. doc_id = str(uuid.uuid4())
  502. hash = helper.generate_text_hash(result['question'])
  503. qa_document.metadata['answer'] = result['answer']
  504. qa_document.metadata['doc_id'] = doc_id
  505. qa_document.metadata['doc_hash'] = hash
  506. qa_documents.append(qa_document)
  507. format_documents.extend(qa_documents)
  508. except Exception as e:
  509. logging.exception(e)
  510. all_qa_documents.extend(format_documents)
  511. def _split_to_documents_for_estimate(self, text_docs: List[Document], splitter: TextSplitter,
  512. processing_rule: DatasetProcessRule) -> List[Document]:
  513. """
  514. Split the text documents into nodes.
  515. """
  516. all_documents = []
  517. for text_doc in text_docs:
  518. # document clean
  519. document_text = self._document_clean(text_doc.page_content, processing_rule)
  520. text_doc.page_content = document_text
  521. # parse document to nodes
  522. documents = splitter.split_documents([text_doc])
  523. split_documents = []
  524. for document in documents:
  525. if document.page_content is None or not document.page_content.strip():
  526. continue
  527. doc_id = str(uuid.uuid4())
  528. hash = helper.generate_text_hash(document.page_content)
  529. document.metadata['doc_id'] = doc_id
  530. document.metadata['doc_hash'] = hash
  531. split_documents.append(document)
  532. all_documents.extend(split_documents)
  533. return all_documents
  534. def _document_clean(self, text: str, processing_rule: DatasetProcessRule) -> str:
  535. """
  536. Clean the document text according to the processing rules.
  537. """
  538. if processing_rule.mode == "automatic":
  539. rules = DatasetProcessRule.AUTOMATIC_RULES
  540. else:
  541. rules = json.loads(processing_rule.rules) if processing_rule.rules else {}
  542. if 'pre_processing_rules' in rules:
  543. pre_processing_rules = rules["pre_processing_rules"]
  544. for pre_processing_rule in pre_processing_rules:
  545. if pre_processing_rule["id"] == "remove_extra_spaces" and pre_processing_rule["enabled"] is True:
  546. # Remove extra spaces
  547. pattern = r'\n{3,}'
  548. text = re.sub(pattern, '\n\n', text)
  549. pattern = r'[\t\f\r\x20\u00a0\u1680\u180e\u2000-\u200a\u202f\u205f\u3000]{2,}'
  550. text = re.sub(pattern, ' ', text)
  551. elif pre_processing_rule["id"] == "remove_urls_emails" and pre_processing_rule["enabled"] is True:
  552. # Remove email
  553. pattern = r'([a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+)'
  554. text = re.sub(pattern, '', text)
  555. # Remove URL
  556. pattern = r'https?://[^\s]+'
  557. text = re.sub(pattern, '', text)
  558. return text
  559. def format_split_text(self, text):
  560. regex = r"Q\d+:\s*(.*?)\s*A\d+:\s*([\s\S]*?)(?=Q|$)"
  561. matches = re.findall(regex, text, re.MULTILINE)
  562. return [
  563. {
  564. "question": q,
  565. "answer": re.sub(r"\n\s*", "\n", a.strip())
  566. }
  567. for q, a in matches if q and a
  568. ]
  569. def _build_index(self, dataset: Dataset, dataset_document: DatasetDocument, documents: List[Document]) -> None:
  570. """
  571. Build the index for the document.
  572. """
  573. vector_index = IndexBuilder.get_index(dataset, 'high_quality')
  574. keyword_table_index = IndexBuilder.get_index(dataset, 'economy')
  575. embedding_model = None
  576. if dataset.indexing_technique == 'high_quality':
  577. embedding_model = ModelFactory.get_embedding_model(
  578. tenant_id=dataset.tenant_id,
  579. model_provider_name=dataset.embedding_model_provider,
  580. model_name=dataset.embedding_model
  581. )
  582. # chunk nodes by chunk size
  583. indexing_start_at = time.perf_counter()
  584. tokens = 0
  585. chunk_size = 100
  586. for i in range(0, len(documents), chunk_size):
  587. # check document is paused
  588. self._check_document_paused_status(dataset_document.id)
  589. chunk_documents = documents[i:i + chunk_size]
  590. if dataset.indexing_technique == 'high_quality' or embedding_model:
  591. tokens += sum(
  592. embedding_model.get_num_tokens(document.page_content)
  593. for document in chunk_documents
  594. )
  595. # save vector index
  596. if vector_index:
  597. vector_index.add_texts(chunk_documents)
  598. # save keyword index
  599. keyword_table_index.add_texts(chunk_documents)
  600. document_ids = [document.metadata['doc_id'] for document in chunk_documents]
  601. db.session.query(DocumentSegment).filter(
  602. DocumentSegment.document_id == dataset_document.id,
  603. DocumentSegment.index_node_id.in_(document_ids),
  604. DocumentSegment.status == "indexing"
  605. ).update({
  606. DocumentSegment.status: "completed",
  607. DocumentSegment.enabled: True,
  608. DocumentSegment.completed_at: datetime.datetime.utcnow()
  609. })
  610. db.session.commit()
  611. indexing_end_at = time.perf_counter()
  612. # update document status to completed
  613. self._update_document_index_status(
  614. document_id=dataset_document.id,
  615. after_indexing_status="completed",
  616. extra_update_params={
  617. DatasetDocument.tokens: tokens,
  618. DatasetDocument.completed_at: datetime.datetime.utcnow(),
  619. DatasetDocument.indexing_latency: indexing_end_at - indexing_start_at,
  620. }
  621. )
  622. def _check_document_paused_status(self, document_id: str):
  623. indexing_cache_key = 'document_{}_is_paused'.format(document_id)
  624. result = redis_client.get(indexing_cache_key)
  625. if result:
  626. raise DocumentIsPausedException()
  627. def _update_document_index_status(self, document_id: str, after_indexing_status: str,
  628. extra_update_params: Optional[dict] = None) -> None:
  629. """
  630. Update the document indexing status.
  631. """
  632. count = DatasetDocument.query.filter_by(id=document_id, is_paused=True).count()
  633. if count > 0:
  634. raise DocumentIsPausedException()
  635. document = DatasetDocument.query.filter_by(id=document_id).first()
  636. if not document:
  637. raise DocumentIsDeletedPausedException()
  638. update_params = {
  639. DatasetDocument.indexing_status: after_indexing_status
  640. }
  641. if extra_update_params:
  642. update_params.update(extra_update_params)
  643. DatasetDocument.query.filter_by(id=document_id).update(update_params)
  644. db.session.commit()
  645. def _update_segments_by_document(self, dataset_document_id: str, update_params: dict) -> None:
  646. """
  647. Update the document segment by document id.
  648. """
  649. DocumentSegment.query.filter_by(document_id=dataset_document_id).update(update_params)
  650. db.session.commit()
  651. def batch_add_segments(self, segments: List[DocumentSegment], dataset: Dataset):
  652. """
  653. Batch add segments index processing
  654. """
  655. documents = []
  656. for segment in segments:
  657. document = Document(
  658. page_content=segment.content,
  659. metadata={
  660. "doc_id": segment.index_node_id,
  661. "doc_hash": segment.index_node_hash,
  662. "document_id": segment.document_id,
  663. "dataset_id": segment.dataset_id,
  664. }
  665. )
  666. documents.append(document)
  667. # save vector index
  668. index = IndexBuilder.get_index(dataset, 'high_quality')
  669. if index:
  670. index.add_texts(documents, duplicate_check=True)
  671. # save keyword index
  672. index = IndexBuilder.get_index(dataset, 'economy')
  673. if index:
  674. index.add_texts(documents)
  675. class DocumentIsPausedException(Exception):
  676. pass
  677. class DocumentIsDeletedPausedException(Exception):
  678. pass