app.py

import os.path
import shutil
import sys
import traceback
import gradio as gr
import json
import random
import time
import uuid
import pandas as pd
from conversation import Conversation
from cache.oss import init_dir, init_oss, upload_oss_file, get_download_link
from cache.cache import data_cache
from prompt_engineering.prompts import *
from kernel import CodeKernel
import yaml
# from config import conv_model, streaming
import atexit
import signal

class app():
    def __init__(self, config_path='config.yaml'):
        print("Load config: ", config_path)
        self.config = yaml.load(open(config_path, 'r'), Loader=yaml.FullLoader)
        self.conv = Conversation(self.config)

        self.conv.programmer.messages = [
            {
                "role": "system",
                "content": SYSTEM_PROMPT
            }
        ]

    def open_board(self):
        return self.conv.show_data()

    def cache_file(self,file):
        try:
            self.oss_file_dir, self.local_cache_dir = init_dir(self.config['cache_dir'])
            init_oss(self.config['oss_endpoint'], self.config['oss_access_key_id'], self.config['oss_access_secret'], self.config['oss_bucket_name'], self.config['expired_time'])
            self.conv.oss_dir, self.conv.local_cache_dir = self.oss_file_dir, self.local_cache_dir
            file_info = upload_oss_file(self.oss_file_dir, file)
            local_file_path = os.path.join(self.local_cache_dir, file_info['file_name'])
            shutil.copy2(file, local_file_path)
            self.conv.add_data(file)
            gen_info = self.conv.my_data_cache.get_description()
            self.conv.programmer.messages[0]["content"] = (PROGRAMMER_PROMPT.format(working_path=self.local_cache_dir) +
                                                f"\nNow, user uploads the datasets in {local_file_path}\n, and here is the general information of the dataset:\n {gen_info}")
            if self.conv.retrieval:
                self.conv.programmer.messages[0]["content"] += KNOWLEDGE_INTEGRATION_SYSTEM
            print(self.conv.programmer.messages[0]["content"])

            # self.conv.run(function_lib)
            # self.conv.messages.append({"role": "assistant", "content": response})
            # chat_history.append((message, response))
        except Exception as e:
            traceback.print_exc()
            print(f"Uploaded file error: {e}")

    # def show_data(self):
    #     return self.conv.show_data()

    def rendering_code(self):
        return self.conv.rendering_code()

    # def report_test(self):
    #     with open("/Users/stephensun/Desktop/pypro/dsagent_ci/cache/conv_cache/2c23cf4b-31a8-4aeb-8689-13b9fbdfb7fc-2024-04-29/programmer_msg.json", "r") as f:
    #         msg = json.load(f)
    #         print(msg)
    #     self.conv.programmer.messages = msg
    #     down_path = self.conv.document_generation()
    #     return [gr.Button(visible=False), gr.DownloadButton(label=f"Download Report", value=down_path, visible=True)]

    def generate_report(self):
        down_path = self.conv.document_generation()
        return [gr.Button(visible=False), gr.DownloadButton(label=f"Download Report", value=down_path, visible=True)]

    def export_code(self):
        down_path = self.conv.export_code()
        return [gr.Button(visible=False), gr.DownloadButton(label=f"Download Notebook", value=down_path, visible=True)]

    def down_report(self):
        return [gr.Button(visible=True), gr.DownloadButton(visible=False)]
    def down_notebook(self):
        return [gr.Button(visible=True), gr.DownloadButton(visible=False)]

    # def human_loop(self, code):
    #     result = self.conv.run_workflow(code)
    #     self.conv.programmer.messages.append({"role": "assistant", "content": response})
    #     return response
    def call_llm(self, message, chat_history, code=None):
        '''
        :param message: input of user
        :param chat_history: the history of whole dialogue. But this variable was only used to show in the UI (not record in LLM)
        :return: chat_history in the Chatbot component
        '''
        if not code:
            self.conv.programmer.messages.append({"role": "user", "content": message})
        else:
            message = code
        response = self.conv.run_workflow(function_lib=None,code=code)
        chat_history.append((message, response))
        self.conv.chat_history = chat_history
        # chat_history = self.conv.messages
        # print(chat_history)
        return "", chat_history  # return "" to let msg box be clear

    # def load_chat_history(self):
    #     # filtered_messages = [msg for msg in self.conv.programmer.messages if isinstance(msg, dict) and msg.get("role") in ["user", "assistant"]]
    #     # # [{"role": "user", "content": "qqq"},{"role": "assistant", "content": "aaa"}]
    #     # chat_history = [(filtered_messages[i].get('content'), filtered_messages[i + 1].get('content')) for i in range(0, len(filtered_messages), 2)]
    #     return self.conv.chat_history

    def save_dialogue(self, chat_history):
        self.conv.save_conv()
        with open(os.path.join(self.local_cache_dir,'system_dialogue.json'), 'w') as f:
            json.dump(chat_history, f, indent=4)
            f.close()
        print(f"Dialogue saved in {os.path.join(self.local_cache_dir,'system_dialogue.json')}.")

    def load_dialogue(self, json_file):
        with open(json_file, 'r') as f:
            chat_history = json.load(f)
            f.close()
        self.conv.chat_history = chat_history
        return chat_history

    def clear_all(self, message, chat_history):
        self.conv.clear()
        return "", []

