datasets_document.py 61 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566
  1. import logging
  2. from argparse import ArgumentTypeError
  3. from typing import Literal, cast
  4. from urllib.parse import quote
  5. from flask import Response, request
  6. from flask_login import current_user
  7. from flask_restx import Resource, marshal, marshal_with, reqparse
  8. from sqlalchemy import asc, desc, select
  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_rate_limit_check,
  28. cloud_edition_billing_resource_check,
  29. setup_required,
  30. )
  31. from core.errors.error import (
  32. LLMBadRequestError,
  33. ModelCurrentlyNotSupportError,
  34. ProviderTokenNotInitError,
  35. QuotaExceededError,
  36. )
  37. from core.indexing_runner import IndexingRunner
  38. from core.model_manager import ModelManager
  39. from core.model_runtime.entities.model_entities import ModelType
  40. from core.model_runtime.errors.invoke import InvokeAuthorizationError
  41. from core.plugin.impl.exc import PluginDaemonClientSideError
  42. from core.rag.extractor.entity.datasource_type import DatasourceType
  43. from core.rag.extractor.entity.extract_setting import ExtractSetting
  44. from core.file import helpers as file_helpers
  45. from extensions.ext_database import db
  46. from extensions.ext_storage import storage
  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 services.dataset_service import DatasetService, DocumentService
  57. from services.document_download_service import DocumentDownloadService
  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) -> list[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. class GetProcessRuleApi(Resource):
  88. @setup_required
  89. @login_required
  90. @account_initialization_required
  91. def get(self):
  92. req_data = request.args
  93. document_id = req_data.get("document_id")
  94. # get default rules
  95. mode = DocumentService.DEFAULT_RULES["mode"]
  96. rules = DocumentService.DEFAULT_RULES["rules"]
  97. limits = DocumentService.DEFAULT_RULES["limits"]
  98. if document_id:
  99. # get the latest process rule
  100. document = db.get_or_404(Document, document_id)
  101. dataset = DatasetService.get_dataset(document.dataset_id)
  102. if not dataset:
  103. raise NotFound("Dataset not found.")
  104. try:
  105. DatasetService.check_dataset_permission(dataset, current_user)
  106. except services.errors.account.NoPermissionError as e:
  107. raise Forbidden(str(e))
  108. # get the latest process rule
  109. dataset_process_rule = (
  110. db.session.query(DatasetProcessRule)
  111. .where(DatasetProcessRule.dataset_id == document.dataset_id)
  112. .order_by(DatasetProcessRule.created_at.desc())
  113. .limit(1)
  114. .one_or_none()
  115. )
  116. if dataset_process_rule:
  117. mode = dataset_process_rule.mode
  118. rules = dataset_process_rule.rules_dict
  119. return {"mode": mode, "rules": rules, "limits": limits}
  120. class DatasetDocumentListApi(Resource):
  121. @setup_required
  122. @login_required
  123. @account_initialization_required
  124. def get(self, dataset_id):
  125. dataset_id = str(dataset_id)
  126. page = request.args.get("page", default=1, type=int)
  127. limit = request.args.get("limit", default=20, type=int)
  128. search = request.args.get("keyword", default=None, type=str)
  129. sort = request.args.get("sort", default="-created_at", type=str)
  130. # "yes", "true", "t", "y", "1" convert to True, while others convert to False.
  131. try:
  132. fetch_val = request.args.get("fetch", default="false")
  133. if isinstance(fetch_val, bool):
  134. fetch = fetch_val
  135. else:
  136. if fetch_val.lower() in ("yes", "true", "t", "y", "1"):
  137. fetch = True
  138. elif fetch_val.lower() in ("no", "false", "f", "n", "0"):
  139. fetch = False
  140. else:
  141. raise ArgumentTypeError(
  142. f"Truthy value expected: got {fetch_val} but expected one of yes/no, true/false, t/f, y/n, 1/0 "
  143. f"(case insensitive)."
  144. )
  145. except (ArgumentTypeError, ValueError, Exception):
  146. fetch = False
  147. dataset = DatasetService.get_dataset(dataset_id)
  148. if not dataset:
  149. raise NotFound("Dataset not found.")
  150. try:
  151. DatasetService.check_dataset_permission(dataset, current_user)
  152. except services.errors.account.NoPermissionError as e:
  153. raise Forbidden(str(e))
  154. query = select(Document).filter_by(dataset_id=str(dataset_id), tenant_id=current_user.current_tenant_id)
  155. if search:
  156. search = f"%{search}%"
  157. query = query.where(Document.name.like(search))
  158. if sort.startswith("-"):
  159. sort_logic = desc
  160. sort = sort[1:]
  161. else:
  162. sort_logic = asc
  163. if sort == "hit_count":
  164. sub_query = (
  165. db.select(DocumentSegment.document_id, db.func.sum(DocumentSegment.hit_count).label("total_hit_count"))
  166. .group_by(DocumentSegment.document_id)
  167. .subquery()
  168. )
  169. query = query.outerjoin(sub_query, sub_query.c.document_id == Document.id).order_by(
  170. sort_logic(db.func.coalesce(sub_query.c.total_hit_count, 0)),
  171. sort_logic(Document.position),
  172. )
  173. elif sort == "created_at":
  174. query = query.order_by(
  175. sort_logic(Document.created_at),
  176. sort_logic(Document.position),
  177. )
  178. else:
  179. query = query.order_by(
  180. desc(Document.created_at),
  181. desc(Document.position),
  182. )
  183. paginated_documents = db.paginate(select=query, page=page, per_page=limit, max_per_page=100, error_out=False)
  184. documents = paginated_documents.items
  185. if fetch:
  186. for document in documents:
  187. completed_segments = (
  188. db.session.query(DocumentSegment)
  189. .where(
  190. DocumentSegment.completed_at.isnot(None),
  191. DocumentSegment.document_id == str(document.id),
  192. DocumentSegment.status != "re_segment",
  193. )
  194. .count()
  195. )
  196. total_segments = (
  197. db.session.query(DocumentSegment)
  198. .where(DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment")
  199. .count()
  200. )
  201. document.completed_segments = completed_segments
  202. document.total_segments = total_segments
  203. data = marshal(documents, document_with_segments_fields)
  204. else:
  205. data = marshal(documents, document_fields)
  206. response = {
  207. "data": data,
  208. "has_more": len(documents) == limit,
  209. "limit": limit,
  210. "total": paginated_documents.total,
  211. "page": page,
  212. }
  213. return response
  214. @setup_required
  215. @login_required
  216. @account_initialization_required
  217. @marshal_with(dataset_and_document_fields)
  218. @cloud_edition_billing_resource_check("vector_space")
  219. @cloud_edition_billing_rate_limit_check("knowledge")
  220. def post(self, dataset_id):
  221. dataset_id = str(dataset_id)
  222. dataset = DatasetService.get_dataset(dataset_id)
  223. if not dataset:
  224. raise NotFound("Dataset not found.")
  225. # The role of the current user in the ta table must be admin, owner, or editor
  226. if not current_user.is_dataset_editor:
  227. raise Forbidden()
  228. try:
  229. DatasetService.check_dataset_permission(dataset, current_user)
  230. except services.errors.account.NoPermissionError as e:
  231. raise Forbidden(str(e))
  232. parser = reqparse.RequestParser()
  233. parser.add_argument(
  234. "indexing_technique", type=str, choices=Dataset.INDEXING_TECHNIQUE_LIST, nullable=False, location="json"
  235. )
  236. parser.add_argument("data_source", type=dict, required=False, location="json")
  237. parser.add_argument("process_rule", type=dict, required=False, location="json")
  238. parser.add_argument("duplicate", type=bool, default=True, nullable=False, location="json")
  239. parser.add_argument("original_document_id", type=str, required=False, location="json")
  240. parser.add_argument("doc_form", type=str, default="text_model", required=False, nullable=False, location="json")
  241. parser.add_argument("retrieval_model", type=dict, required=False, nullable=False, location="json")
  242. parser.add_argument("embedding_model", type=str, required=False, nullable=True, location="json")
  243. parser.add_argument("embedding_model_provider", type=str, required=False, nullable=True, location="json")
  244. parser.add_argument(
  245. "doc_language", type=str, default="English", required=False, nullable=False, location="json"
  246. )
  247. args = parser.parse_args()
  248. knowledge_config = KnowledgeConfig(**args)
  249. if not dataset.indexing_technique and not knowledge_config.indexing_technique:
  250. raise ValueError("indexing_technique is required.")
  251. # validate args
  252. DocumentService.document_create_args_validate(knowledge_config)
  253. try:
  254. documents, batch = DocumentService.save_document_with_dataset_id(dataset, knowledge_config, current_user)
  255. dataset = DatasetService.get_dataset(dataset_id)
  256. except ProviderTokenNotInitError as ex:
  257. raise ProviderNotInitializeError(ex.description)
  258. except QuotaExceededError:
  259. raise ProviderQuotaExceededError()
  260. except ModelCurrentlyNotSupportError:
  261. raise ProviderModelCurrentlyNotSupportError()
  262. return {"dataset": dataset, "documents": documents, "batch": batch}
  263. @setup_required
  264. @login_required
  265. @account_initialization_required
  266. @cloud_edition_billing_rate_limit_check("knowledge")
  267. def delete(self, dataset_id):
  268. dataset_id = str(dataset_id)
  269. dataset = DatasetService.get_dataset(dataset_id)
  270. if dataset is None:
  271. raise NotFound("Dataset not found.")
  272. # check user's model setting
  273. DatasetService.check_dataset_model_setting(dataset)
  274. try:
  275. document_ids = request.args.getlist("document_id")
  276. DocumentService.delete_documents(dataset, document_ids)
  277. except services.errors.document.DocumentIndexingError:
  278. raise DocumentIndexingError("Cannot delete document during indexing.")
  279. return {"result": "success"}, 204
  280. class DatasetInitApi(Resource):
  281. @setup_required
  282. @login_required
  283. @account_initialization_required
  284. @marshal_with(dataset_and_document_fields)
  285. @cloud_edition_billing_resource_check("vector_space")
  286. @cloud_edition_billing_rate_limit_check("knowledge")
  287. def post(self):
  288. # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor
  289. if not current_user.is_dataset_editor:
  290. raise Forbidden()
  291. parser = reqparse.RequestParser()
  292. parser.add_argument(
  293. "indexing_technique",
  294. type=str,
  295. choices=Dataset.INDEXING_TECHNIQUE_LIST,
  296. required=True,
  297. nullable=False,
  298. location="json",
  299. )
  300. parser.add_argument("data_source", type=dict, required=True, nullable=True, location="json")
  301. parser.add_argument("process_rule", type=dict, required=True, nullable=True, location="json")
  302. parser.add_argument("doc_form", type=str, default="text_model", required=False, nullable=False, location="json")
  303. parser.add_argument(
  304. "doc_language", type=str, default="English", required=False, nullable=False, location="json"
  305. )
  306. parser.add_argument("retrieval_model", type=dict, required=False, nullable=False, location="json")
  307. parser.add_argument("embedding_model", type=str, required=False, nullable=True, location="json")
  308. parser.add_argument("embedding_model_provider", type=str, required=False, nullable=True, location="json")
  309. args = parser.parse_args()
  310. knowledge_config = KnowledgeConfig(**args)
  311. if knowledge_config.indexing_technique == "high_quality":
  312. if knowledge_config.embedding_model is None or knowledge_config.embedding_model_provider is None:
  313. raise ValueError("embedding model and embedding model provider are required for high quality indexing.")
  314. try:
  315. model_manager = ModelManager()
  316. model_manager.get_model_instance(
  317. tenant_id=current_user.current_tenant_id,
  318. provider=args["embedding_model_provider"],
  319. model_type=ModelType.TEXT_EMBEDDING,
  320. model=args["embedding_model"],
  321. )
  322. except InvokeAuthorizationError:
  323. raise ProviderNotInitializeError(
  324. "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider."
  325. )
  326. except ProviderTokenNotInitError as ex:
  327. raise ProviderNotInitializeError(ex.description)
  328. # validate args
  329. DocumentService.document_create_args_validate(knowledge_config)
  330. try:
  331. dataset, documents, batch = DocumentService.save_document_without_dataset_id(
  332. tenant_id=current_user.current_tenant_id, knowledge_config=knowledge_config, account=current_user
  333. )
  334. except ProviderTokenNotInitError as ex:
  335. raise ProviderNotInitializeError(ex.description)
  336. except QuotaExceededError:
  337. raise ProviderQuotaExceededError()
  338. except ModelCurrentlyNotSupportError:
  339. raise ProviderModelCurrentlyNotSupportError()
  340. response = {"dataset": dataset, "documents": documents, "batch": batch}
  341. return response
  342. class DocumentIndexingEstimateApi(DocumentResource):
  343. @setup_required
  344. @login_required
  345. @account_initialization_required
  346. def get(self, dataset_id, document_id):
  347. dataset_id = str(dataset_id)
  348. document_id = str(document_id)
  349. document = self.get_document(dataset_id, document_id)
  350. if document.indexing_status in {"completed", "error"}:
  351. raise DocumentAlreadyFinishedError()
  352. data_process_rule = document.dataset_process_rule
  353. data_process_rule_dict = data_process_rule.to_dict()
  354. response = {"tokens": 0, "total_price": 0, "currency": "USD", "total_segments": 0, "preview": []}
  355. if document.data_source_type == "upload_file":
  356. data_source_info = document.data_source_info_dict
  357. if data_source_info and "upload_file_id" in data_source_info:
  358. file_id = data_source_info["upload_file_id"]
  359. file = (
  360. db.session.query(UploadFile)
  361. .where(UploadFile.tenant_id == document.tenant_id, UploadFile.id == file_id)
  362. .first()
  363. )
  364. # raise error if file not found
  365. if not file:
  366. raise NotFound("File not found.")
  367. extract_setting = ExtractSetting(
  368. datasource_type=DatasourceType.FILE.value, upload_file=file, document_model=document.doc_form
  369. )
  370. indexing_runner = IndexingRunner()
  371. try:
  372. estimate_response = indexing_runner.indexing_estimate(
  373. current_user.current_tenant_id,
  374. [extract_setting],
  375. data_process_rule_dict,
  376. document.doc_form,
  377. "English",
  378. dataset_id,
  379. )
  380. return estimate_response.model_dump(), 200
  381. except LLMBadRequestError:
  382. raise ProviderNotInitializeError(
  383. "No Embedding Model available. Please configure a valid provider "
  384. "in the Settings -> Model Provider."
  385. )
  386. except ProviderTokenNotInitError as ex:
  387. raise ProviderNotInitializeError(ex.description)
  388. except PluginDaemonClientSideError as ex:
  389. raise ProviderNotInitializeError(ex.description)
  390. except Exception as e:
  391. raise IndexingEstimateError(str(e))
  392. return response, 200
  393. class DocumentBatchIndexingEstimateApi(DocumentResource):
  394. @setup_required
  395. @login_required
  396. @account_initialization_required
  397. def get(self, dataset_id, batch):
  398. dataset_id = str(dataset_id)
  399. batch = str(batch)
  400. documents = self.get_batch_documents(dataset_id, batch)
  401. if not documents:
  402. return {"tokens": 0, "total_price": 0, "currency": "USD", "total_segments": 0, "preview": []}, 200
  403. data_process_rule = documents[0].dataset_process_rule
  404. data_process_rule_dict = data_process_rule.to_dict()
  405. extract_settings = []
  406. for document in documents:
  407. if document.indexing_status in {"completed", "error"}:
  408. raise DocumentAlreadyFinishedError()
  409. data_source_info = document.data_source_info_dict
  410. if document.data_source_type == "upload_file":
  411. file_id = data_source_info["upload_file_id"]
  412. file_detail = (
  413. db.session.query(UploadFile)
  414. .where(UploadFile.tenant_id == current_user.current_tenant_id, UploadFile.id == file_id)
  415. .first()
  416. )
  417. if file_detail is None:
  418. raise NotFound("File not found.")
  419. extract_setting = ExtractSetting(
  420. datasource_type=DatasourceType.FILE.value, upload_file=file_detail, document_model=document.doc_form
  421. )
  422. extract_settings.append(extract_setting)
  423. elif document.data_source_type == "notion_import":
  424. extract_setting = ExtractSetting(
  425. datasource_type=DatasourceType.NOTION.value,
  426. notion_info={
  427. "notion_workspace_id": data_source_info["notion_workspace_id"],
  428. "notion_obj_id": data_source_info["notion_page_id"],
  429. "notion_page_type": data_source_info["type"],
  430. "tenant_id": current_user.current_tenant_id,
  431. },
  432. document_model=document.doc_form,
  433. )
  434. extract_settings.append(extract_setting)
  435. elif document.data_source_type == "website_crawl":
  436. extract_setting = ExtractSetting(
  437. datasource_type=DatasourceType.WEBSITE.value,
  438. website_info={
  439. "provider": data_source_info["provider"],
  440. "job_id": data_source_info["job_id"],
  441. "url": data_source_info["url"],
  442. "tenant_id": current_user.current_tenant_id,
  443. "mode": data_source_info["mode"],
  444. "only_main_content": data_source_info["only_main_content"],
  445. },
  446. document_model=document.doc_form,
  447. )
  448. extract_settings.append(extract_setting)
  449. else:
  450. raise ValueError("Data source type not support")
  451. indexing_runner = IndexingRunner()
  452. try:
  453. response = indexing_runner.indexing_estimate(
  454. current_user.current_tenant_id,
  455. extract_settings,
  456. data_process_rule_dict,
  457. document.doc_form,
  458. "English",
  459. dataset_id,
  460. )
  461. return response.model_dump(), 200
  462. except LLMBadRequestError:
  463. raise ProviderNotInitializeError(
  464. "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider."
  465. )
  466. except ProviderTokenNotInitError as ex:
  467. raise ProviderNotInitializeError(ex.description)
  468. except PluginDaemonClientSideError as ex:
  469. raise ProviderNotInitializeError(ex.description)
  470. except Exception as e:
  471. raise IndexingEstimateError(str(e))
  472. class DocumentBatchIndexingStatusApi(DocumentResource):
  473. @setup_required
  474. @login_required
  475. @account_initialization_required
  476. def get(self, dataset_id, batch):
  477. dataset_id = str(dataset_id)
  478. batch = str(batch)
  479. documents = self.get_batch_documents(dataset_id, batch)
  480. documents_status = []
  481. for document in documents:
  482. completed_segments = (
  483. db.session.query(DocumentSegment)
  484. .where(
  485. DocumentSegment.completed_at.isnot(None),
  486. DocumentSegment.document_id == str(document.id),
  487. DocumentSegment.status != "re_segment",
  488. )
  489. .count()
  490. )
  491. total_segments = (
  492. db.session.query(DocumentSegment)
  493. .where(DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment")
  494. .count()
  495. )
  496. # Create a dictionary with document attributes and additional fields
  497. document_dict = {
  498. "id": document.id,
  499. "indexing_status": "paused" if document.is_paused else document.indexing_status,
  500. "processing_started_at": document.processing_started_at,
  501. "parsing_completed_at": document.parsing_completed_at,
  502. "cleaning_completed_at": document.cleaning_completed_at,
  503. "splitting_completed_at": document.splitting_completed_at,
  504. "completed_at": document.completed_at,
  505. "paused_at": document.paused_at,
  506. "error": document.error,
  507. "stopped_at": document.stopped_at,
  508. "completed_segments": completed_segments,
  509. "total_segments": total_segments,
  510. }
  511. documents_status.append(marshal(document_dict, document_status_fields))
  512. data = {"data": documents_status}
  513. return data
  514. class DocumentIndexingStatusApi(DocumentResource):
  515. @setup_required
  516. @login_required
  517. @account_initialization_required
  518. def get(self, dataset_id, document_id):
  519. dataset_id = str(dataset_id)
  520. document_id = str(document_id)
  521. document = self.get_document(dataset_id, document_id)
  522. completed_segments = (
  523. db.session.query(DocumentSegment)
  524. .where(
  525. DocumentSegment.completed_at.isnot(None),
  526. DocumentSegment.document_id == str(document_id),
  527. DocumentSegment.status != "re_segment",
  528. )
  529. .count()
  530. )
  531. total_segments = (
  532. db.session.query(DocumentSegment)
  533. .where(DocumentSegment.document_id == str(document_id), DocumentSegment.status != "re_segment")
  534. .count()
  535. )
  536. # Create a dictionary with document attributes and additional fields
  537. document_dict = {
  538. "id": document.id,
  539. "indexing_status": "paused" if document.is_paused else document.indexing_status,
  540. "processing_started_at": document.processing_started_at,
  541. "parsing_completed_at": document.parsing_completed_at,
  542. "cleaning_completed_at": document.cleaning_completed_at,
  543. "splitting_completed_at": document.splitting_completed_at,
  544. "completed_at": document.completed_at,
  545. "paused_at": document.paused_at,
  546. "error": document.error,
  547. "stopped_at": document.stopped_at,
  548. "completed_segments": completed_segments,
  549. "total_segments": total_segments,
  550. }
  551. return marshal(document_dict, document_status_fields)
  552. class DocumentApi(DocumentResource):
  553. METADATA_CHOICES = {"all", "only", "without"}
  554. @setup_required
  555. @login_required
  556. @account_initialization_required
  557. def get(self, dataset_id, document_id):
  558. dataset_id = str(dataset_id)
  559. document_id = str(document_id)
  560. document = self.get_document(dataset_id, document_id)
  561. metadata = request.args.get("metadata", "all")
  562. if metadata not in self.METADATA_CHOICES:
  563. raise InvalidMetadataError(f"Invalid metadata value: {metadata}")
  564. if metadata == "only":
  565. response = {"id": document.id, "doc_type": document.doc_type, "doc_metadata": document.doc_metadata_details}
  566. elif metadata == "without":
  567. dataset_process_rules = DatasetService.get_process_rules(dataset_id)
  568. document_process_rules = document.dataset_process_rule.to_dict()
  569. data_source_info = document.data_source_detail_dict
  570. response = {
  571. "id": document.id,
  572. "position": document.position,
  573. "data_source_type": document.data_source_type,
  574. "data_source_info": data_source_info,
  575. "dataset_process_rule_id": document.dataset_process_rule_id,
  576. "dataset_process_rule": dataset_process_rules,
  577. "document_process_rule": document_process_rules,
  578. "name": document.name,
  579. "created_from": document.created_from,
  580. "created_by": document.created_by,
  581. "created_at": document.created_at.timestamp(),
  582. "tokens": document.tokens,
  583. "indexing_status": document.indexing_status,
  584. "completed_at": int(document.completed_at.timestamp()) if document.completed_at else None,
  585. "updated_at": int(document.updated_at.timestamp()) if document.updated_at else None,
  586. "indexing_latency": document.indexing_latency,
  587. "error": document.error,
  588. "enabled": document.enabled,
  589. "disabled_at": int(document.disabled_at.timestamp()) if document.disabled_at else None,
  590. "disabled_by": document.disabled_by,
  591. "archived": document.archived,
  592. "segment_count": document.segment_count,
  593. "average_segment_length": document.average_segment_length,
  594. "hit_count": document.hit_count,
  595. "display_status": document.display_status,
  596. "doc_form": document.doc_form,
  597. "doc_language": document.doc_language,
  598. }
  599. else:
  600. dataset_process_rules = DatasetService.get_process_rules(dataset_id)
  601. document_process_rules = document.dataset_process_rule.to_dict()
  602. data_source_info = document.data_source_detail_dict
  603. response = {
  604. "id": document.id,
  605. "position": document.position,
  606. "data_source_type": document.data_source_type,
  607. "data_source_info": data_source_info,
  608. "dataset_process_rule_id": document.dataset_process_rule_id,
  609. "dataset_process_rule": dataset_process_rules,
  610. "document_process_rule": document_process_rules,
  611. "name": document.name,
  612. "created_from": document.created_from,
  613. "created_by": document.created_by,
  614. "created_at": document.created_at.timestamp(),
  615. "tokens": document.tokens,
  616. "indexing_status": document.indexing_status,
  617. "completed_at": int(document.completed_at.timestamp()) if document.completed_at else None,
  618. "updated_at": int(document.updated_at.timestamp()) if document.updated_at else None,
  619. "indexing_latency": document.indexing_latency,
  620. "error": document.error,
  621. "enabled": document.enabled,
  622. "disabled_at": int(document.disabled_at.timestamp()) if document.disabled_at else None,
  623. "disabled_by": document.disabled_by,
  624. "archived": document.archived,
  625. "doc_type": document.doc_type,
  626. "doc_metadata": document.doc_metadata_details,
  627. "segment_count": document.segment_count,
  628. "average_segment_length": document.average_segment_length,
  629. "hit_count": document.hit_count,
  630. "display_status": document.display_status,
  631. "doc_form": document.doc_form,
  632. "doc_language": document.doc_language,
  633. }
  634. return response, 200
  635. @setup_required
  636. @login_required
  637. @account_initialization_required
  638. @cloud_edition_billing_rate_limit_check("knowledge")
  639. def delete(self, dataset_id, document_id):
  640. dataset_id = str(dataset_id)
  641. document_id = str(document_id)
  642. dataset = DatasetService.get_dataset(dataset_id)
  643. if dataset is None:
  644. raise NotFound("Dataset not found.")
  645. # check user's model setting
  646. DatasetService.check_dataset_model_setting(dataset)
  647. document = self.get_document(dataset_id, document_id)
  648. try:
  649. DocumentService.delete_document(document)
  650. except services.errors.document.DocumentIndexingError:
  651. raise DocumentIndexingError("Cannot delete document during indexing.")
  652. return {"result": "success"}, 204
  653. class DocumentProcessingApi(DocumentResource):
  654. @setup_required
  655. @login_required
  656. @account_initialization_required
  657. @cloud_edition_billing_rate_limit_check("knowledge")
  658. def patch(self, dataset_id, document_id, action: Literal["pause", "resume"]):
  659. dataset_id = str(dataset_id)
  660. document_id = str(document_id)
  661. document = self.get_document(dataset_id, document_id)
  662. # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor
  663. if not current_user.is_dataset_editor:
  664. raise Forbidden()
  665. if action == "pause":
  666. if document.indexing_status != "indexing":
  667. raise InvalidActionError("Document not in indexing state.")
  668. document.paused_by = current_user.id
  669. document.paused_at = naive_utc_now()
  670. document.is_paused = True
  671. db.session.commit()
  672. elif action == "resume":
  673. if document.indexing_status not in {"paused", "error"}:
  674. raise InvalidActionError("Document not in paused or error state.")
  675. document.paused_by = None
  676. document.paused_at = None
  677. document.is_paused = False
  678. db.session.commit()
  679. return {"result": "success"}, 200
  680. class DocumentMetadataApi(DocumentResource):
  681. @setup_required
  682. @login_required
  683. @account_initialization_required
  684. def put(self, dataset_id, document_id):
  685. dataset_id = str(dataset_id)
  686. document_id = str(document_id)
  687. document = self.get_document(dataset_id, document_id)
  688. req_data = request.get_json()
  689. doc_type = req_data.get("doc_type")
  690. doc_metadata = req_data.get("doc_metadata")
  691. # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor
  692. if not current_user.is_dataset_editor:
  693. raise Forbidden()
  694. if doc_type is None or doc_metadata is None:
  695. raise ValueError("Both doc_type and doc_metadata must be provided.")
  696. if doc_type not in DocumentService.DOCUMENT_METADATA_SCHEMA:
  697. raise ValueError("Invalid doc_type.")
  698. if not isinstance(doc_metadata, dict):
  699. raise ValueError("doc_metadata must be a dictionary.")
  700. metadata_schema: dict = cast(dict, DocumentService.DOCUMENT_METADATA_SCHEMA[doc_type])
  701. document.doc_metadata = {}
  702. if doc_type == "others":
  703. document.doc_metadata = doc_metadata
  704. else:
  705. for key, value_type in metadata_schema.items():
  706. value = doc_metadata.get(key)
  707. if value is not None and isinstance(value, value_type):
  708. document.doc_metadata[key] = value
  709. document.doc_type = doc_type
  710. document.updated_at = naive_utc_now()
  711. db.session.commit()
  712. return {"result": "success", "message": "Document metadata updated."}, 200
  713. class DocumentStatusApi(DocumentResource):
  714. @setup_required
  715. @login_required
  716. @account_initialization_required
  717. @cloud_edition_billing_resource_check("vector_space")
  718. @cloud_edition_billing_rate_limit_check("knowledge")
  719. def patch(self, dataset_id, action: Literal["enable", "disable", "archive", "un_archive"]):
  720. dataset_id = str(dataset_id)
  721. dataset = DatasetService.get_dataset(dataset_id)
  722. if dataset is None:
  723. raise NotFound("Dataset not found.")
  724. # The role of the current user in the ta table must be admin, owner, or editor
  725. if not current_user.is_dataset_editor:
  726. raise Forbidden()
  727. # check user's model setting
  728. DatasetService.check_dataset_model_setting(dataset)
  729. # check user's permission
  730. DatasetService.check_dataset_permission(dataset, current_user)
  731. document_ids = request.args.getlist("document_id")
  732. try:
  733. DocumentService.batch_update_document_status(dataset, document_ids, action, current_user)
  734. except services.errors.document.DocumentIndexingError as e:
  735. raise InvalidActionError(str(e))
  736. except ValueError as e:
  737. raise InvalidActionError(str(e))
  738. except NotFound as e:
  739. raise NotFound(str(e))
  740. return {"result": "success"}, 200
  741. class DocumentPauseApi(DocumentResource):
  742. @setup_required
  743. @login_required
  744. @account_initialization_required
  745. @cloud_edition_billing_rate_limit_check("knowledge")
  746. def patch(self, dataset_id, document_id):
  747. """pause document."""
  748. dataset_id = str(dataset_id)
  749. document_id = str(document_id)
  750. dataset = DatasetService.get_dataset(dataset_id)
  751. if not dataset:
  752. raise NotFound("Dataset not found.")
  753. document = DocumentService.get_document(dataset.id, document_id)
  754. # 404 if document not found
  755. if document is None:
  756. raise NotFound("Document Not Exists.")
  757. # 403 if document is archived
  758. if DocumentService.check_archived(document):
  759. raise ArchivedDocumentImmutableError()
  760. try:
  761. # pause document
  762. DocumentService.pause_document(document)
  763. except services.errors.document.DocumentIndexingError:
  764. raise DocumentIndexingError("Cannot pause completed document.")
  765. return {"result": "success"}, 204
  766. class DocumentRecoverApi(DocumentResource):
  767. @setup_required
  768. @login_required
  769. @account_initialization_required
  770. @cloud_edition_billing_rate_limit_check("knowledge")
  771. def patch(self, dataset_id, document_id):
  772. """recover document."""
  773. dataset_id = str(dataset_id)
  774. document_id = str(document_id)
  775. dataset = DatasetService.get_dataset(dataset_id)
  776. if not dataset:
  777. raise NotFound("Dataset not found.")
  778. document = DocumentService.get_document(dataset.id, document_id)
  779. # 404 if document not found
  780. if document is None:
  781. raise NotFound("Document Not Exists.")
  782. # 403 if document is archived
  783. if DocumentService.check_archived(document):
  784. raise ArchivedDocumentImmutableError()
  785. try:
  786. # pause document
  787. DocumentService.recover_document(document)
  788. except services.errors.document.DocumentIndexingError:
  789. raise DocumentIndexingError("Document is not in paused status.")
  790. return {"result": "success"}, 204
  791. class DocumentRetryApi(DocumentResource):
  792. @setup_required
  793. @login_required
  794. @account_initialization_required
  795. @cloud_edition_billing_rate_limit_check("knowledge")
  796. def post(self, dataset_id):
  797. """retry document."""
  798. parser = reqparse.RequestParser()
  799. parser.add_argument("document_ids", type=list, required=True, nullable=False, location="json")
  800. args = parser.parse_args()
  801. dataset_id = str(dataset_id)
  802. dataset = DatasetService.get_dataset(dataset_id)
  803. retry_documents = []
  804. if not dataset:
  805. raise NotFound("Dataset not found.")
  806. for document_id in args["document_ids"]:
  807. try:
  808. document_id = str(document_id)
  809. document = DocumentService.get_document(dataset.id, document_id)
  810. # 404 if document not found
  811. if document is None:
  812. raise NotFound("Document Not Exists.")
  813. # 403 if document is archived
  814. if DocumentService.check_archived(document):
  815. raise ArchivedDocumentImmutableError()
  816. # 400 if document is completed
  817. if document.indexing_status == "completed":
  818. raise DocumentAlreadyFinishedError()
  819. retry_documents.append(document)
  820. except Exception:
  821. logger.exception("Failed to retry document, document id: %s", document_id)
  822. continue
  823. # retry document
  824. DocumentService.retry_document(dataset_id, retry_documents)
  825. return {"result": "success"}, 204
  826. class DocumentRenameApi(DocumentResource):
  827. @setup_required
  828. @login_required
  829. @account_initialization_required
  830. @marshal_with(document_fields)
  831. def post(self, dataset_id, document_id):
  832. # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator
  833. if not current_user.is_dataset_editor:
  834. raise Forbidden()
  835. dataset = DatasetService.get_dataset(dataset_id)
  836. DatasetService.check_dataset_operator_permission(current_user, dataset)
  837. parser = reqparse.RequestParser()
  838. parser.add_argument("name", type=str, required=True, nullable=False, location="json")
  839. args = parser.parse_args()
  840. try:
  841. document = DocumentService.rename_document(dataset_id, document_id, args["name"])
  842. except services.errors.document.DocumentIndexingError:
  843. raise DocumentIndexingError("Cannot delete document during indexing.")
  844. return document
  845. class WebsiteDocumentSyncApi(DocumentResource):
  846. @setup_required
  847. @login_required
  848. @account_initialization_required
  849. def get(self, dataset_id, document_id):
  850. """sync website document."""
  851. dataset_id = str(dataset_id)
  852. dataset = DatasetService.get_dataset(dataset_id)
  853. if not dataset:
  854. raise NotFound("Dataset not found.")
  855. document_id = str(document_id)
  856. document = DocumentService.get_document(dataset.id, document_id)
  857. if not document:
  858. raise NotFound("Document not found.")
  859. if document.tenant_id != current_user.current_tenant_id:
  860. raise Forbidden("No permission.")
  861. if document.data_source_type != "website_crawl":
  862. raise ValueError("Document is not a website document.")
  863. # 403 if document is archived
  864. if DocumentService.check_archived(document):
  865. raise ArchivedDocumentImmutableError()
  866. # sync document
  867. DocumentService.sync_website_document(dataset_id, document)
  868. return {"result": "success"}, 200
  869. class DocDownloadFileApi(DocumentResource):
  870. @setup_required
  871. @login_required
  872. @account_initialization_required
  873. def get(self, dataset_id, document_id):
  874. """Get upload file download URL for a document.
  875. Returns download information and signed URL for the uploaded file
  876. associated with this document.
  877. """
  878. dataset_id = str(dataset_id)
  879. document_id = str(document_id)
  880. document = self.get_document(dataset_id, document_id)
  881. # Check download permission
  882. has_permission, error_msg = DocumentDownloadService.check_download_permission(
  883. tenant_id=document.tenant_id,
  884. document_id=document_id,
  885. user_id=current_user.id,
  886. )
  887. if not has_permission:
  888. raise Forbidden(error_msg)
  889. # Check if document is based on uploaded file
  890. if document.data_source_type != "upload_file":
  891. raise ValueError(f"Document data source type ({document.data_source_type}) is not upload_file.")
  892. data_source_info = document.data_source_info_dict
  893. if not data_source_info or "upload_file_id" not in data_source_info:
  894. raise ValueError("Upload file id not found in document data source info.")
  895. file_id = data_source_info["upload_file_id"]
  896. upload_file = (
  897. db.session.query(UploadFile)
  898. .where(UploadFile.tenant_id == document.tenant_id, UploadFile.id == file_id)
  899. .first()
  900. )
  901. if not upload_file:
  902. raise NotFound("Upload file not found.")
  903. # Generate signed URL
  904. url = file_helpers.get_signed_file_url(upload_file_id=upload_file.id)
  905. download_url = f"{url}&as_attachment=true"
  906. # Record download
  907. DocumentDownloadService.record_download(
  908. tenant_id=document.tenant_id,
  909. document_id=document_id,
  910. user_id=current_user.id,
  911. user_name=current_user.name,
  912. ip=request.remote_addr or "unknown",
  913. )
  914. return {
  915. "id": upload_file.id,
  916. "name": upload_file.name,
  917. "size": upload_file.size,
  918. "extension": upload_file.extension,
  919. "url": url,
  920. "download_url": download_url,
  921. "mime_type": upload_file.mime_type,
  922. "created_by": upload_file.created_by,
  923. "created_at": upload_file.created_at.timestamp(),
  924. }, 200
  925. class DocumentBatchDownloadApi(DocumentResource):
  926. @setup_required
  927. @login_required
  928. @account_initialization_required
  929. def post(self, dataset_id):
  930. """Batch get upload file download URLs for multiple documents.
  931. Returns download information for multiple documents.
  932. """
  933. dataset_id = str(dataset_id)
  934. # Verify dataset access
  935. dataset = DatasetService.get_dataset(dataset_id)
  936. if not dataset:
  937. return {"error": "Dataset not found"}, 404
  938. try:
  939. DatasetService.check_dataset_permission(dataset, current_user)
  940. except services.errors.account.NoPermissionError as e:
  941. raise Forbidden(str(e))
  942. parser = reqparse.RequestParser()
  943. parser.add_argument("document_ids", type=list, location="json", required=True)
  944. args = parser.parse_args()
  945. document_ids = args["document_ids"]
  946. if not document_ids or not isinstance(document_ids, list):
  947. return {"error": "document_ids must be a non-empty list"}, 400
  948. downloads = []
  949. for document_id in document_ids:
  950. try:
  951. document_id = str(document_id)
  952. document = self.get_document(dataset_id, document_id)
  953. # Check download permission
  954. has_permission, error_msg = DocumentDownloadService.check_download_permission(
  955. tenant_id=document.tenant_id,
  956. document_id=document_id,
  957. user_id=current_user.id,
  958. )
  959. if not has_permission:
  960. downloads.append(
  961. {
  962. "document_id": document_id,
  963. "document_name": document.name,
  964. "download_url": None,
  965. "status": "failed",
  966. "error": error_msg or "Permission denied",
  967. }
  968. )
  969. continue
  970. # Check if document is based on uploaded file
  971. if document.data_source_type != "upload_file":
  972. downloads.append(
  973. {
  974. "document_id": document_id,
  975. "document_name": document.name,
  976. "download_url": None,
  977. "status": "failed",
  978. "error": "Not a file upload",
  979. }
  980. )
  981. continue
  982. data_source_info = document.data_source_info_dict
  983. if not data_source_info or "upload_file_id" not in data_source_info:
  984. downloads.append(
  985. {
  986. "document_id": document_id,
  987. "document_name": document.name,
  988. "download_url": None,
  989. "status": "failed",
  990. "error": "Upload file id not found",
  991. }
  992. )
  993. continue
  994. file_id = data_source_info["upload_file_id"]
  995. upload_file = (
  996. db.session.query(UploadFile)
  997. .where(UploadFile.tenant_id == dataset.tenant_id, UploadFile.id == file_id)
  998. .first()
  999. )
  1000. if not upload_file:
  1001. downloads.append(
  1002. {
  1003. "document_id": document_id,
  1004. "document_name": document.name,
  1005. "download_url": None,
  1006. "status": "failed",
  1007. "error": "Upload file not found",
  1008. }
  1009. )
  1010. continue
  1011. # Generate signed URL
  1012. url = file_helpers.get_signed_file_url(upload_file_id=upload_file.id)
  1013. download_url = f"{url}&as_attachment=true"
  1014. # Record download
  1015. DocumentDownloadService.record_download(
  1016. tenant_id=document.tenant_id,
  1017. document_id=document_id,
  1018. user_id=current_user.id,
  1019. user_name=current_user.name,
  1020. ip=request.remote_addr or "unknown",
  1021. )
  1022. downloads.append(
  1023. {
  1024. "document_id": document_id,
  1025. "document_name": upload_file.name,
  1026. "download_url": download_url,
  1027. "status": "success",
  1028. "error": None,
  1029. }
  1030. )
  1031. except NotFound:
  1032. downloads.append(
  1033. {
  1034. "document_id": document_id,
  1035. "document_name": "",
  1036. "download_url": None,
  1037. "status": "failed",
  1038. "error": "Document not found",
  1039. }
  1040. )
  1041. except Exception as e:
  1042. downloads.append(
  1043. {
  1044. "document_id": document_id,
  1045. "document_name": "",
  1046. "download_url": None,
  1047. "status": "failed",
  1048. "error": str(e),
  1049. }
  1050. )
  1051. return {"downloads": downloads}, 200
  1052. class DocumentFileProxyUrlApi(Resource):
  1053. """API to get signed URL for document file proxy (for KKFileView)"""
  1054. @setup_required
  1055. @login_required
  1056. @account_initialization_required
  1057. def get(self, dataset_id, document_id):
  1058. """Get signed URL for document file proxy
  1059. Args:
  1060. dataset_id: Dataset UUID
  1061. document_id: Document UUID
  1062. Returns:
  1063. JSON with signed URL for file proxy
  1064. """
  1065. from core.file import helpers as file_helpers
  1066. dataset_id_str = str(dataset_id)
  1067. document_id_str = str(document_id)
  1068. dataset = DatasetService.get_dataset(dataset_id_str)
  1069. if not dataset:
  1070. raise NotFound("Dataset not found")
  1071. # Check dataset permission
  1072. DatasetService.check_dataset_permission(dataset, current_user)
  1073. document = (
  1074. db.session.query(Document)
  1075. .where(Document.dataset_id == dataset_id_str, Document.id == document_id_str)
  1076. .first()
  1077. )
  1078. if not document:
  1079. raise NotFound("Document not found")
  1080. # Only support file upload type
  1081. if document.data_source_type != "upload_file":
  1082. raise ValueError("Not a file upload")
  1083. data_source_info = document.data_source_info_dict
  1084. if not data_source_info or "upload_file_id" not in data_source_info:
  1085. raise ValueError("Upload file id not found")
  1086. file_id = data_source_info["upload_file_id"]
  1087. upload_file = (
  1088. db.session.query(UploadFile)
  1089. .where(UploadFile.tenant_id == dataset.tenant_id, UploadFile.id == file_id)
  1090. .first()
  1091. )
  1092. if not upload_file:
  1093. raise NotFound("Upload file not found")
  1094. # Generate secure token
  1095. from urllib.parse import quote
  1096. from configs import dify_config
  1097. token = file_helpers.generate_document_proxy_token(
  1098. dataset_id_str, document_id_str, upload_file.name
  1099. )
  1100. # Build URL in format: /console/api/datasets/file-proxy/{token}/{filename}
  1101. # This format is similar to kkFileView's expected URL format
  1102. # For kkFileView (which may be on a different machine in the LAN):
  1103. # 1. Try INTERNAL_FILES_URL first (for LAN access, e.g., http://192.168.1.100:5001)
  1104. # 2. Fall back to CONSOLE_API_URL (for Docker internal network)
  1105. # 3. Finally use request URL (for development)
  1106. base_url = dify_config.INTERNAL_FILES_URL
  1107. if not base_url:
  1108. base_url = dify_config.CONSOLE_API_URL
  1109. if not base_url:
  1110. # Build base URL from request
  1111. base_url = f"{request.scheme}://{request.host}"
  1112. encoded_filename = quote(upload_file.name)
  1113. signed_url = f"{base_url}/console/api/datasets/file-proxy/{token}/{encoded_filename}"
  1114. return {"url": signed_url}, 200
  1115. class DocumentFileProxyApi(Resource):
  1116. """Proxy endpoint for kkFileView to access document files with token-based auth
  1117. URL format: /console/api/datasets/file-proxy/{token}/{filename}
  1118. This format is similar to kkFileView's expected URL format, e.g.:
  1119. http://127.0.0.1:8080/file/test.txt
  1120. """
  1121. def get(self, token, filename):
  1122. """Get document file stream for kkFileView preview (no login required)
  1123. Args:
  1124. token: Base64-encoded token containing dataset_id, document_id, filename, and signature
  1125. filename: Filename from URL path (for kkFileView to detect file type)
  1126. """
  1127. from urllib.parse import unquote
  1128. from core.file import helpers as file_helpers
  1129. # Decode URL-encoded filename
  1130. decoded_filename = unquote(filename)
  1131. # Verify and decode token
  1132. verified_params = file_helpers.verify_document_proxy_token(token)
  1133. if not verified_params:
  1134. raise Forbidden("Invalid or expired token")
  1135. # Extract verified parameters
  1136. dataset_id_str = verified_params["dataset_id"]
  1137. document_id_str = verified_params["document_id"]
  1138. token_filename = verified_params["filename"]
  1139. # Verify filename matches (security check)
  1140. if token_filename != decoded_filename:
  1141. raise Forbidden("Filename mismatch")
  1142. # Get the file (no authentication required - token is the auth)
  1143. dataset = DatasetService.get_dataset(dataset_id_str)
  1144. if not dataset:
  1145. raise NotFound("Dataset not found")
  1146. document = (
  1147. db.session.query(Document)
  1148. .where(Document.dataset_id == dataset_id_str, Document.id == document_id_str)
  1149. .first()
  1150. )
  1151. if not document:
  1152. raise NotFound("Document not found")
  1153. # Only support file upload type
  1154. if document.data_source_type != "upload_file":
  1155. raise ValueError("Not a file upload")
  1156. data_source_info = document.data_source_info_dict
  1157. if not data_source_info or "upload_file_id" not in data_source_info:
  1158. raise ValueError("Upload file id not found")
  1159. file_id = data_source_info["upload_file_id"]
  1160. upload_file = (
  1161. db.session.query(UploadFile)
  1162. .where(UploadFile.tenant_id == dataset.tenant_id, UploadFile.id == file_id)
  1163. .first()
  1164. )
  1165. if not upload_file:
  1166. raise NotFound("Upload file not found")
  1167. # Ensure filename is properly encoded for Content-Disposition header
  1168. # Use percent-encoding for the filename* parameter (RFC 5987)
  1169. encoded_filename = quote(decoded_filename)
  1170. # Stream file from storage
  1171. try:
  1172. generator = storage.load(upload_file.key, stream=True)
  1173. return Response(
  1174. generator,
  1175. headers={
  1176. "Content-Type": upload_file.mime_type or "application/octet-stream",
  1177. "Content-Disposition": f"inline; filename*=UTF-8''{encoded_filename}",
  1178. },
  1179. )
  1180. except Exception as e:
  1181. raise ValueError(f"Failed to load file: {str(e)}")
  1182. class DocumentDownloadPermissionApi(DocumentResource):
  1183. """API for managing document download permissions"""
  1184. @setup_required
  1185. @login_required
  1186. @account_initialization_required
  1187. def get(self, dataset_id, document_id):
  1188. """Get download permission configuration for a document.
  1189. Returns:
  1190. - is_restricted: Whether download is restricted to specific users
  1191. - users: List of all tenant users with permission status
  1192. """
  1193. dataset_id = str(dataset_id)
  1194. document_id = str(document_id)
  1195. document = self.get_document(dataset_id, document_id)
  1196. # Get allowed user IDs
  1197. allowed_user_ids = DocumentDownloadService.get_allowed_users(
  1198. tenant_id=document.tenant_id,
  1199. document_id=document_id,
  1200. )
  1201. # Get all tenant users
  1202. from models.account import Account, TenantAccountJoin
  1203. tenant_users = (
  1204. db.session.query(Account)
  1205. .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id)
  1206. .filter(TenantAccountJoin.tenant_id == document.tenant_id)
  1207. .all()
  1208. )
  1209. # Build user list with permission status
  1210. users = []
  1211. for user in tenant_users:
  1212. users.append(
  1213. {
  1214. "id": user.id,
  1215. "name": user.name,
  1216. "email": user.email,
  1217. "has_permission": user.id in allowed_user_ids if allowed_user_ids else True,
  1218. }
  1219. )
  1220. return {
  1221. "is_restricted": bool(allowed_user_ids),
  1222. "users": users,
  1223. }, 200
  1224. @setup_required
  1225. @login_required
  1226. @account_initialization_required
  1227. def post(self, dataset_id, document_id):
  1228. """Set download permission for a document.
  1229. Request body:
  1230. - is_restricted: Whether to restrict downloads
  1231. - user_ids: List of user IDs allowed to download (empty if not restricted)
  1232. """
  1233. dataset_id = str(dataset_id)
  1234. document_id = str(document_id)
  1235. document = self.get_document(dataset_id, document_id)
  1236. parser = reqparse.RequestParser()
  1237. parser.add_argument("user_ids", type=list, location="json", required=False, default=[])
  1238. parser.add_argument("is_restricted", type=bool, location="json", required=False, default=False)
  1239. args = parser.parse_args()
  1240. user_ids = args["user_ids"]
  1241. is_restricted = args["is_restricted"]
  1242. # If not restricted, pass empty list
  1243. if not is_restricted:
  1244. user_ids = []
  1245. # Set permission
  1246. DocumentDownloadService.set_allowed_users(
  1247. tenant_id=document.tenant_id,
  1248. document_id=document_id,
  1249. user_ids=user_ids,
  1250. )
  1251. return {"result": "success"}, 200
  1252. class DocumentDownloadLogsApi(DocumentResource):
  1253. """API for retrieving document download logs"""
  1254. @setup_required
  1255. @login_required
  1256. @account_initialization_required
  1257. def get(self, dataset_id, document_id):
  1258. """Get download logs for a document.
  1259. Query params:
  1260. - limit: Maximum number of records to return (default: 100)
  1261. Returns:
  1262. - logs: List of download records
  1263. - total: Total number of downloads
  1264. - unique_users: Number of unique users who downloaded
  1265. """
  1266. dataset_id = str(dataset_id)
  1267. document_id = str(document_id)
  1268. document = self.get_document(dataset_id, document_id)
  1269. parser = reqparse.RequestParser()
  1270. parser.add_argument("limit", type=int, location="args", default=100)
  1271. args = parser.parse_args()
  1272. # Get download logs
  1273. logs = DocumentDownloadService.get_download_logs(
  1274. tenant_id=document.tenant_id,
  1275. document_id=document_id,
  1276. limit=args["limit"],
  1277. )
  1278. # Calculate statistics
  1279. total_downloads = len(logs)
  1280. unique_users = len(set(log["user_id"] for log in logs))
  1281. return {
  1282. "logs": logs,
  1283. "total": total_downloads,
  1284. "unique_users": unique_users,
  1285. }, 200
  1286. api.add_resource(GetProcessRuleApi, "/datasets/process-rule")
  1287. api.add_resource(DatasetDocumentListApi, "/datasets/<uuid:dataset_id>/documents")
  1288. api.add_resource(DatasetInitApi, "/datasets/init")
  1289. api.add_resource(
  1290. DocumentIndexingEstimateApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/indexing-estimate"
  1291. )
  1292. api.add_resource(DocumentBatchIndexingEstimateApi, "/datasets/<uuid:dataset_id>/batch/<string:batch>/indexing-estimate")
  1293. api.add_resource(DocumentBatchIndexingStatusApi, "/datasets/<uuid:dataset_id>/batch/<string:batch>/indexing-status")
  1294. api.add_resource(DocumentIndexingStatusApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/indexing-status")
  1295. api.add_resource(DocumentApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>")
  1296. api.add_resource(
  1297. DocumentProcessingApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/<string:action>"
  1298. )
  1299. api.add_resource(DocumentMetadataApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/metadata")
  1300. api.add_resource(DocumentStatusApi, "/datasets/<uuid:dataset_id>/documents/status/<string:action>/batch")
  1301. api.add_resource(DocumentPauseApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/pause")
  1302. api.add_resource(DocumentRecoverApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/resume")
  1303. api.add_resource(DocumentRetryApi, "/datasets/<uuid:dataset_id>/retry")
  1304. api.add_resource(DocumentRenameApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/rename")
  1305. api.add_resource(DocDownloadFileApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/download")
  1306. api.add_resource(DocumentBatchDownloadApi, "/datasets/<uuid:dataset_id>/documents/batch-download")
  1307. api.add_resource(
  1308. DocumentDownloadPermissionApi,
  1309. "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/download-permission",
  1310. )
  1311. api.add_resource(
  1312. DocumentDownloadLogsApi,
  1313. "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/download-logs",
  1314. )
  1315. api.add_resource(
  1316. DocumentFileProxyUrlApi,
  1317. "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/file-proxy-url",
  1318. )
  1319. api.add_resource(
  1320. DocumentFileProxyApi,
  1321. "/datasets/file-proxy/<string:token>/<path:filename>",
  1322. )
  1323. api.add_resource(WebsiteDocumentSyncApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/website-sync")