AstrBot 4.8.0__py3-none-any.whl → 4.9.1__py3-none-any.whl

astrbot/core/provider/sources/gemini_source.py

@@ -4,6 +4,7 @@ import json
 import logging
 import random
 from collections.abc import AsyncGenerator
+from typing import cast
 
 from google import genai
 from google.genai import types
@@ -126,17 +127,17 @@ class ProviderGoogleGenAI(Provider):
     ) -> types.GenerateContentConfig:
         """准备查询配置"""
         if not modalities:
-            modalities = ["Text"]
+            modalities = ["TEXT"]
 
         # 流式输出不支持图片模态
         if (
             self.provider_settings.get("streaming_response", False)
-            and "Image" in modalities
+            and "IMAGE" in modalities
         ):
             logger.warning("流式输出不支持图片模态，已自动降级为文本模态")
-            modalities = ["Text"]
+            modalities = ["TEXT"]
 
-        tool_list = []
+        tool_list: list[types.Tool] | None = []
         model_name = self.get_model()
         native_coderunner = self.provider_config.get("gm_native_coderunner", False)
         native_search = self.provider_config.get("gm_native_search", False)
@@ -213,7 +214,7 @@ class ProviderGoogleGenAI(Provider):
             logprobs=payloads.get("logprobs"),
             seed=payloads.get("seed"),
             response_modalities=modalities,
-            tools=tool_list,
+            tools=cast(types.ToolListUnion | None, tool_list),
             safety_settings=self.safety_settings if self.safety_settings else None,
             thinking_config=(
                 types.ThinkingConfig(
@@ -257,6 +258,7 @@ class ProviderGoogleGenAI(Provider):
             content_cls: type[types.Content],
         ) -> None:
             if contents and isinstance(contents[-1], content_cls):
+                assert contents[-1].parts is not None
                 contents[-1].parts.extend(part)
             else:
                 contents.append(content_cls(parts=part))
@@ -429,9 +431,9 @@ class ProviderGoogleGenAI(Provider):
             None,
         )
 
-        modalities = ["Text"]
+        modalities = ["TEXT"]
         if self.provider_config.get("gm_resp_image_modal", False):
-            modalities.append("Image")
+            modalities.append("IMAGE")
 
         conversation = self._prepare_conversation(payloads)
         temperature = payloads.get("temperature", 0.7)
@@ -448,7 +450,7 @@ class ProviderGoogleGenAI(Provider):
                 )
                 result = await self.client.models.generate_content(
                     model=self.get_model(),
-                    contents=conversation,
+                    contents=cast(types.ContentListUnion, conversation),
                     config=config,
                 )
                 logger.debug(f"genai result: {result}")
@@ -488,7 +490,7 @@ class ProviderGoogleGenAI(Provider):
                     logger.warning(
                         f"{self.get_model()} 不支持多模态输出，降级为文本模态",
                     )
-                    modalities = ["Text"]
+                    modalities = ["TEXT"]
                 else:
                     raise
                 continue
@@ -524,7 +526,7 @@ class ProviderGoogleGenAI(Provider):
                 )
                 result = await self.client.models.generate_content_stream(
                     model=self.get_model(),
-                    contents=conversation,
+                    contents=cast(types.ContentListUnion, conversation),
                     config=config,
                 )
                 break

astrbot/core/provider/sources/minimax_tts_api_source.py

@@ -87,7 +87,7 @@ class ProviderMiniMaxTTSAPI(TTSProvider):
 
         return json.dumps(dict_body)
 
-    async def _call_tts_stream(self, text: str) -> AsyncIterator[bytes]:
+    async def _call_tts_stream(self, text: str) -> AsyncIterator[str]:
         """进行流式请求"""
         try:
             async with (
@@ -117,7 +117,9 @@ class ProviderMiniMaxTTSAPI(TTSProvider):
                                     data = json.loads(message[6:])
                                     if "extra_info" in data:
                                         continue
-                                    audio = data.get("data", {}).get("audio")
+                                    audio: str | None = data.get("data", {}).get(
+                                        "audio"
+                                    )
                                     if audio is not None:
                                         yield audio
                                 except json.JSONDecodeError:

astrbot/core/provider/sources/openai_embedding_source.py

@@ -30,9 +30,9 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
         embedding = await self.client.embeddings.create(input=text, model=self.model)
         return embedding.data[0].embedding
 
-    async def get_embeddings(self, texts: list[str]) -> list[list[float]]:
+    async def get_embeddings(self, text: list[str]) -> list[list[float]]:
         """批量获取文本的嵌入"""
-        embeddings = await self.client.embeddings.create(input=texts, model=self.model)
+        embeddings = await self.client.embeddings.create(input=text, model=self.model)
         return [item.embedding for item in embeddings.data]
 
     def get_dim(self) -> int:

astrbot/core/provider/sources/openai_source.py

@@ -284,6 +284,10 @@ class ProviderOpenAIOfficial(Provider):
                 if isinstance(tool_call, str):
                     # workaround for #1359
                     tool_call = json.loads(tool_call)
+                if tools is None:
+                    # 工具集未提供
+                    # Should be unreachable
+                    raise Exception("工具集未提供")
                 for tool in tools.func_list:
                     if (
                         tool_call.type == "function"

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

@@ -36,7 +36,7 @@ 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 _get_audio_format(self, file_path):
         # 定义要检测的头部字节

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,