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 116c43be1..d0310d750 100644 --- a/astrbot/core/db/vec_db/faiss_impl/document_storage.py +++ b/astrbot/core/db/vec_db/faiss_impl/document_storage.py @@ -96,36 +96,29 @@ class DocumentStorage: 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' - ) - """, - ), + await self._create_fts5_table(executor, if_not_exists=True) + + is_valid_fts5, has_contentless_delete = await self._inspect_fts5_table( + executor, + ) + if not is_valid_fts5: + logger.warning( + f"Detected incompatible legacy table `{FTS_TABLE_NAME}` in " + f"{self.db_path}; recreating FTS5 table.", ) - 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' - ) - """, - ), + await executor.execute(text(f"DROP TABLE IF EXISTS {FTS_TABLE_NAME}")) + await self._create_fts5_table(executor, if_not_exists=False) + + is_valid_fts5, has_contentless_delete = await self._inspect_fts5_table( + executor, ) - self._fts_contentless_delete = False + if not is_valid_fts5: + raise RuntimeError( + f"Failed to create a valid FTS5 table `{FTS_TABLE_NAME}`", + ) + self.fts5_available = True + self._fts_contentless_delete = has_contentless_delete except Exception as e: self.fts5_available = False self._fts_contentless_delete = False @@ -134,6 +127,69 @@ class DocumentStorage: f"falling back to in-memory BM25 sparse retrieval: {e}", ) + async def _create_fts5_table(self, executor, if_not_exists: bool) -> None: + create_clause = ( + "CREATE VIRTUAL TABLE IF NOT EXISTS" + if if_not_exists + else "CREATE VIRTUAL TABLE" + ) + try: + await executor.execute( + text( + f""" + {create_clause} {FTS_TABLE_NAME} + USING fts5( + search_text, + content='', + contentless_delete=1, + tokenize='unicode61' + ) + """, + ), + ) + except Exception: + await executor.execute( + text( + f""" + {create_clause} {FTS_TABLE_NAME} + USING fts5( + search_text, + content='', + tokenize='unicode61' + ) + """, + ), + ) + + async def _inspect_fts5_table(self, executor) -> tuple[bool, bool]: + schema_result = await executor.execute( + text( + """ + SELECT sql + FROM sqlite_master + WHERE type='table' AND name=:table_name + """, + ), + {"table_name": FTS_TABLE_NAME}, + ) + create_sql = schema_result.scalar_one_or_none() + if not create_sql: + return False, False + + normalized_sql = create_sql.lower() + if "virtual table" not in normalized_sql or "using fts5" not in normalized_sql: + return False, False + + pragma_result = await executor.execute( + text(f"PRAGMA table_info({FTS_TABLE_NAME})"), + ) + columns = {row[1] for row in pragma_result.fetchall()} + if "search_text" not in columns: + return False, False + + normalized_sql_no_whitespace = "".join(normalized_sql.split()) + return True, "contentless_delete=1" in normalized_sql_no_whitespace + async def connect(self) -> None: """Connect to the SQLite database.""" if self.engine is None: diff --git a/tests/unit/test_document_storage_fts.py b/tests/unit/test_document_storage_fts.py index 753c371ef..a7dd32c94 100644 --- a/tests/unit/test_document_storage_fts.py +++ b/tests/unit/test_document_storage_fts.py @@ -1,3 +1,5 @@ +import sqlite3 + import pytest from astrbot.core.db.vec_db.faiss_impl.document_storage import DocumentStorage @@ -73,3 +75,29 @@ async def test_document_storage_fts_delete_skips_missing_fts_row(tmp_path): assert await storage.get_document_by_doc_id("legacy-chunk") is None await storage.close() + + +@pytest.mark.asyncio +async def test_document_storage_fts_recovers_from_legacy_non_fts_table(tmp_path): + db_path = tmp_path / "doc.db" + conn = sqlite3.connect(db_path) + conn.execute("CREATE TABLE documents_fts (rowid INTEGER PRIMARY KEY)") + conn.commit() + conn.close() + + storage = DocumentStorage(str(db_path)) + await storage.initialize() + + assert storage.fts5_available is True + + await storage.insert_document( + doc_id="legacy-fix", + text="legacy fts recovery text", + metadata={"kb_doc_id": "doc-1", "kb_id": "kb-1", "chunk_index": 0}, + ) + results = await storage.search_sparse(["legacy"], limit=10) + + assert results is not None + assert [result["doc_id"] for result in results] == ["legacy-fix"] + + await storage.close()