{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 自定义代理\n",
    "\n",
    "您可能会有一些行为不属于预设的代理。\n",
    "在这种情况下，您可以构建自定义代理。\n",
    "\n",
    "AgentChat 中的所有代理都继承自 {py:class}`~autogen_agentchat.agents.BaseChatAgent` \n",
    "类，并实现以下抽象方法和属性：\n",
    "\n",
    "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages`: 定义代理响应消息的行为的抽象方法。当代理被要求在 {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run` 中提供响应时，会调用此方法。它返回一个 {py:class}`~autogen_agentchat.base.Response` 对象。\n",
    "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_reset`: 将代理重置为初始状态的抽象方法。当代理被要求重置自身时，会调用此方法。\n",
    "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.produced_message_types`: 代理在其响应中可以产生的可能的 {py:class}`~autogen_agentchat.messages.ChatMessage` 消息类型列表。\n",
    "\n",
    "可选地，您可以实现 {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream` 方法来流式传输代理生成的消息。如果未实现此方法，代理\n",
    "将使用 {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream` 的默认实现，\n",
    "该实现调用 {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages` 方法并\n",
    "生成响应中的所有消息。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 倒计时代理\n",
    "\n",
    "在这个例子中，我们创建一个简单的代理，它从给定的数字倒数到零，\n",
    "并生成一个包含当前计数的消息流。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10...\n",
      "9...\n",
      "8...\n",
      "7...\n",
      "6...\n",
      "5...\n",
      "4...\n",
      "3...\n",
      "2...\n",
      "1...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "from typing import AsyncGenerator, List, Sequence\n",
    "\n",
    "from autogen_agentchat.agents import BaseChatAgent\n",
    "from autogen_agentchat.base import Response\n",
    "from autogen_agentchat.messages import AgentMessage, ChatMessage, TextMessage\n",
    "from autogen_core.base import CancellationToken\n",
    "\n",
    "\n",
    "class CountDownAgent(BaseChatAgent):\n",
    "    def __init__(self, name: str, count: int = 10):\n",
    "        super().__init__(name, \"A simple agent that counts down.\")\n",
    "        self._count = count\n",
    "\n",
    "    @property\n",
    "    def produced_message_types(self) -> List[type[ChatMessage]]:\n",
    "        return [TextMessage]\n",
    "\n",
    "    async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:\n",
    "        # Calls the on_messages_stream.\n",
    "        response: Response | None = None\n",
    "        async for message in self.on_messages_stream(messages, cancellation_token):\n",
    "            if isinstance(message, Response):\n",
    "                response = message\n",
    "        assert response is not None\n",
    "        return response\n",
    "\n",
    "    async def on_messages_stream(\n",
    "        self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken\n",
    "    ) -> AsyncGenerator[AgentMessage | Response, None]:\n",
    "        inner_messages: List[AgentMessage] = []\n",
    "        for i in range(self._count, 0, -1):\n",
    "            msg = TextMessage(content=f\"{i}...\", source=self.name)\n",
    "            inner_messages.append(msg)\n",
    "            yield msg\n",
    "        # The response is returned at the end of the stream.\n",
    "        # It contains the final message and all the inner messages.\n",
    "        yield Response(chat_message=TextMessage(content=\"Done!\", source=self.name), inner_messages=inner_messages)\n",
    "\n",
    "    async def on_reset(self, cancellation_token: CancellationToken) -> None:\n",
    "        pass\n",
    "\n",
    "\n",
    "async def run_countdown_agent() -> None:\n",
    "    # Create a countdown agent.\n",
    "    countdown_agent = CountDownAgent(\"countdown\")\n",
    "\n",
    "    # Run the agent with a given task and stream the response.\n",
    "    async for message in countdown_agent.on_messages_stream([], CancellationToken()):\n",
    "        if isinstance(message, Response):\n",
    "            print(message.chat_message.content)\n",
    "        else:\n",
    "            print(message.content)\n",
    "\n",
    "\n",
    "# Use asyncio.run(run_countdown_agent()) when running in a script.\n",
    "await run_countdown_agent()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## UserProxyAgent \n",
    "\n",
    "构建自定义代理的一个常见用例是创建一个作为用户代理的Agent。\n",
    "\n",
    "在下面的例子中，我们展示了如何实现一个 `UserProxyAgent` - 这是一个通过控制台要求用户输入文本，然后将该消息作为响应返回的代理。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hi\n"
     ]
    }
   ],
   "source": [
    "import asyncio\n",
    "from typing import List, Sequence\n",
    "\n",
    "from autogen_agentchat.agents import BaseChatAgent\n",
    "from autogen_agentchat.base import Response\n",
    "from autogen_agentchat.messages import ChatMessage\n",
    "from autogen_core.base import CancellationToken\n",
    "\n",
    "\n",
    "class UserProxyAgent(BaseChatAgent):\n",
    "    def __init__(self, name: str) -> None:\n",
    "        super().__init__(name, \"A human user.\")\n",
    "\n",
    "    @property\n",
    "    def produced_message_types(self) -> List[type[ChatMessage]]:\n",
    "        return [TextMessage]\n",
    "\n",
    "    async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:\n",
    "        user_input = await asyncio.get_event_loop().run_in_executor(None, input, \"Enter your response: \")\n",
    "        return Response(chat_message=TextMessage(content=user_input, source=self.name))\n",
    "\n",
    "    async def on_reset(self, cancellation_token: CancellationToken) -> None:\n",
    "        pass\n",
    "\n",
    "\n",
    "async def run_user_proxy_agent() -> None:\n",
    "    user_proxy_agent = UserProxyAgent(name=\"user_proxy_agent\")\n",
    "    response = await user_proxy_agent.on_messages([], CancellationToken())\n",
    "    print(response.chat_message.content)\n",
    "\n",
    "\n",
    "# Use asyncio.run(run_user_proxy_agent()) when running in a script.\n",
    "await run_user_proxy_agent()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}