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}")