AstrBot 4.3.5__py3-none-any.whl → 4.5.1__py3-none-any.whl

astrbot/core/db/migration/helper.py

@@ -17,8 +17,11 @@ async def check_migration_needed_v4(db_helper: BaseDatabase) -> bool:
     检查是否需要进行数据库迁移
     如果存在 data_v3.db 并且 preference 中没有 migration_done_v4，则需要进行迁移。
     """
-    data_v3_exists = os.path.exists(get_astrbot_data_path())
-    if not data_v3_exists:
+    # 仅当 data 目录下存在旧版本数据（data_v3.db 文件）时才考虑迁移
+    data_dir = get_astrbot_data_path()
+    data_v3_db = os.path.join(data_dir, "data_v3.db")
+
+    if not os.path.exists(data_v3_db):
         return False
     migration_done = await db_helper.get_preference(
         "global", "global", "migration_done_v4"
@@ -32,7 +35,7 @@ async def do_migration_v4(
     db_helper: BaseDatabase,
     platform_id_map: dict[str, dict[str, str]],
     astrbot_config: AstrBotConfig,
-):
+) -> None:
     """
     执行数据库迁移
     迁移旧的 webchat_conversation 表到新的 conversation 表。

astrbot/core/db/migration/migra_45_to_46.py

@@ -0,0 +1,44 @@
+from astrbot.api import logger, sp
+from astrbot.core.astrbot_config_mgr import AstrBotConfigManager
+from astrbot.core.umop_config_router import UmopConfigRouter
+
+
+async def migrate_45_to_46(acm: AstrBotConfigManager, ucr: UmopConfigRouter):
+    abconf_data = acm.abconf_data
+
+    if not isinstance(abconf_data, dict):
+        # should be unreachable
+        logger.warning(
+            f"migrate_45_to_46: abconf_data is not a dict (type={type(abconf_data)}). Value: {abconf_data!r}"
+        )
+        return
+
+    # 如果任何一项带有 umop，则说明需要迁移
+    need_migration = False
+    for conf_id, conf_info in abconf_data.items():
+        if isinstance(conf_info, dict) and "umop" in conf_info:
+            need_migration = True
+            break
+
+    if not need_migration:
+        return
+
+    logger.info("Starting migration from version 4.5 to 4.6")
+
+    # extract umo->conf_id mapping
+    umo_to_conf_id = {}
+    for conf_id, conf_info in abconf_data.items():
+        if isinstance(conf_info, dict) and "umop" in conf_info:
+            umop_ls = conf_info.pop("umop")
+            if not isinstance(umop_ls, list):
+                continue
+            for umo in umop_ls:
+                if isinstance(umo, str) and umo not in umo_to_conf_id:
+                    umo_to_conf_id[umo] = conf_id
+
+    # update the abconf data
+    await sp.global_put("abconf_mapping", abconf_data)
+    # update the umop config router
+    await ucr.update_routing_data(umo_to_conf_id)
+
+    logger.info("Migration from version 45 to 46 completed successfully")

astrbot/core/db/vec_db/base.py

@@ -16,14 +16,42 @@ class BaseVecDB:
         pass
 
     @abc.abstractmethod
-    async def insert(self, content: str, metadata: dict = None, id: str = None) -> int:
+    async def insert(
+        self, content: str, metadata: dict | None = None, id: str | None = None
+    ) -> int:
         """
         插入一条文本和其对应向量，自动生成 ID 并保持一致性。
         """
         ...
 
     @abc.abstractmethod
-    async def retrieve(self, query: str, top_k: int = 5) -> list[Result]:
+    async def insert_batch(
+        self,
+        contents: list[str],
+        metadatas: list[dict] | None = None,
+        ids: list[str] | None = None,
+        batch_size: int = 32,
+        tasks_limit: int = 3,
+        max_retries: int = 3,
+        progress_callback=None,
+    ) -> int:
+        """
+        批量插入文本和其对应向量，自动生成 ID 并保持一致性。
+
+        Args:
+            progress_callback: 进度回调函数，接收参数 (current, total)
+        """
+        ...
+
+    @abc.abstractmethod
+    async def retrieve(
+        self,
+        query: str,
+        top_k: int = 5,
+        fetch_k: int = 20,
+        rerank: bool = False,
+        metadata_filters: dict | None = None,
+    ) -> list[Result]:
         """
         搜索最相似的文档。
         Args:
@@ -44,3 +72,6 @@ class BaseVecDB:
             bool: 删除是否成功
         """
         ...
