feat(kb): [WIP]overhaul knowledge base workflows and retrieval

astrbot/core/db/vec_db/faiss_impl/document_storage.py

@@ -1,19 +1,24 @@
 import json
 import os
+from asyncio import Lock
 from contextlib import asynccontextmanager
 from datetime import datetime
 from pathlib import Path
 
 from sqlalchemy import Column, Text, bindparam
+from sqlalchemy.dialects import sqlite
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
 from sqlalchemy.orm import sessionmaker
+from sqlalchemy.pool import NullPool
+from sqlalchemy.schema import CreateTable
 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,
+    tokenize_text,
 )
 
 FTS_TABLE_NAME = "documents_fts"
@@ -55,50 +60,98 @@ def __init__(self, db_path: str) -> None:
         self._fts_contentless_delete = False
         self._fts_index_ready = False
         self._stopwords: set[str] | None = None
+        self._fts_rebuild_lock = Lock()
 
     async def initialize(self) -> None:
         """Initialize the SQLite database and create the documents table if it doesn't exist."""
         await self.connect()
         async with self.engine.begin() as conn:  # type: ignore
-            # Create tables using SQLModel
-            await conn.run_sync(BaseDocModel.metadata.create_all)
+            await self._ensure_documents_table(conn)
+            await self._ensure_generated_columns(conn)
 
-            try:
-                await conn.execute(
-                    text(
-                        "ALTER TABLE documents ADD COLUMN kb_doc_id TEXT "
-                        "GENERATED ALWAYS AS (json_extract(metadata, '$.kb_doc_id')) STORED",
-                    ),
-                )
-                await conn.execute(
-                    text(
-                        "ALTER TABLE documents ADD COLUMN user_id TEXT "
-                        "GENERATED ALWAYS AS (json_extract(metadata, '$.user_id')) STORED",
-                    ),
-                )
+            await self._initialize_fts5(conn)
+            await conn.commit()
 
-                # Create indexes
-                await conn.execute(
-                    text(
-                        "CREATE INDEX IF NOT EXISTS idx_documents_kb_doc_id ON documents(kb_doc_id)",
-                    ),
-                )
-                await conn.execute(
-                    text(
-                        "CREATE INDEX IF NOT EXISTS idx_documents_user_id ON documents(user_id)",
-                    ),
-                )
-            except BaseException:
-                pass
+    async def _table_columns(self, executor, table_name: str) -> set[str]:
+        result = await executor.execute(text(f"PRAGMA table_xinfo({table_name})"))
+        return {row[1] for row in result.fetchall()}
 
