From aeb1d1946cade60ca2edddeea574cd2525c43fce Mon Sep 17 00:00:00 2001
From: jyong <718720800@qq.com>
Date: Tue, 10 Jun 2025 17:59:14 +0800
Subject: [PATCH] r2

---
 api/fields/dataset_fields.py |  2 ++
 api/models/dataset.py        | 13 +++++++++++++
 2 files changed, 15 insertions(+)

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"