Load and Save state in AgentChat

[victordibia]

Loading and Saving State in AgentChat

This PR adds mechanisms to load/save state in AgentChat.

• Load/Save Agents. Not all agents have state to be saved. Currently only AssistantAgent really keeps state - _model_context. Load and save serves to populate that variable. Current changes adds load_state and save_state ChatAgent (abs methods) and stub for BaseChatAgent. Introduces a BaseState and AssistantAgentState class

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

result = await assistant_agent.run(task="Write a 3 line poem on lake tangayika")
agent_state = await assistant_agent.save_state() 

load state in a new agent

await new_assistant_agent.load_state(agent_state)
result = await new_assistant_agent.run(task="What was the last line of the previous poem you wrote")

The last line of the poem is:  
"Infinite horizon, nature's grand encore."

• Load/Save Teams.
  Currently implemented on the BaseGroupChat class with load_state and save_state methods. Introduces a BaseTeamState
  • agent_names: List[str] = field(default_factory=list)
  • termination_state: Optional[BaseState] = field(default=None)
  • agent_states: Dict[str, BaseState] = field(default_factory=dict)
  • manager_state: BaseGroupChatManagerState = field( default_factory=BaseGroupChatManagerState)
  • state_type: str = field(default="BaseTeamState")

load_state will load each participant agent state in agent_states and termination_state. All teams that inherit from BaseGroupChat should benefit from this OOTB

Note

What this does not do:

• Does not save/load GroupChatManager state yet: GroupChats have GroupChatManager classes that have state but are managed by the runtime. Getting and setting their state e.g., current turn etc is more involved. This implementation does not do this yet, but allows for it in the manager_state member of BaseTeamState.
  This can have some effects - e.g, after you load a team state, it might not know the last speaker to continue from and will begin from the first e.g. in RoundRobinGroupChat
• Custom Agents have to implement their load/save methods

assistant_agent = AssistantAgent(
    name="assistant_agent",
    system_message="You are a helpful assistant",
    model_client=OpenAIChatCompletionClient(
        model="gpt-4o-2024-08-06",
    ),
)
# Define a team
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")
await Console(stream)
team_state = await agent_team.save_state()

from autogen_agentchat.state import BaseTeamState

# convert team state to dict (save to file if needed)
team_state_dict = vars(team_state)
print(team_state_dict)      

# convert back to team state object
team_state = BaseTeamState(**vars(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)

• Save / Load Termination

max_termination = MaxMessageTermination(max_messages=2)
text_termination = TextMentionTermination(text="stop")
termination = text_termination | max_termination
termination_state = await termination.save_state()
print(termination_state)
await termination.load_state(termination_state)

Why are these changes needed?

Related issue number

Closes #4100

Checks

• I've included any doc changes needed for https://microsoft.github.io/autogen/. See https://microsoft.github.io/autogen/docs/Contribute#documentation to build and test documentation locally.
• I've added tests (if relevant) corresponding to the changes introduced in this PR.
• I've made sure all auto checks have passed.

[ekzhu]

I believe right now for team you can serialize the SingleThreadedAgentRuntime by calling the runtime's save_state and load_state. This requires the BaseGroupChatManager and ChatAgentContainer to implement their save_state and load_state methods. For the latter just call the contained ChatAgent's corresponding methods.

This will obviously change once we allow user to customize the runtime the team uses, and a runtime can be shared across multiple teams. However, this will serve as the driving force for that design.

[victordibia]

I believe right now for team you can serialize the SingleThreadedAgentRuntime by calling the runtime's save_state and load_state
. For the latter just call the contained ChatAgent's corresponding methods.

Not sure I understand what you mean above? Can you add some sample code to show what is meant by this?
Also, how does that add up to a developer experience of being able to do team.load_state() and team.save_state()?

[ekzhu]

Each group chat team carries its own runtime currently:

autogen/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py

Lines 67 to 68 in f02aac7

	self._runtime = SingleThreadedAgentRuntime()

The SingleThreadedAgentRuntime supports save_state and load_state for all registered agents:

autogen/python/packages/autogen-core/src/autogen_core/application/_single_threaded_agent_runtime.py

Lines 298 to 299 in f02aac7

	async def save_state(self) -> Mapping[str, Any]:
	state: Dict[str, Dict[str, Any]] = {}

So, from BaseGroupChat, you can implement save_state by calling the self._runtime.save_state() to get all the agent's states, including the group chat manager's.

This requires implementing the save_state and load_state of Core agents:

autogen/python/packages/autogen-core/src/autogen_core/base/_agent.py

Lines 36 to 37 in f02aac7

	async def save_state(self) -> Mapping[str, Any]:
	"""Save the state of the agent. The result must be JSON serializable."""

So this means the ChatAgentContainer and BaseGroupChatManager should implement these methods.

Besides all the agent's state, the team itself also carries some state including the team's id, which is generated in the constructor. This team id should be saved and loaded because it forms the reference to the Core agents in the runtime -- all Core agents' AgentID's key field is the team id.

class BaseGroupChat(Team, ABC):
  ...
  async def save_state(self) -> Mapping[str, Any]:
    if not self._initialized:
        return {"team_id": self._team_id} # If not initialized the only state we have is team id
    if self._is_running:
      raise RuntimeError("The group chat is currently running. It must be stopped before calling save_state()")
    self._is_running = True # Prevent running the team while saving states.
    runtime_state = await self._runtime.save_state()
    self._is_running = False
    return {"runtime_state": runtime_state, "team_id": self._team_id}

There might be other stuff like self._stop_reason etc., but should be trivial to handle.

[victordibia]

Ah, got it. This PR certainly goes in a diff/wrong direction. Do you want to take a stab at the approach you describe above which seems a lot more straightforward (it is still way too under specified for me to explore atm, similar to the original issue).
I’ll close this for now in favor of that PR which I’ll be happy to review.

[victordibia]

Thanks Eric, I tested, looks good.

[lsy641]

Hello. How to use the team save_state feature? Currently, I am not able to use it. I got "AttributeError: 'RoundRobinGroupChat' object has no attribute 'save_state'" in 0.4.0 dev

[ekzhu]

@lsy641 please upgrade to the latest dev release (dev11 as we speak)