import streamlit as st
from langchain.llms import OpenAI
from langchain_experimental.agents.agent_toolkits import create_csv_agent
from langchain.agents import AgentType
from langchain.utilities import PythonREPL

def main():
  st.set_page_config(page_title="Talk to your Data - AI Assistant", page_icon=":robot:")
  st.title('DataCompanion AI Assistant :robot_face:')
  st.subheader(':blue[by Venkat Reddy] :sunglasses:')

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

  if input_csv:
    st.write("Ask any question to your data")
    user_question=st.text_input("Ask a question on your data")
    python_repl=PythonREPL()
    agent=create_csv_agent(OpenAI(temperature=0),
                           input_csv,
                           verbose=True,
                           agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
                           tool=python_repl)
    result=agent.run(user_question)
    st.write(result)


if __name__ == "__main__":
  main()