fz-mirrors/AstrBot
1
0
Fork 0
mirror of https://github.com/AstrBotDevs/AstrBot synced 2026-07-15 17:30:13 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: delete chunk and delete document

Browse Source
This commit is contained in:
Soulter
2025-10-24 13:59:17 +08:00
parent e6d9db9395
commit 2f130ba009
11 changed files with 258 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
astrbot/core/db/vec_db/faiss_impl/document_storage.py
View File

@@ -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()

17
astrbot/core/db/vec_db/faiss_impl/embedding_storage.py
View File

@@ -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):
"""保存索引

43
astrbot/core/db/vec_db/faiss_impl/vec_db.py
View File

@@ -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)

15
astrbot/core/knowledge_base/kb_db_sqlite.py
View File

@@ -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]:

26
astrbot/core/knowledge_base/kb_helper.py
View File

@@ -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,

3
astrbot/core/knowledge_base/kb_mgr.py
View File

@@ -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:

9
astrbot/core/knowledge_base/retrieval/manager.py
View File

@@ -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:

72
astrbot/dashboard/routes/knowledge_base.py
View File

@@ -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__

5
dashboard/src/i18n/locales/zh-CN/features/knowledge-base/document.json
View File

@@ -22,7 +22,10 @@
"preview": "预览",
"search": "搜索分块",
"searchPlaceholder": "输入关键词搜索分块内容...",
"showing": "显示"
"showing": "显示",
"deleteConfirm": "确定要删除该文本块吗?",
"deleteSuccess": "文本块删除成功",
"deleteFailed": "文本块删除失败"
},
"edit": {
"title": "编辑分块",

31
dashboard/src/views/knowledge-base/DocumentDetail.vue
View File

@@ -133,6 +133,14 @@
color="info"
@click="viewChunk(item)"
/>
<!-- 删除 -->
<v-btn
icon="mdi-delete"
variant="text"
size="small"
color="error"
@click="deleteChunk(item)"
/>
</template>
<template #no-data>
@@ -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() || ''

3
dashboard/src/views/knowledge-base/components/DocumentsTab.vue
View File

@@ -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') {