Agent State

The State class is a fundamental component of the UFO agent framework. It represents the current state of the agent and determines the next action and agent to handle the request. Each agent has a specific set of states that define the agent's behavior and workflow.

AgentStatus

The set of states for an agent is defined in the AgentStatus class:

class AgentStatus(Enum):
    """
    The status class for the agent.
    """

    ERROR = "ERROR"
    FINISH = "FINISH"
    CONTINUE = "CONTINUE"
    FAIL = "FAIL"
    PENDING = "PENDING"
    CONFIRM = "CONFIRM"
    SCREENSHOT = "SCREENSHOT"

Each agent implements its own set of AgentStatus to define the states of the agent.

AgentStateManager

The class AgentStateManager manages the state mapping from a string to the corresponding state class. Each state class is registered with the AgentStateManager using the register decorator to associate the state class with a specific agent, e.g.,

@AgentStateManager.register
class SomeAgentState(AgentState):
    """
    The state class for the some agent.
    """

Below is the basic structure of the AgentStateManager class:

class AgentStateManager(ABC, metaclass=SingletonABCMeta):
    """
    A abstract class to manage the states of the agent.
    """

    _state_mapping: Dict[str, Type[AgentState]] = {}

    def __init__(self):
        """
        Initialize the state manager.
        """

        self._state_instance_mapping: Dict[str, AgentState] = {}

    def get_state(self, status: str) -> AgentState:
        """
        Get the state for the status.
        :param status: The status string.
        :return: The state object.
        """

        # Lazy load the state class
        if status not in self._state_instance_mapping:
            state_class = self._state_mapping.get(status)
            if state_class:
                self._state_instance_mapping[status] = state_class()
            else:
                self._state_instance_mapping[status] = self.none_state

        state = self._state_instance_mapping.get(status, self.none_state)

        return state

    def add_state(self, status: str, state: AgentState) -> None:
        """
        Add a new state to the state mapping.
        :param status: The status string.
        :param state: The state object.
        """
        self.state_map[status] = state

    @property
    def state_map(self) -> Dict[str, AgentState]:
        """
        The state mapping of status to state.
        :return: The state mapping.
        """
        return self._state_instance_mapping

    @classmethod
    def register(cls, state_class: Type[AgentState]) -> Type[AgentState]:
        """
        Decorator to register the state class to the state manager.
        :param state_class: The state class to be registered.
        :return: The state class.
        """
        cls._state_mapping[state_class.name()] = state_class
        return state_class

    @property
    @abstractmethod
    def none_state(self) -> AgentState:
        """
        The none state of the state manager.
        """
        pass

AgentState

Each state class inherits from the AgentState class and must implement the method of handle to process the action in the state. In addition, the next_state and next_agent methods are used to determine the next state and agent to handle the transition. Please find below the reference for the State class in UFO.

165
166
167
168
169
170
171
172

@classmethod
@abstractmethod
def agent_class(cls) -> Type[BasicAgent]:
    """
    The class of the agent.
    :return: The class of the agent.
    """
    pass

122
123
124
125
126
127
128
129

@abstractmethod
def handle(self, agent: BasicAgent, context: Optional["Context"] = None) -> None:
    """
    Handle the agent for the current step.
    :param agent: The agent to handle.
    :param context: The context for the agent and session.
    """
    pass

149
150
151
152
153
154
155

@abstractmethod
def is_round_end(self) -> bool:
    """
    Check if the round ends.
    :return: True if the round ends, False otherwise.
    """
    pass

157
158
159
160
161
162
163

@abstractmethod
def is_subtask_end(self) -> bool:
    """
    Check if the subtask ends.
    :return: True if the subtask ends, False otherwise.
    """
    pass

174
175
176
177
178
179
180
181

@classmethod
@abstractmethod
def name(cls) -> str:
    """
    The class name of the state.
    :return: The class name of the state.
    """
    return ""

131
132
133
134
135
136
137
138

@abstractmethod
def next_agent(self, agent: BasicAgent) -> BasicAgent:
    """
    Get the agent for the next step.
    :param agent: The agent for the current step.
    :return: The agent for the next step.
    """
    return agent

140
141
142
143
144
145
146
147

@abstractmethod
def next_state(self, agent: BasicAgent) -> AgentState:
    """
    Get the state for the next step.
    :param agent: The agent for the current step.
    :return: The state for the next step.
    """
    pass