my_app = app()

with gr.Blocks(theme=gr.themes.Soft(), css='front_end/css.css', js='front_end/javascript.js') as demo:
    # get history for chatbot.
    #history = my_app.load_dialogue("system_dialogue.json") # load history in previous conversation by indicate system_dialogue.json
    chatbot = gr.Chatbot(value=my_app.conv.chat_history, height=600, label="LAMBDA", avatar_images=["front_end/user.jpg", "front_end/lambda.jpg"], show_copy_button=True)
    with gr.Group():
        with gr.Row():
            upload_btn = gr.UploadButton(label="Upload Data", file_types=["csv", "xlsx"], scale=1)
            # datasets = gr.File(file_types=["csv", "xlsx"], elem_id="file_upload",height=80)
            msg = gr.Textbox(show_label=False, placeholder="Sent message to LLM", scale=6)
            submit = gr.Button("Submit", scale=1)
    with gr.Row():
        board = gr.Button(value="Show/Update DataFrame", elem_id="df_btn", elem_classes="df_btn")
        export_notebook = gr.Button(value="Notebook")
        down_notebook = gr.DownloadButton("Download Notebook", visible=False)
        generate_report = gr.Button(value="Generate Report")
        down_report = gr.DownloadButton("Download Report",visible=False)

        edit = gr.Button(value="Edit Code",elem_id="ed_btn", elem_classes="ed_btn")
        save = gr.Button(value="Save Dialogue")
        clear = gr.ClearButton(value="Clear All")

    with gr.Group():
        with gr.Row(visible=False, elem_id="ed",elem_classes="ed"):
            code = gr.Code(label="Code",scale=6)
            code_btn = gr.Button("Submit Code",scale=1)
    code_btn.click(fn=my_app.call_llm, inputs=[msg, chatbot, code], outputs=[msg, chatbot])


    df = gr.Dataframe(visible=False,elem_id="df",elem_classes="df")

    upload_btn.upload(fn=my_app.cache_file, inputs=upload_btn)
    if my_app.config['streaming']:
        def chat_streaming(message, chat_history, code=None):
            if not code:
                my_app.conv.programmer.messages.append({"role": "user", "content": message})
            else:
                message = code
            #my_app.conv.stream_workflow(function_lib=None, code=code)
            return "", chat_history + [[message, None]]

        msg.submit(chat_streaming, [msg, chatbot], [msg, chatbot], queue=False).then(
            my_app.conv.stream_workflow, chatbot, chatbot
        )
        submit.click(chat_streaming, [msg, chatbot], [msg, chatbot], queue=False).then(
            my_app.conv.stream_workflow, chatbot, chatbot
        )
    else:
        msg.submit(my_app.call_llm, [msg, chatbot], [msg, chatbot])
        submit.click(my_app.call_llm, [msg, chatbot], [msg, chatbot])
    board.click(my_app.open_board, inputs=[], outputs=df)
    edit.click(my_app.rendering_code, inputs=None, outputs=code)
    export_notebook.click(my_app.export_code, inputs=None, outputs=[export_notebook, down_notebook])
    down_notebook.click(my_app.down_notebook, inputs=None, outputs=[export_notebook, down_notebook])
    generate_report.click(my_app.generate_report,inputs=None,outputs=[generate_report,down_report])
    down_report.click(my_app.down_report,inputs=None,outputs=[generate_report,down_report])
    save.click(my_app.save_dialogue, inputs=chatbot)
    clear.click(fn=my_app.clear_all, inputs=[msg, chatbot], outputs=[msg, chatbot])
    #demo.load(fn=my_app.load_chat_history, inputs=None, outputs=chatbot)

if __name__ == '__main__':
    demo.launch(server_name="0.0.0.0",server_port=8000, debug=True, share=True)