AstrBot 4.7.4__py3-none-any.whl → 4.9.0__py3-none-any.whl

astrbot/core/provider/sources/sensevoice_selfhosted_source.py

@@ -7,6 +7,7 @@ import asyncio
 import os
 import re
 from datetime import datetime
+from typing import cast
 
 from funasr_onnx import SenseVoiceSmall
 from funasr_onnx.utils.postprocess_utils import rich_transcription_postprocess
@@ -32,7 +33,7 @@ class ProviderSenseVoiceSTTSelfHost(STTProvider):
         provider_settings: dict,
     ) -> None:
         super().__init__(provider_config, provider_settings)
-        self.set_model(provider_config.get("stt_model"))
+        self.set_model(provider_config["stt_model"])
         self.model = None
         self.is_emotion = provider_config.get("is_emotion", False)
 
@@ -86,7 +87,9 @@ class ProviderSenseVoiceSTTSelfHost(STTProvider):
             loop = asyncio.get_event_loop()
             res = await loop.run_in_executor(
                 None,  # 使用默认的线程池
-                lambda: self.model(audio_url, language="auto", use_itn=True),
+                lambda: cast(SenseVoiceSmall, self.model)(
+                    audio_url, language="auto", use_itn=True
+                ),
             )
 
             # res = self.model(audio_url, language="auto", use_itn=True)

astrbot/core/provider/sources/vllm_rerank_source.py

@@ -44,6 +44,7 @@ class VLLMRerankProvider(RerankProvider):
         }
         if top_n is not None:
             payload["top_n"] = top_n
+        assert self.client is not None
         async with self.client.post(
             f"{self.base_url}/v1/rerank",
             json=payload,

astrbot/core/provider/sources/whisper_api_source.py

@@ -6,7 +6,10 @@ from openai import NOT_GIVEN, AsyncOpenAI
 from astrbot.core import logger
 from astrbot.core.utils.astrbot_path import get_astrbot_data_path
 from astrbot.core.utils.io import download_file
-from astrbot.core.utils.tencent_record_helper import tencent_silk_to_wav
+from astrbot.core.utils.tencent_record_helper import (
+    convert_to_pcm_wav,
+    tencent_silk_to_wav,
+)
 
 from ..entities import ProviderType
 from ..provider import STTProvider
@@ -33,20 +36,30 @@ class ProviderOpenAIWhisperAPI(STTProvider):
             timeout=provider_config.get("timeout", NOT_GIVEN),
         )
 
-        self.set_model(provider_config.get("model"))
+        self.set_model(provider_config["model"])
 
-    async def _is_silk_file(self, file_path):
+    async def _get_audio_format(self, file_path):
+        # 定义要检测的头部字节
         silk_header = b"SILK"
-        with open(file_path, "rb") as f:
-            file_header = f.read(8)
+        amr_header = b"#!AMR"
+
+        try:
+            with open(file_path, "rb") as f:
+                file_header = f.read(8)
+        except FileNotFoundError:
+            return None
 
         if silk_header in file_header:
-            return True
-        return False
+            return "silk"
+
+        if amr_header in file_header:
+            return "amr"
+        return None
 
     async def get_text(self, audio_url: str) -> str:
         """Only supports mp3, mp4, mpeg, m4a, wav, webm"""
         is_tencent = False
+        output_path = None
 
         if audio_url.startswith("http"):
             if "multimedia.nt.qq.com.cn" in audio_url:
@@ -62,16 +75,35 @@ class ProviderOpenAIWhisperAPI(STTProvider):
             raise FileNotFoundError(f"文件不存在: {audio_url}")
 
         if audio_url.endswith(".amr") or audio_url.endswith(".silk") or is_tencent:
-            is_silk = await self._is_silk_file(audio_url)
-            if is_silk:
-                logger.info("Converting silk file to wav ...")
+            file_format = await self._get_audio_format(audio_url)
+
+            # 判断是否需要转换
+            if file_format in ["silk", "amr"]:
                 temp_dir = os.path.join(get_astrbot_data_path(), "temp")
                 output_path = os.path.join(temp_dir, str(uuid.uuid4()) + ".wav")
