chore(api/controllers): Apply Ruff Formatter. (#7645)

api/controllers/service_api/dataset/document.py

@@ -27,47 +27,40 @@ from services.file_service import FileService
 class DocumentAddByTextApi(DatasetApiResource):
     """Resource for documents."""
 
-    @cloud_edition_billing_resource_check('vector_space', 'dataset')
-    @cloud_edition_billing_resource_check('documents', 'dataset')
+    @cloud_edition_billing_resource_check("vector_space", "dataset")
+    @cloud_edition_billing_resource_check("documents", "dataset")
     def post(self, tenant_id, dataset_id):
         """Create document by text."""
         parser = reqparse.RequestParser()
-        parser.add_argument('name', type=str, required=True, nullable=False, location='json')
-        parser.add_argument('text', type=str, required=True, nullable=False, location='json')
-        parser.add_argument('process_rule', type=dict, required=False, nullable=True, location='json')
-        parser.add_argument('original_document_id', type=str, required=False, location='json')
-        parser.add_argument('doc_form', type=str, default='text_model', required=False, nullable=False, location='json')
-        parser.add_argument('doc_language', type=str, default='English', required=False, nullable=False,
-                            location='json')
-        parser.add_argument('indexing_technique', type=str, choices=Dataset.INDEXING_TECHNIQUE_LIST, nullable=False,
-                            location='json')
-        parser.add_argument('retrieval_model', type=dict, required=False, nullable=False,
-                            location='json')
+        parser.add_argument("name", type=str, required=True, nullable=False, location="json")
+        parser.add_argument("text", type=str, required=True, nullable=False, location="json")
+        parser.add_argument("process_rule", type=dict, required=False, nullable=True, location="json")
+        parser.add_argument("original_document_id", type=str, required=False, location="json")
+        parser.add_argument("doc_form", type=str, default="text_model", required=False, nullable=False, location="json")
+        parser.add_argument(
+            "doc_language", type=str, default="English", required=False, nullable=False, location="json"
+        )
+        parser.add_argument(
+            "indexing_technique", type=str, choices=Dataset.INDEXING_TECHNIQUE_LIST, nullable=False, location="json"
+        )
+        parser.add_argument("retrieval_model", type=dict, required=False, nullable=False, location="json")
         args = parser.parse_args()
         dataset_id = str(dataset_id)
         tenant_id = str(tenant_id)
-        dataset = db.session.query(Dataset).filter(
-            Dataset.tenant_id == tenant_id,
-            Dataset.id == dataset_id
-        ).first()
+        dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first()
 
         if not dataset:
-            raise ValueError('Dataset is not exist.')
+            raise ValueError("Dataset is not exist.")
 
-        if not dataset.indexing_technique and not args['indexing_technique']:
-            raise ValueError('indexing_technique is required.')
+        if not dataset.indexing_technique and not args["indexing_technique"]:
+            raise ValueError("indexing_technique is required.")
 
-        upload_file = FileService.upload_text(args.get('text'), args.get('name'))
+        upload_file = FileService.upload_text(args.get("text"), args.get("name"))
         data_source = {
-            'type': 'upload_file',
-            'info_list': {
-                'data_source_type': 'upload_file',
-                'file_info_list': {
-                    'file_ids': [upload_file.id]
-                }
-            }
+            "type": "upload_file",
+            "info_list": {"data_source_type": "upload_file", "file_info_list": {"file_ids": [upload_file.id]}},
         }
-        args['data_source'] = data_source
+        args["data_source"] = data_source
         # validate args
         DocumentService.document_create_args_validate(args)
 
@@ -76,60 +69,49 @@ class DocumentAddByTextApi(DatasetApiResource):
                 dataset=dataset,
                 document_data=args,
                 account=current_user,
-                dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
-                created_from='api'
+                dataset_process_rule=dataset.latest_process_rule if "process_rule" not in args else None,
+                created_from="api",
             )
         except ProviderTokenNotInitError as ex:
             raise ProviderNotInitializeError(ex.description)
         document = documents[0]
 
