How to implement state save and load for a custom agent built using core inheriting Routed Agent

[cvarun94]

Hi!

I am implementing a group chat consisting a team of agents and a manager agent using autogen_core(0.4+). I want to implement the save and load state methods that is available in the agentchat, however not sure how to do.. I understand that the state is basically the model_context but looking at the Agentchat.AssistantAgent source code I see that they have implemented more complex logic than simply dumping the model_context, also not sure what the model_contex.dump_model() also does?
I have used the GroupChat code in the core documentation as a reference point, using RoutedAgent as base.
Would be pleased to hear all your thoughts on how to implement a save/load state for a custom agent.

Thank You!

[ekzhu]

There are two different serializations:

• Component config: this is what dump_component and load_component is for. It is for serializing the configs that can be used to create the object.
• State save and load: this is what you are looking for. Once you can create an instance of the object e.g., an agent, you want to load a persisted state into it.

In Core API, this means your factory method is responsible for creating the instance, and the save_state and load_state methods are responsible for managing your instance's state.

For model context, built-in ones all have save_state and load_state methods: https://microsoft.github.io/autogen/stable/reference/python/autogen_core.model_context.html#autogen_core.model_context.ChatCompletionContext.save_state. I think you can just use that.

[cvarun94]

Hi Eric!
Thank you for your response.

I went through the source code, context docs and also the doc you had referenced. For the benefit of others I am writing my understanding of states, context and how to implement them, pls correct me if wrong.

Concept

1. Based on my understanding the state of an agent is the context of the agent (LLM context to be precise), and must be any of the ChatCompletionContext (Unbounded, Buffered etc.)
2. In core, when writing an Agent class inheriting the RoutedAgent, we must have the save_state method where to save the context we must return a AssistantAgentState dump.
3. Similarly the base agent class must have a load_state, where we must first validate the state and then load it into the agent's context.
4. I am working on a multiagent system where the runtime has registered 5 agents created using the base agent class.
5. So at any point to save the state of the entire team, I call save_state() on the runtime object which gives me the state the team of agents. I am able to save and load states into the agents successfully.

Observation

1. Even if we are using bufferedChatCompletion context with size 5 for the agent, calling save_state will save all the messages the agent has seen from the start.
2. Using AssistantAgentState from agentchat for allows to save and load the context into the agent context.
3. Calling save_state on the runtime obect is giving a dictionary of states of the different agents registered on that runtime.

Below is the baseAgent class that I have written-

class BaseAgent(RoutedAgent):
    """A groupchat participant agent that uses LLM to respond to queries."""
    def __init__(
        self, 
        description:str, 
        group_broadcast_topic:str, 
        model_client:ChatCompletionClient, 
        sys_ins:str, 
        tool_schema:List[Tool]=[], 
        model_context:ChatCompletionContext|None = None
    ) -> None:
        
        super().__init__(description=description)
        self.description = description
        self.model_client=model_client
        self.sys_ins = sys_ins
        self.tools = tool_schema
        self.group_broadcast_topic = group_broadcast_topic
        self.chat_hist:List[LLMMessage]=[]
        self.model_context = model_context


    async def tool_Exec8r(
        self, 
        call:FunctionCall, 
        cancellation_token:CancellationToken
    ) -> FunctionExecutionResult:
        
        Console().print(f"--Tool call initiated by:{self.id.type}")
        mytool = next((tool for tool in self.tools if tool.name==call.name), None)
        if mytool:
            toolArgs = json.loads(call.arguments)
            toolResult = await mytool.run_json(args=toolArgs, 
                                               cancellation_token=cancellation_token)
            return FunctionExecutionResult(call_id=call.id, 
                                          content=mytool.return_value_as_string(toolResult), 
                                          is_error=False, 
                                          name=mytool.name)
        
    async def save_state(self) -> Mapping[str, Any]:
        """Save the current state of the assistant agent."""
        model_context_state = await self.model_context.save_state()
        return AssistantAgentState(llm_context=model_context_state).model_dump()


    async def load_state(self, state: Mapping[str, Any]) -> None:
        """Load the state of the assistant agent"""
        assistant_agent_state = AssistantAgentState.model_validate(state)
        # Load the model context state.
        await self.model_context.load_state(assistant_agent_state.llm_context)
    
    @message_handler
    async def handle_GroupMessages(self, message:GroupMessages, ctx:MessageContext) -> None:
        if message.source.lower()=="user":
            curr_msg = UserMessage(content=message.content.content, source=message.source)
        else:
            curr_msg = AssistantMessage(content=message.content.content, source=message.source)
        await self.model_context.add_message(message=curr_msg)
        
    @message_handler
    async def handle_Request2Speak(self, message:Request2Speak, ctx:MessageContext) -> None:
        prompt = [SystemMessage(content=self.sys_ins)]+await self.model_context.get_messages()
        model_response = await self.model_client.create(messages=prompt, tools=self.tools, cancellation_token=ctx.cancellation_token)
        Console().print(f"--{self.id.type} Output--\n\n{model_response.content}")
        if isinstance(model_response.content, str): 
            curr_msg = AssistantMessage(content=model_response.content, source=self.id.type)
            await self.model_context.add_message(message=curr_msg)
            await self.publish_message(
                message=GroupMessages(
                    content=curr_msg, source=curr_msg.source
                ), 
                topic_id=DefaultTopicId(
                    type=self.group_broadcast_topic)
            )
        elif isinstance(model_response.content, List) and all(isinstance(call, FunctionCall) for call in model_response.content):
            curr_msg = AssistantMessage(content=model_response.content, source=self.id.type)
            await self.model_context.add_message(message=curr_msg)
            tool_Exec8r_result = await asyncio.gather(*[self.tool_Exec8r(call=call, cancellation_token=ctx.cancellation_token) for call in model_response.content])
            if tool_Exec8r_result:
                curr_msg = FunctionExecutionResultMessage(content=tool_Exec8r_result)
                await self.model_context.add_message(message=curr_msg)
                prompt = [SystemMessage(content=self.sys_ins)]+await self.model_context.get_messages()
                model_response_reflect = await self.model_client.create(messages=prompt, cancellation_token=ctx.cancellation_token)
                curr_msg_reflect = model_response_reflect.content
                if isinstance(curr_msg_reflect, str): 
                    curr_msg = AssistantMessage(content=curr_msg_reflect, source=self.id.type)
                    await self.model_context.add_message(message=curr_msg)
                Console().print(f"--{self.id.type} Output Tool Reflected--\n\n{curr_msg.content}")
                await self.publish_message(
                    message=GroupMessages(
                    content=AssistantMessage(content=curr_msg_reflect, source=self.id.type),
                    source=self.id.type
                ), 
                    topic_id=DefaultTopicId(
                        type=self.group_broadcast_topic)
                )