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

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,

astrbot/core/db/sqlite.py

@@ -1,22 +1,27 @@
 import asyncio
-import typing as T
 import threading
+import typing as T
 from datetime import datetime, timedelta
+
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlmodel import col, delete, desc, func, or_, select, text, update
+
 from astrbot.core.db import BaseDatabase
 from astrbot.core.db.po import (
-    ConversationV2,
-    PlatformStat,
-    PlatformMessageHistory,
     Attachment,
+    ConversationV2,
     Persona,
+    PlatformMessageHistory,
+    PlatformStat,
     Preference,
-    Stats as DeprecatedStats,
-    Platform as DeprecatedPlatformStat,
     SQLModel,
 )
-
-from sqlmodel import select, update, delete, text, func, or_, desc, col
-from sqlalchemy.ext.asyncio import AsyncSession
+from astrbot.core.db.po import (
+    Platform as DeprecatedPlatformStat,
+)
+from astrbot.core.db.po import (
+    Stats as DeprecatedStats,
+)
 
 NOT_GIVEN = T.TypeVar("NOT_GIVEN")
 
@@ -57,7 +62,9 @@ class SQLiteDatabase(BaseDatabase):
             async with session.begin():
                 if timestamp is None:
                     timestamp = datetime.now().replace(
-                        minute=0, second=0, microsecond=0
+                        minute=0,
+                        second=0,
+                        microsecond=0,
                     )
                 current_hour = timestamp
                 await session.execute(
@@ -81,13 +88,13 @@ class SQLiteDatabase(BaseDatabase):
             session: AsyncSession
             result = await session.execute(
                 select(func.count(col(PlatformStat.platform_id))).select_from(
-                    PlatformStat
-                )
+                    PlatformStat,
+                ),
             )
             count = result.scalar_one_or_none()
             return count if count is not None else 0
 
-    async def get_platform_stats(self, offset_sec: int = 86400) -> T.List[PlatformStat]:
+    async def get_platform_stats(self, offset_sec: int = 86400) -> list[PlatformStat]:
         """Get platform statistics within the specified offset in seconds and group by platform_id."""
         async with self.get_db() as session:
             session: AsyncSession
@@ -138,7 +145,7 @@ class SQLiteDatabase(BaseDatabase):
                 select(ConversationV2)
                 .order_by(desc(ConversationV2.created_at))
                 .offset(offset)
-                .limit(page_size)
+                .limit(page_size),
             )
             return result.scalars().all()
 
@@ -157,7 +164,7 @@ class SQLiteDatabase(BaseDatabase):
 
             if platform_ids:
                 base_query = base_query.where(
-                    col(ConversationV2.platform_id).in_(platform_ids)
+                    col(ConversationV2.platform_id).in_(platform_ids),
                 )
             if search_query:
                 search_query = search_query.encode("unicode_escape").decode("utf-8")
@@ -167,16 +174,16 @@ class SQLiteDatabase(BaseDatabase):
                         col(ConversationV2.content).ilike(f"%{search_query}%"),
                         col(ConversationV2.user_id).ilike(f"%{search_query}%"),
                         col(ConversationV2.conversation_id).ilike(f"%{search_query}%"),
-                    )
+                    ),
                 )
             if "message_types" in kwargs and len(kwargs["message_types"]) > 0:
                 for msg_type in kwargs["message_types"]:
                     base_query = base_query.where(
-                        col(ConversationV2.user_id).ilike(f"%:{msg_type}:%")
+                        col(ConversationV2.user_id).ilike(f"%:{msg_type}:%"),
                     )
             if "platforms" in kwargs and len(kwargs["platforms"]) > 0:
                 base_query = base_query.where(
-                    col(ConversationV2.platform_id).in_(kwargs["platforms"])
+                    col(ConversationV2.platform_id).in_(kwargs["platforms"]),
                 )
 
             # Get total count matching the filters
@@ -233,7 +240,7 @@ class SQLiteDatabase(BaseDatabase):
             session: AsyncSession
             async with session.begin():
                 query = update(ConversationV2).where(
-                    col(ConversationV2.conversation_id) == cid
+                    col(ConversationV2.conversation_id) == cid,
                 )
                 values = {}
                 if title is not None:
@@ -243,7 +250,7 @@ class SQLiteDatabase(BaseDatabase):
                 if content is not None:
                     values["content"] = content
                 if not values:
-                    return
+                    return None
                 query = query.values(**values)
                 await session.execute(query)
         return await self.get_conversation_by_id(cid)
@@ -254,8 +261,8 @@ class SQLiteDatabase(BaseDatabase):
             async with session.begin():
                 await session.execute(
                     delete(ConversationV2).where(
-                        col(ConversationV2.conversation_id) == cid
-                    )
+                        col(ConversationV2.conversation_id) == cid,
+                    ),
                 )
 
     async def delete_conversations_by_user_id(self, user_id: str) -> None:
@@ -263,7 +270,9 @@ class SQLiteDatabase(BaseDatabase):
             session: AsyncSession
             async with session.begin():
                 await session.execute(
-                    delete(ConversationV2).where(col(ConversationV2.user_id) == user_id)
+                    delete(ConversationV2).where(
+                        col(ConversationV2.user_id) == user_id
+                    ),
                 )
 
     async def get_session_conversations(
@@ -282,7 +291,7 @@ class SQLiteDatabase(BaseDatabase):
                 select(
                     col(Preference.scope_id).label("session_id"),
                     func.json_extract(Preference.value, "$.val").label(
-                        "conversation_id"
+                        "conversation_id",
                     ),  # type: ignore
                     col(ConversationV2.persona_id).label("persona_id"),
                     col(ConversationV2.title).label("title"),
@@ -295,7 +304,8 @@ class SQLiteDatabase(BaseDatabase):
                     == ConversationV2.conversation_id,
                 )
                 .outerjoin(
-                    Persona, col(ConversationV2.persona_id) == Persona.persona_id
+                    Persona,
+                    col(ConversationV2.persona_id) == Persona.persona_id,
                 )
                 .where(Preference.scope == "umo", Preference.key == "sel_conv_id")
             )
@@ -308,14 +318,14 @@ class SQLiteDatabase(BaseDatabase):
                         col(Preference.scope_id).ilike(search_pattern),
                         col(ConversationV2.title).ilike(search_pattern),
                         col(Persona.persona_id).ilike(search_pattern),
-                    )
+                    ),
                 )
 
             # 平台筛选
             if platform:
                 platform_pattern = f"{platform}:%"
                 base_query = base_query.where(
-                    col(Preference.scope_id).like(platform_pattern)
+                    col(Preference.scope_id).like(platform_pattern),
                 )
 
             # 排序
@@ -336,7 +346,8 @@ class SQLiteDatabase(BaseDatabase):
                     == ConversationV2.conversation_id,
                 )
                 .outerjoin(
-                    Persona, col(ConversationV2.persona_id) == Persona.persona_id
+                    Persona,
+                    col(ConversationV2.persona_id) == Persona.persona_id,
                 )
                 .where(Preference.scope == "umo", Preference.key == "sel_conv_id")
             )
@@ -349,13 +360,13 @@ class SQLiteDatabase(BaseDatabase):
                         col(Preference.scope_id).ilike(search_pattern),
                         col(ConversationV2.title).ilike(search_pattern),
                         col(Persona.persona_id).ilike(search_pattern),
-                    )
+                    ),
                 )
 
             if platform:
                 platform_pattern = f"{platform}:%"
                 count_base_query = count_base_query.where(
-                    col(Preference.scope_id).like(platform_pattern)
+                    col(Preference.scope_id).like(platform_pattern),
                 )
 
             total_result = await session.execute(count_base_query)
@@ -396,7 +407,10 @@ class SQLiteDatabase(BaseDatabase):
                 return new_history
 
     async def delete_platform_message_offset(
-        self, platform_id, user_id, offset_sec=86400
+        self,
+        platform_id,
+        user_id,
+        offset_sec=86400,
     ):
         """Delete platform message history records older than the specified offset."""
         async with self.get_db() as session:
@@ -409,11 +423,15 @@ class SQLiteDatabase(BaseDatabase):
                         col(PlatformMessageHistory.platform_id) == platform_id,
                         col(PlatformMessageHistory.user_id) == user_id,
                         col(PlatformMessageHistory.created_at) < cutoff_time,
-                    )
+                    ),
                 )
 
     async def get_platform_message_history(
-        self, platform_id, user_id, page=1, page_size=20
+        self,
+        platform_id,
+        user_id,
+        page=1,
+        page_size=20,
     ):
         """Get platform message history records."""
         async with self.get_db() as session:
@@ -452,7 +470,11 @@ class SQLiteDatabase(BaseDatabase):
             return result.scalar_one_or_none()
 
     async def insert_persona(
-        self, persona_id, system_prompt, begin_dialogs=None, tools=None
+        self,
+        persona_id,
+        system_prompt,
+        begin_dialogs=None,
+        tools=None,
     ):
         """Insert a new persona record."""
         async with self.get_db() as session:
@@ -484,7 +506,11 @@ class SQLiteDatabase(BaseDatabase):
             return result.scalars().all()
 
     async def update_persona(
-        self, persona_id, system_prompt=None, begin_dialogs=None, tools=NOT_GIVEN
+        self,
+        persona_id,
+        system_prompt=None,
+        begin_dialogs=None,
+        tools=NOT_GIVEN,
     ):
         """Update a persona's system prompt or begin dialogs."""
         async with self.get_db() as session:
@@ -499,7 +525,7 @@ class SQLiteDatabase(BaseDatabase):
                 if tools is not NOT_GIVEN:
                     values["tools"] = tools
                 if not values:
-                    return
+                    return None
                 query = query.values(**values)
                 await session.execute(query)
         return await self.get_persona_by_id(persona_id)
