github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-09-05 19:11:52 +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

23
dbgpt/app/base.py
View File

@@ -31,7 +31,6 @@ def async_db_summary(system_app: SystemApp):
def server_init(param: "WebServerParameters", system_app: SystemApp):
from dbgpt.agent.plugin.commands.command_manage import CommandRegistry
# logger.info(f"args: {args}")
# init config
@@ -43,28 +42,6 @@ def server_init(param: "WebServerParameters", system_app: SystemApp):
# load_native_plugins(cfg)
signal.signal(signal.SIGINT, signal_handler)
# Loader plugins and commands
command_categories = []
# exclude commands
command_categories = [
x for x in command_categories if x not in cfg.disabled_command_categories
]
command_registry = CommandRegistry()
for command_category in command_categories:
command_registry.import_commands(command_category)
cfg.command_registry = command_registry
command_dispaly_commands = [
"dbgpt.agent.plugin.commands.built_in.display_type.show_chart_gen",
"dbgpt.agent.plugin.commands.built_in.display_type.show_table_gen",
"dbgpt.agent.plugin.commands.built_in.display_type.show_text_gen",
]
command_dispaly_registry = CommandRegistry()
for command in command_dispaly_commands:
command_dispaly_registry.import_commands(command)
cfg.command_display = command_dispaly_commands
def _create_model_start_listener(system_app: SystemApp):
def startup_event(wh):

21
dbgpt/app/component_configs.py
View File

@@ -47,6 +47,8 @@ def initialize_components(
)
_initialize_model_cache(system_app)
_initialize_awel(system_app, param)
# Initialize resource manager of agent
_initialize_resource_manager(system_app)
_initialize_agent(system_app)
_initialize_openapi(system_app)
# Register serve apps
@@ -85,6 +87,25 @@ def _initialize_agent(system_app: SystemApp):
initialize_agent(system_app)
def _initialize_resource_manager(system_app: SystemApp):
from dbgpt.agent.expand.resources.dbgpt_tool import list_dbgpt_support_models
from dbgpt.agent.expand.resources.search_tool import baidu_search
from dbgpt.agent.resource.base import ResourceType
from dbgpt.agent.resource.manage import get_resource_manager, initialize_resource
from dbgpt.serve.agent.resource.datasource import DatasourceResource
from dbgpt.serve.agent.resource.knowledge import KnowledgeSpaceRetrieverResource
from dbgpt.serve.agent.resource.plugin import PluginToolPack
initialize_resource(system_app)
rm = get_resource_manager(system_app)
rm.register_resource(DatasourceResource)
rm.register_resource(KnowledgeSpaceRetrieverResource)
rm.register_resource(PluginToolPack, resource_type=ResourceType.Tool)
# Register a search tool
rm.register_resource(resource_instance=baidu_search)
rm.register_resource(resource_instance=list_dbgpt_support_models)
def _initialize_openapi(system_app: SystemApp):
from dbgpt.app.openapi.api_v1.editor.service import EditorService

15
dbgpt/app/dbgpt_server.py
View File