-                await tencent_silk_to_wav(audio_url, output_path)
+
+                if file_format == "silk":
+                    logger.info(
+                        "Converting silk file to wav using tencent_silk_to_wav..."
+                    )
+                    await tencent_silk_to_wav(audio_url, output_path)
+                elif file_format == "amr":
+                    logger.info(
+                        "Converting amr file to wav using convert_to_pcm_wav..."
+                    )
+                    await convert_to_pcm_wav(audio_url, output_path)
+
                 audio_url = output_path
 
         result = await self.client.audio.transcriptions.create(
             model=self.model_name,
-            file=open(audio_url, "rb"),
+            file=("audio.wav", open(audio_url, "rb")),
         )
+
+        # remove temp file
+        if output_path and os.path.exists(output_path):
+            try:
+                os.remove(audio_url)
+            except Exception as e:
+                logger.error(f"Failed to remove temp file {audio_url}: {e}")
         return result.text

astrbot/core/provider/sources/whisper_selfhosted_source.py

@@ -1,6 +1,7 @@
 import asyncio
 import os
 import uuid
+from typing import cast
 
 import whisper
 
@@ -26,7 +27,7 @@ class ProviderOpenAIWhisperSelfHost(STTProvider):
         provider_settings: dict,
     ) -> None:
         super().__init__(provider_config, provider_settings)
-        self.set_model(provider_config.get("model"))
+        self.set_model(provider_config["model"])
         self.model = None
 
     async def initialize(self):
@@ -75,5 +76,8 @@ class ProviderOpenAIWhisperSelfHost(STTProvider):
                 await tencent_silk_to_wav(audio_url, output_path)
                 audio_url = output_path
 
+        if not self.model:
+            raise RuntimeError("Whisper 模型未初始化")
+
         result = await loop.run_in_executor(None, self.model.transcribe, audio_url)
-        return result["text"]
+        return cast(str, result["text"])

astrbot/core/provider/sources/xinference_rerank_source.py

@@ -1,6 +1,11 @@
+from typing import cast
+
 from xinference_client.client.restful.async_restful_client import (
     AsyncClient as Client,
 )
+from xinference_client.client.restful.async_restful_client import (
+    AsyncRESTfulRerankModelHandle,
+)
 
 from astrbot import logger
 
@@ -29,7 +34,7 @@ class XinferenceRerankProvider(RerankProvider):
             False,
         )
         self.client = None
-        self.model = None
+        self.model: AsyncRESTfulRerankModelHandle | None = None
         self.model_uid = None
 
     async def initialize(self):
@@ -65,7 +70,10 @@ class XinferenceRerankProvider(RerankProvider):
                     return
 
             if self.model_uid:
-                self.model = await self.client.get_model(self.model_uid)
+                self.model = cast(
+                    AsyncRESTfulRerankModelHandle,
+                    await self.client.get_model(self.model_uid),
+                )
 
         except Exception as e:
             logger.error(f"Failed to initialize Xinference model: {e}")

astrbot/core/star/context.py

@@ -285,7 +285,7 @@ class Context:
         """获取所有用于 Embedding 任务的 Provider。"""
         return self.provider_manager.embedding_provider_insts
 