-        documents_and_batch_fields = {
-            'document': marshal(document, document_fields),
-            'batch': batch
-        }
+        documents_and_batch_fields = {"document": marshal(document, document_fields), "batch": batch}
         return documents_and_batch_fields, 200
 
 
 class DocumentUpdateByTextApi(DatasetApiResource):
     """Resource for update documents."""
 
-    @cloud_edition_billing_resource_check('vector_space', 'dataset')
+    @cloud_edition_billing_resource_check("vector_space", "dataset")
     def post(self, tenant_id, dataset_id, document_id):
         """Update document by text."""
         parser = reqparse.RequestParser()
-        parser.add_argument('name', type=str, required=False, nullable=True, location='json')
-        parser.add_argument('text', type=str, required=False, nullable=True, location='json')
-        parser.add_argument('process_rule', type=dict, required=False, nullable=True, location='json')
-        parser.add_argument('doc_form', type=str, default='text_model', required=False, nullable=False, location='json')
-        parser.add_argument('doc_language', type=str, default='English', required=False, nullable=False,
-                            location='json')
-        parser.add_argument('retrieval_model', type=dict, required=False, nullable=False,
-                            location='json')
+        parser.add_argument("name", type=str, required=False, nullable=True, location="json")
+        parser.add_argument("text", type=str, required=False, nullable=True, location="json")
+        parser.add_argument("process_rule", type=dict, required=False, nullable=True, location="json")
+        parser.add_argument("doc_form", type=str, default="text_model", required=False, nullable=False, location="json")
+        parser.add_argument(
+            "doc_language", type=str, default="English", required=False, nullable=False, location="json"
+        )
+        parser.add_argument("retrieval_model", type=dict, required=False, nullable=False, location="json")
         args = parser.parse_args()
         dataset_id = str(dataset_id)
         tenant_id = str(tenant_id)
-        dataset = db.session.query(Dataset).filter(
-            Dataset.tenant_id == tenant_id,
-            Dataset.id == dataset_id
-        ).first()
+        dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first()
 
         if not dataset:
-            raise ValueError('Dataset is not exist.')
+            raise ValueError("Dataset is not exist.")
 
-        if args['text']:
-            upload_file = FileService.upload_text(args.get('text'), args.get('name'))
+        if args["text"]:
+            upload_file = FileService.upload_text(args.get("text"), args.get("name"))
             data_source = {
-                'type': 'upload_file',
-                'info_list': {
-                    'data_source_type': 'upload_file',
-                    'file_info_list': {
-                        'file_ids': [upload_file.id]
-                    }
-                }
+                "type": "upload_file",
+                "info_list": {"data_source_type": "upload_file", "file_info_list": {"file_ids": [upload_file.id]}},
             }
-            args['data_source'] = data_source
+            args["data_source"] = data_source
         # validate args
-        args['original_document_id'] = str(document_id)
+        args["original_document_id"] = str(document_id)
         DocumentService.document_create_args_validate(args)
 
         try:
@@ -137,65 +119,53 @@ class DocumentUpdateByTextApi(DatasetApiResource):
                 dataset=dataset,
                 document_data=args,
                 account=current_user,
-                dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
-                created_from='api'
+                dataset_process_rule=dataset.latest_process_rule if "process_rule" not in args else None,
+                created_from="api",
             )
         except ProviderTokenNotInitError as ex:
             raise ProviderNotInitializeError(ex.description)
         document = documents[0]
 
-        documents_and_batch_fields = {
-            'document': marshal(document, document_fields),
-            'batch': batch
-        }
+        documents_and_batch_fields = {"document": marshal(document, document_fields), "batch": batch}
         return documents_and_batch_fields, 200
 
 
 class DocumentAddByFileApi(DatasetApiResource):
     """Resource for documents."""
