From 594ccff9c82f4838faa55c3ab4a8f01747deb01b Mon Sep 17 00:00:00 2001
From: lxfight <1686540385@qq.com>
Date: Sat, 25 Oct 2025 11:56:37 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E6=B7=BB=E5=8A=A0=E6=95=B0=E6=8D=AE?=
 =?UTF-8?q?=E5=BA=93=E8=BF=9E=E6=8E=A5=E6=A3=80=E6=9F=A5=E5=92=8C=E7=9F=A5?=
 =?UTF-8?q?=E8=AF=86=E5=BA=93=E7=BB=88=E6=AD=A2=E5=8A=9F=E8=83=BD=EF=BC=8C?=
 =?UTF-8?q?=E5=A2=9E=E5=BC=BA=E9=94=99=E8=AF=AF=E5=A4=84=E7=90=86=E5=92=8C?=
 =?UTF-8?q?=E6=B8=85=E7=90=86=E9=80=BB=E8=BE=91=EF=BC=8C=E4=BF=AE=E5=A4=8D?=
 =?UTF-8?q?=E7=9F=A5=E8=AF=86=E5=BA=93=E6=97=A0=E6=B3=95=E5=88=A0=E9=99=A4?=
 =?UTF-8?q?=E7=9A=84=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../db/vec_db/faiss_impl/document_storage.py    | 17 ++++++++++++++---
 astrbot/core/knowledge_base/kb_helper.py        |  8 ++++----
 astrbot/core/knowledge_base/kb_mgr.py           | 17 +++++++++++++++++
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/astrbot/core/db/vec_db/faiss_impl/document_storage.py b/astrbot/core/db/vec_db/faiss_impl/document_storage.py
index 1feeb9b92..265c0cc43 100644
--- a/astrbot/core/db/vec_db/faiss_impl/document_storage.py
+++ b/astrbot/core/db/vec_db/faiss_impl/document_storage.py
@@ -7,6 +7,7 @@ from sqlalchemy import Text, Column
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
 from sqlalchemy.orm import sessionmaker
 from sqlmodel import Field, SQLModel, select, col, func, text, MetaData
+from astrbot.core import logger
 
 
 class BaseDocModel(SQLModel, table=False):
@@ -113,7 +114,11 @@ class DocumentStorage:
         Returns:
             list: The list of documents that match the filters.
         """
-        assert self.engine is not None, "Database connection is not initialized."
+        if self.engine is None:
+            logger.warning(
+                "Database connection is not initialized, returning empty result"
+            )
+            return []
 
         async with self.get_session() as session:
             query = select(Document)
@@ -261,7 +266,11 @@ class DocumentStorage:
         Args:
             metadata_filters (dict): The metadata filters to apply.
         """
-        assert self.engine is not None, "Database connection is not initialized."
+        if self.engine is None:
+            logger.warning(
+                "Database connection is not initialized, skipping delete operation"
+            )
+            return
 
         async with self.get_session() as session:
             async with session.begin():
@@ -287,7 +296,9 @@ class DocumentStorage:
         Returns:
             int: The count of documents.
         """
-        assert self.engine is not None, "Database connection is not initialized."
+        if self.engine is None:
+            logger.warning("Database connection is not initialized, returning 0")
+            return 0
 
         async with self.get_session() as session:
             query = select(func.count(col(Document.id)))
diff --git a/astrbot/core/knowledge_base/kb_helper.py b/astrbot/core/knowledge_base/kb_helper.py
index dc9febea9..5c07c93b2 100644
--- a/astrbot/core/knowledge_base/kb_helper.py
+++ b/astrbot/core/knowledge_base/kb_helper.py
@@ -85,12 +85,12 @@ class KBHelper:
         return vec_db
 
     async def delete_vec_db(self):
+        """删除知识库的向量数据库和所有相关文件"""
+        import shutil
+
         await self.terminate()
         if self.kb_dir.exists():
-            for item in self.kb_dir.iterdir():
-                if item.is_file():
-                    item.unlink()
-            self.kb_dir.rmdir()
+            shutil.rmtree(self.kb_dir)
 
     async def terminate(self):
         if self.vec_db:
diff --git a/astrbot/core/knowledge_base/kb_mgr.py b/astrbot/core/knowledge_base/kb_mgr.py
index 0a9c16ad9..606066e90 100644
--- a/astrbot/core/knowledge_base/kb_mgr.py
+++ b/astrbot/core/knowledge_base/kb_mgr.py
@@ -277,3 +277,20 @@ class KnowledgeBaseManager:
             lines.append("")
 
         return "\n".join(lines)
+
+    async def terminate(self):
+        """终止所有知识库实例,关闭数据库连接"""
+        for kb_id, kb_helper in self.kb_insts.items():
+            try:
+                await kb_helper.terminate()
+            except Exception as e:
+                logger.error(f"关闭知识库 {kb_id} 失败: {e}")
+
+        self.kb_insts.clear()
+
+        # 关闭元数据数据库
+        if hasattr(self, "kb_db") and self.kb_db:
+            try:
+                await self.kb_db.close()
+            except Exception as e:
+                logger.error(f"关闭知识库元数据数据库失败: {e}")