yujunjun/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2025-11-29 16:36:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf: add get_all_kb_doc_count func to simplify kb.doc_num updating (#10169)

Browse Source 
### 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
This commit is contained in:
He Wang 2025-09-19 19:11:50 +08:00 committed by GitHub
parent 70ce02faf4
commit 7ccca2143c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 12 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
api/db/init_data.py
View File

@ -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))

12
api/db/services/document_service.py
View File

@ -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()