import streamlit as st
import os
os.environ['OPENAI_API_KEY'] = "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"

from langchain.llms.openai import OpenAI
from langchain.agents.agent_types import AgentType
#from langchain.agents import create_csv_agent
from langchain_experimental.agents import create_csv_agent
import time

def main():
    st.title("Personal Assistant")

    uploaded_file = st.file_uploader("Upload CSV file", type=["csv"])
    if uploaded_file is None:
        st.warning("Please upload a CSV file.")
        st.stop()  # Stop execution if no file uploaded

    llm = OpenAI(temperature=0)
    agent = create_csv_agent(
        llm,
        uploaded_file,
        verbose=False,
        agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    )

    query = st.text_input("What would you like to know?")
    if st.button("Ask"):
        if query.strip() == "":
            st.warning("Please enter a query.")
        else:
            start = time.time()
            answer = agent.run(query)
            end = time.time()
            st.write(answer)
            st.write(f"Answer (took {round(end - start, 2)} s.)")

if __name__ == "__main__":
    main()