mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-10-22 17:39:02 +00:00
87 lines
2.3 KiB
Python
87 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding:utf-8 -*-
|
|
"""Agent manager for managing GPT agents"""
|
|
from __future__ import annotations
|
|
|
|
from pilot.configs.config import Config
|
|
from pilot.singleton import Singleton
|
|
from pilot.configs.config import Config
|
|
from typing import List
|
|
from pilot.model.base import Message
|
|
|
|
|
|
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)
|
|
self.cfg = Config()
|
|
"""Agent manager for managing DB-GPT agents
|
|
In order to compatible auto gpt plugins,
|
|
we use the same template with it.
|
|
|
|
Args: next_keys
|
|
agents
|
|
cfg
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self.next_key = 0
|
|
self.agents = {} #TODO need to define
|
|
self.cfg = Config()
|
|
|
|
# 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
|
|
"""
|
|
|
|
|
|
def message_agent(self, key: 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
|
|
"""
|
|
|
|
|
|
def list_agents(self) -> list[tuple[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: 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
|