datasets.py 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495
  1. # -*- coding:utf-8 -*-
  2. import flask_restful
  3. from flask import request, current_app
  4. from flask_login import current_user
  5. from controllers.console.apikey import api_key_list, api_key_fields
  6. from libs.login import login_required
  7. from flask_restful import Resource, reqparse, marshal, marshal_with
  8. from werkzeug.exceptions import NotFound, Forbidden
  9. import services
  10. from controllers.console import api
  11. from controllers.console.app.error import ProviderNotInitializeError
  12. from controllers.console.datasets.error import DatasetNameDuplicateError
  13. from controllers.console.setup import setup_required
  14. from controllers.console.wraps import account_initialization_required
  15. from core.indexing_runner import IndexingRunner
  16. from core.model_providers.error import LLMBadRequestError, ProviderTokenNotInitError
  17. from core.model_providers.models.entity.model_params import ModelType
  18. from fields.app_fields import related_app_list
  19. from fields.dataset_fields import dataset_detail_fields, dataset_query_detail_fields
  20. from fields.document_fields import document_status_fields
  21. from extensions.ext_database import db
  22. from models.dataset import DocumentSegment, Document
  23. from models.model import UploadFile, ApiToken
  24. from services.dataset_service import DatasetService, DocumentService
  25. from services.provider_service import ProviderService
  26. def _validate_name(name):
  27. if not name or len(name) < 1 or len(name) > 40:
  28. raise ValueError('Name must be between 1 to 40 characters.')
  29. return name
  30. def _validate_description_length(description):
  31. if len(description) > 400:
  32. raise ValueError('Description cannot exceed 400 characters.')
  33. return description
  34. class DatasetListApi(Resource):
  35. @setup_required
  36. @login_required
  37. @account_initialization_required
  38. def get(self):
  39. page = request.args.get('page', default=1, type=int)
  40. limit = request.args.get('limit', default=20, type=int)
  41. ids = request.args.getlist('ids')
  42. provider = request.args.get('provider', default="vendor")
  43. if ids:
  44. datasets, total = DatasetService.get_datasets_by_ids(ids, current_user.current_tenant_id)
  45. else:
  46. datasets, total = DatasetService.get_datasets(page, limit, provider,
  47. current_user.current_tenant_id, current_user)
  48. # check embedding setting
  49. provider_service = ProviderService()
  50. valid_model_list = provider_service.get_valid_model_list(current_user.current_tenant_id,
  51. ModelType.EMBEDDINGS.value)
  52. # if len(valid_model_list) == 0:
  53. # raise ProviderNotInitializeError(
  54. # f"No Embedding Model available. Please configure a valid provider "
  55. # f"in the Settings -> Model Provider.")
  56. model_names = []
  57. for valid_model in valid_model_list:
  58. model_names.append(f"{valid_model['model_name']}:{valid_model['model_provider']['provider_name']}")
  59. data = marshal(datasets, dataset_detail_fields)
  60. for item in data:
  61. if item['indexing_technique'] == 'high_quality':
  62. item_model = f"{item['embedding_model']}:{item['embedding_model_provider']}"
  63. if item_model in model_names:
  64. item['embedding_available'] = True
  65. else:
  66. item['embedding_available'] = False
  67. else:
  68. item['embedding_available'] = True
  69. response = {
  70. 'data': data,
  71. 'has_more': len(datasets) == limit,
  72. 'limit': limit,
  73. 'total': total,
  74. 'page': page
  75. }
  76. return response, 200
  77. @setup_required
  78. @login_required
  79. @account_initialization_required
  80. def post(self):
  81. parser = reqparse.RequestParser()
  82. parser.add_argument('name', nullable=False, required=True,
  83. help='type is required. Name must be between 1 to 40 characters.',
  84. type=_validate_name)
  85. parser.add_argument('indexing_technique', type=str, location='json',
  86. choices=('high_quality', 'economy'),
  87. help='Invalid indexing technique.')
  88. args = parser.parse_args()
  89. # The role of the current user in the ta table must be admin or owner
  90. if current_user.current_tenant.current_role not in ['admin', 'owner']:
  91. raise Forbidden()
  92. try:
  93. dataset = DatasetService.create_empty_dataset(
  94. tenant_id=current_user.current_tenant_id,
  95. name=args['name'],
  96. indexing_technique=args['indexing_technique'],
  97. account=current_user
  98. )
  99. except services.errors.dataset.DatasetNameDuplicateError:
  100. raise DatasetNameDuplicateError()
  101. return marshal(dataset, dataset_detail_fields), 201
  102. class DatasetApi(Resource):
  103. @setup_required
  104. @login_required
  105. @account_initialization_required
  106. def get(self, dataset_id):
  107. dataset_id_str = str(dataset_id)
  108. dataset = DatasetService.get_dataset(dataset_id_str)
  109. if dataset is None:
  110. raise NotFound("Dataset not found.")
  111. try:
  112. DatasetService.check_dataset_permission(
  113. dataset, current_user)
  114. except services.errors.account.NoPermissionError as e:
  115. raise Forbidden(str(e))
  116. data = marshal(dataset, dataset_detail_fields)
  117. # check embedding setting
  118. provider_service = ProviderService()
  119. # get valid model list
  120. valid_model_list = provider_service.get_valid_model_list(current_user.current_tenant_id,
  121. ModelType.EMBEDDINGS.value)
  122. model_names = []
  123. for valid_model in valid_model_list:
  124. model_names.append(f"{valid_model['model_name']}:{valid_model['model_provider']['provider_name']}")
  125. if data['indexing_technique'] == 'high_quality':
  126. item_model = f"{data['embedding_model']}:{data['embedding_model_provider']}"
  127. if item_model in model_names:
  128. data['embedding_available'] = True
  129. else:
  130. data['embedding_available'] = False
  131. else:
  132. data['embedding_available'] = True
  133. return data, 200
  134. @setup_required
  135. @login_required
  136. @account_initialization_required
  137. def patch(self, dataset_id):
  138. dataset_id_str = str(dataset_id)
  139. dataset = DatasetService.get_dataset(dataset_id_str)
  140. if dataset is None:
  141. raise NotFound("Dataset not found.")
  142. # check user's model setting
  143. DatasetService.check_dataset_model_setting(dataset)
  144. parser = reqparse.RequestParser()
  145. parser.add_argument('name', nullable=False,
  146. help='type is required. Name must be between 1 to 40 characters.',
  147. type=_validate_name)
  148. parser.add_argument('description',
  149. location='json', store_missing=False,
  150. type=_validate_description_length)
  151. parser.add_argument('indexing_technique', type=str, location='json',
  152. choices=('high_quality', 'economy'),
  153. help='Invalid indexing technique.')
  154. parser.add_argument('permission', type=str, location='json', choices=(
  155. 'only_me', 'all_team_members'), help='Invalid permission.')
  156. parser.add_argument('retrieval_model', type=dict, location='json', help='Invalid retrieval model.')
  157. args = parser.parse_args()
  158. # The role of the current user in the ta table must be admin or owner
  159. if current_user.current_tenant.current_role not in ['admin', 'owner']:
  160. raise Forbidden()
  161. dataset = DatasetService.update_dataset(
  162. dataset_id_str, args, current_user)
  163. if dataset is None:
  164. raise NotFound("Dataset not found.")
  165. return marshal(dataset, dataset_detail_fields), 200
  166. @setup_required
  167. @login_required
  168. @account_initialization_required
  169. def delete(self, dataset_id):
  170. dataset_id_str = str(dataset_id)
  171. # The role of the current user in the ta table must be admin or owner
  172. if current_user.current_tenant.current_role not in ['admin', 'owner']:
  173. raise Forbidden()
  174. if DatasetService.delete_dataset(dataset_id_str, current_user):
  175. return {'result': 'success'}, 204
  176. else:
  177. raise NotFound("Dataset not found.")
  178. class DatasetQueryApi(Resource):
  179. @setup_required
  180. @login_required
  181. @account_initialization_required
  182. def get(self, dataset_id):
  183. dataset_id_str = str(dataset_id)
  184. dataset = DatasetService.get_dataset(dataset_id_str)
  185. if dataset is None:
  186. raise NotFound("Dataset not found.")
  187. try:
  188. DatasetService.check_dataset_permission(dataset, current_user)
  189. except services.errors.account.NoPermissionError as e:
  190. raise Forbidden(str(e))
  191. page = request.args.get('page', default=1, type=int)
  192. limit = request.args.get('limit', default=20, type=int)
  193. dataset_queries, total = DatasetService.get_dataset_queries(
  194. dataset_id=dataset.id,
  195. page=page,
  196. per_page=limit
  197. )
  198. response = {
  199. 'data': marshal(dataset_queries, dataset_query_detail_fields),
  200. 'has_more': len(dataset_queries) == limit,
  201. 'limit': limit,
  202. 'total': total,
  203. 'page': page
  204. }
  205. return response, 200
  206. class DatasetIndexingEstimateApi(Resource):
  207. @setup_required
  208. @login_required
  209. @account_initialization_required
  210. def post(self):
  211. parser = reqparse.RequestParser()
  212. parser.add_argument('info_list', type=dict, required=True, nullable=True, location='json')
  213. parser.add_argument('process_rule', type=dict, required=True, nullable=True, location='json')
  214. parser.add_argument('indexing_technique', type=str, required=True, nullable=True, location='json')
  215. parser.add_argument('doc_form', type=str, default='text_model', required=False, nullable=False, location='json')
  216. parser.add_argument('dataset_id', type=str, required=False, nullable=False, location='json')
  217. parser.add_argument('doc_language', type=str, default='English', required=False, nullable=False,
  218. location='json')
  219. args = parser.parse_args()
  220. # validate args
  221. DocumentService.estimate_args_validate(args)
  222. if args['info_list']['data_source_type'] == 'upload_file':
  223. file_ids = args['info_list']['file_info_list']['file_ids']
  224. file_details = db.session.query(UploadFile).filter(
  225. UploadFile.tenant_id == current_user.current_tenant_id,
  226. UploadFile.id.in_(file_ids)
  227. ).all()
  228. if file_details is None:
  229. raise NotFound("File not found.")
  230. indexing_runner = IndexingRunner()
  231. try:
  232. response = indexing_runner.file_indexing_estimate(current_user.current_tenant_id, file_details,
  233. args['process_rule'], args['doc_form'],
  234. args['doc_language'], args['dataset_id'],
  235. args['indexing_technique'])
  236. except LLMBadRequestError:
  237. raise ProviderNotInitializeError(
  238. f"No Embedding Model available. Please configure a valid provider "
  239. f"in the Settings -> Model Provider.")
  240. except ProviderTokenNotInitError as ex:
  241. raise ProviderNotInitializeError(ex.description)
  242. elif args['info_list']['data_source_type'] == 'notion_import':
  243. indexing_runner = IndexingRunner()
  244. try:
  245. response = indexing_runner.notion_indexing_estimate(current_user.current_tenant_id,
  246. args['info_list']['notion_info_list'],
  247. args['process_rule'], args['doc_form'],
  248. args['doc_language'], args['dataset_id'],
  249. args['indexing_technique'])
  250. except LLMBadRequestError:
  251. raise ProviderNotInitializeError(
  252. f"No Embedding Model available. Please configure a valid provider "
  253. f"in the Settings -> Model Provider.")
  254. except ProviderTokenNotInitError as ex:
  255. raise ProviderNotInitializeError(ex.description)
  256. else:
  257. raise ValueError('Data source type not support')
  258. return response, 200
  259. class DatasetRelatedAppListApi(Resource):
  260. @setup_required
  261. @login_required
  262. @account_initialization_required
  263. @marshal_with(related_app_list)
  264. def get(self, dataset_id):
  265. dataset_id_str = str(dataset_id)
  266. dataset = DatasetService.get_dataset(dataset_id_str)
  267. if dataset is None:
  268. raise NotFound("Dataset not found.")
  269. try:
  270. DatasetService.check_dataset_permission(dataset, current_user)
  271. except services.errors.account.NoPermissionError as e:
  272. raise Forbidden(str(e))
  273. app_dataset_joins = DatasetService.get_related_apps(dataset.id)
  274. related_apps = []
  275. for app_dataset_join in app_dataset_joins:
  276. app_model = app_dataset_join.app
  277. if app_model:
  278. related_apps.append(app_model)
  279. return {
  280. 'data': related_apps,
  281. 'total': len(related_apps)
  282. }, 200
  283. class DatasetIndexingStatusApi(Resource):
  284. @setup_required
  285. @login_required
  286. @account_initialization_required
  287. def get(self, dataset_id):
  288. dataset_id = str(dataset_id)
  289. documents = db.session.query(Document).filter(
  290. Document.dataset_id == dataset_id,
  291. Document.tenant_id == current_user.current_tenant_id
  292. ).all()
  293. documents_status = []
  294. for document in documents:
  295. completed_segments = DocumentSegment.query.filter(DocumentSegment.completed_at.isnot(None),
  296. DocumentSegment.document_id == str(document.id),
  297. DocumentSegment.status != 're_segment').count()
  298. total_segments = DocumentSegment.query.filter(DocumentSegment.document_id == str(document.id),
  299. DocumentSegment.status != 're_segment').count()
  300. document.completed_segments = completed_segments
  301. document.total_segments = total_segments
  302. documents_status.append(marshal(document, document_status_fields))
  303. data = {
  304. 'data': documents_status
  305. }
  306. return data
  307. class DatasetApiKeyApi(Resource):
  308. max_keys = 10
  309. token_prefix = 'dataset-'
  310. resource_type = 'dataset'
  311. @setup_required
  312. @login_required
  313. @account_initialization_required
  314. @marshal_with(api_key_list)
  315. def get(self):
  316. keys = db.session.query(ApiToken). \
  317. filter(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id). \
  318. all()
  319. return {"items": keys}
  320. @setup_required
  321. @login_required
  322. @account_initialization_required
  323. @marshal_with(api_key_fields)
  324. def post(self):
  325. # The role of the current user in the ta table must be admin or owner
  326. if current_user.current_tenant.current_role not in ['admin', 'owner']:
  327. raise Forbidden()
  328. current_key_count = db.session.query(ApiToken). \
  329. filter(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id). \
  330. count()
  331. if current_key_count >= self.max_keys:
  332. flask_restful.abort(
  333. 400,
  334. message=f"Cannot create more than {self.max_keys} API keys for this resource type.",
  335. code='max_keys_exceeded'
  336. )
  337. key = ApiToken.generate_api_key(self.token_prefix, 24)
  338. api_token = ApiToken()
  339. api_token.tenant_id = current_user.current_tenant_id
  340. api_token.token = key
  341. api_token.type = self.resource_type
  342. db.session.add(api_token)
  343. db.session.commit()
  344. return api_token, 200
  345. class DatasetApiDeleteApi(Resource):
  346. resource_type = 'dataset'
  347. @setup_required
  348. @login_required
  349. @account_initialization_required
  350. def delete(self, api_key_id):
  351. api_key_id = str(api_key_id)
  352. # The role of the current user in the ta table must be admin or owner
  353. if current_user.current_tenant.current_role not in ['admin', 'owner']:
  354. raise Forbidden()
  355. key = db.session.query(ApiToken). \
  356. filter(ApiToken.tenant_id == current_user.current_tenant_id, ApiToken.type == self.resource_type,
  357. ApiToken.id == api_key_id). \
  358. first()
  359. if key is None:
  360. flask_restful.abort(404, message='API key not found')
  361. db.session.query(ApiToken).filter(ApiToken.id == api_key_id).delete()
  362. db.session.commit()
  363. return {'result': 'success'}, 204
  364. class DatasetApiBaseUrlApi(Resource):
  365. @setup_required
  366. @login_required
  367. @account_initialization_required
  368. def get(self):
  369. return {
  370. 'api_base_url': (current_app.config['SERVICE_API_URL'] if current_app.config['SERVICE_API_URL']
  371. else request.host_url.rstrip('/')) + '/v1'
  372. }
  373. class DatasetRetrievalSettingApi(Resource):
  374. @setup_required
  375. @login_required
  376. @account_initialization_required
  377. def get(self):
  378. vector_type = current_app.config['VECTOR_STORE']
  379. if vector_type == 'milvus':
  380. return {
  381. 'retrieval_method': [
  382. 'semantic_search'
  383. ]
  384. }
  385. elif vector_type == 'qdrant' or vector_type == 'weaviate':
  386. return {
  387. 'retrieval_method': [
  388. 'semantic_search', 'full_text_search', 'hybrid_search'
  389. ]
  390. }
  391. else:
  392. raise ValueError("Unsupported vector db type.")
  393. class DatasetRetrievalSettingMockApi(Resource):
  394. @setup_required
  395. @login_required
  396. @account_initialization_required
  397. def get(self, vector_type):
  398. if vector_type == 'milvus':
  399. return {
  400. 'retrieval_method': [
  401. 'semantic_search'
  402. ]
  403. }
  404. elif vector_type == 'qdrant' or vector_type == 'weaviate':
  405. return {
  406. 'retrieval_method': [
  407. 'semantic_search', 'full_text_search', 'hybrid_search'
  408. ]
  409. }
  410. else:
  411. raise ValueError("Unsupported vector db type.")
  412. api.add_resource(DatasetListApi, '/datasets')
  413. api.add_resource(DatasetApi, '/datasets/<uuid:dataset_id>')
  414. api.add_resource(DatasetQueryApi, '/datasets/<uuid:dataset_id>/queries')
  415. api.add_resource(DatasetIndexingEstimateApi, '/datasets/indexing-estimate')
  416. api.add_resource(DatasetRelatedAppListApi, '/datasets/<uuid:dataset_id>/related-apps')
  417. api.add_resource(DatasetIndexingStatusApi, '/datasets/<uuid:dataset_id>/indexing-status')
  418. api.add_resource(DatasetApiKeyApi, '/datasets/api-keys')
  419. api.add_resource(DatasetApiDeleteApi, '/datasets/api-keys/<uuid:api_key_id>')
  420. api.add_resource(DatasetApiBaseUrlApi, '/datasets/api-base-info')
  421. api.add_resource(DatasetRetrievalSettingApi, '/datasets/retrieval-setting')
  422. api.add_resource(DatasetRetrievalSettingMockApi, '/datasets/retrieval-setting/<string:vector_type>')