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

astrbot/core/db/__init__.py

@@ -1,27 +1,27 @@
 import abc
 import datetime
 import typing as T
-from deprecated import deprecated
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
+
+from deprecated import deprecated
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
+from sqlalchemy.orm import sessionmaker
+
 from astrbot.core.db.po import (
-    Stats,
-    PlatformStat,
-    ConversationV2,
-    PlatformMessageHistory,
     Attachment,
+    ConversationV2,
     Persona,
+    PlatformMessageHistory,
+    PlatformStat,
     Preference,
+    Stats,
 )
-from contextlib import asynccontextmanager
-from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
-from sqlalchemy.orm import sessionmaker
 
 
 @dataclass
 class BaseDatabase(abc.ABC):
-    """
-    数据库基类
-    """
+    """数据库基类"""
 
     DATABASE_URL = ""
 
@@ -32,12 +32,13 @@ class BaseDatabase(abc.ABC):
             future=True,
         )
         self.AsyncSessionLocal = sessionmaker(
-            self.engine, class_=AsyncSession, expire_on_commit=False
+            self.engine,
+            class_=AsyncSession,
+            expire_on_commit=False,
         )
 
     async def initialize(self):
         """初始化数据库连接"""
-        pass
 
     @asynccontextmanager
     async def get_db(self) -> T.AsyncGenerator[AsyncSession, None]:
@@ -91,7 +92,9 @@ class BaseDatabase(abc.ABC):
 
     @abc.abstractmethod
     async def get_conversations(
-        self, user_id: str | None = None, platform_id: str | None = None
+        self,
+        user_id: str | None = None,
+        platform_id: str | None = None,
     ) -> list[ConversationV2]:
         """Get all conversations for a specific user and platform_id(optional).
 
@@ -106,7 +109,9 @@ class BaseDatabase(abc.ABC):
 
     @abc.abstractmethod
     async def get_all_conversations(
-        self, page: int = 1, page_size: int = 20
+        self,
+        page: int = 1,
+        page_size: int = 20,
     ) -> list[ConversationV2]:
         """Get all conversations with pagination."""
         ...
@@ -173,7 +178,10 @@ class BaseDatabase(abc.ABC):
 
     @abc.abstractmethod
     async def delete_platform_message_offset(
-        self, platform_id: str, user_id: str, offset_sec: int = 86400
+        self,
+        platform_id: str,
+        user_id: str,
+        offset_sec: int = 86400,
     ) -> None:
         """Delete platform message history records older than the specified offset."""
         ...
@@ -243,7 +251,11 @@ class BaseDatabase(abc.ABC):
 
     @abc.abstractmethod
     async def insert_preference_or_update(
-        self, scope: str, scope_id: str, key: str, value: dict
+        self,
+        scope: str,
+        scope_id: str,
+        key: str,
+        value: dict,
     ) -> Preference:
         """Insert a new preference record."""
         ...
@@ -255,7 +267,10 @@ class BaseDatabase(abc.ABC):
 
     @abc.abstractmethod
     async def get_preferences(
-        self, scope: str, scope_id: str | None = None, key: str | None = None
+        self,
+        scope: str,
+        scope_id: str | None = None,
+        key: str | None = None,
     ) -> list[Preference]:
         """Get all preferences for a specific scope ID or key."""
         ...

astrbot/core/db/migration/helper.py

@@ -1,20 +1,21 @@
 import os
-from astrbot.core.utils.astrbot_path import get_astrbot_data_path
-from astrbot.core.db import BaseDatabase
-from astrbot.core.config import AstrBotConfig
+
 from astrbot.api import logger, sp
+from astrbot.core.config import AstrBotConfig
+from astrbot.core.db import BaseDatabase
+from astrbot.core.utils.astrbot_path import get_astrbot_data_path
+
 from .migra_3_to_4 import (
     migration_conversation_table,
-    migration_platform_table,
-    migration_webchat_data,
     migration_persona_data,
+    migration_platform_table,
     migration_preferences,
+    migration_webchat_data,
 )
 
 
 async def check_migration_needed_v4(db_helper: BaseDatabase) -> bool:
-    """
-    检查是否需要进行数据库迁移
+    """检查是否需要进行数据库迁移
     如果存在 data_v3.db 并且 preference 中没有 migration_done_v4，则需要进行迁移。
     """
     # 仅当 data 目录下存在旧版本数据（data_v3.db 文件）时才考虑迁移
