table_load.py

"""functions"""

def load_file():
    return "Loading..."

def load_xlsx(name):
    import pandas as pd
    
    xls_file = rf'{name}'
    data = pd.read_excel(xls_file)  
    return data

def table_loader(table_file, open_ai_key):
    import os
    from langchain.llms import OpenAI
    from langchain.agents import create_pandas_dataframe_agent    
    from pandas import read_csv
    
    global agent
    if open_ai_key is not None:
        os.environ['OPENAI_API_KEY'] = open_ai_key
    else:
        return "Enter API"
    
    if table_file.name.endswith('.xlsx') or table_file.name.endswith('.xls'):
        data = load_xlsx(table_file.name)
        agent = create_pandas_dataframe_agent(OpenAI(temperature=0), data)
        return "Ready!"
    elif table_file.name.endswith('.csv'):
        data = read_csv(table_file.name)
        agent = create_pandas_dataframe_agent(OpenAI(temperature=0), data)
        return "Ready!"
    else:
        return "Wrong file format! Upload excel file or csv!"
    
def run(query):
    from langchain.callbacks import get_openai_callback
    
    with get_openai_callback() as cb:
        response = (agent.run(query))
        costs = (f"Total Cost (USD): ${cb.total_cost}")
        output = f'{response} \n {costs}'
        return output
    
def respond(message, chat_history):
    import time
    
    bot_message = run(message)
    chat_history.append((message, bot_message))
    time.sleep(0.5)
    return "", chat_history