Implementation:Microsoft Autogen Agent State Models

Key	Value
id	Microsoft_Autogen_Agent_State_Models
source	Microsoft_Autogen
category	State Management

Overview

Description

The Agent State Models module defines a comprehensive hierarchy of Pydantic models for persisting and restoring the state of agents, teams, and group chat managers in the Autogen framework. These models enable checkpoint/restore functionality, allowing conversations and agent states to be saved and resumed later.

The module provides:

BaseState: Root class for all state models with type and version metadata
Agent States: Models for individual agent states (AssistantAgentState, SocietyOfMindAgentState)
Team States: Models for team-level state management (TeamState)
Group Chat Manager States: Models for various group chat orchestration strategies (RoundRobinManagerState, SelectorManagerState, SwarmManagerState, MagenticOneOrchestratorState)
Container States: Models for agent containers (ChatAgentContainerState)

All state models inherit from BaseState and include type discriminators for proper deserialization. The models use Pydantic for validation and serialization, ensuring type safety and easy JSON conversion.

Usage

These state models are used internally by agents and teams to implement the save_state() and load_state() methods defined in the ChatAgent protocol. They enable:

Saving conversation history and agent context to persistent storage
Restoring agent state after crashes or restarts
Implementing checkpoint/rollback mechanisms
Migrating agent state between different execution environments

The type field in each state model acts as a discriminator, allowing the framework to deserialize the correct state model type when loading from JSON or dictionaries.

Code Reference

Source Location

Repository: https://github.com/microsoft/autogen
File Path: /tmp/kapso_repo_2mr4n2g4/python/packages/autogen-agentchat/src/autogen_agentchat/state/_states.py
Lines: 1-80

Signature

class BaseState(BaseModel):
    type: str = Field(default="BaseState")
    version: str = Field(default="1.0.0")


class AssistantAgentState(BaseState):
    llm_context: Mapping[str, Any] = Field(default_factory=lambda: dict([("messages", [])]))
    type: str = Field(default="AssistantAgentState")


class TeamState(BaseState):
    agent_states: Mapping[str, Any] = Field(default_factory=dict)
    type: str = Field(default="TeamState")


class BaseGroupChatManagerState(BaseState):
    message_thread: List[Mapping[str, Any]] = Field(default_factory=list)
    current_turn: int = Field(default=0)
    type: str = Field(default="BaseGroupChatManagerState")


class ChatAgentContainerState(BaseState):
    agent_state: Mapping[str, Any] = Field(default_factory=dict)
    message_buffer: List[Mapping[str, Any]] = Field(default_factory=list)
    type: str = Field(default="ChatAgentContainerState")


class RoundRobinManagerState(BaseGroupChatManagerState):
    next_speaker_index: int = Field(default=0)
    type: str = Field(default="RoundRobinManagerState")


class SelectorManagerState(BaseGroupChatManagerState):
    previous_speaker: Optional[str] = Field(default=None)
    type: str = Field(default="SelectorManagerState")


class SwarmManagerState(BaseGroupChatManagerState):
    current_speaker: str = Field(default="")
    type: str = Field(default="SwarmManagerState")


class MagenticOneOrchestratorState(BaseGroupChatManagerState):
    task: str = Field(default="")
    facts: str = Field(default="")
    plan: str = Field(default="")
    n_rounds: int = Field(default=0)
    n_stalls: int = Field(default=0)
    type: str = Field(default="MagenticOneOrchestratorState")


class SocietyOfMindAgentState(BaseState):
    inner_team_state: Mapping[str, Any] = Field(default_factory=dict)
    type: str = Field(default="SocietyOfMindAgentState")

Import

from autogen_agentchat.state import (
    BaseState,
    AssistantAgentState,
    TeamState,
    BaseGroupChatManagerState,
    ChatAgentContainerState,
    RoundRobinManagerState,
    SelectorManagerState,
    SwarmManagerState,
    MagenticOneOrchestratorState,
    SocietyOfMindAgentState
)

I/O Contract

Base State Model

Field	Type	Default	Description
type	str	"BaseState"	Discriminator for state model type
version	str	"1.0.0"	Version of the state schema

Agent State Models

Model	Additional Fields	Description
AssistantAgentState	llm_context: Mapping[str, Any]	State for LLM-based assistant agents, includes message history
SocietyOfMindAgentState	inner_team_state: Mapping[str, Any]	State for Society of Mind agents with nested team state

Team State Models

Model	Fields	Description
TeamState	agent_states: Mapping[str, Any]	State for a team of agents, mapping agent names to their states

Group Chat Manager State Models