@@ -24,7 +25,9 @@ async def check_migration_needed_v4(db_helper: BaseDatabase) -> bool:
     if not os.path.exists(data_v3_db):
         return False
     migration_done = await db_helper.get_preference(
-        "global", "global", "migration_done_v4"
+        "global",
+        "global",
+        "migration_done_v4",
     )
     if migration_done:
         return False
@@ -36,8 +39,7 @@ async def do_migration_v4(
     platform_id_map: dict[str, dict[str, str]],
     astrbot_config: AstrBotConfig,
 ) -> None:
-    """
-    执行数据库迁移
+    """执行数据库迁移
     迁移旧的 webchat_conversation 表到新的 conversation 表。
     迁移旧的 platform 到新的 platform_stats 表。
     """

astrbot/core/db/migration/migra_3_to_4.py

@@ -1,15 +1,18 @@
-import json
 import datetime
-from .. import BaseDatabase
-from .sqlite_v3 import SQLiteDatabase as SQLiteV3DatabaseV3
-from .shared_preferences_v3 import sp as sp_v3
-from astrbot.core.config.default import DB_PATH
+import json
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+
 from astrbot.api import logger, sp
 from astrbot.core.config import AstrBotConfig
-from astrbot.core.platform.astr_message_event import MessageSesion
-from sqlalchemy.ext.asyncio import AsyncSession
+from astrbot.core.config.default import DB_PATH
 from astrbot.core.db.po import ConversationV2, PlatformMessageHistory
-from sqlalchemy import text
+from astrbot.core.platform.astr_message_event import MessageSesion
+
+from .. import BaseDatabase
+from .shared_preferences_v3 import sp as sp_v3
+from .sqlite_v3 import SQLiteDatabase as SQLiteV3DatabaseV3
 
 """
 1. 迁移旧的 webchat_conversation 表到新的 conversation 表。