@@ -510,7 +536,7 @@ class SQLiteDatabase(BaseDatabase):
             session: AsyncSession
             async with session.begin():
                 await session.execute(
-                    delete(Persona).where(col(Persona.persona_id) == persona_id)
+                    delete(Persona).where(col(Persona.persona_id) == persona_id),
                 )
 
     async def insert_preference_or_update(self, scope, scope_id, key, value):
@@ -529,7 +555,10 @@ class SQLiteDatabase(BaseDatabase):
                     existing_preference.value = value
                 else:
                     new_preference = Preference(
-                        scope=scope, scope_id=scope_id, key=key, value=value
+                        scope=scope,
+                        scope_id=scope_id,
+                        key=key,
+                        value=value,
                     )
                     session.add(new_preference)
                 return existing_preference or new_preference
@@ -568,7 +597,7 @@ class SQLiteDatabase(BaseDatabase):
                         col(Preference.scope) == scope,
                         col(Preference.scope_id) == scope_id,
                         col(Preference.key) == key,
-                    )
+                    ),
                 )
             await session.commit()
 
@@ -581,7 +610,7 @@ class SQLiteDatabase(BaseDatabase):
                     delete(Preference).where(
                         col(Preference.scope) == scope,
                         col(Preference.scope_id) == scope_id,
-                    )
+                    ),
                 )
             await session.commit()
 
@@ -598,7 +627,7 @@ class SQLiteDatabase(BaseDatabase):
                 now = datetime.now()
                 start_time = now - timedelta(seconds=offset_sec)
                 result = await session.execute(
-                    select(PlatformStat).where(PlatformStat.timestamp >= start_time)
+                    select(PlatformStat).where(PlatformStat.timestamp >= start_time),
                 )
                 all_datas = result.scalars().all()
                 deprecated_stats = DeprecatedStats()
@@ -608,7 +637,7 @@ class SQLiteDatabase(BaseDatabase):
                             name=data.platform_id,
                             count=data.count,
                             timestamp=int(data.timestamp.timestamp()),
-                        )
+                        ),
                     )
                 return deprecated_stats
 
@@ -630,7 +659,7 @@ class SQLiteDatabase(BaseDatabase):
             async with self.get_db() as session:
                 session: AsyncSession
                 result = await session.execute(
-                    select(func.sum(PlatformStat.count)).select_from(PlatformStat)
+                    select(func.sum(PlatformStat.count)).select_from(PlatformStat),
                 )
                 total_count = result.scalar_one_or_none()
                 return total_count if total_count is not None else 0
@@ -656,7 +685,7 @@ class SQLiteDatabase(BaseDatabase):
                 result = await session.execute(
                     select(PlatformStat.platform_id, func.sum(PlatformStat.count))
                     .where(PlatformStat.timestamp >= start_time)
-                    .group_by(PlatformStat.platform_id)
+                    .group_by(PlatformStat.platform_id),
                 )
                 grouped_stats = result.all()
                 deprecated_stats = DeprecatedStats()
@@ -666,7 +695,7 @@ class SQLiteDatabase(BaseDatabase):
                             name=platform_id,
                             count=count,
                             timestamp=int(start_time.timestamp()),
-                        )
+                        ),
                     )
                 return deprecated_stats
 

astrbot/core/db/vec_db/base.py

@@ -10,18 +10,16 @@ class Result:
 
 class BaseVecDB:
     async def initialize(self):
-        """
-        初始化向量数据库
-        """
-        pass
+        """初始化向量数据库"""
 
     @abc.abstractmethod
     async def insert(
-        self, content: str, metadata: dict | None = None, id: str | None = None
+        self,
+        content: str,
+        metadata: dict | None = None,
+        id: str | None = None,
     ) -> int:
-        """
-        插入一条文本和其对应向量，自动生成 ID 并保持一致性。
-        """
+        """插入一条文本和其对应向量，自动生成 ID 并保持一致性。"""
         ...
 
     @abc.abstractmethod
@@ -35,11 +33,11 @@ class BaseVecDB:
         max_retries: int = 3,
         progress_callback=None,
     ) -> int:
-        """
-        批量插入文本和其对应向量，自动生成 ID 并保持一致性。
+        """批量插入文本和其对应向量，自动生成 ID 并保持一致性。
 
         Args:
             progress_callback: 进度回调函数，接收参数 (current, total)
+
         """
         ...
 
@@ -52,8 +50,7 @@ class BaseVecDB:
         rerank: bool = False,
         metadata_filters: dict | None = None,
     ) -> list[Result]:
-        """
-        搜索最相似的文档。
+        """搜索最相似的文档。
         Args:
             query (str): 查询文本
             top_k (int): 返回的最相似文档的数量
@@ -64,8 +61,7 @@ class BaseVecDB:
 
     @abc.abstractmethod
     async def delete(self, doc_id: str) -> bool:
-        """
-        删除指定文档。
+        """删除指定文档。
         Args:
             doc_id (str): 要删除的文档 ID
         Returns: