Backup-bdg/OpenHands

1from __future__ import annotations2 3from abc import ABC, abstractmethod4from typing import TYPE_CHECKING5 6if TYPE_CHECKING:7    from openhands.controller.state.state import State8    from openhands.core.config import AgentConfig9    from openhands.events.action import Action10    from openhands.events.action.message import SystemMessageAction11    from openhands.utils.prompt import PromptManager12from litellm import ChatCompletionToolParam13 14from openhands.core.exceptions import (15    AgentAlreadyRegisteredError,16    AgentNotRegisteredError,17)18from openhands.core.logger import openhands_logger as logger19from openhands.events.event import EventSource20from openhands.llm.llm import LLM21from openhands.runtime.plugins import PluginRequirement22 23 24class Agent(ABC):25    DEPRECATED = False26    """27    This abstract base class is an general interface for an agent dedicated to28    executing a specific instruction and allowing human interaction with the29    agent during execution.30    It tracks the execution status and maintains a history of interactions.31    """32 33    _registry: dict[str, type['Agent']] = {}34    sandbox_plugins: list[PluginRequirement] = []35 36    def __init__(37        self,38        llm: LLM,39        config: 'AgentConfig',40    ):41        self.llm = llm42        self.config = config43        self._complete = False44        self._prompt_manager: 'PromptManager' | None = None45        self.mcp_tools: dict[str, ChatCompletionToolParam] = {}46        self.tools: list = []47 48    @property49    def prompt_manager(self) -> 'PromptManager':50        if self._prompt_manager is None:51            raise ValueError(f'Prompt manager not initialized for agent {self.name}')52        return self._prompt_manager53 54    def get_system_message(self) -> 'SystemMessageAction | None':55        """56        Returns a SystemMessageAction containing the system message and tools.57        This will be added to the event stream as the first message.58 59        Returns:60            SystemMessageAction: The system message action with content and tools61            None: If there was an error generating the system message62        """63        # Import here to avoid circular imports64        from openhands.events.action.message import SystemMessageAction65 66        try:67            if not self.prompt_manager:68                logger.warning(69                    f'[{self.name}] Prompt manager not initialized before getting system message'70                )71                return None72 73            system_message = self.prompt_manager.get_system_message()74 75            # Get tools if available76            tools = getattr(self, 'tools', None)77 78            system_message_action = SystemMessageAction(79                content=system_message, tools=tools, agent_class=self.name80            )81            # Set the source attribute82            system_message_action._source = EventSource.AGENT  # type: ignore83 84            return system_message_action85        except Exception as e:86            logger.warning(f'[{self.name}] Failed to generate system message: {e}')87            return None88 89    @property90    def complete(self) -> bool:91        """Indicates whether the current instruction execution is complete.92 93        Returns:94        - complete (bool): True if execution is complete; False otherwise.95        """96        return self._complete97 98    @abstractmethod99    def step(self, state: 'State') -> 'Action':100        """Starts the execution of the assigned instruction. This method should101        be implemented by subclasses to define the specific execution logic.102        """103        pass104 105    def reset(self) -> None:106        """Resets the agent's execution status and clears the history. This method can be used107        to prepare the agent for restarting the instruction or cleaning up before destruction.108 109        """110        # TODO clear history111        self._complete = False112 113        if self.llm:114            self.llm.reset()115 116    @property117    def name(self) -> str:118        return self.__class__.__name__119 120    @classmethod121    def register(cls, name: str, agent_cls: type['Agent']) -> None:122        """Registers an agent class in the registry.123 124        Parameters:125        - name (str): The name to register the class under.126        - agent_cls (Type['Agent']): The class to register.127 128        Raises:129        - AgentAlreadyRegisteredError: If name already registered130        """131        if name in cls._registry:132            raise AgentAlreadyRegisteredError(name)133        cls._registry[name] = agent_cls134 135    @classmethod136    def get_cls(cls, name: str) -> type['Agent']:137        """Retrieves an agent class from the registry.138 139        Parameters:140        - name (str): The name of the class to retrieve141 142        Returns:143        - agent_cls (Type['Agent']): The class registered under the specified name.144 145        Raises:146        - AgentNotRegisteredError: If name not registered147        """148        if name not in cls._registry:149            raise AgentNotRegisteredError(name)150        return cls._registry[name]151 152    @classmethod153    def list_agents(cls) -> list[str]:154        """Retrieves the list of all agent names from the registry.155 156        Raises:157        - AgentNotRegisteredError: If no agent is registered158        """159        if not bool(cls._registry):160            raise AgentNotRegisteredError()161        return list(cls._registry.keys())162 163    def set_mcp_tools(self, mcp_tools: list[dict]) -> None:164        """Sets the list of MCP tools for the agent.165 166        Args:167        - mcp_tools (list[dict]): The list of MCP tools.168        """169        logger.info(170            f'Setting {len(mcp_tools)} MCP tools for agent {self.name}: {[tool["function"]["name"] for tool in mcp_tools]}'171        )172        for tool in mcp_tools:173            _tool = ChatCompletionToolParam(**tool)174            if _tool['function']['name'] in self.mcp_tools:175                logger.warning(176                    f'Tool {_tool["function"]["name"]} already exists, skipping'177                )178                continue179            self.mcp_tools[_tool['function']['name']] = _tool180            self.tools.append(_tool)181        logger.info(182            f'Tools updated for agent {self.name}, total {len(self.tools)}: {[tool["function"]["name"] for tool in self.tools]}'183        )184