-    @cloud_edition_billing_resource_check('vector_space', 'dataset')
-    @cloud_edition_billing_resource_check('documents', 'dataset')
+
+    @cloud_edition_billing_resource_check("vector_space", "dataset")
+    @cloud_edition_billing_resource_check("documents", "dataset")
     def post(self, tenant_id, dataset_id):
         """Create document by upload file."""
         args = {}
-        if 'data' in request.form:
-            args = json.loads(request.form['data'])
-        if 'doc_form' not in args:
-            args['doc_form'] = 'text_model'
-        if 'doc_language' not in args:
-            args['doc_language'] = 'English'
+        if "data" in request.form:
+            args = json.loads(request.form["data"])
+        if "doc_form" not in args:
+            args["doc_form"] = "text_model"
+        if "doc_language" not in args:
+            args["doc_language"] = "English"
         # get dataset info
         dataset_id = str(dataset_id)
         tenant_id = str(tenant_id)
-        dataset = db.session.query(Dataset).filter(
-            Dataset.tenant_id == tenant_id,
-            Dataset.id == dataset_id
-        ).first()
+        dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first()
 
         if not dataset:
-            raise ValueError('Dataset is not exist.')
-        if not dataset.indexing_technique and not args.get('indexing_technique'):
-            raise ValueError('indexing_technique is required.')
+            raise ValueError("Dataset is not exist.")
+        if not dataset.indexing_technique and not args.get("indexing_technique"):
+            raise ValueError("indexing_technique is required.")
 
         # save file info
-        file = request.files['file']
+        file = request.files["file"]
         # check file
-        if 'file' not in request.files:
+        if "file" not in request.files:
             raise NoFileUploadedError()
 
         if len(request.files) > 1:
             raise TooManyFilesError()
 
         upload_file = FileService.upload_file(file, current_user)
-        data_source = {
-            'type': 'upload_file',
-            'info_list': {
-                'file_info_list': {
-                    'file_ids': [upload_file.id]
-                }
-            }
-        }
-        args['data_source'] = data_source
+        data_source = {"type": "upload_file", "info_list": {"file_info_list": {"file_ids": [upload_file.id]}}}
+        args["data_source"] = data_source
         # validate args
         DocumentService.document_create_args_validate(args)
 
@@ -204,63 +174,49 @@ class DocumentAddByFileApi(DatasetApiResource):
                 dataset=dataset,
                 document_data=args,
                 account=dataset.created_by_account,
-                dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
-                created_from='api'
+                dataset_process_rule=dataset.latest_process_rule if "process_rule" not in args else None,
+                created_from="api",
             )
         except ProviderTokenNotInitError as ex:
             raise ProviderNotInitializeError(ex.description)
         document = documents[0]
-        documents_and_batch_fields = {
-            'document': marshal(document, document_fields),
-            'batch': batch
-        }
+        documents_and_batch_fields = {"document": marshal(document, document_fields), "batch": batch}
         return documents_and_batch_fields, 200
 
 
 class DocumentUpdateByFileApi(DatasetApiResource):
     """Resource for update documents."""
 
-    @cloud_edition_billing_resource_check('vector_space', 'dataset')
+    @cloud_edition_billing_resource_check("vector_space", "dataset")
     def post(self, tenant_id, dataset_id, document_id):
         """Update document by upload file."""
         args = {}
-        if 'data' in request.form:
-            args = json.loads(request.form['data'])
-        if 'doc_form' not in args:
-            args['doc_form'] = 'text_model'
-        if 'doc_language' not in args:
-            args['doc_language'] = 'English'
+        if "data" in request.form:
+            args = json.loads(request.form["data"])
+        if "doc_form" not in args:
+            args["doc_form"] = "text_model"
+        if "doc_language" not in args:
+            args["doc_language"] = "English"
 
         # get dataset info
         dataset_id = str(dataset_id)
         tenant_id = str(tenant_id)
