Backup-bdg/OpenHands

1from __future__ import annotations2 3import base644import os5import pickle6from dataclasses import dataclass, field7from enum import Enum8from typing import Any9 10import openhands11from openhands.core.logger import openhands_logger as logger12from openhands.core.schema import AgentState13from openhands.events.action import (14    MessageAction,15)16from openhands.events.action.agent import AgentFinishAction17from openhands.events.event import Event, EventSource18from openhands.llm.metrics import Metrics19from openhands.memory.view import View20from openhands.storage.files import FileStore21from openhands.storage.locations import get_conversation_agent_state_filename22 23 24class TrafficControlState(str, Enum):25    # default state, no rate limiting26    NORMAL = 'normal'27 28    # task paused due to traffic control29    THROTTLING = 'throttling'30 31    # traffic control is temporarily paused32    PAUSED = 'paused'33 34 35RESUMABLE_STATES = [36    AgentState.RUNNING,37    AgentState.PAUSED,38    AgentState.AWAITING_USER_INPUT,39    AgentState.FINISHED,40]41 42 43@dataclass44class State:45    """46    Represents the running state of an agent in the OpenHands system, saving data of its operation and memory.47 48    - Multi-agent/delegate state:49      - store the task (conversation between the agent and the user)50      - the subtask (conversation between an agent and the user or another agent)51      - global and local iterations52      - delegate levels for multi-agent interactions53      - almost stuck state54 55    - Running state of an agent:56      - current agent state (e.g., LOADING, RUNNING, PAUSED)57      - traffic control state for rate limiting58      - confirmation mode59      - the last error encountered60 61    - Data for saving and restoring the agent:62      - save to and restore from a session63      - serialize with pickle and base6464 65    - Save / restore data about message history66      - start and end IDs for events in agent's history67      - summaries and delegate summaries68 69    - Metrics:70      - global metrics for the current task71      - local metrics for the current subtask72 73    - Extra data:74      - additional task-specific data75    """76 77    session_id: str = ''78    # global iteration for the current task79    iteration: int = 080    # local iteration for the current subtask81    local_iteration: int = 082    # max number of iterations for the current task83    max_iterations: int = 10084    confirmation_mode: bool = False85    history: list[Event] = field(default_factory=list)86    inputs: dict = field(default_factory=dict)87    outputs: dict = field(default_factory=dict)88    agent_state: AgentState = AgentState.LOADING89    resume_state: AgentState | None = None90    traffic_control_state: TrafficControlState = TrafficControlState.NORMAL91    # global metrics for the current task92    metrics: Metrics = field(default_factory=Metrics)93    # local metrics for the current subtask94    local_metrics: Metrics = field(default_factory=Metrics)95    # root agent has level 0, and every delegate increases the level by one96    delegate_level: int = 097    # start_id and end_id track the range of events in history98    start_id: int = -199    end_id: int = -1100 101    delegates: dict[tuple[int, int], tuple[str, str]] = field(default_factory=dict)102    # NOTE: This will never be used by the controller, but it can be used by different103    # evaluation tasks to store extra data needed to track the progress/state of the task.104    extra_data: dict[str, Any] = field(default_factory=dict)105    last_error: str = ''106 107    def save_to_session(108        self, sid: str, file_store: FileStore, user_id: str | None109    ) -> None:110        pickled = pickle.dumps(self)111        logger.debug(f'Saving state to session {sid}:{self.agent_state}')112        encoded = base64.b64encode(pickled).decode('utf-8')113        try:114            file_store.write(115                get_conversation_agent_state_filename(sid, user_id), encoded116            )117 118            # see if state is in the old directory on saas/remote use cases and delete it.119            if user_id:120                filename = get_conversation_agent_state_filename(sid)121                try:122                    file_store.delete(filename)123                except Exception:124                    pass125        except Exception as e:126            logger.error(f'Failed to save state to session: {e}')127            raise e128 129    @staticmethod130    def restore_from_session(131        sid: str, file_store: FileStore, user_id: str | None = None132    ) -> 'State':133        """134        Restores the state from the previously saved session.135        """136 137        state: State138        try:139            encoded = file_store.read(140                get_conversation_agent_state_filename(sid, user_id)141            )142            pickled = base64.b64decode(encoded)143            state = pickle.loads(pickled)144        except FileNotFoundError:145            # if user_id is provided, we are in a saas/remote use case146            # and we need to check if the state is in the old directory.147            if user_id:148                filename = get_conversation_agent_state_filename(sid)149                encoded = file_store.read(filename)150                pickled = base64.b64decode(encoded)151                state = pickle.loads(pickled)152            else:153                raise FileNotFoundError(154                    f'Could not restore state from session file for sid: {sid}'155                )156        except Exception as e:157            logger.debug(f'Could not restore state from session: {e}')158            raise e159 160        # update state161        if state.agent_state in RESUMABLE_STATES:162            state.resume_state = state.agent_state163        else:164            state.resume_state = None165 166        # first state after restore167        state.agent_state = AgentState.LOADING168        return state169 170    def __getstate__(self) -> dict:171        # don't pickle history, it will be restored from the event stream172        state = self.__dict__.copy()173        state['history'] = []174 175        # Remove any view caching attributes. They'll be rebuilt frmo the176        # history after that gets reloaded.177        state.pop('_history_checksum', None)178        state.pop('_view', None)179 180        return state181 182    def __setstate__(self, state: dict) -> None:183        self.__dict__.update(state)184 185        # make sure we always have the attribute history186        if not hasattr(self, 'history'):187            self.history = []188 189    def get_current_user_intent(self) -> tuple[str | None, list[str] | None]:190        """Returns the latest user message and image(if provided) that appears after a FinishAction, or the first (the task) if nothing was finished yet."""191        last_user_message = None192        last_user_message_image_urls: list[str] | None = []193        for event in reversed(self.view):194            if isinstance(event, MessageAction) and event.source == 'user':195                last_user_message = event.content196                last_user_message_image_urls = event.image_urls197            elif isinstance(event, AgentFinishAction):198                if last_user_message is not None:199                    return last_user_message, None200 201        return last_user_message, last_user_message_image_urls202 203    def get_last_agent_message(self) -> MessageAction | None:204        for event in reversed(self.view):205            if isinstance(event, MessageAction) and event.source == EventSource.AGENT:206                return event207        return None208 209    def get_last_user_message(self) -> MessageAction | None:210        for event in reversed(self.view):211            if isinstance(event, MessageAction) and event.source == EventSource.USER:212                return event213        return None214 215    def to_llm_metadata(self, agent_name: str) -> dict:216        return {217            'session_id': self.session_id,218            'trace_version': openhands.__version__,219            'tags': [220                f'agent:{agent_name}',221                f'web_host:{os.environ.get("WEB_HOST", "unspecified")}',222                f'openhands_version:{openhands.__version__}',223            ],224        }225 226    @property227    def view(self) -> View:228        # Compute a simple checksum from the history to see if we can re-use any229        # cached view.230        history_checksum = len(self.history)231        old_history_checksum = getattr(self, '_history_checksum', -1)232 233        # If the history has changed, we need to re-create the view and update234        # the caching.235        if history_checksum != old_history_checksum:236            self._history_checksum = history_checksum237            self._view = View.from_events(self.history)238 239        return self._view240