csv-analyst / app.py
namaai's picture
Update app.py
5b3859c verified
import os
import streamlit as st
import pandas as pd
from pandasai import SmartDataframe
# from pandasai.callbacks import BaseCallback
from pandasai.responses.response_parser import ResponseParser
from groq import Groq
from langchain_groq import ChatGroq
os.environ["GROQ_API_KEY"] = "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"
client = Groq()
model = ChatGroq(temperature=0, model_name="llama-3.1-70b-versatile")
# class StreamlitCallback(BaseCallback):
# def __init__(self, container) -> None:
# """Initialize callback handler."""
# self.container = container
# def on_code(self, response: str):
# self.container.code(response)
class StreamlitResponse(ResponseParser):
def __init__(self, context) -> None:
super().__init__(context)
def format_dataframe(self, result):
st.dataframe(result["value"])
return
def format_plot(self, result):
st.image(result["value"])
return
def format_other(self, result):
st.write(result["value"])
return
def read_csv_with_encoding(file):
encodings = ['utf-8', 'latin1', 'iso-8859-1']
for encoding in encodings:
try:
file.seek(0) # Reset file pointer to the beginning
df = pd.read_csv(file, encoding=encoding)
if df.empty:
st.error("The CSV file is empty.")
return None
return df
except UnicodeDecodeError:
st.error(f"UnicodeDecodeError with encoding {encoding}. Trying next encoding.")
except pd.errors.EmptyDataError:
st.error("The CSV file is empty or malformed.")
return None
except Exception as e:
st.error(f"An error occurred: {e}")
return None
return None
st.write("# Chat with Dataset")
uploaded_file = st.file_uploader("Upload a CSV file", type="csv")
if uploaded_file is not None:
df = read_csv_with_encoding(uploaded_file)
with st.expander("🔎 Dataframe Preview"):
st.write(df.tail(3))
query = st.text_area("🗣️ Chat with Dataframe")
container = st.container()
if query:
llm = model
query_engine = SmartDataframe(
df,
config={
"llm": llm,
"response_parser": StreamlitResponse,
},
)
answer = query_engine.chat(query)
st.write(answer)
else:
st.write("Please upload a CSV file to start.")