kkAIGC/SQLBot

1import gradio as gr2import openai3from src.plan.content_moderation import *4from src.plan.get_table_info import *5from src.plan.text_to_sql import *6from src.analysis.sql_observation import *7from src.generate.sql_generation import *8from src.study.answer_evaluation import *9from src.get_api_key import *10from src.get_completion_from_messages import *11from constants.constants import constants12 13# 解包constant中的常量14locals().update(constants)15 16with gr.Blocks(theme='gradio/soft') as demo:17    with gr.Row():18        with gr.Column(scale=4):19            gr.Markdown('''# SQLBot - 一个基于LLMs的SQL学习和工作平台20            **本项目已开源，喜欢该项目请去GitHub点个star✨：** [SQLBot](https://github.com/q734550709/SQLBot/tree/main)''')21        with gr.Column(scale=1):22            apikey_text = gr.Textbox(type='password',label='请输入api_key:')23    with gr.Tab("通过自然语言生成SQL代码"):24        with gr.Row():25            with gr.Column(scale=4):26                chatbot = gr.Chatbot(bubble_full_width=False,show_label=True,show_copy_button=True)27                input_text = gr.Textbox(lines=6,label="请输入查询语句",show_label=True,show_copy_button=True)28                user_button = gr.Button(value='确认输入')29                clear = gr.ClearButton([input_text, chatbot],value="清除对话")30            with gr.Column(scale=1):31                with gr.Accordion(label='模型设置',open=False):32                    model_select = gr.Dropdown(33                                choices=["gpt-3.5-turbo-16k","gpt-3.5-turbo"],34                                value='gpt-3.5-turbo',35                                label='选择模型'36                                )37                    temperature = gr.Slider(minimum=0, maximum=1,value=0,label='temperature')38                    maxtoken = gr.Slider(minimum=0, maximum=3000, value=300,label='max_token')39                hive_select = gr.Dropdown(40                            choices=["hive","presto","mysql"],41                            value='hive',42                            label='选择SQL语言类型'43                            )44                is_current_table = gr.Checkbox(label='根据现有库表查询',value=True)45                data_definition = gr.Textbox(lines=2,46                                                   max_lines=6,47                                                   value=data_scope_definition,48                                                   label="输入涉及的字段定义")49                gr.Examples(50                        examples=question_examples,51                        inputs=input_text,52                        label='可选例子'53                        )54 55        with gr.Row():56            with gr.Accordion(label="相关库表信息：",open=False):57                data_df = gr.Dataframe(58                    headers=database_columns,59                    type="array",60                    row_count=9,61                    col_count=(4, "fixed"),62                    max_rows="infinite",63                    value=database_datalist64                )65 66        user_button.click(fn=process_user_message,67                       inputs=[input_text,68                               chatbot,69                               model_select,70                               temperature,71                               maxtoken,72                               hive_select,73                               is_current_table,74                               data_definition,75                               data_df76                              ],77                       outputs=[input_text,chatbot])78 79    with gr.Tab("对SQL代码进行文字说明"):80        with gr.Row():81 82            with gr.Column(scale=2):83                sql_text = gr.Textbox(lines=21,label="输入SQL代码")84                sql_button = gr.Button('确认输入')85 86            with gr.Column(scale=2):87                output_text = gr.Textbox(lines=24,label="输出",88                                         show_label=True,show_copy_button=True,89                                        )90            with gr.Column(scale=1):91                with gr.Accordion(label='模型设置',open=False):92                    model_select = gr.Dropdown(93                                choices=["gpt-3.5-turbo-16k","gpt-3.5-turbo"],94                                value='gpt-3.5-turbo',95                                label='选择模型'96                                )97                    temperature = gr.Slider(minimum=0, maximum=1,value=0,label='temperature')98                    maxtoken = gr.Slider(minimum=0, maximum=3000, value=300,label='max_token')99                hive_select = gr.Dropdown(100                            choices=["hive","presto","mysql"],101                            value='hive',102                            label='选择SQL语言类型'103                            )104                translate_select = gr.Checkbox(label='需要详细解释')105 106            sql_button.click(fn=function_select,107                           inputs=[sql_text,108                                   model_select,109                                   temperature,110                                   maxtoken,111                                   translate_select112                                  ],113                           outputs=output_text)114 115    with gr.Tab("输入prompt模板生成SQL代码"):116        with gr.Row():117            with gr.Column(scale=2):118                merge_text = gr.Textbox(lines=21,label="prompt模板")119                generation_button = gr.Button('确认输入')120            with gr.Column(scale=2):121                output_text = gr.Textbox(lines=24,label="SQL结果",show_label=True,show_copy_button=True)122            with gr.Column(scale=1):123                with gr.Accordion(label='模型设置',open=False):124                    model_select = gr.Dropdown(125                                choices=["gpt-3.5-turbo-16k","gpt-3.5-turbo"],126                                value='gpt-3.5-turbo',127                                label='选择模型'128                                )129                    temperature = gr.Slider(minimum=0, maximum=1,value=0,label='temperature')130                    maxtoken = gr.Slider(minimum=0, maximum=3000, value=300,label='max_token')131                hive_select = gr.Dropdown(132                            choices=["hive","presto","mysql"],133                            value='hive',134                            label='选择SQL语言类型'135                            )136 137                connect_input = gr.Textbox(lines=2,max_lines=4,label="输入链接词模板")138                select_input = gr.Textbox(lines=2,max_lines=4,label="输入选择词模板")139                filter_input = gr.Textbox(lines=2,max_lines=4,label="输入筛选词模板")140                groupby_input = gr.Textbox(lines=2,max_lines=4,label="输入分组词模板")141 142        gr.Examples(143                    examples=model_example_list,144                    inputs=[connect_input,select_input,filter_input,groupby_input],145                    label='可选例子'146                    )147        connect_input.change(fn=merge_textbox,148                           inputs=[connect_input,149                                   select_input,150                                   filter_input,151                                   groupby_input],152                           outputs=merge_text)153        select_input.change(fn=merge_textbox,154                           inputs=[connect_input,155                                   select_input,156                                   filter_input,157                                   groupby_input],158                           outputs=merge_text)159        filter_input.change(fn=merge_textbox,160                           inputs=[connect_input,161                                   select_input,162                                   filter_input,163                                   groupby_input],164                           outputs=merge_text)165        groupby_input.change(fn=merge_textbox,166                           inputs=[connect_input,167                                   select_input,168                                   filter_input,169                                   groupby_input],170                           outputs=merge_text)171 172        generation_button.click(fn=sql_generation,173                           inputs=[merge_text,174                                   model_select,175                                   temperature,176                                   maxtoken,177                                   hive_select],178                           outputs=output_text)179 180    with gr.Tab("SQL代码练习题"):181        with gr.Row():182            with gr.Column(scale=2):183                title_url_text = gr.Markdown()184                question_text = gr.Textbox(lines=5,max_lines=10,label="题目")185                example_text = gr.Textbox(lines=8,max_lines=10,label="示例")186 187            with gr.Column(scale=4):188                chatbot = gr.Chatbot(bubble_full_width=False,show_label=True,show_copy_button=True)189                input_text = gr.Textbox(lines=6,label="请输入SQL语句",show_label=True,show_copy_button=True)190                user_button = gr.Button(value='确认输入')191                clear = gr.ClearButton([input_text, chatbot],value="清除对话")192 193            with gr.Column(scale=1):194                with gr.Accordion(label='模型设置',open=False):195                    model_select = gr.Dropdown(196                                choices=["gpt-3.5-turbo-16k","gpt-3.5-turbo"],197                                value='gpt-3.5-turbo',198                                label='选择模型'199                                )200                    temperature = gr.Slider(minimum=0, maximum=1,value=0,label='temperature')201                    maxtoken = gr.Slider(minimum=0, maximum=3000, value=300,label='max_token')202                difficulty_select = gr.Dropdown(203                            choices=["简单","中等","困难"],204                            label='选择题目难度'205                            )206                with gr.Accordion(label="参考答案：",open=False):207                    answer_text = gr.Textbox(lines=10,label="答案")208                    answer_explain_text = gr.Markdown()209 210            difficulty_select.change(fn=question_choice,211                               inputs=[difficulty_select],212                               outputs=[title_url_text,213                                        question_text,214                                        example_text,215                                        answer_text,216                                        answer_explain_text])217 218            user_button.click(fn=answer_evaluation,219                              inputs=[input_text,220                                      chatbot,221                                      question_text,222                                      answer_text,223                                      model_select,224                                      temperature,225                                      maxtoken226                                     ],227                              outputs=[input_text,chatbot]228                             )229 230    apikey_text.change(fn=get_api_key,inputs= apikey_text)231 232 233if __name__ == "__main__":234    demo.launch()235