{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# 自定义代理\n\n您可能需要一些行为不符合预设模式的代理。\n在这种情况下，您可以构建自定义代理。\n\nAgentChat中的所有代理都继承自{py:class}`~autogen_agentchat.agents.BaseChatAgent`类，并实现以下抽象方法和属性：\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.BaseChatMessage`消息类型列表。\n\n可选地，您可以实现{py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream`方法来流式传输代理生成的消息。\n此方法由{py:meth}`~autogen_agentchat.agents.BaseChatAgent.run_stream`调用以流式传输消息。\n如果未实现此方法，代理\n将使用{py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream`的默认实现，\n该实现会调用{py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages`方法并\n生成响应中的所有消息。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 倒计时代理\n\n在这个示例中，我们创建一个简单的代理，它从给定数字倒数到零，\n并生成包含当前计数的消息流。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "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 BaseAgentEvent, BaseChatMessage, TextMessage\n",
        "from autogen_core import CancellationToken\n",
        "\n",
        "\n",
        "class CountDownAgent(BaseChatAgent):\n",
        "    def __init__(self, name: str, count: int = 3):\n",
        "        super().__init__(name, \"A simple agent that counts down.\")\n",
        "        self._count = count\n",
        "\n",
        "    @property\n",
        "    def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:\n",
        "        return (TextMessage,)\n",
        "\n",
        "    async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:\n",
        "        # 调用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[BaseChatMessage], cancellation_token: CancellationToken\n",
        "    ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:\n",
        "        inner_messages: List[BaseAgentEvent | BaseChatMessage] = []\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",
        "        # 响应在流结束时返回。它包含最终消息和所有内部消息\n",
        "        # 。\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",
        "    # 创建一个倒计时代理。\n",
        "    countdown_agent = CountDownAgent(\"countdown\")\n",
        "\n",
        "    # 运行代理执行给定任务并流式传输响应。\n",
        "    async for message in countdown_agent.on_messages_stream([], CancellationToken()):\n",
        "        if isinstance(message, Response):\n",
        "            print(message.chat_message)\n",
        "        else:\n",
        "            print(message)\n",
        "\n",
        "\n",
        "# 在脚本中运行时使用 asyncio.run(run_countdown_agent())。\n",
        "await run_countdown_agent()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 算术代理\n\n在这个示例中，我们创建一个代理类，可以对给定整数执行简单算术运算。然后，我们将使用这个代理类的不同实例\n在 {py:class}`~autogen_agentchat.teams.SelectorGroupChat` 中\n通过应用一系列算术运算将给定整数转换为另一个整数。\n\n`ArithmeticAgent` 类接收一个 `operator_func`，该函数接收一个整数并返回一个整数，\n在对该整数应用算术运算之后。\n在其 `on_messages` 方法中，它将 `operator_func` 应用于输入消息中的整数，\n并返回包含结果响应。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from typing import Callable, Sequence\n",
        "\n",
        "from autogen_agentchat.agents import BaseChatAgent\n",
        "from autogen_agentchat.base import Response\n",
        "from autogen_agentchat.conditions import MaxMessageTermination\n",
        "from autogen_agentchat.messages import BaseChatMessage\n",
        "from autogen_agentchat.teams import SelectorGroupChat\n",
        "from autogen_agentchat.ui import Console\n",
        "from autogen_core import CancellationToken\n",
        "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
        "\n",
        "\n",
        "class ArithmeticAgent(BaseChatAgent):\n",
        "    def __init__(self, name: str, description: str, operator_func: Callable[[int], int]) -> None:\n",
        "        super().__init__(name, description=description)\n",
        "        self._operator_func = operator_func\n",
        "        self._message_history: List[BaseChatMessage] = []\n",
        "\n",
        "    @property\n",
        "    def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:\n",
        "        return (TextMessage,)\n",
        "\n",
        "    async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:\n",
        "        # 更新消息历史记录。注意：消息可能为空列表，\n",
        "        # 这意味着代理之前已被选中。\n",
        "        self._message_history.extend(messages)\n",
        "        # 解析最后一条消息中的数字。\n",
        "        assert isinstance(self._message_history[-1], TextMessage)\n",
        "        number = int(self._message_history[-1].content)\n",
        "        # 对数字应用运算符函数。\n",
        "        result = self._operator_func(number)\n",
        "        # 用结果创建新消息。\n",
        "        response_message = TextMessage(content=str(result), source=self.name)\n",
        "        # 更新消息历史记录。\n",
        "        self._message_history.append(response_message)\n",
        "        # 返回响应。\n",
        "        return Response(chat_message=response_message)\n",
        "\n",
        "    async def on_reset(self, cancellation_token: CancellationToken) -> None:\n",
        "        pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "```{note}\n`on_messages` 方法可能会被传入空消息列表调用，这种情况意味着代理之前已被调用过，现在再次被调用时，\n调用者没有发送任何新消息。因此保持代理接收到的先前消息历史记录非常重要，\n并使用该历史记录来生成响应。\n```\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "现在我们可以创建一个包含5个 `ArithmeticAgent` 实例的 {py:class}`~autogen_agentchat.teams.SelectorGroupChat`：\n\n- 一个将输入整数加1的代理，\n- 一个将输入整数减1的代理，\n- 一个将输入整数乘以2的代理，\n- 一个将输入整数除以2并向下取整的代理，\n- 一个保持输入整数不变的代理。\n\n然后我们用这些代理创建一个 {py:class}`~autogen_agentchat.teams.SelectorGroupChat`，\n并设置适当的选择器配置：\n\n- 允许同一个代理连续被选中以支持重复操作，\n- 自定义选择器提示词来使模型响应更适合特定任务。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "---------- user ----------\n",
            "Apply the operations to turn the given number into 25.\n",
            "---------- user ----------\n",
            "10\n",
            "---------- multiply_agent ----------\n",
            "20\n",
            "---------- add_agent ----------\n",
            "21\n",
            "---------- multiply_agent ----------\n",
            "42\n",
            "---------- divide_agent ----------\n",
            "21\n",
            "---------- add_agent ----------\n",
            "22\n",
            "---------- add_agent ----------\n",
            "23\n",
            "---------- add_agent ----------\n",
            "24\n",
            "---------- add_agent ----------\n",
            "25\n",
            "---------- Summary ----------\n",
            "Number of messages: 10\n",
            "Finish reason: Maximum number of messages 10 reached, current message count: 10\n",
            "Total prompt tokens: 0\n",
            "Total completion tokens: 0\n",
            "Duration: 2.40 seconds\n"
          ]
        }
      ],
      "source": [
        "async def run_number_agents() -> None:\n",
        "    # 创建用于数字运算的代理。\n",
        "    add_agent = ArithmeticAgent(\"add_agent\", \"Adds 1 to the number.\", lambda x: x + 1)\n",
        "    multiply_agent = ArithmeticAgent(\"multiply_agent\", \"Multiplies the number by 2.\", lambda x: x * 2)\n",
        "    subtract_agent = ArithmeticAgent(\"subtract_agent\", \"Subtracts 1 from the number.\", lambda x: x - 1)\n",
        "    divide_agent = ArithmeticAgent(\"divide_agent\", \"Divides the number by 2 and rounds down.\", lambda x: x // 2)\n",
        "    identity_agent = ArithmeticAgent(\"identity_agent\", \"Returns the number as is.\", lambda x: x)\n",
        "\n",
        "    # 终止条件是发送10条消息后停止。\n",
        "    termination_condition = MaxMessageTermination(10)\n",
        "\n",
        "    # 创建一个选择器群聊。\n",
        "    selector_group_chat = SelectorGroupChat(\n",
        "        [add_agent, multiply_agent, subtract_agent, divide_agent, identity_agent],\n",
        "        model_client=OpenAIChatCompletionClient(model=\"gpt-4o\"),\n",
        "        termination_condition=termination_condition,\n",
        "        allow_repeated_speaker=True,  # 允许同一个代理多次发言，这对本任务很有必要。\n",
        "        selector_prompt=(\n",
        "            \"Available roles:\\n{roles}\\nTheir job descriptions:\\n{participants}\\n\"\n",
        "            \"Current conversation history:\\n{history}\\n\"\n",
        "            \"Please select the most appropriate role for the next message, and only return the role name.\"\n",
        "        ),\n",
        "    )\n",
        "\n",
        "    # 运行带有给定任务的selector群聊并流式传输响应。\n",
        "    task: List[BaseChatMessage] = [\n",
        "        TextMessage(content=\"Apply the operations to turn the given number into 25.\", source=\"user\"),\n",
        "        TextMessage(content=\"10\", source=\"user\"),\n",
        "    ]\n",
        "    stream = selector_group_chat.run_stream(task=task)\n",
        "    await Console(stream)\n",
        "\n",
        "\n",
        "# 在脚本中运行时使用asyncio.run(run_number_agents())。\n",
        "await run_number_agents()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "从输出中可以看到，代理们通过依次选择应用算术运算的合适代理，成功地将输入整数从10转换为了25。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 在自定义代理中使用自定义模型客户端\n\nAgentChat中{py:class}`~autogen_agentchat.agents.AssistantAgent`预设的一个关键特性是它接受`model_client`参数，并可以在响应消息时使用它。然而，在某些情况下，您可能希望您的代理使用当前不支持的自定义模型客户端（参见[支持的模型客户端](https://microsoft.github.io/autogen/dev/user-guide/core-user-guide/components/model-clients.html)）或自定义模型行为。\n\n您可以通过实现*您的自定义模型客户端*的自定义代理来实现这一点。\n\n在下面的示例中，我们将演示一个直接使用[Google Gemini SDK](https://github.com/googleapis/python-genai)来响应消息的自定义代理。\n\n> **注意：** 您需要安装[Google Gemini SDK](https://github.com/googleapis/python-genai)才能运行此示例。可以使用以下命令安装：\n\n```bash\npip install google-genai\n```\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# !pip install google-genai\n",
        "import os\n",
        "from typing import AsyncGenerator, Sequence\n",
        "\n",
        "from autogen_agentchat.agents import BaseChatAgent\n",
        "from autogen_agentchat.base import Response\n",
        "from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage\n",
        "from autogen_core import CancellationToken\n",
        "from autogen_core.model_context import UnboundedChatCompletionContext\n",
        "from autogen_core.models import AssistantMessage, RequestUsage, UserMessage\n",
        "from google import genai\n",
        "from google.genai import types\n",
        "\n",
        "\n",
        "class GeminiAssistantAgent(BaseChatAgent):\n",
        "    def __init__(\n",
        "        self,\n",
        "        name: str,\n",
        "        description: str = \"An agent that provides assistance with ability to use tools.\",\n",
        "        model: str = \"gemini-1.5-flash-002\",\n",
        "        api_key: str = os.environ[\"GEMINI_API_KEY\"],\n",
        "        system_message: str\n",
        "        | None = \"You are a helpful assistant that can respond to messages. Reply with TERMINATE when the task has been completed.\",\n",
        "    ):\n",
        "        super().__init__(name=name, description=description)\n",
        "        self._model_context = UnboundedChatCompletionContext()\n",
        "        self._model_client = genai.Client(api_key=api_key)\n",
        "        self._system_message = system_message\n",
        "        self._model = model\n",
        "\n",
        "    @property\n",
        "    def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:\n",
        "        return (TextMessage,)\n",
        "\n",
        "    async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:\n",
        "        final_response = None\n",
        "        async for message in self.on_messages_stream(messages, cancellation_token):\n",
        "            if isinstance(message, Response):\n",
        "                final_response = message\n",
        "\n",
        "        if final_response is None:\n",
        "            raise AssertionError(\"The stream should have returned the final result.\")\n",
        "\n",
        "        return final_response\n",
        "\n",
        "    async def on_messages_stream(\n",
        "        self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken\n",
        "    ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:\n",
        "        # 将消息添加到模型上下文\n",
        "        for msg in messages:\n",
        "            await self._model_context.add_message(msg.to_model_message())\n",
        "\n",
        "        # 获取对话历史记录\n",
        "        history = [\n",
        "            (msg.source if hasattr(msg, \"source\") else \"system\")\n",
        "            + \": \"\n",
        "            + (msg.content if isinstance(msg.content, str) else \"\")\n",
        "            + \"\\n\"\n",
        "            for msg in await self._model_context.get_messages()\n",
        "        ]\n",
        "        # 使用Gemini生成响应\n",
        "        response = self._model_client.models.generate_content(\n",
        "            model=self._model,\n",
        "            contents=f\"History: {history}\\nGiven the history, please provide a response\",\n",
        "            config=types.GenerateContentConfig(\n",
        "                system_instruction=self._system_message,\n",
        "                temperature=0.3,\n",
        "            ),\n",
        "        )\n",
        "\n",
        "        # 创建使用元数据\n",
        "        usage = RequestUsage(\n",
        "            prompt_tokens=response.usage_metadata.prompt_token_count,\n",
        "            completion_tokens=response.usage_metadata.candidates_token_count,\n",
        "        )\n",
        "\n",
        "        # 将响应添加到模型上下文\n",
        "        await self._model_context.add_message(AssistantMessage(content=response.text, source=self.name))\n",
        "\n",
        "        # 生成最终响应\n",
        "        yield Response(\n",
        "            chat_message=TextMessage(content=response.text, source=self.name, models_usage=usage),\n",
        "            inner_messages=[],\n",
        "        )\n",
        "\n",
        "    async def on_reset(self, cancellation_token: CancellationToken) -> None:\n",
        "        \"\"\"Reset the assistant by clearing the model context.\"\"\"\n",
        "        await self._model_context.clear()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "---------- user ----------\n",
            "What is the capital of New York?\n",
            "---------- gemini_assistant ----------\n",
            "Albany\n",
            "TERMINATE\n",
            "\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What is the capital of New York?', type='TextMessage'), TextMessage(source='gemini_assistant', models_usage=RequestUsage(prompt_tokens=46, completion_tokens=5), content='Albany\\nTERMINATE\\n', type='TextMessage')], stop_reason=None)"
            ]
          },
          "execution_count": 38,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "gemini_assistant = GeminiAssistantAgent(\"gemini_assistant\")\n",
        "await Console(gemini_assistant.run_stream(task=\"What is the capital of New York?\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "在上面的示例中，我们选择提供 `model`、`api_key` 和 `system_message` 作为参数 - 您可以选择提供模型客户端所需的任何其他参数，或符合您的应用程序设计。\n\n现在，让我们探索如何将这个自定义智能体作为团队的一部分在 AgentChat 中使用。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "---------- user ----------\n",
            "Write a Haiku poem with 4 lines about the fall season.\n",
            "---------- primary ----------\n",
            "Crimson leaves cascade,  \n",
            "Whispering winds sing of change,  \n",
            "Chill wraps the fading,  \n",
            "Nature's quilt, rich and warm.\n",
            "---------- gemini_critic ----------\n",
            "The poem is good, but it has four lines instead of three.  A haiku must have three lines with a 5-7-5 syllable structure.  The content is evocative of autumn, but the form is incorrect.  Please revise to adhere to the haiku's syllable structure.\n",
            "\n",
            "---------- primary ----------\n",
            "Thank you for your feedback! Here’s a revised haiku that follows the 5-7-5 syllable structure:\n",
            "\n",
            "Crimson leaves drift down,  \n",
            "Chill winds whisper through the gold,  \n",
            "Autumn’s breath is near.\n",
            "---------- gemini_critic ----------\n",
            "The revised haiku is much improved.  It correctly follows the 5-7-5 syllable structure and maintains the evocative imagery of autumn.  APPROVE\n",
            "\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Write a Haiku poem with 4 lines about the fall season.', type='TextMessage'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=33, completion_tokens=31), content=\"Crimson leaves cascade,  \\nWhispering winds sing of change,  \\nChill wraps the fading,  \\nNature's quilt, rich and warm.\", type='TextMessage'), TextMessage(source='gemini_critic', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=60), content=\"The poem is good, but it has four lines instead of three.  A haiku must have three lines with a 5-7-5 syllable structure.  The content is evocative of autumn, but the form is incorrect.  Please revise to adhere to the haiku's syllable structure.\\n\", type='TextMessage'), TextMessage(source='primary', models_usage=RequestUsage(prompt_tokens=141, completion_tokens=49), content='Thank you for your feedback! Here’s a revised haiku that follows the 5-7-5 syllable structure:\\n\\nCrimson leaves drift down,  \\nChill winds whisper through the gold,  \\nAutumn’s breath is near.', type='TextMessage'), TextMessage(source='gemini_critic', models_usage=RequestUsage(prompt_tokens=211, completion_tokens=32), content='The revised haiku is much improved.  It correctly follows the 5-7-5 syllable structure and maintains the evocative imagery of autumn.  APPROVE\\n', type='TextMessage')], stop_reason=\"Text 'APPROVE' mentioned\")"
            ]
          },
          "execution_count": 39,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from autogen_agentchat.agents import AssistantAgent\n",
        "from autogen_agentchat.conditions import TextMentionTermination\n",
        "from autogen_agentchat.teams import RoundRobinGroupChat\n",
        "from autogen_agentchat.ui import Console\n",
        "\n",
        "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-mini\")\n",
        "\n",
        "# 创建主智能体\n",
        "primary_agent = AssistantAgent(\n",
        "    \"primary\",\n",
        "    model_client=model_client,\n",
        "    system_message=\"You are a helpful AI assistant.\",\n",
        ")\n",
        "\n",
        "# 基于我们新的 GeminiAssistantAgent 创建评论智能体\n",
        "gemini_critic_agent = GeminiAssistantAgent(\n",
        "    \"gemini_critic\",\n",
        "    system_message=\"Provide constructive feedback. Respond with 'APPROVE' to when your feedbacks are addressed.\",\n",
        ")\n",
        "\n",
        "\n",
        "# 定义一个终止条件，当评论者批准或达到 10 条消息后停止任务\n",
        "termination = TextMentionTermination(\"APPROVE\") | MaxMessageTermination(10)\n",
        "\n",
        "# 创建一个包含主要和评论代理的团队。\n",
        "team = RoundRobinGroupChat([primary_agent, gemini_critic_agent], termination_condition=termination)\n",
        "\n",
        "await Console(team.run_stream(task=\"Write a Haiku poem with 4 lines about the fall season.\"))\n",
        "await model_client.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "在上面的部分中，我们展示了几个非常重要的概念：\n- 我们开发了一个自定义代理，使用Google Gemini SDK来响应消息。\n- 我们展示这个自定义代理可以作为更广泛的AgentChat生态系统的一部分使用 - 在本例中作为{py:class}`~autogen_agentchat.teams.RoundRobinGroupChat`的参与者，只要它继承自{py:class}`~autogen_agentchat.agents.BaseChatAgent`。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 使自定义代理具有声明性\n\nAutogen提供了一个[组件](https://microsoft.github.io/autogen/dev/user-guide/core-user-guide/framework/component-config.html)接口，用于将组件配置序列化为声明性格式。这对于保存和加载配置，以及与他人共享配置非常有用。\n\n我们通过继承`Component`类并实现`_from_config`和`_to_config`方法来实现这一点。\n声明性类可以使用`dump_component`方法序列化为JSON格式，并使用`load_component`方法从JSON格式反序列化。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "from typing import AsyncGenerator, Sequence\n",
        "\n",
        "from autogen_agentchat.agents import BaseChatAgent\n",
        "from autogen_agentchat.base import Response\n",
        "from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage\n",
        "from autogen_core import CancellationToken, Component\n",
        "from pydantic import BaseModel\n",
        "from typing_extensions import Self\n",
        "\n",
        "\n",
        "class GeminiAssistantAgentConfig(BaseModel):\n",
        "    name: str\n",
        "    description: str = \"An agent that provides assistance with ability to use tools.\"\n",
        "    model: str = \"gemini-1.5-flash-002\"\n",
        "    system_message: str | None = None\n",
        "\n",
        "\n",
        "class GeminiAssistantAgent(BaseChatAgent, Component[GeminiAssistantAgentConfig]):  # type: ignore[no-redef]\n",
        "    component_config_schema = GeminiAssistantAgentConfig\n",
        "    # component_provider_override = \"mypackage.agents.GeminiAssistantAgent\"\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        name: str,\n",
        "        description: str = \"An agent that provides assistance with ability to use tools.\",\n",
        "        model: str = \"gemini-1.5-flash-002\",\n",
        "        api_key: str = os.environ[\"GEMINI_API_KEY\"],\n",
        "        system_message: str\n",
        "        | None = \"You are a helpful assistant that can respond to messages. Reply with TERMINATE when the task has been completed.\",\n",
        "    ):\n",
        "        super().__init__(name=name, description=description)\n",
        "        self._model_context = UnboundedChatCompletionContext()\n",
        "        self._model_client = genai.Client(api_key=api_key)\n",
        "        self._system_message = system_message\n",
        "        self._model = model\n",
        "\n",
        "    @property\n",
        "    def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:\n",
        "        return (TextMessage,)\n",
        "\n",
        "    async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:\n",
        "        final_response = None\n",
        "        async for message in self.on_messages_stream(messages, cancellation_token):\n",
        "            if isinstance(message, Response):\n",
        "                final_response = message\n",
        "\n",
        "        if final_response is None:\n",
        "            raise AssertionError(\"The stream should have returned the final result.\")\n",
        "\n",
        "        return final_response\n",
        "\n",
        "    async def on_messages_stream(\n",
        "        self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken\n",
        "    ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:\n",
        "        # 将消息添加到模型上下文\n",
        "        for msg in messages:\n",
        "            await self._model_context.add_message(msg.to_model_message())\n",
        "\n",
        "        # 获取对话历史记录\n",
        "        history = [\n",
        "            (msg.source if hasattr(msg, \"source\") else \"system\")\n",
        "            + \": \"\n",
        "            + (msg.content if isinstance(msg.content, str) else \"\")\n",
        "            + \"\\n\"\n",
        "            for msg in await self._model_context.get_messages()\n",
        "        ]\n",
        "\n",
        "        # 使用Gemini生成响应\n",
        "        response = self._model_client.models.generate_content(\n",
        "            model=self._model,\n",
        "            contents=f\"History: {history}\\nGiven the history, please provide a response\",\n",
        "            config=types.GenerateContentConfig(\n",
        "                system_instruction=self._system_message,\n",
        "                temperature=0.3,\n",
        "            ),\n",
        "        )\n",
        "\n",
        "        # 创建使用元数据\n",
        "        usage = RequestUsage(\n",
        "            prompt_tokens=response.usage_metadata.prompt_token_count,\n",
        "            completion_tokens=response.usage_metadata.candidates_token_count,\n",
        "        )\n",
        "\n",
        "        # 将响应添加到模型上下文\n",
        "        await self._model_context.add_message(AssistantMessage(content=response.text, source=self.name))\n",
        "\n",
        "        # 生成最终响应\n",
        "        yield Response(\n",
        "            chat_message=TextMessage(content=response.text, source=self.name, models_usage=usage),\n",
        "            inner_messages=[],\n",
        "        )\n",
        "\n",
        "    async def on_reset(self, cancellation_token: CancellationToken) -> None:\n",
        "        \"\"\"Reset the assistant by clearing the model context.\"\"\"\n",
        "        await self._model_context.clear()\n",
        "\n",
        "    @classmethod\n",
        "    def _from_config(cls, config: GeminiAssistantAgentConfig) -> Self:\n",
        "        return cls(\n",
        "            name=config.name, description=config.description, model=config.model, system_message=config.system_message\n",
        "        )\n",
        "\n",
        "    def _to_config(self) -> GeminiAssistantAgentConfig:\n",
        "        return GeminiAssistantAgentConfig(\n",
        "            name=self.name,\n",
        "            description=self.description,\n",
        "            model=self._model,\n",
        "            system_message=self._system_message,\n",
        "        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "现在我们已经实现了所需的方法，可以加载自定义代理并将其转储为JSON格式，然后从JSON格式重新加载代理。\n\n> 注意：您需要将`component_provider_override`类变量设置为包含自定义代理类的模块完整路径（例如`mypackage.agents.GeminiAssistantAgent`）。这将被`load_component`方法用来确定如何实例化该类。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{\n",
            "  \"provider\": \"__main__.GeminiAssistantAgent\",\n",
            "  \"component_type\": \"agent\",\n",
            "  \"version\": 1,\n",
            "  \"component_version\": 1,\n",
            "  \"description\": null,\n",
            "  \"label\": \"GeminiAssistantAgent\",\n",
            "  \"config\": {\n",
            "    \"name\": \"gemini_assistant\",\n",
            "    \"description\": \"An agent that provides assistance with ability to use tools.\",\n",
            "    \"model\": \"gemini-1.5-flash-002\",\n",
            "    \"system_message\": \"You are a helpful assistant that can respond to messages. Reply with TERMINATE when the task has been completed.\"\n",
            "  }\n",
            "}\n",
            "<__main__.GeminiAssistantAgent object at 0x11a5c5a90>\n"
          ]
        }
      ],
      "source": [
        "gemini_assistant = GeminiAssistantAgent(\"gemini_assistant\")\n",
        "config = gemini_assistant.dump_component()\n",
        "print(config.model_dump_json(indent=2))\n",
        "loaded_agent = GeminiAssistantAgent.load_component(config)\n",
        "print(loaded_agent)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 后续步骤\n\n到目前为止，我们已经了解了如何创建自定义代理、向代理添加自定义模型客户端以及使自定义代理具有声明性。这个基础示例可以通过以下几种方式进行扩展：\n\n- 扩展Gemini模型客户端以处理类似于{py:class}`~autogen_agentchat.agents.AssistantAgent`类的函数调用功能。参考https://ai.google.dev/gemini-api/docs/function-calling\n- 实现一个包含自定义代理的包，并尝试在[AutoGen Studio](https://microsoft.github.io/autogen/stable/user-guide/autogenstudio-user-guide/index.html)等工具中使用其声明式格式进行实验。\n"
      ]
    }
  ],
  "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.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}