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 2adae69cc..116c43be1 100644 --- a/astrbot/core/db/vec_db/faiss_impl/document_storage.py +++ b/astrbot/core/db/vec_db/faiss_impl/document_storage.py @@ -2,13 +2,22 @@ import json import os from contextlib import asynccontextmanager from datetime import datetime +from pathlib import Path -from sqlalchemy import Column, Text +from sqlalchemy import Column, Text, bindparam from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine from sqlalchemy.orm import sessionmaker from sqlmodel import Field, MetaData, SQLModel, col, func, select, text from astrbot.core import logger +from astrbot.core.knowledge_base.retrieval.tokenizer import ( + build_fts5_or_query, + load_stopwords, + to_fts5_search_text, +) + +FTS_TABLE_NAME = "documents_fts" +FTS_REBUILD_BATCH_SIZE = 1000 class BaseDocModel(SQLModel, table=False): @@ -42,6 +51,10 @@ class DocumentStorage: os.path.dirname(__file__), "sqlite_init.sql", ) + self.fts5_available = False + self._fts_contentless_delete = False + self._fts_index_ready = False + self._stopwords: set[str] | None = None async def initialize(self) -> None: """Initialize the SQLite database and create the documents table if it doesn't exist.""" @@ -78,8 +91,49 @@ class DocumentStorage: except BaseException: pass + await self._initialize_fts5(conn) await conn.commit() + async def _initialize_fts5(self, executor) -> None: + try: + try: + await executor.execute( + text( + f""" + CREATE VIRTUAL TABLE IF NOT EXISTS {FTS_TABLE_NAME} + USING fts5( + search_text, + content='', + contentless_delete=1, + tokenize='unicode61' + ) + """, + ), + ) + self._fts_contentless_delete = True + except Exception: + await executor.execute( + text( + f""" + CREATE VIRTUAL TABLE IF NOT EXISTS {FTS_TABLE_NAME} + USING fts5( + search_text, + content='', + tokenize='unicode61' + ) + """, + ), + ) + self._fts_contentless_delete = False + self.fts5_available = True + except Exception as e: + self.fts5_available = False + self._fts_contentless_delete = False + logger.warning( + f"SQLite FTS5 is unavailable for document storage {self.db_path}; " + f"falling back to in-memory BM25 sparse retrieval: {e}", + ) + async def connect(self) -> None: """Connect to the SQLite database.""" if self.engine is None: @@ -100,6 +154,18 @@ class DocumentStorage: async with self.async_session_maker() as session: # type: ignore yield session + @property + def stopwords(self) -> set[str]: + if self._stopwords is None: + stopwords_path = ( + Path(__file__).parents[3] + / "knowledge_base" + / "retrieval" + / "hit_stopwords.txt" + ) + self._stopwords = load_stopwords(stopwords_path) + return self._stopwords + async def get_documents( self, metadata_filters: dict, @@ -172,6 +238,8 @@ class DocumentStorage: ) session.add(document) await session.flush() # Flush to get the ID + if document.id is not None: + await self._insert_fts_row(session, int(document.id), text) return document.id # type: ignore async def insert_documents_batch( @@ -209,6 +277,7 @@ class DocumentStorage: session.add(document) await session.flush() # Flush to get all IDs + await self._insert_fts_rows_batch(session, documents, texts) return [doc.id for doc in documents] # type: ignore async def delete_document_by_doc_id(self, doc_id: str) -> None: @@ -226,6 +295,8 @@ class DocumentStorage: document = result.scalar_one_or_none() if document: + if document.id is not None: + await self._delete_fts_row(session, int(document.id), document.text) await session.delete(document) async def get_document_by_doc_id(self, doc_id: str): @@ -265,9 +336,13 @@ class DocumentStorage: document = result.scalar_one_or_none() if document: + if document.id is not None: + await self._delete_fts_row(session, int(document.id), document.text) document.text = new_text document.updated_at = datetime.now() session.add(document) + if document.id is not None: + await self._insert_fts_row(session, int(document.id), new_text) async def delete_documents(self, metadata_filters: dict) -> None: """Delete documents by their metadata filters. @@ -293,6 +368,7 @@ class DocumentStorage: result = await session.execute(query) documents = result.scalars().all() + await self._delete_fts_rows_batch(session, documents) for doc in documents: await session.delete(doc) @@ -323,6 +399,286 @@ class DocumentStorage: count = result.scalar_one_or_none() return count if count is not None else 0 + async def ensure_fts_index(self) -> bool: + """Ensure the FTS5 sparse index exists and matches the documents table.""" + if not self.fts5_available: + return False + if self._fts_index_ready: + return True + + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + doc_count = await self._count_documents_in_session(session) + fts_count = await self._count_fts_rows(session) + if doc_count == fts_count: + self._fts_index_ready = True + return True + + logger.info( + f"Rebuilding FTS5 sparse index for {self.db_path}: " + f"documents={doc_count}, fts_rows={fts_count}", + ) + await self.rebuild_fts_index() + return self.fts5_available + + async def rebuild_fts_index(self) -> None: + """Rebuild the contentless FTS5 sparse index from documents.""" + if not self.fts5_available: + return + + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session, session.begin(): + await session.execute(text(f"DROP TABLE IF EXISTS {FTS_TABLE_NAME}")) + await self._initialize_fts5(session) + if not self.fts5_available: + return + + last_id = 0 + while True: + query = ( + select(Document) + .where(col(Document.id) > last_id) + .order_by(col(Document.id)) + .limit(FTS_REBUILD_BATCH_SIZE) + ) + result = await session.execute(query) + documents = result.scalars().all() + if not documents: + break + + await self._insert_fts_rows_batch( + session, + documents, + [doc.text for doc in documents], + ) + last_id = int(documents[-1].id or last_id) + + self._fts_index_ready = True + + async def search_sparse( + self, + query_tokens: list[str], + limit: int, + ) -> list[dict] | None: + """Search chunks using the FTS5 sparse index. + + Returns None when FTS5 is unavailable so callers can fall back to another + sparse retrieval implementation. + """ + if limit <= 0: + return [] + if not await self.ensure_fts_index(): + return None + + match_query = build_fts5_or_query(query_tokens) + if not match_query: + return [] + + async with self.get_session() as session: + try: + result = await session.execute( + text( + f""" + SELECT + d.id AS id, + d.doc_id AS doc_id, + d.text AS text, + d.metadata AS metadata, + d.created_at AS created_at, + d.updated_at AS updated_at, + bm25({FTS_TABLE_NAME}) AS score + FROM {FTS_TABLE_NAME} + JOIN documents d ON d.id = {FTS_TABLE_NAME}.rowid + WHERE {FTS_TABLE_NAME} MATCH :query + ORDER BY score ASC, d.id ASC + LIMIT :limit + """, + ), + {"query": match_query, "limit": int(limit)}, + ) + except Exception as e: + logger.warning( + f"FTS5 sparse search failed for {self.db_path}; " + f"falling back to in-memory BM25: {e}", + ) + self.fts5_available = False + return None + + rows = result.mappings().all() + return [ + { + "id": row["id"], + "doc_id": row["doc_id"], + "text": row["text"], + "metadata": row["metadata"], + "created_at": row["created_at"], + "updated_at": row["updated_at"], + "score": float(row["score"]), + } + for row in rows + ] + + async def _count_documents_in_session(self, session: AsyncSession) -> int: + result = await session.execute(select(func.count(col(Document.id)))) + count = result.scalar_one_or_none() + return int(count or 0) + + async def _count_fts_rows(self, session: AsyncSession) -> int: + result = await session.execute( + text(f"SELECT count(*) FROM {FTS_TABLE_NAME}"), + ) + count = result.scalar_one_or_none() + return int(count or 0) + + async def _insert_fts_row( + self, + session: AsyncSession, + rowid: int, + content: str, + ) -> None: + if not self.fts5_available: + return + + search_text = to_fts5_search_text(content, self.stopwords) + await session.execute( + text( + f""" + INSERT INTO {FTS_TABLE_NAME}(rowid, search_text) + VALUES (:rowid, :search_text) + """, + ), + {"rowid": rowid, "search_text": search_text}, + ) + + async def _insert_fts_rows_batch( + self, + session: AsyncSession, + documents: list[Document], + contents: list[str], + ) -> None: + if not self.fts5_available: + return + + fts_params = [ + { + "rowid": int(doc.id), + "search_text": to_fts5_search_text(content, self.stopwords), + } + for doc, content in zip(documents, contents) + if doc.id is not None + ] + if not fts_params: + return + + await session.execute( + text( + f""" + INSERT INTO {FTS_TABLE_NAME}(rowid, search_text) + VALUES (:rowid, :search_text) + """, + ), + fts_params, + ) + + async def _delete_fts_row( + self, + session: AsyncSession, + rowid: int, + content: str, + ) -> None: + if not self.fts5_available: + return + + if self._fts_contentless_delete: + await session.execute( + text(f"DELETE FROM {FTS_TABLE_NAME} WHERE rowid = :rowid"), + {"rowid": rowid}, + ) + return + + if not await self._fts_row_exists(session, rowid): + return + + search_text = to_fts5_search_text(content, self.stopwords) + await session.execute( + text( + f""" + INSERT INTO {FTS_TABLE_NAME}({FTS_TABLE_NAME}, rowid, search_text) + VALUES ('delete', :rowid, :search_text) + """, + ), + {"rowid": rowid, "search_text": search_text}, + ) + + async def _delete_fts_rows_batch( + self, + session: AsyncSession, + documents: list[Document], + ) -> None: + if not self.fts5_available: + return + + docs_with_ids = [doc for doc in documents if doc.id is not None] + if not docs_with_ids: + return + + if self._fts_contentless_delete: + await session.execute( + text(f"DELETE FROM {FTS_TABLE_NAME} WHERE rowid = :rowid"), + [{"rowid": int(doc.id)} for doc in docs_with_ids if doc.id is not None], + ) + return + + existing_rowids = await self._existing_fts_rowids( + session, + [int(doc.id) for doc in docs_with_ids if doc.id is not None], + ) + fts_params = [ + { + "rowid": int(doc.id), + "search_text": to_fts5_search_text(doc.text, self.stopwords), + } + for doc in docs_with_ids + if doc.id is not None and int(doc.id) in existing_rowids + ] + if not fts_params: + return + + await session.execute( + text( + f""" + INSERT INTO {FTS_TABLE_NAME}({FTS_TABLE_NAME}, rowid, search_text) + VALUES ('delete', :rowid, :search_text) + """, + ), + fts_params, + ) + + async def _fts_row_exists(self, session: AsyncSession, rowid: int) -> bool: + result = await session.execute( + text(f"SELECT 1 FROM {FTS_TABLE_NAME} WHERE rowid = :rowid LIMIT 1"), + {"rowid": rowid}, + ) + return result.scalar_one_or_none() is not None + + async def _existing_fts_rowids( + self, + session: AsyncSession, + rowids: list[int], + ) -> set[int]: + if not rowids: + return set() + + result = await session.execute( + text( + f"SELECT rowid FROM {FTS_TABLE_NAME} WHERE rowid IN :rowids" + ).bindparams(bindparam("rowids", expanding=True)), + {"rowids": rowids}, + ) + return {int(row[0]) for row in result.fetchall()} + async def get_user_ids(self) -> list[str]: """Retrieve all user IDs from the documents table. diff --git a/astrbot/core/knowledge_base/retrieval/__init__.py b/astrbot/core/knowledge_base/retrieval/__init__.py index f5d196cb9..b7c88075d 100644 --- a/astrbot/core/knowledge_base/retrieval/__init__.py +++ b/astrbot/core/knowledge_base/retrieval/__init__.py @@ -1,8 +1,11 @@ """检索模块""" -from .manager import RetrievalManager, RetrievalResult -from .rank_fusion import FusedResult, RankFusion -from .sparse_retriever import SparseResult, SparseRetriever +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .manager import RetrievalManager, RetrievalResult + from .rank_fusion import FusedResult, RankFusion + from .sparse_retriever import SparseResult, SparseRetriever __all__ = [ "FusedResult", @@ -12,3 +15,31 @@ __all__ = [ "SparseResult", "SparseRetriever", ] + + +def __getattr__(name: str): + if name in {"RetrievalManager", "RetrievalResult"}: + from .manager import RetrievalManager, RetrievalResult + + return { + "RetrievalManager": RetrievalManager, + "RetrievalResult": RetrievalResult, + }[name] + + if name in {"FusedResult", "RankFusion"}: + from .rank_fusion import FusedResult, RankFusion + + return { + "FusedResult": FusedResult, + "RankFusion": RankFusion, + }[name] + + if name in {"SparseResult", "SparseRetriever"}: + from .sparse_retriever import SparseResult, SparseRetriever + + return { + "SparseResult": SparseResult, + "SparseRetriever": SparseRetriever, + }[name] + + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/astrbot/core/knowledge_base/retrieval/sparse_retriever.py b/astrbot/core/knowledge_base/retrieval/sparse_retriever.py index 9d485dab2..f06eb5090 100644 --- a/astrbot/core/knowledge_base/retrieval/sparse_retriever.py +++ b/astrbot/core/knowledge_base/retrieval/sparse_retriever.py @@ -8,10 +8,13 @@ import os from dataclasses import dataclass from typing import TYPE_CHECKING -import jieba from rank_bm25 import BM25Okapi from astrbot.core.knowledge_base.kb_db_sqlite import KBSQLiteDatabase +from astrbot.core.knowledge_base.retrieval.tokenizer import ( + load_stopwords, + tokenize_text, +) if TYPE_CHECKING: from astrbot.core.db.vec_db.faiss_impl import FaissVecDB @@ -47,13 +50,9 @@ class SparseRetriever: self.kb_db = kb_db self._index_cache = {} # 缓存 BM25 索引 - with open( + self.hit_stopwords = load_stopwords( os.path.join(os.path.dirname(__file__), "hit_stopwords.txt"), - encoding="utf-8", - ) as f: - self.hit_stopwords = { - word.strip() for word in set(f.read().splitlines()) if word.strip() - } + ) async def retrieve( self, @@ -72,7 +71,52 @@ class SparseRetriever: List[SparseResult]: 检索结果列表 """ - # 1. 获取所有相关块 + fts_results = [] + fallback_kb_ids = [] + query_tokens = tokenize_text(query, self.hit_stopwords) + for kb_id in kb_ids: + vec_db: FaissVecDB | None = kb_options.get(kb_id, {}).get("vec_db") + if not vec_db: + continue + top_k_sparse = kb_options.get(kb_id, {}).get("top_k_sparse", 50) + result = await vec_db.document_storage.search_sparse( + query_tokens=query_tokens, + limit=top_k_sparse, + ) + if result is None: + fallback_kb_ids.append(kb_id) + continue + + for doc in result: + chunk_md = json.loads(doc["metadata"]) + fts_results.append( + SparseResult( + chunk_id=doc["doc_id"], + chunk_index=chunk_md["chunk_index"], + doc_id=chunk_md["kb_doc_id"], + kb_id=kb_id, + content=doc["text"], + score=-float(doc["score"]), + ), + ) + + fallback_results = [] + if fallback_kb_ids: + fallback_results = await self._retrieve_with_bm25( + query=query, + kb_ids=fallback_kb_ids, + kb_options=kb_options, + ) + results = fts_results + fallback_results + results.sort(key=lambda x: x.score, reverse=True) + return results + + async def _retrieve_with_bm25( + self, + query: str, + kb_ids: list[str], + kb_options: dict, + ) -> list[SparseResult]: top_k_sparse = 0 chunks = [] for kb_id in kb_ids: @@ -103,20 +147,13 @@ class SparseRetriever: # 2. 准备文档和索引 corpus = [chunk["text"] for chunk in chunks] - tokenized_corpus = [list(jieba.cut(doc)) for doc in corpus] - tokenized_corpus = [ - [word for word in doc if word not in self.hit_stopwords] - for doc in tokenized_corpus - ] + tokenized_corpus = [tokenize_text(doc, self.hit_stopwords) for doc in corpus] # 3. 构建 BM25 索引 bm25 = BM25Okapi(tokenized_corpus) # 4. 执行检索 - tokenized_query = list(jieba.cut(query)) - tokenized_query = [ - word for word in tokenized_query if word not in self.hit_stopwords - ] + tokenized_query = tokenize_text(query, self.hit_stopwords) scores = bm25.get_scores(tokenized_query) # 5. 排序并返回 Top-K diff --git a/astrbot/core/knowledge_base/retrieval/tokenizer.py b/astrbot/core/knowledge_base/retrieval/tokenizer.py new file mode 100644 index 000000000..1f4f07bb9 --- /dev/null +++ b/astrbot/core/knowledge_base/retrieval/tokenizer.py @@ -0,0 +1,39 @@ +"""Tokenization helpers shared by sparse retrieval indexes.""" + +import re +from pathlib import Path +from re import Pattern + +import jieba + +_TERM_PATTERN: Pattern[str] = re.compile(r"\w", re.UNICODE) + + +def load_stopwords(path: Path | str) -> set[str]: + with Path(path).open(encoding="utf-8") as f: + return {word.strip() for word in set(f.read().splitlines()) if word.strip()} + + +def tokenize_text(text: str, stopwords: set[str]) -> list[str]: + tokens = [] + for token in jieba.cut(text or ""): + token = token.strip() + if not token or token in stopwords: + continue + if not _TERM_PATTERN.search(token): + continue + tokens.append(token) + return tokens + + +def to_fts5_search_text(text: str, stopwords: set[str]) -> str: + return " ".join(tokenize_text(text, stopwords)) + + +def quote_fts5_token(token: str) -> str: + return '"' + token.replace('"', '""') + '"' + + +def build_fts5_or_query(tokens: list[str]) -> str: + quoted_tokens = [quote_fts5_token(token) for token in tokens if token] + return " OR ".join(quoted_tokens) diff --git a/tests/unit/test_document_storage_fts.py b/tests/unit/test_document_storage_fts.py new file mode 100644 index 000000000..753c371ef --- /dev/null +++ b/tests/unit/test_document_storage_fts.py @@ -0,0 +1,75 @@ +import pytest + +from astrbot.core.db.vec_db.faiss_impl.document_storage import DocumentStorage + + +@pytest.mark.asyncio +async def test_document_storage_fts_insert_search_and_delete(tmp_path): + storage = DocumentStorage(str(tmp_path / "doc.db")) + await storage.initialize() + + assert storage.fts5_available is True + + await storage.insert_documents_batch( + doc_ids=["chunk-1", "chunk-2"], + texts=["AstrBot 知识库召回性能优化", "FAISS 向量检索"], + metadatas=[ + {"kb_doc_id": "doc-1", "kb_id": "kb-1", "chunk_index": 0}, + {"kb_doc_id": "doc-1", "kb_id": "kb-1", "chunk_index": 1}, + ], + ) + + results = await storage.search_sparse(["知识库"], limit=10) + + assert results is not None + assert [result["doc_id"] for result in results] == ["chunk-1"] + + await storage.delete_document_by_doc_id("chunk-1") + results = await storage.search_sparse(["知识库"], limit=10) + + assert results == [] + + await storage.close() + + +@pytest.mark.asyncio +async def test_document_storage_fts_rebuilds_existing_documents(tmp_path): + storage = DocumentStorage(str(tmp_path / "doc.db")) + await storage.initialize() + + storage.fts5_available = False + await storage.insert_document( + doc_id="legacy-chunk", + text="legacy 知识库 文本", + metadata={"kb_doc_id": "doc-1", "kb_id": "kb-1", "chunk_index": 0}, + ) + + storage.fts5_available = True + storage._fts_index_ready = False + + results = await storage.search_sparse(["知识库"], limit=10) + + assert results is not None + assert [result["doc_id"] for result in results] == ["legacy-chunk"] + + await storage.close() + + +@pytest.mark.asyncio +async def test_document_storage_fts_delete_skips_missing_fts_row(tmp_path): + storage = DocumentStorage(str(tmp_path / "doc.db")) + await storage.initialize() + + storage.fts5_available = False + await storage.insert_document( + doc_id="legacy-chunk", + text="legacy 知识库 文本", + metadata={"kb_doc_id": "doc-1", "kb_id": "kb-1", "chunk_index": 0}, + ) + + storage.fts5_available = True + await storage.delete_document_by_doc_id("legacy-chunk") + + assert await storage.get_document_by_doc_id("legacy-chunk") is None + + await storage.close() diff --git a/tests/unit/test_sparse_retriever.py b/tests/unit/test_sparse_retriever.py new file mode 100644 index 000000000..11c491b4d --- /dev/null +++ b/tests/unit/test_sparse_retriever.py @@ -0,0 +1,93 @@ +import json +from types import SimpleNamespace + +import pytest + +from astrbot.core.knowledge_base.retrieval.sparse_retriever import SparseRetriever + + +def make_doc(chunk_id: str, text: str, chunk_index: int = 0) -> dict: + return { + "doc_id": chunk_id, + "text": text, + "metadata": json.dumps( + { + "chunk_index": chunk_index, + "kb_doc_id": f"doc-{chunk_index}", + "kb_id": "kb-1", + }, + ), + } + + +class FTSStorage: + def __init__(self): + self.search_sparse_calls = 0 + self.get_documents_calls = 0 + + async def search_sparse(self, query_tokens: list[str], limit: int): + self.search_sparse_calls += 1 + assert query_tokens == ["apple"] + assert limit == 1 + return [ + { + **make_doc("chunk-1", "apple banana", 0), + "score": -1.0, + }, + ] + + async def get_documents(self, *args, **kwargs): + self.get_documents_calls += 1 + return [] + + +class FallbackStorage: + def __init__(self): + self.search_sparse_calls = 0 + self.get_documents_calls = 0 + + async def search_sparse(self, query_tokens: list[str], limit: int): + self.search_sparse_calls += 1 + return None + + async def get_documents(self, metadata_filters: dict, limit: int | None, offset): + self.get_documents_calls += 1 + return [ + make_doc("chunk-1", "apple banana", 0), + make_doc("chunk-2", "orange pear", 1), + make_doc("chunk-3", "grape melon", 2), + ] + + +@pytest.mark.asyncio +async def test_sparse_retriever_uses_fts5_when_available(): + storage = FTSStorage() + vec_db = SimpleNamespace(document_storage=storage) + retriever = SparseRetriever(kb_db=None) + + results = await retriever.retrieve( + query="apple", + kb_ids=["kb-1"], + kb_options={"kb-1": {"vec_db": vec_db, "top_k_sparse": 1}}, + ) + + assert [result.chunk_id for result in results] == ["chunk-1"] + assert storage.search_sparse_calls == 1 + assert storage.get_documents_calls == 0 + + +@pytest.mark.asyncio +async def test_sparse_retriever_falls_back_to_bm25_when_fts5_is_unavailable(): + storage = FallbackStorage() + vec_db = SimpleNamespace(document_storage=storage) + retriever = SparseRetriever(kb_db=None) + + results = await retriever.retrieve( + query="apple", + kb_ids=["kb-1"], + kb_options={"kb-1": {"vec_db": vec_db, "top_k_sparse": 1}}, + ) + + assert [result.chunk_id for result in results] == ["chunk-1"] + assert storage.search_sparse_calls == 1 + assert storage.get_documents_calls == 1