diff --git a/pilot/agent/json_fix_llm.py b/pilot/agent/json_fix_llm.py index 327881a78..075634784 100644 --- a/pilot/agent/json_fix_llm.py +++ b/pilot/agent/json_fix_llm.py @@ -55,54 +55,6 @@ def fix_and_parse_json( logger.error("参数解析错误", e) -def fix_json_using_multiple_techniques(assistant_reply: str) -> Dict[Any, Any]: - """Fix the given JSON string to make it parseable and fully compliant with two techniques. - - Args: - json_string (str): The JSON string to fix. - - Returns: - str: The fixed JSON string. - """ - assistant_reply = assistant_reply.strip() - if assistant_reply.startswith("```json"): - assistant_reply = assistant_reply[7:] - if assistant_reply.endswith("```"): - assistant_reply = assistant_reply[:-3] - try: - return json.loads(assistant_reply) # just check the validity - except json.JSONDecodeError as e: # noqa: E722 - print(f"JSONDecodeError: {e}") - pass - - if assistant_reply.startswith("json "): - assistant_reply = assistant_reply[5:] - assistant_reply = assistant_reply.strip() - try: - return json.loads(assistant_reply) # just check the validity - except json.JSONDecodeError: # noqa: E722 - pass - - # Parse and print Assistant response - assistant_reply_json = fix_and_parse_json(assistant_reply) - logger.debug("Assistant reply JSON: %s", str(assistant_reply_json)) - if assistant_reply_json == {}: - assistant_reply_json = attempt_to_fix_json_by_finding_outermost_brackets( - assistant_reply - ) - - logger.debug("Assistant reply JSON 2: %s", str(assistant_reply_json)) - if assistant_reply_json != {}: - return assistant_reply_json - - logger.error( - "Error: The following AI output couldn't be converted to a JSON:\n", - assistant_reply, - ) - if CFG.speak_mode: - say_text("I have received an invalid JSON response from the OpenAI API.") - - return {} def correct_json(json_to_load: str) -> str: diff --git a/pilot/commands/command.py b/pilot/commands/command.py index 0200ef6cd..8838efff1 100644 --- a/pilot/commands/command.py +++ b/pilot/commands/command.py @@ -4,10 +4,9 @@ import json from typing import Dict -from pilot.agent.json_fix_llm import fix_json_using_multiple_techniques from pilot.commands.exception_not_commands import NotCommands from pilot.configs.config import Config -from pilot.prompts.generator import PromptGenerator +from pilot.prompts.generator import PluginPromptGenerator from pilot.speech import say_text @@ -24,8 +23,8 @@ def _resolve_pathlike_command_args(command_args): def execute_ai_response_json( - prompt: PromptGenerator, - ai_response: str, + prompt: PluginPromptGenerator, + ai_response, user_input: str = None, ) -> str: """ @@ -39,11 +38,8 @@ def execute_ai_response_json( """ cfg = Config() - try: - assistant_reply_json = fix_json_using_multiple_techniques(ai_response) - except (json.JSONDecodeError, ValueError, AttributeError) as e: - raise NotCommands("非可执行命令结构") - command_name, arguments = get_command(assistant_reply_json) + + command_name, arguments = get_command(ai_response) if cfg.speak_mode: say_text(f"I want to execute {command_name}") @@ -71,7 +67,7 @@ def execute_ai_response_json( def execute_command( command_name: str, arguments, - prompt: PromptGenerator, + prompt: PluginPromptGenerator, ): """Execute the command and return the result diff --git a/pilot/common/markdown_text.py b/pilot/common/markdown_text.py index 1d90ba645..1244160fd 100644 --- a/pilot/common/markdown_text.py +++ b/pilot/common/markdown_text.py @@ -1,21 +1,21 @@ -import markdown2 +import markdown2 import pandas as pd + def datas_to_table_html(data): df = pd.DataFrame(data[1:], columns=data[0]) table_style = """""" - html_table = df.to_html(index=False, escape=False) + html_table = df.to_html(index=False, escape=False) html = f"{table_style}{html_table}" return html.replace("\n", " ") - def generate_markdown_table(data): - """\n 生成 Markdown 表格\n data: 一个包含表头和表格内容的二维列表\n """ + """\n 生成 Markdown 表格\n data: 一个包含表头和表格内容的二维列表\n""" # 获取表格列数 num_cols = len(data[0]) # 生成表头 @@ -41,6 +41,7 @@ def generate_markdown_table(data): return table + def generate_htm_table(data): markdown_text = generate_markdown_table(data) html_table = markdown2.markdown(markdown_text, extras=["tables"]) @@ -53,4 +54,4 @@ if __name__ == "__main__": table_style = """""" - print(table_style.replace("\n", " ")) \ No newline at end of file + print(table_style.replace("\n", " ")) diff --git a/pilot/plugins.py b/pilot/common/plugins.py similarity index 100% rename from pilot/plugins.py rename to pilot/common/plugins.py diff --git a/pilot/common/schema.py b/pilot/common/schema.py index f66bba1a6..cd462966c 100644 --- a/pilot/common/schema.py +++ b/pilot/common/schema.py @@ -1,8 +1,9 @@ from enum import auto, Enum from typing import List, Any + class SeparatorStyle(Enum): - SINGLE ="###" + SINGLE = "###" TWO = "" THREE = auto() FOUR = auto() diff --git a/pilot/common/sql_database.py b/pilot/common/sql_database.py index 2c16869d5..2b8d6fe4b 100644 --- a/pilot/common/sql_database.py +++ b/pilot/common/sql_database.py @@ -30,16 +30,16 @@ class Database: """SQLAlchemy wrapper around a database.""" def __init__( - self, - engine, - schema: Optional[str] = None, - metadata: Optional[MetaData] = None, - ignore_tables: Optional[List[str]] = None, - include_tables: Optional[List[str]] = None, - sample_rows_in_table_info: int = 3, - indexes_in_table_info: bool = False, - custom_table_info: Optional[dict] = None, - view_support: bool = False, + self, + engine, + schema: Optional[str] = None, + metadata: Optional[MetaData] = None, + ignore_tables: Optional[List[str]] = None, + include_tables: Optional[List[str]] = None, + sample_rows_in_table_info: int = 3, + indexes_in_table_info: bool = False, + custom_table_info: Optional[dict] = None, + view_support: bool = False, ): """Create engine from database URI.""" self._engine = engine @@ -119,7 +119,7 @@ class Database: @classmethod def from_uri( - cls, database_uri: str, engine_args: Optional[dict] = None, **kwargs: Any + cls, database_uri: str, engine_args: Optional[dict] = None, **kwargs: Any ) -> Database: """Construct a SQLAlchemy engine from URI.""" _engine_args = engine_args or {} @@ -148,7 +148,7 @@ class Database: self._metadata = MetaData() # sql = f"use {db_name}" - sql = text(f'use `{db_name}`') + sql = text(f"use `{db_name}`") session.execute(sql) # 处理表信息数据 @@ -159,13 +159,17 @@ class Database: # tables list if view_support is True self._all_tables = set( self._inspector.get_table_names(schema=db_name) - + (self._inspector.get_view_names(schema=db_name) if self.view_support else []) + + ( + self._inspector.get_view_names(schema=db_name) + if self.view_support + else [] + ) ) return session def get_current_db_name(self, session) -> str: - return session.execute(text('SELECT DATABASE()')).scalar() + return session.execute(text("SELECT DATABASE()")).scalar() def table_simple_info(self, session): _sql = f""" @@ -201,7 +205,7 @@ class Database: tbl for tbl in self._metadata.sorted_tables if tbl.name in set(all_table_names) - and not (self.dialect == "sqlite" and tbl.name.startswith("sqlite_")) + and not (self.dialect == "sqlite" and tbl.name.startswith("sqlite_")) ] tables = [] @@ -214,7 +218,7 @@ class Database: create_table = str(CreateTable(table).compile(self._engine)) table_info = f"{create_table.rstrip()}" has_extra_info = ( - self._indexes_in_table_info or self._sample_rows_in_table_info + self._indexes_in_table_info or self._sample_rows_in_table_info ) if has_extra_info: table_info += "\n\n/*" @@ -303,6 +307,10 @@ class Database: def get_database_list(self): session = self._db_sessions() - cursor = session.execute(text(' show databases;')) + cursor = session.execute(text(" show databases;")) results = cursor.fetchall() - return [d[0] for d in results if d[0] not in ["information_schema", "performance_schema", "sys", "mysql"]] + return [ + d[0] + for d in results + if d[0] not in ["information_schema", "performance_schema", "sys", "mysql"] + ] diff --git a/pilot/configs/ai_config.py b/pilot/configs/ai_config.py deleted file mode 100644 index ed9b4e2f8..000000000 --- a/pilot/configs/ai_config.py +++ /dev/null @@ -1,167 +0,0 @@ -# sourcery skip: do-not-use-staticmethod -""" -A module that contains the AIConfig class object that contains the configuration -""" -from __future__ import annotations - -import os -import platform -from pathlib import Path -from typing import Optional - -import distro -import yaml - -from pilot.configs.config import Config -from pilot.prompts.generator import PromptGenerator -from pilot.prompts.prompt import build_default_prompt_generator - -# Soon this will go in a folder where it remembers more stuff about the run(s) -SAVE_FILE = str(Path(os.getcwd()) / "ai_settings.yaml") - - -class AIConfig: - """ - A class object that contains the configuration information for the AI - - Attributes: - ai_name (str): The name of the AI. - ai_role (str): The description of the AI's role. - ai_goals (list): The list of objectives the AI is supposed to complete. - api_budget (float): The maximum dollar value for API calls (0.0 means infinite) - """ - - def __init__( - self, - ai_name: str = "", - ai_role: str = "", - ai_goals: list | None = None, - api_budget: float = 0.0, - ) -> None: - """ - Initialize a class instance - - Parameters: - ai_name (str): The name of the AI. - ai_role (str): The description of the AI's role. - ai_goals (list): The list of objectives the AI is supposed to complete. - api_budget (float): The maximum dollar value for API calls (0.0 means infinite) - Returns: - None - """ - if ai_goals is None: - ai_goals = [] - self.ai_name = ai_name - self.ai_role = ai_role - self.ai_goals = ai_goals - self.api_budget = api_budget - self.prompt_generator = None - self.command_registry = None - - @staticmethod - def load(config_file: str = SAVE_FILE) -> "AIConfig": - """ - Returns class object with parameters (ai_name, ai_role, ai_goals, api_budget) loaded from - yaml file if yaml file exists, - else returns class with no parameters. - - Parameters: - config_file (int): The path to the config yaml file. - DEFAULT: "../ai_settings.yaml" - - Returns: - cls (object): An instance of given cls object - """ - - try: - with open(config_file, encoding="utf-8") as file: - config_params = yaml.load(file, Loader=yaml.FullLoader) - except FileNotFoundError: - config_params = {} - - ai_name = config_params.get("ai_name", "") - ai_role = config_params.get("ai_role", "") - ai_goals = [ - str(goal).strip("{}").replace("'", "").replace('"', "") - if isinstance(goal, dict) - else str(goal) - for goal in config_params.get("ai_goals", []) - ] - api_budget = config_params.get("api_budget", 0.0) - # type is Type[AIConfig] - return AIConfig(ai_name, ai_role, ai_goals, api_budget) - - def save(self, config_file: str = SAVE_FILE) -> None: - """ - Saves the class parameters to the specified file yaml file path as a yaml file. - - Parameters: - config_file(str): The path to the config yaml file. - DEFAULT: "../ai_settings.yaml" - - Returns: - None - """ - - config = { - "ai_name": self.ai_name, - "ai_role": self.ai_role, - "ai_goals": self.ai_goals, - "api_budget": self.api_budget, - } - with open(config_file, "w", encoding="utf-8") as file: - yaml.dump(config, file, allow_unicode=True) - - def construct_full_prompt( - self, prompt_generator: Optional[PromptGenerator] = None - ) -> str: - """ - Returns a prompt to the user with the class information in an organized fashion. - - Parameters: - None - - Returns: - full_prompt (str): A string containing the initial prompt for the user - including the ai_name, ai_role, ai_goals, and api_budget. - """ - - prompt_start = ( - "Your decisions must always be made independently without" - " seeking user assistance. Play to your strengths as an LLM and pursue" - " simple strategies with no legal complications." - "" - ) - - cfg = Config() - if prompt_generator is None: - prompt_generator = build_default_prompt_generator() - prompt_generator.goals = self.ai_goals - prompt_generator.name = self.ai_name - prompt_generator.role = self.ai_role - prompt_generator.command_registry = self.command_registry - for plugin in cfg.plugins: - if not plugin.can_handle_post_prompt(): - continue - prompt_generator = plugin.post_prompt(prompt_generator) - - if cfg.execute_local_commands: - # add OS info to prompt - os_name = platform.system() - os_info = ( - platform.platform(terse=True) - if os_name != "Linux" - else distro.name(pretty=True) - ) - - prompt_start += f"\nThe OS you are running on is: {os_info}" - - # Construct full prompt - full_prompt = f"You are {prompt_generator.name}, {prompt_generator.role}\n{prompt_start}\n\nGOALS:\n\n" - for i, goal in enumerate(self.ai_goals): - full_prompt += f"{i+1}. {goal}\n" - if self.api_budget > 0.0: - full_prompt += f"\nIt takes money to let you run. Your API budget is ${self.api_budget:.3f}" - self.prompt_generator = prompt_generator - full_prompt += f"\n\n{prompt_generator.generate_prompt_string()}" - return full_prompt diff --git a/pilot/configs/config.py b/pilot/configs/config.py index 88cbf5117..518275f34 100644 --- a/pilot/configs/config.py +++ b/pilot/configs/config.py @@ -47,7 +47,6 @@ class Config(metaclass=Singleton): self.milvus_collection = os.getenv("MILVUS_COLLECTION", "dbgpt") self.milvus_secure = os.getenv("MILVUS_SECURE") == "True" - self.authorise_key = os.getenv("AUTHORISE_COMMAND_KEY", "y") self.exit_key = os.getenv("EXIT_KEY", "n") self.image_provider = os.getenv("IMAGE_PROVIDER", True) @@ -104,8 +103,17 @@ class Config(metaclass=Singleton): self.LOCAL_DB_PASSWORD = os.getenv("LOCAL_DB_PASSWORD", "aa123456") ### TODO Adapt to multiple types of libraries - self.local_db = Database.from_uri("mysql+pymysql://" + self.LOCAL_DB_USER +":"+ self.LOCAL_DB_PASSWORD +"@" +self.LOCAL_DB_HOST + ":" + str(self.LOCAL_DB_PORT) , - engine_args ={"pool_size": 10, "pool_recycle": 3600, "echo": True}) + self.local_db = Database.from_uri( + "mysql+pymysql://" + + self.LOCAL_DB_USER + + ":" + + self.LOCAL_DB_PASSWORD + + "@" + + self.LOCAL_DB_HOST + + ":" + + str(self.LOCAL_DB_PORT), + engine_args={"pool_size": 10, "pool_recycle": 3600, "echo": True}, + ) ### LLM Model Service Configuration self.LLM_MODEL = os.getenv("LLM_MODEL", "vicuna-13b") diff --git a/pilot/connections/base.py b/pilot/connections/base.py index ec41f9273..3905f410d 100644 --- a/pilot/connections/base.py +++ b/pilot/connections/base.py @@ -2,7 +2,34 @@ # -*- coding:utf-8 -*- """We need to design a base class. That other connector can Write with this""" +from abc import ABC, abstractmethod +from pydantic import BaseModel, Extra, Field, root_validator +from typing import Any, Iterable, List, Optional -class BaseConnection: - pass +class BaseConnect(BaseModel, ABC): + type + driver: str + + + def get_session(self, db_name: str): + pass + + + def get_table_names(self) -> Iterable[str]: + pass + + def get_table_info(self, table_names: Optional[List[str]] = None) -> str: + pass + + def get_table_info(self, table_names: Optional[List[str]] = None) -> str: + pass + + def get_index_info(self, table_names: Optional[List[str]] = None) -> str: + pass + + def get_database_list(self): + pass + + def run(self, session, command: str, fetch: str = "all") -> List: + pass \ No newline at end of file diff --git a/pilot/connections/mysql.py b/pilot/connections/mysql.py deleted file mode 100644 index 2f5a1e152..000000000 --- a/pilot/connections/mysql.py +++ /dev/null @@ -1,64 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- - -import pymysql - -class MySQLOperator: - """Connect MySQL Database fetch MetaData For LLM Prompt - Args: - - Usage: - """ - - default_db = ["information_schema", "performance_schema", "sys", "mysql"] - - def __init__(self, user, password, host="localhost", port=3306) -> None: - self.conn = pymysql.connect( - host=host, - user=user, - port=port, - passwd=password, - charset="utf8mb4", - cursorclass=pymysql.cursors.DictCursor, - ) - - def get_schema(self, schema_name): - with self.conn.cursor() as cursor: - _sql = f""" - select concat(table_name, "(" , group_concat(column_name), ")") as schema_info from information_schema.COLUMNS where table_schema="{schema_name}" group by TABLE_NAME; - """ - cursor.execute(_sql) - results = cursor.fetchall() - return results - - - def run_sql(self, db_name:str, sql:str, fetch: str = "all"): - with self.conn.cursor() as cursor: - cursor.execute("USE " + db_name) - cursor.execute(sql) - if fetch == "all": - result = cursor.fetchall() - elif fetch == "one": - result = cursor.fetchone()[0] # type: ignore - else: - raise ValueError("Fetch parameter must be either 'one' or 'all'") - return str(result) - - def get_index(self, schema_name): - pass - - def get_db_list(self): - with self.conn.cursor() as cursor: - _sql = """ - show databases; - """ - cursor.execute(_sql) - results = cursor.fetchall() - - dbs = [ - d["Database"] for d in results if d["Database"] not in self.default_db - ] - return dbs - - def get_meta(self, schema_name): - pass diff --git a/pilot/prompts/generator_new.py b/pilot/connections/nosql/__init__.py similarity index 100% rename from pilot/prompts/generator_new.py rename to pilot/connections/nosql/__init__.py diff --git a/pilot/connections/rdbms/__init__.py b/pilot/connections/rdbms/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pilot/connections/clickhouse.py b/pilot/connections/rdbms/clickhouse.py similarity index 100% rename from pilot/connections/clickhouse.py rename to pilot/connections/rdbms/clickhouse.py diff --git a/pilot/connections/es.py b/pilot/connections/rdbms/es.py similarity index 100% rename from pilot/connections/es.py rename to pilot/connections/rdbms/es.py diff --git a/pilot/connections/mongo.py b/pilot/connections/rdbms/mongo.py similarity index 100% rename from pilot/connections/mongo.py rename to pilot/connections/rdbms/mongo.py diff --git a/pilot/connections/rdbms/mysql.py b/pilot/connections/rdbms/mysql.py new file mode 100644 index 000000000..9d99f3e9b --- /dev/null +++ b/pilot/connections/rdbms/mysql.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +import pymysql +from pilot.connections.rdbms.rdbms_connect import RDBMSDatabase + + +class MySQLConnect(RDBMSDatabase): + """Connect MySQL Database fetch MetaData For LLM Prompt + Args: + Usage: + """ + + type:str = "MySQL" + connect_url = "mysql+pymysql://" + + default_db = ["information_schema", "performance_schema", "sys", "mysql"] + diff --git a/pilot/connections/oracle.py b/pilot/connections/rdbms/oracle.py similarity index 100% rename from pilot/connections/oracle.py rename to pilot/connections/rdbms/oracle.py diff --git a/pilot/connections/postgres.py b/pilot/connections/rdbms/postgres.py similarity index 100% rename from pilot/connections/postgres.py rename to pilot/connections/rdbms/postgres.py diff --git a/pilot/connections/rdbms/rdbms_connect.py b/pilot/connections/rdbms/rdbms_connect.py new file mode 100644 index 000000000..d3cca616f --- /dev/null +++ b/pilot/connections/rdbms/rdbms_connect.py @@ -0,0 +1,318 @@ +from __future__ import annotations + +import warnings +from typing import Any, Iterable, List, Optional +from pydantic import BaseModel, Field, root_validator, validator, Extra +from abc import ABC, abstractmethod +import sqlalchemy +from sqlalchemy import ( + MetaData, + Table, + create_engine, + inspect, + select, + text, +) +from sqlalchemy.engine import CursorResult, Engine +from sqlalchemy.exc import ProgrammingError, SQLAlchemyError +from sqlalchemy.schema import CreateTable +from sqlalchemy.orm import sessionmaker, scoped_session + +from pilot.connections.base import BaseConnect + + +def _format_index(index: sqlalchemy.engine.interfaces.ReflectedIndex) -> str: + return ( + f'Name: {index["name"]}, Unique: {index["unique"]},' + f' Columns: {str(index["column_names"])}' + ) + + +class RDBMSDatabase(BaseConnect): + """SQLAlchemy wrapper around a database.""" + + def __init__( + self, + engine, + schema: Optional[str] = None, + metadata: Optional[MetaData] = None, + ignore_tables: Optional[List[str]] = None, + include_tables: Optional[List[str]] = None, + sample_rows_in_table_info: int = 3, + indexes_in_table_info: bool = False, + custom_table_info: Optional[dict] = None, + view_support: bool = False, + ): + """Create engine from database URI.""" + self._engine = engine + self._schema = schema + if include_tables and ignore_tables: + raise ValueError("Cannot specify both include_tables and ignore_tables") + + self._inspector = inspect(self._engine) + session_factory = sessionmaker(bind=engine) + Session = scoped_session(session_factory) + + self._db_sessions = Session + + self._all_tables = set() + self.view_support = False + self._usable_tables = set() + self._include_tables = set() + self._ignore_tables = set() + self._custom_table_info = set() + self._indexes_in_table_info = set() + self._usable_tables = set() + self._usable_tables = set() + self._sample_rows_in_table_info = set() + # including view support by adding the views as well as tables to the all + # tables list if view_support is True + # self._all_tables = set( + # self._inspector.get_table_names(schema=schema) + # + (self._inspector.get_view_names(schema=schema) if view_support else []) + # ) + + # self._include_tables = set(include_tables) if include_tables else set() + # if self._include_tables: + # missing_tables = self._include_tables - self._all_tables + # if missing_tables: + # raise ValueError( + # f"include_tables {missing_tables} not found in database" + # ) + # self._ignore_tables = set(ignore_tables) if ignore_tables else set() + # if self._ignore_tables: + # missing_tables = self._ignore_tables - self._all_tables + # if missing_tables: + # raise ValueError( + # f"ignore_tables {missing_tables} not found in database" + # ) + # usable_tables = self.get_usable_table_names() + # self._usable_tables = set(usable_tables) if usable_tables else self._all_tables + + # if not isinstance(sample_rows_in_table_info, int): + # raise TypeError("sample_rows_in_table_info must be an integer") + # + # self._sample_rows_in_table_info = sample_rows_in_table_info + # self._indexes_in_table_info = indexes_in_table_info + # + # self._custom_table_info = custom_table_info + # if self._custom_table_info: + # if not isinstance(self._custom_table_info, dict): + # raise TypeError( + # "table_info must be a dictionary with table names as keys and the " + # "desired table info as values" + # ) + # # only keep the tables that are also present in the database + # intersection = set(self._custom_table_info).intersection(self._all_tables) + # self._custom_table_info = dict( + # (table, self._custom_table_info[table]) + # for table in self._custom_table_info + # if table in intersection + # ) + + # self._metadata = metadata or MetaData() + # # # including view support if view_support = true + # self._metadata.reflect( + # views=view_support, + # bind=self._engine, + # only=list(self._usable_tables), + # schema=self._schema, + # ) + + @classmethod + def from_uri( + cls, database_uri: str, engine_args: Optional[dict] = None, **kwargs: Any + ) -> RDBMSDatabase: + """Construct a SQLAlchemy engine from URI.""" + _engine_args = engine_args or {} + return cls(create_engine(database_uri, **_engine_args), **kwargs) + + @property + def dialect(self) -> str: + """Return string representation of dialect to use.""" + return self._engine.dialect.name + + def get_usable_table_names(self) -> Iterable[str]: + """Get names of tables available.""" + if self._include_tables: + return self._include_tables + return self._all_tables - self._ignore_tables + + def get_table_names(self) -> Iterable[str]: + """Get names of tables available.""" + warnings.warn( + "This method is deprecated - please use `get_usable_table_names`." + ) + return self.get_usable_table_names() + + def get_session(self, db_name: str): + session = self._db_sessions() + + self._metadata = MetaData() + # sql = f"use {db_name}" + sql = text(f"use `{db_name}`") + session.execute(sql) + + # 处理表信息数据 + + self._metadata.reflect(bind=self._engine, schema=db_name) + + # including view support by adding the views as well as tables to the all + # tables list if view_support is True + self._all_tables = set( + self._inspector.get_table_names(schema=db_name) + + ( + self._inspector.get_view_names(schema=db_name) + if self.view_support + else [] + ) + ) + + return session + + def get_current_db_name(self, session) -> str: + return session.execute(text("SELECT DATABASE()")).scalar() + + def table_simple_info(self, session): + _sql = f""" + select concat(table_name, "(" , group_concat(column_name), ")") as schema_info from information_schema.COLUMNS where table_schema="{self.get_current_db_name(session)}" group by TABLE_NAME; + """ + cursor = session.execute(text(_sql)) + results = cursor.fetchall() + return results + + @property + def table_info(self) -> str: + """Information about all tables in the database.""" + return self.get_table_info() + + def get_table_info(self, table_names: Optional[List[str]] = None) -> str: + """Get information about specified tables. + + Follows best practices as specified in: Rajkumar et al, 2022 + (https://arxiv.org/abs/2204.00498) + + If `sample_rows_in_table_info`, the specified number of sample rows will be + appended to each table description. This can increase performance as + demonstrated in the paper. + """ + all_table_names = self.get_usable_table_names() + if table_names is not None: + missing_tables = set(table_names).difference(all_table_names) + if missing_tables: + raise ValueError(f"table_names {missing_tables} not found in database") + all_table_names = table_names + + meta_tables = [ + tbl + for tbl in self._metadata.sorted_tables + if tbl.name in set(all_table_names) + and not (self.dialect == "sqlite" and tbl.name.startswith("sqlite_")) + ] + + tables = [] + for table in meta_tables: + if self._custom_table_info and table.name in self._custom_table_info: + tables.append(self._custom_table_info[table.name]) + continue + + # add create table command + create_table = str(CreateTable(table).compile(self._engine)) + table_info = f"{create_table.rstrip()}" + has_extra_info = ( + self._indexes_in_table_info or self._sample_rows_in_table_info + ) + if has_extra_info: + table_info += "\n\n/*" + if self._indexes_in_table_info: + table_info += f"\n{self._get_table_indexes(table)}\n" + if self._sample_rows_in_table_info: + table_info += f"\n{self._get_sample_rows(table)}\n" + if has_extra_info: + table_info += "*/" + tables.append(table_info) + final_str = "\n\n".join(tables) + return final_str + + def _get_sample_rows(self, table: Table) -> str: + # build the select command + command = select(table).limit(self._sample_rows_in_table_info) + + # save the columns in string format + columns_str = "\t".join([col.name for col in table.columns]) + + try: + # get the sample rows + with self._engine.connect() as connection: + sample_rows_result: CursorResult = connection.execute(command) + # shorten values in the sample rows + sample_rows = list( + map(lambda ls: [str(i)[:100] for i in ls], sample_rows_result) + ) + + # save the sample rows in string format + sample_rows_str = "\n".join(["\t".join(row) for row in sample_rows]) + + # in some dialects when there are no rows in the table a + # 'ProgrammingError' is returned + except ProgrammingError: + sample_rows_str = "" + + return ( + f"{self._sample_rows_in_table_info} rows from {table.name} table:\n" + f"{columns_str}\n" + f"{sample_rows_str}" + ) + + def _get_table_indexes(self, table: Table) -> str: + indexes = self._inspector.get_indexes(table.name) + indexes_formatted = "\n".join(map(_format_index, indexes)) + return f"Table Indexes:\n{indexes_formatted}" + + def get_table_info_no_throw(self, table_names: Optional[List[str]] = None) -> str: + """Get information about specified tables.""" + try: + return self.get_table_info(table_names) + except ValueError as e: + """Format the error message""" + return f"Error: {e}" + + def run(self, session, command: str, fetch: str = "all") -> List: + """Execute a SQL command and return a string representing the results.""" + cursor = session.execute(text(command)) + if cursor.returns_rows: + if fetch == "all": + result = cursor.fetchall() + elif fetch == "one": + result = cursor.fetchone()[0] # type: ignore + else: + raise ValueError("Fetch parameter must be either 'one' or 'all'") + field_names = tuple(i[0:] for i in cursor.keys()) + + result = list(result) + result.insert(0, field_names) + return result + + def run_no_throw(self, session, command: str, fetch: str = "all") -> List: + """Execute a SQL command and return a string representing the results. + + If the statement returns rows, a string of the results is returned. + If the statement returns no rows, an empty string is returned. + + If the statement throws an error, the error message is returned. + """ + try: + return self.run(session, command, fetch) + except SQLAlchemyError as e: + """Format the error message""" + return f"Error: {e}" + + def get_database_list(self): + session = self._db_sessions() + cursor = session.execute(text(" show databases;")) + results = cursor.fetchall() + return [ + d[0] + for d in results + if d[0] not in ["information_schema", "performance_schema", "sys", "mysql"] + ] diff --git a/pilot/conversation.py b/pilot/conversation.py index 304469838..ba5ab2701 100644 --- a/pilot/conversation.py +++ b/pilot/conversation.py @@ -44,6 +44,7 @@ class Conversation: skip_next: bool = False conv_id: Any = None last_user_input: Any = None + def get_prompt(self): if self.sep_style == SeparatorStyle.SINGLE: ret = self.system + self.sep diff --git a/pilot/memory/chat_history/base.py b/pilot/memory/chat_history/base.py index b71ad4a5f..8d60eafe7 100644 --- a/pilot/memory/chat_history/base.py +++ b/pilot/memory/chat_history/base.py @@ -1,6 +1,6 @@ from __future__ import annotations -from pydantic import BaseModel, Field, root_validator, validator,Extra +from pydantic import BaseModel, Field, root_validator, validator, Extra from abc import ABC, abstractmethod from typing import ( Any, @@ -17,13 +17,9 @@ from typing import ( from pilot.scene.message import OnceConversation - - - class BaseChatHistoryMemory(ABC): - def __init__(self): - self.conversations:List[OnceConversation] = [] + self.conversations: List[OnceConversation] = [] @abstractmethod def messages(self) -> List[OnceConversation]: # type: ignore @@ -33,8 +29,6 @@ class BaseChatHistoryMemory(ABC): def append(self, message: OnceConversation) -> None: """Append the message to the record in the local file""" - @abstractmethod def clear(self) -> None: """Clear session memory from the local file""" - diff --git a/pilot/memory/chat_history/file_history.py b/pilot/memory/chat_history/file_history.py index a3d53415b..ffdd4169b 100644 --- a/pilot/memory/chat_history/file_history.py +++ b/pilot/memory/chat_history/file_history.py @@ -5,31 +5,33 @@ import datetime from pilot.memory.chat_history.base import BaseChatHistoryMemory from pathlib import Path -from pilot.configs.config import Config -from pilot.scene.message import OnceConversation, conversation_from_dict,conversations_to_dict +from pilot.configs.config import Config +from pilot.scene.message import ( + OnceConversation, + conversation_from_dict, + conversations_to_dict, +) CFG = Config() class FileHistoryMemory(BaseChatHistoryMemory): - def __init__(self, chat_session_id:str): + def __init__(self, chat_session_id: str): now = datetime.datetime.now() date_string = now.strftime("%Y%m%d") path: str = f"{CFG.message_dir}/{date_string}" os.makedirs(path, exist_ok=True) dir_path = Path(path) - self.file_path = Path(dir_path / f"{chat_session_id}.json") + self.file_path = Path(dir_path / f"{chat_session_id}.json") if not self.file_path.exists(): self.file_path.touch() self.file_path.write_text(json.dumps([])) - - def messages(self) -> List[OnceConversation]: items = json.loads(self.file_path.read_text()) - history:List[OnceConversation] = [] + history: List[OnceConversation] = [] for onece in items: messages = conversation_from_dict(onece) history.append(messages) @@ -38,8 +40,10 @@ class FileHistoryMemory(BaseChatHistoryMemory): def append(self, once_message: OnceConversation) -> None: historys = self.messages() historys.append(once_message) - self.file_path.write_text(json.dumps(conversations_to_dict(historys), ensure_ascii=False, indent=4), encoding="UTF-8") + self.file_path.write_text( + json.dumps(conversations_to_dict(historys), ensure_ascii=False, indent=4), + encoding="UTF-8", + ) def clear(self) -> None: self.file_path.write_text(json.dumps([])) - diff --git a/pilot/out_parser/base.py b/pilot/out_parser/base.py index 881c3d034..36ca8eb9c 100644 --- a/pilot/out_parser/base.py +++ b/pilot/out_parser/base.py @@ -13,13 +13,15 @@ from typing import ( TypeVar, Union, ) +from pilot.utils import build_logger +import re from pydantic import BaseModel, Extra, Field, root_validator - +from pilot.configs.model_config import LOGDIR from pilot.prompts.base import PromptValue T = TypeVar("T") - +logger = build_logger("webserver", LOGDIR + "DbChatOutputParser.log") class BaseOutputParser(ABC): """Class to parse the output of an LLM call. @@ -41,16 +43,16 @@ class BaseOutputParser(ABC): text = text.lower() respObj = json.loads(text) - xx = respObj['response'] - xx = xx.strip(b'\x00'.decode()) + xx = respObj["response"] + xx = xx.strip(b"\x00".decode()) respObj_ex = json.loads(xx) - if respObj_ex['error_code'] == 0: - all_text = respObj_ex['text'] + if respObj_ex["error_code"] == 0: + all_text = respObj_ex["text"] ### 解析返回文本，获取AI回复部分 tmpResp = all_text.split(sep) last_index = -1 for i in range(len(tmpResp)): - if tmpResp[i].find('assistant:') != -1: + if tmpResp[i].find("assistant:") != -1: last_index = i ai_response = tmpResp[last_index] ai_response = ai_response.replace("assistant:", "") @@ -60,9 +62,7 @@ class BaseOutputParser(ABC): print("un_stream clear response:{}", ai_response) return ai_response else: - raise ValueError("Model server error!code=" + respObj_ex['error_code']); - - + raise ValueError("Model server error!code=" + respObj_ex["error_code"]) def parse_model_server_out(self, response) -> str: """ @@ -87,7 +87,27 @@ class BaseOutputParser(ABC): Returns: """ - pass + cleaned_output = model_out_text.rstrip() + if "```json" in cleaned_output: + _, cleaned_output = cleaned_output.split("```json") + if "```" in cleaned_output: + cleaned_output, _ = cleaned_output.split("```") + if cleaned_output.startswith("```json"): + cleaned_output = cleaned_output[len("```json"):] + if cleaned_output.startswith("```"): + cleaned_output = cleaned_output[len("```"):] + if cleaned_output.endswith("```"): + cleaned_output = cleaned_output[: -len("```")] + cleaned_output = cleaned_output.strip() + if not cleaned_output.startswith("{") or not cleaned_output.endswith("}"): + logger.info("illegal json processing") + json_pattern = r"{(.+?)}" + m = re.search(json_pattern, cleaned_output) + if m: + cleaned_output = m.group(0) + else: + raise ValueError("model server out not fllow the prompt!") + return cleaned_output def parse_view_response(self, ai_text) -> str: """ @@ -98,7 +118,7 @@ class BaseOutputParser(ABC): Returns: """ - pass + return ai_text def get_format_instructions(self) -> str: """Instructions on how the LLM output should be formatted.""" diff --git a/pilot/prompts/auto_mode_prompt.py b/pilot/prompts/auto_mode_prompt.py deleted file mode 100644 index b47d24a76..000000000 --- a/pilot/prompts/auto_mode_prompt.py +++ /dev/null @@ -1,143 +0,0 @@ -import platform -from typing import Optional - -import distro -import yaml - -from pilot.configs.config import Config -from pilot.prompts.generator import PromptGenerator -from pilot.prompts.prompt import ( - DEFAULT_PROMPT_OHTER, - DEFAULT_TRIGGERING_PROMPT, - build_default_prompt_generator, -) - - -class AutoModePrompt: - """ """ - - def __init__( - self, - ai_goals: list | None = None, - ) -> None: - """ - Initialize a class instance - - Parameters: - ai_name (str): The name of the AI. - ai_role (str): The description of the AI's role. - ai_goals (list): The list of objectives the AI is supposed to complete. - api_budget (float): The maximum dollar value for API calls (0.0 means infinite) - Returns: - None - """ - if ai_goals is None: - ai_goals = [] - self.ai_goals = ai_goals - self.prompt_generator = None - self.command_registry = None - - def construct_follow_up_prompt( - self, - user_input: [str], - last_auto_return: str = None, - prompt_generator: Optional[PromptGenerator] = None, - ) -> str: - """ - Build complete prompt information based on subsequent dialogue information entered by the user - Args: - self: - prompt_generator: - - Returns: - - """ - prompt_start = DEFAULT_PROMPT_OHTER - if prompt_generator is None: - prompt_generator = build_default_prompt_generator() - prompt_generator.goals = user_input - prompt_generator.command_registry = self.command_registry - # 加载插件中可用命令 - cfg = Config() - for plugin in cfg.plugins: - if not plugin.can_handle_post_prompt(): - continue - prompt_generator = plugin.post_prompt(prompt_generator) - - full_prompt = f"{prompt_start}\n\nGOALS:\n\n" - if not self.ai_goals: - self.ai_goals = user_input - for i, goal in enumerate(self.ai_goals): - full_prompt += ( - f"{i+1}.According to the provided Schema information, {goal}\n" - ) - # if last_auto_return == None: - # full_prompt += f"{cfg.last_plugin_return}\n\n" - # else: - # full_prompt += f"{last_auto_return}\n\n" - - full_prompt += f"Constraints:\n\n{DEFAULT_TRIGGERING_PROMPT}\n" - - full_prompt += """Based on the above definition, answer the current goal and ensure that the response meets both the current constraints and the above definition and constraints""" - self.prompt_generator = prompt_generator - return full_prompt - - def construct_first_prompt( - self, - fisrt_message: [str] = [], - db_schemes: str = None, - prompt_generator: Optional[PromptGenerator] = None, - ) -> str: - """ - Build complete prompt information based on the initial dialogue information entered by the user - Args: - self: - prompt_generator: - - Returns: - - """ - prompt_start = ( - "Your decisions must always be made independently without" - " seeking user assistance. Play to your strengths as an LLM and pursue" - " simple strategies with no legal complications." - "" - ) - if prompt_generator is None: - prompt_generator = build_default_prompt_generator() - prompt_generator.goals = fisrt_message - prompt_generator.command_registry = self.command_registry - # 加载插件中可用命令 - cfg = Config() - for plugin in cfg.plugins: - if not plugin.can_handle_post_prompt(): - continue - prompt_generator = plugin.post_prompt(prompt_generator) - if cfg.execute_local_commands: - # add OS info to prompt - os_name = platform.system() - os_info = ( - platform.platform(terse=True) - if os_name != "Linux" - else distro.name(pretty=True) - ) - - prompt_start += f"\nThe OS you are running on is: {os_info}" - - # Construct full prompt - full_prompt = f"{prompt_start}\n\nGOALS:\n\n" - if not self.ai_goals: - self.ai_goals = fisrt_message - for i, goal in enumerate(self.ai_goals): - full_prompt += ( - f"{i+1}.According to the provided Schema information,{goal}\n" - ) - if db_schemes: - full_prompt += f"\nSchema:\n\n" - full_prompt += f"{db_schemes}" - - # if self.api_budget > 0.0: - # full_prompt += f"\nIt takes money to let you run. Your API budget is ${self.api_budget:.3f}" - self.prompt_generator = prompt_generator - full_prompt += f"\n\n{prompt_generator.generate_prompt_string()}" - return full_prompt diff --git a/pilot/prompts/base.py b/pilot/prompts/base.py index bd082000e..12a97e94f 100644 --- a/pilot/prompts/base.py +++ b/pilot/prompts/base.py @@ -1,5 +1,3 @@ - - import json from abc import ABC, abstractmethod from pathlib import Path @@ -8,7 +6,7 @@ from typing import Any, Callable, Dict, List, Mapping, Optional, Set, Union import yaml from pydantic import BaseModel, Extra, Field, root_validator -from pilot.scene.base_message import BaseMessage,HumanMessage,AIMessage, SystemMessage +from pilot.scene.base_message import BaseMessage, HumanMessage, AIMessage, SystemMessage def get_buffer_string( @@ -29,7 +27,6 @@ def get_buffer_string( return "\n".join(string_messages) - class PromptValue(BaseModel, ABC): @abstractmethod def to_string(self) -> str: @@ -39,6 +36,7 @@ class PromptValue(BaseModel, ABC): def to_messages(self) -> List[BaseMessage]: """Return prompt as messages.""" + class ChatPromptValue(PromptValue): messages: List[BaseMessage] @@ -48,4 +46,4 @@ class ChatPromptValue(PromptValue): def to_messages(self) -> List[BaseMessage]: """Return prompt as messages.""" - return self.messages \ No newline at end of file + return self.messages diff --git a/pilot/prompts/generator.py b/pilot/prompts/generator.py index c470ff5a5..22f998a67 100644 --- a/pilot/prompts/generator.py +++ b/pilot/prompts/generator.py @@ -3,7 +3,7 @@ import json from typing import Any, Callable, Dict, List, Optional -class PromptGenerator: +class PluginPromptGenerator: """ A class for generating custom prompt strings based on constraints, commands, resources, and performance evaluations. @@ -133,6 +133,11 @@ class PromptGenerator: else: return "\n".join(f"{i+1}. {item}" for i, item in enumerate(items)) + + def generate_commands_string(self)->str: + return f"{self._generate_numbered_list(self.commands, item_type='command')}" + + def generate_prompt_string(self) -> str: """ Generate a prompt string based on the constraints, commands, resources, diff --git a/pilot/prompts/prompt.py b/pilot/prompts/prompt.py deleted file mode 100644 index d46b69ad5..000000000 --- a/pilot/prompts/prompt.py +++ /dev/null @@ -1,73 +0,0 @@ -from pilot.configs.config import Config -from pilot.prompts.generator import PromptGenerator - -CFG = Config() - -DEFAULT_TRIGGERING_PROMPT = ( - "Determine which next command to use, and respond using the format specified above" -) - -DEFAULT_PROMPT_OHTER = "Previous response was excellent. Please response according to the requirements based on the new goal" - - -def build_default_prompt_generator() -> PromptGenerator: - """ - This function generates a prompt string that includes various constraints, - commands, resources, and performance evaluations. - - Returns: - str: The generated prompt string. - """ - - # Initialize the PromptGenerator object - prompt_generator = PromptGenerator() - - # Add constraints to the PromptGenerator object - # prompt_generator.add_constraint( - # "~4000 word limit for short term memory. Your short term memory is short, so" - # " immediately save important information to files." - # ) - prompt_generator.add_constraint( - "If you are unsure how you previously did something or want to recall past" - " events, thinking about similar events will help you remember." - ) - # prompt_generator.add_constraint("No user assistance") - - prompt_generator.add_constraint("Only output one correct JSON response at a time") - prompt_generator.add_constraint( - 'Exclusively use the commands listed in double quotes e.g. "command name"' - ) - prompt_generator.add_constraint( - "If there is SQL in the args parameter, ensure to use the database and table definitions in Schema, and ensure that the fields and table names are in the definition" - ) - prompt_generator.add_constraint( - "The generated command args need to comply with the definition of the command" - ) - - # Add resources to the PromptGenerator object - # prompt_generator.add_resource( - # "Internet access for searches and information gathering." - # ) - # prompt_generator.add_resource("Long Term memory management.") - # prompt_generator.add_resource( - # "DB-GPT powered Agents for delegation of simple tasks." - # ) - # prompt_generator.add_resource("File output.") - - # Add performance evaluations to the PromptGenerator object - prompt_generator.add_performance_evaluation( - "Continuously review and analyze your actions to ensure you are performing to" - " the best of your abilities." - ) - prompt_generator.add_performance_evaluation( - "Constructively self-criticize your big-picture behavior constantly." - ) - prompt_generator.add_performance_evaluation( - "Reflect on past decisions and strategies to refine your approach." - ) - # prompt_generator.add_performance_evaluation( - # "Every command has a cost, so be smart and efficient. Aim to complete tasks in" - # " the least number of steps." - # ) - # prompt_generator.add_performance_evaluation("Write all code to a file.") - return prompt_generator diff --git a/pilot/prompts/prompt_generator.py b/pilot/prompts/prompt_generator.py index e0ffed4a6..1ec62d5c9 100644 --- a/pilot/prompts/prompt_generator.py +++ b/pilot/prompts/prompt_generator.py @@ -3,8 +3,8 @@ from typing import Any, Callable, Dict, List, Optional class PromptGenerator: """ - generating custom prompt strings based on constraints； - Compatible with AutoGpt Plugin; + generating custom prompt strings based on constraints； + Compatible with AutoGpt Plugin; """ def __init__(self) -> None: @@ -22,8 +22,6 @@ class PromptGenerator: self.role = "AI" self.response_format = None - - def add_command( self, command_label: str, @@ -51,4 +49,4 @@ class PromptGenerator: "args": command_args, "function": function, } - self.commands.append(command) \ No newline at end of file + self.commands.append(command) diff --git a/pilot/prompts/prompt_new.py b/pilot/prompts/prompt_new.py index a79be5171..389b1a33e 100644 --- a/pilot/prompts/prompt_new.py +++ b/pilot/prompts/prompt_new.py @@ -8,6 +8,7 @@ from pilot.common.formatting import formatter from pilot.out_parser.base import BaseOutputParser from pilot.common.schema import SeparatorStyle + def jinja2_formatter(template: str, **kwargs: Any) -> str: """Format a template using jinja2.""" try: @@ -32,22 +33,23 @@ class PromptTemplate(BaseModel, ABC): """A list of the names of the variables the prompt template expects.""" template_scene: str - template_define:str + template_define: str """this template define""" template: str """The prompt template.""" template_format: str = "f-string" """The format of the prompt template. Options are: 'f-string', 'jinja2'.""" - response_format:str + response_format: str """default use stream out""" stream_out: bool = True """""" output_parser: BaseOutputParser = None """""" - sep:str = SeparatorStyle.SINGLE.value + sep: str = SeparatorStyle.SINGLE.value class Config: """Configuration for this pydantic object.""" + arbitrary_types_allowed = True @property @@ -96,10 +98,8 @@ class PromptTemplate(BaseModel, ABC): else: return "\n".join(f"{i+1}. {item}" for i, item in enumerate(items)) - - def format(self, **kwargs: Any) -> str: - """Format the prompt with the inputs. - """ + def format(self, **kwargs: Any) -> str: + """Format the prompt with the inputs.""" kwargs["response"] = json.dumps(self.response_format, indent=4) return DEFAULT_FORMATTER_MAPPING[self.template_format](self.template, **kwargs) diff --git a/pilot/prompts/prompt_template.py b/pilot/prompts/prompt_template.py index ad597c33d..0a014b06e 100644 --- a/pilot/prompts/prompt_template.py +++ b/pilot/prompts/prompt_template.py @@ -207,6 +207,7 @@ class BasePromptTemplate(BaseModel, ABC): else: raise ValueError(f"{save_path} must be json or yaml") + class StringPromptValue(PromptValue): text: str @@ -219,7 +220,6 @@ class StringPromptValue(PromptValue): return [HumanMessage(content=self.text)] - class StringPromptTemplate(BasePromptTemplate, ABC): """String prompt should expose the format method, returning a prompt.""" @@ -360,4 +360,4 @@ class PromptTemplate(StringPromptTemplate): # For backwards compatibility. -Prompt = PromptTemplate \ No newline at end of file +Prompt = PromptTemplate diff --git a/pilot/scene/base.py b/pilot/scene/base.py index 302510f2b..9fcc6fb31 100644 --- a/pilot/scene/base.py +++ b/pilot/scene/base.py @@ -1,8 +1,9 @@ from enum import Enum + class ChatScene(Enum): ChatWithDb = "chat_with_db" ChatExecution = "chat_execution" ChatKnowledge = "chat_default_knowledge" ChatNewKnowledge = "chat_new_knowledge" - ChatNormal = "chat_normal" \ No newline at end of file + ChatNormal = "chat_normal" diff --git a/pilot/scene/base_chat.py b/pilot/scene/base_chat.py index bc36287a6..7a1c77781 100644 --- a/pilot/scene/base_chat.py +++ b/pilot/scene/base_chat.py @@ -1,4 +1,6 @@ from abc import ABC, abstractmethod +import datetime +import traceback from pydantic import BaseModel, Field, root_validator, validator, Extra from typing import ( Any, @@ -20,20 +22,27 @@ from pilot.prompts.prompt_new import PromptTemplate from pilot.memory.chat_history.base import BaseChatHistoryMemory from pilot.memory.chat_history.file_history import FileHistoryMemory -from pilot.configs.model_config import LOGDIR, DATASETS_DIR +from pilot.configs.model_config import LOGDIR, DATASETS_DIR from pilot.utils import ( build_logger, server_error_msg, ) -from pilot.common.schema import SeparatorStyle -from pilot.scene.base import ChatScene +from pilot.scene.base_message import ( + BaseMessage, + SystemMessage, + HumanMessage, + AIMessage, + ViewMessage, +) from pilot.configs.config import Config logger = build_logger("BaseChat", LOGDIR + "BaseChat.log") headers = {"User-Agent": "dbgpt Client"} CFG = Config() -class BaseChat( ABC): - chat_scene:str = None + + +class BaseChat(ABC): + chat_scene: str = None llm_model: Any = None temperature: float = 0.6 max_new_tokens: int = 1024 @@ -42,17 +51,20 @@ class BaseChat( ABC): class Config: """Configuration for this pydantic object.""" + arbitrary_types_allowed = True def __init__(self, chat_mode, chat_session_id, current_user_input): self.chat_session_id = chat_session_id self.chat_mode = chat_mode - self.current_user_input:str = current_user_input + self.current_user_input: str = current_user_input self.llm_model = CFG.LLM_MODEL ### TODO self.memory = FileHistoryMemory(chat_session_id) ### load prompt template - self.prompt_template: PromptTemplate = CFG.prompt_templates[self.chat_mode.value] + self.prompt_template: PromptTemplate = CFG.prompt_templates[ + self.chat_mode.value + ] self.history_message: List[OnceConversation] = [] self.current_message: OnceConversation = OnceConversation() self.current_tokens_used: int = 0 @@ -69,15 +81,163 @@ class BaseChat( ABC): def chat_type(self) -> str: raise NotImplementedError("Not supported for this chat type.") + @abstractmethod + def generate_input_values(self): + pass + + @abstractmethod + def do_with_prompt_response(self, prompt_response): + pass + def call(self): - pass + input_values = self.generate_input_values() + + ### Chat sequence advance + self.current_message.chat_order = len(self.history_message) + 1 + self.current_message.add_user_message(self.current_user_input) + self.current_message.start_date = datetime.datetime.now() + # TODO + self.current_message.tokens = 0 + + current_prompt = self.prompt_template.format(**input_values) + + ### 构建当前对话， 是否安第一次对话prompt构造？ 是否考虑切换库 + if self.history_message: + ## TODO 带历史对话记录的场景需要确定切换库后怎么处理 + logger.info( + f"There are already {len(self.history_message)} rounds of conversations!" + ) + + self.current_message.add_system_message(current_prompt) + + payload = { + "model": self.llm_model, + "prompt": self.generate_llm_text(), + "temperature": float(self.temperature), + "max_new_tokens": int(self.max_new_tokens), + "stop": self.prompt_template.sep, + } + logger.info(f"Requert: \n{payload}") + ai_response_text = "" + try: + if not self.prompt_template.stream_out: + ### 走非流式的模型服务接口 + response = requests.post( + urljoin(CFG.MODEL_SERVER, "generate"), + headers=headers, + json=payload, + timeout=120, + ) + + ### output parse + ai_response_text = ( + self.prompt_template.output_parser.parse_model_server_out(response) + ) + self.current_message.add_ai_message(ai_response_text) + prompt_define_response = self.prompt_template.output_parser.parse_prompt_response(ai_response_text) + + result = self.do_with_prompt_response(prompt_define_response) + + if hasattr(prompt_define_response, "thoughts"): + if prompt_define_response.thoughts.get("speak"): + self.current_message.add_view_message( + self.prompt_template.output_parser.parse_view_response( + prompt_define_response.thoughts.get("speak"), result + ) + ) + elif prompt_define_response.thoughts.get("reasoning"): + self.current_message.add_view_message( + self.prompt_template.output_parser.parse_view_response( + prompt_define_response.thoughts.get("reasoning"), result + ) + ) + else: + self.current_message.add_view_message( + self.prompt_template.output_parser.parse_view_response( + prompt_define_response.thoughts, result + ) + ) + else: + self.current_message.add_view_message( + self.prompt_template.output_parser.parse_view_response( + prompt_define_response, result + ) + ) + else: + response = requests.post( + urljoin(CFG.MODEL_SERVER, "generate_stream"), + headers=headers, + json=payload, + timeout=120, + ) + #TODO + + + + except Exception as e: + print(traceback.format_exc()) + logger.error("model response parase faild！" + str(e)) + self.current_message.add_view_message( + f"""ERROR!{str(e)}\n {ai_response_text} """ + ) + ### 对话记录存储 + self.memory.append(self.current_message) + + def generate_llm_text(self) -> str: + text = self.prompt_template.template_define + self.prompt_template.sep + ### 线处理历史信息 + if len(self.history_message) > self.chat_retention_rounds: + ### 使用历史信息的第一轮和最后一轮数据合并成历史对话记录, 做上下文提示时，用户展示消息需要过滤掉 + for first_message in self.history_message[0].messages: + if not isinstance(first_message, ViewMessage): + text += ( + first_message.type + + ":" + + first_message.content + + self.prompt_template.sep + ) + + index = self.chat_retention_rounds - 1 + for last_message in self.history_message[-index:].messages: + if not isinstance(last_message, ViewMessage): + text += ( + last_message.type + + ":" + + last_message.content + + self.prompt_template.sep + ) + + else: + ### 直接历史记录拼接 + for conversation in self.history_message: + for message in conversation.messages: + if not isinstance(message, ViewMessage): + text += ( + message.type + + ":" + + message.content + + self.prompt_template.sep + ) + + ### current conversation + for now_message in self.current_message.messages: + text += ( + now_message.type + ":" + now_message.content + self.prompt_template.sep + ) + + return text + def chat_show(self): pass + # 暂时为了兼容前端 def current_ai_response(self) -> str: - pass + for message in self.current_message.messages: + if message.type == "view": + return message.content + return None def _load_history(self, session_id: str) -> List[OnceConversation]: """ @@ -88,7 +248,7 @@ class BaseChat( ABC): """ return self.memory.messages() - def generate(self, p)->str: + def generate(self, p) -> str: """ generate context for LLM input Args: diff --git a/pilot/scene/base_message.py b/pilot/scene/base_message.py index 5cd8c4426..56fbb3b20 100644 --- a/pilot/scene/base_message.py +++ b/pilot/scene/base_message.py @@ -15,6 +15,7 @@ from typing import ( from pydantic import BaseModel, Extra, Field, root_validator + class PromptValue(BaseModel, ABC): @abstractmethod def to_string(self) -> str: @@ -37,7 +38,6 @@ class BaseMessage(BaseModel): """Type of the message, used for serialization.""" - class HumanMessage(BaseMessage): """Type of message that is spoken by the human.""" @@ -49,7 +49,6 @@ class HumanMessage(BaseMessage): return "human" - class AIMessage(BaseMessage): """Type of message that is spoken by the AI.""" @@ -81,8 +80,6 @@ class SystemMessage(BaseMessage): return "system" - - class Generation(BaseModel): """Output of a single generation.""" @@ -94,7 +91,6 @@ class Generation(BaseModel): """May include things like reason for finishing (e.g. in OpenAI)""" - class ChatGeneration(Generation): """Output of a single generation.""" @@ -126,7 +122,6 @@ class LLMResult(BaseModel): """For arbitrary LLM provider specific output.""" - def _message_to_dict(message: BaseMessage) -> dict: return {"type": message.type, "data": message.dict()} @@ -149,6 +144,5 @@ def _message_from_dict(message: dict) -> BaseMessage: raise ValueError(f"Got unexpected type: {_type}") - def messages_from_dict(messages: List[dict]) -> List[BaseMessage]: return [_message_from_dict(m) for m in messages] diff --git a/pilot/scene/chat_db/chat.py b/pilot/scene/chat_db/chat.py index 72ad64508..37382b7d2 100644 --- a/pilot/scene/chat_db/chat.py +++ b/pilot/scene/chat_db/chat.py @@ -14,7 +14,13 @@ from sqlalchemy import ( ) from typing import Any, Iterable, List, Optional -from pilot.scene.base_message import BaseMessage, SystemMessage, HumanMessage, AIMessage, ViewMessage +from pilot.scene.base_message import ( + BaseMessage, + SystemMessage, + HumanMessage, + AIMessage, + ViewMessage, +) from pilot.scene.base_chat import BaseChat, logger, headers from pilot.scene.base import ChatScene from pilot.common.sql_database import Database @@ -25,22 +31,24 @@ from pilot.utils import ( build_logger, server_error_msg, ) -from pilot.common.markdown_text import generate_markdown_table, generate_htm_table, datas_to_table_html +from pilot.common.markdown_text import ( + generate_markdown_table, + generate_htm_table, + datas_to_table_html, +) from pilot.scene.chat_db.prompt import chat_db_prompt from pilot.out_parser.base import BaseOutputParser from pilot.scene.chat_db.out_parser import DbChatOutputParser CFG = Config() - class ChatWithDb(BaseChat): chat_scene: str = ChatScene.ChatWithDb.value """Number of results to return from the query""" def __init__(self, chat_session_id, db_name, user_input): - """ - """ + """ """ super().__init__(ChatScene.ChatWithDb, chat_session_id, user_input) if not db_name: raise ValueError(f"{ChatScene.ChatWithDb.value} mode should chose db!") @@ -50,118 +58,126 @@ class ChatWithDb(BaseChat): self.db_connect = self.database.get_session(self.db_name) self.top_k: int = 5 - def call(self) -> str: + def generate_input_values(self): input_values = { - "input": self.current_user_input, - "top_k": str(self.top_k), - "dialect": self.database.dialect, - "table_info": self.database.table_simple_info(self.db_connect), - # "stop": self.sep_style, - } + "input": self.current_user_input, + "top_k": str(self.top_k), + "dialect": self.database.dialect, + "table_info": self.database.table_simple_info(self.db_connect) + } + return input_values - ### Chat sequence advance - self.current_message.chat_order = len(self.history_message) + 1 - self.current_message.add_user_message(self.current_user_input) - self.current_message.start_date = datetime.datetime.now() - # TODO - self.current_message.tokens = 0 - - current_prompt = self.prompt_template.format(**input_values) - - ### 构建当前对话， 是否安第一次对话prompt构造？ 是否考虑切换库 - if self.history_message: - ## TODO 带历史对话记录的场景需要确定切换库后怎么处理 - logger.info(f"There are already {len(self.history_message)} rounds of conversations!") - - self.current_message.add_system_message(current_prompt) - - payload = { - "model": self.llm_model, - "prompt": self.generate_llm_text(), - "temperature": float(self.temperature), - "max_new_tokens": int(self.max_new_tokens), - "stop": self.prompt_template.sep, - } - logger.info(f"Requert: \n{payload}") - ai_response_text = "" - try: - ### 走非流式的模型服务接口 - - response = requests.post(urljoin(CFG.MODEL_SERVER, "generate"), headers=headers, json=payload, timeout=120) - ai_response_text = self.prompt_template.output_parser.parse_model_server_out(response) - self.current_message.add_ai_message(ai_response_text) - prompt_define_response = self.prompt_template.output_parser.parse_prompt_response(ai_response_text) - - result = self.database.run(self.db_connect, prompt_define_response.sql) + def do_with_prompt_response(self, prompt_response): + return self.database.run(self.db_connect, prompt_response.sql) - if hasattr(prompt_define_response, 'thoughts'): - if prompt_define_response.thoughts.get("speak"): - self.current_message.add_view_message( - self.prompt_template.output_parser.parse_view_response(prompt_define_response.thoughts.get("speak"),result)) - elif prompt_define_response.thoughts.get("reasoning"): - self.current_message.add_view_message( - self.prompt_template.output_parser.parse_view_response(prompt_define_response.thoughts.get("reasoning"), result)) - else: - self.current_message.add_view_message( - self.prompt_template.output_parser.parse_view_response(prompt_define_response.thoughts, result)) - else: - self.current_message.add_view_message( - self.prompt_template.output_parser.parse_view_response(prompt_define_response, result)) - - except Exception as e: - print(traceback.format_exc()) - logger.error("model response parase faild！" + str(e)) - self.current_message.add_view_message(f"""ERROR!{str(e)}\n {ai_response_text} """) - ### 对话记录存储 - self.memory.append(self.current_message) + # def call(self) -> str: + # input_values = { + # "input": self.current_user_input, + # "top_k": str(self.top_k), + # "dialect": self.database.dialect, + # "table_info": self.database.table_simple_info(self.db_connect), + # # "stop": self.sep_style, + # } + # + # ### Chat sequence advance + # self.current_message.chat_order = len(self.history_message) + 1 + # self.current_message.add_user_message(self.current_user_input) + # self.current_message.start_date = datetime.datetime.now() + # # TODO + # self.current_message.tokens = 0 + # + # current_prompt = self.prompt_template.format(**input_values) + # + # ### 构建当前对话， 是否安第一次对话prompt构造？ 是否考虑切换库 + # if self.history_message: + # ## TODO 带历史对话记录的场景需要确定切换库后怎么处理 + # logger.info( + # f"There are already {len(self.history_message)} rounds of conversations!" + # ) + # + # self.current_message.add_system_message(current_prompt) + # + # payload = { + # "model": self.llm_model, + # "prompt": self.generate_llm_text(), + # "temperature": float(self.temperature), + # "max_new_tokens": int(self.max_new_tokens), + # "stop": self.prompt_template.sep, + # } + # logger.info(f"Requert: \n{payload}") + # ai_response_text = "" + # try: + # ### 走非流式的模型服务接口 + # + # response = requests.post( + # urljoin(CFG.MODEL_SERVER, "generate"), + # headers=headers, + # json=payload, + # timeout=120, + # ) + # ai_response_text = ( + # self.prompt_template.output_parser.parse_model_server_out(response) + # ) + # self.current_message.add_ai_message(ai_response_text) + # prompt_define_response = ( + # self.prompt_template.output_parser.parse_prompt_response( + # ai_response_text + # ) + # ) + # + # result = self.database.run(self.db_connect, prompt_define_response.sql) + # + # if hasattr(prompt_define_response, "thoughts"): + # if prompt_define_response.thoughts.get("speak"): + # self.current_message.add_view_message( + # self.prompt_template.output_parser.parse_view_response( + # prompt_define_response.thoughts.get("speak"), result + # ) + # ) + # elif prompt_define_response.thoughts.get("reasoning"): + # self.current_message.add_view_message( + # self.prompt_template.output_parser.parse_view_response( + # prompt_define_response.thoughts.get("reasoning"), result + # ) + # ) + # else: + # self.current_message.add_view_message( + # self.prompt_template.output_parser.parse_view_response( + # prompt_define_response.thoughts, result + # ) + # ) + # else: + # self.current_message.add_view_message( + # self.prompt_template.output_parser.parse_view_response( + # prompt_define_response, result + # ) + # ) + # + # except Exception as e: + # print(traceback.format_exc()) + # logger.error("model response parase faild！" + str(e)) + # self.current_message.add_view_message( + # f"""ERROR!{str(e)}\n {ai_response_text} """ + # ) + # ### 对话记录存储 + # self.memory.append(self.current_message) def chat_show(self): ret = [] # 单论对话只能有一次User 记录 和一次 AI 记录 # TODO 推理过程前端展示。。。 for message in self.current_message.messages: - if (isinstance(message, HumanMessage)): + if isinstance(message, HumanMessage): ret[-1][-2] = message.content # 是否展示推理过程 - if (isinstance(message, ViewMessage)): + if isinstance(message, ViewMessage): ret[-1][-1] = message.content return ret - # 暂时为了兼容前端 - def current_ai_response(self) -> str: - for message in self.current_message.messages: - if message.type == 'view': - return message.content - return None - def generate_llm_text(self) -> str: - text = self.prompt_template.template_define + self.prompt_template.sep - ### 线处理历史信息 - if (len(self.history_message) > self.chat_retention_rounds): - ### 使用历史信息的第一轮和最后一轮数据合并成历史对话记录, 做上下文提示时，用户展示消息需要过滤掉 - for first_message in self.history_message[0].messages: - if not isinstance(first_message, ViewMessage): - text += first_message.type + ":" + first_message.content + self.prompt_template.sep - index = self.chat_retention_rounds - 1 - for last_message in self.history_message[-index:].messages: - if not isinstance(last_message, ViewMessage): - text += last_message.type + ":" + last_message.content + self.prompt_template.sep - - else: - ### 直接历史记录拼接 - for conversation in self.history_message: - for message in conversation.messages: - if not isinstance(message, ViewMessage): - text += message.type + ":" + message.content + self.prompt_template.sep - - ### current conversation - for now_message in self.current_message.messages: - text += now_message.type + ":" + now_message.content + self.prompt_template.sep - - return text @property def chat_type(self) -> str: diff --git a/pilot/scene/chat_db/out_parser.py b/pilot/scene/chat_db/out_parser.py index 1d2597f57..307aff680 100644 --- a/pilot/scene/chat_db/out_parser.py +++ b/pilot/scene/chat_db/out_parser.py @@ -1,52 +1,28 @@ import json import re from abc import ABC, abstractmethod -from typing import ( - Dict, - NamedTuple -) +from typing import Dict, NamedTuple import pandas as pd from pilot.utils import build_logger from pilot.out_parser.base import BaseOutputParser, T from pilot.configs.model_config import LOGDIR + class SqlAction(NamedTuple): sql: str thoughts: Dict + logger = build_logger("webserver", LOGDIR + "DbChatOutputParser.log") + + class DbChatOutputParser(BaseOutputParser): + def __init__(self, sep: str, is_stream_out: bool): + super().__init__(sep=sep, is_stream_out=is_stream_out) - def __init__(self, sep:str, is_stream_out: bool): - super().__init__(sep=sep, is_stream_out=is_stream_out ) - - - def parse_model_server_out(self, response) -> str: - return super().parse_model_server_out(response) def parse_prompt_response(self, model_out_text): - cleaned_output = model_out_text.rstrip() - if "```json" in cleaned_output: - _, cleaned_output = cleaned_output.split("```json") - if "```" in cleaned_output: - cleaned_output, _ = cleaned_output.split("```") - if cleaned_output.startswith("```json"): - cleaned_output = cleaned_output[len("```json"):] - if cleaned_output.startswith("```"): - cleaned_output = cleaned_output[len("```"):] - if cleaned_output.endswith("```"): - cleaned_output = cleaned_output[: -len("```")] - cleaned_output = cleaned_output.strip() - if not cleaned_output.startswith("{") or not cleaned_output.endswith("}"): - logger.info("illegal json processing") - json_pattern = r'{(.+?)}' - m = re.search(json_pattern, cleaned_output) - if m: - cleaned_output = m.group(0) - else: - raise ValueError("model server out not fllow the prompt!") - - response = json.loads(cleaned_output) + response = json.loads(super().parse_prompt_response(model_out_text)) sql, thoughts = response["sql"], response["thoughts"] return SqlAction(sql, thoughts) diff --git a/pilot/scene/chat_db/prompt.py b/pilot/scene/chat_db/prompt.py index 8ff1a2b1b..aeaf994c0 100644 --- a/pilot/scene/chat_db/prompt.py +++ b/pilot/scene/chat_db/prompt.py @@ -45,9 +45,15 @@ RESPONSE_FORMAT = { "reasoning": "reasoning", "speak": "thoughts summary to say to user", }, - "sql": "SQL Query to run" + "sql": "SQL Query to run", } +RESPONSE_FORMAT_SIMPLE = { + "thoughts": "thoughts summary to say to user", + "sql": "SQL Query to run", +} + + PROMPT_SEP = SeparatorStyle.SINGLE.value PROMPT_NEED_NEED_STREAM_OUT = False @@ -55,11 +61,13 @@ PROMPT_NEED_NEED_STREAM_OUT = False chat_db_prompt = PromptTemplate( template_scene=ChatScene.ChatWithDb.value, input_variables=["input", "table_info", "dialect", "top_k", "response"], - response_format=json.dumps(RESPONSE_FORMAT, indent=4), + response_format=json.dumps(RESPONSE_FORMAT_SIMPLE, indent=4), template_define=PROMPT_SCENE_DEFINE, template=_DEFAULT_TEMPLATE + PROMPT_SUFFIX + PROMPT_RESPONSE, stream_out=PROMPT_NEED_NEED_STREAM_OUT, - output_parser=DbChatOutputParser(sep=PROMPT_SEP, is_stream_out=PROMPT_NEED_NEED_STREAM_OUT), + output_parser=DbChatOutputParser( + sep=PROMPT_SEP, is_stream_out=PROMPT_NEED_NEED_STREAM_OUT + ), ) CFG.prompt_templates.update({chat_db_prompt.template_scene: chat_db_prompt}) diff --git a/pilot/scene/chat_execution/chat.py b/pilot/scene/chat_execution/chat.py index 5e85c4981..a5abadad0 100644 --- a/pilot/scene/chat_execution/chat.py +++ b/pilot/scene/chat_execution/chat.py @@ -1,26 +1,156 @@ +import requests +import datetime +from urllib.parse import urljoin from typing import List +import traceback from pilot.scene.base_chat import BaseChat, logger, headers from pilot.scene.message import OnceConversation from pilot.scene.base import ChatScene +from pilot.configs.config import Config +from pilot.commands.command import execute_command +from pilot.prompts.generator import PluginPromptGenerator + +CFG = Config() class ChatWithPlugin(BaseChat): - chat_scene: str= ChatScene.ChatExecution.value - def __init__(self, chat_mode, chat_session_id, current_user_input): - super().__init__(chat_mode, chat_session_id, current_user_input) + chat_scene: str = ChatScene.ChatExecution.value + plugins_prompt_generator:PluginPromptGenerator + select_plugin: str = None - def call(self): - super().call() + def __init__(self, chat_mode, chat_session_id, current_user_input, select_plugin:str=None): + super().__init__(chat_mode, chat_session_id, current_user_input) + self.plugins_prompt_generator = PluginPromptGenerator() + self.plugins_prompt_generator.command_registry = self.command_registry + # 加载插件中可用命令 + self.select_plugin = select_plugin + if self.select_plugin: + for plugin in CFG.plugins: + if plugin. + 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) + + + + + def generate_input_values(self): + input_values = { + "input": self.current_user_input, + "constraints": self.__list_to_prompt_str(self.plugins_prompt_generator.constraints), + "commands_infos": self.plugins_prompt_generator.generate_commands_string() + } + return input_values + + def do_with_prompt_response(self, prompt_response): + ## plugin command run + return execute_command(str(prompt_response), self.plugins_prompt_generator) + + + # def call(self): + # input_values = { + # "input": self.current_user_input, + # "constraints": self.__list_to_prompt_str(self.plugins_prompt_generator.constraints), + # "commands_infos": self.__get_comnands_promp_info() + # } + # + # ### Chat sequence advance + # self.current_message.chat_order = len(self.history_message) + 1 + # self.current_message.add_user_message(self.current_user_input) + # self.current_message.start_date = datetime.datetime.now() + # # TODO + # self.current_message.tokens = 0 + # + # current_prompt = self.prompt_template.format(**input_values) + # + # ### 构建当前对话， 是否安第一次对话prompt构造？ 是否考虑切换库 + # if self.history_message: + # ## TODO 带历史对话记录的场景需要确定切换库后怎么处理 + # logger.info( + # f"There are already {len(self.history_message)} rounds of conversations!" + # ) + # + # self.current_message.add_system_message(current_prompt) + # + # payload = { + # "model": self.llm_model, + # "prompt": self.generate_llm_text(), + # "temperature": float(self.temperature), + # "max_new_tokens": int(self.max_new_tokens), + # "stop": self.prompt_template.sep, + # } + # logger.info(f"Requert: \n{payload}") + # ai_response_text = "" + # try: + # ### 走非流式的模型服务接口 + # + # response = requests.post( + # urljoin(CFG.MODEL_SERVER, "generate"), + # headers=headers, + # json=payload, + # timeout=120, + # ) + # ai_response_text = ( + # self.prompt_template.output_parser.parse_model_server_out(response) + # ) + # self.current_message.add_ai_message(ai_response_text) + # prompt_define_response = self.prompt_template.output_parser.parse_prompt_response(ai_response_text) + # + # + # ## plugin command run + # result = execute_command(prompt_define_response, self.plugins_prompt_generator) + # + # if hasattr(prompt_define_response, "thoughts"): + # if prompt_define_response.thoughts.get("speak"): + # self.current_message.add_view_message( + # self.prompt_template.output_parser.parse_view_response( + # prompt_define_response.thoughts.get("speak"), result + # ) + # ) + # elif prompt_define_response.thoughts.get("reasoning"): + # self.current_message.add_view_message( + # self.prompt_template.output_parser.parse_view_response( + # prompt_define_response.thoughts.get("reasoning"), result + # ) + # ) + # else: + # self.current_message.add_view_message( + # self.prompt_template.output_parser.parse_view_response( + # prompt_define_response.thoughts, result + # ) + # ) + # else: + # self.current_message.add_view_message( + # self.prompt_template.output_parser.parse_view_response( + # prompt_define_response, result + # ) + # ) + # + # except Exception as e: + # print(traceback.format_exc()) + # logger.error("model response parase faild！" + str(e)) + # self.current_message.add_view_message( + # f"""ERROR!{str(e)}\n {ai_response_text} """ + # ) + # ### 对话记录存储 + # self.memory.append(self.current_message) def chat_show(self): super().chat_show() - def _load_history(self, session_id: str) -> List[OnceConversation]: - return super()._load_history(session_id) + + def __list_to_prompt_str(list: List) -> str: + if not list: + separator = '\n' + return separator.join(list) + else: + return "" def generate(self, p) -> str: return super().generate(p) @property def chat_type(self) -> str: - return ChatScene.ChatExecution.value \ No newline at end of file + return ChatScene.ChatExecution.value diff --git a/pilot/scene/chat_execution/out_parser.py b/pilot/scene/chat_execution/out_parser.py new file mode 100644 index 000000000..f3f9e683e --- /dev/null +++ b/pilot/scene/chat_execution/out_parser.py @@ -0,0 +1,30 @@ +import json +import re +from abc import ABC, abstractmethod +from typing import Dict, NamedTuple +import pandas as pd +from pilot.utils import build_logger +from pilot.out_parser.base import BaseOutputParser, T +from pilot.configs.model_config import LOGDIR + + +logger = build_logger("webserver", LOGDIR + "DbChatOutputParser.log") + +class PluginAction(NamedTuple): + command: Dict + thoughts: Dict + + + +class PluginChatOutputParser(BaseOutputParser): + + def parse_prompt_response(self, model_out_text) -> T: + response = json.loads(super().parse_prompt_response(model_out_text)) + sql, thoughts = response["command"], response["thoughts"] + return PluginAction(sql, thoughts) + + def parse_view_response(self, ai_text) -> str: + return super().parse_view_response(ai_text) + + def get_format_instructions(self) -> str: + pass diff --git a/pilot/scene/chat_execution/prompt.py b/pilot/scene/chat_execution/prompt.py new file mode 100644 index 000000000..e3469d7c2 --- /dev/null +++ b/pilot/scene/chat_execution/prompt.py @@ -0,0 +1,65 @@ +import json +from pilot.prompts.prompt_new import PromptTemplate +from pilot.configs.config import Config +from pilot.scene.base import ChatScene +from pilot.common.schema import SeparatorStyle + +from pilot.scene.chat_execution.out_parser import PluginChatOutputParser + + +CFG = Config() + +PROMPT_SCENE_DEFINE = """You are an AI designed to solve the user's goals with given commands, please follow the prompts and constraints of the system's input for your answers.Play to your strengths as an LLM and pursue simple strategies with no legal complications.""" + +PROMPT_SUFFIX = """ +Goals: + {input} + +""" + +_DEFAULT_TEMPLATE = """ +Constraints: + Exclusively use the commands listed in double quotes e.g. "command name" + Reflect on past decisions and strategies to refine your approach. + Constructively self-criticize your big-picture behavior constantly. + {constraints} + +Commands: + {commands_infos} +""" + + +PROMPT_RESPONSE = """You must respond in JSON format as following format: +{response} + +Ensure the response is correct json and can be parsed by Python json.loads +""" + +RESPONSE_FORMAT = { + "thoughts": { + "text": "thought", + "reasoning": "reasoning", + "plan": "- short bulleted\n- list that conveys\n- long-term plan", + "criticism": "constructive self-criticism", + "speak": "thoughts summary to say to user", + }, + "command": {"name": "command name", "args": {"arg name": "value"}}, +} + +PROMPT_SEP = SeparatorStyle.SINGLE.value +### Whether the model service is streaming output +PROMPT_NEED_NEED_STREAM_OUT = False + +chat_plugin_prompt = PromptTemplate( + template_scene=ChatScene.ChatExecution.value, + input_variables=["input", "table_info", "dialect", "top_k", "response"], + response_format=json.dumps(RESPONSE_FORMAT, indent=4), + template_define=PROMPT_SCENE_DEFINE, + template=PROMPT_SUFFIX + _DEFAULT_TEMPLATE + PROMPT_RESPONSE, + stream_out=PROMPT_NEED_NEED_STREAM_OUT, + output_parser=PluginChatOutputParser( + sep=PROMPT_SEP, is_stream_out=PROMPT_NEED_NEED_STREAM_OUT + ), +) + +CFG.prompt_templates.update({chat_plugin_prompt.template_scene: chat_plugin_prompt}) diff --git a/pilot/scene/chat_execution/prompt_with_command.py b/pilot/scene/chat_execution/prompt_with_command.py new file mode 100644 index 000000000..e3469d7c2 --- /dev/null +++ b/pilot/scene/chat_execution/prompt_with_command.py @@ -0,0 +1,65 @@ +import json +from pilot.prompts.prompt_new import PromptTemplate +from pilot.configs.config import Config +from pilot.scene.base import ChatScene +from pilot.common.schema import SeparatorStyle + +from pilot.scene.chat_execution.out_parser import PluginChatOutputParser + + +CFG = Config() + +PROMPT_SCENE_DEFINE = """You are an AI designed to solve the user's goals with given commands, please follow the prompts and constraints of the system's input for your answers.Play to your strengths as an LLM and pursue simple strategies with no legal complications.""" + +PROMPT_SUFFIX = """ +Goals: + {input} + +""" + +_DEFAULT_TEMPLATE = """ +Constraints: + Exclusively use the commands listed in double quotes e.g. "command name" + Reflect on past decisions and strategies to refine your approach. + Constructively self-criticize your big-picture behavior constantly. + {constraints} + +Commands: + {commands_infos} +""" + + +PROMPT_RESPONSE = """You must respond in JSON format as following format: +{response} + +Ensure the response is correct json and can be parsed by Python json.loads +""" + +RESPONSE_FORMAT = { + "thoughts": { + "text": "thought", + "reasoning": "reasoning", + "plan": "- short bulleted\n- list that conveys\n- long-term plan", + "criticism": "constructive self-criticism", + "speak": "thoughts summary to say to user", + }, + "command": {"name": "command name", "args": {"arg name": "value"}}, +} + +PROMPT_SEP = SeparatorStyle.SINGLE.value +### Whether the model service is streaming output +PROMPT_NEED_NEED_STREAM_OUT = False + +chat_plugin_prompt = PromptTemplate( + template_scene=ChatScene.ChatExecution.value, + input_variables=["input", "table_info", "dialect", "top_k", "response"], + response_format=json.dumps(RESPONSE_FORMAT, indent=4), + template_define=PROMPT_SCENE_DEFINE, + template=PROMPT_SUFFIX + _DEFAULT_TEMPLATE + PROMPT_RESPONSE, + stream_out=PROMPT_NEED_NEED_STREAM_OUT, + output_parser=PluginChatOutputParser( + sep=PROMPT_SEP, is_stream_out=PROMPT_NEED_NEED_STREAM_OUT + ), +) + +CFG.prompt_templates.update({chat_plugin_prompt.template_scene: chat_plugin_prompt}) diff --git a/pilot/scene/chat_factory.py b/pilot/scene/chat_factory.py index 2d72fc7fe..97c547390 100644 --- a/pilot/scene/chat_factory.py +++ b/pilot/scene/chat_factory.py @@ -1,19 +1,17 @@ - from pilot.scene.base_chat import BaseChat from pilot.singleton import Singleton from pilot.scene.chat_db.chat import ChatWithDb from pilot.scene.chat_execution.chat import ChatWithPlugin -class ChatFactory(metaclass=Singleton): +class ChatFactory(metaclass=Singleton): @staticmethod def get_implementation(chat_mode, **kwargs): - chat_classes = BaseChat.__subclasses__() implementation = None for cls in chat_classes: - if(cls.chat_scene == chat_mode): + if cls.chat_scene == chat_mode: implementation = cls(**kwargs) - if(implementation == None): - raise Exception('Invalid implementation name:' + chat_mode) - return implementation \ No newline at end of file + if implementation == None: + raise Exception("Invalid implementation name:" + chat_mode) + return implementation diff --git a/pilot/scene/message.py b/pilot/scene/message.py index 8dc3eaa3e..0203ec68c 100644 --- a/pilot/scene/message.py +++ b/pilot/scene/message.py @@ -9,12 +9,20 @@ from typing import ( List, ) -from pilot.scene.base_message import BaseMessage, AIMessage, HumanMessage, SystemMessage, ViewMessage, messages_to_dict, messages_from_dict +from pilot.scene.base_message import ( + BaseMessage, + AIMessage, + HumanMessage, + SystemMessage, + ViewMessage, + messages_to_dict, + messages_from_dict, +) class OnceConversation: """ - All the information of a conversation, the current single service in memory, can expand cache and database support distributed services + All the information of a conversation, the current single service in memory, can expand cache and database support distributed services """ def __init__(self): @@ -26,7 +34,9 @@ class OnceConversation: def add_user_message(self, message: str) -> None: """Add a user message to the store""" - has_message = any(isinstance(instance, HumanMessage) for instance in self.messages) + has_message = any( + isinstance(instance, HumanMessage) for instance in self.messages + ) if has_message: raise ValueError("Already Have Human message") self.messages.append(HumanMessage(content=message)) @@ -38,6 +48,7 @@ class OnceConversation: raise ValueError("Already Have Ai message") self.messages.append(AIMessage(content=message)) """ """ + def add_view_message(self, message: str) -> None: """Add an AI message to the store""" @@ -50,7 +61,7 @@ class OnceConversation: def set_start_time(self, datatime: datetime): dt_str = datatime.strftime("%Y-%m-%d %H:%M:%S") - self.start_date = dt_str; + self.start_date = dt_str def clear(self) -> None: """Remove all messages from the store""" @@ -71,7 +82,7 @@ def _conversation_to_dic(once: OnceConversation) -> dict: "start_date": start_str, "cost": once.cost if once.cost else 0, "tokens": once.tokens if once.tokens else 0, - "messages": messages_to_dict(once.messages) + "messages": messages_to_dict(once.messages), } @@ -81,10 +92,10 @@ def conversations_to_dict(conversations: List[OnceConversation]) -> List[dict]: def conversation_from_dict(once: dict) -> OnceConversation: conversation = OnceConversation() - conversation.cost = once.get('cost', 0) - conversation.tokens = once.get('tokens', 0) - conversation.start_date = once.get('start_date', '') - conversation.chat_order = int(once.get('chat_order')) - print(once.get('messages')) - conversation.messages = messages_from_dict(once.get('messages', [])) + conversation.cost = once.get("cost", 0) + conversation.tokens = once.get("tokens", 0) + conversation.start_date = once.get("start_date", "") + conversation.chat_order = int(once.get("chat_order")) + print(once.get("messages")) + conversation.messages = messages_from_dict(once.get("messages", [])) return conversation diff --git a/pilot/server/webserver.py b/pilot/server/webserver.py index 6e7d8b700..8fefdbfff 100644 --- a/pilot/server/webserver.py +++ b/pilot/server/webserver.py @@ -30,7 +30,7 @@ from pilot.configs.model_config import ( LOGDIR, VECTOR_SEARCH_TOP_K, ) -from pilot.connections.mysql import MySQLOperator + from pilot.conversation import ( SeparatorStyle, conv_qa_prompt_template, @@ -39,9 +39,9 @@ from pilot.conversation import ( conversation_types, default_conversation, ) -from pilot.plugins import scan_plugins -from pilot.prompts.auto_mode_prompt import AutoModePrompt -from pilot.prompts.generator import PromptGenerator +from pilot.common.plugins import scan_plugins + +from pilot.prompts.generator import PluginPromptGenerator from pilot.server.gradio_css import code_highlight_css from pilot.server.gradio_patch import Chatbot as grChatbot from pilot.server.vectordb_qa import KnownLedgeBaseQA @@ -95,19 +95,14 @@ def get_simlar(q): def gen_sqlgen_conversation(dbname): - mo = MySQLOperator(**DB_SETTINGS) - message = "" - - schemas = mo.get_schema(dbname) + db_connect = CFG.local_db.get_session(dbname) + schemas = CFG.local_db.table_simple_info(db_connect) for s in schemas: message += s["schema_info"] + ";" return f"数据库{dbname}的Schema信息如下: {message}\n" -def get_database_list(): - mo = MySQLOperator(**DB_SETTINGS) - return mo.get_db_list() get_window_url_params = """ @@ -127,7 +122,6 @@ function() { def load_demo(url_params, request: gr.Request): logger.info(f"load_demo. ip: {request.client.host}. params: {url_params}") - # dbs = get_database_list() dropdown_update = gr.Dropdown.update(visible=True) if dbs: gr.Dropdown.update(choices=dbs) @@ -137,13 +131,15 @@ def load_demo(url_params, request: gr.Request): unique_id = uuid.uuid1() state.conv_id = str(unique_id) - return (state, - dropdown_update, - gr.Chatbot.update(visible=True), - gr.Textbox.update(visible=True), - gr.Button.update(visible=True), - gr.Row.update(visible=True), - gr.Accordion.update(visible=True)) + return ( + state, + dropdown_update, + gr.Chatbot.update(visible=True), + gr.Textbox.update(visible=True), + gr.Button.update(visible=True), + gr.Row.update(visible=True), + gr.Accordion.update(visible=True), + ) def get_conv_log_filename(): @@ -203,30 +199,31 @@ def get_chat_mode(mode, sql_mode, db_selector) -> ChatScene: elif mode == conversation_types["auto_execute_plugin"] and not db_selector: return ChatScene.ChatExecution else: - return ChatScene.ChatNormal + return ChatScene.ChatNormal -def http_bot(state, mode, sql_mode, db_selector, temperature, max_new_tokens, request: gr.Request): +def http_bot( + state, mode, sql_mode, db_selector, temperature, max_new_tokens, request: gr.Request +): logger.info(f"User message send!{state.conv_id},{sql_mode},{db_selector}") start_tstamp = time.time() - scene:ChatScene = get_chat_mode(mode, sql_mode, db_selector) + scene: ChatScene = get_chat_mode(mode, sql_mode, db_selector) print(f"当前对话模式:{scene.value}") model_name = CFG.LLM_MODEL if ChatScene.ChatWithDb == scene: logger.info("基于DB对话走新的模式！") - chat_param ={ + chat_param = { "chat_session_id": state.conv_id, "db_name": db_selector, - "user_input": state.last_user_input + "user_input": state.last_user_input, } - chat: BaseChat = CHAT_FACTORY.get_implementation(scene.value, **chat_param) + chat: BaseChat = CHAT_FACTORY.get_implementation(scene.value, **chat_param) chat.call() - state.messages[-1][-1] = f"{chat.current_ai_response()}" + state.messages[-1][-1] = f"{chat.current_ai_response()}" yield (state, state.to_gradio_chatbot()) + (enable_btn,) * 5 else: - dbname = db_selector # TODO 这里的请求需要拼接现有知识库, 使得其根据现有知识库作答, 所以prompt需要继续优化 if state.skip_next: @@ -242,7 +239,9 @@ def http_bot(state, mode, sql_mode, db_selector, temperature, max_new_tokens, re # prompt 中添加上下文提示, 根据已有知识对话, 上下文提示是否也应该放在第一轮, 还是每一轮都添加上下文? # 如果用户侧的问题跨度很大, 应该每一轮都加提示。 if db_selector: - new_state.append_message(new_state.roles[0], gen_sqlgen_conversation(dbname) + query) + new_state.append_message( + new_state.roles[0], gen_sqlgen_conversation(dbname) + query + ) new_state.append_message(new_state.roles[1], None) else: new_state.append_message(new_state.roles[0], query) @@ -251,7 +250,6 @@ def http_bot(state, mode, sql_mode, db_selector, temperature, max_new_tokens, re new_state.conv_id = uuid.uuid4().hex state = new_state - prompt = state.get_prompt() skip_echo_len = len(prompt.replace("", " ")) + 1 if mode == conversation_types["default_knownledge"] and not db_selector: @@ -263,16 +261,24 @@ def http_bot(state, mode, sql_mode, db_selector, temperature, max_new_tokens, re skip_echo_len = len(prompt.replace("", " ")) + 1 if mode == conversation_types["custome"] and not db_selector: - persist_dir = os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, vector_store_name["vs_name"] + ".vectordb") + persist_dir = os.path.join( + KNOWLEDGE_UPLOAD_ROOT_PATH, vector_store_name["vs_name"] + ".vectordb" + ) print("向量数据库持久化地址: ", persist_dir) - knowledge_embedding_client = KnowledgeEmbedding(file_path="", model_name=LLM_MODEL_CONFIG["sentence-transforms"], vector_store_config={ "vector_store_name": vector_store_name["vs_name"], - "vector_store_path": KNOWLEDGE_UPLOAD_ROOT_PATH}) + knowledge_embedding_client = KnowledgeEmbedding( + file_path="", + model_name=LLM_MODEL_CONFIG["sentence-transforms"], + vector_store_config={ + "vector_store_name": vector_store_name["vs_name"], + "vector_store_path": KNOWLEDGE_UPLOAD_ROOT_PATH, + }, + ) query = state.messages[-2][1] docs = knowledge_embedding_client.similar_search(query, 1) context = [d.page_content for d in docs] prompt_template = PromptTemplate( template=conv_qa_prompt_template, - input_variables=["context", "question"] + input_variables=["context", "question"], ) result = prompt_template.format(context="\n".join(context), question=query) state.messages[-2][1] = result @@ -285,7 +291,9 @@ def http_bot(state, mode, sql_mode, db_selector, temperature, max_new_tokens, re "prompt": prompt, "temperature": float(temperature), "max_new_tokens": int(max_new_tokens), - "stop": state.sep if state.sep_style == SeparatorStyle.SINGLE else state.sep2, + "stop": state.sep + if state.sep_style == SeparatorStyle.SINGLE + else state.sep2, } logger.info(f"Requert: \n{payload}") @@ -295,8 +303,13 @@ def http_bot(state, mode, sql_mode, db_selector, temperature, max_new_tokens, re try: # Stream output - response = requests.post(urljoin(CFG.MODEL_SERVER, "generate_stream"), - headers=headers, json=payload, stream=True, timeout=20) + response = requests.post( + urljoin(CFG.MODEL_SERVER, "generate_stream"), + headers=headers, + json=payload, + stream=True, + timeout=20, + ) for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"): if chunk: data = json.loads(chunk.decode()) @@ -309,12 +322,23 @@ def http_bot(state, mode, sql_mode, db_selector, temperature, max_new_tokens, re output = data["text"] + f" (error_code: {data['error_code']})" state.messages[-1][-1] = output yield (state, state.to_gradio_chatbot()) + ( - disable_btn, disable_btn, disable_btn, enable_btn, enable_btn) + disable_btn, + disable_btn, + disable_btn, + enable_btn, + enable_btn, + ) return except requests.exceptions.RequestException as e: state.messages[-1][-1] = server_error_msg + f" (error_code: 4)" - yield (state, state.to_gradio_chatbot()) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn) + yield (state, state.to_gradio_chatbot()) + ( + disable_btn, + disable_btn, + disable_btn, + enable_btn, + enable_btn, + ) return state.messages[-1][-1] = state.messages[-1][-1][:-1] @@ -405,10 +429,14 @@ def build_single_model_ui(): interactive=True, label="最大输出Token数", ) + + tabs = gr.Tabs() + with tabs: tab_sql = gr.TabItem("SQL生成与诊断", elem_id="SQL") with tab_sql: + print("tab_sql in...") # TODO A selector to choose database with gr.Row(elem_id="db_selector"): db_selector = gr.Dropdown( @@ -423,8 +451,23 @@ def build_single_model_ui(): sql_vs_setting = gr.Markdown("自动执行模式下, DB-GPT可以具备执行SQL、从网络读取知识自动化存储学习的能力") sql_mode.change(fn=change_sql_mode, inputs=sql_mode, outputs=sql_vs_setting) + tab_plugin = gr.TabItem("插件模式", elem_id="PLUGIN") + with tab_plugin: + print("tab_plugin in...") + with gr.Row(elem_id="plugin_selector"): + # TODO + plugin_selector = gr.Dropdown( + label="请选择插件", + choices=[""" [datadance-ddl-excutor]->use datadance deal the ddl task """, """[file-writer]-file read and write """, """ [image-excutor]-> image build"""], + value="datadance-ddl-excutor", + interactive=True, + show_label=True, + ).style(container=False) + + tab_qa = gr.TabItem("知识问答", elem_id="QA") with tab_qa: + print("tab_qa in...") mode = gr.Radio( ["LLM原生对话", "默认知识库对话", "新增知识库对话"], show_label=False, value="LLM原生对话" ) @@ -483,7 +526,7 @@ def build_single_model_ui(): add_text, [state, textbox], [state, chatbot, textbox] + btn_list ).then( http_bot, - [state, mode, sql_mode, db_selector, temperature, max_output_tokens], + [state, mode, sql_mode, db_selector, temperature, max_output_tokens], [state, chatbot] + btn_list, ) @@ -573,7 +616,6 @@ def knowledge_embedding_store(vs_id, files): ) knowledge_embedding_client.knowledge_embedding() - logger.info("knowledge embedding success") return os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, vs_id, vs_id + ".vectordb") diff --git a/pilot/vector_store/connector.py b/pilot/vector_store/connector.py index 06fad00f2..3ff473f1e 100644 --- a/pilot/vector_store/connector.py +++ b/pilot/vector_store/connector.py @@ -1,7 +1,7 @@ from pilot.vector_store.chroma_store import ChromaStore -from pilot.vector_store.milvus_store import MilvusStore +# from pilot.vector_store.milvus_store import MilvusStore -connector = {"Chroma": ChromaStore, "Milvus": MilvusStore} +connector = {"Chroma": ChromaStore, "Milvus": None} class VectorStoreConnector: diff --git a/requirements.txt b/requirements.txt index 19d8ca34e..9595854bb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -54,7 +54,7 @@ gTTS==2.3.1 langchain nltk python-dotenv==1.0.0 -pymilvus==2.2.1 +# pymilvus==2.2.1 vcrpy chromadb markdown2