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

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: