第 5 章 第 2 站：Agent 收信

# 不用元类：每个子类都必须记得加装饰器
class MyAgent(AgentBase):
    @wrap_with_hooks   # 容易忘！
    async def reply(self, *args, **kwargs):
        ...

    @wrap_with_hooks   # 也容易忘！
    async def observe(self, msg):
        ...

# 用元类：定义类时自动包装，无需手动操作
class MyAgent(AgentBase):   # AgentBase 使用了 _AgentMeta 元类
    async def reply(self, *args, **kwargs):   # 自动被包装
        ...

    async def observe(self, msg):   # 也自动被包装
        ...

async def __call__(self, *args: Any, **kwargs: Any) -> Msg:
    """Call the reply function with the given arguments."""
    self._reply_id = shortuuid.uuid()

    reply_msg: Msg | None = None
    try:
        self._reply_task = asyncio.current_task()
        reply_msg = await self.reply(*args, **kwargs)

    # The interruption is triggered by calling the interrupt method
    except asyncio.CancelledError:
        reply_msg = await self.handle_interrupt(*args, **kwargs)

    finally:
        # Broadcast the reply message to all subscribers
        if reply_msg:
            await self._broadcast_to_subscribers(reply_msg)
        self._reply_task = None

    return reply_msg

self._reply_id = shortuuid.uuid()

self._reply_task = asyncio.current_task()

reply_msg = await self.reply(*args, **kwargs)

async def reply(self, *args: Any, **kwargs: Any) -> Msg:
    """The main logic of the agent, which generates a reply based on the
    current state and input arguments."""
    raise NotImplementedError(
        "The reply function is not implemented in "
        f"{self.__class__.__name__} class.",
    )

finally:
    if reply_msg:
        await self._broadcast_to_subscribers(reply_msg)
    self._reply_task = None

class _AgentMeta(type):
    """The agent metaclass that wraps the agent's reply, observe and print
    functions with pre- and post-hooks."""

    def __new__(mcs, name: Any, bases: Any, attrs: Dict) -> Any:
        """Wrap the agent's functions with hooks."""

        for func_name in [
            "reply",
            "print",
            "observe",
        ]:
            if func_name in attrs:
                attrs[func_name] = _wrap_with_hooks(attrs[func_name])

        return super().__new__(mcs, name, bases, attrs)

class AgentBase(StateModule, metaclass=_AgentMeta):

def _wrap_with_hooks(original_func: Callable) -> Callable:

    @wraps(original_func)
    async def async_wrapper(self: AgentBase, *args, **kwargs) -> Any:

        # 1. 防重入：如果已经在 Hook 中，直接执行原函数
        if getattr(self, hook_guard_attr, False):
            return await original_func(self, *args, **kwargs)

        # 2. 把所有参数归一化为关键字参数
        normalized_kwargs = _normalize_to_kwargs(...)

        # 3. 执行所有 pre-hooks（前置钩子）
        for pre_hook in pre_hooks:
            modified_kwargs = await pre_hook(self, normalized_kwargs)
            if modified_kwargs is not None:
                normalized_kwargs = modified_kwargs

        # 4. 执行原函数
        current_output = await original_func(self, **normalized_kwargs)

        # 5. 执行所有 post-hooks（后置钩子）
        for post_hook in post_hooks:
            modified_output = await post_hook(self, normalized_kwargs, output)
            if modified_output is not None:
                current_output = modified_output

        return current_output

    return async_wrapper

# Initialize the instance-level hooks
self._instance_pre_print_hooks = OrderedDict()
self._instance_post_print_hooks = OrderedDict()

self._instance_pre_reply_hooks = OrderedDict()
self._instance_post_reply_hooks = OrderedDict()

self._instance_pre_observe_hooks = OrderedDict()
self._instance_post_observe_hooks = OrderedDict()

if getattr(self, hook_guard_attr, False):
    return await original_func(self, *args, **kwargs)

async def _broadcast_to_subscribers(
    self,
    msg: Msg | list[Msg] | None,
) -> None:
    """Broadcast the message to all subscribers."""
    if msg is None:
        return

    broadcast_msg = self._strip_thinking_blocks(msg)

    for subscribers in self._subscribers.values():
        for subscriber in subscribers:
            await subscriber.observe(broadcast_msg)

broadcast_msg = self._strip_thinking_blocks(msg)

self._subscribers: dict[str, list[AgentBase]] = {}

async def __call__(self, *args: Any, **kwargs: Any) -> Msg:
    """Call the reply function with the given arguments."""
    self._reply_id = shortuuid.uuid()
    print(f"[__call__] reply_id={self._reply_id}, args={args}, kwargs={kwargs}")
    # ...

from agentscope.agent import ReActAgent
from agentscope.memory import InMemoryMemory

agent = ReActAgent(
    name="test",
    sys_prompt="test",
    model=None,
    memory=InMemoryMemory(),
)

# 检查 reply 方法是否被包装
print(type(agent.reply))          # <class 'method'>
print(agent.reply.__wrapped__)    # 如果被 wraps 包装过，可以访问原始函数
print(agent.reply.__name__)       # 'reply'（@wraps 保留了原始名称）

async def __call__(self, *args: Any, **kwargs: Any) -> Msg:
    """Call the reply function with the given arguments."""
    self._reply_id = shortuuid.uuid()
    # 加这一行 ↓
    if args and hasattr(args[0], 'name'):
        print(f"[收信] {self.name} 收到来自 {args[0].name} 的消息")
    reply_msg: Msg | None = None
    # ... 后续代码不变

result = await agent(Msg("user", "北京今天天气怎么样？", "user"))

await agent(Msg("user", "北京天气如何？", "user"))
await agent(Msg("user", "上海呢？", "user"))