+
+    @abc.abstractmethod
+    async def close(self): ...

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

@@ -1,59 +1,224 @@
-import aiosqlite
 import os
+import json
+from datetime import datetime
+from contextlib import asynccontextmanager
+
+from sqlalchemy import Text, Column
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
+from sqlalchemy.orm import sessionmaker
+from sqlmodel import Field, SQLModel, select, col, func, text, MetaData
+from astrbot.core import logger
+
+
+class BaseDocModel(SQLModel, table=False):
+    metadata = MetaData()
+
+
+class Document(BaseDocModel, table=True):
+    """SQLModel for documents table."""
+
+    __tablename__ = "documents"  # type: ignore
+
+    id: int | None = Field(
+        default=None, primary_key=True, sa_column_kwargs={"autoincrement": True}
+    )
+    doc_id: str = Field(nullable=False)
+    text: str = Field(nullable=False)
+    metadata_: str | None = Field(default=None, sa_column=Column("metadata", Text))
+    created_at: datetime | None = Field(default=None)
+    updated_at: datetime | None = Field(default=None)
 
 
 class DocumentStorage:
     def __init__(self, db_path: str):
         self.db_path = db_path
-        self.connection = None
+        self.DATABASE_URL = f"sqlite+aiosqlite:///{db_path}"
+        self.engine: AsyncEngine | None = None
+        self.async_session_maker: sessionmaker | None = None
         self.sqlite_init_path = os.path.join(
             os.path.dirname(__file__), "sqlite_init.sql"
         )
 
     async def initialize(self):
         """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()
-            async with self.connection.cursor() as cursor:
-                with open(self.sqlite_init_path, "r", encoding="utf-8") as f:
-                    sql_script = f.read()
-                await cursor.executescript(sql_script)
-            await self.connection.commit()
-        else:
-            await self.connect()
+        await self.connect()
+        async with self.engine.begin() as conn:  # type: ignore
+            # Create tables using SQLModel
+            await conn.run_sync(BaseDocModel.metadata.create_all)
+
+            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"
+                    )
+                )
+
+                # 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
+
+            await conn.commit()
 
     async def connect(self):
         """Connect to the SQLite database."""
-        self.connection = await aiosqlite.connect(self.db_path)
+        if self.engine is None:
+            self.engine = create_async_engine(
+                self.DATABASE_URL,
+                echo=False,
+                future=True,
+            )
+            self.async_session_maker = sessionmaker(
+                self.engine,  # type: ignore
+                class_=AsyncSession,
+                expire_on_commit=False,
+            )  # type: ignore
+
+    @asynccontextmanager
+    async def get_session(self):
+        """Context manager for database sessions."""
+        async with self.async_session_maker() as session:  # type: ignore
+            yield session
 
-    async def get_documents(self, metadata_filters: dict, ids: list = None):
+    async def get_documents(
+        self,
+        metadata_filters: dict,
+        ids: list | None = None,
+        offset: int | None = 0,
+        limit: int | None = 100,
+    ) -> list[dict]:
         """Retrieve documents by metadata filters and ids.
 
         Args:
             metadata_filters (dict): The metadata filters to apply.
+            ids (list | None): Optional list of document IDs to filter.
+            offset (int | None): Offset for pagination.
+            limit (int | None): Limit for pagination.
+
+        Returns:
+            list: The list of documents that match the filters.
+        """
+        if self.engine is None:
+            logger.warning(
+                "Database connection is not initialized, returning empty result"
+            )
+            return []
+
+        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})
+
+            if ids is not None and len(ids) > 0:
+                valid_ids = [int(i) for i in ids if i != -1]
+                if valid_ids:
+                    query = query.where(col(Document.id).in_(valid_ids))
+
+            if limit is not None:
+                query = query.limit(limit)
+            if offset is not None:
+                query = query.offset(offset)
+
+            result = await session.execute(query)
+            documents = result.scalars().all()
+
+            return [self._document_to_dict(doc) for doc in documents]
+
+    async def insert_document(self, doc_id: str, text: str, metadata: dict) -> int:
+        """Insert a single document and return its integer ID.
+
+        Args:
+            doc_id (str): The document ID (UUID string).
+            text (str): The document text.
+            metadata (dict): The document metadata.
 
         Returns:
