import os
from dotenv import load_dotenv

load_dotenv()

True

管理状态

到目前为止，我们已经讨论了如何在多智能体应用程序中构建组件——智能体、团队、终止条件。在许多情况下，将这些组件的状态保存到磁盘并稍后加载是很有用的。这在 Web 应用程序中尤其有用，因为无状态端点会响应请求，并且需要从持久存储中加载应用程序的状态。

在本笔记本中，我们将讨论如何保存和加载代理、团队和终止条件的状态。

保存和加载代理

save_state()我们可以通过调用方法来获取代理的状态AssistantAgent。

from autogen_agentchat.agents import AssistantAgent  # 引入助手代理类
from autogen_agentchat.conditions import MaxMessageTermination  # 引入最大消息终止条件
from autogen_agentchat.messages import TextMessage  # 引入文本消息类
from autogen_agentchat.teams import RoundRobinGroupChat  # 引入轮询对话团队类
from autogen_agentchat.ui import Console  # 引入控制台 UI 工具
from autogen_core import CancellationToken  # 引入取消令牌，用于控制中止任务
from autogen_ext.models.openai import OpenAIChatCompletionClient  # 引入 OpenAI 聊天模型客户端

# 创建 OpenAI 模型客户端，指定使用 gpt-4o-mini 模型
model_client = OpenAIChatCompletionClient(model="gpt-4o-mini")

# 创建一个助手代理实例，提供系统角色信息和模型客户端
assistant_agent = AssistantAgent(
    name="assistant_agent",  # 代理名称
    system_message="你是一个乐于助人的助手",  # 系统提示信息，用于设定模型行为
    model_client=model_client,  # 使用的模型客户端
)

# 以下代码仅在 Jupyter 或异步环境中直接运行（脚本中应使用 asyncio.run(...) 包裹）

# 构造一条来自用户的文本消息，请求写一首关于青海湖的三行诗
response = await assistant_agent.on_messages(
    [TextMessage(content="请写一首关于青海湖的三行诗", source="user")],
    CancellationToken()  # 创建一个默认的取消令牌
)

# 输出模型生成的回应文本
print(response.chat_message)

# 关闭模型客户端，释放资源
await model_client.close()

source='assistant_agent' models_usage=RequestUsage(prompt_tokens=32, completion_tokens=27) metadata={} created_at=datetime.datetime(2025, 7, 1, 9, 22, 45, 48372, tzinfo=datetime.timezone.utc) content='青海湖畔波光粼，  \n雪山映照天边云，  \n心随白鸟自在行。' type='TextMessage'

# 保存当前助手代理的内部状态（如对话历史、缓存等），返回的是可序列化的对象（如 dict）
agent_state = await assistant_agent.save_state()

# 打印保存下来的代理状态（通常用于调试、持久化或后续恢复）
print(agent_state)

{'type': 'AssistantAgentState', 'version': '1.0.0', 'llm_context': {'messages': [{'content': '请写一首关于青海湖的三行诗', 'source': 'user', 'type': 'UserMessage'}, {'content': '青海湖畔波光粼，  \n雪山映照天边云，  \n心随白鸟自在行。', 'thought': None, 'source': 'assistant_agent', 'type': 'AssistantMessage'}]}}

# 创建一个 OpenAI 聊天模型客户端，使用指定版本的 gpt-4o-mini 模型
model_client = OpenAIChatCompletionClient(model="gpt-4o-mini")

# 创建一个新的助手代理实例，用于恢复之前保存的状态
new_assistant_agent = AssistantAgent(
    name="assistant_agent",  # 助手代理名称
    system_message="你是一个乐于助人的助手",  # 系统提示信息
    model_client=model_client,  # 传入模型客户端
)

# 加载之前保存的代理状态，恢复上下文（如对话历史）
await new_assistant_agent.load_state(agent_state)

# 以下代码适用于 Jupyter 或异步环境（在 .py 脚本中建议使用 asyncio.run(...) 运行）

# 构造用户消息，询问“你之前写的诗最后一句是什么？”
response = await new_assistant_agent.on_messages(
    [TextMessage(content="你之前写的诗的最后一句是什么？", source="user")], 
    CancellationToken()  # 创建取消令牌
)

# 打印助手生成的回复
print(response.chat_message)

# 关闭模型客户端，释放资源
await model_client.close()

source='assistant_agent' models_usage=RequestUsage(prompt_tokens=78, completion_tokens=11) metadata={} created_at=datetime.datetime(2025, 7, 1, 9, 26, 44, 377777, tzinfo=datetime.timezone.utc) content='最后一句是：“心随白鸟自在行。”' type='TextMessage'

对于AssistantAgent，其状态由 model_context 组成。如果您编写自己的自定义代理，请考虑重写save_state()和load_state()方法来自定义行为。默认实现会保存和加载空状态。

