mirror of
https://github.com/csunny/DB-GPT.git
synced 2026-07-28 01:36:01 +00:00
64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
from typing import Dict
|
|
|
|
from pilot.scene.base_chat import BaseChat
|
|
from pilot.scene.base import ChatScene
|
|
from pilot.common.sql_database import Database
|
|
from pilot.configs.config import Config
|
|
from pilot.scene.chat_db.auto_execute.prompt import prompt
|
|
|
|
CFG = Config()
|
|
|
|
|
|
class ChatWithDbAutoExecute(BaseChat):
|
|
chat_scene: str = ChatScene.ChatWithDbExecute.value()
|
|
|
|
"""Number of results to return from the query"""
|
|
|
|
def __init__(self, chat_param: Dict):
|
|
chat_mode = ChatScene.ChatWithDbExecute
|
|
self.db_name = chat_param["select_param"]
|
|
chat_param["chat_mode"] = chat_mode
|
|
""" """
|
|
super().__init__(
|
|
chat_param=chat_param,
|
|
)
|
|
if not self.db_name:
|
|
raise ValueError(
|
|
f"{ChatScene.ChatWithDbExecute.value} mode should chose db!"
|
|
)
|
|
|
|
self.database = CFG.LOCAL_DB_MANAGE.get_connect(self.db_name)
|
|
self.top_k: int = 200
|
|
|
|
def generate_input_values(self):
|
|
try:
|
|
from pilot.summary.db_summary_client import DBSummaryClient
|
|
except ImportError:
|
|
raise ValueError("Could not import DBSummaryClient. ")
|
|
client = DBSummaryClient(system_app=CFG.SYSTEM_APP)
|
|
try:
|
|
table_infos = client.get_db_summary(
|
|
dbname=self.db_name,
|
|
query=self.current_user_input,
|
|
topk=CFG.KNOWLEDGE_SEARCH_TOP_SIZE,
|
|
)
|
|
except Exception as e:
|
|
print("db summary find error!" + str(e))
|
|
table_infos = self.database.table_simple_info()
|
|
if not table_infos:
|
|
table_infos = self.database.table_simple_info()
|
|
|
|
# table_infos = self.database.table_simple_info()
|
|
|
|
input_values = {
|
|
"input": self.current_user_input,
|
|
"top_k": str(self.top_k),
|
|
"dialect": self.database.dialect,
|
|
"table_info": table_infos,
|
|
}
|
|
return input_values
|
|
|
|
def do_action(self, prompt_response):
|
|
print(f"do_action:{prompt_response}")
|
|
return self.database.run(prompt_response.sql)
|