@@ -21,7 +21,12 @@ from dbgpt.app.base import (
# initialize_components import time cost about 0.1s
from dbgpt.app.component_configs import initialize_components
from dbgpt.component import SystemApp
from dbgpt.configs.model_config import EMBEDDING_MODEL_CONFIG, LLM_MODEL_CONFIG, LOGDIR
from dbgpt.configs.model_config import (
EMBEDDING_MODEL_CONFIG,
LLM_MODEL_CONFIG,
LOGDIR,
STATIC_MESSAGE_IMG_PATH,
)
from dbgpt.serve.core import add_exception_handler
from dbgpt.util.fastapi import create_app, replace_router
from dbgpt.util.i18n_utils import _, set_default_language
@@ -88,14 +93,10 @@ def mount_routers(app: FastAPI):
def mount_static_files(app: FastAPI):
from dbgpt.agent.plugin.commands.built_in.display_type import (
static_message_img_path,
)
os.makedirs(static_message_img_path, exist_ok=True)
os.makedirs(STATIC_MESSAGE_IMG_PATH, exist_ok=True)
app.mount(
"/images",
StaticFiles(directory=static_message_img_path, html=True),
StaticFiles(directory=STATIC_MESSAGE_IMG_PATH, html=True),
name="static2",
)
app.mount(

9
dbgpt/app/openapi/api_v1/api_v1.py
View File

@@ -80,13 +80,6 @@ def get_db_list():
return db_params
def plugins_select_info():
plugins_infos: dict = {}
for plugin in CFG.plugins:
plugins_infos.update({f"【{plugin._name}】=>{plugin._description}": plugin._name})
return plugins_infos
def get_db_list_info():
dbs = CFG.local_db_manager.get_db_list()
params: dict = {}
@@ -242,8 +235,6 @@ async def params_list(chat_mode: str = ChatScene.ChatNormal.value()):
return Result.succ(get_db_list())
elif ChatScene.ChatDashboard.value() == chat_mode:
return Result.succ(get_db_list())
elif ChatScene.ChatExecution.value() == chat_mode:
return Result.succ(plugins_select_info())
elif ChatScene.ChatKnowledge.value() == chat_mode:
return Result.succ(knowledge_list())
elif ChatScene.ChatKnowledge.ExtractRefineSummary.value() == chat_mode:

2
dbgpt/app/openapi/api_v1/editor/api_editor_v1.py
View File

@@ -179,7 +179,7 @@ async def editor_chart_run(run_param: dict = Body()):
end_time = time.time() * 1000
sql_run_data: SqlRunData = SqlRunData(
result_info="",
run_cost=(end_time - start_time) / 1000,
run_cost=int((end_time - start_time) / 1000),
colunms=colunms,
values=sql_result,
)

0
dbgpt/app/scene/chat_agent/__init__.py
View File

75
dbgpt/app/scene/chat_agent/chat.py
View File

@@ -1,75 +0,0 @@
import logging
from typing import Dict, List
from dbgpt._private.config import Config
from dbgpt.agent.plugin.commands.command_manage import ApiCall
from dbgpt.agent.plugin.generator import PluginPromptGenerator
from dbgpt.app.scene import BaseChat, ChatScene
from dbgpt.component import ComponentType
from dbgpt.serve.agent.hub.controller import ModulePlugin
from dbgpt.util.tracer import root_tracer, trace
CFG = Config()
logger = logging.getLogger("chat_agent")
class ChatAgent(BaseChat):
"""Chat With Agent through plugin"""
chat_scene: str = ChatScene.ChatAgent.value()
keep_end_rounds = 0
def __init__(self, chat_param: Dict):
"""Chat Agent Module Initialization
Args:
- chat_param: Dict
- chat_session_id: (str) chat session_id
- current_user_input: (str) current user input
- model_name:(str) llm model name
- select_param:(str) agent plugin
"""
if not chat_param["select_param"]:
raise ValueError("Please select a Plugin!")
self.select_plugins = chat_param["select_param"].split(",")
chat_param["chat_mode"] = ChatScene.ChatAgent
super().__init__(chat_param=chat_param)
self.plugins_prompt_generator: PluginPromptGenerator = PluginPromptGenerator()
self.plugins_prompt_generator.set_command_registry(CFG.command_registry)
# load select plugin
agent_module = CFG.SYSTEM_APP.get_component(
ComponentType.PLUGIN_HUB, ModulePlugin
)
self.plugins_prompt_generator = agent_module.load_select_plugin(
self.plugins_prompt_generator, self.select_plugins
)
self.api_call = ApiCall(plugin_generator=self.plugins_prompt_generator)
@trace()
async def generate_input_values(self) -> Dict[str, str]:
input_values = {
"user_goal": self.current_user_input,
"expand_constraints": self.__list_to_prompt_str(
list(self.plugins_prompt_generator.constraints)
),
"tool_list": self.plugins_prompt_generator.generate_commands_string(),
}
return input_values
def stream_plugin_call(self, text):
text = (
text.replace("\\n", " ")
.replace("\n", " ")
.replace("\_", "_")
.replace("\\", " ")
)
with root_tracer.start_span(
"ChatAgent.stream_plugin_call.api_call", metadata={"text": text}
):
return self.api_call.run(text)
def __list_to_prompt_str(self, list: List) -> str:
return "\n".join(f"{i + 1 + 1}. {item}" for i, item in enumerate(list))

23
dbgpt/app/scene/chat_agent/example.py
View File

@@ -1,23 +0,0 @@
from dbgpt.core._private.example_base import ExampleSelector
## Two examples are defined by default
EXAMPLES = [
{
"messages": [
{"type": "human", "data": {"content": "查询xxx", "example": True}},
{
"type": "ai",
"data": {
"content": """{
\"thoughts\": \"thought text\",
\"speak\": \"thoughts summary to say to user\",
\"command\": {\"name\": \"command name\", \"args\": {\"arg name\": \"value\"}},
}""",
"example": True,
},
},
]
},
]
plugin_example = ExampleSelector(examples_record=EXAMPLES, use_example=True)

20
dbgpt/app/scene/chat_agent/out_parser.py
View File

@@ -1,20 +0,0 @@
from typing import Dict, NamedTuple
from dbgpt.core.interface.output_parser import BaseOutputParser
class PluginAction(NamedTuple):
command: Dict
speak: str = ""
thoughts: str = ""
class PluginChatOutputParser(BaseOutputParser):
def parse_view_response(self, speak, data, prompt_response) -> str:
### tool out data to table view
print(f"parse_view_response:{speak},{str(data)}")
view_text = f"##### {speak}" + "\n" + str(data)
return view_text
def get_format_instructions(self) -> str:
pass

82
dbgpt/app/scene/chat_agent/prompt.py
View File

@@ -1,82 +0,0 @@
from dbgpt._private.config import Config
from dbgpt.app.scene import AppScenePromptTemplateAdapter, ChatScene
from dbgpt.app.scene.chat_execution.out_parser import PluginChatOutputParser
from dbgpt.core import ChatPromptTemplate, HumanPromptTemplate, SystemPromptTemplate
CFG = Config()
_PROMPT_SCENE_DEFINE_EN = "You are a universal AI assistant."
_DEFAULT_TEMPLATE_EN = """
You need to analyze the user goals and, under the given constraints, prioritize using one of the following tools to solve the user goals.
Tool list:
{tool_list}
Constraint:
1. After finding the available tools from the tool list given below, please output the following content to use the tool. Please make sure that the following content only appears once in the output result:
<api-call><name>Selected Tool name</name><args><arg1>value</arg1><arg2>value</arg2></args></api-call>
2. Please generate the above call text according to the definition of the corresponding tool in the tool list. The reference case is as follows:
Introduction to tool function: "Tool name", args: "Parameter 1": "<Parameter 1 value description>", "Parameter 2": "<Parameter 2 value description>" Corresponding call text: <api-call>< name>Tool name</name><args><parameter 1>value</parameter 1><parameter 2>value</parameter 2></args></api-call>
3. Generate the call of each tool according to the above constraints. The prompt text for tool use needs to be generated before the tool is used.
4. If the user goals cannot be understood and the intention is unclear, give priority to using search engine tools
5. Parameter content may need to be inferred based on the user's goals, not just extracted from text
6. Constraint conditions and tool information are used as auxiliary information for the reasoning process and should not be expressed in the output content to the user.
{expand_constraints}
User goals:
{user_goal}
"""
_PROMPT_SCENE_DEFINE_ZH = "你是一个通用AI助手！"
_DEFAULT_TEMPLATE_ZH = """
根据用户目标，请一步步思考，如何在满足下面约束条件的前提下，优先使用给出工具回答或者完成用户目标。
约束条件:
1.从下面给定工具列表找到可用的工具后，请输出以下内容用来使用工具, 注意要确保下面内容在输出结果中只出现一次:
<api-call><name>Selected Tool name</name><args><arg1>value</arg1><arg2>value</arg2></args></api-call>
2.请根据工具列表对应工具的定义来生成上述调用文本, 参考案例如下:
工具作用介绍: "工具名称", args: "参数1": "<参数1取值描述>","参数2": "<参数2取值描述>" 对应调用文本:<api-call><name>工具名称</name><args><参数1>value</参数1><参数2>value</参数2></args></api-call>
3.根据上面约束的方式生成每个工具的调用，对于工具使用的提示文本，需要在工具使用前生成
4.如果用户目标无法理解和意图不明确，优先使用搜索引擎工具
5.参数内容可能需要根据用户的目标推理得到，不仅仅是从文本提取
6.约束条件和工具信息作为推理过程的辅助信息，对应内容不要表达在给用户的输出内容中
7.不要把<api-call></api-call>部分内容放在markdown标签里
{expand_constraints}
工具列表:
{tool_list}
用户目标:
{user_goal}
"""
_DEFAULT_TEMPLATE = (
_DEFAULT_TEMPLATE_EN if CFG.LANGUAGE == "en" else _DEFAULT_TEMPLATE_ZH
)
_PROMPT_SCENE_DEFINE = (
_PROMPT_SCENE_DEFINE_EN if CFG.LANGUAGE == "en" else _PROMPT_SCENE_DEFINE_ZH
)
RESPONSE_FORMAT = None
### Whether the model service is streaming output
PROMPT_NEED_STREAM_OUT = True
prompt = ChatPromptTemplate(
messages=[
SystemPromptTemplate.from_template(_PROMPT_SCENE_DEFINE + _DEFAULT_TEMPLATE),
HumanPromptTemplate.from_template("{user_goal}"),
]
)
prompt_adapter = AppScenePromptTemplateAdapter(
prompt=prompt,
template_scene=ChatScene.ChatAgent.value(),
stream_out=PROMPT_NEED_STREAM_OUT,
output_parser=PluginChatOutputParser(is_stream_out=PROMPT_NEED_STREAM_OUT),
need_historical_messages=False,
temperature=1,
)
CFG.prompt_template_registry.register(prompt_adapter, is_default=True)

0
dbgpt/app/scene/chat_dashboard/business_cockpit/__init__.py
View File

5
dbgpt/app/scene/chat_dashboard/prompt.py
View File

@@ -29,14 +29,15 @@ Give the correct {dialect} analysis SQL
4.Carefully check the correctness of the SQL, the SQL must be correct, display method and summary of brief analysis thinking, and respond in the following json format:
{response}
The important thing is: Please make sure to only return the json string, do not add any other content (for direct processing by the program), and the json can be parsed by Python json.loads
5. Please use the same language as the "user"
"""
RESPONSE_FORMAT = [
{
"thoughts": "Current thinking and value of data analysis",
"showcase": "What type of charts to show",
"sql": "data analysis SQL",
"title": "Data Analysis Title",
"showcase": "What type of charts to show",
"thoughts": "Current thinking and value of data analysis",
}
]

4
dbgpt/app/scene/chat_data/chat_excel/excel_analyze/chat.py
View File

@@ -3,7 +3,7 @@ import os
from typing import Dict
from dbgpt._private.config import Config
from dbgpt.agent.plugin.commands.command_manage import ApiCall
from dbgpt.agent.util.api_call import ApiCall
from dbgpt.app.scene import BaseChat, ChatScene
from dbgpt.app.scene.chat_data.chat_excel.excel_learning.chat import ExcelLearning
from dbgpt.app.scene.chat_data.chat_excel.excel_reader import ExcelReader
@@ -45,7 +45,7 @@ class ChatExcel(BaseChat):
KNOWLEDGE_UPLOAD_ROOT_PATH, chat_mode.value(), self.select_param
)
)
self.api_call = ApiCall(display_registry=CFG.command_display)
self.api_call = ApiCall()
super().__init__(chat_param=chat_param)
@trace()

4
dbgpt/app/scene/chat_db/auto_execute/chat.py
View File

@@ -1,7 +1,7 @@
from typing import Dict
from dbgpt._private.config import Config
from dbgpt.agent.plugin.commands.command_manage import ApiCall
from dbgpt.agent.util.api_call import ApiCall
from dbgpt.app.scene import BaseChat, ChatScene
from dbgpt.util.executor_utils import blocking_func_to_async
from dbgpt.util.tracer import root_tracer, trace
@@ -40,7 +40,7 @@ class ChatWithDbAutoExecute(BaseChat):
self.database = CFG.local_db_manager.get_connector(self.db_name)
self.top_k: int = 50
self.api_call = ApiCall(display_registry=CFG.command_display)
self.api_call = ApiCall()
@trace()
async def generate_input_values(self) -> Dict:

0
dbgpt/app/scene/chat_execution/__init__.py
View File

83
dbgpt/app/scene/chat_execution/chat.py
View File

@@ -1,83 +0,0 @@
from typing import Dict, List
from dbgpt._private.config import Config
from dbgpt.agent.plugin.commands.command import execute_command
from dbgpt.agent.plugin.generator import PluginPromptGenerator
from dbgpt.app.scene import BaseChat, ChatScene
from dbgpt.util.tracer import trace
CFG = Config()
class ChatWithPlugin(BaseChat):
"""Chat With Plugin"""
chat_scene: str = ChatScene.ChatExecution.value()
plugins_prompt_generator: PluginPromptGenerator
select_plugin: str = None
def __init__(self, chat_param: Dict):
"""Chat Dashboard Module Initialization
Args:
- chat_param: Dict
- chat_session_id: (str) chat session_id
- current_user_input: (str) current user input
- model_name:(str) llm model name
- select_param:(str) plugin selector
"""
self.plugin_selector = chat_param["select_param"]
chat_param["chat_mode"] = ChatScene.ChatExecution
super().__init__(chat_param=chat_param)
self.plugins_prompt_generator = PluginPromptGenerator()
self.plugins_prompt_generator.set_command_registry(CFG.command_registry)
# 加载插件中可用命令
self.select_plugin = self.plugin_selector
if self.select_plugin:
for plugin in CFG.plugins:
if plugin._name == self.plugin_selector:
if not plugin.can_handle_post_prompt():
continue
self.plugins_prompt_generator = plugin.post_prompt(
self.plugins_prompt_generator
)
else:
for plugin in CFG.plugins:
if not plugin.can_handle_post_prompt():
continue
self.plugins_prompt_generator = plugin.post_prompt(
self.plugins_prompt_generator
)
@trace()
async def generate_input_values(self) -> Dict:
input_values = {
"input": self.current_user_input,
"constraints": self.__list_to_prompt_str(
list(self.plugins_prompt_generator.constraints)
),
"commands_infos": self.plugins_prompt_generator.generate_commands_string(),
}
return input_values
def do_action(self, prompt_response):
print(f"do_action:{prompt_response}")
## plugin command run
return execute_command(
str(prompt_response.command.get("name")),
prompt_response.command.get("args", {}),
self.plugins_prompt_generator,
)
def chat_show(self):
super().chat_show()
def __list_to_prompt_str(self, list: List) -> str:
return "\n".join(f"{i + 1 + 1}. {item}" for i, item in enumerate(list))
def generate(self, p) -> str:
return super().generate(p)
@property
def chat_type(self) -> str:
return ChatScene.ChatExecution.value

23
dbgpt/app/scene/chat_execution/example.py
View File

@@ -1,23 +0,0 @@
from dbgpt.core._private.example_base import ExampleSelector
## Two examples are defined by default
EXAMPLES = [
{
"messages": [
{"type": "human", "data": {"content": "查询xxx", "example": True}},
{
"type": "ai",
"data": {
"content": """{
\"thoughts\": \"thought text\",
\"speak\": \"thoughts summary to say to user\",
\"command\": {\"name\": \"command name\", \"args\": {\"arg name\": \"value\"}},
}""",
"example": True,
},
},
]
},
]
plugin_example = ExampleSelector(examples_record=EXAMPLES, use_example=True)

45
dbgpt/app/scene/chat_execution/out_parser.py
View File

@@ -1,45 +0,0 @@
import json
import logging
from typing import Dict, NamedTuple
from dbgpt.core.interface.output_parser import BaseOutputParser, T
logger = logging.getLogger(__name__)
class PluginAction(NamedTuple):
command: Dict
speak: str = ""
thoughts: str = ""
class PluginChatOutputParser(BaseOutputParser):
def parse_prompt_response(self, model_out_text) -> T:
clean_json_str = super().parse_prompt_response(model_out_text)
print(clean_json_str)
if not clean_json_str:
raise ValueError("model server response not have json!")
try:
response = json.loads(clean_json_str)
except Exception as e:
raise ValueError("model server out not fllow the prompt!")
speak = ""
thoughts = ""
for key in sorted(response):
if key.strip() == "command":
command = response[key]
if key.strip() == "thoughts":
thoughts = response[key]
if key.strip() == "speak":
speak = response[key]
return PluginAction(command, speak, thoughts)
def parse_view_response(self, speak, data, prompt_response) -> str:
### tool out data to table view
print(f"parse_view_response:{speak},{str(data)}")
view_text = f"##### {speak}" + "\n" + str(data)
return view_text
def get_format_instructions(self) -> str:
pass

61
dbgpt/app/scene/chat_execution/prompt.py
View File

@@ -1,61 +0,0 @@
import json
from dbgpt._private.config import Config
from dbgpt.app.scene import AppScenePromptTemplateAdapter, ChatScene
from dbgpt.app.scene.chat_execution.out_parser import PluginChatOutputParser
from dbgpt.core import (
ChatPromptTemplate,
HumanPromptTemplate,
MessagesPlaceholder,
SystemPromptTemplate,
)
CFG = Config()
PROMPT_SCENE_DEFINE = "You are an AI designed to solve the user's goals with given commands, please follow the constraints of the system's input for your answers."
_DEFAULT_TEMPLATE = """
Goals:
{input}
Constraints:
0.Exclusively use the commands listed in double quotes e.g. "command name"
{constraints}
Commands:
{commands_infos}
Please response strictly according to the following json format:
{response}
Ensure the response is correct json and can be parsed by Python json.loads
"""
RESPONSE_FORMAT = {
"thoughts": "thought text",
"speak": "thoughts summary to say to user",
"command": {"name": "command name", "args": {"arg name": "value"}},
}
### Whether the model service is streaming output
PROMPT_NEED_STREAM_OUT = False
prompt = ChatPromptTemplate(
messages=[
SystemPromptTemplate.from_template(
PROMPT_SCENE_DEFINE + _DEFAULT_TEMPLATE,
response_format=json.dumps(RESPONSE_FORMAT, indent=4),
),
MessagesPlaceholder(variable_name="chat_history"),
HumanPromptTemplate.from_template("{input}"),
]
)
prompt_adapter = AppScenePromptTemplateAdapter(
prompt=prompt,
template_scene=ChatScene.ChatExecution.value(),
stream_out=PROMPT_NEED_STREAM_OUT,
output_parser=PluginChatOutputParser(is_stream_out=PROMPT_NEED_STREAM_OUT),
need_historical_messages=False,
)
CFG.prompt_template_registry.register(prompt_adapter, is_default=True)

4
dbgpt/app/scene/chat_factory.py
View File

@@ -8,8 +8,6 @@ class ChatFactory(metaclass=Singleton):
@staticmethod
def get_implementation(chat_mode, **kwargs):
# Lazy loading
from dbgpt.app.scene.chat_agent.chat import ChatAgent
from dbgpt.app.scene.chat_agent.prompt import prompt
from dbgpt.app.scene.chat_dashboard.chat import ChatDashboard
from dbgpt.app.scene.chat_dashboard.prompt import prompt
from dbgpt.app.scene.chat_data.chat_excel.excel_analyze.chat import ChatExcel
@@ -19,8 +17,6 @@ class ChatFactory(metaclass=Singleton):
from dbgpt.app.scene.chat_db.auto_execute.prompt import prompt
from dbgpt.app.scene.chat_db.professional_qa.chat import ChatWithDbQA
from dbgpt.app.scene.chat_db.professional_qa.prompt import prompt
from dbgpt.app.scene.chat_execution.chat import ChatWithPlugin
from dbgpt.app.scene.chat_execution.prompt import prompt
from dbgpt.app.scene.chat_knowledge.extract_entity.chat import ExtractEntity
from dbgpt.app.scene.chat_knowledge.extract_entity.prompt import prompt
from dbgpt.app.scene.chat_knowledge.extract_triplet.chat import ExtractTriplet