table_load.py

"""functions related to tables loading"""

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!"