保存和加载团队

我们可以通过调用团队的方法来获取团队的状态，并通过调用团队的方法save_state来加载它。load_state

当我们召集save_state一个团队时，它会保存团队中所有代理的状态。

我们将首先创建一个RoundRobinGroupChat只有一个代理的简单团队，并要求它写一首诗。

# 创建 OpenAI 聊天模型客户端，使用 gpt-4o-mini
model_client = OpenAIChatCompletionClient(model="gpt-4o-mini")

# 定义一个助手代理（AssistantAgent）
assistant_agent = AssistantAgent(
    name="assistant_agent",  # 代理的名称
    system_message="你是一个乐于助人的助手",  # 系统提示信息
    model_client=model_client,  # 使用上面创建的模型客户端
)

# 定义一个轮流对话团队（RoundRobinGroupChat），只包含一个代理
# 设定终止条件为最多对话 2 轮（即最多 2 条消息）
agent_team = RoundRobinGroupChat(
    [assistant_agent],
    termination_condition=MaxMessageTermination(max_messages=2)
)

# 启动团队任务，生成一个流式响应（stream），任务内容是写一首关于青海湖的三行优美诗歌
stream = agent_team.run_stream(task="写一首关于青海湖的优美三行诗")  

# ⚠️ 注意：如果你是在 .py 脚本中运行，请用 asyncio.run(...) 包裹此处代码
# 在 Jupyter Notebook 或异步环境中可直接 await
await Console(stream)  # 将代理的对话过程实时打印到控制台

# 保存代理团队当前的状态（如对话历史），可用于后续恢复
team_state = await agent_team.save_state()

---------- TextMessage (user) ----------
写一首关于青海湖的优美三行诗
---------- TextMessage (assistant_agent) ----------
青海湖畔水天连，  
碧波荡漾映云烟，  
白鹤翩翩舞梦间。  

---

如果我们重置团队（模拟团队的实例化），并提出问题，我们会发现团队无法完成此操作，因为没有对上次运行的参考。

# 重置代理团队的状态（例如清空对话历史）
await agent_team.reset()

# 重新启动团队对话任务，提问“你刚才写的诗的最后一句是什么？”
# 注意：虽然我们清空了历史，因此代理不一定记得之前的内容
stream = agent_team.run_stream(task="你刚才写的诗的最后一句是什么？") 

# 实时将对话输出到控制台（适用于 Jupyter 或异步环境）
await Console(stream)

---------- TextMessage (user) ----------
你刚才写的诗的最后一句是什么？
---------- TextMessage (assistant_agent) ----------
抱歉，我无法记住之前的对话内容。不过，我很乐意帮你创作一首新诗或者提供一些灵感！你想要什么主题的诗呢？





TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, created_at=datetime.datetime(2025, 7, 1, 9, 36, 11, 688224, tzinfo=datetime.timezone.utc), content='你刚才写的诗的最后一句是什么？', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=31, completion_tokens=39), metadata={}, created_at=datetime.datetime(2025, 7, 1, 9, 36, 29, 209189, tzinfo=datetime.timezone.utc), content='抱歉，我无法记住之前的对话内容。不过，我很乐意帮你创作一首新诗或者提供一些灵感！你想要什么主题的诗呢？', type='TextMessage')], stop_reason='Maximum number of messages 2 reached, current message count: 2')

---

接下来，我们加载团队的状态并提出同样的问题。我们看到团队能够准确地返回他们所写诗歌的最后一行。

# 打印当前代理团队的状态对象（包含对话历史、轮数等信息）
print(team_state)

# 加载之前保存的团队状态，以恢复对话历史等上下文信息
await agent_team.load_state(team_state)

# 加载状态后重新运行团队任务，提问：“你刚才写的诗的最后一句是什么？”
# 因为我们恢复了历史，对话上下文仍在，代理有可能准确记住上次的诗句
stream = agent_team.run_stream(task="你刚才写的诗的最后一句是什么？")

# 将回答实时输出到控制台
await Console(stream)