-    def get_using_provider(self, umo: str | None = None) -> Provider | None:
+    def get_using_provider(self, umo: str | None = None) -> Provider:
         """获取当前使用的用于文本生成任务的 LLM Provider(Chat_Completion 类型)。通过 /provider 指令切换。
 
         Args:
@@ -296,7 +296,7 @@ class Context:
             provider_type=ProviderType.CHAT_COMPLETION,
             umo=umo,
         )
-        if prov and not isinstance(prov, Provider):
+        if not isinstance(prov, Provider):
             raise ValueError("返回的 Provider 不是 Provider 类型")
         return prov
 

astrbot/core/star/register/star_handler.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import re
-from collections.abc import Awaitable, Callable
+from collections.abc import AsyncGenerator, Awaitable, Callable
 from typing import Any
 
 import docstring_parser
@@ -12,6 +12,7 @@ from astrbot.core.agent.handoff import HandoffTool
 from astrbot.core.agent.hooks import BaseAgentRunHooks
 from astrbot.core.agent.tool import FunctionTool
 from astrbot.core.astr_agent_context import AstrAgentContext
+from astrbot.core.message.message_event_result import MessageEventResult
 from astrbot.core.provider.func_tool_manager import PY_TO_JSON_TYPE, SUPPORTED_TYPES
 from astrbot.core.provider.register import llm_tools
 
@@ -28,13 +29,19 @@ from ..filter.regex import RegexFilter
 from ..star_handler import EventType, StarHandlerMetadata, star_handlers_registry
 
 
-def get_handler_full_name(awaitable: Callable[..., Awaitable[Any]]) -> str:
+def get_handler_full_name(
+    awaitable: Callable[..., Awaitable[Any] | AsyncGenerator[Any]],
+) -> str:
     """获取 Handler 的全名"""
     return f"{awaitable.__module__}_{awaitable.__name__}"
 
 
 def get_handler_or_create(
-    handler: Callable[..., Awaitable[Any]],
+    handler: Callable[
+        ...,
+        Awaitable[MessageEventResult | str | None]
+        | AsyncGenerator[MessageEventResult | str | None],
+    ],
     event_type: EventType,
     dont_add=False,
     **kwargs,
@@ -169,6 +176,8 @@ def register_custom_filter(custom_type_filter, *args, **kwargs):
                 for (
                     sub_handle
                 ) in parent_register_commandable.parent_group.sub_command_filters:
+                    if isinstance(sub_handle, CommandGroupFilter):
+                        continue
                     # 所有符合fullname一致的子指令handle添加自定义过滤器。
                     # 不确定是否会有多个子指令有一样的fullname，比如一个方法添加多个command装饰器？
                     sub_handle_md = sub_handle.get_handler_md()
@@ -180,6 +189,8 @@ def register_custom_filter(custom_type_filter, *args, **kwargs):
 
             else:
                 # 裸指令
+                # 确保运行时是可调用的 handler，针对类型检查器添加忽略
+                assert isinstance(awaitable, Callable)
                 handler_md = get_handler_or_create(
                     awaitable,
                     EventType.AdapterMessageEvent,
@@ -237,7 +248,7 @@ class RegisteringCommandable:
 
     group: Callable[..., Callable[..., RegisteringCommandable]] = register_command_group
     command: Callable[..., Callable[..., None]] = register_command
-    custom_filter: Callable[..., Callable[..., None]] = register_custom_filter
+    custom_filter: Callable[..., Callable[..., Any]] = register_custom_filter
 
     def __init__(self, parent_group: CommandGroupFilter):
         self.parent_group = parent_group
@@ -412,7 +423,13 @@ def register_llm_tool(name: str | None = None, **kwargs):
     if kwargs.get("registering_agent"):
         registering_agent = kwargs["registering_agent"]
 
-    def decorator(awaitable: Callable[..., Awaitable[Any]]):
+    def decorator(
+        awaitable: Callable[
+            ...,
+            AsyncGenerator[MessageEventResult | str | None]
+            | Awaitable[MessageEventResult | str | None],
+        ],
+    ):
         llm_tool_name = name_ if name_ else awaitable.__name__
         func_doc = awaitable.__doc__ or ""
         docstring = docstring_parser.parse(func_doc)

astrbot/core/star/star_handler.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
 import enum
-from collections.abc import Awaitable, Callable
+from collections.abc import AsyncGenerator, Awaitable, Callable
 from dataclasses import dataclass, field
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, Literal, TypeVar, overload
 
 from .filter import HandlerFilter
 from .star import star_map
@@ -29,6 +29,84 @@ class StarHandlerRegistry(Generic[T]):
         for handler in self._handlers:
             print(handler.handler_full_name)
 
+    @overload
+    def get_handlers_by_event_type(
+        self,
+        event_type: Literal[EventType.OnAstrBotLoadedEvent],
+        only_activated=True,
+        plugins_name: list[str] | None = None,
+    ) -> list[StarHandlerMetadata[Callable[..., Awaitable[Any]]]]: ...
+
+    @overload
+    def get_handlers_by_event_type(
+        self,
+        event_type: Literal[EventType.OnPlatformLoadedEvent],
+        only_activated=True,
+        plugins_name: list[str] | None = None,
+    ) -> list[StarHandlerMetadata[Callable[..., Awaitable[Any]]]]: ...
+
+    @overload
+    def get_handlers_by_event_type(
+        self,
+        event_type: Literal[EventType.AdapterMessageEvent],
+        only_activated=True,
+        plugins_name: list[str] | None = None,
+    ) -> list[
+        StarHandlerMetadata[Callable[..., Awaitable[Any] | AsyncGenerator[Any]]]
+    ]: ...
+
+    @overload
+    def get_handlers_by_event_type(
+        self,
+        event_type: Literal[EventType.OnLLMRequestEvent],
+        only_activated=True,
+        plugins_name: list[str] | None = None,
+    ) -> list[StarHandlerMetadata[Callable[..., Awaitable[Any]]]]: ...
+
+    @overload
+    def get_handlers_by_event_type(
+        self,
+        event_type: Literal[EventType.OnLLMResponseEvent],
+        only_activated=True,
+        plugins_name: list[str] | None = None,
+    ) -> list[StarHandlerMetadata[Callable[..., Awaitable[Any]]]]: ...
+
+    @overload
+    def get_handlers_by_event_type(
+        self,
+        event_type: Literal[EventType.OnDecoratingResultEvent],
+        only_activated=True,
+        plugins_name: list[str] | None = None,
+    ) -> list[StarHandlerMetadata[Callable[..., Awaitable[Any]]]]: ...
+
+    @overload
+    def get_handlers_by_event_type(
+        self,
+        event_type: Literal[EventType.OnCallingFuncToolEvent],
+        only_activated=True,
+        plugins_name: list[str] | None = None,
+    ) -> list[
+        StarHandlerMetadata[Callable[..., Awaitable[Any] | AsyncGenerator[Any]]]
+    ]: ...
+
+    @overload
+    def get_handlers_by_event_type(
+        self,
+        event_type: Literal[EventType.OnAfterMessageSentEvent],
+        only_activated=True,
+        plugins_name: list[str] | None = None,
+    ) -> list[StarHandlerMetadata[Callable[..., Awaitable[Any]]]]: ...
+
+    @overload
+    def get_handlers_by_event_type(
+        self,
+        event_type: EventType,
+        only_activated=True,
+        plugins_name: list[str] | None = None,
+    ) -> list[
+        StarHandlerMetadata[Callable[..., Awaitable[Any] | AsyncGenerator[Any]]]
+    ]: ...
+
     def get_handlers_by_event_type(
         self,
         event_type: EventType,
@@ -111,8 +189,11 @@ class EventType(enum.Enum):
     OnAfterMessageSentEvent = enum.auto()  # 发送消息后
 
 
+H = TypeVar("H", bound=Callable[..., Any])
+
+
 @dataclass
-class StarHandlerMetadata:
+class StarHandlerMetadata(Generic[H]):
     """描述一个 Star 所注册的某一个 Handler。"""
 
     event_type: EventType
@@ -127,7 +208,7 @@ class StarHandlerMetadata:
     handler_module_path: str
     """Handler 所在的模块路径。"""
 
-    handler: Callable[..., Awaitable[Any]]
+    handler: H
     """Handler 的函数对象，应当是一个异步函数"""
 
     event_filters: list[HandlerFilter]

astrbot/core/updator.py

@@ -71,10 +71,10 @@ class AstrBotUpdator(RepoZipUpdator):
 
     async def check_update(
         self,
-        url: str,
-        current_version: str,
+        url: str | None,
+        current_version: str | None,
         consider_prerelease: bool = True,
-    ) -> ReleaseInfo:
+    ) -> ReleaseInfo | None:
         """检查更新"""
         return await super().check_update(
             self.ASTRBOT_RELEASE_API,

astrbot/core/utils/io.py

@@ -49,7 +49,7 @@ def port_checker(port: int, host: str = "localhost"):
         return False
 
 
-def save_temp_img(img: Image.Image | str) -> str:
+def save_temp_img(img: Image.Image | bytes) -> str:
     temp_dir = os.path.join(get_astrbot_data_path(), "temp")
     # 获得文件创建时间，清除超过 12 小时的
     try:

astrbot/core/utils/session_waiter.py

@@ -20,16 +20,16 @@ class SessionController:
 
     def __init__(self):
         self.future = asyncio.Future()
-        self.current_event: asyncio.Event = None
+        self.current_event: asyncio.Event | None = None
         """当前正在等待的所用的异步事件"""
-        self.ts: float = None
+        self.ts: float | None = None
         """上次保持(keep)开始时的时间"""
-        self.timeout: float | int = None
+        self.timeout: float | int | None = None
         """上次保持(keep)开始时的超时时间"""
 
         self.history_chains: list[list[Comp.BaseMessageComponent]] = []
 
-    def stop(self, error: Exception = None):
+    def stop(self, error: Exception | None = None):
         """立即结束这个会话"""
         if not self.future.done():
             if error:
@@ -53,6 +53,8 @@ class SessionController:
                 self.stop()
                 return
         else:
+            assert self.timeout is not None
+            assert self.ts is not None
             left_timeout = self.timeout - (new_ts - self.ts)
             timeout = left_timeout + timeout
             if timeout <= 0:
@@ -69,7 +71,7 @@ class SessionController:
 
         asyncio.create_task(self._holding(new_event, timeout))  # 开始新的 keep
 
-    async def _holding(self, event: asyncio.Event, timeout: int):
+    async def _holding(self, event: asyncio.Event, timeout: float):
         """等待事件结束或超时"""
         try:
             await asyncio.wait_for(event.wait(), timeout)
@@ -108,7 +110,9 @@ class SessionWaiter:
     ):
         self.session_id = session_id
         self.session_filter = session_filter
-        self.handler: Callable[[str], Awaitable[Any]] | None = None  # 处理函数
+        self.handler: (
+            Callable[[SessionController, AstrMessageEvent], Awaitable[Any]] | None
+        ) = None  # 处理函数
 
         self.session_controller = SessionController()
         self.record_history_chains = record_history_chains
@@ -119,7 +123,7 @@ class SessionWaiter:
 
     async def register_wait(
         self,
-        handler: Callable[[str], Awaitable[Any]],
+        handler: Callable[[SessionController, AstrMessageEvent], Awaitable[Any]],
         timeout: int = 30,
     ) -> Any:
         """等待外部输入并处理"""
@@ -137,7 +141,7 @@ class SessionWaiter:
         finally:
             self._cleanup()
 
-    def _cleanup(self, error: Exception = None):
+    def _cleanup(self, error: Exception | None = None):
         """清理会话"""
         USER_SESSIONS.pop(self.session_id, None)
         try:
@@ -161,6 +165,7 @@ class SessionWaiter:
                     )
                 try:
                     # TODO: 这里使用 create_task，跟踪 task，防止超时后这里 handler 仍然在执行
+                    assert session.handler is not None
                     await session.handler(session.session_controller, event)
                 except Exception as e:
                     session.session_controller.stop(e)
@@ -173,11 +178,13 @@ def session_waiter(timeout: int = 30, record_history_chains: bool = False):
     :param record_history_chain: 是否自动记录历史消息链。可以通过 controller.get_history_chains() 获取。深拷贝。
     """
 
