refactor: document segment query

api/controllers/console/datasets/datasets_document.py

@@ -170,9 +170,9 @@ class DatasetDocumentListApi(Resource):
             raise Forbidden(str(e))
 
         with Session(db.engine) as session:
-            query = session.execute(
-                select(Document).filter_by(dataset_id=str(dataset_id), tenant_id=current_user.current_tenant_id)
-            ).all()
+            query = session.query(Document).filter_by(
+                dataset_id=str(dataset_id), tenant_id=current_user.current_tenant_id
+            )
 
             if search:
                 search = f"%{search}%"
@@ -186,7 +186,9 @@ class DatasetDocumentListApi(Resource):
 
             if sort == "hit_count":
                 sub_query = (
-                db.select(DocumentSegment.document_id, db.func.sum(DocumentSegment.hit_count).label("total_hit_count"))
+                    db.select(
+                        DocumentSegment.document_id, db.func.sum(DocumentSegment.hit_count).label("total_hit_count")
+                    )
                     .group_by(DocumentSegment.document_id)
                     .subquery()
                 )
@@ -209,7 +211,6 @@ class DatasetDocumentListApi(Resource):
             paginated_documents = query.paginate(page=page, per_page=limit, max_per_page=100, error_out=False)
             documents = paginated_documents.items
             if fetch:
-            with Session(db.engine) as session:
                 for document in documents:
                     completed_segments = (
                         session.query(DocumentSegment)