-        dataset = db.session.query(Dataset).filter(
-            Dataset.tenant_id == tenant_id,
-            Dataset.id == dataset_id
-        ).first()
+        dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first()
 
         if not dataset:
-            raise ValueError('Dataset is not exist.')
-        if 'file' in request.files:
+            raise ValueError("Dataset is not exist.")
+        if "file" in request.files:
             # save file info
-            file = request.files['file']
-
+            file = request.files["file"]
 
             if len(request.files) > 1:
                 raise TooManyFilesError()
 
             upload_file = FileService.upload_file(file, current_user)
-            data_source = {
-                'type': 'upload_file',
-                'info_list': {
-                    'file_info_list': {
-                        'file_ids': [upload_file.id]
-                    }
-                }
-            }
-            args['data_source'] = data_source
+            data_source = {"type": "upload_file", "info_list": {"file_info_list": {"file_ids": [upload_file.id]}}}
+            args["data_source"] = data_source
         # validate args
-        args['original_document_id'] = str(document_id)
+        args["original_document_id"] = str(document_id)
         DocumentService.document_create_args_validate(args)
 
         try:
@@ -268,16 +224,13 @@ class DocumentUpdateByFileApi(DatasetApiResource):
                 dataset=dataset,
                 document_data=args,
                 account=dataset.created_by_account,
-                dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
-                created_from='api'
+                dataset_process_rule=dataset.latest_process_rule if "process_rule" not in args else None,
+                created_from="api",
             )
         except ProviderTokenNotInitError as ex:
             raise ProviderNotInitializeError(ex.description)
         document = documents[0]
-        documents_and_batch_fields = {
-            'document': marshal(document, document_fields),
-            'batch': batch
-        }
+        documents_and_batch_fields = {"document": marshal(document, document_fields), "batch": batch}
         return documents_and_batch_fields, 200
 
 
@@ -289,13 +242,10 @@ class DocumentDeleteApi(DatasetApiResource):
         tenant_id = str(tenant_id)
 
         # get dataset info
-        dataset = db.session.query(Dataset).filter(
-            Dataset.tenant_id == tenant_id,
-            Dataset.id == dataset_id
-        ).first()
+        dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first()
 
         if not dataset:
-            raise ValueError('Dataset is not exist.')
+            raise ValueError("Dataset is not exist.")
 
         document = DocumentService.get_document(dataset.id, document_id)
 
@@ -311,44 +261,39 @@ class DocumentDeleteApi(DatasetApiResource):
             # delete document
             DocumentService.delete_document(document)
         except services.errors.document.DocumentIndexingError:
-            raise DocumentIndexingError('Cannot delete document during indexing.')
+            raise DocumentIndexingError("Cannot delete document during indexing.")
 
-        return {'result': 'success'}, 200
+        return {"result": "success"}, 200
 
 
 class DocumentListApi(DatasetApiResource):
     def get(self, tenant_id, dataset_id):
         dataset_id = str(dataset_id)
         tenant_id = str(tenant_id)
-        page = request.args.get('page', default=1, type=int)
-        limit = request.args.get('limit', default=20, type=int)
-        search = request.args.get('keyword', default=None, type=str)
-        dataset = db.session.query(Dataset).filter(
-            Dataset.tenant_id == tenant_id,
-            Dataset.id == dataset_id
-        ).first()
+        page = request.args.get("page", default=1, type=int)
+        limit = request.args.get("limit", default=20, type=int)
+        search = request.args.get("keyword", default=None, type=str)
+        dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first()
         if not dataset:
-            raise NotFound('Dataset not found.')
+            raise NotFound("Dataset not found.")
 
-        query = Document.query.filter_by(
-            dataset_id=str(dataset_id), tenant_id=tenant_id)
+        query = Document.query.filter_by(dataset_id=str(dataset_id), tenant_id=tenant_id)
 
         if search:
