feat(agent): dbgpts support agent (#1417)

dbgpt/agent/__init__.py

@@ -7,6 +7,11 @@
     AgentGenerateContext,
     AgentMessage,
 )
+from .core.agent_manage import (  # noqa: F401
+    AgentManager,
+    get_agent_manager,
+    initialize_agent,
+)
 from .core.base_agent import ConversableAgent  # noqa: F401
 from .core.llm.llm import LLMConfig  # noqa: F401
 from .core.schema import PluginStorageType  # noqa: F401
@@ -20,6 +25,9 @@
     "AgentContext",
     "AgentGenerateContext",
     "AgentMessage",
+    "AgentManager",
+    "initialize_agent",
+    "get_agent_manager",
     "ConversableAgent",
     "Action",
     "ActionOutput",

dbgpt/agent/core/agent_manage.py

@@ -3,14 +3,12 @@
 import logging
 import re
 from collections import defaultdict
-from typing import Dict, List, Type
+from typing import Dict, List, Optional, Set, Tuple, Type, cast
+
+from dbgpt.component import BaseComponent, ComponentType, SystemApp
 
-from ..expand.code_assistant_agent import CodeAssistantAgent
-from ..expand.dashboard_assistant_agent import DashboardAssistantAgent
-from ..expand.data_scientist_agent import DataScientistAgent
-from ..expand.plugin_assistant_agent import PluginAssistantAgent
-from ..expand.summary_assistant_agent import SummaryAssistantAgent
 from .agent import Agent
+from .base_agent import ConversableAgent
 
 logger = logging.getLogger(__name__)
 
@@ -46,62 +44,117 @@ def mentioned_agents(message_content: str, agents: List[Agent]) -> Dict:
     return mentions
 
 
-class AgentManager:
+class AgentManager(BaseComponent):
     """Manages the registration and retrieval of agents."""
 
-    def __init__(self):
-        """Create a new AgentManager."""
-        self._agents = defaultdict()
+    name = ComponentType.AGENT_MANAGER
 
-    def register_agent(self, cls):
+    def __init__(self, system_app: SystemApp):
+        """Create a new AgentManager."""
+        super().__init__(system_app)
+        self.system_app = system_app
+        self._agents: Dict[
+            str, Tuple[Type[ConversableAgent], ConversableAgent]
+        ] = defaultdict()
+
+        self._core_agents: Set[str] = set()
+
+    def init_app(self, system_app: SystemApp):
+        """Initialize the AgentManager."""
+        self.system_app = system_app
+
+    def after_start(self):
+        """Register all agents."""
+        from ..expand.code_assistant_agent import CodeAssistantAgent
+        from ..expand.dashboard_assistant_agent import DashboardAssistantAgent
+        from ..expand.data_scientist_agent import DataScientistAgent
+        from ..expand.plugin_assistant_agent import PluginAssistantAgent
+        from ..expand.summary_assistant_agent import SummaryAssistantAgent
+
+        core_agents = set()
+        core_agents.add(self.register_agent(CodeAssistantAgent))
+        core_agents.add(self.register_agent(DashboardAssistantAgent))
+        core_agents.add(self.register_agent(DataScientistAgent))
+        core_agents.add(self.register_agent(SummaryAssistantAgent))
+        core_agents.add(self.register_agent(PluginAssistantAgent))
+        self._core_agents = core_agents
+
+    def register_agent(
+        self, cls: Type[ConversableAgent], ignore_duplicate: bool = False
+    ) -> str:
         """Register an agent."""
-        self._agents[cls().profile] = cls
-
-    def get_by_name(self, name: str) -> Type[Agent]:
+        inst = cls()
+        profile = inst.get_profile()
+        if profile in self._agents and (
+            profile in self._core_agents or not ignore_duplicate
+        ):
+            raise ValueError(f"Agent:{profile} already register!")
+        self._agents[profile] = (cls, inst)
+        return profile
+
+    def get_by_name(self, name: str) -> Type[ConversableAgent]:
         """Return an agent by name.
 
         Args:
             name (str): The name of the agent to retrieve.
 
         Returns:
-            Type[Agent]: The agent with the given name.
+            Type[ConversableAgent]: The agent with the given name.
 
         Raises:
             ValueError: If the agent with the given name is not registered.
         """
         if name not in self._agents:
             raise ValueError(f"Agent:{name} not register!")
