mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-16 17:47:06 +08:00
chore: initialize
This commit is contained in:
1
src/astrbot_sdk/api/basic/astrbot_config.py
Normal file
1
src/astrbot_sdk/api/basic/astrbot_config.py
Normal file
@@ -0,0 +1 @@
|
||||
class AstrBotConfig(dict): ...
|
||||
221
src/astrbot_sdk/api/basic/conversation_mgr.py
Normal file
221
src/astrbot_sdk/api/basic/conversation_mgr.py
Normal file
@@ -0,0 +1,221 @@
|
||||
# from astr_agent_sdk.message import AssistantMessageSegment, UserMessageSegment
|
||||
|
||||
|
||||
class BaseConversationManager:
|
||||
"""负责管理会话与 LLM 的对话,某个会话当前正在用哪个对话。"""
|
||||
|
||||
async def _trigger_session_deleted(self, unified_msg_origin: str) -> None:
|
||||
"""触发会话删除回调.
|
||||
|
||||
Args:
|
||||
unified_msg_origin: 会话ID
|
||||
|
||||
"""
|
||||
...
|
||||
|
||||
async def new_conversation(
|
||||
self,
|
||||
unified_msg_origin: str,
|
||||
platform_id: str | None = None,
|
||||
content: list[dict] | None = None,
|
||||
title: str | None = None,
|
||||
persona_id: str | None = None,
|
||||
) -> str:
|
||||
"""新建对话,并将当前会话的对话转移到新对话.
|
||||
|
||||
Args:
|
||||
unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id
|
||||
Returns:
|
||||
conversation_id (str): 对话 ID, 是 uuid 格式的字符串
|
||||
|
||||
"""
|
||||
...
|
||||
|
||||
async def switch_conversation(self, unified_msg_origin: str, conversation_id: str):
|
||||
"""切换会话的对话
|
||||
|
||||
Args:
|
||||
unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id
|
||||
conversation_id (str): 对话 ID, 是 uuid 格式的字符串
|
||||
|
||||
"""
|
||||
...
|
||||
|
||||
async def delete_conversation(
|
||||
self,
|
||||
unified_msg_origin: str,
|
||||
conversation_id: str | None = None,
|
||||
):
|
||||
"""删除会话的对话,当 conversation_id 为 None 时删除会话当前的对话
|
||||
|
||||
Args:
|
||||
unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id
|
||||
conversation_id (str): 对话 ID, 是 uuid 格式的字符串
|
||||
|
||||
"""
|
||||
...
|
||||
|
||||
async def delete_conversations_by_user_id(self, unified_msg_origin: str):
|
||||
"""删除会话的所有对话
|
||||
|
||||
Args:
|
||||
unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id
|
||||
|
||||
"""
|
||||
...
|
||||
|
||||
async def get_curr_conversation_id(self, unified_msg_origin: str) -> str | None:
|
||||
"""获取会话当前的对话 ID
|
||||
|
||||
Args:
|
||||
unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id
|
||||
Returns:
|
||||
conversation_id (str): 对话 ID, 是 uuid 格式的字符串
|
||||
|
||||
"""
|
||||
...
|
||||
|
||||
# async def get_conversation(
|
||||
# self,
|
||||
# unified_msg_origin: str,
|
||||
# conversation_id: str,
|
||||
# create_if_not_exists: bool = False,
|
||||
# ) -> Conversation | None:
|
||||
# """获取会话的对话.
|
||||
|
||||
# Args:
|
||||
# unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id
|
||||
# conversation_id (str): 对话 ID, 是 uuid 格式的字符串
|
||||
# create_if_not_exists (bool): 如果对话不存在,是否创建一个新的对话
|
||||
# Returns:
|
||||
# conversation (Conversation): 对话对象
|
||||
|
||||
# """
|
||||
# ...
|
||||
|
||||
# async def get_conversations(
|
||||
# self,
|
||||
# unified_msg_origin: str | None = None,
|
||||
# platform_id: str | None = None,
|
||||
# ) -> list[Conversation]:
|
||||
# """获取对话列表.
|
||||
|
||||
# Args:
|
||||
# unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id,可选
|
||||
# platform_id (str): 平台 ID, 可选参数, 用于过滤对话
|
||||
# Returns:
|
||||
# conversations (List[Conversation]): 对话对象列表
|
||||
|
||||
# """
|
||||
# ...
|
||||
|
||||
# async def get_filtered_conversations(
|
||||
# self,
|
||||
# page: int = 1,
|
||||
# page_size: int = 20,
|
||||
# platform_ids: list[str] | None = None,
|
||||
# search_query: str = "",
|
||||
# **kwargs,
|
||||
# ) -> tuple[list[Conversation], int]:
|
||||
# """获取过滤后的对话列表.
|
||||
|
||||
# Args:
|
||||
# page (int): 页码, 默认为 1
|
||||
# page_size (int): 每页大小, 默认为 20
|
||||
# platform_ids (list[str]): 平台 ID 列表, 可选
|
||||
# search_query (str): 搜索查询字符串, 可选
|
||||
# Returns:
|
||||
# conversations (list[Conversation]): 对话对象列表
|
||||
|
||||
# """
|
||||
# ...
|
||||
|
||||
async def update_conversation(
|
||||
self,
|
||||
unified_msg_origin: str,
|
||||
conversation_id: str | None = None,
|
||||
history: list[dict] | None = None,
|
||||
title: str | None = None,
|
||||
persona_id: str | None = None,
|
||||
) -> None:
|
||||
"""更新会话的对话.
|
||||
|
||||
Args:
|
||||
unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id
|
||||
conversation_id (str): 对话 ID, 是 uuid 格式的字符串
|
||||
history (List[Dict]): 对话历史记录, 是一个字典列表, 每个字典包含 role 和 content 字段
|
||||
|
||||
"""
|
||||
...
|
||||
|
||||
async def update_conversation_title(
|
||||
self,
|
||||
unified_msg_origin: str,
|
||||
title: str,
|
||||
conversation_id: str | None = None,
|
||||
) -> None:
|
||||
"""更新会话的对话标题.
|
||||
|
||||
Args:
|
||||
unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id
|
||||
title (str): 对话标题
|
||||
conversation_id (str): 对话 ID, 是 uuid 格式的字符串
|
||||
Deprecated:
|
||||
Use `update_conversation` with `title` parameter instead.
|
||||
|
||||
"""
|
||||
...
|
||||
|
||||
async def update_conversation_persona_id(
|
||||
self,
|
||||
unified_msg_origin: str,
|
||||
persona_id: str,
|
||||
conversation_id: str | None = None,
|
||||
) -> None:
|
||||
"""更新会话的对话 Persona ID.
|
||||
|
||||
Args:
|
||||
unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id
|
||||
persona_id (str): 对话 Persona ID
|
||||
conversation_id (str): 对话 ID, 是 uuid 格式的字符串
|
||||
Deprecated:
|
||||
Use `update_conversation` with `persona_id` parameter instead.
|
||||
|
||||
"""
|
||||
...
|
||||
|
||||
# async def add_message_pair(
|
||||
# self,
|
||||
# cid: str,
|
||||
# user_message: UserMessageSegment | dict,
|
||||
# assistant_message: AssistantMessageSegment | dict,
|
||||
# ) -> None:
|
||||
# """Add a user-assistant message pair to the conversation history.
|
||||
|
||||
# Args:
|
||||
# cid (str): Conversation ID
|
||||
# user_message (UserMessageSegment | dict): OpenAI-format user message object or dict
|
||||
# assistant_message (AssistantMessageSegment | dict): OpenAI-format assistant message object or dict
|
||||
|
||||
# Raises:
|
||||
# Exception: If the conversation with the given ID is not found
|
||||
# """
|
||||
# ...
|
||||
|
||||
async def get_human_readable_context(
|
||||
self,
|
||||
unified_msg_origin: str,
|
||||
conversation_id: str,
|
||||
page: int = 1,
|
||||
page_size: int = 10,
|
||||
) -> tuple[list[str], int]:
|
||||
"""获取人类可读的上下文.
|
||||
|
||||
Args:
|
||||
unified_msg_origin (str): 统一的消息来源字符串。格式为 platform_name:message_type:session_id
|
||||
conversation_id (str): 对话 ID, 是 uuid 格式的字符串
|
||||
page (int): 页码
|
||||
page_size (int): 每页大小
|
||||
|
||||
"""
|
||||
...
|
||||
2
src/astrbot_sdk/api/components/command.py
Normal file
2
src/astrbot_sdk/api/components/command.py
Normal file
@@ -0,0 +1,2 @@
|
||||
class CommandComponent:
|
||||
pass
|
||||
5
src/astrbot_sdk/api/event/__init__.py
Normal file
5
src/astrbot_sdk/api/event/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
||||
from .astr_message_event import AstrMessageEvent
|
||||
|
||||
__all__ = [
|
||||
"AstrMessageEvent",
|
||||
]
|
||||
370
src/astrbot_sdk/api/event/astr_message_event.py
Normal file
370
src/astrbot_sdk/api/event/astr_message_event.py
Normal file
@@ -0,0 +1,370 @@
|
||||
from __future__ import annotations
|
||||
import typing as T
|
||||
from .astrbot_message import AstrBotMessage, Group
|
||||
from ...api.platform.platform_metadata import PlatformMetadata
|
||||
from ...api.event.message_type import MessageType
|
||||
from ...api.event.message_session import MessageSession
|
||||
from ...api.event.event_result import MessageEventResult
|
||||
from ...api.message.chain import MessageChain
|
||||
from ...api.message.components import BaseMessageComponent
|
||||
from dataclasses import dataclass, field
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class AstrMessageEventModel(BaseModel):
|
||||
message_str: str
|
||||
message_obj: AstrBotMessage
|
||||
platform_meta: PlatformMetadata
|
||||
session_id: str
|
||||
role: T.Literal["admin", "member"] = "member"
|
||||
is_wake: bool = False
|
||||
is_at_or_wake_command: bool = False
|
||||
extras: dict = Field(default_factory=dict)
|
||||
result: MessageEventResult | None = None
|
||||
has_send_oper: bool = False
|
||||
call_llm: bool = False
|
||||
plugins_name: list[str] = Field(default_factory=list)
|
||||
|
||||
@classmethod
|
||||
def from_event(cls, event: AstrMessageEvent) -> AstrMessageEventModel:
|
||||
return cls(
|
||||
message_str=event.message_str,
|
||||
message_obj=event.message_obj,
|
||||
platform_meta=event.platform_meta,
|
||||
session_id=event.session_id,
|
||||
role=event.role,
|
||||
is_wake=event.is_wake,
|
||||
is_at_or_wake_command=event.is_at_or_wake_command,
|
||||
extras=event._extras,
|
||||
result=event._result,
|
||||
has_send_oper=event.has_send_oper,
|
||||
call_llm=event.call_llm,
|
||||
plugins_name=event._plugins_name,
|
||||
)
|
||||
|
||||
def to_event(self) -> AstrMessageEvent:
|
||||
event = AstrMessageEvent(
|
||||
message_str=self.message_str,
|
||||
message_obj=self.message_obj,
|
||||
platform_meta=self.platform_meta,
|
||||
session_id=self.session_id,
|
||||
role=self.role,
|
||||
is_wake=self.is_wake,
|
||||
is_at_or_wake_command=self.is_at_or_wake_command,
|
||||
_extras=self.extras,
|
||||
_result=self.result,
|
||||
has_send_oper=self.has_send_oper,
|
||||
call_llm=self.call_llm,
|
||||
_plugins_name=self.plugins_name,
|
||||
)
|
||||
return event
|
||||
|
||||
|
||||
@dataclass
|
||||
class AstrMessageEvent:
|
||||
message_str: str
|
||||
"""消息的纯文本内容"""
|
||||
|
||||
message_obj: AstrBotMessage
|
||||
"""消息对象"""
|
||||
|
||||
platform_meta: PlatformMetadata
|
||||
"""平台适配器的元信息"""
|
||||
|
||||
session_id: str
|
||||
"""会话 ID"""
|
||||
|
||||
role: T.Literal["admin", "member"] = "member"
|
||||
"""消息发送者的角色,如 "admin", "member" 等"""
|
||||
|
||||
is_wake: bool = False
|
||||
"""是否唤醒(是否通过 WakingStage)"""
|
||||
|
||||
is_at_or_wake_command: bool = False
|
||||
"""是否艾特机器人或通过唤醒命令触发的消息"""
|
||||
|
||||
_extras: dict = field(default_factory=dict)
|
||||
"""存储额外的信息"""
|
||||
|
||||
_result: MessageEventResult | None = None
|
||||
"""消息事件的结果"""
|
||||
|
||||
has_send_oper: bool = False
|
||||
"""是否已经发送过操作"""
|
||||
|
||||
call_llm: bool = False
|
||||
"""是否调用 LLM"""
|
||||
|
||||
_plugins_name: list[str] = field(default_factory=list)
|
||||
"""处理该事件的插件名称列表"""
|
||||
|
||||
def __post_init__(self):
|
||||
self.session = MessageSession(
|
||||
platform_name=self.platform_meta.id,
|
||||
message_type=self.message_obj.type,
|
||||
session_id=self.session_id,
|
||||
)
|
||||
self.unified_msg_origin = str(self.session)
|
||||
self.platform = self.platform_meta # back compatibility
|
||||
|
||||
def get_platform_name(self) -> str:
|
||||
"""
|
||||
获取这个事件所属的平台的类型(如 aiocqhttp, slack, discord 等)。
|
||||
NOTE: 用户可能会同时运行多个相同类型的平台适配器。
|
||||
"""
|
||||
return self.platform_meta.name
|
||||
|
||||
def get_platform_id(self):
|
||||
"""
|
||||
获取这个事件所属的平台的 ID。
|
||||
NOTE: 用户可能会同时运行多个相同类型的平台适配器,但能确定的是 ID 是唯一的。
|
||||
"""
|
||||
return self.platform_meta.id
|
||||
|
||||
def get_message_str(self) -> str:
|
||||
"""获取消息字符串。"""
|
||||
return self.message_str
|
||||
|
||||
def get_messages(self) -> list[BaseMessageComponent]:
|
||||
"""获取消息链。"""
|
||||
return self.message_obj.message
|
||||
|
||||
def get_message_type(self) -> MessageType:
|
||||
"""获取消息类型。"""
|
||||
return self.message_obj.type
|
||||
|
||||
def get_session_id(self) -> str:
|
||||
"""获取会话id。"""
|
||||
return self.session_id
|
||||
|
||||
def get_group_id(self) -> str:
|
||||
"""获取群组id。如果不是群组消息,返回空字符串。"""
|
||||
return self.message_obj.group_id
|
||||
|
||||
def get_self_id(self) -> str:
|
||||
"""获取机器人自身的id。"""
|
||||
return self.message_obj.self_id
|
||||
|
||||
def get_sender_id(self) -> str:
|
||||
"""获取消息发送者的id。"""
|
||||
return self.message_obj.sender.user_id
|
||||
|
||||
def get_sender_name(self) -> str | None:
|
||||
"""获取消息发送者的名称。(可能会返回空字符串)"""
|
||||
return self.message_obj.sender.nickname
|
||||
|
||||
def set_extra(self, key, value):
|
||||
"""设置额外的信息。"""
|
||||
self._extras[key] = value
|
||||
|
||||
def get_extra(self, key: str | None = None, default=None) -> T.Any:
|
||||
"""获取额外的信息。"""
|
||||
if key is None:
|
||||
return self._extras
|
||||
return self._extras.get(key, default)
|
||||
|
||||
def clear_extra(self):
|
||||
"""清除额外的信息。"""
|
||||
self._extras.clear()
|
||||
|
||||
def is_private_chat(self) -> bool:
|
||||
"""是否是私聊。"""
|
||||
return self.message_obj.type.value == (MessageType.FRIEND_MESSAGE).value
|
||||
|
||||
def is_wake_up(self) -> bool:
|
||||
"""是否是唤醒机器人的事件。"""
|
||||
return self.is_wake
|
||||
|
||||
def is_admin(self) -> bool:
|
||||
"""是否是管理员。"""
|
||||
return self.role == "admin"
|
||||
|
||||
# async def send_streaming(
|
||||
# self,
|
||||
# generator: AsyncGenerator[MessageChain, None],
|
||||
# use_fallback: bool = False,
|
||||
# ):
|
||||
# """发送流式消息到消息平台,使用异步生成器。
|
||||
# 目前仅支持: telegram,qq official 私聊。
|
||||
# Fallback仅支持 aiocqhttp。
|
||||
# """
|
||||
# asyncio.create_task(
|
||||
# Metric.upload(msg_event_tick=1, adapter_name=self.platform_meta.name),
|
||||
# )
|
||||
# self._has_send_oper = True
|
||||
|
||||
def set_result(self, result: MessageEventResult | str):
|
||||
"""设置消息事件的结果。
|
||||
|
||||
Note:
|
||||
事件处理器可以通过设置结果来控制事件是否继续传播,并向消息适配器发送消息。
|
||||
|
||||
如果没有设置 `MessageEventResult` 中的 result_type,默认为 CONTINUE。即事件将会继续向后面的 listener 或者 command 传播。
|
||||
|
||||
Example:
|
||||
```
|
||||
async def ban_handler(self, event: AstrMessageEvent):
|
||||
if event.get_sender_id() in self.blacklist:
|
||||
event.set_result(MessageEventResult().set_console_log("由于用户在黑名单,因此消息事件中断处理。")).set_result_type(EventResultType.STOP)
|
||||
return
|
||||
|
||||
async def check_count(self, event: AstrMessageEvent):
|
||||
self.count += 1
|
||||
event.set_result(MessageEventResult().set_console_log("数量已增加", logging.DEBUG).set_result_type(EventResultType.CONTINUE))
|
||||
return
|
||||
```
|
||||
|
||||
"""
|
||||
if isinstance(result, str):
|
||||
result = MessageEventResult().message(result)
|
||||
# 兼容外部插件或调用方传入的 chain=None 的情况,确保为可迭代列表
|
||||
if isinstance(result, MessageEventResult) and result.chain is None:
|
||||
result.chain = []
|
||||
self._result = result
|
||||
|
||||
def stop_event(self):
|
||||
"""终止事件传播。"""
|
||||
if self._result is None:
|
||||
self.set_result(MessageEventResult().stop_event())
|
||||
else:
|
||||
self._result.stop_event()
|
||||
|
||||
def continue_event(self):
|
||||
"""继续事件传播。"""
|
||||
if self._result is None:
|
||||
self.set_result(MessageEventResult().continue_event())
|
||||
else:
|
||||
self._result.continue_event()
|
||||
|
||||
def is_stopped(self) -> bool:
|
||||
"""是否终止事件传播。"""
|
||||
if self._result is None:
|
||||
return False # 默认是继续传播
|
||||
return self._result.is_stopped()
|
||||
|
||||
def should_call_llm(self, call_llm: bool):
|
||||
"""是否在此消息事件中禁止默认的 LLM 请求。
|
||||
|
||||
只会阻止 AstrBot 默认的 LLM 请求链路,不会阻止插件中的 LLM 请求。
|
||||
"""
|
||||
self.call_llm = call_llm
|
||||
|
||||
def get_result(self) -> MessageEventResult | None:
|
||||
"""获取消息事件的结果。"""
|
||||
return self._result
|
||||
|
||||
def clear_result(self):
|
||||
"""清除消息事件的结果。"""
|
||||
self._result = None
|
||||
|
||||
"""消息链相关"""
|
||||
|
||||
def make_result(self) -> MessageEventResult:
|
||||
"""创建一个空的消息事件结果。
|
||||
|
||||
Example:
|
||||
```python
|
||||
# 纯文本回复
|
||||
yield event.make_result().message("Hi")
|
||||
# 发送图片
|
||||
yield event.make_result().url_image("https://example.com/image.jpg")
|
||||
yield event.make_result().file_image("image.jpg")
|
||||
```
|
||||
|
||||
"""
|
||||
return MessageEventResult()
|
||||
|
||||
def plain_result(self, text: str) -> MessageEventResult:
|
||||
"""创建一个空的消息事件结果,只包含一条文本消息。"""
|
||||
return MessageEventResult().message(text)
|
||||
|
||||
def image_result(self, url_or_path: str) -> MessageEventResult:
|
||||
"""创建一个空的消息事件结果,只包含一条图片消息。
|
||||
|
||||
根据开头是否包含 http 来判断是网络图片还是本地图片。
|
||||
"""
|
||||
if url_or_path.startswith("http"):
|
||||
return MessageEventResult().url_image(url_or_path)
|
||||
return MessageEventResult().file_image(url_or_path)
|
||||
|
||||
def chain_result(self, chain: list[BaseMessageComponent]) -> MessageEventResult:
|
||||
"""创建一个空的消息事件结果,包含指定的消息链。"""
|
||||
mer = MessageEventResult()
|
||||
mer.chain = chain
|
||||
return mer
|
||||
|
||||
# """LLM 请求相关"""
|
||||
|
||||
# def request_llm(
|
||||
# self,
|
||||
# prompt: str,
|
||||
# func_tool_manager=None,
|
||||
# session_id: str | None = None,
|
||||
# image_urls: list[str] | None = None,
|
||||
# contexts: list | None = None,
|
||||
# system_prompt: str = "",
|
||||
# conversation: Conversation | None = None,
|
||||
# ) -> ProviderRequest:
|
||||
# """创建一个 LLM 请求。
|
||||
|
||||
# Examples:
|
||||
# ```py
|
||||
# yield event.request_llm(prompt="hi")
|
||||
# ```
|
||||
# prompt: 提示词
|
||||
|
||||
# system_prompt: 系统提示词
|
||||
|
||||
# session_id: 已经过时,留空即可
|
||||
|
||||
# image_urls: 可以是 base64:// 或者 http:// 开头的图片链接,也可以是本地图片路径。
|
||||
|
||||
# contexts: 当指定 contexts 时,将会使用 contexts 作为上下文。如果同时传入了 conversation,将会忽略 conversation。
|
||||
|
||||
# func_tool_manager: 函数工具管理器,用于调用函数工具。用 self.context.get_llm_tool_manager() 获取。
|
||||
|
||||
# conversation: 可选。如果指定,将在指定的对话中进行 LLM 请求。对话的人格会被用于 LLM 请求,并且结果将会被记录到对话中。
|
||||
|
||||
# """
|
||||
# if image_urls is None:
|
||||
# image_urls = []
|
||||
# if contexts is None:
|
||||
# contexts = []
|
||||
# if len(contexts) > 0 and conversation:
|
||||
# conversation = None
|
||||
|
||||
# return ProviderRequest(
|
||||
# prompt=prompt,
|
||||
# session_id=session_id,
|
||||
# image_urls=image_urls,
|
||||
# func_tool=func_tool_manager,
|
||||
# contexts=contexts,
|
||||
# system_prompt=system_prompt,
|
||||
# conversation=conversation,
|
||||
# )
|
||||
|
||||
async def send(self, message: MessageChain):
|
||||
"""发送消息到消息平台。
|
||||
|
||||
Args:
|
||||
message (MessageChain): 消息链,具体使用方式请参考文档。
|
||||
|
||||
"""
|
||||
...
|
||||
|
||||
async def react(self, emoji: str):
|
||||
"""对消息添加表情回应。
|
||||
|
||||
默认实现为发送一条包含该表情的消息。
|
||||
注意:此实现并不一定符合所有平台的原生“表情回应”行为。
|
||||
如需支持平台原生的消息反应功能,请在对应平台的子类中重写本方法。
|
||||
"""
|
||||
...
|
||||
|
||||
async def get_group(self, group_id: str | None = None, **kwargs) -> Group | None:
|
||||
"""获取一个群聊的数据, 如果不填写 group_id: 如果是私聊消息,返回 None。如果是群聊消息,返回当前群聊的数据。
|
||||
|
||||
适配情况:
|
||||
|
||||
- aiocqhttp(OneBotv11)
|
||||
"""
|
||||
98
src/astrbot_sdk/api/event/astrbot_message.py
Normal file
98
src/astrbot_sdk/api/event/astrbot_message.py
Normal file
@@ -0,0 +1,98 @@
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .message_type import MessageType
|
||||
from ..message.components import BaseMessageComponent
|
||||
|
||||
|
||||
@dataclass
|
||||
class MessageMember:
|
||||
user_id: str
|
||||
nickname: str | None = None
|
||||
|
||||
def __str__(self):
|
||||
return (
|
||||
f"User ID: {self.user_id},"
|
||||
f"Nickname: {self.nickname if self.nickname else 'N/A'}"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Group:
|
||||
group_id: str
|
||||
"""群号"""
|
||||
group_name: str | None = None
|
||||
"""群名称"""
|
||||
group_avatar: str | None = None
|
||||
"""群头像"""
|
||||
group_owner: str | None = None
|
||||
"""群主 id"""
|
||||
group_admins: list[str] | None = None
|
||||
"""群管理员 id"""
|
||||
members: list[MessageMember] | None = None
|
||||
"""所有群成员"""
|
||||
|
||||
def __str__(self):
|
||||
return (
|
||||
f"Group ID: {self.group_id}\n"
|
||||
f"Name: {self.group_name if self.group_name else 'N/A'}\n"
|
||||
f"Avatar: {self.group_avatar if self.group_avatar else 'N/A'}\n"
|
||||
f"Owner ID: {self.group_owner if self.group_owner else 'N/A'}\n"
|
||||
f"Admin IDs: {self.group_admins if self.group_admins else 'N/A'}\n"
|
||||
f"Members Len: {len(self.members) if self.members else 0}\n"
|
||||
f"First Member: {self.members[0] if self.members else 'N/A'}\n"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AstrBotMessage:
|
||||
"""AstrBot 的消息对象"""
|
||||
|
||||
type: MessageType
|
||||
"""消息类型"""
|
||||
self_id: str
|
||||
"""机器人自身 ID"""
|
||||
session_id: str
|
||||
"""会话 ID"""
|
||||
message_id: str
|
||||
"""消息 ID"""
|
||||
sender: MessageMember
|
||||
"""发送者"""
|
||||
message: list[BaseMessageComponent]
|
||||
"""消息链组件列表"""
|
||||
message_str: str
|
||||
"""纯文本消息字符串"""
|
||||
raw_message: dict
|
||||
"""原始消息对象"""
|
||||
timestamp: int
|
||||
"""消息时间戳"""
|
||||
group: Group | None = None
|
||||
"""群信息,如果是私聊则为 None"""
|
||||
|
||||
def __init__(self, **kwargs) -> None:
|
||||
self.timestamp = int(time.time())
|
||||
for key, value in kwargs.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return str(self.__dict__)
|
||||
|
||||
@property
|
||||
def group_id(self) -> str:
|
||||
"""向后兼容的 group_id 属性
|
||||
群组id,如果为私聊,则为空
|
||||
"""
|
||||
if self.group:
|
||||
return self.group.group_id
|
||||
return ""
|
||||
|
||||
@group_id.setter
|
||||
def group_id(self, value: str):
|
||||
"""设置 group_id"""
|
||||
if value:
|
||||
if self.group:
|
||||
self.group.group_id = value
|
||||
else:
|
||||
self.group = Group(group_id=value)
|
||||
else:
|
||||
self.group = None
|
||||
93
src/astrbot_sdk/api/event/event_result.py
Normal file
93
src/astrbot_sdk/api/event/event_result.py
Normal file
@@ -0,0 +1,93 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import enum
|
||||
from dataclasses import dataclass, field
|
||||
from typing import AsyncGenerator
|
||||
from ..message.chain import MessageChain
|
||||
|
||||
|
||||
class EventResultType(enum.Enum):
|
||||
"""用于描述事件处理的结果类型。
|
||||
|
||||
Attributes:
|
||||
CONTINUE: 事件将会继续传播
|
||||
STOP: 事件将会终止传播
|
||||
|
||||
"""
|
||||
|
||||
CONTINUE = enum.auto()
|
||||
STOP = enum.auto()
|
||||
|
||||
|
||||
class ResultContentType(enum.Enum):
|
||||
"""用于描述事件结果的内容的类型。"""
|
||||
|
||||
LLM_RESULT = enum.auto()
|
||||
"""调用 LLM 产生的结果"""
|
||||
GENERAL_RESULT = enum.auto()
|
||||
"""普通的消息结果"""
|
||||
STREAMING_RESULT = enum.auto()
|
||||
"""调用 LLM 产生的流式结果"""
|
||||
STREAMING_FINISH = enum.auto()
|
||||
"""流式输出完成"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class MessageEventResult(MessageChain):
|
||||
"""MessageEventResult 描述了一整条消息中带有的所有组件以及事件处理的结果。
|
||||
现代消息平台的一条富文本消息中可能由多个组件构成,如文本、图片、At 等,并且保留了顺序。
|
||||
|
||||
Attributes:
|
||||
`chain` (list): 用于顺序存储各个组件。
|
||||
`use_t2i_` (bool): 用于标记是否使用文本转图片服务。默认为 None,即跟随用户的设置。当设置为 True 时,将会使用文本转图片服务。
|
||||
`result_type` (EventResultType): 事件处理的结果类型。
|
||||
|
||||
"""
|
||||
|
||||
result_type: EventResultType | None = field(
|
||||
default_factory=lambda: EventResultType.CONTINUE,
|
||||
)
|
||||
|
||||
result_content_type: ResultContentType | None = field(
|
||||
default_factory=lambda: ResultContentType.GENERAL_RESULT,
|
||||
)
|
||||
|
||||
# async_stream: AsyncGenerator | None = None
|
||||
# """异步流"""
|
||||
|
||||
def stop_event(self) -> MessageEventResult:
|
||||
"""终止事件传播。"""
|
||||
self.result_type = EventResultType.STOP
|
||||
return self
|
||||
|
||||
def continue_event(self) -> MessageEventResult:
|
||||
"""继续事件传播。"""
|
||||
self.result_type = EventResultType.CONTINUE
|
||||
return self
|
||||
|
||||
def is_stopped(self) -> bool:
|
||||
"""是否终止事件传播。"""
|
||||
return self.result_type == EventResultType.STOP
|
||||
|
||||
def set_async_stream(self, stream: AsyncGenerator) -> MessageEventResult:
|
||||
"""设置异步流。"""
|
||||
self.async_stream = stream
|
||||
return self
|
||||
|
||||
def set_result_content_type(self, typ: ResultContentType) -> MessageEventResult:
|
||||
"""设置事件处理的结果类型。
|
||||
|
||||
Args:
|
||||
result_type (EventResultType): 事件处理的结果类型。
|
||||
|
||||
"""
|
||||
self.result_content_type = typ
|
||||
return self
|
||||
|
||||
def is_llm_result(self) -> bool:
|
||||
"""是否为 LLM 结果。"""
|
||||
return self.result_content_type == ResultContentType.LLM_RESULT
|
||||
|
||||
|
||||
# 为了兼容旧版代码,保留 CommandResult 的别名
|
||||
CommandResult = MessageEventResult
|
||||
19
src/astrbot_sdk/api/event/event_type.py
Normal file
19
src/astrbot_sdk/api/event/event_type.py
Normal file
@@ -0,0 +1,19 @@
|
||||
from __future__ import annotations
|
||||
import enum
|
||||
|
||||
|
||||
class EventType(enum.Enum):
|
||||
"""表示一个 AstrBot 内部事件的类型。如适配器消息事件、LLM 请求事件、发送消息前的事件等
|
||||
|
||||
用于对 Handler 的职能分组。
|
||||
"""
|
||||
|
||||
OnAstrBotLoadedEvent = enum.auto() # AstrBot 加载完成
|
||||
OnPlatformLoadedEvent = enum.auto() # 平台加载完成
|
||||
|
||||
AdapterMessageEvent = enum.auto() # 收到适配器发来的消息
|
||||
OnLLMRequestEvent = enum.auto() # 收到 LLM 请求(可以是用户也可以是插件)
|
||||
OnLLMResponseEvent = enum.auto() # LLM 响应后
|
||||
OnDecoratingResultEvent = enum.auto() # 发送消息前
|
||||
OnCallingFuncToolEvent = enum.auto() # 调用函数工具
|
||||
OnAfterMessageSentEvent = enum.auto() # 发送消息后
|
||||
52
src/astrbot_sdk/api/event/filter.py
Normal file
52
src/astrbot_sdk/api/event/filter.py
Normal file
@@ -0,0 +1,52 @@
|
||||
from ...runtime.stars.filter.custom_filter import CustomFilter
|
||||
from ...runtime.stars.filter.event_message_type import (
|
||||
EventMessageType,
|
||||
EventMessageTypeFilter,
|
||||
)
|
||||
from ...runtime.stars.filter.permission import PermissionType, PermissionTypeFilter
|
||||
from ...runtime.stars.filter.platform_adapter_type import (
|
||||
PlatformAdapterType,
|
||||
PlatformAdapterTypeFilter,
|
||||
)
|
||||
from ...runtime.stars.registry.register import register_after_message_sent as after_message_sent
|
||||
from ...runtime.stars.registry.register import register_command as command
|
||||
from ...runtime.stars.registry.register import register_command_group as command_group
|
||||
from ...runtime.stars.registry.register import register_custom_filter as custom_filter
|
||||
from ...runtime.stars.registry.register import register_event_message_type as event_message_type
|
||||
# from ...runtime.stars.registry.register import register_llm_tool as llm_tool
|
||||
from ...runtime.stars.registry.register import register_on_astrbot_loaded as on_astrbot_loaded
|
||||
from ...runtime.stars.registry.register import (
|
||||
register_on_decorating_result as on_decorating_result,
|
||||
)
|
||||
from ...runtime.stars.registry.register import register_on_llm_request as on_llm_request
|
||||
from ...runtime.stars.registry.register import register_on_llm_response as on_llm_response
|
||||
from ...runtime.stars.registry.register import register_on_platform_loaded as on_platform_loaded
|
||||
from ...runtime.stars.registry.register import register_permission_type as permission_type
|
||||
from ...runtime.stars.registry.register import (
|
||||
register_platform_adapter_type as platform_adapter_type,
|
||||
)
|
||||
from ...runtime.stars.registry.register import register_regex as regex
|
||||
|
||||
__all__ = [
|
||||
"CustomFilter",
|
||||
"EventMessageType",
|
||||
"EventMessageTypeFilter",
|
||||
"PermissionType",
|
||||
"PermissionTypeFilter",
|
||||
"PlatformAdapterType",
|
||||
"PlatformAdapterTypeFilter",
|
||||
"after_message_sent",
|
||||
"command",
|
||||
"command_group",
|
||||
"custom_filter",
|
||||
"event_message_type",
|
||||
# "llm_tool",
|
||||
"on_astrbot_loaded",
|
||||
"on_decorating_result",
|
||||
"on_llm_request",
|
||||
"on_llm_response",
|
||||
"on_platform_loaded",
|
||||
"permission_type",
|
||||
"platform_adapter_type",
|
||||
"regex",
|
||||
]
|
||||
32
src/astrbot_sdk/api/event/message_session.py
Normal file
32
src/astrbot_sdk/api/event/message_session.py
Normal file
@@ -0,0 +1,32 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
from ..event.message_type import MessageType
|
||||
|
||||
|
||||
@dataclass
|
||||
class MessageSession:
|
||||
"""
|
||||
描述一条消息在 AstrBot 中对应的会话的唯一标识。
|
||||
如果您需要实例化 MessageSession,请不要给 platform_id 赋值(或者同时给 platform_name 和 platform_id 赋值相同值)。
|
||||
它会在 __post_init__ 中自动设置为 platform_name 的值。
|
||||
"""
|
||||
|
||||
platform_name: str
|
||||
"""平台适配器实例的唯一标识符。自 AstrBot v4.0.0 起,该字段实际为 platform_id。"""
|
||||
message_type: MessageType
|
||||
session_id: str
|
||||
platform_id: str | None = None
|
||||
|
||||
def __str__(self):
|
||||
return f"{self.platform_id}:{self.message_type.value}:{self.session_id}"
|
||||
|
||||
def __post_init__(self):
|
||||
self.platform_id = self.platform_name
|
||||
|
||||
@staticmethod
|
||||
def from_str(session_str: str):
|
||||
platform_id, message_type, session_id = session_str.split(":")
|
||||
return MessageSession(platform_id, MessageType(message_type), session_id)
|
||||
|
||||
|
||||
MessageSesion = MessageSession # back compatibility
|
||||
7
src/astrbot_sdk/api/event/message_type.py
Normal file
7
src/astrbot_sdk/api/event/message_type.py
Normal file
@@ -0,0 +1,7 @@
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class MessageType(Enum):
|
||||
GROUP_MESSAGE = "GroupMessage" # 群组形式的消息
|
||||
FRIEND_MESSAGE = "FriendMessage" # 私聊、好友等单聊消息
|
||||
OTHER_MESSAGE = "OtherMessage" # 其他类型的消息,如系统消息等
|
||||
136
src/astrbot_sdk/api/message/chain.py
Normal file
136
src/astrbot_sdk/api/message/chain.py
Normal file
@@ -0,0 +1,136 @@
|
||||
from . import components as Comp
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class MessageChain:
|
||||
"""MessageChain 描述了一整条消息中带有的所有组件。
|
||||
现代消息平台的一条富文本消息中可能由多个组件构成,如文本、图片、At 等,并且保留了顺序。
|
||||
|
||||
Attributes:
|
||||
`chain` (list): 用于顺序存储各个组件。
|
||||
`use_t2i_` (bool): 用于标记是否使用文本转图片服务。默认为 None,即跟随用户的设置。当设置为 True 时,将会使用文本转图片服务。
|
||||
|
||||
"""
|
||||
|
||||
chain: list[Comp.BaseMessageComponent] = field(default_factory=list)
|
||||
use_t2i_: bool | None = None # None 为跟随用户设置
|
||||
type: str | None = None
|
||||
"""消息链承载的消息的类型。可选,用于让消息平台区分不同业务场景的消息链。"""
|
||||
|
||||
def message(self, message: str):
|
||||
"""添加一条文本消息到消息链 `chain` 中。
|
||||
|
||||
Example:
|
||||
CommandResult().message("Hello ").message("world!")
|
||||
# 输出 Hello world!
|
||||
|
||||
"""
|
||||
self.chain.append(Comp.Plain(text=message))
|
||||
return self
|
||||
|
||||
def at(self, name: str, qq: str):
|
||||
"""添加一条 At 消息到消息链 `chain` 中。
|
||||
|
||||
Example:
|
||||
CommandResult().at("张三", "12345678910")
|
||||
# 输出 @张三
|
||||
|
||||
"""
|
||||
self.chain.append(Comp.At(user_id=qq, user_name=name))
|
||||
return self
|
||||
|
||||
def at_all(self):
|
||||
"""添加一条 AtAll 消息到消息链 `chain` 中。
|
||||
|
||||
Example:
|
||||
CommandResult().at_all()
|
||||
# 输出 @所有人
|
||||
|
||||
"""
|
||||
self.chain.append(Comp.AtAll())
|
||||
return self
|
||||
|
||||
def error(self, message: str):
|
||||
"""[Deprecated] 添加一条错误消息到消息链 `chain` 中
|
||||
|
||||
Example:
|
||||
CommandResult().error("解析失败")
|
||||
|
||||
"""
|
||||
self.chain.append(Comp.Plain(text=message))
|
||||
return self
|
||||
|
||||
def url_image(self, url: str):
|
||||
"""添加一条图片消息(https 链接)到消息链 `chain` 中。
|
||||
|
||||
Note:
|
||||
如果需要发送本地图片,请使用 `file_image` 方法。
|
||||
|
||||
Example:
|
||||
CommandResult().image("https://example.com/image.jpg")
|
||||
|
||||
"""
|
||||
self.chain.append(Comp.Image(file=url))
|
||||
return self
|
||||
|
||||
def file_image(self, path: str):
|
||||
"""添加一条图片消息(本地文件路径)到消息链 `chain` 中。
|
||||
|
||||
Note:
|
||||
如果需要发送网络图片,请使用 `url_image` 方法。
|
||||
|
||||
Example:
|
||||
CommandResult().file_image("image.jpg")
|
||||
"""
|
||||
self.chain.append(Comp.Image(file=path))
|
||||
return self
|
||||
|
||||
def base64_image(self, base64_str: str):
|
||||
"""添加一条图片消息(base64 编码字符串)到消息链 `chain` 中。
|
||||
|
||||
Example:
|
||||
CommandResult().base64_image("iVBORw0KGgoAAAANSUhEUgAAAAUA...")
|
||||
"""
|
||||
self.chain.append(Comp.Image(file=base64_str))
|
||||
return self
|
||||
|
||||
def use_t2i(self, use_t2i: bool):
|
||||
"""设置是否使用文本转图片服务。
|
||||
|
||||
Args:
|
||||
use_t2i (bool): 是否使用文本转图片服务。默认为 None,即跟随用户的设置。当设置为 True 时,将会使用文本转图片服务。
|
||||
|
||||
"""
|
||||
self.use_t2i_ = use_t2i
|
||||
return self
|
||||
|
||||
def get_plain_text(self) -> str:
|
||||
"""获取纯文本消息。这个方法将获取 chain 中所有 Plain 组件的文本并拼接成一条消息。空格分隔。"""
|
||||
return " ".join(
|
||||
[comp.text for comp in self.chain if isinstance(comp, Comp.Plain)]
|
||||
)
|
||||
|
||||
def squash_plain(self):
|
||||
"""将消息链中的所有 Plain 消息段聚合到第一个 Plain 消息段中。"""
|
||||
if not self.chain:
|
||||
return None
|
||||
|
||||
new_chain = []
|
||||
first_plain = None
|
||||
plain_texts = []
|
||||
|
||||
for comp in self.chain:
|
||||
if isinstance(comp, Comp.Plain):
|
||||
if first_plain is None:
|
||||
first_plain = comp
|
||||
new_chain.append(comp)
|
||||
plain_texts.append(comp.text)
|
||||
else:
|
||||
new_chain.append(comp)
|
||||
|
||||
if first_plain is not None:
|
||||
first_plain.text = "".join(plain_texts)
|
||||
|
||||
self.chain = new_chain
|
||||
return self
|
||||
225
src/astrbot_sdk/api/message/components.py
Normal file
225
src/astrbot_sdk/api/message/components.py
Normal file
@@ -0,0 +1,225 @@
|
||||
from __future__ import annotations
|
||||
from enum import Enum
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Literal
|
||||
|
||||
|
||||
class ComponentType(str, Enum):
|
||||
# Basic Segment Types
|
||||
Plain = "Plain" # plain text message
|
||||
Image = "Image" # image
|
||||
Record = "Record" # audio
|
||||
Video = "Video" # video
|
||||
File = "File" # file attachment
|
||||
|
||||
# IM-specific Segment Types
|
||||
Face = "Face" # Emoji segment for Tencent QQ platform
|
||||
At = "At" # mention a user in IM apps
|
||||
Node = "Node" # a node in a forwarded message
|
||||
Nodes = "Nodes" # a forwarded message consisting of multiple nodes
|
||||
Poke = "Poke" # a poke message for Tencent QQ platform
|
||||
Reply = "Reply" # a reply message segment
|
||||
Forward = "Forward" # a forwarded message segment
|
||||
RPS = "RPS"
|
||||
Dice = "Dice"
|
||||
Shake = "Shake"
|
||||
Share = "Share"
|
||||
Contact = "Contact"
|
||||
Location = "Location"
|
||||
Music = "Music"
|
||||
Json = "Json"
|
||||
Unknown = "Unknown"
|
||||
WechatEmoji = "WechatEmoji"
|
||||
|
||||
|
||||
CompT = ComponentType
|
||||
|
||||
|
||||
class BaseMessageComponent(BaseModel):
|
||||
type: CompT
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Unified dict format"""
|
||||
return self.model_dump()
|
||||
|
||||
|
||||
class Plain(BaseMessageComponent):
|
||||
"""Represents a plain text message segment."""
|
||||
|
||||
type: Literal[CompT.Plain] = CompT.Plain
|
||||
text: str
|
||||
|
||||
|
||||
class Image(BaseMessageComponent):
|
||||
type: Literal[CompT.Image] = CompT.Image
|
||||
file: str
|
||||
"""base64-encoded image data, or file path, or HTTP URL"""
|
||||
|
||||
|
||||
class Record(BaseMessageComponent):
|
||||
type: Literal[CompT.Record] = CompT.Record
|
||||
file: str
|
||||
"""base64-encoded audio data, or file path, or HTTP URL"""
|
||||
|
||||
|
||||
class Video(BaseMessageComponent):
|
||||
type: Literal[CompT.Video] = CompT.Video
|
||||
file: str
|
||||
"""The video file URL."""
|
||||
|
||||
|
||||
class File(BaseMessageComponent):
|
||||
type: Literal[CompT.File] = CompT.File
|
||||
file_name: str
|
||||
mime_type: str | None = None
|
||||
file: str
|
||||
"""The file URL."""
|
||||
|
||||
|
||||
class At(BaseMessageComponent):
|
||||
type: Literal[CompT.At] = CompT.At
|
||||
user_id: str | None = None
|
||||
user_name: str | None = None
|
||||
|
||||
|
||||
class AtAll(At):
|
||||
user_id: str = "all"
|
||||
|
||||
|
||||
class Reply(BaseMessageComponent):
|
||||
type: Literal[CompT.Reply] = CompT.Reply
|
||||
id: str | int
|
||||
"""所引用的消息 ID"""
|
||||
chain: list[BaseMessageComponent] | None = []
|
||||
"""被引用的消息段列表"""
|
||||
sender_id: int | None | str = 0
|
||||
"""被引用的消息对应的发送者的 ID"""
|
||||
sender_nickname: str | None = ""
|
||||
"""被引用的消息对应的发送者的昵称"""
|
||||
time: int | None = 0
|
||||
"""被引用的消息发送时间"""
|
||||
message_str: str | None = ""
|
||||
"""被引用的消息解析后的纯文本消息字符串"""
|
||||
|
||||
|
||||
class Node(BaseMessageComponent):
|
||||
type: Literal[CompT.Node] = CompT.Node
|
||||
sender_id: str
|
||||
nickname: str | None = None
|
||||
content: list[BaseMessageComponent] = Field(default_factory=list)
|
||||
|
||||
|
||||
class Nodes(BaseMessageComponent):
|
||||
type: Literal[CompT.Nodes] = CompT.Nodes
|
||||
nodes: list[Node] = Field(default_factory=list)
|
||||
|
||||
|
||||
class Face(BaseMessageComponent):
|
||||
type: Literal[CompT.Face] = CompT.Face
|
||||
id: int
|
||||
|
||||
|
||||
class RPS(BaseMessageComponent):
|
||||
type: Literal[CompT.RPS] = CompT.RPS
|
||||
|
||||
|
||||
class Dice(BaseMessageComponent):
|
||||
type: Literal[CompT.Dice] = CompT.Dice
|
||||
|
||||
|
||||
class Shake(BaseMessageComponent):
|
||||
type: Literal[CompT.Shake] = CompT.Shake
|
||||
|
||||
|
||||
class Share(BaseMessageComponent):
|
||||
type: Literal[CompT.Share] = CompT.Share
|
||||
url: str
|
||||
title: str
|
||||
content: str | None = ""
|
||||
image: str | None = ""
|
||||
|
||||
|
||||
class Contact(BaseMessageComponent):
|
||||
type: Literal[CompT.Contact] = CompT.Contact
|
||||
_type: str # type 字段冲突
|
||||
id: int | None = 0
|
||||
|
||||
|
||||
class Location(BaseMessageComponent):
|
||||
type: Literal[CompT.Location] = CompT.Location
|
||||
lat: float
|
||||
lon: float
|
||||
title: str | None = ""
|
||||
content: str | None = ""
|
||||
|
||||
|
||||
class Music(BaseMessageComponent):
|
||||
type: Literal[CompT.Music] = CompT.Music
|
||||
_type: str
|
||||
id: int | None = 0
|
||||
url: str | None = ""
|
||||
audio: str | None = ""
|
||||
title: str | None = ""
|
||||
content: str | None = ""
|
||||
image: str | None = ""
|
||||
|
||||
|
||||
class Poke(BaseMessageComponent):
|
||||
type: Literal[CompT.Poke] = CompT.Poke
|
||||
id: int | None = 0
|
||||
qq: int | None = 0
|
||||
|
||||
|
||||
class Forward(BaseMessageComponent):
|
||||
type: Literal[CompT.Forward] = CompT.Forward
|
||||
id: str
|
||||
|
||||
|
||||
class Json(BaseMessageComponent):
|
||||
type: Literal[CompT.Json] = CompT.Json
|
||||
data: dict
|
||||
|
||||
|
||||
class Unknown(BaseMessageComponent):
|
||||
type: Literal[CompT.Unknown] = CompT.Unknown
|
||||
text: str
|
||||
|
||||
|
||||
class WechatEmoji(BaseMessageComponent):
|
||||
type: Literal[CompT.WechatEmoji] = CompT.WechatEmoji
|
||||
md5: str | None = ""
|
||||
md5_len: int | None = 0
|
||||
cdnurl: str | None = ""
|
||||
|
||||
def __init__(self, **_):
|
||||
super().__init__(**_)
|
||||
|
||||
|
||||
ComponentTypes = {
|
||||
# Basic Message Segments
|
||||
"plain": Plain,
|
||||
"text": Plain,
|
||||
"image": Image,
|
||||
"record": Record,
|
||||
"video": Video,
|
||||
"file": File,
|
||||
# IM-specific Message Segments
|
||||
"face": Face,
|
||||
"at": At,
|
||||
"rps": RPS,
|
||||
"dice": Dice,
|
||||
"shake": Shake,
|
||||
"share": Share,
|
||||
"contact": Contact,
|
||||
"location": Location,
|
||||
"music": Music,
|
||||
"reply": Reply,
|
||||
"poke": Poke,
|
||||
"forward": Forward,
|
||||
"node": Node,
|
||||
"nodes": Nodes,
|
||||
"json": Json,
|
||||
"unknown": Unknown,
|
||||
"WechatEmoji": WechatEmoji,
|
||||
}
|
||||
18
src/astrbot_sdk/api/platform/platform_metadata.py
Normal file
18
src/astrbot_sdk/api/platform/platform_metadata.py
Normal file
@@ -0,0 +1,18 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class PlatformMetadata:
|
||||
name: str
|
||||
"""平台的名称,即平台的类型,如 aiocqhttp, discord, slack"""
|
||||
description: str
|
||||
"""平台的描述"""
|
||||
id: str
|
||||
"""平台的唯一标识符,用于配置中识别特定平台"""
|
||||
|
||||
default_config_tmpl: dict | None = None
|
||||
"""平台的默认配置模板"""
|
||||
adapter_display_name: str | None = None
|
||||
"""显示在 WebUI 配置页中的平台名称,如空则是 name"""
|
||||
logo_path: str | None = None
|
||||
"""平台适配器的 logo 文件路径(相对于插件目录)"""
|
||||
0
src/astrbot_sdk/api/star/__init__.py
Normal file
0
src/astrbot_sdk/api/star/__init__.py
Normal file
6
src/astrbot_sdk/api/star/context.py
Normal file
6
src/astrbot_sdk/api/star/context.py
Normal file
@@ -0,0 +1,6 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from ..basic.conversation_mgr import BaseConversationManager
|
||||
|
||||
|
||||
class Context(ABC):
|
||||
conversation_manager: BaseConversationManager
|
||||
59
src/astrbot_sdk/api/star/star.py
Normal file
59
src/astrbot_sdk/api/star/star.py
Normal file
@@ -0,0 +1,59 @@
|
||||
from __future__ import annotations
|
||||
from dataclasses import dataclass, field
|
||||
from ..basic.astrbot_config import AstrBotConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class StarMetadata:
|
||||
"""
|
||||
插件的元数据。
|
||||
当 activated 为 False 时,star_cls 可能为 None,请不要在插件未激活时调用 star_cls 的方法。
|
||||
"""
|
||||
|
||||
name: str | None = None
|
||||
"""插件名"""
|
||||
author: str | None = None
|
||||
"""插件作者"""
|
||||
desc: str | None = None
|
||||
"""插件简介"""
|
||||
version: str | None = None
|
||||
"""插件版本"""
|
||||
repo: str | None = None
|
||||
"""插件仓库地址"""
|
||||
|
||||
# star_cls_type: type[Star] | None = None
|
||||
# """插件的类对象的类型"""
|
||||
|
||||
module_path: str | None = None
|
||||
"""插件的模块路径"""
|
||||
|
||||
# star_cls: Star | None = None
|
||||
# """插件的类对象"""
|
||||
# module: ModuleType | None = None
|
||||
# """插件的模块对象"""
|
||||
|
||||
root_dir_name: str | None = None
|
||||
"""插件的目录名称"""
|
||||
reserved: bool = False
|
||||
"""是否是 AstrBot 的保留插件"""
|
||||
|
||||
activated: bool = True
|
||||
"""是否被激活"""
|
||||
|
||||
config: AstrBotConfig | None = None
|
||||
"""插件配置"""
|
||||
|
||||
star_handler_full_names: list[str] = field(default_factory=list)
|
||||
"""注册的 Handler 的全名列表"""
|
||||
|
||||
display_name: str | None = None
|
||||
"""用于展示的插件名称"""
|
||||
|
||||
logo_path: str | None = None
|
||||
"""插件 Logo 的路径"""
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"Plugin {self.name} ({self.version}) by {self.author}: {self.desc}"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"Plugin {self.name} ({self.version}) by {self.author}: {self.desc}"
|
||||
10
src/astrbot_sdk/runtime/api/context.py
Normal file
10
src/astrbot_sdk/runtime/api/context.py
Normal file
@@ -0,0 +1,10 @@
|
||||
from ...api.star.context import Context as BaseContext
|
||||
from .conversation_mgr import ConversationManager
|
||||
|
||||
|
||||
class Context(BaseContext):
|
||||
def __init__(self, conversation_manager: ConversationManager):
|
||||
self.conversation_manager = conversation_manager
|
||||
|
||||
def _inject_rpc_handlers(self, runner):
|
||||
setattr(self.conversation_manager, "runner", runner)
|
||||
14
src/astrbot_sdk/runtime/api/conversation_mgr.py
Normal file
14
src/astrbot_sdk/runtime/api/conversation_mgr.py
Normal file
@@ -0,0 +1,14 @@
|
||||
from ...api.basic.conversation_mgr import BaseConversationManager
|
||||
from .util import rpc_method
|
||||
|
||||
|
||||
class ConversationManager(BaseConversationManager):
|
||||
@rpc_method
|
||||
async def new_conversation(
|
||||
self,
|
||||
unified_msg_origin: str,
|
||||
platform_id: str | None = None,
|
||||
content: list[dict] | None = None,
|
||||
title: str | None = None,
|
||||
persona_id: str | None = None,
|
||||
) -> str: ...
|
||||
38
src/astrbot_sdk/runtime/api/util.py
Normal file
38
src/astrbot_sdk/runtime/api/util.py
Normal file
@@ -0,0 +1,38 @@
|
||||
import inspect
|
||||
from functools import wraps
|
||||
from typing import Callable
|
||||
from ..star_runner import StarRunner
|
||||
from ..types import CallContextFunctionRequest
|
||||
|
||||
|
||||
def rpc_method(func: Callable) -> Callable:
|
||||
"""sign as an RPC method."""
|
||||
|
||||
@wraps(func)
|
||||
async def wrapper(self, *args, **kwargs):
|
||||
if not hasattr(self, "runner") or not isinstance(self.runner, StarRunner):
|
||||
raise RuntimeError(
|
||||
f"Class {self.__class__.__name__} is not configured for RPC calls."
|
||||
)
|
||||
method_name = f"{self.__class__.__name__}.{func.__name__}"
|
||||
sig = inspect.signature(func)
|
||||
bound_args = sig.bind(self, *args, **kwargs)
|
||||
bound_args.apply_defaults()
|
||||
params = dict(bound_args.arguments)
|
||||
params.pop("self")
|
||||
|
||||
runner: StarRunner = getattr(self, "runner")
|
||||
|
||||
return await runner._call_rpc(
|
||||
CallContextFunctionRequest(
|
||||
jsonrpc="2.0",
|
||||
id=runner._generate_request_id(),
|
||||
method="call_context_function",
|
||||
params=CallContextFunctionRequest.Params(
|
||||
name=method_name,
|
||||
args=params,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
return wrapper
|
||||
36
src/astrbot_sdk/runtime/galaxy.py
Normal file
36
src/astrbot_sdk/runtime/galaxy.py
Normal file
@@ -0,0 +1,36 @@
|
||||
"""
|
||||
VPL means Virtual Star Layer.
|
||||
In the AstrBot 5.0 architecture, VPL is a layer that allows different types of stars to interact with the core system in a standardized way.
|
||||
Currently, AstrBot has two types of stars:
|
||||
1. Legacy Stars: These are the traditional stars that still running in the same runtime as AstrBot core.
|
||||
2. New Stars: These are the modern stars that run in isolated runtime, they communicate with AstrBot core through stdio streams or websocket.
|
||||
|
||||
The VPL module provides the necessary abstractions and interfaces to manage these stars seamlessly,
|
||||
let AstrBot core interact with both types of stars without needing to know the underlying implementation details.
|
||||
"""
|
||||
|
||||
from .stars.virtual import VirtualStar
|
||||
from .stars.new import NewStdioStar, NewWebSocketStar
|
||||
# from .types import StarURI, StarType
|
||||
|
||||
|
||||
class Galaxy:
|
||||
"""Manages the lifecycle and interactions of Virtual Stars (plugins) within AstrBot."""
|
||||
|
||||
vs_map: dict[str, VirtualStar] = {}
|
||||
|
||||
async def connect_to_stdio_star(self, star_name: str, config: dict) -> NewStdioStar:
|
||||
"""Connect to a new-style stdio star given its name."""
|
||||
star = NewStdioStar(**config)
|
||||
await star.initialize()
|
||||
self.vs_map[star_name] = star
|
||||
return star
|
||||
|
||||
async def connect_to_websocket_star(
|
||||
self, star_name: str, config: dict
|
||||
) -> NewWebSocketStar:
|
||||
"""Connect to a new-style websocket star given its name."""
|
||||
star = NewWebSocketStar(**config)
|
||||
await star.initialize()
|
||||
self.vs_map[star_name] = star
|
||||
return star
|
||||
208
src/astrbot_sdk/runtime/rpc/client/README.md
Normal file
208
src/astrbot_sdk/runtime/rpc/client/README.md
Normal file
@@ -0,0 +1,208 @@
|
||||
# JSON-RPC Server Implementation
|
||||
|
||||
This directory contains industry-standard implementations of JSON-RPC 2.0 servers for inter-process communication.
|
||||
|
||||
## Overview
|
||||
|
||||
The implementation follows best practices:
|
||||
|
||||
- **Clean separation of concerns**: Servers handle only communication, not business logic
|
||||
- **Async/await**: Non-blocking I/O for better performance
|
||||
- **Type safety**: Full type hints with Pydantic models
|
||||
- **Error handling**: Proper logging and error propagation
|
||||
- **Resource management**: Clean startup/shutdown lifecycle
|
||||
|
||||
## Architecture
|
||||
|
||||
### Base Class: `JSONRPCServer`
|
||||
|
||||
Abstract base class defining the server interface:
|
||||
|
||||
- `set_message_handler(handler)`: Register a callback for incoming messages
|
||||
- `start()`: Start the server
|
||||
- `stop()`: Stop the server and cleanup
|
||||
- `send_message(message)`: Send a JSON-RPC message
|
||||
|
||||
### STDIO Server: `StdioServer`
|
||||
|
||||
Communicates via standard input/output using line-delimited JSON.
|
||||
|
||||
**Features:**
|
||||
|
||||
- One JSON-RPC message per line
|
||||
- Non-blocking async I/O using executors
|
||||
- Thread-safe write operations with asyncio locks
|
||||
- Graceful EOF handling
|
||||
|
||||
**Use cases:**
|
||||
|
||||
- Plugin subprocess communication
|
||||
- Command-line tools
|
||||
- Pipeline-based architectures
|
||||
|
||||
**Example:**
|
||||
|
||||
```python
|
||||
from astrbot_sdk.runtime.server import StdioServer
|
||||
from astrbot_sdk.runtime.rpc.jsonrpc import JSONRPCMessage
|
||||
|
||||
server = StdioServer()
|
||||
|
||||
def handle_message(message: JSONRPCMessage):
|
||||
# Process the message
|
||||
pass
|
||||
|
||||
server.set_message_handler(handle_message)
|
||||
await server.start()
|
||||
```
|
||||
|
||||
### WebSocket Server: `WebSocketServer`
|
||||
|
||||
Communicates via WebSocket connections.
|
||||
|
||||
**Features:**
|
||||
|
||||
- Single active connection (typical for IPC)
|
||||
- Heartbeat/ping-pong for connection health
|
||||
- Support for text and binary messages
|
||||
- Graceful connection lifecycle management
|
||||
- Built on aiohttp for production readiness
|
||||
|
||||
**Configuration:**
|
||||
|
||||
```python
|
||||
from astrbot_sdk.runtime.server import WebSocketServer
|
||||
|
||||
server = WebSocketServer(
|
||||
host="127.0.0.1",
|
||||
port=8765,
|
||||
path="/rpc",
|
||||
heartbeat=30.0 # seconds, 0 to disable
|
||||
)
|
||||
```
|
||||
|
||||
**Use cases:**
|
||||
|
||||
- Network-based plugin communication
|
||||
- Development/debugging (easier to inspect)
|
||||
- Multiple plugin instances
|
||||
|
||||
## Message Format
|
||||
|
||||
All servers use JSON-RPC 2.0 format:
|
||||
|
||||
**Request:**
|
||||
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "unique-id",
|
||||
"method": "method_name",
|
||||
"params": {"key": "value"}
|
||||
}
|
||||
```
|
||||
|
||||
**Success Response:**
|
||||
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "unique-id",
|
||||
"result": {"data": "response"}
|
||||
}
|
||||
```
|
||||
|
||||
**Error Response:**
|
||||
|
||||
```json
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "unique-id",
|
||||
"error": {
|
||||
"code": -32600,
|
||||
"message": "Invalid Request",
|
||||
"data": null
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Usage Examples
|
||||
|
||||
See the `examples/` directory:
|
||||
|
||||
- `server_stdio_example.py`: STDIO server with echo handler
|
||||
- `server_websocket_example.py`: WebSocket server with echo handler
|
||||
- `client_stdio_test.py`: Test client for STDIO
|
||||
- `client_websocket_test.py`: Test client for WebSocket
|
||||
|
||||
### Running STDIO Example
|
||||
|
||||
Terminal 1 (server):
|
||||
|
||||
```bash
|
||||
python examples/server_stdio_example.py
|
||||
```
|
||||
|
||||
Then type JSON-RPC requests:
|
||||
|
||||
```json
|
||||
{"jsonrpc":"2.0","id":"1","method":"test","params":{"hello":"world"}}
|
||||
```
|
||||
|
||||
Or use the test client:
|
||||
|
||||
```bash
|
||||
python examples/client_stdio_test.py | python examples/server_stdio_example.py
|
||||
```
|
||||
|
||||
### Running WebSocket Example
|
||||
|
||||
Terminal 1 (server):
|
||||
|
||||
```bash
|
||||
python examples/server_websocket_example.py
|
||||
```
|
||||
|
||||
Terminal 2 (client):
|
||||
|
||||
```bash
|
||||
python examples/client_websocket_test.py
|
||||
```
|
||||
|
||||
## Design Principles
|
||||
|
||||
1. **No business logic**: Servers only handle transport and serialization
|
||||
2. **Callback-based**: Use `set_message_handler()` for loose coupling
|
||||
3. **Async-first**: All I/O operations are non-blocking
|
||||
4. **Production-ready**: Proper error handling, logging, and resource cleanup
|
||||
5. **Testable**: Easy to mock and test with custom stdin/stdout
|
||||
|
||||
## Integration with AstrBot SDK
|
||||
|
||||
These servers are designed to be used by the Virtual Plugin Layer (VPL):
|
||||
|
||||
```python
|
||||
# In plugin runtime (subprocess)
|
||||
from astrbot_sdk.runtime.server import StdioServer
|
||||
|
||||
server = StdioServer()
|
||||
server.set_message_handler(handle_core_requests)
|
||||
await server.start()
|
||||
|
||||
# In AstrBot Core
|
||||
# Spawn plugin subprocess with stdio transport
|
||||
# Send JSON-RPC requests to plugin stdin
|
||||
# Receive JSON-RPC responses from plugin stdout
|
||||
```
|
||||
|
||||
## Thread Safety
|
||||
|
||||
- Both servers use `asyncio.Lock` for write operations
|
||||
- Message handlers are called synchronously but can schedule async tasks
|
||||
- Servers must run in an asyncio event loop
|
||||
|
||||
## Error Handling
|
||||
|
||||
- Parse errors are logged but don't crash the server
|
||||
- Connection errors trigger cleanup and can be recovered
|
||||
- User code exceptions in message handlers are contained
|
||||
5
src/astrbot_sdk/runtime/rpc/client/__init__.py
Normal file
5
src/astrbot_sdk/runtime/rpc/client/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
||||
from .base import JSONRPCClient
|
||||
from .stdio import StdioClient
|
||||
from .websocket import WebSocketClient
|
||||
|
||||
__all__ = ["JSONRPCClient", "StdioClient", "WebSocketClient"]
|
||||
14
src/astrbot_sdk/runtime/rpc/client/base.py
Normal file
14
src/astrbot_sdk/runtime/rpc/client/base.py
Normal file
@@ -0,0 +1,14 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC
|
||||
from ..transport import JSONRPCTransport
|
||||
|
||||
|
||||
class JSONRPCClient(JSONRPCTransport, ABC):
|
||||
"""Base class for JSON-RPC clients.
|
||||
|
||||
Handles pure communication (reading/writing JSON-RPC messages).
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
213
src/astrbot_sdk/runtime/rpc/client/stdio.py
Normal file
213
src/astrbot_sdk/runtime/rpc/client/stdio.py
Normal file
@@ -0,0 +1,213 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import subprocess
|
||||
from typing import IO, Any
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from ..jsonrpc import (
|
||||
JSONRPCErrorResponse,
|
||||
JSONRPCMessage,
|
||||
JSONRPCRequest,
|
||||
JSONRPCSuccessResponse,
|
||||
)
|
||||
from .base import JSONRPCClient
|
||||
|
||||
|
||||
class StdioClient(JSONRPCClient):
|
||||
"""JSON-RPC client using standard input/output for communication."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
command: list[str],
|
||||
cwd: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize the STDIO client.
|
||||
|
||||
Args:
|
||||
command: Command to start subprocess (e.g., ['python', 'plugin.py'])
|
||||
cwd: Working directory for subprocess
|
||||
"""
|
||||
super().__init__()
|
||||
self._command = command
|
||||
self._cwd = cwd
|
||||
self._process: subprocess.Popen | None = None
|
||||
self._stdin: IO[Any] | None = None
|
||||
self._stdout: IO[Any] | None = None
|
||||
self._read_task: asyncio.Task | None = None
|
||||
self._write_lock = asyncio.Lock()
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the client and launch subprocess."""
|
||||
if self._running:
|
||||
logger.warning("StdioClient is already running")
|
||||
return
|
||||
|
||||
self._running = True
|
||||
|
||||
# Start subprocess
|
||||
await self._start_subprocess()
|
||||
|
||||
self._read_task = asyncio.create_task(self._read_loop())
|
||||
logger.info("StdioClient started")
|
||||
|
||||
async def _start_subprocess(self) -> None:
|
||||
"""Start the subprocess and connect to its stdio."""
|
||||
logger.info(f"Starting subprocess: {' '.join(self._command)}")
|
||||
|
||||
try:
|
||||
self._process = subprocess.Popen(
|
||||
self._command,
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
cwd=self._cwd,
|
||||
text=True,
|
||||
bufsize=1, # Line buffered
|
||||
)
|
||||
|
||||
# Use subprocess's stdio
|
||||
self._stdin = self._process.stdout # Read from subprocess stdout
|
||||
self._stdout = self._process.stdin # Write to subprocess stdin
|
||||
|
||||
logger.info(f"Subprocess started with PID {self._process.pid}")
|
||||
|
||||
# Start monitoring stderr
|
||||
asyncio.create_task(self._monitor_stderr())
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start subprocess: {e}")
|
||||
raise
|
||||
|
||||
async def _monitor_stderr(self) -> None:
|
||||
"""Monitor subprocess stderr and log output."""
|
||||
if not self._process or not self._process.stderr:
|
||||
return
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
try:
|
||||
while self._running and self._process.poll() is None:
|
||||
line = await loop.run_in_executor(None, self._process.stderr.readline)
|
||||
if line:
|
||||
logger.debug(f"[Subprocess stderr] {line.strip()}")
|
||||
else:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"Error monitoring stderr: {e}")
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the client and terminate subprocess if running."""
|
||||
if not self._running:
|
||||
return
|
||||
|
||||
self._running = False
|
||||
|
||||
# Cancel read task
|
||||
if self._read_task:
|
||||
self._read_task.cancel()
|
||||
try:
|
||||
await self._read_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._read_task = None
|
||||
|
||||
# Terminate subprocess if running
|
||||
if self._process:
|
||||
logger.info("Terminating subprocess...")
|
||||
self._process.terminate()
|
||||
try:
|
||||
self._process.wait(timeout=5.0)
|
||||
logger.info("Subprocess terminated gracefully")
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning("Subprocess did not terminate, killing...")
|
||||
self._process.kill()
|
||||
self._process.wait()
|
||||
logger.info("Subprocess killed")
|
||||
|
||||
self._process = None
|
||||
|
||||
logger.info("StdioClient stopped")
|
||||
|
||||
async def send_message(self, message: JSONRPCMessage) -> None:
|
||||
"""Send a JSON-RPC message to stdout.
|
||||
|
||||
Args:
|
||||
message: The JSON-RPC message to send
|
||||
"""
|
||||
async with self._write_lock:
|
||||
try:
|
||||
json_str = message.model_dump_json(exclude_none=True)
|
||||
await asyncio.get_event_loop().run_in_executor(
|
||||
None, self._write_line, json_str
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to send message: {e}")
|
||||
raise
|
||||
|
||||
def _write_line(self, line: str) -> None:
|
||||
"""Write a line to stdout (synchronous helper)."""
|
||||
if self._stdout:
|
||||
self._stdout.write(line + "\n")
|
||||
self._stdout.flush()
|
||||
|
||||
async def _read_loop(self) -> None:
|
||||
"""Main loop to read messages from stdin."""
|
||||
if not self._stdin:
|
||||
logger.error("No stdin available for reading")
|
||||
return
|
||||
|
||||
logger.debug("Started reading from stdin")
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
try:
|
||||
while self._running:
|
||||
# Read line from stdin in executor to avoid blocking
|
||||
line = await loop.run_in_executor(None, self._stdin.readline)
|
||||
|
||||
if not line:
|
||||
# EOF reached
|
||||
logger.info("EOF reached on stdin")
|
||||
break
|
||||
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
try:
|
||||
# Parse JSON-RPC message
|
||||
message = self._parse_message(line)
|
||||
await self._handle_message(message)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to parse message: {e}, raw line: {line}")
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.debug("Read loop cancelled")
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error in read loop: {e}")
|
||||
finally:
|
||||
logger.debug("Stopped reading from stdin")
|
||||
|
||||
def _parse_message(self, line: str) -> JSONRPCMessage:
|
||||
"""Parse a JSON-RPC message from a string.
|
||||
|
||||
Args:
|
||||
line: JSON string to parse
|
||||
|
||||
Returns:
|
||||
Parsed JSONRPCMessage (Request, SuccessResponse, or ErrorResponse)
|
||||
"""
|
||||
data = json.loads(line)
|
||||
|
||||
# Determine message type based on presence of fields
|
||||
if "method" in data:
|
||||
return JSONRPCRequest.model_validate(data)
|
||||
elif "error" in data:
|
||||
return JSONRPCErrorResponse.model_validate(data)
|
||||
elif "result" in data:
|
||||
return JSONRPCSuccessResponse.model_validate(data)
|
||||
else:
|
||||
raise ValueError(f"Invalid JSON-RPC message: {data}")
|
||||
235
src/astrbot_sdk/runtime/rpc/client/websocket.py
Normal file
235
src/astrbot_sdk/runtime/rpc/client/websocket.py
Normal file
@@ -0,0 +1,235 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import aiohttp
|
||||
from loguru import logger
|
||||
|
||||
from ..jsonrpc import (
|
||||
JSONRPCErrorResponse,
|
||||
JSONRPCMessage,
|
||||
JSONRPCRequest,
|
||||
JSONRPCSuccessResponse,
|
||||
)
|
||||
from .base import JSONRPCClient
|
||||
|
||||
|
||||
class WebSocketClient(JSONRPCClient):
|
||||
"""JSON-RPC client using WebSocket for communication."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: str,
|
||||
heartbeat: float = 30.0,
|
||||
auto_reconnect: bool = True,
|
||||
reconnect_interval: float = 5.0,
|
||||
) -> None:
|
||||
"""Initialize the WebSocket client.
|
||||
|
||||
Args:
|
||||
url: WebSocket server URL (e.g., ws://127.0.0.1:8765/rpc)
|
||||
heartbeat: Heartbeat interval in seconds (0 to disable)
|
||||
auto_reconnect: Whether to automatically reconnect on disconnection
|
||||
reconnect_interval: Interval between reconnection attempts in seconds
|
||||
"""
|
||||
super().__init__()
|
||||
self._url = url
|
||||
self._heartbeat = heartbeat
|
||||
self._auto_reconnect = auto_reconnect
|
||||
self._reconnect_interval = reconnect_interval
|
||||
self._session: aiohttp.ClientSession | None = None
|
||||
self._ws: aiohttp.ClientWebSocketResponse | None = None
|
||||
self._write_lock = asyncio.Lock()
|
||||
self._read_task: asyncio.Task | None = None
|
||||
self._reconnect_task: asyncio.Task | None = None
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Connect to the WebSocket server."""
|
||||
if self._running:
|
||||
logger.warning("WebSocketClient is already running")
|
||||
return
|
||||
|
||||
self._running = True
|
||||
self._session = aiohttp.ClientSession()
|
||||
|
||||
await self._connect()
|
||||
logger.info(f"WebSocketClient started and connected to {self._url}")
|
||||
|
||||
async def _connect(self) -> None:
|
||||
"""Establish WebSocket connection to the server."""
|
||||
try:
|
||||
if not self._session:
|
||||
raise RuntimeError("Session not initialized")
|
||||
|
||||
self._ws = await self._session.ws_connect(
|
||||
self._url,
|
||||
heartbeat=self._heartbeat if self._heartbeat > 0 else None,
|
||||
)
|
||||
logger.info(f"Connected to WebSocket server: {self._url}")
|
||||
|
||||
# Start reading messages
|
||||
self._read_task = asyncio.create_task(self._read_loop())
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to connect to WebSocket server: {e}")
|
||||
if self._auto_reconnect and self._running:
|
||||
logger.info(
|
||||
f"Will retry connection in {self._reconnect_interval} seconds..."
|
||||
)
|
||||
await asyncio.sleep(self._reconnect_interval)
|
||||
if self._running:
|
||||
await self._connect()
|
||||
else:
|
||||
raise
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Disconnect from the WebSocket server and cleanup resources."""
|
||||
if not self._running:
|
||||
return
|
||||
|
||||
self._running = False
|
||||
|
||||
# Cancel reconnection task if running
|
||||
if self._reconnect_task and not self._reconnect_task.done():
|
||||
self._reconnect_task.cancel()
|
||||
try:
|
||||
await self._reconnect_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._reconnect_task = None
|
||||
|
||||
# Cancel read task
|
||||
if self._read_task and not self._read_task.done():
|
||||
self._read_task.cancel()
|
||||
try:
|
||||
await self._read_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._read_task = None
|
||||
|
||||
# Close WebSocket connection
|
||||
if self._ws and not self._ws.closed:
|
||||
await self._ws.close()
|
||||
self._ws = None
|
||||
|
||||
# Close session
|
||||
if self._session and not self._session.closed:
|
||||
await self._session.close()
|
||||
self._session = None
|
||||
|
||||
logger.info("WebSocketClient stopped")
|
||||
|
||||
async def send_message(self, message: JSONRPCMessage) -> None:
|
||||
"""Send a JSON-RPC message through the WebSocket.
|
||||
|
||||
Args:
|
||||
message: The JSON-RPC message to send
|
||||
|
||||
Raises:
|
||||
RuntimeError: If no WebSocket connection is active
|
||||
"""
|
||||
if not self._ws or self._ws.closed:
|
||||
raise RuntimeError("No active WebSocket connection")
|
||||
|
||||
async with self._write_lock:
|
||||
try:
|
||||
json_str = message.model_dump_json(exclude_none=True)
|
||||
await self._ws.send_str(json_str)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to send message: {e}")
|
||||
raise
|
||||
|
||||
async def _read_loop(self) -> None:
|
||||
"""Main loop to read messages from WebSocket."""
|
||||
if not self._ws:
|
||||
logger.error("WebSocket connection not established")
|
||||
return
|
||||
|
||||
logger.debug("Started reading from WebSocket")
|
||||
|
||||
try:
|
||||
async for msg in self._ws:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
try:
|
||||
message = self._parse_message(msg.data)
|
||||
await self._handle_message(message)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to parse message: {e}, raw data: {msg.data}"
|
||||
)
|
||||
|
||||
elif msg.type == aiohttp.WSMsgType.BINARY:
|
||||
try:
|
||||
text = msg.data.decode("utf-8")
|
||||
message = self._parse_message(text)
|
||||
await self._handle_message(message)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to parse binary message: {e}")
|
||||
|
||||
elif msg.type == aiohttp.WSMsgType.ERROR:
|
||||
if self._ws:
|
||||
logger.error(f"WebSocket error: {self._ws.exception()}")
|
||||
break
|
||||
|
||||
elif msg.type in (
|
||||
aiohttp.WSMsgType.CLOSE,
|
||||
aiohttp.WSMsgType.CLOSING,
|
||||
aiohttp.WSMsgType.CLOSED,
|
||||
):
|
||||
logger.debug("WebSocket closing")
|
||||
break
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.debug("Read loop cancelled")
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error in read loop: {e}")
|
||||
finally:
|
||||
logger.debug("Stopped reading from WebSocket")
|
||||
|
||||
# Handle reconnection
|
||||
if self._running and self._auto_reconnect:
|
||||
logger.info("Connection lost, attempting to reconnect...")
|
||||
self._reconnect_task = asyncio.create_task(self._reconnect())
|
||||
|
||||
async def _reconnect(self) -> None:
|
||||
"""Attempt to reconnect to the WebSocket server."""
|
||||
while self._running and self._auto_reconnect:
|
||||
try:
|
||||
logger.info(
|
||||
f"Reconnecting to {self._url} in {self._reconnect_interval} seconds..."
|
||||
)
|
||||
await asyncio.sleep(self._reconnect_interval)
|
||||
|
||||
if not self._running:
|
||||
break
|
||||
|
||||
await self._connect()
|
||||
logger.info("Reconnected successfully")
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Reconnection failed: {e}")
|
||||
# Continue loop to retry
|
||||
|
||||
def _parse_message(self, data: str) -> JSONRPCMessage:
|
||||
"""Parse a JSON-RPC message from a string.
|
||||
|
||||
Args:
|
||||
data: JSON string to parse
|
||||
|
||||
Returns:
|
||||
Parsed JSONRPCMessage (Request, SuccessResponse, or ErrorResponse)
|
||||
"""
|
||||
obj = json.loads(data)
|
||||
|
||||
# Determine message type based on presence of fields
|
||||
if "method" in obj:
|
||||
return JSONRPCRequest.model_validate(obj)
|
||||
elif "error" in obj:
|
||||
return JSONRPCErrorResponse.model_validate(obj)
|
||||
elif "result" in obj:
|
||||
return JSONRPCSuccessResponse.model_validate(obj)
|
||||
else:
|
||||
raise ValueError(f"Invalid JSON-RPC message: {obj}")
|
||||
39
src/astrbot_sdk/runtime/rpc/jsonrpc.py
Normal file
39
src/astrbot_sdk/runtime/rpc/jsonrpc.py
Normal file
@@ -0,0 +1,39 @@
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class _JSONRPCBaseMessage(BaseModel):
|
||||
jsonrpc: Literal["2.0"]
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
class JSONRPCRequest(_JSONRPCBaseMessage):
|
||||
id: str | None = None
|
||||
method: str
|
||||
params: dict[str, Any] = Field(default_factory=dict)
|
||||
"""A request that expects a response."""
|
||||
|
||||
|
||||
class _Result(_JSONRPCBaseMessage):
|
||||
id: str | None
|
||||
|
||||
|
||||
class JSONRPCSuccessResponse(_Result):
|
||||
result: dict[str, Any] = Field(default_factory=dict)
|
||||
"""A successful response to a request."""
|
||||
|
||||
|
||||
class JSONRPCErrorData(BaseModel):
|
||||
code: int
|
||||
message: str
|
||||
data: Any | None = None
|
||||
|
||||
|
||||
class JSONRPCErrorResponse(_Result):
|
||||
error: JSONRPCErrorData
|
||||
"""An error response to a request."""
|
||||
|
||||
|
||||
JSONRPCMessage = JSONRPCRequest | JSONRPCSuccessResponse | JSONRPCErrorResponse
|
||||
9
src/astrbot_sdk/runtime/rpc/server/__init__.py
Normal file
9
src/astrbot_sdk/runtime/rpc/server/__init__.py
Normal file
@@ -0,0 +1,9 @@
|
||||
from .base import JSONRPCServer
|
||||
from .stdio import StdioServer
|
||||
from .websockets import WebSocketServer
|
||||
|
||||
__all__ = [
|
||||
"JSONRPCServer",
|
||||
"StdioServer",
|
||||
"WebSocketServer",
|
||||
]
|
||||
15
src/astrbot_sdk/runtime/rpc/server/base.py
Normal file
15
src/astrbot_sdk/runtime/rpc/server/base.py
Normal file
@@ -0,0 +1,15 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC
|
||||
from ..transport import JSONRPCTransport
|
||||
|
||||
|
||||
class JSONRPCServer(JSONRPCTransport, ABC):
|
||||
"""Base class for JSON-RPC servers.
|
||||
|
||||
Handles pure communication (reading/writing JSON-RPC messages).
|
||||
Server runs in plugin process and receives messages from AstrBot.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
144
src/astrbot_sdk/runtime/rpc/server/stdio.py
Normal file
144
src/astrbot_sdk/runtime/rpc/server/stdio.py
Normal file
@@ -0,0 +1,144 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
from typing import IO, Any
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from ..jsonrpc import (
|
||||
JSONRPCErrorResponse,
|
||||
JSONRPCMessage,
|
||||
JSONRPCRequest,
|
||||
JSONRPCSuccessResponse,
|
||||
)
|
||||
from .base import JSONRPCServer
|
||||
|
||||
|
||||
class StdioServer(JSONRPCServer):
|
||||
"""JSON-RPC server using standard input/output for communication.
|
||||
|
||||
This runs in the plugin process and communicates with AstrBot via stdio.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stdin: IO[Any] | None = None,
|
||||
stdout: IO[Any] | None = None,
|
||||
) -> None:
|
||||
"""Initialize the STDIO server.
|
||||
|
||||
Args:
|
||||
stdin: Input stream to read from (defaults to sys.stdin)
|
||||
stdout: Output stream to write to (defaults to sys.stdout)
|
||||
"""
|
||||
super().__init__()
|
||||
self._stdin = stdin or sys.stdin
|
||||
self._stdout = stdout or sys.stdout
|
||||
self._read_task: asyncio.Task | None = None
|
||||
self._write_lock = asyncio.Lock()
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the server and begin reading from stdin."""
|
||||
if self._running:
|
||||
logger.warning("StdioServer is already running")
|
||||
return
|
||||
|
||||
self._running = True
|
||||
self._read_task = asyncio.create_task(self._read_loop())
|
||||
logger.info("StdioServer started")
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the server and cleanup resources."""
|
||||
if not self._running:
|
||||
return
|
||||
|
||||
self._running = False
|
||||
|
||||
# Cancel read task
|
||||
if self._read_task:
|
||||
self._read_task.cancel()
|
||||
try:
|
||||
await self._read_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._read_task = None
|
||||
|
||||
logger.info("StdioServer stopped")
|
||||
|
||||
async def send_message(self, message: JSONRPCMessage) -> None:
|
||||
"""Send a JSON-RPC message to stdout.
|
||||
|
||||
Args:
|
||||
message: The JSON-RPC message to send
|
||||
"""
|
||||
async with self._write_lock:
|
||||
try:
|
||||
json_str = message.model_dump_json(exclude_none=True)
|
||||
await asyncio.get_event_loop().run_in_executor(
|
||||
None, self._write_line, json_str
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to send message: {e}")
|
||||
raise
|
||||
|
||||
def _write_line(self, line: str) -> None:
|
||||
"""Write a line to stdout (synchronous helper)."""
|
||||
self._stdout.write(line + "\n")
|
||||
self._stdout.flush()
|
||||
|
||||
async def _read_loop(self) -> None:
|
||||
"""Main loop to read messages from stdin."""
|
||||
logger.debug("Started reading from stdin")
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
try:
|
||||
while self._running:
|
||||
# Read line from stdin in executor to avoid blocking
|
||||
line = await loop.run_in_executor(None, self._stdin.readline)
|
||||
|
||||
if not line:
|
||||
# EOF reached
|
||||
logger.info("EOF reached on stdin")
|
||||
break
|
||||
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
try:
|
||||
# Parse JSON-RPC message
|
||||
message = self._parse_message(line)
|
||||
await self._handle_message(message)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to parse message: {e}, raw line: {line}")
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.debug("Read loop cancelled")
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error in read loop: {e}")
|
||||
finally:
|
||||
logger.debug("Stopped reading from stdin")
|
||||
|
||||
def _parse_message(self, line: str) -> JSONRPCMessage:
|
||||
"""Parse a JSON-RPC message from a string.
|
||||
|
||||
Args:
|
||||
line: JSON string to parse
|
||||
|
||||
Returns:
|
||||
Parsed JSONRPCMessage (Request, SuccessResponse, or ErrorResponse)
|
||||
"""
|
||||
data = json.loads(line)
|
||||
|
||||
# Determine message type based on presence of fields
|
||||
if "method" in data:
|
||||
return JSONRPCRequest.model_validate(data)
|
||||
elif "error" in data:
|
||||
return JSONRPCErrorResponse.model_validate(data)
|
||||
elif "result" in data:
|
||||
return JSONRPCSuccessResponse.model_validate(data)
|
||||
else:
|
||||
raise ValueError(f"Invalid JSON-RPC message: {data}")
|
||||
236
src/astrbot_sdk/runtime/rpc/server/websockets.py
Normal file
236
src/astrbot_sdk/runtime/rpc/server/websockets.py
Normal file
@@ -0,0 +1,236 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
from loguru import logger
|
||||
|
||||
from ..jsonrpc import (
|
||||
JSONRPCErrorResponse,
|
||||
JSONRPCMessage,
|
||||
JSONRPCRequest,
|
||||
JSONRPCSuccessResponse,
|
||||
)
|
||||
from .base import JSONRPCServer
|
||||
|
||||
|
||||
class WebSocketServer(JSONRPCServer):
|
||||
"""JSON-RPC server using WebSocket for communication.
|
||||
|
||||
This runs in the plugin process and accepts connections from AstrBot via WebSocket.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "127.0.0.1",
|
||||
port: int = 0, # 0 means auto-assign
|
||||
path: str = "/",
|
||||
heartbeat: float = 30.0,
|
||||
) -> None:
|
||||
"""Initialize the WebSocket server.
|
||||
|
||||
Args:
|
||||
host: Host to bind to
|
||||
port: Port to bind to (0 for auto-assign)
|
||||
path: WebSocket endpoint path
|
||||
heartbeat: Heartbeat interval in seconds (0 to disable)
|
||||
"""
|
||||
super().__init__()
|
||||
self._host = host
|
||||
self._port = port
|
||||
self._path = path
|
||||
self._heartbeat = heartbeat
|
||||
self._app: web.Application | None = None
|
||||
self._runner: web.AppRunner | None = None
|
||||
self._site: web.TCPSite | None = None
|
||||
self._ws: web.WebSocketResponse | None = None
|
||||
self._write_lock = asyncio.Lock()
|
||||
self._actual_port: int | None = None
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the WebSocket server and begin listening for connections."""
|
||||
if self._running:
|
||||
logger.warning("WebSocketServer is already running")
|
||||
return
|
||||
|
||||
self._running = True
|
||||
self._app = web.Application()
|
||||
self._app.router.add_get(self._path, self._handle_websocket)
|
||||
|
||||
self._runner = web.AppRunner(self._app)
|
||||
await self._runner.setup()
|
||||
|
||||
self._site = web.TCPSite(self._runner, self._host, self._port)
|
||||
await self._site.start()
|
||||
|
||||
# Get the actual port (useful when port=0)
|
||||
if self._site._server and hasattr(self._site._server, "sockets"):
|
||||
sockets = getattr(self._site._server, "sockets", None)
|
||||
if sockets:
|
||||
for socket in sockets:
|
||||
self._actual_port = socket.getsockname()[1]
|
||||
break
|
||||
|
||||
logger.info(
|
||||
f"WebSocketServer started on ws://{self._host}:{self._actual_port or self._port}{self._path}"
|
||||
)
|
||||
|
||||
async def _handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
|
||||
"""Handle incoming WebSocket connections.
|
||||
|
||||
Args:
|
||||
request: The aiohttp request object
|
||||
|
||||
Returns:
|
||||
WebSocket response
|
||||
"""
|
||||
ws = web.WebSocketResponse(
|
||||
heartbeat=self._heartbeat if self._heartbeat > 0 else None
|
||||
)
|
||||
await ws.prepare(request)
|
||||
|
||||
# Only allow one connection at a time (typical for plugin IPC)
|
||||
if self._ws and not self._ws.closed:
|
||||
logger.warning(
|
||||
"Rejecting new connection - already have an active connection"
|
||||
)
|
||||
await ws.close(
|
||||
code=1008, message=b"Server already has an active connection"
|
||||
)
|
||||
return ws
|
||||
|
||||
self._ws = ws
|
||||
logger.info(f"WebSocket connection established from {request.remote}")
|
||||
|
||||
try:
|
||||
await self._message_loop(ws)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in WebSocket message loop: {e}")
|
||||
finally:
|
||||
if self._ws == ws:
|
||||
self._ws = None
|
||||
logger.info("WebSocket connection closed")
|
||||
|
||||
return ws
|
||||
|
||||
async def _message_loop(self, ws: web.WebSocketResponse) -> None:
|
||||
"""Main loop to receive messages from WebSocket.
|
||||
|
||||
Args:
|
||||
ws: The WebSocket response object
|
||||
"""
|
||||
async for msg in ws:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
try:
|
||||
message = self._parse_message(msg.data)
|
||||
await self._handle_message(message)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to parse message: {e}, raw data: {msg.data}")
|
||||
|
||||
elif msg.type == aiohttp.WSMsgType.BINARY:
|
||||
try:
|
||||
text = msg.data.decode("utf-8")
|
||||
message = self._parse_message(text)
|
||||
await self._handle_message(message)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to parse binary message: {e}")
|
||||
|
||||
elif msg.type == aiohttp.WSMsgType.ERROR:
|
||||
logger.error(f"WebSocket error: {ws.exception()}")
|
||||
break
|
||||
|
||||
elif msg.type in (
|
||||
aiohttp.WSMsgType.CLOSE,
|
||||
aiohttp.WSMsgType.CLOSING,
|
||||
aiohttp.WSMsgType.CLOSED,
|
||||
):
|
||||
logger.debug("WebSocket closing")
|
||||
break
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the WebSocket server and cleanup resources."""
|
||||
if not self._running:
|
||||
return
|
||||
|
||||
self._running = False
|
||||
|
||||
# Close active WebSocket connection
|
||||
if self._ws and not self._ws.closed:
|
||||
await self._ws.close()
|
||||
self._ws = None
|
||||
|
||||
# Cleanup server
|
||||
if self._site:
|
||||
await self._site.stop()
|
||||
self._site = None
|
||||
|
||||
if self._runner:
|
||||
await self._runner.cleanup()
|
||||
self._runner = None
|
||||
|
||||
self._app = None
|
||||
logger.info("WebSocketServer stopped")
|
||||
|
||||
async def send_message(self, message: JSONRPCMessage) -> None:
|
||||
"""Send a JSON-RPC message through the WebSocket.
|
||||
|
||||
Args:
|
||||
message: The JSON-RPC message to send
|
||||
|
||||
Raises:
|
||||
RuntimeError: If no WebSocket connection is active
|
||||
"""
|
||||
if not self._ws or self._ws.closed:
|
||||
raise RuntimeError("No active WebSocket connection")
|
||||
|
||||
async with self._write_lock:
|
||||
try:
|
||||
json_str = message.model_dump_json(exclude_none=True)
|
||||
await self._ws.send_str(json_str)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to send message: {e}")
|
||||
raise
|
||||
|
||||
@property
|
||||
def port(self) -> int | None:
|
||||
"""Get the actual port the server is listening on.
|
||||
|
||||
Returns:
|
||||
Port number, or None if server is not started
|
||||
"""
|
||||
return self._actual_port or self._port
|
||||
|
||||
@property
|
||||
def url(self) -> str | None:
|
||||
"""Get the WebSocket URL the server is listening on.
|
||||
|
||||
Returns:
|
||||
WebSocket URL, or None if server is not started
|
||||
"""
|
||||
if self._actual_port or self._port:
|
||||
port = self._actual_port or self._port
|
||||
return f"ws://{self._host}:{port}{self._path}"
|
||||
return None
|
||||
|
||||
def _parse_message(self, data: str) -> JSONRPCMessage:
|
||||
"""Parse a JSON-RPC message from a string.
|
||||
|
||||
Args:
|
||||
data: JSON string to parse
|
||||
|
||||
Returns:
|
||||
Parsed JSONRPCMessage (Request, SuccessResponse, or ErrorResponse)
|
||||
"""
|
||||
obj = json.loads(data)
|
||||
|
||||
# Determine message type based on presence of fields
|
||||
if "method" in obj:
|
||||
return JSONRPCRequest.model_validate(obj)
|
||||
elif "error" in obj:
|
||||
return JSONRPCErrorResponse.model_validate(obj)
|
||||
elif "result" in obj:
|
||||
return JSONRPCSuccessResponse.model_validate(obj)
|
||||
else:
|
||||
raise ValueError(f"Invalid JSON-RPC message: {obj}")
|
||||
48
src/astrbot_sdk/runtime/rpc/transport.py
Normal file
48
src/astrbot_sdk/runtime/rpc/transport.py
Normal file
@@ -0,0 +1,48 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Callable, Awaitable
|
||||
|
||||
from .jsonrpc import JSONRPCMessage
|
||||
|
||||
MessageHandler = Callable[[JSONRPCMessage], Awaitable[None]]
|
||||
|
||||
|
||||
class JSONRPCTransport(ABC):
|
||||
"""Base class for JSON-RPC transport layers."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._handler: MessageHandler | None = None
|
||||
self._running = False
|
||||
|
||||
def set_message_handler(self, handler: MessageHandler) -> None:
|
||||
"""Set the handler to be called when a message is received.
|
||||
|
||||
Args:
|
||||
handler: Callback function that receives a JSONRPCMessage
|
||||
"""
|
||||
self._message_handler = handler
|
||||
|
||||
@abstractmethod
|
||||
async def start(self) -> None:
|
||||
"""Start the transport layer."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def stop(self) -> None:
|
||||
"""Stop the transport layer and cleanup resources."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def send_message(self, message: JSONRPCMessage) -> None:
|
||||
"""Send a JSON-RPC message.
|
||||
|
||||
Args:
|
||||
message: The JSON-RPC message to send
|
||||
"""
|
||||
pass
|
||||
|
||||
async def _handle_message(self, message: JSONRPCMessage) -> None:
|
||||
"""Internal method to dispatch received messages to the handler."""
|
||||
if self._message_handler:
|
||||
await self._message_handler(message)
|
||||
102
src/astrbot_sdk/runtime/star_manager.py
Normal file
102
src/astrbot_sdk/runtime/star_manager.py
Normal file
@@ -0,0 +1,102 @@
|
||||
import yaml
|
||||
import importlib
|
||||
import functools
|
||||
from pathlib import Path
|
||||
from loguru import logger
|
||||
from .stars.registry import star_handlers_registry, star_map, star_registry
|
||||
from ..runtime.api.context import Context
|
||||
from ..api.star.star import StarMetadata
|
||||
|
||||
|
||||
class StarManager:
|
||||
def __init__(self, context: Context) -> None:
|
||||
self.context = context
|
||||
|
||||
def discover_star(self, root_dir: Path | None = None):
|
||||
"""
|
||||
Discover star via plugin.yaml.
|
||||
|
||||
Args:
|
||||
root_dir (Path | None): The root directory to search for plugin.yaml. Defaults to None, which means the current working directory.
|
||||
"""
|
||||
if root_dir is None:
|
||||
root_dir = Path.cwd().relative_to(Path.cwd())
|
||||
else:
|
||||
root_dir = Path.cwd().joinpath(root_dir).resolve()
|
||||
path = root_dir / "plugin.yaml"
|
||||
if not path.exists():
|
||||
logger.warning("No plugin.yaml found in the current directory.")
|
||||
return []
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
data = yaml.safe_load(f)
|
||||
|
||||
# Try to find logo.png
|
||||
logo_path = None
|
||||
if Path(root_dir / "logo.png").exists():
|
||||
logo_path = str(root_dir / "logo.png")
|
||||
|
||||
# Validate required fields
|
||||
star_name = data.get("name")
|
||||
if not star_name:
|
||||
logger.error("Plugin name is required in plugin.yaml.")
|
||||
return []
|
||||
|
||||
# Load components
|
||||
components = data.get("components", [])
|
||||
full_name_list = []
|
||||
for comp in components:
|
||||
class_ = comp.get("class", "")
|
||||
print(f"Loading component: {class_}")
|
||||
if not class_:
|
||||
logger.warning(f"Component without class found: {comp}")
|
||||
continue
|
||||
module_path, class_name = class_.rsplit(":", 1)
|
||||
if not module_path:
|
||||
logger.warning(f"Invalid component without module: {comp}")
|
||||
continue
|
||||
# dynamically register the component
|
||||
try:
|
||||
# we need edit the module path to be relative to the root_dir
|
||||
root_dir_dot = str(root_dir).replace("/", ".").lstrip(".")
|
||||
if root_dir_dot:
|
||||
module_path = f"{root_dir_dot}.{module_path}"
|
||||
module_type = importlib.import_module(module_path)
|
||||
logger.info(f"Successfully loaded component module: {module_path}")
|
||||
component_cls = getattr(module_type, class_name)
|
||||
# Instantiate the component with context
|
||||
ccls = component_cls(self.context)
|
||||
|
||||
# add to full name list
|
||||
for h in star_handlers_registry._handlers:
|
||||
if h.handler_full_name.startswith(f"{class_}."):
|
||||
# bind the instance
|
||||
h.handler = functools.partial(h.handler, ccls)
|
||||
full_name_list.append(h.handler_full_name)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load component {module_path}: {e}")
|
||||
continue
|
||||
|
||||
# Register the star metadata
|
||||
star_module_path = f"{star_name}.main"
|
||||
star_metadata = StarMetadata(
|
||||
name=data.get("name"),
|
||||
author=data.get("author"),
|
||||
desc=data.get("desc"),
|
||||
version=data.get("version"),
|
||||
repo=data.get("repo"),
|
||||
module_path=star_module_path,
|
||||
root_dir_name=root_dir.name,
|
||||
reserved=False,
|
||||
star_handler_full_names=full_name_list,
|
||||
display_name=data.get("display_name"),
|
||||
logo_path=logo_path,
|
||||
)
|
||||
star_map[star_module_path] = star_metadata
|
||||
star_registry.append(star_metadata)
|
||||
|
||||
logger.info(f"Discovered {len(star_handlers_registry)} star handlers:")
|
||||
for md in star_handlers_registry:
|
||||
logger.info(
|
||||
f" - {md.handler_full_name} with {len(md.event_filters)} filters"
|
||||
)
|
||||
156
src/astrbot_sdk/runtime/star_runner.py
Normal file
156
src/astrbot_sdk/runtime/star_runner.py
Normal file
@@ -0,0 +1,156 @@
|
||||
import asyncio
|
||||
import inspect
|
||||
from loguru import logger
|
||||
from .rpc.server.base import JSONRPCServer
|
||||
from .stars.registry import star_map, star_handlers_registry
|
||||
from .rpc.jsonrpc import (
|
||||
JSONRPCMessage,
|
||||
JSONRPCRequest,
|
||||
JSONRPCSuccessResponse,
|
||||
JSONRPCErrorResponse,
|
||||
JSONRPCErrorData,
|
||||
)
|
||||
from .types import CallHandlerRequest, HandshakeRequest
|
||||
from ..api.event.astr_message_event import AstrMessageEvent
|
||||
|
||||
|
||||
class StarRunner:
|
||||
def __init__(self, server: JSONRPCServer):
|
||||
self.server = server
|
||||
self._request_id_counter = 0
|
||||
self.pending_requests: dict[str, asyncio.Future] = {}
|
||||
|
||||
def _generate_request_id(self) -> str:
|
||||
self._request_id_counter += 1
|
||||
return str(self._request_id_counter)
|
||||
|
||||
async def _call_rpc(self, message: JSONRPCMessage):
|
||||
if message.id is not None:
|
||||
self.pending_requests[message.id] = asyncio.get_event_loop().create_future()
|
||||
await self.server.send_message(message)
|
||||
if message.id is not None:
|
||||
return await self.pending_requests[message.id]
|
||||
|
||||
async def _handle_messages(self, message: JSONRPCMessage):
|
||||
if isinstance(message, JSONRPCRequest):
|
||||
logger.debug(f"Received RPC request: {message.method}")
|
||||
if message.method == "handshake":
|
||||
payload = {}
|
||||
for star_name, star in star_map.items():
|
||||
payload[star_name] = star.__dict__
|
||||
handlers = []
|
||||
for handler_full_name in star.star_handler_full_names:
|
||||
handler = star_handlers_registry.get_handler_by_full_name(
|
||||
handler_full_name
|
||||
)
|
||||
if handler is None:
|
||||
continue
|
||||
handlers.append(handler.dump_model())
|
||||
payload[star_name]["handlers"] = handlers
|
||||
response = JSONRPCSuccessResponse(
|
||||
jsonrpc="2.0",
|
||||
id=message.id,
|
||||
result=payload,
|
||||
)
|
||||
await self.server.send_message(response)
|
||||
elif message.method == "call_handler":
|
||||
params = CallHandlerRequest.Params.model_validate(message.params)
|
||||
handler_full_name = params.handler_full_name
|
||||
event_model = params.event
|
||||
args = params.args
|
||||
event = event_model.to_event()
|
||||
logger.debug(f"Parsed event: {event}")
|
||||
|
||||
handler = star_handlers_registry.get_handler_by_full_name(
|
||||
handler_full_name
|
||||
)
|
||||
logger.debug(f"Invoking handler: {handler_full_name} with args: {args}")
|
||||
if handler is None:
|
||||
response = JSONRPCErrorResponse(
|
||||
jsonrpc="2.0",
|
||||
id=message.id,
|
||||
error=JSONRPCErrorData(
|
||||
code=-32601,
|
||||
message=f"Handler not found: {handler_full_name}",
|
||||
),
|
||||
)
|
||||
await self.server.send_message(response)
|
||||
else:
|
||||
try:
|
||||
ready_to_call = handler.handler(event, **args)
|
||||
notification = JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
method="handler_stream_start",
|
||||
params={
|
||||
"id": message.id,
|
||||
"handler_full_name": handler_full_name,
|
||||
},
|
||||
)
|
||||
await self.server.send_message(notification)
|
||||
if inspect.iscoroutine(ready_to_call):
|
||||
result = await ready_to_call
|
||||
notification = JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
method="handler_stream_update",
|
||||
params={
|
||||
"id": message.id,
|
||||
"handler_full_name": handler_full_name,
|
||||
"data": result,
|
||||
},
|
||||
)
|
||||
await self.server.send_message(notification)
|
||||
elif inspect.isasyncgen(ready_to_call):
|
||||
try:
|
||||
async for ret in ready_to_call:
|
||||
# Send intermediate results as notifications
|
||||
notification = JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
method="handler_stream_update",
|
||||
params={
|
||||
"id": message.id,
|
||||
"handler_full_name": handler_full_name,
|
||||
"data": ret,
|
||||
},
|
||||
)
|
||||
await self.server.send_message(notification)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error during async generator of handler {handler_full_name}: {e}"
|
||||
)
|
||||
except Exception as e:
|
||||
response = JSONRPCErrorResponse(
|
||||
jsonrpc="2.0",
|
||||
id=message.id,
|
||||
error=JSONRPCErrorData(
|
||||
code=-32000,
|
||||
message=str(e),
|
||||
),
|
||||
)
|
||||
finally:
|
||||
notification = JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
method="handler_stream_end",
|
||||
params={
|
||||
"id": message.id,
|
||||
"handler_full_name": handler_full_name,
|
||||
},
|
||||
)
|
||||
await self.server.send_message(notification)
|
||||
elif isinstance(message, (JSONRPCSuccessResponse, JSONRPCErrorResponse)):
|
||||
if message.id in self.pending_requests:
|
||||
future = self.pending_requests.pop(message.id)
|
||||
if not future.done():
|
||||
future.set_result(message)
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
await self.stop()
|
||||
|
||||
async def run(self):
|
||||
self.server.set_message_handler(handler=self._handle_messages)
|
||||
await self.server.start()
|
||||
|
||||
async def stop(self):
|
||||
await self.server.stop()
|
||||
14
src/astrbot_sdk/runtime/stars/filter/__init__.py
Normal file
14
src/astrbot_sdk/runtime/stars/filter/__init__.py
Normal file
@@ -0,0 +1,14 @@
|
||||
import abc
|
||||
|
||||
from ....api.basic.astrbot_config import AstrBotConfig
|
||||
from ....api.event import AstrMessageEvent
|
||||
|
||||
|
||||
class HandlerFilter(abc.ABC):
|
||||
@abc.abstractmethod
|
||||
def filter(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool:
|
||||
"""是否应当被过滤"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
__all__ = ["AstrBotConfig", "AstrMessageEvent", "HandlerFilter"]
|
||||
218
src/astrbot_sdk/runtime/stars/filter/command.py
Executable file
218
src/astrbot_sdk/runtime/stars/filter/command.py
Executable file
@@ -0,0 +1,218 @@
|
||||
import inspect
|
||||
import re
|
||||
import types
|
||||
import typing
|
||||
from typing import Any
|
||||
|
||||
from ....api.basic.astrbot_config import AstrBotConfig
|
||||
from ....api.event import AstrMessageEvent
|
||||
from ...stars.registry import StarHandlerMetadata
|
||||
from . import HandlerFilter
|
||||
from .custom_filter import CustomFilter
|
||||
|
||||
|
||||
class GreedyStr(str):
|
||||
"""标记指令完成其他参数接收后的所有剩余文本。"""
|
||||
|
||||
|
||||
def unwrap_optional(annotation) -> tuple:
|
||||
"""去掉 Optional[T] / Union[T, None] / T|None,返回 T"""
|
||||
args = typing.get_args(annotation)
|
||||
non_none_args = [a for a in args if a is not type(None)]
|
||||
if len(non_none_args) == 1:
|
||||
return (non_none_args[0],)
|
||||
if len(non_none_args) > 1:
|
||||
return tuple(non_none_args)
|
||||
return ()
|
||||
|
||||
|
||||
# 标准指令受到 wake_prefix 的制约。
|
||||
class CommandFilter(HandlerFilter):
|
||||
"""标准指令过滤器"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
command_name: str,
|
||||
alias: set | None = None,
|
||||
handler_md: StarHandlerMetadata | None = None,
|
||||
parent_command_names: list[str] | None = None,
|
||||
):
|
||||
self.command_name = command_name
|
||||
self.alias = alias if alias else set()
|
||||
self.parent_command_names = (
|
||||
parent_command_names if parent_command_names is not None else [""]
|
||||
)
|
||||
if handler_md:
|
||||
self.init_handler_md(handler_md)
|
||||
self.custom_filter_list: list[CustomFilter] = []
|
||||
|
||||
# Cache for complete command names list
|
||||
self._cmpl_cmd_names: list | None = None
|
||||
|
||||
def print_types(self):
|
||||
parts = []
|
||||
for k, v in self.handler_params.items():
|
||||
if isinstance(v, type):
|
||||
parts.append(f"{k}({v.__name__}),")
|
||||
elif isinstance(v, types.UnionType) or typing.get_origin(v) is typing.Union:
|
||||
parts.append(f"{k}({v}),")
|
||||
else:
|
||||
parts.append(f"{k}({type(v).__name__})={v},")
|
||||
result = "".join(parts).rstrip(",")
|
||||
return result
|
||||
|
||||
def init_handler_md(self, handle_md: StarHandlerMetadata):
|
||||
self.handler_md = handle_md
|
||||
signature = inspect.signature(self.handler_md.handler)
|
||||
self.handler_params = {} # 参数名 -> 参数类型,如果有默认值则为默认值
|
||||
idx = 0
|
||||
for k, v in signature.parameters.items():
|
||||
if idx < 2:
|
||||
# 忽略前两个参数,即 self 和 event
|
||||
idx += 1
|
||||
continue
|
||||
if v.default == inspect.Parameter.empty:
|
||||
self.handler_params[k] = v.annotation
|
||||
else:
|
||||
self.handler_params[k] = v.default
|
||||
|
||||
def get_handler_md(self) -> StarHandlerMetadata:
|
||||
return self.handler_md
|
||||
|
||||
def add_custom_filter(self, custom_filter: CustomFilter):
|
||||
self.custom_filter_list.append(custom_filter)
|
||||
|
||||
def custom_filter_ok(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool:
|
||||
for custom_filter in self.custom_filter_list:
|
||||
if not custom_filter.filter(event, cfg):
|
||||
return False
|
||||
return True
|
||||
|
||||
def validate_and_convert_params(
|
||||
self,
|
||||
params: list[Any],
|
||||
param_type: dict[str, type],
|
||||
) -> dict[str, Any]:
|
||||
"""将参数列表 params 根据 param_type 转换为参数字典。"""
|
||||
result = {}
|
||||
param_items = list(param_type.items())
|
||||
for i, (param_name, param_type_or_default_val) in enumerate(param_items):
|
||||
is_greedy = param_type_or_default_val is GreedyStr
|
||||
|
||||
if is_greedy:
|
||||
# GreedyStr 必须是最后一个参数
|
||||
if i != len(param_items) - 1:
|
||||
raise ValueError(
|
||||
f"参数 '{param_name}' (GreedyStr) 必须是最后一个参数。",
|
||||
)
|
||||
|
||||
# 将剩余的所有部分合并成一个字符串
|
||||
remaining_params = params[i:]
|
||||
result[param_name] = " ".join(remaining_params)
|
||||
break
|
||||
# 没有 GreedyStr 的情况
|
||||
if i >= len(params):
|
||||
if (
|
||||
isinstance(param_type_or_default_val, (type, types.UnionType))
|
||||
or typing.get_origin(param_type_or_default_val) is typing.Union
|
||||
or param_type_or_default_val is inspect.Parameter.empty
|
||||
):
|
||||
# 是类型
|
||||
raise ValueError(
|
||||
f"必要参数缺失。该指令完整参数: {self.print_types()}",
|
||||
)
|
||||
# 是默认值
|
||||
result[param_name] = param_type_or_default_val
|
||||
else:
|
||||
# 尝试强制转换
|
||||
try:
|
||||
if param_type_or_default_val is None:
|
||||
if params[i].isdigit():
|
||||
result[param_name] = int(params[i])
|
||||
else:
|
||||
result[param_name] = params[i]
|
||||
elif isinstance(param_type_or_default_val, str):
|
||||
# 如果 param_type_or_default_val 是字符串,直接赋值
|
||||
result[param_name] = params[i]
|
||||
elif isinstance(param_type_or_default_val, bool):
|
||||
# 处理布尔类型
|
||||
lower_param = str(params[i]).lower()
|
||||
if lower_param in ["true", "yes", "1"]:
|
||||
result[param_name] = True
|
||||
elif lower_param in ["false", "no", "0"]:
|
||||
result[param_name] = False
|
||||
else:
|
||||
raise ValueError(
|
||||
f"参数 {param_name} 必须是布尔值(true/false, yes/no, 1/0)。",
|
||||
)
|
||||
elif isinstance(param_type_or_default_val, int):
|
||||
result[param_name] = int(params[i])
|
||||
elif isinstance(param_type_or_default_val, float):
|
||||
result[param_name] = float(params[i])
|
||||
else:
|
||||
origin = typing.get_origin(param_type_or_default_val)
|
||||
if origin in (typing.Union, types.UnionType):
|
||||
# 注解是联合类型
|
||||
# NOTE: 目前没有处理联合类型嵌套相关的注解写法
|
||||
nn_types = unwrap_optional(param_type_or_default_val)
|
||||
if len(nn_types) == 1:
|
||||
# 只有一个非 NoneType 类型
|
||||
result[param_name] = nn_types[0](params[i])
|
||||
else:
|
||||
# 没有或者有多个非 NoneType 类型,这里我们暂时直接赋值为原始值。
|
||||
# NOTE: 目前还没有做类型校验
|
||||
result[param_name] = params[i]
|
||||
else:
|
||||
result[param_name] = param_type_or_default_val(params[i])
|
||||
except ValueError:
|
||||
raise ValueError(
|
||||
f"参数 {param_name} 类型错误。完整参数: {self.print_types()}",
|
||||
)
|
||||
return result
|
||||
|
||||
def get_complete_command_names(self):
|
||||
if self._cmpl_cmd_names is not None:
|
||||
return self._cmpl_cmd_names
|
||||
self._cmpl_cmd_names = [
|
||||
f"{parent} {cmd}" if parent else cmd
|
||||
for cmd in [self.command_name] + list(self.alias)
|
||||
for parent in self.parent_command_names or [""]
|
||||
]
|
||||
return self._cmpl_cmd_names
|
||||
|
||||
def equals(self, message_str: str) -> bool:
|
||||
for full_cmd in self.get_complete_command_names():
|
||||
if message_str == full_cmd:
|
||||
return True
|
||||
return False
|
||||
|
||||
def filter(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool:
|
||||
if not event.is_at_or_wake_command:
|
||||
return False
|
||||
|
||||
if not self.custom_filter_ok(event, cfg):
|
||||
return False
|
||||
|
||||
# 检查是否以指令开头
|
||||
message_str = re.sub(r"\s+", " ", event.get_message_str().strip())
|
||||
ok = False
|
||||
for full_cmd in self.get_complete_command_names():
|
||||
if message_str.startswith(f"{full_cmd} ") or message_str == full_cmd:
|
||||
ok = True
|
||||
message_str = message_str[len(full_cmd) :].strip()
|
||||
if not ok:
|
||||
return False
|
||||
|
||||
# 分割为列表
|
||||
ls = message_str.split(" ")
|
||||
# 去除空字符串
|
||||
ls = [param for param in ls if param]
|
||||
params = {}
|
||||
try:
|
||||
params = self.validate_and_convert_params(ls, self.handler_params)
|
||||
except ValueError as e:
|
||||
raise e
|
||||
|
||||
event.set_extra("parsed_params", params)
|
||||
|
||||
return True
|
||||
133
src/astrbot_sdk/runtime/stars/filter/command_group.py
Executable file
133
src/astrbot_sdk/runtime/stars/filter/command_group.py
Executable file
@@ -0,0 +1,133 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from ....api.basic.astrbot_config import AstrBotConfig
|
||||
from ....api.event import AstrMessageEvent
|
||||
from . import HandlerFilter
|
||||
from .command import CommandFilter
|
||||
from .custom_filter import CustomFilter
|
||||
|
||||
|
||||
# 指令组受到 wake_prefix 的制约。
|
||||
class CommandGroupFilter(HandlerFilter):
|
||||
def __init__(
|
||||
self,
|
||||
group_name: str,
|
||||
alias: set | None = None,
|
||||
parent_group: CommandGroupFilter | None = None,
|
||||
):
|
||||
self.group_name = group_name
|
||||
self.alias = alias if alias else set()
|
||||
self.sub_command_filters: list[CommandFilter | CommandGroupFilter] = []
|
||||
self.custom_filter_list: list[CustomFilter] = []
|
||||
self.parent_group = parent_group
|
||||
|
||||
# Cache for complete command names list
|
||||
self._cmpl_cmd_names: list | None = None
|
||||
|
||||
def add_sub_command_filter(
|
||||
self,
|
||||
sub_command_filter: CommandFilter | CommandGroupFilter,
|
||||
):
|
||||
self.sub_command_filters.append(sub_command_filter)
|
||||
|
||||
def add_custom_filter(self, custom_filter: CustomFilter):
|
||||
self.custom_filter_list.append(custom_filter)
|
||||
|
||||
def get_complete_command_names(self) -> list[str]:
|
||||
"""遍历父节点获取完整的指令名。
|
||||
|
||||
新版本 v3.4.29 采用预编译指令,不再从指令组递归遍历子指令,因此这个方法是返回包括别名在内的整个指令名列表。
|
||||
"""
|
||||
if self._cmpl_cmd_names is not None:
|
||||
return self._cmpl_cmd_names
|
||||
|
||||
parent_cmd_names = (
|
||||
self.parent_group.get_complete_command_names() if self.parent_group else []
|
||||
)
|
||||
|
||||
if not parent_cmd_names:
|
||||
# 根节点
|
||||
return [self.group_name] + list(self.alias)
|
||||
|
||||
result = []
|
||||
candidates = [self.group_name] + list(self.alias)
|
||||
for parent_cmd_name in parent_cmd_names:
|
||||
for candidate in candidates:
|
||||
result.append(parent_cmd_name + " " + candidate)
|
||||
self._cmpl_cmd_names = result
|
||||
return result
|
||||
|
||||
# 以树的形式打印出来
|
||||
def print_cmd_tree(
|
||||
self,
|
||||
sub_command_filters: list[CommandFilter | CommandGroupFilter],
|
||||
prefix: str = "",
|
||||
event: AstrMessageEvent | None = None,
|
||||
cfg: AstrBotConfig | None = None,
|
||||
) -> str:
|
||||
parts = []
|
||||
for sub_filter in sub_command_filters:
|
||||
if isinstance(sub_filter, CommandFilter):
|
||||
custom_filter_pass = True
|
||||
if event and cfg:
|
||||
custom_filter_pass = sub_filter.custom_filter_ok(event, cfg)
|
||||
if custom_filter_pass:
|
||||
cmd_th = sub_filter.print_types()
|
||||
line = f"{prefix}├── {sub_filter.command_name}"
|
||||
if cmd_th:
|
||||
line += f" ({cmd_th})"
|
||||
else:
|
||||
line += " (无参数指令)"
|
||||
|
||||
if sub_filter.handler_md and sub_filter.handler_md.desc:
|
||||
line += f": {sub_filter.handler_md.desc}"
|
||||
|
||||
parts.append(line + "\n")
|
||||
elif isinstance(sub_filter, CommandGroupFilter):
|
||||
custom_filter_pass = True
|
||||
if event and cfg:
|
||||
custom_filter_pass = sub_filter.custom_filter_ok(event, cfg)
|
||||
if custom_filter_pass:
|
||||
parts.append(f"{prefix}├── {sub_filter.group_name}\n")
|
||||
parts.append(
|
||||
sub_filter.print_cmd_tree(
|
||||
sub_filter.sub_command_filters,
|
||||
prefix + "│ ",
|
||||
event=event,
|
||||
cfg=cfg,
|
||||
)
|
||||
)
|
||||
|
||||
return "".join(parts)
|
||||
|
||||
def custom_filter_ok(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool:
|
||||
for custom_filter in self.custom_filter_list:
|
||||
if not custom_filter.filter(event, cfg):
|
||||
return False
|
||||
return True
|
||||
|
||||
def startswith(self, message_str: str) -> bool:
|
||||
return message_str.startswith(tuple(self.get_complete_command_names()))
|
||||
|
||||
def equals(self, message_str: str) -> bool:
|
||||
return message_str in self.get_complete_command_names()
|
||||
|
||||
def filter(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool:
|
||||
if not event.is_at_or_wake_command:
|
||||
return False
|
||||
|
||||
# 判断当前指令组的自定义过滤器
|
||||
if not self.custom_filter_ok(event, cfg):
|
||||
return False
|
||||
|
||||
if self.equals(event.message_str.strip()):
|
||||
tree = (
|
||||
self.group_name
|
||||
+ "\n"
|
||||
+ self.print_cmd_tree(self.sub_command_filters, event=event, cfg=cfg)
|
||||
)
|
||||
raise ValueError(
|
||||
f"参数不足。{self.group_name} 指令组下有如下指令,请参考:\n" + tree,
|
||||
)
|
||||
|
||||
return self.startswith(event.message_str)
|
||||
61
src/astrbot_sdk/runtime/stars/filter/custom_filter.py
Normal file
61
src/astrbot_sdk/runtime/stars/filter/custom_filter.py
Normal file
@@ -0,0 +1,61 @@
|
||||
from abc import ABCMeta, abstractmethod
|
||||
|
||||
from ....api.basic.astrbot_config import AstrBotConfig
|
||||
from ....api.event import AstrMessageEvent
|
||||
from . import HandlerFilter
|
||||
|
||||
|
||||
class CustomFilterMeta(ABCMeta):
|
||||
def __and__(cls, other):
|
||||
if not issubclass(other, CustomFilter):
|
||||
raise TypeError("Operands must be subclasses of CustomFilter.")
|
||||
return CustomFilterAnd(cls(), other())
|
||||
|
||||
def __or__(cls, other):
|
||||
if not issubclass(other, CustomFilter):
|
||||
raise TypeError("Operands must be subclasses of CustomFilter.")
|
||||
return CustomFilterOr(cls(), other())
|
||||
|
||||
|
||||
class CustomFilter(HandlerFilter, metaclass=CustomFilterMeta):
|
||||
def __init__(self, raise_error: bool = True, **kwargs):
|
||||
self.raise_error = raise_error
|
||||
|
||||
@abstractmethod
|
||||
def filter(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool:
|
||||
"""一个用于重写的自定义Filter"""
|
||||
raise NotImplementedError
|
||||
|
||||
def __or__(self, other):
|
||||
return CustomFilterOr(self, other)
|
||||
|
||||
def __and__(self, other):
|
||||
return CustomFilterAnd(self, other)
|
||||
|
||||
|
||||
class CustomFilterOr(CustomFilter):
|
||||
def __init__(self, filter1: CustomFilter, filter2: CustomFilter):
|
||||
super().__init__()
|
||||
if not isinstance(filter1, (CustomFilter, CustomFilterAnd, CustomFilterOr)):
|
||||
raise ValueError(
|
||||
"CustomFilter lass can only operate with other CustomFilter.",
|
||||
)
|
||||
self.filter1 = filter1
|
||||
self.filter2 = filter2
|
||||
|
||||
def filter(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool:
|
||||
return self.filter1.filter(event, cfg) or self.filter2.filter(event, cfg)
|
||||
|
||||
|
||||
class CustomFilterAnd(CustomFilter):
|
||||
def __init__(self, filter1: CustomFilter, filter2: CustomFilter):
|
||||
super().__init__()
|
||||
if not isinstance(filter1, (CustomFilter, CustomFilterAnd, CustomFilterOr)):
|
||||
raise ValueError(
|
||||
"CustomFilter lass can only operate with other CustomFilter.",
|
||||
)
|
||||
self.filter1 = filter1
|
||||
self.filter2 = filter2
|
||||
|
||||
def filter(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool:
|
||||
return self.filter1.filter(event, cfg) and self.filter2.filter(event, cfg)
|
||||
33
src/astrbot_sdk/runtime/stars/filter/event_message_type.py
Normal file
33
src/astrbot_sdk/runtime/stars/filter/event_message_type.py
Normal file
@@ -0,0 +1,33 @@
|
||||
import enum
|
||||
|
||||
from ....api.basic.astrbot_config import AstrBotConfig
|
||||
from ....api.event import AstrMessageEvent
|
||||
from ....api.event.message_type import MessageType
|
||||
|
||||
from . import HandlerFilter
|
||||
|
||||
|
||||
class EventMessageType(enum.Flag):
|
||||
GROUP_MESSAGE = enum.auto()
|
||||
PRIVATE_MESSAGE = enum.auto()
|
||||
OTHER_MESSAGE = enum.auto()
|
||||
ALL = GROUP_MESSAGE | PRIVATE_MESSAGE | OTHER_MESSAGE
|
||||
|
||||
|
||||
MESSAGE_TYPE_2_EVENT_MESSAGE_TYPE = {
|
||||
MessageType.GROUP_MESSAGE: EventMessageType.GROUP_MESSAGE,
|
||||
MessageType.FRIEND_MESSAGE: EventMessageType.PRIVATE_MESSAGE,
|
||||
MessageType.OTHER_MESSAGE: EventMessageType.OTHER_MESSAGE,
|
||||
}
|
||||
|
||||
|
||||
class EventMessageTypeFilter(HandlerFilter):
|
||||
def __init__(self, event_message_type: EventMessageType):
|
||||
self.event_message_type = event_message_type
|
||||
|
||||
def filter(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool:
|
||||
message_type = event.get_message_type()
|
||||
if message_type in MESSAGE_TYPE_2_EVENT_MESSAGE_TYPE:
|
||||
event_message_type = MESSAGE_TYPE_2_EVENT_MESSAGE_TYPE[message_type]
|
||||
return bool(event_message_type & self.event_message_type)
|
||||
return False
|
||||
29
src/astrbot_sdk/runtime/stars/filter/permission.py
Normal file
29
src/astrbot_sdk/runtime/stars/filter/permission.py
Normal file
@@ -0,0 +1,29 @@
|
||||
import enum
|
||||
|
||||
from ....api.basic.astrbot_config import AstrBotConfig
|
||||
from ....api.event import AstrMessageEvent
|
||||
|
||||
from . import HandlerFilter
|
||||
|
||||
|
||||
class PermissionType(enum.Flag):
|
||||
"""权限类型。当选择 MEMBER,ADMIN 也可以通过。"""
|
||||
|
||||
ADMIN = enum.auto()
|
||||
MEMBER = enum.auto()
|
||||
|
||||
|
||||
class PermissionTypeFilter(HandlerFilter):
|
||||
def __init__(self, permission_type: PermissionType, raise_error: bool = True):
|
||||
self.permission_type = permission_type
|
||||
self.raise_error = raise_error
|
||||
|
||||
def filter(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool:
|
||||
"""过滤器"""
|
||||
if self.permission_type == PermissionType.ADMIN:
|
||||
if not event.is_admin():
|
||||
# event.stop_event()
|
||||
# raise ValueError(f"您 (ID: {event.get_sender_id()}) 没有权限操作管理员指令。")
|
||||
return False
|
||||
|
||||
return True
|
||||
@@ -0,0 +1,71 @@
|
||||
import enum
|
||||
|
||||
from ....api.basic.astrbot_config import AstrBotConfig
|
||||
from ....api.event import AstrMessageEvent
|
||||
|
||||
from . import HandlerFilter
|
||||
|
||||
|
||||
class PlatformAdapterType(enum.Flag):
|
||||
AIOCQHTTP = enum.auto()
|
||||
QQOFFICIAL = enum.auto()
|
||||
TELEGRAM = enum.auto()
|
||||
WECOM = enum.auto()
|
||||
LARK = enum.auto()
|
||||
WECHATPADPRO = enum.auto()
|
||||
DINGTALK = enum.auto()
|
||||
DISCORD = enum.auto()
|
||||
SLACK = enum.auto()
|
||||
KOOK = enum.auto()
|
||||
VOCECHAT = enum.auto()
|
||||
WEIXIN_OFFICIAL_ACCOUNT = enum.auto()
|
||||
SATORI = enum.auto()
|
||||
MISSKEY = enum.auto()
|
||||
ALL = (
|
||||
AIOCQHTTP
|
||||
| QQOFFICIAL
|
||||
| TELEGRAM
|
||||
| WECOM
|
||||
| LARK
|
||||
| WECHATPADPRO
|
||||
| DINGTALK
|
||||
| DISCORD
|
||||
| SLACK
|
||||
| KOOK
|
||||
| VOCECHAT
|
||||
| WEIXIN_OFFICIAL_ACCOUNT
|
||||
| SATORI
|
||||
| MISSKEY
|
||||
)
|
||||
|
||||
|
||||
ADAPTER_NAME_2_TYPE = {
|
||||
"aiocqhttp": PlatformAdapterType.AIOCQHTTP,
|
||||
"qq_official": PlatformAdapterType.QQOFFICIAL,
|
||||
"telegram": PlatformAdapterType.TELEGRAM,
|
||||
"wecom": PlatformAdapterType.WECOM,
|
||||
"lark": PlatformAdapterType.LARK,
|
||||
"dingtalk": PlatformAdapterType.DINGTALK,
|
||||
"discord": PlatformAdapterType.DISCORD,
|
||||
"slack": PlatformAdapterType.SLACK,
|
||||
"kook": PlatformAdapterType.KOOK,
|
||||
"wechatpadpro": PlatformAdapterType.WECHATPADPRO,
|
||||
"vocechat": PlatformAdapterType.VOCECHAT,
|
||||
"weixin_official_account": PlatformAdapterType.WEIXIN_OFFICIAL_ACCOUNT,
|
||||
"satori": PlatformAdapterType.SATORI,
|
||||
"misskey": PlatformAdapterType.MISSKEY,
|
||||
}
|
||||
|
||||
|
||||
class PlatformAdapterTypeFilter(HandlerFilter):
|
||||
def __init__(self, platform_adapter_type_or_str: PlatformAdapterType | str):
|
||||
if isinstance(platform_adapter_type_or_str, str):
|
||||
self.platform_type = ADAPTER_NAME_2_TYPE.get(platform_adapter_type_or_str)
|
||||
else:
|
||||
self.platform_type = platform_adapter_type_or_str
|
||||
|
||||
def filter(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool:
|
||||
adapter_name = event.get_platform_name()
|
||||
if adapter_name in ADAPTER_NAME_2_TYPE and self.platform_type is not None:
|
||||
return bool(ADAPTER_NAME_2_TYPE[adapter_name] & self.platform_type)
|
||||
return False
|
||||
18
src/astrbot_sdk/runtime/stars/filter/regex.py
Normal file
18
src/astrbot_sdk/runtime/stars/filter/regex.py
Normal file
@@ -0,0 +1,18 @@
|
||||
import re
|
||||
|
||||
from ....api.basic.astrbot_config import AstrBotConfig
|
||||
from ....api.event import AstrMessageEvent
|
||||
|
||||
from . import HandlerFilter
|
||||
|
||||
|
||||
# 正则表达式过滤器不会受到 wake_prefix 的制约。
|
||||
class RegexFilter(HandlerFilter):
|
||||
"""正则表达式过滤器"""
|
||||
|
||||
def __init__(self, regex: str):
|
||||
self.regex_str = regex
|
||||
self.regex = re.compile(regex)
|
||||
|
||||
def filter(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool:
|
||||
return bool(self.regex.match(event.get_message_str().strip()))
|
||||
0
src/astrbot_sdk/runtime/stars/legacy.py
Normal file
0
src/astrbot_sdk/runtime/stars/legacy.py
Normal file
594
src/astrbot_sdk/runtime/stars/new.py
Normal file
594
src/astrbot_sdk/runtime/stars/new.py
Normal file
@@ -0,0 +1,594 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from ...api.event.astr_message_event import AstrMessageEvent, AstrMessageEventModel
|
||||
from ...api.star.star import StarMetadata
|
||||
from ..stars.registry import EventType, StarHandlerMetadata
|
||||
from ..rpc.jsonrpc import (
|
||||
JSONRPCErrorData,
|
||||
JSONRPCErrorResponse,
|
||||
JSONRPCMessage,
|
||||
JSONRPCRequest,
|
||||
JSONRPCSuccessResponse,
|
||||
)
|
||||
from ..types import CallHandlerRequest, HandshakeRequest
|
||||
from ..rpc.client import JSONRPCClient
|
||||
from ..rpc.client.stdio import StdioClient
|
||||
from ..rpc.client.websocket import WebSocketClient
|
||||
from .virtual import VirtualStar
|
||||
|
||||
|
||||
class NewStar(VirtualStar):
|
||||
"""NewStar implementation for isolated plugin runtime.
|
||||
|
||||
NewStar runs plugins in separate processes and communicates via JSON-RPC.
|
||||
This provides better isolation, security, and compatibility.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: JSONRPCClient,
|
||||
) -> None:
|
||||
"""Initialize a NewStar instance.
|
||||
|
||||
Args:
|
||||
client: JSON-RPC client for communication
|
||||
"""
|
||||
self._client = client
|
||||
self._metadata: dict[str, StarMetadata] = {}
|
||||
self._handlers: list[StarHandlerMetadata] = []
|
||||
self._request_id_counter = 0
|
||||
self._pending_requests: dict[
|
||||
str, asyncio.Future[dict] | asyncio.Queue[dict]
|
||||
] = {}
|
||||
self._active = False
|
||||
|
||||
# Set up message handler
|
||||
self._client.set_message_handler(self._handle_message)
|
||||
|
||||
def _generate_request_id(self) -> str:
|
||||
"""Generate a unique request ID."""
|
||||
self._request_id_counter += 1
|
||||
return f"req-{self._request_id_counter}"
|
||||
|
||||
async def _handle_message(self, message: JSONRPCMessage) -> None:
|
||||
"""Handle incoming JSON-RPC messages from the plugin.
|
||||
|
||||
Args:
|
||||
message: The received JSON-RPC message
|
||||
"""
|
||||
if isinstance(message, JSONRPCSuccessResponse) or isinstance(
|
||||
message,
|
||||
JSONRPCErrorResponse,
|
||||
):
|
||||
# This is a response to one of our requests
|
||||
request_id = message.id
|
||||
if request_id and request_id in self._pending_requests:
|
||||
pending = self._pending_requests[request_id]
|
||||
|
||||
# Check if it's a Future or Queue
|
||||
if isinstance(pending, asyncio.Future):
|
||||
self._pending_requests.pop(request_id)
|
||||
if isinstance(message, JSONRPCSuccessResponse):
|
||||
if not pending.done():
|
||||
pending.set_result(message.result)
|
||||
else:
|
||||
if not pending.done():
|
||||
pending.set_exception(
|
||||
RuntimeError(
|
||||
f"RPC Error {message.error.code}: {message.error.message}",
|
||||
),
|
||||
)
|
||||
elif isinstance(pending, asyncio.Queue):
|
||||
if isinstance(message, JSONRPCSuccessResponse):
|
||||
logger.debug(
|
||||
f"Streaming handler {request_id} completed successfully"
|
||||
)
|
||||
else:
|
||||
logger.error(
|
||||
f"Streaming handler {request_id} failed: {message.error.message}"
|
||||
)
|
||||
# Put error marker in queue
|
||||
await pending.put(
|
||||
{"_error": True, "message": message.error.message}
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"Received response for unknown request ID: {request_id}"
|
||||
)
|
||||
|
||||
elif isinstance(message, JSONRPCRequest):
|
||||
# Handle notifications from plugin (streaming events or method calls)
|
||||
if message.method in [
|
||||
"handler_stream_start",
|
||||
"handler_stream_update",
|
||||
"handler_stream_end",
|
||||
]:
|
||||
await self._handle_stream_notification(message)
|
||||
else:
|
||||
# Plugin is calling a method on the core
|
||||
asyncio.create_task(self._handle_plugin_request(message))
|
||||
|
||||
async def _handle_plugin_request(self, request: JSONRPCRequest) -> None:
|
||||
"""Handle a JSON-RPC request from the plugin (plugin calling core methods).
|
||||
|
||||
Args:
|
||||
request: The JSON-RPC request from the plugin
|
||||
"""
|
||||
result: dict = {}
|
||||
try:
|
||||
# Handle core methods that plugins might call
|
||||
# For now, we'll implement basic methods
|
||||
method = request.method
|
||||
params = request.params
|
||||
|
||||
if method == "core.log":
|
||||
# Plugin wants to log something
|
||||
level = params.get("level", "info")
|
||||
message = params.get("message", "")
|
||||
getattr(logger, level.lower())(f"[Plugin] {message}")
|
||||
result = {"success": True}
|
||||
|
||||
elif method == "core.send_message":
|
||||
# Plugin wants to send a message
|
||||
# This would integrate with the platform adapter
|
||||
logger.info(f"Plugin requested to send message: {params}")
|
||||
result = {"success": True, "message_id": "mock-msg-id"}
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown method: {method}")
|
||||
|
||||
# Send success response
|
||||
response = JSONRPCSuccessResponse(
|
||||
jsonrpc="2.0",
|
||||
id=request.id,
|
||||
result=result,
|
||||
)
|
||||
await self._client.send_message(response)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error handling plugin request: {e}")
|
||||
# Send error response
|
||||
error_response = JSONRPCErrorResponse(
|
||||
jsonrpc="2.0",
|
||||
id=request.id,
|
||||
error=JSONRPCErrorData(
|
||||
code=-32603,
|
||||
message=str(e),
|
||||
),
|
||||
)
|
||||
await self._client.send_message(error_response)
|
||||
|
||||
async def _handle_stream_notification(self, notification: JSONRPCRequest) -> None:
|
||||
"""Handle streaming notifications from the plugin.
|
||||
|
||||
Args:
|
||||
notification: The streaming notification (handler_stream_start/update/end)
|
||||
"""
|
||||
params = notification.params
|
||||
request_id = params.get("id")
|
||||
|
||||
if not request_id or request_id not in self._pending_requests:
|
||||
logger.warning(
|
||||
f"Received stream notification for unknown request ID: {request_id}"
|
||||
)
|
||||
return
|
||||
|
||||
pending = self._pending_requests.get(request_id)
|
||||
if not isinstance(pending, asyncio.Queue):
|
||||
logger.warning(f"Request {request_id} is not a streaming request")
|
||||
return
|
||||
|
||||
if notification.method == "handler_stream_start":
|
||||
logger.debug(
|
||||
f"Stream started for handler {params.get('handler_full_name')}"
|
||||
)
|
||||
# Optionally put a start marker in the queue
|
||||
# await pending.put({"_stream_start": True})
|
||||
|
||||
elif notification.method == "handler_stream_update":
|
||||
# Put the streamed data into the queue
|
||||
data = params.get("data")
|
||||
logger.debug(f"Stream update for request {request_id}: {data}")
|
||||
if data is not None:
|
||||
await pending.put(data)
|
||||
|
||||
elif notification.method == "handler_stream_end":
|
||||
# Mark the end of the stream
|
||||
logger.debug(f"Stream ended for handler {params.get('handler_full_name')}")
|
||||
# Put a sentinel value to indicate stream end
|
||||
await pending.put({"_stream_end": True})
|
||||
# Clean up the pending request after a short delay to allow queue to be processed
|
||||
asyncio.create_task(self._cleanup_stream_request(request_id))
|
||||
|
||||
async def _cleanup_stream_request(
|
||||
self, request_id: str, delay: float = 1.0
|
||||
) -> None:
|
||||
"""Clean up a streaming request after a delay.
|
||||
|
||||
Args:
|
||||
request_id: The request ID to clean up
|
||||
delay: Delay before cleanup in seconds
|
||||
"""
|
||||
await asyncio.sleep(delay)
|
||||
if request_id in self._pending_requests:
|
||||
self._pending_requests.pop(request_id)
|
||||
logger.debug(f"Cleaned up streaming request {request_id}")
|
||||
|
||||
async def _call_rpc(self, request: JSONRPCRequest) -> dict:
|
||||
"""Call a JSON-RPC method on the plugin and wait for response.
|
||||
|
||||
Args:
|
||||
request: The JSON-RPC request to send
|
||||
|
||||
Returns:
|
||||
The result from the plugin
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the RPC call fails
|
||||
"""
|
||||
# Create a future to wait for the response
|
||||
future: asyncio.Future[dict] = asyncio.Future()
|
||||
|
||||
if request.id is not None:
|
||||
self._pending_requests[request.id] = future
|
||||
|
||||
try:
|
||||
await self._client.send_message(request)
|
||||
# Wait for response with timeout
|
||||
result = await asyncio.wait_for(future, timeout=30.0)
|
||||
return result
|
||||
except asyncio.TimeoutError:
|
||||
if request.id is not None:
|
||||
self._pending_requests.pop(request.id, None)
|
||||
raise RuntimeError(f"RPC call to {request.method} timed out")
|
||||
|
||||
async def _call_rpc_streaming(
|
||||
self,
|
||||
request: JSONRPCRequest,
|
||||
) -> asyncio.Queue[dict]:
|
||||
"""Call a JSON-RPC method on the plugin that returns a stream of results.
|
||||
|
||||
Args:
|
||||
request: The JSON-RPC request to send
|
||||
Returns:
|
||||
An asyncio.Queue that will receive streamed results
|
||||
"""
|
||||
# Create a queue to receive streamed results
|
||||
queue: asyncio.Queue[dict] = asyncio.Queue()
|
||||
|
||||
if request.id is not None:
|
||||
self._pending_requests[request.id] = queue
|
||||
|
||||
try:
|
||||
await self._client.send_message(request)
|
||||
return queue
|
||||
except Exception as e:
|
||||
if request.id is not None:
|
||||
self._pending_requests.pop(request.id, None)
|
||||
raise RuntimeError(f"RPC streaming call to {request.method} failed: {e}")
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Start the plugin process and establish connection."""
|
||||
# Start the client (which may start a subprocess for STDIO)
|
||||
await self._client.start()
|
||||
logger.info("Client started and ready for communication")
|
||||
|
||||
async def handshake(self) -> dict[str, StarMetadata]:
|
||||
"""Perform handshake to retrieve plugin metadata.
|
||||
|
||||
Returns:
|
||||
Plugin metadata including name, version, handlers, etc.
|
||||
"""
|
||||
logger.info("Performing handshake with plugin...")
|
||||
|
||||
result = await self._call_rpc(
|
||||
HandshakeRequest(
|
||||
jsonrpc="2.0", id=self._generate_request_id(), method="handshake"
|
||||
)
|
||||
)
|
||||
|
||||
print(result, result.__class__)
|
||||
|
||||
if isinstance(result, dict):
|
||||
# Parse metadata
|
||||
for star_name, star_info in result.items():
|
||||
handlers_data = star_info.pop("handlers", None)
|
||||
metadata = StarMetadata(**star_info)
|
||||
self._metadata[star_name] = metadata
|
||||
|
||||
# Get handlers
|
||||
self._handlers = []
|
||||
|
||||
for handler_data in handlers_data:
|
||||
handler_meta = StarHandlerMetadata(
|
||||
event_type=EventType(handler_data["event_type"]),
|
||||
handler_full_name=handler_data["handler_full_name"],
|
||||
handler_name=handler_data["handler_name"],
|
||||
handler_module_path=handler_data["handler_module_path"],
|
||||
handler=self._create_handler_proxy(
|
||||
handler_data["handler_full_name"]
|
||||
),
|
||||
event_filters=[],
|
||||
desc=handler_data.get("desc", ""),
|
||||
extras_configs=handler_data.get("extras_configs", {}),
|
||||
)
|
||||
self._handlers.append(handler_meta)
|
||||
|
||||
logger.info(
|
||||
f"Handshake complete: {len(self._metadata)} stars loaded, {self._metadata.keys()}, {len(self._handlers)} handlers registered."
|
||||
)
|
||||
logger.info(f"Registered {len(self._handlers)} handlers")
|
||||
|
||||
return self._metadata
|
||||
raise RuntimeError("Handshake failed: Invalid response from plugin")
|
||||
|
||||
def _create_handler_proxy(self, handler_full_name: str):
|
||||
"""Create a proxy function that calls the handler via RPC.
|
||||
|
||||
Args:
|
||||
handler_full_name: The full name of the handler
|
||||
|
||||
Returns:
|
||||
An async function that proxies calls to the remote handler.
|
||||
The function may return a direct result or an async generator for streaming.
|
||||
"""
|
||||
|
||||
async def handler_proxy(event: AstrMessageEvent, **kwargs):
|
||||
"""Proxy function for remote handler invocation.
|
||||
|
||||
Returns either a direct result or an async generator for streaming handlers.
|
||||
"""
|
||||
request_id = self._generate_request_id()
|
||||
request = CallHandlerRequest(
|
||||
jsonrpc="2.0",
|
||||
id=request_id,
|
||||
method="call_handler",
|
||||
params=CallHandlerRequest.Params(
|
||||
handler_full_name=handler_full_name,
|
||||
event=AstrMessageEventModel.from_event(event),
|
||||
args=kwargs,
|
||||
),
|
||||
)
|
||||
|
||||
# Create a queue for potential streaming response
|
||||
queue: asyncio.Queue[dict] = asyncio.Queue()
|
||||
self._pending_requests[request_id] = queue
|
||||
|
||||
try:
|
||||
# Send the request
|
||||
await self._client.send_message(request)
|
||||
|
||||
# Wait for the first response or stream notification
|
||||
try:
|
||||
# Set a timeout for the first response
|
||||
first_response = await asyncio.wait_for(queue.get(), timeout=30.0)
|
||||
|
||||
# Check what type of response we got
|
||||
if isinstance(first_response, dict):
|
||||
# Check for stream end (empty stream case)
|
||||
if first_response.get("_stream_end"):
|
||||
# Empty stream, return None
|
||||
self._pending_requests.pop(request_id, None)
|
||||
return None
|
||||
|
||||
# Check for error
|
||||
if first_response.get("_error"):
|
||||
self._pending_requests.pop(request_id, None)
|
||||
raise RuntimeError(
|
||||
first_response.get("message", "Unknown error")
|
||||
)
|
||||
|
||||
# Check if this is streaming data or a final result
|
||||
# We peek at the queue to see if more data is coming
|
||||
# If the queue is empty after a short wait, it's a final result
|
||||
try:
|
||||
# Try to get another item with a very short timeout
|
||||
second_response = await asyncio.wait_for(
|
||||
queue.get(), timeout=0.1
|
||||
)
|
||||
# We got a second item, so this is streaming
|
||||
# Create and return the generator
|
||||
return self._create_stream_generator(
|
||||
queue, first_response, second_response
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
# No second item, this might be a final result
|
||||
# But we should check if stream_end arrives shortly
|
||||
try:
|
||||
stream_end = await asyncio.wait_for(
|
||||
queue.get(), timeout=0.5
|
||||
)
|
||||
if isinstance(stream_end, dict) and stream_end.get(
|
||||
"_stream_end"
|
||||
):
|
||||
# This was a single-item stream
|
||||
self._pending_requests.pop(request_id, None)
|
||||
return self._deserialize_result(first_response)
|
||||
else:
|
||||
# More data arrived, it's streaming
|
||||
return self._create_stream_generator(
|
||||
queue, first_response, stream_end
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
# Truly a final result (non-streaming)
|
||||
self._pending_requests.pop(request_id, None)
|
||||
return self._deserialize_result(first_response)
|
||||
else:
|
||||
# Unexpected response type
|
||||
self._pending_requests.pop(request_id, None)
|
||||
return self._deserialize_result(first_response)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
# Timeout waiting for response
|
||||
self._pending_requests.pop(request_id, None)
|
||||
raise RuntimeError(f"RPC call to {handler_full_name} timed out")
|
||||
|
||||
except Exception:
|
||||
# Clean up on error
|
||||
self._pending_requests.pop(request_id, None)
|
||||
raise
|
||||
|
||||
return handler_proxy
|
||||
|
||||
async def _create_stream_generator(
|
||||
self, queue: asyncio.Queue[dict], *initial_items: dict
|
||||
):
|
||||
"""Create an async generator that yields items from the stream queue.
|
||||
|
||||
Args:
|
||||
queue: The queue containing stream items
|
||||
initial_items: Initial items that were already retrieved from the queue
|
||||
|
||||
Yields:
|
||||
Items from the stream
|
||||
"""
|
||||
# Yield any initial items
|
||||
for item in initial_items:
|
||||
if not (isinstance(item, dict) and item.get("_stream_end")):
|
||||
yield self._deserialize_result(item)
|
||||
|
||||
# Continue yielding items from the queue
|
||||
while True:
|
||||
try:
|
||||
item = await queue.get()
|
||||
|
||||
# Check for end marker
|
||||
if isinstance(item, dict) and item.get("_stream_end"):
|
||||
break
|
||||
|
||||
# Check for error marker
|
||||
if isinstance(item, dict) and item.get("_error"):
|
||||
raise RuntimeError(item.get("message", "Stream error"))
|
||||
|
||||
# Yield the item
|
||||
yield self._deserialize_result(item)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
# Generator was cancelled, stop iteration
|
||||
logger.debug("Stream generator cancelled")
|
||||
break
|
||||
|
||||
def _deserialize_result(self, result: Any) -> Any:
|
||||
"""Deserialize result from JSON-RPC response.
|
||||
|
||||
Args:
|
||||
result: The result from the plugin
|
||||
|
||||
Returns:
|
||||
Deserialized result object
|
||||
"""
|
||||
# For now, return as-is
|
||||
# In practice, you might want to reconstruct MessageEventResult etc.
|
||||
return result
|
||||
|
||||
def get_triggered_handlers(
|
||||
self, event: AstrMessageEvent
|
||||
) -> list[StarHandlerMetadata]:
|
||||
"""Get the list of handlers that should be triggered for this event.
|
||||
|
||||
Args:
|
||||
event: The message event
|
||||
|
||||
Returns:
|
||||
List of handler metadata that should handle this event
|
||||
"""
|
||||
# For AdapterMessageEvent, return relevant handlers
|
||||
# This is cached locally, no RPC needed
|
||||
triggered = []
|
||||
|
||||
for handler in self._handlers:
|
||||
if handler.event_type == EventType.AdapterMessageEvent:
|
||||
# In practice, you'd check filters here
|
||||
triggered.append(handler)
|
||||
|
||||
return triggered
|
||||
|
||||
async def call_handler(
|
||||
self,
|
||||
handler: StarHandlerMetadata,
|
||||
event: AstrMessageEvent,
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
"""Call a specific handler in the plugin.
|
||||
|
||||
Args:
|
||||
handler: The handler metadata
|
||||
event: The message event
|
||||
*args: Additional positional arguments
|
||||
**kwargs: Additional keyword arguments
|
||||
|
||||
Returns:
|
||||
Result from the handler
|
||||
"""
|
||||
logger.debug(f"Calling handler: {handler.handler_name}")
|
||||
|
||||
# Call the handler proxy
|
||||
result = await handler.handler(event, *args, **kwargs)
|
||||
return result
|
||||
|
||||
|
||||
class NewStdioStar(NewStar):
|
||||
"""NewStar implementation using STDIO communication.
|
||||
|
||||
This class automatically starts the plugin subprocess and manages its lifecycle.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
plugin_dir: str,
|
||||
python_executable: str = "python",
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize a STDIO-based NewStar.
|
||||
|
||||
Args:
|
||||
plugin_dir: Path to the plugin directory
|
||||
python_executable: Python executable to use (defaults to 'python')
|
||||
main_script: Main script filename (defaults to 'main.py')
|
||||
"""
|
||||
# Construct the command to start the plugin
|
||||
if not os.path.exists(plugin_dir):
|
||||
raise FileNotFoundError(f"Plugin directory not found: {plugin_dir}")
|
||||
|
||||
command = [python_executable, "-m", "astrbot_sdk", "run", "--stdio"]
|
||||
|
||||
# Create StdioClient with subprocess management
|
||||
client = StdioClient(command=command, cwd=plugin_dir)
|
||||
super().__init__(client)
|
||||
|
||||
|
||||
class NewWebSocketStar(NewStar):
|
||||
"""NewStar implementation using WebSocket communication.
|
||||
|
||||
Note: WebSocket-based stars do not start the plugin process.
|
||||
The plugin should be started externally and connect to the specified WebSocket URL.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: str,
|
||||
heartbeat: float = 30.0,
|
||||
reconnect_interval: float = 5.0,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize a WebSocket-based NewStar.
|
||||
|
||||
Args:
|
||||
url: WebSocket server URL that the plugin will connect to
|
||||
heartbeat: Heartbeat interval in seconds
|
||||
reconnect_interval: Interval between reconnection attempts in seconds
|
||||
"""
|
||||
client = WebSocketClient(
|
||||
url=url, heartbeat=heartbeat, reconnect_interval=reconnect_interval
|
||||
)
|
||||
super().__init__(client)
|
||||
self._url = url
|
||||
self._heartbeat = heartbeat
|
||||
self._reconnect_interval = reconnect_interval
|
||||
181
src/astrbot_sdk/runtime/stars/registry/__init__.py
Normal file
181
src/astrbot_sdk/runtime/stars/registry/__init__.py
Normal file
@@ -0,0 +1,181 @@
|
||||
from __future__ import annotations
|
||||
import enum
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Generic, TypeVar
|
||||
from ..filter import HandlerFilter
|
||||
from ....api.star.star import StarMetadata
|
||||
from ....api.star.context import Context as BaseContext
|
||||
|
||||
T = TypeVar("T", bound="StarHandlerMetadata")
|
||||
|
||||
|
||||
class EventType(enum.Enum):
|
||||
"""表示一个 AstrBot 内部事件的类型。如适配器消息事件、LLM 请求事件、发送消息前的事件等
|
||||
|
||||
用于对 Handler 的职能分组。
|
||||
"""
|
||||
|
||||
OnAstrBotLoadedEvent = enum.auto()
|
||||
"""AstrBot 加载完成"""
|
||||
OnPlatformLoadedEvent = enum.auto()
|
||||
"""平台适配器加载完成"""
|
||||
AdapterMessageEvent = enum.auto()
|
||||
"""收到适配器消息事件"""
|
||||
OnLLMRequestEvent = enum.auto()
|
||||
"""LLM 请求前"""
|
||||
OnLLMResponseEvent = enum.auto()
|
||||
"""LLM 响应后"""
|
||||
OnDecoratingResultEvent = enum.auto()
|
||||
"""发送消息前"""
|
||||
OnCallingFuncToolEvent = enum.auto()
|
||||
"""调用函数工具前"""
|
||||
OnAfterMessageSentEvent = enum.auto()
|
||||
"""发送消息后"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class StarHandlerMetadata:
|
||||
"""描述一个 Star 所注册的某一个 Handler。"""
|
||||
|
||||
event_type: EventType
|
||||
"""Handler 的事件类型"""
|
||||
|
||||
handler_full_name: str
|
||||
'''格式为 f"{handler.__module__}_{handler.__name__}"'''
|
||||
|
||||
handler_name: str
|
||||
"""Handler 的名字,也就是方法名"""
|
||||
|
||||
handler_module_path: str
|
||||
"""Handler 所在的模块路径。"""
|
||||
|
||||
handler: Callable[..., Awaitable[Any]]
|
||||
"""Handler 的函数对象,应当是一个异步函数"""
|
||||
|
||||
event_filters: list[HandlerFilter]
|
||||
"""一个适配器消息事件过滤器,用于描述这个 Handler 能够处理、应该处理的适配器消息事件"""
|
||||
|
||||
desc: str = ""
|
||||
"""Handler 的描述信息"""
|
||||
|
||||
extras_configs: dict = field(default_factory=dict)
|
||||
"""插件注册的一些其他的信息, 如 priority 等"""
|
||||
|
||||
def __lt__(self, other: StarHandlerMetadata):
|
||||
"""定义小于运算符以支持优先队列"""
|
||||
return self.extras_configs.get("priority", 0) < other.extras_configs.get(
|
||||
"priority",
|
||||
0,
|
||||
)
|
||||
|
||||
def dump_model(self) -> dict[str, Any]:
|
||||
"""将 Handler 的元数据转换为字典形式,便于序列化。"""
|
||||
p = self.__dict__.copy()
|
||||
p.pop("handler")
|
||||
p.pop("event_filters")
|
||||
return p
|
||||
|
||||
class StarHandlerRegistry(Generic[T]):
|
||||
def __init__(self):
|
||||
self.star_handlers_map: dict[str, StarHandlerMetadata] = {}
|
||||
self._handlers: list[StarHandlerMetadata] = []
|
||||
|
||||
def append(self, handler: StarHandlerMetadata):
|
||||
"""添加一个 Handler,并保持按优先级有序"""
|
||||
if "priority" not in handler.extras_configs:
|
||||
handler.extras_configs["priority"] = 0
|
||||
|
||||
self.star_handlers_map[handler.handler_full_name] = handler
|
||||
self._handlers.append(handler)
|
||||
self._handlers.sort(key=lambda h: -h.extras_configs["priority"])
|
||||
|
||||
def _print_handlers(self):
|
||||
for handler in self._handlers:
|
||||
print(handler.handler_full_name)
|
||||
|
||||
def get_handlers_by_event_type(
|
||||
self,
|
||||
event_type: EventType,
|
||||
only_activated=True,
|
||||
plugins_name: list[str] | None = None,
|
||||
) -> list[StarHandlerMetadata]:
|
||||
handlers = []
|
||||
for handler in self._handlers:
|
||||
# 过滤事件类型
|
||||
if handler.event_type != event_type:
|
||||
continue
|
||||
# 过滤启用状态
|
||||
if only_activated:
|
||||
plugin = star_map.get(handler.handler_module_path)
|
||||
if not (plugin and plugin.activated):
|
||||
continue
|
||||
# 过滤插件白名单
|
||||
if plugins_name is not None and plugins_name != ["*"]:
|
||||
plugin = star_map.get(handler.handler_module_path)
|
||||
if not plugin:
|
||||
continue
|
||||
if (
|
||||
plugin.name not in plugins_name
|
||||
and event_type
|
||||
not in (
|
||||
EventType.OnAstrBotLoadedEvent,
|
||||
EventType.OnPlatformLoadedEvent,
|
||||
)
|
||||
and not plugin.reserved
|
||||
):
|
||||
continue
|
||||
handlers.append(handler)
|
||||
return handlers
|
||||
|
||||
def get_handler_by_full_name(self, full_name: str) -> StarHandlerMetadata | None:
|
||||
return self.star_handlers_map.get(full_name, None)
|
||||
|
||||
def get_handlers_by_module_name(
|
||||
self,
|
||||
module_name: str,
|
||||
) -> list[StarHandlerMetadata]:
|
||||
return [
|
||||
handler
|
||||
for handler in self._handlers
|
||||
if handler.handler_module_path == module_name
|
||||
]
|
||||
|
||||
def clear(self):
|
||||
self.star_handlers_map.clear()
|
||||
self._handlers.clear()
|
||||
|
||||
def remove(self, handler: StarHandlerMetadata):
|
||||
self.star_handlers_map.pop(handler.handler_full_name, None)
|
||||
self._handlers = [h for h in self._handlers if h != handler]
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._handlers)
|
||||
|
||||
def __len__(self):
|
||||
return len(self._handlers)
|
||||
|
||||
|
||||
class Star:
|
||||
"""所有插件的基类。每一个插件都应当继承自这个类,并实现相应的方法。"""
|
||||
|
||||
def __init__(self, context: BaseContext):
|
||||
self.context = context
|
||||
|
||||
def __init_subclass__(cls, **kwargs):
|
||||
super().__init_subclass__(**kwargs)
|
||||
if not star_map.get(cls.__module__):
|
||||
metadata = StarMetadata(
|
||||
# star_cls_type=cls,
|
||||
module_path=cls.__module__,
|
||||
)
|
||||
star_map[cls.__module__] = metadata
|
||||
star_registry.append(metadata)
|
||||
else:
|
||||
# star_map[cls.__module__].star_cls_type = cls
|
||||
star_map[cls.__module__].module_path = cls.__module__
|
||||
|
||||
|
||||
star_handlers_registry = StarHandlerRegistry() # type: ignore
|
||||
star_map: dict[str, StarMetadata] = {}
|
||||
star_registry: list[StarMetadata] = []
|
||||
514
src/astrbot_sdk/runtime/stars/registry/register.py
Normal file
514
src/astrbot_sdk/runtime/stars/registry/register.py
Normal file
@@ -0,0 +1,514 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
# import docstring_parser
|
||||
|
||||
from loguru import logger
|
||||
# from astrbot.core.agent.agent import Agent
|
||||
# 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.provider.register import llm_tools
|
||||
|
||||
from ..filter.command import CommandFilter
|
||||
from ..filter.command_group import CommandGroupFilter
|
||||
from ..filter.custom_filter import CustomFilterAnd, CustomFilterOr
|
||||
from ..filter.event_message_type import EventMessageType, EventMessageTypeFilter
|
||||
from ..filter.permission import PermissionType, PermissionTypeFilter
|
||||
from ..filter.platform_adapter_type import (
|
||||
PlatformAdapterType,
|
||||
PlatformAdapterTypeFilter,
|
||||
)
|
||||
from ..filter.regex import RegexFilter
|
||||
from ..registry import star_handlers_registry, StarHandlerMetadata, EventType
|
||||
|
||||
def get_handler_full_name(awaitable: Callable[..., Awaitable[Any]]) -> str:
|
||||
"""获取 Handler 的全名"""
|
||||
return f"{awaitable.__module__}:{awaitable.__qualname__}"
|
||||
|
||||
|
||||
def get_handler_or_create(
|
||||
handler: Callable[..., Awaitable[Any]],
|
||||
event_type: EventType,
|
||||
dont_add=False,
|
||||
**kwargs,
|
||||
) -> StarHandlerMetadata:
|
||||
"""获取 Handler 或者创建一个新的 Handler"""
|
||||
handler_full_name = get_handler_full_name(handler)
|
||||
md = star_handlers_registry.get_handler_by_full_name(handler_full_name)
|
||||
if md:
|
||||
return md
|
||||
md = StarHandlerMetadata(
|
||||
event_type=event_type,
|
||||
handler_full_name=handler_full_name,
|
||||
handler_name=handler.__name__,
|
||||
handler_module_path=handler.__module__,
|
||||
handler=handler,
|
||||
event_filters=[],
|
||||
)
|
||||
|
||||
# 插件handler的附加额外信息
|
||||
if handler.__doc__:
|
||||
md.desc = handler.__doc__.strip()
|
||||
if "desc" in kwargs:
|
||||
md.desc = kwargs["desc"]
|
||||
del kwargs["desc"]
|
||||
md.extras_configs = kwargs
|
||||
|
||||
if not dont_add:
|
||||
star_handlers_registry.append(md)
|
||||
return md
|
||||
|
||||
|
||||
def register_command(
|
||||
command_name: str | None = None,
|
||||
sub_command: str | None = None,
|
||||
alias: set | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""注册一个 Command."""
|
||||
new_command = None
|
||||
add_to_event_filters = False
|
||||
if isinstance(command_name, RegisteringCommandable):
|
||||
# 子指令
|
||||
if sub_command is not None:
|
||||
parent_command_names = (
|
||||
command_name.parent_group.get_complete_command_names()
|
||||
)
|
||||
new_command = CommandFilter(
|
||||
sub_command,
|
||||
alias,
|
||||
None,
|
||||
parent_command_names=parent_command_names,
|
||||
)
|
||||
command_name.parent_group.add_sub_command_filter(new_command)
|
||||
else:
|
||||
logger.warning(
|
||||
f"注册指令{command_name} 的子指令时未提供 sub_command 参数。",
|
||||
)
|
||||
# 裸指令
|
||||
elif command_name is None:
|
||||
logger.warning("注册裸指令时未提供 command_name 参数。")
|
||||
else:
|
||||
new_command = CommandFilter(command_name, alias, None)
|
||||
add_to_event_filters = True
|
||||
|
||||
def decorator(awaitable):
|
||||
if not add_to_event_filters:
|
||||
kwargs["sub_command"] = (
|
||||
True # 打一个标记,表示这是一个子指令,再 wakingstage 阶段这个 handler 将会直接被跳过(其父指令会接管)
|
||||
)
|
||||
handler_md = get_handler_or_create(
|
||||
awaitable,
|
||||
EventType.AdapterMessageEvent,
|
||||
**kwargs,
|
||||
)
|
||||
if new_command:
|
||||
new_command.init_handler_md(handler_md)
|
||||
handler_md.event_filters.append(new_command)
|
||||
return awaitable
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def register_custom_filter(custom_type_filter, *args, **kwargs):
|
||||
"""注册一个自定义的 CustomFilter
|
||||
|
||||
Args:
|
||||
custom_type_filter: 在裸指令时为CustomFilter对象
|
||||
在指令组时为父指令的RegisteringCommandable对象,即self或者command_group的返回
|
||||
raise_error: 如果没有权限,是否抛出错误到消息平台,并且停止事件传播。默认为 True
|
||||
|
||||
"""
|
||||
add_to_event_filters = False
|
||||
raise_error = True
|
||||
|
||||
# 判断是否是指令组,指令组则添加到指令组的CommandGroupFilter对象中在waking_check的时候一起判断
|
||||
if isinstance(custom_type_filter, RegisteringCommandable):
|
||||
# 子指令, 此时函数为RegisteringCommandable对象的方法,首位参数为RegisteringCommandable对象的self。
|
||||
parent_register_commandable = custom_type_filter
|
||||
custom_filter = args[0]
|
||||
if len(args) > 1:
|
||||
raise_error = args[1]
|
||||
else:
|
||||
# 裸指令
|
||||
add_to_event_filters = True
|
||||
custom_filter = custom_type_filter
|
||||
if args:
|
||||
raise_error = args[0]
|
||||
|
||||
if not isinstance(custom_filter, (CustomFilterAnd, CustomFilterOr)):
|
||||
custom_filter = custom_filter(raise_error)
|
||||
|
||||
def decorator(awaitable):
|
||||
# 裸指令,子指令与指令组的区分,指令组会因为标记跳过wake。
|
||||
if (
|
||||
not add_to_event_filters and isinstance(awaitable, RegisteringCommandable)
|
||||
) or (add_to_event_filters and isinstance(awaitable, RegisteringCommandable)):
|
||||
# 指令组 与 根指令组,添加到本层的grouphandle中一起判断
|
||||
awaitable.parent_group.add_custom_filter(custom_filter)
|
||||
else:
|
||||
handler_md = get_handler_or_create(
|
||||
awaitable,
|
||||
EventType.AdapterMessageEvent,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if not add_to_event_filters and not isinstance(
|
||||
awaitable,
|
||||
RegisteringCommandable,
|
||||
):
|
||||
# 底层子指令
|
||||
handle_full_name = get_handler_full_name(awaitable)
|
||||
for (
|
||||
sub_handle
|
||||
) in parent_register_commandable.parent_group.sub_command_filters:
|
||||
# 所有符合fullname一致的子指令handle添加自定义过滤器。
|
||||
# 不确定是否会有多个子指令有一样的fullname,比如一个方法添加多个command装饰器?
|
||||
sub_handle_md = sub_handle.get_handler_md()
|
||||
if (
|
||||
sub_handle_md
|
||||
and sub_handle_md.handler_full_name == handle_full_name
|
||||
):
|
||||
sub_handle.add_custom_filter(custom_filter)
|
||||
|
||||
else:
|
||||
# 裸指令
|
||||
handler_md = get_handler_or_create(
|
||||
awaitable,
|
||||
EventType.AdapterMessageEvent,
|
||||
**kwargs,
|
||||
)
|
||||
handler_md.event_filters.append(custom_filter)
|
||||
|
||||
return awaitable
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def register_command_group(
|
||||
command_group_name: str | None = None,
|
||||
sub_command: str | None = None,
|
||||
alias: set | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""注册一个 CommandGroup"""
|
||||
new_group = None
|
||||
if isinstance(command_group_name, RegisteringCommandable):
|
||||
# 子指令组
|
||||
if sub_command is None:
|
||||
logger.warning(f"{command_group_name} 指令组的子指令组 sub_command 未指定")
|
||||
else:
|
||||
new_group = CommandGroupFilter(
|
||||
sub_command,
|
||||
alias,
|
||||
parent_group=command_group_name.parent_group,
|
||||
)
|
||||
command_group_name.parent_group.add_sub_command_filter(new_group)
|
||||
# 根指令组
|
||||
elif command_group_name is None:
|
||||
logger.warning("根指令组的名称未指定")
|
||||
else:
|
||||
new_group = CommandGroupFilter(command_group_name, alias)
|
||||
|
||||
def decorator(obj):
|
||||
if new_group:
|
||||
handler_md = get_handler_or_create(
|
||||
obj,
|
||||
EventType.AdapterMessageEvent,
|
||||
**kwargs,
|
||||
)
|
||||
handler_md.event_filters.append(new_group)
|
||||
|
||||
return RegisteringCommandable(new_group)
|
||||
raise ValueError("注册指令组失败。")
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
class RegisteringCommandable:
|
||||
"""用于指令组级联注册"""
|
||||
|
||||
group: Callable[..., Callable[..., RegisteringCommandable]] = register_command_group
|
||||
command: Callable[..., Callable[..., None]] = register_command
|
||||
custom_filter: Callable[..., Callable[..., None]] = register_custom_filter
|
||||
|
||||
def __init__(self, parent_group: CommandGroupFilter):
|
||||
self.parent_group = parent_group
|
||||
|
||||
|
||||
def register_event_message_type(event_message_type: EventMessageType, **kwargs):
|
||||
"""注册一个 EventMessageType"""
|
||||
|
||||
def decorator(awaitable):
|
||||
handler_md = get_handler_or_create(
|
||||
awaitable,
|
||||
EventType.AdapterMessageEvent,
|
||||
**kwargs,
|
||||
)
|
||||
handler_md.event_filters.append(EventMessageTypeFilter(event_message_type))
|
||||
return awaitable
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def register_platform_adapter_type(
|
||||
platform_adapter_type: PlatformAdapterType,
|
||||
**kwargs,
|
||||
):
|
||||
"""注册一个 PlatformAdapterType"""
|
||||
|
||||
def decorator(awaitable):
|
||||
handler_md = get_handler_or_create(awaitable, EventType.AdapterMessageEvent)
|
||||
handler_md.event_filters.append(
|
||||
PlatformAdapterTypeFilter(platform_adapter_type),
|
||||
)
|
||||
return awaitable
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def register_regex(regex: str, **kwargs):
|
||||
"""注册一个 Regex"""
|
||||
|
||||
def decorator(awaitable):
|
||||
handler_md = get_handler_or_create(
|
||||
awaitable,
|
||||
EventType.AdapterMessageEvent,
|
||||
**kwargs,
|
||||
)
|
||||
handler_md.event_filters.append(RegexFilter(regex))
|
||||
return awaitable
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def register_permission_type(permission_type: PermissionType, raise_error: bool = True):
|
||||
"""注册一个 PermissionType
|
||||
|
||||
Args:
|
||||
permission_type: PermissionType
|
||||
raise_error: 如果没有权限,是否抛出错误到消息平台,并且停止事件传播。默认为 True
|
||||
|
||||
"""
|
||||
|
||||
def decorator(awaitable):
|
||||
handler_md = get_handler_or_create(awaitable, EventType.AdapterMessageEvent)
|
||||
handler_md.event_filters.append(
|
||||
PermissionTypeFilter(permission_type, raise_error),
|
||||
)
|
||||
return awaitable
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def register_on_astrbot_loaded(**kwargs):
|
||||
"""当 AstrBot 加载完成时"""
|
||||
|
||||
def decorator(awaitable):
|
||||
_ = get_handler_or_create(awaitable, EventType.OnAstrBotLoadedEvent, **kwargs)
|
||||
return awaitable
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def register_on_platform_loaded(**kwargs):
|
||||
"""当平台加载完成时"""
|
||||
|
||||
def decorator(awaitable):
|
||||
_ = get_handler_or_create(awaitable, EventType.OnPlatformLoadedEvent, **kwargs)
|
||||
return awaitable
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def register_on_llm_request(**kwargs):
|
||||
"""当有 LLM 请求时的事件
|
||||
|
||||
Examples:
|
||||
```py
|
||||
from astrbot.api.provider import ProviderRequest
|
||||
|
||||
@on_llm_request()
|
||||
async def test(self, event: AstrMessageEvent, request: ProviderRequest) -> None:
|
||||
request.system_prompt += "你是一个猫娘..."
|
||||
```
|
||||
|
||||
请务必接收两个参数:event, request
|
||||
|
||||
"""
|
||||
|
||||
def decorator(awaitable):
|
||||
_ = get_handler_or_create(awaitable, EventType.OnLLMRequestEvent, **kwargs)
|
||||
return awaitable
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def register_on_llm_response(**kwargs):
|
||||
"""当有 LLM 请求后的事件
|
||||
|
||||
Examples:
|
||||
```py
|
||||
from astrbot.api.provider import LLMResponse
|
||||
|
||||
@on_llm_response()
|
||||
async def test(self, event: AstrMessageEvent, response: LLMResponse) -> None:
|
||||
...
|
||||
```
|
||||
|
||||
请务必接收两个参数:event, request
|
||||
|
||||
"""
|
||||
|
||||
def decorator(awaitable):
|
||||
_ = get_handler_or_create(awaitable, EventType.OnLLMResponseEvent, **kwargs)
|
||||
return awaitable
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
# def register_llm_tool(name: str | None = None, **kwargs):
|
||||
# """为函数调用(function-calling / tools-use)添加工具。
|
||||
|
||||
# 请务必按照以下格式编写一个工具(包括函数注释,AstrBot 会尝试解析该函数注释)
|
||||
|
||||
# ```
|
||||
# @llm_tool(name="get_weather") # 如果 name 不填,将使用函数名
|
||||
# async def get_weather(event: AstrMessageEvent, location: str):
|
||||
# \'\'\'获取天气信息。
|
||||
|
||||
# Args:
|
||||
# location(string): 地点
|
||||
# \'\'\'
|
||||
# # 处理逻辑
|
||||
# ```
|
||||
|
||||
# 可接受的参数类型有:string, number, object, array, boolean。
|
||||
|
||||
# 返回值:
|
||||
# - 返回 str:结果会被加入下一次 LLM 请求的 prompt 中,用于让 LLM 总结工具返回的结果
|
||||
# - 返回 None:结果不会被加入下一次 LLM 请求的 prompt 中。
|
||||
|
||||
# 可以使用 yield 发送消息、终止事件。
|
||||
|
||||
# 发送消息:请参考文档。
|
||||
|
||||
# 终止事件:
|
||||
# ```
|
||||
# event.stop_event()
|
||||
# yield
|
||||
# ```
|
||||
|
||||
# """
|
||||
# name_ = name
|
||||
# registering_agent = None
|
||||
# if kwargs.get("registering_agent"):
|
||||
# registering_agent = kwargs["registering_agent"]
|
||||
|
||||
# def decorator(awaitable: Callable[..., Awaitable[Any]]):
|
||||
# llm_tool_name = name_ if name_ else awaitable.__name__
|
||||
# func_doc = awaitable.__doc__ or ""
|
||||
# docstring = docstring_parser.parse(func_doc)
|
||||
# args = []
|
||||
# for arg in docstring.params:
|
||||
# args.append(
|
||||
# {
|
||||
# "type": arg.type_name,
|
||||
# "name": arg.arg_name,
|
||||
# "description": arg.description,
|
||||
# },
|
||||
# )
|
||||
# # print(llm_tool_name, registering_agent)
|
||||
# if not registering_agent:
|
||||
# doc_desc = docstring.description.strip() if docstring.description else ""
|
||||
# md = get_handler_or_create(awaitable, EventType.OnCallingFuncToolEvent)
|
||||
# llm_tools.add_func(llm_tool_name, args, doc_desc, md.handler)
|
||||
# else:
|
||||
# assert isinstance(registering_agent, RegisteringAgent)
|
||||
# # print(f"Registering tool {llm_tool_name} for agent", registering_agent._agent.name)
|
||||
# if registering_agent._agent.tools is None:
|
||||
# registering_agent._agent.tools = []
|
||||
|
||||
# desc = docstring.description.strip() if docstring.description else ""
|
||||
# tool = llm_tools.spec_to_func(llm_tool_name, args, desc, awaitable)
|
||||
# registering_agent._agent.tools.append(tool)
|
||||
|
||||
# return awaitable
|
||||
|
||||
# return decorator
|
||||
|
||||
|
||||
# class RegisteringAgent:
|
||||
# """用于 Agent 注册"""
|
||||
|
||||
# def llm_tool(self, *args, **kwargs):
|
||||
# kwargs["registering_agent"] = self
|
||||
# return register_llm_tool(*args, **kwargs)
|
||||
|
||||
# def __init__(self, agent: Agent[AstrAgentContext]):
|
||||
# self._agent = agent
|
||||
|
||||
|
||||
# def register_agent(
|
||||
# name: str,
|
||||
# instruction: str,
|
||||
# tools: list[str | FunctionTool] | None = None,
|
||||
# run_hooks: BaseAgentRunHooks[AstrAgentContext] | None = None,
|
||||
# ):
|
||||
# """注册一个 Agent
|
||||
|
||||
# Args:
|
||||
# name: Agent 的名称
|
||||
# instruction: Agent 的指令
|
||||
# tools: Agent 使用的工具列表
|
||||
# run_hooks: Agent 运行时的钩子函数
|
||||
|
||||
# """
|
||||
# tools_ = tools or []
|
||||
|
||||
# def decorator(awaitable: Callable[..., Awaitable[Any]]):
|
||||
# AstrAgent = Agent[AstrAgentContext]
|
||||
# agent = AstrAgent(
|
||||
# name=name,
|
||||
# instructions=instruction,
|
||||
# tools=tools_,
|
||||
# run_hooks=run_hooks or BaseAgentRunHooks[AstrAgentContext](),
|
||||
# )
|
||||
# handoff_tool = HandoffTool(agent=agent)
|
||||
# handoff_tool.handler = awaitable
|
||||
# llm_tools.func_list.append(handoff_tool)
|
||||
# return RegisteringAgent(agent)
|
||||
|
||||
# return decorator
|
||||
|
||||
|
||||
def register_on_decorating_result(**kwargs):
|
||||
"""在发送消息前的事件"""
|
||||
|
||||
def decorator(awaitable):
|
||||
_ = get_handler_or_create(
|
||||
awaitable,
|
||||
EventType.OnDecoratingResultEvent,
|
||||
**kwargs,
|
||||
)
|
||||
return awaitable
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def register_after_message_sent(**kwargs):
|
||||
"""在消息发送后的事件"""
|
||||
|
||||
def decorator(awaitable):
|
||||
_ = get_handler_or_create(
|
||||
awaitable,
|
||||
EventType.OnAfterMessageSentEvent,
|
||||
**kwargs,
|
||||
)
|
||||
return awaitable
|
||||
|
||||
return decorator
|
||||
121
src/astrbot_sdk/runtime/stars/virtual.py
Normal file
121
src/astrbot_sdk/runtime/stars/virtual.py
Normal file
@@ -0,0 +1,121 @@
|
||||
import typing as T
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from ...api.event.astr_message_event import AstrMessageEvent
|
||||
from ...api.star.star import StarMetadata
|
||||
from .registry import StarHandlerMetadata
|
||||
|
||||
|
||||
class VirtualStar(ABC):
|
||||
"""Abstract base class for virtual plugin implementations.
|
||||
|
||||
VirtualStar defines the interface for plugins that can run in isolated
|
||||
runtime environments (separate processes). It handles the complete lifecycle
|
||||
of a plugin from initialization to shutdown.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def initialize(self) -> None:
|
||||
"""Establish connection and initialize the plugin.
|
||||
|
||||
This method should:
|
||||
- Start the plugin process (if applicable)
|
||||
- Establish communication channels
|
||||
- Wait for the plugin to be ready
|
||||
|
||||
Raises:
|
||||
RuntimeError: If initialization fails
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def handshake(self) -> StarMetadata:
|
||||
"""Perform handshake to retrieve plugin metadata.
|
||||
|
||||
This method should:
|
||||
- Request plugin metadata from the plugin
|
||||
- Cache handler information locally
|
||||
- Validate the plugin's compatibility
|
||||
|
||||
Returns:
|
||||
StarMetadata: Complete plugin metadata including handlers
|
||||
|
||||
Raises:
|
||||
RuntimeError: If handshake fails or times out
|
||||
"""
|
||||
...
|
||||
|
||||
# @abstractmethod
|
||||
# async def turn_on(self) -> None:
|
||||
# """Attach and prepare resources. Only call when the plugin is not active.
|
||||
|
||||
# This method should:
|
||||
# - Activate the plugin
|
||||
# - Initialize any runtime resources
|
||||
# - Prepare the plugin to handle events
|
||||
|
||||
# Raises:
|
||||
# RuntimeError: If activation fails
|
||||
# """
|
||||
# ...
|
||||
|
||||
# @abstractmethod
|
||||
# async def turn_off(self) -> None:
|
||||
# """Detach and clean up resources. Make the plugin inactive.
|
||||
|
||||
# This method should:
|
||||
# - Deactivate the plugin
|
||||
# - Release runtime resources
|
||||
# - Keep the process running but idle
|
||||
|
||||
# Raises:
|
||||
# RuntimeError: If deactivation fails
|
||||
# """
|
||||
# ...
|
||||
|
||||
@abstractmethod
|
||||
def get_triggered_handlers(
|
||||
self,
|
||||
event: AstrMessageEvent,
|
||||
) -> list[StarHandlerMetadata]:
|
||||
"""Get the list of handlers that should be triggered for this event.
|
||||
|
||||
This method uses cached handler metadata to determine which handlers
|
||||
should handle the given event. No RPC calls should be made here.
|
||||
|
||||
Args:
|
||||
event: The message event to check
|
||||
|
||||
Returns:
|
||||
List of handler metadata that match the event
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def call_handler(
|
||||
self,
|
||||
handler: StarHandlerMetadata,
|
||||
event: AstrMessageEvent,
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> T.Any:
|
||||
"""Call a registered handler in the plugin.
|
||||
|
||||
This method should:
|
||||
- Serialize the event and arguments
|
||||
- Call the handler via RPC
|
||||
- Wait for and return the result
|
||||
|
||||
Args:
|
||||
handler: The handler metadata
|
||||
event: The message event
|
||||
*args: Additional positional arguments
|
||||
**kwargs: Additional keyword arguments
|
||||
|
||||
Returns:
|
||||
The result from the handler
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the handler call fails or times out
|
||||
"""
|
||||
...
|
||||
37
src/astrbot_sdk/runtime/start_client.py
Normal file
37
src/astrbot_sdk/runtime/start_client.py
Normal file
@@ -0,0 +1,37 @@
|
||||
from .galaxy import Galaxy
|
||||
from ..api.event import AstrMessageEvent
|
||||
from ..api.event.astrbot_message import AstrBotMessage, MessageMember
|
||||
from ..api.platform.platform_metadata import PlatformMetadata
|
||||
from ..api.event.message_type import MessageType
|
||||
|
||||
async def amain():
|
||||
galaxy = Galaxy()
|
||||
star = await galaxy.connect_to_websocket_star(
|
||||
"hello",
|
||||
{
|
||||
"url": "ws://127.0.0.1:8765",
|
||||
},
|
||||
)
|
||||
print("Connected to websocket star 'hello'")
|
||||
md = await star.handshake()
|
||||
print(f"Handshake metadata: {md}")
|
||||
|
||||
abm = AstrBotMessage()
|
||||
abm.type = MessageType.FRIEND_MESSAGE
|
||||
abm.self_id = "astrbot_123"
|
||||
abm.session_id = "test_session"
|
||||
abm.message_id = "msg_001"
|
||||
abm.message_str = "hello"
|
||||
abm.sender = MessageMember(user_id="user_123", nickname="User123") # Simplified for this example
|
||||
abm.group = None
|
||||
abm.message = []
|
||||
abm.raw_message = {}
|
||||
event = AstrMessageEvent(
|
||||
message_str=abm.message_str,
|
||||
message_obj=abm,
|
||||
platform_meta=PlatformMetadata(
|
||||
name="fake", description="Fake Platform", id="fake_1"
|
||||
),
|
||||
session_id="test_session",
|
||||
)
|
||||
await star.call_handler(star._handlers[0], event)
|
||||
34
src/astrbot_sdk/runtime/start_server.py
Normal file
34
src/astrbot_sdk/runtime/start_server.py
Normal file
@@ -0,0 +1,34 @@
|
||||
import asyncio
|
||||
import signal
|
||||
from .rpc.server import WebSocketServer
|
||||
from .star_runner import StarRunner
|
||||
from .star_manager import StarManager
|
||||
from ..runtime.api.context import Context
|
||||
from ..runtime.api.conversation_mgr import ConversationManager
|
||||
|
||||
|
||||
async def amain():
|
||||
server = WebSocketServer(port=8765)
|
||||
conversation_manager = ConversationManager()
|
||||
context = Context(conversation_manager=conversation_manager)
|
||||
runner = StarRunner(server)
|
||||
context._inject_rpc_handlers(runner=runner)
|
||||
star_manager = StarManager(context=context)
|
||||
star_manager.discover_star()
|
||||
await runner.run()
|
||||
|
||||
# 设置停止事件
|
||||
stop_event = asyncio.Event()
|
||||
|
||||
# 注册信号处理器
|
||||
loop = asyncio.get_running_loop()
|
||||
for sig in (signal.SIGTERM, signal.SIGINT):
|
||||
loop.add_signal_handler(sig, stop_event.set)
|
||||
|
||||
print("Server is running. Press Ctrl+C to stop.")
|
||||
|
||||
try:
|
||||
await stop_event.wait()
|
||||
finally:
|
||||
print("Shutting down...")
|
||||
await server.stop()
|
||||
76
src/astrbot_sdk/runtime/types.py
Normal file
76
src/astrbot_sdk/runtime/types.py
Normal file
@@ -0,0 +1,76 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
from .rpc.jsonrpc import JSONRPCRequest
|
||||
from typing import Any, Literal, Type
|
||||
from ..api.event.astr_message_event import AstrMessageEvent, AstrMessageEventModel
|
||||
|
||||
|
||||
# class StarType(enum.Enum):
|
||||
# LEGACY = "legacy"
|
||||
# STDIO = "stdio"
|
||||
# WEBSOCKET = "websocket"
|
||||
|
||||
|
||||
# class StarURI(BaseModel):
|
||||
# star_type: StarType
|
||||
# namespace: str
|
||||
# plugin_name: str
|
||||
|
||||
# def __str__(self):
|
||||
# return f"astrbot://{self.star_type.value}/{self.namespace}/{self.plugin_name}"
|
||||
|
||||
# @classmethod
|
||||
# def from_str(cls, uri_str: str) -> StarURI:
|
||||
# """Parse a StarURI from a string."""
|
||||
# try:
|
||||
# prefix, rest = uri_str.split("://", 1)
|
||||
# star_type_str, namespace, plugin_name = rest.split("/", 2)
|
||||
# star_type = StarType(star_type_str)
|
||||
# return cls(
|
||||
# star_type=star_type,
|
||||
# namespace=namespace,
|
||||
# plugin_name=plugin_name,
|
||||
# )
|
||||
# except Exception as e:
|
||||
# raise ValueError(f"Invalid StarURI format: {uri_str}") from e
|
||||
|
||||
# def is_new_star(self) -> bool:
|
||||
# """Determine if the Star is a new-style Star (stdio or websocket)."""
|
||||
# return self.star_type in {StarType.STDIO, StarType.WEBSOCKET}
|
||||
|
||||
|
||||
class HandshakeRequest(JSONRPCRequest):
|
||||
class Params(BaseModel):
|
||||
pass
|
||||
|
||||
method: Literal["handshake"]
|
||||
params: Params = Field(default_factory=Params)
|
||||
|
||||
|
||||
class CallHandlerRequest(JSONRPCRequest):
|
||||
class Params(BaseModel):
|
||||
handler_full_name: str
|
||||
event: AstrMessageEventModel
|
||||
args: dict[str, Any] = {}
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def validate_event_data(cls: Type[CallHandlerRequest.Params], data: Any) -> Any:
|
||||
if isinstance(data, dict):
|
||||
event_data = data.get("event")
|
||||
if isinstance(event_data, dict):
|
||||
data["event"] = AstrMessageEventModel.model_validate(event_data)
|
||||
return data
|
||||
|
||||
method: Literal["call_handler"]
|
||||
params: Params | dict = Field(default_factory=dict)
|
||||
|
||||
|
||||
class CallContextFunctionRequest(JSONRPCRequest):
|
||||
class Params(BaseModel):
|
||||
name: str
|
||||
args: dict[str, Any] = {}
|
||||
|
||||
method: Literal["call_context_function"]
|
||||
params: Params | dict = Field(default_factory=dict)
|
||||
81
src/astrbot_sdk/util.py
Normal file
81
src/astrbot_sdk/util.py
Normal file
@@ -0,0 +1,81 @@
|
||||
import aiohttp
|
||||
import certifi
|
||||
import ssl
|
||||
import time
|
||||
from loguru import logger
|
||||
|
||||
|
||||
async def download_file(url: str, path: str, show_progress: bool = False):
|
||||
"""从指定 url 下载文件到指定路径 path"""
|
||||
try:
|
||||
ssl_context = ssl.create_default_context(
|
||||
cafile=certifi.where(),
|
||||
) # 使用 certifi 提供的 CA 证书
|
||||
connector = aiohttp.TCPConnector(ssl=ssl_context)
|
||||
async with aiohttp.ClientSession(
|
||||
trust_env=True,
|
||||
connector=connector,
|
||||
) as session:
|
||||
async with session.get(url, timeout=1800) as resp:
|
||||
if resp.status != 200:
|
||||
raise Exception(f"下载文件失败: {resp.status}")
|
||||
total_size = int(resp.headers.get("content-length", 0))
|
||||
downloaded_size = 0
|
||||
start_time = time.time()
|
||||
if show_progress:
|
||||
print(f"文件大小: {total_size / 1024:.2f} KB | 文件地址: {url}")
|
||||
with open(path, "wb") as f:
|
||||
while True:
|
||||
chunk = await resp.content.read(8192)
|
||||
if not chunk:
|
||||
break
|
||||
f.write(chunk)
|
||||
downloaded_size += len(chunk)
|
||||
if show_progress:
|
||||
elapsed_time = (
|
||||
time.time() - start_time
|
||||
if time.time() - start_time > 0
|
||||
else 1
|
||||
)
|
||||
speed = downloaded_size / 1024 / elapsed_time # KB/s
|
||||
print(
|
||||
f"\r下载进度: {downloaded_size / total_size:.2%} 速度: {speed:.2f} KB/s",
|
||||
end="",
|
||||
)
|
||||
except (aiohttp.ClientConnectorSSLError, aiohttp.ClientConnectorCertificateError):
|
||||
# 关闭SSL验证(仅在证书验证失败时作为fallback)
|
||||
logger.warning(
|
||||
"SSL 证书验证失败,已关闭 SSL 验证(不安全,仅用于临时下载)。请检查目标服务器的证书配置。"
|
||||
)
|
||||
logger.warning(
|
||||
f"SSL certificate verification failed for {url}. "
|
||||
"Falling back to unverified connection (CERT_NONE). "
|
||||
"This is insecure and exposes the application to man-in-the-middle attacks. "
|
||||
"Please investigate certificate issues with the remote server."
|
||||
)
|
||||
ssl_context = ssl.create_default_context()
|
||||
ssl_context.check_hostname = False
|
||||
ssl_context.verify_mode = ssl.CERT_NONE
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(url, ssl=ssl_context, timeout=120) as resp:
|
||||
total_size = int(resp.headers.get("content-length", 0))
|
||||
downloaded_size = 0
|
||||
start_time = time.time()
|
||||
if show_progress:
|
||||
print(f"文件大小: {total_size / 1024:.2f} KB | 文件地址: {url}")
|
||||
with open(path, "wb") as f:
|
||||
while True:
|
||||
chunk = await resp.content.read(8192)
|
||||
if not chunk:
|
||||
break
|
||||
f.write(chunk)
|
||||
downloaded_size += len(chunk)
|
||||
if show_progress:
|
||||
elapsed_time = time.time() - start_time
|
||||
speed = downloaded_size / 1024 / elapsed_time # KB/s
|
||||
print(
|
||||
f"\r下载进度: {downloaded_size / total_size:.2%} 速度: {speed:.2f} KB/s",
|
||||
end="",
|
||||
)
|
||||
if show_progress:
|
||||
print()
|
||||
12
src/handlers/commands/hello.py
Normal file
12
src/handlers/commands/hello.py
Normal file
@@ -0,0 +1,12 @@
|
||||
from astrbot_sdk.api.components.command import CommandComponent
|
||||
from astrbot_sdk.api.event import AstrMessageEvent, filter
|
||||
from astrbot_sdk.api.star.context import Context
|
||||
|
||||
|
||||
class HelloCommand(CommandComponent):
|
||||
def __init__(self, context: Context):
|
||||
self.context = context
|
||||
|
||||
@filter.command("hello")
|
||||
async def hello(self, event: AstrMessageEvent):
|
||||
yield event.plain_result("Hello, Astrbot!")
|
||||
12
src/handlers/listeners/echo.py
Normal file
12
src/handlers/listeners/echo.py
Normal file
@@ -0,0 +1,12 @@
|
||||
from astrbot_sdk.api.components. import ListenerComponent
|
||||
from astrbot.api.v1.event import AstrMessageEvent, filter
|
||||
from astrbot.api.v1.context import Context
|
||||
|
||||
|
||||
class EchoListener(ListenerComponent):
|
||||
def __init__(self, context: Context):
|
||||
super().__init__(context)
|
||||
|
||||
@filter.platform_adapter_type(filter.PlatformAdapterType.ALL)
|
||||
async def on_message(self, event: AstrMessageEvent):
|
||||
yield event.plain_result("Hello, Astrbot!")
|
||||
31
src/handlers/tools/hello.py
Normal file
31
src/handlers/tools/hello.py
Normal file
@@ -0,0 +1,31 @@
|
||||
from mcp.types import CallToolResult
|
||||
from pydantic import Field
|
||||
from pydantic.dataclasses import dataclass
|
||||
|
||||
from astrbot.api.v1.components.agent.tool import FunctionTool
|
||||
from astrbot.api.v1.components.agent import ContextWrapper, AstrAgentContext
|
||||
|
||||
|
||||
@dataclass
|
||||
class HelloWorldTool(FunctionTool):
|
||||
name: str = "hello_world" # 工具名称
|
||||
description: str = "Say hello to the world." # 工具描述
|
||||
parameters: dict = Field(
|
||||
default_factory=lambda: {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"greeting": {
|
||||
"type": "string",
|
||||
"description": "The greeting message.",
|
||||
},
|
||||
},
|
||||
"required": ["greeting"],
|
||||
}
|
||||
) # 工具参数定义,见 OpenAI 官网或 https://json-schema.org/understanding-json-schema/
|
||||
|
||||
async def call(
|
||||
self, context: ContextWrapper[AstrAgentContext], **kwargs
|
||||
) -> str | CallToolResult:
|
||||
# event 在 context.context.event 中可用
|
||||
greeting = kwargs.get("greeting", "Hello")
|
||||
return f"{greeting}, World!" # 也支持 mcp.types.CallToolResult 类型
|
||||
22
src/plugin.yaml
Normal file
22
src/plugin.yaml
Normal file
@@ -0,0 +1,22 @@
|
||||
_schema_version: 2
|
||||
name: astrbot_plugin_helloworld
|
||||
display_name: HelloWorld 插件
|
||||
desc: 一个简单的问候插件示例
|
||||
author: Soulter
|
||||
version: 0.1.0
|
||||
components: # 组件列表,将支持自动生成
|
||||
- class: handlers.commands.hello:HelloCommand
|
||||
type: command
|
||||
name: hello
|
||||
description: 发送问候消息
|
||||
subcommands:
|
||||
- name: wow
|
||||
description: 发送 "Hello, Astrbot!" 消息
|
||||
# - class: handlers.tools.echo:EchoTool
|
||||
# type: llm_tool
|
||||
# name: echo_tool
|
||||
# description: 回显输入的消息
|
||||
# - class: handlers.listeners.echo:EchoListener
|
||||
# type: listener
|
||||
# name: message_logger
|
||||
# description: 监听并记录所有消息
|
||||
6
src/run.py
Normal file
6
src/run.py
Normal file
@@ -0,0 +1,6 @@
|
||||
from astrbot_sdk.runtime.start_server import amain
|
||||
|
||||
import asyncio
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(amain())
|
||||
6
src/run_client.py
Normal file
6
src/run_client.py
Normal file
@@ -0,0 +1,6 @@
|
||||
from astrbot_sdk.runtime.start_client import amain
|
||||
|
||||
import asyncio
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(amain())
|
||||
Reference in New Issue
Block a user