datasets_document.py 42 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048
  1. import logging
  2. from argparse import ArgumentTypeError
  3. from datetime import UTC, datetime
  4. from typing import cast
  5. from flask import request
  6. from flask_login import current_user # type: ignore
  7. from flask_restful import Resource, fields, marshal, marshal_with, reqparse # type: ignore
  8. from sqlalchemy import asc, desc
  9. from werkzeug.exceptions import Forbidden, NotFound
  10. import services
  11. from controllers.console import api
  12. from controllers.console.app.error import (
  13. ProviderModelCurrentlyNotSupportError,
  14. ProviderNotInitializeError,
  15. ProviderQuotaExceededError,
  16. )
  17. from controllers.console.datasets.error import (
  18. ArchivedDocumentImmutableError,
  19. DocumentAlreadyFinishedError,
  20. DocumentIndexingError,
  21. IndexingEstimateError,
  22. InvalidActionError,
  23. InvalidMetadataError,
  24. )
  25. from controllers.console.wraps import (
  26. account_initialization_required,
  27. cloud_edition_billing_resource_check,
  28. setup_required,
  29. )
  30. from core.errors.error import (
  31. LLMBadRequestError,
  32. ModelCurrentlyNotSupportError,
  33. ProviderTokenNotInitError,
  34. QuotaExceededError,
  35. )
  36. from core.indexing_runner import IndexingRunner
  37. from core.model_manager import ModelManager
  38. from core.model_runtime.entities.model_entities import ModelType
  39. from core.model_runtime.errors.invoke import InvokeAuthorizationError
  40. from core.plugin.manager.exc import PluginDaemonClientSideError
  41. from core.rag.extractor.entity.extract_setting import ExtractSetting
  42. from extensions.ext_database import db
  43. from extensions.ext_redis import redis_client
  44. from fields.document_fields import (
  45. dataset_and_document_fields,
  46. document_fields,
  47. document_status_fields,
  48. document_with_segments_fields,
  49. )
  50. from libs.login import login_required
  51. from models import Dataset, DatasetProcessRule, Document, DocumentSegment, UploadFile
  52. from services.dataset_service import DatasetService, DocumentService
  53. from tasks.add_document_to_index_task import add_document_to_index_task
  54. from tasks.remove_document_from_index_task import remove_document_from_index_task
  55. class DocumentResource(Resource):
  56. def get_document(self, dataset_id: str, document_id: str) -> Document:
  57. dataset = DatasetService.get_dataset(dataset_id)
  58. if not dataset:
  59. raise NotFound("Dataset not found.")
  60. try:
  61. DatasetService.check_dataset_permission(dataset, current_user)
  62. except services.errors.account.NoPermissionError as e:
  63. raise Forbidden(str(e))
  64. document = DocumentService.get_document(dataset_id, document_id)
  65. if not document:
  66. raise NotFound("Document not found.")
  67. if document.tenant_id != current_user.current_tenant_id:
  68. raise Forbidden("No permission.")
  69. return document
  70. def get_batch_documents(self, dataset_id: str, batch: str) -> list[Document]:
  71. dataset = DatasetService.get_dataset(dataset_id)
  72. if not dataset:
  73. raise NotFound("Dataset not found.")
  74. try:
  75. DatasetService.check_dataset_permission(dataset, current_user)
  76. except services.errors.account.NoPermissionError as e:
  77. raise Forbidden(str(e))
  78. documents = DocumentService.get_batch_documents(dataset_id, batch)
  79. if not documents:
  80. raise NotFound("Documents not found.")
  81. return documents
  82. class GetProcessRuleApi(Resource):
  83. @setup_required
  84. @login_required
  85. @account_initialization_required
  86. def get(self):
  87. req_data = request.args
  88. document_id = req_data.get("document_id")
  89. # get default rules
  90. mode = DocumentService.DEFAULT_RULES["mode"]
  91. rules = DocumentService.DEFAULT_RULES["rules"]
  92. limits = DocumentService.DEFAULT_RULES["limits"]
  93. if document_id:
  94. # get the latest process rule
  95. document = Document.query.get_or_404(document_id)
  96. dataset = DatasetService.get_dataset(document.dataset_id)
  97. if not dataset:
  98. raise NotFound("Dataset not found.")
  99. try:
  100. DatasetService.check_dataset_permission(dataset, current_user)
  101. except services.errors.account.NoPermissionError as e:
  102. raise Forbidden(str(e))
  103. # get the latest process rule
  104. dataset_process_rule = (
  105. db.session.query(DatasetProcessRule)
  106. .filter(DatasetProcessRule.dataset_id == document.dataset_id)
  107. .order_by(DatasetProcessRule.created_at.desc())
  108. .limit(1)
  109. .one_or_none()
  110. )
  111. if dataset_process_rule:
  112. mode = dataset_process_rule.mode
  113. rules = dataset_process_rule.rules_dict
  114. return {"mode": mode, "rules": rules, "limits": limits}
  115. class DatasetDocumentListApi(Resource):
  116. @setup_required
  117. @login_required
  118. @account_initialization_required
  119. def get(self, dataset_id):
  120. dataset_id = str(dataset_id)
  121. page = request.args.get("page", default=1, type=int)
  122. limit = request.args.get("limit", default=20, type=int)
  123. search = request.args.get("keyword", default=None, type=str)
  124. sort = request.args.get("sort", default="-created_at", type=str)
  125. # "yes", "true", "t", "y", "1" convert to True, while others convert to False.
  126. try:
  127. fetch_val = request.args.get("fetch", default="false")
  128. if isinstance(fetch_val, bool):
  129. fetch = fetch_val
  130. else:
  131. if fetch_val.lower() in ("yes", "true", "t", "y", "1"):
  132. fetch = True
  133. elif fetch_val.lower() in ("no", "false", "f", "n", "0"):
  134. fetch = False
  135. else:
  136. raise ArgumentTypeError(
  137. f"Truthy value expected: got {fetch_val} but expected one of yes/no, true/false, t/f, y/n, 1/0 "
  138. f"(case insensitive)."
  139. )
  140. except (ArgumentTypeError, ValueError, Exception):
  141. fetch = False
  142. dataset = DatasetService.get_dataset(dataset_id)
  143. if not dataset:
  144. raise NotFound("Dataset not found.")
  145. try:
  146. DatasetService.check_dataset_permission(dataset, current_user)
  147. except services.errors.account.NoPermissionError as e:
  148. raise Forbidden(str(e))
  149. query = Document.query.filter_by(dataset_id=str(dataset_id), tenant_id=current_user.current_tenant_id)
  150. if search:
  151. search = f"%{search}%"
  152. query = query.filter(Document.name.like(search))
  153. if sort.startswith("-"):
  154. sort_logic = desc
  155. sort = sort[1:]
  156. else:
  157. sort_logic = asc
  158. if sort == "hit_count":
  159. sub_query = (
  160. db.select(DocumentSegment.document_id, db.func.sum(DocumentSegment.hit_count).label("total_hit_count"))
  161. .group_by(DocumentSegment.document_id)
  162. .subquery()
  163. )
  164. query = query.outerjoin(sub_query, sub_query.c.document_id == Document.id).order_by(
  165. sort_logic(db.func.coalesce(sub_query.c.total_hit_count, 0)),
  166. sort_logic(Document.position),
  167. )
  168. elif sort == "created_at":
  169. query = query.order_by(
  170. sort_logic(Document.created_at),
  171. sort_logic(Document.position),
  172. )
  173. else:
  174. query = query.order_by(
  175. desc(Document.created_at),
  176. desc(Document.position),
  177. )
  178. paginated_documents = query.paginate(page=page, per_page=limit, max_per_page=100, error_out=False)
  179. documents = paginated_documents.items
  180. if fetch:
  181. for document in documents:
  182. completed_segments = DocumentSegment.query.filter(
  183. DocumentSegment.completed_at.isnot(None),
  184. DocumentSegment.document_id == str(document.id),
  185. DocumentSegment.status != "re_segment",
  186. ).count()
  187. total_segments = DocumentSegment.query.filter(
  188. DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment"
  189. ).count()
  190. document.completed_segments = completed_segments
  191. document.total_segments = total_segments
  192. data = marshal(documents, document_with_segments_fields)
  193. else:
  194. data = marshal(documents, document_fields)
  195. response = {
  196. "data": data,
  197. "has_more": len(documents) == limit,
  198. "limit": limit,
  199. "total": paginated_documents.total,
  200. "page": page,
  201. }
  202. return response
  203. documents_and_batch_fields = {"documents": fields.List(fields.Nested(document_fields)), "batch": fields.String}
  204. @setup_required
  205. @login_required
  206. @account_initialization_required
  207. @marshal_with(documents_and_batch_fields)
  208. @cloud_edition_billing_resource_check("vector_space")
  209. def post(self, dataset_id):
  210. dataset_id = str(dataset_id)
  211. dataset = DatasetService.get_dataset(dataset_id)
  212. if not dataset:
  213. raise NotFound("Dataset not found.")
  214. # The role of the current user in the ta table must be admin, owner, or editor
  215. if not current_user.is_dataset_editor:
  216. raise Forbidden()
  217. try:
  218. DatasetService.check_dataset_permission(dataset, current_user)
  219. except services.errors.account.NoPermissionError as e:
  220. raise Forbidden(str(e))
  221. parser = reqparse.RequestParser()
  222. parser.add_argument(
  223. "indexing_technique", type=str, choices=Dataset.INDEXING_TECHNIQUE_LIST, nullable=False, location="json"
  224. )
  225. parser.add_argument("data_source", type=dict, required=False, location="json")
  226. parser.add_argument("process_rule", type=dict, required=False, location="json")
  227. parser.add_argument("duplicate", type=bool, default=True, nullable=False, location="json")
  228. parser.add_argument("original_document_id", type=str, required=False, location="json")
  229. parser.add_argument("doc_form", type=str, default="text_model", required=False, nullable=False, location="json")
  230. parser.add_argument(
  231. "doc_language", type=str, default="English", required=False, nullable=False, location="json"
  232. )
  233. parser.add_argument("retrieval_model", type=dict, required=False, nullable=False, location="json")
  234. args = parser.parse_args()
  235. if not dataset.indexing_technique and not args["indexing_technique"]:
  236. raise ValueError("indexing_technique is required.")
  237. # validate args
  238. DocumentService.document_create_args_validate(args)
  239. try:
  240. documents, batch = DocumentService.save_document_with_dataset_id(dataset, args, current_user)
  241. except ProviderTokenNotInitError as ex:
  242. raise ProviderNotInitializeError(ex.description)
  243. except QuotaExceededError:
  244. raise ProviderQuotaExceededError()
  245. except ModelCurrentlyNotSupportError:
  246. raise ProviderModelCurrentlyNotSupportError()
  247. return {"documents": documents, "batch": batch}
  248. class DatasetInitApi(Resource):
  249. @setup_required
  250. @login_required
  251. @account_initialization_required
  252. @marshal_with(dataset_and_document_fields)
  253. @cloud_edition_billing_resource_check("vector_space")
  254. def post(self):
  255. # The role of the current user in the ta table must be admin, owner, or editor
  256. if not current_user.is_editor:
  257. raise Forbidden()
  258. parser = reqparse.RequestParser()
  259. parser.add_argument(
  260. "indexing_technique",
  261. type=str,
  262. choices=Dataset.INDEXING_TECHNIQUE_LIST,
  263. required=True,
  264. nullable=False,
  265. location="json",
  266. )
  267. parser.add_argument("data_source", type=dict, required=True, nullable=True, location="json")
  268. parser.add_argument("process_rule", type=dict, required=True, nullable=True, location="json")
  269. parser.add_argument("doc_form", type=str, default="text_model", required=False, nullable=False, location="json")
  270. parser.add_argument(
  271. "doc_language", type=str, default="English", required=False, nullable=False, location="json"
  272. )
  273. parser.add_argument("retrieval_model", type=dict, required=False, nullable=False, location="json")
  274. parser.add_argument("embedding_model", type=str, required=False, nullable=True, location="json")
  275. parser.add_argument("embedding_model_provider", type=str, required=False, nullable=True, location="json")
  276. args = parser.parse_args()
  277. # The role of the current user in the ta table must be admin, owner, or editor, or dataset_operator
  278. if not current_user.is_dataset_editor:
  279. raise Forbidden()
  280. if args["indexing_technique"] == "high_quality":
  281. if args["embedding_model"] is None or args["embedding_model_provider"] is None:
  282. raise ValueError("embedding model and embedding model provider are required for high quality indexing.")
  283. try:
  284. model_manager = ModelManager()
  285. model_manager.get_model_instance(
  286. tenant_id=current_user.current_tenant_id,
  287. provider=args["embedding_model_provider"],
  288. model_type=ModelType.TEXT_EMBEDDING,
  289. model=args["embedding_model"],
  290. )
  291. except InvokeAuthorizationError:
  292. raise ProviderNotInitializeError(
  293. "No Embedding Model available. Please configure a valid provider "
  294. "in the Settings -> Model Provider."
  295. )
  296. except ProviderTokenNotInitError as ex:
  297. raise ProviderNotInitializeError(ex.description)
  298. # validate args
  299. DocumentService.document_create_args_validate(args)
  300. try:
  301. dataset, documents, batch = DocumentService.save_document_without_dataset_id(
  302. tenant_id=current_user.current_tenant_id, document_data=args, account=current_user
  303. )
  304. except ProviderTokenNotInitError as ex:
  305. raise ProviderNotInitializeError(ex.description)
  306. except QuotaExceededError:
  307. raise ProviderQuotaExceededError()
  308. except ModelCurrentlyNotSupportError:
  309. raise ProviderModelCurrentlyNotSupportError()
  310. response = {"dataset": dataset, "documents": documents, "batch": batch}
  311. return response
  312. class DocumentIndexingEstimateApi(DocumentResource):
  313. @setup_required
  314. @login_required
  315. @account_initialization_required
  316. def get(self, dataset_id, document_id):
  317. dataset_id = str(dataset_id)
  318. document_id = str(document_id)
  319. document = self.get_document(dataset_id, document_id)
  320. if document.indexing_status in {"completed", "error"}:
  321. raise DocumentAlreadyFinishedError()
  322. data_process_rule = document.dataset_process_rule
  323. data_process_rule_dict = data_process_rule.to_dict()
  324. response = {"tokens": 0, "total_price": 0, "currency": "USD", "total_segments": 0, "preview": []}
  325. if document.data_source_type == "upload_file":
  326. data_source_info = document.data_source_info_dict
  327. if data_source_info and "upload_file_id" in data_source_info:
  328. file_id = data_source_info["upload_file_id"]
  329. file = (
  330. db.session.query(UploadFile)
  331. .filter(UploadFile.tenant_id == document.tenant_id, UploadFile.id == file_id)
  332. .first()
  333. )
  334. # raise error if file not found
  335. if not file:
  336. raise NotFound("File not found.")
  337. extract_setting = ExtractSetting(
  338. datasource_type="upload_file", upload_file=file, document_model=document.doc_form
  339. )
  340. indexing_runner = IndexingRunner()
  341. try:
  342. response = indexing_runner.indexing_estimate(
  343. current_user.current_tenant_id,
  344. [extract_setting],
  345. data_process_rule_dict,
  346. document.doc_form,
  347. "English",
  348. dataset_id,
  349. )
  350. except LLMBadRequestError:
  351. raise ProviderNotInitializeError(
  352. "No Embedding Model available. Please configure a valid provider "
  353. "in the Settings -> Model Provider."
  354. )
  355. except ProviderTokenNotInitError as ex:
  356. raise ProviderNotInitializeError(ex.description)
  357. except PluginDaemonClientSideError as ex:
  358. raise ProviderNotInitializeError(ex.description)
  359. except Exception as e:
  360. raise IndexingEstimateError(str(e))
  361. return response
  362. class DocumentBatchIndexingEstimateApi(DocumentResource):
  363. @setup_required
  364. @login_required
  365. @account_initialization_required
  366. def get(self, dataset_id, batch):
  367. dataset_id = str(dataset_id)
  368. batch = str(batch)
  369. documents = self.get_batch_documents(dataset_id, batch)
  370. response = {"tokens": 0, "total_price": 0, "currency": "USD", "total_segments": 0, "preview": []}
  371. if not documents:
  372. return response
  373. data_process_rule = documents[0].dataset_process_rule
  374. data_process_rule_dict = data_process_rule.to_dict()
  375. info_list = []
  376. extract_settings = []
  377. for document in documents:
  378. if document.indexing_status in {"completed", "error"}:
  379. raise DocumentAlreadyFinishedError()
  380. data_source_info = document.data_source_info_dict
  381. # format document files info
  382. if data_source_info and "upload_file_id" in data_source_info:
  383. file_id = data_source_info["upload_file_id"]
  384. info_list.append(file_id)
  385. # format document notion info
  386. elif (
  387. data_source_info and "notion_workspace_id" in data_source_info and "notion_page_id" in data_source_info
  388. ):
  389. pages = []
  390. page = {"page_id": data_source_info["notion_page_id"], "type": data_source_info["type"]}
  391. pages.append(page)
  392. notion_info = {"workspace_id": data_source_info["notion_workspace_id"], "pages": pages}
  393. info_list.append(notion_info)
  394. if document.data_source_type == "upload_file":
  395. file_id = data_source_info["upload_file_id"]
  396. file_detail = (
  397. db.session.query(UploadFile)
  398. .filter(UploadFile.tenant_id == current_user.current_tenant_id, UploadFile.id == file_id)
  399. .first()
  400. )
  401. if file_detail is None:
  402. raise NotFound("File not found.")
  403. extract_setting = ExtractSetting(
  404. datasource_type="upload_file", upload_file=file_detail, document_model=document.doc_form
  405. )
  406. extract_settings.append(extract_setting)
  407. elif document.data_source_type == "notion_import":
  408. extract_setting = ExtractSetting(
  409. datasource_type="notion_import",
  410. notion_info={
  411. "notion_workspace_id": data_source_info["notion_workspace_id"],
  412. "notion_obj_id": data_source_info["notion_page_id"],
  413. "notion_page_type": data_source_info["type"],
  414. "tenant_id": current_user.current_tenant_id,
  415. },
  416. document_model=document.doc_form,
  417. )
  418. extract_settings.append(extract_setting)
  419. elif document.data_source_type == "website_crawl":
  420. extract_setting = ExtractSetting(
  421. datasource_type="website_crawl",
  422. website_info={
  423. "provider": data_source_info["provider"],
  424. "job_id": data_source_info["job_id"],
  425. "url": data_source_info["url"],
  426. "tenant_id": current_user.current_tenant_id,
  427. "mode": data_source_info["mode"],
  428. "only_main_content": data_source_info["only_main_content"],
  429. },
  430. document_model=document.doc_form,
  431. )
  432. extract_settings.append(extract_setting)
  433. else:
  434. raise ValueError("Data source type not support")
  435. indexing_runner = IndexingRunner()
  436. try:
  437. response = indexing_runner.indexing_estimate(
  438. current_user.current_tenant_id,
  439. extract_settings,
  440. data_process_rule_dict,
  441. document.doc_form,
  442. "English",
  443. dataset_id,
  444. )
  445. except LLMBadRequestError:
  446. raise ProviderNotInitializeError(
  447. "No Embedding Model available. Please configure a valid provider "
  448. "in the Settings -> Model Provider."
  449. )
  450. except ProviderTokenNotInitError as ex:
  451. raise ProviderNotInitializeError(ex.description)
  452. except PluginDaemonClientSideError as ex:
  453. raise ProviderNotInitializeError(ex.description)
  454. except Exception as e:
  455. raise IndexingEstimateError(str(e))
  456. return response
  457. class DocumentBatchIndexingStatusApi(DocumentResource):
  458. @setup_required
  459. @login_required
  460. @account_initialization_required
  461. def get(self, dataset_id, batch):
  462. dataset_id = str(dataset_id)
  463. batch = str(batch)
  464. documents = self.get_batch_documents(dataset_id, batch)
  465. documents_status = []
  466. for document in documents:
  467. completed_segments = DocumentSegment.query.filter(
  468. DocumentSegment.completed_at.isnot(None),
  469. DocumentSegment.document_id == str(document.id),
  470. DocumentSegment.status != "re_segment",
  471. ).count()
  472. total_segments = DocumentSegment.query.filter(
  473. DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment"
  474. ).count()
  475. document.completed_segments = completed_segments
  476. document.total_segments = total_segments
  477. if document.is_paused:
  478. document.indexing_status = "paused"
  479. documents_status.append(marshal(document, document_status_fields))
  480. data = {"data": documents_status}
  481. return data
  482. class DocumentIndexingStatusApi(DocumentResource):
  483. @setup_required
  484. @login_required
  485. @account_initialization_required
  486. def get(self, dataset_id, document_id):
  487. dataset_id = str(dataset_id)
  488. document_id = str(document_id)
  489. document = self.get_document(dataset_id, document_id)
  490. completed_segments = DocumentSegment.query.filter(
  491. DocumentSegment.completed_at.isnot(None),
  492. DocumentSegment.document_id == str(document_id),
  493. DocumentSegment.status != "re_segment",
  494. ).count()
  495. total_segments = DocumentSegment.query.filter(
  496. DocumentSegment.document_id == str(document_id), DocumentSegment.status != "re_segment"
  497. ).count()
  498. document.completed_segments = completed_segments
  499. document.total_segments = total_segments
  500. if document.is_paused:
  501. document.indexing_status = "paused"
  502. return marshal(document, document_status_fields)
  503. class DocumentDetailApi(DocumentResource):
  504. METADATA_CHOICES = {"all", "only", "without"}
  505. @setup_required
  506. @login_required
  507. @account_initialization_required
  508. def get(self, dataset_id, document_id):
  509. dataset_id = str(dataset_id)
  510. document_id = str(document_id)
  511. document = self.get_document(dataset_id, document_id)
  512. metadata = request.args.get("metadata", "all")
  513. if metadata not in self.METADATA_CHOICES:
  514. raise InvalidMetadataError(f"Invalid metadata value: {metadata}")
  515. if metadata == "only":
  516. response = {"id": document.id, "doc_type": document.doc_type, "doc_metadata": document.doc_metadata}
  517. elif metadata == "without":
  518. process_rules = DatasetService.get_process_rules(dataset_id)
  519. data_source_info = document.data_source_detail_dict
  520. response = {
  521. "id": document.id,
  522. "position": document.position,
  523. "data_source_type": document.data_source_type,
  524. "data_source_info": data_source_info,
  525. "dataset_process_rule_id": document.dataset_process_rule_id,
  526. "dataset_process_rule": process_rules,
  527. "name": document.name,
  528. "created_from": document.created_from,
  529. "created_by": document.created_by,
  530. "created_at": document.created_at.timestamp(),
  531. "tokens": document.tokens,
  532. "indexing_status": document.indexing_status,
  533. "completed_at": int(document.completed_at.timestamp()) if document.completed_at else None,
  534. "updated_at": int(document.updated_at.timestamp()) if document.updated_at else None,
  535. "indexing_latency": document.indexing_latency,
  536. "error": document.error,
  537. "enabled": document.enabled,
  538. "disabled_at": int(document.disabled_at.timestamp()) if document.disabled_at else None,
  539. "disabled_by": document.disabled_by,
  540. "archived": document.archived,
  541. "segment_count": document.segment_count,
  542. "average_segment_length": document.average_segment_length,
  543. "hit_count": document.hit_count,
  544. "display_status": document.display_status,
  545. "doc_form": document.doc_form,
  546. "doc_language": document.doc_language,
  547. }
  548. else:
  549. process_rules = DatasetService.get_process_rules(dataset_id)
  550. data_source_info = document.data_source_detail_dict
  551. response = {
  552. "id": document.id,
  553. "position": document.position,
  554. "data_source_type": document.data_source_type,
  555. "data_source_info": data_source_info,
  556. "dataset_process_rule_id": document.dataset_process_rule_id,
  557. "dataset_process_rule": process_rules,
  558. "name": document.name,
  559. "created_from": document.created_from,
  560. "created_by": document.created_by,
  561. "created_at": document.created_at.timestamp(),
  562. "tokens": document.tokens,
  563. "indexing_status": document.indexing_status,
  564. "completed_at": int(document.completed_at.timestamp()) if document.completed_at else None,
  565. "updated_at": int(document.updated_at.timestamp()) if document.updated_at else None,
  566. "indexing_latency": document.indexing_latency,
  567. "error": document.error,
  568. "enabled": document.enabled,
  569. "disabled_at": int(document.disabled_at.timestamp()) if document.disabled_at else None,
  570. "disabled_by": document.disabled_by,
  571. "archived": document.archived,
  572. "doc_type": document.doc_type,
  573. "doc_metadata": document.doc_metadata,
  574. "segment_count": document.segment_count,
  575. "average_segment_length": document.average_segment_length,
  576. "hit_count": document.hit_count,
  577. "display_status": document.display_status,
  578. "doc_form": document.doc_form,
  579. "doc_language": document.doc_language,
  580. }
  581. return response, 200
  582. class DocumentProcessingApi(DocumentResource):
  583. @setup_required
  584. @login_required
  585. @account_initialization_required
  586. def patch(self, dataset_id, document_id, action):
  587. dataset_id = str(dataset_id)
  588. document_id = str(document_id)
  589. document = self.get_document(dataset_id, document_id)
  590. # The role of the current user in the ta table must be admin, owner, or editor
  591. if not current_user.is_editor:
  592. raise Forbidden()
  593. if action == "pause":
  594. if document.indexing_status != "indexing":
  595. raise InvalidActionError("Document not in indexing state.")
  596. document.paused_by = current_user.id
  597. document.paused_at = datetime.now(UTC).replace(tzinfo=None)
  598. document.is_paused = True
  599. db.session.commit()
  600. elif action == "resume":
  601. if document.indexing_status not in {"paused", "error"}:
  602. raise InvalidActionError("Document not in paused or error state.")
  603. document.paused_by = None
  604. document.paused_at = None
  605. document.is_paused = False
  606. db.session.commit()
  607. else:
  608. raise InvalidActionError()
  609. return {"result": "success"}, 200
  610. class DocumentDeleteApi(DocumentResource):
  611. @setup_required
  612. @login_required
  613. @account_initialization_required
  614. def delete(self, dataset_id, document_id):
  615. dataset_id = str(dataset_id)
  616. document_id = str(document_id)
  617. dataset = DatasetService.get_dataset(dataset_id)
  618. if dataset is None:
  619. raise NotFound("Dataset not found.")
  620. # check user's model setting
  621. DatasetService.check_dataset_model_setting(dataset)
  622. document = self.get_document(dataset_id, document_id)
  623. try:
  624. DocumentService.delete_document(document)
  625. except services.errors.document.DocumentIndexingError:
  626. raise DocumentIndexingError("Cannot delete document during indexing.")
  627. return {"result": "success"}, 204
  628. class DocumentMetadataApi(DocumentResource):
  629. @setup_required
  630. @login_required
  631. @account_initialization_required
  632. def put(self, dataset_id, document_id):
  633. dataset_id = str(dataset_id)
  634. document_id = str(document_id)
  635. document = self.get_document(dataset_id, document_id)
  636. req_data = request.get_json()
  637. doc_type = req_data.get("doc_type")
  638. doc_metadata = req_data.get("doc_metadata")
  639. # The role of the current user in the ta table must be admin, owner, or editor
  640. if not current_user.is_editor:
  641. raise Forbidden()
  642. if doc_type is None or doc_metadata is None:
  643. raise ValueError("Both doc_type and doc_metadata must be provided.")
  644. if doc_type not in DocumentService.DOCUMENT_METADATA_SCHEMA:
  645. raise ValueError("Invalid doc_type.")
  646. if not isinstance(doc_metadata, dict):
  647. raise ValueError("doc_metadata must be a dictionary.")
  648. metadata_schema: dict = cast(dict, DocumentService.DOCUMENT_METADATA_SCHEMA[doc_type])
  649. document.doc_metadata = {}
  650. if doc_type == "others":
  651. document.doc_metadata = doc_metadata
  652. else:
  653. for key, value_type in metadata_schema.items():
  654. value = doc_metadata.get(key)
  655. if value is not None and isinstance(value, value_type):
  656. document.doc_metadata[key] = value
  657. document.doc_type = doc_type
  658. document.updated_at = datetime.now(UTC).replace(tzinfo=None)
  659. db.session.commit()
  660. return {"result": "success", "message": "Document metadata updated."}, 200
  661. class DocumentStatusApi(DocumentResource):
  662. @setup_required
  663. @login_required
  664. @account_initialization_required
  665. @cloud_edition_billing_resource_check("vector_space")
  666. def patch(self, dataset_id, document_id, action):
  667. dataset_id = str(dataset_id)
  668. document_id = str(document_id)
  669. dataset = DatasetService.get_dataset(dataset_id)
  670. if dataset is None:
  671. raise NotFound("Dataset not found.")
  672. # The role of the current user in the ta table must be admin, owner, or editor
  673. if not current_user.is_dataset_editor:
  674. raise Forbidden()
  675. # check user's model setting
  676. DatasetService.check_dataset_model_setting(dataset)
  677. # check user's permission
  678. DatasetService.check_dataset_permission(dataset, current_user)
  679. document = self.get_document(dataset_id, document_id)
  680. indexing_cache_key = "document_{}_indexing".format(document.id)
  681. cache_result = redis_client.get(indexing_cache_key)
  682. if cache_result is not None:
  683. raise InvalidActionError("Document is being indexed, please try again later")
  684. if action == "enable":
  685. if document.enabled:
  686. raise InvalidActionError("Document already enabled.")
  687. document.enabled = True
  688. document.disabled_at = None
  689. document.disabled_by = None
  690. document.updated_at = datetime.now(UTC).replace(tzinfo=None)
  691. db.session.commit()
  692. # Set cache to prevent indexing the same document multiple times
  693. redis_client.setex(indexing_cache_key, 600, 1)
  694. add_document_to_index_task.delay(document_id)
  695. return {"result": "success"}, 200
  696. elif action == "disable":
  697. if not document.completed_at or document.indexing_status != "completed":
  698. raise InvalidActionError("Document is not completed.")
  699. if not document.enabled:
  700. raise InvalidActionError("Document already disabled.")
  701. document.enabled = False
  702. document.disabled_at = datetime.now(UTC).replace(tzinfo=None)
  703. document.disabled_by = current_user.id
  704. document.updated_at = datetime.now(UTC).replace(tzinfo=None)
  705. db.session.commit()
  706. # Set cache to prevent indexing the same document multiple times
  707. redis_client.setex(indexing_cache_key, 600, 1)
  708. remove_document_from_index_task.delay(document_id)
  709. return {"result": "success"}, 200
  710. elif action == "archive":
  711. if document.archived:
  712. raise InvalidActionError("Document already archived.")
  713. document.archived = True
  714. document.archived_at = datetime.now(UTC).replace(tzinfo=None)
  715. document.archived_by = current_user.id
  716. document.updated_at = datetime.now(UTC).replace(tzinfo=None)
  717. db.session.commit()
  718. if document.enabled:
  719. # Set cache to prevent indexing the same document multiple times
  720. redis_client.setex(indexing_cache_key, 600, 1)
  721. remove_document_from_index_task.delay(document_id)
  722. return {"result": "success"}, 200
  723. elif action == "un_archive":
  724. if not document.archived:
  725. raise InvalidActionError("Document is not archived.")
  726. document.archived = False
  727. document.archived_at = None
  728. document.archived_by = None
  729. document.updated_at = datetime.now(UTC).replace(tzinfo=None)
  730. db.session.commit()
  731. # Set cache to prevent indexing the same document multiple times
  732. redis_client.setex(indexing_cache_key, 600, 1)
  733. add_document_to_index_task.delay(document_id)
  734. return {"result": "success"}, 200
  735. else:
  736. raise InvalidActionError()
  737. class DocumentPauseApi(DocumentResource):
  738. @setup_required
  739. @login_required
  740. @account_initialization_required
  741. def patch(self, dataset_id, document_id):
  742. """pause document."""
  743. dataset_id = str(dataset_id)
  744. document_id = str(document_id)
  745. dataset = DatasetService.get_dataset(dataset_id)
  746. if not dataset:
  747. raise NotFound("Dataset not found.")
  748. document = DocumentService.get_document(dataset.id, document_id)
  749. # 404 if document not found
  750. if document is None:
  751. raise NotFound("Document Not Exists.")
  752. # 403 if document is archived
  753. if DocumentService.check_archived(document):
  754. raise ArchivedDocumentImmutableError()
  755. try:
  756. # pause document
  757. DocumentService.pause_document(document)
  758. except services.errors.document.DocumentIndexingError:
  759. raise DocumentIndexingError("Cannot pause completed document.")
  760. return {"result": "success"}, 204
  761. class DocumentRecoverApi(DocumentResource):
  762. @setup_required
  763. @login_required
  764. @account_initialization_required
  765. def patch(self, dataset_id, document_id):
  766. """recover document."""
  767. dataset_id = str(dataset_id)
  768. document_id = str(document_id)
  769. dataset = DatasetService.get_dataset(dataset_id)
  770. if not dataset:
  771. raise NotFound("Dataset not found.")
  772. document = DocumentService.get_document(dataset.id, document_id)
  773. # 404 if document not found
  774. if document is None:
  775. raise NotFound("Document Not Exists.")
  776. # 403 if document is archived
  777. if DocumentService.check_archived(document):
  778. raise ArchivedDocumentImmutableError()
  779. try:
  780. # pause document
  781. DocumentService.recover_document(document)
  782. except services.errors.document.DocumentIndexingError:
  783. raise DocumentIndexingError("Document is not in paused status.")
  784. return {"result": "success"}, 204
  785. class DocumentRetryApi(DocumentResource):
  786. @setup_required
  787. @login_required
  788. @account_initialization_required
  789. def post(self, dataset_id):
  790. """retry document."""
  791. parser = reqparse.RequestParser()
  792. parser.add_argument("document_ids", type=list, required=True, nullable=False, location="json")
  793. args = parser.parse_args()
  794. dataset_id = str(dataset_id)
  795. dataset = DatasetService.get_dataset(dataset_id)
  796. retry_documents = []
  797. if not dataset:
  798. raise NotFound("Dataset not found.")
  799. for document_id in args["document_ids"]:
  800. try:
  801. document_id = str(document_id)
  802. document = DocumentService.get_document(dataset.id, document_id)
  803. # 404 if document not found
  804. if document is None:
  805. raise NotFound("Document Not Exists.")
  806. # 403 if document is archived
  807. if DocumentService.check_archived(document):
  808. raise ArchivedDocumentImmutableError()
  809. # 400 if document is completed
  810. if document.indexing_status == "completed":
  811. raise DocumentAlreadyFinishedError()
  812. retry_documents.append(document)
  813. except Exception:
  814. logging.exception(f"Failed to retry document, document id: {document_id}")
  815. continue
  816. # retry document
  817. DocumentService.retry_document(dataset_id, retry_documents)
  818. return {"result": "success"}, 204
  819. class DocumentRenameApi(DocumentResource):
  820. @setup_required
  821. @login_required
  822. @account_initialization_required
  823. @marshal_with(document_fields)
  824. def post(self, dataset_id, document_id):
  825. # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator
  826. if not current_user.is_dataset_editor:
  827. raise Forbidden()
  828. dataset = DatasetService.get_dataset(dataset_id)
  829. DatasetService.check_dataset_operator_permission(current_user, dataset)
  830. parser = reqparse.RequestParser()
  831. parser.add_argument("name", type=str, required=True, nullable=False, location="json")
  832. args = parser.parse_args()
  833. try:
  834. document = DocumentService.rename_document(dataset_id, document_id, args["name"])
  835. except services.errors.document.DocumentIndexingError:
  836. raise DocumentIndexingError("Cannot delete document during indexing.")
  837. return document
  838. class WebsiteDocumentSyncApi(DocumentResource):
  839. @setup_required
  840. @login_required
  841. @account_initialization_required
  842. def get(self, dataset_id, document_id):
  843. """sync website document."""
  844. dataset_id = str(dataset_id)
  845. dataset = DatasetService.get_dataset(dataset_id)
  846. if not dataset:
  847. raise NotFound("Dataset not found.")
  848. document_id = str(document_id)
  849. document = DocumentService.get_document(dataset.id, document_id)
  850. if not document:
  851. raise NotFound("Document not found.")
  852. if document.tenant_id != current_user.current_tenant_id:
  853. raise Forbidden("No permission.")
  854. if document.data_source_type != "website_crawl":
  855. raise ValueError("Document is not a website document.")
  856. # 403 if document is archived
  857. if DocumentService.check_archived(document):
  858. raise ArchivedDocumentImmutableError()
  859. # sync document
  860. DocumentService.sync_website_document(dataset_id, document)
  861. return {"result": "success"}, 200
  862. api.add_resource(GetProcessRuleApi, "/datasets/process-rule")
  863. api.add_resource(DatasetDocumentListApi, "/datasets/<uuid:dataset_id>/documents")
  864. api.add_resource(DatasetInitApi, "/datasets/init")
  865. api.add_resource(
  866. DocumentIndexingEstimateApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/indexing-estimate"
  867. )
  868. api.add_resource(DocumentBatchIndexingEstimateApi, "/datasets/<uuid:dataset_id>/batch/<string:batch>/indexing-estimate")
  869. api.add_resource(DocumentBatchIndexingStatusApi, "/datasets/<uuid:dataset_id>/batch/<string:batch>/indexing-status")
  870. api.add_resource(DocumentIndexingStatusApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/indexing-status")
  871. api.add_resource(DocumentDetailApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>")
  872. api.add_resource(
  873. DocumentProcessingApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/<string:action>"
  874. )
  875. api.add_resource(DocumentDeleteApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>")
  876. api.add_resource(DocumentMetadataApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/metadata")
  877. api.add_resource(DocumentStatusApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/status/<string:action>")
  878. api.add_resource(DocumentPauseApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/pause")
  879. api.add_resource(DocumentRecoverApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/resume")
  880. api.add_resource(DocumentRetryApi, "/datasets/<uuid:dataset_id>/retry")
  881. api.add_resource(DocumentRenameApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/rename")
  882. api.add_resource(WebsiteDocumentSyncApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/website-sync")