-            list: The list of document IDs(primary key, not doc_id) that match the filters.
+            int: The integer ID of the inserted document.
+        """
+        assert self.engine is not None, "Database connection is not initialized."
+
+        async with self.get_session() as session:
+            async with session.begin():
+                document = Document(
+                    doc_id=doc_id,
+                    text=text,
+                    metadata_=json.dumps(metadata),
+                    created_at=datetime.now(),
+                    updated_at=datetime.now(),
+                )
+                session.add(document)
+                await session.flush()  # Flush to get the ID
+                return document.id  # type: ignore
+
+    async def insert_documents_batch(
+        self, doc_ids: list[str], texts: list[str], metadatas: list[dict]
+    ) -> list[int]:
+        """Batch insert documents and return their integer IDs.
+
+        Args:
+            doc_ids (list[str]): List of document IDs (UUID strings).
+            texts (list[str]): List of document texts.
+            metadatas (list[dict]): List of document metadata.
+
+        Returns:
+            list[int]: List of integer IDs of the inserted documents.
+        """
+        assert self.engine is not None, "Database connection is not initialized."
+
+        async with self.get_session() as session:
+            async with session.begin():
+                import json
+
+                documents = []
+                for doc_id, text, metadata in zip(doc_ids, texts, metadatas):
+                    document = Document(
+                        doc_id=doc_id,
+                        text=text,
+                        metadata_=json.dumps(metadata),
+                        created_at=datetime.now(),
+                        updated_at=datetime.now(),
+                    )
+                    documents.append(document)
+                    session.add(document)
+
+                await session.flush()  # Flush to get all IDs
+                return [doc.id for doc in documents]  # type: ignore
+
+    async def delete_document_by_doc_id(self, doc_id: str):
+        """Delete a document by its doc_id.
+
+        Args:
+            doc_id (str): The doc_id of the document to delete.
         """
-        # metadata filter -> SQL WHERE clause
-        where_clauses = []
-        values = []
-        for key, val in metadata_filters.items():
-            where_clauses.append(f"json_extract(metadata, '$.{key}') = ?")
-            values.append(val)
-        if ids is not None and len(ids) > 0:
-            ids = [str(i) for i in ids if i != -1]
-            where_clauses.append("id IN ({})".format(",".join("?" * len(ids))))
-            values.extend(ids)
-        where_sql = " AND ".join(where_clauses) or "1=1"
-
-        result = []
-        async with self.connection.cursor() as cursor:
-            sql = "SELECT * FROM documents WHERE " + where_sql
-            await cursor.execute(sql, values)
-            for row in await cursor.fetchall():
-                result.append(await self.tuple_to_dict(row))
-        return result
+        assert self.engine is not None, "Database connection is not initialized."
+
+        async with self.get_session() as session:
+            async with session.begin():
+                query = select(Document).where(col(Document.doc_id) == doc_id)
+                result = await session.execute(query)
+                document = result.scalar_one_or_none()
+
+                if document:
+                    await session.delete(document)
 
     async def get_document_by_doc_id(self, doc_id: str):
         """Retrieve a document by its doc_id.
@@ -62,28 +227,91 @@ class DocumentStorage:
             doc_id (str): The doc_id of the document to retrieve.
 
         Returns:
-            dict: The document data.
+            dict: The document data or None if not found.
         """
-        async with self.connection.cursor() as cursor:
-            await cursor.execute("SELECT * FROM documents WHERE doc_id = ?", (doc_id,))
-            row = await cursor.fetchone()
-            if row:
-                return await self.tuple_to_dict(row)
-            else:
-                return None
+        assert self.engine is not None, "Database connection is not initialized."
+
+        async with self.get_session() as session:
+            query = select(Document).where(col(Document.doc_id) == doc_id)
+            result = await session.execute(query)
+            document = result.scalar_one_or_none()
+
+            if document:
+                return self._document_to_dict(document)
+            return None
 
     async def update_document_by_doc_id(self, doc_id: str, new_text: str):
-        """Retrieve a document by its doc_id.
+        """Update a document by its doc_id.
 
         Args:
             doc_id (str): The doc_id.
             new_text (str): The new text to update the document with.
         """
