gui.py

import gradio as gr
import os
import shutil
from function_calling_agent import FunctionAgent
from llm_factory import llm_factory
from typing import List, Dict, Any
from dotenv import load_dotenv
import json
from data_loader import extract_metadata
from excel_preparations import ExcelPreparations

load_dotenv()

llm = llm_factory(os.getenv("MODEL_NAME", ""))
calling_agent = FunctionAgent(llm)

# Define the tmp folder
tmp_folder = os.path.join(os.getcwd(), "tmp")

# Ensure tmp folder exists
if not os.path.exists(tmp_folder):
    os.makedirs(tmp_folder)

# Function to clean up the tmp folder after program ends
def cleanup():
    json_path = os.path.join(tmp_folder, "file_mapping.json")
    if os.path.exists(json_path):
        os.remove(json_path)
    if os.path.exists(tmp_folder):
        shutil.rmtree(tmp_folder)
    os.mkdir(tmp_folder)

with gr.Blocks(theme=gr.themes.Ocean(), title="GEDA") as demo:
    gr.HTML("<h1 style='text-align: center;'>GEDA</h1>")
    chatbot: gr.Chatbot = gr.Chatbot(type="messages", height="65vh")
    msg: gr.Textbox = gr.Textbox()
    send_button: gr.Button = gr.Button("Send", variant="primary")
    file_upload: gr.File = gr.File(file_types=[".xls", ".xlsx"], file_count="multiple")
    clear: gr.Button = gr.Button("Clear Chat")

    def handle_file_upload(files):
        yield gr.update(interactive = False), gr.update(interactive = False)
        tmp_files = []
        for file in files:
            src_path = file.name
            dst_path = os.path.join(tmp_folder, os.path.basename(file.name))
            
            # Check if source and destination are the same
            if os.path.abspath(src_path) != os.path.abspath(dst_path):
                shutil.copy(src_path, dst_path)
            tmp_files.append(dst_path.split("/")[-1])

            # Create a dictionary to map filename to its path
            file_mapping = {os.path.basename(file.name): file for file in files}

            # Save the dictionary as a JSON file in the tmp folder
            json_path = os.path.join(tmp_folder, "file_mapping.json")
            with open(json_path, "w") as json_file:
                json.dump(file_mapping, json_file)

        excel_preparation = ExcelPreparations()
        data_frames, info_texts = excel_preparation.read_excel(tmp_files)
        extract_metadata(llm, tmp_files, data_frames, info_texts)
        yield gr.update(interactive = True), gr.update(interactive = True)
        

    def user(
        user_message: str, history: List[Dict[str, Any]]
    ) -> (str, List[Dict[str, Any]]):
        return "", history + [{"role": "user", "content": user_message}]

    def bot(history: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        history.append({"role": "assistant", "content": ""})
        for x in calling_agent(history[:-1]):
            if type(x) == str:  # allow llm answer stream
                history[-1]["content"] += x
            else:  # allow gradio blocks in chat
                history[-1]["content"] = x
            yield history

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    send_button.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )

    clear.click(lambda: None, None, chatbot, queue=False)

    file_upload.upload(
        handle_file_upload, inputs=[file_upload], outputs=[send_button, msg])

def main_gui() -> None:
    demo.launch(favicon_path="./favicon.png")

if __name__ == "__main__":
    try:
        main_gui()
    finally:
        cleanup()  # Delete the tmp folder and file mapping JSON after the function is killed