fix: optimize api /knowledge/space/list (#906)

pilot/server/knowledge/document_db.py

@@ -135,6 +135,31 @@ class KnowledgeDocumentDao(BaseDao):
         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))

pilot/server/knowledge/service.py

@@ -116,8 +116,12 @@ class KnowledgeService:
         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 @@ class KnowledgeService:
             res.gmt_created = space.gmt_created
             res.gmt_modified = space.gmt_modified
             res.context = space.context
-            query = KnowledgeDocumentEntity(space=space.name)
+            res.docs = docs_count.get(space.name, 0)
-            doc_count = knowledge_document_dao.get_knowledge_documents_count(query)
-            res.docs = doc_count
             responses.append(res)
         return responses