diff --git a/pilot/connections/rdbms/base.py b/pilot/connections/rdbms/base.py index b590294ac..43b6a3993 100644 --- a/pilot/connections/rdbms/base.py +++ b/pilot/connections/rdbms/base.py @@ -52,7 +52,18 @@ class RDBMSDatabase(BaseConnect): custom_table_info: Optional[dict] = None, view_support: bool = False, ): - """Create engine from database URI.""" + """Create engine from database URI. + Args: + - engine: Engine sqlalchemy.engine + - schema: Optional[str]. + - metadata: Optional[MetaData] + - ignore_tables: Optional[List[str]] + - include_tables: Optional[List[str]] + - sample_rows_in_table_info: int default:3, + - indexes_in_table_info: bool = False, + - custom_table_info: Optional[dict] = None, + - view_support: bool = False, + """ self._engine = engine self._schema = schema if include_tables and ignore_tables: @@ -92,6 +103,15 @@ class RDBMSDatabase(BaseConnect): engine_args: Optional[dict] = None, **kwargs: Any, ) -> RDBMSDatabase: + """Construct a SQLAlchemy engine from uri database. + Args: + host (str): database host. + port (int): database port. + user (str): database user. + pwd (str): database password. + db_name (str): database name. + engine_args (Optional[dict]):other engine_args. + """ db_url: str = ( cls.driver + "://" diff --git a/pilot/scene/base_chat.py b/pilot/scene/base_chat.py index 243760090..5d685573e 100644 --- a/pilot/scene/base_chat.py +++ b/pilot/scene/base_chat.py @@ -21,6 +21,12 @@ CFG = Config() class BaseChat(ABC): + """DB-GPT Chat Service Base Module + Include: + stream_call():scene + prompt -> stream response + nostream_call():scene + prompt -> nostream response + """ + chat_scene: str = None llm_model: Any = None # By default, keep the last two rounds of conversation records as the context @@ -32,6 +38,14 @@ class BaseChat(ABC): arbitrary_types_allowed = True def __init__(self, chat_param: Dict): + """Chat 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) select param + """ self.chat_session_id = chat_param["chat_session_id"] self.chat_mode = chat_param["chat_mode"] self.current_user_input: str = chat_param["current_user_input"] diff --git a/pilot/scene/chat_agent/chat.py b/pilot/scene/chat_agent/chat.py index 4734c8106..ef37f8448 100644 --- a/pilot/scene/chat_agent/chat.py +++ b/pilot/scene/chat_agent/chat.py @@ -18,10 +18,20 @@ logger = logging.getLogger("chat_agent") class ChatAgent(BaseChat): + """Chat With Agent through plugin""" + chat_scene: str = ChatScene.ChatAgent.value() chat_retention_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(",") diff --git a/pilot/scene/chat_dashboard/chat.py b/pilot/scene/chat_dashboard/chat.py index ba5bdd8ce..7e4433670 100644 --- a/pilot/scene/chat_dashboard/chat.py +++ b/pilot/scene/chat_dashboard/chat.py @@ -19,10 +19,17 @@ CFG = Config() class ChatDashboard(BaseChat): chat_scene: str = ChatScene.ChatDashboard.value() report_name: str - """Number of results to return from the query""" + """Chat Dashboard to generate dashboard chart""" 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) dbname + """ self.db_name = chat_param["select_param"] chat_param["chat_mode"] = ChatScene.ChatDashboard super().__init__(chat_param=chat_param) diff --git a/pilot/scene/chat_data/chat_excel/excel_analyze/chat.py b/pilot/scene/chat_data/chat_excel/excel_analyze/chat.py index 9599c1402..343afc4bd 100644 --- a/pilot/scene/chat_data/chat_excel/excel_analyze/chat.py +++ b/pilot/scene/chat_data/chat_excel/excel_analyze/chat.py @@ -19,10 +19,20 @@ CFG = Config() class ChatExcel(BaseChat): + """a Excel analyzer to analyze Excel Data""" + chat_scene: str = ChatScene.ChatExcel.value() chat_retention_rounds = 1 def __init__(self, chat_param: Dict): + """Chat Excel 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) file path + """ chat_mode = ChatScene.ChatExcel self.select_param = chat_param["select_param"] diff --git a/pilot/scene/chat_db/auto_execute/chat.py b/pilot/scene/chat_db/auto_execute/chat.py index e8e1e2ebd..f92df7a3a 100644 --- a/pilot/scene/chat_db/auto_execute/chat.py +++ b/pilot/scene/chat_db/auto_execute/chat.py @@ -15,6 +15,14 @@ class ChatWithDbAutoExecute(BaseChat): """Number of results to return from the query""" def __init__(self, chat_param: Dict): + """Chat Data 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) dbname + """ chat_mode = ChatScene.ChatWithDbExecute self.db_name = chat_param["select_param"] chat_param["chat_mode"] = chat_mode @@ -31,6 +39,9 @@ class ChatWithDbAutoExecute(BaseChat): self.top_k: int = 200 def generate_input_values(self): + """ + generate input values + """ try: from pilot.summary.db_summary_client import DBSummaryClient except ImportError: diff --git a/pilot/scene/chat_db/professional_qa/chat.py b/pilot/scene/chat_db/professional_qa/chat.py index 39f4052a6..abdfd9f00 100644 --- a/pilot/scene/chat_db/professional_qa/chat.py +++ b/pilot/scene/chat_db/professional_qa/chat.py @@ -12,10 +12,17 @@ CFG = Config() class ChatWithDbQA(BaseChat): chat_scene: str = ChatScene.ChatWithDbQA.value() - """Number of results to return from the query""" + """As a DBA, Chat DB Module, chat with combine DB meta schema """ def __init__(self, chat_param: Dict): - """ """ + """Chat DB 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) dbname + """ self.db_name = chat_param["select_param"] chat_param["chat_mode"] = ChatScene.ChatWithDbQA super().__init__(chat_param=chat_param) diff --git a/pilot/scene/chat_execution/chat.py b/pilot/scene/chat_execution/chat.py index e4c5175a6..c6d7bbe2f 100644 --- a/pilot/scene/chat_execution/chat.py +++ b/pilot/scene/chat_execution/chat.py @@ -11,11 +11,21 @@ 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) diff --git a/pilot/scene/chat_knowledge/v1/chat.py b/pilot/scene/chat_knowledge/v1/chat.py index 672669a1b..9681f13c6 100644 --- a/pilot/scene/chat_knowledge/v1/chat.py +++ b/pilot/scene/chat_knowledge/v1/chat.py @@ -19,10 +19,17 @@ CFG = Config() class ChatKnowledge(BaseChat): chat_scene: str = ChatScene.ChatKnowledge.value() - """Number of results to return from the query""" + """KBQA Chat Module""" def __init__(self, chat_param: Dict): - """ """ + """Chat Knowledge 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) space name + """ from pilot.embedding_engine.embedding_engine import EmbeddingEngine from pilot.embedding_engine.embedding_factory import EmbeddingFactory