AutoGen学习笔记系列（七）Tutorial - Managing State

这篇文章瞄准的是AutoGen框架官方教程中的 Tutorial 章节中的 Managing State 小节，主要介绍了如何对Team内的状态管理，特别是如何 保存 与 加载 状态，这对于Agent系统而言非常重要。

官网链接：https://microsoft.github.io/autogen/stable/user-guide/agentchat-user-guide/tutorial/state.html# ；

【注意】：AutoGen库中有一个函数名存在语意歧义，save_state() 实际是获取状态，至于如何本地化是需要自己完成的，但好在其状态的描述都是json格式，所以这部分很容易实现。

---

Managing State

在之前的文章中详细介绍了如何构建与使用Team、如何控制Team内部轮询节奏等。我们也提到过只要不调用 team.reset() 方法，Team内部的上下文状态是会持续被保留的。那么为什么会需要将这个上下文状态写到本地硬盘上呢？

1. 节省成本：如果你想构建自己私有Team并让其拥有长期记忆，那么保存历史状态就是必须的，否则每次启动后都需要将你和Team的所有历史对话都重新发一遍，token会被快速消耗；
2. 代码测试：在调试过程中难免会抛出异常，为了避免程序被意外中断，保存上下文信息也是非常必要的；
3. 方便迁移：上下文状态相当于一个人的记忆，只要Team结构不变就可以将这个记忆复制很多份在其他地方使用，并行任务的神器；

AutoGen提供了Agent、Team快捷方便的存储方法 save_state() 成员函数，其返回值是一个python的dict类型，能够很容易地转换为json文本，这意味着我们甚至可以手动修改Agent、Team的状态信息。

---

Saving and Loading Agents

对于Agent而言可以使用save_state() 获取 状态，用 load_state() 加载状态。

• 首先写一个demo，在运行后将Agent的状态保存到硬盘上：

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.messages import TextMessage
from autogen_core import CancellationToken
from autogen_ext.models.openai import OpenAIChatCompletionClient
import os, asyncio, jsonos.environ["OPENAI_API_KEY"] = "你的OpenAI API Key"
assistant_agent = AssistantAgent(name="assistant_agent",system_message="You are a helpful assistant",model_client=OpenAIChatCompletionClient(model="gpt-4o-2024-08-06",),
)#----------------------------------------------------#
# 写一个函数用来本地化Agent的状态信息
def persisting_state(file_name:str, state):with open(file_name, "w") as f:json.dump(state, f)# 让Agent执行一小段交互
async def main():response = await assistant_agent.on_messages([TextMessage(content="Write a 3 line poem on lake tangayika", source="user")], CancellationToken())print(response.chat_message.content)print('-' * 50)# 等待运行完成后在获取其状态agent_state = await assistant_agent.save_state()print(agent_state)# 将Agent状态写入本地persisting_state("./agent_state.json", agent_state)asyncio.run(main())

运行结果如下：

$ python demo.py 

当前目录下会生成一个这个Agent的状态信息，里面的messages字段就是Agent和LLM之间的聊天记录，也就是所谓的记忆：

{"type": "AssistantAgentState","version": "1.0.0","llm_context": {"messages": [{"content": "Write a 3 line poem on lake tangayika","source": "user","type": "UserMessage"},{"content": "In waters deep, where ancient secrets keep,  \nTanganyika mirrors the sky's boundless sweep,  \nSilent guardian of tales and dreams, it sleeps.  ","source": "assistant_agent","type": "AssistantMessage"}]}
}

• 然后再写一个demo从硬盘上读取Agent状态：

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.messages import TextMessage
from autogen_core import CancellationToken
from autogen_ext.models.openai import OpenAIChatCompletionClient
import os, asyncio, jsonos.environ["OPENAI_API_KEY"] = "你的OpenAI API Key"#----------------------------------------------------#
# 写一个函数用来加载本地json文件
def load_state(file_path:str) -> dict:with open(file_path, "r") as f:state = json.load(f)return state# 定义一个新的Agent，此时这个Agent还没有记忆，是一个空的状态
new_assistant_agent = AssistantAgent(name="assistant_agent",system_message="You are a helpful assistant",model_client=OpenAIChatCompletionClient(model="gpt-4o-2024-08-06",),
)async def main():# 加载本地Agent状态agent_state = load_state("./agent_state.json")await new_assistant_agent.load_state(agent_state)# 让新的Agent与模型进行一次交互response = await new_assistant_agent.on_messages([TextMessage(content="What was the last line of the previous poem you wrote", source="user")], CancellationToken())print(response.chat_message.content)print('-' * 50)agent_state = await new_assistant_agent.save_state()print(agent_state)asyncio.run(main())

运行结果如下：

$ python demo.py

此时如果你再将这个新的Agent状态保存一下就可以看到下面这些内容：

