github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-09-04 10:34:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(agent): Refactor resource of agents (#1518)

Browse Source
This commit is contained in:
Fangyin Cheng
2024-05-15 09:57:19 +08:00
committed by GitHub
parent db4d318a5f
commit 559affe87d
102 changed files with 2633 additions and 2549 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
dbgpt/agent/core/action/base.py
View File

@@ -27,8 +27,7 @@ from dbgpt._private.pydantic import (
from dbgpt.util.json_utils import find_json_objects
from dbgpt.vis.base import Vis
from ...resource.resource_api import AgentResource, ResourceType
from ...resource.resource_loader import ResourceLoader
from ...resource.base import AgentResource, Resource, ResourceType
T = TypeVar("T", bound=Union[BaseModel, List[BaseModel], None])
@@ -77,11 +76,11 @@ class Action(ABC, Generic[T]):
def __init__(self):
"""Create an action."""
self.resource_loader: Optional[ResourceLoader] = None
self.resource: Optional[Resource] = None
def init_resource_loader(self, resource_loader: Optional[ResourceLoader]):
"""Initialize the resource loader."""
self.resource_loader = resource_loader
def init_resource(self, resource: Optional[Resource]):
"""Initialize the resource."""
self.resource = resource
@property
def resource_need(self) -> Optional[ResourceType]:

2
dbgpt/agent/core/action/blank_action.py
View File

@@ -3,7 +3,7 @@
import logging
from typing import Optional
from ...resource.resource_api import AgentResource
from ...resource.base import AgentResource
from .base import Action, ActionOutput
logger = logging.getLogger(__name__)

2
dbgpt/agent/core/agent.py
View File

@@ -9,7 +9,6 @@ from typing import Any, Dict, List, Optional, Tuple, Union
from dbgpt.core import LLMClient
from dbgpt.util.annotations import PublicAPI
from ..resource.resource_loader import ResourceLoader
from .action.base import ActionOutput
from .memory.agent_memory import AgentMemory
@@ -209,7 +208,6 @@ class AgentGenerateContext:
memory: Optional[AgentMemory] = None
agent_context: Optional[AgentContext] = None
resource_loader: Optional[ResourceLoader] = None
llm_client: Optional[LLMClient] = None
round_index: Optional[int] = None

4
dbgpt/agent/core/agent_manage.py
View File

@@ -68,15 +68,15 @@ class AgentManager(BaseComponent):
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 ..expand.tool_assistant_agent import ToolAssistantAgent
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))
core_agents.add(self.register_agent(ToolAssistantAgent))
self._core_agents = core_agents
def register_agent(

104
dbgpt/agent/core/base_agent.py
View File

@@ -3,16 +3,17 @@
import asyncio
import json
import logging
from typing import Any, Dict, List, Optional, Tuple, Type, cast
from concurrent.futures import Executor, ThreadPoolExecutor
from typing import Any, Callable, Dict, List, Optional, Tuple, Type, cast
from dbgpt._private.pydantic import ConfigDict, Field
from dbgpt.core import LLMClient, ModelMessageRoleType
from dbgpt.util.error_types import LLMChatError
from dbgpt.util.executor_utils import blocking_func_to_async
from dbgpt.util.tracer import SpanType, root_tracer
from dbgpt.util.utils import colored
from ..resource.resource_api import AgentResource, ResourceClient
from ..resource.resource_loader import ResourceLoader
from ..resource.base import Resource
from ..util.llm.llm import LLMConfig, LLMStrategyType
from ..util.llm.llm_client import AIWrapper
from .action.base import Action, ActionOutput
@@ -32,12 +33,15 @@ class ConversableAgent(Role, Agent):
agent_context: Optional[AgentContext] = Field(None, description="Agent context")
actions: List[Action] = Field(default_factory=list)
resources: List[AgentResource] = Field(default_factory=list)
resource: Optional[Resource] = Field(None, description="Resource")
llm_config: Optional[LLMConfig] = None
resource_loader: Optional[ResourceLoader] = None
max_retry_count: int = 3
consecutive_auto_reply_counter: int = 0
llm_client: Optional[AIWrapper] = None
executor: Executor = Field(
default_factory=lambda: ThreadPoolExecutor(max_workers=1),
description="Executor for running tasks",
)
def __init__(self, **kwargs):
"""Create a new agent."""
@@ -58,27 +62,12 @@ class ConversableAgent(Role, Agent):
f"running!"
)
# resource check
for resource in self.resources:
if (
self.resource_loader is None
or self.resource_loader.get_resource_api(
resource.type, check_instance=False
)
is None
):
raise ValueError(
f"Resource {resource.type}:{resource.value} missing resource loader"
f" implementation,unable to read resources!"
)
# action check
if self.actions and len(self.actions) > 0:
have_resource_types = [item.type for item in self.resources]
for action in self.actions:
if (
action.resource_need
and action.resource_need not in have_resource_types
if action.resource_need and (
not self.resource
or not self.resource.get_resource_by_type(action.resource_need)
):
raise ValueError(
f"{self.name}[{self.role}] Missing resources required for "
@@ -112,13 +101,6 @@ class ConversableAgent(Role, Agent):
raise ValueError("Agent context is not initialized！")
return self.agent_context
@property
def not_null_resource_loader(self) -> ResourceLoader:
"""Get the resource loader."""
if not self.resource_loader:
raise ValueError("Resource loader is not initialized！")
return self.resource_loader
@property
def not_null_llm_config(self) -> LLMConfig:
"""Get the LLM config."""
@@ -134,23 +116,32 @@ class ConversableAgent(Role, Agent):
raise ValueError("LLM client is not initialized！")
return llm_client
async def blocking_func_to_async(
self, func: Callable[..., Any], *args, **kwargs
) -> Any:
"""Run a potentially blocking function within an executor."""
if not asyncio.iscoroutinefunction(func):
return await blocking_func_to_async(self.executor, func, *args, **kwargs)
return await func(*args, **kwargs)
async def preload_resource(self) -> None:
"""Preload resources before agent initialization."""
pass
if self.resource:
await self.blocking_func_to_async(self.resource.preload_resource)
async def build(self) -> "ConversableAgent":
"""Build the agent."""
# Preload resources
await self.preload_resource()
# Check if agent is available
self.check_available()
_language = self.not_null_agent_context.language
if _language:
self.language = _language
# Preload resources
await self.preload_resource()
# Initialize resource loader
for action in self.actions:
action.init_resource_loader(self.resource_loader)
action.init_resource(self.resource)
# Initialize LLM Server
if not self.is_human:
@@ -175,13 +166,8 @@ class ConversableAgent(Role, Agent):
raise ValueError("GptsMemory is not supported!")
elif isinstance(target, AgentContext):
self.agent_context = target
elif isinstance(target, ResourceLoader):
self.resource_loader = target
elif isinstance(target, list) and target and len(target) > 0:
if _is_list_of_type(target, Action):
self.actions.extend(target)
elif _is_list_of_type(target, AgentResource):
self.resources = target
elif isinstance(target, Resource):
self.resource = target
elif isinstance(target, AgentMemory):
self.memory = target
return self
@@ -480,12 +466,12 @@ class ConversableAgent(Role, Agent):
last_out: Optional[ActionOutput] = None
for i, action in enumerate(self.actions):
# Select the resources required by acton
need_resource = None
if self.resources and len(self.resources) > 0:
for item in self.resources:
if item.type == action.resource_need:
need_resource = item
break
if action.resource_need and self.resource:
need_resources = self.resource.get_resource_by_type(
action.resource_need
)
else:
need_resources = []
if not message:
raise ValueError("The message content is empty!")
@@ -497,7 +483,7 @@ class ConversableAgent(Role, Agent):
"sender": sender.name if sender else None,
"recipient": self.name,
"reviewer": reviewer.name if reviewer else None,
"need_resource": need_resource.to_dict() if need_resource else None,
"need_resource": need_resources[0].name if need_resources else None,
"rely_action_out": last_out.to_dict() if last_out else None,
"conv_uid": self.not_null_agent_context.conv_id,
"action_index": i,
@@ -506,7 +492,7 @@ class ConversableAgent(Role, Agent):
) as span:
last_out = await action.run(
ai_message=message,
resource=need_resource,
resource=None,
rely_action_out=last_out,
**kwargs,
)
@@ -703,23 +689,11 @@ class ConversableAgent(Role, Agent):
self, question: Optional[str] = None
) -> Dict[str, Any]:
"""Generate the resource variables."""
resource_prompt_list = []
for item in self.resources:
resource_client = self.not_null_resource_loader.get_resource_api(
item.type, ResourceClient
resource_prompt = None
if self.resource:
resource_prompt = await self.resource.get_prompt(
lang=self.language, question=question
)
if not resource_client:
raise ValueError(
f"Resource {item.type}:{item.value} missing resource loader"
f" implementation,unable to read resources!"
)
resource_prompt_list.append(
await resource_client.get_resource_prompt(item, question)
)
resource_prompt = ""
if len(resource_prompt_list) > 0:
resource_prompt = "RESOURCES:" + "\n".join(resource_prompt_list)
out_schema: Optional[str] = ""
if self.actions and len(self.actions) > 0:

6
dbgpt/agent/core/plan/awel/agent_operator.py
View File

@@ -17,6 +17,7 @@ from dbgpt.core.interface.message import ModelMessageRoleType
# TODO: Don't dependent on MixinLLMOperator
from dbgpt.model.operators.llm_operator import MixinLLMOperator
from ....resource.manage import get_resource_manager
from ....util.llm.llm import LLMConfig
from ...agent import Agent, AgentGenerateContext, AgentMessage
from ...agent_manage import get_agent_manager
@@ -228,7 +229,6 @@ class AWELAgentOperator(
silent=input_value.silent,
memory=input_value.memory.structure_clone() if input_value.memory else None,
agent_context=input_value.agent_context,
resource_loader=input_value.resource_loader,
llm_client=input_value.llm_client,
round_index=agent.consecutive_auto_reply_counter,
)
@@ -262,13 +262,13 @@ class AWELAgentOperator(
if self.awel_agent.fixed_subgoal:
kwargs["fixed_subgoal"] = self.awel_agent.fixed_subgoal
resource = get_resource_manager().build_resource(self.awel_agent.resources)
agent = (
await agent_cls(**kwargs)
.bind(input_value.memory)
.bind(llm_config)
.bind(input_value.agent_context)
.bind(self.awel_agent.resources)
.bind(input_value.resource_loader)
.bind(resource)
.build()
)

15
dbgpt/agent/core/plan/awel/agent_operator_resource.py
View File

@@ -12,11 +12,17 @@ from dbgpt.core.awel.flow import (
register_resource,
)
from ....resource.resource_api import AgentResource, ResourceType
from ....resource.base import AgentResource
from ....resource.manage import get_resource_manager
from ....util.llm.llm import LLMConfig, LLMStrategyType
from ...agent_manage import get_agent_manager
def _load_resource_types():
resources = get_resource_manager().get_supported_resources()
return [OptionValue(label=item, name=item, value=item) for item in resources.keys()]
@register_resource(
label="AWEL Agent Resource",
name="agent_operator_resource",
@@ -29,10 +35,7 @@ from ...agent_manage import get_agent_manager
type=str,
optional=True,
default=None,
options=[
OptionValue(label=item.name, name=item.value, value=item.value)
for item in ResourceType
],
options=FunctionDynamicOptions(func=_load_resource_types),
),
Parameter.build_from(
label="Agent Resource Name",
@@ -70,7 +73,7 @@ class AWELAgentResource(AgentResource):
value = values.pop("agent_resource_value")
values["name"] = name
values["type"] = ResourceType(type)
values["type"] = type
values["value"] = value
return values

1
dbgpt/agent/core/plan/awel/team_awel_layout.py
View File

@@ -132,7 +132,6 @@ class AWELBaseManager(ManagerAgent, ABC):
reviewer=reviewer,
memory=self.memory.structure_clone(),
agent_context=self.agent_context,
resource_loader=self.resource_loader,
llm_client=self.not_null_llm_config.llm_client,
)
final_generate_context: AgentGenerateContext = await last_node.call(

2
dbgpt/agent/core/plan/plan_action.py
View File

@@ -6,7 +6,7 @@ from typing import List, Optional
from dbgpt._private.pydantic import BaseModel, Field
from dbgpt.vis.tags.vis_agent_plans import Vis, VisAgentPlans
from ...resource.resource_api import AgentResource
from ...resource.base import AgentResource
from ..action.base import Action, ActionOutput
from ..agent import AgentContext
from ..memory.gpts.base import GptsPlan

7
dbgpt/agent/core/plan/planner_agent.py
View File

@@ -4,6 +4,7 @@ from typing import Any, Dict, List
from dbgpt._private.pydantic import Field
from ...resource.pack import ResourcePack
from ..agent import AgentMessage
from ..base_agent import ConversableAgent
from ..plan.plan_action import PlanAction
@@ -152,9 +153,11 @@ assistants:[
def bind_agents(self, agents: List[ConversableAgent]) -> ConversableAgent:
"""Bind the agents to the planner agent."""
self.agents = agents
resources = []
for agent in self.agents:
if agent.resources and len(agent.resources) > 0:
self.resources.extend(agent.resources)
if agent.resource:
resources.append(agent.resource)
self.resource = ResourcePack(resources)
return self
def prepare_act_param(self) -> Dict[str, Any]:

1
dbgpt/agent/core/plan/team_auto_plan.py
View File

@@ -188,7 +188,6 @@ class AutoPlanChatManager(ManagerAgent):
.bind(self.memory)
.bind(self.agent_context)
.bind(self.llm_config)
.bind(self.resource_loader)
.bind_agents(self.agents)
.build()
)