Coverage for autogpt/agent/agent_manager.py: 13%

1"""Agent manager for managing GPT agents""" 

2from __future__ import annotations 

3 

4from typing import List, Union 

5 

6from autogpt.config.config import Config, Singleton 

7from autogpt.llm_utils import create_chat_completion 

8from autogpt.types.openai import Message 

9 

10 

11class AgentManager(metaclass=Singleton): 

12 """Agent manager for managing GPT agents""" 

13 

14 def __init__(self): 

15 self.next_key = 0 

16 self.agents = {} # key, (task, full_message_history, model) 

17 self.cfg = Config() 

18 

19 # Create new GPT agent 

20 # TODO: Centralise use of create_chat_completion() to globally enforce token limit 

21 

22 def create_agent(self, task: str, prompt: str, model: str) -> tuple[int, str]: 

23 """Create a new agent and return its key 

24 

25 Args: 

26 task: The task to perform 

27 prompt: The prompt to use 

28 model: The model to use 

29 

30 Returns: 

31 The key of the new agent 

32 """ 

33 messages: List[Message] = [ 

34 {"role": "user", "content": prompt}, 

35 ] 

36 for plugin in self.cfg.plugins: 

37 if not plugin.can_handle_pre_instruction(): 

38 continue 

39 if plugin_messages := plugin.pre_instruction(messages): 

40 messages.extend(iter(plugin_messages)) 

41 # Start GPT instance 

42 agent_reply = create_chat_completion( 

43 model=model, 

44 messages=messages, 

45 ) 

46 

47 messages.append({"role": "assistant", "content": agent_reply}) 

48 

49 plugins_reply = "" 

50 for i, plugin in enumerate(self.cfg.plugins): 

51 if not plugin.can_handle_on_instruction(): 

52 continue 

53 if plugin_result := plugin.on_instruction(messages): 

54 sep = "\n" if i else "" 

55 plugins_reply = f"{plugins_reply}{sep}{plugin_result}" 

56 

57 if plugins_reply and plugins_reply != "": 

58 messages.append({"role": "assistant", "content": plugins_reply}) 

59 key = self.next_key 

60 # This is done instead of len(agents) to make keys unique even if agents 

61 # are deleted 

62 self.next_key += 1 

63 

64 self.agents[key] = (task, messages, model) 

65 

66 for plugin in self.cfg.plugins: 

67 if not plugin.can_handle_post_instruction(): 

68 continue 

69 agent_reply = plugin.post_instruction(agent_reply) 

70 

71 return key, agent_reply 

72 

73 def message_agent(self, key: str | int, message: str) -> str: 

74 """Send a message to an agent and return its response 

75 

76 Args: 

77 key: The key of the agent to message 

78 message: The message to send to the agent 

79 

80 Returns: 

81 The agent's response 

82 """ 

83 task, messages, model = self.agents[int(key)] 

84 

85 # Add user message to message history before sending to agent 

86 messages.append({"role": "user", "content": message}) 

87 

88 for plugin in self.cfg.plugins: 

89 if not plugin.can_handle_pre_instruction(): 

90 continue 

91 if plugin_messages := plugin.pre_instruction(messages): 

92 for plugin_message in plugin_messages: 

93 messages.append(plugin_message) 

94 

95 # Start GPT instance 

96 agent_reply = create_chat_completion( 

97 model=model, 

98 messages=messages, 

99 ) 

100 

101 messages.append({"role": "assistant", "content": agent_reply}) 

102 

103 plugins_reply = agent_reply 

104 for i, plugin in enumerate(self.cfg.plugins): 

105 if not plugin.can_handle_on_instruction(): 

106 continue 

107 if plugin_result := plugin.on_instruction(messages): 

108 sep = "\n" if i else "" 

109 plugins_reply = f"{plugins_reply}{sep}{plugin_result}" 

110 # Update full message history 

111 if plugins_reply and plugins_reply != "": 

112 messages.append({"role": "assistant", "content": plugins_reply}) 

113 

114 for plugin in self.cfg.plugins: 

115 if not plugin.can_handle_post_instruction(): 

116 continue 

117 agent_reply = plugin.post_instruction(agent_reply) 

118 

119 return agent_reply 

120 

121 def list_agents(self) -> list[tuple[str | int, str]]: 

122 """Return a list of all agents 

123 

124 Returns: 

125 A list of tuples of the form (key, task) 

126 """ 

127 

128 # Return a list of agent keys and their tasks 

129 return [(key, task) for key, (task, _, _) in self.agents.items()] 

130 

131 def delete_agent(self, key: str | int) -> bool: 

132 """Delete an agent from the agent manager 

133 

134 Args: 

135 key: The key of the agent to delete 

136 

137 Returns: 

138 True if successful, False otherwise 

139 """ 

140 

141 try: 

142 del self.agents[int(key)] 

143 return True 

144 except KeyError: 

145 return False