import streamlit as st
import pandas as pd
from pandasai import SmartDataframe
from pandasai.llm import GooglePalm
import textwrap

# Set up your API key
GOOGLE_API_KEY = "AIzaSyDVITagNdKqvBkWngu9Q6Nywy9WkI4zpak"

# Initialize the LLM
llm = GooglePalm(api_key=GOOGLE_API_KEY)

# Title and description
st.title("Smart Dataframe Explorer")
st.write("Upload your CSV file, explore the data, and ask questions.")

# File uploader
uploaded_file = st.file_uploader("Upload your CSV file", type=["csv"])

if uploaded_file is not None:
    # Load CSV file into a SmartDataframe
    df = pd.read_csv(uploaded_file)
    smart_df = SmartDataframe(df, config={"llm": llm})

    # Display the first three rows by default
    st.subheader("First Three Rows")
    st.write(smart_df.head(3))

    # Text bar for asking questions
    question = st.text_input("Ask a question about your data")

    if question:
        # Get the LLM response
        st.subheader("Answer")
        answer = smart_df.chat(question)
        st.write(textwrap.indent(answer, "> ", predicate=lambda _: True))