agent.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158
  1. from abc import ABC, abstractmethod
  2. from typing import List, Dict, Type
  3. from dataclasses import dataclass
  4. from enum import Enum
  5. from .lib.event import Event
  6. from .lib.command_manager import CommandManager
  7. from .controller import AgentController
  8. class Role(Enum):
  9. SYSTEM = "system" # system message for LLM
  10. USER = "user" # the user
  11. ASSISTANT = "assistant" # the agent
  12. ENVIRONMENT = "environment" # the environment (e.g., bash shell, web browser, etc.)
  13. @dataclass
  14. class Message:
  15. """
  16. This data class represents a message sent by an agent to another agent or user.
  17. """
  18. role: Role
  19. content: str
  20. # TODO: add more fields as needed
  21. def to_dict(self) -> Dict:
  22. """
  23. Converts the message to a dictionary (OpenAI chat-completion format).
  24. Returns:
  25. - message (Dict): A dictionary representation of the message.
  26. """
  27. role = self.role.value
  28. content = self.content
  29. if self.role == Role.ENVIRONMENT:
  30. content = f"Environment Observation:\n{content}"
  31. role = "user" # treat environment messages as user messages
  32. return {"role": role, "content": content}
  33. class Agent(ABC):
  34. """
  35. This abstract base class is an general interface for an agent dedicated to
  36. executing a specific instruction and allowing human interaction with the
  37. agent during execution.
  38. It tracks the execution status and maintains a history of interactions.
  39. :param instruction: The instruction for the agent to execute.
  40. :param workspace_dir: The working directory for the agent.
  41. :param model_name: The litellm name of the model to use for the agent.
  42. :param max_steps: The maximum number of steps to run the agent.
  43. """
  44. _registry: Dict[str, Type['Agent']] = {}
  45. def __init__(
  46. self,
  47. instruction: str,
  48. workspace_dir: str,
  49. model_name: str,
  50. max_steps: int = 100
  51. ):
  52. self.instruction = instruction
  53. self.workspace_dir = workspace_dir
  54. self.model_name = model_name
  55. self.max_steps = max_steps
  56. self._complete = False
  57. self._history: List[Message] = [Message(Role.USER, instruction)]
  58. @property
  59. def complete(self) -> bool:
  60. """
  61. Indicates whether the current instruction execution is complete.
  62. Returns:
  63. - complete (bool): True if execution is complete; False otherwise.
  64. """
  65. return self._complete
  66. @property
  67. def history(self) -> List[str]:
  68. """
  69. Provides the history of interactions or state changes since the instruction was initiated.
  70. Returns:
  71. - history (List[str]): A list of strings representing the history.
  72. """
  73. return self._history
  74. @abstractmethod
  75. def add_event(self, event: Event) -> None:
  76. """
  77. Adds an event to the agent's history.
  78. Parameters:
  79. - event (Event): The event to add to the history.
  80. """
  81. pass
  82. @abstractmethod
  83. def step(self, cmd_mgr: CommandManager) -> Event:
  84. """
  85. Starts the execution of the assigned instruction. This method should
  86. be implemented by subclasses to define the specific execution logic.
  87. """
  88. pass
  89. @abstractmethod
  90. def search_memory(self, query: str) -> List[str]:
  91. """
  92. Searches the agent's memory for information relevant to the given query.
  93. Parameters:
  94. - query (str): The query to search for in the agent's memory.
  95. Returns:
  96. - response (str): The response to the query.
  97. """
  98. pass
  99. def reset(self) -> None:
  100. """
  101. Resets the agent's execution status and clears the history. This method can be used
  102. to prepare the agent for restarting the instruction or cleaning up before destruction.
  103. """
  104. self.instruction = None
  105. self._complete = False
  106. self._history = []
  107. @classmethod
  108. def register(cls, name: str, agent_cls: Type['Agent']):
  109. """
  110. Registers an agent class in the registry.
  111. Parameters:
  112. - name (str): The name to register the class under.
  113. - agent_cls (Type['Agent']): The class to register.
  114. """
  115. if name in cls._registry:
  116. raise ValueError(f"Agent class already registered under '{name}'.")
  117. cls._registry[name] = agent_cls
  118. @classmethod
  119. def get_cls(cls, name: str) -> Type['Agent']:
  120. """
  121. Retrieves an agent class from the registry.
  122. Parameters:
  123. - name (str): The name of the class to retrieve
  124. Returns:
  125. - agent_cls (Type['Agent']): The class registered under the specified name.
  126. """
  127. if name not in cls._registry:
  128. raise ValueError(f"No agent class registered under '{name}'.")
  129. return cls._registry[name]