-        async with self.connection.cursor() as cursor:
-            await cursor.execute(
-                "UPDATE documents SET text = ? WHERE doc_id = ?", (new_text, doc_id)
+        assert self.engine is not None, "Database connection is not initialized."
+
+        async with self.get_session() as session:
+            async with session.begin():
+                query = select(Document).where(col(Document.doc_id) == doc_id)
+                result = await session.execute(query)
+                document = result.scalar_one_or_none()
+
+                if document:
+                    document.text = new_text
+                    document.updated_at = datetime.now()
+                    session.add(document)
+
+    async def delete_documents(self, metadata_filters: dict):
+        """Delete documents by their metadata filters.
+
+        Args:
+            metadata_filters (dict): The metadata filters to apply.
+        """
+        if self.engine is None:
+            logger.warning(
+                "Database connection is not initialized, skipping delete operation"
             )
-            await self.connection.commit()
+            return
+
+        async with self.get_session() as session:
+            async with 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})
+
+                result = await session.execute(query)
+                documents = result.scalars().all()
+
+                for doc in documents:
+                    await session.delete(doc)
+
+    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.
+        """
+        if self.engine is None:
+            logger.warning("Database connection is not initialized, returning 0")
+            return 0
+
+        async with self.get_session() as session:
+            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})
+
+            result = await session.execute(query)
+            count = result.scalar_one_or_none()
+            return count if count is not None else 0
 
     async def get_user_ids(self) -> list[str]:
         """Retrieve all user IDs from the documents table.
@@ -91,11 +319,38 @@ class DocumentStorage:
         Returns:
             list: A list of user IDs.
         """
-        async with self.connection.cursor() as cursor:
-            await cursor.execute("SELECT DISTINCT user_id FROM documents")
-            rows = await cursor.fetchall()
+        assert self.engine is not None, "Database connection is not initialized."
+
+        async with self.get_session() as session:
+            query = text(
+                "SELECT DISTINCT user_id FROM documents WHERE user_id IS NOT NULL"
+            )
+            result = await session.execute(query)
+            rows = result.fetchall()
             return [row[0] for row in rows]
 
+    def _document_to_dict(self, document: Document) -> dict:
+        """Convert a Document model to a dictionary.
+
+        Args:
+            document (Document): The document to convert.
+
+        Returns:
+            dict: The converted dictionary.
+        """
+        return {
+            "id": document.id,
+            "doc_id": document.doc_id,
+            "text": document.text,
+            "metadata": document.metadata_,
+            "created_at": document.created_at.isoformat()
+            if isinstance(document.created_at, datetime)
+            else document.created_at,
+            "updated_at": document.updated_at.isoformat()
+            if isinstance(document.updated_at, datetime)
+            else document.updated_at,
+        }
+
     async def tuple_to_dict(self, row):
         """Convert a tuple to a dictionary.
 
@@ -104,6 +359,8 @@ class DocumentStorage:
 
         Returns:
             dict: The converted dictionary.
+
+        Note: This method is kept for backward compatibility but is no longer used internally.
         """
         return {
             "id": row[0],
@@ -116,6 +373,7 @@ class DocumentStorage:
 
     async def close(self):
         """Close the connection to the SQLite database."""
-        if self.connection:
-            await self.connection.close()
-            self.connection = None
+        if self.engine:
+            await self.engine.dispose()
+            self.engine = None
+            self.async_session_maker = None

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,29 @@ 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 insert_batch(self, vectors: np.ndarray, ids: list[int]):
+        """批量插入向量
+
+        Args:
+            vectors (np.ndarray): 要插入的向量数组
+            ids (list[int]): 向量的ID列表
+        Raises:
+            ValueError: 如果向量的维度与存储的维度不匹配
+        """
+        assert self.index is not None, "FAISS index is not initialized."
+        if vectors.shape[1] != self.dimension:
+            raise ValueError(
+                f"向量维度不匹配, 期望: {self.dimension}, 实际: {vectors.shape[1]}"
+            )
+        self.index.add_with_ids(vectors, np.array(ids))
         await self.save_index()
 
     async def search(self, vector: np.ndarray, k: int) -> tuple:
@@ -46,10 +62,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):
         """保存索引