From 2f130ba00993345c16d988837fb2c16670169a9e Mon Sep 17 00:00:00 2001 From: Soulter <905617992@qq.com> Date: Fri, 24 Oct 2025 13:59:17 +0800 Subject: [PATCH] feat: delete chunk and delete document --- .../db/vec_db/faiss_impl/document_storage.py | 60 ++++++++++++++-- .../db/vec_db/faiss_impl/embedding_storage.py | 17 ++++- astrbot/core/db/vec_db/faiss_impl/vec_db.py | 43 +++++++++-- astrbot/core/knowledge_base/kb_db_sqlite.py | 15 +++- astrbot/core/knowledge_base/kb_helper.py | 26 +++++++ astrbot/core/knowledge_base/kb_mgr.py | 3 +- .../core/knowledge_base/retrieval/manager.py | 9 ++- astrbot/dashboard/routes/knowledge_base.py | 72 +++++++++++++++++-- .../features/knowledge-base/document.json | 5 +- .../views/knowledge-base/DocumentDetail.vue | 31 +++++++- .../components/DocumentsTab.vue | 3 +- 11 files changed, 258 insertions(+), 26 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 7a94b9d91..17c0cb3ae 100644 --- a/astrbot/core/db/vec_db/faiss_impl/document_storage.py +++ b/astrbot/core/db/vec_db/faiss_impl/document_storage.py @@ -14,6 +14,8 @@ class DocumentStorage: """Initialize the SQLite database and create the documents table if it doesn't exist.""" if not os.path.exists(self.db_path): await self.connect() + if not self.connection: + raise RuntimeError("Failed to connect to the database.") async with self.connection.cursor() as cursor: with open(self.sqlite_init_path, "r", encoding="utf-8") as f: sql_script = f.read() @@ -30,8 +32,8 @@ class DocumentStorage: self, metadata_filters: dict, ids: list | None = None, - offset: int = 0, - limit: int = 100, + offset: int | None = 0, + limit: int | None = 100, ) -> list[dict]: """Retrieve documents by metadata filters and ids. @@ -41,6 +43,7 @@ class DocumentStorage: Returns: list: The list of document IDs(primary key, not doc_id) that match the filters. """ + assert self.connection is not None, "Database connection is not initialized." # metadata filter -> SQL WHERE clause where_clauses = [] values = [] @@ -55,8 +58,13 @@ class DocumentStorage: result = [] async with self.connection.cursor() as cursor: - sql = f"SELECT * FROM documents WHERE {where_sql} ORDER BY id LIMIT ? OFFSET ?" - values.extend([limit, offset]) + sql = f"SELECT * FROM documents WHERE {where_sql}" + if limit is not None: + sql += " LIMIT ?" + values.append(limit) + if offset is not None: + sql += " OFFSET ?" + values.append(offset) await cursor.execute(sql, values) for row in await cursor.fetchall(): @@ -72,6 +80,7 @@ class DocumentStorage: Returns: dict: The document data. """ + assert self.connection is not None, "Database connection is not initialized." async with self.connection.cursor() as cursor: await cursor.execute("SELECT * FROM documents WHERE doc_id = ?", (doc_id,)) row = await cursor.fetchone() @@ -87,18 +96,61 @@ class DocumentStorage: doc_id (str): The doc_id. new_text (str): The new text to update the document with. """ + assert self.connection is not None, "Database connection is not initialized." async with self.connection.cursor() as cursor: await cursor.execute( "UPDATE documents SET text = ? WHERE doc_id = ?", (new_text, doc_id) ) await self.connection.commit() + async def delete_documents(self, metadata_filters: dict): + """Delete documents by their metadata filters. + + Args: + metadata_filters (dict): The metadata filters to apply. + """ + assert self.connection is not None, "Database connection is not initialized." + async with self.connection.cursor() as cursor: + where_clauses = [] + values = [] + for key, val in metadata_filters.items(): + where_clauses.append(f"json_extract(metadata, '$.{key}') = ?") + values.append(val) + where_sql = " AND ".join(where_clauses) or "1=1" + await cursor.execute(f"DELETE FROM documents WHERE {where_sql}", values) + await self.connection.commit() + + async def count_documents(self, metadata_filters: dict | None = None) -> int: + """Count documents in the database. + + Args: + metadata_filters (dict | None): Metadata filters to apply. + + Returns: + int: The count of documents. + """ + assert self.connection is not None, "Database connection is not initialized." + async with self.connection.cursor() as cursor: + sql = "SELECT COUNT(*) FROM documents" + values = [] + if metadata_filters: + where_clauses = [] + for key, val in metadata_filters.items(): + where_clauses.append(f"json_extract(metadata, '$.{key}') = ?") + values.append(val) + where_sql = " AND ".join(where_clauses) + sql += f" WHERE {where_sql}" + await cursor.execute(sql, values) + count = await cursor.fetchone() + return count[0] if count else 0 + async def get_user_ids(self) -> list[str]: """Retrieve all user IDs from the documents table. Returns: list: A list of user IDs. """ + assert self.connection is not None, "Database connection is not initialized." async with self.connection.cursor() as cursor: await cursor.execute("SELECT DISTINCT user_id FROM documents") rows = await cursor.fetchall() diff --git a/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py b/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py index 18ee3189c..f29084148 100644 --- a/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py +++ b/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py @@ -9,7 +9,7 @@ import numpy as np class EmbeddingStorage: - def __init__(self, dimension: int, path: str = None): + def __init__(self, dimension: int, path: str | None = None): self.dimension = dimension self.path = path self.index = None @@ -18,7 +18,6 @@ class EmbeddingStorage: else: base_index = faiss.IndexFlatL2(dimension) self.index = faiss.IndexIDMap(base_index) - self.storage = {} async def insert(self, vector: np.ndarray, id: int): """插入向量 @@ -29,12 +28,12 @@ class EmbeddingStorage: Raises: ValueError: 如果向量的维度与存储的维度不匹配 """ + assert self.index is not None, "FAISS index is not initialized." if vector.shape[0] != self.dimension: raise ValueError( f"向量维度不匹配, 期望: {self.dimension}, 实际: {vector.shape[0]}" ) self.index.add_with_ids(vector.reshape(1, -1), np.array([id])) - self.storage[id] = vector await self.save_index() async def search(self, vector: np.ndarray, k: int) -> tuple: @@ -46,10 +45,22 @@ class EmbeddingStorage: Returns: tuple: (距离, 索引) """ + assert self.index is not None, "FAISS index is not initialized." faiss.normalize_L2(vector) distances, indices = self.index.search(vector, k) return distances, indices + async def delete(self, ids: list[int]): + """删除向量 + + Args: + ids (list[int]): 要删除的向量ID列表 + """ + assert self.index is not None, "FAISS index is not initialized." + id_array = np.array(ids, dtype=np.int64) + self.index.remove_ids(id_array) + await self.save_index() + async def save_index(self): """保存索引 diff --git a/astrbot/core/db/vec_db/faiss_impl/vec_db.py b/astrbot/core/db/vec_db/faiss_impl/vec_db.py index 7c2ae1c01..dfbfa3f31 100644 --- a/astrbot/core/db/vec_db/faiss_impl/vec_db.py +++ b/astrbot/core/db/vec_db/faiss_impl/vec_db.py @@ -39,6 +39,9 @@ class FaissVecDB(BaseVecDB): """ 插入一条文本和其对应向量，自动生成 ID 并保持一致性。 """ + assert self.document_storage.connection is not None, ( + "Database connection is not initialized." + ) metadata = metadata or {} str_id = id or str(uuid.uuid4()) # 使用 UUID 作为原始 ID @@ -119,23 +122,49 @@ class FaissVecDB(BaseVecDB): return top_k_results - async def delete(self, doc_id: int): + async def delete(self, doc_id: str): """ - 删除一条文档 + 删除一条文档块（chunk） """ + assert self.document_storage.connection is not None, ( + "Database connection is not initialized." + ) + # 获得对应的 int id + result = await self.document_storage.get_document_by_doc_id(doc_id) + int_id = result["id"] if result else None + if int_id is None: + return await self.document_storage.connection.execute( "DELETE FROM documents WHERE doc_id = ?", (doc_id,) ) + await self.embedding_storage.delete([int_id]) await self.document_storage.connection.commit() async def close(self): await self.document_storage.close() - async def count_documents(self) -> int: + async def count_documents(self, metadata_filter: dict | None = None) -> int: """ 计算文档数量 + + Args: + metadata_filter (dict | None): 元数据过滤器 """ - async with self.document_storage.connection.cursor() as cursor: - await cursor.execute("SELECT COUNT(*) FROM documents") - count = await cursor.fetchone() - return count[0] if count else 0 + assert self.document_storage.connection is not None, ( + "Database connection is not initialized." + ) + count = await self.document_storage.count_documents( + metadata_filters=metadata_filter or {} + ) + return count + + async def delete_documents(self, metadata_filters: dict): + """ + 根据元数据过滤器删除文档 + """ + docs = await self.document_storage.get_documents( + metadata_filters=metadata_filters, offset=None, limit=None + ) + doc_ids: list[int] = [doc["id"] for doc in docs] + await self.embedding_storage.delete(doc_ids) + await self.document_storage.delete_documents(metadata_filters=metadata_filters) diff --git a/astrbot/core/knowledge_base/kb_db_sqlite.py b/astrbot/core/knowledge_base/kb_db_sqlite.py index c49cdf231..77b980a9d 100644 --- a/astrbot/core/knowledge_base/kb_db_sqlite.py +++ b/astrbot/core/knowledge_base/kb_db_sqlite.py @@ -2,7 +2,7 @@ from contextlib import asynccontextmanager from pathlib import Path from sqlmodel import SQLModel, col, desc -from sqlalchemy import text, func, select, update +from sqlalchemy import text, func, select, update, delete from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from astrbot.core import logger @@ -250,6 +250,19 @@ class KBSQLiteDatabase: "knowledge_base": row[1], } + async def delete_document_by_id(self, doc_id: str, vec_db: FaissVecDB): + """删除单个文档及其相关数据""" + # 在知识库表中删除 + async with self.get_db() as session: + async with session.begin(): + # 删除文档记录 + delete_stmt = delete(KBDocument).where(col(KBDocument.doc_id) == doc_id) + await session.execute(delete_stmt) + await session.commit() + + # 在 vec db 中删除相关向量 + await vec_db.delete_documents(metadata_filters={"doc_id": doc_id}) + # ===== 多媒体查询 ===== async def list_media_by_doc(self, doc_id: str) -> list[KBMedia]: diff --git a/astrbot/core/knowledge_base/kb_helper.py b/astrbot/core/knowledge_base/kb_helper.py index 2a9229318..9203c7ed1 100644 --- a/astrbot/core/knowledge_base/kb_helper.py +++ b/astrbot/core/knowledge_base/kb_helper.py @@ -206,6 +206,26 @@ class KBHelper: doc = await self.kb_db.get_document_by_id(doc_id) return doc + async def delete_document(self, doc_id: str): + """删除单个文档及其相关数据""" + await self.kb_db.delete_document_by_id( + doc_id=doc_id, + vec_db=self.vec_db, # type: ignore + ) + await self.kb_db.update_kb_stats( + kb_id=self.kb.kb_id, + vec_db=self.vec_db, # type: ignore + ) + + async def delete_chunk(self, chunk_id: str): + """删除单个文本块及其相关数据""" + vec_db: FaissVecDB = self.vec_db # type: ignore + await vec_db.delete(chunk_id) + await self.kb_db.update_kb_stats( + kb_id=self.kb.kb_id, + vec_db=self.vec_db, # type: ignore + ) + async def get_chunks_by_doc_id( self, doc_id: str, offset: int = 0, limit: int = 100 ) -> list[dict]: @@ -229,6 +249,12 @@ class KBHelper: ) return result + async def get_chunk_count_by_doc_id(self, doc_id: str) -> int: + """获取文档的块数量""" + vec_db: FaissVecDB = self.vec_db # type: ignore + count = await vec_db.count_documents(metadata_filter={"doc_id": doc_id}) + return count + async def _save_media( self, doc_id: str, diff --git a/astrbot/core/knowledge_base/kb_mgr.py b/astrbot/core/knowledge_base/kb_mgr.py index 362a60aa2..47a48af64 100644 --- a/astrbot/core/knowledge_base/kb_mgr.py +++ b/astrbot/core/knowledge_base/kb_mgr.py @@ -190,8 +190,7 @@ class KnowledgeBaseManager: kb.emoji = emoji if embedding_provider_id is not None: kb.embedding_provider_id = embedding_provider_id - if rerank_provider_id is not None: - kb.rerank_provider_id = rerank_provider_id + kb.rerank_provider_id = rerank_provider_id # 允许设置为 None if chunk_size is not None: kb.chunk_size = chunk_size if chunk_overlap is not None: diff --git a/astrbot/core/knowledge_base/retrieval/manager.py b/astrbot/core/knowledge_base/retrieval/manager.py index 63d59e973..7e90cf2f6 100644 --- a/astrbot/core/knowledge_base/retrieval/manager.py +++ b/astrbot/core/knowledge_base/retrieval/manager.py @@ -94,6 +94,7 @@ class RetrievalManager: "top_k_sparse": kb.top_k_sparse or 50, "top_m_final": kb.top_m_final or 5, "vec_db": kb_helper.vec_db, + "rerank_provider_id": kb.rerank_provider_id, } new_kb_ids.append(kb_id) else: @@ -147,7 +148,13 @@ class RetrievalManager: first_rerank = None for kb_id in kb_ids: vec_db: FaissVecDB = kb_options[kb_id]["vec_db"] - if vec_db and vec_db.rerank_provider: + rerank_pi = kb_options[kb_id]["rerank_provider_id"] + if ( + vec_db + and vec_db.rerank_provider + and rerank_pi + and rerank_pi == vec_db.rerank_provider.meta().id + ): first_rerank = vec_db.rerank_provider break if first_rerank and retrieval_results: diff --git a/astrbot/dashboard/routes/knowledge_base.py b/astrbot/dashboard/routes/knowledge_base.py index 27e8e0662..f4f54d80b 100644 --- a/astrbot/dashboard/routes/knowledge_base.py +++ b/astrbot/dashboard/routes/knowledge_base.py @@ -30,8 +30,6 @@ class KnowledgeBaseRoute(Route): # 注册路由 self.routes = { - # 系统管理 - # "/kb/status": ("GET", self.get_kb_status), # 知识库管理 "/kb/list": ("GET", self.list_kbs), "/kb/create": ("POST", self.create_kb), @@ -43,11 +41,10 @@ class KnowledgeBaseRoute(Route): "/kb/document/list": ("GET", self.list_documents), "/kb/document/upload": ("POST", self.upload_document), "/kb/document/get": ("GET", self.get_document), - # "/kb/document/delete": ("POST", self.delete_document), + "/kb/document/delete": ("POST", self.delete_document), # # 块管理 "/kb/chunk/list": ("GET", self.list_chunks), - # "/kb/chunk/get": ("GET", self.get_chunk), - # "/kb/chunk/delete": ("POST", self.delete_chunk), + "/kb/chunk/delete": ("POST", self.delete_chunk), # # 多媒体管理 # "/kb/media/list": ("GET", self.list_media), # "/kb/media/delete": ("POST", self.delete_media), @@ -579,6 +576,70 @@ class KnowledgeBaseRoute(Route): logger.error(traceback.format_exc()) return Response().error(f"获取文档详情失败: {str(e)}").__dict__ + async def delete_document(self): + """删除文档 + + Body: + - kb_id: 知识库 ID (必填) + - doc_id: 文档 ID (必填) + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + + kb_id = data.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + doc_id = data.get("doc_id") + if not doc_id: + return Response().error("缺少参数 doc_id").__dict__ + + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + await kb_helper.delete_document(doc_id) + return Response().ok(message="删除文档成功").__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"删除文档失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"删除文档失败: {str(e)}").__dict__ + + async def delete_chunk(self): + """删除文本块 + + Body: + - kb_id: 知识库 ID (必填) + - chunk_id: 块 ID (必填) + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + + kb_id = data.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + chunk_id = data.get("chunk_id") + if not chunk_id: + return Response().error("缺少参数 chunk_id").__dict__ + + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + await kb_helper.delete_chunk(chunk_id) + return Response().ok(message="删除文本块成功").__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"删除文本块失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"删除文本块失败: {str(e)}").__dict__ + async def list_chunks(self): """获取块列表 @@ -612,6 +673,7 @@ class KnowledgeBaseRoute(Route): "items": chunk_list, "page": page, "page_size": page_size, + "total": await kb_helper.get_chunk_count_by_doc_id(doc_id), } ) .__dict__ diff --git a/dashboard/src/i18n/locales/zh-CN/features/knowledge-base/document.json b/dashboard/src/i18n/locales/zh-CN/features/knowledge-base/document.json index 22781666a..c90c29cc2 100644 --- a/dashboard/src/i18n/locales/zh-CN/features/knowledge-base/document.json +++ b/dashboard/src/i18n/locales/zh-CN/features/knowledge-base/document.json @@ -22,7 +22,10 @@ "preview": "预览", "search": "搜索分块", "searchPlaceholder": "输入关键词搜索分块内容...", - "showing": "显示" + "showing": "显示", + "deleteConfirm": "确定要删除该文本块吗?", + "deleteSuccess": "文本块删除成功", + "deleteFailed": "文本块删除失败" }, "edit": { "title": "编辑分块", diff --git a/dashboard/src/views/knowledge-base/DocumentDetail.vue b/dashboard/src/views/knowledge-base/DocumentDetail.vue index 5e1c22b20..e99d1a75d 100644 --- a/dashboard/src/views/knowledge-base/DocumentDetail.vue +++ b/dashboard/src/views/knowledge-base/DocumentDetail.vue @@ -133,6 +133,14 @@ color="info" @click="viewChunk(item)" /> + +