{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# 状态管理\n\n到目前为止，我们已经讨论了如何在多智能体应用中构建组件——包括智能体、团队和终止条件。在许多情况下，将这些组件的状态保存到磁盘并在之后重新加载会非常有用。这在Web应用中尤其重要，因为无状态的端点需要响应请求，并从持久化存储中加载应用状态。\n\n本笔记本将探讨如何保存和加载智能体、团队以及终止条件的状态。\n\n\n## 保存与加载智能体\n\n我们可以通过调用{py:class}`~autogen_agentchat.agents.AssistantAgent`上的{py:meth}`~autogen_agentchat.agents.AssistantAgent.save_state`方法来获取智能体的状态。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "In Tanganyika's embrace so wide and deep,  \n",
            "Ancient waters cradle secrets they keep,  \n",
            "Echoes of time where horizons sleep.  \n"
          ]
        }
      ],
      "source": [
        "from autogen_agentchat.agents import AssistantAgent\n",
        "from autogen_agentchat.conditions import MaxMessageTermination\n",
        "from autogen_agentchat.messages import TextMessage\n",
        "from autogen_agentchat.teams import RoundRobinGroupChat\n",
        "from autogen_agentchat.ui import Console\n",
        "from autogen_core import CancellationToken\n",
        "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
        "\n",
        "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-2024-08-06\")\n",
        "\n",
        "assistant_agent = AssistantAgent(\n",
        "    name=\"assistant_agent\",\n",
        "    system_message=\"You are a helpful assistant\",\n",
        "    model_client=model_client,\n",
        ")\n",
        "\n",
        "# 在脚本中运行时使用asyncio.run(...)。\n",
        "response = await assistant_agent.on_messages(\n",
        "    [TextMessage(content=\"Write a 3 line poem on lake tangayika\", source=\"user\")], CancellationToken()\n",
        ")\n",
        "print(response.chat_message)\n",
        "await model_client.close()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{'type': 'AssistantAgentState', 'version': '1.0.0', 'llm_messages': [{'content': 'Write a 3 line poem on lake tangayika', 'source': 'user', 'type': 'UserMessage'}, {'content': \"In Tanganyika's embrace so wide and deep,  \\nAncient waters cradle secrets they keep,  \\nEchoes of time where horizons sleep.  \", 'source': 'assistant_agent', 'type': 'AssistantMessage'}]}\n"
          ]
        }
      ],
      "source": [
        "agent_state = await assistant_agent.save_state()\n",
        "print(agent_state)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "The last line of the poem was: \"Echoes of time where horizons sleep.\"\n"
          ]
        }
      ],
      "source": [
        "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-2024-08-06\")\n",
        "\n",
        "new_assistant_agent = AssistantAgent(\n",
        "    name=\"assistant_agent\",\n",
        "    system_message=\"You are a helpful assistant\",\n",
        "    model_client=model_client,\n",
        ")\n",
        "await new_assistant_agent.load_state(agent_state)\n",
        "\n",
        "# 在脚本中运行时使用asyncio.run(...)。\n",
        "response = await new_assistant_agent.on_messages(\n",
        "    [TextMessage(content=\"What was the last line of the previous poem you wrote\", source=\"user\")], CancellationToken()\n",
        ")\n",
        "print(response.chat_message)\n",
        "await model_client.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "```{note}\n对于{py:class}`~autogen_agentchat.agents.AssistantAgent`，其状态包含model_context。\n如果您编写自定义智能体，建议重写{py:meth}`~autogen_agentchat.agents.BaseChatAgent.save_state`和{py:meth}`~autogen_agentchat.agents.BaseChatAgent.load_state`方法来自定义行为。默认实现会保存和加载空状态。\n```\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 保存与加载团队\n\n我们可以通过调用团队的`save_state`方法来获取团队状态，并通过调用团队的`load_state`方法重新加载。\n\n当调用团队的`save_state`方法时，它会保存团队中所有智能体的状态。\n\n我们将从创建一个简单的{py:class}`~autogen_agentchat.teams.RoundRobinGroupChat`团队开始，该团队包含单个智能体，并要求它写一首诗。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "---------- user ----------\n",
            "Write a beautiful poem 3-line about lake tangayika\n",
            "---------- assistant_agent ----------\n",
            "In Tanganyika's gleam, beneath the azure skies,  \n",
            "Whispers of ancient waters, in tranquil guise,  \n",
            "Nature's mirror, where dreams and serenity lie.\n",
            "[Prompt tokens: 29, Completion tokens: 34]\n",
            "---------- Summary ----------\n",
            "Number of messages: 2\n",
            "Finish reason: Maximum number of messages 2 reached, current message count: 2\n",
            "Total prompt tokens: 29\n",
            "Total completion tokens: 34\n",
            "Duration: 0.71 seconds\n"
          ]
        }
      ],
      "source": [
        "model_client = OpenAIChatCompletionClient(model=\"gpt-4o-2024-08-06\")\n",
        "\n",
        "# 定义一个团队。\n",
        "assistant_agent = AssistantAgent(\n",
        "    name=\"assistant_agent\",\n",
        "    system_message=\"You are a helpful assistant\",\n",
        "    model_client=model_client,\n",
        ")\n",
        "agent_team = RoundRobinGroupChat([assistant_agent], termination_condition=MaxMessageTermination(max_messages=2))\n",
        "\n",
        "# 运行团队并将消息流式传输到控制台。\n",
        "stream = agent_team.run_stream(task=\"Write a beautiful poem 3-line about lake tangayika\")\n",
        "\n",
        "# 在脚本中运行时使用 asyncio.run(...)。\n",
        "await Console(stream)\n",
        "\n",
        "# 保存智能体团队的状态。\n",
        "team_state = await agent_team.save_state()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "如果我们重置团队（模拟团队的实例化），并提问`你写的诗的最后一行是什么？`，我们会发现团队无法完成此任务，因为没有对之前运行的引用。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "---------- user ----------\n",
            "What was the last line of the poem you wrote?\n",
            "---------- assistant_agent ----------\n",
            "I'm sorry, but I am unable to recall or access previous interactions, including any specific poem I may have composed in our past conversations. If you like, I can write a new poem for you.\n",
            "[Prompt tokens: 28, Completion tokens: 40]\n",
            "---------- Summary ----------\n",
            "Number of messages: 2\n",
            "Finish reason: Maximum number of messages 2 reached, current message count: 2\n",
            "Total prompt tokens: 28\n",
            "Total completion tokens: 40\n",
            "Duration: 0.70 seconds\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What was the last line of the poem you wrote?', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=28, completion_tokens=40), content=\"I'm sorry, but I am unable to recall or access previous interactions, including any specific poem I may have composed in our past conversations. If you like, I can write a new poem for you.\", type='TextMessage')], stop_reason='Maximum number of messages 2 reached, current message count: 2')"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "await agent_team.reset()\n",
        "stream = agent_team.run_stream(task=\"What was the last line of the poem you wrote?\")\n",
        "await Console(stream)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "接下来，我们加载团队的状态并询问相同的问题。可以看到团队能够准确返回它写的诗的最后一行。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{'type': 'TeamState', 'version': '1.0.0', 'agent_states': {'group_chat_manager/a55364ad-86fd-46ab-9449-dcb5260b1e06': {'type': 'RoundRobinManagerState', 'version': '1.0.0', 'message_thread': [{'source': 'user', 'models_usage': None, 'content': 'Write a beautiful poem 3-line about lake tangayika', 'type': 'TextMessage'}, {'source': 'assistant_agent', 'models_usage': {'prompt_tokens': 29, 'completion_tokens': 34}, 'content': \"In Tanganyika's gleam, beneath the azure skies,  \\nWhispers of ancient waters, in tranquil guise,  \\nNature's mirror, where dreams and serenity lie.\", 'type': 'TextMessage'}], 'current_turn': 0, 'next_speaker_index': 0}, 'collect_output_messages/a55364ad-86fd-46ab-9449-dcb5260b1e06': {}, 'assistant_agent/a55364ad-86fd-46ab-9449-dcb5260b1e06': {'type': 'ChatAgentContainerState', 'version': '1.0.0', 'agent_state': {'type': 'AssistantAgentState', 'version': '1.0.0', 'llm_messages': [{'content': 'Write a beautiful poem 3-line about lake tangayika', 'source': 'user', 'type': 'UserMessage'}, {'content': \"In Tanganyika's gleam, beneath the azure skies,  \\nWhispers of ancient waters, in tranquil guise,  \\nNature's mirror, where dreams and serenity lie.\", 'source': 'assistant_agent', 'type': 'AssistantMessage'}]}, 'message_buffer': []}}, 'team_id': 'a55364ad-86fd-46ab-9449-dcb5260b1e06'}\n",
            "---------- user ----------\n",
            "What was the last line of the poem you wrote?\n",
            "---------- assistant_agent ----------\n",
            "The last line of the poem I wrote is:  \n",
            "\"Nature's mirror, where dreams and serenity lie.\"\n",
            "[Prompt tokens: 86, Completion tokens: 22]\n",
            "---------- Summary ----------\n",
            "Number of messages: 2\n",
            "Finish reason: Maximum number of messages 2 reached, current message count: 2\n",
            "Total prompt tokens: 86\n",
            "Total completion tokens: 22\n",
            "Duration: 0.96 seconds\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What was the last line of the poem you wrote?', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=22), content='The last line of the poem I wrote is:  \\n\"Nature\\'s mirror, where dreams and serenity lie.\"', type='TextMessage')], stop_reason='Maximum number of messages 2 reached, current message count: 2')"
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "print(team_state)\n",
        "\n",
        "# 加载团队状态。\n",
        "await agent_team.load_state(team_state)\n",
        "stream = agent_team.run_stream(task=\"What was the last line of the poem you wrote?\")\n",
        "await Console(stream)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 持久化状态（文件或数据库）\n\n在许多情况下，我们可能希望将团队的状态持久化到磁盘（或数据库）中，并在之后重新加载。状态是一个可以序列化到文件或写入数据库的字典。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "---------- user ----------\n",
            "What was the last line of the poem you wrote?\n",
            "---------- assistant_agent ----------\n",
            "The last line of the poem I wrote is:  \n",
            "\"Nature's mirror, where dreams and serenity lie.\"\n",
            "[Prompt tokens: 86, Completion tokens: 22]\n",
            "---------- Summary ----------\n",
            "Number of messages: 2\n",
            "Finish reason: Maximum number of messages 2 reached, current message count: 2\n",
            "Total prompt tokens: 86\n",
            "Total completion tokens: 22\n",
            "Duration: 0.72 seconds\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What was the last line of the poem you wrote?', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=22), content='The last line of the poem I wrote is:  \\n\"Nature\\'s mirror, where dreams and serenity lie.\"', type='TextMessage')], stop_reason='Maximum number of messages 2 reached, current message count: 2')"
            ]
          },
          "execution_count": 13,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import json\n",
        "\n",
        "# # 将状态保存到磁盘\n",
        "\n",
        "with open(\"coding/team_state.json\", \"w\") as f:\n",
        "    json.dump(team_state, f)\n",
        "\n",
        "# # 从磁盘加载状态\n",
        "with open(\"coding/team_state.json\", \"r\") as f:\n",
        "    team_state = json.load(f)\n",
        "\n",
        "new_agent_team = RoundRobinGroupChat([assistant_agent], termination_condition=MaxMessageTermination(max_messages=2))\n",
        "await new_agent_team.load_state(team_state)\n",
        "stream = new_agent_team.run_stream(task=\"What was the last line of the poem you wrote?\")\n",
        "await Console(stream)\n",
        "await model_client.close()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "agnext",
      "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.11.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}