perf: normalize call_context_function rpc method

This commit is contained in:
Soulter
2025-11-12 20:41:14 +08:00
parent 363fda5a69
commit 0612133cdf
5 changed files with 41 additions and 108 deletions

View File

@@ -1,5 +1,7 @@
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):
@@ -10,12 +12,12 @@ class Context(BaseContext):
self.register_component(self.conversation_manager)
@classmethod
def default_context(cls, runner=None):
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=runner)
conversation_manager = ConversationManager(runner)
return cls(conversation_manager=conversation_manager)

View File

@@ -1,17 +1,16 @@
from ...api.basic.conversation_mgr import BaseConversationManager
from .util import rpc_method
from ..star_runner import StarRunner
class ConversationManager(BaseConversationManager):
def __init__(self, runner=None):
def __init__(self, runner: StarRunner):
"""Initialize ConversationManager.
Args:
runner: Optional StarRunner instance for RPC functionality.
"""
self.runner = runner
@rpc_method(return_model=str)
async def new_conversation(
self,
unified_msg_origin: str,
@@ -19,4 +18,15 @@ class ConversationManager(BaseConversationManager):
content: list[dict] | None = None,
title: str | None = None,
persona_id: str | None = None,
) -> str: ...
) -> 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"]

View File

@@ -1,92 +0,0 @@
import inspect
from functools import wraps
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])
@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,
),
)
)
# TODO: Process result based on _return_model if needed
return result
# Also store on wrapper for easy access
setattr(wrapper, "__return_model__", _return_model)
return wrapper # type: ignore
# 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)

View File

@@ -1,6 +1,7 @@
import asyncio
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 (
@@ -30,6 +31,27 @@ class StarRunner:
if message.id is not None:
return await self.pending_requests[message.id]
async def _call_context_function(
self, method_name: str, params: dict[str, Any]
) -> dict[str, Any]:
result = await self._call_rpc(
JSONRPCRequest(
jsonrpc="2.0",
id=self._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):
if isinstance(message, JSONRPCRequest):
logger.debug(f"Received RPC request: {message.method}")

View File

@@ -30,12 +30,3 @@ class CallHandlerRequest(JSONRPCRequest):
method: Literal["call_handler"]
params: Params | dict = Field(default_factory=dict)
class CallContextFunctionRequest(JSONRPCRequest):
class Params(BaseModel):
name: str
args: dict[str, Any] = {}
method: Literal["call_context_function"]
params: Params | dict = Field(default_factory=dict)