-        return self._agents[name]
+        return self._agents[name][0]
 
     def get_describe_by_name(self, name: str) -> str:
         """Return the description of an agent by name."""
-        return self._agents[name].desc
+        return self._agents[name][1].desc
 
     def all_agents(self):
         """Return a dictionary of all registered agents and their descriptions."""
         result = {}
-        for name, cls in self._agents.items():
-            result[name] = cls.desc
+        for name, value in self._agents.items():
+            result[name] = value[1].desc
         return result
 
     def list_agents(self):
         """Return a list of all registered agents and their descriptions."""
         result = []
-        for name, cls in self._agents.items():
-            instance = cls()
+        for name, value in self._agents.items():
             result.append(
                 {
-                    "name": instance.profile,
-                    "desc": instance.goal,
+                    "name": value[1].profile,
+                    "desc": value[1].goal,
                 }
             )
         return result
 
 
-agent_manager = AgentManager()
+_SYSTEM_APP: Optional[SystemApp] = None
+
 
-agent_manager.register_agent(CodeAssistantAgent)
-agent_manager.register_agent(DashboardAssistantAgent)
-agent_manager.register_agent(DataScientistAgent)
-agent_manager.register_agent(SummaryAssistantAgent)
-agent_manager.register_agent(PluginAssistantAgent)
+def initialize_agent(system_app: SystemApp):
+    """Initialize the agent manager."""
+    global _SYSTEM_APP
+    _SYSTEM_APP = system_app
+    agent_manager = AgentManager(system_app)
+    system_app.register_instance(agent_manager)
+
+
+def get_agent_manager(system_app: Optional[SystemApp] = None) -> AgentManager:
+    """Return the agent manager.
+
+    Args:
+        system_app (Optional[SystemApp], optional): The system app. Defaults to None.
+
+    Returns:
+        AgentManager: The agent manager.
+    """
+    if not _SYSTEM_APP:
+        if not system_app:
+            system_app = SystemApp()
+        initialize_agent(system_app)
+    app = system_app or _SYSTEM_APP
+    return AgentManager.get_instance(cast(SystemApp, app))

dbgpt/agent/core/base_agent.py

@@ -25,7 +25,7 @@
 
 
 class ConversableAgent(Role, Agent):
-    """ConversableAgent is a agent that can communicate with other agents."""
+    """ConversableAgent is an agent that can communicate with other agents."""
 
     agent_context: Optional[AgentContext] = Field(None, description="Agent context")
     actions: List[Action] = Field(default_factory=list)

dbgpt/agent/plan/awel/agent_operator.py

@@ -17,7 +17,7 @@
 from dbgpt.model.operators.llm_operator import MixinLLMOperator
 
 from ...core.agent import Agent, AgentGenerateContext, AgentMessage
-from ...core.agent_manage import agent_manager
+from ...core.agent_manage import get_agent_manager
 from ...core.base_agent import ConversableAgent
 from ...core.llm.llm import LLMConfig
 from .agent_operator_resource import AWELAgent
@@ -244,7 +244,7 @@ async def get_agent(
     ) -> ConversableAgent:
         """Build the agent."""
         # agent build
-        agent_cls: Type[ConversableAgent] = agent_manager.get_by_name(  # type: ignore
+        agent_cls: Type[ConversableAgent] = get_agent_manager().get_by_name(
             self.awel_agent.agent_profile
         )
         llm_config = self.awel_agent.llm_config

dbgpt/agent/plan/awel/agent_operator_resource.py

@@ -11,7 +11,7 @@
     register_resource,
 )
 
-from ...core.agent_manage import agent_manager
+from ...core.agent_manage import get_agent_manager
 from ...core.llm.llm import LLMConfig, LLMStrategyType
 from ...resource.resource_api import AgentResource, ResourceType
 