Model	Base Fields	Additional Fields	Description
BaseGroupChatManagerState	message_thread, current_turn	N/A	Base state for all group chat managers
RoundRobinManagerState	Inherits base fields	next_speaker_index: int	State for RoundRobinGroupChat manager
SelectorManagerState	Inherits base fields	previous_speaker: Optional[str]	State for SelectorGroupChat manager
SwarmManagerState	Inherits base fields	current_speaker: str	State for Swarm manager
MagenticOneOrchestratorState	Inherits base fields	task, facts, plan, n_rounds, n_stalls	State for MagneticOneGroupChat orchestrator

Container State Models

Model	Fields	Description
ChatAgentContainerState	agent_state: Mapping[str, Any], message_buffer: List[Mapping[str, Any]]	State for agent containers in group chats

Usage Examples

Saving and Loading Agent State

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.state import AssistantAgentState


async def save_load_agent_state():
    # Create agent
    agent = AssistantAgent("assistant", model_client=model_client)

    # Process some messages
    response = await agent.on_messages(messages, cancellation_token)

    # Save state
    state_dict = await agent.save_state()

    # Serialize to JSON
    import json
    state_json = json.dumps(state_dict)

    # Later, restore state
    restored_state = json.loads(state_json)
    await agent.load_state(restored_state)

Working with Team State

from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.state import TeamState


async def save_team_state(team: RoundRobinGroupChat):
    # Save entire team state (includes all agent states)
    team_state_dict = await team.save_state()

    # Validate structure
    team_state = TeamState.model_validate(team_state_dict)

    # Access individual agent states
    for agent_name, agent_state in team_state.agent_states.items():
        print(f"Agent {agent_name} state: {agent_state}")

    return team_state_dict

Custom State Serialization

from autogen_agentchat.state import BaseState, AssistantAgentState
from pydantic import Field


class CustomAgentState(BaseState):
    """Custom state model for a specialized agent."""

    custom_field: str = Field(default="")
    custom_counter: int = Field(default=0)
    type: str = Field(default="CustomAgentState")


async def use_custom_state():
    # Create custom state
    state = CustomAgentState(
        custom_field="example",
        custom_counter=42
    )

    # Serialize to dict
    state_dict = state.model_dump()

    # Deserialize
    restored = CustomAgentState.model_validate(state_dict)

    assert restored.custom_field == "example"
    assert restored.custom_counter == 42

State Versioning

from autogen_agentchat.state import BaseState


def check_state_version(state_dict: dict) -> bool:
    """Check if state version is compatible."""
    version = state_dict.get("version", "1.0.0")
    major_version = int(version.split(".")[0])

    # Only compatible with version 1.x.x
    return major_version == 1


async def safe_load_state(agent, state_dict):
    if check_state_version(state_dict):
        await agent.load_state(state_dict)
    else:
        raise ValueError(f"Incompatible state version: {state_dict.get('version')}")

MagenticOne Orchestrator State

from autogen_agentchat.state import MagenticOneOrchestratorState


async def track_orchestrator_progress():
    # Load orchestrator state
    state = MagenticOneOrchestratorState(
        task="Solve complex problem",
        facts="Known information...",
        plan="Step-by-step plan...",
        n_rounds=5,
        n_stalls=1,
        current_turn=10,
        message_thread=[]
    )

    # Check progress
    print(f"Task: {state.task}")
    print(f"Rounds completed: {state.n_rounds}")
    print(f"Stalls encountered: {state.n_stalls}")
    print(f"Current turn: {state.current_turn}")

    # Serialize
    state_dict = state.model_dump()

Discriminated Union Loading

from autogen_agentchat.state import (
    BaseState, AssistantAgentState, RoundRobinManagerState
)


def load_state_by_type(state_dict: dict) -> BaseState:
    """Load appropriate state model based on type discriminator."""
    state_type = state_dict.get("type")

    if state_type == "AssistantAgentState":
        return AssistantAgentState.model_validate(state_dict)
    elif state_type == "RoundRobinManagerState":
        return RoundRobinManagerState.model_validate(state_dict)
    else:
        return BaseState.model_validate(state_dict)

Related Pages

Microsoft_Autogen_ChatAgent_Protocol - Protocol that defines save_state() and load_state() methods
Microsoft_Autogen_AssistantAgent - Agent that uses AssistantAgentState
Microsoft_Autogen_RoundRobinGroupChat - Team that uses RoundRobinManagerState
Microsoft_Autogen_SelectorGroupChat - Team that uses SelectorManagerState
Microsoft_Autogen_SwarmGroupChat - Team that uses SwarmManagerState
Microsoft_Autogen_MagenticOneGroupChat - Team that uses MagenticOneOrchestratorState
Microsoft_Autogen_ChatAgentContainer - Container that uses ChatAgentContainerState
Microsoft_Autogen_SocietyOfMindAgent - Agent that uses SocietyOfMindAgentState
Microsoft_Autogen_BaseState - Root state model class

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment