From e7e4aff667d284fda3a792decf99bc93b31cc8c6 Mon Sep 17 00:00:00 2001
From: Shinexy <643854343@qq.com>
Date: Fri, 8 Dec 2023 14:07:51 +0800
Subject: [PATCH] fix: optimize api /knowledge/space/list (#906)

---
 pilot/server/knowledge/document_db.py | 25 +++++++++++++++++++++++++
 pilot/server/knowledge/service.py     | 10 ++++++----
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/pilot/server/knowledge/document_db.py b/pilot/server/knowledge/document_db.py
index bbe1426d7..d2fc1f582 100644
--- a/pilot/server/knowledge/document_db.py
+++ b/pilot/server/knowledge/document_db.py
@@ -135,6 +135,31 @@ def get_documents(self, query):
         session.close()
         return result
 
+    def get_knowledge_documents_count_bulk(self, space_names):
+        session = self.get_session()
+        """
+        Perform a batch query to count the number of documents for each knowledge space.
+
+        Args:
+            space_names: A list of knowledge space names to query for document counts.
+            session: A SQLAlchemy session object.
+
+        Returns:
+            A dictionary mapping each space name to its document count.
+        """
+        counts_query = (
+            session.query(
+                KnowledgeDocumentEntity.space,
+                func.count(KnowledgeDocumentEntity.id).label("document_count"),
+            )
+            .filter(KnowledgeDocumentEntity.space.in_(space_names))
+            .group_by(KnowledgeDocumentEntity.space)
+        )
+
+        results = counts_query.all()
+        docs_count = {result.space: result.document_count for result in results}
+        return docs_count
+
     def get_knowledge_documents_count(self, query):
         session = self.get_session()
         knowledge_documents = session.query(func.count(KnowledgeDocumentEntity.id))
diff --git a/pilot/server/knowledge/service.py b/pilot/server/knowledge/service.py
index a42396582..b6e844e42 100644
--- a/pilot/server/knowledge/service.py
+++ b/pilot/server/knowledge/service.py
@@ -116,8 +116,12 @@ def get_knowledge_space(self, request: KnowledgeSpaceRequest):
         query = KnowledgeSpaceEntity(
             name=request.name, vector_type=request.vector_type, owner=request.owner
         )
-        responses = []
         spaces = knowledge_space_dao.get_knowledge_space(query)
+        space_names = [space.name for space in spaces]
+        docs_count = knowledge_document_dao.get_knowledge_documents_count_bulk(
+            space_names
+        )
+        responses = []
         for space in spaces:
             res = SpaceQueryResponse()
             res.id = space.id
@@ -128,9 +132,7 @@ def get_knowledge_space(self, request: KnowledgeSpaceRequest):
             res.gmt_created = space.gmt_created
             res.gmt_modified = space.gmt_modified
             res.context = space.context
-            query = KnowledgeDocumentEntity(space=space.name)
-            doc_count = knowledge_document_dao.get_knowledge_documents_count(query)
-            res.docs = doc_count
+            res.docs = docs_count.get(space.name, 0)
             responses.append(res)
         return responses