Merge branch 'AstrBotDevs:master' into cmd_suggestion

Author: elecvoid243

astrbot/core/db/__init__.py

@@ -5,7 +5,10 @@
 from dataclasses import dataclass
 
 from deprecated import deprecated
+from sqlalchemy import event
+from sqlalchemy.engine import make_url
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from sqlalchemy.pool import NullPool
 
 from astrbot.core.db.po import (
     ApiKey,
@@ -28,6 +31,19 @@
 )
 
 
+def _configure_sqlite_connection(dbapi_connection, connection_record) -> None:
+    cursor = dbapi_connection.cursor()
+    try:
+        cursor.execute("PRAGMA journal_mode=WAL")
+        cursor.execute("PRAGMA synchronous=NORMAL")
+        cursor.execute("PRAGMA cache_size=20000")
+        cursor.execute("PRAGMA temp_store=MEMORY")
+        cursor.execute("PRAGMA mmap_size=134217728")
+        cursor.execute("PRAGMA optimize")
+    finally:
+        cursor.close()
+
+
 @dataclass
 class BaseDatabase(abc.ABC):
     """数据库基类"""
@@ -40,14 +56,29 @@ def __init__(self) -> None:
         # second write is attempted.  Setting timeout=30 tells SQLite to
         # wait up to 30 s for the lock, which is enough to ride out brief
         # write bursts from concurrent agent/metrics/session operations.
-        is_sqlite = "sqlite" in self.DATABASE_URL
+        db_url = make_url(self.DATABASE_URL)
+        is_sqlite = db_url.get_backend_name() == "sqlite"
         connect_args = {"timeout": 30} if is_sqlite else {}
+        engine_kwargs = {
+            "echo": False,
+            "future": True,
+            "connect_args": connect_args,
+        }
+        if is_sqlite:
+            # Keep SQLite async engines off SQLAlchemy's default async queue
+            # pool so packaged runtimes don't depend on dialect-specific pool
+            # event support.
+            engine_kwargs["poolclass"] = NullPool
         self.engine = create_async_engine(
             self.DATABASE_URL,
-            echo=False,
-            future=True,
-            connect_args=connect_args,
+            **engine_kwargs,
         )
+        if is_sqlite:
+            event.listen(
+                self.engine.sync_engine,
+                "connect",
+                _configure_sqlite_connection,
+            )
         self.AsyncSessionLocal = async_sessionmaker(
             self.engine,
             class_=AsyncSession,

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

@@ -5,8 +5,11 @@
 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
@@ -34,7 +37,7 @@ class Document(BaseDocModel, table=True):
         primary_key=True,
         sa_column_kwargs={"autoincrement": True},
     )
-    doc_id: str = Field(nullable=False)
+    doc_id: str = Field(nullable=False, unique=True)
     text: str = Field(nullable=False)
     metadata_: str | None = Field(default=None, sa_column=Column("metadata", Text))
     created_at: datetime | None = Field(default=None)
@@ -60,8 +63,7 @@ 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)
 
             try:
                 await conn.execute(
@@ -94,6 +96,56 @@ async def initialize(self) -> None:
             await self._initialize_fts5(conn)
             await conn.commit()
 
+    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:
             await self._create_fts5_table(executor, if_not_exists=True)
@@ -197,6 +249,7 @@ async def connect(self) -> None:
                 self.DATABASE_URL,
                 echo=False,
                 future=True,
+                poolclass=NullPool,
             )
             self.async_session_maker = sessionmaker(
                 self.engine,  # type: ignore

astrbot/core/knowledge_base/kb_db_sqlite.py

@@ -2,8 +2,9 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from sqlalchemy import delete, func, select, text, update
+from sqlalchemy import delete, event, func, select, text, update
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from sqlalchemy.pool import NullPool
 from sqlmodel import col, desc
 
 from astrbot.core import logger
@@ -19,6 +20,19 @@
     from astrbot.core.db.vec_db.faiss_impl import FaissVecDB
 
 
+def _configure_sqlite_connection(dbapi_connection, connection_record) -> None:
+    cursor = dbapi_connection.cursor()
+    try:
+        cursor.execute("PRAGMA journal_mode=WAL")
+        cursor.execute("PRAGMA synchronous=NORMAL")
+        cursor.execute("PRAGMA cache_size=20000")
+        cursor.execute("PRAGMA temp_store=MEMORY")
+        cursor.execute("PRAGMA mmap_size=134217728")
+        cursor.execute("PRAGMA optimize")
+    finally:
+        cursor.close()
+
+
 class KBSQLiteDatabase:
     def __init__(self, db_path: str | None = None) -> None:
         """初始化知识库数据库
@@ -40,8 +54,12 @@ def __init__(self, db_path: str | None = None) -> None:
         self.engine = create_async_engine(
             self.DATABASE_URL,
             echo=False,
-            pool_pre_ping=True,
-            pool_recycle=3600,
+            poolclass=NullPool,
+        )
+        event.listen(
+            self.engine.sync_engine,
+            "connect",
+            _configure_sqlite_connection,
         )
 
         # 创建会话工厂

main.py

@@ -27,6 +27,7 @@
     get_dashboard_version,
     should_use_bundled_dashboard_dist,
 )
+from astrbot.core.utils.runtime_env import is_packaged_desktop_runtime  # noqa: E402
 
 # 将父目录添加到 sys.path
 sys.path.append(Path(__file__).parent.as_posix())
@@ -52,7 +53,10 @@ def check_env() -> None:
         sys.path.insert(0, astrbot_root)
 
     site_packages_path = get_astrbot_site_packages_path()
-    if site_packages_path not in sys.path:
+    if not is_packaged_desktop_runtime() and site_packages_path not in sys.path:
+        # Packaged desktop runtime keeps shared plugin dependencies out of the
+        # global import path so bundled core libraries don't mix with user-
+        # installed wheels from ~/.astrbot/data/site-packages.
         sys.path.append(site_packages_path)
 
     os.makedirs(get_astrbot_config_path(), exist_ok=True)

tests/unit/test_document_storage_fts.py

@@ -1,6 +1,7 @@
 import sqlite3
 
 import pytest
+from sqlalchemy.exc import IntegrityError
 
 from astrbot.core.db.vec_db.faiss_impl.document_storage import DocumentStorage
 
@@ -101,3 +102,38 @@ async def test_document_storage_fts_recovers_from_legacy_non_fts_table(tmp_path)
     assert [result["doc_id"] for result in results] == ["legacy-fix"]
 
     await storage.close()
+
+
+@pytest.mark.asyncio
+async def test_document_storage_adds_unique_doc_id_index_to_existing_table(tmp_path):
+    db_path = tmp_path / "doc.db"
+    conn = sqlite3.connect(db_path)
+    conn.execute(
+        """
+        CREATE TABLE documents (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            doc_id VARCHAR NOT NULL,
+            text VARCHAR NOT NULL,
+            metadata TEXT,
+            created_at DATETIME,
+            updated_at DATETIME
+        )
+        """,
+    )
+    conn.execute(
+        "INSERT INTO documents (doc_id, text) VALUES ('legacy-chunk', 'legacy text')"
+    )
+    conn.commit()
+    conn.close()
+
+    storage = DocumentStorage(str(db_path))
+    await storage.initialize()
+
+    with pytest.raises(IntegrityError):
+        await storage.insert_document(
+            doc_id="legacy-chunk",
+            text="duplicate text",
+            metadata={},
+        )
+
+    await storage.close()