Du kan inte välja fler än 25 ämnen Ämnen måste starta med en bokstav eller siffra, kan innehålla bindestreck ('-') och vara max 35 tecken långa.

datasets.py 31KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818
  1. import flask_restful
  2. from flask import request
  3. from flask_login import current_user
  4. from flask_restful import Resource, marshal, marshal_with, reqparse
  5. from werkzeug.exceptions import Forbidden, NotFound
  6. import services
  7. from configs import dify_config
  8. from controllers.console import api
  9. from controllers.console.apikey import api_key_fields, api_key_list
  10. from controllers.console.app.error import ProviderNotInitializeError
  11. from controllers.console.datasets.error import DatasetInUseError, DatasetNameDuplicateError, IndexingEstimateError
  12. from controllers.console.wraps import (
  13. account_initialization_required,
  14. cloud_edition_billing_rate_limit_check,
  15. enterprise_license_required,
  16. setup_required,
  17. )
  18. from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
  19. from core.indexing_runner import IndexingRunner
  20. from core.model_runtime.entities.model_entities import ModelType
  21. from core.plugin.entities.plugin import ModelProviderID
  22. from core.provider_manager import ProviderManager
  23. from core.rag.datasource.vdb.vector_type import VectorType
  24. from core.rag.extractor.entity.extract_setting import ExtractSetting
  25. from core.rag.retrieval.retrieval_methods import RetrievalMethod
  26. from extensions.ext_database import db
  27. from fields.app_fields import related_app_list
  28. from fields.dataset_fields import dataset_detail_fields, dataset_query_detail_fields
  29. from fields.document_fields import document_status_fields
  30. from libs.login import login_required
  31. from models import ApiToken, Dataset, Document, DocumentSegment, UploadFile
  32. from models.dataset import DatasetPermissionEnum
  33. from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService
  34. def _validate_name(name):
  35. if not name or len(name) < 1 or len(name) > 40:
  36. raise ValueError("Name must be between 1 to 40 characters.")
  37. return name
  38. def _validate_description_length(description):
  39. if len(description) > 400:
  40. raise ValueError("Description cannot exceed 400 characters.")
  41. return description
  42. class DatasetListApi(Resource):
  43. @setup_required
  44. @login_required
  45. @account_initialization_required
  46. @enterprise_license_required
  47. def get(self):
  48. page = request.args.get("page", default=1, type=int)
  49. limit = request.args.get("limit", default=20, type=int)
  50. ids = request.args.getlist("ids")
  51. # provider = request.args.get("provider", default="vendor")
  52. search = request.args.get("keyword", default=None, type=str)
  53. tag_ids = request.args.getlist("tag_ids")
  54. include_all = request.args.get("include_all", default="false").lower() == "true"
  55. if ids:
  56. datasets, total = DatasetService.get_datasets_by_ids(ids, current_user.current_tenant_id)
  57. else:
  58. datasets, total = DatasetService.get_datasets(
  59. page, limit, current_user.current_tenant_id, current_user, search, tag_ids, include_all
  60. )
  61. # check embedding setting
  62. provider_manager = ProviderManager()
  63. configurations = provider_manager.get_configurations(tenant_id=current_user.current_tenant_id)
  64. embedding_models = configurations.get_models(model_type=ModelType.TEXT_EMBEDDING, only_active=True)
  65. model_names = []
  66. for embedding_model in embedding_models:
  67. model_names.append(f"{embedding_model.model}:{embedding_model.provider.provider}")
  68. data = marshal(datasets, dataset_detail_fields)
  69. for item in data:
  70. # convert embedding_model_provider to plugin standard format
  71. if item["indexing_technique"] == "high_quality" and item["embedding_model_provider"]:
  72. item["embedding_model_provider"] = str(ModelProviderID(item["embedding_model_provider"]))
  73. item_model = f"{item['embedding_model']}:{item['embedding_model_provider']}"
  74. if item_model in model_names:
  75. item["embedding_available"] = True
  76. else:
  77. item["embedding_available"] = False
  78. else:
  79. item["embedding_available"] = True
  80. if item.get("permission") == "partial_members":
  81. part_users_list = DatasetPermissionService.get_dataset_partial_member_list(item["id"])
  82. item.update({"partial_member_list": part_users_list})
  83. else:
  84. item.update({"partial_member_list": []})
  85. response = {"data": data, "has_more": len(datasets) == limit, "limit": limit, "total": total, "page": page}
  86. return response, 200
  87. @setup_required
  88. @login_required
  89. @account_initialization_required
  90. @cloud_edition_billing_rate_limit_check("knowledge")
  91. def post(self):
  92. parser = reqparse.RequestParser()
  93. parser.add_argument(
  94. "name",
  95. nullable=False,
  96. required=True,
  97. help="type is required. Name must be between 1 to 40 characters.",
  98. type=_validate_name,
  99. )
  100. parser.add_argument(
  101. "description",
  102. type=str,
  103. nullable=True,
  104. required=False,
  105. default="",
  106. )
  107. parser.add_argument(
  108. "indexing_technique",
  109. type=str,
  110. location="json",
  111. choices=Dataset.INDEXING_TECHNIQUE_LIST,
  112. nullable=True,
  113. help="Invalid indexing technique.",
  114. )
  115. parser.add_argument(
  116. "external_knowledge_api_id",
  117. type=str,
  118. nullable=True,
  119. required=False,
  120. )
  121. parser.add_argument(
  122. "provider",
  123. type=str,
  124. nullable=True,
  125. choices=Dataset.PROVIDER_LIST,
  126. required=False,
  127. default="vendor",
  128. )
  129. parser.add_argument(
  130. "external_knowledge_id",
  131. type=str,
  132. nullable=True,
  133. required=False,
  134. )
  135. args = parser.parse_args()
  136. # The role of the current user in the ta table must be admin, owner, or editor, or dataset_operator
  137. if not current_user.is_dataset_editor:
  138. raise Forbidden()
  139. try:
  140. dataset = DatasetService.create_empty_dataset(
  141. tenant_id=current_user.current_tenant_id,
  142. name=args["name"],
  143. description=args["description"],
  144. indexing_technique=args["indexing_technique"],
  145. account=current_user,
  146. permission=DatasetPermissionEnum.ONLY_ME,
  147. provider=args["provider"],
  148. external_knowledge_api_id=args["external_knowledge_api_id"],
  149. external_knowledge_id=args["external_knowledge_id"],
  150. )
  151. except services.errors.dataset.DatasetNameDuplicateError:
  152. raise DatasetNameDuplicateError()
  153. return marshal(dataset, dataset_detail_fields), 201
  154. class DatasetApi(Resource):
  155. @setup_required
  156. @login_required
  157. @account_initialization_required
  158. def get(self, dataset_id):
  159. dataset_id_str = str(dataset_id)
  160. dataset = DatasetService.get_dataset(dataset_id_str)
  161. if dataset is None:
  162. raise NotFound("Dataset not found.")
  163. try:
  164. DatasetService.check_dataset_permission(dataset, current_user)
  165. except services.errors.account.NoPermissionError as e:
  166. raise Forbidden(str(e))
  167. data = marshal(dataset, dataset_detail_fields)
  168. if dataset.indexing_technique == "high_quality":
  169. if dataset.embedding_model_provider:
  170. provider_id = ModelProviderID(dataset.embedding_model_provider)
  171. data["embedding_model_provider"] = str(provider_id)
  172. if data.get("permission") == "partial_members":
  173. part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)
  174. data.update({"partial_member_list": part_users_list})
  175. # check embedding setting
  176. provider_manager = ProviderManager()
  177. configurations = provider_manager.get_configurations(tenant_id=current_user.current_tenant_id)
  178. embedding_models = configurations.get_models(model_type=ModelType.TEXT_EMBEDDING, only_active=True)
  179. model_names = []
  180. for embedding_model in embedding_models:
  181. model_names.append(f"{embedding_model.model}:{embedding_model.provider.provider}")
  182. if data["indexing_technique"] == "high_quality":
  183. item_model = f"{data['embedding_model']}:{data['embedding_model_provider']}"
  184. if item_model in model_names:
  185. data["embedding_available"] = True
  186. else:
  187. data["embedding_available"] = False
  188. else:
  189. data["embedding_available"] = True
  190. if data.get("permission") == "partial_members":
  191. part_users_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)
  192. data.update({"partial_member_list": part_users_list})
  193. return data, 200
  194. @setup_required
  195. @login_required
  196. @account_initialization_required
  197. @cloud_edition_billing_rate_limit_check("knowledge")
  198. def patch(self, dataset_id):
  199. dataset_id_str = str(dataset_id)
  200. dataset = DatasetService.get_dataset(dataset_id_str)
  201. if dataset is None:
  202. raise NotFound("Dataset not found.")
  203. parser = reqparse.RequestParser()
  204. parser.add_argument(
  205. "name",
  206. nullable=False,
  207. help="type is required. Name must be between 1 to 40 characters.",
  208. type=_validate_name,
  209. )
  210. parser.add_argument("description", location="json", store_missing=False, type=_validate_description_length)
  211. parser.add_argument(
  212. "indexing_technique",
  213. type=str,
  214. location="json",
  215. choices=Dataset.INDEXING_TECHNIQUE_LIST,
  216. nullable=True,
  217. help="Invalid indexing technique.",
  218. )
  219. parser.add_argument(
  220. "permission",
  221. type=str,
  222. location="json",
  223. choices=(DatasetPermissionEnum.ONLY_ME, DatasetPermissionEnum.ALL_TEAM, DatasetPermissionEnum.PARTIAL_TEAM),
  224. help="Invalid permission.",
  225. )
  226. parser.add_argument("embedding_model", type=str, location="json", help="Invalid embedding model.")
  227. parser.add_argument(
  228. "embedding_model_provider", type=str, location="json", help="Invalid embedding model provider."
  229. )
  230. parser.add_argument("retrieval_model", type=dict, location="json", help="Invalid retrieval model.")
  231. parser.add_argument("partial_member_list", type=list, location="json", help="Invalid parent user list.")
  232. parser.add_argument(
  233. "external_retrieval_model",
  234. type=dict,
  235. required=False,
  236. nullable=True,
  237. location="json",
  238. help="Invalid external retrieval model.",
  239. )
  240. parser.add_argument(
  241. "external_knowledge_id",
  242. type=str,
  243. required=False,
  244. nullable=True,
  245. location="json",
  246. help="Invalid external knowledge id.",
  247. )
  248. parser.add_argument(
  249. "external_knowledge_api_id",
  250. type=str,
  251. required=False,
  252. nullable=True,
  253. location="json",
  254. help="Invalid external knowledge api id.",
  255. )
  256. parser.add_argument(
  257. "icon_info",
  258. type=dict,
  259. required=False,
  260. nullable=True,
  261. location="json",
  262. help="Invalid icon info.",
  263. )
  264. args = parser.parse_args()
  265. data = request.get_json()
  266. # check embedding model setting
  267. if (
  268. data.get("indexing_technique") == "high_quality"
  269. and data.get("embedding_model_provider") is not None
  270. and data.get("embedding_model") is not None
  271. ):
  272. DatasetService.check_embedding_model_setting(
  273. dataset.tenant_id, data.get("embedding_model_provider"), data.get("embedding_model")
  274. )
  275. # The role of the current user in the ta table must be admin, owner, editor, or dataset_operator
  276. DatasetPermissionService.check_permission(
  277. current_user, dataset, data.get("permission"), data.get("partial_member_list")
  278. )
  279. dataset = DatasetService.update_dataset(dataset_id_str, args, current_user)
  280. if dataset is None:
  281. raise NotFound("Dataset not found.")
  282. result_data = marshal(dataset, dataset_detail_fields)
  283. tenant_id = current_user.current_tenant_id
  284. if data.get("partial_member_list") and data.get("permission") == "partial_members":
  285. DatasetPermissionService.update_partial_member_list(
  286. tenant_id, dataset_id_str, data.get("partial_member_list")
  287. )
  288. # clear partial member list when permission is only_me or all_team_members
  289. elif (
  290. data.get("permission") == DatasetPermissionEnum.ONLY_ME
  291. or data.get("permission") == DatasetPermissionEnum.ALL_TEAM
  292. ):
  293. DatasetPermissionService.clear_partial_member_list(dataset_id_str)
  294. partial_member_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)
  295. result_data.update({"partial_member_list": partial_member_list})
  296. return result_data, 200
  297. @setup_required
  298. @login_required
  299. @account_initialization_required
  300. @cloud_edition_billing_rate_limit_check("knowledge")
  301. def delete(self, dataset_id):
  302. dataset_id_str = str(dataset_id)
  303. # The role of the current user in the ta table must be admin, owner, or editor
  304. if not current_user.is_editor or current_user.is_dataset_operator:
  305. raise Forbidden()
  306. try:
  307. if DatasetService.delete_dataset(dataset_id_str, current_user):
  308. DatasetPermissionService.clear_partial_member_list(dataset_id_str)
  309. return {"result": "success"}, 204
  310. else:
  311. raise NotFound("Dataset not found.")
  312. except services.errors.dataset.DatasetInUseError:
  313. raise DatasetInUseError()
  314. class DatasetUseCheckApi(Resource):
  315. @setup_required
  316. @login_required
  317. @account_initialization_required
  318. def get(self, dataset_id):
  319. dataset_id_str = str(dataset_id)
  320. dataset_is_using = DatasetService.dataset_use_check(dataset_id_str)
  321. return {"is_using": dataset_is_using}, 200
  322. class DatasetQueryApi(Resource):
  323. @setup_required
  324. @login_required
  325. @account_initialization_required
  326. def get(self, dataset_id):
  327. dataset_id_str = str(dataset_id)
  328. dataset = DatasetService.get_dataset(dataset_id_str)
  329. if dataset is None:
  330. raise NotFound("Dataset not found.")
  331. try:
  332. DatasetService.check_dataset_permission(dataset, current_user)
  333. except services.errors.account.NoPermissionError as e:
  334. raise Forbidden(str(e))
  335. page = request.args.get("page", default=1, type=int)
  336. limit = request.args.get("limit", default=20, type=int)
  337. dataset_queries, total = DatasetService.get_dataset_queries(dataset_id=dataset.id, page=page, per_page=limit)
  338. response = {
  339. "data": marshal(dataset_queries, dataset_query_detail_fields),
  340. "has_more": len(dataset_queries) == limit,
  341. "limit": limit,
  342. "total": total,
  343. "page": page,
  344. }
  345. return response, 200
  346. class DatasetIndexingEstimateApi(Resource):
  347. @setup_required
  348. @login_required
  349. @account_initialization_required
  350. def post(self):
  351. parser = reqparse.RequestParser()
  352. parser.add_argument("info_list", type=dict, required=True, nullable=True, location="json")
  353. parser.add_argument("process_rule", type=dict, required=True, nullable=True, location="json")
  354. parser.add_argument(
  355. "indexing_technique",
  356. type=str,
  357. required=True,
  358. choices=Dataset.INDEXING_TECHNIQUE_LIST,
  359. nullable=True,
  360. location="json",
  361. )
  362. parser.add_argument("doc_form", type=str, default="text_model", required=False, nullable=False, location="json")
  363. parser.add_argument("dataset_id", type=str, required=False, nullable=False, location="json")
  364. parser.add_argument(
  365. "doc_language", type=str, default="English", required=False, nullable=False, location="json"
  366. )
  367. args = parser.parse_args()
  368. # validate args
  369. DocumentService.estimate_args_validate(args)
  370. extract_settings = []
  371. if args["info_list"]["data_source_type"] == "upload_file":
  372. file_ids = args["info_list"]["file_info_list"]["file_ids"]
  373. file_details = (
  374. db.session.query(UploadFile)
  375. .filter(UploadFile.tenant_id == current_user.current_tenant_id, UploadFile.id.in_(file_ids))
  376. .all()
  377. )
  378. if file_details is None:
  379. raise NotFound("File not found.")
  380. if file_details:
  381. for file_detail in file_details:
  382. extract_setting = ExtractSetting(
  383. datasource_type="upload_file", upload_file=file_detail, document_model=args["doc_form"]
  384. )
  385. extract_settings.append(extract_setting)
  386. elif args["info_list"]["data_source_type"] == "notion_import":
  387. notion_info_list = args["info_list"]["notion_info_list"]
  388. for notion_info in notion_info_list:
  389. workspace_id = notion_info["workspace_id"]
  390. for page in notion_info["pages"]:
  391. extract_setting = ExtractSetting(
  392. datasource_type="notion_import",
  393. notion_info={
  394. "notion_workspace_id": workspace_id,
  395. "notion_obj_id": page["page_id"],
  396. "notion_page_type": page["type"],
  397. "tenant_id": current_user.current_tenant_id,
  398. },
  399. document_model=args["doc_form"],
  400. )
  401. extract_settings.append(extract_setting)
  402. elif args["info_list"]["data_source_type"] == "website_crawl":
  403. website_info_list = args["info_list"]["website_info_list"]
  404. for url in website_info_list["urls"]:
  405. extract_setting = ExtractSetting(
  406. datasource_type="website_crawl",
  407. website_info={
  408. "provider": website_info_list["provider"],
  409. "job_id": website_info_list["job_id"],
  410. "url": url,
  411. "tenant_id": current_user.current_tenant_id,
  412. "mode": "crawl",
  413. "only_main_content": website_info_list["only_main_content"],
  414. },
  415. document_model=args["doc_form"],
  416. )
  417. extract_settings.append(extract_setting)
  418. else:
  419. raise ValueError("Data source type not support")
  420. indexing_runner = IndexingRunner()
  421. try:
  422. response = indexing_runner.indexing_estimate(
  423. current_user.current_tenant_id,
  424. extract_settings,
  425. args["process_rule"],
  426. args["doc_form"],
  427. args["doc_language"],
  428. args["dataset_id"],
  429. args["indexing_technique"],
  430. )
  431. except LLMBadRequestError:
  432. raise ProviderNotInitializeError(
  433. "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider."
  434. )
  435. except ProviderTokenNotInitError as ex:
  436. raise ProviderNotInitializeError(ex.description)
  437. except Exception as e:
  438. raise IndexingEstimateError(str(e))
  439. return response.model_dump(), 200
  440. class DatasetRelatedAppListApi(Resource):
  441. @setup_required
  442. @login_required
  443. @account_initialization_required
  444. @marshal_with(related_app_list)
  445. def get(self, dataset_id):
  446. dataset_id_str = str(dataset_id)
  447. dataset = DatasetService.get_dataset(dataset_id_str)
  448. if dataset is None:
  449. raise NotFound("Dataset not found.")
  450. try:
  451. DatasetService.check_dataset_permission(dataset, current_user)
  452. except services.errors.account.NoPermissionError as e:
  453. raise Forbidden(str(e))
  454. app_dataset_joins = DatasetService.get_related_apps(dataset.id)
  455. related_apps = []
  456. for app_dataset_join in app_dataset_joins:
  457. app_model = app_dataset_join.app
  458. if app_model:
  459. related_apps.append(app_model)
  460. return {"data": related_apps, "total": len(related_apps)}, 200
  461. class DatasetIndexingStatusApi(Resource):
  462. @setup_required
  463. @login_required
  464. @account_initialization_required
  465. def get(self, dataset_id):
  466. dataset_id = str(dataset_id)
  467. documents = (
  468. db.session.query(Document)
  469. .filter(Document.dataset_id == dataset_id, Document.tenant_id == current_user.current_tenant_id)
  470. .all()
  471. )
  472. documents_status = []
  473. for document in documents:
  474. completed_segments = (
  475. db.session.query(DocumentSegment)
  476. .filter(
  477. DocumentSegment.completed_at.isnot(None),
  478. DocumentSegment.document_id == str(document.id),
  479. DocumentSegment.status != "re_segment",
  480. )
  481. .count()
  482. )
  483. total_segments = (
  484. db.session.query(DocumentSegment)
  485. .filter(DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment")
  486. .count()
  487. )
  488. # Create a dictionary with document attributes and additional fields
  489. document_dict = {
  490. "id": document.id,
  491. "indexing_status": document.indexing_status,
  492. "processing_started_at": document.processing_started_at,
  493. "parsing_completed_at": document.parsing_completed_at,
  494. "cleaning_completed_at": document.cleaning_completed_at,
  495. "splitting_completed_at": document.splitting_completed_at,
  496. "completed_at": document.completed_at,
  497. "paused_at": document.paused_at,
  498. "error": document.error,
  499. "stopped_at": document.stopped_at,
  500. "completed_segments": completed_segments,
  501. "total_segments": total_segments,
  502. }
  503. documents_status.append(marshal(document_dict, document_status_fields))
  504. data = {"data": documents_status}
  505. return data
  506. class DatasetApiKeyApi(Resource):
  507. max_keys = 10
  508. token_prefix = "dataset-"
  509. resource_type = "dataset"
  510. @setup_required
  511. @login_required
  512. @account_initialization_required
  513. @marshal_with(api_key_list)
  514. def get(self):
  515. keys = (
  516. db.session.query(ApiToken)
  517. .filter(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id)
  518. .all()
  519. )
  520. return {"items": keys}
  521. @setup_required
  522. @login_required
  523. @account_initialization_required
  524. @marshal_with(api_key_fields)
  525. def post(self):
  526. # The role of the current user in the ta table must be admin or owner
  527. if not current_user.is_admin_or_owner:
  528. raise Forbidden()
  529. current_key_count = (
  530. db.session.query(ApiToken)
  531. .filter(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id)
  532. .count()
  533. )
  534. if current_key_count >= self.max_keys:
  535. flask_restful.abort(
  536. 400,
  537. message=f"Cannot create more than {self.max_keys} API keys for this resource type.",
  538. code="max_keys_exceeded",
  539. )
  540. key = ApiToken.generate_api_key(self.token_prefix, 24)
  541. api_token = ApiToken()
  542. api_token.tenant_id = current_user.current_tenant_id
  543. api_token.token = key
  544. api_token.type = self.resource_type
  545. db.session.add(api_token)
  546. db.session.commit()
  547. return api_token, 200
  548. class DatasetApiDeleteApi(Resource):
  549. resource_type = "dataset"
  550. @setup_required
  551. @login_required
  552. @account_initialization_required
  553. def delete(self, api_key_id):
  554. api_key_id = str(api_key_id)
  555. # The role of the current user in the ta table must be admin or owner
  556. if not current_user.is_admin_or_owner:
  557. raise Forbidden()
  558. key = (
  559. db.session.query(ApiToken)
  560. .filter(
  561. ApiToken.tenant_id == current_user.current_tenant_id,
  562. ApiToken.type == self.resource_type,
  563. ApiToken.id == api_key_id,
  564. )
  565. .first()
  566. )
  567. if key is None:
  568. flask_restful.abort(404, message="API key not found")
  569. db.session.query(ApiToken).filter(ApiToken.id == api_key_id).delete()
  570. db.session.commit()
  571. return {"result": "success"}, 204
  572. class DatasetApiBaseUrlApi(Resource):
  573. @setup_required
  574. @login_required
  575. @account_initialization_required
  576. def get(self):
  577. return {"api_base_url": (dify_config.SERVICE_API_URL or request.host_url.rstrip("/")) + "/v1"}
  578. class DatasetRetrievalSettingApi(Resource):
  579. @setup_required
  580. @login_required
  581. @account_initialization_required
  582. def get(self):
  583. vector_type = dify_config.VECTOR_STORE
  584. match vector_type:
  585. case (
  586. VectorType.RELYT
  587. | VectorType.TIDB_VECTOR
  588. | VectorType.CHROMA
  589. | VectorType.PGVECTO_RS
  590. | VectorType.BAIDU
  591. | VectorType.VIKINGDB
  592. | VectorType.UPSTASH
  593. ):
  594. return {"retrieval_method": [RetrievalMethod.SEMANTIC_SEARCH.value]}
  595. case (
  596. VectorType.QDRANT
  597. | VectorType.WEAVIATE
  598. | VectorType.OPENSEARCH
  599. | VectorType.ANALYTICDB
  600. | VectorType.MYSCALE
  601. | VectorType.ORACLE
  602. | VectorType.ELASTICSEARCH
  603. | VectorType.ELASTICSEARCH_JA
  604. | VectorType.PGVECTOR
  605. | VectorType.VASTBASE
  606. | VectorType.TIDB_ON_QDRANT
  607. | VectorType.LINDORM
  608. | VectorType.COUCHBASE
  609. | VectorType.MILVUS
  610. | VectorType.OPENGAUSS
  611. | VectorType.OCEANBASE
  612. | VectorType.TABLESTORE
  613. | VectorType.HUAWEI_CLOUD
  614. | VectorType.TENCENT
  615. ):
  616. return {
  617. "retrieval_method": [
  618. RetrievalMethod.SEMANTIC_SEARCH.value,
  619. RetrievalMethod.FULL_TEXT_SEARCH.value,
  620. RetrievalMethod.HYBRID_SEARCH.value,
  621. ]
  622. }
  623. case _:
  624. raise ValueError(f"Unsupported vector db type {vector_type}.")
  625. class DatasetRetrievalSettingMockApi(Resource):
  626. @setup_required
  627. @login_required
  628. @account_initialization_required
  629. def get(self, vector_type):
  630. match vector_type:
  631. case (
  632. VectorType.MILVUS
  633. | VectorType.RELYT
  634. | VectorType.TIDB_VECTOR
  635. | VectorType.CHROMA
  636. | VectorType.PGVECTO_RS
  637. | VectorType.BAIDU
  638. | VectorType.VIKINGDB
  639. | VectorType.UPSTASH
  640. ):
  641. return {"retrieval_method": [RetrievalMethod.SEMANTIC_SEARCH.value]}
  642. case (
  643. VectorType.QDRANT
  644. | VectorType.WEAVIATE
  645. | VectorType.OPENSEARCH
  646. | VectorType.ANALYTICDB
  647. | VectorType.MYSCALE
  648. | VectorType.ORACLE
  649. | VectorType.ELASTICSEARCH
  650. | VectorType.ELASTICSEARCH_JA
  651. | VectorType.COUCHBASE
  652. | VectorType.PGVECTOR
  653. | VectorType.VASTBASE
  654. | VectorType.LINDORM
  655. | VectorType.OPENGAUSS
  656. | VectorType.OCEANBASE
  657. | VectorType.TABLESTORE
  658. | VectorType.TENCENT
  659. | VectorType.HUAWEI_CLOUD
  660. ):
  661. return {
  662. "retrieval_method": [
  663. RetrievalMethod.SEMANTIC_SEARCH.value,
  664. RetrievalMethod.FULL_TEXT_SEARCH.value,
  665. RetrievalMethod.HYBRID_SEARCH.value,
  666. ]
  667. }
  668. case _:
  669. raise ValueError(f"Unsupported vector db type {vector_type}.")
  670. class DatasetErrorDocs(Resource):
  671. @setup_required
  672. @login_required
  673. @account_initialization_required
  674. def get(self, dataset_id):
  675. dataset_id_str = str(dataset_id)
  676. dataset = DatasetService.get_dataset(dataset_id_str)
  677. if dataset is None:
  678. raise NotFound("Dataset not found.")
  679. results = DocumentService.get_error_documents_by_dataset_id(dataset_id_str)
  680. return {"data": [marshal(item, document_status_fields) for item in results], "total": len(results)}, 200
  681. class DatasetPermissionUserListApi(Resource):
  682. @setup_required
  683. @login_required
  684. @account_initialization_required
  685. def get(self, dataset_id):
  686. dataset_id_str = str(dataset_id)
  687. dataset = DatasetService.get_dataset(dataset_id_str)
  688. if dataset is None:
  689. raise NotFound("Dataset not found.")
  690. try:
  691. DatasetService.check_dataset_permission(dataset, current_user)
  692. except services.errors.account.NoPermissionError as e:
  693. raise Forbidden(str(e))
  694. partial_members_list = DatasetPermissionService.get_dataset_partial_member_list(dataset_id_str)
  695. return {
  696. "data": partial_members_list,
  697. }, 200
  698. class DatasetAutoDisableLogApi(Resource):
  699. @setup_required
  700. @login_required
  701. @account_initialization_required
  702. def get(self, dataset_id):
  703. dataset_id_str = str(dataset_id)
  704. dataset = DatasetService.get_dataset(dataset_id_str)
  705. if dataset is None:
  706. raise NotFound("Dataset not found.")
  707. return DatasetService.get_dataset_auto_disable_logs(dataset_id_str), 200
  708. api.add_resource(DatasetListApi, "/datasets")
  709. api.add_resource(DatasetApi, "/datasets/<uuid:dataset_id>")
  710. api.add_resource(DatasetUseCheckApi, "/datasets/<uuid:dataset_id>/use-check")
  711. api.add_resource(DatasetQueryApi, "/datasets/<uuid:dataset_id>/queries")
  712. api.add_resource(DatasetErrorDocs, "/datasets/<uuid:dataset_id>/error-docs")
  713. api.add_resource(DatasetIndexingEstimateApi, "/datasets/indexing-estimate")
  714. api.add_resource(DatasetRelatedAppListApi, "/datasets/<uuid:dataset_id>/related-apps")
  715. api.add_resource(DatasetIndexingStatusApi, "/datasets/<uuid:dataset_id>/indexing-status")
  716. api.add_resource(DatasetApiKeyApi, "/datasets/api-keys")
  717. api.add_resource(DatasetApiDeleteApi, "/datasets/api-keys/<uuid:api_key_id>")
  718. api.add_resource(DatasetApiBaseUrlApi, "/datasets/api-base-info")
  719. api.add_resource(DatasetRetrievalSettingApi, "/datasets/retrieval-setting")
  720. api.add_resource(DatasetRetrievalSettingMockApi, "/datasets/retrieval-setting/<string:vector_type>")
  721. api.add_resource(DatasetPermissionUserListApi, "/datasets/<uuid:dataset_id>/permission-part-users")
  722. api.add_resource(DatasetAutoDisableLogApi, "/datasets/<uuid:dataset_id>/auto-disable-logs")