dataset optimization

api/controllers/console/datasets/datasets.py

@@ -48,11 +48,13 @@ class DatasetListApi(Resource):
         limit = request.args.get('limit', default=20, type=int)
         ids = request.args.getlist('ids')
         provider = request.args.get('provider', default="vendor")
+        search = request.args.get('keyword', default=None, type=str)
+
         if ids:
             datasets, total = DatasetService.get_datasets_by_ids(ids, current_user.current_tenant_id)
         else:
             datasets, total = DatasetService.get_datasets(page, limit, provider,
-                                                          current_user.current_tenant_id, current_user)
+                                                          current_user.current_tenant_id, current_user, search)
 
         # check embedding setting
         provider_manager = ProviderManager()

api/controllers/service_api/dataset/dataset.py

@@ -26,8 +26,10 @@ class DatasetApi(DatasetApiResource):
         page = request.args.get('page', default=1, type=int)
         limit = request.args.get('limit', default=20, type=int)
         provider = request.args.get('provider', default="vendor")
+        search = request.args.get('keyword', default=None, type=str)
+
         datasets, total = DatasetService.get_datasets(page, limit, provider,
-                                                      tenant_id, current_user)
+                                                      tenant_id, current_user, search)
         # check embedding setting
         provider_manager = ProviderManager()
         configurations = provider_manager.get_configurations(

api/services/dataset_service.py

@@ -51,16 +51,18 @@ from tasks.retry_document_indexing_task import retry_document_indexing_task
 class DatasetService:
 
     @staticmethod
-    def get_datasets(page, per_page, provider="vendor", tenant_id=None, user=None):
+    def get_datasets(page, per_page, provider="vendor", tenant_id=None, user=None, search=None):
         if user:
             permission_filter = db.or_(Dataset.created_by == user.id,
                                        Dataset.permission == 'all_team_members')
         else:
             permission_filter = Dataset.permission == 'all_team_members'
-        datasets = Dataset.query.filter(
+        query = Dataset.query.filter(
             db.and_(Dataset.provider == provider, Dataset.tenant_id == tenant_id, permission_filter)) \
-            .order_by(Dataset.created_at.desc()) \
-            .paginate(
+            .order_by(Dataset.created_at.desc())
+        if search:
+            query = query.filter(db.and_(Dataset.name.ilike(f'%{search}%')))
+        datasets = query.paginate(
                     page=page,
                     per_page=per_page,
                     max_per_page=100,

api/tasks/retry_document_indexing_task.py

@@ -85,5 +85,4 @@ def retry_document_indexing_task(dataset_id: str, document_id: str):
     except Exception as ex:
         logging.info(click.style(str(ex), fg='yellow'))
         redis_client.delete(retry_indexing_cache_key)
-
         pass