@@ -18,7 +21,8 @@ from sqlalchemy import text
 
 
 def get_platform_id(
-    platform_id_map: dict[str, dict[str, str]], old_platform_name: str
+    platform_id_map: dict[str, dict[str, str]],
+    old_platform_name: str,
 ) -> str:
     return platform_id_map.get(
         old_platform_name,
@@ -27,7 +31,8 @@ def get_platform_id(
 
 
 def get_platform_type(
-    platform_id_map: dict[str, dict[str, str]], old_platform_name: str
+    platform_id_map: dict[str, dict[str, str]],
+    old_platform_name: str,
 ) -> str:
     return platform_id_map.get(
         old_platform_name,
@@ -36,13 +41,15 @@ def get_platform_type(
 
 
 async def migration_conversation_table(
-    db_helper: BaseDatabase, platform_id_map: dict[str, dict[str, str]]
+    db_helper: BaseDatabase,
+    platform_id_map: dict[str, dict[str, str]],
 ):
     db_helper_v3 = SQLiteV3DatabaseV3(
-        db_path=DB_PATH.replace("data_v4.db", "data_v3.db")
+        db_path=DB_PATH.replace("data_v4.db", "data_v3.db"),
     )
     conversations, total_cnt = db_helper_v3.get_all_conversations(
-        page=1, page_size=10000000
+        page=1,
+        page_size=10000000,
     )
     logger.info(f"迁移 {total_cnt} 条旧的会话数据到新的表中...")
 
@@ -61,13 +68,14 @@ async def migration_conversation_table(
                     )
                     if not conv:
                         logger.info(
-                            f"未找到该条旧会话对应的具体数据: {conversation}, 跳过。"
+                            f"未找到该条旧会话对应的具体数据: {conversation}, 跳过。",
                         )
                     if ":" not in conv.user_id:
                         continue
                     session = MessageSesion.from_str(session_str=conv.user_id)
                     platform_id = get_platform_id(
-                        platform_id_map, session.platform_name
+                        platform_id_map,
+                        session.platform_name,
                     )
                     session.platform_id = platform_id  # 更新平台名称为新的 ID
                     conv_v2 = ConversationV2(
@@ -90,10 +98,11 @@ async def migration_conversation_table(
 
 
 async def migration_platform_table(
-    db_helper: BaseDatabase, platform_id_map: dict[str, dict[str, str]]
+    db_helper: BaseDatabase,
+    platform_id_map: dict[str, dict[str, str]],
 ):
     db_helper_v3 = SQLiteV3DatabaseV3(
-        db_path=DB_PATH.replace("data_v4.db", "data_v3.db")
+        db_path=DB_PATH.replace("data_v4.db", "data_v3.db"),
     )
     secs_from_2023_4_10_to_now = (
         datetime.datetime.now(datetime.timezone.utc)
@@ -134,10 +143,12 @@ async def migration_platform_table(
                 if cnt == 0:
                     continue
                 platform_id = get_platform_id(
-                    platform_id_map, platform_stats_v3[idx].name
+                    platform_id_map,
+                    platform_stats_v3[idx].name,
                 )
                 platform_type = get_platform_type(
-                    platform_id_map, platform_stats_v3[idx].name
+                    platform_id_map,
+                    platform_stats_v3[idx].name,
                 )
                 try:
                     await dbsession.execute(
@@ -149,7 +160,8 @@ async def migration_platform_table(
                         """),
                         {
                             "timestamp": datetime.datetime.fromtimestamp(
-                                bucket_end, tz=datetime.timezone.utc
+                                bucket_end,
+                                tz=datetime.timezone.utc,
                             ),
                             "platform_id": platform_id,
                             "platform_type": platform_type,
@@ -165,14 +177,16 @@ async def migration_platform_table(
 
 
 async def migration_webchat_data(
-    db_helper: BaseDatabase, platform_id_map: dict[str, dict[str, str]]
+    db_helper: BaseDatabase,
+    platform_id_map: dict[str, dict[str, str]],
 ):
     """迁移 WebChat 的历史记录到新的 PlatformMessageHistory 表中"""
     db_helper_v3 = SQLiteV3DatabaseV3(
-        db_path=DB_PATH.replace("data_v4.db", "data_v3.db")
+        db_path=DB_PATH.replace("data_v4.db", "data_v3.db"),
     )
     conversations, total_cnt = db_helper_v3.get_all_conversations(
-        page=1, page_size=10000000
+        page=1,
+        page_size=10000000,
     )
     logger.info(f"迁移 {total_cnt} 条旧的 WebChat 会话数据到新的表中...")
 
@@ -191,7 +205,7 @@ async def migration_webchat_data(
                     )
                     if not conv:
                         logger.info(
-                            f"未找到该条旧会话对应的具体数据: {conversation}, 跳过。"
+                            f"未找到该条旧会话对应的具体数据: {conversation}, 跳过。",
                         )
                     if ":" in conv.user_id:
                         continue
@@ -218,10 +232,10 @@ async def migration_webchat_data(
 
 
 async def migration_persona_data(
-    db_helper: BaseDatabase, astrbot_config: AstrBotConfig
+    db_helper: BaseDatabase,
+    astrbot_config: AstrBotConfig,
 ):
-    """
-    迁移 Persona 数据到新的表中。
+    """迁移 Persona 数据到新的表中。
     旧的 Persona 数据存储在 preference 中，新的 Persona 数据存储在 persona 表中。
     """
     v3_persona_config: list[dict] = astrbot_config.get("persona", [])
@@ -236,14 +250,15 @@ async def migration_persona_data(
         try:
             begin_dialogs = persona.get("begin_dialogs", [])
             mood_imitation_dialogs = persona.get("mood_imitation_dialogs", [])
-            mood_prompt = ""
+            parts = []
             user_turn = True
             for mood_dialog in mood_imitation_dialogs:
                 if user_turn:
-                    mood_prompt += f"A: {mood_dialog}\n"
+                    parts.append(f"A: {mood_dialog}\n")
                 else:
-                    mood_prompt += f"B: {mood_dialog}\n"
+                    parts.append(f"B: {mood_dialog}\n")
                 user_turn = not user_turn
+            mood_prompt = "".join(parts)
             system_prompt = persona.get("prompt", "")
             if mood_prompt:
                 system_prompt += f"Here are few shots of dialogs, you need to imitate the tone of 'B' in the following dialogs to respond:\n {mood_prompt}"
@@ -253,14 +268,15 @@ async def migration_persona_data(
                 begin_dialogs=begin_dialogs,
             )
             logger.info(
-                f"迁移 Persona {persona['name']}({persona_new.system_prompt[:30]}...) 到新表成功。"
+                f"迁移 Persona {persona['name']}({persona_new.system_prompt[:30]}...) 到新表成功。",
             )
         except Exception as e:
             logger.error(f"解析 Persona 配置失败：{e}")
 
 
 async def migration_preferences(
-    db_helper: BaseDatabase, platform_id_map: dict[str, dict[str, str]]
+    db_helper: BaseDatabase,
+    platform_id_map: dict[str, dict[str, str]],
 ):
     # 1. global scope migration
     keys = [
@@ -329,10 +345,13 @@ async def migration_preferences(
 
             for provider_type, provider_id in perf.items():
                 await sp.put_async(
-                    "umo", str(session), f"provider_perf_{provider_type}", provider_id
+                    "umo",
+                    str(session),
+                    f"provider_perf_{provider_type}",
+                    provider_id,
                 )
             logger.info(
-                f"迁移会话 {umo} 的提供商偏好到新表成功，平台 ID: {platform_id}"
+                f"迁移会话 {umo} 的提供商偏好到新表成功，平台 ID: {platform_id}",
             )
         except Exception as e:
             logger.error(f"迁移会话 {umo} 的提供商偏好失败: {e}", exc_info=True)

astrbot/core/db/migration/migra_45_to_46.py

@@ -9,7 +9,7 @@ async def migrate_45_to_46(acm: AstrBotConfigManager, ucr: UmopConfigRouter):
     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}"
+            f"migrate_45_to_46: abconf_data is not a dict (type={type(abconf_data)}). Value: {abconf_data!r}",
         )
         return
 

astrbot/core/db/migration/shared_preferences_v3.py

@@ -1,6 +1,7 @@
 import json
 import os
 from typing import TypeVar
+
 from astrbot.core.utils.astrbot_path import get_astrbot_data_path
 
 _VT = TypeVar("_VT")
@@ -16,7 +17,7 @@ class SharedPreferences:
     def _load_preferences(self):
         if os.path.exists(self.path):
             try:
-                with open(self.path, "r") as f:
+                with open(self.path) as f:
                     return json.load(f)
             except json.JSONDecodeError:
                 os.remove(self.path)

astrbot/core/db/migration/sqlite_v3.py

@@ -1,8 +1,9 @@
 import sqlite3
 import time
-from astrbot.core.db.po import Platform, Stats
-from typing import Tuple, List, Dict, Any
 from dataclasses import dataclass
+from typing import Any
+
+from astrbot.core.db.po import Platform, Stats
 
 
 @dataclass
@@ -94,7 +95,7 @@ class SQLiteDatabase:
         c.execute(
             """
             PRAGMA table_info(webchat_conversation)
-            """
+            """,
         )
         res = c.fetchall()
         has_title = False
@@ -108,14 +109,14 @@ class SQLiteDatabase:
             c.execute(
                 """
                 ALTER TABLE webchat_conversation ADD COLUMN title TEXT;
-                """
+                """,
             )
             self.conn.commit()
         if not has_persona_id:
             c.execute(
                 """
                 ALTER TABLE webchat_conversation ADD COLUMN persona_id TEXT;
-                """
+                """,
             )
             self.conn.commit()
 
@@ -126,7 +127,7 @@ class SQLiteDatabase:
         conn.text_factory = str
         return conn
 
-    def _exec_sql(self, sql: str, params: Tuple = None):
+    def _exec_sql(self, sql: str, params: tuple = None):
         conn = self.conn
         try:
             c = self.conn.cursor()
@@ -174,7 +175,7 @@ class SQLiteDatabase:
             """
             SELECT * FROM platform
             """
-            + where_clause
+            + where_clause,
         )
 
         platform = []
@@ -194,7 +195,7 @@ class SQLiteDatabase:
         c.execute(
             """
             SELECT SUM(count) FROM platform
-            """
+            """,
         )
         res = c.fetchone()
         c.close()
@@ -214,7 +215,7 @@ class SQLiteDatabase:
             SELECT name, SUM(count), timestamp FROM platform
             """
             + where_clause
-            + " GROUP BY name"
+            + " GROUP BY name",
         )
 
         platform = []
@@ -242,7 +243,7 @@ class SQLiteDatabase:
         c.close()
 
         if not res:
-            return
+            return None
 
         return Conversation(*res)
 
@@ -257,7 +258,7 @@ class SQLiteDatabase:
             (user_id, cid, history, updated_at, created_at),
         )
 
-    def get_conversations(self, user_id: str) -> Tuple:
+    def get_conversations(self, user_id: str) -> tuple:
         try:
             c = self.conn.cursor()
         except sqlite3.ProgrammingError:
@@ -280,7 +281,7 @@ class SQLiteDatabase:
             title = row[3]
             persona_id = row[4]
             conversations.append(
-                Conversation("", cid, "[]", created_at, updated_at, title, persona_id)
+                Conversation("", cid, "[]", created_at, updated_at, title, persona_id),
             )
         return conversations
 
@@ -319,8 +320,10 @@ class SQLiteDatabase:
         )
 
     def get_all_conversations(
-        self, page: int = 1, page_size: int = 20
-    ) -> Tuple[List[Dict[str, Any]], int]:
+        self,
+        page: int = 1,
+        page_size: int = 20,
+    ) -> tuple[list[dict[str, Any]], int]:
         """获取所有对话，支持分页，按更新时间降序排序"""
         try:
             c = self.conn.cursor()
@@ -366,7 +369,7 @@ class SQLiteDatabase:
                         "persona_id": persona_id or "",
                         "created_at": created_at or 0,
                         "updated_at": updated_at or 0,
-                    }
+                    },
                 )
 
             return conversations, total_count
@@ -381,12 +384,12 @@ class SQLiteDatabase:
         self,
         page: int = 1,
         page_size: int = 20,
-        platforms: List[str] = None,
-        message_types: List[str] = None,
-        search_query: str = None,
-        exclude_ids: List[str] = None,
-        exclude_platforms: List[str] = None,
-    ) -> Tuple[List[Dict[str, Any]], int]:
+        platforms: list[str] | None = None,
+        message_types: list[str] | None = None,
+        search_query: str | None = None,
+        exclude_ids: list[str] | None = None,
+        exclude_platforms: list[str] | None = None,
+    ) -> tuple[list[dict[str, Any]], int]:
         """获取筛选后的对话列表"""
         try:
             c = self.conn.cursor()
@@ -422,7 +425,7 @@ class SQLiteDatabase:
             if search_query:
                 search_query = search_query.encode("unicode_escape").decode("utf-8")
                 where_clauses.append(
-                    "(title LIKE ? OR user_id LIKE ? OR cid LIKE ? OR history LIKE ?)"
+                    "(title LIKE ? OR user_id LIKE ? OR cid LIKE ? OR history LIKE ?)",
                 )
                 search_param = f"%{search_query}%"
                 params.extend([search_param, search_param, search_param, search_param])
@@ -482,7 +485,7 @@ class SQLiteDatabase:
                         "persona_id": persona_id or "",
                         "created_at": created_at or 0,
                         "updated_at": updated_at or 0,
-                    }
+                    },
                 )
 
             return conversations, total_count

astrbot/core/db/po.py

@@ -1,15 +1,15 @@
 import uuid
-
-from datetime import datetime, timezone
 from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import TypedDict
+
 from sqlmodel import (
+    JSON,
+    Field,
     SQLModel,
     Text,
-    JSON,
     UniqueConstraint,
-    Field,
 )
-from typing import Optional, TypedDict
 
 
 class PlatformStat(SQLModel, table=True):
@@ -40,7 +40,8 @@ class ConversationV2(SQLModel, table=True):
     __tablename__ = "conversations"
 
     inner_conversation_id: int = Field(
-        primary_key=True, sa_column_kwargs={"autoincrement": True}
+        primary_key=True,
+        sa_column_kwargs={"autoincrement": True},
     )
     conversation_id: str = Field(
         max_length=36,
@@ -50,14 +51,14 @@ class ConversationV2(SQLModel, table=True):
     )
     platform_id: str = Field(nullable=False)
     user_id: str = Field(nullable=False)
-    content: Optional[list] = Field(default=None, sa_type=JSON)
+    content: list | None = Field(default=None, sa_type=JSON)
     created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
     updated_at: datetime = Field(
         default_factory=lambda: datetime.now(timezone.utc),
         sa_column_kwargs={"onupdate": datetime.now(timezone.utc)},
     )
-    title: Optional[str] = Field(default=None, max_length=255)
-    persona_id: Optional[str] = Field(default=None)
+    title: str | None = Field(default=None, max_length=255)
+    persona_id: str | None = Field(default=None)
 
     __table_args__ = (
         UniqueConstraint(
@@ -76,13 +77,15 @@ class Persona(SQLModel, table=True):
     __tablename__ = "personas"
 
     id: int | None = Field(
-        primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None
+        primary_key=True,
+        sa_column_kwargs={"autoincrement": True},
+        default=None,
     )
     persona_id: str = Field(max_length=255, nullable=False)
     system_prompt: str = Field(sa_type=Text, nullable=False)
-    begin_dialogs: Optional[list] = Field(default=None, sa_type=JSON)
+    begin_dialogs: list | None = Field(default=None, sa_type=JSON)
     """a list of strings, each representing a dialog to start with"""
-    tools: Optional[list] = Field(default=None, sa_type=JSON)
+    tools: list | None = Field(default=None, sa_type=JSON)
     """None means use ALL tools for default, empty list means no tools, otherwise a list of tool names."""
     created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
     updated_at: datetime = Field(
@@ -104,7 +107,9 @@ class Preference(SQLModel, table=True):
     __tablename__ = "preferences"
 
     id: int | None = Field(
-        default=None, primary_key=True, sa_column_kwargs={"autoincrement": True}
+        default=None,
+        primary_key=True,
+        sa_column_kwargs={"autoincrement": True},
     )
     scope: str = Field(nullable=False)
     """Scope of the preference, such as 'global', 'umo', 'plugin'."""
@@ -138,13 +143,15 @@ class PlatformMessageHistory(SQLModel, table=True):
     __tablename__ = "platform_message_history"
 
     id: int | None = Field(
-        primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None
+        primary_key=True,
+        sa_column_kwargs={"autoincrement": True},
+        default=None,
     )
     platform_id: str = Field(nullable=False)
     user_id: str = Field(nullable=False)  # An id of group, user in platform
-    sender_id: Optional[str] = Field(default=None)  # ID of the sender in the platform
-    sender_name: Optional[str] = Field(
-        default=None
+    sender_id: str | None = Field(default=None)  # ID of the sender in the platform
+    sender_name: str | None = Field(
+        default=None,
     )  # Name of the sender in the platform
     content: dict = Field(sa_type=JSON, nullable=False)  # a message chain list
     created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
@@ -163,7 +170,9 @@ class Attachment(SQLModel, table=True):
     __tablename__ = "attachments"
 
     inner_attachment_id: int | None = Field(
-        primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None
+        primary_key=True,
+        sa_column_kwargs={"autoincrement": True},
+        default=None,
     )
     attachment_id: str = Field(
         max_length=36,