from pandasai.llm.local_llm import LocalLLM 
import streamlit as st
import pandas as pd
from pandasai import Agent

model = LocalLLM(
    api_base="http://localhost:11434/v1",
    model="llama3"
)

st.title("Data analysis with PandasAI Agent")

# uploaded_file = st.sidebar.file_uploader(
#     "Upload a CSV file",
#     type=["csv"]
# )

# if uploaded_file is not None:
#     data = pd.read_csv(uploaded_file)
#     st.write(data.head(3))

#     agent = Agent(data, config={"llm": model})
#     prompt = st.text_input("Enter your prompt:")

#     if st.button("Generate"):
#         if prompt:
#             with st.spinner("Generating response..."):
#                 st.write(agent.chat(prompt))

data = pd.read_csv('ai4i2020.csv')
st.write(data.head(3))

agent = Agent(data, config={"llm": model})
prompt = st.text_input("Enter your prompt:")

if st.button("Generate"):
    if prompt:
        with st.spinner("Generating response..."):
            st.write(agent.chat(prompt))