Update app.py

app.py

@@ -2,9 +2,10 @@ import streamlit as st
 import pandas as pd
 import matplotlib.pyplot as plt
 import seaborn as sns
+import os
 from io import StringIO
 import openai
-import os
+from groq import Groq
 
 # Groq API key from secrets
 GROQ_API_KEY = os.getenv("HUGGINGFACE_KEY")
@@ -30,27 +31,18 @@ def generate_graph(data, query):
     """Generate a graph based on user query."""
     try:
         fig, ax = plt.subplots(figsize=(10, 6))
-        
-        if "correlation" in query.lower():
-            # Correlation matrix
-            sns.heatmap(data.corr(), annot=True, cmap="coolwarm", ax=ax)
-            st.pyplot(fig)
-        elif "histogram" in query.lower():
-            # Histogram
-            column = st.selectbox("Select a column for the histogram", data.columns)
-            sns.histplot(data[column], kde=True, ax=ax)
-            st.pyplot(fig)
-        elif "bar" in query.lower() and "country" in query.lower():
-            # Bar chart for countries and units sold
-            if 'country' in data.columns and 'units sold' in data.columns:
-                country_data = data[['country', 'units sold']].groupby('country').sum().reset_index()
-                sns.barplot(x='country', y='units sold', data=country_data, ax=ax, color='skyblue')
+
+        if "bar" in query.lower() and "gross sales" in query.lower():
+            # Bar chart for countries and gross sales
+            if 'country' in data.columns and 'gross_sales' in data.columns:
+                country_data = data[['country', 'gross_sales']].groupby('country').sum().reset_index()
+                sns.barplot(x='country', y='gross_sales', data=country_data, ax=ax, color='skyblue')
                 ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha='right')
                 st.pyplot(fig)
             else:
-                st.error("The dataset must contain 'country' and 'units sold' columns.")
+                st.error("The dataset must contain 'country' and 'gross_sales' columns.")
         else:
-            st.error("Unsupported graph type. Try asking for a correlation matrix, histogram, or bar chart.")
+            st.error("Unsupported graph type. Try asking for a bar chart.")
     except Exception as e:
         st.error(f"Error generating graph: {e}")
 
@@ -102,7 +94,7 @@ def main():
             st.write("Dataset preview:", data.head())
 
             # User input for graph generation
-            query = st.text_input("Enter your query (e.g., 'Generate a bar chart for countries and units sold')")
+            query = st.text_input("Enter your query (e.g., 'Generate a bar chart for countries and gross sales')")
 
             if query:
                 # Generate the graph based on the query