{"type": "AssistantAgentState","version": "1.0.0","llm_context": {"messages": [// 这部分是之前Agent的聊天记录{"content": "Write a 3 line poem on lake tangayika","source": "user","type": "UserMessage"},{"content": "In waters deep, where ancient secrets keep,  \nTanganyika mirrors the sky's boundless sweep,  \nSilent guardian of tales and dreams, it sleeps.  ","source": "assistant_agent","type": "AssistantMessage"},// 从这开始是新Agent的聊天记录{"content": "What was the last line of the previous poem you wrote","source": "user","type": "UserMessage"},{"content": "The last line of the previous poem I wrote was:  \nSilent guardian of tales and dreams, it sleeps.","source": "assistant_agent","type": "AssistantMessage"}]}
}

---

Saving and Loading Teams

和上面Agent状态获取与加载一样，Team的状态获取与加载使用的是同名函数save_state()和 load_state()。

• 先写一个demo用来保存team的状态：

from autogen_agentchat.agents import AssistantAgent
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_agentchat.conditions import MaxMessageTermination
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.ui import Console
import os, asyncio, jsonos.environ["OPENAI_API_KEY"] = "你的OpenAI API Key"
def persisting_state(file_name:str, state):with open(file_name, "w") as f:json.dump(state, f)assistant_agent = AssistantAgent(name="assistant_agent",system_message="You are a helpful assistant",model_client=OpenAIChatCompletionClient(model="gpt-4o-2024-08-06",),
)
agent_team = RoundRobinGroupChat([assistant_agent], termination_condition=MaxMessageTermination(max_messages=2))#----------------------------------------------------#
async def main():stream = agent_team.run_stream(task="Write a beautiful poem 3-line about lake tangayika")await Console(stream)# 获取team的状态team_state = await agent_team.save_state()print(team_state)persisting_state("./team_state.json", team_state)asyncio.run(main())

运行结果如下：

$ python demo.py

这里可以发现Team类型的状态要比单个Agent的内容多不少：

{"type": "TeamState","version": "1.0.0","agent_states": {"group_chat_manager/10c60d46-7e4d-4de0-ae64-e2b7eaf97119": {"type": "RoundRobinManagerState","version": "1.0.0","message_thread": [{"source": "user","models_usage": null,"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 embrace, skies mirror deep,  \nWhispering waves cradle secrets to keep,  \nNature's timeless dance, where the heart finds sleep.  ","type": "TextMessage"}],"current_turn": 0,"next_speaker_index": 0},"collect_output_messages/10c60d46-7e4d-4de0-ae64-e2b7eaf97119": {},"assistant_agent/10c60d46-7e4d-4de0-ae64-e2b7eaf97119": {"type": "ChatAgentContainerState","version": "1.0.0","agent_state": {"type": "AssistantAgentState","version": "1.0.0","llm_context": {"messages": [{"content": "Write a beautiful poem 3-line about lake tangayika","source": "user","type": "UserMessage"},{"content": "In Tanganyika's embrace, skies mirror deep,  \nWhispering waves cradle secrets to keep,  \nNature's timeless dance, where the heart finds sleep.  ","source": "assistant_agent","type": "AssistantMessage"}]}},"message_buffer": []}},"team_id": "10c60d46-7e4d-4de0-ae64-e2b7eaf97119"
}

• 再写一个demo用来加载本地保存的状态：

from autogen_agentchat.agents import AssistantAgent
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_agentchat.conditions import MaxMessageTermination
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.ui import Console
import os, asyncio, jsonos.environ["OPENAI_API_KEY"] = "你的OpenAI API Key"
def load_state(file_path:str) -> dict:with open(file_path, "r") as f:state = json.load(f)return stateassistant_agent = AssistantAgent(name="assistant_agent",system_message="You are a helpful assistant",model_client=OpenAIChatCompletionClient(model="gpt-4o-2024-08-06",),
)
agent_team = RoundRobinGroupChat([assistant_agent], termination_condition=MaxMessageTermination(max_messages=2))#----------------------------------------------------#
async def main():# 从本地加载team的状态team_state = load_state("./team_state.json")await agent_team.load_state(team_state)stream = agent_team.run_stream(task="What was the last line of the poem you wrote?")await Console(stream)print('-' * 50)team_state = await agent_team.save_state()print(team_state)asyncio.run(main())

运行结果如下：

$ python demo.py

---

Persisting State (File or Database)

剩下这一小段就是介绍如何将Agent和Team状态写入硬盘并加载，其实上面的示例中已经用到了这个函数，官方教程中没有对其进行封装，我这里直接把上面用到的函数贴在这：

# 将状态写入本地
def persisting_state(file_name:str, state):with open(file_name, "w") as f:json.dump(state, f)# 从本地加载状态
def load_state(file_path:str) -> dict:with open(file_path, "r") as f:state = json.load(f)return state

---

Take a breath

至此，官方教程中的 Tutorial 章节就全部介绍完了，后面将开启官方教程中的进阶章节 Advanced 部分，革命尚未成功，同志仍需努力。