merge file and notion indexing-estimate

api/controllers/console/datasets/datasets.py

@@ -13,7 +13,7 @@ from core.indexing_runner import IndexingRunner
 from libs.helper import TimestampField
 from extensions.ext_database import db
 from models.model import UploadFile
-from services.dataset_service import DatasetService
+from services.dataset_service import DatasetService, DocumentService
 
 dataset_detail_fields = {
     'id': fields.String,
@@ -217,17 +217,30 @@ class DatasetIndexingEstimateApi(Resource):
     @login_required
     @account_initialization_required
     def post(self):
-        segment_rule = request.get_json()
-        file_detail = db.session.query(UploadFile).filter(
-            UploadFile.tenant_id == current_user.current_tenant_id,
-            UploadFile.id == segment_rule["file_id"]
-        ).first()
-
-        if file_detail is None:
-            raise NotFound("File not found.")
-
-        indexing_runner = IndexingRunner()
-        response = indexing_runner.indexing_estimate(file_detail, segment_rule['process_rule'])
+        parser = reqparse.RequestParser()
+        parser.add_argument('info_list', type=list, required=True, nullable=True, location='json')
+        parser.add_argument('process_rule', type=dict, required=True, nullable=True, location='json')
+        args = parser.parse_args()
+        # validate args
+        DocumentService.estimate_args_validate(args)
+        if args['info_list']['data_source_type'] == 'upload_file':
+            file_details = db.session.query(UploadFile).filter(
+                UploadFile.tenant_id == current_user.current_tenant_id,
+                UploadFile.id in args['info_list']['file_info_list']['file_ids']
+            ).all()
+
+            if file_details is None:
+                raise NotFound("File not found.")
+
+            indexing_runner = IndexingRunner()
+            response = indexing_runner.file_indexing_estimate(file_details, args['process_rule'])
+        elif args['info_list']['data_source_type'] == 'notion_import':
+
+            indexing_runner = IndexingRunner()
+            response = indexing_runner.notion_indexing_estimate(args['info_list']['notion_info_list'],
+                                                                args['process_rule'])
+        else:
+            raise ValueError('Data source type not support')
         return response, 200
 
 
@@ -277,5 +290,5 @@ class DatasetRelatedAppListApi(Resource):
 api.add_resource(DatasetListApi, '/datasets')
 api.add_resource(DatasetApi, '/datasets/<uuid:dataset_id>')
 api.add_resource(DatasetQueryApi, '/datasets/<uuid:dataset_id>/queries')
-api.add_resource(DatasetIndexingEstimateApi, '/datasets/file-indexing-estimate')
+api.add_resource(DatasetIndexingEstimateApi, '/datasets/indexing-estimate')
 api.add_resource(DatasetRelatedAppListApi, '/datasets/<uuid:dataset_id>/related-apps')

api/controllers/console/datasets/datasets_document.py

@@ -335,7 +335,8 @@ class DocumentIndexingEstimateApi(DocumentResource):
                     raise NotFound('File not found.')
 
                 indexing_runner = IndexingRunner()
-                response = indexing_runner.indexing_estimate(file, data_process_rule_dict)
+
+                response = indexing_runner.file_indexing_estimate(list(file), data_process_rule_dict)
 
         return response
 

api/core/indexing_runner.py

@@ -168,38 +168,40 @@ class IndexingRunner:
             nodes=nodes
         )
 
-    def indexing_estimate(self, file_detail: UploadFile, tmp_processing_rule: dict) -> dict:
+    def file_indexing_estimate(self, file_details: List[UploadFile], tmp_processing_rule: dict) -> dict:
         """
         Estimate the indexing for the document.
         """
-        # load data from file
-        text_docs = self._load_data_from_file(file_detail)
-
-        processing_rule = DatasetProcessRule(
-            mode=tmp_processing_rule["mode"],
-            rules=json.dumps(tmp_processing_rule["rules"])
-        )
+        tokens = 0
+        preview_texts = []
+        total_segments = 0
+        for file_detail in file_details:
+            # load data from file
+            text_docs = self._load_data_from_file(file_detail)
 
-        # get node parser for splitting
-        node_parser = self._get_node_parser(processing_rule)
+            processing_rule = DatasetProcessRule(
+                mode=tmp_processing_rule["mode"],
+                rules=json.dumps(tmp_processing_rule["rules"])
+            )
 
-        # split to nodes
-        nodes = self._split_to_nodes(
-            text_docs=text_docs,
-            node_parser=node_parser,
-            processing_rule=processing_rule
-        )
+            # get node parser for splitting
+            node_parser = self._get_node_parser(processing_rule)
 
-        tokens = 0
-        preview_texts = []
-        for node in nodes:
-            if len(preview_texts) < 5:
-                preview_texts.append(node.get_text())
+            # split to nodes
+            nodes = self._split_to_nodes(
+                text_docs=text_docs,
+                node_parser=node_parser,
+                processing_rule=processing_rule
+            )
+            total_segments += len(nodes)
+            for node in nodes:
+                if len(preview_texts) < 5:
+                    preview_texts.append(node.get_text())
 
-            tokens += TokenCalculator.get_num_tokens(self.embedding_model_name, node.get_text())
+                tokens += TokenCalculator.get_num_tokens(self.embedding_model_name, node.get_text())
 
         return {
-            "total_segments": len(nodes),
+            "total_segments": total_segments,
             "tokens": tokens,
             "total_price": '{:f}'.format(TokenCalculator.get_token_price(self.embedding_model_name, tokens)),
             "currency": TokenCalculator.get_currency(self.embedding_model_name),

api/services/dataset_service.py

@@ -707,11 +707,11 @@ class DocumentService:
                 raise ValueError("Process rule segmentation max_tokens is invalid")
 
     @classmethod
-    def notion_estimate_args_validate(cls, args: dict):
-        if 'notion_info_list' not in args or not args['notion_info_list']:
-            raise ValueError("Notion info is required")
+    def estimate_args_validate(cls, args: dict):
+        if 'info_list' not in args or not args['info_list']:
+            raise ValueError("Data source info is required")
 
-        if not isinstance(args['notion_info_list'], list):
+        if not isinstance(args['info_list'], dict):
             raise ValueError("Notion info is invalid")
 
         if 'process_rule' not in args or not args['process_rule']: