Update app.py

app.py

@@ -86,27 +86,31 @@ if st.session_state.df is not None and st.session_state.show_preview:
     st.subheader("📂 Dataset Preview")
     st.dataframe(st.session_state.df.head())
 
-"""# Ask GPT-4o for Visualization Suggestions
+
+
+
 def ask_gpt4o_for_visualization(query, df, llm):
     columns = ', '.join(df.columns)
     prompt = f"""
-    Analyze the query and suggest the best visualization.
+    Analyze the query and suggest one or more relevant visualizations.
     Query: "{query}"
     Available Columns: {columns}
-    Respond in this JSON format:
-    {{
-      "chart_type": "bar/box/line/scatter",
-      "x_axis": "column_name",
-      "y_axis": "column_name",
-      "group_by": "optional_column_name"
-    }}
+    Respond in this JSON format (as a list if multiple suggestions):
+    [
+      {{
+        "chart_type": "bar/box/line/scatter",
+        "x_axis": "column_name",
+        "y_axis": "column_name",
+        "group_by": "optional_column_name"
+      }}
+    ]
     """
     response = llm.generate(prompt)
     try:
         return json.loads(response)
     except json.JSONDecodeError:
         st.error("⚠️ GPT-4o failed to generate a valid suggestion.")
-        return None"""
+        return None
 
 def add_stats_to_figure(fig, df, y_axis, chart_type):
     """
@@ -241,6 +245,56 @@ def generate_visualization(suggestion, df):
         st.error(f"⚠️ Failed to generate visualization: {e}")
         return None
 
+def generate_multiple_visualizations(suggestions, df):
+    """
+    Generates one or more visualizations based on GPT-4o's suggestions.
+    Handles both single and multiple suggestions.
+    """
+    visualizations = []
+
+    for suggestion in suggestions:
+        fig = generate_visualization(suggestion, df)
+        if fig:
+            # Apply chart-specific statistics
+            fig = add_stats_to_figure(fig, df, suggestion["y_axis"], suggestion["chart_type"])
+            visualizations.append(fig)
+
+    if not visualizations and suggestions:
+        st.warning("⚠️ No valid visualization found. Displaying the most relevant one.")
+        best_suggestion = suggestions[0]
+        fig = generate_visualization(best_suggestion, df)
+        fig = add_stats_to_figure(fig, df, best_suggestion["y_axis"], best_suggestion["chart_type"])
+        visualizations.append(fig)
+
+    return visualizations
+
+
+def handle_visualization_suggestions(suggestions, df):
+    """
+    Determines whether to generate a single or multiple visualizations.
+    """
+    visualizations = []
+
+    # If multiple suggestions, generate multiple plots
+    if isinstance(suggestions, list) and len(suggestions) > 1:
+        visualizations = generate_multiple_visualizations(suggestions, df)
+    
+    # If only one suggestion, generate a single plot
+    elif isinstance(suggestions, dict) or (isinstance(suggestions, list) and len(suggestions) == 1):
+        suggestion = suggestions[0] if isinstance(suggestions, list) else suggestions
+        fig = generate_visualization(suggestion, df)
+        if fig:
+            visualizations.append(fig)
+    
+    # Handle cases when no visualization could be generated
+    if not visualizations:
+        st.warning("⚠️ Unable to generate any visualization based on the suggestion.")
+
+    # Display all generated visualizations
+    for fig in visualizations:
+        st.plotly_chart(fig, use_container_width=True)
+
+
 
 # Function to create TXT file
 def create_text_report_with_viz_temp(report, conclusion, visualizations):