-            search = f'%{search}%'
+            search = f"%{search}%"
             query = query.filter(Document.name.like(search))
 
         query = query.order_by(desc(Document.created_at))
 
-        paginated_documents = query.paginate(
-            page=page, per_page=limit, max_per_page=100, error_out=False)
+        paginated_documents = query.paginate(page=page, per_page=limit, max_per_page=100, error_out=False)
         documents = paginated_documents.items
 
         response = {
-            'data': marshal(documents, document_fields),
-            'has_more': len(documents) == limit,
-            'limit': limit,
-            'total': paginated_documents.total,
-            'page': page
+            "data": marshal(documents, document_fields),
+            "has_more": len(documents) == limit,
+            "limit": limit,
+            "total": paginated_documents.total,
+            "page": page,
         }
 
         return response
@@ -360,38 +305,36 @@ class DocumentIndexingStatusApi(DatasetApiResource):
         batch = str(batch)
         tenant_id = str(tenant_id)
         # get dataset
-        dataset = db.session.query(Dataset).filter(
-            Dataset.tenant_id == tenant_id,
-            Dataset.id == dataset_id
-        ).first()
+        dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first()
         if not dataset:
-            raise NotFound('Dataset not found.')
+            raise NotFound("Dataset not found.")
         # get documents
         documents = DocumentService.get_batch_documents(dataset_id, batch)
         if not documents:
-            raise NotFound('Documents not found.')
+            raise NotFound("Documents not found.")
         documents_status = []
         for document in documents:
-            completed_segments = DocumentSegment.query.filter(DocumentSegment.completed_at.isnot(None),
-                                                              DocumentSegment.document_id == str(document.id),
-                                                              DocumentSegment.status != 're_segment').count()
-            total_segments = DocumentSegment.query.filter(DocumentSegment.document_id == str(document.id),
-                                                          DocumentSegment.status != 're_segment').count()
+            completed_segments = DocumentSegment.query.filter(
+                DocumentSegment.completed_at.isnot(None),
+                DocumentSegment.document_id == str(document.id),
+                DocumentSegment.status != "re_segment",
+            ).count()
+            total_segments = DocumentSegment.query.filter(
+                DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment"
+            ).count()
             document.completed_segments = completed_segments
             document.total_segments = total_segments
             if document.is_paused:
-                document.indexing_status = 'paused'
+                document.indexing_status = "paused"
             documents_status.append(marshal(document, document_status_fields))
-        data = {
-            'data': documents_status
-        }
+        data = {"data": documents_status}
         return data
 
 
-api.add_resource(DocumentAddByTextApi, '/datasets/<uuid:dataset_id>/document/create_by_text')
-api.add_resource(DocumentAddByFileApi, '/datasets/<uuid:dataset_id>/document/create_by_file')
-api.add_resource(DocumentUpdateByTextApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/update_by_text')
-api.add_resource(DocumentUpdateByFileApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/update_by_file')
-api.add_resource(DocumentDeleteApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>')
-api.add_resource(DocumentListApi, '/datasets/<uuid:dataset_id>/documents')
-api.add_resource(DocumentIndexingStatusApi, '/datasets/<uuid:dataset_id>/documents/<string:batch>/indexing-status')
+api.add_resource(DocumentAddByTextApi, "/datasets/<uuid:dataset_id>/document/create_by_text")
+api.add_resource(DocumentAddByFileApi, "/datasets/<uuid:dataset_id>/document/create_by_file")
+api.add_resource(DocumentUpdateByTextApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/update_by_text")
+api.add_resource(DocumentUpdateByFileApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/update_by_file")
+api.add_resource(DocumentDeleteApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>")
+api.add_resource(DocumentListApi, "/datasets/<uuid:dataset_id>/documents")
+api.add_resource(DocumentIndexingStatusApi, "/datasets/<uuid:dataset_id>/documents/<string:batch>/indexing-status")