From b781207752c2f882e183e0f11777473883b66000 Mon Sep 17 00:00:00 2001
From: Yongteng Lei <yongtengrey@outlook.com>
Date: Fri, 9 May 2025 11:48:54 +0800
Subject: [PATCH] Feat: KB detail supports document total size (#7546)

### What problem does this PR solve?

Kb detail supports return document total size now.

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
---
 api/apps/kb_app.py                  |  1 +
 api/db/services/document_service.py | 16 ++++++++++++++++
 2 files changed, 17 insertions(+)

diff --git a/api/apps/kb_app.py b/api/apps/kb_app.py
index d3c59c698..4758f8222 100644
--- a/api/apps/kb_app.py
+++ b/api/apps/kb_app.py
@@ -152,6 +152,7 @@ def detail():
         if not kb:
             return get_data_error_result(
                 message="Can't find this knowledgebase!")
+        kb["size"] = DocumentService.get_total_size_by_kb_id(kb_id=kb["id"],keywords="", run_status=[], types=[])
         return get_json_result(data=kb)
     except Exception as e:
         return server_error_response(e)
diff --git a/api/db/services/document_service.py b/api/db/services/document_service.py
index 2367c6bc5..b4fb3e87a 100644
--- a/api/db/services/document_service.py
+++ b/api/db/services/document_service.py
@@ -117,6 +117,22 @@ class DocumentService(CommonService):
 
         return count
 
+    @classmethod
+    @DB.connection_context()
+    def get_total_size_by_kb_id(cls, kb_id, keywords="", run_status=[], types=[]):
+        query = cls.model.select(fn.COALESCE(fn.SUM(cls.model.size), 0)).where(
+            cls.model.kb_id == kb_id
+        )
+
+        if keywords:
+            query = query.where(fn.LOWER(cls.model.name).contains(keywords.lower()))
+        if run_status:
+            query = query.where(cls.model.run.in_(run_status))
+        if types:
+            query = query.where(cls.model.type.in_(types))
+
+        return query.scalar() or 0
+
     @classmethod
     @DB.connection_context()
     def insert(cls, doc):