From 7ccca2143c0e69b65ff899d43f957ac4ff88f922 Mon Sep 17 00:00:00 2001 From: He Wang Date: Fri, 19 Sep 2025 19:11:50 +0800 Subject: [PATCH] perf: add get_all_kb_doc_count func to simplify kb.doc_num updating (#10169) ### What problem does this PR solve? Add get_all_kb_doc_count func to simplify kb.doc_num updating. ### Type of change - [x] Performance Improvement --- api/db/init_data.py | 3 ++- api/db/services/document_service.py | 12 ++++++++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/api/db/init_data.py b/api/db/init_data.py index d462a7b2b..166cf3bac 100644 --- a/api/db/init_data.py +++ b/api/db/init_data.py @@ -144,8 +144,9 @@ def init_llm_factory(): except Exception: pass break + doc_count = DocumentService.get_all_kb_doc_count() for kb_id in KnowledgebaseService.get_all_ids(): - KnowledgebaseService.update_document_number_in_init(kb_id=kb_id, doc_num=DocumentService.get_kb_doc_count(kb_id)) + KnowledgebaseService.update_document_number_in_init(kb_id=kb_id, doc_num=doc_count.get(kb_id, 0)) diff --git a/api/db/services/document_service.py b/api/db/services/document_service.py index 23eef474f..ffb373d96 100644 --- a/api/db/services/document_service.py +++ b/api/db/services/document_service.py @@ -660,8 +660,16 @@ class DocumentService(CommonService): @classmethod @DB.connection_context() def get_kb_doc_count(cls, kb_id): - return len(cls.model.select(cls.model.id).where( - cls.model.kb_id == kb_id).dicts()) + return cls.model.select().where(cls.model.kb_id == kb_id).count() + + @classmethod + @DB.connection_context() + def get_all_kb_doc_count(cls): + result = {} + rows = cls.model.select(cls.model.kb_id, fn.COUNT(cls.model.id).alias('count')).group_by(cls.model.kb_id) + for row in rows: + result[row.kb_id] = row.count + return result @classmethod @DB.connection_context()