From cfb6700dcfeafc9c291b45650c9910ea6f184702 Mon Sep 17 00:00:00 2001
From: Wenjie Zhang <xerrors@163.com>
Date: Tue, 24 Jun 2025 16:02:56 +0800
Subject: [PATCH] =?UTF-8?q?refactor:=20=E4=BC=98=E5=8C=96=E7=9F=A5?=
 =?UTF-8?q?=E8=AF=86=E5=BA=93=E8=8E=B7=E5=8F=96=E9=80=BB=E8=BE=91=EF=BC=8C?=
 =?UTF-8?q?=E6=8F=90=E5=8D=87=E6=80=A7=E8=83=BD=E5=92=8C=E6=95=B0=E6=8D=AE?=
 =?UTF-8?q?=E5=A4=84=E7=90=86=E6=95=88=E7=8E=87?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 修改 get_all_databases 方法，避免预加载文件数据，直接构造返回数据，减少查询时间。
- 批量获取集合信息，减少对 Milvus 的单独查询，提高数据获取效率。
---
 src/core/knowledgebase.py | 72 ++++++++++++++++++++++++++++++---------
 1 file changed, 56 insertions(+), 16 deletions(-)

diff --git a/src/core/knowledgebase.py b/src/core/knowledgebase.py
index a3482272..787020c2 100644
--- a/src/core/knowledgebase.py
+++ b/src/core/knowledgebase.py
@@ -106,10 +106,39 @@ class KnowledgeBase:
     def get_all_databases(self):
         """获取所有知识库"""
         with db_manager.get_session_context() as session:
-            databases = session.query(KnowledgeDatabase).options(
-                joinedload(KnowledgeDatabase.files)
-            ).all()
-            return [db.to_dict() for db in databases] # Assuming to_dict handles files correctly
+            # 只查询数据库基本信息，不预加载文件数据
+            databases = session.query(KnowledgeDatabase).all()
+            result = []
+
+            for db in databases:
+                # 直接构造返回数据，避免调用 to_dict 方法中的关联查询
+                db_dict = {
+                    "id": db.id,
+                    "db_id": db.db_id,
+                    "name": db.name,
+                    "description": db.description,
+                    "embed_model": db.embed_model,
+                    "dimension": db.dimension,
+                    "metadata": db.meta_info or {},
+                    "created_at": db.created_at.isoformat() if db.created_at else None
+                }
+
+                # 只查询文件的基本信息，不加载节点数据
+                files = session.query(KnowledgeFile).filter_by(database_id=db.db_id).all()
+                db_dict["files"] = {}
+                for file_obj in files:
+                    db_dict["files"][file_obj.file_id] = {
+                        "file_id": file_obj.file_id,
+                        "filename": file_obj.filename,
+                        "path": file_obj.path,
+                        "type": file_obj.file_type,
+                        "status": file_obj.status,
+                        "created_at": file_obj.created_at.timestamp() if file_obj.created_at else time.time()
+                    }
+
+                result.append(db_dict)
+
+            return result
 
     def get_database_by_id(self, db_id):
         """根据ID获取知识库"""
@@ -310,19 +339,30 @@ class KnowledgeBase:
         assert config.enable_knowledge_base, "知识库未启用"
         databases = self.get_all_databases()
         databases_with_milvus = []
-        for db_data in databases: # db_data is already a dict from to_dict()
-            db_copy = db_data.copy()
-            try:
-                milvus_info = self.get_collection_info(db_copy["db_id"])
-                # Merge Milvus info carefully, avoid overwriting existing keys like 'name', 'description'
-                for k, v in milvus_info.items():
-                    if k not in db_copy or k in ["row_count", "status", "error_message"]: # Milvus specific keys
-                        db_copy[k] = v
-            except Exception as e:
-                logger.warning(f"获取知识库 {db_copy.get('name')} (ID: {db_copy.get('db_id')}) 的Milvus信息失败: {e}")
-                db_copy.update({"row_count": 0, "status": "未连接", "error": str(e)})
 
-            # files should be part of db_copy from to_dict()
+        # 批量获取所有集合信息，避免逐个查询
+        try:
+            all_collections = self.get_collections()
+            collections_dict = {col.get("collection_name", col.get("name")): col for col in all_collections}
+        except Exception as e:
+            logger.warning(f"批量获取Milvus集合信息失败: {e}")
+            collections_dict = {}
+
+        for db_data in databases:
+            db_copy = db_data.copy()
+
+            # 从缓存的集合信息中获取数据，而不是单独查询
+            if db_copy["db_id"] in collections_dict:
+                milvus_info = collections_dict[db_copy["db_id"]]
+                # 只添加Milvus特定的字段，避免覆盖数据库基本信息
+                for k, v in milvus_info.items():
+                    if k in ["row_count", "status", "error_message"]:
+                        db_copy[k] = v
+            else:
+                # 如果集合不存在，设置默认值
+                db_copy.update({"row_count": 0, "status": "未连接"})
+
+            # 统计处理中的文件数量
             db_copy_files = db_copy.get("files", {}).values()
             processing_files_count = sum(1 for file_info in db_copy_files if file_info.get("status") in ["processing", "waiting"])
             if processing_files_count > 0: