from langchain_experimental.agents import create_csv_agent
from langchain.agents.agent_types import AgentType
from langchain.llms import OpenAI
import streamlit as st

def main():
    st.set_page_config(page_title="Ask your CSV")
    st.header("Ask your CSV 📈")
    
    # API key input
    api_key = st.text_input("Enter your OpenAI API key:", type="password")
    
    # File uploader
    csv_file = st.file_uploader("Upload a CSV file", type="csv")
    
    # Process only if both API key and file are provided
    if api_key and csv_file is not None:
        # Create agent with user-provided API key
        try:
            agent = create_csv_agent(
                OpenAI(temperature=0, api_key=api_key), 
                csv_file, 
                verbose=True,
                agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
                allow_dangerous_code=True
            )
            
            # User question input
            user_question = st.text_input("Ask a question about your CSV: ")
            
            # Process question if provided
            if user_question and user_question != "":
                with st.spinner(text="In progress..."):
                    st.write(agent.run(user_question))
        except Exception as e:
            st.error(f"Error: {str(e)}")
    elif not api_key and csv_file is not None:
        st.info("Please enter your OpenAI API key to proceed.")
    elif api_key and csv_file is None:
        st.info("Please upload a CSV file to proceed.")

if __name__ == "__main__":
    main()