Managing State#

So far, we have discussed how to build components in a multi-agent application - agents, teams, termination conditions. In many cases, it is useful to save the state of these components to disk and load them back later. This is particularly useful in a web application where stateless endpoints respond to requests and need to load the state of the application from persistent storage.

In this notebook, we will discuss how to save and load the state of agents, teams, and termination conditions.

Saving and Loading Agents#

We can get the state of an agent by calling save_state() method on an 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
from autogen_core import CancellationToken
from autogen_ext.models.openai import OpenAIChatCompletionClient

assistant_agent = AssistantAgent(
    name="assistant_agent",
    system_message="You are a helpful assistant",
    model_client=OpenAIChatCompletionClient(
        model="gpt-4o-2024-08-06",
        # api_key="YOUR_API_KEY",
    ),
)

# Use asyncio.run(...) when running in a script.
response = await assistant_agent.on_messages(
    [TextMessage(content="Write a 3 line poem on lake tangayika", source="user")], CancellationToken()
)
print(response.chat_message.content)

In Tanganyika's depths so wide and deep,  
Ancient secrets in still waters sleep,  
Ripples tell tales that time longs to keep.  

agent_state = await assistant_agent.save_state()
print(agent_state)

{'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 depths so wide and deep,  \nAncient secrets in still waters sleep,  \nRipples tell tales that time longs to keep.  ", 'source': 'assistant_agent', 'type': 'AssistantMessage'}]}

new_assistant_agent = AssistantAgent(
    name="assistant_agent",
    system_message="You are a helpful assistant",
    model_client=OpenAIChatCompletionClient(
        model="gpt-4o-2024-08-06",
    ),
)
await new_assistant_agent.load_state(agent_state)

# Use asyncio.run(...) when running in a script.
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)

The last line of the poem I wrote was:  
"Ripples tell tales that time longs to keep."

Note

For AssistantAgent, its state consists of the model_context. If your write your own custom agent, consider overriding the save_state() and load_state() methods to customize the behavior. The default implementations save and load an empty state.

Saving and Loading Teams#

We can get the state of a team by calling save_state method on the team and load it back by calling load_state method on the team.

When we call save_state on a team, it saves the state of all the agents in the team.

We will begin by creating a simple RoundRobinGroupChat team with a single agent and ask it to write a poem.

# Define a team.
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))

# Run the team and stream messages to the console.
stream = agent_team.run_stream(task="Write a beautiful poem 3-line about lake tangayika")

# Use asyncio.run(...) when running in a script.
await Console(stream)

# Save the state of the agent team.
team_state = await agent_team.save_state()

---------- user ----------
Write a beautiful poem 3-line about lake tangayika
---------- assistant_agent ----------
In Tanganyika's depths, where light gently weaves,  
Silver reflections dance on ancient water's face,  
Whispered stories of time in the rippling leaves.  
[Prompt tokens: 29, Completion tokens: 36]
---------- Summary ----------
Number of messages: 2
Finish reason: Maximum number of messages 2 reached, current message count: 2
Total prompt tokens: 29
Total completion tokens: 36
Duration: 1.16 seconds

If we reset the team (simulating instantiation of the team), and ask the question What was the last line of the poem you wrote?, we see that the team is unable to accomplish this as there is no reference to the previous run.

await agent_team.reset()
stream = agent_team.run_stream(task="What was the last line of the poem you wrote?")
await Console(stream)

---------- user ----------
What was the last line of the poem you wrote?
---------- assistant_agent ----------
I don't write poems on my own, but I can help create one with you or try to recall a specific poem if you have one in mind. Let me know what you'd like to do!
[Prompt tokens: 28, Completion tokens: 39]
---------- Summary ----------
Number of messages: 2
Finish reason: Maximum number of messages 2 reached, current message count: 2
Total prompt tokens: 28
Total completion tokens: 39
Duration: 0.95 seconds

TaskResult(messages=[TextMessage(source='user', models_usage=None, type='TextMessage', content='What was the last line of the poem you wrote?'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=28, completion_tokens=39), type='TextMessage', content="I don't write poems on my own, but I can help create one with you or try to recall a specific poem if you have one in mind. Let me know what you'd like to do!")], stop_reason='Maximum number of messages 2 reached, current message count: 2')

Next, we load the state of the team and ask the same question. We see that the team is able to accurately return the last line of the poem it wrote.

Note: You can serialize the state of the team to a file and load it back later.

print(team_state)

# Load team state.
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)

{'type': 'TeamState', 'version': '1.0.0', 'agent_states': {'group_chat_manager/c80054be-efb2-4bc7-ba0d-900962092c44': {'type': 'RoundRobinManagerState', 'version': '1.0.0', 'message_thread': [{'source': 'user', 'models_usage': None, 'type': 'TextMessage', 'content': 'Write a beautiful poem 3-line about lake tangayika'}, {'source': 'assistant_agent', 'models_usage': {'prompt_tokens': 29, 'completion_tokens': 36}, 'type': 'TextMessage', 'content': "In Tanganyika's depths, where light gently weaves,  \nSilver reflections dance on ancient water's face,  \nWhispered stories of time in the rippling leaves.  "}], 'current_turn': 0, 'next_speaker_index': 0}, 'collect_output_messages/c80054be-efb2-4bc7-ba0d-900962092c44': {}, 'assistant_agent/c80054be-efb2-4bc7-ba0d-900962092c44': {'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 depths, where light gently weaves,  \nSilver reflections dance on ancient water's face,  \nWhispered stories of time in the rippling leaves.  ", 'source': 'assistant_agent', 'type': 'AssistantMessage'}]}, 'message_buffer': []}}, 'team_id': 'c80054be-efb2-4bc7-ba0d-900962092c44'}
---------- user ----------
What was the last line of the poem you wrote?
---------- assistant_agent ----------
The last line of the poem I wrote was:  
"Whispered stories of time in the rippling leaves."
[Prompt tokens: 88, Completion tokens: 24]
---------- Summary ----------
Number of messages: 2
Finish reason: Maximum number of messages 2 reached, current message count: 2
Total prompt tokens: 88
Total completion tokens: 24
Duration: 0.79 seconds

TaskResult(messages=[TextMessage(source='user', models_usage=None, type='TextMessage', content='What was the last line of the poem you wrote?'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=88, completion_tokens=24), type='TextMessage', content='The last line of the poem I wrote was:  \n"Whispered stories of time in the rippling leaves."')], stop_reason='Maximum number of messages 2 reached, current message count: 2')
On this page
Edit on GitHub
Show Source