Code Editor for app.py

import streamlit as st
import pandas as pd
import os
from openai import OpenAI
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Initialize OpenAI client
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

st.title("CSV Data Analysis Chat")

# File upload
st.write("Upload your CSV file (max 10MB)")
uploaded_file = st.file_uploader("Choose a CSV file", type="csv", accept_multiple_files=False)

# Initialize session state for messages if not exists
if "messages" not in st.session_state:
    st.session_state.messages = []

# Initialize session state for dataframe if not exists
if "df" not in st.session_state:
    st.session_state.df = None

if uploaded_file is not None:
    # Check file size
    if uploaded_file.size > 10 * 1024 * 1024:  # 10MB in bytes
        st.error("File is too large! Please upload a file smaller than 10MB.")
    else:
        # Read and display the CSV
        st.session_state.df = pd.read_csv(uploaded_file)
        st.write("Dataset Preview:")
        st.dataframe(st.session_state.df.head())
        
        # Display dataset info
        st.write("Dataset Info:")
        buffer = []
        st.session_state.df.info(buf=buffer)
        st.text("\n".join(buffer))

# Chat interface
if st.session_state.df is not None:
    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Chat input
    if prompt := st.chat_input("Ask questions about your data"):
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        # Prepare context about the dataset
        df_info = f"""
        Dataset Information:
        - Columns: {', '.join(st.session_state.df.columns.tolist())}
        - Shape: {st.session_state.df.shape}
        - Data sample: {st.session_state.df.head(2).to_string()}
        """

        # Generate OpenAI API response
        with st.chat_message("assistant"):
            messages = [
                {"role": "system", "content": f"You are a data analysis assistant. You'll receive questions about a dataset with the following information:\n{df_info}\nProvide clear and concise answers based on the data structure provided."},
                *st.session_state.messages
            ]
            
            response = client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=[{"role": m["role"], "content": m["content"]} for m in messages],
                stream=True
            )

            # Stream the response
            message = st.write_stream(response)
            
        # Add assistant response to chat history
        st.session_state.messages.append({"role": "assistant", "content": message})

else:
    st.write("Please upload a CSV file to start the conversation.")

# Add a note about the OpenAI API key requirement
st.sidebar.warning("""
This app requires an OpenAI API key to function. 
Please set your OPENAI_API_KEY environment variable or create a .env file with your key.
""")