@@ -118,7 +118,7 @@ def pre_fill(cls, values: Dict[str, Any]) -> Dict[str, Any]:
 def _agent_resource_option_values() -> List[OptionValue]:
     return [
         OptionValue(label=item["name"], name=item["name"], value=item["name"])
-        for item in agent_manager.list_agents()
+        for item in get_agent_manager().list_agents()
     ]
 
 

dbgpt/app/component_configs.py

@@ -47,6 +47,7 @@ def initialize_components(
     )
     _initialize_model_cache(system_app)
     _initialize_awel(system_app, param)
+    _initialize_agent(system_app)
     _initialize_openapi(system_app)
     # Register serve apps
     register_serve_apps(system_app, CFG)
@@ -78,6 +79,12 @@ def _initialize_awel(system_app: SystemApp, param: WebServerParameters):
     initialize_awel(system_app, dag_dirs)
 
 
+def _initialize_agent(system_app: SystemApp):
+    from dbgpt.agent import initialize_agent
+
+    initialize_agent(system_app)
+
+
 def _initialize_openapi(system_app: SystemApp):
     from dbgpt.app.openapi.api_v1.editor.service import EditorService
 

dbgpt/component.py

@@ -86,6 +86,7 @@ class ComponentType(str, Enum):
     AWEL_DAG_MANAGER = "dbgpt_awel_dag_manager"
     UNIFIED_METADATA_DB_MANAGER_FACTORY = "dbgpt_unified_metadata_db_manager_factory"
     CONNECTOR_MANAGER = "dbgpt_connector_manager"
+    AGENT_MANAGER = "dbgpt_agent_manager"
 
 
 _EMPTY_DEFAULT_COMPONENT = "_EMPTY_DEFAULT_COMPONENT"

dbgpt/serve/agent/agents/controller.py

@@ -10,7 +10,7 @@
 
 from dbgpt._private.config import Config
 from dbgpt.agent.core.agent import Agent, AgentContext
-from dbgpt.agent.core.agent_manage import agent_manager
+from dbgpt.agent.core.agent_manage import get_agent_manager
 from dbgpt.agent.core.base_agent import ConversableAgent
 from dbgpt.agent.core.llm.llm import LLMConfig, LLMStrategyType
 from dbgpt.agent.core.schema import Status
@@ -222,7 +222,9 @@ async def agent_team_chat_new(
         self.llm_provider = DefaultLLMClient(worker_manager, auto_convert_message=True)
 
         for record in gpts_app.details:
-            cls: Type[ConversableAgent] = agent_manager.get_by_name(record.agent_name)
+            cls: Type[ConversableAgent] = get_agent_manager().get_by_name(
+                record.agent_name
+            )
             llm_config = LLMConfig(
                 llm_client=self.llm_provider,
                 llm_strategy=LLMStrategyType(record.llm_strategy),
@@ -340,7 +342,7 @@ def gpts_conv_list(self, user_code: str = None, system_app: str = None):
 async def agents_list():
     logger.info("agents_list!")
     try:
-        agents = agent_manager.all_agents()
+        agents = get_agent_manager().all_agents()
         return Result.succ(agents)
     except Exception as e:
         return Result.failed(code="E30001", msg=str(e))

dbgpt/serve/agent/app/controller.py

@@ -3,7 +3,7 @@
 from fastapi import APIRouter
 
 from dbgpt._private.config import Config
-from dbgpt.agent.core.agent_manage import agent_manager
+from dbgpt.agent.core.agent_manage import get_agent_manager
 from dbgpt.agent.core.llm.llm import LLMStrategyType
 from dbgpt.agent.resource.resource_api import ResourceType
 from dbgpt.app.knowledge.api import knowledge_space_service
@@ -63,7 +63,7 @@ async def edit(gpts_app: GptsApp):
 @router.get("/v1/agents/list")
 async def all_agents():
     try:
-        return Result.succ(agent_manager.list_agents())
+        return Result.succ(get_agent_manager().list_agents())
     except Exception as ex:
         return Result.failed(code="E000X", msg=f"query agents error: {ex}")