101 lines
2.8 KiB
Python
101 lines
2.8 KiB
Python
"""Agent manager for managing GPT agents"""
|
|
from typing import List, Tuple, Union
|
|
from autogpt.llm_utils import create_chat_completion
|
|
from autogpt.config.config import Singleton
|
|
|
|
|
|
class AgentManager(metaclass=Singleton):
|
|
"""Agent manager for managing GPT agents"""
|
|
|
|
def __init__(self):
|
|
self.next_key = 0
|
|
self.agents = {} # key, (task, full_message_history, model)
|
|
|
|
# Create new GPT agent
|
|
# TODO: Centralise use of create_chat_completion() to globally enforce token limit
|
|
|
|
def create_agent(self, task: str, prompt: str, model: str) -> Tuple[int, str]:
|
|
"""Create a new agent and return its key
|
|
|
|
Args:
|
|
task: The task to perform
|
|
prompt: The prompt to use
|
|
model: The model to use
|
|
|
|
Returns:
|
|
The key of the new agent
|
|
"""
|
|
messages = [
|
|
{"role": "user", "content": prompt},
|
|
]
|
|
|
|
# Start GPT instance
|
|
agent_reply = create_chat_completion(
|
|
model=model,
|
|
messages=messages,
|
|
)
|
|
|
|
# Update full message history
|
|
messages.append({"role": "assistant", "content": agent_reply})
|
|
|
|
key = self.next_key
|
|
# This is done instead of len(agents) to make keys unique even if agents
|
|
# are deleted
|
|
self.next_key += 1
|
|
|
|
self.agents[key] = (task, messages, model)
|
|
|
|
return key, agent_reply
|
|
|
|
def message_agent(self, key: Union[str, int], message: str) -> str:
|
|
"""Send a message to an agent and return its response
|
|
|
|
Args:
|
|
key: The key of the agent to message
|
|
message: The message to send to the agent
|
|
|
|
Returns:
|
|
The agent's response
|
|
"""
|
|
task, messages, model = self.agents[int(key)]
|
|
|
|
# Add user message to message history before sending to agent
|
|
messages.append({"role": "user", "content": message})
|
|
|
|
# Start GPT instance
|
|
agent_reply = create_chat_completion(
|
|
model=model,
|
|
messages=messages,
|
|
)
|
|
|
|
# Update full message history
|
|
messages.append({"role": "assistant", "content": agent_reply})
|
|
|
|
return agent_reply
|
|
|
|
def list_agents(self) -> List[Tuple[Union[str, int], str]]:
|
|
"""Return a list of all agents
|
|
|
|
Returns:
|
|
A list of tuples of the form (key, task)
|
|
"""
|
|
|
|
# Return a list of agent keys and their tasks
|
|
return [(key, task) for key, (task, _, _) in self.agents.items()]
|
|
|
|
def delete_agent(self, key: Union[str, int]) -> bool:
|
|
"""Delete an agent from the agent manager
|
|
|
|
Args:
|
|
key: The key of the agent to delete
|
|
|
|
Returns:
|
|
True if successful, False otherwise
|
|
"""
|
|
|
|
try:
|
|
del self.agents[int(key)]
|
|
return True
|
|
except KeyError:
|
|
return False
|