Update app.py

app.py

@@ -66,15 +66,27 @@ 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)
+                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.")
         else:
-            st.error("Unsupported graph type. Try asking for a correlation matrix or histogram.")
+            st.error("Unsupported graph type. Try asking for a correlation matrix, histogram, or bar chart.")
     except Exception as e:
         st.error(f"Error generating graph: {e}")