-            await conn.execute(
+    async def _ensure_generated_columns(self, executor) -> None:
+        generated_columns = {
+            "kb_doc_id": "json_extract(metadata, '$.kb_doc_id')",
+            "user_id": "json_extract(metadata, '$.user_id')",
+            "kb_id": "json_extract(metadata, '$.kb_id')",
+        }
+        columns = await self._table_columns(executor, "documents")
+        for column_name, expression in generated_columns.items():
+            if column_name in columns:
+                continue
+            await executor.execute(
                 text(
-                    "CREATE UNIQUE INDEX IF NOT EXISTS idx_documents_doc_id_unique ON documents(doc_id)",
+                    f"ALTER TABLE documents ADD COLUMN {column_name} TEXT "
+                    f"GENERATED ALWAYS AS ({expression}) VIRTUAL",
                 ),
             )
+            columns.add(column_name)
 
-            await self._initialize_fts5(conn)
-            await conn.commit()
+        index_statements = [
+            "CREATE INDEX IF NOT EXISTS idx_documents_kb_doc_id "
+            "ON documents(kb_doc_id)",
+            "CREATE INDEX IF NOT EXISTS idx_documents_user_id ON documents(user_id)",
+            "CREATE INDEX IF NOT EXISTS idx_documents_kb_id ON documents(kb_id)",
+        ]
+        for statement in index_statements:
+            await executor.execute(text(statement))
+
+    async def _ensure_documents_table(self, executor) -> None:
+        """Create the document table from the SQLModel definition."""
+        result = await executor.execute(
+            text(
+                """
+                SELECT 1
+                FROM sqlite_master
+                WHERE type='table' AND name=:table_name
+                LIMIT 1
+                """,
+            ),
+            {"table_name": Document.__tablename__},
+        )
+        if result.scalar_one_or_none() is not None:
+            await self._ensure_doc_id_unique_index(executor)
+            return
+
+        create_table = CreateTable(Document.__table__, if_not_exists=True)  # type: ignore[attr-defined]
+
+        await executor.execute(
+            text(str(create_table.compile(dialect=sqlite.dialect())))
+        )
+        await self._ensure_doc_id_unique_index(executor)
+
+    async def _ensure_doc_id_unique_index(self, executor) -> None:
+        duplicate_result = await executor.execute(
+            text(
+                """
+                SELECT doc_id
+                FROM documents
+                GROUP BY doc_id
+                HAVING COUNT(*) > 1
+                LIMIT 1
+                """,
+            ),
+        )
+        if duplicate_result.scalar_one_or_none() is not None:
+            logger.warning(
+                "Skipping documents.doc_id unique index migration because duplicate "
+                f"doc_id values already exist in {self.db_path}.",
+            )
+            return
+
+        await executor.execute(
+            text(
+                "CREATE UNIQUE INDEX IF NOT EXISTS "
+                "idx_documents_doc_id_unique ON documents(doc_id)",
+            ),
+        )
 
     async def _initialize_fts5(self, executor) -> None:
         try:
@@ -203,6 +256,7 @@ async def connect(self) -> None:
                 self.DATABASE_URL,
                 echo=False,
                 future=True,
+                poolclass=NullPool,
             )
             self.async_session_maker = sessionmaker(
                 self.engine,  # type: ignore
@@ -255,11 +309,11 @@ async def get_documents(
 
         async with self.get_session() as session:
             query = select(Document)
-
-            for key, val in metadata_filters.items():
-                query = query.where(
-                    text(f"json_extract(metadata, '$.{key}') = :filter_{key}"),
-                ).params(**{f"filter_{key}": val})
+            query = await self._apply_metadata_filters(
+                session,
+                query,
+                metadata_filters,
+            )
 
             if ids is not None and len(ids) > 0:
                 valid_ids = [int(i) for i in ids if i != -1]
@@ -421,11 +475,11 @@ async def delete_documents(self, metadata_filters: dict) -> None:
 
         async with self.get_session() as session, session.begin():
             query = select(Document)
-
-            for key, val in metadata_filters.items():
-                query = query.where(
-                    text(f"json_extract(metadata, '$.{key}') = :filter_{key}"),
-                ).params(**{f"filter_{key}": val})
+            query = await self._apply_metadata_filters(
+                session,
+                query,
+                metadata_filters,
+            )
 
             result = await session.execute(query)
             documents = result.scalars().all()
@@ -452,15 +506,144 @@ async def count_documents(self, metadata_filters: dict | None = None) -> int:
             query = select(func.count(col(Document.id)))
 
             if metadata_filters:
-                for key, val in metadata_filters.items():
-                    query = query.where(
-                        text(f"json_extract(metadata, '$.{key}') = :filter_{key}"),
-                    ).params(**{f"filter_{key}": val})
+                query = await self._apply_metadata_filters(
+                    session,
+                    query,
+                    metadata_filters,
+                )
 
             result = await session.execute(query)
             count = result.scalar_one_or_none()
             return count if count is not None else 0
 
+    async def search_documents(
+        self,
+        query_text: str,
+        metadata_filters: dict | None = None,
+        offset: int = 0,
+        limit: int = 100,
+    ) -> tuple[list[dict], int] | None:
+        """Search documents with FTS5 and optional metadata filters.
+
+        Returns None when FTS5 is unavailable so callers can choose whether to
+        fall back to an alternate search strategy.
+        """
+        if limit <= 0:
+            return [], 0
+        if not await self.ensure_fts_index():
+            return None
+
+        match_query = build_fts5_or_query(tokenize_text(query_text, self.stopwords))
+        if not match_query:
+            return [], 0
+
+        metadata_filters = metadata_filters or {}
+        async with self.get_session() as session:
+            filters_sql, filter_params = await self._metadata_filter_sql(
+                session,
+                metadata_filters,
+                table_alias="d",
+            )
+            where_clause = f"{FTS_TABLE_NAME} MATCH :query"
+            if filters_sql:
+                where_clause = f"{where_clause} AND {' AND '.join(filters_sql)}"
+            params = {
+                "query": match_query,
+                "limit": int(limit),
+                "offset": int(offset),
+                **filter_params,
+            }
+            try:
+                count_result = await session.execute(
+                    text(
+                        f"""
+                        SELECT count(*)
+                        FROM {FTS_TABLE_NAME}
+                        JOIN documents d ON d.id = {FTS_TABLE_NAME}.rowid
+                        WHERE {where_clause}
+                        """,
+                    ),
+                    params,
+                )
+                total = int(count_result.scalar_one_or_none() or 0)
+                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 {where_clause}
+                        ORDER BY score ASC, d.id ASC
+                        LIMIT :limit
+                        OFFSET :offset
+                        """,
+                    ),
+                    params,
+                )
+            except Exception as e:
+                logger.warning(
+                    f"FTS5 document search failed for {self.db_path}: {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
+            ], total
+
+    async def _apply_metadata_filters(
+        self,
+        session: AsyncSession,
+        query,
+        metadata_filters: dict,
+    ):
+        filters_sql, params = await self._metadata_filter_sql(
+            session,
+            metadata_filters,
+        )
+        for filter_sql in filters_sql:
+            query = query.where(text(filter_sql))
+        if params:
+            query = query.params(**params)
+        return query
+
+    async def _metadata_filter_sql(
+        self,
+        session: AsyncSession,
+        metadata_filters: dict,
+        table_alias: str | None = None,
+    ) -> tuple[list[str], dict]:
+        columns = await self._table_columns(session, "documents")
+        prefix = f"{table_alias}." if table_alias else ""
+        filters_sql = []
+        params = {}
+        for key, val in metadata_filters.items():
+            if key in {"kb_id", "kb_doc_id", "user_id"} and key in columns:
+                filters_sql.append(f"{prefix}{key} = :filter_{key}")
+            else:
+                filters_sql.append(
+                    f"json_extract({prefix}metadata, '$.{key}') = :filter_{key}"
+                )
+            params[f"filter_{key}"] = val
+        return filters_sql, params
+
     async def ensure_fts_index(self) -> bool:
         """Ensure the FTS5 sparse index exists and matches the documents table."""
         if not self.fts5_available:
@@ -470,22 +653,30 @@ async def ensure_fts_index(self) -> bool:
 
         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
+        async with self._fts_rebuild_lock:
+            if self._fts_index_ready:
                 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 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_unlocked()
+            return self.fts5_available
 
     async def rebuild_fts_index(self) -> None:
         """Rebuild the contentless FTS5 sparse index from documents."""
+        async with self._fts_rebuild_lock:
+            await self._rebuild_fts_index_unlocked()
+
+    async def _rebuild_fts_index_unlocked(self) -> None:
         if not self.fts5_available:
             return
 
@@ -530,7 +721,7 @@ async def search_sparse(
         sparse retrieval implementation.
         """
         if limit <= 0:
-            return []
+            return None
         if not await self.ensure_fts_index():
             return None