datasets.py 30 KB

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