-    def decorator(func: Callable[[str], Awaitable[Any]]):
+    def decorator(
+        func: Callable[[SessionController, AstrMessageEvent], Awaitable[Any]],
+    ):
         @functools.wraps(func)
         async def wrapper(
             event: AstrMessageEvent,
-            session_filter: SessionFilter = None,
+            session_filter: SessionFilter | None = None,
             *args,
             **kwargs,
         ):

astrbot/core/utils/shared_preferences.py

@@ -53,6 +53,38 @@ class SharedPreferences:
         ret = await self.db_helper.get_preferences(scope, scope_id, key)
         return ret
 
+    @overload
+    async def session_get(
+        self,
+        umo: str,
+        key: str,
+        default: _VT = None,
+    ) -> _VT: ...
+
+    @overload
+    async def session_get(
+        self,
+        umo: None,
+        key: str,
+        default: Any = None,
+    ) -> list[Preference]: ...
+
+    @overload
+    async def session_get(
+        self,
+        umo: str,
+        key: None,
+        default: Any = None,
+    ) -> list[Preference]: ...
+
+    @overload
+    async def session_get(
+        self,
+        umo: None,
+        key: None,
+        default: Any = None,
+    ) -> list[Preference]: ...
+
     async def session_get(
         self,
         umo: str | None,

astrbot/core/utils/t2i/__init__.py

@@ -3,11 +3,11 @@ from abc import ABC, abstractmethod
 
 class RenderStrategy(ABC):
     @abstractmethod
-    def render(self, text: str, return_url: bool) -> str:
+    async def render(self, text: str, return_url: bool) -> str:
         pass
 
     @abstractmethod
-    def render_custom_template(
+    async def render_custom_template(
         self,
         tmpl_str: str,
         tmpl_data: dict,