dataset_service.py 52 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173
  1. import json
  2. import logging
  3. import datetime
  4. import time
  5. import random
  6. import uuid
  7. from typing import Optional, List
  8. from flask import current_app
  9. from sqlalchemy import func
  10. from core.index.index import IndexBuilder
  11. from core.model_providers.error import LLMBadRequestError, ProviderTokenNotInitError
  12. from core.model_providers.model_factory import ModelFactory
  13. from extensions.ext_redis import redis_client
  14. from flask_login import current_user
  15. from events.dataset_event import dataset_was_deleted
  16. from events.document_event import document_was_deleted
  17. from extensions.ext_database import db
  18. from libs import helper
  19. from models.account import Account
  20. from models.dataset import Dataset, Document, DatasetQuery, DatasetProcessRule, AppDatasetJoin, DocumentSegment, \
  21. DatasetCollectionBinding
  22. from models.model import UploadFile
  23. from models.source import DataSourceBinding
  24. from services.errors.account import NoPermissionError
  25. from services.errors.dataset import DatasetNameDuplicateError
  26. from services.errors.document import DocumentIndexingError
  27. from services.errors.file import FileNotExistsError
  28. from services.vector_service import VectorService
  29. from tasks.clean_notion_document_task import clean_notion_document_task
  30. from tasks.deal_dataset_vector_index_task import deal_dataset_vector_index_task
  31. from tasks.document_indexing_task import document_indexing_task
  32. from tasks.document_indexing_update_task import document_indexing_update_task
  33. from tasks.create_segment_to_index_task import create_segment_to_index_task
  34. from tasks.update_segment_index_task import update_segment_index_task
  35. from tasks.recover_document_indexing_task import recover_document_indexing_task
  36. from tasks.update_segment_keyword_index_task import update_segment_keyword_index_task
  37. from tasks.delete_segment_from_index_task import delete_segment_from_index_task
  38. class DatasetService:
  39. @staticmethod
  40. def get_datasets(page, per_page, provider="vendor", tenant_id=None, user=None):
  41. if user:
  42. permission_filter = db.or_(Dataset.created_by == user.id,
  43. Dataset.permission == 'all_team_members')
  44. else:
  45. permission_filter = Dataset.permission == 'all_team_members'
  46. datasets = Dataset.query.filter(
  47. db.and_(Dataset.provider == provider, Dataset.tenant_id == tenant_id, permission_filter)) \
  48. .order_by(Dataset.created_at.desc()) \
  49. .paginate(
  50. page=page,
  51. per_page=per_page,
  52. max_per_page=100,
  53. error_out=False
  54. )
  55. return datasets.items, datasets.total
  56. @staticmethod
  57. def get_process_rules(dataset_id):
  58. # get the latest process rule
  59. dataset_process_rule = db.session.query(DatasetProcessRule). \
  60. filter(DatasetProcessRule.dataset_id == dataset_id). \
  61. order_by(DatasetProcessRule.created_at.desc()). \
  62. limit(1). \
  63. one_or_none()
  64. if dataset_process_rule:
  65. mode = dataset_process_rule.mode
  66. rules = dataset_process_rule.rules_dict
  67. else:
  68. mode = DocumentService.DEFAULT_RULES['mode']
  69. rules = DocumentService.DEFAULT_RULES['rules']
  70. return {
  71. 'mode': mode,
  72. 'rules': rules
  73. }
  74. @staticmethod
  75. def get_datasets_by_ids(ids, tenant_id):
  76. datasets = Dataset.query.filter(Dataset.id.in_(ids),
  77. Dataset.tenant_id == tenant_id).paginate(
  78. page=1, per_page=len(ids), max_per_page=len(ids), error_out=False)
  79. return datasets.items, datasets.total
  80. @staticmethod
  81. def create_empty_dataset(tenant_id: str, name: str, indexing_technique: Optional[str], account: Account):
  82. # check if dataset name already exists
  83. if Dataset.query.filter_by(name=name, tenant_id=tenant_id).first():
  84. raise DatasetNameDuplicateError(
  85. f'Dataset with name {name} already exists.')
  86. embedding_model = None
  87. if indexing_technique == 'high_quality':
  88. embedding_model = ModelFactory.get_embedding_model(
  89. tenant_id=tenant_id
  90. )
  91. dataset = Dataset(name=name, indexing_technique=indexing_technique)
  92. # dataset = Dataset(name=name, provider=provider, config=config)
  93. dataset.created_by = account.id
  94. dataset.updated_by = account.id
  95. dataset.tenant_id = tenant_id
  96. dataset.embedding_model_provider = embedding_model.model_provider.provider_name if embedding_model else None
  97. dataset.embedding_model = embedding_model.name if embedding_model else None
  98. db.session.add(dataset)
  99. db.session.commit()
  100. return dataset
  101. @staticmethod
  102. def get_dataset(dataset_id):
  103. dataset = Dataset.query.filter_by(
  104. id=dataset_id
  105. ).first()
  106. if dataset is None:
  107. return None
  108. else:
  109. return dataset
  110. @staticmethod
  111. def check_dataset_model_setting(dataset):
  112. if dataset.indexing_technique == 'high_quality':
  113. try:
  114. ModelFactory.get_embedding_model(
  115. tenant_id=dataset.tenant_id,
  116. model_provider_name=dataset.embedding_model_provider,
  117. model_name=dataset.embedding_model
  118. )
  119. except LLMBadRequestError:
  120. raise ValueError(
  121. f"No Embedding Model available. Please configure a valid provider "
  122. f"in the Settings -> Model Provider.")
  123. except ProviderTokenNotInitError as ex:
  124. raise ValueError(f"The dataset in unavailable, due to: "
  125. f"{ex.description}")
  126. @staticmethod
  127. def update_dataset(dataset_id, data, user):
  128. filtered_data = {k: v for k, v in data.items() if v is not None or k == 'description'}
  129. dataset = DatasetService.get_dataset(dataset_id)
  130. DatasetService.check_dataset_permission(dataset, user)
  131. action = None
  132. if dataset.indexing_technique != data['indexing_technique']:
  133. # if update indexing_technique
  134. if data['indexing_technique'] == 'economy':
  135. action = 'remove'
  136. filtered_data['embedding_model'] = None
  137. filtered_data['embedding_model_provider'] = None
  138. filtered_data['collection_binding_id'] = None
  139. elif data['indexing_technique'] == 'high_quality':
  140. action = 'add'
  141. # get embedding model setting
  142. try:
  143. embedding_model = ModelFactory.get_embedding_model(
  144. tenant_id=current_user.current_tenant_id
  145. )
  146. filtered_data['embedding_model'] = embedding_model.name
  147. filtered_data['embedding_model_provider'] = embedding_model.model_provider.provider_name
  148. dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
  149. embedding_model.model_provider.provider_name,
  150. embedding_model.name
  151. )
  152. filtered_data['collection_binding_id'] = dataset_collection_binding.id
  153. except LLMBadRequestError:
  154. raise ValueError(
  155. f"No Embedding Model available. Please configure a valid provider "
  156. f"in the Settings -> Model Provider.")
  157. except ProviderTokenNotInitError as ex:
  158. raise ValueError(ex.description)
  159. filtered_data['updated_by'] = user.id
  160. filtered_data['updated_at'] = datetime.datetime.now()
  161. dataset.query.filter_by(id=dataset_id).update(filtered_data)
  162. db.session.commit()
  163. if action:
  164. deal_dataset_vector_index_task.delay(dataset_id, action)
  165. return dataset
  166. @staticmethod
  167. def delete_dataset(dataset_id, user):
  168. # todo: cannot delete dataset if it is being processed
  169. dataset = DatasetService.get_dataset(dataset_id)
  170. if dataset is None:
  171. return False
  172. DatasetService.check_dataset_permission(dataset, user)
  173. dataset_was_deleted.send(dataset)
  174. db.session.delete(dataset)
  175. db.session.commit()
  176. return True
  177. @staticmethod
  178. def check_dataset_permission(dataset, user):
  179. if dataset.tenant_id != user.current_tenant_id:
  180. logging.debug(
  181. f'User {user.id} does not have permission to access dataset {dataset.id}')
  182. raise NoPermissionError(
  183. 'You do not have permission to access this dataset.')
  184. if dataset.permission == 'only_me' and dataset.created_by != user.id:
  185. logging.debug(
  186. f'User {user.id} does not have permission to access dataset {dataset.id}')
  187. raise NoPermissionError(
  188. 'You do not have permission to access this dataset.')
  189. @staticmethod
  190. def get_dataset_queries(dataset_id: str, page: int, per_page: int):
  191. dataset_queries = DatasetQuery.query.filter_by(dataset_id=dataset_id) \
  192. .order_by(db.desc(DatasetQuery.created_at)) \
  193. .paginate(
  194. page=page, per_page=per_page, max_per_page=100, error_out=False
  195. )
  196. return dataset_queries.items, dataset_queries.total
  197. @staticmethod
  198. def get_related_apps(dataset_id: str):
  199. return AppDatasetJoin.query.filter(AppDatasetJoin.dataset_id == dataset_id) \
  200. .order_by(db.desc(AppDatasetJoin.created_at)).all()
  201. class DocumentService:
  202. DEFAULT_RULES = {
  203. 'mode': 'custom',
  204. 'rules': {
  205. 'pre_processing_rules': [
  206. {'id': 'remove_extra_spaces', 'enabled': True},
  207. {'id': 'remove_urls_emails', 'enabled': False}
  208. ],
  209. 'segmentation': {
  210. 'delimiter': '\n',
  211. 'max_tokens': 500
  212. }
  213. }
  214. }
  215. DOCUMENT_METADATA_SCHEMA = {
  216. "book": {
  217. "title": str,
  218. "language": str,
  219. "author": str,
  220. "publisher": str,
  221. "publication_date": str,
  222. "isbn": str,
  223. "category": str,
  224. },
  225. "web_page": {
  226. "title": str,
  227. "url": str,
  228. "language": str,
  229. "publish_date": str,
  230. "author/publisher": str,
  231. "topic/keywords": str,
  232. "description": str,
  233. },
  234. "paper": {
  235. "title": str,
  236. "language": str,
  237. "author": str,
  238. "publish_date": str,
  239. "journal/conference_name": str,
  240. "volume/issue/page_numbers": str,
  241. "doi": str,
  242. "topic/keywords": str,
  243. "abstract": str,
  244. },
  245. "social_media_post": {
  246. "platform": str,
  247. "author/username": str,
  248. "publish_date": str,
  249. "post_url": str,
  250. "topic/tags": str,
  251. },
  252. "wikipedia_entry": {
  253. "title": str,
  254. "language": str,
  255. "web_page_url": str,
  256. "last_edit_date": str,
  257. "editor/contributor": str,
  258. "summary/introduction": str,
  259. },
  260. "personal_document": {
  261. "title": str,
  262. "author": str,
  263. "creation_date": str,
  264. "last_modified_date": str,
  265. "document_type": str,
  266. "tags/category": str,
  267. },
  268. "business_document": {
  269. "title": str,
  270. "author": str,
  271. "creation_date": str,
  272. "last_modified_date": str,
  273. "document_type": str,
  274. "department/team": str,
  275. },
  276. "im_chat_log": {
  277. "chat_platform": str,
  278. "chat_participants/group_name": str,
  279. "start_date": str,
  280. "end_date": str,
  281. "summary": str,
  282. },
  283. "synced_from_notion": {
  284. "title": str,
  285. "language": str,
  286. "author/creator": str,
  287. "creation_date": str,
  288. "last_modified_date": str,
  289. "notion_page_link": str,
  290. "category/tags": str,
  291. "description": str,
  292. },
  293. "synced_from_github": {
  294. "repository_name": str,
  295. "repository_description": str,
  296. "repository_owner/organization": str,
  297. "code_filename": str,
  298. "code_file_path": str,
  299. "programming_language": str,
  300. "github_link": str,
  301. "open_source_license": str,
  302. "commit_date": str,
  303. "commit_author": str,
  304. },
  305. "others": dict
  306. }
  307. @staticmethod
  308. def get_document(dataset_id: str, document_id: str) -> Optional[Document]:
  309. document = db.session.query(Document).filter(
  310. Document.id == document_id,
  311. Document.dataset_id == dataset_id
  312. ).first()
  313. return document
  314. @staticmethod
  315. def get_document_by_id(document_id: str) -> Optional[Document]:
  316. document = db.session.query(Document).filter(
  317. Document.id == document_id
  318. ).first()
  319. return document
  320. @staticmethod
  321. def get_document_by_dataset_id(dataset_id: str) -> List[Document]:
  322. documents = db.session.query(Document).filter(
  323. Document.dataset_id == dataset_id,
  324. Document.enabled == True
  325. ).all()
  326. return documents
  327. @staticmethod
  328. def get_batch_documents(dataset_id: str, batch: str) -> List[Document]:
  329. documents = db.session.query(Document).filter(
  330. Document.batch == batch,
  331. Document.dataset_id == dataset_id,
  332. Document.tenant_id == current_user.current_tenant_id
  333. ).all()
  334. return documents
  335. @staticmethod
  336. def get_document_file_detail(file_id: str):
  337. file_detail = db.session.query(UploadFile). \
  338. filter(UploadFile.id == file_id). \
  339. one_or_none()
  340. return file_detail
  341. @staticmethod
  342. def check_archived(document):
  343. if document.archived:
  344. return True
  345. else:
  346. return False
  347. @staticmethod
  348. def delete_document(document):
  349. # trigger document_was_deleted signal
  350. document_was_deleted.send(document.id, dataset_id=document.dataset_id)
  351. db.session.delete(document)
  352. db.session.commit()
  353. @staticmethod
  354. def pause_document(document):
  355. if document.indexing_status not in ["waiting", "parsing", "cleaning", "splitting", "indexing"]:
  356. raise DocumentIndexingError()
  357. # update document to be paused
  358. document.is_paused = True
  359. document.paused_by = current_user.id
  360. document.paused_at = datetime.datetime.utcnow()
  361. db.session.add(document)
  362. db.session.commit()
  363. # set document paused flag
  364. indexing_cache_key = 'document_{}_is_paused'.format(document.id)
  365. redis_client.setnx(indexing_cache_key, "True")
  366. @staticmethod
  367. def recover_document(document):
  368. if not document.is_paused:
  369. raise DocumentIndexingError()
  370. # update document to be recover
  371. document.is_paused = False
  372. document.paused_by = None
  373. document.paused_at = None
  374. db.session.add(document)
  375. db.session.commit()
  376. # delete paused flag
  377. indexing_cache_key = 'document_{}_is_paused'.format(document.id)
  378. redis_client.delete(indexing_cache_key)
  379. # trigger async task
  380. recover_document_indexing_task.delay(document.dataset_id, document.id)
  381. @staticmethod
  382. def get_documents_position(dataset_id):
  383. document = Document.query.filter_by(dataset_id=dataset_id).order_by(Document.position.desc()).first()
  384. if document:
  385. return document.position + 1
  386. else:
  387. return 1
  388. @staticmethod
  389. def save_document_with_dataset_id(dataset: Dataset, document_data: dict,
  390. account: Account, dataset_process_rule: Optional[DatasetProcessRule] = None,
  391. created_from: str = 'web'):
  392. # check document limit
  393. if current_app.config['EDITION'] == 'CLOUD':
  394. if 'original_document_id' not in document_data or not document_data['original_document_id']:
  395. count = 0
  396. if document_data["data_source"]["type"] == "upload_file":
  397. upload_file_list = document_data["data_source"]["info_list"]['file_info_list']['file_ids']
  398. count = len(upload_file_list)
  399. elif document_data["data_source"]["type"] == "notion_import":
  400. notion_info_list = document_data["data_source"]['info_list']['notion_info_list']
  401. for notion_info in notion_info_list:
  402. count = count + len(notion_info['pages'])
  403. documents_count = DocumentService.get_tenant_documents_count()
  404. total_count = documents_count + count
  405. tenant_document_count = int(current_app.config['TENANT_DOCUMENT_COUNT'])
  406. if total_count > tenant_document_count:
  407. raise ValueError(f"over document limit {tenant_document_count}.")
  408. # if dataset is empty, update dataset data_source_type
  409. if not dataset.data_source_type:
  410. dataset.data_source_type = document_data["data_source"]["type"]
  411. if not dataset.indexing_technique:
  412. if 'indexing_technique' not in document_data \
  413. or document_data['indexing_technique'] not in Dataset.INDEXING_TECHNIQUE_LIST:
  414. raise ValueError("Indexing technique is required")
  415. dataset.indexing_technique = document_data["indexing_technique"]
  416. if document_data["indexing_technique"] == 'high_quality':
  417. embedding_model = ModelFactory.get_embedding_model(
  418. tenant_id=dataset.tenant_id
  419. )
  420. dataset.embedding_model = embedding_model.name
  421. dataset.embedding_model_provider = embedding_model.model_provider.provider_name
  422. dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
  423. embedding_model.model_provider.provider_name,
  424. embedding_model.name
  425. )
  426. dataset.collection_binding_id = dataset_collection_binding.id
  427. documents = []
  428. batch = time.strftime('%Y%m%d%H%M%S') + str(random.randint(100000, 999999))
  429. if 'original_document_id' in document_data and document_data["original_document_id"]:
  430. document = DocumentService.update_document_with_dataset_id(dataset, document_data, account)
  431. documents.append(document)
  432. else:
  433. # save process rule
  434. if not dataset_process_rule:
  435. process_rule = document_data["process_rule"]
  436. if process_rule["mode"] == "custom":
  437. dataset_process_rule = DatasetProcessRule(
  438. dataset_id=dataset.id,
  439. mode=process_rule["mode"],
  440. rules=json.dumps(process_rule["rules"]),
  441. created_by=account.id
  442. )
  443. elif process_rule["mode"] == "automatic":
  444. dataset_process_rule = DatasetProcessRule(
  445. dataset_id=dataset.id,
  446. mode=process_rule["mode"],
  447. rules=json.dumps(DatasetProcessRule.AUTOMATIC_RULES),
  448. created_by=account.id
  449. )
  450. db.session.add(dataset_process_rule)
  451. db.session.commit()
  452. position = DocumentService.get_documents_position(dataset.id)
  453. document_ids = []
  454. if document_data["data_source"]["type"] == "upload_file":
  455. upload_file_list = document_data["data_source"]["info_list"]['file_info_list']['file_ids']
  456. for file_id in upload_file_list:
  457. file = db.session.query(UploadFile).filter(
  458. UploadFile.tenant_id == dataset.tenant_id,
  459. UploadFile.id == file_id
  460. ).first()
  461. # raise error if file not found
  462. if not file:
  463. raise FileNotExistsError()
  464. file_name = file.name
  465. data_source_info = {
  466. "upload_file_id": file_id,
  467. }
  468. document = DocumentService.build_document(dataset, dataset_process_rule.id,
  469. document_data["data_source"]["type"],
  470. document_data["doc_form"],
  471. document_data["doc_language"],
  472. data_source_info, created_from, position,
  473. account, file_name, batch)
  474. db.session.add(document)
  475. db.session.flush()
  476. document_ids.append(document.id)
  477. documents.append(document)
  478. position += 1
  479. elif document_data["data_source"]["type"] == "notion_import":
  480. notion_info_list = document_data["data_source"]['info_list']['notion_info_list']
  481. exist_page_ids = []
  482. exist_document = dict()
  483. documents = Document.query.filter_by(
  484. dataset_id=dataset.id,
  485. tenant_id=current_user.current_tenant_id,
  486. data_source_type='notion_import',
  487. enabled=True
  488. ).all()
  489. if documents:
  490. for document in documents:
  491. data_source_info = json.loads(document.data_source_info)
  492. exist_page_ids.append(data_source_info['notion_page_id'])
  493. exist_document[data_source_info['notion_page_id']] = document.id
  494. for notion_info in notion_info_list:
  495. workspace_id = notion_info['workspace_id']
  496. data_source_binding = DataSourceBinding.query.filter(
  497. db.and_(
  498. DataSourceBinding.tenant_id == current_user.current_tenant_id,
  499. DataSourceBinding.provider == 'notion',
  500. DataSourceBinding.disabled == False,
  501. DataSourceBinding.source_info['workspace_id'] == f'"{workspace_id}"'
  502. )
  503. ).first()
  504. if not data_source_binding:
  505. raise ValueError('Data source binding not found.')
  506. for page in notion_info['pages']:
  507. if page['page_id'] not in exist_page_ids:
  508. data_source_info = {
  509. "notion_workspace_id": workspace_id,
  510. "notion_page_id": page['page_id'],
  511. "notion_page_icon": page['page_icon'],
  512. "type": page['type']
  513. }
  514. document = DocumentService.build_document(dataset, dataset_process_rule.id,
  515. document_data["data_source"]["type"],
  516. document_data["doc_form"],
  517. document_data["doc_language"],
  518. data_source_info, created_from, position,
  519. account, page['page_name'], batch)
  520. db.session.add(document)
  521. db.session.flush()
  522. document_ids.append(document.id)
  523. documents.append(document)
  524. position += 1
  525. else:
  526. exist_document.pop(page['page_id'])
  527. # delete not selected documents
  528. if len(exist_document) > 0:
  529. clean_notion_document_task.delay(list(exist_document.values()), dataset.id)
  530. db.session.commit()
  531. # trigger async task
  532. document_indexing_task.delay(dataset.id, document_ids)
  533. return documents, batch
  534. @staticmethod
  535. def build_document(dataset: Dataset, process_rule_id: str, data_source_type: str, document_form: str,
  536. document_language: str, data_source_info: dict, created_from: str, position: int,
  537. account: Account,
  538. name: str, batch: str):
  539. document = Document(
  540. tenant_id=dataset.tenant_id,
  541. dataset_id=dataset.id,
  542. position=position,
  543. data_source_type=data_source_type,
  544. data_source_info=json.dumps(data_source_info),
  545. dataset_process_rule_id=process_rule_id,
  546. batch=batch,
  547. name=name,
  548. created_from=created_from,
  549. created_by=account.id,
  550. doc_form=document_form,
  551. doc_language=document_language
  552. )
  553. return document
  554. @staticmethod
  555. def get_tenant_documents_count():
  556. documents_count = Document.query.filter(Document.completed_at.isnot(None),
  557. Document.enabled == True,
  558. Document.archived == False,
  559. Document.tenant_id == current_user.current_tenant_id).count()
  560. return documents_count
  561. @staticmethod
  562. def update_document_with_dataset_id(dataset: Dataset, document_data: dict,
  563. account: Account, dataset_process_rule: Optional[DatasetProcessRule] = None,
  564. created_from: str = 'web'):
  565. DatasetService.check_dataset_model_setting(dataset)
  566. document = DocumentService.get_document(dataset.id, document_data["original_document_id"])
  567. if document.display_status != 'available':
  568. raise ValueError("Document is not available")
  569. # update document name
  570. if 'name' in document_data and document_data['name']:
  571. document.name = document_data['name']
  572. # save process rule
  573. if 'process_rule' in document_data and document_data['process_rule']:
  574. process_rule = document_data["process_rule"]
  575. if process_rule["mode"] == "custom":
  576. dataset_process_rule = DatasetProcessRule(
  577. dataset_id=dataset.id,
  578. mode=process_rule["mode"],
  579. rules=json.dumps(process_rule["rules"]),
  580. created_by=account.id
  581. )
  582. elif process_rule["mode"] == "automatic":
  583. dataset_process_rule = DatasetProcessRule(
  584. dataset_id=dataset.id,
  585. mode=process_rule["mode"],
  586. rules=json.dumps(DatasetProcessRule.AUTOMATIC_RULES),
  587. created_by=account.id
  588. )
  589. db.session.add(dataset_process_rule)
  590. db.session.commit()
  591. document.dataset_process_rule_id = dataset_process_rule.id
  592. # update document data source
  593. if 'data_source' in document_data and document_data['data_source']:
  594. file_name = ''
  595. data_source_info = {}
  596. if document_data["data_source"]["type"] == "upload_file":
  597. upload_file_list = document_data["data_source"]["info_list"]['file_info_list']['file_ids']
  598. for file_id in upload_file_list:
  599. file = db.session.query(UploadFile).filter(
  600. UploadFile.tenant_id == dataset.tenant_id,
  601. UploadFile.id == file_id
  602. ).first()
  603. # raise error if file not found
  604. if not file:
  605. raise FileNotExistsError()
  606. file_name = file.name
  607. data_source_info = {
  608. "upload_file_id": file_id,
  609. }
  610. elif document_data["data_source"]["type"] == "notion_import":
  611. notion_info_list = document_data["data_source"]['info_list']['notion_info_list']
  612. for notion_info in notion_info_list:
  613. workspace_id = notion_info['workspace_id']
  614. data_source_binding = DataSourceBinding.query.filter(
  615. db.and_(
  616. DataSourceBinding.tenant_id == current_user.current_tenant_id,
  617. DataSourceBinding.provider == 'notion',
  618. DataSourceBinding.disabled == False,
  619. DataSourceBinding.source_info['workspace_id'] == f'"{workspace_id}"'
  620. )
  621. ).first()
  622. if not data_source_binding:
  623. raise ValueError('Data source binding not found.')
  624. for page in notion_info['pages']:
  625. data_source_info = {
  626. "notion_workspace_id": workspace_id,
  627. "notion_page_id": page['page_id'],
  628. "notion_page_icon": page['page_icon'],
  629. "type": page['type']
  630. }
  631. document.data_source_type = document_data["data_source"]["type"]
  632. document.data_source_info = json.dumps(data_source_info)
  633. document.name = file_name
  634. # update document to be waiting
  635. document.indexing_status = 'waiting'
  636. document.completed_at = None
  637. document.processing_started_at = None
  638. document.parsing_completed_at = None
  639. document.cleaning_completed_at = None
  640. document.splitting_completed_at = None
  641. document.updated_at = datetime.datetime.utcnow()
  642. document.created_from = created_from
  643. document.doc_form = document_data['doc_form']
  644. db.session.add(document)
  645. db.session.commit()
  646. # update document segment
  647. update_params = {
  648. DocumentSegment.status: 're_segment'
  649. }
  650. DocumentSegment.query.filter_by(document_id=document.id).update(update_params)
  651. db.session.commit()
  652. # trigger async task
  653. document_indexing_update_task.delay(document.dataset_id, document.id)
  654. return document
  655. @staticmethod
  656. def save_document_without_dataset_id(tenant_id: str, document_data: dict, account: Account):
  657. count = 0
  658. if document_data["data_source"]["type"] == "upload_file":
  659. upload_file_list = document_data["data_source"]["info_list"]['file_info_list']['file_ids']
  660. count = len(upload_file_list)
  661. elif document_data["data_source"]["type"] == "notion_import":
  662. notion_info_list = document_data["data_source"]['info_list']['notion_info_list']
  663. for notion_info in notion_info_list:
  664. count = count + len(notion_info['pages'])
  665. # check document limit
  666. if current_app.config['EDITION'] == 'CLOUD':
  667. documents_count = DocumentService.get_tenant_documents_count()
  668. total_count = documents_count + count
  669. tenant_document_count = int(current_app.config['TENANT_DOCUMENT_COUNT'])
  670. if total_count > tenant_document_count:
  671. raise ValueError(f"All your documents have overed limit {tenant_document_count}.")
  672. embedding_model = None
  673. dataset_collection_binding_id = None
  674. if document_data['indexing_technique'] == 'high_quality':
  675. embedding_model = ModelFactory.get_embedding_model(
  676. tenant_id=tenant_id
  677. )
  678. dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
  679. embedding_model.model_provider.provider_name,
  680. embedding_model.name
  681. )
  682. dataset_collection_binding_id = dataset_collection_binding.id
  683. # save dataset
  684. dataset = Dataset(
  685. tenant_id=tenant_id,
  686. name='',
  687. data_source_type=document_data["data_source"]["type"],
  688. indexing_technique=document_data["indexing_technique"],
  689. created_by=account.id,
  690. embedding_model=embedding_model.name if embedding_model else None,
  691. embedding_model_provider=embedding_model.model_provider.provider_name if embedding_model else None,
  692. collection_binding_id=dataset_collection_binding_id
  693. )
  694. db.session.add(dataset)
  695. db.session.flush()
  696. documents, batch = DocumentService.save_document_with_dataset_id(dataset, document_data, account)
  697. cut_length = 18
  698. cut_name = documents[0].name[:cut_length]
  699. dataset.name = cut_name + '...'
  700. dataset.description = 'useful for when you want to answer queries about the ' + documents[0].name
  701. db.session.commit()
  702. return dataset, documents, batch
  703. @classmethod
  704. def document_create_args_validate(cls, args: dict):
  705. if 'original_document_id' not in args or not args['original_document_id']:
  706. DocumentService.data_source_args_validate(args)
  707. DocumentService.process_rule_args_validate(args)
  708. else:
  709. if ('data_source' not in args and not args['data_source']) \
  710. and ('process_rule' not in args and not args['process_rule']):
  711. raise ValueError("Data source or Process rule is required")
  712. else:
  713. if 'data_source' in args and args['data_source']:
  714. DocumentService.data_source_args_validate(args)
  715. if 'process_rule' in args and args['process_rule']:
  716. DocumentService.process_rule_args_validate(args)
  717. @classmethod
  718. def data_source_args_validate(cls, args: dict):
  719. if 'data_source' not in args or not args['data_source']:
  720. raise ValueError("Data source is required")
  721. if not isinstance(args['data_source'], dict):
  722. raise ValueError("Data source is invalid")
  723. if 'type' not in args['data_source'] or not args['data_source']['type']:
  724. raise ValueError("Data source type is required")
  725. if args['data_source']['type'] not in Document.DATA_SOURCES:
  726. raise ValueError("Data source type is invalid")
  727. if 'info_list' not in args['data_source'] or not args['data_source']['info_list']:
  728. raise ValueError("Data source info is required")
  729. if args['data_source']['type'] == 'upload_file':
  730. if 'file_info_list' not in args['data_source']['info_list'] or not args['data_source']['info_list'][
  731. 'file_info_list']:
  732. raise ValueError("File source info is required")
  733. if args['data_source']['type'] == 'notion_import':
  734. if 'notion_info_list' not in args['data_source']['info_list'] or not args['data_source']['info_list'][
  735. 'notion_info_list']:
  736. raise ValueError("Notion source info is required")
  737. @classmethod
  738. def process_rule_args_validate(cls, args: dict):
  739. if 'process_rule' not in args or not args['process_rule']:
  740. raise ValueError("Process rule is required")
  741. if not isinstance(args['process_rule'], dict):
  742. raise ValueError("Process rule is invalid")
  743. if 'mode' not in args['process_rule'] or not args['process_rule']['mode']:
  744. raise ValueError("Process rule mode is required")
  745. if args['process_rule']['mode'] not in DatasetProcessRule.MODES:
  746. raise ValueError("Process rule mode is invalid")
  747. if args['process_rule']['mode'] == 'automatic':
  748. args['process_rule']['rules'] = {}
  749. else:
  750. if 'rules' not in args['process_rule'] or not args['process_rule']['rules']:
  751. raise ValueError("Process rule rules is required")
  752. if not isinstance(args['process_rule']['rules'], dict):
  753. raise ValueError("Process rule rules is invalid")
  754. if 'pre_processing_rules' not in args['process_rule']['rules'] \
  755. or args['process_rule']['rules']['pre_processing_rules'] is None:
  756. raise ValueError("Process rule pre_processing_rules is required")
  757. if not isinstance(args['process_rule']['rules']['pre_processing_rules'], list):
  758. raise ValueError("Process rule pre_processing_rules is invalid")
  759. unique_pre_processing_rule_dicts = {}
  760. for pre_processing_rule in args['process_rule']['rules']['pre_processing_rules']:
  761. if 'id' not in pre_processing_rule or not pre_processing_rule['id']:
  762. raise ValueError("Process rule pre_processing_rules id is required")
  763. if pre_processing_rule['id'] not in DatasetProcessRule.PRE_PROCESSING_RULES:
  764. raise ValueError("Process rule pre_processing_rules id is invalid")
  765. if 'enabled' not in pre_processing_rule or pre_processing_rule['enabled'] is None:
  766. raise ValueError("Process rule pre_processing_rules enabled is required")
  767. if not isinstance(pre_processing_rule['enabled'], bool):
  768. raise ValueError("Process rule pre_processing_rules enabled is invalid")
  769. unique_pre_processing_rule_dicts[pre_processing_rule['id']] = pre_processing_rule
  770. args['process_rule']['rules']['pre_processing_rules'] = list(unique_pre_processing_rule_dicts.values())
  771. if 'segmentation' not in args['process_rule']['rules'] \
  772. or args['process_rule']['rules']['segmentation'] is None:
  773. raise ValueError("Process rule segmentation is required")
  774. if not isinstance(args['process_rule']['rules']['segmentation'], dict):
  775. raise ValueError("Process rule segmentation is invalid")
  776. if 'separator' not in args['process_rule']['rules']['segmentation'] \
  777. or not args['process_rule']['rules']['segmentation']['separator']:
  778. raise ValueError("Process rule segmentation separator is required")
  779. if not isinstance(args['process_rule']['rules']['segmentation']['separator'], str):
  780. raise ValueError("Process rule segmentation separator is invalid")
  781. if 'max_tokens' not in args['process_rule']['rules']['segmentation'] \
  782. or not args['process_rule']['rules']['segmentation']['max_tokens']:
  783. raise ValueError("Process rule segmentation max_tokens is required")
  784. if not isinstance(args['process_rule']['rules']['segmentation']['max_tokens'], int):
  785. raise ValueError("Process rule segmentation max_tokens is invalid")
  786. @classmethod
  787. def estimate_args_validate(cls, args: dict):
  788. if 'info_list' not in args or not args['info_list']:
  789. raise ValueError("Data source info is required")
  790. if not isinstance(args['info_list'], dict):
  791. raise ValueError("Data info is invalid")
  792. if 'process_rule' not in args or not args['process_rule']:
  793. raise ValueError("Process rule is required")
  794. if not isinstance(args['process_rule'], dict):
  795. raise ValueError("Process rule is invalid")
  796. if 'mode' not in args['process_rule'] or not args['process_rule']['mode']:
  797. raise ValueError("Process rule mode is required")
  798. if args['process_rule']['mode'] not in DatasetProcessRule.MODES:
  799. raise ValueError("Process rule mode is invalid")
  800. if args['process_rule']['mode'] == 'automatic':
  801. args['process_rule']['rules'] = {}
  802. else:
  803. if 'rules' not in args['process_rule'] or not args['process_rule']['rules']:
  804. raise ValueError("Process rule rules is required")
  805. if not isinstance(args['process_rule']['rules'], dict):
  806. raise ValueError("Process rule rules is invalid")
  807. if 'pre_processing_rules' not in args['process_rule']['rules'] \
  808. or args['process_rule']['rules']['pre_processing_rules'] is None:
  809. raise ValueError("Process rule pre_processing_rules is required")
  810. if not isinstance(args['process_rule']['rules']['pre_processing_rules'], list):
  811. raise ValueError("Process rule pre_processing_rules is invalid")
  812. unique_pre_processing_rule_dicts = {}
  813. for pre_processing_rule in args['process_rule']['rules']['pre_processing_rules']:
  814. if 'id' not in pre_processing_rule or not pre_processing_rule['id']:
  815. raise ValueError("Process rule pre_processing_rules id is required")
  816. if pre_processing_rule['id'] not in DatasetProcessRule.PRE_PROCESSING_RULES:
  817. raise ValueError("Process rule pre_processing_rules id is invalid")
  818. if 'enabled' not in pre_processing_rule or pre_processing_rule['enabled'] is None:
  819. raise ValueError("Process rule pre_processing_rules enabled is required")
  820. if not isinstance(pre_processing_rule['enabled'], bool):
  821. raise ValueError("Process rule pre_processing_rules enabled is invalid")
  822. unique_pre_processing_rule_dicts[pre_processing_rule['id']] = pre_processing_rule
  823. args['process_rule']['rules']['pre_processing_rules'] = list(unique_pre_processing_rule_dicts.values())
  824. if 'segmentation' not in args['process_rule']['rules'] \
  825. or args['process_rule']['rules']['segmentation'] is None:
  826. raise ValueError("Process rule segmentation is required")
  827. if not isinstance(args['process_rule']['rules']['segmentation'], dict):
  828. raise ValueError("Process rule segmentation is invalid")
  829. if 'separator' not in args['process_rule']['rules']['segmentation'] \
  830. or not args['process_rule']['rules']['segmentation']['separator']:
  831. raise ValueError("Process rule segmentation separator is required")
  832. if not isinstance(args['process_rule']['rules']['segmentation']['separator'], str):
  833. raise ValueError("Process rule segmentation separator is invalid")
  834. if 'max_tokens' not in args['process_rule']['rules']['segmentation'] \
  835. or not args['process_rule']['rules']['segmentation']['max_tokens']:
  836. raise ValueError("Process rule segmentation max_tokens is required")
  837. if not isinstance(args['process_rule']['rules']['segmentation']['max_tokens'], int):
  838. raise ValueError("Process rule segmentation max_tokens is invalid")
  839. class SegmentService:
  840. @classmethod
  841. def segment_create_args_validate(cls, args: dict, document: Document):
  842. if document.doc_form == 'qa_model':
  843. if 'answer' not in args or not args['answer']:
  844. raise ValueError("Answer is required")
  845. if not args['answer'].strip():
  846. raise ValueError("Answer is empty")
  847. if 'content' not in args or not args['content'] or not args['content'].strip():
  848. raise ValueError("Content is empty")
  849. @classmethod
  850. def create_segment(cls, args: dict, document: Document, dataset: Dataset):
  851. content = args['content']
  852. doc_id = str(uuid.uuid4())
  853. segment_hash = helper.generate_text_hash(content)
  854. tokens = 0
  855. if dataset.indexing_technique == 'high_quality':
  856. embedding_model = ModelFactory.get_embedding_model(
  857. tenant_id=dataset.tenant_id,
  858. model_provider_name=dataset.embedding_model_provider,
  859. model_name=dataset.embedding_model
  860. )
  861. # calc embedding use tokens
  862. tokens = embedding_model.get_num_tokens(content)
  863. max_position = db.session.query(func.max(DocumentSegment.position)).filter(
  864. DocumentSegment.document_id == document.id
  865. ).scalar()
  866. segment_document = DocumentSegment(
  867. tenant_id=current_user.current_tenant_id,
  868. dataset_id=document.dataset_id,
  869. document_id=document.id,
  870. index_node_id=doc_id,
  871. index_node_hash=segment_hash,
  872. position=max_position + 1 if max_position else 1,
  873. content=content,
  874. word_count=len(content),
  875. tokens=tokens,
  876. status='completed',
  877. indexing_at=datetime.datetime.utcnow(),
  878. completed_at=datetime.datetime.utcnow(),
  879. created_by=current_user.id
  880. )
  881. if document.doc_form == 'qa_model':
  882. segment_document.answer = args['answer']
  883. db.session.add(segment_document)
  884. db.session.commit()
  885. # save vector index
  886. try:
  887. VectorService.create_segment_vector(args['keywords'], segment_document, dataset)
  888. except Exception as e:
  889. logging.exception("create segment index failed")
  890. segment_document.enabled = False
  891. segment_document.disabled_at = datetime.datetime.utcnow()
  892. segment_document.status = 'error'
  893. segment_document.error = str(e)
  894. db.session.commit()
  895. segment = db.session.query(DocumentSegment).filter(DocumentSegment.id == segment_document.id).first()
  896. return segment
  897. @classmethod
  898. def multi_create_segment(cls, segments: list, document: Document, dataset: Dataset):
  899. embedding_model = None
  900. if dataset.indexing_technique == 'high_quality':
  901. embedding_model = ModelFactory.get_embedding_model(
  902. tenant_id=dataset.tenant_id,
  903. model_provider_name=dataset.embedding_model_provider,
  904. model_name=dataset.embedding_model
  905. )
  906. max_position = db.session.query(func.max(DocumentSegment.position)).filter(
  907. DocumentSegment.document_id == document.id
  908. ).scalar()
  909. pre_segment_data_list = []
  910. segment_data_list = []
  911. for segment_item in segments:
  912. content = segment_item['content']
  913. doc_id = str(uuid.uuid4())
  914. segment_hash = helper.generate_text_hash(content)
  915. tokens = 0
  916. if dataset.indexing_technique == 'high_quality' and embedding_model:
  917. # calc embedding use tokens
  918. tokens = embedding_model.get_num_tokens(content)
  919. segment_document = DocumentSegment(
  920. tenant_id=current_user.current_tenant_id,
  921. dataset_id=document.dataset_id,
  922. document_id=document.id,
  923. index_node_id=doc_id,
  924. index_node_hash=segment_hash,
  925. position=max_position + 1 if max_position else 1,
  926. content=content,
  927. word_count=len(content),
  928. tokens=tokens,
  929. status='completed',
  930. indexing_at=datetime.datetime.utcnow(),
  931. completed_at=datetime.datetime.utcnow(),
  932. created_by=current_user.id
  933. )
  934. if document.doc_form == 'qa_model':
  935. segment_document.answer = segment_item['answer']
  936. db.session.add(segment_document)
  937. segment_data_list.append(segment_document)
  938. pre_segment_data = {
  939. 'segment': segment_document,
  940. 'keywords': segment_item['keywords']
  941. }
  942. pre_segment_data_list.append(pre_segment_data)
  943. try:
  944. # save vector index
  945. VectorService.multi_create_segment_vector(pre_segment_data_list, dataset)
  946. except Exception as e:
  947. logging.exception("create segment index failed")
  948. for segment_document in segment_data_list:
  949. segment_document.enabled = False
  950. segment_document.disabled_at = datetime.datetime.utcnow()
  951. segment_document.status = 'error'
  952. segment_document.error = str(e)
  953. db.session.commit()
  954. return segment_data_list
  955. @classmethod
  956. def update_segment(cls, args: dict, segment: DocumentSegment, document: Document, dataset: Dataset):
  957. indexing_cache_key = 'segment_{}_indexing'.format(segment.id)
  958. cache_result = redis_client.get(indexing_cache_key)
  959. if cache_result is not None:
  960. raise ValueError("Segment is indexing, please try again later")
  961. try:
  962. content = args['content']
  963. if segment.content == content:
  964. if document.doc_form == 'qa_model':
  965. segment.answer = args['answer']
  966. if args['keywords']:
  967. segment.keywords = args['keywords']
  968. if args['enabled'] is not None:
  969. segment.enabled = args['enabled']
  970. db.session.add(segment)
  971. db.session.commit()
  972. # update segment index task
  973. if args['keywords']:
  974. kw_index = IndexBuilder.get_index(dataset, 'economy')
  975. # delete from keyword index
  976. kw_index.delete_by_ids([segment.index_node_id])
  977. # save keyword index
  978. kw_index.update_segment_keywords_index(segment.index_node_id, segment.keywords)
  979. else:
  980. segment_hash = helper.generate_text_hash(content)
  981. tokens = 0
  982. if dataset.indexing_technique == 'high_quality':
  983. embedding_model = ModelFactory.get_embedding_model(
  984. tenant_id=dataset.tenant_id,
  985. model_provider_name=dataset.embedding_model_provider,
  986. model_name=dataset.embedding_model
  987. )
  988. # calc embedding use tokens
  989. tokens = embedding_model.get_num_tokens(content)
  990. segment.content = content
  991. segment.index_node_hash = segment_hash
  992. segment.word_count = len(content)
  993. segment.tokens = tokens
  994. segment.status = 'completed'
  995. segment.indexing_at = datetime.datetime.utcnow()
  996. segment.completed_at = datetime.datetime.utcnow()
  997. segment.updated_by = current_user.id
  998. segment.updated_at = datetime.datetime.utcnow()
  999. if document.doc_form == 'qa_model':
  1000. segment.answer = args['answer']
  1001. db.session.add(segment)
  1002. db.session.commit()
  1003. # update segment vector index
  1004. VectorService.update_segment_vector(args['keywords'], segment, dataset)
  1005. except Exception as e:
  1006. logging.exception("update segment index failed")
  1007. segment.enabled = False
  1008. segment.disabled_at = datetime.datetime.utcnow()
  1009. segment.status = 'error'
  1010. segment.error = str(e)
  1011. db.session.commit()
  1012. segment = db.session.query(DocumentSegment).filter(DocumentSegment.id == segment.id).first()
  1013. return segment
  1014. @classmethod
  1015. def delete_segment(cls, segment: DocumentSegment, document: Document, dataset: Dataset):
  1016. indexing_cache_key = 'segment_{}_delete_indexing'.format(segment.id)
  1017. cache_result = redis_client.get(indexing_cache_key)
  1018. if cache_result is not None:
  1019. raise ValueError("Segment is deleting.")
  1020. # enabled segment need to delete index
  1021. if segment.enabled:
  1022. # send delete segment index task
  1023. redis_client.setex(indexing_cache_key, 600, 1)
  1024. delete_segment_from_index_task.delay(segment.id, segment.index_node_id, dataset.id, document.id)
  1025. db.session.delete(segment)
  1026. db.session.commit()
  1027. class DatasetCollectionBindingService:
  1028. @classmethod
  1029. def get_dataset_collection_binding(cls, provider_name: str, model_name: str) -> DatasetCollectionBinding:
  1030. dataset_collection_binding = db.session.query(DatasetCollectionBinding). \
  1031. filter(DatasetCollectionBinding.provider_name == provider_name,
  1032. DatasetCollectionBinding.model_name == model_name). \
  1033. order_by(DatasetCollectionBinding.created_at). \
  1034. first()
  1035. if not dataset_collection_binding:
  1036. dataset_collection_binding = DatasetCollectionBinding(
  1037. provider_name=provider_name,
  1038. model_name=model_name,
  1039. collection_name="Vector_index_" + str(uuid.uuid4()).replace("-", "_") + '_Node'
  1040. )
  1041. db.session.add(dataset_collection_binding)
  1042. db.session.flush()
  1043. return dataset_collection_binding