diff --git a/api/fields/dataset_fields.py b/api/fields/dataset_fields.py index 2871b3ec16..aa147331d4 100644 --- a/api/fields/dataset_fields.py +++ b/api/fields/dataset_fields.py @@ -88,6 +88,8 @@ dataset_detail_fields = { "chunk_structure": fields.String, "icon_info": fields.Nested(icon_info_fields), "is_published": fields.Boolean, + "total_documents": fields.Integer, + "total_available_documents": fields.Integer, } dataset_query_detail_fields = { diff --git a/api/models/dataset.py b/api/models/dataset.py index d2fdd5e900..6795f719df 100644 --- a/api/models/dataset.py +++ b/api/models/dataset.py @@ -312,6 +312,19 @@ class DatasetProcessRule(Base): except JSONDecodeError: return None + @property + def total_documents(self): + return db.session.query(func.count(Document.id)).filter(Document.dataset_id == self.dataset_id).scalar() + + @property + def total_available_documents(self): + return db.session.query(func.count(Document.id)).filter( + Document.dataset_id == self.dataset_id, + Document.indexing_status == "completed", + Document.enabled == True, + Document.archived == False, + ).scalar() + class Document(Base): __tablename__ = "documents"