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

datasets_document.py 49KB

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