From 363fda5a695f297fcb5bf1aac9a2887b30fc76bf Mon Sep 17 00:00:00 2001 From: Soulter <905617992@qq.com> Date: Tue, 11 Nov 2025 20:22:13 +0800 Subject: [PATCH] feat: enhance context management with component registration and RPC function handling --- src/astrbot_sdk/api/star/context.py | 46 ++++++++ src/astrbot_sdk/runtime/api/context.py | 3 + .../runtime/api/conversation_mgr.py | 2 +- src/astrbot_sdk/runtime/api/util.py | 108 +++++++++++++----- src/astrbot_sdk/runtime/galaxy.py | 11 +- src/astrbot_sdk/runtime/star_runner.py | 4 +- src/astrbot_sdk/runtime/stars/new.py | 48 +++++++- src/astrbot_sdk/runtime/stars/virtual.py | 3 + src/astrbot_sdk/tests/start_client.py | 31 ++++- test_plugin/commands/hello.py | 2 +- 10 files changed, 215 insertions(+), 43 deletions(-) diff --git a/src/astrbot_sdk/api/star/context.py b/src/astrbot_sdk/api/star/context.py index cac4fcd0f..46c198140 100644 --- a/src/astrbot_sdk/api/star/context.py +++ b/src/astrbot_sdk/api/star/context.py @@ -1,6 +1,52 @@ from abc import ABC +from typing import Any, Callable from ..basic.conversation_mgr import BaseConversationManager class Context(ABC): conversation_manager: BaseConversationManager + + 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 + + def get_registered_function(self, full_name: str) -> Callable | None: + """Get a registered function by its full name. + + Args: + full_name: Full name in format "ComponentClassName.method_name" + + Returns: + The callable function or None if not found + """ + + return self._registered_functions.get(full_name) + + def list_registered_functions(self) -> list[str]: + """List all registered function names. + + Returns: + List of full function names in format "ComponentClassName.method_name" + """ + return list(self._registered_functions.keys()) diff --git a/src/astrbot_sdk/runtime/api/context.py b/src/astrbot_sdk/runtime/api/context.py index 929be8f41..888bb2768 100644 --- a/src/astrbot_sdk/runtime/api/context.py +++ b/src/astrbot_sdk/runtime/api/context.py @@ -4,7 +4,10 @@ from .conversation_mgr import ConversationManager 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=None): diff --git a/src/astrbot_sdk/runtime/api/conversation_mgr.py b/src/astrbot_sdk/runtime/api/conversation_mgr.py index 72b4ecd9b..75f7a96de 100644 --- a/src/astrbot_sdk/runtime/api/conversation_mgr.py +++ b/src/astrbot_sdk/runtime/api/conversation_mgr.py @@ -11,7 +11,7 @@ class ConversationManager(BaseConversationManager): """ self.runner = runner - @rpc_method + @rpc_method(return_model=str) async def new_conversation( self, unified_msg_origin: str, diff --git a/src/astrbot_sdk/runtime/api/util.py b/src/astrbot_sdk/runtime/api/util.py index 7126faf20..c99fdde6e 100644 --- a/src/astrbot_sdk/runtime/api/util.py +++ b/src/astrbot_sdk/runtime/api/util.py @@ -1,38 +1,92 @@ import inspect from functools import wraps -from typing import Callable +from typing import Callable, Type, Any, get_type_hints, TypeVar, overload from ..star_runner import StarRunner from ..types import CallContextFunctionRequest +F = TypeVar("F", bound=Callable[..., Any]) -def rpc_method(func: Callable) -> Callable: - """sign as an RPC method.""" - @wraps(func) - async def wrapper(self, *args, **kwargs): - if not hasattr(self, "runner") or not isinstance(self.runner, StarRunner): - raise RuntimeError( - f"Class {self.__class__.__name__} is not configured for RPC calls." +@overload +def rpc_method(func: F) -> F: ... + + +@overload +def rpc_method(*, return_model: Type[Any] | None = None) -> Callable[[F], F]: ... + + +def rpc_method( + func: F | None = None, *, return_model: Type[Any] | None = None +) -> F | Callable[[F], F]: + """sign as an RPC method. + + Args: + func: The function to decorate + return_model: The expected return type/model (e.g., dict, BaseModel, etc.) + If not specified, will try to infer from type hints + + Examples: + @rpc_method + async def get_config(self) -> dict: + ... + + @rpc_method(return_model=dict) + async def get_config(self): + ... + """ + + def decorator(f: F) -> F: + # Try to get return type from type hints if not explicitly provided + _return_model = return_model + if _return_model is None: + try: + hints = get_type_hints(f) + if "return" in hints: + _return_model = hints["return"] + except Exception: + pass + + # Store return model as function attribute for potential inspection + setattr(f, "__return_model__", _return_model) + + @wraps(f) + async def wrapper(self, *args, **kwargs): + if not hasattr(self, "runner") or not isinstance(self.runner, StarRunner): + raise RuntimeError( + f"Class {self.__class__.__name__} is not configured for RPC calls." + ) + method_name = f"{self.__class__.__name__}.{f.__name__}" + sig = inspect.signature(f) + bound_args = sig.bind(self, *args, **kwargs) + bound_args.apply_defaults() + params = dict(bound_args.arguments) + params.pop("self") + + runner: StarRunner = getattr(self, "runner") + + result = await runner._call_rpc( + CallContextFunctionRequest( + jsonrpc="2.0", + id=runner._generate_request_id(), + method="call_context_function", + params=CallContextFunctionRequest.Params( + name=method_name, + args=params, + ), + ) ) - method_name = f"{self.__class__.__name__}.{func.__name__}" - sig = inspect.signature(func) - bound_args = sig.bind(self, *args, **kwargs) - bound_args.apply_defaults() - params = dict(bound_args.arguments) - params.pop("self") - runner: StarRunner = getattr(self, "runner") + # TODO: Process result based on _return_model if needed + return result - return await runner._call_rpc( - CallContextFunctionRequest( - jsonrpc="2.0", - id=runner._generate_request_id(), - method="call_context_function", - params=CallContextFunctionRequest.Params( - name=method_name, - args=params, - ), - ) - ) + # Also store on wrapper for easy access + setattr(wrapper, "__return_model__", _return_model) + return wrapper # type: ignore - return wrapper + # Support both @rpc_method and @rpc_method(return_model=...) + if func is None: + # Called with arguments: @rpc_method(return_model=...) + return decorator + else: + # Called without arguments: @rpc_method + return decorator(func) diff --git a/src/astrbot_sdk/runtime/galaxy.py b/src/astrbot_sdk/runtime/galaxy.py index 75ca65146..a8b248262 100644 --- a/src/astrbot_sdk/runtime/galaxy.py +++ b/src/astrbot_sdk/runtime/galaxy.py @@ -11,6 +11,7 @@ let AstrBot core interact with both types of stars without needing to know the u from .stars.virtual import VirtualStar from .stars.new import NewStdioStar, NewWebSocketStar +from ..api.star.context import Context # from .types import StarURI, StarType @@ -19,18 +20,20 @@ class Galaxy: vs_map: dict[str, VirtualStar] = {} - async def connect_to_stdio_star(self, star_name: str, config: dict) -> NewStdioStar: + 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(**config) + star = NewStdioStar(context=context, **config) await star.initialize() self.vs_map[star_name] = star return star async def connect_to_websocket_star( - self, star_name: str, config: dict + self, context: Context, star_name: str, config: dict ) -> NewWebSocketStar: """Connect to a new-style websocket star given its name.""" - star = NewWebSocketStar(**config) + star = NewWebSocketStar(context=context, **config) await star.initialize() self.vs_map[star_name] = star return star diff --git a/src/astrbot_sdk/runtime/star_runner.py b/src/astrbot_sdk/runtime/star_runner.py index c31c235a3..464ac0171 100644 --- a/src/astrbot_sdk/runtime/star_runner.py +++ b/src/astrbot_sdk/runtime/star_runner.py @@ -17,13 +17,13 @@ class StarRunner: def __init__(self, server: JSONRPCServer): self.server = server self._request_id_counter = 0 - self.pending_requests: dict[str, asyncio.Future] = {} + self.pending_requests: dict[str, asyncio.Future[JSONRPCMessage]] = {} def _generate_request_id(self) -> str: self._request_id_counter += 1 return str(self._request_id_counter) - async def _call_rpc(self, message: JSONRPCMessage): + async def _call_rpc(self, message: JSONRPCMessage) -> JSONRPCMessage | None: if message.id is not None: self.pending_requests[message.id] = asyncio.get_event_loop().create_future() await self.server.send_message(message) diff --git a/src/astrbot_sdk/runtime/stars/new.py b/src/astrbot_sdk/runtime/stars/new.py index 2daf06cdc..c32ecd656 100644 --- a/src/astrbot_sdk/runtime/stars/new.py +++ b/src/astrbot_sdk/runtime/stars/new.py @@ -33,12 +33,23 @@ class NewStar(VirtualStar): def __init__( self, client: JSONRPCClient, + context: Any = None, ) -> None: """Initialize a NewStar instance. Args: client: JSON-RPC client for communication + context: Context instance for managing managers and their functions """ + # Import here to avoid circular dependency + from ..api.context import Context + + # Initialize context + if context is None: + context = Context.default_context() + + super().__init__(context) + self._client = client self._metadata: dict[str, StarMetadata] = {} self._handlers: list[StarHandlerMetadata] = [] @@ -120,14 +131,34 @@ class NewStar(VirtualStar): Args: request: The JSON-RPC request from the plugin """ - result: dict = {} + result: Any = None try: # Handle core methods that plugins might call method = request.method params = request.params if method == "call_context_function": - logger.debug(f"plugin called call_context_function: {params}") + ctx = self._context + 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 = ctx.get_registered_function(func_full_name) + if func is None: + raise ValueError(f"Function not found: {func_full_name}") + + # Call the function + import inspect + + if inspect.iscoroutinefunction(func): + result = await func(**args) + else: + result = func(**args) + + logger.debug(f"call_context_function result: {result}") else: raise ValueError(f"Unknown method: {method}") @@ -135,7 +166,9 @@ class NewStar(VirtualStar): response = JSONRPCSuccessResponse( jsonrpc="2.0", id=request.id, - result=result, + result={ + "data": result, + }, ) await self._client.send_message(response) @@ -537,6 +570,7 @@ class NewStdioStar(NewStar): self, plugin_dir: str, python_executable: str = "python", + context: Any = None, **kwargs: Any, ) -> None: """Initialize a STDIO-based NewStar. @@ -544,7 +578,7 @@ class NewStdioStar(NewStar): Args: plugin_dir: Path to the plugin directory python_executable: Python executable to use (defaults to 'python') - main_script: Main script filename (defaults to 'main.py') + context: Context instance for managing managers and their functions """ # Construct the command to start the plugin if not os.path.exists(plugin_dir): @@ -554,7 +588,7 @@ class NewStdioStar(NewStar): # Create StdioClient with subprocess management client = StdioClient(command=command, cwd=plugin_dir) - super().__init__(client) + super().__init__(client, context=context) class NewWebSocketStar(NewStar): @@ -569,6 +603,7 @@ class NewWebSocketStar(NewStar): url: str, heartbeat: float = 30.0, reconnect_interval: float = 5.0, + context: Any = None, **kwargs: Any, ) -> None: """Initialize a WebSocket-based NewStar. @@ -577,11 +612,12 @@ class NewWebSocketStar(NewStar): 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) + super().__init__(client, context=context) self._url = url self._heartbeat = heartbeat self._reconnect_interval = reconnect_interval diff --git a/src/astrbot_sdk/runtime/stars/virtual.py b/src/astrbot_sdk/runtime/stars/virtual.py index f4449fc29..acdb90959 100644 --- a/src/astrbot_sdk/runtime/stars/virtual.py +++ b/src/astrbot_sdk/runtime/stars/virtual.py @@ -4,6 +4,7 @@ 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): @@ -13,6 +14,8 @@ class VirtualStar(ABC): 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: diff --git a/src/astrbot_sdk/tests/start_client.py b/src/astrbot_sdk/tests/start_client.py index ea16e3298..232bc35bb 100644 --- a/src/astrbot_sdk/tests/start_client.py +++ b/src/astrbot_sdk/tests/start_client.py @@ -5,12 +5,39 @@ 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( - "hello", - { + context=context, + star_name="hello", + config={ "url": "ws://127.0.0.1:8765", }, ) diff --git a/test_plugin/commands/hello.py b/test_plugin/commands/hello.py index 17bf46d83..1d7dd152e 100644 --- a/test_plugin/commands/hello.py +++ b/test_plugin/commands/hello.py @@ -11,4 +11,4 @@ class HelloCommand(CommandComponent): async def hello(self, event: AstrMessageEvent): ret = await self.context.conversation_manager.new_conversation("hello") print(f"New conversation created: {ret}") - yield event.plain_result("Hello, Astrbot!") + yield event.plain_result(f"Hello, Astrbot! Created conversation ID: {ret}")