{'type': 'TeamState', 'version': '1.0.0', 'agent_states': {'assistant_agent': {'type': 'ChatAgentContainerState', 'version': '1.0.0', 'agent_state': {'type': 'AssistantAgentState', 'version': '1.0.0', 'llm_context': {'messages': [{'content': '写一首关于青海湖的优美三行诗', 'source': 'user', 'type': 'UserMessage'}, {'content': '青海湖畔水天连，  \n碧波荡漾映云烟，  \n白鹤翩翩舞梦间。  ', 'thought': None, 'source': 'assistant_agent', 'type': 'AssistantMessage'}]}}, 'message_buffer': []}, 'RoundRobinGroupChatManager': {'type': 'RoundRobinManagerState', 'version': '1.0.0', 'message_thread': [{'source': 'user', 'models_usage': None, 'metadata': {}, 'created_at': datetime.datetime(2025, 7, 1, 9, 35, 44, 411731, tzinfo=datetime.timezone.utc), 'content': '写一首关于青海湖的优美三行诗', 'type': 'TextMessage'}, {'source': 'assistant_agent', 'models_usage': {'prompt_tokens': 33, 'completion_tokens': 33}, 'metadata': {}, 'created_at': datetime.datetime(2025, 7, 1, 9, 36, 1, 867008, tzinfo=datetime.timezone.utc), 'content': '青海湖畔水天连，  \n碧波荡漾映云烟，  \n白鹤翩翩舞梦间。  ', 'type': 'TextMessage'}], 'current_turn': 0, 'next_speaker_index': 0}}}
---------- TextMessage (user) ----------
你刚才写的诗的最后一句是什么？
---------- TextMessage (assistant_agent) ----------
最后一句是：白鹤翩翩舞梦间。





TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, created_at=datetime.datetime(2025, 7, 1, 9, 41, 58, 607395, tzinfo=datetime.timezone.utc), content='你刚才写的诗的最后一句是什么？', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=15), metadata={}, created_at=datetime.datetime(2025, 7, 1, 9, 42, 16, 696642, tzinfo=datetime.timezone.utc), content='最后一句是：白鹤翩翩舞梦间。', type='TextMessage')], stop_reason='Maximum number of messages 2 reached, current message count: 2')

持久状态（文件或数据库）

很多情况下，我们可能希望将球队状态持久化到磁盘（或数据库），稍后再加载。状态是一个字典，可以序列化到文件或写入数据库。

# 打印当前代理团队的状态对象（包含对话历史、轮数等信息）
print(team_state)

{'type': 'TeamState', 'version': '1.0.0', 'agent_states': {'assistant_agent': {'type': 'ChatAgentContainerState', 'version': '1.0.0', 'agent_state': {'type': 'AssistantAgentState', 'version': '1.0.0', 'llm_context': {'messages': [{'content': '写一首关于青海湖的优美三行诗', 'source': 'user', 'type': 'UserMessage'}, {'content': '青海湖畔水天连，  \n碧波荡漾映云烟，  \n白鹤翩翩舞梦间。  ', 'thought': None, 'source': 'assistant_agent', 'type': 'AssistantMessage'}]}}, 'message_buffer': []}, 'RoundRobinGroupChatManager': {'type': 'RoundRobinManagerState', 'version': '1.0.0', 'message_thread': [{'source': 'user', 'models_usage': None, 'metadata': {}, 'created_at': datetime.datetime(2025, 7, 1, 9, 35, 44, 411731, tzinfo=datetime.timezone.utc), 'content': '写一首关于青海湖的优美三行诗', 'type': 'TextMessage'}, {'source': 'assistant_agent', 'models_usage': {'prompt_tokens': 33, 'completion_tokens': 33}, 'metadata': {}, 'created_at': datetime.datetime(2025, 7, 1, 9, 36, 1, 867008, tzinfo=datetime.timezone.utc), 'content': '青海湖畔水天连，  \n碧波荡漾映云烟，  \n白鹤翩翩舞梦间。  ', 'type': 'TextMessage'}], 'current_turn': 0, 'next_speaker_index': 0}}}

import json

## 将代理团队的状态保存到磁盘文件中（序列化为 JSON 格式）
with open("AutoGen/coding/team_state.json", "w") as f:
    # json.dump(team_state, f)  # 保存状态到指定路径的 JSON 文件
    json.dump(team_state, f, indent=2, ensure_ascii=False, default=str) # 转换 datetime 类型

## 从磁盘读取代理团队的状态（反序列化）
with open("AutoGen/coding/team_state.json", "r") as f:
    team_state = json.load(f)  # 加载之前保存的 JSON 文件内容为 Python 对象

# 创建一个新的代理团队，并设置相同的终止条件（最大消息轮数为 2）
new_agent_team = RoundRobinGroupChat([assistant_agent], termination_condition=MaxMessageTermination(max_messages=2))

# 加载之前保存的状态，以恢复之前的对话上下文
await new_agent_team.load_state(team_state)

# 重新发起对话，提问：“你刚才写的诗的最后一句是什么？”
stream = new_agent_team.run_stream(task="你刚才写的诗的最后一句是什么？")

# 实时将代理生成的内容输出到控制台
await Console(stream)

# 关闭模型客户端，释放资源
await model_client.close()

---------- TextMessage (user) ----------
你刚才写的诗的最后一句是什么？
---------- TextMessage (assistant_agent) ----------
最后一句是：“白鹤翩翩舞梦间。”