From 9ca0e56a8a45e4040cff6603e8ed7e4bef2e9789 Mon Sep 17 00:00:00 2001 From: jyong <718720800@qq.com> Date: Wed, 11 Sep 2024 16:59:19 +0800 Subject: [PATCH] external dataset binding --- api/controllers/console/datasets/external.py | 40 +++++++++++++++++- api/services/external_knowledge_service.py | 44 ++++++++++++++++++-- 2 files changed, 79 insertions(+), 5 deletions(-) diff --git a/api/controllers/console/datasets/external.py b/api/controllers/console/datasets/external.py index 3d9d99ce98..024ed955d4 100644 --- a/api/controllers/console/datasets/external.py +++ b/api/controllers/console/datasets/external.py @@ -73,10 +73,10 @@ class ExternalApiTemplateListApi(Resource): def post(self): parser = reqparse.RequestParser() parser.add_argument('name', nullable=False, required=True, - help='type is required. Name must be between 1 to 100 characters.', + help='Name is required. Name must be between 1 to 100 characters.', type=_validate_name) parser.add_argument('description', nullable=False, required=True, - help='description is required. Description must be between 1 to 400 characters.', + help='Description is required. Description must be between 1 to 400 characters.', type=_validate_description_length) parser.add_argument('settings', type=list, location='json', nullable=False, @@ -215,6 +215,42 @@ class ExternalDatasetInitApi(Resource): return response +class ExternalDatasetCreateApi(Resource): + + @setup_required + @login_required + @account_initialization_required + def post(self): + # The role of the current user in the ta table must be admin, owner, or editor + if not current_user.is_editor: + raise Forbidden() + + parser = reqparse.RequestParser() + parser.add_argument('api_template_id', type=str, required=True, nullable=False, location='json') + parser.add_argument('external_knowledge_id', type=str, required=True, nullable=False, location='json') + parser.add_argument('name', nullable=False, required=True, + help='name is required. Name must be between 1 to 100 characters.', + type=_validate_name) + parser.add_argument('description', type=str, required=True, nullable=True, location='json') + + args = parser.parse_args() + + # The role of the current user in the ta table must be admin, owner, or editor, or dataset_operator + if not current_user.is_dataset_editor: + raise Forbidden() + + try: + dataset = ExternalDatasetService.create_external_dataset( + tenant_id=current_user.current_tenant_id, + user_id=current_user.id, + args=args, + ) + except services.errors.dataset.DatasetNameDuplicateError: + raise DatasetNameDuplicateError() + + return marshal(dataset, dataset_detail_fields), 201 + + api.add_resource(ExternalApiTemplateListApi, '/datasets/external-api-template') api.add_resource(ExternalApiTemplateApi, '/datasets/external-api-template/') api.add_resource(ExternalApiUseCheckApi, '/datasets/external-api-template//use-check') diff --git a/api/services/external_knowledge_service.py b/api/services/external_knowledge_service.py index fa95d9fded..d221c20744 100644 --- a/api/services/external_knowledge_service.py +++ b/api/services/external_knowledge_service.py @@ -25,6 +25,7 @@ from models.dataset import ( from models.model import UploadFile from services.dataset_service import DocumentService from services.entities.external_knowledge_entities.external_knowledge_entities import Authorization, ApiTemplateSetting +from services.errors.dataset import DatasetNameDuplicateError from tasks.external_document_indexing_task import external_document_indexing_task @@ -56,7 +57,6 @@ class ExternalDatasetService: if 'api_key' not in api_settings and not api_settings['api_key']: raise ValueError('api_key is required') - @staticmethod def create_api_template(tenant_id: str, user_id: str, args: dict) -> ExternalApiTemplates: api_template = ExternalApiTemplates( @@ -221,7 +221,7 @@ class ExternalDatasetService: if headers: headers = deepcopy(headers) else: - headers= {} + headers = {} if authorization.type == 'api-key': if authorization.config is None: raise ValueError('authorization config is required') @@ -243,4 +243,42 @@ class ExternalDatasetService: @staticmethod def get_api_template_settings(settings: dict) -> ApiTemplateSetting: - return ApiTemplateSetting.parse_obj(settings) \ No newline at end of file + return ApiTemplateSetting.parse_obj(settings) + + @staticmethod + def create_external_dataset(tenant_id, user_id, args): + # check if dataset name already exists + if Dataset.query.filter_by(name=args.get('name'), tenant_id=tenant_id).first(): + raise DatasetNameDuplicateError(f"Dataset with name {args.get('name')} already exists.") + api_template = ExternalApiTemplates.query.filter_by( + id=args.get('api_template_id'), + tenant_id=tenant_id + ).first() + if api_template is None: + raise ValueError('api template not found') + + dataset = Dataset( + tenant_id=tenant_id, + name=args.get('name'), + description=args.get('description', ''), + provider='external', + created_by=user_id, + ) + + db.session.add(dataset) + db.session.flush() + + external_knowledge_binding = ExternalKnowledgeBindings( + tenant_id=tenant_id, + dataset_id=dataset.id, + external_api_template_id=args.get('api_template_id'), + external_knowledge_id=args.get('external_knowledge_id'), + created_by=user_id, + ) + db.session.add(external_knowledge_binding) + + db.session.commit() + + return dataset + +