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)"
/>
+
+
@@ -318,7 +326,7 @@ const loadChunks = async () => {
})
if (response.data.status === 'ok') {
chunks.value = response.data.data.items || []
- totalChunks.value = response.data.data.items.length
+ totalChunks.value = response.data.data.total || 0
}
} catch (error) {
console.error('Failed to load chunks:', error)
@@ -346,6 +354,27 @@ const viewChunk = (chunk: any) => {
showViewDialog.value = true
}
+// 删除分块
+const deleteChunk = async (chunk: any) => {
+ if (!confirm(t('chunks.deleteConfirm'))) return
+ try {
+ const response = await axios.post('/api/kb/chunk/delete', {
+ chunk_id: chunk.chunk_id,
+ doc_id: docId.value,
+ kb_id: kbId.value
+ })
+ if (response.data.status === 'ok') {
+ showSnackbar(t('chunks.deleteSuccess'))
+ loadChunks()
+ } else {
+ showSnackbar(t('chunks.deleteFailed'), 'error')
+ }
+ } catch (error) {
+ console.error('Failed to delete chunk:', error)
+ showSnackbar(t('chunks.deleteFailed'), 'error')
+ }
+}
+
// 工具函数
const getFileIcon = (fileType: string) => {
const type = fileType?.toLowerCase() || ''
diff --git a/dashboard/src/views/knowledge-base/components/DocumentsTab.vue b/dashboard/src/views/knowledge-base/components/DocumentsTab.vue
index 7dda40363..814443415 100644
--- a/dashboard/src/views/knowledge-base/components/DocumentsTab.vue
+++ b/dashboard/src/views/knowledge-base/components/DocumentsTab.vue
@@ -357,7 +357,8 @@ const deleteDocument = async () => {
deleting.value = true
try {
const response = await axios.post('/api/kb/document/delete', {
- doc_id: deleteTarget.value.doc_id
+ doc_id: deleteTarget.value.doc_id,
+ kb_id: props.kbId
})
if (response.data.status === 'ok') {