mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
delete old sdk (#7)
Co-authored-by: whatevertogo <whatevertogo@users.noreply.github.com>
This commit is contained in:
whatevertogo
@@ -1,168 +0,0 @@
|
||||
# Inspired by MoonshotAI/kosong, credits to MoonshotAI/kosong authors for the original implementation.
|
||||
# License: Apache License 2.0
|
||||
|
||||
from typing import Any, ClassVar, Literal, cast
|
||||
|
||||
from pydantic import BaseModel, GetCoreSchemaHandler
|
||||
from pydantic_core import core_schema
|
||||
|
||||
|
||||
class ContentPart(BaseModel):
|
||||
"""A part of the content in a message."""
|
||||
|
||||
__content_part_registry: ClassVar[dict[str, type["ContentPart"]]] = {}
|
||||
|
||||
type: str
|
||||
|
||||
def __init_subclass__(cls, **kwargs: Any) -> None:
|
||||
super().__init_subclass__(**kwargs)
|
||||
|
||||
invalid_subclass_error_msg = f"ContentPart subclass {cls.__name__} must have a `type` field of type `str`"
|
||||
|
||||
type_value = getattr(cls, "type", None)
|
||||
if type_value is None or not isinstance(type_value, str):
|
||||
raise ValueError(invalid_subclass_error_msg)
|
||||
|
||||
cls.__content_part_registry[type_value] = cls
|
||||
|
||||
@classmethod
|
||||
def __get_pydantic_core_schema__(
|
||||
cls, source_type: Any, handler: GetCoreSchemaHandler
|
||||
) -> core_schema.CoreSchema:
|
||||
# If we're dealing with the base ContentPart class, use custom validation
|
||||
if cls.__name__ == "ContentPart":
|
||||
|
||||
def validate_content_part(value: Any) -> Any:
|
||||
# if it's already an instance of a ContentPart subclass, return it
|
||||
if hasattr(value, "__class__") and issubclass(value.__class__, cls):
|
||||
return value
|
||||
|
||||
# if it's a dict with a type field, dispatch to the appropriate subclass
|
||||
if isinstance(value, dict) and "type" in value:
|
||||
type_value: Any | None = cast(dict[str, Any], value).get("type")
|
||||
if not isinstance(type_value, str):
|
||||
raise ValueError(f"Cannot validate {value} as ContentPart")
|
||||
target_class = cls.__content_part_registry[type_value]
|
||||
return target_class.model_validate(value)
|
||||
|
||||
raise ValueError(f"Cannot validate {value} as ContentPart")
|
||||
|
||||
return core_schema.no_info_plain_validator_function(validate_content_part)
|
||||
|
||||
# for subclasses, use the default schema
|
||||
return handler(source_type)
|
||||
|
||||
|
||||
class TextPart(ContentPart):
|
||||
"""
|
||||
>>> TextPart(text="Hello, world!").model_dump()
|
||||
{'type': 'text', 'text': 'Hello, world!'}
|
||||
"""
|
||||
|
||||
type: str = "text"
|
||||
text: str
|
||||
|
||||
|
||||
class ImageURLPart(ContentPart):
|
||||
"""
|
||||
>>> ImageURLPart(image_url="http://example.com/image.jpg").model_dump()
|
||||
{'type': 'image_url', 'image_url': 'http://example.com/image.jpg'}
|
||||
"""
|
||||
|
||||
class ImageURL(BaseModel):
|
||||
url: str
|
||||
"""The URL of the image, can be data URI scheme like `data:image/png;base64,...`."""
|
||||
id: str | None = None
|
||||
"""The ID of the image, to allow LLMs to distinguish different images."""
|
||||
|
||||
type: str = "image_url"
|
||||
image_url: str
|
||||
|
||||
|
||||
class AudioURLPart(ContentPart):
|
||||
"""
|
||||
>>> AudioURLPart(audio_url=AudioURLPart.AudioURL(url="https://example.com/audio.mp3")).model_dump()
|
||||
{'type': 'audio_url', 'audio_url': {'url': 'https://example.com/audio.mp3', 'id': None}}
|
||||
"""
|
||||
|
||||
class AudioURL(BaseModel):
|
||||
url: str
|
||||
"""The URL of the audio, can be data URI scheme like `data:audio/aac;base64,...`."""
|
||||
id: str | None = None
|
||||
"""The ID of the audio, to allow LLMs to distinguish different audios."""
|
||||
|
||||
type: str = "audio_url"
|
||||
audio_url: AudioURL
|
||||
|
||||
|
||||
class ToolCall(BaseModel):
|
||||
"""
|
||||
A tool call requested by the assistant.
|
||||
|
||||
>>> ToolCall(
|
||||
... id="123",
|
||||
... function=ToolCall.FunctionBody(
|
||||
... name="function",
|
||||
... arguments="{}"
|
||||
... ),
|
||||
... ).model_dump()
|
||||
{'type': 'function', 'id': '123', 'function': {'name': 'function', 'arguments': '{}'}}
|
||||
"""
|
||||
|
||||
class FunctionBody(BaseModel):
|
||||
name: str
|
||||
arguments: str | None
|
||||
|
||||
type: Literal["function"] = "function"
|
||||
|
||||
id: str
|
||||
"""The ID of the tool call."""
|
||||
function: FunctionBody
|
||||
"""The function body of the tool call."""
|
||||
|
||||
|
||||
class ToolCallPart(BaseModel):
|
||||
"""A part of the tool call."""
|
||||
|
||||
arguments_part: str | None = None
|
||||
"""A part of the arguments of the tool call."""
|
||||
|
||||
|
||||
class Message(BaseModel):
|
||||
"""A message in a conversation."""
|
||||
|
||||
role: Literal[
|
||||
"system",
|
||||
"user",
|
||||
"assistant",
|
||||
"tool",
|
||||
]
|
||||
|
||||
content: str | list[ContentPart]
|
||||
"""The content of the message."""
|
||||
|
||||
|
||||
class AssistantMessageSegment(Message):
|
||||
"""A message segment from the assistant."""
|
||||
|
||||
role: Literal["assistant"] = "assistant"
|
||||
tool_calls: list[ToolCall] | list[dict] | None = None
|
||||
|
||||
|
||||
class ToolCallMessageSegment(Message):
|
||||
"""A message segment representing a tool call."""
|
||||
|
||||
role: Literal["tool"] = "tool"
|
||||
tool_call_id: str
|
||||
|
||||
|
||||
class UserMessageSegment(Message):
|
||||
"""A message segment from the user."""
|
||||
|
||||
role: Literal["user"] = "user"
|
||||
|
||||
|
||||
class SystemMessageSegment(Message):
|
||||
"""A message segment from the system."""
|
||||
|
||||
role: Literal["system"] = "system"
|
||||
@@ -1,17 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Generic
|
||||
|
||||
from typing_extensions import TypeVar
|
||||
|
||||
TContext = TypeVar("TContext", default=Any)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ContextWrapper(Generic[TContext]):
|
||||
"""A context for running an agent, which can be used to pass additional data or state."""
|
||||
|
||||
context: TContext
|
||||
tool_call_timeout: int = 60 # Default tool call timeout in seconds
|
||||
|
||||
|
||||
NoContext = ContextWrapper[None]
|
||||
@@ -1,286 +0,0 @@
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any, Generic
|
||||
|
||||
import jsonschema
|
||||
import mcp
|
||||
from deprecated import deprecated
|
||||
from pydantic import model_validator
|
||||
from pydantic.dataclasses import dataclass
|
||||
|
||||
from .run_context import ContextWrapper, TContext
|
||||
|
||||
ParametersType = dict[str, Any]
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolSchema:
|
||||
"""A class representing the schema of a tool for function calling."""
|
||||
|
||||
name: str
|
||||
"""The name of the tool."""
|
||||
|
||||
description: str
|
||||
"""The description of the tool."""
|
||||
|
||||
parameters: ParametersType
|
||||
"""The parameters of the tool, in JSON Schema format."""
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_parameters(self) -> "ToolSchema":
|
||||
jsonschema.validate(
|
||||
self.parameters, jsonschema.Draft202012Validator.META_SCHEMA
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
@dataclass
|
||||
class FunctionTool(ToolSchema, Generic[TContext]):
|
||||
"""A callable tool, for function calling."""
|
||||
|
||||
handler: Callable[..., Awaitable[Any]] | None = None
|
||||
"""a callable that implements the tool's functionality. It should be an async function."""
|
||||
|
||||
handler_module_path: str | None = None
|
||||
"""
|
||||
The module path of the handler function. This is empty when the origin is mcp.
|
||||
This field must be retained, as the handler will be wrapped in functools.partial during initialization,
|
||||
causing the handler's __module__ to be functools
|
||||
"""
|
||||
active: bool = True
|
||||
"""
|
||||
Whether the tool is active. This field is a special field for AstrBot.
|
||||
You can ignore it when integrating with other frameworks.
|
||||
"""
|
||||
|
||||
def __repr__(self):
|
||||
return f"FuncTool(name={self.name}, parameters={self.parameters}, description={self.description})"
|
||||
|
||||
async def call(
|
||||
self, context: ContextWrapper[TContext], **kwargs
|
||||
) -> str | mcp.types.CallToolResult:
|
||||
"""Run the tool with the given arguments. The handler field has priority."""
|
||||
raise NotImplementedError(
|
||||
"FunctionTool.call() must be implemented by subclasses or set a handler."
|
||||
)
|
||||
|
||||
|
||||
class ToolSet:
|
||||
"""A set of function tools that can be used in function calling.
|
||||
|
||||
This class provides methods to add, remove, and retrieve tools, as well as
|
||||
convert the tools to different API formats (OpenAI, Anthropic, Google GenAI).
|
||||
"""
|
||||
|
||||
def __init__(self, tools: list[FunctionTool] | None = None):
|
||||
self.tools: list[FunctionTool] = tools or []
|
||||
|
||||
def empty(self) -> bool:
|
||||
"""Check if the tool set is empty."""
|
||||
return len(self.tools) == 0
|
||||
|
||||
def add_tool(self, tool: FunctionTool):
|
||||
"""Add a tool to the set."""
|
||||
# 检查是否已存在同名工具
|
||||
for i, existing_tool in enumerate(self.tools):
|
||||
if existing_tool.name == tool.name:
|
||||
self.tools[i] = tool
|
||||
return
|
||||
self.tools.append(tool)
|
||||
|
||||
def remove_tool(self, name: str):
|
||||
"""Remove a tool by its name."""
|
||||
self.tools = [tool for tool in self.tools if tool.name != name]
|
||||
|
||||
def get_tool(self, name: str) -> FunctionTool | None:
|
||||
"""Get a tool by its name."""
|
||||
for tool in self.tools:
|
||||
if tool.name == name:
|
||||
return tool
|
||||
return None
|
||||
|
||||
@deprecated(reason="Use add_tool() instead", version="4.0.0")
|
||||
def add_func(
|
||||
self,
|
||||
name: str,
|
||||
func_args: list,
|
||||
desc: str,
|
||||
handler: Callable[..., Awaitable[Any]],
|
||||
):
|
||||
"""Add a function tool to the set."""
|
||||
params = {
|
||||
"type": "object", # hard-coded here
|
||||
"properties": {},
|
||||
}
|
||||
for param in func_args:
|
||||
params["properties"][param["name"]] = {
|
||||
"type": param["type"],
|
||||
"description": param["description"],
|
||||
}
|
||||
_func = FunctionTool(
|
||||
name=name,
|
||||
parameters=params,
|
||||
description=desc,
|
||||
handler=handler,
|
||||
)
|
||||
self.add_tool(_func)
|
||||
|
||||
@deprecated(reason="Use remove_tool() instead", version="4.0.0")
|
||||
def remove_func(self, name: str):
|
||||
"""Remove a function tool by its name."""
|
||||
self.remove_tool(name)
|
||||
|
||||
@deprecated(reason="Use get_tool() instead", version="4.0.0")
|
||||
def get_func(self, name: str) -> FunctionTool | None:
|
||||
"""Get all function tools."""
|
||||
return self.get_tool(name)
|
||||
|
||||
@property
|
||||
def func_list(self) -> list[FunctionTool]:
|
||||
"""Get the list of function tools."""
|
||||
return self.tools
|
||||
|
||||
def openai_schema(self, omit_empty_parameter_field: bool = False) -> list[dict]:
|
||||
"""Convert tools to OpenAI API function calling schema format."""
|
||||
result = []
|
||||
for tool in self.tools:
|
||||
func_def = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tool.name,
|
||||
"description": tool.description,
|
||||
},
|
||||
}
|
||||
|
||||
if (
|
||||
tool.parameters and tool.parameters.get("properties")
|
||||
) or not omit_empty_parameter_field:
|
||||
func_def["function"]["parameters"] = tool.parameters
|
||||
|
||||
result.append(func_def)
|
||||
return result
|
||||
|
||||
def anthropic_schema(self) -> list[dict]:
|
||||
"""Convert tools to Anthropic API format."""
|
||||
result = []
|
||||
for tool in self.tools:
|
||||
input_schema = {"type": "object"}
|
||||
if tool.parameters:
|
||||
input_schema["properties"] = tool.parameters.get("properties", {})
|
||||
input_schema["required"] = tool.parameters.get("required", [])
|
||||
tool_def = {
|
||||
"name": tool.name,
|
||||
"description": tool.description,
|
||||
"input_schema": input_schema,
|
||||
}
|
||||
result.append(tool_def)
|
||||
return result
|
||||
|
||||
def google_schema(self) -> dict:
|
||||
"""Convert tools to Google GenAI API format."""
|
||||
|
||||
def convert_schema(schema: dict) -> dict:
|
||||
"""Convert schema to Gemini API format."""
|
||||
supported_types = {
|
||||
"string",
|
||||
"number",
|
||||
"integer",
|
||||
"boolean",
|
||||
"array",
|
||||
"object",
|
||||
"null",
|
||||
}
|
||||
supported_formats = {
|
||||
"string": {"enum", "date-time"},
|
||||
"integer": {"int32", "int64"},
|
||||
"number": {"float", "double"},
|
||||
}
|
||||
|
||||
if "anyOf" in schema:
|
||||
return {"anyOf": [convert_schema(s) for s in schema["anyOf"]]}
|
||||
|
||||
result = {}
|
||||
|
||||
if "type" in schema and schema["type"] in supported_types:
|
||||
result["type"] = schema["type"]
|
||||
if "format" in schema and schema["format"] in supported_formats.get(
|
||||
result["type"],
|
||||
set(),
|
||||
):
|
||||
result["format"] = schema["format"]
|
||||
else:
|
||||
result["type"] = "null"
|
||||
|
||||
support_fields = {
|
||||
"title",
|
||||
"description",
|
||||
"enum",
|
||||
"minimum",
|
||||
"maximum",
|
||||
"maxItems",
|
||||
"minItems",
|
||||
"nullable",
|
||||
"required",
|
||||
}
|
||||
result.update({k: schema[k] for k in support_fields if k in schema})
|
||||
|
||||
if "properties" in schema:
|
||||
properties = {}
|
||||
for key, value in schema["properties"].items():
|
||||
prop_value = convert_schema(value)
|
||||
if "default" in prop_value:
|
||||
del prop_value["default"]
|
||||
properties[key] = prop_value
|
||||
|
||||
if properties:
|
||||
result["properties"] = properties
|
||||
|
||||
if "items" in schema:
|
||||
result["items"] = convert_schema(schema["items"])
|
||||
|
||||
return result
|
||||
|
||||
tools = []
|
||||
for tool in self.tools:
|
||||
d: dict[str, Any] = {
|
||||
"name": tool.name,
|
||||
"description": tool.description,
|
||||
}
|
||||
if tool.parameters:
|
||||
d["parameters"] = convert_schema(tool.parameters)
|
||||
tools.append(d)
|
||||
|
||||
declarations = {}
|
||||
if tools:
|
||||
declarations["function_declarations"] = tools
|
||||
return declarations
|
||||
|
||||
@deprecated(reason="Use openai_schema() instead", version="4.0.0")
|
||||
def get_func_desc_openai_style(self, omit_empty_parameter_field: bool = False):
|
||||
return self.openai_schema(omit_empty_parameter_field)
|
||||
|
||||
@deprecated(reason="Use anthropic_schema() instead", version="4.0.0")
|
||||
def get_func_desc_anthropic_style(self):
|
||||
return self.anthropic_schema()
|
||||
|
||||
@deprecated(reason="Use google_schema() instead", version="4.0.0")
|
||||
def get_func_desc_google_genai_style(self):
|
||||
return self.google_schema()
|
||||
|
||||
def names(self) -> list[str]:
|
||||
"""获取所有工具的名称列表"""
|
||||
return [tool.name for tool in self.tools]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.tools)
|
||||
|
||||
def __bool__(self):
|
||||
return len(self.tools) > 0
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self.tools)
|
||||
|
||||
def __repr__(self):
|
||||
return f"ToolSet(tools={self.tools})"
|
||||
|
||||
def __str__(self):
|
||||
return f"ToolSet(tools={self.tools})"
|
||||
@@ -1,5 +0,0 @@
|
||||
from .cli.main import cli
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli()
|
||||
@@ -1 +0,0 @@
|
||||
class AstrBotConfig(dict): ...
|
||||
@@ -1,224 +0,0 @@
|
||||
from astr_agent_sdk.message import AssistantMessageSegment, UserMessageSegment
|
||||
from ...api.basic.entities import Conversation
|
||||
|
||||
|
||||
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
|
||||
) -> None:
|
||||
"""切换会话的对话
|
||||
|
||||
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): 每页大小
|
||||
|
||||
"""
|
||||
...
|
||||
@@ -1,23 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class Conversation:
|
||||
"""The conversation entity representing a chat session."""
|
||||
|
||||
platform_id: str
|
||||
"""The platform ID in AstrBot"""
|
||||
user_id: str
|
||||
"""The user ID associated with the conversation."""
|
||||
cid: str
|
||||
"""The conversation ID, in UUID format."""
|
||||
history: str = ""
|
||||
"""The conversation history as a string."""
|
||||
title: str | None = ""
|
||||
"""The title of the conversation. For now, it's only used in WebChat."""
|
||||
persona_id: str | None = ""
|
||||
"""The persona ID associated with the conversation."""
|
||||
created_at: int = 0
|
||||
"""The timestamp when the conversation was created."""
|
||||
updated_at: int = 0
|
||||
"""The timestamp when the conversation was last updated."""
|
||||
@@ -1,2 +0,0 @@
|
||||
class CommandComponent:
|
||||
pass
|
||||
@@ -1,5 +0,0 @@
|
||||
from .astr_message_event import AstrMessageEvent
|
||||
|
||||
__all__ = [
|
||||
"AstrMessageEvent",
|
||||
]
|
||||
@@ -1,370 +0,0 @@
|
||||
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)
|
||||
"""
|
||||
@@ -1,98 +0,0 @@
|
||||
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
|
||||
@@ -1,93 +0,0 @@
|
||||
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
|
||||
@@ -1,19 +0,0 @@
|
||||
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() # 发送消息后
|
||||
@@ -1,65 +0,0 @@
|
||||
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",
|
||||
]
|
||||
@@ -1,32 +0,0 @@
|
||||
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
|
||||
@@ -1,7 +0,0 @@
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class MessageType(Enum):
|
||||
GROUP_MESSAGE = "GroupMessage" # 群组形式的消息
|
||||
FRIEND_MESSAGE = "FriendMessage" # 私聊、好友等单聊消息
|
||||
OTHER_MESSAGE = "OtherMessage" # 其他类型的消息，如系统消息等
|
||||
@@ -1,136 +0,0 @@
|
||||
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
|
||||
@@ -1,225 +0,0 @@
|
||||
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,
|
||||
}
|
||||
@@ -1,18 +0,0 @@
|
||||
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 文件路径（相对于插件目录）"""
|
||||
@@ -1,126 +0,0 @@
|
||||
from __future__ import annotations
|
||||
import json
|
||||
from anthropic.types import Message as AnthropicMessage
|
||||
from google.genai.types import GenerateContentResponse
|
||||
from openai.types.chat.chat_completion import ChatCompletion
|
||||
from dataclasses import dataclass, field
|
||||
from ..message.chain import MessageChain
|
||||
from ..message import components as Comp
|
||||
from typing import Any
|
||||
from astr_agent_sdk.message import ToolCall
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMResponse:
|
||||
role: str
|
||||
"""角色, assistant, tool, err"""
|
||||
result_chain: MessageChain | None = None
|
||||
"""返回的消息链"""
|
||||
tools_call_args: list[dict[str, Any]] = field(default_factory=list)
|
||||
"""工具调用参数"""
|
||||
tools_call_name: list[str] = field(default_factory=list)
|
||||
"""工具调用名称"""
|
||||
tools_call_ids: list[str] = field(default_factory=list)
|
||||
"""工具调用 ID"""
|
||||
|
||||
raw_completion: (
|
||||
ChatCompletion | GenerateContentResponse | AnthropicMessage | None
|
||||
) = None
|
||||
_new_record: dict[str, Any] | None = None
|
||||
|
||||
_completion_text: str = ""
|
||||
|
||||
is_chunk: bool = False
|
||||
"""是否是流式输出的单个 Chunk"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
role: str,
|
||||
completion_text: str = "",
|
||||
result_chain: MessageChain | None = None,
|
||||
tools_call_args: list[dict[str, Any]] | None = None,
|
||||
tools_call_name: list[str] | None = None,
|
||||
tools_call_ids: list[str] | None = None,
|
||||
raw_completion: ChatCompletion
|
||||
| GenerateContentResponse
|
||||
| AnthropicMessage
|
||||
| None = None,
|
||||
_new_record: dict[str, Any] | None = None,
|
||||
is_chunk: bool = False,
|
||||
):
|
||||
"""初始化 LLMResponse
|
||||
|
||||
Args:
|
||||
role (str): 角色, assistant, tool, err
|
||||
completion_text (str, optional): 返回的结果文本，已经过时，推荐使用 result_chain. Defaults to "".
|
||||
result_chain (MessageChain, optional): 返回的消息链. Defaults to None.
|
||||
tools_call_args (List[Dict[str, any]], optional): 工具调用参数. Defaults to None.
|
||||
tools_call_name (List[str], optional): 工具调用名称. Defaults to None.
|
||||
raw_completion (ChatCompletion, optional): 原始响应, OpenAI 格式. Defaults to None.
|
||||
|
||||
"""
|
||||
if tools_call_args is None:
|
||||
tools_call_args = []
|
||||
if tools_call_name is None:
|
||||
tools_call_name = []
|
||||
if tools_call_ids is None:
|
||||
tools_call_ids = []
|
||||
|
||||
self.role = role
|
||||
self.completion_text = completion_text
|
||||
self.result_chain = result_chain
|
||||
self.tools_call_args = tools_call_args
|
||||
self.tools_call_name = tools_call_name
|
||||
self.tools_call_ids = tools_call_ids
|
||||
self.raw_completion = raw_completion
|
||||
self._new_record = _new_record
|
||||
self.is_chunk = is_chunk
|
||||
|
||||
@property
|
||||
def completion_text(self):
|
||||
if self.result_chain:
|
||||
return self.result_chain.get_plain_text()
|
||||
return self._completion_text
|
||||
|
||||
@completion_text.setter
|
||||
def completion_text(self, value):
|
||||
if self.result_chain:
|
||||
self.result_chain.chain = [
|
||||
comp
|
||||
for comp in self.result_chain.chain
|
||||
if not isinstance(comp, Comp.Plain)
|
||||
] # 清空 Plain 组件
|
||||
self.result_chain.chain.insert(0, Comp.Plain(text=value))
|
||||
else:
|
||||
self._completion_text = value
|
||||
|
||||
def to_openai_tool_calls(self) -> list[dict]:
|
||||
"""Convert to OpenAI tool calls format. Deprecated, use to_openai_to_calls_model instead."""
|
||||
ret = []
|
||||
for idx, tool_call_arg in enumerate(self.tools_call_args):
|
||||
ret.append(
|
||||
{
|
||||
"id": self.tools_call_ids[idx],
|
||||
"function": {
|
||||
"name": self.tools_call_name[idx],
|
||||
"arguments": json.dumps(tool_call_arg),
|
||||
},
|
||||
"type": "function",
|
||||
},
|
||||
)
|
||||
return ret
|
||||
|
||||
def to_openai_to_calls_model(self) -> list[ToolCall]:
|
||||
"""The same as to_openai_tool_calls but return pydantic model."""
|
||||
ret = []
|
||||
for idx, tool_call_arg in enumerate(self.tools_call_args):
|
||||
ret.append(
|
||||
ToolCall(
|
||||
id=self.tools_call_ids[idx],
|
||||
function=ToolCall.FunctionBody(
|
||||
name=self.tools_call_name[idx],
|
||||
arguments=json.dumps(tool_call_arg),
|
||||
),
|
||||
),
|
||||
)
|
||||
return ret
|
||||
@@ -1,152 +0,0 @@
|
||||
from abc import ABC
|
||||
from typing import Any, Callable
|
||||
from ..basic.conversation_mgr import BaseConversationManager
|
||||
from astr_agent_sdk.tool import ToolSet, FunctionTool
|
||||
from astr_agent_sdk.message import Message
|
||||
from ..provider.entities import LLMResponse
|
||||
from ..message.chain import MessageChain
|
||||
|
||||
|
||||
class Context(ABC):
|
||||
conversation_manager: BaseConversationManager
|
||||
persona_manager: Any
|
||||
|
||||
def __init__(self):
|
||||
self._registered_managers: dict[str, Any] = {}
|
||||
self._registered_functions: dict[str, Callable] = {}
|
||||
|
||||
def _register_component(self, *components: Any) -> None:
|
||||
"""Register a components instance and its public methods.
|
||||
|
||||
This allows the components's methods to be called via RPC using the pattern:
|
||||
ComponentClassName.method_name
|
||||
|
||||
Args:
|
||||
components: The components instance to register
|
||||
"""
|
||||
for component in components:
|
||||
class_name = component.__class__.__name__
|
||||
self._registered_managers[class_name] = component
|
||||
|
||||
# Register all public methods (not starting with _)
|
||||
for attr_name in dir(component):
|
||||
if not attr_name.startswith("_"):
|
||||
attr = getattr(component, attr_name)
|
||||
if callable(attr):
|
||||
full_name = f"{class_name}.{attr_name}"
|
||||
self._registered_functions[full_name] = attr
|
||||
|
||||
async def llm_generate(
|
||||
self,
|
||||
chat_provider_id: str,
|
||||
prompt: str | None = None,
|
||||
image_urls: list[str] | None = None,
|
||||
tools: ToolSet | None = None,
|
||||
system_prompt: str | None = None,
|
||||
contexts: list[Message] | list[dict] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> LLMResponse:
|
||||
"""Call the LLM to generate a response. The method will not automatically execute tool calls. If you want to use tool calls, please use `tool_loop_agent()`.
|
||||
|
||||
Args:
|
||||
chat_provider_id: The chat provider ID to use.
|
||||
prompt: The prompt to send to the LLM, if `contexts` and `prompt` are both provided, `prompt` will be appended as the last user message
|
||||
image_urls: List of image URLs to include in the prompt, if `contexts` and `prompt` are both provided, `image_urls` will be appended to the last user message
|
||||
tools: ToolSet of tools available to the LLM
|
||||
system_prompt: System prompt to guide the LLM's behavior, if provided, it will always insert as the first system message in the context
|
||||
contexts: context messages for the LLM
|
||||
**kwargs: Additional keyword arguments for LLM generation, OpenAI compatible
|
||||
|
||||
Raises:
|
||||
ChatProviderNotFoundError: If the specified chat provider ID is not found
|
||||
Exception: For other errors during LLM generation
|
||||
"""
|
||||
...
|
||||
|
||||
async def tool_loop_agent(
|
||||
self,
|
||||
chat_provider_id: str,
|
||||
prompt: str | None = None,
|
||||
image_urls: list[str] | None = None,
|
||||
tools: ToolSet | None = None,
|
||||
system_prompt: str | None = None,
|
||||
contexts: list[Message] | list[dict] | None = None,
|
||||
max_steps: int = 30,
|
||||
**kwargs: Any,
|
||||
) -> LLMResponse:
|
||||
"""Run an agent loop that allows the LLM to call tools iteratively until a final answer is produced.
|
||||
|
||||
Args:
|
||||
chat_provider_id: The chat provider ID to use.
|
||||
prompt: The prompt to send to the LLM, if `contexts` and `prompt` are both provided, `prompt` will be appended as the last user message
|
||||
image_urls: List of image URLs to include in the prompt, if `contexts` and `prompt` are both provided, `image_urls` will be appended to the last user message
|
||||
tools: ToolSet of tools available to the LLM
|
||||
system_prompt: System prompt to guide the LLM's behavior, if provided, it will always insert as the first system message in the context
|
||||
contexts: context messages for the LLM
|
||||
max_steps: Maximum number of tool calls before stopping the loop
|
||||
**kwargs: Additional keyword arguments for LLM generation, OpenAI compatible
|
||||
|
||||
Returns:
|
||||
The final LLMResponse after tool calls are completed.
|
||||
|
||||
Raises:
|
||||
ChatProviderNotFoundError: If the specified chat provider ID is not found
|
||||
Exception: For other errors during LLM generation
|
||||
"""
|
||||
...
|
||||
|
||||
async def send_message(
|
||||
self,
|
||||
session: str,
|
||||
message_chain: MessageChain,
|
||||
) -> None:
|
||||
"""Send a message to a user or group.
|
||||
|
||||
Args:
|
||||
session: unified message origin(umo), this can represent a user or group in a specific platform instance
|
||||
message_chain: The MessageChain to send
|
||||
|
||||
Raises:
|
||||
Exception: If sending the message fails
|
||||
"""
|
||||
...
|
||||
|
||||
async def add_llm_tools(self, *tools: FunctionTool) -> None:
|
||||
"""Add tools to the LLM's toolset.
|
||||
|
||||
Args:
|
||||
tools: The FunctionTool instances to add
|
||||
"""
|
||||
...
|
||||
|
||||
async def put_kv_data(
|
||||
self,
|
||||
key: str,
|
||||
value: dict,
|
||||
) -> None:
|
||||
"""Insert a key-value pair data. The data will permanently stored in AstrBot unless user explicitly deleted.
|
||||
|
||||
Args:
|
||||
key: The key to insert
|
||||
value: The value to insert
|
||||
"""
|
||||
...
|
||||
|
||||
async def get_kv_data(self, key: str) -> dict | None:
|
||||
"""Get a value by key from the key-value store.
|
||||
|
||||
Args:
|
||||
key: The key to retrieve
|
||||
|
||||
Returns:
|
||||
The value associated with the key, or None if not found
|
||||
"""
|
||||
...
|
||||
|
||||
async def delete_kv_data(self, key: str) -> None:
|
||||
"""Delete a key-value pair by key.
|
||||
|
||||
Args:
|
||||
key: The key to delete
|
||||
"""
|
||||
...
|
||||
@@ -1,59 +0,0 @@
|
||||
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}"
|
||||
@@ -1,3 +0,0 @@
|
||||
from .main import cli
|
||||
|
||||
__all__ = ["cli"]
|
||||
@@ -1,76 +0,0 @@
|
||||
import asyncio
|
||||
import sys
|
||||
|
||||
import click
|
||||
from loguru import logger
|
||||
|
||||
from ..runtime.serve import run_plugin_worker, run_supervisor, run_websocket_server
|
||||
|
||||
|
||||
def setup_logger(verbose: bool = False):
|
||||
"""Configure loguru for CLI output"""
|
||||
# Remove default handler
|
||||
logger.remove()
|
||||
|
||||
# Add custom handler with CLI-friendly format
|
||||
log_format = (
|
||||
"<green>{time:HH:mm:ss}</green> | "
|
||||
"<level>{level: <8}</level> | "
|
||||
"<level>{message}</level>"
|
||||
)
|
||||
|
||||
level = "DEBUG" if verbose else "INFO"
|
||||
|
||||
logger.add(
|
||||
sys.stderr,
|
||||
format=log_format,
|
||||
level=level,
|
||||
colorize=True,
|
||||
)
|
||||
|
||||
|
||||
@click.group()
|
||||
@click.option("-v", "--verbose", is_flag=True, help="Enable verbose output")
|
||||
@click.pass_context
|
||||
def cli(ctx, verbose):
|
||||
"""AstrBot SDK CLI"""
|
||||
ctx.ensure_object(dict)
|
||||
ctx.obj["verbose"] = verbose
|
||||
setup_logger(verbose)
|
||||
|
||||
|
||||
@cli.command()
|
||||
@click.option(
|
||||
"--plugins-dir",
|
||||
default="plugins",
|
||||
type=click.Path(file_okay=False, dir_okay=True, path_type=str),
|
||||
help="Directory containing plugin folders",
|
||||
)
|
||||
@click.pass_context
|
||||
def run(ctx, plugins_dir: str):
|
||||
"""Start the plugin supervisor over stdio."""
|
||||
logger.info(f"Starting plugin supervisor with plugins dir: {plugins_dir}")
|
||||
asyncio.run(run_supervisor(plugins_dir=plugins_dir))
|
||||
|
||||
|
||||
@cli.command(hidden=True)
|
||||
@click.option(
|
||||
"--plugin-dir",
|
||||
required=True,
|
||||
type=click.Path(file_okay=False, dir_okay=True, path_type=str),
|
||||
)
|
||||
def worker(plugin_dir: str):
|
||||
"""Internal command used by the supervisor to start a worker."""
|
||||
asyncio.run(run_plugin_worker(plugin_dir=plugin_dir))
|
||||
|
||||
|
||||
@cli.command(hidden=True)
|
||||
@click.option("--port", default=8765, help="WebSocket server port", type=int)
|
||||
def websocket(port: int):
|
||||
"""Legacy websocket runtime entrypoint."""
|
||||
logger.info(f"Starting WebSocket server on port {port}...")
|
||||
asyncio.run(run_websocket_server(port=port))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
cli()
|
||||
@@ -1,8 +0,0 @@
|
||||
# AstrBot SDK Runtime Context API
|
||||
|
||||
这个包下存储了暴露给 AstrBot 插件的 Context API 的 RPC 实现。
|
||||
|
||||
## 组件
|
||||
|
||||
- `Context`：这是在实例化插件时，注入到插件中的上下文对象。它封装了插件可以调用的各种功能组件。
|
||||
- `ConversationManager`：这是一个管理对话相关的功能组件。它提供了与对话历史、用户信息等相关的操作接口。
|
||||
@@ -1,23 +0,0 @@
|
||||
from __future__ import annotations
|
||||
from ...api.star.context import Context as BaseContext
|
||||
from .conversation_mgr import ConversationManager
|
||||
from ..star_runner import StarRunner
|
||||
|
||||
|
||||
class Context(BaseContext):
|
||||
def __init__(self, conversation_manager: ConversationManager):
|
||||
super().__init__()
|
||||
self.conversation_manager = conversation_manager
|
||||
# Auto-register the conversation manager
|
||||
self._register_component(self.conversation_manager)
|
||||
|
||||
@classmethod
|
||||
def default_context(cls, runner: StarRunner) -> Context:
|
||||
"""Create a default context instance.
|
||||
|
||||
Args:
|
||||
runner: Optional StarRunner instance to inject into conversation manager.
|
||||
If provided, enables RPC functionality.
|
||||
"""
|
||||
conversation_manager = ConversationManager(runner)
|
||||
return cls(conversation_manager=conversation_manager)
|
||||
@@ -1,140 +0,0 @@
|
||||
from astr_agent_sdk.message import AssistantMessageSegment, UserMessageSegment
|
||||
from astrbot_sdk.api.basic.entities import Conversation
|
||||
from ...api.basic.conversation_mgr import BaseConversationManager
|
||||
from ..star_runner import StarRunner
|
||||
|
||||
|
||||
class ConversationManager(BaseConversationManager):
|
||||
def __init__(self, runner: StarRunner):
|
||||
"""Initialize ConversationManager.
|
||||
|
||||
Args:
|
||||
runner: Optional StarRunner instance for RPC functionality.
|
||||
"""
|
||||
self.runner = runner
|
||||
|
||||
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:
|
||||
result = await self.runner.call_context_function(
|
||||
f"{self.__class__.__name__}.{self.new_conversation.__name__}",
|
||||
{
|
||||
"unified_msg_origin": unified_msg_origin,
|
||||
"platform_id": platform_id,
|
||||
"content": content,
|
||||
"title": title,
|
||||
"persona_id": persona_id,
|
||||
},
|
||||
)
|
||||
return result["data"]
|
||||
|
||||
async def switch_conversation(
|
||||
self, unified_msg_origin: str, conversation_id: str
|
||||
) -> None:
|
||||
await self.runner.call_context_function(
|
||||
f"{self.__class__.__name__}.{self.switch_conversation.__name__}",
|
||||
{
|
||||
"unified_msg_origin": unified_msg_origin,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
async def delete_conversation(
|
||||
self,
|
||||
unified_msg_origin: str,
|
||||
conversation_id: str | None = None,
|
||||
) -> None:
|
||||
await self.runner.call_context_function(
|
||||
f"{self.__class__.__name__}.{self.delete_conversation.__name__}",
|
||||
{
|
||||
"unified_msg_origin": unified_msg_origin,
|
||||
"conversation_id": conversation_id,
|
||||
},
|
||||
)
|
||||
|
||||
async def delete_conversations_by_user_id(self, unified_msg_origin: str) -> None:
|
||||
await self.runner.call_context_function(
|
||||
f"{self.__class__.__name__}.{self.delete_conversations_by_user_id.__name__}",
|
||||
{
|
||||
"unified_msg_origin": unified_msg_origin,
|
||||
},
|
||||
)
|
||||
|
||||
async def get_curr_conversation_id(self, unified_msg_origin: str) -> str | None:
|
||||
result = await self.runner.call_context_function(
|
||||
f"{self.__class__.__name__}.{self.get_curr_conversation_id.__name__}",
|
||||
{
|
||||
"unified_msg_origin": unified_msg_origin,
|
||||
},
|
||||
)
|
||||
return result["data"]
|
||||
|
||||
async def get_conversation(
|
||||
self,
|
||||
unified_msg_origin: str,
|
||||
conversation_id: str,
|
||||
create_if_not_exists: bool = False,
|
||||
) -> Conversation | None:
|
||||
result = await self.runner.call_context_function(
|
||||
f"{self.__class__.__name__}.{self.get_conversation.__name__}",
|
||||
{
|
||||
"unified_msg_origin": unified_msg_origin,
|
||||
"conversation_id": conversation_id,
|
||||
"create_if_not_exists": create_if_not_exists,
|
||||
},
|
||||
)
|
||||
return Conversation(**result["data"]) if result["data"] else None
|
||||
|
||||
async def get_conversations(
|
||||
self, unified_msg_origin: str | None = None, platform_id: str | None = None
|
||||
) -> list[Conversation]:
|
||||
result = await self.runner.call_context_function(
|
||||
f"{self.__class__.__name__}.{self.get_conversations.__name__}",
|
||||
{
|
||||
"unified_msg_origin": unified_msg_origin,
|
||||
"platform_id": platform_id,
|
||||
},
|
||||
)
|
||||
return [Conversation(**conv) for conv in result["data"]]
|
||||
|
||||
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:
|
||||
await self.runner.call_context_function(
|
||||
f"{self.__class__.__name__}.{self.update_conversation.__name__}",
|
||||
{
|
||||
"unified_msg_origin": unified_msg_origin,
|
||||
"conversation_id": conversation_id,
|
||||
"history": history,
|
||||
"title": title,
|
||||
"persona_id": persona_id,
|
||||
},
|
||||
)
|
||||
|
||||
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
|
||||
"""
|
||||
...
|
||||
@@ -1,39 +0,0 @@
|
||||
"""
|
||||
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_star import NewStdioStar, NewWebSocketStar
|
||||
from ..api.star.context import Context
|
||||
# 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, context: Context, star_name: str, config: dict
|
||||
) -> NewStdioStar:
|
||||
"""Connect to a new-style stdio star given its name."""
|
||||
star = NewStdioStar(context=context, **config)
|
||||
await star.initialize()
|
||||
self.vs_map[star_name] = star
|
||||
return star
|
||||
|
||||
async def connect_to_websocket_star(
|
||||
self, context: Context, star_name: str, config: dict
|
||||
) -> NewWebSocketStar:
|
||||
"""Connect to a new-style websocket star given its name."""
|
||||
star = NewWebSocketStar(context=context, **config)
|
||||
await star.initialize()
|
||||
self.vs_map[star_name] = star
|
||||
return star
|
||||
@@ -1,7 +0,0 @@
|
||||
# AstrBor SDK 与 Core 通信的数据交换实现
|
||||
|
||||
这个包下存储了 AstrBot 插件运行时与 AstrBot Core 之间通信的数据交换实现。
|
||||
|
||||
AstrBot SDK 设计了两种传输协议，即 stdio 和 WebSockets，用于实现 AstrBot 插件与 AstrBot Core 之间的双向通信。
|
||||
|
||||
在这两种传输协议之上，我们使用 JSON-RPC 2.0 作为通信的消息格式和调用规范。
|
||||
@@ -1,208 +0,0 @@
|
||||
# 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
|
||||
@@ -1,5 +0,0 @@
|
||||
from .base import JSONRPCClient
|
||||
from .stdio import StdioClient
|
||||
from .websocket import WebSocketClient
|
||||
|
||||
__all__ = ["JSONRPCClient", "StdioClient", "WebSocketClient"]
|
||||
@@ -1,14 +0,0 @@
|
||||
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__()
|
||||
@@ -1,222 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
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,
|
||||
env: dict[str, 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._env = env or os.environ.copy()
|
||||
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,
|
||||
env=self._env,
|
||||
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:
|
||||
if self._stdout:
|
||||
try:
|
||||
self._stdout.close()
|
||||
except Exception:
|
||||
logger.debug("Failed to close subprocess stdin cleanly")
|
||||
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}")
|
||||
@@ -1,235 +0,0 @@
|
||||
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}")
|
||||
@@ -1,39 +0,0 @@
|
||||
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
|
||||
@@ -1,219 +0,0 @@
|
||||
import asyncio
|
||||
from typing import Any
|
||||
from loguru import logger
|
||||
from .jsonrpc import (
|
||||
JSONRPCMessage,
|
||||
JSONRPCRequest,
|
||||
JSONRPCSuccessResponse,
|
||||
JSONRPCErrorResponse,
|
||||
)
|
||||
from .transport import JSONRPCTransport
|
||||
from ..types import (
|
||||
HandlerStreamStartNotification,
|
||||
HandlerStreamUpdateNotification,
|
||||
HandlerStreamEndNotification,
|
||||
)
|
||||
|
||||
|
||||
class RPCRequestHelper:
|
||||
"""Manages RPC communication state and pending requests.
|
||||
|
||||
Supports both single-response and streaming (multi-response) RPC patterns:
|
||||
- Single response: Uses asyncio.Future
|
||||
- Streaming: Uses asyncio.Queue for multiple responses
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._request_id_counter = 0
|
||||
self.pending_requests: dict[
|
||||
str, asyncio.Future[JSONRPCMessage] | asyncio.Queue[Any]
|
||||
] = {}
|
||||
|
||||
def _generate_request_id(self) -> str:
|
||||
"""Generate a unique request ID."""
|
||||
self._request_id_counter += 1
|
||||
return str(self._request_id_counter)
|
||||
|
||||
async def call_rpc(
|
||||
self, transport_impl: JSONRPCTransport, message: JSONRPCMessage
|
||||
) -> JSONRPCMessage | None:
|
||||
"""Send RPC request and wait for a single response.
|
||||
|
||||
Args:
|
||||
transport_impl: The transport to send the message through
|
||||
message: The JSON-RPC request message
|
||||
|
||||
Returns:
|
||||
The JSON-RPC response message, or None if no response expected
|
||||
"""
|
||||
if message.id is None:
|
||||
await transport_impl.send_message(message)
|
||||
return None
|
||||
|
||||
future: asyncio.Future[JSONRPCMessage] = (
|
||||
asyncio.get_event_loop().create_future()
|
||||
)
|
||||
self.pending_requests[message.id] = future
|
||||
await transport_impl.send_message(message)
|
||||
result = await future
|
||||
return result
|
||||
|
||||
async def call_rpc_streaming(
|
||||
self, transport_impl: JSONRPCTransport, message: JSONRPCMessage
|
||||
) -> asyncio.Queue[Any]:
|
||||
"""Send RPC request and expect multiple streaming responses.
|
||||
|
||||
The responses will be delivered via notifications with methods:
|
||||
- handler_stream_start: Stream started
|
||||
- handler_stream_update: New data available
|
||||
- handler_stream_end: Stream completed
|
||||
|
||||
Args:
|
||||
transport_impl: The transport to send the message through
|
||||
message: The JSON-RPC request message
|
||||
|
||||
Returns:
|
||||
An asyncio.Queue that will receive streamed results
|
||||
"""
|
||||
if message.id is None:
|
||||
raise ValueError("Streaming RPC calls require a request ID")
|
||||
|
||||
queue: asyncio.Queue[Any] = asyncio.Queue()
|
||||
self.pending_requests[message.id] = queue
|
||||
|
||||
await transport_impl.send_message(message)
|
||||
return queue
|
||||
|
||||
def resolve_pending_request(
|
||||
self, message: JSONRPCSuccessResponse | JSONRPCErrorResponse
|
||||
):
|
||||
"""Resolve a pending request with a response.
|
||||
|
||||
For single-response requests (Future), sets the result/exception.
|
||||
For streaming requests (Queue), logs completion/error but queue is managed separately.
|
||||
|
||||
Args:
|
||||
message: The JSON-RPC response message
|
||||
"""
|
||||
if message.id not in self.pending_requests:
|
||||
logger.warning(f"Received response for unknown request ID: {message.id}")
|
||||
return
|
||||
|
||||
pending = self.pending_requests[message.id]
|
||||
|
||||
if isinstance(pending, asyncio.Future):
|
||||
# Single response mode
|
||||
self.pending_requests.pop(message.id)
|
||||
if not pending.done():
|
||||
if isinstance(message, JSONRPCSuccessResponse):
|
||||
pending.set_result(message)
|
||||
else:
|
||||
pending.set_exception(
|
||||
RuntimeError(
|
||||
f"RPC Error {message.error.code}: {message.error.message}"
|
||||
)
|
||||
)
|
||||
elif isinstance(pending, asyncio.Queue):
|
||||
# Streaming mode - final response received
|
||||
if isinstance(message, JSONRPCSuccessResponse):
|
||||
logger.debug(f"Streaming request {message.id} completed successfully")
|
||||
else:
|
||||
logger.error(
|
||||
f"Streaming request {message.id} failed: {message.error.message}"
|
||||
)
|
||||
# Put error marker in queue
|
||||
asyncio.create_task(
|
||||
pending.put({"_error": True, "message": message.error.message})
|
||||
)
|
||||
|
||||
async def handle_stream_notification(self, notification: JSONRPCRequest) -> None:
|
||||
"""Handle incoming streaming notifications.
|
||||
|
||||
Processes handler_stream_start/update/end notifications and updates
|
||||
the corresponding queue.
|
||||
|
||||
Args:
|
||||
notification: The streaming notification message
|
||||
|
||||
Raises:
|
||||
ValueError: If the notification method is not a valid stream notification
|
||||
"""
|
||||
# Validate notification method
|
||||
if notification.method not in [
|
||||
"handler_stream_start",
|
||||
"handler_stream_update",
|
||||
"handler_stream_end",
|
||||
]:
|
||||
raise ValueError(
|
||||
f"Invalid stream notification method: {notification.method}"
|
||||
)
|
||||
|
||||
# Extract common parameters
|
||||
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":
|
||||
try:
|
||||
typed_notification = HandlerStreamStartNotification.model_validate(
|
||||
notification.model_dump()
|
||||
)
|
||||
logger.debug(
|
||||
f"Stream started for handler {typed_notification.params.handler_full_name}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Invalid handler_stream_start notification: {e}")
|
||||
# Optionally put a start marker in the queue if needed
|
||||
# await pending.put({"_stream_start": True})
|
||||
|
||||
elif notification.method == "handler_stream_update":
|
||||
try:
|
||||
typed_notification = HandlerStreamUpdateNotification.model_validate(
|
||||
notification.model_dump()
|
||||
)
|
||||
# Put the streamed data into the queue
|
||||
data = typed_notification.params.data
|
||||
logger.debug(f"Stream update for request {request_id}: {data}")
|
||||
if data is not None:
|
||||
await pending.put(data)
|
||||
except Exception as e:
|
||||
logger.error(f"Invalid handler_stream_update notification: {e}")
|
||||
|
||||
elif notification.method == "handler_stream_end":
|
||||
try:
|
||||
typed_notification = HandlerStreamEndNotification.model_validate(
|
||||
notification.model_dump()
|
||||
)
|
||||
logger.debug(
|
||||
f"Stream ended for handler {typed_notification.params.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
|
||||
asyncio.create_task(self._cleanup_stream_request(request_id))
|
||||
except Exception as e:
|
||||
logger.error(f"Invalid handler_stream_end notification: {e}")
|
||||
|
||||
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}")
|
||||
@@ -1,9 +0,0 @@
|
||||
from .base import JSONRPCServer
|
||||
from .stdio import StdioServer
|
||||
from .websockets import WebSocketServer
|
||||
|
||||
__all__ = [
|
||||
"JSONRPCServer",
|
||||
"StdioServer",
|
||||
"WebSocketServer",
|
||||
]
|
||||
@@ -1,15 +0,0 @@
|
||||
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__()
|
||||
@@ -1,152 +0,0 @@
|
||||
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()
|
||||
self._closed_event = asyncio.Event()
|
||||
|
||||
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
|
||||
self._closed_event.set()
|
||||
|
||||
# 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")
|
||||
self._running = False
|
||||
self._closed_event.set()
|
||||
break
|
||||
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
try:
|
||||
# Parse JSON-RPC message
|
||||
message = self._parse_message(line)
|
||||
asyncio.create_task(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:
|
||||
self._closed_event.set()
|
||||
logger.debug("Stopped reading from stdin")
|
||||
|
||||
async def wait_closed(self) -> None:
|
||||
await self._closed_event.wait()
|
||||
|
||||
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}")
|
||||
@@ -1,236 +0,0 @@
|
||||
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)
|
||||
asyncio.create_task(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)
|
||||
asyncio.create_task(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}")
|
||||
@@ -1,48 +0,0 @@
|
||||
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)
|
||||
@@ -1,135 +0,0 @@
|
||||
import asyncio
|
||||
import signal
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import IO, Any
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from .api.context import Context
|
||||
from .rpc.server import StdioServer, WebSocketServer
|
||||
from .star_runner import StarRunner
|
||||
from .stars.star_manager import StarManager
|
||||
from .supervisor import SupervisorRuntime
|
||||
|
||||
|
||||
def _install_signal_handlers(stop_event: asyncio.Event) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
for sig in (signal.SIGTERM, signal.SIGINT):
|
||||
try:
|
||||
loop.add_signal_handler(sig, stop_event.set)
|
||||
except NotImplementedError:
|
||||
logger.debug(f"Signal handlers are not supported for {sig}")
|
||||
|
||||
|
||||
def _prepare_stdio_transport(
|
||||
stdin: IO[Any] | None,
|
||||
stdout: IO[Any] | None,
|
||||
) -> tuple[IO[Any], IO[Any], IO[Any] | None]:
|
||||
if stdin is not None and stdout is not None:
|
||||
return stdin, stdout, None
|
||||
|
||||
transport_stdin = stdin or sys.stdin
|
||||
transport_stdout = stdout or sys.stdout
|
||||
original_stdout = sys.stdout
|
||||
sys.stdout = sys.stderr
|
||||
return transport_stdin, transport_stdout, original_stdout
|
||||
|
||||
|
||||
async def _wait_for_stdio_shutdown(
|
||||
server: StdioServer, stop_event: asyncio.Event
|
||||
) -> None:
|
||||
stop_waiter = asyncio.create_task(stop_event.wait())
|
||||
stdio_waiter = asyncio.create_task(server.wait_closed())
|
||||
done, pending = await asyncio.wait(
|
||||
{stop_waiter, stdio_waiter},
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
for task in done:
|
||||
if task.cancelled():
|
||||
continue
|
||||
task.result()
|
||||
|
||||
|
||||
async def run_websocket_server(
|
||||
host: str = "127.0.0.1",
|
||||
port: int = 8765,
|
||||
path: str = "/",
|
||||
heartbeat_interval: int = 30,
|
||||
plugin_dir: str | Path | None = None,
|
||||
):
|
||||
server = WebSocketServer(
|
||||
port=port, host=host, path=path, heartbeat=heartbeat_interval
|
||||
)
|
||||
runner = StarRunner(server)
|
||||
context = Context.default_context(runner=runner)
|
||||
star_manager = StarManager(context=context)
|
||||
star_manager.discover_star(Path(plugin_dir) if plugin_dir else None)
|
||||
await runner.run()
|
||||
|
||||
stop_event = asyncio.Event()
|
||||
_install_signal_handlers(stop_event)
|
||||
|
||||
logger.info("Server is running. Press Ctrl+C to stop.")
|
||||
|
||||
try:
|
||||
await stop_event.wait()
|
||||
finally:
|
||||
logger.info("Shutting down...")
|
||||
await server.stop()
|
||||
|
||||
|
||||
async def run_supervisor(
|
||||
plugins_dir: str | Path = "plugins",
|
||||
stdin: IO[Any] | None = None,
|
||||
stdout: IO[Any] | None = None,
|
||||
) -> None:
|
||||
transport_stdin, transport_stdout, original_stdout = _prepare_stdio_transport(
|
||||
stdin, stdout
|
||||
)
|
||||
server = StdioServer(stdin=transport_stdin, stdout=transport_stdout)
|
||||
supervisor = SupervisorRuntime(
|
||||
server=server,
|
||||
plugins_dir=Path(plugins_dir),
|
||||
)
|
||||
|
||||
try:
|
||||
await supervisor.start()
|
||||
stop_event = asyncio.Event()
|
||||
_install_signal_handlers(stop_event)
|
||||
logger.info(f"Plugin supervisor is running with plugins dir: {plugins_dir}")
|
||||
await _wait_for_stdio_shutdown(server, stop_event)
|
||||
finally:
|
||||
logger.info("Shutting down plugin supervisor...")
|
||||
await supervisor.stop()
|
||||
if original_stdout is not None:
|
||||
sys.stdout = original_stdout
|
||||
|
||||
|
||||
async def run_plugin_worker(
|
||||
plugin_dir: str | Path,
|
||||
stdin: IO[Any] | None = None,
|
||||
stdout: IO[Any] | None = None,
|
||||
) -> None:
|
||||
transport_stdin, transport_stdout, original_stdout = _prepare_stdio_transport(
|
||||
stdin, stdout
|
||||
)
|
||||
server = StdioServer(stdin=transport_stdin, stdout=transport_stdout)
|
||||
runner = StarRunner(server)
|
||||
context = Context.default_context(runner=runner)
|
||||
star_manager = StarManager(context=context)
|
||||
star_manager.discover_star(Path(plugin_dir))
|
||||
|
||||
try:
|
||||
await runner.run()
|
||||
stop_event = asyncio.Event()
|
||||
_install_signal_handlers(stop_event)
|
||||
logger.info(f"Plugin worker is running for: {plugin_dir}")
|
||||
await _wait_for_stdio_shutdown(server, stop_event)
|
||||
finally:
|
||||
logger.info("Shutting down plugin worker...")
|
||||
await runner.stop()
|
||||
if original_stdout is not None:
|
||||
sys.stdout = original_stdout
|
||||
@@ -1,202 +0,0 @@
|
||||
import inspect
|
||||
from loguru import logger
|
||||
from typing import Any
|
||||
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 .rpc.request_helper import RPCRequestHelper
|
||||
from .types import (
|
||||
CallHandlerRequest,
|
||||
HandlerStreamStartNotification,
|
||||
HandlerStreamUpdateNotification,
|
||||
HandlerStreamEndNotification,
|
||||
)
|
||||
|
||||
|
||||
class HandshakeHandler:
|
||||
"""Handles the handshake protocol to exchange plugin metadata."""
|
||||
|
||||
async def handle(self, message: JSONRPCRequest) -> JSONRPCSuccessResponse:
|
||||
"""Build and return handshake response with plugin metadata."""
|
||||
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
|
||||
|
||||
return JSONRPCSuccessResponse(
|
||||
jsonrpc="2.0",
|
||||
id=message.id,
|
||||
result=payload,
|
||||
)
|
||||
|
||||
|
||||
class HandlerExecutor:
|
||||
"""Executes plugin handlers and manages streaming results."""
|
||||
|
||||
def __init__(self, rpc_request_helper: RPCRequestHelper):
|
||||
self.rpc_request_helper = rpc_request_helper
|
||||
|
||||
async def execute(self, message: JSONRPCRequest, server: JSONRPCServer):
|
||||
"""Execute a handler and stream results back to the caller."""
|
||||
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()
|
||||
|
||||
handler = star_handlers_registry.get_handler_by_full_name(handler_full_name)
|
||||
|
||||
if handler is None:
|
||||
await self._send_error(
|
||||
server, message.id, -32601, f"Handler not found: {handler_full_name}"
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
await self._execute_and_stream(
|
||||
server, message.id, handler_full_name, handler.handler(event, **args)
|
||||
)
|
||||
except Exception as e:
|
||||
await self._send_error(server, message.id, -32000, str(e))
|
||||
|
||||
async def _execute_and_stream(
|
||||
self,
|
||||
server: JSONRPCServer,
|
||||
request_id: str | None,
|
||||
handler_name: str,
|
||||
ready_to_call,
|
||||
):
|
||||
"""Execute handler and stream results."""
|
||||
# Send start notification
|
||||
await server.send_message(
|
||||
HandlerStreamStartNotification(
|
||||
jsonrpc="2.0",
|
||||
method="handler_stream_start",
|
||||
params=HandlerStreamStartNotification.Params(
|
||||
id=request_id,
|
||||
handler_full_name=handler_name,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
if inspect.iscoroutine(ready_to_call):
|
||||
result = await ready_to_call
|
||||
# Send update notification
|
||||
await server.send_message(
|
||||
HandlerStreamUpdateNotification(
|
||||
jsonrpc="2.0",
|
||||
method="handler_stream_update",
|
||||
params=HandlerStreamUpdateNotification.Params(
|
||||
id=request_id,
|
||||
handler_full_name=handler_name,
|
||||
data=result,
|
||||
),
|
||||
)
|
||||
)
|
||||
elif inspect.isasyncgen(ready_to_call):
|
||||
async for ret in ready_to_call:
|
||||
# Send update notification for each item
|
||||
await server.send_message(
|
||||
HandlerStreamUpdateNotification(
|
||||
jsonrpc="2.0",
|
||||
method="handler_stream_update",
|
||||
params=HandlerStreamUpdateNotification.Params(
|
||||
id=request_id,
|
||||
handler_full_name=handler_name,
|
||||
data=ret,
|
||||
),
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during handler {handler_name}: {e}")
|
||||
finally:
|
||||
# Send end notification
|
||||
await server.send_message(
|
||||
HandlerStreamEndNotification(
|
||||
jsonrpc="2.0",
|
||||
method="handler_stream_end",
|
||||
params=HandlerStreamEndNotification.Params(
|
||||
id=request_id,
|
||||
handler_full_name=handler_name,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
async def _send_error(
|
||||
self, server: JSONRPCServer, request_id: str | None, code: int, message: str
|
||||
):
|
||||
"""Send an error response."""
|
||||
response = JSONRPCErrorResponse(
|
||||
jsonrpc="2.0",
|
||||
id=request_id,
|
||||
error=JSONRPCErrorData(code=code, message=message),
|
||||
)
|
||||
await server.send_message(response)
|
||||
|
||||
|
||||
class StarRunner:
|
||||
"""Main runner to handle RPC messages and route them to handlers."""
|
||||
|
||||
def __init__(self, server: JSONRPCServer):
|
||||
self.server = server
|
||||
|
||||
self.rpc_request_helper = RPCRequestHelper()
|
||||
self.handler_executor = HandlerExecutor(self.rpc_request_helper)
|
||||
self.handshake_handler = HandshakeHandler()
|
||||
|
||||
async def call_context_function(
|
||||
self, method_name: str, params: dict[str, Any]
|
||||
) -> dict[str, Any]:
|
||||
result = await self.rpc_request_helper.call_rpc(
|
||||
self.server,
|
||||
JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=self.rpc_request_helper._generate_request_id(),
|
||||
method="call_context_function",
|
||||
params={
|
||||
"name": method_name,
|
||||
"args": params,
|
||||
},
|
||||
),
|
||||
)
|
||||
if isinstance(result, JSONRPCSuccessResponse):
|
||||
return result.result
|
||||
elif isinstance(result, JSONRPCErrorResponse):
|
||||
raise Exception(f"RPC Error {result.error.code}: {result.error.message}")
|
||||
else:
|
||||
raise Exception("Invalid RPC response")
|
||||
|
||||
async def _handle_messages(self, message: JSONRPCMessage):
|
||||
"""Route messages to appropriate handlers."""
|
||||
if isinstance(message, JSONRPCRequest):
|
||||
if message.method == "handshake":
|
||||
response = await self.handshake_handler.handle(message)
|
||||
await self.server.send_message(response)
|
||||
elif message.method == "call_handler":
|
||||
await self.handler_executor.execute(message, self.server)
|
||||
else:
|
||||
logger.warning(f"Unknown method from client: {message.method}")
|
||||
elif isinstance(message, (JSONRPCSuccessResponse, JSONRPCErrorResponse)):
|
||||
self.rpc_request_helper.resolve_pending_request(message)
|
||||
|
||||
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()
|
||||
@@ -1,14 +0,0 @@
|
||||
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"]
|
||||
@@ -1,218 +0,0 @@
|
||||
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
|
||||
@@ -1,133 +0,0 @@
|
||||
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)
|
||||
@@ -1,61 +0,0 @@
|
||||
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)
|
||||
@@ -1,33 +0,0 @@
|
||||
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
|
||||
@@ -1,29 +0,0 @@
|
||||
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
|
||||
@@ -1,71 +0,0 @@
|
||||
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
|
||||
@@ -1,18 +0,0 @@
|
||||
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()))
|
||||
@@ -1,248 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import inspect
|
||||
from pathlib import Path
|
||||
from typing import Any, AsyncGenerator
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from ...api.event.astr_message_event import AstrMessageEvent
|
||||
from ...api.star.star import StarMetadata
|
||||
from .registry import EventType, StarHandlerMetadata
|
||||
from ..rpc.jsonrpc import (
|
||||
JSONRPCErrorResponse,
|
||||
JSONRPCMessage,
|
||||
JSONRPCRequest,
|
||||
JSONRPCSuccessResponse,
|
||||
)
|
||||
from ..rpc.client import JSONRPCClient
|
||||
from ..rpc.client.stdio import StdioClient
|
||||
from ..rpc.client.websocket import WebSocketClient
|
||||
from ..rpc.request_helper import RPCRequestHelper
|
||||
from .virtual import VirtualStar
|
||||
from .new_star_utils import (
|
||||
ClientHandshakeHandler,
|
||||
PluginRequestHandler,
|
||||
HandlerProxyFactory,
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
context: Any,
|
||||
) -> None:
|
||||
"""Initialize a NewStar instance.
|
||||
|
||||
Args:
|
||||
client: JSON-RPC client for communication
|
||||
context: Context instance for managing managers and their functions
|
||||
"""
|
||||
super().__init__(context)
|
||||
|
||||
self._client = client
|
||||
self._metadata: dict[str, StarMetadata] = {}
|
||||
self._handlers: list[StarHandlerMetadata] = []
|
||||
self._active = False
|
||||
|
||||
# Use RPCRequestHelper for managing requests
|
||||
self._rpc_helper = RPCRequestHelper()
|
||||
|
||||
# Initialize specialized handlers
|
||||
self._handshake_handler = ClientHandshakeHandler(self._rpc_helper)
|
||||
self._plugin_request_handler = PluginRequestHandler(context)
|
||||
self._handler_proxy_factory = HandlerProxyFactory(client, self._rpc_helper)
|
||||
|
||||
# Set up message handler
|
||||
self._client.set_message_handler(self._handle_message)
|
||||
|
||||
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,
|
||||
):
|
||||
# Delegate to RPCRequestHelper
|
||||
self._rpc_helper.resolve_pending_request(message)
|
||||
|
||||
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._rpc_helper.handle_stream_notification(message)
|
||||
else:
|
||||
# Plugin is calling a method on the core - delegate to PluginRequestHandler
|
||||
asyncio.create_task(
|
||||
self._plugin_request_handler.handle_request(message, self._client)
|
||||
)
|
||||
|
||||
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.
|
||||
"""
|
||||
# Delegate to ClientHandshakeHandler
|
||||
(
|
||||
self._metadata,
|
||||
self._handlers,
|
||||
) = await self._handshake_handler.perform_handshake(self._client)
|
||||
|
||||
# Set up handler proxies
|
||||
self._handler_proxy_factory.setup_handlers(self._handlers)
|
||||
|
||||
return self._metadata
|
||||
|
||||
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,
|
||||
) -> AsyncGenerator[Any, 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:
|
||||
An async generator yielding results from the handler
|
||||
"""
|
||||
logger.debug(f"Calling handler: {handler.handler_name}")
|
||||
|
||||
# Call the handler proxy
|
||||
assert inspect.isasyncgenfunction(handler.handler), (
|
||||
"Handler proxy must be an async generator function"
|
||||
)
|
||||
async for result in handler.handler(event, **kwargs):
|
||||
yield result
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the NewStar and cleanup resources."""
|
||||
await self._client.stop()
|
||||
logger.info("NewStar client stopped.")
|
||||
|
||||
|
||||
class NewStdioStar(NewStar):
|
||||
"""NewStar implementation using STDIO communication.
|
||||
|
||||
This class automatically starts the plugin subprocess and manages its lifecycle.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
plugins_dir: str,
|
||||
python_executable: str = "python",
|
||||
context: Any = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize a STDIO-based NewStar.
|
||||
|
||||
Args:
|
||||
plugins_dir: Path to the plugins directory
|
||||
python_executable: Python executable to use (defaults to 'python')
|
||||
context: Context instance for managing managers and their functions
|
||||
"""
|
||||
if not os.path.exists(plugins_dir):
|
||||
raise FileNotFoundError(f"Plugins directory not found: {plugins_dir}")
|
||||
|
||||
repo_src_dir = str(Path(__file__).resolve().parents[3])
|
||||
env = os.environ.copy()
|
||||
existing_pythonpath = env.get("PYTHONPATH")
|
||||
env["PYTHONPATH"] = (
|
||||
f"{repo_src_dir}{os.pathsep}{existing_pythonpath}"
|
||||
if existing_pythonpath
|
||||
else repo_src_dir
|
||||
)
|
||||
|
||||
command = [
|
||||
python_executable,
|
||||
"-m",
|
||||
"astrbot_sdk",
|
||||
"run",
|
||||
"--plugins-dir",
|
||||
plugins_dir,
|
||||
]
|
||||
|
||||
# Create StdioClient with subprocess management
|
||||
client = StdioClient(command=command, cwd=plugins_dir, env=env)
|
||||
super().__init__(client, context=context)
|
||||
|
||||
|
||||
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,
|
||||
context: Any = None,
|
||||
**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
|
||||
context: Context instance for managing managers and their functions
|
||||
"""
|
||||
client = WebSocketClient(
|
||||
url=url, heartbeat=heartbeat, reconnect_interval=reconnect_interval
|
||||
)
|
||||
super().__init__(client, context=context)
|
||||
self._url = url
|
||||
self._heartbeat = heartbeat
|
||||
self._reconnect_interval = reconnect_interval
|
||||
@@ -1,266 +0,0 @@
|
||||
import asyncio
|
||||
import inspect
|
||||
from typing import Any, AsyncGenerator
|
||||
from loguru import logger
|
||||
|
||||
from ...api.event.astr_message_event import AstrMessageEvent, AstrMessageEventModel
|
||||
from ...api.star.star import StarMetadata
|
||||
from ...api.star.context import Context
|
||||
from ..rpc.client import JSONRPCClient
|
||||
from ..rpc.request_helper import RPCRequestHelper
|
||||
from ..rpc.jsonrpc import (
|
||||
JSONRPCSuccessResponse,
|
||||
JSONRPCRequest,
|
||||
JSONRPCErrorResponse,
|
||||
JSONRPCErrorData,
|
||||
)
|
||||
from ..types import CallHandlerRequest, HandshakeRequest
|
||||
from .registry import StarHandlerMetadata, EventType
|
||||
|
||||
|
||||
class HandlerProxyFactory:
|
||||
"""Creates proxy functions for remote handler invocation."""
|
||||
|
||||
def __init__(self, client: JSONRPCClient, rpc_helper: RPCRequestHelper):
|
||||
"""Initialize the handler proxy factory.
|
||||
|
||||
Args:
|
||||
client: JSON-RPC client for communication
|
||||
rpc_helper: RPC request helper for making RPC calls
|
||||
"""
|
||||
self._client = client
|
||||
self._rpc_helper = rpc_helper
|
||||
|
||||
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 generator function that proxies calls to the remote handler.
|
||||
"""
|
||||
|
||||
async def handler_proxy(
|
||||
event: AstrMessageEvent, **kwargs
|
||||
) -> AsyncGenerator[Any, None]:
|
||||
"""Proxy function for remote handler invocation.
|
||||
|
||||
Yields results from the remote handler via streaming.
|
||||
"""
|
||||
request_id = self._rpc_helper._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,
|
||||
),
|
||||
)
|
||||
queue = await self._rpc_helper.call_rpc_streaming(self._client, request)
|
||||
|
||||
try:
|
||||
while True:
|
||||
item = await asyncio.wait_for(queue.get(), timeout=30.0)
|
||||
if isinstance(item, dict) and item.get("_stream_end"):
|
||||
break
|
||||
if isinstance(item, dict) and item.get("_error"):
|
||||
raise RuntimeError(item.get("message", "Unknown error"))
|
||||
yield self._deserialize_result(item)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
raise RuntimeError(f"RPC call to {handler_full_name} timed out")
|
||||
|
||||
return handler_proxy
|
||||
|
||||
def setup_handlers(self, handlers: list[StarHandlerMetadata]) -> None:
|
||||
"""Set up handler proxies for all handlers.
|
||||
|
||||
Args:
|
||||
handlers: List of handler metadata to set up
|
||||
"""
|
||||
for handler in handlers:
|
||||
handler.handler = self.create_handler_proxy(handler.handler_full_name)
|
||||
logger.info(f"Set up {len(handlers)} handler proxies")
|
||||
|
||||
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, we might want to reconstruct MessageEventResult etc.
|
||||
return result
|
||||
|
||||
|
||||
class ClientHandshakeHandler:
|
||||
"""Handles the handshake protocol to retrieve plugin metadata."""
|
||||
|
||||
def __init__(self, rpc_helper: RPCRequestHelper):
|
||||
"""Initialize the handshake handler.
|
||||
|
||||
Args:
|
||||
rpc_helper: RPC request helper for making RPC calls
|
||||
"""
|
||||
self._rpc_helper = rpc_helper
|
||||
|
||||
async def perform_handshake(
|
||||
self, client: JSONRPCClient
|
||||
) -> tuple[dict[str, StarMetadata], list[StarHandlerMetadata]]:
|
||||
"""Perform handshake to retrieve plugin metadata.
|
||||
|
||||
Args:
|
||||
client: JSON-RPC client for communication
|
||||
|
||||
Returns:
|
||||
Tuple of (metadata dict, handlers list)
|
||||
|
||||
Raises:
|
||||
RuntimeError: If handshake fails
|
||||
"""
|
||||
logger.info("Performing handshake with plugin...")
|
||||
|
||||
response = await self._rpc_helper.call_rpc(
|
||||
client,
|
||||
HandshakeRequest(
|
||||
jsonrpc="2.0",
|
||||
id=self._rpc_helper._generate_request_id(),
|
||||
method="handshake",
|
||||
),
|
||||
)
|
||||
|
||||
if not isinstance(response, JSONRPCSuccessResponse):
|
||||
raise RuntimeError("Handshake failed: Invalid response from plugin")
|
||||
|
||||
result = response.result
|
||||
|
||||
if not isinstance(result, dict):
|
||||
raise RuntimeError("Handshake failed: Invalid response from plugin")
|
||||
|
||||
metadata_dict: dict[str, StarMetadata] = {}
|
||||
handlers_list: list[StarHandlerMetadata] = []
|
||||
|
||||
# Placeholder handler that will be replaced later
|
||||
def _placeholder_handler(*args, **kwargs):
|
||||
raise NotImplementedError("Handler proxy not set up yet")
|
||||
|
||||
# Parse metadata
|
||||
for star_name, star_info in result.items():
|
||||
handlers_data = star_info.pop("handlers", None)
|
||||
metadata = StarMetadata(**star_info)
|
||||
metadata_dict[star_name] = metadata
|
||||
|
||||
# Parse handlers
|
||||
if handlers_data:
|
||||
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=_placeholder_handler, # Will be replaced by HandlerProxyFactory
|
||||
event_filters=[],
|
||||
desc=handler_data.get("desc", ""),
|
||||
extras_configs=handler_data.get("extras_configs", {}),
|
||||
)
|
||||
handlers_list.append(handler_meta)
|
||||
|
||||
logger.info(
|
||||
f"Handshake complete: {len(metadata_dict)} stars loaded, "
|
||||
f"{metadata_dict.keys()}, {len(handlers_list)} handlers registered."
|
||||
)
|
||||
|
||||
return metadata_dict, handlers_list
|
||||
|
||||
|
||||
class PluginRequestHandler:
|
||||
"""Handles JSON-RPC requests from plugins calling core methods."""
|
||||
|
||||
def __init__(self, context: Context):
|
||||
"""Initialize the plugin request handler.
|
||||
|
||||
Args:
|
||||
context: Context instance for managing managers and their functions
|
||||
"""
|
||||
self._context = context
|
||||
|
||||
async def handle_request(
|
||||
self, request: JSONRPCRequest, client: JSONRPCClient
|
||||
) -> None:
|
||||
"""Handle a JSON-RPC request from the plugin.
|
||||
|
||||
Args:
|
||||
request: The JSON-RPC request from the plugin
|
||||
client: The client to send response back
|
||||
"""
|
||||
result: Any = None
|
||||
try:
|
||||
method = request.method
|
||||
params = request.params or {}
|
||||
|
||||
if method == "call_context_function":
|
||||
result = await self._handle_context_function_call(params)
|
||||
else:
|
||||
raise ValueError(f"Unknown method: {method}")
|
||||
|
||||
# Send success response
|
||||
response = JSONRPCSuccessResponse(
|
||||
jsonrpc="2.0",
|
||||
id=request.id,
|
||||
result={
|
||||
"data": result,
|
||||
},
|
||||
)
|
||||
await 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 client.send_message(error_response)
|
||||
|
||||
async def _handle_context_function_call(self, params: dict) -> Any:
|
||||
"""Handle call_context_function requests.
|
||||
|
||||
Args:
|
||||
params: Request parameters containing function name and args
|
||||
|
||||
Returns:
|
||||
Result from the function call
|
||||
|
||||
Raises:
|
||||
ValueError: If function is not found
|
||||
"""
|
||||
func_full_name = params.get("name", "")
|
||||
args = params.get("args", {})
|
||||
|
||||
logger.debug(
|
||||
f"Plugin called call_context_function: {func_full_name} with args: {args}"
|
||||
)
|
||||
|
||||
# Get the registered function from context
|
||||
func = self._context._registered_functions.get(func_full_name)
|
||||
if func is None:
|
||||
raise ValueError(f"Function not found: {func_full_name}")
|
||||
|
||||
# Call the function
|
||||
if inspect.iscoroutinefunction(func):
|
||||
result = await func(**args)
|
||||
else:
|
||||
result = func(**args)
|
||||
|
||||
logger.debug(f"call_context_function result: {result}")
|
||||
return result
|
||||
@@ -1,182 +0,0 @@
|
||||
from __future__ import annotations
|
||||
import enum
|
||||
from collections.abc import Awaitable, Callable, AsyncGenerator
|
||||
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] | AsyncGenerator[Any, None]]
|
||||
"""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] = []
|
||||
@@ -1,515 +0,0 @@
|
||||
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
|
||||
@@ -1,108 +0,0 @@
|
||||
import yaml
|
||||
import importlib
|
||||
import functools
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from loguru import logger
|
||||
from .registry import star_handlers_registry, star_map, star_registry
|
||||
from ..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()
|
||||
else:
|
||||
root_dir = Path(root_dir).resolve()
|
||||
|
||||
path = root_dir / "plugin.yaml"
|
||||
if not path.exists():
|
||||
logger.warning("No plugin.yaml found in the current directory.")
|
||||
return []
|
||||
|
||||
# Add the plugin directory to sys.path so we can import its modules
|
||||
root_dir_str = str(root_dir)
|
||||
if root_dir_str not in sys.path:
|
||||
sys.path.insert(0, root_dir_str)
|
||||
logger.debug(f"Added {root_dir_str} to sys.path")
|
||||
|
||||
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", "")
|
||||
logger.debug(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:
|
||||
logger.debug(f"Importing module: {module_path}")
|
||||
module_type = importlib.import_module(module_path)
|
||||
logger.debug(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"
|
||||
)
|
||||
@@ -1,125 +0,0 @@
|
||||
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
|
||||
from ...api.star.context import Context
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
def __init__(self, context: Context) -> None:
|
||||
self._context = context
|
||||
|
||||
@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.AsyncGenerator[T.Any, None]:
|
||||
"""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:
|
||||
An async generator yielding results from the handler
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the handler call fails or times out
|
||||
"""
|
||||
...
|
||||
@@ -1,564 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
||||
import yaml
|
||||
from loguru import logger
|
||||
from .rpc.client.stdio import StdioClient
|
||||
from .rpc.jsonrpc import (
|
||||
JSONRPCErrorData,
|
||||
JSONRPCErrorResponse,
|
||||
JSONRPCMessage,
|
||||
JSONRPCRequest,
|
||||
JSONRPCSuccessResponse,
|
||||
)
|
||||
from .rpc.request_helper import RPCRequestHelper
|
||||
from .rpc.server.base import JSONRPCServer
|
||||
from .stars.registry import EventType, StarHandlerMetadata
|
||||
from .types import CallHandlerRequest, HandshakeRequest
|
||||
|
||||
STATE_FILE_NAME = ".astrbot-worker-state.json"
|
||||
|
||||
|
||||
def _venv_python_path(venv_dir: Path) -> Path:
|
||||
if os.name == "nt":
|
||||
return venv_dir / "Scripts" / "python.exe"
|
||||
return venv_dir / "bin" / "python"
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PluginSpec:
|
||||
name: str
|
||||
plugin_dir: Path
|
||||
manifest_path: Path
|
||||
requirements_path: Path
|
||||
python_version: str
|
||||
manifest_data: dict[str, Any]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PluginDiscoveryResult:
|
||||
plugins: list[PluginSpec]
|
||||
skipped_plugins: dict[str, str]
|
||||
|
||||
|
||||
def discover_plugins(plugins_dir: Path) -> PluginDiscoveryResult:
|
||||
plugins_root = plugins_dir.resolve()
|
||||
skipped_plugins: dict[str, str] = {}
|
||||
plugins: list[PluginSpec] = []
|
||||
seen_names: set[str] = set()
|
||||
|
||||
if not plugins_root.exists():
|
||||
logger.warning(f"Plugins directory does not exist: {plugins_root}")
|
||||
return PluginDiscoveryResult([], {})
|
||||
|
||||
for entry in sorted(plugins_root.iterdir()):
|
||||
if not entry.is_dir() or entry.name.startswith("."):
|
||||
continue
|
||||
|
||||
manifest_path = entry / "plugin.yaml"
|
||||
requirements_path = entry / "requirements.txt"
|
||||
if not manifest_path.exists():
|
||||
logger.warning(f"Skipping {entry}: missing plugin.yaml")
|
||||
continue
|
||||
if not requirements_path.exists():
|
||||
logger.warning(f"Skipping {entry}: missing requirements.txt")
|
||||
skipped_plugins[entry.name] = "missing requirements.txt"
|
||||
continue
|
||||
|
||||
try:
|
||||
manifest_data = (
|
||||
yaml.safe_load(manifest_path.read_text(encoding="utf-8")) or {}
|
||||
)
|
||||
except Exception as exc:
|
||||
skipped_plugins[entry.name] = f"failed to parse plugin.yaml: {exc}"
|
||||
continue
|
||||
|
||||
plugin_name = manifest_data.get("name")
|
||||
components = manifest_data.get("components")
|
||||
runtime = manifest_data.get("runtime") or {}
|
||||
python_version = runtime.get("python")
|
||||
|
||||
if not isinstance(plugin_name, str) or not plugin_name:
|
||||
skipped_plugins[entry.name] = "plugin name is required"
|
||||
continue
|
||||
if plugin_name in seen_names:
|
||||
skipped_plugins[plugin_name] = "duplicate plugin name"
|
||||
continue
|
||||
if not isinstance(components, list) or not components:
|
||||
skipped_plugins[plugin_name] = "components must be a non-empty list"
|
||||
continue
|
||||
if not isinstance(python_version, str) or not python_version:
|
||||
skipped_plugins[plugin_name] = "runtime.python is required"
|
||||
continue
|
||||
|
||||
seen_names.add(plugin_name)
|
||||
plugins.append(
|
||||
PluginSpec(
|
||||
name=plugin_name,
|
||||
plugin_dir=entry.resolve(),
|
||||
manifest_path=manifest_path.resolve(),
|
||||
requirements_path=requirements_path.resolve(),
|
||||
python_version=python_version,
|
||||
manifest_data=manifest_data,
|
||||
)
|
||||
)
|
||||
|
||||
return PluginDiscoveryResult(plugins=plugins, skipped_plugins=skipped_plugins)
|
||||
|
||||
|
||||
class PluginEnvironmentManager:
|
||||
def __init__(self, repo_root: Path, uv_binary: str | None = None) -> None:
|
||||
self.repo_root = repo_root.resolve()
|
||||
self.uv_binary = uv_binary or shutil.which("uv")
|
||||
self.cache_dir = self.repo_root / ".uv-cache"
|
||||
|
||||
def prepare_environment(self, plugin: PluginSpec) -> Path:
|
||||
if not self.uv_binary:
|
||||
raise RuntimeError("uv executable not found")
|
||||
|
||||
state_path = plugin.plugin_dir / STATE_FILE_NAME
|
||||
venv_dir = plugin.plugin_dir / ".venv"
|
||||
python_path = _venv_python_path(venv_dir)
|
||||
fingerprint = self._fingerprint(plugin)
|
||||
state = self._load_state(state_path)
|
||||
|
||||
if (
|
||||
not python_path.exists()
|
||||
or not self._matches_python_version(venv_dir, plugin.python_version)
|
||||
or state.get("fingerprint") != fingerprint
|
||||
):
|
||||
self._rebuild(plugin, venv_dir, python_path)
|
||||
self._write_state(state_path, plugin, fingerprint)
|
||||
|
||||
return python_path
|
||||
|
||||
def _rebuild(self, plugin: PluginSpec, venv_dir: Path, python_path: Path) -> None:
|
||||
if venv_dir.exists():
|
||||
shutil.rmtree(venv_dir)
|
||||
|
||||
venv_dir.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._run_command(
|
||||
[
|
||||
self.uv_binary,
|
||||
"venv",
|
||||
"--python",
|
||||
plugin.python_version,
|
||||
"--system-site-packages",
|
||||
"--no-python-downloads",
|
||||
"--no-managed-python",
|
||||
str(venv_dir),
|
||||
],
|
||||
cwd=self.repo_root,
|
||||
command_name=f"create venv for {plugin.name}",
|
||||
)
|
||||
|
||||
requirements_text = plugin.requirements_path.read_text(encoding="utf-8").strip()
|
||||
if not requirements_text:
|
||||
return
|
||||
|
||||
self._run_command(
|
||||
[
|
||||
self.uv_binary,
|
||||
"pip",
|
||||
"install",
|
||||
"--python",
|
||||
str(python_path),
|
||||
"-r",
|
||||
str(plugin.requirements_path),
|
||||
],
|
||||
cwd=plugin.plugin_dir,
|
||||
command_name=f"install requirements for {plugin.name}",
|
||||
)
|
||||
|
||||
def _run_command(
|
||||
self,
|
||||
command: list[str],
|
||||
*,
|
||||
cwd: Path,
|
||||
command_name: str,
|
||||
) -> None:
|
||||
logger.info(f"{command_name}: {' '.join(command)}")
|
||||
process = subprocess.run(
|
||||
command,
|
||||
cwd=str(cwd),
|
||||
env={**os.environ, "UV_CACHE_DIR": str(self.cache_dir)},
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
if process.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"{command_name} failed with exit code {process.returncode}: "
|
||||
f"{process.stderr.strip() or process.stdout.strip()}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _fingerprint(plugin: PluginSpec) -> str:
|
||||
requirements = plugin.requirements_path.read_text(encoding="utf-8")
|
||||
payload = {
|
||||
"python_version": plugin.python_version,
|
||||
"requirements": requirements,
|
||||
}
|
||||
return json.dumps(payload, ensure_ascii=True, sort_keys=True)
|
||||
|
||||
@staticmethod
|
||||
def _load_state(state_path: Path) -> dict[str, Any]:
|
||||
if not state_path.exists():
|
||||
return {}
|
||||
try:
|
||||
data = json.loads(state_path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return {}
|
||||
return data if isinstance(data, dict) else {}
|
||||
|
||||
@staticmethod
|
||||
def _write_state(state_path: Path, plugin: PluginSpec, fingerprint: str) -> None:
|
||||
state_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"plugin": plugin.name,
|
||||
"python_version": plugin.python_version,
|
||||
"fingerprint": fingerprint,
|
||||
},
|
||||
ensure_ascii=True,
|
||||
indent=2,
|
||||
sort_keys=True,
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _matches_python_version(venv_dir: Path, version: str) -> bool:
|
||||
pyvenv_cfg = venv_dir / "pyvenv.cfg"
|
||||
if not pyvenv_cfg.exists():
|
||||
return False
|
||||
try:
|
||||
content = pyvenv_cfg.read_text(encoding="utf-8")
|
||||
except OSError:
|
||||
return False
|
||||
match = re.search(r"version\s*=\s*(\d+\.\d+)\.\d+", content, re.IGNORECASE)
|
||||
return match is not None and match.group(1) == version
|
||||
|
||||
|
||||
class WorkerRuntime:
|
||||
def __init__(
|
||||
self,
|
||||
plugin: PluginSpec,
|
||||
server: JSONRPCServer,
|
||||
repo_root: Path,
|
||||
env_manager: PluginEnvironmentManager,
|
||||
) -> None:
|
||||
self.plugin = plugin
|
||||
self.server = server
|
||||
self.repo_root = repo_root.resolve()
|
||||
self.env_manager = env_manager
|
||||
self.rpc_helper = RPCRequestHelper()
|
||||
self.client: StdioClient | None = None
|
||||
self.raw_handshake: dict[str, Any] = {}
|
||||
self.handlers: list[StarHandlerMetadata] = []
|
||||
self._context_requests: dict[str, str] = {}
|
||||
self._forwarded_call_ids: set[str] = set()
|
||||
|
||||
async def start(self) -> None:
|
||||
python_path = self.env_manager.prepare_environment(self.plugin)
|
||||
repo_src_dir = str(self.repo_root / "src")
|
||||
env = os.environ.copy()
|
||||
existing_pythonpath = env.get("PYTHONPATH")
|
||||
env["PYTHONPATH"] = (
|
||||
f"{repo_src_dir}{os.pathsep}{existing_pythonpath}"
|
||||
if existing_pythonpath
|
||||
else repo_src_dir
|
||||
)
|
||||
|
||||
self.client = StdioClient(
|
||||
command=[
|
||||
str(python_path),
|
||||
"-m",
|
||||
"astrbot_sdk",
|
||||
"worker",
|
||||
"--plugin-dir",
|
||||
str(self.plugin.plugin_dir),
|
||||
],
|
||||
cwd=str(self.plugin.plugin_dir),
|
||||
env=env,
|
||||
)
|
||||
self.client.set_message_handler(self._handle_message)
|
||||
await self.client.start()
|
||||
|
||||
response = await asyncio.wait_for(
|
||||
self.rpc_helper.call_rpc(
|
||||
self.client,
|
||||
HandshakeRequest(
|
||||
jsonrpc="2.0",
|
||||
id=self.rpc_helper._generate_request_id(),
|
||||
method="handshake",
|
||||
),
|
||||
),
|
||||
timeout=60.0,
|
||||
)
|
||||
if not isinstance(response, JSONRPCSuccessResponse):
|
||||
raise RuntimeError(f"Handshake failed for plugin {self.plugin.name}")
|
||||
|
||||
result = response.result
|
||||
if not isinstance(result, dict):
|
||||
raise RuntimeError(
|
||||
f"Invalid handshake payload for plugin {self.plugin.name}"
|
||||
)
|
||||
|
||||
self.raw_handshake = result
|
||||
self.handlers = self._parse_handlers(result)
|
||||
|
||||
async def stop(self) -> None:
|
||||
if self.client is not None:
|
||||
await self.client.stop()
|
||||
|
||||
async def forward_call_handler(self, request: JSONRPCRequest) -> None:
|
||||
if self.client is None:
|
||||
raise RuntimeError(f"Worker for {self.plugin.name} is not running")
|
||||
if request.id is not None:
|
||||
self._forwarded_call_ids.add(str(request.id))
|
||||
await self.client.send_message(request)
|
||||
|
||||
async def handle_context_response(
|
||||
self, message: JSONRPCSuccessResponse | JSONRPCErrorResponse
|
||||
) -> bool:
|
||||
message_id = str(message.id)
|
||||
worker_request_id = self._context_requests.pop(message_id, None)
|
||||
if worker_request_id is None:
|
||||
return False
|
||||
if self.client is None:
|
||||
return True
|
||||
|
||||
if isinstance(message, JSONRPCSuccessResponse):
|
||||
await self.client.send_message(
|
||||
JSONRPCSuccessResponse(
|
||||
jsonrpc="2.0",
|
||||
id=worker_request_id,
|
||||
result=message.result,
|
||||
)
|
||||
)
|
||||
else:
|
||||
await self.client.send_message(
|
||||
JSONRPCErrorResponse(
|
||||
jsonrpc="2.0",
|
||||
id=worker_request_id,
|
||||
error=message.error,
|
||||
)
|
||||
)
|
||||
return True
|
||||
|
||||
async def _handle_message(self, message: JSONRPCMessage) -> None:
|
||||
if isinstance(message, (JSONRPCSuccessResponse, JSONRPCErrorResponse)):
|
||||
if message.id in self.rpc_helper.pending_requests:
|
||||
self.rpc_helper.resolve_pending_request(message)
|
||||
return
|
||||
|
||||
if message.id is not None and str(message.id) in self._forwarded_call_ids:
|
||||
self._forwarded_call_ids.discard(str(message.id))
|
||||
await self.server.send_message(message)
|
||||
return
|
||||
|
||||
if not isinstance(message, JSONRPCRequest):
|
||||
return
|
||||
|
||||
if message.method in [
|
||||
"handler_stream_start",
|
||||
"handler_stream_update",
|
||||
"handler_stream_end",
|
||||
]:
|
||||
await self.server.send_message(message)
|
||||
return
|
||||
|
||||
if message.method != "call_context_function":
|
||||
logger.warning(
|
||||
f"Worker {self.plugin.name} sent unknown request: {message.method}"
|
||||
)
|
||||
return
|
||||
|
||||
supervisor_request_id = (
|
||||
f"ctx:{self.plugin.name}:{message.id}"
|
||||
if message.id is not None
|
||||
else f"ctx:{self.plugin.name}:none"
|
||||
)
|
||||
self._context_requests[supervisor_request_id] = str(message.id)
|
||||
await self.server.send_message(
|
||||
JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
id=supervisor_request_id,
|
||||
method=message.method,
|
||||
params=message.params,
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _parse_handlers(handshake_payload: dict[str, Any]) -> list[StarHandlerMetadata]:
|
||||
handlers: list[StarHandlerMetadata] = []
|
||||
|
||||
def _placeholder_handler(*args, **kwargs):
|
||||
raise NotImplementedError("Worker supervisor does not execute handlers")
|
||||
|
||||
for star_info in handshake_payload.values():
|
||||
handlers_data = star_info.get("handlers") or []
|
||||
for handler_data in handlers_data:
|
||||
handlers.append(
|
||||
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=_placeholder_handler,
|
||||
event_filters=[],
|
||||
desc=handler_data.get("desc", ""),
|
||||
extras_configs=handler_data.get("extras_configs", {}),
|
||||
)
|
||||
)
|
||||
return handlers
|
||||
|
||||
|
||||
class SupervisorRuntime:
|
||||
def __init__(
|
||||
self,
|
||||
server: JSONRPCServer,
|
||||
plugins_dir: Path,
|
||||
*,
|
||||
env_manager: PluginEnvironmentManager | None = None,
|
||||
worker_factory: Callable[
|
||||
[PluginSpec, JSONRPCServer, Path, PluginEnvironmentManager], WorkerRuntime
|
||||
]
|
||||
| None = None,
|
||||
) -> None:
|
||||
self.server = server
|
||||
self.plugins_dir = plugins_dir.resolve()
|
||||
self.repo_root = Path(__file__).resolve().parents[3]
|
||||
self.env_manager = env_manager or PluginEnvironmentManager(self.repo_root)
|
||||
self.worker_factory = worker_factory or WorkerRuntime
|
||||
self.loaded_plugins: list[str] = []
|
||||
self.skipped_plugins: dict[str, str] = {}
|
||||
self._workers_by_name: dict[str, WorkerRuntime] = {}
|
||||
self._handler_to_worker: dict[str, WorkerRuntime] = {}
|
||||
|
||||
async def start(self) -> None:
|
||||
discovery = discover_plugins(self.plugins_dir)
|
||||
self.skipped_plugins = dict(discovery.skipped_plugins)
|
||||
|
||||
for plugin in discovery.plugins:
|
||||
worker = self.worker_factory(
|
||||
plugin,
|
||||
self.server,
|
||||
self.repo_root,
|
||||
self.env_manager,
|
||||
)
|
||||
try:
|
||||
await worker.start()
|
||||
except Exception as exc:
|
||||
self.skipped_plugins[plugin.name] = str(exc)
|
||||
logger.error(f"Failed to start worker for {plugin.name}: {exc}")
|
||||
await worker.stop()
|
||||
continue
|
||||
|
||||
duplicate_handlers = [
|
||||
handler.handler_full_name
|
||||
for handler in worker.handlers
|
||||
if handler.handler_full_name in self._handler_to_worker
|
||||
]
|
||||
if duplicate_handlers:
|
||||
self.skipped_plugins[plugin.name] = (
|
||||
f"duplicate handlers: {', '.join(sorted(duplicate_handlers))}"
|
||||
)
|
||||
await worker.stop()
|
||||
continue
|
||||
|
||||
self._workers_by_name[plugin.name] = worker
|
||||
self.loaded_plugins.append(plugin.name)
|
||||
for handler in worker.handlers:
|
||||
self._handler_to_worker[handler.handler_full_name] = worker
|
||||
|
||||
self.loaded_plugins.sort()
|
||||
self.server.set_message_handler(self._handle_message)
|
||||
await self.server.start()
|
||||
self._log_startup_summary()
|
||||
|
||||
async def stop(self) -> None:
|
||||
for worker in list(self._workers_by_name.values()):
|
||||
await worker.stop()
|
||||
await self.server.stop()
|
||||
|
||||
async def _handle_message(self, message: JSONRPCMessage) -> None:
|
||||
if isinstance(message, JSONRPCRequest):
|
||||
if message.method == "handshake":
|
||||
await self.server.send_message(
|
||||
self._build_handshake_response(message.id)
|
||||
)
|
||||
return
|
||||
if message.method == "call_handler":
|
||||
await self._route_call_handler(message)
|
||||
return
|
||||
logger.warning(f"Unknown method from core: {message.method}")
|
||||
return
|
||||
|
||||
for worker in self._workers_by_name.values():
|
||||
if await worker.handle_context_response(message):
|
||||
return
|
||||
|
||||
logger.warning(f"Received response for unknown request id: {message.id}")
|
||||
|
||||
def _build_handshake_response(
|
||||
self, request_id: str | None
|
||||
) -> JSONRPCSuccessResponse:
|
||||
payload: dict[str, Any] = {}
|
||||
for worker in self._workers_by_name.values():
|
||||
payload.update(worker.raw_handshake)
|
||||
return JSONRPCSuccessResponse(
|
||||
jsonrpc="2.0",
|
||||
id=request_id,
|
||||
result=payload,
|
||||
)
|
||||
|
||||
async def _route_call_handler(self, message: JSONRPCRequest) -> None:
|
||||
try:
|
||||
params = CallHandlerRequest.Params.model_validate(message.params)
|
||||
except Exception as exc:
|
||||
await self.server.send_message(
|
||||
JSONRPCErrorResponse(
|
||||
jsonrpc="2.0",
|
||||
id=message.id,
|
||||
error=JSONRPCErrorData(
|
||||
code=-32602, message=f"Invalid params: {exc}"
|
||||
),
|
||||
)
|
||||
)
|
||||
return
|
||||
|
||||
worker = self._handler_to_worker.get(params.handler_full_name)
|
||||
if worker is None:
|
||||
await self.server.send_message(
|
||||
JSONRPCErrorResponse(
|
||||
jsonrpc="2.0",
|
||||
id=message.id,
|
||||
error=JSONRPCErrorData(
|
||||
code=-32601,
|
||||
message=f"Handler not found: {params.handler_full_name}",
|
||||
),
|
||||
)
|
||||
)
|
||||
return
|
||||
|
||||
await worker.forward_call_handler(message)
|
||||
|
||||
def _log_startup_summary(self) -> None:
|
||||
loaded = ", ".join(self.loaded_plugins) if self.loaded_plugins else "none"
|
||||
logger.info(f"Loaded plugins: {loaded}")
|
||||
if not self.skipped_plugins:
|
||||
logger.info("Skipped plugins: none")
|
||||
return
|
||||
for plugin_name, reason in sorted(self.skipped_plugins.items()):
|
||||
logger.warning(f"Skipped plugin {plugin_name}: {reason}")
|
||||
@@ -1,67 +0,0 @@
|
||||
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 AstrMessageEventModel
|
||||
|
||||
|
||||
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 HandlerStreamStartNotification(JSONRPCRequest):
|
||||
"""Notification sent when a handler stream starts."""
|
||||
|
||||
class Params(BaseModel):
|
||||
id: str | None # The original request ID
|
||||
handler_full_name: str
|
||||
|
||||
method: Literal["handler_stream_start"] = "handler_stream_start"
|
||||
params: Params # type: ignore[assignment]
|
||||
|
||||
|
||||
class HandlerStreamUpdateNotification(JSONRPCRequest):
|
||||
"""Notification sent when a handler stream has new data."""
|
||||
|
||||
class Params(BaseModel):
|
||||
id: str | None # The original request ID
|
||||
handler_full_name: str
|
||||
data: Any # The streamed data
|
||||
|
||||
method: Literal["handler_stream_update"] = "handler_stream_update"
|
||||
params: Params # type: ignore[assignment]
|
||||
|
||||
|
||||
class HandlerStreamEndNotification(JSONRPCRequest):
|
||||
"""Notification sent when a handler stream ends."""
|
||||
|
||||
class Params(BaseModel):
|
||||
id: str | None # The original request ID
|
||||
handler_full_name: str
|
||||
|
||||
method: Literal["handler_stream_end"] = "handler_stream_end"
|
||||
params: Params # type: ignore[assignment]
|
||||
@@ -1,321 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
try:
|
||||
import psutil
|
||||
except ImportError: # pragma: no cover - optional dependency
|
||||
psutil = None
|
||||
|
||||
from astrbot_sdk.api.star.context import Context
|
||||
from astrbot_sdk.runtime.galaxy import Galaxy
|
||||
|
||||
PLUGIN_COUNT = 8
|
||||
TARGET_PYTHON = "3.12"
|
||||
HANDSHAKE_TIMEOUT_SECONDS = 60.0
|
||||
|
||||
|
||||
class BenchmarkContext(Context):
|
||||
pass
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=(
|
||||
"Generate 8 Python 3.12 plugins and measure resource usage for the "
|
||||
"independent worker runtime."
|
||||
)
|
||||
)
|
||||
parser.add_argument(
|
||||
"--python-executable",
|
||||
default=sys.executable,
|
||||
help="Python executable used to launch the supervisor process.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--plugins-dir",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Optional directory to write generated plugins into.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--keep-plugins-dir",
|
||||
action="store_true",
|
||||
help="Keep the generated plugins directory instead of deleting it.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-json",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Optional path to write the benchmark report JSON.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def write_plugin(plugins_dir: Path, index: int) -> None:
|
||||
plugin_name = f"plugin_{index:03d}"
|
||||
command_name = f"bench_{index:03d}"
|
||||
plugin_dir = plugins_dir / plugin_name
|
||||
commands_dir = plugin_dir / "commands"
|
||||
commands_dir.mkdir(parents=True, exist_ok=True)
|
||||
(commands_dir / "__init__.py").write_text("", encoding="utf-8")
|
||||
|
||||
manifest = {
|
||||
"_schema_version": 2,
|
||||
"name": plugin_name,
|
||||
"display_name": plugin_name,
|
||||
"desc": f"Resource benchmark plugin {index}",
|
||||
"author": "codex",
|
||||
"version": "0.1.0",
|
||||
"runtime": {"python": TARGET_PYTHON},
|
||||
"components": [
|
||||
{
|
||||
"class": f"commands.plugin_{index:03d}:BenchmarkCommand{index:03d}",
|
||||
"type": "command",
|
||||
"name": command_name,
|
||||
"description": command_name,
|
||||
}
|
||||
],
|
||||
}
|
||||
(plugin_dir / "plugin.yaml").write_text(
|
||||
yaml.safe_dump(manifest, sort_keys=False),
|
||||
encoding="utf-8",
|
||||
)
|
||||
(plugin_dir / "requirements.txt").write_text("", encoding="utf-8")
|
||||
|
||||
module_source = f"""
|
||||
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 BenchmarkCommand{index:03d}(CommandComponent):
|
||||
def __init__(self, context: Context):
|
||||
self.context = context
|
||||
|
||||
@filter.command("{command_name}")
|
||||
async def handle(self, event: AstrMessageEvent):
|
||||
yield event.plain_result("{plugin_name}:{command_name}")
|
||||
""".strip()
|
||||
(commands_dir / f"plugin_{index:03d}.py").write_text(
|
||||
module_source + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
def _collect_with_psutil(root_pid: int) -> dict[str, Any]:
|
||||
assert psutil is not None
|
||||
root_process = psutil.Process(root_pid)
|
||||
processes = [root_process] + root_process.children(recursive=True)
|
||||
entries: list[dict[str, Any]] = []
|
||||
total_rss = 0
|
||||
|
||||
for process in processes:
|
||||
try:
|
||||
rss = process.memory_info().rss
|
||||
total_rss += rss
|
||||
entries.append(
|
||||
{
|
||||
"pid": process.pid,
|
||||
"name": process.name(),
|
||||
"rss_mb": round(rss / 1024 / 1024, 2),
|
||||
"cmdline": process.cmdline(),
|
||||
}
|
||||
)
|
||||
except (psutil.NoSuchProcess, psutil.AccessDenied):
|
||||
continue
|
||||
|
||||
entries.sort(key=lambda item: item["pid"])
|
||||
return {
|
||||
"collector": "psutil",
|
||||
"process_count": len(entries),
|
||||
"total_rss_mb": round(total_rss / 1024 / 1024, 2),
|
||||
"processes": entries,
|
||||
}
|
||||
|
||||
|
||||
def _collect_with_ps(root_pid: int) -> dict[str, Any]:
|
||||
process = subprocess.run(
|
||||
["ps", "-axo", "pid,ppid,rss,comm"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
children_by_parent: dict[int, list[tuple[int, int, str]]] = {}
|
||||
rss_by_pid: dict[int, int] = {}
|
||||
|
||||
for line in process.stdout.splitlines()[1:]:
|
||||
parts = line.strip().split(None, 3)
|
||||
if len(parts) != 4:
|
||||
continue
|
||||
pid, ppid, rss_kb, command = parts
|
||||
pid_int = int(pid)
|
||||
ppid_int = int(ppid)
|
||||
rss_int = int(rss_kb)
|
||||
rss_by_pid[pid_int] = rss_int
|
||||
children_by_parent.setdefault(ppid_int, []).append((pid_int, rss_int, command))
|
||||
|
||||
queue = [root_pid]
|
||||
seen: set[int] = set()
|
||||
entries: list[dict[str, Any]] = []
|
||||
total_rss = 0
|
||||
|
||||
while queue:
|
||||
pid = queue.pop(0)
|
||||
if pid in seen:
|
||||
continue
|
||||
seen.add(pid)
|
||||
rss_kb = rss_by_pid.get(pid)
|
||||
command = None
|
||||
for siblings in children_by_parent.values():
|
||||
for child_pid, child_rss, child_command in siblings:
|
||||
if child_pid == pid:
|
||||
rss_kb = child_rss
|
||||
command = child_command
|
||||
break
|
||||
if command is not None:
|
||||
break
|
||||
if rss_kb is not None:
|
||||
total_rss += rss_kb * 1024
|
||||
entries.append(
|
||||
{
|
||||
"pid": pid,
|
||||
"name": command or "unknown",
|
||||
"rss_mb": round((rss_kb * 1024) / 1024 / 1024, 2),
|
||||
"cmdline": [command] if command else [],
|
||||
}
|
||||
)
|
||||
for child_pid, _child_rss, _child_command in children_by_parent.get(pid, []):
|
||||
queue.append(child_pid)
|
||||
|
||||
entries.sort(key=lambda item: item["pid"])
|
||||
return {
|
||||
"collector": "ps",
|
||||
"process_count": len(entries),
|
||||
"total_rss_mb": round(total_rss / 1024 / 1024, 2),
|
||||
"processes": entries,
|
||||
}
|
||||
|
||||
|
||||
def collect_process_tree_metrics(root_pid: int) -> dict[str, Any]:
|
||||
if psutil is not None:
|
||||
try:
|
||||
return _collect_with_psutil(root_pid)
|
||||
except (PermissionError, psutil.Error):
|
||||
pass
|
||||
return _collect_with_ps(root_pid)
|
||||
|
||||
|
||||
async def terminate_process(process: Any) -> None:
|
||||
if process is None or process.poll() is not None:
|
||||
return
|
||||
process.terminate()
|
||||
try:
|
||||
await asyncio.to_thread(process.wait, 10.0)
|
||||
except Exception:
|
||||
process.kill()
|
||||
await asyncio.to_thread(process.wait, 10.0)
|
||||
|
||||
|
||||
async def run_benchmark(plugins_dir: Path, python_executable: str) -> dict[str, Any]:
|
||||
for index in range(PLUGIN_COUNT):
|
||||
write_plugin(plugins_dir, index)
|
||||
|
||||
galaxy = Galaxy()
|
||||
context = BenchmarkContext()
|
||||
started_at = time.perf_counter()
|
||||
star = await galaxy.connect_to_stdio_star(
|
||||
context=context,
|
||||
star_name="resource-benchmark",
|
||||
config={
|
||||
"plugins_dir": str(plugins_dir),
|
||||
"python_executable": python_executable,
|
||||
},
|
||||
)
|
||||
connected_at = time.perf_counter()
|
||||
|
||||
client_process = getattr(star._client, "_process", None)
|
||||
metadata: dict[str, Any] = {}
|
||||
handshake_error: str | None = None
|
||||
try:
|
||||
metadata = await asyncio.wait_for(
|
||||
star.handshake(),
|
||||
timeout=HANDSHAKE_TIMEOUT_SECONDS,
|
||||
)
|
||||
except Exception as exc:
|
||||
handshake_error = f"{exc.__class__.__name__}: {exc}"
|
||||
|
||||
measured_at = time.perf_counter()
|
||||
metrics = collect_process_tree_metrics(client_process.pid) if client_process else {}
|
||||
loaded_plugins = sorted(
|
||||
metadata_item.name
|
||||
for metadata_item in metadata.values()
|
||||
if getattr(metadata_item, "name", None)
|
||||
)
|
||||
|
||||
stop_error: str | None = None
|
||||
try:
|
||||
await star.stop()
|
||||
except Exception as exc:
|
||||
stop_error = f"{exc.__class__.__name__}: {exc}"
|
||||
await terminate_process(client_process)
|
||||
|
||||
return {
|
||||
"plugin_count": PLUGIN_COUNT,
|
||||
"target_python": TARGET_PYTHON,
|
||||
"python_executable": python_executable,
|
||||
"loaded_plugin_count": len(loaded_plugins),
|
||||
"loaded_plugins": loaded_plugins,
|
||||
"connect_duration_ms": round((connected_at - started_at) * 1000, 2),
|
||||
"handshake_duration_ms": round((measured_at - connected_at) * 1000, 2),
|
||||
"startup_total_duration_ms": round((measured_at - started_at) * 1000, 2),
|
||||
"handshake_error": handshake_error,
|
||||
"metrics": metrics,
|
||||
"stop_error": stop_error,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
temp_dir: tempfile.TemporaryDirectory[str] | None = None
|
||||
plugins_dir = args.plugins_dir
|
||||
if plugins_dir is None:
|
||||
temp_dir = tempfile.TemporaryDirectory(prefix="astrbot-8-plugin-bench-")
|
||||
plugins_dir = Path(temp_dir.name)
|
||||
else:
|
||||
plugins_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
report = asyncio.run(
|
||||
run_benchmark(
|
||||
plugins_dir=plugins_dir,
|
||||
python_executable=args.python_executable,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
if temp_dir is not None and not args.keep_plugins_dir:
|
||||
temp_dir.cleanup()
|
||||
|
||||
report["plugins_dir"] = str(plugins_dir)
|
||||
|
||||
if args.output_json is not None:
|
||||
args.output_json.write_text(
|
||||
json.dumps(report, ensure_ascii=True, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
print(json.dumps(report, ensure_ascii=True, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,76 +0,0 @@
|
||||
import asyncio
|
||||
from astrbot_sdk.runtime.galaxy import Galaxy
|
||||
from astrbot_sdk.api.event import AstrMessageEvent
|
||||
from astrbot_sdk.api.event.astrbot_message import AstrBotMessage, MessageMember
|
||||
from astrbot_sdk.api.platform.platform_metadata import PlatformMetadata
|
||||
from astrbot_sdk.api.event.message_type import MessageType
|
||||
|
||||
from astrbot_sdk.api.star.context import Context
|
||||
from astrbot_sdk.api.basic.conversation_mgr import BaseConversationManager
|
||||
|
||||
|
||||
class ConversationManager(BaseConversationManager):
|
||||
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:
|
||||
import uuid
|
||||
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
class TestContext(Context):
|
||||
def __init__(self, conversation_manager: ConversationManager):
|
||||
super().__init__()
|
||||
self.conversation_manager = conversation_manager
|
||||
self._register_component(self.conversation_manager)
|
||||
|
||||
|
||||
async def amain():
|
||||
galaxy = Galaxy()
|
||||
conversation_manager = ConversationManager()
|
||||
context = TestContext(conversation_manager)
|
||||
star = await galaxy.connect_to_websocket_star(
|
||||
context=context,
|
||||
star_name="hello",
|
||||
config={
|
||||
"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",
|
||||
)
|
||||
|
||||
async for result in star.call_handler(star._handlers[0], event):
|
||||
print(f"Handler result: {result}")
|
||||
|
||||
await star.stop()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(amain())
|
||||
@@ -1,322 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
from astrbot_sdk.runtime.rpc.jsonrpc import (
|
||||
JSONRPCRequest,
|
||||
JSONRPCSuccessResponse,
|
||||
)
|
||||
from astrbot_sdk.runtime.stars.registry import EventType, StarHandlerMetadata
|
||||
from astrbot_sdk.runtime.supervisor import (
|
||||
PluginEnvironmentManager,
|
||||
PluginSpec,
|
||||
SupervisorRuntime,
|
||||
WorkerRuntime,
|
||||
discover_plugins,
|
||||
)
|
||||
from astrbot_sdk.runtime.types import CallHandlerRequest
|
||||
|
||||
|
||||
def write_plugin(
|
||||
root: Path,
|
||||
folder_name: str,
|
||||
*,
|
||||
plugin_name: str | None = None,
|
||||
python_version: str | None = "3.12",
|
||||
include_requirements: bool = True,
|
||||
) -> Path:
|
||||
plugin_dir = root / folder_name
|
||||
commands_dir = plugin_dir / "commands"
|
||||
commands_dir.mkdir(parents=True, exist_ok=True)
|
||||
(commands_dir / "__init__.py").write_text("", encoding="utf-8")
|
||||
|
||||
manifest: dict[str, Any] = {
|
||||
"_schema_version": 2,
|
||||
"name": plugin_name or folder_name,
|
||||
"display_name": folder_name,
|
||||
"desc": "test plugin",
|
||||
"author": "tester",
|
||||
"version": "0.1.0",
|
||||
"components": [
|
||||
{
|
||||
"class": "commands.sample:SampleCommand",
|
||||
"type": "command",
|
||||
"name": "hello",
|
||||
"description": "hello",
|
||||
}
|
||||
],
|
||||
}
|
||||
if python_version is not None:
|
||||
manifest["runtime"] = {"python": python_version}
|
||||
|
||||
(plugin_dir / "plugin.yaml").write_text(
|
||||
yaml.safe_dump(manifest, sort_keys=False),
|
||||
encoding="utf-8",
|
||||
)
|
||||
if include_requirements:
|
||||
(plugin_dir / "requirements.txt").write_text("", encoding="utf-8")
|
||||
return plugin_dir
|
||||
|
||||
|
||||
class FakeServer:
|
||||
def __init__(self) -> None:
|
||||
self.handler = None
|
||||
self.sent_messages: list[Any] = []
|
||||
|
||||
def set_message_handler(self, handler) -> None:
|
||||
self.handler = handler
|
||||
|
||||
async def start(self) -> None:
|
||||
return None
|
||||
|
||||
async def stop(self) -> None:
|
||||
return None
|
||||
|
||||
async def send_message(self, message) -> None:
|
||||
self.sent_messages.append(message)
|
||||
|
||||
|
||||
class FakeEnvManager(PluginEnvironmentManager):
|
||||
def __init__(self) -> None:
|
||||
self.prepared: list[str] = []
|
||||
|
||||
def prepare_environment(self, plugin: PluginSpec) -> Path:
|
||||
self.prepared.append(plugin.name)
|
||||
return Path("/tmp/fake-python")
|
||||
|
||||
|
||||
class FakeWorkerRuntime(WorkerRuntime):
|
||||
def __init__(
|
||||
self,
|
||||
plugin: PluginSpec,
|
||||
server,
|
||||
repo_root: Path,
|
||||
env_manager: PluginEnvironmentManager,
|
||||
) -> None:
|
||||
self.plugin = plugin
|
||||
self.server = server
|
||||
self.repo_root = repo_root
|
||||
self.env_manager = env_manager
|
||||
self.raw_handshake: dict[str, Any] = {}
|
||||
self.handlers: list[StarHandlerMetadata] = []
|
||||
self.forwarded_requests: list[JSONRPCRequest] = []
|
||||
self.received_context_responses: list[Any] = []
|
||||
self.stopped = False
|
||||
|
||||
async def start(self) -> None:
|
||||
handler_full_name = (
|
||||
f"commands.{self.plugin.name}:SampleCommand.handle_{self.plugin.name}"
|
||||
)
|
||||
self.raw_handshake = {
|
||||
f"{self.plugin.name}.main": {
|
||||
"name": self.plugin.name,
|
||||
"author": "tester",
|
||||
"desc": "test plugin",
|
||||
"version": "0.1.0",
|
||||
"repo": None,
|
||||
"module_path": f"{self.plugin.name}.main",
|
||||
"root_dir_name": self.plugin.plugin_dir.name,
|
||||
"reserved": False,
|
||||
"activated": True,
|
||||
"config": None,
|
||||
"star_handler_full_names": [handler_full_name],
|
||||
"display_name": self.plugin.name,
|
||||
"logo_path": None,
|
||||
"handlers": [
|
||||
{
|
||||
"event_type": EventType.AdapterMessageEvent.value,
|
||||
"handler_full_name": handler_full_name,
|
||||
"handler_name": f"handle_{self.plugin.name}",
|
||||
"handler_module_path": f"commands.{self.plugin.name}",
|
||||
"desc": "",
|
||||
"extras_configs": {},
|
||||
}
|
||||
],
|
||||
}
|
||||
}
|
||||
self.handlers = [
|
||||
StarHandlerMetadata(
|
||||
event_type=EventType.AdapterMessageEvent,
|
||||
handler_full_name=handler_full_name,
|
||||
handler_name=f"handle_{self.plugin.name}",
|
||||
handler_module_path=f"commands.{self.plugin.name}",
|
||||
handler=lambda *args, **kwargs: None,
|
||||
event_filters=[],
|
||||
)
|
||||
]
|
||||
|
||||
async def stop(self) -> None:
|
||||
self.stopped = True
|
||||
|
||||
async def forward_call_handler(self, request: JSONRPCRequest) -> None:
|
||||
self.forwarded_requests.append(request)
|
||||
handler_full_name = self.handlers[0].handler_full_name
|
||||
await self.server.send_message(
|
||||
JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
method="handler_stream_start",
|
||||
params={
|
||||
"id": request.id,
|
||||
"handler_full_name": handler_full_name,
|
||||
},
|
||||
)
|
||||
)
|
||||
await self.server.send_message(
|
||||
JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
method="handler_stream_update",
|
||||
params={
|
||||
"id": request.id,
|
||||
"handler_full_name": handler_full_name,
|
||||
"data": {"plugin": self.plugin.name},
|
||||
},
|
||||
)
|
||||
)
|
||||
await self.server.send_message(
|
||||
JSONRPCRequest(
|
||||
jsonrpc="2.0",
|
||||
method="handler_stream_end",
|
||||
params={
|
||||
"id": request.id,
|
||||
"handler_full_name": handler_full_name,
|
||||
},
|
||||
)
|
||||
)
|
||||
await self.server.send_message(
|
||||
JSONRPCSuccessResponse(
|
||||
jsonrpc="2.0",
|
||||
id=request.id,
|
||||
result={"handled_by": self.plugin.name},
|
||||
)
|
||||
)
|
||||
|
||||
async def handle_context_response(self, message) -> bool:
|
||||
if message.id != f"ctx:{self.plugin.name}:1":
|
||||
return False
|
||||
self.received_context_responses.append(message)
|
||||
return True
|
||||
|
||||
|
||||
class DiscoverPluginsTest(unittest.TestCase):
|
||||
def test_discover_plugins_requires_runtime_python(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
root = Path(temp_dir)
|
||||
write_plugin(root, "plugin_one", plugin_name="plugin_one")
|
||||
write_plugin(
|
||||
root,
|
||||
"plugin_two",
|
||||
plugin_name="plugin_two",
|
||||
python_version=None,
|
||||
)
|
||||
write_plugin(
|
||||
root,
|
||||
"plugin_three",
|
||||
plugin_name="plugin_three",
|
||||
include_requirements=False,
|
||||
)
|
||||
|
||||
discovery = discover_plugins(root)
|
||||
|
||||
self.assertEqual([plugin.name for plugin in discovery.plugins], ["plugin_one"])
|
||||
self.assertIn("plugin_two", discovery.skipped_plugins)
|
||||
self.assertIn("plugin_three", discovery.skipped_plugins)
|
||||
|
||||
|
||||
class SupervisorRuntimeTest(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self) -> None:
|
||||
self.temp_dir = tempfile.TemporaryDirectory()
|
||||
self.plugins_dir = Path(self.temp_dir.name)
|
||||
write_plugin(self.plugins_dir, "plugin_one", plugin_name="plugin_one")
|
||||
write_plugin(self.plugins_dir, "plugin_two", plugin_name="plugin_two")
|
||||
self.server = FakeServer()
|
||||
|
||||
async def asyncTearDown(self) -> None:
|
||||
self.temp_dir.cleanup()
|
||||
|
||||
async def test_handshake_aggregates_workers_and_routes_call_handler(self) -> None:
|
||||
runtime = SupervisorRuntime(
|
||||
server=self.server,
|
||||
plugins_dir=self.plugins_dir,
|
||||
env_manager=FakeEnvManager(),
|
||||
worker_factory=FakeWorkerRuntime,
|
||||
)
|
||||
await runtime.start()
|
||||
|
||||
await self.server.handler(
|
||||
JSONRPCRequest(jsonrpc="2.0", id="handshake-1", method="handshake")
|
||||
)
|
||||
handshake_response = self.server.sent_messages[-1]
|
||||
self.assertIsInstance(handshake_response, JSONRPCSuccessResponse)
|
||||
self.assertEqual(
|
||||
sorted(handshake_response.result.keys()),
|
||||
["plugin_one.main", "plugin_two.main"],
|
||||
)
|
||||
|
||||
handler_full_name = "commands.plugin_two:SampleCommand.handle_plugin_two"
|
||||
await self.server.handler(
|
||||
CallHandlerRequest(
|
||||
jsonrpc="2.0",
|
||||
id="call-1",
|
||||
method="call_handler",
|
||||
params=CallHandlerRequest.Params(
|
||||
handler_full_name=handler_full_name,
|
||||
event={
|
||||
"message_str": "hello",
|
||||
"message_obj": {
|
||||
"type": "FriendMessage",
|
||||
"self_id": "bot",
|
||||
"session_id": "session",
|
||||
"message_id": "message-id",
|
||||
"sender": {"user_id": "user-1", "nickname": "User 1"},
|
||||
"message": [],
|
||||
"message_str": "hello",
|
||||
"raw_message": {},
|
||||
"timestamp": 0,
|
||||
},
|
||||
"platform_meta": {
|
||||
"name": "fake",
|
||||
"description": "fake",
|
||||
"id": "fake-1",
|
||||
},
|
||||
"session_id": "session",
|
||||
"is_at_or_wake_command": True,
|
||||
},
|
||||
args={},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
self.server.sent_messages[-1].result, {"handled_by": "plugin_two"}
|
||||
)
|
||||
await runtime.stop()
|
||||
|
||||
async def test_routes_context_response_back_to_matching_worker(self) -> None:
|
||||
runtime = SupervisorRuntime(
|
||||
server=self.server,
|
||||
plugins_dir=self.plugins_dir,
|
||||
env_manager=FakeEnvManager(),
|
||||
worker_factory=FakeWorkerRuntime,
|
||||
)
|
||||
await runtime.start()
|
||||
|
||||
await self.server.handler(
|
||||
JSONRPCSuccessResponse(
|
||||
jsonrpc="2.0",
|
||||
id="ctx:plugin_one:1",
|
||||
result={"data": "ok"},
|
||||
)
|
||||
)
|
||||
|
||||
worker = runtime._workers_by_name["plugin_one"]
|
||||
self.assertEqual(len(worker.received_context_responses), 1)
|
||||
await runtime.stop()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user